diff --git a/.github/workflows/github_ci_spark_3_5.yaml b/.github/workflows/github_ci_spark_3_5.yaml new file mode 100644 index 0000000..f3dd23f --- /dev/null +++ b/.github/workflows/github_ci_spark_3_5.yaml @@ -0,0 +1,31 @@ +name: Build and Push Spark 3.5 Image + +on: + workflow_dispatch: + +jobs: + build_and_push_spark_image: + name: Build spark base image + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build, tag, and push image to Docker Hub + env: + REGISTRY: ${{ secrets.DOCKERHUB_USERNAME }} + REPOSITORY: spark + IMAGE_TAG: v3.5.0-hadoop3 + run: | + cd $IMAGE_PATH + docker buildx create --name=container --driver=docker-container --use --bootstrap + ./bin/docker-image-tool.sh -r $REGISTRY -t $IMAGE_TAG -p ./kubernetes/dockerfiles/spark/bindings/python/Dockerfile -X build +env: + IMAGE_PATH : "spark-3.5.0-bin-hadoop3" \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/LICENSE b/spark-3.5.0-bin-hadoop3/LICENSE new file mode 100644 index 0000000..0564597 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/LICENSE @@ -0,0 +1,568 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. + +------------------------------------------------------------------------------------ +This project bundles some components that are also licensed under the Apache +License Version 2.0: + +commons-beanutils:commons-beanutils +org.apache.zookeeper:zookeeper +oro:oro +commons-configuration:commons-configuration +commons-digester:commons-digester +com.chuusai:shapeless_2.12 +com.googlecode.javaewah:JavaEWAH +com.twitter:chill-java +com.twitter:chill_2.12 +com.univocity:univocity-parsers +javax.jdo:jdo-api +joda-time:joda-time +net.sf.opencsv:opencsv +org.apache.derby:derby +org.objenesis:objenesis +org.roaringbitmap:RoaringBitmap +org.scalanlp:breeze-macros_2.12 +org.scalanlp:breeze_2.12 +org.typelevel:macro-compat_2.12 +org.yaml:snakeyaml +org.apache.xbean:xbean-asm7-shaded +com.squareup.okhttp3:logging-interceptor +com.squareup.okhttp3:okhttp +com.squareup.okio:okio +org.apache.spark:spark-catalyst_2.12 +org.apache.spark:spark-kvstore_2.12 +org.apache.spark:spark-launcher_2.12 +org.apache.spark:spark-mllib-local_2.12 +org.apache.spark:spark-network-common_2.12 +org.apache.spark:spark-network-shuffle_2.12 +org.apache.spark:spark-sketch_2.12 +org.apache.spark:spark-tags_2.12 +org.apache.spark:spark-unsafe_2.12 +commons-httpclient:commons-httpclient +com.vlkan:flatbuffers +com.ning:compress-lzf +io.airlift:aircompressor +io.dropwizard.metrics:metrics-core +io.dropwizard.metrics:metrics-graphite +io.dropwizard.metrics:metrics-json +io.dropwizard.metrics:metrics-jvm +io.dropwizard.metrics:metrics-jmx +org.iq80.snappy:snappy +com.clearspring.analytics:stream +com.jamesmurty.utils:java-xmlbuilder +commons-codec:commons-codec +commons-collections:commons-collections +io.fabric8:kubernetes-client +io.fabric8:kubernetes-model +io.fabric8:kubernetes-model-common +io.netty:netty-all +net.hydromatic:eigenbase-properties +net.sf.supercsv:super-csv +org.apache.arrow:arrow-format +org.apache.arrow:arrow-memory +org.apache.arrow:arrow-vector +org.apache.commons:commons-crypto +org.apache.commons:commons-lang3 +org.apache.hadoop:hadoop-annotations +org.apache.hadoop:hadoop-auth +org.apache.hadoop:hadoop-client +org.apache.hadoop:hadoop-common +org.apache.hadoop:hadoop-hdfs +org.apache.hadoop:hadoop-hdfs-client +org.apache.hadoop:hadoop-mapreduce-client-app +org.apache.hadoop:hadoop-mapreduce-client-common +org.apache.hadoop:hadoop-mapreduce-client-core +org.apache.hadoop:hadoop-mapreduce-client-jobclient +org.apache.hadoop:hadoop-mapreduce-client-shuffle +org.apache.hadoop:hadoop-yarn-api +org.apache.hadoop:hadoop-yarn-client +org.apache.hadoop:hadoop-yarn-common +org.apache.hadoop:hadoop-yarn-server-common +org.apache.hadoop:hadoop-yarn-server-web-proxy +org.apache.httpcomponents:httpclient +org.apache.httpcomponents:httpcore +org.apache.kerby:kerb-admin +org.apache.kerby:kerb-client +org.apache.kerby:kerb-common +org.apache.kerby:kerb-core +org.apache.kerby:kerb-crypto +org.apache.kerby:kerb-identity +org.apache.kerby:kerb-server +org.apache.kerby:kerb-simplekdc +org.apache.kerby:kerb-util +org.apache.kerby:kerby-asn1 +org.apache.kerby:kerby-config +org.apache.kerby:kerby-pkix +org.apache.kerby:kerby-util +org.apache.kerby:kerby-xdr +org.apache.orc:orc-core +org.apache.orc:orc-mapreduce +org.mortbay.jetty:jetty +org.mortbay.jetty:jetty-util +com.jolbox:bonecp +org.json4s:json4s-ast_2.12 +org.json4s:json4s-core_2.12 +org.json4s:json4s-jackson_2.12 +org.json4s:json4s-scalap_2.12 +com.carrotsearch:hppc +com.fasterxml.jackson.core:jackson-annotations +com.fasterxml.jackson.core:jackson-core +com.fasterxml.jackson.core:jackson-databind +com.fasterxml.jackson.dataformat:jackson-dataformat-yaml +com.fasterxml.jackson.jaxrs:jackson-jaxrs-base +com.fasterxml.jackson.jaxrs:jackson-jaxrs-json-provider +com.fasterxml.jackson.module:jackson-module-jaxb-annotations +com.fasterxml.jackson.module:jackson-module-paranamer +com.fasterxml.jackson.module:jackson-module-scala_2.12 +com.github.mifmif:generex +com.google.code.findbugs:jsr305 +com.google.code.gson:gson +com.google.flatbuffers:flatbuffers-java +com.google.guava:guava +com.google.inject:guice +com.google.inject.extensions:guice-servlet +com.twitter:parquet-hadoop-bundle +commons-cli:commons-cli +commons-dbcp:commons-dbcp +commons-io:commons-io +commons-lang:commons-lang +commons-logging:commons-logging +commons-net:commons-net +commons-pool:commons-pool +io.fabric8:zjsonpatch +javax.inject:javax.inject +javax.validation:validation-api +log4j:apache-log4j-extras +log4j:log4j +net.sf.jpam:jpam +org.apache.avro:avro +org.apache.avro:avro-ipc +org.apache.avro:avro-mapred +org.apache.commons:commons-compress +org.apache.commons:commons-math3 +org.apache.curator:curator-client +org.apache.curator:curator-framework +org.apache.curator:curator-recipes +org.apache.directory.api:api-asn1-api +org.apache.directory.api:api-util +org.apache.directory.server:apacheds-i18n +org.apache.directory.server:apacheds-kerberos-codec +org.apache.htrace:htrace-core +org.apache.ivy:ivy +org.apache.mesos:mesos +org.apache.parquet:parquet-column +org.apache.parquet:parquet-common +org.apache.parquet:parquet-encoding +org.apache.parquet:parquet-format +org.apache.parquet:parquet-hadoop +org.apache.parquet:parquet-jackson +org.apache.thrift:libfb303 +org.apache.thrift:libthrift +org.codehaus.jackson:jackson-core-asl +org.codehaus.jackson:jackson-mapper-asl +org.datanucleus:datanucleus-api-jdo +org.datanucleus:datanucleus-core +org.datanucleus:datanucleus-rdbms +org.lz4:lz4-java +org.xerial.snappy:snappy-java +stax:stax-api +xerces:xercesImpl +org.codehaus.jackson:jackson-jaxrs +org.codehaus.jackson:jackson-xc +org.eclipse.jetty:jetty-client +org.eclipse.jetty:jetty-continuation +org.eclipse.jetty:jetty-http +org.eclipse.jetty:jetty-io +org.eclipse.jetty:jetty-jndi +org.eclipse.jetty:jetty-plus +org.eclipse.jetty:jetty-proxy +org.eclipse.jetty:jetty-security +org.eclipse.jetty:jetty-server +org.eclipse.jetty:jetty-servlet +org.eclipse.jetty:jetty-servlets +org.eclipse.jetty:jetty-util +org.eclipse.jetty:jetty-webapp +org.eclipse.jetty:jetty-xml +org.scala-lang:scala-compiler +org.scala-lang:scala-library +org.scala-lang:scala-reflect +org.scala-lang.modules:scala-parser-combinators_2.12 +org.scala-lang.modules:scala-xml_2.12 +com.github.joshelser:dropwizard-metrics-hadoop-metrics2-reporter +com.zaxxer.HikariCP +org.apache.hive:hive-beeline +org.apache.hive:hive-cli +org.apache.hive:hive-common +org.apache.hive:hive-exec +org.apache.hive:hive-jdbc +org.apache.hive:hive-llap-common +org.apache.hive:hive-metastore +org.apache.hive:hive-serde +org.apache.hive:hive-service-rpc +org.apache.hive:hive-shims-0.23 +org.apache.hive:hive-shims +org.apache.hive:hive-common +org.apache.hive:hive-shims-scheduler +org.apache.hive:hive-storage-api +org.apache.hive:hive-vector-code-gen +org.datanucleus:javax.jdo +com.tdunning:json +org.apache.velocity:velocity +org.apache.yetus:audience-annotations +com.google.cloud.bigdataoss:gcs-connector + +core/src/main/java/org/apache/spark/util/collection/TimSort.java +core/src/main/resources/org/apache/spark/ui/static/bootstrap* +core/src/main/resources/org/apache/spark/ui/static/vis* +docs/js/vendor/bootstrap.js + + +------------------------------------------------------------------------------------ +This product bundles various third-party components under other open source licenses. +This section summarizes those components and their licenses. See licenses-binary/ +for text of these licenses. + + +BSD 2-Clause +------------ + +com.github.luben:zstd-jni +javolution:javolution +com.esotericsoftware:kryo-shaded +com.esotericsoftware:minlog +com.esotericsoftware:reflectasm +org.codehaus.janino:commons-compiler +org.codehaus.janino:janino +jline:jline +org.jodd:jodd-core +com.github.wendykierp:JTransforms +pl.edu.icm:JLargeArrays + + +BSD 3-Clause +------------ + +com.google.protobuf:protobuf-java +dk.brics.automaton:automaton +org.antlr:antlr-runtime +org.antlr:ST4 +org.antlr:stringtemplate +org.antlr:antlr4-runtime +antlr:antlr +com.thoughtworks.paranamer:paranamer +org.fusesource.leveldbjni:leveldbjni-all +net.sourceforge.f2j:arpack_combined_all +xmlenc:xmlenc +net.sf.py4j:py4j +org.jpmml:pmml-model +org.jpmml:pmml-schema +org.threeten:threeten-extra +org.jdom:jdom2 + +python/lib/py4j-*-src.zip +python/pyspark/cloudpickle.py +python/pyspark/join.py +core/src/main/resources/org/apache/spark/ui/static/d3.min.js + +The CSS style for the navigation sidebar of the documentation was originally +submitted by Óscar Nájera for the scikit-learn project. The scikit-learn project +is distributed under the 3-Clause BSD license. + + +MIT License +----------- + +com.microsoft.sqlserver:mssql-jdbc +org.typelevel:spire_2.12 +org.typelevel:spire-macros_2.12 +org.typelevel:spire-platform_2.12 +org.typelevel:spire-util_2.12 +org.typelevel:algebra_2.12:jar +org.typelevel:cats-kernel_2.12 +org.typelevel:machinist_2.12 +net.razorvine:pickle +org.slf4j:jcl-over-slf4j +org.slf4j:jul-to-slf4j +org.slf4j:slf4j-api +org.slf4j:slf4j-log4j12 +com.github.scopt:scopt_2.12 +dev.ludovic.netlib:blas +dev.ludovic.netlib:arpack +dev.ludovic.netlib:lapack + +core/src/main/resources/org/apache/spark/ui/static/dagre-d3.min.js +core/src/main/resources/org/apache/spark/ui/static/*dataTables* +core/src/main/resources/org/apache/spark/ui/static/graphlib-dot.min.js +core/src/main/resources/org/apache/spark/ui/static/jquery* +core/src/main/resources/org/apache/spark/ui/static/sorttable.js +docs/js/vendor/anchor.min.js +docs/js/vendor/jquery* +docs/js/vendor/modernizer* + + +Common Development and Distribution License (CDDL) 1.0 +------------------------------------------------------ + +javax.activation:activation http://www.oracle.com/technetwork/java/javase/tech/index-jsp-138795.html +javax.xml.stream:stax-api https://jcp.org/en/jsr/detail?id=173 +javax.transaction:javax.transaction-api +javax.xml.bind:jaxb-api + + +Common Development and Distribution License (CDDL) 1.1 +------------------------------------------------------ + +javax.el:javax.el-api https://javaee.github.io/uel-ri/ +javax.servlet.jsp:jsp-api +javax.transaction:jta http://www.oracle.com/technetwork/java/index.html +javax.xml.bind:jaxb-api https://github.com/javaee/jaxb-v2 +org.glassfish.hk2:hk2-api https://github.com/javaee/glassfish +org.glassfish.hk2:hk2-locator (same) +org.glassfish.hk2:hk2-utils +org.glassfish.hk2:osgi-resource-locator +org.glassfish.hk2.external:aopalliance-repackaged +org.glassfish.hk2.external:javax.inject +org.glassfish.jersey.bundles.repackaged:jersey-guava +org.glassfish.jersey.containers:jersey-container-servlet +org.glassfish.jersey.containers:jersey-container-servlet-core +org.glassfish.jersey.core:jersey-client +org.glassfish.jersey.core:jersey-common +org.glassfish.jersey.core:jersey-server +org.glassfish.jersey.media:jersey-media-jaxb + + +Eclipse Distribution License (EDL) 1.0 +-------------------------------------- + +org.glassfish.jaxb:jaxb-runtime +jakarta.activation:jakarta.activation-api +jakarta.xml.bind:jakarta.xml.bind-api +com.sun.istack:istack-commons-runtime + + +Eclipse Public License (EPL) 2.0 +-------------------------------- + +jakarta.annotation:jakarta-annotation-api https://projects.eclipse.org/projects/ee4j.ca +jakarta.servlet:jakarta.servlet-api https://projects.eclipse.org/projects/ee4j.servlet +jakarta.ws.rs:jakarta.ws.rs-api https://github.com/eclipse-ee4j/jaxrs-api +org.glassfish.hk2.external:jakarta.inject +com.github.jnr:jnr-posix + + +Public Domain +------------- + +aopalliance:aopalliance +net.iharder:base64 +org.tukaani:xz + + +Creative Commons CC0 1.0 Universal Public Domain Dedication +----------------------------------------------------------- +(see LICENSE-CC0.txt) + +data/mllib/images/kittens/29.5.a_b_EGDP022204.jpg +data/mllib/images/kittens/54893.jpg +data/mllib/images/kittens/DP153539.jpg +data/mllib/images/kittens/DP802813.jpg +data/mllib/images/multi-channel/chr30.4.184.jpg diff --git a/spark-3.5.0-bin-hadoop3/NOTICE b/spark-3.5.0-bin-hadoop3/NOTICE new file mode 100644 index 0000000..95653c6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/NOTICE @@ -0,0 +1,1533 @@ +Apache Spark +Copyright 2014 and onwards The Apache Software Foundation. + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Export Control Notice +--------------------- + +This distribution includes cryptographic software. The country in which you currently reside may have +restrictions on the import, possession, use, and/or re-export to another country, of encryption software. +BEFORE using any encryption software, please check your country's laws, regulations and policies concerning +the import, possession, or use, and re-export of encryption software, to see if this is permitted. See + for more information. + +The U.S. Government Department of Commerce, Bureau of Industry and Security (BIS), has classified this +software as Export Commodity Control Number (ECCN) 5D002.C.1, which includes information security software +using or performing cryptographic functions with asymmetric algorithms. The form and manner of this Apache +Software Foundation distribution makes it eligible for export under the License Exception ENC Technology +Software Unrestricted (TSU) exception (see the BIS Export Administration Regulations, Section 740.13) for +both object code and source code. + +The following provides more details on the included cryptographic software: + +This software uses Apache Commons Crypto (https://commons.apache.org/proper/commons-crypto/) to +support authentication, and encryption and decryption of data sent across the network between +services. + + +// ------------------------------------------------------------------ +// NOTICE file corresponding to the section 4d of The Apache License, +// Version 2.0, in this case for +// ------------------------------------------------------------------ + +Hive Beeline +Copyright 2016 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +Apache Avro +Copyright 2009-2014 The Apache Software Foundation + +This product currently only contains code developed by authors +of specific components, as identified by the source code files; +if such notes are missing files have been created by +Tatu Saloranta. + +For additional credits (generally to people who reported problems) +see CREDITS file. + +Apache Commons Compress +Copyright 2002-2012 The Apache Software Foundation + +This product includes software developed by +The Apache Software Foundation (http://www.apache.org/). + +Apache Avro Mapred API +Copyright 2009-2014 The Apache Software Foundation + +Apache Avro IPC +Copyright 2009-2014 The Apache Software Foundation + +Objenesis +Copyright 2006-2013 Joe Walnes, Henri Tremblay, Leonardo Mesquita + +Apache XBean :: ASM shaded (repackaged) +Copyright 2005-2019 The Apache Software Foundation + +-------------------------------------- + +This product includes software developed at +OW2 Consortium (http://asm.ow2.org/) + +This product includes software developed by The Apache Software +Foundation (http://www.apache.org/). + +The binary distribution of this product bundles binaries of +org.iq80.leveldb:leveldb-api (https://github.com/dain/leveldb), which has the +following notices: +* Copyright 2011 Dain Sundstrom +* Copyright 2011 FuseSource Corp. http://fusesource.com + +The binary distribution of this product bundles binaries of +org.fusesource.hawtjni:hawtjni-runtime (https://github.com/fusesource/hawtjni), +which has the following notices: +* This product includes software developed by FuseSource Corp. + http://fusesource.com +* This product includes software developed at + Progress Software Corporation and/or its subsidiaries or affiliates. +* This product includes software developed by IBM Corporation and others. + +The binary distribution of this product bundles binaries of +Gson 2.2.4, +which has the following notices: + + + The Netty Project + ================= + +Please visit the Netty web site for more information: + + * http://netty.io/ + +Copyright 2014 The Netty Project + +The Netty Project 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. + +Also, please refer to each LICENSE..txt file, which is located in +the 'license' directory of the distribution file, for the license terms of the +components that this product depends on. + +------------------------------------------------------------------------------- +This product contains the extensions to Java Collections Framework which has +been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene: + + * LICENSE: + * license/LICENSE.jsr166y.txt (Public Domain) + * HOMEPAGE: + * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/ + * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/ + +This product contains a modified version of Robert Harder's Public Domain +Base64 Encoder and Decoder, which can be obtained at: + + * LICENSE: + * license/LICENSE.base64.txt (Public Domain) + * HOMEPAGE: + * http://iharder.sourceforge.net/current/java/base64/ + +This product contains a modified portion of 'Webbit', an event based +WebSocket and HTTP server, which can be obtained at: + + * LICENSE: + * license/LICENSE.webbit.txt (BSD License) + * HOMEPAGE: + * https://github.com/joewalnes/webbit + +This product contains a modified portion of 'SLF4J', a simple logging +facade for Java, which can be obtained at: + + * LICENSE: + * license/LICENSE.slf4j.txt (MIT License) + * HOMEPAGE: + * http://www.slf4j.org/ + +This product contains a modified portion of 'Apache Harmony', an open source +Java SE, which can be obtained at: + + * NOTICE: + * license/NOTICE.harmony.txt + * LICENSE: + * license/LICENSE.harmony.txt (Apache License 2.0) + * HOMEPAGE: + * http://archive.apache.org/dist/harmony/ + +This product contains a modified portion of 'jbzip2', a Java bzip2 compression +and decompression library written by Matthew J. Francis. It can be obtained at: + + * LICENSE: + * license/LICENSE.jbzip2.txt (MIT License) + * HOMEPAGE: + * https://code.google.com/p/jbzip2/ + +This product contains a modified portion of 'libdivsufsort', a C API library to construct +the suffix array and the Burrows-Wheeler transformed string for any input string of +a constant-size alphabet written by Yuta Mori. It can be obtained at: + + * LICENSE: + * license/LICENSE.libdivsufsort.txt (MIT License) + * HOMEPAGE: + * https://github.com/y-256/libdivsufsort + +This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM, + which can be obtained at: + + * LICENSE: + * license/LICENSE.jctools.txt (ASL2 License) + * HOMEPAGE: + * https://github.com/JCTools/JCTools + +This product optionally depends on 'JZlib', a re-implementation of zlib in +pure Java, which can be obtained at: + + * LICENSE: + * license/LICENSE.jzlib.txt (BSD style License) + * HOMEPAGE: + * http://www.jcraft.com/jzlib/ + +This product optionally depends on 'Compress-LZF', a Java library for encoding and +decoding data in LZF format, written by Tatu Saloranta. It can be obtained at: + + * LICENSE: + * license/LICENSE.compress-lzf.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/ning/compress + +This product optionally depends on 'lz4', a LZ4 Java compression +and decompression library written by Adrien Grand. It can be obtained at: + + * LICENSE: + * license/LICENSE.lz4.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/jpountz/lz4-java + +This product optionally depends on 'lzma-java', a LZMA Java compression +and decompression library, which can be obtained at: + + * LICENSE: + * license/LICENSE.lzma-java.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/jponge/lzma-java + +This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression +and decompression library written by William Kinney. It can be obtained at: + + * LICENSE: + * license/LICENSE.jfastlz.txt (MIT License) + * HOMEPAGE: + * https://code.google.com/p/jfastlz/ + +This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data +interchange format, which can be obtained at: + + * LICENSE: + * license/LICENSE.protobuf.txt (New BSD License) + * HOMEPAGE: + * https://github.com/google/protobuf + +This product optionally depends on 'Bouncy Castle Crypto APIs' to generate +a temporary self-signed X.509 certificate when the JVM does not provide the +equivalent functionality. It can be obtained at: + + * LICENSE: + * license/LICENSE.bouncycastle.txt (MIT License) + * HOMEPAGE: + * http://www.bouncycastle.org/ + +This product optionally depends on 'Snappy', a compression library produced +by Google Inc, which can be obtained at: + + * LICENSE: + * license/LICENSE.snappy.txt (New BSD License) + * HOMEPAGE: + * https://github.com/google/snappy + +This product optionally depends on 'JBoss Marshalling', an alternative Java +serialization API, which can be obtained at: + + * LICENSE: + * license/LICENSE.jboss-marshalling.txt (GNU LGPL 2.1) + * HOMEPAGE: + * http://www.jboss.org/jbossmarshalling + +This product optionally depends on 'Caliper', Google's micro- +benchmarking framework, which can be obtained at: + + * LICENSE: + * license/LICENSE.caliper.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/google/caliper + +This product optionally depends on 'Apache Commons Logging', a logging +framework, which can be obtained at: + + * LICENSE: + * license/LICENSE.commons-logging.txt (Apache License 2.0) + * HOMEPAGE: + * http://commons.apache.org/logging/ + +This product optionally depends on 'Apache Log4J', a logging framework, which +can be obtained at: + + * LICENSE: + * license/LICENSE.log4j.txt (Apache License 2.0) + * HOMEPAGE: + * http://logging.apache.org/log4j/ + +This product optionally depends on 'Aalto XML', an ultra-high performance +non-blocking XML processor, which can be obtained at: + + * LICENSE: + * license/LICENSE.aalto-xml.txt (Apache License 2.0) + * HOMEPAGE: + * http://wiki.fasterxml.com/AaltoHome + +This product contains a modified version of 'HPACK', a Java implementation of +the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at: + + * LICENSE: + * license/LICENSE.hpack.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/twitter/hpack + +This product contains a modified portion of 'Apache Commons Lang', a Java library +provides utilities for the java.lang API, which can be obtained at: + + * LICENSE: + * license/LICENSE.commons-lang.txt (Apache License 2.0) + * HOMEPAGE: + * https://commons.apache.org/proper/commons-lang/ + + +This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build. + + * LICENSE: + * license/LICENSE.mvn-wrapper.txt (Apache License 2.0) + * HOMEPAGE: + * https://github.com/takari/maven-wrapper + + +The binary distribution of this product bundles binaries of +Commons Codec 1.4, +which has the following notices: + * src/test/org/apache/commons/codec/language/DoubleMetaphoneTest.javacontains test data from http://aspell.net/test/orig/batch0.tab.Copyright (C) 2002 Kevin Atkinson (kevina@gnu.org) + =============================================================================== + The content of package org.apache.commons.codec.language.bm has been translated + from the original php source code available at http://stevemorse.org/phoneticinfo.htm + with permission from the original authors. + Original source copyright:Copyright (c) 2008 Alexander Beider & Stephen P. Morse. + +The binary distribution of this product bundles binaries of +Commons Lang 2.6, +which has the following notices: + * This product includes software from the Spring Framework,under the Apache License 2.0 (see: StringUtils.containsWhitespace()) + +The binary distribution of this product bundles binaries of +Apache Log4j 1.2.17, +which has the following notices: + * ResolverUtil.java + Copyright 2005-2006 Tim Fennell + Dumbster SMTP test server + Copyright 2004 Jason Paul Kitchen + TypeUtil.java + Copyright 2002-2012 Ramnivas Laddad, Juergen Hoeller, Chris Beams + +The binary distribution of this product bundles binaries of +Jetty 6.1.26, +which has the following notices: + * ============================================================== + Jetty Web Container + Copyright 1995-2016 Mort Bay Consulting Pty Ltd. + ============================================================== + + The Jetty Web Container is Copyright Mort Bay Consulting Pty Ltd + unless otherwise noted. + + Jetty is dual licensed under both + + * The Apache 2.0 License + http://www.apache.org/licenses/LICENSE-2.0.html + + and + + * The Eclipse Public 1.0 License + http://www.eclipse.org/legal/epl-v10.html + + Jetty may be distributed under either license. + + ------ + Eclipse + + The following artifacts are EPL. + * org.eclipse.jetty.orbit:org.eclipse.jdt.core + + The following artifacts are EPL and ASL2. + * org.eclipse.jetty.orbit:javax.security.auth.message + + The following artifacts are EPL and CDDL 1.0. + * org.eclipse.jetty.orbit:javax.mail.glassfish + + ------ + Oracle + + The following artifacts are CDDL + GPLv2 with classpath exception. + https://glassfish.dev.java.net/nonav/public/CDDL+GPL.html + + * javax.servlet:javax.servlet-api + * javax.annotation:javax.annotation-api + * javax.transaction:javax.transaction-api + * javax.websocket:javax.websocket-api + + ------ + Oracle OpenJDK + + If ALPN is used to negotiate HTTP/2 connections, then the following + artifacts may be included in the distribution or downloaded when ALPN + module is selected. + + * java.sun.security.ssl + + These artifacts replace/modify OpenJDK classes. The modififications + are hosted at github and both modified and original are under GPL v2 with + classpath exceptions. + http://openjdk.java.net/legal/gplv2+ce.html + + ------ + OW2 + + The following artifacts are licensed by the OW2 Foundation according to the + terms of http://asm.ow2.org/license.html + + org.ow2.asm:asm-commons + org.ow2.asm:asm + + ------ + Apache + + The following artifacts are ASL2 licensed. + + org.apache.taglibs:taglibs-standard-spec + org.apache.taglibs:taglibs-standard-impl + + ------ + MortBay + + The following artifacts are ASL2 licensed. Based on selected classes from + following Apache Tomcat jars, all ASL2 licensed. + + org.mortbay.jasper:apache-jsp + org.apache.tomcat:tomcat-jasper + org.apache.tomcat:tomcat-juli + org.apache.tomcat:tomcat-jsp-api + org.apache.tomcat:tomcat-el-api + org.apache.tomcat:tomcat-jasper-el + org.apache.tomcat:tomcat-api + org.apache.tomcat:tomcat-util-scan + org.apache.tomcat:tomcat-util + + org.mortbay.jasper:apache-el + org.apache.tomcat:tomcat-jasper-el + org.apache.tomcat:tomcat-el-api + + ------ + Mortbay + + The following artifacts are CDDL + GPLv2 with classpath exception. + + https://glassfish.dev.java.net/nonav/public/CDDL+GPL.html + + org.eclipse.jetty.toolchain:jetty-schemas + + ------ + Assorted + + The UnixCrypt.java code implements the one way cryptography used by + Unix systems for simple password protection. Copyright 1996 Aki Yoshida, + modified April 2001 by Iris Van den Broeke, Daniel Deville. + Permission to use, copy, modify and distribute UnixCrypt + for non-commercial or commercial purposes and without fee is + granted provided that the copyright notice appears in all copies./ + +The binary distribution of this product bundles binaries of +Snappy for Java 1.0.4.1, +which has the following notices: + * This product includes software developed by Google + Snappy: http://code.google.com/p/snappy/ (New BSD License) + + This product includes software developed by Apache + PureJavaCrc32C from apache-hadoop-common http://hadoop.apache.org/ + (Apache 2.0 license) + + This library contains statically linked libstdc++. This inclusion is allowed by + "GCC RUntime Library Exception" + http://gcc.gnu.org/onlinedocs/libstdc++/manual/license.html + + == Contributors == + * Tatu Saloranta + * Providing benchmark suite + * Alec Wysoker + * Performance and memory usage improvement + +The binary distribution of this product bundles binaries of +Xerces2 Java Parser 2.9.1, +which has the following notices: + * ========================================================================= + == NOTICE file corresponding to section 4(d) of the Apache License, == + == Version 2.0, in this case for the Apache Xerces Java distribution. == + ========================================================================= + + Apache Xerces Java + Copyright 1999-2007 The Apache Software Foundation + + This product includes software developed at + The Apache Software Foundation (http://www.apache.org/). + + Portions of this software were originally based on the following: + - software copyright (c) 1999, IBM Corporation., http://www.ibm.com. + - software copyright (c) 1999, Sun Microsystems., http://www.sun.com. + - voluntary contributions made by Paul Eng on behalf of the + Apache Software Foundation that were originally developed at iClick, Inc., + software copyright (c) 1999. + +Apache Commons Collections +Copyright 2001-2015 The Apache Software Foundation + +Apache Commons Configuration +Copyright 2001-2008 The Apache Software Foundation + +Apache Jakarta Commons Digester +Copyright 2001-2006 The Apache Software Foundation + +Apache Commons BeanUtils +Copyright 2000-2008 The Apache Software Foundation + +ApacheDS Protocol Kerberos Codec +Copyright 2003-2013 The Apache Software Foundation + +ApacheDS I18n +Copyright 2003-2013 The Apache Software Foundation + +Apache Directory API ASN.1 API +Copyright 2003-2013 The Apache Software Foundation + +Apache Directory LDAP API Utilities +Copyright 2003-2013 The Apache Software Foundation + +Curator Client +Copyright 2011-2015 The Apache Software Foundation + +htrace-core +Copyright 2015 The Apache Software Foundation + + ========================================================================= + == NOTICE file corresponding to section 4(d) of the Apache License, == + == Version 2.0, in this case for the Apache Xerces Java distribution. == + ========================================================================= + + Portions of this software were originally based on the following: + - software copyright (c) 1999, IBM Corporation., http://www.ibm.com. + - software copyright (c) 1999, Sun Microsystems., http://www.sun.com. + - voluntary contributions made by Paul Eng on behalf of the + Apache Software Foundation that were originally developed at iClick, Inc., + software copyright (c) 1999. + +# Jackson JSON processor + +Jackson is a high-performance, Free/Open Source JSON processing library. +It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has +been in development since 2007. +It is currently developed by a community of developers, as well as supported +commercially by FasterXML.com. + +## Licensing + +Jackson core and extension components may licensed under different licenses. +To find the details that apply to this artifact see the accompanying LICENSE file. +For more information, including possible other licensing options, contact +FasterXML.com (http://fasterxml.com). + +## Credits + +A list of contributors may be found from CREDITS file, which is included +in some artifacts (usually source distributions); but is always available +from the source code management (SCM) system project uses. + +Apache HttpCore +Copyright 2005-2017 The Apache Software Foundation + +Curator Recipes +Copyright 2011-2015 The Apache Software Foundation + +Curator Framework +Copyright 2011-2015 The Apache Software Foundation + +Apache Commons Lang +Copyright 2001-2016 The Apache Software Foundation + +This product includes software from the Spring Framework, +under the Apache License 2.0 (see: StringUtils.containsWhitespace()) + +Apache Commons Math +Copyright 2001-2015 The Apache Software Foundation + +This product includes software developed for Orekit by +CS Systèmes d'Information (http://www.c-s.fr/) +Copyright 2010-2012 CS Systèmes d'Information + +Apache log4j +Copyright 2007 The Apache Software Foundation + +# Compress LZF + +This library contains efficient implementation of LZF compression format, +as well as additional helper classes that build on JDK-provided gzip (deflat) +codec. + +Library is licensed under Apache License 2.0, as per accompanying LICENSE file. + +## Credit + +Library has been written by Tatu Saloranta (tatu.saloranta@iki.fi). +It was started at Ning, inc., as an official Open Source process used by +platform backend, but after initial versions has been developed outside of +Ning by supporting community. + +Other contributors include: + +* Jon Hartlaub (first versions of streaming reader/writer; unit tests) +* Cedrik Lime: parallel LZF implementation + +Various community members have contributed bug reports, and suggested minor +fixes; these can be found from file "VERSION.txt" in SCM. + +Apache Commons Net +Copyright 2001-2012 The Apache Software Foundation + + +Jackson core and extension components may be licensed under different licenses. +To find the details that apply to this artifact see the accompanying LICENSE file. +For more information, including possible other licensing options, contact +FasterXML.com (http://fasterxml.com). + +Apache Ivy (TM) +Copyright 2007-2014 The Apache Software Foundation + +Portions of Ivy were originally developed at +Jayasoft SARL (http://www.jayasoft.fr/) +and are licensed to the Apache Software Foundation under the +"Software Grant License Agreement" + +SSH and SFTP support is provided by the JCraft JSch package, +which is open source software, available under +the terms of a BSD style license. +The original software and related information is available +at http://www.jcraft.com/jsch/. + + +ORC Core +Copyright 2013-2018 The Apache Software Foundation + +Apache Commons Lang +Copyright 2001-2011 The Apache Software Foundation + +ORC MapReduce +Copyright 2013-2018 The Apache Software Foundation + +Apache Parquet Format +Copyright 2017 The Apache Software Foundation + +Arrow Vectors +Copyright 2017 The Apache Software Foundation + +Arrow Format +Copyright 2017 The Apache Software Foundation + +Arrow Memory +Copyright 2017 The Apache Software Foundation + +Apache Commons CLI +Copyright 2001-2009 The Apache Software Foundation + +Apache Commons Daemon +Copyright 1999-2019 The Apache Software Foundation + +Google Guice - Extensions - Servlet +Copyright 2006-2011 Google, Inc. + +Apache Commons IO +Copyright 2002-2012 The Apache Software Foundation + +Google Guice - Core Library +Copyright 2006-2011 Google, Inc. + +mesos +Copyright 2017 The Apache Software Foundation + +Apache Parquet Hadoop Bundle (Incubating) +Copyright 2015 The Apache Software Foundation + +Hive Query Language +Copyright 2016 The Apache Software Foundation + +Apache Extras Companion for log4j 1.2. +Copyright 2007 The Apache Software Foundation + +Hive Metastore +Copyright 2016 The Apache Software Foundation + +Apache Commons Logging +Copyright 2003-2013 The Apache Software Foundation + +========================================================================= +== NOTICE file corresponding to section 4(d) of the Apache License, == +== Version 2.0, in this case for the DataNucleus distribution. == +========================================================================= + +=================================================================== +This product includes software developed by many individuals, +including the following: +=================================================================== +Erik Bengtson +Andy Jefferson + +=================================================================== +This product has included contributions from some individuals, +including the following: +=================================================================== + +=================================================================== +This product includes software developed by many individuals, +including the following: +=================================================================== +Andy Jefferson +Erik Bengtson +Joerg von Frantzius +Marco Schulze + +=================================================================== +This product has included contributions from some individuals, +including the following: +=================================================================== +Barry Haddow +Ralph Ullrich +David Ezzio +Brendan de Beer +David Eaves +Martin Taal +Tony Lai +Roland Szabo +Anton Troshin (Timesten) + +=================================================================== +This product also includes software developed by the TJDO project +(http://tjdo.sourceforge.net/). +=================================================================== + +=================================================================== +This product also includes software developed by the Apache Commons project +(http://commons.apache.org/). +=================================================================== + +Apache Commons Pool +Copyright 1999-2009 The Apache Software Foundation + +Apache Commons DBCP +Copyright 2001-2010 The Apache Software Foundation + +Apache Java Data Objects (JDO) +Copyright 2005-2006 The Apache Software Foundation + +Apache Jakarta HttpClient +Copyright 1999-2007 The Apache Software Foundation + +Apache HttpClient +Copyright 1999-2017 The Apache Software Foundation + +Apache Commons Codec +Copyright 2002-2014 The Apache Software Foundation + +src/test/org/apache/commons/codec/language/DoubleMetaphoneTest.java +contains test data from http://aspell.net/test/orig/batch0.tab. +Copyright (C) 2002 Kevin Atkinson (kevina@gnu.org) + +=============================================================================== + +The content of package org.apache.commons.codec.language.bm has been translated +from the original php source code available at http://stevemorse.org/phoneticinfo.htm +with permission from the original authors. +Original source copyright: +Copyright (c) 2008 Alexander Beider & Stephen P. Morse. + +============================================================================= += NOTICE file corresponding to section 4d of the Apache License Version 2.0 = +============================================================================= +This product includes software developed by +Joda.org (http://www.joda.org/). + +=================================================================== +This product has included contributions from some individuals, +including the following: +=================================================================== +Joerg von Frantzius +Thomas Marti +Barry Haddow +Marco Schulze +Ralph Ullrich +David Ezzio +Brendan de Beer +David Eaves +Martin Taal +Tony Lai +Roland Szabo +Marcus Mennemeier +Xuan Baldauf +Eric Sultan + +Apache Thrift +Copyright 2006-2010 The Apache Software Foundation. + +========================================================================= +== NOTICE file corresponding to section 4(d) of the Apache License, +== Version 2.0, in this case for the Apache Derby distribution. +== +== DO NOT EDIT THIS FILE DIRECTLY. IT IS GENERATED +== BY THE buildnotice TARGET IN THE TOP LEVEL build.xml FILE. +== +========================================================================= + +Apache Derby +Copyright 2004-2015 The Apache Software Foundation + +========================================================================= + +Portions of Derby were originally developed by +International Business Machines Corporation and are +licensed to the Apache Software Foundation under the +"Software Grant and Corporate Contribution License Agreement", +informally known as the "Derby CLA". +The following copyright notice(s) were affixed to portions of the code +with which this file is now or was at one time distributed +and are placed here unaltered. + +(C) Copyright 1997,2004 International Business Machines Corporation. All rights reserved. + +(C) Copyright IBM Corp. 2003. + +The portion of the functionTests under 'nist' was originally +developed by the National Institute of Standards and Technology (NIST), +an agency of the United States Department of Commerce, and adapted by +International Business Machines Corporation in accordance with the NIST +Software Acknowledgment and Redistribution document at +http://www.itl.nist.gov/div897/ctg/sql_form.htm + +The JDBC apis for small devices and JDBC3 (under java/stubs/jsr169 and +java/stubs/jdbc3) were produced by trimming sources supplied by the +Apache Harmony project. In addition, the Harmony SerialBlob and +SerialClob implementations are used. The following notice covers the Harmony sources: + +Portions of Harmony were originally developed by +Intel Corporation and are licensed to the Apache Software +Foundation under the "Software Grant and Corporate Contribution +License Agreement", informally known as the "Intel Harmony CLA". + +The Derby build relies on source files supplied by the Apache Felix +project. The following notice covers the Felix files: + + Apache Felix Main + Copyright 2008 The Apache Software Foundation + + I. Included Software + + This product includes software developed at + The Apache Software Foundation (http://www.apache.org/). + Licensed under the Apache License 2.0. + + This product includes software developed at + The OSGi Alliance (http://www.osgi.org/). + Copyright (c) OSGi Alliance (2000, 2007). + Licensed under the Apache License 2.0. + + This product includes software from http://kxml.sourceforge.net. + Copyright (c) 2002,2003, Stefan Haustein, Oberhausen, Rhld., Germany. + Licensed under BSD License. + + II. Used Software + + This product uses software developed at + The OSGi Alliance (http://www.osgi.org/). + Copyright (c) OSGi Alliance (2000, 2007). + Licensed under the Apache License 2.0. + + III. License Summary + - Apache License 2.0 + - BSD License + +The Derby build relies on jar files supplied by the Apache Lucene +project. The following notice covers the Lucene files: + +Apache Lucene +Copyright 2013 The Apache Software Foundation + +Includes software from other Apache Software Foundation projects, +including, but not limited to: + - Apache Ant + - Apache Jakarta Regexp + - Apache Commons + - Apache Xerces + +ICU4J, (under analysis/icu) is licensed under an MIT styles license +and Copyright (c) 1995-2008 International Business Machines Corporation and others + +Some data files (under analysis/icu/src/data) are derived from Unicode data such +as the Unicode Character Database. See http://unicode.org/copyright.html for more +details. + +Brics Automaton (under core/src/java/org/apache/lucene/util/automaton) is +BSD-licensed, created by Anders Møller. See http://www.brics.dk/automaton/ + +The levenshtein automata tables (under core/src/java/org/apache/lucene/util/automaton) were +automatically generated with the moman/finenight FSA library, created by +Jean-Philippe Barrette-LaPierre. This library is available under an MIT license, +see http://sites.google.com/site/rrettesite/moman and +http://bitbucket.org/jpbarrette/moman/overview/ + +The class org.apache.lucene.util.WeakIdentityMap was derived from +the Apache CXF project and is Apache License 2.0. + +The Google Code Prettify is Apache License 2.0. +See http://code.google.com/p/google-code-prettify/ + +JUnit (junit-4.10) is licensed under the Common Public License v. 1.0 +See http://junit.sourceforge.net/cpl-v10.html + +This product includes code (JaspellTernarySearchTrie) from Java Spelling Checkin +g Package (jaspell): http://jaspell.sourceforge.net/ +License: The BSD License (http://www.opensource.org/licenses/bsd-license.php) + +This product includes software developed by the JDOM Project (http://www.jdom.org/) +License: https://raw.githubusercontent.com/hunterhacker/jdom/master/LICENSE.txt + +The snowball stemmers in + analysis/common/src/java/net/sf/snowball +were developed by Martin Porter and Richard Boulton. +The snowball stopword lists in + analysis/common/src/resources/org/apache/lucene/analysis/snowball +were developed by Martin Porter and Richard Boulton. +The full snowball package is available from + http://snowball.tartarus.org/ + +The KStem stemmer in + analysis/common/src/org/apache/lucene/analysis/en +was developed by Bob Krovetz and Sergio Guzman-Lara (CIIR-UMass Amherst) +under the BSD-license. + +The Arabic,Persian,Romanian,Bulgarian, and Hindi analyzers (common) come with a default +stopword list that is BSD-licensed created by Jacques Savoy. These files reside in: +analysis/common/src/resources/org/apache/lucene/analysis/ar/stopwords.txt, +analysis/common/src/resources/org/apache/lucene/analysis/fa/stopwords.txt, +analysis/common/src/resources/org/apache/lucene/analysis/ro/stopwords.txt, +analysis/common/src/resources/org/apache/lucene/analysis/bg/stopwords.txt, +analysis/common/src/resources/org/apache/lucene/analysis/hi/stopwords.txt +See http://members.unine.ch/jacques.savoy/clef/index.html. + +The German,Spanish,Finnish,French,Hungarian,Italian,Portuguese,Russian and Swedish light stemmers +(common) are based on BSD-licensed reference implementations created by Jacques Savoy and +Ljiljana Dolamic. These files reside in: +analysis/common/src/java/org/apache/lucene/analysis/de/GermanLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/de/GermanMinimalStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/es/SpanishLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/fi/FinnishLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/fr/FrenchLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/fr/FrenchMinimalStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/hu/HungarianLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/it/ItalianLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/pt/PortugueseLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/ru/RussianLightStemmer.java +analysis/common/src/java/org/apache/lucene/analysis/sv/SwedishLightStemmer.java + +The Stempel analyzer (stempel) includes BSD-licensed software developed +by the Egothor project http://egothor.sf.net/, created by Leo Galambos, Martin Kvapil, +and Edmond Nolan. + +The Polish analyzer (stempel) comes with a default +stopword list that is BSD-licensed created by the Carrot2 project. The file resides +in stempel/src/resources/org/apache/lucene/analysis/pl/stopwords.txt. +See http://project.carrot2.org/license.html. + +The SmartChineseAnalyzer source code (smartcn) was +provided by Xiaoping Gao and copyright 2009 by www.imdict.net. + +WordBreakTestUnicode_*.java (under modules/analysis/common/src/test/) +is derived from Unicode data such as the Unicode Character Database. +See http://unicode.org/copyright.html for more details. + +The Morfologik analyzer (morfologik) includes BSD-licensed software +developed by Dawid Weiss and Marcin Miłkowski (http://morfologik.blogspot.com/). + +Morfologik uses data from Polish ispell/myspell dictionary +(http://www.sjp.pl/slownik/en/) licenced on the terms of (inter alia) +LGPL and Creative Commons ShareAlike. + +Morfologic includes data from BSD-licensed dictionary of Polish (SGJP) +(http://sgjp.pl/morfeusz/) + +Servlet-api.jar and javax.servlet-*.jar are under the CDDL license, the original +source code for this can be found at http://www.eclipse.org/jetty/downloads.php + +=========================================================================== +Kuromoji Japanese Morphological Analyzer - Apache Lucene Integration +=========================================================================== + +This software includes a binary and/or source version of data from + + mecab-ipadic-2.7.0-20070801 + +which can be obtained from + + http://atilika.com/releases/mecab-ipadic/mecab-ipadic-2.7.0-20070801.tar.gz + +or + + http://jaist.dl.sourceforge.net/project/mecab/mecab-ipadic/2.7.0-20070801/mecab-ipadic-2.7.0-20070801.tar.gz + +=========================================================================== +mecab-ipadic-2.7.0-20070801 Notice +=========================================================================== + +Nara Institute of Science and Technology (NAIST), +the copyright holders, disclaims all warranties with regard to this +software, including all implied warranties of merchantability and +fitness, in no event shall NAIST be liable for +any special, indirect or consequential damages or any damages +whatsoever resulting from loss of use, data or profits, whether in an +action of contract, negligence or other tortuous action, arising out +of or in connection with the use or performance of this software. + +A large portion of the dictionary entries +originate from ICOT Free Software. The following conditions for ICOT +Free Software applies to the current dictionary as well. + +Each User may also freely distribute the Program, whether in its +original form or modified, to any third party or parties, PROVIDED +that the provisions of Section 3 ("NO WARRANTY") will ALWAYS appear +on, or be attached to, the Program, which is distributed substantially +in the same form as set out herein and that such intended +distribution, if actually made, will neither violate or otherwise +contravene any of the laws and regulations of the countries having +jurisdiction over the User or the intended distribution itself. + +NO WARRANTY + +The program was produced on an experimental basis in the course of the +research and development conducted during the project and is provided +to users as so produced on an experimental basis. Accordingly, the +program is provided without any warranty whatsoever, whether express, +implied, statutory or otherwise. The term "warranty" used herein +includes, but is not limited to, any warranty of the quality, +performance, merchantability and fitness for a particular purpose of +the program and the nonexistence of any infringement or violation of +any right of any third party. + +Each user of the program will agree and understand, and be deemed to +have agreed and understood, that there is no warranty whatsoever for +the program and, accordingly, the entire risk arising from or +otherwise connected with the program is assumed by the user. + +Therefore, neither ICOT, the copyright holder, or any other +organization that participated in or was otherwise related to the +development of the program and their respective officials, directors, +officers and other employees shall be held liable for any and all +damages, including, without limitation, general, special, incidental +and consequential damages, arising out of or otherwise in connection +with the use or inability to use the program or any product, material +or result produced or otherwise obtained by using the program, +regardless of whether they have been advised of, or otherwise had +knowledge of, the possibility of such damages at any time during the +project or thereafter. Each user will be deemed to have agreed to the +foregoing by his or her commencement of use of the program. The term +"use" as used herein includes, but is not limited to, the use, +modification, copying and distribution of the program and the +production of secondary products from the program. + +In the case where the program, whether in its original form or +modified, was distributed or delivered to or received by a user from +any person, organization or entity other than ICOT, unless it makes or +grants independently of ICOT any specific warranty to the user in +writing, such person, organization or entity, will also be exempted +from and not be held liable to the user for any such damages as noted +above as far as the program is concerned. + +The Derby build relies on a jar file supplied by the JSON Simple +project, hosted at https://code.google.com/p/json-simple/. +The JSON simple jar file is licensed under the Apache 2.0 License. + +Hive CLI +Copyright 2016 The Apache Software Foundation + +Hive JDBC +Copyright 2016 The Apache Software Foundation + + +Chill is a set of Scala extensions for Kryo. +Copyright 2012 Twitter, Inc. + +Third Party Dependencies: + +Kryo 2.17 +BSD 3-Clause License +http://code.google.com/p/kryo + +Commons-Codec 1.7 +Apache Public License 2.0 +http://hadoop.apache.org + + + +Breeze is distributed under an Apache License V2.0 (See LICENSE) + +=============================================================================== + +Proximal algorithms outlined in Proximal.scala (package breeze.optimize.proximal) +are based on https://github.com/cvxgrp/proximal (see LICENSE for details) and distributed with +Copyright (c) 2014 by Debasish Das (Verizon), all rights reserved. + +=============================================================================== + +QuadraticMinimizer class in package breeze.optimize.proximal is distributed with Copyright (c) +2014, Debasish Das (Verizon), all rights reserved. + +=============================================================================== + +NonlinearMinimizer class in package breeze.optimize.proximal is distributed with Copyright (c) +2015, Debasish Das (Verizon), all rights reserved. + + +stream-lib +Copyright 2016 AddThis + +This product includes software developed by AddThis. + +This product also includes code adapted from: + +Apache Solr (http://lucene.apache.org/solr/) +Copyright 2014 The Apache Software Foundation + +Apache Mahout (http://mahout.apache.org/) +Copyright 2014 The Apache Software Foundation + +scala-xml +Copyright (c) 2002-2019 EPFL +Copyright (c) 2011-2019 Lightbend, Inc. + +scala-xml includes software developed at +LAMP/EPFL (https://lamp.epfl.ch/) and +Lightbend, Inc. (https://www.lightbend.com/). + +Licensed under the Apache License, Version 2.0 (the "License"). +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. + +dropwizard-metrics-hadoop-metrics2-reporter +Copyright 2016 Josh Elser + +Hive Beeline +Copyright 2019 The Apache Software Foundation + +Hive CLI +Copyright 2019 The Apache Software Foundation + +Hive Common +Copyright 2019 The Apache Software Foundation + +Hive JDBC +Copyright 2019 The Apache Software Foundation + +Hive Query Language +Copyright 2019 The Apache Software Foundation + +Hive Llap Common +Copyright 2019 The Apache Software Foundation + +Hive Metastore +Copyright 2019 The Apache Software Foundation + +Hive Serde +Copyright 2019 The Apache Software Foundation + +Hive Service RPC +Copyright 2019 The Apache Software Foundation + +Hive Shims +Copyright 2019 The Apache Software Foundation + +Hive Shims 0.23 +Copyright 2019 The Apache Software Foundation + +Hive Shims Common +Copyright 2019 The Apache Software Foundation + +Hive Shims Scheduler +Copyright 2019 The Apache Software Foundation + +Hive Storage API +Copyright 2018 The Apache Software Foundation + +Hive Vector-Code-Gen Utilities +Copyright 2019 The Apache Software Foundation + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + Copyright 2015-2015 DataNucleus + + Licensed 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. + +Android JSON library +Copyright (C) 2010 The Android Open Source Project + +This product includes software developed by +The Android Open Source Project + +Apache Velocity + +Copyright (C) 2000-2007 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +Apache Yetus - Audience Annotations +Copyright 2015-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +Ehcache V3 +Copyright 2014-2016 Terracotta, Inc. + +The product includes software from the Apache Commons Lang project, +under the Apache License 2.0 (see: org.ehcache.impl.internal.classes.commonslang) + +Apache Geronimo JCache Spec 1.0 +Copyright 2003-2014 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Admin +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Client +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Common +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb core +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Crypto +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Identity +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Server +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerb Simple Kdc +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby-kerb Util +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby ASN1 Project +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby Config +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby PKIX Project +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby Util +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Kerby XDR Project +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Token provider +Copyright 2014-2017 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + + +Metrics +Copyright 2010-2013 Coda Hale and Yammer, Inc. + +This product includes software developed by Coda Hale and Yammer, Inc. + +This product includes code derived from the JSR-166 project (ThreadLocalRandom, Striped64, +LongAdder), which was released with the following comments: + + Written by Doug Lea with assistance from members of JCP JSR-166 + Expert Group and released to the public domain, as explained at + http://creativecommons.org/publicdomain/zero/1.0/ \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/DESCRIPTION b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/DESCRIPTION new file mode 100644 index 0000000..9b7fb88 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/DESCRIPTION @@ -0,0 +1,32 @@ +Package: SparkR +Type: Package +Version: 3.5.0 +Title: R Front End for 'Apache Spark' +Description: Provides an R Front end for 'Apache Spark' . +Authors@R: + person(family = "The Apache Software Foundation", + email = "dev@spark.apache.org", + role = c("aut", "cre", "cph")) +License: Apache License (== 2.0) +URL: https://www.apache.org https://spark.apache.org +BugReports: https://spark.apache.org/contributing.html +SystemRequirements: Java (>= 8, < 18) +Depends: R (>= 3.5), methods +Suggests: knitr, rmarkdown, markdown, testthat, e1071, survival, arrow + (>= 1.0.0) +Collate: 'schema.R' 'generics.R' 'jobj.R' 'column.R' 'group.R' 'RDD.R' + 'pairRDD.R' 'DataFrame.R' 'SQLContext.R' 'WindowSpec.R' + 'backend.R' 'broadcast.R' 'catalog.R' 'client.R' 'context.R' + 'deserialize.R' 'functions.R' 'install.R' 'jvm.R' + 'mllib_classification.R' 'mllib_clustering.R' 'mllib_fpm.R' + 'mllib_recommendation.R' 'mllib_regression.R' 'mllib_stat.R' + 'mllib_tree.R' 'mllib_utils.R' 'serialize.R' 'sparkR.R' + 'stats.R' 'streaming.R' 'types.R' 'utils.R' 'window.R' +RoxygenNote: 7.2.3 +VignetteBuilder: knitr +NeedsCompilation: no +Encoding: UTF-8 +Packaged: 2023-09-09 02:05:55 UTC; spark-rm +Author: The Apache Software Foundation [aut, cre, cph] +Maintainer: The Apache Software Foundation +Built: R 4.3.1; ; 2023-09-09 02:08:15 UTC; unix diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/INDEX b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/INDEX new file mode 100644 index 0000000..fa90c18 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/INDEX @@ -0,0 +1,363 @@ +%<=>% %<=>% +%in% Match a column with given values. +addJobTag Add a tag to be assigned to all the jobs + started by this thread. +AFTSurvivalRegressionModel-class + S4 class that represents a + AFTSurvivalRegressionModel +agg summarize +alias alias +ALSModel-class S4 class that represents an ALSModel +approxQuantile Calculates the approximate quantiles of + numerical columns of a SparkDataFrame +arrange Arrange Rows by Variables +as.data.frame Download data from a SparkDataFrame into a R + data.frame +asc A set of operations working with SparkDataFrame + columns +attach,SparkDataFrame-method + Attach SparkDataFrame to R search path +avg avg +awaitTermination awaitTermination +between between +BisectingKMeansModel-class + S4 class that represents a BisectingKMeansModel +broadcast broadcast +cache Cache +cacheTable Cache Table +cancelJobGroup Cancel active jobs for the specified group +cancelJobsWithTag Cancel active jobs that have the specified tag. +cast Casts the column to a different data type. +checkpoint checkpoint +clearCache Clear Cache +clearJobGroup Clear current job group ID and its description +clearJobTags Clear the current thread's job tags. +coalesce Coalesce +collect Collects all the elements of a SparkDataFrame + and coerces them into an R data.frame. +colnames Column Names of SparkDataFrame +coltypes coltypes +column S4 class that represents a SparkDataFrame + column +column_aggregate_functions + Aggregate functions for Column operations +column_avro_functions Avro processing functions for Column operations +column_collection_functions + Collection functions for Column operations +column_datetime_diff_functions + Date time arithmetic functions for Column + operations +column_math_functions Math functions for Column operations +column_misc_functions Miscellaneous functions for Column operations +column_ml_functions ML functions for Column operations +column_nonaggregate_functions + Non-aggregate functions for Column operations +column_string_functions + String functions for Column operations +column_window_functions + Window functions for Column operations +corr corr +count Count +cov cov +createDataFrame Create a SparkDataFrame +createExternalTable (Deprecated) Create an external table +createOrReplaceTempView + Creates a temporary view using the given name. +createTable Creates a table based on the dataset in a data + source +crossJoin CrossJoin +crosstab Computes a pair-wise frequency table of the + given columns +cube cube +currentCatalog Returns the current default catalog +currentDatabase Returns the current default database +dapply dapply +dapplyCollect dapplyCollect +databaseExists Checks if the database with the specified name + exists. +DecisionTreeClassificationModel-class + S4 class that represents a + DecisionTreeClassificationModel +DecisionTreeRegressionModel-class + S4 class that represents a + DecisionTreeRegressionModel +describe describe +dim Returns the dimensions of SparkDataFrame +distinct Distinct +drop drop +dropDuplicates dropDuplicates +dropFields dropFields +dropna A set of SparkDataFrame functions working with + NA values +dropTempTable (Deprecated) Drop Temporary Table +dropTempView Drops the temporary view with the given view + name in the catalog. +dtypes DataTypes +endsWith endsWith +except except +exceptAll exceptAll +explain Explain +filter Filter +first Return the first row of a SparkDataFrame +fitted Get fitted result from a k-means model +FMClassificationModel-class + S4 class that represents a + FMClassificationModel +FMRegressionModel-class + S4 class that represents a FMRegressionModel +FPGrowthModel-class S4 class that represents a FPGrowthModel +freqItems Finding frequent items for columns, possibly + with false positives +functionExists Checks if the function with the specified name + exists. +gapply gapply +gapplyCollect gapplyCollect +GaussianMixtureModel-class + S4 class that represents a GaussianMixtureModel +GBTClassificationModel-class + S4 class that represents a + GBTClassificationModel +GBTRegressionModel-class + S4 class that represents a GBTRegressionModel +GeneralizedLinearRegressionModel-class + S4 class that represents a generalized linear + model +getDatabase Get the database with the specified name +getFunc Get the function with the specified name +getJobTags Get the tags that are currently set to be + assigned to all the jobs started by this + thread. +getLocalProperty Get a local property set in this thread, or + 'NULL' if it is missing. See + 'setLocalProperty'. +getNumPartitions getNumPartitions +getTable Get the table with the specified name +glm,formula,ANY,SparkDataFrame-method + Generalized Linear Models (R-compliant) +group_by GroupBy +GroupedData-class S4 class that represents a GroupedData +head Head +hint hint +histogram Compute histogram statistics for given column +insertInto insertInto +install.spark Download and Install Apache Spark to a Local + Directory +intersect Intersect +intersectAll intersectAll +isActive isActive +isLocal isLocal +IsotonicRegressionModel-class + S4 class that represents an + IsotonicRegressionModel +isStreaming isStreaming +join Join +KMeansModel-class S4 class that represents a KMeansModel +KSTest-class S4 class that represents an KSTest +last last +lastProgress lastProgress +LDAModel-class S4 class that represents an LDAModel +limit Limit +LinearRegressionModel-class + S4 class that represents a + LinearRegressionModel +LinearSVCModel-class S4 class that represents an LinearSVCModel +listCatalogs Returns a list of catalog available +listColumns Returns a list of columns for the given + table/view in the specified database +listDatabases Returns a list of databases available +listFunctions Returns a list of functions registered in the + specified database +listTables Returns a list of tables or views in the + specified database +localCheckpoint localCheckpoint +LogisticRegressionModel-class + S4 class that represents an + LogisticRegressionModel +merge Merges two data frames +MultilayerPerceptronClassificationModel-class + S4 class that represents a + MultilayerPerceptronClassificationModel +mutate Mutate +NaiveBayesModel-class S4 class that represents a NaiveBayesModel +ncol Returns the number of columns in a + SparkDataFrame +not ! +nrow Returns the number of rows in a SparkDataFrame +orderBy Ordering Columns in a WindowSpec +otherwise otherwise +over over +partitionBy partitionBy +persist Persist +pivot Pivot a column of the GroupedData and perform + the specified aggregation. +PowerIterationClustering-class + S4 class that represents a + PowerIterationClustering +predict Makes predictions from a MLlib model +PrefixSpan-class S4 class that represents a PrefixSpan +printSchema Print Schema of a SparkDataFrame +queryName queryName +RandomForestClassificationModel-class + S4 class that represents a + RandomForestClassificationModel +RandomForestRegressionModel-class + S4 class that represents a + RandomForestRegressionModel +randomSplit randomSplit +rangeBetween rangeBetween +rbind Union two or more SparkDataFrames +read.df Load a SparkDataFrame +read.jdbc Create a SparkDataFrame representing the + database table accessible via JDBC URL +read.json Create a SparkDataFrame from a JSON file. +read.ml Load a fitted MLlib model from the input path. +read.orc Create a SparkDataFrame from an ORC file. +read.parquet Create a SparkDataFrame from a Parquet file. +read.stream Load a streaming SparkDataFrame +read.text Create a SparkDataFrame from a text file. +recoverPartitions Recovers all the partitions in the directory of + a table and update the catalog +refreshByPath Invalidates and refreshes all the cached data + and metadata for SparkDataFrame containing path +refreshTable Invalidates and refreshes all the cached data + and metadata of the given table +registerTempTable (Deprecated) Register Temporary Table +removeJobTag Remove a tag previously added to be assigned to + all the jobs started by this thread. Noop if + such a tag was not added earlier. +rename rename +repartition Repartition +repartitionByRange Repartition by range +rollup rollup +rowsBetween rowsBetween +sample Sample +sampleBy Returns a stratified sample without replacement +saveAsTable Save the contents of the SparkDataFrame to a + data source as a table +schema Get schema object +select Select +selectExpr SelectExpr +setCheckpointDir Set checkpoint directory +setCurrentCatalog Sets the current default catalog +setCurrentDatabase Sets the current default database +setInterruptOnCancel Set the behavior of job cancellation from jobs + started in this thread. +setJobDescription Set a human readable description of the current + job. +setJobGroup Assigns a group ID to all the jobs started by + this thread until the group ID is set to a + different value or cleared. +setLocalProperty Set a local property that affects jobs + submitted from this thread, such as the Spark + fair scheduler pool. +setLogLevel Set new log level +show show +showDF showDF +spark.addFile Add a file or directory to be downloaded with + this Spark job on every node. +spark.als Alternating Least Squares (ALS) for + Collaborative Filtering +spark.assignClusters PowerIterationClustering +spark.bisectingKmeans Bisecting K-Means Clustering Model +spark.decisionTree Decision Tree Model for Regression and + Classification +spark.findFrequentSequentialPatterns + PrefixSpan +spark.fmClassifier Factorization Machines Classification Model +spark.fmRegressor Factorization Machines Regression Model +spark.fpGrowth FP-growth +spark.gaussianMixture Multivariate Gaussian Mixture Model (GMM) +spark.gbt Gradient Boosted Tree Model for Regression and + Classification +spark.getSparkFiles Get the absolute path of a file added through + spark.addFile. +spark.getSparkFilesRootDirectory + Get the root directory that contains files + added through spark.addFile. +spark.glm Generalized Linear Models +spark.isoreg Isotonic Regression Model +spark.kmeans K-Means Clustering Model +spark.kstest (One-Sample) Kolmogorov-Smirnov Test +spark.lapply Run a function over a list of elements, + distributing the computations with Spark +spark.lda Latent Dirichlet Allocation +spark.lm Linear Regression Model +spark.logit Logistic Regression Model +spark.mlp Multilayer Perceptron Classification Model +spark.naiveBayes Naive Bayes Models +spark.randomForest Random Forest Model for Regression and + Classification +spark.survreg Accelerated Failure Time (AFT) Survival + Regression Model +spark.svmLinear Linear SVM Model +SparkDataFrame-class S4 class that represents a SparkDataFrame +sparkR.callJMethod Call Java Methods +sparkR.callJStatic Call Static Java Methods +sparkR.conf Get Runtime Config from the current active + SparkSession +sparkR.init (Deprecated) Initialize a new Spark Context +sparkR.newJObject Create Java Objects +sparkR.session Get the existing SparkSession or initialize a + new SparkSession. +sparkR.session.stop Stop the Spark Session and Spark Context +sparkR.uiWebUrl Get the URL of the SparkUI instance for the + current active SparkSession +sparkR.version Get version of Spark on which this application + is running +sparkRHive.init (Deprecated) Initialize a new HiveContext +sparkRSQL.init (Deprecated) Initialize a new SQLContext +sql SQL Query +startsWith startsWith +status status +stopQuery stopQuery +storageLevel StorageLevel +str Compactly display the structure of a dataset +StreamingQuery-class S4 class that represents a StreamingQuery +structField structField +structType structType +subset Subset +substr substr +summary summary +tableExists Checks if the table with the specified name + exists. +tableNames Table Names +tables Tables +tableToDF Create a SparkDataFrame from a SparkSQL table + or view +take Take the first NUM rows of a SparkDataFrame and + return the results as a R data.frame +timestamp_seconds Date time functions for Column operations +toJSON toJSON +uncacheTable Uncache Table +union Return a new SparkDataFrame containing the + union of rows +unionAll Return a new SparkDataFrame containing the + union of rows. +unionByName Return a new SparkDataFrame containing the + union of rows, matched by column names +unpersist Unpersist +unpivot Unpivot a DataFrame from wide format to long + format. +windowOrderBy windowOrderBy +windowPartitionBy windowPartitionBy +WindowSpec-class S4 class that represents a WindowSpec +with Evaluate a R expression in an environment + constructed from a SparkDataFrame +withColumn WithColumn +withField withField +withWatermark withWatermark +write.df Save the contents of SparkDataFrame to a data + source. +write.jdbc Save the content of SparkDataFrame to an + external database table via JDBC. +write.json Save the contents of SparkDataFrame as a JSON + file +write.ml Saves the MLlib model to the input path +write.orc Save the contents of SparkDataFrame as an ORC + file, preserving the schema. +write.parquet Save the contents of SparkDataFrame as a + Parquet file, preserving the schema. +write.stream Write the streaming SparkDataFrame to a data + source. +write.text Save the content of SparkDataFrame in a text + file at the specified path. diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/Rd.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/Rd.rds new file mode 100644 index 0000000..758e1aa Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/Rd.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/features.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/features.rds new file mode 100644 index 0000000..bb8956c Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/features.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/hsearch.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/hsearch.rds new file mode 100644 index 0000000..c16e943 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/hsearch.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/links.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/links.rds new file mode 100644 index 0000000..f8d10e9 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/links.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/nsInfo.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/nsInfo.rds new file mode 100644 index 0000000..0871f1b Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/nsInfo.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/package.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/package.rds new file mode 100644 index 0000000..e075780 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/package.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/vignette.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/vignette.rds new file mode 100644 index 0000000..927a0a8 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/Meta/vignette.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/NAMESPACE b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/NAMESPACE new file mode 100644 index 0000000..c5668d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/NAMESPACE @@ -0,0 +1,583 @@ +# +# 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. +# + +# Imports from base R +# Do not include stats:: "rpois", "runif" - causes error at runtime +importFrom("methods", "setGeneric", "setMethod", "setOldClass") +importFrom("methods", "is", "new", "signature", "show") +importFrom("stats", "gaussian", "setNames") +importFrom("utils", "download.file", "object.size", "packageVersion", "tail", "untar") + +# Disable native libraries till we figure out how to package it +# See SPARKR-7839 +#useDynLib(SparkR, stringHashCode) + +# S3 methods exported +export("sparkR.session") +export("sparkR.init") +export("sparkR.session.stop") +export("sparkR.stop") +export("sparkR.conf") +export("sparkR.version") +export("sparkR.uiWebUrl") +export("print.jobj") + +export("sparkR.newJObject") +export("sparkR.callJMethod") +export("sparkR.callJStatic") + +export("install.spark") + +export("sparkRSQL.init", + "sparkRHive.init") + +# MLlib integration +exportMethods("glm", + "spark.glm", + "predict", + "summary", + "spark.kmeans", + "fitted", + "spark.mlp", + "spark.naiveBayes", + "spark.survreg", + "spark.lda", + "spark.posterior", + "spark.perplexity", + "spark.isoreg", + "spark.gaussianMixture", + "spark.als", + "spark.kstest", + "spark.logit", + "spark.decisionTree", + "spark.randomForest", + "spark.gbt", + "spark.bisectingKmeans", + "spark.svmLinear", + "spark.fpGrowth", + "spark.freqItemsets", + "spark.associationRules", + "spark.findFrequentSequentialPatterns", + "spark.assignClusters", + "spark.fmClassifier", + "spark.lm", + "spark.fmRegressor") + +# Job group and job tag lifecycle management methods +export("setJobGroup", + "clearJobGroup", + "cancelJobGroup", + "setJobDescription", + "setInterruptOnCancel", + "setLocalProperty", + "getLocalProperty", + "addJobTag", + "removeJobTag", + "getJobTags", + "clearJobTags", + "cancelJobsWithTag") + +# Export Utility methods +export("setLogLevel") + +exportClasses("SparkDataFrame") + +exportMethods("arrange", + "as.data.frame", + "attach", + "broadcast", + "cache", + "checkpoint", + "coalesce", + "collect", + "colnames", + "colnames<-", + "coltypes", + "coltypes<-", + "columns", + "count", + "cov", + "corr", + "covar_samp", + "covar_pop", + "createOrReplaceTempView", + "crossJoin", + "crosstab", + "cube", + "dapply", + "dapplyCollect", + "describe", + "dim", + "distinct", + "drop", + "dropDuplicates", + "dropna", + "dtypes", + "except", + "exceptAll", + "explain", + "fillna", + "filter", + "first", + "freqItems", + "gapply", + "gapplyCollect", + "getNumPartitions", + "group_by", + "groupBy", + "head", + "hint", + "insertInto", + "intersect", + "intersectAll", + "isLocal", + "isStreaming", + "join", + "limit", + "localCheckpoint", + "melt", + "merge", + "mutate", + "na.omit", + "names", + "names<-", + "ncol", + "nrow", + "orderBy", + "persist", + "printSchema", + "randomSplit", + "rbind", + "registerTempTable", + "rename", + "repartition", + "repartitionByRange", + "rollup", + "sample", + "sample_frac", + "sampleBy", + "saveAsTable", + "saveDF", + "schema", + "select", + "selectExpr", + "show", + "showDF", + "storageLevel", + "subset", + "summarize", + "summary", + "take", + "toJSON", + "transform", + "union", + "unionAll", + "unionByName", + "unique", + "unpersist", + "unpivot", + "where", + "with", + "withColumn", + "withColumnRenamed", + "withWatermark", + "write.df", + "write.jdbc", + "write.json", + "write.orc", + "write.parquet", + "write.stream", + "write.text", + "write.ml") + +exportClasses("Column") + +exportMethods("%<=>%", + "%in%", + "abs", + "acos", + "acosh", + "add_months", + "alias", + "approx_count_distinct", + "approxCountDistinct", + "approxQuantile", + "array_aggregate", + "array_contains", + "array_distinct", + "array_except", + "array_exists", + "array_filter", + "array_forall", + "array_intersect", + "array_join", + "array_max", + "array_min", + "array_position", + "array_remove", + "array_repeat", + "array_sort", + "array_to_vector", + "array_transform", + "arrays_overlap", + "array_union", + "arrays_zip", + "arrays_zip_with", + "asc", + "asc_nulls_first", + "asc_nulls_last", + "ascii", + "asin", + "asinh", + "assert_true", + "atan", + "atanh", + "atan2", + "avg", + "base64", + "between", + "bin", + "bit_length", + "bitwise_not", + "bitwiseNOT", + "bround", + "cast", + "cbrt", + "ceil", + "ceiling", + "collect_list", + "collect_set", + "column", + "concat", + "concat_ws", + "contains", + "conv", + "cos", + "cosh", + "cot", + "count", + "count_distinct", + "countDistinct", + "crc32", + "create_array", + "create_map", + "csc", + "current_date", + "current_timestamp", + "hash", + "cume_dist", + "date_add", + "date_format", + "date_sub", + "date_trunc", + "datediff", + "dayofmonth", + "dayofweek", + "dayofyear", + "decode", + "degrees", + "dense_rank", + "desc", + "desc_nulls_first", + "desc_nulls_last", + "dropFields", + "element_at", + "encode", + "endsWith", + "exp", + "explode", + "explode_outer", + "expm1", + "expr", + "factorial", + "first", + "flatten", + "floor", + "format_number", + "format_string", + "from_avro", + "from_csv", + "from_json", + "from_unixtime", + "from_utc_timestamp", + "getField", + "getItem", + "greatest", + "grouping_bit", + "grouping_id", + "hex", + "histogram", + "hour", + "hypot", + "ifelse", + "ilike", + "initcap", + "input_file_name", + "instr", + "isNaN", + "isNotNull", + "isNull", + "is.nan", + "isnan", + "kurtosis", + "lag", + "last", + "last_day", + "lead", + "least", + "length", + "levenshtein", + "like", + "lit", + "ln", + "locate", + "log", + "log10", + "log1p", + "log2", + "lower", + "lpad", + "ltrim", + "make_date", + "map_concat", + "map_entries", + "map_filter", + "map_from_arrays", + "map_from_entries", + "map_keys", + "map_values", + "map_zip_with", + "max", + "max_by", + "md5", + "mean", + "min", + "min_by", + "minute", + "monotonically_increasing_id", + "month", + "months_between", + "n", + "n_distinct", + "nanvl", + "negate", + "negative", + "next_day", + "not", + "nth_value", + "ntile", + "octet_length", + "otherwise", + "over", + "overlay", + "percent_rank", + "percentile_approx", + "pmod", + "posexplode", + "posexplode_outer", + "positive", + "product", + "quarter", + "radians", + "raise_error", + "rand", + "randn", + "rank", + "regexp_extract", + "regexp_replace", + "repeat_string", + "reverse", + "rint", + "rlike", + "round", + "row_number", + "rpad", + "rtrim", + "schema_of_csv", + "schema_of_json", + "sec", + "second", + "sha1", + "sha2", + "shiftLeft", + "shiftleft", + "shiftRight", + "shiftright", + "shiftRightUnsigned", + "shiftrightunsigned", + "shuffle", + "sd", + "sign", + "signum", + "sin", + "sinh", + "size", + "skewness", + "slice", + "sort_array", + "soundex", + "spark_partition_id", + "split_string", + "std", + "stddev", + "stddev_pop", + "stddev_samp", + "struct", + "sqrt", + "startsWith", + "substr", + "substring_index", + "sum", + "sum_distinct", + "sumDistinct", + "tan", + "tanh", + "timestamp_seconds", + "toDegrees", + "toRadians", + "to_avro", + "to_csv", + "to_date", + "to_json", + "to_timestamp", + "to_utc_timestamp", + "translate", + "transform_keys", + "transform_values", + "trim", + "trunc", + "unbase64", + "unhex", + "unix_timestamp", + "upper", + "var", + "variance", + "var_pop", + "var_samp", + "vector_to_array", + "weekofyear", + "when", + "width_bucket", + "window", + "withField", + "xxhash64", + "year") + +exportClasses("GroupedData") +exportMethods("agg") +exportMethods("pivot") + +export("as.DataFrame", + "cacheTable", + "clearCache", + "createDataFrame", + "createExternalTable", + "createTable", + "currentCatalog", + "currentDatabase", + "databaseExists", + "dropTempTable", + "dropTempView", + "functionExists", + "getDatabase", + "getFunc", + "getTable", + "listCatalogs", + "listColumns", + "listDatabases", + "listFunctions", + "listTables", + "loadDF", + "read.df", + "read.jdbc", + "read.json", + "read.orc", + "read.parquet", + "read.stream", + "read.text", + "recoverPartitions", + "refreshByPath", + "refreshTable", + "setCheckpointDir", + "setCurrentCatalog", + "setCurrentDatabase", + "spark.lapply", + "spark.addFile", + "spark.getSparkFilesRootDirectory", + "spark.getSparkFiles", + "sql", + "str", + "tableExists", + "tableToDF", + "tableNames", + "tables", + "uncacheTable", + "print.summary.GeneralizedLinearRegressionModel", + "read.ml", + "print.summary.KSTest", + "print.summary.DecisionTreeRegressionModel", + "print.summary.DecisionTreeClassificationModel", + "print.summary.RandomForestRegressionModel", + "print.summary.RandomForestClassificationModel", + "print.summary.GBTRegressionModel", + "print.summary.GBTClassificationModel") + +export("structField", + "structField.jobj", + "structField.character", + "print.structField", + "structType", + "structType.character", + "structType.jobj", + "structType.structField", + "print.structType") + +exportClasses("WindowSpec") + +export("partitionBy", + "rowsBetween", + "rangeBetween") + +export("windowPartitionBy", + "windowOrderBy") + +exportClasses("StreamingQuery") + +export("awaitTermination", + "isActive", + "lastProgress", + "queryName", + "status", + "stopQuery") + + +S3method(print, jobj) +S3method(print, structField) +S3method(print, structType) +S3method(print, summary.GeneralizedLinearRegressionModel) +S3method(print, summary.KSTest) +S3method(print, summary.DecisionTreeRegressionModel) +S3method(print, summary.DecisionTreeClassificationModel) +S3method(print, summary.RandomForestRegressionModel) +S3method(print, summary.RandomForestClassificationModel) +S3method(print, summary.GBTRegressionModel) +S3method(print, summary.GBTClassificationModel) +S3method(structField, character) +S3method(structField, jobj) +S3method(structType, character) +S3method(structType, jobj) +S3method(structType, structField) diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR new file mode 100644 index 0000000..6686156 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR @@ -0,0 +1,27 @@ +# File share/R/nspackloader.R +# Part of the R package, https://www.R-project.org +# +# Copyright (C) 1995-2012 The R Core Team +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# A copy of the GNU General Public License is available at +# https://www.r-project.org/Licenses/ + +local({ + info <- loadingNamespaceInfo() + pkg <- info$pkgname + ns <- .getNamespace(as.name(pkg)) + if (is.null(ns)) + stop("cannot find namespace environment for ", pkg, domain = NA); + dbbase <- file.path(info$libname, pkg, "R", pkg) + lazyLoad(dbbase, ns, filter = function(n) n != ".__NAMESPACE__.") +}) diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdb b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdb new file mode 100644 index 0000000..e073bfd Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdb differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdx b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdx new file mode 100644 index 0000000..d007444 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/R/SparkR.rdx differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/index.html b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/index.html new file mode 100644 index 0000000..9991db3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/index.html @@ -0,0 +1,29 @@ + + +R: Vignettes and other documentation + + + +
+

Vignettes and other documentation + +

+
+
+[Top] +
+

Vignettes from package 'SparkR'

+ +++++++ + + + + +
SparkR::sparkr-vignettesSparkR - Practical GuideHTMLsourceR code
+
diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.R new file mode 100644 index 0000000..89a25b9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.R @@ -0,0 +1,579 @@ +## ----dynamic-chunk-options, include=FALSE------------------------------------- +# In GitHub lint job, we don't have full JVM build +# SparkR vignette fails to evaluate +GITHUB_ACTIONS <- tolower(Sys.getenv("GITHUB_ACTIONS")) == "true" +EVAL_CHUNK <- !GITHUB_ACTIONS + +## ----setup, include=FALSE, eval=EVAL_CHUNK------------------------------------ +library(knitr) +opts_hooks$set(eval = function(options) { + # override eval to FALSE only on windows + if (.Platform$OS.type == "windows") { + options$eval = FALSE + } + options +}) +r_tmp_dir <- tempdir() +tmp_arg <- paste0("-Djava.io.tmpdir=", r_tmp_dir) +sparkSessionConfig <- list(spark.driver.extraJavaOptions = tmp_arg, + spark.executor.extraJavaOptions = tmp_arg) +old_java_opt <- Sys.getenv("_JAVA_OPTIONS") +Sys.setenv("_JAVA_OPTIONS" = paste("-XX:-UsePerfData", old_java_opt, sep = " ")) + +## ---- message=FALSE, eval=EVAL_CHUNK------------------------------------------ +library(SparkR) + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +# disable eval if java version not supported +override_eval <- tryCatch(!is.numeric(SparkR:::checkJavaVersion()), + error = function(e) { TRUE }, + warning = function(e) { TRUE }) + +if (override_eval) { + opts_hooks$set(eval = function(options) { + options$eval = FALSE + options + }) +} + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +install.spark() +sparkR.session(master = "local[1]", sparkConfig = sparkSessionConfig, enableHiveSupport = FALSE) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +sparkR.session() + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +cars <- cbind(model = rownames(mtcars), mtcars) +carsDF <- createDataFrame(cars) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(carsDF) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsSubDF <- select(carsDF, "model", "mpg", "hp") +carsSubDF <- filter(carsSubDF, carsSubDF$hp >= 200) +head(carsSubDF) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsGPDF <- summarize(groupBy(carsDF, carsDF$gear), count = n(carsDF$gear)) +head(carsGPDF) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsGP <- collect(carsGPDF) +class(carsGP) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +model <- spark.glm(carsDF, mpg ~ wt + cyl) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +summary(model) + +## ---- eval=FALSE-------------------------------------------------------------- +# write.ml(model, path = "/HOME/tmp/mlModel/glmModel") + +## ---- eval=FALSE-------------------------------------------------------------- +# sparkR.session.stop() + +## ---- eval=FALSE-------------------------------------------------------------- +# install.spark() + +## ---- eval=FALSE-------------------------------------------------------------- +# sparkR.session(sparkHome = "/HOME/spark") + +## ---- eval=FALSE-------------------------------------------------------------- +# spark_warehouse_path <- file.path(path.expand('~'), "spark-warehouse") +# sparkR.session(spark.sql.warehouse.dir = spark_warehouse_path) + +## ---- echo=FALSE, tidy = TRUE, eval=EVAL_CHUNK-------------------------------- +paste("Spark", packageVersion("SparkR")) + +## ---- eval=FALSE-------------------------------------------------------------- +# sparkR.session(master = "spark://local:7077") + +## ---- eval=FALSE-------------------------------------------------------------- +# sparkR.session(master = "yarn") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +df <- as.DataFrame(faithful) +head(df) + +## ---- eval=FALSE-------------------------------------------------------------- +# sparkR.session(sparkPackages = "com.databricks:spark-avro_2.12:3.0.0") + +## ---- eval=FALSE-------------------------------------------------------------- +# df <- read.df(csvPath, "csv", header = "true", inferSchema = "true", na.strings = "NA") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +filePath <- paste0(sparkR.conf("spark.home"), + "/examples/src/main/resources/people.json") +readLines(filePath, n = 2L) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +people <- read.df(filePath, "json") +count(people) +head(people) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +printSchema(people) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +people <- read.json(paste0(Sys.getenv("SPARK_HOME"), + c("/examples/src/main/resources/people.json", + "/examples/src/main/resources/people.json"))) +count(people) + +## ---- eval=FALSE-------------------------------------------------------------- +# write.df(people, path = "people.parquet", source = "parquet", mode = "overwrite") + +## ---- eval=FALSE-------------------------------------------------------------- +# sql("CREATE TABLE IF NOT EXISTS src (key INT, value STRING)") +# +# txtPath <- paste0(sparkR.conf("spark.home"), "/examples/src/main/resources/kv1.txt") +# sqlCMD <- sprintf("LOAD DATA LOCAL INPATH '%s' INTO TABLE src", txtPath) +# sql(sqlCMD) +# +# results <- sql("FROM src SELECT key, value") +# +# # results is now a SparkDataFrame +# head(results) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsDF + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +printSchema(carsDF) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(select(carsDF, "mpg")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(filter(carsDF, carsDF$mpg < 20)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +numCyl <- summarize(groupBy(carsDF, carsDF$cyl), count = n(carsDF$cyl)) +head(numCyl) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +mean(cube(carsDF, "cyl", "gear", "am"), "mpg") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +mean(rollup(carsDF, "cyl", "gear", "am"), "mpg") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsDF_km <- carsDF +carsDF_km$kmpg <- carsDF_km$mpg * 1.61 +head(select(carsDF_km, "model", "mpg", "kmpg")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsSubDF <- select(carsDF, "model", "mpg", "cyl") +ws <- orderBy(windowPartitionBy("cyl"), "mpg") +carsRank <- withColumn(carsSubDF, "rank", over(rank(), ws)) +head(carsRank, n = 20L) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +carsSubDF <- select(carsDF, "model", "mpg") +schema <- "model STRING, mpg DOUBLE, kmpg DOUBLE" +out <- dapply(carsSubDF, function(x) { x <- cbind(x, x$mpg * 1.61) }, schema) +head(collect(out)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +out <- dapplyCollect( + carsSubDF, + function(x) { + x <- cbind(x, "kmpg" = x$mpg * 1.61) + }) +head(out, 3) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +schema <- structType(structField("cyl", "double"), structField("max_mpg", "double")) +result <- gapply( + carsDF, + "cyl", + function(key, x) { + y <- data.frame(key, max(x$mpg)) + }, + schema) +head(arrange(result, "max_mpg", decreasing = TRUE)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +result <- gapplyCollect( + carsDF, + "cyl", + function(key, x) { + y <- data.frame(key, max(x$mpg)) + colnames(y) <- c("cyl", "max_mpg") + y + }) +head(result[order(result$max_mpg, decreasing = TRUE), ]) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +costs <- exp(seq(from = log(1), to = log(1000), length.out = 5)) +train <- function(cost) { + stopifnot(requireNamespace("e1071", quietly = TRUE)) + model <- e1071::svm(Species ~ ., data = iris, cost = cost) + summary(model) +} + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +model.summaries <- spark.lapply(costs, train) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +class(model.summaries) + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +ops <- options() +options(max.print=40) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +print(model.summaries[[2]]) + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +options(ops) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +people <- read.df(paste0(sparkR.conf("spark.home"), + "/examples/src/main/resources/people.json"), "json") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +createOrReplaceTempView(people, "people") + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +teenagers <- sql("SELECT name FROM people WHERE age >= 13 AND age <= 19") +head(teenagers) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +splitDF_list <- randomSplit(carsDF, c(0.7, 0.3), seed = 0) +carsDF_train <- splitDF_list[[1]] +carsDF_test <- splitDF_list[[2]] + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +count(carsDF_train) +head(carsDF_train) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +count(carsDF_test) +head(carsDF_test) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +# load training data and create a DataFrame +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# fit a Linear SVM classifier model +model <- spark.svmLinear(training, Survived ~ ., regParam = 0.01, maxIter = 10) +summary(model) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +prediction <- predict(model, training) +head(select(prediction, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +model <- spark.logit(training, Survived ~ ., regParam = 0.04741301) +summary(model) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +fitted <- predict(model, training) +head(select(fitted, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# Note in this case, Spark infers it is multinomial logistic regression, so family = "multinomial" is optional. +model <- spark.logit(training, Class ~ ., regParam = 0.07815179) +summary(model) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# fit a Multilayer Perceptron Classification Model +model <- spark.mlp(training, Survived ~ Age + Sex, blockSize = 128, layers = c(2, 2), solver = "l-bfgs", maxIter = 100, tol = 0.5, stepSize = 1, seed = 1, initialWeights = c( 0, 0, 5, 5, 9, 9)) + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +ops <- options() +options(max.print=5) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +# check the summary of the fitted model +summary(model) + +## ---- include=FALSE, eval=EVAL_CHUNK------------------------------------------ +options(ops) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +# make predictions use the fitted model +predictions <- predict(model, training) +head(select(predictions, predictions$prediction)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +titanic <- as.data.frame(Titanic) +titanicDF <- createDataFrame(titanic[titanic$Freq > 0, -5]) +naiveBayesModel <- spark.naiveBayes(titanicDF, Survived ~ Class + Sex + Age) +summary(naiveBayesModel) +naiveBayesPrediction <- predict(naiveBayesModel, titanicDF) +head(select(naiveBayesPrediction, "Class", "Sex", "Age", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) + +model <- spark.fmClassifier(training, Survived ~ Age + Sex) +summary(model) + +predictions <- predict(model, training) +head(select(predictions, predictions$prediction)) + +## ---- warning=FALSE, eval=EVAL_CHUNK------------------------------------------ +library(survival) +ovarianDF <- createDataFrame(ovarian) +aftModel <- spark.survreg(ovarianDF, Surv(futime, fustat) ~ ecog_ps + rx) +summary(aftModel) +aftPredictions <- predict(aftModel, ovarianDF) +head(aftPredictions) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +gaussianGLM <- spark.glm(carsDF, mpg ~ wt + hp) +summary(gaussianGLM) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +gaussianFitted <- predict(gaussianGLM, carsDF) +head(select(gaussianFitted, "model", "prediction", "mpg", "wt", "hp")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +tweedieGLM1 <- spark.glm(carsDF, mpg ~ wt + hp, family = "tweedie", var.power = 0.0) +summary(tweedieGLM1) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +tweedieGLM2 <- spark.glm(carsDF, mpg ~ wt + hp, family = "tweedie", + var.power = 1.2, link.power = 0.0) +summary(tweedieGLM2) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +y <- c(3.0, 6.0, 8.0, 5.0, 7.0) +x <- c(1.0, 2.0, 3.5, 3.0, 4.0) +w <- rep(1.0, 5) +data <- data.frame(y = y, x = x, w = w) +df <- createDataFrame(data) +isoregModel <- spark.isoreg(df, y ~ x, weightCol = "w") +isoregFitted <- predict(isoregModel, df) +head(select(isoregFitted, "x", "y", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +newDF <- createDataFrame(data.frame(x = c(1.5, 3.2))) +head(predict(isoregModel, newDF)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +model <- spark.lm(carsDF, mpg ~ wt + hp) + +summary(model) +predictions <- predict(model, carsDF) +head(select(predictions, predictions$prediction)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +model <- spark.fmRegressor(carsDF, mpg ~ wt + hp) +summary(model) +predictions <- predict(model, carsDF) +head(select(predictions, predictions$prediction)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +dtModel <- spark.decisionTree(df, Survived ~ ., type = "classification", maxDepth = 2) +summary(dtModel) +predictions <- predict(dtModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +gbtModel <- spark.gbt(df, Survived ~ ., type = "classification", maxDepth = 2, maxIter = 2) +summary(gbtModel) +predictions <- predict(gbtModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +rfModel <- spark.randomForest(df, Survived ~ ., type = "classification", maxDepth = 2, numTrees = 2) +summary(rfModel) +predictions <- predict(rfModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +model <- spark.bisectingKmeans(training, Class ~ Survived, k = 4) +summary(model) +fitted <- predict(model, training) +head(select(fitted, "Class", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +X1 <- data.frame(V1 = rnorm(4), V2 = rnorm(4)) +X2 <- data.frame(V1 = rnorm(6, 3), V2 = rnorm(6, 4)) +data <- rbind(X1, X2) +df <- createDataFrame(data) +gmmModel <- spark.gaussianMixture(df, ~ V1 + V2, k = 2) +summary(gmmModel) +gmmFitted <- predict(gmmModel, df) +head(select(gmmFitted, "V1", "V2", "prediction")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +kmeansModel <- spark.kmeans(carsDF, ~ mpg + hp + wt, k = 3) +summary(kmeansModel) +kmeansPredictions <- predict(kmeansModel, carsDF) +head(select(kmeansPredictions, "model", "mpg", "hp", "wt", "prediction"), n = 20L) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +corpus <- data.frame(features = c( + "1 2 6 0 2 3 1 1 0 0 3", + "1 3 0 1 3 0 0 2 0 0 1", + "1 4 1 0 0 4 9 0 1 2 0", + "2 1 0 3 0 0 5 0 2 3 9", + "3 1 1 9 3 0 2 0 0 1 3", + "4 2 0 3 4 5 1 1 1 4 0", + "2 1 0 3 0 0 5 0 2 2 9", + "1 1 1 9 2 1 2 0 0 1 3", + "4 4 0 3 4 2 1 3 0 0 0", + "2 8 2 0 3 0 2 0 2 7 2", + "1 1 1 9 0 2 2 0 0 3 3", + "4 1 0 0 4 5 1 3 0 1 0")) +corpusDF <- createDataFrame(corpus) +model <- spark.lda(data = corpusDF, k = 5, optimizer = "em") +summary(model) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +posterior <- spark.posterior(model, corpusDF) +head(posterior) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +perplexity <- spark.perplexity(model, corpusDF) +perplexity + +## ---- eval=FALSE-------------------------------------------------------------- +# ratings <- list(list(0, 0, 4.0), list(0, 1, 2.0), list(1, 1, 3.0), list(1, 2, 4.0), +# list(2, 1, 1.0), list(2, 2, 5.0)) +# df <- createDataFrame(ratings, c("user", "item", "rating")) +# model <- spark.als(df, "rating", "user", "item", rank = 10, reg = 0.1, nonnegative = TRUE) + +## ---- eval=FALSE-------------------------------------------------------------- +# stats <- summary(model) +# userFactors <- stats$userFactors +# itemFactors <- stats$itemFactors +# head(userFactors) +# head(itemFactors) + +## ---- eval=FALSE-------------------------------------------------------------- +# predicted <- predict(model, df) +# head(predicted) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +df <- createDataFrame(list(list(0L, 1L, 1.0), list(0L, 2L, 1.0), + list(1L, 2L, 1.0), list(3L, 4L, 1.0), + list(4L, 0L, 0.1)), + schema = c("src", "dst", "weight")) +head(spark.assignClusters(df, initMode = "degree", weightCol = "weight")) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +df <- selectExpr(createDataFrame(data.frame(rawItems = c( + "T,R,U", "T,S", "V,R", "R,U,T,V", "R,S", "V,S,U", "U,R", "S,T", "V,R", "V,U,S", + "T,V,U", "R,V", "T,S", "T,S", "S,T", "S,U", "T,R", "V,R", "S,V", "T,S,U" +))), "split(rawItems, ',') AS items") + +fpm <- spark.fpGrowth(df, minSupport = 0.2, minConfidence = 0.5) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(spark.freqItemsets(fpm)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(spark.associationRules(fpm)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +head(predict(fpm, df)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +df <- createDataFrame(list(list(list(list(1L, 2L), list(3L))), + list(list(list(1L), list(3L, 2L), list(1L, 2L))), + list(list(list(1L, 2L), list(5L))), + list(list(list(6L)))), + schema = c("sequence")) +head(spark.findFrequentSequentialPatterns(df, minSupport = 0.5, maxPatternLength = 5L)) + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +freqStats <- head(select(df, mean(df$Freq), sd(df$Freq))) +freqMean <- freqStats[1] +freqStd <- freqStats[2] + +test <- spark.kstest(df, "Freq", "norm", c(freqMean, freqStd)) +testSummary <- summary(test) +testSummary + +## ---- eval=EVAL_CHUNK--------------------------------------------------------- +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +gaussianGLM <- spark.glm(training, Freq ~ Sex + Age, family = "gaussian") + +# Save and then load a fitted MLlib model +modelPath <- tempfile(pattern = "ml", fileext = ".tmp") +write.ml(gaussianGLM, modelPath) +gaussianGLM2 <- read.ml(modelPath) + +# Check model summary +summary(gaussianGLM2) + +# Check model prediction +gaussianPredictions <- predict(gaussianGLM2, training) +head(gaussianPredictions) + +unlink(modelPath) + +## ---- eval=FALSE-------------------------------------------------------------- +# # Create DataFrame representing the stream of input lines from connection +# lines <- read.stream("socket", host = hostname, port = port) +# +# # Split the lines into words +# words <- selectExpr(lines, "explode(split(value, ' ')) as word") +# +# # Generate running word count +# wordCounts <- count(groupBy(words, "word")) +# +# # Start running the query that prints the running counts to the console +# query <- write.stream(wordCounts, "console", outputMode = "complete") + +## ---- eval=FALSE-------------------------------------------------------------- +# topic <- read.stream("kafka", +# kafka.bootstrap.servers = "host1:port1,host2:port2", +# subscribe = "topic1") +# keyvalue <- selectExpr(topic, "CAST(key AS STRING)", "CAST(value AS STRING)") + +## ---- eval=FALSE-------------------------------------------------------------- +# noAggDF <- select(where(deviceDataStreamingDf, "signal > 10"), "device") +# +# # Print new data to console +# write.stream(noAggDF, "console") +# +# # Write new data to Parquet files +# write.stream(noAggDF, +# "parquet", +# path = "path/to/destination/dir", +# checkpointLocation = "path/to/checkpoint/dir") +# +# # Aggregate +# aggDF <- count(groupBy(noAggDF, "device")) +# +# # Print updated aggregations to console +# write.stream(aggDF, "console", outputMode = "complete") +# +# # Have all the aggregates in an in memory table. The query name will be the table name +# write.stream(aggDF, "memory", queryName = "aggregates", outputMode = "complete") +# +# head(sql("select * from aggregates")) + +## ---- echo=FALSE, eval=EVAL_CHUNK--------------------------------------------- +sparkR.session.stop() + +## ----cleanup, include=FALSE, eval=EVAL_CHUNK---------------------------------- +SparkR:::uninstallDownloadedSpark() + diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.Rmd b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.Rmd new file mode 100644 index 0000000..4a51076 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.Rmd @@ -0,0 +1,1294 @@ +--- +title: "SparkR - Practical Guide" +output: + rmarkdown::html_vignette: + toc: true + toc_depth: 4 +vignette: > + %\VignetteIndexEntry{SparkR - Practical Guide} + %\VignetteEngine{knitr::rmarkdown} + \usepackage[utf8]{inputenc} +--- + + + +```{r dynamic-chunk-options, include=FALSE} +# In GitHub lint job, we don't have full JVM build +# SparkR vignette fails to evaluate +GITHUB_ACTIONS <- tolower(Sys.getenv("GITHUB_ACTIONS")) == "true" +EVAL_CHUNK <- !GITHUB_ACTIONS +``` + + +```{r setup, include=FALSE, eval=EVAL_CHUNK} +library(knitr) +opts_hooks$set(eval = function(options) { + # override eval to FALSE only on windows + if (.Platform$OS.type == "windows") { + options$eval = FALSE + } + options +}) +r_tmp_dir <- tempdir() +tmp_arg <- paste0("-Djava.io.tmpdir=", r_tmp_dir) +sparkSessionConfig <- list(spark.driver.extraJavaOptions = tmp_arg, + spark.executor.extraJavaOptions = tmp_arg) +old_java_opt <- Sys.getenv("_JAVA_OPTIONS") +Sys.setenv("_JAVA_OPTIONS" = paste("-XX:-UsePerfData", old_java_opt, sep = " ")) +``` + +## Overview + +SparkR is an R package that provides a light-weight frontend to use Apache Spark from R. With Spark `r packageVersion("SparkR")`, SparkR provides a distributed data frame implementation that supports data processing operations like selection, filtering, aggregation etc. and distributed machine learning using [MLlib](https://spark.apache.org/mllib/). + +## Getting Started + +We begin with an example running on the local machine and provide an overview of the use of SparkR: data ingestion, data processing and machine learning. + +First, let's load and attach the package. +```{r, message=FALSE, eval=EVAL_CHUNK} +library(SparkR) +``` + +```{r, include=FALSE, eval=EVAL_CHUNK} +# disable eval if java version not supported +override_eval <- tryCatch(!is.numeric(SparkR:::checkJavaVersion()), + error = function(e) { TRUE }, + warning = function(e) { TRUE }) + +if (override_eval) { + opts_hooks$set(eval = function(options) { + options$eval = FALSE + options + }) +} +``` + +`SparkSession` is the entry point into SparkR which connects your R program to a Spark cluster. You can create a `SparkSession` using `sparkR.session` and pass in options such as the application name, any Spark packages depended on, etc. + +We use default settings in which it runs in local mode. It auto downloads Spark package in the background if no previous installation is found. For more details about setup, see [Spark Session](#SetupSparkSession). + +```{r, include=FALSE, eval=EVAL_CHUNK} +install.spark() +sparkR.session(master = "local[1]", sparkConfig = sparkSessionConfig, enableHiveSupport = FALSE) +``` + +```{r, eval=EVAL_CHUNK} +sparkR.session() +``` + +The operations in SparkR are centered around an R class called `SparkDataFrame`. It is a distributed collection of data organized into named columns, which is conceptually equivalent to a table in a relational database or a data frame in R, but with richer optimizations under the hood. + +`SparkDataFrame` can be constructed from a wide array of sources such as: structured data files, tables in Hive, external databases, or existing local R data frames. For example, we create a `SparkDataFrame` from a local R data frame, + +```{r, eval=EVAL_CHUNK} +cars <- cbind(model = rownames(mtcars), mtcars) +carsDF <- createDataFrame(cars) +``` + +We can view the first few rows of the `SparkDataFrame` by `head` or `showDF` function. +```{r, eval=EVAL_CHUNK} +head(carsDF) +``` + +Common data processing operations such as `filter` and `select` are supported on the `SparkDataFrame`. +```{r, eval=EVAL_CHUNK} +carsSubDF <- select(carsDF, "model", "mpg", "hp") +carsSubDF <- filter(carsSubDF, carsSubDF$hp >= 200) +head(carsSubDF) +``` + +SparkR can use many common aggregation functions after grouping. + +```{r, eval=EVAL_CHUNK} +carsGPDF <- summarize(groupBy(carsDF, carsDF$gear), count = n(carsDF$gear)) +head(carsGPDF) +``` + +The results `carsDF` and `carsSubDF` are `SparkDataFrame` objects. To convert back to R `data.frame`, we can use `collect`. **Caution**: This can cause your interactive environment to run out of memory, though, because `collect()` fetches the entire distributed `DataFrame` to your client, which is acting as a Spark driver. +```{r, eval=EVAL_CHUNK} +carsGP <- collect(carsGPDF) +class(carsGP) +``` + +SparkR supports a number of commonly used machine learning algorithms. Under the hood, SparkR uses MLlib to train the model. Users can call `summary` to print a summary of the fitted model, `predict` to make predictions on new data, and `write.ml`/`read.ml` to save/load fitted models. + +SparkR supports a subset of R formula operators for model fitting, including ‘~’, ‘.’, ‘:’, ‘+’, and ‘-‘. We use linear regression as an example. +```{r, eval=EVAL_CHUNK} +model <- spark.glm(carsDF, mpg ~ wt + cyl) +``` + +The result matches that returned by R `glm` function applied to the corresponding `data.frame` `mtcars` of `carsDF`. In fact, for Generalized Linear Model, we specifically expose `glm` for `SparkDataFrame` as well so that the above is equivalent to `model <- glm(mpg ~ wt + cyl, data = carsDF)`. + +```{r, eval=EVAL_CHUNK} +summary(model) +``` + +The model can be saved by `write.ml` and loaded back using `read.ml`. +```{r, eval=FALSE} +write.ml(model, path = "/HOME/tmp/mlModel/glmModel") +``` + +In the end, we can stop Spark Session by running +```{r, eval=FALSE} +sparkR.session.stop() +``` + +## Setup + +### Installation + +Different from many other R packages, to use SparkR, you need an additional installation of Apache Spark. The Spark installation will be used to run a backend process that will compile and execute SparkR programs. + +After installing the SparkR package, you can call `sparkR.session` as explained in the previous section to start and it will check for the Spark installation. If you are working with SparkR from an interactive shell (e.g. R, RStudio) then Spark is downloaded and cached automatically if it is not found. Alternatively, we provide an easy-to-use function `install.spark` for running this manually. If you don't have Spark installed on the computer, you may download it from [Apache Spark Website](https://spark.apache.org/downloads.html). + +```{r, eval=FALSE} +install.spark() +``` + +If you already have Spark installed, you don't have to install again and can pass the `sparkHome` argument to `sparkR.session` to let SparkR know where the existing Spark installation is. + +```{r, eval=FALSE} +sparkR.session(sparkHome = "/HOME/spark") +``` + +### Spark Session {#SetupSparkSession} + + +In addition to `sparkHome`, many other options can be specified in `sparkR.session`. For a complete list, see [Starting up: SparkSession](https://spark.apache.org/docs/latest/sparkr.html#starting-up-sparksession) and [SparkR API doc](https://spark.apache.org/docs/latest/api/R/reference/sparkR.session.html). + +In particular, the following Spark driver properties can be set in `sparkConfig`. + +Property Name | Property group | spark-submit equivalent +---------------- | ------------------ | ---------------------- +`spark.driver.memory` | Application Properties | `--driver-memory` +`spark.driver.extraClassPath` | Runtime Environment | `--driver-class-path` +`spark.driver.extraJavaOptions` | Runtime Environment | `--driver-java-options` +`spark.driver.extraLibraryPath` | Runtime Environment | `--driver-library-path` +`spark.kerberos.keytab` | Application Properties | `--keytab` +`spark.kerberos.principal` | Application Properties | `--principal` + +**For Windows users**: Due to different file prefixes across operating systems, to avoid the issue of potential wrong prefix, a current workaround is to specify `spark.sql.warehouse.dir` when starting the `SparkSession`. + +```{r, eval=FALSE} +spark_warehouse_path <- file.path(path.expand('~'), "spark-warehouse") +sparkR.session(spark.sql.warehouse.dir = spark_warehouse_path) +``` + + +#### Cluster Mode +SparkR can connect to remote Spark clusters. [Cluster Mode Overview](https://spark.apache.org/docs/latest/cluster-overview.html) is a good introduction to different Spark cluster modes. + +When connecting SparkR to a remote Spark cluster, make sure that the Spark version and Hadoop version on the machine match the corresponding versions on the cluster. Current SparkR package is compatible with +```{r, echo=FALSE, tidy = TRUE, eval=EVAL_CHUNK} +paste("Spark", packageVersion("SparkR")) +``` +It should be used both on the local computer and on the remote cluster. + +To connect, pass the URL of the master node to `sparkR.session`. A complete list can be seen in [Spark Master URLs](https://spark.apache.org/docs/latest/submitting-applications.html#master-urls). +For example, to connect to a local standalone Spark master, we can call + +```{r, eval=FALSE} +sparkR.session(master = "spark://local:7077") +``` + +For YARN cluster, SparkR supports the client mode with the master set as "yarn". +```{r, eval=FALSE} +sparkR.session(master = "yarn") +``` +Yarn cluster mode is not supported in the current version. + +## Data Import + +### Local Data Frame +The simplest way is to convert a local R data frame into a `SparkDataFrame`. Specifically we can use `as.DataFrame` or `createDataFrame` and pass in the local R data frame to create a `SparkDataFrame`. As an example, the following creates a `SparkDataFrame` based using the `faithful` dataset from R. +```{r, eval=EVAL_CHUNK} +df <- as.DataFrame(faithful) +head(df) +``` + +### Data Sources +SparkR supports operating on a variety of data sources through the `SparkDataFrame` interface. You can check the Spark SQL Programming Guide for more [specific options](https://spark.apache.org/docs/latest/sql-programming-guide.html#manually-specifying-options) that are available for the built-in data sources. + +The general method for creating `SparkDataFrame` from data sources is `read.df`. This method takes in the path for the file to load and the type of data source, and the currently active Spark Session will be used automatically. SparkR supports reading CSV, JSON and Parquet files natively and through Spark Packages you can find data source connectors for popular file formats like Avro. These packages can be added with `sparkPackages` parameter when initializing SparkSession using `sparkR.session`. + +```{r, eval=FALSE} +sparkR.session(sparkPackages = "com.databricks:spark-avro_2.12:3.0.0") +``` + +We can see how to use data sources using an example CSV input file. For more information please refer to SparkR [read.df](https://spark.apache.org/docs/latest/api/R/reference/read.df.html) API documentation. +```{r, eval=FALSE} +df <- read.df(csvPath, "csv", header = "true", inferSchema = "true", na.strings = "NA") +``` + +The data sources API natively supports JSON formatted input files. Note that the file that is used here is not a typical JSON file. Each line in the file must contain a separate, self-contained valid JSON object. As a consequence, a regular multi-line JSON file will most often fail. + +Let's take a look at the first two lines of the raw JSON file used here. + +```{r, eval=EVAL_CHUNK} +filePath <- paste0(sparkR.conf("spark.home"), + "/examples/src/main/resources/people.json") +readLines(filePath, n = 2L) +``` + +We use `read.df` to read that into a `SparkDataFrame`. + +```{r, eval=EVAL_CHUNK} +people <- read.df(filePath, "json") +count(people) +head(people) +``` + +SparkR automatically infers the schema from the JSON file. +```{r, eval=EVAL_CHUNK} +printSchema(people) +``` + +If we want to read multiple JSON files, `read.json` can be used. +```{r, eval=EVAL_CHUNK} +people <- read.json(paste0(Sys.getenv("SPARK_HOME"), + c("/examples/src/main/resources/people.json", + "/examples/src/main/resources/people.json"))) +count(people) +``` + +The data sources API can also be used to save out `SparkDataFrames` into multiple file formats. For example we can save the `SparkDataFrame` from the previous example to a Parquet file using `write.df`. +```{r, eval=FALSE} +write.df(people, path = "people.parquet", source = "parquet", mode = "overwrite") +``` + +### Hive Tables +You can also create SparkDataFrames from Hive tables. To do this we will need to create a SparkSession with Hive support which can access tables in the Hive MetaStore. Note that Spark should have been built with Hive support and more details can be found in the [SQL Programming Guide](https://spark.apache.org/docs/latest/sql-programming-guide.html). In SparkR, by default it will attempt to create a SparkSession with Hive support enabled (`enableHiveSupport = TRUE`). + +```{r, eval=FALSE} +sql("CREATE TABLE IF NOT EXISTS src (key INT, value STRING)") + +txtPath <- paste0(sparkR.conf("spark.home"), "/examples/src/main/resources/kv1.txt") +sqlCMD <- sprintf("LOAD DATA LOCAL INPATH '%s' INTO TABLE src", txtPath) +sql(sqlCMD) + +results <- sql("FROM src SELECT key, value") + +# results is now a SparkDataFrame +head(results) +``` + + +## Data Processing + +**To dplyr users**: SparkR has similar interface as dplyr in data processing. However, some noticeable differences are worth mentioning in the first place. We use `df` to represent a `SparkDataFrame` and `col` to represent the name of column here. + +1. indicate columns. SparkR uses either a character string of the column name or a Column object constructed with `$` to indicate a column. For example, to select `col` in `df`, we can write `select(df, "col")` or `select(df, df$col)`. + +2. describe conditions. In SparkR, the Column object representation can be inserted into the condition directly, or we can use a character string to describe the condition, without referring to the `SparkDataFrame` used. For example, to select rows with value > 1, we can write `filter(df, df$col > 1)` or `filter(df, "col > 1")`. + +Here are more concrete examples. + +dplyr | SparkR +-------- | --------- +`select(mtcars, mpg, hp)` | `select(carsDF, "mpg", "hp")` +`filter(mtcars, mpg > 20, hp > 100)` | `filter(carsDF, carsDF$mpg > 20, carsDF$hp > 100)` + +Other differences will be mentioned in the specific methods. + +We use the `SparkDataFrame` `carsDF` created above. We can get basic information about the `SparkDataFrame`. +```{r, eval=EVAL_CHUNK} +carsDF +``` + +Print out the schema in tree format. +```{r, eval=EVAL_CHUNK} +printSchema(carsDF) +``` + +### SparkDataFrame Operations + +#### Selecting rows, columns + +SparkDataFrames support a number of functions to do structured data processing. Here we include some basic examples and a complete list can be found in the [API](https://spark.apache.org/docs/latest/api/R/index.html) docs: + +You can also pass in column name as strings. +```{r, eval=EVAL_CHUNK} +head(select(carsDF, "mpg")) +``` + +Filter the SparkDataFrame to only retain rows with mpg less than 20 miles/gallon. +```{r, eval=EVAL_CHUNK} +head(filter(carsDF, carsDF$mpg < 20)) +``` + +#### Grouping, Aggregation + +A common flow of grouping and aggregation is + +1. Use `groupBy` or `group_by` with respect to some grouping variables to create a `GroupedData` object + +2. Feed the `GroupedData` object to `agg` or `summarize` functions, with some provided aggregation functions to compute a number within each group. + +A number of widely used functions are supported to aggregate data after grouping, including `avg`, `count_distinct`, `count`, `first`, `kurtosis`, `last`, `max`, `mean`, `min`, `sd`, `skewness`, `stddev_pop`, `stddev_samp`, `sum_distinct`, `sum`, `var_pop`, `var_samp`, `var`. See the [API doc for aggregate functions](https://spark.apache.org/docs/latest/api/R/reference/column_aggregate_functions.html) linked there. + +For example we can compute a histogram of the number of cylinders in the `mtcars` dataset as shown below. + +```{r, eval=EVAL_CHUNK} +numCyl <- summarize(groupBy(carsDF, carsDF$cyl), count = n(carsDF$cyl)) +head(numCyl) +``` + +Use `cube` or `rollup` to compute subtotals across multiple dimensions. + +```{r, eval=EVAL_CHUNK} +mean(cube(carsDF, "cyl", "gear", "am"), "mpg") +``` + +generates groupings for {(`cyl`, `gear`, `am`), (`cyl`, `gear`), (`cyl`), ()}, while + +```{r, eval=EVAL_CHUNK} +mean(rollup(carsDF, "cyl", "gear", "am"), "mpg") +``` + +generates groupings for all possible combinations of grouping columns. + + +#### Operating on Columns + +SparkR also provides a number of functions that can directly applied to columns for data processing and during aggregation. The example below shows the use of basic arithmetic functions. + +```{r, eval=EVAL_CHUNK} +carsDF_km <- carsDF +carsDF_km$kmpg <- carsDF_km$mpg * 1.61 +head(select(carsDF_km, "model", "mpg", "kmpg")) +``` + + +### Window Functions +A window function is a variation of aggregation function. In simple words, + +* aggregation function: `n` to `1` mapping - returns a single value for a group of entries. Examples include `sum`, `count`, `max`. + +* window function: `n` to `n` mapping - returns one value for each entry in the group, but the value may depend on all the entries of the *group*. Examples include `rank`, `lead`, `lag`. + +Formally, the *group* mentioned above is called the *frame*. Every input row can have a unique frame associated with it and the output of the window function on that row is based on the rows confined in that frame. + +Window functions are often used in conjunction with the following functions: `windowPartitionBy`, `windowOrderBy`, `partitionBy`, `orderBy`, `over`. To illustrate this we next look at an example. + +We still use the `mtcars` dataset. The corresponding `SparkDataFrame` is `carsDF`. Suppose for each number of cylinders, we want to calculate the rank of each car in `mpg` within the group. +```{r, eval=EVAL_CHUNK} +carsSubDF <- select(carsDF, "model", "mpg", "cyl") +ws <- orderBy(windowPartitionBy("cyl"), "mpg") +carsRank <- withColumn(carsSubDF, "rank", over(rank(), ws)) +head(carsRank, n = 20L) +``` + +We explain in detail the above steps. + +* `windowPartitionBy` creates a window specification object `WindowSpec` that defines the partition. It controls which rows will be in the same partition as the given row. In this case, rows with the same value in `cyl` will be put in the same partition. `orderBy` further defines the ordering - the position a given row is in the partition. The resulting `WindowSpec` is returned as `ws`. + +More window specification methods include `rangeBetween`, which can define boundaries of the frame by value, and `rowsBetween`, which can define the boundaries by row indices. + +* `withColumn` appends a Column called `rank` to the `SparkDataFrame`. `over` returns a windowing column. The first argument is usually a Column returned by window function(s) such as `rank()`, `lead(carsDF$wt)`. That calculates the corresponding values according to the partitioned-and-ordered table. + +### User-Defined Function + +In SparkR, we support several kinds of user-defined functions (UDFs). + +#### Apply by Partition + +`dapply` can apply a function to each partition of a `SparkDataFrame`. The function to be applied to each partition of the `SparkDataFrame` should have only one parameter, a `data.frame` corresponding to a partition, and the output should be a `data.frame` as well. Schema specifies the row format of the resulting a `SparkDataFrame`. It must match to data types of returned value. See [here](#DataTypes) for mapping between R and Spark. + +We convert `mpg` to `kmpg` (kilometers per gallon). `carsSubDF` is a `SparkDataFrame` with a subset of `carsDF` columns. + +```{r, eval=EVAL_CHUNK} +carsSubDF <- select(carsDF, "model", "mpg") +schema <- "model STRING, mpg DOUBLE, kmpg DOUBLE" +out <- dapply(carsSubDF, function(x) { x <- cbind(x, x$mpg * 1.61) }, schema) +head(collect(out)) +``` + +Like `dapply`, `dapplyCollect` can apply a function to each partition of a `SparkDataFrame` and collect the result back. The output of the function should be a `data.frame`, but no schema is required in this case. Note that `dapplyCollect` can fail if the output of the UDF on all partitions cannot be pulled into the driver's memory. + +```{r, eval=EVAL_CHUNK} +out <- dapplyCollect( + carsSubDF, + function(x) { + x <- cbind(x, "kmpg" = x$mpg * 1.61) + }) +head(out, 3) +``` + +#### Apply by Group +`gapply` can apply a function to each group of a `SparkDataFrame`. The function is to be applied to each group of the `SparkDataFrame` and should have only two parameters: grouping key and R `data.frame` corresponding to that key. The groups are chosen from `SparkDataFrames` column(s). The output of function should be a `data.frame`. Schema specifies the row format of the resulting `SparkDataFrame`. It must represent R function’s output schema on the basis of Spark data types. The column names of the returned `data.frame` are set by user. See [here](#DataTypes) for mapping between R and Spark. + +```{r, eval=EVAL_CHUNK} +schema <- structType(structField("cyl", "double"), structField("max_mpg", "double")) +result <- gapply( + carsDF, + "cyl", + function(key, x) { + y <- data.frame(key, max(x$mpg)) + }, + schema) +head(arrange(result, "max_mpg", decreasing = TRUE)) +``` + +Like `gapply`, `gapplyCollect` can apply a function to each partition of a `SparkDataFrame` and collect the result back to R `data.frame`. The output of the function should be a `data.frame` but no schema is required in this case. Note that `gapplyCollect` can fail if the output of the UDF on all partitions cannot be pulled into the driver's memory. + +```{r, eval=EVAL_CHUNK} +result <- gapplyCollect( + carsDF, + "cyl", + function(key, x) { + y <- data.frame(key, max(x$mpg)) + colnames(y) <- c("cyl", "max_mpg") + y + }) +head(result[order(result$max_mpg, decreasing = TRUE), ]) +``` + +#### Distribute Local Functions + +Similar to `lapply` in native R, `spark.lapply` runs a function over a list of elements and distributes the computations with Spark. `spark.lapply` works in a manner that is similar to `doParallel` or `lapply` to elements of a list. The results of all the computations should fit in a single machine. If that is not the case you can do something like `df <- createDataFrame(list)` and then use `dapply`. + +We use `svm` in package `e1071` as an example. We use all default settings except for varying costs of constraints violation. `spark.lapply` can train those different models in parallel. + +```{r, eval=EVAL_CHUNK} +costs <- exp(seq(from = log(1), to = log(1000), length.out = 5)) +train <- function(cost) { + stopifnot(requireNamespace("e1071", quietly = TRUE)) + model <- e1071::svm(Species ~ ., data = iris, cost = cost) + summary(model) +} +``` + +Return a list of model's summaries. +```{r, eval=EVAL_CHUNK} +model.summaries <- spark.lapply(costs, train) +``` + +```{r, eval=EVAL_CHUNK} +class(model.summaries) +``` + + +To avoid lengthy display, we only present the partial result of the second fitted model. You are free to inspect other models as well. +```{r, include=FALSE, eval=EVAL_CHUNK} +ops <- options() +options(max.print=40) +``` +```{r, eval=EVAL_CHUNK} +print(model.summaries[[2]]) +``` +```{r, include=FALSE, eval=EVAL_CHUNK} +options(ops) +``` + + +### SQL Queries +A `SparkDataFrame` can also be registered as a temporary view in Spark SQL so that one can run SQL queries over its data. The sql function enables applications to run SQL queries programmatically and returns the result as a `SparkDataFrame`. + +```{r, eval=EVAL_CHUNK} +people <- read.df(paste0(sparkR.conf("spark.home"), + "/examples/src/main/resources/people.json"), "json") +``` + +Register this `SparkDataFrame` as a temporary view. + +```{r, eval=EVAL_CHUNK} +createOrReplaceTempView(people, "people") +``` + +SQL statements can be run using the sql method. +```{r, eval=EVAL_CHUNK} +teenagers <- sql("SELECT name FROM people WHERE age >= 13 AND age <= 19") +head(teenagers) +``` + + +## Machine Learning + +SparkR supports the following machine learning models and algorithms. + +#### Classification + +* Linear Support Vector Machine (SVM) Classifier + +* Logistic Regression + +* Multilayer Perceptron (MLP) + +* Naive Bayes + +* Factorization Machines (FM) Classifier + +#### Regression + +* Accelerated Failure Time (AFT) Survival Model + +* Generalized Linear Model (GLM) + +* Isotonic Regression + +* Linear Regression + +* Factorization Machines (FM) Regressor + +#### Tree - Classification and Regression + +* Decision Tree + +* Gradient-Boosted Trees (GBT) + +* Random Forest + +#### Clustering + +* Bisecting $k$-means + +* Gaussian Mixture Model (GMM) + +* $k$-means Clustering + +* Latent Dirichlet Allocation (LDA) + +* Power Iteration Clustering (PIC) + +#### Collaborative Filtering + +* Alternating Least Squares (ALS) + +#### Frequent Pattern Mining + +* FP-growth +* PrefixSpan + +#### Statistics + +* Kolmogorov-Smirnov Test + +### R Formula + +For most above, SparkR supports **R formula operators**, including `~`, `.`, `:`, `+` and `-` for model fitting. This makes it a similar experience as using R functions. + +### Training and Test Sets + +We can easily split `SparkDataFrame` into random training and test sets by the `randomSplit` function. It returns a list of split `SparkDataFrames` with provided `weights`. We use `carsDF` as an example and want to have about $70%$ training data and $30%$ test data. +```{r, eval=EVAL_CHUNK} +splitDF_list <- randomSplit(carsDF, c(0.7, 0.3), seed = 0) +carsDF_train <- splitDF_list[[1]] +carsDF_test <- splitDF_list[[2]] +``` + +```{r, eval=EVAL_CHUNK} +count(carsDF_train) +head(carsDF_train) +``` + +```{r, eval=EVAL_CHUNK} +count(carsDF_test) +head(carsDF_test) +``` + +### Models and Algorithms + +#### Linear Support Vector Machine (SVM) Classifier + +[Linear Support Vector Machine (SVM)](https://en.wikipedia.org/wiki/Support_vector_machine#Linear_SVM) classifier is an SVM classifier with linear kernels. +This is a binary classifier. We use a simple example to show how to use `spark.svmLinear` +for binary classification. + +```{r, eval=EVAL_CHUNK} +# load training data and create a DataFrame +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# fit a Linear SVM classifier model +model <- spark.svmLinear(training, Survived ~ ., regParam = 0.01, maxIter = 10) +summary(model) +``` + +Predict values on training data +```{r, eval=EVAL_CHUNK} +prediction <- predict(model, training) +head(select(prediction, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) +``` + +#### Logistic Regression + +[Logistic regression](https://en.wikipedia.org/wiki/Logistic_regression) is a widely-used model when the response is categorical. It can be seen as a special case of the [Generalized Linear Predictive Model](https://en.wikipedia.org/wiki/Generalized_linear_model). +We provide `spark.logit` on top of `spark.glm` to support logistic regression with advanced hyper-parameters. +It supports both binary and multiclass classification with elastic-net regularization and feature standardization, similar to `glmnet`. + +We use a simple example to demonstrate `spark.logit` usage. In general, there are three steps of using `spark.logit`: +1). Create a dataframe from a proper data source; 2). Fit a logistic regression model using `spark.logit` with a proper parameter setting; +and 3). Obtain the coefficient matrix of the fitted model using `summary` and use the model for prediction with `predict`. + +Binomial logistic regression +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +model <- spark.logit(training, Survived ~ ., regParam = 0.04741301) +summary(model) +``` + +Predict values on training data +```{r, eval=EVAL_CHUNK} +fitted <- predict(model, training) +head(select(fitted, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) +``` + +Multinomial logistic regression against three classes +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# Note in this case, Spark infers it is multinomial logistic regression, so family = "multinomial" is optional. +model <- spark.logit(training, Class ~ ., regParam = 0.07815179) +summary(model) +``` + +#### Multilayer Perceptron + +Multilayer perceptron classifier (MLPC) is a classifier based on the [feedforward artificial neural network](https://en.wikipedia.org/wiki/Feedforward_neural_network). MLPC consists of multiple layers of nodes. Each layer is fully connected to the next layer in the network. Nodes in the input layer represent the input data. All other nodes map inputs to outputs by a linear combination of the inputs with the node’s weights $w$ and bias $b$ and applying an activation function. This can be written in matrix form for MLPC with $K+1$ layers as follows: +$$ +y(x)=f_K(\ldots f_2(w_2^T f_1(w_1^T x + b_1) + b_2) \ldots + b_K). +$$ + +Nodes in intermediate layers use sigmoid (logistic) function: +$$ +f(z_i) = \frac{1}{1+e^{-z_i}}. +$$ + +Nodes in the output layer use softmax function: +$$ +f(z_i) = \frac{e^{z_i}}{\sum_{k=1}^N e^{z_k}}. +$$ + +The number of nodes $N$ in the output layer corresponds to the number of classes. + +MLPC employs backpropagation for learning the model. We use the logistic loss function for optimization and L-BFGS as an optimization routine. + +`spark.mlp` requires at least two columns in `data`: one named `"label"` and the other one `"features"`. The `"features"` column should be in libSVM-format. + +We use Titanic data set to show how to use `spark.mlp` in classification. +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +# fit a Multilayer Perceptron Classification Model +model <- spark.mlp(training, Survived ~ Age + Sex, blockSize = 128, layers = c(2, 2), solver = "l-bfgs", maxIter = 100, tol = 0.5, stepSize = 1, seed = 1, initialWeights = c( 0, 0, 5, 5, 9, 9)) +``` + +To avoid lengthy display, we only present partial results of the model summary. You can check the full result from your sparkR shell. +```{r, include=FALSE, eval=EVAL_CHUNK} +ops <- options() +options(max.print=5) +``` +```{r, eval=EVAL_CHUNK} +# check the summary of the fitted model +summary(model) +``` +```{r, include=FALSE, eval=EVAL_CHUNK} +options(ops) +``` +```{r, eval=EVAL_CHUNK} +# make predictions use the fitted model +predictions <- predict(model, training) +head(select(predictions, predictions$prediction)) +``` + +#### Naive Bayes + +Naive Bayes model assumes independence among the features. `spark.naiveBayes` fits a [Bernoulli naive Bayes model](https://en.wikipedia.org/wiki/Naive_Bayes_classifier#Bernoulli_naive_Bayes) against a SparkDataFrame. The data should be all categorical. These models are often used for document classification. + +```{r, eval=EVAL_CHUNK} +titanic <- as.data.frame(Titanic) +titanicDF <- createDataFrame(titanic[titanic$Freq > 0, -5]) +naiveBayesModel <- spark.naiveBayes(titanicDF, Survived ~ Class + Sex + Age) +summary(naiveBayesModel) +naiveBayesPrediction <- predict(naiveBayesModel, titanicDF) +head(select(naiveBayesPrediction, "Class", "Sex", "Age", "Survived", "prediction")) +``` + +#### Factorization Machines Classifier + +Factorization Machines for classification problems. + +For background and details about the implementation of factorization machines, +refer to the [Factorization Machines section](https://spark.apache.org/docs/latest/ml-classification-regression.html#factorization-machines). + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) + +model <- spark.fmClassifier(training, Survived ~ Age + Sex) +summary(model) + +predictions <- predict(model, training) +head(select(predictions, predictions$prediction)) +``` + +#### Accelerated Failure Time Survival Model + +Survival analysis studies the expected duration of time until an event happens, and often the relationship with risk factors or treatment taken on the subject. In contrast to standard regression analysis, survival modeling has to deal with special characteristics in the data including non-negative survival time and censoring. + +Accelerated Failure Time (AFT) model is a parametric survival model for censored data that assumes the effect of a covariate is to accelerate or decelerate the life course of an event by some constant. For more information, refer to the Wikipedia page [AFT Model](https://en.wikipedia.org/wiki/Accelerated_failure_time_model) and the references there. Different from a [Proportional Hazards Model](https://en.wikipedia.org/wiki/Proportional_hazards_model) designed for the same purpose, the AFT model is easier to parallelize because each instance contributes to the objective function independently. + +```{r, warning=FALSE, eval=EVAL_CHUNK} +library(survival) +ovarianDF <- createDataFrame(ovarian) +aftModel <- spark.survreg(ovarianDF, Surv(futime, fustat) ~ ecog_ps + rx) +summary(aftModel) +aftPredictions <- predict(aftModel, ovarianDF) +head(aftPredictions) +``` + +#### Generalized Linear Model + +The main function is `spark.glm`. The following families and link functions are supported. The default is gaussian. + +Family | Link Function +------ | --------- +gaussian | identity, log, inverse +binomial | logit, probit, cloglog (complementary log-log) +poisson | log, identity, sqrt +gamma | inverse, identity, log +tweedie | power link function + +There are three ways to specify the `family` argument. + +* Family name as a character string, e.g. `family = "gaussian"`. + +* Family function, e.g. `family = binomial`. + +* Result returned by a family function, e.g. `family = poisson(link = log)`. + +* Note that there are two ways to specify the tweedie family: + a) Set `family = "tweedie"` and specify the `var.power` and `link.power` + b) When package `statmod` is loaded, the tweedie family is specified using the family definition therein, i.e., `tweedie()`. + +For more information regarding the families and their link functions, see the Wikipedia page [Generalized Linear Model](https://en.wikipedia.org/wiki/Generalized_linear_model). + +We use the `mtcars` dataset as an illustration. The corresponding `SparkDataFrame` is `carsDF`. After fitting the model, we print out a summary and see the fitted values by making predictions on the original dataset. We can also pass into a new `SparkDataFrame` of same schema to predict on new data. + +```{r, eval=EVAL_CHUNK} +gaussianGLM <- spark.glm(carsDF, mpg ~ wt + hp) +summary(gaussianGLM) +``` +When doing prediction, a new column called `prediction` will be appended. Let's look at only a subset of columns here. +```{r, eval=EVAL_CHUNK} +gaussianFitted <- predict(gaussianGLM, carsDF) +head(select(gaussianFitted, "model", "prediction", "mpg", "wt", "hp")) +``` + +The following is the same fit using the tweedie family: +```{r, eval=EVAL_CHUNK} +tweedieGLM1 <- spark.glm(carsDF, mpg ~ wt + hp, family = "tweedie", var.power = 0.0) +summary(tweedieGLM1) +``` +We can try other distributions in the tweedie family, for example, a compound Poisson distribution with a log link: +```{r, eval=EVAL_CHUNK} +tweedieGLM2 <- spark.glm(carsDF, mpg ~ wt + hp, family = "tweedie", + var.power = 1.2, link.power = 0.0) +summary(tweedieGLM2) +``` + +#### Isotonic Regression + +`spark.isoreg` fits an [Isotonic Regression](https://en.wikipedia.org/wiki/Isotonic_regression) model against a `SparkDataFrame`. It solves a weighted univariate a regression problem under a complete order constraint. Specifically, given a set of real observed responses $y_1, \ldots, y_n$, corresponding real features $x_1, \ldots, x_n$, and optionally positive weights $w_1, \ldots, w_n$, we want to find a monotone (piecewise linear) function $f$ to minimize +$$ +\ell(f) = \sum_{i=1}^n w_i (y_i - f(x_i))^2. +$$ + +There are a few more arguments that may be useful. + +* `weightCol`: a character string specifying the weight column. + +* `isotonic`: logical value indicating whether the output sequence should be isotonic/increasing (`TRUE`) or antitonic/decreasing (`FALSE`). + +* `featureIndex`: the index of the feature on the right hand side of the formula if it is a vector column (default: 0), no effect otherwise. + +We use an artificial example to show the use. + +```{r, eval=EVAL_CHUNK} +y <- c(3.0, 6.0, 8.0, 5.0, 7.0) +x <- c(1.0, 2.0, 3.5, 3.0, 4.0) +w <- rep(1.0, 5) +data <- data.frame(y = y, x = x, w = w) +df <- createDataFrame(data) +isoregModel <- spark.isoreg(df, y ~ x, weightCol = "w") +isoregFitted <- predict(isoregModel, df) +head(select(isoregFitted, "x", "y", "prediction")) +``` + +In the prediction stage, based on the fitted monotone piecewise function, the rules are: + +* If the prediction input exactly matches a training feature then associated prediction is returned. In case there are multiple predictions with the same feature then one of them is returned. Which one is undefined. + +* If the prediction input is lower or higher than all training features then prediction with lowest or highest feature is returned respectively. In case there are multiple predictions with the same feature then the lowest or highest is returned respectively. + +* If the prediction input falls between two training features then prediction is treated as piecewise linear function and interpolated value is calculated from the predictions of the two closest features. In case there are multiple values with the same feature then the same rules as in previous point are used. + +For example, when the input is $3.2$, the two closest feature values are $3.0$ and $3.5$, then predicted value would be a linear interpolation between the predicted values at $3.0$ and $3.5$. + +```{r, eval=EVAL_CHUNK} +newDF <- createDataFrame(data.frame(x = c(1.5, 3.2))) +head(predict(isoregModel, newDF)) +``` + +#### Linear Regression + +Linear regression model. + +```{r, eval=EVAL_CHUNK} +model <- spark.lm(carsDF, mpg ~ wt + hp) + +summary(model) +predictions <- predict(model, carsDF) +head(select(predictions, predictions$prediction)) +``` + +#### Factorization Machines Regressor + +Factorization Machines for regression problems. + +For background and details about the implementation of factorization machines, +refer to the [Factorization Machines section](https://spark.apache.org/docs/latest/ml-classification-regression.html#factorization-machines). + +```{r, eval=EVAL_CHUNK} +model <- spark.fmRegressor(carsDF, mpg ~ wt + hp) +summary(model) +predictions <- predict(model, carsDF) +head(select(predictions, predictions$prediction)) +``` + +#### Decision Tree + +`spark.decisionTree` fits a [decision tree](https://en.wikipedia.org/wiki/Decision_tree_learning) classification or regression model on a `SparkDataFrame`. +Users can call `summary` to get a summary of the fitted model, `predict` to make predictions, and `write.ml`/`read.ml` to save/load fitted models. + +We use the `Titanic` dataset to train a decision tree and make predictions: + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +dtModel <- spark.decisionTree(df, Survived ~ ., type = "classification", maxDepth = 2) +summary(dtModel) +predictions <- predict(dtModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) +``` + +#### Gradient-Boosted Trees + +`spark.gbt` fits a [gradient-boosted tree](https://en.wikipedia.org/wiki/Gradient_boosting) classification or regression model on a `SparkDataFrame`. +Users can call `summary` to get a summary of the fitted model, `predict` to make predictions, and `write.ml`/`read.ml` to save/load fitted models. + +We use the `Titanic` dataset to train a gradient-boosted tree and make predictions: + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +gbtModel <- spark.gbt(df, Survived ~ ., type = "classification", maxDepth = 2, maxIter = 2) +summary(gbtModel) +predictions <- predict(gbtModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) +``` + +#### Random Forest + +`spark.randomForest` fits a [random forest](https://en.wikipedia.org/wiki/Random_forest) classification or regression model on a `SparkDataFrame`. +Users can call `summary` to get a summary of the fitted model, `predict` to make predictions, and `write.ml`/`read.ml` to save/load fitted models. + +In the following example, we use the `Titanic` dataset to train a random forest and make predictions: + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +rfModel <- spark.randomForest(df, Survived ~ ., type = "classification", maxDepth = 2, numTrees = 2) +summary(rfModel) +predictions <- predict(rfModel, df) +head(select(predictions, "Class", "Sex", "Age", "Freq", "Survived", "prediction")) +``` + +#### Bisecting k-Means + +`spark.bisectingKmeans` is a kind of [hierarchical clustering](https://en.wikipedia.org/wiki/Hierarchical_clustering) using a divisive (or "top-down") approach: all observations start in one cluster, and splits are performed recursively as one moves down the hierarchy. + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +model <- spark.bisectingKmeans(training, Class ~ Survived, k = 4) +summary(model) +fitted <- predict(model, training) +head(select(fitted, "Class", "prediction")) +``` + +#### Gaussian Mixture Model + +`spark.gaussianMixture` fits multivariate [Gaussian Mixture Model](https://en.wikipedia.org/wiki/Mixture_model#Multivariate_Gaussian_mixture_model) (GMM) against a `SparkDataFrame`. [Expectation-Maximization](https://en.wikipedia.org/wiki/Expectation%E2%80%93maximization_algorithm) (EM) is used to approximate the maximum likelihood estimator (MLE) of the model. + +We use a simulated example to demonstrate the usage. +```{r, eval=EVAL_CHUNK} +X1 <- data.frame(V1 = rnorm(4), V2 = rnorm(4)) +X2 <- data.frame(V1 = rnorm(6, 3), V2 = rnorm(6, 4)) +data <- rbind(X1, X2) +df <- createDataFrame(data) +gmmModel <- spark.gaussianMixture(df, ~ V1 + V2, k = 2) +summary(gmmModel) +gmmFitted <- predict(gmmModel, df) +head(select(gmmFitted, "V1", "V2", "prediction")) +``` + +#### k-Means Clustering + +`spark.kmeans` fits a $k$-means clustering model against a `SparkDataFrame`. As an unsupervised learning method, we don't need a response variable. Hence, the left hand side of the R formula should be left blank. The clustering is based only on the variables on the right hand side. + +```{r, eval=EVAL_CHUNK} +kmeansModel <- spark.kmeans(carsDF, ~ mpg + hp + wt, k = 3) +summary(kmeansModel) +kmeansPredictions <- predict(kmeansModel, carsDF) +head(select(kmeansPredictions, "model", "mpg", "hp", "wt", "prediction"), n = 20L) +``` + +#### Latent Dirichlet Allocation + +`spark.lda` fits a [Latent Dirichlet Allocation](https://en.wikipedia.org/wiki/Latent_Dirichlet_allocation) model on a `SparkDataFrame`. It is often used in topic modeling in which topics are inferred from a collection of text documents. LDA can be thought of as a clustering algorithm as follows: + +* Topics correspond to cluster centers, and documents correspond to examples (rows) in a dataset. + +* Topics and documents both exist in a feature space, where feature vectors are vectors of word counts (bag of words). + +* Rather than clustering using a traditional distance, LDA uses a function based on a statistical model of how text documents are generated. + +To use LDA, we need to specify a `features` column in `data` where each entry represents a document. There are two options for the column: + +* character string: This can be a string of the whole document. It will be parsed automatically. Additional stop words can be added in `customizedStopWords`. + +* libSVM: Each entry is a collection of words and will be processed directly. + +Two more functions are provided for the fitted model. + +* `spark.posterior` returns a `SparkDataFrame` containing a column of posterior probabilities vectors named "topicDistribution". + +* `spark.perplexity` returns the log perplexity of given `SparkDataFrame`, or the log perplexity of the training data if missing argument `data`. + +For more information, see the help document `?spark.lda`. + +Let's look an artificial example. +```{r, eval=EVAL_CHUNK} +corpus <- data.frame(features = c( + "1 2 6 0 2 3 1 1 0 0 3", + "1 3 0 1 3 0 0 2 0 0 1", + "1 4 1 0 0 4 9 0 1 2 0", + "2 1 0 3 0 0 5 0 2 3 9", + "3 1 1 9 3 0 2 0 0 1 3", + "4 2 0 3 4 5 1 1 1 4 0", + "2 1 0 3 0 0 5 0 2 2 9", + "1 1 1 9 2 1 2 0 0 1 3", + "4 4 0 3 4 2 1 3 0 0 0", + "2 8 2 0 3 0 2 0 2 7 2", + "1 1 1 9 0 2 2 0 0 3 3", + "4 1 0 0 4 5 1 3 0 1 0")) +corpusDF <- createDataFrame(corpus) +model <- spark.lda(data = corpusDF, k = 5, optimizer = "em") +summary(model) +``` + +```{r, eval=EVAL_CHUNK} +posterior <- spark.posterior(model, corpusDF) +head(posterior) +``` + +```{r, eval=EVAL_CHUNK} +perplexity <- spark.perplexity(model, corpusDF) +perplexity +``` + +#### Alternating Least Squares + +`spark.als` learns latent factors in [collaborative filtering](https://en.wikipedia.org/wiki/Recommender_system#Collaborative_filtering) via [alternating least squares](https://dl.acm.org/doi/10.1109/MC.2009.263). + +There are multiple options that can be configured in `spark.als`, including `rank`, `reg`, and `nonnegative`. For a complete list, refer to the help file. + +```{r, eval=FALSE} +ratings <- list(list(0, 0, 4.0), list(0, 1, 2.0), list(1, 1, 3.0), list(1, 2, 4.0), + list(2, 1, 1.0), list(2, 2, 5.0)) +df <- createDataFrame(ratings, c("user", "item", "rating")) +model <- spark.als(df, "rating", "user", "item", rank = 10, reg = 0.1, nonnegative = TRUE) +``` + +Extract latent factors. +```{r, eval=FALSE} +stats <- summary(model) +userFactors <- stats$userFactors +itemFactors <- stats$itemFactors +head(userFactors) +head(itemFactors) +``` + +Make predictions. + +```{r, eval=FALSE} +predicted <- predict(model, df) +head(predicted) +``` + +#### Power Iteration Clustering + +Power Iteration Clustering (PIC) is a scalable graph clustering algorithm. `spark.assignClusters` method runs the PIC algorithm and returns a cluster assignment for each input vertex. + +```{r, eval=EVAL_CHUNK} +df <- createDataFrame(list(list(0L, 1L, 1.0), list(0L, 2L, 1.0), + list(1L, 2L, 1.0), list(3L, 4L, 1.0), + list(4L, 0L, 0.1)), + schema = c("src", "dst", "weight")) +head(spark.assignClusters(df, initMode = "degree", weightCol = "weight")) +``` + +#### FP-growth + +`spark.fpGrowth` executes FP-growth algorithm to mine frequent itemsets on a `SparkDataFrame`. `itemsCol` should be an array of values. + +```{r, eval=EVAL_CHUNK} +df <- selectExpr(createDataFrame(data.frame(rawItems = c( + "T,R,U", "T,S", "V,R", "R,U,T,V", "R,S", "V,S,U", "U,R", "S,T", "V,R", "V,U,S", + "T,V,U", "R,V", "T,S", "T,S", "S,T", "S,U", "T,R", "V,R", "S,V", "T,S,U" +))), "split(rawItems, ',') AS items") + +fpm <- spark.fpGrowth(df, minSupport = 0.2, minConfidence = 0.5) +``` + +`spark.freqItemsets` method can be used to retrieve a `SparkDataFrame` with the frequent itemsets. + +```{r, eval=EVAL_CHUNK} +head(spark.freqItemsets(fpm)) +``` + +`spark.associationRules` returns a `SparkDataFrame` with the association rules. + +```{r, eval=EVAL_CHUNK} +head(spark.associationRules(fpm)) +``` + +We can make predictions based on the `antecedent`. + +```{r, eval=EVAL_CHUNK} +head(predict(fpm, df)) +``` + +#### PrefixSpan + +`spark.findFrequentSequentialPatterns` method can be used to find the complete set of frequent sequential patterns in the input sequences of itemsets. + +```{r, eval=EVAL_CHUNK} +df <- createDataFrame(list(list(list(list(1L, 2L), list(3L))), + list(list(list(1L), list(3L, 2L), list(1L, 2L))), + list(list(list(1L, 2L), list(5L))), + list(list(list(6L)))), + schema = c("sequence")) +head(spark.findFrequentSequentialPatterns(df, minSupport = 0.5, maxPatternLength = 5L)) +``` + +#### Kolmogorov-Smirnov Test + +`spark.kstest` runs a two-sided, one-sample [Kolmogorov-Smirnov (KS) test](https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Smirnov_test). +Given a `SparkDataFrame`, the test compares continuous data in a given column `testCol` with the theoretical distribution +specified by parameter `nullHypothesis`. +Users can call `summary` to get a summary of the test results. + +In the following example, we test whether the `Titanic` dataset's `Freq` column +follows a normal distribution. We set the parameters of the normal distribution using +the mean and standard deviation of the sample. + +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +df <- createDataFrame(t) +freqStats <- head(select(df, mean(df$Freq), sd(df$Freq))) +freqMean <- freqStats[1] +freqStd <- freqStats[2] + +test <- spark.kstest(df, "Freq", "norm", c(freqMean, freqStd)) +testSummary <- summary(test) +testSummary +``` + + +### Model Persistence +The following example shows how to save/load an ML model in SparkR. +```{r, eval=EVAL_CHUNK} +t <- as.data.frame(Titanic) +training <- createDataFrame(t) +gaussianGLM <- spark.glm(training, Freq ~ Sex + Age, family = "gaussian") + +# Save and then load a fitted MLlib model +modelPath <- tempfile(pattern = "ml", fileext = ".tmp") +write.ml(gaussianGLM, modelPath) +gaussianGLM2 <- read.ml(modelPath) + +# Check model summary +summary(gaussianGLM2) + +# Check model prediction +gaussianPredictions <- predict(gaussianGLM2, training) +head(gaussianPredictions) + +unlink(modelPath) +``` + + +## Structured Streaming + +SparkR supports the Structured Streaming API. + +You can check the Structured Streaming Programming Guide for [an introduction](https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html#programming-model) to its programming model and basic concepts. + +### Simple Source and Sink + +Spark has a few built-in input sources. As an example, to test with a socket source reading text into words and displaying the computed word counts: + +```{r, eval=FALSE} +# Create DataFrame representing the stream of input lines from connection +lines <- read.stream("socket", host = hostname, port = port) + +# Split the lines into words +words <- selectExpr(lines, "explode(split(value, ' ')) as word") + +# Generate running word count +wordCounts <- count(groupBy(words, "word")) + +# Start running the query that prints the running counts to the console +query <- write.stream(wordCounts, "console", outputMode = "complete") +``` + +### Kafka Source + +It is simple to read data from Kafka. For more information, see [Input Sources](https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html#input-sources) supported by Structured Streaming. + +```{r, eval=FALSE} +topic <- read.stream("kafka", + kafka.bootstrap.servers = "host1:port1,host2:port2", + subscribe = "topic1") +keyvalue <- selectExpr(topic, "CAST(key AS STRING)", "CAST(value AS STRING)") +``` + +### Operations and Sinks + +Most of the common operations on `SparkDataFrame` are supported for streaming, including selection, projection, and aggregation. Once you have defined the final result, to start the streaming computation, you will call the `write.stream` method setting a sink and `outputMode`. + +A streaming `SparkDataFrame` can be written for debugging to the console, to a temporary in-memory table, or for further processing in a fault-tolerant manner to a File Sink in different formats. + +```{r, eval=FALSE} +noAggDF <- select(where(deviceDataStreamingDf, "signal > 10"), "device") + +# Print new data to console +write.stream(noAggDF, "console") + +# Write new data to Parquet files +write.stream(noAggDF, + "parquet", + path = "path/to/destination/dir", + checkpointLocation = "path/to/checkpoint/dir") + +# Aggregate +aggDF <- count(groupBy(noAggDF, "device")) + +# Print updated aggregations to console +write.stream(aggDF, "console", outputMode = "complete") + +# Have all the aggregates in an in memory table. The query name will be the table name +write.stream(aggDF, "memory", queryName = "aggregates", outputMode = "complete") + +head(sql("select * from aggregates")) +``` + + +## Advanced Topics + +### SparkR Object Classes + +There are three main object classes in SparkR you may be working with. + +* `SparkDataFrame`: the central component of SparkR. It is an S4 class representing distributed collection of data organized into named columns, which is conceptually equivalent to a table in a relational database or a data frame in R. It has two slots `sdf` and `env`. + + `sdf` stores a reference to the corresponding Spark Dataset in the Spark JVM backend. + + `env` saves the meta-information of the object such as `isCached`. + + It can be created by data import methods or by transforming an existing `SparkDataFrame`. We can manipulate `SparkDataFrame` by numerous data processing functions and feed that into machine learning algorithms. + +* `Column`: an S4 class representing a column of `SparkDataFrame`. The slot `jc` saves a reference to the corresponding `Column` object in the Spark JVM backend. + + It can be obtained from a `SparkDataFrame` by `$` operator, e.g., `df$col`. More often, it is used together with other functions, for example, with `select` to select particular columns, with `filter` and constructed conditions to select rows, with aggregation functions to compute aggregate statistics for each group. + +* `GroupedData`: an S4 class representing grouped data created by `groupBy` or by transforming other `GroupedData`. Its `sgd` slot saves a reference to a `RelationalGroupedDataset` object in the backend. + + This is often an intermediate object with group information and followed up by aggregation operations. + +### Architecture + +A complete description of architecture can be seen in the references, in particular the paper *SparkR: Scaling R Programs with Spark*. + +Under the hood of SparkR is Spark SQL engine. This avoids the overheads of running interpreted R code, and the optimized SQL execution engine in Spark uses structural information about data and computation flow to perform a bunch of optimizations to speed up the computation. + +The main method calls of actual computation happen in the Spark JVM of the driver. We have a socket-based SparkR API that allows us to invoke functions on the JVM from R. We use a SparkR JVM backend that listens on a Netty-based socket server. + +Two kinds of RPCs are supported in the SparkR JVM backend: method invocation and creating new objects. Method invocation can be done in two ways. + +* `sparkR.callJMethod` takes a reference to an existing Java object and a list of arguments to be passed on to the method. + +* `sparkR.callJStatic` takes a class name for static method and a list of arguments to be passed on to the method. + +The arguments are serialized using our custom wire format which is then deserialized on the JVM side. We then use Java reflection to invoke the appropriate method. + +To create objects, `sparkR.newJObject` is used and then similarly the appropriate constructor is invoked with provided arguments. + +Finally, we use a new R class `jobj` that refers to a Java object existing in the backend. These references are tracked on the Java side and are automatically garbage collected when they go out of scope on the R side. + +## Appendix + +### R and Spark Data Types {#DataTypes} + +R | Spark +----------- | ------------- +byte | byte +integer | integer +float | float +double | double +numeric | double +character | string +string | string +binary | binary +raw | binary +logical | boolean +POSIXct | timestamp +POSIXlt | timestamp +Date | date +array | array +list | array +env | map + +## References + +* [Spark Cluster Mode Overview](https://spark.apache.org/docs/latest/cluster-overview.html) + +* [Submitting Spark Applications](https://spark.apache.org/docs/latest/submitting-applications.html) + +* [Machine Learning Library Guide (MLlib)](https://spark.apache.org/docs/latest/ml-guide.html) + +* [SparkR: Scaling R Programs with Spark](https://people.csail.mit.edu/matei/papers/2016/sigmod_sparkr.pdf), Shivaram Venkataraman, Zongheng Yang, Davies Liu, Eric Liang, Hossein Falaki, Xiangrui Meng, Reynold Xin, Ali Ghodsi, Michael Franklin, Ion Stoica, and Matei Zaharia. SIGMOD 2016. June 2016. + +```{r, echo=FALSE, eval=EVAL_CHUNK} +sparkR.session.stop() +``` + +```{r cleanup, include=FALSE, eval=EVAL_CHUNK} +SparkR:::uninstallDownloadedSpark() +``` diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.html b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.html new file mode 100644 index 0000000..2e842ef --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/doc/sparkr-vignettes.html @@ -0,0 +1,2309 @@ + + + + + + + + + + + + + + +SparkR - Practical Guide + + + + + + + + + + + + + + + + + + + + + + + + + +

SparkR - Practical Guide

+ + + + + +
+

Overview

+

SparkR is an R package that provides a light-weight frontend to use Apache Spark from R. With Spark 3.5.0, SparkR provides a distributed data frame implementation that supports data processing operations like selection, filtering, aggregation etc. and distributed machine learning using MLlib.

+
+
+

Getting Started

+

We begin with an example running on the local machine and provide an overview of the use of SparkR: data ingestion, data processing and machine learning.

+

First, let’s load and attach the package.

+ +

SparkSession is the entry point into SparkR which connects your R program to a Spark cluster. You can create a SparkSession using sparkR.session and pass in options such as the application name, any Spark packages depended on, etc.

+

We use default settings in which it runs in local mode. It auto downloads Spark package in the background if no previous installation is found. For more details about setup, see Spark Session.

+ +
## Java ref type org.apache.spark.sql.SparkSession id 1
+

The operations in SparkR are centered around an R class called SparkDataFrame. It is a distributed collection of data organized into named columns, which is conceptually equivalent to a table in a relational database or a data frame in R, but with richer optimizations under the hood.

+

SparkDataFrame can be constructed from a wide array of sources such as: structured data files, tables in Hive, external databases, or existing local R data frames. For example, we create a SparkDataFrame from a local R data frame,

+ +

We can view the first few rows of the SparkDataFrame by head or showDF function.

+ +
##               model  mpg cyl disp  hp drat    wt  qsec vs am gear carb
+## 1         Mazda RX4 21.0   6  160 110 3.90 2.620 16.46  0  1    4    4
+## 2     Mazda RX4 Wag 21.0   6  160 110 3.90 2.875 17.02  0  1    4    4
+## 3        Datsun 710 22.8   4  108  93 3.85 2.320 18.61  1  1    4    1
+## 4    Hornet 4 Drive 21.4   6  258 110 3.08 3.215 19.44  1  0    3    1
+## 5 Hornet Sportabout 18.7   8  360 175 3.15 3.440 17.02  0  0    3    2
+## 6           Valiant 18.1   6  225 105 2.76 3.460 20.22  1  0    3    1
+

Common data processing operations such as filter and select are supported on the SparkDataFrame.

+ +
##                 model  mpg  hp
+## 1          Duster 360 14.3 245
+## 2  Cadillac Fleetwood 10.4 205
+## 3 Lincoln Continental 10.4 215
+## 4   Chrysler Imperial 14.7 230
+## 5          Camaro Z28 13.3 245
+## 6      Ford Pantera L 15.8 264
+

SparkR can use many common aggregation functions after grouping.

+ +
##   gear count
+## 1    4    12
+## 2    3    15
+## 3    5     5
+

The results carsDF and carsSubDF are SparkDataFrame objects. To convert back to R data.frame, we can use collect. Caution: This can cause your interactive environment to run out of memory, though, because collect() fetches the entire distributed DataFrame to your client, which is acting as a Spark driver.

+ +
## [1] "data.frame"
+

SparkR supports a number of commonly used machine learning algorithms. Under the hood, SparkR uses MLlib to train the model. Users can call summary to print a summary of the fitted model, predict to make predictions on new data, and write.ml/read.ml to save/load fitted models.

+

SparkR supports a subset of R formula operators for model fitting, including ‘~’, ‘.’, ‘:’, ‘+’, and ‘-‘. We use linear regression as an example.

+ +

The result matches that returned by R glm function applied to the corresponding data.frame mtcars of carsDF. In fact, for Generalized Linear Model, we specifically expose glm for SparkDataFrame as well so that the above is equivalent to model <- glm(mpg ~ wt + cyl, data = carsDF).

+ +
## 
+## Deviance Residuals: 
+## (Note: These are approximate quantiles with relative error <= 0.01)
+##     Min       1Q   Median       3Q      Max  
+## -4.2893  -1.7085  -0.4713   1.5729   6.1004  
+## 
+## Coefficients:
+##              Estimate  Std. Error  t value    Pr(>|t|)
+## (Intercept)   39.6863     1.71498  23.1409  0.00000000
+## wt            -3.1910     0.75691  -4.2158  0.00022202
+## cyl           -1.5078     0.41469  -3.6360  0.00106428
+## 
+## (Dispersion parameter for gaussian family taken to be 6.592137)
+## 
+##     Null deviance: 1126.05  on 31  degrees of freedom
+## Residual deviance:  191.17  on 29  degrees of freedom
+## AIC: 156
+## 
+## Number of Fisher Scoring iterations: 1
+

The model can be saved by write.ml and loaded back using read.ml.

+ +

In the end, we can stop Spark Session by running

+ +
+
+

Setup

+
+

Installation

+

Different from many other R packages, to use SparkR, you need an additional installation of Apache Spark. The Spark installation will be used to run a backend process that will compile and execute SparkR programs.

+

After installing the SparkR package, you can call sparkR.session as explained in the previous section to start and it will check for the Spark installation. If you are working with SparkR from an interactive shell (e.g. R, RStudio) then Spark is downloaded and cached automatically if it is not found. Alternatively, we provide an easy-to-use function install.spark for running this manually. If you don’t have Spark installed on the computer, you may download it from Apache Spark Website.

+ +

If you already have Spark installed, you don’t have to install again and can pass the sparkHome argument to sparkR.session to let SparkR know where the existing Spark installation is.

+ +
+
+

Spark Session

+

In addition to sparkHome, many other options can be specified in sparkR.session. For a complete list, see Starting up: SparkSession and SparkR API doc.

+

In particular, the following Spark driver properties can be set in sparkConfig.

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Property NameProperty groupspark-submit equivalent
spark.driver.memoryApplication Properties--driver-memory
spark.driver.extraClassPathRuntime Environment--driver-class-path
spark.driver.extraJavaOptionsRuntime Environment--driver-java-options
spark.driver.extraLibraryPathRuntime Environment--driver-library-path
spark.kerberos.keytabApplication Properties--keytab
spark.kerberos.principalApplication Properties--principal
+

For Windows users: Due to different file prefixes across operating systems, to avoid the issue of potential wrong prefix, a current workaround is to specify spark.sql.warehouse.dir when starting the SparkSession.

+ +
+

Cluster Mode

+

SparkR can connect to remote Spark clusters. Cluster Mode Overview is a good introduction to different Spark cluster modes.

+

When connecting SparkR to a remote Spark cluster, make sure that the Spark version and Hadoop version on the machine match the corresponding versions on the cluster. Current SparkR package is compatible with

+
## [1] "Spark 3.5.0"
+

It should be used both on the local computer and on the remote cluster.

+

To connect, pass the URL of the master node to sparkR.session. A complete list can be seen in Spark Master URLs. For example, to connect to a local standalone Spark master, we can call

+ +

For YARN cluster, SparkR supports the client mode with the master set as “yarn”.

+ +

Yarn cluster mode is not supported in the current version.

+
+
+
+
+

Data Import

+
+

Local Data Frame

+

The simplest way is to convert a local R data frame into a SparkDataFrame. Specifically we can use as.DataFrame or createDataFrame and pass in the local R data frame to create a SparkDataFrame. As an example, the following creates a SparkDataFrame based using the faithful dataset from R.

+ +
##   eruptions waiting
+## 1     3.600      79
+## 2     1.800      54
+## 3     3.333      74
+## 4     2.283      62
+## 5     4.533      85
+## 6     2.883      55
+
+
+

Data Sources

+

SparkR supports operating on a variety of data sources through the SparkDataFrame interface. You can check the Spark SQL Programming Guide for more specific options that are available for the built-in data sources.

+

The general method for creating SparkDataFrame from data sources is read.df. This method takes in the path for the file to load and the type of data source, and the currently active Spark Session will be used automatically. SparkR supports reading CSV, JSON and Parquet files natively and through Spark Packages you can find data source connectors for popular file formats like Avro. These packages can be added with sparkPackages parameter when initializing SparkSession using sparkR.session.

+ +

We can see how to use data sources using an example CSV input file. For more information please refer to SparkR read.df API documentation.

+ +

The data sources API natively supports JSON formatted input files. Note that the file that is used here is not a typical JSON file. Each line in the file must contain a separate, self-contained valid JSON object. As a consequence, a regular multi-line JSON file will most often fail.

+

Let’s take a look at the first two lines of the raw JSON file used here.

+ +
## [1] "{\"name\":\"Michael\"}"          "{\"name\":\"Andy\", \"age\":30}"
+

We use read.df to read that into a SparkDataFrame.

+ +
## [1] 3
+ +
##   age    name
+## 1  NA Michael
+## 2  30    Andy
+## 3  19  Justin
+

SparkR automatically infers the schema from the JSON file.

+ +
## root
+##  |-- age: long (nullable = true)
+##  |-- name: string (nullable = true)
+

If we want to read multiple JSON files, read.json can be used.

+ +
## [1] 6
+

The data sources API can also be used to save out SparkDataFrames into multiple file formats. For example we can save the SparkDataFrame from the previous example to a Parquet file using write.df.

+ +
+
+

Hive Tables

+

You can also create SparkDataFrames from Hive tables. To do this we will need to create a SparkSession with Hive support which can access tables in the Hive MetaStore. Note that Spark should have been built with Hive support and more details can be found in the SQL Programming Guide. In SparkR, by default it will attempt to create a SparkSession with Hive support enabled (enableHiveSupport = TRUE).

+ +
+
+
+

Data Processing

+

To dplyr users: SparkR has similar interface as dplyr in data processing. However, some noticeable differences are worth mentioning in the first place. We use df to represent a SparkDataFrame and col to represent the name of column here.

+
    +
  1. indicate columns. SparkR uses either a character string of the column name or a Column object constructed with $ to indicate a column. For example, to select col in df, we can write select(df, "col") or select(df, df$col).

  2. +
  3. describe conditions. In SparkR, the Column object representation can be inserted into the condition directly, or we can use a character string to describe the condition, without referring to the SparkDataFrame used. For example, to select rows with value > 1, we can write filter(df, df$col > 1) or filter(df, "col > 1").

  4. +
+

Here are more concrete examples.

+ ++++ + + + + + + + + + + + + + + + + +
dplyrSparkR
select(mtcars, mpg, hp)select(carsDF, "mpg", "hp")
filter(mtcars, mpg > 20, hp > 100)filter(carsDF, carsDF$mpg > 20, carsDF$hp > 100)
+

Other differences will be mentioned in the specific methods.

+

We use the SparkDataFrame carsDF created above. We can get basic information about the SparkDataFrame.

+ +
## SparkDataFrame[model:string, mpg:double, cyl:double, disp:double, hp:double, drat:double, wt:double, qsec:double, vs:double, am:double, gear:double, carb:double]
+

Print out the schema in tree format.

+ +
## root
+##  |-- model: string (nullable = true)
+##  |-- mpg: double (nullable = true)
+##  |-- cyl: double (nullable = true)
+##  |-- disp: double (nullable = true)
+##  |-- hp: double (nullable = true)
+##  |-- drat: double (nullable = true)
+##  |-- wt: double (nullable = true)
+##  |-- qsec: double (nullable = true)
+##  |-- vs: double (nullable = true)
+##  |-- am: double (nullable = true)
+##  |-- gear: double (nullable = true)
+##  |-- carb: double (nullable = true)
+
+

SparkDataFrame Operations

+
+

Selecting rows, columns

+

SparkDataFrames support a number of functions to do structured data processing. Here we include some basic examples and a complete list can be found in the API docs:

+

You can also pass in column name as strings.

+ +
##    mpg
+## 1 21.0
+## 2 21.0
+## 3 22.8
+## 4 21.4
+## 5 18.7
+## 6 18.1
+

Filter the SparkDataFrame to only retain rows with mpg less than 20 miles/gallon.

+ +
##               model  mpg cyl  disp  hp drat   wt  qsec vs am gear carb
+## 1 Hornet Sportabout 18.7   8 360.0 175 3.15 3.44 17.02  0  0    3    2
+## 2           Valiant 18.1   6 225.0 105 2.76 3.46 20.22  1  0    3    1
+## 3        Duster 360 14.3   8 360.0 245 3.21 3.57 15.84  0  0    3    4
+## 4          Merc 280 19.2   6 167.6 123 3.92 3.44 18.30  1  0    4    4
+## 5         Merc 280C 17.8   6 167.6 123 3.92 3.44 18.90  1  0    4    4
+## 6        Merc 450SE 16.4   8 275.8 180 3.07 4.07 17.40  0  0    3    3
+
+
+

Grouping, Aggregation

+

A common flow of grouping and aggregation is

+
    +
  1. Use groupBy or group_by with respect to some grouping variables to create a GroupedData object

  2. +
  3. Feed the GroupedData object to agg or summarize functions, with some provided aggregation functions to compute a number within each group.

  4. +
+

A number of widely used functions are supported to aggregate data after grouping, including avg, count_distinct, count, first, kurtosis, last, max, mean, min, sd, skewness, stddev_pop, stddev_samp, sum_distinct, sum, var_pop, var_samp, var. See the API doc for aggregate functions linked there.

+

For example we can compute a histogram of the number of cylinders in the mtcars dataset as shown below.

+ +
##   cyl count
+## 1   8    14
+## 2   4    11
+## 3   6     7
+

Use cube or rollup to compute subtotals across multiple dimensions.

+ +
## SparkDataFrame[cyl:double, gear:double, am:double, avg(mpg):double]
+

generates groupings for {(cyl, gear, am), (cyl, gear), (cyl), ()}, while

+ +
## SparkDataFrame[cyl:double, gear:double, am:double, avg(mpg):double]
+

generates groupings for all possible combinations of grouping columns.

+
+
+

Operating on Columns

+

SparkR also provides a number of functions that can directly applied to columns for data processing and during aggregation. The example below shows the use of basic arithmetic functions.

+ +
##               model  mpg   kmpg
+## 1         Mazda RX4 21.0 33.810
+## 2     Mazda RX4 Wag 21.0 33.810
+## 3        Datsun 710 22.8 36.708
+## 4    Hornet 4 Drive 21.4 34.454
+## 5 Hornet Sportabout 18.7 30.107
+## 6           Valiant 18.1 29.141
+
+
+
+

Window Functions

+

A window function is a variation of aggregation function. In simple words,

+
    +
  • aggregation function: n to 1 mapping - returns a single value for a group of entries. Examples include sum, count, max.

  • +
  • window function: n to n mapping - returns one value for each entry in the group, but the value may depend on all the entries of the group. Examples include rank, lead, lag.

  • +
+

Formally, the group mentioned above is called the frame. Every input row can have a unique frame associated with it and the output of the window function on that row is based on the rows confined in that frame.

+

Window functions are often used in conjunction with the following functions: windowPartitionBy, windowOrderBy, partitionBy, orderBy, over. To illustrate this we next look at an example.

+

We still use the mtcars dataset. The corresponding SparkDataFrame is carsDF. Suppose for each number of cylinders, we want to calculate the rank of each car in mpg within the group.

+ +
##                  model  mpg cyl rank
+## 1           Volvo 142E 21.4   4    1
+## 2        Toyota Corona 21.5   4    2
+## 3           Datsun 710 22.8   4    3
+## 4             Merc 230 22.8   4    3
+## 5            Merc 240D 24.4   4    5
+## 6        Porsche 914-2 26.0   4    6
+## 7            Fiat X1-9 27.3   4    7
+## 8          Honda Civic 30.4   4    8
+## 9         Lotus Europa 30.4   4    8
+## 10            Fiat 128 32.4   4   10
+## 11      Toyota Corolla 33.9   4   11
+## 12           Merc 280C 17.8   6    1
+## 13             Valiant 18.1   6    2
+## 14            Merc 280 19.2   6    3
+## 15        Ferrari Dino 19.7   6    4
+## 16           Mazda RX4 21.0   6    5
+## 17       Mazda RX4 Wag 21.0   6    5
+## 18      Hornet 4 Drive 21.4   6    7
+## 19  Cadillac Fleetwood 10.4   8    1
+## 20 Lincoln Continental 10.4   8    1
+

We explain in detail the above steps.

+
    +
  • windowPartitionBy creates a window specification object WindowSpec that defines the partition. It controls which rows will be in the same partition as the given row. In this case, rows with the same value in cyl will be put in the same partition. orderBy further defines the ordering - the position a given row is in the partition. The resulting WindowSpec is returned as ws.
  • +
+

More window specification methods include rangeBetween, which can define boundaries of the frame by value, and rowsBetween, which can define the boundaries by row indices.

+
    +
  • withColumn appends a Column called rank to the SparkDataFrame. over returns a windowing column. The first argument is usually a Column returned by window function(s) such as rank(), lead(carsDF$wt). That calculates the corresponding values according to the partitioned-and-ordered table.
  • +
+
+
+

User-Defined Function

+

In SparkR, we support several kinds of user-defined functions (UDFs).

+
+

Apply by Partition

+

dapply can apply a function to each partition of a SparkDataFrame. The function to be applied to each partition of the SparkDataFrame should have only one parameter, a data.frame corresponding to a partition, and the output should be a data.frame as well. Schema specifies the row format of the resulting a SparkDataFrame. It must match to data types of returned value. See here for mapping between R and Spark.

+

We convert mpg to kmpg (kilometers per gallon). carsSubDF is a SparkDataFrame with a subset of carsDF columns.

+ +
##               model  mpg   kmpg
+## 1         Mazda RX4 21.0 33.810
+## 2     Mazda RX4 Wag 21.0 33.810
+## 3        Datsun 710 22.8 36.708
+## 4    Hornet 4 Drive 21.4 34.454
+## 5 Hornet Sportabout 18.7 30.107
+## 6           Valiant 18.1 29.141
+

Like dapply, dapplyCollect can apply a function to each partition of a SparkDataFrame and collect the result back. The output of the function should be a data.frame, but no schema is required in this case. Note that dapplyCollect can fail if the output of the UDF on all partitions cannot be pulled into the driver’s memory.

+ +
##           model  mpg   kmpg
+## 1     Mazda RX4 21.0 33.810
+## 2 Mazda RX4 Wag 21.0 33.810
+## 3    Datsun 710 22.8 36.708
+
+
+

Apply by Group

+

gapply can apply a function to each group of a SparkDataFrame. The function is to be applied to each group of the SparkDataFrame and should have only two parameters: grouping key and R data.frame corresponding to that key. The groups are chosen from SparkDataFrames column(s). The output of function should be a data.frame. Schema specifies the row format of the resulting SparkDataFrame. It must represent R function’s output schema on the basis of Spark data types. The column names of the returned data.frame are set by user. See here for mapping between R and Spark.

+ +
##   cyl max_mpg
+## 1   4    33.9
+## 2   6    21.4
+## 3   8    19.2
+

Like gapply, gapplyCollect can apply a function to each partition of a SparkDataFrame and collect the result back to R data.frame. The output of the function should be a data.frame but no schema is required in this case. Note that gapplyCollect can fail if the output of the UDF on all partitions cannot be pulled into the driver’s memory.

+ +
##   cyl max_mpg
+## 1   4    33.9
+## 2   6    21.4
+## 3   8    19.2
+
+
+

Distribute Local Functions

+

Similar to lapply in native R, spark.lapply runs a function over a list of elements and distributes the computations with Spark. spark.lapply works in a manner that is similar to doParallel or lapply to elements of a list. The results of all the computations should fit in a single machine. If that is not the case you can do something like df <- createDataFrame(list) and then use dapply.

+

We use svm in package e1071 as an example. We use all default settings except for varying costs of constraints violation. spark.lapply can train those different models in parallel.

+ +

Return a list of model’s summaries.

+ + +
## [1] "list"
+

To avoid lengthy display, we only present the partial result of the second fitted model. You are free to inspect other models as well.

+ +
## $call
+## svm(formula = Species ~ ., data = iris, cost = cost)
+## 
+## $type
+## [1] 0
+## 
+## $kernel
+## [1] 2
+## 
+## $cost
+## [1] 5.623413
+## 
+## $degree
+## [1] 3
+## 
+## $gamma
+## [1] 0.25
+## 
+## $coef0
+## [1] 0
+## 
+## $nu
+## [1] 0.5
+## 
+## $epsilon
+## [1] 0.1
+## 
+## $sparse
+## [1] FALSE
+## 
+## $scaled
+## [1] TRUE TRUE TRUE TRUE
+## 
+## $x.scale
+## $x.scale$`scaled:center`
+## Sepal.Length  Sepal.Width Petal.Length  Petal.Width 
+##     5.843333     3.057333     3.758000     1.199333 
+## 
+## $x.scale$`scaled:scale`
+## Sepal.Length  Sepal.Width Petal.Length  Petal.Width 
+##    0.8280661    0.4358663    1.7652982    0.7622377 
+## 
+## 
+## $y.scale
+## NULL
+## 
+## $nclasses
+## [1] 3
+## 
+## $levels
+## [1] "setosa"     "versicolor" "virginica" 
+## 
+## $tot.nSV
+## [1] 35
+## 
+## $nSV
+## [1]  6 15 14
+## 
+## $labels
+## [1] 1 2 3
+## 
+## $SV
+##     Sepal.Length Sepal.Width Petal.Length Petal.Width
+## 14   -1.86378030 -0.13153881   -1.5056946  -1.4422448
+## 16   -0.17309407  3.08045544   -1.2791040  -1.0486668
+## 21   -0.53538397  0.78617383   -1.1658087  -1.3110521
+## 23   -1.50149039  1.24503015   -1.5623422  -1.3110521
+## 24   -0.89767388  0.55674567   -1.1658087  -0.9174741
+## 42   -1.62225369 -1.73753594   -1.3923993  -1.1798595
+## 51    1.39682886  0.32731751    0.5336209   0.2632600
+## 53    1.27606556  0.09788935    0.6469162   0.3944526
+## 54   -0.41462067 -1.73753594    0.1370873   0.1320673
+## 55    0.79301235 -0.59039513    0.4769732   0.3944526
+##  [ reached getOption("max.print") -- omitted 25 rows ]
+## 
+## $index
+##  [1]  14  16  21  23  24  42  51  53  54  55  58  61  69  71  73  78  79  84  85
+## [20]  86  99 107 111 119 120 124 127 128 130 132 134 135 139 149 150
+## 
+## $rho
+## [1] -0.10346530  0.12160294 -0.09540346
+## 
+## $compprob
+## [1] FALSE
+## 
+## $probA
+## NULL
+## 
+## $probB
+## NULL
+## 
+## $sigma
+## NULL
+## 
+## $coefs
+##              [,1]        [,2]
+##  [1,]  0.00000000  0.06561739
+##  [2,]  0.76813720  0.93378721
+##  [3,]  0.00000000  0.12123270
+##  [4,]  0.00000000  0.31170741
+##  [5,]  1.11614066  0.46397392
+##  [6,]  1.88141600  1.10392128
+##  [7,] -0.55872622  0.00000000
+##  [8,]  0.00000000  5.62341325
+##  [9,]  0.00000000  0.27711792
+## [10,]  0.00000000  5.28440007
+## [11,] -1.06596713  0.00000000
+## [12,] -0.57076709  1.09019756
+## [13,] -0.03365904  5.62341325
+## [14,]  0.00000000  5.62341325
+## [15,]  0.00000000  5.62341325
+## [16,]  0.00000000  5.62341325
+## [17,]  0.00000000  4.70398738
+## [18,]  0.00000000  5.62341325
+## [19,]  0.00000000  4.97981371
+## [20,] -0.77497987  0.00000000
+##  [ reached getOption("max.print") -- omitted 15 rows ]
+## 
+## $na.action
+## NULL
+## 
+## $fitted
+##      1      2      3      4      5      6      7      8      9     10     11 
+## setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa 
+##     12     13     14     15     16     17     18     19     20     21     22 
+## setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa 
+##     23     24     25     26     27     28     29     30     31     32     33 
+## setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa setosa 
+##     34     35     36     37     38     39     40 
+## setosa setosa setosa setosa setosa setosa setosa 
+##  [ reached getOption("max.print") -- omitted 110 entries ]
+## Levels: setosa versicolor virginica
+## 
+## $decision.values
+##     setosa/versicolor setosa/virginica versicolor/virginica
+## 1           1.1911739        1.0908424            1.1275805
+## 2           1.1336557        1.0619543            1.3260964
+## 3           1.2085065        1.0698101            1.0511345
+## 4           1.1646153        1.0505915            1.0806874
+## 5           1.1880814        1.0950348            0.9542815
+## 6           1.0990761        1.0984626            0.9326361
+## 7           1.1573474        1.0343287            0.9726843
+## 8           1.1851598        1.0815750            1.2206802
+## 9           1.1673499        1.0406734            0.8837945
+## 10          1.1629911        1.0560925            1.2430067
+## 11          1.1339282        1.0803946            1.0338357
+## 12          1.1724182        1.0641469            1.1190423
+## 13          1.1827355        1.0667956            1.1414844
+##  [ reached getOption("max.print") -- omitted 137 rows ]
+## 
+## $terms
+## Species ~ Sepal.Length + Sepal.Width + Petal.Length + Petal.Width
+## attr(,"variables")
+## list(Species, Sepal.Length, Sepal.Width, Petal.Length, Petal.Width)
+## attr(,"factors")
+##              Sepal.Length Sepal.Width Petal.Length Petal.Width
+## Species                 0           0            0           0
+## Sepal.Length            1           0            0           0
+## Sepal.Width             0           1            0           0
+## Petal.Length            0           0            1           0
+## Petal.Width             0           0            0           1
+## attr(,"term.labels")
+## [1] "Sepal.Length" "Sepal.Width"  "Petal.Length" "Petal.Width" 
+## attr(,"order")
+## [1] 1 1 1 1
+## attr(,"intercept")
+## [1] 0
+## attr(,"response")
+## [1] 1
+## attr(,".Environment")
+## <environment: 0x55f721df1dc0>
+## attr(,"predvars")
+## list(Species, Sepal.Length, Sepal.Width, Petal.Length, Petal.Width)
+## attr(,"dataClasses")
+##      Species Sepal.Length  Sepal.Width Petal.Length  Petal.Width 
+##     "factor"    "numeric"    "numeric"    "numeric"    "numeric" 
+## 
+## attr(,"class")
+## [1] "summary.svm"
+
+
+
+

SQL Queries

+

A SparkDataFrame can also be registered as a temporary view in Spark SQL so that one can run SQL queries over its data. The sql function enables applications to run SQL queries programmatically and returns the result as a SparkDataFrame.

+ +

Register this SparkDataFrame as a temporary view.

+ +

SQL statements can be run using the sql method.

+ +
##     name
+## 1 Justin
+
+
+
+

Machine Learning

+

SparkR supports the following machine learning models and algorithms.

+
+

Classification

+
    +
  • Linear Support Vector Machine (SVM) Classifier

  • +
  • Logistic Regression

  • +
  • Multilayer Perceptron (MLP)

  • +
  • Naive Bayes

  • +
  • Factorization Machines (FM) Classifier

  • +
+
+
+

Regression

+
    +
  • Accelerated Failure Time (AFT) Survival Model

  • +
  • Generalized Linear Model (GLM)

  • +
  • Isotonic Regression

  • +
  • Linear Regression

  • +
  • Factorization Machines (FM) Regressor

  • +
+
+
+

Tree - Classification and Regression

+
    +
  • Decision Tree

  • +
  • Gradient-Boosted Trees (GBT)

  • +
  • Random Forest

  • +
+
+
+

Clustering

+
    +
  • Bisecting \(k\)-means

  • +
  • Gaussian Mixture Model (GMM)

  • +
  • \(k\)-means Clustering

  • +
  • Latent Dirichlet Allocation (LDA)

  • +
  • Power Iteration Clustering (PIC)

  • +
+
+
+

Collaborative Filtering

+
    +
  • Alternating Least Squares (ALS)
  • +
+
+
+

Frequent Pattern Mining

+
    +
  • FP-growth
  • +
  • PrefixSpan
  • +
+
+
+

Statistics

+
    +
  • Kolmogorov-Smirnov Test
  • +
+
+
+

R Formula

+

For most above, SparkR supports R formula operators, including ~, ., :, + and - for model fitting. This makes it a similar experience as using R functions.

+
+
+

Training and Test Sets

+

We can easily split SparkDataFrame into random training and test sets by the randomSplit function. It returns a list of split SparkDataFrames with provided weights. We use carsDF as an example and want to have about \(70%\) training data and \(30%\) test data.

+ + +
## [1] 24
+ +
##                model  mpg cyl disp  hp drat    wt  qsec vs am gear carb
+## 1 Cadillac Fleetwood 10.4   8  472 205 2.93 5.250 17.98  0  0    3    4
+## 2         Camaro Z28 13.3   8  350 245 3.73 3.840 15.41  0  0    3    4
+## 3  Chrysler Imperial 14.7   8  440 230 3.23 5.345 17.42  0  0    3    4
+## 4   Dodge Challenger 15.5   8  318 150 2.76 3.520 16.87  0  0    3    2
+## 5         Duster 360 14.3   8  360 245 3.21 3.570 15.84  0  0    3    4
+## 6       Ferrari Dino 19.7   6  145 175 3.62 2.770 15.50  0  1    5    6
+ +
## [1] 8
+ +
##            model  mpg cyl  disp  hp drat    wt  qsec vs am gear carb
+## 1    AMC Javelin 15.2   8 304.0 150 3.15 3.435 17.30  0  0    3    2
+## 2     Datsun 710 22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1
+## 3       Fiat 128 32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1
+## 4      Merc 240D 24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2
+## 5       Merc 280 19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
+## 6 Toyota Corolla 33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1
+
+
+

Models and Algorithms

+
+

Linear Support Vector Machine (SVM) Classifier

+

Linear Support Vector Machine (SVM) classifier is an SVM classifier with linear kernels. This is a binary classifier. We use a simple example to show how to use spark.svmLinear for binary classification.

+ +
## $coefficients
+##                 Estimate
+## (Intercept)  0.993131388
+## Class_1st   -0.386500359
+## Class_2nd   -0.622627816
+## Class_3rd   -0.204446602
+## Sex_Female  -0.589950309
+## Age_Adult    0.741676902
+## Freq        -0.006582887
+## 
+## $numClasses
+## [1] 2
+## 
+## $numFeatures
+## [1] 6
+

Predict values on training data

+ +
##   Class    Sex   Age Freq Survived prediction
+## 1   1st   Male Child    0       No        Yes
+## 2   2nd   Male Child    0       No        Yes
+## 3   3rd   Male Child   35       No        Yes
+## 4  Crew   Male Child    0       No        Yes
+## 5   1st Female Child    0       No        Yes
+## 6   2nd Female Child    0       No         No
+
+
+

Logistic Regression

+

Logistic regression is a widely-used model when the response is categorical. It can be seen as a special case of the Generalized Linear Predictive Model. We provide spark.logit on top of spark.glm to support logistic regression with advanced hyper-parameters. It supports both binary and multiclass classification with elastic-net regularization and feature standardization, similar to glmnet.

+

We use a simple example to demonstrate spark.logit usage. In general, there are three steps of using spark.logit: 1). Create a dataframe from a proper data source; 2). Fit a logistic regression model using spark.logit with a proper parameter setting; and 3). Obtain the coefficient matrix of the fitted model using summary and use the model for prediction with predict.

+

Binomial logistic regression

+ +
## $coefficients
+##                  Estimate
+## (Intercept)  0.2255014282
+## Class_1st   -0.1338856652
+## Class_2nd   -0.1479826947
+## Class_3rd    0.0005674937
+## Sex_Female  -0.2011183871
+## Age_Adult    0.3263186885
+## Freq        -0.0033111157
+

Predict values on training data

+ +
##   Class    Sex   Age Freq Survived prediction
+## 1   1st   Male Child    0       No        Yes
+## 2   2nd   Male Child    0       No        Yes
+## 3   3rd   Male Child   35       No        Yes
+## 4  Crew   Male Child    0       No        Yes
+## 5   1st Female Child    0       No         No
+## 6   2nd Female Child    0       No         No
+

Multinomial logistic regression against three classes

+ +
## $coefficients
+##                      1st          2nd          3rd         Crew
+## (Intercept)  0.051662845  0.062998145 -0.039083689 -0.075577300
+## Sex_Female  -0.088030587 -0.102528148  0.059233106  0.131325629
+## Age_Adult    0.141935316  0.169492058 -0.102562719 -0.208864654
+## Survived_No  0.052721020  0.057980057 -0.029408423 -0.081292653
+## Freq        -0.001555912 -0.001970377  0.001303836  0.002222453
+
+
+

Multilayer Perceptron

+

Multilayer perceptron classifier (MLPC) is a classifier based on the feedforward artificial neural network. MLPC consists of multiple layers of nodes. Each layer is fully connected to the next layer in the network. Nodes in the input layer represent the input data. All other nodes map inputs to outputs by a linear combination of the inputs with the node’s weights \(w\) and bias \(b\) and applying an activation function. This can be written in matrix form for MLPC with \(K+1\) layers as follows: \[ +y(x)=f_K(\ldots f_2(w_2^T f_1(w_1^T x + b_1) + b_2) \ldots + b_K). +\]

+

Nodes in intermediate layers use sigmoid (logistic) function: \[ +f(z_i) = \frac{1}{1+e^{-z_i}}. +\]

+

Nodes in the output layer use softmax function: \[ +f(z_i) = \frac{e^{z_i}}{\sum_{k=1}^N e^{z_k}}. +\]

+

The number of nodes \(N\) in the output layer corresponds to the number of classes.

+

MLPC employs backpropagation for learning the model. We use the logistic loss function for optimization and L-BFGS as an optimization routine.

+

spark.mlp requires at least two columns in data: one named "label" and the other one "features". The "features" column should be in libSVM-format.

+

We use Titanic data set to show how to use spark.mlp in classification.

+ +

To avoid lengthy display, we only present partial results of the model summary. You can check the full result from your sparkR shell.

+ +
## $numOfInputs
+## [1] 2
+## 
+## $numOfOutputs
+## [1] 2
+## 
+## $layers
+## [1] 2 2
+## 
+## $weights
+## $weights[[1]]
+## [1] 0
+## 
+## $weights[[2]]
+## [1] 0
+## 
+## $weights[[3]]
+## [1] 5
+## 
+## $weights[[4]]
+## [1] 5
+## 
+## $weights[[5]]
+## [1] 9
+## 
+## $weights[[6]]
+## [1] 9
+ +
##   prediction
+## 1         No
+## 2         No
+## 3         No
+## 4         No
+## 5         No
+## 6         No
+
+
+

Naive Bayes

+

Naive Bayes model assumes independence among the features. spark.naiveBayes fits a Bernoulli naive Bayes model against a SparkDataFrame. The data should be all categorical. These models are often used for document classification.

+ +
## $apriori
+##            Yes        No
+## [1,] 0.5769231 0.4230769
+## 
+## $tables
+##     Class_3rd Class_1st Class_2nd Sex_Female Age_Adult
+## Yes 0.3125    0.3125    0.3125    0.5        0.5625   
+## No  0.4166667 0.25      0.25      0.5        0.75
+ +
##   Class    Sex   Age Survived prediction
+## 1   3rd   Male Child       No        Yes
+## 2   3rd Female Child       No        Yes
+## 3   1st   Male Adult       No        Yes
+## 4   2nd   Male Adult       No        Yes
+## 5   3rd   Male Adult       No         No
+## 6  Crew   Male Adult       No        Yes
+
+
+

Factorization Machines Classifier

+

Factorization Machines for classification problems.

+

For background and details about the implementation of factorization machines, refer to the Factorization Machines section.

+ +
## $coefficients
+##                 Estimate
+## (Intercept) 0.0064275991
+## Age_Adult   0.0001294448
+## Sex_Female  0.0001294448
+## 
+## $factors
+##            [,1]        [,2]      [,3]       [,4]        [,5]        [,6]
+## [1,] -0.3256224  0.11912568 0.1460235  0.1620567  0.13153516  0.06403695
+## [2,] -0.1382155 -0.03658261 0.1717808 -0.1602241 -0.08446129 -0.19287098
+##             [,7]        [,8]
+## [1,] -0.03292446 -0.05166818
+## [2,]  0.19252571  0.06237194
+## 
+## $numClasses
+## [1] 2
+## 
+## $numFeatures
+## [1] 2
+## 
+## $factorSize
+## [1] 8
+ +
##   prediction
+## 1        Yes
+## 2        Yes
+## 3        Yes
+## 4        Yes
+## 5        Yes
+## 6        Yes
+
+
+

Accelerated Failure Time Survival Model

+

Survival analysis studies the expected duration of time until an event happens, and often the relationship with risk factors or treatment taken on the subject. In contrast to standard regression analysis, survival modeling has to deal with special characteristics in the data including non-negative survival time and censoring.

+

Accelerated Failure Time (AFT) model is a parametric survival model for censored data that assumes the effect of a covariate is to accelerate or decelerate the life course of an event by some constant. For more information, refer to the Wikipedia page AFT Model and the references there. Different from a Proportional Hazards Model designed for the same purpose, the AFT model is easier to parallelize because each instance contributes to the objective function independently.

+ +
## $coefficients
+##                  Value
+## (Intercept)  6.8966910
+## ecog_ps     -0.3850414
+## rx           0.5286455
+## Log(scale)  -0.1234429
+ +
##   futime fustat     age resid_ds rx ecog_ps label prediction
+## 1     59      1 72.3315        2  1       1    59   1141.724
+## 2    115      1 74.4932        2  1       1   115   1141.724
+## 3    156      1 66.4658        2  1       2   156    776.855
+## 4    421      0 53.3644        2  2       1   421   1937.087
+## 5    431      1 50.3397        2  1       1   431   1141.724
+## 6    448      0 56.4301        1  1       2   448    776.855
+
+
+

Generalized Linear Model

+

The main function is spark.glm. The following families and link functions are supported. The default is gaussian.

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
FamilyLink Function
gaussianidentity, log, inverse
binomiallogit, probit, cloglog (complementary log-log)
poissonlog, identity, sqrt
gammainverse, identity, log
tweediepower link function
+

There are three ways to specify the family argument.

+
    +
  • Family name as a character string, e.g. family = "gaussian".

  • +
  • Family function, e.g. family = binomial.

  • +
  • Result returned by a family function, e.g. family = poisson(link = log).

  • +
  • Note that there are two ways to specify the tweedie family: +
      +
    1. Set family = "tweedie" and specify the var.power and link.power
    2. +
    3. When package statmod is loaded, the tweedie family is specified using the family definition therein, i.e., tweedie().
    4. +
  • +
+

For more information regarding the families and their link functions, see the Wikipedia page Generalized Linear Model.

+

We use the mtcars dataset as an illustration. The corresponding SparkDataFrame is carsDF. After fitting the model, we print out a summary and see the fitted values by making predictions on the original dataset. We can also pass into a new SparkDataFrame of same schema to predict on new data.

+ +
## 
+## Deviance Residuals: 
+## (Note: These are approximate quantiles with relative error <= 0.01)
+##     Min       1Q   Median       3Q      Max  
+## -3.9410  -1.6499  -0.3267   1.0373   5.8538  
+## 
+## Coefficients:
+##               Estimate  Std. Error  t value    Pr(>|t|)
+## (Intercept)  37.227270   1.5987875  23.2847  0.0000e+00
+## wt           -3.877831   0.6327335  -6.1287  1.1196e-06
+## hp           -0.031773   0.0090297  -3.5187  1.4512e-03
+## 
+## (Dispersion parameter for gaussian family taken to be 6.725785)
+## 
+##     Null deviance: 1126.05  on 31  degrees of freedom
+## Residual deviance:  195.05  on 29  degrees of freedom
+## AIC: 156.7
+## 
+## Number of Fisher Scoring iterations: 1
+

When doing prediction, a new column called prediction will be appended. Let’s look at only a subset of columns here.

+ +
##               model prediction  mpg    wt  hp
+## 1         Mazda RX4   23.57233 21.0 2.620 110
+## 2     Mazda RX4 Wag   22.58348 21.0 2.875 110
+## 3        Datsun 710   25.27582 22.8 2.320  93
+## 4    Hornet 4 Drive   21.26502 21.4 3.215 110
+## 5 Hornet Sportabout   18.32727 18.7 3.440 175
+## 6           Valiant   20.47382 18.1 3.460 105
+

The following is the same fit using the tweedie family:

+ +
## 
+## Deviance Residuals: 
+## (Note: These are approximate quantiles with relative error <= 0.01)
+##     Min       1Q   Median       3Q      Max  
+## -3.9410  -1.6499  -0.3267   1.0373   5.8538  
+## 
+## Coefficients:
+##               Estimate  Std. Error  t value    Pr(>|t|)
+## (Intercept)  37.227270   1.5987875  23.2847  0.0000e+00
+## wt           -3.877831   0.6327335  -6.1287  1.1196e-06
+## hp           -0.031773   0.0090297  -3.5187  1.4512e-03
+## 
+## (Dispersion parameter for tweedie family taken to be 6.725785)
+## 
+##     Null deviance: 1126.05  on 31  degrees of freedom
+## Residual deviance:  195.05  on 29  degrees of freedom
+## AIC: 156.7
+## 
+## Number of Fisher Scoring iterations: 1
+

We can try other distributions in the tweedie family, for example, a compound Poisson distribution with a log link:

+ +
## 
+## Deviance Residuals: 
+## (Note: These are approximate quantiles with relative error <= 0.01)
+##      Min        1Q    Median        3Q       Max  
+## -0.58074  -0.25335  -0.09892   0.18608   0.82717  
+## 
+## Coefficients:
+##                Estimate  Std. Error  t value    Pr(>|t|)
+## (Intercept)   3.8500849  0.06698272  57.4788  0.0000e+00
+## wt           -0.2018426  0.02897283  -6.9666  1.1691e-07
+## hp           -0.0016248  0.00041603  -3.9054  5.1697e-04
+## 
+## (Dispersion parameter for tweedie family taken to be 0.1340111)
+## 
+##     Null deviance: 29.8820  on 31  degrees of freedom
+## Residual deviance:  3.7739  on 29  degrees of freedom
+## AIC: NA
+## 
+## Number of Fisher Scoring iterations: 4
+
+
+

Isotonic Regression

+

spark.isoreg fits an Isotonic Regression model against a SparkDataFrame. It solves a weighted univariate a regression problem under a complete order constraint. Specifically, given a set of real observed responses \(y_1, \ldots, y_n\), corresponding real features \(x_1, \ldots, x_n\), and optionally positive weights \(w_1, \ldots, w_n\), we want to find a monotone (piecewise linear) function \(f\) to minimize \[ +\ell(f) = \sum_{i=1}^n w_i (y_i - f(x_i))^2. +\]

+

There are a few more arguments that may be useful.

+
    +
  • weightCol: a character string specifying the weight column.

  • +
  • isotonic: logical value indicating whether the output sequence should be isotonic/increasing (TRUE) or antitonic/decreasing (FALSE).

  • +
  • featureIndex: the index of the feature on the right hand side of the formula if it is a vector column (default: 0), no effect otherwise.

  • +
+

We use an artificial example to show the use.

+ +
##     x y prediction
+## 1 1.0 3        3.0
+## 2 2.0 6        5.5
+## 3 3.5 8        7.5
+## 4 3.0 5        5.5
+## 5 4.0 7        7.5
+

In the prediction stage, based on the fitted monotone piecewise function, the rules are:

+
    +
  • If the prediction input exactly matches a training feature then associated prediction is returned. In case there are multiple predictions with the same feature then one of them is returned. Which one is undefined.

  • +
  • If the prediction input is lower or higher than all training features then prediction with lowest or highest feature is returned respectively. In case there are multiple predictions with the same feature then the lowest or highest is returned respectively.

  • +
  • If the prediction input falls between two training features then prediction is treated as piecewise linear function and interpolated value is calculated from the predictions of the two closest features. In case there are multiple values with the same feature then the same rules as in previous point are used.

  • +
+

For example, when the input is \(3.2\), the two closest feature values are \(3.0\) and \(3.5\), then predicted value would be a linear interpolation between the predicted values at \(3.0\) and \(3.5\).

+ +
##     x prediction
+## 1 1.5       4.25
+## 2 3.2       6.30
+
+
+

Linear Regression

+

Linear regression model.

+ +
## $coefficients
+##                Estimate
+## (Intercept) 37.22727012
+## wt          -3.87783074
+## hp          -0.03177295
+## 
+## $numFeatures
+## [1] 2
+ +
##   prediction
+## 1   23.57233
+## 2   22.58348
+## 3   25.27582
+## 4   21.26502
+## 5   18.32727
+## 6   20.47382
+
+
+

Factorization Machines Regressor

+

Factorization Machines for regression problems.

+

For background and details about the implementation of factorization machines, refer to the Factorization Machines section.

+ +
## $coefficients
+##              Estimate
+## (Intercept) 0.1518559
+## wt          3.6472555
+## hp          2.8026828
+## 
+## $factors
+##            [,1]       [,2]       [,3]       [,4]      [,5]       [,6]
+## [1,]  0.1424420 -0.1178110 -0.3970272 -0.4696695  0.400288  0.3690930
+## [2,] -0.1626185  0.1512138  0.3690435  0.4076975 -0.625752 -0.3715109
+##             [,7]       [,8]
+## [1,]  0.03472468 -0.1703219
+## [2,] -0.02109148 -0.2006249
+## 
+## $numFeatures
+## [1] 2
+## 
+## $factorSize
+## [1] 8
+ +
##   prediction
+## 1  106.70996
+## 2   87.07526
+## 3  111.07931
+## 4   60.89565
+## 5   61.81374
+## 6   40.70095
+
+
+

Decision Tree

+

spark.decisionTree fits a decision tree classification or regression model on a SparkDataFrame. Users can call summary to get a summary of the fitted model, predict to make predictions, and write.ml/read.ml to save/load fitted models.

+

We use the Titanic dataset to train a decision tree and make predictions:

+ +
## Formula:  Survived ~ .
+## Number of features:  6
+## Features:  Class_1st Class_2nd Class_3rd Sex_Female Age_Adult Freq
+## Feature importances:  (6,[5],[1.0])
+## Max Depth:  2
+##  DecisionTreeClassificationModel: uid=dtc_1a036ac3b7f5, depth=2, numNodes=5, numClasses=2, numFeatures=6
+##   If (feature 5 <= 4.5)
+##    Predict: 0.0
+##   Else (feature 5 > 4.5)
+##    If (feature 5 <= 84.5)
+##     Predict: 1.0
+##    Else (feature 5 > 84.5)
+##     Predict: 0.0
+## 
+ +
##   Class    Sex   Age Freq Survived prediction
+## 1   1st   Male Child    0       No         No
+## 2   2nd   Male Child    0       No         No
+## 3   3rd   Male Child   35       No        Yes
+## 4  Crew   Male Child    0       No         No
+## 5   1st Female Child    0       No         No
+## 6   2nd Female Child    0       No         No
+
+
+

Gradient-Boosted Trees

+

spark.gbt fits a gradient-boosted tree classification or regression model on a SparkDataFrame. Users can call summary to get a summary of the fitted model, predict to make predictions, and write.ml/read.ml to save/load fitted models.

+

We use the Titanic dataset to train a gradient-boosted tree and make predictions:

+ +
## Formula:  Survived ~ .
+## Number of features:  6
+## Features:  Class_1st Class_2nd Class_3rd Sex_Female Age_Adult Freq
+## Feature importances:  (6,[1,2,5],[0.03336902858878361,0.16099525743106016,0.8056357139801562])
+## Max Depth:  2
+## Number of trees:  2
+## Tree weights:  1 0.1
+##  GBTClassificationModel: uid = gbtc_b9a7ebafe868, numTrees=2, numClasses=2, numFeatures=6
+##   Tree 0 (weight 1.0):
+##     If (feature 5 <= 4.5)
+##      If (feature 1 in {1.0})
+##       Predict: -1.0
+##      Else (feature 1 not in {1.0})
+##       Predict: -0.3333333333333333
+##     Else (feature 5 > 4.5)
+##      If (feature 5 <= 84.5)
+##       Predict: 0.5714285714285714
+##      Else (feature 5 > 84.5)
+##       Predict: -0.42857142857142855
+##   Tree 1 (weight 0.1):
+##     If (feature 2 in {1.0})
+##      If (feature 5 <= 15.5)
+##       Predict: 0.9671846896296403
+##      Else (feature 5 > 15.5)
+##       Predict: -1.0857923804083338
+##     Else (feature 2 not in {1.0})
+##      If (feature 5 <= 13.5)
+##       Predict: -0.08651035613926407
+##      Else (feature 5 > 13.5)
+##       Predict: 0.6566673506774614
+## 
+ +
##   Class    Sex   Age Freq Survived prediction
+## 1   1st   Male Child    0       No         No
+## 2   2nd   Male Child    0       No         No
+## 3   3rd   Male Child   35       No        Yes
+## 4  Crew   Male Child    0       No         No
+## 5   1st Female Child    0       No         No
+## 6   2nd Female Child    0       No         No
+
+
+

Random Forest

+

spark.randomForest fits a random forest classification or regression model on a SparkDataFrame. Users can call summary to get a summary of the fitted model, predict to make predictions, and write.ml/read.ml to save/load fitted models.

+

In the following example, we use the Titanic dataset to train a random forest and make predictions:

+ +
## Formula:  Survived ~ .
+## Number of features:  6
+## Features:  Class_1st Class_2nd Class_3rd Sex_Female Age_Adult Freq
+## Feature importances:  (6,[3,4,5],[0.17058779274099098,0.09676977311565654,0.7326424341433525])
+## Max Depth:  2
+## Number of trees:  2
+## Tree weights:  1 1
+##  RandomForestClassificationModel: uid=rfc_a4079206d906, numTrees=2, numClasses=2, numFeatures=6
+##   Tree 0 (weight 1.0):
+##     If (feature 4 in {0.0})
+##      If (feature 3 in {0.0})
+##       Predict: 0.0
+##      Else (feature 3 not in {0.0})
+##       Predict: 1.0
+##     Else (feature 4 not in {0.0})
+##      If (feature 5 <= 13.5)
+##       Predict: 0.0
+##      Else (feature 5 > 13.5)
+##       Predict: 1.0
+##   Tree 1 (weight 1.0):
+##     If (feature 5 <= 84.5)
+##      If (feature 5 <= 4.5)
+##       Predict: 0.0
+##      Else (feature 5 > 4.5)
+##       Predict: 1.0
+##     Else (feature 5 > 84.5)
+##      Predict: 0.0
+## 
+ +
##   Class    Sex   Age Freq Survived prediction
+## 1   1st   Male Child    0       No         No
+## 2   2nd   Male Child    0       No         No
+## 3   3rd   Male Child   35       No        Yes
+## 4  Crew   Male Child    0       No         No
+## 5   1st Female Child    0       No         No
+## 6   2nd Female Child    0       No         No
+
+
+

Bisecting k-Means

+

spark.bisectingKmeans is a kind of hierarchical clustering using a divisive (or “top-down”) approach: all observations start in one cluster, and splits are performed recursively as one moves down the hierarchy.

+ +
## $k
+## [1] 4
+## 
+## $coefficients
+##   Survived_No
+## 1 0          
+## 2 1          
+## 3 0          
+## 4 1          
+## 
+## $size
+## $size[[1]]
+## [1] 16
+## 
+## $size[[2]]
+## [1] 16
+## 
+## $size[[3]]
+## [1] 0
+## 
+## $size[[4]]
+## [1] 0
+## 
+## 
+## $cluster
+## SparkDataFrame[prediction:int]
+## 
+## $is.loaded
+## [1] FALSE
+ +
##   Class prediction
+## 1   1st          1
+## 2   2nd          1
+## 3   3rd          1
+## 4  Crew          1
+## 5   1st          1
+## 6   2nd          1
+
+
+

Gaussian Mixture Model

+

spark.gaussianMixture fits multivariate Gaussian Mixture Model (GMM) against a SparkDataFrame. Expectation-Maximization (EM) is used to approximate the maximum likelihood estimator (MLE) of the model.

+

We use a simulated example to demonstrate the usage.

+ +
## $lambda
+## [1] 0.486168 0.513832
+## 
+## $mu
+## $mu[[1]]
+## [1] 1.424505 1.995314
+## 
+## $mu[[2]]
+## [1] 1.922353 2.469800
+## 
+## 
+## $sigma
+## $sigma[[1]]
+##      [,1]     [,2]    
+## [1,] 3.130571 2.557531
+## [2,] 2.557531 3.903907
+## 
+## $sigma[[2]]
+##      [,1]     [,2]    
+## [1,] 2.952933 2.382424
+## [2,] 2.382424 3.763825
+## 
+## 
+## $loglik
+## [1] -37.05907
+## 
+## $posterior
+## SparkDataFrame[posterior:array<double>]
+## 
+## $is.loaded
+## [1] FALSE
+ +
##            V1          V2 prediction
+## 1 -1.04010580  0.44731013          0
+## 2  1.03343749 -0.80795279          0
+## 3  0.07209263  0.75728370          0
+## 4 -0.81326509  0.03299667          0
+## 5  3.06240539  5.25819335          1
+## 6  3.23457954  2.18585177          1
+
+
+

k-Means Clustering

+

spark.kmeans fits a \(k\)-means clustering model against a SparkDataFrame. As an unsupervised learning method, we don’t need a response variable. Hence, the left hand side of the R formula should be left blank. The clustering is based only on the variables on the right hand side.

+ +
## $k
+## [1] 3
+## 
+## $coefficients
+##        mpg        hp       wt
+## 1 24.22353  93.52941 2.599588
+## 2 15.80000 178.50000 3.926400
+## 3 14.62000 263.80000 3.899000
+## 
+## $size
+## $size[[1]]
+## [1] 17
+## 
+## $size[[2]]
+## [1] 10
+## 
+## $size[[3]]
+## [1] 5
+## 
+## 
+## $cluster
+## SparkDataFrame[prediction:int]
+## 
+## $is.loaded
+## [1] FALSE
+## 
+## $clusterSize
+## [1] 3
+ +
##                  model  mpg  hp    wt prediction
+## 1            Mazda RX4 21.0 110 2.620          0
+## 2        Mazda RX4 Wag 21.0 110 2.875          0
+## 3           Datsun 710 22.8  93 2.320          0
+## 4       Hornet 4 Drive 21.4 110 3.215          0
+## 5    Hornet Sportabout 18.7 175 3.440          1
+## 6              Valiant 18.1 105 3.460          0
+## 7           Duster 360 14.3 245 3.570          2
+## 8            Merc 240D 24.4  62 3.190          0
+## 9             Merc 230 22.8  95 3.150          0
+## 10            Merc 280 19.2 123 3.440          0
+## 11           Merc 280C 17.8 123 3.440          0
+## 12          Merc 450SE 16.4 180 4.070          1
+## 13          Merc 450SL 17.3 180 3.730          1
+## 14         Merc 450SLC 15.2 180 3.780          1
+## 15  Cadillac Fleetwood 10.4 205 5.250          1
+## 16 Lincoln Continental 10.4 215 5.424          1
+## 17   Chrysler Imperial 14.7 230 5.345          2
+## 18            Fiat 128 32.4  66 2.200          0
+## 19         Honda Civic 30.4  52 1.615          0
+## 20      Toyota Corolla 33.9  65 1.835          0
+
+
+

Latent Dirichlet Allocation

+

spark.lda fits a Latent Dirichlet Allocation model on a SparkDataFrame. It is often used in topic modeling in which topics are inferred from a collection of text documents. LDA can be thought of as a clustering algorithm as follows:

+
    +
  • Topics correspond to cluster centers, and documents correspond to examples (rows) in a dataset.

  • +
  • Topics and documents both exist in a feature space, where feature vectors are vectors of word counts (bag of words).

  • +
  • Rather than clustering using a traditional distance, LDA uses a function based on a statistical model of how text documents are generated.

  • +
+

To use LDA, we need to specify a features column in data where each entry represents a document. There are two options for the column:

+
    +
  • character string: This can be a string of the whole document. It will be parsed automatically. Additional stop words can be added in customizedStopWords.

  • +
  • libSVM: Each entry is a collection of words and will be processed directly.

  • +
+

Two more functions are provided for the fitted model.

+
    +
  • spark.posterior returns a SparkDataFrame containing a column of posterior probabilities vectors named “topicDistribution”.

  • +
  • spark.perplexity returns the log perplexity of given SparkDataFrame, or the log perplexity of the training data if missing argument data.

  • +
+

For more information, see the help document ?spark.lda.

+

Let’s look an artificial example.

+ +
## $docConcentration
+## [1] 11 11 11 11 11
+## 
+## $topicConcentration
+## [1] 1.1
+## 
+## $logLikelihood
+## [1] -353.2948
+## 
+## $logPerplexity
+## [1] 2.676476
+## 
+## $isDistributed
+## [1] TRUE
+## 
+## $vocabSize
+## [1] 10
+## 
+## $topics
+## SparkDataFrame[topic:int, term:array<string>, termWeights:array<double>]
+## 
+## $vocabulary
+##  [1] "0" "1" "2" "3" "4" "9" "5" "8" "7" "6"
+## 
+## $trainingLogLikelihood
+## [1] -239.5629
+## 
+## $logPrior
+## [1] -980.2974
+ +
##                features                                     topicDistribution
+## 1 1 2 6 0 2 3 1 1 0 0 3 0.1972168, 0.1986649, 0.2022016, 0.2006567, 0.2012600
+## 2 1 3 0 1 3 0 0 2 0 0 1 0.1989973, 0.1988744, 0.2015969, 0.2006370, 0.1998944
+## 3 1 4 1 0 0 4 9 0 1 2 0 0.2020597, 0.2026093, 0.1968839, 0.1987286, 0.1997187
+## 4 2 1 0 3 0 0 5 0 2 3 9 0.2004066, 0.1981940, 0.2013016, 0.2006305, 0.1994673
+## 5 3 1 1 9 3 0 2 0 0 1 3 0.1971473, 0.1983960, 0.2023580, 0.2011575, 0.2009413
+## 6 4 2 0 3 4 5 1 1 1 4 0 0.2020226, 0.2041832, 0.1955394, 0.1997222, 0.1985325
+ +
## [1] 2.676476
+
+ +
+

Power Iteration Clustering

+

Power Iteration Clustering (PIC) is a scalable graph clustering algorithm. spark.assignClusters method runs the PIC algorithm and returns a cluster assignment for each input vertex.

+ +
##   id cluster
+## 1  4       1
+## 2  0       0
+## 3  1       0
+## 4  3       1
+## 5  2       0
+
+
+

FP-growth

+

spark.fpGrowth executes FP-growth algorithm to mine frequent itemsets on a SparkDataFrame. itemsCol should be an array of values.

+ +

spark.freqItemsets method can be used to retrieve a SparkDataFrame with the frequent itemsets.

+ +
##   items freq
+## 1     R    9
+## 2     U    8
+## 3  U, T    4
+## 4  U, V    4
+## 5  U, S    4
+## 6     T   10
+

spark.associationRules returns a SparkDataFrame with the association rules.

+ +
##   antecedent consequent confidence     lift support
+## 1          V          R  0.5555556 1.234568    0.25
+## 2          S          T  0.5454545 1.090909    0.30
+## 3          T          S  0.6000000 1.090909    0.30
+## 4          R          V  0.5555556 1.234568    0.25
+## 5          U          T  0.5000000 1.000000    0.20
+## 6          U          V  0.5000000 1.111111    0.20
+

We can make predictions based on the antecedent.

+ +
##        items prediction
+## 1    T, R, U       S, V
+## 2       T, S       NULL
+## 3       V, R       NULL
+## 4 R, U, T, V          S
+## 5       R, S       T, V
+## 6    V, S, U       R, T
+
+
+

PrefixSpan

+

spark.findFrequentSequentialPatterns method can be used to find the complete set of frequent sequential patterns in the input sequences of itemsets.

+ +
##   sequence freq
+## 1        1    3
+## 2        3    2
+## 3        2    3
+## 4     1, 2    3
+## 5     1, 3    2
+
+
+

Kolmogorov-Smirnov Test

+

spark.kstest runs a two-sided, one-sample Kolmogorov-Smirnov (KS) test. Given a SparkDataFrame, the test compares continuous data in a given column testCol with the theoretical distribution specified by parameter nullHypothesis. Users can call summary to get a summary of the test results.

+

In the following example, we test whether the Titanic dataset’s Freq column follows a normal distribution. We set the parameters of the normal distribution using the mean and standard deviation of the sample.

+ +
## Kolmogorov-Smirnov test summary:
+## degrees of freedom = 0 
+## statistic = 0.3065126710255011 
+## pValue = 0.0036336792155329256 
+## Very strong presumption against null hypothesis: Sample follows theoretical distribution.
+
+
+
+

Model Persistence

+

The following example shows how to save/load an ML model in SparkR.

+ +
## 
+## Saved-loaded model does not support output 'Deviance Residuals'.
+## 
+## Coefficients:
+##              Estimate  Std. Error  t value   Pr(>|t|)
+## (Intercept)    46.219      35.994   1.2841  0.2092846
+## Sex_Female    -78.812      41.562  -1.8962  0.0679311
+## Age_Adult     123.938      41.562   2.9820  0.0057522
+## 
+## (Dispersion parameter for gaussian family taken to be 13819.52)
+## 
+##     Null deviance: 573341  on 31  degrees of freedom
+## Residual deviance: 400766  on 29  degrees of freedom
+## AIC: 400.7
+## 
+## Number of Fisher Scoring iterations: 1
+ +
##   Class    Sex   Age Survived Freq label prediction
+## 1   1st   Male Child       No    0     0   46.21875
+## 2   2nd   Male Child       No    0     0   46.21875
+## 3   3rd   Male Child       No   35    35   46.21875
+## 4  Crew   Male Child       No    0     0   46.21875
+## 5   1st Female Child       No    0     0  -32.59375
+## 6   2nd Female Child       No    0     0  -32.59375
+ +
+
+
+

Structured Streaming

+

SparkR supports the Structured Streaming API.

+

You can check the Structured Streaming Programming Guide for an introduction to its programming model and basic concepts.

+ + +
+

Operations and Sinks

+

Most of the common operations on SparkDataFrame are supported for streaming, including selection, projection, and aggregation. Once you have defined the final result, to start the streaming computation, you will call the write.stream method setting a sink and outputMode.

+

A streaming SparkDataFrame can be written for debugging to the console, to a temporary in-memory table, or for further processing in a fault-tolerant manner to a File Sink in different formats.

+ +
+
+
+

Advanced Topics

+
+

SparkR Object Classes

+

There are three main object classes in SparkR you may be working with.

+
    +
  • SparkDataFrame: the central component of SparkR. It is an S4 class representing distributed collection of data organized into named columns, which is conceptually equivalent to a table in a relational database or a data frame in R. It has two slots sdf and env. +
      +
    • sdf stores a reference to the corresponding Spark Dataset in the Spark JVM backend.
    • +
    • env saves the meta-information of the object such as isCached.
    • +
    +

    It can be created by data import methods or by transforming an existing SparkDataFrame. We can manipulate SparkDataFrame by numerous data processing functions and feed that into machine learning algorithms.

  • +
  • Column: an S4 class representing a column of SparkDataFrame. The slot jc saves a reference to the corresponding Column object in the Spark JVM backend.

    +

    It can be obtained from a SparkDataFrame by $ operator, e.g., df$col. More often, it is used together with other functions, for example, with select to select particular columns, with filter and constructed conditions to select rows, with aggregation functions to compute aggregate statistics for each group.

  • +
  • GroupedData: an S4 class representing grouped data created by groupBy or by transforming other GroupedData. Its sgd slot saves a reference to a RelationalGroupedDataset object in the backend.

    +

    This is often an intermediate object with group information and followed up by aggregation operations.

  • +
+
+
+

Architecture

+

A complete description of architecture can be seen in the references, in particular the paper SparkR: Scaling R Programs with Spark.

+

Under the hood of SparkR is Spark SQL engine. This avoids the overheads of running interpreted R code, and the optimized SQL execution engine in Spark uses structural information about data and computation flow to perform a bunch of optimizations to speed up the computation.

+

The main method calls of actual computation happen in the Spark JVM of the driver. We have a socket-based SparkR API that allows us to invoke functions on the JVM from R. We use a SparkR JVM backend that listens on a Netty-based socket server.

+

Two kinds of RPCs are supported in the SparkR JVM backend: method invocation and creating new objects. Method invocation can be done in two ways.

+
    +
  • sparkR.callJMethod takes a reference to an existing Java object and a list of arguments to be passed on to the method.

  • +
  • sparkR.callJStatic takes a class name for static method and a list of arguments to be passed on to the method.

  • +
+

The arguments are serialized using our custom wire format which is then deserialized on the JVM side. We then use Java reflection to invoke the appropriate method.

+

To create objects, sparkR.newJObject is used and then similarly the appropriate constructor is invoked with provided arguments.

+

Finally, we use a new R class jobj that refers to a Java object existing in the backend. These references are tracked on the Java side and are automatically garbage collected when they go out of scope on the R side.

+
+
+
+

Appendix

+
+

R and Spark Data Types

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
RSpark
bytebyte
integerinteger
floatfloat
doubledouble
numericdouble
characterstring
stringstring
binarybinary
rawbinary
logicalboolean
POSIXcttimestamp
POSIXlttimestamp
Datedate
arrayarray
listarray
envmap
+
+
+
+

References

+ +
+ + + + + + + + + + + diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/AnIndex b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/AnIndex new file mode 100644 index 0000000..faeb2f6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/AnIndex @@ -0,0 +1,1099 @@ +!,Column-method not +$ select +$,SparkDataFrame-method select +$<- select +$<-,SparkDataFrame-method select +%<=>% eq_null_safe +%<=>%,Column-method eq_null_safe +%in% match +%in%,Column-method match +abs column_math_functions +abs,Column-method column_math_functions +acos column_math_functions +acos,Column-method column_math_functions +acosh column_math_functions +acosh,Column-method column_math_functions +addJobTag addJobTAg +add_months column_datetime_diff_functions +add_months,Column,numeric-method column_datetime_diff_functions +AFTSurvivalRegressionModel-class AFTSurvivalRegressionModel-class +agg summarize +agg,GroupedData-method summarize +agg,SparkDataFrame-method summarize +alias alias +alias,Column-method alias +alias,SparkDataFrame-method alias +ALSModel-class ALSModel-class +approxCountDistinct column_aggregate_functions +approxCountDistinct,Column-method column_aggregate_functions +approxQuantile approxQuantile +approxQuantile,SparkDataFrame,character,numeric,numeric-method approxQuantile +approx_count_distinct column_aggregate_functions +approx_count_distinct,Column-method column_aggregate_functions +arrange arrange +arrange,SparkDataFrame,character-method arrange +arrange,SparkDataFrame,Column-method arrange +arrays_overlap column_collection_functions +arrays_overlap,Column,Column-method column_collection_functions +arrays_overlap,Column-method column_collection_functions +arrays_zip column_collection_functions +arrays_zip,Column-method column_collection_functions +arrays_zip_with column_collection_functions +arrays_zip_with,characterOrColumn,characterOrColumn,function-method column_collection_functions +array_aggregate column_collection_functions +array_aggregate,characterOrColumn,Column,function-method column_collection_functions +array_contains column_collection_functions +array_contains,Column-method column_collection_functions +array_distinct column_collection_functions +array_distinct,Column-method column_collection_functions +array_except column_collection_functions +array_except,Column,Column-method column_collection_functions +array_except,Column-method column_collection_functions +array_exists column_collection_functions +array_exists,characterOrColumn,function-method column_collection_functions +array_filter column_collection_functions +array_filter,characterOrColumn,function-method column_collection_functions +array_forall column_collection_functions +array_forall,characterOrColumn,function-method column_collection_functions +array_intersect column_collection_functions +array_intersect,Column,Column-method column_collection_functions +array_intersect,Column-method column_collection_functions +array_join column_collection_functions +array_join,Column,character-method column_collection_functions +array_join,Column-method column_collection_functions +array_max column_collection_functions +array_max,Column-method column_collection_functions +array_min column_collection_functions +array_min,Column-method column_collection_functions +array_position column_collection_functions +array_position,Column-method column_collection_functions +array_remove column_collection_functions +array_remove,Column-method column_collection_functions +array_repeat column_collection_functions +array_repeat,Column,numericOrColumn-method column_collection_functions +array_sort column_collection_functions +array_sort,Column-method column_collection_functions +array_to_vector column_ml_functions +array_to_vector,Column-method column_ml_functions +array_transform column_collection_functions +array_transform,characterOrColumn,characterOrColumn,function-method column_collection_functions +array_transform,characterOrColumn,function-method column_collection_functions +array_union column_collection_functions +array_union,Column,Column-method column_collection_functions +array_union,Column-method column_collection_functions +as.data.frame as.data.frame +as.data.frame,SparkDataFrame-method as.data.frame +as.DataFrame createDataFrame +asc columnfunctions +ascii column_string_functions +ascii,Column-method column_string_functions +asc_nulls_first columnfunctions +asc_nulls_last columnfunctions +asin column_math_functions +asin,Column-method column_math_functions +asinh column_math_functions +asinh,Column-method column_math_functions +assert_true column_misc_functions +assert_true,Column-method column_misc_functions +associationRules,FPGrowthModel-method spark.fpGrowth +atan column_math_functions +atan,Column-method column_math_functions +atan2 column_math_functions +atan2,Column-method column_math_functions +atanh column_math_functions +atanh,Column-method column_math_functions +attach attach +attach,SparkDataFrame-method attach +avg avg +avg,Column-method avg +awaitTermination awaitTermination +awaitTermination,StreamingQuery-method awaitTermination +base64 column_string_functions +base64,Column-method column_string_functions +between between +between,Column-method between +bin column_math_functions +bin,Column-method column_math_functions +BisectingKMeansModel-class BisectingKMeansModel-class +bitwiseNOT column_nonaggregate_functions +bitwiseNOT,Column-method column_nonaggregate_functions +bitwise_not column_nonaggregate_functions +bitwise_not,Column-method column_nonaggregate_functions +bit_length column_string_functions +bit_length,Column-method column_string_functions +broadcast broadcast +broadcast,SparkDataFrame-method broadcast +bround column_math_functions +bround,Column-method column_math_functions +cache cache +cache,SparkDataFrame-method cache +cacheTable cacheTable +cancelJobGroup cancelJobGroup +cancelJobsWithTag cancelJobsWithTag +cast cast +cast,Column-method cast +cbrt column_math_functions +cbrt,Column-method column_math_functions +ceil column_math_functions +ceil,Column-method column_math_functions +ceiling column_math_functions +ceiling,Column-method column_math_functions +checkpoint checkpoint +checkpoint,SparkDataFrame-method checkpoint +clearCache clearCache +clearJobGroup clearJobGroup +clearJobTags clearJobTags +coalesce coalesce +coalesce,Column-method column_nonaggregate_functions +coalesce,SparkDataFrame-method coalesce +collect collect +collect,SparkDataFrame-method collect +collect_list column_aggregate_functions +collect_list,Column-method column_aggregate_functions +collect_set column_aggregate_functions +collect_set,Column-method column_aggregate_functions +colnames columns +colnames,SparkDataFrame-method columns +colnames<- columns +colnames<-,SparkDataFrame-method columns +coltypes coltypes +coltypes,SparkDataFrame-method coltypes +coltypes<- coltypes +coltypes<-,SparkDataFrame,character-method coltypes +column column +column,character-method column +column,jobj-method column +Column-class column +columnfunctions columnfunctions +columns columns +columns,SparkDataFrame-method columns +column_aggregate_functions column_aggregate_functions +column_avro_functions column_avro_functions +column_collection_functions column_collection_functions +column_datetime_diff_functions column_datetime_diff_functions +column_datetime_functions column_datetime_functions +column_math_functions column_math_functions +column_misc_functions column_misc_functions +column_ml_functions column_ml_functions +column_nonaggregate_functions column_nonaggregate_functions +column_string_functions column_string_functions +column_window_functions column_window_functions +concat column_collection_functions +concat,Column-method column_collection_functions +concat_ws column_string_functions +concat_ws,character,Column-method column_string_functions +contains columnfunctions +conv column_math_functions +conv,Column,numeric,numeric-method column_math_functions +corr corr +corr,Column-method corr +corr,SparkDataFrame-method corr +cos column_math_functions +cos,Column-method column_math_functions +cosh column_math_functions +cosh,Column-method column_math_functions +cot column_math_functions +cot,Column-method column_math_functions +count count +count,Column-method count +count,GroupedData-method count +count,SparkDataFrame-method nrow +countDistinct column_aggregate_functions +countDistinct,Column-method column_aggregate_functions +count_distinct column_aggregate_functions +count_distinct,Column-method column_aggregate_functions +cov cov +cov,characterOrColumn-method cov +cov,SparkDataFrame-method cov +covar_pop cov +covar_pop,characterOrColumn,characterOrColumn-method cov +covar_samp cov +covar_samp,characterOrColumn,characterOrColumn-method cov +crc32 column_misc_functions +crc32,Column-method column_misc_functions +createDataFrame createDataFrame +createExternalTable createExternalTable-deprecated +createOrReplaceTempView createOrReplaceTempView +createOrReplaceTempView,SparkDataFrame,character-method createOrReplaceTempView +createTable createTable +create_array column_nonaggregate_functions +create_array,Column-method column_nonaggregate_functions +create_lambda create_lambda +create_map column_nonaggregate_functions +create_map,Column-method column_nonaggregate_functions +crossJoin crossJoin +crossJoin,SparkDataFrame,SparkDataFrame-method crossJoin +crosstab crosstab +crosstab,SparkDataFrame,character,character-method crosstab +csc column_math_functions +csc,Column-method column_math_functions +cube cube +cube,SparkDataFrame-method cube +cume_dist column_window_functions +cume_dist,missing-method column_window_functions +currentCatalog currentCatalog +currentDatabase currentDatabase +current_date column_datetime_functions +current_date,missing-method column_datetime_functions +current_timestamp column_datetime_functions +current_timestamp,missing-method column_datetime_functions +dapply dapply +dapply,SparkDataFrame,function,characterOrstructType-method dapply +dapplyCollect dapplyCollect +dapplyCollect,SparkDataFrame,function-method dapplyCollect +databaseExists databaseExists +datediff column_datetime_diff_functions +datediff,Column-method column_datetime_diff_functions +date_add column_datetime_diff_functions +date_add,Column,numeric-method column_datetime_diff_functions +date_format column_datetime_diff_functions +date_format,Column,character-method column_datetime_diff_functions +date_sub column_datetime_diff_functions +date_sub,Column,numeric-method column_datetime_diff_functions +date_trunc column_datetime_functions +date_trunc,character,Column-method column_datetime_functions +dayofmonth column_datetime_functions +dayofmonth,Column-method column_datetime_functions +dayofweek column_datetime_functions +dayofweek,Column-method column_datetime_functions +dayofyear column_datetime_functions +dayofyear,Column-method column_datetime_functions +DecisionTreeClassificationModel-class DecisionTreeClassificationModel-class +DecisionTreeRegressionModel-class DecisionTreeRegressionModel-class +decode column_string_functions +decode,Column,character-method column_string_functions +degrees column_math_functions +degrees,Column-method column_math_functions +dense_rank column_window_functions +dense_rank,missing-method column_window_functions +desc columnfunctions +describe describe +describe,SparkDataFrame,ANY-method describe +describe,SparkDataFrame,character-method describe +describe,SparkDataFrame-method describe +desc_nulls_first columnfunctions +desc_nulls_last columnfunctions +dim dim +dim,SparkDataFrame-method dim +distinct distinct +distinct,SparkDataFrame-method distinct +drop drop +drop,ANY,ANY-method drop +drop,SparkDataFrame,characterOrColumn-method drop +dropDuplicates dropDuplicates +dropDuplicates,SparkDataFrame-method dropDuplicates +dropFields dropFields +dropFields,Column-method dropFields +dropna nafunctions +dropna,SparkDataFrame-method nafunctions +dropTempTable dropTempTable-deprecated +dropTempView dropTempView +dtypes dtypes +dtypes,SparkDataFrame-method dtypes +element_at column_collection_functions +element_at,Column-method column_collection_functions +encode column_string_functions +encode,Column,character-method column_string_functions +endsWith endsWith +endsWith,Column-method endsWith +except except +except,SparkDataFrame,SparkDataFrame-method except +exceptAll exceptAll +exceptAll,SparkDataFrame,SparkDataFrame-method exceptAll +exp column_math_functions +exp,Column-method column_math_functions +explain explain +explain,SparkDataFrame-method explain +explain,StreamingQuery-method explain +explode column_collection_functions +explode,Column-method column_collection_functions +explode_outer column_collection_functions +explode_outer,Column-method column_collection_functions +expm1 column_math_functions +expm1,Column-method column_math_functions +expr column_nonaggregate_functions +expr,character-method column_nonaggregate_functions +factorial column_math_functions +factorial,Column-method column_math_functions +fillna nafunctions +fillna,SparkDataFrame-method nafunctions +filter filter +filter,SparkDataFrame,characterOrColumn-method filter +findFrequentSequentialPatterns,PrefixSpan,SparkDataFrame-method spark.prefixSpan +first first +first,characterOrColumn-method first +first,SparkDataFrame-method first +fitted fitted +fitted,BisectingKMeansModel-method spark.bisectingKmeans +fitted,KMeansModel-method fitted +flatten column_collection_functions +flatten,Column-method column_collection_functions +floor column_math_functions +floor,Column-method column_math_functions +FMClassificationModel-class FMClassificationModel-class +FMRegressionModel-class FMRegressionModel-class +format_number column_string_functions +format_number,Column,numeric-method column_string_functions +format_string column_string_functions +format_string,character,Column-method column_string_functions +FPGrowthModel-class FPGrowthModel-class +freqItems freqItems +freqItems,SparkDataFrame,character-method freqItems +freqItemsets,FPGrowthModel-method spark.fpGrowth +from_avro column_avro_functions +from_avro,characterOrColumn-method column_avro_functions +from_avro,Column-method column_avro_functions +from_csv column_collection_functions +from_csv,Column,characterOrstructTypeOrColumn-method column_collection_functions +from_json column_collection_functions +from_json,Column,characterOrstructTypeOrColumn-method column_collection_functions +from_unixtime column_datetime_functions +from_unixtime,Column-method column_datetime_functions +from_utc_timestamp column_datetime_diff_functions +from_utc_timestamp,Column,character-method column_datetime_diff_functions +functionExists functionExists +gapply gapply +gapply,GroupedData-method gapply +gapply,SparkDataFrame-method gapply +gapplyCollect gapplyCollect +gapplyCollect,GroupedData-method gapplyCollect +gapplyCollect,SparkDataFrame-method gapplyCollect +GaussianMixtureModel-class GaussianMixtureModel-class +GBTClassificationModel-class GBTClassificationModel-class +GBTRegressionModel-class GBTRegressionModel-class +GeneralizedLinearRegressionModel-class GeneralizedLinearRegressionModel-class +getDatabase getDatabase +getField columnfunctions +getFunc getFunc +getItem columnfunctions +getJobTags getJobTags +getLocalProperty getLocalProperty +getNumPartitions getNumPartitions +getNumPartitions,SparkDataFrame-method getNumPartitions +getTable getTable +glm glm +glm,formula,ANY,SparkDataFrame-method glm +greatest column_nonaggregate_functions +greatest,Column-method column_nonaggregate_functions +groupBy groupBy +groupBy,SparkDataFrame-method groupBy +groupedData GroupedData +GroupedData-class GroupedData +grouping_bit column_aggregate_functions +grouping_bit,Column-method column_aggregate_functions +grouping_id column_aggregate_functions +grouping_id,Column-method column_aggregate_functions +group_by groupBy +group_by,SparkDataFrame-method groupBy +hash column_misc_functions +hash,Column-method column_misc_functions +hashCode hashCode +head head +head,SparkDataFrame-method head +hex column_math_functions +hex,Column-method column_math_functions +hint hint +hint,SparkDataFrame,character-method hint +histogram histogram +histogram,SparkDataFrame,characterOrColumn-method histogram +hour column_datetime_functions +hour,Column-method column_datetime_functions +hypot column_math_functions +hypot,Column-method column_math_functions +ifelse column_nonaggregate_functions +ifelse,Column-method column_nonaggregate_functions +ilike columnfunctions +initcap column_string_functions +initcap,Column-method column_string_functions +input_file_name column_nonaggregate_functions +input_file_name,missing-method column_nonaggregate_functions +insertInto insertInto +insertInto,SparkDataFrame,character-method insertInto +install.spark install.spark +instr column_string_functions +instr,Column,character-method column_string_functions +intersect intersect +intersect,SparkDataFrame,SparkDataFrame-method intersect +intersectAll intersectAll +intersectAll,SparkDataFrame,SparkDataFrame-method intersectAll +invoke_higher_order_function invoke_higher_order_function +is.nan column_nonaggregate_functions +is.nan,Column-method column_nonaggregate_functions +isActive isActive +isActive,StreamingQuery-method isActive +isLocal isLocal +isLocal,SparkDataFrame-method isLocal +isNaN columnfunctions +isnan column_nonaggregate_functions +isnan,Column-method column_nonaggregate_functions +isNotNull columnfunctions +isNull columnfunctions +IsotonicRegressionModel-class IsotonicRegressionModel-class +isStreaming isStreaming +isStreaming,SparkDataFrame-method isStreaming +join join +join,SparkDataFrame,SparkDataFrame-method join +KMeansModel-class KMeansModel-class +KSTest-class KSTest-class +kurtosis column_aggregate_functions +kurtosis,Column-method column_aggregate_functions +lag column_window_functions +lag,characterOrColumn-method column_window_functions +last last +last,characterOrColumn-method last +lastProgress lastProgress +lastProgress,StreamingQuery-method lastProgress +last_day column_datetime_functions +last_day,Column-method column_datetime_functions +LDAModel-class LDAModel-class +lead column_window_functions +lead,characterOrColumn,numeric-method column_window_functions +least column_nonaggregate_functions +least,Column-method column_nonaggregate_functions +length column_string_functions +length,Column-method column_string_functions +levenshtein column_string_functions +levenshtein,Column-method column_string_functions +like columnfunctions +limit limit +limit,SparkDataFrame,numeric-method limit +LinearRegressionModel-class LinearRegressionModel-class +LinearSVCModel-class LinearSVCModel-class +listCatalogs listCatalogs +listColumns listColumns +listDatabases listDatabases +listFunctions listFunctions +listTables listTables +lit column_nonaggregate_functions +lit,ANY-method column_nonaggregate_functions +ln column_math_functions +ln,Column-method column_math_functions +loadDF read.df +localCheckpoint localCheckpoint +localCheckpoint,SparkDataFrame-method localCheckpoint +locate column_string_functions +locate,character,Column-method column_string_functions +log column_math_functions +log,Column-method column_math_functions +log10 column_math_functions +log10,Column-method column_math_functions +log1p column_math_functions +log1p,Column-method column_math_functions +log2 column_math_functions +log2,Column-method column_math_functions +LogisticRegressionModel-class LogisticRegressionModel-class +lower column_string_functions +lower,Column-method column_string_functions +lpad column_string_functions +lpad,Column,numeric,character-method column_string_functions +ltrim column_string_functions +ltrim,Column,character-method column_string_functions +ltrim,Column,missing-method column_string_functions +make_date column_datetime_functions +make_date,Column,Column,Column-method column_datetime_functions +make_date,Column-method column_datetime_functions +map_concat column_collection_functions +map_concat,Column-method column_collection_functions +map_entries column_collection_functions +map_entries,Column-method column_collection_functions +map_filter column_collection_functions +map_filter,characterOrColumn,function-method column_collection_functions +map_from_arrays column_collection_functions +map_from_arrays,Column,Column-method column_collection_functions +map_from_arrays,Column-method column_collection_functions +map_from_entries column_collection_functions +map_from_entries,Column-method column_collection_functions +map_keys column_collection_functions +map_keys,Column-method column_collection_functions +map_values column_collection_functions +map_values,Column-method column_collection_functions +map_zip_with column_collection_functions +map_zip_with,characterOrColumn,characterOrColumn,function-method column_collection_functions +max column_aggregate_functions +max,Column-method column_aggregate_functions +max_by column_aggregate_functions +max_by,Column,Column-method column_aggregate_functions +max_by,Column-method column_aggregate_functions +md5 column_misc_functions +md5,Column-method column_misc_functions +mean column_aggregate_functions +mean,Column-method column_aggregate_functions +melt unpivot +melt,SparkDataFrame,ANY,ANY,character,character-method unpivot +merge merge +merge,SparkDataFrame,SparkDataFrame-method merge +min column_aggregate_functions +min,Column-method column_aggregate_functions +minute column_datetime_functions +minute,Column-method column_datetime_functions +min_by column_aggregate_functions +min_by,Column,Column-method column_aggregate_functions +min_by,Column-method column_aggregate_functions +monotonically_increasing_id column_nonaggregate_functions +monotonically_increasing_id,missing-method column_nonaggregate_functions +month column_datetime_functions +month,Column-method column_datetime_functions +months_between column_datetime_diff_functions +months_between,Column-method column_datetime_diff_functions +MultilayerPerceptronClassificationModel-class MultilayerPerceptronClassificationModel-class +mutate mutate +mutate,SparkDataFrame-method mutate +n count +n,Column-method count +na.omit nafunctions +na.omit,SparkDataFrame-method nafunctions +NaiveBayesModel-class NaiveBayesModel-class +names columns +names,SparkDataFrame-method columns +names<- columns +names<-,SparkDataFrame-method columns +nanvl column_nonaggregate_functions +nanvl,Column-method column_nonaggregate_functions +ncol ncol +ncol,SparkDataFrame-method ncol +negate column_nonaggregate_functions +negate,Column-method column_nonaggregate_functions +negative column_nonaggregate_functions +negative,Column-method column_nonaggregate_functions +next_day column_datetime_diff_functions +next_day,Column,character-method column_datetime_diff_functions +not not +not,Column-method not +nrow nrow +nrow,SparkDataFrame-method nrow +nth_value column_window_functions +nth_value,characterOrColumn,numeric-method column_window_functions +nth_value,characterOrColumn-method column_window_functions +ntile column_window_functions +ntile,numeric-method column_window_functions +n_distinct column_aggregate_functions +n_distinct,Column-method column_aggregate_functions +octet_length column_string_functions +octet_length,Column-method column_string_functions +orderBy orderBy +orderBy,SparkDataFrame,characterOrColumn-method arrange +orderBy,WindowSpec,character-method orderBy +orderBy,WindowSpec,Column-method orderBy +otherwise otherwise +otherwise,Column-method otherwise +over over +over,Column,WindowSpec-method over +overlay column_string_functions +overlay,Column,Column,numericOrColumn-method column_string_functions +overlay,Column-method,numericOrColumn-method column_string_functions +partitionBy partitionBy +partitionBy,WindowSpec-method partitionBy +percentile_approx column_aggregate_functions +percentile_approx,characterOrColumn,numericOrColumn-method column_aggregate_functions +percentile_approx,Column-method column_aggregate_functions +percent_rank column_window_functions +percent_rank,missing-method column_window_functions +persist persist +persist,SparkDataFrame,character-method persist +pivot pivot +pivot,GroupedData,character-method pivot +pmod column_math_functions +pmod,Column-method column_math_functions +posexplode column_collection_functions +posexplode,Column-method column_collection_functions +posexplode_outer column_collection_functions +posexplode_outer,Column-method column_collection_functions +positive column_nonaggregate_functions +positive,Column-method column_nonaggregate_functions +PowerIterationClustering-class PowerIterationClustering-class +predict predict +predict,AFTSurvivalRegressionModel-method spark.survreg +predict,ALSModel-method spark.als +predict,BisectingKMeansModel-method spark.bisectingKmeans +predict,DecisionTreeClassificationModel-method spark.decisionTree +predict,DecisionTreeRegressionModel-method spark.decisionTree +predict,FMClassificationModel,SparkDataFrame-method spark.fmClassifier +predict,FMClassificationModel-method spark.fmClassifier +predict,FMRegressionModel,SparkDataFrame-method spark.fmRegressor +predict,FMRegressionModel-method spark.fmRegressor +predict,FPGrowthModel-method spark.fpGrowth +predict,GaussianMixtureModel,SparkDataFrame-method spark.gaussianMixture +predict,GaussianMixtureModel-method spark.gaussianMixture +predict,GBTClassificationModel-method spark.gbt +predict,GBTRegressionModel-method spark.gbt +predict,GeneralizedLinearRegressionModel-method spark.glm +predict,IsotonicRegressionModel,SparkDataFrame-method spark.isoreg +predict,IsotonicRegressionModel-method spark.isoreg +predict,KMeansModel-method spark.kmeans +predict,LinearRegressionModel,SparkDataFrame-method spark.lm +predict,LinearRegressionModel-method spark.lm +predict,LinearSVCModel,SparkDataFrame-method spark.svmLinear +predict,LinearSVCModel-method spark.svmLinear +predict,LogisticRegressionModel,SparkDataFrame-method spark.logit +predict,LogisticRegressionModel-method spark.logit +predict,MultilayerPerceptronClassificationModel-method spark.mlp +predict,NaiveBayesModel-method spark.naiveBayes +predict,RandomForestClassificationModel-method spark.randomForest +predict,RandomForestRegressionModel-method spark.randomForest +PrefixSpan-class PrefixSpan-class +print.jobj print.jobj +print.structField print.structField +print.structType print.structType +print.summary.DecisionTreeClassificationModel spark.decisionTree +print.summary.DecisionTreeRegressionModel spark.decisionTree +print.summary.GBTClassificationModel spark.gbt +print.summary.GBTRegressionModel spark.gbt +print.summary.GeneralizedLinearRegressionModel spark.glm +print.summary.KSTest spark.kstest +print.summary.RandomForestClassificationModel spark.randomForest +print.summary.RandomForestRegressionModel spark.randomForest +printSchema printSchema +printSchema,SparkDataFrame-method printSchema +product column_aggregate_functions +product,Column-method column_aggregate_functions +quarter column_datetime_functions +quarter,Column-method column_datetime_functions +queryName queryName +queryName,StreamingQuery-method queryName +radians column_math_functions +radians,Column-method column_math_functions +raise_error column_misc_functions +raise_error,characterOrColumn-method column_misc_functions +rand column_nonaggregate_functions +rand,missing-method column_nonaggregate_functions +rand,numeric-method column_nonaggregate_functions +randn column_nonaggregate_functions +randn,missing-method column_nonaggregate_functions +randn,numeric-method column_nonaggregate_functions +RandomForestClassificationModel-class RandomForestClassificationModel-class +RandomForestRegressionModel-class RandomForestRegressionModel-class +randomSplit randomSplit +randomSplit,SparkDataFrame,numeric-method randomSplit +rangeBetween rangeBetween +rangeBetween,WindowSpec,numeric,numeric-method rangeBetween +rank column_window_functions +rank,ANY-method column_window_functions +rank,missing-method column_window_functions +rbind rbind +rbind,SparkDataFrame-method rbind +read.df read.df +read.jdbc read.jdbc +read.json read.json +read.ml read.ml +read.orc read.orc +read.parquet read.parquet +read.stream read.stream +read.text read.text +recoverPartitions recoverPartitions +refreshByPath refreshByPath +refreshTable refreshTable +regexp_extract column_string_functions +regexp_extract,Column,character,numeric-method column_string_functions +regexp_replace column_string_functions +regexp_replace,Column,character,character-method column_string_functions +registerTempTable registerTempTable-deprecated +registerTempTable,SparkDataFrame,character-method registerTempTable-deprecated +removeJobTag removeJobTAg +rename rename +rename,SparkDataFrame-method rename +repartition repartition +repartition,SparkDataFrame-method repartition +repartitionByRange repartitionByRange +repartitionByRange,SparkDataFrame-method repartitionByRange +repeat_string column_string_functions +repeat_string,Column,numeric-method column_string_functions +repeat_string,Column-method column_string_functions +reverse column_collection_functions +reverse,Column-method column_collection_functions +rint column_math_functions +rint,Column-method column_math_functions +rlike columnfunctions +rollup rollup +rollup,SparkDataFrame-method rollup +round column_math_functions +round,Column-method column_math_functions +rowsBetween rowsBetween +rowsBetween,WindowSpec,numeric,numeric-method rowsBetween +row_number column_window_functions +row_number,missing-method column_window_functions +rpad column_string_functions +rpad,Column,numeric,character-method column_string_functions +rtrim column_string_functions +rtrim,Column,character-method column_string_functions +rtrim,Column,missing-method column_string_functions +sample sample +sample,SparkDataFrame-method sample +sampleBy sampleBy +sampleBy,SparkDataFrame,character,list,numeric-method sampleBy +sample_frac sample +sample_frac,SparkDataFrame-method sample +saveAsTable saveAsTable +saveAsTable,SparkDataFrame,character-method saveAsTable +saveDF write.df +saveDF,SparkDataFrame,character-method write.df +schema schema +schema,SparkDataFrame-method schema +schema_of_csv column_collection_functions +schema_of_csv,characterOrColumn-method column_collection_functions +schema_of_json column_collection_functions +schema_of_json,characterOrColumn-method column_collection_functions +sd column_aggregate_functions +sd,Column-method column_aggregate_functions +sec column_math_functions +sec,Column-method column_math_functions +second column_datetime_functions +second,Column-method column_datetime_functions +select select +select,SparkDataFrame,character-method select +select,SparkDataFrame,Column-method select +select,SparkDataFrame,list-method select +selectExpr selectExpr +selectExpr,SparkDataFrame,character-method selectExpr +setCheckpointDir setCheckpointDir +setCurrentCatalog setCurrentCatalog +setCurrentDatabase setCurrentDatabase +setInterruptOnCancel setInterruptOnCancel +setJobDescription setJobDescription +setJobGroup setJobGroup +setLocalProperty setLocalProperty +setLogLevel setLogLevel +sha1 column_misc_functions +sha1,Column-method column_misc_functions +sha2 column_misc_functions +sha2,Column,numeric-method column_misc_functions +shiftLeft column_math_functions +shiftleft column_math_functions +shiftLeft,Column,numeric-method column_math_functions +shiftleft,Column,numeric-method column_math_functions +shiftRight column_math_functions +shiftright column_math_functions +shiftRight,Column,numeric-method column_math_functions +shiftright,Column,numeric-method column_math_functions +shiftRightUnsigned column_math_functions +shiftrightunsigned column_math_functions +shiftRightUnsigned,Column,numeric-method column_math_functions +shiftrightunsigned,Column,numeric-method column_math_functions +show show +show,Column-method show +show,GroupedData-method show +show,SparkDataFrame-method show +show,StreamingQuery-method show +show,WindowSpec-method show +showDF showDF +showDF,SparkDataFrame-method showDF +shuffle column_collection_functions +shuffle,Column-method column_collection_functions +sign column_math_functions +sign,Column-method column_math_functions +signum column_math_functions +signum,Column-method column_math_functions +sin column_math_functions +sin,Column-method column_math_functions +sinh column_math_functions +sinh,Column-method column_math_functions +size column_collection_functions +size,Column-method column_collection_functions +skewness column_aggregate_functions +skewness,Column-method column_aggregate_functions +slice column_collection_functions +slice,Column-method column_collection_functions +sort_array column_collection_functions +sort_array,Column-method column_collection_functions +soundex column_string_functions +soundex,Column-method column_string_functions +spark.addFile spark.addFile +spark.als spark.als +spark.als,SparkDataFrame-method spark.als +spark.assignClusters spark.powerIterationClustering +spark.assignClusters,SparkDataFrame-method spark.powerIterationClustering +spark.associationRules spark.fpGrowth +spark.associationRules,FPGrowthModel-method spark.fpGrowth +spark.bisectingKmeans spark.bisectingKmeans +spark.bisectingKmeans,SparkDataFrame,formula-method spark.bisectingKmeans +spark.decisionTree spark.decisionTree +spark.decisionTree,SparkDataFrame,formula-method spark.decisionTree +spark.findFrequentSequentialPatterns spark.prefixSpan +spark.findFrequentSequentialPatterns,SparkDataFrame-method spark.prefixSpan +spark.fmClassifier spark.fmClassifier +spark.fmClassifier,SparkDataFrame,formula-method spark.fmClassifier +spark.fmRegressor spark.fmRegressor +spark.fmRegressor,SparkDataFrame,formula-method spark.fmRegressor +spark.fpGrowth spark.fpGrowth +spark.fpGrowth,SparkDataFrame-method spark.fpGrowth +spark.freqItemsets spark.fpGrowth +spark.freqItemsets,FPGrowthModel-method spark.fpGrowth +spark.gaussianMixture spark.gaussianMixture +spark.gaussianMixture,SparkDataFrame,formula-method spark.gaussianMixture +spark.gbt spark.gbt +spark.gbt,SparkDataFrame,formula-method spark.gbt +spark.getSparkFiles spark.getSparkFiles +spark.getSparkFilesRootDirectory spark.getSparkFilesRootDirectory +spark.glm spark.glm +spark.glm,SparkDataFrame,formula-method spark.glm +spark.isoreg spark.isoreg +spark.isoreg,SparkDataFrame,formula-method spark.isoreg +spark.kmeans spark.kmeans +spark.kmeans,SparkDataFrame,formula-method spark.kmeans +spark.kstest spark.kstest +spark.kstest,SparkDataFrame-method spark.kstest +spark.lapply spark.lapply +spark.lda spark.lda +spark.lda,SparkDataFrame-method spark.lda +spark.lm spark.lm +spark.lm,SparkDataFrame,formula-method spark.lm +spark.logit spark.logit +spark.logit,SparkDataFrame,formula-method spark.logit +spark.mlp spark.mlp +spark.mlp,SparkDataFrame,formula-method spark.mlp +spark.naiveBayes spark.naiveBayes +spark.naiveBayes,SparkDataFrame,formula-method spark.naiveBayes +spark.perplexity spark.lda +spark.perplexity,LDAModel,SparkDataFrame-method spark.lda +spark.perplexity,LDAModel-method spark.lda +spark.posterior spark.lda +spark.posterior,LDAModel,SparkDataFrame-method spark.lda +spark.randomForest spark.randomForest +spark.randomForest,SparkDataFrame,formula-method spark.randomForest +spark.survreg spark.survreg +spark.survreg,SparkDataFrame,formula-method spark.survreg +spark.svmLinear spark.svmLinear +spark.svmLinear,SparkDataFrame,formula-method spark.svmLinear +SparkDataFrame-class SparkDataFrame +sparkR.callJMethod sparkR.callJMethod +sparkR.callJStatic sparkR.callJStatic +sparkR.conf sparkR.conf +sparkR.init sparkR.init-deprecated +sparkR.newJObject sparkR.newJObject +sparkR.session sparkR.session +sparkR.session.stop sparkR.session.stop +sparkR.stop sparkR.session.stop +sparkR.uiWebUrl sparkR.uiWebUrl +sparkR.version sparkR.version +sparkRHive.init sparkRHive.init-deprecated +sparkRSQL.init sparkRSQL.init-deprecated +spark_partition_id column_nonaggregate_functions +spark_partition_id,missing-method column_nonaggregate_functions +split_string column_string_functions +split_string,Column,character-method column_string_functions +split_string,Column-method column_string_functions +sql sql +sqrt column_math_functions +sqrt,Column-method column_math_functions +startsWith startsWith +startsWith,Column-method startsWith +status status +status,StreamingQuery-method status +std column_aggregate_functions +std,Column-method column_aggregate_functions +stddev column_aggregate_functions +stddev,Column-method column_aggregate_functions +stddev_pop column_aggregate_functions +stddev_pop,Column-method column_aggregate_functions +stddev_samp column_aggregate_functions +stddev_samp,Column-method column_aggregate_functions +stopQuery stopQuery +stopQuery,StreamingQuery-method stopQuery +storageLevel storageLevel +storageLevel,SparkDataFrame-method storageLevel +str str +str,SparkDataFrame-method str +StreamingQuery-class StreamingQuery +struct column_nonaggregate_functions +struct,characterOrColumn-method column_nonaggregate_functions +structField structField +structField.character structField +structField.jobj structField +structType structType +structType.character structType +structType.jobj structType +structType.structField structType +subset subset +subset,SparkDataFrame-method subset +substr substr +substr,Column-method substr +substring_index column_string_functions +substring_index,Column,character,numeric-method column_string_functions +sum column_aggregate_functions +sum,Column-method column_aggregate_functions +sumDistinct column_aggregate_functions +sumDistinct,Column-method column_aggregate_functions +summarize summarize +summarize,GroupedData-method summarize +summarize,SparkDataFrame-method summarize +summary summary +summary,AFTSurvivalRegressionModel-method spark.survreg +summary,ALSModel-method spark.als +summary,BisectingKMeansModel-method spark.bisectingKmeans +summary,DecisionTreeClassificationModel-method spark.decisionTree +summary,DecisionTreeRegressionModel-method spark.decisionTree +summary,FMClassificationModel-method spark.fmClassifier +summary,FMRegressionModel-method spark.fmRegressor +summary,GaussianMixtureModel-method spark.gaussianMixture +summary,GBTClassificationModel-method spark.gbt +summary,GBTRegressionModel-method spark.gbt +summary,GeneralizedLinearRegressionModel-method spark.glm +summary,IsotonicRegressionModel-method spark.isoreg +summary,KMeansModel-method spark.kmeans +summary,KSTest-method spark.kstest +summary,LDAModel-method spark.lda +summary,LinearRegressionModel-method spark.lm +summary,LinearSVCModel-method spark.svmLinear +summary,LogisticRegressionModel-method spark.logit +summary,MultilayerPerceptronClassificationModel-method spark.mlp +summary,NaiveBayesModel-method spark.naiveBayes +summary,RandomForestClassificationModel-method spark.randomForest +summary,RandomForestRegressionModel-method spark.randomForest +summary,SparkDataFrame-method summary +sum_distinct column_aggregate_functions +sum_distinct,Column-method column_aggregate_functions +tableExists tableExists +tableNames tableNames +tables tables +tableToDF tableToDF +take take +take,SparkDataFrame,numeric-method take +tan column_math_functions +tan,Column-method column_math_functions +tanh column_math_functions +tanh,Column-method column_math_functions +timestamp_seconds column_datetime_functions +timestamp_seconds,Column-method column_datetime_functions +toDegrees column_math_functions +toDegrees,Column-method column_math_functions +toJSON toJSON +toJSON,SparkDataFrame-method toJSON +toRadians column_math_functions +toRadians,Column-method column_math_functions +to_avro column_avro_functions +to_avro,characterOrColumn-method column_avro_functions +to_avro,Column-method column_avro_functions +to_csv column_collection_functions +to_csv,Column-method column_collection_functions +to_date column_datetime_functions +to_date,Column,character-method column_datetime_functions +to_date,Column,missing-method column_datetime_functions +to_json column_collection_functions +to_json,Column-method column_collection_functions +to_timestamp column_datetime_functions +to_timestamp,Column,character-method column_datetime_functions +to_timestamp,Column,missing-method column_datetime_functions +to_utc_timestamp column_datetime_diff_functions +to_utc_timestamp,Column,character-method column_datetime_diff_functions +transform mutate +transform,SparkDataFrame-method mutate +transform_keys column_collection_functions +transform_keys,characterOrColumn,function-method column_collection_functions +transform_values column_collection_functions +transform_values,characterOrColumn,function-method column_collection_functions +translate column_string_functions +translate,Column,character,character-method column_string_functions +trim column_string_functions +trim,Column,character-method column_string_functions +trim,Column,missing-method column_string_functions +trunc column_datetime_functions +trunc,Column-method column_datetime_functions +unbase64 column_string_functions +unbase64,Column-method column_string_functions +uncacheTable uncacheTable +unhex column_math_functions +unhex,Column-method column_math_functions +union union +union,SparkDataFrame,SparkDataFrame-method union +unionAll unionAll +unionAll,SparkDataFrame,SparkDataFrame-method unionAll +unionByName unionByName +unionByName,SparkDataFrame,SparkDataFrame-method unionByName +unique distinct +unique,SparkDataFrame-method distinct +unix_timestamp column_datetime_functions +unix_timestamp,Column,character-method column_datetime_functions +unix_timestamp,Column,missing-method column_datetime_functions +unix_timestamp,missing,missing-method column_datetime_functions +unpersist unpersist +unpersist,SparkDataFrame-method unpersist +unpivot unpivot +unpivot,SparkDataFrame,ANY,ANY,character,character-method unpivot +unresolved_named_lambda_var unresolved_named_lambda_var +upper column_string_functions +upper,Column-method column_string_functions +var column_aggregate_functions +var,Column-method column_aggregate_functions +variance column_aggregate_functions +variance,Column-method column_aggregate_functions +var_pop column_aggregate_functions +var_pop,Column-method column_aggregate_functions +var_samp column_aggregate_functions +var_samp,Column-method column_aggregate_functions +vector_to_array column_ml_functions +vector_to_array,Column-method column_ml_functions +weekofyear column_datetime_functions +weekofyear,Column-method column_datetime_functions +when column_nonaggregate_functions +when,Column-method column_nonaggregate_functions +where filter +where,SparkDataFrame,characterOrColumn-method filter +width_bucket column_math_functions +width_bucket,Column,Column,Column,Column-method column_math_functions +width_bucket,Column-method column_math_functions +window column_datetime_functions +window,Column-method column_datetime_functions +windowOrderBy windowOrderBy +windowOrderBy,character-method windowOrderBy +windowOrderBy,Column-method windowOrderBy +windowPartitionBy windowPartitionBy +windowPartitionBy,character-method windowPartitionBy +windowPartitionBy,Column-method windowPartitionBy +WindowSpec-class WindowSpec +with with +with,SparkDataFrame-method with +withColumn withColumn +withColumn,SparkDataFrame,character-method withColumn +withColumnRenamed rename +withColumnRenamed,SparkDataFrame,character,character-method rename +withField withField +withField,Column,character,Column-method withField +withField,Column-method withField +withWatermark withWatermark +withWatermark,SparkDataFrame,character,character-method withWatermark +write.df write.df +write.df,SparkDataFrame-method write.df +write.jdbc write.jdbc +write.jdbc,SparkDataFrame,character,character-method write.jdbc +write.json write.json +write.json,SparkDataFrame,character-method write.json +write.ml write.ml +write.ml,AFTSurvivalRegressionModel,character-method spark.survreg +write.ml,ALSModel,character-method spark.als +write.ml,BisectingKMeansModel,character-method spark.bisectingKmeans +write.ml,DecisionTreeClassificationModel,character-method spark.decisionTree +write.ml,DecisionTreeRegressionModel,character-method spark.decisionTree +write.ml,FMClassificationModel,character-method spark.fmClassifier +write.ml,FMRegressionModel,character-method spark.fmRegressor +write.ml,FPGrowthModel,character-method spark.fpGrowth +write.ml,GaussianMixtureModel,character-method spark.gaussianMixture +write.ml,GBTClassificationModel,character-method spark.gbt +write.ml,GBTRegressionModel,character-method spark.gbt +write.ml,GeneralizedLinearRegressionModel,character-method spark.glm +write.ml,IsotonicRegressionModel,character-method spark.isoreg +write.ml,KMeansModel,character-method spark.kmeans +write.ml,LDAModel,character-method spark.lda +write.ml,LinearRegressionModel,character-method spark.lm +write.ml,LinearSVCModel,character-method spark.svmLinear +write.ml,LogisticRegressionModel,character-method spark.logit +write.ml,MultilayerPerceptronClassificationModel,character-method spark.mlp +write.ml,NaiveBayesModel,character-method spark.naiveBayes +write.ml,RandomForestClassificationModel,character-method spark.randomForest +write.ml,RandomForestRegressionModel,character-method spark.randomForest +write.orc write.orc +write.orc,SparkDataFrame,character-method write.orc +write.parquet write.parquet +write.parquet,SparkDataFrame,character-method write.parquet +write.stream write.stream +write.stream,SparkDataFrame-method write.stream +write.text write.text +write.text,SparkDataFrame,character-method write.text +xxhash64 column_misc_functions +xxhash64,Column-method column_misc_functions +year column_datetime_functions +year,Column-method column_datetime_functions +[ subset +[,SparkDataFrame-method subset +[[ subset +[[,SparkDataFrame,numericOrcharacter-method subset +[[<- subset +[[<-,SparkDataFrame,numericOrcharacter-method subset diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdb b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdb new file mode 100644 index 0000000..70c15f2 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdb differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdx b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdx new file mode 100644 index 0000000..c595e2f Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/SparkR.rdx differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/aliases.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/aliases.rds new file mode 100644 index 0000000..035b1b5 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/aliases.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/paths.rds b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/paths.rds new file mode 100644 index 0000000..4d7658f Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/help/paths.rds differ diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/00Index.html b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/00Index.html new file mode 100644 index 0000000..9dd836a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/00Index.html @@ -0,0 +1,2227 @@ + + +R: R Front End for 'Apache Spark' + + + +
+

R Front End for 'Apache Spark' + +

+
+
+[Up] +[Top] +

Documentation for package ‘SparkR’ version 3.5.0

+ + + +

Help Pages

+ + +

+A +B +C +D +E +F +G +H +I +J +K +L +M +N +O +P +Q +R +S +T +U +V +W +X +Y +misc +

+ + +

-- A --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
absMath functions for Column operations
abs-methodMath functions for Column operations
acosMath functions for Column operations
acos-methodMath functions for Column operations
acoshMath functions for Column operations
acosh-methodMath functions for Column operations
addJobTagAdd a tag to be assigned to all the jobs started by this thread.
add_monthsDate time arithmetic functions for Column operations
add_months-methodDate time arithmetic functions for Column operations
AFTSurvivalRegressionModel-classS4 class that represents a AFTSurvivalRegressionModel
aggsummarize
agg-methodsummarize
aliasalias
alias-methodalias
ALSModel-classS4 class that represents an ALSModel
approxCountDistinctAggregate functions for Column operations
approxCountDistinct-methodAggregate functions for Column operations
approxQuantileCalculates the approximate quantiles of numerical columns of a SparkDataFrame
approxQuantile-methodCalculates the approximate quantiles of numerical columns of a SparkDataFrame
approx_count_distinctAggregate functions for Column operations
approx_count_distinct-methodAggregate functions for Column operations
arrangeArrange Rows by Variables
arrange-methodArrange Rows by Variables
arrays_overlapCollection functions for Column operations
arrays_overlap-methodCollection functions for Column operations
arrays_zipCollection functions for Column operations
arrays_zip-methodCollection functions for Column operations
arrays_zip_withCollection functions for Column operations
arrays_zip_with-methodCollection functions for Column operations
array_aggregateCollection functions for Column operations
array_aggregate-methodCollection functions for Column operations
array_containsCollection functions for Column operations
array_contains-methodCollection functions for Column operations
array_distinctCollection functions for Column operations
array_distinct-methodCollection functions for Column operations
array_exceptCollection functions for Column operations
array_except-methodCollection functions for Column operations
array_existsCollection functions for Column operations
array_exists-methodCollection functions for Column operations
array_filterCollection functions for Column operations
array_filter-methodCollection functions for Column operations
array_forallCollection functions for Column operations
array_forall-methodCollection functions for Column operations
array_intersectCollection functions for Column operations
array_intersect-methodCollection functions for Column operations
array_joinCollection functions for Column operations
array_join-methodCollection functions for Column operations
array_maxCollection functions for Column operations
array_max-methodCollection functions for Column operations
array_minCollection functions for Column operations
array_min-methodCollection functions for Column operations
array_positionCollection functions for Column operations
array_position-methodCollection functions for Column operations
array_removeCollection functions for Column operations
array_remove-methodCollection functions for Column operations
array_repeatCollection functions for Column operations
array_repeat-methodCollection functions for Column operations
array_sortCollection functions for Column operations
array_sort-methodCollection functions for Column operations
array_to_vectorML functions for Column operations
array_to_vector-methodML functions for Column operations
array_transformCollection functions for Column operations
array_transform-methodCollection functions for Column operations
array_unionCollection functions for Column operations
array_union-methodCollection functions for Column operations
as.data.frameDownload data from a SparkDataFrame into a R data.frame
as.data.frame-methodDownload data from a SparkDataFrame into a R data.frame
as.DataFrameCreate a SparkDataFrame
ascA set of operations working with SparkDataFrame columns
asciiString functions for Column operations
ascii-methodString functions for Column operations
asc_nulls_firstA set of operations working with SparkDataFrame columns
asc_nulls_lastA set of operations working with SparkDataFrame columns
asinMath functions for Column operations
asin-methodMath functions for Column operations
asinhMath functions for Column operations
asinh-methodMath functions for Column operations
assert_trueMiscellaneous functions for Column operations
assert_true-methodMiscellaneous functions for Column operations
associationRules-methodFP-growth
atanMath functions for Column operations
atan-methodMath functions for Column operations
atan2Math functions for Column operations
atan2-methodMath functions for Column operations
atanhMath functions for Column operations
atanh-methodMath functions for Column operations
attachAttach SparkDataFrame to R search path
attach-methodAttach SparkDataFrame to R search path
avgavg
avg-methodavg
awaitTerminationawaitTermination
awaitTermination-methodawaitTermination
+ +

-- B --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
base64String functions for Column operations
base64-methodString functions for Column operations
betweenbetween
between-methodbetween
binMath functions for Column operations
bin-methodMath functions for Column operations
BisectingKMeansModel-classS4 class that represents a BisectingKMeansModel
bitwiseNOTNon-aggregate functions for Column operations
bitwiseNOT-methodNon-aggregate functions for Column operations
bitwise_notNon-aggregate functions for Column operations
bitwise_not-methodNon-aggregate functions for Column operations
bit_lengthString functions for Column operations
bit_length-methodString functions for Column operations
broadcastbroadcast
broadcast-methodbroadcast
broundMath functions for Column operations
bround-methodMath functions for Column operations
+ +

-- C --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
cacheCache
cache-methodCache
cacheTableCache Table
cancelJobGroupCancel active jobs for the specified group
cancelJobsWithTagCancel active jobs that have the specified tag.
castCasts the column to a different data type.
cast-methodCasts the column to a different data type.
cbrtMath functions for Column operations
cbrt-methodMath functions for Column operations
ceilMath functions for Column operations
ceil-methodMath functions for Column operations
ceilingMath functions for Column operations
ceiling-methodMath functions for Column operations
checkpointcheckpoint
checkpoint-methodcheckpoint
clearCacheClear Cache
clearJobGroupClear current job group ID and its description
clearJobTagsClear the current thread's job tags.
coalesceCoalesce
coalesce-methodCoalesce
coalesce-methodNon-aggregate functions for Column operations
collectCollects all the elements of a SparkDataFrame and coerces them into an R data.frame.
collect-methodCollects all the elements of a SparkDataFrame and coerces them into an R data.frame.
collect_listAggregate functions for Column operations
collect_list-methodAggregate functions for Column operations
collect_setAggregate functions for Column operations
collect_set-methodAggregate functions for Column operations
colnamesColumn Names of SparkDataFrame
colnames-methodColumn Names of SparkDataFrame
colnames<-Column Names of SparkDataFrame
colnames<--methodColumn Names of SparkDataFrame
coltypescoltypes
coltypes-methodcoltypes
coltypes<-coltypes
coltypes<--methodcoltypes
columnS4 class that represents a SparkDataFrame column
Column-classS4 class that represents a SparkDataFrame column
column-methodS4 class that represents a SparkDataFrame column
columnfunctionsA set of operations working with SparkDataFrame columns
columnsColumn Names of SparkDataFrame
columns-methodColumn Names of SparkDataFrame
column_aggregate_functionsAggregate functions for Column operations
column_avro_functionsAvro processing functions for Column operations
column_collection_functionsCollection functions for Column operations
column_datetime_diff_functionsDate time arithmetic functions for Column operations
column_datetime_functionsDate time functions for Column operations
column_math_functionsMath functions for Column operations
column_misc_functionsMiscellaneous functions for Column operations
column_ml_functionsML functions for Column operations
column_nonaggregate_functionsNon-aggregate functions for Column operations
column_string_functionsString functions for Column operations
column_window_functionsWindow functions for Column operations
concatCollection functions for Column operations
concat-methodCollection functions for Column operations
concat_wsString functions for Column operations
concat_ws-methodString functions for Column operations
containsA set of operations working with SparkDataFrame columns
convMath functions for Column operations
conv-methodMath functions for Column operations
corrcorr
corr-methodcorr
cosMath functions for Column operations
cos-methodMath functions for Column operations
coshMath functions for Column operations
cosh-methodMath functions for Column operations
cotMath functions for Column operations
cot-methodMath functions for Column operations
countCount
count-methodCount
count-methodReturns the number of rows in a SparkDataFrame
countDistinctAggregate functions for Column operations
countDistinct-methodAggregate functions for Column operations
count_distinctAggregate functions for Column operations
count_distinct-methodAggregate functions for Column operations
covcov
cov-methodcov
covar_popcov
covar_pop-methodcov
covar_sampcov
covar_samp-methodcov
crc32Miscellaneous functions for Column operations
crc32-methodMiscellaneous functions for Column operations
createDataFrameCreate a SparkDataFrame
createExternalTable(Deprecated) Create an external table
createOrReplaceTempViewCreates a temporary view using the given name.
createOrReplaceTempView-methodCreates a temporary view using the given name.
createTableCreates a table based on the dataset in a data source
create_arrayNon-aggregate functions for Column operations
create_array-methodNon-aggregate functions for Column operations
create_mapNon-aggregate functions for Column operations
create_map-methodNon-aggregate functions for Column operations
crossJoinCrossJoin
crossJoin-methodCrossJoin
crosstabComputes a pair-wise frequency table of the given columns
crosstab-methodComputes a pair-wise frequency table of the given columns
cscMath functions for Column operations
csc-methodMath functions for Column operations
cubecube
cube-methodcube
cume_distWindow functions for Column operations
cume_dist-methodWindow functions for Column operations
currentCatalogReturns the current default catalog
currentDatabaseReturns the current default database
current_dateDate time functions for Column operations
current_date-methodDate time functions for Column operations
current_timestampDate time functions for Column operations
current_timestamp-methodDate time functions for Column operations
+ +

-- D --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
dapplydapply
dapply-methoddapply
dapplyCollectdapplyCollect
dapplyCollect-methoddapplyCollect
databaseExistsChecks if the database with the specified name exists.
datediffDate time arithmetic functions for Column operations
datediff-methodDate time arithmetic functions for Column operations
date_addDate time arithmetic functions for Column operations
date_add-methodDate time arithmetic functions for Column operations
date_formatDate time arithmetic functions for Column operations
date_format-methodDate time arithmetic functions for Column operations
date_subDate time arithmetic functions for Column operations
date_sub-methodDate time arithmetic functions for Column operations
date_truncDate time functions for Column operations
date_trunc-methodDate time functions for Column operations
dayofmonthDate time functions for Column operations
dayofmonth-methodDate time functions for Column operations
dayofweekDate time functions for Column operations
dayofweek-methodDate time functions for Column operations
dayofyearDate time functions for Column operations
dayofyear-methodDate time functions for Column operations
DecisionTreeClassificationModel-classS4 class that represents a DecisionTreeClassificationModel
DecisionTreeRegressionModel-classS4 class that represents a DecisionTreeRegressionModel
decodeString functions for Column operations
decode-methodString functions for Column operations
degreesMath functions for Column operations
degrees-methodMath functions for Column operations
dense_rankWindow functions for Column operations
dense_rank-methodWindow functions for Column operations
descA set of operations working with SparkDataFrame columns
describedescribe
describe-methoddescribe
desc_nulls_firstA set of operations working with SparkDataFrame columns
desc_nulls_lastA set of operations working with SparkDataFrame columns
dimReturns the dimensions of SparkDataFrame
dim-methodReturns the dimensions of SparkDataFrame
distinctDistinct
distinct-methodDistinct
dropdrop
drop-methoddrop
dropDuplicatesdropDuplicates
dropDuplicates-methoddropDuplicates
dropFieldsdropFields
dropFields-methoddropFields
dropnaA set of SparkDataFrame functions working with NA values
dropna-methodA set of SparkDataFrame functions working with NA values
dropTempTable(Deprecated) Drop Temporary Table
dropTempViewDrops the temporary view with the given view name in the catalog.
dtypesDataTypes
dtypes-methodDataTypes
+ +

-- E --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
element_atCollection functions for Column operations
element_at-methodCollection functions for Column operations
encodeString functions for Column operations
encode-methodString functions for Column operations
endsWithendsWith
endsWith-methodendsWith
exceptexcept
except-methodexcept
exceptAllexceptAll
exceptAll-methodexceptAll
expMath functions for Column operations
exp-methodMath functions for Column operations
explainExplain
explain-methodExplain
explodeCollection functions for Column operations
explode-methodCollection functions for Column operations
explode_outerCollection functions for Column operations
explode_outer-methodCollection functions for Column operations
expm1Math functions for Column operations
expm1-methodMath functions for Column operations
exprNon-aggregate functions for Column operations
expr-methodNon-aggregate functions for Column operations
+ +

-- F --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
factorialMath functions for Column operations
factorial-methodMath functions for Column operations
fillnaA set of SparkDataFrame functions working with NA values
fillna-methodA set of SparkDataFrame functions working with NA values
filterFilter
filter-methodFilter
findFrequentSequentialPatterns-methodPrefixSpan
firstReturn the first row of a SparkDataFrame
first-methodReturn the first row of a SparkDataFrame
fittedGet fitted result from a k-means model
fitted-methodGet fitted result from a k-means model
fitted-methodBisecting K-Means Clustering Model
flattenCollection functions for Column operations
flatten-methodCollection functions for Column operations
floorMath functions for Column operations
floor-methodMath functions for Column operations
FMClassificationModel-classS4 class that represents a FMClassificationModel
FMRegressionModel-classS4 class that represents a FMRegressionModel
format_numberString functions for Column operations
format_number-methodString functions for Column operations
format_stringString functions for Column operations
format_string-methodString functions for Column operations
FPGrowthModel-classS4 class that represents a FPGrowthModel
freqItemsFinding frequent items for columns, possibly with false positives
freqItems-methodFinding frequent items for columns, possibly with false positives
freqItemsets-methodFP-growth
from_avroAvro processing functions for Column operations
from_avro-methodAvro processing functions for Column operations
from_csvCollection functions for Column operations
from_csv-methodCollection functions for Column operations
from_jsonCollection functions for Column operations
from_json-methodCollection functions for Column operations
from_unixtimeDate time functions for Column operations
from_unixtime-methodDate time functions for Column operations
from_utc_timestampDate time arithmetic functions for Column operations
from_utc_timestamp-methodDate time arithmetic functions for Column operations
functionExistsChecks if the function with the specified name exists.
+ +

-- G --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
gapplygapply
gapply-methodgapply
gapplyCollectgapplyCollect
gapplyCollect-methodgapplyCollect
GaussianMixtureModel-classS4 class that represents a GaussianMixtureModel
GBTClassificationModel-classS4 class that represents a GBTClassificationModel
GBTRegressionModel-classS4 class that represents a GBTRegressionModel
GeneralizedLinearRegressionModel-classS4 class that represents a generalized linear model
getDatabaseGet the database with the specified name
getFieldA set of operations working with SparkDataFrame columns
getFuncGet the function with the specified name
getItemA set of operations working with SparkDataFrame columns
getJobTagsGet the tags that are currently set to be assigned to all the jobs started by this thread.
getLocalPropertyGet a local property set in this thread, or 'NULL' if it is missing. See 'setLocalProperty'.
getNumPartitionsgetNumPartitions
getNumPartitions-methodgetNumPartitions
getTableGet the table with the specified name
glmGeneralized Linear Models (R-compliant)
glm-methodGeneralized Linear Models (R-compliant)
greatestNon-aggregate functions for Column operations
greatest-methodNon-aggregate functions for Column operations
groupByGroupBy
groupBy-methodGroupBy
groupedDataS4 class that represents a GroupedData
GroupedData-classS4 class that represents a GroupedData
grouping_bitAggregate functions for Column operations
grouping_bit-methodAggregate functions for Column operations
grouping_idAggregate functions for Column operations
grouping_id-methodAggregate functions for Column operations
group_byGroupBy
group_by-methodGroupBy
+ +

-- H --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
hashMiscellaneous functions for Column operations
hash-methodMiscellaneous functions for Column operations
headHead
head-methodHead
hexMath functions for Column operations
hex-methodMath functions for Column operations
hinthint
hint-methodhint
histogramCompute histogram statistics for given column
histogram-methodCompute histogram statistics for given column
hourDate time functions for Column operations
hour-methodDate time functions for Column operations
hypotMath functions for Column operations
hypot-methodMath functions for Column operations
+ +

-- I --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ifelseNon-aggregate functions for Column operations
ifelse-methodNon-aggregate functions for Column operations
ilikeA set of operations working with SparkDataFrame columns
initcapString functions for Column operations
initcap-methodString functions for Column operations
input_file_nameNon-aggregate functions for Column operations
input_file_name-methodNon-aggregate functions for Column operations
insertIntoinsertInto
insertInto-methodinsertInto
install.sparkDownload and Install Apache Spark to a Local Directory
instrString functions for Column operations
instr-methodString functions for Column operations
intersectIntersect
intersect-methodIntersect
intersectAllintersectAll
intersectAll-methodintersectAll
is.nanNon-aggregate functions for Column operations
is.nan-methodNon-aggregate functions for Column operations
isActiveisActive
isActive-methodisActive
isLocalisLocal
isLocal-methodisLocal
isNaNA set of operations working with SparkDataFrame columns
isnanNon-aggregate functions for Column operations
isnan-methodNon-aggregate functions for Column operations
isNotNullA set of operations working with SparkDataFrame columns
isNullA set of operations working with SparkDataFrame columns
IsotonicRegressionModel-classS4 class that represents an IsotonicRegressionModel
isStreamingisStreaming
isStreaming-methodisStreaming
+ +

-- J --

+ + + + + + +
joinJoin
join-methodJoin
+ +

-- K --

+ + + + + + + + + + +
KMeansModel-classS4 class that represents a KMeansModel
KSTest-classS4 class that represents an KSTest
kurtosisAggregate functions for Column operations
kurtosis-methodAggregate functions for Column operations
+ +

-- L --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
lagWindow functions for Column operations
lag-methodWindow functions for Column operations
lastlast
last-methodlast
lastProgresslastProgress
lastProgress-methodlastProgress
last_dayDate time functions for Column operations
last_day-methodDate time functions for Column operations
LDAModel-classS4 class that represents an LDAModel
leadWindow functions for Column operations
lead-methodWindow functions for Column operations
leastNon-aggregate functions for Column operations
least-methodNon-aggregate functions for Column operations
lengthString functions for Column operations
length-methodString functions for Column operations
levenshteinString functions for Column operations
levenshtein-methodString functions for Column operations
likeA set of operations working with SparkDataFrame columns
limitLimit
limit-methodLimit
LinearRegressionModel-classS4 class that represents a LinearRegressionModel
LinearSVCModel-classS4 class that represents an LinearSVCModel
listCatalogsReturns a list of catalog available
listColumnsReturns a list of columns for the given table/view in the specified database
listDatabasesReturns a list of databases available
listFunctionsReturns a list of functions registered in the specified database
listTablesReturns a list of tables or views in the specified database
litNon-aggregate functions for Column operations
lit-methodNon-aggregate functions for Column operations
lnMath functions for Column operations
ln-methodMath functions for Column operations
loadDFLoad a SparkDataFrame
localCheckpointlocalCheckpoint
localCheckpoint-methodlocalCheckpoint
locateString functions for Column operations
locate-methodString functions for Column operations
logMath functions for Column operations
log-methodMath functions for Column operations
log10Math functions for Column operations
log10-methodMath functions for Column operations
log1pMath functions for Column operations
log1p-methodMath functions for Column operations
log2Math functions for Column operations
log2-methodMath functions for Column operations
LogisticRegressionModel-classS4 class that represents an LogisticRegressionModel
lowerString functions for Column operations
lower-methodString functions for Column operations
lpadString functions for Column operations
lpad-methodString functions for Column operations
ltrimString functions for Column operations
ltrim-methodString functions for Column operations
+ +

-- M --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
make_dateDate time functions for Column operations
make_date-methodDate time functions for Column operations
map_concatCollection functions for Column operations
map_concat-methodCollection functions for Column operations
map_entriesCollection functions for Column operations
map_entries-methodCollection functions for Column operations
map_filterCollection functions for Column operations
map_filter-methodCollection functions for Column operations
map_from_arraysCollection functions for Column operations
map_from_arrays-methodCollection functions for Column operations
map_from_entriesCollection functions for Column operations
map_from_entries-methodCollection functions for Column operations
map_keysCollection functions for Column operations
map_keys-methodCollection functions for Column operations
map_valuesCollection functions for Column operations
map_values-methodCollection functions for Column operations
map_zip_withCollection functions for Column operations
map_zip_with-methodCollection functions for Column operations
maxAggregate functions for Column operations
max-methodAggregate functions for Column operations
max_byAggregate functions for Column operations
max_by-methodAggregate functions for Column operations
md5Miscellaneous functions for Column operations
md5-methodMiscellaneous functions for Column operations
meanAggregate functions for Column operations
mean-methodAggregate functions for Column operations
meltUnpivot a DataFrame from wide format to long format.
melt-methodUnpivot a DataFrame from wide format to long format.
mergeMerges two data frames
merge-methodMerges two data frames
minAggregate functions for Column operations
min-methodAggregate functions for Column operations
minuteDate time functions for Column operations
minute-methodDate time functions for Column operations
min_byAggregate functions for Column operations
min_by-methodAggregate functions for Column operations
monotonically_increasing_idNon-aggregate functions for Column operations
monotonically_increasing_id-methodNon-aggregate functions for Column operations
monthDate time functions for Column operations
month-methodDate time functions for Column operations
months_betweenDate time arithmetic functions for Column operations
months_between-methodDate time arithmetic functions for Column operations
MultilayerPerceptronClassificationModel-classS4 class that represents a MultilayerPerceptronClassificationModel
mutateMutate
mutate-methodMutate
+ +

-- N --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
nCount
n-methodCount
na.omitA set of SparkDataFrame functions working with NA values
na.omit-methodA set of SparkDataFrame functions working with NA values
NaiveBayesModel-classS4 class that represents a NaiveBayesModel
namesColumn Names of SparkDataFrame
names-methodColumn Names of SparkDataFrame
names<-Column Names of SparkDataFrame
names<--methodColumn Names of SparkDataFrame
nanvlNon-aggregate functions for Column operations
nanvl-methodNon-aggregate functions for Column operations
ncolReturns the number of columns in a SparkDataFrame
ncol-methodReturns the number of columns in a SparkDataFrame
negateNon-aggregate functions for Column operations
negate-methodNon-aggregate functions for Column operations
negativeNon-aggregate functions for Column operations
negative-methodNon-aggregate functions for Column operations
next_dayDate time arithmetic functions for Column operations
next_day-methodDate time arithmetic functions for Column operations
not!
not-method!
nrowReturns the number of rows in a SparkDataFrame
nrow-methodReturns the number of rows in a SparkDataFrame
nth_valueWindow functions for Column operations
nth_value-methodWindow functions for Column operations
ntileWindow functions for Column operations
ntile-methodWindow functions for Column operations
n_distinctAggregate functions for Column operations
n_distinct-methodAggregate functions for Column operations
+ +

-- O --

+ + + + + + + + + + + + + + + + + + + + + + + + +
octet_lengthString functions for Column operations
octet_length-methodString functions for Column operations
orderByOrdering Columns in a WindowSpec
orderBy-methodArrange Rows by Variables
orderBy-methodOrdering Columns in a WindowSpec
otherwiseotherwise
otherwise-methodotherwise
overover
over-methodover
overlayString functions for Column operations
overlay-methodString functions for Column operations
+ +

-- P --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
partitionBypartitionBy
partitionBy-methodpartitionBy
percentile_approxAggregate functions for Column operations
percentile_approx-methodAggregate functions for Column operations
percent_rankWindow functions for Column operations
percent_rank-methodWindow functions for Column operations
persistPersist
persist-methodPersist
pivotPivot a column of the GroupedData and perform the specified aggregation.
pivot-methodPivot a column of the GroupedData and perform the specified aggregation.
pmodMath functions for Column operations
pmod-methodMath functions for Column operations
posexplodeCollection functions for Column operations
posexplode-methodCollection functions for Column operations
posexplode_outerCollection functions for Column operations
posexplode_outer-methodCollection functions for Column operations
positiveNon-aggregate functions for Column operations
positive-methodNon-aggregate functions for Column operations
PowerIterationClustering-classS4 class that represents a PowerIterationClustering
predictMakes predictions from a MLlib model
predict-methodAlternating Least Squares (ALS) for Collaborative Filtering
predict-methodBisecting K-Means Clustering Model
predict-methodDecision Tree Model for Regression and Classification
predict-methodFactorization Machines Classification Model
predict-methodFactorization Machines Regression Model
predict-methodFP-growth
predict-methodMultivariate Gaussian Mixture Model (GMM)
predict-methodGradient Boosted Tree Model for Regression and Classification
predict-methodGeneralized Linear Models
predict-methodIsotonic Regression Model
predict-methodK-Means Clustering Model
predict-methodLinear Regression Model
predict-methodLogistic Regression Model
predict-methodMultilayer Perceptron Classification Model
predict-methodNaive Bayes Models
predict-methodRandom Forest Model for Regression and Classification
predict-methodAccelerated Failure Time (AFT) Survival Regression Model
predict-methodLinear SVM Model
PrefixSpan-classS4 class that represents a PrefixSpan
print.summary.DecisionTreeClassificationModelDecision Tree Model for Regression and Classification
print.summary.DecisionTreeRegressionModelDecision Tree Model for Regression and Classification
print.summary.GBTClassificationModelGradient Boosted Tree Model for Regression and Classification
print.summary.GBTRegressionModelGradient Boosted Tree Model for Regression and Classification
print.summary.GeneralizedLinearRegressionModelGeneralized Linear Models
print.summary.KSTest(One-Sample) Kolmogorov-Smirnov Test
print.summary.RandomForestClassificationModelRandom Forest Model for Regression and Classification
print.summary.RandomForestRegressionModelRandom Forest Model for Regression and Classification
printSchemaPrint Schema of a SparkDataFrame
printSchema-methodPrint Schema of a SparkDataFrame
productAggregate functions for Column operations
product-methodAggregate functions for Column operations
+ +

-- Q --

+ + + + + + + + + + +
quarterDate time functions for Column operations
quarter-methodDate time functions for Column operations
queryNamequeryName
queryName-methodqueryName
+ +

-- R --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
radiansMath functions for Column operations
radians-methodMath functions for Column operations
raise_errorMiscellaneous functions for Column operations
raise_error-methodMiscellaneous functions for Column operations
randNon-aggregate functions for Column operations
rand-methodNon-aggregate functions for Column operations
randnNon-aggregate functions for Column operations
randn-methodNon-aggregate functions for Column operations
RandomForestClassificationModel-classS4 class that represents a RandomForestClassificationModel
RandomForestRegressionModel-classS4 class that represents a RandomForestRegressionModel
randomSplitrandomSplit
randomSplit-methodrandomSplit
rangeBetweenrangeBetween
rangeBetween-methodrangeBetween
rankWindow functions for Column operations
rank-methodWindow functions for Column operations
rbindUnion two or more SparkDataFrames
rbind-methodUnion two or more SparkDataFrames
read.dfLoad a SparkDataFrame
read.jdbcCreate a SparkDataFrame representing the database table accessible via JDBC URL
read.jsonCreate a SparkDataFrame from a JSON file.
read.mlLoad a fitted MLlib model from the input path.
read.orcCreate a SparkDataFrame from an ORC file.
read.parquetCreate a SparkDataFrame from a Parquet file.
read.streamLoad a streaming SparkDataFrame
read.textCreate a SparkDataFrame from a text file.
recoverPartitionsRecovers all the partitions in the directory of a table and update the catalog
refreshByPathInvalidates and refreshes all the cached data and metadata for SparkDataFrame containing path
refreshTableInvalidates and refreshes all the cached data and metadata of the given table
regexp_extractString functions for Column operations
regexp_extract-methodString functions for Column operations
regexp_replaceString functions for Column operations
regexp_replace-methodString functions for Column operations
registerTempTable(Deprecated) Register Temporary Table
registerTempTable-method(Deprecated) Register Temporary Table
removeJobTagRemove a tag previously added to be assigned to all the jobs started by this thread. Noop if such a tag was not added earlier.
renamerename
rename-methodrename
repartitionRepartition
repartition-methodRepartition
repartitionByRangeRepartition by range
repartitionByRange-methodRepartition by range
repeat_stringString functions for Column operations
repeat_string-methodString functions for Column operations
reverseCollection functions for Column operations
reverse-methodCollection functions for Column operations
rintMath functions for Column operations
rint-methodMath functions for Column operations
rlikeA set of operations working with SparkDataFrame columns
rolluprollup
rollup-methodrollup
roundMath functions for Column operations
round-methodMath functions for Column operations
rowsBetweenrowsBetween
rowsBetween-methodrowsBetween
row_numberWindow functions for Column operations
row_number-methodWindow functions for Column operations
rpadString functions for Column operations
rpad-methodString functions for Column operations
rtrimString functions for Column operations
rtrim-methodString functions for Column operations
+ +

-- S --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
sampleSample
sample-methodSample
sampleByReturns a stratified sample without replacement
sampleBy-methodReturns a stratified sample without replacement
sample_fracSample
sample_frac-methodSample
saveAsTableSave the contents of the SparkDataFrame to a data source as a table
saveAsTable-methodSave the contents of the SparkDataFrame to a data source as a table
saveDFSave the contents of SparkDataFrame to a data source.
saveDF-methodSave the contents of SparkDataFrame to a data source.
schemaGet schema object
schema-methodGet schema object
schema_of_csvCollection functions for Column operations
schema_of_csv-methodCollection functions for Column operations
schema_of_jsonCollection functions for Column operations
schema_of_json-methodCollection functions for Column operations
sdAggregate functions for Column operations
sd-methodAggregate functions for Column operations
secMath functions for Column operations
sec-methodMath functions for Column operations
secondDate time functions for Column operations
second-methodDate time functions for Column operations
selectSelect
select-methodSelect
selectExprSelectExpr
selectExpr-methodSelectExpr
setCheckpointDirSet checkpoint directory
setCurrentCatalogSets the current default catalog
setCurrentDatabaseSets the current default database
setInterruptOnCancelSet the behavior of job cancellation from jobs started in this thread.
setJobDescriptionSet a human readable description of the current job.
setJobGroupAssigns a group ID to all the jobs started by this thread until the group ID is set to a different value or cleared.
setLocalPropertySet a local property that affects jobs submitted from this thread, such as the Spark fair scheduler pool.
setLogLevelSet new log level
sha1Miscellaneous functions for Column operations
sha1-methodMiscellaneous functions for Column operations
sha2Miscellaneous functions for Column operations
sha2-methodMiscellaneous functions for Column operations
shiftLeftMath functions for Column operations
shiftleftMath functions for Column operations
shiftLeft-methodMath functions for Column operations
shiftleft-methodMath functions for Column operations
shiftRightMath functions for Column operations
shiftrightMath functions for Column operations
shiftRight-methodMath functions for Column operations
shiftright-methodMath functions for Column operations
shiftRightUnsignedMath functions for Column operations
shiftrightunsignedMath functions for Column operations
shiftRightUnsigned-methodMath functions for Column operations
shiftrightunsigned-methodMath functions for Column operations
showshow
show-methodshow
showDFshowDF
showDF-methodshowDF
shuffleCollection functions for Column operations
shuffle-methodCollection functions for Column operations
signMath functions for Column operations
sign-methodMath functions for Column operations
signumMath functions for Column operations
signum-methodMath functions for Column operations
sinMath functions for Column operations
sin-methodMath functions for Column operations
sinhMath functions for Column operations
sinh-methodMath functions for Column operations
sizeCollection functions for Column operations
size-methodCollection functions for Column operations
skewnessAggregate functions for Column operations
skewness-methodAggregate functions for Column operations
sliceCollection functions for Column operations
slice-methodCollection functions for Column operations
sort_arrayCollection functions for Column operations
sort_array-methodCollection functions for Column operations
soundexString functions for Column operations
soundex-methodString functions for Column operations
spark.addFileAdd a file or directory to be downloaded with this Spark job on every node.
spark.alsAlternating Least Squares (ALS) for Collaborative Filtering
spark.als-methodAlternating Least Squares (ALS) for Collaborative Filtering
spark.assignClustersPowerIterationClustering
spark.assignClusters-methodPowerIterationClustering
spark.associationRulesFP-growth
spark.associationRules-methodFP-growth
spark.bisectingKmeansBisecting K-Means Clustering Model
spark.bisectingKmeans-methodBisecting K-Means Clustering Model
spark.decisionTreeDecision Tree Model for Regression and Classification
spark.decisionTree-methodDecision Tree Model for Regression and Classification
spark.findFrequentSequentialPatternsPrefixSpan
spark.findFrequentSequentialPatterns-methodPrefixSpan
spark.fmClassifierFactorization Machines Classification Model
spark.fmClassifier-methodFactorization Machines Classification Model
spark.fmRegressorFactorization Machines Regression Model
spark.fmRegressor-methodFactorization Machines Regression Model
spark.fpGrowthFP-growth
spark.fpGrowth-methodFP-growth
spark.freqItemsetsFP-growth
spark.freqItemsets-methodFP-growth
spark.gaussianMixtureMultivariate Gaussian Mixture Model (GMM)
spark.gaussianMixture-methodMultivariate Gaussian Mixture Model (GMM)
spark.gbtGradient Boosted Tree Model for Regression and Classification
spark.gbt-methodGradient Boosted Tree Model for Regression and Classification
spark.getSparkFilesGet the absolute path of a file added through spark.addFile.
spark.getSparkFilesRootDirectoryGet the root directory that contains files added through spark.addFile.
spark.glmGeneralized Linear Models
spark.glm-methodGeneralized Linear Models
spark.isoregIsotonic Regression Model
spark.isoreg-methodIsotonic Regression Model
spark.kmeansK-Means Clustering Model
spark.kmeans-methodK-Means Clustering Model
spark.kstest(One-Sample) Kolmogorov-Smirnov Test
spark.kstest-method(One-Sample) Kolmogorov-Smirnov Test
spark.lapplyRun a function over a list of elements, distributing the computations with Spark
spark.ldaLatent Dirichlet Allocation
spark.lda-methodLatent Dirichlet Allocation
spark.lmLinear Regression Model
spark.lm-methodLinear Regression Model
spark.logitLogistic Regression Model
spark.logit-methodLogistic Regression Model
spark.mlpMultilayer Perceptron Classification Model
spark.mlp-methodMultilayer Perceptron Classification Model
spark.naiveBayesNaive Bayes Models
spark.naiveBayes-methodNaive Bayes Models
spark.perplexityLatent Dirichlet Allocation
spark.perplexity-methodLatent Dirichlet Allocation
spark.posteriorLatent Dirichlet Allocation
spark.posterior-methodLatent Dirichlet Allocation
spark.randomForestRandom Forest Model for Regression and Classification
spark.randomForest-methodRandom Forest Model for Regression and Classification
spark.survregAccelerated Failure Time (AFT) Survival Regression Model
spark.survreg-methodAccelerated Failure Time (AFT) Survival Regression Model
spark.svmLinearLinear SVM Model
spark.svmLinear-methodLinear SVM Model
SparkDataFrame-classS4 class that represents a SparkDataFrame
sparkR.callJMethodCall Java Methods
sparkR.callJStaticCall Static Java Methods
sparkR.confGet Runtime Config from the current active SparkSession
sparkR.init(Deprecated) Initialize a new Spark Context
sparkR.newJObjectCreate Java Objects
sparkR.sessionGet the existing SparkSession or initialize a new SparkSession.
sparkR.session.stopStop the Spark Session and Spark Context
sparkR.stopStop the Spark Session and Spark Context
sparkR.uiWebUrlGet the URL of the SparkUI instance for the current active SparkSession
sparkR.versionGet version of Spark on which this application is running
sparkRHive.init(Deprecated) Initialize a new HiveContext
sparkRSQL.init(Deprecated) Initialize a new SQLContext
spark_partition_idNon-aggregate functions for Column operations
spark_partition_id-methodNon-aggregate functions for Column operations
split_stringString functions for Column operations
split_string-methodString functions for Column operations
sqlSQL Query
sqrtMath functions for Column operations
sqrt-methodMath functions for Column operations
startsWithstartsWith
startsWith-methodstartsWith
statusstatus
status-methodstatus
stdAggregate functions for Column operations
std-methodAggregate functions for Column operations
stddevAggregate functions for Column operations
stddev-methodAggregate functions for Column operations
stddev_popAggregate functions for Column operations
stddev_pop-methodAggregate functions for Column operations
stddev_sampAggregate functions for Column operations
stddev_samp-methodAggregate functions for Column operations
stopQuerystopQuery
stopQuery-methodstopQuery
storageLevelStorageLevel
storageLevel-methodStorageLevel
strCompactly display the structure of a dataset
str-methodCompactly display the structure of a dataset
StreamingQuery-classS4 class that represents a StreamingQuery
structNon-aggregate functions for Column operations
struct-methodNon-aggregate functions for Column operations
structFieldstructField
structField.characterstructField
structField.jobjstructField
structTypestructType
structType.characterstructType
structType.jobjstructType
structType.structFieldstructType
subsetSubset
subset-methodSubset
substrsubstr
substr-methodsubstr
substring_indexString functions for Column operations
substring_index-methodString functions for Column operations
sumAggregate functions for Column operations
sum-methodAggregate functions for Column operations
sumDistinctAggregate functions for Column operations
sumDistinct-methodAggregate functions for Column operations
summarizesummarize
summarize-methodsummarize
summarysummary
summary-methodAlternating Least Squares (ALS) for Collaborative Filtering
summary-methodBisecting K-Means Clustering Model
summary-methodDecision Tree Model for Regression and Classification
summary-methodFactorization Machines Classification Model
summary-methodFactorization Machines Regression Model
summary-methodMultivariate Gaussian Mixture Model (GMM)
summary-methodGradient Boosted Tree Model for Regression and Classification
summary-methodGeneralized Linear Models
summary-methodIsotonic Regression Model
summary-methodK-Means Clustering Model
summary-method(One-Sample) Kolmogorov-Smirnov Test
summary-methodLatent Dirichlet Allocation
summary-methodLinear Regression Model
summary-methodLogistic Regression Model
summary-methodMultilayer Perceptron Classification Model
summary-methodNaive Bayes Models
summary-methodRandom Forest Model for Regression and Classification
summary-methodAccelerated Failure Time (AFT) Survival Regression Model
summary-methodLinear SVM Model
summary-methodsummary
sum_distinctAggregate functions for Column operations
sum_distinct-methodAggregate functions for Column operations
+ +

-- T --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
tableExistsChecks if the table with the specified name exists.
tableNamesTable Names
tablesTables
tableToDFCreate a SparkDataFrame from a SparkSQL table or view
takeTake the first NUM rows of a SparkDataFrame and return the results as a R data.frame
take-methodTake the first NUM rows of a SparkDataFrame and return the results as a R data.frame
tanMath functions for Column operations
tan-methodMath functions for Column operations
tanhMath functions for Column operations
tanh-methodMath functions for Column operations
timestamp_secondsDate time functions for Column operations
timestamp_seconds-methodDate time functions for Column operations
toDegreesMath functions for Column operations
toDegrees-methodMath functions for Column operations
toJSONtoJSON
toJSON-methodtoJSON
toRadiansMath functions for Column operations
toRadians-methodMath functions for Column operations
to_avroAvro processing functions for Column operations
to_avro-methodAvro processing functions for Column operations
to_csvCollection functions for Column operations
to_csv-methodCollection functions for Column operations
to_dateDate time functions for Column operations
to_date-methodDate time functions for Column operations
to_jsonCollection functions for Column operations
to_json-methodCollection functions for Column operations
to_timestampDate time functions for Column operations
to_timestamp-methodDate time functions for Column operations
to_utc_timestampDate time arithmetic functions for Column operations
to_utc_timestamp-methodDate time arithmetic functions for Column operations
transformMutate
transform-methodMutate
transform_keysCollection functions for Column operations
transform_keys-methodCollection functions for Column operations
transform_valuesCollection functions for Column operations
transform_values-methodCollection functions for Column operations
translateString functions for Column operations
translate-methodString functions for Column operations
trimString functions for Column operations
trim-methodString functions for Column operations
truncDate time functions for Column operations
trunc-methodDate time functions for Column operations
+ +

-- U --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
unbase64String functions for Column operations
unbase64-methodString functions for Column operations
uncacheTableUncache Table
unhexMath functions for Column operations
unhex-methodMath functions for Column operations
unionReturn a new SparkDataFrame containing the union of rows
union-methodReturn a new SparkDataFrame containing the union of rows
unionAllReturn a new SparkDataFrame containing the union of rows.
unionAll-methodReturn a new SparkDataFrame containing the union of rows.
unionByNameReturn a new SparkDataFrame containing the union of rows, matched by column names
unionByName-methodReturn a new SparkDataFrame containing the union of rows, matched by column names
uniqueDistinct
unique-methodDistinct
unix_timestampDate time functions for Column operations
unix_timestamp-methodDate time functions for Column operations
unpersistUnpersist
unpersist-methodUnpersist
unpivotUnpivot a DataFrame from wide format to long format.
unpivot-methodUnpivot a DataFrame from wide format to long format.
upperString functions for Column operations
upper-methodString functions for Column operations
+ +

-- V --

+ + + + + + + + + + + + + + + + + + + + + + +
varAggregate functions for Column operations
var-methodAggregate functions for Column operations
varianceAggregate functions for Column operations
variance-methodAggregate functions for Column operations
var_popAggregate functions for Column operations
var_pop-methodAggregate functions for Column operations
var_sampAggregate functions for Column operations
var_samp-methodAggregate functions for Column operations
vector_to_arrayML functions for Column operations
vector_to_array-methodML functions for Column operations
+ +

-- W --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
weekofyearDate time functions for Column operations
weekofyear-methodDate time functions for Column operations
whenNon-aggregate functions for Column operations
when-methodNon-aggregate functions for Column operations
whereFilter
where-methodFilter
width_bucketMath functions for Column operations
width_bucket-methodMath functions for Column operations
windowDate time functions for Column operations
window-methodDate time functions for Column operations
windowOrderBywindowOrderBy
windowOrderBy-methodwindowOrderBy
windowPartitionBywindowPartitionBy
windowPartitionBy-methodwindowPartitionBy
WindowSpec-classS4 class that represents a WindowSpec
withEvaluate a R expression in an environment constructed from a SparkDataFrame
with-methodEvaluate a R expression in an environment constructed from a SparkDataFrame
withColumnWithColumn
withColumn-methodWithColumn
withColumnRenamedrename
withColumnRenamed-methodrename
withFieldwithField
withField-methodwithField
withWatermarkwithWatermark
withWatermark-methodwithWatermark
write.dfSave the contents of SparkDataFrame to a data source.
write.df-methodSave the contents of SparkDataFrame to a data source.
write.jdbcSave the content of SparkDataFrame to an external database table via JDBC.
write.jdbc-methodSave the content of SparkDataFrame to an external database table via JDBC.
write.jsonSave the contents of SparkDataFrame as a JSON file
write.json-methodSave the contents of SparkDataFrame as a JSON file
write.mlSaves the MLlib model to the input path
write.ml-methodAlternating Least Squares (ALS) for Collaborative Filtering
write.ml-methodBisecting K-Means Clustering Model
write.ml-methodDecision Tree Model for Regression and Classification
write.ml-methodFactorization Machines Classification Model
write.ml-methodFactorization Machines Regression Model
write.ml-methodFP-growth
write.ml-methodMultivariate Gaussian Mixture Model (GMM)
write.ml-methodGradient Boosted Tree Model for Regression and Classification
write.ml-methodGeneralized Linear Models
write.ml-methodIsotonic Regression Model
write.ml-methodK-Means Clustering Model
write.ml-methodLatent Dirichlet Allocation
write.ml-methodLinear Regression Model
write.ml-methodLogistic Regression Model
write.ml-methodMultilayer Perceptron Classification Model
write.ml-methodNaive Bayes Models
write.ml-methodRandom Forest Model for Regression and Classification
write.ml-methodAccelerated Failure Time (AFT) Survival Regression Model
write.ml-methodLinear SVM Model
write.orcSave the contents of SparkDataFrame as an ORC file, preserving the schema.
write.orc-methodSave the contents of SparkDataFrame as an ORC file, preserving the schema.
write.parquetSave the contents of SparkDataFrame as a Parquet file, preserving the schema.
write.parquet-methodSave the contents of SparkDataFrame as a Parquet file, preserving the schema.
write.streamWrite the streaming SparkDataFrame to a data source.
write.stream-methodWrite the streaming SparkDataFrame to a data source.
write.textSave the content of SparkDataFrame in a text file at the specified path.
write.text-methodSave the content of SparkDataFrame in a text file at the specified path.
+ +

-- X --

+ + + + + + +
xxhash64Miscellaneous functions for Column operations
xxhash64-methodMiscellaneous functions for Column operations
+ +

-- Y --

+ + + + + + +
yearDate time functions for Column operations
year-methodDate time functions for Column operations
+ +

-- misc --

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
!-method!
$Select
$-methodSelect
$<-Select
$<--methodSelect
%<=>%%<=>%
%<=>%-method%<=>%
%in%Match a column with given values.
%in%-methodMatch a column with given values.
[Subset
[-methodSubset
[[Subset
[[-methodSubset
[[<-Subset
[[<--methodSubset
+
diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/R.css b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/R.css new file mode 100644 index 0000000..c228909 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/html/R.css @@ -0,0 +1,130 @@ +@media screen { + .container { + padding-right: 10px; + padding-left: 10px; + margin-right: auto; + margin-left: auto; + max-width: 900px; + } +} + +.rimage img { /* from knitr - for examples and demos */ + width: 96%; + margin-left: 2%; +} + +.katex { font-size: 1.1em; } + +code { + color: inherit; + background: inherit; +} + +body { + line-height: 1.4; + background: white; + color: black; +} + +a:link { + background: white; + color: blue; +} + +a:visited { + background: white; + color: rgb(50%, 0%, 50%); +} + +h1 { + background: white; + color: rgb(55%, 55%, 55%); + font-family: monospace; + font-size: 1.4em; /* x-large; */ + text-align: center; +} + +h2 { + background: white; + color: rgb(40%, 40%, 40%); + font-family: monospace; + font-size: 1.2em; /* large; */ + text-align: center; +} + +h3 { + background: white; + color: rgb(40%, 40%, 40%); + font-family: monospace; + font-size: 1.2em; /* large; */ +} + +h4 { + background: white; + color: rgb(40%, 40%, 40%); + font-family: monospace; + font-style: italic; + font-size: 1.2em; /* large; */ +} + +h5 { + background: white; + color: rgb(40%, 40%, 40%); + font-family: monospace; +} + +h6 { + background: white; + color: rgb(40%, 40%, 40%); + font-family: monospace; + font-style: italic; +} + +img.toplogo { + width: 4em; + vertical-align: middle; +} + +img.arrow { + width: 30px; + height: 30px; + border: 0; +} + +span.acronym { + font-size: small; +} + +span.env { + font-family: monospace; +} + +span.file { + font-family: monospace; +} + +span.option{ + font-family: monospace; +} + +span.pkg { + font-weight: bold; +} + +span.samp{ + font-family: monospace; +} + +div.vignettes a:hover { + background: rgb(85%, 85%, 85%); +} + +tr { + vertical-align: top; +} + +span.rlang { + font-family: Courier New, Courier; + color: #666666; +} + diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/general.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/general.R new file mode 100644 index 0000000..8c75c19 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/general.R @@ -0,0 +1,23 @@ +# +# 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. +# + +.First <- function() { + packageDir <- Sys.getenv("SPARKR_PACKAGE_DIR") + dirs <- strsplit(packageDir, ",")[[1]] + .libPaths(c(dirs, .libPaths())) + Sys.setenv(NOAWT = 1) +} diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/shell.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/shell.R new file mode 100644 index 0000000..ffedb30 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/profile/shell.R @@ -0,0 +1,49 @@ +# +# 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. +# + +.First <- function() { + home <- Sys.getenv("SPARK_HOME") + .libPaths(c(file.path(home, "R", "lib"), .libPaths())) + Sys.setenv(NOAWT = 1) + + # Make sure SparkR package is the last loaded one + old <- getOption("defaultPackages") + options(defaultPackages = c(old, "SparkR")) + + spark <- SparkR::sparkR.session() + assign("spark", spark, envir = .GlobalEnv) + sc <- SparkR:::callJStatic("org.apache.spark.sql.api.r.SQLUtils", "getJavaSparkContext", spark) + assign("sc", sc, envir = .GlobalEnv) + sparkVer <- SparkR:::callJMethod(sc, "version") + cat("\nWelcome to") + cat("\n") + cat(" ____ __", "\n") + cat(" / __/__ ___ _____/ /__", "\n") + cat(" _\\ \\/ _ \\/ _ `/ __/ '_/", "\n") + cat(" /___/ .__/\\_,_/_/ /_/\\_\\") + if (nchar(sparkVer) == 0) { + cat("\n") + } else { + cat(" version", sparkVer, "\n") + } + cat(" /_/", "\n") + cat("\n") + + cat("\nSparkSession Web UI available at", SparkR::sparkR.uiWebUrl()) + cat("\nSparkSession available as 'spark'(master = ", unlist(SparkR::sparkR.conf("spark.master")), + ", app id = ", unlist(SparkR::sparkR.conf("spark.app.id")), ").", "\n", sep = "") +} diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/tests/testthat/test_basic.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/tests/testthat/test_basic.R new file mode 100644 index 0000000..80df3d8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/tests/testthat/test_basic.R @@ -0,0 +1,100 @@ +# +# 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. +# + +context("basic tests for CRAN") + +test_that("create DataFrame from list or data.frame", { + tryCatch(checkJavaVersion(), + error = function(e) { skip("error on Java check") }, + warning = function(e) { skip("warning on Java check") }) + + sparkR.session(master = sparkRTestMaster, enableHiveSupport = FALSE, + sparkConfig = sparkRTestConfig) + + i <- 4 + df <- createDataFrame(data.frame(dummy = 1:i)) + expect_equal(count(df), i) + + l <- list(list(a = 1, b = 2), list(a = 3, b = 4)) + df <- createDataFrame(l) + expect_equal(columns(df), c("a", "b")) + + a <- 1:3 + b <- c("a", "b", "c") + ldf <- data.frame(a, b) + df <- createDataFrame(ldf) + expect_equal(columns(df), c("a", "b")) + expect_equal(dtypes(df), list(c("a", "int"), c("b", "string"))) + expect_equal(count(df), 3) + ldf2 <- collect(df) + expect_equal(ldf$a, ldf2$a) + + mtcarsdf <- createDataFrame(mtcars) + expect_equivalent(collect(mtcarsdf), mtcars) + + bytes <- as.raw(c(1, 2, 3)) + df <- createDataFrame(list(list(bytes))) + expect_equal(collect(df)[[1]][[1]], bytes) + + sparkR.session.stop() +}) + +test_that("spark.glm and predict", { + tryCatch(checkJavaVersion(), + error = function(e) { skip("error on Java check") }, + warning = function(e) { skip("warning on Java check") }) + + sparkR.session(master = sparkRTestMaster, enableHiveSupport = FALSE, + sparkConfig = sparkRTestConfig) + + training <- suppressWarnings(createDataFrame(iris)) + # gaussian family + model <- spark.glm(training, Sepal_Width ~ Sepal_Length + Species) + prediction <- predict(model, training) + expect_equal(typeof(take(select(prediction, "prediction"), 1)$prediction), "double") + vals <- collect(select(prediction, "prediction")) + rVals <- predict(glm(Sepal.Width ~ Sepal.Length + Species, data = iris), iris) + expect_true(all(abs(rVals - vals) < 1e-6), rVals - vals) + + # Gamma family + x <- runif(100, -1, 1) + y <- rgamma(100, rate = 10 / exp(0.5 + 1.2 * x), shape = 10) + df <- as.DataFrame(as.data.frame(list(x = x, y = y))) + model <- glm(y ~ x, family = Gamma, df) + out <- capture.output(print(summary(model))) + expect_true(any(grepl("Dispersion parameter for gamma family", out))) + + # tweedie family + model <- spark.glm(training, Sepal_Width ~ Sepal_Length + Species, + family = "tweedie", var.power = 1.2, link.power = 0.0) + prediction <- predict(model, training) + expect_equal(typeof(take(select(prediction, "prediction"), 1)$prediction), "double") + vals <- collect(select(prediction, "prediction")) + + # manual calculation of the R predicted values to avoid dependence on statmod + #' library(statmod) + #' rModel <- glm(Sepal.Width ~ Sepal.Length + Species, data = iris, + #' family = tweedie(var.power = 1.2, link.power = 0.0)) + #' print(coef(rModel)) + + rCoef <- c(0.6455409, 0.1169143, -0.3224752, -0.3282174) + rVals <- exp(as.numeric(model.matrix(Sepal.Width ~ Sepal.Length + Species, + data = iris) %*% rCoef)) + expect_true(all(abs(rVals - vals) < 1e-5), rVals - vals) + + sparkR.session.stop() +}) diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/daemon.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/daemon.R new file mode 100644 index 0000000..4589bb9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/daemon.R @@ -0,0 +1,102 @@ +# +# 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. +# + +# Worker daemon + +rLibDir <- Sys.getenv("SPARKR_RLIBDIR") +connectionTimeout <- as.integer(Sys.getenv("SPARKR_BACKEND_CONNECTION_TIMEOUT", "6000")) +dirs <- strsplit(rLibDir, ",")[[1]] +script <- file.path(dirs[[1]], "SparkR", "worker", "worker.R") + +# preload SparkR package, speedup worker +.libPaths(c(dirs, .libPaths())) +suppressPackageStartupMessages(library(SparkR)) + +port <- as.integer(Sys.getenv("SPARKR_WORKER_PORT")) +inputCon <- socketConnection( + port = port, open = "wb", blocking = TRUE, timeout = connectionTimeout) + +SparkR:::doServerAuth(inputCon, Sys.getenv("SPARKR_WORKER_SECRET")) + +# Waits indefinitely for a socket connection by default. +selectTimeout <- NULL + +while (TRUE) { + ready <- socketSelect(list(inputCon), timeout = selectTimeout) + + # Note that the children should be terminated in the parent. If each child terminates + # itself, it appears that the resource is not released properly, that causes an unexpected + # termination of this daemon due to, for example, running out of file descriptors + # (see SPARK-21093). Therefore, the current implementation tries to retrieve children + # that are exited (but not terminated) and then sends a kill signal to terminate them properly + # in the parent. + # + # There are two paths that it attempts to send a signal to terminate the children in the parent. + # + # 1. Every second if any socket connection is not available and if there are child workers + # running. + # 2. Right after a socket connection is available. + # + # In other words, the parent attempts to send the signal to the children every second if + # any worker is running or right before launching other worker children from the following + # new socket connection. + + # The process IDs of exited children are returned below. + children <- parallel:::selectChildren(timeout = 0) + + if (is.integer(children)) { + lapply(children, function(child) { + # This should be the PIDs of exited children. Otherwise, this returns raw bytes if any data + # was sent from this child. In this case, we discard it. + pid <- parallel:::readChild(child) + if (is.integer(pid)) { + # This checks if the data from this child is the same pid of this selected child. + if (child == pid) { + # If so, we terminate this child. + tools::pskill(child, tools::SIGUSR1) + } + } + }) + } else if (is.null(children)) { + # If it is NULL, there are no children. Waits indefinitely for a socket connection. + selectTimeout <- NULL + } + + if (ready) { + port <- SparkR:::readInt(inputCon) + # There is a small chance that it could be interrupted by signal, retry one time + if (length(port) == 0) { + port <- SparkR:::readInt(inputCon) + if (length(port) == 0) { + cat("quitting daemon\n") + quit(save = "no") + } + } + p <- parallel:::mcfork() + if (inherits(p, "masterProcess")) { + # Reach here because this is a child process. + close(inputCon) + Sys.setenv(SPARKR_WORKER_PORT = port) + try(source(script)) + # Note that this mcexit does not fully terminate this child. + parallel:::mcexit(0L) + } else { + # Forking succeeded and we need to check if they finished their jobs every second. + selectTimeout <- 1 + } + } +} diff --git a/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/worker.R b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/worker.R new file mode 100644 index 0000000..7fc4680 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/R/lib/SparkR/worker/worker.R @@ -0,0 +1,292 @@ +# +# 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. +# + +# Worker class + +# Get current system time +currentTimeSecs <- function() { + as.numeric(Sys.time()) +} + +# Get elapsed time +elapsedSecs <- function() { + proc.time()[3] +} + +compute <- function(mode, partition, serializer, deserializer, key, + colNames, computeFunc, inputData) { + if (mode > 0) { + if (deserializer == "row") { + # Transform the list of rows into a data.frame + # Note that the optional argument stringsAsFactors for rbind is + # available since R 3.2.4. So we set the global option here. + oldOpt <- getOption("stringsAsFactors") + options(stringsAsFactors = FALSE) + + # Handle binary data types + if ("raw" %in% sapply(inputData[[1]], class)) { + inputData <- SparkR:::rbindRaws(inputData) + } else { + inputData <- do.call(rbind.data.frame, inputData) + } + + options(stringsAsFactors = oldOpt) + + names(inputData) <- colNames + } else { + # Check to see if inputData is a valid data.frame + stopifnot(deserializer == "byte" || deserializer == "arrow") + stopifnot(is.data.frame(inputData)) + } + + if (mode == 2) { + output <- computeFunc(key, inputData) + } else { + output <- computeFunc(inputData) + } + if (serializer == "row") { + # Transform the result data.frame back to a list of rows + output <- split(output, seq(nrow(output))) + } else { + # Serialize the output to a byte array + stopifnot(serializer == "byte" || serializer == "arrow") + } + } else { + output <- computeFunc(partition, inputData) + } + return(output) +} + +outputResult <- function(serializer, output, outputCon) { + if (serializer == "byte") { + SparkR:::writeRawSerialize(outputCon, output) + } else if (serializer == "row") { + SparkR:::writeRowSerialize(outputCon, output) + } else if (serializer == "arrow") { + SparkR:::writeSerializeInArrow(outputCon, output) + } else { + # write lines one-by-one with flag + lapply(output, function(line) SparkR:::writeString(outputCon, line)) + } +} + +# Constants +specialLengths <- list(END_OF_STREAM = 0L, TIMING_DATA = -1L) + +# Timing R process boot +bootTime <- currentTimeSecs() +bootElap <- elapsedSecs() + +rLibDir <- Sys.getenv("SPARKR_RLIBDIR") +connectionTimeout <- as.integer(Sys.getenv("SPARKR_BACKEND_CONNECTION_TIMEOUT", "6000")) +dirs <- strsplit(rLibDir, ",")[[1]] +# Set libPaths to include SparkR package as loadNamespace needs this +# TODO: Figure out if we can avoid this by not loading any objects that require +# SparkR namespace +.libPaths(c(dirs, .libPaths())) +suppressPackageStartupMessages(library(SparkR)) + +port <- as.integer(Sys.getenv("SPARKR_WORKER_PORT")) +inputCon <- socketConnection( + port = port, blocking = TRUE, open = "wb", timeout = connectionTimeout) +SparkR:::doServerAuth(inputCon, Sys.getenv("SPARKR_WORKER_SECRET")) + +outputCon <- socketConnection( + port = port, blocking = TRUE, open = "wb", timeout = connectionTimeout) +SparkR:::doServerAuth(outputCon, Sys.getenv("SPARKR_WORKER_SECRET")) + +# read the index of the current partition inside the RDD +partition <- SparkR:::readInt(inputCon) + +deserializer <- SparkR:::readString(inputCon) +serializer <- SparkR:::readString(inputCon) + +# Include packages as required +packageNames <- unserialize(SparkR:::readRaw(inputCon)) +for (pkg in packageNames) { + suppressPackageStartupMessages(library(as.character(pkg), character.only = TRUE)) +} + +# read function dependencies +funcLen <- SparkR:::readInt(inputCon) +computeFunc <- unserialize(SparkR:::readRawLen(inputCon, funcLen)) +env <- environment(computeFunc) +parent.env(env) <- .GlobalEnv # Attach under global environment. + +# Timing init envs for computing +initElap <- elapsedSecs() + +# Read and set broadcast variables +numBroadcastVars <- SparkR:::readInt(inputCon) +if (numBroadcastVars > 0) { + for (bcast in seq(1:numBroadcastVars)) { + bcastId <- SparkR:::readInt(inputCon) + value <- unserialize(SparkR:::readRaw(inputCon)) + SparkR:::setBroadcastValue(bcastId, value) + } +} + +# Timing broadcast +broadcastElap <- elapsedSecs() +# Initial input timing +inputElap <- broadcastElap + +# If -1: read as normal RDD; if >= 0, treat as pairwise RDD and treat the int +# as number of partitions to create. +numPartitions <- SparkR:::readInt(inputCon) + +# 0 - RDD mode, 1 - dapply mode, 2 - gapply mode +mode <- SparkR:::readInt(inputCon) + +if (mode > 0) { + colNames <- SparkR:::readObject(inputCon) +} + +isEmpty <- SparkR:::readInt(inputCon) +computeInputElapsDiff <- 0 +outputComputeElapsDiff <- 0 + +if (isEmpty != 0) { + if (numPartitions == -1) { + if (deserializer == "byte") { + # Now read as many characters as described in funcLen + data <- SparkR:::readDeserialize(inputCon) + } else if (deserializer == "string") { + data <- as.list(readLines(inputCon)) + } else if (deserializer == "row" && mode == 2) { + dataWithKeys <- SparkR:::readMultipleObjectsWithKeys(inputCon) + keys <- dataWithKeys$keys + data <- dataWithKeys$data + } else if (deserializer == "row") { + data <- SparkR:::readMultipleObjects(inputCon) + } else if (deserializer == "arrow" && mode == 2) { + dataWithKeys <- SparkR:::readDeserializeWithKeysInArrow(inputCon) + keys <- dataWithKeys$keys + data <- dataWithKeys$data + } else if (deserializer == "arrow" && mode == 1) { + data <- SparkR:::readDeserializeInArrow(inputCon) + # See https://stat.ethz.ch/pipermail/r-help/2010-September/252046.html + # rbind.fill might be an alternative to make it faster if plyr is installed. + # Also, note that, 'dapply' applies a function to each partition. + data <- do.call("rbind", data) + } + + # Timing reading input data for execution + inputElap <- elapsedSecs() + if (mode > 0) { + if (mode == 1) { + output <- compute(mode, partition, serializer, deserializer, NULL, + colNames, computeFunc, data) + } else { + # gapply mode + outputs <- list() + for (i in seq_len(length(data))) { + # Timing reading input data for execution + computeStart <- elapsedSecs() + output <- compute(mode, partition, serializer, deserializer, keys[[i]], + colNames, computeFunc, data[[i]]) + computeElap <- elapsedSecs() + if (serializer == "arrow") { + outputs[[length(outputs) + 1L]] <- output + } else { + outputResult(serializer, output, outputCon) + outputComputeElapsDiff <- outputComputeElapsDiff + (elapsedSecs() - computeElap) + } + computeInputElapsDiff <- computeInputElapsDiff + (computeElap - computeStart) + } + + if (serializer == "arrow") { + # See https://stat.ethz.ch/pipermail/r-help/2010-September/252046.html + # rbind.fill might be an alternative to make it faster if plyr is installed. + outputStart <- elapsedSecs() + combined <- do.call("rbind", outputs) + SparkR:::writeSerializeInArrow(outputCon, combined) + outputComputeElapsDiff <- elapsedSecs() - outputStart + } + } + } else { + output <- compute(mode, partition, serializer, deserializer, NULL, + colNames, computeFunc, data) + } + if (mode != 2) { + # Not a gapply mode + computeElap <- elapsedSecs() + outputResult(serializer, output, outputCon) + outputElap <- elapsedSecs() + computeInputElapsDiff <- computeElap - inputElap + outputComputeElapsDiff <- outputElap - computeElap + } + } else { + if (deserializer == "byte") { + # Now read as many characters as described in funcLen + data <- SparkR:::readDeserialize(inputCon) + } else if (deserializer == "string") { + data <- readLines(inputCon) + } else if (deserializer == "row") { + data <- SparkR:::readMultipleObjects(inputCon) + } + # Timing reading input data for execution + inputElap <- elapsedSecs() + + res <- new.env() + + # Step 1: hash the data to an environment + hashTupleToEnvir <- function(tuple) { + # NOTE: execFunction is the hash function here + hashVal <- computeFunc(tuple[[1]]) + bucket <- as.character(hashVal %% numPartitions) + acc <- res[[bucket]] + # Create a new accumulator + if (is.null(acc)) { + acc <- SparkR:::initAccumulator() + } + SparkR:::addItemToAccumulator(acc, tuple) + res[[bucket]] <- acc + } + invisible(lapply(data, hashTupleToEnvir)) + # Timing computing + computeElap <- elapsedSecs() + + # Step 2: write out all of the environment as key-value pairs. + for (name in ls(res)) { + SparkR:::writeInt(outputCon, 2L) + SparkR:::writeInt(outputCon, as.integer(name)) + # Truncate the accumulator list to the number of elements we have + length(res[[name]]$data) <- res[[name]]$counter + SparkR:::writeRawSerialize(outputCon, res[[name]]$data) + } + # Timing output + outputElap <- elapsedSecs() + computeInputElapsDiff <- computeElap - inputElap + outputComputeElapsDiff <- outputElap - computeElap + } +} + +# Report timing +SparkR:::writeInt(outputCon, specialLengths$TIMING_DATA) +SparkR:::writeDouble(outputCon, bootTime) +SparkR:::writeDouble(outputCon, initElap - bootElap) # init +SparkR:::writeDouble(outputCon, broadcastElap - initElap) # broadcast +SparkR:::writeDouble(outputCon, inputElap - broadcastElap) # input +SparkR:::writeDouble(outputCon, computeInputElapsDiff) # compute +SparkR:::writeDouble(outputCon, outputComputeElapsDiff) # output + +# End of output +SparkR:::writeInt(outputCon, specialLengths$END_OF_STREAM) + +close(outputCon) +close(inputCon) diff --git a/spark-3.5.0-bin-hadoop3/R/lib/sparkr.zip b/spark-3.5.0-bin-hadoop3/R/lib/sparkr.zip new file mode 100644 index 0000000..4bccfc4 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/R/lib/sparkr.zip differ diff --git a/spark-3.5.0-bin-hadoop3/README.md b/spark-3.5.0-bin-hadoop3/README.md new file mode 100644 index 0000000..310df41 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/README.md @@ -0,0 +1,125 @@ +# Apache Spark + +Spark is a unified analytics engine for large-scale data processing. It provides +high-level APIs in Scala, Java, Python, and R, and an optimized engine that +supports general computation graphs for data analysis. It also supports a +rich set of higher-level tools including Spark SQL for SQL and DataFrames, +pandas API on Spark for pandas workloads, MLlib for machine learning, GraphX for graph processing, +and Structured Streaming for stream processing. + + + +[![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_main.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_main.yml) +[![AppVeyor Build](https://img.shields.io/appveyor/ci/ApacheSoftwareFoundation/spark/master.svg?style=plastic&logo=appveyor)](https://ci.appveyor.com/project/ApacheSoftwareFoundation/spark) +[![PySpark Coverage](https://codecov.io/gh/apache/spark/branch/master/graph/badge.svg)](https://codecov.io/gh/apache/spark) +[![PyPI Downloads](https://static.pepy.tech/personalized-badge/pyspark?period=month&units=international_system&left_color=black&right_color=orange&left_text=PyPI%20downloads)](https://pypi.org/project/pyspark/) + + +## Online Documentation + +You can find the latest Spark documentation, including a programming +guide, on the [project web page](https://spark.apache.org/documentation.html). +This README file only contains basic setup instructions. + +## Building Spark + +Spark is built using [Apache Maven](https://maven.apache.org/). +To build Spark and its example programs, run: + +```bash +./build/mvn -DskipTests clean package +``` + +(You do not need to do this if you downloaded a pre-built package.) + +More detailed documentation is available from the project site, at +["Building Spark"](https://spark.apache.org/docs/latest/building-spark.html). + +For general development tips, including info on developing Spark using an IDE, see ["Useful Developer Tools"](https://spark.apache.org/developer-tools.html). + +## Interactive Scala Shell + +The easiest way to start using Spark is through the Scala shell: + +```bash +./bin/spark-shell +``` + +Try the following command, which should return 1,000,000,000: + +```scala +scala> spark.range(1000 * 1000 * 1000).count() +``` + +## Interactive Python Shell + +Alternatively, if you prefer Python, you can use the Python shell: + +```bash +./bin/pyspark +``` + +And run the following command, which should also return 1,000,000,000: + +```python +>>> spark.range(1000 * 1000 * 1000).count() +``` + +## Example Programs + +Spark also comes with several sample programs in the `examples` directory. +To run one of them, use `./bin/run-example [params]`. For example: + +```bash +./bin/run-example SparkPi +``` + +will run the Pi example locally. + +You can set the MASTER environment variable when running examples to submit +examples to a cluster. This can be a mesos:// or spark:// URL, +"yarn" to run on YARN, and "local" to run +locally with one thread, or "local[N]" to run locally with N threads. You +can also use an abbreviated class name if the class is in the `examples` +package. For instance: + +```bash +MASTER=spark://host:7077 ./bin/run-example SparkPi +``` + +Many of the example programs print usage help if no params are given. + +## Running Tests + +Testing first requires [building Spark](#building-spark). Once Spark is built, tests +can be run using: + +```bash +./dev/run-tests +``` + +Please see the guidance on how to +[run tests for a module, or individual tests](https://spark.apache.org/developer-tools.html#individual-tests). + +There is also a Kubernetes integration test, see resource-managers/kubernetes/integration-tests/README.md + +## A Note About Hadoop Versions + +Spark uses the Hadoop core library to talk to HDFS and other Hadoop-supported +storage systems. Because the protocols have changed in different versions of +Hadoop, you must build Spark against the same version that your cluster runs. + +Please refer to the build documentation at +["Specifying the Hadoop Version and Enabling YARN"](https://spark.apache.org/docs/latest/building-spark.html#specifying-the-hadoop-version-and-enabling-yarn) +for detailed guidance on building for a particular distribution of Hadoop, including +building for particular Hive and Hive Thriftserver distributions. + +## Configuration + +Please refer to the [Configuration Guide](https://spark.apache.org/docs/latest/configuration.html) +in the online documentation for an overview on how to configure Spark. + +## Contributing + +Please review the [Contribution to Spark guide](https://spark.apache.org/contributing.html) +for information on how to get started contributing to the project. diff --git a/spark-3.5.0-bin-hadoop3/RELEASE b/spark-3.5.0-bin-hadoop3/RELEASE new file mode 100644 index 0000000..f4a45f1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/RELEASE @@ -0,0 +1,2 @@ +Spark 3.5.0 (git revision ce5ddad9903) built for Hadoop 3.3.4 +Build flags: -B -Pmesos -Pyarn -Pkubernetes -Psparkr -Pscala-2.12 -Phadoop-3 -Phive -Phive-thriftserver diff --git a/spark-3.5.0-bin-hadoop3/bin/beeline b/spark-3.5.0-bin-hadoop3/bin/beeline new file mode 100644 index 0000000..0585346 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/beeline @@ -0,0 +1,32 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# +# Shell script for starting BeeLine + +# Enter posix mode for bash +set -o posix + +# Figure out if SPARK_HOME is set +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +CLASS="org.apache.hive.beeline.BeeLine" +exec "${SPARK_HOME}/bin/spark-class" $CLASS "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/beeline.cmd b/spark-3.5.0-bin-hadoop3/bin/beeline.cmd new file mode 100644 index 0000000..288059a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/beeline.cmd @@ -0,0 +1,22 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-class.cmd" org.apache.hive.beeline.BeeLine %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/docker-image-tool.sh b/spark-3.5.0-bin-hadoop3/bin/docker-image-tool.sh new file mode 100644 index 0000000..a137a2f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/docker-image-tool.sh @@ -0,0 +1,336 @@ +#!/usr/bin/env bash + +# 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. +# + +# This script builds and pushes docker images when run from a release of Spark +# with Kubernetes support. + +function error { + echo "$@" 1>&2 + exit 1 +} + +if [ -z "${SPARK_HOME}" ]; then + SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi +. "${SPARK_HOME}/bin/load-spark-env.sh" + +CTX_DIR="$SPARK_HOME/target/tmp/docker" + +function is_dev_build { + [ ! -f "$SPARK_HOME/RELEASE" ] +} + +function cleanup_ctx_dir { + if is_dev_build; then + rm -rf "$CTX_DIR" + fi +} + +trap cleanup_ctx_dir EXIT + +function image_ref { + local image="$1" + local add_repo="${2:-1}" + if [ $add_repo = 1 ] && [ -n "$REPO" ]; then + image="$REPO/$image" + fi + if [ -n "$TAG" ]; then + image="$image:$TAG" + fi + echo "$image" +} + +function docker_push { + local image_name="$1" + if [ ! -z $(docker images -q "$(image_ref ${image_name})") ]; then + docker push "$(image_ref ${image_name})" + if [ $? -ne 0 ]; then + error "Failed to push $image_name Docker image." + fi + else + echo "$(image_ref ${image_name}) image not found. Skipping push for this image." + fi +} + +function resolve_file { + local FILE=$1 + if [ -n "$FILE" ]; then + local DIR=$(dirname $FILE) + DIR=$(cd $DIR && pwd) + FILE="${DIR}/$(basename $FILE)" + fi + echo $FILE +} + +# Create a smaller build context for docker in dev builds to make the build faster. Docker +# uploads all of the current directory to the daemon, and it can get pretty big with dev +# builds that contain test log files and other artifacts. +# +# Three build contexts are created, one for each image: base, pyspark, and sparkr. For them +# to have the desired effect, the docker command needs to be executed inside the appropriate +# context directory. +# +# Note: docker does not support symlinks in the build context. +function create_dev_build_context {( + set -e + local BASE_CTX="$CTX_DIR/base" + mkdir -p "$BASE_CTX/kubernetes" + cp -r "resource-managers/kubernetes/docker/src/main/dockerfiles" \ + "$BASE_CTX/kubernetes/dockerfiles" + + cp -r "assembly/target/scala-$SPARK_SCALA_VERSION/jars" "$BASE_CTX/jars" + cp -r "resource-managers/kubernetes/integration-tests/tests" \ + "$BASE_CTX/kubernetes/tests" + + mkdir "$BASE_CTX/examples" + cp -r "examples/src" "$BASE_CTX/examples/src" + # Copy just needed examples jars instead of everything. + mkdir "$BASE_CTX/examples/jars" + for i in examples/target/scala-$SPARK_SCALA_VERSION/jars/*; do + if [ ! -f "$BASE_CTX/jars/$(basename $i)" ]; then + cp $i "$BASE_CTX/examples/jars" + fi + done + + for other in bin sbin data; do + cp -r "$other" "$BASE_CTX/$other" + done + + local PYSPARK_CTX="$CTX_DIR/pyspark" + mkdir -p "$PYSPARK_CTX/kubernetes" + cp -r "resource-managers/kubernetes/docker/src/main/dockerfiles" \ + "$PYSPARK_CTX/kubernetes/dockerfiles" + mkdir "$PYSPARK_CTX/python" + cp -r "python/lib" "$PYSPARK_CTX/python/lib" + cp -r "python/pyspark" "$PYSPARK_CTX/python/pyspark" + + local R_CTX="$CTX_DIR/sparkr" + mkdir -p "$R_CTX/kubernetes" + cp -r "resource-managers/kubernetes/docker/src/main/dockerfiles" \ + "$R_CTX/kubernetes/dockerfiles" + cp -r "R" "$R_CTX/R" +)} + +function img_ctx_dir { + if is_dev_build; then + echo "$CTX_DIR/$1" + else + echo "$SPARK_HOME" + fi +} + +function build { + local BUILD_ARGS + local SPARK_ROOT="$SPARK_HOME" + + if is_dev_build; then + create_dev_build_context || error "Failed to create docker build context." + SPARK_ROOT="$CTX_DIR/base" + fi + + # Verify that the Docker image content directory is present + if [ ! -d "$SPARK_ROOT/kubernetes/dockerfiles" ]; then + error "Cannot find docker image. This script must be run from a runnable distribution of Apache Spark." + fi + + # Verify that Spark has actually been built/is a runnable distribution + # i.e. the Spark JARs that the Docker files will place into the image are present + local TOTAL_JARS=$(ls $SPARK_ROOT/jars/spark-* | wc -l) + TOTAL_JARS=$(( $TOTAL_JARS )) + if [ "${TOTAL_JARS}" -eq 0 ]; then + error "Cannot find Spark JARs. This script assumes that Apache Spark has first been built locally or this is a runnable distribution." + fi + + local BUILD_ARGS=(${BUILD_PARAMS}) + + # If a custom SPARK_UID was set add it to build arguments + if [ -n "$SPARK_UID" ]; then + BUILD_ARGS+=(--build-arg spark_uid=$SPARK_UID) + fi + + local BINDING_BUILD_ARGS=( + ${BUILD_ARGS[@]} + --build-arg + base_img=$(image_ref spark) + ) + + local BASEDOCKERFILE=${BASEDOCKERFILE:-"kubernetes/dockerfiles/spark/Dockerfile"} + local PYDOCKERFILE=${PYDOCKERFILE:-false} + local RDOCKERFILE=${RDOCKERFILE:-false} + local ARCHS=${ARCHS:-"--platform linux/amd64,linux/arm64"} + + (cd $(img_ctx_dir base) && docker build $NOCACHEARG "${BUILD_ARGS[@]}" \ + -t $(image_ref spark) \ + -f "$BASEDOCKERFILE" .) + if [ $? -ne 0 ]; then + error "Failed to build Spark JVM Docker image, please refer to Docker build output for details." + fi + if [ "${CROSS_BUILD}" != "false" ]; then + (cd $(img_ctx_dir base) && docker buildx build $ARCHS $NOCACHEARG "${BUILD_ARGS[@]}" --push --provenance=false \ + -t $(image_ref spark) \ + -f "$BASEDOCKERFILE" .) + fi + + if [ "${PYDOCKERFILE}" != "false" ]; then + (cd $(img_ctx_dir pyspark) && docker build $NOCACHEARG "${BINDING_BUILD_ARGS[@]}" \ + -t $(image_ref spark-py) \ + -f "$PYDOCKERFILE" .) + if [ $? -ne 0 ]; then + error "Failed to build PySpark Docker image, please refer to Docker build output for details." + fi + if [ "${CROSS_BUILD}" != "false" ]; then + (cd $(img_ctx_dir pyspark) && docker buildx build $ARCHS $NOCACHEARG "${BINDING_BUILD_ARGS[@]}" --push --provenance=false \ + -t $(image_ref spark-py) \ + -f "$PYDOCKERFILE" .) + fi + fi + + if [ "${RDOCKERFILE}" != "false" ]; then + (cd $(img_ctx_dir sparkr) && docker build $NOCACHEARG "${BINDING_BUILD_ARGS[@]}" \ + -t $(image_ref spark-r) \ + -f "$RDOCKERFILE" .) + if [ $? -ne 0 ]; then + error "Failed to build SparkR Docker image, please refer to Docker build output for details." + fi + if [ "${CROSS_BUILD}" != "false" ]; then + (cd $(img_ctx_dir sparkr) && docker buildx build $ARCHS $NOCACHEARG "${BINDING_BUILD_ARGS[@]}" --push --provenance=false \ + -t $(image_ref spark-r) \ + -f "$RDOCKERFILE" .) + fi + fi +} + +function push { + docker_push "spark" + docker_push "spark-py" + docker_push "spark-r" +} + +function usage { + cat </dev/null; then + error "Cannot find minikube." + fi + if ! minikube status 1>/dev/null; then + error "Cannot contact minikube. Make sure it's running." + fi + eval $(minikube docker-env --shell bash) + ;; + u) SPARK_UID=${OPTARG};; + esac +done + +case "${@: -1}" in + build) + build + ;; + push) + if [ -z "$REPO" ]; then + usage + exit 1 + fi + push + ;; + *) + usage + exit 1 + ;; +esac diff --git a/spark-3.5.0-bin-hadoop3/bin/find-spark-home b/spark-3.5.0-bin-hadoop3/bin/find-spark-home new file mode 100644 index 0000000..462b538 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/find-spark-home @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Attempts to find a proper value for SPARK_HOME. Should be included using "source" directive. + +FIND_SPARK_HOME_PYTHON_SCRIPT="$(cd "$(dirname "$0")"; pwd)/find_spark_home.py" + +# Short circuit if the user already has this set. +if [ ! -z "${SPARK_HOME}" ]; then + exit 0 +elif [ ! -f "$FIND_SPARK_HOME_PYTHON_SCRIPT" ]; then + # If we are not in the same directory as find_spark_home.py we are not pip installed so we don't + # need to search the different Python directories for a Spark installation. + # Note only that, if the user has pip installed PySpark but is directly calling pyspark-shell or + # spark-submit in another directory we want to use that version of PySpark rather than the + # pip installed version of PySpark. + export SPARK_HOME="$(cd "$(dirname "$0")"/..; pwd)" +else + # We are pip installed, use the Python script to resolve a reasonable SPARK_HOME + # Default to standard python3 interpreter unless told otherwise + if [[ -z "$PYSPARK_DRIVER_PYTHON" ]]; then + PYSPARK_DRIVER_PYTHON="${PYSPARK_PYTHON:-"python3"}" + fi + export SPARK_HOME=$($PYSPARK_DRIVER_PYTHON "$FIND_SPARK_HOME_PYTHON_SCRIPT") +fi diff --git a/spark-3.5.0-bin-hadoop3/bin/find-spark-home.cmd b/spark-3.5.0-bin-hadoop3/bin/find-spark-home.cmd new file mode 100644 index 0000000..3149d05 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/find-spark-home.cmd @@ -0,0 +1,60 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Path to Python script finding SPARK_HOME +set FIND_SPARK_HOME_PYTHON_SCRIPT=%~dp0find_spark_home.py + +rem Default to standard python3 interpreter unless told otherwise +set PYTHON_RUNNER=python3 +rem If PYSPARK_DRIVER_PYTHON is set, it overwrites the python version +if not "x%PYSPARK_DRIVER_PYTHON%"=="x" ( + set PYTHON_RUNNER=%PYSPARK_DRIVER_PYTHON% +) +rem If PYSPARK_PYTHON is set, it overwrites the python version +if not "x%PYSPARK_PYTHON%"=="x" ( + set PYTHON_RUNNER=%PYSPARK_PYTHON% +) + +rem If there is python installed, trying to use the root dir as SPARK_HOME +where %PYTHON_RUNNER% > nul 2>&1 +if %ERRORLEVEL% neq 0 ( + if not exist %PYTHON_RUNNER% ( + if "x%SPARK_HOME%"=="x" ( + echo Missing Python executable '%PYTHON_RUNNER%', defaulting to '%~dp0..' for SPARK_HOME ^ +environment variable. Please install Python or specify the correct Python executable in ^ +PYSPARK_DRIVER_PYTHON or PYSPARK_PYTHON environment variable to detect SPARK_HOME safely. + set SPARK_HOME=%~dp0.. + ) + ) +) + +rem Only attempt to find SPARK_HOME if it is not set. +if "x%SPARK_HOME%"=="x" ( + if not exist "%FIND_SPARK_HOME_PYTHON_SCRIPT%" ( + rem If we are not in the same directory as find_spark_home.py we are not pip installed so we don't + rem need to search the different Python directories for a Spark installation. + rem Note only that, if the user has pip installed PySpark but is directly calling pyspark-shell or + rem spark-submit in another directory we want to use that version of PySpark rather than the + rem pip installed version of PySpark. + set SPARK_HOME=%~dp0.. + ) else ( + rem We are pip installed, use the Python script to resolve a reasonable SPARK_HOME + for /f "delims=" %%i in ('%PYTHON_RUNNER% "%FIND_SPARK_HOME_PYTHON_SCRIPT%"') do set SPARK_HOME=%%i + ) +) diff --git a/spark-3.5.0-bin-hadoop3/bin/load-spark-env.cmd b/spark-3.5.0-bin-hadoop3/bin/load-spark-env.cmd new file mode 100644 index 0000000..5692af5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/load-spark-env.cmd @@ -0,0 +1,62 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This script loads spark-env.cmd if it exists, and ensures it is only loaded once. +rem spark-env.cmd is loaded from SPARK_CONF_DIR if set, or within the current directory's +rem conf\ subdirectory. + +if not defined SPARK_ENV_LOADED ( + set SPARK_ENV_LOADED=1 + + if not defined SPARK_CONF_DIR ( + set SPARK_CONF_DIR=%~dp0..\conf + ) + + call :LoadSparkEnv +) + +rem Setting SPARK_SCALA_VERSION if not already set. + +set SCALA_VERSION_1=2.13 +set SCALA_VERSION_2=2.12 + +set ASSEMBLY_DIR1="%SPARK_HOME%\assembly\target\scala-%SCALA_VERSION_1%" +set ASSEMBLY_DIR2="%SPARK_HOME%\assembly\target\scala-%SCALA_VERSION_2%" +set ENV_VARIABLE_DOC=https://spark.apache.org/docs/latest/configuration.html#environment-variables + +if not defined SPARK_SCALA_VERSION ( + if exist %ASSEMBLY_DIR2% if exist %ASSEMBLY_DIR1% ( + echo Presence of build for multiple Scala versions detected ^(%ASSEMBLY_DIR1% and %ASSEMBLY_DIR2%^). + echo Remove one of them or, set SPARK_SCALA_VERSION=%SCALA_VERSION_1% in spark-env.cmd. + echo Visit %ENV_VARIABLE_DOC% for more details about setting environment variables in spark-env.cmd. + echo Either clean one of them or, set SPARK_SCALA_VERSION in spark-env.cmd. + exit 1 + ) + if exist %ASSEMBLY_DIR1% ( + set SPARK_SCALA_VERSION=%SCALA_VERSION_1% + ) else ( + set SPARK_SCALA_VERSION=%SCALA_VERSION_2% + ) +) +exit /b 0 + +:LoadSparkEnv +if exist "%SPARK_CONF_DIR%\spark-env.cmd" ( + call "%SPARK_CONF_DIR%\spark-env.cmd" +) diff --git a/spark-3.5.0-bin-hadoop3/bin/load-spark-env.sh b/spark-3.5.0-bin-hadoop3/bin/load-spark-env.sh new file mode 100644 index 0000000..fc5e881 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/load-spark-env.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# This script loads spark-env.sh if it exists, and ensures it is only loaded once. +# spark-env.sh is loaded from SPARK_CONF_DIR if set, or within the current directory's +# conf/ subdirectory. + +# Figure out where Spark is installed +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +SPARK_ENV_SH="spark-env.sh" +if [ -z "$SPARK_ENV_LOADED" ]; then + export SPARK_ENV_LOADED=1 + + export SPARK_CONF_DIR="${SPARK_CONF_DIR:-"${SPARK_HOME}"/conf}" + + SPARK_ENV_SH="${SPARK_CONF_DIR}/${SPARK_ENV_SH}" + if [[ -f "${SPARK_ENV_SH}" ]]; then + # Promote all variable declarations to environment (exported) variables + set -a + . ${SPARK_ENV_SH} + set +a + fi +fi + +# Setting SPARK_SCALA_VERSION if not already set. + +if [ -z "$SPARK_SCALA_VERSION" ]; then + SCALA_VERSION_1=2.13 + SCALA_VERSION_2=2.12 + + ASSEMBLY_DIR_1="${SPARK_HOME}/assembly/target/scala-${SCALA_VERSION_1}" + ASSEMBLY_DIR_2="${SPARK_HOME}/assembly/target/scala-${SCALA_VERSION_2}" + ENV_VARIABLE_DOC="https://spark.apache.org/docs/latest/configuration.html#environment-variables" + if [[ -d "$ASSEMBLY_DIR_1" && -d "$ASSEMBLY_DIR_2" ]]; then + echo "Presence of build for multiple Scala versions detected ($ASSEMBLY_DIR_1 and $ASSEMBLY_DIR_2)." 1>&2 + echo "Remove one of them or, export SPARK_SCALA_VERSION=$SCALA_VERSION_1 in ${SPARK_ENV_SH}." 1>&2 + echo "Visit ${ENV_VARIABLE_DOC} for more details about setting environment variables in spark-env.sh." 1>&2 + exit 1 + fi + + if [[ -d "$ASSEMBLY_DIR_1" ]]; then + export SPARK_SCALA_VERSION=${SCALA_VERSION_1} + else + export SPARK_SCALA_VERSION=${SCALA_VERSION_2} + fi +fi + +# Append jline option to enable the Beeline process to run in background. +if [[ ( ! $(ps -o stat= -p $$) =~ "+" ) && ! ( -p /dev/stdin ) ]]; then + export SPARK_BEELINE_OPTS="$SPARK_BEELINE_OPTS -Djline.terminal=jline.UnsupportedTerminal" +fi diff --git a/spark-3.5.0-bin-hadoop3/bin/pyspark b/spark-3.5.0-bin-hadoop3/bin/pyspark new file mode 100644 index 0000000..1ae28b1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/pyspark @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +source "${SPARK_HOME}"/bin/load-spark-env.sh +export _SPARK_CMD_USAGE="Usage: ./bin/pyspark [options]" + +# In Spark 2.0, IPYTHON and IPYTHON_OPTS are removed and pyspark fails to launch if either option +# is set in the user's environment. Instead, users should set PYSPARK_DRIVER_PYTHON=ipython +# to use IPython and set PYSPARK_DRIVER_PYTHON_OPTS to pass options when starting the Python driver +# (e.g. PYSPARK_DRIVER_PYTHON_OPTS='notebook'). This supports full customization of the IPython +# and executor Python executables. + +# Fail noisily if removed options are set +if [[ -n "$IPYTHON" || -n "$IPYTHON_OPTS" ]]; then + echo "Error in pyspark startup:" + echo "IPYTHON and IPYTHON_OPTS are removed in Spark 2.0+. Remove these from the environment and set PYSPARK_DRIVER_PYTHON and PYSPARK_DRIVER_PYTHON_OPTS instead." + exit 1 +fi + +# Default to standard python3 interpreter unless told otherwise +if [[ -z "$PYSPARK_PYTHON" ]]; then + PYSPARK_PYTHON=python3 +fi +if [[ -z "$PYSPARK_DRIVER_PYTHON" ]]; then + PYSPARK_DRIVER_PYTHON=$PYSPARK_PYTHON +fi +export PYSPARK_PYTHON +export PYSPARK_DRIVER_PYTHON +export PYSPARK_DRIVER_PYTHON_OPTS + +# Add the PySpark classes to the Python path: +export PYTHONPATH="${SPARK_HOME}/python/:$PYTHONPATH" +export PYTHONPATH="${SPARK_HOME}/python/lib/py4j-0.10.9.7-src.zip:$PYTHONPATH" + +# Load the PySpark shell.py script when ./pyspark is used interactively: +export OLD_PYTHONSTARTUP="$PYTHONSTARTUP" +export PYTHONSTARTUP="${SPARK_HOME}/python/pyspark/shell.py" + +# For pyspark tests +if [[ -n "$SPARK_TESTING" ]]; then + unset YARN_CONF_DIR + unset HADOOP_CONF_DIR + export PYTHONHASHSEED=0 + exec "$PYSPARK_DRIVER_PYTHON" -m "$@" + exit +fi + +exec "${SPARK_HOME}"/bin/spark-submit pyspark-shell-main --name "PySparkShell" "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/pyspark.cmd b/spark-3.5.0-bin-hadoop3/bin/pyspark.cmd new file mode 100644 index 0000000..3dcf1d4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/pyspark.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running PySpark. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0pyspark2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/pyspark2.cmd b/spark-3.5.0-bin-hadoop3/bin/pyspark2.cmd new file mode 100644 index 0000000..232813b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/pyspark2.cmd @@ -0,0 +1,38 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +call "%SPARK_HOME%\bin\load-spark-env.cmd" +set _SPARK_CMD_USAGE=Usage: bin\pyspark.cmd [options] + +rem Figure out which Python to use. +if "x%PYSPARK_DRIVER_PYTHON%"=="x" ( + set PYSPARK_DRIVER_PYTHON=python + if not [%PYSPARK_PYTHON%] == [] set PYSPARK_DRIVER_PYTHON=%PYSPARK_PYTHON% +) + +set PYTHONPATH=%SPARK_HOME%\python;%PYTHONPATH% +set PYTHONPATH=%SPARK_HOME%\python\lib\py4j-0.10.9.7-src.zip;%PYTHONPATH% + +set OLD_PYTHONSTARTUP=%PYTHONSTARTUP% +set PYTHONSTARTUP=%SPARK_HOME%\python\pyspark\shell.py + +call "%SPARK_HOME%\bin\spark-submit2.cmd" pyspark-shell-main --name "PySparkShell" %* diff --git a/spark-3.5.0-bin-hadoop3/bin/run-example b/spark-3.5.0-bin-hadoop3/bin/run-example new file mode 100644 index 0000000..4ba5399 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/run-example @@ -0,0 +1,25 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +export _SPARK_CMD_USAGE="Usage: ./bin/run-example [options] example-class [example args]" +exec "${SPARK_HOME}"/bin/spark-submit run-example "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/run-example.cmd b/spark-3.5.0-bin-hadoop3/bin/run-example.cmd new file mode 100644 index 0000000..2dd396e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/run-example.cmd @@ -0,0 +1,27 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +set _SPARK_CMD_USAGE=Usage: .\bin\run-example [options] example-class [example args] + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-submit.cmd" run-example %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-class b/spark-3.5.0-bin-hadoop3/bin/spark-class new file mode 100644 index 0000000..fc343ca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-class @@ -0,0 +1,112 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +. "${SPARK_HOME}"/bin/load-spark-env.sh + +# Find the java binary +if [ -n "${JAVA_HOME}" ]; then + RUNNER="${JAVA_HOME}/bin/java" +else + if [ "$(command -v java)" ]; then + RUNNER="java" + else + echo "JAVA_HOME is not set" >&2 + exit 1 + fi +fi + +# Find Spark jars. +if [ -d "${SPARK_HOME}/jars" ]; then + SPARK_JARS_DIR="${SPARK_HOME}/jars" +else + SPARK_JARS_DIR="${SPARK_HOME}/assembly/target/scala-$SPARK_SCALA_VERSION/jars" +fi + +if [ ! -d "$SPARK_JARS_DIR" ] && [ -z "$SPARK_TESTING$SPARK_SQL_TESTING" ]; then + echo "Failed to find Spark jars directory ($SPARK_JARS_DIR)." 1>&2 + echo "You need to build Spark with the target \"package\" before running this program." 1>&2 + exit 1 +else + LAUNCH_CLASSPATH="$SPARK_JARS_DIR/*" +fi + +# Add the launcher build dir to the classpath if requested. +if [ -n "$SPARK_PREPEND_CLASSES" ]; then + LAUNCH_CLASSPATH="${SPARK_HOME}/launcher/target/scala-$SPARK_SCALA_VERSION/classes:$LAUNCH_CLASSPATH" +fi + +# For tests +if [[ -n "$SPARK_TESTING" ]]; then + unset YARN_CONF_DIR + unset HADOOP_CONF_DIR +fi + +# The launcher library will print arguments separated by a NULL character, to allow arguments with +# characters that would be otherwise interpreted by the shell. Read that in a while loop, populating +# an array that will be used to exec the final command. +# +# The exit code of the launcher is appended to the output, so the parent shell removes it from the +# command array and checks the value to see if the launcher succeeded. +build_command() { + "$RUNNER" -Xmx128m $SPARK_LAUNCHER_OPTS -cp "$LAUNCH_CLASSPATH" org.apache.spark.launcher.Main "$@" + printf "%d\0" $? +} + +# Turn off posix mode since it does not allow process substitution +set +o posix +CMD=() +DELIM=$'\n' +CMD_START_FLAG="false" +while IFS= read -d "$DELIM" -r _ARG; do + ARG=${_ARG//$'\r'} + if [ "$CMD_START_FLAG" == "true" ]; then + CMD+=("$ARG") + else + if [ "$ARG" == $'\0' ]; then + # After NULL character is consumed, change the delimiter and consume command string. + DELIM='' + CMD_START_FLAG="true" + elif [ "$ARG" != "" ]; then + echo "$ARG" + fi + fi +done < <(build_command "$@") + +COUNT=${#CMD[@]} +LAST=$((COUNT - 1)) +LAUNCHER_EXIT_CODE=${CMD[$LAST]} + +# Certain JVM failures result in errors being printed to stdout (instead of stderr), which causes +# the code that parses the output of the launcher to get confused. In those cases, check if the +# exit code is an integer, and if it's not, handle it as a special error case. +if ! [[ $LAUNCHER_EXIT_CODE =~ ^[0-9]+$ ]]; then + echo "${CMD[@]}" | head -n-1 1>&2 + exit 1 +fi + +if [ $LAUNCHER_EXIT_CODE != 0 ]; then + exit $LAUNCHER_EXIT_CODE +fi + +CMD=("${CMD[@]:0:$LAST}") +exec "${CMD[@]}" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-class.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-class.cmd new file mode 100644 index 0000000..b22536a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-class.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running a Spark class. To avoid polluting +rem the environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-class2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-class2.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-class2.cmd new file mode 100644 index 0000000..800ec0c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-class2.cmd @@ -0,0 +1,79 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +call "%SPARK_HOME%\bin\load-spark-env.cmd" + +rem Test that an argument was given +if "x%1"=="x" ( + echo Usage: spark-class ^ [^] + exit /b 1 +) + +rem Find Spark jars. +if exist "%SPARK_HOME%\jars" ( + set SPARK_JARS_DIR=%SPARK_HOME%\jars +) else ( + set SPARK_JARS_DIR=%SPARK_HOME%\assembly\target\scala-%SPARK_SCALA_VERSION%\jars +) + +if not exist "%SPARK_JARS_DIR%" ( + echo Failed to find Spark jars directory. + echo You need to build Spark before running this program. + exit /b 1 +) + +set LAUNCH_CLASSPATH=%SPARK_JARS_DIR%\* + +rem Add the launcher build dir to the classpath if requested. +if not "x%SPARK_PREPEND_CLASSES%"=="x" ( + set LAUNCH_CLASSPATH="%SPARK_HOME%\launcher\target\scala-%SPARK_SCALA_VERSION%\classes;%LAUNCH_CLASSPATH%" +) + +rem Figure out where java is. +set RUNNER=java +if not "x%JAVA_HOME%"=="x" ( + set RUNNER=%JAVA_HOME%\bin\java +) else ( + where /q "%RUNNER%" + if ERRORLEVEL 1 ( + echo Java not found and JAVA_HOME environment variable is not set. + echo Install Java and set JAVA_HOME to point to the Java installation directory. + exit /b 1 + ) +) + +rem The launcher library prints the command to be executed in a single line suitable for being +rem executed by the batch interpreter. So read all the output of the launcher into a variable. +:gen +set LAUNCHER_OUTPUT=%temp%\spark-class-launcher-output-%RANDOM%.txt +rem SPARK-28302: %RANDOM% would return the same number if we call it instantly after last call, +rem so we should make it sure to generate unique file to avoid process collision of writing into +rem the same file concurrently. +if exist %LAUNCHER_OUTPUT% goto :gen +rem unset SHELL to indicate non-bash environment to launcher/Main +set SHELL= +"%RUNNER%" -Xmx128m -cp "%LAUNCH_CLASSPATH%" org.apache.spark.launcher.Main %* > %LAUNCHER_OUTPUT% +for /f "tokens=*" %%i in (%LAUNCHER_OUTPUT%) do ( + set SPARK_CMD=%%i +) +del %LAUNCHER_OUTPUT% +%SPARK_CMD% diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-connect-shell b/spark-3.5.0-bin-hadoop3/bin/spark-connect-shell new file mode 100644 index 0000000..9026c81 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-connect-shell @@ -0,0 +1,27 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# The shell script to start a spark-shell with spark connect enabled. + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +# This requires building the spark with `-Pconnect`, e,g, `build/sbt -Pconnect package` +exec "${SPARK_HOME}"/bin/spark-shell --conf spark.plugins=org.apache.spark.sql.connect.SparkConnectPlugin "$@" \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-shell b/spark-3.5.0-bin-hadoop3/bin/spark-shell new file mode 100644 index 0000000..e920137 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-shell @@ -0,0 +1,98 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# +# Shell script for starting the Spark Shell REPL + +cygwin=false +case "$(uname)" in + CYGWIN*) cygwin=true;; +esac + +# Enter posix mode for bash +set -o posix + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +export _SPARK_CMD_USAGE="Usage: ./bin/spark-shell [options] + +Scala REPL options: + -I preload , enforcing line-by-line interpretation" + +# SPARK-4161: scala does not assume use of the java classpath, +# so we need to add the "-Dscala.usejavacp=true" flag manually. We +# do this specifically for the Spark shell because the scala REPL +# has its own class loader, and any additional classpath specified +# through spark.driver.extraClassPath is not automatically propagated. +SPARK_SUBMIT_OPTS="$SPARK_SUBMIT_OPTS -Dscala.usejavacp=true" + +function main() { + if $cygwin; then + # Workaround for issue involving JLine and Cygwin + # (see http://sourceforge.net/p/jline/bugs/40/). + # If you're using the Mintty terminal emulator in Cygwin, may need to set the + # "Backspace sends ^H" setting in "Keys" section of the Mintty options + # (see https://github.com/sbt/sbt/issues/562). + stty -icanon min 1 -echo > /dev/null 2>&1 + export SPARK_SUBMIT_OPTS="$SPARK_SUBMIT_OPTS -Djline.terminal=unix" + "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@" + stty icanon echo > /dev/null 2>&1 + else + export SPARK_SUBMIT_OPTS + "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@" + fi +} + +# Copy restore-TTY-on-exit functions from Scala script so spark-shell exits properly even in +# binary distribution of Spark where Scala is not installed +exit_status=127 +saved_stty="" + +# restore stty settings (echo in particular) +function restoreSttySettings() { + stty $saved_stty + saved_stty="" +} + +function onExit() { + if [[ "$saved_stty" != "" ]]; then + restoreSttySettings + fi + exit $exit_status +} + +# to reenable echo if we are interrupted before completing. +trap onExit INT + +# save terminal settings +saved_stty=$(stty -g 2>/dev/null) +# clear on error so we don't later try to restore them +if [[ ! $? ]]; then + saved_stty="" +fi + +main "$@" + +# record the exit status lest it be overwritten: +# then reenable echo and propagate the code. +exit_status=$? +onExit + diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-shell.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-shell.cmd new file mode 100644 index 0000000..e734f13 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-shell.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running Spark shell. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-shell2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-shell2.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-shell2.cmd new file mode 100644 index 0000000..549bf43 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-shell2.cmd @@ -0,0 +1,43 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +set LF=^ + + +rem two empty lines are required +set _SPARK_CMD_USAGE=Usage: .\bin\spark-shell.cmd [options]^%LF%%LF%^%LF%%LF%^ +Scala REPL options:^%LF%%LF%^ + -I ^ preload ^, enforcing line-by-line interpretation + +rem SPARK-4161: scala does not assume use of the java classpath, +rem so we need to add the "-Dscala.usejavacp=true" flag manually. We +rem do this specifically for the Spark shell because the scala REPL +rem has its own class loader, and any additional classpath specified +rem through spark.driver.extraClassPath is not automatically propagated. +if "x%SPARK_SUBMIT_OPTS%"=="x" ( + set SPARK_SUBMIT_OPTS=-Dscala.usejavacp=true + goto run_shell +) +set SPARK_SUBMIT_OPTS="%SPARK_SUBMIT_OPTS% -Dscala.usejavacp=true" + +:run_shell +"%SPARK_HOME%\bin\spark-submit2.cmd" --class org.apache.spark.repl.Main --name "Spark shell" %* diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-sql b/spark-3.5.0-bin-hadoop3/bin/spark-sql new file mode 100644 index 0000000..b08b944 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-sql @@ -0,0 +1,25 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +export _SPARK_CMD_USAGE="Usage: ./bin/spark-sql [options] [cli option]" +exec "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.sql.hive.thriftserver.SparkSQLCLIDriver "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-sql.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-sql.cmd new file mode 100644 index 0000000..919e321 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-sql.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running SparkSQL. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-sql2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-sql2.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-sql2.cmd new file mode 100644 index 0000000..c34a3c5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-sql2.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +set _SPARK_CMD_USAGE=Usage: .\bin\spark-sql [options] [cli option] + +call "%SPARK_HOME%\bin\spark-submit2.cmd" --class org.apache.spark.sql.hive.thriftserver.SparkSQLCLIDriver %* diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-submit b/spark-3.5.0-bin-hadoop3/bin/spark-submit new file mode 100644 index 0000000..4e9d361 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-submit @@ -0,0 +1,27 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +# disable randomized hash for string in Python 3.3+ +export PYTHONHASHSEED=0 + +exec "${SPARK_HOME}"/bin/spark-class org.apache.spark.deploy.SparkSubmit "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-submit.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-submit.cmd new file mode 100644 index 0000000..da62a87 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-submit.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running Spark submit. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0spark-submit2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/spark-submit2.cmd b/spark-3.5.0-bin-hadoop3/bin/spark-submit2.cmd new file mode 100644 index 0000000..49e350f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/spark-submit2.cmd @@ -0,0 +1,27 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running Spark submit. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem disable randomized hash for string in Python 3.3+ +set PYTHONHASHSEED=0 + +set CLASS=org.apache.spark.deploy.SparkSubmit +"%~dp0spark-class2.cmd" %CLASS% %* diff --git a/spark-3.5.0-bin-hadoop3/bin/sparkR b/spark-3.5.0-bin-hadoop3/bin/sparkR new file mode 100644 index 0000000..8ecc755 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/sparkR @@ -0,0 +1,26 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + source "$(dirname "$0")"/find-spark-home +fi + +source "${SPARK_HOME}"/bin/load-spark-env.sh +export _SPARK_CMD_USAGE="Usage: ./bin/sparkR [options]" +exec "${SPARK_HOME}"/bin/spark-submit sparkr-shell-main "$@" diff --git a/spark-3.5.0-bin-hadoop3/bin/sparkR.cmd b/spark-3.5.0-bin-hadoop3/bin/sparkR.cmd new file mode 100644 index 0000000..fcd172b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/sparkR.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running SparkR. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +rem The outermost quotes are used to prevent Windows command line parse error +rem when there are some quotes in parameters, see SPARK-21877. +cmd /V /E /C ""%~dp0sparkR2.cmd" %*" diff --git a/spark-3.5.0-bin-hadoop3/bin/sparkR2.cmd b/spark-3.5.0-bin-hadoop3/bin/sparkR2.cmd new file mode 100644 index 0000000..446f0c3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/bin/sparkR2.cmd @@ -0,0 +1,25 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Figure out where the Spark framework is installed +call "%~dp0find-spark-home.cmd" + +call "%SPARK_HOME%\bin\load-spark-env.cmd" +set _SPARK_CMD_USAGE=Usage: .\bin\sparkR [options] +call "%SPARK_HOME%\bin\spark-submit2.cmd" sparkr-shell-main %* diff --git a/spark-3.5.0-bin-hadoop3/conf/fairscheduler.xml.template b/spark-3.5.0-bin-hadoop3/conf/fairscheduler.xml.template new file mode 100644 index 0000000..385b2e7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/fairscheduler.xml.template @@ -0,0 +1,31 @@ + + + + + + + FAIR + 1 + 2 + + + FIFO + 2 + 3 + + diff --git a/spark-3.5.0-bin-hadoop3/conf/log4j2.properties.template b/spark-3.5.0-bin-hadoop3/conf/log4j2.properties.template new file mode 100644 index 0000000..ab96e03 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/log4j2.properties.template @@ -0,0 +1,69 @@ +# +# 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. +# + +# Set everything to be logged to the console +rootLogger.level = info +rootLogger.appenderRef.stdout.ref = console + +# In the pattern layout configuration below, we specify an explicit `%ex` conversion +# pattern for logging Throwables. If this was omitted, then (by default) Log4J would +# implicitly add an `%xEx` conversion pattern which logs stacktraces with additional +# class packaging information. That extra information can sometimes add a substantial +# performance overhead, so we disable it in our default logging config. +# For more information, see SPARK-39361. +appender.console.type = Console +appender.console.name = console +appender.console.target = SYSTEM_ERR +appender.console.layout.type = PatternLayout +appender.console.layout.pattern = %d{yy/MM/dd HH:mm:ss} %p %c{1}: %m%n%ex + +# Set the default spark-shell/spark-sql log level to WARN. When running the +# spark-shell/spark-sql, the log level for these classes is used to overwrite +# the root logger's log level, so that the user can have different defaults +# for the shell and regular Spark apps. +logger.repl.name = org.apache.spark.repl.Main +logger.repl.level = warn + +logger.thriftserver.name = org.apache.spark.sql.hive.thriftserver.SparkSQLCLIDriver +logger.thriftserver.level = warn + +# Settings to quiet third party logs that are too verbose +logger.jetty1.name = org.sparkproject.jetty +logger.jetty1.level = warn +logger.jetty2.name = org.sparkproject.jetty.util.component.AbstractLifeCycle +logger.jetty2.level = error +logger.replexprTyper.name = org.apache.spark.repl.SparkIMain$exprTyper +logger.replexprTyper.level = info +logger.replSparkILoopInterpreter.name = org.apache.spark.repl.SparkILoop$SparkILoopInterpreter +logger.replSparkILoopInterpreter.level = info +logger.parquet1.name = org.apache.parquet +logger.parquet1.level = error +logger.parquet2.name = parquet +logger.parquet2.level = error + +# SPARK-9183: Settings to avoid annoying messages when looking up nonexistent UDFs in SparkSQL with Hive support +logger.RetryingHMSHandler.name = org.apache.hadoop.hive.metastore.RetryingHMSHandler +logger.RetryingHMSHandler.level = fatal +logger.FunctionRegistry.name = org.apache.hadoop.hive.ql.exec.FunctionRegistry +logger.FunctionRegistry.level = error + +# For deploying Spark ThriftServer +# SPARK-34128: Suppress undesirable TTransportException warnings involved in THRIFT-4805 +appender.console.filter.1.type = RegexFilter +appender.console.filter.1.regex = .*Thrift error occurred during processing of message.* +appender.console.filter.1.onMatch = deny +appender.console.filter.1.onMismatch = neutral diff --git a/spark-3.5.0-bin-hadoop3/conf/metrics.properties.template b/spark-3.5.0-bin-hadoop3/conf/metrics.properties.template new file mode 100644 index 0000000..f52d33f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/metrics.properties.template @@ -0,0 +1,210 @@ +# +# 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. +# + +# syntax: [instance].sink|source.[name].[options]=[value] + +# This file configures Spark's internal metrics system. The metrics system is +# divided into instances which correspond to internal components. +# Each instance can be configured to report its metrics to one or more sinks. +# Accepted values for [instance] are "master", "worker", "executor", "driver", +# and "applications". A wildcard "*" can be used as an instance name, in +# which case all instances will inherit the supplied property. +# +# Within an instance, a "source" specifies a particular set of grouped metrics. +# there are two kinds of sources: +# 1. Spark internal sources, like MasterSource, WorkerSource, etc, which will +# collect a Spark component's internal state. Each instance is paired with a +# Spark source that is added automatically. +# 2. Common sources, like JvmSource, which will collect low level state. +# These can be added through configuration options and are then loaded +# using reflection. +# +# A "sink" specifies where metrics are delivered to. Each instance can be +# assigned one or more sinks. +# +# The sink|source field specifies whether the property relates to a sink or +# source. +# +# The [name] field specifies the name of source or sink. +# +# The [options] field is the specific property of this source or sink. The +# source or sink is responsible for parsing this property. +# +# Notes: +# 1. To add a new sink, set the "class" option to a fully qualified class +# name (see examples below). +# 2. Some sinks involve a polling period. The minimum allowed polling period +# is 1 second. +# 3. Wildcard properties can be overridden by more specific properties. +# For example, master.sink.console.period takes precedence over +# *.sink.console.period. +# 4. A metrics specific configuration +# "spark.metrics.conf=${SPARK_HOME}/conf/metrics.properties" should be +# added to Java properties using -Dspark.metrics.conf=xxx if you want to +# customize metrics system. You can also put the file in ${SPARK_HOME}/conf +# and it will be loaded automatically. +# 5. The MetricsServlet sink is added by default as a sink in the master, +# worker and driver, and you can send HTTP requests to the "/metrics/json" +# endpoint to get a snapshot of all the registered metrics in JSON format. +# For master, requests to the "/metrics/master/json" and +# "/metrics/applications/json" endpoints can be sent separately to get +# metrics snapshots of the master instance and applications. This +# MetricsServlet does not have to be configured. +# 6. The metrics system can also be configured using Spark configuration +# parameters. The relevant parameter names are formed by adding the +# prefix "spark.metrics.conf." to the configuration entries detailed in +# this file (see examples below). + +## List of available common sources and their properties. + +# org.apache.spark.metrics.source.JvmSource +# Note: Currently, JvmSource is the only available common source. +# It can be added to an instance by setting the "class" option to its +# fully qualified class name (see examples below). + +## List of available sinks and their properties. + +# org.apache.spark.metrics.sink.ConsoleSink +# Name: Default: Description: +# period 10 Poll period +# unit seconds Unit of the poll period + +# org.apache.spark.metrics.sink.CSVSink +# Name: Default: Description: +# period 10 Poll period +# unit seconds Unit of the poll period +# directory /tmp Where to store CSV files + +# org.apache.spark.metrics.sink.GangliaSink +# Name: Default: Description: +# host NONE Hostname or multicast group of the Ganglia server, +# must be set +# port NONE Port of the Ganglia server(s), must be set +# period 10 Poll period +# unit seconds Unit of the poll period +# ttl 1 TTL of messages sent by Ganglia +# dmax 0 Lifetime in seconds of metrics (0 never expired) +# mode multicast Ganglia network mode ('unicast' or 'multicast') + +# org.apache.spark.metrics.sink.JmxSink + +# org.apache.spark.metrics.sink.MetricsServlet +# Name: Default: Description: +# path VARIES* Path prefix from the web server root +# sample false Whether to show entire set of samples for histograms +# ('false' or 'true') +# +# * Default path is /metrics/json for all instances except the master. The +# master has two paths: +# /metrics/applications/json # App information +# /metrics/master/json # Master information + +# org.apache.spark.metrics.sink.PrometheusServlet +# Name: Default: Description: +# path VARIES* Path prefix from the web server root +# +# * Default path is /metrics/prometheus for all instances except the master. The +# master has two paths: +# /metrics/applications/prometheus # App information +# /metrics/master/prometheus # Master information + +# org.apache.spark.metrics.sink.GraphiteSink +# Name: Default: Description: +# host NONE Hostname of the Graphite server, must be set +# port NONE Port of the Graphite server, must be set +# period 10 Poll period +# unit seconds Unit of the poll period +# prefix EMPTY STRING Prefix to prepend to every metric's name +# protocol tcp Protocol ("tcp" or "udp") to use +# regex NONE Optional filter to send only metrics matching this regex string + +# org.apache.spark.metrics.sink.StatsdSink +# Name: Default: Description: +# host 127.0.0.1 Hostname or IP of StatsD server +# port 8125 Port of StatsD server +# period 10 Poll period +# unit seconds Units of poll period +# prefix EMPTY STRING Prefix to prepend to metric name + +## Examples +# Enable JmxSink for all instances by class name +#*.sink.jmx.class=org.apache.spark.metrics.sink.JmxSink + +# Enable ConsoleSink for all instances by class name +#*.sink.console.class=org.apache.spark.metrics.sink.ConsoleSink + +# Enable StatsdSink for all instances by class name +#*.sink.statsd.class=org.apache.spark.metrics.sink.StatsdSink +#*.sink.statsd.prefix=spark + +# Polling period for the ConsoleSink +#*.sink.console.period=10 +# Unit of the polling period for the ConsoleSink +#*.sink.console.unit=seconds + +# Polling period for the ConsoleSink specific for the master instance +#master.sink.console.period=15 +# Unit of the polling period for the ConsoleSink specific for the master +# instance +#master.sink.console.unit=seconds + +# Enable CsvSink for all instances by class name +#*.sink.csv.class=org.apache.spark.metrics.sink.CsvSink + +# Polling period for the CsvSink +#*.sink.csv.period=1 +# Unit of the polling period for the CsvSink +#*.sink.csv.unit=minutes + +# Polling directory for CsvSink +#*.sink.csv.directory=/tmp/ + +# Polling period for the CsvSink specific for the worker instance +#worker.sink.csv.period=10 +# Unit of the polling period for the CsvSink specific for the worker instance +#worker.sink.csv.unit=minutes + +# Enable Slf4jSink for all instances by class name +#*.sink.slf4j.class=org.apache.spark.metrics.sink.Slf4jSink + +# Polling period for the Slf4JSink +#*.sink.slf4j.period=1 +# Unit of the polling period for the Slf4jSink +#*.sink.slf4j.unit=minutes + +# Example configuration for Graphite sink +#*.sink.graphite.class=org.apache.spark.metrics.sink.GraphiteSink +#*.sink.graphite.host= +#*.sink.graphite.port= +#*.sink.graphite.period=10 +#*.sink.graphite.unit=seconds +#*.sink.graphite.prefix= + +# Enable JvmSource for instance master, worker, driver and executor +#master.source.jvm.class=org.apache.spark.metrics.source.JvmSource + +#worker.source.jvm.class=org.apache.spark.metrics.source.JvmSource + +#driver.source.jvm.class=org.apache.spark.metrics.source.JvmSource + +#executor.source.jvm.class=org.apache.spark.metrics.source.JvmSource + +# Example configuration for PrometheusServlet +#*.sink.prometheusServlet.class=org.apache.spark.metrics.sink.PrometheusServlet +#*.sink.prometheusServlet.path=/metrics/prometheus +#master.sink.prometheusServlet.path=/metrics/master/prometheus +#applications.sink.prometheusServlet.path=/metrics/applications/prometheus diff --git a/spark-3.5.0-bin-hadoop3/conf/spark-defaults.conf.template b/spark-3.5.0-bin-hadoop3/conf/spark-defaults.conf.template new file mode 100644 index 0000000..19cba6e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/spark-defaults.conf.template @@ -0,0 +1,27 @@ +# +# 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. +# + +# Default system properties included when running spark-submit. +# This is useful for setting default environmental settings. + +# Example: +# spark.master spark://master:7077 +# spark.eventLog.enabled true +# spark.eventLog.dir hdfs://namenode:8021/directory +# spark.serializer org.apache.spark.serializer.KryoSerializer +# spark.driver.memory 5g +# spark.executor.extraJavaOptions -XX:+PrintGCDetails -Dkey=value -Dnumbers="one two three" diff --git a/spark-3.5.0-bin-hadoop3/conf/spark-env.sh.template b/spark-3.5.0-bin-hadoop3/conf/spark-env.sh.template new file mode 100644 index 0000000..e949199 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/spark-env.sh.template @@ -0,0 +1,81 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# This file is sourced when running various Spark programs. +# Copy it as spark-env.sh and edit that to configure Spark for your site. + +# Options read when launching programs locally with +# ./bin/run-example or ./bin/spark-submit +# - HADOOP_CONF_DIR, to point Spark towards Hadoop configuration files +# - SPARK_LOCAL_IP, to set the IP address Spark binds to on this node +# - SPARK_PUBLIC_DNS, to set the public dns name of the driver program + +# Options read by executors and drivers running inside the cluster +# - SPARK_LOCAL_IP, to set the IP address Spark binds to on this node +# - SPARK_PUBLIC_DNS, to set the public DNS name of the driver program +# - SPARK_LOCAL_DIRS, storage directories to use on this node for shuffle and RDD data +# - MESOS_NATIVE_JAVA_LIBRARY, to point to your libmesos.so if you use Mesos + +# Options read in any mode +# - SPARK_CONF_DIR, Alternate conf dir. (Default: ${SPARK_HOME}/conf) +# - SPARK_EXECUTOR_CORES, Number of cores for the executors (Default: 1). +# - SPARK_EXECUTOR_MEMORY, Memory per Executor (e.g. 1000M, 2G) (Default: 1G) +# - SPARK_DRIVER_MEMORY, Memory for Driver (e.g. 1000M, 2G) (Default: 1G) + +# Options read in any cluster manager using HDFS +# - HADOOP_CONF_DIR, to point Spark towards Hadoop configuration files + +# Options read in YARN client/cluster mode +# - YARN_CONF_DIR, to point Spark towards YARN configuration files when you use YARN + +# Options for the daemons used in the standalone deploy mode +# - SPARK_MASTER_HOST, to bind the master to a different IP address or hostname +# - SPARK_MASTER_PORT / SPARK_MASTER_WEBUI_PORT, to use non-default ports for the master +# - SPARK_MASTER_OPTS, to set config properties only for the master (e.g. "-Dx=y") +# - SPARK_WORKER_CORES, to set the number of cores to use on this machine +# - SPARK_WORKER_MEMORY, to set how much total memory workers have to give executors (e.g. 1000m, 2g) +# - SPARK_WORKER_PORT / SPARK_WORKER_WEBUI_PORT, to use non-default ports for the worker +# - SPARK_WORKER_DIR, to set the working directory of worker processes +# - SPARK_WORKER_OPTS, to set config properties only for the worker (e.g. "-Dx=y") +# - SPARK_DAEMON_MEMORY, to allocate to the master, worker and history server themselves (default: 1g). +# - SPARK_HISTORY_OPTS, to set config properties only for the history server (e.g. "-Dx=y") +# - SPARK_SHUFFLE_OPTS, to set config properties only for the external shuffle service (e.g. "-Dx=y") +# - SPARK_DAEMON_JAVA_OPTS, to set config properties for all daemons (e.g. "-Dx=y") +# - SPARK_DAEMON_CLASSPATH, to set the classpath for all daemons +# - SPARK_PUBLIC_DNS, to set the public dns name of the master or workers + +# Options for launcher +# - SPARK_LAUNCHER_OPTS, to set config properties and Java options for the launcher (e.g. "-Dx=y") + +# Generic options for the daemons used in the standalone deploy mode +# - SPARK_CONF_DIR Alternate conf dir. (Default: ${SPARK_HOME}/conf) +# - SPARK_LOG_DIR Where log files are stored. (Default: ${SPARK_HOME}/logs) +# - SPARK_LOG_MAX_FILES Max log files of Spark daemons can rotate to. Default is 5. +# - SPARK_PID_DIR Where the pid file is stored. (Default: /tmp) +# - SPARK_IDENT_STRING A string representing this instance of spark. (Default: $USER) +# - SPARK_NICENESS The scheduling priority for daemons. (Default: 0) +# - SPARK_NO_DAEMONIZE Run the proposed command in the foreground. It will not output a PID file. +# Options for native BLAS, like Intel MKL, OpenBLAS, and so on. +# You might get better performance to enable these options if using native BLAS (see SPARK-21305). +# - MKL_NUM_THREADS=1 Disable multi-threading of Intel MKL +# - OPENBLAS_NUM_THREADS=1 Disable multi-threading of OpenBLAS + +# Options for beeline +# - SPARK_BEELINE_OPTS, to set config properties only for the beeline cli (e.g. "-Dx=y") +# - SPARK_BEELINE_MEMORY, Memory for beeline (e.g. 1000M, 2G) (Default: 1G) diff --git a/spark-3.5.0-bin-hadoop3/conf/workers.template b/spark-3.5.0-bin-hadoop3/conf/workers.template new file mode 100644 index 0000000..be42a63 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/conf/workers.template @@ -0,0 +1,19 @@ +# +# 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. +# + +# A Spark Worker will be started on each of the machines listed below. +localhost \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/README.md b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/README.md new file mode 100644 index 0000000..df9af41 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/README.md @@ -0,0 +1,5 @@ +The CRCs for a specific file are stored in a text file with the same name (excluding the original extension). + +The CRCs are calculated for data chunks of `32768 bytes` (individual CRCs) and are newline delimited. + +The CRCs were calculated using https://simplycalc.com/crc32-file.php \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/junitLargeJar.txt b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/junitLargeJar.txt new file mode 100644 index 0000000..3e89631 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/junitLargeJar.txt @@ -0,0 +1,12 @@ +902183889 +2415704507 +1084811487 +1951510 +1158852476 +2003120166 +3026803842 +3850244775 +3409267044 +652109216 +104029242 +3019434266 \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/smallJar.txt b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/smallJar.txt new file mode 100644 index 0000000..df32adc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/artifact-tests/crc/smallJar.txt @@ -0,0 +1 @@ +1631702900 \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/data/artifact-tests/junitLargeJar.jar b/spark-3.5.0-bin-hadoop3/data/artifact-tests/junitLargeJar.jar new file mode 100644 index 0000000..6da55d8 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/artifact-tests/junitLargeJar.jar differ diff --git a/spark-3.5.0-bin-hadoop3/data/artifact-tests/smallJar.jar b/spark-3.5.0-bin-hadoop3/data/artifact-tests/smallJar.jar new file mode 100644 index 0000000..3c4930e Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/artifact-tests/smallJar.jar differ diff --git a/spark-3.5.0-bin-hadoop3/data/graphx/followers.txt b/spark-3.5.0-bin-hadoop3/data/graphx/followers.txt new file mode 100644 index 0000000..7bb8e90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/graphx/followers.txt @@ -0,0 +1,8 @@ +2 1 +4 1 +1 2 +6 3 +7 3 +7 6 +6 7 +3 7 diff --git a/spark-3.5.0-bin-hadoop3/data/graphx/users.txt b/spark-3.5.0-bin-hadoop3/data/graphx/users.txt new file mode 100644 index 0000000..982d19d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/graphx/users.txt @@ -0,0 +1,7 @@ +1,BarackObama,Barack Obama +2,ladygaga,Goddess of Love +3,jeresig,John Resig +4,justinbieber,Justin Bieber +6,matei_zaharia,Matei Zaharia +7,odersky,Martin Odersky +8,anonsys diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/als/sample_movielens_ratings.txt b/spark-3.5.0-bin-hadoop3/data/mllib/als/sample_movielens_ratings.txt new file mode 100644 index 0000000..0889142 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/als/sample_movielens_ratings.txt @@ -0,0 +1,1501 @@ +0::2::3::1424380312 +0::3::1::1424380312 +0::5::2::1424380312 +0::9::4::1424380312 +0::11::1::1424380312 +0::12::2::1424380312 +0::15::1::1424380312 +0::17::1::1424380312 +0::19::1::1424380312 +0::21::1::1424380312 +0::23::1::1424380312 +0::26::3::1424380312 +0::27::1::1424380312 +0::28::1::1424380312 +0::29::1::1424380312 +0::30::1::1424380312 +0::31::1::1424380312 +0::34::1::1424380312 +0::37::1::1424380312 +0::41::2::1424380312 +0::44::1::1424380312 +0::45::2::1424380312 +0::46::1::1424380312 +0::47::1::1424380312 +0::48::1::1424380312 +0::50::1::1424380312 +0::51::1::1424380312 +0::54::1::1424380312 +0::55::1::1424380312 +0::59::2::1424380312 +0::61::2::1424380312 +0::64::1::1424380312 +0::67::1::1424380312 +0::68::1::1424380312 +0::69::1::1424380312 +0::71::1::1424380312 +0::72::1::1424380312 +0::77::2::1424380312 +0::79::1::1424380312 +0::83::1::1424380312 +0::87::1::1424380312 +0::89::2::1424380312 +0::91::3::1424380312 +0::92::4::1424380312 +0::94::1::1424380312 +0::95::2::1424380312 +0::96::1::1424380312 +0::98::1::1424380312 +0::99::1::1424380312 +1::2::2::1424380312 +1::3::1::1424380312 +1::4::2::1424380312 +1::6::1::1424380312 +1::9::3::1424380312 +1::12::1::1424380312 +1::13::1::1424380312 +1::14::1::1424380312 +1::16::1::1424380312 +1::19::1::1424380312 +1::21::3::1424380312 +1::27::1::1424380312 +1::28::3::1424380312 +1::33::1::1424380312 +1::36::2::1424380312 +1::37::1::1424380312 +1::40::1::1424380312 +1::41::2::1424380312 +1::43::1::1424380312 +1::44::1::1424380312 +1::47::1::1424380312 +1::50::1::1424380312 +1::54::1::1424380312 +1::56::2::1424380312 +1::57::1::1424380312 +1::58::1::1424380312 +1::60::1::1424380312 +1::62::4::1424380312 +1::63::1::1424380312 +1::67::1::1424380312 +1::68::4::1424380312 +1::70::2::1424380312 +1::72::1::1424380312 +1::73::1::1424380312 +1::74::2::1424380312 +1::76::1::1424380312 +1::77::3::1424380312 +1::78::1::1424380312 +1::81::1::1424380312 +1::82::1::1424380312 +1::85::3::1424380312 +1::86::2::1424380312 +1::88::2::1424380312 +1::91::1::1424380312 +1::92::2::1424380312 +1::93::1::1424380312 +1::94::2::1424380312 +1::96::1::1424380312 +1::97::1::1424380312 +2::4::3::1424380312 +2::6::1::1424380312 +2::8::5::1424380312 +2::9::1::1424380312 +2::10::1::1424380312 +2::12::3::1424380312 +2::13::1::1424380312 +2::15::2::1424380312 +2::18::2::1424380312 +2::19::4::1424380312 +2::22::1::1424380312 +2::26::1::1424380312 +2::28::1::1424380312 +2::34::4::1424380312 +2::35::1::1424380312 +2::37::5::1424380312 +2::38::1::1424380312 +2::39::5::1424380312 +2::40::4::1424380312 +2::47::1::1424380312 +2::50::1::1424380312 +2::52::2::1424380312 +2::54::1::1424380312 +2::55::1::1424380312 +2::57::2::1424380312 +2::58::2::1424380312 +2::59::1::1424380312 +2::61::1::1424380312 +2::62::1::1424380312 +2::64::1::1424380312 +2::65::1::1424380312 +2::66::3::1424380312 +2::68::1::1424380312 +2::71::3::1424380312 +2::76::1::1424380312 +2::77::1::1424380312 +2::78::1::1424380312 +2::80::1::1424380312 +2::83::5::1424380312 +2::85::1::1424380312 +2::87::2::1424380312 +2::88::1::1424380312 +2::89::4::1424380312 +2::90::1::1424380312 +2::92::4::1424380312 +2::93::5::1424380312 +3::0::1::1424380312 +3::1::1::1424380312 +3::2::1::1424380312 +3::7::3::1424380312 +3::8::3::1424380312 +3::9::1::1424380312 +3::14::1::1424380312 +3::15::1::1424380312 +3::16::1::1424380312 +3::18::4::1424380312 +3::19::1::1424380312 +3::24::3::1424380312 +3::26::1::1424380312 +3::29::3::1424380312 +3::33::1::1424380312 +3::34::3::1424380312 +3::35::1::1424380312 +3::36::3::1424380312 +3::37::1::1424380312 +3::38::2::1424380312 +3::43::1::1424380312 +3::44::1::1424380312 +3::46::1::1424380312 +3::47::1::1424380312 +3::51::5::1424380312 +3::52::3::1424380312 +3::56::1::1424380312 +3::58::1::1424380312 +3::60::3::1424380312 +3::62::1::1424380312 +3::65::2::1424380312 +3::66::1::1424380312 +3::67::1::1424380312 +3::68::2::1424380312 +3::70::1::1424380312 +3::72::2::1424380312 +3::76::3::1424380312 +3::79::3::1424380312 +3::80::4::1424380312 +3::81::1::1424380312 +3::83::1::1424380312 +3::84::1::1424380312 +3::86::1::1424380312 +3::87::2::1424380312 +3::88::4::1424380312 +3::89::1::1424380312 +3::91::1::1424380312 +3::94::3::1424380312 +4::1::1::1424380312 +4::6::1::1424380312 +4::8::1::1424380312 +4::9::1::1424380312 +4::10::1::1424380312 +4::11::1::1424380312 +4::12::1::1424380312 +4::13::1::1424380312 +4::14::2::1424380312 +4::15::1::1424380312 +4::17::1::1424380312 +4::20::1::1424380312 +4::22::1::1424380312 +4::23::1::1424380312 +4::24::1::1424380312 +4::29::4::1424380312 +4::30::1::1424380312 +4::31::1::1424380312 +4::34::1::1424380312 +4::35::1::1424380312 +4::36::1::1424380312 +4::39::2::1424380312 +4::40::3::1424380312 +4::41::4::1424380312 +4::43::2::1424380312 +4::44::1::1424380312 +4::45::1::1424380312 +4::46::1::1424380312 +4::47::1::1424380312 +4::49::2::1424380312 +4::50::1::1424380312 +4::51::1::1424380312 +4::52::4::1424380312 +4::54::1::1424380312 +4::55::1::1424380312 +4::60::3::1424380312 +4::61::1::1424380312 +4::62::4::1424380312 +4::63::3::1424380312 +4::65::1::1424380312 +4::67::2::1424380312 +4::69::1::1424380312 +4::70::4::1424380312 +4::71::1::1424380312 +4::73::1::1424380312 +4::78::1::1424380312 +4::84::1::1424380312 +4::85::1::1424380312 +4::87::3::1424380312 +4::88::3::1424380312 +4::89::2::1424380312 +4::96::1::1424380312 +4::97::1::1424380312 +4::98::1::1424380312 +4::99::1::1424380312 +5::0::1::1424380312 +5::1::1::1424380312 +5::4::1::1424380312 +5::5::1::1424380312 +5::8::1::1424380312 +5::9::3::1424380312 +5::10::2::1424380312 +5::13::3::1424380312 +5::15::1::1424380312 +5::19::1::1424380312 +5::20::3::1424380312 +5::21::2::1424380312 +5::23::3::1424380312 +5::27::1::1424380312 +5::28::1::1424380312 +5::29::1::1424380312 +5::31::1::1424380312 +5::36::3::1424380312 +5::38::2::1424380312 +5::39::1::1424380312 +5::42::1::1424380312 +5::48::3::1424380312 +5::49::4::1424380312 +5::50::3::1424380312 +5::51::1::1424380312 +5::52::1::1424380312 +5::54::1::1424380312 +5::55::5::1424380312 +5::56::3::1424380312 +5::58::1::1424380312 +5::60::1::1424380312 +5::61::1::1424380312 +5::64::3::1424380312 +5::65::2::1424380312 +5::68::4::1424380312 +5::70::1::1424380312 +5::71::1::1424380312 +5::72::1::1424380312 +5::74::1::1424380312 +5::79::1::1424380312 +5::81::2::1424380312 +5::84::1::1424380312 +5::85::1::1424380312 +5::86::1::1424380312 +5::88::1::1424380312 +5::90::4::1424380312 +5::91::2::1424380312 +5::95::2::1424380312 +5::99::1::1424380312 +6::0::1::1424380312 +6::1::1::1424380312 +6::2::3::1424380312 +6::5::1::1424380312 +6::6::1::1424380312 +6::9::1::1424380312 +6::10::1::1424380312 +6::15::2::1424380312 +6::16::2::1424380312 +6::17::1::1424380312 +6::18::1::1424380312 +6::20::1::1424380312 +6::21::1::1424380312 +6::22::1::1424380312 +6::24::1::1424380312 +6::25::5::1424380312 +6::26::1::1424380312 +6::28::1::1424380312 +6::30::1::1424380312 +6::33::1::1424380312 +6::38::1::1424380312 +6::39::1::1424380312 +6::43::4::1424380312 +6::44::1::1424380312 +6::45::1::1424380312 +6::48::1::1424380312 +6::49::1::1424380312 +6::50::1::1424380312 +6::53::1::1424380312 +6::54::1::1424380312 +6::55::1::1424380312 +6::56::1::1424380312 +6::58::4::1424380312 +6::59::1::1424380312 +6::60::1::1424380312 +6::61::3::1424380312 +6::63::3::1424380312 +6::66::1::1424380312 +6::67::3::1424380312 +6::68::1::1424380312 +6::69::1::1424380312 +6::71::2::1424380312 +6::73::1::1424380312 +6::75::1::1424380312 +6::77::1::1424380312 +6::79::1::1424380312 +6::81::1::1424380312 +6::84::1::1424380312 +6::85::3::1424380312 +6::86::1::1424380312 +6::87::1::1424380312 +6::88::1::1424380312 +6::89::1::1424380312 +6::91::2::1424380312 +6::94::1::1424380312 +6::95::2::1424380312 +6::96::1::1424380312 +7::1::1::1424380312 +7::2::2::1424380312 +7::3::1::1424380312 +7::4::1::1424380312 +7::7::1::1424380312 +7::10::1::1424380312 +7::11::2::1424380312 +7::14::2::1424380312 +7::15::1::1424380312 +7::16::1::1424380312 +7::18::1::1424380312 +7::21::1::1424380312 +7::22::1::1424380312 +7::23::1::1424380312 +7::25::5::1424380312 +7::26::1::1424380312 +7::29::4::1424380312 +7::30::1::1424380312 +7::31::3::1424380312 +7::32::1::1424380312 +7::33::1::1424380312 +7::35::1::1424380312 +7::37::2::1424380312 +7::39::3::1424380312 +7::40::2::1424380312 +7::42::2::1424380312 +7::44::1::1424380312 +7::45::2::1424380312 +7::47::4::1424380312 +7::48::1::1424380312 +7::49::1::1424380312 +7::53::1::1424380312 +7::54::1::1424380312 +7::55::1::1424380312 +7::56::1::1424380312 +7::59::1::1424380312 +7::61::2::1424380312 +7::62::3::1424380312 +7::63::2::1424380312 +7::66::1::1424380312 +7::67::3::1424380312 +7::74::1::1424380312 +7::75::1::1424380312 +7::76::3::1424380312 +7::77::1::1424380312 +7::81::1::1424380312 +7::82::1::1424380312 +7::84::2::1424380312 +7::85::4::1424380312 +7::86::1::1424380312 +7::92::2::1424380312 +7::96::1::1424380312 +7::97::1::1424380312 +7::98::1::1424380312 +8::0::1::1424380312 +8::2::4::1424380312 +8::3::2::1424380312 +8::4::2::1424380312 +8::5::1::1424380312 +8::7::1::1424380312 +8::9::1::1424380312 +8::11::1::1424380312 +8::15::1::1424380312 +8::18::1::1424380312 +8::19::1::1424380312 +8::21::1::1424380312 +8::29::5::1424380312 +8::31::3::1424380312 +8::33::1::1424380312 +8::35::1::1424380312 +8::36::1::1424380312 +8::40::2::1424380312 +8::44::1::1424380312 +8::45::1::1424380312 +8::50::1::1424380312 +8::51::1::1424380312 +8::52::5::1424380312 +8::53::5::1424380312 +8::54::1::1424380312 +8::55::1::1424380312 +8::56::1::1424380312 +8::58::4::1424380312 +8::60::3::1424380312 +8::62::4::1424380312 +8::64::1::1424380312 +8::67::3::1424380312 +8::69::1::1424380312 +8::71::1::1424380312 +8::72::3::1424380312 +8::77::3::1424380312 +8::78::1::1424380312 +8::79::1::1424380312 +8::83::1::1424380312 +8::85::5::1424380312 +8::86::1::1424380312 +8::88::1::1424380312 +8::90::1::1424380312 +8::92::2::1424380312 +8::95::4::1424380312 +8::96::3::1424380312 +8::97::1::1424380312 +8::98::1::1424380312 +8::99::1::1424380312 +9::2::3::1424380312 +9::3::1::1424380312 +9::4::1::1424380312 +9::5::1::1424380312 +9::6::1::1424380312 +9::7::5::1424380312 +9::9::1::1424380312 +9::12::1::1424380312 +9::14::3::1424380312 +9::15::1::1424380312 +9::19::1::1424380312 +9::21::1::1424380312 +9::22::1::1424380312 +9::24::1::1424380312 +9::25::1::1424380312 +9::26::1::1424380312 +9::30::3::1424380312 +9::32::4::1424380312 +9::35::2::1424380312 +9::36::2::1424380312 +9::37::2::1424380312 +9::38::1::1424380312 +9::39::1::1424380312 +9::43::3::1424380312 +9::49::5::1424380312 +9::50::3::1424380312 +9::53::1::1424380312 +9::54::1::1424380312 +9::58::1::1424380312 +9::59::1::1424380312 +9::60::1::1424380312 +9::61::1::1424380312 +9::63::3::1424380312 +9::64::3::1424380312 +9::68::1::1424380312 +9::69::1::1424380312 +9::70::3::1424380312 +9::71::1::1424380312 +9::73::2::1424380312 +9::75::1::1424380312 +9::77::2::1424380312 +9::81::2::1424380312 +9::82::1::1424380312 +9::83::1::1424380312 +9::84::1::1424380312 +9::86::1::1424380312 +9::87::4::1424380312 +9::88::1::1424380312 +9::90::3::1424380312 +9::94::2::1424380312 +9::95::3::1424380312 +9::97::2::1424380312 +9::98::1::1424380312 +10::0::3::1424380312 +10::2::4::1424380312 +10::4::3::1424380312 +10::7::1::1424380312 +10::8::1::1424380312 +10::10::1::1424380312 +10::13::2::1424380312 +10::14::1::1424380312 +10::16::2::1424380312 +10::17::1::1424380312 +10::18::1::1424380312 +10::21::1::1424380312 +10::22::1::1424380312 +10::24::1::1424380312 +10::25::3::1424380312 +10::28::1::1424380312 +10::35::1::1424380312 +10::36::1::1424380312 +10::37::1::1424380312 +10::38::1::1424380312 +10::39::1::1424380312 +10::40::4::1424380312 +10::41::2::1424380312 +10::42::3::1424380312 +10::43::1::1424380312 +10::49::3::1424380312 +10::50::1::1424380312 +10::51::1::1424380312 +10::52::1::1424380312 +10::55::2::1424380312 +10::56::1::1424380312 +10::58::1::1424380312 +10::63::1::1424380312 +10::66::1::1424380312 +10::67::2::1424380312 +10::68::1::1424380312 +10::75::1::1424380312 +10::77::1::1424380312 +10::79::1::1424380312 +10::86::1::1424380312 +10::89::3::1424380312 +10::90::1::1424380312 +10::97::1::1424380312 +10::98::1::1424380312 +11::0::1::1424380312 +11::6::2::1424380312 +11::9::1::1424380312 +11::10::1::1424380312 +11::11::1::1424380312 +11::12::1::1424380312 +11::13::4::1424380312 +11::16::1::1424380312 +11::18::5::1424380312 +11::19::4::1424380312 +11::20::1::1424380312 +11::21::1::1424380312 +11::22::1::1424380312 +11::23::5::1424380312 +11::25::1::1424380312 +11::27::5::1424380312 +11::30::5::1424380312 +11::32::5::1424380312 +11::35::3::1424380312 +11::36::2::1424380312 +11::37::2::1424380312 +11::38::4::1424380312 +11::39::1::1424380312 +11::40::1::1424380312 +11::41::1::1424380312 +11::43::2::1424380312 +11::45::1::1424380312 +11::47::1::1424380312 +11::48::5::1424380312 +11::50::4::1424380312 +11::51::3::1424380312 +11::59::1::1424380312 +11::61::1::1424380312 +11::62::1::1424380312 +11::64::1::1424380312 +11::66::4::1424380312 +11::67::1::1424380312 +11::69::5::1424380312 +11::70::1::1424380312 +11::71::3::1424380312 +11::72::3::1424380312 +11::75::3::1424380312 +11::76::1::1424380312 +11::77::1::1424380312 +11::78::1::1424380312 +11::79::5::1424380312 +11::80::3::1424380312 +11::81::4::1424380312 +11::82::1::1424380312 +11::86::1::1424380312 +11::88::1::1424380312 +11::89::1::1424380312 +11::90::4::1424380312 +11::94::2::1424380312 +11::97::3::1424380312 +11::99::1::1424380312 +12::2::1::1424380312 +12::4::1::1424380312 +12::6::1::1424380312 +12::7::3::1424380312 +12::8::1::1424380312 +12::14::1::1424380312 +12::15::2::1424380312 +12::16::4::1424380312 +12::17::5::1424380312 +12::18::2::1424380312 +12::21::1::1424380312 +12::22::2::1424380312 +12::23::3::1424380312 +12::24::1::1424380312 +12::25::1::1424380312 +12::27::5::1424380312 +12::30::2::1424380312 +12::31::4::1424380312 +12::35::5::1424380312 +12::38::1::1424380312 +12::41::1::1424380312 +12::44::2::1424380312 +12::45::1::1424380312 +12::50::4::1424380312 +12::51::1::1424380312 +12::52::1::1424380312 +12::53::1::1424380312 +12::54::1::1424380312 +12::56::2::1424380312 +12::57::1::1424380312 +12::60::1::1424380312 +12::63::1::1424380312 +12::64::5::1424380312 +12::66::3::1424380312 +12::67::1::1424380312 +12::70::1::1424380312 +12::72::1::1424380312 +12::74::1::1424380312 +12::75::1::1424380312 +12::77::1::1424380312 +12::78::1::1424380312 +12::79::3::1424380312 +12::82::2::1424380312 +12::83::1::1424380312 +12::84::1::1424380312 +12::85::1::1424380312 +12::86::1::1424380312 +12::87::1::1424380312 +12::88::1::1424380312 +12::91::3::1424380312 +12::92::1::1424380312 +12::94::4::1424380312 +12::95::2::1424380312 +12::96::1::1424380312 +12::98::2::1424380312 +13::0::1::1424380312 +13::3::1::1424380312 +13::4::2::1424380312 +13::5::1::1424380312 +13::6::1::1424380312 +13::12::1::1424380312 +13::14::2::1424380312 +13::15::1::1424380312 +13::17::1::1424380312 +13::18::3::1424380312 +13::20::1::1424380312 +13::21::1::1424380312 +13::22::1::1424380312 +13::26::1::1424380312 +13::27::1::1424380312 +13::29::3::1424380312 +13::31::1::1424380312 +13::33::1::1424380312 +13::40::2::1424380312 +13::43::2::1424380312 +13::44::1::1424380312 +13::45::1::1424380312 +13::49::1::1424380312 +13::51::1::1424380312 +13::52::2::1424380312 +13::53::3::1424380312 +13::54::1::1424380312 +13::62::1::1424380312 +13::63::2::1424380312 +13::64::1::1424380312 +13::68::1::1424380312 +13::71::1::1424380312 +13::72::3::1424380312 +13::73::1::1424380312 +13::74::3::1424380312 +13::77::2::1424380312 +13::78::1::1424380312 +13::79::2::1424380312 +13::83::3::1424380312 +13::85::1::1424380312 +13::86::1::1424380312 +13::87::2::1424380312 +13::88::2::1424380312 +13::90::1::1424380312 +13::93::4::1424380312 +13::94::1::1424380312 +13::98::1::1424380312 +13::99::1::1424380312 +14::1::1::1424380312 +14::3::3::1424380312 +14::4::1::1424380312 +14::5::1::1424380312 +14::6::1::1424380312 +14::7::1::1424380312 +14::9::1::1424380312 +14::10::1::1424380312 +14::11::1::1424380312 +14::12::1::1424380312 +14::13::1::1424380312 +14::14::3::1424380312 +14::15::1::1424380312 +14::16::1::1424380312 +14::17::1::1424380312 +14::20::1::1424380312 +14::21::1::1424380312 +14::24::1::1424380312 +14::25::2::1424380312 +14::27::1::1424380312 +14::28::1::1424380312 +14::29::5::1424380312 +14::31::3::1424380312 +14::34::1::1424380312 +14::36::1::1424380312 +14::37::2::1424380312 +14::39::2::1424380312 +14::40::1::1424380312 +14::44::1::1424380312 +14::45::1::1424380312 +14::47::3::1424380312 +14::48::1::1424380312 +14::49::1::1424380312 +14::51::1::1424380312 +14::52::5::1424380312 +14::53::3::1424380312 +14::54::1::1424380312 +14::55::1::1424380312 +14::56::1::1424380312 +14::62::4::1424380312 +14::63::5::1424380312 +14::67::3::1424380312 +14::68::1::1424380312 +14::69::3::1424380312 +14::71::1::1424380312 +14::72::4::1424380312 +14::73::1::1424380312 +14::76::5::1424380312 +14::79::1::1424380312 +14::82::1::1424380312 +14::83::1::1424380312 +14::88::1::1424380312 +14::93::3::1424380312 +14::94::1::1424380312 +14::95::2::1424380312 +14::96::4::1424380312 +14::98::1::1424380312 +15::0::1::1424380312 +15::1::4::1424380312 +15::2::1::1424380312 +15::5::2::1424380312 +15::6::1::1424380312 +15::7::1::1424380312 +15::13::1::1424380312 +15::14::1::1424380312 +15::15::1::1424380312 +15::17::2::1424380312 +15::19::2::1424380312 +15::22::2::1424380312 +15::23::2::1424380312 +15::25::1::1424380312 +15::26::3::1424380312 +15::27::1::1424380312 +15::28::2::1424380312 +15::29::1::1424380312 +15::32::1::1424380312 +15::33::2::1424380312 +15::34::1::1424380312 +15::35::2::1424380312 +15::36::1::1424380312 +15::37::1::1424380312 +15::39::1::1424380312 +15::42::1::1424380312 +15::46::5::1424380312 +15::48::2::1424380312 +15::50::2::1424380312 +15::51::1::1424380312 +15::52::1::1424380312 +15::58::1::1424380312 +15::62::1::1424380312 +15::64::3::1424380312 +15::65::2::1424380312 +15::72::1::1424380312 +15::73::1::1424380312 +15::74::1::1424380312 +15::79::1::1424380312 +15::80::1::1424380312 +15::81::1::1424380312 +15::82::2::1424380312 +15::85::1::1424380312 +15::87::1::1424380312 +15::91::2::1424380312 +15::96::1::1424380312 +15::97::1::1424380312 +15::98::3::1424380312 +16::2::1::1424380312 +16::5::3::1424380312 +16::6::2::1424380312 +16::7::1::1424380312 +16::9::1::1424380312 +16::12::1::1424380312 +16::14::1::1424380312 +16::15::1::1424380312 +16::19::1::1424380312 +16::21::2::1424380312 +16::29::4::1424380312 +16::30::2::1424380312 +16::32::1::1424380312 +16::34::1::1424380312 +16::36::1::1424380312 +16::38::1::1424380312 +16::46::1::1424380312 +16::47::3::1424380312 +16::48::1::1424380312 +16::49::1::1424380312 +16::50::1::1424380312 +16::51::5::1424380312 +16::54::5::1424380312 +16::55::1::1424380312 +16::56::2::1424380312 +16::57::1::1424380312 +16::60::1::1424380312 +16::63::2::1424380312 +16::65::1::1424380312 +16::67::1::1424380312 +16::72::1::1424380312 +16::74::1::1424380312 +16::80::1::1424380312 +16::81::1::1424380312 +16::82::1::1424380312 +16::85::5::1424380312 +16::86::1::1424380312 +16::90::5::1424380312 +16::91::1::1424380312 +16::93::1::1424380312 +16::94::3::1424380312 +16::95::2::1424380312 +16::96::3::1424380312 +16::98::3::1424380312 +16::99::1::1424380312 +17::2::1::1424380312 +17::3::1::1424380312 +17::6::1::1424380312 +17::10::4::1424380312 +17::11::1::1424380312 +17::13::2::1424380312 +17::17::5::1424380312 +17::19::1::1424380312 +17::20::5::1424380312 +17::22::4::1424380312 +17::28::1::1424380312 +17::29::1::1424380312 +17::33::1::1424380312 +17::34::1::1424380312 +17::35::2::1424380312 +17::37::1::1424380312 +17::38::1::1424380312 +17::45::1::1424380312 +17::46::5::1424380312 +17::47::1::1424380312 +17::49::3::1424380312 +17::51::1::1424380312 +17::55::5::1424380312 +17::56::3::1424380312 +17::57::1::1424380312 +17::58::1::1424380312 +17::59::1::1424380312 +17::60::1::1424380312 +17::63::1::1424380312 +17::66::1::1424380312 +17::68::4::1424380312 +17::69::1::1424380312 +17::70::1::1424380312 +17::72::1::1424380312 +17::73::3::1424380312 +17::78::1::1424380312 +17::79::1::1424380312 +17::82::2::1424380312 +17::84::1::1424380312 +17::90::5::1424380312 +17::91::3::1424380312 +17::92::1::1424380312 +17::93::1::1424380312 +17::94::4::1424380312 +17::95::2::1424380312 +17::97::1::1424380312 +18::1::1::1424380312 +18::4::3::1424380312 +18::5::2::1424380312 +18::6::1::1424380312 +18::7::1::1424380312 +18::10::1::1424380312 +18::11::4::1424380312 +18::12::2::1424380312 +18::13::1::1424380312 +18::15::1::1424380312 +18::18::1::1424380312 +18::20::1::1424380312 +18::21::2::1424380312 +18::22::1::1424380312 +18::23::2::1424380312 +18::25::1::1424380312 +18::26::1::1424380312 +18::27::1::1424380312 +18::28::5::1424380312 +18::29::1::1424380312 +18::31::1::1424380312 +18::32::1::1424380312 +18::36::1::1424380312 +18::38::5::1424380312 +18::39::5::1424380312 +18::40::1::1424380312 +18::42::1::1424380312 +18::43::1::1424380312 +18::44::4::1424380312 +18::46::1::1424380312 +18::47::1::1424380312 +18::48::1::1424380312 +18::51::2::1424380312 +18::55::1::1424380312 +18::56::1::1424380312 +18::57::1::1424380312 +18::62::1::1424380312 +18::63::1::1424380312 +18::66::3::1424380312 +18::67::1::1424380312 +18::70::1::1424380312 +18::75::1::1424380312 +18::76::3::1424380312 +18::77::1::1424380312 +18::80::3::1424380312 +18::81::3::1424380312 +18::82::1::1424380312 +18::83::5::1424380312 +18::84::1::1424380312 +18::97::1::1424380312 +18::98::1::1424380312 +18::99::2::1424380312 +19::0::1::1424380312 +19::1::1::1424380312 +19::2::1::1424380312 +19::4::1::1424380312 +19::6::2::1424380312 +19::11::1::1424380312 +19::12::1::1424380312 +19::14::1::1424380312 +19::23::1::1424380312 +19::26::1::1424380312 +19::31::1::1424380312 +19::32::4::1424380312 +19::33::1::1424380312 +19::34::1::1424380312 +19::37::1::1424380312 +19::38::1::1424380312 +19::41::1::1424380312 +19::43::1::1424380312 +19::45::1::1424380312 +19::48::1::1424380312 +19::49::1::1424380312 +19::50::2::1424380312 +19::53::2::1424380312 +19::54::3::1424380312 +19::55::1::1424380312 +19::56::2::1424380312 +19::58::1::1424380312 +19::61::1::1424380312 +19::62::1::1424380312 +19::63::1::1424380312 +19::64::1::1424380312 +19::65::1::1424380312 +19::69::2::1424380312 +19::72::1::1424380312 +19::74::3::1424380312 +19::76::1::1424380312 +19::78::1::1424380312 +19::79::1::1424380312 +19::81::1::1424380312 +19::82::1::1424380312 +19::84::1::1424380312 +19::86::1::1424380312 +19::87::2::1424380312 +19::90::4::1424380312 +19::93::1::1424380312 +19::94::4::1424380312 +19::95::2::1424380312 +19::96::1::1424380312 +19::98::4::1424380312 +20::0::1::1424380312 +20::1::1::1424380312 +20::2::2::1424380312 +20::4::2::1424380312 +20::6::1::1424380312 +20::8::1::1424380312 +20::12::1::1424380312 +20::21::2::1424380312 +20::22::5::1424380312 +20::24::2::1424380312 +20::25::1::1424380312 +20::26::1::1424380312 +20::29::2::1424380312 +20::30::2::1424380312 +20::32::2::1424380312 +20::39::1::1424380312 +20::40::1::1424380312 +20::41::2::1424380312 +20::45::2::1424380312 +20::48::1::1424380312 +20::50::1::1424380312 +20::51::3::1424380312 +20::53::3::1424380312 +20::55::1::1424380312 +20::57::2::1424380312 +20::60::1::1424380312 +20::61::1::1424380312 +20::64::1::1424380312 +20::66::1::1424380312 +20::70::2::1424380312 +20::72::1::1424380312 +20::73::2::1424380312 +20::75::4::1424380312 +20::76::1::1424380312 +20::77::4::1424380312 +20::78::1::1424380312 +20::79::1::1424380312 +20::84::2::1424380312 +20::85::2::1424380312 +20::88::3::1424380312 +20::89::1::1424380312 +20::90::3::1424380312 +20::91::1::1424380312 +20::92::2::1424380312 +20::93::1::1424380312 +20::94::4::1424380312 +20::97::1::1424380312 +21::0::1::1424380312 +21::2::4::1424380312 +21::3::1::1424380312 +21::7::2::1424380312 +21::11::1::1424380312 +21::12::1::1424380312 +21::13::1::1424380312 +21::14::3::1424380312 +21::17::1::1424380312 +21::19::1::1424380312 +21::20::1::1424380312 +21::21::1::1424380312 +21::22::1::1424380312 +21::23::1::1424380312 +21::24::1::1424380312 +21::27::1::1424380312 +21::29::5::1424380312 +21::30::2::1424380312 +21::38::1::1424380312 +21::40::2::1424380312 +21::43::3::1424380312 +21::44::1::1424380312 +21::45::1::1424380312 +21::46::1::1424380312 +21::48::1::1424380312 +21::51::1::1424380312 +21::53::5::1424380312 +21::54::1::1424380312 +21::55::1::1424380312 +21::56::1::1424380312 +21::58::3::1424380312 +21::59::3::1424380312 +21::64::1::1424380312 +21::66::1::1424380312 +21::68::1::1424380312 +21::71::1::1424380312 +21::73::1::1424380312 +21::74::4::1424380312 +21::80::1::1424380312 +21::81::1::1424380312 +21::83::1::1424380312 +21::84::1::1424380312 +21::85::3::1424380312 +21::87::4::1424380312 +21::89::2::1424380312 +21::92::2::1424380312 +21::96::3::1424380312 +21::99::1::1424380312 +22::0::1::1424380312 +22::3::2::1424380312 +22::5::2::1424380312 +22::6::2::1424380312 +22::9::1::1424380312 +22::10::1::1424380312 +22::11::1::1424380312 +22::13::1::1424380312 +22::14::1::1424380312 +22::16::1::1424380312 +22::18::3::1424380312 +22::19::1::1424380312 +22::22::5::1424380312 +22::25::1::1424380312 +22::26::1::1424380312 +22::29::3::1424380312 +22::30::5::1424380312 +22::32::4::1424380312 +22::33::1::1424380312 +22::35::1::1424380312 +22::36::3::1424380312 +22::37::1::1424380312 +22::40::1::1424380312 +22::41::3::1424380312 +22::44::1::1424380312 +22::45::2::1424380312 +22::48::1::1424380312 +22::51::5::1424380312 +22::55::1::1424380312 +22::56::2::1424380312 +22::60::3::1424380312 +22::61::1::1424380312 +22::62::4::1424380312 +22::63::1::1424380312 +22::65::1::1424380312 +22::66::1::1424380312 +22::68::4::1424380312 +22::69::4::1424380312 +22::70::3::1424380312 +22::71::1::1424380312 +22::74::5::1424380312 +22::75::5::1424380312 +22::78::1::1424380312 +22::80::3::1424380312 +22::81::1::1424380312 +22::82::1::1424380312 +22::84::1::1424380312 +22::86::1::1424380312 +22::87::3::1424380312 +22::88::5::1424380312 +22::90::2::1424380312 +22::92::3::1424380312 +22::95::2::1424380312 +22::96::2::1424380312 +22::98::4::1424380312 +22::99::1::1424380312 +23::0::1::1424380312 +23::2::1::1424380312 +23::4::1::1424380312 +23::6::2::1424380312 +23::10::4::1424380312 +23::12::1::1424380312 +23::13::4::1424380312 +23::14::1::1424380312 +23::15::1::1424380312 +23::18::4::1424380312 +23::22::2::1424380312 +23::23::4::1424380312 +23::24::1::1424380312 +23::25::1::1424380312 +23::26::1::1424380312 +23::27::5::1424380312 +23::28::1::1424380312 +23::29::1::1424380312 +23::30::4::1424380312 +23::32::5::1424380312 +23::33::2::1424380312 +23::36::3::1424380312 +23::37::1::1424380312 +23::38::1::1424380312 +23::39::1::1424380312 +23::43::1::1424380312 +23::48::5::1424380312 +23::49::5::1424380312 +23::50::4::1424380312 +23::53::1::1424380312 +23::55::5::1424380312 +23::57::1::1424380312 +23::59::1::1424380312 +23::60::1::1424380312 +23::61::1::1424380312 +23::64::4::1424380312 +23::65::5::1424380312 +23::66::2::1424380312 +23::67::1::1424380312 +23::68::3::1424380312 +23::69::1::1424380312 +23::72::1::1424380312 +23::73::3::1424380312 +23::77::1::1424380312 +23::82::2::1424380312 +23::83::1::1424380312 +23::84::1::1424380312 +23::85::1::1424380312 +23::87::3::1424380312 +23::88::1::1424380312 +23::95::2::1424380312 +23::97::1::1424380312 +24::4::1::1424380312 +24::6::3::1424380312 +24::7::1::1424380312 +24::10::2::1424380312 +24::12::1::1424380312 +24::15::1::1424380312 +24::19::1::1424380312 +24::24::1::1424380312 +24::27::3::1424380312 +24::30::5::1424380312 +24::31::1::1424380312 +24::32::3::1424380312 +24::33::1::1424380312 +24::37::1::1424380312 +24::39::1::1424380312 +24::40::1::1424380312 +24::42::1::1424380312 +24::43::3::1424380312 +24::45::2::1424380312 +24::46::1::1424380312 +24::47::1::1424380312 +24::48::1::1424380312 +24::49::1::1424380312 +24::50::1::1424380312 +24::52::5::1424380312 +24::57::1::1424380312 +24::59::4::1424380312 +24::63::4::1424380312 +24::65::1::1424380312 +24::66::1::1424380312 +24::67::1::1424380312 +24::68::3::1424380312 +24::69::5::1424380312 +24::71::1::1424380312 +24::72::4::1424380312 +24::77::4::1424380312 +24::78::1::1424380312 +24::80::1::1424380312 +24::82::1::1424380312 +24::84::1::1424380312 +24::86::1::1424380312 +24::87::1::1424380312 +24::88::2::1424380312 +24::89::1::1424380312 +24::90::5::1424380312 +24::91::1::1424380312 +24::92::1::1424380312 +24::94::2::1424380312 +24::95::1::1424380312 +24::96::5::1424380312 +24::98::1::1424380312 +24::99::1::1424380312 +25::1::3::1424380312 +25::2::1::1424380312 +25::7::1::1424380312 +25::9::1::1424380312 +25::12::3::1424380312 +25::16::3::1424380312 +25::17::1::1424380312 +25::18::1::1424380312 +25::20::1::1424380312 +25::22::1::1424380312 +25::23::1::1424380312 +25::26::2::1424380312 +25::29::1::1424380312 +25::30::1::1424380312 +25::31::2::1424380312 +25::33::4::1424380312 +25::34::3::1424380312 +25::35::2::1424380312 +25::36::1::1424380312 +25::37::1::1424380312 +25::40::1::1424380312 +25::41::1::1424380312 +25::43::1::1424380312 +25::47::4::1424380312 +25::50::1::1424380312 +25::51::1::1424380312 +25::53::1::1424380312 +25::56::1::1424380312 +25::58::2::1424380312 +25::64::2::1424380312 +25::67::2::1424380312 +25::68::1::1424380312 +25::70::1::1424380312 +25::71::4::1424380312 +25::73::1::1424380312 +25::74::1::1424380312 +25::76::1::1424380312 +25::79::1::1424380312 +25::82::1::1424380312 +25::84::2::1424380312 +25::85::1::1424380312 +25::91::3::1424380312 +25::92::1::1424380312 +25::94::1::1424380312 +25::95::1::1424380312 +25::97::2::1424380312 +26::0::1::1424380312 +26::1::1::1424380312 +26::2::1::1424380312 +26::3::1::1424380312 +26::4::4::1424380312 +26::5::2::1424380312 +26::6::3::1424380312 +26::7::5::1424380312 +26::13::3::1424380312 +26::14::1::1424380312 +26::16::1::1424380312 +26::18::3::1424380312 +26::20::1::1424380312 +26::21::3::1424380312 +26::22::5::1424380312 +26::23::5::1424380312 +26::24::5::1424380312 +26::27::1::1424380312 +26::31::1::1424380312 +26::35::1::1424380312 +26::36::4::1424380312 +26::40::1::1424380312 +26::44::1::1424380312 +26::45::2::1424380312 +26::47::1::1424380312 +26::48::1::1424380312 +26::49::3::1424380312 +26::50::2::1424380312 +26::52::1::1424380312 +26::54::4::1424380312 +26::55::1::1424380312 +26::57::3::1424380312 +26::58::1::1424380312 +26::61::1::1424380312 +26::62::2::1424380312 +26::66::1::1424380312 +26::68::4::1424380312 +26::71::1::1424380312 +26::73::4::1424380312 +26::76::1::1424380312 +26::81::3::1424380312 +26::85::1::1424380312 +26::86::3::1424380312 +26::88::5::1424380312 +26::91::1::1424380312 +26::94::5::1424380312 +26::95::1::1424380312 +26::96::1::1424380312 +26::97::1::1424380312 +27::0::1::1424380312 +27::9::1::1424380312 +27::10::1::1424380312 +27::18::4::1424380312 +27::19::3::1424380312 +27::20::1::1424380312 +27::22::2::1424380312 +27::24::2::1424380312 +27::25::1::1424380312 +27::27::3::1424380312 +27::28::1::1424380312 +27::29::1::1424380312 +27::31::1::1424380312 +27::33::3::1424380312 +27::40::1::1424380312 +27::42::1::1424380312 +27::43::1::1424380312 +27::44::3::1424380312 +27::45::1::1424380312 +27::51::3::1424380312 +27::52::1::1424380312 +27::55::3::1424380312 +27::57::1::1424380312 +27::59::1::1424380312 +27::60::1::1424380312 +27::61::1::1424380312 +27::64::1::1424380312 +27::66::3::1424380312 +27::68::1::1424380312 +27::70::1::1424380312 +27::71::2::1424380312 +27::72::1::1424380312 +27::75::3::1424380312 +27::78::1::1424380312 +27::80::3::1424380312 +27::82::1::1424380312 +27::83::3::1424380312 +27::86::1::1424380312 +27::87::2::1424380312 +27::90::1::1424380312 +27::91::1::1424380312 +27::92::1::1424380312 +27::93::1::1424380312 +27::94::2::1424380312 +27::95::1::1424380312 +27::98::1::1424380312 +28::0::3::1424380312 +28::1::1::1424380312 +28::2::4::1424380312 +28::3::1::1424380312 +28::6::1::1424380312 +28::7::1::1424380312 +28::12::5::1424380312 +28::13::2::1424380312 +28::14::1::1424380312 +28::15::1::1424380312 +28::17::1::1424380312 +28::19::3::1424380312 +28::20::1::1424380312 +28::23::3::1424380312 +28::24::3::1424380312 +28::27::1::1424380312 +28::29::1::1424380312 +28::33::1::1424380312 +28::34::1::1424380312 +28::36::1::1424380312 +28::38::2::1424380312 +28::39::2::1424380312 +28::44::1::1424380312 +28::45::1::1424380312 +28::49::4::1424380312 +28::50::1::1424380312 +28::52::1::1424380312 +28::54::1::1424380312 +28::56::1::1424380312 +28::57::3::1424380312 +28::58::1::1424380312 +28::59::1::1424380312 +28::60::1::1424380312 +28::62::3::1424380312 +28::63::1::1424380312 +28::65::1::1424380312 +28::75::1::1424380312 +28::78::1::1424380312 +28::81::5::1424380312 +28::82::4::1424380312 +28::83::1::1424380312 +28::85::1::1424380312 +28::88::2::1424380312 +28::89::4::1424380312 +28::90::1::1424380312 +28::92::5::1424380312 +28::94::1::1424380312 +28::95::2::1424380312 +28::98::1::1424380312 +28::99::1::1424380312 +29::3::1::1424380312 +29::4::1::1424380312 +29::5::1::1424380312 +29::7::2::1424380312 +29::9::1::1424380312 +29::10::3::1424380312 +29::11::1::1424380312 +29::13::3::1424380312 +29::14::1::1424380312 +29::15::1::1424380312 +29::17::3::1424380312 +29::19::3::1424380312 +29::22::3::1424380312 +29::23::4::1424380312 +29::25::1::1424380312 +29::29::1::1424380312 +29::31::1::1424380312 +29::32::4::1424380312 +29::33::2::1424380312 +29::36::2::1424380312 +29::38::3::1424380312 +29::39::1::1424380312 +29::42::1::1424380312 +29::46::5::1424380312 +29::49::3::1424380312 +29::51::2::1424380312 +29::59::1::1424380312 +29::61::1::1424380312 +29::62::1::1424380312 +29::67::1::1424380312 +29::68::3::1424380312 +29::69::1::1424380312 +29::70::1::1424380312 +29::74::1::1424380312 +29::75::1::1424380312 +29::79::2::1424380312 +29::80::1::1424380312 +29::81::2::1424380312 +29::83::1::1424380312 +29::85::1::1424380312 +29::86::1::1424380312 +29::90::4::1424380312 +29::93::1::1424380312 +29::94::4::1424380312 +29::97::1::1424380312 +29::99::1::1424380312 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/als/test.data b/spark-3.5.0-bin-hadoop3/data/mllib/als/test.data new file mode 100644 index 0000000..e476cc2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/als/test.data @@ -0,0 +1,16 @@ +1,1,5.0 +1,2,1.0 +1,3,5.0 +1,4,1.0 +2,1,5.0 +2,2,1.0 +2,3,5.0 +2,4,1.0 +3,1,1.0 +3,2,5.0 +3,3,1.0 +3,4,5.0 +4,1,1.0 +4,2,5.0 +4,3,1.0 +4,4,5.0 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/gmm_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/gmm_data.txt new file mode 100644 index 0000000..934ee4a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/gmm_data.txt @@ -0,0 +1,2000 @@ + 2.59470454e+00 2.12298217e+00 + 1.15807024e+00 -1.46498723e-01 + 2.46206638e+00 6.19556894e-01 + -5.54845070e-01 -7.24700066e-01 + -3.23111426e+00 -1.42579084e+00 + 3.02978115e+00 7.87121753e-01 + 1.97365907e+00 1.15914704e+00 + -6.44852101e+00 -3.18154314e+00 + 1.30963349e+00 1.62866434e-01 + 4.26482541e+00 2.15547996e+00 + 3.79927257e+00 1.50572445e+00 + 4.17452609e-01 -6.74032760e-01 + 4.21117627e-01 4.45590255e-01 + -2.80425571e+00 -7.77150554e-01 + 2.55928797e+00 7.03954218e-01 + 1.32554059e+00 -9.46663152e-01 + -3.39691439e+00 -1.49005743e+00 + -2.26542270e-01 3.60052515e-02 + 1.04994198e+00 5.29825685e-01 + -1.51566882e+00 -1.86264432e-01 + -3.27928172e-01 -7.60859110e-01 + -3.18054866e-01 3.97719805e-01 + 1.65579418e-01 -3.47232033e-01 + 6.47162333e-01 4.96059961e-02 + -2.80776647e-01 4.79418757e-01 + 7.45069752e-01 1.20790281e-01 + 2.13604102e-01 1.59542555e-01 + -3.08860224e+00 -1.43259870e+00 + 8.97066497e-01 1.10206801e+00 + -2.23918874e-01 -1.07267267e+00 + 2.51525708e+00 2.84761973e-01 + 9.98052532e-01 1.08333783e+00 + 1.76705588e+00 8.18866778e-01 + 5.31555163e-02 -1.90111151e-01 + -2.17405059e+00 7.21854582e-02 + -2.13772505e+00 -3.62010387e-01 + 2.95974057e+00 1.31602381e+00 + 2.74053561e+00 1.61781757e+00 + 6.68135448e-01 2.86586009e-01 + 2.82323739e+00 1.74437257e+00 + 8.11540288e-01 5.50744478e-01 + 4.10050897e-01 5.10668402e-03 + 9.58626136e-01 -3.49633680e-01 + 4.66599798e+00 1.49964894e+00 + 4.94507794e-01 2.58928077e-01 + -2.36029742e+00 -1.61042909e+00 + -4.99306804e-01 -8.04984769e-01 + 1.07448510e+00 9.39605828e-01 + -1.80448949e+00 -1.05983264e+00 + -3.22353821e-01 1.73612093e-01 + 1.85418702e+00 1.15640643e+00 + 6.93794163e-01 6.59993560e-01 + 1.99399102e+00 1.44547123e+00 + 3.38866124e+00 1.23379290e+00 + -4.24067720e+00 -1.22264282e+00 + 6.03230201e-02 2.95232729e-01 + -3.59341813e+00 -7.17453726e-01 + 4.87447372e-01 -2.00733911e-01 + 1.20149195e+00 4.07880197e-01 + -2.13331464e+00 -4.58518077e-01 + -3.84091083e+00 -1.71553950e+00 + -5.37279250e-01 2.64822629e-02 + -2.10155227e+00 -1.32558103e+00 + -1.71318897e+00 -7.12098563e-01 + -1.46280695e+00 -1.84868337e-01 + -3.59785325e+00 -1.54832434e+00 + -5.77528081e-01 -5.78580857e-01 + 3.14734283e-01 5.80184639e-01 + -2.71164714e+00 -1.19379432e+00 + 1.09634489e+00 7.20143887e-01 + -3.05527722e+00 -1.47774064e+00 + 6.71753586e-01 7.61350020e-01 + 3.98294144e+00 1.54166484e+00 + -3.37220384e+00 -2.21332064e+00 + 1.81222914e+00 7.41212752e-01 + 2.71458282e-01 1.36329078e-01 + -3.97815359e-01 1.16766886e-01 + -1.70192814e+00 -9.75851571e-01 + -3.46803804e+00 -1.09965988e+00 + -1.69649627e+00 -5.76045801e-01 + -1.02485636e-01 -8.81841246e-01 + -3.24194667e-02 2.55429276e-01 + -2.75343168e+00 -1.51366320e+00 + -2.78676702e+00 -5.22360489e-01 + 1.70483164e+00 1.19769805e+00 + 4.92022579e-01 3.24944706e-01 + 2.48768464e+00 1.00055363e+00 + 4.48786400e-01 7.63902870e-01 + 2.93862696e+00 1.73809968e+00 + -3.55019305e+00 -1.97875558e+00 + 1.74270784e+00 6.90229224e-01 + 5.13391994e-01 4.58374016e-01 + 1.78379499e+00 9.08026381e-01 + 1.75814147e+00 7.41449784e-01 + -2.30687792e-01 3.91009729e-01 + 3.92271353e+00 1.44006290e+00 + 2.93361679e-01 -4.99886375e-03 + 2.47902690e-01 -7.49542503e-01 + -3.97675355e-01 1.36824887e-01 + 3.56535953e+00 1.15181329e+00 + 3.22425301e+00 1.28702383e+00 + -2.94192478e-01 -2.42382557e-01 + 8.02068864e-01 -1.51671475e-01 + 8.54133530e-01 -4.89514885e-02 + -1.64316316e-01 -5.34642346e-01 + -6.08485405e-01 -2.10332352e-01 + -2.18940059e+00 -1.07024952e+00 + -1.71586960e+00 -2.83333492e-02 + 1.70200448e-01 -3.28031178e-01 + -1.97210346e+00 -5.39948532e-01 + 2.19500160e+00 1.05697170e+00 + -1.76239935e+00 -1.09377438e+00 + 1.68314744e+00 6.86491164e-01 + -2.99852288e+00 -1.46619067e+00 + -2.23769560e+00 -9.15008355e-01 + 9.46887516e-01 5.58410503e-01 + 5.02153123e-01 1.63851235e-01 + -9.70297062e-01 3.14625374e-01 + -1.29405593e+00 -8.20994131e-01 + 2.72516079e+00 7.85839947e-01 + 1.45788024e+00 3.37487353e-01 + -4.36292749e-01 -5.42150480e-01 + 2.21304711e+00 1.25254042e+00 + -1.20810271e-01 4.79632898e-01 + -3.30884511e+00 -1.50607586e+00 + -6.55882455e+00 -1.94231256e+00 + -3.17033630e+00 -9.94678930e-01 + 1.42043617e+00 7.28808957e-01 + -1.57546099e+00 -1.10320497e+00 + -3.22748754e+00 -1.64174579e+00 + 2.96776017e-03 -3.16191512e-02 + -2.25986054e+00 -6.13123197e-01 + 2.49434243e+00 7.73069183e-01 + 9.08494049e-01 -1.53926853e-01 + -2.80559090e+00 -1.37474221e+00 + 4.75224286e-01 2.53153674e-01 + 4.37644006e+00 8.49116998e-01 + 2.27282959e+00 6.16568202e-01 + 1.16006880e+00 1.65832798e-01 + -1.67163193e+00 -1.22555386e+00 + -1.38231118e+00 -7.29575504e-01 + -3.49922750e+00 -2.26446675e+00 + -3.73780110e-01 -1.90657869e-01 + 1.68627679e+00 1.05662987e+00 + -3.28891792e+00 -1.11080334e+00 + -2.59815798e+00 -1.51410198e+00 + -2.61203309e+00 -6.00143552e-01 + 6.58964943e-01 4.47216094e-01 + -2.26711381e+00 -7.26512923e-01 + -5.31429009e-02 -1.97925341e-02 + 3.19749807e+00 9.20425476e-01 + -1.37595787e+00 -6.58062732e-01 + 8.09900278e-01 -3.84286160e-01 + -5.07741280e+00 -1.97683808e+00 + -2.99764250e+00 -1.50753777e+00 + -9.87671815e-01 -4.63255889e-01 + 1.65390765e+00 6.73806615e-02 + 5.51252659e+00 2.69842267e+00 + -2.23724309e+00 -4.77624004e-01 + 4.99726228e+00 1.74690949e+00 + 1.75859162e-01 -1.49350995e-01 + 4.13382789e+00 1.31735161e+00 + 2.69058117e+00 4.87656923e-01 + 1.07180318e+00 1.01426954e+00 + 3.37216869e+00 1.05955377e+00 + -2.95006781e+00 -1.57048303e+00 + -2.46401648e+00 -8.37056374e-01 + 1.19012962e-01 7.54702770e-01 + 3.34142539e+00 4.81938295e-01 + 2.92643913e+00 1.04301050e+00 + 2.89697751e+00 1.37551442e+00 + -1.03094242e+00 2.20903962e-01 + -5.13914589e+00 -2.23355387e+00 + -8.81680780e-01 1.83590000e-01 + 2.82334775e+00 1.26650464e+00 + -2.81042540e-01 -3.26370240e-01 + 2.97995487e+00 8.34569452e-01 + -1.39857135e+00 -1.15798385e+00 + 4.27186506e+00 9.04253702e-01 + 6.98684517e-01 7.91167305e-01 + 3.52233095e+00 1.29976473e+00 + 2.21448029e+00 2.73213379e-01 + -3.13505683e-01 -1.20593774e-01 + 3.70571571e+00 1.06220876e+00 + 9.83881041e-01 5.67713803e-01 + -2.17897705e+00 2.52925205e-01 + 1.38734039e+00 4.61287066e-01 + -1.41181602e+00 -1.67248955e-02 + -1.69974639e+00 -7.17812071e-01 + -2.01005793e-01 -7.49662056e-01 + 1.69016336e+00 3.24687979e-01 + -2.03250179e+00 -2.76108460e-01 + 3.68776848e-01 4.12536941e-01 + 7.66238259e-01 -1.84750637e-01 + -2.73989147e-01 -1.72817250e-01 + -2.18623745e+00 -2.10906798e-01 + -1.39795625e-01 3.26066094e-02 + -2.73826912e-01 -6.67586097e-02 + -1.57880654e+00 -4.99395900e-01 + 4.55950908e+00 2.29410489e+00 + -7.36479631e-01 -1.57861857e-01 + 1.92082888e+00 1.05843391e+00 + 4.29192810e+00 1.38127810e+00 + 1.61852879e+00 1.95871986e-01 + -1.95027403e+00 -5.22448168e-01 + -1.67446281e+00 -9.41497162e-01 + 6.07097859e-01 3.44178029e-01 + -3.44004683e+00 -1.49258461e+00 + 2.72114752e+00 6.00728991e-01 + 8.80685522e-01 -2.53243336e-01 + 1.39254928e+00 3.42988512e-01 + 1.14194836e-01 -8.57945694e-02 + -1.49387332e+00 -7.60860481e-01 + -1.98053285e+00 -4.86039865e-01 + 3.56008568e+00 1.08438692e+00 + 2.27833961e-01 1.09441881e+00 + -1.16716710e+00 -6.54778242e-01 + 2.02156613e+00 5.42075758e-01 + 1.08429178e+00 -7.67420693e-01 + 6.63058455e-01 4.61680991e-01 + -1.06201537e+00 1.38862846e-01 + 3.08701875e+00 8.32580273e-01 + -4.96558108e-01 -2.47031257e-01 + 7.95109987e-01 7.59314147e-02 + -3.39903524e-01 8.71565566e-03 + 8.68351357e-01 4.78358641e-01 + 1.48750819e+00 7.63257420e-01 + -4.51224101e-01 -4.44056898e-01 + -3.02734750e-01 -2.98487961e-01 + 5.46846609e-01 7.02377629e-01 + 1.65129778e+00 3.74008231e-01 + -7.43336512e-01 3.95723531e-01 + -5.88446605e-01 -6.47520211e-01 + 3.58613167e+00 1.95024937e+00 + 3.11718883e+00 8.37984715e-01 + 1.80919244e+00 9.62644986e-01 + 5.43856371e-02 -5.86297543e-01 + -1.95186766e+00 -1.02624212e-01 + 8.95628057e-01 5.91812281e-01 + 4.97691627e-02 5.31137156e-01 + -1.07633113e+00 -2.47392788e-01 + -1.17257986e+00 -8.68528265e-01 + -8.19227665e-02 5.80579434e-03 + -2.86409787e-01 1.95812924e-01 + 1.10582671e+00 7.42853240e-01 + 4.06429774e+00 1.06557476e+00 + -3.42521792e+00 -7.74327139e-01 + 1.28468671e+00 6.20431661e-01 + 6.01201008e-01 -1.16799728e-01 + -1.85058727e-01 -3.76235293e-01 + 5.44083324e+00 2.98490868e+00 + 2.69273070e+00 7.83901153e-01 + 1.88938036e-01 -4.83222152e-01 + 1.05667256e+00 -2.57003165e-01 + 2.99711662e-01 -4.33131912e-01 + 7.73689216e-02 -1.78738364e-01 + 9.58326279e-01 6.38325706e-01 + -3.97727049e-01 2.27314759e-01 + 3.36098175e+00 1.12165237e+00 + 1.77804871e+00 6.46961933e-01 + -2.86945546e+00 -1.00395518e+00 + 3.03494815e+00 7.51814612e-01 + -1.43658194e+00 -3.55432244e-01 + -3.08455105e+00 -1.51535106e+00 + -1.55841975e+00 3.93454820e-02 + 7.96073412e-01 -3.11036969e-01 + -9.84125401e-01 -1.02064649e+00 + -7.75688143e+00 -3.65219926e+00 + 1.53816429e+00 7.65926670e-01 + -4.92712738e-01 2.32244240e-02 + -1.93166919e+00 -1.07701304e+00 + 2.03029875e-02 -7.54055699e-01 + 2.52177489e+00 1.01544979e+00 + 3.65109048e-01 -9.48328494e-01 + -1.28849143e-01 2.51947174e-01 + -1.02428075e+00 -9.37767116e-01 + -3.04179748e+00 -9.97926994e-01 + -2.51986980e+00 -1.69117413e+00 + -1.24900838e+00 -4.16179917e-01 + 2.77943992e+00 1.22842327e+00 + -4.37434557e+00 -1.70182693e+00 + -1.60019319e+00 -4.18345639e-01 + -1.67613646e+00 -9.44087262e-01 + -9.00843245e-01 8.26378089e-02 + 3.29770621e-01 -9.07870444e-01 + -2.84650535e+00 -9.00155396e-01 + 1.57111705e+00 7.07432268e-01 + 1.24948552e+00 1.04812849e-01 + 1.81440558e+00 9.53545082e-01 + -1.74915794e+00 -1.04606288e+00 + 1.20593269e+00 -1.12607147e-02 + 1.36004919e-01 -1.09828044e+00 + 2.57480693e-01 3.34941541e-01 + 7.78775385e-01 -5.32494732e-01 + -1.79155126e+00 -6.29994129e-01 + -1.75706839e+00 -8.35100126e-01 + 4.29512012e-01 7.81426910e-02 + 3.08349370e-01 -1.27359861e-01 + 1.05560329e+00 4.55150640e-01 + 1.95662574e+00 1.17593217e+00 + 8.77376632e-01 6.57866662e-01 + 7.71311255e-01 9.15134334e-02 + -6.36978275e+00 -2.55874241e+00 + -2.98335339e+00 -1.59567024e+00 + -3.67104587e-01 1.85315291e-01 + 1.95347407e+00 -7.15503113e-02 + 8.45556363e-01 6.51256415e-02 + 9.42868521e-01 3.56647624e-01 + 2.99321875e+00 1.07505254e+00 + -2.91030538e-01 -3.77637183e-01 + 1.62870918e+00 3.37563671e-01 + 2.05773173e-01 3.43337416e-01 + -8.40879199e-01 -1.35600767e-01 + 1.38101624e+00 5.99253495e-01 + -6.93715607e+00 -2.63580662e+00 + -1.04423404e+00 -8.32865050e-01 + 1.33448476e+00 1.04863475e+00 + 6.01675207e-01 1.98585194e-01 + 2.31233993e+00 7.98628331e-01 + 1.85201313e-01 -1.76070247e+00 + 1.92006354e+00 8.45737582e-01 + 1.06320415e+00 2.93426068e-01 + -1.20360141e+00 -1.00301288e+00 + 1.95926629e+00 6.26643532e-01 + 6.04483978e-02 5.72643059e-01 + -1.04568563e+00 -5.91021496e-01 + 2.62300678e+00 9.50997831e-01 + -4.04610275e-01 3.73150879e-01 + 2.26371902e+00 8.73627529e-01 + 2.12545313e+00 7.90640352e-01 + 7.72181917e-03 1.65718952e-02 + 1.00422340e-01 -2.05562936e-01 + -1.22989802e+00 -1.01841681e-01 + 3.09064082e+00 1.04288010e+00 + 5.18274167e+00 1.34749259e+00 + -8.32075153e-01 -1.97592029e-01 + 3.84126764e-02 5.58171345e-01 + 4.99560727e-01 -4.26154438e-02 + 4.79071151e+00 2.19728942e+00 + -2.78437968e+00 -1.17812590e+00 + -2.22804226e+00 -4.31174255e-01 + 8.50762292e-01 -1.06445261e-01 + 1.10812830e+00 -2.59118812e-01 + -2.91450155e-01 6.42802679e-01 + -1.38631532e-01 -5.88585623e-01 + -5.04120983e-01 -2.17094915e-01 + 3.41410820e+00 1.67897767e+00 + -2.23697326e+00 -6.62735244e-01 + -3.55961064e-01 -1.27647226e-01 + -3.55568274e+00 -2.49011369e+00 + -8.77586408e-01 -9.38268065e-03 + 1.52382384e-01 -5.62155760e-01 + 1.55885574e-01 1.07617069e-01 + -8.37129973e-01 -5.22259081e-01 + -2.92741750e+00 -1.35049428e+00 + -3.54670781e-01 5.69205952e-02 + 2.21030255e+00 1.34689986e+00 + 1.60787722e+00 5.75984706e-01 + 1.32294221e+00 5.31577509e-01 + 7.05672928e-01 3.34241244e-01 + 1.41406179e+00 1.15783408e+00 + -6.92172228e-01 -2.84817896e-01 + 3.28358655e-01 -2.66910083e-01 + 1.68013644e-01 -4.28016549e-02 + 2.07365974e+00 7.76496211e-01 + -3.92974907e-01 2.46796730e-01 + -5.76078636e-01 3.25676963e-01 + -1.82547204e-01 -5.06410543e-01 + 3.04754906e+00 1.16174496e+00 + -3.01090632e+00 -1.09195183e+00 + -1.44659696e+00 -6.87838682e-01 + 2.11395861e+00 9.10495785e-01 + 1.40962871e+00 1.13568678e+00 + -1.66653234e-01 -2.10012503e-01 + 3.17456029e+00 9.74502922e-01 + 2.15944820e+00 8.62807189e-01 + -3.45418719e+00 -1.33647548e+00 + -3.41357732e+00 -8.47048920e-01 + -3.06702448e-01 -6.64280634e-01 + -2.86930714e-01 -1.35268264e-01 + -3.15835557e+00 -5.43439253e-01 + 2.49541440e-01 -4.71733570e-01 + 2.71933912e+00 4.13308399e-01 + -2.43787038e+00 -1.08050547e+00 + -4.90234490e-01 -6.64069865e-01 + 8.99524451e-02 5.76180541e-01 + 5.00500404e+00 2.12125521e+00 + -1.73107940e-01 -2.28506575e-02 + 5.44938858e-01 -1.29523352e-01 + 5.13526842e+00 1.68785993e+00 + 1.70228304e+00 1.02601138e+00 + 3.58957507e+00 1.54396196e+00 + 1.85615738e+00 4.92916197e-01 + 2.55772147e+00 7.88438908e-01 + -1.57008279e+00 -4.17377300e-01 + -1.42548604e+00 -3.63684860e-01 + -8.52026118e-01 2.72052686e-01 + -5.10563077e+00 -2.35665994e+00 + -2.95517031e+00 -1.84945297e+00 + -2.91947959e+00 -1.66016784e+00 + -4.21462387e+00 -1.41131535e+00 + 6.59901121e-01 4.87156314e-01 + -9.75352532e-01 -4.50231285e-01 + -5.94084444e-01 -1.16922670e+00 + 7.50554615e-01 -9.83692552e-01 + 1.07054926e+00 2.77143030e-01 + -3.88079578e-01 -4.17737309e-02 + -9.59373733e-01 -8.85454886e-01 + -7.53560665e-02 -5.16223870e-02 + 9.84108158e-01 -5.89290700e-02 + 1.87272961e-01 -4.34238391e-01 + 6.86509981e-01 -3.15116460e-01 + -1.07762538e+00 6.58984161e-02 + 6.09266592e-01 6.91808473e-02 + -8.30529954e-01 -7.00454791e-01 + -9.13179464e-01 -6.31712891e-01 + 7.68744851e-01 1.09840676e+00 + -1.07606690e+00 -8.78390282e-01 + -1.71038184e+00 -5.73606033e-01 + 8.75982765e-01 3.66343143e-01 + -7.04919009e-01 -8.49182590e-01 + -1.00274668e+00 -7.99573611e-01 + -1.05562848e+00 -5.84060076e-01 + 4.03490015e+00 1.28679206e+00 + -3.53484804e+00 -1.71381255e+00 + 2.31527363e-01 1.04179397e-01 + -3.58592392e-02 3.74895739e-01 + 3.92253428e+00 1.81852726e+00 + -7.27384249e-01 -6.45605128e-01 + 4.65678097e+00 2.41379899e+00 + 1.16750534e+00 7.60718205e-01 + 1.15677059e+00 7.96225550e-01 + -1.42920261e+00 -4.66946295e-01 + 3.71148192e+00 1.88060191e+00 + 2.44052407e+00 3.84472199e-01 + -1.64535035e+00 -8.94530036e-01 + -3.69608753e+00 -1.36402754e+00 + 2.24419208e+00 9.69744889e-01 + 2.54822427e+00 1.22613039e+00 + 3.77484909e-01 -5.98521878e-01 + -3.61521175e+00 -1.11123912e+00 + 3.28113127e+00 1.52551775e+00 + -3.51030902e+00 -1.53913980e+00 + -2.44874505e+00 -6.30246005e-01 + -3.42516153e-01 -5.07352665e-01 + 1.09110502e+00 6.36821628e-01 + -2.49434967e+00 -8.02827146e-01 + 1.41763139e+00 -3.46591820e-01 + 1.61108619e+00 5.93871102e-01 + 3.97371717e+00 1.35552499e+00 + -1.33437177e+00 -2.83908670e-01 + -1.41606483e+00 -1.76402601e-01 + 2.23945322e-01 -1.77157065e-01 + 2.60271569e+00 2.40778251e-01 + -2.82213895e-02 1.98255474e-01 + 4.20727940e+00 1.31490863e+00 + 3.36944889e+00 1.57566635e+00 + 3.53049396e+00 1.73579350e+00 + -1.29170202e+00 -1.64196290e+00 + 9.27295604e-01 9.98808036e-01 + 1.75321843e-01 -2.83267817e-01 + -2.19069578e+00 -1.12814358e+00 + 1.66606031e+00 7.68006933e-01 + -7.13826035e-01 5.20881684e-02 + -3.43821888e+00 -2.36137021e+00 + -5.93210310e-01 1.21843813e-01 + -4.09800822e+00 -1.39893953e+00 + 2.74110954e+00 1.52728606e+00 + 1.72652512e+00 -1.25435113e-01 + 1.97722357e+00 6.40667481e-01 + 4.18635780e-01 3.57018509e-01 + -1.78303569e+00 -2.11864764e-01 + -3.52809366e+00 -2.58794450e-01 + -4.72407090e+00 -1.63870734e+00 + 1.73917807e+00 8.73251829e-01 + 4.37979356e-01 8.49210569e-01 + 3.93791881e+00 1.76269490e+00 + 2.79065411e+00 1.04019042e+00 + -8.47426142e-01 -3.40136892e-01 + -4.24389181e+00 -1.80253120e+00 + -1.86675870e+00 -7.64558265e-01 + 9.46212675e-01 -7.77681445e-02 + -2.82448462e+00 -1.33592449e+00 + -2.57938567e+00 -1.56554690e+00 + -2.71615767e+00 -6.27667233e-01 + -1.55999166e+00 -5.81013466e-01 + -4.24696864e-01 -7.44673250e-01 + 1.67592970e+00 7.68164292e-01 + 8.48455216e-01 -6.05681126e-01 + 6.12575454e+00 1.65607584e+00 + 1.38207327e+00 2.39261863e-01 + 3.13364450e+00 1.17154698e+00 + 1.71694858e+00 1.26744905e+00 + -1.61746367e+00 -8.80098073e-01 + -8.52196756e-01 -9.27299728e-01 + -1.51562462e-01 -8.36552490e-02 + -7.04792753e-01 -1.24726713e-02 + -3.35265757e+00 -1.82176312e+00 + 3.32173170e-01 -1.33405580e-01 + 4.95841013e-01 4.58292712e-01 + 1.57713955e+00 7.79272991e-01 + 2.09743109e+00 9.23542557e-01 + 3.90450311e-03 -8.42873164e-01 + 2.59519038e+00 7.56479591e-01 + -5.77643976e-01 -2.36401904e-01 + -5.22310654e-01 1.34187830e-01 + -2.22096086e+00 -7.75507719e-01 + 1.35907831e+00 7.80197510e-01 + 3.80355868e+00 1.16983476e+00 + 3.82746596e+00 1.31417718e+00 + 3.30451183e+00 1.55398159e+00 + -3.42917814e-01 -8.62281222e-02 + -2.59093020e+00 -9.29883526e-01 + 1.40928562e+00 1.08398346e+00 + 1.54400137e-01 3.35881092e-01 + 1.59171586e+00 1.18855802e+00 + -5.25164002e-01 -1.03104220e-01 + 2.20067959e+00 1.37074713e+00 + 6.97860830e-01 6.27718548e-01 + -4.59743507e-01 1.36061163e-01 + -1.04691963e-01 -2.16271727e-01 + -1.08905573e+00 -5.95510769e-01 + -1.00826983e+00 -5.38509162e-02 + -3.16402719e+00 -1.33414216e+00 + 1.47870874e-01 1.75234619e-01 + -2.57078234e-01 7.03316889e-02 + 1.81073945e+00 4.26901462e-01 + 2.65476530e+00 6.74217273e-01 + 1.27539811e+00 6.22914081e-01 + -3.76750499e-01 -1.20629449e+00 + 1.00177595e+00 -1.40660091e-01 + -2.98919265e+00 -1.65145013e+00 + -2.21557682e+00 -8.11123452e-01 + -3.22635378e+00 -1.65639056e+00 + -2.72868553e+00 -1.02812087e+00 + 1.26042797e+00 8.49005248e-01 + -9.38318534e-01 -9.87588651e-01 + 3.38013194e-01 -1.00237461e-01 + 1.91175691e+00 8.48716369e-01 + 4.30244344e-01 6.05539915e-02 + 2.21783435e+00 3.03268204e-01 + 1.78019576e+00 1.27377108e+00 + 1.59733274e+00 4.40674687e-02 + 3.97428484e+00 2.20881566e+00 + -2.41108677e+00 -6.01410418e-01 + -2.50796499e+00 -5.71169866e-01 + -3.71957427e+00 -1.38195726e+00 + -1.57992670e+00 1.32068593e-01 + -1.35278851e+00 -6.39349270e-01 + 1.23075932e+00 2.40445409e-01 + 1.35606530e+00 4.33180078e-01 + 9.60968518e-02 2.26734255e-01 + 6.22975063e-01 5.03431915e-02 + -1.47624851e+00 -3.60568238e-01 + -2.49337808e+00 -1.15083052e+00 + 2.15717792e+00 1.03071559e+00 + -3.07814376e-02 1.38700314e-02 + 4.52049499e-02 -4.86409775e-01 + 2.58231061e+00 1.14327809e-01 + 1.10999138e+00 -5.18568405e-01 + -2.19426443e-01 -5.37505538e-01 + -4.44740298e-01 6.78099955e-01 + 4.03379080e+00 1.49825720e+00 + -5.13182408e-01 -4.90201950e-01 + -6.90139716e-01 1.63875126e-01 + -8.17281461e-01 2.32155064e-01 + -2.92357619e-01 -8.02573544e-01 + -1.80769841e+00 -7.58907326e-01 + 2.16981590e+00 1.06728873e+00 + 1.98995203e-01 -6.84176682e-02 + -2.39546753e+00 -2.92873789e-01 + -4.24251021e+00 -1.46255564e+00 + -5.01411291e-01 -5.95712813e-03 + 2.68085809e+00 1.42883780e+00 + -4.13289873e+00 -1.62729388e+00 + 1.87957843e+00 3.63341638e-01 + -1.15270744e+00 -3.03563774e-01 + -4.43994248e+00 -2.97323905e+00 + -7.17067733e-01 -7.08349542e-01 + -3.28870393e+00 -1.19263863e+00 + -7.55325944e-01 -5.12703329e-01 + -2.07291938e+00 -2.65025085e-01 + -7.50073814e-01 -1.70771041e-01 + -8.77381404e-01 -5.47417325e-01 + -5.33725862e-01 5.15837119e-01 + 8.45056431e-01 2.82125560e-01 + -1.59598637e+00 -1.38743235e+00 + 1.41362902e+00 1.06407789e+00 + 1.02584504e+00 -3.68219466e-01 + -1.04644488e+00 -1.48769392e-01 + 2.66990191e+00 8.57633492e-01 + -1.84251857e+00 -9.82430175e-01 + 9.71404204e-01 -2.81934209e-01 + -2.50177989e+00 -9.21260335e-01 + -1.31060074e+00 -5.84488113e-01 + -2.12129400e-01 -3.06244708e-02 + -5.28933882e+00 -2.50663129e+00 + 1.90220541e+00 1.08662918e+00 + -3.99366086e-02 -6.87178973e-01 + -4.93417342e-01 4.37354182e-01 + 2.13494486e+00 1.37679569e+00 + 2.18396765e+00 5.81023868e-01 + -3.07866587e+00 -1.45384974e+00 + 6.10894119e-01 -4.17050124e-01 + -1.88766952e+00 -8.86160058e-01 + 3.34527253e+00 1.78571260e+00 + 6.87769059e-01 -5.01157336e-01 + 2.60470837e+00 1.45853560e+00 + -6.49315691e-01 -9.16112805e-01 + -1.29817687e+00 -2.15924339e-01 + -1.20100409e-03 -4.03137422e-01 + -1.36471594e+00 -6.93266356e-01 + 1.38682062e+00 7.15131598e-01 + 2.47830103e+00 1.24862305e+00 + -2.78288147e+00 -1.03329235e+00 + -7.33443403e-01 -6.11041652e-01 + -4.12745671e-01 -5.96133390e-02 + -2.58632336e+00 -4.51557058e-01 + -1.16570367e+00 -1.27065510e+00 + 2.76187104e+00 2.21895451e-01 + -3.80443767e+00 -1.66319902e+00 + 9.84658633e-01 6.81475569e-01 + 9.33814584e-01 -4.89335563e-02 + -4.63427997e-01 1.72989539e-01 + 1.82401546e+00 3.60164021e-01 + -5.36521077e-01 -8.08691351e-01 + -1.37367030e+00 -1.02126160e+00 + -3.70310682e+00 -1.19840844e+00 + -1.51894242e+00 -3.89510223e-01 + -3.67347940e-01 -3.25540516e-02 + -1.00988595e+00 1.82802194e-01 + 2.01622795e+00 7.86367901e-01 + 1.02440231e+00 8.79780360e-01 + -3.05971480e+00 -8.40901527e-01 + 2.73909457e+00 1.20558628e+00 + 2.39559056e+00 1.10786694e+00 + 1.65471544e+00 7.33824651e-01 + 2.18546787e+00 6.41168955e-01 + 1.47152266e+00 3.91839132e-01 + 1.45811155e+00 5.21820495e-01 + -4.27531469e-02 -3.52343068e-03 + -9.54948010e-01 -1.52313876e-01 + 7.57151215e-01 -5.68728854e-03 + -8.46205751e-01 -7.54580229e-01 + 4.14493548e+00 1.45532780e+00 + 4.58688968e-01 -4.54012803e-02 + -1.49295381e+00 -4.57471758e-01 + 1.80020351e+00 8.13724973e-01 + -5.82727738e+00 -2.18269581e+00 + -2.09017809e+00 -1.18305177e+00 + -2.31628303e+00 -7.21600235e-01 + -8.09679091e-01 -1.49101752e-01 + 8.88005605e-01 8.57940857e-01 + -1.44148219e+00 -3.10926299e-01 + 3.68828186e-01 -3.08848059e-01 + -6.63267389e-01 -8.58950139e-02 + -1.14702569e+00 -6.32147854e-01 + -1.51741715e+00 -8.53330564e-01 + -1.33903718e+00 -1.45875547e-01 + 4.12485387e+00 1.85620435e+00 + -2.42353639e+00 -2.92669850e-01 + 1.88708583e+00 9.35984730e-01 + 2.15585179e+00 6.30469051e-01 + -1.13627973e-01 -1.62554045e-01 + 2.04540494e+00 1.36599834e+00 + 2.81591381e+00 1.60897941e+00 + 3.02736260e-02 3.83255815e-03 + 7.97634013e-02 -2.82035099e-01 + -3.24607473e-01 -5.30065956e-01 + -3.91862894e+00 -1.94083334e+00 + 1.56360901e+00 7.93882743e-01 + -1.03905772e+00 6.25590229e-01 + 2.54746492e+00 1.64233560e+00 + -4.80774423e-01 -8.92298032e-02 + 9.06979990e-02 1.05020427e+00 + -2.47521290e+00 -1.78275982e-01 + -3.91871729e-01 3.80285423e-01 + 1.00658382e+00 4.58947483e-01 + 4.68102941e-01 1.02992741e+00 + 4.44242568e-01 2.89870239e-01 + 3.29684452e+00 1.44677474e+00 + -2.24983007e+00 -9.65574499e-01 + -3.54453926e-01 -3.99020325e-01 + -3.87429665e+00 -1.90079739e+00 + 2.02656674e+00 1.12444894e+00 + 3.77011621e+00 1.43200852e+00 + 1.61259275e+00 4.65417399e-01 + 2.28725434e+00 6.79181395e-01 + 2.75421009e+00 2.27327345e+00 + -2.40894409e+00 -1.03926359e+00 + 1.52996651e-01 -2.73373046e-02 + -2.63218977e+00 -7.22802821e-01 + 2.77688169e+00 1.15310186e+00 + 1.18832341e+00 4.73457165e-01 + -2.35536326e+00 -1.08034554e+00 + -5.84221627e-01 1.03505984e-02 + 2.96730300e+00 1.33478306e+00 + -8.61947692e-01 6.09137051e-02 + 8.22343921e-01 -8.14155286e-02 + 1.75809015e+00 1.07921470e+00 + 1.19501279e+00 1.05309972e+00 + -1.75901792e+00 9.75320161e-02 + 1.64398635e+00 9.54384323e-01 + -2.21878052e-01 -3.64847144e-01 + -2.03128968e+00 -8.57866419e-01 + 1.86750633e+00 7.08524487e-01 + 8.03972976e-01 3.47404314e-01 + 3.41203749e+00 1.39810900e+00 + 4.22397681e-01 -6.41440488e-01 + -4.88493360e+00 -1.58967816e+00 + -1.67649284e-01 -1.08485915e-01 + 2.11489023e+00 1.50506158e+00 + -1.81639929e+00 -3.85542192e-01 + 2.24044819e-01 -1.45100577e-01 + -3.39262411e+00 -1.44394324e+00 + 1.68706599e+00 2.29199618e-01 + -1.94093257e+00 -1.65975814e-01 + 8.28143367e-01 5.92109281e-01 + -8.29587998e-01 -9.57130831e-01 + -1.50011401e+00 -8.36802092e-01 + 2.40770449e+00 9.32820177e-01 + 7.41391309e-02 3.12878473e-01 + 1.87745264e-01 6.19231425e-01 + 9.57622692e-01 -2.20640033e-01 + 3.18479243e+00 1.02986233e+00 + 2.43133846e+00 8.41302677e-01 + -7.09963834e-01 1.99718943e-01 + -2.88253498e-01 -3.62772094e-01 + 5.14052574e+00 1.79304595e+00 + -3.27930993e+00 -1.29177973e+00 + -1.16723536e+00 1.29519656e-01 + 1.04801056e+00 3.41508300e-01 + -3.99256195e+00 -2.51176471e+00 + -7.62824318e-01 -6.84242153e-01 + 2.71524986e-02 5.35157164e-02 + 3.26430102e+00 1.34887262e+00 + -1.72357766e+00 -4.94524388e-01 + -3.81149536e+00 -1.28121944e+00 + 3.36919354e+00 1.10672075e+00 + -3.14841757e+00 -7.10713767e-01 + -3.16463676e+00 -7.58558435e-01 + -2.44745969e+00 -1.08816514e+00 + 2.79173264e-01 -2.19652051e-02 + 4.15309883e-01 6.07502790e-01 + -9.51007417e-01 -5.83976336e-01 + -1.47929839e+00 -8.39850409e-01 + 2.38335703e+00 6.16055149e-01 + -7.47749031e-01 -5.56164928e-01 + -3.65643622e-01 -5.06684411e-01 + -1.76634163e+00 -7.86382097e-01 + 6.76372222e-01 -3.06592181e-01 + -1.33505058e+00 -1.18301441e-01 + 3.59660179e+00 2.00424178e+00 + -7.88912762e-02 8.71956146e-02 + 1.22656397e+00 1.18149583e+00 + 4.24919729e+00 1.20082355e+00 + 2.94607456e+00 1.00676505e+00 + 7.46061275e-02 4.41761753e-02 + -2.47738025e-02 1.92737701e-01 + -2.20509316e-01 -3.79163193e-01 + -3.50222190e-01 3.58727299e-01 + -3.64788014e+00 -1.36107312e+00 + 3.56062799e+00 9.27032742e-01 + 1.04317289e+00 6.08035970e-01 + 4.06718718e-01 3.00628051e-01 + 4.33158086e+00 2.25860714e+00 + 2.13917145e-01 -1.72757967e-01 + -1.40637998e+00 -1.14119465e+00 + 3.61554872e+00 1.87797348e+00 + 1.01726871e+00 5.70255097e-01 + -7.04902551e-01 2.16444147e-01 + -2.51492186e+00 -8.52997369e-01 + 1.85097530e+00 1.15124496e+00 + -8.67569714e-01 -3.05682432e-01 + 8.07550858e-01 5.88901608e-01 + 1.85186755e-01 -1.94589367e-01 + -1.23378238e+00 -7.84128347e-01 + -1.22713161e+00 -4.21218235e-01 + 2.97751165e-01 2.81055275e-01 + 4.77703554e+00 1.66265524e+00 + 2.51549669e+00 7.49980674e-01 + 2.76510822e-01 1.40456909e-01 + 1.98740905e+00 -1.79608212e-01 + 9.35429145e-01 8.44344180e-01 + -1.20854492e+00 -5.00598453e-01 + 2.29936219e+00 8.10236668e-01 + 6.92555544e-01 -2.65891331e-01 + -1.58050994e+00 2.31237821e-01 + -1.50864880e+00 -9.49661690e-01 + -1.27689206e+00 -7.18260016e-01 + -3.12517127e+00 -1.75587113e+00 + 8.16062912e-02 -6.56551804e-01 + -5.02479939e-01 -4.67162543e-01 + -5.47435788e+00 -2.47799576e+00 + 1.95872901e-02 5.80874076e-01 + -1.59064958e+00 -6.34554756e-01 + -3.77521478e+00 -1.74301790e+00 + 5.89628224e-01 8.55736553e-01 + -1.81903543e+00 -7.50011008e-01 + 1.38557775e+00 3.71490991e-01 + 9.70032652e-01 -7.11356016e-01 + 2.63539625e-01 -4.20994771e-01 + 2.12154222e+00 8.19081400e-01 + -6.56977937e-01 -1.37810098e-01 + 8.91309581e-01 2.77864361e-01 + -7.43693195e-01 -1.46293770e-01 + 2.24447769e+00 4.00911438e-01 + -2.25169262e-01 2.04148801e-02 + 1.68744684e+00 9.47573007e-01 + 2.73086373e-01 3.30877195e-01 + 5.54294414e+00 2.14198009e+00 + -8.49238733e-01 3.65603298e-02 + 2.39685712e+00 1.17951039e+00 + -2.58230528e+00 -5.52116673e-01 + 2.79785277e+00 2.88833717e-01 + -1.96576188e-01 1.11652123e+00 + -4.69383301e-01 1.96496282e-01 + -1.95011845e+00 -6.15235169e-01 + 1.03379890e-02 2.33701239e-01 + 4.18933607e-01 2.77939814e-01 + -1.18473337e+00 -4.10051126e-01 + -7.61499744e-01 -1.43658094e+00 + -1.65586092e+00 -3.41615303e-01 + -5.58523700e-02 -5.21837080e-01 + -2.40331088e+00 -2.64521583e-01 + 2.24925206e+00 6.79843335e-02 + 1.46360479e+00 1.04271443e+00 + -3.09255443e+00 -1.82548953e+00 + 2.11325841e+00 1.14996627e+00 + -8.70657797e-01 1.02461839e-01 + -5.71056521e-01 9.71232588e-02 + -3.37870752e+00 -1.54091877e+00 + 1.03907189e+00 -1.35661392e-01 + 8.40057486e-01 6.12172413e-02 + -1.30998234e+00 -1.34077226e+00 + 7.53744974e-01 1.49447350e-01 + 9.13995056e-01 -1.81227962e-01 + 2.28386229e-01 3.74498520e-01 + 2.54829151e-01 -2.88802704e-01 + 1.61709009e+00 2.09319193e-01 + -1.12579380e+00 -5.95955338e-01 + -2.69610726e+00 -2.76222736e-01 + -2.63773329e+00 -7.84491970e-01 + -2.62167427e+00 -1.54792874e+00 + -4.80639856e-01 -1.30582102e-01 + -1.26130891e+00 -8.86841840e-01 + -1.24951950e+00 -1.18182622e+00 + -1.40107574e+00 -9.13695575e-01 + 4.99872179e-01 4.69014702e-01 + -2.03550193e-02 -1.48859738e-01 + -1.50189069e+00 -2.97714278e-02 + -2.07846113e+00 -7.29937809e-01 + -5.50576792e-01 -7.03151525e-01 + -3.88069238e+00 -1.63215295e+00 + 2.97032988e+00 6.43571144e-01 + -1.85999273e-01 1.18107620e+00 + 1.79249709e+00 6.65356160e-01 + 2.68842472e+00 1.35703255e+00 + 1.07675417e+00 1.39845588e-01 + 8.01226349e-01 2.11392275e-01 + 9.64329379e-01 3.96146195e-01 + -8.22529511e-01 1.96080831e-01 + 1.92481841e+00 4.62985744e-01 + 3.69756927e-01 3.77135799e-01 + 1.19807835e+00 8.87715050e-01 + -1.01363587e+00 -2.48151636e-01 + 8.53071010e-01 4.96887868e-01 + -3.41120553e+00 -1.35401843e+00 + -2.64787381e+00 -1.08690563e+00 + -1.11416759e+00 -4.43848915e-01 + 1.46242648e+00 6.17106076e-02 + -7.52968881e-01 -9.20972209e-01 + -1.22492228e+00 -5.40327617e-01 + 1.08001827e+00 5.29593785e-01 + -2.58706464e-01 1.13022085e-01 + -4.27394011e-01 1.17864354e-02 + -3.20728413e+00 -1.71224737e-01 + 1.71398530e+00 8.68885893e-01 + 2.12067866e+00 1.45092772e+00 + 4.32782616e-01 -3.34117769e-01 + 7.80084374e-01 -1.35100217e-01 + -2.05547729e+00 -4.70217750e-01 + 2.38379736e+00 1.09186058e+00 + -2.80825477e+00 -1.03320187e+00 + 2.63434576e+00 1.15671733e+00 + -1.60936214e+00 1.91843035e-01 + -5.02298769e+00 -2.32820708e+00 + 1.90349195e+00 1.45215416e+00 + 3.00232888e-01 3.24412586e-01 + -2.46503943e+00 -1.19550010e+00 + 1.06304233e+00 2.20136246e-01 + -2.99101388e+00 -1.58299318e+00 + 2.30071719e+00 1.12881362e+00 + -2.37587247e+00 -8.08298336e-01 + 7.27006308e-01 3.80828984e-01 + 2.61199061e+00 1.56473491e+00 + 8.33936357e-01 -1.42189425e-01 + 3.13291605e+00 1.77771210e+00 + 2.21917371e+00 5.68427075e-01 + 2.38867649e+00 9.06637262e-01 + -6.92959466e+00 -3.57682881e+00 + 2.57904824e+00 5.93959108e-01 + 2.71452670e+00 1.34436199e+00 + 4.39988761e+00 2.13124672e+00 + 5.71783077e-01 5.08346173e-01 + -3.65399429e+00 -1.18192861e+00 + 4.46176453e-01 3.75685594e-02 + -2.97501495e+00 -1.69459236e+00 + 1.60855728e+00 9.20930014e-01 + -1.44270290e+00 -1.93922306e-01 + 1.67624229e+00 1.66233866e+00 + -1.42579598e+00 -1.44990145e-01 + 1.19923176e+00 4.58490278e-01 + -9.00068460e-01 5.09701825e-02 + -1.69391694e+00 -7.60070300e-01 + -1.36576440e+00 -5.24244256e-01 + -1.03016748e+00 -3.44625878e-01 + 2.40519313e+00 1.09947587e+00 + 1.50365433e+00 1.06464802e+00 + -1.07609727e+00 -3.68897187e-01 + 2.44969069e+00 1.28486192e+00 + -1.25610307e+00 -1.14644789e+00 + 2.05962899e+00 4.31162369e-01 + -7.15886908e-01 -6.11587804e-02 + -6.92354119e-01 -7.85019920e-01 + -1.63016508e+00 -5.96944975e-01 + 1.90352536e+00 1.28197457e+00 + -4.01535243e+00 -1.81934488e+00 + -1.07534435e+00 -2.10544784e-01 + 3.25500866e-01 7.69603661e-01 + 2.18443365e+00 6.59773335e-01 + 8.80856790e-01 6.39505913e-01 + -2.23956372e-01 -4.65940132e-01 + -1.06766519e+00 -5.38388505e-03 + 7.25556863e-01 -2.91123488e-01 + -4.69451411e-01 7.89182650e-02 + 2.58146587e+00 1.29653243e+00 + 1.53747468e-01 7.69239075e-01 + -4.61152262e-01 -4.04151413e-01 + 1.48183517e+00 8.10079506e-01 + -1.83402614e+00 -1.36939322e+00 + 1.49315501e+00 7.95225425e-01 + 1.41922346e+00 1.05582774e-01 + 1.57473493e-01 9.70795657e-01 + -2.67603254e+00 -7.48562280e-01 + -8.49156216e-01 -6.05762529e-03 + 1.12944274e+00 3.67741591e-01 + 1.94228071e-01 5.28188141e-01 + -3.65610158e-01 4.05851838e-01 + -1.98839111e+00 -1.38452764e+00 + 2.73765752e+00 8.24150530e-01 + 7.63728641e-01 3.51617707e-01 + 5.78307267e+00 1.68103612e+00 + 2.27547227e+00 3.60876164e-01 + -3.50681697e+00 -1.74429984e+00 + 4.01241184e+00 1.26227829e+00 + 2.44946343e+00 9.06119057e-01 + -2.96638941e+00 -9.01532322e-01 + 1.11267643e+00 -3.43333381e-01 + -6.61868994e-01 -3.44666391e-01 + -8.34917179e-01 5.69478372e-01 + -1.91888454e+00 -3.03791075e-01 + 1.50397636e+00 8.31961240e-01 + 6.12260198e+00 2.16851807e+00 + 1.34093127e+00 8.86649385e-01 + 1.48748519e+00 8.26273697e-01 + 7.62243068e-01 2.64841396e-01 + -2.17604986e+00 -3.54219958e-01 + 2.64708640e-01 -4.38136718e-02 + 1.44725372e+00 1.18499914e-01 + -6.71259446e-01 -1.19526851e-01 + 2.40134595e-01 -8.90042323e-02 + -3.57238199e+00 -1.23166201e+00 + -3.77626645e+00 -1.19533443e+00 + -3.81101035e-01 -4.94160532e-01 + -3.02758757e+00 -1.18436066e+00 + 2.59116298e-01 1.38023047e+00 + 4.17900116e+00 1.12065959e+00 + 1.54598848e+00 2.89806755e-01 + 1.00656475e+00 1.76974511e-01 + -4.15730234e-01 -6.22681694e-01 + -6.00903565e-01 -1.43256959e-01 + -6.03652508e-01 -5.09936379e-01 + -1.94096658e+00 -9.48789544e-01 + -1.74464105e+00 -8.50491590e-01 + 1.17652544e+00 1.88118317e+00 + 2.35507776e+00 1.44000205e+00 + 2.63067924e+00 1.06692988e+00 + 2.88805386e+00 1.23924715e+00 + 8.27595008e-01 5.75364692e-01 + 3.91384216e-01 9.72781920e-02 + -1.03866816e+00 -1.37567768e+00 + -1.34777969e+00 -8.40266025e-02 + -4.12904508e+00 -1.67618340e+00 + 1.27918111e+00 3.52085961e-01 + 4.15361174e-01 6.28896189e-01 + -7.00539496e-01 4.80447955e-02 + -1.62332639e+00 -5.98236485e-01 + 1.45957300e+00 1.00305154e+00 + -3.06875603e+00 -1.25897545e+00 + -1.94708176e+00 4.85143006e-01 + 3.55744156e+00 -1.07468822e+00 + 1.21602223e+00 1.28768827e-01 + 1.89093098e+00 -4.70835659e-01 + -6.55759125e+00 2.70114082e+00 + 8.96843535e-01 -3.98115252e-01 + 4.13450429e+00 -2.32069236e+00 + 2.37764218e+00 -1.09098890e+00 + -1.11388901e+00 6.27083097e-01 + -6.34116929e-01 4.62816387e-01 + 2.90203079e+00 -1.33589143e+00 + 3.17457598e+00 -5.13575945e-01 + -1.76362299e+00 5.71820693e-01 + 1.66103362e+00 -8.99466249e-01 + -2.53947433e+00 8.40084780e-01 + 4.36631397e-01 7.24234261e-02 + -1.87589394e+00 5.08529113e-01 + 4.49563965e+00 -9.43365992e-01 + 1.78876299e+00 -1.27076149e+00 + -1.16269107e-01 -4.55078316e-01 + 1.92966079e+00 -8.05371385e-01 + 2.20632583e+00 -9.00919345e-01 + 1.52387824e+00 -4.82391996e-01 + 8.04004564e-01 -2.73650595e-01 + -7.75326067e-01 1.07469566e+00 + 1.83226282e+00 -4.52173344e-01 + 1.25079758e-01 -3.52895417e-02 + -9.90957437e-01 8.55993130e-01 + 1.71623322e+00 -7.08691667e-01 + -2.86175924e+00 6.75160955e-01 + -8.40817853e-01 -1.00361809e-01 + 1.33393000e+00 -4.65788123e-01 + 5.29394114e-01 -5.44881619e-02 + -8.07435599e-01 8.27353370e-01 + -4.33165824e+00 1.97299638e+00 + 1.26452422e+00 -8.34070486e-01 + 1.45996394e-02 2.97736043e-01 + -1.64489287e+00 6.72839598e-01 + -5.74234578e+00 3.20975117e+00 + 2.13841341e-02 3.64514015e-01 + 6.68084924e+00 -2.27464254e+00 + -3.22881590e+00 8.01879324e-01 + 3.02534313e-01 -4.56222796e-01 + -5.84520734e+00 1.95678162e+00 + 2.81515232e+00 -1.72101318e+00 + -2.39620908e-01 2.69145522e-01 + -7.41669691e-01 -2.30283281e-01 + -2.15682714e+00 3.45313021e-01 + 1.23475788e+00 -7.32276553e-01 + -1.71816113e-01 1.20419560e-02 + 1.89174235e+00 2.27435901e-01 + -3.64511114e-01 1.72260361e-02 + -3.24143860e+00 6.50125817e-01 + -2.25707409e+00 5.66970751e-01 + 1.03901456e+00 -1.00588433e+00 + -5.09159710e+00 1.58736109e+00 + 1.45534075e+00 -5.83787452e-01 + 4.28879587e+00 -1.58006866e+00 + 8.52384427e-01 -1.11042299e+00 + 4.51431615e+00 -2.63844265e+00 + -4.33042648e+00 1.86497078e+00 + -2.13568046e+00 5.82559743e-01 + -4.42568887e+00 1.26131214e+00 + 3.15821315e+00 -1.61515905e+00 + -3.14125204e+00 8.49604386e-01 + 6.54152300e-01 -2.04624711e-01 + -3.73374317e-01 9.94187820e-02 + -3.96177282e+00 1.27245623e+00 + 9.59825199e-01 -1.15547861e+00 + 3.56902055e+00 -1.46591091e+00 + 1.55433633e-02 6.93544345e-01 + 1.15684646e+00 -4.99836352e-01 + 3.11824573e+00 -4.75900506e-01 + -8.61706369e-01 -3.50774059e-01 + 9.89057391e-01 -7.16878802e-01 + -4.94787870e+00 2.09137481e+00 + 1.37777347e+00 -1.34946349e+00 + -1.13161577e+00 8.05114754e-01 + 8.12020675e-01 -1.04849421e+00 + 4.73783881e+00 -2.26718812e+00 + 8.99579366e-01 -8.89764451e-02 + 4.78524868e+00 -2.25795843e+00 + 1.75164590e+00 -1.73822209e-01 + 1.30204590e+00 -7.26724717e-01 + -7.26526403e-01 -5.23925361e-02 + 2.01255351e+00 -1.69965366e+00 + 9.87852740e-01 -4.63577220e-01 + 2.45957762e+00 -1.29278962e+00 + -3.13817948e+00 1.64433038e+00 + -1.76302159e+00 9.62784302e-01 + -1.91106331e+00 5.81460008e-01 + -3.30883001e+00 1.30378978e+00 + 5.54376450e-01 3.78814272e-01 + 1.09982111e+00 -1.47969612e+00 + -2.61300705e-02 -1.42573464e-01 + -2.22096157e+00 7.75684440e-01 + 1.70319323e+00 -2.89738444e-01 + -1.43223842e+00 6.39284281e-01 + 2.34360959e-01 -1.64379268e-01 + -2.67147991e+00 9.46548086e-01 + 1.51131425e+00 -4.91594395e-01 + -2.48446856e+00 1.01286123e+00 + 1.50534658e-01 -2.94620246e-01 + -1.66966792e+00 1.67755508e+00 + -1.50094241e+00 3.30163095e-01 + 2.27681194e+00 -1.08064317e+00 + 2.05122965e+00 -1.15165939e+00 + -4.23509309e-01 -6.56906167e-02 + 1.80084023e+00 -1.07228556e+00 + -2.65769521e+00 1.18023206e+00 + 2.02852676e+00 -8.06793574e-02 + -4.49544185e+00 2.68200163e+00 + -7.50043216e-01 1.17079331e+00 + 6.80060893e-02 3.99055351e-01 + -3.83634635e+00 1.38406887e+00 + 3.24858545e-01 -9.25273218e-02 + -2.19895100e+00 1.47819500e+00 + -3.61569522e-01 -1.03188739e-01 + 1.12180375e-01 -9.52696354e-02 + -1.31477803e+00 1.79900570e-01 + 2.39573628e+00 -6.09739269e-01 + -1.00135700e+00 6.02837296e-01 + -4.11994589e+00 2.49599192e+00 + -1.54196236e-01 -4.84921951e-01 + 5.92569908e-01 -1.87310359e-01 + 3.85407741e+00 -1.50979925e+00 + 5.17802528e+00 -2.26032607e+00 + -1.37018916e+00 1.87111822e-01 + 8.46682996e-01 -3.56676331e-01 + -1.17559949e+00 5.29057734e-02 + -5.56475671e-02 6.79049243e-02 + 1.07851745e+00 -5.14535101e-01 + -2.71622446e+00 1.00151846e+00 + -1.08477208e+00 8.81391054e-01 + 5.50755824e-01 -5.20577727e-02 + 4.70885495e+00 -2.04220397e+00 + -1.87375336e-01 -6.16962830e-02 + 3.52097100e-01 2.21163550e-01 + 7.07929984e-01 -1.75827590e-01 + -1.22149219e+00 1.83084346e-01 + 2.58247412e+00 -6.15914898e-01 + -6.01206182e-01 -2.29832987e-01 + 9.83360449e-01 -3.75870060e-01 + -3.20027685e+00 1.35467480e+00 + 1.79178978e+00 -1.38531981e+00 + -3.30376867e-01 -1.16250192e-01 + -1.89053055e+00 5.68463567e-01 + -4.20604849e+00 1.65429681e+00 + -1.01185529e+00 1.92801240e-01 + -6.18819882e-01 5.42206996e-01 + -5.08091672e+00 2.61598591e+00 + -2.62570344e+00 2.51590658e+00 + 3.05577906e+00 -1.49090609e+00 + 2.77609677e+00 -1.37681378e+00 + -7.93515301e-02 4.28072744e-01 + -2.08359471e+00 8.94334295e-01 + 2.20163801e+00 4.01127167e-02 + -1.18145785e-01 -2.06822464e-01 + -2.74788298e-01 2.96250607e-01 + 1.59613555e+00 -3.87246203e-01 + -3.82971472e-01 -3.39716093e-02 + -4.20311307e-02 3.88529510e-01 + 1.52128574e+00 -9.33138876e-01 + -9.06584458e-01 -2.75016094e-02 + 3.56216834e+00 -9.99384622e-01 + 2.11964220e+00 -9.98749118e-02 + 4.01203480e+00 -2.03032745e+00 + -1.24171557e+00 1.97596725e-01 + -1.57230455e+00 4.14126609e-01 + -1.85484741e+00 5.40041563e-01 + 1.76329831e+00 -6.95967734e-01 + -2.29439232e-01 5.08669245e-01 + -5.45124276e+00 2.26907549e+00 + -5.71364288e-02 5.04476476e-01 + 3.12468018e+00 -1.46358879e+00 + 8.20017359e-01 6.51949028e-01 + -1.33977500e+00 2.83634232e-04 + -1.83311685e+00 1.23947117e+00 + 6.31205922e-01 1.19792164e-02 + -2.21967834e+00 6.94056232e-01 + -1.41693842e+00 9.93526233e-01 + -7.58885703e-01 6.78547347e-01 + 3.60239086e+00 -1.08644935e+00 + 6.72217073e-02 3.00036011e-02 + -3.42680958e-01 -3.48049352e-01 + 1.87546079e+00 -4.78018246e-01 + 7.00485821e-01 -3.52905383e-01 + -8.54580948e-01 8.17330861e-01 + 8.19123706e-01 -5.73927281e-01 + 2.70855639e-01 -3.08940052e-01 + -1.05059952e+00 3.27873168e-01 + 1.08282999e+00 4.84559349e-02 + -7.89899220e-01 1.22291138e+00 + -2.87939816e+00 7.17403497e-01 + -2.08429452e+00 8.87409226e-01 + 1.58409232e+00 -4.74123532e-01 + 1.26882735e+00 1.59162510e-01 + -2.53782993e+00 6.18253491e-01 + -8.92757445e-01 3.35979011e-01 + 1.31867900e+00 -1.17355054e+00 + 1.14918879e-01 -5.35184038e-01 + -1.70288738e-01 5.35868087e-02 + 4.21355121e-01 5.41848690e-02 + 2.07926943e+00 -5.72538144e-01 + 4.08788970e-01 3.77655777e-01 + -3.39631381e+00 9.84216764e-01 + 2.94170163e+00 -1.83120916e+00 + -7.94798752e-01 7.39889052e-01 + 1.46555463e+00 -4.62275563e-01 + 2.57255955e+00 -1.04671434e+00 + 8.45042540e-01 -1.96952892e-01 + -3.23526646e+00 1.60049846e+00 + 3.21948565e+00 -8.88376674e-01 + 1.43005104e+00 -9.21561086e-01 + 8.82360506e-01 2.98403872e-01 + -8.91168097e-01 1.01319072e+00 + -5.13215241e-01 -2.47182649e-01 + -1.35759444e+00 7.07450608e-02 + -4.04550983e+00 2.23534867e+00 + 1.39348883e+00 3.81637747e-01 + -2.85676418e+00 1.53240862e+00 + -1.37183120e+00 6.37977425e-02 + -3.88195859e+00 1.73887145e+00 + 1.19509776e+00 -6.25013512e-01 + -2.80062734e+00 1.79840585e+00 + 1.96558429e+00 -4.70997234e-01 + 1.93111352e+00 -9.70318441e-01 + 3.57991190e+00 -1.65065116e+00 + 2.12831714e+00 -1.11531708e+00 + -3.95661018e-01 -8.54339904e-02 + -2.41630441e+00 1.65166304e+00 + 7.55412624e-01 -1.53453579e-01 + -1.77043450e+00 1.39928715e+00 + -9.32631260e-01 8.73649199e-01 + 1.53342205e+00 -8.39569765e-01 + -6.29846924e-02 1.25023084e-01 + 3.31509049e+00 -1.10733235e+00 + -2.18957109e+00 3.07376993e-01 + -2.35740747e+00 6.47437564e-01 + -2.22142438e+00 8.47318938e-01 + -6.51401147e-01 3.48398562e-01 + 2.75763095e+00 -1.21390708e+00 + 1.12550484e+00 -5.61412847e-01 + -5.65053161e-01 6.74365205e-02 + 1.68952456e+00 -6.57566096e-01 + 8.95598401e-01 3.96738993e-01 + -1.86537066e+00 9.44129208e-01 + -2.59933294e+00 2.57423247e-01 + -6.59598267e-01 1.91828851e-02 + -2.64506676e+00 8.41783205e-01 + -1.25911802e+00 5.52425066e-01 + -1.39754507e+00 3.73689222e-01 + 5.49550729e-02 1.35071215e+00 + 3.31874811e+00 -1.05682424e+00 + 3.63159604e+00 -1.42864695e+00 + -4.45944617e+00 1.42889446e+00 + 5.87314342e-01 -4.88892988e-01 + -7.26130820e-01 1.51936106e-01 + -1.79246441e+00 6.05888105e-01 + -5.50948207e-01 6.21443081e-01 + -3.17246063e-01 1.77213880e-01 + -2.00098937e+00 1.23799074e+00 + 4.33790961e+00 -1.08490465e+00 + -2.03114114e+00 1.31613237e+00 + -6.29216542e+00 1.92406317e+00 + -1.60265624e+00 8.87947500e-01 + 8.64465062e-01 -8.37416270e-01 + -2.14273937e+00 8.05485900e-01 + -2.36844256e+00 6.17915124e-01 + -1.40429636e+00 6.78296866e-01 + 9.99019988e-01 -5.84297572e-01 + 7.38824546e-01 1.68838678e-01 + 1.45681238e+00 3.04641461e-01 + 2.15914949e+00 -3.43089227e-01 + -1.23895930e+00 1.05339864e-01 + -1.23162264e+00 6.46629863e-01 + 2.28183862e+00 -9.24157063e-01 + -4.29615882e-01 5.69130863e-01 + -1.37449121e+00 -9.12032183e-01 + -7.33890904e-01 -3.91865471e-02 + 8.41400661e-01 -4.76002200e-01 + -1.73349274e-01 -6.84143467e-02 + 3.16042891e+00 -1.32651856e+00 + -3.78244609e+00 2.38619718e+00 + -3.69634380e+00 2.22368561e+00 + 1.83766344e+00 -1.65675953e+00 + -1.63206002e+00 1.19484469e+00 + 3.68480064e-01 -5.70764494e-01 + 3.61982479e-01 1.04274409e-01 + 2.48863048e+00 -1.13285542e+00 + -2.81896488e+00 9.47958768e-01 + 5.74952901e-01 -2.75959392e-01 + 3.72783275e-01 -3.48937848e-01 + 1.95935716e+00 -1.06750415e+00 + 5.19357531e+00 -2.32070803e+00 + 4.09246149e+00 -1.89976700e+00 + -3.36666087e-01 8.17645057e-02 + 1.85453493e-01 3.76913151e-01 + -3.06458262e+00 1.34106402e+00 + -3.13796566e+00 7.00485099e-01 + 1.42964058e+00 -1.35536932e-01 + -1.23440423e-01 4.60094177e-02 + -2.86753037e+00 -5.21724160e-02 + 2.67113726e+00 -1.83746924e+00 + -1.35335062e+00 1.28238073e+00 + -2.43569899e+00 1.25998539e+00 + 1.26036740e-01 -2.35416844e-01 + -1.35725745e+00 7.37788491e-01 + -3.80897538e-01 3.30757889e-01 + 6.58694434e-01 -1.07566603e+00 + 2.11273640e+00 -9.02260632e-01 + 4.00755057e-01 -2.49229150e-02 + -1.80095812e+00 9.73099742e-01 + -2.68408372e+00 1.63737364e+00 + -2.66079826e+00 7.47289412e-01 + -9.92321439e-02 -1.49331396e-01 + 4.45678251e+00 -1.80352394e+00 + 1.35962915e+00 -1.31554389e+00 + -7.76601417e-01 -9.66173523e-02 + 1.68096348e+00 -6.27235133e-01 + 1.53081227e-01 -3.54216830e-01 + -1.54913095e+00 3.43689269e-01 + 5.29187357e-02 -6.73916964e-01 + -2.06606084e+00 8.34784242e-01 + 1.73701179e+00 -6.06467340e-01 + 1.55856757e+00 -2.58642780e-01 + 1.04349101e+00 -4.43027348e-01 + -1.02397719e+00 1.01308824e+00 + -2.13860204e-01 -4.73347361e-01 + -2.59004955e+00 1.43367853e+00 + 7.98457679e-01 2.18621627e-02 + -1.32974762e+00 4.61802208e-01 + 3.21419359e-01 2.30723316e-02 + 2.87201888e-02 6.24566672e-02 + -1.22261418e+00 6.02340363e-01 + 1.28750335e+00 -3.34839548e-02 + -9.67952623e-01 4.34470505e-01 + 2.02850324e+00 -9.05160255e-01 + -4.13946010e+00 2.33779091e+00 + -4.47508806e-01 3.06440495e-01 + -3.91543394e+00 1.68251022e+00 + -6.45193001e-01 5.29781162e-01 + -2.15518916e-02 5.07278355e-01 + -2.83356868e+00 1.00670227e+00 + 1.82989749e+00 -1.37329222e+00 + -1.09330213e+00 1.08560688e+00 + 1.90533722e+00 -1.28905879e+00 + 2.33986084e+00 2.30642626e-02 + 8.01940220e-01 -1.63986962e+00 + -4.23415165e+00 2.07530423e+00 + 9.33382522e-01 -7.62917211e-01 + -1.84033954e+00 1.07469401e+00 + -2.81938669e+00 1.07342024e+00 + -7.05169988e-01 2.13124943e-01 + 5.09598137e-01 1.32725493e-01 + -2.34558226e+00 8.62383168e-01 + -1.70322072e+00 2.70893796e-01 + 1.23652660e+00 -7.53216034e-02 + 2.84660646e+00 -3.48178304e-02 + 2.50250128e+00 -1.27770855e+00 + -1.00279469e+00 8.77194218e-01 + -4.34674121e-02 -2.12091350e-01 + -5.84151289e-01 1.50382340e-01 + -1.79024013e+00 4.24972808e-01 + -1.23434666e+00 -8.85546570e-02 + 1.36575412e+00 -6.42639880e-01 + -1.98429947e+00 2.27650336e-01 + 2.36253589e+00 -1.51340773e+00 + 8.79157643e-01 6.84142159e-01 + -2.18577755e+00 2.76526200e-01 + -3.55473434e-01 8.29976561e-01 + 1.16442595e+00 -5.97699411e-01 + -7.35528097e-01 2.40318183e-01 + -1.73702631e-01 7.33788663e-02 + -1.40451745e+00 3.24899628e-01 + -2.05434385e+00 5.68123738e-01 + 8.47876642e-01 -5.74224294e-01 + -6.91955602e-01 1.26009087e+00 + 2.56574498e+00 -1.15602581e+00 + 3.93306545e+00 -1.38398209e+00 + -2.73230251e+00 4.89062581e-01 + -1.04315474e+00 6.06335547e-01 + 1.23231431e+00 -4.46675065e-01 + -3.93035285e+00 1.43287651e+00 + -1.02132111e+00 9.58919791e-01 + -1.49425352e+00 1.06456165e+00 + -6.26485337e-01 1.03791402e+00 + -6.61772998e-01 2.63275425e-01 + -1.80940386e+00 5.70767403e-01 + 9.83720450e-01 -1.39449756e-01 + -2.24619662e+00 9.01044870e-01 + 8.94343014e-01 5.31038678e-02 + 1.95518199e-01 -2.81343295e-01 + -2.30533019e-01 -1.74478106e-01 + -2.01550361e+00 5.55958010e-01 + -4.36281469e+00 1.94374226e+00 + -5.18530457e+00 2.89278357e+00 + 2.67289101e+00 -2.98511449e-01 + -1.53566179e+00 -1.00588944e-01 + -6.09943217e-02 -1.56986047e-01 + -5.22146452e+00 1.66209208e+00 + -3.69777478e+00 2.26154873e+00 + 2.24607181e-01 -4.86934960e-01 + 2.49909450e+00 -1.03033370e+00 + -1.07841120e+00 8.22388054e-01 + -3.20697089e+00 1.09536143e+00 + 3.43524232e+00 -1.47289362e+00 + -5.65784134e-01 4.60365175e-01 + -1.76714734e+00 1.57752346e-01 + -7.77620365e-01 5.60153443e-01 + 6.34399352e-01 -5.22339836e-01 + 2.91011875e+00 -9.72623380e-01 + -1.19286824e+00 6.32370253e-01 + -2.18327609e-01 8.23953181e-01 + 3.42430842e-01 1.37098055e-01 + 1.28658034e+00 -9.11357320e-01 + 2.06914465e+00 -6.67556382e-01 + -6.69451020e-01 -6.38605102e-01 + -2.09312398e+00 1.16743634e+00 + -3.63778357e+00 1.91919157e+00 + 8.74685911e-01 -1.09931208e+00 + -3.91496791e+00 1.00808357e+00 + 1.29621330e+00 -8.32239802e-01 + 9.00222045e-01 -1.31159793e+00 + -1.12242062e+00 1.98517079e-01 + -3.71932852e-01 1.31667093e-01 + -2.23829610e+00 1.26328346e+00 + -2.08365062e+00 9.93385336e-01 + -1.91082720e+00 7.45866855e-01 + 4.38024917e+00 -2.05901118e+00 + -2.28872886e+00 6.85279335e-01 + 1.01274497e-01 -3.26227153e-01 + -5.04447572e-01 -3.18619513e-01 + 1.28537006e+00 -1.04573551e+00 + -7.83175212e-01 1.54791645e-01 + -3.89239175e+00 1.60017929e+00 + -8.87877111e-01 -1.04968005e-01 + 9.32215179e-01 -5.58691113e-01 + -6.44977127e-01 -2.23018375e-01 + 1.10141900e+00 -1.00666432e+00 + 2.92755687e-01 -1.45480350e-01 + 7.73580681e-01 -2.21150567e-01 + -1.40873709e+00 7.61548044e-01 + -8.89031805e-01 -3.48542923e-01 + 4.16844267e-01 -2.39914494e-01 + -4.64265832e-01 7.29581138e-01 + 1.99835179e+00 -7.70542813e-01 + 4.20523191e-02 -2.18783563e-01 + -6.32611758e-01 -3.09926115e-01 + 6.82912198e-02 -8.48327050e-01 + 1.92425229e+00 -1.37876951e+00 + 3.49461782e+00 -1.88354255e+00 + -3.25209026e+00 1.49809395e+00 + 6.59273182e-01 -2.37435654e-01 + -1.15517300e+00 8.46134387e-01 + 1.26756151e+00 -4.58988026e-01 + -3.99178418e+00 2.04153008e+00 + 7.05687841e-01 -6.83433306e-01 + -1.61997342e+00 8.16577004e-01 + -3.89750399e-01 4.29753250e-01 + -2.53026432e-01 4.92861432e-01 + -3.16788324e+00 4.44285524e-01 + -7.86248901e-01 1.12753716e+00 + -3.02351433e+00 1.28419015e+00 + -1.30131355e+00 1.71226678e+00 + -4.08843475e+00 1.62063214e+00 + -3.09209403e+00 1.19958520e+00 + 1.49102271e+00 -1.11834864e+00 + -3.18059348e+00 5.74587042e-01 + 2.06054867e+00 3.25797860e-03 + -3.50999200e+00 2.02412428e+00 + -8.26610023e-01 3.46528211e-01 + 2.00546034e+00 -4.07333110e-01 + -9.69941653e-01 4.80953753e-01 + 4.47925660e+00 -2.33127314e+00 + 2.03845790e+00 -9.90439915e-01 + -1.11349191e+00 4.31183918e-01 + -4.03628396e+00 1.68509679e+00 + -1.48177601e+00 7.74322088e-01 + 3.07369385e+00 -9.57465886e-01 + 2.39011286e+00 -6.44506921e-01 + 2.91561991e+00 -8.78627328e-01 + 1.10212733e+00 -4.21637388e-01 + 5.31985231e-01 -6.17445696e-01 + -6.82340929e-01 -2.93529716e-01 + 1.94290679e+00 -4.64268634e-01 + 1.92262116e+00 -7.93142835e-01 + 4.73762800e+00 -1.63654174e+00 + -3.17848641e+00 8.05791391e-01 + 4.08739432e+00 -1.80816807e+00 + -7.60648826e-01 1.24216138e-01 + -2.24716400e+00 7.90020937e-01 + 1.64284052e+00 -7.18784070e-01 + 1.04410012e-01 -7.11195880e-02 + 2.18268225e+00 -7.01767831e-01 + 2.06218013e+00 -8.70251746e-01 + -1.35266581e+00 7.08456358e-01 + -1.38157779e+00 5.14401086e-01 + -3.28326008e+00 1.20988399e+00 + 8.85358917e-01 -8.12213495e-01 + -2.34067500e+00 3.67657353e-01 + 3.96878127e+00 -1.66841450e+00 + 1.36518053e+00 -8.33436812e-01 + 5.25771988e-01 -5.06121987e-01 + -2.25948361e+00 1.30663765e+00 + -2.57662070e+00 6.32114628e-01 + -3.43134685e+00 2.38106008e+00 + 2.31571924e+00 -1.56566818e+00 + -2.95397202e+00 1.05661888e+00 + -1.35331242e+00 6.76383411e-01 + 1.40977132e+00 -1.17775938e+00 + 1.52561996e+00 -9.83147176e-01 + 2.26550832e+00 -2.10464123e-02 + 6.23371684e-01 -5.30768122e-01 + -4.42356624e-01 9.72226986e-01 + 2.31517901e+00 -1.08468105e+00 + 1.97236640e+00 -1.42016619e+00 + 3.18618687e+00 -1.45056343e+00 + -2.75880360e+00 5.40254980e-01 + -1.92916581e+00 1.45029864e-01 + 1.90022524e+00 -6.03805754e-01 + -1.05446211e+00 5.74361752e-01 + 1.45990390e+00 -9.28233993e-01 + 5.14960557e+00 -2.07564096e+00 + -7.53104842e-01 1.55876958e-01 + 8.09490983e-02 -8.58886384e-02 + -1.56894969e+00 4.53497227e-01 + 1.36944658e-01 5.60670875e-01 + -5.32635329e-01 4.40309945e-01 + 1.32507853e+00 -5.83670099e-01 + 1.20676031e+00 -8.02296831e-01 + -3.65023422e+00 1.17211368e+00 + 1.53393850e+00 -6.17771312e-01 + -3.99977129e+00 1.71415137e+00 + 5.70705058e-01 -4.60771539e-01 + -2.20608002e+00 1.07866596e+00 + -1.09040244e+00 6.77441076e-01 + -5.09886482e-01 -1.97282128e-01 + -1.58062785e+00 6.18333697e-01 + -1.53295020e+00 4.02168701e-01 + -5.18580598e-01 2.25767177e-01 + 1.59514316e+00 -2.54983617e-01 + -5.91938655e+00 2.68223782e+00 + 2.84200509e+00 -1.04685313e+00 + 1.31298664e+00 -1.16672614e+00 + -2.36660033e+00 1.81359460e+00 + 6.94163290e-02 3.76658816e-01 + 2.33973934e+00 -8.33173023e-01 + -8.24640389e-01 7.83717285e-01 + -1.02888281e+00 1.04680766e+00 + 1.34750745e+00 -5.89568160e-01 + -2.48761231e+00 7.44199284e-01 + -1.04501559e+00 4.72326911e-01 + -3.14610089e+00 1.89843692e+00 + 2.13003416e-01 5.76633620e-01 + -1.69239608e+00 5.66070021e-01 + 1.80491280e+00 -9.31701080e-01 + -6.94362572e-02 6.96026587e-01 + 1.36502578e+00 -6.85599000e-02 + -7.76764337e-01 3.64328661e-01 + -2.67322167e+00 6.80150021e-01 + 1.84338485e+00 -1.18487494e+00 + 2.88009231e+00 -1.25700411e+00 + 1.17114433e+00 -7.69727080e-01 + 2.11576167e+00 2.81502116e-01 + -1.51470088e+00 2.61553540e-01 + 1.18923669e-01 -1.17890202e-01 + 4.48359786e+00 -1.81427466e+00 + -1.27055948e+00 9.92388998e-01 + -8.00276606e-01 9.11326621e-02 + 7.51764024e-01 -1.03676498e-01 + 1.35769348e-01 -2.11470084e-01 + 2.50731332e+00 -1.12418270e+00 + -2.49752781e-01 7.81224033e-02 + -6.23037902e-01 3.16599691e-01 + -3.93772902e+00 1.37195391e+00 + 1.74256361e+00 -1.12363582e+00 + -1.49737281e+00 5.98828310e-01 + 7.75592115e-01 -4.64733802e-01 + -2.26027693e+00 1.36991118e+00 + -1.62849836e+00 7.36899107e-01 + 2.36850751e+00 -9.32126872e-01 + 5.86169745e+00 -2.49342512e+00 + -5.37092226e-01 1.23821274e+00 + 2.80535867e+00 -1.93363302e+00 + -1.77638106e+00 9.10050276e-01 + 3.02692018e+00 -1.60774676e+00 + 1.97833084e+00 -1.50636531e+00 + 9.09168906e-01 -8.83799359e-01 + 2.39769655e+00 -7.56977869e-01 + 1.47283981e+00 -1.06749890e+00 + 2.92060943e-01 -6.07040605e-01 + -2.09278201e+00 7.71858590e-01 + 7.10015905e-01 -5.42768432e-01 + -2.16826169e-01 1.56897896e-01 + 4.56288247e+00 -2.08912680e+00 + -6.63374020e-01 6.67325183e-01 + 1.80564442e+00 -9.76366134e-01 + 3.28720168e+00 -4.66575145e-01 + -1.60463695e-01 -2.58428153e-01 + 1.78590750e+00 -3.96427146e-01 + 2.75950306e+00 -1.82102856e+00 + -1.18234310e+00 6.28073320e-01 + 4.11415835e+00 -2.33551216e+00 + 1.38721004e+00 -2.77450622e-01 + -2.94903545e+00 1.74813352e+00 + 8.67290400e-01 -6.51667894e-01 + 2.70022274e+00 -8.11832480e-01 + -2.06766146e+00 8.24047249e-01 + 3.90717142e+00 -1.20155758e+00 + -2.95102809e+00 1.36667968e+00 + 6.08815147e+00 -2.60737974e+00 + 2.78576476e+00 -7.86628755e-01 + -3.26258407e+00 1.09302450e+00 + 1.59849422e+00 -1.09705202e+00 + -2.50600710e-01 1.63243175e-01 + -4.90477087e-01 -4.57729572e-01 + -1.24837181e+00 3.22157840e-01 + -2.46341049e+00 1.06517849e+00 + 9.62880751e-01 4.56962496e-01 + 3.99964487e-01 2.07472802e-01 + 6.36657705e-01 -3.46400942e-02 + 4.91231407e-02 -1.40289235e-02 + -4.66683524e-02 -3.72326100e-01 + -5.22049702e-01 -1.70440260e-01 + 5.27062938e-01 -2.32628395e-01 + -2.69440318e+00 1.18914874e+00 + 3.65087539e+00 -1.53427267e+00 + -1.16546364e-01 4.93245392e-02 + 7.55931384e-01 -3.02980139e-01 + 2.06338745e+00 -6.24841225e-01 + 1.31177908e-01 7.29338183e-01 + 1.48021784e+00 -6.39509896e-01 + -5.98656707e-01 2.84525503e-01 + -2.18611080e+00 1.79549812e+00 + -2.91673624e+00 2.15772237e-01 + -8.95591350e-01 7.68250538e-01 + 1.36139762e+00 -1.93845144e-01 + 5.45730414e+00 -2.28114404e+00 + 3.22747247e-01 9.33582332e-01 + -1.46384504e+00 1.12801186e-01 + 4.26728166e-01 -2.33481242e-01 + -1.41327270e+00 8.16103740e-01 + -2.53998067e-01 1.44906646e-01 + -1.32436467e+00 1.87556361e-01 + -3.77313086e+00 1.32896038e+00 + 3.77651731e+00 -1.76548043e+00 + -2.45297093e+00 1.32571926e+00 + -6.55900588e-01 3.56921462e-01 + 9.25558722e-01 -4.51988954e-01 + 1.20732231e+00 -3.02821614e-01 + 3.72660154e-01 -1.89365208e-01 + -1.77090939e+00 9.18087975e-01 + 3.01127567e-01 2.67965829e-01 + -1.76708900e+00 4.62069259e-01 + -2.71812099e+00 1.57233508e+00 + -5.35297633e-01 4.99231535e-01 + 1.50507631e+00 -9.85763646e-01 + 3.00424787e+00 -1.29837562e+00 + -4.99311105e-01 3.91086482e-01 + 1.30125207e+00 -1.26247924e-01 + 4.01699483e-01 -4.46909391e-01 + -1.33635257e+00 5.12068703e-01 + 1.39229757e+00 -9.10974858e-01 + -1.74229508e+00 1.49475978e+00 + -1.21489414e+00 4.04193753e-01 + -3.36537605e-01 -6.74335427e-01 + -2.79186828e-01 8.48314720e-01 + -2.03080140e+00 1.66599815e+00 + -3.53064281e-01 -7.68582906e-04 + -5.30305657e+00 2.91091546e+00 + -1.20049972e+00 8.26578358e-01 + 2.95906989e-01 2.40215920e-01 + -1.42955534e+00 4.63480310e-01 + -1.87856619e+00 8.21459385e-01 + -2.71124720e+00 1.80246843e+00 + -3.06933780e+00 1.22235760e+00 + 5.21935582e-01 -1.27298218e+00 + -1.34175797e+00 7.69018937e-01 + -1.81962785e+00 1.15528991e+00 + -3.99227550e-01 2.93821598e-01 + 1.22533179e+00 -4.73846323e-01 + -2.08068359e-01 -1.75039817e-01 + -2.03068526e+00 1.50370503e+00 + -3.27606113e+00 1.74906330e+00 + -4.37802587e-01 -2.26956048e-01 + -7.69774213e-02 -3.54922468e-01 + 6.47160749e-02 -2.07334721e-01 + -1.37791524e+00 4.43766709e-01 + 3.29846803e+00 -1.04060799e+00 + -3.63704046e+00 1.05800226e+00 + -1.26716116e+00 1.13077353e+00 + 1.98549075e+00 -1.31864807e+00 + 1.85159500e+00 -5.78629560e-01 + -1.55295206e+00 1.23655857e+00 + 6.76026255e-01 9.18824125e-02 + 1.23418960e+00 -4.68162027e-01 + 2.43186642e+00 -9.22422440e-01 + -3.18729701e+00 1.77582673e+00 + -4.02945613e+00 1.14303496e+00 + -1.92694576e-01 1.03301431e-01 + 1.89554730e+00 -4.60128096e-01 + -2.55626581e+00 1.16057084e+00 + 6.89144365e-01 -9.94982900e-01 + -4.44680606e+00 2.19751983e+00 + -3.15196193e+00 1.18762993e+00 + -1.17434977e+00 1.04534656e+00 + 8.58386984e-02 -1.03947487e+00 + 3.33354973e-01 5.54813610e-01 + -9.37631808e-01 3.33450150e-01 + -2.50232471e+00 5.39720635e-01 + 1.03611949e+00 -7.16304095e-02 + -2.05556816e-02 -3.28992265e-01 + -2.24176201e+00 1.13077506e+00 + 4.53583688e+00 -1.10710212e+00 + 4.77389762e-01 -8.99445512e-01 + -2.69075551e+00 6.83176866e-01 + -2.21779724e+00 1.16916849e+00 + -1.09669056e+00 2.10044765e-01 + -8.45367920e-01 -8.45951423e-02 + 4.37558941e-01 -6.95904256e-01 + 1.84884195e+00 -1.71205136e-01 + -8.36371957e-01 5.62862478e-01 + 1.27786531e+00 -1.33362147e+00 + 2.90684492e+00 -7.49892184e-01 + -3.38652716e+00 1.51180670e+00 + -1.30945978e+00 7.09261928e-01 + -7.50471924e-01 -5.24637889e-01 + 1.18580718e+00 -9.97943971e-04 + -7.55395645e+00 3.19273590e+00 + 1.72822535e+00 -1.20996962e+00 + 5.67374320e-01 6.19573416e-01 + -2.99163781e+00 1.79721534e+00 + 1.49862187e+00 -6.05631846e-02 + 1.79503506e+00 -4.90419706e-01 + 3.85626054e+00 -1.95396324e+00 + -9.39188410e-01 7.96498057e-01 + 2.91986664e+00 -1.29392724e+00 + -1.54265750e+00 6.40727933e-01 + 1.14919794e+00 1.20834257e-01 + 2.00936817e+00 -1.53728359e+00 + 3.72468420e+00 -1.38704612e+00 + -1.27794802e+00 3.48543179e-01 + 3.63294077e-01 5.70623314e-01 + 1.49381016e+00 -6.04500534e-01 + 2.98912256e+00 -1.72295726e+00 + -1.80833817e+00 2.94907625e-01 + -3.19669622e+00 1.31888700e+00 + 1.45889401e+00 -8.88448639e-01 + -2.80045388e+00 1.01207060e+00 + -4.78379567e+00 1.48646520e+00 + 2.25510003e+00 -7.13372461e-01 + -9.74441433e-02 -2.17766373e-01 + 2.64468496e-01 -3.60842698e-01 + -5.98821713e+00 3.20197892e+00 + 2.67030213e-01 -5.36386416e-01 + 2.24546960e+00 -8.13464649e-01 + -4.89171414e-01 3.86255031e-01 + -7.45713706e-01 6.29800380e-01 + -3.30460503e-01 3.85127284e-01 + -4.19588147e+00 1.52793198e+00 + 5.42078582e-01 -2.61642741e-02 + 4.24938513e-01 -5.72936751e-01 + 2.82717288e+00 -6.75355024e-01 + -1.44741788e+00 5.03578028e-01 + -1.65547573e+00 7.76444277e-01 + 2.20361170e+00 -1.40835680e+00 + -3.69540235e+00 2.32953767e+00 + -1.41909357e-01 2.28989778e-01 + 1.92838879e+00 -8.72525737e-01 + 1.40708100e+00 -6.81849638e-02 + 1.24988112e+00 -1.39470590e-01 + -2.39435855e+00 7.26587655e-01 + 7.03985028e-01 4.85403277e-02 + 4.05214529e+00 -9.16928318e-01 + 3.74198837e-01 -5.04192358e-01 + -8.43374127e-01 2.36064018e-01 + -3.32253349e-01 7.47840055e-01 + -6.03725210e+00 1.95173337e+00 + 4.60829865e+00 -1.51191309e+00 + -1.46247098e+00 1.11140916e+00 + -9.60111157e-01 -1.23189114e-01 + -7.49613187e-01 4.53614129e-01 + -5.77838219e-01 2.07366469e-02 + 8.07652950e-01 -5.16272662e-01 + -6.02556049e-01 5.05318649e-01 + -1.28712445e-01 2.57836512e-01 + -5.27662820e+00 2.11790737e+00 + 5.40819308e+00 -2.15366022e+00 + 9.37742513e-02 -1.60221751e-01 + 4.55902865e+00 -1.24646307e+00 + -9.06582589e-01 1.92928110e-01 + 2.99928996e+00 -8.04301218e-01 + -3.24317381e+00 1.80076061e+00 + 3.20421743e-01 8.76524679e-01 + -5.29606705e-01 -3.16717696e-01 + -1.77264560e+00 7.52686776e-01 + -1.51706824e+00 8.43755103e-01 + 1.52759111e+00 -7.86814243e-01 + 4.74845617e-01 4.21319700e-01 + 6.97829149e-01 -8.15664881e-01 + 3.09564973e+00 -1.06202469e+00 + 2.95320379e+00 -1.98963943e+00 + -4.23033224e+00 1.41013338e+00 + 1.48576206e+00 8.02908511e-02 + 4.52041627e+00 -2.04620399e+00 + 6.58403922e-01 -7.60781799e-01 + 2.10667543e-01 1.15241731e-01 + 1.77702583e+00 -8.10271859e-01 + 2.41277385e+00 -1.46972042e+00 + 1.50685525e+00 -1.99272545e-01 + 7.61665522e-01 -4.11276152e-01 + 1.18352312e+00 -9.59908608e-01 + -3.32031305e-01 8.07500132e-02 + 1.16813118e+00 -1.73095194e-01 + 1.18363346e+00 -5.41565052e-01 + 5.17702179e-01 -7.62442035e-01 + 4.57401006e-01 -1.45951115e-02 + 1.49377115e-01 2.99571605e-01 + 1.40399453e+00 -1.30160353e+00 + 5.26231567e-01 3.52783752e-01 + -1.91136514e+00 4.24228635e-01 + 1.74156701e+00 -9.92076776e-01 + -4.89323391e+00 2.32483507e+00 + 2.54011209e+00 -8.80366295e-01 + -5.56925706e-01 1.48842026e-01 + -2.35904668e+00 9.60474853e-01 + 1.42216971e+00 -4.67062761e-01 + -1.10809680e+00 7.68684300e-01 + 4.09674726e+00 -1.90795680e+00 + -2.23048923e+00 9.03812542e-01 + 6.57025763e-01 1.36514871e-01 + 2.10944145e+00 -9.78897838e-02 + 1.22552525e+00 -2.50303867e-01 + 2.84620103e-01 -5.30164020e-01 + -2.13562585e+00 1.03503056e+00 + 1.32414902e-01 -8.14190240e-03 + -5.82433561e-01 3.21020292e-01 + -5.06473247e-01 3.11530419e-01 + 1.57162465e+00 -1.20763919e+00 + -1.43155284e+00 -2.51203698e-02 + -1.47093713e+00 -1.39620999e-01 + -2.65765643e+00 1.06091403e+00 + 2.45992927e+00 -5.88815836e-01 + -1.28440162e+00 -1.99377398e-01 + 6.11257504e-01 -3.73577401e-01 + -3.46606103e-01 6.06081290e-01 + 3.76687505e+00 -8.80181424e-01 + -1.03725103e+00 1.45177517e+00 + 2.76659936e+00 -1.09361320e+00 + -3.61311296e+00 9.75032455e-01 + 3.22878655e+00 -9.69497365e-01 + 1.43560379e+00 -5.52524585e-01 + 2.94042153e+00 -1.79747037e+00 + 1.30739580e+00 2.47989248e-01 + -4.05056982e-01 1.22831715e+00 + -2.25827421e+00 2.30604626e-01 + 3.69262926e-01 4.32714650e-02 + -5.52064063e-01 6.07806340e-01 + 7.03325987e+00 -2.17956730e+00 + -2.37823835e-01 -8.28068639e-01 + -4.84279888e-01 5.67765194e-01 + -3.15863410e+00 1.02241617e+00 + -3.39561593e+00 1.36876374e+00 + -2.78482934e+00 6.81641104e-01 + -4.37604334e+00 2.23826340e+00 + -2.54049692e+00 8.22676745e-01 + 3.73264822e+00 -9.93498732e-01 + -3.49536064e+00 1.84771519e+00 + 9.81801604e-01 -5.21278776e-01 + 1.52996831e+00 -1.27386206e+00 + -9.23490293e-01 5.29099482e-01 + -2.76999461e+00 9.24831872e-01 + -3.30029834e-01 -2.49645555e-01 + -1.71156166e+00 5.44940854e-01 + -2.37009487e+00 5.83826982e-01 + -3.03216865e+00 1.04922722e+00 + -2.19539936e+00 1.37558730e+00 + 1.15350207e+00 -6.15318535e-01 + 4.62011792e+00 -2.46714517e+00 + 1.52627952e-02 -1.00618283e-01 + -1.10399342e+00 4.87413533e-01 + 3.55448194e+00 -9.10394190e-01 + -5.21890321e+00 2.44710745e+00 + 1.54289749e+00 -6.54269311e-01 + 2.67935674e+00 -9.92758863e-01 + 1.05801310e+00 2.60054285e-02 + 1.52509097e+00 -4.08768600e-01 + 3.27576917e+00 -1.28769406e+00 + 1.71008412e-01 -2.68739994e-01 + -9.83351344e-04 7.02495897e-02 + -7.60795056e-03 1.61968285e-01 + -1.80620472e+00 4.24934471e-01 + 2.32023297e-02 -2.57284559e-01 + 3.98219478e-01 -4.65361935e-01 + 6.63476988e-01 -3.29823196e-02 + 4.00154707e+00 -1.01792211e+00 + -1.50286870e+00 9.46875359e-01 + -2.22717585e+00 7.50636195e-01 + -3.47381508e-01 -6.51596975e-01 + 2.08076453e+00 -8.22800165e-01 + 2.05099963e+00 -4.00868250e-01 + 3.52576988e-02 -2.54418565e-01 + 1.57342042e+00 -7.62166492e-02 + -1.47019722e+00 3.40861172e-01 + -1.21156090e+00 3.21891246e-01 + 3.79729047e+00 -1.54350764e+00 + 1.26459678e-02 6.99203693e-01 + 1.53974177e-01 4.68643204e-01 + -1.73923561e-01 -1.26229768e-01 + 4.54644993e+00 -2.13951783e+00 + 1.46022547e-01 -4.57084165e-01 + 6.50048037e+00 -2.78872609e+00 + -1.51934912e+00 1.03216768e+00 + -3.06483575e+00 1.81101446e+00 + -2.38212125e+00 9.19559042e-01 + -1.81319611e+00 8.10545112e-01 + 1.70951294e+00 -6.10712680e-01 + 1.67974156e+00 -1.51241453e+00 + -5.94795113e+00 2.56893813e+00 + 3.62633110e-01 -7.46965304e-01 + -2.44042594e+00 8.52761797e-01 + 3.32412550e+00 -1.28439899e+00 + 4.74860766e+00 -1.72821964e+00 + 1.29072541e+00 -8.24872902e-01 + -1.69450702e+00 4.09600876e-01 + 1.29705411e+00 1.22300809e-01 + -2.63597613e+00 8.55612913e-01 + 9.28467301e-01 -2.63550114e-02 + 2.44670264e+00 -4.10123002e-01 + 1.06408206e+00 -5.03361942e-01 + 5.12384049e-02 -1.27116595e-02 + -1.06731272e+00 -1.76205029e-01 + -9.45454582e-01 3.74404917e-01 + 2.54343689e+00 -7.13810545e-01 + -2.54460335e+00 1.31590265e+00 + 1.89864233e+00 -3.98436339e-01 + -1.93990133e+00 6.01474630e-01 + -1.35938824e+00 4.00751788e-01 + 2.38567018e+00 -6.13904880e-01 + 2.18748050e-01 2.62631712e-01 + -2.01388788e+00 1.41474031e+00 + 2.74014581e+00 -1.27448105e+00 + -2.13828583e+00 1.13616144e+00 + 5.98730932e+00 -2.53430080e+00 + -1.72872795e+00 1.53702057e+00 + -2.53263962e+00 1.27342410e+00 + 1.34326968e+00 -1.99395088e-01 + 3.83352666e-01 -1.25683065e-01 + -2.35630657e+00 5.54116983e-01 + -1.94900838e+00 5.76270178e-01 + -1.36699108e+00 -3.40904824e-01 + -2.34727346e+00 -1.93054940e-02 + -3.82779777e+00 1.83025664e+00 + -4.31602080e+00 9.21605705e-01 + 5.54098133e-01 2.33991419e-01 + -4.53591188e+00 1.99833353e+00 + -3.92715909e+00 1.83231482e+00 + 3.91344440e-01 -1.11355111e-01 + 3.48576363e+00 -1.41379449e+00 + -1.42858690e+00 3.84532286e-01 + 1.79519859e+00 -9.23486448e-01 + 8.49691242e-01 -1.76551084e-01 + 1.53618138e+00 8.23835015e-02 + 5.91476520e-02 3.88296940e-02 + 1.44837346e+00 -7.24097604e-01 + -6.79008418e-01 4.04078097e-01 + 2.87555510e+00 -9.51825076e-01 + -1.12379101e+00 2.93457714e-01 + 1.45263980e+00 -6.01960544e-01 + -2.55741621e-01 9.26233518e-01 + 3.54570714e+00 -1.41521877e+00 + -1.61542388e+00 6.57844512e-01 + -3.22844269e-01 3.02823546e-01 + 1.03523913e+00 -6.92730711e-01 + 1.11084909e+00 -3.50823642e-01 + 3.41268693e+00 -1.90865862e+00 + 7.67062858e-01 -9.48792160e-01 + -5.49798016e+00 1.71139960e+00 + 1.14865798e+00 -6.12669150e-01 + -2.18256680e+00 7.78634462e-01 + 4.78857389e+00 -2.55555085e+00 + -1.85555569e+00 8.04311615e-01 + -4.22278799e+00 2.01162524e+00 + -1.56556149e+00 1.54353907e+00 + -3.11527864e+00 1.65973526e+00 + 2.66342611e+00 -1.20449402e+00 + 1.57635314e+00 -1.48716308e-01 + -6.35606865e-01 2.59701180e-01 + 1.02431976e+00 -6.76929904e-01 + 1.12973772e+00 1.49473892e-02 + -9.12758116e-01 2.21533933e-01 + -2.98014470e+00 1.71651189e+00 + 2.74016965e+00 -9.47893923e-01 + -3.47830591e+00 1.34941430e+00 + 1.74757562e+00 -3.72503752e-01 + 5.55820383e-01 -6.47992466e-01 + -1.19871928e+00 9.82429151e-01 + -2.53040133e+00 2.10671307e+00 + -1.94085605e+00 1.38938137e+00 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/license.txt b/spark-3.5.0-bin-hadoop3/data/mllib/images/license.txt new file mode 100644 index 0000000..052f302 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/images/license.txt @@ -0,0 +1,13 @@ +The images in the folder "kittens" are under the creative commons CC0 license, or no rights reserved: +https://creativecommons.org/share-your-work/public-domain/cc0/ +The images are taken from: +https://ccsearch.creativecommons.org/image/detail/WZnbJSJ2-dzIDiuUUdto3Q== +https://ccsearch.creativecommons.org/image/detail/_TlKu_rm_QrWlR0zthQTXA== +https://ccsearch.creativecommons.org/image/detail/OPNnHJb6q37rSZ5o_L5JHQ== +https://ccsearch.creativecommons.org/image/detail/B2CVP_j5KjwZm7UAVJ3Hvw== + +The chr30.4.184.jpg and grayscale.jpg images are also under the CC0 license, taken from: +https://ccsearch.creativecommons.org/image/detail/8eO_qqotBfEm2UYxirLntw== + +The image under "multi-channel" directory is under the CC BY-SA 4.0 license cropped from: +https://en.wikipedia.org/wiki/Alpha_compositing#/media/File:Hue_alpha_falloff.png diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/29.5.a_b_EGDP022204.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/29.5.a_b_EGDP022204.jpg new file mode 100644 index 0000000..435e7df Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/29.5.a_b_EGDP022204.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/54893.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/54893.jpg new file mode 100644 index 0000000..825630c Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/54893.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP153539.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP153539.jpg new file mode 100644 index 0000000..571efe9 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP153539.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP802813.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP802813.jpg new file mode 100644 index 0000000..2d12359 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/DP802813.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/not-image.txt b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/not-image.txt new file mode 100644 index 0000000..283e5e9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/kittens/not-image.txt @@ -0,0 +1 @@ +not an image diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/license.txt b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/license.txt new file mode 100644 index 0000000..052f302 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/license.txt @@ -0,0 +1,13 @@ +The images in the folder "kittens" are under the creative commons CC0 license, or no rights reserved: +https://creativecommons.org/share-your-work/public-domain/cc0/ +The images are taken from: +https://ccsearch.creativecommons.org/image/detail/WZnbJSJ2-dzIDiuUUdto3Q== +https://ccsearch.creativecommons.org/image/detail/_TlKu_rm_QrWlR0zthQTXA== +https://ccsearch.creativecommons.org/image/detail/OPNnHJb6q37rSZ5o_L5JHQ== +https://ccsearch.creativecommons.org/image/detail/B2CVP_j5KjwZm7UAVJ3Hvw== + +The chr30.4.184.jpg and grayscale.jpg images are also under the CC0 license, taken from: +https://ccsearch.creativecommons.org/image/detail/8eO_qqotBfEm2UYxirLntw== + +The image under "multi-channel" directory is under the CC BY-SA 4.0 license cropped from: +https://en.wikipedia.org/wiki/Alpha_compositing#/media/File:Hue_alpha_falloff.png diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA.png b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA.png new file mode 100644 index 0000000..a944c6c Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA.png differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA_alpha_60.png b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA_alpha_60.png new file mode 100644 index 0000000..913637c Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/BGRA_alpha_60.png differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/chr30.4.184.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/chr30.4.184.jpg new file mode 100644 index 0000000..7068b97 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/chr30.4.184.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/grayscale.jpg b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/grayscale.jpg new file mode 100644 index 0000000..621cdd1 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/data/mllib/images/origin/multi-channel/grayscale.jpg differ diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/kmeans_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/kmeans_data.txt new file mode 100644 index 0000000..338664f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/kmeans_data.txt @@ -0,0 +1,6 @@ +0.0 0.0 0.0 +0.1 0.1 0.1 +0.2 0.2 0.2 +9.0 9.0 9.0 +9.1 9.1 9.1 +9.2 9.2 9.2 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/pagerank_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/pagerank_data.txt new file mode 100644 index 0000000..95755ab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/pagerank_data.txt @@ -0,0 +1,6 @@ +1 2 +1 3 +1 4 +2 1 +3 1 +4 1 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/pic_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/pic_data.txt new file mode 100644 index 0000000..fcfef8c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/pic_data.txt @@ -0,0 +1,19 @@ +0 1 1.0 +0 2 1.0 +0 3 1.0 +1 2 1.0 +1 3 1.0 +2 3 1.0 +3 4 0.1 +4 5 1.0 +4 15 1.0 +5 6 1.0 +6 7 1.0 +7 8 1.0 +8 9 1.0 +9 10 1.0 +10 11 1.0 +11 12 1.0 +12 13 1.0 +13 14 1.0 +14 15 1.0 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/ridge-data/lpsa.data b/spark-3.5.0-bin-hadoop3/data/mllib/ridge-data/lpsa.data new file mode 100644 index 0000000..fdd16e3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/ridge-data/lpsa.data @@ -0,0 +1,67 @@ +-0.4307829,-1.63735562648104 -2.00621178480549 -1.86242597251066 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +-0.1625189,-1.98898046126935 -0.722008756122123 -0.787896192088153 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +-0.1625189,-1.57881887548545 -2.1887840293994 1.36116336875686 -1.02470580167082 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.155348103855541 +-0.1625189,-2.16691708463163 -0.807993896938655 -0.787896192088153 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +0.3715636,-0.507874475300631 -0.458834049396776 -0.250631301876899 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +0.7654678,-2.03612849966376 -0.933954647105133 -1.86242597251066 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +0.8544153,-0.557312518810673 -0.208756571683607 -0.787896192088153 0.990146852537193 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.2669476,-0.929360463147704 -0.0578991819441687 0.152317365781542 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.2669476,-2.28833047634983 -0.0706369432557794 -0.116315079324086 0.80409888772376 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.2669476,0.223498042876113 -1.41471935455355 -0.116315079324086 -1.02470580167082 -0.522940888712441 -0.29928234305568 0.342627053981254 0.199211097885341 +1.3480731,0.107785900236813 -1.47221551299731 0.420949810887169 -1.02470580167082 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.687186906466865 +1.446919,0.162180092313795 -1.32557369901905 0.286633588334355 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.4701758,-1.49795329918548 -0.263601072284232 0.823898478545609 0.788388310173035 -0.522940888712441 -0.29928234305568 0.342627053981254 0.199211097885341 +1.4929041,0.796247055396743 0.0476559407005752 0.286633588334355 -1.02470580167082 -0.522940888712441 0.394013435896129 -1.04215728919298 -0.864466507337306 +1.5581446,-1.62233848461465 -0.843294091975396 -3.07127197548598 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.5993876,-0.990720665490831 0.458513517212311 0.823898478545609 1.07379746308195 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.6389967,-0.171901281967138 -0.489197399065355 -0.65357996953534 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.6956156,-1.60758252338831 -0.590700340358265 -0.65357996953534 -0.619561070667254 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +1.7137979,0.366273918511144 -0.414014962912583 -0.116315079324086 0.232904453212813 -0.522940888712441 0.971228997418125 0.342627053981254 1.26288870310799 +1.8000583,-0.710307384579833 0.211731938156277 0.152317365781542 -1.02470580167082 -0.522940888712441 -0.442797990776478 0.342627053981254 1.61744790484887 +1.8484548,-0.262791728113881 -1.16708345615721 0.420949810887169 0.0846342590816532 -0.522940888712441 0.163172393491611 0.342627053981254 1.97200710658975 +1.8946169,0.899043117369237 -0.590700340358265 0.152317365781542 -1.02470580167082 -0.522940888712441 1.28643254437683 -1.04215728919298 -0.864466507337306 +1.9242487,-0.903451690500615 1.07659722048274 0.152317365781542 1.28380453408541 -0.522940888712441 -0.442797990776478 -1.04215728919298 -0.864466507337306 +2.008214,-0.0633337899773081 -1.38088970920094 0.958214701098423 0.80409888772376 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +2.0476928,-1.15393789990757 -0.961853075398404 -0.116315079324086 -1.02470580167082 -0.522940888712441 -0.442797990776478 -1.04215728919298 -0.864466507337306 +2.1575593,0.0620203721138446 0.0657973885499142 1.22684714620405 -0.468824786336838 -0.522940888712441 1.31421001659859 1.72741139715549 -0.332627704725983 +2.1916535,-0.75731027755674 -2.92717970468456 0.018001143228728 -1.02470580167082 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.332627704725983 +2.2137539,1.11226993252773 1.06484916245061 0.555266033439982 0.877691038550889 1.89254797819741 1.43890404648442 0.342627053981254 0.376490698755783 +2.2772673,-0.468768642850639 -1.43754788774533 -1.05652863719378 0.576050411655607 -0.522940888712441 0.0120483832567209 0.342627053981254 -0.687186906466865 +2.2975726,-0.618884859896728 -1.1366360750781 -0.519263746982526 -1.02470580167082 -0.522940888712441 -0.863171185425945 3.11219574032972 1.97200710658975 +2.3272777,-0.651431999123483 0.55329161145762 -0.250631301876899 1.11210019001038 -0.522940888712441 -0.179808625688859 -1.04215728919298 -0.864466507337306 +2.5217206,0.115499102435224 -0.512233676577595 0.286633588334355 1.13650173283446 -0.522940888712441 -0.179808625688859 0.342627053981254 -0.155348103855541 +2.5533438,0.266341329949937 -0.551137885443386 -0.384947524429713 0.354857790686005 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.332627704725983 +2.5687881,1.16902610257751 0.855491905752846 2.03274448152093 1.22628985326088 1.89254797819741 2.02833774827712 3.11219574032972 2.68112551007152 +2.6567569,-0.218972367124187 0.851192298581141 0.555266033439982 -1.02470580167082 -0.522940888712441 -0.863171185425945 0.342627053981254 0.908329501367106 +2.677591,0.263121415733908 1.4142681068416 0.018001143228728 1.35980653053822 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +2.7180005,-0.0704736333296423 1.52000996595417 0.286633588334355 1.39364261119802 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.332627704725983 +2.7942279,-0.751957286017338 0.316843561689933 -1.99674219506348 0.911736065044475 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +2.8063861,-0.685277652430997 1.28214038482516 0.823898478545609 0.232904453212813 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.155348103855541 +2.8124102,-0.244991501432929 0.51882005949686 -0.384947524429713 0.823246560137838 -0.522940888712441 -0.863171185425945 0.342627053981254 0.553770299626224 +2.8419982,-0.75731027755674 2.09041984898851 1.22684714620405 1.53428167116843 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +2.8535925,1.20962937075363 -0.242882661178889 1.09253092365124 -1.02470580167082 -0.522940888712441 1.24263233939889 3.11219574032972 2.50384590920108 +2.9204698,0.570886990493502 0.58243883987948 0.555266033439982 1.16006887775962 -0.522940888712441 1.07357183940747 0.342627053981254 1.61744790484887 +2.9626924,0.719758684343624 0.984970304132004 1.09253092365124 1.52137230773457 -0.522940888712441 -0.179808625688859 0.342627053981254 -0.509907305596424 +2.9626924,-1.52406140158064 1.81975700990333 0.689582255992796 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +2.9729753,-0.132431544081234 2.68769877553723 1.09253092365124 1.53428167116843 -0.522940888712441 -0.442797990776478 0.342627053981254 -0.687186906466865 +3.0130809,0.436161292804989 -0.0834447307428255 -0.519263746982526 -1.02470580167082 1.89254797819741 1.07357183940747 0.342627053981254 1.26288870310799 +3.0373539,-0.161195191984091 -0.671900359186746 1.7641120364153 1.13650173283446 -0.522940888712441 -0.863171185425945 0.342627053981254 0.0219314970149 +3.2752562,1.39927182372944 0.513852869452676 0.689582255992796 -1.02470580167082 1.89254797819741 1.49394503405693 0.342627053981254 -0.155348103855541 +3.3375474,1.51967002306341 -0.852203755696565 0.555266033439982 -0.104527297798983 1.89254797819741 1.85927724828569 0.342627053981254 0.908329501367106 +3.3928291,0.560725834706224 1.87867703391426 1.09253092365124 1.39364261119802 -0.522940888712441 0.486423065822545 0.342627053981254 1.26288870310799 +3.4355988,1.00765532502814 1.69426310090641 1.89842825896812 1.53428167116843 -0.522940888712441 -0.863171185425945 0.342627053981254 -0.509907305596424 +3.4578927,1.10152996153577 -0.10927271844907 0.689582255992796 -1.02470580167082 1.89254797819741 1.97630171771485 0.342627053981254 1.61744790484887 +3.5160131,0.100001934217311 -1.30380956369388 0.286633588334355 0.316555063757567 -0.522940888712441 0.28786643052924 0.342627053981254 0.553770299626224 +3.5307626,0.987291634724086 -0.36279314978779 -0.922212414640967 0.232904453212813 -0.522940888712441 1.79270085261407 0.342627053981254 1.26288870310799 +3.5652984,1.07158528137575 0.606453149641961 1.7641120364153 -0.432854616994416 1.89254797819741 0.528504607720369 0.342627053981254 0.199211097885341 +3.5876769,0.180156323255198 0.188987436375017 -0.519263746982526 1.09956763075594 -0.522940888712441 0.708239632330506 0.342627053981254 0.199211097885341 +3.6309855,1.65687973755377 -0.256675483533719 0.018001143228728 -1.02470580167082 1.89254797819741 1.79270085261407 0.342627053981254 1.26288870310799 +3.6800909,0.5720085322365 0.239854450210939 -0.787896192088153 1.0605418233138 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +3.7123518,0.323806133438225 -0.606717660886078 -0.250631301876899 -1.02470580167082 1.89254797819741 0.342907418101747 0.342627053981254 0.199211097885341 +3.9843437,1.23668206715898 2.54220539083611 0.152317365781542 -1.02470580167082 1.89254797819741 1.89037692416194 0.342627053981254 1.26288870310799 +3.993603,0.180156323255198 0.154448192444669 1.62979581386249 0.576050411655607 1.89254797819741 0.708239632330506 0.342627053981254 1.79472750571931 +4.029806,1.60906277046565 1.10378605019827 0.555266033439982 -1.02470580167082 -0.522940888712441 -0.863171185425945 -1.04215728919298 -0.864466507337306 +4.1295508,1.0036214996026 0.113496885050331 -0.384947524429713 0.860016436332751 1.89254797819741 -0.863171185425945 0.342627053981254 -0.332627704725983 +4.3851468,1.25591974271076 0.577607033774471 0.555266033439982 -1.02470580167082 1.89254797819741 1.07357183940747 0.342627053981254 1.26288870310799 +4.6844434,2.09650591351268 0.625488598331018 -2.66832330782754 -1.02470580167082 1.89254797819741 1.67954222367555 0.342627053981254 0.553770299626224 +5.477509,1.30028987435881 0.338383613253713 0.555266033439982 1.00481276295349 1.89254797819741 1.24263233939889 0.342627053981254 1.97200710658975 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_binary_classification_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_binary_classification_data.txt new file mode 100644 index 0000000..861c70c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_binary_classification_data.txt @@ -0,0 +1,100 @@ +0 128:51 129:159 130:253 131:159 132:50 155:48 156:238 157:252 158:252 159:252 160:237 182:54 183:227 184:253 185:252 186:239 187:233 188:252 189:57 190:6 208:10 209:60 210:224 211:252 212:253 213:252 214:202 215:84 216:252 217:253 218:122 236:163 237:252 238:252 239:252 240:253 241:252 242:252 243:96 244:189 245:253 246:167 263:51 264:238 265:253 266:253 267:190 268:114 269:253 270:228 271:47 272:79 273:255 274:168 290:48 291:238 292:252 293:252 294:179 295:12 296:75 297:121 298:21 301:253 302:243 303:50 317:38 318:165 319:253 320:233 321:208 322:84 329:253 330:252 331:165 344:7 345:178 346:252 347:240 348:71 349:19 350:28 357:253 358:252 359:195 372:57 373:252 374:252 375:63 385:253 386:252 387:195 400:198 401:253 402:190 413:255 414:253 415:196 427:76 428:246 429:252 430:112 441:253 442:252 443:148 455:85 456:252 457:230 458:25 467:7 468:135 469:253 470:186 471:12 483:85 484:252 485:223 494:7 495:131 496:252 497:225 498:71 511:85 512:252 513:145 521:48 522:165 523:252 524:173 539:86 540:253 541:225 548:114 549:238 550:253 551:162 567:85 568:252 569:249 570:146 571:48 572:29 573:85 574:178 575:225 576:253 577:223 578:167 579:56 595:85 596:252 597:252 598:252 599:229 600:215 601:252 602:252 603:252 604:196 605:130 623:28 624:199 625:252 626:252 627:253 628:252 629:252 630:233 631:145 652:25 653:128 654:252 655:253 656:252 657:141 658:37 +1 159:124 160:253 161:255 162:63 186:96 187:244 188:251 189:253 190:62 214:127 215:251 216:251 217:253 218:62 241:68 242:236 243:251 244:211 245:31 246:8 268:60 269:228 270:251 271:251 272:94 296:155 297:253 298:253 299:189 323:20 324:253 325:251 326:235 327:66 350:32 351:205 352:253 353:251 354:126 378:104 379:251 380:253 381:184 382:15 405:80 406:240 407:251 408:193 409:23 432:32 433:253 434:253 435:253 436:159 460:151 461:251 462:251 463:251 464:39 487:48 488:221 489:251 490:251 491:172 515:234 516:251 517:251 518:196 519:12 543:253 544:251 545:251 546:89 570:159 571:255 572:253 573:253 574:31 597:48 598:228 599:253 600:247 601:140 602:8 625:64 626:251 627:253 628:220 653:64 654:251 655:253 656:220 681:24 682:193 683:253 684:220 +1 125:145 126:255 127:211 128:31 152:32 153:237 154:253 155:252 156:71 180:11 181:175 182:253 183:252 184:71 209:144 210:253 211:252 212:71 236:16 237:191 238:253 239:252 240:71 264:26 265:221 266:253 267:252 268:124 269:31 293:125 294:253 295:252 296:252 297:108 322:253 323:252 324:252 325:108 350:255 351:253 352:253 353:108 378:253 379:252 380:252 381:108 406:253 407:252 408:252 409:108 434:253 435:252 436:252 437:108 462:255 463:253 464:253 465:170 490:253 491:252 492:252 493:252 494:42 518:149 519:252 520:252 521:252 522:144 546:109 547:252 548:252 549:252 550:144 575:218 576:253 577:253 578:255 579:35 603:175 604:252 605:252 606:253 607:35 631:73 632:252 633:252 634:253 635:35 659:31 660:211 661:252 662:253 663:35 +1 153:5 154:63 155:197 181:20 182:254 183:230 184:24 209:20 210:254 211:254 212:48 237:20 238:254 239:255 240:48 265:20 266:254 267:254 268:57 293:20 294:254 295:254 296:108 321:16 322:239 323:254 324:143 350:178 351:254 352:143 378:178 379:254 380:143 406:178 407:254 408:162 434:178 435:254 436:240 462:113 463:254 464:240 490:83 491:254 492:245 493:31 518:79 519:254 520:246 521:38 547:214 548:254 549:150 575:144 576:241 577:8 603:144 604:240 605:2 631:144 632:254 633:82 659:230 660:247 661:40 687:168 688:209 689:31 +1 152:1 153:168 154:242 155:28 180:10 181:228 182:254 183:100 209:190 210:254 211:122 237:83 238:254 239:162 265:29 266:254 267:248 268:25 293:29 294:255 295:254 296:103 321:29 322:254 323:254 324:109 349:29 350:254 351:254 352:109 377:29 378:254 379:254 380:109 405:29 406:255 407:254 408:109 433:29 434:254 435:254 436:109 461:29 462:254 463:254 464:63 489:29 490:254 491:254 492:28 517:29 518:254 519:254 520:28 545:29 546:254 547:254 548:35 573:29 574:254 575:254 576:109 601:6 602:212 603:254 604:109 630:203 631:254 632:178 658:155 659:254 660:190 686:32 687:199 688:104 +0 130:64 131:253 132:255 133:63 157:96 158:205 159:251 160:253 161:205 162:111 163:4 184:96 185:189 186:251 187:251 188:253 189:251 190:251 191:31 209:16 210:64 211:223 212:244 213:251 214:251 215:211 216:213 217:251 218:251 219:31 236:80 237:181 238:251 239:253 240:251 241:251 242:251 243:94 244:96 245:251 246:251 247:31 263:92 264:253 265:253 266:253 267:255 268:253 269:253 270:253 271:95 272:96 273:253 274:253 275:31 290:92 291:236 292:251 293:243 294:220 295:233 296:251 297:251 298:243 299:82 300:96 301:251 302:251 303:31 317:80 318:253 319:251 320:251 321:188 323:96 324:251 325:251 326:109 328:96 329:251 330:251 331:31 344:96 345:240 346:253 347:243 348:188 349:42 351:96 352:204 353:109 354:4 356:12 357:197 358:251 359:31 372:221 373:251 374:253 375:121 379:36 380:23 385:190 386:251 387:31 399:48 400:234 401:253 413:191 414:253 415:31 426:44 427:221 428:251 429:251 440:12 441:197 442:251 443:31 454:190 455:251 456:251 457:251 468:96 469:251 470:251 471:31 482:190 483:251 484:251 485:113 495:40 496:234 497:251 498:219 499:23 510:190 511:251 512:251 513:94 522:40 523:217 524:253 525:231 526:47 538:191 539:253 540:253 541:253 548:12 549:174 550:253 551:253 552:219 553:39 566:67 567:236 568:251 569:251 570:191 571:190 572:111 573:72 574:190 575:191 576:197 577:251 578:243 579:121 580:39 595:63 596:236 597:251 598:253 599:251 600:251 601:251 602:251 603:253 604:251 605:188 606:94 624:27 625:129 626:253 627:251 628:251 629:251 630:251 631:229 632:168 633:15 654:95 655:212 656:251 657:211 658:94 659:59 +1 159:121 160:254 161:136 186:13 187:230 188:253 189:248 190:99 213:4 214:118 215:253 216:253 217:225 218:42 241:61 242:253 243:253 244:253 245:74 268:32 269:206 270:253 271:253 272:186 273:9 296:211 297:253 298:253 299:239 300:69 324:254 325:253 326:253 327:133 351:142 352:255 353:253 354:186 355:8 378:149 379:229 380:254 381:207 382:21 405:54 406:229 407:253 408:254 409:105 433:152 434:254 435:254 436:213 437:26 460:112 461:251 462:253 463:253 464:26 487:29 488:212 489:253 490:250 491:149 514:36 515:214 516:253 517:253 518:137 542:75 543:253 544:253 545:253 546:59 570:93 571:253 572:253 573:189 574:17 598:224 599:253 600:253 601:84 625:43 626:235 627:253 628:126 629:1 653:99 654:248 655:253 656:119 682:225 683:235 684:49 +1 100:166 101:222 102:55 128:197 129:254 130:218 131:5 155:29 156:249 157:254 158:254 159:9 183:45 184:254 185:254 186:174 187:2 210:4 211:164 212:254 213:254 214:85 238:146 239:254 240:254 241:254 242:85 265:101 266:245 267:254 268:254 269:254 270:85 292:97 293:248 294:254 295:204 296:254 297:254 298:85 315:12 316:59 317:98 318:151 319:237 320:254 321:254 322:109 323:35 324:254 325:254 326:85 343:41 344:216 345:254 346:254 347:239 348:153 349:37 350:4 351:32 352:254 353:254 354:85 372:7 373:44 374:44 375:30 379:32 380:254 381:254 382:96 407:19 408:230 409:254 410:174 436:197 437:254 438:110 464:197 465:254 466:85 492:197 493:253 494:63 515:37 516:54 517:54 518:45 519:26 520:84 521:221 522:84 523:21 524:31 525:162 526:78 540:6 541:41 542:141 543:244 544:254 545:254 546:248 547:236 548:254 549:254 550:254 551:233 552:239 553:254 554:138 567:23 568:167 569:254 570:254 571:254 572:254 573:229 574:228 575:185 576:138 577:138 578:138 579:138 580:138 581:138 582:44 595:113 596:254 597:254 598:254 599:179 600:64 601:5 623:32 624:209 625:183 626:97 +0 155:53 156:255 157:253 158:253 159:253 160:124 183:180 184:253 185:251 186:251 187:251 188:251 189:145 190:62 209:32 210:217 211:241 212:253 213:251 214:251 215:251 216:251 217:253 218:107 237:37 238:251 239:251 240:253 241:251 242:251 243:251 244:251 245:253 246:107 265:166 266:251 267:251 268:253 269:251 270:96 271:148 272:251 273:253 274:107 291:73 292:253 293:253 294:253 295:253 296:130 299:110 300:253 301:255 302:108 319:73 320:251 321:251 322:251 323:251 327:109 328:251 329:253 330:107 347:202 348:251 349:251 350:251 351:225 354:6 355:129 356:251 357:253 358:107 375:150 376:251 377:251 378:251 379:71 382:115 383:251 384:251 385:253 386:107 403:253 404:251 405:251 406:173 407:20 410:217 411:251 412:251 413:253 414:107 430:182 431:255 432:253 433:216 438:218 439:253 440:253 441:182 457:63 458:221 459:253 460:251 461:215 465:84 466:236 467:251 468:251 469:77 485:109 486:251 487:253 488:251 489:215 492:11 493:160 494:251 495:251 496:96 513:109 514:251 515:253 516:251 517:137 520:150 521:251 522:251 523:251 524:71 541:109 542:251 543:253 544:251 545:35 547:130 548:253 549:251 550:251 551:173 552:20 569:110 570:253 571:255 572:253 573:98 574:150 575:253 576:255 577:253 578:164 597:109 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:35 625:93 626:241 627:253 628:251 629:251 630:251 631:251 632:216 633:112 634:5 654:103 655:253 656:251 657:251 658:251 659:251 683:124 684:251 685:225 686:71 687:71 +0 128:73 129:253 130:227 131:73 132:21 156:73 157:251 158:251 159:251 160:174 182:16 183:166 184:228 185:251 186:251 187:251 188:122 210:62 211:220 212:253 213:251 214:251 215:251 216:251 217:79 238:79 239:231 240:253 241:251 242:251 243:251 244:251 245:232 246:77 264:145 265:253 266:253 267:253 268:255 269:253 270:253 271:253 272:253 273:255 274:108 292:144 293:251 294:251 295:251 296:253 297:168 298:107 299:169 300:251 301:253 302:189 303:20 318:27 319:89 320:236 321:251 322:235 323:215 324:164 325:15 326:6 327:129 328:251 329:253 330:251 331:35 345:47 346:211 347:253 348:251 349:251 350:142 354:37 355:251 356:251 357:253 358:251 359:35 373:109 374:251 375:253 376:251 377:251 378:142 382:11 383:148 384:251 385:253 386:251 387:164 400:11 401:150 402:253 403:255 404:211 405:25 410:11 411:150 412:253 413:255 414:211 415:25 428:140 429:251 430:251 431:253 432:107 438:37 439:251 440:251 441:211 442:46 456:190 457:251 458:251 459:253 460:128 461:5 466:37 467:251 468:251 469:51 484:115 485:251 486:251 487:253 488:188 489:20 492:32 493:109 494:129 495:251 496:173 497:103 512:217 513:251 514:251 515:201 516:30 520:73 521:251 522:251 523:251 524:71 540:166 541:253 542:253 543:255 544:149 545:73 546:150 547:253 548:255 549:253 550:253 551:143 568:140 569:251 570:251 571:253 572:251 573:251 574:251 575:251 576:253 577:251 578:230 579:61 596:190 597:251 598:251 599:253 600:251 601:251 602:251 603:251 604:242 605:215 606:55 624:21 625:189 626:251 627:253 628:251 629:251 630:251 631:173 632:103 653:31 654:200 655:253 656:251 657:96 658:71 659:20 +1 155:178 156:255 157:105 182:6 183:188 184:253 185:216 186:14 210:14 211:202 212:253 213:253 214:23 238:12 239:199 240:253 241:128 242:6 266:42 267:253 268:253 269:158 294:42 295:253 296:253 297:158 322:155 323:253 324:253 325:158 350:160 351:253 352:253 353:147 378:160 379:253 380:253 381:41 405:17 406:225 407:253 408:235 409:31 433:24 434:253 435:253 436:176 461:24 462:253 463:253 464:176 489:24 490:253 491:253 492:176 517:24 518:253 519:253 520:176 545:24 546:253 547:253 548:162 573:46 574:253 575:253 576:59 601:142 602:253 603:253 604:59 629:142 630:253 631:253 632:59 657:142 658:253 659:202 660:8 685:87 686:253 687:139 +0 154:46 155:105 156:254 157:254 158:254 159:254 160:255 161:239 162:41 180:37 181:118 182:222 183:254 184:253 185:253 186:253 187:253 188:253 189:253 190:211 191:54 207:14 208:200 209:253 210:253 211:254 212:253 213:253 214:253 215:253 216:253 217:253 218:253 219:116 233:16 234:160 235:236 236:253 237:253 238:253 239:254 240:253 241:253 242:246 243:229 244:253 245:253 246:253 247:116 261:99 262:253 263:253 264:253 265:253 266:253 267:254 268:253 269:253 270:213 271:99 272:253 273:253 274:253 275:116 288:25 289:194 290:253 291:253 292:253 293:253 294:131 295:97 296:169 297:253 298:93 299:99 300:253 301:253 302:253 303:116 316:206 317:253 318:253 319:251 320:233 321:127 322:9 324:18 325:38 326:3 327:15 328:171 329:253 330:253 331:116 343:55 344:240 345:253 346:253 347:233 355:31 356:186 357:253 358:253 359:116 371:176 372:253 373:253 374:253 375:127 383:99 384:253 385:253 386:253 387:116 399:176 400:253 401:253 402:131 403:9 411:99 412:253 413:253 414:253 415:116 426:119 427:254 428:254 429:232 430:75 440:158 441:254 442:254 443:117 454:118 455:253 456:253 457:154 468:156 469:253 470:253 471:116 482:118 483:253 484:253 485:154 496:156 497:253 498:253 499:116 509:46 510:222 511:253 512:253 513:154 522:7 523:116 524:246 525:253 526:180 527:9 538:118 539:253 540:253 541:154 550:116 551:253 552:253 553:253 554:174 566:118 567:253 568:253 569:154 577:110 578:246 579:253 580:253 581:240 582:67 594:118 595:253 596:253 597:238 598:215 599:49 600:20 601:20 602:20 603:66 604:215 605:241 606:253 607:245 608:233 609:64 622:82 623:229 624:253 625:253 626:253 627:253 628:253 629:253 630:253 631:254 632:253 633:253 634:240 635:107 651:176 652:253 653:253 654:253 655:253 656:253 657:253 658:253 659:254 660:253 661:253 662:108 679:40 680:239 681:253 682:253 683:253 684:253 685:253 686:253 687:254 688:161 689:57 690:4 +0 152:56 153:105 154:220 155:254 156:63 178:18 179:166 180:233 181:253 182:253 183:253 184:236 185:209 186:209 187:209 188:77 189:18 206:84 207:253 208:253 209:253 210:253 211:253 212:254 213:253 214:253 215:253 216:253 217:172 218:8 233:57 234:238 235:253 236:253 237:253 238:253 239:253 240:254 241:253 242:253 243:253 244:253 245:253 246:119 260:14 261:238 262:253 263:253 264:253 265:253 266:253 267:253 268:179 269:196 270:253 271:253 272:253 273:253 274:238 275:12 288:33 289:253 290:253 291:253 292:253 293:253 294:248 295:134 297:18 298:83 299:237 300:253 301:253 302:253 303:14 316:164 317:253 318:253 319:253 320:253 321:253 322:128 327:57 328:119 329:214 330:253 331:94 343:57 344:248 345:253 346:253 347:253 348:126 349:14 350:4 357:179 358:253 359:248 360:56 371:175 372:253 373:253 374:240 375:190 376:28 385:179 386:253 387:253 388:173 399:209 400:253 401:253 402:178 413:92 414:253 415:253 416:208 427:211 428:254 429:254 430:179 442:135 443:255 444:209 455:209 456:253 457:253 458:90 470:134 471:253 472:208 483:209 484:253 485:253 486:178 497:2 498:142 499:253 500:208 511:209 512:253 513:253 514:214 515:35 525:30 526:253 527:253 528:208 539:165 540:253 541:253 542:253 543:215 544:36 553:163 554:253 555:253 556:164 567:18 568:172 569:253 570:253 571:253 572:214 573:127 574:7 580:72 581:232 582:253 583:171 584:17 596:8 597:182 598:253 599:253 600:253 601:253 602:162 603:56 607:64 608:240 609:253 610:253 611:14 625:7 626:173 627:253 628:253 629:253 630:253 631:245 632:241 633:239 634:239 635:246 636:253 637:225 638:14 639:1 654:18 655:59 656:138 657:224 658:253 659:253 660:254 661:253 662:253 663:253 664:240 665:96 685:37 686:104 687:192 688:255 689:253 690:253 691:182 692:73 +1 130:7 131:176 132:254 133:224 158:51 159:253 160:253 161:223 185:4 186:170 187:253 188:253 189:214 213:131 214:253 215:253 216:217 217:39 241:209 242:253 243:253 244:134 268:75 269:240 270:253 271:239 272:26 296:184 297:253 298:245 299:63 323:142 324:255 325:253 326:185 350:62 351:229 352:254 353:242 354:73 377:54 378:229 379:253 380:254 381:105 405:152 406:254 407:254 408:213 409:26 432:32 433:243 434:253 435:253 436:115 459:2 460:142 461:253 462:253 463:155 487:30 488:253 489:253 490:232 491:55 515:75 516:253 517:253 518:164 542:72 543:232 544:253 545:189 546:17 570:224 571:253 572:253 573:163 597:43 598:235 599:253 600:253 601:195 602:21 625:28 626:231 627:253 628:253 629:184 630:14 654:225 655:253 656:253 657:75 +0 155:21 156:176 157:253 158:253 159:124 182:105 183:176 184:251 185:251 186:251 187:251 188:105 208:58 209:217 210:241 211:253 212:251 213:251 214:251 215:251 216:243 217:113 218:5 235:63 236:231 237:251 238:251 239:253 240:251 241:251 242:251 243:251 244:253 245:251 246:113 263:144 264:251 265:251 266:251 267:253 268:251 269:251 270:251 271:251 272:253 273:251 274:215 290:125 291:253 292:253 293:253 294:253 295:255 296:253 297:253 298:253 299:253 300:255 301:253 302:227 303:42 318:253 319:251 320:251 321:251 322:251 323:253 324:251 325:251 326:251 327:251 328:253 329:251 330:251 331:142 345:27 346:253 347:251 348:251 349:235 350:241 351:253 352:251 353:246 354:137 355:35 356:98 357:251 358:251 359:236 360:61 372:47 373:211 374:253 375:251 376:235 377:82 378:103 379:253 380:251 381:137 384:73 385:251 386:251 387:251 388:71 399:27 400:211 401:251 402:253 403:251 404:86 407:72 408:71 409:10 412:73 413:251 414:251 415:173 416:20 427:89 428:253 429:253 430:255 431:253 432:35 440:73 441:253 442:253 443:253 444:72 454:84 455:236 456:251 457:251 458:253 459:251 460:138 468:73 469:251 470:251 471:251 472:71 481:63 482:236 483:251 484:251 485:251 486:227 487:251 488:246 489:138 490:11 494:16 495:37 496:228 497:251 498:246 499:137 500:10 509:73 510:251 511:251 512:251 513:173 514:42 515:142 516:142 517:142 518:41 522:109 523:251 524:253 525:251 526:137 537:73 538:251 539:251 540:173 541:20 549:27 550:211 551:251 552:253 553:147 554:10 565:73 566:253 567:253 568:143 575:21 576:176 577:253 578:253 579:253 593:73 594:251 595:251 596:205 597:144 603:176 604:251 605:251 606:188 607:107 621:62 622:236 623:251 624:251 625:251 626:218 627:217 628:217 629:217 630:217 631:253 632:230 633:189 634:20 650:83 651:158 652:251 653:251 654:253 655:251 656:251 657:251 658:251 659:253 660:107 679:37 680:251 681:251 682:253 683:251 684:251 685:251 686:122 687:72 688:30 +1 151:68 152:45 153:131 154:131 155:131 156:101 157:68 158:92 159:44 187:19 188:170 211:29 212:112 213:89 215:40 216:222 239:120 240:254 241:251 242:127 243:40 244:222 267:197 268:254 269:254 270:91 271:40 272:222 294:64 295:247 296:254 297:236 298:50 299:40 300:107 322:184 323:254 324:254 325:91 327:6 328:14 350:203 351:254 352:254 353:71 377:23 378:218 379:254 380:254 381:71 405:113 406:254 407:255 408:239 409:53 433:210 434:254 435:254 436:195 460:62 461:242 462:254 463:241 464:88 468:28 488:86 489:254 490:254 491:189 495:28 496:104 516:106 517:254 518:254 519:168 523:40 524:91 544:216 545:254 546:245 547:51 551:35 552:80 572:216 573:254 574:102 599:55 600:239 601:254 602:52 627:166 628:254 629:210 630:23 655:223 656:252 657:104 683:223 684:169 +0 125:29 126:170 127:255 128:255 129:141 151:29 152:198 153:255 154:255 155:255 156:226 157:255 158:86 178:141 179:255 180:255 181:170 182:29 184:86 185:255 186:255 187:141 204:29 205:226 206:255 207:198 208:57 213:226 214:255 215:255 216:226 217:114 231:29 232:255 233:255 234:114 241:141 242:170 243:114 244:255 245:255 246:141 259:226 260:255 261:170 269:29 270:57 273:141 274:255 275:226 286:57 287:255 288:170 302:114 303:255 304:198 314:226 315:255 331:170 332:255 333:57 342:255 343:226 360:255 361:170 370:255 371:170 388:114 389:198 398:255 399:226 416:86 417:255 426:198 427:255 444:86 445:255 454:114 455:255 456:57 472:86 473:255 482:29 483:255 484:226 500:141 501:255 511:170 512:255 513:170 528:226 529:198 539:29 540:226 541:255 542:170 555:29 556:255 557:114 568:29 569:226 570:255 571:141 582:57 583:226 584:226 598:141 599:255 600:255 601:170 602:86 607:29 608:86 609:226 610:255 611:226 612:29 627:86 628:198 629:255 630:255 631:255 632:255 633:255 634:255 635:255 636:255 637:255 638:141 639:29 657:29 658:114 659:170 660:170 661:170 662:170 663:170 664:86 +0 153:203 154:254 155:252 156:252 157:252 158:214 159:51 160:20 180:62 181:221 182:252 183:250 184:250 185:250 186:252 187:250 188:160 189:20 207:62 208:211 209:250 210:252 211:250 212:250 213:250 214:252 215:250 216:250 217:49 234:41 235:221 236:250 237:250 238:252 239:250 240:250 241:250 242:252 243:250 244:128 245:10 262:254 263:252 264:252 265:252 266:254 267:252 268:252 269:252 270:254 271:252 272:252 273:90 290:150 291:190 292:250 293:250 294:252 295:250 296:250 297:169 298:171 299:250 300:250 301:250 302:82 318:31 319:191 320:250 321:250 322:252 323:189 324:100 325:20 326:172 327:250 328:250 329:250 330:80 346:213 347:250 348:250 349:250 350:212 351:29 354:252 355:250 356:250 357:250 374:92 375:252 376:252 377:252 382:51 383:252 384:252 385:252 386:203 401:82 402:252 403:250 404:250 405:169 410:132 411:250 412:250 413:250 414:121 428:92 429:231 430:252 431:250 432:159 433:20 438:252 439:250 440:250 441:250 456:30 457:211 458:252 459:250 460:221 461:40 466:90 467:250 468:250 469:250 470:163 484:31 485:213 486:254 487:232 488:80 494:92 495:252 496:252 497:212 498:163 512:151 513:250 514:252 515:149 522:252 523:250 524:250 525:49 540:60 541:221 542:252 543:210 544:60 550:252 551:250 552:250 553:49 569:202 570:252 571:250 572:221 573:40 576:123 577:202 578:252 579:250 580:250 581:49 596:123 597:243 598:255 599:252 600:252 601:252 602:254 603:252 604:252 605:252 606:254 607:252 608:100 625:121 626:171 627:250 628:250 629:250 630:252 631:250 632:250 633:250 634:252 635:250 636:100 654:20 655:160 656:250 657:250 658:252 659:250 660:250 661:250 662:252 663:189 664:40 683:20 684:170 685:250 686:252 687:250 688:128 689:49 690:49 691:29 +1 98:64 99:191 100:70 125:68 126:243 127:253 128:249 129:63 152:30 153:223 154:253 155:253 156:247 157:41 179:73 180:238 181:253 182:253 183:253 184:242 206:73 207:236 208:253 209:253 210:253 211:253 212:242 234:182 235:253 236:253 237:191 238:247 239:253 240:149 262:141 263:253 264:143 265:86 266:249 267:253 268:122 290:9 291:36 292:7 293:14 294:233 295:253 296:122 322:230 323:253 324:122 350:230 351:253 352:122 378:231 379:255 380:123 406:230 407:253 408:52 433:61 434:245 435:253 461:98 462:253 463:253 468:35 469:12 489:98 490:253 491:253 494:9 495:142 496:233 497:146 517:190 518:253 519:253 520:128 521:7 522:99 523:253 524:253 525:180 544:29 545:230 546:253 547:253 548:252 549:210 550:253 551:253 552:253 553:140 571:28 572:207 573:253 574:253 575:253 576:254 577:253 578:253 579:235 580:70 581:9 599:126 600:253 601:253 602:253 603:253 604:254 605:253 606:168 607:19 627:79 628:253 629:253 630:201 631:190 632:132 633:63 634:5 +1 125:26 126:240 127:72 153:25 154:238 155:208 182:209 183:226 184:14 210:209 211:254 212:43 238:175 239:254 240:128 266:63 267:254 268:204 294:107 295:254 296:204 322:88 323:254 324:204 350:55 351:254 352:204 378:126 379:254 380:204 406:126 407:254 408:189 434:169 435:254 436:121 462:209 463:254 464:193 490:209 491:254 492:111 517:22 518:235 519:254 520:37 545:137 546:254 547:227 548:16 573:205 574:255 575:185 601:205 602:254 603:125 629:205 630:254 631:125 657:111 658:212 659:43 +0 155:62 156:91 157:213 158:255 159:228 160:91 161:12 182:70 183:230 184:253 185:253 186:253 187:253 188:253 189:152 190:7 210:246 211:253 212:253 213:253 214:253 215:253 216:253 217:253 218:106 237:21 238:247 239:253 240:253 241:253 242:253 243:253 244:253 245:208 246:24 265:156 266:253 267:253 268:253 269:253 270:253 271:253 272:253 273:195 292:88 293:238 294:253 295:253 296:253 297:221 298:253 299:253 300:253 301:195 320:230 321:253 322:253 323:253 324:198 325:40 326:177 327:253 328:253 329:195 346:56 347:156 348:251 349:253 350:189 351:182 352:15 354:86 355:240 356:253 357:210 358:28 374:213 375:253 376:253 377:156 378:3 383:205 384:253 385:253 386:106 401:121 402:252 403:253 404:135 405:3 411:46 412:253 413:253 414:106 428:28 429:212 430:253 431:248 432:23 439:42 440:253 441:253 442:106 456:197 457:253 458:234 459:70 467:42 468:253 469:253 470:106 483:11 484:202 485:253 486:187 495:58 496:253 497:210 498:27 511:107 512:253 513:253 514:40 522:53 523:227 524:253 525:195 539:107 540:253 541:253 542:40 549:47 550:227 551:253 552:231 553:58 567:107 568:253 569:253 570:40 575:5 576:131 577:222 578:253 579:231 580:59 595:14 596:204 597:253 598:226 599:222 600:73 601:58 602:58 603:170 604:253 605:253 606:227 607:58 624:197 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:238 634:58 652:33 653:179 654:241 655:253 656:253 657:253 658:253 659:250 660:116 661:14 682:75 683:179 684:253 685:151 686:89 687:86 +1 157:42 158:228 159:253 160:253 185:144 186:251 187:251 188:251 212:89 213:236 214:251 215:235 216:215 239:79 240:253 241:251 242:251 243:142 267:180 268:253 269:251 270:251 271:142 294:32 295:202 296:255 297:253 298:216 322:109 323:251 324:253 325:251 326:112 349:6 350:129 351:251 352:253 353:127 354:5 377:37 378:251 379:251 380:253 381:107 405:166 406:251 407:251 408:201 409:30 432:42 433:228 434:253 435:253 460:144 461:251 462:251 463:147 487:63 488:236 489:251 490:251 491:71 515:150 516:251 517:251 518:204 519:41 543:253 544:251 545:251 546:142 571:255 572:253 573:164 598:105 599:253 600:251 601:35 626:180 627:253 628:251 629:35 654:180 655:253 656:251 657:35 682:180 683:253 684:251 685:35 +1 128:62 129:254 130:213 156:102 157:253 158:252 159:102 160:20 184:102 185:254 186:253 187:254 188:50 212:102 213:253 214:252 215:253 216:50 240:102 241:254 242:253 243:254 244:50 268:142 269:253 270:252 271:253 272:50 295:51 296:253 297:254 298:253 299:224 300:20 323:132 324:252 325:253 326:252 327:162 351:173 352:253 353:254 354:253 355:102 378:82 379:253 380:252 381:253 382:252 383:61 406:203 407:254 408:253 409:254 410:233 433:41 434:243 435:253 436:252 437:253 438:111 461:132 462:253 463:254 464:253 465:203 488:41 489:253 490:252 491:253 492:252 493:40 515:11 516:213 517:254 518:253 519:254 520:151 543:92 544:252 545:253 546:252 547:192 548:50 570:21 571:214 572:253 573:255 574:253 575:41 598:142 599:253 600:252 601:253 602:171 625:113 626:253 627:255 628:253 629:203 630:40 653:30 654:131 655:233 656:111 +0 154:28 155:195 156:254 157:254 158:254 159:254 160:254 161:255 162:61 181:6 182:191 183:253 184:253 185:253 186:253 187:253 188:253 189:253 190:60 208:26 209:190 210:253 211:253 212:253 213:253 214:240 215:191 216:242 217:253 218:60 235:15 236:187 237:253 238:253 239:253 240:253 241:253 242:200 244:211 245:253 246:60 262:22 263:66 264:253 265:253 266:253 267:253 268:241 269:209 270:44 271:23 272:218 273:253 274:60 290:124 291:253 292:253 293:253 294:253 295:253 296:182 299:131 300:253 301:253 302:60 318:38 319:217 320:253 321:253 322:244 323:111 324:37 327:131 328:253 329:253 330:60 346:124 347:253 348:253 349:253 350:165 354:22 355:182 356:253 357:253 358:60 374:124 375:253 376:253 377:240 378:45 382:53 383:253 384:253 385:249 386:58 401:16 402:168 403:253 404:216 405:45 410:53 411:253 412:253 413:138 429:159 430:253 431:253 432:147 438:53 439:253 440:253 441:138 456:136 457:252 458:253 459:227 460:5 466:53 467:253 468:243 469:101 484:140 485:253 486:253 487:124 494:156 495:253 496:218 511:13 512:164 513:253 514:142 515:5 521:32 522:233 523:253 524:218 539:62 540:253 541:253 542:130 548:37 549:203 550:253 551:253 552:127 567:62 568:253 569:253 570:147 571:36 572:36 573:36 574:36 575:151 576:222 577:253 578:245 579:127 580:8 595:34 596:202 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:200 624:140 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:248 633:235 634:65 652:87 653:173 654:253 655:253 656:253 657:253 658:253 659:253 660:182 681:14 682:78 683:96 684:253 685:253 686:253 687:137 688:56 +0 123:8 124:76 125:202 126:254 127:255 128:163 129:37 130:2 150:13 151:182 152:253 153:253 154:253 155:253 156:253 157:253 158:23 177:15 178:179 179:253 180:253 181:212 182:91 183:218 184:253 185:253 186:179 187:109 205:105 206:253 207:253 208:160 209:35 210:156 211:253 212:253 213:253 214:253 215:250 216:113 232:19 233:212 234:253 235:253 236:88 237:121 238:253 239:233 240:128 241:91 242:245 243:253 244:248 245:114 260:104 261:253 262:253 263:110 264:2 265:142 266:253 267:90 270:26 271:199 272:253 273:248 274:63 287:1 288:173 289:253 290:253 291:29 293:84 294:228 295:39 299:72 300:251 301:253 302:215 303:29 315:36 316:253 317:253 318:203 319:13 328:82 329:253 330:253 331:170 343:36 344:253 345:253 346:164 356:11 357:198 358:253 359:184 360:6 371:36 372:253 373:253 374:82 385:138 386:253 387:253 388:35 399:128 400:253 401:253 402:47 413:48 414:253 415:253 416:35 427:154 428:253 429:253 430:47 441:48 442:253 443:253 444:35 455:102 456:253 457:253 458:99 469:48 470:253 471:253 472:35 483:36 484:253 485:253 486:164 496:16 497:208 498:253 499:211 500:17 511:32 512:244 513:253 514:175 515:4 524:44 525:253 526:253 527:156 540:171 541:253 542:253 543:29 551:30 552:217 553:253 554:188 555:19 568:171 569:253 570:253 571:59 578:60 579:217 580:253 581:253 582:70 596:78 597:253 598:253 599:231 600:48 604:26 605:128 606:249 607:253 608:244 609:94 610:15 624:8 625:151 626:253 627:253 628:234 629:101 630:121 631:219 632:229 633:253 634:253 635:201 636:80 653:38 654:232 655:253 656:253 657:253 658:253 659:253 660:253 661:253 662:201 663:66 +0 127:68 128:254 129:255 130:254 131:107 153:11 154:176 155:230 156:253 157:253 158:253 159:212 180:28 181:197 182:253 183:253 184:253 185:253 186:253 187:229 188:107 189:14 208:194 209:253 210:253 211:253 212:253 213:253 214:253 215:253 216:253 217:53 235:69 236:241 237:253 238:253 239:253 240:253 241:241 242:186 243:253 244:253 245:195 262:10 263:161 264:253 265:253 266:253 267:246 268:40 269:57 270:231 271:253 272:253 273:195 290:140 291:253 292:253 293:253 294:253 295:154 297:25 298:253 299:253 300:253 301:195 318:213 319:253 320:253 321:253 322:135 323:8 325:3 326:128 327:253 328:253 329:195 345:77 346:238 347:253 348:253 349:253 350:7 354:116 355:253 356:253 357:195 372:11 373:165 374:253 375:253 376:231 377:70 378:1 382:78 383:237 384:253 385:195 400:33 401:253 402:253 403:253 404:182 411:200 412:253 413:195 428:98 429:253 430:253 431:253 432:24 439:42 440:253 441:195 456:197 457:253 458:253 459:253 460:24 467:163 468:253 469:195 484:197 485:253 486:253 487:189 488:13 494:53 495:227 496:253 497:121 512:197 513:253 514:253 515:114 521:21 522:227 523:253 524:231 525:27 540:197 541:253 542:253 543:114 547:5 548:131 549:143 550:253 551:231 552:59 568:197 569:253 570:253 571:236 572:73 573:58 574:217 575:223 576:253 577:253 578:253 579:174 596:197 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:253 607:48 624:149 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:182 634:15 635:3 652:12 653:168 654:253 655:253 656:253 657:253 658:253 659:248 660:89 661:23 +1 157:85 158:255 159:103 160:1 185:205 186:253 187:253 188:30 213:205 214:253 215:253 216:30 240:44 241:233 242:253 243:244 244:27 268:135 269:253 270:253 271:100 296:153 297:253 298:240 299:76 323:12 324:208 325:253 326:166 351:69 352:253 353:253 354:142 378:14 379:110 380:253 381:235 382:33 406:63 407:223 408:235 409:130 434:186 435:253 436:235 437:37 461:17 462:145 463:253 464:231 465:35 489:69 490:220 491:231 492:123 516:18 517:205 518:253 519:176 520:27 543:17 544:125 545:253 546:185 547:39 571:71 572:214 573:231 574:41 599:167 600:253 601:225 602:33 626:72 627:205 628:207 629:14 653:30 654:249 655:233 656:49 681:32 682:253 683:89 +1 126:94 127:132 154:250 155:250 156:4 182:250 183:254 184:95 210:250 211:254 212:95 238:250 239:254 240:95 266:250 267:254 268:95 294:250 295:254 296:95 322:250 323:254 324:95 350:250 351:254 352:95 378:250 379:254 380:95 405:77 406:254 407:250 408:19 433:96 434:254 435:249 461:53 462:253 463:252 464:43 490:250 491:251 492:32 517:85 518:254 519:249 545:96 546:254 547:249 573:83 574:254 575:250 576:14 602:250 603:254 604:95 630:250 631:255 632:95 658:132 659:254 660:95 +1 124:32 125:253 126:31 152:32 153:251 154:149 180:32 181:251 182:188 208:32 209:251 210:188 236:32 237:251 238:228 239:59 264:32 265:253 266:253 267:95 292:28 293:236 294:251 295:114 321:127 322:251 323:251 349:127 350:251 351:251 377:48 378:232 379:251 406:223 407:253 408:159 434:221 435:251 436:158 462:142 463:251 464:158 490:64 491:251 492:242 493:55 518:64 519:251 520:253 521:161 546:64 547:253 548:255 549:221 574:16 575:181 576:253 577:220 603:79 604:253 605:236 606:63 632:213 633:251 634:126 660:96 661:251 662:126 +1 129:39 130:254 131:255 132:254 133:140 157:136 158:253 159:253 160:228 161:67 184:6 185:227 186:253 187:253 188:58 211:29 212:188 213:253 214:253 215:253 216:17 239:95 240:253 241:253 242:253 243:157 244:8 266:3 267:107 268:253 269:253 270:245 271:77 294:29 295:253 296:253 297:240 298:100 322:141 323:253 324:253 325:215 349:129 350:248 351:253 352:253 353:215 377:151 378:253 379:253 380:253 381:144 405:151 406:253 407:253 408:253 409:27 431:3 432:102 433:242 434:253 435:253 436:110 437:3 459:97 460:253 461:253 462:253 463:214 464:55 487:207 488:253 489:253 490:253 491:158 515:67 516:253 517:253 518:253 519:158 543:207 544:253 545:253 546:240 547:88 571:207 572:253 573:253 574:224 598:32 599:217 600:253 601:253 602:224 626:141 627:253 628:253 629:253 630:133 654:36 655:219 656:253 657:140 658:10 +0 123:59 124:55 149:71 150:192 151:254 152:250 153:147 154:17 176:123 177:247 178:253 179:254 180:253 181:253 182:196 183:79 184:176 185:175 186:175 187:124 188:48 203:87 204:247 205:247 206:176 207:95 208:102 209:117 210:243 211:237 212:192 213:232 214:253 215:253 216:245 217:152 218:6 230:23 231:229 232:253 233:138 238:219 239:58 241:95 242:118 243:80 244:230 245:254 246:196 247:30 258:120 259:254 260:205 261:8 266:114 272:38 273:255 274:254 275:155 276:5 286:156 287:253 288:92 301:61 302:235 303:253 304:102 314:224 315:253 316:78 330:117 331:253 332:196 333:18 342:254 343:253 344:78 358:9 359:211 360:253 361:73 370:254 371:253 372:78 387:175 388:253 389:155 398:194 399:254 400:101 415:79 416:254 417:155 426:112 427:253 428:211 429:9 443:73 444:251 445:200 454:41 455:241 456:253 457:87 471:25 472:240 473:253 483:147 484:253 485:227 486:47 499:94 500:253 501:200 511:5 512:193 513:253 514:230 515:76 527:175 528:253 529:155 540:31 541:219 542:254 543:255 544:126 545:18 553:14 554:149 555:254 556:244 557:45 569:21 570:158 571:254 572:253 573:226 574:162 575:118 576:96 577:20 578:20 579:73 580:118 581:224 582:253 583:247 584:85 598:30 599:155 600:253 601:253 602:253 603:253 604:254 605:253 606:253 607:253 608:253 609:254 610:247 611:84 627:5 628:27 629:117 630:206 631:244 632:229 633:213 634:213 635:213 636:176 637:117 638:32 659:45 660:23 +1 128:58 129:139 156:247 157:247 158:25 183:121 184:253 185:156 186:3 211:133 212:253 213:145 238:11 239:227 240:253 241:145 266:7 267:189 268:253 269:145 294:35 295:252 296:253 297:145 322:146 323:252 324:253 325:131 350:146 351:252 352:253 353:13 378:146 379:252 380:253 381:13 406:147 407:253 408:255 409:13 434:146 435:252 436:253 437:13 462:146 463:252 464:253 465:13 490:146 491:252 492:253 493:13 517:22 518:230 519:252 520:221 521:9 545:22 546:230 547:252 548:133 574:146 575:252 576:133 602:146 603:252 604:120 630:146 631:252 658:146 659:252 +1 129:28 130:247 131:255 132:165 156:47 157:221 158:252 159:252 160:164 184:177 185:252 186:252 187:252 188:164 212:177 213:252 214:252 215:223 216:78 240:177 241:252 242:252 243:197 267:114 268:236 269:252 270:235 271:42 294:5 295:148 296:252 297:252 298:230 321:14 322:135 323:252 324:252 325:252 326:230 349:78 350:252 351:252 352:252 353:252 354:162 377:78 378:252 379:252 380:252 381:252 382:9 405:78 406:252 407:252 408:252 409:252 410:9 432:32 433:200 434:252 435:252 436:252 437:105 438:3 459:10 460:218 461:252 462:252 463:252 464:105 465:8 487:225 488:252 489:252 490:252 491:240 492:69 514:44 515:237 516:252 517:252 518:228 519:85 541:59 542:218 543:252 544:252 545:225 546:93 568:65 569:208 570:252 571:252 572:252 573:175 596:133 597:252 598:252 599:252 600:225 601:68 624:133 625:252 626:252 627:244 628:54 652:133 653:252 654:252 655:48 +0 156:13 157:6 181:10 182:77 183:145 184:253 185:190 186:67 207:11 208:77 209:193 210:252 211:252 212:253 213:252 214:238 215:157 216:71 217:26 233:10 234:78 235:193 236:252 237:252 238:252 239:252 240:253 241:252 242:252 243:252 244:252 245:228 246:128 247:49 248:5 259:6 260:78 261:194 262:252 263:252 264:252 265:252 266:252 267:252 268:253 269:217 270:192 271:232 272:252 273:252 274:252 275:252 276:135 277:3 286:4 287:147 288:252 289:252 290:252 291:252 292:252 293:252 294:252 295:252 296:175 297:26 299:40 300:145 301:235 302:252 303:252 304:252 305:104 314:208 315:252 316:252 317:252 318:252 319:252 320:252 321:133 322:48 323:48 329:71 330:236 331:252 332:252 333:230 342:253 343:185 344:170 345:252 346:252 347:252 348:173 349:22 358:102 359:252 360:252 361:252 370:24 371:141 372:243 373:252 374:252 375:186 376:5 386:8 387:220 388:252 389:252 398:70 399:247 400:252 401:252 402:165 403:37 414:81 415:251 416:252 417:194 426:255 427:253 428:253 429:251 430:69 441:39 442:231 443:253 444:253 445:127 454:253 455:252 456:249 457:127 468:6 469:147 470:252 471:252 472:190 473:5 482:253 483:252 484:216 495:7 496:145 497:252 498:252 499:252 500:69 510:253 511:252 512:223 513:16 522:25 523:185 524:252 525:252 526:252 527:107 528:8 538:167 539:252 540:252 541:181 542:18 549:105 550:191 551:252 552:252 553:235 554:151 555:10 566:37 567:221 568:252 569:252 570:210 571:193 572:96 573:73 574:130 575:188 576:194 577:227 578:252 579:252 580:235 581:128 595:97 596:220 597:252 598:252 599:252 600:252 601:252 602:252 603:252 604:253 605:252 606:252 607:236 608:70 624:40 625:174 626:252 627:252 628:252 629:252 630:252 631:252 632:253 633:197 634:138 635:29 653:5 654:23 655:116 656:143 657:143 658:143 659:143 660:24 661:10 +0 127:28 128:164 129:254 130:233 131:148 132:11 154:3 155:164 156:254 157:234 158:225 159:254 160:204 182:91 183:254 184:235 185:48 186:32 187:166 188:251 189:92 208:33 209:111 210:214 211:205 212:49 215:24 216:216 217:210 235:34 236:217 237:254 238:254 239:211 244:87 245:237 246:43 262:34 263:216 264:254 265:254 266:252 267:243 268:61 272:38 273:248 274:182 290:171 291:254 292:184 293:205 294:175 295:36 301:171 302:227 317:28 318:234 319:190 320:13 321:193 322:157 329:124 330:238 331:26 345:140 346:254 347:131 349:129 350:157 357:124 358:254 359:95 373:201 374:238 375:56 377:70 378:103 385:124 386:254 387:148 400:62 401:255 402:210 413:150 414:254 415:122 428:86 429:254 430:201 431:15 440:28 441:237 442:246 443:44 456:128 457:254 458:143 468:34 469:243 470:227 484:62 485:254 486:210 496:58 497:249 498:179 512:30 513:240 514:210 524:207 525:254 526:64 541:216 542:231 543:34 551:129 552:248 553:170 554:9 569:131 570:254 571:170 577:17 578:129 579:248 580:225 581:24 597:50 598:245 599:245 600:184 601:106 602:106 603:106 604:133 605:231 606:254 607:244 608:53 626:67 627:249 628:254 629:254 630:254 631:254 632:254 633:251 634:193 635:40 655:38 656:157 657:248 658:166 659:166 660:139 661:57 +0 129:105 130:255 131:219 132:67 133:67 134:52 156:20 157:181 158:253 159:253 160:253 161:253 162:226 163:69 182:4 183:129 184:206 185:253 186:253 187:253 188:253 189:253 190:253 191:130 209:9 210:141 211:253 212:253 213:253 214:253 215:253 216:253 217:253 218:253 219:166 220:20 237:134 238:253 239:253 240:253 241:253 242:253 243:253 244:253 245:253 246:253 247:253 248:65 262:2 263:83 264:207 265:246 266:253 267:253 268:253 269:253 270:253 271:249 272:234 273:247 274:253 275:253 276:65 290:83 291:253 292:253 293:253 294:253 295:253 296:189 297:253 298:253 299:205 301:179 302:253 303:253 304:65 317:85 318:234 319:253 320:253 321:253 322:253 323:157 324:26 325:164 326:151 327:83 329:179 330:253 331:253 332:65 344:65 345:237 346:253 347:253 348:253 349:67 350:36 351:14 353:15 354:12 357:179 358:253 359:253 360:65 371:4 372:141 373:253 374:253 375:221 376:158 377:23 385:179 386:253 387:253 388:65 399:129 400:253 401:253 402:241 403:62 412:72 413:226 414:253 415:175 416:24 426:119 427:247 428:253 429:253 430:206 439:8 440:134 441:253 442:253 443:130 454:132 455:253 456:253 457:194 458:27 467:125 468:253 469:253 470:253 471:130 481:45 482:213 483:253 484:253 485:112 493:70 494:170 495:247 496:253 497:253 498:89 499:43 509:67 510:253 511:253 512:196 513:55 514:9 520:8 521:131 522:253 523:253 524:253 525:86 526:1 537:67 538:253 539:253 540:253 541:253 542:129 546:43 547:114 548:134 549:253 550:253 551:231 552:139 553:41 565:20 566:167 567:253 568:253 569:253 570:247 571:179 572:179 573:179 574:206 575:253 576:253 577:253 578:253 579:72 594:103 595:240 596:253 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:244 605:119 606:8 607:1 623:107 624:253 625:253 626:253 627:253 628:253 629:253 630:253 631:175 632:111 651:3 652:121 653:253 654:253 655:253 656:253 657:253 658:182 659:24 +0 125:22 126:183 127:252 128:254 129:252 130:252 131:252 132:76 151:85 152:85 153:168 154:250 155:250 156:252 157:250 158:250 159:250 160:250 161:71 163:43 164:85 165:14 178:107 179:252 180:250 181:250 182:250 183:250 184:252 185:250 186:250 187:250 188:250 189:210 191:127 192:250 193:146 205:114 206:237 207:252 208:250 209:250 210:250 211:250 212:252 213:250 214:250 215:250 216:250 217:210 219:127 220:250 221:250 232:107 233:237 234:250 235:252 236:250 237:250 238:250 239:74 240:41 241:41 242:41 243:41 244:217 245:34 247:127 248:250 249:250 259:15 260:148 261:252 262:252 263:254 264:238 265:105 275:128 276:252 277:252 286:15 287:140 288:250 289:250 290:250 291:167 292:111 303:127 304:250 305:250 314:43 315:250 316:250 317:250 318:250 331:127 332:250 333:250 342:183 343:250 344:250 345:250 346:110 358:57 359:210 360:250 361:250 370:252 371:250 372:250 373:110 374:7 386:85 387:250 388:250 389:250 398:254 399:252 400:252 401:83 414:86 415:252 416:252 417:217 426:252 427:250 428:250 429:138 430:14 441:15 442:140 443:250 444:250 445:41 454:252 455:250 456:250 457:250 458:41 469:43 470:250 471:250 472:250 473:41 482:252 483:250 484:250 485:250 486:181 497:183 498:250 499:250 500:250 501:41 510:76 511:250 512:250 513:250 514:250 524:177 525:252 526:250 527:250 528:110 529:7 538:36 539:224 540:252 541:252 542:252 543:219 544:43 545:43 546:43 547:7 549:15 550:43 551:183 552:252 553:255 554:252 555:126 567:85 568:250 569:250 570:250 571:252 572:250 573:250 574:250 575:111 576:86 577:140 578:250 579:250 580:250 581:252 582:222 583:83 595:42 596:188 597:250 598:250 599:252 600:250 601:250 602:250 603:250 604:252 605:250 606:250 607:250 608:250 609:126 610:83 624:127 625:250 626:250 627:252 628:250 629:250 630:250 631:250 632:252 633:250 634:250 635:137 636:83 652:21 653:41 654:217 655:252 656:250 657:250 658:250 659:250 660:217 661:41 662:41 663:14 +1 155:114 156:206 157:25 183:238 184:252 185:55 211:222 212:252 213:55 239:113 240:252 241:55 267:113 268:252 269:55 295:255 296:253 297:56 323:253 324:176 325:6 350:32 351:253 352:233 353:43 378:140 379:253 380:195 381:19 406:140 407:253 408:167 433:29 434:253 435:141 461:29 462:252 463:140 489:29 490:252 491:140 517:29 518:252 519:140 545:29 546:252 547:140 573:169 574:253 575:79 601:169 602:252 628:76 629:234 630:141 656:197 657:233 658:37 684:197 685:223 +1 127:73 128:253 129:253 130:63 155:115 156:252 157:252 158:144 183:217 184:252 185:252 186:144 210:63 211:237 212:252 213:252 214:144 238:109 239:252 240:252 241:252 266:109 267:252 268:252 269:252 294:109 295:252 296:252 297:252 322:191 323:252 324:252 325:252 349:145 350:255 351:253 352:253 353:253 376:32 377:237 378:253 379:252 380:252 381:210 404:37 405:252 406:253 407:252 408:252 409:108 432:37 433:252 434:253 435:252 436:252 437:108 460:21 461:207 462:255 463:253 464:253 465:108 489:144 490:253 491:252 492:252 493:108 516:27 517:221 518:253 519:252 520:252 521:108 544:16 545:190 546:253 547:252 548:252 549:108 573:145 574:255 575:253 576:253 577:253 601:144 602:253 603:252 604:252 605:210 629:144 630:253 631:252 632:252 633:108 657:62 658:253 659:252 660:252 661:108 +1 120:85 121:253 122:132 123:9 147:82 148:241 149:251 150:251 151:128 175:175 176:251 177:251 178:251 179:245 180:121 203:13 204:204 205:251 206:251 207:251 208:245 209:107 232:39 233:251 234:251 235:251 236:251 237:167 238:22 260:15 261:155 262:251 263:251 264:251 265:251 266:177 289:15 290:157 291:248 292:251 293:251 294:251 295:165 319:214 320:251 321:251 322:251 323:212 324:78 325:24 347:109 348:251 349:251 350:251 351:253 352:251 353:170 354:10 375:5 376:57 377:162 378:251 379:253 380:251 381:251 382:18 405:106 406:239 407:255 408:253 409:253 410:213 434:105 435:253 436:251 437:251 438:230 439:72 463:253 464:251 465:251 466:251 467:221 468:67 491:72 492:251 493:251 494:251 495:251 496:96 519:36 520:199 521:251 522:251 523:251 524:155 525:15 548:45 549:204 550:251 551:251 552:251 553:157 577:161 578:249 579:251 580:251 581:248 582:147 606:233 607:251 608:251 609:251 610:173 634:233 635:251 636:251 637:251 638:173 662:53 663:131 664:251 665:251 666:173 +1 126:15 127:200 128:255 129:90 154:42 155:254 156:254 157:173 182:42 183:254 184:254 185:199 210:26 211:237 212:254 213:221 214:12 239:213 240:254 241:231 242:17 267:213 268:254 269:199 295:213 296:254 297:199 323:213 324:254 325:96 350:20 351:232 352:254 353:33 378:84 379:254 380:229 381:17 406:168 407:254 408:203 433:8 434:217 435:254 436:187 461:84 462:254 463:254 464:48 489:195 490:254 491:254 492:37 516:20 517:233 518:254 519:212 520:4 544:132 545:254 546:254 547:82 571:9 572:215 573:254 574:254 575:116 576:46 599:55 600:254 601:254 602:254 603:254 604:121 627:113 628:254 629:254 630:254 631:254 632:40 655:12 656:163 657:254 658:185 659:58 660:1 +0 182:32 183:57 184:57 185:57 186:57 187:57 188:57 189:57 208:67 209:185 210:229 211:252 212:252 213:252 214:253 215:252 216:252 217:252 218:185 219:66 234:13 235:188 236:246 237:252 238:253 239:252 240:252 241:252 242:241 243:139 244:177 245:252 246:253 247:246 248:187 249:13 261:26 262:255 263:253 264:244 265:175 266:101 274:126 275:244 276:253 277:153 288:82 289:243 290:253 291:214 292:81 303:169 304:252 305:252 315:19 316:215 317:252 318:206 319:56 331:169 332:252 333:252 343:157 344:252 345:252 346:13 359:169 360:252 361:151 370:41 371:253 372:253 373:128 386:92 387:253 388:206 389:13 398:166 399:252 400:196 401:9 414:216 415:252 416:142 426:253 427:252 428:168 441:89 442:253 443:208 444:13 454:253 455:252 456:68 468:38 469:225 470:253 471:96 482:254 483:253 484:56 495:45 496:229 497:253 498:151 510:253 511:252 512:81 522:70 523:225 524:252 525:227 538:216 539:252 540:168 548:29 549:134 550:253 551:252 552:186 553:31 566:91 567:252 568:243 569:125 573:51 574:114 575:113 576:210 577:252 578:253 579:151 580:19 595:157 596:253 597:253 598:254 599:253 600:253 601:253 602:254 603:253 604:244 605:175 606:51 623:19 624:122 625:196 626:197 627:221 628:196 629:196 630:197 631:121 632:56 655:25 +0 127:42 128:235 129:255 130:84 153:15 154:132 155:208 156:253 157:253 158:171 159:108 180:6 181:177 182:253 183:253 184:253 185:253 186:253 187:242 188:110 208:151 209:253 210:253 211:253 212:253 213:253 214:253 215:253 216:139 235:48 236:208 237:253 238:253 239:253 240:253 241:253 242:253 243:253 244:139 263:85 264:253 265:253 266:253 267:253 268:236 269:156 270:184 271:253 272:148 273:6 290:7 291:141 292:253 293:253 294:253 295:253 296:27 298:170 299:253 300:253 301:74 318:19 319:253 320:253 321:253 322:253 323:253 324:27 326:170 327:253 328:253 329:74 345:16 346:186 347:253 348:253 349:253 350:242 351:105 352:4 354:170 355:253 356:253 357:94 358:1 373:141 374:253 375:253 376:253 377:242 378:100 382:170 383:253 384:253 385:253 386:8 401:141 402:253 403:253 404:253 405:224 410:170 411:253 412:253 413:253 414:8 428:12 429:158 430:253 431:253 432:230 433:51 438:18 439:237 440:253 441:253 442:8 456:76 457:253 458:253 459:218 460:61 467:236 468:253 469:253 470:8 484:76 485:253 486:253 487:168 495:110 496:253 497:132 498:3 512:76 513:253 514:253 515:168 521:20 522:174 523:239 524:147 525:5 539:5 540:155 541:253 542:253 543:168 548:102 549:170 550:253 551:253 552:139 567:3 568:128 569:253 570:253 571:228 572:179 573:179 574:179 575:179 576:245 577:253 578:253 579:219 580:41 596:76 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:253 607:163 624:39 625:199 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:253 634:170 635:9 653:36 654:219 655:253 656:253 657:253 658:253 659:253 660:224 661:65 662:22 +1 156:202 157:253 158:69 184:253 185:252 186:121 212:253 213:252 214:69 240:253 241:252 242:69 267:106 268:253 269:231 270:37 295:179 296:255 297:196 322:17 323:234 324:253 325:92 350:93 351:252 352:253 353:92 378:93 379:252 380:253 381:92 406:93 407:252 408:232 409:8 434:208 435:253 436:116 462:207 463:252 464:116 490:207 491:252 492:32 517:57 518:244 519:252 545:122 546:252 547:252 573:185 574:253 575:253 601:184 602:252 603:252 629:101 630:252 631:252 657:13 658:173 659:252 660:43 686:9 687:232 688:116 +1 156:73 157:253 158:253 159:253 160:124 184:73 185:251 186:251 187:251 188:251 212:99 213:251 214:251 215:251 216:225 240:253 241:251 242:251 243:251 244:71 266:79 267:180 268:253 269:251 270:251 271:173 272:20 294:110 295:253 296:255 297:253 298:216 322:109 323:251 324:253 325:251 326:215 350:109 351:251 352:253 353:251 354:215 378:109 379:251 380:253 381:251 382:137 406:109 407:251 408:253 409:251 410:35 433:37 434:253 435:253 436:255 437:253 438:35 461:140 462:251 463:251 464:253 465:168 466:15 488:125 489:246 490:251 491:251 492:190 493:15 516:144 517:251 518:251 519:251 520:180 543:53 544:221 545:251 546:251 547:251 548:51 571:125 572:253 573:253 574:253 575:201 598:105 599:253 600:251 601:251 602:188 603:30 626:180 627:253 628:251 629:251 630:142 654:180 655:253 656:251 657:235 658:82 682:180 683:253 684:251 685:215 +1 124:111 125:255 126:48 152:162 153:253 154:237 155:63 180:206 181:253 182:253 183:183 208:87 209:217 210:253 211:205 237:90 238:253 239:238 240:60 265:37 266:225 267:253 268:89 294:206 295:253 296:159 322:206 323:253 324:226 350:206 351:253 352:226 378:206 379:253 380:226 406:206 407:253 408:226 434:206 435:253 436:226 462:206 463:253 464:226 490:206 491:253 492:226 518:206 519:253 520:237 521:45 546:206 547:253 548:253 549:109 574:173 575:253 576:253 577:109 602:69 603:253 604:253 605:109 630:64 631:248 632:253 633:109 659:112 660:253 661:109 +0 99:70 100:255 101:165 102:114 127:122 128:253 129:253 130:253 131:120 155:165 156:253 157:253 158:253 159:234 160:52 183:99 184:253 185:253 186:253 187:253 188:228 189:26 209:60 210:168 211:238 212:202 213:174 214:253 215:253 216:253 217:127 235:91 236:81 237:1 238:215 239:128 240:28 241:12 242:181 243:253 244:253 245:175 246:3 262:18 263:204 264:253 265:77 270:7 271:253 272:253 273:253 274:54 289:54 290:248 291:253 292:253 293:143 298:1 299:127 300:253 301:253 302:188 317:104 318:253 319:253 320:253 321:20 327:81 328:249 329:253 330:191 345:192 346:253 347:253 348:218 349:5 356:203 357:253 358:208 359:21 372:56 373:237 374:253 375:250 376:100 384:104 385:253 386:253 387:75 400:76 401:253 402:253 403:224 412:119 413:253 414:253 415:75 428:80 429:253 430:253 431:103 439:4 440:241 441:253 442:218 443:32 456:213 457:253 458:253 459:103 467:125 468:253 469:253 470:191 484:213 485:253 486:253 487:103 494:3 495:176 496:253 497:253 498:135 512:213 513:253 514:253 515:103 521:9 522:162 523:253 524:253 525:226 526:37 540:179 541:253 542:253 543:135 548:46 549:157 550:253 551:253 552:253 553:63 568:23 569:188 570:253 571:249 572:179 573:179 574:179 575:179 576:233 577:253 578:253 579:233 580:156 581:10 597:51 598:235 599:253 600:253 601:253 602:253 603:253 604:253 605:251 606:232 607:120 626:16 627:124 628:253 629:253 630:253 631:253 632:152 633:104 +1 124:29 125:197 126:255 127:84 152:85 153:251 154:253 155:83 180:86 181:253 182:254 183:253 208:85 209:251 210:253 211:251 236:86 237:253 238:254 239:253 240:169 264:85 265:251 266:253 267:251 268:168 292:86 293:253 294:254 295:253 296:169 320:28 321:196 322:253 323:251 324:168 349:169 350:254 351:253 352:169 377:168 378:253 379:251 380:168 405:169 406:254 407:253 408:169 433:168 434:253 435:251 436:168 462:254 463:253 464:254 465:139 490:253 491:251 492:253 493:251 518:254 519:253 520:254 521:253 522:57 546:253 547:251 548:253 549:251 550:168 574:198 575:253 576:254 577:253 578:114 602:85 603:251 604:253 605:251 630:85 631:253 632:254 633:253 658:28 659:83 660:196 661:83 +1 159:31 160:210 161:253 162:163 187:198 188:252 189:252 190:162 213:10 214:86 215:242 216:252 217:252 218:66 241:164 242:252 243:252 244:252 245:188 246:8 268:53 269:242 270:252 271:252 272:225 273:14 296:78 297:252 298:252 299:252 300:204 323:56 324:231 325:252 326:252 327:212 328:35 351:157 352:252 353:252 354:252 355:37 377:8 378:132 379:253 380:252 381:252 382:230 383:24 405:45 406:252 407:253 408:252 409:154 410:55 427:7 428:55 433:107 434:253 435:255 436:228 437:53 454:15 455:24 456:23 460:110 461:242 462:252 463:228 464:59 482:57 483:83 487:88 488:247 489:252 490:252 491:140 514:15 515:189 516:252 517:252 518:252 542:74 543:252 544:252 545:238 546:90 570:178 571:252 572:252 573:189 597:40 598:217 599:252 600:252 601:59 625:75 626:252 627:252 628:252 629:85 630:61 653:62 654:239 655:252 656:156 657:14 682:178 683:252 684:14 +1 131:159 132:255 133:122 158:167 159:228 160:253 161:121 185:64 186:236 187:251 188:205 189:110 212:48 213:158 214:251 215:251 216:178 217:39 240:190 241:251 242:251 243:251 267:96 268:253 269:253 270:253 271:153 295:194 296:251 297:251 298:211 299:74 322:80 323:174 324:251 325:251 326:140 327:47 349:16 350:181 351:253 352:251 353:219 354:23 377:64 378:251 379:253 380:251 381:204 382:19 405:223 406:253 407:255 408:233 409:48 431:20 432:174 433:244 434:251 435:253 436:109 437:31 459:96 460:189 461:251 462:251 463:126 464:31 486:24 487:106 488:251 489:235 490:188 491:100 514:96 515:251 516:251 517:228 518:59 542:255 543:253 544:253 545:213 546:36 569:100 570:253 571:251 572:251 573:85 574:23 596:32 597:127 598:253 599:235 600:126 601:15 624:104 625:251 626:253 627:240 628:79 652:83 653:193 654:253 655:220 +0 153:92 154:191 155:178 156:253 157:242 158:141 159:104 160:29 180:26 181:253 182:252 183:252 184:252 185:253 186:252 187:252 188:252 189:108 190:19 206:57 207:123 208:222 209:253 210:252 211:252 212:252 213:168 214:224 215:252 216:252 217:253 218:84 233:176 234:243 235:252 236:252 237:253 238:252 239:252 240:252 242:19 243:153 244:252 245:253 246:209 247:25 259:10 260:128 261:255 262:253 263:244 264:225 265:114 266:194 267:253 268:178 272:163 273:254 274:253 275:168 287:85 288:252 289:253 290:189 291:56 294:19 295:133 296:9 300:38 301:253 302:252 303:168 314:19 315:191 316:252 317:194 318:19 329:253 330:252 331:234 332:22 342:107 343:252 344:252 345:13 357:253 358:252 359:252 360:128 370:169 371:253 372:241 385:141 386:253 387:253 388:140 397:19 398:225 399:252 400:139 413:66 414:252 415:252 416:139 425:29 426:252 427:252 428:52 441:29 442:252 443:252 444:139 453:29 454:252 455:252 456:28 469:29 470:252 471:252 472:40 481:141 482:253 483:253 484:91 497:154 498:253 499:168 509:66 510:252 511:252 512:165 525:253 526:252 527:168 537:19 538:224 539:252 540:252 552:126 553:253 554:252 555:80 566:169 567:252 568:252 569:214 570:38 579:126 580:249 581:253 582:151 583:6 594:26 595:223 596:253 597:254 598:253 599:128 600:29 604:13 605:41 606:216 607:253 608:253 609:226 610:38 623:122 624:252 625:253 626:252 627:252 628:252 629:169 630:169 631:169 632:206 633:253 634:252 635:252 636:202 637:38 651:19 652:56 653:168 654:224 655:252 656:252 657:253 658:252 659:252 660:252 661:253 662:233 663:130 664:6 682:94 683:139 684:190 685:153 686:252 687:164 688:139 689:28 690:22 +1 128:53 129:250 130:255 131:25 156:167 157:253 158:253 159:25 182:3 183:123 184:247 185:253 186:253 187:25 210:9 211:253 212:253 213:253 214:253 215:25 238:9 239:253 240:253 241:253 242:253 243:25 266:9 267:253 268:253 269:253 270:180 271:13 294:9 295:253 296:253 297:253 298:104 322:9 323:253 324:253 325:253 326:104 350:15 351:253 352:253 353:253 354:104 378:184 379:253 380:253 381:228 382:68 406:184 407:253 408:253 409:182 433:103 434:251 435:253 436:253 437:12 461:106 462:253 463:253 464:253 465:8 488:24 489:238 490:253 491:253 492:253 493:8 516:27 517:253 518:253 519:253 520:253 521:8 544:27 545:253 546:253 547:253 548:253 549:8 572:27 573:253 574:253 575:253 576:177 577:4 600:160 601:253 602:253 603:253 604:87 628:202 629:253 630:253 631:219 632:54 656:81 657:253 658:247 659:51 +0 122:63 123:176 124:253 125:253 126:159 127:113 128:63 150:140 151:253 152:252 153:252 154:252 155:252 156:241 157:100 158:66 177:54 178:227 179:253 180:252 181:252 182:252 183:252 184:253 185:252 186:239 187:181 188:57 204:38 205:224 206:252 207:253 208:226 209:246 210:252 211:252 212:253 213:252 214:252 215:252 216:252 217:108 218:3 232:57 233:252 234:252 235:253 236:27 237:88 238:112 239:112 240:112 241:112 242:142 243:252 244:252 245:253 246:152 247:31 260:198 261:253 262:253 263:79 270:32 271:153 272:253 273:255 274:253 275:196 287:76 288:246 289:252 290:127 299:3 300:106 301:253 302:252 303:214 304:28 315:194 316:252 317:252 318:112 329:143 330:252 331:252 332:193 343:225 344:252 345:217 346:37 357:38 358:234 359:252 360:223 370:63 371:240 372:252 373:84 386:146 387:252 388:223 398:114 399:253 400:228 401:47 414:147 415:253 416:253 417:112 426:159 427:252 428:195 442:225 443:252 444:252 445:112 454:253 455:252 456:195 470:225 471:252 472:230 473:25 482:159 483:252 484:202 485:10 497:92 498:243 499:252 500:208 510:113 511:252 512:252 513:161 524:79 525:253 526:252 527:220 528:37 538:114 539:253 540:253 541:253 542:174 543:63 550:26 551:128 552:253 553:255 554:253 555:133 566:12 567:228 568:252 569:252 570:252 571:241 572:100 573:85 574:76 576:85 577:131 578:231 579:252 580:252 581:253 582:129 583:6 595:97 596:208 597:252 598:252 599:253 600:252 601:252 602:246 603:197 604:253 605:252 606:252 607:252 608:220 609:133 610:6 624:19 625:99 626:239 627:253 628:252 629:252 630:252 631:252 632:253 633:252 634:245 635:223 636:99 654:63 655:112 656:112 657:221 658:252 659:252 660:253 661:127 662:87 +0 153:12 154:136 155:254 156:255 157:195 158:115 159:3 180:6 181:175 182:253 183:196 184:160 185:252 186:253 187:15 208:130 209:253 210:234 211:4 213:27 214:205 215:232 216:40 235:54 236:246 237:253 238:68 242:24 243:243 244:106 262:3 263:134 264:235 265:99 266:4 271:132 272:247 273:77 290:56 291:253 292:62 299:23 300:233 301:129 318:179 319:183 320:4 328:182 329:220 345:21 346:232 347:59 356:95 357:232 358:21 373:128 374:183 385:228 386:85 401:187 402:124 413:228 414:186 429:187 430:124 441:228 442:104 457:187 458:124 469:169 470:184 485:187 486:124 497:203 498:150 513:187 514:124 524:10 525:220 526:39 541:187 542:155 552:111 553:201 569:129 570:228 571:7 579:12 580:181 581:76 598:234 599:166 600:9 606:24 607:209 608:106 626:139 627:250 628:167 629:11 630:2 631:11 632:11 633:129 634:227 635:90 636:11 655:95 656:247 657:253 658:178 659:253 660:253 661:244 662:86 684:47 685:175 686:253 687:232 688:149 689:40 +1 128:255 129:253 130:57 156:253 157:251 158:225 159:56 183:169 184:254 185:253 186:254 187:84 211:168 212:253 213:251 214:253 215:83 238:85 239:253 240:254 241:253 242:169 266:85 267:251 268:253 269:251 270:56 294:141 295:253 296:254 297:253 322:253 323:251 324:253 325:251 350:254 351:253 352:254 353:253 378:253 379:251 380:253 381:251 406:254 407:253 408:254 409:196 433:114 434:253 435:251 436:253 437:83 461:169 462:254 463:253 464:226 465:56 489:168 490:253 491:251 492:168 516:85 517:253 518:254 519:253 544:85 545:251 546:253 547:251 572:254 573:253 574:254 575:253 600:253 601:251 602:253 603:251 628:254 629:253 630:254 631:253 656:139 657:251 658:253 659:138 +0 151:23 152:167 153:208 154:254 155:255 156:129 157:19 179:151 180:253 181:253 182:253 183:253 184:253 185:209 186:26 207:181 208:253 209:253 210:253 211:227 212:181 213:253 214:207 215:22 235:227 236:253 237:253 238:253 239:92 240:38 241:226 242:253 243:129 244:2 263:193 264:253 265:253 266:248 267:62 269:50 270:253 271:253 272:45 291:170 292:253 293:253 294:135 297:12 298:208 299:253 300:119 318:16 319:232 320:253 321:253 322:21 326:60 327:253 328:185 346:164 347:253 348:253 349:224 350:14 354:14 355:217 356:247 357:62 373:3 374:193 375:253 376:250 377:64 383:199 384:253 385:179 401:67 402:253 403:253 404:205 411:98 412:253 413:188 429:151 430:253 431:245 432:43 439:63 440:250 441:188 457:151 458:253 459:243 468:244 469:222 470:22 485:151 486:253 487:217 496:244 497:253 498:115 512:3 513:195 514:253 515:134 524:156 525:253 526:150 541:140 542:253 543:134 552:239 553:253 554:139 569:44 570:253 571:134 579:53 580:246 581:237 582:32 597:8 598:200 599:229 600:40 606:25 607:225 608:253 609:188 626:120 627:250 628:230 629:58 630:17 632:12 633:42 634:213 635:253 636:238 637:84 655:151 656:253 657:253 658:217 659:179 660:206 661:253 662:253 663:196 664:118 683:18 684:58 685:145 686:152 687:253 688:214 689:145 690:74 691:7 +1 130:24 131:150 132:233 133:38 156:14 157:89 158:253 159:254 160:254 161:71 183:78 184:203 185:254 186:254 187:254 188:232 189:77 190:54 191:8 209:12 210:155 211:240 212:254 213:223 214:76 215:254 216:254 217:254 218:254 219:68 235:3 236:101 237:216 238:254 239:227 240:122 241:26 242:110 243:254 244:254 245:254 246:184 247:100 262:46 263:222 264:254 265:254 266:179 267:48 270:181 271:254 272:254 273:146 274:6 288:2 289:145 290:248 291:254 292:182 293:111 294:4 297:3 298:250 299:254 300:206 301:3 315:6 316:144 317:254 318:254 319:171 325:125 326:254 327:252 328:80 342:6 343:142 344:254 345:179 346:95 347:4 352:61 353:246 354:254 355:150 370:64 371:254 372:177 373:14 380:124 381:254 382:246 383:32 398:108 399:97 400:15 407:24 408:226 409:254 410:116 435:177 436:255 437:254 438:5 463:196 464:254 465:99 466:1 490:3 491:199 492:254 493:79 518:129 519:254 520:254 521:23 546:178 547:254 548:192 549:8 550:3 551:43 573:11 574:198 575:254 576:128 577:66 578:130 579:225 595:137 596:202 597:106 598:84 599:84 600:84 601:112 602:254 603:254 604:254 605:254 606:212 607:151 623:172 624:254 625:254 626:254 627:254 628:254 629:254 630:254 631:254 632:254 633:162 634:75 651:12 652:106 653:177 654:254 655:254 656:254 657:235 658:135 659:100 660:17 661:2 +0 125:120 126:253 127:253 128:63 151:38 152:131 153:246 154:252 155:252 156:203 157:15 179:222 180:252 181:252 182:252 183:252 184:166 185:38 205:4 206:107 207:253 208:252 209:252 210:252 211:252 212:253 213:224 214:137 215:26 233:107 234:252 235:253 236:252 237:220 238:128 239:252 240:253 241:252 242:252 243:239 244:140 261:170 262:253 263:255 264:168 267:79 268:192 269:253 270:253 271:253 272:253 273:255 274:90 288:51 289:243 290:252 291:215 292:33 296:12 297:74 298:233 299:252 300:252 301:253 302:195 303:19 316:166 317:252 318:252 319:31 326:43 327:149 328:195 329:253 330:252 331:177 332:19 343:57 344:234 345:252 346:252 357:237 358:252 359:252 360:180 361:13 371:85 372:252 373:252 374:173 385:50 386:237 387:252 388:252 389:112 399:226 400:253 401:240 402:63 414:163 415:253 416:253 417:112 426:38 427:234 428:252 429:176 442:85 443:252 444:252 445:158 454:113 455:252 456:252 457:84 470:19 471:209 472:252 473:252 482:207 483:252 484:252 485:84 498:10 499:203 500:252 501:236 510:253 511:252 512:252 513:84 526:85 527:252 528:252 529:112 538:114 539:253 540:253 541:146 553:51 554:159 555:253 556:240 557:63 566:75 567:243 568:252 569:249 570:146 579:57 580:85 581:238 582:252 583:252 584:99 595:116 596:252 597:252 598:252 599:198 600:197 601:165 602:57 603:57 604:57 605:182 606:197 607:234 608:252 609:253 610:233 611:164 612:19 623:28 624:84 625:180 626:252 627:253 628:252 629:252 630:252 631:252 632:253 633:252 634:252 635:252 636:252 637:225 638:71 653:13 654:112 655:253 656:252 657:252 658:252 659:252 660:253 661:252 662:252 663:157 664:112 +1 127:155 128:253 129:126 155:253 156:251 157:141 158:4 183:253 184:251 185:251 186:31 211:253 212:251 213:251 214:31 239:253 240:251 241:251 242:31 267:255 268:253 269:253 270:31 293:8 294:131 295:253 296:251 297:235 298:27 321:64 322:251 323:253 324:251 325:126 349:64 350:251 351:253 352:251 353:126 377:64 378:251 379:253 380:251 381:126 405:64 406:253 407:255 408:221 433:182 434:251 435:253 436:200 460:64 461:236 462:251 463:253 464:62 487:8 488:158 489:251 490:251 491:169 492:8 515:32 516:251 517:251 518:251 519:158 543:32 544:253 545:253 546:253 547:159 571:32 572:251 573:251 574:251 575:39 599:32 600:251 601:251 602:251 627:32 628:251 629:251 630:251 631:100 655:32 656:251 657:251 658:251 +0 101:88 102:127 103:5 126:19 127:58 128:20 129:14 130:217 131:19 152:7 153:146 154:247 155:253 156:235 157:27 158:84 159:81 180:126 181:253 182:164 183:19 184:15 187:156 188:9 208:214 209:222 210:34 215:234 216:58 235:59 236:254 237:116 243:235 244:58 263:141 264:251 265:72 271:151 272:140 291:224 292:233 299:136 300:223 319:254 320:218 327:136 328:253 347:254 348:135 355:136 356:253 374:23 375:255 376:114 383:137 384:231 402:98 403:254 404:122 411:136 412:155 430:98 431:254 432:106 439:166 440:155 458:98 459:254 460:128 467:234 468:193 486:98 487:254 488:135 494:61 495:248 496:118 515:255 516:238 517:18 521:13 522:224 523:254 524:58 543:201 544:253 545:128 546:2 548:5 549:150 550:253 551:167 552:9 571:18 572:226 573:253 574:49 575:31 576:156 577:253 578:228 579:13 600:147 601:253 602:243 603:241 604:254 605:227 606:43 628:5 629:126 630:245 631:253 632:231 633:46 +0 127:37 128:141 129:156 130:156 131:194 132:194 133:47 153:11 154:132 155:239 156:253 157:253 158:253 159:253 160:254 161:181 180:25 181:172 182:253 183:235 184:167 185:78 186:93 187:174 188:254 189:247 190:54 207:26 208:210 209:253 210:237 211:90 216:201 217:253 218:78 235:192 236:253 237:237 238:58 244:156 245:253 246:78 262:141 263:254 264:235 265:53 269:19 270:5 272:156 273:254 274:78 289:46 290:254 291:253 292:92 296:17 297:226 298:217 299:49 300:148 301:253 302:78 317:165 318:254 319:239 320:24 324:20 325:253 326:253 327:58 328:18 329:115 330:24 344:37 345:248 346:254 347:91 352:2 353:117 354:250 355:163 356:91 372:77 373:253 374:254 375:39 382:196 383:253 384:173 400:159 401:254 402:218 403:15 410:77 411:254 412:255 413:61 428:234 429:253 430:113 438:21 439:226 440:254 441:135 455:25 456:240 457:253 458:68 467:195 468:254 469:135 483:79 484:253 485:253 495:195 496:254 497:135 511:79 512:253 513:253 514:76 523:195 524:254 525:99 540:212 541:254 542:209 543:9 550:10 551:209 552:196 553:15 568:54 569:253 570:254 571:137 572:36 576:2 577:20 578:168 579:253 580:60 596:28 597:235 598:254 599:253 600:199 601:124 602:79 603:79 604:167 605:253 606:253 607:185 608:30 625:15 626:117 627:217 628:253 629:253 630:253 631:254 632:253 633:240 634:109 635:12 655:27 656:126 657:208 658:253 659:193 660:147 661:40 +0 154:32 155:134 156:218 157:254 158:254 159:254 160:217 161:84 176:44 177:208 178:215 179:156 180:35 181:119 182:236 183:246 184:136 185:91 186:69 187:151 188:249 189:246 190:78 203:44 204:230 205:254 206:254 207:254 208:254 209:254 210:196 211:48 216:60 217:224 218:210 219:24 231:118 232:254 233:202 234:19 235:201 236:254 237:181 238:9 245:35 246:233 247:168 259:193 260:223 261:34 263:59 264:163 265:236 266:15 274:140 275:205 276:8 286:60 287:254 288:176 293:38 302:54 303:237 304:80 314:59 315:254 316:93 331:131 332:200 342:59 343:240 344:24 359:79 360:214 370:59 371:234 387:67 388:248 389:54 398:59 399:234 416:235 417:58 426:60 427:235 443:79 444:255 445:59 454:59 455:251 456:66 471:79 472:250 473:54 482:59 483:254 484:108 499:146 500:214 510:5 511:203 512:187 513:3 526:4 527:188 528:199 539:118 540:254 541:57 554:96 555:254 556:117 567:16 568:237 569:224 570:14 581:14 582:187 583:206 584:8 596:88 597:252 598:186 599:16 608:16 609:187 610:252 611:125 625:100 626:254 627:237 628:94 629:24 635:13 636:214 637:254 638:166 653:3 654:57 655:215 656:248 657:241 658:235 659:197 660:137 661:137 662:137 663:231 664:238 665:155 666:25 684:57 685:155 686:246 687:254 688:254 689:254 690:254 691:147 692:36 +1 124:102 125:252 126:252 127:41 152:102 153:250 154:250 155:202 180:102 181:250 182:250 183:232 184:91 208:102 209:250 210:250 211:212 212:29 236:102 237:252 238:252 239:254 240:150 264:102 265:250 266:250 267:252 268:149 292:102 293:250 294:250 295:252 296:149 320:102 321:250 322:250 323:252 324:231 325:80 349:152 350:252 351:254 352:252 353:100 377:151 378:250 379:252 380:250 381:100 405:151 406:250 407:252 408:250 409:100 433:151 434:250 435:252 436:250 437:100 461:123 462:243 463:254 464:252 465:100 490:202 491:252 492:250 493:100 518:80 519:252 520:250 521:190 522:30 547:252 548:250 549:250 550:49 575:255 576:252 577:252 578:252 579:214 580:31 603:171 604:250 605:250 606:250 607:252 608:190 609:40 631:20 632:160 633:250 634:250 635:252 636:250 637:100 660:20 661:170 662:250 663:212 664:49 665:20 +0 124:20 125:121 126:197 127:253 128:64 151:23 152:200 153:252 154:252 155:252 156:184 157:6 178:25 179:197 180:252 181:252 182:252 183:252 184:253 185:228 186:107 187:15 205:26 206:196 207:252 208:252 209:252 210:252 211:252 212:253 213:252 214:252 215:219 216:178 217:21 233:186 234:252 235:238 236:94 237:67 238:224 239:217 240:53 241:109 242:245 243:252 244:252 245:213 246:63 260:98 261:242 262:252 263:101 266:39 267:31 270:109 271:128 272:241 273:252 274:207 275:97 287:17 288:230 289:252 290:241 291:56 300:109 301:252 302:252 303:229 304:17 314:13 315:192 316:252 317:243 318:96 328:25 329:127 330:252 331:252 332:120 342:121 343:252 344:252 345:165 357:125 358:252 359:252 360:246 361:70 370:190 371:252 372:252 373:39 385:26 386:210 387:252 388:252 389:119 398:255 399:253 400:159 414:22 415:209 416:253 417:183 426:253 427:252 428:103 443:34 444:252 445:252 454:253 455:252 456:26 471:27 472:252 473:252 482:253 483:252 484:168 485:13 499:70 500:252 501:209 510:147 511:252 512:252 513:75 526:68 527:233 528:252 529:119 538:121 539:252 540:252 541:189 542:40 552:15 553:82 554:231 555:252 556:214 557:31 566:38 567:135 568:248 569:252 570:231 571:145 572:41 573:41 574:41 575:41 576:20 577:24 578:37 579:83 580:194 581:252 582:252 583:212 584:33 596:83 597:213 598:252 599:252 600:252 601:252 602:252 603:252 604:204 605:213 606:243 607:252 608:252 609:252 610:212 611:34 625:34 626:140 627:238 628:248 629:252 630:252 631:252 632:253 633:252 634:252 635:241 636:238 637:238 638:75 656:82 657:119 658:119 659:119 660:120 661:119 662:119 663:19 +1 127:20 128:254 129:255 130:37 155:19 156:253 157:253 158:134 183:19 184:253 185:253 186:246 187:125 211:76 212:253 213:253 214:253 215:158 239:207 240:253 241:253 242:253 243:158 267:207 268:253 269:253 270:253 271:158 294:48 295:223 296:253 297:253 298:243 299:106 322:141 323:253 324:253 325:253 326:113 349:65 350:237 351:253 352:253 353:253 354:36 377:76 378:253 379:253 380:253 381:253 382:36 405:76 406:253 407:253 408:253 409:253 410:36 433:76 434:253 435:253 436:253 437:118 438:4 460:4 461:148 462:253 463:253 464:253 465:103 488:10 489:253 490:253 491:253 492:253 493:103 516:10 517:253 518:253 519:253 520:173 521:7 544:10 545:253 546:253 547:253 548:168 572:143 573:253 574:253 575:239 576:49 600:198 601:253 602:253 603:234 615:140 628:198 629:253 630:253 631:234 656:198 657:253 658:253 659:234 +0 235:40 236:37 238:7 239:77 240:137 241:136 242:136 243:136 244:136 245:40 246:6 261:16 262:135 263:254 264:233 266:152 267:215 268:96 269:140 270:155 271:118 272:230 273:254 274:158 275:68 288:19 289:164 290:254 291:114 294:235 295:140 301:99 302:230 303:254 304:186 305:14 315:70 316:226 317:242 318:121 322:104 323:195 324:38 330:33 331:179 332:253 333:140 342:41 343:241 344:198 345:43 359:24 360:209 361:223 370:164 371:250 372:66 388:136 389:253 398:254 399:158 416:136 417:215 426:255 427:76 442:5 443:127 444:246 445:133 454:254 455:122 469:5 470:150 471:247 472:91 473:9 482:254 483:165 495:13 496:79 497:194 498:216 499:84 510:111 511:251 512:87 519:16 520:25 521:40 522:107 523:186 524:213 525:117 526:25 538:14 539:185 540:235 541:142 542:23 546:91 547:157 548:231 549:207 550:126 551:49 569:143 570:195 571:255 572:254 573:254 574:244 575:157 576:76 599:39 600:39 601:39 602:33 +1 128:166 129:255 130:187 131:6 156:165 157:253 158:253 159:13 183:15 184:191 185:253 186:253 187:13 211:49 212:253 213:253 214:253 215:13 239:141 240:253 241:253 242:169 243:4 266:4 267:189 268:253 269:249 270:53 294:69 295:253 296:253 297:246 322:69 323:253 324:253 325:246 350:118 351:253 352:253 353:124 378:206 379:253 380:231 381:21 405:66 406:241 407:253 408:199 433:105 434:253 435:253 436:89 460:3 461:228 462:253 463:252 464:86 488:111 489:253 490:253 491:205 516:166 517:253 518:253 519:75 543:43 544:249 545:253 546:193 547:9 570:4 571:160 572:253 573:253 574:184 598:37 599:253 600:253 601:253 602:88 626:140 627:253 628:253 629:186 630:18 654:14 655:253 656:253 657:27 +1 128:117 129:128 155:2 156:199 157:127 183:81 184:254 185:87 211:116 212:254 213:48 239:175 240:234 241:18 266:5 267:230 268:168 294:80 295:255 296:142 322:80 323:255 324:142 350:80 351:251 352:57 378:129 379:239 406:164 407:209 433:28 434:245 435:159 461:64 462:254 463:144 489:84 490:254 491:80 517:143 518:254 519:30 544:3 545:225 546:200 572:48 573:254 574:174 600:48 601:254 602:174 628:93 629:254 630:129 656:53 657:234 658:41 +1 129:159 130:142 156:11 157:220 158:141 184:78 185:254 186:141 212:111 213:254 214:109 240:196 241:221 242:15 267:26 268:221 269:159 295:63 296:254 297:159 323:178 324:254 325:93 350:7 351:191 352:254 353:97 378:42 379:255 380:254 381:41 406:42 407:254 408:195 409:10 434:141 435:255 436:78 461:11 462:202 463:254 464:59 489:86 490:254 491:254 492:59 517:142 518:254 519:248 520:52 545:142 546:254 547:195 573:142 574:254 575:164 601:142 602:254 603:77 629:142 630:254 631:131 657:77 658:172 659:5 +0 124:66 125:254 126:254 127:58 128:60 129:59 130:59 131:50 151:73 152:233 153:253 154:253 155:148 156:254 157:253 158:253 159:232 160:73 179:156 180:253 181:253 182:253 183:117 184:255 185:253 186:253 187:253 188:223 189:176 190:162 205:37 206:116 207:246 208:253 209:180 210:18 211:4 212:18 213:109 214:241 215:253 216:253 217:253 218:236 219:28 233:235 234:253 235:253 236:245 237:107 242:109 243:170 244:253 245:253 246:253 247:174 261:235 262:253 263:253 264:233 271:15 272:156 273:253 274:253 275:223 276:72 287:10 288:156 289:250 290:253 291:253 292:67 300:99 301:253 302:253 303:253 304:127 305:5 315:118 316:253 317:253 318:253 319:204 320:26 328:68 329:223 330:253 331:253 332:253 333:57 342:32 343:191 344:253 345:253 346:253 347:97 357:156 358:253 359:253 360:253 361:57 370:59 371:253 372:253 373:253 374:253 375:97 385:36 386:224 387:253 388:253 389:57 398:60 399:254 400:255 401:254 402:156 413:37 414:226 415:254 416:254 417:58 426:59 427:253 428:253 429:253 430:154 441:156 442:253 443:253 444:253 445:57 454:59 455:253 456:253 457:253 458:154 469:156 470:253 471:253 472:253 473:57 482:59 483:253 484:253 485:253 486:246 487:90 496:16 497:171 498:253 499:253 500:231 501:49 510:59 511:253 512:253 513:253 514:253 515:156 516:91 524:99 525:253 526:253 527:222 528:71 538:59 539:253 540:253 541:253 542:253 543:253 544:245 545:109 551:145 552:194 553:253 554:253 555:174 566:9 567:38 568:174 569:251 570:253 571:253 572:253 573:241 574:215 575:215 576:217 577:215 578:215 579:250 580:253 581:253 582:221 583:26 597:235 598:253 599:253 600:253 601:253 602:253 603:253 604:254 605:253 606:253 607:253 608:253 609:204 610:26 625:108 626:116 627:200 628:253 629:253 630:253 631:253 632:254 633:253 634:253 635:253 636:199 637:44 655:36 656:57 657:118 658:253 659:253 660:58 661:57 662:57 663:57 664:35 +1 129:101 130:222 131:84 157:225 158:252 159:84 184:89 185:246 186:208 187:19 212:128 213:252 214:195 239:79 240:253 241:252 242:195 267:141 268:255 269:253 270:133 294:26 295:240 296:253 297:252 298:55 322:60 323:252 324:253 325:154 326:12 349:7 350:178 351:252 352:253 353:27 377:57 378:252 379:252 380:253 381:27 405:57 406:253 407:253 408:204 409:15 433:104 434:252 435:252 436:94 460:19 461:209 462:252 463:252 488:101 489:252 490:252 491:157 516:225 517:252 518:252 519:112 544:226 545:253 546:240 547:63 572:225 573:252 574:223 600:225 601:252 602:223 628:225 629:252 630:242 631:75 656:146 657:252 658:236 659:50 +0 124:41 125:254 126:254 127:157 128:34 129:34 130:218 131:255 132:206 133:34 134:18 151:53 152:238 153:252 154:252 155:252 156:252 157:252 158:252 159:252 160:252 161:252 162:162 163:26 178:66 179:220 180:252 181:252 182:252 183:209 184:153 185:223 186:252 187:252 188:252 189:252 190:252 191:98 206:166 207:252 208:252 209:252 210:252 211:141 213:85 214:230 215:252 216:252 217:252 218:252 219:98 234:166 235:252 236:252 237:252 238:252 239:141 242:73 243:102 244:252 245:252 246:252 247:98 262:166 263:252 264:252 265:252 266:191 267:30 271:5 272:97 273:252 274:252 275:220 276:51 289:123 290:245 291:252 292:252 293:202 294:14 300:56 301:252 302:252 303:252 304:65 316:18 317:154 318:252 319:252 320:241 328:56 329:252 330:252 331:252 332:65 343:21 344:146 345:252 346:252 347:252 348:241 356:56 357:252 358:252 359:252 360:65 371:67 372:252 373:252 374:252 375:252 376:241 384:56 385:252 386:252 387:252 388:65 399:67 400:252 401:252 402:252 403:252 404:116 412:56 413:252 414:252 415:252 416:65 427:67 428:252 429:252 430:252 431:252 432:20 440:56 441:252 442:252 443:252 444:65 455:67 456:252 457:252 458:252 459:87 460:4 468:56 469:252 470:252 471:124 472:11 483:67 484:252 485:252 486:252 487:54 494:19 495:236 496:245 497:252 498:252 499:98 511:67 512:252 513:252 514:252 515:97 516:5 521:39 522:219 523:252 524:252 525:252 526:252 527:98 539:67 540:252 541:252 542:252 543:252 544:102 545:89 546:89 547:89 548:89 549:203 550:252 551:252 552:252 553:252 554:209 555:64 567:67 568:252 569:252 570:252 571:252 572:252 573:252 574:252 575:252 576:252 577:252 578:252 579:252 580:226 581:130 582:68 595:67 596:252 597:252 598:252 599:252 600:252 601:252 602:252 603:252 604:252 605:252 606:252 607:239 608:77 623:17 624:65 625:163 626:252 627:252 628:252 629:252 630:252 631:252 632:252 633:252 634:96 635:59 653:17 654:176 655:252 656:252 657:252 658:252 659:155 660:32 661:32 662:6 +0 96:56 97:247 98:121 124:24 125:242 126:245 127:122 153:231 154:253 155:253 156:104 157:12 181:90 182:253 183:253 184:254 185:221 186:120 187:120 188:85 206:67 207:75 208:36 209:11 210:56 211:222 212:254 213:253 214:253 215:253 216:245 217:207 218:36 233:86 234:245 235:249 236:105 239:44 240:224 241:230 242:253 243:253 244:253 245:253 246:214 247:10 260:8 261:191 262:253 263:143 269:29 270:119 271:119 272:158 273:253 274:253 275:94 288:15 289:253 290:226 291:48 300:4 301:183 302:253 303:248 304:56 316:42 317:253 318:178 329:179 330:253 331:184 332:14 344:164 345:253 346:178 357:179 358:253 359:163 371:61 372:254 373:254 374:179 384:76 385:254 386:254 387:164 399:60 400:253 401:253 402:178 411:29 412:206 413:253 414:253 415:40 427:60 428:253 429:253 430:178 439:120 440:253 441:253 442:245 443:13 455:60 456:253 457:253 458:178 467:120 468:253 469:239 470:63 483:60 484:253 485:253 486:178 494:14 495:238 496:253 497:179 511:18 512:190 513:253 514:231 515:70 521:43 522:184 523:253 524:253 525:74 540:86 541:253 542:253 543:239 544:134 545:8 548:56 549:163 550:253 551:253 552:213 553:35 568:16 569:253 570:253 571:253 572:253 573:240 574:239 575:239 576:247 577:253 578:253 579:210 580:27 596:4 597:59 598:204 599:253 600:253 601:253 602:253 603:253 604:254 605:253 606:250 607:110 626:31 627:122 628:253 629:253 630:253 631:253 632:255 633:217 634:98 +0 125:19 126:164 127:253 128:255 129:253 130:118 131:59 132:36 153:78 154:251 155:251 156:253 157:251 158:251 159:251 160:199 161:45 180:14 181:198 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:204 190:26 208:5 209:117 210:251 211:251 212:243 213:212 214:239 215:251 216:251 217:251 218:218 236:95 237:251 238:251 239:251 240:120 242:175 243:251 244:251 245:251 246:231 263:97 264:237 265:251 266:251 267:251 270:67 271:240 272:251 273:251 274:243 275:108 290:8 291:163 292:251 293:251 294:240 295:81 299:68 300:251 301:251 302:251 303:179 304:9 317:13 318:145 319:251 320:251 321:226 322:80 327:39 328:251 329:251 330:251 331:251 332:115 345:144 346:251 347:251 348:251 349:173 355:18 356:167 357:251 358:251 359:251 360:115 373:233 374:251 375:251 376:251 377:173 384:98 385:251 386:251 387:251 388:115 400:176 401:253 402:253 403:216 404:179 412:99 413:253 414:253 415:253 416:116 427:55 428:210 429:251 430:251 431:96 440:98 441:251 442:251 443:214 444:62 455:117 456:251 457:251 458:251 459:96 467:28 468:204 469:251 470:237 471:53 482:55 483:241 484:251 485:251 486:160 487:7 494:28 495:222 496:251 497:251 498:231 510:59 511:251 512:251 513:251 514:153 520:23 521:98 522:204 523:251 524:251 525:251 526:156 538:59 539:251 540:251 541:251 542:153 546:85 547:155 548:179 549:251 550:251 551:251 552:251 553:154 554:15 566:59 567:251 568:251 569:251 570:236 571:214 572:214 573:214 574:234 575:251 576:253 577:251 578:251 579:248 580:156 581:15 594:41 595:209 596:251 597:251 598:251 599:251 600:251 601:251 602:251 603:251 604:253 605:251 606:196 607:146 623:54 624:115 625:241 626:251 627:251 628:251 629:251 630:251 631:251 632:253 633:187 634:35 653:83 654:251 655:251 656:251 657:251 658:251 659:101 660:57 661:31 +1 129:232 130:255 131:107 156:58 157:244 158:253 159:106 184:95 185:253 186:253 187:106 212:95 213:253 214:253 215:106 240:95 241:253 242:249 243:69 268:144 269:253 270:192 295:97 296:233 297:253 298:66 323:195 324:253 325:253 326:5 350:38 351:232 352:253 353:182 354:2 377:10 378:160 379:253 380:231 381:53 405:42 406:253 407:253 408:158 433:141 434:253 435:253 436:115 460:75 461:245 462:253 463:183 464:4 487:1 488:147 489:253 490:251 491:58 515:20 516:253 517:253 518:180 543:202 544:253 545:226 546:27 571:243 572:253 573:212 598:85 599:251 600:253 601:173 626:209 627:253 628:244 629:57 654:169 655:253 656:174 +1 127:63 128:128 129:2 155:63 156:254 157:123 183:63 184:254 185:179 211:63 212:254 213:179 239:63 240:254 241:179 267:142 268:254 269:179 295:187 296:254 297:158 323:187 324:254 325:55 350:68 351:235 352:254 353:55 378:181 379:254 380:254 381:55 406:181 407:254 408:202 409:14 434:181 435:254 436:186 462:181 463:254 464:146 490:181 491:254 492:62 518:181 519:254 520:62 546:181 547:254 548:62 574:181 575:255 576:62 602:181 603:254 604:241 605:52 630:181 631:254 632:222 633:30 658:181 659:224 660:34 +1 130:131 131:255 132:184 133:15 157:99 158:247 159:253 160:182 161:15 185:124 186:253 187:253 188:253 189:38 212:9 213:171 214:253 215:253 216:140 217:1 240:47 241:253 242:253 243:251 244:117 267:43 268:219 269:253 270:253 271:153 295:78 296:253 297:253 298:253 299:84 323:97 324:253 325:253 326:244 327:74 350:69 351:243 352:253 353:253 354:183 377:10 378:168 379:253 380:253 381:215 382:34 405:31 406:253 407:253 408:253 409:129 433:107 434:253 435:253 436:242 437:67 460:24 461:204 462:253 463:253 464:187 488:95 489:253 490:253 491:201 492:25 516:239 517:253 518:253 519:176 543:119 544:251 545:253 546:253 547:138 570:30 571:212 572:253 573:252 574:165 575:8 598:193 599:253 600:253 601:222 626:193 627:253 628:253 629:189 654:193 655:253 656:201 657:27 +0 125:57 126:255 127:253 128:198 129:85 153:168 154:253 155:251 156:253 157:251 158:169 159:56 180:86 181:253 182:254 183:253 184:254 185:253 186:254 187:253 188:57 208:197 209:251 210:253 211:251 212:253 213:251 214:253 215:251 216:225 217:56 235:169 236:255 237:253 238:226 239:56 241:114 242:254 243:253 244:254 245:84 262:57 263:224 264:253 265:251 266:56 270:139 271:251 272:253 273:83 290:141 291:253 292:255 293:84 298:57 299:225 300:254 301:196 318:253 319:251 320:253 321:83 327:168 328:253 329:83 345:169 346:254 347:253 348:169 355:169 356:254 357:253 358:169 373:168 374:253 375:251 376:56 383:168 384:253 385:251 386:56 401:169 402:254 403:84 412:254 413:253 429:168 430:253 431:83 440:253 441:251 456:29 457:197 458:254 459:84 467:169 468:254 469:196 484:85 485:251 486:253 487:83 494:57 495:224 496:253 497:83 512:57 513:225 514:254 515:139 521:57 522:141 523:253 524:254 525:84 541:168 542:253 543:251 544:169 545:56 547:114 548:169 549:224 550:253 551:251 552:253 553:83 569:169 570:254 571:253 572:254 573:253 574:254 575:253 576:254 577:253 578:254 579:253 580:226 581:56 597:56 598:253 599:251 600:253 601:251 602:253 603:251 604:253 605:251 606:253 607:251 608:56 626:169 627:225 628:254 629:253 630:254 631:253 632:254 633:253 634:226 635:56 655:56 656:253 657:251 658:253 659:251 660:84 661:83 662:56 +0 127:12 128:105 129:224 130:255 131:247 132:22 155:131 156:254 157:254 158:243 159:252 160:76 182:131 183:225 184:254 185:224 186:48 187:136 208:13 209:109 210:252 211:254 212:254 213:254 214:197 215:76 235:9 236:181 237:254 238:254 239:240 240:229 241:237 242:254 243:252 244:152 245:21 262:9 263:143 264:254 265:254 266:226 267:36 269:22 270:138 271:254 272:254 273:188 289:13 290:181 291:254 292:254 293:250 294:64 298:2 299:53 300:236 301:252 302:131 317:102 318:254 319:254 320:254 321:111 328:56 329:243 330:251 331:42 344:30 345:186 346:254 347:254 348:206 349:29 357:199 358:254 359:91 372:92 373:254 374:254 375:237 376:13 385:134 386:254 387:91 400:133 401:254 402:254 403:126 413:134 414:250 415:17 428:187 429:254 430:237 431:23 441:200 442:183 456:187 457:254 458:213 467:2 468:134 469:252 470:101 484:183 485:254 486:133 495:14 496:254 497:234 498:34 512:92 513:254 514:161 522:84 523:204 524:254 525:56 540:92 541:254 542:229 549:85 550:252 551:252 552:188 553:11 568:56 569:252 570:229 575:3 576:53 577:235 578:253 579:166 597:224 598:245 599:130 600:68 601:68 602:134 603:214 604:254 605:254 606:159 625:141 626:254 627:254 628:254 629:254 630:254 631:254 632:233 633:95 634:3 653:14 654:152 655:254 656:254 657:254 658:186 659:157 660:53 +1 130:226 131:247 132:55 157:99 158:248 159:254 160:230 161:30 185:125 186:254 187:254 188:254 189:38 213:125 214:254 215:254 216:212 217:24 240:18 241:223 242:254 243:252 244:118 268:24 269:254 270:254 271:239 295:27 296:195 297:254 298:254 299:93 323:78 324:254 325:254 326:246 327:74 351:158 352:254 353:254 354:185 378:41 379:239 380:254 381:254 382:43 405:22 406:218 407:254 408:254 409:167 410:9 433:32 434:254 435:254 436:254 437:130 460:24 461:187 462:254 463:254 464:234 465:16 488:189 489:254 490:254 491:254 492:128 515:64 516:247 517:254 518:255 519:219 520:42 543:139 544:254 545:254 546:222 547:40 570:30 571:213 572:254 573:235 574:45 598:194 599:254 600:254 601:223 626:194 627:254 628:254 629:190 654:194 655:254 656:202 657:27 +1 130:166 131:253 132:124 133:53 158:140 159:251 160:251 161:180 185:125 186:246 187:251 188:251 189:51 212:32 213:190 214:251 215:251 216:251 217:103 240:21 241:174 242:251 243:251 244:251 268:73 269:176 270:253 271:253 272:201 296:149 297:251 298:251 299:251 300:71 323:27 324:228 325:251 326:251 327:157 328:10 351:180 352:253 353:251 354:251 355:142 377:27 378:180 379:231 380:253 381:251 382:96 383:41 405:89 406:253 407:253 408:255 409:211 410:25 433:217 434:251 435:251 436:253 437:107 460:21 461:221 462:251 463:251 464:242 465:92 487:32 488:190 489:251 490:251 491:251 492:103 515:202 516:251 517:251 518:251 519:122 542:53 543:255 544:253 545:253 546:221 547:51 570:180 571:253 572:251 573:251 574:142 598:180 599:253 600:251 601:251 602:142 626:180 627:253 628:251 629:157 630:82 654:180 655:253 656:147 657:10 +1 129:17 130:206 131:229 132:44 157:2 158:125 159:254 160:123 185:95 186:254 187:254 188:123 212:78 213:240 214:254 215:254 216:123 240:100 241:254 242:254 243:254 244:123 267:2 268:129 269:254 270:254 271:220 272:20 295:9 296:254 297:254 298:254 299:123 322:22 323:179 324:254 325:254 326:254 327:49 350:83 351:254 352:254 353:254 354:183 355:19 378:136 379:254 380:254 381:254 382:139 404:3 405:111 406:252 407:254 408:254 409:232 410:45 432:67 433:254 434:254 435:254 436:216 437:40 459:14 460:192 461:254 462:254 463:254 464:140 486:23 487:192 488:254 489:254 490:254 491:246 514:77 515:254 516:254 517:255 518:241 519:100 541:65 542:235 543:254 544:254 545:254 546:172 568:30 569:238 570:254 571:254 572:254 573:219 574:26 596:34 597:254 598:254 599:254 600:216 601:41 624:34 625:254 626:254 627:254 628:188 652:12 653:170 654:254 655:254 656:82 +1 130:218 131:253 132:124 157:84 158:236 159:251 160:251 184:63 185:236 186:251 187:251 188:122 212:73 213:251 214:251 215:251 216:173 240:202 241:251 242:251 243:251 244:71 267:53 268:255 269:253 270:253 271:253 272:72 295:180 296:253 297:251 298:251 299:188 300:30 323:180 324:253 325:251 326:251 327:142 350:47 351:211 352:253 353:251 354:235 355:82 377:27 378:211 379:251 380:253 381:251 382:215 405:89 406:253 407:253 408:255 409:253 410:164 433:217 434:251 435:251 436:253 437:168 438:15 460:21 461:221 462:251 463:251 464:253 465:107 487:32 488:190 489:251 490:251 491:251 492:221 493:61 515:73 516:251 517:251 518:251 519:251 520:180 543:255 544:253 545:253 546:253 547:201 570:105 571:253 572:251 573:251 574:251 575:71 598:180 599:253 600:251 601:246 602:137 603:10 626:180 627:253 628:251 629:215 654:180 655:253 656:251 657:86 +1 124:102 125:180 126:1 152:140 153:254 154:130 180:140 181:254 182:204 208:140 209:254 210:204 236:72 237:254 238:204 264:25 265:231 266:250 267:135 292:11 293:211 294:254 295:222 321:101 322:254 323:250 324:15 349:96 350:254 351:254 352:95 377:2 378:251 379:254 380:95 405:2 406:251 407:254 408:95 433:96 434:254 435:254 436:95 461:53 462:253 463:254 464:139 490:250 491:254 492:235 493:27 518:201 519:254 520:254 521:128 546:80 547:254 548:254 549:139 574:65 575:254 576:254 577:139 602:150 603:254 604:254 605:139 630:229 631:254 632:254 633:43 658:52 659:196 660:168 661:9 +0 128:87 129:208 130:249 155:27 156:212 157:254 158:195 182:118 183:225 184:254 185:254 186:232 187:147 188:46 209:115 210:248 211:254 212:254 213:254 214:254 215:254 216:230 217:148 218:12 236:18 237:250 238:254 239:245 240:226 241:254 242:254 243:254 244:254 245:254 246:148 263:92 264:205 265:254 266:250 267:101 268:20 269:194 270:254 271:254 272:254 273:254 274:229 275:53 291:152 292:254 293:254 294:94 297:14 298:124 299:187 300:254 301:254 302:254 303:213 318:95 319:252 320:254 321:206 322:15 327:3 328:6 329:51 330:231 331:254 332:94 345:50 346:246 347:254 348:254 349:20 358:200 359:254 360:96 372:21 373:184 374:254 375:254 376:147 377:2 386:200 387:254 388:96 400:177 401:254 402:254 403:218 404:33 413:16 414:211 415:254 416:96 427:11 428:219 429:254 430:251 431:92 441:84 442:254 443:232 444:44 455:101 456:254 457:254 458:141 469:162 470:254 471:231 472:42 483:235 484:254 485:227 486:42 496:51 497:238 498:254 499:213 511:235 512:254 513:199 524:160 525:254 526:229 527:52 539:235 540:254 541:199 549:10 550:84 551:150 552:253 553:254 554:147 567:235 568:254 569:213 570:20 575:17 576:63 577:158 578:254 579:254 580:254 581:155 582:12 595:122 596:248 597:254 598:204 599:98 600:42 601:177 602:180 603:200 604:254 605:254 606:253 607:213 608:82 609:10 624:203 625:254 626:254 627:254 628:254 629:254 630:254 631:254 632:251 633:219 634:94 652:35 653:221 654:254 655:254 656:254 657:254 658:254 659:217 660:95 +1 126:134 127:230 154:133 155:231 156:10 182:133 183:253 184:96 210:133 211:253 212:96 238:133 239:253 240:183 266:133 267:253 268:217 294:133 295:253 296:217 322:133 323:253 324:217 350:133 351:253 352:217 378:133 379:253 380:217 406:134 407:254 408:218 434:133 435:253 436:159 462:133 463:253 464:199 490:156 491:253 492:96 518:254 519:247 520:73 546:254 547:248 548:74 573:99 574:254 575:245 576:64 600:89 601:230 602:254 603:125 627:140 628:251 629:253 630:243 631:10 655:114 656:242 657:195 658:69 +1 125:29 126:85 127:255 128:139 153:197 154:251 155:253 156:251 181:254 182:253 183:254 184:253 209:253 210:251 211:253 212:251 237:254 238:253 239:254 240:253 265:253 266:251 267:253 268:138 293:254 294:253 295:254 296:196 321:253 322:251 323:253 324:196 349:254 350:253 351:254 352:84 377:253 378:251 379:253 380:196 405:254 406:253 407:254 408:253 433:253 434:251 435:253 436:251 461:254 462:253 463:254 464:253 489:253 490:251 491:253 492:251 517:254 518:253 519:254 520:253 545:253 546:251 547:253 548:251 573:254 574:253 575:254 576:253 601:253 602:251 603:253 604:251 629:57 630:225 631:254 632:253 658:56 659:253 660:251 +1 125:149 126:255 127:254 128:58 153:215 154:253 155:183 156:2 180:41 181:232 182:253 183:181 208:92 209:253 210:253 211:181 236:92 237:253 238:253 239:181 264:92 265:253 266:253 267:181 292:92 293:253 294:253 295:181 320:92 321:253 322:253 323:181 348:92 349:253 350:253 351:181 376:92 377:253 378:253 379:181 404:92 405:253 406:253 407:181 432:92 433:253 434:253 435:181 460:92 461:253 462:253 463:181 488:31 489:228 490:253 491:181 517:198 518:253 519:228 520:54 545:33 546:226 547:253 548:195 549:7 574:199 575:253 576:253 577:75 602:34 603:218 604:253 605:228 606:117 607:14 608:12 631:33 632:219 633:253 634:253 635:253 636:211 660:32 661:123 662:149 663:230 664:41 +1 130:79 131:203 132:141 157:51 158:240 159:240 160:140 185:88 186:252 187:252 188:140 213:197 214:252 215:252 216:140 241:197 242:252 243:252 244:140 268:147 269:253 270:253 271:253 295:38 296:234 297:252 298:242 299:89 323:113 324:252 325:252 326:223 350:16 351:207 352:252 353:252 354:129 377:16 378:203 379:253 380:252 381:220 382:37 405:29 406:253 407:255 408:253 409:56 432:19 433:181 434:252 435:253 436:176 437:6 460:166 461:252 462:252 463:228 464:52 487:10 488:203 489:252 490:252 491:126 514:63 515:178 516:252 517:252 518:173 542:114 543:253 544:253 545:225 570:238 571:252 572:252 573:99 596:7 597:135 598:253 599:252 600:176 601:19 624:29 625:252 626:253 627:252 628:55 652:13 653:189 654:253 655:204 656:25 +1 126:94 127:254 128:75 154:166 155:253 156:231 182:208 183:253 184:147 210:208 211:253 212:116 238:208 239:253 240:168 266:146 267:254 268:222 294:166 295:253 296:116 322:208 323:253 324:116 350:166 351:253 352:158 378:145 379:253 380:231 406:209 407:254 408:169 434:187 435:253 436:168 462:93 463:253 464:116 490:93 491:253 492:116 518:93 519:253 520:116 546:94 547:254 548:179 549:11 574:93 575:253 576:246 577:101 602:145 603:253 604:255 605:92 630:93 631:253 632:246 633:59 658:93 659:253 660:74 +0 127:46 128:105 129:254 130:254 131:224 132:59 133:59 134:9 155:196 156:254 157:253 158:253 159:253 160:253 161:253 162:128 182:96 183:235 184:254 185:253 186:253 187:253 188:253 189:253 190:247 191:122 208:4 209:101 210:244 211:253 212:254 213:234 214:241 215:253 216:253 217:253 218:253 219:186 220:18 236:96 237:253 238:253 239:253 240:232 241:83 242:109 243:170 244:253 245:253 246:253 247:253 248:116 264:215 265:253 266:253 267:253 268:196 271:40 272:253 273:253 274:253 275:253 276:116 290:8 291:141 292:247 293:253 294:253 295:237 296:29 299:6 300:38 301:171 302:253 303:253 304:116 317:13 318:146 319:253 320:253 321:253 322:253 323:57 329:156 330:253 331:253 332:116 345:40 346:253 347:253 348:253 349:253 350:178 351:27 357:156 358:253 359:253 360:116 372:136 373:204 374:253 375:253 376:253 377:192 378:27 385:156 386:253 387:253 388:116 399:28 400:195 401:254 402:254 403:254 404:250 405:135 412:99 413:255 414:254 415:254 416:117 427:118 428:253 429:253 430:253 431:253 432:142 439:19 440:170 441:253 442:253 443:216 444:62 454:42 455:212 456:253 457:253 458:253 459:253 460:38 466:124 467:188 468:253 469:253 470:253 471:174 482:59 483:253 484:253 485:253 486:237 487:93 488:3 491:31 492:40 493:130 494:247 495:253 496:253 497:253 498:204 499:13 510:59 511:253 512:253 513:253 514:154 518:54 519:218 520:254 521:253 522:253 523:253 524:253 525:253 526:38 538:59 539:253 540:253 541:253 542:215 543:156 544:156 545:156 546:209 547:253 548:255 549:253 550:253 551:253 552:192 553:97 554:15 566:55 567:242 568:253 569:253 570:253 571:253 572:253 573:253 574:253 575:253 576:254 577:253 578:253 579:204 580:23 595:118 596:253 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:254 605:216 606:174 607:13 623:54 624:116 625:243 626:253 627:253 628:253 629:253 630:253 631:146 632:117 633:62 653:53 654:132 655:253 656:253 657:192 658:57 659:13 +1 125:42 126:232 127:254 128:58 153:86 154:253 155:253 156:58 181:86 182:253 183:253 184:58 209:206 210:253 211:253 212:58 237:215 238:253 239:253 240:58 265:215 266:253 267:253 268:58 293:215 294:253 295:253 296:58 321:215 322:253 323:253 324:58 349:215 350:253 351:253 352:58 377:215 378:253 379:253 380:58 405:215 406:253 407:253 408:58 433:188 434:253 435:253 436:85 461:86 462:253 463:253 464:200 465:12 489:29 490:223 491:253 492:253 493:151 518:209 519:253 520:253 521:194 546:128 547:253 548:253 549:200 550:8 574:32 575:213 576:253 577:253 578:152 579:6 603:32 604:221 605:253 606:253 607:153 608:5 632:90 633:215 634:253 635:253 636:151 661:59 662:253 663:253 664:84 +1 156:60 157:229 158:38 184:187 185:254 186:78 211:121 212:252 213:254 214:78 239:197 240:254 241:206 242:6 267:197 268:254 269:202 294:27 295:218 296:233 297:62 322:117 323:254 324:195 350:203 351:254 352:195 377:64 378:244 379:254 380:195 405:79 406:254 407:255 408:161 433:79 434:254 435:254 436:65 461:79 462:254 463:241 464:52 489:79 490:254 491:189 517:79 518:254 519:189 545:79 546:254 547:189 573:79 574:254 575:189 601:79 602:254 603:194 604:5 629:35 630:219 631:254 632:72 658:34 659:223 660:195 687:129 688:195 +1 101:11 102:150 103:72 129:37 130:251 131:71 157:63 158:251 159:71 185:217 186:251 187:71 213:217 214:251 215:71 240:145 241:253 242:253 243:72 267:42 268:206 269:251 270:251 271:71 295:99 296:251 297:251 298:251 299:71 323:253 324:251 325:251 326:251 327:71 350:130 351:253 352:251 353:251 354:251 355:71 377:110 378:253 379:255 380:253 381:253 382:253 383:72 405:109 406:251 407:253 408:251 409:251 410:188 411:30 433:109 434:251 435:253 436:251 437:246 438:123 460:16 461:170 462:251 463:253 464:251 465:215 488:37 489:251 490:251 491:253 492:251 493:86 516:218 517:253 518:253 519:255 520:253 521:35 543:84 544:236 545:251 546:251 547:253 548:168 549:15 571:144 572:251 573:251 574:251 575:190 576:15 599:144 600:251 601:251 602:251 603:180 626:53 627:221 628:251 629:251 630:251 631:180 +0 127:45 128:254 129:254 130:254 131:148 132:24 133:9 154:43 155:254 156:252 157:252 158:252 159:252 160:252 161:121 162:13 181:58 182:237 183:254 184:252 185:252 186:252 187:252 188:252 189:252 190:68 208:69 209:224 210:252 211:254 212:252 213:252 214:252 215:252 216:252 217:252 218:135 219:17 235:75 236:216 237:252 238:252 239:254 240:231 241:168 242:252 243:252 244:252 245:252 246:252 247:45 262:77 263:212 264:252 265:252 266:252 267:242 268:93 269:32 270:114 271:177 272:252 273:252 274:252 275:158 276:12 289:75 290:212 291:252 292:252 293:252 294:252 295:231 299:116 300:252 301:252 302:252 303:252 304:21 316:69 317:216 318:252 319:252 320:252 321:252 322:252 323:62 327:116 328:252 329:252 330:252 331:252 332:21 344:93 345:252 346:252 347:252 348:252 349:252 350:62 355:21 356:158 357:252 358:252 359:252 360:21 371:64 372:239 373:252 374:252 375:252 376:252 377:252 378:21 384:139 385:252 386:252 387:252 388:21 398:5 399:87 400:254 401:254 402:254 403:254 404:237 405:41 411:11 412:150 413:254 414:254 415:254 416:22 425:5 426:85 427:252 428:252 429:252 430:252 431:222 432:55 439:116 440:252 441:252 442:252 443:214 444:18 453:24 454:252 455:252 456:252 457:252 458:252 459:91 466:26 467:153 468:252 469:252 470:252 471:45 481:24 482:252 483:252 484:252 485:252 486:252 487:91 492:18 493:93 494:151 495:252 496:252 497:252 498:184 499:28 509:24 510:252 511:252 512:252 513:252 514:252 515:164 516:116 517:116 518:116 519:117 520:141 521:252 522:252 523:252 524:252 525:252 526:68 537:24 538:252 539:252 540:252 541:252 542:252 543:252 544:252 545:252 546:252 547:254 548:252 549:252 550:252 551:252 552:252 553:163 554:31 565:9 566:121 567:252 568:252 569:252 570:252 571:252 572:252 573:252 574:252 575:254 576:252 577:252 578:252 579:178 580:91 581:33 594:13 595:119 596:252 597:252 598:252 599:252 600:252 601:252 602:252 603:254 604:252 605:252 606:184 607:37 623:13 624:121 625:252 626:252 627:252 628:252 629:252 630:252 631:254 632:214 633:45 634:28 652:8 653:21 654:21 655:169 656:252 657:252 658:41 659:22 660:18 +0 125:218 126:253 127:253 128:255 129:149 130:62 151:42 152:144 153:236 154:251 155:251 156:253 157:251 158:236 159:144 160:144 179:99 180:251 181:251 182:251 183:225 184:253 185:251 186:251 187:251 188:251 189:166 190:16 206:79 207:253 208:251 209:251 210:204 211:41 212:143 213:205 214:251 215:251 216:251 217:253 218:169 219:15 233:79 234:231 235:253 236:251 237:225 238:41 241:41 242:226 243:251 244:251 245:253 246:251 247:164 260:37 261:253 262:253 263:255 264:253 265:35 271:79 272:232 273:255 274:253 275:227 276:42 288:140 289:251 290:251 291:253 292:168 293:15 300:77 301:253 302:251 303:251 304:142 315:21 316:221 317:251 318:251 319:164 320:15 329:227 330:251 331:251 332:236 333:61 342:32 343:190 344:251 345:251 346:251 357:73 358:251 359:251 360:251 361:71 370:73 371:251 372:251 373:251 374:251 385:73 386:251 387:251 388:251 389:71 398:73 399:253 400:253 401:253 402:201 413:73 414:253 415:253 416:253 417:72 426:176 427:251 428:251 429:251 430:71 441:73 442:251 443:251 444:251 445:71 454:253 455:251 456:251 457:157 458:10 469:73 470:251 471:251 472:251 473:71 482:253 483:251 484:251 485:142 497:150 498:251 499:251 500:204 501:41 510:124 511:251 512:251 513:220 514:180 524:130 525:253 526:251 527:225 528:41 538:73 539:253 540:253 541:253 542:253 543:73 544:73 545:10 549:42 550:73 551:150 552:253 553:255 554:253 555:216 566:31 567:189 568:251 569:251 570:251 571:253 572:251 573:159 574:144 575:144 576:145 577:206 578:251 579:251 580:251 581:253 582:168 583:92 595:20 596:195 597:251 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:251 607:251 608:225 609:164 610:15 624:21 625:142 626:220 627:253 628:251 629:251 630:251 631:251 632:253 633:251 634:251 635:204 636:41 654:51 655:72 656:174 657:251 658:251 659:251 660:253 661:147 662:71 663:41 +0 127:60 128:96 129:96 130:48 153:16 154:171 155:228 156:253 157:251 158:220 159:51 160:32 181:127 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:80 207:24 208:182 209:236 210:251 211:211 212:189 213:236 214:251 215:251 216:251 217:242 218:193 234:100 235:194 236:251 237:251 238:211 239:35 241:71 242:173 243:251 244:251 245:253 246:240 247:158 248:19 261:64 262:253 263:255 264:253 265:205 266:19 271:40 272:218 273:255 274:253 275:253 276:91 288:16 289:186 290:251 291:253 292:247 293:110 300:39 301:233 302:251 303:251 304:188 315:16 316:189 317:251 318:251 319:205 320:110 329:48 330:220 331:251 332:220 333:48 343:72 344:251 345:251 346:251 347:158 358:51 359:251 360:251 361:232 371:190 372:251 373:251 374:251 375:59 386:32 387:251 388:251 389:251 398:96 399:253 400:253 401:253 402:95 414:32 415:253 416:253 417:193 426:214 427:251 428:251 429:204 430:23 442:52 443:251 444:251 445:94 454:253 455:251 456:251 457:109 469:48 470:221 471:251 472:219 473:47 482:253 483:251 484:251 485:70 497:234 498:251 499:251 500:188 510:253 511:251 512:251 513:188 523:40 524:158 525:253 526:251 527:172 528:70 539:191 540:253 541:253 542:253 543:96 544:24 549:12 550:174 551:253 552:253 553:255 554:221 567:71 568:251 569:251 570:251 571:253 572:205 573:190 574:190 575:190 576:191 577:197 578:251 579:251 580:231 581:221 582:93 595:16 596:126 597:236 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:251 607:140 608:47 625:67 626:188 627:189 628:188 629:188 630:188 631:188 632:189 633:188 634:109 635:4 +0 126:32 127:202 128:255 129:253 130:253 131:175 132:21 152:84 153:144 154:190 155:251 156:253 157:251 158:251 159:251 160:174 176:6 177:37 178:166 179:218 180:236 181:251 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:156 204:115 205:251 206:251 207:253 208:251 209:251 210:251 211:251 212:253 213:251 214:251 215:251 216:251 217:180 231:105 232:241 233:251 234:251 235:253 236:251 237:251 238:251 239:122 240:72 241:71 242:71 243:148 244:251 245:180 258:73 259:253 260:253 261:253 262:253 263:202 264:253 265:253 266:143 286:31 287:189 288:251 289:251 290:251 291:31 292:189 293:251 294:142 314:63 315:236 316:251 317:251 318:96 320:124 321:246 322:142 330:21 331:166 332:21 342:73 343:251 344:251 345:251 346:71 349:217 350:142 357:32 358:190 359:251 360:142 370:73 371:251 372:251 373:251 374:71 377:217 378:142 385:73 386:251 387:251 388:142 398:73 399:253 400:253 401:253 402:72 405:156 406:103 413:73 414:253 415:253 416:253 417:72 426:73 427:251 428:251 429:251 430:174 441:73 442:251 443:251 444:251 445:71 454:73 455:251 456:251 457:251 458:251 469:73 470:251 471:251 472:251 473:71 482:42 483:205 484:251 485:251 486:251 487:79 497:73 498:251 499:251 500:251 501:71 511:41 512:226 513:251 514:251 515:232 516:77 525:73 526:251 527:251 528:251 529:71 540:166 541:253 542:253 543:255 544:253 545:227 546:73 547:21 553:125 554:253 555:253 556:143 568:16 569:169 570:251 571:253 572:251 573:251 574:251 575:174 576:105 579:63 580:144 581:253 582:251 583:251 584:142 597:15 598:35 599:253 600:251 601:251 602:251 603:251 604:243 605:217 606:217 607:231 608:251 609:253 610:251 611:220 612:20 627:143 628:142 629:236 630:251 631:251 632:253 633:251 634:251 635:251 636:251 637:253 638:251 639:137 657:61 658:71 659:200 660:253 661:251 662:251 663:251 664:251 665:201 666:71 667:10 +1 130:218 131:170 132:108 157:32 158:227 159:252 160:232 185:129 186:252 187:252 188:252 212:1 213:253 214:252 215:252 216:168 240:144 241:253 242:252 243:236 244:62 268:144 269:253 270:252 271:215 296:144 297:253 298:252 299:112 323:21 324:206 325:253 326:252 327:71 351:99 352:253 353:255 354:119 378:63 379:242 380:252 381:253 382:35 406:94 407:252 408:252 409:154 410:10 433:145 434:237 435:252 436:252 461:255 462:253 463:253 464:108 487:11 488:155 489:253 490:252 491:179 492:15 514:11 515:150 516:252 517:253 518:200 519:20 542:73 543:252 544:252 545:253 546:97 569:47 570:233 571:253 572:253 596:1 597:149 598:252 599:252 600:252 624:1 625:252 626:252 627:246 628:132 652:1 653:169 654:252 655:132 +1 130:116 131:255 132:123 157:29 158:213 159:253 160:122 185:189 186:253 187:253 188:122 213:189 214:253 215:253 216:122 241:189 242:253 243:253 244:122 267:2 268:114 269:243 270:253 271:186 272:19 295:100 296:253 297:253 298:253 299:48 323:172 324:253 325:253 326:253 327:48 351:172 352:253 353:253 354:182 355:19 378:133 379:251 380:253 381:175 382:4 405:107 406:251 407:253 408:253 409:65 432:26 433:194 434:253 435:253 436:214 437:40 459:105 460:205 461:253 462:253 463:125 464:40 487:139 488:253 489:253 490:253 491:81 514:41 515:231 516:253 517:253 518:159 519:16 541:65 542:155 543:253 544:253 545:172 546:4 569:124 570:253 571:253 572:253 573:98 597:124 598:253 599:253 600:214 601:41 624:22 625:207 626:253 627:253 628:139 653:124 654:253 655:162 656:9 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_fpgrowth.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_fpgrowth.txt new file mode 100644 index 0000000..c451583 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_fpgrowth.txt @@ -0,0 +1,6 @@ +r z h k p +z y x w v u t s +s x o n r +x z y m t s q e +z +x z y r q t p diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_isotonic_regression_libsvm_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_isotonic_regression_libsvm_data.txt new file mode 100644 index 0000000..f39fe02 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_isotonic_regression_libsvm_data.txt @@ -0,0 +1,100 @@ +0.24579296 1:0.01 +0.28505864 1:0.02 +0.31208567 1:0.03 +0.35900051 1:0.04 +0.35747068 1:0.05 +0.16675166 1:0.06 +0.17491076 1:0.07 +0.04181540 1:0.08 +0.04793473 1:0.09 +0.03926568 1:0.10 +0.12952575 1:0.11 +0.00000000 1:0.12 +0.01376849 1:0.13 +0.13105558 1:0.14 +0.08873024 1:0.15 +0.12595614 1:0.16 +0.15247323 1:0.17 +0.25956145 1:0.18 +0.20040796 1:0.19 +0.19581846 1:0.20 +0.15757267 1:0.21 +0.13717491 1:0.22 +0.19020908 1:0.23 +0.19581846 1:0.24 +0.20091790 1:0.25 +0.16879143 1:0.26 +0.18510964 1:0.27 +0.20040796 1:0.28 +0.29576747 1:0.29 +0.43396226 1:0.30 +0.53391127 1:0.31 +0.52116267 1:0.32 +0.48546660 1:0.33 +0.49209587 1:0.34 +0.54156043 1:0.35 +0.59765426 1:0.36 +0.56144824 1:0.37 +0.58592555 1:0.38 +0.52983172 1:0.39 +0.50178480 1:0.40 +0.52626211 1:0.41 +0.58286588 1:0.42 +0.64660887 1:0.43 +0.68077511 1:0.44 +0.74298827 1:0.45 +0.64864865 1:0.46 +0.67261601 1:0.47 +0.65782764 1:0.48 +0.69811321 1:0.49 +0.63029067 1:0.50 +0.61601224 1:0.51 +0.63233044 1:0.52 +0.65323814 1:0.53 +0.65323814 1:0.54 +0.67363590 1:0.55 +0.67006629 1:0.56 +0.51555329 1:0.57 +0.50892402 1:0.58 +0.33299337 1:0.59 +0.36206017 1:0.60 +0.43090260 1:0.61 +0.45996940 1:0.62 +0.56348802 1:0.63 +0.54920959 1:0.64 +0.48393677 1:0.65 +0.48495665 1:0.66 +0.46965834 1:0.67 +0.45181030 1:0.68 +0.45843957 1:0.69 +0.47118817 1:0.70 +0.51555329 1:0.71 +0.58031617 1:0.72 +0.55481897 1:0.73 +0.56297807 1:0.74 +0.56603774 1:0.75 +0.57929628 1:0.76 +0.64762876 1:0.77 +0.66241713 1:0.78 +0.69301377 1:0.79 +0.65119837 1:0.80 +0.68332483 1:0.81 +0.66598674 1:0.82 +0.73890872 1:0.83 +0.73992861 1:0.84 +0.84242733 1:0.85 +0.91330954 1:0.86 +0.88016318 1:0.87 +0.90719021 1:0.88 +0.93115757 1:0.89 +0.93115757 1:0.90 +0.91942886 1:0.91 +0.92911780 1:0.92 +0.95665477 1:0.93 +0.95002550 1:0.94 +0.96940337 1:0.95 +1.00000000 1:0.96 +0.89801122 1:0.97 +0.90311066 1:0.98 +0.90362060 1:0.99 +0.83477817 1:1.0 \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_kmeans_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_kmeans_data.txt new file mode 100644 index 0000000..5001377 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_kmeans_data.txt @@ -0,0 +1,6 @@ +0 1:0.0 2:0.0 3:0.0 +1 1:0.1 2:0.1 3:0.1 +2 1:0.2 2:0.2 3:0.2 +3 1:9.0 2:9.0 3:9.0 +4 1:9.1 2:9.1 3:9.1 +5 1:9.2 2:9.2 3:9.2 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_data.txt new file mode 100644 index 0000000..2e76702 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_data.txt @@ -0,0 +1,12 @@ +1 2 6 0 2 3 1 1 0 0 3 +1 3 0 1 3 0 0 2 0 0 1 +1 4 1 0 0 4 9 0 1 2 0 +2 1 0 3 0 0 5 0 2 3 9 +3 1 1 9 3 0 2 0 0 1 3 +4 2 0 3 4 5 1 1 1 4 0 +2 1 0 3 0 0 5 0 2 2 9 +1 1 1 9 2 1 2 0 0 1 3 +4 4 0 3 4 2 1 3 0 0 0 +2 8 2 0 3 0 2 0 2 7 2 +1 1 1 9 0 2 2 0 0 3 3 +4 1 0 0 4 5 1 3 0 1 0 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_libsvm_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_libsvm_data.txt new file mode 100644 index 0000000..bf118d7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_lda_libsvm_data.txt @@ -0,0 +1,12 @@ +0 1:1 2:2 3:6 4:0 5:2 6:3 7:1 8:1 9:0 10:0 11:3 +1 1:1 2:3 3:0 4:1 5:3 6:0 7:0 8:2 9:0 10:0 11:1 +2 1:1 2:4 3:1 4:0 5:0 6:4 7:9 8:0 9:1 10:2 11:0 +3 1:2 2:1 3:0 4:3 5:0 6:0 7:5 8:0 9:2 10:3 11:9 +4 1:3 2:1 3:1 4:9 5:3 6:0 7:2 8:0 9:0 10:1 11:3 +5 1:4 2:2 3:0 4:3 5:4 6:5 7:1 8:1 9:1 10:4 11:0 +6 1:2 2:1 3:0 4:3 5:0 6:0 7:5 8:0 9:2 10:2 11:9 +7 1:1 2:1 3:1 4:9 5:2 6:1 7:2 8:0 9:0 10:1 11:3 +8 1:4 2:4 3:0 4:3 5:4 6:2 7:1 8:3 9:0 10:0 11:0 +9 1:2 2:8 3:2 4:0 5:3 6:0 7:2 8:0 9:2 10:7 11:2 +10 1:1 2:1 3:1 4:9 5:0 6:2 7:2 8:0 9:0 10:3 11:3 +11 1:4 2:1 3:0 4:0 5:4 6:5 7:1 8:3 9:0 10:1 11:0 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_libsvm_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_libsvm_data.txt new file mode 100644 index 0000000..861c70c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_libsvm_data.txt @@ -0,0 +1,100 @@ +0 128:51 129:159 130:253 131:159 132:50 155:48 156:238 157:252 158:252 159:252 160:237 182:54 183:227 184:253 185:252 186:239 187:233 188:252 189:57 190:6 208:10 209:60 210:224 211:252 212:253 213:252 214:202 215:84 216:252 217:253 218:122 236:163 237:252 238:252 239:252 240:253 241:252 242:252 243:96 244:189 245:253 246:167 263:51 264:238 265:253 266:253 267:190 268:114 269:253 270:228 271:47 272:79 273:255 274:168 290:48 291:238 292:252 293:252 294:179 295:12 296:75 297:121 298:21 301:253 302:243 303:50 317:38 318:165 319:253 320:233 321:208 322:84 329:253 330:252 331:165 344:7 345:178 346:252 347:240 348:71 349:19 350:28 357:253 358:252 359:195 372:57 373:252 374:252 375:63 385:253 386:252 387:195 400:198 401:253 402:190 413:255 414:253 415:196 427:76 428:246 429:252 430:112 441:253 442:252 443:148 455:85 456:252 457:230 458:25 467:7 468:135 469:253 470:186 471:12 483:85 484:252 485:223 494:7 495:131 496:252 497:225 498:71 511:85 512:252 513:145 521:48 522:165 523:252 524:173 539:86 540:253 541:225 548:114 549:238 550:253 551:162 567:85 568:252 569:249 570:146 571:48 572:29 573:85 574:178 575:225 576:253 577:223 578:167 579:56 595:85 596:252 597:252 598:252 599:229 600:215 601:252 602:252 603:252 604:196 605:130 623:28 624:199 625:252 626:252 627:253 628:252 629:252 630:233 631:145 652:25 653:128 654:252 655:253 656:252 657:141 658:37 +1 159:124 160:253 161:255 162:63 186:96 187:244 188:251 189:253 190:62 214:127 215:251 216:251 217:253 218:62 241:68 242:236 243:251 244:211 245:31 246:8 268:60 269:228 270:251 271:251 272:94 296:155 297:253 298:253 299:189 323:20 324:253 325:251 326:235 327:66 350:32 351:205 352:253 353:251 354:126 378:104 379:251 380:253 381:184 382:15 405:80 406:240 407:251 408:193 409:23 432:32 433:253 434:253 435:253 436:159 460:151 461:251 462:251 463:251 464:39 487:48 488:221 489:251 490:251 491:172 515:234 516:251 517:251 518:196 519:12 543:253 544:251 545:251 546:89 570:159 571:255 572:253 573:253 574:31 597:48 598:228 599:253 600:247 601:140 602:8 625:64 626:251 627:253 628:220 653:64 654:251 655:253 656:220 681:24 682:193 683:253 684:220 +1 125:145 126:255 127:211 128:31 152:32 153:237 154:253 155:252 156:71 180:11 181:175 182:253 183:252 184:71 209:144 210:253 211:252 212:71 236:16 237:191 238:253 239:252 240:71 264:26 265:221 266:253 267:252 268:124 269:31 293:125 294:253 295:252 296:252 297:108 322:253 323:252 324:252 325:108 350:255 351:253 352:253 353:108 378:253 379:252 380:252 381:108 406:253 407:252 408:252 409:108 434:253 435:252 436:252 437:108 462:255 463:253 464:253 465:170 490:253 491:252 492:252 493:252 494:42 518:149 519:252 520:252 521:252 522:144 546:109 547:252 548:252 549:252 550:144 575:218 576:253 577:253 578:255 579:35 603:175 604:252 605:252 606:253 607:35 631:73 632:252 633:252 634:253 635:35 659:31 660:211 661:252 662:253 663:35 +1 153:5 154:63 155:197 181:20 182:254 183:230 184:24 209:20 210:254 211:254 212:48 237:20 238:254 239:255 240:48 265:20 266:254 267:254 268:57 293:20 294:254 295:254 296:108 321:16 322:239 323:254 324:143 350:178 351:254 352:143 378:178 379:254 380:143 406:178 407:254 408:162 434:178 435:254 436:240 462:113 463:254 464:240 490:83 491:254 492:245 493:31 518:79 519:254 520:246 521:38 547:214 548:254 549:150 575:144 576:241 577:8 603:144 604:240 605:2 631:144 632:254 633:82 659:230 660:247 661:40 687:168 688:209 689:31 +1 152:1 153:168 154:242 155:28 180:10 181:228 182:254 183:100 209:190 210:254 211:122 237:83 238:254 239:162 265:29 266:254 267:248 268:25 293:29 294:255 295:254 296:103 321:29 322:254 323:254 324:109 349:29 350:254 351:254 352:109 377:29 378:254 379:254 380:109 405:29 406:255 407:254 408:109 433:29 434:254 435:254 436:109 461:29 462:254 463:254 464:63 489:29 490:254 491:254 492:28 517:29 518:254 519:254 520:28 545:29 546:254 547:254 548:35 573:29 574:254 575:254 576:109 601:6 602:212 603:254 604:109 630:203 631:254 632:178 658:155 659:254 660:190 686:32 687:199 688:104 +0 130:64 131:253 132:255 133:63 157:96 158:205 159:251 160:253 161:205 162:111 163:4 184:96 185:189 186:251 187:251 188:253 189:251 190:251 191:31 209:16 210:64 211:223 212:244 213:251 214:251 215:211 216:213 217:251 218:251 219:31 236:80 237:181 238:251 239:253 240:251 241:251 242:251 243:94 244:96 245:251 246:251 247:31 263:92 264:253 265:253 266:253 267:255 268:253 269:253 270:253 271:95 272:96 273:253 274:253 275:31 290:92 291:236 292:251 293:243 294:220 295:233 296:251 297:251 298:243 299:82 300:96 301:251 302:251 303:31 317:80 318:253 319:251 320:251 321:188 323:96 324:251 325:251 326:109 328:96 329:251 330:251 331:31 344:96 345:240 346:253 347:243 348:188 349:42 351:96 352:204 353:109 354:4 356:12 357:197 358:251 359:31 372:221 373:251 374:253 375:121 379:36 380:23 385:190 386:251 387:31 399:48 400:234 401:253 413:191 414:253 415:31 426:44 427:221 428:251 429:251 440:12 441:197 442:251 443:31 454:190 455:251 456:251 457:251 468:96 469:251 470:251 471:31 482:190 483:251 484:251 485:113 495:40 496:234 497:251 498:219 499:23 510:190 511:251 512:251 513:94 522:40 523:217 524:253 525:231 526:47 538:191 539:253 540:253 541:253 548:12 549:174 550:253 551:253 552:219 553:39 566:67 567:236 568:251 569:251 570:191 571:190 572:111 573:72 574:190 575:191 576:197 577:251 578:243 579:121 580:39 595:63 596:236 597:251 598:253 599:251 600:251 601:251 602:251 603:253 604:251 605:188 606:94 624:27 625:129 626:253 627:251 628:251 629:251 630:251 631:229 632:168 633:15 654:95 655:212 656:251 657:211 658:94 659:59 +1 159:121 160:254 161:136 186:13 187:230 188:253 189:248 190:99 213:4 214:118 215:253 216:253 217:225 218:42 241:61 242:253 243:253 244:253 245:74 268:32 269:206 270:253 271:253 272:186 273:9 296:211 297:253 298:253 299:239 300:69 324:254 325:253 326:253 327:133 351:142 352:255 353:253 354:186 355:8 378:149 379:229 380:254 381:207 382:21 405:54 406:229 407:253 408:254 409:105 433:152 434:254 435:254 436:213 437:26 460:112 461:251 462:253 463:253 464:26 487:29 488:212 489:253 490:250 491:149 514:36 515:214 516:253 517:253 518:137 542:75 543:253 544:253 545:253 546:59 570:93 571:253 572:253 573:189 574:17 598:224 599:253 600:253 601:84 625:43 626:235 627:253 628:126 629:1 653:99 654:248 655:253 656:119 682:225 683:235 684:49 +1 100:166 101:222 102:55 128:197 129:254 130:218 131:5 155:29 156:249 157:254 158:254 159:9 183:45 184:254 185:254 186:174 187:2 210:4 211:164 212:254 213:254 214:85 238:146 239:254 240:254 241:254 242:85 265:101 266:245 267:254 268:254 269:254 270:85 292:97 293:248 294:254 295:204 296:254 297:254 298:85 315:12 316:59 317:98 318:151 319:237 320:254 321:254 322:109 323:35 324:254 325:254 326:85 343:41 344:216 345:254 346:254 347:239 348:153 349:37 350:4 351:32 352:254 353:254 354:85 372:7 373:44 374:44 375:30 379:32 380:254 381:254 382:96 407:19 408:230 409:254 410:174 436:197 437:254 438:110 464:197 465:254 466:85 492:197 493:253 494:63 515:37 516:54 517:54 518:45 519:26 520:84 521:221 522:84 523:21 524:31 525:162 526:78 540:6 541:41 542:141 543:244 544:254 545:254 546:248 547:236 548:254 549:254 550:254 551:233 552:239 553:254 554:138 567:23 568:167 569:254 570:254 571:254 572:254 573:229 574:228 575:185 576:138 577:138 578:138 579:138 580:138 581:138 582:44 595:113 596:254 597:254 598:254 599:179 600:64 601:5 623:32 624:209 625:183 626:97 +0 155:53 156:255 157:253 158:253 159:253 160:124 183:180 184:253 185:251 186:251 187:251 188:251 189:145 190:62 209:32 210:217 211:241 212:253 213:251 214:251 215:251 216:251 217:253 218:107 237:37 238:251 239:251 240:253 241:251 242:251 243:251 244:251 245:253 246:107 265:166 266:251 267:251 268:253 269:251 270:96 271:148 272:251 273:253 274:107 291:73 292:253 293:253 294:253 295:253 296:130 299:110 300:253 301:255 302:108 319:73 320:251 321:251 322:251 323:251 327:109 328:251 329:253 330:107 347:202 348:251 349:251 350:251 351:225 354:6 355:129 356:251 357:253 358:107 375:150 376:251 377:251 378:251 379:71 382:115 383:251 384:251 385:253 386:107 403:253 404:251 405:251 406:173 407:20 410:217 411:251 412:251 413:253 414:107 430:182 431:255 432:253 433:216 438:218 439:253 440:253 441:182 457:63 458:221 459:253 460:251 461:215 465:84 466:236 467:251 468:251 469:77 485:109 486:251 487:253 488:251 489:215 492:11 493:160 494:251 495:251 496:96 513:109 514:251 515:253 516:251 517:137 520:150 521:251 522:251 523:251 524:71 541:109 542:251 543:253 544:251 545:35 547:130 548:253 549:251 550:251 551:173 552:20 569:110 570:253 571:255 572:253 573:98 574:150 575:253 576:255 577:253 578:164 597:109 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:35 625:93 626:241 627:253 628:251 629:251 630:251 631:251 632:216 633:112 634:5 654:103 655:253 656:251 657:251 658:251 659:251 683:124 684:251 685:225 686:71 687:71 +0 128:73 129:253 130:227 131:73 132:21 156:73 157:251 158:251 159:251 160:174 182:16 183:166 184:228 185:251 186:251 187:251 188:122 210:62 211:220 212:253 213:251 214:251 215:251 216:251 217:79 238:79 239:231 240:253 241:251 242:251 243:251 244:251 245:232 246:77 264:145 265:253 266:253 267:253 268:255 269:253 270:253 271:253 272:253 273:255 274:108 292:144 293:251 294:251 295:251 296:253 297:168 298:107 299:169 300:251 301:253 302:189 303:20 318:27 319:89 320:236 321:251 322:235 323:215 324:164 325:15 326:6 327:129 328:251 329:253 330:251 331:35 345:47 346:211 347:253 348:251 349:251 350:142 354:37 355:251 356:251 357:253 358:251 359:35 373:109 374:251 375:253 376:251 377:251 378:142 382:11 383:148 384:251 385:253 386:251 387:164 400:11 401:150 402:253 403:255 404:211 405:25 410:11 411:150 412:253 413:255 414:211 415:25 428:140 429:251 430:251 431:253 432:107 438:37 439:251 440:251 441:211 442:46 456:190 457:251 458:251 459:253 460:128 461:5 466:37 467:251 468:251 469:51 484:115 485:251 486:251 487:253 488:188 489:20 492:32 493:109 494:129 495:251 496:173 497:103 512:217 513:251 514:251 515:201 516:30 520:73 521:251 522:251 523:251 524:71 540:166 541:253 542:253 543:255 544:149 545:73 546:150 547:253 548:255 549:253 550:253 551:143 568:140 569:251 570:251 571:253 572:251 573:251 574:251 575:251 576:253 577:251 578:230 579:61 596:190 597:251 598:251 599:253 600:251 601:251 602:251 603:251 604:242 605:215 606:55 624:21 625:189 626:251 627:253 628:251 629:251 630:251 631:173 632:103 653:31 654:200 655:253 656:251 657:96 658:71 659:20 +1 155:178 156:255 157:105 182:6 183:188 184:253 185:216 186:14 210:14 211:202 212:253 213:253 214:23 238:12 239:199 240:253 241:128 242:6 266:42 267:253 268:253 269:158 294:42 295:253 296:253 297:158 322:155 323:253 324:253 325:158 350:160 351:253 352:253 353:147 378:160 379:253 380:253 381:41 405:17 406:225 407:253 408:235 409:31 433:24 434:253 435:253 436:176 461:24 462:253 463:253 464:176 489:24 490:253 491:253 492:176 517:24 518:253 519:253 520:176 545:24 546:253 547:253 548:162 573:46 574:253 575:253 576:59 601:142 602:253 603:253 604:59 629:142 630:253 631:253 632:59 657:142 658:253 659:202 660:8 685:87 686:253 687:139 +0 154:46 155:105 156:254 157:254 158:254 159:254 160:255 161:239 162:41 180:37 181:118 182:222 183:254 184:253 185:253 186:253 187:253 188:253 189:253 190:211 191:54 207:14 208:200 209:253 210:253 211:254 212:253 213:253 214:253 215:253 216:253 217:253 218:253 219:116 233:16 234:160 235:236 236:253 237:253 238:253 239:254 240:253 241:253 242:246 243:229 244:253 245:253 246:253 247:116 261:99 262:253 263:253 264:253 265:253 266:253 267:254 268:253 269:253 270:213 271:99 272:253 273:253 274:253 275:116 288:25 289:194 290:253 291:253 292:253 293:253 294:131 295:97 296:169 297:253 298:93 299:99 300:253 301:253 302:253 303:116 316:206 317:253 318:253 319:251 320:233 321:127 322:9 324:18 325:38 326:3 327:15 328:171 329:253 330:253 331:116 343:55 344:240 345:253 346:253 347:233 355:31 356:186 357:253 358:253 359:116 371:176 372:253 373:253 374:253 375:127 383:99 384:253 385:253 386:253 387:116 399:176 400:253 401:253 402:131 403:9 411:99 412:253 413:253 414:253 415:116 426:119 427:254 428:254 429:232 430:75 440:158 441:254 442:254 443:117 454:118 455:253 456:253 457:154 468:156 469:253 470:253 471:116 482:118 483:253 484:253 485:154 496:156 497:253 498:253 499:116 509:46 510:222 511:253 512:253 513:154 522:7 523:116 524:246 525:253 526:180 527:9 538:118 539:253 540:253 541:154 550:116 551:253 552:253 553:253 554:174 566:118 567:253 568:253 569:154 577:110 578:246 579:253 580:253 581:240 582:67 594:118 595:253 596:253 597:238 598:215 599:49 600:20 601:20 602:20 603:66 604:215 605:241 606:253 607:245 608:233 609:64 622:82 623:229 624:253 625:253 626:253 627:253 628:253 629:253 630:253 631:254 632:253 633:253 634:240 635:107 651:176 652:253 653:253 654:253 655:253 656:253 657:253 658:253 659:254 660:253 661:253 662:108 679:40 680:239 681:253 682:253 683:253 684:253 685:253 686:253 687:254 688:161 689:57 690:4 +0 152:56 153:105 154:220 155:254 156:63 178:18 179:166 180:233 181:253 182:253 183:253 184:236 185:209 186:209 187:209 188:77 189:18 206:84 207:253 208:253 209:253 210:253 211:253 212:254 213:253 214:253 215:253 216:253 217:172 218:8 233:57 234:238 235:253 236:253 237:253 238:253 239:253 240:254 241:253 242:253 243:253 244:253 245:253 246:119 260:14 261:238 262:253 263:253 264:253 265:253 266:253 267:253 268:179 269:196 270:253 271:253 272:253 273:253 274:238 275:12 288:33 289:253 290:253 291:253 292:253 293:253 294:248 295:134 297:18 298:83 299:237 300:253 301:253 302:253 303:14 316:164 317:253 318:253 319:253 320:253 321:253 322:128 327:57 328:119 329:214 330:253 331:94 343:57 344:248 345:253 346:253 347:253 348:126 349:14 350:4 357:179 358:253 359:248 360:56 371:175 372:253 373:253 374:240 375:190 376:28 385:179 386:253 387:253 388:173 399:209 400:253 401:253 402:178 413:92 414:253 415:253 416:208 427:211 428:254 429:254 430:179 442:135 443:255 444:209 455:209 456:253 457:253 458:90 470:134 471:253 472:208 483:209 484:253 485:253 486:178 497:2 498:142 499:253 500:208 511:209 512:253 513:253 514:214 515:35 525:30 526:253 527:253 528:208 539:165 540:253 541:253 542:253 543:215 544:36 553:163 554:253 555:253 556:164 567:18 568:172 569:253 570:253 571:253 572:214 573:127 574:7 580:72 581:232 582:253 583:171 584:17 596:8 597:182 598:253 599:253 600:253 601:253 602:162 603:56 607:64 608:240 609:253 610:253 611:14 625:7 626:173 627:253 628:253 629:253 630:253 631:245 632:241 633:239 634:239 635:246 636:253 637:225 638:14 639:1 654:18 655:59 656:138 657:224 658:253 659:253 660:254 661:253 662:253 663:253 664:240 665:96 685:37 686:104 687:192 688:255 689:253 690:253 691:182 692:73 +1 130:7 131:176 132:254 133:224 158:51 159:253 160:253 161:223 185:4 186:170 187:253 188:253 189:214 213:131 214:253 215:253 216:217 217:39 241:209 242:253 243:253 244:134 268:75 269:240 270:253 271:239 272:26 296:184 297:253 298:245 299:63 323:142 324:255 325:253 326:185 350:62 351:229 352:254 353:242 354:73 377:54 378:229 379:253 380:254 381:105 405:152 406:254 407:254 408:213 409:26 432:32 433:243 434:253 435:253 436:115 459:2 460:142 461:253 462:253 463:155 487:30 488:253 489:253 490:232 491:55 515:75 516:253 517:253 518:164 542:72 543:232 544:253 545:189 546:17 570:224 571:253 572:253 573:163 597:43 598:235 599:253 600:253 601:195 602:21 625:28 626:231 627:253 628:253 629:184 630:14 654:225 655:253 656:253 657:75 +0 155:21 156:176 157:253 158:253 159:124 182:105 183:176 184:251 185:251 186:251 187:251 188:105 208:58 209:217 210:241 211:253 212:251 213:251 214:251 215:251 216:243 217:113 218:5 235:63 236:231 237:251 238:251 239:253 240:251 241:251 242:251 243:251 244:253 245:251 246:113 263:144 264:251 265:251 266:251 267:253 268:251 269:251 270:251 271:251 272:253 273:251 274:215 290:125 291:253 292:253 293:253 294:253 295:255 296:253 297:253 298:253 299:253 300:255 301:253 302:227 303:42 318:253 319:251 320:251 321:251 322:251 323:253 324:251 325:251 326:251 327:251 328:253 329:251 330:251 331:142 345:27 346:253 347:251 348:251 349:235 350:241 351:253 352:251 353:246 354:137 355:35 356:98 357:251 358:251 359:236 360:61 372:47 373:211 374:253 375:251 376:235 377:82 378:103 379:253 380:251 381:137 384:73 385:251 386:251 387:251 388:71 399:27 400:211 401:251 402:253 403:251 404:86 407:72 408:71 409:10 412:73 413:251 414:251 415:173 416:20 427:89 428:253 429:253 430:255 431:253 432:35 440:73 441:253 442:253 443:253 444:72 454:84 455:236 456:251 457:251 458:253 459:251 460:138 468:73 469:251 470:251 471:251 472:71 481:63 482:236 483:251 484:251 485:251 486:227 487:251 488:246 489:138 490:11 494:16 495:37 496:228 497:251 498:246 499:137 500:10 509:73 510:251 511:251 512:251 513:173 514:42 515:142 516:142 517:142 518:41 522:109 523:251 524:253 525:251 526:137 537:73 538:251 539:251 540:173 541:20 549:27 550:211 551:251 552:253 553:147 554:10 565:73 566:253 567:253 568:143 575:21 576:176 577:253 578:253 579:253 593:73 594:251 595:251 596:205 597:144 603:176 604:251 605:251 606:188 607:107 621:62 622:236 623:251 624:251 625:251 626:218 627:217 628:217 629:217 630:217 631:253 632:230 633:189 634:20 650:83 651:158 652:251 653:251 654:253 655:251 656:251 657:251 658:251 659:253 660:107 679:37 680:251 681:251 682:253 683:251 684:251 685:251 686:122 687:72 688:30 +1 151:68 152:45 153:131 154:131 155:131 156:101 157:68 158:92 159:44 187:19 188:170 211:29 212:112 213:89 215:40 216:222 239:120 240:254 241:251 242:127 243:40 244:222 267:197 268:254 269:254 270:91 271:40 272:222 294:64 295:247 296:254 297:236 298:50 299:40 300:107 322:184 323:254 324:254 325:91 327:6 328:14 350:203 351:254 352:254 353:71 377:23 378:218 379:254 380:254 381:71 405:113 406:254 407:255 408:239 409:53 433:210 434:254 435:254 436:195 460:62 461:242 462:254 463:241 464:88 468:28 488:86 489:254 490:254 491:189 495:28 496:104 516:106 517:254 518:254 519:168 523:40 524:91 544:216 545:254 546:245 547:51 551:35 552:80 572:216 573:254 574:102 599:55 600:239 601:254 602:52 627:166 628:254 629:210 630:23 655:223 656:252 657:104 683:223 684:169 +0 125:29 126:170 127:255 128:255 129:141 151:29 152:198 153:255 154:255 155:255 156:226 157:255 158:86 178:141 179:255 180:255 181:170 182:29 184:86 185:255 186:255 187:141 204:29 205:226 206:255 207:198 208:57 213:226 214:255 215:255 216:226 217:114 231:29 232:255 233:255 234:114 241:141 242:170 243:114 244:255 245:255 246:141 259:226 260:255 261:170 269:29 270:57 273:141 274:255 275:226 286:57 287:255 288:170 302:114 303:255 304:198 314:226 315:255 331:170 332:255 333:57 342:255 343:226 360:255 361:170 370:255 371:170 388:114 389:198 398:255 399:226 416:86 417:255 426:198 427:255 444:86 445:255 454:114 455:255 456:57 472:86 473:255 482:29 483:255 484:226 500:141 501:255 511:170 512:255 513:170 528:226 529:198 539:29 540:226 541:255 542:170 555:29 556:255 557:114 568:29 569:226 570:255 571:141 582:57 583:226 584:226 598:141 599:255 600:255 601:170 602:86 607:29 608:86 609:226 610:255 611:226 612:29 627:86 628:198 629:255 630:255 631:255 632:255 633:255 634:255 635:255 636:255 637:255 638:141 639:29 657:29 658:114 659:170 660:170 661:170 662:170 663:170 664:86 +0 153:203 154:254 155:252 156:252 157:252 158:214 159:51 160:20 180:62 181:221 182:252 183:250 184:250 185:250 186:252 187:250 188:160 189:20 207:62 208:211 209:250 210:252 211:250 212:250 213:250 214:252 215:250 216:250 217:49 234:41 235:221 236:250 237:250 238:252 239:250 240:250 241:250 242:252 243:250 244:128 245:10 262:254 263:252 264:252 265:252 266:254 267:252 268:252 269:252 270:254 271:252 272:252 273:90 290:150 291:190 292:250 293:250 294:252 295:250 296:250 297:169 298:171 299:250 300:250 301:250 302:82 318:31 319:191 320:250 321:250 322:252 323:189 324:100 325:20 326:172 327:250 328:250 329:250 330:80 346:213 347:250 348:250 349:250 350:212 351:29 354:252 355:250 356:250 357:250 374:92 375:252 376:252 377:252 382:51 383:252 384:252 385:252 386:203 401:82 402:252 403:250 404:250 405:169 410:132 411:250 412:250 413:250 414:121 428:92 429:231 430:252 431:250 432:159 433:20 438:252 439:250 440:250 441:250 456:30 457:211 458:252 459:250 460:221 461:40 466:90 467:250 468:250 469:250 470:163 484:31 485:213 486:254 487:232 488:80 494:92 495:252 496:252 497:212 498:163 512:151 513:250 514:252 515:149 522:252 523:250 524:250 525:49 540:60 541:221 542:252 543:210 544:60 550:252 551:250 552:250 553:49 569:202 570:252 571:250 572:221 573:40 576:123 577:202 578:252 579:250 580:250 581:49 596:123 597:243 598:255 599:252 600:252 601:252 602:254 603:252 604:252 605:252 606:254 607:252 608:100 625:121 626:171 627:250 628:250 629:250 630:252 631:250 632:250 633:250 634:252 635:250 636:100 654:20 655:160 656:250 657:250 658:252 659:250 660:250 661:250 662:252 663:189 664:40 683:20 684:170 685:250 686:252 687:250 688:128 689:49 690:49 691:29 +1 98:64 99:191 100:70 125:68 126:243 127:253 128:249 129:63 152:30 153:223 154:253 155:253 156:247 157:41 179:73 180:238 181:253 182:253 183:253 184:242 206:73 207:236 208:253 209:253 210:253 211:253 212:242 234:182 235:253 236:253 237:191 238:247 239:253 240:149 262:141 263:253 264:143 265:86 266:249 267:253 268:122 290:9 291:36 292:7 293:14 294:233 295:253 296:122 322:230 323:253 324:122 350:230 351:253 352:122 378:231 379:255 380:123 406:230 407:253 408:52 433:61 434:245 435:253 461:98 462:253 463:253 468:35 469:12 489:98 490:253 491:253 494:9 495:142 496:233 497:146 517:190 518:253 519:253 520:128 521:7 522:99 523:253 524:253 525:180 544:29 545:230 546:253 547:253 548:252 549:210 550:253 551:253 552:253 553:140 571:28 572:207 573:253 574:253 575:253 576:254 577:253 578:253 579:235 580:70 581:9 599:126 600:253 601:253 602:253 603:253 604:254 605:253 606:168 607:19 627:79 628:253 629:253 630:201 631:190 632:132 633:63 634:5 +1 125:26 126:240 127:72 153:25 154:238 155:208 182:209 183:226 184:14 210:209 211:254 212:43 238:175 239:254 240:128 266:63 267:254 268:204 294:107 295:254 296:204 322:88 323:254 324:204 350:55 351:254 352:204 378:126 379:254 380:204 406:126 407:254 408:189 434:169 435:254 436:121 462:209 463:254 464:193 490:209 491:254 492:111 517:22 518:235 519:254 520:37 545:137 546:254 547:227 548:16 573:205 574:255 575:185 601:205 602:254 603:125 629:205 630:254 631:125 657:111 658:212 659:43 +0 155:62 156:91 157:213 158:255 159:228 160:91 161:12 182:70 183:230 184:253 185:253 186:253 187:253 188:253 189:152 190:7 210:246 211:253 212:253 213:253 214:253 215:253 216:253 217:253 218:106 237:21 238:247 239:253 240:253 241:253 242:253 243:253 244:253 245:208 246:24 265:156 266:253 267:253 268:253 269:253 270:253 271:253 272:253 273:195 292:88 293:238 294:253 295:253 296:253 297:221 298:253 299:253 300:253 301:195 320:230 321:253 322:253 323:253 324:198 325:40 326:177 327:253 328:253 329:195 346:56 347:156 348:251 349:253 350:189 351:182 352:15 354:86 355:240 356:253 357:210 358:28 374:213 375:253 376:253 377:156 378:3 383:205 384:253 385:253 386:106 401:121 402:252 403:253 404:135 405:3 411:46 412:253 413:253 414:106 428:28 429:212 430:253 431:248 432:23 439:42 440:253 441:253 442:106 456:197 457:253 458:234 459:70 467:42 468:253 469:253 470:106 483:11 484:202 485:253 486:187 495:58 496:253 497:210 498:27 511:107 512:253 513:253 514:40 522:53 523:227 524:253 525:195 539:107 540:253 541:253 542:40 549:47 550:227 551:253 552:231 553:58 567:107 568:253 569:253 570:40 575:5 576:131 577:222 578:253 579:231 580:59 595:14 596:204 597:253 598:226 599:222 600:73 601:58 602:58 603:170 604:253 605:253 606:227 607:58 624:197 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:238 634:58 652:33 653:179 654:241 655:253 656:253 657:253 658:253 659:250 660:116 661:14 682:75 683:179 684:253 685:151 686:89 687:86 +1 157:42 158:228 159:253 160:253 185:144 186:251 187:251 188:251 212:89 213:236 214:251 215:235 216:215 239:79 240:253 241:251 242:251 243:142 267:180 268:253 269:251 270:251 271:142 294:32 295:202 296:255 297:253 298:216 322:109 323:251 324:253 325:251 326:112 349:6 350:129 351:251 352:253 353:127 354:5 377:37 378:251 379:251 380:253 381:107 405:166 406:251 407:251 408:201 409:30 432:42 433:228 434:253 435:253 460:144 461:251 462:251 463:147 487:63 488:236 489:251 490:251 491:71 515:150 516:251 517:251 518:204 519:41 543:253 544:251 545:251 546:142 571:255 572:253 573:164 598:105 599:253 600:251 601:35 626:180 627:253 628:251 629:35 654:180 655:253 656:251 657:35 682:180 683:253 684:251 685:35 +1 128:62 129:254 130:213 156:102 157:253 158:252 159:102 160:20 184:102 185:254 186:253 187:254 188:50 212:102 213:253 214:252 215:253 216:50 240:102 241:254 242:253 243:254 244:50 268:142 269:253 270:252 271:253 272:50 295:51 296:253 297:254 298:253 299:224 300:20 323:132 324:252 325:253 326:252 327:162 351:173 352:253 353:254 354:253 355:102 378:82 379:253 380:252 381:253 382:252 383:61 406:203 407:254 408:253 409:254 410:233 433:41 434:243 435:253 436:252 437:253 438:111 461:132 462:253 463:254 464:253 465:203 488:41 489:253 490:252 491:253 492:252 493:40 515:11 516:213 517:254 518:253 519:254 520:151 543:92 544:252 545:253 546:252 547:192 548:50 570:21 571:214 572:253 573:255 574:253 575:41 598:142 599:253 600:252 601:253 602:171 625:113 626:253 627:255 628:253 629:203 630:40 653:30 654:131 655:233 656:111 +0 154:28 155:195 156:254 157:254 158:254 159:254 160:254 161:255 162:61 181:6 182:191 183:253 184:253 185:253 186:253 187:253 188:253 189:253 190:60 208:26 209:190 210:253 211:253 212:253 213:253 214:240 215:191 216:242 217:253 218:60 235:15 236:187 237:253 238:253 239:253 240:253 241:253 242:200 244:211 245:253 246:60 262:22 263:66 264:253 265:253 266:253 267:253 268:241 269:209 270:44 271:23 272:218 273:253 274:60 290:124 291:253 292:253 293:253 294:253 295:253 296:182 299:131 300:253 301:253 302:60 318:38 319:217 320:253 321:253 322:244 323:111 324:37 327:131 328:253 329:253 330:60 346:124 347:253 348:253 349:253 350:165 354:22 355:182 356:253 357:253 358:60 374:124 375:253 376:253 377:240 378:45 382:53 383:253 384:253 385:249 386:58 401:16 402:168 403:253 404:216 405:45 410:53 411:253 412:253 413:138 429:159 430:253 431:253 432:147 438:53 439:253 440:253 441:138 456:136 457:252 458:253 459:227 460:5 466:53 467:253 468:243 469:101 484:140 485:253 486:253 487:124 494:156 495:253 496:218 511:13 512:164 513:253 514:142 515:5 521:32 522:233 523:253 524:218 539:62 540:253 541:253 542:130 548:37 549:203 550:253 551:253 552:127 567:62 568:253 569:253 570:147 571:36 572:36 573:36 574:36 575:151 576:222 577:253 578:245 579:127 580:8 595:34 596:202 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:200 624:140 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:248 633:235 634:65 652:87 653:173 654:253 655:253 656:253 657:253 658:253 659:253 660:182 681:14 682:78 683:96 684:253 685:253 686:253 687:137 688:56 +0 123:8 124:76 125:202 126:254 127:255 128:163 129:37 130:2 150:13 151:182 152:253 153:253 154:253 155:253 156:253 157:253 158:23 177:15 178:179 179:253 180:253 181:212 182:91 183:218 184:253 185:253 186:179 187:109 205:105 206:253 207:253 208:160 209:35 210:156 211:253 212:253 213:253 214:253 215:250 216:113 232:19 233:212 234:253 235:253 236:88 237:121 238:253 239:233 240:128 241:91 242:245 243:253 244:248 245:114 260:104 261:253 262:253 263:110 264:2 265:142 266:253 267:90 270:26 271:199 272:253 273:248 274:63 287:1 288:173 289:253 290:253 291:29 293:84 294:228 295:39 299:72 300:251 301:253 302:215 303:29 315:36 316:253 317:253 318:203 319:13 328:82 329:253 330:253 331:170 343:36 344:253 345:253 346:164 356:11 357:198 358:253 359:184 360:6 371:36 372:253 373:253 374:82 385:138 386:253 387:253 388:35 399:128 400:253 401:253 402:47 413:48 414:253 415:253 416:35 427:154 428:253 429:253 430:47 441:48 442:253 443:253 444:35 455:102 456:253 457:253 458:99 469:48 470:253 471:253 472:35 483:36 484:253 485:253 486:164 496:16 497:208 498:253 499:211 500:17 511:32 512:244 513:253 514:175 515:4 524:44 525:253 526:253 527:156 540:171 541:253 542:253 543:29 551:30 552:217 553:253 554:188 555:19 568:171 569:253 570:253 571:59 578:60 579:217 580:253 581:253 582:70 596:78 597:253 598:253 599:231 600:48 604:26 605:128 606:249 607:253 608:244 609:94 610:15 624:8 625:151 626:253 627:253 628:234 629:101 630:121 631:219 632:229 633:253 634:253 635:201 636:80 653:38 654:232 655:253 656:253 657:253 658:253 659:253 660:253 661:253 662:201 663:66 +0 127:68 128:254 129:255 130:254 131:107 153:11 154:176 155:230 156:253 157:253 158:253 159:212 180:28 181:197 182:253 183:253 184:253 185:253 186:253 187:229 188:107 189:14 208:194 209:253 210:253 211:253 212:253 213:253 214:253 215:253 216:253 217:53 235:69 236:241 237:253 238:253 239:253 240:253 241:241 242:186 243:253 244:253 245:195 262:10 263:161 264:253 265:253 266:253 267:246 268:40 269:57 270:231 271:253 272:253 273:195 290:140 291:253 292:253 293:253 294:253 295:154 297:25 298:253 299:253 300:253 301:195 318:213 319:253 320:253 321:253 322:135 323:8 325:3 326:128 327:253 328:253 329:195 345:77 346:238 347:253 348:253 349:253 350:7 354:116 355:253 356:253 357:195 372:11 373:165 374:253 375:253 376:231 377:70 378:1 382:78 383:237 384:253 385:195 400:33 401:253 402:253 403:253 404:182 411:200 412:253 413:195 428:98 429:253 430:253 431:253 432:24 439:42 440:253 441:195 456:197 457:253 458:253 459:253 460:24 467:163 468:253 469:195 484:197 485:253 486:253 487:189 488:13 494:53 495:227 496:253 497:121 512:197 513:253 514:253 515:114 521:21 522:227 523:253 524:231 525:27 540:197 541:253 542:253 543:114 547:5 548:131 549:143 550:253 551:231 552:59 568:197 569:253 570:253 571:236 572:73 573:58 574:217 575:223 576:253 577:253 578:253 579:174 596:197 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:253 607:48 624:149 625:253 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:182 634:15 635:3 652:12 653:168 654:253 655:253 656:253 657:253 658:253 659:248 660:89 661:23 +1 157:85 158:255 159:103 160:1 185:205 186:253 187:253 188:30 213:205 214:253 215:253 216:30 240:44 241:233 242:253 243:244 244:27 268:135 269:253 270:253 271:100 296:153 297:253 298:240 299:76 323:12 324:208 325:253 326:166 351:69 352:253 353:253 354:142 378:14 379:110 380:253 381:235 382:33 406:63 407:223 408:235 409:130 434:186 435:253 436:235 437:37 461:17 462:145 463:253 464:231 465:35 489:69 490:220 491:231 492:123 516:18 517:205 518:253 519:176 520:27 543:17 544:125 545:253 546:185 547:39 571:71 572:214 573:231 574:41 599:167 600:253 601:225 602:33 626:72 627:205 628:207 629:14 653:30 654:249 655:233 656:49 681:32 682:253 683:89 +1 126:94 127:132 154:250 155:250 156:4 182:250 183:254 184:95 210:250 211:254 212:95 238:250 239:254 240:95 266:250 267:254 268:95 294:250 295:254 296:95 322:250 323:254 324:95 350:250 351:254 352:95 378:250 379:254 380:95 405:77 406:254 407:250 408:19 433:96 434:254 435:249 461:53 462:253 463:252 464:43 490:250 491:251 492:32 517:85 518:254 519:249 545:96 546:254 547:249 573:83 574:254 575:250 576:14 602:250 603:254 604:95 630:250 631:255 632:95 658:132 659:254 660:95 +1 124:32 125:253 126:31 152:32 153:251 154:149 180:32 181:251 182:188 208:32 209:251 210:188 236:32 237:251 238:228 239:59 264:32 265:253 266:253 267:95 292:28 293:236 294:251 295:114 321:127 322:251 323:251 349:127 350:251 351:251 377:48 378:232 379:251 406:223 407:253 408:159 434:221 435:251 436:158 462:142 463:251 464:158 490:64 491:251 492:242 493:55 518:64 519:251 520:253 521:161 546:64 547:253 548:255 549:221 574:16 575:181 576:253 577:220 603:79 604:253 605:236 606:63 632:213 633:251 634:126 660:96 661:251 662:126 +1 129:39 130:254 131:255 132:254 133:140 157:136 158:253 159:253 160:228 161:67 184:6 185:227 186:253 187:253 188:58 211:29 212:188 213:253 214:253 215:253 216:17 239:95 240:253 241:253 242:253 243:157 244:8 266:3 267:107 268:253 269:253 270:245 271:77 294:29 295:253 296:253 297:240 298:100 322:141 323:253 324:253 325:215 349:129 350:248 351:253 352:253 353:215 377:151 378:253 379:253 380:253 381:144 405:151 406:253 407:253 408:253 409:27 431:3 432:102 433:242 434:253 435:253 436:110 437:3 459:97 460:253 461:253 462:253 463:214 464:55 487:207 488:253 489:253 490:253 491:158 515:67 516:253 517:253 518:253 519:158 543:207 544:253 545:253 546:240 547:88 571:207 572:253 573:253 574:224 598:32 599:217 600:253 601:253 602:224 626:141 627:253 628:253 629:253 630:133 654:36 655:219 656:253 657:140 658:10 +0 123:59 124:55 149:71 150:192 151:254 152:250 153:147 154:17 176:123 177:247 178:253 179:254 180:253 181:253 182:196 183:79 184:176 185:175 186:175 187:124 188:48 203:87 204:247 205:247 206:176 207:95 208:102 209:117 210:243 211:237 212:192 213:232 214:253 215:253 216:245 217:152 218:6 230:23 231:229 232:253 233:138 238:219 239:58 241:95 242:118 243:80 244:230 245:254 246:196 247:30 258:120 259:254 260:205 261:8 266:114 272:38 273:255 274:254 275:155 276:5 286:156 287:253 288:92 301:61 302:235 303:253 304:102 314:224 315:253 316:78 330:117 331:253 332:196 333:18 342:254 343:253 344:78 358:9 359:211 360:253 361:73 370:254 371:253 372:78 387:175 388:253 389:155 398:194 399:254 400:101 415:79 416:254 417:155 426:112 427:253 428:211 429:9 443:73 444:251 445:200 454:41 455:241 456:253 457:87 471:25 472:240 473:253 483:147 484:253 485:227 486:47 499:94 500:253 501:200 511:5 512:193 513:253 514:230 515:76 527:175 528:253 529:155 540:31 541:219 542:254 543:255 544:126 545:18 553:14 554:149 555:254 556:244 557:45 569:21 570:158 571:254 572:253 573:226 574:162 575:118 576:96 577:20 578:20 579:73 580:118 581:224 582:253 583:247 584:85 598:30 599:155 600:253 601:253 602:253 603:253 604:254 605:253 606:253 607:253 608:253 609:254 610:247 611:84 627:5 628:27 629:117 630:206 631:244 632:229 633:213 634:213 635:213 636:176 637:117 638:32 659:45 660:23 +1 128:58 129:139 156:247 157:247 158:25 183:121 184:253 185:156 186:3 211:133 212:253 213:145 238:11 239:227 240:253 241:145 266:7 267:189 268:253 269:145 294:35 295:252 296:253 297:145 322:146 323:252 324:253 325:131 350:146 351:252 352:253 353:13 378:146 379:252 380:253 381:13 406:147 407:253 408:255 409:13 434:146 435:252 436:253 437:13 462:146 463:252 464:253 465:13 490:146 491:252 492:253 493:13 517:22 518:230 519:252 520:221 521:9 545:22 546:230 547:252 548:133 574:146 575:252 576:133 602:146 603:252 604:120 630:146 631:252 658:146 659:252 +1 129:28 130:247 131:255 132:165 156:47 157:221 158:252 159:252 160:164 184:177 185:252 186:252 187:252 188:164 212:177 213:252 214:252 215:223 216:78 240:177 241:252 242:252 243:197 267:114 268:236 269:252 270:235 271:42 294:5 295:148 296:252 297:252 298:230 321:14 322:135 323:252 324:252 325:252 326:230 349:78 350:252 351:252 352:252 353:252 354:162 377:78 378:252 379:252 380:252 381:252 382:9 405:78 406:252 407:252 408:252 409:252 410:9 432:32 433:200 434:252 435:252 436:252 437:105 438:3 459:10 460:218 461:252 462:252 463:252 464:105 465:8 487:225 488:252 489:252 490:252 491:240 492:69 514:44 515:237 516:252 517:252 518:228 519:85 541:59 542:218 543:252 544:252 545:225 546:93 568:65 569:208 570:252 571:252 572:252 573:175 596:133 597:252 598:252 599:252 600:225 601:68 624:133 625:252 626:252 627:244 628:54 652:133 653:252 654:252 655:48 +0 156:13 157:6 181:10 182:77 183:145 184:253 185:190 186:67 207:11 208:77 209:193 210:252 211:252 212:253 213:252 214:238 215:157 216:71 217:26 233:10 234:78 235:193 236:252 237:252 238:252 239:252 240:253 241:252 242:252 243:252 244:252 245:228 246:128 247:49 248:5 259:6 260:78 261:194 262:252 263:252 264:252 265:252 266:252 267:252 268:253 269:217 270:192 271:232 272:252 273:252 274:252 275:252 276:135 277:3 286:4 287:147 288:252 289:252 290:252 291:252 292:252 293:252 294:252 295:252 296:175 297:26 299:40 300:145 301:235 302:252 303:252 304:252 305:104 314:208 315:252 316:252 317:252 318:252 319:252 320:252 321:133 322:48 323:48 329:71 330:236 331:252 332:252 333:230 342:253 343:185 344:170 345:252 346:252 347:252 348:173 349:22 358:102 359:252 360:252 361:252 370:24 371:141 372:243 373:252 374:252 375:186 376:5 386:8 387:220 388:252 389:252 398:70 399:247 400:252 401:252 402:165 403:37 414:81 415:251 416:252 417:194 426:255 427:253 428:253 429:251 430:69 441:39 442:231 443:253 444:253 445:127 454:253 455:252 456:249 457:127 468:6 469:147 470:252 471:252 472:190 473:5 482:253 483:252 484:216 495:7 496:145 497:252 498:252 499:252 500:69 510:253 511:252 512:223 513:16 522:25 523:185 524:252 525:252 526:252 527:107 528:8 538:167 539:252 540:252 541:181 542:18 549:105 550:191 551:252 552:252 553:235 554:151 555:10 566:37 567:221 568:252 569:252 570:210 571:193 572:96 573:73 574:130 575:188 576:194 577:227 578:252 579:252 580:235 581:128 595:97 596:220 597:252 598:252 599:252 600:252 601:252 602:252 603:252 604:253 605:252 606:252 607:236 608:70 624:40 625:174 626:252 627:252 628:252 629:252 630:252 631:252 632:253 633:197 634:138 635:29 653:5 654:23 655:116 656:143 657:143 658:143 659:143 660:24 661:10 +0 127:28 128:164 129:254 130:233 131:148 132:11 154:3 155:164 156:254 157:234 158:225 159:254 160:204 182:91 183:254 184:235 185:48 186:32 187:166 188:251 189:92 208:33 209:111 210:214 211:205 212:49 215:24 216:216 217:210 235:34 236:217 237:254 238:254 239:211 244:87 245:237 246:43 262:34 263:216 264:254 265:254 266:252 267:243 268:61 272:38 273:248 274:182 290:171 291:254 292:184 293:205 294:175 295:36 301:171 302:227 317:28 318:234 319:190 320:13 321:193 322:157 329:124 330:238 331:26 345:140 346:254 347:131 349:129 350:157 357:124 358:254 359:95 373:201 374:238 375:56 377:70 378:103 385:124 386:254 387:148 400:62 401:255 402:210 413:150 414:254 415:122 428:86 429:254 430:201 431:15 440:28 441:237 442:246 443:44 456:128 457:254 458:143 468:34 469:243 470:227 484:62 485:254 486:210 496:58 497:249 498:179 512:30 513:240 514:210 524:207 525:254 526:64 541:216 542:231 543:34 551:129 552:248 553:170 554:9 569:131 570:254 571:170 577:17 578:129 579:248 580:225 581:24 597:50 598:245 599:245 600:184 601:106 602:106 603:106 604:133 605:231 606:254 607:244 608:53 626:67 627:249 628:254 629:254 630:254 631:254 632:254 633:251 634:193 635:40 655:38 656:157 657:248 658:166 659:166 660:139 661:57 +0 129:105 130:255 131:219 132:67 133:67 134:52 156:20 157:181 158:253 159:253 160:253 161:253 162:226 163:69 182:4 183:129 184:206 185:253 186:253 187:253 188:253 189:253 190:253 191:130 209:9 210:141 211:253 212:253 213:253 214:253 215:253 216:253 217:253 218:253 219:166 220:20 237:134 238:253 239:253 240:253 241:253 242:253 243:253 244:253 245:253 246:253 247:253 248:65 262:2 263:83 264:207 265:246 266:253 267:253 268:253 269:253 270:253 271:249 272:234 273:247 274:253 275:253 276:65 290:83 291:253 292:253 293:253 294:253 295:253 296:189 297:253 298:253 299:205 301:179 302:253 303:253 304:65 317:85 318:234 319:253 320:253 321:253 322:253 323:157 324:26 325:164 326:151 327:83 329:179 330:253 331:253 332:65 344:65 345:237 346:253 347:253 348:253 349:67 350:36 351:14 353:15 354:12 357:179 358:253 359:253 360:65 371:4 372:141 373:253 374:253 375:221 376:158 377:23 385:179 386:253 387:253 388:65 399:129 400:253 401:253 402:241 403:62 412:72 413:226 414:253 415:175 416:24 426:119 427:247 428:253 429:253 430:206 439:8 440:134 441:253 442:253 443:130 454:132 455:253 456:253 457:194 458:27 467:125 468:253 469:253 470:253 471:130 481:45 482:213 483:253 484:253 485:112 493:70 494:170 495:247 496:253 497:253 498:89 499:43 509:67 510:253 511:253 512:196 513:55 514:9 520:8 521:131 522:253 523:253 524:253 525:86 526:1 537:67 538:253 539:253 540:253 541:253 542:129 546:43 547:114 548:134 549:253 550:253 551:231 552:139 553:41 565:20 566:167 567:253 568:253 569:253 570:247 571:179 572:179 573:179 574:206 575:253 576:253 577:253 578:253 579:72 594:103 595:240 596:253 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:244 605:119 606:8 607:1 623:107 624:253 625:253 626:253 627:253 628:253 629:253 630:253 631:175 632:111 651:3 652:121 653:253 654:253 655:253 656:253 657:253 658:182 659:24 +0 125:22 126:183 127:252 128:254 129:252 130:252 131:252 132:76 151:85 152:85 153:168 154:250 155:250 156:252 157:250 158:250 159:250 160:250 161:71 163:43 164:85 165:14 178:107 179:252 180:250 181:250 182:250 183:250 184:252 185:250 186:250 187:250 188:250 189:210 191:127 192:250 193:146 205:114 206:237 207:252 208:250 209:250 210:250 211:250 212:252 213:250 214:250 215:250 216:250 217:210 219:127 220:250 221:250 232:107 233:237 234:250 235:252 236:250 237:250 238:250 239:74 240:41 241:41 242:41 243:41 244:217 245:34 247:127 248:250 249:250 259:15 260:148 261:252 262:252 263:254 264:238 265:105 275:128 276:252 277:252 286:15 287:140 288:250 289:250 290:250 291:167 292:111 303:127 304:250 305:250 314:43 315:250 316:250 317:250 318:250 331:127 332:250 333:250 342:183 343:250 344:250 345:250 346:110 358:57 359:210 360:250 361:250 370:252 371:250 372:250 373:110 374:7 386:85 387:250 388:250 389:250 398:254 399:252 400:252 401:83 414:86 415:252 416:252 417:217 426:252 427:250 428:250 429:138 430:14 441:15 442:140 443:250 444:250 445:41 454:252 455:250 456:250 457:250 458:41 469:43 470:250 471:250 472:250 473:41 482:252 483:250 484:250 485:250 486:181 497:183 498:250 499:250 500:250 501:41 510:76 511:250 512:250 513:250 514:250 524:177 525:252 526:250 527:250 528:110 529:7 538:36 539:224 540:252 541:252 542:252 543:219 544:43 545:43 546:43 547:7 549:15 550:43 551:183 552:252 553:255 554:252 555:126 567:85 568:250 569:250 570:250 571:252 572:250 573:250 574:250 575:111 576:86 577:140 578:250 579:250 580:250 581:252 582:222 583:83 595:42 596:188 597:250 598:250 599:252 600:250 601:250 602:250 603:250 604:252 605:250 606:250 607:250 608:250 609:126 610:83 624:127 625:250 626:250 627:252 628:250 629:250 630:250 631:250 632:252 633:250 634:250 635:137 636:83 652:21 653:41 654:217 655:252 656:250 657:250 658:250 659:250 660:217 661:41 662:41 663:14 +1 155:114 156:206 157:25 183:238 184:252 185:55 211:222 212:252 213:55 239:113 240:252 241:55 267:113 268:252 269:55 295:255 296:253 297:56 323:253 324:176 325:6 350:32 351:253 352:233 353:43 378:140 379:253 380:195 381:19 406:140 407:253 408:167 433:29 434:253 435:141 461:29 462:252 463:140 489:29 490:252 491:140 517:29 518:252 519:140 545:29 546:252 547:140 573:169 574:253 575:79 601:169 602:252 628:76 629:234 630:141 656:197 657:233 658:37 684:197 685:223 +1 127:73 128:253 129:253 130:63 155:115 156:252 157:252 158:144 183:217 184:252 185:252 186:144 210:63 211:237 212:252 213:252 214:144 238:109 239:252 240:252 241:252 266:109 267:252 268:252 269:252 294:109 295:252 296:252 297:252 322:191 323:252 324:252 325:252 349:145 350:255 351:253 352:253 353:253 376:32 377:237 378:253 379:252 380:252 381:210 404:37 405:252 406:253 407:252 408:252 409:108 432:37 433:252 434:253 435:252 436:252 437:108 460:21 461:207 462:255 463:253 464:253 465:108 489:144 490:253 491:252 492:252 493:108 516:27 517:221 518:253 519:252 520:252 521:108 544:16 545:190 546:253 547:252 548:252 549:108 573:145 574:255 575:253 576:253 577:253 601:144 602:253 603:252 604:252 605:210 629:144 630:253 631:252 632:252 633:108 657:62 658:253 659:252 660:252 661:108 +1 120:85 121:253 122:132 123:9 147:82 148:241 149:251 150:251 151:128 175:175 176:251 177:251 178:251 179:245 180:121 203:13 204:204 205:251 206:251 207:251 208:245 209:107 232:39 233:251 234:251 235:251 236:251 237:167 238:22 260:15 261:155 262:251 263:251 264:251 265:251 266:177 289:15 290:157 291:248 292:251 293:251 294:251 295:165 319:214 320:251 321:251 322:251 323:212 324:78 325:24 347:109 348:251 349:251 350:251 351:253 352:251 353:170 354:10 375:5 376:57 377:162 378:251 379:253 380:251 381:251 382:18 405:106 406:239 407:255 408:253 409:253 410:213 434:105 435:253 436:251 437:251 438:230 439:72 463:253 464:251 465:251 466:251 467:221 468:67 491:72 492:251 493:251 494:251 495:251 496:96 519:36 520:199 521:251 522:251 523:251 524:155 525:15 548:45 549:204 550:251 551:251 552:251 553:157 577:161 578:249 579:251 580:251 581:248 582:147 606:233 607:251 608:251 609:251 610:173 634:233 635:251 636:251 637:251 638:173 662:53 663:131 664:251 665:251 666:173 +1 126:15 127:200 128:255 129:90 154:42 155:254 156:254 157:173 182:42 183:254 184:254 185:199 210:26 211:237 212:254 213:221 214:12 239:213 240:254 241:231 242:17 267:213 268:254 269:199 295:213 296:254 297:199 323:213 324:254 325:96 350:20 351:232 352:254 353:33 378:84 379:254 380:229 381:17 406:168 407:254 408:203 433:8 434:217 435:254 436:187 461:84 462:254 463:254 464:48 489:195 490:254 491:254 492:37 516:20 517:233 518:254 519:212 520:4 544:132 545:254 546:254 547:82 571:9 572:215 573:254 574:254 575:116 576:46 599:55 600:254 601:254 602:254 603:254 604:121 627:113 628:254 629:254 630:254 631:254 632:40 655:12 656:163 657:254 658:185 659:58 660:1 +0 182:32 183:57 184:57 185:57 186:57 187:57 188:57 189:57 208:67 209:185 210:229 211:252 212:252 213:252 214:253 215:252 216:252 217:252 218:185 219:66 234:13 235:188 236:246 237:252 238:253 239:252 240:252 241:252 242:241 243:139 244:177 245:252 246:253 247:246 248:187 249:13 261:26 262:255 263:253 264:244 265:175 266:101 274:126 275:244 276:253 277:153 288:82 289:243 290:253 291:214 292:81 303:169 304:252 305:252 315:19 316:215 317:252 318:206 319:56 331:169 332:252 333:252 343:157 344:252 345:252 346:13 359:169 360:252 361:151 370:41 371:253 372:253 373:128 386:92 387:253 388:206 389:13 398:166 399:252 400:196 401:9 414:216 415:252 416:142 426:253 427:252 428:168 441:89 442:253 443:208 444:13 454:253 455:252 456:68 468:38 469:225 470:253 471:96 482:254 483:253 484:56 495:45 496:229 497:253 498:151 510:253 511:252 512:81 522:70 523:225 524:252 525:227 538:216 539:252 540:168 548:29 549:134 550:253 551:252 552:186 553:31 566:91 567:252 568:243 569:125 573:51 574:114 575:113 576:210 577:252 578:253 579:151 580:19 595:157 596:253 597:253 598:254 599:253 600:253 601:253 602:254 603:253 604:244 605:175 606:51 623:19 624:122 625:196 626:197 627:221 628:196 629:196 630:197 631:121 632:56 655:25 +0 127:42 128:235 129:255 130:84 153:15 154:132 155:208 156:253 157:253 158:171 159:108 180:6 181:177 182:253 183:253 184:253 185:253 186:253 187:242 188:110 208:151 209:253 210:253 211:253 212:253 213:253 214:253 215:253 216:139 235:48 236:208 237:253 238:253 239:253 240:253 241:253 242:253 243:253 244:139 263:85 264:253 265:253 266:253 267:253 268:236 269:156 270:184 271:253 272:148 273:6 290:7 291:141 292:253 293:253 294:253 295:253 296:27 298:170 299:253 300:253 301:74 318:19 319:253 320:253 321:253 322:253 323:253 324:27 326:170 327:253 328:253 329:74 345:16 346:186 347:253 348:253 349:253 350:242 351:105 352:4 354:170 355:253 356:253 357:94 358:1 373:141 374:253 375:253 376:253 377:242 378:100 382:170 383:253 384:253 385:253 386:8 401:141 402:253 403:253 404:253 405:224 410:170 411:253 412:253 413:253 414:8 428:12 429:158 430:253 431:253 432:230 433:51 438:18 439:237 440:253 441:253 442:8 456:76 457:253 458:253 459:218 460:61 467:236 468:253 469:253 470:8 484:76 485:253 486:253 487:168 495:110 496:253 497:132 498:3 512:76 513:253 514:253 515:168 521:20 522:174 523:239 524:147 525:5 539:5 540:155 541:253 542:253 543:168 548:102 549:170 550:253 551:253 552:139 567:3 568:128 569:253 570:253 571:228 572:179 573:179 574:179 575:179 576:245 577:253 578:253 579:219 580:41 596:76 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:253 605:253 606:253 607:163 624:39 625:199 626:253 627:253 628:253 629:253 630:253 631:253 632:253 633:253 634:170 635:9 653:36 654:219 655:253 656:253 657:253 658:253 659:253 660:224 661:65 662:22 +1 156:202 157:253 158:69 184:253 185:252 186:121 212:253 213:252 214:69 240:253 241:252 242:69 267:106 268:253 269:231 270:37 295:179 296:255 297:196 322:17 323:234 324:253 325:92 350:93 351:252 352:253 353:92 378:93 379:252 380:253 381:92 406:93 407:252 408:232 409:8 434:208 435:253 436:116 462:207 463:252 464:116 490:207 491:252 492:32 517:57 518:244 519:252 545:122 546:252 547:252 573:185 574:253 575:253 601:184 602:252 603:252 629:101 630:252 631:252 657:13 658:173 659:252 660:43 686:9 687:232 688:116 +1 156:73 157:253 158:253 159:253 160:124 184:73 185:251 186:251 187:251 188:251 212:99 213:251 214:251 215:251 216:225 240:253 241:251 242:251 243:251 244:71 266:79 267:180 268:253 269:251 270:251 271:173 272:20 294:110 295:253 296:255 297:253 298:216 322:109 323:251 324:253 325:251 326:215 350:109 351:251 352:253 353:251 354:215 378:109 379:251 380:253 381:251 382:137 406:109 407:251 408:253 409:251 410:35 433:37 434:253 435:253 436:255 437:253 438:35 461:140 462:251 463:251 464:253 465:168 466:15 488:125 489:246 490:251 491:251 492:190 493:15 516:144 517:251 518:251 519:251 520:180 543:53 544:221 545:251 546:251 547:251 548:51 571:125 572:253 573:253 574:253 575:201 598:105 599:253 600:251 601:251 602:188 603:30 626:180 627:253 628:251 629:251 630:142 654:180 655:253 656:251 657:235 658:82 682:180 683:253 684:251 685:215 +1 124:111 125:255 126:48 152:162 153:253 154:237 155:63 180:206 181:253 182:253 183:183 208:87 209:217 210:253 211:205 237:90 238:253 239:238 240:60 265:37 266:225 267:253 268:89 294:206 295:253 296:159 322:206 323:253 324:226 350:206 351:253 352:226 378:206 379:253 380:226 406:206 407:253 408:226 434:206 435:253 436:226 462:206 463:253 464:226 490:206 491:253 492:226 518:206 519:253 520:237 521:45 546:206 547:253 548:253 549:109 574:173 575:253 576:253 577:109 602:69 603:253 604:253 605:109 630:64 631:248 632:253 633:109 659:112 660:253 661:109 +0 99:70 100:255 101:165 102:114 127:122 128:253 129:253 130:253 131:120 155:165 156:253 157:253 158:253 159:234 160:52 183:99 184:253 185:253 186:253 187:253 188:228 189:26 209:60 210:168 211:238 212:202 213:174 214:253 215:253 216:253 217:127 235:91 236:81 237:1 238:215 239:128 240:28 241:12 242:181 243:253 244:253 245:175 246:3 262:18 263:204 264:253 265:77 270:7 271:253 272:253 273:253 274:54 289:54 290:248 291:253 292:253 293:143 298:1 299:127 300:253 301:253 302:188 317:104 318:253 319:253 320:253 321:20 327:81 328:249 329:253 330:191 345:192 346:253 347:253 348:218 349:5 356:203 357:253 358:208 359:21 372:56 373:237 374:253 375:250 376:100 384:104 385:253 386:253 387:75 400:76 401:253 402:253 403:224 412:119 413:253 414:253 415:75 428:80 429:253 430:253 431:103 439:4 440:241 441:253 442:218 443:32 456:213 457:253 458:253 459:103 467:125 468:253 469:253 470:191 484:213 485:253 486:253 487:103 494:3 495:176 496:253 497:253 498:135 512:213 513:253 514:253 515:103 521:9 522:162 523:253 524:253 525:226 526:37 540:179 541:253 542:253 543:135 548:46 549:157 550:253 551:253 552:253 553:63 568:23 569:188 570:253 571:249 572:179 573:179 574:179 575:179 576:233 577:253 578:253 579:233 580:156 581:10 597:51 598:235 599:253 600:253 601:253 602:253 603:253 604:253 605:251 606:232 607:120 626:16 627:124 628:253 629:253 630:253 631:253 632:152 633:104 +1 124:29 125:197 126:255 127:84 152:85 153:251 154:253 155:83 180:86 181:253 182:254 183:253 208:85 209:251 210:253 211:251 236:86 237:253 238:254 239:253 240:169 264:85 265:251 266:253 267:251 268:168 292:86 293:253 294:254 295:253 296:169 320:28 321:196 322:253 323:251 324:168 349:169 350:254 351:253 352:169 377:168 378:253 379:251 380:168 405:169 406:254 407:253 408:169 433:168 434:253 435:251 436:168 462:254 463:253 464:254 465:139 490:253 491:251 492:253 493:251 518:254 519:253 520:254 521:253 522:57 546:253 547:251 548:253 549:251 550:168 574:198 575:253 576:254 577:253 578:114 602:85 603:251 604:253 605:251 630:85 631:253 632:254 633:253 658:28 659:83 660:196 661:83 +1 159:31 160:210 161:253 162:163 187:198 188:252 189:252 190:162 213:10 214:86 215:242 216:252 217:252 218:66 241:164 242:252 243:252 244:252 245:188 246:8 268:53 269:242 270:252 271:252 272:225 273:14 296:78 297:252 298:252 299:252 300:204 323:56 324:231 325:252 326:252 327:212 328:35 351:157 352:252 353:252 354:252 355:37 377:8 378:132 379:253 380:252 381:252 382:230 383:24 405:45 406:252 407:253 408:252 409:154 410:55 427:7 428:55 433:107 434:253 435:255 436:228 437:53 454:15 455:24 456:23 460:110 461:242 462:252 463:228 464:59 482:57 483:83 487:88 488:247 489:252 490:252 491:140 514:15 515:189 516:252 517:252 518:252 542:74 543:252 544:252 545:238 546:90 570:178 571:252 572:252 573:189 597:40 598:217 599:252 600:252 601:59 625:75 626:252 627:252 628:252 629:85 630:61 653:62 654:239 655:252 656:156 657:14 682:178 683:252 684:14 +1 131:159 132:255 133:122 158:167 159:228 160:253 161:121 185:64 186:236 187:251 188:205 189:110 212:48 213:158 214:251 215:251 216:178 217:39 240:190 241:251 242:251 243:251 267:96 268:253 269:253 270:253 271:153 295:194 296:251 297:251 298:211 299:74 322:80 323:174 324:251 325:251 326:140 327:47 349:16 350:181 351:253 352:251 353:219 354:23 377:64 378:251 379:253 380:251 381:204 382:19 405:223 406:253 407:255 408:233 409:48 431:20 432:174 433:244 434:251 435:253 436:109 437:31 459:96 460:189 461:251 462:251 463:126 464:31 486:24 487:106 488:251 489:235 490:188 491:100 514:96 515:251 516:251 517:228 518:59 542:255 543:253 544:253 545:213 546:36 569:100 570:253 571:251 572:251 573:85 574:23 596:32 597:127 598:253 599:235 600:126 601:15 624:104 625:251 626:253 627:240 628:79 652:83 653:193 654:253 655:220 +0 153:92 154:191 155:178 156:253 157:242 158:141 159:104 160:29 180:26 181:253 182:252 183:252 184:252 185:253 186:252 187:252 188:252 189:108 190:19 206:57 207:123 208:222 209:253 210:252 211:252 212:252 213:168 214:224 215:252 216:252 217:253 218:84 233:176 234:243 235:252 236:252 237:253 238:252 239:252 240:252 242:19 243:153 244:252 245:253 246:209 247:25 259:10 260:128 261:255 262:253 263:244 264:225 265:114 266:194 267:253 268:178 272:163 273:254 274:253 275:168 287:85 288:252 289:253 290:189 291:56 294:19 295:133 296:9 300:38 301:253 302:252 303:168 314:19 315:191 316:252 317:194 318:19 329:253 330:252 331:234 332:22 342:107 343:252 344:252 345:13 357:253 358:252 359:252 360:128 370:169 371:253 372:241 385:141 386:253 387:253 388:140 397:19 398:225 399:252 400:139 413:66 414:252 415:252 416:139 425:29 426:252 427:252 428:52 441:29 442:252 443:252 444:139 453:29 454:252 455:252 456:28 469:29 470:252 471:252 472:40 481:141 482:253 483:253 484:91 497:154 498:253 499:168 509:66 510:252 511:252 512:165 525:253 526:252 527:168 537:19 538:224 539:252 540:252 552:126 553:253 554:252 555:80 566:169 567:252 568:252 569:214 570:38 579:126 580:249 581:253 582:151 583:6 594:26 595:223 596:253 597:254 598:253 599:128 600:29 604:13 605:41 606:216 607:253 608:253 609:226 610:38 623:122 624:252 625:253 626:252 627:252 628:252 629:169 630:169 631:169 632:206 633:253 634:252 635:252 636:202 637:38 651:19 652:56 653:168 654:224 655:252 656:252 657:253 658:252 659:252 660:252 661:253 662:233 663:130 664:6 682:94 683:139 684:190 685:153 686:252 687:164 688:139 689:28 690:22 +1 128:53 129:250 130:255 131:25 156:167 157:253 158:253 159:25 182:3 183:123 184:247 185:253 186:253 187:25 210:9 211:253 212:253 213:253 214:253 215:25 238:9 239:253 240:253 241:253 242:253 243:25 266:9 267:253 268:253 269:253 270:180 271:13 294:9 295:253 296:253 297:253 298:104 322:9 323:253 324:253 325:253 326:104 350:15 351:253 352:253 353:253 354:104 378:184 379:253 380:253 381:228 382:68 406:184 407:253 408:253 409:182 433:103 434:251 435:253 436:253 437:12 461:106 462:253 463:253 464:253 465:8 488:24 489:238 490:253 491:253 492:253 493:8 516:27 517:253 518:253 519:253 520:253 521:8 544:27 545:253 546:253 547:253 548:253 549:8 572:27 573:253 574:253 575:253 576:177 577:4 600:160 601:253 602:253 603:253 604:87 628:202 629:253 630:253 631:219 632:54 656:81 657:253 658:247 659:51 +0 122:63 123:176 124:253 125:253 126:159 127:113 128:63 150:140 151:253 152:252 153:252 154:252 155:252 156:241 157:100 158:66 177:54 178:227 179:253 180:252 181:252 182:252 183:252 184:253 185:252 186:239 187:181 188:57 204:38 205:224 206:252 207:253 208:226 209:246 210:252 211:252 212:253 213:252 214:252 215:252 216:252 217:108 218:3 232:57 233:252 234:252 235:253 236:27 237:88 238:112 239:112 240:112 241:112 242:142 243:252 244:252 245:253 246:152 247:31 260:198 261:253 262:253 263:79 270:32 271:153 272:253 273:255 274:253 275:196 287:76 288:246 289:252 290:127 299:3 300:106 301:253 302:252 303:214 304:28 315:194 316:252 317:252 318:112 329:143 330:252 331:252 332:193 343:225 344:252 345:217 346:37 357:38 358:234 359:252 360:223 370:63 371:240 372:252 373:84 386:146 387:252 388:223 398:114 399:253 400:228 401:47 414:147 415:253 416:253 417:112 426:159 427:252 428:195 442:225 443:252 444:252 445:112 454:253 455:252 456:195 470:225 471:252 472:230 473:25 482:159 483:252 484:202 485:10 497:92 498:243 499:252 500:208 510:113 511:252 512:252 513:161 524:79 525:253 526:252 527:220 528:37 538:114 539:253 540:253 541:253 542:174 543:63 550:26 551:128 552:253 553:255 554:253 555:133 566:12 567:228 568:252 569:252 570:252 571:241 572:100 573:85 574:76 576:85 577:131 578:231 579:252 580:252 581:253 582:129 583:6 595:97 596:208 597:252 598:252 599:253 600:252 601:252 602:246 603:197 604:253 605:252 606:252 607:252 608:220 609:133 610:6 624:19 625:99 626:239 627:253 628:252 629:252 630:252 631:252 632:253 633:252 634:245 635:223 636:99 654:63 655:112 656:112 657:221 658:252 659:252 660:253 661:127 662:87 +0 153:12 154:136 155:254 156:255 157:195 158:115 159:3 180:6 181:175 182:253 183:196 184:160 185:252 186:253 187:15 208:130 209:253 210:234 211:4 213:27 214:205 215:232 216:40 235:54 236:246 237:253 238:68 242:24 243:243 244:106 262:3 263:134 264:235 265:99 266:4 271:132 272:247 273:77 290:56 291:253 292:62 299:23 300:233 301:129 318:179 319:183 320:4 328:182 329:220 345:21 346:232 347:59 356:95 357:232 358:21 373:128 374:183 385:228 386:85 401:187 402:124 413:228 414:186 429:187 430:124 441:228 442:104 457:187 458:124 469:169 470:184 485:187 486:124 497:203 498:150 513:187 514:124 524:10 525:220 526:39 541:187 542:155 552:111 553:201 569:129 570:228 571:7 579:12 580:181 581:76 598:234 599:166 600:9 606:24 607:209 608:106 626:139 627:250 628:167 629:11 630:2 631:11 632:11 633:129 634:227 635:90 636:11 655:95 656:247 657:253 658:178 659:253 660:253 661:244 662:86 684:47 685:175 686:253 687:232 688:149 689:40 +1 128:255 129:253 130:57 156:253 157:251 158:225 159:56 183:169 184:254 185:253 186:254 187:84 211:168 212:253 213:251 214:253 215:83 238:85 239:253 240:254 241:253 242:169 266:85 267:251 268:253 269:251 270:56 294:141 295:253 296:254 297:253 322:253 323:251 324:253 325:251 350:254 351:253 352:254 353:253 378:253 379:251 380:253 381:251 406:254 407:253 408:254 409:196 433:114 434:253 435:251 436:253 437:83 461:169 462:254 463:253 464:226 465:56 489:168 490:253 491:251 492:168 516:85 517:253 518:254 519:253 544:85 545:251 546:253 547:251 572:254 573:253 574:254 575:253 600:253 601:251 602:253 603:251 628:254 629:253 630:254 631:253 656:139 657:251 658:253 659:138 +0 151:23 152:167 153:208 154:254 155:255 156:129 157:19 179:151 180:253 181:253 182:253 183:253 184:253 185:209 186:26 207:181 208:253 209:253 210:253 211:227 212:181 213:253 214:207 215:22 235:227 236:253 237:253 238:253 239:92 240:38 241:226 242:253 243:129 244:2 263:193 264:253 265:253 266:248 267:62 269:50 270:253 271:253 272:45 291:170 292:253 293:253 294:135 297:12 298:208 299:253 300:119 318:16 319:232 320:253 321:253 322:21 326:60 327:253 328:185 346:164 347:253 348:253 349:224 350:14 354:14 355:217 356:247 357:62 373:3 374:193 375:253 376:250 377:64 383:199 384:253 385:179 401:67 402:253 403:253 404:205 411:98 412:253 413:188 429:151 430:253 431:245 432:43 439:63 440:250 441:188 457:151 458:253 459:243 468:244 469:222 470:22 485:151 486:253 487:217 496:244 497:253 498:115 512:3 513:195 514:253 515:134 524:156 525:253 526:150 541:140 542:253 543:134 552:239 553:253 554:139 569:44 570:253 571:134 579:53 580:246 581:237 582:32 597:8 598:200 599:229 600:40 606:25 607:225 608:253 609:188 626:120 627:250 628:230 629:58 630:17 632:12 633:42 634:213 635:253 636:238 637:84 655:151 656:253 657:253 658:217 659:179 660:206 661:253 662:253 663:196 664:118 683:18 684:58 685:145 686:152 687:253 688:214 689:145 690:74 691:7 +1 130:24 131:150 132:233 133:38 156:14 157:89 158:253 159:254 160:254 161:71 183:78 184:203 185:254 186:254 187:254 188:232 189:77 190:54 191:8 209:12 210:155 211:240 212:254 213:223 214:76 215:254 216:254 217:254 218:254 219:68 235:3 236:101 237:216 238:254 239:227 240:122 241:26 242:110 243:254 244:254 245:254 246:184 247:100 262:46 263:222 264:254 265:254 266:179 267:48 270:181 271:254 272:254 273:146 274:6 288:2 289:145 290:248 291:254 292:182 293:111 294:4 297:3 298:250 299:254 300:206 301:3 315:6 316:144 317:254 318:254 319:171 325:125 326:254 327:252 328:80 342:6 343:142 344:254 345:179 346:95 347:4 352:61 353:246 354:254 355:150 370:64 371:254 372:177 373:14 380:124 381:254 382:246 383:32 398:108 399:97 400:15 407:24 408:226 409:254 410:116 435:177 436:255 437:254 438:5 463:196 464:254 465:99 466:1 490:3 491:199 492:254 493:79 518:129 519:254 520:254 521:23 546:178 547:254 548:192 549:8 550:3 551:43 573:11 574:198 575:254 576:128 577:66 578:130 579:225 595:137 596:202 597:106 598:84 599:84 600:84 601:112 602:254 603:254 604:254 605:254 606:212 607:151 623:172 624:254 625:254 626:254 627:254 628:254 629:254 630:254 631:254 632:254 633:162 634:75 651:12 652:106 653:177 654:254 655:254 656:254 657:235 658:135 659:100 660:17 661:2 +0 125:120 126:253 127:253 128:63 151:38 152:131 153:246 154:252 155:252 156:203 157:15 179:222 180:252 181:252 182:252 183:252 184:166 185:38 205:4 206:107 207:253 208:252 209:252 210:252 211:252 212:253 213:224 214:137 215:26 233:107 234:252 235:253 236:252 237:220 238:128 239:252 240:253 241:252 242:252 243:239 244:140 261:170 262:253 263:255 264:168 267:79 268:192 269:253 270:253 271:253 272:253 273:255 274:90 288:51 289:243 290:252 291:215 292:33 296:12 297:74 298:233 299:252 300:252 301:253 302:195 303:19 316:166 317:252 318:252 319:31 326:43 327:149 328:195 329:253 330:252 331:177 332:19 343:57 344:234 345:252 346:252 357:237 358:252 359:252 360:180 361:13 371:85 372:252 373:252 374:173 385:50 386:237 387:252 388:252 389:112 399:226 400:253 401:240 402:63 414:163 415:253 416:253 417:112 426:38 427:234 428:252 429:176 442:85 443:252 444:252 445:158 454:113 455:252 456:252 457:84 470:19 471:209 472:252 473:252 482:207 483:252 484:252 485:84 498:10 499:203 500:252 501:236 510:253 511:252 512:252 513:84 526:85 527:252 528:252 529:112 538:114 539:253 540:253 541:146 553:51 554:159 555:253 556:240 557:63 566:75 567:243 568:252 569:249 570:146 579:57 580:85 581:238 582:252 583:252 584:99 595:116 596:252 597:252 598:252 599:198 600:197 601:165 602:57 603:57 604:57 605:182 606:197 607:234 608:252 609:253 610:233 611:164 612:19 623:28 624:84 625:180 626:252 627:253 628:252 629:252 630:252 631:252 632:253 633:252 634:252 635:252 636:252 637:225 638:71 653:13 654:112 655:253 656:252 657:252 658:252 659:252 660:253 661:252 662:252 663:157 664:112 +1 127:155 128:253 129:126 155:253 156:251 157:141 158:4 183:253 184:251 185:251 186:31 211:253 212:251 213:251 214:31 239:253 240:251 241:251 242:31 267:255 268:253 269:253 270:31 293:8 294:131 295:253 296:251 297:235 298:27 321:64 322:251 323:253 324:251 325:126 349:64 350:251 351:253 352:251 353:126 377:64 378:251 379:253 380:251 381:126 405:64 406:253 407:255 408:221 433:182 434:251 435:253 436:200 460:64 461:236 462:251 463:253 464:62 487:8 488:158 489:251 490:251 491:169 492:8 515:32 516:251 517:251 518:251 519:158 543:32 544:253 545:253 546:253 547:159 571:32 572:251 573:251 574:251 575:39 599:32 600:251 601:251 602:251 627:32 628:251 629:251 630:251 631:100 655:32 656:251 657:251 658:251 +0 101:88 102:127 103:5 126:19 127:58 128:20 129:14 130:217 131:19 152:7 153:146 154:247 155:253 156:235 157:27 158:84 159:81 180:126 181:253 182:164 183:19 184:15 187:156 188:9 208:214 209:222 210:34 215:234 216:58 235:59 236:254 237:116 243:235 244:58 263:141 264:251 265:72 271:151 272:140 291:224 292:233 299:136 300:223 319:254 320:218 327:136 328:253 347:254 348:135 355:136 356:253 374:23 375:255 376:114 383:137 384:231 402:98 403:254 404:122 411:136 412:155 430:98 431:254 432:106 439:166 440:155 458:98 459:254 460:128 467:234 468:193 486:98 487:254 488:135 494:61 495:248 496:118 515:255 516:238 517:18 521:13 522:224 523:254 524:58 543:201 544:253 545:128 546:2 548:5 549:150 550:253 551:167 552:9 571:18 572:226 573:253 574:49 575:31 576:156 577:253 578:228 579:13 600:147 601:253 602:243 603:241 604:254 605:227 606:43 628:5 629:126 630:245 631:253 632:231 633:46 +0 127:37 128:141 129:156 130:156 131:194 132:194 133:47 153:11 154:132 155:239 156:253 157:253 158:253 159:253 160:254 161:181 180:25 181:172 182:253 183:235 184:167 185:78 186:93 187:174 188:254 189:247 190:54 207:26 208:210 209:253 210:237 211:90 216:201 217:253 218:78 235:192 236:253 237:237 238:58 244:156 245:253 246:78 262:141 263:254 264:235 265:53 269:19 270:5 272:156 273:254 274:78 289:46 290:254 291:253 292:92 296:17 297:226 298:217 299:49 300:148 301:253 302:78 317:165 318:254 319:239 320:24 324:20 325:253 326:253 327:58 328:18 329:115 330:24 344:37 345:248 346:254 347:91 352:2 353:117 354:250 355:163 356:91 372:77 373:253 374:254 375:39 382:196 383:253 384:173 400:159 401:254 402:218 403:15 410:77 411:254 412:255 413:61 428:234 429:253 430:113 438:21 439:226 440:254 441:135 455:25 456:240 457:253 458:68 467:195 468:254 469:135 483:79 484:253 485:253 495:195 496:254 497:135 511:79 512:253 513:253 514:76 523:195 524:254 525:99 540:212 541:254 542:209 543:9 550:10 551:209 552:196 553:15 568:54 569:253 570:254 571:137 572:36 576:2 577:20 578:168 579:253 580:60 596:28 597:235 598:254 599:253 600:199 601:124 602:79 603:79 604:167 605:253 606:253 607:185 608:30 625:15 626:117 627:217 628:253 629:253 630:253 631:254 632:253 633:240 634:109 635:12 655:27 656:126 657:208 658:253 659:193 660:147 661:40 +0 154:32 155:134 156:218 157:254 158:254 159:254 160:217 161:84 176:44 177:208 178:215 179:156 180:35 181:119 182:236 183:246 184:136 185:91 186:69 187:151 188:249 189:246 190:78 203:44 204:230 205:254 206:254 207:254 208:254 209:254 210:196 211:48 216:60 217:224 218:210 219:24 231:118 232:254 233:202 234:19 235:201 236:254 237:181 238:9 245:35 246:233 247:168 259:193 260:223 261:34 263:59 264:163 265:236 266:15 274:140 275:205 276:8 286:60 287:254 288:176 293:38 302:54 303:237 304:80 314:59 315:254 316:93 331:131 332:200 342:59 343:240 344:24 359:79 360:214 370:59 371:234 387:67 388:248 389:54 398:59 399:234 416:235 417:58 426:60 427:235 443:79 444:255 445:59 454:59 455:251 456:66 471:79 472:250 473:54 482:59 483:254 484:108 499:146 500:214 510:5 511:203 512:187 513:3 526:4 527:188 528:199 539:118 540:254 541:57 554:96 555:254 556:117 567:16 568:237 569:224 570:14 581:14 582:187 583:206 584:8 596:88 597:252 598:186 599:16 608:16 609:187 610:252 611:125 625:100 626:254 627:237 628:94 629:24 635:13 636:214 637:254 638:166 653:3 654:57 655:215 656:248 657:241 658:235 659:197 660:137 661:137 662:137 663:231 664:238 665:155 666:25 684:57 685:155 686:246 687:254 688:254 689:254 690:254 691:147 692:36 +1 124:102 125:252 126:252 127:41 152:102 153:250 154:250 155:202 180:102 181:250 182:250 183:232 184:91 208:102 209:250 210:250 211:212 212:29 236:102 237:252 238:252 239:254 240:150 264:102 265:250 266:250 267:252 268:149 292:102 293:250 294:250 295:252 296:149 320:102 321:250 322:250 323:252 324:231 325:80 349:152 350:252 351:254 352:252 353:100 377:151 378:250 379:252 380:250 381:100 405:151 406:250 407:252 408:250 409:100 433:151 434:250 435:252 436:250 437:100 461:123 462:243 463:254 464:252 465:100 490:202 491:252 492:250 493:100 518:80 519:252 520:250 521:190 522:30 547:252 548:250 549:250 550:49 575:255 576:252 577:252 578:252 579:214 580:31 603:171 604:250 605:250 606:250 607:252 608:190 609:40 631:20 632:160 633:250 634:250 635:252 636:250 637:100 660:20 661:170 662:250 663:212 664:49 665:20 +0 124:20 125:121 126:197 127:253 128:64 151:23 152:200 153:252 154:252 155:252 156:184 157:6 178:25 179:197 180:252 181:252 182:252 183:252 184:253 185:228 186:107 187:15 205:26 206:196 207:252 208:252 209:252 210:252 211:252 212:253 213:252 214:252 215:219 216:178 217:21 233:186 234:252 235:238 236:94 237:67 238:224 239:217 240:53 241:109 242:245 243:252 244:252 245:213 246:63 260:98 261:242 262:252 263:101 266:39 267:31 270:109 271:128 272:241 273:252 274:207 275:97 287:17 288:230 289:252 290:241 291:56 300:109 301:252 302:252 303:229 304:17 314:13 315:192 316:252 317:243 318:96 328:25 329:127 330:252 331:252 332:120 342:121 343:252 344:252 345:165 357:125 358:252 359:252 360:246 361:70 370:190 371:252 372:252 373:39 385:26 386:210 387:252 388:252 389:119 398:255 399:253 400:159 414:22 415:209 416:253 417:183 426:253 427:252 428:103 443:34 444:252 445:252 454:253 455:252 456:26 471:27 472:252 473:252 482:253 483:252 484:168 485:13 499:70 500:252 501:209 510:147 511:252 512:252 513:75 526:68 527:233 528:252 529:119 538:121 539:252 540:252 541:189 542:40 552:15 553:82 554:231 555:252 556:214 557:31 566:38 567:135 568:248 569:252 570:231 571:145 572:41 573:41 574:41 575:41 576:20 577:24 578:37 579:83 580:194 581:252 582:252 583:212 584:33 596:83 597:213 598:252 599:252 600:252 601:252 602:252 603:252 604:204 605:213 606:243 607:252 608:252 609:252 610:212 611:34 625:34 626:140 627:238 628:248 629:252 630:252 631:252 632:253 633:252 634:252 635:241 636:238 637:238 638:75 656:82 657:119 658:119 659:119 660:120 661:119 662:119 663:19 +1 127:20 128:254 129:255 130:37 155:19 156:253 157:253 158:134 183:19 184:253 185:253 186:246 187:125 211:76 212:253 213:253 214:253 215:158 239:207 240:253 241:253 242:253 243:158 267:207 268:253 269:253 270:253 271:158 294:48 295:223 296:253 297:253 298:243 299:106 322:141 323:253 324:253 325:253 326:113 349:65 350:237 351:253 352:253 353:253 354:36 377:76 378:253 379:253 380:253 381:253 382:36 405:76 406:253 407:253 408:253 409:253 410:36 433:76 434:253 435:253 436:253 437:118 438:4 460:4 461:148 462:253 463:253 464:253 465:103 488:10 489:253 490:253 491:253 492:253 493:103 516:10 517:253 518:253 519:253 520:173 521:7 544:10 545:253 546:253 547:253 548:168 572:143 573:253 574:253 575:239 576:49 600:198 601:253 602:253 603:234 615:140 628:198 629:253 630:253 631:234 656:198 657:253 658:253 659:234 +0 235:40 236:37 238:7 239:77 240:137 241:136 242:136 243:136 244:136 245:40 246:6 261:16 262:135 263:254 264:233 266:152 267:215 268:96 269:140 270:155 271:118 272:230 273:254 274:158 275:68 288:19 289:164 290:254 291:114 294:235 295:140 301:99 302:230 303:254 304:186 305:14 315:70 316:226 317:242 318:121 322:104 323:195 324:38 330:33 331:179 332:253 333:140 342:41 343:241 344:198 345:43 359:24 360:209 361:223 370:164 371:250 372:66 388:136 389:253 398:254 399:158 416:136 417:215 426:255 427:76 442:5 443:127 444:246 445:133 454:254 455:122 469:5 470:150 471:247 472:91 473:9 482:254 483:165 495:13 496:79 497:194 498:216 499:84 510:111 511:251 512:87 519:16 520:25 521:40 522:107 523:186 524:213 525:117 526:25 538:14 539:185 540:235 541:142 542:23 546:91 547:157 548:231 549:207 550:126 551:49 569:143 570:195 571:255 572:254 573:254 574:244 575:157 576:76 599:39 600:39 601:39 602:33 +1 128:166 129:255 130:187 131:6 156:165 157:253 158:253 159:13 183:15 184:191 185:253 186:253 187:13 211:49 212:253 213:253 214:253 215:13 239:141 240:253 241:253 242:169 243:4 266:4 267:189 268:253 269:249 270:53 294:69 295:253 296:253 297:246 322:69 323:253 324:253 325:246 350:118 351:253 352:253 353:124 378:206 379:253 380:231 381:21 405:66 406:241 407:253 408:199 433:105 434:253 435:253 436:89 460:3 461:228 462:253 463:252 464:86 488:111 489:253 490:253 491:205 516:166 517:253 518:253 519:75 543:43 544:249 545:253 546:193 547:9 570:4 571:160 572:253 573:253 574:184 598:37 599:253 600:253 601:253 602:88 626:140 627:253 628:253 629:186 630:18 654:14 655:253 656:253 657:27 +1 128:117 129:128 155:2 156:199 157:127 183:81 184:254 185:87 211:116 212:254 213:48 239:175 240:234 241:18 266:5 267:230 268:168 294:80 295:255 296:142 322:80 323:255 324:142 350:80 351:251 352:57 378:129 379:239 406:164 407:209 433:28 434:245 435:159 461:64 462:254 463:144 489:84 490:254 491:80 517:143 518:254 519:30 544:3 545:225 546:200 572:48 573:254 574:174 600:48 601:254 602:174 628:93 629:254 630:129 656:53 657:234 658:41 +1 129:159 130:142 156:11 157:220 158:141 184:78 185:254 186:141 212:111 213:254 214:109 240:196 241:221 242:15 267:26 268:221 269:159 295:63 296:254 297:159 323:178 324:254 325:93 350:7 351:191 352:254 353:97 378:42 379:255 380:254 381:41 406:42 407:254 408:195 409:10 434:141 435:255 436:78 461:11 462:202 463:254 464:59 489:86 490:254 491:254 492:59 517:142 518:254 519:248 520:52 545:142 546:254 547:195 573:142 574:254 575:164 601:142 602:254 603:77 629:142 630:254 631:131 657:77 658:172 659:5 +0 124:66 125:254 126:254 127:58 128:60 129:59 130:59 131:50 151:73 152:233 153:253 154:253 155:148 156:254 157:253 158:253 159:232 160:73 179:156 180:253 181:253 182:253 183:117 184:255 185:253 186:253 187:253 188:223 189:176 190:162 205:37 206:116 207:246 208:253 209:180 210:18 211:4 212:18 213:109 214:241 215:253 216:253 217:253 218:236 219:28 233:235 234:253 235:253 236:245 237:107 242:109 243:170 244:253 245:253 246:253 247:174 261:235 262:253 263:253 264:233 271:15 272:156 273:253 274:253 275:223 276:72 287:10 288:156 289:250 290:253 291:253 292:67 300:99 301:253 302:253 303:253 304:127 305:5 315:118 316:253 317:253 318:253 319:204 320:26 328:68 329:223 330:253 331:253 332:253 333:57 342:32 343:191 344:253 345:253 346:253 347:97 357:156 358:253 359:253 360:253 361:57 370:59 371:253 372:253 373:253 374:253 375:97 385:36 386:224 387:253 388:253 389:57 398:60 399:254 400:255 401:254 402:156 413:37 414:226 415:254 416:254 417:58 426:59 427:253 428:253 429:253 430:154 441:156 442:253 443:253 444:253 445:57 454:59 455:253 456:253 457:253 458:154 469:156 470:253 471:253 472:253 473:57 482:59 483:253 484:253 485:253 486:246 487:90 496:16 497:171 498:253 499:253 500:231 501:49 510:59 511:253 512:253 513:253 514:253 515:156 516:91 524:99 525:253 526:253 527:222 528:71 538:59 539:253 540:253 541:253 542:253 543:253 544:245 545:109 551:145 552:194 553:253 554:253 555:174 566:9 567:38 568:174 569:251 570:253 571:253 572:253 573:241 574:215 575:215 576:217 577:215 578:215 579:250 580:253 581:253 582:221 583:26 597:235 598:253 599:253 600:253 601:253 602:253 603:253 604:254 605:253 606:253 607:253 608:253 609:204 610:26 625:108 626:116 627:200 628:253 629:253 630:253 631:253 632:254 633:253 634:253 635:253 636:199 637:44 655:36 656:57 657:118 658:253 659:253 660:58 661:57 662:57 663:57 664:35 +1 129:101 130:222 131:84 157:225 158:252 159:84 184:89 185:246 186:208 187:19 212:128 213:252 214:195 239:79 240:253 241:252 242:195 267:141 268:255 269:253 270:133 294:26 295:240 296:253 297:252 298:55 322:60 323:252 324:253 325:154 326:12 349:7 350:178 351:252 352:253 353:27 377:57 378:252 379:252 380:253 381:27 405:57 406:253 407:253 408:204 409:15 433:104 434:252 435:252 436:94 460:19 461:209 462:252 463:252 488:101 489:252 490:252 491:157 516:225 517:252 518:252 519:112 544:226 545:253 546:240 547:63 572:225 573:252 574:223 600:225 601:252 602:223 628:225 629:252 630:242 631:75 656:146 657:252 658:236 659:50 +0 124:41 125:254 126:254 127:157 128:34 129:34 130:218 131:255 132:206 133:34 134:18 151:53 152:238 153:252 154:252 155:252 156:252 157:252 158:252 159:252 160:252 161:252 162:162 163:26 178:66 179:220 180:252 181:252 182:252 183:209 184:153 185:223 186:252 187:252 188:252 189:252 190:252 191:98 206:166 207:252 208:252 209:252 210:252 211:141 213:85 214:230 215:252 216:252 217:252 218:252 219:98 234:166 235:252 236:252 237:252 238:252 239:141 242:73 243:102 244:252 245:252 246:252 247:98 262:166 263:252 264:252 265:252 266:191 267:30 271:5 272:97 273:252 274:252 275:220 276:51 289:123 290:245 291:252 292:252 293:202 294:14 300:56 301:252 302:252 303:252 304:65 316:18 317:154 318:252 319:252 320:241 328:56 329:252 330:252 331:252 332:65 343:21 344:146 345:252 346:252 347:252 348:241 356:56 357:252 358:252 359:252 360:65 371:67 372:252 373:252 374:252 375:252 376:241 384:56 385:252 386:252 387:252 388:65 399:67 400:252 401:252 402:252 403:252 404:116 412:56 413:252 414:252 415:252 416:65 427:67 428:252 429:252 430:252 431:252 432:20 440:56 441:252 442:252 443:252 444:65 455:67 456:252 457:252 458:252 459:87 460:4 468:56 469:252 470:252 471:124 472:11 483:67 484:252 485:252 486:252 487:54 494:19 495:236 496:245 497:252 498:252 499:98 511:67 512:252 513:252 514:252 515:97 516:5 521:39 522:219 523:252 524:252 525:252 526:252 527:98 539:67 540:252 541:252 542:252 543:252 544:102 545:89 546:89 547:89 548:89 549:203 550:252 551:252 552:252 553:252 554:209 555:64 567:67 568:252 569:252 570:252 571:252 572:252 573:252 574:252 575:252 576:252 577:252 578:252 579:252 580:226 581:130 582:68 595:67 596:252 597:252 598:252 599:252 600:252 601:252 602:252 603:252 604:252 605:252 606:252 607:239 608:77 623:17 624:65 625:163 626:252 627:252 628:252 629:252 630:252 631:252 632:252 633:252 634:96 635:59 653:17 654:176 655:252 656:252 657:252 658:252 659:155 660:32 661:32 662:6 +0 96:56 97:247 98:121 124:24 125:242 126:245 127:122 153:231 154:253 155:253 156:104 157:12 181:90 182:253 183:253 184:254 185:221 186:120 187:120 188:85 206:67 207:75 208:36 209:11 210:56 211:222 212:254 213:253 214:253 215:253 216:245 217:207 218:36 233:86 234:245 235:249 236:105 239:44 240:224 241:230 242:253 243:253 244:253 245:253 246:214 247:10 260:8 261:191 262:253 263:143 269:29 270:119 271:119 272:158 273:253 274:253 275:94 288:15 289:253 290:226 291:48 300:4 301:183 302:253 303:248 304:56 316:42 317:253 318:178 329:179 330:253 331:184 332:14 344:164 345:253 346:178 357:179 358:253 359:163 371:61 372:254 373:254 374:179 384:76 385:254 386:254 387:164 399:60 400:253 401:253 402:178 411:29 412:206 413:253 414:253 415:40 427:60 428:253 429:253 430:178 439:120 440:253 441:253 442:245 443:13 455:60 456:253 457:253 458:178 467:120 468:253 469:239 470:63 483:60 484:253 485:253 486:178 494:14 495:238 496:253 497:179 511:18 512:190 513:253 514:231 515:70 521:43 522:184 523:253 524:253 525:74 540:86 541:253 542:253 543:239 544:134 545:8 548:56 549:163 550:253 551:253 552:213 553:35 568:16 569:253 570:253 571:253 572:253 573:240 574:239 575:239 576:247 577:253 578:253 579:210 580:27 596:4 597:59 598:204 599:253 600:253 601:253 602:253 603:253 604:254 605:253 606:250 607:110 626:31 627:122 628:253 629:253 630:253 631:253 632:255 633:217 634:98 +0 125:19 126:164 127:253 128:255 129:253 130:118 131:59 132:36 153:78 154:251 155:251 156:253 157:251 158:251 159:251 160:199 161:45 180:14 181:198 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:204 190:26 208:5 209:117 210:251 211:251 212:243 213:212 214:239 215:251 216:251 217:251 218:218 236:95 237:251 238:251 239:251 240:120 242:175 243:251 244:251 245:251 246:231 263:97 264:237 265:251 266:251 267:251 270:67 271:240 272:251 273:251 274:243 275:108 290:8 291:163 292:251 293:251 294:240 295:81 299:68 300:251 301:251 302:251 303:179 304:9 317:13 318:145 319:251 320:251 321:226 322:80 327:39 328:251 329:251 330:251 331:251 332:115 345:144 346:251 347:251 348:251 349:173 355:18 356:167 357:251 358:251 359:251 360:115 373:233 374:251 375:251 376:251 377:173 384:98 385:251 386:251 387:251 388:115 400:176 401:253 402:253 403:216 404:179 412:99 413:253 414:253 415:253 416:116 427:55 428:210 429:251 430:251 431:96 440:98 441:251 442:251 443:214 444:62 455:117 456:251 457:251 458:251 459:96 467:28 468:204 469:251 470:237 471:53 482:55 483:241 484:251 485:251 486:160 487:7 494:28 495:222 496:251 497:251 498:231 510:59 511:251 512:251 513:251 514:153 520:23 521:98 522:204 523:251 524:251 525:251 526:156 538:59 539:251 540:251 541:251 542:153 546:85 547:155 548:179 549:251 550:251 551:251 552:251 553:154 554:15 566:59 567:251 568:251 569:251 570:236 571:214 572:214 573:214 574:234 575:251 576:253 577:251 578:251 579:248 580:156 581:15 594:41 595:209 596:251 597:251 598:251 599:251 600:251 601:251 602:251 603:251 604:253 605:251 606:196 607:146 623:54 624:115 625:241 626:251 627:251 628:251 629:251 630:251 631:251 632:253 633:187 634:35 653:83 654:251 655:251 656:251 657:251 658:251 659:101 660:57 661:31 +1 129:232 130:255 131:107 156:58 157:244 158:253 159:106 184:95 185:253 186:253 187:106 212:95 213:253 214:253 215:106 240:95 241:253 242:249 243:69 268:144 269:253 270:192 295:97 296:233 297:253 298:66 323:195 324:253 325:253 326:5 350:38 351:232 352:253 353:182 354:2 377:10 378:160 379:253 380:231 381:53 405:42 406:253 407:253 408:158 433:141 434:253 435:253 436:115 460:75 461:245 462:253 463:183 464:4 487:1 488:147 489:253 490:251 491:58 515:20 516:253 517:253 518:180 543:202 544:253 545:226 546:27 571:243 572:253 573:212 598:85 599:251 600:253 601:173 626:209 627:253 628:244 629:57 654:169 655:253 656:174 +1 127:63 128:128 129:2 155:63 156:254 157:123 183:63 184:254 185:179 211:63 212:254 213:179 239:63 240:254 241:179 267:142 268:254 269:179 295:187 296:254 297:158 323:187 324:254 325:55 350:68 351:235 352:254 353:55 378:181 379:254 380:254 381:55 406:181 407:254 408:202 409:14 434:181 435:254 436:186 462:181 463:254 464:146 490:181 491:254 492:62 518:181 519:254 520:62 546:181 547:254 548:62 574:181 575:255 576:62 602:181 603:254 604:241 605:52 630:181 631:254 632:222 633:30 658:181 659:224 660:34 +1 130:131 131:255 132:184 133:15 157:99 158:247 159:253 160:182 161:15 185:124 186:253 187:253 188:253 189:38 212:9 213:171 214:253 215:253 216:140 217:1 240:47 241:253 242:253 243:251 244:117 267:43 268:219 269:253 270:253 271:153 295:78 296:253 297:253 298:253 299:84 323:97 324:253 325:253 326:244 327:74 350:69 351:243 352:253 353:253 354:183 377:10 378:168 379:253 380:253 381:215 382:34 405:31 406:253 407:253 408:253 409:129 433:107 434:253 435:253 436:242 437:67 460:24 461:204 462:253 463:253 464:187 488:95 489:253 490:253 491:201 492:25 516:239 517:253 518:253 519:176 543:119 544:251 545:253 546:253 547:138 570:30 571:212 572:253 573:252 574:165 575:8 598:193 599:253 600:253 601:222 626:193 627:253 628:253 629:189 654:193 655:253 656:201 657:27 +0 125:57 126:255 127:253 128:198 129:85 153:168 154:253 155:251 156:253 157:251 158:169 159:56 180:86 181:253 182:254 183:253 184:254 185:253 186:254 187:253 188:57 208:197 209:251 210:253 211:251 212:253 213:251 214:253 215:251 216:225 217:56 235:169 236:255 237:253 238:226 239:56 241:114 242:254 243:253 244:254 245:84 262:57 263:224 264:253 265:251 266:56 270:139 271:251 272:253 273:83 290:141 291:253 292:255 293:84 298:57 299:225 300:254 301:196 318:253 319:251 320:253 321:83 327:168 328:253 329:83 345:169 346:254 347:253 348:169 355:169 356:254 357:253 358:169 373:168 374:253 375:251 376:56 383:168 384:253 385:251 386:56 401:169 402:254 403:84 412:254 413:253 429:168 430:253 431:83 440:253 441:251 456:29 457:197 458:254 459:84 467:169 468:254 469:196 484:85 485:251 486:253 487:83 494:57 495:224 496:253 497:83 512:57 513:225 514:254 515:139 521:57 522:141 523:253 524:254 525:84 541:168 542:253 543:251 544:169 545:56 547:114 548:169 549:224 550:253 551:251 552:253 553:83 569:169 570:254 571:253 572:254 573:253 574:254 575:253 576:254 577:253 578:254 579:253 580:226 581:56 597:56 598:253 599:251 600:253 601:251 602:253 603:251 604:253 605:251 606:253 607:251 608:56 626:169 627:225 628:254 629:253 630:254 631:253 632:254 633:253 634:226 635:56 655:56 656:253 657:251 658:253 659:251 660:84 661:83 662:56 +0 127:12 128:105 129:224 130:255 131:247 132:22 155:131 156:254 157:254 158:243 159:252 160:76 182:131 183:225 184:254 185:224 186:48 187:136 208:13 209:109 210:252 211:254 212:254 213:254 214:197 215:76 235:9 236:181 237:254 238:254 239:240 240:229 241:237 242:254 243:252 244:152 245:21 262:9 263:143 264:254 265:254 266:226 267:36 269:22 270:138 271:254 272:254 273:188 289:13 290:181 291:254 292:254 293:250 294:64 298:2 299:53 300:236 301:252 302:131 317:102 318:254 319:254 320:254 321:111 328:56 329:243 330:251 331:42 344:30 345:186 346:254 347:254 348:206 349:29 357:199 358:254 359:91 372:92 373:254 374:254 375:237 376:13 385:134 386:254 387:91 400:133 401:254 402:254 403:126 413:134 414:250 415:17 428:187 429:254 430:237 431:23 441:200 442:183 456:187 457:254 458:213 467:2 468:134 469:252 470:101 484:183 485:254 486:133 495:14 496:254 497:234 498:34 512:92 513:254 514:161 522:84 523:204 524:254 525:56 540:92 541:254 542:229 549:85 550:252 551:252 552:188 553:11 568:56 569:252 570:229 575:3 576:53 577:235 578:253 579:166 597:224 598:245 599:130 600:68 601:68 602:134 603:214 604:254 605:254 606:159 625:141 626:254 627:254 628:254 629:254 630:254 631:254 632:233 633:95 634:3 653:14 654:152 655:254 656:254 657:254 658:186 659:157 660:53 +1 130:226 131:247 132:55 157:99 158:248 159:254 160:230 161:30 185:125 186:254 187:254 188:254 189:38 213:125 214:254 215:254 216:212 217:24 240:18 241:223 242:254 243:252 244:118 268:24 269:254 270:254 271:239 295:27 296:195 297:254 298:254 299:93 323:78 324:254 325:254 326:246 327:74 351:158 352:254 353:254 354:185 378:41 379:239 380:254 381:254 382:43 405:22 406:218 407:254 408:254 409:167 410:9 433:32 434:254 435:254 436:254 437:130 460:24 461:187 462:254 463:254 464:234 465:16 488:189 489:254 490:254 491:254 492:128 515:64 516:247 517:254 518:255 519:219 520:42 543:139 544:254 545:254 546:222 547:40 570:30 571:213 572:254 573:235 574:45 598:194 599:254 600:254 601:223 626:194 627:254 628:254 629:190 654:194 655:254 656:202 657:27 +1 130:166 131:253 132:124 133:53 158:140 159:251 160:251 161:180 185:125 186:246 187:251 188:251 189:51 212:32 213:190 214:251 215:251 216:251 217:103 240:21 241:174 242:251 243:251 244:251 268:73 269:176 270:253 271:253 272:201 296:149 297:251 298:251 299:251 300:71 323:27 324:228 325:251 326:251 327:157 328:10 351:180 352:253 353:251 354:251 355:142 377:27 378:180 379:231 380:253 381:251 382:96 383:41 405:89 406:253 407:253 408:255 409:211 410:25 433:217 434:251 435:251 436:253 437:107 460:21 461:221 462:251 463:251 464:242 465:92 487:32 488:190 489:251 490:251 491:251 492:103 515:202 516:251 517:251 518:251 519:122 542:53 543:255 544:253 545:253 546:221 547:51 570:180 571:253 572:251 573:251 574:142 598:180 599:253 600:251 601:251 602:142 626:180 627:253 628:251 629:157 630:82 654:180 655:253 656:147 657:10 +1 129:17 130:206 131:229 132:44 157:2 158:125 159:254 160:123 185:95 186:254 187:254 188:123 212:78 213:240 214:254 215:254 216:123 240:100 241:254 242:254 243:254 244:123 267:2 268:129 269:254 270:254 271:220 272:20 295:9 296:254 297:254 298:254 299:123 322:22 323:179 324:254 325:254 326:254 327:49 350:83 351:254 352:254 353:254 354:183 355:19 378:136 379:254 380:254 381:254 382:139 404:3 405:111 406:252 407:254 408:254 409:232 410:45 432:67 433:254 434:254 435:254 436:216 437:40 459:14 460:192 461:254 462:254 463:254 464:140 486:23 487:192 488:254 489:254 490:254 491:246 514:77 515:254 516:254 517:255 518:241 519:100 541:65 542:235 543:254 544:254 545:254 546:172 568:30 569:238 570:254 571:254 572:254 573:219 574:26 596:34 597:254 598:254 599:254 600:216 601:41 624:34 625:254 626:254 627:254 628:188 652:12 653:170 654:254 655:254 656:82 +1 130:218 131:253 132:124 157:84 158:236 159:251 160:251 184:63 185:236 186:251 187:251 188:122 212:73 213:251 214:251 215:251 216:173 240:202 241:251 242:251 243:251 244:71 267:53 268:255 269:253 270:253 271:253 272:72 295:180 296:253 297:251 298:251 299:188 300:30 323:180 324:253 325:251 326:251 327:142 350:47 351:211 352:253 353:251 354:235 355:82 377:27 378:211 379:251 380:253 381:251 382:215 405:89 406:253 407:253 408:255 409:253 410:164 433:217 434:251 435:251 436:253 437:168 438:15 460:21 461:221 462:251 463:251 464:253 465:107 487:32 488:190 489:251 490:251 491:251 492:221 493:61 515:73 516:251 517:251 518:251 519:251 520:180 543:255 544:253 545:253 546:253 547:201 570:105 571:253 572:251 573:251 574:251 575:71 598:180 599:253 600:251 601:246 602:137 603:10 626:180 627:253 628:251 629:215 654:180 655:253 656:251 657:86 +1 124:102 125:180 126:1 152:140 153:254 154:130 180:140 181:254 182:204 208:140 209:254 210:204 236:72 237:254 238:204 264:25 265:231 266:250 267:135 292:11 293:211 294:254 295:222 321:101 322:254 323:250 324:15 349:96 350:254 351:254 352:95 377:2 378:251 379:254 380:95 405:2 406:251 407:254 408:95 433:96 434:254 435:254 436:95 461:53 462:253 463:254 464:139 490:250 491:254 492:235 493:27 518:201 519:254 520:254 521:128 546:80 547:254 548:254 549:139 574:65 575:254 576:254 577:139 602:150 603:254 604:254 605:139 630:229 631:254 632:254 633:43 658:52 659:196 660:168 661:9 +0 128:87 129:208 130:249 155:27 156:212 157:254 158:195 182:118 183:225 184:254 185:254 186:232 187:147 188:46 209:115 210:248 211:254 212:254 213:254 214:254 215:254 216:230 217:148 218:12 236:18 237:250 238:254 239:245 240:226 241:254 242:254 243:254 244:254 245:254 246:148 263:92 264:205 265:254 266:250 267:101 268:20 269:194 270:254 271:254 272:254 273:254 274:229 275:53 291:152 292:254 293:254 294:94 297:14 298:124 299:187 300:254 301:254 302:254 303:213 318:95 319:252 320:254 321:206 322:15 327:3 328:6 329:51 330:231 331:254 332:94 345:50 346:246 347:254 348:254 349:20 358:200 359:254 360:96 372:21 373:184 374:254 375:254 376:147 377:2 386:200 387:254 388:96 400:177 401:254 402:254 403:218 404:33 413:16 414:211 415:254 416:96 427:11 428:219 429:254 430:251 431:92 441:84 442:254 443:232 444:44 455:101 456:254 457:254 458:141 469:162 470:254 471:231 472:42 483:235 484:254 485:227 486:42 496:51 497:238 498:254 499:213 511:235 512:254 513:199 524:160 525:254 526:229 527:52 539:235 540:254 541:199 549:10 550:84 551:150 552:253 553:254 554:147 567:235 568:254 569:213 570:20 575:17 576:63 577:158 578:254 579:254 580:254 581:155 582:12 595:122 596:248 597:254 598:204 599:98 600:42 601:177 602:180 603:200 604:254 605:254 606:253 607:213 608:82 609:10 624:203 625:254 626:254 627:254 628:254 629:254 630:254 631:254 632:251 633:219 634:94 652:35 653:221 654:254 655:254 656:254 657:254 658:254 659:217 660:95 +1 126:134 127:230 154:133 155:231 156:10 182:133 183:253 184:96 210:133 211:253 212:96 238:133 239:253 240:183 266:133 267:253 268:217 294:133 295:253 296:217 322:133 323:253 324:217 350:133 351:253 352:217 378:133 379:253 380:217 406:134 407:254 408:218 434:133 435:253 436:159 462:133 463:253 464:199 490:156 491:253 492:96 518:254 519:247 520:73 546:254 547:248 548:74 573:99 574:254 575:245 576:64 600:89 601:230 602:254 603:125 627:140 628:251 629:253 630:243 631:10 655:114 656:242 657:195 658:69 +1 125:29 126:85 127:255 128:139 153:197 154:251 155:253 156:251 181:254 182:253 183:254 184:253 209:253 210:251 211:253 212:251 237:254 238:253 239:254 240:253 265:253 266:251 267:253 268:138 293:254 294:253 295:254 296:196 321:253 322:251 323:253 324:196 349:254 350:253 351:254 352:84 377:253 378:251 379:253 380:196 405:254 406:253 407:254 408:253 433:253 434:251 435:253 436:251 461:254 462:253 463:254 464:253 489:253 490:251 491:253 492:251 517:254 518:253 519:254 520:253 545:253 546:251 547:253 548:251 573:254 574:253 575:254 576:253 601:253 602:251 603:253 604:251 629:57 630:225 631:254 632:253 658:56 659:253 660:251 +1 125:149 126:255 127:254 128:58 153:215 154:253 155:183 156:2 180:41 181:232 182:253 183:181 208:92 209:253 210:253 211:181 236:92 237:253 238:253 239:181 264:92 265:253 266:253 267:181 292:92 293:253 294:253 295:181 320:92 321:253 322:253 323:181 348:92 349:253 350:253 351:181 376:92 377:253 378:253 379:181 404:92 405:253 406:253 407:181 432:92 433:253 434:253 435:181 460:92 461:253 462:253 463:181 488:31 489:228 490:253 491:181 517:198 518:253 519:228 520:54 545:33 546:226 547:253 548:195 549:7 574:199 575:253 576:253 577:75 602:34 603:218 604:253 605:228 606:117 607:14 608:12 631:33 632:219 633:253 634:253 635:253 636:211 660:32 661:123 662:149 663:230 664:41 +1 130:79 131:203 132:141 157:51 158:240 159:240 160:140 185:88 186:252 187:252 188:140 213:197 214:252 215:252 216:140 241:197 242:252 243:252 244:140 268:147 269:253 270:253 271:253 295:38 296:234 297:252 298:242 299:89 323:113 324:252 325:252 326:223 350:16 351:207 352:252 353:252 354:129 377:16 378:203 379:253 380:252 381:220 382:37 405:29 406:253 407:255 408:253 409:56 432:19 433:181 434:252 435:253 436:176 437:6 460:166 461:252 462:252 463:228 464:52 487:10 488:203 489:252 490:252 491:126 514:63 515:178 516:252 517:252 518:173 542:114 543:253 544:253 545:225 570:238 571:252 572:252 573:99 596:7 597:135 598:253 599:252 600:176 601:19 624:29 625:252 626:253 627:252 628:55 652:13 653:189 654:253 655:204 656:25 +1 126:94 127:254 128:75 154:166 155:253 156:231 182:208 183:253 184:147 210:208 211:253 212:116 238:208 239:253 240:168 266:146 267:254 268:222 294:166 295:253 296:116 322:208 323:253 324:116 350:166 351:253 352:158 378:145 379:253 380:231 406:209 407:254 408:169 434:187 435:253 436:168 462:93 463:253 464:116 490:93 491:253 492:116 518:93 519:253 520:116 546:94 547:254 548:179 549:11 574:93 575:253 576:246 577:101 602:145 603:253 604:255 605:92 630:93 631:253 632:246 633:59 658:93 659:253 660:74 +0 127:46 128:105 129:254 130:254 131:224 132:59 133:59 134:9 155:196 156:254 157:253 158:253 159:253 160:253 161:253 162:128 182:96 183:235 184:254 185:253 186:253 187:253 188:253 189:253 190:247 191:122 208:4 209:101 210:244 211:253 212:254 213:234 214:241 215:253 216:253 217:253 218:253 219:186 220:18 236:96 237:253 238:253 239:253 240:232 241:83 242:109 243:170 244:253 245:253 246:253 247:253 248:116 264:215 265:253 266:253 267:253 268:196 271:40 272:253 273:253 274:253 275:253 276:116 290:8 291:141 292:247 293:253 294:253 295:237 296:29 299:6 300:38 301:171 302:253 303:253 304:116 317:13 318:146 319:253 320:253 321:253 322:253 323:57 329:156 330:253 331:253 332:116 345:40 346:253 347:253 348:253 349:253 350:178 351:27 357:156 358:253 359:253 360:116 372:136 373:204 374:253 375:253 376:253 377:192 378:27 385:156 386:253 387:253 388:116 399:28 400:195 401:254 402:254 403:254 404:250 405:135 412:99 413:255 414:254 415:254 416:117 427:118 428:253 429:253 430:253 431:253 432:142 439:19 440:170 441:253 442:253 443:216 444:62 454:42 455:212 456:253 457:253 458:253 459:253 460:38 466:124 467:188 468:253 469:253 470:253 471:174 482:59 483:253 484:253 485:253 486:237 487:93 488:3 491:31 492:40 493:130 494:247 495:253 496:253 497:253 498:204 499:13 510:59 511:253 512:253 513:253 514:154 518:54 519:218 520:254 521:253 522:253 523:253 524:253 525:253 526:38 538:59 539:253 540:253 541:253 542:215 543:156 544:156 545:156 546:209 547:253 548:255 549:253 550:253 551:253 552:192 553:97 554:15 566:55 567:242 568:253 569:253 570:253 571:253 572:253 573:253 574:253 575:253 576:254 577:253 578:253 579:204 580:23 595:118 596:253 597:253 598:253 599:253 600:253 601:253 602:253 603:253 604:254 605:216 606:174 607:13 623:54 624:116 625:243 626:253 627:253 628:253 629:253 630:253 631:146 632:117 633:62 653:53 654:132 655:253 656:253 657:192 658:57 659:13 +1 125:42 126:232 127:254 128:58 153:86 154:253 155:253 156:58 181:86 182:253 183:253 184:58 209:206 210:253 211:253 212:58 237:215 238:253 239:253 240:58 265:215 266:253 267:253 268:58 293:215 294:253 295:253 296:58 321:215 322:253 323:253 324:58 349:215 350:253 351:253 352:58 377:215 378:253 379:253 380:58 405:215 406:253 407:253 408:58 433:188 434:253 435:253 436:85 461:86 462:253 463:253 464:200 465:12 489:29 490:223 491:253 492:253 493:151 518:209 519:253 520:253 521:194 546:128 547:253 548:253 549:200 550:8 574:32 575:213 576:253 577:253 578:152 579:6 603:32 604:221 605:253 606:253 607:153 608:5 632:90 633:215 634:253 635:253 636:151 661:59 662:253 663:253 664:84 +1 156:60 157:229 158:38 184:187 185:254 186:78 211:121 212:252 213:254 214:78 239:197 240:254 241:206 242:6 267:197 268:254 269:202 294:27 295:218 296:233 297:62 322:117 323:254 324:195 350:203 351:254 352:195 377:64 378:244 379:254 380:195 405:79 406:254 407:255 408:161 433:79 434:254 435:254 436:65 461:79 462:254 463:241 464:52 489:79 490:254 491:189 517:79 518:254 519:189 545:79 546:254 547:189 573:79 574:254 575:189 601:79 602:254 603:194 604:5 629:35 630:219 631:254 632:72 658:34 659:223 660:195 687:129 688:195 +1 101:11 102:150 103:72 129:37 130:251 131:71 157:63 158:251 159:71 185:217 186:251 187:71 213:217 214:251 215:71 240:145 241:253 242:253 243:72 267:42 268:206 269:251 270:251 271:71 295:99 296:251 297:251 298:251 299:71 323:253 324:251 325:251 326:251 327:71 350:130 351:253 352:251 353:251 354:251 355:71 377:110 378:253 379:255 380:253 381:253 382:253 383:72 405:109 406:251 407:253 408:251 409:251 410:188 411:30 433:109 434:251 435:253 436:251 437:246 438:123 460:16 461:170 462:251 463:253 464:251 465:215 488:37 489:251 490:251 491:253 492:251 493:86 516:218 517:253 518:253 519:255 520:253 521:35 543:84 544:236 545:251 546:251 547:253 548:168 549:15 571:144 572:251 573:251 574:251 575:190 576:15 599:144 600:251 601:251 602:251 603:180 626:53 627:221 628:251 629:251 630:251 631:180 +0 127:45 128:254 129:254 130:254 131:148 132:24 133:9 154:43 155:254 156:252 157:252 158:252 159:252 160:252 161:121 162:13 181:58 182:237 183:254 184:252 185:252 186:252 187:252 188:252 189:252 190:68 208:69 209:224 210:252 211:254 212:252 213:252 214:252 215:252 216:252 217:252 218:135 219:17 235:75 236:216 237:252 238:252 239:254 240:231 241:168 242:252 243:252 244:252 245:252 246:252 247:45 262:77 263:212 264:252 265:252 266:252 267:242 268:93 269:32 270:114 271:177 272:252 273:252 274:252 275:158 276:12 289:75 290:212 291:252 292:252 293:252 294:252 295:231 299:116 300:252 301:252 302:252 303:252 304:21 316:69 317:216 318:252 319:252 320:252 321:252 322:252 323:62 327:116 328:252 329:252 330:252 331:252 332:21 344:93 345:252 346:252 347:252 348:252 349:252 350:62 355:21 356:158 357:252 358:252 359:252 360:21 371:64 372:239 373:252 374:252 375:252 376:252 377:252 378:21 384:139 385:252 386:252 387:252 388:21 398:5 399:87 400:254 401:254 402:254 403:254 404:237 405:41 411:11 412:150 413:254 414:254 415:254 416:22 425:5 426:85 427:252 428:252 429:252 430:252 431:222 432:55 439:116 440:252 441:252 442:252 443:214 444:18 453:24 454:252 455:252 456:252 457:252 458:252 459:91 466:26 467:153 468:252 469:252 470:252 471:45 481:24 482:252 483:252 484:252 485:252 486:252 487:91 492:18 493:93 494:151 495:252 496:252 497:252 498:184 499:28 509:24 510:252 511:252 512:252 513:252 514:252 515:164 516:116 517:116 518:116 519:117 520:141 521:252 522:252 523:252 524:252 525:252 526:68 537:24 538:252 539:252 540:252 541:252 542:252 543:252 544:252 545:252 546:252 547:254 548:252 549:252 550:252 551:252 552:252 553:163 554:31 565:9 566:121 567:252 568:252 569:252 570:252 571:252 572:252 573:252 574:252 575:254 576:252 577:252 578:252 579:178 580:91 581:33 594:13 595:119 596:252 597:252 598:252 599:252 600:252 601:252 602:252 603:254 604:252 605:252 606:184 607:37 623:13 624:121 625:252 626:252 627:252 628:252 629:252 630:252 631:254 632:214 633:45 634:28 652:8 653:21 654:21 655:169 656:252 657:252 658:41 659:22 660:18 +0 125:218 126:253 127:253 128:255 129:149 130:62 151:42 152:144 153:236 154:251 155:251 156:253 157:251 158:236 159:144 160:144 179:99 180:251 181:251 182:251 183:225 184:253 185:251 186:251 187:251 188:251 189:166 190:16 206:79 207:253 208:251 209:251 210:204 211:41 212:143 213:205 214:251 215:251 216:251 217:253 218:169 219:15 233:79 234:231 235:253 236:251 237:225 238:41 241:41 242:226 243:251 244:251 245:253 246:251 247:164 260:37 261:253 262:253 263:255 264:253 265:35 271:79 272:232 273:255 274:253 275:227 276:42 288:140 289:251 290:251 291:253 292:168 293:15 300:77 301:253 302:251 303:251 304:142 315:21 316:221 317:251 318:251 319:164 320:15 329:227 330:251 331:251 332:236 333:61 342:32 343:190 344:251 345:251 346:251 357:73 358:251 359:251 360:251 361:71 370:73 371:251 372:251 373:251 374:251 385:73 386:251 387:251 388:251 389:71 398:73 399:253 400:253 401:253 402:201 413:73 414:253 415:253 416:253 417:72 426:176 427:251 428:251 429:251 430:71 441:73 442:251 443:251 444:251 445:71 454:253 455:251 456:251 457:157 458:10 469:73 470:251 471:251 472:251 473:71 482:253 483:251 484:251 485:142 497:150 498:251 499:251 500:204 501:41 510:124 511:251 512:251 513:220 514:180 524:130 525:253 526:251 527:225 528:41 538:73 539:253 540:253 541:253 542:253 543:73 544:73 545:10 549:42 550:73 551:150 552:253 553:255 554:253 555:216 566:31 567:189 568:251 569:251 570:251 571:253 572:251 573:159 574:144 575:144 576:145 577:206 578:251 579:251 580:251 581:253 582:168 583:92 595:20 596:195 597:251 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:251 607:251 608:225 609:164 610:15 624:21 625:142 626:220 627:253 628:251 629:251 630:251 631:251 632:253 633:251 634:251 635:204 636:41 654:51 655:72 656:174 657:251 658:251 659:251 660:253 661:147 662:71 663:41 +0 127:60 128:96 129:96 130:48 153:16 154:171 155:228 156:253 157:251 158:220 159:51 160:32 181:127 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:80 207:24 208:182 209:236 210:251 211:211 212:189 213:236 214:251 215:251 216:251 217:242 218:193 234:100 235:194 236:251 237:251 238:211 239:35 241:71 242:173 243:251 244:251 245:253 246:240 247:158 248:19 261:64 262:253 263:255 264:253 265:205 266:19 271:40 272:218 273:255 274:253 275:253 276:91 288:16 289:186 290:251 291:253 292:247 293:110 300:39 301:233 302:251 303:251 304:188 315:16 316:189 317:251 318:251 319:205 320:110 329:48 330:220 331:251 332:220 333:48 343:72 344:251 345:251 346:251 347:158 358:51 359:251 360:251 361:232 371:190 372:251 373:251 374:251 375:59 386:32 387:251 388:251 389:251 398:96 399:253 400:253 401:253 402:95 414:32 415:253 416:253 417:193 426:214 427:251 428:251 429:204 430:23 442:52 443:251 444:251 445:94 454:253 455:251 456:251 457:109 469:48 470:221 471:251 472:219 473:47 482:253 483:251 484:251 485:70 497:234 498:251 499:251 500:188 510:253 511:251 512:251 513:188 523:40 524:158 525:253 526:251 527:172 528:70 539:191 540:253 541:253 542:253 543:96 544:24 549:12 550:174 551:253 552:253 553:255 554:221 567:71 568:251 569:251 570:251 571:253 572:205 573:190 574:190 575:190 576:191 577:197 578:251 579:251 580:231 581:221 582:93 595:16 596:126 597:236 598:251 599:253 600:251 601:251 602:251 603:251 604:253 605:251 606:251 607:140 608:47 625:67 626:188 627:189 628:188 629:188 630:188 631:188 632:189 633:188 634:109 635:4 +0 126:32 127:202 128:255 129:253 130:253 131:175 132:21 152:84 153:144 154:190 155:251 156:253 157:251 158:251 159:251 160:174 176:6 177:37 178:166 179:218 180:236 181:251 182:251 183:251 184:253 185:251 186:251 187:251 188:251 189:156 204:115 205:251 206:251 207:253 208:251 209:251 210:251 211:251 212:253 213:251 214:251 215:251 216:251 217:180 231:105 232:241 233:251 234:251 235:253 236:251 237:251 238:251 239:122 240:72 241:71 242:71 243:148 244:251 245:180 258:73 259:253 260:253 261:253 262:253 263:202 264:253 265:253 266:143 286:31 287:189 288:251 289:251 290:251 291:31 292:189 293:251 294:142 314:63 315:236 316:251 317:251 318:96 320:124 321:246 322:142 330:21 331:166 332:21 342:73 343:251 344:251 345:251 346:71 349:217 350:142 357:32 358:190 359:251 360:142 370:73 371:251 372:251 373:251 374:71 377:217 378:142 385:73 386:251 387:251 388:142 398:73 399:253 400:253 401:253 402:72 405:156 406:103 413:73 414:253 415:253 416:253 417:72 426:73 427:251 428:251 429:251 430:174 441:73 442:251 443:251 444:251 445:71 454:73 455:251 456:251 457:251 458:251 469:73 470:251 471:251 472:251 473:71 482:42 483:205 484:251 485:251 486:251 487:79 497:73 498:251 499:251 500:251 501:71 511:41 512:226 513:251 514:251 515:232 516:77 525:73 526:251 527:251 528:251 529:71 540:166 541:253 542:253 543:255 544:253 545:227 546:73 547:21 553:125 554:253 555:253 556:143 568:16 569:169 570:251 571:253 572:251 573:251 574:251 575:174 576:105 579:63 580:144 581:253 582:251 583:251 584:142 597:15 598:35 599:253 600:251 601:251 602:251 603:251 604:243 605:217 606:217 607:231 608:251 609:253 610:251 611:220 612:20 627:143 628:142 629:236 630:251 631:251 632:253 633:251 634:251 635:251 636:251 637:253 638:251 639:137 657:61 658:71 659:200 660:253 661:251 662:251 663:251 664:251 665:201 666:71 667:10 +1 130:218 131:170 132:108 157:32 158:227 159:252 160:232 185:129 186:252 187:252 188:252 212:1 213:253 214:252 215:252 216:168 240:144 241:253 242:252 243:236 244:62 268:144 269:253 270:252 271:215 296:144 297:253 298:252 299:112 323:21 324:206 325:253 326:252 327:71 351:99 352:253 353:255 354:119 378:63 379:242 380:252 381:253 382:35 406:94 407:252 408:252 409:154 410:10 433:145 434:237 435:252 436:252 461:255 462:253 463:253 464:108 487:11 488:155 489:253 490:252 491:179 492:15 514:11 515:150 516:252 517:253 518:200 519:20 542:73 543:252 544:252 545:253 546:97 569:47 570:233 571:253 572:253 596:1 597:149 598:252 599:252 600:252 624:1 625:252 626:252 627:246 628:132 652:1 653:169 654:252 655:132 +1 130:116 131:255 132:123 157:29 158:213 159:253 160:122 185:189 186:253 187:253 188:122 213:189 214:253 215:253 216:122 241:189 242:253 243:253 244:122 267:2 268:114 269:243 270:253 271:186 272:19 295:100 296:253 297:253 298:253 299:48 323:172 324:253 325:253 326:253 327:48 351:172 352:253 353:253 354:182 355:19 378:133 379:251 380:253 381:175 382:4 405:107 406:251 407:253 408:253 409:65 432:26 433:194 434:253 435:253 436:214 437:40 459:105 460:205 461:253 462:253 463:125 464:40 487:139 488:253 489:253 490:253 491:81 514:41 515:231 516:253 517:253 518:159 519:16 541:65 542:155 543:253 544:253 545:172 546:4 569:124 570:253 571:253 572:253 573:98 597:124 598:253 599:253 600:214 601:41 624:22 625:207 626:253 627:253 628:139 653:124 654:253 655:162 656:9 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_linear_regression_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_linear_regression_data.txt new file mode 100644 index 0000000..9aaaa42 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_linear_regression_data.txt @@ -0,0 +1,501 @@ +-9.490009878824548 1:0.4551273600657362 2:0.36644694351969087 3:-0.38256108933468047 4:-0.4458430198517267 5:0.33109790358914726 6:0.8067445293443565 7:-0.2624341731773887 8:-0.44850386111659524 9:-0.07269284838169332 10:0.5658035575800715 +0.2577820163584905 1:0.8386555657374337 2:-0.1270180511534269 3:0.499812362510895 4:-0.22686625128130267 5:-0.6452430441812433 6:0.18869982177936828 7:-0.5804648622673358 8:0.651931743775642 9:-0.6555641246242951 10:0.17485476357259122 +-4.438869807456516 1:0.5025608135349202 2:0.14208069682973434 3:0.16004976900412138 4:0.505019897181302 5:-0.9371635223468384 6:-0.2841601610457427 7:0.6355938616712786 8:-0.1646249064941625 9:0.9480713629917628 10:0.42681251564645817 +-19.782762789614537 1:-0.0388509668871313 2:-0.4166870051763918 3:0.8997202693189332 4:0.6409836467726933 5:0.273289095712564 6:-0.26175701211620517 7:-0.2794902492677298 8:-0.1306778297187794 9:-0.08536581111046115 10:-0.05462315824828923 +-7.966593841555266 1:-0.06195495876886281 2:0.6546448480299902 3:-0.6979368909424835 4:0.6677324708883314 5:-0.07938725467767771 6:-0.43885601665437957 7:-0.608071585153688 8:-0.6414531182501653 9:0.7313735926547045 10:-0.026818676347611925 +-7.896274316726144 1:-0.15805658673794265 2:0.26573958270655806 3:0.3997172901343442 4:-0.3693430998846541 5:0.14324061105995334 6:-0.25797542063247825 7:0.7436291919296774 8:0.6114618853239959 9:0.2324273700703574 10:-0.25128128782199144 +-8.464803554195287 1:0.39449745853945895 2:0.817229160415142 3:-0.6077058562362969 4:0.6182496334554788 5:0.2558665508269453 6:-0.07320145794330979 7:-0.38884168866510227 8:0.07981886851873865 9:0.27022202891277614 10:-0.7474843534024693 +2.1214592666251364 1:-0.005346215048158909 2:-0.9453716674280683 3:-0.9270309666195007 4:-0.032312290091389695 5:0.31010676221964206 6:-0.20846743965751569 7:0.8803449313707621 8:-0.23077831216541722 9:0.29246395759528565 10:0.5409312755478819 +1.0720117616524107 1:0.7880855916368177 2:0.19767407429003536 3:0.9520689432368168 4:-0.845829774129496 5:0.5502413918543512 6:-0.44235539500246457 7:0.7984106594591154 8:-0.2523277127589152 9:-0.1373808897290778 10:-0.3353514432305029 +-13.772441561702871 1:-0.3697050572653644 2:-0.11452811582755928 3:-0.807098168238352 4:0.4903066124307711 5:-0.6582805242342049 6:0.6107814398427647 7:-0.7204208094262783 8:-0.8141063661170889 9:-0.9459402662357332 10:0.09666938346350307 +-5.082010756207233 1:-0.43560342773870375 2:0.9349906440170221 3:0.8090021580031235 4:-0.3121157071110545 5:-0.9718883630945336 6:0.6191882496201251 7:0.0429886073795116 8:0.670311110015402 9:0.16692329718223786 10:0.37649213869502973 +7.887786536531237 1:0.11276440263810383 2:-0.7684997525607482 3:0.1770172737885798 4:0.7902845707138706 5:0.2529503304079441 6:-0.23483801763662826 7:0.8072501895004851 8:0.6673992021927047 9:-0.4796127376677324 10:0.9244724404994455 +14.323146365332388 1:-0.2049276879687938 2:0.1470694373531216 3:-0.48366999792166787 4:0.643491115907358 5:0.3183669486383729 6:0.22821350958477082 7:-0.023605251086149304 8:-0.2770587742156372 9:0.47596326458377436 10:0.7107229819632654 +-20.057482615789212 1:-0.3205057828114841 2:0.51605972926996 3:0.45215640988181516 4:0.01712446974606241 5:0.5508198371849293 6:-0.2478254241316491 7:0.7256483175955235 8:0.39418662792516 9:-0.6797384914236382 10:0.6001217520150142 +-0.8995693247765151 1:0.4508991072414843 2:0.589749448443134 3:0.6464818311502738 4:0.7005669004769028 5:0.9699584106930381 6:-0.7417466269908464 7:0.22818964839784495 8:0.08574936236270037 9:-0.6945765138377225 10:0.06915201979238828 +-19.16829262296376 1:0.09798746565879424 2:-0.34288007110901964 3:0.440249350802451 4:-0.22440768392359534 5:-0.9695067570891225 6:-0.7942032659310758 7:-0.792286205517398 8:-0.6535487038528798 9:0.7952676470618951 10:-0.1622831617066689 +5.601801561245534 1:0.6949189734965766 2:-0.32697929564739403 3:-0.15359663581829275 4:-0.8951865090520432 5:0.2057889391931318 6:-0.6676656789571533 7:-0.03553655732400762 8:0.14550349954571096 9:0.034600542078191854 10:0.4223352065067103 +-3.2256352187273354 1:0.35278245969741096 2:0.7022211035026023 3:0.5686638754605697 4:-0.4202155290448111 5:-0.26102723928249216 6:0.010688215941416779 7:-0.4311544807877927 8:0.9500151672991208 9:0.14380635780710693 10:-0.7549354840975826 +1.5299675726687754 1:-0.13079299081883855 2:0.0983382230287082 3:0.15347083875928424 4:0.45507300685816965 5:0.1921083467305864 6:0.6361110540492223 7:0.7675261182370992 8:-0.2543488202081907 9:0.2927051050236915 10:0.680182444769418 +-0.250102447941961 1:-0.8062832278617296 2:0.8266289890474885 3:0.22684501241708888 4:0.1726291966578266 5:-0.6778773666126594 6:0.9993906921393696 7:0.1789490173139363 8:0.5584053824232391 9:0.03495894704368174 10:-0.8505720014852347 +12.792267926563595 1:-0.008461200645088818 2:-0.648273596036564 3:-0.005334477339629995 4:0.3781469006858833 5:0.30565234666790686 6:-0.2822867492866177 7:0.10175120738413801 8:0.5342432888482425 9:0.05146513075475534 10:-0.6459729964194652 +6.082192787194888 1:0.42519013450094767 2:0.09441503345243984 3:-0.07898439043103522 4:-0.32207498048636474 5:-0.9180071861219266 6:0.5951317320731633 7:0.41000814588717693 8:-0.3926260640533046 9:0.2789036768568971 10:0.13163692286014528 +-7.481405271455238 1:0.03324842612749346 2:0.07055844751995122 3:-0.47199515597021113 4:-0.682690342465275 5:0.3983414713797069 6:-0.2136729393256811 7:-0.09066563475481249 8:-0.4640338194317184 9:-0.03513782089224482 10:-0.1711809802758364 +6.739533816100517 1:0.1774546460228057 2:-0.6783644553523549 3:-0.47871398278230504 4:0.02272121490463097 5:-0.5047649289302389 6:0.26479596144873896 7:-0.32045436544054096 8:0.3113047940487379 9:0.6269418147567556 10:0.9710114516962312 +3.780807062175497 1:0.01715676997104909 2:0.8975962429865936 3:-0.46594560920034134 4:0.2873623499953055 5:0.8894362304584083 6:0.17973981232418468 7:0.49105791400707743 8:-0.7359842740294882 9:0.38941133808001127 10:-0.7151884777228046 +4.564039393483412 1:0.07478785545033317 2:-0.8672651994084235 3:0.450599300176334 4:0.35104802298560056 5:0.6797318185095045 6:-0.03891997518827006 7:-0.33208695871398675 8:0.6166574577055226 9:0.5730212324012205 10:-0.4194925751047054 +-0.3195679646035633 1:0.054527683864544096 2:-0.15591931640565093 3:0.9266742559542833 4:0.888522581905147 5:0.6576203900699167 6:0.6417770212400336 7:0.7509788029052338 8:-0.3104974571382815 9:0.7234744267051683 10:-0.15869049651427103 +11.290452658023497 1:0.20173310976772196 2:0.8657502566551409 3:0.9325160601080682 4:0.24570884032596263 5:-0.6546108813337841 6:-0.14020032028377583 7:-0.8825687891702743 8:-0.21420166926412865 9:-0.8600275184792756 10:-0.7990574622230739 +-4.003499192090455 1:0.8325875503351796 2:-0.5956350140619129 3:0.12598048009007923 4:0.12340188733473134 5:-0.839435659309717 6:-0.16623481818728414 7:0.12028795301041662 8:-0.7994713170657952 9:0.2216721974907896 10:0.8407561415075087 +-19.872991038068406 1:-0.9325810772922609 2:-0.6411471147334535 3:0.9949216290375054 4:0.483048267470493 5:-0.8736297429070232 6:-0.36222771685582544 7:0.26397860162786957 8:0.45527588775737704 9:-0.9424989711186325 10:0.6251162293059616 +10.502762149373098 1:-0.2307778924009991 2:0.6977871128979924 3:0.022830408261390822 4:0.6257738824362347 5:0.9770979848265122 6:0.09985730624684575 7:-0.9755858424230182 8:-0.689969833240031 9:-0.7294587311376761 10:0.3496326193951331 +-14.328978509075442 1:0.37929821892417404 2:0.8402056881660709 3:-0.1806835799958202 4:0.766314307210441 5:0.865876513623024 6:-0.7113501219432434 7:-0.0932956557986735 8:-0.7042025810921411 9:0.47530696925672267 10:-0.4629102077669889 +-16.26143027545273 1:-0.9309578475799722 2:0.7591795880911123 3:0.06296957473213705 4:0.786790093290086 5:-0.9527998391625465 6:-0.08573982501921895 7:-0.3812232026687308 8:-0.6890669703685022 9:0.25415911467755015 10:-0.07664746267502509 +11.772544195529013 1:0.3614756404325046 2:0.14508027508253818 3:0.23042774014795753 4:0.4164348685332022 5:0.4109091750657461 6:0.03853098236933272 7:0.38911994885223145 8:-0.5031309357181766 9:-0.596467768575587 10:0.17884522225228028 +14.697703557439503 1:0.24508864174863 2:0.7576193329655578 3:0.09030511120334461 4:0.9537528991778741 5:-0.7224092160621338 6:-0.34089385162121943 7:0.6924170720838818 8:0.32912306214891784 9:-0.4064624712125904 10:-0.5344662061201593 +-13.976130931152703 1:0.5891192531479754 2:0.29862103742464274 3:-0.36153976712796343 4:-0.6552669564323226 5:-0.22672513691161766 6:0.3001336202535376 7:0.34490251346382617 8:0.2072633053920192 9:-0.5659371284058774 10:0.49599636156628835 +-14.762758252931127 1:0.31302496164254223 2:-0.6062773982342133 3:-0.9874007658402217 4:-0.6214904627601421 5:-0.11421073677207683 6:-0.5850843421161205 7:0.1250679146774638 8:-0.7108170726393621 9:-0.6888351241194393 10:0.6077343683084389 +-3.300641320608255 1:-0.1407178879203672 2:0.12960233233004925 3:-0.4236196478321872 4:0.7903078296084356 5:-0.8755754953628643 6:-0.2062360260394529 7:-0.045680124889026175 8:0.783182093429277 9:-0.02995737262668463 10:-0.33382351650328435 +-15.72351561304857 1:-0.1802575775708093 2:-0.991006951265341 3:-0.9107951763247621 4:0.9069820084047908 5:-0.12691921206803047 6:-0.7087012119383593 7:-0.9179510577925369 8:0.18480349982718325 9:-0.4478459144114004 10:-0.5560585660624608 +-22.949825936196074 1:0.4797855980916854 2:0.01997502546020402 3:-0.8827928315487465 4:0.2755107907750989 5:0.015544482147298977 6:0.9652687138748801 7:0.6622667860970648 8:-0.7708138539912186 9:0.17728148663006627 10:0.47818190728952925 +12.092431628826905 1:0.1358843437335564 2:0.03643446587894239 3:-0.31070823939673287 4:0.5283033206569152 5:0.3469111543845367 6:-0.5162518174930761 7:0.24270234207184016 8:0.7352292800096338 9:0.8860322286740037 10:0.6748068653962045 +-23.51088409032297 1:-0.4683538422180036 2:0.1469540185936138 3:0.9113612952591796 4:-0.9838482669789823 5:0.4506466371133697 6:0.6456121712599778 7:0.8264783725578371 8:0.562664168655115 9:-0.8299281852090683 10:0.40690300256653256 +5.998186124881712 1:-0.9781302074883151 2:0.32984303335155785 3:0.7303430847899663 4:0.841481297188956 5:0.05580773881989276 6:0.7130788298702062 7:-0.218087116119847 8:-0.9889494995220598 9:0.9182854134226501 10:-0.7501751701020942 +9.852316338642547 1:0.146854160091757 2:-0.3611508707370965 3:0.3517016971654914 4:0.6187697988029395 5:-0.010768583697787548 6:0.5236725885871243 7:0.5945666964145524 8:-0.009180562740628506 9:-0.44474762415618274 10:0.41852743519493685 +-5.313930756588526 1:-0.6304209277071555 2:-0.37010359785263813 3:-0.3194739026510125 4:-0.750533359080716 5:0.45500303301733114 6:-0.012727544364283805 7:-0.43941651856862274 8:0.927108876532093 9:-0.24164903158058149 10:0.44134972919002124 +-4.2775224863223915 1:-0.35785764991284363 2:0.942797043714243 3:0.4539569191274251 4:-0.6944903010994341 5:-0.08357221983075225 6:0.4433049548665855 7:-0.5488972050023557 8:-0.24014623658145773 9:-0.6178118485382511 10:-0.4575463952834564 +-10.57769830424322 1:0.22693864400257335 2:-0.041639691095668674 3:0.9948726461115123 4:-0.7450471554938383 5:-0.1114847126717804 6:-0.27881184842402673 7:0.029766812446276214 8:-0.3727649352432578 9:-0.7791732805568077 10:0.9425576681069683 +-0.8430338600258201 1:0.4607090007225536 2:-0.6079961642969514 3:-0.5671626932935381 4:0.12784576080614185 5:-0.30766031989910236 6:-0.21232963505711555 7:0.3310463755850872 8:-0.6807682731528943 9:0.7826634145951483 10:0.0608057623636995 +13.450586257053727 1:-0.2697769964284986 2:0.07743737732312428 3:-0.8459687499864881 4:0.6091901514177853 5:-0.9464815428211699 6:0.15780407422581533 7:-0.28552052619478996 8:-0.27500859181806403 9:-0.7207541548282903 10:0.05215593729084533 +20.358241877831016 1:0.29768927445620164 2:-0.5379390525163252 3:0.6591913001003027 4:0.6635992348010928 5:0.3786594651413009 6:-0.7217135278882543 7:0.9634013908615768 8:0.03961253903778861 9:0.1335121312144949 10:0.7933944303463509 +9.800993960518852 1:0.39896823489212285 2:0.30948413101894023 3:0.08568060094378493 4:-0.7454513450113371 5:0.8054125831421357 6:-0.24464240413169347 7:-0.18294406588625112 8:-0.883455504399858 9:0.2468431033653562 10:-0.708151566382103 +-21.432387764165806 1:-0.4785033857256795 2:0.520350718059089 3:-0.2988515012130126 4:-0.46260150057299754 5:0.5394344995663083 6:0.39320468081626836 7:0.1890560923345248 8:0.13123799325264507 9:0.43613839380760355 10:0.39541998419731494 +-4.090570760187878 1:0.3909705814857716 2:0.9830271975811611 3:0.672523651785939 4:0.0035177223850744177 5:0.567082732451311 6:-0.2620454326881394 7:0.46622578556708105 8:0.646246879249865 9:0.4263175536668733 10:0.8982696975276223 +3.7459201216906926 1:-0.9480167656870653 2:-4.888270196095057E-4 3:0.48226844071577646 4:-0.23706663537631645 5:0.22420266627462127 6:0.2981747607694978 7:0.3893425967975348 8:0.6302701381298614 9:-0.21909113816064196 10:0.8371697958140494 +9.767952084958061 1:-0.2300790371078303 2:-0.4457883630748676 3:0.28710853302295325 4:0.7112839743052013 5:-0.8765858382640623 6:-0.6470779468607217 7:0.4369262584371727 8:-0.7175412028407337 9:0.5506733477278882 10:0.5393007189573547 +6.9802839308913365 1:0.21769855012808215 2:0.8653818331675485 3:0.2322943113578111 4:0.3760591265797468 5:0.06554014167292377 6:0.6866096712933549 7:0.866929973115441 8:-0.6462263417217329 9:0.2507247465275353 10:-0.7005877782050307 +16.014720800069103 1:0.6058055248984549 2:0.048517868234337014 3:-0.15744912875924877 4:0.32598079708869365 5:-0.587791997223768 6:-0.4636187312118474 7:0.7771908559246068 8:-0.349403853888719 9:0.229800030145503 10:-0.674614818934488 +8.417571532985823 1:-0.21164946152466801 2:-0.9981936663594053 3:0.8611869575187896 4:0.11100891297254312 5:-0.7406067304729631 6:-0.7613837395522254 7:-0.9617573325708704 8:0.5697426971647488 9:-0.5830879716990833 10:0.5951448538064159 +-12.491442077546413 1:-0.19172117564625735 2:-0.12421304883392126 3:0.7095605786791346 4:0.6401582292398038 5:-0.9347790209840108 6:0.6592209285686903 7:0.702282297844389 8:-0.22765902007749528 9:-0.17746922342943816 10:0.7196663432778121 +-8.605713514762092 1:0.36490454976480846 2:0.6991204480538957 3:0.6546945560337121 4:-0.032324845758738174 5:0.2453935969836043 6:0.5363119225093116 7:0.6266741350524205 8:-0.2132266305382322 9:-0.308105870487996 10:-0.08219413867616465 +-10.35591860037468 1:-0.014204168485027147 2:-0.7077035677144325 3:0.024004217785642767 4:0.818971992516166 5:0.9081305263471056 6:0.808854493237229 7:-0.6474336785461867 8:-0.32559288177031465 9:-0.32850453072496055 10:-0.7035310416695784 +3.605002621628445 1:0.6085817977516599 2:0.8101072412357928 3:0.7697891508923966 4:-0.5738750389864677 5:-0.734314989863889 6:-0.7879014492215499 7:0.6884442838920775 8:-0.46131231930402383 9:-0.7730585954271005 10:-0.7819874019145132 +12.30435312415091 1:0.3283668768730639 2:-0.18316686990068187 3:0.3955614099142126 4:0.8450470350842108 5:0.3958042901611589 6:0.6578475571960676 7:-0.4395488558075096 8:0.15720430113495376 9:-0.5318362828977672 10:0.45585285255232044 +9.020048819638827 1:-0.5986521145193395 2:0.3266542215286443 3:-0.09911773729611917 4:-0.21478254478908676 5:0.6546175049764293 6:-0.1414796368932345 7:0.25802631337510085 8:-0.6773828562539816 9:-0.22038193899258718 10:-0.17465737306657902 +14.854262978981406 1:0.5293763924477841 2:-0.24658868331583683 3:0.8268631648872109 4:0.8969207203400265 5:0.03933229861213983 6:-0.6212951181360529 7:-0.36695460282178205 8:-0.5468014636386027 9:-0.3419492829414976 10:-0.8273314086998671 +5.658665647926016 1:0.9543096383762801 2:0.13230023957687176 3:-0.3071929861496465 4:-0.3646067841449696 5:0.6979929890816723 6:-0.20721664168809228 7:0.6676482547655365 8:0.944757051233543 9:0.024377296173674567 10:-0.9413728609667691 +-6.930603551528371 1:0.09198647857985232 2:-0.3685113649452161 3:-0.2361728930325453 4:0.3674268130607439 5:0.27385598384498344 6:-0.7151900241735676 7:0.3310154476154119 8:-0.24328111897361682 9:0.2511378679668912 10:-0.35825141175578934 +13.361196783041926 1:0.11676665169094824 2:-0.49968608916548307 3:0.9941342810313298 4:-0.17858967215374988 5:0.1993744673440312 6:0.14596837574280297 7:-0.8245495433125194 8:-0.5637934691545672 9:-0.8589185806222286 10:-0.4923216901915597 +-3.280508467210429 1:-0.9917770074538397 2:-0.1547651813493751 3:0.621733177563484 4:0.7303326279246298 5:-0.0786900332560696 6:0.9107127797641994 7:0.7104513024299466 8:-0.32858522942354407 9:0.17013652749847386 10:0.27656984316288824 +11.13509519160867 1:0.6874932143640391 2:-0.46610293161038907 3:0.8744681017967024 4:0.40900365224695956 5:-0.49770054448432055 6:-0.0635770754462921 7:-0.5705387648707747 8:-0.577988250149829 9:-0.8099463063934682 10:0.42132700180827354 +-11.857350365429426 1:-0.24607974991258308 2:-0.943388538022258 3:0.8679112109377674 4:0.7779951176637694 5:-0.5802336023276593 6:-0.9093352471884992 7:0.29337797938742316 8:0.498519874589175 9:0.3493034812120912 10:-0.07307210651399076 +11.421632138263703 1:0.3911519359353859 2:-0.8154393787235621 3:0.47194271125243237 4:0.14014792298759593 5:-0.3589345913619957 6:0.7887695409762479 7:0.49962792312858895 8:-0.6402670146359797 9:-0.2314041601683119 10:-0.798901341175887 +5.194792012146463 1:0.810279303469398 2:-0.9772756877199589 3:-0.20925958437085557 4:0.8797562461102444 5:0.3211532423260066 6:0.25250279470783754 7:0.14387831263435813 8:-0.021466789385169882 9:0.18909293657271564 10:-0.5981349964027893 +12.242677118499806 1:0.3565715672082048 2:0.7366743237221687 3:0.1922233582434527 4:-0.3551925780624561 5:0.5290849503909634 6:0.7744214641246749 7:0.7277215028580597 8:-0.590440215391044 9:0.7427328184290733 10:-0.6231904162251609 +3.496172341296411 1:0.5028717258135624 2:-0.5838871888624848 3:-0.5540116561110324 4:0.8502487679795261 5:-0.7983061034328727 6:-0.3853123296389005 7:-0.1493800684643869 8:0.6008798629354264 9:-0.32299062155495406 10:-0.5827019502242026 +-15.437384793431217 1:0.41994681418237345 2:0.7106426870657483 3:-0.45211033467567696 4:-0.7272406549392239 5:-0.35736594496490737 6:0.4764507578985955 7:-0.5249912641281373 8:0.8562010912051132 9:0.45927621623833637 10:-0.3701817429794385 +5.490036861541498 1:0.8414999442459015 2:0.9273442862476728 3:-0.054654787893199774 4:-0.23126134156257327 5:-0.9155048245317694 6:0.25750538376376975 7:-0.8470916763665326 8:0.9105674676753848 9:0.5026028522378054 10:-0.06650501561108468 +-1.074065343287859 1:0.37484830603001607 2:-0.9858854245832975 3:0.007159356555897611 4:0.8172796295244154 5:0.519147377529164 6:0.8211049991970722 7:0.9901658817979146 8:-0.026519560032641998 9:-0.2328762488733862 10:0.43161994187258035 +2.0482082496444622 1:0.24940246021565793 2:0.47248358864259177 3:0.23833814894291105 4:-0.3449172512379757 5:0.7412869866239866 6:0.1351422898741914 7:-0.003784141556894216 8:-0.6321917152754075 9:0.8246267827865776 10:0.5057520480449009 +16.709794859608397 1:-0.5977424405191092 2:-0.13991362149785713 3:0.613487896720806 4:-0.37681525320218157 5:-0.4369592282569783 6:0.4702242879506955 7:0.07498463532645339 8:-0.9942304127133292 9:0.41304209196175257 10:0.6799250665519481 +4.598881854940949 1:-0.41212838137243835 2:0.6737124633791323 3:0.8376369191216593 4:0.2848328781926128 5:-0.17960265353296 6:0.0035488712665472377 7:-0.8355355482928055 8:-0.7439716673142398 9:-0.009043467128117433 10:0.7423272515054122 +9.566038608555402 1:-0.662329643040616 2:0.4727113884417973 3:-0.15734218732411365 4:-0.3950754785173889 5:0.13837083076070011 6:0.633261314089351 7:0.9929998062307679 8:-0.4639028424346423 9:-0.073992579817449 10:0.3413166410117088 +1.629198477883475 1:-0.2875719791707101 2:0.9395753700232541 3:-0.45090801750966314 4:-0.384528069378699 5:-0.35937736478702753 6:0.9597102694501136 7:-0.6898325123180971 8:-0.11436012866371303 9:-0.5330550575952768 10:0.24688769932037258 +-7.374620970147229 1:0.16864051681940984 2:-0.08391828256018252 3:-0.8184503043836224 4:0.5461252511055263 5:0.7264676659099087 6:-0.9654384426822686 7:-0.8537533138667612 8:0.9189716013058653 9:-0.03449322582531389 10:0.5490329745887035 +-0.5741704240890674 1:0.9392753294760656 2:-0.5579682000156501 3:-0.8083270703362093 4:-0.7022804026958895 5:-0.30426803430649896 6:0.8211432527140852 7:-0.8101343265051797 8:-0.0945946325760949 9:0.49546915718101814 10:0.5184327698839013 +12.583032451116004 1:0.20496323995364651 2:0.5082017540304999 3:0.2428646053751764 4:0.7101854338863274 5:-0.9619925264660094 6:0.4610134502825909 7:-0.5620669052678122 8:0.6766614078376236 9:-0.7169693435782278 10:-0.14362322382035164 +-10.489157123372898 1:-0.7441633083637054 2:0.07069898351187809 3:-0.47119552972566336 4:-0.43970155900871344 5:0.43192289605353973 6:-0.0798550143899397 7:0.2111188135787776 8:0.9101748615761336 9:-0.4079984876629721 10:-0.8101424982394589 +-3.811365493249739 1:0.7250263461647963 2:0.22182621035333838 3:-0.12735342714215725 4:0.26222861719040624 5:0.3928174057935714 6:0.817131411734006 7:-0.056109765698795 8:0.7908779197353637 9:-0.06768319505245768 10:0.4107045608924882 +-7.604636483513961 1:0.876751634787073 2:0.04037085575852295 3:0.18142385658771398 4:0.38350565074271903 5:-0.30937664332011905 6:-0.9544807672006823 7:0.008643477632712449 8:-0.27676843472226276 9:-0.12938540988602476 10:-0.2929762262661819 +-1.9889499615051784 1:-0.4243149295090465 2:0.22578711943818686 3:0.662530786460152 4:0.28592235843136105 5:0.4170345231441832 6:0.9290881132120887 7:0.5332443368002588 8:-0.33248958421809927 9:0.16273139830495942 10:0.6899022585936985 +-1.99891354174786 1:-0.1732078452611825 2:0.2361029542296429 3:-0.8455867017505336 4:0.31638672033240867 5:-0.648387667144986 6:-0.7647886103837449 7:0.6910155501192978 8:-0.2665663102538198 9:-0.5980899570876459 10:-0.9165896495676276 +9.74348630903265 1:0.18934450539532244 2:-0.715110505416745 3:-0.453777527810155 4:0.2743741252197758 5:-0.8439310405443103 6:-0.533835190276116 7:-0.5911710854054728 8:0.21026462628920695 9:-0.45849607678093585 10:0.1327074179200407 +20.221961806051706 1:0.624731930687735 2:-0.39914395421723015 3:0.781887900750925 4:0.5442619051596436 5:0.16651193067479153 6:0.9064846121246533 7:-0.3643159594276202 8:-0.5182065337246469 9:-0.6785628247191553 10:0.7111152852903913 +20.456947955410897 1:-0.21923785332346513 2:0.11340668617783778 3:0.7397883986253251 4:-0.11748081084695605 5:0.06314872700777197 6:-0.7124574845946587 7:0.18043581960897104 8:-0.09023925260092103 9:-0.7256417560118238 10:-0.5038088673851804 +12.241006086129564 1:-0.15271598143132215 2:0.9038942665552285 3:-0.6168514099878155 4:-0.12219038322317011 5:0.5402785935596728 6:0.4059744401803913 7:0.258870596734184 8:0.3190881033039108 9:0.2372469007313076 10:0.367188299614863 +3.980473021620311 1:-0.9025895351376971 2:-0.03333947011476446 3:-0.8220776066161464 4:0.449117985679933 5:0.9970519437779266 6:0.27430911004640457 7:0.039081352882204046 8:-0.8621514950929796 9:-0.569587565933642 10:-0.9118346349929578 +-13.420594775890757 1:0.3697979495309094 2:0.07383664120111888 3:0.7199366131785143 4:0.2118625428869032 5:-0.9015976323216077 6:-0.5298395275757712 7:-0.9517419542156635 8:0.39554920787574743 9:-0.3721957439110324 10:-0.4750272836396878 +-1.052659359353786 1:0.02106845330888185 2:0.7571245678782959 3:0.8034228830223251 4:0.32968340513846917 5:-0.6510386482911554 6:0.2710115488605187 7:-0.1319580272290235 8:0.932600992666184 9:0.8260461527035414 10:-0.8507648952138052 +9.813440129324034 1:0.41048687946340134 2:0.9384639988086239 3:0.4569555844323441 4:-0.3084729082645552 5:-0.7299010284877061 6:-0.6925012997779212 7:-0.6798013915257548 8:-0.504368104320321 9:-0.6234398059664716 10:0.8633407902005543 +-2.8942782378157714 1:0.5546381825677706 2:0.7959405841824887 3:0.584699836289184 4:-0.5726371777829862 5:-0.2827976152663936 6:0.138034013875719 7:-0.2935080791661324 8:-0.5323479091625714 9:0.6837641044797451 10:0.5986680812032501 +8.562937733537664 1:0.14753220510180776 2:-0.31591341855048327 3:-0.748545617199091 4:0.3251888821665734 5:0.8228589483149358 6:0.046659706976506676 7:-0.35049927996132624 8:0.2953170004605874 9:-0.6429374177050204 10:0.4624083116836044 +13.413187970975178 1:-0.7229883396779724 2:0.8876940454894067 3:-0.033794226589695775 4:0.46700071356381523 5:0.1599557295166274 6:-0.8944619785248653 7:-0.1258464584151997 8:-0.8797551785991506 9:-0.14408879184669354 10:0.11020655997336015 +-5.491389764900794 1:-0.366507395597937 2:0.630480481240723 3:-0.16600801981741609 4:0.09842042773854076 5:0.30129535029579047 6:0.14102166298628882 7:-0.28131788612036623 8:0.49635295715686234 9:0.0625636989631968 10:-0.41748132718912 +-10.29566593602992 1:-0.7898597726154271 2:-0.05425577320946573 3:0.5992645759265662 4:-0.4716868549309716 5:-0.020137302700854676 6:0.6216515277233232 7:-0.7295510954484412 8:-0.41443875567123967 9:-0.610576632050404 10:-0.9515988311377204 +7.084732852050431 1:0.9990215581592679 2:-0.9868954542412269 3:0.49133473382040704 4:0.7697599878561228 5:-0.34668939907967267 6:0.9777705993519483 7:0.4449043102759509 8:0.9812971199646168 9:0.6666598587737487 10:0.14398842572598514 +0.23715467505851734 1:0.21628799185444336 2:-0.4526390568867018 3:0.6558486691929235 4:0.13730688681492142 5:0.23076986155942736 6:0.7020484017619715 7:-0.12077999528458938 8:0.8306084972447003 9:-0.49337323198621563 10:-0.8270028152572872 +1.1552619549601455 1:-0.48202394020369277 2:-0.6274878708695264 3:-0.27623674153600697 4:-0.5312153415813432 5:-0.030820182786174044 6:-0.5893370965577813 7:0.6666315120904487 8:-0.36482991729570036 9:0.6065771813692735 10:0.05831057330788525 +-0.20433879835596253 1:-0.4702220250018212 2:0.9123705796362889 3:-0.2045657170490376 4:-0.18922063450309534 5:-0.31431213362503163 6:0.4150130060120387 7:0.34016193625941127 8:0.8391374136299805 9:0.6884250315764333 10:-0.7916408854251566 +-9.751622607785082 1:-0.0014232315621649505 2:-0.1284246813729939 3:0.5228953023175369 4:0.9688522449007109 5:-0.7857721219549156 6:-0.7812922263391038 7:-0.5916136652814756 8:0.793988610184206 9:0.7982949061274296 10:-0.592785473963741 +-22.837460416919342 1:-0.17363144173810174 2:-0.3340314573781735 3:0.9351424971322297 4:-0.6430601902397572 5:-0.13363305808148818 6:-0.42446359566938585 7:-0.4093070316761178 8:-0.9302259781839204 9:0.47004365892170585 10:-0.6231289889808045 +-3.6318714209289436 1:-0.8296410705737971 2:-0.6056572341069668 3:-0.2975417404042737 4:0.07134138175064741 5:-0.8966463747179154 6:-0.4112675899658855 7:0.7908013478009401 8:0.407396254566472 9:0.9227769302156879 10:0.12418427404473764 +-3.8909712376010583 1:-0.6552751548581366 2:-0.5641921108932855 3:-0.6340486345063014 4:-0.5441069121131075 5:0.908720622198947 6:-0.026054643814348077 7:0.03706191653058433 8:-0.6672524338819317 9:0.7958274915288801 10:-0.19029619970124023 +-10.600130341909033 1:-0.7457695999520562 2:-0.3739453132549577 3:0.01327423342620393 4:-0.08482897201178563 5:0.84573456086082 6:0.6279927575103963 7:0.014494803555804125 8:0.9420647557771027 9:-0.13484113287285893 10:0.3037405853352888 +-12.094351278535258 1:0.9982796018306028 2:0.8354271779265348 3:0.46284321795736116 4:0.07693347919601745 5:-0.4753440408996932 6:-0.47098252868073787 7:0.4810729184846003 8:-0.6136990339205741 9:-0.6715833036640317 10:-0.6247058955319091 +9.936399360181602 1:0.7330323083522969 2:0.47204204993669197 3:0.3850471475752122 4:0.21483460195167958 5:0.3806220122265147 6:0.6336993433402796 7:-0.47987416364572 8:-0.195509010865196 9:-0.6561820282562041 10:-0.45300480439842894 +-4.706701061062994 1:-0.847895844561626 2:-0.29946646506145114 3:0.25432868082106497 4:0.1352958872054535 5:-0.8803017974303002 6:-0.3675110562764785 7:0.10888496324899721 8:0.9620737605396772 9:-0.031046632561323895 10:-0.09466883461500908 +5.101614991255809 1:-0.5174248135588373 2:0.14394061894828014 3:0.5613709266711013 4:-0.5678634944326011 5:0.930216209978763 6:-0.6204727890080077 7:0.4133141749872311 8:0.6262685035917408 9:0.03382924477926896 10:-0.15231139191832854 +-8.772667465932606 1:-0.2117605577769197 2:-0.4283897136887762 3:0.44686767473401035 4:-0.5507826261358746 5:0.237124956028401 6:0.6413157520982717 7:0.2409214827604571 8:-0.8505503638033114 9:-0.9811997368468401 10:-0.9499963936664035 +-11.615775265015627 1:0.8782018665273386 2:-0.9751473570197167 3:0.6307050068521085 4:0.7012721336851997 5:0.21318736263512283 6:0.024885128053773853 7:-0.4580644243558505 8:0.1318650007251434 9:-0.9306090092992167 10:-0.5688746770986652 +19.64829023536192 1:0.14426537998360645 2:0.3557716894181753 3:-0.8577143134654568 4:0.5288643233801469 5:0.9231529738221469 6:0.975999712077738 7:0.24700404691888678 8:0.10206517527052283 9:-0.10041951294847062 10:-0.9412918491876225 +2.7409415438025486 1:-0.7404936009304737 2:-0.9792071376296605 3:-0.49510748520932113 4:0.9538460112904268 5:-0.5075114153141447 6:-0.5890791308058669 7:-0.775366087491284 8:0.4983912525892249 9:-0.2976197956132913 10:0.6791258030468514 +-4.394658158733604 1:-0.41628618754613345 2:-0.1472602552309057 3:0.44136102233464025 4:0.011882653940414434 5:-0.6559502840386595 6:-0.4022529016339016 7:0.048402312931387526 8:0.8753776623326166 9:-0.8528247288266961 10:0.6593783978826002 +1.1915739133607073 1:-0.7840827624854878 2:-0.4860418508208426 3:-0.7418773161179972 4:0.129874781837924 5:-0.22631682294184796 6:0.47794208013755024 7:0.5532183426143056 8:0.11879859459306741 9:0.09927630694484524 10:-0.19268618891399636 +2.156192215438919 1:0.44325986644475646 2:-0.6057278708888592 3:0.3943381582091723 4:0.6560336238050575 5:-0.9651308100517204 6:-0.2358219003943678 7:-0.04143043460232465 8:0.8623951169233035 9:-0.4933545255502605 10:0.8990427200454263 +-1.1009750789589774 1:-0.4515707618788496 2:-0.745936099912899 3:0.41307003181926794 4:0.6127760492402428 5:0.9250878169732681 6:-0.6778628527469126 7:0.42794190420905753 8:0.4943969797578971 9:0.7762709104958854 10:-0.6932349268610041 +10.04434496594037 1:-0.0995467494040092 2:-0.7766769414838959 3:-0.6608009972582911 4:0.7075788021090594 5:0.5208396359138381 6:-0.09724033794207299 7:-0.743087245352148 8:0.765372791789753 9:0.3788699859744704 10:-0.6898257995055466 +8.038039859115667 1:-0.5321510657026671 2:0.5571925538006008 3:0.747268102801854 4:0.09079641165917596 5:0.25861122989509266 6:-0.9948187479498878 7:-0.9665136866462685 8:-0.3904629432867681 9:-0.9975425877998279 10:0.32024289816988416 +5.14371929922303 1:-0.4829199170694627 2:-0.5713285263827719 3:-0.43889652467111184 4:0.18478247261988967 5:-0.27374063120041225 6:-0.8069125377696931 7:-0.15497746743367058 8:0.32448521325998714 9:-0.39397735035206227 10:0.08184957956614292 +-1.6848276484379352 1:-0.39250489761445895 2:0.02730338852529557 3:0.9916055514435305 4:-0.07571433435055064 5:0.19024527726403728 6:0.6385182319185971 7:0.32480605537471297 8:0.5807543325220577 9:-0.35642510103381153 10:-0.9060482769392468 +-11.640549677888826 1:0.03707410390488852 2:0.2527049166981137 3:0.4114872952854447 4:-0.8508977901757795 5:-0.42791544663481895 6:-0.9864047295390463 7:0.6023685964407528 8:0.12018443688097036 9:-0.36816249877130414 10:-0.9583147535652901 +11.672104494601319 1:-0.2416258355340175 2:0.6737553249072334 3:0.9041602191361382 4:-0.2123232797997281 5:-0.008255188002961988 6:-0.5151894064136904 7:-0.7341877977528246 8:0.624625272218277 9:-0.6261434804192929 10:-0.5710586715741532 +-2.2960192492344627 1:-0.7457768645184579 2:-0.5954998103421847 3:0.5428846769211537 4:-0.6176587961491775 5:0.46222150678166574 6:0.7852238239427731 7:-0.3614580530629148 8:-0.325840253127059 9:-0.5660596710348922 10:-0.8060263366626401 +5.428302298615722 1:0.8774286357993033 2:-0.23110126319781088 3:0.6264134914476072 4:-0.143015582616014 5:0.350109539755298 6:-0.147747167834422 7:0.05020570422182824 8:-0.5469605849960337 9:0.951112567977048 10:-0.34800121380288185 +-17.32672073267595 1:0.31374599099683476 2:-0.36270498808879115 3:0.7456203273799138 4:0.046239858938568856 5:-0.030136501929084014 6:-0.06596637210739509 7:-0.46829487815816484 8:-0.2054839116368734 9:-0.7006480295111763 10:-0.6886047709544985 +7.314490512652487 1:0.8745354279105222 2:-0.9270067504840309 3:0.965218170323435 4:0.12808957052353698 5:-0.5309399625085234 6:-0.5968520990090951 7:-0.667403236513185 8:0.08154410986660832 9:0.33025488397543934 10:0.03406708067839537 +4.687373993408297 1:0.6731426721418288 2:-0.7111023070261273 3:-0.9849054116048603 4:-0.12831346258317322 5:-0.04095946352836921 6:0.6967001556166801 7:0.8479895229743999 8:-0.35600791972899404 9:0.5005979045264868 10:0.6421341979636503 +-6.82923852156868 1:-0.04849233571020073 2:-0.8505855619911602 3:0.2927180954190314 4:0.5780268040086791 5:-0.22322207765417268 6:-0.8436513934568071 7:-0.3906240514635124 8:0.7258714963093444 9:-0.21695043530813085 10:0.8049335285918169 +-8.24622879369294 1:0.12154833675098842 2:-0.26446415445316673 3:-0.06653791221669247 4:-0.7920694887292259 5:0.6128791496627621 6:-0.6927179137980173 7:-0.24584418172709932 8:0.3557416365779935 9:0.22868636757755234 10:-0.8288196322549064 +-5.090863544403131 1:-0.1535668648046895 2:-0.59868738365189 3:-0.8822518703008675 4:-0.1790505106198006 5:0.9467581256591948 6:-0.0661313762905984 7:0.31263046332923694 8:-0.03628894224569357 9:0.8969599435828515 10:-0.05386674051170348 +-15.780685032623301 1:-0.2568492063716883 2:0.7740976197426315 3:-0.7829158104387535 4:0.8578846037465748 5:-0.6111039318672586 6:-0.26939268282639306 7:0.3659136640533909 8:-0.8205938562638555 9:-0.24945505706767923 10:-0.935948184861368 +-3.9916779937384743 1:0.22925954469403154 2:0.18159238246979537 3:0.05607027262862396 4:-0.3376037702047998 5:-0.10630000583678934 6:-0.7148277241201622 7:-0.08327294541727137 8:0.6532439360618307 9:0.34352364313237294 10:-0.21028242388807916 +8.798748248458631 1:0.509058184822212 2:-0.17532831457577935 3:-0.6387880909085213 4:-0.966194650702529 5:0.7829797328120436 6:0.5470735549914605 7:-0.38312745239682333 8:-0.8890923931840893 9:0.6823342859396513 10:0.9231260597729121 +14.341273640964873 1:0.6996156678090684 2:0.8612833977834464 3:0.9131301694042417 4:0.5199385192744859 5:-0.32605907950755086 6:-0.9816465962348846 7:-0.5939885763232406 8:-0.7730924566676425 9:0.6367821449954114 10:0.10873812383881054 +9.75855501262469 1:0.2933324921347933 2:-0.4652534314332506 3:-0.2940640558090537 4:0.9883453215038367 5:-0.042460731786114314 6:-0.15438550895912062 7:-0.11182397625560592 8:0.7425954283250873 9:0.5063859049644963 10:0.3012211854180429 +7.695200921242407 1:0.3554353390157281 2:0.08707592690448718 3:-0.10815435665633877 4:0.05524046679762784 5:0.8000157491787581 6:0.3756193347272323 7:-0.18659830666742527 8:-0.08168623764933125 9:-0.2551379303720174 10:0.8560030587463281 +26.903524792043335 1:-0.4672678144441864 2:0.868381965588082 3:-0.04748335609643428 4:-0.0908285508827269 5:-0.22436865911994275 6:-0.953965287326564 7:0.40644848732968164 8:-0.33391575325981115 9:0.008337907338700212 10:-0.45597904754961416 +9.87318781117539 1:0.7310287890171705 2:-0.38300115058116324 3:0.5492682498036086 4:0.552016070316655 5:0.3715022458396897 6:-0.3888040017277252 7:0.21348231125683648 8:0.23219558685722874 9:-0.6271161253492459 10:-0.009137052604519136 +7.6930514050666625 1:0.48603550488592284 2:-0.9218820771919889 3:0.17678612698428053 4:0.5110501870908806 5:0.5817010201164554 6:0.4488707800038747 7:0.4977618637956498 8:0.1683214570038094 9:0.17237242672259323 10:-0.5276084644007359 +3.155413914311745 1:0.04582517188512947 2:-0.9060800653779759 3:0.049786270132956556 4:-0.4236784487542993 5:0.6293910028372613 6:-0.7370237043436467 7:-0.4599678991281728 8:0.5317111095323057 9:0.0029525239228334055 10:0.9294876800738165 +-10.18815737519111 1:-0.9023553189306839 2:0.08434165073970856 3:0.7163931103395633 4:0.41749986495957914 5:-0.8190972970472759 6:-0.9996126872234177 7:0.1779075727741255 8:0.18212754689351862 9:0.24628508239298963 10:0.667589863190412 +18.585731475373457 1:-0.8399129036462931 2:-0.10024819268489127 3:-0.5011350892733817 4:-0.7299256348863585 5:-0.9412022985072928 6:-0.245064895931544 7:-0.1032512650854267 8:0.9943863256441088 9:-0.6429371028855466 10:0.062299742931960056 +8.998359297106072 1:-0.16850226855111905 2:0.7262839202089402 3:-0.04876255055071854 4:0.8948164957242868 5:-0.10720585418953132 6:0.2622719447841948 7:0.26433837506661373 8:-0.5143449147399106 9:0.17444585508955002 10:-0.813182163328944 +13.032424230011074 1:0.4014766166181287 2:-0.1710502754125871 3:-0.309850483152607 4:0.255642456909988 5:0.9949117714165621 6:0.12553772251510864 7:0.6412602805648968 8:-0.6225679446416825 9:-0.15867011477056936 10:-0.4970695349036196 +-6.931030745953174 1:0.5151452174260762 2:0.10077064818539072 3:0.9515221270405545 4:-0.21551878535257907 5:0.29152528087481366 6:-0.10995497026133605 7:-0.7872786530801681 8:0.9909149980139627 9:-0.6044617953251021 10:0.4135285912642448 +15.538062451207367 1:-0.493569696351595 2:0.7280914440594639 3:-0.5399160539735497 4:0.5688018985826291 5:0.8296550361854862 6:-0.3519274619833537 7:-0.5536583684230114 8:-0.9648774930921231 9:-0.2649670832738824 10:-0.2337289004188019 +9.499729032920945 1:0.22017490770298553 2:0.7693082799289328 3:-0.7645745307823122 4:-0.4243400515554365 5:-0.7065281515163817 6:-0.9488470141298047 7:-0.7888781431404843 8:-0.38027758953310964 9:0.11329243985448345 10:-0.5636550498916204 +-0.6039115764951412 1:0.3128791250125589 2:0.4690308315665288 3:-0.9819748103687955 4:0.28931283693913223 5:-0.6283983933456656 6:-0.10795935596621975 7:0.7785831799196448 8:0.4453768248295542 9:0.4055410615499917 10:-0.581108383985806 +9.682301463907875 1:0.5039970331368235 2:-0.008965105921562966 3:-0.5415225380115054 4:0.4677111860370293 5:-0.3854089758945243 6:-0.8468317339287676 7:-0.29258253017713587 8:0.7361173598968789 9:0.5722561668394952 10:0.8524030171340933 +-2.8752191903928064 1:-0.45407356732165205 2:0.6563221064539377 3:-0.8938366926767671 4:0.6028173420234533 5:0.6792881349943096 6:-0.6295604812779405 7:-0.21641416912497213 8:-0.8703620515028858 9:-0.3397362922228042 10:-0.0966947467107604 +-28.046018037776633 1:0.9493308195854675 2:0.3285214661535252 3:0.749300278016316 4:-0.006681618268088219 5:0.2936055273341429 6:0.0044706790416966236 7:0.5006172205470896 8:0.38751814960349473 9:0.6069735922707928 10:-0.794612882855285 +2.8752582614589373 1:-0.9443232811926943 2:0.3153126492983107 3:0.6423843271417344 4:-0.09528333043829118 5:-0.2318773828230698 6:0.32597909562645766 7:0.42808555740416065 8:0.2895959316734451 9:-0.5106491076955746 10:-0.2917418155655722 +-2.203945173593806 1:-0.13844025039418084 2:-0.024638102806725293 3:0.6114514176076162 4:-0.6939316676972749 5:-0.036549673716341324 6:0.0942395290460385 7:0.7943411369475493 8:0.7025693796408046 9:-0.21822635487138853 10:-0.6045250179827362 +-5.070655299509993 1:-0.8035156105848074 2:-0.5344928236067734 3:0.6105404604447127 4:-0.7538635525543969 5:0.9836765037886612 6:-0.5700253195942724 7:0.9232380985458313 8:-0.26374377078100464 9:0.9079431466301682 10:0.8404281771949533 +-2.540181413836895 1:0.220453181647285 2:-0.03105792440486077 3:-0.17131282366411926 4:-0.41800060634660485 5:-0.1477564564540963 6:0.055537469469941536 7:-0.8092076926316594 8:-0.29815112444525727 9:-0.20030580647762464 10:0.337865838755971 +19.341342586351033 1:-0.32052868280788616 2:0.954507993011956 3:0.38642226954792824 4:0.9240442034609888 5:-0.46077559741256824 6:-0.7991393493824104 7:0.9396232321156679 8:-0.2486930151964184 9:-0.6256485833035617 10:0.14861843824730103 +0.31398559122529757 1:-0.4684215762946897 2:0.07873308388585198 3:-0.3589594894052015 4:0.14284662079329458 5:-0.8936272055527841 6:0.5647217242826741 7:0.49613233215723507 8:-0.501698787526992 9:-0.46710107378968724 10:0.898517179577361 +12.243117462926584 1:-0.8147610562690222 2:0.21104006948075482 3:0.42405323019132957 4:-0.667965573810795 5:-0.267026607469405 6:0.7949752815579358 7:-0.07664414977654532 8:-0.6023087644686556 9:-0.659375887511856 10:0.459735946423397 +-4.623091296763939 1:0.08405646515942733 2:-0.40009448092691446 3:-0.39831245310544094 4:0.8794137836499942 5:-0.04788565812369017 6:-0.06763019434549333 7:0.41324877265674065 8:0.39746868847324146 9:-0.986729367280818 10:0.7001677710291752 +-5.782162271139417 1:0.29127970805530157 2:0.6712715787317827 3:0.27575757044478477 4:0.31525054647682804 5:0.6905016168465983 6:-0.5195319089267731 7:-0.06598129860341295 8:-0.5142554034519407 9:-0.11546331150946942 10:-0.2934524891698944 +-9.892155927826222 1:-0.7048583334456604 2:-0.46598491327111247 3:-0.02034722477413209 4:-0.663294196316965 5:0.4485329128582778 6:0.4553619594861118 7:0.7154814909138205 8:0.7532937661147989 9:0.020693077287389894 10:-0.23131986644633207 +0.5422764698408844 1:-0.1513298744027669 2:-0.4591544337339648 3:-0.7192219559850723 4:0.21236658135317632 5:0.12050445497328166 6:-0.42411528242712127 7:-0.15103925528861595 8:0.9032115729799512 9:-0.9228817525021624 10:0.2604090001033641 +4.187800872274017 1:0.3084355607627949 2:0.7029638272178733 3:0.34098344122299573 4:-0.5818421369891376 5:0.1332356708082485 6:0.22671316744441716 7:-0.6750469752494854 8:-0.4065302428716193 9:-0.48213803977370073 10:0.17918596677210186 +4.487701812297124 1:0.8352061350259052 2:0.2757393215770836 3:0.9310504392364667 4:0.519503546762708 5:0.5270245209143005 6:-0.08181154800488488 7:0.5148324302455536 8:-0.6680946101511949 9:0.7574060703813035 10:-0.4721334895419935 +-5.150140984417682 1:0.8113709439821006 2:0.21939305063309278 3:0.02109986546311826 4:0.07450107676582762 5:0.723883853128624 6:0.5392035186380486 7:-0.1382740221237464 8:0.9990201540159807 9:0.10429329766137108 10:-0.1365266408862309 +-6.544633229269576 1:-0.08278037549320039 2:0.6982730989138761 3:0.9090685953368327 4:0.6754092061339365 5:0.5889199822482736 6:0.020678619551471433 7:0.47605785660672084 8:-0.49926771127869873 9:-0.28380077002944093 10:0.5282319276258469 +7.216836352055753 1:-0.8510680074642156 2:0.42611818262128476 3:0.593607821624947 4:0.5635067468583634 5:0.2121930523769171 6:0.2708063180622071 7:-0.31491113345871735 8:0.005990053407278095 9:0.8985259402559085 10:-0.44549339042232794 +20.874246167942125 1:-0.53010692413621 2:-0.9897084749945524 3:-0.9083978261828305 4:-0.15581655583739495 5:0.9974035542095165 6:0.9894717992956665 7:-0.7287287537245402 8:0.06425127137526943 9:-0.06684164745938337 10:-0.3600621883071937 +-6.556192430758147 1:-0.7655958349167471 2:-0.08083170734199419 3:-0.8540636958251198 4:-0.09994429443696973 5:0.1734809016500265 6:-0.29563180244063325 7:0.2158497607364409 8:-0.6071644305523003 9:0.8063426715403785 10:0.47092299197899345 +7.252748885335252 1:-0.36403312429467216 2:0.1237451136826817 3:-0.5756427605741237 4:0.7612833636750866 5:0.9350628314096134 6:-0.012087843264624754 7:-0.03742573515965031 8:-0.05210460803183037 9:-0.5333214800203341 10:-0.013320030179712505 +-9.2679651250406 1:-0.5057250557539077 2:-0.41655319851679495 3:0.1897431234740683 4:-0.038318717640150046 5:0.9136495575471062 6:-0.8890525036858237 7:0.40859501498633377 8:-0.8746985847539293 9:-0.005836984002720369 10:0.7838036026237987 +-15.732088272239245 1:-0.8546867577633044 2:-0.3003980324850013 3:0.49649883896876834 4:0.710496747220617 5:0.5848510480601048 6:0.5714826756665468 7:0.5487975165953451 8:0.5654333402837335 9:0.863539315599626 10:-0.9699410102494574 +-0.20412431312519014 1:0.13323548063028934 2:-0.3030177580658542 3:-0.6358920925969869 4:0.3729380701923921 5:-0.8131818118430312 6:0.11567152703716288 7:-0.3645508535812394 8:-0.5487213252460876 9:0.5605886387366952 10:-0.8400308993051686 +10.445759684895373 1:-0.92707280355555 2:-0.9307772570299944 3:-0.11971873660640964 4:0.5140245291069254 5:0.5751145648836897 6:-0.43850910073502347 7:-0.7872208869913662 8:-0.3087975452145404 9:-0.4645849758749403 10:-0.1563641826381328 +3.349708377102383 1:-0.6334394121009499 2:-0.9008086683014112 3:-0.2678892493467009 4:0.7626514243443427 5:0.6406493676995701 6:0.3669245573649391 7:-0.052050629941784665 8:0.6713394117904852 9:-0.11458974566378233 10:-0.25949626043219576 +-23.487440120936512 1:-0.5195354431261132 2:0.8080357948412571 3:0.8498613208566037 4:0.044766977500795946 5:-0.9031972948753286 6:0.284006053218262 7:0.9640004956647206 8:-0.04090127960289358 9:0.44190479952918427 10:-0.7359820144913463 +-11.827072996392571 1:0.9409739656166973 2:0.17053032210347996 3:-0.5735271206214345 4:0.2713064952443933 5:-0.11725988807909005 6:0.34413389399753047 7:-0.2987734110474076 8:-0.5436538528015331 9:-0.06578668798680076 10:0.7901644743575837 +-3.650649176738987 1:0.9665344025238449 2:0.1395514751689353 3:0.954697162791015 4:0.2093601878355722 5:-0.42841737775246336 6:-0.02877209657213764 7:-0.8382526163632971 8:-0.03773878779258388 9:-0.3751775119106411 10:0.6477987464528951 +0.21915863046310957 1:0.25143109618049353 2:-0.06463696557011112 3:-0.3324862332340037 4:-0.7184623449423757 5:-0.8897217937178385 6:-0.7336278194091297 7:0.8547631637534296 8:-0.7582613025929346 9:0.9080481791309838 10:0.9427850135311773 +4.813247597584681 1:-0.4564689661727537 2:-0.4315414033069003 3:0.09676404446694242 4:0.6024645727173434 5:0.20466090997530606 6:-0.09432916868838737 7:0.6402934161890248 8:0.741842551426011 9:-0.343937669190693 10:0.308871619426873 +-3.0700825038127206 1:0.660084046469162 2:-0.02317305725931229 3:0.7567569356692221 4:0.2528834502236612 5:-0.3935091635208898 6:-0.9965507922509653 7:0.9065754202428946 8:0.6988037588300844 9:0.05145737657924321 10:0.4697377584426863 +9.762542323725354 1:-0.036129448543738896 2:-0.8252508992030534 3:-0.752854859129851 4:-0.9027424488033049 5:-0.4763092428375775 6:0.4832492121777574 7:-0.2935697977919014 8:-0.9197908986231211 9:0.8914359296658816 10:0.8688484670974876 +6.690913813146277 1:-0.7649833946109403 2:0.0419327356721928 3:0.5420954694310764 4:-0.7373259510045522 5:-0.9187577877864708 6:0.6431180783847401 7:-0.6272529754533058 8:-0.43356486537110106 9:0.16848266440424364 10:0.3129700315745716 +21.325049167466855 1:-0.36392795201361383 2:0.846518905511275 3:-0.26361421923150097 4:0.5140384860444887 5:-0.9147771624497878 6:-0.22044646197773576 7:0.14099760779666948 8:-0.546631395802236 9:-0.4345465263406878 10:-0.2759894364167672 +0.41237529640734055 1:0.05016964684797287 2:0.21708512805176072 3:-0.9444942733586354 4:-0.6118772896807114 5:-0.18053631846913665 6:-0.06752556529755416 7:-0.0081819952134361 8:-0.7774039956687315 9:-0.5548994336153177 10:0.7510833121912588 +-15.056482974542433 1:0.6012054064354875 2:-0.6127014811673221 3:-0.8356741843949218 4:0.19830469052767397 5:-0.07726493085289698 6:-0.5756891943805014 7:-0.49010583357941884 8:0.7493759119974515 9:-0.7828994218436376 10:0.6154265137741459 +-2.109441044710089 1:-0.5757976103755722 2:0.3686657403505862 3:0.5418762444017706 4:-0.5896052565388463 5:-0.1000712585735879 6:-0.8114188394866342 7:-0.5863884932327266 8:0.28289838755838015 9:0.5378646921099333 10:0.5063780890366179 +-5.249715067336168 1:0.6828022788286754 2:0.9044668986252975 3:-0.6010464361571437 4:0.8416122052398811 5:-0.9846446498408039 6:-0.3777762313579811 7:0.5763775880953983 8:-0.07608009385213488 9:-0.29576023599575474 10:0.8845728751981716 +6.907770824878343 1:-0.9751352215365647 2:-0.8297271715190588 3:-0.07240311280415779 4:0.4796310183582191 5:0.358213469979769 6:0.4628020211207058 7:-0.9753405605972942 8:-0.765583403709019 9:0.5623611232648877 10:-0.14725965272406616 +-9.299021854126096 1:0.8784076266914045 2:-0.8314918563417382 3:0.8701529449600536 4:-0.8070129727442199 5:0.07396877198841345 6:0.0040889707225901795 7:0.40529205456687145 8:0.6412485325027342 9:0.1443450351498905 10:0.404997568726581 +10.95643670126225 1:-0.37321642594676097 2:0.47766490569544473 3:0.9201313123144423 4:-0.649393433578801 5:-0.9084894063674787 6:-0.2547160991750408 7:0.7674649994523459 8:0.646056370118979 9:0.6014100713287893 10:-0.15130291862509182 +-2.6397202393123336 1:0.3285252466844373 2:-0.2714281159811125 3:-0.5869561846815805 4:-0.5643935541712441 5:-0.7285201267315389 6:0.6502951074428092 7:0.8611880383193904 8:0.6380425291162128 9:0.5118538704085516 10:0.4012684110865874 +12.521131042032012 1:0.4843931319727355 2:0.06440853455169626 3:-0.6151259240105509 4:-0.4180928328467284 5:-0.4607061773323424 6:0.8104775289268906 7:0.3284199695768064 8:0.8425028998495565 9:-0.34822319854822825 10:0.1969239149176112 +-16.151349351277112 1:0.7787909191620395 2:-0.5219981442072688 3:-0.7230569918898555 4:-0.05707801168212101 5:-0.8134225824740247 6:0.09950684183685454 7:0.6261274830059296 8:-0.9502006765164366 9:-0.6724983095526844 10:-0.600347212281825 +-5.039628433467326 1:0.7680701397575322 2:0.7956844224408437 3:0.4131717201035916 4:-0.3127895385265915 5:0.7226571953995224 6:-0.06845863083031967 7:-0.1007291660029832 8:-0.9130249132342207 9:-0.4605180615474036 10:0.42093879298156 +9.007596502870785 1:-0.6562175566238462 2:0.08420074013217049 3:0.589801949672486 4:-0.11964901133703987 5:-0.8145711913860048 6:0.43854302140351065 7:0.5992967124729605 8:0.253745043289755 9:-0.10742030998120033 10:-0.5993228348160153 +-12.41094640284016 1:0.31035917086763765 2:-0.8053417167237813 3:0.5754655536186164 4:-0.3645388095106201 5:-0.9135176753316416 6:-0.8690739610562535 7:-0.14039224825138197 8:-0.7112835675593987 9:0.25762942117230825 10:-0.9483300117501923 +-12.130353212287929 1:-0.41404309625298485 2:-0.7695984204591535 3:-0.44569447239245275 4:-0.3168863099965644 5:-0.26669244730409036 6:-0.33484042698895755 7:-0.41062396946367685 8:-0.09075804785640385 9:0.8511367190902208 10:0.021918606255194595 +-15.375857723312297 1:-0.9794952880997945 2:-0.9547237660069134 3:0.2460912345929791 4:0.3174335823329406 5:-0.23758562926743054 6:-0.113610303129287 7:0.18292675847568063 8:-0.9656446754474337 9:-0.58300134324846 10:-0.6689602908128025 +-6.397510534969392 1:0.440780662587545 2:-0.03737991637410243 3:0.9506435891605849 4:0.8177486462589998 5:-0.2917628929963241 6:0.42365289098031034 7:-0.4280555544979745 8:-0.18388426864865903 9:0.5057230088452542 10:-0.1699163749308643 +-9.789294452221961 1:-0.25066699970459694 2:0.1772977344415987 3:0.5913498268900952 4:0.6293756431864681 5:-0.6430441015863757 6:-0.7238519180293621 7:0.13639541626580498 8:-0.6620281401715837 9:-0.9515237061912034 10:-0.4333426289849791 +-13.15333560636553 1:0.3491978525665129 2:-0.4067353159374012 3:-0.8677040612253524 4:-0.5757086910974862 5:-0.3186886816681207 6:-0.06609938943414573 7:-0.5419747642754873 8:0.9632759660044383 9:0.2673520823110991 10:0.36463236596724546 +2.2307697392937795 1:0.12285527276472785 2:0.8938323722714365 3:-0.16995870341610209 4:-0.3298643049714254 5:0.16781582791954253 6:0.42381594687105895 7:0.9245288214717629 8:-0.08709025093361689 9:-0.14227085487682722 10:-0.2888302862659746 +5.892885365508635 1:0.10116053019915738 2:-0.41641547074900154 3:-0.3750004290914961 4:-0.5619470211369917 5:0.33343039544460384 6:0.46766042657994733 7:-0.6937940929321615 8:0.7044604392055189 9:0.8879353764416567 10:-0.5490902425042639 +-16.692207021311106 1:0.9117919458569854 2:0.628599902089868 3:-0.29426892743208954 4:-0.7936280881977256 5:0.8429787263741186 6:0.7932494418330283 7:0.31956207523432667 8:0.9890773145202636 9:-0.7936494627564858 10:0.9917688731048739 +10.454641756541454 1:0.3490213088098768 2:0.6103387992494194 3:0.6753935651135747 4:-0.39560763769937934 5:-0.3042308221531884 6:-0.9283481899557042 7:-0.7904038212853011 8:0.37488335848537346 9:-0.296477977723397 10:0.30894819444660304 +0.08978797103855778 1:-0.13445409764877803 2:-0.6404150831493631 3:-0.24740260669490133 4:0.031151119464385646 5:0.9207882173498612 6:-0.6146471129497393 7:-0.9736175690408087 8:-0.2673180325645341 9:0.5800384183301572 10:0.479811220263183 +1.7362099941626894 1:0.5171681395917551 2:0.6177735922313075 3:0.6446678302226738 4:-0.5731769722311459 5:-0.2686270617709168 6:-0.6048534221658814 7:0.7002124303669326 8:-0.1479765297345712 9:0.009254061109394307 10:-0.31519081920853287 +-1.0349488340235453 1:0.612980711993536 2:0.05771318707554962 3:-0.10821368362160744 4:-0.8755559420458141 5:0.42566546089913326 6:-0.7966341558699277 7:-0.45253617234374466 8:-0.8289517557653971 9:-0.8968075137250837 10:-0.6325457096866376 +0.10157453780074743 1:0.9143592240573388 2:0.06419631741815457 3:-0.9961326744227916 4:-0.47174548800139715 5:-0.0821464027819967 6:-0.5495006555498168 7:-0.5627911401420294 8:-0.43426056724099005 9:0.892026786364895 10:-0.23546485121284055 +-12.92222310337042 1:0.218687524173371 2:0.013626751799176162 3:-0.8372219908323961 4:0.6197296846266354 5:0.7429130827811232 6:0.48009972886541896 7:-0.35667717521227904 8:0.18337067878780533 9:-0.22935396092245197 10:0.4076715024284059 +22.923352376063196 1:-0.7522075505725567 2:-0.20686029838909326 3:-0.1386664769095396 4:0.157117595808127 5:0.9462377653889174 6:0.9182504509330662 7:0.18170057747293833 8:0.27735387813088863 9:-0.6355799944714868 10:0.9764849106195284 +-6.132450015997121 1:0.2822534275343054 2:0.2625905791399692 3:-0.02565260641304068 4:0.4891221076432757 5:-0.6426178913585772 6:-0.8999539149461033 7:0.12659507663825287 8:0.5889572439755832 9:0.49107548332672857 10:0.47595749470452 +-9.173693798406978 1:0.4430245286298278 2:0.9923116639471541 3:-0.5609082824097824 4:-0.36542266258313916 5:-0.5814039716882617 6:0.20413852042674874 7:0.6097541611931963 8:0.5743002479324253 9:0.4735459963431561 10:-0.053969823043886755 +-5.814408490931223 1:-0.9206287328000513 2:-0.48298486023273157 3:-0.8477202041890262 4:0.5801385102362351 5:0.7146074564553095 6:-0.5987672678579339 7:0.6829077928212723 8:-0.3522788540815065 9:0.7729595638821951 10:0.5264904880591215 +6.474329501040298 1:0.6914309300550991 2:-0.4507700505202725 3:0.713821440501512 4:0.41599059910235847 5:0.507160951750409 6:0.8635615811574222 7:-0.6235518270244333 8:-0.5336201820384283 9:-0.7989630679361768 10:0.837293162455248 +6.984517471584806 1:0.16745919469723392 2:0.018033079961716103 3:-0.7339201095541323 4:0.17042828693740697 5:0.4493471632580528 6:-0.8938445962323078 7:-0.3206968104792325 8:-0.616617071238893 9:0.9327878222034172 10:-0.6575294247048245 +-12.479280211451497 1:0.9769767754725367 2:0.7706430863248943 3:-0.4506244622476816 4:0.12921761745628713 5:-0.0697804449658812 6:-0.7702703569987461 7:0.017734558413919688 8:0.7216294158911261 9:0.42547357862241886 10:-0.9001915116155741 +2.8363866587728186 1:0.11478724114928918 2:-0.4679790550082039 3:0.2344912687736711 4:0.5524878060045462 5:0.5252859884051309 6:0.5080674087215156 7:0.5010449021825665 8:0.048046765816400105 9:0.06654581719548891 10:-0.5801934713347348 +4.186809777233374 1:-0.02335342201396018 2:0.9035437912091193 3:-0.9283585631882163 4:0.454351316397237 5:-0.6948564428085262 6:0.11495485234890368 7:-0.23683956078769963 8:0.6442534752881419 9:-0.013866407845647188 10:0.23369602940650736 +2.8235031660626415 1:0.5609344938188046 2:0.3449103464885612 3:0.03972169049525687 4:0.31858762565827137 5:0.4409953589124853 6:0.22836189275697016 7:-0.1497811991899889 8:-0.23248048920679265 9:-0.30066618281100177 10:-0.9247232456911632 +6.96223432848425 1:-0.8160398553437558 2:-0.8212180893749699 3:0.7728655115832999 4:0.02387973088796369 5:-0.043499804905828166 6:-0.6997726250046865 7:-0.8686633773265577 8:-0.12597318402253976 9:0.967018116368416 10:0.5951339624149812 +4.669684795838683 1:-0.32226903644852833 2:0.5465858078942492 3:0.5228467793266189 4:-0.013157722224545143 5:0.5810668818928995 6:-0.1372653090293532 7:0.6446157527288279 8:-0.06005754873230629 9:0.014302180040152379 10:0.43474245441042636 +16.112744845653285 1:0.37257742858083365 2:0.19398954512844124 3:-0.11860882189887478 4:0.6492510749703395 5:-0.41273736981203313 6:0.18643017041815835 7:0.29136917186214384 8:0.47602883023389 9:0.7126916980867937 10:0.48462508659691483 +-9.196003366226202 1:-0.7263358951920722 2:-0.8503799288093836 3:-0.3120563620589105 4:0.3925562655164563 5:0.027666662972283484 6:-0.35173134138805406 7:-0.32703527910354757 8:0.3060102722285065 9:0.8609161725740202 10:0.33394557004432923 +1.242972458167591 1:-0.9029238804456814 2:-0.6392681059531908 3:0.8940879647942577 4:-0.8807357173896475 5:-0.13628130467470512 6:-0.5487534785116224 7:-0.40270307148061346 8:0.09152108686997096 9:-0.20745066734844642 10:-0.20624830574384978 +3.453659210660726 1:0.2710596844435682 2:0.6510497900145247 3:-0.2899158136103117 4:-0.13531811694554707 5:0.6965847786422426 6:0.9105343028780231 7:-0.007340232468413754 8:0.7672537187738411 9:0.3538906829188173 10:0.35387524540947646 +-0.48115211266405217 1:-0.17943755364759517 2:-0.1384979591151625 3:0.8425773648797268 4:-0.43234064993405097 5:0.919754442523921 6:0.8390197802990036 7:0.43890653121452683 8:-0.7647648217789051 9:0.14770258954363835 10:-0.6681813635676657 +6.965069440749298 1:-0.9158261471030473 2:0.5228494114644282 3:-0.07760531122743153 4:0.6154296244963067 5:0.5231830145381096 6:0.4892535590799165 7:0.1987053183082137 8:0.9995670294711712 9:-0.2020375688074112 10:-0.7853579334836087 +-1.6896486293598596 1:0.4638529147853421 2:0.0953805943546191 3:0.8506904243225251 4:-0.028262644692445438 5:-0.9462342015500664 6:-0.6934738957112123 7:0.601125018257533 8:-0.04871041957758315 9:-0.015245062056267411 10:0.6119856200040805 +-1.763729644326212 1:0.5376618752928528 2:0.8062119856717131 3:0.44996834959923593 4:0.9917728248530817 5:0.5974717482179492 6:-0.406972851600659 7:-0.8523198502065281 8:-0.3076377139692321 9:0.9099974915864462 10:-0.43374966692373484 +9.012829566937228 1:0.6885456531832366 2:-0.0631164354373237 3:0.8394182300770314 4:0.7207913383891218 5:0.4715324450375691 6:-0.34417503908167757 7:-0.31448279255342126 8:-0.020591617987411936 9:-0.37668573574418107 10:-0.6528048324896532 +-15.951512565794573 1:-0.6112828771933607 2:0.4867007149846869 3:0.863494046941478 4:-0.7292072742454481 5:0.6338749652624007 6:0.5980798993978542 7:-0.5119002889878654 8:0.8394383182101366 9:-0.1412423080445726 10:-0.15838730884968655 +-0.29622788243318465 1:-0.9436253326661384 2:0.2907259958032098 3:-0.1530538226933904 4:-0.6174176535420375 5:0.8209632215649141 6:0.5060548803172731 7:0.8212448453211292 8:0.33506684706740386 9:-0.5408309869188785 10:-0.8105966349150977 +-7.683213587039055 1:0.2525015766703558 2:0.6417869320191234 3:-0.7569571597336913 4:0.5265130776924394 5:-0.03992944660560949 6:0.18292946303778823 7:0.4286344960738724 8:0.9158523573288766 9:0.5039796366711773 10:0.27660486075533797 +3.9061298856792797 1:-0.6501789225392032 2:-0.6040685518173872 3:-0.6448094322678659 4:-0.2019498832769746 5:-0.5302977370883424 6:-0.010754341856880067 7:0.8791702222974846 8:-0.2283571791337704 9:0.4726320486679656 10:0.3413255179758332 +12.928385148211825 1:0.7793178379505685 2:-0.5207562047491976 3:0.37253320760898934 4:0.7540757518052998 5:-0.679378421540417 6:-0.11966022036636881 7:-0.4317798870297489 8:-0.004211291952602059 9:0.39024653887361693 10:0.45391057946097146 +5.787566514603203 1:-0.20596730554338039 2:-0.8840796727164746 3:-0.749416279057892 4:-0.5511023306046077 5:0.9941631901218697 6:-0.09907966722992234 7:0.701617914811792 8:0.9696055014561289 9:-0.7083648075748707 10:0.5781111533720358 +5.701262468657861 1:-0.7066995012593675 2:-0.6756815056791965 3:-0.5720277255842998 4:-0.09218662060241067 5:0.21494136076896653 6:-0.37012884573008153 7:-0.6828277646796448 8:-0.10038134655965236 9:-0.46253754509583356 10:-0.20813933595648115 +0.9473494330088033 1:0.6876806675510589 2:-0.9530860102792402 3:-0.4043172626863887 4:0.6696455505098386 5:0.17863581804857254 6:0.1944646561635497 7:-0.5283662172535679 8:0.4872263841818012 9:-0.2882651789318431 10:-0.06293411605141874 +-2.6834375589185675 1:-0.22376759986120187 2:0.36555755546798885 3:-0.5223502955721961 4:-0.20702347869224624 5:-0.7745351063999764 6:0.22879328233099971 7:-0.5440007473902635 8:-0.6959483071829207 9:-0.131433881760733 10:0.2764225554693165 +-3.2766108642276146 1:0.0304613976530983 2:-0.3148062986719251 3:0.24950420590071953 4:0.7152023826801459 5:0.9656885739650887 6:-0.3210562623763835 7:-0.7305896664502614 8:-0.49074917893875836 9:0.7802670253347352 10:0.8667409958355992 +-1.1838791995691869 1:0.06642047806096318 2:0.5336148776806793 3:-0.6199614859883396 4:-0.15342280723497237 5:0.8407250402808968 6:0.7060811811107444 7:-0.2913182140909305 8:-0.5925203360011633 9:0.22644925021629692 10:0.42395071889002467 +-1.5856680515554806 1:-0.8724712788102853 2:0.11445744032031424 3:0.5483166457680566 4:0.9469521544884028 5:0.2541682828467746 6:-0.436750733871873 7:-0.9001249399695319 8:-0.7555793441458385 9:0.06946992897983018 10:0.9724148045760346 +-13.039928064104615 1:-0.558607026518148 2:-0.7356765018678253 3:-0.7547644426290201 4:-0.24898664843938745 5:-0.3606374046883567 6:0.5836652368902306 7:0.8497678666873467 8:0.21331875915717635 9:0.3558733809635668 10:0.9642603628738968 +-17.428674570939506 1:0.8562209225926345 2:0.7077202100653552 3:0.7449487615498371 4:0.4648122665228682 5:0.20867633509077188 6:0.08516406450475422 7:0.22426604902631664 8:-0.5503074163123833 9:-0.40653248591627533 10:-0.34680731694527833 +13.886853032969585 1:-0.6354915752033683 2:-0.9132338112681755 3:-0.4816479770266455 4:0.5448417181244594 5:-0.6250746297187781 6:0.7410618768880199 7:-0.18029029550083675 8:0.777358236920447 9:0.9625064189449102 10:0.048040935468046 +15.61684729251139 1:0.2980237970192188 2:-0.8160931971814265 3:-0.29649852157138445 4:0.3896688599904572 5:-0.17552110506337826 6:0.8721328328445139 7:0.48984799668438916 8:0.9984496052876473 9:0.9665885195526289 10:0.8966559812150274 +10.33625540376971 1:0.09939495068155724 2:0.9790332181038015 3:0.9483428886275702 4:-0.5717299810793317 5:0.4876405069057712 6:0.163962913892302 7:-0.4095537988924203 8:0.8608269751255508 9:0.010028680058212114 10:0.9095786494455713 +9.706032970113723 1:0.7687898546315146 2:-0.9825109379412285 3:-0.5423211794439926 4:-0.3099509487314134 5:-0.11561305536236333 6:0.9012327035409926 7:0.5257495475790148 8:-0.33804422025989433 9:-0.144428735681567 10:0.28019332199039604 +6.189043888072968 1:0.13246655756059478 2:-0.751192382628302 3:0.2233421456265161 4:-0.1933575076984373 5:0.8681727702736863 6:-0.7656847407654899 7:0.1033145549916572 8:0.33909210370257403 9:-0.22241363302770267 10:-0.14479004187830435 +-8.680225911784335 1:-0.07718769939880432 2:0.6702228057326558 3:0.6647810334933819 4:-0.05115658747070784 5:-0.850780588302118 6:-0.040961453376221924 7:-0.8407690297644956 8:0.33775829053563156 9:-0.45421556034898547 10:0.8238500771967823 +-9.42898793151394 1:0.8925906426831107 2:-0.6771269725125597 3:-0.11635105688280678 4:-0.7266044201050157 5:-0.6902918845825077 6:-0.5911234800910024 7:0.49395074569300657 8:0.43660804414878274 9:0.8736983081269782 10:-0.8001177058312081 +8.486245765579415 1:0.5614295382716652 2:0.3972427851719582 3:-0.276268504977494 4:0.7803448249454739 5:-0.358957923558495 6:0.3477822689529795 7:-0.7944805581842691 8:0.8356932134547437 9:-0.4783293647580624 10:-0.2522633417723845 +-1.8722161156986976 1:0.11831037290857482 2:-0.7309091607574014 3:-0.7339122716951587 4:0.2046641765436359 5:-0.9914679283125301 6:0.13518339528098555 7:-0.9760821540963867 8:-0.6080636193563043 9:0.3890502262427238 10:0.33864957953815145 +0.5122357093733743 1:-0.9555852441641726 2:0.4754771858792488 3:0.3743376249200432 4:-0.2651772997462427 5:-0.7915484529586028 6:-0.7575915279708862 7:-0.10432268807273859 8:0.021604934223709238 9:-0.6458011732912265 10:0.40773716196391674 +-18.845922472898582 1:-0.6031480148285926 2:-0.8736524730197766 3:-0.311456616524979 4:0.420921703897325 5:-0.2904011177124777 6:0.6683252350591937 7:-0.3436202976676894 8:0.5023604359385605 9:-0.33056149241985633 10:0.5168854058825227 +6.492106438811399 1:0.7824832256885428 2:0.6105456307389117 3:-0.0436873997963223 4:0.46730493583332855 5:0.2057529813440686 6:0.5738310686722767 7:0.6307964411259019 8:0.6208424783086652 9:0.8931894299284251 10:0.7164648197763028 +-1.6472226859532182 1:0.8854767145642171 2:-0.8175744681485637 3:-0.14894858038610903 4:0.9667400540136402 5:-0.3575837217508149 6:-0.9211342680517054 7:-0.956785876301889 8:0.6558217028031554 9:0.8014538160668165 10:-0.9475520920917395 +0.185861229793925 1:-0.8181719548530746 2:0.9990094335332504 3:-0.8195848911987829 4:0.6991933015233858 5:0.07295718417836583 6:0.5968996100546737 7:0.4871410306452193 8:0.2980483098540927 9:0.779953293728507 10:-0.7978867112395516 +-5.973450525185694 1:-0.975435413991927 2:-0.7832951303253313 3:0.5098999023442101 4:0.46795978867990007 5:0.2538986807863044 6:-0.8182887550010198 7:0.8335391734637112 8:0.4286082996234335 9:-0.1726765956719154 10:0.7649845978453362 +-12.773226999251197 1:-0.383327656965585 2:-0.9439560491389036 3:0.25039001869622446 4:-0.9342091044843222 5:0.8711023711291135 6:-0.6027135241543655 7:0.9456874780319795 8:-0.243290468946338 9:0.625765915285031 10:0.5160550067618355 +24.290551295953957 1:-0.8368553572749229 2:-0.5859456648150321 3:0.873779532007048 4:0.7462623178738954 5:-0.08133011570245352 6:0.36767541461776676 7:-0.33129619282275047 8:0.6104289727615573 9:0.9416581563055089 10:0.18201841676606856 +14.490247980976621 1:-0.4765937762114507 2:0.16430711839945555 3:-0.526776940706293 4:-0.6802269991653915 5:0.40748236413299344 6:-0.500290139207977 7:-0.31915972151663885 8:-0.4586068416002418 9:-0.15572660263944127 10:-0.32925702602833073 +8.377230871265601 1:0.44141613060964846 2:0.1582267687752743 3:0.8760950367284166 4:0.40434058393690364 5:-0.7063758409891474 6:-0.616055773516162 7:0.996372393127579 8:0.6142084876085476 9:-0.528320587432094 10:-0.2815909691094802 +-3.2987560995836653 1:-0.4600479783378091 2:-0.04201794336103326 3:-0.8934505203905587 4:-0.44991326751905536 5:-0.5220579476363783 6:0.46060949186328703 7:0.9169289030735643 8:-0.022458426893944283 9:0.08100795210565637 10:0.5726732415540354 +0.3422568955736137 1:-0.9888686059817204 2:0.22752298580182706 3:-0.5048696915520232 4:-0.059433420464226616 5:0.7823831512651716 6:0.9865977573980389 7:0.9164100011124972 8:-0.3638554550863984 9:0.3038282907667611 10:0.4652367033461571 +-8.24116881862084 1:0.7565819250331731 2:-0.3733277500524168 3:-0.8841150081071696 4:-0.922282989989148 5:-0.041520813551309876 6:0.8615967014876558 7:0.8474207144091339 8:-0.7518437864641427 9:0.45076605239968837 10:-0.48912984167595375 +-4.367083147104942 1:-0.276459380002813 2:-0.957555271384241 3:-0.3761632810202544 4:-0.3897414804149022 5:-0.3133861519856074 6:0.0777990809172171 7:0.6638552243422928 8:-0.3477312155364247 9:0.5934885465182675 10:-0.5238903641193555 +1.9280240152322783 1:-0.40051093785549696 2:0.5070348672240661 3:0.7506759969575532 4:0.5042104954516786 5:0.9959688260926507 6:0.4657024999761399 7:0.910611131925299 8:0.9836517468598804 9:-0.6263172749113686 10:0.16955852322929155 +8.918138317441574 1:-0.22407391224687023 2:0.5545084933214972 3:0.6335932367683528 4:-0.2786481116648991 5:-0.9549992830441785 6:-0.5577873948545062 7:-0.960657200286197 8:0.3709573488946196 9:-0.9191180485753339 10:0.5033478020271929 +-5.657796797481157 1:0.6359910361030725 2:-0.1742637774815281 3:0.39699327107265137 4:-0.9841991491194473 5:-0.622093571871533 6:-0.5433497301426455 7:-0.6731178481686009 8:0.930615153085582 9:-0.3065877908950827 10:-0.5456093749639228 +8.697079562319692 1:0.4815820396629933 2:0.1173457441514223 3:0.7313645402039386 4:0.3354835387237334 5:-0.10300554535074702 6:0.5116687640761355 7:-0.8850803659104614 8:0.10654026377571157 9:-0.864976708975602 10:0.01345035085413615 +0.033954684723234596 1:0.6703241653088159 2:-0.13447915740201166 3:0.026022550037831937 4:-0.5145659862194116 5:-0.6963587636078901 6:0.652083884947352 7:0.22644722530715278 8:0.2671580129293405 9:0.9659035105360283 10:0.9547989197693989 +7.359108382166921 1:-0.6855762478384229 2:-0.7543318537260015 3:0.4772611975128618 4:-0.5588002332845741 5:-0.24271386844336496 6:-0.28595644325868896 7:0.8732728098501104 8:-0.8026384804471058 9:0.7589508830210041 10:-0.9992933613402135 +4.953597303754355 1:0.8915633023548608 2:0.04688596266450751 3:-0.26866754730613374 4:0.16694236975718102 5:0.23465297255622608 6:0.36488427850844407 7:-0.06717041145276781 8:0.9470029805221898 9:0.32483835237272674 10:-0.7892521260150298 +0.683536559775105 1:-0.32176084249781556 2:0.5446298870866526 3:0.4095848716057642 4:-0.42579711490120187 5:0.4482850543749355 6:-0.0982243826242506 7:-0.9190317048427039 8:0.06234509402976718 9:0.21327512416175054 10:-0.38023673796734525 +-28.571478869743427 1:-0.4597184465402242 2:-0.5489429386926741 3:0.33422914572951634 4:-0.15992695377395516 5:-0.7310003311728188 6:0.18241063863467488 7:-0.48385214010599453 8:0.08139879039334552 9:-0.8401239538877046 10:-0.8896372220209929 +-19.884560774273424 1:0.4619217451285318 2:0.28157115824800005 3:-0.3829811521605375 4:0.5802544015450464 5:0.1117061271473403 6:-0.8926034502584623 7:-0.34862293810401956 8:0.2733254857260612 9:0.6514176550598809 10:-0.02758604919357066 +-17.494200356883344 1:-0.4218585945316018 2:0.15566399304488754 3:-0.164665303422032 4:-0.8579743106885072 5:0.5651453461779163 6:-0.6582935645654426 7:-0.40838717556437576 8:-0.19258926475033356 9:0.9864284520934183 10:0.7156150246487265 +-15.86200932757056 1:-0.6341453831788726 2:-0.9259180639727085 3:0.302702923864538 4:0.749555004323947 5:-0.7932989575334761 6:-0.5620972938631934 7:0.020542041027870717 8:0.11610338700447698 9:-0.7912600154897766 10:0.5108307672038874 +9.027804254487519 1:0.1746878011084212 2:-0.5872807344913673 3:0.6018547246457264 4:0.5106104933121229 5:0.7329523371170135 6:-0.40058771577765895 7:-0.48753463550174025 8:0.34308791976318 9:0.3407668956765344 10:0.5964472848798394 +15.949172086880687 1:-0.7790584545657173 2:-0.017224094786103317 3:-0.0974907790179953 4:-0.10287391996036166 5:0.6007953354774878 6:-0.7032497754397848 7:-0.36068070856329437 8:0.021391994204512432 9:-0.6509100388083549 10:-0.5410899936281377 +-6.151586699415245 1:-0.5318094974022525 2:-0.830796057445983 3:0.603828597318087 4:0.6660892552257192 5:-0.18529748408390523 6:-0.47166833767648986 7:0.592915541856605 8:0.9944601563352204 9:-0.6981606574244703 10:0.34942553665003584 +2.010398523297265 1:-0.9293899922307269 2:-0.07588009904844029 3:-0.8500855420709359 4:0.12191867923536615 5:-0.528778681165414 6:0.3117086447237414 7:-0.4222963938187163 8:-0.03247894950300623 9:-0.05387792412717962 10:0.4053568741659812 +-6.749023248121471 1:-0.9875370165216966 2:0.7137693455001415 3:-0.2510160963160164 4:0.8732150877079123 5:0.49658934612905314 6:-0.9817012857861731 7:-0.2045309437850289 8:0.7562713668333418 9:-0.6787434327188155 10:-0.6147932888026117 +4.452639829999693 1:-0.35256148944834176 2:0.7581152951164591 3:-0.37755890552299265 4:0.9480813371197343 5:-0.3419340388717347 6:0.3487602851799074 7:-0.5576726724270562 8:0.4899696188087421 9:0.563074979676983 10:0.7865891460062227 +-4.938733988900586 1:-0.4108386466193119 2:0.3287655432069885 3:-0.5853553038038923 4:-0.6480591422742821 5:-0.4787998161299789 6:-0.5828003484675421 7:0.42835744317623003 8:0.8378098987706633 9:-0.5645180498703375 10:0.28981512694646705 +-3.373242544176224 1:0.04989033652617936 2:0.6575826440927308 3:-0.24028051935833128 4:-0.6649808138961095 5:-0.6530198970442704 6:-0.19331254127919362 7:-0.6743004878881749 8:-0.7214986105015062 9:-0.30648035516261385 10:-0.6455097687924254 +-3.2843694575334834 1:-0.3548536057581908 2:0.7350125943559394 3:-0.3635282827378974 4:-0.8552820154885781 5:0.9140879208466111 6:0.21870365067770892 7:-0.17738543429561382 8:-0.052851966578491005 9:-0.36066059517759097 10:-0.9020765799355679 +-3.277146077677404 1:0.910961221014513 2:0.4302525202590246 3:0.11079959840001119 4:-0.3614188274820125 5:0.5080231397310961 6:0.013940825892631237 7:0.33583012240022403 8:0.5008797094229163 9:-0.663083147090173 10:-0.0865028013627418 +-0.202246147968096 1:-0.4929308143227653 2:0.8374300027105082 3:0.08763999085193186 4:-0.499738438136623 5:0.5926071511295365 6:-0.5135396038023627 7:0.6946715869746543 8:-0.5184428793490325 9:0.21753085495829239 10:-0.33796308746585235 +-7.1237150573506955 1:-0.8506203499039495 2:-0.6581804183622855 3:0.6484205342724825 4:0.013914696389758285 5:-0.6214530117645831 6:-0.011163110491807293 7:-0.6025372583334574 8:-0.0371573886520411 9:-0.7933455929226487 10:-0.38653838674273455 +6.298226129171093 1:0.7304191211928768 2:0.8128475475660479 3:-0.03161148630216015 4:-0.6018899317958344 5:0.19277055729934367 6:0.3002272616310928 7:0.949169758830406 8:-0.1011823256970481 9:0.16093341376629966 10:0.9596833606094763 +14.906594657519511 1:0.5053240355803015 2:0.6775698974866082 3:-0.6194771000646291 4:-0.02876927004033525 5:-0.5481504206112477 6:-0.9239150546263386 7:0.471216755072994 8:-0.0027794620943384363 9:-0.8954411386878227 10:0.8991742143686698 +2.1710965297686267 1:0.4578509053930304 2:0.9270194505165124 3:0.22470373699901236 4:0.21526179917432753 5:0.5299563895862103 6:-0.5824108997775908 7:0.03801922095671095 8:-0.5164033454609385 9:0.4370246809487237 10:0.6514133050988229 +15.05806598279517 1:0.48645077410559057 2:0.7821442063987365 3:0.1943681666933883 4:0.8289246958621577 5:-0.08034311437806041 6:0.03709694472527203 7:-0.895481297246602 8:-0.42921579749551664 9:0.5447075872378688 10:0.844397849728866 +-0.4683784136986876 1:-0.5083135683360327 2:0.626070365769088 3:-0.8737725909401557 4:0.725622293853621 5:0.0018794384199978253 6:-0.9343604622552886 7:0.6655593328822609 8:0.47501755618845753 9:0.8388618477210947 10:-0.5143806767304449 +5.823027255871114 1:0.08635467091841886 2:0.6314532702073175 3:0.8862069437865836 4:0.6542025864928516 5:-0.6846784290231471 6:0.048487096050569445 7:0.30828004933669395 8:-0.49438881988995687 9:0.5706936923061823 10:0.037705651885639346 +7.03779380408974 1:-0.07193682621291098 2:-0.5816975957307158 3:-0.8426927090342973 4:-0.37504851992255306 5:0.4473129018316815 6:0.3101938194888525 7:0.6160050428837607 8:-0.913998555949695 9:0.40461966540531313 10:-0.7581141330823786 +-9.770500546345563 1:-0.31358873581579894 2:0.11771478839130278 3:-0.3404842110585631 4:-0.0604362797252429 5:0.2159524972176814 6:-0.24737863017398087 7:-0.8541428610709716 8:-0.06753562283135062 9:-0.11567537916769255 10:-0.5606246203677223 +20.000154367451547 1:-0.344717847914646 2:0.8454969480099985 3:-0.58856299370874 4:0.5884510299634649 5:0.49162879631128553 6:0.7958075013181658 7:0.7781911267315837 8:-0.6780885011989877 9:0.9797694629597928 10:-0.1872163682079866 +-6.239848349456753 1:0.9132793720646253 2:0.1680340663118458 3:0.01740115925682284 4:-0.26580395408599133 5:0.28551914590761074 6:-0.9939706142381568 7:-0.8740927279520219 8:-0.8731218126652498 9:-0.10993630739903892 10:-0.3069565039708746 +-4.173072569004537 1:0.7864835254860851 2:-0.5614522227484218 3:-0.7718396381376464 4:0.49508673889127985 5:0.24030155936964714 6:0.8080778221819038 7:0.05395496402881128 8:-0.3045148076729973 9:-0.6134406357458853 10:0.7447268183581948 +-11.328415936777782 1:-0.10183127796258096 2:0.5689039487721601 3:-0.07015335898840225 4:0.23254189629731292 5:-0.3226974656715038 6:0.2859450214054784 7:-0.4916677058012495 8:-0.27564895614732055 9:-0.9416483232894219 10:-0.7472248333434015 +8.719164753818454 1:-0.8231424386390782 2:-0.03953537069863633 3:-0.3271580541537027 4:0.892192314973022 5:-0.6759017192358232 6:-0.419591686354591 7:-0.23967385135363606 8:0.936992531568956 9:-0.12946409158671512 10:-0.9082863469271643 +22.31738046492344 1:0.37030851555335365 2:-0.06654751559177563 3:-0.5759425437665169 4:0.9179952251152963 5:0.8628921839116359 6:0.8421952184405965 7:0.9625804174561126 8:-0.03075332253237728 9:0.12227386374957994 10:-0.6243390357793757 +-1.189108450798179 1:0.5681776913545951 2:0.46049028271139436 3:-0.366463711956754 4:0.025856437432560275 5:0.7547565372954261 6:0.5506193192167212 7:-0.6279807084274867 8:-0.38698884324386107 9:-0.9885778854008227 10:0.7814740172261654 +2.8767042393531965 1:-0.6841229745503388 2:0.6252203895646273 3:-0.6737644654353572 4:-0.7321040107741059 5:0.3162570540986238 6:0.6211089085315002 7:-0.33984617437403464 8:0.1227089818682312 9:0.04586594421613177 10:-0.4679977358965799 +2.783332151730615 1:-0.39148258540779013 2:-0.3037233649803406 3:0.7955133548911926 4:-0.1729544208044842 5:-0.18247049275020033 6:-0.1315085429729259 7:-4.447133918370483E-4 8:-0.805837119503338 9:0.11574866650006688 10:0.8517519041042676 +-8.99205564094827 1:-0.45501536967706535 2:-0.35829694693457914 3:0.775695048377375 4:-0.25331195582275745 5:0.15524612858817055 6:0.7400717904631442 7:0.8382485596668376 8:-0.5619009369436814 9:0.4386801597659249 10:0.09960232210246622 +-9.808386702564658 1:-0.987404834666963 2:-0.6732308850750186 3:0.5528285725528492 4:-0.8796302275267409 5:0.30705569958232193 6:0.8635312232105203 7:-0.14033675947074187 8:0.5516086773506235 9:-0.7487899106678442 10:0.8851518933134919 +4.948281656077033 1:0.4331269064492329 2:0.4628446087354616 3:0.33730748244242537 4:0.3473124014683382 5:-0.1707966473106064 6:0.8558057784524846 7:0.1390312032172829 8:-0.7918343112673001 9:-0.85993782695915 10:0.33563174747577107 +10.791261476321019 1:-0.5417345768902055 2:-0.06334901799780424 3:0.027652223245870466 4:-0.9881487640651161 5:-0.19441123027957707 6:0.40295156581142355 7:-0.8315553696517317 8:0.11405283165483926 9:0.5377980570161418 10:-0.24581620554740824 +-0.7287230169119936 1:0.33985587202063283 2:0.6841261099887705 3:-0.9441564997438197 4:0.28660913255058906 5:-0.7597915572726905 6:-0.8535957517473378 7:0.609134673753593 8:0.29636368731717977 9:0.05791523580926916 10:0.5589907965230858 +-26.805483428483072 1:0.4572552704218824 2:-0.576096954000229 3:-0.20809839485012915 4:0.9140086345619809 5:-0.5922981637492224 6:-0.8969369345510854 7:0.3741080343476908 8:-0.01854004246308416 9:0.07834089512221243 10:0.3838413057880994 +-16.71909683360509 1:-0.24375714099465773 2:-0.11915875769929496 3:-0.3741442802364221 4:-0.3812947578178094 5:-0.7032156297055756 6:-0.18339122712542388 7:-0.8634662520461855 8:-0.714561692659166 9:0.020558676493369177 10:0.22804428969949986 +-8.822357870425154 1:0.39332200105884363 2:0.5652370435795515 3:0.6220479966351453 4:-0.018976695481651484 5:-0.6868425195058918 6:0.2029750380170401 7:-0.5550873767310935 8:0.16864133648532342 9:-0.008843355054633628 10:0.6472547984399621 +0.36392761004065594 1:-0.9059630492963144 2:-0.41039282402227384 3:-0.006673269562094131 4:-0.4989314017618798 5:-0.17726034513032318 6:0.037764439388023874 7:0.30703957185016595 8:-0.09040426404909185 9:0.38661451965066274 10:0.1630571642147851 +7.415902871490132 1:0.188586850708651 2:-0.33013604761672566 3:0.6667976416858177 4:0.8537064956198137 5:0.03971370422819254 6:-0.43229195778759966 7:-0.9607154505216515 8:0.8413204878098277 9:0.40010565279599897 10:0.7306602852367441 +-4.129456164370826 1:-0.7967510984807558 2:0.545111159425699 3:0.16038228447433012 4:0.6311115528116698 5:-0.01985759480036542 6:-0.9516543115476572 7:0.18022912194075458 8:-0.2177157123823752 9:-0.5433158910016767 10:-0.4603867691069983 +-9.211066571082247 1:-0.3611235296125135 2:0.1402619601475985 3:-0.23132525512647795 4:0.5534401725834837 5:-0.34978585787763206 6:-0.24147682088922773 7:0.8089009287617064 8:-0.09075864922490862 9:-0.05759391404550773 10:0.3371306765964468 +6.52392916461972 1:0.19122050285976044 2:-0.625453376800498 3:-0.26804961781489856 4:0.9669297468261109 5:0.9142504122291741 6:0.7678963028488108 7:-0.6852943621882759 8:0.5898129788981794 9:-0.6580947533327339 10:0.46875109532259396 +-12.46765638103286 1:0.35148385951742633 2:-0.5206883134357769 3:0.35436280451876345 4:-0.8837833467474128 5:0.3433887284719144 6:0.3914771858025621 7:-0.17813796710416252 8:0.6553344538056296 9:0.3721548243590813 10:0.9442185832979726 +-4.937258492902948 1:0.9150659354384785 2:-0.17085510578573548 3:0.8233227233543232 4:0.2539669132090434 5:0.18955049451212935 6:-0.2833188558310358 7:-0.48483747414616496 8:0.8917378487725669 9:-0.13169122011498646 10:0.9815059855284158 +-0.5233425797210233 1:0.4238363705720569 2:-0.18363058784066522 3:0.2949874786744968 4:0.12235592695567354 5:-0.9746310186182559 6:-0.8990867637441311 7:-0.8580982328464586 8:-0.7930887027205957 9:0.16757307988090275 10:0.988861929608575 +-11.904986902675114 1:-0.3692990475534952 2:0.32166293883244323 3:0.3401547722249436 4:0.10009747375878408 5:0.7598877208920192 6:0.2853003389082669 7:0.22880221701675074 8:0.4521491122351502 9:0.33222018268933895 10:-0.9500018867461919 +8.324969054805921 1:-0.48086111720736513 2:0.3705524122401185 3:0.43635448766342133 4:0.6544321903349255 5:0.059000747296945155 6:0.3328036763371236 7:0.9609146376298034 8:0.5943082361322021 9:-0.3074246170581105 10:-0.6763916655761453 +0.21701641918233017 1:-0.29449708766806304 2:0.040640346437143426 3:-0.6524819533513639 4:0.37482287233702394 5:-0.29800608396043216 6:-0.537030944860492 7:0.2862394027536084 8:-0.3783043133672048 9:-0.5292179323972728 10:-0.09583783955916791 +-6.84977373580439 1:0.825136109631339 2:-0.5722868691442817 3:0.11048134523744757 4:-0.5946054293068455 5:0.28061485657354823 6:0.9135611623885838 7:0.35590421873954603 8:0.8943562249941011 9:0.4183378981109729 10:0.5714160298247304 +-11.039347808253828 1:-0.9620263418414967 2:0.22669065740934724 3:-0.7378036492234086 4:-0.4460191511609126 5:-0.2594476006347024 6:-0.989879976130936 7:0.762096015449097 8:0.6983868222083149 9:0.8729993459982626 10:0.3426647417451305 +-5.882860061103163 1:0.5247178959769465 2:-0.6217169944869176 3:-0.13640714414758315 4:0.6608201052790283 5:0.5789945243704264 6:-0.12686057623612612 7:0.7277882307863026 8:-0.47949544949858236 9:0.9781208432412936 10:-0.8980068284379361 +23.52945433069272 1:-0.12339549394875426 2:-0.6769524283089239 3:0.9324962870874394 4:0.28956947294105206 5:-0.2957355479338608 6:0.7504385350771912 7:-0.8769262306643106 8:0.41591311300668155 9:-0.7694611231426498 10:0.9885110924181837 +19.043184423383824 1:-0.13783178628851878 2:-0.853631844645959 3:-0.12792415583066052 4:0.6936898387576049 5:0.8488563282318959 6:-0.6530521292304581 7:0.27832187660440666 8:0.09838048719062442 9:-0.5913230087557231 10:0.260839433107553 +6.83105883806984 1:-0.9085282656519695 2:0.65203708247844 3:-0.687580071985604 4:-0.045008726377529173 5:0.4762107922777967 6:0.15939259525248506 7:-0.46363191848939334 8:-0.25856682230410266 9:0.313842004143269 10:0.5042938214484851 +-9.409197719620593 1:-0.34356198962701945 2:-0.06381545064099514 3:-0.9332814619122063 4:-0.2629675367181199 5:-0.03876014002851913 6:-0.4606936151803749 7:0.49272969757318563 8:0.5550196351479111 9:-0.1758425343811718 10:0.20285868144226837 +-1.3101852978323116 1:-0.3740821549570985 2:-0.9788976137554464 3:-0.6078739734947245 4:-0.8007745980271539 5:0.7381298546055934 6:0.7407750458109124 7:-0.7711351008178868 8:-0.9895256155202141 9:0.35793767138197174 10:0.6589909255086295 +0.5180809608973377 1:0.19289850282287446 2:0.6301214514538145 3:-0.15311307199521518 4:-0.8607670552113709 5:-0.46422067276745316 6:-0.29812862604449464 7:0.519464836430044 8:-0.9480450997338103 9:0.973503038633444 10:-0.7843880226794626 +1.9947872601406775 1:-0.15799682110486057 2:0.22645891561571352 3:0.3141842574216682 4:-0.36086019480721676 5:-0.1429373936064291 6:0.8097261636650581 7:0.11764088861630029 8:-0.9151998265501957 9:0.6536711690904891 10:-0.17232697113157425 +12.352290000973428 1:0.8176113135335772 2:0.39342616792621987 3:0.44011948797971234 4:-0.4412435869837865 5:-0.24509203724837314 6:0.8636655043434542 7:-0.4251583124505798 8:0.2068056615503988 9:-0.3501114760443049 10:-0.23701353324739483 +-2.891643319177732 1:0.7722403010820704 2:0.7994121584045861 3:0.18520464815273208 4:0.7273575609391227 5:-0.3758589216283552 6:-0.7598404862373955 7:0.5748649410179301 8:0.6897988099260968 9:0.5638920860629713 10:-0.992567809902162 +4.803737144054077 1:-0.7367711178556622 2:0.07370548192399351 3:-0.5510509754264419 4:0.11949095653894504 5:-0.7723751845800411 6:0.6450480728551136 7:-0.9508825019800493 8:-0.3250395411575804 9:-0.24913562167143777 10:-0.3617439870343031 +5.051689886526102 1:-0.09854955786627007 2:0.5298224514703289 3:-0.014996634675966236 4:-0.4462048687049027 5:0.22912790083984547 6:-0.513533454471272 7:0.1452771069237353 8:0.371152210841464 9:0.9204732090987018 10:0.7472990716905279 +3.8591142298280476 1:0.7532169023970261 2:0.8291433156934658 3:0.9255891263525324 4:0.3248663809949248 5:0.9905320652281553 6:-0.10383453745167626 7:0.8519246838852608 8:0.6024015353989258 9:-0.06958036249881938 10:0.5862142389541998 +11.30005914221598 1:0.026411858067972194 2:-0.6968445330429607 3:-0.8194566946165238 4:-0.12780659247925996 5:0.8406393783194903 6:-0.24617182945415128 7:0.30199973460219853 8:0.6062457235841974 9:-0.19314055910416927 10:-0.48313233883372964 +-10.288657252388708 1:-0.7388306404020344 2:0.07753617971873439 3:-0.5735498713988352 4:0.2183581175474576 5:-0.873572721679176 6:-0.8788755575751708 7:0.7087858362905568 8:0.7126712562404713 9:-0.7607334319316799 10:-0.4627367552114916 +4.895250842405817 1:0.9772954128558484 2:0.6020087399988574 3:0.16946626176056134 4:-0.011334492807484997 5:-0.5391845039589362 6:-0.4315843612118535 7:0.9065130011032458 8:-0.4860160207844919 9:0.0921755607946162 10:-0.022200673265013515 +1.0479421939727227 1:-0.055436367433274514 2:-0.6710483362647659 3:0.9222786043047919 4:-0.22005981623386184 5:-0.8141845044113469 6:-0.31766631447334226 7:0.6067696845798944 8:-0.1445661385071555 9:0.9172271611227454 10:-0.8079554780561127 +-9.754451457291598 1:0.533713237587885 2:0.6499588942067549 3:-0.49188790503368285 4:-0.6925119436487435 5:0.3345265979579788 6:-0.8117849521672496 7:0.9312055115656304 8:0.3273803451149724 9:0.7567478475677727 10:-0.6256676928549367 +5.869027126482974 1:0.7273823383600513 2:-0.2519813990388706 3:-0.8239584025397881 4:-0.13749750031735974 5:0.6142824732416132 6:0.6251630800232315 7:-0.6138240706157267 8:0.7210396245391326 9:-0.41832155201953714 10:-0.8965988320689853 +9.14234252751227 1:0.7295320896113133 2:0.6150271212503227 3:-0.9785024737101733 4:0.30006672036705506 5:0.11703528191771406 6:0.2971639460196238 7:-0.7920108995168815 8:0.32649036066184567 9:0.03522428067355543 10:-0.1766251898148803 +-5.643698771141404 1:0.27360638280623983 2:-0.6124401810442446 3:0.24950528730210886 4:0.09920211684887548 5:0.7187490549286091 6:0.6212724115415782 7:0.5864634211269566 8:0.114951165007104 9:0.44859258949094283 10:-0.3768352371578665 +12.781643819428492 1:0.9144335582094396 2:-0.4579872615218674 3:-0.6521934534632468 4:0.4462086111316512 5:0.240360283350179 6:0.23974046479581124 7:0.4840439971437822 8:-0.7250363120037027 9:-0.29769496257362094 10:-0.3382859512018359 +8.393556738722923 1:-0.8263387132502396 2:0.9434824094966923 3:0.1607861709872136 4:0.15217100448798782 5:-0.6517945935711484 6:-3.354731073326178E-4 7:0.07846631386981562 8:0.687844846942889 9:0.9277854407325892 10:-0.8855380268588307 +-15.348871155379253 1:-0.5734707274250155 2:-0.2526008551945753 3:0.23752094195309925 4:-0.7074613963298721 5:0.4674168537545218 6:-0.3198997855552628 7:-0.10415974108745596 8:0.5616912699671224 9:0.43742425558560694 10:0.19732530755184596 +13.138260063721448 1:-0.9415220143797984 2:0.6015431361268124 3:0.38898046240229545 4:-0.5750448371021175 5:-0.5803995196333898 6:0.11772198725731342 7:0.7512685244060366 8:-0.6683465740662857 9:0.9515652825318053 10:-0.32405935964523547 +-26.736207182601724 1:-0.47083104147202404 2:0.28748860067800597 3:0.007399318769021113 4:-0.8189013750589702 5:-0.5156633937248272 6:-0.9906928746525896 7:-0.8848419810272337 8:0.2197280161306785 9:0.12855082514870197 10:-0.7862803985146845 +-20.212077258958672 1:0.5609065808412279 2:-0.9201904391147984 3:0.908305865183735 4:0.9255146658282842 5:0.6871419344095282 6:0.4201876217923466 7:-0.42906289792612684 8:0.5787691868233418 9:0.7260522064761288 10:0.28251641556690554 +-0.44652227528840105 1:0.37640618494870504 2:-0.20012451052963542 3:0.9420894309510319 4:0.4218728633972739 5:0.5551974480349577 6:0.07615991810462619 7:-0.12409220462011294 8:-0.22212591926375946 9:0.21160498862483723 10:-0.6092792830633924 +-1.9481059746438067 1:-0.43820030250217457 2:-0.6836588417639442 3:0.733018205278934 4:-0.6564348753121718 5:0.7333385435136448 6:-0.5577457688360317 7:-0.31035811050608975 8:-0.7189201447768139 9:-0.7629842028723994 10:0.7179459779331092 +1.1951162998609508 1:0.19541555859727744 2:-0.4796785506546435 3:0.14123852670749248 4:0.7161847585887089 5:-0.2502765085719578 6:0.8815667909545981 7:-0.6418691905513725 8:0.49600147195728783 9:-0.3091837674381053 10:0.4320162841463153 +-8.99125390483227 1:-0.01183888602092864 2:-0.5901829024081027 3:-0.4343074406380647 4:-0.40450313056290166 5:0.05269590196351448 6:0.733631212862198 7:0.9575176715505025 8:0.5974628692830348 9:-0.20284241796038271 10:0.9577348510907686 +-7.955533026930219 1:0.6104830760481679 2:0.5915483572646505 3:0.3275427350991458 4:0.48361434056132424 5:-0.9466590639056058 6:-0.24662428438925743 7:0.9856361456534972 8:0.9434155212648045 9:0.3466736921968707 10:0.12927980558284102 +-12.500773785355054 1:0.5733321361720694 2:0.39154119830075085 3:-0.9347116355607772 4:0.0920586614926524 5:-0.6959457183810456 6:0.2136579936466858 7:0.17595268059814395 8:0.8828168055200465 9:0.18934277314853398 10:0.7565908584660754 +-11.43180236554046 1:0.082018621904135 2:0.9074181204118958 3:0.46125595008850273 4:0.40328845936169966 5:0.7803064691948824 6:0.20802011482729377 7:-0.41368899649077284 8:-0.8997565495498339 9:-0.1880483213318005 10:-0.15538597634233264 +-5.055293333055445 1:0.4442675297698402 2:0.19045719972922193 3:0.4877438951288897 4:0.7984474402420494 5:0.3251350777349489 6:-0.18676050499673869 7:-0.2701840041572374 8:0.4486609996458524 9:0.5403637876036615 10:-0.8971614841211264 +1.0276485382241776 1:0.7953696703382547 2:-0.3245779681908927 3:-0.3507435626548021 4:0.9510986059491036 5:-0.8655491074076527 6:0.20729233888498677 7:-0.43078300089533594 8:0.19504657032168216 9:-0.3173814102187291 10:-0.042479969052890754 +9.690201571311908 1:0.16852987139559206 2:-0.2514893273405625 3:-0.9993240281686275 4:-0.2166013247997891 5:0.33294165754921234 6:-0.5824203831560628 7:-0.15253642946648616 8:0.3547892367555441 9:-0.047604356104869794 10:0.9229112136183077 +2.2591036039970347 1:-0.9919593184325572 2:0.6323551392201245 3:-0.20815293136790447 4:-0.002395046469600759 5:-0.5015903362190326 6:-0.16698803749234048 7:0.7901657583805675 8:0.33755402936964973 9:-0.3707337678548108 10:0.6995480653730146 +1.5130881908855742 1:0.973710432688613 2:0.6518972988019702 3:-0.16491318496856833 4:-0.6066757853095415 5:0.8762371591845273 6:-0.9056066630820714 7:-0.3388079327070965 8:0.3934146060660142 9:-0.8756168865642253 10:0.9522427911640303 +4.023618949132531 1:-0.14974626191548301 2:-0.5874962377709136 3:0.6780439909311404 4:-0.37291203746764356 5:0.08104034602232169 6:-0.4706923395029945 7:-0.8924577368048239 8:-0.3363784341297067 9:-0.4139746050396018 10:-0.5107600309932907 +-2.8674162893420965 1:-0.7554383289076523 2:-0.1355597928418868 3:-0.3891904246986413 4:0.43949832438341785 5:-0.43859957095446833 6:0.37548094528561093 7:-0.5228633291549518 8:0.24169710795100352 9:0.7131753590746546 10:0.03458176767001042 +4.661164232198611 1:-0.12738868751385546 2:0.9446285809821182 3:-0.17981416859193433 4:-0.7535879975625193 5:-0.08594548726529161 6:-0.9983154486609989 7:-0.7272748852665216 8:-0.8197811039616518 9:0.5177610923333253 10:-0.6180731281817853 +-0.12347625601866746 1:0.10820547757674692 2:0.1825421454873002 3:-0.3412486258429426 4:-0.14925445930975534 5:-0.6594599831395103 6:0.9552502376248448 7:-0.7875626067291472 8:0.3854984181307912 9:0.014303876202374832 10:-0.7300443667550689 +14.546296184422973 1:0.2459523985646046 2:0.9434777073825811 3:0.2112745925235362 4:0.7730688005214974 5:-0.13727994893203732 6:0.6140037510172511 7:0.7545298281668846 8:0.7814551909982614 9:0.0026683642139069264 10:0.5633973602849358 +-19.66731861537172 1:0.9353590082406811 2:0.8768609458072838 3:0.9618210554140587 4:0.12103715737151921 5:-0.7691766106953688 6:-0.4220229608873225 7:-0.18117247651928658 8:-0.14333978019692784 9:-0.31512358142857066 10:0.4022153556528465 +18.84119697288412 1:0.4423204637505467 2:-0.4364821709544735 3:0.3935363893778452 4:-0.7750286735195999 5:-0.6981814766625978 6:0.6889512553826111 7:0.3646791168217727 8:0.0023536025493677837 9:-0.08378048150085249 10:-0.05659381771155503 +17.40329212914592 1:0.9155980216177384 2:-0.35593866074295355 3:0.44775710780914824 4:-0.42914421567532357 5:-0.2734430718503955 6:-0.8937042912745483 7:-0.3143761936611371 8:0.07805814979426184 9:-0.31386151509289784 10:0.6202932236456253 +-19.402336030214553 1:0.462288625222409 2:-0.902975525942725 3:0.7442695642729447 4:0.3802724233363486 5:0.4068685903786069 6:-0.5054707879424198 7:-0.8686166000900748 8:-0.014710838968344575 9:-0.1362606460134499 10:0.8444452252816472 +-3.855123203007599 1:0.5072557393175969 2:0.4626973233672753 3:-0.20910077161652119 4:0.9431415515135266 5:-0.1293690767585638 6:-0.2033835058111637 7:0.501429131658198 8:0.175133281735671 9:-0.6091682952201736 10:0.543010689352589 +1.493768355655548 1:-0.7772812666041105 2:-0.7743738591348672 3:-0.2848754060915175 4:0.3336846848765145 5:0.6219572132443736 6:-0.11144657683793624 7:0.7606913325884337 8:0.8547085151723017 9:-0.31728444617771134 10:-0.4668474022688931 +-17.803626188664516 1:0.5176340000264179 2:0.23048377874011128 3:0.6162746928601832 4:0.16908590014785418 5:0.9695207469685181 6:-0.34713218673384705 7:0.8526833760069625 8:0.9895592279649763 9:0.8805561957342884 10:-0.43452438291417894 +1.4060200157931342 1:-0.41964471941333525 2:0.7738486114171979 3:-0.0964606192284374 4:-0.25351781452566025 5:-0.21065389913054244 6:-0.40490416354122916 7:-0.7696501777959646 8:-0.7710488116813146 9:-0.6777228721053572 10:-0.09381158095961428 +-17.026492264209548 1:0.8367805314799452 2:0.1559190443625338 3:0.048200110551483544 4:-0.7340083467235765 5:0.2661150265782781 6:0.3881661781792165 7:0.9485287302765621 8:0.7201540574376382 9:0.8509234862656003 10:0.9658114866648093 +8.729450606651499 1:0.6404862166906327 2:0.16516090922657822 3:0.29013117743588057 4:0.37056732180613317 5:-0.3376494575302882 6:0.9012625630650577 7:-0.42150978319487 8:-0.05630249989686087 9:0.706104255632954 10:0.01935884085365225 +-5.516822117602276 1:-0.5718348423045241 2:-0.2145777722920088 3:-0.09307467998835195 4:-0.7311274103678378 5:0.5272184003067053 6:-0.00528176138162495 7:0.2852826178935919 8:0.6180999884045897 9:-0.7526372151008776 10:0.20416472532830543 +13.001541259752251 1:-0.5137703877272299 2:-0.15452359837207896 3:-0.25657600903152744 4:-0.9773110735601165 5:0.0718147980090178 6:0.18965211809311744 7:0.7795354990363292 8:0.21976898743223638 9:-0.20364089221752524 10:0.33822332985943304 +18.443388694564348 1:-0.9278344397401963 2:0.2678538727090136 3:-0.46932389854374734 4:0.8494176173177825 5:0.45765527018197694 6:0.20546395745879287 7:-0.199860294349123 8:0.47798730134403256 9:-0.2279771893187592 10:-0.30836118564314274 +8.952089112152663 1:-0.7371671220953286 2:0.8160149639986789 3:-0.026630089188139028 4:0.5931015267817183 5:-0.12216243475451294 6:0.161290795125286 7:0.7423016751095652 8:-0.5212872902985852 9:5.606147011660845E-5 10:-0.409626733921443 +-3.7062463981908027 1:0.5633514321449928 2:0.9914900963311462 3:0.1867799930236702 4:-0.15960235736142847 5:0.1204791067384241 6:-0.7733281422620872 7:-0.887447048141158 8:0.7931515335800692 9:0.732289882696125 10:-0.034992898370363124 +-10.58331129986813 1:0.6627003739767989 2:0.10688718810947728 3:-0.49230090744757216 4:0.8936580036513948 5:0.012227929286241057 6:-0.1442038886014838 7:0.9203452040795139 8:-0.20719832624131262 9:0.29561869366253335 10:-0.08597725084864649 +9.818996211259908 1:0.580133516885796 2:0.07422424429848573 3:0.33438634998226924 4:0.26054797992533696 5:-0.8771304726537796 6:-0.9710990591964794 7:-0.1869287393875041 8:-0.6167738073093247 9:0.34401921428837245 10:0.6737600514607418 +-11.87816749996684 1:-0.7193071334885193 2:0.5247127705364141 3:-0.02978727198197606 4:0.18353223007701058 5:0.40350110058596944 6:-0.36002841871228686 7:-0.20781535546501528 8:0.5517883176456557 9:-0.9938027872744732 10:0.6245061418135955 +-12.198096564661412 1:0.27542314155961156 2:0.3459734388741733 3:-0.38737776987446937 4:0.6244101669171684 5:-0.7801218302490938 6:0.20444733666197523 7:-0.5667599464182904 8:-0.9462131580071358 9:0.5576565405741785 10:-0.9307557040059242 +-3.6610413123521357 1:0.045569951437504086 2:0.32203961277046145 3:-0.04228927426053675 4:-0.9435304938416831 5:0.3750509710699601 6:0.21298970117620142 7:0.5491054691791977 8:0.33695088608872203 9:-0.9923500858828505 10:-0.6402707119893463 +3.782742149409224 1:0.7795250611996376 2:0.43296979846218275 3:-0.6481485005937841 4:0.3235717281667645 5:-0.8067382770768907 6:-0.06740397503468509 7:-0.2835017205434338 8:-0.5875853498478532 9:-0.25699561837680585 10:0.7813561594373908 +-17.065399625876015 1:-0.01772446594568744 2:0.563282914714494 3:0.14232420381013955 4:0.031667902604941345 5:-0.7815348482900619 6:0.3657733497576803 7:0.7208326162626688 8:-0.7863253120180662 9:0.3329194167867533 10:0.6175752945608013 +16.23248797654815 1:0.2615647748812251 2:-0.6631801348538622 3:0.6420349382574477 4:-0.31980528388089846 5:0.38021930887251365 6:-0.060298437830818896 7:-0.8911652782989568 8:0.3424617259589986 9:-0.8515350749364614 10:-0.42354709676980207 +-5.015963911416578 1:-0.07890564237014686 2:-0.09864377281008885 3:-0.13139943914680408 4:0.6610949669857866 5:0.06777579108221987 6:-0.26586245727222835 7:0.17443498956808612 8:-0.3129854922817781 9:-0.37913757211269505 10:0.7627186373372121 +22.647750304177556 1:-0.03666997412165163 2:0.49691867674483814 3:-0.45898559472166967 4:-0.09932248891016404 5:0.05692910907689508 6:-0.5006743461081364 7:0.9992936758550379 8:0.8252525466172065 9:0.9431711015127009 10:-0.4891497061921315 +-3.731112242951253 1:0.44353490207818513 2:0.23112032838224117 3:0.4697682541445527 4:-0.7507514828346664 5:-0.06323257550543837 6:0.0997091431243109 7:0.9394036761509628 8:0.4103869738859962 9:0.6564209227640914 10:-0.5427466755921158 +0.6761872737225261 1:-0.30051626190360503 2:-0.26699232020158803 3:0.8668758741279379 4:-0.40325291744583347 5:-0.9756425738484267 6:-0.5116398654634617 7:0.16424789009043073 8:0.8034099442414044 9:0.8554935001446193 10:0.42747702930667497 +8.449247195197387 1:-0.6588765973399024 2:0.2502285196526799 3:-0.20481547024283087 4:0.3770725284683252 5:-0.169707887761277 6:-0.0804075502584003 7:-0.3580757176408007 8:-0.6042549664471129 9:0.360349278976142 10:0.15899650901110962 +27.111027963108548 1:0.7106841652047162 2:0.6853699382312817 3:-0.8076297545289823 4:0.7932321056591545 5:-0.8011085095234463 6:-0.7017292726737878 7:0.10568649778064154 8:-0.40755358264969255 9:-0.061008981132773865 10:0.08895972651409556 +27.78383192005107 1:-0.8378790218922778 2:-0.6651002504721837 3:0.021049638665430415 4:0.32994334871293196 5:-0.7981304887988308 6:-0.2947962117284566 7:0.9739408711845776 8:0.9442893181893954 9:0.010541491359981059 10:0.8332791453382604 +15.700710963871254 1:-0.538773982400854 2:-0.5966426806845984 3:0.14570292467314627 4:-0.5937791901212952 5:0.7779251136963325 6:0.9962962075803357 7:-0.4774083823748394 8:-0.02528476957876369 9:-0.17305036341254398 10:-0.6013841506503688 +-12.558575788856189 1:0.03250364930617211 2:-0.6723950859659307 3:0.7090474884514901 4:0.25034305882632735 5:0.7036774024093582 6:0.622650236684523 7:0.5776881238206741 8:0.7999754726258337 9:0.21332972563833508 10:0.33849062947231645 +6.2776776518215955 1:-0.009605588630256623 2:0.5786496865369053 3:0.9208276908400748 4:-0.9477397424337148 5:0.6306053656362194 6:0.5396434662389846 7:-0.9841930450269964 8:0.5492682920407823 9:-0.020767248025529206 10:-0.8684655435686472 +6.424586997399564 1:0.861374923392324 2:0.8356037964367176 3:-0.7173479824827564 4:-0.6309584820438245 5:0.16136758138471285 6:-0.7485184163431866 7:-0.006053583829132236 8:-0.8762221084691306 9:0.19195377669247726 10:0.07259634302552964 +-9.64772485466405 1:0.7568015336230662 2:-0.4221524485756756 3:0.011711847664269248 4:0.7387065048724242 5:-0.04347512566745104 6:0.06642100869974654 7:-0.6993705848315939 8:0.16312217088045422 9:-0.11975577990989916 10:-0.6188717473788392 +3.8183706502283647 1:-0.7226937936463145 2:-0.5462756960199258 3:-0.39158419906610664 4:0.014310440945434433 5:-0.9950315917350652 6:-0.1844037449550875 7:0.9023517651879036 8:0.7948752060508435 9:-0.6792702010973877 10:0.40730074403235617 +1.1585019476700562 1:0.5575546848694 2:0.8997032130006739 3:0.6088643323129037 4:0.4872893656051758 5:-0.03977520372748922 6:0.3202565433572042 7:-0.31231768645537206 8:-0.6861153669592381 9:-0.08561643820383291 10:0.522243657731251 +-8.18651039877047 1:-0.809069379967462 2:-0.04827229852445103 3:0.19963602092982624 4:0.2568971171641006 5:-0.0015346733366310428 6:-0.6104625526166494 7:0.7746715041233412 8:-0.7343750018341593 9:-0.49272635466510106 10:-0.8115191199688623 +-3.377690136019927 1:-0.9408187510685164 2:0.9654993263332854 3:-0.16725010447984268 4:0.2574069587853294 5:-0.6930506968932861 6:0.11124762075550176 7:0.39145805505914866 8:0.2906495128462767 9:-0.27454907309824916 10:0.9001175309434777 +12.692571815413245 1:0.7404426710258791 2:0.9060576634778448 3:0.7023712021897308 4:-0.9808126157768493 5:0.03447666475715194 6:-0.4146339211599541 7:-0.7329651749553896 8:-0.2696019807317358 9:-0.9885367164723897 10:-0.8540304023043486 +2.5111054050889354 1:0.7448154454968356 2:-0.7532143233138027 3:-0.9724617436335079 4:0.662620399592766 5:0.45517204589358307 6:0.37409736074838684 7:0.337245076577648 8:0.50951903847353 9:0.2590369923587328 10:-0.3248257475117191 +-8.300340493749207 1:0.5504850435404609 2:0.5077232940244447 3:0.778859307357816 4:0.2601916883813373 5:-0.0032275666062382413 6:0.039752927221862855 7:0.19468432568826755 8:-0.2859531554546477 9:-0.4113477962970582 10:0.43272011953041667 +5.904938653193952 1:0.6622293273002955 2:0.6428891633785236 3:0.6999663090423285 4:0.9132698742913088 5:-0.3960072336866507 6:-0.14500922264286054 7:-0.4390171033743564 8:0.002067106212897185 9:-0.6079874251539117 10:-0.7131416109696531 +5.004048239623824 1:0.7212309895357449 2:0.3425199843383353 3:-0.7290323633040705 4:-0.5563097960397918 5:-0.7577898297822001 6:0.647883070472203 7:-0.23710559062843073 8:0.34398507133293954 9:-0.5440251617348038 10:-0.2971638032112218 +6.21255598077158 1:0.2498685983586959 2:-0.2586857335205359 3:-0.6380810501916263 4:0.17008841621855852 5:0.9485802018202867 6:-0.2580306792121272 7:0.032916516140567786 8:0.32950951532163675 9:-0.9291915084526683 10:0.8454021164786922 +-3.741044592262687 1:0.763300390779396 2:-0.1832552896771813 3:-0.39361907876758573 4:0.9050768615040607 5:-0.8850093869496836 6:0.9302208653737598 7:-0.12972094056755412 8:-0.459442486378308 9:0.5044112394875107 10:0.1399067554681861 +7.378402183384303 1:-0.27686808475610114 2:0.12735524561214606 3:0.5216635958678004 4:-0.9418584785460469 5:0.20441570818728771 6:-0.35073421178920583 7:0.7847501694079704 8:0.3222999552829353 9:0.21025696511089764 10:-0.5813710201294744 +-7.1500991588127265 1:-0.1945259148773102 2:-0.4089845159829022 3:-0.1971859124232922 4:0.9531447983295496 5:0.07996455700202221 6:0.17013529724757648 7:-0.2442095218739362 8:-0.8564146371721229 9:-0.5843910532907555 10:-0.33846471424918767 +-4.288417758202577 1:0.020710986120182184 2:-0.7450564238727908 3:0.3674992023059285 4:0.46737461414601555 5:0.9411702705113052 6:-0.7257365059912877 7:0.5813280037560231 8:-0.01567531846894843 9:0.24734195293533467 10:0.6516001002566887 +5.916426037500391 1:0.8260000862135342 2:-0.11324162495165968 3:0.13061304369435334 4:0.5762591624576425 5:0.548049763999644 6:-0.9751599851764361 7:0.02828821483057764 8:-0.4113286027346803 9:0.8912856976307486 10:-0.8470910204808244 +2.431004294471012 1:0.14088576701299083 2:-0.45104190898994734 3:0.29891134031619115 4:0.955503074037666 5:0.15962522624750242 6:0.7664481093046553 7:0.051697815479792686 8:-0.3471787155014081 9:-0.8007151537631465 10:-0.5598899500902301 +-16.08565904102149 1:0.3946137229565083 2:0.8443779319638349 3:0.5116855547320893 4:-0.5319339991982652 5:0.26564506849312797 6:0.18905397829944448 7:0.1976357098053687 8:0.15505612242632538 9:-0.935633748308776 10:-0.9782957013204887 +18.058440348477184 1:0.8402487524597533 2:-0.6200725197687718 3:-0.6158487677192792 4:0.0709328308135515 5:0.7501256905495493 6:0.38092209802839583 7:-0.8192579128383128 8:-0.9304002828581583 9:-0.6570300818845025 10:-0.5252554781538985 +-1.0026720160736349 1:0.46122079684901474 2:-0.7609201036934166 3:-0.9372178059537293 4:-0.25391036498391006 5:-0.7487429157699828 6:0.38024314675291637 7:0.21886059803198576 8:0.027516853267765207 9:0.33483464322377765 10:0.618580130027746 +-2.6688695419207162 1:-0.8775911623423445 2:-0.6647410420697879 3:0.05948516302547313 4:0.7278526664475804 5:-0.011366224409705028 6:0.33475665968289436 7:-0.6386120399761575 8:0.39609772177595115 9:-0.7872076290319412 10:-0.6195857302948329 +-13.867087895158768 1:-0.9114780602695882 2:0.7997695296649912 3:0.8337252417804881 4:-0.7927267913881113 5:0.6863829853181673 6:0.4162562153517635 7:0.2659922421074139 8:-0.551994669040742 9:-0.6403900338772157 10:-0.8680387717518072 +7.826011095515239 1:-0.2881951904396949 2:-0.19317071325391022 3:-0.06581062483451183 4:-0.6074074436315555 5:-0.9434740067975405 6:0.9426572655575483 7:-0.1812629432036228 8:0.39425575292939863 9:0.5065890539615039 10:0.8969825696966649 +1.4213836206303339 1:0.6996840540120932 2:0.1283999569152492 3:-0.2537375462472613 4:0.24772110606788456 5:0.9040210381745799 6:0.47062010977660207 7:0.9697678931927365 8:-0.9215764371674713 9:-0.27541598110075793 10:0.44277003247067803 +-0.973650798730175 1:-0.2121645467631068 2:-0.6770222508071349 3:-0.5733067523949165 4:0.27979529516037105 5:0.7128588235545461 6:-0.9208763636184307 7:0.14128337151047532 8:-0.002851660400375433 9:0.6943908711123281 10:-0.9201922993121072 +-0.17500848560451965 1:-0.6015070903427717 2:0.7815998200409671 3:-0.9932006200204946 4:-0.3303953411379028 5:-0.3329917860768894 6:-0.2822852019877604 7:0.6834785385197197 8:-0.6458607648553825 9:-0.06171476054995373 10:0.11421513352405444 +-15.310980589416289 1:-0.35290763483001486 2:-0.7263565311032778 3:-0.8688987069582226 4:-0.991098319894185 5:0.7029028082332363 6:-0.20251284356518684 7:-0.10928416773360117 8:0.307764663956116 9:0.6423143148384418 10:-0.15527637175127107 +3.260298266762908 1:-0.7817510582064782 2:0.45336200757318257 3:-0.15365670773321338 4:0.5063951567230205 5:-0.7102867196895872 6:-0.48050036620725955 7:0.9838016675169072 8:0.07854601230194436 9:-0.18953694857147863 10:0.19370072527454107 +3.846123583197846 1:0.6665586449040093 2:-0.2894063530813835 3:0.29965348483445386 4:0.23590344101670313 5:-0.7456743720187828 6:-0.4680876353446175 7:0.8106301610699425 8:0.691280702194663 9:-0.6060141408622055 10:0.34018639920235194 +-10.945919657782932 1:0.7669971723591666 2:0.38702771863552776 3:-0.6664311930513411 4:-0.2817072090916286 5:-0.16955916900934387 6:-0.9425831315444453 7:0.5685476711649924 8:-0.20782258743798265 9:0.015213591474494637 10:0.8183723865760859 +9.820049725467145 1:0.9582163993327679 2:0.7503905881505508 3:0.6255110430336392 4:0.6522701954798096 5:0.09248037700932144 6:-0.2833482854986902 7:-0.9841968940607242 8:-0.9343780716625845 9:-0.605526104070818 10:0.6000165028195326 +11.398715935456183 1:0.6605086903456443 2:0.14675454515266395 3:-0.7880053589830274 4:-0.8570785944515658 5:-0.4317693974151271 6:-0.12244918233307645 7:0.9808241653220866 8:0.5455853515046201 9:0.6870972425676756 10:0.7427686762232875 +-7.846310147695936 1:0.4355817642106965 2:0.7659504362110916 3:-0.3784171977305315 4:-0.5675896574776877 5:-0.20116390539973938 6:0.8775467546326667 7:-0.2824903364469842 8:0.7470660314619617 9:0.8967783051712528 10:0.7133700339519966 +-1.3847391232663768 1:0.3707613476850027 2:0.6931092598460797 3:-0.7701621508103305 4:-0.5679366502518555 5:-0.7234356749703683 6:-0.8059255104944509 7:-0.8307993875388229 8:0.6133975694770035 9:-0.7399749904168824 10:-0.1534990394513953 +16.93981662267873 1:0.6552665678625891 2:0.023248457840923775 3:-0.6850641408327465 4:0.7129790774369389 5:0.04166304042825364 6:-0.7160289667702797 7:-0.4733073680976494 8:0.2720897719417634 9:0.05850741911975099 10:0.34427554125371174 +2.8497179990245116 1:0.6664937514484015 2:0.3343796939204209 3:0.2611910348746209 4:-0.13658810351647 5:-0.5821801257591224 6:0.9854683468621908 7:-0.21396555404689188 8:-0.5923272173716836 9:-0.5674796199927252 10:-0.5681633547764235 +4.981807952389501 1:0.7517426071091595 2:0.7029291090701855 3:0.7126619831046563 4:-0.9982007415355478 5:-0.7743343367502893 6:-0.9048858749551119 7:-0.8243783842398396 8:0.4936163270697016 9:-0.6835495591484724 10:0.8412758607464845 +8.508637575729951 1:0.6837354268578517 2:-0.435346907350056 3:0.6597448795477736 4:0.8870204157376871 5:-0.6938576101541436 6:0.9199495715292882 7:0.33119640706964293 8:-0.6181273221979411 9:0.12929034268333317 10:0.6855150395247027 +14.369378079132883 1:-0.9489372180887643 2:-0.6577177233364067 3:0.543899463531252 4:0.5411152154119976 5:0.43733244485250733 6:0.5927084968109424 7:0.6100068837998656 8:0.9392735722529637 9:-0.9806701698603073 10:0.3984176141500082 +-6.456944198081549 1:0.8380442392342373 2:0.05166133486184443 3:-0.25864153418691704 4:-0.9506672344106888 5:0.5227275493542325 6:-0.03899736644563956 7:0.7660133053649136 8:-0.9375236703284806 9:-0.37213210747743175 10:0.0560768367274771 +-10.041353112580456 1:0.5293717914660876 2:-0.35874932480194044 3:0.14403824250820763 4:-0.4106496629336782 5:-0.794648717231762 6:-0.4369956159772408 7:0.8273613210141495 8:0.9212255384858874 9:0.00409867676727993 10:-0.23796544184855795 +-6.606325361718908 1:0.2765102732490652 2:0.10184669160432525 3:-0.9406443798496789 4:-0.46661976112717896 5:-0.5836573778289609 6:0.1308554421925976 7:0.05232199712543473 8:-0.4965370542771641 9:-0.3695836654343949 10:0.4874427445939513 +-15.359544879832677 1:-0.8253830145927283 2:0.29683545543963885 3:-0.9790356574071053 4:0.33749594518426473 5:-0.449483349548623 6:0.1740013774913005 7:0.5737323257916764 8:0.20159372721320645 9:-0.1812760896634873 10:-0.17652712339895738 +2.1801769966756845 1:0.3664130766917151 2:-0.1929450967547921 3:-0.7834945448457515 4:-0.03806442314852432 5:-0.6167622313628849 6:0.34919852301325394 7:-0.785891329691004 8:-0.5704062599527768 9:0.9846140894872721 10:-0.548571249100203 +-2.7006646885251415 1:-0.48505178676353067 2:0.06347121974094883 3:-0.3704723119141229 4:0.7407080276548548 5:0.06713252857406937 6:-0.2103524488773294 7:-0.9402467715192988 8:-0.8555624501612784 9:0.6244760190429901 10:-0.9038885681517279 +0.2105613019270259 1:-0.17125223509187282 2:-0.23328463772140529 3:-0.6497773470047024 4:0.33111604806115524 5:0.7944287248398398 6:0.5163977380074081 7:-0.025715995643062595 8:0.11762566041047462 9:0.9938658554834845 10:0.5363394203614278 +-0.6433952980357234 1:-0.905126800719938 2:0.5826442985002787 3:-0.8207546276288018 4:-0.0773547002692121 5:-0.6420058913410687 6:-0.9290787206193325 7:0.21829202840889095 8:-0.7752845890678082 9:0.4533233304372326 10:0.5457315861825041 +5.622874731146287 1:0.5486636398086722 2:-0.21867854114956642 3:0.13260110994566032 4:-0.024868470628895967 5:0.9246597814546305 6:0.07490395250443149 7:-0.21327567620097132 8:-0.33970581204395867 9:-0.19408398882121713 10:0.9757334811378136 +-18.27521356600463 1:-0.489685764918109 2:0.6832314342743568 3:0.9115808714640257 4:-4.680515344936964E-4 5:0.03760860984717218 6:0.4344127744883004 7:-0.30019645809377127 8:-0.48339658188341783 9:-0.5488933834939806 10:-0.4735052851773165 +5.518650144654079 1:-0.16881374315243192 2:0.22747702179774354 3:-0.8555270909193926 4:-0.6914231522703247 5:0.03618437407657238 6:-0.8404831131806643 7:0.16378525699004887 8:-0.333895928854854 9:0.23026574917978326 10:0.9409087845740918 +2.5599738684677646 1:-0.24371170373626905 2:-0.1752613047793694 3:-0.7930324885557696 4:0.17288443448968627 5:0.7233942014077801 6:0.47222694561171963 7:0.7878187692414558 8:-0.6520011755878357 9:-0.9952507460157223 10:-0.32951026378415094 +-8.508663400554862 1:0.9194236423060742 2:0.9517284917259223 3:-0.18723709334016392 4:-0.24913001260985546 5:0.8818286401027424 6:0.13661210218384512 7:-0.40792517201812983 8:-0.33132907984544957 9:-0.49137388288628703 10:-0.3273925353006979 +-10.233439586953153 1:0.0960128812383454 2:-0.8611756848964027 3:0.11807312551418647 4:-0.24570750746947145 5:-0.047826307143366886 6:-0.717269426008625 7:-0.2841658181308486 8:-0.31500935950449516 9:0.23183474949267713 10:-0.512986169560546 +-6.3459370724834265 1:0.9537835418930307 2:0.4598546399405288 3:-0.257013655072986 4:-0.29185820894937575 5:-0.6843688281544562 6:0.8343952028925479 7:-0.9656517094615942 8:-0.447440560943553 9:-0.9510349521362857 10:0.5918946980259567 +1.114406550703455 1:-0.5721838436595965 2:0.1201917297381252 3:-0.5253701290141362 4:-0.5874011312890843 5:0.7893580092022578 6:-0.18012813622584134 7:0.4781905737504004 8:-4.6732390143988667E-4 9:-0.7965374182885014 10:-0.8515444146742359 +8.688243146888663 1:0.2245581140502393 2:-0.0697600364101425 3:-0.7661833153629154 4:-0.2289151515902894 5:-0.5643191391300282 6:0.08069861795512168 7:-0.9670317635091523 8:0.14826752863715287 9:0.9325364047311011 10:0.4071178661803092 +14.896035572185347 1:0.20630949870309911 2:-0.5738578325975092 3:0.5664829389128903 4:0.3732752326637825 5:0.04079303403038881 6:-0.6604984910400766 7:0.15136076091734352 8:-0.6244939282579305 9:-0.5236288549540624 10:0.47284992666739023 +4.396558596072123 1:0.5565602414172521 2:0.1444095747909111 3:0.028227502879770272 4:0.38297378287943773 5:-0.26739745457451725 6:-0.708209627997985 7:0.7604483272526881 8:0.8072075261139096 9:0.11460574885028274 10:-0.07669406807610635 +1.7457141275341528 1:0.3668576517164046 2:-0.5352200081463954 3:0.5853385976871426 4:-0.4482551060006992 5:-0.5676795208498786 6:0.8043295590331514 7:-0.02160829797068753 8:0.42281303847010454 9:0.027894531623162466 10:-0.541120112980032 +-15.334767479922341 1:-0.036676500783341615 2:0.804758241454594 3:-0.0642091078911513 4:0.1402705435750966 5:-0.9215322030628859 6:0.7951173116514345 7:-0.994819896842561 8:0.2382406912119326 9:0.6634166177958731 10:0.7623222578718651 +5.017247792012723 1:-0.5925393497160352 2:0.48506599831456443 3:-0.5079795649118319 4:0.6668553329827696 5:-0.1103174867779837 6:0.7048535526809607 7:-0.9819230894106692 8:0.19609620625274982 9:0.5173985272313828 10:-0.11269849619148875 +6.201510810634532 1:-0.6802942101330738 2:0.898957584078176 3:0.853293387559251 4:0.6089336185656065 5:-0.9352626288322801 6:0.3208583332890447 7:-0.964481544931127 8:-0.8294773786068643 9:-0.8817311989413614 10:0.5165364663580934 +19.174935630244647 1:-0.20026105252200788 2:0.7276178994821614 3:0.7748716685190951 4:-0.7423420145576229 5:0.13147770471985032 6:-0.8382015712894606 7:0.021760992104270294 8:-0.24586987823702944 9:-0.05958177281299326 10:0.47347236224860834 +-14.822152909751189 1:0.7255660700197897 2:-0.22751988933383926 3:-0.08409197084114317 4:0.072750455428638 5:0.1841692073989072 6:-0.33838406658716513 7:-0.44701963574290526 8:0.5031210959133143 9:0.09640858549693743 10:0.9857351194637847 +-6.310082095945472 1:-0.7692076133438608 2:0.8533601511731044 3:0.676268298275629 4:-0.783895030001512 5:-0.8195462819549715 6:0.3963101354895673 7:-0.6254922461977397 8:-0.7521135990258581 9:-0.8032003997516024 10:0.8388672800826487 +8.853802632714807 1:0.46950948246522195 2:-0.6148693581037883 3:0.028739220735170656 4:-0.024281643566285815 5:-0.3495458137792231 6:-0.12347196435522867 7:0.5253894065203333 8:0.5100713458262918 9:0.63975795701667 10:0.08644353314625053 +-10.293714040655924 1:-0.17971950768550893 2:-0.6621720204354751 3:0.888036885802737 4:-0.04977483590350751 5:-0.8964991391283221 6:0.6873490822438724 7:0.42369087852118836 8:0.48972554317650663 9:0.8617233178519317 10:-0.8348331836605276 +0.23985611568891863 1:0.050526696983213215 2:0.8544297176525815 3:0.8586358519997579 4:-0.021299752441110487 5:0.2606696929560939 6:-0.39446486150105997 7:-0.4166234435381613 8:-0.6097643266459343 9:0.46633996256010146 10:-0.22521646199731027 +21.57719950299147 1:-0.5878491135126271 2:0.802134056970349 3:-0.5471017580843434 4:0.6067966843473331 5:-0.691712219323007 6:0.7814323754276735 7:0.31689445927290016 8:-0.1668780061940922 9:0.5285692389527452 10:0.8027091025203246 +-0.7836538830323514 1:0.5766794801558166 2:0.8281463568384935 3:0.5087453132796032 4:0.5212853344036532 5:0.6294700781054074 6:-0.9385097739886943 7:-0.13127371407538302 8:0.9845390503404141 9:-0.7224166213906742 10:-0.11155327354295896 +6.710413649604831 1:-0.6919803228062729 2:-0.6526904017578161 3:-0.34211291948607014 4:0.9094842803341618 5:-0.9454398661995895 6:0.3780766512494227 7:0.5823385348738088 8:0.8817830051841733 9:-0.514843382774189 10:0.32579701113259296 +5.384747201245483 1:-0.9661857672086316 2:-0.519769534339731 3:-0.4466396856529564 4:-0.4370113024678448 5:-0.6397400687811474 6:0.08225309277403725 7:-0.25936524603970756 8:-0.1711463274766858 9:-0.42848099098115755 10:-0.8096854737357237 +7.688509532916731 1:0.3892872094452817 2:-0.13306620868059982 3:-0.932974891205117 4:-0.8921357494146682 5:0.4806996560679244 6:-0.21500288444218696 7:-0.8911268070046585 8:-0.9510264953215406 9:0.1899740993687098 10:-0.43944320580463536 +2.2546997585565296 1:-0.5963883101717473 2:-0.01115153603404151 3:0.8781871380140298 4:0.7736250964135891 5:-0.7325745711528668 6:0.2518631794989008 7:0.5760249284318746 8:0.8690107952725199 9:0.02320853138646095 10:0.08570951531344417 +5.597710012706039 1:-0.5323512235815979 2:0.03366944321271936 3:0.5143537675853551 4:0.28471250955283445 5:0.4012202634439719 6:0.12032039285431151 7:-0.08108716844967812 8:0.30231384371011294 9:0.03259115565303028 10:0.9567467516929173 +-12.977848725392104 1:-0.5908891529017144 2:-0.7678208242918028 3:0.8512434510178621 4:-0.14910196410347298 5:0.6250260229199651 6:0.5393378705290228 7:-0.9573580597625002 8:-0.864881502860934 9:0.4175735160503429 10:0.4872169215922426 +10.35887243981476 1:-0.09126023790482862 2:0.18852634121926526 3:-0.13523918100503107 4:0.8333842692409983 5:-0.6015442103644761 6:0.5347736461652235 7:-0.823489760471118 8:0.5562688292037381 9:-0.807478561291906 10:-0.666881464988351 +0.4250502150408626 1:0.7771717566171905 2:-0.8729202752916785 3:-0.25782888805127024 4:-0.13605474993771205 5:0.5911781118120025 6:-0.8444023967853633 7:0.6787302541469229 8:-0.5444299313083194 9:0.356121883138657 10:-0.8845333845080687 +-0.8743487925900991 1:-0.9087681208947878 2:-0.292625136739453 3:-0.35113758823291774 4:-0.705933223571676 5:-0.6882289471031144 6:0.8350131255297044 7:-0.7659016065609232 8:0.11400114955653207 9:-0.9466143658505732 10:-0.5033643125229932 +-5.615143641864686 1:-0.6688289820084299 2:-0.4623159855015393 3:0.012827807007503855 4:-0.44521264878006117 5:-0.5563111031201406 6:-0.6065295981983794 7:0.3806712426786838 8:-0.11317152118817408 9:0.507896127467435 10:-0.8487801189674464 +-0.1829397047693725 1:0.09377558075225512 2:0.5774384503027374 3:-0.7104684187448009 4:-0.07285914169135976 5:-0.8797920488335114 6:0.6099615504974201 7:-0.8047440624324915 8:-0.6877856114263066 9:0.5843004021777447 10:0.5190581455348131 +18.479680552020344 1:0.9635517137863321 2:0.9954507816218203 3:0.11959899129360774 4:0.3753283274192787 5:-0.9386713095183621 6:0.0926833703812433 7:0.48003949462701323 8:0.9432769781973132 9:-0.9637036991931129 10:-0.4064407447273508 +1.3850645873427236 1:0.14476184437006356 2:-0.11280617018445871 3:-0.4385084538142101 4:-0.5961619435136434 5:0.419554626795412 6:-0.5047767472761191 7:0.457180284958592 8:-0.9129360314541999 9:-0.6320022059786656 10:-0.44989608519659363 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_movielens_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_movielens_data.txt new file mode 100644 index 0000000..f0eee19 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_movielens_data.txt @@ -0,0 +1,1501 @@ +0::2::3 +0::3::1 +0::5::2 +0::9::4 +0::11::1 +0::12::2 +0::15::1 +0::17::1 +0::19::1 +0::21::1 +0::23::1 +0::26::3 +0::27::1 +0::28::1 +0::29::1 +0::30::1 +0::31::1 +0::34::1 +0::37::1 +0::41::2 +0::44::1 +0::45::2 +0::46::1 +0::47::1 +0::48::1 +0::50::1 +0::51::1 +0::54::1 +0::55::1 +0::59::2 +0::61::2 +0::64::1 +0::67::1 +0::68::1 +0::69::1 +0::71::1 +0::72::1 +0::77::2 +0::79::1 +0::83::1 +0::87::1 +0::89::2 +0::91::3 +0::92::4 +0::94::1 +0::95::2 +0::96::1 +0::98::1 +0::99::1 +1::2::2 +1::3::1 +1::4::2 +1::6::1 +1::9::3 +1::12::1 +1::13::1 +1::14::1 +1::16::1 +1::19::1 +1::21::3 +1::27::1 +1::28::3 +1::33::1 +1::36::2 +1::37::1 +1::40::1 +1::41::2 +1::43::1 +1::44::1 +1::47::1 +1::50::1 +1::54::1 +1::56::2 +1::57::1 +1::58::1 +1::60::1 +1::62::4 +1::63::1 +1::67::1 +1::68::4 +1::70::2 +1::72::1 +1::73::1 +1::74::2 +1::76::1 +1::77::3 +1::78::1 +1::81::1 +1::82::1 +1::85::3 +1::86::2 +1::88::2 +1::91::1 +1::92::2 +1::93::1 +1::94::2 +1::96::1 +1::97::1 +2::4::3 +2::6::1 +2::8::5 +2::9::1 +2::10::1 +2::12::3 +2::13::1 +2::15::2 +2::18::2 +2::19::4 +2::22::1 +2::26::1 +2::28::1 +2::34::4 +2::35::1 +2::37::5 +2::38::1 +2::39::5 +2::40::4 +2::47::1 +2::50::1 +2::52::2 +2::54::1 +2::55::1 +2::57::2 +2::58::2 +2::59::1 +2::61::1 +2::62::1 +2::64::1 +2::65::1 +2::66::3 +2::68::1 +2::71::3 +2::76::1 +2::77::1 +2::78::1 +2::80::1 +2::83::5 +2::85::1 +2::87::2 +2::88::1 +2::89::4 +2::90::1 +2::92::4 +2::93::5 +3::0::1 +3::1::1 +3::2::1 +3::7::3 +3::8::3 +3::9::1 +3::14::1 +3::15::1 +3::16::1 +3::18::4 +3::19::1 +3::24::3 +3::26::1 +3::29::3 +3::33::1 +3::34::3 +3::35::1 +3::36::3 +3::37::1 +3::38::2 +3::43::1 +3::44::1 +3::46::1 +3::47::1 +3::51::5 +3::52::3 +3::56::1 +3::58::1 +3::60::3 +3::62::1 +3::65::2 +3::66::1 +3::67::1 +3::68::2 +3::70::1 +3::72::2 +3::76::3 +3::79::3 +3::80::4 +3::81::1 +3::83::1 +3::84::1 +3::86::1 +3::87::2 +3::88::4 +3::89::1 +3::91::1 +3::94::3 +4::1::1 +4::6::1 +4::8::1 +4::9::1 +4::10::1 +4::11::1 +4::12::1 +4::13::1 +4::14::2 +4::15::1 +4::17::1 +4::20::1 +4::22::1 +4::23::1 +4::24::1 +4::29::4 +4::30::1 +4::31::1 +4::34::1 +4::35::1 +4::36::1 +4::39::2 +4::40::3 +4::41::4 +4::43::2 +4::44::1 +4::45::1 +4::46::1 +4::47::1 +4::49::2 +4::50::1 +4::51::1 +4::52::4 +4::54::1 +4::55::1 +4::60::3 +4::61::1 +4::62::4 +4::63::3 +4::65::1 +4::67::2 +4::69::1 +4::70::4 +4::71::1 +4::73::1 +4::78::1 +4::84::1 +4::85::1 +4::87::3 +4::88::3 +4::89::2 +4::96::1 +4::97::1 +4::98::1 +4::99::1 +5::0::1 +5::1::1 +5::4::1 +5::5::1 +5::8::1 +5::9::3 +5::10::2 +5::13::3 +5::15::1 +5::19::1 +5::20::3 +5::21::2 +5::23::3 +5::27::1 +5::28::1 +5::29::1 +5::31::1 +5::36::3 +5::38::2 +5::39::1 +5::42::1 +5::48::3 +5::49::4 +5::50::3 +5::51::1 +5::52::1 +5::54::1 +5::55::5 +5::56::3 +5::58::1 +5::60::1 +5::61::1 +5::64::3 +5::65::2 +5::68::4 +5::70::1 +5::71::1 +5::72::1 +5::74::1 +5::79::1 +5::81::2 +5::84::1 +5::85::1 +5::86::1 +5::88::1 +5::90::4 +5::91::2 +5::95::2 +5::99::1 +6::0::1 +6::1::1 +6::2::3 +6::5::1 +6::6::1 +6::9::1 +6::10::1 +6::15::2 +6::16::2 +6::17::1 +6::18::1 +6::20::1 +6::21::1 +6::22::1 +6::24::1 +6::25::5 +6::26::1 +6::28::1 +6::30::1 +6::33::1 +6::38::1 +6::39::1 +6::43::4 +6::44::1 +6::45::1 +6::48::1 +6::49::1 +6::50::1 +6::53::1 +6::54::1 +6::55::1 +6::56::1 +6::58::4 +6::59::1 +6::60::1 +6::61::3 +6::63::3 +6::66::1 +6::67::3 +6::68::1 +6::69::1 +6::71::2 +6::73::1 +6::75::1 +6::77::1 +6::79::1 +6::81::1 +6::84::1 +6::85::3 +6::86::1 +6::87::1 +6::88::1 +6::89::1 +6::91::2 +6::94::1 +6::95::2 +6::96::1 +7::1::1 +7::2::2 +7::3::1 +7::4::1 +7::7::1 +7::10::1 +7::11::2 +7::14::2 +7::15::1 +7::16::1 +7::18::1 +7::21::1 +7::22::1 +7::23::1 +7::25::5 +7::26::1 +7::29::4 +7::30::1 +7::31::3 +7::32::1 +7::33::1 +7::35::1 +7::37::2 +7::39::3 +7::40::2 +7::42::2 +7::44::1 +7::45::2 +7::47::4 +7::48::1 +7::49::1 +7::53::1 +7::54::1 +7::55::1 +7::56::1 +7::59::1 +7::61::2 +7::62::3 +7::63::2 +7::66::1 +7::67::3 +7::74::1 +7::75::1 +7::76::3 +7::77::1 +7::81::1 +7::82::1 +7::84::2 +7::85::4 +7::86::1 +7::92::2 +7::96::1 +7::97::1 +7::98::1 +8::0::1 +8::2::4 +8::3::2 +8::4::2 +8::5::1 +8::7::1 +8::9::1 +8::11::1 +8::15::1 +8::18::1 +8::19::1 +8::21::1 +8::29::5 +8::31::3 +8::33::1 +8::35::1 +8::36::1 +8::40::2 +8::44::1 +8::45::1 +8::50::1 +8::51::1 +8::52::5 +8::53::5 +8::54::1 +8::55::1 +8::56::1 +8::58::4 +8::60::3 +8::62::4 +8::64::1 +8::67::3 +8::69::1 +8::71::1 +8::72::3 +8::77::3 +8::78::1 +8::79::1 +8::83::1 +8::85::5 +8::86::1 +8::88::1 +8::90::1 +8::92::2 +8::95::4 +8::96::3 +8::97::1 +8::98::1 +8::99::1 +9::2::3 +9::3::1 +9::4::1 +9::5::1 +9::6::1 +9::7::5 +9::9::1 +9::12::1 +9::14::3 +9::15::1 +9::19::1 +9::21::1 +9::22::1 +9::24::1 +9::25::1 +9::26::1 +9::30::3 +9::32::4 +9::35::2 +9::36::2 +9::37::2 +9::38::1 +9::39::1 +9::43::3 +9::49::5 +9::50::3 +9::53::1 +9::54::1 +9::58::1 +9::59::1 +9::60::1 +9::61::1 +9::63::3 +9::64::3 +9::68::1 +9::69::1 +9::70::3 +9::71::1 +9::73::2 +9::75::1 +9::77::2 +9::81::2 +9::82::1 +9::83::1 +9::84::1 +9::86::1 +9::87::4 +9::88::1 +9::90::3 +9::94::2 +9::95::3 +9::97::2 +9::98::1 +10::0::3 +10::2::4 +10::4::3 +10::7::1 +10::8::1 +10::10::1 +10::13::2 +10::14::1 +10::16::2 +10::17::1 +10::18::1 +10::21::1 +10::22::1 +10::24::1 +10::25::3 +10::28::1 +10::35::1 +10::36::1 +10::37::1 +10::38::1 +10::39::1 +10::40::4 +10::41::2 +10::42::3 +10::43::1 +10::49::3 +10::50::1 +10::51::1 +10::52::1 +10::55::2 +10::56::1 +10::58::1 +10::63::1 +10::66::1 +10::67::2 +10::68::1 +10::75::1 +10::77::1 +10::79::1 +10::86::1 +10::89::3 +10::90::1 +10::97::1 +10::98::1 +11::0::1 +11::6::2 +11::9::1 +11::10::1 +11::11::1 +11::12::1 +11::13::4 +11::16::1 +11::18::5 +11::19::4 +11::20::1 +11::21::1 +11::22::1 +11::23::5 +11::25::1 +11::27::5 +11::30::5 +11::32::5 +11::35::3 +11::36::2 +11::37::2 +11::38::4 +11::39::1 +11::40::1 +11::41::1 +11::43::2 +11::45::1 +11::47::1 +11::48::5 +11::50::4 +11::51::3 +11::59::1 +11::61::1 +11::62::1 +11::64::1 +11::66::4 +11::67::1 +11::69::5 +11::70::1 +11::71::3 +11::72::3 +11::75::3 +11::76::1 +11::77::1 +11::78::1 +11::79::5 +11::80::3 +11::81::4 +11::82::1 +11::86::1 +11::88::1 +11::89::1 +11::90::4 +11::94::2 +11::97::3 +11::99::1 +12::2::1 +12::4::1 +12::6::1 +12::7::3 +12::8::1 +12::14::1 +12::15::2 +12::16::4 +12::17::5 +12::18::2 +12::21::1 +12::22::2 +12::23::3 +12::24::1 +12::25::1 +12::27::5 +12::30::2 +12::31::4 +12::35::5 +12::38::1 +12::41::1 +12::44::2 +12::45::1 +12::50::4 +12::51::1 +12::52::1 +12::53::1 +12::54::1 +12::56::2 +12::57::1 +12::60::1 +12::63::1 +12::64::5 +12::66::3 +12::67::1 +12::70::1 +12::72::1 +12::74::1 +12::75::1 +12::77::1 +12::78::1 +12::79::3 +12::82::2 +12::83::1 +12::84::1 +12::85::1 +12::86::1 +12::87::1 +12::88::1 +12::91::3 +12::92::1 +12::94::4 +12::95::2 +12::96::1 +12::98::2 +13::0::1 +13::3::1 +13::4::2 +13::5::1 +13::6::1 +13::12::1 +13::14::2 +13::15::1 +13::17::1 +13::18::3 +13::20::1 +13::21::1 +13::22::1 +13::26::1 +13::27::1 +13::29::3 +13::31::1 +13::33::1 +13::40::2 +13::43::2 +13::44::1 +13::45::1 +13::49::1 +13::51::1 +13::52::2 +13::53::3 +13::54::1 +13::62::1 +13::63::2 +13::64::1 +13::68::1 +13::71::1 +13::72::3 +13::73::1 +13::74::3 +13::77::2 +13::78::1 +13::79::2 +13::83::3 +13::85::1 +13::86::1 +13::87::2 +13::88::2 +13::90::1 +13::93::4 +13::94::1 +13::98::1 +13::99::1 +14::1::1 +14::3::3 +14::4::1 +14::5::1 +14::6::1 +14::7::1 +14::9::1 +14::10::1 +14::11::1 +14::12::1 +14::13::1 +14::14::3 +14::15::1 +14::16::1 +14::17::1 +14::20::1 +14::21::1 +14::24::1 +14::25::2 +14::27::1 +14::28::1 +14::29::5 +14::31::3 +14::34::1 +14::36::1 +14::37::2 +14::39::2 +14::40::1 +14::44::1 +14::45::1 +14::47::3 +14::48::1 +14::49::1 +14::51::1 +14::52::5 +14::53::3 +14::54::1 +14::55::1 +14::56::1 +14::62::4 +14::63::5 +14::67::3 +14::68::1 +14::69::3 +14::71::1 +14::72::4 +14::73::1 +14::76::5 +14::79::1 +14::82::1 +14::83::1 +14::88::1 +14::93::3 +14::94::1 +14::95::2 +14::96::4 +14::98::1 +15::0::1 +15::1::4 +15::2::1 +15::5::2 +15::6::1 +15::7::1 +15::13::1 +15::14::1 +15::15::1 +15::17::2 +15::19::2 +15::22::2 +15::23::2 +15::25::1 +15::26::3 +15::27::1 +15::28::2 +15::29::1 +15::32::1 +15::33::2 +15::34::1 +15::35::2 +15::36::1 +15::37::1 +15::39::1 +15::42::1 +15::46::5 +15::48::2 +15::50::2 +15::51::1 +15::52::1 +15::58::1 +15::62::1 +15::64::3 +15::65::2 +15::72::1 +15::73::1 +15::74::1 +15::79::1 +15::80::1 +15::81::1 +15::82::2 +15::85::1 +15::87::1 +15::91::2 +15::96::1 +15::97::1 +15::98::3 +16::2::1 +16::5::3 +16::6::2 +16::7::1 +16::9::1 +16::12::1 +16::14::1 +16::15::1 +16::19::1 +16::21::2 +16::29::4 +16::30::2 +16::32::1 +16::34::1 +16::36::1 +16::38::1 +16::46::1 +16::47::3 +16::48::1 +16::49::1 +16::50::1 +16::51::5 +16::54::5 +16::55::1 +16::56::2 +16::57::1 +16::60::1 +16::63::2 +16::65::1 +16::67::1 +16::72::1 +16::74::1 +16::80::1 +16::81::1 +16::82::1 +16::85::5 +16::86::1 +16::90::5 +16::91::1 +16::93::1 +16::94::3 +16::95::2 +16::96::3 +16::98::3 +16::99::1 +17::2::1 +17::3::1 +17::6::1 +17::10::4 +17::11::1 +17::13::2 +17::17::5 +17::19::1 +17::20::5 +17::22::4 +17::28::1 +17::29::1 +17::33::1 +17::34::1 +17::35::2 +17::37::1 +17::38::1 +17::45::1 +17::46::5 +17::47::1 +17::49::3 +17::51::1 +17::55::5 +17::56::3 +17::57::1 +17::58::1 +17::59::1 +17::60::1 +17::63::1 +17::66::1 +17::68::4 +17::69::1 +17::70::1 +17::72::1 +17::73::3 +17::78::1 +17::79::1 +17::82::2 +17::84::1 +17::90::5 +17::91::3 +17::92::1 +17::93::1 +17::94::4 +17::95::2 +17::97::1 +18::1::1 +18::4::3 +18::5::2 +18::6::1 +18::7::1 +18::10::1 +18::11::4 +18::12::2 +18::13::1 +18::15::1 +18::18::1 +18::20::1 +18::21::2 +18::22::1 +18::23::2 +18::25::1 +18::26::1 +18::27::1 +18::28::5 +18::29::1 +18::31::1 +18::32::1 +18::36::1 +18::38::5 +18::39::5 +18::40::1 +18::42::1 +18::43::1 +18::44::4 +18::46::1 +18::47::1 +18::48::1 +18::51::2 +18::55::1 +18::56::1 +18::57::1 +18::62::1 +18::63::1 +18::66::3 +18::67::1 +18::70::1 +18::75::1 +18::76::3 +18::77::1 +18::80::3 +18::81::3 +18::82::1 +18::83::5 +18::84::1 +18::97::1 +18::98::1 +18::99::2 +19::0::1 +19::1::1 +19::2::1 +19::4::1 +19::6::2 +19::11::1 +19::12::1 +19::14::1 +19::23::1 +19::26::1 +19::31::1 +19::32::4 +19::33::1 +19::34::1 +19::37::1 +19::38::1 +19::41::1 +19::43::1 +19::45::1 +19::48::1 +19::49::1 +19::50::2 +19::53::2 +19::54::3 +19::55::1 +19::56::2 +19::58::1 +19::61::1 +19::62::1 +19::63::1 +19::64::1 +19::65::1 +19::69::2 +19::72::1 +19::74::3 +19::76::1 +19::78::1 +19::79::1 +19::81::1 +19::82::1 +19::84::1 +19::86::1 +19::87::2 +19::90::4 +19::93::1 +19::94::4 +19::95::2 +19::96::1 +19::98::4 +20::0::1 +20::1::1 +20::2::2 +20::4::2 +20::6::1 +20::8::1 +20::12::1 +20::21::2 +20::22::5 +20::24::2 +20::25::1 +20::26::1 +20::29::2 +20::30::2 +20::32::2 +20::39::1 +20::40::1 +20::41::2 +20::45::2 +20::48::1 +20::50::1 +20::51::3 +20::53::3 +20::55::1 +20::57::2 +20::60::1 +20::61::1 +20::64::1 +20::66::1 +20::70::2 +20::72::1 +20::73::2 +20::75::4 +20::76::1 +20::77::4 +20::78::1 +20::79::1 +20::84::2 +20::85::2 +20::88::3 +20::89::1 +20::90::3 +20::91::1 +20::92::2 +20::93::1 +20::94::4 +20::97::1 +21::0::1 +21::2::4 +21::3::1 +21::7::2 +21::11::1 +21::12::1 +21::13::1 +21::14::3 +21::17::1 +21::19::1 +21::20::1 +21::21::1 +21::22::1 +21::23::1 +21::24::1 +21::27::1 +21::29::5 +21::30::2 +21::38::1 +21::40::2 +21::43::3 +21::44::1 +21::45::1 +21::46::1 +21::48::1 +21::51::1 +21::53::5 +21::54::1 +21::55::1 +21::56::1 +21::58::3 +21::59::3 +21::64::1 +21::66::1 +21::68::1 +21::71::1 +21::73::1 +21::74::4 +21::80::1 +21::81::1 +21::83::1 +21::84::1 +21::85::3 +21::87::4 +21::89::2 +21::92::2 +21::96::3 +21::99::1 +22::0::1 +22::3::2 +22::5::2 +22::6::2 +22::9::1 +22::10::1 +22::11::1 +22::13::1 +22::14::1 +22::16::1 +22::18::3 +22::19::1 +22::22::5 +22::25::1 +22::26::1 +22::29::3 +22::30::5 +22::32::4 +22::33::1 +22::35::1 +22::36::3 +22::37::1 +22::40::1 +22::41::3 +22::44::1 +22::45::2 +22::48::1 +22::51::5 +22::55::1 +22::56::2 +22::60::3 +22::61::1 +22::62::4 +22::63::1 +22::65::1 +22::66::1 +22::68::4 +22::69::4 +22::70::3 +22::71::1 +22::74::5 +22::75::5 +22::78::1 +22::80::3 +22::81::1 +22::82::1 +22::84::1 +22::86::1 +22::87::3 +22::88::5 +22::90::2 +22::92::3 +22::95::2 +22::96::2 +22::98::4 +22::99::1 +23::0::1 +23::2::1 +23::4::1 +23::6::2 +23::10::4 +23::12::1 +23::13::4 +23::14::1 +23::15::1 +23::18::4 +23::22::2 +23::23::4 +23::24::1 +23::25::1 +23::26::1 +23::27::5 +23::28::1 +23::29::1 +23::30::4 +23::32::5 +23::33::2 +23::36::3 +23::37::1 +23::38::1 +23::39::1 +23::43::1 +23::48::5 +23::49::5 +23::50::4 +23::53::1 +23::55::5 +23::57::1 +23::59::1 +23::60::1 +23::61::1 +23::64::4 +23::65::5 +23::66::2 +23::67::1 +23::68::3 +23::69::1 +23::72::1 +23::73::3 +23::77::1 +23::82::2 +23::83::1 +23::84::1 +23::85::1 +23::87::3 +23::88::1 +23::95::2 +23::97::1 +24::4::1 +24::6::3 +24::7::1 +24::10::2 +24::12::1 +24::15::1 +24::19::1 +24::24::1 +24::27::3 +24::30::5 +24::31::1 +24::32::3 +24::33::1 +24::37::1 +24::39::1 +24::40::1 +24::42::1 +24::43::3 +24::45::2 +24::46::1 +24::47::1 +24::48::1 +24::49::1 +24::50::1 +24::52::5 +24::57::1 +24::59::4 +24::63::4 +24::65::1 +24::66::1 +24::67::1 +24::68::3 +24::69::5 +24::71::1 +24::72::4 +24::77::4 +24::78::1 +24::80::1 +24::82::1 +24::84::1 +24::86::1 +24::87::1 +24::88::2 +24::89::1 +24::90::5 +24::91::1 +24::92::1 +24::94::2 +24::95::1 +24::96::5 +24::98::1 +24::99::1 +25::1::3 +25::2::1 +25::7::1 +25::9::1 +25::12::3 +25::16::3 +25::17::1 +25::18::1 +25::20::1 +25::22::1 +25::23::1 +25::26::2 +25::29::1 +25::30::1 +25::31::2 +25::33::4 +25::34::3 +25::35::2 +25::36::1 +25::37::1 +25::40::1 +25::41::1 +25::43::1 +25::47::4 +25::50::1 +25::51::1 +25::53::1 +25::56::1 +25::58::2 +25::64::2 +25::67::2 +25::68::1 +25::70::1 +25::71::4 +25::73::1 +25::74::1 +25::76::1 +25::79::1 +25::82::1 +25::84::2 +25::85::1 +25::91::3 +25::92::1 +25::94::1 +25::95::1 +25::97::2 +26::0::1 +26::1::1 +26::2::1 +26::3::1 +26::4::4 +26::5::2 +26::6::3 +26::7::5 +26::13::3 +26::14::1 +26::16::1 +26::18::3 +26::20::1 +26::21::3 +26::22::5 +26::23::5 +26::24::5 +26::27::1 +26::31::1 +26::35::1 +26::36::4 +26::40::1 +26::44::1 +26::45::2 +26::47::1 +26::48::1 +26::49::3 +26::50::2 +26::52::1 +26::54::4 +26::55::1 +26::57::3 +26::58::1 +26::61::1 +26::62::2 +26::66::1 +26::68::4 +26::71::1 +26::73::4 +26::76::1 +26::81::3 +26::85::1 +26::86::3 +26::88::5 +26::91::1 +26::94::5 +26::95::1 +26::96::1 +26::97::1 +27::0::1 +27::9::1 +27::10::1 +27::18::4 +27::19::3 +27::20::1 +27::22::2 +27::24::2 +27::25::1 +27::27::3 +27::28::1 +27::29::1 +27::31::1 +27::33::3 +27::40::1 +27::42::1 +27::43::1 +27::44::3 +27::45::1 +27::51::3 +27::52::1 +27::55::3 +27::57::1 +27::59::1 +27::60::1 +27::61::1 +27::64::1 +27::66::3 +27::68::1 +27::70::1 +27::71::2 +27::72::1 +27::75::3 +27::78::1 +27::80::3 +27::82::1 +27::83::3 +27::86::1 +27::87::2 +27::90::1 +27::91::1 +27::92::1 +27::93::1 +27::94::2 +27::95::1 +27::98::1 +28::0::3 +28::1::1 +28::2::4 +28::3::1 +28::6::1 +28::7::1 +28::12::5 +28::13::2 +28::14::1 +28::15::1 +28::17::1 +28::19::3 +28::20::1 +28::23::3 +28::24::3 +28::27::1 +28::29::1 +28::33::1 +28::34::1 +28::36::1 +28::38::2 +28::39::2 +28::44::1 +28::45::1 +28::49::4 +28::50::1 +28::52::1 +28::54::1 +28::56::1 +28::57::3 +28::58::1 +28::59::1 +28::60::1 +28::62::3 +28::63::1 +28::65::1 +28::75::1 +28::78::1 +28::81::5 +28::82::4 +28::83::1 +28::85::1 +28::88::2 +28::89::4 +28::90::1 +28::92::5 +28::94::1 +28::95::2 +28::98::1 +28::99::1 +29::3::1 +29::4::1 +29::5::1 +29::7::2 +29::9::1 +29::10::3 +29::11::1 +29::13::3 +29::14::1 +29::15::1 +29::17::3 +29::19::3 +29::22::3 +29::23::4 +29::25::1 +29::29::1 +29::31::1 +29::32::4 +29::33::2 +29::36::2 +29::38::3 +29::39::1 +29::42::1 +29::46::5 +29::49::3 +29::51::2 +29::59::1 +29::61::1 +29::62::1 +29::67::1 +29::68::3 +29::69::1 +29::70::1 +29::74::1 +29::75::1 +29::79::2 +29::80::1 +29::81::2 +29::83::1 +29::85::1 +29::86::1 +29::90::4 +29::93::1 +29::94::4 +29::97::1 +29::99::1 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_multiclass_classification_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_multiclass_classification_data.txt new file mode 100644 index 0000000..a0d7f90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_multiclass_classification_data.txt @@ -0,0 +1,150 @@ +1 1:-0.222222 2:0.5 3:-0.762712 4:-0.833333 +1 1:-0.555556 2:0.25 3:-0.864407 4:-0.916667 +1 1:-0.722222 2:-0.166667 3:-0.864407 4:-0.833333 +1 1:-0.722222 2:0.166667 3:-0.694915 4:-0.916667 +0 1:0.166667 2:-0.416667 3:0.457627 4:0.5 +1 1:-0.833333 3:-0.864407 4:-0.916667 +2 1:-1.32455e-07 2:-0.166667 3:0.220339 4:0.0833333 +2 1:-1.32455e-07 2:-0.333333 3:0.0169491 4:-4.03573e-08 +1 1:-0.5 2:0.75 3:-0.830508 4:-1 +0 1:0.611111 3:0.694915 4:0.416667 +0 1:0.222222 2:-0.166667 3:0.423729 4:0.583333 +1 1:-0.722222 2:-0.166667 3:-0.864407 4:-1 +1 1:-0.5 2:0.166667 3:-0.864407 4:-0.916667 +2 1:-0.222222 2:-0.333333 3:0.0508474 4:-4.03573e-08 +2 1:-0.0555556 2:-0.833333 3:0.0169491 4:-0.25 +2 1:-0.166667 2:-0.416667 3:-0.0169491 4:-0.0833333 +1 1:-0.944444 3:-0.898305 4:-0.916667 +2 1:-0.277778 2:-0.583333 3:-0.0169491 4:-0.166667 +0 1:0.111111 2:-0.333333 3:0.38983 4:0.166667 +2 1:-0.222222 2:-0.166667 3:0.0847457 4:-0.0833333 +0 1:0.166667 2:-0.333333 3:0.559322 4:0.666667 +1 1:-0.611111 2:0.0833333 3:-0.864407 4:-0.916667 +2 1:-0.333333 2:-0.583333 3:0.0169491 4:-4.03573e-08 +0 1:0.555555 2:-0.166667 3:0.661017 4:0.666667 +2 1:0.166667 3:0.186441 4:0.166667 +2 1:0.111111 2:-0.75 3:0.152542 4:-4.03573e-08 +2 1:0.166667 2:-0.25 3:0.118644 4:-4.03573e-08 +0 1:-0.0555556 2:-0.833333 3:0.355932 4:0.166667 +0 1:-0.277778 2:-0.333333 3:0.322034 4:0.583333 +2 1:-0.222222 2:-0.5 3:-0.152542 4:-0.25 +2 1:-0.111111 3:0.288136 4:0.416667 +2 1:-0.0555556 2:-0.25 3:0.186441 4:0.166667 +2 1:0.333333 2:-0.166667 3:0.355932 4:0.333333 +1 1:-0.611111 2:0.25 3:-0.898305 4:-0.833333 +0 1:0.166667 2:-0.333333 3:0.559322 4:0.75 +0 1:0.111111 2:-0.25 3:0.559322 4:0.416667 +0 1:0.833333 2:-0.166667 3:0.898305 4:0.666667 +2 1:-0.277778 2:-0.166667 3:0.186441 4:0.166667 +0 1:-0.666667 2:-0.583333 3:0.186441 4:0.333333 +1 1:-0.666667 2:-0.0833334 3:-0.830508 4:-1 +1 1:-0.166667 2:0.666667 3:-0.932203 4:-0.916667 +0 1:0.0555554 2:-0.333333 3:0.288136 4:0.416667 +1 1:-0.666667 2:-0.0833334 3:-0.830508 4:-1 +1 1:-0.833333 2:0.166667 3:-0.864407 4:-0.833333 +0 1:0.0555554 2:0.166667 3:0.491525 4:0.833333 +0 1:0.722222 2:-0.333333 3:0.728813 4:0.5 +2 1:-0.166667 2:-0.416667 3:0.0508474 4:-0.25 +2 1:0.5 3:0.254237 4:0.0833333 +0 1:0.111111 2:-0.583333 3:0.355932 4:0.5 +1 1:-0.944444 2:-0.166667 3:-0.898305 4:-0.916667 +2 1:0.277778 2:-0.25 3:0.220339 4:-4.03573e-08 +0 1:0.666667 2:-0.25 3:0.79661 4:0.416667 +0 1:0.111111 2:0.0833333 3:0.694915 4:1 +0 1:0.444444 3:0.59322 4:0.833333 +2 1:-0.0555556 2:0.166667 3:0.186441 4:0.25 +1 1:-0.833333 2:0.333333 3:-1 4:-0.916667 +1 1:-0.555556 2:0.416667 3:-0.830508 4:-0.75 +2 1:-0.333333 2:-0.5 3:0.152542 4:-0.0833333 +1 1:-1 2:-0.166667 3:-0.966102 4:-1 +1 1:-0.333333 2:0.25 3:-0.898305 4:-0.916667 +2 1:0.388889 2:-0.333333 3:0.288136 4:0.0833333 +2 1:0.277778 2:-0.166667 3:0.152542 4:0.0833333 +0 1:0.333333 2:0.0833333 3:0.59322 4:0.666667 +1 1:-0.777778 3:-0.79661 4:-0.916667 +1 1:-0.444444 2:0.416667 3:-0.830508 4:-0.916667 +0 1:0.222222 2:-0.166667 3:0.627119 4:0.75 +1 1:-0.555556 2:0.5 3:-0.79661 4:-0.916667 +1 1:-0.555556 2:0.5 3:-0.694915 4:-0.75 +2 1:-1.32455e-07 2:-0.25 3:0.254237 4:0.0833333 +1 1:-0.5 2:0.25 3:-0.830508 4:-0.916667 +0 1:0.166667 3:0.457627 4:0.833333 +2 1:0.444444 2:-0.0833334 3:0.322034 4:0.166667 +0 1:0.111111 2:0.166667 3:0.559322 4:0.916667 +1 1:-0.611111 2:0.25 3:-0.79661 4:-0.583333 +0 1:0.388889 3:0.661017 4:0.833333 +1 1:-0.722222 2:0.166667 3:-0.79661 4:-0.916667 +1 1:-0.722222 2:-0.0833334 3:-0.79661 4:-0.916667 +1 1:-0.555556 2:0.166667 3:-0.830508 4:-0.916667 +2 1:-0.666667 2:-0.666667 3:-0.220339 4:-0.25 +2 1:-0.611111 2:-0.75 3:-0.220339 4:-0.25 +2 1:0.0555554 2:-0.833333 3:0.186441 4:0.166667 +0 1:-0.166667 2:-0.416667 3:0.38983 4:0.5 +0 1:0.611111 2:0.333333 3:0.728813 4:1 +2 1:0.0555554 2:-0.25 3:0.118644 4:-4.03573e-08 +1 1:-0.666667 2:-0.166667 3:-0.864407 4:-0.916667 +1 1:-0.833333 2:-0.0833334 3:-0.830508 4:-0.916667 +0 1:0.611111 2:-0.166667 3:0.627119 4:0.25 +0 1:0.888889 2:0.5 3:0.932203 4:0.75 +2 1:0.222222 2:-0.333333 3:0.220339 4:0.166667 +1 1:-0.555556 2:0.25 3:-0.864407 4:-0.833333 +0 1:-1.32455e-07 2:-0.166667 3:0.322034 4:0.416667 +0 1:-1.32455e-07 2:-0.5 3:0.559322 4:0.0833333 +1 1:-0.611111 3:-0.932203 4:-0.916667 +1 1:-0.333333 2:0.833333 3:-0.864407 4:-0.916667 +0 1:-0.166667 2:-0.333333 3:0.38983 4:0.916667 +2 1:-0.333333 2:-0.666667 3:-0.0847458 4:-0.25 +2 1:-0.0555556 2:-0.416667 3:0.38983 4:0.25 +1 1:-0.388889 2:0.416667 3:-0.830508 4:-0.916667 +0 1:0.444444 2:-0.0833334 3:0.38983 4:0.833333 +1 1:-0.611111 2:0.333333 3:-0.864407 4:-0.916667 +0 1:0.111111 2:-0.416667 3:0.322034 4:0.416667 +0 1:0.166667 2:-0.0833334 3:0.525424 4:0.416667 +2 1:0.333333 2:-0.0833334 3:0.152542 4:0.0833333 +0 1:-0.0555556 2:-0.166667 3:0.288136 4:0.416667 +0 1:-0.166667 2:-0.416667 3:0.38983 4:0.5 +1 1:-0.611111 2:0.166667 3:-0.830508 4:-0.916667 +0 1:0.888889 2:-0.166667 3:0.728813 4:0.833333 +2 1:-0.277778 2:-0.25 3:-0.118644 4:-4.03573e-08 +2 1:-0.222222 2:-0.333333 3:0.186441 4:-4.03573e-08 +0 1:0.333333 2:-0.583333 3:0.627119 4:0.416667 +0 1:0.444444 2:-0.0833334 3:0.491525 4:0.666667 +2 1:-0.222222 2:-0.25 3:0.0847457 4:-4.03573e-08 +1 1:-0.611111 2:0.166667 3:-0.79661 4:-0.75 +2 1:-0.277778 2:-0.166667 3:0.0508474 4:-4.03573e-08 +0 1:1 2:0.5 3:0.830508 4:0.583333 +2 1:-0.333333 2:-0.666667 3:-0.0508475 4:-0.166667 +2 1:-0.277778 2:-0.416667 3:0.0847457 4:-4.03573e-08 +0 1:0.888889 2:-0.333333 3:0.932203 4:0.583333 +2 1:-0.111111 2:-0.166667 3:0.0847457 4:0.166667 +2 1:0.111111 2:-0.583333 3:0.322034 4:0.166667 +0 1:0.333333 2:0.0833333 3:0.59322 4:1 +0 1:0.222222 2:-0.166667 3:0.525424 4:0.416667 +1 1:-0.555556 2:0.5 3:-0.830508 4:-0.833333 +0 1:-0.111111 2:-0.166667 3:0.38983 4:0.416667 +0 1:0.888889 2:-0.5 3:1 4:0.833333 +1 1:-0.388889 2:0.583333 3:-0.898305 4:-0.75 +2 1:0.111111 2:0.0833333 3:0.254237 4:0.25 +0 1:0.333333 2:-0.166667 3:0.423729 4:0.833333 +1 1:-0.388889 2:0.166667 3:-0.762712 4:-0.916667 +0 1:0.333333 2:-0.0833334 3:0.559322 4:0.916667 +2 1:-0.333333 2:-0.75 3:0.0169491 4:-4.03573e-08 +1 1:-0.222222 2:1 3:-0.830508 4:-0.75 +1 1:-0.388889 2:0.583333 3:-0.762712 4:-0.75 +2 1:-0.611111 2:-1 3:-0.152542 4:-0.25 +2 1:-1.32455e-07 2:-0.333333 3:0.254237 4:-0.0833333 +2 1:-0.5 2:-0.416667 3:-0.0169491 4:0.0833333 +1 1:-0.888889 2:-0.75 3:-0.898305 4:-0.833333 +1 1:-0.666667 2:-0.0833334 3:-0.830508 4:-1 +2 1:-0.555556 2:-0.583333 3:-0.322034 4:-0.166667 +2 1:-0.166667 2:-0.5 3:0.0169491 4:-0.0833333 +1 1:-0.555556 2:0.0833333 3:-0.762712 4:-0.666667 +1 1:-0.777778 3:-0.898305 4:-0.916667 +0 1:0.388889 2:-0.166667 3:0.525424 4:0.666667 +0 1:0.222222 3:0.38983 4:0.583333 +2 1:0.333333 2:-0.0833334 3:0.254237 4:0.166667 +2 1:-0.388889 2:-0.166667 3:0.186441 4:0.166667 +0 1:-0.222222 2:-0.583333 3:0.355932 4:0.583333 +1 1:-0.611111 2:-0.166667 3:-0.79661 4:-0.916667 +1 1:-0.944444 2:-0.25 3:-0.864407 4:-0.916667 +1 1:-0.388889 2:0.166667 3:-0.830508 4:-0.75 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/sample_svm_data.txt b/spark-3.5.0-bin-hadoop3/data/mllib/sample_svm_data.txt new file mode 100644 index 0000000..7ab30bd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/sample_svm_data.txt @@ -0,0 +1,322 @@ +1 0 2.52078447201548 0 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 12.72816758217773 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 0 0 0 0 4.745052855503306 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 0 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 0 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 0 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 0 0 2.004684436494304 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 6.857275130999357 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 10.4087817597473 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 12.72816758217773 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 4.745052855503306 0 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 0 2.52078447201548 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +0 0 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 2.061393766919624 0 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 4.745052855503306 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 0 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 6.857275130999357 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 4.745052855503306 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 0 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 0 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +0 0 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 2.52078447201548 0 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +0 0 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 0 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 6.857275130999357 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +0 0 2.52078447201548 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 0 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 0 2.52078447201548 0 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 0 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 2.52078447201548 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 4.745052855503306 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 0 0 0 0 0 0 0 2.122974378789621 0 0 0 0 12.72816758217773 10.4087817597473 12.72816758217773 17.97228742438751 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 0 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 2.52078447201548 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 0 0 0 0 0 0 0 0 0 0 +1 0 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 0 0 0 6.857275130999357 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 4.745052855503306 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 6.857275130999357 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 0 4.745052855503306 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 0 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +0 0 2.52078447201548 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 4.745052855503306 2.004684436494304 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 0 2.52078447201548 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 0 0 0 2.000347299268466 2.122974378789621 0 0 6.857275130999357 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 2.619965104088255 0 0 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 0 0 2.061393766919624 0 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 2.061393766919624 0 0 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 6.857275130999357 0 0 0 0 0 +1 2.857738033247042 0 2.061393766919624 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 2.000347299268466 0 0 0 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 2.52078447201548 2.061393766919624 2.619965104088255 0 0 2.000347299268466 0 0 0 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 0 0 0 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 4.745052855503306 0 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 2.52078447201548 0 0 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 2.857738033247042 2.52078447201548 0 2.619965104088255 0 2.004684436494304 0 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 2.061393766919624 0 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 10.4087817597473 0 0 +0 2.857738033247042 0 2.061393766919624 2.619965104088255 0 2.004684436494304 0 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +0 2.857738033247042 0 0 2.619965104088255 0 2.004684436494304 2.000347299268466 2.122974378789621 0 0 0 0 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 2.000347299268466 0 2.228387042742021 2.228387042742023 0 0 0 0 0 0 +1 0 2.52078447201548 0 2.619965104088255 0 0 0 0 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 +1 2.857738033247042 0 0 2.619965104088255 0 0 0 2.122974378789621 2.228387042742021 2.228387042742023 0 2.055002875864414 0 0 0 0 diff --git a/spark-3.5.0-bin-hadoop3/data/mllib/streaming_kmeans_data_test.txt b/spark-3.5.0-bin-hadoop3/data/mllib/streaming_kmeans_data_test.txt new file mode 100644 index 0000000..649a0d6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/mllib/streaming_kmeans_data_test.txt @@ -0,0 +1,2 @@ +(1.0), [1.7, 0.4, 0.9] +(2.0), [2.2, 1.8, 0.0] diff --git a/spark-3.5.0-bin-hadoop3/data/streaming/AFINN-111.txt b/spark-3.5.0-bin-hadoop3/data/streaming/AFINN-111.txt new file mode 100644 index 0000000..0f6fb8e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/data/streaming/AFINN-111.txt @@ -0,0 +1,2477 @@ +abandon -2 +abandoned -2 +abandons -2 +abducted -2 +abduction -2 +abductions -2 +abhor -3 +abhorred -3 +abhorrent -3 +abhors -3 +abilities 2 +ability 2 +aboard 1 +absentee -1 +absentees -1 +absolve 2 +absolved 2 +absolves 2 +absolving 2 +absorbed 1 +abuse -3 +abused -3 +abuses -3 +abusive -3 +accept 1 +accepted 1 +accepting 1 +accepts 1 +accident -2 +accidental -2 +accidentally -2 +accidents -2 +accomplish 2 +accomplished 2 +accomplishes 2 +accusation -2 +accusations -2 +accuse -2 +accused -2 +accuses -2 +accusing -2 +ache -2 +achievable 1 +aching -2 +acquit 2 +acquits 2 +acquitted 2 +acquitting 2 +acrimonious -3 +active 1 +adequate 1 +admire 3 +admired 3 +admires 3 +admiring 3 +admit -1 +admits -1 +admitted -1 +admonish -2 +admonished -2 +adopt 1 +adopts 1 +adorable 3 +adore 3 +adored 3 +adores 3 +advanced 1 +advantage 2 +advantages 2 +adventure 2 +adventures 2 +adventurous 2 +affected -1 +affection 3 +affectionate 3 +afflicted -1 +affronted -1 +afraid -2 +aggravate -2 +aggravated -2 +aggravates -2 +aggravating -2 +aggression -2 +aggressions -2 +aggressive -2 +aghast -2 +agog 2 +agonise -3 +agonised -3 +agonises -3 +agonising -3 +agonize -3 +agonized -3 +agonizes -3 +agonizing -3 +agree 1 +agreeable 2 +agreed 1 +agreement 1 +agrees 1 +alarm -2 +alarmed -2 +alarmist -2 +alarmists -2 +alas -1 +alert -1 +alienation -2 +alive 1 +allergic -2 +allow 1 +alone -2 +amaze 2 +amazed 2 +amazes 2 +amazing 4 +ambitious 2 +ambivalent -1 +amuse 3 +amused 3 +amusement 3 +amusements 3 +anger -3 +angers -3 +angry -3 +anguish -3 +anguished -3 +animosity -2 +annoy -2 +annoyance -2 +annoyed -2 +annoying -2 +annoys -2 +antagonistic -2 +anti -1 +anticipation 1 +anxiety -2 +anxious -2 +apathetic -3 +apathy -3 +apeshit -3 +apocalyptic -2 +apologise -1 +apologised -1 +apologises -1 +apologising -1 +apologize -1 +apologized -1 +apologizes -1 +apologizing -1 +apology -1 +appalled -2 +appalling -2 +appease 2 +appeased 2 +appeases 2 +appeasing 2 +applaud 2 +applauded 2 +applauding 2 +applauds 2 +applause 2 +appreciate 2 +appreciated 2 +appreciates 2 +appreciating 2 +appreciation 2 +apprehensive -2 +approval 2 +approved 2 +approves 2 +ardent 1 +arrest -2 +arrested -3 +arrests -2 +arrogant -2 +ashame -2 +ashamed -2 +ass -4 +assassination -3 +assassinations -3 +asset 2 +assets 2 +assfucking -4 +asshole -4 +astonished 2 +astound 3 +astounded 3 +astounding 3 +astoundingly 3 +astounds 3 +attack -1 +attacked -1 +attacking -1 +attacks -1 +attract 1 +attracted 1 +attracting 2 +attraction 2 +attractions 2 +attracts 1 +audacious 3 +authority 1 +avert -1 +averted -1 +averts -1 +avid 2 +avoid -1 +avoided -1 +avoids -1 +await -1 +awaited -1 +awaits -1 +award 3 +awarded 3 +awards 3 +awesome 4 +awful -3 +awkward -2 +axe -1 +axed -1 +backed 1 +backing 2 +backs 1 +bad -3 +badass -3 +badly -3 +bailout -2 +bamboozle -2 +bamboozled -2 +bamboozles -2 +ban -2 +banish -1 +bankrupt -3 +bankster -3 +banned -2 +bargain 2 +barrier -2 +bastard -5 +bastards -5 +battle -1 +battles -1 +beaten -2 +beatific 3 +beating -1 +beauties 3 +beautiful 3 +beautifully 3 +beautify 3 +belittle -2 +belittled -2 +beloved 3 +benefit 2 +benefits 2 +benefitted 2 +benefitting 2 +bereave -2 +bereaved -2 +bereaves -2 +bereaving -2 +best 3 +betray -3 +betrayal -3 +betrayed -3 +betraying -3 +betrays -3 +better 2 +bias -1 +biased -2 +big 1 +bitch -5 +bitches -5 +bitter -2 +bitterly -2 +bizarre -2 +blah -2 +blame -2 +blamed -2 +blames -2 +blaming -2 +bless 2 +blesses 2 +blessing 3 +blind -1 +bliss 3 +blissful 3 +blithe 2 +block -1 +blockbuster 3 +blocked -1 +blocking -1 +blocks -1 +bloody -3 +blurry -2 +boastful -2 +bold 2 +boldly 2 +bomb -1 +boost 1 +boosted 1 +boosting 1 +boosts 1 +bore -2 +bored -2 +boring -3 +bother -2 +bothered -2 +bothers -2 +bothersome -2 +boycott -2 +boycotted -2 +boycotting -2 +boycotts -2 +brainwashing -3 +brave 2 +breakthrough 3 +breathtaking 5 +bribe -3 +bright 1 +brightest 2 +brightness 1 +brilliant 4 +brisk 2 +broke -1 +broken -1 +brooding -2 +bullied -2 +bullshit -4 +bully -2 +bullying -2 +bummer -2 +buoyant 2 +burden -2 +burdened -2 +burdening -2 +burdens -2 +calm 2 +calmed 2 +calming 2 +calms 2 +can't stand -3 +cancel -1 +cancelled -1 +cancelling -1 +cancels -1 +cancer -1 +capable 1 +captivated 3 +care 2 +carefree 1 +careful 2 +carefully 2 +careless -2 +cares 2 +cashing in -2 +casualty -2 +catastrophe -3 +catastrophic -4 +cautious -1 +celebrate 3 +celebrated 3 +celebrates 3 +celebrating 3 +censor -2 +censored -2 +censors -2 +certain 1 +chagrin -2 +chagrined -2 +challenge -1 +chance 2 +chances 2 +chaos -2 +chaotic -2 +charged -3 +charges -2 +charm 3 +charming 3 +charmless -3 +chastise -3 +chastised -3 +chastises -3 +chastising -3 +cheat -3 +cheated -3 +cheater -3 +cheaters -3 +cheats -3 +cheer 2 +cheered 2 +cheerful 2 +cheering 2 +cheerless -2 +cheers 2 +cheery 3 +cherish 2 +cherished 2 +cherishes 2 +cherishing 2 +chic 2 +childish -2 +chilling -1 +choke -2 +choked -2 +chokes -2 +choking -2 +clarifies 2 +clarity 2 +clash -2 +classy 3 +clean 2 +cleaner 2 +clear 1 +cleared 1 +clearly 1 +clears 1 +clever 2 +clouded -1 +clueless -2 +cock -5 +cocksucker -5 +cocksuckers -5 +cocky -2 +coerced -2 +collapse -2 +collapsed -2 +collapses -2 +collapsing -2 +collide -1 +collides -1 +colliding -1 +collision -2 +collisions -2 +colluding -3 +combat -1 +combats -1 +comedy 1 +comfort 2 +comfortable 2 +comforting 2 +comforts 2 +commend 2 +commended 2 +commit 1 +commitment 2 +commits 1 +committed 1 +committing 1 +compassionate 2 +compelled 1 +competent 2 +competitive 2 +complacent -2 +complain -2 +complained -2 +complains -2 +comprehensive 2 +conciliate 2 +conciliated 2 +conciliates 2 +conciliating 2 +condemn -2 +condemnation -2 +condemned -2 +condemns -2 +confidence 2 +confident 2 +conflict -2 +conflicting -2 +conflictive -2 +conflicts -2 +confuse -2 +confused -2 +confusing -2 +congrats 2 +congratulate 2 +congratulation 2 +congratulations 2 +consent 2 +consents 2 +consolable 2 +conspiracy -3 +constrained -2 +contagion -2 +contagions -2 +contagious -1 +contempt -2 +contemptuous -2 +contemptuously -2 +contend -1 +contender -1 +contending -1 +contentious -2 +contestable -2 +controversial -2 +controversially -2 +convince 1 +convinced 1 +convinces 1 +convivial 2 +cool 1 +cool stuff 3 +cornered -2 +corpse -1 +costly -2 +courage 2 +courageous 2 +courteous 2 +courtesy 2 +cover-up -3 +coward -2 +cowardly -2 +coziness 2 +cramp -1 +crap -3 +crash -2 +crazier -2 +craziest -2 +crazy -2 +creative 2 +crestfallen -2 +cried -2 +cries -2 +crime -3 +criminal -3 +criminals -3 +crisis -3 +critic -2 +criticism -2 +criticize -2 +criticized -2 +criticizes -2 +criticizing -2 +critics -2 +cruel -3 +cruelty -3 +crush -1 +crushed -2 +crushes -1 +crushing -1 +cry -1 +crying -2 +cunt -5 +curious 1 +curse -1 +cut -1 +cute 2 +cuts -1 +cutting -1 +cynic -2 +cynical -2 +cynicism -2 +damage -3 +damages -3 +damn -4 +damned -4 +damnit -4 +danger -2 +daredevil 2 +daring 2 +darkest -2 +darkness -1 +dauntless 2 +dead -3 +deadlock -2 +deafening -1 +dear 2 +dearly 3 +death -2 +debonair 2 +debt -2 +deceit -3 +deceitful -3 +deceive -3 +deceived -3 +deceives -3 +deceiving -3 +deception -3 +decisive 1 +dedicated 2 +defeated -2 +defect -3 +defects -3 +defender 2 +defenders 2 +defenseless -2 +defer -1 +deferring -1 +defiant -1 +deficit -2 +degrade -2 +degraded -2 +degrades -2 +dehumanize -2 +dehumanized -2 +dehumanizes -2 +dehumanizing -2 +deject -2 +dejected -2 +dejecting -2 +dejects -2 +delay -1 +delayed -1 +delight 3 +delighted 3 +delighting 3 +delights 3 +demand -1 +demanded -1 +demanding -1 +demands -1 +demonstration -1 +demoralized -2 +denied -2 +denier -2 +deniers -2 +denies -2 +denounce -2 +denounces -2 +deny -2 +denying -2 +depressed -2 +depressing -2 +derail -2 +derailed -2 +derails -2 +deride -2 +derided -2 +derides -2 +deriding -2 +derision -2 +desirable 2 +desire 1 +desired 2 +desirous 2 +despair -3 +despairing -3 +despairs -3 +desperate -3 +desperately -3 +despondent -3 +destroy -3 +destroyed -3 +destroying -3 +destroys -3 +destruction -3 +destructive -3 +detached -1 +detain -2 +detained -2 +detention -2 +determined 2 +devastate -2 +devastated -2 +devastating -2 +devoted 3 +diamond 1 +dick -4 +dickhead -4 +die -3 +died -3 +difficult -1 +diffident -2 +dilemma -1 +dipshit -3 +dire -3 +direful -3 +dirt -2 +dirtier -2 +dirtiest -2 +dirty -2 +disabling -1 +disadvantage -2 +disadvantaged -2 +disappear -1 +disappeared -1 +disappears -1 +disappoint -2 +disappointed -2 +disappointing -2 +disappointment -2 +disappointments -2 +disappoints -2 +disaster -2 +disasters -2 +disastrous -3 +disbelieve -2 +discard -1 +discarded -1 +discarding -1 +discards -1 +disconsolate -2 +disconsolation -2 +discontented -2 +discord -2 +discounted -1 +discouraged -2 +discredited -2 +disdain -2 +disgrace -2 +disgraced -2 +disguise -1 +disguised -1 +disguises -1 +disguising -1 +disgust -3 +disgusted -3 +disgusting -3 +disheartened -2 +dishonest -2 +disillusioned -2 +disinclined -2 +disjointed -2 +dislike -2 +dismal -2 +dismayed -2 +disorder -2 +disorganized -2 +disoriented -2 +disparage -2 +disparaged -2 +disparages -2 +disparaging -2 +displeased -2 +dispute -2 +disputed -2 +disputes -2 +disputing -2 +disqualified -2 +disquiet -2 +disregard -2 +disregarded -2 +disregarding -2 +disregards -2 +disrespect -2 +disrespected -2 +disruption -2 +disruptions -2 +disruptive -2 +dissatisfied -2 +distort -2 +distorted -2 +distorting -2 +distorts -2 +distract -2 +distracted -2 +distraction -2 +distracts -2 +distress -2 +distressed -2 +distresses -2 +distressing -2 +distrust -3 +distrustful -3 +disturb -2 +disturbed -2 +disturbing -2 +disturbs -2 +dithering -2 +dizzy -1 +dodging -2 +dodgy -2 +does not work -3 +dolorous -2 +dont like -2 +doom -2 +doomed -2 +doubt -1 +doubted -1 +doubtful -1 +doubting -1 +doubts -1 +douche -3 +douchebag -3 +downcast -2 +downhearted -2 +downside -2 +drag -1 +dragged -1 +drags -1 +drained -2 +dread -2 +dreaded -2 +dreadful -3 +dreading -2 +dream 1 +dreams 1 +dreary -2 +droopy -2 +drop -1 +drown -2 +drowned -2 +drowns -2 +drunk -2 +dubious -2 +dud -2 +dull -2 +dumb -3 +dumbass -3 +dump -1 +dumped -2 +dumps -1 +dupe -2 +duped -2 +dysfunction -2 +eager 2 +earnest 2 +ease 2 +easy 1 +ecstatic 4 +eerie -2 +eery -2 +effective 2 +effectively 2 +elated 3 +elation 3 +elegant 2 +elegantly 2 +embarrass -2 +embarrassed -2 +embarrasses -2 +embarrassing -2 +embarrassment -2 +embittered -2 +embrace 1 +emergency -2 +empathetic 2 +emptiness -1 +empty -1 +enchanted 2 +encourage 2 +encouraged 2 +encouragement 2 +encourages 2 +endorse 2 +endorsed 2 +endorsement 2 +endorses 2 +enemies -2 +enemy -2 +energetic 2 +engage 1 +engages 1 +engrossed 1 +enjoy 2 +enjoying 2 +enjoys 2 +enlighten 2 +enlightened 2 +enlightening 2 +enlightens 2 +ennui -2 +enrage -2 +enraged -2 +enrages -2 +enraging -2 +enrapture 3 +enslave -2 +enslaved -2 +enslaves -2 +ensure 1 +ensuring 1 +enterprising 1 +entertaining 2 +enthral 3 +enthusiastic 3 +entitled 1 +entrusted 2 +envies -1 +envious -2 +envy -1 +envying -1 +erroneous -2 +error -2 +errors -2 +escape -1 +escapes -1 +escaping -1 +esteemed 2 +ethical 2 +euphoria 3 +euphoric 4 +eviction -1 +evil -3 +exaggerate -2 +exaggerated -2 +exaggerates -2 +exaggerating -2 +exasperated 2 +excellence 3 +excellent 3 +excite 3 +excited 3 +excitement 3 +exciting 3 +exclude -1 +excluded -2 +exclusion -1 +exclusive 2 +excuse -1 +exempt -1 +exhausted -2 +exhilarated 3 +exhilarates 3 +exhilarating 3 +exonerate 2 +exonerated 2 +exonerates 2 +exonerating 2 +expand 1 +expands 1 +expel -2 +expelled -2 +expelling -2 +expels -2 +exploit -2 +exploited -2 +exploiting -2 +exploits -2 +exploration 1 +explorations 1 +expose -1 +exposed -1 +exposes -1 +exposing -1 +extend 1 +extends 1 +exuberant 4 +exultant 3 +exultantly 3 +fabulous 4 +fad -2 +fag -3 +faggot -3 +faggots -3 +fail -2 +failed -2 +failing -2 +fails -2 +failure -2 +failures -2 +fainthearted -2 +fair 2 +faith 1 +faithful 3 +fake -3 +fakes -3 +faking -3 +fallen -2 +falling -1 +falsified -3 +falsify -3 +fame 1 +fan 3 +fantastic 4 +farce -1 +fascinate 3 +fascinated 3 +fascinates 3 +fascinating 3 +fascist -2 +fascists -2 +fatalities -3 +fatality -3 +fatigue -2 +fatigued -2 +fatigues -2 +fatiguing -2 +favor 2 +favored 2 +favorite 2 +favorited 2 +favorites 2 +favors 2 +fear -2 +fearful -2 +fearing -2 +fearless 2 +fearsome -2 +fed up -3 +feeble -2 +feeling 1 +felonies -3 +felony -3 +fervent 2 +fervid 2 +festive 2 +fiasco -3 +fidgety -2 +fight -1 +fine 2 +fire -2 +fired -2 +firing -2 +fit 1 +fitness 1 +flagship 2 +flees -1 +flop -2 +flops -2 +flu -2 +flustered -2 +focused 2 +fond 2 +fondness 2 +fool -2 +foolish -2 +fools -2 +forced -1 +foreclosure -2 +foreclosures -2 +forget -1 +forgetful -2 +forgive 1 +forgiving 1 +forgotten -1 +fortunate 2 +frantic -1 +fraud -4 +frauds -4 +fraudster -4 +fraudsters -4 +fraudulence -4 +fraudulent -4 +free 1 +freedom 2 +frenzy -3 +fresh 1 +friendly 2 +fright -2 +frightened -2 +frightening -3 +frikin -2 +frisky 2 +frowning -1 +frustrate -2 +frustrated -2 +frustrates -2 +frustrating -2 +frustration -2 +ftw 3 +fuck -4 +fucked -4 +fucker -4 +fuckers -4 +fuckface -4 +fuckhead -4 +fucking -4 +fucktard -4 +fud -3 +fuked -4 +fuking -4 +fulfill 2 +fulfilled 2 +fulfills 2 +fuming -2 +fun 4 +funeral -1 +funerals -1 +funky 2 +funnier 4 +funny 4 +furious -3 +futile 2 +gag -2 +gagged -2 +gain 2 +gained 2 +gaining 2 +gains 2 +gallant 3 +gallantly 3 +gallantry 3 +generous 2 +genial 3 +ghost -1 +giddy -2 +gift 2 +glad 3 +glamorous 3 +glamourous 3 +glee 3 +gleeful 3 +gloom -1 +gloomy -2 +glorious 2 +glory 2 +glum -2 +god 1 +goddamn -3 +godsend 4 +good 3 +goodness 3 +grace 1 +gracious 3 +grand 3 +grant 1 +granted 1 +granting 1 +grants 1 +grateful 3 +gratification 2 +grave -2 +gray -1 +great 3 +greater 3 +greatest 3 +greed -3 +greedy -2 +green wash -3 +green washing -3 +greenwash -3 +greenwasher -3 +greenwashers -3 +greenwashing -3 +greet 1 +greeted 1 +greeting 1 +greetings 2 +greets 1 +grey -1 +grief -2 +grieved -2 +gross -2 +growing 1 +growth 2 +guarantee 1 +guilt -3 +guilty -3 +gullibility -2 +gullible -2 +gun -1 +ha 2 +hacked -1 +haha 3 +hahaha 3 +hahahah 3 +hail 2 +hailed 2 +hapless -2 +haplessness -2 +happiness 3 +happy 3 +hard -1 +hardier 2 +hardship -2 +hardy 2 +harm -2 +harmed -2 +harmful -2 +harming -2 +harms -2 +harried -2 +harsh -2 +harsher -2 +harshest -2 +hate -3 +hated -3 +haters -3 +hates -3 +hating -3 +haunt -1 +haunted -2 +haunting 1 +haunts -1 +havoc -2 +healthy 2 +heartbreaking -3 +heartbroken -3 +heartfelt 3 +heaven 2 +heavenly 4 +heavyhearted -2 +hell -4 +help 2 +helpful 2 +helping 2 +helpless -2 +helps 2 +hero 2 +heroes 2 +heroic 3 +hesitant -2 +hesitate -2 +hid -1 +hide -1 +hides -1 +hiding -1 +highlight 2 +hilarious 2 +hindrance -2 +hoax -2 +homesick -2 +honest 2 +honor 2 +honored 2 +honoring 2 +honour 2 +honoured 2 +honouring 2 +hooligan -2 +hooliganism -2 +hooligans -2 +hope 2 +hopeful 2 +hopefully 2 +hopeless -2 +hopelessness -2 +hopes 2 +hoping 2 +horrendous -3 +horrible -3 +horrific -3 +horrified -3 +hostile -2 +huckster -2 +hug 2 +huge 1 +hugs 2 +humerous 3 +humiliated -3 +humiliation -3 +humor 2 +humorous 2 +humour 2 +humourous 2 +hunger -2 +hurrah 5 +hurt -2 +hurting -2 +hurts -2 +hypocritical -2 +hysteria -3 +hysterical -3 +hysterics -3 +idiot -3 +idiotic -3 +ignorance -2 +ignorant -2 +ignore -1 +ignored -2 +ignores -1 +ill -2 +illegal -3 +illiteracy -2 +illness -2 +illnesses -2 +imbecile -3 +immobilized -1 +immortal 2 +immune 1 +impatient -2 +imperfect -2 +importance 2 +important 2 +impose -1 +imposed -1 +imposes -1 +imposing -1 +impotent -2 +impress 3 +impressed 3 +impresses 3 +impressive 3 +imprisoned -2 +improve 2 +improved 2 +improvement 2 +improves 2 +improving 2 +inability -2 +inaction -2 +inadequate -2 +incapable -2 +incapacitated -2 +incensed -2 +incompetence -2 +incompetent -2 +inconsiderate -2 +inconvenience -2 +inconvenient -2 +increase 1 +increased 1 +indecisive -2 +indestructible 2 +indifference -2 +indifferent -2 +indignant -2 +indignation -2 +indoctrinate -2 +indoctrinated -2 +indoctrinates -2 +indoctrinating -2 +ineffective -2 +ineffectively -2 +infatuated 2 +infatuation 2 +infected -2 +inferior -2 +inflamed -2 +influential 2 +infringement -2 +infuriate -2 +infuriated -2 +infuriates -2 +infuriating -2 +inhibit -1 +injured -2 +injury -2 +injustice -2 +innovate 1 +innovates 1 +innovation 1 +innovative 2 +inquisition -2 +inquisitive 2 +insane -2 +insanity -2 +insecure -2 +insensitive -2 +insensitivity -2 +insignificant -2 +insipid -2 +inspiration 2 +inspirational 2 +inspire 2 +inspired 2 +inspires 2 +inspiring 3 +insult -2 +insulted -2 +insulting -2 +insults -2 +intact 2 +integrity 2 +intelligent 2 +intense 1 +interest 1 +interested 2 +interesting 2 +interests 1 +interrogated -2 +interrupt -2 +interrupted -2 +interrupting -2 +interruption -2 +interrupts -2 +intimidate -2 +intimidated -2 +intimidates -2 +intimidating -2 +intimidation -2 +intricate 2 +intrigues 1 +invincible 2 +invite 1 +inviting 1 +invulnerable 2 +irate -3 +ironic -1 +irony -1 +irrational -1 +irresistible 2 +irresolute -2 +irresponsible 2 +irreversible -1 +irritate -3 +irritated -3 +irritating -3 +isolated -1 +itchy -2 +jackass -4 +jackasses -4 +jailed -2 +jaunty 2 +jealous -2 +jeopardy -2 +jerk -3 +jesus 1 +jewel 1 +jewels 1 +jocular 2 +join 1 +joke 2 +jokes 2 +jolly 2 +jovial 2 +joy 3 +joyful 3 +joyfully 3 +joyless -2 +joyous 3 +jubilant 3 +jumpy -1 +justice 2 +justifiably 2 +justified 2 +keen 1 +kill -3 +killed -3 +killing -3 +kills -3 +kind 2 +kinder 2 +kiss 2 +kudos 3 +lack -2 +lackadaisical -2 +lag -1 +lagged -2 +lagging -2 +lags -2 +lame -2 +landmark 2 +laugh 1 +laughed 1 +laughing 1 +laughs 1 +laughting 1 +launched 1 +lawl 3 +lawsuit -2 +lawsuits -2 +lazy -1 +leak -1 +leaked -1 +leave -1 +legal 1 +legally 1 +lenient 1 +lethargic -2 +lethargy -2 +liar -3 +liars -3 +libelous -2 +lied -2 +lifesaver 4 +lighthearted 1 +like 2 +liked 2 +likes 2 +limitation -1 +limited -1 +limits -1 +litigation -1 +litigious -2 +lively 2 +livid -2 +lmao 4 +lmfao 4 +loathe -3 +loathed -3 +loathes -3 +loathing -3 +lobby -2 +lobbying -2 +lol 3 +lonely -2 +lonesome -2 +longing -1 +loom -1 +loomed -1 +looming -1 +looms -1 +loose -3 +looses -3 +loser -3 +losing -3 +loss -3 +lost -3 +lovable 3 +love 3 +loved 3 +lovelies 3 +lovely 3 +loving 2 +lowest -1 +loyal 3 +loyalty 3 +luck 3 +luckily 3 +lucky 3 +lugubrious -2 +lunatic -3 +lunatics -3 +lurk -1 +lurking -1 +lurks -1 +mad -3 +maddening -3 +made-up -1 +madly -3 +madness -3 +mandatory -1 +manipulated -1 +manipulating -1 +manipulation -1 +marvel 3 +marvelous 3 +marvels 3 +masterpiece 4 +masterpieces 4 +matter 1 +matters 1 +mature 2 +meaningful 2 +meaningless -2 +medal 3 +mediocrity -3 +meditative 1 +melancholy -2 +menace -2 +menaced -2 +mercy 2 +merry 3 +mess -2 +messed -2 +messing up -2 +methodical 2 +mindless -2 +miracle 4 +mirth 3 +mirthful 3 +mirthfully 3 +misbehave -2 +misbehaved -2 +misbehaves -2 +misbehaving -2 +mischief -1 +mischiefs -1 +miserable -3 +misery -2 +misgiving -2 +misinformation -2 +misinformed -2 +misinterpreted -2 +misleading -3 +misread -1 +misreporting -2 +misrepresentation -2 +miss -2 +missed -2 +missing -2 +mistake -2 +mistaken -2 +mistakes -2 +mistaking -2 +misunderstand -2 +misunderstanding -2 +misunderstands -2 +misunderstood -2 +moan -2 +moaned -2 +moaning -2 +moans -2 +mock -2 +mocked -2 +mocking -2 +mocks -2 +mongering -2 +monopolize -2 +monopolized -2 +monopolizes -2 +monopolizing -2 +moody -1 +mope -1 +moping -1 +moron -3 +motherfucker -5 +motherfucking -5 +motivate 1 +motivated 2 +motivating 2 +motivation 1 +mourn -2 +mourned -2 +mournful -2 +mourning -2 +mourns -2 +mumpish -2 +murder -2 +murderer -2 +murdering -3 +murderous -3 +murders -2 +myth -1 +n00b -2 +naive -2 +nasty -3 +natural 1 +naïve -2 +needy -2 +negative -2 +negativity -2 +neglect -2 +neglected -2 +neglecting -2 +neglects -2 +nerves -1 +nervous -2 +nervously -2 +nice 3 +nifty 2 +niggas -5 +nigger -5 +no -1 +no fun -3 +noble 2 +noisy -1 +nonsense -2 +noob -2 +nosey -2 +not good -2 +not working -3 +notorious -2 +novel 2 +numb -1 +nuts -3 +obliterate -2 +obliterated -2 +obnoxious -3 +obscene -2 +obsessed 2 +obsolete -2 +obstacle -2 +obstacles -2 +obstinate -2 +odd -2 +offend -2 +offended -2 +offender -2 +offending -2 +offends -2 +offline -1 +oks 2 +ominous 3 +once-in-a-lifetime 3 +opportunities 2 +opportunity 2 +oppressed -2 +oppressive -2 +optimism 2 +optimistic 2 +optionless -2 +outcry -2 +outmaneuvered -2 +outrage -3 +outraged -3 +outreach 2 +outstanding 5 +overjoyed 4 +overload -1 +overlooked -1 +overreact -2 +overreacted -2 +overreaction -2 +overreacts -2 +oversell -2 +overselling -2 +oversells -2 +oversimplification -2 +oversimplified -2 +oversimplifies -2 +oversimplify -2 +overstatement -2 +overstatements -2 +overweight -1 +oxymoron -1 +pain -2 +pained -2 +panic -3 +panicked -3 +panics -3 +paradise 3 +paradox -1 +pardon 2 +pardoned 2 +pardoning 2 +pardons 2 +parley -1 +passionate 2 +passive -1 +passively -1 +pathetic -2 +pay -1 +peace 2 +peaceful 2 +peacefully 2 +penalty -2 +pensive -1 +perfect 3 +perfected 2 +perfectly 3 +perfects 2 +peril -2 +perjury -3 +perpetrator -2 +perpetrators -2 +perplexed -2 +persecute -2 +persecuted -2 +persecutes -2 +persecuting -2 +perturbed -2 +pesky -2 +pessimism -2 +pessimistic -2 +petrified -2 +phobic -2 +picturesque 2 +pileup -1 +pique -2 +piqued -2 +piss -4 +pissed -4 +pissing -3 +piteous -2 +pitied -1 +pity -2 +playful 2 +pleasant 3 +please 1 +pleased 3 +pleasure 3 +poised -2 +poison -2 +poisoned -2 +poisons -2 +pollute -2 +polluted -2 +polluter -2 +polluters -2 +pollutes -2 +poor -2 +poorer -2 +poorest -2 +popular 3 +positive 2 +positively 2 +possessive -2 +postpone -1 +postponed -1 +postpones -1 +postponing -1 +poverty -1 +powerful 2 +powerless -2 +praise 3 +praised 3 +praises 3 +praising 3 +pray 1 +praying 1 +prays 1 +prblm -2 +prblms -2 +prepared 1 +pressure -1 +pressured -2 +pretend -1 +pretending -1 +pretends -1 +pretty 1 +prevent -1 +prevented -1 +preventing -1 +prevents -1 +prick -5 +prison -2 +prisoner -2 +prisoners -2 +privileged 2 +proactive 2 +problem -2 +problems -2 +profiteer -2 +progress 2 +prominent 2 +promise 1 +promised 1 +promises 1 +promote 1 +promoted 1 +promotes 1 +promoting 1 +propaganda -2 +prosecute -1 +prosecuted -2 +prosecutes -1 +prosecution -1 +prospect 1 +prospects 1 +prosperous 3 +protect 1 +protected 1 +protects 1 +protest -2 +protesters -2 +protesting -2 +protests -2 +proud 2 +proudly 2 +provoke -1 +provoked -1 +provokes -1 +provoking -1 +pseudoscience -3 +punish -2 +punished -2 +punishes -2 +punitive -2 +pushy -1 +puzzled -2 +quaking -2 +questionable -2 +questioned -1 +questioning -1 +racism -3 +racist -3 +racists -3 +rage -2 +rageful -2 +rainy -1 +rant -3 +ranter -3 +ranters -3 +rants -3 +rape -4 +rapist -4 +rapture 2 +raptured 2 +raptures 2 +rapturous 4 +rash -2 +ratified 2 +reach 1 +reached 1 +reaches 1 +reaching 1 +reassure 1 +reassured 1 +reassures 1 +reassuring 2 +rebellion -2 +recession -2 +reckless -2 +recommend 2 +recommended 2 +recommends 2 +redeemed 2 +refuse -2 +refused -2 +refusing -2 +regret -2 +regretful -2 +regrets -2 +regretted -2 +regretting -2 +reject -1 +rejected -1 +rejecting -1 +rejects -1 +rejoice 4 +rejoiced 4 +rejoices 4 +rejoicing 4 +relaxed 2 +relentless -1 +reliant 2 +relieve 1 +relieved 2 +relieves 1 +relieving 2 +relishing 2 +remarkable 2 +remorse -2 +repulse -1 +repulsed -2 +rescue 2 +rescued 2 +rescues 2 +resentful -2 +resign -1 +resigned -1 +resigning -1 +resigns -1 +resolute 2 +resolve 2 +resolved 2 +resolves 2 +resolving 2 +respected 2 +responsible 2 +responsive 2 +restful 2 +restless -2 +restore 1 +restored 1 +restores 1 +restoring 1 +restrict -2 +restricted -2 +restricting -2 +restriction -2 +restricts -2 +retained -1 +retard -2 +retarded -2 +retreat -1 +revenge -2 +revengeful -2 +revered 2 +revive 2 +revives 2 +reward 2 +rewarded 2 +rewarding 2 +rewards 2 +rich 2 +ridiculous -3 +rig -1 +rigged -1 +right direction 3 +rigorous 3 +rigorously 3 +riot -2 +riots -2 +risk -2 +risks -2 +rob -2 +robber -2 +robed -2 +robing -2 +robs -2 +robust 2 +rofl 4 +roflcopter 4 +roflmao 4 +romance 2 +rotfl 4 +rotflmfao 4 +rotflol 4 +ruin -2 +ruined -2 +ruining -2 +ruins -2 +sabotage -2 +sad -2 +sadden -2 +saddened -2 +sadly -2 +safe 1 +safely 1 +safety 1 +salient 1 +sappy -1 +sarcastic -2 +satisfied 2 +save 2 +saved 2 +scam -2 +scams -2 +scandal -3 +scandalous -3 +scandals -3 +scapegoat -2 +scapegoats -2 +scare -2 +scared -2 +scary -2 +sceptical -2 +scold -2 +scoop 3 +scorn -2 +scornful -2 +scream -2 +screamed -2 +screaming -2 +screams -2 +screwed -2 +screwed up -3 +scumbag -4 +secure 2 +secured 2 +secures 2 +sedition -2 +seditious -2 +seduced -1 +self-confident 2 +self-deluded -2 +selfish -3 +selfishness -3 +sentence -2 +sentenced -2 +sentences -2 +sentencing -2 +serene 2 +severe -2 +sexy 3 +shaky -2 +shame -2 +shamed -2 +shameful -2 +share 1 +shared 1 +shares 1 +shattered -2 +shit -4 +shithead -4 +shitty -3 +shock -2 +shocked -2 +shocking -2 +shocks -2 +shoot -1 +short-sighted -2 +short-sightedness -2 +shortage -2 +shortages -2 +shrew -4 +shy -1 +sick -2 +sigh -2 +significance 1 +significant 1 +silencing -1 +silly -1 +sincere 2 +sincerely 2 +sincerest 2 +sincerity 2 +sinful -3 +singleminded -2 +skeptic -2 +skeptical -2 +skepticism -2 +skeptics -2 +slam -2 +slash -2 +slashed -2 +slashes -2 +slashing -2 +slavery -3 +sleeplessness -2 +slick 2 +slicker 2 +slickest 2 +sluggish -2 +slut -5 +smart 1 +smarter 2 +smartest 2 +smear -2 +smile 2 +smiled 2 +smiles 2 +smiling 2 +smog -2 +sneaky -1 +snub -2 +snubbed -2 +snubbing -2 +snubs -2 +sobering 1 +solemn -1 +solid 2 +solidarity 2 +solution 1 +solutions 1 +solve 1 +solved 1 +solves 1 +solving 1 +somber -2 +some kind 0 +son-of-a-bitch -5 +soothe 3 +soothed 3 +soothing 3 +sophisticated 2 +sore -1 +sorrow -2 +sorrowful -2 +sorry -1 +spam -2 +spammer -3 +spammers -3 +spamming -2 +spark 1 +sparkle 3 +sparkles 3 +sparkling 3 +speculative -2 +spirit 1 +spirited 2 +spiritless -2 +spiteful -2 +splendid 3 +sprightly 2 +squelched -1 +stab -2 +stabbed -2 +stable 2 +stabs -2 +stall -2 +stalled -2 +stalling -2 +stamina 2 +stampede -2 +startled -2 +starve -2 +starved -2 +starves -2 +starving -2 +steadfast 2 +steal -2 +steals -2 +stereotype -2 +stereotyped -2 +stifled -1 +stimulate 1 +stimulated 1 +stimulates 1 +stimulating 2 +stingy -2 +stolen -2 +stop -1 +stopped -1 +stopping -1 +stops -1 +stout 2 +straight 1 +strange -1 +strangely -1 +strangled -2 +strength 2 +strengthen 2 +strengthened 2 +strengthening 2 +strengthens 2 +stressed -2 +stressor -2 +stressors -2 +stricken -2 +strike -1 +strikers -2 +strikes -1 +strong 2 +stronger 2 +strongest 2 +struck -1 +struggle -2 +struggled -2 +struggles -2 +struggling -2 +stubborn -2 +stuck -2 +stunned -2 +stunning 4 +stupid -2 +stupidly -2 +suave 2 +substantial 1 +substantially 1 +subversive -2 +success 2 +successful 3 +suck -3 +sucks -3 +suffer -2 +suffering -2 +suffers -2 +suicidal -2 +suicide -2 +suing -2 +sulking -2 +sulky -2 +sullen -2 +sunshine 2 +super 3 +superb 5 +superior 2 +support 2 +supported 2 +supporter 1 +supporters 1 +supporting 1 +supportive 2 +supports 2 +survived 2 +surviving 2 +survivor 2 +suspect -1 +suspected -1 +suspecting -1 +suspects -1 +suspend -1 +suspended -1 +suspicious -2 +swear -2 +swearing -2 +swears -2 +sweet 2 +swift 2 +swiftly 2 +swindle -3 +swindles -3 +swindling -3 +sympathetic 2 +sympathy 2 +tard -2 +tears -2 +tender 2 +tense -2 +tension -1 +terrible -3 +terribly -3 +terrific 4 +terrified -3 +terror -3 +terrorize -3 +terrorized -3 +terrorizes -3 +thank 2 +thankful 2 +thanks 2 +thorny -2 +thoughtful 2 +thoughtless -2 +threat -2 +threaten -2 +threatened -2 +threatening -2 +threatens -2 +threats -2 +thrilled 5 +thwart -2 +thwarted -2 +thwarting -2 +thwarts -2 +timid -2 +timorous -2 +tired -2 +tits -2 +tolerant 2 +toothless -2 +top 2 +tops 2 +torn -2 +torture -4 +tortured -4 +tortures -4 +torturing -4 +totalitarian -2 +totalitarianism -2 +tout -2 +touted -2 +touting -2 +touts -2 +tragedy -2 +tragic -2 +tranquil 2 +trap -1 +trapped -2 +trauma -3 +traumatic -3 +travesty -2 +treason -3 +treasonous -3 +treasure 2 +treasures 2 +trembling -2 +tremulous -2 +tricked -2 +trickery -2 +triumph 4 +triumphant 4 +trouble -2 +troubled -2 +troubles -2 +true 2 +trust 1 +trusted 2 +tumor -2 +twat -5 +ugly -3 +unacceptable -2 +unappreciated -2 +unapproved -2 +unaware -2 +unbelievable -1 +unbelieving -1 +unbiased 2 +uncertain -1 +unclear -1 +uncomfortable -2 +unconcerned -2 +unconfirmed -1 +unconvinced -1 +uncredited -1 +undecided -1 +underestimate -1 +underestimated -1 +underestimates -1 +underestimating -1 +undermine -2 +undermined -2 +undermines -2 +undermining -2 +undeserving -2 +undesirable -2 +uneasy -2 +unemployment -2 +unequal -1 +unequaled 2 +unethical -2 +unfair -2 +unfocused -2 +unfulfilled -2 +unhappy -2 +unhealthy -2 +unified 1 +unimpressed -2 +unintelligent -2 +united 1 +unjust -2 +unlovable -2 +unloved -2 +unmatched 1 +unmotivated -2 +unprofessional -2 +unresearched -2 +unsatisfied -2 +unsecured -2 +unsettled -1 +unsophisticated -2 +unstable -2 +unstoppable 2 +unsupported -2 +unsure -1 +untarnished 2 +unwanted -2 +unworthy -2 +upset -2 +upsets -2 +upsetting -2 +uptight -2 +urgent -1 +useful 2 +usefulness 2 +useless -2 +uselessness -2 +vague -2 +validate 1 +validated 1 +validates 1 +validating 1 +verdict -1 +verdicts -1 +vested 1 +vexation -2 +vexing -2 +vibrant 3 +vicious -2 +victim -3 +victimize -3 +victimized -3 +victimizes -3 +victimizing -3 +victims -3 +vigilant 3 +vile -3 +vindicate 2 +vindicated 2 +vindicates 2 +vindicating 2 +violate -2 +violated -2 +violates -2 +violating -2 +violence -3 +violent -3 +virtuous 2 +virulent -2 +vision 1 +visionary 3 +visioning 1 +visions 1 +vitality 3 +vitamin 1 +vitriolic -3 +vivacious 3 +vociferous -1 +vulnerability -2 +vulnerable -2 +walkout -2 +walkouts -2 +wanker -3 +want 1 +war -2 +warfare -2 +warm 1 +warmth 2 +warn -2 +warned -2 +warning -3 +warnings -3 +warns -2 +waste -1 +wasted -2 +wasting -2 +wavering -1 +weak -2 +weakness -2 +wealth 3 +wealthy 2 +weary -2 +weep -2 +weeping -2 +weird -2 +welcome 2 +welcomed 2 +welcomes 2 +whimsical 1 +whitewash -3 +whore -4 +wicked -2 +widowed -1 +willingness 2 +win 4 +winner 4 +winning 4 +wins 4 +winwin 3 +wish 1 +wishes 1 +wishing 1 +withdrawal -3 +woebegone -2 +woeful -3 +won 3 +wonderful 4 +woo 3 +woohoo 3 +wooo 4 +woow 4 +worn -1 +worried -3 +worry -3 +worrying -3 +worse -3 +worsen -3 +worsened -3 +worsening -3 +worsens -3 +worshiped 3 +worst -3 +worth 2 +worthless -2 +worthy 2 +wow 4 +wowow 4 +wowww 4 +wrathful -3 +wreck -2 +wrong -2 +wronged -2 +wtf -4 +yeah 1 +yearning 1 +yeees 2 +yes 1 +youthful 2 +yucky -2 +yummy 3 +zealot -2 +zealots -2 +zealous 2 \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/examples/jars/scopt_2.12-3.7.1.jar b/spark-3.5.0-bin-hadoop3/examples/jars/scopt_2.12-3.7.1.jar new file mode 100644 index 0000000..991486d Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/examples/jars/scopt_2.12-3.7.1.jar differ diff --git a/spark-3.5.0-bin-hadoop3/examples/jars/spark-examples_2.12-3.5.0.jar b/spark-3.5.0-bin-hadoop3/examples/jars/spark-examples_2.12-3.5.0.jar new file mode 100644 index 0000000..ec899ff Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/examples/jars/spark-examples_2.12-3.5.0.jar differ diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaHdfsLR.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaHdfsLR.java new file mode 100644 index 0000000..362bd44 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaHdfsLR.java @@ -0,0 +1,159 @@ +/* + * 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. + */ + +package org.apache.spark.examples; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.function.Function; +import org.apache.spark.api.java.function.Function2; +import org.apache.spark.sql.SparkSession; + +import java.io.Serializable; +import java.util.Arrays; +import java.util.Random; +import java.util.regex.Pattern; + +/** + * Logistic regression based classification. + * + * This is an example implementation for learning how to use Spark. For more conventional use, + * please refer to org.apache.spark.ml.classification.LogisticRegression. + */ +public final class JavaHdfsLR { + + private static final int D = 10; // Number of dimensions + private static final Random rand = new Random(42); + + static void showWarning() { + String warning = "WARN: This is a naive implementation of Logistic Regression " + + "and is given as an example!\n" + + "Please use org.apache.spark.ml.classification.LogisticRegression " + + "for more conventional use."; + System.err.println(warning); + } + + static class DataPoint implements Serializable { + DataPoint(double[] x, double y) { + this.x = x; + this.y = y; + } + + double[] x; + double y; + } + + static class ParsePoint implements Function { + private static final Pattern SPACE = Pattern.compile(" "); + + @Override + public DataPoint call(String line) { + String[] tok = SPACE.split(line); + double y = Double.parseDouble(tok[0]); + double[] x = new double[D]; + for (int i = 0; i < D; i++) { + x[i] = Double.parseDouble(tok[i + 1]); + } + return new DataPoint(x, y); + } + } + + static class VectorSum implements Function2 { + @Override + public double[] call(double[] a, double[] b) { + double[] result = new double[D]; + for (int j = 0; j < D; j++) { + result[j] = a[j] + b[j]; + } + return result; + } + } + + static class ComputeGradient implements Function { + private final double[] weights; + + ComputeGradient(double[] weights) { + this.weights = weights; + } + + @Override + public double[] call(DataPoint p) { + double[] gradient = new double[D]; + for (int i = 0; i < D; i++) { + double dot = dot(weights, p.x); + gradient[i] = (1 / (1 + Math.exp(-p.y * dot)) - 1) * p.y * p.x[i]; + } + return gradient; + } + } + + public static double dot(double[] a, double[] b) { + double x = 0; + for (int i = 0; i < D; i++) { + x += a[i] * b[i]; + } + return x; + } + + public static void printWeights(double[] a) { + System.out.println(Arrays.toString(a)); + } + + public static void main(String[] args) { + + if (args.length < 2) { + System.err.println("Usage: JavaHdfsLR "); + System.exit(1); + } + + showWarning(); + + SparkSession spark = SparkSession + .builder() + .appName("JavaHdfsLR") + .getOrCreate(); + + JavaRDD lines = spark.read().textFile(args[0]).javaRDD(); + JavaRDD points = lines.map(new ParsePoint()).cache(); + int ITERATIONS = Integer.parseInt(args[1]); + + // Initialize w to a random value + double[] w = new double[D]; + for (int i = 0; i < D; i++) { + w[i] = 2 * rand.nextDouble() - 1; + } + + System.out.print("Initial w: "); + printWeights(w); + + for (int i = 1; i <= ITERATIONS; i++) { + System.out.println("On iteration " + i); + + double[] gradient = points.map( + new ComputeGradient(w) + ).reduce(new VectorSum()); + + for (int j = 0; j < D; j++) { + w[j] -= gradient[j]; + } + + } + + System.out.print("Final w: "); + printWeights(w); + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaLogQuery.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaLogQuery.java new file mode 100644 index 0000000..259843f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaLogQuery.java @@ -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. + */ + +package org.apache.spark.examples; + +import scala.Tuple2; +import scala.Tuple3; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.sql.SparkSession; + +import java.io.Serializable; +import java.util.Arrays; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Executes a roll up-style query against Apache logs. + * + * Usage: JavaLogQuery [logFile] + */ +public final class JavaLogQuery { + + public static final List exampleApacheLogs = Arrays.asList( + "10.10.10.10 - \"FRED\" [18/Jan/2013:17:56:07 +1100] \"GET http://images.com/2013/Generic.jpg " + + "HTTP/1.1\" 304 315 \"http://referall.com/\" \"Mozilla/4.0 (compatible; MSIE 7.0; " + + "Windows NT 5.1; GTB7.4; .NET CLR 2.0.50727; .NET CLR 3.0.04506.30; .NET CLR 3.0.04506.648; " + + ".NET CLR 3.5.21022; .NET CLR 3.0.4506.2152; .NET CLR 1.0.3705; .NET CLR 1.1.4322; .NET CLR " + + "3.5.30729; Release=ARP)\" \"UD-1\" - \"image/jpeg\" \"whatever\" 0.350 \"-\" - \"\" 265 923 934 \"\" " + + "62.24.11.25 images.com 1358492167 - Whatup", + "10.10.10.10 - \"FRED\" [18/Jan/2013:18:02:37 +1100] \"GET http://images.com/2013/Generic.jpg " + + "HTTP/1.1\" 304 306 \"http:/referall.com\" \"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; " + + "GTB7.4; .NET CLR 2.0.50727; .NET CLR 3.0.04506.30; .NET CLR 3.0.04506.648; .NET CLR " + + "3.5.21022; .NET CLR 3.0.4506.2152; .NET CLR 1.0.3705; .NET CLR 1.1.4322; .NET CLR " + + "3.5.30729; Release=ARP)\" \"UD-1\" - \"image/jpeg\" \"whatever\" 0.352 \"-\" - \"\" 256 977 988 \"\" " + + "0 73.23.2.15 images.com 1358492557 - Whatup"); + + public static final Pattern apacheLogRegex = Pattern.compile( + "^([\\d.]+) (\\S+) (\\S+) \\[([\\w\\d:/]+\\s[+\\-]\\d{4})\\] \"(.+?)\" (\\d{3}) ([\\d\\-]+) \"([^\"]+)\" \"([^\"]+)\".*"); + + /** Tracks the total query count and number of aggregate bytes for a particular group. */ + public static class Stats implements Serializable { + + private final int count; + private final int numBytes; + + public Stats(int count, int numBytes) { + this.count = count; + this.numBytes = numBytes; + } + public Stats merge(Stats other) { + return new Stats(count + other.count, numBytes + other.numBytes); + } + + @Override + public String toString() { + return String.format("bytes=%s\tn=%s", numBytes, count); + } + } + + public static Tuple3 extractKey(String line) { + Matcher m = apacheLogRegex.matcher(line); + if (m.find()) { + String ip = m.group(1); + String user = m.group(3); + String query = m.group(5); + if (!user.equalsIgnoreCase("-")) { + return new Tuple3<>(ip, user, query); + } + } + return new Tuple3<>(null, null, null); + } + + public static Stats extractStats(String line) { + Matcher m = apacheLogRegex.matcher(line); + if (m.find()) { + int bytes = Integer.parseInt(m.group(7)); + return new Stats(1, bytes); + } else { + return new Stats(1, 0); + } + } + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaLogQuery") + .getOrCreate(); + + JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext()); + + JavaRDD dataSet = (args.length == 1) ? jsc.textFile(args[0]) : jsc.parallelize(exampleApacheLogs); + + JavaPairRDD, Stats> extracted = + dataSet.mapToPair(s -> new Tuple2<>(extractKey(s), extractStats(s))); + + JavaPairRDD, Stats> counts = extracted.reduceByKey(Stats::merge); + + List, Stats>> output = counts.collect(); + for (Tuple2 t : output) { + System.out.println(t._1() + "\t" + t._2()); + } + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaPageRank.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaPageRank.java new file mode 100644 index 0000000..b5b4703 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaPageRank.java @@ -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. + */ + +package org.apache.spark.examples; + +import java.util.ArrayList; +import java.util.List; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import com.google.common.collect.Iterables; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.function.Function2; +import org.apache.spark.sql.SparkSession; + +/** + * Computes the PageRank of URLs from an input file. Input file should + * be in format of: + * URL neighbor URL + * URL neighbor URL + * URL neighbor URL + * ... + * where URL and their neighbors are separated by space(s). + * + * This is an example implementation for learning how to use Spark. For more conventional use, + * please refer to org.apache.spark.graphx.lib.PageRank + * + * Example Usage: + *
+ * bin/run-example JavaPageRank data/mllib/pagerank_data.txt 10
+ * 
+ */ +public final class JavaPageRank { + private static final Pattern SPACES = Pattern.compile("\\s+"); + + static void showWarning() { + String warning = "WARN: This is a naive implementation of PageRank " + + "and is given as an example! \n" + + "Please use the PageRank implementation found in " + + "org.apache.spark.graphx.lib.PageRank for more conventional use."; + System.err.println(warning); + } + + private static class Sum implements Function2 { + @Override + public Double call(Double a, Double b) { + return a + b; + } + } + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaPageRank "); + System.exit(1); + } + + showWarning(); + + SparkSession spark = SparkSession + .builder() + .appName("JavaPageRank") + .getOrCreate(); + + // Loads in input file. It should be in format of: + // URL neighbor URL + // URL neighbor URL + // URL neighbor URL + // ... + JavaRDD lines = spark.read().textFile(args[0]).javaRDD(); + + // Loads all URLs from input file and initialize their neighbors. + JavaPairRDD> links = lines.mapToPair(s -> { + String[] parts = SPACES.split(s); + return new Tuple2<>(parts[0], parts[1]); + }).distinct().groupByKey().cache(); + + // Loads all URLs with other URL(s) link to from input file and initialize ranks of them to one. + JavaPairRDD ranks = links.mapValues(rs -> 1.0); + + // Calculates and updates URL ranks continuously using PageRank algorithm. + for (int current = 0; current < Integer.parseInt(args[1]); current++) { + // Calculates URL contributions to the rank of other URLs. + JavaPairRDD contribs = links.join(ranks).values() + .flatMapToPair(s -> { + int urlCount = Iterables.size(s._1()); + List> results = new ArrayList<>(); + for (String n : s._1) { + results.add(new Tuple2<>(n, s._2() / urlCount)); + } + return results.iterator(); + }); + + // Re-calculates URL ranks based on neighbor contributions. + ranks = contribs.reduceByKey(new Sum()).mapValues(sum -> 0.15 + sum * 0.85); + } + + // Collects all URL ranks and dump them to console. + List> output = ranks.collect(); + for (Tuple2 tuple : output) { + System.out.println(tuple._1() + " has rank: " + tuple._2() + "."); + } + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaSparkPi.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaSparkPi.java new file mode 100644 index 0000000..37bd8ff --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaSparkPi.java @@ -0,0 +1,60 @@ +/* + * 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. + */ + +package org.apache.spark.examples; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.sql.SparkSession; + +import java.util.ArrayList; +import java.util.List; + +/** + * Computes an approximation to pi + * Usage: JavaSparkPi [partitions] + */ +public final class JavaSparkPi { + + public static void main(String[] args) throws Exception { + SparkSession spark = SparkSession + .builder() + .appName("JavaSparkPi") + .getOrCreate(); + + JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext()); + + int slices = (args.length == 1) ? Integer.parseInt(args[0]) : 2; + int n = 100000 * slices; + List l = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + l.add(i); + } + + JavaRDD dataSet = jsc.parallelize(l, slices); + + int count = dataSet.map(integer -> { + double x = Math.random() * 2 - 1; + double y = Math.random() * 2 - 1; + return (x * x + y * y <= 1) ? 1 : 0; + }).reduce((integer, integer2) -> integer + integer2); + + System.out.println("Pi is roughly " + 4.0 * count / n); + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaStatusTrackerDemo.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaStatusTrackerDemo.java new file mode 100644 index 0000000..b0ebedf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaStatusTrackerDemo.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples; + +import org.apache.spark.SparkJobInfo; +import org.apache.spark.SparkStageInfo; +import org.apache.spark.api.java.JavaFutureAction; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.function.Function; +import org.apache.spark.sql.SparkSession; + +import java.util.Arrays; +import java.util.List; + +/** + * Example of using Spark's status APIs from Java. + */ +public final class JavaStatusTrackerDemo { + + public static final String APP_NAME = "JavaStatusAPIDemo"; + + public static final class IdentityWithDelay implements Function { + @Override + public T call(T x) throws Exception { + Thread.sleep(2 * 1000); // 2 seconds + return x; + } + } + + public static void main(String[] args) throws Exception { + SparkSession spark = SparkSession + .builder() + .appName(APP_NAME) + .getOrCreate(); + + JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext()); + + // Example of implementing a progress reporter for a simple job. + JavaRDD rdd = jsc.parallelize(Arrays.asList(1, 2, 3, 4, 5), 5).map( + new IdentityWithDelay<>()); + JavaFutureAction> jobFuture = rdd.collectAsync(); + while (!jobFuture.isDone()) { + Thread.sleep(1000); // 1 second + List jobIds = jobFuture.jobIds(); + if (jobIds.isEmpty()) { + continue; + } + int currentJobId = jobIds.get(jobIds.size() - 1); + SparkJobInfo jobInfo = jsc.statusTracker().getJobInfo(currentJobId); + SparkStageInfo stageInfo = jsc.statusTracker().getStageInfo(jobInfo.stageIds()[0]); + System.out.println(stageInfo.numTasks() + " tasks total: " + stageInfo.numActiveTasks() + + " active, " + stageInfo.numCompletedTasks() + " complete"); + } + + System.out.println("Job results are: " + jobFuture.get()); + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaTC.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaTC.java new file mode 100644 index 0000000..7e8df69 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaTC.java @@ -0,0 +1,98 @@ +/* + * 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. + */ + +package org.apache.spark.examples; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; + +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.function.PairFunction; +import org.apache.spark.sql.SparkSession; + +/** + * Transitive closure on a graph, implemented in Java. + * Usage: JavaTC [partitions] + */ +public final class JavaTC { + + private static final int numEdges = 200; + private static final int numVertices = 100; + private static final Random rand = new Random(42); + + static List> generateGraph() { + Set> edges = new HashSet<>(numEdges); + while (edges.size() < numEdges) { + int from = rand.nextInt(numVertices); + int to = rand.nextInt(numVertices); + Tuple2 e = new Tuple2<>(from, to); + if (from != to) { + edges.add(e); + } + } + return new ArrayList<>(edges); + } + + static class ProjectFn implements PairFunction>, + Integer, Integer> { + static final ProjectFn INSTANCE = new ProjectFn(); + + @Override + public Tuple2 call(Tuple2> triple) { + return new Tuple2<>(triple._2()._2(), triple._2()._1()); + } + } + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaTC") + .getOrCreate(); + + JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext()); + + int slices = (args.length > 0) ? Integer.parseInt(args[0]): 2; + JavaPairRDD tc = jsc.parallelizePairs(generateGraph(), slices).cache(); + + // Linear transitive closure: each round grows paths by one edge, + // by joining the graph's edges with the already-discovered paths. + // e.g. join the path (y, z) from the TC with the edge (x, y) from + // the graph to obtain the path (x, z). + + // Because join() joins on keys, the edges are stored in reversed order. + JavaPairRDD edges = tc.mapToPair(e -> new Tuple2<>(e._2(), e._1())); + + long oldCount; + long nextCount = tc.count(); + do { + oldCount = nextCount; + // Perform the join, obtaining an RDD of (y, (z, x)) pairs, + // then project the result to obtain the new (x, z) paths. + tc = tc.union(tc.join(edges).mapToPair(ProjectFn.INSTANCE)).distinct().cache(); + nextCount = tc.count(); + } while (nextCount != oldCount); + + System.out.println("TC has " + tc.count() + " edges."); + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaWordCount.java new file mode 100644 index 0000000..f1ce1e9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/JavaWordCount.java @@ -0,0 +1,59 @@ +/* + * 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. + */ + +package org.apache.spark.examples; + +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.sql.SparkSession; + +import java.util.Arrays; +import java.util.List; +import java.util.regex.Pattern; + +public final class JavaWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + + if (args.length < 1) { + System.err.println("Usage: JavaWordCount "); + System.exit(1); + } + + SparkSession spark = SparkSession + .builder() + .appName("JavaWordCount") + .getOrCreate(); + + JavaRDD lines = spark.read().textFile(args[0]).javaRDD(); + + JavaRDD words = lines.flatMap(s -> Arrays.asList(SPACE.split(s)).iterator()); + + JavaPairRDD ones = words.mapToPair(s -> new Tuple2<>(s, 1)); + + JavaPairRDD counts = ones.reduceByKey((i1, i2) -> i1 + i2); + + List> output = counts.collect(); + for (Tuple2 tuple : output) { + System.out.println(tuple._1() + ": " + tuple._2()); + } + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java new file mode 100644 index 0000000..7c741ff --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java @@ -0,0 +1,82 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.regression.AFTSurvivalRegression; +import org.apache.spark.ml.regression.AFTSurvivalRegressionModel; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +/** + * An example demonstrating AFTSurvivalRegression. + * Run with + *
+ * bin/run-example ml.JavaAFTSurvivalRegressionExample
+ * 
+ */ +public class JavaAFTSurvivalRegressionExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaAFTSurvivalRegressionExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(1.218, 1.0, Vectors.dense(1.560, -0.605)), + RowFactory.create(2.949, 0.0, Vectors.dense(0.346, 2.158)), + RowFactory.create(3.627, 0.0, Vectors.dense(1.380, 0.231)), + RowFactory.create(0.273, 1.0, Vectors.dense(0.520, 1.151)), + RowFactory.create(4.199, 0.0, Vectors.dense(0.795, -0.226)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("label", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("censor", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset training = spark.createDataFrame(data, schema); + double[] quantileProbabilities = new double[]{0.3, 0.6}; + AFTSurvivalRegression aft = new AFTSurvivalRegression() + .setQuantileProbabilities(quantileProbabilities) + .setQuantilesCol("quantiles"); + + AFTSurvivalRegressionModel model = aft.fit(training); + + // Print the coefficients, intercept and scale parameter for AFT survival regression + System.out.println("Coefficients: " + model.coefficients()); + System.out.println("Intercept: " + model.intercept()); + System.out.println("Scale: " + model.scale()); + model.transform(training).show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaALSExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaALSExample.java new file mode 100644 index 0000000..b8d2c9f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaALSExample.java @@ -0,0 +1,136 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.io.Serializable; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.recommendation.ALS; +import org.apache.spark.ml.recommendation.ALSModel; +// $example off$ + +public class JavaALSExample { + + // $example on$ + public static class Rating implements Serializable { + private int userId; + private int movieId; + private float rating; + private long timestamp; + + public Rating() {} + + public Rating(int userId, int movieId, float rating, long timestamp) { + this.userId = userId; + this.movieId = movieId; + this.rating = rating; + this.timestamp = timestamp; + } + + public int getUserId() { + return userId; + } + + public int getMovieId() { + return movieId; + } + + public float getRating() { + return rating; + } + + public long getTimestamp() { + return timestamp; + } + + public static Rating parseRating(String str) { + String[] fields = str.split("::"); + if (fields.length != 4) { + throw new IllegalArgumentException("Each line must contain 4 fields"); + } + int userId = Integer.parseInt(fields[0]); + int movieId = Integer.parseInt(fields[1]); + float rating = Float.parseFloat(fields[2]); + long timestamp = Long.parseLong(fields[3]); + return new Rating(userId, movieId, rating, timestamp); + } + } + // $example off$ + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaALSExample") + .getOrCreate(); + + // $example on$ + JavaRDD ratingsRDD = spark + .read().textFile("data/mllib/als/sample_movielens_ratings.txt").javaRDD() + .map(Rating::parseRating); + Dataset ratings = spark.createDataFrame(ratingsRDD, Rating.class); + Dataset[] splits = ratings.randomSplit(new double[]{0.8, 0.2}); + Dataset training = splits[0]; + Dataset test = splits[1]; + + // Build the recommendation model using ALS on the training data + ALS als = new ALS() + .setMaxIter(5) + .setRegParam(0.01) + .setUserCol("userId") + .setItemCol("movieId") + .setRatingCol("rating"); + ALSModel model = als.fit(training); + + // Evaluate the model by computing the RMSE on the test data + // Note we set cold start strategy to 'drop' to ensure we don't get NaN evaluation metrics + model.setColdStartStrategy("drop"); + Dataset predictions = model.transform(test); + + RegressionEvaluator evaluator = new RegressionEvaluator() + .setMetricName("rmse") + .setLabelCol("rating") + .setPredictionCol("prediction"); + double rmse = evaluator.evaluate(predictions); + System.out.println("Root-mean-square error = " + rmse); + + // Generate top 10 movie recommendations for each user + Dataset userRecs = model.recommendForAllUsers(10); + // Generate top 10 user recommendations for each movie + Dataset movieRecs = model.recommendForAllItems(10); + + // Generate top 10 movie recommendations for a specified set of users + Dataset users = ratings.select(als.getUserCol()).distinct().limit(3); + Dataset userSubsetRecs = model.recommendForUserSubset(users, 10); + // Generate top 10 user recommendations for a specified set of movies + Dataset movies = ratings.select(als.getItemCol()).distinct().limit(3); + Dataset movieSubSetRecs = model.recommendForItemSubset(movies, 10); + // $example off$ + userRecs.show(); + movieRecs.show(); + userSubsetRecs.show(); + movieSubSetRecs.show(); + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBinarizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBinarizerExample.java new file mode 100644 index 0000000..3090d8f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBinarizerExample.java @@ -0,0 +1,68 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.Binarizer; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaBinarizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaBinarizerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, 0.1), + RowFactory.create(1, 0.8), + RowFactory.create(2, 0.2) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("feature", DataTypes.DoubleType, false, Metadata.empty()) + }); + Dataset continuousDataFrame = spark.createDataFrame(data, schema); + + Binarizer binarizer = new Binarizer() + .setInputCol("feature") + .setOutputCol("binarized_feature") + .setThreshold(0.5); + + Dataset binarizedDataFrame = binarizer.transform(continuousDataFrame); + + System.out.println("Binarizer output with Threshold = " + binarizer.getThreshold()); + binarizedDataFrame.show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBisectingKMeansExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBisectingKMeansExample.java new file mode 100644 index 0000000..f517dc3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBisectingKMeansExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.clustering.BisectingKMeans; +import org.apache.spark.ml.clustering.BisectingKMeansModel; +import org.apache.spark.ml.evaluation.ClusteringEvaluator; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + + +/** + * An example demonstrating bisecting k-means clustering. + * Run with + *
+ * bin/run-example ml.JavaBisectingKMeansExample
+ * 
+ */ +public class JavaBisectingKMeansExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaBisectingKMeansExample") + .getOrCreate(); + + // $example on$ + // Loads data. + Dataset dataset = spark.read().format("libsvm").load("data/mllib/sample_kmeans_data.txt"); + + // Trains a bisecting k-means model. + BisectingKMeans bkm = new BisectingKMeans().setK(2).setSeed(1); + BisectingKMeansModel model = bkm.fit(dataset); + + // Make predictions + Dataset predictions = model.transform(dataset); + + // Evaluate clustering by computing Silhouette score + ClusteringEvaluator evaluator = new ClusteringEvaluator(); + + double silhouette = evaluator.evaluate(predictions); + System.out.println("Silhouette with squared euclidean distance = " + silhouette); + + // Shows the result. + System.out.println("Cluster Centers: "); + Vector[] centers = model.clusterCenters(); + for (Vector center : centers) { + System.out.println(center); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketedRandomProjectionLSHExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketedRandomProjectionLSHExample.java new file mode 100644 index 0000000..ff917b7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketedRandomProjectionLSHExample.java @@ -0,0 +1,110 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.BucketedRandomProjectionLSH; +import org.apache.spark.ml.feature.BucketedRandomProjectionLSHModel; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import static org.apache.spark.sql.functions.col; +// $example off$ + +/** + * An example demonstrating BucketedRandomProjectionLSH. + * Run with: + * bin/run-example ml.JavaBucketedRandomProjectionLSHExample + */ +public class JavaBucketedRandomProjectionLSHExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaBucketedRandomProjectionLSHExample") + .getOrCreate(); + + // $example on$ + List dataA = Arrays.asList( + RowFactory.create(0, Vectors.dense(1.0, 1.0)), + RowFactory.create(1, Vectors.dense(1.0, -1.0)), + RowFactory.create(2, Vectors.dense(-1.0, -1.0)), + RowFactory.create(3, Vectors.dense(-1.0, 1.0)) + ); + + List dataB = Arrays.asList( + RowFactory.create(4, Vectors.dense(1.0, 0.0)), + RowFactory.create(5, Vectors.dense(-1.0, 0.0)), + RowFactory.create(6, Vectors.dense(0.0, 1.0)), + RowFactory.create(7, Vectors.dense(0.0, -1.0)) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset dfA = spark.createDataFrame(dataA, schema); + Dataset dfB = spark.createDataFrame(dataB, schema); + + Vector key = Vectors.dense(1.0, 0.0); + + BucketedRandomProjectionLSH mh = new BucketedRandomProjectionLSH() + .setBucketLength(2.0) + .setNumHashTables(3) + .setInputCol("features") + .setOutputCol("hashes"); + + BucketedRandomProjectionLSHModel model = mh.fit(dfA); + + // Feature Transformation + System.out.println("The hashed dataset where hashed values are stored in the column 'hashes':"); + model.transform(dfA).show(); + + // Compute the locality sensitive hashes for the input rows, then perform approximate + // similarity join. + // We could avoid computing hashes by passing in the already-transformed dataset, e.g. + // `model.approxSimilarityJoin(transformedA, transformedB, 1.5)` + System.out.println("Approximately joining dfA and dfB on distance smaller than 1.5:"); + model.approxSimilarityJoin(dfA, dfB, 1.5, "EuclideanDistance") + .select(col("datasetA.id").alias("idA"), + col("datasetB.id").alias("idB"), + col("EuclideanDistance")).show(); + + // Compute the locality sensitive hashes for the input rows, then perform approximate nearest + // neighbor search. + // We could avoid computing hashes by passing in the already-transformed dataset, e.g. + // `model.approxNearestNeighbors(transformedA, key, 2)` + System.out.println("Approximately searching dfA for 2 nearest neighbors of the key:"); + model.approxNearestNeighbors(dfA, key, 2).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketizerExample.java new file mode 100644 index 0000000..3e49bf0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaBucketizerExample.java @@ -0,0 +1,116 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.Bucketizer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +/** + * An example for Bucketizer. + * Run with + *
+ * bin/run-example ml.JavaBucketizerExample
+ * 
+ */ +public class JavaBucketizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaBucketizerExample") + .getOrCreate(); + + // $example on$ + double[] splits = {Double.NEGATIVE_INFINITY, -0.5, 0.0, 0.5, Double.POSITIVE_INFINITY}; + + List data = Arrays.asList( + RowFactory.create(-999.9), + RowFactory.create(-0.5), + RowFactory.create(-0.3), + RowFactory.create(0.0), + RowFactory.create(0.2), + RowFactory.create(999.9) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("features", DataTypes.DoubleType, false, Metadata.empty()) + }); + Dataset dataFrame = spark.createDataFrame(data, schema); + + Bucketizer bucketizer = new Bucketizer() + .setInputCol("features") + .setOutputCol("bucketedFeatures") + .setSplits(splits); + + // Transform original data into its bucket index. + Dataset bucketedData = bucketizer.transform(dataFrame); + + System.out.println("Bucketizer output with " + (bucketizer.getSplits().length-1) + " buckets"); + bucketedData.show(); + // $example off$ + + // $example on$ + // Bucketize multiple columns at one pass. + double[][] splitsArray = { + {Double.NEGATIVE_INFINITY, -0.5, 0.0, 0.5, Double.POSITIVE_INFINITY}, + {Double.NEGATIVE_INFINITY, -0.3, 0.0, 0.3, Double.POSITIVE_INFINITY} + }; + + List data2 = Arrays.asList( + RowFactory.create(-999.9, -999.9), + RowFactory.create(-0.5, -0.2), + RowFactory.create(-0.3, -0.1), + RowFactory.create(0.0, 0.0), + RowFactory.create(0.2, 0.4), + RowFactory.create(999.9, 999.9) + ); + StructType schema2 = new StructType(new StructField[]{ + new StructField("features1", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("features2", DataTypes.DoubleType, false, Metadata.empty()) + }); + Dataset dataFrame2 = spark.createDataFrame(data2, schema2); + + Bucketizer bucketizer2 = new Bucketizer() + .setInputCols(new String[] {"features1", "features2"}) + .setOutputCols(new String[] {"bucketedFeatures1", "bucketedFeatures2"}) + .setSplitsArray(splitsArray); + // Transform original data into its bucket index. + Dataset bucketedData2 = bucketizer2.transform(dataFrame2); + + System.out.println("Bucketizer output with [" + + (bucketizer2.getSplitsArray()[0].length-1) + ", " + + (bucketizer2.getSplitsArray()[1].length-1) + "] buckets for each input column"); + bucketedData2.show(); + // $example off$ + + spark.stop(); + } +} + + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSqSelectorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSqSelectorExample.java new file mode 100644 index 0000000..7373896 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSqSelectorExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.ChiSqSelector; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaChiSqSelectorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaChiSqSelectorExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(7, Vectors.dense(0.0, 0.0, 18.0, 1.0), 1.0), + RowFactory.create(8, Vectors.dense(0.0, 1.0, 12.0, 0.0), 0.0), + RowFactory.create(9, Vectors.dense(1.0, 0.0, 15.0, 0.1), 0.0) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()), + new StructField("clicked", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + ChiSqSelector selector = new ChiSqSelector() + .setNumTopFeatures(1) + .setFeaturesCol("features") + .setLabelCol("clicked") + .setOutputCol("selectedFeatures"); + + Dataset result = selector.fit(df).transform(df); + + System.out.println("ChiSqSelector output with top " + selector.getNumTopFeatures() + + " features selected"); + result.show(); + + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSquareTestExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSquareTestExample.java new file mode 100644 index 0000000..4b39350 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaChiSquareTestExample.java @@ -0,0 +1,75 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.stat.ChiSquareTest; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example for Chi-square hypothesis testing. + * Run with + *
+ * bin/run-example ml.JavaChiSquareTestExample
+ * 
+ */ +public class JavaChiSquareTestExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaChiSquareTestExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0.0, Vectors.dense(0.5, 10.0)), + RowFactory.create(0.0, Vectors.dense(1.5, 20.0)), + RowFactory.create(1.0, Vectors.dense(1.5, 30.0)), + RowFactory.create(0.0, Vectors.dense(3.5, 30.0)), + RowFactory.create(0.0, Vectors.dense(3.5, 40.0)), + RowFactory.create(1.0, Vectors.dense(3.5, 40.0)) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("label", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()), + }); + + Dataset df = spark.createDataFrame(data, schema); + Row r = ChiSquareTest.test(df, "features", "label").head(); + System.out.println("pValues: " + r.get(0).toString()); + System.out.println("degreesOfFreedom: " + r.getList(1).toString()); + System.out.println("statistics: " + r.get(2).toString()); + + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCorrelationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCorrelationExample.java new file mode 100644 index 0000000..2a6d62a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCorrelationExample.java @@ -0,0 +1,72 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.stat.Correlation; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example for computing correlation matrix. + * Run with + *
+ * bin/run-example ml.JavaCorrelationExample
+ * 
+ */ +public class JavaCorrelationExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaCorrelationExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Vectors.sparse(4, new int[]{0, 3}, new double[]{1.0, -2.0})), + RowFactory.create(Vectors.dense(4.0, 5.0, 0.0, 3.0)), + RowFactory.create(Vectors.dense(6.0, 7.0, 0.0, 8.0)), + RowFactory.create(Vectors.sparse(4, new int[]{0, 3}, new double[]{9.0, 1.0})) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("features", new VectorUDT(), false, Metadata.empty()), + }); + + Dataset df = spark.createDataFrame(data, schema); + Row r1 = Correlation.corr(df, "features").head(); + System.out.println("Pearson correlation matrix:\n" + r1.get(0).toString()); + + Row r2 = Correlation.corr(df, "features", "spearman").head(); + System.out.println("Spearman correlation matrix:\n" + r2.get(0).toString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCountVectorizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCountVectorizerExample.java new file mode 100644 index 0000000..ac2a86c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaCountVectorizerExample.java @@ -0,0 +1,69 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.CountVectorizer; +import org.apache.spark.ml.feature.CountVectorizerModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +public class JavaCountVectorizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaCountVectorizerExample") + .getOrCreate(); + + // $example on$ + // Input data: Each row is a bag of words from a sentence or document. + List data = Arrays.asList( + RowFactory.create(Arrays.asList("a", "b", "c")), + RowFactory.create(Arrays.asList("a", "b", "b", "c", "a")) + ); + StructType schema = new StructType(new StructField [] { + new StructField("text", new ArrayType(DataTypes.StringType, true), false, Metadata.empty()) + }); + Dataset df = spark.createDataFrame(data, schema); + + // fit a CountVectorizerModel from the corpus + CountVectorizerModel cvModel = new CountVectorizer() + .setInputCol("text") + .setOutputCol("feature") + .setVocabSize(3) + .setMinDF(2) + .fit(df); + + // alternatively, define CountVectorizerModel with a-priori vocabulary + CountVectorizerModel cvm = new CountVectorizerModel(new String[]{"a", "b", "c"}) + .setInputCol("text") + .setOutputCol("feature"); + + cvModel.transform(df).show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDCTExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDCTExample.java new file mode 100644 index 0000000..04546d2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDCTExample.java @@ -0,0 +1,68 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.DCT; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaDCTExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaDCTExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Vectors.dense(0.0, 1.0, -2.0, 3.0)), + RowFactory.create(Vectors.dense(-1.0, 2.0, 4.0, -7.0)), + RowFactory.create(Vectors.dense(14.0, -2.0, -5.0, 1.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("features", new VectorUDT(), false, Metadata.empty()), + }); + Dataset df = spark.createDataFrame(data, schema); + + DCT dct = new DCT() + .setInputCol("features") + .setOutputCol("featuresDCT") + .setInverse(false); + + Dataset dctDf = dct.transform(df); + + dctDf.select("featuresDCT").show(false); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java new file mode 100644 index 0000000..f7e1144 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java @@ -0,0 +1,104 @@ +/* + * 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. + */ +// scalastyle:off println +package org.apache.spark.examples.ml; +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.DecisionTreeClassifier; +import org.apache.spark.ml.classification.DecisionTreeClassificationModel; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.ml.feature.*; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaDecisionTreeClassificationExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaDecisionTreeClassificationExample") + .getOrCreate(); + + // $example on$ + // Load the data stored in LIBSVM format as a DataFrame. + Dataset data = spark + .read() + .format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + // Index labels, adding metadata to the label column. + // Fit on whole dataset to include all labels in index. + StringIndexerModel labelIndexer = new StringIndexer() + .setInputCol("label") + .setOutputCol("indexedLabel") + .fit(data); + + // Automatically identify categorical features, and index them. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) // features with > 4 distinct values are treated as continuous. + .fit(data); + + // Split the data into training and test sets (30% held out for testing). + Dataset[] splits = data.randomSplit(new double[]{0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a DecisionTree model. + DecisionTreeClassifier dt = new DecisionTreeClassifier() + .setLabelCol("indexedLabel") + .setFeaturesCol("indexedFeatures"); + + // Convert indexed labels back to original labels. + IndexToString labelConverter = new IndexToString() + .setInputCol("prediction") + .setOutputCol("predictedLabel") + .setLabels(labelIndexer.labelsArray()[0]); + + // Chain indexers and tree in a Pipeline. + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[]{labelIndexer, featureIndexer, dt, labelConverter}); + + // Train model. This also runs the indexers. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("predictedLabel", "label", "features").show(5); + + // Select (prediction, true label) and compute test error. + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setLabelCol("indexedLabel") + .setPredictionCol("prediction") + .setMetricName("accuracy"); + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test Error = " + (1.0 - accuracy)); + + DecisionTreeClassificationModel treeModel = + (DecisionTreeClassificationModel) (model.stages()[2]); + System.out.println("Learned classification tree model:\n" + treeModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java new file mode 100644 index 0000000..cffb713 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java @@ -0,0 +1,89 @@ +/* + * 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. + */ +// scalastyle:off println +package org.apache.spark.examples.ml; +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.feature.VectorIndexer; +import org.apache.spark.ml.feature.VectorIndexerModel; +import org.apache.spark.ml.regression.DecisionTreeRegressionModel; +import org.apache.spark.ml.regression.DecisionTreeRegressor; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaDecisionTreeRegressionExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaDecisionTreeRegressionExample") + .getOrCreate(); + // $example on$ + // Load the data stored in LIBSVM format as a DataFrame. + Dataset data = spark.read().format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + // Automatically identify categorical features, and index them. + // Set maxCategories so features with > 4 distinct values are treated as continuous. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) + .fit(data); + + // Split the data into training and test sets (30% held out for testing). + Dataset[] splits = data.randomSplit(new double[]{0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a DecisionTree model. + DecisionTreeRegressor dt = new DecisionTreeRegressor() + .setFeaturesCol("indexedFeatures"); + + // Chain indexer and tree in a Pipeline. + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[]{featureIndexer, dt}); + + // Train model. This also runs the indexer. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("label", "features").show(5); + + // Select (prediction, true label) and compute test error. + RegressionEvaluator evaluator = new RegressionEvaluator() + .setLabelCol("label") + .setPredictionCol("prediction") + .setMetricName("rmse"); + double rmse = evaluator.evaluate(predictions); + System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse); + + DecisionTreeRegressionModel treeModel = + (DecisionTreeRegressionModel) (model.stages()[1]); + System.out.println("Learned regression tree model:\n" + treeModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDocument.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDocument.java new file mode 100644 index 0000000..6459dab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaDocument.java @@ -0,0 +1,43 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import java.io.Serializable; + +/** + * Unlabeled instance type, Spark SQL can infer schema from Java Beans. + */ +@SuppressWarnings("serial") +public class JavaDocument implements Serializable { + + private long id; + private String text; + + public JavaDocument(long id, String text) { + this.id = id; + this.text = text; + } + + public long getId() { + return this.id; + } + + public String getText() { + return this.text; + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java new file mode 100644 index 0000000..d2e70c2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.ElementwiseProduct; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaElementwiseProductExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaElementwiseProductExample") + .getOrCreate(); + + // $example on$ + // Create some vector data; also works for sparse vectors + List data = Arrays.asList( + RowFactory.create("a", Vectors.dense(1.0, 2.0, 3.0)), + RowFactory.create("b", Vectors.dense(4.0, 5.0, 6.0)) + ); + + List fields = new ArrayList<>(2); + fields.add(DataTypes.createStructField("id", DataTypes.StringType, false)); + fields.add(DataTypes.createStructField("vector", new VectorUDT(), false)); + + StructType schema = DataTypes.createStructType(fields); + + Dataset dataFrame = spark.createDataFrame(data, schema); + + Vector transformingVector = Vectors.dense(0.0, 1.0, 2.0); + + ElementwiseProduct transformer = new ElementwiseProduct() + .setScalingVec(transformingVector) + .setInputCol("vector") + .setOutputCol("transformedVector"); + + // Batch transform the vectors to create new column: + transformer.transform(dataFrame).show(); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaEstimatorTransformerParamExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaEstimatorTransformerParamExample.java new file mode 100644 index 0000000..9e07a0c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaEstimatorTransformerParamExample.java @@ -0,0 +1,117 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.classification.LogisticRegressionModel; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.param.ParamMap; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * Java example for Estimator, Transformer, and Param. + */ +public class JavaEstimatorTransformerParamExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaEstimatorTransformerParamExample") + .getOrCreate(); + + // $example on$ + // Prepare training data. + List dataTraining = Arrays.asList( + RowFactory.create(1.0, Vectors.dense(0.0, 1.1, 0.1)), + RowFactory.create(0.0, Vectors.dense(2.0, 1.0, -1.0)), + RowFactory.create(0.0, Vectors.dense(2.0, 1.3, 1.0)), + RowFactory.create(1.0, Vectors.dense(0.0, 1.2, -0.5)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("label", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset training = spark.createDataFrame(dataTraining, schema); + + // Create a LogisticRegression instance. This instance is an Estimator. + LogisticRegression lr = new LogisticRegression(); + // Print out the parameters, documentation, and any default values. + System.out.println("LogisticRegression parameters:\n" + lr.explainParams() + "\n"); + + // We may set parameters using setter methods. + lr.setMaxIter(10).setRegParam(0.01); + + // Learn a LogisticRegression model. This uses the parameters stored in lr. + LogisticRegressionModel model1 = lr.fit(training); + // Since model1 is a Model (i.e., a Transformer produced by an Estimator), + // we can view the parameters it used during fit(). + // This prints the parameter (name: value) pairs, where names are unique IDs for this + // LogisticRegression instance. + System.out.println("Model 1 was fit using parameters: " + model1.parent().extractParamMap()); + + // We may alternatively specify parameters using a ParamMap. + ParamMap paramMap = new ParamMap() + .put(lr.maxIter().w(20)) // Specify 1 Param. + .put(lr.maxIter(), 30) // This overwrites the original maxIter. + .put(lr.regParam().w(0.1), lr.threshold().w(0.55)); // Specify multiple Params. + + // One can also combine ParamMaps. + ParamMap paramMap2 = new ParamMap() + .put(lr.probabilityCol().w("myProbability")); // Change output column name + ParamMap paramMapCombined = paramMap.$plus$plus(paramMap2); + + // Now learn a new model using the paramMapCombined parameters. + // paramMapCombined overrides all parameters set earlier via lr.set* methods. + LogisticRegressionModel model2 = lr.fit(training, paramMapCombined); + System.out.println("Model 2 was fit using parameters: " + model2.parent().extractParamMap()); + + // Prepare test documents. + List dataTest = Arrays.asList( + RowFactory.create(1.0, Vectors.dense(-1.0, 1.5, 1.3)), + RowFactory.create(0.0, Vectors.dense(3.0, 2.0, -0.1)), + RowFactory.create(1.0, Vectors.dense(0.0, 2.2, -1.5)) + ); + Dataset test = spark.createDataFrame(dataTest, schema); + + // Make predictions on test documents using the Transformer.transform() method. + // LogisticRegression.transform will only use the 'features' column. + // Note that model2.transform() outputs a 'myProbability' column instead of the usual + // 'probability' column since we renamed the lr.probabilityCol parameter previously. + Dataset results = model2.transform(test); + Dataset rows = results.select("features", "label", "myProbability", "prediction"); + for (Row r: rows.collectAsList()) { + System.out.println("(" + r.get(0) + ", " + r.get(1) + ") -> prob=" + r.get(2) + + ", prediction=" + r.get(3)); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMClassifierExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMClassifierExample.java new file mode 100644 index 0000000..2f11082 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMClassifierExample.java @@ -0,0 +1,105 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.FMClassificationModel; +import org.apache.spark.ml.classification.FMClassifier; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.ml.feature.*; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaFMClassifierExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaFMClassifierExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark + .read() + .format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + // Index labels, adding metadata to the label column. + // Fit on whole dataset to include all labels in index. + StringIndexerModel labelIndexer = new StringIndexer() + .setInputCol("label") + .setOutputCol("indexedLabel") + .fit(data); + // Scale features. + MinMaxScalerModel featureScaler = new MinMaxScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures") + .fit(data); + + // Split the data into training and test sets (30% held out for testing) + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a FM model. + FMClassifier fm = new FMClassifier() + .setLabelCol("indexedLabel") + .setFeaturesCol("scaledFeatures") + .setStepSize(0.001); + + // Convert indexed labels back to original labels. + IndexToString labelConverter = new IndexToString() + .setInputCol("prediction") + .setOutputCol("predictedLabel") + .setLabels(labelIndexer.labelsArray()[0]); + + // Create a Pipeline. + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {labelIndexer, featureScaler, fm, labelConverter}); + + // Train model. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("predictedLabel", "label", "features").show(5); + + // Select (prediction, true label) and compute test accuracy. + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setLabelCol("indexedLabel") + .setPredictionCol("prediction") + .setMetricName("accuracy"); + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test Accuracy = " + accuracy); + + FMClassificationModel fmModel = (FMClassificationModel)(model.stages()[2]); + System.out.println("Factors: " + fmModel.factors()); + System.out.println("Linear: " + fmModel.linear()); + System.out.println("Intercept: " + fmModel.intercept()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMRegressorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMRegressorExample.java new file mode 100644 index 0000000..3ade40f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFMRegressorExample.java @@ -0,0 +1,90 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.feature.MinMaxScaler; +import org.apache.spark.ml.feature.MinMaxScalerModel; +import org.apache.spark.ml.regression.FMRegressionModel; +import org.apache.spark.ml.regression.FMRegressor; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaFMRegressorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaFMRegressorExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + // Scale features. + MinMaxScalerModel featureScaler = new MinMaxScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures") + .fit(data); + + // Split the data into training and test sets (30% held out for testing). + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a FM model. + FMRegressor fm = new FMRegressor() + .setLabelCol("label") + .setFeaturesCol("scaledFeatures") + .setStepSize(0.001); + + // Create a Pipeline. + Pipeline pipeline = new Pipeline().setStages(new PipelineStage[] {featureScaler, fm}); + + // Train model. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("prediction", "label", "features").show(5); + + // Select (prediction, true label) and compute test error. + RegressionEvaluator evaluator = new RegressionEvaluator() + .setLabelCol("label") + .setPredictionCol("prediction") + .setMetricName("rmse"); + double rmse = evaluator.evaluate(predictions); + System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse); + + FMRegressionModel fmModel = (FMRegressionModel)(model.stages()[1]); + System.out.println("Factors: " + fmModel.factors()); + System.out.println("Linear: " + fmModel.linear()); + System.out.println("Intercept: " + fmModel.intercept()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFPGrowthExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFPGrowthExample.java new file mode 100644 index 0000000..717ec21 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFPGrowthExample.java @@ -0,0 +1,77 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.fpm.FPGrowth; +import org.apache.spark.ml.fpm.FPGrowthModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example demonstrating FPGrowth. + * Run with + *
+ * bin/run-example ml.JavaFPGrowthExample
+ * 
+ */ +public class JavaFPGrowthExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaFPGrowthExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Arrays.asList("1 2 5".split(" "))), + RowFactory.create(Arrays.asList("1 2 3 5".split(" "))), + RowFactory.create(Arrays.asList("1 2".split(" "))) + ); + StructType schema = new StructType(new StructField[]{ new StructField( + "items", new ArrayType(DataTypes.StringType, true), false, Metadata.empty()) + }); + Dataset itemsDF = spark.createDataFrame(data, schema); + + FPGrowthModel model = new FPGrowth() + .setItemsCol("items") + .setMinSupport(0.5) + .setMinConfidence(0.6) + .fit(itemsDF); + + // Display frequent itemsets. + model.freqItemsets().show(); + + // Display generated association rules. + model.associationRules().show(); + + // transform examines the input items against all the association rules and summarize the + // consequents as prediction + model.transform(itemsDF).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFeatureHasherExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFeatureHasherExample.java new file mode 100644 index 0000000..9730d42 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaFeatureHasherExample.java @@ -0,0 +1,69 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.FeatureHasher; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaFeatureHasherExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaFeatureHasherExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(2.2, true, "1", "foo"), + RowFactory.create(3.3, false, "2", "bar"), + RowFactory.create(4.4, false, "3", "baz"), + RowFactory.create(5.5, false, "4", "foo") + ); + StructType schema = new StructType(new StructField[]{ + new StructField("real", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("bool", DataTypes.BooleanType, false, Metadata.empty()), + new StructField("stringNum", DataTypes.StringType, false, Metadata.empty()), + new StructField("string", DataTypes.StringType, false, Metadata.empty()) + }); + Dataset dataset = spark.createDataFrame(data, schema); + + FeatureHasher hasher = new FeatureHasher() + .setInputCols(new String[]{"real", "bool", "stringNum", "string"}) + .setOutputCol("features"); + + Dataset featurized = hasher.transform(dataset); + + featurized.show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGaussianMixtureExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGaussianMixtureExample.java new file mode 100644 index 0000000..72bd5d0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGaussianMixtureExample.java @@ -0,0 +1,64 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.clustering.GaussianMixture; +import org.apache.spark.ml.clustering.GaussianMixtureModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + + +/** + * An example demonstrating Gaussian Mixture Model. + * Run with + *
+ * bin/run-example ml.JavaGaussianMixtureExample
+ * 
+ */ +public class JavaGaussianMixtureExample { + + public static void main(String[] args) { + + // Creates a SparkSession + SparkSession spark = SparkSession + .builder() + .appName("JavaGaussianMixtureExample") + .getOrCreate(); + + // $example on$ + // Loads data + Dataset dataset = spark.read().format("libsvm").load("data/mllib/sample_kmeans_data.txt"); + + // Trains a GaussianMixture model + GaussianMixture gmm = new GaussianMixture() + .setK(2); + GaussianMixtureModel model = gmm.fit(dataset); + + // Output the parameters of the mixture model + for (int i = 0; i < model.getK(); i++) { + System.out.printf("Gaussian %d:\nweight=%f\nmu=%s\nsigma=\n%s\n\n", + i, model.weights()[i], model.gaussians()[i].mean(), model.gaussians()[i].cov()); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGeneralizedLinearRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGeneralizedLinearRegressionExample.java new file mode 100644 index 0000000..3f072d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGeneralizedLinearRegressionExample.java @@ -0,0 +1,83 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.ml.regression.GeneralizedLinearRegression; +import org.apache.spark.ml.regression.GeneralizedLinearRegressionModel; +import org.apache.spark.ml.regression.GeneralizedLinearRegressionTrainingSummary; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * An example demonstrating generalized linear regression. + * Run with + *
+ * bin/run-example ml.JavaGeneralizedLinearRegressionExample
+ * 
+ */ + +public class JavaGeneralizedLinearRegressionExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaGeneralizedLinearRegressionExample") + .getOrCreate(); + + // $example on$ + // Load training data + Dataset dataset = spark.read().format("libsvm") + .load("data/mllib/sample_linear_regression_data.txt"); + + GeneralizedLinearRegression glr = new GeneralizedLinearRegression() + .setFamily("gaussian") + .setLink("identity") + .setMaxIter(10) + .setRegParam(0.3); + + // Fit the model + GeneralizedLinearRegressionModel model = glr.fit(dataset); + + // Print the coefficients and intercept for generalized linear regression model + System.out.println("Coefficients: " + model.coefficients()); + System.out.println("Intercept: " + model.intercept()); + + // Summarize the model over the training set and print out some metrics + GeneralizedLinearRegressionTrainingSummary summary = model.summary(); + System.out.println("Coefficient Standard Errors: " + + Arrays.toString(summary.coefficientStandardErrors())); + System.out.println("T Values: " + Arrays.toString(summary.tValues())); + System.out.println("P Values: " + Arrays.toString(summary.pValues())); + System.out.println("Dispersion: " + summary.dispersion()); + System.out.println("Null Deviance: " + summary.nullDeviance()); + System.out.println("Residual Degree Of Freedom Null: " + summary.residualDegreeOfFreedomNull()); + System.out.println("Deviance: " + summary.deviance()); + System.out.println("Residual Degree Of Freedom: " + summary.residualDegreeOfFreedom()); + System.out.println("AIC: " + summary.aic()); + System.out.println("Deviance Residuals: "); + summary.residuals().show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java new file mode 100644 index 0000000..807027a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java @@ -0,0 +1,105 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.GBTClassificationModel; +import org.apache.spark.ml.classification.GBTClassifier; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.ml.feature.*; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaGradientBoostedTreeClassifierExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaGradientBoostedTreeClassifierExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark + .read() + .format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + // Index labels, adding metadata to the label column. + // Fit on whole dataset to include all labels in index. + StringIndexerModel labelIndexer = new StringIndexer() + .setInputCol("label") + .setOutputCol("indexedLabel") + .fit(data); + // Automatically identify categorical features, and index them. + // Set maxCategories so features with > 4 distinct values are treated as continuous. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) + .fit(data); + + // Split the data into training and test sets (30% held out for testing) + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a GBT model. + GBTClassifier gbt = new GBTClassifier() + .setLabelCol("indexedLabel") + .setFeaturesCol("indexedFeatures") + .setMaxIter(10); + + // Convert indexed labels back to original labels. + IndexToString labelConverter = new IndexToString() + .setInputCol("prediction") + .setOutputCol("predictedLabel") + .setLabels(labelIndexer.labelsArray()[0]); + + // Chain indexers and GBT in a Pipeline. + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {labelIndexer, featureIndexer, gbt, labelConverter}); + + // Train model. This also runs the indexers. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("predictedLabel", "label", "features").show(5); + + // Select (prediction, true label) and compute test error. + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setLabelCol("indexedLabel") + .setPredictionCol("prediction") + .setMetricName("accuracy"); + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test Error = " + (1.0 - accuracy)); + + GBTClassificationModel gbtModel = (GBTClassificationModel)(model.stages()[2]); + System.out.println("Learned classification GBT model:\n" + gbtModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java new file mode 100644 index 0000000..769b5c3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java @@ -0,0 +1,90 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.feature.VectorIndexer; +import org.apache.spark.ml.feature.VectorIndexerModel; +import org.apache.spark.ml.regression.GBTRegressionModel; +import org.apache.spark.ml.regression.GBTRegressor; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaGradientBoostedTreeRegressorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaGradientBoostedTreeRegressorExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + // Automatically identify categorical features, and index them. + // Set maxCategories so features with > 4 distinct values are treated as continuous. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) + .fit(data); + + // Split the data into training and test sets (30% held out for testing). + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a GBT model. + GBTRegressor gbt = new GBTRegressor() + .setLabelCol("label") + .setFeaturesCol("indexedFeatures") + .setMaxIter(10); + + // Chain indexer and GBT in a Pipeline. + Pipeline pipeline = new Pipeline().setStages(new PipelineStage[] {featureIndexer, gbt}); + + // Train model. This also runs the indexer. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("prediction", "label", "features").show(5); + + // Select (prediction, true label) and compute test error. + RegressionEvaluator evaluator = new RegressionEvaluator() + .setLabelCol("label") + .setPredictionCol("prediction") + .setMetricName("rmse"); + double rmse = evaluator.evaluate(predictions); + System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse); + + GBTRegressionModel gbtModel = (GBTRegressionModel)(model.stages()[1]); + System.out.println("Learned regression GBT model:\n" + gbtModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaImputerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaImputerExample.java new file mode 100644 index 0000000..ac40ccd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaImputerExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.Imputer; +import org.apache.spark.ml.feature.ImputerModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +import static org.apache.spark.sql.types.DataTypes.*; + +/** + * An example demonstrating Imputer. + * Run with: + * bin/run-example ml.JavaImputerExample + */ +public class JavaImputerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaImputerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(1.0, Double.NaN), + RowFactory.create(2.0, Double.NaN), + RowFactory.create(Double.NaN, 3.0), + RowFactory.create(4.0, 4.0), + RowFactory.create(5.0, 5.0) + ); + StructType schema = new StructType(new StructField[]{ + createStructField("a", DoubleType, false), + createStructField("b", DoubleType, false) + }); + Dataset df = spark.createDataFrame(data, schema); + + Imputer imputer = new Imputer() + .setInputCols(new String[]{"a", "b"}) + .setOutputCols(new String[]{"out_a", "out_b"}); + + ImputerModel model = imputer.fit(df); + model.transform(df).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIndexToStringExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIndexToStringExample.java new file mode 100644 index 0000000..6965512 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIndexToStringExample.java @@ -0,0 +1,87 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.attribute.Attribute; +import org.apache.spark.ml.feature.IndexToString; +import org.apache.spark.ml.feature.StringIndexer; +import org.apache.spark.ml.feature.StringIndexerModel; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaIndexToStringExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaIndexToStringExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, "a"), + RowFactory.create(1, "b"), + RowFactory.create(2, "c"), + RowFactory.create(3, "a"), + RowFactory.create(4, "a"), + RowFactory.create(5, "c") + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("category", DataTypes.StringType, false, Metadata.empty()) + }); + Dataset df = spark.createDataFrame(data, schema); + + StringIndexerModel indexer = new StringIndexer() + .setInputCol("category") + .setOutputCol("categoryIndex") + .fit(df); + Dataset indexed = indexer.transform(df); + + System.out.println("Transformed string column '" + indexer.getInputCol() + "' " + + "to indexed column '" + indexer.getOutputCol() + "'"); + indexed.show(); + + StructField inputColSchema = indexed.schema().apply(indexer.getOutputCol()); + System.out.println("StringIndexer will store labels in output column metadata: " + + Attribute.fromStructField(inputColSchema).toString() + "\n"); + + IndexToString converter = new IndexToString() + .setInputCol("categoryIndex") + .setOutputCol("originalCategory"); + Dataset converted = converter.transform(indexed); + + System.out.println("Transformed indexed column '" + converter.getInputCol() + "' back to " + + "original string column '" + converter.getOutputCol() + "' using labels in metadata"); + converted.select("id", "categoryIndex", "originalCategory").show(); + + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaInteractionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaInteractionExample.java new file mode 100644 index 0000000..3684a87 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaInteractionExample.java @@ -0,0 +1,87 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.ml.feature.Interaction; +import org.apache.spark.ml.feature.VectorAssembler; +import org.apache.spark.sql.*; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import java.util.Arrays; +import java.util.List; + +// $example on$ +// $example off$ + +public class JavaInteractionExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaInteractionExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(1, 1, 2, 3, 8, 4, 5), + RowFactory.create(2, 4, 3, 8, 7, 9, 8), + RowFactory.create(3, 6, 1, 9, 2, 3, 6), + RowFactory.create(4, 10, 8, 6, 9, 4, 5), + RowFactory.create(5, 9, 2, 7, 10, 7, 3), + RowFactory.create(6, 1, 1, 4, 2, 8, 4) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id1", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id2", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id3", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id4", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id5", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id6", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("id7", DataTypes.IntegerType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + VectorAssembler assembler1 = new VectorAssembler() + .setInputCols(new String[]{"id2", "id3", "id4"}) + .setOutputCol("vec1"); + + Dataset assembled1 = assembler1.transform(df); + + VectorAssembler assembler2 = new VectorAssembler() + .setInputCols(new String[]{"id5", "id6", "id7"}) + .setOutputCol("vec2"); + + Dataset assembled2 = assembler2.transform(assembled1).select("id1", "vec1", "vec2"); + + Interaction interaction = new Interaction() + .setInputCols(new String[]{"id1","vec1","vec2"}) + .setOutputCol("interactedCol"); + + Dataset interacted = interaction.transform(assembled2); + + interacted.show(false); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIsotonicRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIsotonicRegressionExample.java new file mode 100644 index 0000000..a7de8e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaIsotonicRegressionExample.java @@ -0,0 +1,62 @@ +/* + * 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. + */ +package org.apache.spark.examples.ml; + +// $example on$ + +import org.apache.spark.ml.regression.IsotonicRegression; +import org.apache.spark.ml.regression.IsotonicRegressionModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * An example demonstrating IsotonicRegression. + * Run with + *
+ * bin/run-example ml.JavaIsotonicRegressionExample
+ * 
+ */ +public class JavaIsotonicRegressionExample { + + public static void main(String[] args) { + // Create a SparkSession. + SparkSession spark = SparkSession + .builder() + .appName("JavaIsotonicRegressionExample") + .getOrCreate(); + + // $example on$ + // Loads data. + Dataset dataset = spark.read().format("libsvm") + .load("data/mllib/sample_isotonic_regression_libsvm_data.txt"); + + // Trains an isotonic regression model. + IsotonicRegression ir = new IsotonicRegression(); + IsotonicRegressionModel model = ir.fit(dataset); + + System.out.println("Boundaries in increasing order: " + model.boundaries() + "\n"); + System.out.println("Predictions associated with the boundaries: " + model.predictions() + "\n"); + + // Makes predictions. + model.transform(dataset).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java new file mode 100644 index 0000000..dc4b0bc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.clustering.KMeansModel; +import org.apache.spark.ml.clustering.KMeans; +import org.apache.spark.ml.evaluation.ClusteringEvaluator; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + + +/** + * An example demonstrating k-means clustering. + * Run with + *
+ * bin/run-example ml.JavaKMeansExample
+ * 
+ */ +public class JavaKMeansExample { + + public static void main(String[] args) { + // Create a SparkSession. + SparkSession spark = SparkSession + .builder() + .appName("JavaKMeansExample") + .getOrCreate(); + + // $example on$ + // Loads data. + Dataset dataset = spark.read().format("libsvm").load("data/mllib/sample_kmeans_data.txt"); + + // Trains a k-means model. + KMeans kmeans = new KMeans().setK(2).setSeed(1L); + KMeansModel model = kmeans.fit(dataset); + + // Make predictions + Dataset predictions = model.transform(dataset); + + // Evaluate clustering by computing Silhouette score + ClusteringEvaluator evaluator = new ClusteringEvaluator(); + + double silhouette = evaluator.evaluate(predictions); + System.out.println("Silhouette with squared euclidean distance = " + silhouette); + + // Shows the result. + Vector[] centers = model.clusterCenters(); + System.out.println("Cluster Centers: "); + for (Vector center: centers) { + System.out.println(center); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java new file mode 100644 index 0000000..0e5d005 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java @@ -0,0 +1,69 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; +// $example on$ +import org.apache.spark.ml.clustering.LDA; +import org.apache.spark.ml.clustering.LDAModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +/** + * An example demonstrating LDA. + * Run with + *
+ * bin/run-example ml.JavaLDAExample
+ * 
+ */ +public class JavaLDAExample { + + public static void main(String[] args) { + // Creates a SparkSession + SparkSession spark = SparkSession + .builder() + .appName("JavaLDAExample") + .getOrCreate(); + + // $example on$ + // Loads data. + Dataset dataset = spark.read().format("libsvm") + .load("data/mllib/sample_lda_libsvm_data.txt"); + + // Trains a LDA model. + LDA lda = new LDA().setK(10).setMaxIter(10); + LDAModel model = lda.fit(dataset); + + double ll = model.logLikelihood(dataset); + double lp = model.logPerplexity(dataset); + System.out.println("The lower bound on the log likelihood of the entire corpus: " + ll); + System.out.println("The upper bound on perplexity: " + lp); + + // Describe topics. + Dataset topics = model.describeTopics(3); + System.out.println("The topics described by their top-weighted terms:"); + topics.show(false); + + // Shows the result. + Dataset transformed = model.transform(dataset); + transformed.show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLabeledDocument.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLabeledDocument.java new file mode 100644 index 0000000..68d1caf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLabeledDocument.java @@ -0,0 +1,38 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import java.io.Serializable; + +/** + * Labeled instance type, Spark SQL can infer schema from Java Beans. + */ +@SuppressWarnings("serial") +public class JavaLabeledDocument extends JavaDocument implements Serializable { + + private double label; + + public JavaLabeledDocument(long id, String text, double label) { + super(id, text); + this.label = label; + } + + public double getLabel() { + return this.label; + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java new file mode 100644 index 0000000..a561b6d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java @@ -0,0 +1,65 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.regression.LinearRegression; +import org.apache.spark.ml.regression.LinearRegressionModel; +import org.apache.spark.ml.regression.LinearRegressionTrainingSummary; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaLinearRegressionWithElasticNetExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaLinearRegressionWithElasticNetExample") + .getOrCreate(); + + // $example on$ + // Load training data. + Dataset training = spark.read().format("libsvm") + .load("data/mllib/sample_linear_regression_data.txt"); + + LinearRegression lr = new LinearRegression() + .setMaxIter(10) + .setRegParam(0.3) + .setElasticNetParam(0.8); + + // Fit the model. + LinearRegressionModel lrModel = lr.fit(training); + + // Print the coefficients and intercept for linear regression. + System.out.println("Coefficients: " + + lrModel.coefficients() + " Intercept: " + lrModel.intercept()); + + // Summarize the model over the training set and print out some metrics. + LinearRegressionTrainingSummary trainingSummary = lrModel.summary(); + System.out.println("numIterations: " + trainingSummary.totalIterations()); + System.out.println("objectiveHistory: " + Vectors.dense(trainingSummary.objectiveHistory())); + trainingSummary.residuals().show(); + System.out.println("RMSE: " + trainingSummary.rootMeanSquaredError()); + System.out.println("r2: " + trainingSummary.r2()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearSVCExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearSVCExample.java new file mode 100644 index 0000000..a18ed1d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearSVCExample.java @@ -0,0 +1,54 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.LinearSVC; +import org.apache.spark.ml.classification.LinearSVCModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaLinearSVCExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaLinearSVCExample") + .getOrCreate(); + + // $example on$ + // Load training data + Dataset training = spark.read().format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + LinearSVC lsvc = new LinearSVC() + .setMaxIter(10) + .setRegParam(0.1); + + // Fit the model + LinearSVCModel lsvcModel = lsvc.fit(training); + + // Print the coefficients and intercept for LinearSVC + System.out.println("Coefficients: " + + lsvcModel.coefficients() + " Intercept: " + lsvcModel.intercept()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java new file mode 100644 index 0000000..1529da1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java @@ -0,0 +1,77 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.BinaryLogisticRegressionTrainingSummary; +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.classification.LogisticRegressionModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.functions; +// $example off$ + +public class JavaLogisticRegressionSummaryExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaLogisticRegressionSummaryExample") + .getOrCreate(); + + // Load training data + Dataset training = spark.read().format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + LogisticRegression lr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.3) + .setElasticNetParam(0.8); + + // Fit the model + LogisticRegressionModel lrModel = lr.fit(training); + + // $example on$ + // Extract the summary from the returned LogisticRegressionModel instance trained in the earlier + // example + BinaryLogisticRegressionTrainingSummary trainingSummary = lrModel.binarySummary(); + + // Obtain the loss per iteration. + double[] objectiveHistory = trainingSummary.objectiveHistory(); + for (double lossPerIteration : objectiveHistory) { + System.out.println(lossPerIteration); + } + + // Obtain the receiver-operating characteristic as a dataframe and areaUnderROC. + Dataset roc = trainingSummary.roc(); + roc.show(); + roc.select("FPR").show(); + System.out.println(trainingSummary.areaUnderROC()); + + // Get the threshold corresponding to the maximum F-Measure and rerun LogisticRegression with + // this selected threshold. + Dataset fMeasure = trainingSummary.fMeasureByThreshold(); + double maxFMeasure = fMeasure.select(functions.max("F-Measure")).head().getDouble(0); + double bestThreshold = fMeasure.where(fMeasure.col("F-Measure").equalTo(maxFMeasure)) + .select("threshold").head().getDouble(0); + lrModel.setThreshold(bestThreshold); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java new file mode 100644 index 0000000..4cdec21 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java @@ -0,0 +1,69 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.classification.LogisticRegressionModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaLogisticRegressionWithElasticNetExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaLogisticRegressionWithElasticNetExample") + .getOrCreate(); + + // $example on$ + // Load training data + Dataset training = spark.read().format("libsvm") + .load("data/mllib/sample_libsvm_data.txt"); + + LogisticRegression lr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.3) + .setElasticNetParam(0.8); + + // Fit the model + LogisticRegressionModel lrModel = lr.fit(training); + + // Print the coefficients and intercept for logistic regression + System.out.println("Coefficients: " + + lrModel.coefficients() + " Intercept: " + lrModel.intercept()); + + // We can also use the multinomial family for binary classification + LogisticRegression mlr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.3) + .setElasticNetParam(0.8) + .setFamily("multinomial"); + + // Fit the model + LogisticRegressionModel mlrModel = mlr.fit(training); + + // Print the coefficients and intercepts for logistic regression with multinomial family + System.out.println("Multinomial coefficients: " + lrModel.coefficientMatrix() + + "\nMultinomial intercepts: " + mlrModel.interceptVector()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMaxAbsScalerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMaxAbsScalerExample.java new file mode 100644 index 0000000..9f1ce46 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMaxAbsScalerExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.MaxAbsScaler; +import org.apache.spark.ml.feature.MaxAbsScalerModel; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ +import org.apache.spark.sql.SparkSession; + +public class JavaMaxAbsScalerExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaMaxAbsScalerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, Vectors.dense(1.0, 0.1, -8.0)), + RowFactory.create(1, Vectors.dense(2.0, 1.0, -4.0)), + RowFactory.create(2, Vectors.dense(4.0, 10.0, 8.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset dataFrame = spark.createDataFrame(data, schema); + + MaxAbsScaler scaler = new MaxAbsScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures"); + + // Compute summary statistics and generate MaxAbsScalerModel + MaxAbsScalerModel scalerModel = scaler.fit(dataFrame); + + // rescale each feature to range [-1, 1]. + Dataset scaledData = scalerModel.transform(dataFrame); + scaledData.select("features", "scaledFeatures").show(); + // $example off$ + + spark.stop(); + } + +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinHashLSHExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinHashLSHExample.java new file mode 100644 index 0000000..e164598 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinHashLSHExample.java @@ -0,0 +1,111 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.MinHashLSH; +import org.apache.spark.ml.feature.MinHashLSHModel; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import static org.apache.spark.sql.functions.col; +// $example off$ + +/** + * An example demonstrating MinHashLSH. + * Run with: + * bin/run-example ml.JavaMinHashLSHExample + */ +public class JavaMinHashLSHExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaMinHashLSHExample") + .getOrCreate(); + + // $example on$ + List dataA = Arrays.asList( + RowFactory.create(0, Vectors.sparse(6, new int[]{0, 1, 2}, new double[]{1.0, 1.0, 1.0})), + RowFactory.create(1, Vectors.sparse(6, new int[]{2, 3, 4}, new double[]{1.0, 1.0, 1.0})), + RowFactory.create(2, Vectors.sparse(6, new int[]{0, 2, 4}, new double[]{1.0, 1.0, 1.0})) + ); + + List dataB = Arrays.asList( + RowFactory.create(0, Vectors.sparse(6, new int[]{1, 3, 5}, new double[]{1.0, 1.0, 1.0})), + RowFactory.create(1, Vectors.sparse(6, new int[]{2, 3, 5}, new double[]{1.0, 1.0, 1.0})), + RowFactory.create(2, Vectors.sparse(6, new int[]{1, 2, 4}, new double[]{1.0, 1.0, 1.0})) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset dfA = spark.createDataFrame(dataA, schema); + Dataset dfB = spark.createDataFrame(dataB, schema); + + int[] indices = {1, 3}; + double[] values = {1.0, 1.0}; + Vector key = Vectors.sparse(6, indices, values); + + MinHashLSH mh = new MinHashLSH() + .setNumHashTables(5) + .setInputCol("features") + .setOutputCol("hashes"); + + MinHashLSHModel model = mh.fit(dfA); + + // Feature Transformation + System.out.println("The hashed dataset where hashed values are stored in the column 'hashes':"); + model.transform(dfA).show(); + + // Compute the locality sensitive hashes for the input rows, then perform approximate + // similarity join. + // We could avoid computing hashes by passing in the already-transformed dataset, e.g. + // `model.approxSimilarityJoin(transformedA, transformedB, 0.6)` + System.out.println("Approximately joining dfA and dfB on Jaccard distance smaller than 0.6:"); + model.approxSimilarityJoin(dfA, dfB, 0.6, "JaccardDistance") + .select(col("datasetA.id").alias("idA"), + col("datasetB.id").alias("idB"), + col("JaccardDistance")).show(); + + // Compute the locality sensitive hashes for the input rows, then perform approximate nearest + // neighbor search. + // We could avoid computing hashes by passing in the already-transformed dataset, e.g. + // `model.approxNearestNeighbors(transformedA, key, 2)` + // It may return less than 2 rows when not enough approximate near-neighbor candidates are + // found. + System.out.println("Approximately searching dfA for 2 nearest neighbors of the key:"); + model.approxNearestNeighbors(dfA, key, 2).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinMaxScalerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinMaxScalerExample.java new file mode 100644 index 0000000..2757af8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMinMaxScalerExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.MinMaxScaler; +import org.apache.spark.ml.feature.MinMaxScalerModel; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaMinMaxScalerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaMinMaxScalerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, Vectors.dense(1.0, 0.1, -1.0)), + RowFactory.create(1, Vectors.dense(2.0, 1.1, 1.0)), + RowFactory.create(2, Vectors.dense(3.0, 10.1, 3.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset dataFrame = spark.createDataFrame(data, schema); + + MinMaxScaler scaler = new MinMaxScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures"); + + // Compute summary statistics and generate MinMaxScalerModel + MinMaxScalerModel scalerModel = scaler.fit(dataFrame); + + // rescale each feature to range [min, max]. + Dataset scaledData = scalerModel.transform(dataFrame); + System.out.println("Features scaled to range: [" + scaler.getMin() + ", " + + scaler.getMax() + "]"); + scaledData.select("features", "scaledFeatures").show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaCrossValidationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaCrossValidationExample.java new file mode 100644 index 0000000..d973279 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaCrossValidationExample.java @@ -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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +// $example off$ + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator; +import org.apache.spark.ml.feature.HashingTF; +import org.apache.spark.ml.feature.Tokenizer; +import org.apache.spark.ml.param.ParamMap; +import org.apache.spark.ml.tuning.CrossValidator; +import org.apache.spark.ml.tuning.CrossValidatorModel; +import org.apache.spark.ml.tuning.ParamGridBuilder; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * Java example for Model Selection via Cross Validation. + */ +public class JavaModelSelectionViaCrossValidationExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaModelSelectionViaCrossValidationExample") + .getOrCreate(); + + // $example on$ + // Prepare training documents, which are labeled. + Dataset training = spark.createDataFrame(Arrays.asList( + new JavaLabeledDocument(0L, "a b c d e spark", 1.0), + new JavaLabeledDocument(1L, "b d", 0.0), + new JavaLabeledDocument(2L,"spark f g h", 1.0), + new JavaLabeledDocument(3L, "hadoop mapreduce", 0.0), + new JavaLabeledDocument(4L, "b spark who", 1.0), + new JavaLabeledDocument(5L, "g d a y", 0.0), + new JavaLabeledDocument(6L, "spark fly", 1.0), + new JavaLabeledDocument(7L, "was mapreduce", 0.0), + new JavaLabeledDocument(8L, "e spark program", 1.0), + new JavaLabeledDocument(9L, "a e c l", 0.0), + new JavaLabeledDocument(10L, "spark compile", 1.0), + new JavaLabeledDocument(11L, "hadoop software", 0.0) + ), JavaLabeledDocument.class); + + // Configure an ML pipeline, which consists of three stages: tokenizer, hashingTF, and lr. + Tokenizer tokenizer = new Tokenizer() + .setInputCol("text") + .setOutputCol("words"); + HashingTF hashingTF = new HashingTF() + .setNumFeatures(1000) + .setInputCol(tokenizer.getOutputCol()) + .setOutputCol("features"); + LogisticRegression lr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.01); + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {tokenizer, hashingTF, lr}); + + // We use a ParamGridBuilder to construct a grid of parameters to search over. + // With 3 values for hashingTF.numFeatures and 2 values for lr.regParam, + // this grid will have 3 x 2 = 6 parameter settings for CrossValidator to choose from. + ParamMap[] paramGrid = new ParamGridBuilder() + .addGrid(hashingTF.numFeatures(), new int[] {10, 100, 1000}) + .addGrid(lr.regParam(), new double[] {0.1, 0.01}) + .build(); + + // We now treat the Pipeline as an Estimator, wrapping it in a CrossValidator instance. + // This will allow us to jointly choose parameters for all Pipeline stages. + // A CrossValidator requires an Estimator, a set of Estimator ParamMaps, and an Evaluator. + // Note that the evaluator here is a BinaryClassificationEvaluator and its default metric + // is areaUnderROC. + CrossValidator cv = new CrossValidator() + .setEstimator(pipeline) + .setEvaluator(new BinaryClassificationEvaluator()) + .setEstimatorParamMaps(paramGrid) + .setNumFolds(2) // Use 3+ in practice + .setParallelism(2); // Evaluate up to 2 parameter settings in parallel + + // Run cross-validation, and choose the best set of parameters. + CrossValidatorModel cvModel = cv.fit(training); + + // Prepare test documents, which are unlabeled. + Dataset test = spark.createDataFrame(Arrays.asList( + new JavaDocument(4L, "spark i j k"), + new JavaDocument(5L, "l m n"), + new JavaDocument(6L, "mapreduce spark"), + new JavaDocument(7L, "apache hadoop") + ), JavaDocument.class); + + // Make predictions on test documents. cvModel uses the best model found (lrModel). + Dataset predictions = cvModel.transform(test); + for (Row r : predictions.select("id", "text", "probability", "prediction").collectAsList()) { + System.out.println("(" + r.get(0) + ", " + r.get(1) + ") --> prob=" + r.get(2) + + ", prediction=" + r.get(3)); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaTrainValidationSplitExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaTrainValidationSplitExample.java new file mode 100644 index 0000000..2ef8bea --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaModelSelectionViaTrainValidationSplitExample.java @@ -0,0 +1,88 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.param.ParamMap; +import org.apache.spark.ml.regression.LinearRegression; +import org.apache.spark.ml.tuning.ParamGridBuilder; +import org.apache.spark.ml.tuning.TrainValidationSplit; +import org.apache.spark.ml.tuning.TrainValidationSplitModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * Java example demonstrating model selection using TrainValidationSplit. + * + * Run with + * {{{ + * bin/run-example ml.JavaModelSelectionViaTrainValidationSplitExample + * }}} + */ +public class JavaModelSelectionViaTrainValidationSplitExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaModelSelectionViaTrainValidationSplitExample") + .getOrCreate(); + + // $example on$ + Dataset data = spark.read().format("libsvm") + .load("data/mllib/sample_linear_regression_data.txt"); + + // Prepare training and test data. + Dataset[] splits = data.randomSplit(new double[] {0.9, 0.1}, 12345); + Dataset training = splits[0]; + Dataset test = splits[1]; + + LinearRegression lr = new LinearRegression(); + + // We use a ParamGridBuilder to construct a grid of parameters to search over. + // TrainValidationSplit will try all combinations of values and determine best model using + // the evaluator. + ParamMap[] paramGrid = new ParamGridBuilder() + .addGrid(lr.regParam(), new double[] {0.1, 0.01}) + .addGrid(lr.fitIntercept()) + .addGrid(lr.elasticNetParam(), new double[] {0.0, 0.5, 1.0}) + .build(); + + // In this case the estimator is simply the linear regression. + // A TrainValidationSplit requires an Estimator, a set of Estimator ParamMaps, and an Evaluator. + TrainValidationSplit trainValidationSplit = new TrainValidationSplit() + .setEstimator(lr) + .setEvaluator(new RegressionEvaluator()) + .setEstimatorParamMaps(paramGrid) + .setTrainRatio(0.8) // 80% for training and the remaining 20% for validation + .setParallelism(2); // Evaluate up to 2 parameter settings in parallel + + // Run train validation split, and choose the best set of parameters. + TrainValidationSplitModel model = trainValidationSplit.fit(training); + + // Make predictions on test data. model is the model with combination of parameters + // that performed best. + model.transform(test) + .select("features", "label", "prediction") + .show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMulticlassLogisticRegressionWithElasticNetExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMulticlassLogisticRegressionWithElasticNetExample.java new file mode 100644 index 0000000..801a82c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMulticlassLogisticRegressionWithElasticNetExample.java @@ -0,0 +1,117 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.classification.LogisticRegressionModel; +import org.apache.spark.ml.classification.LogisticRegressionTrainingSummary; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaMulticlassLogisticRegressionWithElasticNetExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaMulticlassLogisticRegressionWithElasticNetExample") + .getOrCreate(); + + // $example on$ + // Load training data + Dataset training = spark.read().format("libsvm") + .load("data/mllib/sample_multiclass_classification_data.txt"); + + LogisticRegression lr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.3) + .setElasticNetParam(0.8); + + // Fit the model + LogisticRegressionModel lrModel = lr.fit(training); + + // Print the coefficients and intercept for multinomial logistic regression + System.out.println("Coefficients: \n" + + lrModel.coefficientMatrix() + " \nIntercept: " + lrModel.interceptVector()); + LogisticRegressionTrainingSummary trainingSummary = lrModel.summary(); + + // Obtain the loss per iteration. + double[] objectiveHistory = trainingSummary.objectiveHistory(); + for (double lossPerIteration : objectiveHistory) { + System.out.println(lossPerIteration); + } + + // for multiclass, we can inspect metrics on a per-label basis + System.out.println("False positive rate by label:"); + int i = 0; + double[] fprLabel = trainingSummary.falsePositiveRateByLabel(); + for (double fpr : fprLabel) { + System.out.println("label " + i + ": " + fpr); + i++; + } + + System.out.println("True positive rate by label:"); + i = 0; + double[] tprLabel = trainingSummary.truePositiveRateByLabel(); + for (double tpr : tprLabel) { + System.out.println("label " + i + ": " + tpr); + i++; + } + + System.out.println("Precision by label:"); + i = 0; + double[] precLabel = trainingSummary.precisionByLabel(); + for (double prec : precLabel) { + System.out.println("label " + i + ": " + prec); + i++; + } + + System.out.println("Recall by label:"); + i = 0; + double[] recLabel = trainingSummary.recallByLabel(); + for (double rec : recLabel) { + System.out.println("label " + i + ": " + rec); + i++; + } + + System.out.println("F-measure by label:"); + i = 0; + double[] fLabel = trainingSummary.fMeasureByLabel(); + for (double f : fLabel) { + System.out.println("label " + i + ": " + f); + i++; + } + + double accuracy = trainingSummary.accuracy(); + double falsePositiveRate = trainingSummary.weightedFalsePositiveRate(); + double truePositiveRate = trainingSummary.weightedTruePositiveRate(); + double fMeasure = trainingSummary.weightedFMeasure(); + double precision = trainingSummary.weightedPrecision(); + double recall = trainingSummary.weightedRecall(); + System.out.println("Accuracy: " + accuracy); + System.out.println("FPR: " + falsePositiveRate); + System.out.println("TPR: " + truePositiveRate); + System.out.println("F-measure: " + fMeasure); + System.out.println("Precision: " + precision); + System.out.println("Recall: " + recall); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java new file mode 100644 index 0000000..43db41c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java @@ -0,0 +1,76 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel; +import org.apache.spark.ml.classification.MultilayerPerceptronClassifier; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +// $example off$ + +/** + * An example for Multilayer Perceptron Classification. + */ +public class JavaMultilayerPerceptronClassifierExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaMultilayerPerceptronClassifierExample") + .getOrCreate(); + + // $example on$ + // Load training data + String path = "data/mllib/sample_multiclass_classification_data.txt"; + Dataset dataFrame = spark.read().format("libsvm").load(path); + + // Split the data into train and test + Dataset[] splits = dataFrame.randomSplit(new double[]{0.6, 0.4}, 1234L); + Dataset train = splits[0]; + Dataset test = splits[1]; + + // specify layers for the neural network: + // input layer of size 4 (features), two intermediate of size 5 and 4 + // and output of size 3 (classes) + int[] layers = new int[] {4, 5, 4, 3}; + + // create the trainer and set its parameters + MultilayerPerceptronClassifier trainer = new MultilayerPerceptronClassifier() + .setLayers(layers) + .setBlockSize(128) + .setSeed(1234L) + .setMaxIter(100); + + // train the model + MultilayerPerceptronClassificationModel model = trainer.fit(train); + + // compute accuracy on the test set + Dataset result = model.transform(test); + Dataset predictionAndLabels = result.select("prediction", "label"); + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setMetricName("accuracy"); + + System.out.println("Test set accuracy = " + evaluator.evaluate(predictionAndLabels)); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNGramExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNGramExample.java new file mode 100644 index 0000000..5427e46 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNGramExample.java @@ -0,0 +1,66 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.NGram; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaNGramExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaNGramExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, Arrays.asList("Hi", "I", "heard", "about", "Spark")), + RowFactory.create(1, Arrays.asList("I", "wish", "Java", "could", "use", "case", "classes")), + RowFactory.create(2, Arrays.asList("Logistic", "regression", "models", "are", "neat")) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField( + "words", DataTypes.createArrayType(DataTypes.StringType), false, Metadata.empty()) + }); + + Dataset wordDataFrame = spark.createDataFrame(data, schema); + + NGram ngramTransformer = new NGram().setN(2).setInputCol("words").setOutputCol("ngrams"); + + Dataset ngramDataFrame = ngramTransformer.transform(wordDataFrame); + ngramDataFrame.select("ngrams").show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNaiveBayesExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNaiveBayesExample.java new file mode 100644 index 0000000..be578dc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNaiveBayesExample.java @@ -0,0 +1,70 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.NaiveBayes; +import org.apache.spark.ml.classification.NaiveBayesModel; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +/** + * An example for Naive Bayes Classification. + */ +public class JavaNaiveBayesExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaNaiveBayesExample") + .getOrCreate(); + + // $example on$ + // Load training data + Dataset dataFrame = + spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + // Split the data into train and test + Dataset[] splits = dataFrame.randomSplit(new double[]{0.6, 0.4}, 1234L); + Dataset train = splits[0]; + Dataset test = splits[1]; + + // create the trainer and set its parameters + NaiveBayes nb = new NaiveBayes(); + + // train the model + NaiveBayesModel model = nb.fit(train); + + // Select example rows to display. + Dataset predictions = model.transform(test); + predictions.show(); + + // compute accuracy on the test set + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setLabelCol("label") + .setPredictionCol("prediction") + .setMetricName("accuracy"); + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test set accuracy = " + accuracy); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNormalizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNormalizerExample.java new file mode 100644 index 0000000..f878c42 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaNormalizerExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.Normalizer; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaNormalizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaNormalizerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, Vectors.dense(1.0, 0.1, -8.0)), + RowFactory.create(1, Vectors.dense(2.0, 1.0, -4.0)), + RowFactory.create(2, Vectors.dense(4.0, 10.0, 8.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + Dataset dataFrame = spark.createDataFrame(data, schema); + + // Normalize each Vector using $L^1$ norm. + Normalizer normalizer = new Normalizer() + .setInputCol("features") + .setOutputCol("normFeatures") + .setP(1.0); + + Dataset l1NormData = normalizer.transform(dataFrame); + l1NormData.show(); + + // Normalize each Vector using $L^\infty$ norm. + Dataset lInfNormData = + normalizer.transform(dataFrame, normalizer.p().w(Double.POSITIVE_INFINITY)); + lInfNormData.show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java new file mode 100644 index 0000000..4b49beb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.OneHotEncoder; +import org.apache.spark.ml.feature.OneHotEncoderModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaOneHotEncoderExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaOneHotEncoderExample") + .getOrCreate(); + + // Note: categorical features are usually first encoded with StringIndexer + // $example on$ + List data = Arrays.asList( + RowFactory.create(0.0, 1.0), + RowFactory.create(1.0, 0.0), + RowFactory.create(2.0, 1.0), + RowFactory.create(0.0, 2.0), + RowFactory.create(0.0, 1.0), + RowFactory.create(2.0, 0.0) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("categoryIndex1", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("categoryIndex2", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + OneHotEncoder encoder = new OneHotEncoder() + .setInputCols(new String[] {"categoryIndex1", "categoryIndex2"}) + .setOutputCols(new String[] {"categoryVec1", "categoryVec2"}); + + OneHotEncoderModel model = encoder.fit(df); + Dataset encoded = model.transform(df); + encoded.show(); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java new file mode 100644 index 0000000..82fb540 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java @@ -0,0 +1,84 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.classification.OneVsRest; +import org.apache.spark.ml.classification.OneVsRestModel; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + + +/** + * An example of Multiclass to Binary Reduction with One Vs Rest, + * using Logistic Regression as the base classifier. + * Run with + *
+ * bin/run-example ml.JavaOneVsRestExample
+ * 
+ */ +public class JavaOneVsRestExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaOneVsRestExample") + .getOrCreate(); + + // $example on$ + // load data file. + Dataset inputData = spark.read().format("libsvm") + .load("data/mllib/sample_multiclass_classification_data.txt"); + + // generate the train/test split. + Dataset[] tmp = inputData.randomSplit(new double[]{0.8, 0.2}); + Dataset train = tmp[0]; + Dataset test = tmp[1]; + + // configure the base classifier. + LogisticRegression classifier = new LogisticRegression() + .setMaxIter(10) + .setTol(1E-6) + .setFitIntercept(true); + + // instantiate the One Vs Rest Classifier. + OneVsRest ovr = new OneVsRest().setClassifier(classifier); + + // train the multiclass model. + OneVsRestModel ovrModel = ovr.fit(train); + + // score the model on test data. + Dataset predictions = ovrModel.transform(test) + .select("prediction", "label"); + + // obtain evaluator. + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setMetricName("accuracy"); + + // compute the classification error on test data. + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test Error = " + (1 - accuracy)); + // $example off$ + + spark.stop(); + } + +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java new file mode 100644 index 0000000..6951a65 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java @@ -0,0 +1,70 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.PCA; +import org.apache.spark.ml.feature.PCAModel; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaPCAExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaPCAExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Vectors.sparse(5, new int[]{1, 3}, new double[]{1.0, 7.0})), + RowFactory.create(Vectors.dense(2.0, 0.0, 3.0, 4.0, 5.0)), + RowFactory.create(Vectors.dense(4.0, 0.0, 0.0, 6.0, 7.0)) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("features", new VectorUDT(), false, Metadata.empty()), + }); + + Dataset df = spark.createDataFrame(data, schema); + + PCAModel pca = new PCA() + .setInputCol("features") + .setOutputCol("pcaFeatures") + .setK(3) + .fit(df); + + Dataset result = pca.transform(df).select("pcaFeatures"); + result.show(false); + // $example off$ + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPipelineExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPipelineExample.java new file mode 100644 index 0000000..4ccd8f6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPipelineExample.java @@ -0,0 +1,88 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.LogisticRegression; +import org.apache.spark.ml.feature.HashingTF; +import org.apache.spark.ml.feature.Tokenizer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ +import org.apache.spark.sql.SparkSession; + +/** + * Java example for simple text document 'Pipeline'. + */ +public class JavaPipelineExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaPipelineExample") + .getOrCreate(); + + // $example on$ + // Prepare training documents, which are labeled. + Dataset training = spark.createDataFrame(Arrays.asList( + new JavaLabeledDocument(0L, "a b c d e spark", 1.0), + new JavaLabeledDocument(1L, "b d", 0.0), + new JavaLabeledDocument(2L, "spark f g h", 1.0), + new JavaLabeledDocument(3L, "hadoop mapreduce", 0.0) + ), JavaLabeledDocument.class); + + // Configure an ML pipeline, which consists of three stages: tokenizer, hashingTF, and lr. + Tokenizer tokenizer = new Tokenizer() + .setInputCol("text") + .setOutputCol("words"); + HashingTF hashingTF = new HashingTF() + .setNumFeatures(1000) + .setInputCol(tokenizer.getOutputCol()) + .setOutputCol("features"); + LogisticRegression lr = new LogisticRegression() + .setMaxIter(10) + .setRegParam(0.001); + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {tokenizer, hashingTF, lr}); + + // Fit the pipeline to training documents. + PipelineModel model = pipeline.fit(training); + + // Prepare test documents, which are unlabeled. + Dataset test = spark.createDataFrame(Arrays.asList( + new JavaDocument(4L, "spark i j k"), + new JavaDocument(5L, "l m n"), + new JavaDocument(6L, "spark hadoop spark"), + new JavaDocument(7L, "apache hadoop") + ), JavaDocument.class); + + // Make predictions on test documents. + Dataset predictions = model.transform(test); + for (Row r : predictions.select("id", "text", "probability", "prediction").collectAsList()) { + System.out.println("(" + r.get(0) + ", " + r.get(1) + ") --> prob=" + r.get(2) + + ", prediction=" + r.get(3)); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPolynomialExpansionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPolynomialExpansionExample.java new file mode 100644 index 0000000..43c636c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPolynomialExpansionExample.java @@ -0,0 +1,66 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.PolynomialExpansion; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaPolynomialExpansionExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaPolynomialExpansionExample") + .getOrCreate(); + + // $example on$ + PolynomialExpansion polyExpansion = new PolynomialExpansion() + .setInputCol("features") + .setOutputCol("polyFeatures") + .setDegree(3); + + List data = Arrays.asList( + RowFactory.create(Vectors.dense(2.0, 1.0)), + RowFactory.create(Vectors.dense(0.0, 0.0)), + RowFactory.create(Vectors.dense(3.0, -1.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("features", new VectorUDT(), false, Metadata.empty()), + }); + Dataset df = spark.createDataFrame(data, schema); + + Dataset polyDF = polyExpansion.transform(df); + polyDF.show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPowerIterationClusteringExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPowerIterationClusteringExample.java new file mode 100644 index 0000000..5186563 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPowerIterationClusteringExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.clustering.PowerIterationClustering; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaPowerIterationClusteringExample { + public static void main(String[] args) { + // Create a SparkSession. + SparkSession spark = SparkSession + .builder() + .appName("JavaPowerIterationClustering") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0L, 1L, 1.0), + RowFactory.create(0L, 2L, 1.0), + RowFactory.create(1L, 2L, 1.0), + RowFactory.create(3L, 4L, 1.0), + RowFactory.create(4L, 0L, 0.1) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("src", DataTypes.LongType, false, Metadata.empty()), + new StructField("dst", DataTypes.LongType, false, Metadata.empty()), + new StructField("weight", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + PowerIterationClustering model = new PowerIterationClustering() + .setK(2) + .setMaxIter(10) + .setInitMode("degree") + .setWeightCol("weight"); + + Dataset result = model.assignClusters(df); + result.show(false); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPrefixSpanExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPrefixSpanExample.java new file mode 100644 index 0000000..891f306 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaPrefixSpanExample.java @@ -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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.fpm.PrefixSpan; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example demonstrating PrefixSpan. + * Run with + *
+ * bin/run-example ml.JavaPrefixSpanExample
+ * 
+ */ +public class JavaPrefixSpanExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaPrefixSpanExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3))), + RowFactory.create(Arrays.asList(Arrays.asList(1), Arrays.asList(3, 2), Arrays.asList(1,2))), + RowFactory.create(Arrays.asList(Arrays.asList(1, 2), Arrays.asList(5))), + RowFactory.create(Arrays.asList(Arrays.asList(6))) + ); + StructType schema = new StructType(new StructField[]{ new StructField( + "sequence", new ArrayType(new ArrayType(DataTypes.IntegerType, true), true), + false, Metadata.empty()) + }); + Dataset sequenceDF = spark.createDataFrame(data, schema); + + PrefixSpan prefixSpan = new PrefixSpan().setMinSupport(0.5).setMaxPatternLength(5); + + // Finding frequent sequential patterns + prefixSpan.findFrequentSequentialPatterns(sequenceDF).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaQuantileDiscretizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaQuantileDiscretizerExample.java new file mode 100644 index 0000000..43cc30c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaQuantileDiscretizerExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.QuantileDiscretizer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaQuantileDiscretizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaQuantileDiscretizerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, 18.0), + RowFactory.create(1, 19.0), + RowFactory.create(2, 8.0), + RowFactory.create(3, 5.0), + RowFactory.create(4, 2.2) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("hour", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + // $example off$ + // Output of QuantileDiscretizer for such small datasets can depend on the number of + // partitions. Here we force a single partition to ensure consistent results. + // Note this is not necessary for normal use cases + df = df.repartition(1); + // $example on$ + QuantileDiscretizer discretizer = new QuantileDiscretizer() + .setInputCol("hour") + .setOutputCol("result") + .setNumBuckets(3); + + Dataset result = discretizer.fit(df).transform(df); + result.show(false); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRFormulaExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRFormulaExample.java new file mode 100644 index 0000000..428067e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRFormulaExample.java @@ -0,0 +1,68 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.RFormula; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import static org.apache.spark.sql.types.DataTypes.*; +// $example off$ + +public class JavaRFormulaExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaRFormulaExample") + .getOrCreate(); + + // $example on$ + StructType schema = createStructType(new StructField[]{ + createStructField("id", IntegerType, false), + createStructField("country", StringType, false), + createStructField("hour", IntegerType, false), + createStructField("clicked", DoubleType, false) + }); + + List data = Arrays.asList( + RowFactory.create(7, "US", 18, 1.0), + RowFactory.create(8, "CA", 12, 0.0), + RowFactory.create(9, "NZ", 15, 0.0) + ); + + Dataset dataset = spark.createDataFrame(data, schema); + RFormula formula = new RFormula() + .setFormula("clicked ~ country + hour") + .setFeaturesCol("features") + .setLabelCol("label"); + Dataset output = formula.fit(dataset).transform(dataset); + output.select("features", "label").show(); + // $example off$ + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java new file mode 100644 index 0000000..8863f18 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java @@ -0,0 +1,101 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.classification.RandomForestClassificationModel; +import org.apache.spark.ml.classification.RandomForestClassifier; +import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator; +import org.apache.spark.ml.feature.*; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaRandomForestClassifierExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaRandomForestClassifierExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + // Index labels, adding metadata to the label column. + // Fit on whole dataset to include all labels in index. + StringIndexerModel labelIndexer = new StringIndexer() + .setInputCol("label") + .setOutputCol("indexedLabel") + .fit(data); + // Automatically identify categorical features, and index them. + // Set maxCategories so features with > 4 distinct values are treated as continuous. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) + .fit(data); + + // Split the data into training and test sets (30% held out for testing) + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a RandomForest model. + RandomForestClassifier rf = new RandomForestClassifier() + .setLabelCol("indexedLabel") + .setFeaturesCol("indexedFeatures"); + + // Convert indexed labels back to original labels. + IndexToString labelConverter = new IndexToString() + .setInputCol("prediction") + .setOutputCol("predictedLabel") + .setLabels(labelIndexer.labelsArray()[0]); + + // Chain indexers and forest in a Pipeline + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {labelIndexer, featureIndexer, rf, labelConverter}); + + // Train model. This also runs the indexers. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("predictedLabel", "label", "features").show(5); + + // Select (prediction, true label) and compute test error + MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator() + .setLabelCol("indexedLabel") + .setPredictionCol("prediction") + .setMetricName("accuracy"); + double accuracy = evaluator.evaluate(predictions); + System.out.println("Test Error = " + (1.0 - accuracy)); + + RandomForestClassificationModel rfModel = (RandomForestClassificationModel)(model.stages()[2]); + System.out.println("Learned classification forest model:\n" + rfModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java new file mode 100644 index 0000000..a707845 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java @@ -0,0 +1,90 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import org.apache.spark.ml.Pipeline; +import org.apache.spark.ml.PipelineModel; +import org.apache.spark.ml.PipelineStage; +import org.apache.spark.ml.evaluation.RegressionEvaluator; +import org.apache.spark.ml.feature.VectorIndexer; +import org.apache.spark.ml.feature.VectorIndexerModel; +import org.apache.spark.ml.regression.RandomForestRegressionModel; +import org.apache.spark.ml.regression.RandomForestRegressor; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off$ + +public class JavaRandomForestRegressorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaRandomForestRegressorExample") + .getOrCreate(); + + // $example on$ + // Load and parse the data file, converting it to a DataFrame. + Dataset data = spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + // Automatically identify categorical features, and index them. + // Set maxCategories so features with > 4 distinct values are treated as continuous. + VectorIndexerModel featureIndexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexedFeatures") + .setMaxCategories(4) + .fit(data); + + // Split the data into training and test sets (30% held out for testing) + Dataset[] splits = data.randomSplit(new double[] {0.7, 0.3}); + Dataset trainingData = splits[0]; + Dataset testData = splits[1]; + + // Train a RandomForest model. + RandomForestRegressor rf = new RandomForestRegressor() + .setLabelCol("label") + .setFeaturesCol("indexedFeatures"); + + // Chain indexer and forest in a Pipeline + Pipeline pipeline = new Pipeline() + .setStages(new PipelineStage[] {featureIndexer, rf}); + + // Train model. This also runs the indexer. + PipelineModel model = pipeline.fit(trainingData); + + // Make predictions. + Dataset predictions = model.transform(testData); + + // Select example rows to display. + predictions.select("prediction", "label", "features").show(5); + + // Select (prediction, true label) and compute test error + RegressionEvaluator evaluator = new RegressionEvaluator() + .setLabelCol("label") + .setPredictionCol("prediction") + .setMetricName("rmse"); + double rmse = evaluator.evaluate(predictions); + System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse); + + RandomForestRegressionModel rfModel = (RandomForestRegressionModel)(model.stages()[1]); + System.out.println("Learned regression forest model:\n" + rfModel.toDebugString()); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRobustScalerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRobustScalerExample.java new file mode 100644 index 0000000..475d046 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaRobustScalerExample.java @@ -0,0 +1,57 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import org.apache.spark.ml.feature.RobustScaler; +import org.apache.spark.ml.feature.RobustScalerModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ + +public class JavaRobustScalerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaRobustScalerExample") + .getOrCreate(); + + // $example on$ + Dataset dataFrame = + spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + RobustScaler scaler = new RobustScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures") + .setWithScaling(true) + .setWithCentering(false) + .setLower(0.25) + .setUpper(0.75); + + // Compute summary statistics by fitting the RobustScaler + RobustScalerModel scalerModel = scaler.fit(dataFrame); + + // Transform each feature to have unit quantile range. + Dataset scaledData = scalerModel.transform(dataFrame); + scaledData.show(); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSQLTransformerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSQLTransformerExample.java new file mode 100644 index 0000000..2a3d62d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSQLTransformerExample.java @@ -0,0 +1,59 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.SQLTransformer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +public class JavaSQLTransformerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaSQLTransformerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, 1.0, 3.0), + RowFactory.create(2, 2.0, 5.0) + ); + StructType schema = new StructType(new StructField [] { + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("v1", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("v2", DataTypes.DoubleType, false, Metadata.empty()) + }); + Dataset df = spark.createDataFrame(data, schema); + + SQLTransformer sqlTrans = new SQLTransformer().setStatement( + "SELECT *, (v1 + v2) AS v3, (v1 * v2) AS v4 FROM __THIS__"); + + sqlTrans.transform(df).show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStandardScalerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStandardScalerExample.java new file mode 100644 index 0000000..08ea285 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStandardScalerExample.java @@ -0,0 +1,55 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import org.apache.spark.ml.feature.StandardScaler; +import org.apache.spark.ml.feature.StandardScalerModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ + +public class JavaStandardScalerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaStandardScalerExample") + .getOrCreate(); + + // $example on$ + Dataset dataFrame = + spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + StandardScaler scaler = new StandardScaler() + .setInputCol("features") + .setOutputCol("scaledFeatures") + .setWithStd(true) + .setWithMean(false); + + // Compute summary statistics by fitting the StandardScaler + StandardScalerModel scalerModel = scaler.fit(dataFrame); + + // Normalize each feature to have unit standard deviation. + Dataset scaledData = scalerModel.transform(dataFrame); + scaledData.show(); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStopWordsRemoverExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStopWordsRemoverExample.java new file mode 100644 index 0000000..94ead62 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStopWordsRemoverExample.java @@ -0,0 +1,64 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.StopWordsRemover; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaStopWordsRemoverExample { + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaStopWordsRemoverExample") + .getOrCreate(); + + // $example on$ + StopWordsRemover remover = new StopWordsRemover() + .setInputCol("raw") + .setOutputCol("filtered"); + + List data = Arrays.asList( + RowFactory.create(Arrays.asList("I", "saw", "the", "red", "balloon")), + RowFactory.create(Arrays.asList("Mary", "had", "a", "little", "lamb")) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField( + "raw", DataTypes.createArrayType(DataTypes.StringType), false, Metadata.empty()) + }); + + Dataset dataset = spark.createDataFrame(data, schema); + remover.transform(dataset).show(false); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStringIndexerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStringIndexerExample.java new file mode 100644 index 0000000..cf9747a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaStringIndexerExample.java @@ -0,0 +1,68 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.StringIndexer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import static org.apache.spark.sql.types.DataTypes.*; +// $example off$ + +public class JavaStringIndexerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaStringIndexerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, "a"), + RowFactory.create(1, "b"), + RowFactory.create(2, "c"), + RowFactory.create(3, "a"), + RowFactory.create(4, "a"), + RowFactory.create(5, "c") + ); + StructType schema = new StructType(new StructField[]{ + createStructField("id", IntegerType, false), + createStructField("category", StringType, false) + }); + Dataset df = spark.createDataFrame(data, schema); + + StringIndexer indexer = new StringIndexer() + .setInputCol("category") + .setOutputCol("categoryIndex"); + + Dataset indexed = indexer.fit(df).transform(df); + indexed.show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSummarizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSummarizerExample.java new file mode 100644 index 0000000..e9b8436 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaSummarizerExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.*; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.stat.Summarizer; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaSummarizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaSummarizerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(Vectors.dense(2.0, 3.0, 5.0), 1.0), + RowFactory.create(Vectors.dense(4.0, 6.0, 7.0), 2.0) + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("features", new VectorUDT(), false, Metadata.empty()), + new StructField("weight", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + Row result1 = df.select(Summarizer.metrics("mean", "variance") + .summary(new Column("features"), new Column("weight")).as("summary")) + .select("summary.mean", "summary.variance").first(); + System.out.println("with weight: mean = " + result1.getAs(0).toString() + + ", variance = " + result1.getAs(1).toString()); + + Row result2 = df.select( + Summarizer.mean(new Column("features")), + Summarizer.variance(new Column("features")) + ).first(); + System.out.println("without weight: mean = " + result2.getAs(0).toString() + + ", variance = " + result2.getAs(1).toString()); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTfIdfExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTfIdfExample.java new file mode 100644 index 0000000..b740cd0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTfIdfExample.java @@ -0,0 +1,78 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.HashingTF; +import org.apache.spark.ml.feature.IDF; +import org.apache.spark.ml.feature.IDFModel; +import org.apache.spark.ml.feature.Tokenizer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off$ + +public class JavaTfIdfExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaTfIdfExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0.0, "Hi I heard about Spark"), + RowFactory.create(0.0, "I wish Java could use case classes"), + RowFactory.create(1.0, "Logistic regression models are neat") + ); + StructType schema = new StructType(new StructField[]{ + new StructField("label", DataTypes.DoubleType, false, Metadata.empty()), + new StructField("sentence", DataTypes.StringType, false, Metadata.empty()) + }); + Dataset sentenceData = spark.createDataFrame(data, schema); + + Tokenizer tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words"); + Dataset wordsData = tokenizer.transform(sentenceData); + + int numFeatures = 20; + HashingTF hashingTF = new HashingTF() + .setInputCol("words") + .setOutputCol("rawFeatures") + .setNumFeatures(numFeatures); + + Dataset featurizedData = hashingTF.transform(wordsData); + // alternatively, CountVectorizer can also be used to get term frequency vectors + + IDF idf = new IDF().setInputCol("rawFeatures").setOutputCol("features"); + IDFModel idfModel = idf.fit(featurizedData); + + Dataset rescaledData = idfModel.transform(featurizedData); + rescaledData.select("label", "features").show(); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java new file mode 100644 index 0000000..091f7d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java @@ -0,0 +1,87 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import scala.collection.mutable.Seq; + +import org.apache.spark.ml.feature.RegexTokenizer; +import org.apache.spark.ml.feature.Tokenizer; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +// col("...") is preferable to df.col("...") +import static org.apache.spark.sql.functions.call_udf; +import static org.apache.spark.sql.functions.col; +// $example off$ + +public class JavaTokenizerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaTokenizerExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(0, "Hi I heard about Spark"), + RowFactory.create(1, "I wish Java could use case classes"), + RowFactory.create(2, "Logistic,regression,models,are,neat") + ); + + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("sentence", DataTypes.StringType, false, Metadata.empty()) + }); + + Dataset sentenceDataFrame = spark.createDataFrame(data, schema); + + Tokenizer tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words"); + + RegexTokenizer regexTokenizer = new RegexTokenizer() + .setInputCol("sentence") + .setOutputCol("words") + .setPattern("\\W"); // alternatively .setPattern("\\w+").setGaps(false); + + spark.udf().register( + "countTokens", (Seq words) -> words.size(), DataTypes.IntegerType); + + Dataset tokenized = tokenizer.transform(sentenceDataFrame); + tokenized.select("sentence", "words") + .withColumn("tokens", call_udf("countTokens", col("words"))) + .show(false); + + Dataset regexTokenized = regexTokenizer.transform(sentenceDataFrame); + regexTokenized.select("sentence", "words") + .withColumn("tokens", call_udf("countTokens", col("words"))) + .show(false); + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaUnivariateFeatureSelectorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaUnivariateFeatureSelectorExample.java new file mode 100644 index 0000000..748262f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaUnivariateFeatureSelectorExample.java @@ -0,0 +1,84 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.UnivariateFeatureSelector; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example for UnivariateFeatureSelector. + * Run with + *
+ * bin/run-example ml.JavaUnivariateFeatureSelectorExample
+ * 
+ */ +public class JavaUnivariateFeatureSelectorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaUnivariateFeatureSelectorExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(1, Vectors.dense(1.7, 4.4, 7.6, 5.8, 9.6, 2.3), 3.0), + RowFactory.create(2, Vectors.dense(8.8, 7.3, 5.7, 7.3, 2.2, 4.1), 2.0), + RowFactory.create(3, Vectors.dense(1.2, 9.5, 2.5, 3.1, 8.7, 2.5), 3.0), + RowFactory.create(4, Vectors.dense(3.7, 9.2, 6.1, 4.1, 7.5, 3.8), 2.0), + RowFactory.create(5, Vectors.dense(8.9, 5.2, 7.8, 8.3, 5.2, 3.0), 4.0), + RowFactory.create(6, Vectors.dense(7.9, 8.5, 9.2, 4.0, 9.4, 2.1), 4.0) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()), + new StructField("label", DataTypes.DoubleType, false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + UnivariateFeatureSelector selector = new UnivariateFeatureSelector() + .setFeatureType("continuous") + .setLabelType("categorical") + .setSelectionMode("numTopFeatures") + .setSelectionThreshold(1) + .setFeaturesCol("features") + .setLabelCol("label") + .setOutputCol("selectedFeatures"); + + Dataset result = selector.fit(df).transform(df); + + System.out.println("UnivariateFeatureSelector output with top " + + selector.getSelectionThreshold() + " features selected using f_classif"); + result.show(); + + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVarianceThresholdSelectorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVarianceThresholdSelectorExample.java new file mode 100644 index 0000000..5820a95 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVarianceThresholdSelectorExample.java @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.VarianceThresholdSelector; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +// $example off$ + +/** + * An example for VarianceThresholdSelector. + * Run with + *
+ * bin/run-example ml.JavaVarianceThresholdSelectorExample
+ * 
+ */ +public class JavaVarianceThresholdSelectorExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaVarianceThresholdSelectorExample") + .getOrCreate(); + + // $example on$ + List data = Arrays.asList( + RowFactory.create(1, Vectors.dense(6.0, 7.0, 0.0, 7.0, 6.0, 0.0)), + RowFactory.create(2, Vectors.dense(0.0, 9.0, 6.0, 0.0, 5.0, 9.0)), + RowFactory.create(3, Vectors.dense(0.0, 9.0, 3.0, 0.0, 5.0, 5.0)), + RowFactory.create(4, Vectors.dense(0.0, 9.0, 8.0, 5.0, 6.0, 4.0)), + RowFactory.create(5, Vectors.dense(8.0, 9.0, 6.0, 5.0, 4.0, 4.0)), + RowFactory.create(6, Vectors.dense(8.0, 9.0, 6.0, 0.0, 0.0, 0.0)) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("features", new VectorUDT(), false, Metadata.empty()) + }); + + Dataset df = spark.createDataFrame(data, schema); + + VarianceThresholdSelector selector = new VarianceThresholdSelector() + .setVarianceThreshold(8.0) + .setFeaturesCol("features") + .setOutputCol("selectedFeatures"); + + Dataset result = selector.fit(df).transform(df); + + System.out.println("Output: Features with variance lower than " + + selector.getVarianceThreshold() + " are removed."); + result.show(); + + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorAssemblerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorAssemblerExample.java new file mode 100644 index 0000000..384e09c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorAssemblerExample.java @@ -0,0 +1,66 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.ml.feature.VectorAssembler; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +import static org.apache.spark.sql.types.DataTypes.*; +// $example off$ + +public class JavaVectorAssemblerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaVectorAssemblerExample") + .getOrCreate(); + + // $example on$ + StructType schema = createStructType(new StructField[]{ + createStructField("id", IntegerType, false), + createStructField("hour", IntegerType, false), + createStructField("mobile", DoubleType, false), + createStructField("userFeatures", new VectorUDT(), false), + createStructField("clicked", DoubleType, false) + }); + Row row = RowFactory.create(0, 18, 1.0, Vectors.dense(0.0, 10.0, 0.5), 1.0); + Dataset dataset = spark.createDataFrame(Arrays.asList(row), schema); + + VectorAssembler assembler = new VectorAssembler() + .setInputCols(new String[]{"hour", "mobile", "userFeatures"}) + .setOutputCol("features"); + + Dataset output = assembler.transform(dataset); + System.out.println("Assembled columns 'hour', 'mobile', 'userFeatures' to vector column " + + "'features'"); + output.select("features", "clicked").show(false); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorIndexerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorIndexerExample.java new file mode 100644 index 0000000..dd9d757 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorIndexerExample.java @@ -0,0 +1,61 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Map; + +import org.apache.spark.ml.feature.VectorIndexer; +import org.apache.spark.ml.feature.VectorIndexerModel; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off$ + +public class JavaVectorIndexerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaVectorIndexerExample") + .getOrCreate(); + + // $example on$ + Dataset data = spark.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt"); + + VectorIndexer indexer = new VectorIndexer() + .setInputCol("features") + .setOutputCol("indexed") + .setMaxCategories(10); + VectorIndexerModel indexerModel = indexer.fit(data); + + Map> categoryMaps = indexerModel.javaCategoryMaps(); + System.out.print("Chose " + categoryMaps.size() + " categorical features:"); + + for (Integer feature : categoryMaps.keySet()) { + System.out.print(" " + feature); + } + System.out.println(); + + // Create new column "indexed" with categorical values transformed to indices + Dataset indexedData = indexerModel.transform(data); + indexedData.show(); + // $example off$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSizeHintExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSizeHintExample.java new file mode 100644 index 0000000..d649a2c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSizeHintExample.java @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.ml.feature.VectorAssembler; +import org.apache.spark.ml.feature.VectorSizeHint; +import org.apache.spark.ml.linalg.VectorUDT; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import static org.apache.spark.sql.types.DataTypes.*; +// $example off$ + +public class JavaVectorSizeHintExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaVectorSizeHintExample") + .getOrCreate(); + + // $example on$ + StructType schema = createStructType(new StructField[]{ + createStructField("id", IntegerType, false), + createStructField("hour", IntegerType, false), + createStructField("mobile", DoubleType, false), + createStructField("userFeatures", new VectorUDT(), false), + createStructField("clicked", DoubleType, false) + }); + Row row0 = RowFactory.create(0, 18, 1.0, Vectors.dense(0.0, 10.0, 0.5), 1.0); + Row row1 = RowFactory.create(0, 18, 1.0, Vectors.dense(0.0, 10.0), 0.0); + Dataset dataset = spark.createDataFrame(Arrays.asList(row0, row1), schema); + + VectorSizeHint sizeHint = new VectorSizeHint() + .setInputCol("userFeatures") + .setHandleInvalid("skip") + .setSize(3); + + Dataset datasetWithSize = sizeHint.transform(dataset); + System.out.println("Rows where 'userFeatures' is not the right size are filtered out"); + datasetWithSize.show(false); + + VectorAssembler assembler = new VectorAssembler() + .setInputCols(new String[]{"hour", "mobile", "userFeatures"}) + .setOutputCol("features"); + + // This dataframe can be used by downstream transformers as before + Dataset output = assembler.transform(datasetWithSize); + System.out.println("Assembled columns 'hour', 'mobile', 'userFeatures' to vector column " + + "'features'"); + output.select("features", "clicked").show(false); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSlicerExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSlicerExample.java new file mode 100644 index 0000000..1ae48be --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaVectorSlicerExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +import org.apache.spark.sql.SparkSession; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.attribute.Attribute; +import org.apache.spark.ml.attribute.AttributeGroup; +import org.apache.spark.ml.attribute.NumericAttribute; +import org.apache.spark.ml.feature.VectorSlicer; +import org.apache.spark.ml.linalg.Vectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.types.*; +// $example off$ + +public class JavaVectorSlicerExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaVectorSlicerExample") + .getOrCreate(); + + // $example on$ + Attribute[] attrs = { + NumericAttribute.defaultAttr().withName("f1"), + NumericAttribute.defaultAttr().withName("f2"), + NumericAttribute.defaultAttr().withName("f3") + }; + AttributeGroup group = new AttributeGroup("userFeatures", attrs); + + List data = Arrays.asList( + RowFactory.create(Vectors.sparse(3, new int[]{0, 1}, new double[]{-2.0, 2.3})), + RowFactory.create(Vectors.dense(-2.0, 2.3, 0.0)) + ); + + Dataset dataset = + spark.createDataFrame(data, (new StructType()).add(group.toStructField())); + + VectorSlicer vectorSlicer = new VectorSlicer() + .setInputCol("userFeatures").setOutputCol("features"); + + vectorSlicer.setIndices(new int[]{1}).setNames(new String[]{"f3"}); + // or slicer.setIndices(new int[]{1, 2}), or slicer.setNames(new String[]{"f2", "f3"}) + + Dataset output = vectorSlicer.transform(dataset); + output.show(false); + // $example off$ + + spark.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaWord2VecExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaWord2VecExample.java new file mode 100644 index 0000000..fc9b459 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/ml/JavaWord2VecExample.java @@ -0,0 +1,72 @@ +/* + * 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. + */ + +package org.apache.spark.examples.ml; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.ml.feature.Word2Vec; +import org.apache.spark.ml.feature.Word2VecModel; +import org.apache.spark.ml.linalg.Vector; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.*; +// $example off$ + +public class JavaWord2VecExample { + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("JavaWord2VecExample") + .getOrCreate(); + + // $example on$ + // Input data: Each row is a bag of words from a sentence or document. + List data = Arrays.asList( + RowFactory.create(Arrays.asList("Hi I heard about Spark".split(" "))), + RowFactory.create(Arrays.asList("I wish Java could use case classes".split(" "))), + RowFactory.create(Arrays.asList("Logistic regression models are neat".split(" "))) + ); + StructType schema = new StructType(new StructField[]{ + new StructField("text", new ArrayType(DataTypes.StringType, true), false, Metadata.empty()) + }); + Dataset documentDF = spark.createDataFrame(data, schema); + + // Learn a mapping from words to Vectors. + Word2Vec word2Vec = new Word2Vec() + .setInputCol("text") + .setOutputCol("result") + .setVectorSize(3) + .setMinCount(0); + + Word2VecModel model = word2Vec.fit(documentDF); + Dataset result = model.transform(documentDF); + + for (Row row : result.collectAsList()) { + List text = row.getList(0); + Vector vector = (Vector) row.get(1); + System.out.println("Text: " + text + " => \nVector: " + vector + "\n"); + } + // $example off$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaALS.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaALS.java new file mode 100644 index 0000000..95a430f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaALS.java @@ -0,0 +1,90 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.function.Function; + +import org.apache.spark.mllib.recommendation.ALS; +import org.apache.spark.mllib.recommendation.MatrixFactorizationModel; +import org.apache.spark.mllib.recommendation.Rating; + +import java.util.Arrays; +import java.util.regex.Pattern; + +import scala.Tuple2; + +/** + * Example using MLlib ALS from Java. + */ +public final class JavaALS { + + static class ParseRating implements Function { + private static final Pattern COMMA = Pattern.compile(","); + + @Override + public Rating call(String line) { + String[] tok = COMMA.split(line); + int x = Integer.parseInt(tok[0]); + int y = Integer.parseInt(tok[1]); + double rating = Double.parseDouble(tok[2]); + return new Rating(x, y, rating); + } + } + + static class FeaturesToString implements Function, String> { + @Override + public String call(Tuple2 element) { + return element._1() + "," + Arrays.toString(element._2()); + } + } + + public static void main(String[] args) { + + if (args.length < 4) { + System.err.println( + "Usage: JavaALS []"); + System.exit(1); + } + SparkConf sparkConf = new SparkConf().setAppName("JavaALS"); + int rank = Integer.parseInt(args[1]); + int iterations = Integer.parseInt(args[2]); + String outputDir = args[3]; + int blocks = -1; + if (args.length == 5) { + blocks = Integer.parseInt(args[4]); + } + + JavaSparkContext sc = new JavaSparkContext(sparkConf); + JavaRDD lines = sc.textFile(args[0]); + + JavaRDD ratings = lines.map(new ParseRating()); + + MatrixFactorizationModel model = ALS.train(ratings.rdd(), rank, iterations, 0.01, blocks); + + model.userFeatures().toJavaRDD().map(new FeaturesToString()).saveAsTextFile( + outputDir + "/userFeatures"); + model.productFeatures().toJavaRDD().map(new FeaturesToString()).saveAsTextFile( + outputDir + "/productFeatures"); + System.out.println("Final user/product features written to " + outputDir); + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java new file mode 100644 index 0000000..5f43603 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java @@ -0,0 +1,58 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.fpm.AssociationRules; +import org.apache.spark.mllib.fpm.FPGrowth; +import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset; +// $example off$ + +import org.apache.spark.SparkConf; + +public class JavaAssociationRulesExample { + + public static void main(String[] args) { + + SparkConf sparkConf = new SparkConf().setAppName("JavaAssociationRulesExample"); + JavaSparkContext sc = new JavaSparkContext(sparkConf); + + // $example on$ + JavaRDD> freqItemsets = sc.parallelize(Arrays.asList( + new FreqItemset<>(new String[] {"a"}, 15L), + new FreqItemset<>(new String[] {"b"}, 35L), + new FreqItemset<>(new String[] {"a", "b"}, 12L) + )); + + AssociationRules arules = new AssociationRules() + .setMinConfidence(0.8); + JavaRDD> results = arules.run(freqItemsets); + + for (AssociationRules.Rule rule : results.collect()) { + System.out.println( + rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence()); + } + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java new file mode 100644 index 0000000..b9d0313 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java @@ -0,0 +1,102 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.classification.LogisticRegressionModel; +import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS; +import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; + +public class JavaBinaryClassificationMetricsExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("Java Binary Classification Metrics Example"); + SparkContext sc = new SparkContext(conf); + // $example on$ + String path = "data/mllib/sample_binary_classification_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD(); + + // Split initial RDD into two... [60% training data, 40% testing data]. + JavaRDD[] splits = + data.randomSplit(new double[]{0.6, 0.4}, 11L); + JavaRDD training = splits[0].cache(); + JavaRDD test = splits[1]; + + // Run training algorithm to build the model. + LogisticRegressionModel model = new LogisticRegressionWithLBFGS() + .setNumClasses(2) + .run(training.rdd()); + + // Clear the prediction threshold so the model will return probabilities + model.clearThreshold(); + + // Compute raw scores on the test set. + JavaPairRDD predictionAndLabels = test.mapToPair(p -> + new Tuple2<>(model.predict(p.features()), p.label())); + + // Get evaluation metrics. + BinaryClassificationMetrics metrics = + new BinaryClassificationMetrics(predictionAndLabels.rdd()); + + // Precision by threshold + JavaRDD> precision = metrics.precisionByThreshold().toJavaRDD(); + System.out.println("Precision by threshold: " + precision.collect()); + + // Recall by threshold + JavaRDD recall = metrics.recallByThreshold().toJavaRDD(); + System.out.println("Recall by threshold: " + recall.collect()); + + // F Score by threshold + JavaRDD f1Score = metrics.fMeasureByThreshold().toJavaRDD(); + System.out.println("F1 Score by threshold: " + f1Score.collect()); + + JavaRDD f2Score = metrics.fMeasureByThreshold(2.0).toJavaRDD(); + System.out.println("F2 Score by threshold: " + f2Score.collect()); + + // Precision-recall curve + JavaRDD prc = metrics.pr().toJavaRDD(); + System.out.println("Precision-recall curve: " + prc.collect()); + + // Thresholds + JavaRDD thresholds = precision.map(t -> Double.parseDouble(t._1().toString())); + + // ROC Curve + JavaRDD roc = metrics.roc().toJavaRDD(); + System.out.println("ROC curve: " + roc.collect()); + + // AUPRC + System.out.println("Area under precision-recall curve = " + metrics.areaUnderPR()); + + // AUROC + System.out.println("Area under ROC = " + metrics.areaUnderROC()); + + // Save and load model + model.save(sc, "target/tmp/LogisticRegressionModel"); + LogisticRegressionModel.load(sc, "target/tmp/LogisticRegressionModel"); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBisectingKMeansExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBisectingKMeansExample.java new file mode 100644 index 0000000..f878b55 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaBisectingKMeansExample.java @@ -0,0 +1,66 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; +// $example off$ +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.clustering.BisectingKMeans; +import org.apache.spark.mllib.clustering.BisectingKMeansModel; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +// $example off$ + +/** + * Java example for bisecting k-means clustering. + */ +public class JavaBisectingKMeansExample { + public static void main(String[] args) { + SparkConf sparkConf = new SparkConf().setAppName("JavaBisectingKMeansExample"); + JavaSparkContext sc = new JavaSparkContext(sparkConf); + + // $example on$ + List localData = Arrays.asList( + Vectors.dense(0.1, 0.1), Vectors.dense(0.3, 0.3), + Vectors.dense(10.1, 10.1), Vectors.dense(10.3, 10.3), + Vectors.dense(20.1, 20.1), Vectors.dense(20.3, 20.3), + Vectors.dense(30.1, 30.1), Vectors.dense(30.3, 30.3) + ); + JavaRDD data = sc.parallelize(localData, 2); + + BisectingKMeans bkm = new BisectingKMeans() + .setK(4); + BisectingKMeansModel model = bkm.run(data); + + System.out.println("Compute Cost: " + model.computeCost(data)); + + Vector[] clusterCenters = model.clusterCenters(); + for (int i = 0; i < clusterCenters.length; i++) { + Vector clusterCenter = clusterCenters[i]; + System.out.println("Cluster Center " + i + ": " + clusterCenter); + } + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaChiSqSelectorExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaChiSqSelectorExample.java new file mode 100644 index 0000000..ce354af --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaChiSqSelectorExample.java @@ -0,0 +1,65 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.feature.ChiSqSelector; +import org.apache.spark.mllib.feature.ChiSqSelectorModel; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +public class JavaChiSqSelectorExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaChiSqSelectorExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + JavaRDD points = MLUtils.loadLibSVMFile(jsc.sc(), + "data/mllib/sample_libsvm_data.txt").toJavaRDD().cache(); + + // Discretize data in 16 equal bins since ChiSqSelector requires categorical features + // Although features are doubles, the ChiSqSelector treats each unique value as a category + JavaRDD discretizedData = points.map(lp -> { + double[] discretizedFeatures = new double[lp.features().size()]; + for (int i = 0; i < lp.features().size(); ++i) { + discretizedFeatures[i] = Math.floor(lp.features().apply(i) / 16); + } + return new LabeledPoint(lp.label(), Vectors.dense(discretizedFeatures)); + }); + + // Create ChiSqSelector that will select top 50 of 692 features + ChiSqSelector selector = new ChiSqSelector(50); + // Create ChiSqSelector model (selecting features) + ChiSqSelectorModel transformer = selector.fit(discretizedData.rdd()); + // Filter the top 50 features from each feature vector + JavaRDD filteredData = discretizedData.map(lp -> + new LabeledPoint(lp.label(), transformer.transform(lp.features()))); + // $example off$ + + System.out.println("filtered data: "); + filteredData.foreach(System.out::println); + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaCorrelationsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaCorrelationsExample.java new file mode 100644 index 0000000..9bd858b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaCorrelationsExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaDoubleRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.linalg.Matrix; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.stat.Statistics; +// $example off$ + +public class JavaCorrelationsExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaCorrelationsExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + JavaDoubleRDD seriesX = jsc.parallelizeDoubles( + Arrays.asList(1.0, 2.0, 3.0, 3.0, 5.0)); // a series + + // must have the same number of partitions and cardinality as seriesX + JavaDoubleRDD seriesY = jsc.parallelizeDoubles( + Arrays.asList(11.0, 22.0, 33.0, 33.0, 555.0)); + + // compute the correlation using Pearson's method. Enter "spearman" for Spearman's method. + // If a method is not specified, Pearson's method will be used by default. + double correlation = Statistics.corr(seriesX.srdd(), seriesY.srdd(), "pearson"); + System.out.println("Correlation is: " + correlation); + + // note that each Vector is a row and not a column + JavaRDD data = jsc.parallelize( + Arrays.asList( + Vectors.dense(1.0, 10.0, 100.0), + Vectors.dense(2.0, 20.0, 200.0), + Vectors.dense(5.0, 33.0, 366.0) + ) + ); + + // calculate the correlation matrix using Pearson's method. + // Use "spearman" for Spearman's method. + // If a method is not specified, Pearson's method will be used by default. + Matrix correlMatrix = Statistics.corr(data.rdd(), "pearson"); + System.out.println(correlMatrix.toString()); + // $example off$ + + jsc.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java new file mode 100644 index 0000000..032c168 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.DecisionTree; +import org.apache.spark.mllib.tree.model.DecisionTreeModel; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +class JavaDecisionTreeClassificationExample { + + public static void main(String[] args) { + + // $example on$ + SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTreeClassificationExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Set parameters. + // Empty categoricalFeaturesInfo indicates all features are continuous. + int numClasses = 2; + Map categoricalFeaturesInfo = new HashMap<>(); + String impurity = "gini"; + int maxDepth = 5; + int maxBins = 32; + + // Train a DecisionTree model for classification. + DecisionTreeModel model = DecisionTree.trainClassifier(trainingData, numClasses, + categoricalFeaturesInfo, impurity, maxDepth, maxBins); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testErr = + predictionAndLabel.filter(pl -> !pl._1().equals(pl._2())).count() / (double) testData.count(); + + System.out.println("Test Error: " + testErr); + System.out.println("Learned classification tree model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myDecisionTreeClassificationModel"); + DecisionTreeModel sameModel = DecisionTreeModel + .load(jsc.sc(), "target/tmp/myDecisionTreeClassificationModel"); + // $example off$ + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java new file mode 100644 index 0000000..f222c38 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.DecisionTree; +import org.apache.spark.mllib.tree.model.DecisionTreeModel; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +class JavaDecisionTreeRegressionExample { + + public static void main(String[] args) { + + // $example on$ + SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTreeRegressionExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Set parameters. + // Empty categoricalFeaturesInfo indicates all features are continuous. + Map categoricalFeaturesInfo = new HashMap<>(); + String impurity = "variance"; + int maxDepth = 5; + int maxBins = 32; + + // Train a DecisionTree model. + DecisionTreeModel model = DecisionTree.trainRegressor(trainingData, + categoricalFeaturesInfo, impurity, maxDepth, maxBins); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testMSE = predictionAndLabel.mapToDouble(pl -> { + double diff = pl._1() - pl._2(); + return diff * diff; + }).mean(); + System.out.println("Test Mean Squared Error: " + testMSE); + System.out.println("Learned regression tree model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myDecisionTreeRegressionModel"); + DecisionTreeModel sameModel = DecisionTreeModel + .load(jsc.sc(), "target/tmp/myDecisionTreeRegressionModel"); + // $example off$ + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaElementwiseProductExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaElementwiseProductExample.java new file mode 100644 index 0000000..2d45c61 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaElementwiseProductExample.java @@ -0,0 +1,59 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +// $example off$ + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.feature.ElementwiseProduct; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +// $example off$ + +public class JavaElementwiseProductExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaElementwiseProductExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // Create some vector data; also works for sparse vectors + JavaRDD data = jsc.parallelize(Arrays.asList( + Vectors.dense(1.0, 2.0, 3.0), Vectors.dense(4.0, 5.0, 6.0))); + Vector transformingVector = Vectors.dense(0.0, 1.0, 2.0); + ElementwiseProduct transformer = new ElementwiseProduct(transformingVector); + + // Batch transform and per-row transform give the same results: + JavaRDD transformedData = transformer.transform(data); + JavaRDD transformedData2 = data.map(transformer::transform); + // $example off$ + + System.out.println("transformedData: "); + transformedData.foreach(System.out::println); + + System.out.println("transformedData2: "); + transformedData2.foreach(System.out::println); + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGaussianMixtureExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGaussianMixtureExample.java new file mode 100644 index 0000000..5792e5a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGaussianMixtureExample.java @@ -0,0 +1,68 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; + +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.clustering.GaussianMixture; +import org.apache.spark.mllib.clustering.GaussianMixtureModel; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +// $example off$ + +public class JavaGaussianMixtureExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaGaussianMixtureExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // Load and parse data + String path = "data/mllib/gmm_data.txt"; + JavaRDD data = jsc.textFile(path); + JavaRDD parsedData = data.map(s -> { + String[] sarray = s.trim().split(" "); + double[] values = new double[sarray.length]; + for (int i = 0; i < sarray.length; i++) { + values[i] = Double.parseDouble(sarray[i]); + } + return Vectors.dense(values); + }); + parsedData.cache(); + + // Cluster the data into two classes using GaussianMixture + GaussianMixtureModel gmm = new GaussianMixture().setK(2).run(parsedData.rdd()); + + // Save and load GaussianMixtureModel + gmm.save(jsc.sc(), "target/org/apache/spark/JavaGaussianMixtureExample/GaussianMixtureModel"); + GaussianMixtureModel sameModel = GaussianMixtureModel.load(jsc.sc(), + "target/org.apache.spark.JavaGaussianMixtureExample/GaussianMixtureModel"); + + // Output the parameters of the mixture model + for (int j = 0; j < gmm.k(); j++) { + System.out.printf("weight=%f\nmu=%s\nsigma=\n%s\n", + gmm.weights()[j], gmm.gaussians()[j].mu(), gmm.gaussians()[j].sigma()); + } + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java new file mode 100644 index 0000000..521ee96 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java @@ -0,0 +1,81 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.GradientBoostedTrees; +import org.apache.spark.mllib.tree.configuration.BoostingStrategy; +import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +public class JavaGradientBoostingClassificationExample { + public static void main(String[] args) { + // $example on$ + SparkConf sparkConf = new SparkConf() + .setAppName("JavaGradientBoostedTreesClassificationExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Train a GradientBoostedTrees model. + // The defaultParams for Classification use LogLoss by default. + BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Classification"); + boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice. + boostingStrategy.getTreeStrategy().setNumClasses(2); + boostingStrategy.getTreeStrategy().setMaxDepth(5); + // Empty categoricalFeaturesInfo indicates all features are continuous. + Map categoricalFeaturesInfo = new HashMap<>(); + boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo); + + GradientBoostedTreesModel model = GradientBoostedTrees.train(trainingData, boostingStrategy); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testErr = + predictionAndLabel.filter(pl -> !pl._1().equals(pl._2())).count() / (double) testData.count(); + System.out.println("Test Error: " + testErr); + System.out.println("Learned classification GBT model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myGradientBoostingClassificationModel"); + GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(jsc.sc(), + "target/tmp/myGradientBoostingClassificationModel"); + // $example off$ + + jsc.stop(); + } + +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java new file mode 100644 index 0000000..b345d19 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java @@ -0,0 +1,80 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.GradientBoostedTrees; +import org.apache.spark.mllib.tree.configuration.BoostingStrategy; +import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +public class JavaGradientBoostingRegressionExample { + public static void main(String[] args) { + // $example on$ + SparkConf sparkConf = new SparkConf() + .setAppName("JavaGradientBoostedTreesRegressionExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Train a GradientBoostedTrees model. + // The defaultParams for Regression use SquaredError by default. + BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Regression"); + boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice. + boostingStrategy.getTreeStrategy().setMaxDepth(5); + // Empty categoricalFeaturesInfo indicates all features are continuous. + Map categoricalFeaturesInfo = new HashMap<>(); + boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo); + + GradientBoostedTreesModel model = GradientBoostedTrees.train(trainingData, boostingStrategy); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testMSE = predictionAndLabel.mapToDouble(pl -> { + double diff = pl._1() - pl._2(); + return diff * diff; + }).mean(); + System.out.println("Test Mean Squared Error: " + testMSE); + System.out.println("Learned regression GBT model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myGradientBoostingRegressionModel"); + GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(jsc.sc(), + "target/tmp/myGradientBoostingRegressionModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingExample.java new file mode 100644 index 0000000..2732736 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingExample.java @@ -0,0 +1,84 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; + +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.linalg.Matrices; +import org.apache.spark.mllib.linalg.Matrix; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.stat.Statistics; +import org.apache.spark.mllib.stat.test.ChiSqTestResult; +// $example off$ + +public class JavaHypothesisTestingExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaHypothesisTestingExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // a vector composed of the frequencies of events + Vector vec = Vectors.dense(0.1, 0.15, 0.2, 0.3, 0.25); + + // compute the goodness of fit. If a second vector to test against is not supplied + // as a parameter, the test runs against a uniform distribution. + ChiSqTestResult goodnessOfFitTestResult = Statistics.chiSqTest(vec); + // summary of the test including the p-value, degrees of freedom, test statistic, + // the method used, and the null hypothesis. + System.out.println(goodnessOfFitTestResult + "\n"); + + // Create a contingency matrix ((1.0, 2.0), (3.0, 4.0), (5.0, 6.0)) + Matrix mat = Matrices.dense(3, 2, new double[]{1.0, 3.0, 5.0, 2.0, 4.0, 6.0}); + + // conduct Pearson's independence test on the input contingency matrix + ChiSqTestResult independenceTestResult = Statistics.chiSqTest(mat); + // summary of the test including the p-value, degrees of freedom... + System.out.println(independenceTestResult + "\n"); + + // an RDD of labeled points + JavaRDD obs = jsc.parallelize( + Arrays.asList( + new LabeledPoint(1.0, Vectors.dense(1.0, 0.0, 3.0)), + new LabeledPoint(1.0, Vectors.dense(1.0, 2.0, 0.0)), + new LabeledPoint(-1.0, Vectors.dense(-1.0, 0.0, -0.5)) + ) + ); + + // The contingency table is constructed from the raw (label, feature) pairs and used to conduct + // the independence test. Returns an array containing the ChiSquaredTestResult for every feature + // against the label. + ChiSqTestResult[] featureTestResults = Statistics.chiSqTest(obs.rdd()); + int i = 1; + for (ChiSqTestResult result : featureTestResults) { + System.out.println("Column " + i + ":"); + System.out.println(result + "\n"); // summary of the test + i++; + } + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingKolmogorovSmirnovTestExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingKolmogorovSmirnovTestExample.java new file mode 100644 index 0000000..fe611c9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaHypothesisTestingKolmogorovSmirnovTestExample.java @@ -0,0 +1,49 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaDoubleRDD; +import org.apache.spark.mllib.stat.Statistics; +import org.apache.spark.mllib.stat.test.KolmogorovSmirnovTestResult; +// $example off$ + +public class JavaHypothesisTestingKolmogorovSmirnovTestExample { + public static void main(String[] args) { + + SparkConf conf = + new SparkConf().setAppName("JavaHypothesisTestingKolmogorovSmirnovTestExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + JavaDoubleRDD data = jsc.parallelizeDoubles(Arrays.asList(0.1, 0.15, 0.2, 0.3, 0.25)); + KolmogorovSmirnovTestResult testResult = + Statistics.kolmogorovSmirnovTest(data, "norm", 0.0, 1.0); + // summary of the test including the p-value, test statistic, and null hypothesis + // if our p-value indicates significance, we can reject the null hypothesis + System.out.println(testResult); + // $example off$ + + jsc.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java new file mode 100644 index 0000000..adebafe --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ +package org.apache.spark.examples.mllib; + +// $example on$ + +import scala.Tuple2; +import scala.Tuple3; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.regression.IsotonicRegression; +import org.apache.spark.mllib.regression.IsotonicRegressionModel; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ +import org.apache.spark.SparkConf; + +public class JavaIsotonicRegressionExample { + public static void main(String[] args) { + SparkConf sparkConf = new SparkConf().setAppName("JavaIsotonicRegressionExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + // $example on$ + JavaRDD data = MLUtils.loadLibSVMFile( + jsc.sc(), "data/mllib/sample_isotonic_regression_libsvm_data.txt").toJavaRDD(); + + // Create label, feature, weight tuples from input data with weight set to default value 1.0. + JavaRDD> parsedData = data.map(point -> + new Tuple3<>(point.label(), point.features().apply(0), 1.0)); + + // Split data into training (60%) and test (40%) sets. + JavaRDD>[] splits = + parsedData.randomSplit(new double[]{0.6, 0.4}, 11L); + JavaRDD> training = splits[0]; + JavaRDD> test = splits[1]; + + // Create isotonic regression model from training data. + // Isotonic parameter defaults to true so it is only shown for demonstration + IsotonicRegressionModel model = new IsotonicRegression().setIsotonic(true).run(training); + + // Create tuples of predicted and real labels. + JavaPairRDD predictionAndLabel = test.mapToPair(point -> + new Tuple2<>(model.predict(point._2()), point._1())); + + // Calculate mean squared error between predicted and real labels. + double meanSquaredError = predictionAndLabel.mapToDouble(pl -> { + double diff = pl._1() - pl._2(); + return diff * diff; + }).mean(); + System.out.println("Mean Squared Error = " + meanSquaredError); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myIsotonicRegressionModel"); + IsotonicRegressionModel sameModel = + IsotonicRegressionModel.load(jsc.sc(), "target/tmp/myIsotonicRegressionModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKMeansExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKMeansExample.java new file mode 100644 index 0000000..f172756 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKMeansExample.java @@ -0,0 +1,75 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; + +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.clustering.KMeans; +import org.apache.spark.mllib.clustering.KMeansModel; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +// $example off$ + +public class JavaKMeansExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaKMeansExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // Load and parse data + String path = "data/mllib/kmeans_data.txt"; + JavaRDD data = jsc.textFile(path); + JavaRDD parsedData = data.map(s -> { + String[] sarray = s.split(" "); + double[] values = new double[sarray.length]; + for (int i = 0; i < sarray.length; i++) { + values[i] = Double.parseDouble(sarray[i]); + } + return Vectors.dense(values); + }); + parsedData.cache(); + + // Cluster the data into two classes using KMeans + int numClusters = 2; + int numIterations = 20; + KMeansModel clusters = KMeans.train(parsedData.rdd(), numClusters, numIterations); + + System.out.println("Cluster centers:"); + for (Vector center: clusters.clusterCenters()) { + System.out.println(" " + center); + } + double cost = clusters.computeCost(parsedData.rdd()); + System.out.println("Cost: " + cost); + + // Evaluate clustering by computing Within Set Sum of Squared Errors + double WSSSE = clusters.computeCost(parsedData.rdd()); + System.out.println("Within Set Sum of Squared Errors = " + WSSSE); + + // Save and load model + clusters.save(jsc.sc(), "target/org/apache/spark/JavaKMeansExample/KMeansModel"); + KMeansModel sameModel = KMeansModel.load(jsc.sc(), + "target/org/apache/spark/JavaKMeansExample/KMeansModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKernelDensityEstimationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKernelDensityEstimationExample.java new file mode 100644 index 0000000..41de0d9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaKernelDensityEstimationExample.java @@ -0,0 +1,53 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.stat.KernelDensity; +// $example off$ + +public class JavaKernelDensityEstimationExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaKernelDensityEstimationExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // an RDD of sample data + JavaRDD data = jsc.parallelize( + Arrays.asList(1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 6.0, 7.0, 8.0, 9.0, 9.0)); + + // Construct the density estimator with the sample data + // and a standard deviation for the Gaussian kernels + KernelDensity kd = new KernelDensity().setSample(data).setBandwidth(3.0); + + // Find density estimates for the given values + double[] densities = kd.estimate(new double[]{-1.0, 2.0, 5.0}); + + System.out.println(Arrays.toString(densities)); + // $example off$ + + jsc.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java new file mode 100644 index 0000000..3fdc03a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java @@ -0,0 +1,101 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; + +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.classification.LogisticRegressionModel; +import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.optimization.*; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; +// $example off$ + +public class JavaLBFGSExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("L-BFGS Example"); + SparkContext sc = new SparkContext(conf); + + // $example on$ + String path = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD(); + int numFeatures = data.take(1).get(0).features().size(); + + // Split initial RDD into two... [60% training data, 40% testing data]. + JavaRDD trainingInit = data.sample(false, 0.6, 11L); + JavaRDD test = data.subtract(trainingInit); + + // Append 1 into the training data as intercept. + JavaPairRDD training = data.mapToPair(p -> + new Tuple2<>(p.label(), MLUtils.appendBias(p.features()))); + training.cache(); + + // Run training algorithm to build the model. + int numCorrections = 10; + double convergenceTol = 1e-4; + int maxNumIterations = 20; + double regParam = 0.1; + Vector initialWeightsWithIntercept = Vectors.dense(new double[numFeatures + 1]); + + Tuple2 result = LBFGS.runLBFGS( + training.rdd(), + new LogisticGradient(), + new SquaredL2Updater(), + numCorrections, + convergenceTol, + maxNumIterations, + regParam, + initialWeightsWithIntercept); + Vector weightsWithIntercept = result._1(); + double[] loss = result._2(); + + LogisticRegressionModel model = new LogisticRegressionModel( + Vectors.dense(Arrays.copyOf(weightsWithIntercept.toArray(), weightsWithIntercept.size() - 1)), + (weightsWithIntercept.toArray())[weightsWithIntercept.size() - 1]); + + // Clear the default threshold. + model.clearThreshold(); + + // Compute raw scores on the test set. + JavaPairRDD scoreAndLabels = test.mapToPair(p -> + new Tuple2<>(model.predict(p.features()), p.label())); + + // Get evaluation metrics. + BinaryClassificationMetrics metrics = + new BinaryClassificationMetrics(scoreAndLabels.rdd()); + double auROC = metrics.areaUnderROC(); + + System.out.println("Loss of each step in training process"); + for (double l : loss) { + System.out.println(l); + } + System.out.println("Area under ROC = " + auROC); + // $example off$ + + sc.stop(); + } +} + diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLatentDirichletAllocationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLatentDirichletAllocationExample.java new file mode 100644 index 0000000..887edf8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLatentDirichletAllocationExample.java @@ -0,0 +1,82 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.clustering.DistributedLDAModel; +import org.apache.spark.mllib.clustering.LDA; +import org.apache.spark.mllib.clustering.LDAModel; +import org.apache.spark.mllib.linalg.Matrix; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +// $example off$ + +public class JavaLatentDirichletAllocationExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaKLatentDirichletAllocationExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + // Load and parse the data + String path = "data/mllib/sample_lda_data.txt"; + JavaRDD data = jsc.textFile(path); + JavaRDD parsedData = data.map(s -> { + String[] sarray = s.trim().split(" "); + double[] values = new double[sarray.length]; + for (int i = 0; i < sarray.length; i++) { + values[i] = Double.parseDouble(sarray[i]); + } + return Vectors.dense(values); + }); + // Index documents with unique IDs + JavaPairRDD corpus = + JavaPairRDD.fromJavaRDD(parsedData.zipWithIndex().map(Tuple2::swap)); + corpus.cache(); + + // Cluster the documents into three topics using LDA + LDAModel ldaModel = new LDA().setK(3).run(corpus); + + // Output topics. Each is a distribution over words (matching word count vectors) + System.out.println("Learned topics (as distributions over vocab of " + ldaModel.vocabSize() + + " words):"); + Matrix topics = ldaModel.topicsMatrix(); + for (int topic = 0; topic < 3; topic++) { + System.out.print("Topic " + topic + ":"); + for (int word = 0; word < ldaModel.vocabSize(); word++) { + System.out.print(" " + topics.apply(word, topic)); + } + System.out.println(); + } + + ldaModel.save(jsc.sc(), + "target/org/apache/spark/JavaLatentDirichletAllocationExample/LDAModel"); + DistributedLDAModel sameModel = DistributedLDAModel.load(jsc.sc(), + "target/org/apache/spark/JavaLatentDirichletAllocationExample/LDAModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLogisticRegressionWithLBFGSExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLogisticRegressionWithLBFGSExample.java new file mode 100644 index 0000000..26b8a6e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaLogisticRegressionWithLBFGSExample.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.classification.LogisticRegressionModel; +import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS; +import org.apache.spark.mllib.evaluation.MulticlassMetrics; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +/** + * Example for LogisticRegressionWithLBFGS. + */ +public class JavaLogisticRegressionWithLBFGSExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("JavaLogisticRegressionWithLBFGSExample"); + SparkContext sc = new SparkContext(conf); + // $example on$ + String path = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD(); + + // Split initial RDD into two... [60% training data, 40% testing data]. + JavaRDD[] splits = data.randomSplit(new double[] {0.6, 0.4}, 11L); + JavaRDD training = splits[0].cache(); + JavaRDD test = splits[1]; + + // Run training algorithm to build the model. + LogisticRegressionModel model = new LogisticRegressionWithLBFGS() + .setNumClasses(10) + .run(training.rdd()); + + // Compute raw scores on the test set. + JavaPairRDD predictionAndLabels = test.mapToPair(p -> + new Tuple2<>(model.predict(p.features()), p.label())); + + // Get evaluation metrics. + MulticlassMetrics metrics = new MulticlassMetrics(predictionAndLabels.rdd()); + double accuracy = metrics.accuracy(); + System.out.println("Accuracy = " + accuracy); + + // Save and load model + model.save(sc, "target/tmp/javaLogisticRegressionWithLBFGSModel"); + LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc, + "target/tmp/javaLogisticRegressionWithLBFGSModel"); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java new file mode 100644 index 0000000..bc99dc0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java @@ -0,0 +1,80 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.evaluation.MultilabelMetrics; +import org.apache.spark.SparkConf; +// $example off$ + +public class JavaMultiLabelClassificationMetricsExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("Multilabel Classification Metrics Example"); + JavaSparkContext sc = new JavaSparkContext(conf); + // $example on$ + List> data = Arrays.asList( + new Tuple2<>(new double[]{0.0, 1.0}, new double[]{0.0, 2.0}), + new Tuple2<>(new double[]{0.0, 2.0}, new double[]{0.0, 1.0}), + new Tuple2<>(new double[]{}, new double[]{0.0}), + new Tuple2<>(new double[]{2.0}, new double[]{2.0}), + new Tuple2<>(new double[]{2.0, 0.0}, new double[]{2.0, 0.0}), + new Tuple2<>(new double[]{0.0, 1.0, 2.0}, new double[]{0.0, 1.0}), + new Tuple2<>(new double[]{1.0}, new double[]{1.0, 2.0}) + ); + JavaRDD> scoreAndLabels = sc.parallelize(data); + + // Instantiate metrics object + MultilabelMetrics metrics = new MultilabelMetrics(scoreAndLabels.rdd()); + + // Summary stats + System.out.format("Recall = %f\n", metrics.recall()); + System.out.format("Precision = %f\n", metrics.precision()); + System.out.format("F1 measure = %f\n", metrics.f1Measure()); + System.out.format("Accuracy = %f\n", metrics.accuracy()); + + // Stats by labels + for (int i = 0; i < metrics.labels().length - 1; i++) { + System.out.format("Class %1.1f precision = %f\n", metrics.labels()[i], metrics.precision( + metrics.labels()[i])); + System.out.format("Class %1.1f recall = %f\n", metrics.labels()[i], metrics.recall( + metrics.labels()[i])); + System.out.format("Class %1.1f F1 score = %f\n", metrics.labels()[i], metrics.f1Measure( + metrics.labels()[i])); + } + + // Micro stats + System.out.format("Micro recall = %f\n", metrics.microRecall()); + System.out.format("Micro precision = %f\n", metrics.microPrecision()); + System.out.format("Micro F1 measure = %f\n", metrics.microF1Measure()); + + // Hamming loss + System.out.format("Hamming loss = %f\n", metrics.hammingLoss()); + + // Subset accuracy + System.out.format("Subset accuracy = %f\n", metrics.subsetAccuracy()); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java new file mode 100644 index 0000000..0367038 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java @@ -0,0 +1,90 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.classification.LogisticRegressionModel; +import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS; +import org.apache.spark.mllib.evaluation.MulticlassMetrics; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +import org.apache.spark.mllib.linalg.Matrix; +// $example off$ +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; + +public class JavaMulticlassClassificationMetricsExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("Multi class Classification Metrics Example"); + SparkContext sc = new SparkContext(conf); + // $example on$ + String path = "data/mllib/sample_multiclass_classification_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD(); + + // Split initial RDD into two... [60% training data, 40% testing data]. + JavaRDD[] splits = data.randomSplit(new double[]{0.6, 0.4}, 11L); + JavaRDD training = splits[0].cache(); + JavaRDD test = splits[1]; + + // Run training algorithm to build the model. + LogisticRegressionModel model = new LogisticRegressionWithLBFGS() + .setNumClasses(3) + .run(training.rdd()); + + // Compute raw scores on the test set. + JavaPairRDD predictionAndLabels = test.mapToPair(p -> + new Tuple2<>(model.predict(p.features()), p.label())); + + // Get evaluation metrics. + MulticlassMetrics metrics = new MulticlassMetrics(predictionAndLabels.rdd()); + + // Confusion matrix + Matrix confusion = metrics.confusionMatrix(); + System.out.println("Confusion matrix: \n" + confusion); + + // Overall statistics + System.out.println("Accuracy = " + metrics.accuracy()); + + // Stats by labels + for (int i = 0; i < metrics.labels().length; i++) { + System.out.format("Class %f precision = %f\n", metrics.labels()[i],metrics.precision( + metrics.labels()[i])); + System.out.format("Class %f recall = %f\n", metrics.labels()[i], metrics.recall( + metrics.labels()[i])); + System.out.format("Class %f F1 score = %f\n", metrics.labels()[i], metrics.fMeasure( + metrics.labels()[i])); + } + + //Weighted stats + System.out.format("Weighted precision = %f\n", metrics.weightedPrecision()); + System.out.format("Weighted recall = %f\n", metrics.weightedRecall()); + System.out.format("Weighted F1 score = %f\n", metrics.weightedFMeasure()); + System.out.format("Weighted false positive rate = %f\n", metrics.weightedFalsePositiveRate()); + + // Save and load model + model.save(sc, "target/tmp/LogisticRegressionModel"); + LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc, + "target/tmp/LogisticRegressionModel"); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java new file mode 100644 index 0000000..d80dbe8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java @@ -0,0 +1,55 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import scala.Tuple2; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.classification.NaiveBayes; +import org.apache.spark.mllib.classification.NaiveBayesModel; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ +import org.apache.spark.SparkConf; + +public class JavaNaiveBayesExample { + public static void main(String[] args) { + SparkConf sparkConf = new SparkConf().setAppName("JavaNaiveBayesExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + // $example on$ + String path = "data/mllib/sample_libsvm_data.txt"; + JavaRDD inputData = MLUtils.loadLibSVMFile(jsc.sc(), path).toJavaRDD(); + JavaRDD[] tmp = inputData.randomSplit(new double[]{0.6, 0.4}); + JavaRDD training = tmp[0]; // training set + JavaRDD test = tmp[1]; // test set + NaiveBayesModel model = NaiveBayes.train(training.rdd(), 1.0); + JavaPairRDD predictionAndLabel = + test.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double accuracy = + predictionAndLabel.filter(pl -> pl._1().equals(pl._2())).count() / (double) test.count(); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myNaiveBayesModel"); + NaiveBayesModel sameModel = NaiveBayesModel.load(jsc.sc(), "target/tmp/myNaiveBayesModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPCAExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPCAExample.java new file mode 100644 index 0000000..0a7dc62 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPCAExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; +// $example off$ + +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.linalg.Matrix; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.linalg.distributed.RowMatrix; +// $example off$ + +/** + * Example for compute principal components on a 'RowMatrix'. + */ +public class JavaPCAExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("PCA Example"); + SparkContext sc = new SparkContext(conf); + JavaSparkContext jsc = JavaSparkContext.fromSparkContext(sc); + + // $example on$ + List data = Arrays.asList( + Vectors.sparse(5, new int[] {1, 3}, new double[] {1.0, 7.0}), + Vectors.dense(2.0, 0.0, 3.0, 4.0, 5.0), + Vectors.dense(4.0, 0.0, 0.0, 6.0, 7.0) + ); + + JavaRDD rows = jsc.parallelize(data); + + // Create a RowMatrix from JavaRDD. + RowMatrix mat = new RowMatrix(rows.rdd()); + + // Compute the top 4 principal components. + // Principal components are stored in a local dense matrix. + Matrix pc = mat.computePrincipalComponents(4); + + // Project the rows to the linear space spanned by the top 4 principal components. + RowMatrix projected = mat.multiply(pc); + // $example off$ + Vector[] collectPartitions = (Vector[])projected.rows().collect(); + System.out.println("Projected vector of principal component:"); + for (Vector vector : collectPartitions) { + System.out.println("\t" + vector); + } + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPowerIterationClusteringExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPowerIterationClusteringExample.java new file mode 100644 index 0000000..e2260e1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPowerIterationClusteringExample.java @@ -0,0 +1,60 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import java.util.Arrays; + +import scala.Tuple3; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import org.apache.spark.mllib.clustering.PowerIterationClustering; +import org.apache.spark.mllib.clustering.PowerIterationClusteringModel; +// $example off$ + +/** + * Java example for graph clustering using power iteration clustering (PIC). + */ +public class JavaPowerIterationClusteringExample { + public static void main(String[] args) { + SparkConf sparkConf = new SparkConf().setAppName("JavaPowerIterationClusteringExample"); + JavaSparkContext sc = new JavaSparkContext(sparkConf); + + // $example on$ + JavaRDD> similarities = sc.parallelize(Arrays.asList( + new Tuple3<>(0L, 1L, 0.9), + new Tuple3<>(1L, 2L, 0.9), + new Tuple3<>(2L, 3L, 0.9), + new Tuple3<>(3L, 4L, 0.1), + new Tuple3<>(4L, 5L, 0.9))); + + PowerIterationClustering pic = new PowerIterationClustering() + .setK(2) + .setMaxIterations(10); + PowerIterationClusteringModel model = pic.run(similarities); + + for (PowerIterationClustering.Assignment a: model.assignments().toJavaRDD().collect()) { + System.out.println(a.id() + " -> " + a.cluster()); + } + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java new file mode 100644 index 0000000..1634075 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java @@ -0,0 +1,57 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; +// $example off$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import org.apache.spark.mllib.fpm.PrefixSpan; +import org.apache.spark.mllib.fpm.PrefixSpanModel; +// $example off$ +import org.apache.spark.SparkConf; + +public class JavaPrefixSpanExample { + + public static void main(String[] args) { + + SparkConf sparkConf = new SparkConf().setAppName("JavaPrefixSpanExample"); + JavaSparkContext sc = new JavaSparkContext(sparkConf); + + // $example on$ + JavaRDD>> sequences = sc.parallelize(Arrays.asList( + Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3)), + Arrays.asList(Arrays.asList(1), Arrays.asList(3, 2), Arrays.asList(1, 2)), + Arrays.asList(Arrays.asList(1, 2), Arrays.asList(5)), + Arrays.asList(Arrays.asList(6)) + ), 2); + PrefixSpan prefixSpan = new PrefixSpan() + .setMinSupport(0.5) + .setMaxPatternLength(5); + PrefixSpanModel model = prefixSpan.run(sequences); + for (PrefixSpan.FreqSequence freqSeq: model.freqSequences().toJavaRDD().collect()) { + System.out.println(freqSeq.javaSequence() + ", " + freqSeq.freq()); + } + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java new file mode 100644 index 0000000..7ee3c0a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java @@ -0,0 +1,80 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.RandomForest; +import org.apache.spark.mllib.tree.model.RandomForestModel; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +public class JavaRandomForestClassificationExample { + public static void main(String[] args) { + // $example on$ + SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestClassificationExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Train a RandomForest model. + // Empty categoricalFeaturesInfo indicates all features are continuous. + int numClasses = 2; + Map categoricalFeaturesInfo = new HashMap<>(); + int numTrees = 3; // Use more in practice. + String featureSubsetStrategy = "auto"; // Let the algorithm choose. + String impurity = "gini"; + int maxDepth = 5; + int maxBins = 32; + int seed = 12345; + + RandomForestModel model = RandomForest.trainClassifier(trainingData, numClasses, + categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins, + seed); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testErr = + predictionAndLabel.filter(pl -> !pl._1().equals(pl._2())).count() / (double) testData.count(); + System.out.println("Test Error: " + testErr); + System.out.println("Learned classification forest model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myRandomForestClassificationModel"); + RandomForestModel sameModel = RandomForestModel.load(jsc.sc(), + "target/tmp/myRandomForestClassificationModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java new file mode 100644 index 0000000..4a0f55f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java @@ -0,0 +1,80 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.HashMap; +import java.util.Map; + +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.tree.RandomForest; +import org.apache.spark.mllib.tree.model.RandomForestModel; +import org.apache.spark.mllib.util.MLUtils; +import org.apache.spark.SparkConf; +// $example off$ + +public class JavaRandomForestRegressionExample { + public static void main(String[] args) { + // $example on$ + SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestRegressionExample"); + JavaSparkContext jsc = new JavaSparkContext(sparkConf); + // Load and parse the data file. + String datapath = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD(); + // Split the data into training and test sets (30% held out for testing) + JavaRDD[] splits = data.randomSplit(new double[]{0.7, 0.3}); + JavaRDD trainingData = splits[0]; + JavaRDD testData = splits[1]; + + // Set parameters. + // Empty categoricalFeaturesInfo indicates all features are continuous. + Map categoricalFeaturesInfo = new HashMap<>(); + int numTrees = 3; // Use more in practice. + String featureSubsetStrategy = "auto"; // Let the algorithm choose. + String impurity = "variance"; + int maxDepth = 4; + int maxBins = 32; + int seed = 12345; + // Train a RandomForest model. + RandomForestModel model = RandomForest.trainRegressor(trainingData, + categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins, seed); + + // Evaluate model on test instances and compute test error + JavaPairRDD predictionAndLabel = + testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label())); + double testMSE = predictionAndLabel.mapToDouble(pl -> { + double diff = pl._1() - pl._2(); + return diff * diff; + }).mean(); + System.out.println("Test Mean Squared Error: " + testMSE); + System.out.println("Learned regression forest model:\n" + model.toDebugString()); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myRandomForestRegressionModel"); + RandomForestModel sameModel = RandomForestModel.load(jsc.sc(), + "target/tmp/myRandomForestRegressionModel"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java new file mode 100644 index 0000000..de398de --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java @@ -0,0 +1,142 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.*; + +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.evaluation.RegressionMetrics; +import org.apache.spark.mllib.evaluation.RankingMetrics; +import org.apache.spark.mllib.recommendation.ALS; +import org.apache.spark.mllib.recommendation.MatrixFactorizationModel; +import org.apache.spark.mllib.recommendation.Rating; +// $example off$ +import org.apache.spark.SparkConf; + +public class JavaRankingMetricsExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("Java Ranking Metrics Example"); + JavaSparkContext sc = new JavaSparkContext(conf); + // $example on$ + String path = "data/mllib/sample_movielens_data.txt"; + JavaRDD data = sc.textFile(path); + JavaRDD ratings = data.map(line -> { + String[] parts = line.split("::"); + return new Rating(Integer.parseInt(parts[0]), Integer.parseInt(parts[1]), Double + .parseDouble(parts[2]) - 2.5); + }); + ratings.cache(); + + // Train an ALS model + MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), 10, 10, 0.01); + + // Get top 10 recommendations for every user and scale ratings from 0 to 1 + JavaRDD> userRecs = model.recommendProductsForUsers(10).toJavaRDD(); + JavaRDD> userRecsScaled = userRecs.map(t -> { + Rating[] scaledRatings = new Rating[t._2().length]; + for (int i = 0; i < scaledRatings.length; i++) { + double newRating = Math.max(Math.min(t._2()[i].rating(), 1.0), 0.0); + scaledRatings[i] = new Rating(t._2()[i].user(), t._2()[i].product(), newRating); + } + return new Tuple2<>(t._1(), scaledRatings); + }); + JavaPairRDD userRecommended = JavaPairRDD.fromJavaRDD(userRecsScaled); + + // Map ratings to 1 or 0, 1 indicating a movie that should be recommended + JavaRDD binarizedRatings = ratings.map(r -> { + double binaryRating; + if (r.rating() > 0.0) { + binaryRating = 1.0; + } else { + binaryRating = 0.0; + } + return new Rating(r.user(), r.product(), binaryRating); + }); + + // Group ratings by common user + JavaPairRDD> userMovies = binarizedRatings.groupBy(Rating::user); + + // Get true relevant documents from all user ratings + JavaPairRDD> userMoviesList = userMovies.mapValues(docs -> { + List products = new ArrayList<>(); + for (Rating r : docs) { + if (r.rating() > 0.0) { + products.add(r.product()); + } + } + return products; + }); + + // Extract the product id from each recommendation + JavaPairRDD> userRecommendedList = userRecommended.mapValues(docs -> { + List products = new ArrayList<>(); + for (Rating r : docs) { + products.add(r.product()); + } + return products; + }); + JavaRDD, List>> relevantDocs = userMoviesList.join( + userRecommendedList).values(); + + // Instantiate the metrics object + RankingMetrics metrics = RankingMetrics.of(relevantDocs); + + // Precision, NDCG and Recall at k + Integer[] kVector = {1, 3, 5}; + for (Integer k : kVector) { + System.out.format("Precision at %d = %f\n", k, metrics.precisionAt(k)); + System.out.format("NDCG at %d = %f\n", k, metrics.ndcgAt(k)); + System.out.format("Recall at %d = %f\n", k, metrics.recallAt(k)); + } + + // Mean average precision + System.out.format("Mean average precision = %f\n", metrics.meanAveragePrecision()); + + //Mean average precision at k + System.out.format("Mean average precision at 2 = %f\n", metrics.meanAveragePrecisionAt(2)); + + // Evaluate the model using numerical ratings and regression metrics + JavaRDD> userProducts = + ratings.map(r -> new Tuple2<>(r.user(), r.product())); + + JavaPairRDD, Object> predictions = JavaPairRDD.fromJavaRDD( + model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD().map(r -> + new Tuple2<>(new Tuple2<>(r.user(), r.product()), r.rating()))); + JavaRDD> ratesAndPreds = + JavaPairRDD.fromJavaRDD(ratings.map(r -> + new Tuple2, Object>( + new Tuple2<>(r.user(), r.product()), + r.rating()) + )).join(predictions).values(); + + // Create regression metrics object + RegressionMetrics regressionMetrics = new RegressionMetrics(ratesAndPreds.rdd()); + + // Root mean squared error + System.out.format("RMSE = %f\n", regressionMetrics.rootMeanSquaredError()); + + // R-squared + System.out.format("R-squared = %f\n", regressionMetrics.r2()); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java new file mode 100644 index 0000000..1ee68da --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java @@ -0,0 +1,75 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.*; +import org.apache.spark.mllib.recommendation.ALS; +import org.apache.spark.mllib.recommendation.MatrixFactorizationModel; +import org.apache.spark.mllib.recommendation.Rating; +import org.apache.spark.SparkConf; +// $example off$ + +public class JavaRecommendationExample { + public static void main(String[] args) { + // $example on$ + SparkConf conf = new SparkConf().setAppName("Java Collaborative Filtering Example"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // Load and parse the data + String path = "data/mllib/als/test.data"; + JavaRDD data = jsc.textFile(path); + JavaRDD ratings = data.map(s -> { + String[] sarray = s.split(","); + return new Rating(Integer.parseInt(sarray[0]), + Integer.parseInt(sarray[1]), + Double.parseDouble(sarray[2])); + }); + + // Build the recommendation model using ALS + int rank = 10; + int numIterations = 10; + MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), rank, numIterations, 0.01); + + // Evaluate the model on rating data + JavaRDD> userProducts = + ratings.map(r -> new Tuple2<>(r.user(), r.product())); + JavaPairRDD, Double> predictions = JavaPairRDD.fromJavaRDD( + model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD() + .map(r -> new Tuple2<>(new Tuple2<>(r.user(), r.product()), r.rating())) + ); + JavaRDD> ratesAndPreds = JavaPairRDD.fromJavaRDD( + ratings.map(r -> new Tuple2<>(new Tuple2<>(r.user(), r.product()), r.rating()))) + .join(predictions).values(); + double MSE = ratesAndPreds.mapToDouble(pair -> { + double err = pair._1() - pair._2(); + return err * err; + }).mean(); + System.out.println("Mean Squared Error = " + MSE); + + // Save and load model + model.save(jsc.sc(), "target/tmp/myCollaborativeFilter"); + MatrixFactorizationModel sameModel = MatrixFactorizationModel.load(jsc.sc(), + "target/tmp/myCollaborativeFilter"); + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVDExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVDExample.java new file mode 100644 index 0000000..802be39 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVDExample.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; +// $example off$ + +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; +// $example on$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.linalg.Matrix; +import org.apache.spark.mllib.linalg.SingularValueDecomposition; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.linalg.distributed.RowMatrix; +// $example off$ + +/** + * Example for SingularValueDecomposition. + */ +public class JavaSVDExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("SVD Example"); + SparkContext sc = new SparkContext(conf); + JavaSparkContext jsc = JavaSparkContext.fromSparkContext(sc); + + // $example on$ + List data = Arrays.asList( + Vectors.sparse(5, new int[] {1, 3}, new double[] {1.0, 7.0}), + Vectors.dense(2.0, 0.0, 3.0, 4.0, 5.0), + Vectors.dense(4.0, 0.0, 0.0, 6.0, 7.0) + ); + + JavaRDD rows = jsc.parallelize(data); + + // Create a RowMatrix from JavaRDD. + RowMatrix mat = new RowMatrix(rows.rdd()); + + // Compute the top 5 singular values and corresponding singular vectors. + SingularValueDecomposition svd = mat.computeSVD(5, true, 1.0E-9d); + RowMatrix U = svd.U(); // The U factor is a RowMatrix. + Vector s = svd.s(); // The singular values are stored in a local dense vector. + Matrix V = svd.V(); // The V factor is a local dense matrix. + // $example off$ + Vector[] collectPartitions = (Vector[]) U.rows().collect(); + System.out.println("U factor is:"); + for (Vector vector : collectPartitions) { + System.out.println("\t" + vector); + } + System.out.println("Singular values are: " + s); + System.out.println("V factor is:\n" + V); + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVMWithSGDExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVMWithSGDExample.java new file mode 100644 index 0000000..866a221 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSVMWithSGDExample.java @@ -0,0 +1,75 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.SparkContext; + +// $example on$ +import scala.Tuple2; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.classification.SVMModel; +import org.apache.spark.mllib.classification.SVMWithSGD; +import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics; +import org.apache.spark.mllib.regression.LabeledPoint; +import org.apache.spark.mllib.util.MLUtils; +// $example off$ + +/** + * Example for SVMWithSGD. + */ +public class JavaSVMWithSGDExample { + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("JavaSVMWithSGDExample"); + SparkContext sc = new SparkContext(conf); + // $example on$ + String path = "data/mllib/sample_libsvm_data.txt"; + JavaRDD data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD(); + + // Split initial RDD into two... [60% training data, 40% testing data]. + JavaRDD training = data.sample(false, 0.6, 11L); + training.cache(); + JavaRDD test = data.subtract(training); + + // Run training algorithm to build the model. + int numIterations = 100; + SVMModel model = SVMWithSGD.train(training.rdd(), numIterations); + + // Clear the default threshold. + model.clearThreshold(); + + // Compute raw scores on the test set. + JavaRDD> scoreAndLabels = test.map(p -> + new Tuple2<>(model.predict(p.features()), p.label())); + + // Get evaluation metrics. + BinaryClassificationMetrics metrics = + new BinaryClassificationMetrics(JavaRDD.toRDD(scoreAndLabels)); + double auROC = metrics.areaUnderROC(); + + System.out.println("Area under ROC = " + auROC); + + // Save and load model + model.save(sc, "target/tmp/javaSVMWithSGDModel"); + SVMModel sameModel = SVMModel.load(sc, "target/tmp/javaSVMWithSGDModel"); + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java new file mode 100644 index 0000000..f9198e7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java @@ -0,0 +1,63 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.mllib.fpm.AssociationRules; +import org.apache.spark.mllib.fpm.FPGrowth; +import org.apache.spark.mllib.fpm.FPGrowthModel; +// $example off$ + +import org.apache.spark.SparkConf; + +public class JavaSimpleFPGrowth { + + public static void main(String[] args) { + SparkConf conf = new SparkConf().setAppName("FP-growth Example"); + JavaSparkContext sc = new JavaSparkContext(conf); + + // $example on$ + JavaRDD data = sc.textFile("data/mllib/sample_fpgrowth.txt"); + + JavaRDD> transactions = data.map(line -> Arrays.asList(line.split(" "))); + + FPGrowth fpg = new FPGrowth() + .setMinSupport(0.2) + .setNumPartitions(10); + FPGrowthModel model = fpg.run(transactions); + + for (FPGrowth.FreqItemset itemset: model.freqItemsets().toJavaRDD().collect()) { + System.out.println("[" + itemset.javaItems() + "], " + itemset.freq()); + } + + double minConfidence = 0.8; + for (AssociationRules.Rule rule + : model.generateAssociationRules(minConfidence).toJavaRDD().collect()) { + System.out.println( + rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence()); + } + // $example off$ + + sc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStratifiedSamplingExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStratifiedSamplingExample.java new file mode 100644 index 0000000..b06b2cc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStratifiedSamplingExample.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import com.google.common.collect.ImmutableMap; +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; + +// $example on$ +import java.util.*; + +import scala.Tuple2; + +import org.apache.spark.api.java.JavaPairRDD; +// $example off$ + +public class JavaStratifiedSamplingExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaStratifiedSamplingExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + List> list = Arrays.asList( + new Tuple2<>(1, 'a'), + new Tuple2<>(1, 'b'), + new Tuple2<>(2, 'c'), + new Tuple2<>(2, 'd'), + new Tuple2<>(2, 'e'), + new Tuple2<>(3, 'f') + ); + + JavaPairRDD data = jsc.parallelizePairs(list); + + // specify the exact fraction desired from each key Map + ImmutableMap fractions = ImmutableMap.of(1, 0.1, 2, 0.6, 3, 0.3); + + // Get an approximate sample from each stratum + JavaPairRDD approxSample = data.sampleByKey(false, fractions); + // Get an exact sample from each stratum + JavaPairRDD exactSample = data.sampleByKeyExact(false, fractions); + // $example off$ + + System.out.println("approxSample size is " + approxSample.collect().size()); + for (Tuple2 t : approxSample.collect()) { + System.out.println(t._1() + " " + t._2()); + } + + System.out.println("exactSample size is " + exactSample.collect().size()); + for (Tuple2 t : exactSample.collect()) { + System.out.println(t._1() + " " + t._2()); + } + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStreamingTestExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStreamingTestExample.java new file mode 100644 index 0000000..4be702c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaStreamingTestExample.java @@ -0,0 +1,104 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +// $example on$ +import org.apache.spark.mllib.stat.test.BinarySample; +import org.apache.spark.mllib.stat.test.StreamingTest; +import org.apache.spark.mllib.stat.test.StreamingTestResult; +// $example off$ +import org.apache.spark.SparkConf; +import org.apache.spark.streaming.Duration; +import org.apache.spark.streaming.Seconds; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; +import org.apache.spark.util.Utils; + + +/** + * Perform streaming testing using Welch's 2-sample t-test on a stream of data, where the data + * stream arrives as text files in a directory. Stops when the two groups are statistically + * significant (p-value < 0.05) or after a user-specified timeout in number of batches is exceeded. + * + * The rows of the text files must be in the form `Boolean, Double`. For example: + * false, -3.92 + * true, 99.32 + * + * Usage: + * JavaStreamingTestExample + * + * To run on your local machine using the directory `dataDir` with 5 seconds between each batch and + * a timeout after 100 insignificant batches, call: + * $ bin/run-example mllib.JavaStreamingTestExample dataDir 5 100 + * + * As you add text files to `dataDir` the significance test wil continually update every + * `batchDuration` seconds until the test becomes significant (p-value < 0.05) or the number of + * batches processed exceeds `numBatchesTimeout`. + */ +public class JavaStreamingTestExample { + + private static int timeoutCounter = 0; + + public static void main(String[] args) throws Exception { + if (args.length != 3) { + System.err.println("Usage: JavaStreamingTestExample " + + " "); + System.exit(1); + } + + String dataDir = args[0]; + Duration batchDuration = Seconds.apply(Long.parseLong(args[1])); + int numBatchesTimeout = Integer.parseInt(args[2]); + + SparkConf conf = new SparkConf().setMaster("local").setAppName("StreamingTestExample"); + JavaStreamingContext ssc = new JavaStreamingContext(conf, batchDuration); + + ssc.checkpoint(Utils.createTempDir(System.getProperty("java.io.tmpdir"), "spark").toString()); + + // $example on$ + JavaDStream data = ssc.textFileStream(dataDir).map(line -> { + String[] ts = line.split(","); + boolean label = Boolean.parseBoolean(ts[0]); + double value = Double.parseDouble(ts[1]); + return new BinarySample(label, value); + }); + + StreamingTest streamingTest = new StreamingTest() + .setPeacePeriod(0) + .setWindowSize(0) + .setTestMethod("welch"); + + JavaDStream out = streamingTest.registerStream(data); + out.print(); + // $example off$ + + // Stop processing if test becomes significant or we time out + timeoutCounter = numBatchesTimeout; + + out.foreachRDD(rdd -> { + timeoutCounter -= 1; + boolean anySignificant = !rdd.filter(v -> v.pValue() < 0.05).isEmpty(); + if (timeoutCounter <= 0 || anySignificant) { + rdd.context().stop(); + } + }); + + ssc.start(); + ssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSummaryStatisticsExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSummaryStatisticsExample.java new file mode 100644 index 0000000..278706b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/mllib/JavaSummaryStatisticsExample.java @@ -0,0 +1,56 @@ +/* + * 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. + */ + +package org.apache.spark.examples.mllib; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +// $example on$ +import java.util.Arrays; + +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.mllib.linalg.Vector; +import org.apache.spark.mllib.linalg.Vectors; +import org.apache.spark.mllib.stat.MultivariateStatisticalSummary; +import org.apache.spark.mllib.stat.Statistics; +// $example off$ + +public class JavaSummaryStatisticsExample { + public static void main(String[] args) { + + SparkConf conf = new SparkConf().setAppName("JavaSummaryStatisticsExample"); + JavaSparkContext jsc = new JavaSparkContext(conf); + + // $example on$ + JavaRDD mat = jsc.parallelize( + Arrays.asList( + Vectors.dense(1.0, 10.0, 100.0), + Vectors.dense(2.0, 20.0, 200.0), + Vectors.dense(3.0, 30.0, 300.0) + ) + ); // an RDD of Vectors + + // Compute column summary statistics. + MultivariateStatisticalSummary summary = Statistics.colStats(mat.rdd()); + System.out.println(summary.mean()); // a dense vector containing the mean value for each column + System.out.println(summary.variance()); // column-wise variance + System.out.println(summary.numNonzeros()); // number of nonzeros in each column + // $example off$ + + jsc.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSQLDataSourceExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSQLDataSourceExample.java new file mode 100644 index 0000000..c096054 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSQLDataSourceExample.java @@ -0,0 +1,499 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql; + +// $example on:schema_merging$ +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +// $example off:schema_merging$ +import java.util.Properties; + +// $example on:basic_parquet_example$ +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Encoders; +// $example on:schema_merging$ +// $example on:json_dataset$ +// $example on:csv_dataset$ +// $example on:text_dataset$ +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off:text_dataset$ +// $example off:csv_dataset$ +// $example off:json_dataset$ +// $example off:schema_merging$ +// $example off:basic_parquet_example$ +import org.apache.spark.sql.SparkSession; + +public class JavaSQLDataSourceExample { + + // $example on:schema_merging$ + public static class Square implements Serializable { + private int value; + private int square; + + // Getters and setters... + // $example off:schema_merging$ + public int getValue() { + return value; + } + + public void setValue(int value) { + this.value = value; + } + + public int getSquare() { + return square; + } + + public void setSquare(int square) { + this.square = square; + } + // $example on:schema_merging$ + } + // $example off:schema_merging$ + + // $example on:schema_merging$ + public static class Cube implements Serializable { + private int value; + private int cube; + + // Getters and setters... + // $example off:schema_merging$ + public int getValue() { + return value; + } + + public void setValue(int value) { + this.value = value; + } + + public int getCube() { + return cube; + } + + public void setCube(int cube) { + this.cube = cube; + } + // $example on:schema_merging$ + } + // $example off:schema_merging$ + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("Java Spark SQL data sources example") + .config("spark.some.config.option", "some-value") + .getOrCreate(); + + runBasicDataSourceExample(spark); + runGenericFileSourceOptionsExample(spark); + runBasicParquetExample(spark); + runParquetSchemaMergingExample(spark); + runJsonDatasetExample(spark); + runCsvDatasetExample(spark); + runTextDatasetExample(spark); + runJdbcDatasetExample(spark); + + spark.stop(); + } + + private static void runGenericFileSourceOptionsExample(SparkSession spark) { + // $example on:ignore_corrupt_files$ + // enable ignore corrupt files via the data source option + // dir1/file3.json is corrupt from parquet's view + Dataset testCorruptDF0 = spark.read().option("ignoreCorruptFiles", "true").parquet( + "examples/src/main/resources/dir1/", + "examples/src/main/resources/dir1/dir2/"); + testCorruptDF0.show(); + // +-------------+ + // | file| + // +-------------+ + // |file1.parquet| + // |file2.parquet| + // +-------------+ + + // enable ignore corrupt files via the configuration + spark.sql("set spark.sql.files.ignoreCorruptFiles=true"); + // dir1/file3.json is corrupt from parquet's view + Dataset testCorruptDF1 = spark.read().parquet( + "examples/src/main/resources/dir1/", + "examples/src/main/resources/dir1/dir2/"); + testCorruptDF1.show(); + // +-------------+ + // | file| + // +-------------+ + // |file1.parquet| + // |file2.parquet| + // +-------------+ + // $example off:ignore_corrupt_files$ + // $example on:recursive_file_lookup$ + Dataset recursiveLoadedDF = spark.read().format("parquet") + .option("recursiveFileLookup", "true") + .load("examples/src/main/resources/dir1"); + recursiveLoadedDF.show(); + // +-------------+ + // | file| + // +-------------+ + // |file1.parquet| + // |file2.parquet| + // +-------------+ + // $example off:recursive_file_lookup$ + spark.sql("set spark.sql.files.ignoreCorruptFiles=false"); + // $example on:load_with_path_glob_filter$ + Dataset testGlobFilterDF = spark.read().format("parquet") + .option("pathGlobFilter", "*.parquet") // json file should be filtered out + .load("examples/src/main/resources/dir1"); + testGlobFilterDF.show(); + // +-------------+ + // | file| + // +-------------+ + // |file1.parquet| + // +-------------+ + // $example off:load_with_path_glob_filter$ + // $example on:load_with_modified_time_filter$ + Dataset beforeFilterDF = spark.read().format("parquet") + // Only load files modified before 7/1/2020 at 05:30 + .option("modifiedBefore", "2020-07-01T05:30:00") + // Only load files modified after 6/1/2020 at 05:30 + .option("modifiedAfter", "2020-06-01T05:30:00") + // Interpret both times above relative to CST timezone + .option("timeZone", "CST") + .load("examples/src/main/resources/dir1"); + beforeFilterDF.show(); + // +-------------+ + // | file| + // +-------------+ + // |file1.parquet| + // +-------------+ + // $example off:load_with_modified_time_filter$ + } + + private static void runBasicDataSourceExample(SparkSession spark) { + // $example on:generic_load_save_functions$ + Dataset usersDF = spark.read().load("examples/src/main/resources/users.parquet"); + usersDF.select("name", "favorite_color").write().save("namesAndFavColors.parquet"); + // $example off:generic_load_save_functions$ + // $example on:manual_load_options$ + Dataset peopleDF = + spark.read().format("json").load("examples/src/main/resources/people.json"); + peopleDF.select("name", "age").write().format("parquet").save("namesAndAges.parquet"); + // $example off:manual_load_options$ + // $example on:manual_load_options_csv$ + Dataset peopleDFCsv = spark.read().format("csv") + .option("sep", ";") + .option("inferSchema", "true") + .option("header", "true") + .load("examples/src/main/resources/people.csv"); + // $example off:manual_load_options_csv$ + // $example on:manual_save_options_orc$ + usersDF.write().format("orc") + .option("orc.bloom.filter.columns", "favorite_color") + .option("orc.dictionary.key.threshold", "1.0") + .option("orc.column.encoding.direct", "name") + .save("users_with_options.orc"); + // $example off:manual_save_options_orc$ + // $example on:manual_save_options_parquet$ + usersDF.write().format("parquet") + .option("parquet.bloom.filter.enabled#favorite_color", "true") + .option("parquet.bloom.filter.expected.ndv#favorite_color", "1000000") + .option("parquet.enable.dictionary", "true") + .option("parquet.page.write-checksum.enabled", "false") + .save("users_with_options.parquet"); + // $example off:manual_save_options_parquet$ + // $example on:direct_sql$ + Dataset sqlDF = + spark.sql("SELECT * FROM parquet.`examples/src/main/resources/users.parquet`"); + // $example off:direct_sql$ + // $example on:write_sorting_and_bucketing$ + peopleDF.write().bucketBy(42, "name").sortBy("age").saveAsTable("people_bucketed"); + // $example off:write_sorting_and_bucketing$ + // $example on:write_partitioning$ + usersDF + .write() + .partitionBy("favorite_color") + .format("parquet") + .save("namesPartByColor.parquet"); + // $example off:write_partitioning$ + // $example on:write_partition_and_bucket$ + usersDF + .write() + .partitionBy("favorite_color") + .bucketBy(42, "name") + .saveAsTable("users_partitioned_bucketed"); + // $example off:write_partition_and_bucket$ + + spark.sql("DROP TABLE IF EXISTS people_bucketed"); + spark.sql("DROP TABLE IF EXISTS users_partitioned_bucketed"); + } + + private static void runBasicParquetExample(SparkSession spark) { + // $example on:basic_parquet_example$ + Dataset peopleDF = spark.read().json("examples/src/main/resources/people.json"); + + // DataFrames can be saved as Parquet files, maintaining the schema information + peopleDF.write().parquet("people.parquet"); + + // Read in the Parquet file created above. + // Parquet files are self-describing so the schema is preserved + // The result of loading a parquet file is also a DataFrame + Dataset parquetFileDF = spark.read().parquet("people.parquet"); + + // Parquet files can also be used to create a temporary view and then used in SQL statements + parquetFileDF.createOrReplaceTempView("parquetFile"); + Dataset namesDF = spark.sql("SELECT name FROM parquetFile WHERE age BETWEEN 13 AND 19"); + Dataset namesDS = namesDF.map( + (MapFunction) row -> "Name: " + row.getString(0), + Encoders.STRING()); + namesDS.show(); + // +------------+ + // | value| + // +------------+ + // |Name: Justin| + // +------------+ + // $example off:basic_parquet_example$ + } + + private static void runParquetSchemaMergingExample(SparkSession spark) { + // $example on:schema_merging$ + List squares = new ArrayList<>(); + for (int value = 1; value <= 5; value++) { + Square square = new Square(); + square.setValue(value); + square.setSquare(value * value); + squares.add(square); + } + + // Create a simple DataFrame, store into a partition directory + Dataset squaresDF = spark.createDataFrame(squares, Square.class); + squaresDF.write().parquet("data/test_table/key=1"); + + List cubes = new ArrayList<>(); + for (int value = 6; value <= 10; value++) { + Cube cube = new Cube(); + cube.setValue(value); + cube.setCube(value * value * value); + cubes.add(cube); + } + + // Create another DataFrame in a new partition directory, + // adding a new column and dropping an existing column + Dataset cubesDF = spark.createDataFrame(cubes, Cube.class); + cubesDF.write().parquet("data/test_table/key=2"); + + // Read the partitioned table + Dataset mergedDF = spark.read().option("mergeSchema", true).parquet("data/test_table"); + mergedDF.printSchema(); + + // The final schema consists of all 3 columns in the Parquet files together + // with the partitioning column appeared in the partition directory paths + // root + // |-- value: int (nullable = true) + // |-- square: int (nullable = true) + // |-- cube: int (nullable = true) + // |-- key: int (nullable = true) + // $example off:schema_merging$ + } + + private static void runJsonDatasetExample(SparkSession spark) { + // $example on:json_dataset$ + // A JSON dataset is pointed to by path. + // The path can be either a single text file or a directory storing text files + Dataset people = spark.read().json("examples/src/main/resources/people.json"); + + // The inferred schema can be visualized using the printSchema() method + people.printSchema(); + // root + // |-- age: long (nullable = true) + // |-- name: string (nullable = true) + + // Creates a temporary view using the DataFrame + people.createOrReplaceTempView("people"); + + // SQL statements can be run by using the sql methods provided by spark + Dataset namesDF = spark.sql("SELECT name FROM people WHERE age BETWEEN 13 AND 19"); + namesDF.show(); + // +------+ + // | name| + // +------+ + // |Justin| + // +------+ + + // Alternatively, a DataFrame can be created for a JSON dataset represented by + // a Dataset storing one JSON object per string. + List jsonData = Arrays.asList( + "{\"name\":\"Yin\",\"address\":{\"city\":\"Columbus\",\"state\":\"Ohio\"}}"); + Dataset anotherPeopleDataset = spark.createDataset(jsonData, Encoders.STRING()); + Dataset anotherPeople = spark.read().json(anotherPeopleDataset); + anotherPeople.show(); + // +---------------+----+ + // | address|name| + // +---------------+----+ + // |[Columbus,Ohio]| Yin| + // +---------------+----+ + // $example off:json_dataset$ + } + + private static void runCsvDatasetExample(SparkSession spark) { + // $example on:csv_dataset$ + // A CSV dataset is pointed to by path. + // The path can be either a single CSV file or a directory of CSV files + String path = "examples/src/main/resources/people.csv"; + + Dataset df = spark.read().csv(path); + df.show(); + // +------------------+ + // | _c0| + // +------------------+ + // | name;age;job| + // |Jorge;30;Developer| + // | Bob;32;Developer| + // +------------------+ + + // Read a csv with delimiter, the default delimiter is "," + Dataset df2 = spark.read().option("delimiter", ";").csv(path); + df2.show(); + // +-----+---+---------+ + // | _c0|_c1| _c2| + // +-----+---+---------+ + // | name|age| job| + // |Jorge| 30|Developer| + // | Bob| 32|Developer| + // +-----+---+---------+ + + // Read a csv with delimiter and a header + Dataset df3 = spark.read().option("delimiter", ";").option("header", "true").csv(path); + df3.show(); + // +-----+---+---------+ + // | name|age| job| + // +-----+---+---------+ + // |Jorge| 30|Developer| + // | Bob| 32|Developer| + // +-----+---+---------+ + + // You can also use options() to use multiple options + java.util.Map optionsMap = new java.util.HashMap(); + optionsMap.put("delimiter",";"); + optionsMap.put("header","true"); + Dataset df4 = spark.read().options(optionsMap).csv(path); + + // "output" is a folder which contains multiple csv files and a _SUCCESS file. + df3.write().csv("output"); + + // Read all files in a folder, please make sure only CSV files should present in the folder. + String folderPath = "examples/src/main/resources"; + Dataset df5 = spark.read().csv(folderPath); + df5.show(); + // Wrong schema because non-CSV files are read + // +-----------+ + // | _c0| + // +-----------+ + // |238val_238| + // | 86val_86| + // |311val_311| + // | 27val_27| + // |165val_165| + // +-----------+ + + // $example off:csv_dataset$ + } + + private static void runTextDatasetExample(SparkSession spark) { + // $example on:text_dataset$ + // A text dataset is pointed to by path. + // The path can be either a single text file or a directory of text files + String path = "examples/src/main/resources/people.txt"; + + Dataset df1 = spark.read().text(path); + df1.show(); + // +-----------+ + // | value| + // +-----------+ + // |Michael, 29| + // | Andy, 30| + // | Justin, 19| + // +-----------+ + + // You can use 'lineSep' option to define the line separator. + // The line separator handles all `\r`, `\r\n` and `\n` by default. + Dataset df2 = spark.read().option("lineSep", ",").text(path); + df2.show(); + // +-----------+ + // | value| + // +-----------+ + // | Michael| + // | 29\nAndy| + // | 30\nJustin| + // | 19\n| + // +-----------+ + + // You can also use 'wholetext' option to read each input file as a single row. + Dataset df3 = spark.read().option("wholetext", "true").text(path); + df3.show(); + // +--------------------+ + // | value| + // +--------------------+ + // |Michael, 29\nAndy...| + // +--------------------+ + + // "output" is a folder which contains multiple text files and a _SUCCESS file. + df1.write().text("output"); + + // You can specify the compression format using the 'compression' option. + df1.write().option("compression", "gzip").text("output_compressed"); + + // $example off:text_dataset$ + } + + private static void runJdbcDatasetExample(SparkSession spark) { + // $example on:jdbc_dataset$ + // Note: JDBC loading and saving can be achieved via either the load/save or jdbc methods + // Loading data from a JDBC source + Dataset jdbcDF = spark.read() + .format("jdbc") + .option("url", "jdbc:postgresql:dbserver") + .option("dbtable", "schema.tablename") + .option("user", "username") + .option("password", "password") + .load(); + + Properties connectionProperties = new Properties(); + connectionProperties.put("user", "username"); + connectionProperties.put("password", "password"); + Dataset jdbcDF2 = spark.read() + .jdbc("jdbc:postgresql:dbserver", "schema.tablename", connectionProperties); + + // Saving data to a JDBC source + jdbcDF.write() + .format("jdbc") + .option("url", "jdbc:postgresql:dbserver") + .option("dbtable", "schema.tablename") + .option("user", "username") + .option("password", "password") + .save(); + + jdbcDF2.write() + .jdbc("jdbc:postgresql:dbserver", "schema.tablename", connectionProperties); + + // Specifying create table column data types on write + jdbcDF.write() + .option("createTableColumnTypes", "name CHAR(64), comments VARCHAR(1024)") + .jdbc("jdbc:postgresql:dbserver", "schema.tablename", connectionProperties); + // $example off:jdbc_dataset$ + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSparkSQLExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSparkSQLExample.java new file mode 100644 index 0000000..86a9045 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaSparkSQLExample.java @@ -0,0 +1,344 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql; + +// $example on:programmatic_schema$ +import java.util.ArrayList; +import java.util.List; +// $example off:programmatic_schema$ +// $example on:create_ds$ +import java.util.Arrays; +import java.util.Collections; +import java.io.Serializable; +// $example off:create_ds$ + +// $example on:schema_inferring$ +// $example on:programmatic_schema$ +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.function.Function; +// $example off:programmatic_schema$ +// $example on:create_ds$ +import org.apache.spark.api.java.function.MapFunction; +// $example on:create_df$ +// $example on:run_sql$ +// $example on:programmatic_schema$ +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +// $example off:programmatic_schema$ +// $example off:create_df$ +// $example off:run_sql$ +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; +// $example off:create_ds$ +// $example off:schema_inferring$ +import org.apache.spark.sql.RowFactory; +// $example on:init_session$ +import org.apache.spark.sql.SparkSession; +// $example off:init_session$ +// $example on:programmatic_schema$ +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +// $example off:programmatic_schema$ +import org.apache.spark.sql.AnalysisException; + +// $example on:untyped_ops$ +// col("...") is preferable to df.col("...") +import static org.apache.spark.sql.functions.col; +// $example off:untyped_ops$ + +public class JavaSparkSQLExample { + // $example on:create_ds$ + public static class Person implements Serializable { + private String name; + private long age; + + public String getName() { + return name; + } + + public void setName(String name) { + this.name = name; + } + + public long getAge() { + return age; + } + + public void setAge(long age) { + this.age = age; + } + } + // $example off:create_ds$ + + public static void main(String[] args) throws AnalysisException { + // $example on:init_session$ + SparkSession spark = SparkSession + .builder() + .appName("Java Spark SQL basic example") + .config("spark.some.config.option", "some-value") + .getOrCreate(); + // $example off:init_session$ + + runBasicDataFrameExample(spark); + runDatasetCreationExample(spark); + runInferSchemaExample(spark); + runProgrammaticSchemaExample(spark); + + spark.stop(); + } + + private static void runBasicDataFrameExample(SparkSession spark) throws AnalysisException { + // $example on:create_df$ + Dataset df = spark.read().json("examples/src/main/resources/people.json"); + + // Displays the content of the DataFrame to stdout + df.show(); + // +----+-------+ + // | age| name| + // +----+-------+ + // |null|Michael| + // | 30| Andy| + // | 19| Justin| + // +----+-------+ + // $example off:create_df$ + + // $example on:untyped_ops$ + // Print the schema in a tree format + df.printSchema(); + // root + // |-- age: long (nullable = true) + // |-- name: string (nullable = true) + + // Select only the "name" column + df.select("name").show(); + // +-------+ + // | name| + // +-------+ + // |Michael| + // | Andy| + // | Justin| + // +-------+ + + // Select everybody, but increment the age by 1 + df.select(col("name"), col("age").plus(1)).show(); + // +-------+---------+ + // | name|(age + 1)| + // +-------+---------+ + // |Michael| null| + // | Andy| 31| + // | Justin| 20| + // +-------+---------+ + + // Select people older than 21 + df.filter(col("age").gt(21)).show(); + // +---+----+ + // |age|name| + // +---+----+ + // | 30|Andy| + // +---+----+ + + // Count people by age + df.groupBy("age").count().show(); + // +----+-----+ + // | age|count| + // +----+-----+ + // | 19| 1| + // |null| 1| + // | 30| 1| + // +----+-----+ + // $example off:untyped_ops$ + + // $example on:run_sql$ + // Register the DataFrame as a SQL temporary view + df.createOrReplaceTempView("people"); + + Dataset sqlDF = spark.sql("SELECT * FROM people"); + sqlDF.show(); + // +----+-------+ + // | age| name| + // +----+-------+ + // |null|Michael| + // | 30| Andy| + // | 19| Justin| + // +----+-------+ + // $example off:run_sql$ + + // $example on:global_temp_view$ + // Register the DataFrame as a global temporary view + df.createGlobalTempView("people"); + + // Global temporary view is tied to a system preserved database `global_temp` + spark.sql("SELECT * FROM global_temp.people").show(); + // +----+-------+ + // | age| name| + // +----+-------+ + // |null|Michael| + // | 30| Andy| + // | 19| Justin| + // +----+-------+ + + // Global temporary view is cross-session + spark.newSession().sql("SELECT * FROM global_temp.people").show(); + // +----+-------+ + // | age| name| + // +----+-------+ + // |null|Michael| + // | 30| Andy| + // | 19| Justin| + // +----+-------+ + // $example off:global_temp_view$ + } + + private static void runDatasetCreationExample(SparkSession spark) { + // $example on:create_ds$ + // Create an instance of a Bean class + Person person = new Person(); + person.setName("Andy"); + person.setAge(32); + + // Encoders are created for Java beans + Encoder personEncoder = Encoders.bean(Person.class); + Dataset javaBeanDS = spark.createDataset( + Collections.singletonList(person), + personEncoder + ); + javaBeanDS.show(); + // +---+----+ + // |age|name| + // +---+----+ + // | 32|Andy| + // +---+----+ + + // Encoders for most common types are provided in class Encoders + Encoder longEncoder = Encoders.LONG(); + Dataset primitiveDS = spark.createDataset(Arrays.asList(1L, 2L, 3L), longEncoder); + Dataset transformedDS = primitiveDS.map( + (MapFunction) value -> value + 1L, + longEncoder); + transformedDS.collect(); // Returns [2, 3, 4] + + // DataFrames can be converted to a Dataset by providing a class. Mapping based on name + String path = "examples/src/main/resources/people.json"; + Dataset peopleDS = spark.read().json(path).as(personEncoder); + peopleDS.show(); + // +----+-------+ + // | age| name| + // +----+-------+ + // |null|Michael| + // | 30| Andy| + // | 19| Justin| + // +----+-------+ + // $example off:create_ds$ + } + + private static void runInferSchemaExample(SparkSession spark) { + // $example on:schema_inferring$ + // Create an RDD of Person objects from a text file + JavaRDD peopleRDD = spark.read() + .textFile("examples/src/main/resources/people.txt") + .javaRDD() + .map(line -> { + String[] parts = line.split(","); + Person person = new Person(); + person.setName(parts[0]); + person.setAge(Integer.parseInt(parts[1].trim())); + return person; + }); + + // Apply a schema to an RDD of JavaBeans to get a DataFrame + Dataset peopleDF = spark.createDataFrame(peopleRDD, Person.class); + // Register the DataFrame as a temporary view + peopleDF.createOrReplaceTempView("people"); + + // SQL statements can be run by using the sql methods provided by spark + Dataset teenagersDF = spark.sql("SELECT name FROM people WHERE age BETWEEN 13 AND 19"); + + // The columns of a row in the result can be accessed by field index + Encoder stringEncoder = Encoders.STRING(); + Dataset teenagerNamesByIndexDF = teenagersDF.map( + (MapFunction) row -> "Name: " + row.getString(0), + stringEncoder); + teenagerNamesByIndexDF.show(); + // +------------+ + // | value| + // +------------+ + // |Name: Justin| + // +------------+ + + // or by field name + Dataset teenagerNamesByFieldDF = teenagersDF.map( + (MapFunction) row -> "Name: " + row.getAs("name"), + stringEncoder); + teenagerNamesByFieldDF.show(); + // +------------+ + // | value| + // +------------+ + // |Name: Justin| + // +------------+ + // $example off:schema_inferring$ + } + + private static void runProgrammaticSchemaExample(SparkSession spark) { + // $example on:programmatic_schema$ + // Create an RDD + JavaRDD peopleRDD = spark.sparkContext() + .textFile("examples/src/main/resources/people.txt", 1) + .toJavaRDD(); + + // The schema is encoded in a string + String schemaString = "name age"; + + // Generate the schema based on the string of schema + List fields = new ArrayList<>(); + for (String fieldName : schemaString.split(" ")) { + StructField field = DataTypes.createStructField(fieldName, DataTypes.StringType, true); + fields.add(field); + } + StructType schema = DataTypes.createStructType(fields); + + // Convert records of the RDD (people) to Rows + JavaRDD rowRDD = peopleRDD.map((Function) record -> { + String[] attributes = record.split(","); + return RowFactory.create(attributes[0], attributes[1].trim()); + }); + + // Apply the schema to the RDD + Dataset peopleDataFrame = spark.createDataFrame(rowRDD, schema); + + // Creates a temporary view using the DataFrame + peopleDataFrame.createOrReplaceTempView("people"); + + // SQL can be run over a temporary view created using DataFrames + Dataset results = spark.sql("SELECT name FROM people"); + + // The results of SQL queries are DataFrames and support all the normal RDD operations + // The columns of a row in the result can be accessed by field index or by field name + Dataset namesDS = results.map( + (MapFunction) row -> "Name: " + row.getString(0), + Encoders.STRING()); + namesDS.show(); + // +-------------+ + // | value| + // +-------------+ + // |Name: Michael| + // | Name: Andy| + // | Name: Justin| + // +-------------+ + // $example off:programmatic_schema$ + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedScalar.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedScalar.java new file mode 100644 index 0000000..e420368 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedScalar.java @@ -0,0 +1,90 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql; + +// $example on:udf_scalar$ +import org.apache.spark.sql.*; +import org.apache.spark.sql.api.java.UDF1; +import org.apache.spark.sql.expressions.UserDefinedFunction; +import static org.apache.spark.sql.functions.udf; +import org.apache.spark.sql.types.DataTypes; +// $example off:udf_scalar$ + +public class JavaUserDefinedScalar { + + public static void main(String[] args) { + + // $example on:udf_scalar$ + SparkSession spark = SparkSession + .builder() + .appName("Java Spark SQL UDF scalar example") + .getOrCreate(); + + // Define and register a zero-argument non-deterministic UDF + // UDF is deterministic by default, i.e. produces the same result for the same input. + UserDefinedFunction random = udf( + () -> Math.random(), DataTypes.DoubleType + ); + random.asNondeterministic(); + spark.udf().register("random", random); + spark.sql("SELECT random()").show(); + // +-------+ + // |UDF() | + // +-------+ + // |xxxxxxx| + // +-------+ + + // Define and register a one-argument UDF + spark.udf().register("plusOne", + (UDF1) x -> x + 1, DataTypes.IntegerType); + spark.sql("SELECT plusOne(5)").show(); + // +----------+ + // |plusOne(5)| + // +----------+ + // | 6| + // +----------+ + + // Define and register a two-argument UDF + UserDefinedFunction strLen = udf( + (String s, Integer x) -> s.length() + x, DataTypes.IntegerType + ); + spark.udf().register("strLen", strLen); + spark.sql("SELECT strLen('test', 1)").show(); + // +------------+ + // |UDF(test, 1)| + // +------------+ + // | 5| + // +------------+ + + // UDF in a WHERE clause + spark.udf().register("oneArgFilter", + (UDF1) x -> x > 5, DataTypes.BooleanType); + spark.range(1, 10).createOrReplaceTempView("test"); + spark.sql("SELECT * FROM test WHERE oneArgFilter(id)").show(); + // +---+ + // | id| + // +---+ + // | 6| + // | 7| + // | 8| + // | 9| + // +---+ + + // $example off:udf_scalar$ + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedTypedAggregation.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedTypedAggregation.java new file mode 100644 index 0000000..8926210 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedTypedAggregation.java @@ -0,0 +1,166 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql; + +// $example on:typed_custom_aggregation$ +import java.io.Serializable; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.TypedColumn; +import org.apache.spark.sql.expressions.Aggregator; +// $example off:typed_custom_aggregation$ + +public class JavaUserDefinedTypedAggregation { + + // $example on:typed_custom_aggregation$ + public static class Employee implements Serializable { + private String name; + private long salary; + + // Constructors, getters, setters... + // $example off:typed_custom_aggregation$ + public String getName() { + return name; + } + + public void setName(String name) { + this.name = name; + } + + public long getSalary() { + return salary; + } + + public void setSalary(long salary) { + this.salary = salary; + } + // $example on:typed_custom_aggregation$ + } + + public static class Average implements Serializable { + private long sum; + private long count; + + // Constructors, getters, setters... + // $example off:typed_custom_aggregation$ + public Average() { + } + + public Average(long sum, long count) { + this.sum = sum; + this.count = count; + } + + public long getSum() { + return sum; + } + + public void setSum(long sum) { + this.sum = sum; + } + + public long getCount() { + return count; + } + + public void setCount(long count) { + this.count = count; + } + // $example on:typed_custom_aggregation$ + } + + public static class MyAverage extends Aggregator { + // A zero value for this aggregation. Should satisfy the property that any b + zero = b + @Override + public Average zero() { + return new Average(0L, 0L); + } + // Combine two values to produce a new value. For performance, the function may modify `buffer` + // and return it instead of constructing a new object + @Override + public Average reduce(Average buffer, Employee employee) { + long newSum = buffer.getSum() + employee.getSalary(); + long newCount = buffer.getCount() + 1; + buffer.setSum(newSum); + buffer.setCount(newCount); + return buffer; + } + // Merge two intermediate values + @Override + public Average merge(Average b1, Average b2) { + long mergedSum = b1.getSum() + b2.getSum(); + long mergedCount = b1.getCount() + b2.getCount(); + b1.setSum(mergedSum); + b1.setCount(mergedCount); + return b1; + } + // Transform the output of the reduction + @Override + public Double finish(Average reduction) { + return ((double) reduction.getSum()) / reduction.getCount(); + } + // Specifies the Encoder for the intermediate value type + @Override + public Encoder bufferEncoder() { + return Encoders.bean(Average.class); + } + // Specifies the Encoder for the final output value type + @Override + public Encoder outputEncoder() { + return Encoders.DOUBLE(); + } + } + // $example off:typed_custom_aggregation$ + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("Java Spark SQL user-defined Datasets aggregation example") + .getOrCreate(); + + // $example on:typed_custom_aggregation$ + Encoder employeeEncoder = Encoders.bean(Employee.class); + String path = "examples/src/main/resources/employees.json"; + Dataset ds = spark.read().json(path).as(employeeEncoder); + ds.show(); + // +-------+------+ + // | name|salary| + // +-------+------+ + // |Michael| 3000| + // | Andy| 4500| + // | Justin| 3500| + // | Berta| 4000| + // +-------+------+ + + MyAverage myAverage = new MyAverage(); + // Convert the function to a `TypedColumn` and give it a name + TypedColumn averageSalary = myAverage.toColumn().name("average_salary"); + Dataset result = ds.select(averageSalary); + result.show(); + // +--------------+ + // |average_salary| + // +--------------+ + // | 3750.0| + // +--------------+ + // $example off:typed_custom_aggregation$ + spark.stop(); + } + +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedUntypedAggregation.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedUntypedAggregation.java new file mode 100644 index 0000000..6bfec79 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/JavaUserDefinedUntypedAggregation.java @@ -0,0 +1,142 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql; + +// $example on:untyped_custom_aggregation$ +import java.io.Serializable; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.expressions.Aggregator; +import org.apache.spark.sql.functions; +// $example off:untyped_custom_aggregation$ + +public class JavaUserDefinedUntypedAggregation { + + // $example on:untyped_custom_aggregation$ + public static class Average implements Serializable { + private long sum; + private long count; + + // Constructors, getters, setters... + // $example off:typed_custom_aggregation$ + public Average() { + } + + public Average(long sum, long count) { + this.sum = sum; + this.count = count; + } + + public long getSum() { + return sum; + } + + public void setSum(long sum) { + this.sum = sum; + } + + public long getCount() { + return count; + } + + public void setCount(long count) { + this.count = count; + } + // $example on:typed_custom_aggregation$ + } + + public static class MyAverage extends Aggregator { + // A zero value for this aggregation. Should satisfy the property that any b + zero = b + @Override + public Average zero() { + return new Average(0L, 0L); + } + // Combine two values to produce a new value. For performance, the function may modify `buffer` + // and return it instead of constructing a new object + @Override + public Average reduce(Average buffer, Long data) { + long newSum = buffer.getSum() + data; + long newCount = buffer.getCount() + 1; + buffer.setSum(newSum); + buffer.setCount(newCount); + return buffer; + } + // Merge two intermediate values + @Override + public Average merge(Average b1, Average b2) { + long mergedSum = b1.getSum() + b2.getSum(); + long mergedCount = b1.getCount() + b2.getCount(); + b1.setSum(mergedSum); + b1.setCount(mergedCount); + return b1; + } + // Transform the output of the reduction + @Override + public Double finish(Average reduction) { + return ((double) reduction.getSum()) / reduction.getCount(); + } + // Specifies the Encoder for the intermediate value type + @Override + public Encoder bufferEncoder() { + return Encoders.bean(Average.class); + } + // Specifies the Encoder for the final output value type + @Override + public Encoder outputEncoder() { + return Encoders.DOUBLE(); + } + } + // $example off:untyped_custom_aggregation$ + + public static void main(String[] args) { + SparkSession spark = SparkSession + .builder() + .appName("Java Spark SQL user-defined DataFrames aggregation example") + .getOrCreate(); + + // $example on:untyped_custom_aggregation$ + // Register the function to access it + spark.udf().register("myAverage", functions.udaf(new MyAverage(), Encoders.LONG())); + + Dataset df = spark.read().json("examples/src/main/resources/employees.json"); + df.createOrReplaceTempView("employees"); + df.show(); + // +-------+------+ + // | name|salary| + // +-------+------+ + // |Michael| 3000| + // | Andy| 4500| + // | Justin| 3500| + // | Berta| 4000| + // +-------+------+ + + Dataset result = spark.sql("SELECT myAverage(salary) as average_salary FROM employees"); + result.show(); + // +--------------+ + // |average_salary| + // +--------------+ + // | 3750.0| + // +--------------+ + // $example off:untyped_custom_aggregation$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/hive/JavaSparkHiveExample.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/hive/JavaSparkHiveExample.java new file mode 100644 index 0000000..575a463 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/hive/JavaSparkHiveExample.java @@ -0,0 +1,129 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql.hive; + +// $example on:spark_hive$ +import java.io.File; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; + +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +// $example off:spark_hive$ + +public class JavaSparkHiveExample { + + // $example on:spark_hive$ + public static class Record implements Serializable { + private int key; + private String value; + + public int getKey() { + return key; + } + + public void setKey(int key) { + this.key = key; + } + + public String getValue() { + return value; + } + + public void setValue(String value) { + this.value = value; + } + } + // $example off:spark_hive$ + + public static void main(String[] args) { + // $example on:spark_hive$ + // warehouseLocation points to the default location for managed databases and tables + String warehouseLocation = new File("spark-warehouse").getAbsolutePath(); + SparkSession spark = SparkSession + .builder() + .appName("Java Spark Hive Example") + .config("spark.sql.warehouse.dir", warehouseLocation) + .enableHiveSupport() + .getOrCreate(); + + spark.sql("CREATE TABLE IF NOT EXISTS src (key INT, value STRING) USING hive"); + spark.sql("LOAD DATA LOCAL INPATH 'examples/src/main/resources/kv1.txt' INTO TABLE src"); + + // Queries are expressed in HiveQL + spark.sql("SELECT * FROM src").show(); + // +---+-------+ + // |key| value| + // +---+-------+ + // |238|val_238| + // | 86| val_86| + // |311|val_311| + // ... + + // Aggregation queries are also supported. + spark.sql("SELECT COUNT(*) FROM src").show(); + // +--------+ + // |count(1)| + // +--------+ + // | 500 | + // +--------+ + + // The results of SQL queries are themselves DataFrames and support all normal functions. + Dataset sqlDF = spark.sql("SELECT key, value FROM src WHERE key < 10 ORDER BY key"); + + // The items in DataFrames are of type Row, which lets you to access each column by ordinal. + Dataset stringsDS = sqlDF.map( + (MapFunction) row -> "Key: " + row.get(0) + ", Value: " + row.get(1), + Encoders.STRING()); + stringsDS.show(); + // +--------------------+ + // | value| + // +--------------------+ + // |Key: 0, Value: val_0| + // |Key: 0, Value: val_0| + // |Key: 0, Value: val_0| + // ... + + // You can also use DataFrames to create temporary views within a SparkSession. + List records = new ArrayList<>(); + for (int key = 1; key < 100; key++) { + Record record = new Record(); + record.setKey(key); + record.setValue("val_" + key); + records.add(record); + } + Dataset recordsDF = spark.createDataFrame(records, Record.class); + recordsDF.createOrReplaceTempView("records"); + + // Queries can then join DataFrames data with data stored in Hive. + spark.sql("SELECT * FROM records r JOIN src s ON r.key = s.key").show(); + // +---+------+---+------+ + // |key| value|key| value| + // +---+------+---+------+ + // | 2| val_2| 2| val_2| + // | 2| val_2| 2| val_2| + // | 4| val_4| 4| val_4| + // ... + // $example off:spark_hive$ + + spark.stop(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredComplexSessionization.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredComplexSessionization.java new file mode 100644 index 0000000..4c1642c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredComplexSessionization.java @@ -0,0 +1,461 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql.streaming; + +import java.io.Serializable; +import java.sql.Timestamp; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import java.util.Spliterator; +import java.util.Spliterators; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import java.util.stream.StreamSupport; + +import org.apache.spark.api.java.function.FlatMapGroupsWithStateFunction; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.GroupState; +import org.apache.spark.sql.streaming.GroupStateTimeout; +import org.apache.spark.sql.streaming.OutputMode; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.types.StructType; + +import static org.apache.spark.sql.types.DataTypes.StringType; +import static org.apache.spark.sql.types.DataTypes.TimestampType; +import static org.apache.spark.sql.functions.*; + +/** + * Sessionize events in UTF8 encoded, '\n' delimited text received from the network. + * Each line composes an event, and the line should match to the json format. + *

+ * The schema of the event is following: + * - user_id: String + * - event_type: String + * - timestamp: Long + *

+ * The supported types are following: + * - NEW_EVENT + * - CLOSE_SESSION + *

+ * This example focuses to demonstrate the complex sessionization which uses two conditions + * on closing session; conditions are following: + * - No further event is provided for the user ID within 5 seconds + * - An event having CLOSE_SESSION as event_type is provided for the user ID + *

+ * Usage: JavaStructuredComplexSessionization + * and describe the TCP server that Structured Streaming + * would connect to receive data. + *

+ * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example sql.streaming.JavaStructuredComplexSessionization + * localhost 9999` + *

+ * Here's a set of events for example: + * + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 13} + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 10} + * {"user_id": "user1", "event_type": "CLOSE_SESSION", "timestamp": 15} + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 17} + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 19} + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 29} + * + * {"user_id": "user2", "event_type": "NEW_EVENT", "timestamp": 45} + * + * {"user_id": "user1", "event_type": "NEW_EVENT", "timestamp": 65} + * + * and results (the output can be split across micro-batches): + * + * +-----+----------+---------+ + * | id|durationMs|numEvents| + * +-----+----------+---------+ + * |user1| 5000| 3| + * |user1| 7000| 2| + * |user1| 5000| 1| + * |user2| 5000| 1| + * +-----+----------+---------+ + * (The last event is not reflected into output due to watermark.) + *

+ * Note that there're three different sessions for 'user1'. The events in first two sessions + * are occurred within gap duration for nearest events, but they don't compose a single session + * due to the event of CLOSE_SESSION. + *

+ * Also note that the implementation is simplified one. This example doesn't address + * - UPDATE MODE (the semantic is not clear for session window with event time processing) + * - partial merge (events in session which are earlier than watermark can be aggregated) + * - other possible optimizations (especially the implementation is ported from Scala example) + */ +public final class JavaStructuredComplexSessionization { + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaStructuredComplexSessionization "); + System.exit(1); + } + + String host = args[0]; + int port = Integer.parseInt(args[1]); + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredComplexSessionization") + .getOrCreate(); + + // Create DataFrame representing the stream of input lines from connection to host:port + Dataset lines = spark + .readStream() + .format("socket") + .option("host", host) + .option("port", port) + .option("includeTimestamp", true) + .load(); + + StructType jsonSchema = new StructType() + .add("user_id", StringType) + .add("event_type", StringType) + .add("timestamp", TimestampType); + + long gapDuration = 5 * 1000; // 5 seconds + + // Parse the line into event, as described in classdoc. + Dataset events = lines + .select(from_json(col("value"), jsonSchema).as("event")) + .selectExpr("event.user_id AS user_id", "event.event_type AS event_type", + "event.timestamp AS timestamp") + .withWatermark("timestamp", "10 seconds"); + + // Sessionize the events. Track number of events, start and end timestamps of session, + // and report session when session is closed. + FlatMapGroupsWithStateFunction stateUpdateFunc = + new FlatMapGroupsWithStateFunction() { + private Iterator handleEvict(String userId, GroupState state) { + Sessions sessions = state.get(); + + List evicted = new ArrayList<>(); + List kept = new ArrayList<>(); + + // we sorted sessions by timestamp + sessions.getSessions().forEach(session -> { + if (session.endTime().getTime() < state.getCurrentWatermarkMs()) { + evicted.add(session); + } else { + kept.add(session); + } + }); + + if (kept.isEmpty()) { + state.remove(); + } else { + state.update(Sessions.newInstance(kept)); + // trigger timeout at the end time of the first session + state.setTimeoutTimestamp(kept.get(0).endTime().getTime()); + } + + return evicted.stream() + .map(sessionAcc -> Session.newInstance( + userId, + sessionAcc.endTime().getTime() - sessionAcc.startTime().getTime(), + sessionAcc.getEvents().size())) + .iterator(); + } + + private void mergeSessions(List sessionAccs, GroupState state) { + // we sorted sessionAccs by timestamp + + int curIdx = 0; + while (curIdx < sessionAccs.size() - 1) { + SessionAcc curSession = sessionAccs.get(curIdx); + SessionAcc nextSession = sessionAccs.get(curIdx + 1); + + // Current session and next session can be merged + if (curSession.endTime().getTime() > nextSession.startTime().getTime()) { + List accumulatedEvents = new ArrayList<>(curSession.getEvents()); + accumulatedEvents.addAll(nextSession.getEvents()); + accumulatedEvents.sort( + Comparator.comparingLong(e -> e.getStartTimestamp().getTime())); + + List newSessions = new ArrayList<>(); + List eventsForCurSession = new ArrayList<>(); + for (SessionEvent event : accumulatedEvents) { + eventsForCurSession.add(event); + if (event.eventType == EventTypes.CLOSE_SESSION) { + SessionAcc newSessionAcc = SessionAcc.newInstance(eventsForCurSession); + newSessions.add(newSessionAcc); + eventsForCurSession = new ArrayList<>(); + } + } + if (!eventsForCurSession.isEmpty()) { + SessionAcc newSessionAcc = SessionAcc.newInstance(eventsForCurSession); + newSessions.add(newSessionAcc); + } + + // replace current session and next session with new session(s) + sessionAccs.remove(curIdx + 1); + sessionAccs.set(curIdx, newSessions.get(0)); + if (newSessions.size() > 1) { + sessionAccs.addAll(curIdx + 1, + newSessions.stream().skip(1).collect(Collectors.toList())); + } + + // move the cursor to the last new session(s) + curIdx += newSessions.size() - 1; + } else { + // move to the next session + curIdx++; + } + } + + // update state + state.update(Sessions.newInstance(sessionAccs)); + } + + @Override + public Iterator call( + String userId, Iterator events, GroupState state) { + + if (state.hasTimedOut() && state.exists()) { + return handleEvict(userId, state); + } + + // convert each event as individual session + Stream stream = StreamSupport.stream( + Spliterators.spliteratorUnknownSize( + events, + Spliterator.ORDERED), false); + List sessionsFromEvents = stream.map(r -> { + SessionEvent event = SessionEvent.newInstance(userId, r.getString(1), + r.getTimestamp(2), gapDuration); + return SessionAcc.newInstance(event); + }).collect(Collectors.toList()); + + if (sessionsFromEvents.isEmpty()) { + return Collections.emptyIterator(); + } + + // sort sessions via start timestamp + List allSessions = new ArrayList<>(sessionsFromEvents); + if (state.exists()) { + allSessions.addAll(state.get().getSessions()); + } + allSessions.sort(Comparator.comparingLong(s -> s.startTime().getTime())); + + // merge sessions + mergeSessions(allSessions, state); + + // we still need to handle eviction here + return handleEvict(userId, state); + } + }; + + Dataset sessionUpdates = events + .groupByKey((MapFunction) event -> event.getString(0), Encoders.STRING()) + .flatMapGroupsWithState( + stateUpdateFunc, + OutputMode.Append(), + Encoders.bean(Sessions.class), + Encoders.bean(Session.class), + GroupStateTimeout.EventTimeTimeout()); + + // Start running the query that prints the session updates to the console + StreamingQuery query = sessionUpdates + .writeStream() + .outputMode("append") + .format("console") + .start(); + + query.awaitTermination(); + } + + public static class Sessions { + private List sessions; + + public List getSessions() { + return sessions; + } + + public void setSessions(List sessions) { + // `sessions` should not be empty, and be sorted by start time + if (sessions.isEmpty()) { + throw new IllegalArgumentException("events should not be empty!"); + } + + List sorted = new ArrayList<>(sessions); + sorted.sort(Comparator.comparingLong(session -> session.startTime().getTime())); + + this.sessions = sorted; + } + + public static Sessions newInstance(List sessions) { + Sessions instance = new Sessions(); + instance.setSessions(sessions); + return instance; + } + } + + public enum EventTypes { + NEW_EVENT, CLOSE_SESSION; + } + + public static class SessionEvent implements Serializable { + private String userId; + private EventTypes eventType; + private Timestamp startTimestamp; + private Timestamp endTimestamp; + + public String getUserId() { + return userId; + } + + public void setUserId(String userId) { + this.userId = userId; + } + + public EventTypes getEventType() { + return eventType; + } + + public void setEventType(EventTypes eventType) { + this.eventType = eventType; + } + + public Timestamp getStartTimestamp() { + return startTimestamp; + } + + public void setStartTimestamp(Timestamp startTimestamp) { + this.startTimestamp = startTimestamp; + } + + public Timestamp getEndTimestamp() { + return endTimestamp; + } + + public void setEndTimestamp(Timestamp endTimestamp) { + this.endTimestamp = endTimestamp; + } + + public static SessionEvent newInstance(String userId, String eventTypeStr, + Timestamp startTimestamp, long gapDuration) { + SessionEvent instance = new SessionEvent(); + instance.setUserId(userId); + instance.setEventType(EventTypes.valueOf(eventTypeStr)); + instance.setStartTimestamp(startTimestamp); + + if (instance.getEventType() == EventTypes.CLOSE_SESSION) { + instance.setEndTimestamp(instance.getStartTimestamp()); + } else { + instance.setEndTimestamp( + new Timestamp(instance.getStartTimestamp().getTime() + gapDuration)); + } + + return instance; + } + } + + public static class SessionAcc implements Serializable { + private List events; + + public Timestamp startTime() { + return events.get(0).startTimestamp; + } + + public Timestamp endTime() { + return events.get(events.size() - 1).getEndTimestamp(); + } + + public List getEvents() { + return events; + } + + public void setEvents(List events) { + // `events` should not be empty, and be sorted by start time + if (events.isEmpty()) { + throw new IllegalArgumentException("events should not be empty!"); + } + + List sorted = new ArrayList<>(events); + sorted.sort(Comparator.comparingLong(event -> event.startTimestamp.getTime())); + + boolean eventCloseSessionExistBeforeLastEvent = sorted + .stream() + .limit(sorted.size() - 1) + .anyMatch(e -> e.eventType == EventTypes.CLOSE_SESSION); + + if (eventCloseSessionExistBeforeLastEvent) { + throw new IllegalStateException("CLOSE_SESSION event cannot be placed except " + + "the last event!"); + } + + this.events = sorted; + } + + public static SessionAcc newInstance(SessionEvent event) { + return newInstance(Collections.singletonList(event)); + } + + public static SessionAcc newInstance(List events) { + SessionAcc instance = new SessionAcc(); + instance.setEvents(events); + return instance; + } + } + + public static class Session implements Serializable { + private String id; + private long duration; + private int numEvents; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public long getDuration() { + return duration; + } + + public void setDuration(long duration) { + this.duration = duration; + } + + public int getNumEvents() { + return numEvents; + } + + public void setNumEvents(int numEvents) { + this.numEvents = numEvents; + } + + public static Session newInstance(String id, long duration, int numEvents) { + Session instance = new Session(); + instance.setId(id); + instance.setDuration(duration); + instance.setNumEvents(numEvents); + return instance; + } + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKafkaWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKafkaWordCount.java new file mode 100644 index 0000000..4e02719 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKafkaWordCount.java @@ -0,0 +1,92 @@ +/* + * 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. + */ + +package org.apache.spark.examples.sql.streaming; + +import org.apache.spark.api.java.function.FlatMapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; + +import java.util.Arrays; + +/** + * Consumes messages from one or more topics in Kafka and does wordcount. + * Usage: JavaStructuredKafkaWordCount + * The Kafka "bootstrap.servers" configuration. A + * comma-separated list of host:port. + * There are three kinds of type, i.e. 'assign', 'subscribe', + * 'subscribePattern'. + * |- Specific TopicPartitions to consume. Json string + * | {"topicA":[0,1],"topicB":[2,4]}. + * |- The topic list to subscribe. A comma-separated list of + * | topics. + * |- The pattern used to subscribe to topic(s). + * | Java regex string. + * |- Only one of "assign, "subscribe" or "subscribePattern" options can be + * | specified for Kafka source. + * Different value format depends on the value of 'subscribe-type'. + * + * Example: + * `$ bin/run-example \ + * sql.streaming.JavaStructuredKafkaWordCount host1:port1,host2:port2 \ + * subscribe topic1,topic2` + */ +public final class JavaStructuredKafkaWordCount { + + public static void main(String[] args) throws Exception { + if (args.length < 3) { + System.err.println("Usage: JavaStructuredKafkaWordCount " + + " "); + System.exit(1); + } + + String bootstrapServers = args[0]; + String subscribeType = args[1]; + String topics = args[2]; + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredKafkaWordCount") + .getOrCreate(); + + // Create DataSet representing the stream of input lines from kafka + Dataset lines = spark + .readStream() + .format("kafka") + .option("kafka.bootstrap.servers", bootstrapServers) + .option(subscribeType, topics) + .load() + .selectExpr("CAST(value AS STRING)") + .as(Encoders.STRING()); + + // Generate running word count + Dataset wordCounts = lines.flatMap( + (FlatMapFunction) x -> Arrays.asList(x.split(" ")).iterator(), + Encoders.STRING()).groupBy("value").count(); + + // Start running the query that prints the running counts to the console + StreamingQuery query = wordCounts.writeStream() + .outputMode("complete") + .format("console") + .start(); + + query.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKerberizedKafkaWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKerberizedKafkaWordCount.java new file mode 100644 index 0000000..ba0c373 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredKerberizedKafkaWordCount.java @@ -0,0 +1,129 @@ +/* + * 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. + */ + +package org.apache.spark.examples.sql.streaming; + +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.spark.api.java.function.FlatMapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; + +import java.util.Arrays; + +/** + * Consumes messages from one or more topics in Kafka and does wordcount. + * Usage: JavaStructuredKerberizedKafkaWordCount + * The Kafka "bootstrap.servers" configuration. A + * comma-separated list of host:port. + * There are three kinds of type, i.e. 'assign', 'subscribe', + * 'subscribePattern'. + * |- Specific TopicPartitions to consume. Json string + * | {"topicA":[0,1],"topicB":[2,4]}. + * |- The topic list to subscribe. A comma-separated list of + * | topics. + * |- The pattern used to subscribe to topic(s). + * | Java regex string. + * |- Only one of "assign, "subscribe" or "subscribePattern" options can be + * | specified for Kafka source. + * Different value format depends on the value of 'subscribe-type'. + * + * Example: + * Yarn client: + * $ bin/run-example --files ${jaas_path}/kafka_jaas.conf,${keytab_path}/kafka.service.keytab \ + * --driver-java-options "-Djava.security.auth.login.config=${path}/kafka_driver_jaas.conf" \ + * --conf \ + * "spark.executor.extraJavaOptions=-Djava.security.auth.login.config=./kafka_jaas.conf" \ + * --master yarn + * sql.streaming.JavaStructuredKerberizedKafkaWordCount broker1-host:port,broker2-host:port \ + * subscribe topic1,topic2 + * Yarn cluster: + * $ bin/run-example --files \ + * ${jaas_path}/kafka_jaas.conf,${keytab_path}/kafka.service.keytab,${krb5_path}/krb5.conf \ + * --driver-java-options \ + * "-Djava.security.auth.login.config=./kafka_jaas.conf \ + * -Djava.security.krb5.conf=./krb5.conf" \ + * --conf \ + * "spark.executor.extraJavaOptions=-Djava.security.auth.login.config=./kafka_jaas.conf" \ + * --master yarn --deploy-mode cluster \ + * sql.streaming.JavaStructuredKerberizedKafkaWordCount broker1-host:port,broker2-host:port \ + * subscribe topic1,topic2 + * + * kafka_jaas.conf can manually create, template as: + * KafkaClient { + * com.sun.security.auth.module.Krb5LoginModule required + * keyTab="./kafka.service.keytab" + * useKeyTab=true + * storeKey=true + * useTicketCache=false + * serviceName="kafka" + * principal="kafka/host@EXAMPLE.COM"; + * }; + * kafka_driver_jaas.conf (used by yarn client) and kafka_jaas.conf are basically the same + * except for some differences at 'keyTab'. In kafka_driver_jaas.conf, 'keyTab' should be + * "${keytab_path}/kafka.service.keytab". + * In addition, for IBM JVMs, please use 'com.ibm.security.auth.module.Krb5LoginModule' + * instead of 'com.sun.security.auth.module.Krb5LoginModule'. + * + * Note that this example uses SASL_PLAINTEXT for simplicity; however, + * SASL_PLAINTEXT has no SSL encryption and likely be less secure. Please consider + * using SASL_SSL in production. + */ +public final class JavaStructuredKerberizedKafkaWordCount { + public static void main(String[] args) throws Exception { + if (args.length < 3) { + System.err.println("Usage: JavaStructuredKerberizedKafkaWordCount " + + " "); + System.exit(1); + } + + String bootstrapServers = args[0]; + String subscribeType = args[1]; + String topics = args[2]; + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredKerberizedKafkaWordCount") + .getOrCreate(); + + // Create DataSet representing the stream of input lines from kafka + Dataset lines = spark + .readStream() + .format("kafka") + .option("kafka.bootstrap.servers", bootstrapServers) + .option(subscribeType, topics) + .option("kafka.security.protocol", SecurityProtocol.SASL_PLAINTEXT.name) + .load() + .selectExpr("CAST(value AS STRING)") + .as(Encoders.STRING()); + + // Generate running word count + Dataset wordCounts = lines.flatMap( + (FlatMapFunction) x -> Arrays.asList(x.split(" ")).iterator(), + Encoders.STRING()).groupBy("value").count(); + + // Start running the query that prints the running counts to the console + StreamingQuery query = wordCounts.writeStream() + .outputMode("complete") + .format("console") + .start(); + + query.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCount.java new file mode 100644 index 0000000..3af7869 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCount.java @@ -0,0 +1,78 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql.streaming; + +import org.apache.spark.api.java.function.FlatMapFunction; +import org.apache.spark.sql.*; +import org.apache.spark.sql.streaming.StreamingQuery; + +import java.util.Arrays; + +/** + * Counts words in UTF8 encoded, '\n' delimited text received from the network. + * + * Usage: JavaStructuredNetworkWordCount + * and describe the TCP server that Structured Streaming + * would connect to receive data. + * + * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example sql.streaming.JavaStructuredNetworkWordCount + * localhost 9999` + */ +public final class JavaStructuredNetworkWordCount { + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaStructuredNetworkWordCount "); + System.exit(1); + } + + String host = args[0]; + int port = Integer.parseInt(args[1]); + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredNetworkWordCount") + .getOrCreate(); + + // Create DataFrame representing the stream of input lines from connection to host:port + Dataset lines = spark + .readStream() + .format("socket") + .option("host", host) + .option("port", port) + .load(); + + // Split the lines into words + Dataset words = lines.as(Encoders.STRING()).flatMap( + (FlatMapFunction) x -> Arrays.asList(x.split(" ")).iterator(), + Encoders.STRING()); + + // Generate running word count + Dataset wordCounts = words.groupBy("value").count(); + + // Start running the query that prints the running counts to the console + StreamingQuery query = wordCounts.writeStream() + .outputMode("complete") + .format("console") + .start(); + + query.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCountWindowed.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCountWindowed.java new file mode 100644 index 0000000..93ec5e2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredNetworkWordCountWindowed.java @@ -0,0 +1,112 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql.streaming; + +import org.apache.spark.api.java.function.FlatMapFunction; +import org.apache.spark.sql.*; +import org.apache.spark.sql.streaming.StreamingQuery; +import scala.Tuple2; + +import java.sql.Timestamp; +import java.util.ArrayList; +import java.util.List; + +/** + * Counts words in UTF8 encoded, '\n' delimited text received from the network over a + * sliding window of configurable duration. Each line from the network is tagged + * with a timestamp that is used to determine the windows into which it falls. + * + * Usage: JavaStructuredNetworkWordCountWindowed + * [] + * and describe the TCP server that Structured Streaming + * would connect to receive data. + * gives the size of window, specified as integer number of seconds + * gives the amount of time successive windows are offset from one another, + * given in the same units as above. should be less than or equal to + * . If the two are equal, successive windows have no overlap. If + * is not provided, it defaults to . + * + * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example sql.streaming.JavaStructuredNetworkWordCountWindowed + * localhost 9999 []` + * + * One recommended , pair is 10, 5 + */ +public final class JavaStructuredNetworkWordCountWindowed { + + public static void main(String[] args) throws Exception { + if (args.length < 3) { + System.err.println("Usage: JavaStructuredNetworkWordCountWindowed " + + " []"); + System.exit(1); + } + + String host = args[0]; + int port = Integer.parseInt(args[1]); + int windowSize = Integer.parseInt(args[2]); + int slideSize = (args.length == 3) ? windowSize : Integer.parseInt(args[3]); + if (slideSize > windowSize) { + System.err.println(" must be less than or equal to "); + } + String windowDuration = windowSize + " seconds"; + String slideDuration = slideSize + " seconds"; + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredNetworkWordCountWindowed") + .getOrCreate(); + + // Create DataFrame representing the stream of input lines from connection to host:port + Dataset lines = spark + .readStream() + .format("socket") + .option("host", host) + .option("port", port) + .option("includeTimestamp", true) + .load(); + + // Split the lines into words, retaining timestamps + Dataset words = lines + .as(Encoders.tuple(Encoders.STRING(), Encoders.TIMESTAMP())) + .flatMap((FlatMapFunction, Tuple2>) t -> { + List> result = new ArrayList<>(); + for (String word : t._1.split(" ")) { + result.add(new Tuple2<>(word, t._2)); + } + return result.iterator(); + }, + Encoders.tuple(Encoders.STRING(), Encoders.TIMESTAMP()) + ).toDF("word", "timestamp"); + + // Group the data by window and word and compute the count of each group + Dataset windowedCounts = words.groupBy( + functions.window(words.col("timestamp"), windowDuration, slideDuration), + words.col("word") + ).count().orderBy("window"); + + // Start running the query that prints the windowed word counts to the console + StreamingQuery query = windowedCounts.writeStream() + .outputMode("complete") + .format("console") + .option("truncate", "false") + .start(); + + query.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredSessionization.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredSessionization.java new file mode 100644 index 0000000..eb7ce11 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/sql/streaming/JavaStructuredSessionization.java @@ -0,0 +1,86 @@ +/* + * 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. + */ +package org.apache.spark.examples.sql.streaming; + +import org.apache.spark.sql.*; +import org.apache.spark.sql.streaming.StreamingQuery; + +import static org.apache.spark.sql.functions.*; + +/** + * Counts words in UTF8 encoded, '\n' delimited text received from the network. + *

+ * Usage: JavaStructuredSessionization + * and describe the TCP server that Structured Streaming + * would connect to receive data. + *

+ * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example sql.streaming.JavaStructuredSessionization + * localhost 9999` + */ +public final class JavaStructuredSessionization { + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaStructuredSessionization "); + System.exit(1); + } + + String host = args[0]; + int port = Integer.parseInt(args[1]); + + SparkSession spark = SparkSession + .builder() + .appName("JavaStructuredSessionization") + .getOrCreate(); + + // Create DataFrame representing the stream of input lines from connection to host:port + Dataset lines = spark + .readStream() + .format("socket") + .option("host", host) + .option("port", port) + .option("includeTimestamp", true) + .load(); + + // Split the lines into words, retaining timestamps + // split() splits each line into an array, and explode() turns the array into multiple rows + // treat words as sessionId of events + Dataset events = lines + .selectExpr("explode(split(value, ' ')) AS sessionId", "timestamp AS eventTime"); + + // Sessionize the events. Track number of events, start and end timestamps of session, + // and report session updates. + Dataset sessionUpdates = events + .groupBy(session_window(col("eventTime"), "10 seconds").as("session"), col("sessionId")) + .agg(count("*").as("numEvents")) + .selectExpr("sessionId", "CAST(session.start AS LONG)", "CAST(session.end AS LONG)", + "CAST(session.end AS LONG) - CAST(session.start AS LONG) AS durationMs", + "numEvents"); + + // Start running the query that prints the session updates to the console + StreamingQuery query = sessionUpdates + .writeStream() + .outputMode("update") + .format("console") + .start(); + + query.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java new file mode 100644 index 0000000..f84a197 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java @@ -0,0 +1,134 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import com.google.common.io.Closeables; + +import org.apache.spark.SparkConf; +import org.apache.spark.storage.StorageLevel; +import org.apache.spark.streaming.Duration; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaPairDStream; +import org.apache.spark.streaming.api.java.JavaReceiverInputDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; +import org.apache.spark.streaming.receiver.Receiver; +import scala.Tuple2; + +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.net.ConnectException; +import java.net.Socket; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.regex.Pattern; + +/** + * Custom Receiver that receives data over a socket. Received bytes is interpreted as + * text and \n delimited lines are considered as records. They are then counted and printed. + * + * Usage: JavaCustomReceiver + * is the Spark master URL. In local mode, should be 'local[n]' with n > 1. + * and of the TCP server that Spark Streaming would connect to receive data. + * + * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example org.apache.spark.examples.streaming.JavaCustomReceiver localhost 9999` + */ + +public class JavaCustomReceiver extends Receiver { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaCustomReceiver "); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + // Create the context with a 1 second batch size + SparkConf sparkConf = new SparkConf().setAppName("JavaCustomReceiver"); + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, new Duration(1000)); + + // Create an input stream with the custom receiver on target ip:port and count the + // words in input stream of \n delimited text (e.g. generated by 'nc') + JavaReceiverInputDStream lines = ssc.receiverStream( + new JavaCustomReceiver(args[0], Integer.parseInt(args[1]))); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + JavaPairDStream wordCounts = words.mapToPair(s -> new Tuple2<>(s, 1)) + .reduceByKey((i1, i2) -> i1 + i2); + + wordCounts.print(); + ssc.start(); + ssc.awaitTermination(); + } + + // ============= Receiver code that receives data over a socket ============== + + String host = null; + int port = -1; + + public JavaCustomReceiver(String host_ , int port_) { + super(StorageLevel.MEMORY_AND_DISK_2()); + host = host_; + port = port_; + } + + @Override + public void onStart() { + // Start the thread that receives data over a connection + new Thread(this::receive).start(); + } + + @Override + public void onStop() { + // There is nothing much to do as the thread calling receive() + // is designed to stop by itself isStopped() returns false + } + + /** Create a socket connection and receive data until receiver is stopped */ + private void receive() { + try { + Socket socket = null; + BufferedReader reader = null; + try { + // connect to the server + socket = new Socket(host, port); + reader = new BufferedReader( + new InputStreamReader(socket.getInputStream(), StandardCharsets.UTF_8)); + // Until stopped or connection broken continue reading + String userInput; + while (!isStopped() && (userInput = reader.readLine()) != null) { + System.out.println("Received data '" + userInput + "'"); + store(userInput); + } + } finally { + Closeables.close(reader, /* swallowIOException = */ true); + Closeables.close(socket, /* swallowIOException = */ true); + } + // Restart in an attempt to connect again when server is active again + restart("Trying to connect again"); + } catch(ConnectException ce) { + // restart if could not connect to server + restart("Could not connect", ce); + } catch(Throwable t) { + restart("Error receiving data", t); + } + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java new file mode 100644 index 0000000..748bf58 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java @@ -0,0 +1,98 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Arrays; +import java.util.Map; +import java.util.Set; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.common.serialization.StringDeserializer; + +import org.apache.spark.SparkConf; +import org.apache.spark.streaming.api.java.*; +import org.apache.spark.streaming.kafka010.ConsumerStrategies; +import org.apache.spark.streaming.kafka010.KafkaUtils; +import org.apache.spark.streaming.kafka010.LocationStrategies; +import org.apache.spark.streaming.Durations; + +/** + * Consumes messages from one or more topics in Kafka and does wordcount. + * Usage: JavaDirectKafkaWordCount + * is a list of one or more Kafka brokers + * is a consumer group name to consume from topics + * is a list of one or more kafka topics to consume from + * + * Example: + * $ bin/run-example streaming.JavaDirectKafkaWordCount broker1-host:port,broker2-host:port \ + * consumer-group topic1,topic2 + */ + +public final class JavaDirectKafkaWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 3) { + System.err.println("Usage: JavaDirectKafkaWordCount \n" + + " is a list of one or more Kafka brokers\n" + + " is a consumer group name to consume from topics\n" + + " is a list of one or more kafka topics to consume from\n\n"); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + String brokers = args[0]; + String groupId = args[1]; + String topics = args[2]; + + // Create context with a 2 seconds batch interval + SparkConf sparkConf = new SparkConf().setAppName("JavaDirectKafkaWordCount"); + JavaStreamingContext jssc = new JavaStreamingContext(sparkConf, Durations.seconds(2)); + + Set topicsSet = new HashSet<>(Arrays.asList(topics.split(","))); + Map kafkaParams = new HashMap<>(); + kafkaParams.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokers); + kafkaParams.put(ConsumerConfig.GROUP_ID_CONFIG, groupId); + kafkaParams.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + kafkaParams.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + + // Create direct kafka stream with brokers and topics + JavaInputDStream> messages = KafkaUtils.createDirectStream( + jssc, + LocationStrategies.PreferConsistent(), + ConsumerStrategies.Subscribe(topicsSet, kafkaParams)); + + // Get the lines, split them into words, count the words and print + JavaDStream lines = messages.map(ConsumerRecord::value); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + JavaPairDStream wordCounts = words.mapToPair(s -> new Tuple2<>(s, 1)) + .reduceByKey((i1, i2) -> i1 + i2); + wordCounts.print(); + + // Start the computation + jssc.start(); + jssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKerberizedKafkaWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKerberizedKafkaWordCount.java new file mode 100644 index 0000000..9c327c1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKerberizedKafkaWordCount.java @@ -0,0 +1,140 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Arrays; +import java.util.Map; +import java.util.Set; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import org.apache.kafka.clients.CommonClientConfigs; +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.common.serialization.StringDeserializer; + +import org.apache.spark.SparkConf; +import org.apache.spark.streaming.api.java.*; +import org.apache.spark.streaming.kafka010.ConsumerStrategies; +import org.apache.spark.streaming.kafka010.KafkaUtils; +import org.apache.spark.streaming.kafka010.LocationStrategies; +import org.apache.spark.streaming.Durations; + +/** + * Consumes messages from one or more topics in Kafka and does wordcount. + * Usage: JavaDirectKerberizedKafkaWordCount + * is a list of one or more Kafka brokers + * is a consumer group name to consume from topics + * is a list of one or more kafka topics to consume from + * + * Example: + * Yarn client: + * $ bin/run-example --files ${jaas_path}/kafka_jaas.conf,${keytab_path}/kafka.service.keytab \ + * --driver-java-options "-Djava.security.auth.login.config=${path}/kafka_driver_jaas.conf" \ + * --conf \ + * "spark.executor.extraJavaOptions=-Djava.security.auth.login.config=./kafka_jaas.conf" \ + * --master yarn + * streaming.JavaDirectKerberizedKafkaWordCount broker1-host:port,broker2-host:port \ + * consumer-group topic1,topic2 + * Yarn cluster: + * $ bin/run-example --files \ + * ${jaas_path}/kafka_jaas.conf,${keytab_path}/kafka.service.keytab,${krb5_path}/krb5.conf \ + * --driver-java-options \ + * "-Djava.security.auth.login.config=./kafka_jaas.conf \ + * -Djava.security.krb5.conf=./krb5.conf" \ + * --conf \ + * "spark.executor.extraJavaOptions=-Djava.security.auth.login.config=./kafka_jaas.conf" \ + * --master yarn --deploy-mode cluster \ + * streaming.JavaDirectKerberizedKafkaWordCount broker1-host:port,broker2-host:port \ + * consumer-group topic1,topic2 + * + * kafka_jaas.conf can manually create, template as: + * KafkaClient { + * com.sun.security.auth.module.Krb5LoginModule required + * keyTab="./kafka.service.keytab" + * useKeyTab=true + * storeKey=true + * useTicketCache=false + * serviceName="kafka" + * principal="kafka/host@EXAMPLE.COM"; + * }; + * kafka_driver_jaas.conf (used by yarn client) and kafka_jaas.conf are basically the same + * except for some differences at 'keyTab'. In kafka_driver_jaas.conf, 'keyTab' should be + * "${keytab_path}/kafka.service.keytab". + * In addition, for IBM JVMs, please use 'com.ibm.security.auth.module.Krb5LoginModule' + * instead of 'com.sun.security.auth.module.Krb5LoginModule'. + * + * Note that this example uses SASL_PLAINTEXT for simplicity; however, + * SASL_PLAINTEXT has no SSL encryption and likely be less secure. Please consider + * using SASL_SSL in production. + */ + +public final class JavaDirectKerberizedKafkaWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 3) { + System.err.println( + "Usage: JavaDirectKerberizedKafkaWordCount \n" + + " is a list of one or more Kafka brokers\n" + + " is a consumer group name to consume from topics\n" + + " is a list of one or more kafka topics to consume from\n\n"); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + String brokers = args[0]; + String groupId = args[1]; + String topics = args[2]; + + // Create context with a 2 seconds batch interval + SparkConf sparkConf = new SparkConf().setAppName("JavaDirectKerberizedKafkaWordCount"); + JavaStreamingContext jssc = new JavaStreamingContext(sparkConf, Durations.seconds(2)); + + Set topicsSet = new HashSet<>(Arrays.asList(topics.split(","))); + Map kafkaParams = new HashMap<>(); + kafkaParams.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokers); + kafkaParams.put(ConsumerConfig.GROUP_ID_CONFIG, groupId); + kafkaParams.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + kafkaParams.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + kafkaParams.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, + SecurityProtocol.SASL_PLAINTEXT.name); + + // Create direct kafka stream with brokers and topics + JavaInputDStream> messages = KafkaUtils.createDirectStream( + jssc, + LocationStrategies.PreferConsistent(), + ConsumerStrategies.Subscribe(topicsSet, kafkaParams)); + + // Get the lines, split them into words, count the words and print + JavaDStream lines = messages.map(ConsumerRecord::value); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + JavaPairDStream wordCounts = words.mapToPair(s -> new Tuple2<>(s, 1)) + .reduceByKey((i1, i2) -> i1 + i2); + wordCounts.print(); + + // Start the computation + jssc.start(); + jssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaNetworkWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaNetworkWordCount.java new file mode 100644 index 0000000..d56134b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaNetworkWordCount.java @@ -0,0 +1,73 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.Arrays; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.StorageLevels; +import org.apache.spark.streaming.Durations; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaPairDStream; +import org.apache.spark.streaming.api.java.JavaReceiverInputDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; + +/** + * Counts words in UTF8 encoded, '\n' delimited text received from the network every second. + * + * Usage: JavaNetworkWordCount + * and describe the TCP server that Spark Streaming would connect to receive data. + * + * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example org.apache.spark.examples.streaming.JavaNetworkWordCount localhost 9999` + */ +public final class JavaNetworkWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaNetworkWordCount "); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + // Create the context with a 1 second batch size + SparkConf sparkConf = new SparkConf().setAppName("JavaNetworkWordCount"); + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, Durations.seconds(1)); + + // Create a JavaReceiverInputDStream on target ip:port and count the + // words in input stream of \n delimited text (e.g. generated by 'nc') + // Note that no duplication in storage level only for running locally. + // Replication necessary in distributed scenario for fault tolerance. + JavaReceiverInputDStream lines = ssc.socketTextStream( + args[0], Integer.parseInt(args[1]), StorageLevels.MEMORY_AND_DISK_SER); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + JavaPairDStream wordCounts = words.mapToPair(s -> new Tuple2<>(s, 1)) + .reduceByKey((i1, i2) -> i1 + i2); + + wordCounts.print(); + ssc.start(); + ssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaQueueStream.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaQueueStream.java new file mode 100644 index 0000000..e86f8ab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaQueueStream.java @@ -0,0 +1,71 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.ArrayList; +import java.util.LinkedList; +import java.util.List; +import java.util.Queue; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.streaming.Duration; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaPairDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; + +public final class JavaQueueStream { + private JavaQueueStream() { + } + + public static void main(String[] args) throws Exception { + + StreamingExamples.setStreamingLogLevels(); + SparkConf sparkConf = new SparkConf().setAppName("JavaQueueStream"); + + // Create the context + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, new Duration(1000)); + + // Create the queue through which RDDs can be pushed to + // a QueueInputDStream + + // Create and push some RDDs into the queue + List list = new ArrayList<>(); + for (int i = 0; i < 1000; i++) { + list.add(i); + } + + Queue> rddQueue = new LinkedList<>(); + for (int i = 0; i < 30; i++) { + rddQueue.add(ssc.sparkContext().parallelize(list)); + } + + // Create the QueueInputDStream and use it do some processing + JavaDStream inputStream = ssc.queueStream(rddQueue); + JavaPairDStream mappedStream = inputStream.mapToPair( + i -> new Tuple2<>(i % 10, 1)); + JavaPairDStream reducedStream = mappedStream.reduceByKey( + (i1, i2) -> i1 + i2); + + reducedStream.print(); + ssc.start(); + ssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecord.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecord.java new file mode 100644 index 0000000..e63697a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecord.java @@ -0,0 +1,31 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +/** Java Bean class to be used with the example JavaSqlNetworkWordCount. */ +public class JavaRecord implements java.io.Serializable { + private String word; + + public String getWord() { + return word; + } + + public void setWord(String word) { + this.word = word; + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java new file mode 100644 index 0000000..0c11c40 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java @@ -0,0 +1,191 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.io.File; +import java.nio.charset.Charset; +import java.util.Arrays; +import java.util.List; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import com.google.common.io.Files; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.function.*; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.streaming.Durations; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaPairDStream; +import org.apache.spark.streaming.api.java.JavaReceiverInputDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; +import org.apache.spark.util.LongAccumulator; + +/** + * Use this singleton to get or register a Broadcast variable. + */ +class JavaWordExcludeList { + + private static volatile Broadcast> instance = null; + + public static Broadcast> getInstance(JavaSparkContext jsc) { + if (instance == null) { + synchronized (JavaWordExcludeList.class) { + if (instance == null) { + List wordExcludeList = Arrays.asList("a", "b", "c"); + instance = jsc.broadcast(wordExcludeList); + } + } + } + return instance; + } +} + +/** + * Use this singleton to get or register an Accumulator. + */ +class JavaDroppedWordsCounter { + + private static volatile LongAccumulator instance = null; + + public static LongAccumulator getInstance(JavaSparkContext jsc) { + if (instance == null) { + synchronized (JavaDroppedWordsCounter.class) { + if (instance == null) { + instance = jsc.sc().longAccumulator("DroppedWordsCounter"); + } + } + } + return instance; + } +} + +/** + * Counts words in text encoded with UTF8 received from the network every second. This example also + * shows how to use lazily instantiated singleton instances for Accumulator and Broadcast so that + * they can be registered on driver failures. + * + * Usage: JavaRecoverableNetworkWordCount + * and describe the TCP server that Spark Streaming would connect to receive + * data. directory to HDFS-compatible file system which checkpoint data + * file to which the word counts will be appended + * + * and must be absolute paths + * + * To run this on your local machine, you need to first run a Netcat server + * + * `$ nc -lk 9999` + * + * and run the example as + * + * `$ ./bin/run-example org.apache.spark.examples.streaming.JavaRecoverableNetworkWordCount \ + * localhost 9999 ~/checkpoint/ ~/out` + * + * If the directory ~/checkpoint/ does not exist (e.g. running for the first time), it will create + * a new StreamingContext (will print "Creating new context" to the console). Otherwise, if + * checkpoint data exists in ~/checkpoint/, then it will create StreamingContext from + * the checkpoint data. + * + * Refer to the online documentation for more details. + */ +public final class JavaRecoverableNetworkWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + private static JavaStreamingContext createContext(String ip, + int port, + String checkpointDirectory, + String outputPath) { + + // If you do not see this printed, that means the StreamingContext has been loaded + // from the new checkpoint + System.out.println("Creating new context"); + File outputFile = new File(outputPath); + if (outputFile.exists()) { + outputFile.delete(); + } + SparkConf sparkConf = new SparkConf().setAppName("JavaRecoverableNetworkWordCount"); + // Create the context with a 1 second batch size + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, Durations.seconds(1)); + ssc.checkpoint(checkpointDirectory); + + // Create a socket stream on target ip:port and count the + // words in input stream of \n delimited text (e.g. generated by 'nc') + JavaReceiverInputDStream lines = ssc.socketTextStream(ip, port); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + JavaPairDStream wordCounts = words.mapToPair(s -> new Tuple2<>(s, 1)) + .reduceByKey((i1, i2) -> i1 + i2); + + wordCounts.foreachRDD((rdd, time) -> { + // Get or register the excludeList Broadcast + Broadcast> excludeList = + JavaWordExcludeList.getInstance(new JavaSparkContext(rdd.context())); + // Get or register the droppedWordsCounter Accumulator + LongAccumulator droppedWordsCounter = + JavaDroppedWordsCounter.getInstance(new JavaSparkContext(rdd.context())); + // Use excludeList to drop words and use droppedWordsCounter to count them + String counts = rdd.filter(wordCount -> { + if (excludeList.value().contains(wordCount._1())) { + droppedWordsCounter.add(wordCount._2()); + return false; + } else { + return true; + } + }).collect().toString(); + String output = "Counts at time " + time + " " + counts; + System.out.println(output); + System.out.println("Dropped " + droppedWordsCounter.value() + " word(s) totally"); + System.out.println("Appending to " + outputFile.getAbsolutePath()); + Files.append(output + "\n", outputFile, Charset.defaultCharset()); + }); + + return ssc; + } + + public static void main(String[] args) throws Exception { + if (args.length != 4) { + System.err.println("You arguments were " + Arrays.asList(args)); + System.err.println( + "Usage: JavaRecoverableNetworkWordCount \n" + + " . and describe the TCP server that Spark\n" + + " Streaming would connect to receive data. directory to\n" + + " HDFS-compatible file system which checkpoint data file to which\n" + + " the word counts will be appended\n" + + "\n" + + "In local mode, should be 'local[n]' with n > 1\n" + + "Both and must be absolute paths"); + System.exit(1); + } + + String ip = args[0]; + int port = Integer.parseInt(args[1]); + String checkpointDirectory = args[2]; + String outputPath = args[3]; + + // Function to create JavaStreamingContext without any output operations + // (used to detect the new context) + Function0 createContextFunc = + () -> createContext(ip, port, checkpointDirectory, outputPath); + + JavaStreamingContext ssc = + JavaStreamingContext.getOrCreate(checkpointDirectory, createContextFunc); + ssc.start(); + ssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaSqlNetworkWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaSqlNetworkWordCount.java new file mode 100644 index 0000000..5d30698 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaSqlNetworkWordCount.java @@ -0,0 +1,108 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.Arrays; +import java.util.regex.Pattern; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.api.java.StorageLevels; +import org.apache.spark.streaming.Durations; +import org.apache.spark.streaming.api.java.JavaDStream; +import org.apache.spark.streaming.api.java.JavaReceiverInputDStream; +import org.apache.spark.streaming.api.java.JavaStreamingContext; + +/** + * Use DataFrames and SQL to count words in UTF8 encoded, '\n' delimited text received from the + * network every second. + * + * Usage: JavaSqlNetworkWordCount + * and describe the TCP server that Spark Streaming would connect to receive data. + * + * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example org.apache.spark.examples.streaming.JavaSqlNetworkWordCount localhost 9999` + */ +public final class JavaSqlNetworkWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaNetworkWordCount "); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + // Create the context with a 1 second batch size + SparkConf sparkConf = new SparkConf().setAppName("JavaSqlNetworkWordCount"); + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, Durations.seconds(1)); + + // Create a JavaReceiverInputDStream on target ip:port and count the + // words in input stream of \n delimited text (e.g. generated by 'nc') + // Note that no duplication in storage level only for running locally. + // Replication necessary in distributed scenario for fault tolerance. + JavaReceiverInputDStream lines = ssc.socketTextStream( + args[0], Integer.parseInt(args[1]), StorageLevels.MEMORY_AND_DISK_SER); + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + + // Convert RDDs of the words DStream to DataFrame and run SQL query + words.foreachRDD((rdd, time) -> { + SparkSession spark = JavaSparkSessionSingleton.getInstance(rdd.context().getConf()); + + // Convert JavaRDD[String] to JavaRDD[bean class] to DataFrame + JavaRDD rowRDD = rdd.map(word -> { + JavaRecord record = new JavaRecord(); + record.setWord(word); + return record; + }); + Dataset wordsDataFrame = spark.createDataFrame(rowRDD, JavaRecord.class); + + // Creates a temporary view using the DataFrame + wordsDataFrame.createOrReplaceTempView("words"); + + // Do word count on table using SQL and print it + Dataset wordCountsDataFrame = + spark.sql("select word, count(*) as total from words group by word"); + System.out.println("========= " + time + "========="); + wordCountsDataFrame.show(); + }); + + ssc.start(); + ssc.awaitTermination(); + } +} + +/** Lazily instantiated singleton instance of SparkSession */ +class JavaSparkSessionSingleton { + private static transient SparkSession instance = null; + public static SparkSession getInstance(SparkConf sparkConf) { + if (instance == null) { + instance = SparkSession + .builder() + .config(sparkConf) + .getOrCreate(); + } + return instance; + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java new file mode 100644 index 0000000..ec1185a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java @@ -0,0 +1,94 @@ +/* + * 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. + */ + +package org.apache.spark.examples.streaming; + +import java.util.Arrays; +import java.util.List; +import java.util.regex.Pattern; + +import scala.Tuple2; + +import org.apache.spark.SparkConf; +import org.apache.spark.api.java.function.*; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.Optional; +import org.apache.spark.api.java.StorageLevels; +import org.apache.spark.streaming.Durations; +import org.apache.spark.streaming.State; +import org.apache.spark.streaming.StateSpec; +import org.apache.spark.streaming.api.java.*; + +/** + * Counts words cumulatively in UTF8 encoded, '\n' delimited text received from the network every + * second starting with initial value of word count. + * Usage: JavaStatefulNetworkWordCount + * and describe the TCP server that Spark Streaming would connect to receive + * data. + *

+ * To run this on your local machine, you need to first run a Netcat server + * `$ nc -lk 9999` + * and then run the example + * `$ bin/run-example + * org.apache.spark.examples.streaming.JavaStatefulNetworkWordCount localhost 9999` + */ +public class JavaStatefulNetworkWordCount { + private static final Pattern SPACE = Pattern.compile(" "); + + public static void main(String[] args) throws Exception { + if (args.length < 2) { + System.err.println("Usage: JavaStatefulNetworkWordCount "); + System.exit(1); + } + + StreamingExamples.setStreamingLogLevels(); + + // Create the context with a 1 second batch size + SparkConf sparkConf = new SparkConf().setAppName("JavaStatefulNetworkWordCount"); + JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, Durations.seconds(1)); + ssc.checkpoint("."); + + // Initial state RDD input to mapWithState + List> tuples = + Arrays.asList(new Tuple2<>("hello", 1), new Tuple2<>("world", 1)); + JavaPairRDD initialRDD = ssc.sparkContext().parallelizePairs(tuples); + + JavaReceiverInputDStream lines = ssc.socketTextStream( + args[0], Integer.parseInt(args[1]), StorageLevels.MEMORY_AND_DISK_SER_2); + + JavaDStream words = lines.flatMap(x -> Arrays.asList(SPACE.split(x)).iterator()); + + JavaPairDStream wordsDstream = words.mapToPair(s -> new Tuple2<>(s, 1)); + + // Update the cumulative count function + Function3, State, Tuple2> mappingFunc = + (word, one, state) -> { + int sum = one.orElse(0) + (state.exists() ? state.get() : 0); + Tuple2 output = new Tuple2<>(word, sum); + state.update(sum); + return output; + }; + + // DStream made of get cumulative counts that get updated in every batch + JavaMapWithStateDStream> stateDstream = + wordsDstream.mapWithState(StateSpec.function(mappingFunc).initialState(initialRDD)); + + stateDstream.print(); + ssc.start(); + ssc.awaitTermination(); + } +} diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/python/__init__.py b/spark-3.5.0-bin-hadoop3/examples/src/main/python/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/python/__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/spark-3.5.0-bin-hadoop3/examples/src/main/python/als.py b/spark-3.5.0-bin-hadoop3/examples/src/main/python/als.py new file mode 100644 index 0000000..5bd1807 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/python/als.py @@ -0,0 +1,106 @@ +# +# 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. +# + +""" +This is an example implementation of ALS for learning how to use Spark. Please refer to +pyspark.ml.recommendation.ALS for more conventional use. + +This example requires numpy (http://www.numpy.org/) +""" +import sys + +import numpy as np +from numpy.random import rand +from numpy import matrix +from pyspark.sql import SparkSession + +LAMBDA = 0.01 # regularization +np.random.seed(42) + + +def rmse(R: np.ndarray, ms: np.ndarray, us: np.ndarray) -> np.float64: + diff = R - ms * us.T + return np.sqrt(np.sum(np.power(diff, 2)) / (M * U)) + + +def update(i: int, mat: np.ndarray, ratings: np.ndarray) -> np.ndarray: + uu = mat.shape[0] + ff = mat.shape[1] + + XtX = mat.T * mat + Xty = mat.T * ratings[i, :].T + + for j in range(ff): + XtX[j, j] += LAMBDA * uu + + return np.linalg.solve(XtX, Xty) + + +if __name__ == "__main__": + + """ + Usage: als [M] [U] [F] [iterations] [partitions]" + """ + + print("""WARN: This is a naive implementation of ALS and is given as an + example. Please use pyspark.ml.recommendation.ALS for more + conventional use.""", file=sys.stderr) + + spark = SparkSession\ + .builder\ + .appName("PythonALS")\ + .getOrCreate() + + sc = spark.sparkContext + + M = int(sys.argv[1]) if len(sys.argv) > 1 else 100 + U = int(sys.argv[2]) if len(sys.argv) > 2 else 500 + F = int(sys.argv[3]) if len(sys.argv) > 3 else 10 + ITERATIONS = int(sys.argv[4]) if len(sys.argv) > 4 else 5 + partitions = int(sys.argv[5]) if len(sys.argv) > 5 else 2 + + print("Running ALS with M=%d, U=%d, F=%d, iters=%d, partitions=%d\n" % + (M, U, F, ITERATIONS, partitions)) + + R = matrix(rand(M, F)) * matrix(rand(U, F).T) + ms: matrix = matrix(rand(M, F)) + us: matrix = matrix(rand(U, F)) + + Rb = sc.broadcast(R) + msb = sc.broadcast(ms) + usb = sc.broadcast(us) + + for i in range(ITERATIONS): + ms_ = sc.parallelize(range(M), partitions) \ + .map(lambda x: update(x, usb.value, Rb.value)) \ + .collect() + # collect() returns a list, so array ends up being + # a 3-d array, we take the first 2 dims for the matrix + ms = matrix(np.array(ms_)[:, :, 0]) + msb = sc.broadcast(ms) + + us_ = sc.parallelize(range(U), partitions) \ + .map(lambda x: update(x, msb.value, Rb.value.T)) \ + .collect() + us = matrix(np.array(us_)[:, :, 0]) + usb = sc.broadcast(us) + + error = rmse(R, ms, us) + print("Iteration %d:" % i) + print("\nRMSE: %5.4f\n" % error) + + spark.stop() diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/python/avro_inputformat.py b/spark-3.5.0-bin-hadoop3/examples/src/main/python/avro_inputformat.py new file mode 100644 index 0000000..e303860 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/python/avro_inputformat.py @@ -0,0 +1,91 @@ +# +# 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. +# + +""" +Read data file users.avro in local Spark distro: + +$ cd $SPARK_HOME +$ ./bin/spark-submit --driver-class-path /path/to/example/jar \ +> ./examples/src/main/python/avro_inputformat.py \ +> examples/src/main/resources/users.avro +{u'favorite_color': None, u'name': u'Alyssa', u'favorite_numbers': [3, 9, 15, 20]} +{u'favorite_color': u'red', u'name': u'Ben', u'favorite_numbers': []} + +To read name and favorite_color fields only, specify the following reader schema: + +$ cat examples/src/main/resources/user.avsc +{"namespace": "example.avro", + "type": "record", + "name": "User", + "fields": [ + {"name": "name", "type": "string"}, + {"name": "favorite_color", "type": ["string", "null"]} + ] +} + +$ ./bin/spark-submit --driver-class-path /path/to/example/jar \ +> ./examples/src/main/python/avro_inputformat.py \ +> examples/src/main/resources/users.avro examples/src/main/resources/user.avsc +{u'favorite_color': None, u'name': u'Alyssa'} +{u'favorite_color': u'red', u'name': u'Ben'} +""" +import sys +from typing import Any, Tuple + +from functools import reduce +from pyspark.rdd import RDD +from pyspark.sql import SparkSession + +if __name__ == "__main__": + if len(sys.argv) != 2 and len(sys.argv) != 3: + print(""" + Usage: avro_inputformat [reader_schema_file] + + Run with example jar: + ./bin/spark-submit --driver-class-path /path/to/example/jar \ + /path/to/examples/avro_inputformat.py [reader_schema_file] + Assumes you have Avro data stored in . Reader schema can be optionally specified + in [reader_schema_file]. + """, file=sys.stderr) + sys.exit(-1) + + path = sys.argv[1] + + spark = SparkSession\ + .builder\ + .appName("AvroKeyInputFormat")\ + .getOrCreate() + + sc = spark.sparkContext + + conf = None + if len(sys.argv) == 3: + schema_rdd = sc.textFile(sys.argv[2], 1).collect() + conf = {"avro.schema.input.key": reduce(lambda x, y: x + y, schema_rdd)} + + avro_rdd: RDD[Tuple[Any, None]] = sc.newAPIHadoopFile( + path, + "org.apache.avro.mapreduce.AvroKeyInputFormat", + "org.apache.avro.mapred.AvroKey", + "org.apache.hadoop.io.NullWritable", + keyConverter="org.apache.spark.examples.pythonconverters.AvroWrapperToJavaConverter", + conf=conf) + output = avro_rdd.map(lambda x: x[0]).collect() + for k in output: + print(k) + + spark.stop() diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/python/kmeans.py b/spark-3.5.0-bin-hadoop3/examples/src/main/python/kmeans.py new file mode 100644 index 0000000..dc82827 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/python/kmeans.py @@ -0,0 +1,85 @@ +# +# 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. +# + +""" +The K-means algorithm written from scratch against PySpark. In practice, +one may prefer to use the KMeans algorithm in ML, as shown in +examples/src/main/python/ml/kmeans_example.py. + +This example requires NumPy (http://www.numpy.org/). +""" +import sys +from typing import List + +import numpy as np +from pyspark.sql import SparkSession + + +def parseVector(line: str) -> np.ndarray: + return np.array([float(x) for x in line.split(' ')]) + + +def closestPoint(p: np.ndarray, centers: List[np.ndarray]) -> int: + bestIndex = 0 + closest = float("+inf") + for i in range(len(centers)): + tempDist = np.sum((p - centers[i]) ** 2) + if tempDist < closest: + closest = tempDist + bestIndex = i + return bestIndex + + +if __name__ == "__main__": + + if len(sys.argv) != 4: + print("Usage: kmeans ", file=sys.stderr) + sys.exit(-1) + + print("""WARN: This is a naive implementation of KMeans Clustering and is given + as an example! Please refer to examples/src/main/python/ml/kmeans_example.py for an + example on how to use ML's KMeans implementation.""", file=sys.stderr) + + spark = SparkSession\ + .builder\ + .appName("PythonKMeans")\ + .getOrCreate() + + lines = spark.read.text(sys.argv[1]).rdd.map(lambda r: r[0]) + data = lines.map(parseVector).cache() + K = int(sys.argv[2]) + convergeDist = float(sys.argv[3]) + + kPoints = data.takeSample(False, K, 1) + tempDist = 1.0 + + while tempDist > convergeDist: + closest = data.map( + lambda p: (closestPoint(p, kPoints), (p, 1))) + pointStats = closest.reduceByKey( + lambda p1_c1, p2_c2: (p1_c1[0] + p2_c2[0], p1_c1[1] + p2_c2[1])) + newPoints = pointStats.map( + lambda st: (st[0], st[1][0] / st[1][1])).collect() + + tempDist = sum(np.sum((kPoints[iK] - p) ** 2) for (iK, p) in newPoints) + + for (iK, p) in newPoints: + kPoints[iK] = p + + print("Final centers: " + str(kPoints)) + + spark.stop() diff --git a/spark-3.5.0-bin-hadoop3/examples/src/main/python/logistic_regression.py b/spark-3.5.0-bin-hadoop3/examples/src/main/python/logistic_regression.py new file mode 100644 index 0000000..9645af6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/examples/src/main/python/logistic_regression.py @@ -0,0 +1,89 @@ +# +# 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. +# + +""" +A logistic regression implementation that uses NumPy (http://www.numpy.org) +to act on batches of input data using efficient matrix operations. + +In practice, one may prefer to use the LogisticRegression algorithm in +ML, as shown in examples/src/main/python/ml/logistic_regression_with_elastic_net.py. +""" +import sys +from typing import Iterable, List + + +import numpy as np +from pyspark.sql import SparkSession + + +D = 10 # Number of dimensions + + +# Read a batch of points from the input file into a NumPy matrix object. We operate on batches to +# make further computations faster. +# The data file contains lines of the form

MOZILLA PUBLIC LICENSE
Version + 1.1 +

+


+
+

1. Definitions. +

    1.0.1. "Commercial Use" means distribution or otherwise making the + Covered Code available to a third party. +

    1.1. ''Contributor'' means each entity that creates or contributes + to the creation of Modifications. +

    1.2. ''Contributor Version'' means the combination of the Original + Code, prior Modifications used by a Contributor, and the Modifications made by + that particular Contributor. +

    1.3. ''Covered Code'' means the Original Code or Modifications or + the combination of the Original Code and Modifications, in each case including + portions thereof. +

    1.4. ''Electronic Distribution Mechanism'' means a mechanism + generally accepted in the software development community for the electronic + transfer of data. +

    1.5. ''Executable'' means Covered Code in any form other than Source + Code. +

    1.6. ''Initial Developer'' means the individual or entity identified + as the Initial Developer in the Source Code notice required by Exhibit + A. +

    1.7. ''Larger Work'' means a work which combines Covered Code or + portions thereof with code not governed by the terms of this License. +

    1.8. ''License'' means this document. +

    1.8.1. "Licensable" means having the right to grant, to the maximum + extent possible, whether at the time of the initial grant or subsequently + acquired, any and all of the rights conveyed herein. +

    1.9. ''Modifications'' means any addition to or deletion from the + substance or structure of either the Original Code or any previous + Modifications. When Covered Code is released as a series of files, a + Modification is: +

      A. Any addition to or deletion from the contents of a file + containing Original Code or previous Modifications. +

      B. Any new file that contains any part of the Original Code or + previous Modifications.
       

    1.10. ''Original Code'' +means Source Code of computer software code which is described in the Source +Code notice required by Exhibit A as Original Code, and which, at the +time of its release under this License is not already Covered Code governed by +this License. +

    1.10.1. "Patent Claims" means any patent claim(s), now owned or + hereafter acquired, including without limitation,  method, process, and + apparatus claims, in any patent Licensable by grantor. +

    1.11. ''Source Code'' means the preferred form of the Covered Code + for making modifications to it, including all modules it contains, plus any + associated interface definition files, scripts used to control compilation and + installation of an Executable, or source code differential comparisons against + either the Original Code or another well known, available Covered Code of the + Contributor's choice. The Source Code can be in a compressed or archival form, + provided the appropriate decompression or de-archiving software is widely + available for no charge. +

    1.12. "You'' (or "Your")  means an individual or a legal entity + exercising rights under, and complying with all of the terms of, this License + or a future version of this License issued under Section 6.1. For legal + entities, "You'' includes any entity which controls, is controlled by, or is + under common control with You. For purposes of this definition, "control'' + means (a) the power, direct or indirect, to cause the direction or management + of such entity, whether by contract or otherwise, or (b) ownership of more + than fifty percent (50%) of the outstanding shares or beneficial ownership of + such entity.

2. Source Code License. +
    2.1. The Initial Developer Grant.
    The Initial Developer hereby + grants You a world-wide, royalty-free, non-exclusive license, subject to third + party intellectual property claims: +
      (a)  under intellectual property rights (other than + patent or trademark) Licensable by Initial Developer to use, reproduce, + modify, display, perform, sublicense and distribute the Original Code (or + portions thereof) with or without Modifications, and/or as part of a Larger + Work; and +

      (b) under Patents Claims infringed by the making, using or selling + of Original Code, to make, have made, use, practice, sell, and offer for + sale, and/or otherwise dispose of the Original Code (or portions thereof). +

        +
        (c) the licenses granted in this Section 2.1(a) and (b) + are effective on the date Initial Developer first distributes Original Code + under the terms of this License. +

        (d) Notwithstanding Section 2.1(b) above, no patent license is + granted: 1) for code that You delete from the Original Code; 2) separate + from the Original Code;  or 3) for infringements caused by: i) the + modification of the Original Code or ii) the combination of the Original + Code with other software or devices.
         

      2.2. Contributor + Grant.
      Subject to third party intellectual property claims, each + Contributor hereby grants You a world-wide, royalty-free, non-exclusive + license +

        (a)  under intellectual property rights (other + than patent or trademark) Licensable by Contributor, to use, reproduce, + modify, display, perform, sublicense and distribute the Modifications + created by such Contributor (or portions thereof) either on an unmodified + basis, with other Modifications, as Covered Code and/or as part of a Larger + Work; and +

        (b) under Patent Claims infringed by the making, using, or selling + of  Modifications made by that Contributor either alone and/or in combination with its Contributor Version (or portions of such + combination), to make, use, sell, offer for sale, have made, and/or + otherwise dispose of: 1) Modifications made by that Contributor (or portions + thereof); and 2) the combination of  Modifications made by that + Contributor with its Contributor Version (or portions of such + combination). +

        (c) the licenses granted in Sections 2.2(a) and 2.2(b) are + effective on the date Contributor first makes Commercial Use of the Covered + Code. +

        (d)    Notwithstanding Section 2.2(b) above, no + patent license is granted: 1) for any code that Contributor has deleted from + the Contributor Version; 2)  separate from the Contributor + Version;  3)  for infringements caused by: i) third party + modifications of Contributor Version or ii)  the combination of + Modifications made by that Contributor with other software  (except as + part of the Contributor Version) or other devices; or 4) under Patent Claims + infringed by Covered Code in the absence of Modifications made by that + Contributor.

    +


    3. Distribution Obligations. +

      3.1. Application of License.
      The Modifications which You create + or to which You contribute are governed by the terms of this License, + including without limitation Section 2.2. The Source Code version of + Covered Code may be distributed only under the terms of this License or a + future version of this License released under Section 6.1, and You must + include a copy of this License with every copy of the Source Code You + distribute. You may not offer or impose any terms on any Source Code version + that alters or restricts the applicable version of this License or the + recipients' rights hereunder. However, You may include an additional document + offering the additional rights described in Section 3.5. +

      3.2. Availability of Source Code.
      Any Modification which You + create or to which You contribute must be made available in Source Code form + under the terms of this License either on the same media as an Executable + version or via an accepted Electronic Distribution Mechanism to anyone to whom + you made an Executable version available; and if made available via Electronic + Distribution Mechanism, must remain available for at least twelve (12) months + after the date it initially became available, or at least six (6) months after + a subsequent version of that particular Modification has been made available + to such recipients. You are responsible for ensuring that the Source Code + version remains available even if the Electronic Distribution Mechanism is + maintained by a third party. +

      3.3. Description of Modifications.
      You must cause all Covered + Code to which You contribute to contain a file documenting the changes You + made to create that Covered Code and the date of any change. You must include + a prominent statement that the Modification is derived, directly or + indirectly, from Original Code provided by the Initial Developer and including + the name of the Initial Developer in (a) the Source Code, and (b) in any + notice in an Executable version or related documentation in which You describe + the origin or ownership of the Covered Code. +

      3.4. Intellectual Property Matters +

        (a) Third Party Claims.
        If Contributor has knowledge that a + license under a third party's intellectual property rights is required to + exercise the rights granted by such Contributor under Sections 2.1 or 2.2, + Contributor must include a text file with the Source Code distribution + titled "LEGAL'' which describes the claim and the party making the claim in + sufficient detail that a recipient will know whom to contact. If Contributor + obtains such knowledge after the Modification is made available as described + in Section 3.2, Contributor shall promptly modify the LEGAL file in all + copies Contributor makes available thereafter and shall take other steps + (such as notifying appropriate mailing lists or newsgroups) reasonably + calculated to inform those who received the Covered Code that new knowledge + has been obtained. +

        (b) Contributor APIs.
        If Contributor's Modifications include + an application programming interface and Contributor has knowledge of patent + licenses which are reasonably necessary to implement that API, Contributor + must also include this information in the LEGAL file. +
         

                +(c)    Representations. +
        Contributor represents that, except as disclosed pursuant to Section + 3.4(a) above, Contributor believes that Contributor's Modifications are + Contributor's original creation(s) and/or Contributor has sufficient rights + to grant the rights conveyed by this License.
      +


      3.5. Required Notices.
      You must duplicate the notice in + Exhibit A in each file of the Source Code.  If it is not possible + to put such notice in a particular Source Code file due to its structure, then + You must include such notice in a location (such as a relevant directory) + where a user would be likely to look for such a notice.  If You created + one or more Modification(s) You may add your name as a Contributor to the + notice described in Exhibit A.  You must also duplicate this + License in any documentation for the Source Code where You describe + recipients' rights or ownership rights relating to Covered Code.  You may + choose to offer, and to charge a fee for, warranty, support, indemnity or + liability obligations to one or more recipients of Covered Code. However, You + may do so only on Your own behalf, and not on behalf of the Initial Developer + or any Contributor. You must make it absolutely clear than any such warranty, + support, indemnity or liability obligation is offered by You alone, and You + hereby agree to indemnify the Initial Developer and every Contributor for any + liability incurred by the Initial Developer or such Contributor as a result of + warranty, support, indemnity or liability terms You offer. +

      3.6. Distribution of Executable Versions.
      You may distribute + Covered Code in Executable form only if the requirements of Section + 3.1-3.5 have been met for that Covered Code, and if You include a + notice stating that the Source Code version of the Covered Code is available + under the terms of this License, including a description of how and where You + have fulfilled the obligations of Section 3.2. The notice must be + conspicuously included in any notice in an Executable version, related + documentation or collateral in which You describe recipients' rights relating + to the Covered Code. You may distribute the Executable version of Covered Code + or ownership rights under a license of Your choice, which may contain terms + different from this License, provided that You are in compliance with the + terms of this License and that the license for the Executable version does not + attempt to limit or alter the recipient's rights in the Source Code version + from the rights set forth in this License. If You distribute the Executable + version under a different license You must make it absolutely clear that any + terms which differ from this License are offered by You alone, not by the + Initial Developer or any Contributor. You hereby agree to indemnify the + Initial Developer and every Contributor for any liability incurred by the + Initial Developer or such Contributor as a result of any such terms You offer. + +

      3.7. Larger Works.
      You may create a Larger Work by combining + Covered Code with other code not governed by the terms of this License and + distribute the Larger Work as a single product. In such a case, You must make + sure the requirements of this License are fulfilled for the Covered + Code.

    4. Inability to Comply Due to Statute or Regulation. +
      If it is impossible for You to comply with any of the terms of this + License with respect to some or all of the Covered Code due to statute, + judicial order, or regulation then You must: (a) comply with the terms of this + License to the maximum extent possible; and (b) describe the limitations and + the code they affect. Such description must be included in the LEGAL file + described in Section 3.4 and must be included with all distributions of + the Source Code. Except to the extent prohibited by statute or regulation, + such description must be sufficiently detailed for a recipient of ordinary + skill to be able to understand it.
    5. Application of this License. +
      This License applies to code to which the Initial Developer has attached + the notice in Exhibit A and to related Covered Code.
    6. Versions + of the License. +
      6.1. New Versions.
      Netscape Communications Corporation + (''Netscape'') may publish revised and/or new versions of the License from + time to time. Each version will be given a distinguishing version number. +

      6.2. Effect of New Versions.
      Once Covered Code has been + published under a particular version of the License, You may always continue + to use it under the terms of that version. You may also choose to use such + Covered Code under the terms of any subsequent version of the License + published by Netscape. No one other than Netscape has the right to modify the + terms applicable to Covered Code created under this License. +

      6.3. Derivative Works.
      If You create or use a modified version + of this License (which you may only do in order to apply it to code which is + not already Covered Code governed by this License), You must (a) rename Your + license so that the phrases ''Mozilla'', ''MOZILLAPL'', ''MOZPL'', + ''Netscape'', "MPL", ''NPL'' or any confusingly similar phrase do not appear + in your license (except to note that your license differs from this License) + and (b) otherwise make it clear that Your version of the license contains + terms which differ from the Mozilla Public License and Netscape Public + License. (Filling in the name of the Initial Developer, Original Code or + Contributor in the notice described in Exhibit A shall not of + themselves be deemed to be modifications of this License.)

    7. + DISCLAIMER OF WARRANTY. +
      COVERED CODE IS PROVIDED UNDER THIS LICENSE ON AN "AS IS'' BASIS, WITHOUT + WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, WITHOUT + LIMITATION, WARRANTIES THAT THE COVERED CODE IS FREE OF DEFECTS, MERCHANTABLE, + FIT FOR A PARTICULAR PURPOSE OR NON-INFRINGING. THE ENTIRE RISK AS TO THE + QUALITY AND PERFORMANCE OF THE COVERED CODE IS WITH YOU. SHOULD ANY COVERED + CODE PROVE DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL DEVELOPER OR ANY + OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY SERVICING, REPAIR OR + CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL PART OF THIS + LICENSE. NO USE OF ANY COVERED CODE IS AUTHORIZED HEREUNDER EXCEPT UNDER THIS + DISCLAIMER.
    8. TERMINATION. +
      8.1.  This License and the rights granted hereunder will + terminate automatically if You fail to comply with terms herein and fail to + cure such breach within 30 days of becoming aware of the breach. All + sublicenses to the Covered Code which are properly granted shall survive any + termination of this License. Provisions which, by their nature, must remain in + effect beyond the termination of this License shall survive. +

      8.2.  If You initiate litigation by asserting a patent + infringement claim (excluding declatory judgment actions) against Initial + Developer or a Contributor (the Initial Developer or Contributor against whom + You file such action is referred to as "Participant")  alleging that: +

      (a)  such Participant's Contributor Version directly or + indirectly infringes any patent, then any and all rights granted by such + Participant to You under Sections 2.1 and/or 2.2 of this License shall, upon + 60 days notice from Participant terminate prospectively, unless if within 60 + days after receipt of notice You either: (i)  agree in writing to pay + Participant a mutually agreeable reasonable royalty for Your past and future + use of Modifications made by such Participant, or (ii) withdraw Your + litigation claim with respect to the Contributor Version against such + Participant.  If within 60 days of notice, a reasonable royalty and + payment arrangement are not mutually agreed upon in writing by the parties or + the litigation claim is not withdrawn, the rights granted by Participant to + You under Sections 2.1 and/or 2.2 automatically terminate at the expiration of + the 60 day notice period specified above. +

      (b)  any software, hardware, or device, other than such + Participant's Contributor Version, directly or indirectly infringes any + patent, then any rights granted to You by such Participant under Sections + 2.1(b) and 2.2(b) are revoked effective as of the date You first made, used, + sold, distributed, or had made, Modifications made by that Participant. +

      8.3.  If You assert a patent infringement claim against + Participant alleging that such Participant's Contributor Version directly or + indirectly infringes any patent where such claim is resolved (such as by + license or settlement) prior to the initiation of patent infringement + litigation, then the reasonable value of the licenses granted by such + Participant under Sections 2.1 or 2.2 shall be taken into account in + determining the amount or value of any payment or license. +

      8.4.  In the event of termination under Sections 8.1 or 8.2 + above,  all end user license agreements (excluding distributors and + resellers) which have been validly granted by You or any distributor hereunder + prior to termination shall survive termination.

    9. LIMITATION OF + LIABILITY. +
      UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL THEORY, WHETHER TORT (INCLUDING + NEGLIGENCE), CONTRACT, OR OTHERWISE, SHALL YOU, THE INITIAL DEVELOPER, ANY + OTHER CONTRIBUTOR, OR ANY DISTRIBUTOR OF COVERED CODE, OR ANY SUPPLIER OF ANY + OF SUCH PARTIES, BE LIABLE TO ANY PERSON FOR ANY INDIRECT, SPECIAL, + INCIDENTAL, OR CONSEQUENTIAL DAMAGES OF ANY CHARACTER INCLUDING, WITHOUT + LIMITATION, DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER FAILURE OR + MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR LOSSES, EVEN IF SUCH + PARTY SHALL HAVE BEEN INFORMED OF THE POSSIBILITY OF SUCH DAMAGES. THIS + LIMITATION OF LIABILITY SHALL NOT APPLY TO LIABILITY FOR DEATH OR PERSONAL + INJURY RESULTING FROM SUCH PARTY'S NEGLIGENCE TO THE EXTENT APPLICABLE LAW + PROHIBITS SUCH LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR + LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION AND + LIMITATION MAY NOT APPLY TO YOU.
    10. U.S. GOVERNMENT END USERS. +
      The Covered Code is a ''commercial item,'' as that term is defined in 48 + C.F.R. 2.101 (Oct. 1995), consisting of ''commercial computer software'' and + ''commercial computer software documentation,'' as such terms are used in 48 + C.F.R. 12.212 (Sept. 1995). Consistent with 48 C.F.R. 12.212 and 48 C.F.R. + 227.7202-1 through 227.7202-4 (June 1995), all U.S. Government End Users + acquire Covered Code with only those rights set forth herein.
    11. + MISCELLANEOUS. +
      This License represents the complete agreement concerning subject matter + hereof. If any provision of this License is held to be unenforceable, such + provision shall be reformed only to the extent necessary to make it + enforceable. This License shall be governed by California law provisions + (except to the extent applicable law, if any, provides otherwise), excluding + its conflict-of-law provisions. With respect to disputes in which at least one + party is a citizen of, or an entity chartered or registered to do business in + the United States of America, any litigation relating to this License shall be + subject to the jurisdiction of the Federal Courts of the Northern District of + California, with venue lying in Santa Clara County, California, with the + losing party responsible for costs, including without limitation, court costs + and reasonable attorneys' fees and expenses. The application of the United + Nations Convention on Contracts for the International Sale of Goods is + expressly excluded. Any law or regulation which provides that the language of + a contract shall be construed against the drafter shall not apply to this + License.
    12. RESPONSIBILITY FOR CLAIMS. +
      As between Initial Developer and the Contributors, each party is + responsible for claims and damages arising, directly or indirectly, out of its + utilization of rights under this License and You agree to work with Initial + Developer and Contributors to distribute such responsibility on an equitable + basis. Nothing herein is intended or shall be deemed to constitute any + admission of liability.
    13. MULTIPLE-LICENSED CODE. +
      Initial Developer may designate portions of the Covered Code as + "Multiple-Licensed".  "Multiple-Licensed" means that the Initial + Developer permits you to utilize portions of the Covered Code under Your + choice of the MPL or the alternative licenses, if any, specified by the + Initial Developer in the file described in Exhibit A.
    +


    EXHIBIT A -Mozilla Public License. +

      The contents of this file are subject to the Mozilla Public License + Version 1.1 (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.mozilla.org/MPL/ +

      Software distributed under the License is distributed on an "AS IS" basis, + WITHOUT WARRANTY OF
      ANY KIND, either express or implied. See the License + for the specific language governing rights and
      limitations under the + License. +

      The Original Code is Javassist. +

      The Initial Developer of the Original Code is Shigeru Chiba. + Portions created by the Initial Developer are
        + Copyright (C) 1999- Shigeru Chiba. All Rights Reserved. +

      Contributor(s): __Bill Burke, Jason T. Greene______________. + +

      Alternatively, the contents of this software may be used under the + terms of the GNU Lesser General Public License Version 2.1 or later + (the "LGPL"), or the Apache License Version 2.0 (the "AL"), + in which case the provisions of the LGPL or the AL are applicable + instead of those above. If you wish to allow use of your version of + this software only under the terms of either the LGPL or the AL, and not to allow others to + use your version of this software under the terms of the MPL, indicate + your decision by deleting the provisions above and replace them with + the notice and other provisions required by the LGPL or the AL. If you do not + delete the provisions above, a recipient may use your version of this + software under the terms of any one of the MPL, the LGPL or the AL. + +

    + + \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javax-transaction-transaction-api.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javax-transaction-transaction-api.txt new file mode 100644 index 0000000..3d48d9c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javax-transaction-transaction-api.txt @@ -0,0 +1,119 @@ +COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL) Version 1.0 + +1. Definitions. + +1.1. Contributor means each individual or entity that creates or contributes to the creation of Modifications. + +1.2. Contributor Version means the combination of the Original Software, prior Modifications used by a Contributor (if any), and the Modifications made by that particular Contributor. + +1.3. Covered Software means (a) the Original Software, or (b) Modifications, or (c) the combination of files containing Original Software with files containing Modifications, in each case including portions thereof. + +1.4. Executable means the Covered Software in any form other than Source Code. + +1.5. Initial Developer means the individual or entity that first makes Original Software available under this License. + +1.6. Larger Work means a work which combines Covered Software or portions thereof with code not governed by the terms of this License. + +1.7. License means this document. + +1.8. Licensable means having the right to grant, to the maximum extent possible, whether at the time of the initial grant or subsequently acquired, any and all of the rights conveyed herein. + +1.9. Modifications means the Source Code and Executable form of any of the following: + +A. Any file that results from an addition to, deletion from or modification of the contents of a file containing Original Software or previous Modifications; + +B. Any new file that contains any part of the Original Software or previous Modification; or + +C. Any new file that is contributed or otherwise made available under the terms of this License. + +1.10. Original Software means the Source Code and Executable form of computer software code that is originally released under this License. + +1.11. Patent Claims means any patent claim(s), now owned or hereafter acquired, including without limitation, method, process, and apparatus claims, in any patent Licensable by grantor. + +1.12. Source Code means (a) the common form of computer software code in which modifications are made and (b) associated documentation included in or with such code. + +1.13. You (or Your) means an individual or a legal entity exercising rights under, and complying with all of the terms of, this License. For legal entities, You includes any entity which controls, is controlled by, or is under common control with You. For purposes of this definition, control means (a) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (b) ownership of more than fifty percent (50%) of the outstanding shares or beneficial ownership of such entity. + +2. License Grants. + +2.1. The Initial Developer Grant. +Conditioned upon Your compliance with Section 3.1 below and subject to third party intellectual property claims, the Initial Developer hereby grants You a world-wide, royalty-free, non-exclusive license: +(a) under intellectual property rights (other than patent or trademark) Licensable by Initial Developer, to use, reproduce, modify, display, perform, sublicense and distribute the Original Software (or portions thereof), with or without Modifications, and/or as part of a Larger Work; and +(b) under Patent Claims infringed by the making, using or selling of Original Software, to make, have made, use, practice, sell, and offer for sale, and/or otherwise dispose of the Original Software (or portions thereof). +(c) The licenses granted in Sections 2.1(a) and (b) are effective on the date Initial Developer first distributes or otherwise makes the Original Software available to a third party under the terms of this License. +(d) Notwithstanding Section 2.1(b) above, no patent license is granted: (1) for code that You delete from the Original Software, or (2) for infringements caused by: (i) the modification of the Original Software, or (ii) the combination of the Original Software with other software or devices. + +2.2. Contributor Grant. +Conditioned upon Your compliance with Section 3.1 below and subject to third party intellectual property claims, each Contributor hereby grants You a world-wide, royalty-free, non-exclusive license: +(a) under intellectual property rights (other than patent or trademark) Licensable by Contributor to use, reproduce, modify, display, perform, sublicense and distribute the Modifications created by such Contributor (or portions thereof), either on an unmodified basis, with other Modifications, as Covered Software and/or as part of a Larger Work; and +(b) under Patent Claims infringed by the making, using, or selling of Modifications made by that Contributor either alone and/or in combination with its Contributor Version (or portions of such combination), to make, use, sell, offer for sale, have made, and/or otherwise dispose of: (1) Modifications made by that Contributor (or portions thereof); and (2) the combination of Modifications made by that Contributor with its Contributor Version (or portions of such combination). +(c) The licenses granted in Sections 2.2(a) and 2.2(b) are effective on the date Contributor first distributes or otherwise makes the Modifications available to a third party. +(d) Notwithstanding Section 2.2(b) above, no patent license is granted: (1) for any code that Contributor has deleted from the Contributor Version; (2) for infringements caused by: (i) third party modifications of Contributor Version, or (ii) the combination of Modifications made by that Contributor with other software (except as part of the Contributor Version) or other devices; or (3) under Patent Claims infringed by Covered Software in the absence of Modifications made by that Contributor. + +3. Distribution Obligations. + +3.1. Availability of Source Code. + +Any Covered Software that You distribute or otherwise make available in Executable form must also be made available in Source Code form and that Source Code form must be distributed only under the terms of this License. You must include a copy of this License with every copy of the Source Code form of the Covered Software You distribute or otherwise make available. You must inform recipients of any such Covered Software in Executable form as to how they can obtain such Covered Software in Source Code form in a reasonable manner on or through a medium customarily used for software exchange. + +3.2. Modifications. + +The Modifications that You create or to which You contribute are governed by the terms of this License. You represent that You believe Your Modifications are Your original creation(s) and/or You have sufficient rights to grant the rights conveyed by this License. + +3.3. Required Notices. +You must include a notice in each of Your Modifications that identifies You as the Contributor of the Modification. You may not remove or alter any copyright, patent or trademark notices contained within the Covered Software, or any notices of licensing or any descriptive text giving attribution to any Contributor or the Initial Developer. + +3.4. Application of Additional Terms. +You may not offer or impose any terms on any Covered Software in Source Code form that alters or restricts the applicable version of this License or the recipients rights hereunder. You may choose to offer, and to charge a fee for, warranty, support, indemnity or liability obligations to one or more recipients of Covered Software. However, you may do so only on Your own behalf, and not on behalf of the Initial Developer or any Contributor. You must make it absolutely clear that any such warranty, support, indemnity or liability obligation is offered by You alone, and You hereby agree to indemnify the Initial Developer and every Contributor for any liability incurred by the Initial Developer or such Contributor as a result of warranty, support, indemnity or liability terms You offer. + +3.5. Distribution of Executable Versions. +You may distribute the Executable form of the Covered Software under the terms of this License or under the terms of a license of Your choice, which may contain terms different from this License, provided that You are in compliance with the terms of this License and that the license for the Executable form does not attempt to limit or alter the recipients rights in the Source Code form from the rights set forth in this License. If You distribute the Covered Software in Executable form under a different license, You must make it absolutely clear that any terms which differ from this License are offered by You alone, not by the Initial Developer or Contributor. You hereby agree to indemnify the Initial Developer and every Contributor for any liability incurred by the Initial Developer or such Contributor as a result of any such terms You offer. + +3.6. Larger Works. +You may create a Larger Work by combining Covered Software with other code not governed by the terms of this License and distribute the Larger Work as a single product. In such a case, You must make sure the requirements of this License are fulfilled for the Covered Software. + +4. Versions of the License. + +4.1. New Versions. +Sun Microsystems, Inc. is the initial license steward and may publish revised and/or new versions of this License from time to time. Each version will be given a distinguishing version number. Except as provided in Section 4.3, no one other than the license steward has the right to modify this License. + +4.2. Effect of New Versions. + +You may always continue to use, distribute or otherwise make the Covered Software available under the terms of the version of the License under which You originally received the Covered Software. If the Initial Developer includes a notice in the Original Software prohibiting it from being distributed or otherwise made available under any subsequent version of the License, You must distribute and make the Covered Software available under the terms of the version of the License under which You originally received the Covered Software. Otherwise, You may also choose to use, distribute or otherwise make the Covered Software available under the terms of any subsequent version of the License published by the license steward. +4.3. Modified Versions. + +When You are an Initial Developer and You want to create a new license for Your Original Software, You may create and use a modified version of this License if You: (a) rename the license and remove any references to the name of the license steward (except to note that the license differs from this License); and (b) otherwise make it clear that the license contains terms which differ from this License. + +5. DISCLAIMER OF WARRANTY. + +COVERED SOFTWARE IS PROVIDED UNDER THIS LICENSE ON AN AS IS BASIS, WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, WITHOUT LIMITATION, WARRANTIES THAT THE COVERED SOFTWARE IS FREE OF DEFECTS, MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE OR NON-INFRINGING. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE COVERED SOFTWARE IS WITH YOU. SHOULD ANY COVERED SOFTWARE PROVE DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL DEVELOPER OR ANY OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY SERVICING, REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL PART OF THIS LICENSE. NO USE OF ANY COVERED SOFTWARE IS AUTHORIZED HEREUNDER EXCEPT UNDER THIS DISCLAIMER. + +6. TERMINATION. + +6.1. This License and the rights granted hereunder will terminate automatically if You fail to comply with terms herein and fail to cure such breach within 30 days of becoming aware of the breach. Provisions which, by their nature, must remain in effect beyond the termination of this License shall survive. + +6.2. If You assert a patent infringement claim (excluding declaratory judgment actions) against Initial Developer or a Contributor (the Initial Developer or Contributor against whom You assert such claim is referred to as Participant) alleging that the Participant Software (meaning the Contributor Version where the Participant is a Contributor or the Original Software where the Participant is the Initial Developer) directly or indirectly infringes any patent, then any and all rights granted directly or indirectly to You by such Participant, the Initial Developer (if the Initial Developer is not the Participant) and all Contributors under Sections 2.1 and/or 2.2 of this License shall, upon 60 days notice from Participant terminate prospectively and automatically at the expiration of such 60 day notice period, unless if within such 60 day period You withdraw Your claim with respect to the Participant Software against such Participant either unilaterally or pursuant to a written agreement with Participant. + +6.3. In the event of termination under Sections 6.1 or 6.2 above, all end user licenses that have been validly granted by You or any distributor hereunder prior to termination (excluding licenses granted to You by any distributor) shall survive termination. + +7. LIMITATION OF LIABILITY. + +UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL THEORY, WHETHER TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE, SHALL YOU, THE INITIAL DEVELOPER, ANY OTHER CONTRIBUTOR, OR ANY DISTRIBUTOR OF COVERED SOFTWARE, OR ANY SUPPLIER OF ANY OF SUCH PARTIES, BE LIABLE TO ANY PERSON FOR ANY INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES OF ANY CHARACTER INCLUDING, WITHOUT LIMITATION, DAMAGES FOR LOST PROFITS, LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR LOSSES, EVEN IF SUCH PARTY SHALL HAVE BEEN INFORMED OF THE POSSIBILITY OF SUCH DAMAGES. THIS LIMITATION OF LIABILITY SHALL NOT APPLY TO LIABILITY FOR DEATH OR PERSONAL INJURY RESULTING FROM SUCH PARTYS NEGLIGENCE TO THE EXTENT APPLICABLE LAW PROHIBITS SUCH LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION AND LIMITATION MAY NOT APPLY TO YOU. + +8. U.S. GOVERNMENT END USERS. + +The Covered Software is a commercial item, as that term is defined in 48 C.F.R. 2.101 (Oct. 1995), consisting of commercial computer software (as that term is defined at 48 C.F.R. 252.227-7014(a)(1)) and commercial computer software documentation as such terms are used in 48 C.F.R. 12.212 (Sept. 1995). Consistent with 48 C.F.R. 12.212 and 48 C.F.R. 227.7202-1 through 227.7202-4 (June 1995), all U.S. Government End Users acquire Covered Software with only those rights set forth herein. This U.S. Government Rights clause is in lieu of, and supersedes, any other FAR, DFAR, or other clause or provision that addresses Government rights in computer software under this License. + +9. MISCELLANEOUS. + +This License represents the complete agreement concerning subject matter hereof. If any provision of this License is held to be unenforceable, such provision shall be reformed only to the extent necessary to make it enforceable. This License shall be governed by the law of the jurisdiction specified in a notice contained within the Original Software (except to the extent applicable law, if any, provides otherwise), excluding such jurisdictions conflict-of-law provisions. Any litigation relating to this License shall be subject to the jurisdiction of the courts located in the jurisdiction and venue specified in a notice contained within the Original Software, with the losing party responsible for costs, including, without limitation, court costs and reasonable attorneys fees and expenses. The application of the United Nations Convention on Contracts for the International Sale of Goods is expressly excluded. Any law or regulation which provides that the language of a contract shall be construed against the drafter shall not apply to this License. You agree that You alone are responsible for compliance with the United States export administration regulations (and the export control laws and regulation of any other countries) when You use, distribute or otherwise make available any Covered Software. + +10. RESPONSIBILITY FOR CLAIMS. + +As between Initial Developer and the Contributors, each party is responsible for claims and damages arising, directly or indirectly, out of its utilization of rights under this License and You agree to work with Initial Developer and Contributors to distribute such responsibility on an equitable basis. Nothing herein is intended or shall be deemed to constitute any admission of liability. + +NOTICE PURSUANT TO SECTION 9 OF THE COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL) +The GlassFish code released under the CDDL shall be governed by the laws of the State of California (excluding conflict-of-law provisions). Any litigation relating to this License shall be subject to the jurisdiction of the Federal Courts of the Northern District of California and the state courts of the State of California, with venue lying in Santa Clara County, California. + + + diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javolution.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javolution.txt new file mode 100644 index 0000000..b64af4d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-javolution.txt @@ -0,0 +1,27 @@ +/* + * Javolution - Java(tm) Solution for Real-Time and Embedded Systems + * Copyright (c) 2012, Javolution (http://javolution.org/) + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF + * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING + * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jaxb-runtime.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jaxb-runtime.txt new file mode 100644 index 0000000..02319e9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jaxb-runtime.txt @@ -0,0 +1,11 @@ +Eclipse Distribution License - v 1.0 +Copyright (c) 2007, Eclipse Foundation, Inc. and its licensors. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + +Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. +Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. +Neither the name of the Eclipse Foundation, Inc. nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jline.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jline.txt new file mode 100644 index 0000000..2ec539d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jline.txt @@ -0,0 +1,32 @@ +Copyright (c) 2002-2006, Marc Prud'hommeaux +All rights reserved. + +Redistribution and use in source and binary forms, with or +without modification, are permitted provided that the following +conditions are met: + +Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + +Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with +the distribution. + +Neither the name of JLine nor the names of its contributors +may be used to endorse or promote products derived from this +software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, +BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY +AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO +EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, +OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED +AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING +IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED +OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jodd.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jodd.txt new file mode 100644 index 0000000..cc6b458 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jodd.txt @@ -0,0 +1,24 @@ +Copyright (c) 2003-present, Jodd Team (https://jodd.org) +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, +this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in the +documentation and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-join.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-join.txt new file mode 100644 index 0000000..1d91609 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-join.txt @@ -0,0 +1,30 @@ +Copyright (c) 2011, Douban Inc. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + + * Neither the name of the Douban Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jquery.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jquery.txt new file mode 100644 index 0000000..4593054 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jquery.txt @@ -0,0 +1,20 @@ +Copyright JS Foundation and other contributors, https://js.foundation/ + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-json-formatter.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-json-formatter.txt new file mode 100644 index 0000000..5193348 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-json-formatter.txt @@ -0,0 +1,6 @@ +Copyright 2014 Mohsen Azimi + +Licensed 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. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jsp-api.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jsp-api.txt new file mode 100644 index 0000000..68076ad --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-jsp-api.txt @@ -0,0 +1,759 @@ +COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL) Version 1.1 + +1. Definitions. + + 1.1. "Contributor" means each individual or entity that creates or + contributes to the creation of Modifications. + + 1.2. "Contributor Version" means the combination of the Original + Software, prior Modifications used by a Contributor (if any), and + the Modifications made by that particular Contributor. + + 1.3. "Covered Software" means (a) the Original Software, or (b) + Modifications, or (c) the combination of files containing Original + Software with files containing Modifications, in each case including + portions thereof. + + 1.4. "Executable" means the Covered Software in any form other than + Source Code. + + 1.5. "Initial Developer" means the individual or entity that first + makes Original Software available under this License. + + 1.6. "Larger Work" means a work which combines Covered Software or + portions thereof with code not governed by the terms of this License. + + 1.7. "License" means this document. + + 1.8. "Licensable" means having the right to grant, to the maximum + extent possible, whether at the time of the initial grant or + subsequently acquired, any and all of the rights conveyed herein. + + 1.9. "Modifications" means the Source Code and Executable form of + any of the following: + + A. Any file that results from an addition to, deletion from or + modification of the contents of a file containing Original Software + or previous Modifications; + + B. Any new file that contains any part of the Original Software or + previous Modification; or + + C. Any new file that is contributed or otherwise made available + under the terms of this License. + + 1.10. "Original Software" means the Source Code and Executable form + of computer software code that is originally released under this + License. + + 1.11. "Patent Claims" means any patent claim(s), now owned or + hereafter acquired, including without limitation, method, process, + and apparatus claims, in any patent Licensable by grantor. + + 1.12. "Source Code" means (a) the common form of computer software + code in which modifications are made and (b) associated + documentation included in or with such code. + + 1.13. "You" (or "Your") means an individual or a legal entity + exercising rights under, and complying with all of the terms of, + this License. For legal entities, "You" includes any entity which + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants. + + 2.1. The Initial Developer Grant. + + Conditioned upon Your compliance with Section 3.1 below and subject + to third party intellectual property claims, the Initial Developer + hereby grants You a world-wide, royalty-free, non-exclusive license: + + (a) under intellectual property rights (other than patent or + trademark) Licensable by Initial Developer, to use, reproduce, + modify, display, perform, sublicense and distribute the Original + Software (or portions thereof), with or without Modifications, + and/or as part of a Larger Work; and + + (b) under Patent Claims infringed by the making, using or selling of + Original Software, to make, have made, use, practice, sell, and + offer for sale, and/or otherwise dispose of the Original Software + (or portions thereof). + + (c) The licenses granted in Sections 2.1(a) and (b) are effective on + the date Initial Developer first distributes or otherwise makes the + Original Software available to a third party under the terms of this + License. + + (d) Notwithstanding Section 2.1(b) above, no patent license is + granted: (1) for code that You delete from the Original Software, or + (2) for infringements caused by: (i) the modification of the + Original Software, or (ii) the combination of the Original Software + with other software or devices. + + 2.2. Contributor Grant. + + Conditioned upon Your compliance with Section 3.1 below and subject + to third party intellectual property claims, each Contributor hereby + grants You a world-wide, royalty-free, non-exclusive license: + + (a) under intellectual property rights (other than patent or + trademark) Licensable by Contributor to use, reproduce, modify, + display, perform, sublicense and distribute the Modifications + created by such Contributor (or portions thereof), either on an + unmodified basis, with other Modifications, as Covered Software + and/or as part of a Larger Work; and + + (b) under Patent Claims infringed by the making, using, or selling + of Modifications made by that Contributor either alone and/or in + combination with its Contributor Version (or portions of such + combination), to make, use, sell, offer for sale, have made, and/or + otherwise dispose of: (1) Modifications made by that Contributor (or + portions thereof); and (2) the combination of Modifications made by + that Contributor with its Contributor Version (or portions of such + combination). + + (c) The licenses granted in Sections 2.2(a) and 2.2(b) are effective + on the date Contributor first distributes or otherwise makes the + Modifications available to a third party. + + (d) Notwithstanding Section 2.2(b) above, no patent license is + granted: (1) for any code that Contributor has deleted from the + Contributor Version; (2) for infringements caused by: (i) third + party modifications of Contributor Version, or (ii) the combination + of Modifications made by that Contributor with other software + (except as part of the Contributor Version) or other devices; or (3) + under Patent Claims infringed by Covered Software in the absence of + Modifications made by that Contributor. + +3. Distribution Obligations. + + 3.1. Availability of Source Code. + + Any Covered Software that You distribute or otherwise make available + in Executable form must also be made available in Source Code form + and that Source Code form must be distributed only under the terms + of this License. You must include a copy of this License with every + copy of the Source Code form of the Covered Software You distribute + or otherwise make available. You must inform recipients of any such + Covered Software in Executable form as to how they can obtain such + Covered Software in Source Code form in a reasonable manner on or + through a medium customarily used for software exchange. + + 3.2. Modifications. + + The Modifications that You create or to which You contribute are + governed by the terms of this License. You represent that You + believe Your Modifications are Your original creation(s) and/or You + have sufficient rights to grant the rights conveyed by this License. + + 3.3. Required Notices. + + You must include a notice in each of Your Modifications that + identifies You as the Contributor of the Modification. You may not + remove or alter any copyright, patent or trademark notices contained + within the Covered Software, or any notices of licensing or any + descriptive text giving attribution to any Contributor or the + Initial Developer. + + 3.4. Application of Additional Terms. + + You may not offer or impose any terms on any Covered Software in + Source Code form that alters or restricts the applicable version of + this License or the recipients' rights hereunder. You may choose to + offer, and to charge a fee for, warranty, support, indemnity or + liability obligations to one or more recipients of Covered Software. + However, you may do so only on Your own behalf, and not on behalf of + the Initial Developer or any Contributor. You must make it + absolutely clear that any such warranty, support, indemnity or + liability obligation is offered by You alone, and You hereby agree + to indemnify the Initial Developer and every Contributor for any + liability incurred by the Initial Developer or such Contributor as a + result of warranty, support, indemnity or liability terms You offer. + + 3.5. Distribution of Executable Versions. + + You may distribute the Executable form of the Covered Software under + the terms of this License or under the terms of a license of Your + choice, which may contain terms different from this License, + provided that You are in compliance with the terms of this License + and that the license for the Executable form does not attempt to + limit or alter the recipient's rights in the Source Code form from + the rights set forth in this License. If You distribute the Covered + Software in Executable form under a different license, You must make + it absolutely clear that any terms which differ from this License + are offered by You alone, not by the Initial Developer or + Contributor. You hereby agree to indemnify the Initial Developer and + every Contributor for any liability incurred by the Initial + Developer or such Contributor as a result of any such terms You offer. + + 3.6. Larger Works. + + You may create a Larger Work by combining Covered Software with + other code not governed by the terms of this License and distribute + the Larger Work as a single product. In such a case, You must make + sure the requirements of this License are fulfilled for the Covered + Software. + +4. Versions of the License. + + 4.1. New Versions. + + Oracle is the initial license steward and may publish revised and/or + new versions of this License from time to time. Each version will be + given a distinguishing version number. Except as provided in Section + 4.3, no one other than the license steward has the right to modify + this License. + + 4.2. Effect of New Versions. + + You may always continue to use, distribute or otherwise make the + Covered Software available under the terms of the version of the + License under which You originally received the Covered Software. If + the Initial Developer includes a notice in the Original Software + prohibiting it from being distributed or otherwise made available + under any subsequent version of the License, You must distribute and + make the Covered Software available under the terms of the version + of the License under which You originally received the Covered + Software. Otherwise, You may also choose to use, distribute or + otherwise make the Covered Software available under the terms of any + subsequent version of the License published by the license steward. + + 4.3. Modified Versions. + + When You are an Initial Developer and You want to create a new + license for Your Original Software, You may create and use a + modified version of this License if You: (a) rename the license and + remove any references to the name of the license steward (except to + note that the license differs from this License); and (b) otherwise + make it clear that the license contains terms which differ from this + License. + +5. DISCLAIMER OF WARRANTY. + + COVERED SOFTWARE IS PROVIDED UNDER THIS LICENSE ON AN "AS IS" BASIS, + WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, + INCLUDING, WITHOUT LIMITATION, WARRANTIES THAT THE COVERED SOFTWARE + IS FREE OF DEFECTS, MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE OR + NON-INFRINGING. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF + THE COVERED SOFTWARE IS WITH YOU. SHOULD ANY COVERED SOFTWARE PROVE + DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL DEVELOPER OR ANY + OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY SERVICING, + REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN + ESSENTIAL PART OF THIS LICENSE. NO USE OF ANY COVERED SOFTWARE IS + AUTHORIZED HEREUNDER EXCEPT UNDER THIS DISCLAIMER. + +6. TERMINATION. + + 6.1. This License and the rights granted hereunder will terminate + automatically if You fail to comply with terms herein and fail to + cure such breach within 30 days of becoming aware of the breach. + Provisions which, by their nature, must remain in effect beyond the + termination of this License shall survive. + + 6.2. If You assert a patent infringement claim (excluding + declaratory judgment actions) against Initial Developer or a + Contributor (the Initial Developer or Contributor against whom You + assert such claim is referred to as "Participant") alleging that the + Participant Software (meaning the Contributor Version where the + Participant is a Contributor or the Original Software where the + Participant is the Initial Developer) directly or indirectly + infringes any patent, then any and all rights granted directly or + indirectly to You by such Participant, the Initial Developer (if the + Initial Developer is not the Participant) and all Contributors under + Sections 2.1 and/or 2.2 of this License shall, upon 60 days notice + from Participant terminate prospectively and automatically at the + expiration of such 60 day notice period, unless if within such 60 + day period You withdraw Your claim with respect to the Participant + Software against such Participant either unilaterally or pursuant to + a written agreement with Participant. + + 6.3. If You assert a patent infringement claim against Participant + alleging that the Participant Software directly or indirectly + infringes any patent where such claim is resolved (such as by + license or settlement) prior to the initiation of patent + infringement litigation, then the reasonable value of the licenses + granted by such Participant under Sections 2.1 or 2.2 shall be taken + into account in determining the amount or value of any payment or + license. + + 6.4. In the event of termination under Sections 6.1 or 6.2 above, + all end user licenses that have been validly granted by You or any + distributor hereunder prior to termination (excluding licenses + granted to You by any distributor) shall survive termination. + +7. LIMITATION OF LIABILITY. + + UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL THEORY, WHETHER TORT + (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE, SHALL YOU, THE + INITIAL DEVELOPER, ANY OTHER CONTRIBUTOR, OR ANY DISTRIBUTOR OF + COVERED SOFTWARE, OR ANY SUPPLIER OF ANY OF SUCH PARTIES, BE LIABLE + TO ANY PERSON FOR ANY INDIRECT, SPECIAL, INCIDENTAL, OR + CONSEQUENTIAL DAMAGES OF ANY CHARACTER INCLUDING, WITHOUT + LIMITATION, DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER + FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR + LOSSES, EVEN IF SUCH PARTY SHALL HAVE BEEN INFORMED OF THE + POSSIBILITY OF SUCH DAMAGES. THIS LIMITATION OF LIABILITY SHALL NOT + APPLY TO LIABILITY FOR DEATH OR PERSONAL INJURY RESULTING FROM SUCH + PARTY'S NEGLIGENCE TO THE EXTENT APPLICABLE LAW PROHIBITS SUCH + LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR + LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION + AND LIMITATION MAY NOT APPLY TO YOU. + +8. U.S. GOVERNMENT END USERS. + + The Covered Software is a "commercial item," as that term is defined + in 48 C.F.R. 2.101 (Oct. 1995), consisting of "commercial computer + software" (as that term is defined at 48 C.F.R. § + 252.227-7014(a)(1)) and "commercial computer software documentation" + as such terms are used in 48 C.F.R. 12.212 (Sept. 1995). Consistent + with 48 C.F.R. 12.212 and 48 C.F.R. 227.7202-1 through 227.7202-4 + (June 1995), all U.S. Government End Users acquire Covered Software + with only those rights set forth herein. This U.S. Government Rights + clause is in lieu of, and supersedes, any other FAR, DFAR, or other + clause or provision that addresses Government rights in computer + software under this License. + +9. MISCELLANEOUS. + + This License represents the complete agreement concerning subject + matter hereof. If any provision of this License is held to be + unenforceable, such provision shall be reformed only to the extent + necessary to make it enforceable. This License shall be governed by + the law of the jurisdiction specified in a notice contained within + the Original Software (except to the extent applicable law, if any, + provides otherwise), excluding such jurisdiction's conflict-of-law + provisions. Any litigation relating to this License shall be subject + to the jurisdiction of the courts located in the jurisdiction and + venue specified in a notice contained within the Original Software, + with the losing party responsible for costs, including, without + limitation, court costs and reasonable attorneys' fees and expenses. + The application of the United Nations Convention on Contracts for + the International Sale of Goods is expressly excluded. Any law or + regulation which provides that the language of a contract shall be + construed against the drafter shall not apply to this License. You + agree that You alone are responsible for compliance with the United + States export administration regulations (and the export control + laws and regulation of any other countries) when You use, distribute + or otherwise make available any Covered Software. + +10. RESPONSIBILITY FOR CLAIMS. + + As between Initial Developer and the Contributors, each party is + responsible for claims and damages arising, directly or indirectly, + out of its utilization of rights under this License and You agree to + work with Initial Developer and Contributors to distribute such + responsibility on an equitable basis. Nothing herein is intended or + shall be deemed to constitute any admission of liability. + +------------------------------------------------------------------------ + +NOTICE PURSUANT TO SECTION 9 OF THE COMMON DEVELOPMENT AND DISTRIBUTION +LICENSE (CDDL) + +The code released under the CDDL shall be governed by the laws of the +State of California (excluding conflict-of-law provisions). Any +litigation relating to this License shall be subject to the jurisdiction +of the Federal Courts of the Northern District of California and the +state courts of the State of California, with venue lying in Santa Clara +County, California. + + + + The GNU General Public License (GPL) Version 2, June 1991 + +Copyright (C) 1989, 1991 Free Software Foundation, Inc. +51 Franklin Street, Fifth Floor +Boston, MA 02110-1335 +USA + +Everyone is permitted to copy and distribute verbatim copies +of this license document, but changing it is not allowed. + +Preamble + +The licenses for most software are designed to take away your freedom to +share and change it. By contrast, the GNU General Public License is +intended to guarantee your freedom to share and change free software--to +make sure the software is free for all its users. This General Public +License applies to most of the Free Software Foundation's software and +to any other program whose authors commit to using it. (Some other Free +Software Foundation software is covered by the GNU Library General +Public License instead.) You can apply it to your programs, too. + +When we speak of free software, we are referring to freedom, not price. +Our General Public Licenses are designed to make sure that you have the +freedom to distribute copies of free software (and charge for this +service if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs; and that you know you can do these things. + +To protect your rights, we need to make restrictions that forbid anyone +to deny you these rights or to ask you to surrender the rights. These +restrictions translate to certain responsibilities for you if you +distribute copies of the software, or if you modify it. + +For example, if you distribute copies of such a program, whether gratis +or for a fee, you must give the recipients all the rights that you have. +You must make sure that they, too, receive or can get the source code. +And you must show them these terms so they know their rights. + +We protect your rights with two steps: (1) copyright the software, and +(2) offer you this license which gives you legal permission to copy, +distribute and/or modify the software. + +Also, for each author's protection and ours, we want to make certain +that everyone understands that there is no warranty for this free +software. If the software is modified by someone else and passed on, we +want its recipients to know that what they have is not the original, so +that any problems introduced by others will not reflect on the original +authors' reputations. + +Finally, any free program is threatened constantly by software patents. +We wish to avoid the danger that redistributors of a free program will +individually obtain patent licenses, in effect making the program +proprietary. To prevent this, we have made it clear that any patent must +be licensed for everyone's free use or not licensed at all. + +The precise terms and conditions for copying, distribution and +modification follow. + +TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + +0. This License applies to any program or other work which contains a +notice placed by the copyright holder saying it may be distributed under +the terms of this General Public License. The "Program", below, refers +to any such program or work, and a "work based on the Program" means +either the Program or any derivative work under copyright law: that is +to say, a work containing the Program or a portion of it, either +verbatim or with modifications and/or translated into another language. +(Hereinafter, translation is included without limitation in the term +"modification".) Each licensee is addressed as "you". + +Activities other than copying, distribution and modification are not +covered by this License; they are outside its scope. The act of running +the Program is not restricted, and the output from the Program is +covered only if its contents constitute a work based on the Program +(independent of having been made by running the Program). Whether that +is true depends on what the Program does. + +1. You may copy and distribute verbatim copies of the Program's source +code as you receive it, in any medium, provided that you conspicuously +and appropriately publish on each copy an appropriate copyright notice +and disclaimer of warranty; keep intact all the notices that refer to +this License and to the absence of any warranty; and give any other +recipients of the Program a copy of this License along with the Program. + +You may charge a fee for the physical act of transferring a copy, and +you may at your option offer warranty protection in exchange for a fee. + +2. You may modify your copy or copies of the Program or any portion of +it, thus forming a work based on the Program, and copy and distribute +such modifications or work under the terms of Section 1 above, provided +that you also meet all of these conditions: + + a) You must cause the modified files to carry prominent notices + stating that you changed the files and the date of any change. + + b) You must cause any work that you distribute or publish, that in + whole or in part contains or is derived from the Program or any part + thereof, to be licensed as a whole at no charge to all third parties + under the terms of this License. + + c) If the modified program normally reads commands interactively + when run, you must cause it, when started running for such + interactive use in the most ordinary way, to print or display an + announcement including an appropriate copyright notice and a notice + that there is no warranty (or else, saying that you provide a + warranty) and that users may redistribute the program under these + conditions, and telling the user how to view a copy of this License. + (Exception: if the Program itself is interactive but does not + normally print such an announcement, your work based on the Program + is not required to print an announcement.) + +These requirements apply to the modified work as a whole. If +identifiable sections of that work are not derived from the Program, and +can be reasonably considered independent and separate works in +themselves, then this License, and its terms, do not apply to those +sections when you distribute them as separate works. But when you +distribute the same sections as part of a whole which is a work based on +the Program, the distribution of the whole must be on the terms of this +License, whose permissions for other licensees extend to the entire +whole, and thus to each and every part regardless of who wrote it. + +Thus, it is not the intent of this section to claim rights or contest +your rights to work written entirely by you; rather, the intent is to +exercise the right to control the distribution of derivative or +collective works based on the Program. + +In addition, mere aggregation of another work not based on the Program +with the Program (or with a work based on the Program) on a volume of a +storage or distribution medium does not bring the other work under the +scope of this License. + +3. You may copy and distribute the Program (or a work based on it, +under Section 2) in object code or executable form under the terms of +Sections 1 and 2 above provided that you also do one of the following: + + a) Accompany it with the complete corresponding machine-readable + source code, which must be distributed under the terms of Sections 1 + and 2 above on a medium customarily used for software interchange; or, + + b) Accompany it with a written offer, valid for at least three + years, to give any third party, for a charge no more than your cost + of physically performing source distribution, a complete + machine-readable copy of the corresponding source code, to be + distributed under the terms of Sections 1 and 2 above on a medium + customarily used for software interchange; or, + + c) Accompany it with the information you received as to the offer to + distribute corresponding source code. (This alternative is allowed + only for noncommercial distribution and only if you received the + program in object code or executable form with such an offer, in + accord with Subsection b above.) + +The source code for a work means the preferred form of the work for +making modifications to it. For an executable work, complete source code +means all the source code for all modules it contains, plus any +associated interface definition files, plus the scripts used to control +compilation and installation of the executable. However, as a special +exception, the source code distributed need not include anything that is +normally distributed (in either source or binary form) with the major +components (compiler, kernel, and so on) of the operating system on +which the executable runs, unless that component itself accompanies the +executable. + +If distribution of executable or object code is made by offering access +to copy from a designated place, then offering equivalent access to copy +the source code from the same place counts as distribution of the source +code, even though third parties are not compelled to copy the source +along with the object code. + +4. You may not copy, modify, sublicense, or distribute the Program +except as expressly provided under this License. Any attempt otherwise +to copy, modify, sublicense or distribute the Program is void, and will +automatically terminate your rights under this License. However, parties +who have received copies, or rights, from you under this License will +not have their licenses terminated so long as such parties remain in +full compliance. + +5. You are not required to accept this License, since you have not +signed it. However, nothing else grants you permission to modify or +distribute the Program or its derivative works. These actions are +prohibited by law if you do not accept this License. Therefore, by +modifying or distributing the Program (or any work based on the +Program), you indicate your acceptance of this License to do so, and all +its terms and conditions for copying, distributing or modifying the +Program or works based on it. + +6. Each time you redistribute the Program (or any work based on the +Program), the recipient automatically receives a license from the +original licensor to copy, distribute or modify the Program subject to +these terms and conditions. You may not impose any further restrictions +on the recipients' exercise of the rights granted herein. You are not +responsible for enforcing compliance by third parties to this License. + +7. If, as a consequence of a court judgment or allegation of patent +infringement or for any other reason (not limited to patent issues), +conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot distribute +so as to satisfy simultaneously your obligations under this License and +any other pertinent obligations, then as a consequence you may not +distribute the Program at all. For example, if a patent license would +not permit royalty-free redistribution of the Program by all those who +receive copies directly or indirectly through you, then the only way you +could satisfy both it and this License would be to refrain entirely from +distribution of the Program. + +If any portion of this section is held invalid or unenforceable under +any particular circumstance, the balance of the section is intended to +apply and the section as a whole is intended to apply in other +circumstances. + +It is not the purpose of this section to induce you to infringe any +patents or other property right claims or to contest validity of any +such claims; this section has the sole purpose of protecting the +integrity of the free software distribution system, which is implemented +by public license practices. Many people have made generous +contributions to the wide range of software distributed through that +system in reliance on consistent application of that system; it is up to +the author/donor to decide if he or she is willing to distribute +software through any other system and a licensee cannot impose that choice. + +This section is intended to make thoroughly clear what is believed to be +a consequence of the rest of this License. + +8. If the distribution and/or use of the Program is restricted in +certain countries either by patents or by copyrighted interfaces, the +original copyright holder who places the Program under this License may +add an explicit geographical distribution limitation excluding those +countries, so that distribution is permitted only in or among countries +not thus excluded. In such case, this License incorporates the +limitation as if written in the body of this License. + +9. The Free Software Foundation may publish revised and/or new +versions of the General Public License from time to time. Such new +versions will be similar in spirit to the present version, but may +differ in detail to address new problems or concerns. + +Each version is given a distinguishing version number. If the Program +specifies a version number of this License which applies to it and "any +later version", you have the option of following the terms and +conditions either of that version or of any later version published by +the Free Software Foundation. If the Program does not specify a version +number of this License, you may choose any version ever published by the +Free Software Foundation. + +10. If you wish to incorporate parts of the Program into other free +programs whose distribution conditions are different, write to the +author to ask for permission. For software which is copyrighted by the +Free Software Foundation, write to the Free Software Foundation; we +sometimes make exceptions for this. Our decision will be guided by the +two goals of preserving the free status of all derivatives of our free +software and of promoting the sharing and reuse of software generally. + +NO WARRANTY + +11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO +WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. +EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR +OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, +EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE +ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH +YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL +NECESSARY SERVICING, REPAIR OR CORRECTION. + +12. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN +WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY +AND/OR REDISTRIBUTE THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR +DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL +DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM +(INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED +INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF +THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR +OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +END OF TERMS AND CONDITIONS + +How to Apply These Terms to Your New Programs + +If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + +To do so, attach the following notices to the program. It is safest to +attach them to the start of each source file to most effectively convey +the exclusion of warranty; and each file should have at least the +"copyright" line and a pointer to where the full notice is found. + + One line to give the program's name and a brief idea of what it does. + Copyright (C) + + This program is free software; you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation; either version 2 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, but + WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program; if not, write to the Free Software + Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335 USA + +Also add information on how to contact you by electronic and paper mail. + +If the program is interactive, make it output a short notice like this +when it starts in an interactive mode: + + Gnomovision version 69, Copyright (C) year name of author + Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type + `show w'. This is free software, and you are welcome to redistribute + it under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the +appropriate parts of the General Public License. Of course, the commands +you use may be called something other than `show w' and `show c'; they +could even be mouse-clicks or menu items--whatever suits your program. + +You should also get your employer (if you work as a programmer) or your +school, if any, to sign a "copyright disclaimer" for the program, if +necessary. Here is a sample; alter the names: + + Yoyodyne, Inc., hereby disclaims all copyright interest in the + program `Gnomovision' (which makes passes at compilers) written by + James Hacker. + + signature of Ty Coon, 1 April 1989 + Ty Coon, President of Vice + +This General Public License does not permit incorporating your program +into proprietary programs. If your program is a subroutine library, you +may consider it more useful to permit linking proprietary applications +with the library. If this is what you want to do, use the GNU Library +General Public License instead of this License. + +# + +Certain source files distributed by Oracle America, Inc. and/or its +affiliates are subject to the following clarification and special +exception to the GPLv2, based on the GNU Project exception for its +Classpath libraries, known as the GNU Classpath Exception, but only +where Oracle has expressly included in the particular source file's +header the words "Oracle designates this particular file as subject to +the "Classpath" exception as provided by Oracle in the LICENSE file +that accompanied this code." + +You should also note that Oracle includes multiple, independent +programs in this software package. Some of those programs are provided +under licenses deemed incompatible with the GPLv2 by the Free Software +Foundation and others. For example, the package includes programs +licensed under the Apache License, Version 2.0. Such programs are +licensed to you under their original licenses. + +Oracle facilitates your further distribution of this package by adding +the Classpath Exception to the necessary parts of its GPLv2 code, which +permits you to use that code in combination with other independent +modules not licensed under the GPLv2. However, note that this would +not permit you to commingle code under an incompatible license with +Oracle's GPLv2 licensed code by, for example, cutting and pasting such +code into a file also containing Oracle's GPLv2 licensed code and then +distributing the result. Additionally, if you were to remove the +Classpath Exception from any of the files to which it applies and +distribute the result, you would likely be required to license some or +all of the other code in that distribution under the GPLv2 as well, and +since the GPLv2 is incompatible with the license terms of some items +included in the distribution by Oracle, removing the Classpath +Exception could therefore effectively compromise your ability to +further distribute the package. + +Proceed with caution and we recommend that you obtain the advice of a +lawyer skilled in open source matters before removing the Classpath +Exception or making modifications to this package which may +subsequently be redistributed and/or involve the use of third party +software. + +CLASSPATH EXCEPTION +Linking this library statically or dynamically with other modules is +making a combined work based on this library. Thus, the terms and +conditions of the GNU General Public License version 2 cover the whole +combination. + +As a special exception, the copyright holders of this library give you +permission to link this library with independent modules to produce an +executable, regardless of the license terms of these independent +modules, and to copy and distribute the resulting executable under +terms of your choice, provided that you also meet, for each linked +independent module, the terms and conditions of the license of that +module. An independent module is a module which is not derived from or +based on this library. If you modify this library, you may extend this +exception to your version of the library, but you are not obligated to +do so. If you do not wish to do so, delete this exception statement +from your version. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-kryo.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-kryo.txt new file mode 100644 index 0000000..3f6a160 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-kryo.txt @@ -0,0 +1,10 @@ +Copyright (c) 2008, Nathan Sweet +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. + * Neither the name of Esoteric Software nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-leveldbjni.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-leveldbjni.txt new file mode 100644 index 0000000..b4dabb9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-leveldbjni.txt @@ -0,0 +1,27 @@ +Copyright (c) 2011 FuseSource Corp. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + * Neither the name of FuseSource Corp. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-machinist.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-machinist.txt new file mode 100644 index 0000000..68cc3a3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-machinist.txt @@ -0,0 +1,19 @@ +Copyright (c) 2011-2014 Erik Osheim, Tom Switzer + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-matchMedia-polyfill.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-matchMedia-polyfill.txt new file mode 100644 index 0000000..2fd0bc2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-matchMedia-polyfill.txt @@ -0,0 +1 @@ +matchMedia() polyfill - Test a CSS media type/query in JS. Authors & copyright (c) 2012: Scott Jehl, Paul Irish, Nicholas Zakas. Dual MIT/BSD license \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-minlog.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-minlog.txt new file mode 100644 index 0000000..3f6a160 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-minlog.txt @@ -0,0 +1,10 @@ +Copyright (c) 2008, Nathan Sweet +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. + * Neither the name of Esoteric Software nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-modernizr.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-modernizr.txt new file mode 100644 index 0000000..2bf24b9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-modernizr.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-mustache.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-mustache.txt new file mode 100644 index 0000000..038cbb9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-mustache.txt @@ -0,0 +1,11 @@ +The MIT License + +Copyright (c) 2009 Chris Wanstrath (Ruby) +Copyright (c) 2010-2014 Jan Lehnardt (JavaScript) +Copyright (c) 2010-2015 The mustache.js community + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-paranamer.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-paranamer.txt new file mode 100644 index 0000000..fca1847 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-paranamer.txt @@ -0,0 +1,28 @@ +[ ParaNamer used to be 'Pubic Domain', but since it includes a small piece of ASM it is now the same license as that: BSD ] + + Copyright (c) 2006 Paul Hammant & ThoughtWorks Inc + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions + are met: + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + 2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + 3. Neither the name of the copyright holders nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE + LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF + THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pmml-model.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pmml-model.txt new file mode 100644 index 0000000..69411d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pmml-model.txt @@ -0,0 +1,10 @@ +Copyright (c) 2009, University of Tartu +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. +3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-protobuf.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-protobuf.txt new file mode 100644 index 0000000..97a6e3d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-protobuf.txt @@ -0,0 +1,32 @@ +Copyright 2008 Google Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + * Neither the name of Google Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +Code generated by the Protocol Buffer compiler is owned by the owner +of the input file used when generating it. This code is not +standalone and requires a support library to be linked with it. This +support library is itself covered by the above license. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-py4j.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-py4j.txt new file mode 100644 index 0000000..70af3e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-py4j.txt @@ -0,0 +1,27 @@ +Copyright (c) 2009-2011, Barthelemy Dagenais All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +- Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +- Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +- The name of the author may not be used to endorse or promote products +derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. + diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pyrolite.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pyrolite.txt new file mode 100644 index 0000000..9457c7a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-pyrolite.txt @@ -0,0 +1,28 @@ + +Pyro - Python Remote Objects +Software License, copyright, and disclaimer + + Pyro is Copyright (c) by Irmen de Jong (irmen@razorvine.net). + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in + all copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + + +This is the "MIT Software License" which is OSI-certified, and GPL-compatible. +See http://www.opensource.org/licenses/mit-license.php + diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-re2j.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-re2j.txt new file mode 100644 index 0000000..0dc3cd7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-re2j.txt @@ -0,0 +1,32 @@ +This is a work derived from Russ Cox's RE2 in Go, whose license +http://golang.org/LICENSE is as follows: + +Copyright (c) 2009 The Go Authors. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + + * Neither the name of Google Inc. nor the names of its contributors + may be used to endorse or promote products derived from this + software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-reflectasm.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-reflectasm.txt new file mode 100644 index 0000000..3f6a160 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-reflectasm.txt @@ -0,0 +1,10 @@ +Copyright (c) 2008, Nathan Sweet +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. + * Neither the name of Esoteric Software nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-respond.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-respond.txt new file mode 100644 index 0000000..dea4ff9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-respond.txt @@ -0,0 +1,22 @@ +Copyright (c) 2012 Scott Jehl + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated documentation +files (the "Software"), to deal in the Software without +restriction, including without limitation the rights to use, +copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES +OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT +HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, +WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sbt-launch-lib.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sbt-launch-lib.txt new file mode 100644 index 0000000..3b9156b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sbt-launch-lib.txt @@ -0,0 +1,26 @@ +// Generated from http://www.opensource.org/licenses/bsd-license.php +Copyright (c) 2011, Paul Phillips. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + * Neither the name of the author nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, +EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-scopt.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-scopt.txt new file mode 100644 index 0000000..e92e9b5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-scopt.txt @@ -0,0 +1,9 @@ +This project is licensed under the MIT license. + +Copyright (c) scopt contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-slf4j.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-slf4j.txt new file mode 100644 index 0000000..6548cd3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-slf4j.txt @@ -0,0 +1,21 @@ +Copyright (c) 2004-2013 QOS.ch + All rights reserved. + + Permission is hereby granted, free of charge, to any person obtaining + a copy of this software and associated documentation files (the + "Software"), to deal in the Software without restriction, including + without limitation the rights to use, copy, modify, merge, publish, + distribute, sublicense, and/or sell copies of the Software, and to + permit persons to whom the Software is furnished to do so, subject to + the following conditions: + + The above copyright notice and this permission notice shall be + included in all copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND + NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE + LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION + OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sorttable.js.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sorttable.js.txt new file mode 100644 index 0000000..b31a5b2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-sorttable.js.txt @@ -0,0 +1,16 @@ +Copyright (c) 1997-2007 Stuart Langridge + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-spire.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-spire.txt new file mode 100644 index 0000000..40af774 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-spire.txt @@ -0,0 +1,19 @@ +Copyright (c) 2011-2012 Erik Osheim, Tom Switzer + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-vis-timeline.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-vis-timeline.txt new file mode 100644 index 0000000..e9f98c7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-vis-timeline.txt @@ -0,0 +1,23 @@ +vis-timeline +https://visjs.github.io/vis-timeline/ + +Create a fully customizable, interactive timeline with items and ranges. + +@version 7.7.2 +@date 2023-03-22T11:14:31.874Z + +@copyright (c) 2011-2017 Almende B.V, http://almende.com +@copyright (c) 2017-2019 visjs contributors, https://github.com/visjs + +@license +vis.js is dual licensed under both + + 1. The Apache 2.0 License + http://www.apache.org/licenses/LICENSE-2.0 + + and + + 2. The MIT License + http://opensource.org/licenses/MIT + +vis.js may be distributed under either license. diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-xmlenc.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-xmlenc.txt new file mode 100644 index 0000000..3a70c9b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-xmlenc.txt @@ -0,0 +1,27 @@ +Copyright 2003-2005, Ernst de Haan +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd-jni.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd-jni.txt new file mode 100644 index 0000000..32c6bbd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd-jni.txt @@ -0,0 +1,26 @@ +Zstd-jni: JNI bindings to Zstd Library + +Copyright (c) 2015-2016, Luben Karavelov/ All rights reserved. + +BSD License + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd.txt b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd.txt new file mode 100644 index 0000000..a793a80 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/licenses/LICENSE-zstd.txt @@ -0,0 +1,30 @@ +BSD License + +For Zstandard software + +Copyright (c) 2016-present, Facebook, Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name Facebook nor the names of its contributors may be used to + endorse or promote products derived from this software without specific + prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/spark-3.5.0-bin-hadoop3/python/.coveragerc b/spark-3.5.0-bin-hadoop3/python/.coveragerc new file mode 100644 index 0000000..b3339cd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/.coveragerc @@ -0,0 +1,21 @@ +# +# 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. +# + +[run] +branch = true +parallel = true +data_file = ${COVERAGE_DIR}/coverage_data/coverage diff --git a/spark-3.5.0-bin-hadoop3/python/.gitignore b/spark-3.5.0-bin-hadoop3/python/.gitignore new file mode 100644 index 0000000..52128cf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/.gitignore @@ -0,0 +1,5 @@ +*.pyc +docs/_build/ +pyspark.egg-info +build/ +dist/ diff --git a/spark-3.5.0-bin-hadoop3/python/MANIFEST.in b/spark-3.5.0-bin-hadoop3/python/MANIFEST.in new file mode 100644 index 0000000..862d62b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/MANIFEST.in @@ -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. + +global-exclude *.py[cod] __pycache__ .DS_Store +recursive-include deps/jars *.jar +graft deps/bin +recursive-include deps/sbin spark-config.sh spark-daemon.sh start-history-server.sh stop-history-server.sh +recursive-include deps/data *.data *.txt +recursive-include deps/licenses *.txt +recursive-include deps/examples *.py +recursive-include lib *.zip +recursive-include pyspark *.pyi py.typed +include README.md diff --git a/spark-3.5.0-bin-hadoop3/python/README.md b/spark-3.5.0-bin-hadoop3/python/README.md new file mode 100644 index 0000000..088054e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/README.md @@ -0,0 +1,33 @@ +# Apache Spark + +Spark is a unified analytics engine for large-scale data processing. It provides +high-level APIs in Scala, Java, Python, and R, and an optimized engine that +supports general computation graphs for data analysis. It also supports a +rich set of higher-level tools including Spark SQL for SQL and DataFrames, +pandas API on Spark for pandas workloads, MLlib for machine learning, GraphX for graph processing, +and Structured Streaming for stream processing. + + + +## Online Documentation + +You can find the latest Spark documentation, including a programming +guide, on the [project web page](https://spark.apache.org/documentation.html) + + +## Python Packaging + +This README file only contains basic information related to pip installed PySpark. +This packaging is currently experimental and may change in future versions (although we will do our best to keep compatibility). +Using PySpark requires the Spark JARs, and if you are building this from source please see the builder instructions at +["Building Spark"](https://spark.apache.org/docs/latest/building-spark.html). + +The Python packaging for Spark is not intended to replace all of the other use cases. This Python packaged version of Spark is suitable for interacting with an existing cluster (be it Spark standalone, YARN, or Mesos) - but does not contain the tools required to set up your own standalone Spark cluster. You can download the full version of Spark from the [Apache Spark downloads page](https://spark.apache.org/downloads.html). + + +**NOTE:** If you are using this with a Spark standalone cluster you must ensure that the version (including minor version) matches or you may experience odd errors. + +## Python Requirements + +At its core PySpark depends on Py4J, but some additional sub-packages have their own extra requirements for some features (including numpy, pandas, and pyarrow). +See also [Dependencies](https://spark.apache.org/docs/latest/api/python/getting_started/install.html#dependencies) for production, and [dev/requirements.txt](https://github.com/apache/spark/blob/master/dev/requirements.txt) for development. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/Makefile b/spark-3.5.0-bin-hadoop3/python/docs/Makefile new file mode 100644 index 0000000..c959620 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/Makefile @@ -0,0 +1,35 @@ +# 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. + +# Minimal makefile for Sphinx documentation + +# You can set these variables from the command line. +SPHINXOPTS ?= "-W" +SPHINXBUILD ?= sphinx-build +SOURCEDIR ?= source +BUILDDIR ?= build + +export PYTHONPATH=$(realpath ..):$(realpath ../lib/py4j-0.10.9.7-src.zip) + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/spark-3.5.0-bin-hadoop3/python/docs/make.bat b/spark-3.5.0-bin-hadoop3/python/docs/make.bat new file mode 100644 index 0000000..ad72e8f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/make.bat @@ -0,0 +1,23 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem This is the entry point for running Sphinx documentation. To avoid polluting the +rem environment, it just launches a new cmd to do the real work. + +cmd /V /E /C %~dp0make2.bat %* diff --git a/spark-3.5.0-bin-hadoop3/python/docs/make2.bat b/spark-3.5.0-bin-hadoop3/python/docs/make2.bat new file mode 100644 index 0000000..41e33cd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/make2.bat @@ -0,0 +1,51 @@ +@echo off + +rem +rem Licensed to the Apache Software Foundation (ASF) under one or more +rem contributor license agreements. See the NOTICE file distributed with +rem this work for additional information regarding copyright ownership. +rem The ASF licenses this file to You under the Apache License, Version 2.0 +rem (the "License"); you may not use this file except in compliance with +rem the License. You may obtain a copy of the License at +rem +rem http://www.apache.org/licenses/LICENSE-2.0 +rem +rem Unless required by applicable law or agreed to in writing, software +rem distributed under the License is distributed on an "AS IS" BASIS, +rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +rem See the License for the specific language governing permissions and +rem limitations under the License. +rem + +rem Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +set PYTHONPATH=..;..\lib\py4j-0.10.9.7-src.zip + +if "%1" == "" goto help + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 +) + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/_static/copybutton.js b/spark-3.5.0-bin-hadoop3/python/docs/source/_static/copybutton.js new file mode 100644 index 0000000..896faad --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/_static/copybutton.js @@ -0,0 +1,67 @@ +// Copyright 2014 PSF. Licensed under the PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +// File originates from the cpython source found in Doc/tools/sphinxext/static/copybutton.js + +$(document).ready(function() { + /* Add a [>>>] button on the top-right corner of code samples to hide + * the >>> and ... prompts and the output and thus make the code + * copyable. */ + var div = $('.highlight-python .highlight,' + + '.highlight-default .highlight,' + + '.highlight-python3 .highlight') + var pre = div.find('pre'); + + // get the styles from the current theme + pre.parent().parent().css('position', 'relative'); + var hide_text = 'Hide the prompts and output'; + var show_text = 'Show the prompts and output'; + var border_width = pre.css('border-top-width'); + var border_style = pre.css('border-top-style'); + var border_color = pre.css('border-top-color'); + var button_styles = { + 'cursor':'pointer', 'position': 'absolute', 'top': '0', 'right': '0', + 'border-color': border_color, 'border-style': border_style, + 'border-width': border_width, 'color': border_color, 'text-size': '75%', + 'font-family': 'monospace', 'padding-left': '0.2em', 'padding-right': '0.2em', + 'border-radius': '0 3px 0 0', + 'user-select': 'none' + } + + // create and add the button to all the code blocks that contain >>> + div.each(function(index) { + var jthis = $(this); + if (jthis.find('.gp').length > 0) { + var button = $('>>>'); + button.css(button_styles) + button.attr('title', hide_text); + button.data('hidden', 'false'); + jthis.prepend(button); + } + // tracebacks (.gt) contain bare text elements that need to be + // wrapped in a span to work with .nextUntil() (see later) + jthis.find('pre:has(.gt)').contents().filter(function() { + return ((this.nodeType == 3) && (this.data.trim().length > 0)); + }).wrap(''); + }); + + // define the behavior of the button when it's clicked + $('.copybutton').click(function(e){ + e.preventDefault(); + var button = $(this); + if (button.data('hidden') === 'false') { + // hide the code output + button.parent().find('.go, .gp, .gt').hide(); + button.next('pre').find('.gt').nextUntil('.gp, .go').css('visibility', 'hidden'); + button.css('text-decoration', 'line-through'); + button.attr('title', show_text); + button.data('hidden', 'true'); + } else { + // show the code output + button.parent().find('.go, .gp, .gt').show(); + button.next('pre').find('.gt').nextUntil('.gp, .go').css('visibility', 'visible'); + button.css('text-decoration', 'none'); + button.attr('title', hide_text); + button.data('hidden', 'false'); + } + }); +}); + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/_static/css/pyspark.css b/spark-3.5.0-bin-hadoop3/python/docs/source/_static/css/pyspark.css new file mode 100644 index 0000000..89b7c65 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/_static/css/pyspark.css @@ -0,0 +1,97 @@ +/* + * 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. + * + */ + +/* PySpark style CSS overwrite */ + +/* Lato font (from its parent) does not make any difference when it's bold. Defaults to 'Source Sans Pro' */ +body { + font-family:"Source Sans Pro",sans-serif!important; +} + +h1,h2 { + color:#1B5162!important; +} + +h3 { + color: #555555 +} + +/* Top menu */ +#navbar-main { + background: #1B5162!important; + box-shadow: 0px 2px 4px rgba(0, 0, 0, 0.11); +} + +#navbar-main-elements li.nav-item a { + color: rgba(255, 255, 255, 0.8); +} + +#navbar-main-elements li.active a { + font-weight: 600; + color: #FFFFFF!important; +} + +.col-9 { + flex: 0 0 80%; + max-width: 80%; +} + +/* Left panel size */ +@media (min-width: 768px) { + .col-md-3 { + flex: 0 0 20%; + max-width: 20%; + } +} + +/* Top menu right button */ +.navbar-toggler { + color:rgba(255,255,255,.5)!important; + border-color:rgba(255,255,255,.5)!important; +} + +.navbar-toggler-icon { + background-image:url("data:image/svg+xml;charset=utf-8,%3Csvg xmlns='http://www.w3.org/2000/svg' width='30' height='30'%3E%3Cpath stroke='rgba(255,255,255,.5)' stroke-linecap='round' stroke-miterlimit='10' stroke-width='2' d='M4 7h22M4 15h22M4 23h22'/%3E%3C/svg%3E")!important; +} + +/* Left bar list colors */ +.bd-sidebar .nav>.active>a { + color: #1B5162!important; +} + +.bd-sidebar .nav>li>a:hover { + color: #1B5162!important; +} + +.bd-sidebar .nav>.active:hover>a,.bd-sidebar .nav>.active>a { + color: #1B5162!important; +} + +u.bd-sidebar .nav>li>ul>.active:hover>a,.bd-sidebar .nav>li>ul>.active>a { + color: #1B5162!important; +} + +/* Right bar list colors */ +.toc-entry>.nav-link.active { + color: #1B5162!important; + border-left: 2px solid #1B5162!important; +} + +.spec_table tr, td, th { + border-top: none!important; +} diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class.rst new file mode 100644 index 0000000..b5f6267 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class.rst @@ -0,0 +1,38 @@ +.. 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. + +.. Workaround to avoid documenting __init__. + +{% extends "!autosummary/class.rst" %} + +{% if '__init__' in methods %} +{% set caught_result = methods.remove('__init__') %} +{% endif %} + +{% block methods %} +{% if methods %} + + .. rubric:: Methods + + .. autosummary:: + {% for item in methods %} + ~{{ name }}.{{ item }} + {%- endfor %} + +{% endif %} +{% endblock %} + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class_with_docs.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class_with_docs.rst new file mode 100644 index 0000000..7c37b83 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/_templates/autosummary/class_with_docs.rst @@ -0,0 +1,79 @@ +.. 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. + + +{{ objname }} +{{ underline }} + +.. currentmodule:: {{ module }} + +.. autoclass:: {{ objname }} + + {% if '__init__' in methods %} + {% set caught_result = methods.remove('__init__') %} + {% endif %} + + {% block methods_summary %} + {% if methods %} + + .. rubric:: Methods + + .. autosummary:: + {% for item in methods %} + ~{{ name }}.{{ item }} + {%- endfor %} + + {% endif %} + {% endblock %} + + {% block attributes_summary %} + {% if attributes %} + + .. rubric:: Attributes + + .. autosummary:: + {% for item in attributes %} + ~{{ name }}.{{ item }} + {%- endfor %} + + {% endif %} + {% endblock %} + + {% block methods_documentation %} + {% if methods %} + + .. rubric:: Methods Documentation + + {% for item in methods %} + .. automethod:: {{ item }} + {%- endfor %} + + {% endif %} + {% endblock %} + + {% block attributes_documentation %} + {% if attributes %} + + .. rubric:: Attributes Documentation + + {% for item in attributes %} + .. autoattribute:: {{ item }} + {%- endfor %} + + {% endif %} + {% endblock %} + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/conf.py b/spark-3.5.0-bin-hadoop3/python/docs/source/conf.py new file mode 100644 index 0000000..38c3310 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/conf.py @@ -0,0 +1,417 @@ +# -*- coding: utf-8 -*- +# +# pyspark documentation build configuration file, created by +# sphinx-quickstart on Thu Aug 28 15:17:47 2014. +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +import sys +import os +import shutil +import errno + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +sys.path.insert(0, os.path.abspath('.')) + +# generate user_guide/pandas_on_spark/supported_pandas_api.rst +from pyspark.pandas.supported_api_gen import generate_supported_api + +output_rst_file_path = ( + "%s/user_guide/pandas_on_spark/supported_pandas_api.rst" + % os.path.dirname(os.path.abspath(__file__)) +) +generate_supported_api(output_rst_file_path) + +# Remove previously generated rst files. Ignore errors just in case it stops +# generating whole docs. +shutil.rmtree( + "%s/reference/api" % os.path.dirname(os.path.abspath(__file__)), ignore_errors=True) +shutil.rmtree( + "%s/reference/pyspark.pandas/api" % os.path.dirname(os.path.abspath(__file__)), + ignore_errors=True) +try: + os.mkdir("%s/reference/api" % os.path.dirname(os.path.abspath(__file__))) +except OSError as e: + if e.errno != errno.EEXIST: + raise +try: + os.mkdir("%s/reference/pyspark.pandas/api" % os.path.dirname( + os.path.abspath(__file__))) +except OSError as e: + if e.errno != errno.EEXIST: + raise + +# -- General configuration ------------------------------------------------ + +# If your documentation needs a minimal Sphinx version, state it here. +needs_sphinx = '1.2' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.viewcode', + 'sphinx.ext.mathjax', + 'sphinx.ext.autosummary', + 'nbsphinx', # Converts Jupyter Notebook to reStructuredText files for Sphinx. + # For ipython directive in reStructuredText files. It is generated by the notebook. + 'IPython.sphinxext.ipython_console_highlighting', + 'numpydoc', # handle NumPy documentation formatted docstrings. + 'sphinx_plotly_directive', # For visualize plot result +] + +# plotly plot directive +plotly_include_source = True +plotly_html_show_formats = False +plotly_html_show_source_link = False +plotly_pre_code = """import numpy as np +import pandas as pd +import pyspark.pandas as ps""" + +numpydoc_show_class_members = False + +# Links used globally in the RST files. +# These are defined here to allow link substitutions dynamically. +rst_epilog = """ +.. |binder| replace:: Live Notebook +.. _binder: https://mybinder.org/v2/gh/apache/spark/{0}?filepath=python%2Fdocs%2Fsource%2Fgetting_started%2Fquickstart_df.ipynb +.. |binder_df| replace:: Live Notebook: DataFrame +.. _binder_df: https://mybinder.org/v2/gh/apache/spark/{0}?filepath=python%2Fdocs%2Fsource%2Fgetting_started%2Fquickstart_df.ipynb +.. |binder_ps| replace:: Live Notebook: pandas API on Spark +.. _binder_ps: https://mybinder.org/v2/gh/apache/spark/{0}?filepath=python%2Fdocs%2Fsource%2Fgetting_started%2Fquickstart_ps.ipynb +.. |binder_connect| replace:: Live Notebook: Spark Connect +.. _binder_connect: https://mybinder.org/v2/gh/apache/spark/{0}?filepath=python%2Fdocs%2Fsource%2Fgetting_started%2Fquickstart_connect.ipynb +.. |examples| replace:: Examples +.. _examples: https://github.com/apache/spark/tree/{0}/examples/src/main/python +.. |downloading| replace:: Downloading +.. _downloading: https://spark.apache.org/docs/{1}/building-spark.html +.. |building_spark| replace:: Building Spark +.. _building_spark: https://spark.apache.org/docs/{1}/#downloading +""".format( + os.environ.get("GIT_HASH", "master"), + os.environ.get("RELEASE_VERSION", "latest"), +) + +# Add any paths that contain templates here, relative to this directory. +templates_path = ['_templates'] + +# The suffix of source filenames. +source_suffix = '.rst' + +# The encoding of source files. +#source_encoding = 'utf-8-sig' + +# The master toctree document. +master_doc = 'index' + +# General information about the project. +project = 'PySpark' +copyright = '' + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. +# +# The short X.Y version. +version = 'master' +# The full version, including alpha/beta/rc tags. +release = os.environ.get('RELEASE_VERSION', version) + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +#language = None + +# There are two options for replacing |today|: either, you set today to some +# non-false value, then it is used: +#today = '' +# Else, today_fmt is used as the format for a strftime call. +#today_fmt = '%B %d, %Y' + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +exclude_patterns = ['_build', '.DS_Store', '**.ipynb_checkpoints'] + +# The reST default role (used for this markup: `text`) to use for all +# documents. +#default_role = None + +# If true, '()' will be appended to :func: etc. cross-reference text. +#add_function_parentheses = True + +# If true, the current module name will be prepended to all description +# unit titles (such as .. function::). +#add_module_names = True + +# If true, sectionauthor and moduleauthor directives will be shown in the +# output. They are ignored by default. +#show_authors = False + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = 'sphinx' + +# A list of ignored prefixes for module index sorting. +#modindex_common_prefix = [] + +# If true, keep warnings as "system message" paragraphs in the built documents. +#keep_warnings = False + +# -- Options for autodoc -------------------------------------------------- + +# Look at the first line of the docstring for function and method signatures. +autodoc_docstring_signature = True +autosummary_generate = True + +# -- Options for HTML output ---------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +html_theme = 'pydata_sphinx_theme' + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +#html_theme_options = {} + +# Add any paths that contain custom themes here, relative to this directory. +#html_theme_path = [] + +# The name for this set of Sphinx documents. If None, it defaults to +# " v documentation". +#html_title = None + +# A shorter title for the navigation bar. Default is the same as html_title. +#html_short_title = None + +# The name of an image file (relative to this directory) to place at the top +# of the sidebar. +html_logo = "../../../docs/img/spark-logo-reverse.png" + +# The name of an image file (within the static path) to use as a favicon of the +# docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 +# pixels large. +#html_favicon = None + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ['_static'] + +html_css_files = [ + 'css/pyspark.css', +] + +# Add any extra paths that contain custom files (such as robots.txt or +# .htaccess) here, relative to this directory. These files are copied +# directly to the root of the documentation. +#html_extra_path = [] + +# If not '', a 'Last updated on:' timestamp is inserted at every page bottom, +# using the given strftime format. +#html_last_updated_fmt = '%b %d, %Y' + +# If true, SmartyPants will be used to convert quotes and dashes to +# typographically correct entities. +#html_use_smartypants = True + +# Custom sidebar templates, maps document names to template names. +#html_sidebars = {} + +# Additional templates that should be rendered to pages; maps page names to +# template names. +#html_additional_pages = {} + +# If false, no module index is generated. +html_domain_indices = False + +# If false, no index is generated. +html_use_index = False + +# If true, the index is split into individual pages for each letter. +#html_split_index = False + +# If true, links to the reST sources are added to the pages. +#html_show_sourcelink = True + +# If true, "Created using Sphinx" is shown in the HTML footer. Default is True. +#html_show_sphinx = True + +# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. +#html_show_copyright = True + +# If true, an OpenSearch description file will be output, and all pages will +# contain a tag referring to it. The value of this option must be the +# base URL from which the finished HTML is served. +#html_use_opensearch = '' + +# This is the file name suffix for HTML files (e.g. ".xhtml"). +#html_file_suffix = None + +# Output file base name for HTML help builder. +htmlhelp_basename = 'pysparkdoc' + +# The base URL which points to the root of the HTML documentation. +html_baseurl = 'https://spark.apache.org/docs/latest/api/python' + +# -- Options for LaTeX output --------------------------------------------- + +latex_elements = { +# The paper size ('letterpaper' or 'a4paper'). +#'papersize': 'letterpaper', + +# The font size ('10pt', '11pt' or '12pt'). +#'pointsize': '10pt', + +# Additional stuff for the LaTeX preamble. +#'preamble': '', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + ('index', 'pyspark.tex', 'pyspark Documentation', + 'Author', 'manual'), +] + +# The name of an image file (relative to this directory) to place at the top of +# the title page. +#latex_logo = None + +# For "manual" documents, if this is true, then toplevel headings are parts, +# not chapters. +#latex_use_parts = False + +# If true, show page references after internal links. +#latex_show_pagerefs = False + +# If true, show URL addresses after external links. +#latex_show_urls = False + +# Documents to append as an appendix to all manuals. +#latex_appendices = [] + +# If false, no module index is generated. +#latex_domain_indices = True + + +# -- Options for manual page output --------------------------------------- + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + ('index', 'pyspark', 'pyspark Documentation', + ['Author'], 1) +] + +# If true, show URL addresses after external links. +#man_show_urls = False + + +# -- Options for Texinfo output ------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ('index', 'pyspark', 'pyspark Documentation', + 'Author', 'pyspark', 'One line description of project.', + 'Miscellaneous'), +] + +# Documents to append as an appendix to all manuals. +#texinfo_appendices = [] + +# If false, no module index is generated. +#texinfo_domain_indices = True + +# How to display URL addresses: 'footnote', 'no', or 'inline'. +#texinfo_show_urls = 'footnote' + +# If true, do not generate a @detailmenu in the "Top" node's menu. +#texinfo_no_detailmenu = False + + +# -- Options for Epub output ---------------------------------------------- + +# Bibliographic Dublin Core info. +epub_title = 'pyspark' +epub_author = 'Author' +epub_publisher = 'Author' +epub_copyright = '2014, Author' + +# The basename for the epub file. It defaults to the project name. +#epub_basename = 'pyspark' + +# The HTML theme for the epub output. Since the default themes are not optimized +# for small screen space, using the same theme for HTML and epub output is +# usually not wise. This defaults to 'epub', a theme designed to save visual +# space. +#epub_theme = 'epub' + +# The language of the text. It defaults to the language option +# or en if the language is not set. +#epub_language = '' + +# The scheme of the identifier. Typical schemes are ISBN or URL. +#epub_scheme = '' + +# The unique identifier of the text. This can be an ISBN number +# or the project homepage. +#epub_identifier = '' + +# A unique identification for the text. +#epub_uid = '' + +# A tuple containing the cover image and cover page html template filenames. +#epub_cover = () + +# A sequence of (type, uri, title) tuples for the guide element of content.opf. +#epub_guide = () + +# HTML files that should be inserted before the pages created by sphinx. +# The format is a list of tuples containing the path and title. +#epub_pre_files = [] + +# HTML files that should be inserted after the pages created by sphinx. +# The format is a list of tuples containing the path and title. +#epub_post_files = [] + +# A list of files that should not be packed into the epub file. +epub_exclude_files = ['search.html'] + +# The depth of the table of contents in toc.ncx. +#epub_tocdepth = 3 + +# Allow duplicate toc entries. +#epub_tocdup = True + +# Choose between 'default' and 'includehidden'. +#epub_tocscope = 'default' + +# Fix unsupported image types using the PIL. +#epub_fix_images = False + +# Scale large images. +#epub_max_image_width = 0 + +# How to display URL addresses: 'footnote', 'no', or 'inline'. +#epub_show_urls = 'inline' + +# If false, no index is generated. +#epub_use_index = True +def setup(app): + # The app.add_javascript() is deprecated. + getattr(app, "add_js_file", getattr(app, "add_javascript", None))('copybutton.js') + +# Skip sample endpoint link (not expected to resolve) +linkcheck_ignore = [r'https://kinesis.us-east-1.amazonaws.com'] diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/contributing.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/contributing.rst new file mode 100644 index 0000000..92fde4b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/contributing.rst @@ -0,0 +1,335 @@ +.. 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. + +======================= +Contributing to PySpark +======================= + +There are many types of contribution, for example, helping other users, testing releases, reviewing changes, +documentation contribution, bug reporting, JIRA maintenance, code changes, etc. +These are documented at `the general guidelines `_. +This page focuses on PySpark and includes additional details specifically for PySpark. + + +Contributing by Testing Releases +-------------------------------- + +Before the official release, PySpark release candidates are shared in the `dev@spark.apache.org `_ mailing list to vote on. +This release candidates can be easily installed via pip. For example, in case of Spark 3.0.0 RC1, you can install as below: + +.. code-block:: bash + + pip install https://dist.apache.org/repos/dist/dev/spark/v3.0.0-rc1-bin/pyspark-3.0.0.tar.gz + +The link for release files such as ``https://dist.apache.org/repos/dist/dev/spark/v3.0.0-rc1-bin`` can be found in the vote thread. + +Testing and verifying users' existing workloads against release candidates is one of the vital contributions to PySpark. +It prevents breaking users' existing workloads before the official release. +When there is an issue such as a regression, correctness problem or performance degradation worth enough to drop the release candidate, +usually the release candidate is dropped and the community focuses on fixing it to include in the next release candidate. + + +Contributing Documentation Changes +---------------------------------- + +The release documentation is located under Spark's `docs `_ directory. +`README.md `_ describes the required dependencies and steps +to generate the documentations. Usually, PySpark documentation is tested with the command below +under the `docs `_ directory: + +.. code-block:: bash + + SKIP_SCALADOC=1 SKIP_RDOC=1 SKIP_SQLDOC=1 bundle exec jekyll serve --watch + +PySpark uses Sphinx to generate its release PySpark documentation. Therefore, if you want to build only PySpark documentation alone, +you can build under `python/docs `_ directory by: + +.. code-block:: bash + + make html + +It generates the corresponding HTMLs under ``python/docs/build/html``. + +Lastly, please make sure that the new APIs are documented by manually adding methods and/or classes at the corresponding RST files +under ``python/docs/source/reference``. Otherwise, they would not be documented in PySpark documentation. + + +Preparing to Contribute Code Changes +------------------------------------ + +Before starting to work on codes in PySpark, it is recommended to read `the general guidelines `_. +Additionally, there are a couple of additional notes to keep in mind when contributing to codes in PySpark: + +* Be Pythonic + See `The Zen of Python `_. + +* Match APIs with Scala and Java sides + Apache Spark is an unified engine that provides a consistent API layer. In general, the APIs are consistently supported across other languages. + +* PySpark-specific APIs can be accepted + As long as they are Pythonic and do not conflict with other existent APIs, it is fine to raise a API request, for example, decorator usage of UDFs. + +* Adjust the corresponding type hints if you extend or modify public API + See `Contributing and Maintaining Type Hints`_ for details. + +If you are fixing pandas API on Spark (``pyspark.pandas``) package, please consider the design principles below: + +* Return pandas-on-Spark data structure for big data, and pandas data structure for small data + Often developers face the question whether a particular function should return a pandas-on-Spark DataFrame/Series, or a pandas DataFrame/Series. The principle is: if the returned object can be large, use a pandas-on-Spark DataFrame/Series. If the data is bound to be small, use a pandas DataFrame/Series. For example, ``DataFrame.dtypes`` return a pandas Series, because the number of columns in a DataFrame is bounded and small, whereas ``DataFrame.head()`` or ``Series.unique()`` returns a pandas-on-Spark DataFrame/Series, because the resulting object can be large. + +* Provide discoverable APIs for common data science tasks + At the risk of overgeneralization, there are two API design approaches: the first focuses on providing APIs for common tasks; the second starts with abstractions, and enables users to accomplish their tasks by composing primitives. While the world is not black and white, pandas takes more of the former approach, while Spark has taken more of the latter. + + One example is value count (count by some key column), one of the most common operations in data science. pandas ``DataFrame.value_counts()`` returns the result in sorted order, which in 90% of the cases is what users prefer when exploring data, whereas Spark's does not sort, which is more desirable when building data pipelines, as users can accomplish the pandas behavior by adding an explicit ``orderBy``. + + Similar to pandas, pandas API on Spark should also lean more towards the former, providing discoverable APIs for common data science tasks. In most cases, this principle is well taken care of by simply implementing pandas' APIs. However, there will be circumstances in which pandas' APIs don't address a specific need, e.g. plotting for big data. + +* Guardrails to prevent users from shooting themselves in the foot + Certain operations in pandas are prohibitively expensive as data scales, and we don't want to give users the illusion that they can rely on such operations in pandas API on Spark. That is to say, methods implemented in pandas API on Spark should be safe to perform by default on large datasets. As a result, the following capabilities are not implemented in pandas API on Spark: + + * Capabilities that are fundamentally not parallelizable: e.g. imperatively looping over each element + * Capabilities that require materializing the entire working set in a single node's memory. This is why we do not implement `pandas.DataFrame.to_xarray `_. Another example is the ``_repr_html_`` call caps the total number of records shown to a maximum of 1000, to prevent users from blowing up their driver node simply by typing the name of the DataFrame in a notebook. + + A few exceptions, however, exist. One common pattern with "big data science" is that while the initial dataset is large, the working set becomes smaller as the analysis goes deeper. For example, data scientists often perform aggregation on datasets and want to then convert the aggregated dataset to some local data structure. To help data scientists, we offer the following: + + * ``DataFrame.to_pandas``: returns a pandas DataFrame (pandas-on-Spark only) + * ``DataFrame.to_numpy``: returns a numpy array, works with both pandas and pandas API on Spark + + Note that it is clear from the names that these functions return some local data structure that would require materializing data in a single node's memory. For these functions, we also explicitly document them with a warning note that the resulting data structure must be small. + + +Environment Setup +----------------- + +Prerequisite +~~~~~~~~~~~~ + +PySpark development requires to build Spark that needs a proper JDK installed, etc. See `Building Spark `_ for more details. + +Note that if you intend to contribute to Spark Connect in Python, ``buf`` version ``1.24.0`` is required, see `Buf Installation `_ for more details. + +Conda +~~~~~ + +If you are using Conda, the development environment can be set as follows. + +.. code-block:: bash + + # Python 3.8+ is required + conda create --name pyspark-dev-env python=3.9 + conda activate pyspark-dev-env + pip install --upgrade -r dev/requirements.txt + +Once it is set up, make sure you switch to `pyspark-dev-env` before starting the development: + +.. code-block:: bash + + conda activate pyspark-dev-env + +Now, you can start developing and `running the tests `_. + +pip +~~~ + +With Python 3.8+, pip can be used as below to install and set up the development environment. + +.. code-block:: bash + + pip install --upgrade -r dev/requirements.txt + +Now, you can start developing and `running the tests `_. + + +Contributing and Maintaining Type Hints +---------------------------------------- + +PySpark type hints are inlined, to take advantage of static type checking. + +As a rule of thumb, only public API is annotated. + +Annotations should, when possible: + +* Reflect expectations of the underlying JVM API, to help avoid type related failures outside Python interpreter. +* In case of conflict between too broad (``Any``) and too narrow argument annotations, prefer the latter as one, as long as it is covering most of the typical use cases. +* Indicate nonsensical combinations of arguments using ``@overload`` annotations. For example, to indicate that ``*Col`` and ``*Cols`` arguments are mutually exclusive: + + .. code-block:: python + + @overload + def __init__( + self, + *, + threshold: float = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ... + ) -> None: ... + @overload + def __init__( + self, + *, + thresholds: Optional[List[float]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ... + ) -> None: ... + +* Be compatible with the current stable MyPy release. + + +Complex supporting type definitions, should be placed in dedicated ``_typing.pyi`` stubs. See for example `pyspark.sql._typing.pyi `_. + +Annotations can be validated using ``dev/lint-python`` script or by invoking mypy directly: + +.. code-block:: bash + + mypy --config python/mypy.ini python/pyspark + + +Code and Docstring Guide +------------------------ + +Code Conventions +~~~~~~~~~~~~~~~~ + +Please follow the style of the existing codebase as is, which is virtually PEP 8 with one exception: lines can be up +to 100 characters in length, not 79. + +Note that: + +* the method and variable names in PySpark are the similar case is ``threading`` library in Python itself where the APIs were inspired by Java. PySpark also follows `camelCase` for exposed APIs that match with Scala and Java. + +* In contrast, ``functions.py`` uses `snake_case` in order to make APIs SQL (and Python) friendly. + +* In addition, pandas-on-Spark (``pyspark.pandas``) also uses `snake_case` because this package is free from API consistency with other languages. + +PySpark leverages linters such as `pycodestyle `_ and `flake8 `_, which ``dev/lint-python`` runs. Therefore, make sure to run that script to double check. + + +Docstring Conventions +~~~~~~~~~~~~~~~~~~~~~ + +PySpark follows `NumPy documentation style `_. + + +Doctest Conventions +~~~~~~~~~~~~~~~~~~~ + +In general, doctests should be grouped logically by separating a newline. + +For instance, the first block is for the statements for preparation, the second block is for using the function with a specific argument, +and third block is for another argument. As a example, please refer `DataFrame.rsub `_ in pandas. + +These blocks should be consistently separated in PySpark doctests, and more doctests should be added if the coverage of the doctests or the number of examples to show is not enough. + + +Contributing Error and Exception +-------------------------------- + +.. currentmodule:: pyspark.errors + +To throw a standardized user-facing error or exception, developers should specify the error class and message parameters rather than an arbitrary error message. + + +Usage +~~~~~ + +1. Check if an appropriate error class already exists in `error_classes.py`. + If true, use the error class and skip to step 3. +2. Add a new class to `error_classes.py`; keep in mind the invariants below. +3. Check if the exception type already extends `PySparkException`. + If true, skip to step 5. +4. Mix `PySparkException` into the exception. +5. Throw the exception with the error class and message parameters. + + +**Before** + +Throw with arbitrary error message: + +.. code-block:: python + + raise ValueError("Problem A because B") + + +**After** + +`error_classes.py` + +.. code-block:: python + + "PROBLEM_BECAUSE": { + "message": ["Problem because "] + } + +`exceptions.py` + +.. code-block:: python + + class PySparkTestError(PySparkException): + def __init__(self, error_class: str, message_parameters: Dict[str, str]): + super().__init__(error_class=error_class, message_parameters=message_parameters) + + def getMessageParameters(self) -> Optional[Dict[str, str]]: + return super().getMessageParameters() + +Throw with error class and message parameters: + +.. code-block:: python + + raise PySparkTestError("PROBLEM_BECAUSE", {"problem": "A", "cause": "B"}) + + +Access fields +~~~~~~~~~~~~~ + +To access error fields, catch exceptions that extend :class:`PySparkException` and access to error class with :func:`PySparkException.getErrorClass`. + +.. code-block:: python + + try: + ... + except PySparkException as pe: + if pe.getErrorClass() == "PROBLEM_BECAUSE": + ... + + +Fields +~~~~~~ + +**Error class** + +Error classes are a succinct, human-readable representation of the error category. + +An uncategorized errors can be assigned to a legacy error class with the prefix `_LEGACY_ERROR_TEMP_` and an unused sequential number, for instance `_LEGACY_ERROR_TEMP_0053`. + +Invariants: + +* Unique + +* Consistent across releases + +* Sorted alphabetically + +**Message** + +Error messages provide a descriptive, human-readable representation of the error. +The message format accepts string parameters via the C-style printf syntax. + +The quality of the error message should match the `Apache Spark Error Message Guidelines `_ + +Invariants: + +* Unique diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/debugging.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/debugging.rst new file mode 100644 index 0000000..ef848a6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/debugging.rst @@ -0,0 +1,676 @@ +.. 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. + +================= +Debugging PySpark +================= + +PySpark uses Spark as an engine. PySpark uses `Py4J `_ to leverage Spark to submit and computes the jobs. + +On the driver side, PySpark communicates with the driver on JVM by using `Py4J `_. +When :class:`pyspark.sql.SparkSession` or :class:`pyspark.SparkContext` is created and initialized, PySpark launches a JVM +to communicate. + +On the executor side, Python workers execute and handle Python native functions or data. They are not launched if +a PySpark application does not require interaction between Python workers and JVMs. They are lazily launched only when +Python native functions or data have to be handled, for example, when you execute pandas UDFs or +PySpark RDD APIs. + +This page focuses on debugging Python side of PySpark on both driver and executor sides instead of focusing on debugging +with JVM. Profiling and debugging JVM is described at `Useful Developer Tools `_. + +Note that, + +- If you are running locally, you can directly debug the driver side via using your IDE without the remote debug feature. Setting PySpark with IDEs is documented `here `__. +- *There are many other ways of debugging PySpark applications*. For example, you can remotely debug by using the open source `Remote Debugger `_ instead of using PyCharm Professional documented here. + + +Remote Debugging (PyCharm Professional) +--------------------------------------- + +This section describes remote debugging on both driver and executor sides within a single machine to demonstrate easily. +The ways of debugging PySpark on the executor side is different from doing in the driver. Therefore, they will be demonstrated respectively. +In order to debug PySpark applications on other machines, please refer to the full instructions that are specific +to PyCharm, documented `here `_. + +Firstly, choose **Edit Configuration...** from the *Run* menu. It opens the **Run/Debug Configurations dialog**. +You have to click ``+`` configuration on the toolbar, and from the list of available configurations, select **Python Debug Server**. +Enter the name of this new configuration, for example, ``MyRemoteDebugger`` and also specify the port number, for example ``12345``. + +.. image:: ../../../../docs/img/pyspark-remote-debug1.png + :alt: PyCharm remote debugger setting + +| After that, you should install the corresponding version of the ``pydevd-pycharm`` package in all the machines which will connect to your PyCharm debugger. In the previous dialog, it shows the command to install. + +.. code-block:: text + + pip install pydevd-pycharm~= + +Driver Side +~~~~~~~~~~~ + +To debug on the driver side, your application should be able to connect to the debugging server. Copy and paste the codes +with ``pydevd_pycharm.settrace`` to the top of your PySpark script. Suppose the script name is ``app.py``: + +.. code-block:: bash + + echo "#======================Copy and paste from the previous dialog=========================== + import pydevd_pycharm + pydevd_pycharm.settrace('localhost', port=12345, stdoutToServer=True, stderrToServer=True) + #======================================================================================== + # Your PySpark application codes: + from pyspark.sql import SparkSession + spark = SparkSession.builder.getOrCreate() + spark.range(10).show()" > app.py + +Start to debug with your ``MyRemoteDebugger``. + +.. image:: ../../../../docs/img/pyspark-remote-debug2.png + :alt: PyCharm run remote debugger + +| After that, submit your application. This will connect to your PyCharm debugging server and enable you to debug on the driver side remotely. + +.. code-block:: bash + + spark-submit app.py + +Executor Side +~~~~~~~~~~~~~ + +To debug on the executor side, prepare a Python file as below in your current working directory. + +.. code-block:: bash + + echo "from pyspark import daemon, worker + def remote_debug_wrapped(*args, **kwargs): + #======================Copy and paste from the previous dialog=========================== + import pydevd_pycharm + pydevd_pycharm.settrace('localhost', port=12345, stdoutToServer=True, stderrToServer=True) + #======================================================================================== + worker.main(*args, **kwargs) + daemon.worker_main = remote_debug_wrapped + if __name__ == '__main__': + daemon.manager()" > remote_debug.py + +You will use this file as the Python worker in your PySpark applications by using the ``spark.python.daemon.module`` configuration. +Run the ``pyspark`` shell with the configuration below: + +.. code-block:: bash + + pyspark --conf spark.python.daemon.module=remote_debug + +Now you're ready to remotely debug. Start to debug with your ``MyRemoteDebugger``. + +.. image:: ../../../../docs/img/pyspark-remote-debug2.png + :alt: PyCharm run remote debugger + +| After that, run a job that creates Python workers, for example, as below: + +.. code-block:: python + + spark.range(10).repartition(1).rdd.map(lambda x: x).collect() + + +Checking Resource Usage (``top`` and ``ps``) +-------------------------------------------- + +The Python processes on the driver and executor can be checked via typical ways such as ``top`` and ``ps`` commands. + +Driver Side +~~~~~~~~~~~ + +On the driver side, you can get the process id from your PySpark shell easily as below to know the process id and resources. + +.. code-block:: python + + >>> import os; os.getpid() + 18482 + +.. code-block:: bash + + ps -fe 18482 + +.. code-block:: text + + UID PID PPID C STIME TTY TIME CMD + 000 18482 12345 0 0:00PM ttys001 0:00.00 /.../python + +Executor Side +~~~~~~~~~~~~~ + +To check on the executor side, you can simply ``grep`` them to figure out the process +ids and relevant resources because Python workers are forked from ``pyspark.daemon``. + +.. code-block:: bash + + ps -fe | grep pyspark.daemon + +.. code-block:: text + + 000 12345 1 0 0:00PM ttys000 0:00.00 /.../python -m pyspark.daemon + 000 12345 1 0 0:00PM ttys000 0:00.00 /.../python -m pyspark.daemon + 000 12345 1 0 0:00PM ttys000 0:00.00 /.../python -m pyspark.daemon + 000 12345 1 0 0:00PM ttys000 0:00.00 /.../python -m pyspark.daemon + ... + + +Profiling Memory Usage (Memory Profiler) +---------------------------------------- + +`memory_profiler `_ is one of the profilers that allow you to +check the memory usage line by line. + +Driver Side +~~~~~~~~~~~ + +Unless you are running your driver program in another machine (e.g., YARN cluster mode), this useful tool can be used +to debug the memory usage on driver side easily. Suppose your PySpark script name is ``profile_memory.py``. +You can profile it as below. + +.. code-block:: bash + + echo "from pyspark.sql import SparkSession + #===Your function should be decorated with @profile=== + from memory_profiler import profile + @profile + #===================================================== + def my_func(): + session = SparkSession.builder.getOrCreate() + df = session.range(10000) + return df.collect() + if __name__ == '__main__': + my_func()" > profile_memory.py + +.. code-block:: bash + + python -m memory_profiler profile_memory.py + +.. code-block:: text + + Filename: profile_memory.py + + Line # Mem usage Increment Line Contents + ================================================ + ... + 6 def my_func(): + 7 51.5 MiB 0.6 MiB session = SparkSession.builder.getOrCreate() + 8 51.5 MiB 0.0 MiB df = session.range(10000) + 9 54.4 MiB 2.8 MiB return df.collect() + +Python/Pandas UDF +~~~~~~~~~~~~~~~~~ + +PySpark provides remote `memory_profiler `_ for +Python/Pandas UDFs, which can be enabled by setting ``spark.python.profile.memory`` configuration to ``true``. That +can be used on editors with line numbers such as Jupyter notebooks. An example on a Jupyter notebook is as shown below. + +.. code-block:: bash + + pyspark --conf spark.python.profile.memory=true + + +.. code-block:: python + + from pyspark.sql.functions import pandas_udf + df = spark.range(10) + + @pandas_udf("long") + def add1(x): + return x + 1 + + added = df.select(add1("id")) + added.show() + sc.show_profiles() + + +The result profile is as shown below. + +.. code-block:: text + + ============================================================ + Profile of UDF + ============================================================ + Filename: ... + + Line # Mem usage Increment Occurrences Line Contents + ============================================================= + 4 974.0 MiB 974.0 MiB 10 @pandas_udf("long") + 5 def add1(x): + 6 974.4 MiB 0.4 MiB 10 return x + 1 + +The UDF IDs can be seen in the query plan, for example, ``add1(...)#2L`` in ``ArrowEvalPython`` as shown below. + +.. code-block:: python + + added.explain() + + +.. code-block:: text + + == Physical Plan == + *(2) Project [pythonUDF0#11L AS add1(id)#3L] + +- ArrowEvalPython [add1(id#0L)#2L], [pythonUDF0#11L], 200 + +- *(1) Range (0, 10, step=1, splits=16) + +This feature is not supported with registered UDFs or UDFs with iterators as inputs/outputs. + + +Identifying Hot Loops (Python Profilers) +---------------------------------------- + +`Python Profilers `_ are useful built-in features in Python itself. These +provide deterministic profiling of Python programs with a lot of useful statistics. This section describes how to use it on +both driver and executor sides in order to identify expensive or hot code paths. + +Driver Side +~~~~~~~~~~~ + +To use this on driver side, you can use it as you would do for regular Python programs because PySpark on driver side is a +regular Python process unless you are running your driver program in another machine (e.g., YARN cluster mode). + +.. code-block:: bash + + echo "from pyspark.sql import SparkSession + spark = SparkSession.builder.getOrCreate() + spark.range(10).show()" > app.py + +.. code-block:: bash + + python -m cProfile app.py + +.. code-block:: text + + ... + 129215 function calls (125446 primitive calls) in 5.926 seconds + + Ordered by: standard name + + ncalls tottime percall cumtime percall filename:lineno(function) + 1198/405 0.001 0.000 0.083 0.000 :1009(_handle_fromlist) + 561 0.001 0.000 0.001 0.000 :103(release) + 276 0.000 0.000 0.000 0.000 :143(__init__) + 276 0.000 0.000 0.002 0.000 :147(__enter__) + ... + +Executor Side +~~~~~~~~~~~~~ + +To use this on executor side, PySpark provides remote `Python Profilers `_ for +executor side, which can be enabled by setting ``spark.python.profile`` configuration to ``true``. + +.. code-block:: bash + + pyspark --conf spark.python.profile=true + + +.. code-block:: python + + >>> rdd = sc.parallelize(range(100)).map(str) + >>> rdd.count() + 100 + >>> sc.show_profiles() + ============================================================ + Profile of RDD + ============================================================ + 728 function calls (692 primitive calls) in 0.004 seconds + + Ordered by: internal time, cumulative time + + ncalls tottime percall cumtime percall filename:lineno(function) + 12 0.001 0.000 0.001 0.000 serializers.py:210(load_stream) + 12 0.000 0.000 0.000 0.000 {built-in method _pickle.dumps} + 12 0.000 0.000 0.001 0.000 serializers.py:252(dump_stream) + 12 0.000 0.000 0.001 0.000 context.py:506(f) + ... + +Python/Pandas UDF +~~~~~~~~~~~~~~~~~ + +To use this on Python/Pandas UDFs, PySpark provides remote `Python Profilers `_ for +Python/Pandas UDFs, which can be enabled by setting ``spark.python.profile`` configuration to ``true``. + +.. code-block:: bash + + pyspark --conf spark.python.profile=true + + +.. code-block:: python + + >>> from pyspark.sql.functions import pandas_udf + >>> df = spark.range(10) + >>> @pandas_udf("long") + ... def add1(x): + ... return x + 1 + ... + >>> added = df.select(add1("id")) + + >>> added.show() + +--------+ + |add1(id)| + +--------+ + ... + +--------+ + + >>> sc.show_profiles() + ============================================================ + Profile of UDF + ============================================================ + 2300 function calls (2270 primitive calls) in 0.006 seconds + + Ordered by: internal time, cumulative time + + ncalls tottime percall cumtime percall filename:lineno(function) + 10 0.001 0.000 0.005 0.001 series.py:5515(_arith_method) + 10 0.001 0.000 0.001 0.000 _ufunc_config.py:425(__init__) + 10 0.000 0.000 0.000 0.000 {built-in method _operator.add} + 10 0.000 0.000 0.002 0.000 series.py:315(__init__) + ... + +The UDF IDs can be seen in the query plan, for example, ``add1(...)#2L`` in ``ArrowEvalPython`` below. + +.. code-block:: python + + >>> added.explain() + == Physical Plan == + *(2) Project [pythonUDF0#11L AS add1(id)#3L] + +- ArrowEvalPython [add1(id#0L)#2L], [pythonUDF0#11L], 200 + +- *(1) Range (0, 10, step=1, splits=16) + + +This feature is not supported with registered UDFs. + +Common Exceptions / Errors +-------------------------- + +PySpark SQL +~~~~~~~~~~~ + +**AnalysisException** + +``AnalysisException`` is raised when failing to analyze a SQL query plan. + +Example: + +.. code-block:: python + + >>> df = spark.range(1) + >>> df['bad_key'] + Traceback (most recent call last): + ... + pyspark.errors.exceptions.AnalysisException: Cannot resolve column name "bad_key" among (id) + +Solution: + +.. code-block:: python + + >>> df['id'] + Column<'id'> + +**ParseException** + +``ParseException`` is raised when failing to parse a SQL command. + +Example: + +.. code-block:: python + + >>> spark.sql("select * 1") + Traceback (most recent call last): + ... + pyspark.errors.exceptions.ParseException: + [PARSE_SYNTAX_ERROR] Syntax error at or near '1': extra input '1'.(line 1, pos 9) + + == SQL == + select * 1 + ---------^^^ + +Solution: + +.. code-block:: python + + >>> spark.sql("select *") + DataFrame[] + +**IllegalArgumentException** + +``IllegalArgumentException`` is raised when passing an illegal or inappropriate argument. + +Example: + +.. code-block:: python + + >>> spark.range(1).sample(-1.0) + Traceback (most recent call last): + ... + pyspark.errors.exceptions.IllegalArgumentException: requirement failed: Sampling fraction (-1.0) must be on interval [0, 1] without replacement + +Solution: + +.. code-block:: python + + >>> spark.range(1).sample(1.0) + DataFrame[id: bigint] + +**PythonException** + +``PythonException`` is thrown from Python workers. + +You can see the type of exception that was thrown from the Python worker and its stack trace, as ``TypeError`` below. + +Example: + +.. code-block:: python + + >>> import pyspark.sql.functions as sf + >>> from pyspark.sql.functions import udf + >>> def f(x): + ... return sf.abs(x) + ... + >>> spark.range(-1, 1).withColumn("abs", udf(f)("id")).collect() + 22/04/12 14:52:31 ERROR Executor: Exception in task 7.0 in stage 37.0 (TID 232) + org.apache.spark.api.python.PythonException: Traceback (most recent call last): + ... + TypeError: Invalid argument, not a string or column: -1 of type . For column literals, use 'lit', 'array', 'struct' or 'create_map' function. + +Solution: + +.. code-block:: python + + >>> def f(x): + ... return abs(x) + ... + >>> spark.range(-1, 1).withColumn("abs", udf(f)("id")).collect() + [Row(id=-1, abs='1'), Row(id=0, abs='0')] + +**StreamingQueryException** + +``StreamingQueryException`` is raised when failing a StreamingQuery. Most often, it is thrown from Python workers, that wrap it as a ``PythonException``. + +Example: + +.. code-block:: python + + >>> sdf = spark.readStream.format("text").load("python/test_support/sql/streaming") + >>> from pyspark.sql.functions import col, udf + >>> bad_udf = udf(lambda x: 1 / 0) + >>> (sdf.select(bad_udf(col("value"))).writeStream.format("memory").queryName("q1").start()).processAllAvailable() + Traceback (most recent call last): + ... + org.apache.spark.api.python.PythonException: Traceback (most recent call last): + File "", line 1, in + ZeroDivisionError: division by zero + ... + pyspark.errors.exceptions.StreamingQueryException: [STREAM_FAILED] Query [id = 74eb53a8-89bd-49b0-9313-14d29eed03aa, runId = 9f2d5cf6-a373-478d-b718-2c2b6d8a0f24] terminated with exception: Job aborted + +Solution: + +Fix the StreamingQuery and re-execute the workflow. + +**SparkUpgradeException** + +``SparkUpgradeException`` is thrown because of Spark upgrade. + +Example: + +.. code-block:: python + + >>> from pyspark.sql.functions import to_date, unix_timestamp, from_unixtime + >>> df = spark.createDataFrame([("2014-31-12",)], ["date_str"]) + >>> df2 = df.select("date_str", to_date(from_unixtime(unix_timestamp("date_str", "yyyy-dd-aa")))) + >>> df2.collect() + Traceback (most recent call last): + ... + pyspark.sql.utils.SparkUpgradeException: You may get a different result due to the upgrading to Spark >= 3.0: Fail to recognize 'yyyy-dd-aa' pattern in the DateTimeFormatter. 1) You can set spark.sql.legacy.timeParserPolicy to LEGACY to restore the behavior before Spark 3.0. 2) You can form a valid datetime pattern with the guide from https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html + +Solution: + +.. code-block:: python + + >>> spark.conf.set("spark.sql.legacy.timeParserPolicy", "LEGACY") + >>> df2 = df.select("date_str", to_date(from_unixtime(unix_timestamp("date_str", "yyyy-dd-aa")))) + >>> df2.collect() + [Row(date_str='2014-31-12', to_date(from_unixtime(unix_timestamp(date_str, yyyy-dd-aa), yyyy-MM-dd HH:mm:ss))=None)] + +pandas API on Spark +~~~~~~~~~~~~~~~~~~~ + +There are specific common exceptions / errors in pandas API on Spark. + +**ValueError: Cannot combine the series or dataframe because it comes from a different dataframe** + +Operations involving more than one series or dataframes raises a ``ValueError`` if ``compute.ops_on_diff_frames`` is disabled (disabled by default). Such operations may be expensive due to joining of underlying Spark frames. So users should be aware of the cost and enable that flag only when necessary. + +Exception: + +.. code-block:: python + + >>> ps.Series([1, 2]) + ps.Series([3, 4]) + Traceback (most recent call last): + ... + ValueError: Cannot combine the series or dataframe because it comes from a different dataframe. In order to allow this operation, enable 'compute.ops_on_diff_frames' option. + + +Solution: + +.. code-block:: python + + >>> with ps.option_context('compute.ops_on_diff_frames', True): + ... ps.Series([1, 2]) + ps.Series([3, 4]) + ... + 0 4 + 1 6 + dtype: int64 + +**RuntimeError: Result vector from pandas_udf was not the required length** + +Exception: + +.. code-block:: python + + >>> def f(x) -> ps.Series[np.int32]: + ... return x[:-1] + ... + >>> ps.DataFrame({"x":[1, 2], "y":[3, 4]}).transform(f) + 22/04/12 13:46:39 ERROR Executor: Exception in task 2.0 in stage 16.0 (TID 88) + org.apache.spark.api.python.PythonException: Traceback (most recent call last): + ... + RuntimeError: Result vector from pandas_udf was not the required length: expected 1, got 0 + +Solution: + +.. code-block:: python + + >>> def f(x) -> ps.Series[np.int32]: + ... return x + ... + >>> ps.DataFrame({"x":[1, 2], "y":[3, 4]}).transform(f) + x y + 0 1 3 + 1 2 4 + +Py4j +~~~~ + +**Py4JJavaError** + +``Py4JJavaError`` is raised when an exception occurs in the Java client code. +You can see the type of exception that was thrown on the Java side and its stack trace, as ``java.lang.NullPointerException`` below. + +Example: + +.. code-block:: python + + >>> spark.sparkContext._jvm.java.lang.String(None) + Traceback (most recent call last): + ... + py4j.protocol.Py4JJavaError: An error occurred while calling None.java.lang.String. + : java.lang.NullPointerException + .. + +Solution: + +.. code-block:: python + + >>> spark.sparkContext._jvm.java.lang.String("x") + 'x' + +**Py4JError** + +``Py4JError`` is raised when any other error occurs such as when the Python client program tries to access an object that no longer exists on the Java side. + +Example: + +.. code-block:: python + + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.regression import LinearRegression + >>> df = spark.createDataFrame( + ... [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ... ["label", "weight", "features"], + ... ) + >>> lr = LinearRegression( + ... maxIter=1, regParam=0.0, solver="normal", weightCol="weight", fitIntercept=False + ... ) + >>> model = lr.fit(df) + >>> model + LinearRegressionModel: uid=LinearRegression_eb7bc1d4bf25, numFeatures=1 + >>> model.__del__() + >>> model + Traceback (most recent call last): + ... + py4j.protocol.Py4JError: An error occurred while calling o531.toString. Trace: + py4j.Py4JException: Target Object ID does not exist for this gateway :o531 + ... + +Solution: + +Access an object that exists on the Java side. + +**Py4JNetworkError** + +``Py4JNetworkError`` is raised when a problem occurs during network transfer (e.g., connection lost). In this case, we shall debug the network and rebuild the connection. + +Stack Traces +------------ + +There are Spark configurations to control stack traces: + +- ``spark.sql.execution.pyspark.udf.simplifiedTraceback.enabled`` is true by default to simplify traceback from Python UDFs. + +- ``spark.sql.pyspark.jvmStacktrace.enabled`` is false by default to hide JVM stacktrace and to show a Python-friendly exception only. + +Spark configurations above are independent from log level settings. Control log levels through :meth:`pyspark.SparkContext.setLogLevel`. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/errors.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/errors.rst new file mode 100644 index 0000000..f48c224 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/errors.rst @@ -0,0 +1,92 @@ +.. 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. + +======================== +Error classes in PySpark +======================== + +This is a list of common, named error classes returned by PySpark which are defined at `error_classes.py `_. + +When writing PySpark errors, developers must use an error class from the list. If an appropriate error class is not available, add a new one into the list. For more information, please refer to `Contributing Error and Exception `_. + ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| Error class | Error message | ++============================================================+==============================================================================================================+ +| ARGUMENT_REQUIRED | Argument `` is required when . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| COLUMN_IN_LIST | `` does not allow a Column in a list. | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| DISALLOWED_TYPE_FOR_CONTAINER | Argument ``(`type`: ) should only contain a type in [], got | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN | Function `` should return Column, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_BOOL | Argument `` should be a bool, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_LIST_OR_STR_OR_TUPLE | Argument `` should be a bool, dict, float, int, str or tuple, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_STR | Argument `` should be a bool, dict, float, int or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_BOOL_OR_LIST | Argument `` should be a bool or list, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_BOOL_OR_STR | Argument `` should be a bool or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN | Argument `` should be a Column, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN_OR_DATATYPE_OR_STR | Argument `` should be a Column, str or DataType, but got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN_OR_FLOAT_OR_INT_OR_LIST_OR_STR | Argument `` should be a column, float, integer, list or string, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN_OR_INT | Argument `` should be a Column or int, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN_OR_INT_OR_STR | Argument `` should be a Column, int or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_COLUMN_OR_STR | Argument `` should be a Column or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_DATAFRAME | Argument `` should be a DataFrame, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_DATATYPE_OR_STR | Argument `` should be a DataType or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_DICT | Argument `` should be a dict, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_EXPRESSION | Argument should be a Expression, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_FLOAT_OR_INT | Argument should be a float or int, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_FLOAT_OR_INT_OR_LIST_OR_STR | Argument should be a float, int, list or str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_INT | Argument should be an int, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_ITERABLE | is not iterable. | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_LIST_OR_STR_OR_TUPLE | Argument should be a list, str or tuple, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_LIST_OR_TUPLE | Argument should be a list or tuple, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_SAME_TYPE | Argument and should be the same type, got and . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_STR | Argument should be a str, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| NOT_WINDOWSPEC | Argument should be a WindowSpec, got . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| UNSUPPORTED_NUMPY_ARRAY_SCALAR | The type of array scalar '' is not supported. | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION | Function should use only POSITIONAL or POSITIONAL OR KEYWORD arguments. | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION | Function should take between 1 and 3 arguments, but provided function takes . | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| WRONG_NUM_COLUMNS | Function should take at least columns. | ++------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/index.rst new file mode 100644 index 0000000..50150e0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/index.rst @@ -0,0 +1,30 @@ +.. 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. + +=========== +Development +=========== + +.. toctree:: + :maxdepth: 2 + + contributing + testing + debugging + setting_ide + errors + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/setting_ide.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/setting_ide.rst new file mode 100644 index 0000000..6e8f014 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/setting_ide.rst @@ -0,0 +1,62 @@ +.. 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. + +=============== +Setting up IDEs +=============== + + +PyCharm +------- + +This section describes how to setup PySpark on PyCharm. +It guides step by step to the process of downloading the source code from GitHub and running the test code successfully. + +Firstly, download the Spark source code from GitHub using git url. You can download the source code by simply using ``git clone`` command as shown below. +If you want to download the code from any forked repository rather than Spark original repository, please change the url properly. + +.. code-block:: bash + + git clone https://github.com/apache/spark.git + +When the download is completed, go to the ``spark`` directory and build the package. +SBT build is generally much faster than Maven. More details about the build are documented `here `_. + +.. code-block:: bash + + build/sbt package + +After building is finished, run PyCharm and select the path ``spark/python``. + +.. image:: ../../../../docs/img/pycharm-with-pyspark1.png + :alt: Select the Spark path + + +Let's go to the path ``python/pyspark/tests`` in PyCharm and try to run the any test like ``test_join.py``. +You might can see the ``KeyError: 'SPARK_HOME'`` because the environment variable has not been set yet. + +Go **Run -> Edit Configurations**, and set the environment variables as below. +Please make sure to specify your own path for ``SPARK_HOME`` rather than ``/.../spark``. After completing the variable, click **OK** to apply the changes. + +.. image:: ../../../../docs/img/pycharm-with-pyspark2.png + :alt: Setting up SPARK_HOME + + +Once ``SPARK_HOME`` is set properly, you'll be able to run the tests properly as below: + +.. image:: ../../../../docs/img/pycharm-with-pyspark3.png + :alt: Running tests properly diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/development/testing.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/development/testing.rst new file mode 100644 index 0000000..46a11ed --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/development/testing.rst @@ -0,0 +1,89 @@ +.. 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. + +=============== +Testing PySpark +=============== + +In order to run PySpark tests, you should build Spark itself first via Maven or SBT. For example, + +.. code-block:: bash + + build/mvn -DskipTests clean package + +.. code-block:: bash + + build/sbt -Phive clean package + + +After that, the PySpark test cases can be run via using ``python/run-tests``. For example, + +.. code-block:: bash + + python/run-tests --python-executable=python3 + +Note that you may set ``OBJC_DISABLE_INITIALIZE_FORK_SAFETY`` environment variable to ``YES`` if you are running tests on Mac OS. + +Please see the guidance on how to `build Spark `_, +`run tests for a module, or individual tests `_. + + +Running Individual PySpark Tests +-------------------------------- + +You can run a specific test via using ``python/run-tests``, for example, as below: + +.. code-block:: bash + + python/run-tests --testnames pyspark.sql.tests.test_arrow + +Please refer to `Testing PySpark `_ for more details. + + +Running Tests using GitHub Actions +---------------------------------- + +You can run the full PySpark tests by using GitHub Actions in your own forked GitHub +repository with a few clicks. Please refer to +`Running tests in your forked repository using GitHub Actions `_ for more details. + + +Running Tests for Spark Connect +------------------------------- + +Running Tests for Python Client +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +In order to test the changes in Protobuf definitions, for example, at +`spark/connector/connect/common/src/main/protobuf/spark/connect `_, +you should regenerate Python Protobuf client first by running ``dev/connect-gen-protos.sh``. + + +Running PySpark Shell with Python Client +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For Apache Spark you locally built: + +.. code-block:: bash + + bin/pyspark --remote "local[*]" + +For the Apache Spark release: + +.. code-block:: bash + + bin/pyspark --remote "local[*]" --packages org.apache.spark:spark-connect_2.12:3.4.0 diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/index.rst new file mode 100644 index 0000000..5f6d306 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/index.rst @@ -0,0 +1,43 @@ +.. 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. + + +=============== +Getting Started +=============== + +This page summarizes the basic steps required to setup and get started with PySpark. +There are more guides shared with other languages such as +`Quick Start `_ in Programming Guides +at `the Spark documentation `_. + +There are live notebooks where you can try PySpark out without any other step: + +* |binder_df|_ +* |binder_connect|_ +* |binder_ps|_ + +The list below is the contents of this quickstart page: + +.. toctree:: + :maxdepth: 2 + + install + quickstart_df + quickstart_connect + quickstart_ps + testing_pyspark diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/install.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/install.rst new file mode 100644 index 0000000..6822285 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/install.rst @@ -0,0 +1,170 @@ +.. 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. + +============ +Installation +============ + +PySpark is included in the official releases of Spark available in the `Apache Spark website `_. +For Python users, PySpark also provides ``pip`` installation from PyPI. This is usually for local usage or as +a client to connect to a cluster instead of setting up a cluster itself. + +This page includes instructions for installing PySpark by using pip, Conda, downloading manually, +and building from the source. + + +Python Versions Supported +------------------------- + +Python 3.8 and above. + + +Using PyPI +---------- + +PySpark installation using `PyPI `_ is as follows: + +.. code-block:: bash + + pip install pyspark + +If you want to install extra dependencies for a specific component, you can install it as below: + +.. code-block:: bash + + # Spark SQL + pip install pyspark[sql] + # pandas API on Spark + pip install pyspark[pandas_on_spark] plotly # to plot your data, you can install plotly together. + # Spark Connect + pip install pyspark[connect] + +For PySpark with/without a specific Hadoop version, you can install it by using ``PYSPARK_HADOOP_VERSION`` environment variables as below: + +.. code-block:: bash + + PYSPARK_HADOOP_VERSION=3 pip install pyspark + +The default distribution uses Hadoop 3.3 and Hive 2.3. If users specify different versions of Hadoop, the pip installation automatically +downloads a different version and uses it in PySpark. Downloading it can take a while depending on +the network and the mirror chosen. ``PYSPARK_RELEASE_MIRROR`` can be set to manually choose the mirror for faster downloading. + +.. code-block:: bash + + PYSPARK_RELEASE_MIRROR=http://mirror.apache-kr.org PYSPARK_HADOOP_VERSION=3 pip install + +It is recommended to use ``-v`` option in ``pip`` to track the installation and download status. + +.. code-block:: bash + + PYSPARK_HADOOP_VERSION=3 pip install pyspark -v + +Supported values in ``PYSPARK_HADOOP_VERSION`` are: + +- ``without``: Spark pre-built with user-provided Apache Hadoop +- ``3``: Spark pre-built for Apache Hadoop 3.3 and later (default) + +Note that this installation of PySpark with/without a specific Hadoop version is experimental. It can change or be removed between minor releases. + + +Using Conda +----------- + +Conda is an open-source package management and environment management system (developed by +`Anaconda `_), which is best installed through +`Miniconda `_ or `Miniforge `_. +The tool is both cross-platform and language agnostic, and in practice, conda can replace both +`pip `_ and `virtualenv `_. + +Conda uses so-called channels to distribute packages, and together with the default channels by +Anaconda itself, the most important channel is `conda-forge `_, which +is the community-driven packaging effort that is the most extensive & the most current (and also +serves as the upstream for the Anaconda channels in most cases). + +To create a new conda environment from your terminal and activate it, proceed as shown below: + +.. code-block:: bash + + conda create -n pyspark_env + conda activate pyspark_env + +After activating the environment, use the following command to install pyspark, +a python version of your choice, as well as other packages you want to use in +the same session as pyspark (you can install in several steps too). + +.. code-block:: bash + + conda install -c conda-forge pyspark # can also add "python=3.8 some_package [etc.]" here + +Note that `PySpark for conda `_ is maintained +separately by the community; while new versions generally get packaged quickly, the +availability through conda(-forge) is not directly in sync with the PySpark release cycle. + +While using pip in a conda environment is technically feasible (with the same command as +`above <#using-pypi>`_), this approach is `discouraged `_, +because pip does not interoperate with conda. + +For a short summary about useful conda commands, see their +`cheat sheet `_. + + +Manually Downloading +-------------------- + +PySpark is included in the distributions available at the `Apache Spark website `_. +You can download a distribution you want from the site. After that, uncompress the tar file into the directory where you want +to install Spark, for example, as below: + +.. parsed-literal:: + + tar xzvf spark-\ |release|\-bin-hadoop3.tgz + +Ensure the ``SPARK_HOME`` environment variable points to the directory where the tar file has been extracted. +Update ``PYTHONPATH`` environment variable such that it can find the PySpark and Py4J under ``SPARK_HOME/python/lib``. +One example of doing this is shown below: + +.. parsed-literal:: + + cd spark-\ |release|\-bin-hadoop3 + export SPARK_HOME=`pwd` + export PYTHONPATH=$(ZIPS=("$SPARK_HOME"/python/lib/*.zip); IFS=:; echo "${ZIPS[*]}"):$PYTHONPATH + + +Installing from Source +---------------------- + +To install PySpark from source, refer to |building_spark|_. + + +Dependencies +------------ +========================== ========================= ====================================================================================== +Package Supported version Note +========================== ========================= ====================================================================================== +`py4j` >=0.10.9.7 Required +`pandas` >=1.0.5 Required for pandas API on Spark and Spark Connect; Optional for Spark SQL +`pyarrow` >=4.0.0 Required for pandas API on Spark and Spark Connect; Optional for Spark SQL +`numpy` >=1.15 Required for pandas API on Spark and MLLib DataFrame-based API; Optional for Spark SQL +`grpcio` >=1.48,<1.57 Required for Spark Connect +`grpcio-status` >=1.48,<1.57 Required for Spark Connect +`googleapis-common-protos` ==1.56.4 Required for Spark Connect +========================== ========================= ====================================================================================== + +Note that PySpark requires Java 8 or later with ``JAVA_HOME`` properly set. +If using JDK 11, set ``-Dio.netty.tryReflectionSetAccessible=true`` for Arrow related features and refer +to |downloading|_. + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_connect.ipynb b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_connect.ipynb new file mode 100644 index 0000000..15a2ab7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_connect.ipynb @@ -0,0 +1,142 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Quickstart: Spark Connect\n", + "\n", + "Spark Connect introduced a decoupled client-server architecture for Spark that allows remote connectivity to Spark clusters using the [DataFrame API](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.html?highlight=dataframe#pyspark.sql.DataFrame).\n", + "\n", + "This notebook walks through a simple step-by-step example of how to use Spark Connect to build any type of application that needs to leverage the power of Spark when working with data.\n", + "\n", + "Spark Connect includes both client and server components and we will show you how to set up and use both." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch Spark server with Spark Connect\n", + "\n", + "To launch Spark with support for Spark Connect sessions, run the `start-connect-server.sh` script." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "!$HOME/sbin/start-connect-server.sh --packages org.apache.spark:spark-connect_2.12:$SPARK_VERSION" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Connect to Spark Connect server\n", + "\n", + "Now that the Spark server is running, we can connect to it remotely using Spark Connect. We do this by creating a remote Spark session on the client where our application runs. Before we can do that, we need to make sure to stop the existing regular Spark session because it cannot coexist with the remote Spark Connect session we are about to create." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "from pyspark.sql import SparkSession\n", + "\n", + "SparkSession.builder.master(\"local[*]\").getOrCreate().stop()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The command we used above to launch the server configured Spark to run as `localhost:15002`. So now we can create a remote Spark session on the client using the following command." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "spark = SparkSession.builder.remote(\"sc://localhost:15002\").getOrCreate()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create DataFrame\n", + "\n", + "Once the remote Spark session is created successfully, it can be used the same way as a regular Spark session. Therefore, you can create a DataFrame with the following command." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+\n", + "| a| b| c| d| e|\n", + "+---+---+-------+----------+-------------------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|\n", + "| 2|3.0|string2|2000-02-01|2000-01-02 12:00:00|\n", + "| 4|5.0|string3|2000-03-01|2000-01-03 12:00:00|\n", + "+---+---+-------+----------+-------------------+\n", + "\n" + ] + } + ], + "source": [ + "from datetime import datetime, date\n", + "from pyspark.sql import Row\n", + "\n", + "df = spark.createDataFrame([\n", + " Row(a=1, b=2., c='string1', d=date(2000, 1, 1), e=datetime(2000, 1, 1, 12, 0)),\n", + " Row(a=2, b=3., c='string2', d=date(2000, 2, 1), e=datetime(2000, 1, 2, 12, 0)),\n", + " Row(a=4, b=5., c='string3', d=date(2000, 3, 1), e=datetime(2000, 1, 3, 12, 0))\n", + "])\n", + "df.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "See 'Live Notebook: DataFrame' at [the quickstart page](https://spark.apache.org/docs/latest/api/python/getting_started/index.html) for more detail usage of DataFrame API." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.11" + }, + "name": "quickstart", + "notebookId": 1927513300154480 + }, + "nbformat": 4, + "nbformat_minor": 1 +} diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_df.ipynb b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_df.ipynb new file mode 100644 index 0000000..f1c04c8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_df.ipynb @@ -0,0 +1,1144 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Quickstart: DataFrame\n", + "\n", + "This is a short introduction and quickstart for the PySpark DataFrame API. PySpark DataFrames are lazily evaluated. They are implemented on top of [RDD](https://spark.apache.org/docs/latest/rdd-programming-guide.html#overview)s. When Spark [transforms](https://spark.apache.org/docs/latest/rdd-programming-guide.html#transformations) data, it does not immediately compute the transformation but plans how to compute later. When [actions](https://spark.apache.org/docs/latest/rdd-programming-guide.html#actions) such as `collect()` are explicitly called, the computation starts.\n", + "This notebook shows the basic usages of the DataFrame, geared mainly for new users. You can run the latest version of these examples by yourself in 'Live Notebook: DataFrame' at [the quickstart page](https://spark.apache.org/docs/latest/api/python/getting_started/index.html).\n", + "\n", + "There is also other useful information in Apache Spark documentation site, see the latest version of [Spark SQL and DataFrames](https://spark.apache.org/docs/latest/sql-programming-guide.html), [RDD Programming Guide](https://spark.apache.org/docs/latest/rdd-programming-guide.html), [Structured Streaming Programming Guide](https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html), [Spark Streaming Programming Guide](https://spark.apache.org/docs/latest/streaming-programming-guide.html) and [Machine Learning Library (MLlib) Guide](https://spark.apache.org/docs/latest/ml-guide.html).\n", + "\n", + "PySpark applications start with initializing `SparkSession` which is the entry point of PySpark as below. In case of running it in PySpark shell via pyspark executable, the shell automatically creates the session in the variable spark for users." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from pyspark.sql import SparkSession\n", + "\n", + "spark = SparkSession.builder.getOrCreate()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## DataFrame Creation\n", + "\n", + "A PySpark DataFrame can be created via `pyspark.sql.SparkSession.createDataFrame` typically by passing a list of lists, tuples, dictionaries and `pyspark.sql.Row`s, a [pandas DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) and an RDD consisting of such a list.\n", + "`pyspark.sql.SparkSession.createDataFrame` takes the `schema` argument to specify the schema of the DataFrame. When it is omitted, PySpark infers the corresponding schema by taking a sample from the data.\n", + "\n", + "Firstly, you can create a PySpark DataFrame from a list of rows" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "DataFrame[a: bigint, b: double, c: string, d: date, e: timestamp]" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from datetime import datetime, date\n", + "import pandas as pd\n", + "from pyspark.sql import Row\n", + "\n", + "df = spark.createDataFrame([\n", + " Row(a=1, b=2., c='string1', d=date(2000, 1, 1), e=datetime(2000, 1, 1, 12, 0)),\n", + " Row(a=2, b=3., c='string2', d=date(2000, 2, 1), e=datetime(2000, 1, 2, 12, 0)),\n", + " Row(a=4, b=5., c='string3', d=date(2000, 3, 1), e=datetime(2000, 1, 3, 12, 0))\n", + "])\n", + "df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a PySpark DataFrame with an explicit schema." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "DataFrame[a: bigint, b: double, c: string, d: date, e: timestamp]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = spark.createDataFrame([\n", + " (1, 2., 'string1', date(2000, 1, 1), datetime(2000, 1, 1, 12, 0)),\n", + " (2, 3., 'string2', date(2000, 2, 1), datetime(2000, 1, 2, 12, 0)),\n", + " (3, 4., 'string3', date(2000, 3, 1), datetime(2000, 1, 3, 12, 0))\n", + "], schema='a long, b double, c string, d date, e timestamp')\n", + "df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a PySpark DataFrame from a pandas DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "DataFrame[a: bigint, b: double, c: string, d: date, e: timestamp]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pandas_df = pd.DataFrame({\n", + " 'a': [1, 2, 3],\n", + " 'b': [2., 3., 4.],\n", + " 'c': ['string1', 'string2', 'string3'],\n", + " 'd': [date(2000, 1, 1), date(2000, 2, 1), date(2000, 3, 1)],\n", + " 'e': [datetime(2000, 1, 1, 12, 0), datetime(2000, 1, 2, 12, 0), datetime(2000, 1, 3, 12, 0)]\n", + "})\n", + "df = spark.createDataFrame(pandas_df)\n", + "df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The DataFrames created above all have the same results and schema." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+\n", + "| a| b| c| d| e|\n", + "+---+---+-------+----------+-------------------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|\n", + "| 2|3.0|string2|2000-02-01|2000-01-02 12:00:00|\n", + "| 3|4.0|string3|2000-03-01|2000-01-03 12:00:00|\n", + "+---+---+-------+----------+-------------------+\n", + "\n", + "root\n", + " |-- a: long (nullable = true)\n", + " |-- b: double (nullable = true)\n", + " |-- c: string (nullable = true)\n", + " |-- d: date (nullable = true)\n", + " |-- e: timestamp (nullable = true)\n", + "\n" + ] + } + ], + "source": [ + "# All DataFrames above result same.\n", + "df.show()\n", + "df.printSchema()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Viewing Data\n", + "\n", + "The top rows of a DataFrame can be displayed using `DataFrame.show()`." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+\n", + "| a| b| c| d| e|\n", + "+---+---+-------+----------+-------------------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|\n", + "+---+---+-------+----------+-------------------+\n", + "only showing top 1 row\n", + "\n" + ] + } + ], + "source": [ + "df.show(1)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Alternatively, you can enable `spark.sql.repl.eagerEval.enabled` configuration for the eager evaluation of PySpark DataFrame in notebooks such as Jupyter. The number of rows to show can be controlled via `spark.sql.repl.eagerEval.maxNumRows` configuration." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "
    abcde
    12.0string12000-01-012000-01-01 12:00:00
    23.0string22000-02-012000-01-02 12:00:00
    34.0string32000-03-012000-01-03 12:00:00
    \n" + ], + "text/plain": [ + "DataFrame[a: bigint, b: double, c: string, d: date, e: timestamp]" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "spark.conf.set('spark.sql.repl.eagerEval.enabled', True)\n", + "df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The rows can also be shown vertically. This is useful when rows are too long to show horizontally." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "-RECORD 0------------------\n", + " a | 1 \n", + " b | 2.0 \n", + " c | string1 \n", + " d | 2000-01-01 \n", + " e | 2000-01-01 12:00:00 \n", + "only showing top 1 row\n", + "\n" + ] + } + ], + "source": [ + "df.show(1, vertical=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see the DataFrame's schema and column names as follows:" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['a', 'b', 'c', 'd', 'e']" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.columns" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "root\n", + " |-- a: long (nullable = true)\n", + " |-- b: double (nullable = true)\n", + " |-- c: string (nullable = true)\n", + " |-- d: date (nullable = true)\n", + " |-- e: timestamp (nullable = true)\n", + "\n" + ] + } + ], + "source": [ + "df.printSchema()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Show the summary of the DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-------+---+---+-------+\n", + "|summary| a| b| c|\n", + "+-------+---+---+-------+\n", + "| count| 3| 3| 3|\n", + "| mean|2.0|3.0| null|\n", + "| stddev|1.0|1.0| null|\n", + "| min| 1|2.0|string1|\n", + "| max| 3|4.0|string3|\n", + "+-------+---+---+-------+\n", + "\n" + ] + } + ], + "source": [ + "df.select(\"a\", \"b\", \"c\").describe().show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`DataFrame.collect()` collects the distributed data to the driver side as the local data in Python. Note that this can throw an out-of-memory error when the dataset is too large to fit in the driver side because it collects all the data from executors to the driver side." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[Row(a=1, b=2.0, c='string1', d=datetime.date(2000, 1, 1), e=datetime.datetime(2000, 1, 1, 12, 0)),\n", + " Row(a=2, b=3.0, c='string2', d=datetime.date(2000, 2, 1), e=datetime.datetime(2000, 1, 2, 12, 0)),\n", + " Row(a=3, b=4.0, c='string3', d=datetime.date(2000, 3, 1), e=datetime.datetime(2000, 1, 3, 12, 0))]" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.collect()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In order to avoid throwing an out-of-memory exception, use `DataFrame.take()` or `DataFrame.tail()`." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[Row(a=1, b=2.0, c='string1', d=datetime.date(2000, 1, 1), e=datetime.datetime(2000, 1, 1, 12, 0))]" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.take(1)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "PySpark DataFrame also provides the conversion back to a [pandas DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) to leverage pandas API. Note that `toPandas` also collects all data into the driver side that can easily cause an out-of-memory-error when the data is too large to fit into the driver side." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    abcde
    012.0string12000-01-012000-01-01 12:00:00
    123.0string22000-02-012000-01-02 12:00:00
    234.0string32000-03-012000-01-03 12:00:00
    \n", + "
    " + ], + "text/plain": [ + " a b c d e\n", + "0 1 2.0 string1 2000-01-01 2000-01-01 12:00:00\n", + "1 2 3.0 string2 2000-02-01 2000-01-02 12:00:00\n", + "2 3 4.0 string3 2000-03-01 2000-01-03 12:00:00" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.toPandas()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Selecting and Accessing Data\n", + "\n", + "PySpark DataFrame is lazily evaluated and simply selecting a column does not trigger the computation but it returns a `Column` instance." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Column" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.a" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In fact, most of column-wise operations return `Column`s." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from pyspark.sql import Column\n", + "from pyspark.sql.functions import upper\n", + "\n", + "type(df.c) == type(upper(df.c)) == type(df.c.isNull())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These `Column`s can be used to select the columns from a DataFrame. For example, `DataFrame.select()` takes the `Column` instances that returns another DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-------+\n", + "| c|\n", + "+-------+\n", + "|string1|\n", + "|string2|\n", + "|string3|\n", + "+-------+\n", + "\n" + ] + } + ], + "source": [ + "df.select(df.c).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Assign new `Column` instance." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+-------+\n", + "| a| b| c| d| e|upper_c|\n", + "+---+---+-------+----------+-------------------+-------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|STRING1|\n", + "| 2|3.0|string2|2000-02-01|2000-01-02 12:00:00|STRING2|\n", + "| 3|4.0|string3|2000-03-01|2000-01-03 12:00:00|STRING3|\n", + "+---+---+-------+----------+-------------------+-------+\n", + "\n" + ] + } + ], + "source": [ + "df.withColumn('upper_c', upper(df.c)).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To select a subset of rows, use `DataFrame.filter()`." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+\n", + "| a| b| c| d| e|\n", + "+---+---+-------+----------+-------------------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|\n", + "+---+---+-------+----------+-------------------+\n", + "\n" + ] + } + ], + "source": [ + "df.filter(df.a == 1).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Applying a Function\n", + "\n", + "PySpark supports various UDFs and APIs to allow users to execute Python native functions. See also the latest [Pandas UDFs](https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#pandas-udfs-aka-vectorized-udfs) and [Pandas Function APIs](https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#pandas-function-apis). For instance, the example below allows users to directly use the APIs in [a pandas Series](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.html) within Python native function." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+------------------+\n", + "|pandas_plus_one(a)|\n", + "+------------------+\n", + "| 2|\n", + "| 3|\n", + "| 4|\n", + "+------------------+\n", + "\n" + ] + } + ], + "source": [ + "import pandas as pd\n", + "from pyspark.sql.functions import pandas_udf\n", + "\n", + "@pandas_udf('long')\n", + "def pandas_plus_one(series: pd.Series) -> pd.Series:\n", + " # Simply plus one by using pandas Series.\n", + " return series + 1\n", + "\n", + "df.select(pandas_plus_one(df.a)).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Another example is `DataFrame.mapInPandas` which allows users directly use the APIs in a [pandas DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) without any restrictions such as the result length." + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+---+-------+----------+-------------------+\n", + "| a| b| c| d| e|\n", + "+---+---+-------+----------+-------------------+\n", + "| 1|2.0|string1|2000-01-01|2000-01-01 12:00:00|\n", + "+---+---+-------+----------+-------------------+\n", + "\n" + ] + } + ], + "source": [ + "def pandas_filter_func(iterator):\n", + " for pandas_df in iterator:\n", + " yield pandas_df[pandas_df.a == 1]\n", + "\n", + "df.mapInPandas(pandas_filter_func, schema=df.schema).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Grouping Data\n", + "\n", + "PySpark DataFrame also provides a way of handling grouped data by using the common approach, split-apply-combine strategy.\n", + "It groups the data by a certain condition applies a function to each group and then combines them back to the DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+------+---+---+\n", + "|color| fruit| v1| v2|\n", + "+-----+------+---+---+\n", + "| red|banana| 1| 10|\n", + "| blue|banana| 2| 20|\n", + "| red|carrot| 3| 30|\n", + "| blue| grape| 4| 40|\n", + "| red|carrot| 5| 50|\n", + "|black|carrot| 6| 60|\n", + "| red|banana| 7| 70|\n", + "| red| grape| 8| 80|\n", + "+-----+------+---+---+\n", + "\n" + ] + } + ], + "source": [ + "df = spark.createDataFrame([\n", + " ['red', 'banana', 1, 10], ['blue', 'banana', 2, 20], ['red', 'carrot', 3, 30],\n", + " ['blue', 'grape', 4, 40], ['red', 'carrot', 5, 50], ['black', 'carrot', 6, 60],\n", + " ['red', 'banana', 7, 70], ['red', 'grape', 8, 80]], schema=['color', 'fruit', 'v1', 'v2'])\n", + "df.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Grouping and then applying the `avg()` function to the resulting groups." + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+-------+-------+\n", + "|color|avg(v1)|avg(v2)|\n", + "+-----+-------+-------+\n", + "| red| 4.8| 48.0|\n", + "|black| 6.0| 60.0|\n", + "| blue| 3.0| 30.0|\n", + "+-----+-------+-------+\n", + "\n" + ] + } + ], + "source": [ + "df.groupby('color').avg().show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can also apply a Python native function against each group by using pandas API." + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+------+---+---+\n", + "|color| fruit| v1| v2|\n", + "+-----+------+---+---+\n", + "| red|banana| -3| 10|\n", + "| red|carrot| -1| 30|\n", + "| red|carrot| 0| 50|\n", + "| red|banana| 2| 70|\n", + "| red| grape| 3| 80|\n", + "|black|carrot| 0| 60|\n", + "| blue|banana| -1| 20|\n", + "| blue| grape| 1| 40|\n", + "+-----+------+---+---+\n", + "\n" + ] + } + ], + "source": [ + "def plus_mean(pandas_df):\n", + " return pandas_df.assign(v1=pandas_df.v1 - pandas_df.v1.mean())\n", + "\n", + "df.groupby('color').applyInPandas(plus_mean, schema=df.schema).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Co-grouping and applying a function." + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+--------+---+---+---+\n", + "| time| id| v1| v2|\n", + "+--------+---+---+---+\n", + "|20000101| 1|1.0| x|\n", + "|20000102| 1|3.0| x|\n", + "|20000101| 2|2.0| y|\n", + "|20000102| 2|4.0| y|\n", + "+--------+---+---+---+\n", + "\n" + ] + } + ], + "source": [ + "df1 = spark.createDataFrame(\n", + " [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)],\n", + " ('time', 'id', 'v1'))\n", + "\n", + "df2 = spark.createDataFrame(\n", + " [(20000101, 1, 'x'), (20000101, 2, 'y')],\n", + " ('time', 'id', 'v2'))\n", + "\n", + "def merge_ordered(l, r):\n", + " return pd.merge_ordered(l, r)\n", + "\n", + "df1.groupby('id').cogroup(df2.groupby('id')).applyInPandas(\n", + " merge_ordered, schema='time int, id int, v1 double, v2 string').show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Getting Data In/Out\n", + "\n", + "CSV is straightforward and easy to use. Parquet and ORC are efficient and compact file formats to read and write faster.\n", + "\n", + "There are many other data sources available in PySpark such as JDBC, text, binaryFile, Avro, etc. See also the latest [Spark SQL, DataFrames and Datasets Guide](https://spark.apache.org/docs/latest/sql-programming-guide.html) in Apache Spark documentation." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### CSV" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+------+---+---+\n", + "|color| fruit| v1| v2|\n", + "+-----+------+---+---+\n", + "| red|banana| 1| 10|\n", + "| blue|banana| 2| 20|\n", + "| red|carrot| 3| 30|\n", + "| blue| grape| 4| 40|\n", + "| red|carrot| 5| 50|\n", + "|black|carrot| 6| 60|\n", + "| red|banana| 7| 70|\n", + "| red| grape| 8| 80|\n", + "+-----+------+---+---+\n", + "\n" + ] + } + ], + "source": [ + "df.write.csv('foo.csv', header=True)\n", + "spark.read.csv('foo.csv', header=True).show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parquet" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+------+---+---+\n", + "|color| fruit| v1| v2|\n", + "+-----+------+---+---+\n", + "| red|banana| 1| 10|\n", + "| blue|banana| 2| 20|\n", + "| red|carrot| 3| 30|\n", + "| blue| grape| 4| 40|\n", + "| red|carrot| 5| 50|\n", + "|black|carrot| 6| 60|\n", + "| red|banana| 7| 70|\n", + "| red| grape| 8| 80|\n", + "+-----+------+---+---+\n", + "\n" + ] + } + ], + "source": [ + "df.write.parquet('bar.parquet')\n", + "spark.read.parquet('bar.parquet').show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### ORC" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----+------+---+---+\n", + "|color| fruit| v1| v2|\n", + "+-----+------+---+---+\n", + "| red|banana| 1| 10|\n", + "| blue|banana| 2| 20|\n", + "| red|carrot| 3| 30|\n", + "| blue| grape| 4| 40|\n", + "| red|carrot| 5| 50|\n", + "|black|carrot| 6| 60|\n", + "| red|banana| 7| 70|\n", + "| red| grape| 8| 80|\n", + "+-----+------+---+---+\n", + "\n" + ] + } + ], + "source": [ + "df.write.orc('zoo.orc')\n", + "spark.read.orc('zoo.orc').show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Working with SQL\n", + "\n", + "DataFrame and Spark SQL share the same execution engine so they can be interchangeably used seamlessly. For example, you can register the DataFrame as a table and run a SQL easily as below:" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+--------+\n", + "|count(1)|\n", + "+--------+\n", + "| 8|\n", + "+--------+\n", + "\n" + ] + } + ], + "source": [ + "df.createOrReplaceTempView(\"tableA\")\n", + "spark.sql(\"SELECT count(*) from tableA\").show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In addition, UDFs can be registered and invoked in SQL out of the box:" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----------+\n", + "|add_one(v1)|\n", + "+-----------+\n", + "| 2|\n", + "| 3|\n", + "| 4|\n", + "| 5|\n", + "| 6|\n", + "| 7|\n", + "| 8|\n", + "| 9|\n", + "+-----------+\n", + "\n" + ] + } + ], + "source": [ + "@pandas_udf(\"integer\")\n", + "def add_one(s: pd.Series) -> pd.Series:\n", + " return s + 1\n", + "\n", + "spark.udf.register(\"add_one\", add_one)\n", + "spark.sql(\"SELECT add_one(v1) FROM tableA\").show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These SQL expressions can directly be mixed and used as PySpark columns." + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+-----------+\n", + "|add_one(v1)|\n", + "+-----------+\n", + "| 2|\n", + "| 3|\n", + "| 4|\n", + "| 5|\n", + "| 6|\n", + "| 7|\n", + "| 8|\n", + "| 9|\n", + "+-----------+\n", + "\n", + "+--------------+\n", + "|(count(1) > 0)|\n", + "+--------------+\n", + "| true|\n", + "+--------------+\n", + "\n" + ] + } + ], + "source": [ + "from pyspark.sql.functions import expr\n", + "\n", + "df.selectExpr('add_one(v1)').show()\n", + "df.select(expr('count(*)') > 0).show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.10" + }, + "name": "quickstart", + "notebookId": 1927513300154480 + }, + "nbformat": 4, + "nbformat_minor": 1 +} diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_ps.ipynb b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_ps.ipynb new file mode 100644 index 0000000..dc47bdf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/quickstart_ps.ipynb @@ -0,0 +1,14489 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Quickstart: Pandas API on Spark\n", + "\n", + "This is a short introduction to pandas API on Spark, geared mainly for new users. This notebook shows you some key differences between pandas and pandas API on Spark. You can run this examples by yourself in 'Live Notebook: pandas API on Spark' at [the quickstart page](https://spark.apache.org/docs/latest/api/python/getting_started/index.html).\n", + "\n", + "Customarily, we import pandas API on Spark as follows:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import numpy as np\n", + "import pyspark.pandas as ps\n", + "from pyspark.sql import SparkSession" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Object Creation\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Creating a pandas-on-Spark Series by passing a list of values, letting pandas API on Spark create a default integer index:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "s = ps.Series([1, 3, 5, np.nan, 6, 8])" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0 1.0\n", + "1 3.0\n", + "2 5.0\n", + "3 NaN\n", + "4 6.0\n", + "5 8.0\n", + "dtype: float64" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "s" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Creating a pandas-on-Spark DataFrame by passing a dict of objects that can be converted to series-like." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = ps.DataFrame(\n", + " {'a': [1, 2, 3, 4, 5, 6],\n", + " 'b': [100, 200, 300, 400, 500, 600],\n", + " 'c': [\"one\", \"two\", \"three\", \"four\", \"five\", \"six\"]},\n", + " index=[10, 20, 30, 40, 50, 60])" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    abc
    101100one
    202200two
    303300three
    404400four
    505500five
    606600six
    \n", + "
    " + ], + "text/plain": [ + " a b c\n", + "10 1 100 one\n", + "20 2 200 two\n", + "30 3 300 three\n", + "40 4 400 four\n", + "50 5 500 five\n", + "60 6 600 six" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Creating a pandas DataFrame by passing a numpy array, with a datetime index and labeled columns:" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "dates = pd.date_range('20130101', periods=6)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "DatetimeIndex(['2013-01-01', '2013-01-02', '2013-01-03', '2013-01-04',\n", + " '2013-01-05', '2013-01-06'],\n", + " dtype='datetime64[ns]', freq='D')" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "dates" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "pdf = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    2013-01-010.912558-0.795645-0.2891150.187606
    2013-01-02-0.059703-1.2338970.316625-1.226828
    2013-01-030.332871-1.262010-0.434844-0.579920
    2013-01-040.924016-1.022019-0.405249-1.036021
    2013-01-05-0.772209-1.2280990.0689010.896679
    2013-01-061.485582-0.709306-0.202637-0.248766
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "2013-01-01 0.912558 -0.795645 -0.289115 0.187606\n", + "2013-01-02 -0.059703 -1.233897 0.316625 -1.226828\n", + "2013-01-03 0.332871 -1.262010 -0.434844 -0.579920\n", + "2013-01-04 0.924016 -1.022019 -0.405249 -1.036021\n", + "2013-01-05 -0.772209 -1.228099 0.068901 0.896679\n", + "2013-01-06 1.485582 -0.709306 -0.202637 -0.248766" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pdf" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, this pandas DataFrame can be converted to a pandas-on-Spark DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = ps.from_pandas(pdf)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "pyspark.pandas.frame.DataFrame" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(psdf)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It looks and behaves the same as a pandas DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "scrolled": true + }, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    2013-01-010.912558-0.795645-0.2891150.187606
    2013-01-02-0.059703-1.2338970.316625-1.226828
    2013-01-030.332871-1.262010-0.434844-0.579920
    2013-01-040.924016-1.022019-0.405249-1.036021
    2013-01-05-0.772209-1.2280990.0689010.896679
    2013-01-061.485582-0.709306-0.202637-0.248766
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "2013-01-01 0.912558 -0.795645 -0.289115 0.187606\n", + "2013-01-02 -0.059703 -1.233897 0.316625 -1.226828\n", + "2013-01-03 0.332871 -1.262010 -0.434844 -0.579920\n", + "2013-01-04 0.924016 -1.022019 -0.405249 -1.036021\n", + "2013-01-05 -0.772209 -1.228099 0.068901 0.896679\n", + "2013-01-06 1.485582 -0.709306 -0.202637 -0.248766" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Also, it is possible to create a pandas-on-Spark DataFrame from Spark DataFrame easily.\n", + "\n", + "Creating a Spark DataFrame from pandas DataFrame" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "spark = SparkSession.builder.getOrCreate()" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "sdf = spark.createDataFrame(pdf)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+--------------------+-------------------+--------------------+--------------------+\n", + "| A| B| C| D|\n", + "+--------------------+-------------------+--------------------+--------------------+\n", + "| 0.91255803205208|-0.7956452608556638|-0.28911463069772175| 0.18760566615081622|\n", + "|-0.05970271470242...| -1.233896949308984| 0.3166246451758431| -1.2268284000402265|\n", + "| 0.33287106947536615|-1.2620100816441786| -0.4348444277082644| -0.5799199651437185|\n", + "| 0.9240158461589916|-1.0220190956326003| -0.4052488880650239| -1.0360212104348547|\n", + "| -0.7722090016558953|-1.2280986385313222| 0.0689011451939635| 0.8966790729426755|\n", + "| 1.4855822995785612|-0.7093056426018517| -0.2026366848847041|-0.24876619876451092|\n", + "+--------------------+-------------------+--------------------+--------------------+\n", + "\n" + ] + } + ], + "source": [ + "sdf.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Creating pandas-on-Spark DataFrame from Spark DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = sdf.pandas_api()" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    00.912558-0.795645-0.2891150.187606
    1-0.059703-1.2338970.316625-1.226828
    20.332871-1.262010-0.434844-0.579920
    30.924016-1.022019-0.405249-1.036021
    4-0.772209-1.2280990.0689010.896679
    51.485582-0.709306-0.202637-0.248766
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 0.912558 -0.795645 -0.289115 0.187606\n", + "1 -0.059703 -1.233897 0.316625 -1.226828\n", + "2 0.332871 -1.262010 -0.434844 -0.579920\n", + "3 0.924016 -1.022019 -0.405249 -1.036021\n", + "4 -0.772209 -1.228099 0.068901 0.896679\n", + "5 1.485582 -0.709306 -0.202637 -0.248766" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Having specific [dtypes](http://pandas.pydata.org/pandas-docs/stable/basics.html#basics-dtypes) . Types that are common to both Spark and pandas are currently supported." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "A float64\n", + "B float64\n", + "C float64\n", + "D float64\n", + "dtype: object" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.dtypes" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is how to show top rows from the frame below.\n", + "\n", + "Note that the data in a Spark dataframe does not preserve the natural order by default. The natural order can be preserved by setting `compute.ordered_head` option but it causes a performance overhead with sorting internally." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    00.912558-0.795645-0.2891150.187606
    1-0.059703-1.2338970.316625-1.226828
    20.332871-1.262010-0.434844-0.579920
    30.924016-1.022019-0.405249-1.036021
    4-0.772209-1.2280990.0689010.896679
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 0.912558 -0.795645 -0.289115 0.187606\n", + "1 -0.059703 -1.233897 0.316625 -1.226828\n", + "2 0.332871 -1.262010 -0.434844 -0.579920\n", + "3 0.924016 -1.022019 -0.405249 -1.036021\n", + "4 -0.772209 -1.228099 0.068901 0.896679" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Displaying the index, columns, and the underlying numpy data." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Int64Index([0, 1, 2, 3, 4, 5], dtype='int64')" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.index" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Index(['A', 'B', 'C', 'D'], dtype='object')" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.columns" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([[ 0.91255803, -0.79564526, -0.28911463, 0.18760567],\n", + " [-0.05970271, -1.23389695, 0.31662465, -1.2268284 ],\n", + " [ 0.33287107, -1.26201008, -0.43484443, -0.57991997],\n", + " [ 0.92401585, -1.0220191 , -0.40524889, -1.03602121],\n", + " [-0.772209 , -1.22809864, 0.06890115, 0.89667907],\n", + " [ 1.4855823 , -0.70930564, -0.20263668, -0.2487662 ]])" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.to_numpy()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Showing a quick statistic summary of your data" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    count6.0000006.0000006.0000006.000000
    mean0.470519-1.041829-0.157720-0.334542
    std0.8094280.2415110.2945200.793014
    min-0.772209-1.262010-0.434844-1.226828
    25%-0.059703-1.233897-0.405249-1.036021
    50%0.332871-1.228099-0.289115-0.579920
    75%0.924016-0.7956450.0689010.187606
    max1.485582-0.7093060.3166250.896679
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "count 6.000000 6.000000 6.000000 6.000000\n", + "mean 0.470519 -1.041829 -0.157720 -0.334542\n", + "std 0.809428 0.241511 0.294520 0.793014\n", + "min -0.772209 -1.262010 -0.434844 -1.226828\n", + "25% -0.059703 -1.233897 -0.405249 -1.036021\n", + "50% 0.332871 -1.228099 -0.289115 -0.579920\n", + "75% 0.924016 -0.795645 0.068901 0.187606\n", + "max 1.485582 -0.709306 0.316625 0.896679" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.describe()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Transposing your data" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    012345
    A0.912558-0.0597030.3328710.924016-0.7722091.485582
    B-0.795645-1.233897-1.262010-1.022019-1.228099-0.709306
    C-0.2891150.316625-0.434844-0.4052490.068901-0.202637
    D0.187606-1.226828-0.579920-1.0360210.896679-0.248766
    \n", + "
    " + ], + "text/plain": [ + " 0 1 2 3 4 5\n", + "A 0.912558 -0.059703 0.332871 0.924016 -0.772209 1.485582\n", + "B -0.795645 -1.233897 -1.262010 -1.022019 -1.228099 -0.709306\n", + "C -0.289115 0.316625 -0.434844 -0.405249 0.068901 -0.202637\n", + "D 0.187606 -1.226828 -0.579920 -1.036021 0.896679 -0.248766" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.T" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Sorting by its index" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    51.485582-0.709306-0.202637-0.248766
    4-0.772209-1.2280990.0689010.896679
    30.924016-1.022019-0.405249-1.036021
    20.332871-1.262010-0.434844-0.579920
    1-0.059703-1.2338970.316625-1.226828
    00.912558-0.795645-0.2891150.187606
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "5 1.485582 -0.709306 -0.202637 -0.248766\n", + "4 -0.772209 -1.228099 0.068901 0.896679\n", + "3 0.924016 -1.022019 -0.405249 -1.036021\n", + "2 0.332871 -1.262010 -0.434844 -0.579920\n", + "1 -0.059703 -1.233897 0.316625 -1.226828\n", + "0 0.912558 -0.795645 -0.289115 0.187606" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.sort_index(ascending=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Sorting by value" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    20.332871-1.262010-0.434844-0.579920
    1-0.059703-1.2338970.316625-1.226828
    4-0.772209-1.2280990.0689010.896679
    30.924016-1.022019-0.405249-1.036021
    00.912558-0.795645-0.2891150.187606
    51.485582-0.709306-0.202637-0.248766
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "2 0.332871 -1.262010 -0.434844 -0.579920\n", + "1 -0.059703 -1.233897 0.316625 -1.226828\n", + "4 -0.772209 -1.228099 0.068901 0.896679\n", + "3 0.924016 -1.022019 -0.405249 -1.036021\n", + "0 0.912558 -0.795645 -0.289115 0.187606\n", + "5 1.485582 -0.709306 -0.202637 -0.248766" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.sort_values(by='B')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Missing Data\n", + "Pandas API on Spark primarily uses the value `np.nan` to represent missing data. It is by default not included in computations. \n" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": {}, + "outputs": [], + "source": [ + "pdf1 = pdf.reindex(index=dates[0:4], columns=list(pdf.columns) + ['E'])" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [], + "source": [ + "pdf1.loc[dates[0]:dates[1], 'E'] = 1" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [], + "source": [ + "psdf1 = ps.from_pandas(pdf1)" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCDE
    2013-01-010.912558-0.795645-0.2891150.1876061.0
    2013-01-02-0.059703-1.2338970.316625-1.2268281.0
    2013-01-030.332871-1.262010-0.434844-0.579920NaN
    2013-01-040.924016-1.022019-0.405249-1.036021NaN
    \n", + "
    " + ], + "text/plain": [ + " A B C D E\n", + "2013-01-01 0.912558 -0.795645 -0.289115 0.187606 1.0\n", + "2013-01-02 -0.059703 -1.233897 0.316625 -1.226828 1.0\n", + "2013-01-03 0.332871 -1.262010 -0.434844 -0.579920 NaN\n", + "2013-01-04 0.924016 -1.022019 -0.405249 -1.036021 NaN" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf1" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To drop any rows that have missing data." + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCDE
    2013-01-010.912558-0.795645-0.2891150.1876061.0
    2013-01-02-0.059703-1.2338970.316625-1.2268281.0
    \n", + "
    " + ], + "text/plain": [ + " A B C D E\n", + "2013-01-01 0.912558 -0.795645 -0.289115 0.187606 1.0\n", + "2013-01-02 -0.059703 -1.233897 0.316625 -1.226828 1.0" + ] + }, + "execution_count": 31, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf1.dropna(how='any')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Filling missing data." + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCDE
    2013-01-010.912558-0.795645-0.2891150.1876061.0
    2013-01-02-0.059703-1.2338970.316625-1.2268281.0
    2013-01-030.332871-1.262010-0.434844-0.5799205.0
    2013-01-040.924016-1.022019-0.405249-1.0360215.0
    \n", + "
    " + ], + "text/plain": [ + " A B C D E\n", + "2013-01-01 0.912558 -0.795645 -0.289115 0.187606 1.0\n", + "2013-01-02 -0.059703 -1.233897 0.316625 -1.226828 1.0\n", + "2013-01-03 0.332871 -1.262010 -0.434844 -0.579920 5.0\n", + "2013-01-04 0.924016 -1.022019 -0.405249 -1.036021 5.0" + ] + }, + "execution_count": 32, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf1.fillna(value=5)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Operations" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Stats\n", + "Performing a descriptive statistic:" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "A 0.470519\n", + "B -1.041829\n", + "C -0.157720\n", + "D -0.334542\n", + "dtype: float64" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.mean()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Spark Configurations\n", + "\n", + "Various configurations in PySpark could be applied internally in pandas API on Spark.\n", + "For example, you can enable Arrow optimization to hugely speed up internal pandas conversion. See also PySpark Usage Guide for Pandas with Apache Arrow in PySpark documentation." + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "metadata": {}, + "outputs": [], + "source": [ + "prev = spark.conf.get(\"spark.sql.execution.arrow.pyspark.enabled\") # Keep its default value.\n", + "ps.set_option(\"compute.default_index_type\", \"distributed\") # Use default index prevent overhead.\n", + "import warnings\n", + "warnings.filterwarnings(\"ignore\") # Ignore warnings coming from Arrow optimizations." + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "900 ms ± 186 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" + ] + } + ], + "source": [ + "spark.conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", True)\n", + "%timeit ps.range(300000).to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3.08 s ± 227 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" + ] + } + ], + "source": [ + "spark.conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", False)\n", + "%timeit ps.range(300000).to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": {}, + "outputs": [], + "source": [ + "ps.reset_option(\"compute.default_index_type\")\n", + "spark.conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", prev) # Set its default value back." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Grouping\n", + "By “group by” we are referring to a process involving one or more of the following steps:\n", + "\n", + "- Splitting the data into groups based on some criteria\n", + "- Applying a function to each group independently\n", + "- Combining the results into a data structure" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = ps.DataFrame({'A': ['foo', 'bar', 'foo', 'bar',\n", + " 'foo', 'bar', 'foo', 'foo'],\n", + " 'B': ['one', 'one', 'two', 'three',\n", + " 'two', 'two', 'one', 'three'],\n", + " 'C': np.random.randn(8),\n", + " 'D': np.random.randn(8)})" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    0fooone1.039632-0.571950
    1barone0.9720891.085353
    2footwo-1.931621-2.579164
    3barthree-0.654371-0.340704
    4footwo-0.1570800.893736
    5bartwo0.8827950.024978
    6fooone-0.1493840.201667
    7foothree-1.3551360.693883
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 foo one 1.039632 -0.571950\n", + "1 bar one 0.972089 1.085353\n", + "2 foo two -1.931621 -2.579164\n", + "3 bar three -0.654371 -0.340704\n", + "4 foo two -0.157080 0.893736\n", + "5 bar two 0.882795 0.024978\n", + "6 foo one -0.149384 0.201667\n", + "7 foo three -1.355136 0.693883" + ] + }, + "execution_count": 39, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Grouping and then applying the [sum()](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.groupby.GroupBy.sum.html) function to the resulting groups." + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    CD
    A
    bar1.2005130.769627
    foo-2.553589-1.361828
    \n", + "
    " + ], + "text/plain": [ + " C D\n", + "A \n", + "bar 1.200513 0.769627\n", + "foo -2.553589 -1.361828" + ] + }, + "execution_count": 40, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.groupby('A').sum()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Grouping by multiple columns forms a hierarchical index, and again we can apply the sum function." + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    CD
    AB
    fooone0.890248-0.370283
    two-2.088701-1.685428
    barthree-0.654371-0.340704
    foothree-1.3551360.693883
    bartwo0.8827950.024978
    one0.9720891.085353
    \n", + "
    " + ], + "text/plain": [ + " C D\n", + "A B \n", + "foo one 0.890248 -0.370283\n", + " two -2.088701 -1.685428\n", + "bar three -0.654371 -0.340704\n", + "foo three -1.355136 0.693883\n", + "bar two 0.882795 0.024978\n", + " one 0.972089 1.085353" + ] + }, + "execution_count": 41, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.groupby(['A', 'B']).sum()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Plotting" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "metadata": {}, + "outputs": [], + "source": [ + "pser = pd.Series(np.random.randn(1000),\n", + " index=pd.date_range('1/1/2000', periods=1000))" + ] + }, + { + "cell_type": "code", + "execution_count": 43, + "metadata": {}, + "outputs": [], + "source": [ + "psser = ps.Series(pser)" + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": {}, + "outputs": [], + "source": [ + "psser = psser.cummax()" + ] + }, + { + "cell_type": "code", + "execution_count": 45, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + " \n", + " " + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.plotly.v1+json": { + "config": { + "plotlyServerURL": "https://plot.ly" + }, + "data": [ + { + "hovertemplate": "variable=0
    index=%{x}
    value=%{y}", + "legendgroup": "0", + "line": { + "color": "#636efa", + "dash": "solid" + }, + "mode": "lines", + "name": "0", + "orientation": "v", + "showlegend": true, + "type": "scatter", + "x": [ + "2000-01-01T00:00:00", + "2000-01-02T00:00:00", + "2000-01-03T00:00:00", + "2000-01-04T00:00:00", + "2000-01-05T00:00:00", + "2000-01-06T00:00:00", + "2000-01-07T00:00:00", + "2000-01-08T00:00:00", + "2000-01-09T00:00:00", + "2000-01-10T00:00:00", + "2000-01-11T00:00:00", + "2000-01-12T00:00:00", + "2000-01-13T00:00:00", + "2000-01-14T00:00:00", + "2000-01-15T00:00:00", + "2000-01-16T00:00:00", + "2000-01-17T00:00:00", + "2000-01-18T00:00:00", + "2000-01-19T00:00:00", + "2000-01-20T00:00:00", + "2000-01-21T00:00:00", + "2000-01-22T00:00:00", + "2000-01-23T00:00:00", + "2000-01-24T00:00:00", + "2000-01-25T00:00:00", + "2000-01-26T00:00:00", + "2000-01-27T00:00:00", + "2000-01-28T00:00:00", + "2000-01-29T00:00:00", + "2000-01-30T00:00:00", + "2000-01-31T00:00:00", + "2000-02-01T00:00:00", + "2000-02-02T00:00:00", + "2000-02-03T00:00:00", + "2000-02-04T00:00:00", + "2000-02-05T00:00:00", + "2000-02-06T00:00:00", + "2000-02-07T00:00:00", + "2000-02-08T00:00:00", + "2000-02-09T00:00:00", + "2000-02-10T00:00:00", + "2000-02-11T00:00:00", + "2000-02-12T00:00:00", + "2000-02-13T00:00:00", + "2000-02-14T00:00:00", + "2000-02-15T00:00:00", + "2000-02-16T00:00:00", + "2000-02-17T00:00:00", + "2000-02-18T00:00:00", + "2000-02-19T00:00:00", + "2000-02-20T00:00:00", + "2000-02-21T00:00:00", + "2000-02-22T00:00:00", + "2000-02-23T00:00:00", + "2000-02-24T00:00:00", + "2000-02-25T00:00:00", + "2000-02-26T00:00:00", + "2000-02-27T00:00:00", + "2000-02-28T00:00:00", + "2000-02-29T00:00:00", + "2000-03-01T00:00:00", + "2000-03-02T00:00:00", + "2000-03-03T00:00:00", + "2000-03-04T00:00:00", + "2000-03-05T00:00:00", + "2000-03-06T00:00:00", + "2000-03-07T00:00:00", + "2000-03-08T00:00:00", + "2000-03-09T00:00:00", + "2000-03-10T00:00:00", + "2000-03-11T00:00:00", + "2000-03-12T00:00:00", + "2000-03-13T00:00:00", + "2000-03-14T00:00:00", + "2000-03-15T00:00:00", + "2000-03-16T00:00:00", + "2000-03-17T00:00:00", + "2000-03-18T00:00:00", + "2000-03-19T00:00:00", + "2000-03-20T00:00:00", + "2000-03-21T00:00:00", + "2000-03-22T00:00:00", + "2000-03-23T00:00:00", + "2000-03-24T00:00:00", + "2000-03-25T00:00:00", + "2000-03-26T00:00:00", + "2000-03-27T00:00:00", + "2000-03-28T00:00:00", + "2000-03-29T00:00:00", + "2000-03-30T00:00:00", + "2000-03-31T00:00:00", + "2000-04-01T00:00:00", + "2000-04-02T00:00:00", + "2000-04-03T00:00:00", + "2000-04-04T00:00:00", + "2000-04-05T00:00:00", + "2000-04-06T00:00:00", + "2000-04-07T00:00:00", + "2000-04-08T00:00:00", + "2000-04-09T00:00:00", + "2000-04-10T00:00:00", + "2000-04-11T00:00:00", + "2000-04-12T00:00:00", + "2000-04-13T00:00:00", + "2000-04-14T00:00:00", + "2000-04-15T00:00:00", + "2000-04-16T00:00:00", + "2000-04-17T00:00:00", + "2000-04-18T00:00:00", + "2000-04-19T00:00:00", + "2000-04-20T00:00:00", + "2000-04-21T00:00:00", + "2000-04-22T00:00:00", + "2000-04-23T00:00:00", + "2000-04-24T00:00:00", + "2000-04-25T00:00:00", + "2000-04-26T00:00:00", + "2000-04-27T00:00:00", + "2000-04-28T00:00:00", + "2000-04-29T00:00:00", + "2000-04-30T00:00:00", + "2000-05-01T00:00:00", + "2000-05-02T00:00:00", + "2000-05-03T00:00:00", + "2000-05-04T00:00:00", + "2000-05-05T00:00:00", + "2000-05-06T00:00:00", + "2000-05-07T00:00:00", + "2000-05-08T00:00:00", + "2000-05-09T00:00:00", + "2000-05-10T00:00:00", + "2000-05-11T00:00:00", + "2000-05-12T00:00:00", + "2000-05-13T00:00:00", + "2000-05-14T00:00:00", + "2000-05-15T00:00:00", + "2000-05-16T00:00:00", + "2000-05-17T00:00:00", + "2000-05-18T00:00:00", + "2000-05-19T00:00:00", + "2000-05-20T00:00:00", + "2000-05-21T00:00:00", + "2000-05-22T00:00:00", + "2000-05-23T00:00:00", + "2000-05-24T00:00:00", + "2000-05-25T00:00:00", + "2000-05-26T00:00:00", + "2000-05-27T00:00:00", + "2000-05-28T00:00:00", + "2000-05-29T00:00:00", + "2000-05-30T00:00:00", + "2000-05-31T00:00:00", + "2000-06-01T00:00:00", + "2000-06-02T00:00:00", + "2000-06-03T00:00:00", + "2000-06-04T00:00:00", + "2000-06-05T00:00:00", + "2000-06-06T00:00:00", + "2000-06-07T00:00:00", + "2000-06-08T00:00:00", + "2000-06-09T00:00:00", + "2000-06-10T00:00:00", + "2000-06-11T00:00:00", + "2000-06-12T00:00:00", + "2000-06-13T00:00:00", + "2000-06-14T00:00:00", + "2000-06-15T00:00:00", + "2000-06-16T00:00:00", + "2000-06-17T00:00:00", + "2000-06-18T00:00:00", + "2000-06-19T00:00:00", + "2000-06-20T00:00:00", + "2000-06-21T00:00:00", + "2000-06-22T00:00:00", + "2000-06-23T00:00:00", + "2000-06-24T00:00:00", + "2000-06-25T00:00:00", + "2000-06-26T00:00:00", + "2000-06-27T00:00:00", + "2000-06-28T00:00:00", + "2000-06-29T00:00:00", + "2000-06-30T00:00:00", + "2000-07-01T00:00:00", + "2000-07-02T00:00:00", + "2000-07-03T00:00:00", + "2000-07-04T00:00:00", + "2000-07-05T00:00:00", + "2000-07-06T00:00:00", + "2000-07-07T00:00:00", + "2000-07-08T00:00:00", + "2000-07-09T00:00:00", + "2000-07-10T00:00:00", + "2000-07-11T00:00:00", + "2000-07-12T00:00:00", + "2000-07-13T00:00:00", + "2000-07-14T00:00:00", + "2000-07-15T00:00:00", + "2000-07-16T00:00:00", + "2000-07-17T00:00:00", + "2000-07-18T00:00:00", + "2000-07-19T00:00:00", + "2000-07-20T00:00:00", + "2000-07-21T00:00:00", + "2000-07-22T00:00:00", + "2000-07-23T00:00:00", + "2000-07-24T00:00:00", + "2000-07-25T00:00:00", + "2000-07-26T00:00:00", + "2000-07-27T00:00:00", + "2000-07-28T00:00:00", + "2000-07-29T00:00:00", + "2000-07-30T00:00:00", + "2000-07-31T00:00:00", + "2000-08-01T00:00:00", + "2000-08-02T00:00:00", + "2000-08-03T00:00:00", + "2000-08-04T00:00:00", + "2000-08-05T00:00:00", + "2000-08-06T00:00:00", + "2000-08-07T00:00:00", + "2000-08-08T00:00:00", + "2000-08-09T00:00:00", + "2000-08-10T00:00:00", + "2000-08-11T00:00:00", + "2000-08-12T00:00:00", + "2000-08-13T00:00:00", + "2000-08-14T00:00:00", + "2000-08-15T00:00:00", + "2000-08-16T00:00:00", + "2000-08-17T00:00:00", + "2000-08-18T00:00:00", + "2000-08-19T00:00:00", + "2000-08-20T00:00:00", + "2000-08-21T00:00:00", + "2000-08-22T00:00:00", + "2000-08-23T00:00:00", + "2000-08-24T00:00:00", + "2000-08-25T00:00:00", + "2000-08-26T00:00:00", + "2000-08-27T00:00:00", + "2000-08-28T00:00:00", + "2000-08-29T00:00:00", + "2000-08-30T00:00:00", + "2000-08-31T00:00:00", + "2000-09-01T00:00:00", + "2000-09-02T00:00:00", + "2000-09-03T00:00:00", + "2000-09-04T00:00:00", + "2000-09-05T00:00:00", + "2000-09-06T00:00:00", + "2000-09-07T00:00:00", + "2000-09-08T00:00:00", + "2000-09-09T00:00:00", + "2000-09-10T00:00:00", + "2000-09-11T00:00:00", + "2000-09-12T00:00:00", + "2000-09-13T00:00:00", + "2000-09-14T00:00:00", + "2000-09-15T00:00:00", + "2000-09-16T00:00:00", + "2000-09-17T00:00:00", + "2000-09-18T00:00:00", + "2000-09-19T00:00:00", + "2000-09-20T00:00:00", + "2000-09-21T00:00:00", + "2000-09-22T00:00:00", + "2000-09-23T00:00:00", + "2000-09-24T00:00:00", + "2000-09-25T00:00:00", + "2000-09-26T00:00:00", + "2000-09-27T00:00:00", + "2000-09-28T00:00:00", + "2000-09-29T00:00:00", + "2000-09-30T00:00:00", + "2000-10-01T00:00:00", + "2000-10-02T00:00:00", + "2000-10-03T00:00:00", + "2000-10-04T00:00:00", + "2000-10-05T00:00:00", + "2000-10-06T00:00:00", + "2000-10-07T00:00:00", + "2000-10-08T00:00:00", + "2000-10-09T00:00:00", + "2000-10-10T00:00:00", + "2000-10-11T00:00:00", + "2000-10-12T00:00:00", + "2000-10-13T00:00:00", + "2000-10-14T00:00:00", + "2000-10-15T00:00:00", + "2000-10-16T00:00:00", + "2000-10-17T00:00:00", + "2000-10-18T00:00:00", + "2000-10-19T00:00:00", + "2000-10-20T00:00:00", + "2000-10-21T00:00:00", + "2000-10-22T00:00:00", + "2000-10-23T00:00:00", + "2000-10-24T00:00:00", + "2000-10-25T00:00:00", + "2000-10-26T00:00:00", + "2000-10-27T00:00:00", + "2000-10-28T00:00:00", + "2000-10-29T00:00:00", + "2000-10-30T00:00:00", + "2000-10-31T00:00:00", + "2000-11-01T00:00:00", + "2000-11-02T00:00:00", + "2000-11-03T00:00:00", + "2000-11-04T00:00:00", + "2000-11-05T00:00:00", + "2000-11-06T00:00:00", + "2000-11-07T00:00:00", + "2000-11-08T00:00:00", + "2000-11-09T00:00:00", + "2000-11-10T00:00:00", + "2000-11-11T00:00:00", + "2000-11-12T00:00:00", + "2000-11-13T00:00:00", + "2000-11-14T00:00:00", + "2000-11-15T00:00:00", + "2000-11-16T00:00:00", + "2000-11-17T00:00:00", + "2000-11-18T00:00:00", + "2000-11-19T00:00:00", + "2000-11-20T00:00:00", + "2000-11-21T00:00:00", + "2000-11-22T00:00:00", + "2000-11-23T00:00:00", + "2000-11-24T00:00:00", + "2000-11-25T00:00:00", + "2000-11-26T00:00:00", + "2000-11-27T00:00:00", + "2000-11-28T00:00:00", + "2000-11-29T00:00:00", + "2000-11-30T00:00:00", + "2000-12-01T00:00:00", + "2000-12-02T00:00:00", + "2000-12-03T00:00:00", + "2000-12-04T00:00:00", + "2000-12-05T00:00:00", + "2000-12-06T00:00:00", + "2000-12-07T00:00:00", + "2000-12-08T00:00:00", + "2000-12-09T00:00:00", + "2000-12-10T00:00:00", + "2000-12-11T00:00:00", + "2000-12-12T00:00:00", + "2000-12-13T00:00:00", + "2000-12-14T00:00:00", + "2000-12-15T00:00:00", + "2000-12-16T00:00:00", + "2000-12-17T00:00:00", + "2000-12-18T00:00:00", + "2000-12-19T00:00:00", + "2000-12-20T00:00:00", + "2000-12-21T00:00:00", + "2000-12-22T00:00:00", + "2000-12-23T00:00:00", + "2000-12-24T00:00:00", + "2000-12-25T00:00:00", + "2000-12-26T00:00:00", + "2000-12-27T00:00:00", + "2000-12-28T00:00:00", + "2000-12-29T00:00:00", + "2000-12-30T00:00:00", + "2000-12-31T00:00:00", + "2001-01-01T00:00:00", + "2001-01-02T00:00:00", + "2001-01-03T00:00:00", + "2001-01-04T00:00:00", + "2001-01-05T00:00:00", + "2001-01-06T00:00:00", + "2001-01-07T00:00:00", + "2001-01-08T00:00:00", + "2001-01-09T00:00:00", + "2001-01-10T00:00:00", + "2001-01-11T00:00:00", + "2001-01-12T00:00:00", + "2001-01-13T00:00:00", + "2001-01-14T00:00:00", + "2001-01-15T00:00:00", + "2001-01-16T00:00:00", + "2001-01-17T00:00:00", + "2001-01-18T00:00:00", + "2001-01-19T00:00:00", + "2001-01-20T00:00:00", + "2001-01-21T00:00:00", + "2001-01-22T00:00:00", + "2001-01-23T00:00:00", + "2001-01-24T00:00:00", + "2001-01-25T00:00:00", + "2001-01-26T00:00:00", + "2001-01-27T00:00:00", + "2001-01-28T00:00:00", + "2001-01-29T00:00:00", + "2001-01-30T00:00:00", + "2001-01-31T00:00:00", + "2001-02-01T00:00:00", + "2001-02-02T00:00:00", + "2001-02-03T00:00:00", + "2001-02-04T00:00:00", + "2001-02-05T00:00:00", + "2001-02-06T00:00:00", + "2001-02-07T00:00:00", + "2001-02-08T00:00:00", + "2001-02-09T00:00:00", + "2001-02-10T00:00:00", + "2001-02-11T00:00:00", + "2001-02-12T00:00:00", + "2001-02-13T00:00:00", + "2001-02-14T00:00:00", + "2001-02-15T00:00:00", + "2001-02-16T00:00:00", + "2001-02-17T00:00:00", + "2001-02-18T00:00:00", + "2001-02-19T00:00:00", + "2001-02-20T00:00:00", + "2001-02-21T00:00:00", + "2001-02-22T00:00:00", + "2001-02-23T00:00:00", + "2001-02-24T00:00:00", + "2001-02-25T00:00:00", + "2001-02-26T00:00:00", + "2001-02-27T00:00:00", + "2001-02-28T00:00:00", + "2001-03-01T00:00:00", + "2001-03-02T00:00:00", + "2001-03-03T00:00:00", + "2001-03-04T00:00:00", + "2001-03-05T00:00:00", + "2001-03-06T00:00:00", + "2001-03-07T00:00:00", + "2001-03-08T00:00:00", + "2001-03-09T00:00:00", + "2001-03-10T00:00:00", + "2001-03-11T00:00:00", + "2001-03-12T00:00:00", + "2001-03-13T00:00:00", + "2001-03-14T00:00:00", + "2001-03-15T00:00:00", + "2001-03-16T00:00:00", + "2001-03-17T00:00:00", + "2001-03-18T00:00:00", + "2001-03-19T00:00:00", + "2001-03-20T00:00:00", + "2001-03-21T00:00:00", + "2001-03-22T00:00:00", + "2001-03-23T00:00:00", + "2001-03-24T00:00:00", + "2001-03-25T00:00:00", + "2001-03-26T00:00:00", + "2001-03-27T00:00:00", + "2001-03-28T00:00:00", + "2001-03-29T00:00:00", + "2001-03-30T00:00:00", + "2001-03-31T00:00:00", + "2001-04-01T00:00:00", + "2001-04-02T00:00:00", + "2001-04-03T00:00:00", + "2001-04-04T00:00:00", + "2001-04-05T00:00:00", + "2001-04-06T00:00:00", + "2001-04-07T00:00:00", + "2001-04-08T00:00:00", + "2001-04-09T00:00:00", + "2001-04-10T00:00:00", + "2001-04-11T00:00:00", + "2001-04-12T00:00:00", + "2001-04-13T00:00:00", + "2001-04-14T00:00:00", + "2001-04-15T00:00:00", + "2001-04-16T00:00:00", + "2001-04-17T00:00:00", + "2001-04-18T00:00:00", + "2001-04-19T00:00:00", + "2001-04-20T00:00:00", + "2001-04-21T00:00:00", + "2001-04-22T00:00:00", + "2001-04-23T00:00:00", + "2001-04-24T00:00:00", + "2001-04-25T00:00:00", + "2001-04-26T00:00:00", + "2001-04-27T00:00:00", + "2001-04-28T00:00:00", + "2001-04-29T00:00:00", + "2001-04-30T00:00:00", + "2001-05-01T00:00:00", + "2001-05-02T00:00:00", + "2001-05-03T00:00:00", + "2001-05-04T00:00:00", + "2001-05-05T00:00:00", + "2001-05-06T00:00:00", + "2001-05-07T00:00:00", + "2001-05-08T00:00:00", + "2001-05-09T00:00:00", + "2001-05-10T00:00:00", + "2001-05-11T00:00:00", + "2001-05-12T00:00:00", + "2001-05-13T00:00:00", + "2001-05-14T00:00:00", + "2001-05-15T00:00:00", + "2001-05-16T00:00:00", + "2001-05-17T00:00:00", + "2001-05-18T00:00:00", + "2001-05-19T00:00:00", + "2001-05-20T00:00:00", + "2001-05-21T00:00:00", + "2001-05-22T00:00:00", + "2001-05-23T00:00:00", + "2001-05-24T00:00:00", + "2001-05-25T00:00:00", + "2001-05-26T00:00:00", + "2001-05-27T00:00:00", + "2001-05-28T00:00:00", + "2001-05-29T00:00:00", + "2001-05-30T00:00:00", + "2001-05-31T00:00:00", + "2001-06-01T00:00:00", + "2001-06-02T00:00:00", + "2001-06-03T00:00:00", + "2001-06-04T00:00:00", + "2001-06-05T00:00:00", + "2001-06-06T00:00:00", + "2001-06-07T00:00:00", + "2001-06-08T00:00:00", + "2001-06-09T00:00:00", + "2001-06-10T00:00:00", + "2001-06-11T00:00:00", + "2001-06-12T00:00:00", + "2001-06-13T00:00:00", + "2001-06-14T00:00:00", + "2001-06-15T00:00:00", + "2001-06-16T00:00:00", + "2001-06-17T00:00:00", + "2001-06-18T00:00:00", + "2001-06-19T00:00:00", + "2001-06-20T00:00:00", + "2001-06-21T00:00:00", + "2001-06-22T00:00:00", + "2001-06-23T00:00:00", + "2001-06-24T00:00:00", + "2001-06-25T00:00:00", + "2001-06-26T00:00:00", + "2001-06-27T00:00:00", + "2001-06-28T00:00:00", + "2001-06-29T00:00:00", + "2001-06-30T00:00:00", + "2001-07-01T00:00:00", + "2001-07-02T00:00:00", + "2001-07-03T00:00:00", + "2001-07-04T00:00:00", + "2001-07-05T00:00:00", + "2001-07-06T00:00:00", + "2001-07-07T00:00:00", + "2001-07-08T00:00:00", + "2001-07-09T00:00:00", + "2001-07-10T00:00:00", + "2001-07-11T00:00:00", + "2001-07-12T00:00:00", + "2001-07-13T00:00:00", + "2001-07-14T00:00:00", + "2001-07-15T00:00:00", + "2001-07-16T00:00:00", + "2001-07-17T00:00:00", + "2001-07-18T00:00:00", + "2001-07-19T00:00:00", + "2001-07-20T00:00:00", + "2001-07-21T00:00:00", + "2001-07-22T00:00:00", + "2001-07-23T00:00:00", + "2001-07-24T00:00:00", + "2001-07-25T00:00:00", + "2001-07-26T00:00:00", + "2001-07-27T00:00:00", + "2001-07-28T00:00:00", + "2001-07-29T00:00:00", + "2001-07-30T00:00:00", + "2001-07-31T00:00:00", + "2001-08-01T00:00:00", + "2001-08-02T00:00:00", + "2001-08-03T00:00:00", + "2001-08-04T00:00:00", + "2001-08-05T00:00:00", + "2001-08-06T00:00:00", + "2001-08-07T00:00:00", + "2001-08-08T00:00:00", + "2001-08-09T00:00:00", + "2001-08-10T00:00:00", + "2001-08-11T00:00:00", + "2001-08-12T00:00:00", + "2001-08-13T00:00:00", + "2001-08-14T00:00:00", + "2001-08-15T00:00:00", + "2001-08-16T00:00:00", + "2001-08-17T00:00:00", + "2001-08-18T00:00:00", + "2001-08-19T00:00:00", + "2001-08-20T00:00:00", + "2001-08-21T00:00:00", + "2001-08-22T00:00:00", + "2001-08-23T00:00:00", + "2001-08-24T00:00:00", + "2001-08-25T00:00:00", + "2001-08-26T00:00:00", + "2001-08-27T00:00:00", + "2001-08-28T00:00:00", + "2001-08-29T00:00:00", + "2001-08-30T00:00:00", + "2001-08-31T00:00:00", + "2001-09-01T00:00:00", + "2001-09-02T00:00:00", + "2001-09-03T00:00:00", + "2001-09-04T00:00:00", + "2001-09-05T00:00:00", + "2001-09-06T00:00:00", + "2001-09-07T00:00:00", + "2001-09-08T00:00:00", + "2001-09-09T00:00:00", + "2001-09-10T00:00:00", + "2001-09-11T00:00:00", + "2001-09-12T00:00:00", + "2001-09-13T00:00:00", + "2001-09-14T00:00:00", + "2001-09-15T00:00:00", + "2001-09-16T00:00:00", + "2001-09-17T00:00:00", + "2001-09-18T00:00:00", + "2001-09-19T00:00:00", + "2001-09-20T00:00:00", + "2001-09-21T00:00:00", + "2001-09-22T00:00:00", + "2001-09-23T00:00:00", + "2001-09-24T00:00:00", + "2001-09-25T00:00:00", + "2001-09-26T00:00:00", + "2001-09-27T00:00:00", + "2001-09-28T00:00:00", + "2001-09-29T00:00:00", + "2001-09-30T00:00:00", + "2001-10-01T00:00:00", + "2001-10-02T00:00:00", + "2001-10-03T00:00:00", + "2001-10-04T00:00:00", + "2001-10-05T00:00:00", + "2001-10-06T00:00:00", + "2001-10-07T00:00:00", + "2001-10-08T00:00:00", + "2001-10-09T00:00:00", + "2001-10-10T00:00:00", + "2001-10-11T00:00:00", + "2001-10-12T00:00:00", + "2001-10-13T00:00:00", + "2001-10-14T00:00:00", + "2001-10-15T00:00:00", + "2001-10-16T00:00:00", + "2001-10-17T00:00:00", + "2001-10-18T00:00:00", + "2001-10-19T00:00:00", + "2001-10-20T00:00:00", + "2001-10-21T00:00:00", + "2001-10-22T00:00:00", + "2001-10-23T00:00:00", + "2001-10-24T00:00:00", + "2001-10-25T00:00:00", + "2001-10-26T00:00:00", + "2001-10-27T00:00:00", + "2001-10-28T00:00:00", + "2001-10-29T00:00:00", + "2001-10-30T00:00:00", + "2001-10-31T00:00:00", + "2001-11-01T00:00:00", + "2001-11-02T00:00:00", + "2001-11-03T00:00:00", + "2001-11-04T00:00:00", + "2001-11-05T00:00:00", + "2001-11-06T00:00:00", + "2001-11-07T00:00:00", + "2001-11-08T00:00:00", + "2001-11-09T00:00:00", + "2001-11-10T00:00:00", + "2001-11-11T00:00:00", + "2001-11-12T00:00:00", + "2001-11-13T00:00:00", + "2001-11-14T00:00:00", + "2001-11-15T00:00:00", + "2001-11-16T00:00:00", + "2001-11-17T00:00:00", + "2001-11-18T00:00:00", + "2001-11-19T00:00:00", + "2001-11-20T00:00:00", + "2001-11-21T00:00:00", + "2001-11-22T00:00:00", + "2001-11-23T00:00:00", + "2001-11-24T00:00:00", + "2001-11-25T00:00:00", + "2001-11-26T00:00:00", + "2001-11-27T00:00:00", + "2001-11-28T00:00:00", + "2001-11-29T00:00:00", + "2001-11-30T00:00:00", + "2001-12-01T00:00:00", + "2001-12-02T00:00:00", + "2001-12-03T00:00:00", + "2001-12-04T00:00:00", + "2001-12-05T00:00:00", + "2001-12-06T00:00:00", + "2001-12-07T00:00:00", + "2001-12-08T00:00:00", + "2001-12-09T00:00:00", + "2001-12-10T00:00:00", + "2001-12-11T00:00:00", + "2001-12-12T00:00:00", + "2001-12-13T00:00:00", + "2001-12-14T00:00:00", + "2001-12-15T00:00:00", + "2001-12-16T00:00:00", + "2001-12-17T00:00:00", + "2001-12-18T00:00:00", + "2001-12-19T00:00:00", + "2001-12-20T00:00:00", + "2001-12-21T00:00:00", + "2001-12-22T00:00:00", + "2001-12-23T00:00:00", + "2001-12-24T00:00:00", + "2001-12-25T00:00:00", + "2001-12-26T00:00:00", + "2001-12-27T00:00:00", + "2001-12-28T00:00:00", + "2001-12-29T00:00:00", + "2001-12-30T00:00:00", + "2001-12-31T00:00:00", + "2002-01-01T00:00:00", + "2002-01-02T00:00:00", + "2002-01-03T00:00:00", + "2002-01-04T00:00:00", + "2002-01-05T00:00:00", + "2002-01-06T00:00:00", + "2002-01-07T00:00:00", + "2002-01-08T00:00:00", + "2002-01-09T00:00:00", + "2002-01-10T00:00:00", + "2002-01-11T00:00:00", + "2002-01-12T00:00:00", + "2002-01-13T00:00:00", + "2002-01-14T00:00:00", + "2002-01-15T00:00:00", + "2002-01-16T00:00:00", + "2002-01-17T00:00:00", + "2002-01-18T00:00:00", + "2002-01-19T00:00:00", + "2002-01-20T00:00:00", + "2002-01-21T00:00:00", + "2002-01-22T00:00:00", + "2002-01-23T00:00:00", + "2002-01-24T00:00:00", + "2002-01-25T00:00:00", + "2002-01-26T00:00:00", + "2002-01-27T00:00:00", + "2002-01-28T00:00:00", + "2002-01-29T00:00:00", + "2002-01-30T00:00:00", + "2002-01-31T00:00:00", + "2002-02-01T00:00:00", + "2002-02-02T00:00:00", + "2002-02-03T00:00:00", + "2002-02-04T00:00:00", + "2002-02-05T00:00:00", + "2002-02-06T00:00:00", + "2002-02-07T00:00:00", + "2002-02-08T00:00:00", + "2002-02-09T00:00:00", + "2002-02-10T00:00:00", + "2002-02-11T00:00:00", + "2002-02-12T00:00:00", + "2002-02-13T00:00:00", + "2002-02-14T00:00:00", + "2002-02-15T00:00:00", + "2002-02-16T00:00:00", + "2002-02-17T00:00:00", + "2002-02-18T00:00:00", + "2002-02-19T00:00:00", + "2002-02-20T00:00:00", + "2002-02-21T00:00:00", + "2002-02-22T00:00:00", + "2002-02-23T00:00:00", + "2002-02-24T00:00:00", + "2002-02-25T00:00:00", + "2002-02-26T00:00:00", + "2002-02-27T00:00:00", + "2002-02-28T00:00:00", + "2002-03-01T00:00:00", + "2002-03-02T00:00:00", + "2002-03-03T00:00:00", + "2002-03-04T00:00:00", + "2002-03-05T00:00:00", + "2002-03-06T00:00:00", + "2002-03-07T00:00:00", + "2002-03-08T00:00:00", + "2002-03-09T00:00:00", + "2002-03-10T00:00:00", + "2002-03-11T00:00:00", + "2002-03-12T00:00:00", + "2002-03-13T00:00:00", + "2002-03-14T00:00:00", + "2002-03-15T00:00:00", + "2002-03-16T00:00:00", + "2002-03-17T00:00:00", + "2002-03-18T00:00:00", + "2002-03-19T00:00:00", + "2002-03-20T00:00:00", + "2002-03-21T00:00:00", + "2002-03-22T00:00:00", + "2002-03-23T00:00:00", + "2002-03-24T00:00:00", + "2002-03-25T00:00:00", + "2002-03-26T00:00:00", + "2002-03-27T00:00:00", + "2002-03-28T00:00:00", + "2002-03-29T00:00:00", + "2002-03-30T00:00:00", + "2002-03-31T00:00:00", + "2002-04-01T00:00:00", + "2002-04-02T00:00:00", + "2002-04-03T00:00:00", + "2002-04-04T00:00:00", + "2002-04-05T00:00:00", + "2002-04-06T00:00:00", + "2002-04-07T00:00:00", + "2002-04-08T00:00:00", + "2002-04-09T00:00:00", + "2002-04-10T00:00:00", + "2002-04-11T00:00:00", + "2002-04-12T00:00:00", + "2002-04-13T00:00:00", + "2002-04-14T00:00:00", + "2002-04-15T00:00:00", + "2002-04-16T00:00:00", + "2002-04-17T00:00:00", + "2002-04-18T00:00:00", + "2002-04-19T00:00:00", + "2002-04-20T00:00:00", + "2002-04-21T00:00:00", + "2002-04-22T00:00:00", + "2002-04-23T00:00:00", + "2002-04-24T00:00:00", + "2002-04-25T00:00:00", + "2002-04-26T00:00:00", + "2002-04-27T00:00:00", + "2002-04-28T00:00:00", + "2002-04-29T00:00:00", + "2002-04-30T00:00:00", + "2002-05-01T00:00:00", + "2002-05-02T00:00:00", + "2002-05-03T00:00:00", + "2002-05-04T00:00:00", + "2002-05-05T00:00:00", + "2002-05-06T00:00:00", + "2002-05-07T00:00:00", + "2002-05-08T00:00:00", + "2002-05-09T00:00:00", + "2002-05-10T00:00:00", + "2002-05-11T00:00:00", + "2002-05-12T00:00:00", + "2002-05-13T00:00:00", + "2002-05-14T00:00:00", + "2002-05-15T00:00:00", + "2002-05-16T00:00:00", + "2002-05-17T00:00:00", + "2002-05-18T00:00:00", + "2002-05-19T00:00:00", + "2002-05-20T00:00:00", + "2002-05-21T00:00:00", + "2002-05-22T00:00:00", + "2002-05-23T00:00:00", + "2002-05-24T00:00:00", + "2002-05-25T00:00:00", + "2002-05-26T00:00:00", + "2002-05-27T00:00:00", + "2002-05-28T00:00:00", + "2002-05-29T00:00:00", + "2002-05-30T00:00:00", + "2002-05-31T00:00:00", + "2002-06-01T00:00:00", + "2002-06-02T00:00:00", + "2002-06-03T00:00:00", + "2002-06-04T00:00:00", + "2002-06-05T00:00:00", + "2002-06-06T00:00:00", + "2002-06-07T00:00:00", + "2002-06-08T00:00:00", + "2002-06-09T00:00:00", + "2002-06-10T00:00:00", + "2002-06-11T00:00:00", + "2002-06-12T00:00:00", + "2002-06-13T00:00:00", + "2002-06-14T00:00:00", + "2002-06-15T00:00:00", + "2002-06-16T00:00:00", + "2002-06-17T00:00:00", + "2002-06-18T00:00:00", + "2002-06-19T00:00:00", + "2002-06-20T00:00:00", + "2002-06-21T00:00:00", + "2002-06-22T00:00:00", + "2002-06-23T00:00:00", + "2002-06-24T00:00:00", + "2002-06-25T00:00:00", + "2002-06-26T00:00:00", + "2002-06-27T00:00:00", + "2002-06-28T00:00:00", + "2002-06-29T00:00:00", + "2002-06-30T00:00:00", + "2002-07-01T00:00:00", + "2002-07-02T00:00:00", + "2002-07-03T00:00:00", + "2002-07-04T00:00:00", + "2002-07-05T00:00:00", + "2002-07-06T00:00:00", + "2002-07-07T00:00:00", + "2002-07-08T00:00:00", + "2002-07-09T00:00:00", + "2002-07-10T00:00:00", + "2002-07-11T00:00:00", + "2002-07-12T00:00:00", + "2002-07-13T00:00:00", + "2002-07-14T00:00:00", + "2002-07-15T00:00:00", + "2002-07-16T00:00:00", + "2002-07-17T00:00:00", + "2002-07-18T00:00:00", + "2002-07-19T00:00:00", + "2002-07-20T00:00:00", + "2002-07-21T00:00:00", + "2002-07-22T00:00:00", + "2002-07-23T00:00:00", + "2002-07-24T00:00:00", + "2002-07-25T00:00:00", + "2002-07-26T00:00:00", + "2002-07-27T00:00:00", + "2002-07-28T00:00:00", + "2002-07-29T00:00:00", + "2002-07-30T00:00:00", + "2002-07-31T00:00:00", + "2002-08-01T00:00:00", + "2002-08-02T00:00:00", + "2002-08-03T00:00:00", + "2002-08-04T00:00:00", + "2002-08-05T00:00:00", + "2002-08-06T00:00:00", + "2002-08-07T00:00:00", + "2002-08-08T00:00:00", + "2002-08-09T00:00:00", + "2002-08-10T00:00:00", + "2002-08-11T00:00:00", + "2002-08-12T00:00:00", + "2002-08-13T00:00:00", + "2002-08-14T00:00:00", + "2002-08-15T00:00:00", + "2002-08-16T00:00:00", + "2002-08-17T00:00:00", + "2002-08-18T00:00:00", + "2002-08-19T00:00:00", + "2002-08-20T00:00:00", + "2002-08-21T00:00:00", + "2002-08-22T00:00:00", + "2002-08-23T00:00:00", + "2002-08-24T00:00:00", + "2002-08-25T00:00:00", + "2002-08-26T00:00:00", + "2002-08-27T00:00:00", + "2002-08-28T00:00:00", + "2002-08-29T00:00:00", + "2002-08-30T00:00:00", + "2002-08-31T00:00:00", + "2002-09-01T00:00:00", + "2002-09-02T00:00:00", + "2002-09-03T00:00:00", + "2002-09-04T00:00:00", + "2002-09-05T00:00:00", + "2002-09-06T00:00:00", + "2002-09-07T00:00:00", + "2002-09-08T00:00:00", + "2002-09-09T00:00:00", + "2002-09-10T00:00:00", + "2002-09-11T00:00:00", + "2002-09-12T00:00:00", + "2002-09-13T00:00:00", + "2002-09-14T00:00:00", + "2002-09-15T00:00:00", + "2002-09-16T00:00:00", + "2002-09-17T00:00:00", + "2002-09-18T00:00:00", + "2002-09-19T00:00:00", + "2002-09-20T00:00:00", + "2002-09-21T00:00:00", + "2002-09-22T00:00:00", + "2002-09-23T00:00:00", + "2002-09-24T00:00:00", + "2002-09-25T00:00:00", + "2002-09-26T00:00:00" + ], + "xaxis": "x", + "y": [ + -0.17417233140788102, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.2165336303720558, + 1.3085708703350725, + 1.3085708703350725, + 1.3085708703350725, + 1.3085708703350725, + 1.3085708703350725, + 1.3085708703350725, + 1.3085708703350725, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 1.9067812837762355, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.664478597876851, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 2.8909311065893863, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196, + 3.1217690859399196 + ], + "yaxis": "y" + } + ], + "layout": { + "legend": { + "title": { + "text": "variable" + }, + "tracegroupgap": 0 + }, + "margin": { + "t": 60 + }, + "template": { + "data": { + "bar": [ + { + "error_x": { + "color": "#2a3f5f" + }, + "error_y": { + "color": "#2a3f5f" + }, + "marker": { + "line": { + "color": "#E5ECF6", + "width": 0.5 + }, + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "bar" + } + ], + "barpolar": [ + { + "marker": { + "line": { + "color": "#E5ECF6", + "width": 0.5 + }, + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "barpolar" + } + ], + "carpet": [ + { + "aaxis": { + "endlinecolor": "#2a3f5f", + "gridcolor": "white", + "linecolor": "white", + "minorgridcolor": "white", + "startlinecolor": "#2a3f5f" + }, + "baxis": { + "endlinecolor": "#2a3f5f", + "gridcolor": "white", + "linecolor": "white", + "minorgridcolor": "white", + "startlinecolor": "#2a3f5f" + }, + "type": "carpet" + } + ], + "choropleth": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "choropleth" + } + ], + "contour": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "contour" + } + ], + "contourcarpet": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "contourcarpet" + } + ], + "heatmap": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "heatmap" + } + ], + "heatmapgl": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "heatmapgl" + } + ], + "histogram": [ + { + "marker": { + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "histogram" + } + ], + "histogram2d": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "histogram2d" + } + ], + "histogram2dcontour": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "histogram2dcontour" + } + ], + "mesh3d": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "mesh3d" + } + ], + "parcoords": [ + { + "line": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "parcoords" + } + ], + "pie": [ + { + "automargin": true, + "type": "pie" + } + ], + "scatter": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatter" + } + ], + "scatter3d": [ + { + "line": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatter3d" + } + ], + "scattercarpet": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattercarpet" + } + ], + "scattergeo": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattergeo" + } + ], + "scattergl": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattergl" + } + ], + "scattermapbox": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattermapbox" + } + ], + "scatterpolar": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterpolar" + } + ], + "scatterpolargl": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterpolargl" + } + ], + "scatterternary": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterternary" + } + ], + "surface": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "surface" + } + ], + "table": [ + { + "cells": { + "fill": { + "color": "#EBF0F8" + }, + "line": { + "color": "white" + } + }, + "header": { + "fill": { + "color": "#C8D4E3" + }, + "line": { + "color": "white" + } + }, + "type": "table" + } + ] + }, + "layout": { + "annotationdefaults": { + "arrowcolor": "#2a3f5f", + "arrowhead": 0, + "arrowwidth": 1 + }, + "autotypenumbers": "strict", + "coloraxis": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "colorscale": { + "diverging": [ + [ + 0, + "#8e0152" + ], + [ + 0.1, + "#c51b7d" + ], + [ + 0.2, + "#de77ae" + ], + [ + 0.3, + "#f1b6da" + ], + [ + 0.4, + "#fde0ef" + ], + [ + 0.5, + "#f7f7f7" + ], + [ + 0.6, + "#e6f5d0" + ], + [ + 0.7, + "#b8e186" + ], + [ + 0.8, + "#7fbc41" + ], + [ + 0.9, + "#4d9221" + ], + [ + 1, + "#276419" + ] + ], + "sequential": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "sequentialminus": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ] + }, + "colorway": [ + "#636efa", + "#EF553B", + "#00cc96", + "#ab63fa", + "#FFA15A", + "#19d3f3", + "#FF6692", + "#B6E880", + "#FF97FF", + "#FECB52" + ], + "font": { + "color": "#2a3f5f" + }, + "geo": { + "bgcolor": "white", + "lakecolor": "white", + "landcolor": "#E5ECF6", + "showlakes": true, + "showland": true, + "subunitcolor": "white" + }, + "hoverlabel": { + "align": "left" + }, + "hovermode": "closest", + "mapbox": { + "style": "light" + }, + "paper_bgcolor": "white", + "plot_bgcolor": "#E5ECF6", + "polar": { + "angularaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "bgcolor": "#E5ECF6", + "radialaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + } + }, + "scene": { + "xaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + }, + "yaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + }, + "zaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + } + }, + "shapedefaults": { + "line": { + "color": "#2a3f5f" + } + }, + "ternary": { + "aaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "baxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "bgcolor": "#E5ECF6", + "caxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + } + }, + "title": { + "x": 0.05 + }, + "xaxis": { + "automargin": true, + "gridcolor": "white", + "linecolor": "white", + "ticks": "", + "title": { + "standoff": 15 + }, + "zerolinecolor": "white", + "zerolinewidth": 2 + }, + "yaxis": { + "automargin": true, + "gridcolor": "white", + "linecolor": "white", + "ticks": "", + "title": { + "standoff": 15 + }, + "zerolinecolor": "white", + "zerolinewidth": 2 + } + } + }, + "xaxis": { + "anchor": "y", + "domain": [ + 0, + 1 + ], + "title": { + "text": "index" + } + }, + "yaxis": { + "anchor": "x", + "domain": [ + 0, + 1 + ], + "title": { + "text": "value" + } + } + } + }, + "text/html": [ + "
    " + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "psser.plot()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "On a DataFrame, the [plot()](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.plot.html) method is a convenience to plot all of the columns with labels:" + ] + }, + { + "cell_type": "code", + "execution_count": 46, + "metadata": {}, + "outputs": [], + "source": [ + "pdf = pd.DataFrame(np.random.randn(1000, 4), index=pser.index,\n", + " columns=['A', 'B', 'C', 'D'])" + ] + }, + { + "cell_type": "code", + "execution_count": 47, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = ps.from_pandas(pdf)" + ] + }, + { + "cell_type": "code", + "execution_count": 48, + "metadata": {}, + "outputs": [], + "source": [ + "psdf = psdf.cummax()" + ] + }, + { + "cell_type": "code", + "execution_count": 49, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.plotly.v1+json": { + "config": { + "plotlyServerURL": "https://plot.ly" + }, + "data": [ + { + "hovertemplate": "variable=A
    index=%{x}
    value=%{y}", + "legendgroup": "A", + "line": { + "color": "#636efa", + "dash": "solid" + }, + "mode": "lines", + "name": "A", + "showlegend": true, + "type": "scattergl", + "x": [ + "2000-01-01T00:00:00", + "2000-01-02T00:00:00", + "2000-01-03T00:00:00", + "2000-01-04T00:00:00", + "2000-01-05T00:00:00", + "2000-01-06T00:00:00", + "2000-01-07T00:00:00", + "2000-01-08T00:00:00", + "2000-01-09T00:00:00", + "2000-01-10T00:00:00", + "2000-01-11T00:00:00", + "2000-01-12T00:00:00", + "2000-01-13T00:00:00", + "2000-01-14T00:00:00", + "2000-01-15T00:00:00", + "2000-01-16T00:00:00", + "2000-01-17T00:00:00", + "2000-01-18T00:00:00", + "2000-01-19T00:00:00", + "2000-01-20T00:00:00", + "2000-01-21T00:00:00", + "2000-01-22T00:00:00", + "2000-01-23T00:00:00", + "2000-01-24T00:00:00", + "2000-01-25T00:00:00", + "2000-01-26T00:00:00", + "2000-01-27T00:00:00", + "2000-01-28T00:00:00", + "2000-01-29T00:00:00", + "2000-01-30T00:00:00", + "2000-01-31T00:00:00", + "2000-02-01T00:00:00", + "2000-02-02T00:00:00", + "2000-02-03T00:00:00", + "2000-02-04T00:00:00", + "2000-02-05T00:00:00", + "2000-02-06T00:00:00", + "2000-02-07T00:00:00", + "2000-02-08T00:00:00", + "2000-02-09T00:00:00", + "2000-02-10T00:00:00", + "2000-02-11T00:00:00", + "2000-02-12T00:00:00", + "2000-02-13T00:00:00", + "2000-02-14T00:00:00", + "2000-02-15T00:00:00", + "2000-02-16T00:00:00", + "2000-02-17T00:00:00", + "2000-02-18T00:00:00", + "2000-02-19T00:00:00", + "2000-02-20T00:00:00", + "2000-02-21T00:00:00", + "2000-02-22T00:00:00", + "2000-02-23T00:00:00", + "2000-02-24T00:00:00", + "2000-02-25T00:00:00", + "2000-02-26T00:00:00", + "2000-02-27T00:00:00", + "2000-02-28T00:00:00", + "2000-02-29T00:00:00", + "2000-03-01T00:00:00", + "2000-03-02T00:00:00", + "2000-03-03T00:00:00", + "2000-03-04T00:00:00", + "2000-03-05T00:00:00", + "2000-03-06T00:00:00", + "2000-03-07T00:00:00", + "2000-03-08T00:00:00", + "2000-03-09T00:00:00", + "2000-03-10T00:00:00", + "2000-03-11T00:00:00", + "2000-03-12T00:00:00", + "2000-03-13T00:00:00", + "2000-03-14T00:00:00", + "2000-03-15T00:00:00", + "2000-03-16T00:00:00", + "2000-03-17T00:00:00", + "2000-03-18T00:00:00", + "2000-03-19T00:00:00", + "2000-03-20T00:00:00", + "2000-03-21T00:00:00", + "2000-03-22T00:00:00", + "2000-03-23T00:00:00", + "2000-03-24T00:00:00", + "2000-03-25T00:00:00", + "2000-03-26T00:00:00", + "2000-03-27T00:00:00", + "2000-03-28T00:00:00", + "2000-03-29T00:00:00", + "2000-03-30T00:00:00", + "2000-03-31T00:00:00", + "2000-04-01T00:00:00", + "2000-04-02T00:00:00", + "2000-04-03T00:00:00", + "2000-04-04T00:00:00", + "2000-04-05T00:00:00", + "2000-04-06T00:00:00", + "2000-04-07T00:00:00", + "2000-04-08T00:00:00", + "2000-04-09T00:00:00", + "2000-04-10T00:00:00", + "2000-04-11T00:00:00", + "2000-04-12T00:00:00", + "2000-04-13T00:00:00", + "2000-04-14T00:00:00", + "2000-04-15T00:00:00", + "2000-04-16T00:00:00", + "2000-04-17T00:00:00", + "2000-04-18T00:00:00", + "2000-04-19T00:00:00", + "2000-04-20T00:00:00", + "2000-04-21T00:00:00", + "2000-04-22T00:00:00", + "2000-04-23T00:00:00", + "2000-04-24T00:00:00", + "2000-04-25T00:00:00", + "2000-04-26T00:00:00", + "2000-04-27T00:00:00", + "2000-04-28T00:00:00", + "2000-04-29T00:00:00", + "2000-04-30T00:00:00", + "2000-05-01T00:00:00", + "2000-05-02T00:00:00", + "2000-05-03T00:00:00", + "2000-05-04T00:00:00", + "2000-05-05T00:00:00", + "2000-05-06T00:00:00", + "2000-05-07T00:00:00", + "2000-05-08T00:00:00", + "2000-05-09T00:00:00", + "2000-05-10T00:00:00", + "2000-05-11T00:00:00", + "2000-05-12T00:00:00", + "2000-05-13T00:00:00", + "2000-05-14T00:00:00", + "2000-05-15T00:00:00", + "2000-05-16T00:00:00", + "2000-05-17T00:00:00", + "2000-05-18T00:00:00", + "2000-05-19T00:00:00", + "2000-05-20T00:00:00", + "2000-05-21T00:00:00", + "2000-05-22T00:00:00", + "2000-05-23T00:00:00", + "2000-05-24T00:00:00", + "2000-05-25T00:00:00", + "2000-05-26T00:00:00", + "2000-05-27T00:00:00", + "2000-05-28T00:00:00", + "2000-05-29T00:00:00", + "2000-05-30T00:00:00", + "2000-05-31T00:00:00", + "2000-06-01T00:00:00", + "2000-06-02T00:00:00", + "2000-06-03T00:00:00", + "2000-06-04T00:00:00", + "2000-06-05T00:00:00", + "2000-06-06T00:00:00", + "2000-06-07T00:00:00", + "2000-06-08T00:00:00", + "2000-06-09T00:00:00", + "2000-06-10T00:00:00", + "2000-06-11T00:00:00", + "2000-06-12T00:00:00", + "2000-06-13T00:00:00", + "2000-06-14T00:00:00", + "2000-06-15T00:00:00", + "2000-06-16T00:00:00", + "2000-06-17T00:00:00", + "2000-06-18T00:00:00", + "2000-06-19T00:00:00", + "2000-06-20T00:00:00", + "2000-06-21T00:00:00", + "2000-06-22T00:00:00", + "2000-06-23T00:00:00", + "2000-06-24T00:00:00", + "2000-06-25T00:00:00", + "2000-06-26T00:00:00", + "2000-06-27T00:00:00", + "2000-06-28T00:00:00", + "2000-06-29T00:00:00", + "2000-06-30T00:00:00", + "2000-07-01T00:00:00", + "2000-07-02T00:00:00", + "2000-07-03T00:00:00", + "2000-07-04T00:00:00", + "2000-07-05T00:00:00", + "2000-07-06T00:00:00", + "2000-07-07T00:00:00", + "2000-07-08T00:00:00", + "2000-07-09T00:00:00", + "2000-07-10T00:00:00", + "2000-07-11T00:00:00", + "2000-07-12T00:00:00", + "2000-07-13T00:00:00", + "2000-07-14T00:00:00", + "2000-07-15T00:00:00", + "2000-07-16T00:00:00", + "2000-07-17T00:00:00", + "2000-07-18T00:00:00", + "2000-07-19T00:00:00", + "2000-07-20T00:00:00", + "2000-07-21T00:00:00", + "2000-07-22T00:00:00", + "2000-07-23T00:00:00", + "2000-07-24T00:00:00", + "2000-07-25T00:00:00", + "2000-07-26T00:00:00", + "2000-07-27T00:00:00", + "2000-07-28T00:00:00", + "2000-07-29T00:00:00", + "2000-07-30T00:00:00", + "2000-07-31T00:00:00", + "2000-08-01T00:00:00", + "2000-08-02T00:00:00", + "2000-08-03T00:00:00", + "2000-08-04T00:00:00", + "2000-08-05T00:00:00", + "2000-08-06T00:00:00", + "2000-08-07T00:00:00", + "2000-08-08T00:00:00", + "2000-08-09T00:00:00", + "2000-08-10T00:00:00", + "2000-08-11T00:00:00", + "2000-08-12T00:00:00", + "2000-08-13T00:00:00", + "2000-08-14T00:00:00", + "2000-08-15T00:00:00", + "2000-08-16T00:00:00", + "2000-08-17T00:00:00", + "2000-08-18T00:00:00", + "2000-08-19T00:00:00", + "2000-08-20T00:00:00", + "2000-08-21T00:00:00", + "2000-08-22T00:00:00", + "2000-08-23T00:00:00", + "2000-08-24T00:00:00", + "2000-08-25T00:00:00", + "2000-08-26T00:00:00", + "2000-08-27T00:00:00", + "2000-08-28T00:00:00", + "2000-08-29T00:00:00", + "2000-08-30T00:00:00", + "2000-08-31T00:00:00", + "2000-09-01T00:00:00", + "2000-09-02T00:00:00", + "2000-09-03T00:00:00", + "2000-09-04T00:00:00", + "2000-09-05T00:00:00", + "2000-09-06T00:00:00", + "2000-09-07T00:00:00", + "2000-09-08T00:00:00", + "2000-09-09T00:00:00", + "2000-09-10T00:00:00", + "2000-09-11T00:00:00", + "2000-09-12T00:00:00", + "2000-09-13T00:00:00", + "2000-09-14T00:00:00", + "2000-09-15T00:00:00", + "2000-09-16T00:00:00", + "2000-09-17T00:00:00", + "2000-09-18T00:00:00", + "2000-09-19T00:00:00", + "2000-09-20T00:00:00", + "2000-09-21T00:00:00", + "2000-09-22T00:00:00", + "2000-09-23T00:00:00", + "2000-09-24T00:00:00", + "2000-09-25T00:00:00", + "2000-09-26T00:00:00", + "2000-09-27T00:00:00", + "2000-09-28T00:00:00", + "2000-09-29T00:00:00", + "2000-09-30T00:00:00", + "2000-10-01T00:00:00", + "2000-10-02T00:00:00", + "2000-10-03T00:00:00", + "2000-10-04T00:00:00", + "2000-10-05T00:00:00", + "2000-10-06T00:00:00", + "2000-10-07T00:00:00", + "2000-10-08T00:00:00", + "2000-10-09T00:00:00", + "2000-10-10T00:00:00", + "2000-10-11T00:00:00", + "2000-10-12T00:00:00", + "2000-10-13T00:00:00", + "2000-10-14T00:00:00", + "2000-10-15T00:00:00", + "2000-10-16T00:00:00", + "2000-10-17T00:00:00", + "2000-10-18T00:00:00", + "2000-10-19T00:00:00", + "2000-10-20T00:00:00", + "2000-10-21T00:00:00", + "2000-10-22T00:00:00", + "2000-10-23T00:00:00", + "2000-10-24T00:00:00", + "2000-10-25T00:00:00", + "2000-10-26T00:00:00", + "2000-10-27T00:00:00", + "2000-10-28T00:00:00", + "2000-10-29T00:00:00", + "2000-10-30T00:00:00", + "2000-10-31T00:00:00", + "2000-11-01T00:00:00", + "2000-11-02T00:00:00", + "2000-11-03T00:00:00", + "2000-11-04T00:00:00", + "2000-11-05T00:00:00", + "2000-11-06T00:00:00", + "2000-11-07T00:00:00", + "2000-11-08T00:00:00", + "2000-11-09T00:00:00", + "2000-11-10T00:00:00", + "2000-11-11T00:00:00", + "2000-11-12T00:00:00", + "2000-11-13T00:00:00", + "2000-11-14T00:00:00", + "2000-11-15T00:00:00", + "2000-11-16T00:00:00", + "2000-11-17T00:00:00", + "2000-11-18T00:00:00", + "2000-11-19T00:00:00", + "2000-11-20T00:00:00", + "2000-11-21T00:00:00", + "2000-11-22T00:00:00", + "2000-11-23T00:00:00", + "2000-11-24T00:00:00", + "2000-11-25T00:00:00", + "2000-11-26T00:00:00", + "2000-11-27T00:00:00", + "2000-11-28T00:00:00", + "2000-11-29T00:00:00", + "2000-11-30T00:00:00", + "2000-12-01T00:00:00", + "2000-12-02T00:00:00", + "2000-12-03T00:00:00", + "2000-12-04T00:00:00", + "2000-12-05T00:00:00", + "2000-12-06T00:00:00", + "2000-12-07T00:00:00", + "2000-12-08T00:00:00", + "2000-12-09T00:00:00", + "2000-12-10T00:00:00", + "2000-12-11T00:00:00", + "2000-12-12T00:00:00", + "2000-12-13T00:00:00", + "2000-12-14T00:00:00", + "2000-12-15T00:00:00", + "2000-12-16T00:00:00", + "2000-12-17T00:00:00", + "2000-12-18T00:00:00", + "2000-12-19T00:00:00", + "2000-12-20T00:00:00", + "2000-12-21T00:00:00", + "2000-12-22T00:00:00", + "2000-12-23T00:00:00", + "2000-12-24T00:00:00", + "2000-12-25T00:00:00", + "2000-12-26T00:00:00", + "2000-12-27T00:00:00", + "2000-12-28T00:00:00", + "2000-12-29T00:00:00", + "2000-12-30T00:00:00", + "2000-12-31T00:00:00", + "2001-01-01T00:00:00", + "2001-01-02T00:00:00", + "2001-01-03T00:00:00", + "2001-01-04T00:00:00", + "2001-01-05T00:00:00", + "2001-01-06T00:00:00", + "2001-01-07T00:00:00", + "2001-01-08T00:00:00", + "2001-01-09T00:00:00", + "2001-01-10T00:00:00", + "2001-01-11T00:00:00", + "2001-01-12T00:00:00", + "2001-01-13T00:00:00", + "2001-01-14T00:00:00", + "2001-01-15T00:00:00", + "2001-01-16T00:00:00", + "2001-01-17T00:00:00", + "2001-01-18T00:00:00", + "2001-01-19T00:00:00", + "2001-01-20T00:00:00", + "2001-01-21T00:00:00", + "2001-01-22T00:00:00", + "2001-01-23T00:00:00", + "2001-01-24T00:00:00", + "2001-01-25T00:00:00", + "2001-01-26T00:00:00", + "2001-01-27T00:00:00", + "2001-01-28T00:00:00", + "2001-01-29T00:00:00", + "2001-01-30T00:00:00", + "2001-01-31T00:00:00", + "2001-02-01T00:00:00", + "2001-02-02T00:00:00", + "2001-02-03T00:00:00", + "2001-02-04T00:00:00", + "2001-02-05T00:00:00", + "2001-02-06T00:00:00", + "2001-02-07T00:00:00", + "2001-02-08T00:00:00", + "2001-02-09T00:00:00", + "2001-02-10T00:00:00", + "2001-02-11T00:00:00", + "2001-02-12T00:00:00", + "2001-02-13T00:00:00", + "2001-02-14T00:00:00", + "2001-02-15T00:00:00", + "2001-02-16T00:00:00", + "2001-02-17T00:00:00", + "2001-02-18T00:00:00", + "2001-02-19T00:00:00", + "2001-02-20T00:00:00", + "2001-02-21T00:00:00", + "2001-02-22T00:00:00", + "2001-02-23T00:00:00", + "2001-02-24T00:00:00", + "2001-02-25T00:00:00", + "2001-02-26T00:00:00", + "2001-02-27T00:00:00", + "2001-02-28T00:00:00", + "2001-03-01T00:00:00", + "2001-03-02T00:00:00", + "2001-03-03T00:00:00", + "2001-03-04T00:00:00", + "2001-03-05T00:00:00", + "2001-03-06T00:00:00", + "2001-03-07T00:00:00", + "2001-03-08T00:00:00", + "2001-03-09T00:00:00", + "2001-03-10T00:00:00", + "2001-03-11T00:00:00", + "2001-03-12T00:00:00", + "2001-03-13T00:00:00", + "2001-03-14T00:00:00", + "2001-03-15T00:00:00", + "2001-03-16T00:00:00", + "2001-03-17T00:00:00", + "2001-03-18T00:00:00", + "2001-03-19T00:00:00", + "2001-03-20T00:00:00", + "2001-03-21T00:00:00", + "2001-03-22T00:00:00", + "2001-03-23T00:00:00", + "2001-03-24T00:00:00", + "2001-03-25T00:00:00", + "2001-03-26T00:00:00", + "2001-03-27T00:00:00", + "2001-03-28T00:00:00", + "2001-03-29T00:00:00", + "2001-03-30T00:00:00", + "2001-03-31T00:00:00", + "2001-04-01T00:00:00", + "2001-04-02T00:00:00", + "2001-04-03T00:00:00", + "2001-04-04T00:00:00", + "2001-04-05T00:00:00", + "2001-04-06T00:00:00", + "2001-04-07T00:00:00", + "2001-04-08T00:00:00", + "2001-04-09T00:00:00", + "2001-04-10T00:00:00", + "2001-04-11T00:00:00", + "2001-04-12T00:00:00", + "2001-04-13T00:00:00", + "2001-04-14T00:00:00", + "2001-04-15T00:00:00", + "2001-04-16T00:00:00", + "2001-04-17T00:00:00", + "2001-04-18T00:00:00", + "2001-04-19T00:00:00", + "2001-04-20T00:00:00", + "2001-04-21T00:00:00", + "2001-04-22T00:00:00", + "2001-04-23T00:00:00", + "2001-04-24T00:00:00", + "2001-04-25T00:00:00", + "2001-04-26T00:00:00", + "2001-04-27T00:00:00", + "2001-04-28T00:00:00", + "2001-04-29T00:00:00", + "2001-04-30T00:00:00", + "2001-05-01T00:00:00", + "2001-05-02T00:00:00", + "2001-05-03T00:00:00", + "2001-05-04T00:00:00", + "2001-05-05T00:00:00", + "2001-05-06T00:00:00", + "2001-05-07T00:00:00", + "2001-05-08T00:00:00", + "2001-05-09T00:00:00", + "2001-05-10T00:00:00", + "2001-05-11T00:00:00", + "2001-05-12T00:00:00", + "2001-05-13T00:00:00", + "2001-05-14T00:00:00", + "2001-05-15T00:00:00", + "2001-05-16T00:00:00", + "2001-05-17T00:00:00", + "2001-05-18T00:00:00", + "2001-05-19T00:00:00", + "2001-05-20T00:00:00", + "2001-05-21T00:00:00", + "2001-05-22T00:00:00", + "2001-05-23T00:00:00", + "2001-05-24T00:00:00", + "2001-05-25T00:00:00", + "2001-05-26T00:00:00", + "2001-05-27T00:00:00", + "2001-05-28T00:00:00", + "2001-05-29T00:00:00", + "2001-05-30T00:00:00", + "2001-05-31T00:00:00", + "2001-06-01T00:00:00", + "2001-06-02T00:00:00", + "2001-06-03T00:00:00", + "2001-06-04T00:00:00", + "2001-06-05T00:00:00", + "2001-06-06T00:00:00", + "2001-06-07T00:00:00", + "2001-06-08T00:00:00", + "2001-06-09T00:00:00", + "2001-06-10T00:00:00", + "2001-06-11T00:00:00", + "2001-06-12T00:00:00", + "2001-06-13T00:00:00", + "2001-06-14T00:00:00", + "2001-06-15T00:00:00", + "2001-06-16T00:00:00", + "2001-06-17T00:00:00", + "2001-06-18T00:00:00", + "2001-06-19T00:00:00", + "2001-06-20T00:00:00", + "2001-06-21T00:00:00", + "2001-06-22T00:00:00", + "2001-06-23T00:00:00", + "2001-06-24T00:00:00", + "2001-06-25T00:00:00", + "2001-06-26T00:00:00", + "2001-06-27T00:00:00", + "2001-06-28T00:00:00", + "2001-06-29T00:00:00", + "2001-06-30T00:00:00", + "2001-07-01T00:00:00", + "2001-07-02T00:00:00", + "2001-07-03T00:00:00", + "2001-07-04T00:00:00", + "2001-07-05T00:00:00", + "2001-07-06T00:00:00", + "2001-07-07T00:00:00", + "2001-07-08T00:00:00", + "2001-07-09T00:00:00", + "2001-07-10T00:00:00", + "2001-07-11T00:00:00", + "2001-07-12T00:00:00", + "2001-07-13T00:00:00", + "2001-07-14T00:00:00", + "2001-07-15T00:00:00", + "2001-07-16T00:00:00", + "2001-07-17T00:00:00", + "2001-07-18T00:00:00", + "2001-07-19T00:00:00", + "2001-07-20T00:00:00", + "2001-07-21T00:00:00", + "2001-07-22T00:00:00", + "2001-07-23T00:00:00", + "2001-07-24T00:00:00", + "2001-07-25T00:00:00", + "2001-07-26T00:00:00", + "2001-07-27T00:00:00", + "2001-07-28T00:00:00", + "2001-07-29T00:00:00", + "2001-07-30T00:00:00", + "2001-07-31T00:00:00", + "2001-08-01T00:00:00", + "2001-08-02T00:00:00", + "2001-08-03T00:00:00", + "2001-08-04T00:00:00", + "2001-08-05T00:00:00", + "2001-08-06T00:00:00", + "2001-08-07T00:00:00", + "2001-08-08T00:00:00", + "2001-08-09T00:00:00", + "2001-08-10T00:00:00", + "2001-08-11T00:00:00", + "2001-08-12T00:00:00", + "2001-08-13T00:00:00", + "2001-08-14T00:00:00", + "2001-08-15T00:00:00", + "2001-08-16T00:00:00", + "2001-08-17T00:00:00", + "2001-08-18T00:00:00", + "2001-08-19T00:00:00", + "2001-08-20T00:00:00", + "2001-08-21T00:00:00", + "2001-08-22T00:00:00", + "2001-08-23T00:00:00", + "2001-08-24T00:00:00", + "2001-08-25T00:00:00", + "2001-08-26T00:00:00", + "2001-08-27T00:00:00", + "2001-08-28T00:00:00", + "2001-08-29T00:00:00", + "2001-08-30T00:00:00", + "2001-08-31T00:00:00", + "2001-09-01T00:00:00", + "2001-09-02T00:00:00", + "2001-09-03T00:00:00", + "2001-09-04T00:00:00", + "2001-09-05T00:00:00", + "2001-09-06T00:00:00", + "2001-09-07T00:00:00", + "2001-09-08T00:00:00", + "2001-09-09T00:00:00", + "2001-09-10T00:00:00", + "2001-09-11T00:00:00", + "2001-09-12T00:00:00", + "2001-09-13T00:00:00", + "2001-09-14T00:00:00", + "2001-09-15T00:00:00", + "2001-09-16T00:00:00", + "2001-09-17T00:00:00", + "2001-09-18T00:00:00", + "2001-09-19T00:00:00", + "2001-09-20T00:00:00", + "2001-09-21T00:00:00", + "2001-09-22T00:00:00", + "2001-09-23T00:00:00", + "2001-09-24T00:00:00", + "2001-09-25T00:00:00", + "2001-09-26T00:00:00", + "2001-09-27T00:00:00", + "2001-09-28T00:00:00", + "2001-09-29T00:00:00", + "2001-09-30T00:00:00", + "2001-10-01T00:00:00", + "2001-10-02T00:00:00", + "2001-10-03T00:00:00", + "2001-10-04T00:00:00", + "2001-10-05T00:00:00", + "2001-10-06T00:00:00", + "2001-10-07T00:00:00", + "2001-10-08T00:00:00", + "2001-10-09T00:00:00", + "2001-10-10T00:00:00", + "2001-10-11T00:00:00", + "2001-10-12T00:00:00", + "2001-10-13T00:00:00", + "2001-10-14T00:00:00", + "2001-10-15T00:00:00", + "2001-10-16T00:00:00", + "2001-10-17T00:00:00", + "2001-10-18T00:00:00", + "2001-10-19T00:00:00", + "2001-10-20T00:00:00", + "2001-10-21T00:00:00", + "2001-10-22T00:00:00", + "2001-10-23T00:00:00", + "2001-10-24T00:00:00", + "2001-10-25T00:00:00", + "2001-10-26T00:00:00", + "2001-10-27T00:00:00", + "2001-10-28T00:00:00", + "2001-10-29T00:00:00", + "2001-10-30T00:00:00", + "2001-10-31T00:00:00", + "2001-11-01T00:00:00", + "2001-11-02T00:00:00", + "2001-11-03T00:00:00", + "2001-11-04T00:00:00", + "2001-11-05T00:00:00", + "2001-11-06T00:00:00", + "2001-11-07T00:00:00", + "2001-11-08T00:00:00", + "2001-11-09T00:00:00", + "2001-11-10T00:00:00", + "2001-11-11T00:00:00", + "2001-11-12T00:00:00", + "2001-11-13T00:00:00", + "2001-11-14T00:00:00", + "2001-11-15T00:00:00", + "2001-11-16T00:00:00", + "2001-11-17T00:00:00", + "2001-11-18T00:00:00", + "2001-11-19T00:00:00", + "2001-11-20T00:00:00", + "2001-11-21T00:00:00", + "2001-11-22T00:00:00", + "2001-11-23T00:00:00", + "2001-11-24T00:00:00", + "2001-11-25T00:00:00", + "2001-11-26T00:00:00", + "2001-11-27T00:00:00", + "2001-11-28T00:00:00", + "2001-11-29T00:00:00", + "2001-11-30T00:00:00", + "2001-12-01T00:00:00", + "2001-12-02T00:00:00", + "2001-12-03T00:00:00", + "2001-12-04T00:00:00", + "2001-12-05T00:00:00", + "2001-12-06T00:00:00", + "2001-12-07T00:00:00", + "2001-12-08T00:00:00", + "2001-12-09T00:00:00", + "2001-12-10T00:00:00", + "2001-12-11T00:00:00", + "2001-12-12T00:00:00", + "2001-12-13T00:00:00", + "2001-12-14T00:00:00", + "2001-12-15T00:00:00", + "2001-12-16T00:00:00", + "2001-12-17T00:00:00", + "2001-12-18T00:00:00", + "2001-12-19T00:00:00", + "2001-12-20T00:00:00", + "2001-12-21T00:00:00", + "2001-12-22T00:00:00", + "2001-12-23T00:00:00", + "2001-12-24T00:00:00", + "2001-12-25T00:00:00", + "2001-12-26T00:00:00", + "2001-12-27T00:00:00", + "2001-12-28T00:00:00", + "2001-12-29T00:00:00", + "2001-12-30T00:00:00", + "2001-12-31T00:00:00", + "2002-01-01T00:00:00", + "2002-01-02T00:00:00", + "2002-01-03T00:00:00", + "2002-01-04T00:00:00", + "2002-01-05T00:00:00", + "2002-01-06T00:00:00", + "2002-01-07T00:00:00", + "2002-01-08T00:00:00", + "2002-01-09T00:00:00", + "2002-01-10T00:00:00", + "2002-01-11T00:00:00", + "2002-01-12T00:00:00", + "2002-01-13T00:00:00", + "2002-01-14T00:00:00", + "2002-01-15T00:00:00", + "2002-01-16T00:00:00", + "2002-01-17T00:00:00", + "2002-01-18T00:00:00", + "2002-01-19T00:00:00", + "2002-01-20T00:00:00", + "2002-01-21T00:00:00", + "2002-01-22T00:00:00", + "2002-01-23T00:00:00", + "2002-01-24T00:00:00", + "2002-01-25T00:00:00", + "2002-01-26T00:00:00", + "2002-01-27T00:00:00", + "2002-01-28T00:00:00", + "2002-01-29T00:00:00", + "2002-01-30T00:00:00", + "2002-01-31T00:00:00", + "2002-02-01T00:00:00", + "2002-02-02T00:00:00", + "2002-02-03T00:00:00", + "2002-02-04T00:00:00", + "2002-02-05T00:00:00", + "2002-02-06T00:00:00", + "2002-02-07T00:00:00", + "2002-02-08T00:00:00", + "2002-02-09T00:00:00", + "2002-02-10T00:00:00", + "2002-02-11T00:00:00", + "2002-02-12T00:00:00", + "2002-02-13T00:00:00", + "2002-02-14T00:00:00", + "2002-02-15T00:00:00", + "2002-02-16T00:00:00", + "2002-02-17T00:00:00", + "2002-02-18T00:00:00", + "2002-02-19T00:00:00", + "2002-02-20T00:00:00", + "2002-02-21T00:00:00", + "2002-02-22T00:00:00", + "2002-02-23T00:00:00", + "2002-02-24T00:00:00", + "2002-02-25T00:00:00", + "2002-02-26T00:00:00", + "2002-02-27T00:00:00", + "2002-02-28T00:00:00", + "2002-03-01T00:00:00", + "2002-03-02T00:00:00", + "2002-03-03T00:00:00", + "2002-03-04T00:00:00", + "2002-03-05T00:00:00", + "2002-03-06T00:00:00", + "2002-03-07T00:00:00", + "2002-03-08T00:00:00", + "2002-03-09T00:00:00", + "2002-03-10T00:00:00", + "2002-03-11T00:00:00", + "2002-03-12T00:00:00", + "2002-03-13T00:00:00", + "2002-03-14T00:00:00", + "2002-03-15T00:00:00", + "2002-03-16T00:00:00", + "2002-03-17T00:00:00", + "2002-03-18T00:00:00", + "2002-03-19T00:00:00", + "2002-03-20T00:00:00", + "2002-03-21T00:00:00", + "2002-03-22T00:00:00", + "2002-03-23T00:00:00", + "2002-03-24T00:00:00", + "2002-03-25T00:00:00", + "2002-03-26T00:00:00", + "2002-03-27T00:00:00", + "2002-03-28T00:00:00", + "2002-03-29T00:00:00", + "2002-03-30T00:00:00", + "2002-03-31T00:00:00", + "2002-04-01T00:00:00", + "2002-04-02T00:00:00", + "2002-04-03T00:00:00", + "2002-04-04T00:00:00", + "2002-04-05T00:00:00", + "2002-04-06T00:00:00", + "2002-04-07T00:00:00", + "2002-04-08T00:00:00", + "2002-04-09T00:00:00", + "2002-04-10T00:00:00", + "2002-04-11T00:00:00", + "2002-04-12T00:00:00", + "2002-04-13T00:00:00", + "2002-04-14T00:00:00", + "2002-04-15T00:00:00", + "2002-04-16T00:00:00", + "2002-04-17T00:00:00", + "2002-04-18T00:00:00", + "2002-04-19T00:00:00", + "2002-04-20T00:00:00", + "2002-04-21T00:00:00", + "2002-04-22T00:00:00", + "2002-04-23T00:00:00", + "2002-04-24T00:00:00", + "2002-04-25T00:00:00", + "2002-04-26T00:00:00", + "2002-04-27T00:00:00", + "2002-04-28T00:00:00", + "2002-04-29T00:00:00", + "2002-04-30T00:00:00", + "2002-05-01T00:00:00", + "2002-05-02T00:00:00", + "2002-05-03T00:00:00", + "2002-05-04T00:00:00", + "2002-05-05T00:00:00", + "2002-05-06T00:00:00", + "2002-05-07T00:00:00", + "2002-05-08T00:00:00", + "2002-05-09T00:00:00", + "2002-05-10T00:00:00", + "2002-05-11T00:00:00", + "2002-05-12T00:00:00", + "2002-05-13T00:00:00", + "2002-05-14T00:00:00", + "2002-05-15T00:00:00", + "2002-05-16T00:00:00", + "2002-05-17T00:00:00", + "2002-05-18T00:00:00", + "2002-05-19T00:00:00", + "2002-05-20T00:00:00", + "2002-05-21T00:00:00", + "2002-05-22T00:00:00", + "2002-05-23T00:00:00", + "2002-05-24T00:00:00", + "2002-05-25T00:00:00", + "2002-05-26T00:00:00", + "2002-05-27T00:00:00", + "2002-05-28T00:00:00", + "2002-05-29T00:00:00", + "2002-05-30T00:00:00", + "2002-05-31T00:00:00", + "2002-06-01T00:00:00", + "2002-06-02T00:00:00", + "2002-06-03T00:00:00", + "2002-06-04T00:00:00", + "2002-06-05T00:00:00", + "2002-06-06T00:00:00", + "2002-06-07T00:00:00", + "2002-06-08T00:00:00", + "2002-06-09T00:00:00", + "2002-06-10T00:00:00", + "2002-06-11T00:00:00", + "2002-06-12T00:00:00", + "2002-06-13T00:00:00", + "2002-06-14T00:00:00", + "2002-06-15T00:00:00", + "2002-06-16T00:00:00", + "2002-06-17T00:00:00", + "2002-06-18T00:00:00", + "2002-06-19T00:00:00", + "2002-06-20T00:00:00", + "2002-06-21T00:00:00", + "2002-06-22T00:00:00", + "2002-06-23T00:00:00", + "2002-06-24T00:00:00", + "2002-06-25T00:00:00", + "2002-06-26T00:00:00", + "2002-06-27T00:00:00", + "2002-06-28T00:00:00", + "2002-06-29T00:00:00", + "2002-06-30T00:00:00", + "2002-07-01T00:00:00", + "2002-07-02T00:00:00", + "2002-07-03T00:00:00", + "2002-07-04T00:00:00", + "2002-07-05T00:00:00", + "2002-07-06T00:00:00", + "2002-07-07T00:00:00", + "2002-07-08T00:00:00", + "2002-07-09T00:00:00", + "2002-07-10T00:00:00", + "2002-07-11T00:00:00", + "2002-07-12T00:00:00", + "2002-07-13T00:00:00", + "2002-07-14T00:00:00", + "2002-07-15T00:00:00", + "2002-07-16T00:00:00", + "2002-07-17T00:00:00", + "2002-07-18T00:00:00", + "2002-07-19T00:00:00", + "2002-07-20T00:00:00", + "2002-07-21T00:00:00", + "2002-07-22T00:00:00", + "2002-07-23T00:00:00", + "2002-07-24T00:00:00", + "2002-07-25T00:00:00", + "2002-07-26T00:00:00", + "2002-07-27T00:00:00", + "2002-07-28T00:00:00", + "2002-07-29T00:00:00", + "2002-07-30T00:00:00", + "2002-07-31T00:00:00", + "2002-08-01T00:00:00", + "2002-08-02T00:00:00", + "2002-08-03T00:00:00", + "2002-08-04T00:00:00", + "2002-08-05T00:00:00", + "2002-08-06T00:00:00", + "2002-08-07T00:00:00", + "2002-08-08T00:00:00", + "2002-08-09T00:00:00", + "2002-08-10T00:00:00", + "2002-08-11T00:00:00", + "2002-08-12T00:00:00", + "2002-08-13T00:00:00", + "2002-08-14T00:00:00", + "2002-08-15T00:00:00", + "2002-08-16T00:00:00", + "2002-08-17T00:00:00", + "2002-08-18T00:00:00", + "2002-08-19T00:00:00", + "2002-08-20T00:00:00", + "2002-08-21T00:00:00", + "2002-08-22T00:00:00", + "2002-08-23T00:00:00", + "2002-08-24T00:00:00", + "2002-08-25T00:00:00", + "2002-08-26T00:00:00", + "2002-08-27T00:00:00", + "2002-08-28T00:00:00", + "2002-08-29T00:00:00", + "2002-08-30T00:00:00", + "2002-08-31T00:00:00", + "2002-09-01T00:00:00", + "2002-09-02T00:00:00", + "2002-09-03T00:00:00", + "2002-09-04T00:00:00", + "2002-09-05T00:00:00", + "2002-09-06T00:00:00", + "2002-09-07T00:00:00", + "2002-09-08T00:00:00", + "2002-09-09T00:00:00", + "2002-09-10T00:00:00", + "2002-09-11T00:00:00", + "2002-09-12T00:00:00", + "2002-09-13T00:00:00", + "2002-09-14T00:00:00", + "2002-09-15T00:00:00", + "2002-09-16T00:00:00", + "2002-09-17T00:00:00", + "2002-09-18T00:00:00", + "2002-09-19T00:00:00", + "2002-09-20T00:00:00", + "2002-09-21T00:00:00", + "2002-09-22T00:00:00", + "2002-09-23T00:00:00", + "2002-09-24T00:00:00", + "2002-09-25T00:00:00", + "2002-09-26T00:00:00" + ], + "xaxis": "x", + "y": [ + -1.1870965907666513, + 0.331740732048968, + 0.331740732048968, + 0.6216201483866751, + 0.6216201483866751, + 2.1691978004162578, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493, + 2.7557379611927493 + ], + "yaxis": "y" + }, + { + "hovertemplate": "variable=B
    index=%{x}
    value=%{y}", + "legendgroup": "B", + "line": { + "color": "#EF553B", + "dash": "solid" + }, + "mode": "lines", + "name": "B", + "showlegend": true, + "type": "scattergl", + "x": [ + "2000-01-01T00:00:00", + "2000-01-02T00:00:00", + "2000-01-03T00:00:00", + "2000-01-04T00:00:00", + "2000-01-05T00:00:00", + "2000-01-06T00:00:00", + "2000-01-07T00:00:00", + "2000-01-08T00:00:00", + "2000-01-09T00:00:00", + "2000-01-10T00:00:00", + "2000-01-11T00:00:00", + "2000-01-12T00:00:00", + "2000-01-13T00:00:00", + "2000-01-14T00:00:00", + "2000-01-15T00:00:00", + "2000-01-16T00:00:00", + "2000-01-17T00:00:00", + "2000-01-18T00:00:00", + "2000-01-19T00:00:00", + "2000-01-20T00:00:00", + "2000-01-21T00:00:00", + "2000-01-22T00:00:00", + "2000-01-23T00:00:00", + "2000-01-24T00:00:00", + "2000-01-25T00:00:00", + "2000-01-26T00:00:00", + "2000-01-27T00:00:00", + "2000-01-28T00:00:00", + "2000-01-29T00:00:00", + "2000-01-30T00:00:00", + "2000-01-31T00:00:00", + "2000-02-01T00:00:00", + "2000-02-02T00:00:00", + "2000-02-03T00:00:00", + "2000-02-04T00:00:00", + "2000-02-05T00:00:00", + "2000-02-06T00:00:00", + "2000-02-07T00:00:00", + "2000-02-08T00:00:00", + "2000-02-09T00:00:00", + "2000-02-10T00:00:00", + "2000-02-11T00:00:00", + "2000-02-12T00:00:00", + "2000-02-13T00:00:00", + "2000-02-14T00:00:00", + "2000-02-15T00:00:00", + "2000-02-16T00:00:00", + "2000-02-17T00:00:00", + "2000-02-18T00:00:00", + "2000-02-19T00:00:00", + "2000-02-20T00:00:00", + "2000-02-21T00:00:00", + "2000-02-22T00:00:00", + "2000-02-23T00:00:00", + "2000-02-24T00:00:00", + "2000-02-25T00:00:00", + "2000-02-26T00:00:00", + "2000-02-27T00:00:00", + "2000-02-28T00:00:00", + "2000-02-29T00:00:00", + "2000-03-01T00:00:00", + "2000-03-02T00:00:00", + "2000-03-03T00:00:00", + "2000-03-04T00:00:00", + "2000-03-05T00:00:00", + "2000-03-06T00:00:00", + "2000-03-07T00:00:00", + "2000-03-08T00:00:00", + "2000-03-09T00:00:00", + "2000-03-10T00:00:00", + "2000-03-11T00:00:00", + "2000-03-12T00:00:00", + "2000-03-13T00:00:00", + "2000-03-14T00:00:00", + "2000-03-15T00:00:00", + "2000-03-16T00:00:00", + "2000-03-17T00:00:00", + "2000-03-18T00:00:00", + "2000-03-19T00:00:00", + "2000-03-20T00:00:00", + "2000-03-21T00:00:00", + "2000-03-22T00:00:00", + "2000-03-23T00:00:00", + "2000-03-24T00:00:00", + "2000-03-25T00:00:00", + "2000-03-26T00:00:00", + "2000-03-27T00:00:00", + "2000-03-28T00:00:00", + "2000-03-29T00:00:00", + "2000-03-30T00:00:00", + "2000-03-31T00:00:00", + "2000-04-01T00:00:00", + "2000-04-02T00:00:00", + "2000-04-03T00:00:00", + "2000-04-04T00:00:00", + "2000-04-05T00:00:00", + "2000-04-06T00:00:00", + "2000-04-07T00:00:00", + "2000-04-08T00:00:00", + "2000-04-09T00:00:00", + "2000-04-10T00:00:00", + "2000-04-11T00:00:00", + "2000-04-12T00:00:00", + "2000-04-13T00:00:00", + "2000-04-14T00:00:00", + "2000-04-15T00:00:00", + "2000-04-16T00:00:00", + "2000-04-17T00:00:00", + "2000-04-18T00:00:00", + "2000-04-19T00:00:00", + "2000-04-20T00:00:00", + "2000-04-21T00:00:00", + "2000-04-22T00:00:00", + "2000-04-23T00:00:00", + "2000-04-24T00:00:00", + "2000-04-25T00:00:00", + "2000-04-26T00:00:00", + "2000-04-27T00:00:00", + "2000-04-28T00:00:00", + "2000-04-29T00:00:00", + "2000-04-30T00:00:00", + "2000-05-01T00:00:00", + "2000-05-02T00:00:00", + "2000-05-03T00:00:00", + "2000-05-04T00:00:00", + "2000-05-05T00:00:00", + "2000-05-06T00:00:00", + "2000-05-07T00:00:00", + "2000-05-08T00:00:00", + "2000-05-09T00:00:00", + "2000-05-10T00:00:00", + "2000-05-11T00:00:00", + "2000-05-12T00:00:00", + "2000-05-13T00:00:00", + "2000-05-14T00:00:00", + "2000-05-15T00:00:00", + "2000-05-16T00:00:00", + "2000-05-17T00:00:00", + "2000-05-18T00:00:00", + "2000-05-19T00:00:00", + "2000-05-20T00:00:00", + "2000-05-21T00:00:00", + "2000-05-22T00:00:00", + "2000-05-23T00:00:00", + "2000-05-24T00:00:00", + "2000-05-25T00:00:00", + "2000-05-26T00:00:00", + "2000-05-27T00:00:00", + "2000-05-28T00:00:00", + "2000-05-29T00:00:00", + "2000-05-30T00:00:00", + "2000-05-31T00:00:00", + "2000-06-01T00:00:00", + "2000-06-02T00:00:00", + "2000-06-03T00:00:00", + "2000-06-04T00:00:00", + "2000-06-05T00:00:00", + "2000-06-06T00:00:00", + "2000-06-07T00:00:00", + "2000-06-08T00:00:00", + "2000-06-09T00:00:00", + "2000-06-10T00:00:00", + "2000-06-11T00:00:00", + "2000-06-12T00:00:00", + "2000-06-13T00:00:00", + "2000-06-14T00:00:00", + "2000-06-15T00:00:00", + "2000-06-16T00:00:00", + "2000-06-17T00:00:00", + "2000-06-18T00:00:00", + "2000-06-19T00:00:00", + "2000-06-20T00:00:00", + "2000-06-21T00:00:00", + "2000-06-22T00:00:00", + "2000-06-23T00:00:00", + "2000-06-24T00:00:00", + "2000-06-25T00:00:00", + "2000-06-26T00:00:00", + "2000-06-27T00:00:00", + "2000-06-28T00:00:00", + "2000-06-29T00:00:00", + "2000-06-30T00:00:00", + "2000-07-01T00:00:00", + "2000-07-02T00:00:00", + "2000-07-03T00:00:00", + "2000-07-04T00:00:00", + "2000-07-05T00:00:00", + "2000-07-06T00:00:00", + "2000-07-07T00:00:00", + "2000-07-08T00:00:00", + "2000-07-09T00:00:00", + "2000-07-10T00:00:00", + "2000-07-11T00:00:00", + "2000-07-12T00:00:00", + "2000-07-13T00:00:00", + "2000-07-14T00:00:00", + "2000-07-15T00:00:00", + "2000-07-16T00:00:00", + "2000-07-17T00:00:00", + "2000-07-18T00:00:00", + "2000-07-19T00:00:00", + "2000-07-20T00:00:00", + "2000-07-21T00:00:00", + "2000-07-22T00:00:00", + "2000-07-23T00:00:00", + "2000-07-24T00:00:00", + "2000-07-25T00:00:00", + "2000-07-26T00:00:00", + "2000-07-27T00:00:00", + "2000-07-28T00:00:00", + "2000-07-29T00:00:00", + "2000-07-30T00:00:00", + "2000-07-31T00:00:00", + "2000-08-01T00:00:00", + "2000-08-02T00:00:00", + "2000-08-03T00:00:00", + "2000-08-04T00:00:00", + "2000-08-05T00:00:00", + "2000-08-06T00:00:00", + "2000-08-07T00:00:00", + "2000-08-08T00:00:00", + "2000-08-09T00:00:00", + "2000-08-10T00:00:00", + "2000-08-11T00:00:00", + "2000-08-12T00:00:00", + "2000-08-13T00:00:00", + "2000-08-14T00:00:00", + "2000-08-15T00:00:00", + "2000-08-16T00:00:00", + "2000-08-17T00:00:00", + "2000-08-18T00:00:00", + "2000-08-19T00:00:00", + "2000-08-20T00:00:00", + "2000-08-21T00:00:00", + "2000-08-22T00:00:00", + "2000-08-23T00:00:00", + "2000-08-24T00:00:00", + "2000-08-25T00:00:00", + "2000-08-26T00:00:00", + "2000-08-27T00:00:00", + "2000-08-28T00:00:00", + "2000-08-29T00:00:00", + "2000-08-30T00:00:00", + "2000-08-31T00:00:00", + "2000-09-01T00:00:00", + "2000-09-02T00:00:00", + "2000-09-03T00:00:00", + "2000-09-04T00:00:00", + "2000-09-05T00:00:00", + "2000-09-06T00:00:00", + "2000-09-07T00:00:00", + "2000-09-08T00:00:00", + "2000-09-09T00:00:00", + "2000-09-10T00:00:00", + "2000-09-11T00:00:00", + "2000-09-12T00:00:00", + "2000-09-13T00:00:00", + "2000-09-14T00:00:00", + "2000-09-15T00:00:00", + "2000-09-16T00:00:00", + "2000-09-17T00:00:00", + "2000-09-18T00:00:00", + "2000-09-19T00:00:00", + "2000-09-20T00:00:00", + "2000-09-21T00:00:00", + "2000-09-22T00:00:00", + "2000-09-23T00:00:00", + "2000-09-24T00:00:00", + "2000-09-25T00:00:00", + "2000-09-26T00:00:00", + "2000-09-27T00:00:00", + "2000-09-28T00:00:00", + "2000-09-29T00:00:00", + "2000-09-30T00:00:00", + "2000-10-01T00:00:00", + "2000-10-02T00:00:00", + "2000-10-03T00:00:00", + "2000-10-04T00:00:00", + "2000-10-05T00:00:00", + "2000-10-06T00:00:00", + "2000-10-07T00:00:00", + "2000-10-08T00:00:00", + "2000-10-09T00:00:00", + "2000-10-10T00:00:00", + "2000-10-11T00:00:00", + "2000-10-12T00:00:00", + "2000-10-13T00:00:00", + "2000-10-14T00:00:00", + "2000-10-15T00:00:00", + "2000-10-16T00:00:00", + "2000-10-17T00:00:00", + "2000-10-18T00:00:00", + "2000-10-19T00:00:00", + "2000-10-20T00:00:00", + "2000-10-21T00:00:00", + "2000-10-22T00:00:00", + "2000-10-23T00:00:00", + "2000-10-24T00:00:00", + "2000-10-25T00:00:00", + "2000-10-26T00:00:00", + "2000-10-27T00:00:00", + "2000-10-28T00:00:00", + "2000-10-29T00:00:00", + "2000-10-30T00:00:00", + "2000-10-31T00:00:00", + "2000-11-01T00:00:00", + "2000-11-02T00:00:00", + "2000-11-03T00:00:00", + "2000-11-04T00:00:00", + "2000-11-05T00:00:00", + "2000-11-06T00:00:00", + "2000-11-07T00:00:00", + "2000-11-08T00:00:00", + "2000-11-09T00:00:00", + "2000-11-10T00:00:00", + "2000-11-11T00:00:00", + "2000-11-12T00:00:00", + "2000-11-13T00:00:00", + "2000-11-14T00:00:00", + "2000-11-15T00:00:00", + "2000-11-16T00:00:00", + "2000-11-17T00:00:00", + "2000-11-18T00:00:00", + "2000-11-19T00:00:00", + "2000-11-20T00:00:00", + "2000-11-21T00:00:00", + "2000-11-22T00:00:00", + "2000-11-23T00:00:00", + "2000-11-24T00:00:00", + "2000-11-25T00:00:00", + "2000-11-26T00:00:00", + "2000-11-27T00:00:00", + "2000-11-28T00:00:00", + "2000-11-29T00:00:00", + "2000-11-30T00:00:00", + "2000-12-01T00:00:00", + "2000-12-02T00:00:00", + "2000-12-03T00:00:00", + "2000-12-04T00:00:00", + "2000-12-05T00:00:00", + "2000-12-06T00:00:00", + "2000-12-07T00:00:00", + "2000-12-08T00:00:00", + "2000-12-09T00:00:00", + "2000-12-10T00:00:00", + "2000-12-11T00:00:00", + "2000-12-12T00:00:00", + "2000-12-13T00:00:00", + "2000-12-14T00:00:00", + "2000-12-15T00:00:00", + "2000-12-16T00:00:00", + "2000-12-17T00:00:00", + "2000-12-18T00:00:00", + "2000-12-19T00:00:00", + "2000-12-20T00:00:00", + "2000-12-21T00:00:00", + "2000-12-22T00:00:00", + "2000-12-23T00:00:00", + "2000-12-24T00:00:00", + "2000-12-25T00:00:00", + "2000-12-26T00:00:00", + "2000-12-27T00:00:00", + "2000-12-28T00:00:00", + "2000-12-29T00:00:00", + "2000-12-30T00:00:00", + "2000-12-31T00:00:00", + "2001-01-01T00:00:00", + "2001-01-02T00:00:00", + "2001-01-03T00:00:00", + "2001-01-04T00:00:00", + "2001-01-05T00:00:00", + "2001-01-06T00:00:00", + "2001-01-07T00:00:00", + "2001-01-08T00:00:00", + "2001-01-09T00:00:00", + "2001-01-10T00:00:00", + "2001-01-11T00:00:00", + "2001-01-12T00:00:00", + "2001-01-13T00:00:00", + "2001-01-14T00:00:00", + "2001-01-15T00:00:00", + "2001-01-16T00:00:00", + "2001-01-17T00:00:00", + "2001-01-18T00:00:00", + "2001-01-19T00:00:00", + "2001-01-20T00:00:00", + "2001-01-21T00:00:00", + "2001-01-22T00:00:00", + "2001-01-23T00:00:00", + "2001-01-24T00:00:00", + "2001-01-25T00:00:00", + "2001-01-26T00:00:00", + "2001-01-27T00:00:00", + "2001-01-28T00:00:00", + "2001-01-29T00:00:00", + "2001-01-30T00:00:00", + "2001-01-31T00:00:00", + "2001-02-01T00:00:00", + "2001-02-02T00:00:00", + "2001-02-03T00:00:00", + "2001-02-04T00:00:00", + "2001-02-05T00:00:00", + "2001-02-06T00:00:00", + "2001-02-07T00:00:00", + "2001-02-08T00:00:00", + "2001-02-09T00:00:00", + "2001-02-10T00:00:00", + "2001-02-11T00:00:00", + "2001-02-12T00:00:00", + "2001-02-13T00:00:00", + "2001-02-14T00:00:00", + "2001-02-15T00:00:00", + "2001-02-16T00:00:00", + "2001-02-17T00:00:00", + "2001-02-18T00:00:00", + "2001-02-19T00:00:00", + "2001-02-20T00:00:00", + "2001-02-21T00:00:00", + "2001-02-22T00:00:00", + "2001-02-23T00:00:00", + "2001-02-24T00:00:00", + "2001-02-25T00:00:00", + "2001-02-26T00:00:00", + "2001-02-27T00:00:00", + "2001-02-28T00:00:00", + "2001-03-01T00:00:00", + "2001-03-02T00:00:00", + "2001-03-03T00:00:00", + "2001-03-04T00:00:00", + "2001-03-05T00:00:00", + "2001-03-06T00:00:00", + "2001-03-07T00:00:00", + "2001-03-08T00:00:00", + "2001-03-09T00:00:00", + "2001-03-10T00:00:00", + "2001-03-11T00:00:00", + "2001-03-12T00:00:00", + "2001-03-13T00:00:00", + "2001-03-14T00:00:00", + "2001-03-15T00:00:00", + "2001-03-16T00:00:00", + "2001-03-17T00:00:00", + "2001-03-18T00:00:00", + "2001-03-19T00:00:00", + "2001-03-20T00:00:00", + "2001-03-21T00:00:00", + "2001-03-22T00:00:00", + "2001-03-23T00:00:00", + "2001-03-24T00:00:00", + "2001-03-25T00:00:00", + "2001-03-26T00:00:00", + "2001-03-27T00:00:00", + "2001-03-28T00:00:00", + "2001-03-29T00:00:00", + "2001-03-30T00:00:00", + "2001-03-31T00:00:00", + "2001-04-01T00:00:00", + "2001-04-02T00:00:00", + "2001-04-03T00:00:00", + "2001-04-04T00:00:00", + "2001-04-05T00:00:00", + "2001-04-06T00:00:00", + "2001-04-07T00:00:00", + "2001-04-08T00:00:00", + "2001-04-09T00:00:00", + "2001-04-10T00:00:00", + "2001-04-11T00:00:00", + "2001-04-12T00:00:00", + "2001-04-13T00:00:00", + "2001-04-14T00:00:00", + "2001-04-15T00:00:00", + "2001-04-16T00:00:00", + "2001-04-17T00:00:00", + "2001-04-18T00:00:00", + "2001-04-19T00:00:00", + "2001-04-20T00:00:00", + "2001-04-21T00:00:00", + "2001-04-22T00:00:00", + "2001-04-23T00:00:00", + "2001-04-24T00:00:00", + "2001-04-25T00:00:00", + "2001-04-26T00:00:00", + "2001-04-27T00:00:00", + "2001-04-28T00:00:00", + "2001-04-29T00:00:00", + "2001-04-30T00:00:00", + "2001-05-01T00:00:00", + "2001-05-02T00:00:00", + "2001-05-03T00:00:00", + "2001-05-04T00:00:00", + "2001-05-05T00:00:00", + "2001-05-06T00:00:00", + "2001-05-07T00:00:00", + "2001-05-08T00:00:00", + "2001-05-09T00:00:00", + "2001-05-10T00:00:00", + "2001-05-11T00:00:00", + "2001-05-12T00:00:00", + "2001-05-13T00:00:00", + "2001-05-14T00:00:00", + "2001-05-15T00:00:00", + "2001-05-16T00:00:00", + "2001-05-17T00:00:00", + "2001-05-18T00:00:00", + "2001-05-19T00:00:00", + "2001-05-20T00:00:00", + "2001-05-21T00:00:00", + "2001-05-22T00:00:00", + "2001-05-23T00:00:00", + "2001-05-24T00:00:00", + "2001-05-25T00:00:00", + "2001-05-26T00:00:00", + "2001-05-27T00:00:00", + "2001-05-28T00:00:00", + "2001-05-29T00:00:00", + "2001-05-30T00:00:00", + "2001-05-31T00:00:00", + "2001-06-01T00:00:00", + "2001-06-02T00:00:00", + "2001-06-03T00:00:00", + "2001-06-04T00:00:00", + "2001-06-05T00:00:00", + "2001-06-06T00:00:00", + "2001-06-07T00:00:00", + "2001-06-08T00:00:00", + "2001-06-09T00:00:00", + "2001-06-10T00:00:00", + "2001-06-11T00:00:00", + "2001-06-12T00:00:00", + "2001-06-13T00:00:00", + "2001-06-14T00:00:00", + "2001-06-15T00:00:00", + "2001-06-16T00:00:00", + "2001-06-17T00:00:00", + "2001-06-18T00:00:00", + "2001-06-19T00:00:00", + "2001-06-20T00:00:00", + "2001-06-21T00:00:00", + "2001-06-22T00:00:00", + "2001-06-23T00:00:00", + "2001-06-24T00:00:00", + "2001-06-25T00:00:00", + "2001-06-26T00:00:00", + "2001-06-27T00:00:00", + "2001-06-28T00:00:00", + "2001-06-29T00:00:00", + "2001-06-30T00:00:00", + "2001-07-01T00:00:00", + "2001-07-02T00:00:00", + "2001-07-03T00:00:00", + "2001-07-04T00:00:00", + "2001-07-05T00:00:00", + "2001-07-06T00:00:00", + "2001-07-07T00:00:00", + "2001-07-08T00:00:00", + "2001-07-09T00:00:00", + "2001-07-10T00:00:00", + "2001-07-11T00:00:00", + "2001-07-12T00:00:00", + "2001-07-13T00:00:00", + "2001-07-14T00:00:00", + "2001-07-15T00:00:00", + "2001-07-16T00:00:00", + "2001-07-17T00:00:00", + "2001-07-18T00:00:00", + "2001-07-19T00:00:00", + "2001-07-20T00:00:00", + "2001-07-21T00:00:00", + "2001-07-22T00:00:00", + "2001-07-23T00:00:00", + "2001-07-24T00:00:00", + "2001-07-25T00:00:00", + "2001-07-26T00:00:00", + "2001-07-27T00:00:00", + "2001-07-28T00:00:00", + "2001-07-29T00:00:00", + "2001-07-30T00:00:00", + "2001-07-31T00:00:00", + "2001-08-01T00:00:00", + "2001-08-02T00:00:00", + "2001-08-03T00:00:00", + "2001-08-04T00:00:00", + "2001-08-05T00:00:00", + "2001-08-06T00:00:00", + "2001-08-07T00:00:00", + "2001-08-08T00:00:00", + "2001-08-09T00:00:00", + "2001-08-10T00:00:00", + "2001-08-11T00:00:00", + "2001-08-12T00:00:00", + "2001-08-13T00:00:00", + "2001-08-14T00:00:00", + "2001-08-15T00:00:00", + "2001-08-16T00:00:00", + "2001-08-17T00:00:00", + "2001-08-18T00:00:00", + "2001-08-19T00:00:00", + "2001-08-20T00:00:00", + "2001-08-21T00:00:00", + "2001-08-22T00:00:00", + "2001-08-23T00:00:00", + "2001-08-24T00:00:00", + "2001-08-25T00:00:00", + "2001-08-26T00:00:00", + "2001-08-27T00:00:00", + "2001-08-28T00:00:00", + "2001-08-29T00:00:00", + "2001-08-30T00:00:00", + "2001-08-31T00:00:00", + "2001-09-01T00:00:00", + "2001-09-02T00:00:00", + "2001-09-03T00:00:00", + "2001-09-04T00:00:00", + "2001-09-05T00:00:00", + "2001-09-06T00:00:00", + "2001-09-07T00:00:00", + "2001-09-08T00:00:00", + "2001-09-09T00:00:00", + "2001-09-10T00:00:00", + "2001-09-11T00:00:00", + "2001-09-12T00:00:00", + "2001-09-13T00:00:00", + "2001-09-14T00:00:00", + "2001-09-15T00:00:00", + "2001-09-16T00:00:00", + "2001-09-17T00:00:00", + "2001-09-18T00:00:00", + "2001-09-19T00:00:00", + "2001-09-20T00:00:00", + "2001-09-21T00:00:00", + "2001-09-22T00:00:00", + "2001-09-23T00:00:00", + "2001-09-24T00:00:00", + "2001-09-25T00:00:00", + "2001-09-26T00:00:00", + "2001-09-27T00:00:00", + "2001-09-28T00:00:00", + "2001-09-29T00:00:00", + "2001-09-30T00:00:00", + "2001-10-01T00:00:00", + "2001-10-02T00:00:00", + "2001-10-03T00:00:00", + "2001-10-04T00:00:00", + "2001-10-05T00:00:00", + "2001-10-06T00:00:00", + "2001-10-07T00:00:00", + "2001-10-08T00:00:00", + "2001-10-09T00:00:00", + "2001-10-10T00:00:00", + "2001-10-11T00:00:00", + "2001-10-12T00:00:00", + "2001-10-13T00:00:00", + "2001-10-14T00:00:00", + "2001-10-15T00:00:00", + "2001-10-16T00:00:00", + "2001-10-17T00:00:00", + "2001-10-18T00:00:00", + "2001-10-19T00:00:00", + "2001-10-20T00:00:00", + "2001-10-21T00:00:00", + "2001-10-22T00:00:00", + "2001-10-23T00:00:00", + "2001-10-24T00:00:00", + "2001-10-25T00:00:00", + "2001-10-26T00:00:00", + "2001-10-27T00:00:00", + "2001-10-28T00:00:00", + "2001-10-29T00:00:00", + "2001-10-30T00:00:00", + "2001-10-31T00:00:00", + "2001-11-01T00:00:00", + "2001-11-02T00:00:00", + "2001-11-03T00:00:00", + "2001-11-04T00:00:00", + "2001-11-05T00:00:00", + "2001-11-06T00:00:00", + "2001-11-07T00:00:00", + "2001-11-08T00:00:00", + "2001-11-09T00:00:00", + "2001-11-10T00:00:00", + "2001-11-11T00:00:00", + "2001-11-12T00:00:00", + "2001-11-13T00:00:00", + "2001-11-14T00:00:00", + "2001-11-15T00:00:00", + "2001-11-16T00:00:00", + "2001-11-17T00:00:00", + "2001-11-18T00:00:00", + "2001-11-19T00:00:00", + "2001-11-20T00:00:00", + "2001-11-21T00:00:00", + "2001-11-22T00:00:00", + "2001-11-23T00:00:00", + "2001-11-24T00:00:00", + "2001-11-25T00:00:00", + "2001-11-26T00:00:00", + "2001-11-27T00:00:00", + "2001-11-28T00:00:00", + "2001-11-29T00:00:00", + "2001-11-30T00:00:00", + "2001-12-01T00:00:00", + "2001-12-02T00:00:00", + "2001-12-03T00:00:00", + "2001-12-04T00:00:00", + "2001-12-05T00:00:00", + "2001-12-06T00:00:00", + "2001-12-07T00:00:00", + "2001-12-08T00:00:00", + "2001-12-09T00:00:00", + "2001-12-10T00:00:00", + "2001-12-11T00:00:00", + "2001-12-12T00:00:00", + "2001-12-13T00:00:00", + "2001-12-14T00:00:00", + "2001-12-15T00:00:00", + "2001-12-16T00:00:00", + "2001-12-17T00:00:00", + "2001-12-18T00:00:00", + "2001-12-19T00:00:00", + "2001-12-20T00:00:00", + "2001-12-21T00:00:00", + "2001-12-22T00:00:00", + "2001-12-23T00:00:00", + "2001-12-24T00:00:00", + "2001-12-25T00:00:00", + "2001-12-26T00:00:00", + "2001-12-27T00:00:00", + "2001-12-28T00:00:00", + "2001-12-29T00:00:00", + "2001-12-30T00:00:00", + "2001-12-31T00:00:00", + "2002-01-01T00:00:00", + "2002-01-02T00:00:00", + "2002-01-03T00:00:00", + "2002-01-04T00:00:00", + "2002-01-05T00:00:00", + "2002-01-06T00:00:00", + "2002-01-07T00:00:00", + "2002-01-08T00:00:00", + "2002-01-09T00:00:00", + "2002-01-10T00:00:00", + "2002-01-11T00:00:00", + "2002-01-12T00:00:00", + "2002-01-13T00:00:00", + "2002-01-14T00:00:00", + "2002-01-15T00:00:00", + "2002-01-16T00:00:00", + "2002-01-17T00:00:00", + "2002-01-18T00:00:00", + "2002-01-19T00:00:00", + "2002-01-20T00:00:00", + "2002-01-21T00:00:00", + "2002-01-22T00:00:00", + "2002-01-23T00:00:00", + "2002-01-24T00:00:00", + "2002-01-25T00:00:00", + "2002-01-26T00:00:00", + "2002-01-27T00:00:00", + "2002-01-28T00:00:00", + "2002-01-29T00:00:00", + "2002-01-30T00:00:00", + "2002-01-31T00:00:00", + "2002-02-01T00:00:00", + "2002-02-02T00:00:00", + "2002-02-03T00:00:00", + "2002-02-04T00:00:00", + "2002-02-05T00:00:00", + "2002-02-06T00:00:00", + "2002-02-07T00:00:00", + "2002-02-08T00:00:00", + "2002-02-09T00:00:00", + "2002-02-10T00:00:00", + "2002-02-11T00:00:00", + "2002-02-12T00:00:00", + "2002-02-13T00:00:00", + "2002-02-14T00:00:00", + "2002-02-15T00:00:00", + "2002-02-16T00:00:00", + "2002-02-17T00:00:00", + "2002-02-18T00:00:00", + "2002-02-19T00:00:00", + "2002-02-20T00:00:00", + "2002-02-21T00:00:00", + "2002-02-22T00:00:00", + "2002-02-23T00:00:00", + "2002-02-24T00:00:00", + "2002-02-25T00:00:00", + "2002-02-26T00:00:00", + "2002-02-27T00:00:00", + "2002-02-28T00:00:00", + "2002-03-01T00:00:00", + "2002-03-02T00:00:00", + "2002-03-03T00:00:00", + "2002-03-04T00:00:00", + "2002-03-05T00:00:00", + "2002-03-06T00:00:00", + "2002-03-07T00:00:00", + "2002-03-08T00:00:00", + "2002-03-09T00:00:00", + "2002-03-10T00:00:00", + "2002-03-11T00:00:00", + "2002-03-12T00:00:00", + "2002-03-13T00:00:00", + "2002-03-14T00:00:00", + "2002-03-15T00:00:00", + "2002-03-16T00:00:00", + "2002-03-17T00:00:00", + "2002-03-18T00:00:00", + "2002-03-19T00:00:00", + "2002-03-20T00:00:00", + "2002-03-21T00:00:00", + "2002-03-22T00:00:00", + "2002-03-23T00:00:00", + "2002-03-24T00:00:00", + "2002-03-25T00:00:00", + "2002-03-26T00:00:00", + "2002-03-27T00:00:00", + "2002-03-28T00:00:00", + "2002-03-29T00:00:00", + "2002-03-30T00:00:00", + "2002-03-31T00:00:00", + "2002-04-01T00:00:00", + "2002-04-02T00:00:00", + "2002-04-03T00:00:00", + "2002-04-04T00:00:00", + "2002-04-05T00:00:00", + "2002-04-06T00:00:00", + "2002-04-07T00:00:00", + "2002-04-08T00:00:00", + "2002-04-09T00:00:00", + "2002-04-10T00:00:00", + "2002-04-11T00:00:00", + "2002-04-12T00:00:00", + "2002-04-13T00:00:00", + "2002-04-14T00:00:00", + "2002-04-15T00:00:00", + "2002-04-16T00:00:00", + "2002-04-17T00:00:00", + "2002-04-18T00:00:00", + "2002-04-19T00:00:00", + "2002-04-20T00:00:00", + "2002-04-21T00:00:00", + "2002-04-22T00:00:00", + "2002-04-23T00:00:00", + "2002-04-24T00:00:00", + "2002-04-25T00:00:00", + "2002-04-26T00:00:00", + "2002-04-27T00:00:00", + "2002-04-28T00:00:00", + "2002-04-29T00:00:00", + "2002-04-30T00:00:00", + "2002-05-01T00:00:00", + "2002-05-02T00:00:00", + "2002-05-03T00:00:00", + "2002-05-04T00:00:00", + "2002-05-05T00:00:00", + "2002-05-06T00:00:00", + "2002-05-07T00:00:00", + "2002-05-08T00:00:00", + "2002-05-09T00:00:00", + "2002-05-10T00:00:00", + "2002-05-11T00:00:00", + "2002-05-12T00:00:00", + "2002-05-13T00:00:00", + "2002-05-14T00:00:00", + "2002-05-15T00:00:00", + "2002-05-16T00:00:00", + "2002-05-17T00:00:00", + "2002-05-18T00:00:00", + "2002-05-19T00:00:00", + "2002-05-20T00:00:00", + "2002-05-21T00:00:00", + "2002-05-22T00:00:00", + "2002-05-23T00:00:00", + "2002-05-24T00:00:00", + "2002-05-25T00:00:00", + "2002-05-26T00:00:00", + "2002-05-27T00:00:00", + "2002-05-28T00:00:00", + "2002-05-29T00:00:00", + "2002-05-30T00:00:00", + "2002-05-31T00:00:00", + "2002-06-01T00:00:00", + "2002-06-02T00:00:00", + "2002-06-03T00:00:00", + "2002-06-04T00:00:00", + "2002-06-05T00:00:00", + "2002-06-06T00:00:00", + "2002-06-07T00:00:00", + "2002-06-08T00:00:00", + "2002-06-09T00:00:00", + "2002-06-10T00:00:00", + "2002-06-11T00:00:00", + "2002-06-12T00:00:00", + "2002-06-13T00:00:00", + "2002-06-14T00:00:00", + "2002-06-15T00:00:00", + "2002-06-16T00:00:00", + "2002-06-17T00:00:00", + "2002-06-18T00:00:00", + "2002-06-19T00:00:00", + "2002-06-20T00:00:00", + "2002-06-21T00:00:00", + "2002-06-22T00:00:00", + "2002-06-23T00:00:00", + "2002-06-24T00:00:00", + "2002-06-25T00:00:00", + "2002-06-26T00:00:00", + "2002-06-27T00:00:00", + "2002-06-28T00:00:00", + "2002-06-29T00:00:00", + "2002-06-30T00:00:00", + "2002-07-01T00:00:00", + "2002-07-02T00:00:00", + "2002-07-03T00:00:00", + "2002-07-04T00:00:00", + "2002-07-05T00:00:00", + "2002-07-06T00:00:00", + "2002-07-07T00:00:00", + "2002-07-08T00:00:00", + "2002-07-09T00:00:00", + "2002-07-10T00:00:00", + "2002-07-11T00:00:00", + "2002-07-12T00:00:00", + "2002-07-13T00:00:00", + "2002-07-14T00:00:00", + "2002-07-15T00:00:00", + "2002-07-16T00:00:00", + "2002-07-17T00:00:00", + "2002-07-18T00:00:00", + "2002-07-19T00:00:00", + "2002-07-20T00:00:00", + "2002-07-21T00:00:00", + "2002-07-22T00:00:00", + "2002-07-23T00:00:00", + "2002-07-24T00:00:00", + "2002-07-25T00:00:00", + "2002-07-26T00:00:00", + "2002-07-27T00:00:00", + "2002-07-28T00:00:00", + "2002-07-29T00:00:00", + "2002-07-30T00:00:00", + "2002-07-31T00:00:00", + "2002-08-01T00:00:00", + "2002-08-02T00:00:00", + "2002-08-03T00:00:00", + "2002-08-04T00:00:00", + "2002-08-05T00:00:00", + "2002-08-06T00:00:00", + "2002-08-07T00:00:00", + "2002-08-08T00:00:00", + "2002-08-09T00:00:00", + "2002-08-10T00:00:00", + "2002-08-11T00:00:00", + "2002-08-12T00:00:00", + "2002-08-13T00:00:00", + "2002-08-14T00:00:00", + "2002-08-15T00:00:00", + "2002-08-16T00:00:00", + "2002-08-17T00:00:00", + "2002-08-18T00:00:00", + "2002-08-19T00:00:00", + "2002-08-20T00:00:00", + "2002-08-21T00:00:00", + "2002-08-22T00:00:00", + "2002-08-23T00:00:00", + "2002-08-24T00:00:00", + "2002-08-25T00:00:00", + "2002-08-26T00:00:00", + "2002-08-27T00:00:00", + "2002-08-28T00:00:00", + "2002-08-29T00:00:00", + "2002-08-30T00:00:00", + "2002-08-31T00:00:00", + "2002-09-01T00:00:00", + "2002-09-02T00:00:00", + "2002-09-03T00:00:00", + "2002-09-04T00:00:00", + "2002-09-05T00:00:00", + "2002-09-06T00:00:00", + "2002-09-07T00:00:00", + "2002-09-08T00:00:00", + "2002-09-09T00:00:00", + "2002-09-10T00:00:00", + "2002-09-11T00:00:00", + "2002-09-12T00:00:00", + "2002-09-13T00:00:00", + "2002-09-14T00:00:00", + "2002-09-15T00:00:00", + "2002-09-16T00:00:00", + "2002-09-17T00:00:00", + "2002-09-18T00:00:00", + "2002-09-19T00:00:00", + "2002-09-20T00:00:00", + "2002-09-21T00:00:00", + "2002-09-22T00:00:00", + "2002-09-23T00:00:00", + "2002-09-24T00:00:00", + "2002-09-25T00:00:00", + "2002-09-26T00:00:00" + ], + "xaxis": "x", + "y": [ + -0.1346448450209322, + 0.16621832502053047, + 0.43945019008041913, + 0.43945019008041913, + 0.43945019008041913, + 1.069183120163687, + 1.069183120163687, + 1.069183120163687, + 1.069183120163687, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 1.5087320230682002, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.017421403368495, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.0981239060579906, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.483775070915523, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.710882653336446, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 2.8233161008927845, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.02876548148906, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671, + 3.064434309818671 + ], + "yaxis": "y" + }, + { + "hovertemplate": "variable=C
    index=%{x}
    value=%{y}", + "legendgroup": "C", + "line": { + "color": "#00cc96", + "dash": "solid" + }, + "mode": "lines", + "name": "C", + "showlegend": true, + "type": "scattergl", + "x": [ + "2000-01-01T00:00:00", + "2000-01-02T00:00:00", + "2000-01-03T00:00:00", + "2000-01-04T00:00:00", + "2000-01-05T00:00:00", + "2000-01-06T00:00:00", + "2000-01-07T00:00:00", + "2000-01-08T00:00:00", + "2000-01-09T00:00:00", + "2000-01-10T00:00:00", + "2000-01-11T00:00:00", + "2000-01-12T00:00:00", + "2000-01-13T00:00:00", + "2000-01-14T00:00:00", + "2000-01-15T00:00:00", + "2000-01-16T00:00:00", + "2000-01-17T00:00:00", + "2000-01-18T00:00:00", + "2000-01-19T00:00:00", + "2000-01-20T00:00:00", + "2000-01-21T00:00:00", + "2000-01-22T00:00:00", + "2000-01-23T00:00:00", + "2000-01-24T00:00:00", + "2000-01-25T00:00:00", + "2000-01-26T00:00:00", + "2000-01-27T00:00:00", + "2000-01-28T00:00:00", + "2000-01-29T00:00:00", + "2000-01-30T00:00:00", + "2000-01-31T00:00:00", + "2000-02-01T00:00:00", + "2000-02-02T00:00:00", + "2000-02-03T00:00:00", + "2000-02-04T00:00:00", + "2000-02-05T00:00:00", + "2000-02-06T00:00:00", + "2000-02-07T00:00:00", + "2000-02-08T00:00:00", + "2000-02-09T00:00:00", + "2000-02-10T00:00:00", + "2000-02-11T00:00:00", + "2000-02-12T00:00:00", + "2000-02-13T00:00:00", + "2000-02-14T00:00:00", + "2000-02-15T00:00:00", + "2000-02-16T00:00:00", + "2000-02-17T00:00:00", + "2000-02-18T00:00:00", + "2000-02-19T00:00:00", + "2000-02-20T00:00:00", + "2000-02-21T00:00:00", + "2000-02-22T00:00:00", + "2000-02-23T00:00:00", + "2000-02-24T00:00:00", + "2000-02-25T00:00:00", + "2000-02-26T00:00:00", + "2000-02-27T00:00:00", + "2000-02-28T00:00:00", + "2000-02-29T00:00:00", + "2000-03-01T00:00:00", + "2000-03-02T00:00:00", + "2000-03-03T00:00:00", + "2000-03-04T00:00:00", + "2000-03-05T00:00:00", + "2000-03-06T00:00:00", + "2000-03-07T00:00:00", + "2000-03-08T00:00:00", + "2000-03-09T00:00:00", + "2000-03-10T00:00:00", + "2000-03-11T00:00:00", + "2000-03-12T00:00:00", + "2000-03-13T00:00:00", + "2000-03-14T00:00:00", + "2000-03-15T00:00:00", + "2000-03-16T00:00:00", + "2000-03-17T00:00:00", + "2000-03-18T00:00:00", + "2000-03-19T00:00:00", + "2000-03-20T00:00:00", + "2000-03-21T00:00:00", + "2000-03-22T00:00:00", + "2000-03-23T00:00:00", + "2000-03-24T00:00:00", + "2000-03-25T00:00:00", + "2000-03-26T00:00:00", + "2000-03-27T00:00:00", + "2000-03-28T00:00:00", + "2000-03-29T00:00:00", + "2000-03-30T00:00:00", + "2000-03-31T00:00:00", + "2000-04-01T00:00:00", + "2000-04-02T00:00:00", + "2000-04-03T00:00:00", + "2000-04-04T00:00:00", + "2000-04-05T00:00:00", + "2000-04-06T00:00:00", + "2000-04-07T00:00:00", + "2000-04-08T00:00:00", + "2000-04-09T00:00:00", + "2000-04-10T00:00:00", + "2000-04-11T00:00:00", + "2000-04-12T00:00:00", + "2000-04-13T00:00:00", + "2000-04-14T00:00:00", + "2000-04-15T00:00:00", + "2000-04-16T00:00:00", + "2000-04-17T00:00:00", + "2000-04-18T00:00:00", + "2000-04-19T00:00:00", + "2000-04-20T00:00:00", + "2000-04-21T00:00:00", + "2000-04-22T00:00:00", + "2000-04-23T00:00:00", + "2000-04-24T00:00:00", + "2000-04-25T00:00:00", + "2000-04-26T00:00:00", + "2000-04-27T00:00:00", + "2000-04-28T00:00:00", + "2000-04-29T00:00:00", + "2000-04-30T00:00:00", + "2000-05-01T00:00:00", + "2000-05-02T00:00:00", + "2000-05-03T00:00:00", + "2000-05-04T00:00:00", + "2000-05-05T00:00:00", + "2000-05-06T00:00:00", + "2000-05-07T00:00:00", + "2000-05-08T00:00:00", + "2000-05-09T00:00:00", + "2000-05-10T00:00:00", + "2000-05-11T00:00:00", + "2000-05-12T00:00:00", + "2000-05-13T00:00:00", + "2000-05-14T00:00:00", + "2000-05-15T00:00:00", + "2000-05-16T00:00:00", + "2000-05-17T00:00:00", + "2000-05-18T00:00:00", + "2000-05-19T00:00:00", + "2000-05-20T00:00:00", + "2000-05-21T00:00:00", + "2000-05-22T00:00:00", + "2000-05-23T00:00:00", + "2000-05-24T00:00:00", + "2000-05-25T00:00:00", + "2000-05-26T00:00:00", + "2000-05-27T00:00:00", + "2000-05-28T00:00:00", + "2000-05-29T00:00:00", + "2000-05-30T00:00:00", + "2000-05-31T00:00:00", + "2000-06-01T00:00:00", + "2000-06-02T00:00:00", + "2000-06-03T00:00:00", + "2000-06-04T00:00:00", + "2000-06-05T00:00:00", + "2000-06-06T00:00:00", + "2000-06-07T00:00:00", + "2000-06-08T00:00:00", + "2000-06-09T00:00:00", + "2000-06-10T00:00:00", + "2000-06-11T00:00:00", + "2000-06-12T00:00:00", + "2000-06-13T00:00:00", + "2000-06-14T00:00:00", + "2000-06-15T00:00:00", + "2000-06-16T00:00:00", + "2000-06-17T00:00:00", + "2000-06-18T00:00:00", + "2000-06-19T00:00:00", + "2000-06-20T00:00:00", + "2000-06-21T00:00:00", + "2000-06-22T00:00:00", + "2000-06-23T00:00:00", + "2000-06-24T00:00:00", + "2000-06-25T00:00:00", + "2000-06-26T00:00:00", + "2000-06-27T00:00:00", + "2000-06-28T00:00:00", + "2000-06-29T00:00:00", + "2000-06-30T00:00:00", + "2000-07-01T00:00:00", + "2000-07-02T00:00:00", + "2000-07-03T00:00:00", + "2000-07-04T00:00:00", + "2000-07-05T00:00:00", + "2000-07-06T00:00:00", + "2000-07-07T00:00:00", + "2000-07-08T00:00:00", + "2000-07-09T00:00:00", + "2000-07-10T00:00:00", + "2000-07-11T00:00:00", + "2000-07-12T00:00:00", + "2000-07-13T00:00:00", + "2000-07-14T00:00:00", + "2000-07-15T00:00:00", + "2000-07-16T00:00:00", + "2000-07-17T00:00:00", + "2000-07-18T00:00:00", + "2000-07-19T00:00:00", + "2000-07-20T00:00:00", + "2000-07-21T00:00:00", + "2000-07-22T00:00:00", + "2000-07-23T00:00:00", + "2000-07-24T00:00:00", + "2000-07-25T00:00:00", + "2000-07-26T00:00:00", + "2000-07-27T00:00:00", + "2000-07-28T00:00:00", + "2000-07-29T00:00:00", + "2000-07-30T00:00:00", + "2000-07-31T00:00:00", + "2000-08-01T00:00:00", + "2000-08-02T00:00:00", + "2000-08-03T00:00:00", + "2000-08-04T00:00:00", + "2000-08-05T00:00:00", + "2000-08-06T00:00:00", + "2000-08-07T00:00:00", + "2000-08-08T00:00:00", + "2000-08-09T00:00:00", + "2000-08-10T00:00:00", + "2000-08-11T00:00:00", + "2000-08-12T00:00:00", + "2000-08-13T00:00:00", + "2000-08-14T00:00:00", + "2000-08-15T00:00:00", + "2000-08-16T00:00:00", + "2000-08-17T00:00:00", + "2000-08-18T00:00:00", + "2000-08-19T00:00:00", + "2000-08-20T00:00:00", + "2000-08-21T00:00:00", + "2000-08-22T00:00:00", + "2000-08-23T00:00:00", + "2000-08-24T00:00:00", + "2000-08-25T00:00:00", + "2000-08-26T00:00:00", + "2000-08-27T00:00:00", + "2000-08-28T00:00:00", + "2000-08-29T00:00:00", + "2000-08-30T00:00:00", + "2000-08-31T00:00:00", + "2000-09-01T00:00:00", + "2000-09-02T00:00:00", + "2000-09-03T00:00:00", + "2000-09-04T00:00:00", + "2000-09-05T00:00:00", + "2000-09-06T00:00:00", + "2000-09-07T00:00:00", + "2000-09-08T00:00:00", + "2000-09-09T00:00:00", + "2000-09-10T00:00:00", + "2000-09-11T00:00:00", + "2000-09-12T00:00:00", + "2000-09-13T00:00:00", + "2000-09-14T00:00:00", + "2000-09-15T00:00:00", + "2000-09-16T00:00:00", + "2000-09-17T00:00:00", + "2000-09-18T00:00:00", + "2000-09-19T00:00:00", + "2000-09-20T00:00:00", + "2000-09-21T00:00:00", + "2000-09-22T00:00:00", + "2000-09-23T00:00:00", + "2000-09-24T00:00:00", + "2000-09-25T00:00:00", + "2000-09-26T00:00:00", + "2000-09-27T00:00:00", + "2000-09-28T00:00:00", + "2000-09-29T00:00:00", + "2000-09-30T00:00:00", + "2000-10-01T00:00:00", + "2000-10-02T00:00:00", + "2000-10-03T00:00:00", + "2000-10-04T00:00:00", + "2000-10-05T00:00:00", + "2000-10-06T00:00:00", + "2000-10-07T00:00:00", + "2000-10-08T00:00:00", + "2000-10-09T00:00:00", + "2000-10-10T00:00:00", + "2000-10-11T00:00:00", + "2000-10-12T00:00:00", + "2000-10-13T00:00:00", + "2000-10-14T00:00:00", + "2000-10-15T00:00:00", + "2000-10-16T00:00:00", + "2000-10-17T00:00:00", + "2000-10-18T00:00:00", + "2000-10-19T00:00:00", + "2000-10-20T00:00:00", + "2000-10-21T00:00:00", + "2000-10-22T00:00:00", + "2000-10-23T00:00:00", + "2000-10-24T00:00:00", + "2000-10-25T00:00:00", + "2000-10-26T00:00:00", + "2000-10-27T00:00:00", + "2000-10-28T00:00:00", + "2000-10-29T00:00:00", + "2000-10-30T00:00:00", + "2000-10-31T00:00:00", + "2000-11-01T00:00:00", + "2000-11-02T00:00:00", + "2000-11-03T00:00:00", + "2000-11-04T00:00:00", + "2000-11-05T00:00:00", + "2000-11-06T00:00:00", + "2000-11-07T00:00:00", + "2000-11-08T00:00:00", + "2000-11-09T00:00:00", + "2000-11-10T00:00:00", + "2000-11-11T00:00:00", + "2000-11-12T00:00:00", + "2000-11-13T00:00:00", + "2000-11-14T00:00:00", + "2000-11-15T00:00:00", + "2000-11-16T00:00:00", + "2000-11-17T00:00:00", + "2000-11-18T00:00:00", + "2000-11-19T00:00:00", + "2000-11-20T00:00:00", + "2000-11-21T00:00:00", + "2000-11-22T00:00:00", + "2000-11-23T00:00:00", + "2000-11-24T00:00:00", + "2000-11-25T00:00:00", + "2000-11-26T00:00:00", + "2000-11-27T00:00:00", + "2000-11-28T00:00:00", + "2000-11-29T00:00:00", + "2000-11-30T00:00:00", + "2000-12-01T00:00:00", + "2000-12-02T00:00:00", + "2000-12-03T00:00:00", + "2000-12-04T00:00:00", + "2000-12-05T00:00:00", + "2000-12-06T00:00:00", + "2000-12-07T00:00:00", + "2000-12-08T00:00:00", + "2000-12-09T00:00:00", + "2000-12-10T00:00:00", + "2000-12-11T00:00:00", + "2000-12-12T00:00:00", + "2000-12-13T00:00:00", + "2000-12-14T00:00:00", + "2000-12-15T00:00:00", + "2000-12-16T00:00:00", + "2000-12-17T00:00:00", + "2000-12-18T00:00:00", + "2000-12-19T00:00:00", + "2000-12-20T00:00:00", + "2000-12-21T00:00:00", + "2000-12-22T00:00:00", + "2000-12-23T00:00:00", + "2000-12-24T00:00:00", + "2000-12-25T00:00:00", + "2000-12-26T00:00:00", + "2000-12-27T00:00:00", + "2000-12-28T00:00:00", + "2000-12-29T00:00:00", + "2000-12-30T00:00:00", + "2000-12-31T00:00:00", + "2001-01-01T00:00:00", + "2001-01-02T00:00:00", + "2001-01-03T00:00:00", + "2001-01-04T00:00:00", + "2001-01-05T00:00:00", + "2001-01-06T00:00:00", + "2001-01-07T00:00:00", + "2001-01-08T00:00:00", + "2001-01-09T00:00:00", + "2001-01-10T00:00:00", + "2001-01-11T00:00:00", + "2001-01-12T00:00:00", + "2001-01-13T00:00:00", + "2001-01-14T00:00:00", + "2001-01-15T00:00:00", + "2001-01-16T00:00:00", + "2001-01-17T00:00:00", + "2001-01-18T00:00:00", + "2001-01-19T00:00:00", + "2001-01-20T00:00:00", + "2001-01-21T00:00:00", + "2001-01-22T00:00:00", + "2001-01-23T00:00:00", + "2001-01-24T00:00:00", + "2001-01-25T00:00:00", + "2001-01-26T00:00:00", + "2001-01-27T00:00:00", + "2001-01-28T00:00:00", + "2001-01-29T00:00:00", + "2001-01-30T00:00:00", + "2001-01-31T00:00:00", + "2001-02-01T00:00:00", + "2001-02-02T00:00:00", + "2001-02-03T00:00:00", + "2001-02-04T00:00:00", + "2001-02-05T00:00:00", + "2001-02-06T00:00:00", + "2001-02-07T00:00:00", + "2001-02-08T00:00:00", + "2001-02-09T00:00:00", + "2001-02-10T00:00:00", + "2001-02-11T00:00:00", + "2001-02-12T00:00:00", + "2001-02-13T00:00:00", + "2001-02-14T00:00:00", + "2001-02-15T00:00:00", + "2001-02-16T00:00:00", + "2001-02-17T00:00:00", + "2001-02-18T00:00:00", + "2001-02-19T00:00:00", + "2001-02-20T00:00:00", + "2001-02-21T00:00:00", + "2001-02-22T00:00:00", + "2001-02-23T00:00:00", + "2001-02-24T00:00:00", + "2001-02-25T00:00:00", + "2001-02-26T00:00:00", + "2001-02-27T00:00:00", + "2001-02-28T00:00:00", + "2001-03-01T00:00:00", + "2001-03-02T00:00:00", + "2001-03-03T00:00:00", + "2001-03-04T00:00:00", + "2001-03-05T00:00:00", + "2001-03-06T00:00:00", + "2001-03-07T00:00:00", + "2001-03-08T00:00:00", + "2001-03-09T00:00:00", + "2001-03-10T00:00:00", + "2001-03-11T00:00:00", + "2001-03-12T00:00:00", + "2001-03-13T00:00:00", + "2001-03-14T00:00:00", + "2001-03-15T00:00:00", + "2001-03-16T00:00:00", + "2001-03-17T00:00:00", + "2001-03-18T00:00:00", + "2001-03-19T00:00:00", + "2001-03-20T00:00:00", + "2001-03-21T00:00:00", + "2001-03-22T00:00:00", + "2001-03-23T00:00:00", + "2001-03-24T00:00:00", + "2001-03-25T00:00:00", + "2001-03-26T00:00:00", + "2001-03-27T00:00:00", + "2001-03-28T00:00:00", + "2001-03-29T00:00:00", + "2001-03-30T00:00:00", + "2001-03-31T00:00:00", + "2001-04-01T00:00:00", + "2001-04-02T00:00:00", + "2001-04-03T00:00:00", + "2001-04-04T00:00:00", + "2001-04-05T00:00:00", + "2001-04-06T00:00:00", + "2001-04-07T00:00:00", + "2001-04-08T00:00:00", + "2001-04-09T00:00:00", + "2001-04-10T00:00:00", + "2001-04-11T00:00:00", + "2001-04-12T00:00:00", + "2001-04-13T00:00:00", + "2001-04-14T00:00:00", + "2001-04-15T00:00:00", + "2001-04-16T00:00:00", + "2001-04-17T00:00:00", + "2001-04-18T00:00:00", + "2001-04-19T00:00:00", + "2001-04-20T00:00:00", + "2001-04-21T00:00:00", + "2001-04-22T00:00:00", + "2001-04-23T00:00:00", + "2001-04-24T00:00:00", + "2001-04-25T00:00:00", + "2001-04-26T00:00:00", + "2001-04-27T00:00:00", + "2001-04-28T00:00:00", + "2001-04-29T00:00:00", + "2001-04-30T00:00:00", + "2001-05-01T00:00:00", + "2001-05-02T00:00:00", + "2001-05-03T00:00:00", + "2001-05-04T00:00:00", + "2001-05-05T00:00:00", + "2001-05-06T00:00:00", + "2001-05-07T00:00:00", + "2001-05-08T00:00:00", + "2001-05-09T00:00:00", + "2001-05-10T00:00:00", + "2001-05-11T00:00:00", + "2001-05-12T00:00:00", + "2001-05-13T00:00:00", + "2001-05-14T00:00:00", + "2001-05-15T00:00:00", + "2001-05-16T00:00:00", + "2001-05-17T00:00:00", + "2001-05-18T00:00:00", + "2001-05-19T00:00:00", + "2001-05-20T00:00:00", + "2001-05-21T00:00:00", + "2001-05-22T00:00:00", + "2001-05-23T00:00:00", + "2001-05-24T00:00:00", + "2001-05-25T00:00:00", + "2001-05-26T00:00:00", + "2001-05-27T00:00:00", + "2001-05-28T00:00:00", + "2001-05-29T00:00:00", + "2001-05-30T00:00:00", + "2001-05-31T00:00:00", + "2001-06-01T00:00:00", + "2001-06-02T00:00:00", + "2001-06-03T00:00:00", + "2001-06-04T00:00:00", + "2001-06-05T00:00:00", + "2001-06-06T00:00:00", + "2001-06-07T00:00:00", + "2001-06-08T00:00:00", + "2001-06-09T00:00:00", + "2001-06-10T00:00:00", + "2001-06-11T00:00:00", + "2001-06-12T00:00:00", + "2001-06-13T00:00:00", + "2001-06-14T00:00:00", + "2001-06-15T00:00:00", + "2001-06-16T00:00:00", + "2001-06-17T00:00:00", + "2001-06-18T00:00:00", + "2001-06-19T00:00:00", + "2001-06-20T00:00:00", + "2001-06-21T00:00:00", + "2001-06-22T00:00:00", + "2001-06-23T00:00:00", + "2001-06-24T00:00:00", + "2001-06-25T00:00:00", + "2001-06-26T00:00:00", + "2001-06-27T00:00:00", + "2001-06-28T00:00:00", + "2001-06-29T00:00:00", + "2001-06-30T00:00:00", + "2001-07-01T00:00:00", + "2001-07-02T00:00:00", + "2001-07-03T00:00:00", + "2001-07-04T00:00:00", + "2001-07-05T00:00:00", + "2001-07-06T00:00:00", + "2001-07-07T00:00:00", + "2001-07-08T00:00:00", + "2001-07-09T00:00:00", + "2001-07-10T00:00:00", + "2001-07-11T00:00:00", + "2001-07-12T00:00:00", + "2001-07-13T00:00:00", + "2001-07-14T00:00:00", + "2001-07-15T00:00:00", + "2001-07-16T00:00:00", + "2001-07-17T00:00:00", + "2001-07-18T00:00:00", + "2001-07-19T00:00:00", + "2001-07-20T00:00:00", + "2001-07-21T00:00:00", + "2001-07-22T00:00:00", + "2001-07-23T00:00:00", + "2001-07-24T00:00:00", + "2001-07-25T00:00:00", + "2001-07-26T00:00:00", + "2001-07-27T00:00:00", + "2001-07-28T00:00:00", + "2001-07-29T00:00:00", + "2001-07-30T00:00:00", + "2001-07-31T00:00:00", + "2001-08-01T00:00:00", + "2001-08-02T00:00:00", + "2001-08-03T00:00:00", + "2001-08-04T00:00:00", + "2001-08-05T00:00:00", + "2001-08-06T00:00:00", + "2001-08-07T00:00:00", + "2001-08-08T00:00:00", + "2001-08-09T00:00:00", + "2001-08-10T00:00:00", + "2001-08-11T00:00:00", + "2001-08-12T00:00:00", + "2001-08-13T00:00:00", + "2001-08-14T00:00:00", + "2001-08-15T00:00:00", + "2001-08-16T00:00:00", + "2001-08-17T00:00:00", + "2001-08-18T00:00:00", + "2001-08-19T00:00:00", + "2001-08-20T00:00:00", + "2001-08-21T00:00:00", + "2001-08-22T00:00:00", + "2001-08-23T00:00:00", + "2001-08-24T00:00:00", + "2001-08-25T00:00:00", + "2001-08-26T00:00:00", + "2001-08-27T00:00:00", + "2001-08-28T00:00:00", + "2001-08-29T00:00:00", + "2001-08-30T00:00:00", + "2001-08-31T00:00:00", + "2001-09-01T00:00:00", + "2001-09-02T00:00:00", + "2001-09-03T00:00:00", + "2001-09-04T00:00:00", + "2001-09-05T00:00:00", + "2001-09-06T00:00:00", + "2001-09-07T00:00:00", + "2001-09-08T00:00:00", + "2001-09-09T00:00:00", + "2001-09-10T00:00:00", + "2001-09-11T00:00:00", + "2001-09-12T00:00:00", + "2001-09-13T00:00:00", + "2001-09-14T00:00:00", + "2001-09-15T00:00:00", + "2001-09-16T00:00:00", + "2001-09-17T00:00:00", + "2001-09-18T00:00:00", + "2001-09-19T00:00:00", + "2001-09-20T00:00:00", + "2001-09-21T00:00:00", + "2001-09-22T00:00:00", + "2001-09-23T00:00:00", + "2001-09-24T00:00:00", + "2001-09-25T00:00:00", + "2001-09-26T00:00:00", + "2001-09-27T00:00:00", + "2001-09-28T00:00:00", + "2001-09-29T00:00:00", + "2001-09-30T00:00:00", + "2001-10-01T00:00:00", + "2001-10-02T00:00:00", + "2001-10-03T00:00:00", + "2001-10-04T00:00:00", + "2001-10-05T00:00:00", + "2001-10-06T00:00:00", + "2001-10-07T00:00:00", + "2001-10-08T00:00:00", + "2001-10-09T00:00:00", + "2001-10-10T00:00:00", + "2001-10-11T00:00:00", + "2001-10-12T00:00:00", + "2001-10-13T00:00:00", + "2001-10-14T00:00:00", + "2001-10-15T00:00:00", + "2001-10-16T00:00:00", + "2001-10-17T00:00:00", + "2001-10-18T00:00:00", + "2001-10-19T00:00:00", + "2001-10-20T00:00:00", + "2001-10-21T00:00:00", + "2001-10-22T00:00:00", + "2001-10-23T00:00:00", + "2001-10-24T00:00:00", + "2001-10-25T00:00:00", + "2001-10-26T00:00:00", + "2001-10-27T00:00:00", + "2001-10-28T00:00:00", + "2001-10-29T00:00:00", + "2001-10-30T00:00:00", + "2001-10-31T00:00:00", + "2001-11-01T00:00:00", + "2001-11-02T00:00:00", + "2001-11-03T00:00:00", + "2001-11-04T00:00:00", + "2001-11-05T00:00:00", + "2001-11-06T00:00:00", + "2001-11-07T00:00:00", + "2001-11-08T00:00:00", + "2001-11-09T00:00:00", + "2001-11-10T00:00:00", + "2001-11-11T00:00:00", + "2001-11-12T00:00:00", + "2001-11-13T00:00:00", + "2001-11-14T00:00:00", + "2001-11-15T00:00:00", + "2001-11-16T00:00:00", + "2001-11-17T00:00:00", + "2001-11-18T00:00:00", + "2001-11-19T00:00:00", + "2001-11-20T00:00:00", + "2001-11-21T00:00:00", + "2001-11-22T00:00:00", + "2001-11-23T00:00:00", + "2001-11-24T00:00:00", + "2001-11-25T00:00:00", + "2001-11-26T00:00:00", + "2001-11-27T00:00:00", + "2001-11-28T00:00:00", + "2001-11-29T00:00:00", + "2001-11-30T00:00:00", + "2001-12-01T00:00:00", + "2001-12-02T00:00:00", + "2001-12-03T00:00:00", + "2001-12-04T00:00:00", + "2001-12-05T00:00:00", + "2001-12-06T00:00:00", + "2001-12-07T00:00:00", + "2001-12-08T00:00:00", + "2001-12-09T00:00:00", + "2001-12-10T00:00:00", + "2001-12-11T00:00:00", + "2001-12-12T00:00:00", + "2001-12-13T00:00:00", + "2001-12-14T00:00:00", + "2001-12-15T00:00:00", + "2001-12-16T00:00:00", + "2001-12-17T00:00:00", + "2001-12-18T00:00:00", + "2001-12-19T00:00:00", + "2001-12-20T00:00:00", + "2001-12-21T00:00:00", + "2001-12-22T00:00:00", + "2001-12-23T00:00:00", + "2001-12-24T00:00:00", + "2001-12-25T00:00:00", + "2001-12-26T00:00:00", + "2001-12-27T00:00:00", + "2001-12-28T00:00:00", + "2001-12-29T00:00:00", + "2001-12-30T00:00:00", + "2001-12-31T00:00:00", + "2002-01-01T00:00:00", + "2002-01-02T00:00:00", + "2002-01-03T00:00:00", + "2002-01-04T00:00:00", + "2002-01-05T00:00:00", + "2002-01-06T00:00:00", + "2002-01-07T00:00:00", + "2002-01-08T00:00:00", + "2002-01-09T00:00:00", + "2002-01-10T00:00:00", + "2002-01-11T00:00:00", + "2002-01-12T00:00:00", + "2002-01-13T00:00:00", + "2002-01-14T00:00:00", + "2002-01-15T00:00:00", + "2002-01-16T00:00:00", + "2002-01-17T00:00:00", + "2002-01-18T00:00:00", + "2002-01-19T00:00:00", + "2002-01-20T00:00:00", + "2002-01-21T00:00:00", + "2002-01-22T00:00:00", + "2002-01-23T00:00:00", + "2002-01-24T00:00:00", + "2002-01-25T00:00:00", + "2002-01-26T00:00:00", + "2002-01-27T00:00:00", + "2002-01-28T00:00:00", + "2002-01-29T00:00:00", + "2002-01-30T00:00:00", + "2002-01-31T00:00:00", + "2002-02-01T00:00:00", + "2002-02-02T00:00:00", + "2002-02-03T00:00:00", + "2002-02-04T00:00:00", + "2002-02-05T00:00:00", + "2002-02-06T00:00:00", + "2002-02-07T00:00:00", + "2002-02-08T00:00:00", + "2002-02-09T00:00:00", + "2002-02-10T00:00:00", + "2002-02-11T00:00:00", + "2002-02-12T00:00:00", + "2002-02-13T00:00:00", + "2002-02-14T00:00:00", + "2002-02-15T00:00:00", + "2002-02-16T00:00:00", + "2002-02-17T00:00:00", + "2002-02-18T00:00:00", + "2002-02-19T00:00:00", + "2002-02-20T00:00:00", + "2002-02-21T00:00:00", + "2002-02-22T00:00:00", + "2002-02-23T00:00:00", + "2002-02-24T00:00:00", + "2002-02-25T00:00:00", + "2002-02-26T00:00:00", + "2002-02-27T00:00:00", + "2002-02-28T00:00:00", + "2002-03-01T00:00:00", + "2002-03-02T00:00:00", + "2002-03-03T00:00:00", + "2002-03-04T00:00:00", + "2002-03-05T00:00:00", + "2002-03-06T00:00:00", + "2002-03-07T00:00:00", + "2002-03-08T00:00:00", + "2002-03-09T00:00:00", + "2002-03-10T00:00:00", + "2002-03-11T00:00:00", + "2002-03-12T00:00:00", + "2002-03-13T00:00:00", + "2002-03-14T00:00:00", + "2002-03-15T00:00:00", + "2002-03-16T00:00:00", + "2002-03-17T00:00:00", + "2002-03-18T00:00:00", + "2002-03-19T00:00:00", + "2002-03-20T00:00:00", + "2002-03-21T00:00:00", + "2002-03-22T00:00:00", + "2002-03-23T00:00:00", + "2002-03-24T00:00:00", + "2002-03-25T00:00:00", + "2002-03-26T00:00:00", + "2002-03-27T00:00:00", + "2002-03-28T00:00:00", + "2002-03-29T00:00:00", + "2002-03-30T00:00:00", + "2002-03-31T00:00:00", + "2002-04-01T00:00:00", + "2002-04-02T00:00:00", + "2002-04-03T00:00:00", + "2002-04-04T00:00:00", + "2002-04-05T00:00:00", + "2002-04-06T00:00:00", + "2002-04-07T00:00:00", + "2002-04-08T00:00:00", + "2002-04-09T00:00:00", + "2002-04-10T00:00:00", + "2002-04-11T00:00:00", + "2002-04-12T00:00:00", + "2002-04-13T00:00:00", + "2002-04-14T00:00:00", + "2002-04-15T00:00:00", + "2002-04-16T00:00:00", + "2002-04-17T00:00:00", + "2002-04-18T00:00:00", + "2002-04-19T00:00:00", + "2002-04-20T00:00:00", + "2002-04-21T00:00:00", + "2002-04-22T00:00:00", + "2002-04-23T00:00:00", + "2002-04-24T00:00:00", + "2002-04-25T00:00:00", + "2002-04-26T00:00:00", + "2002-04-27T00:00:00", + "2002-04-28T00:00:00", + "2002-04-29T00:00:00", + "2002-04-30T00:00:00", + "2002-05-01T00:00:00", + "2002-05-02T00:00:00", + "2002-05-03T00:00:00", + "2002-05-04T00:00:00", + "2002-05-05T00:00:00", + "2002-05-06T00:00:00", + "2002-05-07T00:00:00", + "2002-05-08T00:00:00", + "2002-05-09T00:00:00", + "2002-05-10T00:00:00", + "2002-05-11T00:00:00", + "2002-05-12T00:00:00", + "2002-05-13T00:00:00", + "2002-05-14T00:00:00", + "2002-05-15T00:00:00", + "2002-05-16T00:00:00", + "2002-05-17T00:00:00", + "2002-05-18T00:00:00", + "2002-05-19T00:00:00", + "2002-05-20T00:00:00", + "2002-05-21T00:00:00", + "2002-05-22T00:00:00", + "2002-05-23T00:00:00", + "2002-05-24T00:00:00", + "2002-05-25T00:00:00", + "2002-05-26T00:00:00", + "2002-05-27T00:00:00", + "2002-05-28T00:00:00", + "2002-05-29T00:00:00", + "2002-05-30T00:00:00", + "2002-05-31T00:00:00", + "2002-06-01T00:00:00", + "2002-06-02T00:00:00", + "2002-06-03T00:00:00", + "2002-06-04T00:00:00", + "2002-06-05T00:00:00", + "2002-06-06T00:00:00", + "2002-06-07T00:00:00", + "2002-06-08T00:00:00", + "2002-06-09T00:00:00", + "2002-06-10T00:00:00", + "2002-06-11T00:00:00", + "2002-06-12T00:00:00", + "2002-06-13T00:00:00", + "2002-06-14T00:00:00", + "2002-06-15T00:00:00", + "2002-06-16T00:00:00", + "2002-06-17T00:00:00", + "2002-06-18T00:00:00", + "2002-06-19T00:00:00", + "2002-06-20T00:00:00", + "2002-06-21T00:00:00", + "2002-06-22T00:00:00", + "2002-06-23T00:00:00", + "2002-06-24T00:00:00", + "2002-06-25T00:00:00", + "2002-06-26T00:00:00", + "2002-06-27T00:00:00", + "2002-06-28T00:00:00", + "2002-06-29T00:00:00", + "2002-06-30T00:00:00", + "2002-07-01T00:00:00", + "2002-07-02T00:00:00", + "2002-07-03T00:00:00", + "2002-07-04T00:00:00", + "2002-07-05T00:00:00", + "2002-07-06T00:00:00", + "2002-07-07T00:00:00", + "2002-07-08T00:00:00", + "2002-07-09T00:00:00", + "2002-07-10T00:00:00", + "2002-07-11T00:00:00", + "2002-07-12T00:00:00", + "2002-07-13T00:00:00", + "2002-07-14T00:00:00", + "2002-07-15T00:00:00", + "2002-07-16T00:00:00", + "2002-07-17T00:00:00", + "2002-07-18T00:00:00", + "2002-07-19T00:00:00", + "2002-07-20T00:00:00", + "2002-07-21T00:00:00", + "2002-07-22T00:00:00", + "2002-07-23T00:00:00", + "2002-07-24T00:00:00", + "2002-07-25T00:00:00", + "2002-07-26T00:00:00", + "2002-07-27T00:00:00", + "2002-07-28T00:00:00", + "2002-07-29T00:00:00", + "2002-07-30T00:00:00", + "2002-07-31T00:00:00", + "2002-08-01T00:00:00", + "2002-08-02T00:00:00", + "2002-08-03T00:00:00", + "2002-08-04T00:00:00", + "2002-08-05T00:00:00", + "2002-08-06T00:00:00", + "2002-08-07T00:00:00", + "2002-08-08T00:00:00", + "2002-08-09T00:00:00", + "2002-08-10T00:00:00", + "2002-08-11T00:00:00", + "2002-08-12T00:00:00", + "2002-08-13T00:00:00", + "2002-08-14T00:00:00", + "2002-08-15T00:00:00", + "2002-08-16T00:00:00", + "2002-08-17T00:00:00", + "2002-08-18T00:00:00", + "2002-08-19T00:00:00", + "2002-08-20T00:00:00", + "2002-08-21T00:00:00", + "2002-08-22T00:00:00", + "2002-08-23T00:00:00", + "2002-08-24T00:00:00", + "2002-08-25T00:00:00", + "2002-08-26T00:00:00", + "2002-08-27T00:00:00", + "2002-08-28T00:00:00", + "2002-08-29T00:00:00", + "2002-08-30T00:00:00", + "2002-08-31T00:00:00", + "2002-09-01T00:00:00", + "2002-09-02T00:00:00", + "2002-09-03T00:00:00", + "2002-09-04T00:00:00", + "2002-09-05T00:00:00", + "2002-09-06T00:00:00", + "2002-09-07T00:00:00", + "2002-09-08T00:00:00", + "2002-09-09T00:00:00", + "2002-09-10T00:00:00", + "2002-09-11T00:00:00", + "2002-09-12T00:00:00", + "2002-09-13T00:00:00", + "2002-09-14T00:00:00", + "2002-09-15T00:00:00", + "2002-09-16T00:00:00", + "2002-09-17T00:00:00", + "2002-09-18T00:00:00", + "2002-09-19T00:00:00", + "2002-09-20T00:00:00", + "2002-09-21T00:00:00", + "2002-09-22T00:00:00", + "2002-09-23T00:00:00", + "2002-09-24T00:00:00", + "2002-09-25T00:00:00", + "2002-09-26T00:00:00" + ], + "xaxis": "x", + "y": [ + 0.37709406610993934, + 0.37709406610993934, + 0.37709406610993934, + 1.1901802430468873, + 1.1901802430468873, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.3956419488700784, + 1.663318451084788, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.2675467062035444, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.613117343341243, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836, + 2.880226274210836 + ], + "yaxis": "y" + }, + { + "hovertemplate": "variable=D
    index=%{x}
    value=%{y}", + "legendgroup": "D", + "line": { + "color": "#ab63fa", + "dash": "solid" + }, + "mode": "lines", + "name": "D", + "showlegend": true, + "type": "scattergl", + "x": [ + "2000-01-01T00:00:00", + "2000-01-02T00:00:00", + "2000-01-03T00:00:00", + "2000-01-04T00:00:00", + "2000-01-05T00:00:00", + "2000-01-06T00:00:00", + "2000-01-07T00:00:00", + "2000-01-08T00:00:00", + "2000-01-09T00:00:00", + "2000-01-10T00:00:00", + "2000-01-11T00:00:00", + "2000-01-12T00:00:00", + "2000-01-13T00:00:00", + "2000-01-14T00:00:00", + "2000-01-15T00:00:00", + "2000-01-16T00:00:00", + "2000-01-17T00:00:00", + "2000-01-18T00:00:00", + "2000-01-19T00:00:00", + "2000-01-20T00:00:00", + "2000-01-21T00:00:00", + "2000-01-22T00:00:00", + "2000-01-23T00:00:00", + "2000-01-24T00:00:00", + "2000-01-25T00:00:00", + "2000-01-26T00:00:00", + "2000-01-27T00:00:00", + "2000-01-28T00:00:00", + "2000-01-29T00:00:00", + "2000-01-30T00:00:00", + "2000-01-31T00:00:00", + "2000-02-01T00:00:00", + "2000-02-02T00:00:00", + "2000-02-03T00:00:00", + "2000-02-04T00:00:00", + "2000-02-05T00:00:00", + "2000-02-06T00:00:00", + "2000-02-07T00:00:00", + "2000-02-08T00:00:00", + "2000-02-09T00:00:00", + "2000-02-10T00:00:00", + "2000-02-11T00:00:00", + "2000-02-12T00:00:00", + "2000-02-13T00:00:00", + "2000-02-14T00:00:00", + "2000-02-15T00:00:00", + "2000-02-16T00:00:00", + "2000-02-17T00:00:00", + "2000-02-18T00:00:00", + "2000-02-19T00:00:00", + "2000-02-20T00:00:00", + "2000-02-21T00:00:00", + "2000-02-22T00:00:00", + "2000-02-23T00:00:00", + "2000-02-24T00:00:00", + "2000-02-25T00:00:00", + "2000-02-26T00:00:00", + "2000-02-27T00:00:00", + "2000-02-28T00:00:00", + "2000-02-29T00:00:00", + "2000-03-01T00:00:00", + "2000-03-02T00:00:00", + "2000-03-03T00:00:00", + "2000-03-04T00:00:00", + "2000-03-05T00:00:00", + "2000-03-06T00:00:00", + "2000-03-07T00:00:00", + "2000-03-08T00:00:00", + "2000-03-09T00:00:00", + "2000-03-10T00:00:00", + "2000-03-11T00:00:00", + "2000-03-12T00:00:00", + "2000-03-13T00:00:00", + "2000-03-14T00:00:00", + "2000-03-15T00:00:00", + "2000-03-16T00:00:00", + "2000-03-17T00:00:00", + "2000-03-18T00:00:00", + "2000-03-19T00:00:00", + "2000-03-20T00:00:00", + "2000-03-21T00:00:00", + "2000-03-22T00:00:00", + "2000-03-23T00:00:00", + "2000-03-24T00:00:00", + "2000-03-25T00:00:00", + "2000-03-26T00:00:00", + "2000-03-27T00:00:00", + "2000-03-28T00:00:00", + "2000-03-29T00:00:00", + "2000-03-30T00:00:00", + "2000-03-31T00:00:00", + "2000-04-01T00:00:00", + "2000-04-02T00:00:00", + "2000-04-03T00:00:00", + "2000-04-04T00:00:00", + "2000-04-05T00:00:00", + "2000-04-06T00:00:00", + "2000-04-07T00:00:00", + "2000-04-08T00:00:00", + "2000-04-09T00:00:00", + "2000-04-10T00:00:00", + "2000-04-11T00:00:00", + "2000-04-12T00:00:00", + "2000-04-13T00:00:00", + "2000-04-14T00:00:00", + "2000-04-15T00:00:00", + "2000-04-16T00:00:00", + "2000-04-17T00:00:00", + "2000-04-18T00:00:00", + "2000-04-19T00:00:00", + "2000-04-20T00:00:00", + "2000-04-21T00:00:00", + "2000-04-22T00:00:00", + "2000-04-23T00:00:00", + "2000-04-24T00:00:00", + "2000-04-25T00:00:00", + "2000-04-26T00:00:00", + "2000-04-27T00:00:00", + "2000-04-28T00:00:00", + "2000-04-29T00:00:00", + "2000-04-30T00:00:00", + "2000-05-01T00:00:00", + "2000-05-02T00:00:00", + "2000-05-03T00:00:00", + "2000-05-04T00:00:00", + "2000-05-05T00:00:00", + "2000-05-06T00:00:00", + "2000-05-07T00:00:00", + "2000-05-08T00:00:00", + "2000-05-09T00:00:00", + "2000-05-10T00:00:00", + "2000-05-11T00:00:00", + "2000-05-12T00:00:00", + "2000-05-13T00:00:00", + "2000-05-14T00:00:00", + "2000-05-15T00:00:00", + "2000-05-16T00:00:00", + "2000-05-17T00:00:00", + "2000-05-18T00:00:00", + "2000-05-19T00:00:00", + "2000-05-20T00:00:00", + "2000-05-21T00:00:00", + "2000-05-22T00:00:00", + "2000-05-23T00:00:00", + "2000-05-24T00:00:00", + "2000-05-25T00:00:00", + "2000-05-26T00:00:00", + "2000-05-27T00:00:00", + "2000-05-28T00:00:00", + "2000-05-29T00:00:00", + "2000-05-30T00:00:00", + "2000-05-31T00:00:00", + "2000-06-01T00:00:00", + "2000-06-02T00:00:00", + "2000-06-03T00:00:00", + "2000-06-04T00:00:00", + "2000-06-05T00:00:00", + "2000-06-06T00:00:00", + "2000-06-07T00:00:00", + "2000-06-08T00:00:00", + "2000-06-09T00:00:00", + "2000-06-10T00:00:00", + "2000-06-11T00:00:00", + "2000-06-12T00:00:00", + "2000-06-13T00:00:00", + "2000-06-14T00:00:00", + "2000-06-15T00:00:00", + "2000-06-16T00:00:00", + "2000-06-17T00:00:00", + "2000-06-18T00:00:00", + "2000-06-19T00:00:00", + "2000-06-20T00:00:00", + "2000-06-21T00:00:00", + "2000-06-22T00:00:00", + "2000-06-23T00:00:00", + "2000-06-24T00:00:00", + "2000-06-25T00:00:00", + "2000-06-26T00:00:00", + "2000-06-27T00:00:00", + "2000-06-28T00:00:00", + "2000-06-29T00:00:00", + "2000-06-30T00:00:00", + "2000-07-01T00:00:00", + "2000-07-02T00:00:00", + "2000-07-03T00:00:00", + "2000-07-04T00:00:00", + "2000-07-05T00:00:00", + "2000-07-06T00:00:00", + "2000-07-07T00:00:00", + "2000-07-08T00:00:00", + "2000-07-09T00:00:00", + "2000-07-10T00:00:00", + "2000-07-11T00:00:00", + "2000-07-12T00:00:00", + "2000-07-13T00:00:00", + "2000-07-14T00:00:00", + "2000-07-15T00:00:00", + "2000-07-16T00:00:00", + "2000-07-17T00:00:00", + "2000-07-18T00:00:00", + "2000-07-19T00:00:00", + "2000-07-20T00:00:00", + "2000-07-21T00:00:00", + "2000-07-22T00:00:00", + "2000-07-23T00:00:00", + "2000-07-24T00:00:00", + "2000-07-25T00:00:00", + "2000-07-26T00:00:00", + "2000-07-27T00:00:00", + "2000-07-28T00:00:00", + "2000-07-29T00:00:00", + "2000-07-30T00:00:00", + "2000-07-31T00:00:00", + "2000-08-01T00:00:00", + "2000-08-02T00:00:00", + "2000-08-03T00:00:00", + "2000-08-04T00:00:00", + "2000-08-05T00:00:00", + "2000-08-06T00:00:00", + "2000-08-07T00:00:00", + "2000-08-08T00:00:00", + "2000-08-09T00:00:00", + "2000-08-10T00:00:00", + "2000-08-11T00:00:00", + "2000-08-12T00:00:00", + "2000-08-13T00:00:00", + "2000-08-14T00:00:00", + "2000-08-15T00:00:00", + "2000-08-16T00:00:00", + "2000-08-17T00:00:00", + "2000-08-18T00:00:00", + "2000-08-19T00:00:00", + "2000-08-20T00:00:00", + "2000-08-21T00:00:00", + "2000-08-22T00:00:00", + "2000-08-23T00:00:00", + "2000-08-24T00:00:00", + "2000-08-25T00:00:00", + "2000-08-26T00:00:00", + "2000-08-27T00:00:00", + "2000-08-28T00:00:00", + "2000-08-29T00:00:00", + "2000-08-30T00:00:00", + "2000-08-31T00:00:00", + "2000-09-01T00:00:00", + "2000-09-02T00:00:00", + "2000-09-03T00:00:00", + "2000-09-04T00:00:00", + "2000-09-05T00:00:00", + "2000-09-06T00:00:00", + "2000-09-07T00:00:00", + "2000-09-08T00:00:00", + "2000-09-09T00:00:00", + "2000-09-10T00:00:00", + "2000-09-11T00:00:00", + "2000-09-12T00:00:00", + "2000-09-13T00:00:00", + "2000-09-14T00:00:00", + "2000-09-15T00:00:00", + "2000-09-16T00:00:00", + "2000-09-17T00:00:00", + "2000-09-18T00:00:00", + "2000-09-19T00:00:00", + "2000-09-20T00:00:00", + "2000-09-21T00:00:00", + "2000-09-22T00:00:00", + "2000-09-23T00:00:00", + "2000-09-24T00:00:00", + "2000-09-25T00:00:00", + "2000-09-26T00:00:00", + "2000-09-27T00:00:00", + "2000-09-28T00:00:00", + "2000-09-29T00:00:00", + "2000-09-30T00:00:00", + "2000-10-01T00:00:00", + "2000-10-02T00:00:00", + "2000-10-03T00:00:00", + "2000-10-04T00:00:00", + "2000-10-05T00:00:00", + "2000-10-06T00:00:00", + "2000-10-07T00:00:00", + "2000-10-08T00:00:00", + "2000-10-09T00:00:00", + "2000-10-10T00:00:00", + "2000-10-11T00:00:00", + "2000-10-12T00:00:00", + "2000-10-13T00:00:00", + "2000-10-14T00:00:00", + "2000-10-15T00:00:00", + "2000-10-16T00:00:00", + "2000-10-17T00:00:00", + "2000-10-18T00:00:00", + "2000-10-19T00:00:00", + "2000-10-20T00:00:00", + "2000-10-21T00:00:00", + "2000-10-22T00:00:00", + "2000-10-23T00:00:00", + "2000-10-24T00:00:00", + "2000-10-25T00:00:00", + "2000-10-26T00:00:00", + "2000-10-27T00:00:00", + "2000-10-28T00:00:00", + "2000-10-29T00:00:00", + "2000-10-30T00:00:00", + "2000-10-31T00:00:00", + "2000-11-01T00:00:00", + "2000-11-02T00:00:00", + "2000-11-03T00:00:00", + "2000-11-04T00:00:00", + "2000-11-05T00:00:00", + "2000-11-06T00:00:00", + "2000-11-07T00:00:00", + "2000-11-08T00:00:00", + "2000-11-09T00:00:00", + "2000-11-10T00:00:00", + "2000-11-11T00:00:00", + "2000-11-12T00:00:00", + "2000-11-13T00:00:00", + "2000-11-14T00:00:00", + "2000-11-15T00:00:00", + "2000-11-16T00:00:00", + "2000-11-17T00:00:00", + "2000-11-18T00:00:00", + "2000-11-19T00:00:00", + "2000-11-20T00:00:00", + "2000-11-21T00:00:00", + "2000-11-22T00:00:00", + "2000-11-23T00:00:00", + "2000-11-24T00:00:00", + "2000-11-25T00:00:00", + "2000-11-26T00:00:00", + "2000-11-27T00:00:00", + "2000-11-28T00:00:00", + "2000-11-29T00:00:00", + "2000-11-30T00:00:00", + "2000-12-01T00:00:00", + "2000-12-02T00:00:00", + "2000-12-03T00:00:00", + "2000-12-04T00:00:00", + "2000-12-05T00:00:00", + "2000-12-06T00:00:00", + "2000-12-07T00:00:00", + "2000-12-08T00:00:00", + "2000-12-09T00:00:00", + "2000-12-10T00:00:00", + "2000-12-11T00:00:00", + "2000-12-12T00:00:00", + "2000-12-13T00:00:00", + "2000-12-14T00:00:00", + "2000-12-15T00:00:00", + "2000-12-16T00:00:00", + "2000-12-17T00:00:00", + "2000-12-18T00:00:00", + "2000-12-19T00:00:00", + "2000-12-20T00:00:00", + "2000-12-21T00:00:00", + "2000-12-22T00:00:00", + "2000-12-23T00:00:00", + "2000-12-24T00:00:00", + "2000-12-25T00:00:00", + "2000-12-26T00:00:00", + "2000-12-27T00:00:00", + "2000-12-28T00:00:00", + "2000-12-29T00:00:00", + "2000-12-30T00:00:00", + "2000-12-31T00:00:00", + "2001-01-01T00:00:00", + "2001-01-02T00:00:00", + "2001-01-03T00:00:00", + "2001-01-04T00:00:00", + "2001-01-05T00:00:00", + "2001-01-06T00:00:00", + "2001-01-07T00:00:00", + "2001-01-08T00:00:00", + "2001-01-09T00:00:00", + "2001-01-10T00:00:00", + "2001-01-11T00:00:00", + "2001-01-12T00:00:00", + "2001-01-13T00:00:00", + "2001-01-14T00:00:00", + "2001-01-15T00:00:00", + "2001-01-16T00:00:00", + "2001-01-17T00:00:00", + "2001-01-18T00:00:00", + "2001-01-19T00:00:00", + "2001-01-20T00:00:00", + "2001-01-21T00:00:00", + "2001-01-22T00:00:00", + "2001-01-23T00:00:00", + "2001-01-24T00:00:00", + "2001-01-25T00:00:00", + "2001-01-26T00:00:00", + "2001-01-27T00:00:00", + "2001-01-28T00:00:00", + "2001-01-29T00:00:00", + "2001-01-30T00:00:00", + "2001-01-31T00:00:00", + "2001-02-01T00:00:00", + "2001-02-02T00:00:00", + "2001-02-03T00:00:00", + "2001-02-04T00:00:00", + "2001-02-05T00:00:00", + "2001-02-06T00:00:00", + "2001-02-07T00:00:00", + "2001-02-08T00:00:00", + "2001-02-09T00:00:00", + "2001-02-10T00:00:00", + "2001-02-11T00:00:00", + "2001-02-12T00:00:00", + "2001-02-13T00:00:00", + "2001-02-14T00:00:00", + "2001-02-15T00:00:00", + "2001-02-16T00:00:00", + "2001-02-17T00:00:00", + "2001-02-18T00:00:00", + "2001-02-19T00:00:00", + "2001-02-20T00:00:00", + "2001-02-21T00:00:00", + "2001-02-22T00:00:00", + "2001-02-23T00:00:00", + "2001-02-24T00:00:00", + "2001-02-25T00:00:00", + "2001-02-26T00:00:00", + "2001-02-27T00:00:00", + "2001-02-28T00:00:00", + "2001-03-01T00:00:00", + "2001-03-02T00:00:00", + "2001-03-03T00:00:00", + "2001-03-04T00:00:00", + "2001-03-05T00:00:00", + "2001-03-06T00:00:00", + "2001-03-07T00:00:00", + "2001-03-08T00:00:00", + "2001-03-09T00:00:00", + "2001-03-10T00:00:00", + "2001-03-11T00:00:00", + "2001-03-12T00:00:00", + "2001-03-13T00:00:00", + "2001-03-14T00:00:00", + "2001-03-15T00:00:00", + "2001-03-16T00:00:00", + "2001-03-17T00:00:00", + "2001-03-18T00:00:00", + "2001-03-19T00:00:00", + "2001-03-20T00:00:00", + "2001-03-21T00:00:00", + "2001-03-22T00:00:00", + "2001-03-23T00:00:00", + "2001-03-24T00:00:00", + "2001-03-25T00:00:00", + "2001-03-26T00:00:00", + "2001-03-27T00:00:00", + "2001-03-28T00:00:00", + "2001-03-29T00:00:00", + "2001-03-30T00:00:00", + "2001-03-31T00:00:00", + "2001-04-01T00:00:00", + "2001-04-02T00:00:00", + "2001-04-03T00:00:00", + "2001-04-04T00:00:00", + "2001-04-05T00:00:00", + "2001-04-06T00:00:00", + "2001-04-07T00:00:00", + "2001-04-08T00:00:00", + "2001-04-09T00:00:00", + "2001-04-10T00:00:00", + "2001-04-11T00:00:00", + "2001-04-12T00:00:00", + "2001-04-13T00:00:00", + "2001-04-14T00:00:00", + "2001-04-15T00:00:00", + "2001-04-16T00:00:00", + "2001-04-17T00:00:00", + "2001-04-18T00:00:00", + "2001-04-19T00:00:00", + "2001-04-20T00:00:00", + "2001-04-21T00:00:00", + "2001-04-22T00:00:00", + "2001-04-23T00:00:00", + "2001-04-24T00:00:00", + "2001-04-25T00:00:00", + "2001-04-26T00:00:00", + "2001-04-27T00:00:00", + "2001-04-28T00:00:00", + "2001-04-29T00:00:00", + "2001-04-30T00:00:00", + "2001-05-01T00:00:00", + "2001-05-02T00:00:00", + "2001-05-03T00:00:00", + "2001-05-04T00:00:00", + "2001-05-05T00:00:00", + "2001-05-06T00:00:00", + "2001-05-07T00:00:00", + "2001-05-08T00:00:00", + "2001-05-09T00:00:00", + "2001-05-10T00:00:00", + "2001-05-11T00:00:00", + "2001-05-12T00:00:00", + "2001-05-13T00:00:00", + "2001-05-14T00:00:00", + "2001-05-15T00:00:00", + "2001-05-16T00:00:00", + "2001-05-17T00:00:00", + "2001-05-18T00:00:00", + "2001-05-19T00:00:00", + "2001-05-20T00:00:00", + "2001-05-21T00:00:00", + "2001-05-22T00:00:00", + "2001-05-23T00:00:00", + "2001-05-24T00:00:00", + "2001-05-25T00:00:00", + "2001-05-26T00:00:00", + "2001-05-27T00:00:00", + "2001-05-28T00:00:00", + "2001-05-29T00:00:00", + "2001-05-30T00:00:00", + "2001-05-31T00:00:00", + "2001-06-01T00:00:00", + "2001-06-02T00:00:00", + "2001-06-03T00:00:00", + "2001-06-04T00:00:00", + "2001-06-05T00:00:00", + "2001-06-06T00:00:00", + "2001-06-07T00:00:00", + "2001-06-08T00:00:00", + "2001-06-09T00:00:00", + "2001-06-10T00:00:00", + "2001-06-11T00:00:00", + "2001-06-12T00:00:00", + "2001-06-13T00:00:00", + "2001-06-14T00:00:00", + "2001-06-15T00:00:00", + "2001-06-16T00:00:00", + "2001-06-17T00:00:00", + "2001-06-18T00:00:00", + "2001-06-19T00:00:00", + "2001-06-20T00:00:00", + "2001-06-21T00:00:00", + "2001-06-22T00:00:00", + "2001-06-23T00:00:00", + "2001-06-24T00:00:00", + "2001-06-25T00:00:00", + "2001-06-26T00:00:00", + "2001-06-27T00:00:00", + "2001-06-28T00:00:00", + "2001-06-29T00:00:00", + "2001-06-30T00:00:00", + "2001-07-01T00:00:00", + "2001-07-02T00:00:00", + "2001-07-03T00:00:00", + "2001-07-04T00:00:00", + "2001-07-05T00:00:00", + "2001-07-06T00:00:00", + "2001-07-07T00:00:00", + "2001-07-08T00:00:00", + "2001-07-09T00:00:00", + "2001-07-10T00:00:00", + "2001-07-11T00:00:00", + "2001-07-12T00:00:00", + "2001-07-13T00:00:00", + "2001-07-14T00:00:00", + "2001-07-15T00:00:00", + "2001-07-16T00:00:00", + "2001-07-17T00:00:00", + "2001-07-18T00:00:00", + "2001-07-19T00:00:00", + "2001-07-20T00:00:00", + "2001-07-21T00:00:00", + "2001-07-22T00:00:00", + "2001-07-23T00:00:00", + "2001-07-24T00:00:00", + "2001-07-25T00:00:00", + "2001-07-26T00:00:00", + "2001-07-27T00:00:00", + "2001-07-28T00:00:00", + "2001-07-29T00:00:00", + "2001-07-30T00:00:00", + "2001-07-31T00:00:00", + "2001-08-01T00:00:00", + "2001-08-02T00:00:00", + "2001-08-03T00:00:00", + "2001-08-04T00:00:00", + "2001-08-05T00:00:00", + "2001-08-06T00:00:00", + "2001-08-07T00:00:00", + "2001-08-08T00:00:00", + "2001-08-09T00:00:00", + "2001-08-10T00:00:00", + "2001-08-11T00:00:00", + "2001-08-12T00:00:00", + "2001-08-13T00:00:00", + "2001-08-14T00:00:00", + "2001-08-15T00:00:00", + "2001-08-16T00:00:00", + "2001-08-17T00:00:00", + "2001-08-18T00:00:00", + "2001-08-19T00:00:00", + "2001-08-20T00:00:00", + "2001-08-21T00:00:00", + "2001-08-22T00:00:00", + "2001-08-23T00:00:00", + "2001-08-24T00:00:00", + "2001-08-25T00:00:00", + "2001-08-26T00:00:00", + "2001-08-27T00:00:00", + "2001-08-28T00:00:00", + "2001-08-29T00:00:00", + "2001-08-30T00:00:00", + "2001-08-31T00:00:00", + "2001-09-01T00:00:00", + "2001-09-02T00:00:00", + "2001-09-03T00:00:00", + "2001-09-04T00:00:00", + "2001-09-05T00:00:00", + "2001-09-06T00:00:00", + "2001-09-07T00:00:00", + "2001-09-08T00:00:00", + "2001-09-09T00:00:00", + "2001-09-10T00:00:00", + "2001-09-11T00:00:00", + "2001-09-12T00:00:00", + "2001-09-13T00:00:00", + "2001-09-14T00:00:00", + "2001-09-15T00:00:00", + "2001-09-16T00:00:00", + "2001-09-17T00:00:00", + "2001-09-18T00:00:00", + "2001-09-19T00:00:00", + "2001-09-20T00:00:00", + "2001-09-21T00:00:00", + "2001-09-22T00:00:00", + "2001-09-23T00:00:00", + "2001-09-24T00:00:00", + "2001-09-25T00:00:00", + "2001-09-26T00:00:00", + "2001-09-27T00:00:00", + "2001-09-28T00:00:00", + "2001-09-29T00:00:00", + "2001-09-30T00:00:00", + "2001-10-01T00:00:00", + "2001-10-02T00:00:00", + "2001-10-03T00:00:00", + "2001-10-04T00:00:00", + "2001-10-05T00:00:00", + "2001-10-06T00:00:00", + "2001-10-07T00:00:00", + "2001-10-08T00:00:00", + "2001-10-09T00:00:00", + "2001-10-10T00:00:00", + "2001-10-11T00:00:00", + "2001-10-12T00:00:00", + "2001-10-13T00:00:00", + "2001-10-14T00:00:00", + "2001-10-15T00:00:00", + "2001-10-16T00:00:00", + "2001-10-17T00:00:00", + "2001-10-18T00:00:00", + "2001-10-19T00:00:00", + "2001-10-20T00:00:00", + "2001-10-21T00:00:00", + "2001-10-22T00:00:00", + "2001-10-23T00:00:00", + "2001-10-24T00:00:00", + "2001-10-25T00:00:00", + "2001-10-26T00:00:00", + "2001-10-27T00:00:00", + "2001-10-28T00:00:00", + "2001-10-29T00:00:00", + "2001-10-30T00:00:00", + "2001-10-31T00:00:00", + "2001-11-01T00:00:00", + "2001-11-02T00:00:00", + "2001-11-03T00:00:00", + "2001-11-04T00:00:00", + "2001-11-05T00:00:00", + "2001-11-06T00:00:00", + "2001-11-07T00:00:00", + "2001-11-08T00:00:00", + "2001-11-09T00:00:00", + "2001-11-10T00:00:00", + "2001-11-11T00:00:00", + "2001-11-12T00:00:00", + "2001-11-13T00:00:00", + "2001-11-14T00:00:00", + "2001-11-15T00:00:00", + "2001-11-16T00:00:00", + "2001-11-17T00:00:00", + "2001-11-18T00:00:00", + "2001-11-19T00:00:00", + "2001-11-20T00:00:00", + "2001-11-21T00:00:00", + "2001-11-22T00:00:00", + "2001-11-23T00:00:00", + "2001-11-24T00:00:00", + "2001-11-25T00:00:00", + "2001-11-26T00:00:00", + "2001-11-27T00:00:00", + "2001-11-28T00:00:00", + "2001-11-29T00:00:00", + "2001-11-30T00:00:00", + "2001-12-01T00:00:00", + "2001-12-02T00:00:00", + "2001-12-03T00:00:00", + "2001-12-04T00:00:00", + "2001-12-05T00:00:00", + "2001-12-06T00:00:00", + "2001-12-07T00:00:00", + "2001-12-08T00:00:00", + "2001-12-09T00:00:00", + "2001-12-10T00:00:00", + "2001-12-11T00:00:00", + "2001-12-12T00:00:00", + "2001-12-13T00:00:00", + "2001-12-14T00:00:00", + "2001-12-15T00:00:00", + "2001-12-16T00:00:00", + "2001-12-17T00:00:00", + "2001-12-18T00:00:00", + "2001-12-19T00:00:00", + "2001-12-20T00:00:00", + "2001-12-21T00:00:00", + "2001-12-22T00:00:00", + "2001-12-23T00:00:00", + "2001-12-24T00:00:00", + "2001-12-25T00:00:00", + "2001-12-26T00:00:00", + "2001-12-27T00:00:00", + "2001-12-28T00:00:00", + "2001-12-29T00:00:00", + "2001-12-30T00:00:00", + "2001-12-31T00:00:00", + "2002-01-01T00:00:00", + "2002-01-02T00:00:00", + "2002-01-03T00:00:00", + "2002-01-04T00:00:00", + "2002-01-05T00:00:00", + "2002-01-06T00:00:00", + "2002-01-07T00:00:00", + "2002-01-08T00:00:00", + "2002-01-09T00:00:00", + "2002-01-10T00:00:00", + "2002-01-11T00:00:00", + "2002-01-12T00:00:00", + "2002-01-13T00:00:00", + "2002-01-14T00:00:00", + "2002-01-15T00:00:00", + "2002-01-16T00:00:00", + "2002-01-17T00:00:00", + "2002-01-18T00:00:00", + "2002-01-19T00:00:00", + "2002-01-20T00:00:00", + "2002-01-21T00:00:00", + "2002-01-22T00:00:00", + "2002-01-23T00:00:00", + "2002-01-24T00:00:00", + "2002-01-25T00:00:00", + "2002-01-26T00:00:00", + "2002-01-27T00:00:00", + "2002-01-28T00:00:00", + "2002-01-29T00:00:00", + "2002-01-30T00:00:00", + "2002-01-31T00:00:00", + "2002-02-01T00:00:00", + "2002-02-02T00:00:00", + "2002-02-03T00:00:00", + "2002-02-04T00:00:00", + "2002-02-05T00:00:00", + "2002-02-06T00:00:00", + "2002-02-07T00:00:00", + "2002-02-08T00:00:00", + "2002-02-09T00:00:00", + "2002-02-10T00:00:00", + "2002-02-11T00:00:00", + "2002-02-12T00:00:00", + "2002-02-13T00:00:00", + "2002-02-14T00:00:00", + "2002-02-15T00:00:00", + "2002-02-16T00:00:00", + "2002-02-17T00:00:00", + "2002-02-18T00:00:00", + "2002-02-19T00:00:00", + "2002-02-20T00:00:00", + "2002-02-21T00:00:00", + "2002-02-22T00:00:00", + "2002-02-23T00:00:00", + "2002-02-24T00:00:00", + "2002-02-25T00:00:00", + "2002-02-26T00:00:00", + "2002-02-27T00:00:00", + "2002-02-28T00:00:00", + "2002-03-01T00:00:00", + "2002-03-02T00:00:00", + "2002-03-03T00:00:00", + "2002-03-04T00:00:00", + "2002-03-05T00:00:00", + "2002-03-06T00:00:00", + "2002-03-07T00:00:00", + "2002-03-08T00:00:00", + "2002-03-09T00:00:00", + "2002-03-10T00:00:00", + "2002-03-11T00:00:00", + "2002-03-12T00:00:00", + "2002-03-13T00:00:00", + "2002-03-14T00:00:00", + "2002-03-15T00:00:00", + "2002-03-16T00:00:00", + "2002-03-17T00:00:00", + "2002-03-18T00:00:00", + "2002-03-19T00:00:00", + "2002-03-20T00:00:00", + "2002-03-21T00:00:00", + "2002-03-22T00:00:00", + "2002-03-23T00:00:00", + "2002-03-24T00:00:00", + "2002-03-25T00:00:00", + "2002-03-26T00:00:00", + "2002-03-27T00:00:00", + "2002-03-28T00:00:00", + "2002-03-29T00:00:00", + "2002-03-30T00:00:00", + "2002-03-31T00:00:00", + "2002-04-01T00:00:00", + "2002-04-02T00:00:00", + "2002-04-03T00:00:00", + "2002-04-04T00:00:00", + "2002-04-05T00:00:00", + "2002-04-06T00:00:00", + "2002-04-07T00:00:00", + "2002-04-08T00:00:00", + "2002-04-09T00:00:00", + "2002-04-10T00:00:00", + "2002-04-11T00:00:00", + "2002-04-12T00:00:00", + "2002-04-13T00:00:00", + "2002-04-14T00:00:00", + "2002-04-15T00:00:00", + "2002-04-16T00:00:00", + "2002-04-17T00:00:00", + "2002-04-18T00:00:00", + "2002-04-19T00:00:00", + "2002-04-20T00:00:00", + "2002-04-21T00:00:00", + "2002-04-22T00:00:00", + "2002-04-23T00:00:00", + "2002-04-24T00:00:00", + "2002-04-25T00:00:00", + "2002-04-26T00:00:00", + "2002-04-27T00:00:00", + "2002-04-28T00:00:00", + "2002-04-29T00:00:00", + "2002-04-30T00:00:00", + "2002-05-01T00:00:00", + "2002-05-02T00:00:00", + "2002-05-03T00:00:00", + "2002-05-04T00:00:00", + "2002-05-05T00:00:00", + "2002-05-06T00:00:00", + "2002-05-07T00:00:00", + "2002-05-08T00:00:00", + "2002-05-09T00:00:00", + "2002-05-10T00:00:00", + "2002-05-11T00:00:00", + "2002-05-12T00:00:00", + "2002-05-13T00:00:00", + "2002-05-14T00:00:00", + "2002-05-15T00:00:00", + "2002-05-16T00:00:00", + "2002-05-17T00:00:00", + "2002-05-18T00:00:00", + "2002-05-19T00:00:00", + "2002-05-20T00:00:00", + "2002-05-21T00:00:00", + "2002-05-22T00:00:00", + "2002-05-23T00:00:00", + "2002-05-24T00:00:00", + "2002-05-25T00:00:00", + "2002-05-26T00:00:00", + "2002-05-27T00:00:00", + "2002-05-28T00:00:00", + "2002-05-29T00:00:00", + "2002-05-30T00:00:00", + "2002-05-31T00:00:00", + "2002-06-01T00:00:00", + "2002-06-02T00:00:00", + "2002-06-03T00:00:00", + "2002-06-04T00:00:00", + "2002-06-05T00:00:00", + "2002-06-06T00:00:00", + "2002-06-07T00:00:00", + "2002-06-08T00:00:00", + "2002-06-09T00:00:00", + "2002-06-10T00:00:00", + "2002-06-11T00:00:00", + "2002-06-12T00:00:00", + "2002-06-13T00:00:00", + "2002-06-14T00:00:00", + "2002-06-15T00:00:00", + "2002-06-16T00:00:00", + "2002-06-17T00:00:00", + "2002-06-18T00:00:00", + "2002-06-19T00:00:00", + "2002-06-20T00:00:00", + "2002-06-21T00:00:00", + "2002-06-22T00:00:00", + "2002-06-23T00:00:00", + "2002-06-24T00:00:00", + "2002-06-25T00:00:00", + "2002-06-26T00:00:00", + "2002-06-27T00:00:00", + "2002-06-28T00:00:00", + "2002-06-29T00:00:00", + "2002-06-30T00:00:00", + "2002-07-01T00:00:00", + "2002-07-02T00:00:00", + "2002-07-03T00:00:00", + "2002-07-04T00:00:00", + "2002-07-05T00:00:00", + "2002-07-06T00:00:00", + "2002-07-07T00:00:00", + "2002-07-08T00:00:00", + "2002-07-09T00:00:00", + "2002-07-10T00:00:00", + "2002-07-11T00:00:00", + "2002-07-12T00:00:00", + "2002-07-13T00:00:00", + "2002-07-14T00:00:00", + "2002-07-15T00:00:00", + "2002-07-16T00:00:00", + "2002-07-17T00:00:00", + "2002-07-18T00:00:00", + "2002-07-19T00:00:00", + "2002-07-20T00:00:00", + "2002-07-21T00:00:00", + "2002-07-22T00:00:00", + "2002-07-23T00:00:00", + "2002-07-24T00:00:00", + "2002-07-25T00:00:00", + "2002-07-26T00:00:00", + "2002-07-27T00:00:00", + "2002-07-28T00:00:00", + "2002-07-29T00:00:00", + "2002-07-30T00:00:00", + "2002-07-31T00:00:00", + "2002-08-01T00:00:00", + "2002-08-02T00:00:00", + "2002-08-03T00:00:00", + "2002-08-04T00:00:00", + "2002-08-05T00:00:00", + "2002-08-06T00:00:00", + "2002-08-07T00:00:00", + "2002-08-08T00:00:00", + "2002-08-09T00:00:00", + "2002-08-10T00:00:00", + "2002-08-11T00:00:00", + "2002-08-12T00:00:00", + "2002-08-13T00:00:00", + "2002-08-14T00:00:00", + "2002-08-15T00:00:00", + "2002-08-16T00:00:00", + "2002-08-17T00:00:00", + "2002-08-18T00:00:00", + "2002-08-19T00:00:00", + "2002-08-20T00:00:00", + "2002-08-21T00:00:00", + "2002-08-22T00:00:00", + "2002-08-23T00:00:00", + "2002-08-24T00:00:00", + "2002-08-25T00:00:00", + "2002-08-26T00:00:00", + "2002-08-27T00:00:00", + "2002-08-28T00:00:00", + "2002-08-29T00:00:00", + "2002-08-30T00:00:00", + "2002-08-31T00:00:00", + "2002-09-01T00:00:00", + "2002-09-02T00:00:00", + "2002-09-03T00:00:00", + "2002-09-04T00:00:00", + "2002-09-05T00:00:00", + "2002-09-06T00:00:00", + "2002-09-07T00:00:00", + "2002-09-08T00:00:00", + "2002-09-09T00:00:00", + "2002-09-10T00:00:00", + "2002-09-11T00:00:00", + "2002-09-12T00:00:00", + "2002-09-13T00:00:00", + "2002-09-14T00:00:00", + "2002-09-15T00:00:00", + "2002-09-16T00:00:00", + "2002-09-17T00:00:00", + "2002-09-18T00:00:00", + "2002-09-19T00:00:00", + "2002-09-20T00:00:00", + "2002-09-21T00:00:00", + "2002-09-22T00:00:00", + "2002-09-23T00:00:00", + "2002-09-24T00:00:00", + "2002-09-25T00:00:00", + "2002-09-26T00:00:00" + ], + "xaxis": "x", + "y": [ + -0.6272166884627877, + -0.6272166884627877, + 0.3659704255900808, + 0.3659704255900808, + 0.3659704255900808, + 0.3659704255900808, + 1.0458675267067983, + 1.0458675267067983, + 1.0458675267067983, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.5569328413694534, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9356353306523484, + 1.9608776827034715, + 1.9608776827034715, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 2.6646718087031096, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873, + 3.2630067735774873 + ], + "yaxis": "y" + } + ], + "layout": { + "legend": { + "title": { + "text": "variable" + }, + "tracegroupgap": 0 + }, + "margin": { + "t": 60 + }, + "template": { + "data": { + "bar": [ + { + "error_x": { + "color": "#2a3f5f" + }, + "error_y": { + "color": "#2a3f5f" + }, + "marker": { + "line": { + "color": "#E5ECF6", + "width": 0.5 + }, + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "bar" + } + ], + "barpolar": [ + { + "marker": { + "line": { + "color": "#E5ECF6", + "width": 0.5 + }, + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "barpolar" + } + ], + "carpet": [ + { + "aaxis": { + "endlinecolor": "#2a3f5f", + "gridcolor": "white", + "linecolor": "white", + "minorgridcolor": "white", + "startlinecolor": "#2a3f5f" + }, + "baxis": { + "endlinecolor": "#2a3f5f", + "gridcolor": "white", + "linecolor": "white", + "minorgridcolor": "white", + "startlinecolor": "#2a3f5f" + }, + "type": "carpet" + } + ], + "choropleth": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "choropleth" + } + ], + "contour": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "contour" + } + ], + "contourcarpet": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "contourcarpet" + } + ], + "heatmap": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "heatmap" + } + ], + "heatmapgl": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "heatmapgl" + } + ], + "histogram": [ + { + "marker": { + "pattern": { + "fillmode": "overlay", + "size": 10, + "solidity": 0.2 + } + }, + "type": "histogram" + } + ], + "histogram2d": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "histogram2d" + } + ], + "histogram2dcontour": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "histogram2dcontour" + } + ], + "mesh3d": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "type": "mesh3d" + } + ], + "parcoords": [ + { + "line": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "parcoords" + } + ], + "pie": [ + { + "automargin": true, + "type": "pie" + } + ], + "scatter": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatter" + } + ], + "scatter3d": [ + { + "line": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatter3d" + } + ], + "scattercarpet": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattercarpet" + } + ], + "scattergeo": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattergeo" + } + ], + "scattergl": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattergl" + } + ], + "scattermapbox": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scattermapbox" + } + ], + "scatterpolar": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterpolar" + } + ], + "scatterpolargl": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterpolargl" + } + ], + "scatterternary": [ + { + "marker": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "type": "scatterternary" + } + ], + "surface": [ + { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + }, + "colorscale": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "type": "surface" + } + ], + "table": [ + { + "cells": { + "fill": { + "color": "#EBF0F8" + }, + "line": { + "color": "white" + } + }, + "header": { + "fill": { + "color": "#C8D4E3" + }, + "line": { + "color": "white" + } + }, + "type": "table" + } + ] + }, + "layout": { + "annotationdefaults": { + "arrowcolor": "#2a3f5f", + "arrowhead": 0, + "arrowwidth": 1 + }, + "autotypenumbers": "strict", + "coloraxis": { + "colorbar": { + "outlinewidth": 0, + "ticks": "" + } + }, + "colorscale": { + "diverging": [ + [ + 0, + "#8e0152" + ], + [ + 0.1, + "#c51b7d" + ], + [ + 0.2, + "#de77ae" + ], + [ + 0.3, + "#f1b6da" + ], + [ + 0.4, + "#fde0ef" + ], + [ + 0.5, + "#f7f7f7" + ], + [ + 0.6, + "#e6f5d0" + ], + [ + 0.7, + "#b8e186" + ], + [ + 0.8, + "#7fbc41" + ], + [ + 0.9, + "#4d9221" + ], + [ + 1, + "#276419" + ] + ], + "sequential": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ], + "sequentialminus": [ + [ + 0, + "#0d0887" + ], + [ + 0.1111111111111111, + "#46039f" + ], + [ + 0.2222222222222222, + "#7201a8" + ], + [ + 0.3333333333333333, + "#9c179e" + ], + [ + 0.4444444444444444, + "#bd3786" + ], + [ + 0.5555555555555556, + "#d8576b" + ], + [ + 0.6666666666666666, + "#ed7953" + ], + [ + 0.7777777777777778, + "#fb9f3a" + ], + [ + 0.8888888888888888, + "#fdca26" + ], + [ + 1, + "#f0f921" + ] + ] + }, + "colorway": [ + "#636efa", + "#EF553B", + "#00cc96", + "#ab63fa", + "#FFA15A", + "#19d3f3", + "#FF6692", + "#B6E880", + "#FF97FF", + "#FECB52" + ], + "font": { + "color": "#2a3f5f" + }, + "geo": { + "bgcolor": "white", + "lakecolor": "white", + "landcolor": "#E5ECF6", + "showlakes": true, + "showland": true, + "subunitcolor": "white" + }, + "hoverlabel": { + "align": "left" + }, + "hovermode": "closest", + "mapbox": { + "style": "light" + }, + "paper_bgcolor": "white", + "plot_bgcolor": "#E5ECF6", + "polar": { + "angularaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "bgcolor": "#E5ECF6", + "radialaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + } + }, + "scene": { + "xaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + }, + "yaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + }, + "zaxis": { + "backgroundcolor": "#E5ECF6", + "gridcolor": "white", + "gridwidth": 2, + "linecolor": "white", + "showbackground": true, + "ticks": "", + "zerolinecolor": "white" + } + }, + "shapedefaults": { + "line": { + "color": "#2a3f5f" + } + }, + "ternary": { + "aaxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "baxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + }, + "bgcolor": "#E5ECF6", + "caxis": { + "gridcolor": "white", + "linecolor": "white", + "ticks": "" + } + }, + "title": { + "x": 0.05 + }, + "xaxis": { + "automargin": true, + "gridcolor": "white", + "linecolor": "white", + "ticks": "", + "title": { + "standoff": 15 + }, + "zerolinecolor": "white", + "zerolinewidth": 2 + }, + "yaxis": { + "automargin": true, + "gridcolor": "white", + "linecolor": "white", + "ticks": "", + "title": { + "standoff": 15 + }, + "zerolinecolor": "white", + "zerolinewidth": 2 + } + } + }, + "xaxis": { + "anchor": "y", + "domain": [ + 0, + 1 + ], + "title": { + "text": "index" + } + }, + "yaxis": { + "anchor": "x", + "domain": [ + 0, + 1 + ], + "title": { + "text": "value" + } + } + } + }, + "text/html": [ + "
    " + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "psdf.plot()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For more details, [Plotting](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#plotting) documentation." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Getting data in/out" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### CSV\n", + "\n", + "CSV is straightforward and easy to use. See [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.to_csv.html) to write a CSV file and [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.read_csv.html) to read a CSV file." + ] + }, + { + "cell_type": "code", + "execution_count": 50, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    0-1.187097-0.1346450.377094-0.627217
    10.3317410.1662180.377094-0.627217
    20.3317410.4394500.3770940.365970
    30.6216200.4394501.1901800.365970
    40.6216200.4394501.1901800.365970
    52.1691981.0691831.3956420.365970
    62.7557381.0691831.3956421.045868
    72.7557381.0691831.3956421.045868
    82.7557381.0691831.3956421.045868
    92.7557381.5087321.3956421.556933
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 -1.187097 -0.134645 0.377094 -0.627217\n", + "1 0.331741 0.166218 0.377094 -0.627217\n", + "2 0.331741 0.439450 0.377094 0.365970\n", + "3 0.621620 0.439450 1.190180 0.365970\n", + "4 0.621620 0.439450 1.190180 0.365970\n", + "5 2.169198 1.069183 1.395642 0.365970\n", + "6 2.755738 1.069183 1.395642 1.045868\n", + "7 2.755738 1.069183 1.395642 1.045868\n", + "8 2.755738 1.069183 1.395642 1.045868\n", + "9 2.755738 1.508732 1.395642 1.556933" + ] + }, + "execution_count": 50, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.to_csv('foo.csv')\n", + "ps.read_csv('foo.csv').head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parquet\n", + "\n", + "Parquet is an efficient and compact file format to read and write faster. See [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.to_parquet.html) to write a Parquet file and [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.read_parquet.html) to read a Parquet file." + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    0-1.187097-0.1346450.377094-0.627217
    10.3317410.1662180.377094-0.627217
    20.3317410.4394500.3770940.365970
    30.6216200.4394501.1901800.365970
    40.6216200.4394501.1901800.365970
    52.1691981.0691831.3956420.365970
    62.7557381.0691831.3956421.045868
    72.7557381.0691831.3956421.045868
    82.7557381.0691831.3956421.045868
    92.7557381.5087321.3956421.556933
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 -1.187097 -0.134645 0.377094 -0.627217\n", + "1 0.331741 0.166218 0.377094 -0.627217\n", + "2 0.331741 0.439450 0.377094 0.365970\n", + "3 0.621620 0.439450 1.190180 0.365970\n", + "4 0.621620 0.439450 1.190180 0.365970\n", + "5 2.169198 1.069183 1.395642 0.365970\n", + "6 2.755738 1.069183 1.395642 1.045868\n", + "7 2.755738 1.069183 1.395642 1.045868\n", + "8 2.755738 1.069183 1.395642 1.045868\n", + "9 2.755738 1.508732 1.395642 1.556933" + ] + }, + "execution_count": 51, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.to_parquet('bar.parquet')\n", + "ps.read_parquet('bar.parquet').head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Spark IO\n", + "\n", + "In addition, pandas API on Spark fully supports Spark's various datasources such as ORC and an external datasource. See [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.to_orc.html) to write it to the specified datasource and [here](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.read_orc.html) to read it from the datasource." + ] + }, + { + "cell_type": "code", + "execution_count": 52, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
    \n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    ABCD
    0-1.187097-0.1346450.377094-0.627217
    10.3317410.1662180.377094-0.627217
    20.3317410.4394500.3770940.365970
    30.6216200.4394501.1901800.365970
    40.6216200.4394501.1901800.365970
    52.1691981.0691831.3956420.365970
    62.7557381.0691831.3956421.045868
    72.7557381.0691831.3956421.045868
    82.7557381.0691831.3956421.045868
    92.7557381.5087321.3956421.556933
    \n", + "
    " + ], + "text/plain": [ + " A B C D\n", + "0 -1.187097 -0.134645 0.377094 -0.627217\n", + "1 0.331741 0.166218 0.377094 -0.627217\n", + "2 0.331741 0.439450 0.377094 0.365970\n", + "3 0.621620 0.439450 1.190180 0.365970\n", + "4 0.621620 0.439450 1.190180 0.365970\n", + "5 2.169198 1.069183 1.395642 0.365970\n", + "6 2.755738 1.069183 1.395642 1.045868\n", + "7 2.755738 1.069183 1.395642 1.045868\n", + "8 2.755738 1.069183 1.395642 1.045868\n", + "9 2.755738 1.508732 1.395642 1.556933" + ] + }, + "execution_count": 52, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "psdf.to_spark_io('zoo.orc', format=\"orc\")\n", + "ps.read_spark_io('zoo.orc', format=\"orc\").head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "See the [Input/Output](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/io.html) documentation for more details." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.10" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/testing_pyspark.ipynb b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/testing_pyspark.ipynb new file mode 100644 index 0000000..268ace0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/getting_started/testing_pyspark.ipynb @@ -0,0 +1,485 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4ee2125b-f889-47e6-9c3d-8bd63a253683", + "metadata": {}, + "source": [ + "# Testing PySpark\n", + "\n", + "This guide is a reference for writing robust tests for PySpark code.\n", + "\n", + "To view the docs for PySpark test utils, see here. To see the code for PySpark built-in test utils, check out the Spark repository here. To see the JIRA board tickets for the PySpark test framework, see here." + ] + }, + { + "cell_type": "markdown", + "id": "0e8ee4b6-9544-45e1-8a91-e71ed8ef8b9d", + "metadata": {}, + "source": [ + "## Build a PySpark Application\n", + "Here is an example for how to start a PySpark application. Feel free to skip to the next section, “Testing your PySpark Application,” if you already have an application you’re ready to test.\n", + "\n", + "First, start your Spark Session." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "9af4a35b-17e8-4e45-816b-34c14c5902f7", + "metadata": {}, + "outputs": [], + "source": [ + "from pyspark.sql import SparkSession \n", + "from pyspark.sql.functions import col \n", + "\n", + "# Create a SparkSession \n", + "spark = SparkSession.builder.appName(\"Testing PySpark Example\").getOrCreate() " + ] + }, + { + "cell_type": "markdown", + "id": "4a4c6efe-91f5-4e18-b4b2-b0401c2368e4", + "metadata": {}, + "source": [ + "Next, create a DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "3b483dd8-3a76-41c6-9206-301d7ef314d6", + "metadata": {}, + "outputs": [], + "source": [ + "sample_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}] \n", + "\n", + "df = spark.createDataFrame(sample_data)" + ] + }, + { + "cell_type": "markdown", + "id": "e0f44333-0e08-470b-9fa2-38f59e3dbd63", + "metadata": {}, + "source": [ + "Now, let’s define and apply a transformation function to our DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a6c0b766-af5f-4e1d-acf8-887d7cf0b0b2", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "+---+--------+\n", + "|age| name|\n", + "+---+--------+\n", + "| 30| John D.|\n", + "| 25|Alice G.|\n", + "| 35| Bob T.|\n", + "| 28| Eve A.|\n", + "+---+--------+\n", + "\n" + ] + } + ], + "source": [ + "from pyspark.sql.functions import col, regexp_replace\n", + "\n", + "# Remove additional spaces in name\n", + "def remove_extra_spaces(df, column_name):\n", + " # Remove extra spaces from the specified column\n", + " df_transformed = df.withColumn(column_name, regexp_replace(col(column_name), \"\\\\s+\", \" \"))\n", + " \n", + " return df_transformed\n", + "\n", + "transformed_df = remove_extra_spaces(df, \"name\")\n", + "\n", + "transformed_df.show()" + ] + }, + { + "cell_type": "markdown", + "id": "530beaa6-aabf-43a1-ad2b-361f267e9608", + "metadata": {}, + "source": [ + "## Testing your PySpark Application\n", + "Now let’s test our PySpark transformation function. \n", + "\n", + "One option is to simply eyeball the resulting DataFrame. However, this can be impractical for large DataFrame or input sizes.\n", + "\n", + "A better way is to write tests. Here are some examples of how we can test our code. The examples below apply for Spark 3.5 and above versions.\n", + "\n", + "Note that these examples are not exhaustive, as there are many other test framework alternatives which you can use instead of `unittest` or `pytest`. The built-in PySpark testing util functions are standalone, meaning they can be compatible with any test framework or CI test pipeline.\n" + ] + }, + { + "cell_type": "markdown", + "id": "d84a9fc1-9768-4af4-bfbf-e832f23334dc", + "metadata": {}, + "source": [ + "### Option 1: Using Only PySpark Built-in Test Utility Functions\n", + "\n", + "For simple ad-hoc validation cases, PySpark testing utils like `assertDataFrameEqual` and `assertSchemaEqual` can be used in a standalone context.\n", + "You could easily test PySpark code in a notebook session. For example, say you want to assert equality between two DataFrames:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "8e533732-ee40-4cd0-9669-8eb92973908a", + "metadata": {}, + "outputs": [], + "source": [ + "import pyspark.testing\n", + "from pyspark.testing.utils import assertDataFrameEqual\n", + "\n", + "# Example 1\n", + "df1 = spark.createDataFrame(data=[(\"1\", 1000), (\"2\", 3000)], schema=[\"id\", \"amount\"])\n", + "df2 = spark.createDataFrame(data=[(\"1\", 1000), (\"2\", 3000)], schema=[\"id\", \"amount\"])\n", + "assertDataFrameEqual(df1, df2) # pass, DataFrames are identical" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "2d77a6be-1e50-4c1a-8a44-85cf7dcec3f3", + "metadata": {}, + "outputs": [], + "source": [ + "# Example 2\n", + "df1 = spark.createDataFrame(data=[(\"1\", 0.1), (\"2\", 3.23)], schema=[\"id\", \"amount\"])\n", + "df2 = spark.createDataFrame(data=[(\"1\", 0.109), (\"2\", 3.23)], schema=[\"id\", \"amount\"])\n", + "assertDataFrameEqual(df1, df2, rtol=1e-1) # pass, DataFrames are approx equal by rtol" + ] + }, + { + "cell_type": "markdown", + "id": "76ade5f2-4a1f-4601-9d2a-80da9da950ff", + "metadata": {}, + "source": [ + "You can also simply compare two DataFrame schemas:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "74393af5-40fb-4d04-87cb-265971ffe6d0", + "metadata": {}, + "outputs": [], + "source": [ + "from pyspark.testing.utils import assertSchemaEqual\n", + "from pyspark.sql.types import StructType, StructField, ArrayType, DoubleType\n", + "\n", + "s1 = StructType([StructField(\"names\", ArrayType(DoubleType(), True), True)])\n", + "s2 = StructType([StructField(\"names\", ArrayType(DoubleType(), True), True)])\n", + "\n", + "assertSchemaEqual(s1, s2) # pass, schemas are identical" + ] + }, + { + "cell_type": "markdown", + "id": "c67be105-f6b1-4083-ad11-9e819331eae8", + "metadata": {}, + "source": [ + "### Option 2: Using [Unit Test](https://docs.python.org/3/library/unittest.html)\n", + "For more complex testing scenarios, you may want to use a testing framework.\n", + "\n", + "One of the most popular testing framework options is unit tests. Let’s walk through how you can use the built-in Python `unittest` library to write PySpark tests. For more information about the `unittest` library, see here: https://docs.python.org/3/library/unittest.html. \n", + "\n", + "First, you will need a Spark session. You can use the `@classmethod` decorator from the `unittest` package to take care of setting up and tearing down a Spark session." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "54093761-0b49-4aee-baec-2d29bcf13f9f", + "metadata": {}, + "outputs": [], + "source": [ + "import unittest\n", + "\n", + "class PySparkTestCase(unittest.TestCase):\n", + " @classmethod\n", + " def setUpClass(cls):\n", + " cls.spark = SparkSession.builder.appName(\"Testing PySpark Example\").getOrCreate() \n", + "\n", + " \n", + " @classmethod\n", + " def tearDownClass(cls):\n", + " cls.spark.stop()" + ] + }, + { + "cell_type": "markdown", + "id": "3de27500-8526-412e-bf09-6927a760c5d7", + "metadata": {}, + "source": [ + "Now let’s write a `unittest` class." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "34feb5e1-944f-4f6b-9c5f-3b0bf68c7d05", + "metadata": {}, + "outputs": [], + "source": [ + "from pyspark.testing.utils import assertDataFrameEqual\n", + "\n", + "class TestTranformation(PySparkTestCase):\n", + " def test_single_space(self):\n", + " sample_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}] \n", + " \n", + " # Create a Spark DataFrame\n", + " original_df = spark.createDataFrame(sample_data)\n", + " \n", + " # Apply the transformation function from before\n", + " transformed_df = remove_extra_spaces(original_df, \"name\")\n", + " \n", + " expected_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}]\n", + " \n", + " expected_df = spark.createDataFrame(expected_data)\n", + " \n", + " assertDataFrameEqual(transformed_df, expected_df)\n" + ] + }, + { + "cell_type": "markdown", + "id": "319a690f-71bd-4886-bd3a-424e866525c2", + "metadata": {}, + "source": [ + "When run, `unittest` will pick up all functions with a name beginning with “test.”" + ] + }, + { + "cell_type": "markdown", + "id": "7d79e53d-cc1e-4fdf-a069-478337bed83d", + "metadata": {}, + "source": [ + "### Option 3: Using [Pytest](https://docs.pytest.org/en/7.1.x/contents.html)\n", + "\n", + "We can also write our tests with `pytest`, which is one of the most popular Python testing frameworks. For more information about `pytest`, see the docs here: https://docs.pytest.org/en/7.1.x/contents.html.\n", + "\n", + "Using a `pytest` fixture allows us to share a spark session across tests, tearing it down when the tests are complete." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "60a4f304-1911-4b4d-8ed9-00ecc8b0890b", + "metadata": {}, + "outputs": [], + "source": [ + "import pytest\n", + "\n", + "@pytest.fixture\n", + "def spark_fixture():\n", + " spark = SparkSession.builder.appName(\"Testing PySpark Example\").getOrCreate()\n", + " yield spark" + ] + }, + { + "cell_type": "markdown", + "id": "fcb4e26a-9bfc-48a5-8aca-538697d66642", + "metadata": {}, + "source": [ + "We can then define our tests like this:" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "fa5db3a1-7305-44b7-ab84-f5ed55fd2ba9", + "metadata": {}, + "outputs": [], + "source": [ + "import pytest\n", + "from pyspark.testing.utils import assertDataFrameEqual\n", + "\n", + "def test_single_space(spark_fixture):\n", + " sample_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}] \n", + " \n", + " # Create a Spark DataFrame\n", + " original_df = spark.createDataFrame(sample_data)\n", + " \n", + " # Apply the transformation function from before\n", + " transformed_df = remove_extra_spaces(original_df, \"name\")\n", + " \n", + " expected_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}]\n", + " \n", + " expected_df = spark.createDataFrame(expected_data)\n", + "\n", + " assertDataFrameEqual(transformed_df, expected_df)" + ] + }, + { + "cell_type": "markdown", + "id": "0fc3f394-3260-4e42-82cf-1a7edc859151", + "metadata": {}, + "source": [ + "When you run your test file with the `pytest` command, it will pick up all functions that have their name beginning with “test.”" + ] + }, + { + "cell_type": "markdown", + "id": "d8f50eee-5d0b-4719-b505-1b3ff05c16e8", + "metadata": {}, + "source": [ + "## Putting It All Together!\n", + "\n", + "Let’s see all the steps together, in a Unit Test example." + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "a2ea9dec-0ac0-4c23-8770-d6cc226d2e97", + "metadata": {}, + "outputs": [], + "source": [ + "# pkg/etl.py\n", + "import unittest\n", + "\n", + "from pyspark.sql import SparkSession \n", + "from pyspark.sql.functions import col\n", + "from pyspark.sql.functions import regexp_replace\n", + "from pyspark.testing.utils import assertDataFrameEqual\n", + "\n", + "# Create a SparkSession \n", + "spark = SparkSession.builder.appName(\"Sample PySpark ETL\").getOrCreate() \n", + "\n", + "sample_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}] \n", + "\n", + "df = spark.createDataFrame(sample_data)\n", + "\n", + "# Define DataFrame transformation function\n", + "def remove_extra_spaces(df, column_name):\n", + " # Remove extra spaces from the specified column using regexp_replace\n", + " df_transformed = df.withColumn(column_name, regexp_replace(col(column_name), \"\\\\s+\", \" \"))\n", + "\n", + " return df_transformed" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "248aede2-feb9-4828-bd9c-8e25e6b194ab", + "metadata": {}, + "outputs": [], + "source": [ + "# pkg/test_etl.py\n", + "import unittest\n", + "\n", + "from pyspark.sql import SparkSession \n", + "\n", + "# Define unit test base class\n", + "class PySparkTestCase(unittest.TestCase):\n", + " @classmethod\n", + " def setUpClass(cls):\n", + " cls.spark = SparkSession.builder.appName(\"Sample PySpark ETL\").getOrCreate() \n", + "\n", + " @classmethod\n", + " def tearDownClass(cls):\n", + " cls.spark.stop()\n", + " \n", + "# Define unit test\n", + "class TestTranformation(PySparkTestCase):\n", + " def test_single_space(self):\n", + " sample_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}] \n", + " \n", + " # Create a Spark DataFrame\n", + " original_df = spark.createDataFrame(sample_data)\n", + " \n", + " # Apply the transformation function from before\n", + " transformed_df = remove_extra_spaces(original_df, \"name\")\n", + " \n", + " expected_data = [{\"name\": \"John D.\", \"age\": 30}, \n", + " {\"name\": \"Alice G.\", \"age\": 25}, \n", + " {\"name\": \"Bob T.\", \"age\": 35}, \n", + " {\"name\": \"Eve A.\", \"age\": 28}]\n", + " \n", + " expected_df = spark.createDataFrame(expected_data)\n", + " \n", + " assertDataFrameEqual(transformed_df, expected_df)" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "a77df5b2-f32e-4d8c-a64b-0078dfa21217", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Ran 1 test in 1.734s\n", + "\n", + "OK\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "unittest.main(argv=[''], verbosity=0, exit=False)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "jupyter-oss-env", + "language": "python", + "name": "jupyter-oss-env" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/index.rst new file mode 100644 index 0000000..b323374 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/index.rst @@ -0,0 +1,166 @@ +.. 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. + +.. PySpark documentation master file + +================= +PySpark Overview +================= + +**Date**: |today| **Version**: |release| + +**Useful links**: +|binder|_ | `GitHub `_ | `Issues `_ | |examples|_ | `Community `_ + +PySpark is the Python API for Apache Spark. It enables you to perform real-time, +large-scale data processing in a distributed environment using Python. It also provides a PySpark +shell for interactively analyzing your data. + +PySpark combines Python's learnability and ease of use with the power of Apache Spark +to enable processing and analysis of data at any size for everyone familiar with Python. + +PySpark supports all of Spark's features such as Spark SQL, +DataFrames, Structured Streaming, Machine Learning (MLlib) and Spark Core. + +.. list-table:: + :widths: 10 20 20 20 20 10 + :header-rows: 0 + :class: borderless spec_table + + * - + - .. image:: ../../../docs/img/pyspark-spark_sql_and_dataframes.png + :target: reference/pyspark.sql/index.html + :width: 100% + :alt: Spark SQL + - .. image:: ../../../docs/img/pyspark-pandas_api_on_spark.png + :target: reference/pyspark.pandas/index.html + :width: 100% + :alt: Pandas API on Spark + - .. image:: ../../../docs/img/pyspark-structured_streaming.png + :target: reference/pyspark.ss/index.html + :width: 100% + :alt: Streaming + - .. image:: ../../../docs/img/pyspark-machine_learning.png + :target: reference/pyspark.ml.html + :width: 100% + :alt: Machine Learning + - + +.. list-table:: + :widths: 10 80 10 + :header-rows: 0 + :class: borderless spec_table + + * - + - .. image:: ../../../docs/img/pyspark-spark_core_and_rdds.png + :target: reference/pyspark.html + :width: 100% + :alt: Spark Core and RDDs + - + +.. _Index Page - Spark SQL and DataFrames: + +**Spark SQL and DataFrames** + +Spark SQL is Apache Spark's module for working with structured data. +It allows you to seamlessly mix SQL queries with Spark programs. +With PySpark DataFrames you can efficiently read, write, transform, +and analyze data using Python and SQL. +Whether you use Python or SQL, the same underlying execution +engine is used so you will always leverage the full power of Spark. + +- :ref:`/getting_started/quickstart_df.ipynb` +- |binder_df|_ +- :ref:`Spark SQL API Reference` + +**Pandas API on Spark** + +Pandas API on Spark allows you to scale your pandas workload to any size +by running it distributed across multiple nodes. If you are already familiar +with pandas and want to leverage Spark for big data, pandas API on Spark makes +you immediately productive and lets you migrate your applications without modifying the code. +You can have a single codebase that works both with pandas (tests, smaller datasets) +and with Spark (production, distributed datasets) and you can switch between the +pandas API and the Pandas API on Spark easily and without overhead. + +Pandas API on Spark aims to make the transition from pandas to Spark easy but +if you are new to Spark or deciding which API to use, we recommend using PySpark +(see :ref:`Spark SQL and DataFrames `). + +- :ref:`/getting_started/quickstart_ps.ipynb` +- |binder_ps|_ +- :ref:`Pandas API on Spark Reference` + +.. _Index Page - Structured Streaming: + +**Structured Streaming** + +Structured Streaming is a scalable and fault-tolerant stream processing engine built on the Spark SQL engine. +You can express your streaming computation the same way you would express a batch computation on static data. +The Spark SQL engine will take care of running it incrementally and continuously and updating the final result +as streaming data continues to arrive. + +- `Structured Streaming Programming Guide `_ +- :ref:`Structured Streaming API Reference` + +**Machine Learning (MLlib)** + +Built on top of Spark, MLlib is a scalable machine learning library that provides +a uniform set of high-level APIs that help users create and tune practical machine +learning pipelines. + +- `Machine Learning Library (MLlib) Programming Guide `_ +- :ref:`Machine Learning (MLlib) API Reference` + +**Spark Core and RDDs** + +Spark Core is the underlying general execution engine for the Spark platform that all +other functionality is built on top of. It provides RDDs (Resilient Distributed Datasets) +and in-memory computing capabilities. + +Note that the RDD API is a low-level API which can be difficult to use and you do not get +the benefit of Spark's automatic query optimization capabilities. +We recommend using DataFrames (see :ref:`Spark SQL and DataFrames ` above) +instead of RDDs as it allows you to express what you want more easily and lets Spark automatically +construct the most efficient query for you. + +- :ref:`Spark Core API Reference` + +**Spark Streaming (Legacy)** + +Spark Streaming is an extension of the core Spark API that enables scalable, +high-throughput, fault-tolerant stream processing of live data streams. + +Note that Spark Streaming is the previous generation of Spark's streaming engine. +It is a legacy project and it is no longer being updated. +There is a newer and easier to use streaming engine in Spark called +:ref:`Structured Streaming ` which you +should use for your streaming applications and pipelines. + +- `Spark Streaming Programming Guide (Legacy) `_ +- :ref:`Spark Streaming API Reference (Legacy)` + +.. toctree:: + :maxdepth: 2 + :hidden: + + Overview + getting_started/index + user_guide/index + reference/index + development/index + migration_guide/index diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/index.rst new file mode 100644 index 0000000..b5ccaee --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/index.rst @@ -0,0 +1,51 @@ +.. 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. + + +================ +Migration Guides +================ + +This page includes links to guides that will help you migrate to PySpark. + +If you are upgrading from an older to a newer version of PySpark, refer to +the following page for differences between versions: + +- :ref:`Upgrading PySpark` + +.. toctree:: + :maxdepth: 2 + :hidden: + + pyspark_upgrade + +If you are migrating code from `Koalas `_ to PySpark, +this guide is for you: + +.. toctree:: + :maxdepth: 2 + + koalas_to_pyspark + +A lot of content in other migration guides can also be helpful when migrating to newer PySpark versions because PySpark internally shares the same components. +Please also refer to the following migration guides: + +- `Migration Guide: SQL, Datasets and DataFrame `_ +- `Migration Guide: Structured Streaming `_ +- `Migration Guide: MLlib (Machine Learning) `_ +- `Migration Guide: Spark Core `_ + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/koalas_to_pyspark.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/koalas_to_pyspark.rst new file mode 100644 index 0000000..c1c1d1e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/koalas_to_pyspark.rst @@ -0,0 +1,39 @@ +.. 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. + + +============================================ +Migrating from Koalas to pandas API on Spark +============================================ + +* The package name to import should be changed to ``pyspark.pandas`` from ``databricks.koalas``. + + .. code-block:: python + + # import databricks.koalas as ks + import pyspark.pandas as ps + +* ``DataFrame.koalas`` in Koalas DataFrame was renamed to ``DataFrame.pandas_on_spark`` in pandas-on-Spark DataFrame. ``DataFrame.koalas`` was kept for compatibility reasons but deprecated as of Spark 3.2. + ``DataFrame.koalas`` will be removed in the future releases. + +* Monkey-patched ``DataFrame.to_koalas`` in PySpark DataFrame was renamed to ``DataFrame.pandas_api`` in PySpark DataFrame. ``DataFrame.to_koalas`` was kept for compatibility reasons. + ``DataFrame.to_koalas`` will be removed in the future releases. + +* Monkey-patched ``DataFrame.to_pandas_on_spark`` in PySpark DataFrame was renamed to ``DataFrame.pandas_api`` in PySpark DataFrame. ``DataFrame.to_pandas_on_spark`` was kept for compatibility reasons but deprecated as of Spark 3.3. + ``DataFrame.to_pandas_on_spark`` will be removed in the future releases. + +* ``databricks.koalas.__version__`` was removed. ``pyspark.__version__`` should be used instead. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/pyspark_upgrade.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/pyspark_upgrade.rst new file mode 100644 index 0000000..7513d64 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/migration_guide/pyspark_upgrade.rst @@ -0,0 +1,113 @@ +.. 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. + +================== +Upgrading PySpark +================== + +Upgrading from PySpark 3.3 to 3.4 +--------------------------------- + +* In Spark 3.4, the schema of an array column is inferred by merging the schemas of all elements in the array. To restore the previous behavior where the schema is only inferred from the first element, you can set ``spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled`` to ``true``. +* In Spark 3.4, if Pandas on Spark API ``Groupby.apply``'s ``func`` parameter return type is not specified and ``compute.shortcut_limit`` is set to 0, the sampling rows will be set to 2 (ensure sampling rows always >= 2) to make sure infer schema is accurate. +* In Spark 3.4, if Pandas on Spark API ``Index.insert`` is out of bounds, will raise IndexError with ``index {} is out of bounds for axis 0 with size {}`` to follow pandas 1.4 behavior. +* In Spark 3.4, the series name will be preserved in Pandas on Spark API ``Series.mode`` to follow pandas 1.4 behavior. +* In Spark 3.4, the Pandas on Spark API ``Index.__setitem__`` will first to check ``value`` type is ``Column`` type to avoid raising unexpected ``ValueError`` in ``is_list_like`` like `Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions.`. +* In Spark 3.4, the Pandas on Spark API ``astype('category')`` will also refresh ``categories.dtype`` according to original data ``dtype`` to follow pandas 1.4 behavior. +* In Spark 3.4, the Pandas on Spark API supports groupby positional indexing in ``GroupBy.head`` and ``GroupBy.tail`` to follow pandas 1.4. Negative arguments now work correctly and result in ranges relative to the end and start of each group, Previously, negative arguments returned empty frames. +* In Spark 3.4, the infer schema process of ``groupby.apply`` in Pandas on Spark, will first infer the pandas type to ensure the accuracy of the pandas ``dtype`` as much as possible. +* In Spark 3.4, the ``Series.concat`` sort parameter will be respected to follow pandas 1.4 behaviors. +* In Spark 3.4, the ``DataFrame.__setitem__`` will make a copy and replace pre-existing arrays, which will NOT be over-written to follow pandas 1.4 behaviors. +* In Spark 3.4, the ``SparkSession.sql`` and the Pandas on Spark API ``sql`` have got new parameter ``args`` which provides binding of named parameters to their SQL literals. +* In Spark 3.4, Pandas API on Spark follows for the pandas 2.0, and some APIs were deprecated or removed in Spark 3.4 according to the changes made in pandas 2.0. Please refer to the [release notes of pandas](https://pandas.pydata.org/docs/dev/whatsnew/) for more details. +* In Spark 3.4, the custom monkey-patch of ``collections.namedtuple`` was removed, and ``cloudpickle`` was used by default. To restore the previous behavior for any relevant pickling issue of ``collections.namedtuple``, set ``PYSPARK_ENABLE_NAMEDTUPLE_PATCH`` environment variable to ``1``. + + +Upgrading from PySpark 3.2 to 3.3 +--------------------------------- + +* In Spark 3.3, the ``pyspark.pandas.sql`` method follows [the standard Python string formatter](https://docs.python.org/3/library/string.html#format-string-syntax). To restore the previous behavior, set ``PYSPARK_PANDAS_SQL_LEGACY`` environment variable to ``1``. +* In Spark 3.3, the ``drop`` method of pandas API on Spark DataFrame supports dropping rows by ``index``, and sets dropping by index instead of column by default. +* In Spark 3.3, PySpark upgrades Pandas version, the new minimum required version changes from 0.23.2 to 1.0.5. +* In Spark 3.3, the ``repr`` return values of SQL DataTypes have been changed to yield an object with the same value when passed to ``eval``. + + +Upgrading from PySpark 3.1 to 3.2 +--------------------------------- + +* In Spark 3.2, the PySpark methods from sql, ml, spark_on_pandas modules raise the ``TypeError`` instead of ``ValueError`` when are applied to an param of inappropriate type. +* In Spark 3.2, the traceback from Python UDFs, pandas UDFs and pandas function APIs are simplified by default without the traceback from the internal Python workers. In Spark 3.1 or earlier, the traceback from Python workers was printed out. To restore the behavior before Spark 3.2, you can set ``spark.sql.execution.pyspark.udf.simplifiedTraceback.enabled`` to ``false``. +* In Spark 3.2, pinned thread mode is enabled by default to map each Python thread to the corresponding JVM thread. Previously, + one JVM thread could be reused for multiple Python threads, which resulted in one JVM thread local being shared to multiple Python threads. + Also, note that now ``pyspark.InheritableThread`` or ``pyspark.inheritable_thread_target`` is recommended to use together for a Python thread + to properly inherit the inheritable attributes such as local properties in a JVM thread, and to avoid a potential resource leak issue. + To restore the behavior before Spark 3.2, you can set ``PYSPARK_PIN_THREAD`` environment variable to ``false``. + + +Upgrading from PySpark 2.4 to 3.0 +--------------------------------- + +* In Spark 3.0, PySpark requires a pandas version of 0.23.2 or higher to use pandas related functionality, such as ``toPandas``, ``createDataFrame`` from pandas DataFrame, and so on. +* In Spark 3.0, PySpark requires a PyArrow version of 0.12.1 or higher to use PyArrow related functionality, such as ``pandas_udf``, ``toPandas`` and ``createDataFrame`` with "spark.sql.execution.arrow.enabled=true", etc. +* In PySpark, when creating a ``SparkSession`` with ``SparkSession.builder.getOrCreate()``, if there is an existing ``SparkContext``, the builder was trying to update the ``SparkConf`` of the existing ``SparkContext`` with configurations specified to the builder, but the ``SparkContext`` is shared by all ``SparkSession`` s, so we should not update them. In 3.0, the builder comes to not update the configurations. This is the same behavior as Java/Scala API in 2.3 and above. If you want to update them, you need to update them prior to creating a ``SparkSession``. +* In PySpark, when Arrow optimization is enabled, if Arrow version is higher than 0.11.0, Arrow can perform safe type conversion when converting pandas.Series to an Arrow array during serialization. Arrow raises errors when detecting unsafe type conversions like overflow. You enable it by setting ``spark.sql.execution.pandas.convertToArrowArraySafely`` to true. The default setting is false. PySpark behavior for Arrow versions is illustrated in the following table: + + ======================================= ================ ========================= + PyArrow version Integer overflow Floating point truncation + ======================================= ================ ========================= + 0.11.0 and below Raise error Silently allows + > 0.11.0, arrowSafeTypeConversion=false Silent overflow Silently allows + > 0.11.0, arrowSafeTypeConversion=true Raise error Raise error + ======================================= ================ ========================= + +* In Spark 3.0, ``createDataFrame(..., verifySchema=True)`` validates LongType as well in PySpark. Previously, LongType was not verified and resulted in None in case the value overflows. To restore this behavior, verifySchema can be set to False to disable the validation. +* As of Spark 3.0, ``Row`` field names are no longer sorted alphabetically when constructing with named arguments for Python versions 3.6 and above, and the order of fields will match that as entered. To enable sorted fields by default, as in Spark 2.4, set the environment variable ``PYSPARK_ROW_FIELD_SORTING_ENABLED`` to true for both executors and driver - this environment variable must be consistent on all executors and driver; otherwise, it may cause failures or incorrect answers. For Python versions less than 3.6, the field names will be sorted alphabetically as the only option. +* In Spark 3.0, ``pyspark.ml.param.shared.Has*`` mixins do not provide any ``set*(self, value)`` setter methods anymore, use the respective ``self.set(self.*, value)`` instead. See `SPARK-29093 `_ for details. + + +Upgrading from PySpark 2.3 to 2.4 +--------------------------------- + +* In PySpark, when Arrow optimization is enabled, previously ``toPandas`` just failed when Arrow optimization is unable to be used whereas ``createDataFrame`` from Pandas DataFrame allowed the fallback to non-optimization. Now, both ``toPandas`` and ``createDataFrame`` from Pandas DataFrame allow the fallback by default, which can be switched off by ``spark.sql.execution.arrow.fallback.enabled``. + + +Upgrading from PySpark 2.3.0 to 2.3.1 and above +----------------------------------------------- + +* As of version 2.3.1 Arrow functionality, including ``pandas_udf`` and ``toPandas()``/``createDataFrame()`` with ``spark.sql.execution.arrow.enabled`` set to ``True``, has been marked as experimental. These are still evolving and not currently recommended for use in production. + + +Upgrading from PySpark 2.2 to 2.3 +--------------------------------- + +* In PySpark, now we need Pandas 0.19.2 or upper if you want to use Pandas related functionalities, such as ``toPandas``, ``createDataFrame`` from Pandas DataFrame, etc. +* In PySpark, the behavior of timestamp values for Pandas related functionalities was changed to respect session timezone. If you want to use the old behavior, you need to set a configuration ``spark.sql.execution.pandas.respectSessionTimeZone`` to False. See `SPARK-22395 `_ for details. +* In PySpark, ``na.fill()`` or ``fillna`` also accepts boolean and replaces nulls with booleans. In prior Spark versions, PySpark just ignores it and returns the original Dataset/DataFrame. +* In PySpark, ``df.replace`` does not allow to omit value when ``to_replace`` is not a dictionary. Previously, value could be omitted in the other cases and had None by default, which is counterintuitive and error-prone. + + +Upgrading from PySpark 1.4 to 1.5 +--------------------------------- + +* Resolution of strings to columns in Python now supports using dots (.) to qualify the column or access nested values. For example ``df['table.column.nestedField']``. However, this means that if your column name contains any dots you must now escape them using backticks (e.g., ``table.`column.with.dots`.nested``). +* DataFrame.withColumn method in PySpark supports adding a new column or replacing existing columns of the same name. + + +Upgrading from PySpark 1.0-1.2 to 1.3 +------------------------------------- + +* When using DataTypes in Python you will need to construct them (i.e. ``StringType()``) instead of referencing a singleton. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/index.rst new file mode 100644 index 0000000..ed3eb4d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/index.rst @@ -0,0 +1,39 @@ +.. 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. + + +============= +API Reference +============= + +This page lists an overview of all public PySpark modules, classes, functions and methods. + +Pandas API on Spark follows the API specifications of latest pandas release. + +.. toctree:: + :maxdepth: 2 + + pyspark.sql/index + pyspark.pandas/index + pyspark.ss/index + pyspark.ml + pyspark.streaming + pyspark.mllib + pyspark + pyspark.resource + pyspark.errors + pyspark.testing diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.errors.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.errors.rst new file mode 100644 index 0000000..13db9bd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.errors.rst @@ -0,0 +1,53 @@ +.. 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. + + +====== +Errors +====== + +Classes +------- + +.. currentmodule:: pyspark.errors + +.. autosummary:: + :toctree: api/ + + PySparkException + AnalysisException + TempTableAlreadyExistsException + ParseException + IllegalArgumentException + StreamingQueryException + QueryExecutionException + PythonException + UnknownException + SparkUpgradeException + + +Methods +------- + +.. currentmodule:: pyspark.errors + +.. autosummary:: + :toctree: api/ + + PySparkException.getErrorClass + PySparkException.getMessageParameters + PySparkException.getSqlState diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ml.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ml.rst new file mode 100644 index 0000000..965cbe7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ml.rst @@ -0,0 +1,384 @@ +.. 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. + + +MLlib (DataFrame-based) +======================= + +Pipeline APIs +------------- + +.. currentmodule:: pyspark.ml + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Transformer + UnaryTransformer + Estimator + Model + Predictor + PredictionModel + Pipeline + PipelineModel + + +Parameters +---------- + +.. currentmodule:: pyspark.ml.param + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Param + Params + TypeConverters + + +Feature +------- + +.. currentmodule:: pyspark.ml.feature + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Binarizer + BucketedRandomProjectionLSH + BucketedRandomProjectionLSHModel + Bucketizer + ChiSqSelector + ChiSqSelectorModel + CountVectorizer + CountVectorizerModel + DCT + ElementwiseProduct + FeatureHasher + HashingTF + IDF + IDFModel + Imputer + ImputerModel + IndexToString + Interaction + MaxAbsScaler + MaxAbsScalerModel + MinHashLSH + MinHashLSHModel + MinMaxScaler + MinMaxScalerModel + NGram + Normalizer + OneHotEncoder + OneHotEncoderModel + PCA + PCAModel + PolynomialExpansion + QuantileDiscretizer + RobustScaler + RobustScalerModel + RegexTokenizer + RFormula + RFormulaModel + SQLTransformer + StandardScaler + StandardScalerModel + StopWordsRemover + StringIndexer + StringIndexerModel + Tokenizer + UnivariateFeatureSelector + UnivariateFeatureSelectorModel + VarianceThresholdSelector + VarianceThresholdSelectorModel + VectorAssembler + VectorIndexer + VectorIndexerModel + VectorSizeHint + VectorSlicer + Word2Vec + Word2VecModel + + +Classification +-------------- + +.. currentmodule:: pyspark.ml.classification + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + LinearSVC + LinearSVCModel + LinearSVCSummary + LinearSVCTrainingSummary + LogisticRegression + LogisticRegressionModel + LogisticRegressionSummary + LogisticRegressionTrainingSummary + BinaryLogisticRegressionSummary + BinaryLogisticRegressionTrainingSummary + DecisionTreeClassifier + DecisionTreeClassificationModel + GBTClassifier + GBTClassificationModel + RandomForestClassifier + RandomForestClassificationModel + RandomForestClassificationSummary + RandomForestClassificationTrainingSummary + BinaryRandomForestClassificationSummary + BinaryRandomForestClassificationTrainingSummary + NaiveBayes + NaiveBayesModel + MultilayerPerceptronClassifier + MultilayerPerceptronClassificationModel + MultilayerPerceptronClassificationSummary + MultilayerPerceptronClassificationTrainingSummary + OneVsRest + OneVsRestModel + FMClassifier + FMClassificationModel + FMClassificationSummary + FMClassificationTrainingSummary + + +Clustering +---------- + +.. currentmodule:: pyspark.ml.clustering + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + BisectingKMeans + BisectingKMeansModel + BisectingKMeansSummary + KMeans + KMeansModel + KMeansSummary + GaussianMixture + GaussianMixtureModel + GaussianMixtureSummary + LDA + LDAModel + LocalLDAModel + DistributedLDAModel + PowerIterationClustering + + +Functions +--------- + +.. currentmodule:: pyspark.ml.functions + +.. autosummary:: + :toctree: api/ + + array_to_vector + vector_to_array + predict_batch_udf + + +Vector and Matrix +----------------- + +.. currentmodule:: pyspark.ml.linalg + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Vector + DenseVector + SparseVector + Vectors + Matrix + DenseMatrix + SparseMatrix + Matrices + + +Recommendation +-------------- + +.. currentmodule:: pyspark.ml.recommendation + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + ALS + ALSModel + + +Regression +---------- + +.. currentmodule:: pyspark.ml.regression + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + AFTSurvivalRegression + AFTSurvivalRegressionModel + DecisionTreeRegressor + DecisionTreeRegressionModel + GBTRegressor + GBTRegressionModel + GeneralizedLinearRegression + GeneralizedLinearRegressionModel + GeneralizedLinearRegressionSummary + GeneralizedLinearRegressionTrainingSummary + IsotonicRegression + IsotonicRegressionModel + LinearRegression + LinearRegressionModel + LinearRegressionSummary + LinearRegressionTrainingSummary + RandomForestRegressor + RandomForestRegressionModel + FMRegressor + FMRegressionModel + + +Statistics +---------- + +.. currentmodule:: pyspark.ml.stat + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + ChiSquareTest + Correlation + KolmogorovSmirnovTest + MultivariateGaussian + Summarizer + SummaryBuilder + + +Tuning +------ + +.. currentmodule:: pyspark.ml.tuning + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + ParamGridBuilder + CrossValidator + CrossValidatorModel + TrainValidationSplit + TrainValidationSplitModel + + +Evaluation +---------- + +.. currentmodule:: pyspark.ml.evaluation + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Evaluator + BinaryClassificationEvaluator + RegressionEvaluator + MulticlassClassificationEvaluator + MultilabelClassificationEvaluator + ClusteringEvaluator + RankingEvaluator + + +Frequency Pattern Mining +---------------------------- + +.. currentmodule:: pyspark.ml.fpm + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + FPGrowth + FPGrowthModel + PrefixSpan + + +Image +----- + +.. currentmodule:: pyspark.ml.image + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + ImageSchema + _ImageSchema + + +Distributor +----------- + + +.. currentmodule:: pyspark.ml.torch.distributor + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + TorchDistributor + +.. currentmodule:: pyspark.ml.deepspeed.deepspeed_distributor +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + DeepspeedTorchDistributor + + +Utilities +--------- + +.. currentmodule:: pyspark.ml.util + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + BaseReadWrite + DefaultParamsReadable + DefaultParamsReader + DefaultParamsWritable + DefaultParamsWriter + GeneralMLWriter + HasTrainingSummary + Identifiable + MLReadable + MLReader + MLWritable + MLWriter + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.mllib.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.mllib.rst new file mode 100644 index 0000000..12fc479 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.mllib.rst @@ -0,0 +1,254 @@ +.. 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. + + +MLlib (RDD-based) +================= + +Classification +-------------- + +.. currentmodule:: pyspark.mllib.classification + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + LogisticRegressionModel + LogisticRegressionWithSGD + LogisticRegressionWithLBFGS + SVMModel + SVMWithSGD + NaiveBayesModel + NaiveBayes + StreamingLogisticRegressionWithSGD + + +Clustering +---------- + +.. currentmodule:: pyspark.mllib.clustering + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + + BisectingKMeansModel + BisectingKMeans + KMeansModel + KMeans + GaussianMixtureModel + GaussianMixture + PowerIterationClusteringModel + PowerIterationClustering + StreamingKMeans + StreamingKMeansModel + LDA + LDAModel + + +Evaluation +---------- + +.. currentmodule:: pyspark.mllib.evaluation + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + BinaryClassificationMetrics + RegressionMetrics + MulticlassMetrics + RankingMetrics + + +Feature +------- + +.. currentmodule:: pyspark.mllib.feature + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Normalizer + StandardScalerModel + StandardScaler + HashingTF + IDFModel + IDF + Word2Vec + Word2VecModel + ChiSqSelector + ChiSqSelectorModel + ElementwiseProduct + + +Frequency Pattern Mining +------------------------ + +.. currentmodule:: pyspark.mllib.fpm + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + FPGrowth + FPGrowthModel + PrefixSpan + PrefixSpanModel + + +Vector and Matrix +----------------- + +.. currentmodule:: pyspark.mllib.linalg + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Vector + DenseVector + SparseVector + Vectors + Matrix + DenseMatrix + SparseMatrix + Matrices + QRDecomposition + + +Distributed Representation +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.mllib.linalg.distributed + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + BlockMatrix + CoordinateMatrix + DistributedMatrix + IndexedRow + IndexedRowMatrix + MatrixEntry + RowMatrix + SingularValueDecomposition + + +Random +------ + +.. currentmodule:: pyspark.mllib.random + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + RandomRDDs + + +Recommendation +-------------- + +.. currentmodule:: pyspark.mllib.recommendation + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + MatrixFactorizationModel + ALS + Rating + + +Regression +---------- + +.. currentmodule:: pyspark.mllib.regression + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + LabeledPoint + LinearModel + LinearRegressionModel + LinearRegressionWithSGD + RidgeRegressionModel + RidgeRegressionWithSGD + LassoModel + LassoWithSGD + IsotonicRegressionModel + IsotonicRegression + StreamingLinearAlgorithm + StreamingLinearRegressionWithSGD + + +Statistics +---------- + +.. currentmodule:: pyspark.mllib.stat + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + Statistics + MultivariateStatisticalSummary + ChiSqTestResult + MultivariateGaussian + KernelDensity + ChiSqTestResult + KolmogorovSmirnovTestResult + + +Tree +---- + +.. currentmodule:: pyspark.mllib.tree + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + DecisionTreeModel + DecisionTree + RandomForestModel + RandomForest + GradientBoostedTreesModel + GradientBoostedTrees + + +Utilities +--------- + +.. currentmodule:: pyspark.mllib.util + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + JavaLoader + JavaSaveable + LinearDataGenerator + Loader + MLUtils + Saveable diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/extensions.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/extensions.rst new file mode 100644 index 0000000..8e59f24 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/extensions.rst @@ -0,0 +1,39 @@ +.. 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. + + +.. _api.extensions: + +========== +Extensions +========== +.. currentmodule:: pyspark.pandas.extensions + +Accessors +--------- + +Accessors can be written and registered with pandas-on-Spark Dataframes, Series, and +Index objects. Accessors allow developers to extend the functionality of +pandas-on-Spark objects seamlessly by writing arbitrary classes and methods which are +then wrapped in one of the following decorators. + +.. autosummary:: + :toctree: api/ + + register_dataframe_accessor + register_series_accessor + register_index_accessor diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/frame.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/frame.rst new file mode 100644 index 0000000..a8d1141 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/frame.rst @@ -0,0 +1,349 @@ +.. 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. + + +.. _api.dataframe: + +========= +DataFrame +========= +.. currentmodule:: pyspark.pandas + +Constructor +----------- +.. autosummary:: + :toctree: api/ + + DataFrame + +Attributes and underlying data +------------------------------ + +.. autosummary:: + :toctree: api/ + + DataFrame.index + DataFrame.info + DataFrame.columns + DataFrame.empty + +.. autosummary:: + :toctree: api/ + + DataFrame.dtypes + DataFrame.shape + DataFrame.axes + DataFrame.ndim + DataFrame.size + DataFrame.select_dtypes + DataFrame.values + +Conversion +---------- +.. autosummary:: + :toctree: api/ + + DataFrame.copy + DataFrame.isna + DataFrame.astype + DataFrame.isnull + DataFrame.notna + DataFrame.notnull + DataFrame.bool + +Indexing, iteration +------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.at + DataFrame.iat + DataFrame.head + DataFrame.idxmax + DataFrame.idxmin + DataFrame.loc + DataFrame.iloc + DataFrame.insert + DataFrame.items + DataFrame.iteritems + DataFrame.iterrows + DataFrame.itertuples + DataFrame.keys + DataFrame.pop + DataFrame.tail + DataFrame.xs + DataFrame.get + DataFrame.where + DataFrame.mask + DataFrame.query + +Binary operator functions +------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.add + DataFrame.radd + DataFrame.div + DataFrame.rdiv + DataFrame.truediv + DataFrame.rtruediv + DataFrame.mul + DataFrame.rmul + DataFrame.sub + DataFrame.rsub + DataFrame.pow + DataFrame.rpow + DataFrame.mod + DataFrame.rmod + DataFrame.floordiv + DataFrame.rfloordiv + DataFrame.lt + DataFrame.gt + DataFrame.le + DataFrame.ge + DataFrame.ne + DataFrame.eq + DataFrame.dot + DataFrame.combine_first + +Function application, GroupBy & Window +-------------------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.apply + DataFrame.applymap + DataFrame.pipe + DataFrame.agg + DataFrame.aggregate + DataFrame.groupby + DataFrame.rolling + DataFrame.expanding + DataFrame.transform + +.. _api.dataframe.stats: + +Computations / Descriptive Stats +-------------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.abs + DataFrame.all + DataFrame.any + DataFrame.clip + DataFrame.corr + DataFrame.corrwith + DataFrame.count + DataFrame.cov + DataFrame.describe + DataFrame.ewm + DataFrame.kurt + DataFrame.kurtosis + DataFrame.mad + DataFrame.max + DataFrame.mean + DataFrame.min + DataFrame.median + DataFrame.mode + DataFrame.pct_change + DataFrame.prod + DataFrame.product + DataFrame.quantile + DataFrame.rank + DataFrame.nunique + DataFrame.sem + DataFrame.skew + DataFrame.sum + DataFrame.std + DataFrame.var + DataFrame.cummin + DataFrame.cummax + DataFrame.cumsum + DataFrame.cumprod + DataFrame.round + DataFrame.diff + DataFrame.eval + +Reindexing / Selection / Label manipulation +------------------------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.add_prefix + DataFrame.add_suffix + DataFrame.align + DataFrame.at_time + DataFrame.between_time + DataFrame.drop + DataFrame.droplevel + DataFrame.drop_duplicates + DataFrame.duplicated + DataFrame.equals + DataFrame.filter + DataFrame.first + DataFrame.head + DataFrame.last + DataFrame.reindex + DataFrame.reindex_like + DataFrame.rename + DataFrame.rename_axis + DataFrame.reset_index + DataFrame.set_index + DataFrame.swapaxes + DataFrame.swaplevel + DataFrame.take + DataFrame.isin + DataFrame.sample + DataFrame.truncate + +.. _api.dataframe.missing: + +Missing data handling +--------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.backfill + DataFrame.dropna + DataFrame.fillna + DataFrame.replace + DataFrame.bfill + DataFrame.ffill + DataFrame.interpolate + DataFrame.pad + +Reshaping, sorting, transposing +------------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.pivot_table + DataFrame.pivot + DataFrame.sort_index + DataFrame.sort_values + DataFrame.nlargest + DataFrame.nsmallest + DataFrame.stack + DataFrame.unstack + DataFrame.melt + DataFrame.explode + DataFrame.squeeze + DataFrame.T + DataFrame.transpose + +Combining / joining / merging +----------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.append + DataFrame.assign + DataFrame.merge + DataFrame.join + DataFrame.update + +Time series-related +------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.resample + DataFrame.shift + DataFrame.first_valid_index + DataFrame.last_valid_index + +Serialization / IO / Conversion +------------------------------- +.. autosummary:: + :toctree: api/ + + DataFrame.from_dict + DataFrame.from_records + DataFrame.to_table + DataFrame.to_delta + DataFrame.to_parquet + DataFrame.to_spark_io + DataFrame.to_csv + DataFrame.to_orc + DataFrame.to_pandas + DataFrame.to_html + DataFrame.to_numpy + DataFrame.to_spark + DataFrame.to_string + DataFrame.to_json + DataFrame.to_dict + DataFrame.to_excel + DataFrame.to_clipboard + DataFrame.to_markdown + DataFrame.to_records + DataFrame.to_latex + DataFrame.style + +Spark-related +------------- +``DataFrame.spark`` provides features that does not exist in pandas but +in Spark. These can be accessed by ``DataFrame.spark.``. + +.. autosummary:: + :toctree: api/ + + DataFrame.spark.frame + DataFrame.spark.cache + DataFrame.spark.persist + DataFrame.spark.hint + DataFrame.spark.to_table + DataFrame.spark.to_spark_io + DataFrame.spark.apply + DataFrame.spark.repartition + DataFrame.spark.coalesce + +.. _api.dataframe.plot: + +Plotting +-------- +``DataFrame.plot`` is both a callable method and a namespace attribute for +specific plotting methods of the form ``DataFrame.plot.``. + +.. autosummary:: + :toctree: api/ + + DataFrame.plot + DataFrame.plot.area + DataFrame.plot.barh + DataFrame.plot.bar + DataFrame.plot.hist + DataFrame.plot.box + DataFrame.plot.line + DataFrame.plot.pie + DataFrame.plot.scatter + DataFrame.plot.density + DataFrame.hist + DataFrame.boxplot + DataFrame.kde + +Pandas-on-Spark specific +------------------------ +``DataFrame.pandas_on_spark`` provides pandas-on-Spark specific features that exists only in pandas API on Spark. +These can be accessed by ``DataFrame.pandas_on_spark.``. + +.. autosummary:: + :toctree: api/ + + DataFrame.pandas_on_spark.apply_batch + DataFrame.pandas_on_spark.transform_batch diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/general_functions.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/general_functions.rst new file mode 100644 index 0000000..a67c0a8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/general_functions.rst @@ -0,0 +1,77 @@ +.. 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. + + +.. _api.general_functions: + +================= +General functions +================= +.. currentmodule:: pyspark.pandas + +Working with options +-------------------- + +.. autosummary:: + :toctree: api/ + + reset_option + get_option + set_option + option_context + +Data manipulations and SQL +-------------------------- +.. autosummary:: + :toctree: api/ + + melt + merge + merge_asof + get_dummies + concat + sql + broadcast + +Top-level missing data +---------------------- + +.. autosummary:: + :toctree: api/ + + isna + isnull + notna + notnull + +Top-level dealing with numeric data +----------------------------------- + +.. autosummary:: + :toctree: api/ + + to_numeric + +Top-level dealing with datetimelike data +---------------------------------------- +.. autosummary:: + :toctree: api/ + + to_datetime + date_range + to_timedelta + timedelta_range diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/groupby.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/groupby.rst new file mode 100644 index 0000000..da1579f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/groupby.rst @@ -0,0 +1,112 @@ +.. 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. + + +.. _api.groupby: + +======= +GroupBy +======= +.. currentmodule:: pyspark.pandas + +GroupBy objects are returned by groupby calls: :func:`DataFrame.groupby`, :func:`Series.groupby`, etc. + +.. currentmodule:: pyspark.pandas.groupby + + +Indexing, iteration +------------------- +.. autosummary:: + :toctree: api/ + + GroupBy.get_group + +Function application +-------------------- +.. autosummary:: + :toctree: api/ + + GroupBy.apply + GroupBy.transform + +The following methods are available only for `DataFrameGroupBy` objects. + +.. autosummary:: + :toctree: api/ + + DataFrameGroupBy.agg + DataFrameGroupBy.aggregate + +Computations / Descriptive Stats +-------------------------------- +.. autosummary:: + :toctree: api/ + + GroupBy.all + GroupBy.any + GroupBy.count + GroupBy.cumcount + GroupBy.cummax + GroupBy.cummin + GroupBy.cumprod + GroupBy.cumsum + GroupBy.ewm + GroupBy.filter + GroupBy.first + GroupBy.last + GroupBy.mad + GroupBy.max + GroupBy.mean + GroupBy.median + GroupBy.min + GroupBy.nth + GroupBy.prod + GroupBy.rank + GroupBy.sem + GroupBy.std + GroupBy.sum + GroupBy.var + GroupBy.nunique + GroupBy.quantile + GroupBy.size + GroupBy.diff + GroupBy.idxmax + GroupBy.idxmin + GroupBy.fillna + GroupBy.bfill + GroupBy.ffill + GroupBy.head + GroupBy.backfill + GroupBy.shift + GroupBy.tail + +The following methods are available only for `DataFrameGroupBy` objects. + +.. autosummary:: + :toctree: api/ + + DataFrameGroupBy.describe + +The following methods are available only for `SeriesGroupBy` objects. + +.. autosummary:: + :toctree: api/ + + SeriesGroupBy.nsmallest + SeriesGroupBy.nlargest + SeriesGroupBy.value_counts + SeriesGroupBy.unique diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/index.rst new file mode 100644 index 0000000..96da086 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/index.rst @@ -0,0 +1,37 @@ +.. 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. + + +=================== +Pandas API on Spark +=================== + +This page gives an overview of all public pandas API on Spark. + +.. toctree:: + :maxdepth: 2 + + io + general_functions + series + frame + indexing + window + groupby + resampling + ml + extensions diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/indexing.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/indexing.rst new file mode 100644 index 0000000..15539fa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/indexing.rst @@ -0,0 +1,404 @@ +.. 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. + + +.. _api.indexing: + +============= +Index objects +============= + +Index +----- +.. currentmodule:: pyspark.pandas + +.. autosummary:: + :toctree: api/ + + Index + +Properties +~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.is_monotonic + Index.is_monotonic_increasing + Index.is_monotonic_decreasing + Index.is_unique + Index.has_duplicates + Index.hasnans + Index.dtype + Index.inferred_type + Index.is_all_dates + Index.shape + Index.name + Index.names + Index.ndim + Index.size + Index.nlevels + Index.empty + Index.T + Index.values + +Modifying and computations +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autosummary:: + :toctree: api/ + + Index.all + Index.any + Index.argmin + Index.argmax + Index.copy + Index.delete + Index.equals + Index.factorize + Index.identical + Index.insert + Index.is_boolean + Index.is_categorical + Index.is_floating + Index.is_integer + Index.is_interval + Index.is_numeric + Index.is_object + Index.drop + Index.drop_duplicates + Index.min + Index.max + Index.map + Index.rename + Index.repeat + Index.take + Index.unique + Index.nunique + Index.value_counts + +Compatibility with MultiIndex +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autosummary:: + :toctree: api/ + + Index.set_names + Index.droplevel + +Missing Values +~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.fillna + Index.dropna + Index.isna + Index.notna + +Conversion +~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.astype + Index.item + Index.to_list + Index.to_series + Index.to_frame + Index.view + Index.to_numpy + +Spark-related +------------- +``Index.spark`` provides features that does not exist in pandas but +in Spark. These can be accessed by ``Index.spark.``. + +.. autosummary:: + :toctree: api/ + + Index.spark.column + Index.spark.transform + +Sorting +~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.sort_values + +Time-specific operations +~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.shift + +Combining / joining / set operations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.append + Index.intersection + Index.union + Index.difference + Index.symmetric_difference + +Selecting +~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Index.asof + Index.isin + +.. _api.numeric: + +Numeric Index +------------- +.. autosummary:: + :toctree: api/ + + Int64Index + Float64Index + +.. _api.categorical: + +CategoricalIndex +---------------- +.. autosummary:: + :toctree: api/ + + CategoricalIndex + +Categorical components +~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + CategoricalIndex.codes + CategoricalIndex.categories + CategoricalIndex.ordered + CategoricalIndex.rename_categories + CategoricalIndex.reorder_categories + CategoricalIndex.add_categories + CategoricalIndex.remove_categories + CategoricalIndex.remove_unused_categories + CategoricalIndex.set_categories + CategoricalIndex.as_ordered + CategoricalIndex.as_unordered + CategoricalIndex.map + +.. _api.multiindex: + +MultiIndex +---------- +.. autosummary:: + :toctree: api/ + + MultiIndex + +MultiIndex Constructors +~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.from_arrays + MultiIndex.from_tuples + MultiIndex.from_product + MultiIndex.from_frame + +MultiIndex Properties +~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.has_duplicates + MultiIndex.hasnans + MultiIndex.inferred_type + MultiIndex.is_all_dates + MultiIndex.shape + MultiIndex.names + MultiIndex.ndim + MultiIndex.empty + MultiIndex.T + MultiIndex.size + MultiIndex.nlevels + MultiIndex.levshape + MultiIndex.values + MultiIndex.dtypes + +MultiIndex components +~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.swaplevel + +MultiIndex components +~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.droplevel + +MultiIndex Missing Values +~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.fillna + MultiIndex.dropna + +MultiIndex Modifying and computations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.equals + MultiIndex.equal_levels + MultiIndex.identical + MultiIndex.insert + MultiIndex.drop + MultiIndex.copy + MultiIndex.delete + MultiIndex.rename + MultiIndex.repeat + MultiIndex.take + MultiIndex.unique + MultiIndex.min + MultiIndex.max + MultiIndex.value_counts + +MultiIndex Combining / joining / set operations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.append + MultiIndex.intersection + MultiIndex.union + MultiIndex.difference + MultiIndex.symmetric_difference + +MultiIndex Conversion +~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.astype + MultiIndex.item + MultiIndex.to_list + MultiIndex.to_series + MultiIndex.to_frame + MultiIndex.view + MultiIndex.to_numpy + +MultiIndex Spark-related +------------------------ +``MultiIndex.spark`` provides features that does not exist in pandas but +in Spark. These can be accessed by ``MultiIndex.spark.``. + +.. autosummary:: + :toctree: api/ + + MultiIndex.spark.data_type + MultiIndex.spark.column + MultiIndex.spark.transform + +MultiIndex Sorting +~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + MultiIndex.sort_values + +.. _api.datetimes: + +DatatimeIndex +------------- +.. autosummary:: + :toctree: api/ + + DatetimeIndex + +Time/date components +~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + DatetimeIndex.year + DatetimeIndex.month + DatetimeIndex.day + DatetimeIndex.hour + DatetimeIndex.minute + DatetimeIndex.second + DatetimeIndex.microsecond + DatetimeIndex.week + DatetimeIndex.weekofyear + DatetimeIndex.dayofweek + DatetimeIndex.day_of_week + DatetimeIndex.weekday + DatetimeIndex.dayofyear + DatetimeIndex.day_of_year + DatetimeIndex.quarter + DatetimeIndex.is_month_start + DatetimeIndex.is_month_end + DatetimeIndex.is_quarter_start + DatetimeIndex.is_quarter_end + DatetimeIndex.is_year_start + DatetimeIndex.is_year_end + DatetimeIndex.is_leap_year + DatetimeIndex.daysinmonth + DatetimeIndex.days_in_month + +Selecting +~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + DatetimeIndex.indexer_between_time + DatetimeIndex.indexer_at_time + +Time-specific operations +~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + DatetimeIndex.normalize + DatetimeIndex.strftime + DatetimeIndex.round + DatetimeIndex.floor + DatetimeIndex.ceil + DatetimeIndex.month_name + DatetimeIndex.day_name + +TimedeltaIndex +-------------- +.. autosummary:: + :toctree: api/ + + TimedeltaIndex + +Components +~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + TimedeltaIndex.days + TimedeltaIndex.seconds + TimedeltaIndex.microseconds diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/io.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/io.rst new file mode 100644 index 0000000..b39a4e8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/io.rst @@ -0,0 +1,121 @@ +.. 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. + + +.. _api.io: + +============ +Input/Output +============ +.. currentmodule:: pyspark.pandas + + +Data Generator +-------------- +.. autosummary:: + :toctree: api/ + + range + +Spark Metastore Table +--------------------- +.. autosummary:: + :toctree: api/ + + read_table + DataFrame.to_table + +Delta Lake +---------- +.. autosummary:: + :toctree: api/ + + read_delta + DataFrame.to_delta + +Parquet +------- +.. autosummary:: + :toctree: api/ + + read_parquet + DataFrame.to_parquet + +ORC +------- +.. autosummary:: + :toctree: api/ + + read_orc + DataFrame.to_orc + +Generic Spark I/O +----------------- +.. autosummary:: + :toctree: api/ + + read_spark_io + DataFrame.to_spark_io + +Flat File / CSV +--------------- +.. autosummary:: + :toctree: api/ + + read_csv + DataFrame.to_csv + +Clipboard +--------- +.. autosummary:: + :toctree: api/ + + read_clipboard + DataFrame.to_clipboard + +Excel +----- +.. autosummary:: + :toctree: api/ + + read_excel + DataFrame.to_excel + +JSON +---- +.. autosummary:: + :toctree: api/ + + read_json + DataFrame.to_json + +HTML +---- +.. autosummary:: + :toctree: api/ + + read_html + DataFrame.to_html + +SQL +--- +.. autosummary:: + :toctree: api/ + + read_sql_table + read_sql_query + read_sql diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/ml.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/ml.rst new file mode 100644 index 0000000..96dd166 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/ml.rst @@ -0,0 +1,42 @@ +.. 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. + + +.. _api.ml: + +========================== +Machine Learning utilities +========================== +.. currentmodule:: pyspark.pandas.mlflow + +MLflow +------ + +Arbitrary MLflow models can be used with pandas-on-Spark Dataframes, +provided they implement the 'pyfunc' flavor. This is the case +for most frameworks supported by MLflow (scikit-learn, pytorch, +tensorflow, ...). See comprehensive examples in +:func:`load_model` for more information. + +.. note:: + The MLflow package must be installed in order to use this module. + +.. autosummary:: + :toctree: api/ + + PythonModelWrapper + load_model diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/resampling.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/resampling.rst new file mode 100644 index 0000000..dcb9df2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/resampling.rst @@ -0,0 +1,41 @@ +.. 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. + + +.. _api.resampling: + +========== +Resampling +========== +.. currentmodule:: pyspark.pandas + +Resampler objects are returned by ``.resample`` calls: :func:`DataFrame.resample`, :func:`Series.resample`, etc. + +.. currentmodule:: pyspark.pandas.resample + + +Computations / descriptive stats +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + Resampler.max + Resampler.mean + Resampler.min + Resampler.std + Resampler.sum + Resampler.var diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/series.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/series.rst new file mode 100644 index 0000000..a011959 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/series.rst @@ -0,0 +1,485 @@ +.. 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. + + +.. _api.series: + +====== +Series +====== +.. currentmodule:: pyspark.pandas + +Constructor +----------- +.. autosummary:: + :toctree: api/ + + Series + +Attributes +---------- + +.. autosummary:: + :toctree: api/ + + Series.index + Series.dtype + Series.dtypes + Series.ndim + Series.name + Series.shape + Series.axes + Series.size + Series.empty + Series.T + Series.hasnans + Series.values + +Conversion +---------- +.. autosummary:: + :toctree: api/ + + Series.astype + Series.copy + Series.bool + +Indexing, iteration +------------------- +.. autosummary:: + :toctree: api/ + + Series.at + Series.iat + Series.loc + Series.iloc + Series.keys + Series.pop + Series.items + Series.iteritems + Series.item + Series.xs + Series.get + +Binary operator functions +------------------------- + +.. autosummary:: + :toctree: api/ + + Series.add + Series.div + Series.mul + Series.radd + Series.rdiv + Series.rmul + Series.rsub + Series.rtruediv + Series.sub + Series.truediv + Series.pow + Series.rpow + Series.mod + Series.rmod + Series.floordiv + Series.rfloordiv + Series.divmod + Series.rdivmod + Series.combine_first + Series.lt + Series.gt + Series.le + Series.ge + Series.ne + Series.eq + Series.product + Series.dot + +Function application, GroupBy & Window +-------------------------------------- +.. autosummary:: + :toctree: api/ + + Series.apply + Series.agg + Series.aggregate + Series.transform + Series.map + Series.groupby + Series.rolling + Series.expanding + Series.pipe + +.. _api.series.stats: + +Computations / Descriptive Stats +-------------------------------- +.. autosummary:: + :toctree: api/ + + Series.abs + Series.all + Series.any + Series.autocorr + Series.between + Series.clip + Series.corr + Series.count + Series.cov + Series.cummax + Series.cummin + Series.cumsum + Series.cumprod + Series.describe + Series.ewm + Series.filter + Series.kurt + Series.mad + Series.max + Series.mean + Series.min + Series.mode + Series.nlargest + Series.nsmallest + Series.pct_change + Series.prod + Series.nunique + Series.is_unique + Series.quantile + Series.rank + Series.sem + Series.skew + Series.std + Series.sum + Series.median + Series.var + Series.kurtosis + Series.unique + Series.value_counts + Series.round + Series.diff + Series.is_monotonic + Series.is_monotonic_increasing + Series.is_monotonic_decreasing + +Reindexing / Selection / Label manipulation +------------------------------------------- +.. autosummary:: + :toctree: api/ + + Series.align + Series.drop + Series.droplevel + Series.drop_duplicates + Series.duplicated + Series.equals + Series.add_prefix + Series.add_suffix + Series.first + Series.head + Series.idxmax + Series.idxmin + Series.isin + Series.last + Series.rename + Series.rename_axis + Series.reindex + Series.reindex_like + Series.reset_index + Series.sample + Series.searchsorted + Series.swaplevel + Series.swapaxes + Series.take + Series.tail + Series.where + Series.mask + Series.truncate + +Missing data handling +--------------------- +.. autosummary:: + :toctree: api/ + + Series.backfill + Series.bfill + Series.isna + Series.isnull + Series.notna + Series.notnull + Series.pad + Series.dropna + Series.fillna + Series.interpolate + +Reshaping, sorting, transposing +------------------------------- +.. autosummary:: + :toctree: api/ + + Series.argsort + Series.argmin + Series.argmax + Series.sort_index + Series.sort_values + Series.unstack + Series.explode + Series.repeat + Series.squeeze + Series.factorize + +Combining / joining / merging +----------------------------- +.. autosummary:: + :toctree: api/ + + Series.append + Series.compare + Series.replace + Series.update + +Time series-related +------------------- + +.. autosummary:: + :toctree: api/ + + Series.asof + Series.resample + Series.shift + Series.first_valid_index + Series.last_valid_index + Series.at_time + Series.between_time + +Spark-related +------------- +``Series.spark`` provides features that does not exist in pandas but +in Spark. These can be accessed by ``Series.spark.``. + +.. autosummary:: + :toctree: api/ + + Series.spark.column + Series.spark.transform + Series.spark.apply + +Accessors +--------- + +Pandas API on Spark provides dtype-specific methods under various accessors. +These are separate namespaces within :class:`Series` that only apply +to specific data types. + +=========== =========================== +Data Type Accessor +=========== =========================== +Datetime :ref:`dt ` +String :ref:`str ` +Categorical :ref:`cat ` +=========== =========================== + +.. _api.series.dt: + +Date Time Handling +------------------ + +``Series.dt`` can be used to access the values of the series as +datetimelike and return several properties. +These can be accessed like ``Series.dt.``. + +Datetime Properties +~~~~~~~~~~~~~~~~~~~ + +.. autosummary:: + :toctree: api/ + + Series.dt.date + Series.dt.year + Series.dt.month + Series.dt.day + Series.dt.hour + Series.dt.minute + Series.dt.second + Series.dt.microsecond + Series.dt.week + Series.dt.weekofyear + Series.dt.dayofweek + Series.dt.weekday + Series.dt.dayofyear + Series.dt.quarter + Series.dt.is_month_start + Series.dt.is_month_end + Series.dt.is_quarter_start + Series.dt.is_quarter_end + Series.dt.is_year_start + Series.dt.is_year_end + Series.dt.is_leap_year + Series.dt.daysinmonth + Series.dt.days_in_month + +Datetime Methods +~~~~~~~~~~~~~~~~ + +.. autosummary:: + :toctree: api/ + + Series.dt.normalize + Series.dt.strftime + Series.dt.round + Series.dt.floor + Series.dt.ceil + Series.dt.month_name + Series.dt.day_name + +.. _api.series.str: + +String Handling +--------------- + +``Series.str`` can be used to access the values of the series as +strings and apply several methods to it. These can be accessed +like ``Series.str.``. + +.. autosummary:: + :toctree: api/ + + Series.str.capitalize + Series.str.cat + Series.str.center + Series.str.contains + Series.str.count + Series.str.decode + Series.str.encode + Series.str.endswith + Series.str.extract + Series.str.extractall + Series.str.find + Series.str.findall + Series.str.get + Series.str.get_dummies + Series.str.index + Series.str.isalnum + Series.str.isalpha + Series.str.isdigit + Series.str.isspace + Series.str.islower + Series.str.isupper + Series.str.istitle + Series.str.isnumeric + Series.str.isdecimal + Series.str.join + Series.str.len + Series.str.ljust + Series.str.lower + Series.str.lstrip + Series.str.match + Series.str.normalize + Series.str.pad + Series.str.partition + Series.str.repeat + Series.str.replace + Series.str.rfind + Series.str.rindex + Series.str.rjust + Series.str.rpartition + Series.str.rsplit + Series.str.rstrip + Series.str.slice + Series.str.slice_replace + Series.str.split + Series.str.startswith + Series.str.strip + Series.str.swapcase + Series.str.title + Series.str.translate + Series.str.upper + Series.str.wrap + Series.str.zfill + +.. _api.series.cat: + +Categorical accessor +-------------------- +Categorical-dtype specific methods and attributes are available under +the ``Series.cat`` accessor. + +.. autosummary:: + :toctree: api/ + + Series.cat.categories + Series.cat.ordered + Series.cat.codes + Series.cat.rename_categories + Series.cat.reorder_categories + Series.cat.add_categories + Series.cat.remove_categories + Series.cat.remove_unused_categories + Series.cat.set_categories + Series.cat.as_ordered + Series.cat.as_unordered + +.. _api.series.plot: + +Plotting +------------------------------- +``Series.plot`` is both a callable method and a namespace attribute for +specific plotting methods of the form ``Series.plot.``. + +.. autosummary:: + :toctree: api/ + + Series.plot + Series.plot.area + Series.plot.bar + Series.plot.barh + Series.plot.box + Series.plot.density + Series.plot.hist + Series.plot.line + Series.plot.pie + Series.plot.kde + Series.hist + +Serialization / IO / Conversion +------------------------------- +.. autosummary:: + :toctree: api/ + + Series.to_pandas + Series.to_numpy + Series.to_list + Series.to_string + Series.to_dict + Series.to_clipboard + Series.to_latex + Series.to_markdown + Series.to_json + Series.to_csv + Series.to_excel + Series.to_frame + +Pandas-on-Spark specific +------------------------ +``Series.pandas_on_spark`` provides pandas-on-Spark specific features that exists only in pandas API on Spark. +These can be accessed by ``Series.pandas_on_spark.``. + +.. autosummary:: + :toctree: api/ + + Series.pandas_on_spark.transform_batch + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/window.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/window.rst new file mode 100644 index 0000000..c840be3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.pandas/window.rst @@ -0,0 +1,62 @@ +.. 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. + + +====== +Window +====== +.. currentmodule:: pyspark.pandas.window + +Rolling objects are returned by ``.rolling`` calls: :func:`pyspark.pandas.DataFrame.rolling`, :func:`pyspark.pandas.Series.rolling`, etc. + +Expanding objects are returned by ``.expanding`` calls: :func:`pyspark.pandas.DataFrame.expanding`, :func:`pyspark.pandas.Series.expanding`, etc. + +ExponentialMoving objects are returned by ``.ewm`` calls: :func:`pyspark.pandas.DataFrame.ewm`, :func:`pyspark.pandas.Series.ewm`, etc. + +Standard moving window functions +-------------------------------- + +.. autosummary:: + :toctree: api/ + + Rolling.count + Rolling.sum + Rolling.min + Rolling.max + Rolling.mean + Rolling.quantile + +Standard expanding window functions +----------------------------------- + +.. autosummary:: + :toctree: api/ + + Expanding.count + Expanding.sum + Expanding.min + Expanding.max + Expanding.mean + Expanding.quantile + +Exponential moving window functions +----------------------------------- + +.. autosummary:: + :toctree: api/ + + ExponentialMoving.mean diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.resource.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.resource.rst new file mode 100644 index 0000000..a1d885c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.resource.rst @@ -0,0 +1,38 @@ +.. 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. + + +=================== +Resource Management +=================== + +Core Classes +------------ + +.. currentmodule:: pyspark.resource + +.. autosummary:: + :toctree: api/ + + ResourceInformation + ResourceProfile + ResourceProfileBuilder + ExecutorResourceRequest + ExecutorResourceRequests + TaskResourceRequest + TaskResourceRequests + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.rst new file mode 100644 index 0000000..9a6fbb6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.rst @@ -0,0 +1,297 @@ +.. 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. + + +========== +Spark Core +========== + +Public Classes +-------------- + +.. currentmodule:: pyspark + +.. autosummary:: + :toctree: api/ + + SparkContext + RDD + Broadcast + Accumulator + AccumulatorParam + SparkConf + SparkFiles + StorageLevel + TaskContext + RDDBarrier + BarrierTaskContext + BarrierTaskInfo + InheritableThread + util.VersionUtils + +Spark Context APIs +------------------ + +.. currentmodule:: pyspark + +.. autosummary:: + :toctree: api/ + + SparkContext.PACKAGE_EXTENSIONS + SparkContext.accumulator + SparkContext.addArchive + SparkContext.addFile + SparkContext.addJobTag + SparkContext.addPyFile + SparkContext.applicationId + SparkContext.binaryFiles + SparkContext.binaryRecords + SparkContext.broadcast + SparkContext.cancelAllJobs + SparkContext.cancelJobGroup + SparkContext.cancelJobsWithTag + SparkContext.clearJobTags + SparkContext.defaultMinPartitions + SparkContext.defaultParallelism + SparkContext.dump_profiles + SparkContext.emptyRDD + SparkContext.getCheckpointDir + SparkContext.getConf + SparkContext.getJobTags + SparkContext.getLocalProperty + SparkContext.getOrCreate + SparkContext.hadoopFile + SparkContext.hadoopRDD + SparkContext.listArchives + SparkContext.listFiles + SparkContext.newAPIHadoopFile + SparkContext.newAPIHadoopRDD + SparkContext.parallelize + SparkContext.pickleFile + SparkContext.range + SparkContext.resources + SparkContext.removeJobTag + SparkContext.runJob + SparkContext.sequenceFile + SparkContext.setCheckpointDir + SparkContext.setInterruptOnCancel + SparkContext.setJobDescription + SparkContext.setJobGroup + SparkContext.setLocalProperty + SparkContext.setLogLevel + SparkContext.setSystemProperty + SparkContext.show_profiles + SparkContext.sparkUser + SparkContext.startTime + SparkContext.statusTracker + SparkContext.stop + SparkContext.textFile + SparkContext.uiWebUrl + SparkContext.union + SparkContext.version + SparkContext.wholeTextFiles + + +RDD APIs +-------- + +.. currentmodule:: pyspark + +.. autosummary:: + :toctree: api/ + + RDD.aggregate + RDD.aggregateByKey + RDD.barrier + RDD.cache + RDD.cartesian + RDD.checkpoint + RDD.cleanShuffleDependencies + RDD.coalesce + RDD.cogroup + RDD.collect + RDD.collectAsMap + RDD.collectWithJobGroup + RDD.combineByKey + RDD.context + RDD.count + RDD.countApprox + RDD.countApproxDistinct + RDD.countByKey + RDD.countByValue + RDD.distinct + RDD.filter + RDD.first + RDD.flatMap + RDD.flatMapValues + RDD.fold + RDD.foldByKey + RDD.foreach + RDD.foreachPartition + RDD.fullOuterJoin + RDD.getCheckpointFile + RDD.getNumPartitions + RDD.getResourceProfile + RDD.getStorageLevel + RDD.glom + RDD.groupBy + RDD.groupByKey + RDD.groupWith + RDD.histogram + RDD.id + RDD.intersection + RDD.isCheckpointed + RDD.isEmpty + RDD.isLocallyCheckpointed + RDD.join + RDD.keyBy + RDD.keys + RDD.leftOuterJoin + RDD.localCheckpoint + RDD.lookup + RDD.map + RDD.mapPartitions + RDD.mapPartitionsWithIndex + RDD.mapPartitionsWithSplit + RDD.mapValues + RDD.max + RDD.mean + RDD.meanApprox + RDD.min + RDD.name + RDD.partitionBy + RDD.persist + RDD.pipe + RDD.randomSplit + RDD.reduce + RDD.reduceByKey + RDD.reduceByKeyLocally + RDD.repartition + RDD.repartitionAndSortWithinPartitions + RDD.rightOuterJoin + RDD.sample + RDD.sampleByKey + RDD.sampleStdev + RDD.sampleVariance + RDD.saveAsHadoopDataset + RDD.saveAsHadoopFile + RDD.saveAsNewAPIHadoopDataset + RDD.saveAsNewAPIHadoopFile + RDD.saveAsPickleFile + RDD.saveAsSequenceFile + RDD.saveAsTextFile + RDD.setName + RDD.sortBy + RDD.sortByKey + RDD.stats + RDD.stdev + RDD.subtract + RDD.subtractByKey + RDD.sum + RDD.sumApprox + RDD.take + RDD.takeOrdered + RDD.takeSample + RDD.toDebugString + RDD.toLocalIterator + RDD.top + RDD.treeAggregate + RDD.treeReduce + RDD.union + RDD.unpersist + RDD.values + RDD.variance + RDD.withResources + RDD.zip + RDD.zipWithIndex + RDD.zipWithUniqueId + + +Broadcast and Accumulator +------------------------- + +.. currentmodule:: pyspark + +.. autosummary:: + :toctree: api/ + + Broadcast.destroy + Broadcast.dump + Broadcast.load + Broadcast.load_from_path + Broadcast.unpersist + Broadcast.value + Accumulator.add + Accumulator.value + AccumulatorParam.addInPlace + AccumulatorParam.zero + + +Management +---------- + +.. currentmodule:: pyspark + +.. autosummary:: + :toctree: api/ + + inheritable_thread_target + SparkConf.contains + SparkConf.get + SparkConf.getAll + SparkConf.set + SparkConf.setAll + SparkConf.setAppName + SparkConf.setExecutorEnv + SparkConf.setIfMissing + SparkConf.setMaster + SparkConf.setSparkHome + SparkConf.toDebugString + SparkFiles.get + SparkFiles.getRootDirectory + StorageLevel.DISK_ONLY + StorageLevel.DISK_ONLY_2 + StorageLevel.DISK_ONLY_3 + StorageLevel.MEMORY_AND_DISK + StorageLevel.MEMORY_AND_DISK_2 + StorageLevel.MEMORY_AND_DISK_DESER + StorageLevel.MEMORY_ONLY + StorageLevel.MEMORY_ONLY_2 + StorageLevel.OFF_HEAP + TaskContext.attemptNumber + TaskContext.cpus + TaskContext.get + TaskContext.getLocalProperty + TaskContext.partitionId + TaskContext.resources + TaskContext.stageId + TaskContext.taskAttemptId + RDDBarrier.mapPartitions + RDDBarrier.mapPartitionsWithIndex + BarrierTaskContext.allGather + BarrierTaskContext.attemptNumber + BarrierTaskContext.barrier + BarrierTaskContext.cpus + BarrierTaskContext.get + BarrierTaskContext.getLocalProperty + BarrierTaskContext.getTaskInfos + BarrierTaskContext.partitionId + BarrierTaskContext.resources + BarrierTaskContext.stageId + BarrierTaskContext.taskAttemptId + util.VersionUtils.majorMinorVersion + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/avro.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/avro.rst new file mode 100644 index 0000000..b6de88d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/avro.rst @@ -0,0 +1,28 @@ +.. 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. + + +==== +Avro +==== +.. currentmodule:: pyspark.sql.avro.functions + +.. autosummary:: + :toctree: api/ + + from_avro + to_avro diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/catalog.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/catalog.rst new file mode 100644 index 0000000..742af10 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/catalog.rst @@ -0,0 +1,54 @@ +.. 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. + + +======= +Catalog +======= + +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + Catalog.cacheTable + Catalog.clearCache + Catalog.createExternalTable + Catalog.createTable + Catalog.currentCatalog + Catalog.currentDatabase + Catalog.databaseExists + Catalog.dropGlobalTempView + Catalog.dropTempView + Catalog.functionExists + Catalog.getDatabase + Catalog.getFunction + Catalog.getTable + Catalog.isCached + Catalog.listCatalogs + Catalog.listColumns + Catalog.listDatabases + Catalog.listFunctions + Catalog.listTables + Catalog.recoverPartitions + Catalog.refreshByPath + Catalog.refreshTable + Catalog.registerFunction + Catalog.setCurrentCatalog + Catalog.setCurrentDatabase + Catalog.tableExists + Catalog.uncacheTable diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/column.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/column.rst new file mode 100644 index 0000000..b897b5c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/column.rst @@ -0,0 +1,60 @@ +.. 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. + + +====== +Column +====== +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + Column.__getattr__ + Column.__getitem__ + Column.alias + Column.asc + Column.asc_nulls_first + Column.asc_nulls_last + Column.astype + Column.between + Column.bitwiseAND + Column.bitwiseOR + Column.bitwiseXOR + Column.cast + Column.contains + Column.desc + Column.desc_nulls_first + Column.desc_nulls_last + Column.dropFields + Column.endswith + Column.eqNullSafe + Column.getField + Column.getItem + Column.ilike + Column.isNotNull + Column.isNull + Column.isin + Column.like + Column.name + Column.otherwise + Column.over + Column.rlike + Column.startswith + Column.substr + Column.when + Column.withField diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/configuration.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/configuration.rst new file mode 100644 index 0000000..7a5c104 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/configuration.rst @@ -0,0 +1,27 @@ +.. 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. + + +============= +Configuration +============= +.. currentmodule:: pyspark.sql.conf + +.. autosummary:: + :toctree: api/ + + RuntimeConfig diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/core_classes.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/core_classes.rst new file mode 100644 index 0000000..3cf1968 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/core_classes.rst @@ -0,0 +1,44 @@ +.. 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. + + +============ +Core Classes +============ +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + SparkSession + Catalog + DataFrame + Column + Observation + Row + GroupedData + PandasCogroupedOps + DataFrameNaFunctions + DataFrameStatFunctions + Window + DataFrameReader + DataFrameWriter + DataFrameWriterV2 + UDFRegistration + UDTFRegistration + udf.UserDefinedFunction + udtf.UserDefinedTableFunction diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/data_types.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/data_types.rst new file mode 100644 index 0000000..60c6b92 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/data_types.rst @@ -0,0 +1,50 @@ +.. 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. + + +========== +Data Types +========== +.. currentmodule:: pyspark.sql.types + +.. autosummary:: + :template: autosummary/class_with_docs.rst + :toctree: api/ + + ArrayType + BinaryType + BooleanType + ByteType + DataType + DateType + DecimalType + DoubleType + FloatType + IntegerType + LongType + MapType + NullType + ShortType + StringType + CharType + VarcharType + StructField + StructType + TimestampType + TimestampNTZType + DayTimeIntervalType + YearMonthIntervalType diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/dataframe.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/dataframe.rst new file mode 100644 index 0000000..1d6712b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/dataframe.rst @@ -0,0 +1,141 @@ +.. 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. + + +========= +DataFrame +========= + +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + DataFrame.__getattr__ + DataFrame.__getitem__ + DataFrame.agg + DataFrame.alias + DataFrame.approxQuantile + DataFrame.cache + DataFrame.checkpoint + DataFrame.coalesce + DataFrame.colRegex + DataFrame.collect + DataFrame.columns + DataFrame.corr + DataFrame.count + DataFrame.cov + DataFrame.createGlobalTempView + DataFrame.createOrReplaceGlobalTempView + DataFrame.createOrReplaceTempView + DataFrame.createTempView + DataFrame.crossJoin + DataFrame.crosstab + DataFrame.cube + DataFrame.describe + DataFrame.distinct + DataFrame.drop + DataFrame.dropDuplicates + DataFrame.dropDuplicatesWithinWatermark + DataFrame.drop_duplicates + DataFrame.dropna + DataFrame.dtypes + DataFrame.exceptAll + DataFrame.explain + DataFrame.fillna + DataFrame.filter + DataFrame.first + DataFrame.foreach + DataFrame.foreachPartition + DataFrame.freqItems + DataFrame.groupBy + DataFrame.head + DataFrame.hint + DataFrame.inputFiles + DataFrame.intersect + DataFrame.intersectAll + DataFrame.isEmpty + DataFrame.isLocal + DataFrame.isStreaming + DataFrame.join + DataFrame.limit + DataFrame.localCheckpoint + DataFrame.mapInPandas + DataFrame.mapInArrow + DataFrame.melt + DataFrame.na + DataFrame.observe + DataFrame.offset + DataFrame.orderBy + DataFrame.persist + DataFrame.printSchema + DataFrame.randomSplit + DataFrame.rdd + DataFrame.registerTempTable + DataFrame.repartition + DataFrame.repartitionByRange + DataFrame.replace + DataFrame.rollup + DataFrame.sameSemantics + DataFrame.sample + DataFrame.sampleBy + DataFrame.schema + DataFrame.select + DataFrame.selectExpr + DataFrame.semanticHash + DataFrame.show + DataFrame.sort + DataFrame.sortWithinPartitions + DataFrame.sparkSession + DataFrame.stat + DataFrame.storageLevel + DataFrame.subtract + DataFrame.summary + DataFrame.tail + DataFrame.take + DataFrame.to + DataFrame.toDF + DataFrame.toJSON + DataFrame.toLocalIterator + DataFrame.toPandas + DataFrame.to_pandas_on_spark + DataFrame.transform + DataFrame.union + DataFrame.unionAll + DataFrame.unionByName + DataFrame.unpersist + DataFrame.unpivot + DataFrame.where + DataFrame.withColumn + DataFrame.withColumns + DataFrame.withColumnRenamed + DataFrame.withColumnsRenamed + DataFrame.withMetadata + DataFrame.withWatermark + DataFrame.write + DataFrame.writeStream + DataFrame.writeTo + DataFrame.pandas_api + DataFrameNaFunctions.drop + DataFrameNaFunctions.fill + DataFrameNaFunctions.replace + DataFrameStatFunctions.approxQuantile + DataFrameStatFunctions.corr + DataFrameStatFunctions.cov + DataFrameStatFunctions.crosstab + DataFrameStatFunctions.freqItems + DataFrameStatFunctions.sampleBy diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/functions.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/functions.rst new file mode 100644 index 0000000..cd4f79e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/functions.rst @@ -0,0 +1,540 @@ +.. 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. + + +========= +Functions +========= +.. currentmodule:: pyspark.sql.functions + +A collections of builtin functions available for DataFrame operations. +From Apache Spark 3.5.0, all functions support Spark Connect. + +Normal Functions +---------------- +.. autosummary:: + :toctree: api/ + + col + column + lit + broadcast + coalesce + input_file_name + isnan + isnull + monotonically_increasing_id + named_struct + nanvl + rand + randn + spark_partition_id + when + bitwise_not + bitwiseNOT + expr + greatest + least + + +Math Functions +-------------- +.. autosummary:: + :toctree: api/ + + sqrt + abs + acos + acosh + asin + asinh + atan + atanh + atan2 + bin + cbrt + ceil + ceiling + conv + cos + cosh + cot + csc + e + exp + expm1 + factorial + floor + hex + unhex + hypot + ln + log + log10 + log1p + log2 + negate + negative + pi + pmod + positive + pow + power + rint + round + bround + sec + shiftleft + shiftright + shiftrightunsigned + sign + signum + sin + sinh + tan + tanh + toDegrees + try_add + try_avg + try_divide + try_multiply + try_subtract + try_sum + try_to_binary + try_to_number + degrees + toRadians + radians + width_bucket + + +Datetime Functions +------------------ +.. autosummary:: + :toctree: api/ + + add_months + convert_timezone + curdate + current_date + current_timestamp + current_timezone + date_add + date_diff + date_format + date_from_unix_date + date_sub + date_trunc + dateadd + datediff + day + date_part + datepart + dayofmonth + dayofweek + dayofyear + extract + second + weekofyear + year + quarter + month + last_day + localtimestamp + make_dt_interval + make_interval + make_timestamp + make_timestamp_ltz + make_timestamp_ntz + make_ym_interval + minute + months_between + next_day + hour + make_date + now + from_unixtime + unix_timestamp + to_unix_timestamp + to_timestamp + to_timestamp_ltz + to_timestamp_ntz + to_date + trunc + from_utc_timestamp + to_utc_timestamp + weekday + window + session_window + timestamp_micros + timestamp_millis + timestamp_seconds + try_to_timestamp + unix_date + unix_micros + unix_millis + unix_seconds + window_time + + +Collection Functions +-------------------- +.. autosummary:: + :toctree: api/ + + array + array_contains + arrays_overlap + array_join + create_map + slice + concat + array_position + element_at + array_append + array_size + array_sort + array_insert + array_remove + array_prepend + array_distinct + array_intersect + array_union + array_except + array_compact + transform + exists + forall + filter + aggregate + zip_with + transform_keys + transform_values + map_filter + map_from_arrays + map_zip_with + explode + explode_outer + posexplode + posexplode_outer + inline + inline_outer + get + get_json_object + json_tuple + from_json + schema_of_json + to_json + json_array_length + json_object_keys + size + cardinality + struct + sort_array + array_max + array_min + shuffle + reverse + flatten + sequence + array_repeat + map_contains_key + map_keys + map_values + map_entries + map_from_entries + arrays_zip + map_concat + from_csv + schema_of_csv + str_to_map + to_csv + try_element_at + + +Partition Transformation Functions +---------------------------------- +.. autosummary:: + :toctree: api/ + + years + months + days + hours + bucket + + +Aggregate Functions +------------------- +.. autosummary:: + :toctree: api/ + + any_value + approxCountDistinct + approx_count_distinct + approx_percentile + array_agg + avg + bit_and + bit_or + bit_xor + bool_and + bool_or + collect_list + collect_set + corr + count + count_distinct + countDistinct + count_min_sketch + count_if + covar_pop + covar_samp + every + first + first_value + grouping + grouping_id + histogram_numeric + hll_sketch_agg + hll_union_agg + kurtosis + last + last_value + max + max_by + mean + median + min + min_by + mode + percentile + percentile_approx + product + reduce + regr_avgx + regr_avgy + regr_count + regr_intercept + regr_r2 + regr_slope + regr_sxx + regr_sxy + regr_syy + skewness + some + std + stddev + stddev_pop + stddev_samp + sum + sum_distinct + sumDistinct + var_pop + var_samp + variance + + +Window Functions +---------------- +.. autosummary:: + :toctree: api/ + + cume_dist + dense_rank + lag + lead + nth_value + ntile + percent_rank + rank + row_number + + +Sort Functions +-------------- +.. autosummary:: + :toctree: api/ + + asc + asc_nulls_first + asc_nulls_last + desc + desc_nulls_first + desc_nulls_last + + +String Functions +---------------- +.. autosummary:: + :toctree: api/ + + ascii + base64 + bit_length + btrim + char + character_length + char_length + concat_ws + contains + decode + elt + encode + endswith + find_in_set + format_number + format_string + ilike + initcap + instr + lcase + length + like + lower + left + levenshtein + locate + lpad + ltrim + mask + octet_length + parse_url + position + printf + rlike + regexp + regexp_like + regexp_count + regexp_extract + regexp_extract_all + regexp_replace + regexp_substr + regexp_instr + replace + right + ucase + unbase64 + rpad + repeat + rtrim + soundex + split + split_part + startswith + substr + substring + substring_index + overlay + sentences + to_binary + to_char + to_number + to_varchar + translate + trim + upper + url_decode + url_encode + + +Bitwise Functions +----------------- +.. autosummary:: + :toctree: api/ + + bit_count + bit_get + getbit + + +Call Functions +-------------- +.. autosummary:: + :toctree: api/ + + call_function + call_udf + pandas_udf + udf + udtf + unwrap_udt + +Misc Functions +-------------- +.. autosummary:: + :toctree: api/ + + aes_decrypt + aes_encrypt + bitmap_bit_position + bitmap_bucket_number + bitmap_construct_agg + bitmap_count + bitmap_or_agg + current_catalog + current_database + current_schema + current_user + input_file_block_length + input_file_block_start + md5 + sha + sha1 + sha2 + crc32 + hash + xxhash64 + assert_true + raise_error + reflect + hll_sketch_estimate + hll_union + java_method + stack + try_aes_decrypt + typeof + user + version + +Predicate Functions +------------------- +.. autosummary:: + :toctree: api/ + + equal_null + ifnull + isnotnull + nullif + nvl + nvl2 + +Xml Functions +-------------- +.. autosummary:: + :toctree: api/ + + xpath + xpath_boolean + xpath_double + xpath_float + xpath_int + xpath_long + xpath_number + xpath_short + xpath_string + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/grouping.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/grouping.rst new file mode 100644 index 0000000..42de155 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/grouping.rst @@ -0,0 +1,39 @@ +.. 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. + + +======== +Grouping +======== +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + GroupedData.agg + GroupedData.apply + GroupedData.applyInPandas + GroupedData.applyInPandasWithState + GroupedData.avg + GroupedData.cogroup + GroupedData.count + GroupedData.max + GroupedData.mean + GroupedData.min + GroupedData.pivot + GroupedData.sum + PandasCogroupedOps.applyInPandas diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/index.rst new file mode 100644 index 0000000..233c8b2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/index.rst @@ -0,0 +1,44 @@ +.. 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. + + +========= +Spark SQL +========= + +This page gives an overview of all public Spark SQL API. + +.. toctree:: + :maxdepth: 2 + + core_classes + spark_session + configuration + io + dataframe + column + data_types + row + functions + window + grouping + catalog + avro + observation + udf + udtf + protobuf diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/io.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/io.rst new file mode 100644 index 0000000..e687ca0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/io.rst @@ -0,0 +1,65 @@ +.. 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. + + +============ +Input/Output +============ +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + DataFrameReader.csv + DataFrameReader.format + DataFrameReader.jdbc + DataFrameReader.json + DataFrameReader.load + DataFrameReader.option + DataFrameReader.options + DataFrameReader.orc + DataFrameReader.parquet + DataFrameReader.schema + DataFrameReader.table + DataFrameReader.text + DataFrameWriter.bucketBy + DataFrameWriter.csv + DataFrameWriter.format + DataFrameWriter.insertInto + DataFrameWriter.jdbc + DataFrameWriter.json + DataFrameWriter.mode + DataFrameWriter.option + DataFrameWriter.options + DataFrameWriter.orc + DataFrameWriter.parquet + DataFrameWriter.partitionBy + DataFrameWriter.save + DataFrameWriter.saveAsTable + DataFrameWriter.sortBy + DataFrameWriter.text + DataFrameWriterV2.using + DataFrameWriterV2.option + DataFrameWriterV2.options + DataFrameWriterV2.tableProperty + DataFrameWriterV2.partitionedBy + DataFrameWriterV2.create + DataFrameWriterV2.replace + DataFrameWriterV2.createOrReplace + DataFrameWriterV2.append + DataFrameWriterV2.overwrite + DataFrameWriterV2.overwritePartitions diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/observation.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/observation.rst new file mode 100644 index 0000000..52867ed --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/observation.rst @@ -0,0 +1,27 @@ +.. 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. + + +=========== +Observation +=========== +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + Observation.get diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/protobuf.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/protobuf.rst new file mode 100644 index 0000000..0ba3d56 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/protobuf.rst @@ -0,0 +1,28 @@ +.. 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. + + +======== +Protobuf +======== +.. currentmodule:: pyspark.sql.protobuf.functions + +.. autosummary:: + :toctree: api/ + + from_protobuf + to_protobuf diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/row.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/row.rst new file mode 100644 index 0000000..1234b8d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/row.rst @@ -0,0 +1,27 @@ +.. 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. + + +=== +Row +=== +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + Row.asDict diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/spark_session.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/spark_session.rst new file mode 100644 index 0000000..74315a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/spark_session.rst @@ -0,0 +1,73 @@ +.. 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. + + +============= +Spark Session +============= +.. currentmodule:: pyspark.sql + +The entry point to programming Spark with the Dataset and DataFrame API. +To create a Spark session, you should use ``SparkSession.builder`` attribute. +See also :class:`SparkSession`. + +.. autosummary:: + :toctree: api/ + + SparkSession.active + SparkSession.builder.appName + SparkSession.builder.config + SparkSession.builder.enableHiveSupport + SparkSession.builder.getOrCreate + SparkSession.builder.master + SparkSession.builder.remote + SparkSession.catalog + SparkSession.conf + SparkSession.createDataFrame + SparkSession.getActiveSession + SparkSession.newSession + SparkSession.range + SparkSession.read + SparkSession.readStream + SparkSession.sparkContext + SparkSession.sql + SparkSession.stop + SparkSession.streams + SparkSession.table + SparkSession.udf + SparkSession.udtf + SparkSession.version + + +Spark Connect Only +------------------ + +.. autosummary:: + :toctree: api/ + + SparkSession.builder.create + SparkSession.addArtifact + SparkSession.addArtifacts + SparkSession.copyFromLocalToFs + SparkSession.client + SparkSession.interruptAll + SparkSession.interruptTag + SparkSession.interruptOperation + SparkSession.addTag + SparkSession.removeTag + SparkSession.getTags + SparkSession.clearTags diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udf.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udf.rst new file mode 100644 index 0000000..a3d27c4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udf.rst @@ -0,0 +1,32 @@ +.. 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. + + +=== +UDF +=== + +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + udf.UserDefinedFunction.asNondeterministic + udf.UserDefinedFunction.returnType + UDFRegistration.register + UDFRegistration.registerJavaFunction + UDFRegistration.registerJavaUDAF diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udtf.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udtf.rst new file mode 100644 index 0000000..64400a2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/udtf.rst @@ -0,0 +1,30 @@ +.. 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. + + +==== +UDTF +==== + +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + udtf.UserDefinedTableFunction.asDeterministic + udtf.UserDefinedTableFunction.returnType + UDTFRegistration.register diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/window.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/window.rst new file mode 100644 index 0000000..3625164 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.sql/window.rst @@ -0,0 +1,38 @@ +.. 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. + + +====== +Window +====== + +.. currentmodule:: pyspark.sql + +.. autosummary:: + :toctree: api/ + + Window.currentRow + Window.orderBy + Window.partitionBy + Window.rangeBetween + Window.rowsBetween + Window.unboundedFollowing + Window.unboundedPreceding + WindowSpec.orderBy + WindowSpec.partitionBy + WindowSpec.rangeBetween + WindowSpec.rowsBetween diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/core_classes.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/core_classes.rst new file mode 100644 index 0000000..10c2211 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/core_classes.rst @@ -0,0 +1,32 @@ +.. 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. + + +============ +Core Classes +============ + +.. currentmodule:: pyspark.sql.streaming + +.. autosummary:: + :toctree: api/ + + DataStreamReader + DataStreamWriter + StreamingQuery + StreamingQueryManager + StreamingQueryListener diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/index.rst new file mode 100644 index 0000000..2cb0b12 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/index.rst @@ -0,0 +1,30 @@ +.. 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. + + +==================== +Structured Streaming +==================== + +This page gives an overview of all public Structed Streaming API. + +.. toctree:: + :maxdepth: 2 + + core_classes + io + query_management diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/io.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/io.rst new file mode 100644 index 0000000..7a20777 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/io.rst @@ -0,0 +1,49 @@ +.. 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. + + +============ +Input/Output +============ + +.. currentmodule:: pyspark.sql.streaming + +.. autosummary:: + :toctree: api/ + + DataStreamReader.csv + DataStreamReader.format + DataStreamReader.json + DataStreamReader.load + DataStreamReader.option + DataStreamReader.options + DataStreamReader.orc + DataStreamReader.parquet + DataStreamReader.schema + DataStreamReader.table + DataStreamReader.text + DataStreamWriter.foreach + DataStreamWriter.foreachBatch + DataStreamWriter.format + DataStreamWriter.option + DataStreamWriter.options + DataStreamWriter.outputMode + DataStreamWriter.partitionBy + DataStreamWriter.queryName + DataStreamWriter.start + DataStreamWriter.toTable + DataStreamWriter.trigger diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/query_management.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/query_management.rst new file mode 100644 index 0000000..2a8a84d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.ss/query_management.rst @@ -0,0 +1,45 @@ +.. 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. + + +================ +Query Management +================ + +.. currentmodule:: pyspark.sql.streaming + +.. autosummary:: + :toctree: api/ + + StreamingQuery.awaitTermination + StreamingQuery.exception + StreamingQuery.explain + StreamingQuery.id + StreamingQuery.isActive + StreamingQuery.lastProgress + StreamingQuery.name + StreamingQuery.processAllAvailable + StreamingQuery.recentProgress + StreamingQuery.runId + StreamingQuery.status + StreamingQuery.stop + StreamingQueryManager.active + StreamingQueryManager.addListener + StreamingQueryManager.awaitAnyTermination + StreamingQueryManager.get + StreamingQueryManager.removeListener + StreamingQueryManager.resetTerminated diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.streaming.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.streaming.rst new file mode 100644 index 0000000..b8bf7e3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.streaming.rst @@ -0,0 +1,130 @@ +.. 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. + + +======================== +Spark Streaming (Legacy) +======================== + +Core Classes +------------ + +.. currentmodule:: pyspark.streaming + +.. autosummary:: + :toctree: api/ + + StreamingContext + DStream + + +Streaming Management +-------------------- + +.. currentmodule:: pyspark.streaming + +.. autosummary:: + :toctree: api/ + + StreamingContext.addStreamingListener + StreamingContext.awaitTermination + StreamingContext.awaitTerminationOrTimeout + StreamingContext.checkpoint + StreamingContext.getActive + StreamingContext.getActiveOrCreate + StreamingContext.getOrCreate + StreamingContext.remember + StreamingContext.sparkContext + StreamingContext.start + StreamingContext.stop + StreamingContext.transform + StreamingContext.union + + +Input and Output +---------------- + +.. autosummary:: + :toctree: api/ + + StreamingContext.binaryRecordsStream + StreamingContext.queueStream + StreamingContext.socketTextStream + StreamingContext.textFileStream + DStream.pprint + DStream.saveAsTextFiles + + +Transformations and Actions +--------------------------- + +.. currentmodule:: pyspark.streaming + +.. autosummary:: + :toctree: api/ + + DStream.cache + DStream.checkpoint + DStream.cogroup + DStream.combineByKey + DStream.context + DStream.count + DStream.countByValue + DStream.countByValueAndWindow + DStream.countByWindow + DStream.filter + DStream.flatMap + DStream.flatMapValues + DStream.foreachRDD + DStream.fullOuterJoin + DStream.glom + DStream.groupByKey + DStream.groupByKeyAndWindow + DStream.join + DStream.leftOuterJoin + DStream.map + DStream.mapPartitions + DStream.mapPartitionsWithIndex + DStream.mapValues + DStream.partitionBy + DStream.persist + DStream.reduce + DStream.reduceByKey + DStream.reduceByKeyAndWindow + DStream.reduceByWindow + DStream.repartition + DStream.rightOuterJoin + DStream.slice + DStream.transform + DStream.transformWith + DStream.union + DStream.updateStateByKey + DStream.window + + +Kinesis +------- + +.. currentmodule:: pyspark.streaming.kinesis + +.. autosummary:: + :toctree: api/ + + KinesisUtils.createStream + InitialPositionInStream.LATEST + InitialPositionInStream.TRIM_HORIZON + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.testing.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.testing.rst new file mode 100644 index 0000000..96b0c72 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/reference/pyspark.testing.rst @@ -0,0 +1,30 @@ +.. 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. + + +======= +Testing +======= + +.. currentmodule:: pyspark.testing + +.. autosummary:: + :toctree: api/ + + assertDataFrameEqual + assertPandasOnSparkEqual + assertSchemaEqual diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/arrow_pandas.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/arrow_pandas.rst new file mode 100644 index 0000000..60c11b7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/arrow_pandas.rst @@ -0,0 +1,24 @@ +.. 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. + +:orphan: + +======================= +Apache Arrow in PySpark +======================= + +This page has been moved to `Apache Arrow in PySpark `_. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/index.rst new file mode 100644 index 0000000..67f8c8d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/index.rst @@ -0,0 +1,37 @@ +.. 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. + + +=========== +User Guides +=========== + +PySpark specific user guides are available here: + +.. toctree:: + :maxdepth: 2 + + python_packaging + sql/index + pandas_on_spark/index + +There are also basic programming guides covering multiple languages available in +`the Spark documentation `_, including these: + +- `Spark SQL, DataFrames and Datasets Guide `_ +- `Structured Streaming Programming Guide `_ +- `Machine Learning Library (MLlib) Guide `_ diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/best_practices.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/best_practices.rst new file mode 100644 index 0000000..14c04aa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/best_practices.rst @@ -0,0 +1,331 @@ +.. 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. + + +============== +Best Practices +============== + +Leverage PySpark APIs +--------------------- + +Pandas API on Spark uses Spark under the hood; therefore, many features and performance optimizations are available +in pandas API on Spark as well. Leverage and combine those cutting-edge features with pandas API on Spark. + +Existing Spark context and Spark sessions are used out of the box in pandas API on Spark. If you already have your own +configured Spark context or sessions running, pandas API on Spark uses them. + +If there is no Spark context or session running in your environment (e.g., ordinary Python interpreter), +such configurations can be set to ``SparkContext`` and/or ``SparkSession``. +Once a Spark context and/or session is created, pandas API on Spark can use this context and/or session automatically. +For example, if you want to configure the executor memory in Spark, you can do as below: + +.. code-block:: python + + from pyspark import SparkConf, SparkContext + conf = SparkConf() + conf.set('spark.executor.memory', '2g') + # Pandas API on Spark automatically uses this Spark context with the configurations set. + SparkContext(conf=conf) + + import pyspark.pandas as ps + ... + +Another common configuration might be Arrow optimization in PySpark. In case of SQL configuration, +it can be set into Spark session as below: + +.. code-block:: python + + from pyspark.sql import SparkSession + builder = SparkSession.builder.appName("pandas-on-spark") + builder = builder.config("spark.sql.execution.arrow.pyspark.enabled", "true") + # Pandas API on Spark automatically uses this Spark session with the configurations set. + builder.getOrCreate() + + import pyspark.pandas as ps + ... + +All Spark features such as history server, web UI and deployment modes can be used as are with pandas API on Spark. +If you are interested in performance tuning, please see also `Tuning Spark `_. + + +Check execution plans +--------------------- + +Expensive operations can be predicted by leveraging PySpark API `DataFrame.spark.explain()` +before the actual computation since pandas API on Spark is based on lazy execution. For example, see below. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'id': range(10)}) + >>> psdf = psdf[psdf.id > 5] + >>> psdf.spark.explain() + == Physical Plan == + *(1) Filter (id#1L > 5) + +- *(1) Scan ExistingRDD[__index_level_0__#0L,id#1L] + + +Whenever you are not sure about such cases, you can check the actual execution plans and +foresee the expensive cases. + +Even though pandas API on Spark tries its best to optimize and reduce such shuffle operations by leveraging Spark +optimizers, it is best to avoid shuffling in the application side whenever possible. + + +Use checkpoint +-------------- + +After a bunch of operations on pandas API on Spark objects, the underlying Spark planner can slow down due to the huge and complex plan. +If the Spark plan becomes huge or it takes the planning long time, ``DataFrame.spark.checkpoint()`` +or ``DataFrame.spark.local_checkpoint()`` would be helpful. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'id': range(10)}) + >>> psdf = psdf[psdf.id > 5] + >>> psdf['id'] = psdf['id'] + (10 * psdf['id'] + psdf['id']) + >>> psdf = psdf.groupby('id').head(2) + >>> psdf.spark.explain() + == Physical Plan == + *(3) Project [__index_level_0__#0L, id#31L] + +- *(3) Filter (isnotnull(__row_number__#44) AND (__row_number__#44 <= 2)) + +- Window [row_number() windowspecdefinition(__groupkey_0__#36L, __natural_order__#16L ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS __row_number__#44], [__groupkey_0__#36L], [__natural_order__#16L ASC NULLS FIRST] + +- *(2) Sort [__groupkey_0__#36L ASC NULLS FIRST, __natural_order__#16L ASC NULLS FIRST], false, 0 + +- Exchange hashpartitioning(__groupkey_0__#36L, 200), true, [id=#33] + +- *(1) Project [__index_level_0__#0L, (id#1L + ((id#1L * 10) + id#1L)) AS __groupkey_0__#36L, (id#1L + ((id#1L * 10) + id#1L)) AS id#31L, __natural_order__#16L] + +- *(1) Project [__index_level_0__#0L, id#1L, monotonically_increasing_id() AS __natural_order__#16L] + +- *(1) Filter (id#1L > 5) + +- *(1) Scan ExistingRDD[__index_level_0__#0L,id#1L] + + >>> psdf = psdf.spark.local_checkpoint() # or psdf.spark.checkpoint() + >>> psdf.spark.explain() + == Physical Plan == + *(1) Project [__index_level_0__#0L, id#31L] + +- *(1) Scan ExistingRDD[__index_level_0__#0L,id#31L,__natural_order__#59L] + +As you can see, the previous Spark plan is dropped and starts with a simple plan. +The result of the previous DataFrame is stored in the configured file system when calling ``DataFrame.spark.checkpoint()``, +or in the executor when calling ``DataFrame.spark.local_checkpoint()``. + + +Avoid shuffling +--------------- + +Some operations such as ``sort_values`` are more difficult to do in a parallel or distributed +environment than in in-memory on a single machine because it needs to send data to other nodes, +and exchange the data across multiple nodes via networks. See the example below. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'id': range(10)}).sort_values(by="id") + >>> psdf.spark.explain() + == Physical Plan == + *(2) Sort [id#9L ASC NULLS LAST], true, 0 + +- Exchange rangepartitioning(id#9L ASC NULLS LAST, 200), true, [id=#18] + +- *(1) Scan ExistingRDD[__index_level_0__#8L,id#9L] + +As you can see, it requires ``Exchange`` which requires a shuffle and it is likely expensive. + + +Avoid computation on single partition +------------------------------------- + +Another common case is the computation on a single partition. Currently, some APIs such as +`DataFrame.rank `_ +use PySpark’s Window without specifying partition specification. This moves all data into a single +partition in a single machine and could cause serious performance degradation. +Such APIs should be avoided for very large datasets. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'id': range(10)}) + >>> psdf.rank().spark.explain() + == Physical Plan == + *(4) Project [__index_level_0__#16L, id#24] + +- Window [avg(cast(_w0#26 as bigint)) windowspecdefinition(id#17L, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS id#24], [id#17L] + +- *(3) Project [__index_level_0__#16L, _w0#26, id#17L] + +- Window [row_number() windowspecdefinition(id#17L ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS _w0#26], [id#17L ASC NULLS FIRST] + +- *(2) Sort [id#17L ASC NULLS FIRST], false, 0 + +- Exchange SinglePartition, true, [id=#48] + +- *(1) Scan ExistingRDD[__index_level_0__#16L,id#17L] + +Instead, use +`GroupBy.rank `_ +as it is less expensive because data can be distributed and computed for each group. + + +Avoid reserved column names +--------------------------- + +Columns with leading ``__`` and trailing ``__`` are reserved in pandas API on Spark. To handle internal behaviors for, such as, index, +pandas API on Spark uses some internal columns. Therefore, it is discouraged to use such column names and they are not guaranteed to work. + + +Do not use duplicated column names +---------------------------------- + +It is disallowed to use duplicated column names because Spark SQL does not allow this in general. Pandas API on Spark inherits +this behavior. For instance, see below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'a': [1, 2], 'b':[3, 4]}) + >>> psdf.columns = ["a", "a"] + ... + Reference 'a' is ambiguous, could be: a, a.; + +Additionally, it is strongly discouraged to use case sensitive column names. Pandas API on Spark disallows it by default. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'a': [1, 2], 'A':[3, 4]}) + ... + Reference 'a' is ambiguous, could be: a, a.; + +However, you can turn on ``spark.sql.caseSensitive`` in Spark configuration to enable it for use at your own risk. + +.. code-block:: python + + >>> from pyspark.sql import SparkSession + >>> builder = SparkSession.builder.appName("pandas-on-spark") + >>> builder = builder.config("spark.sql.caseSensitive", "true") + >>> builder.getOrCreate() + + >>> import pyspark.pandas as ps + >>> psdf = ps.DataFrame({'a': [1, 2], 'A':[3, 4]}) + >>> psdf + a A + 0 1 3 + 1 2 4 + + +Specify the index column in conversion from Spark DataFrame to pandas-on-Spark DataFrame +---------------------------------------------------------------------------------------- + +When pandas-on-Spark Dataframe is converted from Spark DataFrame, it loses the index information, which results in using +the default index in pandas API on Spark DataFrame. The default index is inefficient in general comparing to explicitly specifying +the index column. Specify the index column whenever possible. + +See `working with PySpark `_ + + +Use ``distributed`` or ``distributed-sequence`` default index +------------------------------------------------------------- + +One common issue that pandas-on-Spark users face is the slow performance due to the default index. Pandas API on Spark attaches +a default index when the index is unknown, for example, Spark DataFrame is directly converted to pandas-on-Spark DataFrame. + +Note that ``sequence`` requires the computation on a single partition which is discouraged. If you plan +to handle large data in production, make it distributed by configuring the default index to ``distributed`` or +``distributed-sequence`` . + +See `Default Index Type `_ for more details about configuring default index. + + +Reduce the operations on different DataFrame/Series +--------------------------------------------------- + +Pandas API on Spark disallows the operations on different DataFrames (or Series) by default to prevent expensive operations. +It internally performs a join operation which can be expensive in general, which is discouraged. Whenever possible, +this operation should be avoided. + +See `Operations on different DataFrames `_ for more details. + + +Use pandas API on Spark directly whenever possible +--------------------------------------------------- + +Although pandas API on Spark has most of the pandas-equivalent APIs, there are several APIs not implemented yet or explicitly unsupported. + +As an example, pandas API on Spark does not implement ``__iter__()`` to prevent users from collecting all data into the client (driver) side from the whole cluster. +Unfortunately, many external APIs such as Python built-in functions such as min, max, sum, etc. require the given argument to be iterable. +In case of pandas, it works properly out of the box as below: + +.. code-block:: python + + >>> import pandas as pd + >>> max(pd.Series([1, 2, 3])) + 3 + >>> min(pd.Series([1, 2, 3])) + 1 + >>> sum(pd.Series([1, 2, 3])) + 6 + +pandas dataset lives in the single machine, and is naturally iterable locally within the same machine. +However, pandas-on-Spark dataset lives across multiple machines, and they are computed in a distributed manner. +It is difficult to be locally iterable and it is very likely users collect the entire data into the client side without knowing it. +Therefore, it is best to stick to using pandas-on-Spark APIs. +The examples above can be converted as below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.Series([1, 2, 3]).max() + 3 + >>> ps.Series([1, 2, 3]).min() + 1 + >>> ps.Series([1, 2, 3]).sum() + 6 + +Another common pattern from pandas users might be to rely on list comprehension or generator expression. +However, it also assumes the dataset is locally iterable under the hood. +Therefore, it works seamlessly in pandas as below: + +.. code-block:: python + + >>> import pandas as pd + >>> data = [] + >>> countries = ['London', 'New York', 'Helsinki'] + >>> pser = pd.Series([20., 21., 12.], index=countries) + >>> for temperature in pser: + ... assert temperature > 0 + ... if temperature > 1000: + ... temperature = None + ... data.append(temperature ** 2) + ... + >>> pd.Series(data, index=countries) + London 400.0 + New York 441.0 + Helsinki 144.0 + dtype: float64 + +However, for pandas API on Spark it does not work for the same reason above. +The example above can be also changed to directly using pandas-on-Spark APIs as below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> import numpy as np + >>> countries = ['London', 'New York', 'Helsinki'] + >>> psser = ps.Series([20., 21., 12.], index=countries) + >>> def square(temperature) -> np.float64: + ... assert temperature > 0 + ... if temperature > 1000: + ... temperature = None + ... return temperature ** 2 + ... + >>> psser.apply(square) + London 400.0 + New York 441.0 + Helsinki 144.0 + dtype: float64 diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/faq.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/faq.rst new file mode 100644 index 0000000..6f38a76 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/faq.rst @@ -0,0 +1,57 @@ +.. 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. + + +=== +FAQ +=== + +Should I use PySpark's DataFrame API or pandas API on Spark? +------------------------------------------------------------ + +If you are already familiar with pandas and want to leverage Spark for big data, we recommend using pandas API on Spark. +If you are learning Spark from the ground up, we recommend you start with PySpark's API. + +Does pandas API on Spark support Structured Streaming? +------------------------------------------------------ + +No, pandas API on Spark does not support Structured Streaming officially. + +As a workaround, you can use pandas-on-Spark APIs with `foreachBatch` in Structured Streaming which allows batch APIs: + +.. code-block:: python + + >>> def func(batch_df, batch_id): + ... pandas_on_spark_df = ps.DataFrame(batch_df) + ... pandas_on_spark_df['a'] = 1 + ... print(pandas_on_spark_df) + + >>> spark.readStream.format("rate").load().writeStream.foreachBatch(func).start() + timestamp value a + 0 2020-02-21 09:49:37.574 4 1 + timestamp value a + 0 2020-02-21 09:49:38.574 5 1 + ... + +How is pandas API on Spark different from Dask? +----------------------------------------------- + +Different projects have different focuses. Spark is already deployed in virtually every +organization, and often is the primary interface to the massive amount of data stored in data lakes. +pandas API on Spark was inspired by Dask, and aims to make the transition from pandas to Spark easy for data +scientists. + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/from_to_dbms.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/from_to_dbms.rst new file mode 100644 index 0000000..b33a88f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/from_to_dbms.rst @@ -0,0 +1,125 @@ +.. 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/to other DBMSes +==================== +.. currentmodule:: pyspark.pandas + + +The APIs interacting with other DBMSes in pandas API on Spark are slightly different from the ones in pandas +because pandas API on Spark leverages JDBC APIs in PySpark to read and write from/to other DBMSes. + +The APIs to read/write from/to external DBMSes are as follows: + +.. autosummary:: + + read_sql_table + read_sql_query + read_sql + +.. + TODO: we should implement and document `DataFrame.to_sql`. + +pandas-on-Spark needs a canonical JDBC URL for ``con``, and is able to take extra keyword arguments for `the options in PySpark JDBC APIs `_: + +.. code-block:: python + + ps.read_sql(..., dbtable="...", driver="", keytab="", ...) + + +Reading and writing DataFrames +------------------------------ + +In the example below, you will read and write a table in SQLite. + +Firstly, create the ``example`` database as below via Python's SQLite library. This will be read to pandas-on-Spark later: + +.. code-block:: python + + import sqlite3 + + con = sqlite3.connect('example.db') + cur = con.cursor() + # Create table + cur.execute( + '''CREATE TABLE stocks + (date text, trans text, symbol text, qty real, price real)''') + # Insert a row of data + cur.execute("INSERT INTO stocks VALUES ('2006-01-05','BUY','RHAT',100,35.14)") + # Save (commit) the changes + con.commit() + con.close() + +Pandas API on Spark requires a JDBC driver to read so it requires the driver for your particular database to be on the Spark's classpath. For SQLite JDBC driver, you can download it, for example, as below: + +.. code-block:: bash + + curl -O https://repo1.maven.org/maven2/org/xerial/sqlite-jdbc/3.34.0/sqlite-jdbc-3.34.0.jar + +After that, you should add it into your Spark session first. Once you add it, pandas API on Spark will automatically detect the Spark session and leverage it. + +.. code-block:: python + + import os + + from pyspark.sql import SparkSession + + (SparkSession.builder + .master("local") + .appName("SQLite JDBC") + .config( + "spark.jars", + "{}/sqlite-jdbc-3.34.0.jar".format(os.getcwd())) + .config( + "spark.driver.extraClassPath", + "{}/sqlite-jdbc-3.34.0.jar".format(os.getcwd())) + .getOrCreate()) + +Now, you're ready to read the table: + +.. code-block:: python + + import pyspark.pandas as ps + + df = ps.read_sql("stocks", con="jdbc:sqlite:{}/example.db".format(os.getcwd())) + df + +.. code-block:: text + + date trans symbol qty price + 0 2006-01-05 BUY RHAT 100.0 35.14 + +You can also write it back to the ``stocks`` table as below: + +.. + TODO: switch to use DataFrame.to_sql in the example + +.. code-block:: python + + df.price += 1 + df.spark.to_spark_io( + format="jdbc", mode="append", + dbtable="stocks", url="jdbc:sqlite:{}/example.db".format(os.getcwd())) + ps.read_sql("stocks", con="jdbc:sqlite:{}/example.db".format(os.getcwd())) + +.. code-block:: text + + date trans symbol qty price + 0 2006-01-05 BUY RHAT 100.0 35.14 + 1 2006-01-05 BUY RHAT 100.0 36.14 diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/index.rst new file mode 100644 index 0000000..daf69e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/index.rst @@ -0,0 +1,34 @@ +.. 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. + + +=================== +Pandas API on Spark +=================== + +.. toctree:: + :maxdepth: 2 + + options + pandas_pyspark + transform_apply + types + typehints + from_to_dbms + best_practices + supported_pandas_api + faq diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/options.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/options.rst new file mode 100644 index 0000000..92b572b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/options.rst @@ -0,0 +1,316 @@ +.. 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. + + +==================== +Options and settings +==================== +.. currentmodule:: pyspark.pandas + +Pandas API on Spark has an options system that lets you customize some aspects of its behaviour, +display-related options being those the user is most likely to adjust. + +Options have a full "dotted-style", case-insensitive name (e.g. ``display.max_rows``). +You can get/set options directly as attributes of the top-level ``options`` attribute: + + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.options.display.max_rows + 1000 + >>> ps.options.display.max_rows = 10 + >>> ps.options.display.max_rows + 10 + +The API is composed of 3 relevant functions, available directly from the ``pandas_on_spark`` +namespace: + +* :func:`get_option` / :func:`set_option` - get/set the value of a single option. +* :func:`reset_option` - reset one or more options to their default value. + +**Note:** Developers can check out `pyspark.pandas/config.py `_ for more information. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.get_option("display.max_rows") + 1000 + >>> ps.set_option("display.max_rows", 101) + >>> ps.get_option("display.max_rows") + 101 + + +Getting and setting options +--------------------------- + +As described above, :func:`get_option` and :func:`set_option` +are available from the ``pandas_on_spark`` namespace. To change an option, call +``set_option('option name', new_value)``. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.get_option('compute.max_rows') + 1000 + >>> ps.set_option('compute.max_rows', 2000) + >>> ps.get_option('compute.max_rows') + 2000 + +All options also have a default value, and you can use ``reset_option`` to do just that: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.reset_option("display.max_rows") + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.get_option("display.max_rows") + 1000 + >>> ps.set_option("display.max_rows", 999) + >>> ps.get_option("display.max_rows") + 999 + >>> ps.reset_option("display.max_rows") + >>> ps.get_option("display.max_rows") + 1000 + +``option_context`` context manager has been exposed through +the top-level API, allowing you to execute code with given option values. Option values +are restored automatically when you exit the `with` block: + +.. code-block:: python + + >>> with ps.option_context("display.max_rows", 10, "compute.max_rows", 5): + ... print(ps.get_option("display.max_rows")) + ... print(ps.get_option("compute.max_rows")) + 10 + 5 + >>> print(ps.get_option("display.max_rows")) + >>> print(ps.get_option("compute.max_rows")) + 1000 + 1000 + + +Operations on different DataFrames +---------------------------------- + +Pandas API on Spark disallows the operations on different DataFrames (or Series) by default to prevent expensive +operations. It internally performs a join operation which can be expensive in general. + +This can be enabled by setting `compute.ops_on_diff_frames` to `True` to allow such cases. +See the examples below. + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.set_option('compute.ops_on_diff_frames', True) + >>> psdf1 = ps.range(5) + >>> psdf2 = ps.DataFrame({'id': [5, 4, 3]}) + >>> (psdf1 - psdf2).sort_index() + id + 0 -5.0 + 1 -3.0 + 2 -1.0 + 3 NaN + 4 NaN + >>> ps.reset_option('compute.ops_on_diff_frames') + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.set_option('compute.ops_on_diff_frames', True) + >>> psdf = ps.range(5) + >>> psser_a = ps.Series([1, 2, 3, 4]) + >>> # 'psser_a' is not from 'psdf' DataFrame. So it is considered as a Series not from 'psdf'. + >>> psdf['new_col'] = psser_a + >>> psdf + id new_col + 0 0 1.0 + 1 1 2.0 + 3 3 4.0 + 2 2 3.0 + 4 4 NaN + >>> ps.reset_option('compute.ops_on_diff_frames') + + +Default Index type +------------------ + +In the pandas API on Spark, the default index is used in several cases, for instance, +when Spark DataFrame is converted into pandas-on-Spark DataFrame. In this case, internally pandas API on Spark attaches a +default index into pandas-on-Spark DataFrame. + +There are several types of the default index that can be configured by `compute.default_index_type` as below: + +**sequence**: It implements a sequence that increases one by one, by PySpark's Window function without +specifying a partition. Therefore, it can end up with a whole partition in a single node. +This index type should be avoided when the data is large. See the example below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.set_option('compute.default_index_type', 'sequence') + >>> psdf = ps.range(3) + >>> ps.reset_option('compute.default_index_type') + >>> psdf.index + Int64Index([0, 1, 2], dtype='int64') + +This is conceptually equivalent to the PySpark example as below: + +.. code-block:: python + + >>> from pyspark.sql import functions as sf, Window + >>> import pyspark.pandas as ps + >>> spark_df = ps.range(3).to_spark() + >>> sequential_index = sf.row_number().over( + ... Window.orderBy(sf.monotonically_increasing_id().asc())) - 1 + >>> spark_df.select(sequential_index).rdd.map(lambda r: r[0]).collect() + [0, 1, 2] + +**distributed-sequence** (default): It implements a sequence that increases one by one, by group-by and +group-map approach in a distributed manner. It still generates the sequential index globally. +If the default index must be the sequence in a large dataset, this +index has to be used. See the example below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.set_option('compute.default_index_type', 'distributed-sequence') + >>> psdf = ps.range(3) + >>> ps.reset_option('compute.default_index_type') + >>> psdf.index + Int64Index([0, 1, 2], dtype='int64') + +This is conceptually equivalent to the PySpark example as below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> spark_df = ps.range(3).to_spark() + >>> spark_df.rdd.zipWithIndex().map(lambda p: p[1]).collect() + [0, 1, 2] + +**distributed**: It implements a monotonically increasing sequence simply by using +PySpark's `monotonically_increasing_id` function in a fully distributed manner. The +values are indeterministic. If the index does not have to be a sequence that increases +one by one, this index should be used. Performance-wise, this index almost does not +have any penalty compared to other index types. See the example below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> ps.set_option('compute.default_index_type', 'distributed') + >>> psdf = ps.range(3) + >>> ps.reset_option('compute.default_index_type') + >>> psdf.index + Int64Index([25769803776, 60129542144, 94489280512], dtype='int64') + +This is conceptually equivalent to the PySpark example as below: + +.. code-block:: python + + >>> from pyspark.sql import functions as sf + >>> import pyspark.pandas as ps + >>> spark_df = ps.range(3).to_spark() + >>> spark_df.select(sf.monotonically_increasing_id()) \ + ... .rdd.map(lambda r: r[0]).collect() + [25769803776, 60129542144, 94489280512] + +.. warning:: + It is very unlikely for this type of index to be used for computing two + different dataframes because it is not guaranteed to have the same indexes in two dataframes. + If you use this default index and turn on `compute.ops_on_diff_frames`, the result + from the operations between two different DataFrames will likely be an unexpected + output due to the indeterministic index values. + + +Available options +----------------- + +=============================== ======================= ===================================================== +Option Default Description +=============================== ======================= ===================================================== +display.max_rows 1000 This sets the maximum number of rows pandas-on-Spark + should output when printing out various output. For + example, this value determines the number of rows to + be shown at the repr() in a dataframe. Set `None` to + unlimit the input length. Default is 1000. +compute.max_rows 1000 'compute.max_rows' sets the limit of the current + pandas-on-Spark DataFrame. Set `None` to unlimit the + input length. When the limit is set, it is executed + by the shortcut by collecting the data into the + driver, and then using the pandas API. If the limit + is unset, the operation is executed by PySpark. + Default is 1000. +compute.shortcut_limit 1000 'compute.shortcut_limit' sets the limit for a + shortcut. It computes specified number of rows and + use its schema. When the dataframe length is larger + than this limit, pandas-on-Spark uses PySpark to + compute. +compute.ops_on_diff_frames False This determines whether or not to operate between two + different dataframes. For example, 'combine_frames' + function internally performs a join operation which + can be expensive in general. So, if + `compute.ops_on_diff_frames` variable is not True, + that method throws an exception. +compute.default_index_type 'distributed-sequence' This sets the default index type: sequence, + distributed and distributed-sequence. +compute.default_index_cache 'MEMORY_AND_DISK_SER' This sets the default storage level for temporary + RDDs cached in distributed-sequence indexing: 'NONE', + 'DISK_ONLY', 'DISK_ONLY_2', 'DISK_ONLY_3', + 'MEMORY_ONLY', 'MEMORY_ONLY_2', 'MEMORY_ONLY_SER', + 'MEMORY_ONLY_SER_2', 'MEMORY_AND_DISK', + 'MEMORY_AND_DISK_2', 'MEMORY_AND_DISK_SER', + 'MEMORY_AND_DISK_SER_2', 'OFF_HEAP', + 'LOCAL_CHECKPOINT'. +compute.ordered_head False 'compute.ordered_head' sets whether or not to operate + head with natural ordering. pandas-on-Spark does not + guarantee the row ordering so `head` could return + some rows from distributed partitions. If + 'compute.ordered_head' is set to True, pandas-on- + Spark performs natural ordering beforehand, but it + will cause a performance overhead. +compute.eager_check True 'compute.eager_check' sets whether or not to launch + some Spark jobs just for the sake of validation. If + 'compute.eager_check' is set to True, pandas-on-Spark + performs the validation beforehand, but it will cause + a performance overhead. Otherwise, pandas-on-Spark + skip the validation and will be slightly different + from pandas. Affected APIs: `Series.dot`, + `Series.asof`, `Series.compare`, + `FractionalExtensionOps.astype`, + `IntegralExtensionOps.astype`, + `FractionalOps.astype`, `DecimalOps.astype`, `skipna + of statistical functions`. +compute.isin_limit 80 'compute.isin_limit' sets the limit for filtering by + 'Column.isin(list)'. If the length of the ‘list’ is + above the limit, broadcast join is used instead for + better performance. +plotting.max_rows 1000 'plotting.max_rows' sets the visual limit on top-n- + based plots such as `plot.bar` and `plot.pie`. If it + is set to 1000, the first 1000 data points will be + used for plotting. Default is 1000. +plotting.sample_ratio None 'plotting.sample_ratio' sets the proportion of data + that will be plotted for sample-based plots such as + `plot.line` and `plot.area`. This option defaults to + 'plotting.max_rows' option. +plotting.backend 'plotly' Backend to use for plotting. Default is plotly. + Supports any package that has a top-level `.plot` + method. Known options are: [matplotlib, plotly]. +=============================== ======================= ===================================================== diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/pandas_pyspark.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/pandas_pyspark.rst new file mode 100644 index 0000000..79bca27 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/pandas_pyspark.rst @@ -0,0 +1,136 @@ +.. 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/to pandas and PySpark DataFrames +===================================== + +.. currentmodule:: pyspark.pandas + +Users from pandas and/or PySpark face API compatibility issue sometimes when they +work with pandas API on Spark. Since pandas API on Spark does not target 100% compatibility of both pandas and +PySpark, users need to do some workaround to port their pandas and/or PySpark codes or +get familiar with pandas API on Spark in this case. This page aims to describe it. + + +pandas +------ + +pandas users can access the full pandas API by calling :func:`DataFrame.to_pandas`. +pandas-on-Spark DataFrame and pandas DataFrame are similar. However, the former is distributed +and the latter is in a single machine. When converting to each other, the data is +transferred between multiple machines and the single client machine. + +For example, if you need to call ``pandas_df.values`` of pandas DataFrame, you can do +as below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> + >>> psdf = ps.range(10) + >>> pdf = psdf.to_pandas() + >>> pdf.values + array([[0], + [1], + [2], + [3], + [4], + [5], + [6], + [7], + [8], + [9]]) + +pandas DataFrame can be a pandas-on-Spark DataFrame easily as below: + +.. code-block:: python + + >>> ps.from_pandas(pdf) + id + 0 0 + 1 1 + 2 2 + 3 3 + 4 4 + 5 5 + 6 6 + 7 7 + 8 8 + 9 9 + +Note that converting pandas-on-Spark DataFrame to pandas requires to collect all the data into the client machine; therefore, +if possible, it is recommended to use pandas API on Spark or PySpark APIs instead. + + +PySpark +------- + +PySpark users can access the full PySpark APIs by calling :func:`DataFrame.to_spark`. +pandas-on-Spark DataFrame and Spark DataFrame are virtually interchangeable. + +For example, if you need to call ``spark_df.filter(...)`` of Spark DataFrame, you can do +as below: + +.. code-block:: python + + >>> import pyspark.pandas as ps + >>> + >>> psdf = ps.range(10) + >>> sdf = psdf.to_spark().filter("id > 5") + >>> sdf.show() + +---+ + | id| + +---+ + | 6| + | 7| + | 8| + | 9| + +---+ + +Spark DataFrame can be a pandas-on-Spark DataFrame easily as below: + +.. code-block:: python + + >>> sdf.pandas_api() + id + 0 6 + 1 7 + 2 8 + 3 9 + +However, note that a new default index is created when pandas-on-Spark DataFrame is created from +Spark DataFrame. See `Default Index Type `_. In order to avoid this overhead, specify the column +to use as an index when possible. + +.. code-block:: python + + >>> # Create a pandas-on-Spark DataFrame with an explicit index. + ... psdf = ps.DataFrame({'id': range(10)}, index=range(10)) + >>> # Keep the explicit index. + ... sdf = psdf.to_spark(index_col='index') + >>> # Call Spark APIs + ... sdf = sdf.filter("id > 5") + >>> # Uses the explicit index to avoid to create default index. + ... sdf.pandas_api(index_col='index') + id + index + 6 6 + 7 7 + 8 8 + 9 9 diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/transform_apply.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/transform_apply.rst new file mode 100644 index 0000000..0778006 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/transform_apply.rst @@ -0,0 +1,139 @@ +.. 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. + + +============================== +Transform and apply a function +============================== + +.. NOTE: the images are stored at https://github.com/koalas/issues/1443. Feel free to edit and/or add. + +.. currentmodule:: pyspark.pandas + +There are many APIs that allow users to apply a function against pandas-on-Spark DataFrame such as +:func:`DataFrame.transform`, :func:`DataFrame.apply`, :func:`DataFrame.pandas_on_spark.transform_batch`, +:func:`DataFrame.pandas_on_spark.apply_batch`, :func:`Series.pandas_on_spark.transform_batch`, etc. Each has a distinct +purpose and works differently internally. This section describes the differences among +them where users are confused often. + +``transform`` and ``apply`` +--------------------------- + +The main difference between :func:`DataFrame.transform` and :func:`DataFrame.apply` is that the former requires +to return the same length of the input and the latter does not require this. See the example below: + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]}) + >>> def pandas_plus(pser): + ... return pser + 1 # should always return the same length as input. + ... + >>> psdf.transform(pandas_plus) + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[5,6,7]}) + >>> def pandas_plus(pser): + ... return pser[pser % 2 == 1] # allows an arbitrary length + ... + >>> psdf.apply(pandas_plus) + +In this case, each function takes a pandas Series, and the pandas API on Spark computes the functions in a distributed manner as below. + +.. image:: ../../../../../docs/img/pyspark-pandas_on_spark-transform_apply1.png + :alt: transform and apply + :align: center + :width: 550 + +In the case of 'column' axis, the function takes each row as a pandas Series. + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]}) + >>> def pandas_plus(pser): + ... return sum(pser) # allows an arbitrary length + ... + >>> psdf.apply(pandas_plus, axis='columns') + +The example above calculates the summation of each row as a pandas Series. See below: + +.. image:: ../../../../../docs/img/pyspark-pandas_on_spark-transform_apply2.png + :alt: apply axis + :align: center + :width: 600 + +In the examples above, the type hints were not used for simplicity but it is encouraged to use them to avoid performance penalty. +Please refer to the API documentations. + + +``pandas_on_spark.transform_batch`` and ``pandas_on_spark.apply_batch`` +----------------------------------------------------------------------- + +In :func:`DataFrame.pandas_on_spark.transform_batch`, :func:`DataFrame.pandas_on_spark.apply_batch`, :func:`Series.pandas_on_spark.transform_batch`, etc., the ``batch`` +postfix means each chunk in pandas-on-Spark DataFrame or Series. The APIs slice the pandas-on-Spark DataFrame or Series, and +then apply the given function with pandas DataFrame or Series as input and output. See the examples below: + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]}) + >>> def pandas_plus(pdf): + ... return pdf + 1 # should always return the same length as input. + ... + >>> psdf.pandas_on_spark.transform_batch(pandas_plus) + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]}) + >>> def pandas_plus(pdf): + ... return pdf[pdf.a > 1] # allow arbitrary length + ... + >>> psdf.pandas_on_spark.apply_batch(pandas_plus) + +The functions in both examples take a pandas DataFrame as a chunk of pandas-on-Spark DataFrame, and output a pandas DataFrame. +Pandas API on Spark combines the pandas DataFrames as a pandas-on-Spark DataFrame. + +Note that :func:`DataFrame.pandas_on_spark.transform_batch` has the length restriction - the length of input and output should be +the same - whereas :func:`DataFrame.pandas_on_spark.apply_batch` does not. However, it is important to know that +the output belongs to the same DataFrame when :func:`DataFrame.pandas_on_spark.transform_batch` returns a Series, and +you can avoid a shuffle by the operations between different DataFrames. In case of :func:`DataFrame.pandas_on_spark.apply_batch`, its output is always +treated as though it belongs to a new different DataFrame. See also +`Operations on different DataFrames `_ for more details. + +.. image:: ../../../../../docs/img/pyspark-pandas_on_spark-transform_apply3.png + :alt: pandas_on_spark.transform_batch and pandas_on_spark.apply_batch in Frame + :align: center + :width: 650 + +In case of :func:`Series.pandas_on_spark.transform_batch`, it is also similar with :func:`DataFrame.pandas_on_spark.transform_batch`; however, it takes +a pandas Series as a chunk of pandas-on-Spark Series. + +.. code-block:: python + + >>> psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]}) + >>> def pandas_plus(pser): + ... return pser + 1 # should always return the same length as input. + ... + >>> psdf.a.pandas_on_spark.transform_batch(pandas_plus) + +Under the hood, each batch of pandas-on-Spark Series is split to multiple pandas Series, and each function computes on that as below: + +.. image:: ../../../../../docs/img/pyspark-pandas_on_spark-transform_apply4.png + :alt: pandas_on_spark.transform_batch in Series + :width: 350 + :align: center + +There are more details such as the type inference and preventing its performance penalty. Please refer to the API documentations. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/typehints.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/typehints.rst new file mode 100644 index 0000000..1405baa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/typehints.rst @@ -0,0 +1,264 @@ +.. 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. + + +================================= +Type Hints in Pandas API on Spark +================================= + +.. currentmodule:: pyspark.pandas + +Pandas API on Spark, by default, infers the schema by taking some top records from the output, +in particular, when you use APIs that allow users to apply a function against pandas-on-Spark DataFrame +such as :func:`DataFrame.transform`, :func:`DataFrame.apply`, :func:`DataFrame.pandas_on_spark.apply_batch`, +:func:`DataFrame.pandas_on_spark.apply_batch`, :func:`Series.pandas_on_spark.apply_batch`, etc. + +However, this is potentially expensive. If there are several expensive operations such as a shuffle +in the upstream of the execution plan, pandas API on Spark will end up with executing the Spark job twice, once +for schema inference, and once for processing actual data with the schema. + +To avoid the consequences, pandas API on Spark has its own type hinting style to specify the schema to avoid +schema inference. Pandas API on Spark understands the type hints specified in the return type and converts it +as a Spark schema for pandas UDFs used internally. The way of type hinting has evolved over time. + +This chapter covers the recommended way and the supported ways in detail. + +.. note:: + The variadic generics support is experimental and unstable in pandas API on Spark. + The way of typing can change between minor releases without a warning. + See also `PEP 646 `_ for variadic generics in Python. + + +pandas-on-Spark DataFrame and Pandas DataFrame +---------------------------------------------- + +In the early pandas-on-Spark version, it was introduced to specify a type hint in the function in order to use +it as a Spark schema. As an example, you can specify the return type hint as below by using pandas-on-Spark +:class:`DataFrame`. + +.. code-block:: python + + >>> def pandas_div(pdf) -> ps.DataFrame[float, float]: + ... # pdf is a pandas DataFrame. + ... return pdf[['B', 'C']] / pdf[['B', 'C']] + ... + >>> df = ps.DataFrame({'A': ['a', 'a', 'b'], 'B': [1, 2, 3], 'C': [4, 6, 5]}) + >>> df.groupby('A').apply(pandas_div) + +Notice that the function ``pandas_div`` actually takes and outputs a pandas DataFrame instead of +pandas-on-Spark :class:`DataFrame`. So, technically the correct types should be of pandas. + +With Python 3.8+, you can specify the type hints by using pandas instances as follows: + +.. code-block:: python + + >>> def pandas_div(pdf) -> pd.DataFrame[float, float]: + ... # pdf is a pandas DataFrame. + ... return pdf[['B', 'C']] / pdf[['B', 'C']] + ... + >>> df = ps.DataFrame({'A': ['a', 'a', 'b'], 'B': [1, 2, 3], 'C': [4, 6, 5]}) + >>> df.groupby('A').apply(pandas_div) + +Likewise, pandas Series can be also used as a type hints: + +.. code-block:: python + + >>> def sqrt(x) -> pd.Series[float]: + ... return np.sqrt(x) + ... + >>> df = ps.DataFrame([[4, 9]] * 3, columns=['A', 'B']) + >>> df.apply(sqrt, axis=0) + +Currently, both pandas API on Spark and pandas instances can be used to specify the type hints; however, pandas-on-Spark +plans to move gradually towards using pandas instances only as the stability becomes proven. + + +Type Hinting with Names +----------------------- + +This approach is to overcome the limitations in the existing type +hinting especially for DataFrame. When you use a DataFrame as the return type hint, for example, +``DataFrame[int, int]``, there is no way to specify the names of each Series. In the old way, pandas API on Spark just generates +the column names as ``c#`` and this easily leads users to lose or forget the Series mappings. See the example below: + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[int, int]: + ... pdf['A'] = pdf.id + 1 + ... return pdf + ... + >>> ps.range(5).pandas_on_spark.apply_batch(transform) + +.. code-block:: bash + + c0 c1 + 0 0 1 + 1 1 2 + 2 2 3 + 3 3 4 + 4 4 5 + +The new style of type hinting in pandas API on Spark is similar to the regular Python type hints in variables. The Series name +is specified as a string, and the type is specified after a colon. The following example shows a simple case with +the Series names, ``id`` and ``A``, and ``int`` types respectively. + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame["id": int, "A": int]: + ... pdf['A'] = pdf.id + 1 + ... return pdf + ... + >>> ps.range(5).pandas_on_spark.apply_batch(transform) + +.. code-block:: bash + + id A + 0 0 1 + 1 1 2 + 2 2 3 + 3 3 4 + 4 4 5 + +In addition, pandas API on Spark also dynamically supports ``dtype`` instances and the column index in pandas so that users can +programmatically generate the return type and schema. + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + .. zip(sample.columns, sample.dtypes)]: + ... return pdf + 1 + ... + >>> psdf.pandas_on_spark.apply_batch(transform) + +Likewise, ``dtype`` instances from pandas DataFrame can be used alone and let pandas API on Spark generate column names. + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[sample.dtypes]: + ... return pdf + 1 + ... + >>> psdf.pandas_on_spark.apply_batch(transform) + + +Type Hinting with Index +----------------------- + +When you omit index types in the type hints, pandas API on Spark attaches the default index (`compute.default_index_type`), +and it loses the index column and information from the original data. The default index sometimes requires to have an +expensive computation such as shuffle so it is best to specify the index type together. + + +Index +~~~~~ + +With the pandas DataFrames below: + +.. code-block:: python + + >>> pdf = pd.DataFrame({'id': range(5)}) + >>> sample = pdf.copy() + >>> sample["a"] = sample.id + 1 + +The ways below are allowed for a regular index: + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[int, [int, int]]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... sample.index.dtype, sample.dtypes]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... ("idxA", int), [("id", int), ("a", int)]]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... (sample.index.name, sample.index.dtype), + ... zip(sample.columns, sample.dtypes)]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + + +MultiIndex +~~~~~~~~~~ + +With the pandas DataFrames below: + + >>> midx = pd.MultiIndex.from_arrays( + ... [(1, 1, 2), (1.5, 4.5, 7.5)], + ... names=("int", "float")) + >>> pdf = pd.DataFrame(range(3), index=midx, columns=["id"]) + >>> sample = pdf.copy() + >>> sample["a"] = sample.id + 1 + +The ways below are allowed for multi-index: + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[[int, float], [int, int]]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... sample.index.dtypes, sample.dtypes]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... [("int", int), ("float", float)], + ... [("id", int), ("a", int)]]: + ... pdf["a"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) + +.. code-block:: python + + >>> def transform(pdf) -> pd.DataFrame[ + ... zip(sample.index.names, sample.index.dtypes), + ... zip(sample.columns, sample.dtypes)]: + ... pdf["A"] = pdf.id + 1 + ... return pdf + ... + >>> ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform) diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/types.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/types.rst new file mode 100644 index 0000000..a806410 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/pandas_on_spark/types.rst @@ -0,0 +1,242 @@ +.. 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. + + +=================================== +Type Support in Pandas API on Spark +=================================== + +.. currentmodule:: pyspark.pandas + +In this chapter, we will briefly show you how data types change when converting pandas-on-Spark DataFrame from/to PySpark DataFrame or pandas DataFrame. + + +Type casting between PySpark and pandas API on Spark +---------------------------------------------------- + +When converting a pandas-on-Spark DataFrame from/to PySpark DataFrame, the data types are automatically casted to the appropriate type. + +The example below shows how data types are casted from PySpark DataFrame to pandas-on-Spark DataFrame. + +.. code-block:: python + + # 1. Create a PySpark DataFrame + >>> sdf = spark.createDataFrame([ + ... (1, Decimal(1.0), 1., 1., 1, 1, 1, datetime(2020, 10, 27), "1", True, datetime(2020, 10, 27)), + ... ], 'tinyint tinyint, decimal decimal, float float, double double, integer integer, long long, short short, timestamp timestamp, string string, boolean boolean, date date') + + # 2. Check the PySpark data types + >>> sdf + DataFrame[tinyint: tinyint, decimal: decimal(10,0), float: float, double: double, integer: int, long: bigint, short: smallint, timestamp: timestamp, string: string, boolean: boolean, date: date] + + # 3. Convert PySpark DataFrame to pandas-on-Spark DataFrame + >>> psdf = sdf.pandas_api() + + # 4. Check the pandas-on-Spark data types + >>> psdf.dtypes + tinyint int8 + decimal object + float float32 + double float64 + integer int32 + long int64 + short int16 + timestamp datetime64[ns] + string object + boolean bool + date object + dtype: object + + +The example below shows how data types are casted from pandas-on-Spark DataFrame to PySpark DataFrame. + +.. code-block:: python + + # 1. Create a pandas-on-Spark DataFrame + >>> psdf = ps.DataFrame({"int8": [1], "bool": [True], "float32": [1.0], "float64": [1.0], "int32": [1], "int64": [1], "int16": [1], "datetime": [datetime.datetime(2020, 10, 27)], "object_string": ["1"], "object_decimal": [decimal.Decimal("1.1")], "object_date": [datetime.date(2020, 10, 27)]}) + + # 2. Type casting by using `astype` + >>> psdf['int8'] = psdf['int8'].astype('int8') + >>> psdf['int16'] = psdf['int16'].astype('int16') + >>> psdf['int32'] = psdf['int32'].astype('int32') + >>> psdf['float32'] = psdf['float32'].astype('float32') + + # 3. Check the pandas-on-Spark data types + >>> psdf.dtypes + int8 int8 + bool bool + float32 float32 + float64 float64 + int32 int32 + int64 int64 + int16 int16 + datetime datetime64[ns] + object_string object + object_decimal object + object_date object + dtype: object + + # 4. Convert pandas-on-Spark DataFrame to PySpark DataFrame + >>> sdf = psdf.to_spark() + + # 5. Check the PySpark data types + >>> sdf + DataFrame[int8: tinyint, bool: boolean, float32: float, float64: double, int32: int, int64: bigint, int16: smallint, datetime: timestamp, object_string: string, object_decimal: decimal(2,1), object_date: date] + + +Type casting between pandas and pandas API on Spark +--------------------------------------------------- + +When converting pandas-on-Spark DataFrame to pandas DataFrame, the data types are basically the same as pandas. + +.. code-block:: python + + # Convert pandas-on-Spark DataFrame to pandas DataFrame + >>> pdf = psdf.to_pandas() + + # Check the pandas data types + >>> pdf.dtypes + int8 int8 + bool bool + float32 float32 + float64 float64 + int32 int32 + int64 int64 + int16 int16 + datetime datetime64[ns] + object_string object + object_decimal object + object_date object + dtype: object + + +However, there are several data types only provided by pandas. + +.. code-block:: python + + # pd.Catrgorical type is not supported in pandas API on Spark yet. + >>> ps.Series([pd.Categorical([1, 2, 3])]) + Traceback (most recent call last): + ... + pyarrow.lib.ArrowInvalid: Could not convert [1, 2, 3] + Categories (3, int64): [1, 2, 3] with type Categorical: did not recognize Python value type when inferring an Arrow data type + + +These kinds of pandas specific data types below are not currently supported in the pandas API on Spark but planned to be supported. + +* pd.Timedelta +* pd.Categorical +* pd.CategoricalDtype + + +The pandas specific data types below are not planned to be supported in the pandas API on Spark yet. + +* pd.SparseDtype +* pd.DatetimeTZDtype +* pd.UInt*Dtype +* pd.BooleanDtype +* pd.StringDtype + + +Internal type mapping +--------------------- + +The table below shows which NumPy data types are matched to which PySpark data types internally in the pandas API on Spark. + +============= ======================= +NumPy PySpark +============= ======================= +np.character BinaryType +np.bytes\_ BinaryType +np.string\_ BinaryType +np.int8 ByteType +np.byte ByteType +np.int16 ShortType +np.int32 IntegerType +np.int64 LongType +np.float32 FloatType +np.float64 DoubleType +np.unicode\_ StringType +np.datetime64 TimestampType +np.ndarray ArrayType(StringType()) +============= ======================= + + +The table below shows which Python data types are matched to which PySpark data types internally in pandas API on Spark. + +================= =================== +Python PySpark +================= =================== +bytes BinaryType +int LongType +float DoubleType +str StringType +bool BooleanType +datetime.datetime TimestampType +datetime.date DateType +decimal.Decimal DecimalType(38, 18) +================= =================== + +For decimal type, pandas API on Spark uses Spark's system default precision and scale. + +You can check this mapping by using the `as_spark_type` function. + +.. code-block:: python + + >>> import typing + >>> import numpy as np + >>> from pyspark.pandas.typedef import as_spark_type + + >>> as_spark_type(int) + LongType + + >>> as_spark_type(np.int32) + IntegerType + + >>> as_spark_type(typing.List[float]) + ArrayType(DoubleType,true) + + +You can also check the underlying PySpark data type of `Series` or schema of `DataFrame` by using Spark accessor. + +.. code-block:: python + + >>> ps.Series([0.3, 0.1, 0.8]).spark.data_type + DoubleType + + >>> ps.Series(["welcome", "to", "pandas-on-Spark"]).spark.data_type + StringType + + >>> ps.Series([[False, True, False]]).spark.data_type + ArrayType(BooleanType,true) + + >>> ps.DataFrame({"d": [0.3, 0.1, 0.8], "s": ["welcome", "to", "pandas-on-Spark"], "b": [False, True, False]}).spark.print_schema() + root + |-- d: double (nullable = false) + |-- s: string (nullable = false) + |-- b: boolean (nullable = false) + +.. note:: + + Pandas API on Spark currently does not support multiple types of data in a single column. + + .. code-block:: python + + >>> ps.Series([1, "A"]) + Traceback (most recent call last): + ... + TypeError: an integer is required (got type str) diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/python_packaging.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/python_packaging.rst new file mode 100644 index 0000000..0284313 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/python_packaging.rst @@ -0,0 +1,253 @@ +.. 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. + + +========================= +Python Package Management +========================= + +When you want to run your PySpark application on a cluster such as YARN, Kubernetes, Mesos, etc., you need to make +sure that your code and all used libraries are available on the executors. + +As an example, let's say you may want to run the `Pandas UDF examples `_. +As it uses pyarrow as an underlying implementation we need to make sure to have pyarrow installed on each executor +on the cluster. Otherwise you may get errors such as ``ModuleNotFoundError: No module named 'pyarrow'``. + +Here is the script ``app.py`` from the previous example that will be executed on the cluster: + +.. code-block:: python + + import pandas as pd + from pyspark.sql.functions import pandas_udf + from pyspark.sql import SparkSession + + def main(spark): + df = spark.createDataFrame( + [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], + ("id", "v")) + + @pandas_udf("double") + def mean_udf(v: pd.Series) -> float: + return v.mean() + + print(df.groupby("id").agg(mean_udf(df['v'])).collect()) + + + if __name__ == "__main__": + main(SparkSession.builder.getOrCreate()) + + +There are multiple ways to manage Python dependencies in the cluster: + +- Using PySpark Native Features +- Using Conda +- Using Virtualenv +- Using PEX + + +Using PySpark Native Features +----------------------------- + +PySpark allows to upload Python files (``.py``), zipped Python packages (``.zip``), and Egg files (``.egg``) +to the executors by one of the following: + +- Setting the configuration setting ``spark.submit.pyFiles`` +- Setting ``--py-files`` option in Spark scripts +- Directly calling :meth:`pyspark.SparkContext.addPyFile` in applications + +This is a straightforward method to ship additional custom Python code to the cluster. You can just add individual files or zip whole +packages and upload them. Using :meth:`pyspark.SparkContext.addPyFile` allows you to upload code even after having started your job. + +However, it does not allow to add packages built as `Wheels `_ and therefore +does not allow to include dependencies with native code. + + +Using Conda +----------- + +`Conda `_ is one of the most widely-used Python package management systems. PySpark users can directly +use a Conda environment to ship their third-party Python packages by leveraging +`conda-pack `_ which is a command line tool creating +relocatable Conda environments. + +The example below creates a Conda environment to use on both the driver and executor and packs +it into an archive file. This archive file captures the Conda environment for Python and stores +both Python interpreter and all its relevant dependencies. + +.. code-block:: bash + + conda create -y -n pyspark_conda_env -c conda-forge pyarrow pandas conda-pack + conda activate pyspark_conda_env + conda pack -f -o pyspark_conda_env.tar.gz + +After that, you can ship it together with scripts or in the code by using the ``--archives`` option +or ``spark.archives`` configuration (``spark.yarn.dist.archives`` in YARN). It automatically unpacks the archive on executors. + +In the case of a ``spark-submit`` script, you can use it as follows: + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python # Do not set in cluster modes. + export PYSPARK_PYTHON=./environment/bin/python + spark-submit --archives pyspark_conda_env.tar.gz#environment app.py + +Note that ``PYSPARK_DRIVER_PYTHON`` above should not be set for cluster modes in YARN or Kubernetes. + +If you're on a regular Python shell or notebook, you can try it as shown below: + +.. code-block:: python + + import os + from pyspark.sql import SparkSession + from app import main + + os.environ['PYSPARK_PYTHON'] = "./environment/bin/python" + spark = SparkSession.builder.config( + "spark.archives", # 'spark.yarn.dist.archives' in YARN. + "pyspark_conda_env.tar.gz#environment").getOrCreate() + main(spark) + +For a pyspark shell: + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python + export PYSPARK_PYTHON=./environment/bin/python + pyspark --archives pyspark_conda_env.tar.gz#environment + + +Using Virtualenv +---------------- + +`Virtualenv `_ is a Python tool to create isolated Python environments. +Since Python 3.3, a subset of its features has been integrated into Python as a standard library under +the `venv `_ module. PySpark users can use virtualenv to manage +Python dependencies in their clusters by using `venv-pack `_ +in a similar way as conda-pack. + +A virtual environment to use on both driver and executor can be created as demonstrated below. +It packs the current virtual environment to an archive file, and it contains both Python interpreter and the dependencies. +However, it requires all nodes in a cluster to have the same Python interpreter installed because +`venv-pack packs Python interpreter as a symbolic link `_. + + +.. code-block:: bash + + python -m venv pyspark_venv + source pyspark_venv/bin/activate + pip install pyarrow pandas venv-pack + venv-pack -o pyspark_venv.tar.gz + +You can directly pass/unpack the archive file and enable the environment on executors by leveraging +the ``--archives`` option or ``spark.archives`` configuration (``spark.yarn.dist.archives`` in YARN). + +For ``spark-submit``, you can use it by running the command as follows. Also, notice that +``PYSPARK_DRIVER_PYTHON`` has to be unset in Kubernetes or YARN cluster modes. + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python # Do not set in cluster modes. + export PYSPARK_PYTHON=./environment/bin/python + spark-submit --archives pyspark_venv.tar.gz#environment app.py + +For regular Python shells or notebooks: + +.. code-block:: bash + + import os + from pyspark.sql import SparkSession + from app import main + + os.environ['PYSPARK_PYTHON'] = "./environment/bin/python" + spark = SparkSession.builder.config( + "spark.archives", # 'spark.yarn.dist.archives' in YARN. + "pyspark_venv.tar.gz#environment").getOrCreate() + main(spark) + +In the case of a pyspark shell: + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python + export PYSPARK_PYTHON=./environment/bin/python + pyspark --archives pyspark_venv.tar.gz#environment + + +Using PEX +--------- + +PySpark can also use `PEX `_ to ship the Python packages +together. PEX is a tool that creates a self-contained Python environment. This is similar +to Conda or virtualenv, but a ``.pex`` file is executable by itself. + +The following example creates a ``.pex`` file for the driver and executor to use. +The file contains the Python dependencies specified with the ``pex`` command. + +.. code-block:: bash + + pip install pyarrow pandas pex + pex pyspark pyarrow pandas -o pyspark_pex_env.pex + +This file behaves similarly with a regular Python interpreter. + +.. code-block:: bash + + ./pyspark_pex_env.pex -c "import pandas; print(pandas.__version__)" + 1.1.5 + +However, ``.pex`` file does not include a Python interpreter itself under the hood so all +nodes in a cluster should have the same Python interpreter installed. + +In order to transfer and use the ``.pex`` file in a cluster, you should ship it via the +``spark.files`` configuration (``spark.yarn.dist.files`` in YARN) or ``--files`` option because they are regular files instead +of directories or archive files. + +For application submission, you run the commands as shown below. +Note that ``PYSPARK_DRIVER_PYTHON`` should not be set for cluster modes in YARN or Kubernetes. + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python # Do not set in cluster modes. + export PYSPARK_PYTHON=./pyspark_pex_env.pex + spark-submit --files pyspark_pex_env.pex app.py + +For regular Python shells or notebooks: + +.. code-block:: python + + import os + from pyspark.sql import SparkSession + from app import main + + os.environ['PYSPARK_PYTHON'] = "./pyspark_pex_env.pex" + spark = SparkSession.builder.config( + "spark.files", # 'spark.yarn.dist.files' in YARN. + "pyspark_pex_env.pex").getOrCreate() + main(spark) + +For the interactive pyspark shell, the commands are almost the same: + +.. code-block:: bash + + export PYSPARK_DRIVER_PYTHON=python + export PYSPARK_PYTHON=./pyspark_pex_env.pex + pyspark --files pyspark_pex_env.pex + +An end-to-end Docker example for deploying a standalone PySpark with ``SparkSession.builder`` and PEX +can be found `here `_ +- it uses cluster-pack, a library on top of PEX that automatizes the intermediate step of having +to create & upload the PEX manually. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/arrow_pandas.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/arrow_pandas.rst new file mode 100644 index 0000000..e9355cd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/arrow_pandas.rst @@ -0,0 +1,429 @@ +.. 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. + +======================= +Apache Arrow in PySpark +======================= + +.. currentmodule:: pyspark.sql + +Apache Arrow is an in-memory columnar data format that is used in Spark to efficiently transfer +data between JVM and Python processes. This currently is most beneficial to Python users that +work with Pandas/NumPy data. Its usage is not automatic and might require some minor +changes to configuration or code to take full advantage and ensure compatibility. This guide will +give a high-level description of how to use Arrow in Spark and highlight any differences when +working with Arrow-enabled data. + +Ensure PyArrow Installed +------------------------ + +To use Apache Arrow in PySpark, `the recommended version of PyArrow `_ +should be installed. +If you install PySpark using pip, then PyArrow can be brought in as an extra dependency of the +SQL module with the command ``pip install pyspark[sql]``. Otherwise, you must ensure that PyArrow +is installed and available on all cluster nodes. +You can install it using pip or conda from the conda-forge channel. See PyArrow +`installation `_ for details. + +Enabling for Conversion to/from Pandas +-------------------------------------- + +Arrow is available as an optimization when converting a Spark DataFrame to a Pandas DataFrame +using the call :meth:`DataFrame.toPandas` and when creating a Spark DataFrame from a Pandas DataFrame with +:meth:`SparkSession.createDataFrame`. To use Arrow when executing these calls, users need to first set +the Spark configuration ``spark.sql.execution.arrow.pyspark.enabled`` to ``true``. This is disabled by default. + +In addition, optimizations enabled by ``spark.sql.execution.arrow.pyspark.enabled`` could fallback automatically +to non-Arrow optimization implementation if an error occurs before the actual computation within Spark. +This can be controlled by ``spark.sql.execution.arrow.pyspark.fallback.enabled``. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 37-52 + :dedent: 4 + +Using the above optimizations with Arrow will produce the same results as when Arrow is not +enabled. + +Note that even with Arrow, :meth:`DataFrame.toPandas` results in the collection of all records in the +DataFrame to the driver program and should be done on a small subset of the data. Not all Spark +data types are currently supported and an error can be raised if a column has an unsupported type. +If an error occurs during :meth:`SparkSession.createDataFrame`, Spark will fall back to create the +DataFrame without Arrow. + +Pandas UDFs (a.k.a. Vectorized UDFs) +------------------------------------ + +.. currentmodule:: pyspark.sql.functions + +Pandas UDFs are user defined functions that are executed by Spark using +Arrow to transfer data and Pandas to work with the data, which allows vectorized operations. A Pandas +UDF is defined using the :meth:`pandas_udf` as a decorator or to wrap the function, and no additional +configuration is required. A Pandas UDF behaves as a regular PySpark function API in general. + +Before Spark 3.0, Pandas UDFs used to be defined with ``pyspark.sql.functions.PandasUDFType``. From Spark 3.0 +with Python 3.6+, you can also use `Python type hints `_. +Using Python type hints is preferred and using ``pyspark.sql.functions.PandasUDFType`` will be deprecated in +the future release. + +.. currentmodule:: pyspark.sql.types + +Note that the type hint should use ``pandas.Series`` in all cases but there is one variant +that ``pandas.DataFrame`` should be used for its input or output type hint instead when the input +or output column is of :class:`StructType`. The following example shows a Pandas UDF which takes long +column, string column and struct column, and outputs a struct column. It requires the function to +specify the type hints of ``pandas.Series`` and ``pandas.DataFrame`` as below: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 56-80 + :dedent: 4 + +In the following sections, it describes the combinations of the supported type hints. For simplicity, +``pandas.DataFrame`` variant is omitted. + +Series to Series +~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql.functions + +The type hint can be expressed as ``pandas.Series``, ... -> ``pandas.Series``. + +By using :func:`pandas_udf` with the function having such type hints above, it creates a Pandas UDF where the given +function takes one or more ``pandas.Series`` and outputs one ``pandas.Series``. The output of the function should +always be of the same length as the input. Internally, PySpark will execute a Pandas UDF by splitting +columns into batches and calling the function for each batch as a subset of the data, then concatenating +the results together. + +The following example shows how to create this Pandas UDF that computes the product of 2 columns. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 84-114 + :dedent: 4 + +For detailed usage, please see :func:`pandas_udf`. + +Iterator of Series to Iterator of Series +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql.functions + +The type hint can be expressed as ``Iterator[pandas.Series]`` -> ``Iterator[pandas.Series]``. + +By using :func:`pandas_udf` with the function having such type hints above, it creates a Pandas UDF where the given +function takes an iterator of ``pandas.Series`` and outputs an iterator of ``pandas.Series``. The +length of the entire output from the function should be the same length of the entire input; therefore, it can +prefetch the data from the input iterator as long as the lengths are the same. +In this case, the created Pandas UDF requires one input column when the Pandas UDF is called. To use +multiple input columns, a different type hint is required. See Iterator of Multiple Series to Iterator +of Series. + +It is also useful when the UDF execution requires initializing some states although internally it works +identically as Series to Series case. The pseudocode below illustrates the example. + +.. code-block:: python + + @pandas_udf("long") + def calculate(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: + # Do some expensive initialization with a state + state = very_expensive_initialization() + for x in iterator: + # Use that state for the whole iterator. + yield calculate_with_state(x, state) + + df.select(calculate("value")).show() + +The following example shows how to create this Pandas UDF: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 118-140 + :dedent: 4 + +For detailed usage, please see :func:`pandas_udf`. + +Iterator of Multiple Series to Iterator of Series +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql.functions + +The type hint can be expressed as ``Iterator[Tuple[pandas.Series, ...]]`` -> ``Iterator[pandas.Series]``. + +By using :func:`pandas_udf` with the function having such type hints above, it creates a Pandas UDF where the +given function takes an iterator of a tuple of multiple ``pandas.Series`` and outputs an iterator of ``pandas.Series``. +In this case, the created pandas UDF requires multiple input columns as many as the series in the tuple +when the Pandas UDF is called. Otherwise, it has the same characteristics and restrictions as the Iterator of Series +to Iterator of Series case. + +The following example shows how to create this Pandas UDF: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 144-167 + :dedent: 4 + +For detailed usage, please see :func:`pandas_udf`. + +Series to Scalar +~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql.functions + +The type hint can be expressed as ``pandas.Series``, ... -> ``Any``. + +By using :func:`pandas_udf` with the function having such type hints above, it creates a Pandas UDF similar +to PySpark's aggregate functions. The given function takes `pandas.Series` and returns a scalar value. +The return type should be a primitive data type, and the returned scalar can be either a python +primitive type, e.g., ``int`` or ``float`` or a numpy data type, e.g., ``numpy.int64`` or ``numpy.float64``. +``Any`` should ideally be a specific scalar type accordingly. + +.. currentmodule:: pyspark.sql + +This UDF can be also used with :meth:`GroupedData.agg` and `Window`. +It defines an aggregation from one or more ``pandas.Series`` to a scalar value, where each ``pandas.Series`` +represents a column within the group or window. + +Note that this type of UDF does not support partial aggregation and all data for a group or window +will be loaded into memory. Also, only unbounded window is supported with Grouped aggregate Pandas +UDFs currently. The following example shows how to use this type of UDF to compute mean with a group-by +and window operations: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 171-212 + :dedent: 4 + +.. currentmodule:: pyspark.sql.functions + +For detailed usage, please see :func:`pandas_udf`. + +Pandas Function APIs +-------------------- + +.. currentmodule:: pyspark.sql + +Pandas Function APIs can directly apply a Python native function against the whole :class:`DataFrame` by +using Pandas instances. Internally it works similarly with Pandas UDFs by using Arrow to transfer +data and Pandas to work with the data, which allows vectorized operations. However, a Pandas Function +API behaves as a regular API under PySpark :class:`DataFrame` instead of :class:`Column`, and Python type hints in Pandas +Functions APIs are optional and do not affect how it works internally at this moment although they +might be required in the future. + +.. currentmodule:: pyspark.sql.functions + +From Spark 3.0, grouped map pandas UDF is now categorized as a separate Pandas Function API, +``DataFrame.groupby().applyInPandas()``. It is still possible to use it with ``pyspark.sql.functions.PandasUDFType`` +and ``DataFrame.groupby().apply()`` as it was; however, it is preferred to use +``DataFrame.groupby().applyInPandas()`` directly. Using ``pyspark.sql.functions.PandasUDFType`` will be deprecated +in the future. + +Grouped Map +~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql + +Grouped map operations with Pandas instances are supported by ``DataFrame.groupby().applyInPandas()`` +which requires a Python function that takes a ``pandas.DataFrame`` and return another ``pandas.DataFrame``. +It maps each group to each ``pandas.DataFrame`` in the Python function. + +This API implements the "split-apply-combine" pattern which consists of three steps: + +* Split the data into groups by using :meth:`DataFrame.groupBy`. + +* Apply a function on each group. The input and output of the function are both ``pandas.DataFrame``. The input data contains all the rows and columns for each group. + +* Combine the results into a new PySpark :class:`DataFrame`. + +To use ``DataFrame.groupBy().applyInPandas()``, the user needs to define the following: + +* A Python function that defines the computation for each group. + +* A ``StructType`` object or a string that defines the schema of the output PySpark :class:`DataFrame`. + +The column labels of the returned ``pandas.DataFrame`` must either match the field names in the +defined output schema if specified as strings, or match the field data types by position if not +strings, e.g. integer indices. See `pandas.DataFrame `_ +on how to label columns when constructing a ``pandas.DataFrame``. + +Note that all data for a group will be loaded into memory before the function is applied. This can +lead to out of memory exceptions, especially if the group sizes are skewed. The configuration for +`maxRecordsPerBatch `_ is not applied on groups and it is up to the user +to ensure that the grouped data will fit into the available memory. + +The following example shows how to use ``DataFrame.groupby().applyInPandas()`` to subtract the mean from each value +in the group. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 216-234 + :dedent: 4 + +For detailed usage, please see please see :meth:`GroupedData.applyInPandas` + +Map +~~~ + +Map operations with Pandas instances are supported by :meth:`DataFrame.mapInPandas` which maps an iterator +of ``pandas.DataFrame``\s to another iterator of ``pandas.DataFrame``\s that represents the current +PySpark :class:`DataFrame` and returns the result as a PySpark :class:`DataFrame`. The function takes and outputs +an iterator of ``pandas.DataFrame``. It can return the output of arbitrary length in contrast to some +Pandas UDFs although internally it works similarly with Series to Series Pandas UDF. + +The following example shows how to use :meth:`DataFrame.mapInPandas`: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 238-249 + :dedent: 4 + +For detailed usage, please see :meth:`DataFrame.mapInPandas`. + +Co-grouped Map +~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql + +Co-grouped map operations with Pandas instances are supported by ``DataFrame.groupby().cogroup().applyInPandas()`` which +allows two PySpark :class:`DataFrame`\s to be cogrouped by a common key and then a Python function applied to each +cogroup. It consists of the following steps: + +* Shuffle the data such that the groups of each dataframe which share a key are cogrouped together. + +* Apply a function to each cogroup. The input of the function is two ``pandas.DataFrame`` (with an optional tuple representing the key). The output of the function is a ``pandas.DataFrame``. + +* Combine the ``pandas.DataFrame``\s from all groups into a new PySpark :class:`DataFrame`. + +To use ``groupBy().cogroup().applyInPandas()``, the user needs to define the following: + +* A Python function that defines the computation for each cogroup. + +* A ``StructType`` object or a string that defines the schema of the output PySpark :class:`DataFrame`. + +The column labels of the returned ``pandas.DataFrame`` must either match the field names in the +defined output schema if specified as strings, or match the field data types by position if not +strings, e.g. integer indices. See `pandas.DataFrame `_. +on how to label columns when constructing a ``pandas.DataFrame``. + +Note that all data for a cogroup will be loaded into memory before the function is applied. This can lead to out of +memory exceptions, especially if the group sizes are skewed. The configuration for `maxRecordsPerBatch `_ +is not applied and it is up to the user to ensure that the cogrouped data will fit into the available memory. + +The following example shows how to use ``DataFrame.groupby().cogroup().applyInPandas()`` to perform an asof join between two datasets. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 253-275 + :dedent: 4 + + +For detailed usage, please see :meth:`PandasCogroupedOps.applyInPandas` + +Arrow Python UDFs +----------------- + +Arrow Python UDFs are user defined functions that are executed row-by-row, utilizing Arrow for efficient batch data +transfer and serialization. To define an Arrow Python UDF, you can use the :meth:`udf` decorator or wrap the function +with the :meth:`udf` method, ensuring the ``useArrow`` parameter is set to True. Additionally, you can enable Arrow +optimization for Python UDFs throughout the entire SparkSession by setting the Spark configuration ``spark.sql +.execution.pythonUDF.arrow.enabled`` to true. It's important to note that the Spark configuration takes effect only +when ``useArrow`` is either not set or set to None. + +The type hints for Arrow Python UDFs should be specified in the same way as for default, pickled Python UDFs. + +Here's an example that demonstrates the usage of both a default, pickled Python UDF and an Arrow Python UDF: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py + :language: python + :lines: 279-297 + :dedent: 4 + +Compared to the default, pickled Python UDFs, Arrow Python UDFs provide a more coherent type coercion mechanism. UDF +type coercion poses challenges when the Python instances returned by UDFs do not align with the user-specified +return type. The default, pickled Python UDFs' type coercion has certain limitations, such as relying on None as a +fallback for type mismatches, leading to potential ambiguity and data loss. Additionally, converting date, datetime, +and tuples to strings can yield ambiguous results. Arrow Python UDFs, on the other hand, leverage Arrow's +capabilities to standardize type coercion and address these issues effectively. + +Usage Notes +----------- + +Supported SQL Types +~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql.types + +Currently, all Spark SQL data types are supported by Arrow-based conversion except +:class:`ArrayType` of :class:`TimestampType`. +:class:`MapType` and :class:`ArrayType` of nested :class:`StructType` are only supported +when using PyArrow 2.0.0 and above. + +Setting Arrow Batch Size +~~~~~~~~~~~~~~~~~~~~~~~~ + +Data partitions in Spark are converted into Arrow record batches, which can temporarily lead to +high memory usage in the JVM. To avoid possible out of memory exceptions, the size of the Arrow +record batches can be adjusted by setting the conf ``spark.sql.execution.arrow.maxRecordsPerBatch`` +to an integer that will determine the maximum number of rows for each batch. The default value is +10,000 records per batch. If the number of columns is large, the value should be adjusted +accordingly. Using this limit, each data partition will be made into 1 or more record batches for +processing. + +Timestamp with Time Zone Semantics +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: pyspark.sql + +Spark internally stores timestamps as UTC values, and timestamp data that is brought in without +a specified time zone is converted as local time to UTC with microsecond resolution. When timestamp +data is exported or displayed in Spark, the session time zone is used to localize the timestamp +values. The session time zone is set with the configuration ``spark.sql.session.timeZone`` and will +default to the JVM system local time zone if not set. Pandas uses a ``datetime64`` type with nanosecond +resolution, ``datetime64[ns]``, with optional time zone on a per-column basis. + +When timestamp data is transferred from Spark to Pandas it will be converted to nanoseconds +and each column will be converted to the Spark session time zone then localized to that time +zone, which removes the time zone and displays values as local time. This will occur +when calling :meth:`DataFrame.toPandas()` or ``pandas_udf`` with timestamp columns. + +When timestamp data is transferred from Pandas to Spark, it will be converted to UTC microseconds. This +occurs when calling :meth:`SparkSession.createDataFrame` with a Pandas DataFrame or when returning a timestamp from a +``pandas_udf``. These conversions are done automatically to ensure Spark will have data in the +expected format, so it is not necessary to do any of these conversions yourself. Any nanosecond +values will be truncated. + +Note that a standard UDF (non-Pandas) will load timestamp data as Python datetime objects, which is +different from a Pandas timestamp. It is recommended to use Pandas time series functionality when +working with timestamps in ``pandas_udf``\s to get the best performance, see +`here `_ for details. + +Recommended Pandas and PyArrow Versions +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For usage with pyspark.sql, the minimum supported versions of Pandas is 1.0.5 and PyArrow is 4.0.0. +Higher versions may be used, however, compatibility and data correctness can not be guaranteed and should +be verified by the user. + +Setting Arrow ``self_destruct`` for memory savings +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Since Spark 3.2, the Spark configuration ``spark.sql.execution.arrow.pyspark.selfDestruct.enabled`` can be used to enable PyArrow's ``self_destruct`` feature, which can save memory when creating a Pandas DataFrame via ``toPandas`` by freeing Arrow-allocated memory while building the Pandas DataFrame. +This option is experimental, and some operations may fail on the resulting Pandas DataFrame due to immutable backing arrays. +Typically, you would see the error ``ValueError: buffer source array is read-only``. +Newer versions of Pandas may fix these errors by improving support for such cases. +You can work around this error by copying the column(s) beforehand. +Additionally, this conversion may be slower because it is single-threaded. diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/index.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/index.rst new file mode 100644 index 0000000..c0369de --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/index.rst @@ -0,0 +1,28 @@ +.. 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. + + +========= +Spark SQL +========= + +.. toctree:: + :maxdepth: 2 + + arrow_pandas + python_udtf + diff --git a/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/python_udtf.rst b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/python_udtf.rst new file mode 100644 index 0000000..0e58391 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/docs/source/user_guide/sql/python_udtf.rst @@ -0,0 +1,233 @@ +.. 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. + +=========================================== +Python User-defined Table Functions (UDTFs) +=========================================== + +Spark 3.5 introduces the Python user-defined table function (UDTF), a new type of user-defined function. +Unlike scalar functions that return a single result value from each call, each UDTF is invoked in +the ``FROM`` clause of a query and returns an entire table as output. +Each UDTF call can accept zero or more arguments. +These arguments can either be scalar expressions or table arguments that represent entire input tables. + +Implementing a Python UDTF +-------------------------- + +.. currentmodule:: pyspark.sql.functions + +To implement a Python UDTF, you first need to define a class implementing the methods: + +.. code-block:: python + + class PythonUDTF: + + def __init__(self) -> None: + """ + Initializes the user-defined table function (UDTF). This is optional. + + This method serves as the default constructor and is called once when the + UDTF is instantiated on the executor side. + + Any class fields assigned in this method will be available for subsequent + calls to the `eval` and `terminate` methods. This class instance will remain + alive until all rows in the current partition have been consumed by the `eval` + method. + + Notes + ----- + - This method does not accept any extra arguments. Only the default + constructor is supported. + - You cannot create or reference the Spark session within the UDTF. Any + attempt to do so will result in a serialization error. + """ + ... + + def eval(self, *args: Any) -> Iterator[Any]: + """ + Evaluates the function using the given input arguments. + + This method is required and must be implemented. + + Argument Mapping: + - Each provided scalar expression maps to exactly one value in the + `*args` list. + - Each provided table argument maps to a pyspark.sql.Row object containing + the columns in the order they appear in the provided input table, + and with the names computed by the query analyzer. + + This method is called on every input row, and can produce zero or more + output rows. Each element in the output tuple corresponds to one column + specified in the return type of the UDTF. + + Parameters + ---------- + *args : Any + Arbitrary positional arguments representing the input to the UDTF. + + Yields + ------ + tuple + A tuple representing a single row in the UDTF result table. + Yield as many times as needed to produce multiple rows. + + Notes + ----- + - The result of the function must be a tuple representing a single row + in the UDTF result table. + - UDTFs currently do not accept keyword arguments during the function call. + + Examples + -------- + eval that returns one row and one column for each input. + + >>> def eval(self, x: int): + ... yield (x, ) + + eval that returns two rows and two columns for each input. + + >>> def eval(self, x: int, y: int): + ... yield (x + y, x - y) + ... yield (y + x, y - x) + """ + ... + + def terminate(self) -> Iterator[Any]: + """ + Called when the UDTF has processed all input rows. + + This method is optional to implement and is useful for performing any + cleanup or finalization operations after the UDTF has finished processing + all rows. It can also be used to yield additional rows if needed. + Table functions that consume all rows in the entire input partition + and then compute and return the entire output table can do so from + this method as well (please be mindful of memory usage when doing + this). + + Yields + ------ + tuple + A tuple representing a single row in the UDTF result table. + Yield this if you want to return additional rows during termination. + + Examples + -------- + >>> def terminate(self) -> Iterator[Any]: + >>> yield "done", None + """ + ... + + +The return type of the UDTF defines the schema of the table it outputs. +It must be either a ``StructType``, for example ``StructType().add("c1", StringType())`` +or a DDL string representing a struct type, for example ``c1: string``. + +**Example of UDTF Class Implementation** + +Here is a simple example of a UDTF class implementation: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 36-40 + :dedent: 4 + + +**Instantiating a UDTF with the ``udtf`` Decorator** + +To make use of the UDTF, you'll first need to instantiate it using the ``@udtf`` decorator: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 42-55 + :dedent: 4 + + +**Instantiating a UDTF with the ``udtf`` Function** + +An alternative way to create a UDTF is to use the :func:`udtf` function: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 60-77 + :dedent: 4 + +For more detailed usage, please see :func:`udtf`. + + +Registering and Using Python UDTFs in SQL +----------------------------------------- + +Python UDTFs can also be registered and used in SQL queries. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 82-116 + :dedent: 4 + + +Arrow Optimization +------------------ +Apache Arrow is an in-memory columnar data format used in Spark to efficiently transfer +data between Java and Python processes. Apache Arrow is disabled by default for Python UDTFs. + +Arrow can improve performance when each input row generates a large result table from the UDTF. + +To enable Arrow optimization, set the ``spark.sql.execution.pythonUDTF.arrow.enabled`` +configuration to ``true``. You can also enable it by specifying the ``useArrow`` parameter +when declaring the UDTF. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 121-126 + :dedent: 4 + + +For more details, please see `Apache Arrow in PySpark <../arrow_pandas.rst>`_. + + +TABLE input argument +~~~~~~~~~~~~~~~~~~~~ +Python UDTFs can also take a TABLE as input argument, and it can be used in conjunction +with scalar input arguments. +By default, you are allowed to have only one TABLE argument as input, primarily for +performance reasons. If you need to have more than one TABLE input argument, +you can enable this by setting the ``spark.sql.tvf.allowMultipleTableArguments.enabled`` +configuration to ``true``. + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 191-210 + :dedent: 4 + + +More Examples +------------- + +A Python UDTF that expands date ranges into individual dates: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 131-152 + :dedent: 4 + + +A Python UDTF with ``__init__`` and ``terminate``: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 157-186 + :dedent: 4 diff --git a/spark-3.5.0-bin-hadoop3/python/lib/PY4J_LICENSE.txt b/spark-3.5.0-bin-hadoop3/python/lib/PY4J_LICENSE.txt new file mode 100644 index 0000000..a70279c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/lib/PY4J_LICENSE.txt @@ -0,0 +1,27 @@ + +Copyright (c) 2009-2011, Barthelemy Dagenais All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +- Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +- Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +- The name of the author may not be used to endorse or promote products +derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/spark-3.5.0-bin-hadoop3/python/lib/py4j-0.10.9.7-src.zip b/spark-3.5.0-bin-hadoop3/python/lib/py4j-0.10.9.7-src.zip new file mode 100644 index 0000000..6abba4e Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/lib/py4j-0.10.9.7-src.zip differ diff --git a/spark-3.5.0-bin-hadoop3/python/lib/pyspark.zip b/spark-3.5.0-bin-hadoop3/python/lib/pyspark.zip new file mode 100644 index 0000000..9071c7e Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/lib/pyspark.zip differ diff --git a/spark-3.5.0-bin-hadoop3/python/mypy.ini b/spark-3.5.0-bin-hadoop3/python/mypy.ini new file mode 100644 index 0000000..3443af9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/mypy.ini @@ -0,0 +1,171 @@ +; +; 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. +; + +[mypy] +strict_optional = True +no_implicit_optional = True +disallow_untyped_defs = True +show_error_codes = True +warn_unused_ignores = True +warn_redundant_casts = True +namespace_packages = True + +[mypy-pyspark.sql.connect.proto.*] +ignore_errors = True + +; Allow untyped def in internal modules + +[mypy-pyspark.daemon] +disallow_untyped_defs = False + +[mypy-pyspark.find_spark_home] +disallow_untyped_defs = False + +[mypy-pyspark._globals] +disallow_untyped_defs = False + +[mypy-pyspark.install] +disallow_untyped_defs = False + +[mypy-pyspark.java_gateway] +disallow_untyped_defs = False + +[mypy-pyspark.join] +disallow_untyped_defs = False + +[mypy-pyspark.rddsampler] +disallow_untyped_defs = False + +[mypy-pyspark.serializers] +disallow_untyped_defs = False + +[mypy-pyspark.shuffle] +disallow_untyped_defs = False + +[mypy-pyspark.streaming.util] +disallow_untyped_defs = False + +[mypy-pyspark.sql.pandas.serializers] +disallow_untyped_defs = False + +[mypy-pyspark.sql.pandas._typing.protocols.*] +disallow_untyped_defs = False + +[mypy-pyspark.pandas.missing.*] +disallow_untyped_defs = False + +[mypy-pyspark.pandas.plot.*] +disallow_untyped_defs = False + +[mypy-pyspark.pandas.usage_logging.*] +disallow_untyped_defs = False + +[mypy-pyspark.traceback_utils] +disallow_untyped_defs = False + +[mypy-pyspark.worker] +disallow_untyped_defs = False + +; Allow untyped def and disable certain error codes in examples + +[mypy-python.sql.udtf] +disallow_untyped_defs = False +disable_error_code = attr-defined,arg-type,call-arg,union-attr + +; Ignore errors in tests + +[mypy-pyspark.ml.tests.*] +ignore_errors = True + +[mypy-pyspark.ml.torch.tests.*] +ignore_errors = True + +[mypy-pyspark.ml.deepspeed.tests.*] +ignore_errors = True + +[mypy-pyspark.mllib.tests.*] +ignore_errors = True + +[mypy-pyspark.resource.tests.*] +ignore_errors = True + +[mypy-pyspark.streaming.tests.*] +ignore_errors = True + +[mypy-pyspark.sql.tests.*] +ignore_errors = True + +[mypy-pyspark.pandas.tests.*] +ignore_errors = True + +[mypy-pyspark.tests.*] +ignore_errors = True + +[mypy-pyspark.testing.*] +ignore_errors = True + +[mypy-pyspark.errors.tests.*] +ignore_errors = True + +; Allow non-strict optional for pyspark.pandas + +[mypy-pyspark.pandas.*] +strict_optional = False + +; Ignore errors in embedded third party code + +[mypy-pyspark.cloudpickle.*] +ignore_errors = True + +; Ignore missing imports for external untyped packages + +[mypy-py4j.*] +ignore_missing_imports = True + +[mypy-scipy.*] +ignore_missing_imports = True + +[mypy-pyarrow.*] +ignore_missing_imports = True + +[mypy-psutil.*] +ignore_missing_imports = True + +[mypy-matplotlib.*] +ignore_missing_imports = True + +[mypy-plotly.*] +ignore_missing_imports = True + +[mypy-mlflow.*] +ignore_missing_imports = True + +[mypy-sklearn.*] +ignore_missing_imports = True + +[mypy-tabulate.*] +ignore_missing_imports = True + +[mypy-google.protobuf.*] +ignore_missing_imports = True + +[mypy-grpc.*] +ignore_missing_imports = True + +; Ignore errors for proto generated code +[mypy-pyspark.sql.connect.proto.*, pyspark.sql.connect.proto] +ignore_errors = True diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/__init__.py new file mode 100644 index 0000000..b8bca77 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/__init__.py @@ -0,0 +1,173 @@ +# +# 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. +# + +""" +PySpark is the Python API for Spark. + +Public classes: + + - :class:`SparkContext`: + Main entry point for Spark functionality. + - :class:`RDD`: + A Resilient Distributed Dataset (RDD), the basic abstraction in Spark. + - :class:`Broadcast`: + A broadcast variable that gets reused across tasks. + - :class:`Accumulator`: + An "add-only" shared variable that tasks can only add values to. + - :class:`SparkConf`: + For configuring Spark. + - :class:`SparkFiles`: + Access files shipped with jobs. + - :class:`StorageLevel`: + Finer-grained cache persistence levels. + - :class:`TaskContext`: + Information about the current running task, available on the workers and experimental. + - :class:`RDDBarrier`: + Wraps an RDD under a barrier stage for barrier execution. + - :class:`BarrierTaskContext`: + A :class:`TaskContext` that provides extra info and tooling for barrier execution. + - :class:`BarrierTaskInfo`: + Information about a barrier task. + - :class:`InheritableThread`: + A inheritable thread to use in Spark when the pinned thread mode is on. +""" + +from functools import wraps +import types +from typing import cast, Any, Callable, Optional, TypeVar, Union +from warnings import filterwarnings + +filterwarnings( + "ignore", message="distutils Version classes are deprecated. Use packaging.version instead." +) + +from pyspark.conf import SparkConf +from pyspark.rdd import RDD, RDDBarrier +from pyspark.files import SparkFiles +from pyspark.status import StatusTracker, SparkJobInfo, SparkStageInfo +from pyspark.util import InheritableThread, inheritable_thread_target +from pyspark.storagelevel import StorageLevel +from pyspark.accumulators import Accumulator, AccumulatorParam +from pyspark.broadcast import Broadcast +from pyspark.serializers import MarshalSerializer, CPickleSerializer +from pyspark.taskcontext import TaskContext, BarrierTaskContext, BarrierTaskInfo +from pyspark.profiler import Profiler, BasicProfiler +from pyspark.version import __version__ +from pyspark._globals import _NoValue # noqa: F401 + +_F = TypeVar("_F", bound=Callable) + + +def since(version: Union[str, float]) -> Callable[[_F], _F]: + """ + A decorator that annotates a function to append the version of Spark the function was added. + """ + import re + + indent_p = re.compile(r"\n( +)") + + def deco(f: _F) -> _F: + assert f.__doc__ is not None + + indents = indent_p.findall(f.__doc__) + indent = " " * (min(len(m) for m in indents) if indents else 0) + f.__doc__ = f.__doc__.rstrip() + "\n\n%s.. versionadded:: %s" % (indent, version) + return f + + return deco + + +def copy_func( + f: _F, + name: Optional[str] = None, + sinceversion: Optional[Union[str, float]] = None, + doc: Optional[str] = None, +) -> _F: + """ + Returns a function with same code, globals, defaults, closure, and + name (or provide a new name). + """ + # See + # http://stackoverflow.com/questions/6527633/how-can-i-make-a-deepcopy-of-a-function-in-python + assert isinstance(f, types.FunctionType) + + fn = types.FunctionType( + f.__code__, + f.__globals__, + name or f.__name__, + f.__defaults__, + f.__closure__, + ) + # in case f was given attrs (note this dict is a shallow copy): + fn.__dict__.update(f.__dict__) + if doc is not None: + fn.__doc__ = doc + if sinceversion is not None: + fn = since(sinceversion)(fn) + return cast(_F, fn) + + +def keyword_only(func: _F) -> _F: + """ + A decorator that forces keyword arguments in the wrapped method + and saves actual input keyword arguments in `_input_kwargs`. + + Notes + ----- + Should only be used to wrap a method where first arg is `self` + """ + + @wraps(func) + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if len(args) > 0: + raise TypeError("Method %s forces keyword arguments." % func.__name__) + self._input_kwargs = kwargs + return func(self, **kwargs) + + return cast(_F, wrapper) + + +# To avoid circular dependencies +from pyspark.context import SparkContext + +# for back compatibility +from pyspark.sql import SQLContext, HiveContext, Row # noqa: F401 + +__all__ = [ + "SparkConf", + "SparkContext", + "SparkFiles", + "RDD", + "StorageLevel", + "Broadcast", + "Accumulator", + "AccumulatorParam", + "MarshalSerializer", + "CPickleSerializer", + "StatusTracker", + "SparkJobInfo", + "SparkStageInfo", + "Profiler", + "BasicProfiler", + "TaskContext", + "RDDBarrier", + "BarrierTaskContext", + "BarrierTaskInfo", + "InheritableThread", + "inheritable_thread_target", + "__version__", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/_globals.py b/spark-3.5.0-bin-hadoop3/python/pyspark/_globals.py new file mode 100644 index 0000000..1b8c827 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/_globals.py @@ -0,0 +1,71 @@ +# +# 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. +# + +""" +Module defining global singleton classes. + +This module raises a RuntimeError if an attempt to reload it is made. In that +way the identities of the classes defined here are fixed and will remain so +even if pyspark itself is reloaded. In particular, a function like the following +will still work correctly after pyspark is reloaded: + + def foo(arg=pyspark._NoValue): + if arg is pyspark._NoValue: + ... + +See gh-7844 for a discussion of the reload problem that motivated this module. + +Note that this approach is taken after from NumPy. +""" + +__ALL__ = ["_NoValue"] + + +# Disallow reloading this module so as to preserve the identities of the +# classes defined here. +if "_is_loaded" in globals(): + raise RuntimeError("Reloading pyspark._globals is not allowed") +_is_loaded = True + + +class _NoValueType: + """Special keyword value. + + The instance of this class may be used as the default value assigned to a + deprecated keyword in order to check if it has been given a user defined + value. + + This class was copied from NumPy. + """ + + __instance = None + + def __new__(cls): + # ensure that only one instance exists + if not cls.__instance: + cls.__instance = super(_NoValueType, cls).__new__(cls) + return cls.__instance + + # needed for python 2 to preserve identity through a pickle + def __reduce__(self): + return (self.__class__, ()) + + def __repr__(self): + return "" + + +_NoValue = _NoValueType() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/_typing.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/_typing.pyi new file mode 100644 index 0000000..6cc0926 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/_typing.pyi @@ -0,0 +1,43 @@ +# +# 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 Callable, Iterable, Sized, TypeVar, Union +from typing_extensions import Literal, Protocol + +from numpy import int32, int64, float32, float64, ndarray + +F = TypeVar("F", bound=Callable) +T_co = TypeVar("T_co", covariant=True) + +PrimitiveType = Union[bool, float, int, str] + +NonUDFType = Literal[0] + +class SupportsIAdd(Protocol): + def __iadd__(self, other: SupportsIAdd) -> SupportsIAdd: ... + +class SupportsOrdering(Protocol): + def __lt__(self, other: SupportsOrdering) -> bool: ... + +class SizedIterable(Protocol, Sized, Iterable[T_co]): ... + +S = TypeVar("S", bound=SupportsOrdering) + +NumberOrArray = TypeVar( + "NumberOrArray", float, int, complex, int32, int64, float32, float64, ndarray +) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/accumulators.py b/spark-3.5.0-bin-hadoop3/python/pyspark/accumulators.py new file mode 100644 index 0000000..a95bd9d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/accumulators.py @@ -0,0 +1,341 @@ +# +# 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 sys +import select +import struct +import socketserver as SocketServer +import threading +from typing import Callable, Dict, Generic, Tuple, Type, TYPE_CHECKING, TypeVar, Union + +from pyspark.serializers import read_int, CPickleSerializer +from pyspark.errors import PySparkRuntimeError + +if TYPE_CHECKING: + from pyspark._typing import SupportsIAdd # noqa: F401 + import socketserver.BaseRequestHandler # type: ignore[import] + + +__all__ = ["Accumulator", "AccumulatorParam"] + +T = TypeVar("T") +U = TypeVar("U", bound="SupportsIAdd") + +pickleSer = CPickleSerializer() + +# 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"] = {} + + +def _deserialize_accumulator( + aid: int, zero_value: T, accum_param: "AccumulatorParam[T]" +) -> "Accumulator[T]": + from pyspark.accumulators import _accumulatorRegistry + + # If this certain accumulator was deserialized, don't overwrite it. + if aid in _accumulatorRegistry: + return _accumulatorRegistry[aid] + else: + accum = Accumulator(aid, zero_value, accum_param) + accum._deserialized = True + _accumulatorRegistry[aid] = accum + return accum + + +class Accumulator(Generic[T]): + + """ + A shared variable that can be accumulated, i.e., has a commutative and associative "add" + operation. Worker tasks on a Spark cluster can add values to an Accumulator with the `+=` + operator, but only the driver program is allowed to access its value, using `value`. + Updates from the workers get propagated automatically to the driver program. + + While :class:`SparkContext` supports accumulators for primitive data types like :class:`int` and + :class:`float`, users can also define accumulators for custom types by providing a custom + :py:class:`AccumulatorParam` object. Refer to its doctest for an example. + + Examples + -------- + >>> a = sc.accumulator(1) + >>> a.value + 1 + >>> a.value = 2 + >>> a.value + 2 + >>> a += 5 + >>> a.value + 7 + >>> sc.accumulator(1.0).value + 1.0 + >>> sc.accumulator(1j).value + 1j + >>> rdd = sc.parallelize([1,2,3]) + >>> def f(x): + ... global a + ... a += x + ... + >>> rdd.foreach(f) + >>> a.value + 13 + >>> b = sc.accumulator(0) + >>> def g(x): + ... b.add(x) + ... + >>> rdd.foreach(g) + >>> b.value + 6 + + >>> rdd.map(lambda x: a.value).collect() # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + Py4JJavaError: ... + + >>> def h(x): + ... global a + ... a.value = 7 + ... + >>> rdd.foreach(h) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + Py4JJavaError: ... + + >>> sc.accumulator([1.0, 2.0, 3.0]) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TypeError: ... + """ + + def __init__(self, aid: int, value: T, accum_param: "AccumulatorParam[T]"): + """Create a new Accumulator with a given initial value and AccumulatorParam object""" + from pyspark.accumulators import _accumulatorRegistry + + self.aid = aid + self.accum_param = accum_param + self._value = value + self._deserialized = False + _accumulatorRegistry[aid] = self + + def __reduce__( + self, + ) -> Tuple[ + Callable[[int, T, "AccumulatorParam[T]"], "Accumulator[T]"], + Tuple[int, T, "AccumulatorParam[T]"], + ]: + """Custom serialization; saves the zero value from our AccumulatorParam""" + param = self.accum_param + return (_deserialize_accumulator, (self.aid, param.zero(self._value), param)) + + @property + def value(self) -> T: + """Get the accumulator's value; only usable in driver program""" + if self._deserialized: + raise PySparkRuntimeError( + error_class="VALUE_NOT_ACCESSIBLE", + message_parameters={ + "value": "Accumulator.value", + }, + ) + return self._value + + @value.setter + def value(self, value: T) -> None: + """Sets the accumulator's value; only usable in driver program""" + if self._deserialized: + raise PySparkRuntimeError( + error_class="VALUE_NOT_ACCESSIBLE", + message_parameters={ + "value": "Accumulator.value", + }, + ) + self._value = value + + def add(self, term: T) -> None: + """Adds a term to this accumulator's value""" + self._value = self.accum_param.addInPlace(self._value, term) + + def __iadd__(self, term: T) -> "Accumulator[T]": + """The += operator; adds a term to this accumulator's value""" + self.add(term) + return self + + def __str__(self) -> str: + return str(self._value) + + def __repr__(self) -> str: + return "Accumulator" % (self.aid, self._value) + + +class AccumulatorParam(Generic[T]): + + """ + Helper object that defines how to accumulate values of a given type. + + Examples + -------- + >>> from pyspark.accumulators import AccumulatorParam + >>> class VectorAccumulatorParam(AccumulatorParam): + ... def zero(self, value): + ... return [0.0] * len(value) + ... def addInPlace(self, val1, val2): + ... for i in range(len(val1)): + ... val1[i] += val2[i] + ... return val1 + >>> va = sc.accumulator([1.0, 2.0, 3.0], VectorAccumulatorParam()) + >>> va.value + [1.0, 2.0, 3.0] + >>> def g(x): + ... global va + ... va += [x] * 3 + ... + >>> rdd = sc.parallelize([1,2,3]) + >>> rdd.foreach(g) + >>> va.value + [7.0, 8.0, 9.0] + """ + + def zero(self, value: T) -> T: + """ + Provide a "zero value" for the type, compatible in dimensions with the + provided `value` (e.g., a zero vector) + """ + raise NotImplementedError + + def addInPlace(self, value1: T, value2: T) -> T: + """ + Add two values of the accumulator's data type, returning a new value; + for efficiency, can also update `value1` in place and return it. + """ + raise NotImplementedError + + +class AddingAccumulatorParam(AccumulatorParam[U]): + + """ + An AccumulatorParam that uses the + operators to add values. Designed for simple types + such as integers, floats, and lists. Requires the zero value for the underlying type + as a parameter. + """ + + def __init__(self, zero_value: U): + self.zero_value = zero_value + + def zero(self, value: U) -> U: + return self.zero_value + + def addInPlace(self, value1: U, value2: U) -> U: + value1 += value2 # type: ignore[operator] + return value1 + + +# Singleton accumulator params for some standard types +INT_ACCUMULATOR_PARAM = AddingAccumulatorParam(0) # type: ignore[type-var] +FLOAT_ACCUMULATOR_PARAM = AddingAccumulatorParam(0.0) # type: ignore[type-var] +COMPLEX_ACCUMULATOR_PARAM = AddingAccumulatorParam(0.0j) # type: ignore[type-var] + + +class _UpdateRequestHandler(SocketServer.StreamRequestHandler): + + """ + This handler will keep polling updates from the same socket until the + server is shutdown. + """ + + def handle(self) -> None: + from pyspark.accumulators import _accumulatorRegistry + + auth_token = self.server.auth_token # type: ignore[attr-defined] + + def poll(func: Callable[[], bool]) -> None: + while not self.server.server_shutdown: # type: ignore[attr-defined] + # Poll every 1 second for new data -- don't block in case of shutdown. + r, _, _ = select.select([self.rfile], [], [], 1) + if self.rfile in r and func(): + break + + def accum_updates() -> bool: + num_updates = read_int(self.rfile) + for _ in range(num_updates): + (aid, update) = pickleSer._read_with_length(self.rfile) + _accumulatorRegistry[aid] += update + # Write a byte in acknowledgement + self.wfile.write(struct.pack("!b", 1)) + return False + + def authenticate_and_accum_updates() -> bool: + received_token: Union[bytes, str] = self.rfile.read(len(auth_token)) + if isinstance(received_token, bytes): + received_token = received_token.decode("utf-8") + if received_token == auth_token: + accum_updates() + # we've authenticated, we can break out of the first loop now + return True + else: + raise ValueError( + "The value of the provided token to the AccumulatorServer is not correct." + ) + + # first we keep polling till we've received the authentication token + poll(authenticate_and_accum_updates) + # now we've authenticated, don't need to check for the token anymore + poll(accum_updates) + + +class AccumulatorServer(SocketServer.TCPServer): + def __init__( + self, + server_address: Tuple[str, int], + RequestHandlerClass: Type["socketserver.BaseRequestHandler"], + auth_token: str, + ): + SocketServer.TCPServer.__init__(self, server_address, RequestHandlerClass) + self.auth_token = auth_token + + """ + A simple TCP server that intercepts shutdown() in order to interrupt + our continuous polling on the handler. + """ + server_shutdown = False + + def shutdown(self) -> None: + self.server_shutdown = True + SocketServer.TCPServer.shutdown(self) + self.server_close() + + +def _start_update_server(auth_token: str) -> AccumulatorServer: + """Start a TCP server to receive accumulator updates in a daemon thread, and returns it""" + server = AccumulatorServer(("localhost", 0), _UpdateRequestHandler, auth_token) + thread = threading.Thread(target=server.serve_forever) + thread.daemon = True + thread.start() + return server + + +if __name__ == "__main__": + import doctest + + from pyspark.context import SparkContext + + globs = globals().copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + globs["sc"] = SparkContext("local", "test") + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/broadcast.py b/spark-3.5.0-bin-hadoop3/python/pyspark/broadcast.py new file mode 100644 index 0000000..1f2b326 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/broadcast.py @@ -0,0 +1,383 @@ +# +# 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 gc +import os +import sys +from tempfile import NamedTemporaryFile +import threading +import pickle +from typing import ( + overload, + Any, + BinaryIO, + Callable, + Dict, + Generic, + IO, + Iterator, + Optional, + Tuple, + TypeVar, + TYPE_CHECKING, + Union, +) + +from pyspark.java_gateway import local_connect_and_auth +from pyspark.serializers import ChunkedStream, pickle_protocol +from pyspark.util import print_exec +from pyspark.errors import PySparkRuntimeError + +if TYPE_CHECKING: + from pyspark import SparkContext + + +__all__ = ["Broadcast"] + +T = TypeVar("T") + + +# Holds broadcasted data received from Java, keyed by its id. +_broadcastRegistry: Dict[int, "Broadcast[Any]"] = {} + + +def _from_id(bid: int) -> "Broadcast[Any]": + from pyspark.broadcast import _broadcastRegistry + + if bid not in _broadcastRegistry: + raise PySparkRuntimeError( + error_class="BROADCAST_VARIABLE_NOT_LOADED", + message_parameters={ + "variable": str(bid), + }, + ) + return _broadcastRegistry[bid] + + +class Broadcast(Generic[T]): + + """ + A broadcast variable created with :meth:`SparkContext.broadcast`. + Access its value through :attr:`value`. + + Examples + -------- + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + >>> b.value + [1, 2, 3, 4, 5] + >>> spark.sparkContext.parallelize([0, 0]).flatMap(lambda x: b.value).collect() + [1, 2, 3, 4, 5, 1, 2, 3, 4, 5] + >>> b.unpersist() + + >>> large_broadcast = spark.sparkContext.broadcast(range(10000)) + """ + + @overload # On driver + def __init__( + self: "Broadcast[T]", + sc: "SparkContext", + value: T, + pickle_registry: "BroadcastPickleRegistry", + ): + ... + + @overload # On worker without decryption server + def __init__(self: "Broadcast[Any]", *, path: str): + ... + + @overload # On worker with decryption server + def __init__(self: "Broadcast[Any]", *, sock_file: str): + ... + + def __init__( # type: ignore[misc] + self, + sc: Optional["SparkContext"] = None, + value: Optional[T] = None, + pickle_registry: Optional["BroadcastPickleRegistry"] = None, + path: Optional[str] = None, + sock_file: Optional[BinaryIO] = None, + ): + """ + Should not be called directly by users -- use :meth:`SparkContext.broadcast` + instead. + """ + if sc is not None: + # we're on the driver. We want the pickled data to end up in a file (maybe encrypted) + f = NamedTemporaryFile(delete=False, dir=sc._temp_dir) + self._path = f.name + self._sc: Optional["SparkContext"] = sc + assert sc._jvm is not None + self._python_broadcast = sc._jvm.PythonRDD.setupBroadcast(self._path) + broadcast_out: Union[ChunkedStream, IO[bytes]] + if sc._encryption_enabled: + # with encryption, we ask the jvm to do the encryption for us, we send it data + # over a socket + port, auth_secret = self._python_broadcast.setupEncryptionServer() + (encryption_sock_file, _) = local_connect_and_auth(port, auth_secret) + broadcast_out = ChunkedStream(encryption_sock_file, 8192) + else: + # no encryption, we can just write pickled data directly to the file from python + broadcast_out = f + self.dump(value, broadcast_out) # type: ignore[arg-type] + if sc._encryption_enabled: + self._python_broadcast.waitTillDataReceived() + self._jbroadcast = sc._jsc.broadcast(self._python_broadcast) + self._pickle_registry = pickle_registry + else: + # we're on an executor + self._jbroadcast = None + self._sc = None + self._python_broadcast = None + if sock_file is not None: + # the jvm is doing decryption for us. Read the value + # immediately from the sock_file + self._value = self.load(sock_file) + else: + # the jvm just dumps the pickled data in path -- we'll unpickle lazily when + # the value is requested + assert path is not None + self._path = path + + def dump(self, value: T, f: BinaryIO) -> None: + """ + Write a pickled representation of value to the open file or socket. + The protocol pickle is HIGHEST_PROTOCOL. + + Parameters + ---------- + value : T + Value to write. + + f : :class:`BinaryIO` + File or socket where the pickled value will be stored. + + Examples + -------- + >>> import os + >>> import tempfile + + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + + Write a pickled representation of `b` to the open temp file. + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "test.txt") + ... with open(path, "wb") as f: + ... b.dump(b.value, f) + """ + try: + pickle.dump(value, f, pickle_protocol) + except pickle.PickleError: + raise + except Exception as e: + msg = "Could not serialize broadcast: %s: %s" % (e.__class__.__name__, str(e)) + print_exec(sys.stderr) + raise pickle.PicklingError(msg) + f.close() + + def load_from_path(self, path: str) -> T: + """ + Read the pickled representation of an object from the open file and + return the reconstituted object hierarchy specified therein. + + Parameters + ---------- + path : str + File path where reads the pickled value. + + Returns + ------- + T + The object hierarchy specified therein reconstituted + from the pickled representation of an object. + + Examples + -------- + >>> import os + >>> import tempfile + + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + >>> c = spark.sparkContext.broadcast(1) + + Read the pickled representation of value from temp file. + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "test.txt") + ... with open(path, "wb") as f: + ... b.dump(b.value, f) + ... c.load_from_path(path) + [1, 2, 3, 4, 5] + """ + with open(path, "rb", 1 << 20) as f: + return self.load(f) + + def load(self, file: BinaryIO) -> T: + """ + Read a pickled representation of value from the open file or socket. + + Parameters + ---------- + file : :class:`BinaryIO` + File or socket where the pickled value will be read. + + Returns + ------- + T + The object hierarchy specified therein reconstituted + from the pickled representation of an object. + + Examples + -------- + >>> import os + >>> import tempfile + + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + >>> c = spark.sparkContext.broadcast(1) + + Read the pickled representation of value from the open temp file. + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "test.txt") + ... with open(path, "wb") as f: + ... b.dump(b.value, f) + ... with open(path, "rb") as f: + ... c.load(f) + [1, 2, 3, 4, 5] + """ + gc.disable() + try: + return pickle.load(file) + finally: + gc.enable() + + @property + def value(self) -> T: + """Return the broadcasted value""" + if not hasattr(self, "_value") and self._path is not None: + # we only need to decrypt it here when encryption is enabled and + # if its on the driver, since executor decryption is handled already + if self._sc is not None and self._sc._encryption_enabled: + port, auth_secret = self._python_broadcast.setupDecryptionServer() + (decrypted_sock_file, _) = local_connect_and_auth(port, auth_secret) + self._python_broadcast.waitTillBroadcastDataSent() + return self.load(decrypted_sock_file) + else: + self._value = self.load_from_path(self._path) + return self._value + + def unpersist(self, blocking: bool = False) -> None: + """ + Delete cached copies of this broadcast on the executors. If the + broadcast is used after this is called, it will need to be + re-sent to each executor. + + Parameters + ---------- + blocking : bool, optional, default False + Whether to block until unpersisting has completed. + + Examples + -------- + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + + Delete cached copies of this broadcast on the executors + + >>> b.unpersist() + """ + if self._jbroadcast is None: + raise PySparkRuntimeError( + error_class="INVALID_BROADCAST_OPERATION", + message_parameters={"operation": "unpersisted"}, + ) + self._jbroadcast.unpersist(blocking) + + def destroy(self, blocking: bool = False) -> None: + """ + Destroy all data and metadata related to this broadcast variable. + Use this with caution; once a broadcast variable has been destroyed, + it cannot be used again. + + .. versionchanged:: 3.0.0 + Added optional argument `blocking` to specify whether to block until all + blocks are deleted. + + Parameters + ---------- + blocking : bool, optional, default False + Whether to block until unpersisting has completed. + + Examples + -------- + >>> b = spark.sparkContext.broadcast([1, 2, 3, 4, 5]) + + Destroy all data and metadata related to this broadcast variable + + >>> b.destroy() + """ + if self._jbroadcast is None: + raise PySparkRuntimeError( + error_class="INVALID_BROADCAST_OPERATION", + message_parameters={"operation": "destroyed"}, + ) + self._jbroadcast.destroy(blocking) + os.unlink(self._path) + + def __reduce__(self) -> Tuple[Callable[[int], "Broadcast[T]"], Tuple[int]]: + if self._jbroadcast is None: + raise PySparkRuntimeError( + error_class="INVALID_BROADCAST_OPERATION", + message_parameters={"operation": "serialized"}, + ) + assert self._pickle_registry is not None + self._pickle_registry.add(self) + return _from_id, (self._jbroadcast.id(),) + + +class BroadcastPickleRegistry(threading.local): + """Thread-local registry for broadcast variables that have been pickled""" + + def __init__(self) -> None: + self.__dict__.setdefault("_registry", set()) + + def __iter__(self) -> Iterator[Broadcast[Any]]: + for bcast in self._registry: + yield bcast + + def add(self, bcast: Broadcast[Any]) -> None: + self._registry.add(bcast) + + def clear(self) -> None: + self._registry.clear() + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.broadcast + + globs = pyspark.broadcast.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("broadcast tests").getOrCreate() + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod(pyspark.broadcast, globs=globs) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/__init__.py new file mode 100644 index 0000000..af35a0a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/__init__.py @@ -0,0 +1,8 @@ +from pyspark.cloudpickle.cloudpickle import * # noqa +from pyspark.cloudpickle.cloudpickle_fast import CloudPickler, dumps, dump # noqa + +# Conform to the convention used by python serialization libraries, which +# expose their Pickler subclass at top-level under the "Pickler" name. +Pickler = CloudPickler + +__version__ = '2.2.1' diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle.py b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle.py new file mode 100644 index 0000000..317be69 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle.py @@ -0,0 +1,948 @@ +""" +This class is defined to override standard pickle functionality + +The goals of it follow: +-Serialize lambdas and nested functions to compiled byte code +-Deal with main module correctly +-Deal with other non-serializable objects + +It does not include an unpickler, as standard python unpickling suffices. + +This module was extracted from the `cloud` package, developed by `PiCloud, Inc. +`_. + +Copyright (c) 2012, Regents of the University of California. +Copyright (c) 2009 `PiCloud, Inc. `_. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the University of California, Berkeley nor the + names of its contributors may be used to endorse or promote + products derived from this software without specific prior written + permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +""" + +import builtins +import dis +import opcode +import platform +import sys +import types +import weakref +import uuid +import threading +import typing +import warnings + +from .compat import pickle +from collections import OrderedDict +from typing import ClassVar, Generic, Union, Tuple, Callable +from pickle import _getattribute +from importlib._bootstrap import _find_spec + +try: # pragma: no branch + import typing_extensions as _typing_extensions + from typing_extensions import Literal, Final +except ImportError: + _typing_extensions = Literal = Final = None + +if sys.version_info >= (3, 8): + from types import CellType +else: + def f(): + a = 1 + + def g(): + return a + return g + CellType = type(f().__closure__[0]) + + +# cloudpickle is meant for inter process communication: we expect all +# communicating processes to run the same Python version hence we favor +# communication speed over compatibility: +DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL + +# Names of modules whose resources should be treated as dynamic. +_PICKLE_BY_VALUE_MODULES = set() + +# Track the provenance of reconstructed dynamic classes to make it possible to +# reconstruct instances from the matching singleton class definition when +# appropriate and preserve the usual "isinstance" semantics of Python objects. +_DYNAMIC_CLASS_TRACKER_BY_CLASS = weakref.WeakKeyDictionary() +_DYNAMIC_CLASS_TRACKER_BY_ID = weakref.WeakValueDictionary() +_DYNAMIC_CLASS_TRACKER_LOCK = threading.Lock() + +PYPY = platform.python_implementation() == "PyPy" + +builtin_code_type = None +if PYPY: + # builtin-code objects only exist in pypy + builtin_code_type = type(float.__new__.__code__) + +_extract_code_globals_cache = weakref.WeakKeyDictionary() + + +def _get_or_create_tracker_id(class_def): + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_tracker_id = _DYNAMIC_CLASS_TRACKER_BY_CLASS.get(class_def) + if class_tracker_id is None: + class_tracker_id = uuid.uuid4().hex + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + _DYNAMIC_CLASS_TRACKER_BY_ID[class_tracker_id] = class_def + return class_tracker_id + + +def _lookup_class_or_track(class_tracker_id, class_def): + if class_tracker_id is not None: + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_def = _DYNAMIC_CLASS_TRACKER_BY_ID.setdefault( + class_tracker_id, class_def) + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + return class_def + + +def register_pickle_by_value(module): + """Register a module to make it functions and classes picklable by value. + + By default, functions and classes that are attributes of an importable + module are to be pickled by reference, that is relying on re-importing + the attribute from the module at load time. + + If `register_pickle_by_value(module)` is called, all its functions and + classes are subsequently to be pickled by value, meaning that they can + be loaded in Python processes where the module is not importable. + + This is especially useful when developing a module in a distributed + execution environment: restarting the client Python process with the new + source code is enough: there is no need to re-install the new version + of the module on all the worker nodes nor to restart the workers. + + Note: this feature is considered experimental. See the cloudpickle + README.md file for more details and limitations. + """ + if not isinstance(module, types.ModuleType): + raise ValueError( + f"Input should be a module object, got {str(module)} instead" + ) + # In the future, cloudpickle may need a way to access any module registered + # for pickling by value in order to introspect relative imports inside + # functions pickled by value. (see + # https://github.com/cloudpipe/cloudpickle/pull/417#issuecomment-873684633). + # This access can be ensured by checking that module is present in + # sys.modules at registering time and assuming that it will still be in + # there when accessed during pickling. Another alternative would be to + # store a weakref to the module. Even though cloudpickle does not implement + # this introspection yet, in order to avoid a possible breaking change + # later, we still enforce the presence of module inside sys.modules. + if module.__name__ not in sys.modules: + raise ValueError( + f"{module} was not imported correctly, have you used an " + f"`import` statement to access it?" + ) + _PICKLE_BY_VALUE_MODULES.add(module.__name__) + + +def unregister_pickle_by_value(module): + """Unregister that the input module should be pickled by value.""" + if not isinstance(module, types.ModuleType): + raise ValueError( + f"Input should be a module object, got {str(module)} instead" + ) + if module.__name__ not in _PICKLE_BY_VALUE_MODULES: + raise ValueError(f"{module} is not registered for pickle by value") + else: + _PICKLE_BY_VALUE_MODULES.remove(module.__name__) + + +def list_registry_pickle_by_value(): + return _PICKLE_BY_VALUE_MODULES.copy() + + +def _is_registered_pickle_by_value(module): + module_name = module.__name__ + if module_name in _PICKLE_BY_VALUE_MODULES: + return True + while True: + parent_name = module_name.rsplit(".", 1)[0] + if parent_name == module_name: + break + if parent_name in _PICKLE_BY_VALUE_MODULES: + return True + module_name = parent_name + return False + + +def _whichmodule(obj, name): + """Find the module an object belongs to. + + This function differs from ``pickle.whichmodule`` in two ways: + - it does not mangle the cases where obj's module is __main__ and obj was + not found in any module. + - Errors arising during module introspection are ignored, as those errors + are considered unwanted side effects. + """ + if sys.version_info[:2] < (3, 7) and isinstance(obj, typing.TypeVar): # pragma: no branch # noqa + # Workaround bug in old Python versions: prior to Python 3.7, + # T.__module__ would always be set to "typing" even when the TypeVar T + # would be defined in a different module. + if name is not None and getattr(typing, name, None) is obj: + # Built-in TypeVar defined in typing such as AnyStr + return 'typing' + else: + # User defined or third-party TypeVar: __module__ attribute is + # irrelevant, thus trigger a exhaustive search for obj in all + # modules. + module_name = None + else: + module_name = getattr(obj, '__module__', None) + + if module_name is not None: + return module_name + # Protect the iteration by using a copy of sys.modules against dynamic + # modules that trigger imports of other modules upon calls to getattr or + # other threads importing at the same time. + for module_name, module in sys.modules.copy().items(): + # Some modules such as coverage can inject non-module objects inside + # sys.modules + if ( + module_name == '__main__' or + module is None or + not isinstance(module, types.ModuleType) + ): + continue + try: + if _getattribute(module, name)[0] is obj: + return module_name + except Exception: + pass + return None + + +def _should_pickle_by_reference(obj, name=None): + """Test whether an function or a class should be pickled by reference + + Pickling by reference means by that the object (typically a function or a + class) is an attribute of a module that is assumed to be importable in the + target Python environment. Loading will therefore rely on importing the + module and then calling `getattr` on it to access the function or class. + + Pickling by reference is the only option to pickle functions and classes + in the standard library. In cloudpickle the alternative option is to + pickle by value (for instance for interactively or locally defined + functions and classes or for attributes of modules that have been + explicitly registered to be pickled by value. + """ + if isinstance(obj, types.FunctionType) or issubclass(type(obj), type): + module_and_name = _lookup_module_and_qualname(obj, name=name) + if module_and_name is None: + return False + module, name = module_and_name + return not _is_registered_pickle_by_value(module) + + elif isinstance(obj, types.ModuleType): + # We assume that sys.modules is primarily used as a cache mechanism for + # the Python import machinery. Checking if a module has been added in + # is sys.modules therefore a cheap and simple heuristic to tell us + # whether we can assume that a given module could be imported by name + # in another Python process. + if _is_registered_pickle_by_value(obj): + return False + return obj.__name__ in sys.modules + else: + raise TypeError( + "cannot check importability of {} instances".format( + type(obj).__name__) + ) + + +def _lookup_module_and_qualname(obj, name=None): + if name is None: + name = getattr(obj, '__qualname__', None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, '__name__', None) + + module_name = _whichmodule(obj, name) + + if module_name is None: + # In this case, obj.__module__ is None AND obj was not found in any + # imported module. obj is thus treated as dynamic. + return None + + if module_name == "__main__": + return None + + # Note: if module_name is in sys.modules, the corresponding module is + # assumed importable at unpickling time. See #357 + module = sys.modules.get(module_name, None) + if module is None: + # The main reason why obj's module would not be imported is that this + # module has been dynamically created, using for example + # types.ModuleType. The other possibility is that module was removed + # from sys.modules after obj was created/imported. But this case is not + # supported, as the standard pickle does not support it either. + return None + + try: + obj2, parent = _getattribute(module, name) + except AttributeError: + # obj was not found inside the module it points to + return None + if obj2 is not obj: + return None + return module, name + + +def _extract_code_globals(co): + """ + Find all globals names read or written to by codeblock co + """ + out_names = _extract_code_globals_cache.get(co) + if out_names is None: + # We use a dict with None values instead of a set to get a + # deterministic order (assuming Python 3.6+) and avoid introducing + # non-deterministic pickle bytes as a results. + out_names = {name: None for name in _walk_global_ops(co)} + + # Declaring a function inside another one using the "def ..." + # syntax generates a constant code object corresponding to the one + # of the nested function's As the nested function may itself need + # global variables, we need to introspect its code, extract its + # globals, (look for code object in it's co_consts attribute..) and + # add the result to code_globals + if co.co_consts: + for const in co.co_consts: + if isinstance(const, types.CodeType): + out_names.update(_extract_code_globals(const)) + + _extract_code_globals_cache[co] = out_names + + return out_names + + +def _find_imported_submodules(code, top_level_dependencies): + """ + Find currently imported submodules used by a function. + + Submodules used by a function need to be detected and referenced for the + function to work correctly at depickling time. Because submodules can be + referenced as attribute of their parent package (``package.submodule``), we + need a special introspection technique that does not rely on GLOBAL-related + opcodes to find references of them in a code object. + + Example: + ``` + import concurrent.futures + import cloudpickle + def func(): + x = concurrent.futures.ThreadPoolExecutor + if __name__ == '__main__': + cloudpickle.dumps(func) + ``` + The globals extracted by cloudpickle in the function's state include the + concurrent package, but not its submodule (here, concurrent.futures), which + is the module used by func. Find_imported_submodules will detect the usage + of concurrent.futures. Saving this module alongside with func will ensure + that calling func once depickled does not fail due to concurrent.futures + not being imported + """ + + subimports = [] + # check if any known dependency is an imported package + for x in top_level_dependencies: + if (isinstance(x, types.ModuleType) and + hasattr(x, '__package__') and x.__package__): + # check if the package has any currently loaded sub-imports + prefix = x.__name__ + '.' + # A concurrent thread could mutate sys.modules, + # make sure we iterate over a copy to avoid exceptions + for name in list(sys.modules): + # Older versions of pytest will add a "None" module to + # sys.modules. + if name is not None and name.startswith(prefix): + # check whether the function can address the sub-module + tokens = set(name[len(prefix):].split('.')) + if not tokens - set(code.co_names): + subimports.append(sys.modules[name]) + return subimports + + +def cell_set(cell, value): + """Set the value of a closure cell. + + The point of this function is to set the cell_contents attribute of a cell + after its creation. This operation is necessary in case the cell contains a + reference to the function the cell belongs to, as when calling the + function's constructor + ``f = types.FunctionType(code, globals, name, argdefs, closure)``, + closure will not be able to contain the yet-to-be-created f. + + In Python3.7, cell_contents is writeable, so setting the contents of a cell + can be done simply using + >>> cell.cell_contents = value + + In earlier Python3 versions, the cell_contents attribute of a cell is read + only, but this limitation can be worked around by leveraging the Python 3 + ``nonlocal`` keyword. + + In Python2 however, this attribute is read only, and there is no + ``nonlocal`` keyword. For this reason, we need to come up with more + complicated hacks to set this attribute. + + The chosen approach is to create a function with a STORE_DEREF opcode, + which sets the content of a closure variable. Typically: + + >>> def inner(value): + ... lambda: cell # the lambda makes cell a closure + ... cell = value # cell is a closure, so this triggers a STORE_DEREF + + (Note that in Python2, A STORE_DEREF can never be triggered from an inner + function. The function g for example here + >>> def f(var): + ... def g(): + ... var += 1 + ... return g + + will not modify the closure variable ``var```inplace, but instead try to + load a local variable var and increment it. As g does not assign the local + variable ``var`` any initial value, calling f(1)() will fail at runtime.) + + Our objective is to set the value of a given cell ``cell``. So we need to + somewhat reference our ``cell`` object into the ``inner`` function so that + this object (and not the smoke cell of the lambda function) gets affected + by the STORE_DEREF operation. + + In inner, ``cell`` is referenced as a cell variable (an enclosing variable + that is referenced by the inner function). If we create a new function + cell_set with the exact same code as ``inner``, but with ``cell`` marked as + a free variable instead, the STORE_DEREF will be applied on its closure - + ``cell``, which we can specify explicitly during construction! The new + cell_set variable thus actually sets the contents of a specified cell! + + Note: we do not make use of the ``nonlocal`` keyword to set the contents of + a cell in early python3 versions to limit possible syntax errors in case + test and checker libraries decide to parse the whole file. + """ + + if sys.version_info[:2] >= (3, 7): # pragma: no branch + cell.cell_contents = value + else: + _cell_set = types.FunctionType( + _cell_set_template_code, {}, '_cell_set', (), (cell,),) + _cell_set(value) + + +def _make_cell_set_template_code(): + def _cell_set_factory(value): + lambda: cell + cell = value + + co = _cell_set_factory.__code__ + + _cell_set_template_code = types.CodeType( + co.co_argcount, + co.co_kwonlyargcount, # Python 3 only argument + co.co_nlocals, + co.co_stacksize, + co.co_flags, + co.co_code, + co.co_consts, + co.co_names, + co.co_varnames, + co.co_filename, + co.co_name, + co.co_firstlineno, + co.co_lnotab, + co.co_cellvars, # co_freevars is initialized with co_cellvars + (), # co_cellvars is made empty + ) + return _cell_set_template_code + + +if sys.version_info[:2] < (3, 7): + _cell_set_template_code = _make_cell_set_template_code() + +# relevant opcodes +STORE_GLOBAL = opcode.opmap['STORE_GLOBAL'] +DELETE_GLOBAL = opcode.opmap['DELETE_GLOBAL'] +LOAD_GLOBAL = opcode.opmap['LOAD_GLOBAL'] +GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL) +HAVE_ARGUMENT = dis.HAVE_ARGUMENT +EXTENDED_ARG = dis.EXTENDED_ARG + + +_BUILTIN_TYPE_NAMES = {} +for k, v in types.__dict__.items(): + if type(v) is type: + _BUILTIN_TYPE_NAMES[v] = k + + +def _builtin_type(name): + if name == "ClassType": # pragma: no cover + # Backward compat to load pickle files generated with cloudpickle + # < 1.3 even if loading pickle files from older versions is not + # officially supported. + return type + return getattr(types, name) + + +def _walk_global_ops(code): + """ + Yield referenced name for all global-referencing instructions in *code*. + """ + for instr in dis.get_instructions(code): + op = instr.opcode + if op in GLOBAL_OPS: + yield instr.argval + + +def _extract_class_dict(cls): + """Retrieve a copy of the dict of a class without the inherited methods""" + clsdict = dict(cls.__dict__) # copy dict proxy to a dict + if len(cls.__bases__) == 1: + inherited_dict = cls.__bases__[0].__dict__ + else: + inherited_dict = {} + for base in reversed(cls.__bases__): + inherited_dict.update(base.__dict__) + to_remove = [] + for name, value in clsdict.items(): + try: + base_value = inherited_dict[name] + if value is base_value: + to_remove.append(name) + except KeyError: + pass + for name in to_remove: + clsdict.pop(name) + return clsdict + + +if sys.version_info[:2] < (3, 7): # pragma: no branch + def _is_parametrized_type_hint(obj): + # This is very cheap but might generate false positives. So try to + # narrow it down is good as possible. + type_module = getattr(type(obj), '__module__', None) + from_typing_extensions = type_module == 'typing_extensions' + from_typing = type_module == 'typing' + + # general typing Constructs + is_typing = getattr(obj, '__origin__', None) is not None + + # typing_extensions.Literal + is_literal = ( + (getattr(obj, '__values__', None) is not None) + and from_typing_extensions + ) + + # typing_extensions.Final + is_final = ( + (getattr(obj, '__type__', None) is not None) + and from_typing_extensions + ) + + # typing.ClassVar + is_classvar = ( + (getattr(obj, '__type__', None) is not None) and from_typing + ) + + # typing.Union/Tuple for old Python 3.5 + is_union = getattr(obj, '__union_params__', None) is not None + is_tuple = getattr(obj, '__tuple_params__', None) is not None + is_callable = ( + getattr(obj, '__result__', None) is not None and + getattr(obj, '__args__', None) is not None + ) + return any((is_typing, is_literal, is_final, is_classvar, is_union, + is_tuple, is_callable)) + + def _create_parametrized_type_hint(origin, args): + return origin[args] +else: + _is_parametrized_type_hint = None + _create_parametrized_type_hint = None + + +def parametrized_type_hint_getinitargs(obj): + # The distorted type check sematic for typing construct becomes: + # ``type(obj) is type(TypeHint)``, which means "obj is a + # parametrized TypeHint" + if type(obj) is type(Literal): # pragma: no branch + initargs = (Literal, obj.__values__) + elif type(obj) is type(Final): # pragma: no branch + initargs = (Final, obj.__type__) + elif type(obj) is type(ClassVar): + initargs = (ClassVar, obj.__type__) + elif type(obj) is type(Generic): + initargs = (obj.__origin__, obj.__args__) + elif type(obj) is type(Union): + initargs = (Union, obj.__args__) + elif type(obj) is type(Tuple): + initargs = (Tuple, obj.__args__) + elif type(obj) is type(Callable): + (*args, result) = obj.__args__ + if len(args) == 1 and args[0] is Ellipsis: + args = Ellipsis + else: + args = list(args) + initargs = (Callable, (args, result)) + else: # pragma: no cover + raise pickle.PicklingError( + f"Cloudpickle Error: Unknown type {type(obj)}" + ) + return initargs + + +# Tornado support + +def is_tornado_coroutine(func): + """ + Return whether *func* is a Tornado coroutine function. + Running coroutines are not supported. + """ + if 'tornado.gen' not in sys.modules: + return False + gen = sys.modules['tornado.gen'] + if not hasattr(gen, "is_coroutine_function"): + # Tornado version is too old + return False + return gen.is_coroutine_function(func) + + +def _rebuild_tornado_coroutine(func): + from tornado import gen + return gen.coroutine(func) + + +# including pickles unloading functions in this namespace +load = pickle.load +loads = pickle.loads + + +def subimport(name): + # We cannot do simply: `return __import__(name)`: Indeed, if ``name`` is + # the name of a submodule, __import__ will return the top-level root module + # of this submodule. For instance, __import__('os.path') returns the `os` + # module. + __import__(name) + return sys.modules[name] + + +def dynamic_subimport(name, vars): + mod = types.ModuleType(name) + mod.__dict__.update(vars) + mod.__dict__['__builtins__'] = builtins.__dict__ + return mod + + +def _gen_ellipsis(): + return Ellipsis + + +def _gen_not_implemented(): + return NotImplemented + + +def _get_cell_contents(cell): + try: + return cell.cell_contents + except ValueError: + # sentinel used by ``_fill_function`` which will leave the cell empty + return _empty_cell_value + + +def instance(cls): + """Create a new instance of a class. + + Parameters + ---------- + cls : type + The class to create an instance of. + + Returns + ------- + instance : cls + A new instance of ``cls``. + """ + return cls() + + +@instance +class _empty_cell_value: + """sentinel for empty closures + """ + @classmethod + def __reduce__(cls): + return cls.__name__ + + +def _fill_function(*args): + """Fills in the rest of function data into the skeleton function object + + The skeleton itself is create by _make_skel_func(). + """ + if len(args) == 2: + func = args[0] + state = args[1] + elif len(args) == 5: + # Backwards compat for cloudpickle v0.4.0, after which the `module` + # argument was introduced + func = args[0] + keys = ['globals', 'defaults', 'dict', 'closure_values'] + state = dict(zip(keys, args[1:])) + elif len(args) == 6: + # Backwards compat for cloudpickle v0.4.1, after which the function + # state was passed as a dict to the _fill_function it-self. + func = args[0] + keys = ['globals', 'defaults', 'dict', 'module', 'closure_values'] + state = dict(zip(keys, args[1:])) + else: + raise ValueError(f'Unexpected _fill_value arguments: {args!r}') + + # - At pickling time, any dynamic global variable used by func is + # serialized by value (in state['globals']). + # - At unpickling time, func's __globals__ attribute is initialized by + # first retrieving an empty isolated namespace that will be shared + # with other functions pickled from the same original module + # by the same CloudPickler instance and then updated with the + # content of state['globals'] to populate the shared isolated + # namespace with all the global variables that are specifically + # referenced for this function. + func.__globals__.update(state['globals']) + + func.__defaults__ = state['defaults'] + func.__dict__ = state['dict'] + if 'annotations' in state: + func.__annotations__ = state['annotations'] + if 'doc' in state: + func.__doc__ = state['doc'] + if 'name' in state: + func.__name__ = state['name'] + if 'module' in state: + func.__module__ = state['module'] + if 'qualname' in state: + func.__qualname__ = state['qualname'] + if 'kwdefaults' in state: + func.__kwdefaults__ = state['kwdefaults'] + # _cloudpickle_subimports is a set of submodules that must be loaded for + # the pickled function to work correctly at unpickling time. Now that these + # submodules are depickled (hence imported), they can be removed from the + # object's state (the object state only served as a reference holder to + # these submodules) + if '_cloudpickle_submodules' in state: + state.pop('_cloudpickle_submodules') + + cells = func.__closure__ + if cells is not None: + for cell, value in zip(cells, state['closure_values']): + if value is not _empty_cell_value: + cell_set(cell, value) + + return func + + +def _make_function(code, globals, name, argdefs, closure): + # Setting __builtins__ in globals is needed for nogil CPython. + globals["__builtins__"] = __builtins__ + return types.FunctionType(code, globals, name, argdefs, closure) + + +def _make_empty_cell(): + if False: + # trick the compiler into creating an empty cell in our lambda + cell = None + raise AssertionError('this route should not be executed') + + return (lambda: cell).__closure__[0] + + +def _make_cell(value=_empty_cell_value): + cell = _make_empty_cell() + if value is not _empty_cell_value: + cell_set(cell, value) + return cell + + +def _make_skel_func(code, cell_count, base_globals=None): + """ Creates a skeleton function object that contains just the provided + code and the correct number of cells in func_closure. All other + func attributes (e.g. func_globals) are empty. + """ + # This function is deprecated and should be removed in cloudpickle 1.7 + warnings.warn( + "A pickle file created using an old (<=1.4.1) version of cloudpickle " + "is currently being loaded. This is not supported by cloudpickle and " + "will break in cloudpickle 1.7", category=UserWarning + ) + # This is backward-compatibility code: for cloudpickle versions between + # 0.5.4 and 0.7, base_globals could be a string or None. base_globals + # should now always be a dictionary. + if base_globals is None or isinstance(base_globals, str): + base_globals = {} + + base_globals['__builtins__'] = __builtins__ + + closure = ( + tuple(_make_empty_cell() for _ in range(cell_count)) + if cell_count >= 0 else + None + ) + return types.FunctionType(code, base_globals, None, None, closure) + + +def _make_skeleton_class(type_constructor, name, bases, type_kwargs, + class_tracker_id, extra): + """Build dynamic class with an empty __dict__ to be filled once memoized + + If class_tracker_id is not None, try to lookup an existing class definition + matching that id. If none is found, track a newly reconstructed class + definition under that id so that other instances stemming from the same + class id will also reuse this class definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + skeleton_class = types.new_class( + name, bases, {'metaclass': type_constructor}, + lambda ns: ns.update(type_kwargs) + ) + return _lookup_class_or_track(class_tracker_id, skeleton_class) + + +def _rehydrate_skeleton_class(skeleton_class, class_dict): + """Put attributes from `class_dict` back on `skeleton_class`. + + See CloudPickler.save_dynamic_class for more info. + """ + registry = None + for attrname, attr in class_dict.items(): + if attrname == "_abc_impl": + registry = attr + else: + setattr(skeleton_class, attrname, attr) + if registry is not None: + for subclass in registry: + skeleton_class.register(subclass) + + return skeleton_class + + +def _make_skeleton_enum(bases, name, qualname, members, module, + class_tracker_id, extra): + """Build dynamic enum with an empty __dict__ to be filled once memoized + + The creation of the enum class is inspired by the code of + EnumMeta._create_. + + If class_tracker_id is not None, try to lookup an existing enum definition + matching that id. If none is found, track a newly reconstructed enum + definition under that id so that other instances stemming from the same + class id will also reuse this enum definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + # enums always inherit from their base Enum class at the last position in + # the list of base classes: + enum_base = bases[-1] + metacls = enum_base.__class__ + classdict = metacls.__prepare__(name, bases) + + for member_name, member_value in members.items(): + classdict[member_name] = member_value + enum_class = metacls.__new__(metacls, name, bases, classdict) + enum_class.__module__ = module + enum_class.__qualname__ = qualname + + return _lookup_class_or_track(class_tracker_id, enum_class) + + +def _make_typevar(name, bound, constraints, covariant, contravariant, + class_tracker_id): + tv = typing.TypeVar( + name, *constraints, bound=bound, + covariant=covariant, contravariant=contravariant + ) + if class_tracker_id is not None: + return _lookup_class_or_track(class_tracker_id, tv) + else: # pragma: nocover + # Only for Python 3.5.3 compat. + return tv + + +def _decompose_typevar(obj): + return ( + obj.__name__, obj.__bound__, obj.__constraints__, + obj.__covariant__, obj.__contravariant__, + _get_or_create_tracker_id(obj), + ) + + +def _typevar_reduce(obj): + # TypeVar instances require the module information hence why we + # are not using the _should_pickle_by_reference directly + module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__) + + if module_and_name is None: + return (_make_typevar, _decompose_typevar(obj)) + elif _is_registered_pickle_by_value(module_and_name[0]): + return (_make_typevar, _decompose_typevar(obj)) + + return (getattr, module_and_name) + + +def _get_bases(typ): + if '__orig_bases__' in getattr(typ, '__dict__', {}): + # For generic types (see PEP 560) + # Note that simply checking `hasattr(typ, '__orig_bases__')` is not + # correct. Subclasses of a fully-parameterized generic class does not + # have `__orig_bases__` defined, but `hasattr(typ, '__orig_bases__')` + # will return True because it's defined in the base class. + bases_attr = '__orig_bases__' + else: + # For regular class objects + bases_attr = '__bases__' + return getattr(typ, bases_attr) + + +def _make_dict_keys(obj, is_ordered=False): + if is_ordered: + return OrderedDict.fromkeys(obj).keys() + else: + return dict.fromkeys(obj).keys() + + +def _make_dict_values(obj, is_ordered=False): + if is_ordered: + return OrderedDict((i, _) for i, _ in enumerate(obj)).values() + else: + return {i: _ for i, _ in enumerate(obj)}.values() + + +def _make_dict_items(obj, is_ordered=False): + if is_ordered: + return OrderedDict(obj).items() + else: + return obj.items() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle_fast.py b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle_fast.py new file mode 100644 index 0000000..ee1f4b8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/cloudpickle_fast.py @@ -0,0 +1,844 @@ +""" +New, fast version of the CloudPickler. + +This new CloudPickler class can now extend the fast C Pickler instead of the +previous Python implementation of the Pickler class. Because this functionality +is only available for Python versions 3.8+, a lot of backward-compatibility +code is also removed. + +Note that the C Pickler subclassing API is CPython-specific. Therefore, some +guards present in cloudpickle.py that were written to handle PyPy specificities +are not present in cloudpickle_fast.py +""" +import _collections_abc +import abc +import copyreg +import io +import itertools +import logging +import sys +import struct +import types +import weakref +import typing + +from enum import Enum +from collections import ChainMap, OrderedDict + +from .compat import pickle, Pickler +from .cloudpickle import ( + _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL, + _find_imported_submodules, _get_cell_contents, _should_pickle_by_reference, + _builtin_type, _get_or_create_tracker_id, _make_skeleton_class, + _make_skeleton_enum, _extract_class_dict, dynamic_subimport, subimport, + _typevar_reduce, _get_bases, _make_cell, _make_empty_cell, CellType, + _is_parametrized_type_hint, PYPY, cell_set, + parametrized_type_hint_getinitargs, _create_parametrized_type_hint, + builtin_code_type, + _make_dict_keys, _make_dict_values, _make_dict_items, _make_function, +) + + +if pickle.HIGHEST_PROTOCOL >= 5: + # Shorthands similar to pickle.dump/pickle.dumps + + def dump(obj, file, protocol=None, buffer_callback=None): + """Serialize obj as bytes streamed into file + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python. + """ + CloudPickler( + file, protocol=protocol, buffer_callback=buffer_callback + ).dump(obj) + + def dumps(obj, protocol=None, buffer_callback=None): + """Serialize obj as a string of bytes allocated in memory + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python. + """ + with io.BytesIO() as file: + cp = CloudPickler( + file, protocol=protocol, buffer_callback=buffer_callback + ) + cp.dump(obj) + return file.getvalue() + +else: + # Shorthands similar to pickle.dump/pickle.dumps + def dump(obj, file, protocol=None): + """Serialize obj as bytes streamed into file + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python. + """ + CloudPickler(file, protocol=protocol).dump(obj) + + def dumps(obj, protocol=None): + """Serialize obj as a string of bytes allocated in memory + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python. + """ + with io.BytesIO() as file: + cp = CloudPickler(file, protocol=protocol) + cp.dump(obj) + return file.getvalue() + + +load, loads = pickle.load, pickle.loads + + +# COLLECTION OF OBJECTS __getnewargs__-LIKE METHODS +# ------------------------------------------------- + +def _class_getnewargs(obj): + type_kwargs = {} + if "__module__" in obj.__dict__: + type_kwargs["__module__"] = obj.__module__ + + __dict__ = obj.__dict__.get('__dict__', None) + if isinstance(__dict__, property): + type_kwargs['__dict__'] = __dict__ + + return (type(obj), obj.__name__, _get_bases(obj), type_kwargs, + _get_or_create_tracker_id(obj), None) + + +def _enum_getnewargs(obj): + members = {e.name: e.value for e in obj} + return (obj.__bases__, obj.__name__, obj.__qualname__, members, + obj.__module__, _get_or_create_tracker_id(obj), None) + + +# COLLECTION OF OBJECTS RECONSTRUCTORS +# ------------------------------------ +def _file_reconstructor(retval): + return retval + + +# COLLECTION OF OBJECTS STATE GETTERS +# ----------------------------------- +def _function_getstate(func): + # - Put func's dynamic attributes (stored in func.__dict__) in state. These + # attributes will be restored at unpickling time using + # f.__dict__.update(state) + # - Put func's members into slotstate. Such attributes will be restored at + # unpickling time by iterating over slotstate and calling setattr(func, + # slotname, slotvalue) + slotstate = { + "__name__": func.__name__, + "__qualname__": func.__qualname__, + "__annotations__": func.__annotations__, + "__kwdefaults__": func.__kwdefaults__, + "__defaults__": func.__defaults__, + "__module__": func.__module__, + "__doc__": func.__doc__, + "__closure__": func.__closure__, + } + + f_globals_ref = _extract_code_globals(func.__code__) + f_globals = {k: func.__globals__[k] for k in f_globals_ref if k in + func.__globals__} + + closure_values = ( + list(map(_get_cell_contents, func.__closure__)) + if func.__closure__ is not None else () + ) + + # Extract currently-imported submodules used by func. Storing these modules + # in a smoke _cloudpickle_subimports attribute of the object's state will + # trigger the side effect of importing these modules at unpickling time + # (which is necessary for func to work correctly once depickled) + slotstate["_cloudpickle_submodules"] = _find_imported_submodules( + func.__code__, itertools.chain(f_globals.values(), closure_values)) + slotstate["__globals__"] = f_globals + + state = func.__dict__ + return state, slotstate + + +def _class_getstate(obj): + clsdict = _extract_class_dict(obj) + clsdict.pop('__weakref__', None) + + if issubclass(type(obj), abc.ABCMeta): + # If obj is an instance of an ABCMeta subclass, don't pickle the + # cache/negative caches populated during isinstance/issubclass + # checks, but pickle the list of registered subclasses of obj. + clsdict.pop('_abc_cache', None) + clsdict.pop('_abc_negative_cache', None) + clsdict.pop('_abc_negative_cache_version', None) + registry = clsdict.pop('_abc_registry', None) + if registry is None: + # in Python3.7+, the abc caches and registered subclasses of a + # class are bundled into the single _abc_impl attribute + clsdict.pop('_abc_impl', None) + (registry, _, _, _) = abc._get_dump(obj) + + clsdict["_abc_impl"] = [subclass_weakref() + for subclass_weakref in registry] + else: + # In the above if clause, registry is a set of weakrefs -- in + # this case, registry is a WeakSet + clsdict["_abc_impl"] = [type_ for type_ in registry] + + if "__slots__" in clsdict: + # pickle string length optimization: member descriptors of obj are + # created automatically from obj's __slots__ attribute, no need to + # save them in obj's state + if isinstance(obj.__slots__, str): + clsdict.pop(obj.__slots__) + else: + for k in obj.__slots__: + clsdict.pop(k, None) + + clsdict.pop('__dict__', None) # unpicklable property object + + return (clsdict, {}) + + +def _enum_getstate(obj): + clsdict, slotstate = _class_getstate(obj) + + members = {e.name: e.value for e in obj} + # Cleanup the clsdict that will be passed to _rehydrate_skeleton_class: + # Those attributes are already handled by the metaclass. + for attrname in ["_generate_next_value_", "_member_names_", + "_member_map_", "_member_type_", + "_value2member_map_"]: + clsdict.pop(attrname, None) + for member in members: + clsdict.pop(member) + # Special handling of Enum subclasses + return clsdict, slotstate + + +# COLLECTIONS OF OBJECTS REDUCERS +# ------------------------------- +# A reducer is a function taking a single argument (obj), and that returns a +# tuple with all the necessary data to re-construct obj. Apart from a few +# exceptions (list, dict, bytes, int, etc.), a reducer is necessary to +# correctly pickle an object. +# While many built-in objects (Exceptions objects, instances of the "object" +# class, etc), are shipped with their own built-in reducer (invoked using +# obj.__reduce__), some do not. The following methods were created to "fill +# these holes". + +def _code_reduce(obj): + """codeobject reducer""" + # If you are not sure about the order of arguments, take a look at help + # of the specific type from types, for example: + # >>> from types import CodeType + # >>> help(CodeType) + if hasattr(obj, "co_exceptiontable"): # pragma: no branch + # Python 3.11 and later: there are some new attributes + # related to the enhanced exceptions. + args = ( + obj.co_argcount, obj.co_posonlyargcount, + obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, + obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, + obj.co_varnames, obj.co_filename, obj.co_name, obj.co_qualname, + obj.co_firstlineno, obj.co_linetable, obj.co_exceptiontable, + obj.co_freevars, obj.co_cellvars, + ) + elif hasattr(obj, "co_linetable"): # pragma: no branch + # Python 3.10 and later: obj.co_lnotab is deprecated and constructor + # expects obj.co_linetable instead. + args = ( + obj.co_argcount, obj.co_posonlyargcount, + obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, + obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, + obj.co_varnames, obj.co_filename, obj.co_name, + obj.co_firstlineno, obj.co_linetable, obj.co_freevars, + obj.co_cellvars + ) + elif hasattr(obj, "co_nmeta"): # pragma: no cover + # "nogil" Python: modified attributes from 3.9 + args = ( + obj.co_argcount, obj.co_posonlyargcount, + obj.co_kwonlyargcount, obj.co_nlocals, obj.co_framesize, + obj.co_ndefaultargs, obj.co_nmeta, + obj.co_flags, obj.co_code, obj.co_consts, + obj.co_varnames, obj.co_filename, obj.co_name, + obj.co_firstlineno, obj.co_lnotab, obj.co_exc_handlers, + obj.co_jump_table, obj.co_freevars, obj.co_cellvars, + obj.co_free2reg, obj.co_cell2reg + ) + elif hasattr(obj, "co_posonlyargcount"): + # Backward compat for 3.9 and older + args = ( + obj.co_argcount, obj.co_posonlyargcount, + obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, + obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, + obj.co_varnames, obj.co_filename, obj.co_name, + obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, + obj.co_cellvars + ) + else: + # Backward compat for even older versions of Python + args = ( + obj.co_argcount, obj.co_kwonlyargcount, obj.co_nlocals, + obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, + obj.co_names, obj.co_varnames, obj.co_filename, + obj.co_name, obj.co_firstlineno, obj.co_lnotab, + obj.co_freevars, obj.co_cellvars + ) + return types.CodeType, args + + +def _cell_reduce(obj): + """Cell (containing values of a function's free variables) reducer""" + try: + obj.cell_contents + except ValueError: # cell is empty + return _make_empty_cell, () + else: + return _make_cell, (obj.cell_contents, ) + + +def _classmethod_reduce(obj): + orig_func = obj.__func__ + return type(obj), (orig_func,) + + +def _file_reduce(obj): + """Save a file""" + import io + + if not hasattr(obj, "name") or not hasattr(obj, "mode"): + raise pickle.PicklingError( + "Cannot pickle files that do not map to an actual file" + ) + if obj is sys.stdout: + return getattr, (sys, "stdout") + if obj is sys.stderr: + return getattr, (sys, "stderr") + if obj is sys.stdin: + raise pickle.PicklingError("Cannot pickle standard input") + if obj.closed: + raise pickle.PicklingError("Cannot pickle closed files") + if hasattr(obj, "isatty") and obj.isatty(): + raise pickle.PicklingError( + "Cannot pickle files that map to tty objects" + ) + if "r" not in obj.mode and "+" not in obj.mode: + raise pickle.PicklingError( + "Cannot pickle files that are not opened for reading: %s" + % obj.mode + ) + + name = obj.name + + retval = io.StringIO() + + try: + # Read the whole file + curloc = obj.tell() + obj.seek(0) + contents = obj.read() + obj.seek(curloc) + except IOError as e: + raise pickle.PicklingError( + "Cannot pickle file %s as it cannot be read" % name + ) from e + retval.write(contents) + retval.seek(curloc) + + retval.name = name + return _file_reconstructor, (retval,) + + +def _getset_descriptor_reduce(obj): + return getattr, (obj.__objclass__, obj.__name__) + + +def _mappingproxy_reduce(obj): + return types.MappingProxyType, (dict(obj),) + + +def _memoryview_reduce(obj): + return bytes, (obj.tobytes(),) + + +def _module_reduce(obj): + if _should_pickle_by_reference(obj): + return subimport, (obj.__name__,) + else: + # Some external libraries can populate the "__builtins__" entry of a + # module's `__dict__` with unpicklable objects (see #316). For that + # reason, we do not attempt to pickle the "__builtins__" entry, and + # restore a default value for it at unpickling time. + state = obj.__dict__.copy() + state.pop('__builtins__', None) + return dynamic_subimport, (obj.__name__, state) + + +def _method_reduce(obj): + return (types.MethodType, (obj.__func__, obj.__self__)) + + +def _logger_reduce(obj): + return logging.getLogger, (obj.name,) + + +def _root_logger_reduce(obj): + return logging.getLogger, () + + +def _property_reduce(obj): + return property, (obj.fget, obj.fset, obj.fdel, obj.__doc__) + + +def _weakset_reduce(obj): + return weakref.WeakSet, (list(obj),) + + +def _dynamic_class_reduce(obj): + """ + Save a class that can't be stored as module global. + + This method is used to serialize classes that are defined inside + functions, or that otherwise can't be serialized as attribute lookups + from global modules. + """ + if Enum is not None and issubclass(obj, Enum): + return ( + _make_skeleton_enum, _enum_getnewargs(obj), _enum_getstate(obj), + None, None, _class_setstate + ) + else: + return ( + _make_skeleton_class, _class_getnewargs(obj), _class_getstate(obj), + None, None, _class_setstate + ) + + +def _class_reduce(obj): + """Select the reducer depending on the dynamic nature of the class obj""" + if obj is type(None): # noqa + return type, (None,) + elif obj is type(Ellipsis): + return type, (Ellipsis,) + elif obj is type(NotImplemented): + return type, (NotImplemented,) + elif obj in _BUILTIN_TYPE_NAMES: + return _builtin_type, (_BUILTIN_TYPE_NAMES[obj],) + elif not _should_pickle_by_reference(obj): + return _dynamic_class_reduce(obj) + return NotImplemented + + +def _dict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj), ) + + +def _dict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj), ) + + +def _dict_items_reduce(obj): + return _make_dict_items, (dict(obj), ) + + +def _odict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj), True) + + +def _odict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj), True) + + +def _odict_items_reduce(obj): + return _make_dict_items, (dict(obj), True) + + +# COLLECTIONS OF OBJECTS STATE SETTERS +# ------------------------------------ +# state setters are called at unpickling time, once the object is created and +# it has to be updated to how it was at unpickling time. + + +def _function_setstate(obj, state): + """Update the state of a dynamic function. + + As __closure__ and __globals__ are readonly attributes of a function, we + cannot rely on the native setstate routine of pickle.load_build, that calls + setattr on items of the slotstate. Instead, we have to modify them inplace. + """ + state, slotstate = state + obj.__dict__.update(state) + + obj_globals = slotstate.pop("__globals__") + obj_closure = slotstate.pop("__closure__") + # _cloudpickle_subimports is a set of submodules that must be loaded for + # the pickled function to work correctly at unpickling time. Now that these + # submodules are depickled (hence imported), they can be removed from the + # object's state (the object state only served as a reference holder to + # these submodules) + slotstate.pop("_cloudpickle_submodules") + + obj.__globals__.update(obj_globals) + obj.__globals__["__builtins__"] = __builtins__ + + if obj_closure is not None: + for i, cell in enumerate(obj_closure): + try: + value = cell.cell_contents + except ValueError: # cell is empty + continue + cell_set(obj.__closure__[i], value) + + for k, v in slotstate.items(): + setattr(obj, k, v) + + +def _class_setstate(obj, state): + state, slotstate = state + registry = None + for attrname, attr in state.items(): + if attrname == "_abc_impl": + registry = attr + else: + setattr(obj, attrname, attr) + if registry is not None: + for subclass in registry: + obj.register(subclass) + + return obj + + +class CloudPickler(Pickler): + # set of reducers defined and used by cloudpickle (private) + _dispatch_table = {} + _dispatch_table[classmethod] = _classmethod_reduce + _dispatch_table[io.TextIOWrapper] = _file_reduce + _dispatch_table[logging.Logger] = _logger_reduce + _dispatch_table[logging.RootLogger] = _root_logger_reduce + _dispatch_table[memoryview] = _memoryview_reduce + _dispatch_table[property] = _property_reduce + _dispatch_table[staticmethod] = _classmethod_reduce + _dispatch_table[CellType] = _cell_reduce + _dispatch_table[types.CodeType] = _code_reduce + _dispatch_table[types.GetSetDescriptorType] = _getset_descriptor_reduce + _dispatch_table[types.ModuleType] = _module_reduce + _dispatch_table[types.MethodType] = _method_reduce + _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce + _dispatch_table[weakref.WeakSet] = _weakset_reduce + _dispatch_table[typing.TypeVar] = _typevar_reduce + _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce + _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce + _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce + _dispatch_table[type(OrderedDict().keys())] = _odict_keys_reduce + _dispatch_table[type(OrderedDict().values())] = _odict_values_reduce + _dispatch_table[type(OrderedDict().items())] = _odict_items_reduce + _dispatch_table[abc.abstractmethod] = _classmethod_reduce + _dispatch_table[abc.abstractclassmethod] = _classmethod_reduce + _dispatch_table[abc.abstractstaticmethod] = _classmethod_reduce + _dispatch_table[abc.abstractproperty] = _property_reduce + + dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table) + + # function reducers are defined as instance methods of CloudPickler + # objects, as they rely on a CloudPickler attribute (globals_ref) + def _dynamic_function_reduce(self, func): + """Reduce a function that is not pickleable via attribute lookup.""" + newargs = self._function_getnewargs(func) + state = _function_getstate(func) + return (_make_function, newargs, state, None, None, + _function_setstate) + + def _function_reduce(self, obj): + """Reducer for function objects. + + If obj is a top-level attribute of a file-backed module, this + reducer returns NotImplemented, making the CloudPickler fallback to + traditional _pickle.Pickler routines to save obj. Otherwise, it reduces + obj using a custom cloudpickle reducer designed specifically to handle + dynamic functions. + + As opposed to cloudpickle.py, There no special handling for builtin + pypy functions because cloudpickle_fast is CPython-specific. + """ + if _should_pickle_by_reference(obj): + return NotImplemented + else: + return self._dynamic_function_reduce(obj) + + def _function_getnewargs(self, func): + code = func.__code__ + + # base_globals represents the future global namespace of func at + # unpickling time. Looking it up and storing it in + # CloudpiPickler.globals_ref allow functions sharing the same globals + # at pickling time to also share them once unpickled, at one condition: + # since globals_ref is an attribute of a CloudPickler instance, and + # that a new CloudPickler is created each time pickle.dump or + # pickle.dumps is called, functions also need to be saved within the + # same invocation of cloudpickle.dump/cloudpickle.dumps (for example: + # cloudpickle.dumps([f1, f2])). There is no such limitation when using + # CloudPickler.dump, as long as the multiple invocations are bound to + # the same CloudPickler. + base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) + + if base_globals == {}: + # Add module attributes used to resolve relative imports + # instructions inside func. + for k in ["__package__", "__name__", "__path__", "__file__"]: + if k in func.__globals__: + base_globals[k] = func.__globals__[k] + + # Do not bind the free variables before the function is created to + # avoid infinite recursion. + if func.__closure__ is None: + closure = None + else: + closure = tuple( + _make_empty_cell() for _ in range(len(code.co_freevars))) + + return code, base_globals, None, None, closure + + def dump(self, obj): + try: + return Pickler.dump(self, obj) + except RuntimeError as e: + if len(e.args) > 0 and "recursion" in e.args[0]: + msg = ( + "Could not pickle object as excessively deep recursion " + "required." + ) + raise pickle.PicklingError(msg) from e + else: + raise + + if pickle.HIGHEST_PROTOCOL >= 5: + def __init__(self, file, protocol=None, buffer_callback=None): + if protocol is None: + protocol = DEFAULT_PROTOCOL + Pickler.__init__( + self, file, protocol=protocol, buffer_callback=buffer_callback + ) + # map functions __globals__ attribute ids, to ensure that functions + # sharing the same global namespace at pickling time also share + # their global namespace at unpickling time. + self.globals_ref = {} + self.proto = int(protocol) + else: + def __init__(self, file, protocol=None): + if protocol is None: + protocol = DEFAULT_PROTOCOL + Pickler.__init__(self, file, protocol=protocol) + # map functions __globals__ attribute ids, to ensure that functions + # sharing the same global namespace at pickling time also share + # their global namespace at unpickling time. + self.globals_ref = {} + assert hasattr(self, 'proto') + + if pickle.HIGHEST_PROTOCOL >= 5 and not PYPY: + # Pickler is the C implementation of the CPython pickler and therefore + # we rely on reduce_override method to customize the pickler behavior. + + # `CloudPickler.dispatch` is only left for backward compatibility - note + # that when using protocol 5, `CloudPickler.dispatch` is not an + # extension of `Pickler.dispatch` dictionary, because CloudPickler + # subclasses the C-implemented Pickler, which does not expose a + # `dispatch` attribute. Earlier versions of the protocol 5 CloudPickler + # used `CloudPickler.dispatch` as a class-level attribute storing all + # reducers implemented by cloudpickle, but the attribute name was not a + # great choice given the meaning of `CloudPickler.dispatch` when + # `CloudPickler` extends the pure-python pickler. + dispatch = dispatch_table + + # Implementation of the reducer_override callback, in order to + # efficiently serialize dynamic functions and classes by subclassing + # the C-implemented Pickler. + # TODO: decorrelate reducer_override (which is tied to CPython's + # implementation - would it make sense to backport it to pypy? - and + # pickle's protocol 5 which is implementation agnostic. Currently, the + # availability of both notions coincide on CPython's pickle and the + # pickle5 backport, but it may not be the case anymore when pypy + # implements protocol 5 + + def reducer_override(self, obj): + """Type-agnostic reducing callback for function and classes. + + For performance reasons, subclasses of the C _pickle.Pickler class + cannot register custom reducers for functions and classes in the + dispatch_table. Reducer for such types must instead implemented in + the special reducer_override method. + + Note that method will be called for any object except a few + builtin-types (int, lists, dicts etc.), which differs from reducers + in the Pickler's dispatch_table, each of them being invoked for + objects of a specific type only. + + This property comes in handy for classes: although most classes are + instances of the ``type`` metaclass, some of them can be instances + of other custom metaclasses (such as enum.EnumMeta for example). In + particular, the metaclass will likely not be known in advance, and + thus cannot be special-cased using an entry in the dispatch_table. + reducer_override, among other things, allows us to register a + reducer that will be called for any class, independently of its + type. + + + Notes: + + * reducer_override has the priority over dispatch_table-registered + reducers. + * reducer_override can be used to fix other limitations of + cloudpickle for other types that suffered from type-specific + reducers, such as Exceptions. See + https://github.com/cloudpipe/cloudpickle/issues/248 + """ + if sys.version_info[:2] < (3, 7) and _is_parametrized_type_hint(obj): # noqa # pragma: no branch + return ( + _create_parametrized_type_hint, + parametrized_type_hint_getinitargs(obj) + ) + t = type(obj) + try: + is_anyclass = issubclass(t, type) + except TypeError: # t is not a class (old Boost; see SF #502085) + is_anyclass = False + + if is_anyclass: + return _class_reduce(obj) + elif isinstance(obj, types.FunctionType): + return self._function_reduce(obj) + else: + # fallback to save_global, including the Pickler's + # dispatch_table + return NotImplemented + + else: + # When reducer_override is not available, hack the pure-Python + # Pickler's types.FunctionType and type savers. Note: the type saver + # must override Pickler.save_global, because pickle.py contains a + # hard-coded call to save_global when pickling meta-classes. + dispatch = Pickler.dispatch.copy() + + def _save_reduce_pickle5(self, func, args, state=None, listitems=None, + dictitems=None, state_setter=None, obj=None): + save = self.save + write = self.write + self.save_reduce( + func, args, state=None, listitems=listitems, + dictitems=dictitems, obj=obj + ) + # backport of the Python 3.8 state_setter pickle operations + save(state_setter) + save(obj) # simple BINGET opcode as obj is already memoized. + save(state) + write(pickle.TUPLE2) + # Trigger a state_setter(obj, state) function call. + write(pickle.REDUCE) + # The purpose of state_setter is to carry-out an + # inplace modification of obj. We do not care about what the + # method might return, so its output is eventually removed from + # the stack. + write(pickle.POP) + + def save_global(self, obj, name=None, pack=struct.pack): + """ + Save a "global". + + The name of this method is somewhat misleading: all types get + dispatched here. + """ + if obj is type(None): # noqa + return self.save_reduce(type, (None,), obj=obj) + elif obj is type(Ellipsis): + return self.save_reduce(type, (Ellipsis,), obj=obj) + elif obj is type(NotImplemented): + return self.save_reduce(type, (NotImplemented,), obj=obj) + elif obj in _BUILTIN_TYPE_NAMES: + return self.save_reduce( + _builtin_type, (_BUILTIN_TYPE_NAMES[obj],), obj=obj) + + if sys.version_info[:2] < (3, 7) and _is_parametrized_type_hint(obj): # noqa # pragma: no branch + # Parametrized typing constructs in Python < 3.7 are not + # compatible with type checks and ``isinstance`` semantics. For + # this reason, it is easier to detect them using a + # duck-typing-based check (``_is_parametrized_type_hint``) than + # to populate the Pickler's dispatch with type-specific savers. + self.save_reduce( + _create_parametrized_type_hint, + parametrized_type_hint_getinitargs(obj), + obj=obj + ) + elif name is not None: + Pickler.save_global(self, obj, name=name) + elif not _should_pickle_by_reference(obj, name=name): + self._save_reduce_pickle5(*_dynamic_class_reduce(obj), obj=obj) + else: + Pickler.save_global(self, obj, name=name) + dispatch[type] = save_global + + def save_function(self, obj, name=None): + """ Registered with the dispatch to handle all function types. + + Determines what kind of function obj is (e.g. lambda, defined at + interactive prompt, etc) and handles the pickling appropriately. + """ + if _should_pickle_by_reference(obj, name=name): + return Pickler.save_global(self, obj, name=name) + elif PYPY and isinstance(obj.__code__, builtin_code_type): + return self.save_pypy_builtin_func(obj) + else: + return self._save_reduce_pickle5( + *self._dynamic_function_reduce(obj), obj=obj + ) + + def save_pypy_builtin_func(self, obj): + """Save pypy equivalent of builtin functions. + PyPy does not have the concept of builtin-functions. Instead, + builtin-functions are simple function instances, but with a + builtin-code attribute. + Most of the time, builtin functions should be pickled by attribute. + But PyPy has flaky support for __qualname__, so some builtin + functions such as float.__new__ will be classified as dynamic. For + this reason only, we created this special routine. Because + builtin-functions are not expected to have closure or globals, + there is no additional hack (compared the one already implemented + in pickle) to protect ourselves from reference cycles. A simple + (reconstructor, newargs, obj.__dict__) tuple is save_reduced. Note + also that PyPy improved their support for __qualname__ in v3.6, so + this routing should be removed when cloudpickle supports only PyPy + 3.6 and later. + """ + rv = (types.FunctionType, (obj.__code__, {}, obj.__name__, + obj.__defaults__, obj.__closure__), + obj.__dict__) + self.save_reduce(*rv, obj=obj) + + dispatch[types.FunctionType] = save_function diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/compat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/compat.py new file mode 100644 index 0000000..5e9b527 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/cloudpickle/compat.py @@ -0,0 +1,18 @@ +import sys + + +if sys.version_info < (3, 8): + try: + import pickle5 as pickle # noqa: F401 + from pickle5 import Pickler # noqa: F401 + except ImportError: + import pickle # noqa: F401 + + # Use the Python pickler for old CPython versions + from pickle import _Pickler as Pickler # noqa: F401 +else: + import pickle # noqa: F401 + + # Pickler will the C implementation in CPython and the Python + # implementation in PyPy + from pickle import Pickler # noqa: F401 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/conf.py new file mode 100644 index 0000000..dd203e2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/conf.py @@ -0,0 +1,276 @@ +# +# 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. +# + +__all__ = ["SparkConf"] + +import sys +from typing import Dict, List, Optional, Tuple, cast, overload + +from py4j.java_gateway import JVMView, JavaObject +from pyspark.errors import PySparkRuntimeError + + +class SparkConf: + """ + Configuration for a Spark application. Used to set various Spark + parameters as key-value pairs. + + Most of the time, you would create a SparkConf object with + ``SparkConf()``, which will load values from `spark.*` Java system + properties as well. In this case, any parameters you set directly on + the :class:`SparkConf` object take priority over system properties. + + For unit tests, you can also call ``SparkConf(false)`` to skip + loading external settings and get the same configuration no matter + what the system properties are. + + All setter methods in this class support chaining. For example, + you can write ``conf.setMaster("local").setAppName("My app")``. + + Parameters + ---------- + loadDefaults : bool + whether to load values from Java system properties (True by default) + _jvm : class:`py4j.java_gateway.JVMView` + internal parameter used to pass a handle to the + Java VM; does not need to be set by users + _jconf : class:`py4j.java_gateway.JavaObject` + Optionally pass in an existing SparkConf handle + to use its parameters + + Notes + ----- + Once a SparkConf object is passed to Spark, it is cloned + and can no longer be modified by the user. + + Examples + -------- + >>> from pyspark.conf import SparkConf + >>> from pyspark.context import SparkContext + >>> conf = SparkConf() + >>> conf.setMaster("local").setAppName("My app") + + >>> conf.get("spark.master") + 'local' + >>> conf.get("spark.app.name") + 'My app' + >>> sc = SparkContext(conf=conf) + >>> sc.master + 'local' + >>> sc.appName + 'My app' + >>> sc.sparkHome is None + True + + >>> conf = SparkConf(loadDefaults=False) + >>> conf.setSparkHome("/path") + + >>> conf.get("spark.home") + '/path' + >>> conf.setExecutorEnv("VAR1", "value1") + + >>> conf.setExecutorEnv(pairs = [("VAR3", "value3"), ("VAR4", "value4")]) + + >>> conf.get("spark.executorEnv.VAR1") + 'value1' + >>> print(conf.toDebugString()) + spark.executorEnv.VAR1=value1 + spark.executorEnv.VAR3=value3 + spark.executorEnv.VAR4=value4 + spark.home=/path + >>> for p in sorted(conf.getAll(), key=lambda p: p[0]): + ... print(p) + ('spark.executorEnv.VAR1', 'value1') + ('spark.executorEnv.VAR3', 'value3') + ('spark.executorEnv.VAR4', 'value4') + ('spark.home', '/path') + >>> conf._jconf.setExecutorEnv("VAR5", "value5") + JavaObject id... + >>> print(conf.toDebugString()) + spark.executorEnv.VAR1=value1 + spark.executorEnv.VAR3=value3 + spark.executorEnv.VAR4=value4 + spark.executorEnv.VAR5=value5 + spark.home=/path + """ + + _jconf: Optional[JavaObject] + _conf: Optional[Dict[str, str]] + + def __init__( + self, + loadDefaults: bool = True, + _jvm: Optional[JVMView] = None, + _jconf: Optional[JavaObject] = None, + ): + """ + Create a new Spark configuration. + """ + if _jconf: + self._jconf = _jconf + else: + from pyspark.context import SparkContext + + _jvm = _jvm or SparkContext._jvm + + if _jvm is not None: + # JVM is created, so create self._jconf directly through JVM + self._jconf = _jvm.SparkConf(loadDefaults) + self._conf = None + else: + # JVM is not created, so store data in self._conf first + self._jconf = None + self._conf = {} + + def set(self, key: str, value: str) -> "SparkConf": + """Set a configuration property.""" + # Try to set self._jconf first if JVM is created, set self._conf if JVM is not created yet. + if self._jconf is not None: + self._jconf.set(key, str(value)) + else: + assert self._conf is not None + self._conf[key] = str(value) + return self + + def setIfMissing(self, key: str, value: str) -> "SparkConf": + """Set a configuration property, if not already set.""" + if self.get(key) is None: + self.set(key, value) + return self + + def setMaster(self, value: str) -> "SparkConf": + """Set master URL to connect to.""" + self.set("spark.master", value) + return self + + def setAppName(self, value: str) -> "SparkConf": + """Set application name.""" + self.set("spark.app.name", value) + return self + + def setSparkHome(self, value: str) -> "SparkConf": + """Set path where Spark is installed on worker nodes.""" + self.set("spark.home", value) + return self + + @overload + def setExecutorEnv(self, key: str, value: str) -> "SparkConf": + ... + + @overload + def setExecutorEnv(self, *, pairs: List[Tuple[str, str]]) -> "SparkConf": + ... + + def setExecutorEnv( + self, + key: Optional[str] = None, + value: Optional[str] = None, + pairs: Optional[List[Tuple[str, str]]] = None, + ) -> "SparkConf": + """Set an environment variable to be passed to executors.""" + if (key is not None and pairs is not None) or (key is None and pairs is None): + raise PySparkRuntimeError( + error_class="KEY_VALUE_PAIR_REQUIRED", + message_parameters={}, + ) + elif key is not None: + self.set("spark.executorEnv.{}".format(key), cast(str, value)) + elif pairs is not None: + for (k, v) in pairs: + self.set("spark.executorEnv.{}".format(k), v) + return self + + def setAll(self, pairs: List[Tuple[str, str]]) -> "SparkConf": + """ + Set multiple parameters, passed as a list of key-value pairs. + + Parameters + ---------- + pairs : iterable of tuples + list of key-value pairs to set + """ + for (k, v) in pairs: + self.set(k, v) + return self + + @overload + def get(self, key: str) -> Optional[str]: + ... + + @overload + def get(self, key: str, defaultValue: None) -> Optional[str]: + ... + + @overload + def get(self, key: str, defaultValue: str) -> str: + ... + + def get(self, key: str, defaultValue: Optional[str] = None) -> Optional[str]: + """Get the configured value for some key, or return a default otherwise.""" + if defaultValue is None: # Py4J doesn't call the right get() if we pass None + if self._jconf is not None: + if not self._jconf.contains(key): + return None + return self._jconf.get(key) + else: + assert self._conf is not None + return self._conf.get(key, None) + else: + if self._jconf is not None: + return self._jconf.get(key, defaultValue) + else: + assert self._conf is not None + return self._conf.get(key, defaultValue) + + def getAll(self) -> List[Tuple[str, str]]: + """Get all values as a list of key-value pairs.""" + if self._jconf is not None: + return [(elem._1(), elem._2()) for elem in cast(JavaObject, self._jconf).getAll()] + else: + assert self._conf is not None + return list(self._conf.items()) + + def contains(self, key: str) -> bool: + """Does this configuration contain a given key?""" + if self._jconf is not None: + return self._jconf.contains(key) + else: + assert self._conf is not None + return key in self._conf + + def toDebugString(self) -> str: + """ + Returns a printable version of the configuration, as a list of + key=value pairs, one per line. + """ + if self._jconf is not None: + return self._jconf.toDebugString() + else: + assert self._conf is not None + return "\n".join("%s=%s" % (k, v) for k, v in self._conf.items()) + + +def _test() -> None: + import doctest + + (failure_count, test_count) = doctest.testmod(optionflags=doctest.ELLIPSIS) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/context.py new file mode 100644 index 0000000..81adb6c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/context.py @@ -0,0 +1,2601 @@ +# +# 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 os +import shutil +import signal +import sys +import threading +import warnings +import importlib +from threading import RLock +from tempfile import NamedTemporaryFile +from types import TracebackType +from typing import ( + Any, + Callable, + cast, + ClassVar, + Dict, + Iterable, + List, + NoReturn, + Optional, + Sequence, + Tuple, + Type, + TYPE_CHECKING, + TypeVar, + Set, +) + +from py4j.java_collections import JavaMap +from py4j.protocol import Py4JError + +from pyspark import accumulators +from pyspark.accumulators import Accumulator +from pyspark.broadcast import Broadcast, BroadcastPickleRegistry +from pyspark.conf import SparkConf +from pyspark.files import SparkFiles +from pyspark.java_gateway import launch_gateway, local_connect_and_auth +from pyspark.serializers import ( + CPickleSerializer, + BatchedSerializer, + Serializer, + UTF8Deserializer, + PairDeserializer, + AutoBatchedSerializer, + NoOpSerializer, + ChunkedStream, +) +from pyspark.storagelevel import StorageLevel +from pyspark.resource.information import ResourceInformation +from pyspark.rdd import RDD, _load_from_socket +from pyspark.taskcontext import TaskContext +from pyspark.traceback_utils import CallSite, first_spark_call +from pyspark.status import StatusTracker +from pyspark.profiler import ProfilerCollector, BasicProfiler, UDFBasicProfiler, MemoryProfiler +from pyspark.errors import PySparkRuntimeError +from py4j.java_gateway import is_instance_of, JavaGateway, JavaObject, JVMView + +if TYPE_CHECKING: + from pyspark.accumulators import AccumulatorParam + +__all__ = ["SparkContext"] + + +# These are special default configs for PySpark, they will overwrite +# the default ones for Spark if they are not configured by user. +DEFAULT_CONFIGS: Dict[str, Any] = { + "spark.serializer.objectStreamReset": 100, + "spark.rdd.compress": True, +} + +T = TypeVar("T") +U = TypeVar("U") + + +class SparkContext: + + """ + Main entry point for Spark functionality. A SparkContext represents the + connection to a Spark cluster, and can be used to create :class:`RDD` and + broadcast variables on that cluster. + + When you create a new SparkContext, at least the master and app name should + be set, either through the named parameters here or through `conf`. + + Parameters + ---------- + master : str, optional + Cluster URL to connect to (e.g. mesos://host:port, spark://host:port, local[4]). + appName : str, optional + A name for your job, to display on the cluster web UI. + sparkHome : str, optional + Location where Spark is installed on cluster nodes. + pyFiles : list, optional + Collection of .zip or .py files to send to the cluster + and add to PYTHONPATH. These can be paths on the local file + system or HDFS, HTTP, HTTPS, or FTP URLs. + environment : dict, optional + A dictionary of environment variables to set on + worker nodes. + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. Set 1 to disable batching, 0 to automatically choose + the batch size based on object sizes, or -1 to use an unlimited + batch size + serializer : :class:`Serializer`, optional, default :class:`CPickleSerializer` + The serializer for RDDs. + conf : :class:`SparkConf`, optional + An object setting Spark properties. + gateway : class:`py4j.java_gateway.JavaGateway`, optional + Use an existing gateway and JVM, otherwise a new JVM + will be instantiated. This is only used internally. + jsc : class:`py4j.java_gateway.JavaObject`, optional + The JavaSparkContext instance. This is only used internally. + profiler_cls : type, optional, default :class:`BasicProfiler` + A class of custom Profiler used to do profiling + udf_profiler_cls : type, optional, default :class:`UDFBasicProfiler` + A class of custom Profiler used to do udf profiling + + Notes + ----- + Only one :class:`SparkContext` should be active per JVM. You must `stop()` + the active :class:`SparkContext` before creating a new one. + + :class:`SparkContext` instance is not supported to share across multiple + processes out of the box, and PySpark does not guarantee multi-processing execution. + Use threads instead for concurrent processing purpose. + + Examples + -------- + >>> from pyspark.context import SparkContext + >>> sc = SparkContext('local', 'test') + >>> sc2 = SparkContext('local', 'test2') # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: ... + """ + + _gateway: ClassVar[Optional[JavaGateway]] = None + _jvm: ClassVar[Optional[JVMView]] = None + _next_accum_id = 0 + _active_spark_context: ClassVar[Optional["SparkContext"]] = None + _lock = RLock() + _python_includes: Optional[ + List[str] + ] = None # zip and egg files that need to be added to PYTHONPATH + serializer: Serializer + profiler_collector: ProfilerCollector + + PACKAGE_EXTENSIONS: Iterable[str] = (".zip", ".egg", ".jar") + + def __init__( + self, + master: Optional[str] = None, + appName: Optional[str] = None, + sparkHome: Optional[str] = None, + pyFiles: Optional[List[str]] = None, + environment: Optional[Dict[str, Any]] = None, + batchSize: int = 0, + serializer: "Serializer" = CPickleSerializer(), + conf: Optional[SparkConf] = None, + gateway: Optional[JavaGateway] = None, + jsc: Optional[JavaObject] = None, + profiler_cls: Type[BasicProfiler] = BasicProfiler, + udf_profiler_cls: Type[UDFBasicProfiler] = UDFBasicProfiler, + memory_profiler_cls: Type[MemoryProfiler] = MemoryProfiler, + ): + if "SPARK_CONNECT_MODE_ENABLED" in os.environ and "SPARK_LOCAL_REMOTE" not in os.environ: + raise PySparkRuntimeError( + error_class="CONTEXT_UNAVAILABLE_FOR_REMOTE_CLIENT", + message_parameters={}, + ) + + if conf is None or conf.get("spark.executor.allowSparkContext", "false").lower() != "true": + # In order to prevent SparkContext from being created in executors. + SparkContext._assert_on_driver() + + self._callsite = first_spark_call() or CallSite(None, None, None) + if gateway is not None and gateway.gateway_parameters.auth_token is None: + raise ValueError( + "You are trying to pass an insecure Py4j gateway to Spark. This" + " is not allowed as it is a security risk." + ) + + SparkContext._ensure_initialized(self, gateway=gateway, conf=conf) + try: + self._do_init( + master, + appName, + sparkHome, + pyFiles, + environment, + batchSize, + serializer, + conf, + jsc, + profiler_cls, + udf_profiler_cls, + memory_profiler_cls, + ) + except BaseException: + # If an error occurs, clean up in order to allow future SparkContext creation: + self.stop() + raise + + def _do_init( + self, + master: Optional[str], + appName: Optional[str], + sparkHome: Optional[str], + pyFiles: Optional[List[str]], + environment: Optional[Dict[str, Any]], + batchSize: int, + serializer: Serializer, + conf: Optional[SparkConf], + jsc: JavaObject, + profiler_cls: Type[BasicProfiler] = BasicProfiler, + udf_profiler_cls: Type[UDFBasicProfiler] = UDFBasicProfiler, + memory_profiler_cls: Type[MemoryProfiler] = MemoryProfiler, + ) -> None: + self.environment = environment or {} + # java gateway must have been launched at this point. + if conf is not None and conf._jconf is not None: + # conf has been initialized in JVM properly, so use conf directly. This represents the + # scenario that JVM has been launched before SparkConf is created (e.g. SparkContext is + # created and then stopped, and we create a new SparkConf and new SparkContext again) + self._conf = conf + else: + self._conf = SparkConf(_jvm=SparkContext._jvm) + if conf is not None: + for k, v in conf.getAll(): + self._conf.set(k, v) + + self._batchSize = batchSize # -1 represents an unlimited batch size + self._unbatched_serializer = serializer + if batchSize == 0: + self.serializer = AutoBatchedSerializer(self._unbatched_serializer) + else: + self.serializer = BatchedSerializer(self._unbatched_serializer, batchSize) + + # Set any parameters passed directly to us on the conf + if master: + self._conf.setMaster(master) + if appName: + self._conf.setAppName(appName) + if sparkHome: + self._conf.setSparkHome(sparkHome) + if environment: + for key, value in environment.items(): + self._conf.setExecutorEnv(key, value) + for key, value in DEFAULT_CONFIGS.items(): + self._conf.setIfMissing(key, value) + + # Check that we have at least the required parameters + if not self._conf.contains("spark.master"): + raise PySparkRuntimeError( + error_class="MASTER_URL_NOT_SET", + message_parameters={}, + ) + if not self._conf.contains("spark.app.name"): + raise PySparkRuntimeError( + error_class="APPLICATION_NAME_NOT_SET", + message_parameters={}, + ) + + # Read back our properties from the conf in case we loaded some of them from + # the classpath or an external config file + self.master = self._conf.get("spark.master") + self.appName = self._conf.get("spark.app.name") + self.sparkHome = self._conf.get("spark.home", None) + + for (k, v) in self._conf.getAll(): + if k.startswith("spark.executorEnv."): + varName = k[len("spark.executorEnv.") :] + self.environment[varName] = v + + self.environment["PYTHONHASHSEED"] = os.environ.get("PYTHONHASHSEED", "0") + + # Create the Java SparkContext through Py4J + self._jsc = jsc or self._initialize_context(self._conf._jconf) + # Reset the SparkConf to the one actually used by the SparkContext in JVM. + self._conf = SparkConf(_jconf=self._jsc.sc().conf()) + + # Create a single Accumulator in Java that we'll send all our updates through; + # they will be passed back to us through a TCP server + assert self._gateway is not None + auth_token = self._gateway.gateway_parameters.auth_token + start_update_server = accumulators._start_update_server + self._accumulatorServer = start_update_server(auth_token) + (host, port) = self._accumulatorServer.server_address + assert self._jvm is not None + self._javaAccumulator = self._jvm.PythonAccumulatorV2(host, port, auth_token) + self._jsc.sc().register(self._javaAccumulator) + + # If encryption is enabled, we need to setup a server in the jvm to read broadcast + # data via a socket. + # scala's mangled names w/ $ in them require special treatment. + self._encryption_enabled = self._jvm.PythonUtils.isEncryptionEnabled(self._jsc) + os.environ["SPARK_AUTH_SOCKET_TIMEOUT"] = str( + self._jvm.PythonUtils.getPythonAuthSocketTimeout(self._jsc) + ) + os.environ["SPARK_BUFFER_SIZE"] = str(self._jvm.PythonUtils.getSparkBufferSize(self._jsc)) + + self.pythonExec = os.environ.get("PYSPARK_PYTHON", "python3") + self.pythonVer = "%d.%d" % sys.version_info[:2] + + # Broadcast's __reduce__ method stores Broadcast instances here. + # This allows other code to determine which Broadcast instances have + # been pickled, so it can determine which Java broadcast objects to + # send. + self._pickled_broadcast_vars = BroadcastPickleRegistry() + + SparkFiles._sc = self + root_dir = SparkFiles.getRootDirectory() + sys.path.insert(1, root_dir) + + # Deploy any code dependencies specified in the constructor + self._python_includes = list() + for path in pyFiles or []: + self.addPyFile(path) + + # Deploy code dependencies set by spark-submit; these will already have been added + # with SparkContext.addFile, so we just need to add them to the PYTHONPATH + for path in self._conf.get("spark.submit.pyFiles", "").split(","): + if path != "": + (dirname, filename) = os.path.split(path) + try: + filepath = os.path.join(SparkFiles.getRootDirectory(), filename) + if not os.path.exists(filepath): + # In case of YARN with shell mode, 'spark.submit.pyFiles' files are + # not added via SparkContext.addFile. Here we check if the file exists, + # try to copy and then add it to the path. See SPARK-21945. + shutil.copyfile(path, filepath) + if filename[-4:].lower() in self.PACKAGE_EXTENSIONS: + self._python_includes.append(filename) + sys.path.insert(1, filepath) + except Exception: + warnings.warn( + "Failed to add file [%s] specified in 'spark.submit.pyFiles' to " + "Python path:\n %s" % (path, "\n ".join(sys.path)), + RuntimeWarning, + ) + + # Create a temporary directory inside spark.local.dir: + assert self._jvm is not None + local_dir = self._jvm.org.apache.spark.util.Utils.getLocalDir(self._jsc.sc().conf()) + self._temp_dir = self._jvm.org.apache.spark.util.Utils.createTempDir( + local_dir, "pyspark" + ).getAbsolutePath() + + # profiling stats collected for each PythonRDD + if ( + self._conf.get("spark.python.profile", "false") == "true" + or self._conf.get("spark.python.profile.memory", "false") == "true" + ): + dump_path = self._conf.get("spark.python.profile.dump", None) + self.profiler_collector = ProfilerCollector( + profiler_cls, udf_profiler_cls, memory_profiler_cls, dump_path + ) + else: + self.profiler_collector = None # type: ignore[assignment] + + # create a signal handler which would be invoked on receiving SIGINT + def signal_handler(signal: Any, frame: Any) -> NoReturn: + self.cancelAllJobs() + raise KeyboardInterrupt() + + # see http://stackoverflow.com/questions/23206787/ + if isinstance( + threading.current_thread(), threading._MainThread # type: ignore[attr-defined] + ): + signal.signal(signal.SIGINT, signal_handler) + + def __repr__(self) -> str: + return "".format( + master=self.master, + appName=self.appName, + ) + + def _repr_html_(self) -> str: + return """ +
    +

    SparkContext

    + +

    Spark UI

    + +
    +
    Version
    +
    v{sc.version}
    +
    Master
    +
    {sc.master}
    +
    AppName
    +
    {sc.appName}
    +
    +
    + """.format( + sc=self + ) + + def _initialize_context(self, jconf: JavaObject) -> JavaObject: + """ + Initialize SparkContext in function to allow subclass specific initialization + """ + assert self._jvm is not None + return self._jvm.JavaSparkContext(jconf) + + @classmethod + def _ensure_initialized( + cls, + instance: Optional["SparkContext"] = None, + gateway: Optional[JavaGateway] = None, + conf: Optional[SparkConf] = None, + ) -> None: + """ + Checks whether a SparkContext is initialized or not. + Throws error if a SparkContext is already running. + """ + with SparkContext._lock: + if not SparkContext._gateway: + SparkContext._gateway = gateway or launch_gateway(conf) + SparkContext._jvm = SparkContext._gateway.jvm + + if instance: + if ( + SparkContext._active_spark_context + and SparkContext._active_spark_context != instance + ): + currentMaster = SparkContext._active_spark_context.master + currentAppName = SparkContext._active_spark_context.appName + callsite = SparkContext._active_spark_context._callsite + + # Raise error if there is already a running Spark context + raise ValueError( + "Cannot run multiple SparkContexts at once; " + "existing SparkContext(app=%s, master=%s)" + " created by %s at %s:%s " + % ( + currentAppName, + currentMaster, + callsite.function, + callsite.file, + callsite.linenum, + ) + ) + else: + SparkContext._active_spark_context = instance + + def __getnewargs__(self) -> NoReturn: + # This method is called when attempting to pickle SparkContext, which is always an error: + raise PySparkRuntimeError( + error_class="CONTEXT_ONLY_VALID_ON_DRIVER", + message_parameters={}, + ) + + def __enter__(self) -> "SparkContext": + """ + Enable 'with SparkContext(...) as sc: app(sc)' syntax. + """ + return self + + def __exit__( + self, + type: Optional[Type[BaseException]], + value: Optional[BaseException], + trace: Optional[TracebackType], + ) -> None: + """ + Enable 'with SparkContext(...) as sc: app' syntax. + + Specifically stop the context on exit of the with block. + """ + self.stop() + + @classmethod + def getOrCreate(cls, conf: Optional[SparkConf] = None) -> "SparkContext": + """ + Get or instantiate a :class:`SparkContext` and register it as a singleton object. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + conf : :class:`SparkConf`, optional + :class:`SparkConf` that will be used for initialization of the :class:`SparkContext`. + + Returns + ------- + :class:`SparkContext` + current :class:`SparkContext`, or a new one if it wasn't created before the function + call. + + Examples + -------- + >>> SparkContext.getOrCreate() + + """ + with SparkContext._lock: + if SparkContext._active_spark_context is None: + SparkContext(conf=conf or SparkConf()) + assert SparkContext._active_spark_context is not None + return SparkContext._active_spark_context + + def setLogLevel(self, logLevel: str) -> None: + """ + Control our logLevel. This overrides any user-defined log settings. + Valid log levels include: ALL, DEBUG, ERROR, FATAL, INFO, OFF, TRACE, WARN + + .. versionadded:: 1.4.0 + + Parameters + ---------- + logLevel : str + The desired log level as a string. + + Examples + -------- + >>> sc.setLogLevel("WARN") # doctest :+SKIP + """ + self._jsc.setLogLevel(logLevel) + + @classmethod + def setSystemProperty(cls, key: str, value: str) -> None: + """ + Set a Java system property, such as `spark.executor.memory`. This must + be invoked before instantiating :class:`SparkContext`. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + key : str + The key of a new Java system property. + value : str + The value of a new Java system property. + """ + SparkContext._ensure_initialized() + assert SparkContext._jvm is not None + SparkContext._jvm.java.lang.System.setProperty(key, value) + + @property + def version(self) -> str: + """ + The version of Spark on which this application is running. + + .. versionadded:: 1.1.0 + + Examples + -------- + >>> _ = sc.version + """ + return self._jsc.version() + + @property + def applicationId(self) -> str: + """ + A unique identifier for the Spark application. + Its format depends on the scheduler implementation. + + * in case of local spark app something like 'local-1433865536131' + * in case of YARN something like 'application_1433865536131_34483' + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> sc.applicationId # doctest: +ELLIPSIS + 'local-...' + """ + return self._jsc.sc().applicationId() + + @property + def uiWebUrl(self) -> Optional[str]: + """Return the URL of the SparkUI instance started by this :class:`SparkContext` + + .. versionadded:: 2.1.0 + + Notes + ----- + When the web ui is disabled, e.g., by ``spark.ui.enabled`` set to ``False``, + it returns ``None``. + + Examples + -------- + >>> sc.uiWebUrl + 'http://...' + """ + jurl = self._jsc.sc().uiWebUrl() + return jurl.get() if jurl.nonEmpty() else None + + @property + def startTime(self) -> int: + """Return the epoch time when the :class:`SparkContext` was started. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> _ = sc.startTime + """ + return self._jsc.startTime() + + @property + def defaultParallelism(self) -> int: + """ + Default level of parallelism to use when not given by user (e.g. for reduce tasks) + + .. versionadded:: 0.7.0 + + Examples + -------- + >>> sc.defaultParallelism > 0 + True + """ + return self._jsc.sc().defaultParallelism() + + @property + def defaultMinPartitions(self) -> int: + """ + Default min number of partitions for Hadoop RDDs when not given by user + + .. versionadded:: 1.1.0 + + Examples + -------- + >>> sc.defaultMinPartitions > 0 + True + """ + return self._jsc.sc().defaultMinPartitions() + + def stop(self) -> None: + """ + Shut down the :class:`SparkContext`. + + .. versionadded:: 0.7.0 + """ + if getattr(self, "_jsc", None): + try: + self._jsc.stop() + except Py4JError: + # Case: SPARK-18523 + warnings.warn( + "Unable to cleanly shutdown Spark JVM process." + " It is possible that the process has crashed," + " been killed or may also be in a zombie state.", + RuntimeWarning, + ) + finally: + self._jsc = None + if getattr(self, "_accumulatorServer", None): + self._accumulatorServer.shutdown() + self._accumulatorServer = None # type: ignore[assignment] + with SparkContext._lock: + SparkContext._active_spark_context = None + + def emptyRDD(self) -> RDD[Any]: + """ + Create an :class:`RDD` that has no partitions or elements. + + .. versionadded:: 1.5.0 + + Returns + ------- + :class:`RDD` + An empty RDD + + Examples + -------- + >>> sc.emptyRDD() + EmptyRDD... + >>> sc.emptyRDD().count() + 0 + """ + return RDD(self._jsc.emptyRDD(), self, NoOpSerializer()) + + def range( + self, start: int, end: Optional[int] = None, step: int = 1, numSlices: Optional[int] = None + ) -> RDD[int]: + """ + Create a new RDD of int containing elements from `start` to `end` + (exclusive), increased by `step` every element. Can be called the same + way as python's built-in range() function. If called with a single argument, + the argument is interpreted as `end`, and `start` is set to 0. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + start : int + the start value + end : int, optional + the end value (exclusive) + step : int, optional, default 1 + the incremental step + numSlices : int, optional + the number of partitions of the new RDD + + Returns + ------- + :class:`RDD` + An RDD of int + + See Also + -------- + :meth:`pyspark.sql.SparkSession.range` + + Examples + -------- + >>> sc.range(5).collect() + [0, 1, 2, 3, 4] + >>> sc.range(2, 4).collect() + [2, 3] + >>> sc.range(1, 7, 2).collect() + [1, 3, 5] + + Generate RDD with a negative step + + >>> sc.range(5, 0, -1).collect() + [5, 4, 3, 2, 1] + >>> sc.range(0, 5, -1).collect() + [] + + Control the number of partitions + + >>> sc.range(5, numSlices=1).getNumPartitions() + 1 + >>> sc.range(5, numSlices=10).getNumPartitions() + 10 + """ + if end is None: + end = start + start = 0 + + return self.parallelize(range(start, end, step), numSlices) + + def parallelize(self, c: Iterable[T], numSlices: Optional[int] = None) -> RDD[T]: + """ + Distribute a local Python collection to form an RDD. Using range + is recommended if the input represents a range for performance. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + c : :class:`collections.abc.Iterable` + iterable collection to distribute + numSlices : int, optional + the number of partitions of the new RDD + + Returns + ------- + :class:`RDD` + RDD representing distributed collection. + + Examples + -------- + >>> sc.parallelize([0, 2, 3, 4, 6], 5).glom().collect() + [[0], [2], [3], [4], [6]] + >>> sc.parallelize(range(0, 6, 2), 5).glom().collect() + [[], [0], [], [2], [4]] + + Deal with a list of strings. + + >>> strings = ["a", "b", "c"] + >>> sc.parallelize(strings, 2).glom().collect() + [['a'], ['b', 'c']] + """ + numSlices = int(numSlices) if numSlices is not None else self.defaultParallelism + if isinstance(c, range): + size = len(c) + if size == 0: + return self.parallelize([], numSlices) + step = c[1] - c[0] if size > 1 else 1 # type: ignore[index] + start0 = c[0] # type: ignore[index] + + def getStart(split: int) -> int: + assert numSlices is not None + return start0 + int((split * size / numSlices)) * step + + def f(split: int, iterator: Iterable[T]) -> Iterable: + # it's an empty iterator here but we need this line for triggering the + # logic of signal handling in FramedSerializer.load_stream, for instance, + # SpecialLengths.END_OF_DATA_SECTION in _read_with_length. Since + # FramedSerializer.load_stream produces a generator, the control should + # at least be in that function once. Here we do it by explicitly converting + # the empty iterator to a list, thus make sure worker reuse takes effect. + # See more details in SPARK-26549. + assert len(list(iterator)) == 0 + return range(getStart(split), getStart(split + 1), step) + + return self.parallelize([], numSlices).mapPartitionsWithIndex(f) + + # Make sure we distribute data evenly if it's smaller than self.batchSize + if "__len__" not in dir(c): + c = list(c) # Make it a list so we can compute its length + batchSize = max( + 1, min(len(c) // numSlices, self._batchSize or 1024) # type: ignore[arg-type] + ) + serializer = BatchedSerializer(self._unbatched_serializer, batchSize) + + def reader_func(temp_filename: str) -> JavaObject: + assert self._jvm is not None + return self._jvm.PythonRDD.readRDDFromFile(self._jsc, temp_filename, numSlices) + + def createRDDServer() -> JavaObject: + assert self._jvm is not None + return self._jvm.PythonParallelizeServer(self._jsc.sc(), numSlices) + + jrdd = self._serialize_to_jvm(c, serializer, reader_func, createRDDServer) + return RDD(jrdd, self, serializer) + + def _serialize_to_jvm( + self, + data: Iterable[T], + serializer: Serializer, + reader_func: Callable, + server_func: Callable, + ) -> JavaObject: + """ + Using Py4J to send a large dataset to the jvm is slow, so we use either a file + or a socket if we have encryption enabled. + + Examples + -------- + data + object to be serialized + serializer : class:`pyspark.serializers.Serializer` + reader_func : function + A function which takes a filename and reads in the data in the jvm and + returns a JavaRDD. Only used when encryption is disabled. + server_func : function + A function which creates a SocketAuthServer in the JVM to + accept the serialized data, for use when encryption is enabled. + """ + if self._encryption_enabled: + # with encryption, we open a server in java and send the data directly + server = server_func() + (sock_file, _) = local_connect_and_auth(server.port(), server.secret()) + chunked_out = ChunkedStream(sock_file, 8192) + serializer.dump_stream(data, chunked_out) + chunked_out.close() + # this call will block until the server has read all the data and processed it (or + # throws an exception) + r = server.getResult() + return r + else: + # without encryption, we serialize to a file, and we read the file in java and + # parallelize from there. + tempFile = NamedTemporaryFile(delete=False, dir=self._temp_dir) + try: + try: + serializer.dump_stream(data, tempFile) + finally: + tempFile.close() + return reader_func(tempFile.name) + finally: + # we eagerly reads the file so we can delete right after. + os.unlink(tempFile.name) + + def pickleFile(self, name: str, minPartitions: Optional[int] = None) -> RDD[Any]: + """ + Load an RDD previously saved using :meth:`RDD.saveAsPickleFile` method. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + name : str + directory to the input data files, the path can be comma separated + paths as a list of inputs + minPartitions : int, optional + suggested minimum number of partitions for the resulting RDD + + Returns + ------- + :class:`RDD` + RDD representing unpickled data from the file(s). + + See Also + -------- + :meth:`RDD.saveAsPickleFile` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary pickled file + ... path1 = os.path.join(d, "pickled1") + ... sc.parallelize(range(10)).saveAsPickleFile(path1, 3) + ... + ... # Write another temporary pickled file + ... path2 = os.path.join(d, "pickled2") + ... sc.parallelize(range(-10, -5)).saveAsPickleFile(path2, 3) + ... + ... # Load picked file + ... collected1 = sorted(sc.pickleFile(path1, 3).collect()) + ... collected2 = sorted(sc.pickleFile(path2, 4).collect()) + ... + ... # Load two picked files together + ... collected3 = sorted(sc.pickleFile('{},{}'.format(path1, path2), 5).collect()) + + >>> collected1 + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + >>> collected2 + [-10, -9, -8, -7, -6] + >>> collected3 + [-10, -9, -8, -7, -6, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + """ + minPartitions = minPartitions or self.defaultMinPartitions + return RDD(self._jsc.objectFile(name, minPartitions), self) + + def textFile( + self, name: str, minPartitions: Optional[int] = None, use_unicode: bool = True + ) -> RDD[str]: + """ + Read a text file from HDFS, a local file system (available on all + nodes), or any Hadoop-supported file system URI, and return it as an + RDD of Strings. The text files must be encoded as UTF-8. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + name : str + directory to the input data files, the path can be comma separated + paths as a list of inputs + minPartitions : int, optional + suggested minimum number of partitions for the resulting RDD + use_unicode : bool, default True + If `use_unicode` is False, the strings will be kept as `str` (encoding + as `utf-8`), which is faster and smaller than unicode. + + .. versionadded:: 1.2.0 + + Returns + ------- + :class:`RDD` + RDD representing text data from the file(s). + + See Also + -------- + :meth:`RDD.saveAsTextFile` + :meth:`SparkContext.wholeTextFiles` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... path1 = os.path.join(d, "text1") + ... path2 = os.path.join(d, "text2") + ... + ... # Write a temporary text file + ... sc.parallelize(["x", "y", "z"]).saveAsTextFile(path1) + ... + ... # Write another temporary text file + ... sc.parallelize(["aa", "bb", "cc"]).saveAsTextFile(path2) + ... + ... # Load text file + ... collected1 = sorted(sc.textFile(path1, 3).collect()) + ... collected2 = sorted(sc.textFile(path2, 4).collect()) + ... + ... # Load two text files together + ... collected3 = sorted(sc.textFile('{},{}'.format(path1, path2), 5).collect()) + + >>> collected1 + ['x', 'y', 'z'] + >>> collected2 + ['aa', 'bb', 'cc'] + >>> collected3 + ['aa', 'bb', 'cc', 'x', 'y', 'z'] + """ + minPartitions = minPartitions or min(self.defaultParallelism, 2) + return RDD(self._jsc.textFile(name, minPartitions), self, UTF8Deserializer(use_unicode)) + + def wholeTextFiles( + self, path: str, minPartitions: Optional[int] = None, use_unicode: bool = True + ) -> RDD[Tuple[str, str]]: + """ + Read a directory of text files from HDFS, a local file system + (available on all nodes), or any Hadoop-supported file system + URI. Each file is read as a single record and returned in a + key-value pair, where the key is the path of each file, the + value is the content of each file. + The text files must be encoded as UTF-8. + + .. versionadded:: 1.0.0 + + For example, if you have the following files: + + .. code-block:: text + + hdfs://a-hdfs-path/part-00000 + hdfs://a-hdfs-path/part-00001 + ... + hdfs://a-hdfs-path/part-nnnnn + + Do ``rdd = sparkContext.wholeTextFiles("hdfs://a-hdfs-path")``, + then ``rdd`` contains: + + .. code-block:: text + + (a-hdfs-path/part-00000, its content) + (a-hdfs-path/part-00001, its content) + ... + (a-hdfs-path/part-nnnnn, its content) + + Parameters + ---------- + path : str + directory to the input data files, the path can be comma separated + paths as a list of inputs + minPartitions : int, optional + suggested minimum number of partitions for the resulting RDD + use_unicode : bool, default True + If `use_unicode` is False, the strings will be kept as `str` (encoding + as `utf-8`), which is faster and smaller than unicode. + + .. versionadded:: 1.2.0 + + Returns + ------- + :class:`RDD` + RDD representing path-content pairs from the file(s). + + Notes + ----- + Small files are preferred, as each file will be loaded fully in memory. + + See Also + -------- + :meth:`RDD.saveAsTextFile` + :meth:`SparkContext.textFile` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary text file + ... with open(os.path.join(d, "1.txt"), "w") as f: + ... _ = f.write("123") + ... + ... # Write another temporary text file + ... with open(os.path.join(d, "2.txt"), "w") as f: + ... _ = f.write("xyz") + ... + ... collected = sorted(sc.wholeTextFiles(d).collect()) + >>> collected + [('.../1.txt', '123'), ('.../2.txt', 'xyz')] + """ + minPartitions = minPartitions or self.defaultMinPartitions + return RDD( + self._jsc.wholeTextFiles(path, minPartitions), + self, + PairDeserializer(UTF8Deserializer(use_unicode), UTF8Deserializer(use_unicode)), + ) + + def binaryFiles(self, path: str, minPartitions: Optional[int] = None) -> RDD[Tuple[str, bytes]]: + """ + Read a directory of binary files from HDFS, a local file system + (available on all nodes), or any Hadoop-supported file system URI + as a byte array. Each file is read as a single record and returned + in a key-value pair, where the key is the path of each file, the + value is the content of each file. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + path : str + directory to the input data files, the path can be comma separated + paths as a list of inputs + minPartitions : int, optional + suggested minimum number of partitions for the resulting RDD + + Returns + ------- + :class:`RDD` + RDD representing path-content pairs from the file(s). + + Notes + ----- + Small files are preferred, large file is also allowable, but may cause bad performance. + + See Also + -------- + :meth:`SparkContext.binaryRecords` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary binary file + ... with open(os.path.join(d, "1.bin"), "wb") as f1: + ... _ = f1.write(b"binary data I") + ... + ... # Write another temporary binary file + ... with open(os.path.join(d, "2.bin"), "wb") as f2: + ... _ = f2.write(b"binary data II") + ... + ... collected = sorted(sc.binaryFiles(d).collect()) + + >>> collected + [('.../1.bin', b'binary data I'), ('.../2.bin', b'binary data II')] + """ + minPartitions = minPartitions or self.defaultMinPartitions + return RDD( + self._jsc.binaryFiles(path, minPartitions), + self, + PairDeserializer(UTF8Deserializer(), NoOpSerializer()), + ) + + def binaryRecords(self, path: str, recordLength: int) -> RDD[bytes]: + """ + Load data from a flat binary file, assuming each record is a set of numbers + with the specified numerical format (see ByteBuffer), and the number of + bytes per record is constant. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + path : str + Directory to the input data files + recordLength : int + The length at which to split the records + + Returns + ------- + :class:`RDD` + RDD of data with values, represented as byte arrays + + See Also + -------- + :meth:`SparkContext.binaryFiles` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary file + ... with open(os.path.join(d, "1.bin"), "w") as f: + ... for i in range(3): + ... _ = f.write("%04d" % i) + ... + ... # Write another file + ... with open(os.path.join(d, "2.bin"), "w") as f: + ... for i in [-1, -2, -10]: + ... _ = f.write("%04d" % i) + ... + ... collected = sorted(sc.binaryRecords(d, 4).collect()) + + >>> collected + [b'-001', b'-002', b'-010', b'0000', b'0001', b'0002'] + """ + return RDD(self._jsc.binaryRecords(path, recordLength), self, NoOpSerializer()) + + def _dictToJavaMap(self, d: Optional[Dict[str, str]]) -> JavaMap: + assert self._jvm is not None + jm = self._jvm.java.util.HashMap() + if not d: + d = {} + for k, v in d.items(): + jm[k] = v + return jm + + def sequenceFile( + self, + path: str, + keyClass: Optional[str] = None, + valueClass: Optional[str] = None, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + minSplits: Optional[int] = None, + batchSize: int = 0, + ) -> RDD[Tuple[T, U]]: + """ + Read a Hadoop SequenceFile with arbitrary key and value Writable class from HDFS, + a local file system (available on all nodes), or any Hadoop-supported file system URI. + The mechanism is as follows: + + 1. A Java RDD is created from the SequenceFile or other InputFormat, and the key + and value Writable classes + 2. Serialization is attempted via Pickle pickling + 3. If this fails, the fallback is to call 'toString' on each key and value + 4. :class:`CPickleSerializer` is used to deserialize pickled objects on the Python side + + .. versionadded:: 1.3.0 + + Parameters + ---------- + path : str + path to sequencefile + keyClass: str, optional + fully qualified classname of key Writable class (e.g. "org.apache.hadoop.io.Text") + valueClass : str, optional + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.LongWritable") + keyConverter : str, optional + fully qualified name of a function returning key WritableConverter + valueConverter : str, optional + fully qualifiedname of a function returning value WritableConverter + minSplits : int, optional + minimum splits in dataset (default min(2, sc.defaultParallelism)) + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. (default 0, choose batchSize automatically) + + Returns + ------- + :class:`RDD` + RDD of tuples of key and corresponding value + + See Also + -------- + :meth:`RDD.saveAsSequenceFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsHadoopFile` + :meth:`SparkContext.newAPIHadoopFile` + :meth:`SparkContext.hadoopFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the class of output format + + >>> output_format_class = "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "hadoop_file") + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, {3.0: "bb"}), (2, {1.0: "aa"}), (3, {2.0: "dd"})]) + ... rdd.saveAsNewAPIHadoopFile(path, output_format_class) + ... + ... collected = sorted(sc.sequenceFile(path).collect()) + + >>> collected + [(1, {3.0: 'bb'}), (2, {1.0: 'aa'}), (3, {2.0: 'dd'})] + """ + minSplits = minSplits or min(self.defaultParallelism, 2) + assert self._jvm is not None + jrdd = self._jvm.PythonRDD.sequenceFile( + self._jsc, + path, + keyClass, + valueClass, + keyConverter, + valueConverter, + minSplits, + batchSize, + ) + return RDD(jrdd, self) + + def newAPIHadoopFile( + self, + path: str, + inputFormatClass: str, + keyClass: str, + valueClass: str, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + batchSize: int = 0, + ) -> RDD[Tuple[T, U]]: + """ + Read a 'new API' Hadoop InputFormat with arbitrary key and value class from HDFS, + a local file system (available on all nodes), or any Hadoop-supported file system URI. + The mechanism is the same as for meth:`SparkContext.sequenceFile`. + + A Hadoop configuration can be passed in as a Python dict. This will be converted into a + Configuration in Java + + .. versionadded:: 1.1.0 + + Parameters + ---------- + path : str + path to Hadoop file + inputFormatClass : str + fully qualified classname of Hadoop InputFormat + (e.g. "org.apache.hadoop.mapreduce.lib.input.TextInputFormat") + keyClass : str + fully qualified classname of key Writable class + (e.g. "org.apache.hadoop.io.Text") + valueClass : str + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.LongWritable") + keyConverter : str, optional + fully qualified name of a function returning key WritableConverter + None by default + valueConverter : str, optional + fully qualified name of a function returning value WritableConverter + None by default + conf : dict, optional + Hadoop configuration, passed in as a dict + None by default + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. (default 0, choose batchSize automatically) + + Returns + ------- + :class:`RDD` + RDD of tuples of key and corresponding value + + See Also + -------- + :meth:`RDD.saveAsSequenceFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsHadoopFile` + :meth:`SparkContext.sequenceFile` + :meth:`SparkContext.hadoopFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "new_hadoop_file") + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsNewAPIHadoopFile(path, output_format_class, key_class, value_class) + ... + ... loaded = sc.newAPIHadoopFile(path, input_format_class, key_class, value_class) + ... collected = sorted(loaded.collect()) + + >>> collected + [(1, ''), (1, 'a'), (3, 'x')] + """ + jconf = self._dictToJavaMap(conf) + assert self._jvm is not None + jrdd = self._jvm.PythonRDD.newAPIHadoopFile( + self._jsc, + path, + inputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + batchSize, + ) + return RDD(jrdd, self) + + def newAPIHadoopRDD( + self, + inputFormatClass: str, + keyClass: str, + valueClass: str, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + batchSize: int = 0, + ) -> RDD[Tuple[T, U]]: + """ + Read a 'new API' Hadoop InputFormat with arbitrary key and value class, from an arbitrary + Hadoop configuration, which is passed in as a Python dict. + This will be converted into a Configuration in Java. + The mechanism is the same as for meth:`SparkContext.sequenceFile`. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + inputFormatClass : str + fully qualified classname of Hadoop InputFormat + (e.g. "org.apache.hadoop.mapreduce.lib.input.TextInputFormat") + keyClass : str + fully qualified classname of key Writable class (e.g. "org.apache.hadoop.io.Text") + valueClass : str + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.LongWritable") + keyConverter : str, optional + fully qualified name of a function returning key WritableConverter + (None by default) + valueConverter : str, optional + fully qualified name of a function returning value WritableConverter + (None by default) + conf : dict, optional + Hadoop configuration, passed in as a dict (None by default) + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. (default 0, choose batchSize automatically) + + Returns + ------- + :class:`RDD` + RDD of tuples of key and corresponding value + + See Also + -------- + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsHadoopDataset` + :meth:`SparkContext.hadoopRDD` + :meth:`SparkContext.hadoopFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "new_hadoop_file") + ... + ... # Create the conf for writing + ... write_conf = { + ... "mapreduce.job.outputformat.class": (output_format_class), + ... "mapreduce.job.output.key.class": key_class, + ... "mapreduce.job.output.value.class": value_class, + ... "mapreduce.output.fileoutputformat.outputdir": path, + ... } + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsNewAPIHadoopDataset(conf=write_conf) + ... + ... # Create the conf for reading + ... read_conf = {"mapreduce.input.fileinputformat.inputdir": path} + ... + ... loaded = sc.newAPIHadoopRDD(input_format_class, + ... key_class, value_class, conf=read_conf) + ... collected = sorted(loaded.collect()) + + >>> collected + [(1, ''), (1, 'a'), (3, 'x')] + """ + jconf = self._dictToJavaMap(conf) + assert self._jvm is not None + jrdd = self._jvm.PythonRDD.newAPIHadoopRDD( + self._jsc, + inputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + batchSize, + ) + return RDD(jrdd, self) + + def hadoopFile( + self, + path: str, + inputFormatClass: str, + keyClass: str, + valueClass: str, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + batchSize: int = 0, + ) -> RDD[Tuple[T, U]]: + """ + Read an 'old' Hadoop InputFormat with arbitrary key and value class from HDFS, + a local file system (available on all nodes), or any Hadoop-supported file system URI. + The mechanism is the same as for meth:`SparkContext.sequenceFile`. + + .. versionadded:: 1.1.0 + + A Hadoop configuration can be passed in as a Python dict. This will be converted into a + Configuration in Java. + + Parameters + ---------- + path : str + path to Hadoop file + inputFormatClass : str + fully qualified classname of Hadoop InputFormat + (e.g. "org.apache.hadoop.mapreduce.lib.input.TextInputFormat") + keyClass : str + fully qualified classname of key Writable class (e.g. "org.apache.hadoop.io.Text") + valueClass : str + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.LongWritable") + keyConverter : str, optional + fully qualified name of a function returning key WritableConverter + valueConverter : str, optional + fully qualified name of a function returning value WritableConverter + conf : dict, optional + Hadoop configuration, passed in as a dict + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. (default 0, choose batchSize automatically) + + Returns + ------- + :class:`RDD` + RDD of tuples of key and corresponding value + + See Also + -------- + :meth:`RDD.saveAsSequenceFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsHadoopFile` + :meth:`SparkContext.newAPIHadoopFile` + :meth:`SparkContext.hadoopRDD` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapred.TextOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapred.TextInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "old_hadoop_file") + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsHadoopFile(path, output_format_class, key_class, value_class) + ... + ... loaded = sc.hadoopFile(path, input_format_class, key_class, value_class) + ... collected = sorted(loaded.collect()) + + >>> collected + [(0, '1\\t'), (0, '1\\ta'), (0, '3\\tx')] + """ + jconf = self._dictToJavaMap(conf) + assert self._jvm is not None + jrdd = self._jvm.PythonRDD.hadoopFile( + self._jsc, + path, + inputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + batchSize, + ) + return RDD(jrdd, self) + + def hadoopRDD( + self, + inputFormatClass: str, + keyClass: str, + valueClass: str, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + batchSize: int = 0, + ) -> RDD[Tuple[T, U]]: + """ + Read an 'old' Hadoop InputFormat with arbitrary key and value class, from an arbitrary + Hadoop configuration, which is passed in as a Python dict. + This will be converted into a Configuration in Java. + The mechanism is the same as for meth:`SparkContext.sequenceFile`. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + inputFormatClass : str + fully qualified classname of Hadoop InputFormat + (e.g. "org.apache.hadoop.mapreduce.lib.input.TextInputFormat") + keyClass : str + fully qualified classname of key Writable class (e.g. "org.apache.hadoop.io.Text") + valueClass : str + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.LongWritable") + keyConverter : str, optional + fully qualified name of a function returning key WritableConverter + valueConverter : str, optional + fully qualified name of a function returning value WritableConverter + conf : dict, optional + Hadoop configuration, passed in as a dict + batchSize : int, optional, default 0 + The number of Python objects represented as a single + Java object. (default 0, choose batchSize automatically) + + Returns + ------- + :class:`RDD` + RDD of tuples of key and corresponding value + + See Also + -------- + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsHadoopDataset` + :meth:`SparkContext.newAPIHadoopRDD` + :meth:`SparkContext.hadoopFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapred.TextOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapred.TextInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "old_hadoop_file") + ... + ... # Create the conf for writing + ... write_conf = { + ... "mapred.output.format.class": output_format_class, + ... "mapreduce.job.output.key.class": key_class, + ... "mapreduce.job.output.value.class": value_class, + ... "mapreduce.output.fileoutputformat.outputdir": path, + ... } + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsHadoopDataset(conf=write_conf) + ... + ... # Create the conf for reading + ... read_conf = {"mapreduce.input.fileinputformat.inputdir": path} + ... + ... loaded = sc.hadoopRDD(input_format_class, key_class, value_class, conf=read_conf) + ... collected = sorted(loaded.collect()) + + >>> collected + [(0, '1\\t'), (0, '1\\ta'), (0, '3\\tx')] + """ + jconf = self._dictToJavaMap(conf) + assert self._jvm is not None + jrdd = self._jvm.PythonRDD.hadoopRDD( + self._jsc, + inputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + batchSize, + ) + return RDD(jrdd, self) + + def _checkpointFile(self, name: str, input_deserializer: PairDeserializer) -> RDD: + jrdd = self._jsc.checkpointFile(name) + return RDD(jrdd, self, input_deserializer) + + def union(self, rdds: List[RDD[T]]) -> RDD[T]: + """ + Build the union of a list of RDDs. + + This supports unions() of RDDs with different serialized formats, + although this forces them to be reserialized using the default + serializer: + + .. versionadded:: 0.7.0 + + See Also + -------- + :meth:`RDD.union` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # generate a text RDD + ... with open(os.path.join(d, "union-text.txt"), "w") as f: + ... _ = f.write("Hello") + ... text_rdd = sc.textFile(d) + ... + ... # generate another RDD + ... parallelized = sc.parallelize(["World!"]) + ... + ... unioned = sorted(sc.union([text_rdd, parallelized]).collect()) + + >>> unioned + ['Hello', 'World!'] + """ + first_jrdd_deserializer = rdds[0]._jrdd_deserializer + if any(x._jrdd_deserializer != first_jrdd_deserializer for x in rdds): + rdds = [x._reserialize() for x in rdds] + gw = SparkContext._gateway + assert gw is not None + jvm = SparkContext._jvm + assert jvm is not None + jrdd_cls = jvm.org.apache.spark.api.java.JavaRDD + jpair_rdd_cls = jvm.org.apache.spark.api.java.JavaPairRDD + jdouble_rdd_cls = jvm.org.apache.spark.api.java.JavaDoubleRDD + if is_instance_of(gw, rdds[0]._jrdd, jrdd_cls): + cls = jrdd_cls + elif is_instance_of(gw, rdds[0]._jrdd, jpair_rdd_cls): + cls = jpair_rdd_cls + elif is_instance_of(gw, rdds[0]._jrdd, jdouble_rdd_cls): + cls = jdouble_rdd_cls + else: + cls_name = rdds[0]._jrdd.getClass().getCanonicalName() + raise TypeError("Unsupported Java RDD class %s" % cls_name) + jrdds = gw.new_array(cls, len(rdds)) + for i in range(0, len(rdds)): + jrdds[i] = rdds[i]._jrdd + return RDD(self._jsc.union(jrdds), self, rdds[0]._jrdd_deserializer) + + def broadcast(self, value: T) -> "Broadcast[T]": + """ + Broadcast a read-only variable to the cluster, returning a :class:`Broadcast` + object for reading it in distributed functions. The variable will + be sent to each cluster only once. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + value : T + value to broadcast to the Spark nodes + + Returns + ------- + :class:`Broadcast` + :class:`Broadcast` object, a read-only variable cached on each machine + + Examples + -------- + >>> mapping = {1: 10001, 2: 10002} + >>> bc = sc.broadcast(mapping) + + >>> rdd = sc.range(5) + >>> rdd2 = rdd.map(lambda i: bc.value[i] if i in bc.value else -1) + >>> rdd2.collect() + [-1, 10001, 10002, -1, -1] + + >>> bc.destroy() + """ + return Broadcast(self, value, self._pickled_broadcast_vars) + + def accumulator( + self, value: T, accum_param: Optional["AccumulatorParam[T]"] = None + ) -> "Accumulator[T]": + """ + Create an :class:`Accumulator` with the given initial value, using a given + :class:`AccumulatorParam` helper object to define how to add values of the + data type if provided. Default AccumulatorParams are used for integers + and floating-point numbers if you do not provide one. For other types, + a custom AccumulatorParam can be used. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + value : T + initialized value + accum_param : :class:`pyspark.AccumulatorParam`, optional + helper object to define how to add values + + Returns + ------- + :class:`Accumulator` + `Accumulator` object, a shared variable that can be accumulated + + Examples + -------- + >>> acc = sc.accumulator(9) + >>> acc.value + 9 + >>> acc += 1 + >>> acc.value + 10 + + Accumulator object can be accumulated in RDD operations: + + >>> rdd = sc.range(5) + >>> def f(x): + ... global acc + ... acc += 1 + ... + >>> rdd.foreach(f) + >>> acc.value + 15 + """ + if accum_param is None: + if isinstance(value, int): + accum_param = cast("AccumulatorParam[T]", accumulators.INT_ACCUMULATOR_PARAM) + elif isinstance(value, float): + accum_param = cast("AccumulatorParam[T]", accumulators.FLOAT_ACCUMULATOR_PARAM) + elif isinstance(value, complex): + accum_param = cast("AccumulatorParam[T]", accumulators.COMPLEX_ACCUMULATOR_PARAM) + else: + raise TypeError("No default accumulator param for type %s" % type(value)) + SparkContext._next_accum_id += 1 + return Accumulator(SparkContext._next_accum_id - 1, value, accum_param) + + def addFile(self, path: str, recursive: bool = False) -> None: + """ + Add a file to be downloaded with this Spark job on every node. + The `path` passed can be either a local file, a file in HDFS + (or other Hadoop-supported filesystems), or an HTTP, HTTPS or + FTP URI. + + To access the file in Spark jobs, use :meth:`SparkFiles.get` with the + filename to find its download location. + + A directory can be given if the recursive option is set to True. + Currently directories are only supported for Hadoop-supported filesystems. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + path : str + can be either a local file, a file in HDFS (or other Hadoop-supported + filesystems), or an HTTP, HTTPS or FTP URI. To access the file in Spark jobs, + use :meth:`SparkFiles.get` to find its download location. + recursive : bool, default False + whether to recursively add files in the input directory + + See Also + -------- + :meth:`SparkContext.listFiles` + :meth:`SparkContext.addPyFile` + :meth:`SparkFiles.get` + + Notes + ----- + A path can be added only once. Subsequent additions of the same path are ignored. + + Examples + -------- + >>> import os + >>> import tempfile + >>> from pyspark import SparkFiles + + >>> with tempfile.TemporaryDirectory() as d: + ... path1 = os.path.join(d, "test1.txt") + ... with open(path1, "w") as f: + ... _ = f.write("100") + ... + ... path2 = os.path.join(d, "test2.txt") + ... with open(path2, "w") as f: + ... _ = f.write("200") + ... + ... sc.addFile(path1) + ... file_list1 = sorted(sc.listFiles) + ... + ... sc.addFile(path2) + ... file_list2 = sorted(sc.listFiles) + ... + ... # add path2 twice, this addition will be ignored + ... sc.addFile(path2) + ... file_list3 = sorted(sc.listFiles) + ... + ... def func(iterator): + ... with open(SparkFiles.get("test1.txt")) as f: + ... mul = int(f.readline()) + ... return [x * mul for x in iterator] + ... + ... collected = sc.parallelize([1, 2, 3, 4]).mapPartitions(func).collect() + + >>> file_list1 + ['file:/.../test1.txt'] + >>> file_list2 + ['file:/.../test1.txt', 'file:/.../test2.txt'] + >>> file_list3 + ['file:/.../test1.txt', 'file:/.../test2.txt'] + >>> collected + [100, 200, 300, 400] + """ + self._jsc.sc().addFile(path, recursive) + + @property + def listFiles(self) -> List[str]: + """Returns a list of file paths that are added to resources. + + .. versionadded:: 3.4.0 + + See Also + -------- + :meth:`SparkContext.addFile` + """ + return list( + self._jvm.scala.collection.JavaConverters.seqAsJavaList( # type: ignore[union-attr] + self._jsc.sc().listFiles() + ) + ) + + def addPyFile(self, path: str) -> None: + """ + Add a .py or .zip dependency for all tasks to be executed on this + SparkContext in the future. The `path` passed can be either a local + file, a file in HDFS (or other Hadoop-supported filesystems), or an + HTTP, HTTPS or FTP URI. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + path : str + can be either a .py file or .zip dependency. + + See Also + -------- + :meth:`SparkContext.addFile` + + Notes + ----- + A path can be added only once. Subsequent additions of the same path are ignored. + """ + self.addFile(path) + (dirname, filename) = os.path.split(path) # dirname may be directory or HDFS/S3 prefix + if filename[-4:].lower() in self.PACKAGE_EXTENSIONS: + assert self._python_includes is not None + self._python_includes.append(filename) + # for tests in local mode + sys.path.insert(1, os.path.join(SparkFiles.getRootDirectory(), filename)) + + importlib.invalidate_caches() + + def addArchive(self, path: str) -> None: + """ + Add an archive to be downloaded with this Spark job on every node. + The `path` passed can be either a local file, a file in HDFS + (or other Hadoop-supported filesystems), or an HTTP, HTTPS or + FTP URI. + + To access the file in Spark jobs, use :meth:`SparkFiles.get` with the + filename to find its download/unpacked location. The given path should + be one of .zip, .tar, .tar.gz, .tgz and .jar. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + path : str + can be either a local file, a file in HDFS (or other Hadoop-supported + filesystems), or an HTTP, HTTPS or FTP URI. To access the file in Spark jobs, + use :meth:`SparkFiles.get` to find its download location. + + See Also + -------- + :meth:`SparkContext.listArchives` + :meth:`SparkFiles.get` + + Notes + ----- + A path can be added only once. Subsequent additions of the same path are ignored. + This API is experimental. + + Examples + -------- + Creates a zipped file that contains a text file written '100'. + + >>> import os + >>> import tempfile + >>> import zipfile + >>> from pyspark import SparkFiles + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "test.txt") + ... with open(path, "w") as f: + ... _ = f.write("100") + ... + ... zip_path1 = os.path.join(d, "test1.zip") + ... with zipfile.ZipFile(zip_path1, "w", zipfile.ZIP_DEFLATED) as z: + ... z.write(path, os.path.basename(path)) + ... + ... zip_path2 = os.path.join(d, "test2.zip") + ... with zipfile.ZipFile(zip_path2, "w", zipfile.ZIP_DEFLATED) as z: + ... z.write(path, os.path.basename(path)) + ... + ... sc.addArchive(zip_path1) + ... arch_list1 = sorted(sc.listArchives) + ... + ... sc.addArchive(zip_path2) + ... arch_list2 = sorted(sc.listArchives) + ... + ... # add zip_path2 twice, this addition will be ignored + ... sc.addArchive(zip_path2) + ... arch_list3 = sorted(sc.listArchives) + ... + ... def func(iterator): + ... with open("%s/test.txt" % SparkFiles.get("test1.zip")) as f: + ... mul = int(f.readline()) + ... return [x * mul for x in iterator] + ... + ... collected = sc.parallelize([1, 2, 3, 4]).mapPartitions(func).collect() + + >>> arch_list1 + ['file:/.../test1.zip'] + >>> arch_list2 + ['file:/.../test1.zip', 'file:/.../test2.zip'] + >>> arch_list3 + ['file:/.../test1.zip', 'file:/.../test2.zip'] + >>> collected + [100, 200, 300, 400] + """ + self._jsc.sc().addArchive(path) + + @property + def listArchives(self) -> List[str]: + """Returns a list of archive paths that are added to resources. + + .. versionadded:: 3.4.0 + + See Also + -------- + :meth:`SparkContext.addArchive` + """ + return list( + self._jvm.scala.collection.JavaConverters.seqAsJavaList( # type: ignore[union-attr] + self._jsc.sc().listArchives() + ) + ) + + def setCheckpointDir(self, dirName: str) -> None: + """ + Set the directory under which RDDs are going to be checkpointed. The + directory must be an HDFS path if running on a cluster. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + dirName : str + path to the directory where checkpoint files will be stored + (must be HDFS path if running in cluster) + + See Also + -------- + :meth:`SparkContext.getCheckpointDir` + :meth:`RDD.checkpoint` + :meth:`RDD.getCheckpointFile` + """ + self._jsc.sc().setCheckpointDir(dirName) + + def getCheckpointDir(self) -> Optional[str]: + """ + Return the directory where RDDs are checkpointed. Returns None if no + checkpoint directory has been set. + + .. versionadded:: 3.1.0 + + See Also + -------- + :meth:`SparkContext.setCheckpointDir` + :meth:`RDD.checkpoint` + :meth:`RDD.getCheckpointFile` + """ + if not self._jsc.sc().getCheckpointDir().isEmpty(): + return self._jsc.sc().getCheckpointDir().get() + return None + + def _getJavaStorageLevel(self, storageLevel: StorageLevel) -> JavaObject: + """ + Returns a Java StorageLevel based on a pyspark.StorageLevel. + """ + if not isinstance(storageLevel, StorageLevel): + raise TypeError("storageLevel must be of type pyspark.StorageLevel") + assert self._jvm is not None + newStorageLevel = self._jvm.org.apache.spark.storage.StorageLevel + return newStorageLevel( + storageLevel.useDisk, + storageLevel.useMemory, + storageLevel.useOffHeap, + storageLevel.deserialized, + storageLevel.replication, + ) + + def setJobGroup(self, groupId: str, description: str, interruptOnCancel: bool = False) -> None: + """ + Assigns a group ID to all the jobs started by this thread until the group ID is set to a + different value or cleared. + + Often, a unit of execution in an application consists of multiple Spark actions or jobs. + Application programmers can use this method to group all those jobs together and give a + group description. Once set, the Spark web UI will associate such jobs with this group. + + The application can use :meth:`SparkContext.cancelJobGroup` to cancel all + running jobs in this group. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + groupId : str + The group ID to assign. + description : str + The description to set for the job group. + interruptOnCancel : bool, optional, default False + whether to interrupt jobs on job cancellation. + + Notes + ----- + If interruptOnCancel is set to true for the job group, then job cancellation will result + in Thread.interrupt() being called on the job's executor threads. This is useful to help + ensure that the tasks are actually stopped in a timely manner, but is off by default due + to HDFS-1208, where HDFS may respond to Thread.interrupt() by marking nodes as dead. + + If you run jobs in parallel, use :class:`pyspark.InheritableThread` for thread + local inheritance. + + See Also + -------- + :meth:`SparkContext.cancelJobGroup` + + Examples + -------- + >>> import threading + >>> from time import sleep + >>> from pyspark import InheritableThread + >>> result = "Not Set" + >>> lock = threading.Lock() + >>> def map_func(x): + ... sleep(100) + ... raise RuntimeError("Task should have been cancelled") + ... + >>> def start_job(x): + ... global result + ... try: + ... sc.setJobGroup("job_to_cancel", "some description") + ... result = sc.parallelize(range(x)).map(map_func).collect() + ... except Exception as e: + ... result = "Cancelled" + ... lock.release() + ... + >>> def stop_job(): + ... sleep(5) + ... sc.cancelJobGroup("job_to_cancel") + ... + >>> suppress = lock.acquire() + >>> suppress = InheritableThread(target=start_job, args=(10,)).start() + >>> suppress = InheritableThread(target=stop_job).start() + >>> suppress = lock.acquire() + >>> print(result) + Cancelled + """ + self._jsc.setJobGroup(groupId, description, interruptOnCancel) + + def setInterruptOnCancel(self, interruptOnCancel: bool) -> None: + """ + Set the behavior of job cancellation from jobs started in this thread. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + interruptOnCancel : bool + If true, then job cancellation will result in ``Thread.interrupt()`` + being called on the job's executor threads. This is useful to help ensure that + the tasks are actually stopped in a timely manner, but is off by default due to + HDFS-1208, where HDFS may respond to ``Thread.interrupt()`` by marking nodes as dead. + + See Also + -------- + :meth:`SparkContext.addJobTag` + :meth:`SparkContext.removeJobTag` + :meth:`SparkContext.cancelAllJobs` + :meth:`SparkContext.cancelJobGroup` + :meth:`SparkContext.cancelJobsWithTag` + """ + self._jsc.setInterruptOnCancel(interruptOnCancel) + + def addJobTag(self, tag: str) -> None: + """ + Add a tag to be assigned to all the jobs started by this thread. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + tag : str + The tag to be added. Cannot contain ',' (comma) character. + + See Also + -------- + :meth:`SparkContext.removeJobTag` + :meth:`SparkContext.getJobTags` + :meth:`SparkContext.clearJobTags` + :meth:`SparkContext.cancelJobsWithTag` + :meth:`SparkContext.setInterruptOnCancel` + + Examples + -------- + >>> import threading + >>> from time import sleep + >>> from pyspark import InheritableThread + >>> sc.setInterruptOnCancel(interruptOnCancel=True) + >>> result = "Not Set" + >>> lock = threading.Lock() + >>> def map_func(x): + ... sleep(100) + ... raise RuntimeError("Task should have been cancelled") + ... + >>> def start_job(x): + ... global result + ... try: + ... sc.addJobTag("job_to_cancel") + ... result = sc.parallelize(range(x)).map(map_func).collect() + ... except Exception as e: + ... result = "Cancelled" + ... lock.release() + ... + >>> def stop_job(): + ... sleep(5) + ... sc.cancelJobsWithTag("job_to_cancel") + ... + >>> suppress = lock.acquire() + >>> suppress = InheritableThread(target=start_job, args=(10,)).start() + >>> suppress = InheritableThread(target=stop_job).start() + >>> suppress = lock.acquire() + >>> print(result) + Cancelled + >>> sc.clearJobTags() + """ + self._jsc.addJobTag(tag) + + def removeJobTag(self, tag: str) -> None: + """ + Remove a tag previously added to be assigned to all the jobs started by this thread. + Noop if such a tag was not added earlier. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + tag : str + The tag to be removed. Cannot contain ',' (comma) character. + + See Also + -------- + :meth:`SparkContext.addJobTag` + :meth:`SparkContext.getJobTags` + :meth:`SparkContext.clearJobTags` + :meth:`SparkContext.cancelJobsWithTag` + :meth:`SparkContext.setInterruptOnCancel` + + Examples + -------- + >>> sc.addJobTag("job_to_cancel1") + >>> sc.addJobTag("job_to_cancel2") + >>> sc.getJobTags() + {'job_to_cancel1', 'job_to_cancel2'} + >>> sc.removeJobTag("job_to_cancel1") + >>> sc.getJobTags() + {'job_to_cancel2'} + >>> sc.clearJobTags() + """ + self._jsc.removeJobTag(tag) + + def getJobTags(self) -> Set[str]: + """ + Get the tags that are currently set to be assigned to all the jobs started by this thread. + + .. versionadded:: 3.5.0 + + Returns + ------- + set of str + the tags that are currently set to be assigned to all the jobs started by this thread. + + See Also + -------- + :meth:`SparkContext.addJobTag` + :meth:`SparkContext.removeJobTag` + :meth:`SparkContext.clearJobTags` + :meth:`SparkContext.cancelJobsWithTag` + :meth:`SparkContext.setInterruptOnCancel` + + Examples + -------- + >>> sc.addJobTag("job_to_cancel") + >>> sc.getJobTags() + {'job_to_cancel'} + >>> sc.clearJobTags() + """ + return self._jsc.getJobTags() + + def clearJobTags(self) -> None: + """ + Clear the current thread's job tags. + + .. versionadded:: 3.5.0 + + See Also + -------- + :meth:`SparkContext.addJobTag` + :meth:`SparkContext.removeJobTag` + :meth:`SparkContext.getJobTags` + :meth:`SparkContext.cancelJobsWithTag` + :meth:`SparkContext.setInterruptOnCancel` + + Examples + -------- + >>> sc.addJobTag("job_to_cancel") + >>> sc.clearJobTags() + >>> sc.getJobTags() + set() + """ + self._jsc.clearJobTags() + + def setLocalProperty(self, key: str, value: str) -> None: + """ + Set a local property that affects jobs submitted from this thread, such as the + Spark fair scheduler pool. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + key : str + The key of the local property to set. + value : str + The value of the local property to set. + + See Also + -------- + :meth:`SparkContext.getLocalProperty` + + Notes + ----- + If you run jobs in parallel, use :class:`pyspark.InheritableThread` for thread + local inheritance. + """ + self._jsc.setLocalProperty(key, value) + + def getLocalProperty(self, key: str) -> Optional[str]: + """ + Get a local property set in this thread, or null if it is missing. See + :meth:`setLocalProperty`. + + .. versionadded:: 1.0.0 + + See Also + -------- + :meth:`SparkContext.setLocalProperty` + """ + return self._jsc.getLocalProperty(key) + + def setJobDescription(self, value: str) -> None: + """ + Set a human readable description of the current job. + + .. versionadded:: 2.3.0 + + Parameters + ---------- + value : str + The job description to set. + + Notes + ----- + If you run jobs in parallel, use :class:`pyspark.InheritableThread` for thread + local inheritance. + """ + self._jsc.setJobDescription(value) + + def sparkUser(self) -> str: + """ + Get SPARK_USER for user who is running SparkContext. + + .. versionadded:: 1.0.0 + """ + return self._jsc.sc().sparkUser() + + def cancelJobGroup(self, groupId: str) -> None: + """ + Cancel active jobs for the specified group. See :meth:`SparkContext.setJobGroup`. + for more information. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + groupId : str + The group ID to cancel the job. + + See Also + -------- + :meth:`SparkContext.setJobGroup` + """ + self._jsc.sc().cancelJobGroup(groupId) + + def cancelJobsWithTag(self, tag: str) -> None: + """ + Cancel active jobs that have the specified tag. See + :meth:`SparkContext.addJobTag`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + tag : str + The tag to be cancelled. Cannot contain ',' (comma) character. + + See Also + -------- + :meth:`SparkContext.addJobTag` + :meth:`SparkContext.removeJobTag` + :meth:`SparkContext.getJobTags` + :meth:`SparkContext.clearJobTags` + :meth:`SparkContext.setInterruptOnCancel` + """ + return self._jsc.cancelJobsWithTag(tag) + + def cancelAllJobs(self) -> None: + """ + Cancel all jobs that have been scheduled or are running. + + .. versionadded:: 1.1.0 + + See Also + -------- + :meth:`SparkContext.cancelJobGroup` + :meth:`SparkContext.cancelJobsWithTag` + :meth:`SparkContext.runJob` + """ + self._jsc.sc().cancelAllJobs() + + def statusTracker(self) -> StatusTracker: + """ + Return :class:`StatusTracker` object + + .. versionadded:: 1.4.0 + """ + return StatusTracker(self._jsc.statusTracker()) + + def runJob( + self, + rdd: RDD[T], + partitionFunc: Callable[[Iterable[T]], Iterable[U]], + partitions: Optional[Sequence[int]] = None, + allowLocal: bool = False, + ) -> List[U]: + """ + Executes the given partitionFunc on the specified set of partitions, + returning the result as an array of elements. + + If 'partitions' is not specified, this will run over all partitions. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + rdd : :class:`RDD` + target RDD to run tasks on + partitionFunc : function + a function to run on each partition of the RDD + partitions : list, optional + set of partitions to run on; some jobs may not want to compute on all + partitions of the target RDD, e.g. for operations like `first` + allowLocal : bool, default False + this parameter takes no effect + + Returns + ------- + list + results of specified partitions + + See Also + -------- + :meth:`SparkContext.cancelAllJobs` + + Examples + -------- + >>> myRDD = sc.parallelize(range(6), 3) + >>> sc.runJob(myRDD, lambda part: [x * x for x in part]) + [0, 1, 4, 9, 16, 25] + + >>> myRDD = sc.parallelize(range(6), 3) + >>> sc.runJob(myRDD, lambda part: [x * x for x in part], [0, 2], True) + [0, 1, 16, 25] + """ + if partitions is None: + partitions = list(range(rdd._jrdd.partitions().size())) + + # Implementation note: This is implemented as a mapPartitions followed + # by runJob() in order to avoid having to pass a Python lambda into + # SparkContext#runJob. + mappedRDD = rdd.mapPartitions(partitionFunc) + assert self._jvm is not None + sock_info = self._jvm.PythonRDD.runJob(self._jsc.sc(), mappedRDD._jrdd, partitions) + return list(_load_from_socket(sock_info, mappedRDD._jrdd_deserializer)) + + def show_profiles(self) -> None: + """Print the profile stats to stdout + + .. versionadded:: 1.2.0 + + See Also + -------- + :meth:`SparkContext.dump_profiles` + """ + if self.profiler_collector is not None: + self.profiler_collector.show_profiles() + else: + raise PySparkRuntimeError( + error_class="INCORRECT_CONF_FOR_PROFILE", + message_parameters={}, + ) + + def dump_profiles(self, path: str) -> None: + """Dump the profile stats into directory `path` + + .. versionadded:: 1.2.0 + + See Also + -------- + :meth:`SparkContext.show_profiles` + """ + if self.profiler_collector is not None: + self.profiler_collector.dump_profiles(path) + else: + raise PySparkRuntimeError( + error_class="INCORRECT_CONF_FOR_PROFILE", + message_parameters={}, + ) + + def getConf(self) -> SparkConf: + """Return a copy of this SparkContext's configuration :class:`SparkConf`. + + .. versionadded:: 2.1.0 + """ + conf = SparkConf() + conf.setAll(self._conf.getAll()) + return conf + + @property + def resources(self) -> Dict[str, ResourceInformation]: + """ + Return the resource information of this :class:`SparkContext`. + A resource could be a GPU, FPGA, etc. + + .. versionadded:: 3.0.0 + """ + resources = {} + jresources = self._jsc.resources() + for x in jresources: + name = jresources[x].name() + jaddresses = jresources[x].addresses() + addrs = [addr for addr in jaddresses] + resources[name] = ResourceInformation(name, addrs) + return resources + + @staticmethod + def _assert_on_driver() -> None: + """ + Called to ensure that SparkContext is created only on the Driver. + + Throws an exception if a SparkContext is about to be created in executors. + """ + if TaskContext.get() is not None: + raise PySparkRuntimeError( + error_class="CONTEXT_ONLY_VALID_ON_DRIVER", + message_parameters={}, + ) + + +def _test() -> None: + import doctest + from pyspark import SparkConf + + globs = globals().copy() + conf = SparkConf().set("spark.ui.enabled", "True") + globs["sc"] = SparkContext("local[4]", "context tests", conf=conf) + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/daemon.py b/spark-3.5.0-bin-hadoop3/python/pyspark/daemon.py new file mode 100644 index 0000000..81b6481 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/daemon.py @@ -0,0 +1,218 @@ +# +# 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 numbers +import os +import signal +import select +import socket +import sys +import traceback +import time +import gc +from errno import EINTR, EAGAIN +from socket import AF_INET, AF_INET6, SOCK_STREAM, SOMAXCONN +from signal import SIGHUP, SIGTERM, SIGCHLD, SIG_DFL, SIG_IGN, SIGINT + +from pyspark.worker import main as worker_main +from pyspark.serializers import read_int, write_int, write_with_length, UTF8Deserializer + + +def compute_real_exit_code(exit_code): + # SystemExit's code can be integer or string, but os._exit only accepts integers + if isinstance(exit_code, numbers.Integral): + return exit_code + else: + return 1 + + +def worker(sock, authenticated): + """ + Called by a worker process after the fork(). + """ + signal.signal(SIGHUP, SIG_DFL) + signal.signal(SIGCHLD, SIG_DFL) + signal.signal(SIGTERM, SIG_DFL) + # restore the handler for SIGINT, + # it's useful for debugging (show the stacktrace before exit) + signal.signal(SIGINT, signal.default_int_handler) + + # Read the socket using fdopen instead of socket.makefile() because the latter + # seems to be very slow; note that we need to dup() the file descriptor because + # otherwise writes also cause a seek that makes us miss data on the read side. + buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536)) + infile = os.fdopen(os.dup(sock.fileno()), "rb", buffer_size) + outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size) + + if not authenticated: + client_secret = UTF8Deserializer().loads(infile) + if os.environ["PYTHON_WORKER_FACTORY_SECRET"] == client_secret: + write_with_length("ok".encode("utf-8"), outfile) + outfile.flush() + else: + write_with_length("err".encode("utf-8"), outfile) + outfile.flush() + sock.close() + return 1 + + exit_code = 0 + try: + worker_main(infile, outfile) + except SystemExit as exc: + exit_code = compute_real_exit_code(exc.code) + finally: + try: + outfile.flush() + except Exception: + pass + return exit_code + + +def manager(): + # Create a new process group to corral our children + os.setpgid(0, 0) + + # Create a listening socket on the loopback interface + if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true": + listen_sock = socket.socket(AF_INET6, SOCK_STREAM) + listen_sock.bind(("::1", 0, 0, 0)) + listen_sock.listen(max(1024, SOMAXCONN)) + listen_host, listen_port, _, _ = listen_sock.getsockname() + else: + listen_sock = socket.socket(AF_INET, SOCK_STREAM) + listen_sock.bind(("127.0.0.1", 0)) + listen_sock.listen(max(1024, SOMAXCONN)) + listen_host, listen_port = listen_sock.getsockname() + + # re-open stdin/stdout in 'wb' mode + stdin_bin = os.fdopen(sys.stdin.fileno(), "rb", 4) + stdout_bin = os.fdopen(sys.stdout.fileno(), "wb", 4) + write_int(listen_port, stdout_bin) + stdout_bin.flush() + + def shutdown(code): + signal.signal(SIGTERM, SIG_DFL) + # Send SIGHUP to notify workers of shutdown + os.kill(0, SIGHUP) + sys.exit(code) + + def handle_sigterm(*args): + shutdown(1) + + signal.signal(SIGTERM, handle_sigterm) # Gracefully exit on SIGTERM + signal.signal(SIGHUP, SIG_IGN) # Don't die on SIGHUP + signal.signal(SIGCHLD, SIG_IGN) + + reuse = os.environ.get("SPARK_REUSE_WORKER") + + # Initialization complete + try: + while True: + try: + ready_fds = select.select([0, listen_sock], [], [], 1)[0] + except select.error as ex: + if ex[0] == EINTR: + continue + else: + raise + + if 0 in ready_fds: + try: + worker_pid = read_int(stdin_bin) + except EOFError: + # Spark told us to exit by closing stdin + shutdown(0) + try: + os.kill(worker_pid, signal.SIGKILL) + except OSError: + pass # process already died + + if listen_sock in ready_fds: + try: + sock, _ = listen_sock.accept() + except OSError as e: + if e.errno == EINTR: + continue + raise + + # Launch a worker process + try: + pid = os.fork() + except OSError as e: + if e.errno in (EAGAIN, EINTR): + time.sleep(1) + pid = os.fork() # error here will shutdown daemon + else: + outfile = sock.makefile(mode="wb") + write_int(e.errno, outfile) # Signal that the fork failed + outfile.flush() + outfile.close() + sock.close() + continue + + if pid == 0: + # in child process + listen_sock.close() + + # It should close the standard input in the child process so that + # Python native function executions stay intact. + # + # Note that if we just close the standard input (file descriptor 0), + # the lowest file descriptor (file descriptor 0) will be allocated, + # later when other file descriptors should happen to open. + # + # Therefore, here we redirects it to '/dev/null' by duplicating + # another file descriptor for '/dev/null' to the standard input (0). + # See SPARK-26175. + devnull = open(os.devnull, "r") + os.dup2(devnull.fileno(), 0) + devnull.close() + + try: + # Acknowledge that the fork was successful + outfile = sock.makefile(mode="wb") + write_int(os.getpid(), outfile) + outfile.flush() + outfile.close() + authenticated = False + while True: + code = worker(sock, authenticated) + if code == 0: + authenticated = True + if not reuse or code: + # wait for closing + try: + while sock.recv(1024): + pass + except Exception: + pass + break + gc.collect() + except BaseException: + traceback.print_exc() + os._exit(1) + else: + os._exit(0) + else: + sock.close() + + finally: + shutdown(1) + + +if __name__ == "__main__": + manager() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/__init__.py new file mode 100644 index 0000000..a9bcb97 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/__init__.py @@ -0,0 +1,70 @@ +# +# 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. +# + +""" +PySpark exceptions. +""" +from pyspark.errors.exceptions.base import ( # noqa: F401 + PySparkException, + AnalysisException, + TempTableAlreadyExistsException, + ParseException, + IllegalArgumentException, + ArithmeticException, + UnsupportedOperationException, + ArrayIndexOutOfBoundsException, + DateTimeException, + NumberFormatException, + StreamingQueryException, + QueryExecutionException, + PythonException, + UnknownException, + SparkRuntimeException, + SparkUpgradeException, + PySparkTypeError, + PySparkValueError, + PySparkAttributeError, + PySparkRuntimeError, + PySparkAssertionError, + PySparkNotImplementedError, +) + + +__all__ = [ + "PySparkException", + "AnalysisException", + "TempTableAlreadyExistsException", + "ParseException", + "IllegalArgumentException", + "ArithmeticException", + "UnsupportedOperationException", + "ArrayIndexOutOfBoundsException", + "DateTimeException", + "NumberFormatException", + "StreamingQueryException", + "QueryExecutionException", + "PythonException", + "UnknownException", + "SparkRuntimeException", + "SparkUpgradeException", + "PySparkTypeError", + "PySparkValueError", + "PySparkAttributeError", + "PySparkRuntimeError", + "PySparkAssertionError", + "PySparkNotImplementedError", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/error_classes.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/error_classes.py new file mode 100644 index 0000000..0fbe489 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/error_classes.py @@ -0,0 +1,889 @@ +# +# 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 json + + +ERROR_CLASSES_JSON = """ +{ + "APPLICATION_NAME_NOT_SET" : { + "message" : [ + "An application name must be set in your configuration." + ] + }, + "ARGUMENT_REQUIRED": { + "message": [ + "Argument `` is required when ." + ] + }, + "ATTRIBUTE_NOT_CALLABLE" : { + "message" : [ + "Attribute `` in provided object `` is not callable." + ] + }, + "ATTRIBUTE_NOT_SUPPORTED" : { + "message" : [ + "Attribute `` is not supported." + ] + }, + "AXIS_LENGTH_MISMATCH" : { + "message" : [ + "Length mismatch: Expected axis has element, new values have elements." + ] + }, + "BROADCAST_VARIABLE_NOT_LOADED": { + "message": [ + "Broadcast variable `` not loaded." + ] + }, + "CALL_BEFORE_INITIALIZE": { + "message": [ + "Not supported to call `` before initialize ." + ] + }, + "CANNOT_ACCEPT_OBJECT_IN_TYPE": { + "message": [ + "`` can not accept object `` in type ``." + ] + }, + "CANNOT_ACCESS_TO_DUNDER": { + "message": [ + "Dunder(double underscore) attribute is for internal use only." + ] + }, + "CANNOT_APPLY_IN_FOR_COLUMN": { + "message": [ + "Cannot apply 'in' operator against a column: please use 'contains' in a string column or 'array_contains' function for an array column." + ] + }, + "CANNOT_BE_EMPTY": { + "message": [ + "At least one must be specified." + ] + }, + "CANNOT_BE_NONE": { + "message": [ + "Argument `` can not be None." + ] + }, + "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_TYPE": { + "message": [ + "Cannot convert into ." + ] + }, + "CANNOT_DETERMINE_TYPE": { + "message": [ + "Some of types cannot be determined after inferring." + ] + }, + "CANNOT_GET_BATCH_ID": { + "message": [ + "Could not get batch id from ." + ] + }, + "CANNOT_INFER_ARRAY_TYPE": { + "message": [ + "Can not infer Array Type from an list with None as the first element." + ] + }, + "CANNOT_INFER_EMPTY_SCHEMA": { + "message": [ + "Can not infer schema from empty dataset." + ] + }, + "CANNOT_INFER_SCHEMA_FOR_TYPE": { + "message": [ + "Can not infer schema for type: ``." + ] + }, + "CANNOT_INFER_TYPE_FOR_FIELD": { + "message": [ + "Unable to infer the type of the field ``." + ] + }, + "CANNOT_MERGE_TYPE": { + "message": [ + "Can not merge type `` and ``." + ] + }, + "CANNOT_OPEN_SOCKET": { + "message": [ + "Can not open socket: ." + ] + }, + "CANNOT_PARSE_DATATYPE": { + "message": [ + "Unable to parse datatype. ." + ] + }, + "CANNOT_PROVIDE_METADATA": { + "message": [ + "metadata can only be provided for a single column." + ] + }, + "CANNOT_SET_TOGETHER": { + "message": [ + " should not be set together." + ] + }, + "CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF": { + "message": [ + "returnType can not be specified when `` is a user-defined function, but got ." + ] + }, + "COLUMN_IN_LIST": { + "message": [ + "`` does not allow a Column in a list." + ] + }, + "CONTEXT_ONLY_VALID_ON_DRIVER" : { + "message" : [ + "It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063." + ] + }, + "CONTEXT_UNAVAILABLE_FOR_REMOTE_CLIENT" : { + "message" : [ + "Remote client cannot create a SparkContext. Create SparkSession instead." + ] + }, + "DIFFERENT_PANDAS_DATAFRAME" : { + "message" : [ + "DataFrames are not almost equal:", + "Left:", + "", + "", + "Right:", + "", + "" + ] + }, + "DIFFERENT_PANDAS_INDEX" : { + "message" : [ + "Indices are not almost equal:", + "Left:", + "", + "", + "Right:", + "", + "" + ] + }, + "DIFFERENT_PANDAS_MULTIINDEX" : { + "message" : [ + "MultiIndices are not almost equal:", + "Left:", + "", + "", + "Right:", + "", + "" + ] + }, + "DIFFERENT_PANDAS_SERIES" : { + "message" : [ + "Series are not almost equal:", + "Left:", + "", + "", + "Right:", + "", + "" + ] + }, + "DIFFERENT_ROWS" : { + "message" : [ + "" + ] + }, + "DIFFERENT_SCHEMA" : { + "message" : [ + "Schemas do not match.", + "--- actual", + "+++ expected", + "" + ] + }, + "DISALLOWED_TYPE_FOR_CONTAINER" : { + "message" : [ + "Argument ``(type: ) should only contain a type in [], got " + ] + }, + "DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT" : { + "message" : [ + "Duplicated field names in Arrow Struct are not allowed, got " + ] + }, + "HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN" : { + "message" : [ + "Function `` should return Column, got ." + ] + }, + "INCORRECT_CONF_FOR_PROFILE" : { + "message" : [ + "`spark.python.profile` or `spark.python.profile.memory` configuration", + " must be set to `true` to enable Python profile." + ] + }, + "INVALID_ARROW_UDTF_RETURN_TYPE" : { + "message" : [ + "The return type of the arrow-optimized Python UDTF should be of type 'pandas.DataFrame', but the '' method returned a value of type with value: ." + ] + }, + "INVALID_BROADCAST_OPERATION": { + "message": [ + "Broadcast can only be in driver." + ] + }, + "INVALID_CALL_ON_UNRESOLVED_OBJECT": { + "message": [ + "Invalid call to `` on unresolved object." + ] + }, + "INVALID_CONNECT_URL" : { + "message" : [ + "Invalid URL for Spark Connect: " + ] + }, + "INVALID_ITEM_FOR_CONTAINER": { + "message": [ + "All items in `` should be in , got ." + ] + }, + "INVALID_NDARRAY_DIMENSION": { + "message": [ + "NumPy array input should be of dimensions." + ] + }, + "INVALID_PANDAS_UDF" : { + "message" : [ + "Invalid function: " + ] + }, + "INVALID_PANDAS_UDF_TYPE" : { + "message" : [ + "`` should be one the values from PandasUDFType, got " + ] + }, + "INVALID_RETURN_TYPE_FOR_PANDAS_UDF": { + "message": [ + "Pandas UDF should return StructType for , got ." + ] + }, + "INVALID_TIMEOUT_TIMESTAMP" : { + "message" : [ + "Timeout timestamp () cannot be earlier than the current watermark ()." + ] + }, + "INVALID_TYPE" : { + "message" : [ + "Argument `` should not be a ." + ] + }, + "INVALID_TYPENAME_CALL" : { + "message" : [ + "StructField does not have typeName. Use typeName on its type explicitly instead." + ] + }, + "INVALID_TYPE_DF_EQUALITY_ARG" : { + "message" : [ + "Expected type for `` but got type ." + ] + }, + "INVALID_UDF_EVAL_TYPE" : { + "message" : [ + "Eval type for UDF must be ." + ] + }, + "INVALID_UDTF_EVAL_TYPE" : { + "message" : [ + "The eval type for the UDTF '' is invalid. It must be one of ." + ] + }, + "INVALID_UDTF_HANDLER_TYPE" : { + "message" : [ + "The UDTF is invalid. The function handler must be a class, but got ''. Please provide a class as the function handler." + ] + }, + "INVALID_UDTF_NO_EVAL" : { + "message" : [ + "The UDTF '' is invalid. It does not implement the required 'eval' method. Please implement the 'eval' method in '' and try again." + ] + }, + "INVALID_WHEN_USAGE": { + "message": [ + "when() can only be applied on a Column previously generated by when() function, and cannot be applied once otherwise() is applied." + ] + }, + "INVALID_WINDOW_BOUND_TYPE" : { + "message" : [ + "Invalid window bound type: ." + ] + }, + "JAVA_GATEWAY_EXITED" : { + "message" : [ + "Java gateway process exited before sending its port number." + ] + }, + "JVM_ATTRIBUTE_NOT_SUPPORTED" : { + "message" : [ + "Attribute `` is not supported in Spark Connect as it depends on the JVM. If you need to use this attribute, do not use Spark Connect when creating your session." + ] + }, + "KEY_VALUE_PAIR_REQUIRED" : { + "message" : [ + "Key-value pair or a list of pairs is required." + ] + }, + "LENGTH_SHOULD_BE_THE_SAME" : { + "message" : [ + " and should be of the same length, got and ." + ] + }, + "MASTER_URL_NOT_SET" : { + "message" : [ + "A master URL must be set in your configuration." + ] + }, + "MISSING_LIBRARY_FOR_PROFILER" : { + "message" : [ + "Install the 'memory_profiler' library in the cluster to enable memory profiling." + ] + }, + "MISSING_VALID_PLAN" : { + "message" : [ + "Argument to does not contain a valid plan." + ] + }, + "MIXED_TYPE_REPLACEMENT" : { + "message" : [ + "Mixed type replacements are not supported." + ] + }, + "NEGATIVE_VALUE" : { + "message" : [ + "Value for `` must be greater than or equal to 0, got ''." + ] + }, + "NOT_BOOL" : { + "message" : [ + "Argument `` should be a bool, got ." + ] + }, + "NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_LIST_OR_STR_OR_TUPLE" : { + "message" : [ + "Argument `` should be a bool, dict, float, int, str or tuple, got ." + ] + }, + "NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_STR" : { + "message" : [ + "Argument `` should be a bool, dict, float, int or str, got ." + ] + }, + "NOT_BOOL_OR_FLOAT_OR_INT" : { + "message" : [ + "Argument `` should be a bool, float or str, got ." + ] + }, + "NOT_BOOL_OR_FLOAT_OR_INT_OR_LIST_OR_NONE_OR_STR_OR_TUPLE" : { + "message" : [ + "Argument `` should be a bool, float, int, list, None, str or tuple, got ." + ] + }, + "NOT_BOOL_OR_FLOAT_OR_INT_OR_STR" : { + "message" : [ + "Argument `` should be a bool, float, int or str, got ." + ] + }, + "NOT_BOOL_OR_LIST" : { + "message" : [ + "Argument `` should be a bool or list, got ." + ] + }, + "NOT_BOOL_OR_STR" : { + "message" : [ + "Argument `` should be a bool or str, got ." + ] + }, + "NOT_CALLABLE" : { + "message" : [ + "Argument `` should be a callable, got ." + ] + }, + "NOT_COLUMN" : { + "message" : [ + "Argument `` should be a Column, got ." + ] + }, + "NOT_COLUMN_OR_DATATYPE_OR_STR" : { + "message" : [ + "Argument `` should be a Column, str or DataType, but got ." + ] + }, + "NOT_COLUMN_OR_FLOAT_OR_INT_OR_LIST_OR_STR" : { + "message" : [ + "Argument `` should be a column, float, integer, list or string, got ." + ] + }, + "NOT_COLUMN_OR_INT" : { + "message" : [ + "Argument `` should be a Column or int, got ." + ] + }, + "NOT_COLUMN_OR_INT_OR_LIST_OR_STR_OR_TUPLE" : { + "message" : [ + "Argument `` should be a Column, int, list, str or tuple, got ." + ] + }, + "NOT_COLUMN_OR_INT_OR_STR" : { + "message" : [ + "Argument `` should be a Column, int or str, got ." + ] + }, + "NOT_COLUMN_OR_LIST_OR_STR" : { + "message" : [ + "Argument `` should be a Column, list or str, got ." + ] + }, + "NOT_COLUMN_OR_STR" : { + "message" : [ + "Argument `` should be a Column or str, got ." + ] + }, + "NOT_DATAFRAME" : { + "message" : [ + "Argument `` should be a DataFrame, got ." + ] + }, + "NOT_DATATYPE_OR_STR" : { + "message" : [ + "Argument `` should be a DataType or str, got ." + ] + }, + "NOT_DICT" : { + "message" : [ + "Argument `` should be a dict, got ." + ] + }, + "NOT_EXPRESSION" : { + "message" : [ + "Argument `` should be a Expression, got ." + ] + }, + "NOT_FLOAT_OR_INT" : { + "message" : [ + "Argument `` should be a float or int, got ." + ] + }, + "NOT_FLOAT_OR_INT_OR_LIST_OR_STR" : { + "message" : [ + "Argument `` should be a float, int, list or str, got ." + ] + }, + "NOT_IMPLEMENTED" : { + "message" : [ + " is not implemented." + ] + }, + "NOT_INSTANCE_OF" : { + "message" : [ + " is not an instance of type ." + ] + }, + "NOT_INT" : { + "message" : [ + "Argument `` should be an int, got ." + ] + }, + "NOT_INT_OR_SLICE_OR_STR" : { + "message" : [ + "Argument `` should be an int, slice or str, got ." + ] + }, + "NOT_IN_BARRIER_STAGE" : { + "message" : [ + "It is not in a barrier stage." + ] + }, + "NOT_ITERABLE" : { + "message" : [ + " is not iterable." + ] + }, + "NOT_LIST" : { + "message" : [ + "Argument `` should be a list, got ." + ] + }, + "NOT_LIST_OF_COLUMN" : { + "message" : [ + "Argument `` should be a list[Column]." + ] + }, + "NOT_LIST_OF_COLUMN_OR_STR" : { + "message" : [ + "Argument `` should be a list[Column]." + ] + }, + "NOT_LIST_OF_FLOAT_OR_INT" : { + "message" : [ + "Argument `` should be a list[float, int], got ." + ] + }, + "NOT_LIST_OF_STR" : { + "message" : [ + "Argument `` should be a list[str], got ." + ] + }, + "NOT_LIST_OR_NONE_OR_STRUCT" : { + "message" : [ + "Argument `` should be a list, None or StructType, got ." + ] + }, + "NOT_LIST_OR_STR_OR_TUPLE" : { + "message" : [ + "Argument `` should be a list, str or tuple, got ." + ] + }, + "NOT_LIST_OR_TUPLE" : { + "message" : [ + "Argument `` should be a list or tuple, got ." + ] + }, + "NOT_NUMERIC_COLUMNS" : { + "message" : [ + "Numeric aggregation function can only be applied on numeric columns, got ." + ] + }, + "NOT_OBSERVATION_OR_STR" : { + "message" : [ + "Argument `` should be a Observation or str, got ." + ] + }, + "NOT_SAME_TYPE" : { + "message" : [ + "Argument `` and `` should be the same type, got and ." + ] + }, + "NOT_STR" : { + "message" : [ + "Argument `` should be a str, got ." + ] + }, + "NOT_STR_OR_LIST_OF_RDD" : { + "message" : [ + "Argument `` should be a str or list[RDD], got ." + ] + }, + "NOT_STR_OR_STRUCT" : { + "message" : [ + "Argument `` should be a str or structType, got ." + ] + }, + "NOT_WINDOWSPEC" : { + "message" : [ + "Argument `` should be a WindowSpec, got ." + ] + }, + "NO_ACTIVE_OR_DEFAULT_SESSION" : { + "message" : [ + "No active or default Spark session found. Please create a new Spark session before running the code." + ] + }, + "NO_ACTIVE_SESSION" : { + "message" : [ + "No active Spark session found. Please create a new Spark session before running the code." + ] + }, + "ONLY_ALLOWED_FOR_SINGLE_COLUMN" : { + "message" : [ + "Argument `` can only be provided for a single column." + ] + }, + "ONLY_ALLOW_SINGLE_TRIGGER" : { + "message" : [ + "Only a single trigger is allowed." + ] + }, + "PIPE_FUNCTION_EXITED" : { + "message" : [ + "Pipe function `` exited with error code ." + ] + }, + "PYTHON_HASH_SEED_NOT_SET" : { + "message" : [ + "Randomness of hash of string should be disabled via PYTHONHASHSEED." + ] + }, + "PYTHON_VERSION_MISMATCH" : { + "message" : [ + "Python in worker has different version than that in driver , PySpark cannot run with different minor versions.", + "Please check environment variables PYSPARK_PYTHON and PYSPARK_DRIVER_PYTHON are correctly set." + ] + }, + "RDD_TRANSFORM_ONLY_VALID_ON_DRIVER" : { + "message" : [ + "It appears that you are attempting to broadcast an RDD or reference an RDD from an ", + "action or transformation. RDD transformations and actions can only be invoked by the ", + "driver, not inside of other transformations; for example, ", + "rdd1.map(lambda x: rdd2.values.count() * x) is invalid because the values ", + "transformation and count action cannot be performed inside of the rdd1.map ", + "transformation. For more information, see SPARK-5063." + ] + }, + "RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF" : { + "message" : [ + "Column names of the returned pandas.DataFrame do not match specified schema." + ] + }, + "RESULT_LENGTH_MISMATCH_FOR_PANDAS_UDF" : { + "message" : [ + "Number of columns of the returned pandas.DataFrame doesn't match specified schema. Expected: Actual: " + ] + }, + "RESULT_LENGTH_MISMATCH_FOR_SCALAR_ITER_PANDAS_UDF" : { + "message" : [ + "The length of output in Scalar iterator pandas UDF should be the same with the input's; however, the length of output was and the length of input was ." + ] + }, + "SCHEMA_MISMATCH_FOR_PANDAS_UDF" : { + "message" : [ + "Result vector from pandas_udf 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." + ] + }, + "SESSION_NOT_SAME" : { + "message" : [ + "Both Datasets must belong to the same SparkSession." + ] + }, + "SESSION_OR_CONTEXT_EXISTS" : { + "message" : [ + "There should not be an existing Spark Session or Spark Context." + ] + }, + "SHOULD_NOT_DATAFRAME": { + "message": [ + "Argument `` should not be a DataFrame." + ] + }, + "SLICE_WITH_STEP" : { + "message" : [ + "Slice with step is not supported." + ] + }, + "STATE_NOT_EXISTS" : { + "message" : [ + "State is either not defined or has already been removed." + ] + }, + "STOP_ITERATION_OCCURRED" : { + "message" : [ + "Caught StopIteration thrown from user's code; failing the task: " + ] + }, + "STOP_ITERATION_OCCURRED_FROM_SCALAR_ITER_PANDAS_UDF" : { + "message" : [ + "pandas iterator UDF should exhaust the input iterator." + ] + }, + "TOO_MANY_VALUES" : { + "message" : [ + "Expected values for ``, got ." + ] + }, + "UDF_RETURN_TYPE" : { + "message" : [ + "Return type of the user-defined function should be , but is ." + ] + }, + "UDTF_ARROW_TYPE_CAST_ERROR" : { + "message" : [ + "Cannot convert the output value of the column '' with type '' to the specified return type of the column: ''. Please check if the data types match and try again." + ] + }, + "UDTF_EXEC_ERROR" : { + "message" : [ + "User defined table function encountered an error in the '' method: " + ] + }, + "UDTF_INVALID_OUTPUT_ROW_TYPE" : { + "message" : [ + "The type of an individual output row in the '' method of the UDTF is invalid. Each row should be a tuple, list, or dict, but got ''. Please make sure that the output rows are of the correct type." + ] + }, + "UDTF_RETURN_NOT_ITERABLE" : { + "message" : [ + "The return value of the '' method of the UDTF is invalid. It should be an iterable (e.g., generator or list), but got ''. Please make sure that the UDTF returns one of these types." + ] + }, + "UDTF_RETURN_SCHEMA_MISMATCH" : { + "message" : [ + "The number of columns in the result does not match the specified schema. Expected column count: , Actual column count: . Please make sure the values returned by the '' method have the same number of columns as specified in the output schema." + ] + }, + "UDTF_RETURN_TYPE_MISMATCH" : { + "message" : [ + "Mismatch in return type for the UDTF ''. Expected a 'StructType', but got ''. Please ensure the return type is a correctly formatted StructType." + ] + }, + "UDTF_SERIALIZATION_ERROR" : { + "message" : [ + "Cannot serialize the UDTF '': " + ] + }, + "UNEXPECTED_RESPONSE_FROM_SERVER" : { + "message" : [ + "Unexpected response from iterator server." + ] + }, + "UNEXPECTED_TUPLE_WITH_STRUCT" : { + "message" : [ + "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'." + ] + }, + "UNKNOWN_RESPONSE" : { + "message" : [ + "Unknown response: ." + ] + }, + "UNSUPPORTED_DATA_TYPE" : { + "message" : [ + "Unsupported DataType ``." + ] + }, + "UNSUPPORTED_DATA_TYPE_FOR_ARROW" : { + "message" : [ + "Single data type is not supported with Arrow." + ] + }, + "UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION" : { + "message" : [ + " is not supported in conversion to Arrow." + ] + }, + "UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION" : { + "message" : [ + " is only supported with pyarrow 2.0.0 and above." + ] + }, + "UNSUPPORTED_JOIN_TYPE" : { + "message" : [ + "Unsupported join type: . Supported join types include: \\"inner\\", \\"outer\\", \\"full\\", \\"fullouter\\", \\"full_outer\\", \\"leftouter\\", \\"left\\", \\"left_outer\\", \\"rightouter\\", \\"right\\", \\"right_outer\\", \\"leftsemi\\", \\"left_semi\\", \\"semi\\", \\"leftanti\\", \\"left_anti\\", \\"anti\\", \\"cross\\"." + ] + }, + "UNSUPPORTED_LITERAL" : { + "message" : [ + "Unsupported Literal ''." + ] + }, + "UNSUPPORTED_NUMPY_ARRAY_SCALAR" : { + "message" : [ + "The type of array scalar '' is not supported." + ] + }, + "UNSUPPORTED_OPERATION" : { + "message" : [ + " is not supported." + ] + }, + "UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION" : { + "message" : [ + "Function `` should use only POSITIONAL or POSITIONAL OR KEYWORD arguments." + ] + }, + "UNSUPPORTED_SIGNATURE" : { + "message" : [ + "Unsupported signature: ." + ] + }, + "UNSUPPORTED_WITH_ARROW_OPTIMIZATION" : { + "message" : [ + " is not supported with Arrow optimization enabled in Python UDFs. Disable 'spark.sql.execution.pythonUDF.arrow.enabled' to workaround.." + ] + }, + "VALUE_NOT_ACCESSIBLE": { + "message": [ + "Value `` cannot be accessed inside tasks." + ] + }, + "VALUE_NOT_ANY_OR_ALL" : { + "message" : [ + "Value for `` must be 'any' or 'all', got ''." + ] + }, + "VALUE_NOT_BETWEEN" : { + "message" : [ + "Value for `` must be between and ." + ] + }, + "VALUE_NOT_NON_EMPTY_STR" : { + "message" : [ + "Value for `` must be a non empty string, got ''." + ] + }, + "VALUE_NOT_PEARSON" : { + "message" : [ + "Value for `` only supports the 'pearson', got ''." + ] + }, + "VALUE_NOT_POSITIVE" : { + "message" : [ + "Value for `` must be positive, got ''." + ] + }, + "VALUE_NOT_TRUE" : { + "message" : [ + "Value for `` must be True, got ''." + ] + }, + "VALUE_OUT_OF_BOUND" : { + "message" : [ + "Value for `` must be greater than or less than , got " + ] + }, + "WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION" : { + "message" : [ + "Function `` should take between 1 and 3 arguments, but provided function takes ." + ] + }, + "WRONG_NUM_COLUMNS" : { + "message" : [ + "Function `` should take at least columns." + ] + } +} +""" + +ERROR_CLASSES_MAP = json.loads(ERROR_CLASSES_JSON) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/__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/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/base.py new file mode 100644 index 0000000..fd1c07c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/base.py @@ -0,0 +1,228 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Dict, Optional, cast + +from pyspark.errors.utils import ErrorClassesReader + + +class PySparkException(Exception): + """ + Base Exception for handling errors generated from PySpark. + """ + + def __init__( + self, + message: Optional[str] = None, + error_class: Optional[str] = None, + message_parameters: Optional[Dict[str, str]] = None, + ): + # `message` vs `error_class` & `message_parameters` are mutually exclusive. + assert (message is not None and (error_class is None and message_parameters is None)) or ( + message is None and (error_class is not None and message_parameters is not None) + ) + + self.error_reader = ErrorClassesReader() + + if message is None: + self.message = self.error_reader.get_error_message( + cast(str, error_class), cast(Dict[str, str], message_parameters) + ) + else: + self.message = message + + self.error_class = error_class + self.message_parameters = message_parameters + + def getErrorClass(self) -> Optional[str]: + """ + Returns an error class as a string. + + .. versionadded:: 3.4.0 + + See Also + -------- + :meth:`PySparkException.getMessageParameters` + :meth:`PySparkException.getSqlState` + """ + return self.error_class + + def getMessageParameters(self) -> Optional[Dict[str, str]]: + """ + Returns a message parameters as a dictionary. + + .. versionadded:: 3.4.0 + + See Also + -------- + :meth:`PySparkException.getErrorClass` + :meth:`PySparkException.getSqlState` + """ + return self.message_parameters + + def getSqlState(self) -> None: + """ + Returns an SQLSTATE as a string. + + Errors generated in Python have no SQLSTATE, so it always returns None. + + .. versionadded:: 3.4.0 + + See Also + -------- + :meth:`PySparkException.getErrorClass` + :meth:`PySparkException.getMessageParameters` + """ + return None + + def __str__(self) -> str: + if self.getErrorClass() is not None: + return f"[{self.getErrorClass()}] {self.message}" + else: + return self.message + + +class AnalysisException(PySparkException): + """ + Failed to analyze a SQL query plan. + """ + + +class SessionNotSameException(PySparkException): + """ + Performed the same operation on different SparkSession. + """ + + +class TempTableAlreadyExistsException(AnalysisException): + """ + Failed to create temp view since it is already exists. + """ + + +class ParseException(AnalysisException): + """ + Failed to parse a SQL command. + """ + + +class IllegalArgumentException(PySparkException): + """ + Passed an illegal or inappropriate argument. + """ + + +class ArithmeticException(PySparkException): + """ + Arithmetic exception thrown from Spark with an error class. + """ + + +class UnsupportedOperationException(PySparkException): + """ + Unsupported operation exception thrown from Spark with an error class. + """ + + +class ArrayIndexOutOfBoundsException(PySparkException): + """ + Array index out of bounds exception thrown from Spark with an error class. + """ + + +class DateTimeException(PySparkException): + """ + Datetime exception thrown from Spark with an error class. + """ + + +class NumberFormatException(IllegalArgumentException): + """ + Number format exception thrown from Spark with an error class. + """ + + +class StreamingQueryException(PySparkException): + """ + Exception that stopped a :class:`StreamingQuery`. + """ + + +class QueryExecutionException(PySparkException): + """ + Failed to execute a query. + """ + + +class PythonException(PySparkException): + """ + Exceptions thrown from Python workers. + """ + + +class SparkRuntimeException(PySparkException): + """ + Runtime exception thrown from Spark with an error class. + """ + + +class SparkUpgradeException(PySparkException): + """ + Exception thrown because of Spark upgrade. + """ + + +class UnknownException(PySparkException): + """ + None of the above exceptions. + """ + + +class PySparkValueError(PySparkException, ValueError): + """ + Wrapper class for ValueError to support error classes. + """ + + +class PySparkTypeError(PySparkException, TypeError): + """ + Wrapper class for TypeError to support error classes. + """ + + +class PySparkAttributeError(PySparkException, AttributeError): + """ + Wrapper class for AttributeError to support error classes. + """ + + +class PySparkRuntimeError(PySparkException, RuntimeError): + """ + Wrapper class for RuntimeError to support error classes. + """ + + +class PySparkAssertionError(PySparkException, AssertionError): + """ + Wrapper class for AssertionError to support error classes. + """ + + +class PySparkNotImplementedError(PySparkException, NotImplementedError): + """ + Wrapper class for NotImplementedError to support error classes. + """ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/captured.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/captured.py new file mode 100644 index 0000000..d62b7d2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/captured.py @@ -0,0 +1,307 @@ +# +# 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 contextlib import contextmanager +from typing import Any, Callable, Dict, Iterator, Optional, cast + +import py4j +from py4j.protocol import Py4JJavaError +from py4j.java_gateway import is_instance_of + +from pyspark import SparkContext +from pyspark.errors.exceptions.base import ( + AnalysisException as BaseAnalysisException, + IllegalArgumentException as BaseIllegalArgumentException, + ArithmeticException as BaseArithmeticException, + UnsupportedOperationException as BaseUnsupportedOperationException, + ArrayIndexOutOfBoundsException as BaseArrayIndexOutOfBoundsException, + DateTimeException as BaseDateTimeException, + NumberFormatException as BaseNumberFormatException, + ParseException as BaseParseException, + PySparkException, + PythonException as BasePythonException, + QueryExecutionException as BaseQueryExecutionException, + SparkRuntimeException as BaseSparkRuntimeException, + SparkUpgradeException as BaseSparkUpgradeException, + StreamingQueryException as BaseStreamingQueryException, + UnknownException as BaseUnknownException, +) + + +class CapturedException(PySparkException): + def __init__( + self, + desc: Optional[str] = None, + stackTrace: Optional[str] = None, + cause: Optional[Py4JJavaError] = None, + origin: Optional[Py4JJavaError] = None, + ): + # desc & stackTrace vs origin are mutually exclusive. + # cause is optional. + assert (origin is not None and desc is None and stackTrace is None) or ( + origin is None and desc is not None and stackTrace is not None + ) + + self.desc = desc if desc is not None else cast(Py4JJavaError, origin).getMessage() + assert SparkContext._jvm is not None + self.stackTrace = ( + stackTrace + if stackTrace is not None + else (SparkContext._jvm.org.apache.spark.util.Utils.exceptionString(origin)) + ) + self.cause = convert_exception(cause) if cause is not None else None + if self.cause is None and origin is not None and origin.getCause() is not None: + self.cause = convert_exception(origin.getCause()) + self._origin = origin + + def __str__(self) -> str: + assert SparkContext._jvm is not None + + jvm = SparkContext._jvm + + # SPARK-42752: default to True to see issues with initialization + debug_enabled = True + try: + sql_conf = jvm.org.apache.spark.sql.internal.SQLConf.get() + debug_enabled = sql_conf.pysparkJVMStacktraceEnabled() + except BaseException: + pass + + desc = self.desc + if debug_enabled: + desc = desc + "\n\nJVM stacktrace:\n%s" % self.stackTrace + return str(desc) + + def getErrorClass(self) -> Optional[str]: + assert SparkContext._gateway is not None + + gw = SparkContext._gateway + if self._origin is not None and is_instance_of( + gw, self._origin, "org.apache.spark.SparkThrowable" + ): + return self._origin.getErrorClass() + else: + return None + + def getMessageParameters(self) -> Optional[Dict[str, str]]: + assert SparkContext._gateway is not None + + gw = SparkContext._gateway + if self._origin is not None and is_instance_of( + gw, self._origin, "org.apache.spark.SparkThrowable" + ): + return self._origin.getMessageParameters() + else: + return None + + def getSqlState(self) -> Optional[str]: # type: ignore[override] + assert SparkContext._gateway is not None + gw = SparkContext._gateway + if self._origin is not None and is_instance_of( + gw, self._origin, "org.apache.spark.SparkThrowable" + ): + return self._origin.getSqlState() + else: + return None + + +def convert_exception(e: Py4JJavaError) -> CapturedException: + assert e is not None + assert SparkContext._jvm is not None + assert SparkContext._gateway is not None + + jvm = SparkContext._jvm + gw = SparkContext._gateway + + if is_instance_of(gw, e, "org.apache.spark.sql.catalyst.parser.ParseException"): + return ParseException(origin=e) + # Order matters. ParseException inherits AnalysisException. + elif is_instance_of(gw, e, "org.apache.spark.sql.AnalysisException"): + return AnalysisException(origin=e) + elif is_instance_of(gw, e, "org.apache.spark.sql.streaming.StreamingQueryException"): + return StreamingQueryException(origin=e) + elif is_instance_of(gw, e, "org.apache.spark.sql.execution.QueryExecutionException"): + return QueryExecutionException(origin=e) + # Order matters. NumberFormatException inherits IllegalArgumentException. + elif is_instance_of(gw, e, "java.lang.NumberFormatException"): + return NumberFormatException(origin=e) + elif is_instance_of(gw, e, "java.lang.IllegalArgumentException"): + return IllegalArgumentException(origin=e) + elif is_instance_of(gw, e, "java.lang.ArithmeticException"): + return ArithmeticException(origin=e) + elif is_instance_of(gw, e, "java.lang.UnsupportedOperationException"): + return UnsupportedOperationException(origin=e) + elif is_instance_of(gw, e, "java.lang.ArrayIndexOutOfBoundsException"): + return ArrayIndexOutOfBoundsException(origin=e) + elif is_instance_of(gw, e, "java.time.DateTimeException"): + return DateTimeException(origin=e) + elif is_instance_of(gw, e, "org.apache.spark.SparkRuntimeException"): + return SparkRuntimeException(origin=e) + elif is_instance_of(gw, e, "org.apache.spark.SparkUpgradeException"): + return SparkUpgradeException(origin=e) + + c: Py4JJavaError = e.getCause() + stacktrace: str = jvm.org.apache.spark.util.Utils.exceptionString(e) + if c is not None and ( + is_instance_of(gw, c, "org.apache.spark.api.python.PythonException") + # To make sure this only catches Python UDFs. + and any( + map( + lambda v: "org.apache.spark.sql.execution.python" in v.toString(), c.getStackTrace() + ) + ) + ): + msg = ( + "\n An exception was thrown from the Python worker. " + "Please see the stack trace below.\n%s" % c.getMessage() + ) + return PythonException(msg, stacktrace) + + return UnknownException(desc=e.toString(), stackTrace=stacktrace, cause=c) + + +def capture_sql_exception(f: Callable[..., Any]) -> Callable[..., Any]: + def deco(*a: Any, **kw: Any) -> Any: + try: + return f(*a, **kw) + except Py4JJavaError as e: + converted = convert_exception(e.java_exception) + if not isinstance(converted, UnknownException): + # Hide where the exception came from that shows a non-Pythonic + # JVM exception message. + raise converted from None + else: + raise + + return deco + + +@contextmanager +def unwrap_spark_exception() -> Iterator[Any]: + assert SparkContext._gateway is not None + + gw = SparkContext._gateway + try: + yield + except Py4JJavaError as e: + je: Py4JJavaError = e.java_exception + if je is not None and is_instance_of(gw, je, "org.apache.spark.SparkException"): + converted = convert_exception(je.getCause()) + if not isinstance(converted, UnknownException): + raise converted from None + raise + + +def install_exception_handler() -> None: + """ + Hook an exception handler into Py4j, which could capture some SQL exceptions in Java. + + When calling Java API, it will call `get_return_value` to parse the returned object. + If any exception happened in JVM, the result will be Java exception object, it raise + py4j.protocol.Py4JJavaError. We replace the original `get_return_value` with one that + could capture the Java exception and throw a Python one (with the same error message). + + It's idempotent, could be called multiple times. + """ + original = py4j.protocol.get_return_value + # The original `get_return_value` is not patched, it's idempotent. + patched = capture_sql_exception(original) + # only patch the one used in py4j.java_gateway (call Java API) + py4j.java_gateway.get_return_value = patched + + +class AnalysisException(CapturedException, BaseAnalysisException): + """ + Failed to analyze a SQL query plan. + """ + + +class ParseException(AnalysisException, BaseParseException): + """ + Failed to parse a SQL command. + """ + + +class IllegalArgumentException(CapturedException, BaseIllegalArgumentException): + """ + Passed an illegal or inappropriate argument. + """ + + +class StreamingQueryException(CapturedException, BaseStreamingQueryException): + """ + Exception that stopped a :class:`StreamingQuery`. + """ + + +class QueryExecutionException(CapturedException, BaseQueryExecutionException): + """ + Failed to execute a query. + """ + + +class PythonException(CapturedException, BasePythonException): + """ + Exceptions thrown from Python workers. + """ + + +class ArithmeticException(CapturedException, BaseArithmeticException): + """ + Arithmetic exception. + """ + + +class UnsupportedOperationException(CapturedException, BaseUnsupportedOperationException): + """ + Unsupported operation exception. + """ + + +class ArrayIndexOutOfBoundsException(CapturedException, BaseArrayIndexOutOfBoundsException): + """ + Array index out of bounds exception. + """ + + +class DateTimeException(CapturedException, BaseDateTimeException): + """ + Datetime exception. + """ + + +class NumberFormatException(IllegalArgumentException, BaseNumberFormatException): + """ + Number format exception. + """ + + +class SparkRuntimeException(CapturedException, BaseSparkRuntimeException): + """ + Runtime exception. + """ + + +class SparkUpgradeException(CapturedException, BaseSparkUpgradeException): + """ + Exception thrown because of Spark upgrade. + """ + + +class UnknownException(CapturedException, BaseUnknownException): + """ + None of the above exceptions. + """ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/connect.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/connect.py new file mode 100644 index 0000000..48b2130 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/exceptions/connect.py @@ -0,0 +1,190 @@ +# +# 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 json +from typing import Dict, Optional, TYPE_CHECKING + + +from pyspark.errors.exceptions.base import ( + AnalysisException as BaseAnalysisException, + IllegalArgumentException as BaseIllegalArgumentException, + ArithmeticException as BaseArithmeticException, + UnsupportedOperationException as BaseUnsupportedOperationException, + ArrayIndexOutOfBoundsException as BaseArrayIndexOutOfBoundsException, + DateTimeException as BaseDateTimeException, + NumberFormatException as BaseNumberFormatException, + ParseException as BaseParseException, + PySparkException, + PythonException as BasePythonException, + StreamingQueryException as BaseStreamingQueryException, + QueryExecutionException as BaseQueryExecutionException, + SparkRuntimeException as BaseSparkRuntimeException, + SparkUpgradeException as BaseSparkUpgradeException, +) + +if TYPE_CHECKING: + from google.rpc.error_details_pb2 import ErrorInfo + + +class SparkConnectException(PySparkException): + """ + Exception thrown from Spark Connect. + """ + + +def convert_exception(info: "ErrorInfo", message: str) -> SparkConnectException: + classes = [] + if "classes" in info.metadata: + classes = json.loads(info.metadata["classes"]) + + if "stackTrace" in info.metadata: + stackTrace = info.metadata["stackTrace"] + message += f"\n\nJVM stacktrace:\n{stackTrace}" + + if "org.apache.spark.sql.catalyst.parser.ParseException" in classes: + return ParseException(message) + # Order matters. ParseException inherits AnalysisException. + elif "org.apache.spark.sql.AnalysisException" in classes: + return AnalysisException(message) + elif "org.apache.spark.sql.streaming.StreamingQueryException" in classes: + return StreamingQueryException(message) + elif "org.apache.spark.sql.execution.QueryExecutionException" in classes: + return QueryExecutionException(message) + # Order matters. NumberFormatException inherits IllegalArgumentException. + elif "java.lang.NumberFormatException" in classes: + return NumberFormatException(message) + elif "java.lang.IllegalArgumentException" in classes: + return IllegalArgumentException(message) + elif "java.lang.ArithmeticException" in classes: + return ArithmeticException(message) + elif "java.lang.UnsupportedOperationException" in classes: + return UnsupportedOperationException(message) + elif "java.lang.ArrayIndexOutOfBoundsException" in classes: + return ArrayIndexOutOfBoundsException(message) + elif "java.time.DateTimeException" in classes: + return DateTimeException(message) + elif "org.apache.spark.SparkRuntimeException" in classes: + return SparkRuntimeException(message) + elif "org.apache.spark.SparkUpgradeException" in classes: + return SparkUpgradeException(message) + elif "org.apache.spark.api.python.PythonException" in classes: + return PythonException( + "\n An exception was thrown from the Python worker. " + "Please see the stack trace below.\n%s" % message + ) + else: + return SparkConnectGrpcException(message, reason=info.reason) + + +class SparkConnectGrpcException(SparkConnectException): + """ + Base class to handle the errors from GRPC. + """ + + def __init__( + self, + message: Optional[str] = None, + error_class: Optional[str] = None, + message_parameters: Optional[Dict[str, str]] = None, + reason: Optional[str] = None, + ) -> None: + self.message = message # type: ignore[assignment] + if reason is not None: + self.message = f"({reason}) {self.message}" + + super().__init__( + message=self.message, + error_class=error_class, + message_parameters=message_parameters, + ) + + +class AnalysisException(SparkConnectGrpcException, BaseAnalysisException): + """ + Failed to analyze a SQL query plan, thrown from Spark Connect. + """ + + +class ParseException(AnalysisException, BaseParseException): + """ + Failed to parse a SQL command, thrown from Spark Connect. + """ + + +class IllegalArgumentException(SparkConnectGrpcException, BaseIllegalArgumentException): + """ + Passed an illegal or inappropriate argument, thrown from Spark Connect. + """ + + +class StreamingQueryException(SparkConnectGrpcException, BaseStreamingQueryException): + """ + Exception that stopped a :class:`StreamingQuery` thrown from Spark Connect. + """ + + +class QueryExecutionException(SparkConnectGrpcException, BaseQueryExecutionException): + """ + Failed to execute a query, thrown from Spark Connect. + """ + + +class PythonException(SparkConnectGrpcException, BasePythonException): + """ + Exceptions thrown from Spark Connect. + """ + + +class ArithmeticException(SparkConnectGrpcException, BaseArithmeticException): + """ + Arithmetic exception thrown from Spark Connect. + """ + + +class UnsupportedOperationException(SparkConnectGrpcException, BaseUnsupportedOperationException): + """ + Unsupported operation exception thrown from Spark Connect. + """ + + +class ArrayIndexOutOfBoundsException(SparkConnectGrpcException, BaseArrayIndexOutOfBoundsException): + """ + Array index out of bounds exception thrown from Spark Connect. + """ + + +class DateTimeException(SparkConnectGrpcException, BaseDateTimeException): + """ + Datetime exception thrown from Spark Connect. + """ + + +class NumberFormatException(IllegalArgumentException, BaseNumberFormatException): + """ + Number format exception thrown from Spark Connect. + """ + + +class SparkRuntimeException(SparkConnectGrpcException, BaseSparkRuntimeException): + """ + Runtime exception thrown from Spark Connect. + """ + + +class SparkUpgradeException(SparkConnectGrpcException, BaseSparkUpgradeException): + """ + Exception thrown because of Spark upgrade from Spark Connect. + """ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/test_errors.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/test_errors.py new file mode 100644 index 0000000..4e743bf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/tests/test_errors.py @@ -0,0 +1,60 @@ +# -*- encoding: utf-8 -*- +# +# 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 json +import unittest + +from pyspark.errors.error_classes import ERROR_CLASSES_JSON +from pyspark.errors.utils import ErrorClassesReader + + +class ErrorsTest(unittest.TestCase): + def test_error_classes_sorted(self): + # Test error classes is sorted alphabetically + error_reader = ErrorClassesReader() + error_class_names = list(error_reader.error_info_map.keys()) + for i in range(len(error_class_names) - 1): + self.assertTrue( + error_class_names[i] < error_class_names[i + 1], + f"Error class [{error_class_names[i]}] should place" + f"after [{error_class_names[i + 1]}]", + ) + + def test_error_classes_duplicated(self): + # Test error classes is not duplicated + def detect_duplication(pairs): + error_classes_json = {} + for name, message in pairs: + self.assertTrue(name not in error_classes_json, f"Duplicate error class: {name}") + error_classes_json[name] = message + return error_classes_json + + json.loads(ERROR_CLASSES_JSON, object_pairs_hook=detect_duplication) + + +if __name__ == "__main__": + import unittest + from pyspark.errors.tests.test_errors import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/errors/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/utils.py new file mode 100644 index 0000000..a4894dc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/errors/utils.py @@ -0,0 +1,116 @@ +# +# 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 re +from typing import Dict + +from pyspark.errors.error_classes import ERROR_CLASSES_MAP + + +class ErrorClassesReader: + """ + A reader to load error information from error_classes.py. + """ + + def __init__(self) -> None: + self.error_info_map = ERROR_CLASSES_MAP + + def get_error_message(self, error_class: str, message_parameters: Dict[str, str]) -> str: + """ + Returns the completed error message by applying message parameters to the message template. + """ + message_template = self.get_message_template(error_class) + # Verify message parameters. + message_parameters_from_template = re.findall("<([a-zA-Z0-9_-]+)>", message_template) + assert set(message_parameters_from_template) == set(message_parameters), ( + f"Undefined error message parameter for error class: {error_class}. " + f"Parameters: {message_parameters}" + ) + table = str.maketrans("<>", "{}") + + return message_template.translate(table).format(**message_parameters) + + def get_message_template(self, error_class: str) -> str: + """ + Returns the message template for corresponding error class from error_classes.py. + + For example, + when given `error_class` is "EXAMPLE_ERROR_CLASS", + and corresponding error class in error_classes.py looks like the below: + + .. code-block:: python + + "EXAMPLE_ERROR_CLASS" : { + "message" : [ + "Problem because of ." + ] + } + + In this case, this function returns: + "Problem because of ." + + For sub error class, when given `error_class` is "EXAMPLE_ERROR_CLASS.SUB_ERROR_CLASS", + and corresponding error class in error_classes.py looks like the below: + + .. code-block:: python + + "EXAMPLE_ERROR_CLASS" : { + "message" : [ + "Problem because of ." + ], + "sub_class" : { + "SUB_ERROR_CLASS" : { + "message" : [ + "Do to fix the problem." + ] + } + } + } + + In this case, this function returns: + "Problem because . Do to fix the problem." + """ + error_classes = error_class.split(".") + len_error_classes = len(error_classes) + assert len_error_classes in (1, 2) + + # Generate message template for main error class. + main_error_class = error_classes[0] + if main_error_class in self.error_info_map: + main_error_class_info_map = self.error_info_map[main_error_class] + else: + raise ValueError(f"Cannot find main error class '{main_error_class}'") + + main_message_template = "\n".join(main_error_class_info_map["message"]) + + has_sub_class = len_error_classes == 2 + + if not has_sub_class: + message_template = main_message_template + else: + # Generate message template for sub error class if exists. + sub_error_class = error_classes[1] + main_error_class_subclass_info_map = main_error_class_info_map["sub_class"] + if sub_error_class in main_error_class_subclass_info_map: + sub_error_class_info_map = main_error_class_subclass_info_map[sub_error_class] + else: + raise ValueError(f"Cannot find sub error class '{sub_error_class}'") + + sub_message_template = "\n".join(sub_error_class_info_map["message"]) + message_template = main_message_template + " " + sub_message_template + + return message_template diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/files.py b/spark-3.5.0-bin-hadoop3/python/pyspark/files.py new file mode 100644 index 0000000..8044cf4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/files.py @@ -0,0 +1,165 @@ +# +# 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 os + + +__all__ = ["SparkFiles"] + +from typing import cast, ClassVar, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from pyspark import SparkContext + + +class SparkFiles: + + """ + Resolves paths to files added through :meth:`SparkContext.addFile`. + + SparkFiles contains only classmethods; users should not create SparkFiles + instances. + """ + + _root_directory: ClassVar[Optional[str]] = None + _is_running_on_worker: ClassVar[bool] = False + _sc: ClassVar[Optional["SparkContext"]] = None + + def __init__(self) -> None: + raise NotImplementedError("Do not construct SparkFiles objects") + + @classmethod + def get(cls, filename: str) -> str: + """ + Get the absolute path of a file added through + :meth:`SparkContext.addFile` or :meth:`SparkContext.addPyFile`. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + filename : str + file that are added to resources + + Returns + ------- + str + the absolute path of the file + + See Also + -------- + :meth:`SparkFiles.getRootDirectory` + :meth:`SparkContext.addFile` + :meth:`SparkContext.addPyFile` + :meth:`SparkContext.listFiles` + + Examples + -------- + >>> import os + >>> import tempfile + >>> from pyspark import SparkFiles + + >>> with tempfile.TemporaryDirectory() as d: + ... path1 = os.path.join(d, "test.txt") + ... with open(path1, "w") as f: + ... _ = f.write("100") + ... + ... sc.addFile(path1) + ... file_list1 = sorted(sc.listFiles) + ... + ... def func1(iterator): + ... path = SparkFiles.get("test.txt") + ... assert path.startswith(SparkFiles.getRootDirectory()) + ... return [path] + ... + ... path_list1 = sc.parallelize([1, 2, 3, 4]).mapPartitions(func1).collect() + ... + ... path2 = os.path.join(d, "test.py") + ... with open(path2, "w") as f: + ... _ = f.write("import pyspark") + ... + ... # py files + ... sc.addPyFile(path2) + ... file_list2 = sorted(sc.listFiles) + ... + ... def func2(iterator): + ... path = SparkFiles.get("test.py") + ... assert path.startswith(SparkFiles.getRootDirectory()) + ... return [path] + ... + ... path_list2 = sc.parallelize([1, 2, 3, 4]).mapPartitions(func2).collect() + >>> file_list1 + ['file:/.../test.txt'] + >>> set(path_list1) + {'.../test.txt'} + >>> file_list2 + ['file:/.../test.py', 'file:/.../test.txt'] + >>> set(path_list2) + {'.../test.py'} + """ + path = os.path.join(SparkFiles.getRootDirectory(), filename) + return os.path.abspath(path) + + @classmethod + def getRootDirectory(cls) -> str: + """ + Get the root directory that contains files added through + :meth:`SparkContext.addFile` or :meth:`SparkContext.addPyFile`. + + .. versionadded:: 0.7.0 + + Returns + ------- + str + the root directory that contains files added to resources + + See Also + -------- + :meth:`SparkFiles.get` + :meth:`SparkContext.addFile` + :meth:`SparkContext.addPyFile` + + Examples + -------- + >>> from pyspark.files import SparkFiles + >>> SparkFiles.getRootDirectory() # doctest: +SKIP + '.../spark-a904728e-08d3-400c-a872-cfd82fd6dcd2/userFiles-648cf6d6-bb2c-4f53-82bd-e658aba0c5de' + """ + if cls._is_running_on_worker: + return cast(str, cls._root_directory) + else: + # This will have to change if we support multiple SparkContexts: + assert cls._sc is not None + assert cls._sc._jvm is not None + return cls._sc._jvm.org.apache.spark.SparkFiles.getRootDirectory() + + +def _test() -> None: + import doctest + import sys + from pyspark import SparkContext + + globs = globals().copy() + globs["sc"] = SparkContext("local[2]", "files tests") + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/find_spark_home.py b/spark-3.5.0-bin-hadoop3/python/pyspark/find_spark_home.py new file mode 100644 index 0000000..a2226f8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/find_spark_home.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 + +# +# 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. +# + +# This script attempt to determine the correct setting for SPARK_HOME given +# that Spark may have been installed on the system with pip. + +import os +import sys + + +def _find_spark_home(): + """Find the SPARK_HOME.""" + # If the environment has SPARK_HOME set trust it. + if "SPARK_HOME" in os.environ: + return os.environ["SPARK_HOME"] + + def is_spark_home(path): + """Takes a path and returns true if the provided path could be a reasonable SPARK_HOME""" + return os.path.isfile(os.path.join(path, "bin/spark-submit")) and ( + os.path.isdir(os.path.join(path, "jars")) + or os.path.isdir(os.path.join(path, "assembly")) + ) + + # Spark distribution can be downloaded when PYSPARK_HADOOP_VERSION environment variable is set. + # We should look up this directory first, see also SPARK-32017. + spark_dist_dir = "spark-distribution" + paths = [ + "../", # When we're in spark/python. + ] + + if "__file__" in globals(): + paths += [ + # Two case belows are valid when the current script is called as a library. + os.path.join(os.path.dirname(os.path.realpath(__file__)), spark_dist_dir), + os.path.dirname(os.path.realpath(__file__)), + ] + + # Add the path of the PySpark module if it exists + import_error_raised = False + from importlib.util import find_spec + + try: + module_home = os.path.dirname(find_spec("pyspark").origin) + paths.append(os.path.join(module_home, spark_dist_dir)) + paths.append(module_home) + # If we are installed in edit mode also look two dirs up + # Downloading different versions are not supported in edit mode. + paths.append(os.path.join(module_home, "../../")) + except ImportError: + # Not pip installed no worries + import_error_raised = True + + # Normalize the paths + paths = [os.path.abspath(p) for p in paths] + + try: + return next(path for path in paths if is_spark_home(path)) + except StopIteration: + print("Could not find valid SPARK_HOME while searching {0}".format(paths), file=sys.stderr) + if import_error_raised: + print( + "\nDid you install PySpark via a package manager such as pip or Conda? If so,\n" + "PySpark was not found in your Python environment. It is possible your\n" + "Python environment does not properly bind with your package manager.\n" + "\nPlease check your default 'python' and if you set PYSPARK_PYTHON and/or\n" + "PYSPARK_DRIVER_PYTHON environment variables, and see if you can import\n" + "PySpark, for example, 'python -c 'import pyspark'.\n" + "\nIf you cannot import, you can install by using the Python executable directly,\n" + "for example, 'python -m pip install pyspark [--user]'. Otherwise, you can also\n" + "explicitly set the Python executable, that has PySpark installed, to\n" + "PYSPARK_PYTHON or PYSPARK_DRIVER_PYTHON environment variables, for example,\n" + "'PYSPARK_PYTHON=python3 pyspark'.\n", + file=sys.stderr, + ) + sys.exit(-1) + + +if __name__ == "__main__": + print(_find_spark_home()) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/install.py b/spark-3.5.0-bin-hadoop3/python/pyspark/install.py new file mode 100644 index 0000000..90b0150 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/install.py @@ -0,0 +1,203 @@ +# +# 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 os +import re +import tarfile +import traceback +import urllib.request +from shutil import rmtree + +# NOTE that we shouldn't import pyspark here because this is used in +# setup.py, and assume there's no PySpark imported. + +DEFAULT_HADOOP = "hadoop3" +DEFAULT_HIVE = "hive2.3" +SUPPORTED_HADOOP_VERSIONS = ["hadoop3", "without-hadoop"] +SUPPORTED_HIVE_VERSIONS = ["hive2.3"] +UNSUPPORTED_COMBINATIONS = [] # type: ignore + + +def checked_package_name(spark_version, hadoop_version, hive_version): + """ + Check the generated package name, here we need to use the final hadoop version. + """ + return "%s-bin-%s" % (spark_version, hadoop_version) + + +def checked_versions(spark_version, hadoop_version, hive_version): + """ + Check the valid combinations of supported versions in Spark distributions. + + Parameters + ---------- + spark_version : str + Spark version. It should be X.X.X such as '3.0.0' or spark-3.0.0. + hadoop_version : str + Hadoop version. It should be X such as '2' or 'hadoop2'. + 'without' and 'without-hadoop' are supported as special keywords for Hadoop free + distribution. + hive_version : str + Hive version. It should be X.X such as '2.3' or 'hive2.3'. + + Parameters + ---------- + tuple + fully-qualified versions of Spark, Hadoop and Hive in a tuple. + For example, spark-3.2.0, hadoop3 and hive2.3. + """ + if re.match("^[0-9]+\\.[0-9]+\\.[0-9]+$", spark_version): + spark_version = "spark-%s" % spark_version + if not spark_version.startswith("spark-"): + raise RuntimeError( + "Spark version should start with 'spark-' prefix; however, " "got %s" % spark_version + ) + + if hadoop_version == "without": + hadoop_version = "without-hadoop" + elif re.match("^[0-9]+$", hadoop_version): + hadoop_version = "hadoop%s" % hadoop_version + + if hadoop_version not in SUPPORTED_HADOOP_VERSIONS: + raise RuntimeError( + "Spark distribution of %s is not supported. Hadoop version should be " + "one of [%s]" % (hadoop_version, ", ".join(SUPPORTED_HADOOP_VERSIONS)) + ) + + if re.match("^[0-9]+\\.[0-9]+$", hive_version): + hive_version = "hive%s" % hive_version + + if hive_version not in SUPPORTED_HIVE_VERSIONS: + raise RuntimeError( + "Spark distribution of %s is not supported. Hive version should be " + "one of [%s]" % (hive_version, ", ".join(SUPPORTED_HADOOP_VERSIONS)) + ) + + return spark_version, convert_old_hadoop_version(spark_version, hadoop_version), hive_version + + +def convert_old_hadoop_version(spark_version, hadoop_version): + # check if Spark version <= 3.2, if so, convert hadoop3 to hadoop3.2 and hadoop2 to hadoop2.7 + version_dict = { + "hadoop3": "hadoop3.2", + "hadoop2": "hadoop2.7", + "without": "without", + "without-hadoop": "without-hadoop", + } + spark_version_parts = re.search("^spark-([0-9]+)\\.([0-9]+)\\.[0-9]+$", spark_version) + spark_major_version = int(spark_version_parts.group(1)) + spark_minor_version = int(spark_version_parts.group(2)) + if spark_major_version < 3 or (spark_major_version == 3 and spark_minor_version <= 2): + hadoop_version = version_dict[hadoop_version] + return hadoop_version + + +def install_spark(dest, spark_version, hadoop_version, hive_version): + """ + Installs Spark that corresponds to the given Hadoop version in the current + library directory. + + Parameters + ---------- + dest : str + The location to download and install the Spark. + spark_version : str + Spark version. It should be spark-X.X.X form. + hadoop_version : str + Hadoop version. It should be hadoopX.X + such as 'hadoop2.7' or 'without-hadoop'. + hive_version : str + Hive version. It should be hiveX.X such as 'hive2.3'. + """ + + package_name = checked_package_name(spark_version, hadoop_version, hive_version) + package_local_path = os.path.join(dest, "%s.tgz" % package_name) + if "PYSPARK_RELEASE_MIRROR" in os.environ: + sites = [os.environ["PYSPARK_RELEASE_MIRROR"]] + else: + sites = get_preferred_mirrors() + print("Trying to download Spark %s from [%s]" % (spark_version, ", ".join(sites))) + + pretty_pkg_name = "%s for Hadoop %s" % ( + spark_version, + "Free build" if hadoop_version == "without" else hadoop_version, + ) + + for site in sites: + os.makedirs(dest, exist_ok=True) + url = "%s/spark/%s/%s.tgz" % (site, spark_version, package_name) + + tar = None + try: + print("Downloading %s from:\n- %s" % (pretty_pkg_name, url)) + download_to_file(urllib.request.urlopen(url), package_local_path) + + print("Installing to %s" % dest) + tar = tarfile.open(package_local_path, "r:gz") + for member in tar.getmembers(): + if member.name == package_name: + # Skip the root directory. + continue + member.name = os.path.relpath(member.name, package_name + os.path.sep) + tar.extract(member, dest) + return + except Exception: + print("Failed to download %s from %s:" % (pretty_pkg_name, url)) + traceback.print_exc() + rmtree(dest, ignore_errors=True) + finally: + if tar is not None: + tar.close() + if os.path.exists(package_local_path): + os.remove(package_local_path) + raise IOError("Unable to download %s." % pretty_pkg_name) + + +def get_preferred_mirrors(): + mirror_urls = [] + for _ in range(3): + try: + response = urllib.request.urlopen( + "https://www.apache.org/dyn/closer.lua?preferred=true" + ) + mirror_urls.append(response.read().decode("utf-8")) + except Exception: + # If we can't get a mirror URL, skip it. No retry. + pass + + default_sites = [ + "https://dlcdn.apache.org/", + "https://archive.apache.org/dist", + "https://dist.apache.org/repos/dist/release", + ] + return list(set(mirror_urls)) + [x for x in default_sites if x not in mirror_urls] + + +def download_to_file(response, path, chunk_size=1024 * 1024): + total_size = int(response.info().get("Content-Length").strip()) + bytes_so_far = 0 + + with open(path, mode="wb") as dest: + while True: + chunk = response.read(chunk_size) + bytes_so_far += len(chunk) + if not chunk: + break + dest.write(chunk) + print( + "Downloaded %d of %d bytes (%0.2f%%)" + % (bytes_so_far, total_size, round(float(bytes_so_far) / total_size * 100, 2)) + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/instrumentation_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/instrumentation_utils.py new file mode 100644 index 0000000..256c090 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/instrumentation_utils.py @@ -0,0 +1,190 @@ +# -*- coding: utf-8 -*- +# +# 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 functools +import inspect +import threading +import importlib +import time +from types import ModuleType +from typing import Tuple, Union, List, Callable, Any, Type + + +__all__: List[str] = [] + +_local = threading.local() + + +def _wrap_function(class_name: str, function_name: str, func: Callable, logger: Any) -> Callable: + + signature = inspect.signature(func) + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if hasattr(_local, "logging") and _local.logging: + # no need to log since this should be internal call. + return func(*args, **kwargs) + _local.logging = True + try: + start = time.perf_counter() + try: + res = func(*args, **kwargs) + logger.log_success( + class_name, function_name, time.perf_counter() - start, signature + ) + return res + except Exception as ex: + logger.log_failure( + class_name, function_name, ex, time.perf_counter() - start, signature + ) + raise + finally: + _local.logging = False + + return wrapper + + +def _wrap_property(class_name: str, property_name: str, prop: Any, logger: Any) -> Any: + @property # type: ignore[misc] + def wrapper(self: Any) -> Any: + if hasattr(_local, "logging") and _local.logging: + # no need to log since this should be internal call. + return prop.fget(self) + _local.logging = True + try: + start = time.perf_counter() + try: + res = prop.fget(self) + logger.log_success(class_name, property_name, time.perf_counter() - start) + return res + except Exception as ex: + logger.log_failure(class_name, property_name, ex, time.perf_counter() - start) + raise + finally: + _local.logging = False + + wrapper.__doc__ = prop.__doc__ + + if prop.fset is not None: + wrapper = wrapper.setter( # type: ignore[attr-defined] + _wrap_function(class_name, prop.fset.__name__, prop.fset, logger) + ) + + return wrapper + + +def _wrap_missing_function( + class_name: str, function_name: str, func: Callable, original: Any, logger: Any +) -> Any: + + if not hasattr(original, function_name): + return func + + signature = inspect.signature(getattr(original, function_name)) + + is_deprecated = func.__name__ == "deprecated_function" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return func(*args, **kwargs) + finally: + logger.log_missing(class_name, function_name, is_deprecated, signature) + + return wrapper + + +def _wrap_missing_property(class_name: str, property_name: str, prop: Any, logger: Any) -> Any: + + is_deprecated = prop.fget.__name__ == "deprecated_property" + + @property # type: ignore[misc] + def wrapper(self: Any) -> Any: + try: + return prop.fget(self) + finally: + logger.log_missing(class_name, property_name, is_deprecated) + + return wrapper + + +def _attach( + logger_module: Union[str, ModuleType], + modules: List[ModuleType], + classes: List[Type[Any]], + missings: List[Tuple[Type[Any], Type[Any]]], +) -> None: + if isinstance(logger_module, str): + logger_module = importlib.import_module(logger_module) + + logger = getattr(logger_module, "get_logger")() + + special_functions = set( + [ + "__init__", + "__repr__", + "__str__", + "_repr_html_", + "__len__", + "__getitem__", + "__setitem__", + "__getattr__", + "__enter__", + "__exit__", + ] + ) + + # Modules + for target_module in modules: + target_name = target_module.__name__.split(".")[-1] + for name in getattr(target_module, "__all__"): + func = getattr(target_module, name) + if not inspect.isfunction(func): + continue + setattr(target_module, name, _wrap_function(target_name, name, func, logger)) + + # Classes + for target_class in classes: + for name, func in inspect.getmembers(target_class, inspect.isfunction): + if name.startswith("_") and name not in special_functions: + continue + try: + isstatic = isinstance(inspect.getattr_static(target_class, name), staticmethod) + except AttributeError: + isstatic = False + wrapped_function = _wrap_function(target_class.__name__, name, func, logger) + setattr( + target_class, name, staticmethod(wrapped_function) if isstatic else wrapped_function + ) + + for name, prop in inspect.getmembers(target_class, lambda o: isinstance(o, property)): + if name.startswith("_"): + continue + setattr(target_class, name, _wrap_property(target_class.__name__, name, prop, logger)) + + # Missings + for original, missing in missings: + for name, func in inspect.getmembers(missing, inspect.isfunction): + setattr( + missing, + name, + _wrap_missing_function(original.__name__, name, func, original, logger), + ) + + for name, prop in inspect.getmembers(missing, lambda o: isinstance(o, property)): + setattr(missing, name, _wrap_missing_property(original.__name__, name, prop, logger)) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/java_gateway.py b/spark-3.5.0-bin-hadoop3/python/pyspark/java_gateway.py new file mode 100644 index 0000000..691120a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/java_gateway.py @@ -0,0 +1,248 @@ +# +# 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 atexit +import os +import signal +import shlex +import shutil +import socket +import platform +import tempfile +import time +from subprocess import Popen, PIPE + +from py4j.java_gateway import java_import, JavaGateway, JavaObject, GatewayParameters +from py4j.clientserver import ClientServer, JavaParameters, PythonParameters +from pyspark.find_spark_home import _find_spark_home +from pyspark.serializers import read_int, write_with_length, UTF8Deserializer +from pyspark.errors import PySparkRuntimeError + + +def launch_gateway(conf=None, popen_kwargs=None): + """ + launch jvm gateway + + Parameters + ---------- + conf : :py:class:`pyspark.SparkConf` + spark configuration passed to spark-submit + popen_kwargs : dict + Dictionary of kwargs to pass to Popen when spawning + the py4j JVM. This is a developer feature intended for use in + customizing how pyspark interacts with the py4j JVM (e.g., capturing + stdout/stderr). + + Returns + ------- + ClientServer or JavaGateway + """ + if "PYSPARK_GATEWAY_PORT" in os.environ: + gateway_port = int(os.environ["PYSPARK_GATEWAY_PORT"]) + gateway_secret = os.environ["PYSPARK_GATEWAY_SECRET"] + # Process already exists + proc = None + else: + SPARK_HOME = _find_spark_home() + # Launch the Py4j gateway using Spark's run command so that we pick up the + # proper classpath and settings from spark-env.sh + on_windows = platform.system() == "Windows" + script = "./bin/spark-submit.cmd" if on_windows else "./bin/spark-submit" + command = [os.path.join(SPARK_HOME, script)] + if conf: + for k, v in conf.getAll(): + command += ["--conf", "%s=%s" % (k, v)] + submit_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + if os.environ.get("SPARK_TESTING"): + submit_args = " ".join(["--conf spark.ui.enabled=false", submit_args]) + command = command + shlex.split(submit_args) + + # Create a temporary directory where the gateway server should write the connection + # information. + conn_info_dir = tempfile.mkdtemp() + try: + fd, conn_info_file = tempfile.mkstemp(dir=conn_info_dir) + os.close(fd) + os.unlink(conn_info_file) + + env = dict(os.environ) + env["_PYSPARK_DRIVER_CONN_INFO_PATH"] = conn_info_file + + # Launch the Java gateway. + popen_kwargs = {} if popen_kwargs is None else popen_kwargs + # We open a pipe to stdin so that the Java gateway can die when the pipe is broken + popen_kwargs["stdin"] = PIPE + # We always set the necessary environment variables. + popen_kwargs["env"] = env + if not on_windows: + # Don't send ctrl-c / SIGINT to the Java gateway: + def preexec_func(): + signal.signal(signal.SIGINT, signal.SIG_IGN) + + popen_kwargs["preexec_fn"] = preexec_func + proc = Popen(command, **popen_kwargs) + else: + # preexec_fn not supported on Windows + proc = Popen(command, **popen_kwargs) + + # Wait for the file to appear, or for the process to exit, whichever happens first. + while not proc.poll() and not os.path.isfile(conn_info_file): + time.sleep(0.1) + + if not os.path.isfile(conn_info_file): + raise PySparkRuntimeError( + error_class="JAVA_GATEWAY_EXITED", + message_parameters={}, + ) + + with open(conn_info_file, "rb") as info: + gateway_port = read_int(info) + gateway_secret = UTF8Deserializer().loads(info) + finally: + shutil.rmtree(conn_info_dir) + + # In Windows, ensure the Java child processes do not linger after Python has exited. + # In UNIX-based systems, the child process can kill itself on broken pipe (i.e. when + # the parent process' stdin sends an EOF). In Windows, however, this is not possible + # because java.lang.Process reads directly from the parent process' stdin, contending + # with any opportunity to read an EOF from the parent. Note that this is only best + # effort and will not take effect if the python process is violently terminated. + if on_windows: + # In Windows, the child process here is "spark-submit.cmd", not the JVM itself + # (because the UNIX "exec" command is not available). This means we cannot simply + # call proc.kill(), which kills only the "spark-submit.cmd" process but not the + # JVMs. Instead, we use "taskkill" with the tree-kill option "/t" to terminate all + # child processes in the tree (http://technet.microsoft.com/en-us/library/bb491009.aspx) + def killChild(): + Popen(["cmd", "/c", "taskkill", "/f", "/t", "/pid", str(proc.pid)]) + + atexit.register(killChild) + + # Connect to the gateway (or client server to pin the thread between JVM and Python) + if os.environ.get("PYSPARK_PIN_THREAD", "true").lower() == "true": + gateway = ClientServer( + java_parameters=JavaParameters( + port=gateway_port, auth_token=gateway_secret, auto_convert=True + ), + python_parameters=PythonParameters(port=0, eager_load=False), + ) + else: + gateway = JavaGateway( + gateway_parameters=GatewayParameters( + port=gateway_port, auth_token=gateway_secret, auto_convert=True + ) + ) + + # Store a reference to the Popen object for use by the caller (e.g., in reading stdout/stderr) + gateway.proc = proc + + # Import the classes used by PySpark + java_import(gateway.jvm, "org.apache.spark.SparkConf") + java_import(gateway.jvm, "org.apache.spark.api.java.*") + java_import(gateway.jvm, "org.apache.spark.api.python.*") + java_import(gateway.jvm, "org.apache.spark.ml.python.*") + java_import(gateway.jvm, "org.apache.spark.mllib.api.python.*") + java_import(gateway.jvm, "org.apache.spark.resource.*") + # TODO(davies): move into sql + java_import(gateway.jvm, "org.apache.spark.sql.*") + java_import(gateway.jvm, "org.apache.spark.sql.api.python.*") + java_import(gateway.jvm, "org.apache.spark.sql.hive.*") + java_import(gateway.jvm, "scala.Tuple2") + + return gateway + + +def _do_server_auth(conn, auth_secret): + """ + Performs the authentication protocol defined by the SocketAuthHelper class on the given + file-like object 'conn'. + """ + write_with_length(auth_secret.encode("utf-8"), conn) + conn.flush() + reply = UTF8Deserializer().loads(conn) + if reply != "ok": + conn.close() + raise PySparkRuntimeError( + error_class="UNEXPECTED_RESPONSE_FROM_SERVER", + message_parameters={}, + ) + + +def local_connect_and_auth(port, auth_secret): + """ + Connect to local host, authenticate with it, and return a (sockfile,sock) for that connection. + Handles IPV4 & IPV6, does some error handling. + + Parameters + ---------- + port : str or int or None + auth_secret : str + + Returns + ------- + tuple + with (sockfile, sock) + """ + sock = None + errors = [] + # Support for both IPv4 and IPv6. + addr = "127.0.0.1" + if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true": + addr = "::1" + for res in socket.getaddrinfo(addr, port, socket.AF_UNSPEC, socket.SOCK_STREAM): + af, socktype, proto, _, sa = res + try: + sock = socket.socket(af, socktype, proto) + sock.settimeout(int(os.environ.get("SPARK_AUTH_SOCKET_TIMEOUT", 15))) + sock.connect(sa) + sockfile = sock.makefile("rwb", int(os.environ.get("SPARK_BUFFER_SIZE", 65536))) + _do_server_auth(sockfile, auth_secret) + return (sockfile, sock) + except socket.error as e: + emsg = str(e) + errors.append("tried to connect to %s, but an error occurred: %s" % (sa, emsg)) + sock.close() + sock = None + raise PySparkRuntimeError( + error_class="CANNOT_OPEN_SOCKET", + message_parameters={ + "errors": str(errors), + }, + ) + + +def ensure_callback_server_started(gw): + """ + Start callback server if not already started. The callback server is needed if the Java + driver process needs to callback into the Python driver process to execute Python code. + """ + + # getattr will fallback to JVM, so we cannot test by hasattr() + if "_callback_server" not in gw.__dict__ or gw._callback_server is None: + gw.callback_server_parameters.eager_load = True + gw.callback_server_parameters.daemonize = True + gw.callback_server_parameters.daemonize_connections = True + gw.callback_server_parameters.port = 0 + gw.start_callback_server(gw.callback_server_parameters) + cbport = gw._callback_server.server_socket.getsockname()[1] + gw._callback_server.port = cbport + # gateway with real port + gw._python_proxy_port = gw._callback_server.port + # get the GatewayServer object in JVM by ID + jgws = JavaObject("GATEWAY_SERVER", gw._gateway_client) + # update the port of CallbackClient with real port + jgws.resetCallbackClient(jgws.getCallbackClient().getAddress(), gw._python_proxy_port) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/join.py b/spark-3.5.0-bin-hadoop3/python/pyspark/join.py new file mode 100644 index 0000000..040c946 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/join.py @@ -0,0 +1,118 @@ +""" +Copyright (c) 2011, Douban Inc. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + + * Neither the name of the Douban Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +""" + +from pyspark.resultiterable import ResultIterable +from functools import reduce + + +def _do_python_join(rdd, other, numPartitions, dispatch): + vs = rdd.mapValues(lambda v: (1, v)) + ws = other.mapValues(lambda v: (2, v)) + return vs.union(ws).groupByKey(numPartitions).flatMapValues(lambda x: dispatch(x.__iter__())) + + +def python_join(rdd, other, numPartitions): + def dispatch(seq): + vbuf, wbuf = [], [] + for (n, v) in seq: + if n == 1: + vbuf.append(v) + elif n == 2: + wbuf.append(v) + return ((v, w) for v in vbuf for w in wbuf) + + return _do_python_join(rdd, other, numPartitions, dispatch) + + +def python_right_outer_join(rdd, other, numPartitions): + def dispatch(seq): + vbuf, wbuf = [], [] + for (n, v) in seq: + if n == 1: + vbuf.append(v) + elif n == 2: + wbuf.append(v) + if not vbuf: + vbuf.append(None) + return ((v, w) for v in vbuf for w in wbuf) + + return _do_python_join(rdd, other, numPartitions, dispatch) + + +def python_left_outer_join(rdd, other, numPartitions): + def dispatch(seq): + vbuf, wbuf = [], [] + for (n, v) in seq: + if n == 1: + vbuf.append(v) + elif n == 2: + wbuf.append(v) + if not wbuf: + wbuf.append(None) + return ((v, w) for v in vbuf for w in wbuf) + + return _do_python_join(rdd, other, numPartitions, dispatch) + + +def python_full_outer_join(rdd, other, numPartitions): + def dispatch(seq): + vbuf, wbuf = [], [] + for (n, v) in seq: + if n == 1: + vbuf.append(v) + elif n == 2: + wbuf.append(v) + if not vbuf: + vbuf.append(None) + if not wbuf: + wbuf.append(None) + return ((v, w) for v in vbuf for w in wbuf) + + return _do_python_join(rdd, other, numPartitions, dispatch) + + +def python_cogroup(rdds, numPartitions): + def make_mapper(i): + return lambda v: (i, v) + + vrdds = [rdd.mapValues(make_mapper(i)) for i, rdd in enumerate(rdds)] + union_vrdds = reduce(lambda acc, other: acc.union(other), vrdds) + rdd_len = len(vrdds) + + def dispatch(seq): + bufs = [[] for _ in range(rdd_len)] + for n, v in seq: + bufs[n].append(v) + return tuple(ResultIterable(vs) for vs in bufs) + + return union_vrdds.groupByKey(numPartitions).mapValues(dispatch) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/__init__.py new file mode 100644 index 0000000..167cf33 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/__init__.py @@ -0,0 +1,71 @@ +# +# 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. +# + +""" +DataFrame-based machine learning APIs to let users quickly assemble and configure practical +machine learning pipelines. +""" +from pyspark.ml.base import ( + Estimator, + Model, + Predictor, + PredictionModel, + Transformer, + UnaryTransformer, +) +from pyspark.ml.pipeline import Pipeline, PipelineModel +from pyspark.ml import ( + classification, + clustering, + evaluation, + feature, + fpm, + image, + recommendation, + regression, + stat, + tuning, + util, + linalg, + param, +) +from pyspark.ml.torch.distributor import TorchDistributor + +__all__ = [ + "Transformer", + "UnaryTransformer", + "Estimator", + "Model", + "Predictor", + "PredictionModel", + "Pipeline", + "PipelineModel", + "classification", + "clustering", + "evaluation", + "feature", + "fpm", + "image", + "recommendation", + "regression", + "stat", + "tuning", + "util", + "linalg", + "param", + "TorchDistributor", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/_typing.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/_typing.pyi new file mode 100644 index 0000000..12d831f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/_typing.pyi @@ -0,0 +1,84 @@ +# +# 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 Any, Dict, List, TypeVar, Tuple, Union +from typing_extensions import Literal + +from numpy import ndarray + +import pyspark.ml.base +import pyspark.ml.param +import pyspark.ml.util +from pyspark.ml.linalg import Vector +import pyspark.ml.wrapper +from py4j.java_gateway import JavaObject + +ParamMap = Dict[pyspark.ml.param.Param, Any] +PipelineStage = Union[pyspark.ml.base.Estimator, pyspark.ml.base.Transformer] + +T = TypeVar("T") +P = TypeVar("P", bound=pyspark.ml.param.Params) +M = TypeVar("M", bound=pyspark.ml.base.Transformer) +JM = TypeVar("JM", bound=pyspark.ml.wrapper.JavaTransformer) +C = TypeVar("C", bound=type) + +JavaObjectOrPickleDump = Union[JavaObject, bytearray, bytes] +BinaryClassificationEvaluatorMetricType = Union[Literal["areaUnderROC"], Literal["areaUnderPR"]] +RegressionEvaluatorMetricType = Union[ + Literal["rmse"], Literal["mse"], Literal["r2"], Literal["mae"], Literal["var"] +] +MulticlassClassificationEvaluatorMetricType = Union[ + Literal["f1"], + Literal["accuracy"], + Literal["weightedPrecision"], + Literal["weightedRecall"], + Literal["weightedTruePositiveRate"], + Literal["weightedFalsePositiveRate"], + Literal["weightedFMeasure"], + Literal["truePositiveRateByLabel"], + Literal["falsePositiveRateByLabel"], + Literal["precisionByLabel"], + Literal["recallByLabel"], + Literal["fMeasureByLabel"], +] +MultilabelClassificationEvaluatorMetricType = Union[ + Literal["subsetAccuracy"], + Literal["accuracy"], + Literal["hammingLoss"], + Literal["precision"], + Literal["recall"], + Literal["f1Measure"], + Literal["precisionByLabel"], + Literal["recallByLabel"], + Literal["f1MeasureByLabel"], + Literal["microPrecision"], + Literal["microRecall"], + Literal["microF1Measure"], +] +ClusteringEvaluatorMetricType = Literal["silhouette"] +ClusteringEvaluatorDistanceMeasureType = Union[Literal["squaredEuclidean"], Literal["cosine"]] + +RankingEvaluatorMetricType = Union[ + Literal["meanAveragePrecision"], + Literal["meanAveragePrecisionAtK"], + Literal["precisionAtK"], + Literal["ndcgAtK"], + Literal["recallAtK"], +] + +VectorLike = Union[ndarray, Vector, List[float], Tuple[float, ...]] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/base.py new file mode 100644 index 0000000..b94358d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/base.py @@ -0,0 +1,414 @@ +# +# 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 + +import copy +import threading + +from typing import ( + Any, + Callable, + Generic, + Iterator, + List, + Optional, + Sequence, + Tuple, + TypeVar, + Union, + cast, + overload, + TYPE_CHECKING, +) + +from pyspark import since +from pyspark.ml.param import P +from pyspark.ml.common import inherit_doc +from pyspark.ml.param.shared import ( + HasInputCol, + HasOutputCol, + HasLabelCol, + HasFeaturesCol, + HasPredictionCol, + Params, +) +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.functions import udf +from pyspark.sql.types import DataType, StructField, StructType + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + +T = TypeVar("T") +M = TypeVar("M", bound="Transformer") + + +class _FitMultipleIterator(Generic[M]): + """ + Used by default implementation of Estimator.fitMultiple to produce models in a thread safe + iterator. This class handles the simple case of fitMultiple where each param map should be + fit independently. + + Parameters + ---------- + fitSingleModel : function + Callable[[int], Transformer] which fits an estimator to a dataset. + `fitSingleModel` may be called up to `numModels` times, with a unique index each time. + Each call to `fitSingleModel` with an index should return the Model associated with + that index. + numModel : int + Number of models this iterator should produce. + + Notes + ----- + See :py:meth:`Estimator.fitMultiple` for more info. + """ + + def __init__(self, fitSingleModel: Callable[[int], M], numModels: int): + """ """ + self.fitSingleModel = fitSingleModel + self.numModel = numModels + self.counter = 0 + self.lock = threading.Lock() + + def __iter__(self) -> Iterator[Tuple[int, M]]: + return self + + def __next__(self) -> Tuple[int, M]: + with self.lock: + index = self.counter + if index >= self.numModel: + raise StopIteration("No models remaining.") + self.counter += 1 + return index, self.fitSingleModel(index) + + def next(self) -> Tuple[int, M]: + """For python2 compatibility.""" + return self.__next__() + + +@inherit_doc +class Estimator(Params, Generic[M], metaclass=ABCMeta): + """ + Abstract class for estimators that fit models to data. + + .. versionadded:: 1.3.0 + """ + + @abstractmethod + def _fit(self, dataset: 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 fitMultiple( + self, dataset: DataFrame, paramMaps: Sequence["ParamMap"] + ) -> Iterator[Tuple[int, M]]: + """ + Fits a model to the input dataset for each param map in `paramMaps`. + + .. versionadded:: 2.3.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset. + paramMaps : :py:class:`collections.abc.Sequence` + A Sequence of param maps. + + Returns + ------- + :py:class:`_FitMultipleIterator` + A thread safe iterable which contains one model for each param map. Each + call to `next(modelIterator)` will return `(index, model)` where model was fit + using `paramMaps[index]`. `index` values may not be sequential. + """ + estimator = self.copy() + + def fitSingleModel(index: int) -> M: + return estimator.fit(dataset, paramMaps[index]) + + return _FitMultipleIterator(fitSingleModel, len(paramMaps)) + + @overload + def fit(self, dataset: DataFrame, params: Optional["ParamMap"] = ...) -> M: + ... + + @overload + def fit( + self, dataset: DataFrame, params: Union[List["ParamMap"], Tuple["ParamMap"]] + ) -> List[M]: + ... + + def fit( + self, + dataset: DataFrame, + params: Optional[Union["ParamMap", List["ParamMap"], Tuple["ParamMap"]]] = None, + ) -> Union[M, List[M]]: + """ + Fits a model to the input dataset with optional parameters. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset. + params : dict or list or tuple, optional + an optional param map that overrides embedded params. If a list/tuple of + param maps is given, this calls fit on each param map and returns a list of + models. + + Returns + ------- + :py:class:`Transformer` or a list of :py:class:`Transformer` + fitted model(s) + """ + if params is None: + params = dict() + if isinstance(params, (list, tuple)): + models: List[Optional[M]] = [None] * len(params) + for index, model in self.fitMultiple(dataset, params): + models[index] = model + return cast(List[M], models) + elif 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) + ) + + +@inherit_doc +class Transformer(Params, metaclass=ABCMeta): + """ + Abstract class for transformers that transform one dataset into another. + + .. versionadded:: 1.3.0 + """ + + @abstractmethod + def _transform(self, dataset: DataFrame) -> DataFrame: + """ + Transforms the input dataset. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + transformed dataset + """ + raise NotImplementedError() + + def transform(self, dataset: DataFrame, params: Optional["ParamMap"] = None) -> DataFrame: + """ + Transforms the input dataset with optional parameters. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + params : dict, optional + an optional param map that overrides embedded params. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + transformed dataset + """ + if params is None: + params = dict() + if isinstance(params, dict): + if params: + return self.copy(params)._transform(dataset) + else: + return self._transform(dataset) + else: + raise TypeError("Params must be a param map but got %s." % type(params)) + + +@inherit_doc +class Model(Transformer, metaclass=ABCMeta): + """ + Abstract class for models that are fitted by estimators. + + .. versionadded:: 1.4.0 + """ + + pass + + +@inherit_doc +class UnaryTransformer(HasInputCol, HasOutputCol, Transformer): + """ + Abstract class for transformers that take one input column, apply transformation, + and output the result as a new column. + + .. versionadded:: 2.3.0 + """ + + def setInputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @abstractmethod + def createTransformFunc(self) -> Callable[..., Any]: + """ + Creates the transform function using the given param map. The input param map already takes + account of the embedded param map. So the param values should be determined + solely by the input param map. + """ + raise NotImplementedError() + + @abstractmethod + def outputDataType(self) -> DataType: + """ + Returns the data type of the output column. + """ + raise NotImplementedError() + + @abstractmethod + def validateInputType(self, inputType: DataType) -> None: + """ + Validates the input type. Throw an exception if it is invalid. + """ + raise NotImplementedError() + + def transformSchema(self, schema: StructType) -> StructType: + inputType = schema[self.getInputCol()].dataType + self.validateInputType(inputType) + if self.getOutputCol() in schema.names: + raise ValueError("Output column %s already exists." % self.getOutputCol()) + outputFields = copy.copy(schema.fields) + outputFields.append(StructField(self.getOutputCol(), self.outputDataType(), nullable=False)) + return StructType(outputFields) + + def _transform(self, dataset: DataFrame) -> DataFrame: + self.transformSchema(dataset.schema) + transformUDF = udf(self.createTransformFunc(), self.outputDataType()) + transformedDataset = dataset.withColumn( + self.getOutputCol(), transformUDF(dataset[self.getInputCol()]) + ) + return transformedDataset + + +@inherit_doc +class _PredictorParams(HasLabelCol, HasFeaturesCol, HasPredictionCol): + """ + Params for :py:class:`Predictor` and :py:class:`PredictorModel`. + + .. versionadded:: 3.0.0 + """ + + pass + + +@inherit_doc +class Predictor(Estimator[M], _PredictorParams, metaclass=ABCMeta): + """ + Estimator for prediction tasks (regression and classification). + """ + + @since("3.0.0") + def setLabelCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + @since("3.0.0") + def setFeaturesCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + +@inherit_doc +class PredictionModel(Model, _PredictorParams, Generic[T], metaclass=ABCMeta): + """ + Model for prediction tasks (regression and classification). + """ + + @since("3.0.0") + def setFeaturesCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @property + @abstractmethod + @since("2.1.0") + def numFeatures(self) -> int: + """ + Returns the number of features the model was trained on. If unknown, returns -1 + """ + raise NotImplementedError() + + @abstractmethod + @since("3.0.0") + def predict(self, value: T) -> float: + """ + Predict label for the given features. + """ + raise NotImplementedError() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/classification.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/classification.py new file mode 100644 index 0000000..81d7a1d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/classification.py @@ -0,0 +1,4332 @@ +# +# 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 os +import operator +import sys +import uuid +import warnings +from abc import ABCMeta, abstractmethod +from multiprocessing.pool import ThreadPool + +from typing import ( + Any, + Dict, + Generic, + Iterable, + List, + Optional, + Type, + TypeVar, + Union, + cast, + overload, + TYPE_CHECKING, +) + +from pyspark import keyword_only, since, SparkContext, inheritable_thread_target +from pyspark.ml import Estimator, Predictor, PredictionModel, Model +from pyspark.ml.param.shared import ( + HasRawPredictionCol, + HasProbabilityCol, + HasThresholds, + HasRegParam, + HasMaxIter, + HasFitIntercept, + HasTol, + HasStandardization, + HasWeightCol, + HasAggregationDepth, + HasThreshold, + HasBlockSize, + HasMaxBlockSizeInMB, + Param, + Params, + TypeConverters, + HasElasticNetParam, + HasSeed, + HasStepSize, + HasSolver, + HasParallelism, +) +from pyspark.ml.tree import ( + _DecisionTreeModel, + _DecisionTreeParams, + _TreeEnsembleModel, + _RandomForestParams, + _GBTParams, + _HasVarianceImpurity, + _TreeClassifierParams, +) +from pyspark.ml.regression import _FactorizationMachinesParams, DecisionTreeRegressionModel +from pyspark.ml.base import _PredictorParams +from pyspark.ml.util import ( + DefaultParamsReader, + DefaultParamsWriter, + JavaMLReadable, + JavaMLReader, + JavaMLWritable, + JavaMLWriter, + MLReader, + MLReadable, + MLWriter, + MLWritable, + HasTrainingSummary, +) +from pyspark.ml.wrapper import JavaParams, JavaPredictor, JavaPredictionModel, JavaWrapper +from pyspark.ml.common import inherit_doc +from pyspark.ml.linalg import Matrix, Vector, Vectors, VectorUDT +from pyspark.sql import DataFrame, Row +from pyspark.sql.functions import udf, when +from pyspark.sql.types import ArrayType, DoubleType +from pyspark.storagelevel import StorageLevel + + +if TYPE_CHECKING: + from pyspark.ml._typing import P, ParamMap + from py4j.java_gateway import JavaObject + + +T = TypeVar("T") +JPM = TypeVar("JPM", bound=JavaPredictionModel) +CM = TypeVar("CM", bound="ClassificationModel") + +__all__ = [ + "LinearSVC", + "LinearSVCModel", + "LinearSVCSummary", + "LinearSVCTrainingSummary", + "LogisticRegression", + "LogisticRegressionModel", + "LogisticRegressionSummary", + "LogisticRegressionTrainingSummary", + "BinaryLogisticRegressionSummary", + "BinaryLogisticRegressionTrainingSummary", + "DecisionTreeClassifier", + "DecisionTreeClassificationModel", + "GBTClassifier", + "GBTClassificationModel", + "RandomForestClassifier", + "RandomForestClassificationModel", + "RandomForestClassificationSummary", + "RandomForestClassificationTrainingSummary", + "BinaryRandomForestClassificationSummary", + "BinaryRandomForestClassificationTrainingSummary", + "NaiveBayes", + "NaiveBayesModel", + "MultilayerPerceptronClassifier", + "MultilayerPerceptronClassificationModel", + "MultilayerPerceptronClassificationSummary", + "MultilayerPerceptronClassificationTrainingSummary", + "OneVsRest", + "OneVsRestModel", + "FMClassifier", + "FMClassificationModel", + "FMClassificationSummary", + "FMClassificationTrainingSummary", +] + + +class _ClassifierParams(HasRawPredictionCol, _PredictorParams): + """ + Classifier Params for classification tasks. + + .. versionadded:: 3.0.0 + """ + + pass + + +@inherit_doc +class Classifier(Predictor[CM], _ClassifierParams, Generic[CM], metaclass=ABCMeta): + """ + Classifier for classification tasks. + Classes are indexed {0, 1, ..., numClasses - 1}. + """ + + @since("3.0.0") + def setRawPredictionCol(self: "P", value: str) -> "P": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + +@inherit_doc +class ClassificationModel(PredictionModel, _ClassifierParams, metaclass=ABCMeta): + """ + Model produced by a ``Classifier``. + Classes are indexed {0, 1, ..., numClasses - 1}. + """ + + @since("3.0.0") + def setRawPredictionCol(self: "P", value: str) -> "P": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + @property + @abstractmethod + @since("2.1.0") + def numClasses(self) -> int: + """ + Number of classes (values which the label can take). + """ + raise NotImplementedError() + + @abstractmethod + @since("3.0.0") + def predictRaw(self, value: Vector) -> Vector: + """ + Raw prediction for each possible label. + """ + raise NotImplementedError() + + +class _ProbabilisticClassifierParams(HasProbabilityCol, HasThresholds, _ClassifierParams): + """ + Params for :py:class:`ProbabilisticClassifier` and + :py:class:`ProbabilisticClassificationModel`. + + .. versionadded:: 3.0.0 + """ + + pass + + +@inherit_doc +class ProbabilisticClassifier(Classifier, _ProbabilisticClassifierParams, metaclass=ABCMeta): + """ + Probabilistic Classifier for classification tasks. + """ + + @since("3.0.0") + def setProbabilityCol(self: "P", value: str) -> "P": + """ + Sets the value of :py:attr:`probabilityCol`. + """ + return self._set(probabilityCol=value) + + @since("3.0.0") + def setThresholds(self: "P", value: List[float]) -> "P": + """ + Sets the value of :py:attr:`thresholds`. + """ + return self._set(thresholds=value) + + +@inherit_doc +class ProbabilisticClassificationModel( + ClassificationModel, _ProbabilisticClassifierParams, metaclass=ABCMeta +): + """ + Model produced by a ``ProbabilisticClassifier``. + """ + + @since("3.0.0") + def setProbabilityCol(self: CM, value: str) -> CM: + """ + Sets the value of :py:attr:`probabilityCol`. + """ + return self._set(probabilityCol=value) + + @since("3.0.0") + def setThresholds(self: CM, value: List[float]) -> CM: + """ + Sets the value of :py:attr:`thresholds`. + """ + return self._set(thresholds=value) + + @abstractmethod + @since("3.0.0") + def predictProbability(self, value: Vector) -> Vector: + """ + Predict the probability of each class given the features. + """ + raise NotImplementedError() + + +@inherit_doc +class _JavaClassifier(Classifier, JavaPredictor[JPM], Generic[JPM], metaclass=ABCMeta): + """ + Java Classifier for classification tasks. + Classes are indexed {0, 1, ..., numClasses - 1}. + """ + + @since("3.0.0") + def setRawPredictionCol(self: "P", value: str) -> "P": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + +@inherit_doc +class _JavaClassificationModel(ClassificationModel, JavaPredictionModel[T]): + """ + Java Model produced by a ``Classifier``. + Classes are indexed {0, 1, ..., numClasses - 1}. + To be mixed in with :class:`pyspark.ml.JavaModel` + """ + + @property + @since("2.1.0") + def numClasses(self) -> int: + """ + Number of classes (values which the label can take). + """ + return self._call_java("numClasses") + + @since("3.0.0") + def predictRaw(self, value: Vector) -> Vector: + """ + Raw prediction for each possible label. + """ + return self._call_java("predictRaw", value) + + +@inherit_doc +class _JavaProbabilisticClassifier( + ProbabilisticClassifier, _JavaClassifier[JPM], Generic[JPM], metaclass=ABCMeta +): + """ + Java Probabilistic Classifier for classification tasks. + """ + + pass + + +@inherit_doc +class _JavaProbabilisticClassificationModel( + ProbabilisticClassificationModel, _JavaClassificationModel[T] +): + """ + Java Model produced by a ``ProbabilisticClassifier``. + """ + + @since("3.0.0") + def predictProbability(self, value: Vector) -> Vector: + """ + Predict the probability of each class given the features. + """ + return self._call_java("predictProbability", value) + + +@inherit_doc +class _ClassificationSummary(JavaWrapper): + """ + Abstraction for multiclass classification results for a given model. + + .. versionadded:: 3.1.0 + """ + + @property + @since("3.1.0") + def predictions(self) -> DataFrame: + """ + Dataframe outputted by the model's `transform` method. + """ + return self._call_java("predictions") + + @property + @since("3.1.0") + def predictionCol(self) -> str: + """ + Field in "predictions" which gives the prediction of each class. + """ + return self._call_java("predictionCol") + + @property + @since("3.1.0") + def labelCol(self) -> str: + """ + Field in "predictions" which gives the true label of each + instance. + """ + return self._call_java("labelCol") + + @property + @since("3.1.0") + def weightCol(self) -> str: + """ + Field in "predictions" which gives the weight of each instance + as a vector. + """ + return self._call_java("weightCol") + + @property + def labels(self) -> List[str]: + """ + Returns the sequence of labels in ascending order. This order matches the order used + in metrics which are specified as arrays over labels, e.g., truePositiveRateByLabel. + + .. versionadded:: 3.1.0 + + Notes + ----- + In most cases, it will be values {0.0, 1.0, ..., numClasses-1}, However, if the + training set is missing a label, then all of the arrays over labels + (e.g., from truePositiveRateByLabel) will be of length numClasses-1 instead of the + expected numClasses. + """ + return self._call_java("labels") + + @property + @since("3.1.0") + def truePositiveRateByLabel(self) -> List[float]: + """ + Returns true positive rate for each label (category). + """ + return self._call_java("truePositiveRateByLabel") + + @property + @since("3.1.0") + def falsePositiveRateByLabel(self) -> List[float]: + """ + Returns false positive rate for each label (category). + """ + return self._call_java("falsePositiveRateByLabel") + + @property + @since("3.1.0") + def precisionByLabel(self) -> List[float]: + """ + Returns precision for each label (category). + """ + return self._call_java("precisionByLabel") + + @property + @since("3.1.0") + def recallByLabel(self) -> List[float]: + """ + Returns recall for each label (category). + """ + return self._call_java("recallByLabel") + + @since("3.1.0") + def fMeasureByLabel(self, beta: float = 1.0) -> List[float]: + """ + Returns f-measure for each label (category). + """ + return self._call_java("fMeasureByLabel", beta) + + @property + @since("3.1.0") + def accuracy(self) -> float: + """ + Returns accuracy. + (equals to the total number of correctly classified instances + out of the total number of instances.) + """ + return self._call_java("accuracy") + + @property + @since("3.1.0") + def weightedTruePositiveRate(self) -> float: + """ + Returns weighted true positive rate. + (equals to precision, recall and f-measure) + """ + return self._call_java("weightedTruePositiveRate") + + @property + @since("3.1.0") + def weightedFalsePositiveRate(self) -> float: + """ + Returns weighted false positive rate. + """ + return self._call_java("weightedFalsePositiveRate") + + @property + @since("3.1.0") + def weightedRecall(self) -> float: + """ + Returns weighted averaged recall. + (equals to precision, recall and f-measure) + """ + return self._call_java("weightedRecall") + + @property + @since("3.1.0") + def weightedPrecision(self) -> float: + """ + Returns weighted averaged precision. + """ + return self._call_java("weightedPrecision") + + @since("3.1.0") + def weightedFMeasure(self, beta: float = 1.0) -> float: + """ + Returns weighted averaged f-measure. + """ + return self._call_java("weightedFMeasure", beta) + + +@inherit_doc +class _TrainingSummary(JavaWrapper): + """ + Abstraction for Training results. + + .. versionadded:: 3.1.0 + """ + + @property + @since("3.1.0") + def objectiveHistory(self) -> List[float]: + """ + Objective function (scaled loss + regularization) at each + iteration. It contains one more element, the initial state, + than number of iterations. + """ + return self._call_java("objectiveHistory") + + @property + @since("3.1.0") + def totalIterations(self) -> int: + """ + Number of training iterations until termination. + """ + return self._call_java("totalIterations") + + +@inherit_doc +class _BinaryClassificationSummary(_ClassificationSummary): + """ + Binary classification results for a given model. + + .. versionadded:: 3.1.0 + """ + + @property + @since("3.1.0") + def scoreCol(self) -> str: + """ + Field in "predictions" which gives the probability or raw prediction + of each class as a vector. + """ + return self._call_java("scoreCol") + + @property + def roc(self) -> DataFrame: + """ + Returns the receiver operating characteristic (ROC) curve, + which is a Dataframe having two fields (FPR, TPR) with + (0.0, 0.0) prepended and (1.0, 1.0) appended to it. + + .. versionadded:: 3.1.0 + + Notes + ----- + `Wikipedia reference `_ + """ + return self._call_java("roc") + + @property + @since("3.1.0") + def areaUnderROC(self) -> float: + """ + Computes the area under the receiver operating characteristic + (ROC) curve. + """ + return self._call_java("areaUnderROC") + + @property + @since("3.1.0") + def pr(self) -> DataFrame: + """ + Returns the precision-recall curve, which is a Dataframe + containing two fields recall, precision with (0.0, 1.0) prepended + to it. + """ + return self._call_java("pr") + + @property + @since("3.1.0") + def fMeasureByThreshold(self) -> DataFrame: + """ + Returns a dataframe with two fields (threshold, F-Measure) curve + with beta = 1.0. + """ + return self._call_java("fMeasureByThreshold") + + @property + @since("3.1.0") + def precisionByThreshold(self) -> DataFrame: + """ + Returns a dataframe with two fields (threshold, precision) curve. + Every possible probability obtained in transforming the dataset + are used as thresholds used in calculating the precision. + """ + return self._call_java("precisionByThreshold") + + @property + @since("3.1.0") + def recallByThreshold(self) -> DataFrame: + """ + Returns a dataframe with two fields (threshold, recall) curve. + Every possible probability obtained in transforming the dataset + are used as thresholds used in calculating the recall. + """ + return self._call_java("recallByThreshold") + + +class _LinearSVCParams( + _ClassifierParams, + HasRegParam, + HasMaxIter, + HasFitIntercept, + HasTol, + HasStandardization, + HasWeightCol, + HasAggregationDepth, + HasThreshold, + HasMaxBlockSizeInMB, +): + """ + Params for :py:class:`LinearSVC` and :py:class:`LinearSVCModel`. + + .. versionadded:: 3.0.0 + """ + + threshold: Param[float] = Param( + Params._dummy(), + "threshold", + "The threshold in binary classification applied to the linear model" + " prediction. This threshold can be any real number, where Inf will make" + " all predictions 0.0 and -Inf will make all predictions 1.0.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any) -> None: + super(_LinearSVCParams, self).__init__(*args) + self._setDefault( + maxIter=100, + regParam=0.0, + tol=1e-6, + fitIntercept=True, + standardization=True, + threshold=0.0, + aggregationDepth=2, + maxBlockSizeInMB=0.0, + ) + + +@inherit_doc +class LinearSVC( + _JavaClassifier["LinearSVCModel"], + _LinearSVCParams, + JavaMLWritable, + JavaMLReadable["LinearSVC"], +): + """ + This binary classifier optimizes the Hinge Loss using the OWLQN optimizer. + Only supports L2 regularization currently. + + .. versionadded:: 2.2.0 + + Notes + ----- + `Linear SVM Classifier `_ + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.ml.linalg import Vectors + >>> df = sc.parallelize([ + ... Row(label=1.0, features=Vectors.dense(1.0, 1.0, 1.0)), + ... Row(label=0.0, features=Vectors.dense(1.0, 2.0, 3.0))]).toDF() + >>> svm = LinearSVC() + >>> svm.getMaxIter() + 100 + >>> svm.setMaxIter(5) + LinearSVC... + >>> svm.getMaxIter() + 5 + >>> svm.getRegParam() + 0.0 + >>> svm.setRegParam(0.01) + LinearSVC... + >>> svm.getRegParam() + 0.01 + >>> model = svm.fit(df) + >>> model.setPredictionCol("newPrediction") + LinearSVCModel... + >>> model.getPredictionCol() + 'newPrediction' + >>> model.setThreshold(0.5) + LinearSVCModel... + >>> model.getThreshold() + 0.5 + >>> model.getMaxBlockSizeInMB() + 0.0 + >>> model.coefficients + DenseVector([0.0, -1.0319, -0.5159]) + >>> model.intercept + 2.579645978780695 + >>> model.numClasses + 2 + >>> model.numFeatures + 3 + >>> test0 = sc.parallelize([Row(features=Vectors.dense(-1.0, -1.0, -1.0))]).toDF() + >>> model.predict(test0.head().features) + 1.0 + >>> model.predictRaw(test0.head().features) + DenseVector([-4.1274, 4.1274]) + >>> result = model.transform(test0).head() + >>> result.newPrediction + 1.0 + >>> result.rawPrediction + DenseVector([-4.1274, 4.1274]) + >>> svm_path = temp_path + "/svm" + >>> svm.save(svm_path) + >>> svm2 = LinearSVC.load(svm_path) + >>> svm2.getMaxIter() + 5 + >>> model_path = temp_path + "/svm_model" + >>> model.save(model_path) + >>> model2 = LinearSVCModel.load(model_path) + >>> model.coefficients[0] == model2.coefficients[0] + True + >>> model.intercept == model2.intercept + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + tol: float = 1e-6, + rawPredictionCol: str = "rawPrediction", + fitIntercept: bool = True, + standardization: bool = True, + threshold: float = 0.0, + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + maxBlockSizeInMB: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, tol=1e-6, rawPredictionCol="rawPrediction", \ + fitIntercept=True, standardization=True, threshold=0.0, weightCol=None, \ + aggregationDepth=2, maxBlockSizeInMB=0.0): + """ + super(LinearSVC, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.LinearSVC", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.2.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + tol: float = 1e-6, + rawPredictionCol: str = "rawPrediction", + fitIntercept: bool = True, + standardization: bool = True, + threshold: float = 0.0, + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + maxBlockSizeInMB: float = 0.0, + ) -> "LinearSVC": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, tol=1e-6, rawPredictionCol="rawPrediction", \ + fitIntercept=True, standardization=True, threshold=0.0, weightCol=None, \ + aggregationDepth=2, maxBlockSizeInMB=0.0): + Sets params for Linear SVM Classifier. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "LinearSVCModel": + return LinearSVCModel(java_model) + + @since("2.2.0") + def setMaxIter(self, value: int) -> "LinearSVC": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.2.0") + def setRegParam(self, value: float) -> "LinearSVC": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + @since("2.2.0") + def setTol(self, value: float) -> "LinearSVC": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("2.2.0") + def setFitIntercept(self, value: bool) -> "LinearSVC": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + @since("2.2.0") + def setStandardization(self, value: bool) -> "LinearSVC": + """ + Sets the value of :py:attr:`standardization`. + """ + return self._set(standardization=value) + + @since("2.2.0") + def setThreshold(self, value: float) -> "LinearSVC": + """ + Sets the value of :py:attr:`threshold`. + """ + return self._set(threshold=value) + + @since("2.2.0") + def setWeightCol(self, value: str) -> "LinearSVC": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("2.2.0") + def setAggregationDepth(self, value: int) -> "LinearSVC": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + @since("3.1.0") + def setMaxBlockSizeInMB(self, value: float) -> "LinearSVC": + """ + Sets the value of :py:attr:`maxBlockSizeInMB`. + """ + return self._set(maxBlockSizeInMB=value) + + +class LinearSVCModel( + _JavaClassificationModel[Vector], + _LinearSVCParams, + JavaMLWritable, + JavaMLReadable["LinearSVCModel"], + HasTrainingSummary["LinearSVCTrainingSummary"], +): + """ + Model fitted by LinearSVC. + + .. versionadded:: 2.2.0 + """ + + @since("3.0.0") + def setThreshold(self, value: float) -> "LinearSVCModel": + """ + Sets the value of :py:attr:`threshold`. + """ + return self._set(threshold=value) + + @property + @since("2.2.0") + def coefficients(self) -> Vector: + """ + Model coefficients of Linear SVM Classifier. + """ + return self._call_java("coefficients") + + @property + @since("2.2.0") + def intercept(self) -> float: + """ + Model intercept of Linear SVM Classifier. + """ + return self._call_java("intercept") + + @since("3.1.0") + def summary(self) -> "LinearSVCTrainingSummary": + """ + Gets summary (accuracy/precision/recall, objective history, total iterations) of model + trained on the training set. An exception is thrown if `trainingSummary is None`. + """ + if self.hasSummary: + return LinearSVCTrainingSummary(super(LinearSVCModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "LinearSVCSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_lsvc_summary = self._call_java("evaluate", dataset) + return LinearSVCSummary(java_lsvc_summary) + + +class LinearSVCSummary(_BinaryClassificationSummary): + """ + Abstraction for LinearSVC Results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class LinearSVCTrainingSummary(LinearSVCSummary, _TrainingSummary): + """ + Abstraction for LinearSVC Training results. + + .. versionadded:: 3.1.0 + """ + + pass + + +class _LogisticRegressionParams( + _ProbabilisticClassifierParams, + HasRegParam, + HasElasticNetParam, + HasMaxIter, + HasFitIntercept, + HasTol, + HasStandardization, + HasWeightCol, + HasAggregationDepth, + HasThreshold, + HasMaxBlockSizeInMB, +): + """ + Params for :py:class:`LogisticRegression` and :py:class:`LogisticRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + threshold: Param[float] = Param( + Params._dummy(), + "threshold", + "Threshold in binary classification prediction, in range [0, 1]." + + " If threshold and thresholds are both set, they must match." + + "e.g. if threshold is p, then thresholds must be equal to [1-p, p].", + typeConverter=TypeConverters.toFloat, + ) + + family: Param[str] = Param( + Params._dummy(), + "family", + "The name of family which is a description of the label distribution to " + + "be used in the model. Supported options: auto, binomial, multinomial", + typeConverter=TypeConverters.toString, + ) + + lowerBoundsOnCoefficients: Param[Matrix] = Param( + Params._dummy(), + "lowerBoundsOnCoefficients", + "The lower bounds on coefficients if fitting under bound " + "constrained optimization. The bound matrix must be " + "compatible with the shape " + "(1, number of features) for binomial regression, or " + "(number of classes, number of features) " + "for multinomial regression.", + typeConverter=TypeConverters.toMatrix, + ) + + upperBoundsOnCoefficients: Param[Matrix] = Param( + Params._dummy(), + "upperBoundsOnCoefficients", + "The upper bounds on coefficients if fitting under bound " + "constrained optimization. The bound matrix must be " + "compatible with the shape " + "(1, number of features) for binomial regression, or " + "(number of classes, number of features) " + "for multinomial regression.", + typeConverter=TypeConverters.toMatrix, + ) + + lowerBoundsOnIntercepts: Param[Vector] = Param( + Params._dummy(), + "lowerBoundsOnIntercepts", + "The lower bounds on intercepts if fitting under bound " + "constrained optimization. The bounds vector size must be" + "equal with 1 for binomial regression, or the number of" + "lasses for multinomial regression.", + typeConverter=TypeConverters.toVector, + ) + + upperBoundsOnIntercepts: Param[Vector] = Param( + Params._dummy(), + "upperBoundsOnIntercepts", + "The upper bounds on intercepts if fitting under bound " + "constrained optimization. The bound vector size must be " + "equal with 1 for binomial regression, or the number of " + "classes for multinomial regression.", + typeConverter=TypeConverters.toVector, + ) + + def __init__(self, *args: Any): + super(_LogisticRegressionParams, self).__init__(*args) + self._setDefault( + maxIter=100, regParam=0.0, tol=1e-6, threshold=0.5, family="auto", maxBlockSizeInMB=0.0 + ) + + @since("1.4.0") + def setThreshold(self: "P", value: float) -> "P": + """ + Sets the value of :py:attr:`threshold`. + Clears value of :py:attr:`thresholds` if it has been set. + """ + self._set(threshold=value) + self.clear(self.thresholds) # type: ignore[attr-defined] + return self + + @since("1.4.0") + def getThreshold(self) -> float: + """ + Get threshold for binary classification. + + If :py:attr:`thresholds` is set with length 2 (i.e., binary classification), + this returns the equivalent threshold: + :math:`\\frac{1}{1 + \\frac{thresholds(0)}{thresholds(1)}}`. + Otherwise, returns :py:attr:`threshold` if set or its default value if unset. + """ + self._checkThresholdConsistency() + if self.isSet(self.thresholds): + ts = self.getOrDefault(self.thresholds) + if len(ts) != 2: + raise ValueError( + "Logistic Regression getThreshold only applies to" + + " binary classification, but thresholds has length != 2." + + " thresholds: {ts}".format(ts=ts) + ) + return 1.0 / (1.0 + ts[0] / ts[1]) + else: + return self.getOrDefault(self.threshold) + + @since("1.5.0") + def setThresholds(self: "P", value: List[float]) -> "P": + """ + Sets the value of :py:attr:`thresholds`. + Clears value of :py:attr:`threshold` if it has been set. + """ + self._set(thresholds=value) + self.clear(self.threshold) # type: ignore[attr-defined] + return self + + @since("1.5.0") + def getThresholds(self) -> List[float]: + """ + If :py:attr:`thresholds` is set, return its value. + Otherwise, if :py:attr:`threshold` is set, return the equivalent thresholds for binary + classification: (1-threshold, threshold). + If neither are set, throw an error. + """ + self._checkThresholdConsistency() + if not self.isSet(self.thresholds) and self.isSet(self.threshold): + t = self.getOrDefault(self.threshold) + return [1.0 - t, t] + else: + return self.getOrDefault(self.thresholds) + + def _checkThresholdConsistency(self) -> None: + if self.isSet(self.threshold) and self.isSet(self.thresholds): + ts = self.getOrDefault(self.thresholds) + if len(ts) != 2: + raise ValueError( + "Logistic Regression getThreshold only applies to" + + " binary classification, but thresholds has length != 2." + + " thresholds: {0}".format(str(ts)) + ) + t = 1.0 / (1.0 + ts[0] / ts[1]) + t2 = self.getOrDefault(self.threshold) + if abs(t2 - t) >= 1e-5: + raise ValueError( + "Logistic Regression getThreshold found inconsistent values for" + + " threshold (%g) and thresholds (equivalent to %g)" % (t2, t) + ) + + @since("2.1.0") + def getFamily(self) -> str: + """ + Gets the value of :py:attr:`family` or its default value. + """ + return self.getOrDefault(self.family) + + @since("2.3.0") + def getLowerBoundsOnCoefficients(self) -> Matrix: + """ + Gets the value of :py:attr:`lowerBoundsOnCoefficients` + """ + return self.getOrDefault(self.lowerBoundsOnCoefficients) + + @since("2.3.0") + def getUpperBoundsOnCoefficients(self) -> Matrix: + """ + Gets the value of :py:attr:`upperBoundsOnCoefficients` + """ + return self.getOrDefault(self.upperBoundsOnCoefficients) + + @since("2.3.0") + def getLowerBoundsOnIntercepts(self) -> Vector: + """ + Gets the value of :py:attr:`lowerBoundsOnIntercepts` + """ + return self.getOrDefault(self.lowerBoundsOnIntercepts) + + @since("2.3.0") + def getUpperBoundsOnIntercepts(self) -> Vector: + """ + Gets the value of :py:attr:`upperBoundsOnIntercepts` + """ + return self.getOrDefault(self.upperBoundsOnIntercepts) + + +@inherit_doc +class LogisticRegression( + _JavaProbabilisticClassifier["LogisticRegressionModel"], + _LogisticRegressionParams, + JavaMLWritable, + JavaMLReadable["LogisticRegression"], +): + """ + Logistic regression. + This class supports multinomial logistic (softmax) and binomial logistic regression. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.ml.linalg import Vectors + >>> bdf = sc.parallelize([ + ... Row(label=1.0, weight=1.0, features=Vectors.dense(0.0, 5.0)), + ... Row(label=0.0, weight=2.0, features=Vectors.dense(1.0, 2.0)), + ... Row(label=1.0, weight=3.0, features=Vectors.dense(2.0, 1.0)), + ... Row(label=0.0, weight=4.0, features=Vectors.dense(3.0, 3.0))]).toDF() + >>> blor = LogisticRegression(weightCol="weight") + >>> blor.getRegParam() + 0.0 + >>> blor.setRegParam(0.01) + LogisticRegression... + >>> blor.getRegParam() + 0.01 + >>> blor.setMaxIter(10) + LogisticRegression... + >>> blor.getMaxIter() + 10 + >>> blor.clear(blor.maxIter) + >>> blorModel = blor.fit(bdf) + >>> blorModel.setFeaturesCol("features") + LogisticRegressionModel... + >>> blorModel.setProbabilityCol("newProbability") + LogisticRegressionModel... + >>> blorModel.getProbabilityCol() + 'newProbability' + >>> blorModel.getMaxBlockSizeInMB() + 0.0 + >>> blorModel.setThreshold(0.1) + LogisticRegressionModel... + >>> blorModel.getThreshold() + 0.1 + >>> blorModel.coefficients + DenseVector([-1.080..., -0.646...]) + >>> blorModel.intercept + 3.112... + >>> blorModel.evaluate(bdf).accuracy == blorModel.summary.accuracy + True + >>> data_path = "data/mllib/sample_multiclass_classification_data.txt" + >>> mdf = spark.read.format("libsvm").load(data_path) + >>> mlor = LogisticRegression(regParam=0.1, elasticNetParam=1.0, family="multinomial") + >>> mlorModel = mlor.fit(mdf) + >>> mlorModel.coefficientMatrix + SparseMatrix(3, 4, [0, 1, 2, 3], [3, 2, 1], [1.87..., -2.75..., -0.50...], 1) + >>> mlorModel.interceptVector + DenseVector([0.04..., -0.42..., 0.37...]) + >>> test0 = sc.parallelize([Row(features=Vectors.dense(-1.0, 1.0))]).toDF() + >>> blorModel.predict(test0.head().features) + 1.0 + >>> blorModel.predictRaw(test0.head().features) + DenseVector([-3.54..., 3.54...]) + >>> blorModel.predictProbability(test0.head().features) + DenseVector([0.028, 0.972]) + >>> result = blorModel.transform(test0).head() + >>> result.prediction + 1.0 + >>> result.newProbability + DenseVector([0.02..., 0.97...]) + >>> result.rawPrediction + DenseVector([-3.54..., 3.54...]) + >>> test1 = sc.parallelize([Row(features=Vectors.sparse(2, [0], [1.0]))]).toDF() + >>> blorModel.transform(test1).head().prediction + 1.0 + >>> blor.setParams("vector") + Traceback (most recent call last): + ... + TypeError: Method setParams forces keyword arguments. + >>> lr_path = temp_path + "/lr" + >>> blor.save(lr_path) + >>> lr2 = LogisticRegression.load(lr_path) + >>> lr2.getRegParam() + 0.01 + >>> model_path = temp_path + "/lr_model" + >>> blorModel.save(model_path) + >>> model2 = LogisticRegressionModel.load(model_path) + >>> blorModel.coefficients[0] == model2.coefficients[0] + True + >>> blorModel.intercept == model2.intercept + True + >>> model2 + LogisticRegressionModel: uid=..., numClasses=2, numFeatures=2 + >>> blorModel.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @overload + def __init__( + self, + *, + featuresCol: str = ..., + labelCol: str = ..., + predictionCol: str = ..., + maxIter: int = ..., + regParam: float = ..., + elasticNetParam: float = ..., + tol: float = ..., + fitIntercept: bool = ..., + threshold: float = ..., + probabilityCol: str = ..., + rawPredictionCol: str = ..., + standardization: bool = ..., + weightCol: Optional[str] = ..., + aggregationDepth: int = ..., + family: str = ..., + lowerBoundsOnCoefficients: Optional[Matrix] = ..., + upperBoundsOnCoefficients: Optional[Matrix] = ..., + lowerBoundsOnIntercepts: Optional[Vector] = ..., + upperBoundsOnIntercepts: Optional[Vector] = ..., + maxBlockSizeInMB: float = ..., + ): + ... + + @overload + def __init__( + self, + *, + featuresCol: str = ..., + labelCol: str = ..., + predictionCol: str = ..., + maxIter: int = ..., + regParam: float = ..., + elasticNetParam: float = ..., + tol: float = ..., + fitIntercept: bool = ..., + thresholds: Optional[List[float]] = ..., + probabilityCol: str = ..., + rawPredictionCol: str = ..., + standardization: bool = ..., + weightCol: Optional[str] = ..., + aggregationDepth: int = ..., + family: str = ..., + lowerBoundsOnCoefficients: Optional[Matrix] = ..., + upperBoundsOnCoefficients: Optional[Matrix] = ..., + lowerBoundsOnIntercepts: Optional[Vector] = ..., + upperBoundsOnIntercepts: Optional[Vector] = ..., + maxBlockSizeInMB: float = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + elasticNetParam: float = 0.0, + tol: float = 1e-6, + fitIntercept: bool = True, + threshold: float = 0.5, + thresholds: Optional[List[float]] = None, + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + standardization: bool = True, + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + family: str = "auto", + lowerBoundsOnCoefficients: Optional[Matrix] = None, + upperBoundsOnCoefficients: Optional[Matrix] = None, + lowerBoundsOnIntercepts: Optional[Vector] = None, + upperBoundsOnIntercepts: Optional[Vector] = None, + maxBlockSizeInMB: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \ + threshold=0.5, thresholds=None, probabilityCol="probability", \ + rawPredictionCol="rawPrediction", standardization=True, weightCol=None, \ + aggregationDepth=2, family="auto", \ + lowerBoundsOnCoefficients=None, upperBoundsOnCoefficients=None, \ + lowerBoundsOnIntercepts=None, upperBoundsOnIntercepts=None, \ + maxBlockSizeInMB=0.0): + If the threshold and thresholds Params are both set, they must be equivalent. + """ + super(LogisticRegression, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.LogisticRegression", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + self._checkThresholdConsistency() + + @overload + def setParams( + self, + *, + featuresCol: str = ..., + labelCol: str = ..., + predictionCol: str = ..., + maxIter: int = ..., + regParam: float = ..., + elasticNetParam: float = ..., + tol: float = ..., + fitIntercept: bool = ..., + threshold: float = ..., + probabilityCol: str = ..., + rawPredictionCol: str = ..., + standardization: bool = ..., + weightCol: Optional[str] = ..., + aggregationDepth: int = ..., + family: str = ..., + lowerBoundsOnCoefficients: Optional[Matrix] = ..., + upperBoundsOnCoefficients: Optional[Matrix] = ..., + lowerBoundsOnIntercepts: Optional[Vector] = ..., + upperBoundsOnIntercepts: Optional[Vector] = ..., + maxBlockSizeInMB: float = ..., + ) -> "LogisticRegression": + ... + + @overload + def setParams( + self, + *, + featuresCol: str = ..., + labelCol: str = ..., + predictionCol: str = ..., + maxIter: int = ..., + regParam: float = ..., + elasticNetParam: float = ..., + tol: float = ..., + fitIntercept: bool = ..., + thresholds: Optional[List[float]] = ..., + probabilityCol: str = ..., + rawPredictionCol: str = ..., + standardization: bool = ..., + weightCol: Optional[str] = ..., + aggregationDepth: int = ..., + family: str = ..., + lowerBoundsOnCoefficients: Optional[Matrix] = ..., + upperBoundsOnCoefficients: Optional[Matrix] = ..., + lowerBoundsOnIntercepts: Optional[Vector] = ..., + upperBoundsOnIntercepts: Optional[Vector] = ..., + maxBlockSizeInMB: float = ..., + ) -> "LogisticRegression": + ... + + @keyword_only + @since("1.3.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + elasticNetParam: float = 0.0, + tol: float = 1e-6, + fitIntercept: bool = True, + threshold: float = 0.5, + thresholds: Optional[List[float]] = None, + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + standardization: bool = True, + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + family: str = "auto", + lowerBoundsOnCoefficients: Optional[Matrix] = None, + upperBoundsOnCoefficients: Optional[Matrix] = None, + lowerBoundsOnIntercepts: Optional[Vector] = None, + upperBoundsOnIntercepts: Optional[Vector] = None, + maxBlockSizeInMB: float = 0.0, + ) -> "LogisticRegression": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \ + threshold=0.5, thresholds=None, probabilityCol="probability", \ + rawPredictionCol="rawPrediction", standardization=True, weightCol=None, \ + aggregationDepth=2, family="auto", \ + lowerBoundsOnCoefficients=None, upperBoundsOnCoefficients=None, \ + lowerBoundsOnIntercepts=None, upperBoundsOnIntercepts=None, \ + maxBlockSizeInMB=0.0): + Sets params for logistic regression. + If the threshold and thresholds Params are both set, they must be equivalent. + """ + kwargs = self._input_kwargs + self._set(**kwargs) + self._checkThresholdConsistency() + return self + + def _create_model(self, java_model: "JavaObject") -> "LogisticRegressionModel": + return LogisticRegressionModel(java_model) + + @since("2.1.0") + def setFamily(self, value: str) -> "LogisticRegression": + """ + Sets the value of :py:attr:`family`. + """ + return self._set(family=value) + + @since("2.3.0") + def setLowerBoundsOnCoefficients(self, value: Matrix) -> "LogisticRegression": + """ + Sets the value of :py:attr:`lowerBoundsOnCoefficients` + """ + return self._set(lowerBoundsOnCoefficients=value) + + @since("2.3.0") + def setUpperBoundsOnCoefficients(self, value: Matrix) -> "LogisticRegression": + """ + Sets the value of :py:attr:`upperBoundsOnCoefficients` + """ + return self._set(upperBoundsOnCoefficients=value) + + @since("2.3.0") + def setLowerBoundsOnIntercepts(self, value: Vector) -> "LogisticRegression": + """ + Sets the value of :py:attr:`lowerBoundsOnIntercepts` + """ + return self._set(lowerBoundsOnIntercepts=value) + + @since("2.3.0") + def setUpperBoundsOnIntercepts(self, value: Vector) -> "LogisticRegression": + """ + Sets the value of :py:attr:`upperBoundsOnIntercepts` + """ + return self._set(upperBoundsOnIntercepts=value) + + def setMaxIter(self, value: int) -> "LogisticRegression": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + def setRegParam(self, value: float) -> "LogisticRegression": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + def setTol(self, value: float) -> "LogisticRegression": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + def setElasticNetParam(self, value: float) -> "LogisticRegression": + """ + Sets the value of :py:attr:`elasticNetParam`. + """ + return self._set(elasticNetParam=value) + + def setFitIntercept(self, value: bool) -> "LogisticRegression": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + def setStandardization(self, value: bool) -> "LogisticRegression": + """ + Sets the value of :py:attr:`standardization`. + """ + return self._set(standardization=value) + + def setWeightCol(self, value: str) -> "LogisticRegression": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + def setAggregationDepth(self, value: int) -> "LogisticRegression": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + @since("3.1.0") + def setMaxBlockSizeInMB(self, value: float) -> "LogisticRegression": + """ + Sets the value of :py:attr:`maxBlockSizeInMB`. + """ + return self._set(maxBlockSizeInMB=value) + + +class LogisticRegressionModel( + _JavaProbabilisticClassificationModel[Vector], + _LogisticRegressionParams, + JavaMLWritable, + JavaMLReadable["LogisticRegressionModel"], + HasTrainingSummary["LogisticRegressionTrainingSummary"], +): + """ + Model fitted by LogisticRegression. + + .. versionadded:: 1.3.0 + """ + + @property + @since("2.0.0") + def coefficients(self) -> Vector: + """ + Model coefficients of binomial logistic regression. + An exception is thrown in the case of multinomial logistic regression. + """ + return self._call_java("coefficients") + + @property + @since("1.4.0") + def intercept(self) -> float: + """ + Model intercept of binomial logistic regression. + An exception is thrown in the case of multinomial logistic regression. + """ + return self._call_java("intercept") + + @property + @since("2.1.0") + def coefficientMatrix(self) -> Matrix: + """ + Model coefficients. + """ + return self._call_java("coefficientMatrix") + + @property + @since("2.1.0") + def interceptVector(self) -> Vector: + """ + Model intercept. + """ + return self._call_java("interceptVector") + + @property + @since("2.0.0") + def summary(self) -> "LogisticRegressionTrainingSummary": + """ + Gets summary (accuracy/precision/recall, objective history, total iterations) of model + trained on the training set. An exception is thrown if `trainingSummary is None`. + """ + if self.hasSummary: + if self.numClasses <= 2: + return BinaryLogisticRegressionTrainingSummary( + super(LogisticRegressionModel, self).summary + ) + else: + return LogisticRegressionTrainingSummary( + super(LogisticRegressionModel, self).summary + ) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "LogisticRegressionSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_blr_summary = self._call_java("evaluate", dataset) + if self.numClasses <= 2: + return BinaryLogisticRegressionSummary(java_blr_summary) + else: + return LogisticRegressionSummary(java_blr_summary) + + +class LogisticRegressionSummary(_ClassificationSummary): + """ + Abstraction for Logistic Regression Results for a given model. + + .. versionadded:: 2.0.0 + """ + + @property + @since("2.0.0") + def probabilityCol(self) -> str: + """ + Field in "predictions" which gives the probability + of each class as a vector. + """ + return self._call_java("probabilityCol") + + @property + @since("2.0.0") + def featuresCol(self) -> str: + """ + Field in "predictions" which gives the features of each instance + as a vector. + """ + return self._call_java("featuresCol") + + +@inherit_doc +class LogisticRegressionTrainingSummary(LogisticRegressionSummary, _TrainingSummary): + """ + Abstraction for multinomial Logistic Regression Training results. + + .. versionadded:: 2.0.0 + """ + + pass + + +@inherit_doc +class BinaryLogisticRegressionSummary(_BinaryClassificationSummary, LogisticRegressionSummary): + """ + Binary Logistic regression results for a given model. + + .. versionadded:: 2.0.0 + """ + + pass + + +@inherit_doc +class BinaryLogisticRegressionTrainingSummary( + BinaryLogisticRegressionSummary, LogisticRegressionTrainingSummary +): + """ + Binary Logistic regression training results for a given model. + + .. versionadded:: 2.0.0 + """ + + pass + + +@inherit_doc +class _DecisionTreeClassifierParams(_DecisionTreeParams, _TreeClassifierParams): + """ + Params for :py:class:`DecisionTreeClassifier` and :py:class:`DecisionTreeClassificationModel`. + """ + + def __init__(self, *args: Any): + super(_DecisionTreeClassifierParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + checkpointInterval=10, + impurity="gini", + leafCol="", + minWeightFractionPerNode=0.0, + ) + + +@inherit_doc +class DecisionTreeClassifier( + _JavaProbabilisticClassifier["DecisionTreeClassificationModel"], + _DecisionTreeClassifierParams, + JavaMLWritable, + JavaMLReadable["DecisionTreeClassifier"], +): + """ + `Decision tree `_ + learning algorithm for classification. + It supports both binary and multiclass labels, as well as both continuous and categorical + features. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.feature import StringIndexer + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> stringIndexer = StringIndexer(inputCol="label", outputCol="indexed") + >>> si_model = stringIndexer.fit(df) + >>> td = si_model.transform(df) + >>> dt = DecisionTreeClassifier(maxDepth=2, labelCol="indexed", leafCol="leafId") + >>> model = dt.fit(td) + >>> model.getLabelCol() + 'indexed' + >>> model.setFeaturesCol("features") + DecisionTreeClassificationModel... + >>> model.numNodes + 3 + >>> model.depth + 1 + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> model.numFeatures + 1 + >>> model.numClasses + 2 + >>> print(model.toDebugString) + DecisionTreeClassificationModel...depth=1, numNodes=3... + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> model.predictRaw(test0.head().features) + DenseVector([1.0, 0.0]) + >>> model.predictProbability(test0.head().features) + DenseVector([1.0, 0.0]) + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> result.probability + DenseVector([1.0, 0.0]) + >>> result.rawPrediction + DenseVector([1.0, 0.0]) + >>> result.leafId + 0.0 + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 1.0 + >>> dtc_path = temp_path + "/dtc" + >>> dt.save(dtc_path) + >>> dt2 = DecisionTreeClassifier.load(dtc_path) + >>> dt2.getMaxDepth() + 2 + >>> model_path = temp_path + "/dtc_model" + >>> model.save(model_path) + >>> model2 = DecisionTreeClassificationModel.load(model_path) + >>> model.featureImportances == model2.featureImportances + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> df3 = spark.createDataFrame([ + ... (1.0, 0.2, Vectors.dense(1.0)), + ... (1.0, 0.8, Vectors.dense(1.0)), + ... (0.0, 1.0, Vectors.sparse(1, [], []))], ["label", "weight", "features"]) + >>> si3 = StringIndexer(inputCol="label", outputCol="indexed") + >>> si_model3 = si3.fit(df3) + >>> td3 = si_model3.transform(df3) + >>> dt3 = DecisionTreeClassifier(maxDepth=2, weightCol="weight", labelCol="indexed") + >>> model3 = dt3.fit(td3) + >>> print(model3.toDebugString) + DecisionTreeClassificationModel...depth=1, numNodes=3... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "gini", + seed: Optional[int] = None, + weightCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, impurity="gini", \ + seed=None, weightCol=None, leafCol="", minWeightFractionPerNode=0.0) + """ + super(DecisionTreeClassifier, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.DecisionTreeClassifier", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "gini", + seed: Optional[int] = None, + weightCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + ) -> "DecisionTreeClassifier": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, impurity="gini", \ + seed=None, weightCol=None, leafCol="", minWeightFractionPerNode=0.0) + Sets params for the DecisionTreeClassifier. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "DecisionTreeClassificationModel": + return DecisionTreeClassificationModel(java_model) + + def setMaxDepth(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + def setMaxBins(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + def setMinInstancesPerNode(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + def setMinInfoGain(self, value: float) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + def setMaxMemoryInMB(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + def setCacheNodeIds(self, value: bool) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setCheckpointInterval(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + def setSeed(self, value: int) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "DecisionTreeClassifier": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + +@inherit_doc +class DecisionTreeClassificationModel( + _DecisionTreeModel, + _JavaProbabilisticClassificationModel[Vector], + _DecisionTreeClassifierParams, + JavaMLWritable, + JavaMLReadable["DecisionTreeClassificationModel"], +): + """ + Model fitted by DecisionTreeClassifier. + + .. versionadded:: 1.4.0 + """ + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + This generalizes the idea of "Gini" importance to other losses, + following the explanation of Gini importance from "Random Forests" documentation + by Leo Breiman and Adele Cutler, and following the implementation from scikit-learn. + + This feature importance is calculated as follows: + - importance(feature j) = sum (over nodes which split on feature j) of the gain, + where gain is scaled by the number of instances passing through node + - Normalize importances for tree to sum to 1. + + .. versionadded:: 2.0.0 + + Notes + ----- + Feature importance for single decision trees can have high variance due to + correlated predictor variables. Consider using a :py:class:`RandomForestClassifier` + to determine feature importance instead. + """ + return self._call_java("featureImportances") + + +@inherit_doc +class _RandomForestClassifierParams(_RandomForestParams, _TreeClassifierParams): + """ + Params for :py:class:`RandomForestClassifier` and :py:class:`RandomForestClassificationModel`. + """ + + def __init__(self, *args: Any): + super(_RandomForestClassifierParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + checkpointInterval=10, + impurity="gini", + numTrees=20, + featureSubsetStrategy="auto", + subsamplingRate=1.0, + leafCol="", + minWeightFractionPerNode=0.0, + bootstrap=True, + ) + + +@inherit_doc +class RandomForestClassifier( + _JavaProbabilisticClassifier["RandomForestClassificationModel"], + _RandomForestClassifierParams, + JavaMLWritable, + JavaMLReadable["RandomForestClassifier"], +): + """ + `Random Forest `_ + learning algorithm for classification. + It supports both binary and multiclass labels, as well as both continuous and categorical + features. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> import numpy + >>> from numpy import allclose + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.feature import StringIndexer + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> stringIndexer = StringIndexer(inputCol="label", outputCol="indexed") + >>> si_model = stringIndexer.fit(df) + >>> td = si_model.transform(df) + >>> rf = RandomForestClassifier(numTrees=3, maxDepth=2, labelCol="indexed", seed=42, + ... leafCol="leafId") + >>> rf.getMinWeightFractionPerNode() + 0.0 + >>> model = rf.fit(td) + >>> model.getLabelCol() + 'indexed' + >>> model.setFeaturesCol("features") + RandomForestClassificationModel... + >>> model.setRawPredictionCol("newRawPrediction") + RandomForestClassificationModel... + >>> model.getBootstrap() + True + >>> model.getRawPredictionCol() + 'newRawPrediction' + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> allclose(model.treeWeights, [1.0, 1.0, 1.0]) + True + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> model.predictRaw(test0.head().features) + DenseVector([2.0, 0.0]) + >>> model.predictProbability(test0.head().features) + DenseVector([1.0, 0.0]) + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> numpy.argmax(result.probability) + 0 + >>> numpy.argmax(result.newRawPrediction) + 0 + >>> result.leafId + DenseVector([0.0, 0.0, 0.0]) + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 1.0 + >>> model.trees + [DecisionTreeClassificationModel...depth=..., DecisionTreeClassificationModel...] + >>> rfc_path = temp_path + "/rfc" + >>> rf.save(rfc_path) + >>> rf2 = RandomForestClassifier.load(rfc_path) + >>> rf2.getNumTrees() + 3 + >>> model_path = temp_path + "/rfc_model" + >>> model.save(model_path) + >>> model2 = RandomForestClassificationModel.load(model_path) + >>> model.featureImportances == model2.featureImportances + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "gini", + numTrees: int = 20, + featureSubsetStrategy: str = "auto", + seed: Optional[int] = None, + subsamplingRate: float = 1.0, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + bootstrap: Optional[bool] = True, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, impurity="gini", \ + numTrees=20, featureSubsetStrategy="auto", seed=None, subsamplingRate=1.0, \ + leafCol="", minWeightFractionPerNode=0.0, weightCol=None, bootstrap=True) + """ + super(RandomForestClassifier, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.RandomForestClassifier", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "gini", + numTrees: int = 20, + featureSubsetStrategy: str = "auto", + seed: Optional[int] = None, + subsamplingRate: float = 1.0, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + bootstrap: Optional[bool] = True, + ) -> "RandomForestClassifier": + """ + setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, seed=None, \ + impurity="gini", numTrees=20, featureSubsetStrategy="auto", subsamplingRate=1.0, \ + leafCol="", minWeightFractionPerNode=0.0, weightCol=None, bootstrap=True) + Sets params for linear classification. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "RandomForestClassificationModel": + return RandomForestClassificationModel(java_model) + + def setMaxDepth(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + def setMaxBins(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + def setMinInstancesPerNode(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + def setMinInfoGain(self, value: float) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + def setMaxMemoryInMB(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + def setCacheNodeIds(self, value: bool) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setNumTrees(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`numTrees`. + """ + return self._set(numTrees=value) + + @since("3.0.0") + def setBootstrap(self, value: bool) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`bootstrap`. + """ + return self._set(bootstrap=value) + + @since("1.4.0") + def setSubsamplingRate(self, value: float) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`subsamplingRate`. + """ + return self._set(subsamplingRate=value) + + @since("2.4.0") + def setFeatureSubsetStrategy(self, value: str) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`featureSubsetStrategy`. + """ + return self._set(featureSubsetStrategy=value) + + def setSeed(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def setCheckpointInterval(self, value: int) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + +class RandomForestClassificationModel( + _TreeEnsembleModel, + _JavaProbabilisticClassificationModel[Vector], + _RandomForestClassifierParams, + JavaMLWritable, + JavaMLReadable["RandomForestClassificationModel"], + HasTrainingSummary["RandomForestClassificationTrainingSummary"], +): + """ + Model fitted by RandomForestClassifier. + + .. versionadded:: 1.4.0 + """ + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + Each feature's importance is the average of its importance across all trees in the ensemble + The importance vector is normalized to sum to 1. This method is suggested by Hastie et al. + (Hastie, Tibshirani, Friedman. "The Elements of Statistical Learning, 2nd Edition." 2001.) + and follows the implementation from scikit-learn. + + .. versionadded:: 2.0.0 + + See Also + -------- + DecisionTreeClassificationModel.featureImportances + """ + return self._call_java("featureImportances") + + @property + @since("2.0.0") + def trees(self) -> List[DecisionTreeClassificationModel]: + """Trees in this ensemble. Warning: These have null parent Estimators.""" + return [DecisionTreeClassificationModel(m) for m in list(self._call_java("trees"))] + + @property + @since("3.1.0") + def summary(self) -> "RandomForestClassificationTrainingSummary": + """ + Gets summary (accuracy/precision/recall, objective history, total iterations) of model + trained on the training set. An exception is thrown if `trainingSummary is None`. + """ + if self.hasSummary: + if self.numClasses <= 2: + return BinaryRandomForestClassificationTrainingSummary( + super(RandomForestClassificationModel, self).summary + ) + else: + return RandomForestClassificationTrainingSummary( + super(RandomForestClassificationModel, self).summary + ) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate( + self, dataset: DataFrame + ) -> Union["BinaryRandomForestClassificationSummary", "RandomForestClassificationSummary"]: + """ + Evaluates the model on a test dataset. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_rf_summary = self._call_java("evaluate", dataset) + if self.numClasses <= 2: + return BinaryRandomForestClassificationSummary(java_rf_summary) + else: + return RandomForestClassificationSummary(java_rf_summary) + + +class RandomForestClassificationSummary(_ClassificationSummary): + """ + Abstraction for RandomForestClassification Results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class RandomForestClassificationTrainingSummary( + RandomForestClassificationSummary, _TrainingSummary +): + """ + Abstraction for RandomForestClassificationTraining Training results. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class BinaryRandomForestClassificationSummary(_BinaryClassificationSummary): + """ + BinaryRandomForestClassification results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class BinaryRandomForestClassificationTrainingSummary( + BinaryRandomForestClassificationSummary, RandomForestClassificationTrainingSummary +): + """ + BinaryRandomForestClassification training results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +class _GBTClassifierParams(_GBTParams, _HasVarianceImpurity): + """ + Params for :py:class:`GBTClassifier` and :py:class:`GBTClassifierModel`. + + .. versionadded:: 3.0.0 + """ + + supportedLossTypes: List[str] = ["logistic"] + + lossType: Param[str] = Param( + Params._dummy(), + "lossType", + "Loss function which GBT tries to minimize (case-insensitive). " + + "Supported options: " + + ", ".join(supportedLossTypes), + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_GBTClassifierParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + checkpointInterval=10, + lossType="logistic", + maxIter=20, + stepSize=0.1, + subsamplingRate=1.0, + impurity="variance", + featureSubsetStrategy="all", + validationTol=0.01, + leafCol="", + minWeightFractionPerNode=0.0, + ) + + @since("1.4.0") + def getLossType(self) -> str: + """ + Gets the value of lossType or its default value. + """ + return self.getOrDefault(self.lossType) + + +@inherit_doc +class GBTClassifier( + _JavaProbabilisticClassifier["GBTClassificationModel"], + _GBTClassifierParams, + JavaMLWritable, + JavaMLReadable["GBTClassifier"], +): + """ + `Gradient-Boosted Trees (GBTs) `_ + learning algorithm for classification. + It supports binary labels, as well as both continuous and categorical features. + + .. versionadded:: 1.4.0 + + Notes + ----- + Multiclass labels are not currently supported. + + The implementation is based upon: J.H. Friedman. "Stochastic Gradient Boosting." 1999. + + Gradient Boosting vs. TreeBoost: + + - This implementation is for Stochastic Gradient Boosting, not for TreeBoost. + - Both algorithms learn tree ensembles by minimizing loss functions. + - TreeBoost (Friedman, 1999) additionally modifies the outputs at tree leaf nodes + based on the loss function, whereas the original gradient boosting method does not. + - We expect to implement TreeBoost in the future: + `SPARK-4240 `_ + + Examples + -------- + >>> from numpy import allclose + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.feature import StringIndexer + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> stringIndexer = StringIndexer(inputCol="label", outputCol="indexed") + >>> si_model = stringIndexer.fit(df) + >>> td = si_model.transform(df) + >>> gbt = GBTClassifier(maxIter=5, maxDepth=2, labelCol="indexed", seed=42, + ... leafCol="leafId") + >>> gbt.setMaxIter(5) + GBTClassifier... + >>> gbt.setMinWeightFractionPerNode(0.049) + GBTClassifier... + >>> gbt.getMaxIter() + 5 + >>> gbt.getFeatureSubsetStrategy() + 'all' + >>> model = gbt.fit(td) + >>> model.getLabelCol() + 'indexed' + >>> model.setFeaturesCol("features") + GBTClassificationModel... + >>> model.setThresholds([0.3, 0.7]) + GBTClassificationModel... + >>> model.getThresholds() + [0.3, 0.7] + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> allclose(model.treeWeights, [1.0, 0.1, 0.1, 0.1, 0.1]) + True + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> model.predictRaw(test0.head().features) + DenseVector([1.1697, -1.1697]) + >>> model.predictProbability(test0.head().features) + DenseVector([0.9121, 0.0879]) + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> result.leafId + DenseVector([0.0, 0.0, 0.0, 0.0, 0.0]) + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 1.0 + >>> model.totalNumNodes + 15 + >>> print(model.toDebugString) + GBTClassificationModel...numTrees=5... + >>> gbtc_path = temp_path + "gbtc" + >>> gbt.save(gbtc_path) + >>> gbt2 = GBTClassifier.load(gbtc_path) + >>> gbt2.getMaxDepth() + 2 + >>> model_path = temp_path + "gbtc_model" + >>> model.save(model_path) + >>> model2 = GBTClassificationModel.load(model_path) + >>> model.featureImportances == model2.featureImportances + True + >>> model.treeWeights == model2.treeWeights + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> model.trees + [DecisionTreeRegressionModel...depth=..., DecisionTreeRegressionModel...] + >>> validation = spark.createDataFrame([(0.0, Vectors.dense(-1.0),)], + ... ["indexed", "features"]) + >>> model.evaluateEachIteration(validation) + [0.25..., 0.23..., 0.21..., 0.19..., 0.18...] + >>> model.numClasses + 2 + >>> gbt = gbt.setValidationIndicatorCol("validationIndicator") + >>> gbt.getValidationIndicatorCol() + 'validationIndicator' + >>> gbt.getValidationTol() + 0.01 + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + lossType: str = "logistic", + maxIter: int = 20, + stepSize: float = 0.1, + seed: Optional[int] = None, + subsamplingRate: float = 1.0, + impurity: str = "variance", + featureSubsetStrategy: str = "all", + validationTol: float = 0.01, + validationIndicatorCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + lossType="logistic", maxIter=20, stepSize=0.1, seed=None, subsamplingRate=1.0, \ + impurity="variance", featureSubsetStrategy="all", validationTol=0.01, \ + validationIndicatorCol=None, leafCol="", minWeightFractionPerNode=0.0, \ + weightCol=None) + """ + super(GBTClassifier, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.GBTClassifier", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + lossType: str = "logistic", + maxIter: int = 20, + stepSize: float = 0.1, + seed: Optional[int] = None, + subsamplingRate: float = 1.0, + impurity: str = "variance", + featureSubsetStrategy: str = "all", + validationTol: float = 0.01, + validationIndicatorCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + ) -> "GBTClassifier": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + lossType="logistic", maxIter=20, stepSize=0.1, seed=None, subsamplingRate=1.0, \ + impurity="variance", featureSubsetStrategy="all", validationTol=0.01, \ + validationIndicatorCol=None, leafCol="", minWeightFractionPerNode=0.0, \ + weightCol=None) + Sets params for Gradient Boosted Tree Classification. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "GBTClassificationModel": + return GBTClassificationModel(java_model) + + def setMaxDepth(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + def setMaxBins(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + def setMinInstancesPerNode(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + def setMinInfoGain(self, value: float) -> "GBTClassifier": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + def setMaxMemoryInMB(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + def setCacheNodeIds(self, value: bool) -> "GBTClassifier": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setLossType(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`lossType`. + """ + return self._set(lossType=value) + + @since("1.4.0") + def setSubsamplingRate(self, value: float) -> "GBTClassifier": + """ + Sets the value of :py:attr:`subsamplingRate`. + """ + return self._set(subsamplingRate=value) + + @since("2.4.0") + def setFeatureSubsetStrategy(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`featureSubsetStrategy`. + """ + return self._set(featureSubsetStrategy=value) + + @since("3.0.0") + def setValidationIndicatorCol(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`validationIndicatorCol`. + """ + return self._set(validationIndicatorCol=value) + + @since("1.4.0") + def setMaxIter(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("1.4.0") + def setCheckpointInterval(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + @since("1.4.0") + def setSeed(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("1.4.0") + def setStepSize(self, value: int) -> "GBTClassifier": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "GBTClassifier": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + +class GBTClassificationModel( + _TreeEnsembleModel, + _JavaProbabilisticClassificationModel[Vector], + _GBTClassifierParams, + JavaMLWritable, + JavaMLReadable["GBTClassificationModel"], +): + """ + Model fitted by GBTClassifier. + + .. versionadded:: 1.4.0 + """ + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + Each feature's importance is the average of its importance across all trees in the ensemble + The importance vector is normalized to sum to 1. This method is suggested by Hastie et al. + (Hastie, Tibshirani, Friedman. "The Elements of Statistical Learning, 2nd Edition." 2001.) + and follows the implementation from scikit-learn. + + .. versionadded:: 2.0.0 + + See Also + -------- + DecisionTreeClassificationModel.featureImportances + """ + return self._call_java("featureImportances") + + @property + @since("2.0.0") + def trees(self) -> List[DecisionTreeRegressionModel]: + """Trees in this ensemble. Warning: These have null parent Estimators.""" + return [DecisionTreeRegressionModel(m) for m in list(self._call_java("trees"))] + + def evaluateEachIteration(self, dataset: DataFrame) -> List[float]: + """ + Method to compute error or loss for every iteration of gradient boosting. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + return self._call_java("evaluateEachIteration", dataset) + + +class _NaiveBayesParams(_PredictorParams, HasWeightCol): + """ + Params for :py:class:`NaiveBayes` and :py:class:`NaiveBayesModel`. + + .. versionadded:: 3.0.0 + """ + + smoothing: Param[float] = Param( + Params._dummy(), + "smoothing", + "The smoothing parameter, should be >= 0, " + "default is 1.0", + typeConverter=TypeConverters.toFloat, + ) + modelType: Param[str] = Param( + Params._dummy(), + "modelType", + "The model type which is a string " + + "(case-sensitive). Supported options: multinomial (default), bernoulli " + + "and gaussian.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_NaiveBayesParams, self).__init__(*args) + self._setDefault(smoothing=1.0, modelType="multinomial") + + @since("1.5.0") + def getSmoothing(self) -> float: + """ + Gets the value of smoothing or its default value. + """ + return self.getOrDefault(self.smoothing) + + @since("1.5.0") + def getModelType(self) -> str: + """ + Gets the value of modelType or its default value. + """ + return self.getOrDefault(self.modelType) + + +@inherit_doc +class NaiveBayes( + _JavaProbabilisticClassifier["NaiveBayesModel"], + _NaiveBayesParams, + HasThresholds, + HasWeightCol, + JavaMLWritable, + JavaMLReadable["NaiveBayes"], +): + """ + Naive Bayes Classifiers. + It supports both Multinomial and Bernoulli NB. `Multinomial NB \ + `_ + can handle finitely supported discrete data. For example, by converting documents into + TF-IDF vectors, it can be used for document classification. By making every vector a + binary (0/1) data, it can also be used as `Bernoulli NB \ + `_. + + The input feature values for Multinomial NB and Bernoulli NB must be nonnegative. + Since 3.0.0, it supports Complement NB which is an adaptation of the Multinomial NB. + Specifically, Complement NB uses statistics from the complement of each class to compute + the model's coefficients. The inventors of Complement NB show empirically that the parameter + estimates for CNB are more stable than those for Multinomial NB. Like Multinomial NB, the + input feature values for Complement NB must be nonnegative. + Since 3.0.0, it also supports `Gaussian NB \ + `_. + which can handle continuous data. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... Row(label=0.0, weight=0.1, features=Vectors.dense([0.0, 0.0])), + ... Row(label=0.0, weight=0.5, features=Vectors.dense([0.0, 1.0])), + ... Row(label=1.0, weight=1.0, features=Vectors.dense([1.0, 0.0]))]) + >>> nb = NaiveBayes(smoothing=1.0, modelType="multinomial", weightCol="weight") + >>> model = nb.fit(df) + >>> model.setFeaturesCol("features") + NaiveBayesModel... + >>> model.getSmoothing() + 1.0 + >>> model.pi + DenseVector([-0.81..., -0.58...]) + >>> model.theta + DenseMatrix(2, 2, [-0.91..., -0.51..., -0.40..., -1.09...], 1) + >>> model.sigma + DenseMatrix(0, 0, [...], ...) + >>> test0 = sc.parallelize([Row(features=Vectors.dense([1.0, 0.0]))]).toDF() + >>> model.predict(test0.head().features) + 1.0 + >>> model.predictRaw(test0.head().features) + DenseVector([-1.72..., -0.99...]) + >>> model.predictProbability(test0.head().features) + DenseVector([0.32..., 0.67...]) + >>> result = model.transform(test0).head() + >>> result.prediction + 1.0 + >>> result.probability + DenseVector([0.32..., 0.67...]) + >>> result.rawPrediction + DenseVector([-1.72..., -0.99...]) + >>> test1 = sc.parallelize([Row(features=Vectors.sparse(2, [0], [1.0]))]).toDF() + >>> model.transform(test1).head().prediction + 1.0 + >>> nb_path = temp_path + "/nb" + >>> nb.save(nb_path) + >>> nb2 = NaiveBayes.load(nb_path) + >>> nb2.getSmoothing() + 1.0 + >>> model_path = temp_path + "/nb_model" + >>> model.save(model_path) + >>> model2 = NaiveBayesModel.load(model_path) + >>> model.pi == model2.pi + True + >>> model.theta == model2.theta + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> nb = nb.setThresholds([0.01, 10.00]) + >>> model3 = nb.fit(df) + >>> result = model3.transform(test0).head() + >>> result.prediction + 0.0 + >>> nb3 = NaiveBayes().setModelType("gaussian") + >>> model4 = nb3.fit(df) + >>> model4.getModelType() + 'gaussian' + >>> model4.sigma + DenseMatrix(2, 2, [0.0, 0.25, 0.0, 0.0], 1) + >>> nb5 = NaiveBayes(smoothing=1.0, modelType="complement", weightCol="weight") + >>> model5 = nb5.fit(df) + >>> model5.getModelType() + 'complement' + >>> model5.theta + DenseMatrix(2, 2, [...], 1) + >>> model5.sigma + DenseMatrix(0, 0, [...], ...) + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + smoothing: float = 1.0, + modelType: str = "multinomial", + thresholds: Optional[List[float]] = None, + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", smoothing=1.0, \ + modelType="multinomial", thresholds=None, weightCol=None) + """ + super(NaiveBayes, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.NaiveBayes", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + smoothing: float = 1.0, + modelType: str = "multinomial", + thresholds: Optional[List[float]] = None, + weightCol: Optional[str] = None, + ) -> "NaiveBayes": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", smoothing=1.0, \ + modelType="multinomial", thresholds=None, weightCol=None) + Sets params for Naive Bayes. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "NaiveBayesModel": + return NaiveBayesModel(java_model) + + @since("1.5.0") + def setSmoothing(self, value: float) -> "NaiveBayes": + """ + Sets the value of :py:attr:`smoothing`. + """ + return self._set(smoothing=value) + + @since("1.5.0") + def setModelType(self, value: str) -> "NaiveBayes": + """ + Sets the value of :py:attr:`modelType`. + """ + return self._set(modelType=value) + + def setWeightCol(self, value: str) -> "NaiveBayes": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + +class NaiveBayesModel( + _JavaProbabilisticClassificationModel[Vector], + _NaiveBayesParams, + JavaMLWritable, + JavaMLReadable["NaiveBayesModel"], +): + """ + Model fitted by NaiveBayes. + + .. versionadded:: 1.5.0 + """ + + @property + @since("2.0.0") + def pi(self) -> Vector: + """ + log of class priors. + """ + return self._call_java("pi") + + @property + @since("2.0.0") + def theta(self) -> Matrix: + """ + log of class conditional probabilities. + """ + return self._call_java("theta") + + @property + @since("3.0.0") + def sigma(self) -> Matrix: + """ + variance of each feature. + """ + return self._call_java("sigma") + + +class _MultilayerPerceptronParams( + _ProbabilisticClassifierParams, + HasSeed, + HasMaxIter, + HasTol, + HasStepSize, + HasSolver, + HasBlockSize, +): + """ + Params for :py:class:`MultilayerPerceptronClassifier`. + + .. versionadded:: 3.0.0 + """ + + layers: Param[List[int]] = Param( + Params._dummy(), + "layers", + "Sizes of layers from input layer to output layer " + + "E.g., Array(780, 100, 10) means 780 inputs, one hidden layer with 100 " + + "neurons and output layer of 10 neurons.", + typeConverter=TypeConverters.toListInt, + ) + solver: Param[str] = Param( + Params._dummy(), + "solver", + "The solver algorithm for optimization. Supported " + "options: l-bfgs, gd.", + typeConverter=TypeConverters.toString, + ) + initialWeights: Param[Vector] = Param( + Params._dummy(), + "initialWeights", + "The initial weights of the model.", + typeConverter=TypeConverters.toVector, + ) + + def __init__(self, *args: Any): + super(_MultilayerPerceptronParams, self).__init__(*args) + self._setDefault(maxIter=100, tol=1e-6, blockSize=128, stepSize=0.03, solver="l-bfgs") + + @since("1.6.0") + def getLayers(self) -> List[int]: + """ + Gets the value of layers or its default value. + """ + return self.getOrDefault(self.layers) + + @since("2.0.0") + def getInitialWeights(self) -> Vector: + """ + Gets the value of initialWeights or its default value. + """ + return self.getOrDefault(self.initialWeights) + + +@inherit_doc +class MultilayerPerceptronClassifier( + _JavaProbabilisticClassifier["MultilayerPerceptronClassificationModel"], + _MultilayerPerceptronParams, + JavaMLWritable, + JavaMLReadable["MultilayerPerceptronClassifier"], +): + """ + Classifier trainer based on the Multilayer Perceptron. + Each layer has sigmoid activation function, output layer has softmax. + Number of inputs has to be equal to the size of feature vectors. + Number of outputs has to be equal to the total number of labels. + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (0.0, Vectors.dense([0.0, 0.0])), + ... (1.0, Vectors.dense([0.0, 1.0])), + ... (1.0, Vectors.dense([1.0, 0.0])), + ... (0.0, Vectors.dense([1.0, 1.0]))], ["label", "features"]) + >>> mlp = MultilayerPerceptronClassifier(layers=[2, 2, 2], seed=123) + >>> mlp.setMaxIter(100) + MultilayerPerceptronClassifier... + >>> mlp.getMaxIter() + 100 + >>> mlp.getBlockSize() + 128 + >>> mlp.setBlockSize(1) + MultilayerPerceptronClassifier... + >>> mlp.getBlockSize() + 1 + >>> model = mlp.fit(df) + >>> model.setFeaturesCol("features") + MultilayerPerceptronClassificationModel... + >>> model.getMaxIter() + 100 + >>> model.getLayers() + [2, 2, 2] + >>> model.weights.size + 12 + >>> testDF = spark.createDataFrame([ + ... (Vectors.dense([1.0, 0.0]),), + ... (Vectors.dense([0.0, 0.0]),)], ["features"]) + >>> model.predict(testDF.head().features) + 1.0 + >>> model.predictRaw(testDF.head().features) + DenseVector([-16.208, 16.344]) + >>> model.predictProbability(testDF.head().features) + DenseVector([0.0, 1.0]) + >>> model.transform(testDF).select("features", "prediction").show() + +---------+----------+ + | features|prediction| + +---------+----------+ + |[1.0,0.0]| 1.0| + |[0.0,0.0]| 0.0| + +---------+----------+ + ... + >>> mlp_path = temp_path + "/mlp" + >>> mlp.save(mlp_path) + >>> mlp2 = MultilayerPerceptronClassifier.load(mlp_path) + >>> mlp2.getBlockSize() + 1 + >>> model_path = temp_path + "/mlp_model" + >>> model.save(model_path) + >>> model2 = MultilayerPerceptronClassificationModel.load(model_path) + >>> model.getLayers() == model2.getLayers() + True + >>> model.weights == model2.weights + True + >>> model.transform(testDF).take(1) == model2.transform(testDF).take(1) + True + >>> mlp2 = mlp2.setInitialWeights(list(range(0, 12))) + >>> model3 = mlp2.fit(df) + >>> model3.weights != model2.weights + True + >>> model3.getLayers() == model.getLayers() + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + tol: float = 1e-6, + seed: Optional[int] = None, + layers: Optional[List[int]] = None, + blockSize: int = 128, + stepSize: float = 0.03, + solver: str = "l-bfgs", + initialWeights: Optional[Vector] = None, + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, tol=1e-6, seed=None, layers=None, blockSize=128, stepSize=0.03, \ + solver="l-bfgs", initialWeights=None, probabilityCol="probability", \ + rawPredictionCol="rawPrediction") + """ + super(MultilayerPerceptronClassifier, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.MultilayerPerceptronClassifier", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + tol: float = 1e-6, + seed: Optional[int] = None, + layers: Optional[List[int]] = None, + blockSize: int = 128, + stepSize: float = 0.03, + solver: str = "l-bfgs", + initialWeights: Optional[Vector] = None, + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + ) -> "MultilayerPerceptronClassifier": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, tol=1e-6, seed=None, layers=None, blockSize=128, stepSize=0.03, \ + solver="l-bfgs", initialWeights=None, probabilityCol="probability", \ + rawPredictionCol="rawPrediction"): + Sets params for MultilayerPerceptronClassifier. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "MultilayerPerceptronClassificationModel": + return MultilayerPerceptronClassificationModel(java_model) + + @since("1.6.0") + def setLayers(self, value: List[int]) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`layers`. + """ + return self._set(layers=value) + + @since("1.6.0") + def setBlockSize(self, value: int) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`blockSize`. + """ + return self._set(blockSize=value) + + @since("2.0.0") + def setInitialWeights(self, value: Vector) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`initialWeights`. + """ + return self._set(initialWeights=value) + + def setMaxIter(self, value: int) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + def setSeed(self, value: int) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def setTol(self, value: float) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("2.0.0") + def setStepSize(self, value: float) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + def setSolver(self, value: str) -> "MultilayerPerceptronClassifier": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + +class MultilayerPerceptronClassificationModel( + _JavaProbabilisticClassificationModel[Vector], + _MultilayerPerceptronParams, + JavaMLWritable, + JavaMLReadable["MultilayerPerceptronClassificationModel"], + HasTrainingSummary["MultilayerPerceptronClassificationTrainingSummary"], +): + """ + Model fitted by MultilayerPerceptronClassifier. + + .. versionadded:: 1.6.0 + """ + + @property + @since("2.0.0") + def weights(self) -> Vector: + """ + the weights of layers. + """ + return self._call_java("weights") + + @since("3.1.0") + def summary(self) -> "MultilayerPerceptronClassificationTrainingSummary": + """ + Gets summary (accuracy/precision/recall, objective history, total iterations) of model + trained on the training set. An exception is thrown if `trainingSummary is None`. + """ + if self.hasSummary: + return MultilayerPerceptronClassificationTrainingSummary( + super(MultilayerPerceptronClassificationModel, self).summary + ) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "MultilayerPerceptronClassificationSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_mlp_summary = self._call_java("evaluate", dataset) + return MultilayerPerceptronClassificationSummary(java_mlp_summary) + + +class MultilayerPerceptronClassificationSummary(_ClassificationSummary): + """ + Abstraction for MultilayerPerceptronClassifier Results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class MultilayerPerceptronClassificationTrainingSummary( + MultilayerPerceptronClassificationSummary, _TrainingSummary +): + """ + Abstraction for MultilayerPerceptronClassifier Training results. + + .. versionadded:: 3.1.0 + """ + + pass + + +class _OneVsRestParams(_ClassifierParams, HasWeightCol): + """ + Params for :py:class:`OneVsRest` and :py:class:`OneVsRestModelModel`. + """ + + classifier: Param[Classifier] = Param(Params._dummy(), "classifier", "base binary classifier") + + @since("2.0.0") + def getClassifier(self) -> Classifier: + """ + Gets the value of classifier or its default value. + """ + return self.getOrDefault(self.classifier) + + +@inherit_doc +class OneVsRest( + Estimator["OneVsRestModel"], + _OneVsRestParams, + HasParallelism, + MLReadable["OneVsRest"], + MLWritable, + Generic[CM], +): + """ + Reduction of Multiclass Classification to Binary Classification. + Performs reduction using one against all strategy. + For a multiclass classification with k classes, train k models (one per class). + Each example is scored against all k models and the model with highest score + is picked to label the example. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.ml.linalg import Vectors + >>> data_path = "data/mllib/sample_multiclass_classification_data.txt" + >>> df = spark.read.format("libsvm").load(data_path) + >>> lr = LogisticRegression(regParam=0.01) + >>> ovr = OneVsRest(classifier=lr) + >>> ovr.getRawPredictionCol() + 'rawPrediction' + >>> ovr.setPredictionCol("newPrediction") + OneVsRest... + >>> model = ovr.fit(df) + >>> model.models[0].coefficients + DenseVector([0.5..., -1.0..., 3.4..., 4.2...]) + >>> model.models[1].coefficients + DenseVector([-2.1..., 3.1..., -2.6..., -2.3...]) + >>> model.models[2].coefficients + DenseVector([0.3..., -3.4..., 1.0..., -1.1...]) + >>> [x.intercept for x in model.models] + [-2.7..., -2.5..., -1.3...] + >>> test0 = sc.parallelize([Row(features=Vectors.dense(-1.0, 0.0, 1.0, 1.0))]).toDF() + >>> model.transform(test0).head().newPrediction + 0.0 + >>> test1 = sc.parallelize([Row(features=Vectors.sparse(4, [0], [1.0]))]).toDF() + >>> model.transform(test1).head().newPrediction + 2.0 + >>> test2 = sc.parallelize([Row(features=Vectors.dense(0.5, 0.4, 0.3, 0.2))]).toDF() + >>> model.transform(test2).head().newPrediction + 0.0 + >>> model_path = temp_path + "/ovr_model" + >>> model.save(model_path) + >>> model2 = OneVsRestModel.load(model_path) + >>> model2.transform(test0).head().newPrediction + 0.0 + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> model.transform(test2).columns + ['features', 'rawPrediction', 'newPrediction'] + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + rawPredictionCol: str = "rawPrediction", + classifier: Optional[Classifier[CM]] = None, + weightCol: Optional[str] = None, + parallelism: int = 1, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + rawPredictionCol="rawPrediction", classifier=None, weightCol=None, parallelism=1): + """ + super(OneVsRest, self).__init__() + self._setDefault(parallelism=1) + kwargs = self._input_kwargs + self._set(**kwargs) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + rawPredictionCol: str = "rawPrediction", + classifier: Optional[Classifier[CM]] = None, + weightCol: Optional[str] = None, + parallelism: int = 1, + ) -> "OneVsRest": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + rawPredictionCol="rawPrediction", classifier=None, weightCol=None, parallelism=1): + Sets params for OneVsRest. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setClassifier(self, value: Classifier[CM]) -> "OneVsRest": + """ + Sets the value of :py:attr:`classifier`. + """ + return self._set(classifier=value) + + def setLabelCol(self, value: str) -> "OneVsRest": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def setFeaturesCol(self, value: str) -> "OneVsRest": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setPredictionCol(self, value: str) -> "OneVsRest": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + def setRawPredictionCol(self, value: str) -> "OneVsRest": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + def setWeightCol(self, value: str) -> "OneVsRest": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + def setParallelism(self, value: int) -> "OneVsRest": + """ + Sets the value of :py:attr:`parallelism`. + """ + return self._set(parallelism=value) + + def _fit(self, dataset: DataFrame) -> "OneVsRestModel": + labelCol = self.getLabelCol() + featuresCol = self.getFeaturesCol() + predictionCol = self.getPredictionCol() + classifier = self.getClassifier() + + numClasses = ( + int(cast(Row, dataset.agg({labelCol: "max"}).head())["max(" + labelCol + ")"]) + 1 + ) + + weightCol = None + if self.isDefined(self.weightCol) and self.getWeightCol(): + if isinstance(classifier, HasWeightCol): + weightCol = self.getWeightCol() + else: + warnings.warn( + "weightCol is ignored, " "as it is not supported by {} now.".format(classifier) + ) + + if weightCol: + multiclassLabeled = dataset.select(labelCol, featuresCol, weightCol) + else: + multiclassLabeled = dataset.select(labelCol, featuresCol) + + # persist if underlying dataset is not persistent. + handlePersistence = dataset.storageLevel == StorageLevel(False, False, False, False) + if handlePersistence: + multiclassLabeled.persist(StorageLevel.MEMORY_AND_DISK) + + def trainSingleClass(index: int) -> CM: + binaryLabelCol = "mc2b$" + str(index) + trainingDataset = multiclassLabeled.withColumn( + binaryLabelCol, + when(multiclassLabeled[labelCol] == float(index), 1.0).otherwise(0.0), + ) + paramMap = dict( + [ + (classifier.labelCol, binaryLabelCol), + (classifier.featuresCol, featuresCol), + (classifier.predictionCol, predictionCol), + ] + ) + if weightCol: + paramMap[cast(HasWeightCol, classifier).weightCol] = weightCol + return classifier.fit(trainingDataset, paramMap) + + pool = ThreadPool(processes=min(self.getParallelism(), numClasses)) + + models = pool.map(inheritable_thread_target(trainSingleClass), range(numClasses)) + + if handlePersistence: + multiclassLabeled.unpersist() + + return self._copyValues(OneVsRestModel(models=models)) + + def copy(self, extra: Optional["ParamMap"] = None) -> "OneVsRest": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This creates a deep copy of the embedded paramMap, + and copies the embedded and extra parameters over. + + .. versionadded:: 2.0.0 + + Examples + -------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`OneVsRest` + Copy of this instance + """ + if extra is None: + extra = dict() + newOvr = Params.copy(self, extra) + if self.isSet(self.classifier): + newOvr.setClassifier(self.getClassifier().copy(extra)) + return newOvr + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "OneVsRest": + """ + Given a Java OneVsRest, create and return a Python wrapper of it. + Used for ML persistence. + """ + featuresCol = java_stage.getFeaturesCol() + labelCol = java_stage.getLabelCol() + predictionCol = java_stage.getPredictionCol() + rawPredictionCol = java_stage.getRawPredictionCol() + classifier: Classifier = JavaParams._from_java(java_stage.getClassifier()) + parallelism = java_stage.getParallelism() + py_stage = cls( + featuresCol=featuresCol, + labelCol=labelCol, + predictionCol=predictionCol, + rawPredictionCol=rawPredictionCol, + classifier=classifier, + parallelism=parallelism, + ) + if java_stage.isDefined(java_stage.getParam("weightCol")): + py_stage.setWeightCol(java_stage.getWeightCol()) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java OneVsRest. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.classification.OneVsRest", self.uid + ) + _java_obj.setClassifier(cast(_JavaClassifier, self.getClassifier())._to_java()) + _java_obj.setParallelism(self.getParallelism()) + _java_obj.setFeaturesCol(self.getFeaturesCol()) + _java_obj.setLabelCol(self.getLabelCol()) + _java_obj.setPredictionCol(self.getPredictionCol()) + if self.isDefined(self.weightCol) and self.getWeightCol(): + _java_obj.setWeightCol(self.getWeightCol()) + _java_obj.setRawPredictionCol(self.getRawPredictionCol()) + return _java_obj + + @classmethod + def read(cls) -> "OneVsRestReader": + return OneVsRestReader(cls) + + def write(self) -> MLWriter: + if isinstance(self.getClassifier(), JavaMLWritable): + return JavaMLWriter(self) # type: ignore[arg-type] + else: + return OneVsRestWriter(self) + + +class _OneVsRestSharedReadWrite: + @staticmethod + def saveImpl( + instance: Union[OneVsRest, "OneVsRestModel"], + sc: SparkContext, + path: str, + extraMetadata: Optional[Dict[str, Any]] = None, + ) -> None: + skipParams = ["classifier"] + jsonParams = DefaultParamsWriter.extractJsonParams(instance, skipParams) + DefaultParamsWriter.saveMetadata( + instance, path, sc, paramMap=jsonParams, extraMetadata=extraMetadata + ) + classifierPath = os.path.join(path, "classifier") + cast(MLWritable, instance.getClassifier()).save(classifierPath) + + @staticmethod + def loadClassifier(path: str, sc: SparkContext) -> Union[OneVsRest, "OneVsRestModel"]: + classifierPath = os.path.join(path, "classifier") + return DefaultParamsReader.loadParamsInstance(classifierPath, sc) + + @staticmethod + def validateParams(instance: Union[OneVsRest, "OneVsRestModel"]) -> None: + elems_to_check: List[Params] = [instance.getClassifier()] + if isinstance(instance, OneVsRestModel): + elems_to_check.extend(instance.models) + + for elem in elems_to_check: + if not isinstance(elem, MLWritable): + raise ValueError( + f"OneVsRest write will fail because it contains {elem.uid} " + f"which is not writable." + ) + + +@inherit_doc +class OneVsRestReader(MLReader[OneVsRest]): + def __init__(self, cls: Type[OneVsRest]) -> None: + super(OneVsRestReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> OneVsRest: + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + classifier = cast(Classifier, _OneVsRestSharedReadWrite.loadClassifier(path, self.sc)) + ova: OneVsRest = OneVsRest(classifier=classifier)._resetUid(metadata["uid"]) + DefaultParamsReader.getAndSetParams(ova, metadata, skipParams=["classifier"]) + return ova + + +@inherit_doc +class OneVsRestWriter(MLWriter): + def __init__(self, instance: OneVsRest): + super(OneVsRestWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _OneVsRestSharedReadWrite.validateParams(self.instance) + _OneVsRestSharedReadWrite.saveImpl(self.instance, self.sc, path) + + +class OneVsRestModel( + Model, + _OneVsRestParams, + MLReadable["OneVsRestModel"], + MLWritable, +): + """ + Model fitted by OneVsRest. + This stores the models resulting from training k binary classifiers: one for each class. + Each example is scored against all k models, and the model with the highest score + is picked to label the example. + + .. versionadded:: 2.0.0 + """ + + def setFeaturesCol(self, value: str) -> "OneVsRestModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setPredictionCol(self, value: str) -> "OneVsRestModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + def setRawPredictionCol(self, value: str) -> "OneVsRestModel": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + def __init__(self, models: List[ClassificationModel]): + super(OneVsRestModel, self).__init__() + self.models = models + if not isinstance(models[0], JavaMLWritable): + return + # set java instance + java_models = [cast(_JavaClassificationModel, model)._to_java() for model in self.models] + sc = SparkContext._active_spark_context + assert sc is not None and sc._gateway is not None + + java_models_array = JavaWrapper._new_java_array( + java_models, sc._gateway.jvm.org.apache.spark.ml.classification.ClassificationModel + ) + # TODO: need to set metadata + metadata = JavaParams._new_java_obj("org.apache.spark.sql.types.Metadata") + self._java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.classification.OneVsRestModel", + self.uid, + metadata.empty(), + java_models_array, + ) + + def _transform(self, dataset: DataFrame) -> DataFrame: + # determine the input columns: these need to be passed through + origCols = dataset.columns + + # add an accumulator column to store predictions of all the models + accColName = "mbc$acc" + str(uuid.uuid4()) + initUDF = udf(lambda _: [], ArrayType(DoubleType())) + newDataset = dataset.withColumn(accColName, initUDF(dataset[origCols[0]])) + + # persist if underlying dataset is not persistent. + handlePersistence = dataset.storageLevel == StorageLevel(False, False, False, False) + if handlePersistence: + newDataset.persist(StorageLevel.MEMORY_AND_DISK) + + # update the accumulator column with the result of prediction of models + aggregatedDataset = newDataset + for index, model in enumerate(self.models): + rawPredictionCol = self.getRawPredictionCol() + + columns = origCols + [rawPredictionCol, accColName] + + # add temporary column to store intermediate scores and update + tmpColName = "mbc$tmp" + str(uuid.uuid4()) + updateUDF = udf( + lambda predictions, prediction: predictions + [prediction.tolist()[1]], + ArrayType(DoubleType()), + ) + transformedDataset = model.transform(aggregatedDataset).select(*columns) + updatedDataset = transformedDataset.withColumn( + tmpColName, + updateUDF(transformedDataset[accColName], transformedDataset[rawPredictionCol]), + ) + newColumns = origCols + [tmpColName] + + # switch out the intermediate column with the accumulator column + aggregatedDataset = updatedDataset.select(*newColumns).withColumnRenamed( + tmpColName, accColName + ) + + if handlePersistence: + newDataset.unpersist() + + if self.getRawPredictionCol(): + + def func(predictions: Iterable[float]) -> Vector: + predArray: List[float] = [] + for x in predictions: + predArray.append(x) + return Vectors.dense(predArray) + + rawPredictionUDF = udf(func, VectorUDT()) + aggregatedDataset = aggregatedDataset.withColumn( + self.getRawPredictionCol(), rawPredictionUDF(aggregatedDataset[accColName]) + ) + + if self.getPredictionCol(): + # output the index of the classifier with highest confidence as prediction + labelUDF = udf( + lambda predictions: float( + max(enumerate(predictions), key=operator.itemgetter(1))[0] + ), + DoubleType(), + ) + aggregatedDataset = aggregatedDataset.withColumn( + self.getPredictionCol(), labelUDF(aggregatedDataset[accColName]) + ) + return aggregatedDataset.drop(accColName) + + def copy(self, extra: Optional["ParamMap"] = None) -> "OneVsRestModel": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This creates a deep copy of the embedded paramMap, + and copies the embedded and extra parameters over. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`OneVsRestModel` + Copy of this instance + """ + if extra is None: + extra = dict() + newModel = Params.copy(self, extra) + newModel.models = [model.copy(extra) for model in self.models] + return newModel + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "OneVsRestModel": + """ + Given a Java OneVsRestModel, create and return a Python wrapper of it. + Used for ML persistence. + """ + featuresCol = java_stage.getFeaturesCol() + labelCol = java_stage.getLabelCol() + predictionCol = java_stage.getPredictionCol() + classifier: Classifier = JavaParams._from_java(java_stage.getClassifier()) + models: List[ClassificationModel] = [ + JavaParams._from_java(model) for model in java_stage.models() + ] + py_stage = cls(models=models).setPredictionCol(predictionCol).setFeaturesCol(featuresCol) + py_stage._set(labelCol=labelCol) + if java_stage.isDefined(java_stage.getParam("weightCol")): + py_stage._set(weightCol=java_stage.getWeightCol()) + py_stage._set(classifier=classifier) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java OneVsRestModel. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._gateway is not None + + java_models = [cast(_JavaClassificationModel, model)._to_java() for model in self.models] + java_models_array = JavaWrapper._new_java_array( + java_models, sc._gateway.jvm.org.apache.spark.ml.classification.ClassificationModel + ) + metadata = JavaParams._new_java_obj("org.apache.spark.sql.types.Metadata") + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.classification.OneVsRestModel", + self.uid, + metadata.empty(), + java_models_array, + ) + _java_obj.set("classifier", cast(_JavaClassifier, self.getClassifier())._to_java()) + _java_obj.set("featuresCol", self.getFeaturesCol()) + _java_obj.set("labelCol", self.getLabelCol()) + _java_obj.set("predictionCol", self.getPredictionCol()) + if self.isDefined(self.weightCol) and self.getWeightCol(): + _java_obj.set("weightCol", self.getWeightCol()) + return _java_obj + + @classmethod + def read(cls) -> "OneVsRestModelReader": + return OneVsRestModelReader(cls) + + def write(self) -> MLWriter: + if all( + map( + lambda elem: isinstance(elem, JavaMLWritable), + [self.getClassifier()] + self.models, # type: ignore[operator] + ) + ): + return JavaMLWriter(self) # type: ignore[arg-type] + else: + return OneVsRestModelWriter(self) + + +@inherit_doc +class OneVsRestModelReader(MLReader[OneVsRestModel]): + def __init__(self, cls: Type[OneVsRestModel]): + super(OneVsRestModelReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> OneVsRestModel: + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + classifier = _OneVsRestSharedReadWrite.loadClassifier(path, self.sc) + numClasses = metadata["numClasses"] + subModels = [None] * numClasses + for idx in range(numClasses): + subModelPath = os.path.join(path, f"model_{idx}") + subModels[idx] = DefaultParamsReader.loadParamsInstance(subModelPath, self.sc) + ovaModel = OneVsRestModel(cast(List[ClassificationModel], subModels))._resetUid( + metadata["uid"] + ) + ovaModel.set(ovaModel.classifier, classifier) + DefaultParamsReader.getAndSetParams(ovaModel, metadata, skipParams=["classifier"]) + return ovaModel + + +@inherit_doc +class OneVsRestModelWriter(MLWriter): + def __init__(self, instance: OneVsRestModel): + super(OneVsRestModelWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _OneVsRestSharedReadWrite.validateParams(self.instance) + instance = self.instance + numClasses = len(instance.models) + extraMetadata = {"numClasses": numClasses} + _OneVsRestSharedReadWrite.saveImpl(instance, self.sc, path, extraMetadata=extraMetadata) + for idx in range(numClasses): + subModelPath = os.path.join(path, f"model_{idx}") + cast(MLWritable, instance.models[idx]).save(subModelPath) + + +@inherit_doc +class FMClassifier( + _JavaProbabilisticClassifier["FMClassificationModel"], + _FactorizationMachinesParams, + JavaMLWritable, + JavaMLReadable["FMClassifier"], +): + """ + Factorization Machines learning algorithm for classification. + + Solver supports: + + * gd (normal mini-batch gradient descent) + * adamW (default) + + .. versionadded:: 3.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.classification import FMClassifier + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> fm = FMClassifier(factorSize=2) + >>> fm.setSeed(11) + FMClassifier... + >>> model = fm.fit(df) + >>> model.getMaxIter() + 100 + >>> test0 = spark.createDataFrame([ + ... (Vectors.dense(-1.0),), + ... (Vectors.dense(0.5),), + ... (Vectors.dense(1.0),), + ... (Vectors.dense(2.0),)], ["features"]) + >>> model.predictRaw(test0.head().features) + DenseVector([22.13..., -22.13...]) + >>> model.predictProbability(test0.head().features) + DenseVector([1.0, 0.0]) + >>> model.transform(test0).select("features", "probability").show(10, False) + +--------+------------------------------------------+ + |features|probability | + +--------+------------------------------------------+ + |[-1.0] |[0.9999999997574736,2.425264676902229E-10]| + |[0.5] |[0.47627851732981163,0.5237214826701884] | + |[1.0] |[5.491554426243495E-4,0.9994508445573757] | + |[2.0] |[2.005766663870645E-10,0.9999999997994233]| + +--------+------------------------------------------+ + ... + >>> model.intercept + -7.316665276826291 + >>> model.linear + DenseVector([14.8232]) + >>> model.factors + DenseMatrix(1, 2, [0.0163, -0.0051], 1) + >>> model_path = temp_path + "/fm_model" + >>> model.save(model_path) + >>> model2 = FMClassificationModel.load(model_path) + >>> model2.intercept + -7.316665276826291 + >>> model2.linear + DenseVector([14.8232]) + >>> model2.factors + DenseMatrix(1, 2, [0.0163, -0.0051], 1) + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + factorSize: int = 8, + fitIntercept: bool = True, + fitLinear: bool = True, + regParam: float = 0.0, + miniBatchFraction: float = 1.0, + initStd: float = 0.01, + maxIter: int = 100, + stepSize: float = 1.0, + tol: float = 1e-6, + solver: str = "adamW", + thresholds: Optional[List[float]] = None, + seed: Optional[int] = None, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + factorSize=8, fitIntercept=True, fitLinear=True, regParam=0.0, \ + miniBatchFraction=1.0, initStd=0.01, maxIter=100, stepSize=1.0, \ + tol=1e-6, solver="adamW", thresholds=None, seed=None) + """ + super(FMClassifier, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.classification.FMClassifier", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + rawPredictionCol: str = "rawPrediction", + factorSize: int = 8, + fitIntercept: bool = True, + fitLinear: bool = True, + regParam: float = 0.0, + miniBatchFraction: float = 1.0, + initStd: float = 0.01, + maxIter: int = 100, + stepSize: float = 1.0, + tol: float = 1e-6, + solver: str = "adamW", + thresholds: Optional[List[float]] = None, + seed: Optional[int] = None, + ) -> "FMClassifier": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + probabilityCol="probability", rawPredictionCol="rawPrediction", \ + factorSize=8, fitIntercept=True, fitLinear=True, regParam=0.0, \ + miniBatchFraction=1.0, initStd=0.01, maxIter=100, stepSize=1.0, \ + tol=1e-6, solver="adamW", thresholds=None, seed=None) + Sets Params for FMClassifier. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "FMClassificationModel": + return FMClassificationModel(java_model) + + @since("3.0.0") + def setFactorSize(self, value: int) -> "FMClassifier": + """ + Sets the value of :py:attr:`factorSize`. + """ + return self._set(factorSize=value) + + @since("3.0.0") + def setFitLinear(self, value: bool) -> "FMClassifier": + """ + Sets the value of :py:attr:`fitLinear`. + """ + return self._set(fitLinear=value) + + @since("3.0.0") + def setMiniBatchFraction(self, value: float) -> "FMClassifier": + """ + Sets the value of :py:attr:`miniBatchFraction`. + """ + return self._set(miniBatchFraction=value) + + @since("3.0.0") + def setInitStd(self, value: float) -> "FMClassifier": + """ + Sets the value of :py:attr:`initStd`. + """ + return self._set(initStd=value) + + @since("3.0.0") + def setMaxIter(self, value: int) -> "FMClassifier": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("3.0.0") + def setStepSize(self, value: float) -> "FMClassifier": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + @since("3.0.0") + def setTol(self, value: float) -> "FMClassifier": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("3.0.0") + def setSolver(self, value: str) -> "FMClassifier": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + @since("3.0.0") + def setSeed(self, value: int) -> "FMClassifier": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setFitIntercept(self, value: bool) -> "FMClassifier": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + @since("3.0.0") + def setRegParam(self, value: float) -> "FMClassifier": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + +class FMClassificationModel( + _JavaProbabilisticClassificationModel[Vector], + _FactorizationMachinesParams, + JavaMLWritable, + JavaMLReadable["FMClassificationModel"], + HasTrainingSummary, +): + """ + Model fitted by :class:`FMClassifier`. + + .. versionadded:: 3.0.0 + """ + + @property + @since("3.0.0") + def intercept(self) -> float: + """ + Model intercept. + """ + return self._call_java("intercept") + + @property + @since("3.0.0") + def linear(self) -> Vector: + """ + Model linear term. + """ + return self._call_java("linear") + + @property + @since("3.0.0") + def factors(self) -> Matrix: + """ + Model factor term. + """ + return self._call_java("factors") + + @since("3.1.0") + def summary(self) -> "FMClassificationTrainingSummary": + """ + Gets summary (accuracy/precision/recall, objective history, total iterations) of model + trained on the training set. An exception is thrown if `trainingSummary is None`. + """ + if self.hasSummary: + return FMClassificationTrainingSummary(super(FMClassificationModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "FMClassificationSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on. + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_fm_summary = self._call_java("evaluate", dataset) + return FMClassificationSummary(java_fm_summary) + + +class FMClassificationSummary(_BinaryClassificationSummary): + """ + Abstraction for FMClassifier Results for a given model. + + .. versionadded:: 3.1.0 + """ + + pass + + +@inherit_doc +class FMClassificationTrainingSummary(FMClassificationSummary, _TrainingSummary): + """ + Abstraction for FMClassifier Training results. + + .. versionadded:: 3.1.0 + """ + + pass + + +if __name__ == "__main__": + import doctest + import pyspark.ml.classification + from pyspark.sql import SparkSession + + globs = pyspark.ml.classification.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.classification tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/clustering.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/clustering.py new file mode 100644 index 0000000..4110878 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/clustering.py @@ -0,0 +1,2188 @@ +# +# 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 sys +import warnings + +from typing import Any, Dict, List, Optional, TYPE_CHECKING + +import numpy as np + +from pyspark import since, keyword_only +from pyspark.ml.param.shared import ( + HasMaxIter, + HasFeaturesCol, + HasSeed, + HasPredictionCol, + HasAggregationDepth, + HasWeightCol, + HasTol, + HasProbabilityCol, + HasDistanceMeasure, + HasCheckpointInterval, + HasSolver, + HasMaxBlockSizeInMB, + Param, + Params, + TypeConverters, +) +from pyspark.ml.util import ( + JavaMLWritable, + JavaMLReadable, + GeneralJavaMLWritable, + HasTrainingSummary, + SparkContext, +) +from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams, JavaWrapper +from pyspark.ml.common import inherit_doc, _java2py +from pyspark.ml.stat import MultivariateGaussian +from pyspark.sql import DataFrame +from pyspark.ml.linalg import Vector, Matrix + +if TYPE_CHECKING: + from pyspark.ml._typing import M + from py4j.java_gateway import JavaObject + + +__all__ = [ + "BisectingKMeans", + "BisectingKMeansModel", + "BisectingKMeansSummary", + "KMeans", + "KMeansModel", + "KMeansSummary", + "GaussianMixture", + "GaussianMixtureModel", + "GaussianMixtureSummary", + "LDA", + "LDAModel", + "LocalLDAModel", + "DistributedLDAModel", + "PowerIterationClustering", +] + + +class ClusteringSummary(JavaWrapper): + """ + Clustering results for a given model. + + .. versionadded:: 2.1.0 + """ + + @property + @since("2.1.0") + def predictionCol(self) -> str: + """ + Name for column of predicted clusters in `predictions`. + """ + return self._call_java("predictionCol") + + @property + @since("2.1.0") + def predictions(self) -> DataFrame: + """ + DataFrame produced by the model's `transform` method. + """ + return self._call_java("predictions") + + @property + @since("2.1.0") + def featuresCol(self) -> str: + """ + Name for column of features in `predictions`. + """ + return self._call_java("featuresCol") + + @property + @since("2.1.0") + def k(self) -> int: + """ + The number of clusters the model was trained with. + """ + return self._call_java("k") + + @property + @since("2.1.0") + def cluster(self) -> DataFrame: + """ + DataFrame of predicted cluster centers for each training data point. + """ + return self._call_java("cluster") + + @property + @since("2.1.0") + def clusterSizes(self) -> List[int]: + """ + Size of (number of data points in) each cluster. + """ + return self._call_java("clusterSizes") + + @property + @since("2.4.0") + def numIter(self) -> int: + """ + Number of iterations. + """ + return self._call_java("numIter") + + +@inherit_doc +class _GaussianMixtureParams( + HasMaxIter, + HasFeaturesCol, + HasSeed, + HasPredictionCol, + HasProbabilityCol, + HasTol, + HasAggregationDepth, + HasWeightCol, +): + """ + Params for :py:class:`GaussianMixture` and :py:class:`GaussianMixtureModel`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "Number of independent Gaussians in the mixture model. " + "Must be > 1.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self, *args: Any): + super(_GaussianMixtureParams, self).__init__(*args) + self._setDefault(k=2, tol=0.01, maxIter=100, aggregationDepth=2) + + @since("2.0.0") + def getK(self) -> int: + """ + Gets the value of `k` + """ + return self.getOrDefault(self.k) + + +class GaussianMixtureModel( + JavaModel, + _GaussianMixtureParams, + JavaMLWritable, + JavaMLReadable["GaussianMixtureModel"], + HasTrainingSummary["GaussianMixtureSummary"], +): + """ + Model fitted by GaussianMixture. + + .. versionadded:: 2.0.0 + """ + + @since("3.0.0") + def setFeaturesCol(self, value: str) -> "GaussianMixtureModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "GaussianMixtureModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("3.0.0") + def setProbabilityCol(self, value: str) -> "GaussianMixtureModel": + """ + Sets the value of :py:attr:`probabilityCol`. + """ + return self._set(probabilityCol=value) + + @property + @since("2.0.0") + def weights(self) -> List[float]: + """ + Weight for each Gaussian distribution in the mixture. + This is a multinomial probability distribution over the k Gaussians, + where weights[i] is the weight for Gaussian i, and weights sum to 1. + """ + return self._call_java("weights") + + @property + @since("3.0.0") + def gaussians(self) -> List[MultivariateGaussian]: + """ + Array of :py:class:`MultivariateGaussian` where gaussians[i] represents + the Multivariate Gaussian (Normal) Distribution for Gaussian i + """ + sc = SparkContext._active_spark_context + assert sc is not None and self._java_obj is not None + + jgaussians = self._java_obj.gaussians() + return [ + MultivariateGaussian(_java2py(sc, jgaussian.mean()), _java2py(sc, jgaussian.cov())) + for jgaussian in jgaussians + ] + + @property + @since("2.0.0") + def gaussiansDF(self) -> DataFrame: + """ + Retrieve Gaussian distributions as a DataFrame. + Each row represents a Gaussian Distribution. + The DataFrame has two columns: mean (Vector) and cov (Matrix). + """ + return self._call_java("gaussiansDF") + + @property + @since("2.1.0") + def summary(self) -> "GaussianMixtureSummary": + """ + Gets summary (cluster assignments, cluster sizes) of the model trained on the + training set. An exception is thrown if no summary exists. + """ + if self.hasSummary: + return GaussianMixtureSummary(super(GaussianMixtureModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + @since("3.0.0") + def predict(self, value: Vector) -> int: + """ + Predict label for the given features. + """ + return self._call_java("predict", value) + + @since("3.0.0") + def predictProbability(self, value: Vector) -> Vector: + """ + Predict probability for the given features. + """ + return self._call_java("predictProbability", value) + + +@inherit_doc +class GaussianMixture( + JavaEstimator[GaussianMixtureModel], + _GaussianMixtureParams, + JavaMLWritable, + JavaMLReadable["GaussianMixture"], +): + """ + GaussianMixture clustering. + This class performs expectation maximization for multivariate Gaussian + Mixture Models (GMMs). A GMM represents a composite distribution of + independent Gaussian distributions with associated "mixing" weights + specifying each's contribution to the composite. + + Given a set of sample points, this class will maximize the log-likelihood + for a mixture of k Gaussians, iterating until the log-likelihood changes by + less than convergenceTol, or until it has reached the max number of iterations. + While this process is generally guaranteed to converge, it is not guaranteed + to find a global optimum. + + .. versionadded:: 2.0.0 + + Notes + ----- + For high-dimensional data (with many features), this algorithm may perform poorly. + This is due to high-dimensional data (a) making it difficult to cluster at all + (based on statistical/theoretical arguments) and (b) numerical issues with + Gaussian distributions. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + + >>> data = [(Vectors.dense([-0.1, -0.05 ]),), + ... (Vectors.dense([-0.01, -0.1]),), + ... (Vectors.dense([0.9, 0.8]),), + ... (Vectors.dense([0.75, 0.935]),), + ... (Vectors.dense([-0.83, -0.68]),), + ... (Vectors.dense([-0.91, -0.76]),)] + >>> df = spark.createDataFrame(data, ["features"]) + >>> gm = GaussianMixture(k=3, tol=0.0001, seed=10) + >>> gm.getMaxIter() + 100 + >>> gm.setMaxIter(30) + GaussianMixture... + >>> gm.getMaxIter() + 30 + >>> model = gm.fit(df) + >>> model.getAggregationDepth() + 2 + >>> model.getFeaturesCol() + 'features' + >>> model.setPredictionCol("newPrediction") + GaussianMixtureModel... + >>> model.predict(df.head().features) + 2 + >>> model.predictProbability(df.head().features) + DenseVector([0.0, 0.0, 1.0]) + >>> model.hasSummary + True + >>> summary = model.summary + >>> summary.k + 3 + >>> summary.clusterSizes + [2, 2, 2] + >>> weights = model.weights + >>> len(weights) + 3 + >>> gaussians = model.gaussians + >>> len(gaussians) + 3 + >>> gaussians[0].mean + DenseVector([0.825, 0.8675]) + >>> gaussians[0].cov + DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], 0) + >>> gaussians[1].mean + DenseVector([-0.87, -0.72]) + >>> gaussians[1].cov + DenseMatrix(2, 2, [0.0016, 0.0016, 0.0016, 0.0016], 0) + >>> gaussians[2].mean + DenseVector([-0.055, -0.075]) + >>> gaussians[2].cov + DenseMatrix(2, 2, [0.002, -0.0011, -0.0011, 0.0006], 0) + >>> model.gaussiansDF.select("mean").head() + Row(mean=DenseVector([0.825, 0.8675])) + >>> model.gaussiansDF.select("cov").head() + Row(cov=DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], False)) + >>> transformed = model.transform(df).select("features", "newPrediction") + >>> rows = transformed.collect() + >>> rows[4].newPrediction == rows[5].newPrediction + True + >>> rows[2].newPrediction == rows[3].newPrediction + True + >>> gmm_path = temp_path + "/gmm" + >>> gm.save(gmm_path) + >>> gm2 = GaussianMixture.load(gmm_path) + >>> gm2.getK() + 3 + >>> model_path = temp_path + "/gmm_model" + >>> model.save(model_path) + >>> model2 = GaussianMixtureModel.load(model_path) + >>> model2.hasSummary + False + >>> model2.weights == model.weights + True + >>> model2.gaussians[0].mean == model.gaussians[0].mean + True + >>> model2.gaussians[0].cov == model.gaussians[0].cov + True + >>> model2.gaussians[1].mean == model.gaussians[1].mean + True + >>> model2.gaussians[1].cov == model.gaussians[1].cov + True + >>> model2.gaussians[2].mean == model.gaussians[2].mean + True + >>> model2.gaussians[2].cov == model.gaussians[2].cov + True + >>> model2.gaussiansDF.select("mean").head() + Row(mean=DenseVector([0.825, 0.8675])) + >>> model2.gaussiansDF.select("cov").head() + Row(cov=DenseMatrix(2, 2, [0.0056, -0.0051, -0.0051, 0.0046], False)) + >>> model.transform(df).take(1) == model2.transform(df).take(1) + True + >>> gm2.setWeightCol("weight") + GaussianMixture... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + k: int = 2, + probabilityCol: str = "probability", + tol: float = 0.01, + maxIter: int = 100, + seed: Optional[int] = None, + aggregationDepth: int = 2, + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \ + probabilityCol="probability", tol=0.01, maxIter=100, seed=None, \ + aggregationDepth=2, weightCol=None) + """ + super(GaussianMixture, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.clustering.GaussianMixture", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "GaussianMixtureModel": + return GaussianMixtureModel(java_model) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + k: int = 2, + probabilityCol: str = "probability", + tol: float = 0.01, + maxIter: int = 100, + seed: Optional[int] = None, + aggregationDepth: int = 2, + weightCol: Optional[str] = None, + ) -> "GaussianMixture": + """ + setParams(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \ + probabilityCol="probability", tol=0.01, maxIter=100, seed=None, \ + aggregationDepth=2, weightCol=None) + + Sets params for GaussianMixture. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setK(self, value: int) -> "GaussianMixture": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + @since("2.0.0") + def setMaxIter(self, value: int) -> "GaussianMixture": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.0.0") + def setFeaturesCol(self, value: str) -> "GaussianMixture": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("2.0.0") + def setPredictionCol(self, value: str) -> "GaussianMixture": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("2.0.0") + def setProbabilityCol(self, value: str) -> "GaussianMixture": + """ + Sets the value of :py:attr:`probabilityCol`. + """ + return self._set(probabilityCol=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "GaussianMixture": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("2.0.0") + def setSeed(self, value: int) -> "GaussianMixture": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("2.0.0") + def setTol(self, value: float) -> "GaussianMixture": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("3.0.0") + def setAggregationDepth(self, value: int) -> "GaussianMixture": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + +class GaussianMixtureSummary(ClusteringSummary): + """ + Gaussian mixture clustering results for a given model. + + .. versionadded:: 2.1.0 + """ + + @property + @since("2.1.0") + def probabilityCol(self) -> str: + """ + Name for column of predicted probability of each cluster in `predictions`. + """ + return self._call_java("probabilityCol") + + @property + @since("2.1.0") + def probability(self) -> DataFrame: + """ + DataFrame of probabilities of each cluster for each training data point. + """ + return self._call_java("probability") + + @property + @since("2.2.0") + def logLikelihood(self) -> float: + """ + Total log-likelihood for this model on the given data. + """ + return self._call_java("logLikelihood") + + +class KMeansSummary(ClusteringSummary): + """ + Summary of KMeans. + + .. versionadded:: 2.1.0 + """ + + @property + @since("2.4.0") + def trainingCost(self) -> float: + """ + K-means cost (sum of squared distances to the nearest centroid for all points in the + training dataset). This is equivalent to sklearn's inertia. + """ + return self._call_java("trainingCost") + + +@inherit_doc +class _KMeansParams( + HasMaxIter, + HasFeaturesCol, + HasSeed, + HasPredictionCol, + HasTol, + HasDistanceMeasure, + HasWeightCol, + HasSolver, + HasMaxBlockSizeInMB, +): + """ + Params for :py:class:`KMeans` and :py:class:`KMeansModel`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "The number of clusters to create. Must be > 1.", + typeConverter=TypeConverters.toInt, + ) + initMode: Param[str] = Param( + Params._dummy(), + "initMode", + 'The initialization algorithm. This can be either "random" to ' + + 'choose random points as initial cluster centers, or "k-means||" ' + + "to use a parallel variant of k-means++", + typeConverter=TypeConverters.toString, + ) + initSteps: Param[int] = Param( + Params._dummy(), + "initSteps", + "The number of steps for k-means|| " + "initialization mode. Must be > 0.", + typeConverter=TypeConverters.toInt, + ) + solver: Param[str] = Param( + Params._dummy(), + "solver", + "The solver algorithm for optimization. Supported " + "options: auto, row, block.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_KMeansParams, self).__init__(*args) + self._setDefault( + k=2, + initMode="k-means||", + initSteps=2, + tol=1e-4, + maxIter=20, + distanceMeasure="euclidean", + solver="auto", + maxBlockSizeInMB=0.0, + ) + + @since("1.5.0") + def getK(self) -> int: + """ + Gets the value of `k` + """ + return self.getOrDefault(self.k) + + @since("1.5.0") + def getInitMode(self) -> str: + """ + Gets the value of `initMode` + """ + return self.getOrDefault(self.initMode) + + @since("1.5.0") + def getInitSteps(self) -> int: + """ + Gets the value of `initSteps` + """ + return self.getOrDefault(self.initSteps) + + +class KMeansModel( + JavaModel, + _KMeansParams, + GeneralJavaMLWritable, + JavaMLReadable["KMeansModel"], + HasTrainingSummary["KMeansSummary"], +): + """ + Model fitted by KMeans. + + .. versionadded:: 1.5.0 + """ + + @since("3.0.0") + def setFeaturesCol(self, value: str) -> "KMeansModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "KMeansModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("1.5.0") + def clusterCenters(self) -> List[np.ndarray]: + """Get the cluster centers, represented as a list of NumPy arrays.""" + return [c.toArray() for c in self._call_java("clusterCenters")] + + @property + @since("2.1.0") + def summary(self) -> KMeansSummary: + """ + Gets summary (cluster assignments, cluster sizes) of the model trained on the + training set. An exception is thrown if no summary exists. + """ + if self.hasSummary: + return KMeansSummary(super(KMeansModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + @since("3.0.0") + def predict(self, value: Vector) -> int: + """ + Predict label for the given features. + """ + return self._call_java("predict", value) + + +@inherit_doc +class KMeans(JavaEstimator[KMeansModel], _KMeansParams, JavaMLWritable, JavaMLReadable["KMeans"]): + """ + K-means clustering with a k-means++ like initialization mode + (the k-means|| algorithm by Bahmani et al). + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> data = [(Vectors.dense([0.0, 0.0]), 2.0), (Vectors.dense([1.0, 1.0]), 2.0), + ... (Vectors.dense([9.0, 8.0]), 2.0), (Vectors.dense([8.0, 9.0]), 2.0)] + >>> df = spark.createDataFrame(data, ["features", "weighCol"]) + >>> kmeans = KMeans(k=2) + >>> kmeans.setSeed(1) + KMeans... + >>> kmeans.setWeightCol("weighCol") + KMeans... + >>> kmeans.setMaxIter(10) + KMeans... + >>> kmeans.getMaxIter() + 10 + >>> kmeans.clear(kmeans.maxIter) + >>> kmeans.getSolver() + 'auto' + >>> model = kmeans.fit(df) + >>> model.getMaxBlockSizeInMB() + 0.0 + >>> model.getDistanceMeasure() + 'euclidean' + >>> model.setPredictionCol("newPrediction") + KMeansModel... + >>> model.predict(df.head().features) + 0 + >>> centers = model.clusterCenters() + >>> len(centers) + 2 + >>> transformed = model.transform(df).select("features", "newPrediction") + >>> rows = transformed.collect() + >>> rows[0].newPrediction == rows[1].newPrediction + True + >>> rows[2].newPrediction == rows[3].newPrediction + True + >>> model.hasSummary + True + >>> summary = model.summary + >>> summary.k + 2 + >>> summary.clusterSizes + [2, 2] + >>> summary.trainingCost + 4.0 + >>> kmeans_path = temp_path + "/kmeans" + >>> kmeans.save(kmeans_path) + >>> kmeans2 = KMeans.load(kmeans_path) + >>> kmeans2.getK() + 2 + >>> model_path = temp_path + "/kmeans_model" + >>> model.save(model_path) + >>> model2 = KMeansModel.load(model_path) + >>> model2.hasSummary + False + >>> model.clusterCenters()[0] == model2.clusterCenters()[0] + array([ True, True], dtype=bool) + >>> model.clusterCenters()[1] == model2.clusterCenters()[1] + array([ True, True], dtype=bool) + >>> model.transform(df).take(1) == model2.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + k: int = 2, + initMode: str = "k-means||", + initSteps: int = 2, + tol: float = 1e-4, + maxIter: int = 20, + seed: Optional[int] = None, + distanceMeasure: str = "euclidean", + weightCol: Optional[str] = None, + solver: str = "auto", + maxBlockSizeInMB: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \ + initMode="k-means||", initSteps=2, tol=1e-4, maxIter=20, seed=None, \ + distanceMeasure="euclidean", weightCol=None, solver="auto", \ + maxBlockSizeInMB=0.0) + """ + super(KMeans, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.clustering.KMeans", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> KMeansModel: + return KMeansModel(java_model) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + k: int = 2, + initMode: str = "k-means||", + initSteps: int = 2, + tol: float = 1e-4, + maxIter: int = 20, + seed: Optional[int] = None, + distanceMeasure: str = "euclidean", + weightCol: Optional[str] = None, + solver: str = "auto", + maxBlockSizeInMB: float = 0.0, + ) -> "KMeans": + """ + setParams(self, \\*, featuresCol="features", predictionCol="prediction", k=2, \ + initMode="k-means||", initSteps=2, tol=1e-4, maxIter=20, seed=None, \ + distanceMeasure="euclidean", weightCol=None, solver="auto", \ + maxBlockSizeInMB=0.0) + + Sets params for KMeans. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.5.0") + def setK(self, value: int) -> "KMeans": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + @since("1.5.0") + def setInitMode(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`initMode`. + """ + return self._set(initMode=value) + + @since("1.5.0") + def setInitSteps(self, value: int) -> "KMeans": + """ + Sets the value of :py:attr:`initSteps`. + """ + return self._set(initSteps=value) + + @since("2.4.0") + def setDistanceMeasure(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`distanceMeasure`. + """ + return self._set(distanceMeasure=value) + + @since("1.5.0") + def setMaxIter(self, value: int) -> "KMeans": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("1.5.0") + def setFeaturesCol(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("1.5.0") + def setPredictionCol(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("1.5.0") + def setSeed(self, value: int) -> "KMeans": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("1.5.0") + def setTol(self, value: float) -> "KMeans": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("3.4.0") + def setSolver(self, value: str) -> "KMeans": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + @since("3.4.0") + def setMaxBlockSizeInMB(self, value: float) -> "KMeans": + """ + Sets the value of :py:attr:`maxBlockSizeInMB`. + """ + return self._set(maxBlockSizeInMB=value) + + +@inherit_doc +class _BisectingKMeansParams( + HasMaxIter, + HasFeaturesCol, + HasSeed, + HasPredictionCol, + HasDistanceMeasure, + HasWeightCol, +): + """ + Params for :py:class:`BisectingKMeans` and :py:class:`BisectingKMeansModel`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "The desired number of leaf clusters. Must be > 1.", + typeConverter=TypeConverters.toInt, + ) + minDivisibleClusterSize: Param[float] = Param( + Params._dummy(), + "minDivisibleClusterSize", + "The minimum number of points (if >= 1.0) or the minimum " + + "proportion of points (if < 1.0) of a divisible cluster.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_BisectingKMeansParams, self).__init__(*args) + self._setDefault(maxIter=20, k=4, minDivisibleClusterSize=1.0) + + @since("2.0.0") + def getK(self) -> int: + """ + Gets the value of `k` or its default value. + """ + return self.getOrDefault(self.k) + + @since("2.0.0") + def getMinDivisibleClusterSize(self) -> float: + """ + Gets the value of `minDivisibleClusterSize` or its default value. + """ + return self.getOrDefault(self.minDivisibleClusterSize) + + +class BisectingKMeansModel( + JavaModel, + _BisectingKMeansParams, + JavaMLWritable, + JavaMLReadable["BisectingKMeansModel"], + HasTrainingSummary["BisectingKMeansSummary"], +): + """ + Model fitted by BisectingKMeans. + + .. versionadded:: 2.0.0 + """ + + @since("3.0.0") + def setFeaturesCol(self, value: str) -> "BisectingKMeansModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "BisectingKMeansModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("2.0.0") + def clusterCenters(self) -> List[np.ndarray]: + """Get the cluster centers, represented as a list of NumPy arrays.""" + return [c.toArray() for c in self._call_java("clusterCenters")] + + @since("2.0.0") + def computeCost(self, dataset: DataFrame) -> float: + """ + Computes the sum of squared distances between the input points + and their corresponding cluster centers. + + .. deprecated:: 3.0.0 + It will be removed in future versions. Use :py:class:`ClusteringEvaluator` instead. + You can also get the cost on the training dataset in the summary. + """ + warnings.warn( + "Deprecated in 3.0.0. It will be removed in future versions. Use " + "ClusteringEvaluator instead. You can also get the cost on the training " + "dataset in the summary.", + FutureWarning, + ) + return self._call_java("computeCost", dataset) + + @property + @since("2.1.0") + def summary(self) -> "BisectingKMeansSummary": + """ + Gets summary (cluster assignments, cluster sizes) of the model trained on the + training set. An exception is thrown if no summary exists. + """ + if self.hasSummary: + return BisectingKMeansSummary(super(BisectingKMeansModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + @since("3.0.0") + def predict(self, value: Vector) -> int: + """ + Predict label for the given features. + """ + return self._call_java("predict", value) + + +@inherit_doc +class BisectingKMeans( + JavaEstimator[BisectingKMeansModel], + _BisectingKMeansParams, + JavaMLWritable, + JavaMLReadable["BisectingKMeans"], +): + """ + A bisecting k-means algorithm based on the paper "A comparison of document clustering + techniques" by Steinbach, Karypis, and Kumar, with modification to fit Spark. + The algorithm starts from a single cluster that contains all points. + Iteratively it finds divisible clusters on the bottom level and bisects each of them using + k-means, until there are `k` leaf clusters in total or no leaf clusters are divisible. + The bisecting steps of clusters on the same level are grouped together to increase parallelism. + If bisecting all divisible clusters on the bottom level would result more than `k` leaf + clusters, larger clusters get higher priority. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> data = [(Vectors.dense([0.0, 0.0]), 2.0), (Vectors.dense([1.0, 1.0]), 2.0), + ... (Vectors.dense([9.0, 8.0]), 2.0), (Vectors.dense([8.0, 9.0]), 2.0)] + >>> df = spark.createDataFrame(data, ["features", "weighCol"]) + >>> bkm = BisectingKMeans(k=2, minDivisibleClusterSize=1.0) + >>> bkm.setMaxIter(10) + BisectingKMeans... + >>> bkm.getMaxIter() + 10 + >>> bkm.clear(bkm.maxIter) + >>> bkm.setSeed(1) + BisectingKMeans... + >>> bkm.setWeightCol("weighCol") + BisectingKMeans... + >>> bkm.getSeed() + 1 + >>> bkm.clear(bkm.seed) + >>> model = bkm.fit(df) + >>> model.getMaxIter() + 20 + >>> model.setPredictionCol("newPrediction") + BisectingKMeansModel... + >>> model.predict(df.head().features) + 0 + >>> centers = model.clusterCenters() + >>> len(centers) + 2 + >>> model.computeCost(df) + 2.0 + >>> model.hasSummary + True + >>> summary = model.summary + >>> summary.k + 2 + >>> summary.clusterSizes + [2, 2] + >>> summary.trainingCost + 4.000... + >>> transformed = model.transform(df).select("features", "newPrediction") + >>> rows = transformed.collect() + >>> rows[0].newPrediction == rows[1].newPrediction + True + >>> rows[2].newPrediction == rows[3].newPrediction + True + >>> bkm_path = temp_path + "/bkm" + >>> bkm.save(bkm_path) + >>> bkm2 = BisectingKMeans.load(bkm_path) + >>> bkm2.getK() + 2 + >>> bkm2.getDistanceMeasure() + 'euclidean' + >>> model_path = temp_path + "/bkm_model" + >>> model.save(model_path) + >>> model2 = BisectingKMeansModel.load(model_path) + >>> model2.hasSummary + False + >>> model.clusterCenters()[0] == model2.clusterCenters()[0] + array([ True, True], dtype=bool) + >>> model.clusterCenters()[1] == model2.clusterCenters()[1] + array([ True, True], dtype=bool) + >>> model.transform(df).take(1) == model2.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + maxIter: int = 20, + seed: Optional[int] = None, + k: int = 4, + minDivisibleClusterSize: float = 1.0, + distanceMeasure: str = "euclidean", + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, featuresCol="features", predictionCol="prediction", maxIter=20, \ + seed=None, k=4, minDivisibleClusterSize=1.0, distanceMeasure="euclidean", \ + weightCol=None) + """ + super(BisectingKMeans, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.clustering.BisectingKMeans", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + predictionCol: str = "prediction", + maxIter: int = 20, + seed: Optional[int] = None, + k: int = 4, + minDivisibleClusterSize: float = 1.0, + distanceMeasure: str = "euclidean", + weightCol: Optional[str] = None, + ) -> "BisectingKMeans": + """ + setParams(self, \\*, featuresCol="features", predictionCol="prediction", maxIter=20, \ + seed=None, k=4, minDivisibleClusterSize=1.0, distanceMeasure="euclidean", \ + weightCol=None) + Sets params for BisectingKMeans. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setK(self, value: int) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + @since("2.0.0") + def setMinDivisibleClusterSize(self, value: float) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`minDivisibleClusterSize`. + """ + return self._set(minDivisibleClusterSize=value) + + @since("2.4.0") + def setDistanceMeasure(self, value: str) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`distanceMeasure`. + """ + return self._set(distanceMeasure=value) + + @since("2.0.0") + def setMaxIter(self, value: int) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.0.0") + def setFeaturesCol(self, value: str) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("2.0.0") + def setPredictionCol(self, value: str) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("2.0.0") + def setSeed(self, value: int) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "BisectingKMeans": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + def _create_model(self, java_model: "JavaObject") -> BisectingKMeansModel: + return BisectingKMeansModel(java_model) + + +class BisectingKMeansSummary(ClusteringSummary): + """ + Bisecting KMeans clustering results for a given model. + + .. versionadded:: 2.1.0 + """ + + @property + @since("3.0.0") + def trainingCost(self) -> float: + """ + Sum of squared distances to the nearest centroid for all points in the training dataset. + This is equivalent to sklearn's inertia. + """ + return self._call_java("trainingCost") + + +@inherit_doc +class _LDAParams(HasMaxIter, HasFeaturesCol, HasSeed, HasCheckpointInterval): + """ + Params for :py:class:`LDA` and :py:class:`LDAModel`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "The number of topics (clusters) to infer. Must be > 1.", + typeConverter=TypeConverters.toInt, + ) + optimizer: Param[str] = Param( + Params._dummy(), + "optimizer", + "Optimizer or inference algorithm used to estimate the LDA model. " + "Supported: online, em", + typeConverter=TypeConverters.toString, + ) + learningOffset: Param[float] = Param( + Params._dummy(), + "learningOffset", + "A (positive) learning parameter that downweights early iterations." + " Larger values make early iterations count less", + typeConverter=TypeConverters.toFloat, + ) + learningDecay: Param[float] = Param( + Params._dummy(), + "learningDecay", + "Learning rate, set as an" + "exponential decay rate. This should be between (0.5, 1.0] to " + "guarantee asymptotic convergence.", + typeConverter=TypeConverters.toFloat, + ) + subsamplingRate: Param[float] = Param( + Params._dummy(), + "subsamplingRate", + "Fraction of the corpus to be sampled and used in each iteration " + "of mini-batch gradient descent, in range (0, 1].", + typeConverter=TypeConverters.toFloat, + ) + optimizeDocConcentration: Param[bool] = Param( + Params._dummy(), + "optimizeDocConcentration", + "Indicates whether the docConcentration (Dirichlet parameter " + "for document-topic distribution) will be optimized during " + "training.", + typeConverter=TypeConverters.toBoolean, + ) + docConcentration: Param[List[float]] = Param( + Params._dummy(), + "docConcentration", + 'Concentration parameter (commonly named "alpha") for the ' + 'prior placed on documents\' distributions over topics ("theta").', + typeConverter=TypeConverters.toListFloat, + ) + topicConcentration: Param[float] = Param( + Params._dummy(), + "topicConcentration", + 'Concentration parameter (commonly named "beta" or "eta") for ' + "the prior placed on topic' distributions over terms.", + typeConverter=TypeConverters.toFloat, + ) + topicDistributionCol: Param[str] = Param( + Params._dummy(), + "topicDistributionCol", + "Output column with estimates of the topic mixture distribution " + 'for each document (often called "theta" in the literature). ' + "Returns a vector of zeros for an empty document.", + typeConverter=TypeConverters.toString, + ) + keepLastCheckpoint: Param[bool] = Param( + Params._dummy(), + "keepLastCheckpoint", + "(For EM optimizer) If using checkpointing, this indicates whether" + " to keep the last checkpoint. If false, then the checkpoint will be" + " deleted. Deleting the checkpoint can cause failures if a data" + " partition is lost, so set this bit with care.", + TypeConverters.toBoolean, + ) + + def __init__(self, *args: Any): + super(_LDAParams, self).__init__(*args) + self._setDefault( + maxIter=20, + checkpointInterval=10, + k=10, + optimizer="online", + learningOffset=1024.0, + learningDecay=0.51, + subsamplingRate=0.05, + optimizeDocConcentration=True, + topicDistributionCol="topicDistribution", + keepLastCheckpoint=True, + ) + + @since("2.0.0") + def getK(self) -> int: + """ + Gets the value of :py:attr:`k` or its default value. + """ + return self.getOrDefault(self.k) + + @since("2.0.0") + def getOptimizer(self) -> str: + """ + Gets the value of :py:attr:`optimizer` or its default value. + """ + return self.getOrDefault(self.optimizer) + + @since("2.0.0") + def getLearningOffset(self) -> float: + """ + Gets the value of :py:attr:`learningOffset` or its default value. + """ + return self.getOrDefault(self.learningOffset) + + @since("2.0.0") + def getLearningDecay(self) -> float: + """ + Gets the value of :py:attr:`learningDecay` or its default value. + """ + return self.getOrDefault(self.learningDecay) + + @since("2.0.0") + def getSubsamplingRate(self) -> float: + """ + Gets the value of :py:attr:`subsamplingRate` or its default value. + """ + return self.getOrDefault(self.subsamplingRate) + + @since("2.0.0") + def getOptimizeDocConcentration(self) -> bool: + """ + Gets the value of :py:attr:`optimizeDocConcentration` or its default value. + """ + return self.getOrDefault(self.optimizeDocConcentration) + + @since("2.0.0") + def getDocConcentration(self) -> List[float]: + """ + Gets the value of :py:attr:`docConcentration` or its default value. + """ + return self.getOrDefault(self.docConcentration) + + @since("2.0.0") + def getTopicConcentration(self) -> float: + """ + Gets the value of :py:attr:`topicConcentration` or its default value. + """ + return self.getOrDefault(self.topicConcentration) + + @since("2.0.0") + def getTopicDistributionCol(self) -> str: + """ + Gets the value of :py:attr:`topicDistributionCol` or its default value. + """ + return self.getOrDefault(self.topicDistributionCol) + + @since("2.0.0") + def getKeepLastCheckpoint(self) -> bool: + """ + Gets the value of :py:attr:`keepLastCheckpoint` or its default value. + """ + return self.getOrDefault(self.keepLastCheckpoint) + + +@inherit_doc +class LDAModel(JavaModel, _LDAParams): + """ + Latent Dirichlet Allocation (LDA) model. + This abstraction permits for different underlying representations, + including local and distributed data structures. + + .. versionadded:: 2.0.0 + """ + + @since("3.0.0") + def setFeaturesCol(self: "M", value: str) -> "M": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setSeed(self: "M", value: int) -> "M": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setTopicDistributionCol(self: "M", value: str) -> "M": + """ + Sets the value of :py:attr:`topicDistributionCol`. + """ + return self._set(topicDistributionCol=value) + + @since("2.0.0") + def isDistributed(self) -> bool: + """ + Indicates whether this instance is of type DistributedLDAModel + """ + return self._call_java("isDistributed") + + @since("2.0.0") + def vocabSize(self) -> int: + """Vocabulary size (number of terms or words in the vocabulary)""" + return self._call_java("vocabSize") + + @since("2.0.0") + def topicsMatrix(self) -> Matrix: + """ + Inferred topics, where each topic is represented by a distribution over terms. + This is a matrix of size vocabSize x k, where each column is a topic. + No guarantees are given about the ordering of the topics. + + .. warning:: If this model is actually a :py:class:`DistributedLDAModel` + instance produced by the Expectation-Maximization ("em") `optimizer`, + then this method could involve collecting a large amount of data + to the driver (on the order of vocabSize x k). + """ + return self._call_java("topicsMatrix") + + @since("2.0.0") + def logLikelihood(self, dataset: DataFrame) -> float: + """ + Calculates a lower bound on the log likelihood of the entire corpus. + See Equation (16) in the Online LDA paper (Hoffman et al., 2010). + + .. warning:: If this model is an instance of :py:class:`DistributedLDAModel` (produced when + :py:attr:`optimizer` is set to "em"), this involves collecting a large + :py:func:`topicsMatrix` to the driver. This implementation may be changed in the future. + """ + return self._call_java("logLikelihood", dataset) + + @since("2.0.0") + def logPerplexity(self, dataset: DataFrame) -> float: + """ + Calculate an upper bound on perplexity. (Lower is better.) + See Equation (16) in the Online LDA paper (Hoffman et al., 2010). + + .. warning:: If this model is an instance of :py:class:`DistributedLDAModel` (produced when + :py:attr:`optimizer` is set to "em"), this involves collecting a large + :py:func:`topicsMatrix` to the driver. This implementation may be changed in the future. + """ + return self._call_java("logPerplexity", dataset) + + @since("2.0.0") + def describeTopics(self, maxTermsPerTopic: int = 10) -> DataFrame: + """ + Return the topics described by their top-weighted terms. + """ + return self._call_java("describeTopics", maxTermsPerTopic) + + @since("2.0.0") + def estimatedDocConcentration(self) -> Vector: + """ + Value for :py:attr:`LDA.docConcentration` estimated from data. + If Online LDA was used and :py:attr:`LDA.optimizeDocConcentration` was set to false, + then this returns the fixed (given) value for the :py:attr:`LDA.docConcentration` parameter. + """ + return self._call_java("estimatedDocConcentration") + + +@inherit_doc +class DistributedLDAModel(LDAModel, JavaMLReadable["DistributedLDAModel"], JavaMLWritable): + """ + Distributed model fitted by :py:class:`LDA`. + This type of model is currently only produced by Expectation-Maximization (EM). + + This model stores the inferred topics, the full training dataset, and the topic distribution + for each training document. + + .. versionadded:: 2.0.0 + """ + + @since("2.0.0") + def toLocal(self) -> "LocalLDAModel": + """ + Convert this distributed model to a local representation. This discards info about the + training dataset. + + .. warning:: This involves collecting a large :py:func:`topicsMatrix` to the driver. + """ + model = LocalLDAModel(self._call_java("toLocal")) + + # SPARK-10931: Temporary fix to be removed once LDAModel defines Params + model._create_params_from_java() + model._transfer_params_from_java() + + return model + + @since("2.0.0") + def trainingLogLikelihood(self) -> float: + """ + Log likelihood of the observed tokens in the training set, + given the current parameter estimates: + log P(docs | topics, topic distributions for docs, Dirichlet hyperparameters) + + Notes + ----- + - This excludes the prior; for that, use :py:func:`logPrior`. + - Even with :py:func:`logPrior`, this is NOT the same as the data log likelihood given + the hyperparameters. + - This is computed from the topic distributions computed during training. If you call + :py:func:`logLikelihood` on the same training dataset, the topic distributions + will be computed again, possibly giving different results. + """ + return self._call_java("trainingLogLikelihood") + + @since("2.0.0") + def logPrior(self) -> float: + """ + Log probability of the current parameter estimate: + log P(topics, topic distributions for docs | alpha, eta) + """ + return self._call_java("logPrior") + + def getCheckpointFiles(self) -> List[str]: + """ + If using checkpointing and :py:attr:`LDA.keepLastCheckpoint` is set to true, then there may + be saved checkpoint files. This method is provided so that users can manage those files. + + .. versionadded:: 2.0.0 + + Returns + ------- + list + List of checkpoint files from training + + Notes + ----- + Removing the checkpoints can cause failures if a partition is lost and is needed + by certain :py:class:`DistributedLDAModel` methods. Reference counting will clean up + the checkpoints when this model and derivative data go out of scope. + """ + return self._call_java("getCheckpointFiles") + + +@inherit_doc +class LocalLDAModel(LDAModel, JavaMLReadable["LocalLDAModel"], JavaMLWritable): + """ + Local (non-distributed) model fitted by :py:class:`LDA`. + This model stores the inferred topics only; it does not store info about the training dataset. + + .. versionadded:: 2.0.0 + """ + + pass + + +@inherit_doc +class LDA(JavaEstimator[LDAModel], _LDAParams, JavaMLReadable["LDA"], JavaMLWritable): + """ + Latent Dirichlet Allocation (LDA), a topic model designed for text documents. + + Terminology: + + - "term" = "word": an element of the vocabulary + - "token": instance of a term appearing in a document + - "topic": multinomial distribution over terms representing some concept + - "document": one piece of text, corresponding to one row in the input data + + Original LDA paper (journal version): + Blei, Ng, and Jordan. "Latent Dirichlet Allocation." JMLR, 2003. + + Input data (featuresCol): + LDA is given a collection of documents as input data, via the featuresCol parameter. + Each document is specified as a :py:class:`Vector` of length vocabSize, where each entry is the + count for the corresponding term (word) in the document. Feature transformers such as + :py:class:`pyspark.ml.feature.Tokenizer` and :py:class:`pyspark.ml.feature.CountVectorizer` + can be useful for converting text to word count vectors. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors, SparseVector + >>> from pyspark.ml.clustering import LDA + >>> df = spark.createDataFrame([[1, Vectors.dense([0.0, 1.0])], + ... [2, SparseVector(2, {0: 1.0})],], ["id", "features"]) + >>> lda = LDA(k=2, seed=1, optimizer="em") + >>> lda.setMaxIter(10) + LDA... + >>> lda.getMaxIter() + 10 + >>> lda.clear(lda.maxIter) + >>> model = lda.fit(df) + >>> model.setSeed(1) + DistributedLDAModel... + >>> model.getTopicDistributionCol() + 'topicDistribution' + >>> model.isDistributed() + True + >>> localModel = model.toLocal() + >>> localModel.isDistributed() + False + >>> model.vocabSize() + 2 + >>> model.describeTopics().show() + +-----+-----------+--------------------+ + |topic|termIndices| termWeights| + +-----+-----------+--------------------+ + | 0| [1, 0]|[0.50401530077160...| + | 1| [0, 1]|[0.50401530077160...| + +-----+-----------+--------------------+ + ... + >>> model.topicsMatrix() + DenseMatrix(2, 2, [0.496, 0.504, 0.504, 0.496], 0) + >>> lda_path = temp_path + "/lda" + >>> lda.save(lda_path) + >>> sameLDA = LDA.load(lda_path) + >>> distributed_model_path = temp_path + "/lda_distributed_model" + >>> model.save(distributed_model_path) + >>> sameModel = DistributedLDAModel.load(distributed_model_path) + >>> local_model_path = temp_path + "/lda_local_model" + >>> localModel.save(local_model_path) + >>> sameLocalModel = LocalLDAModel.load(local_model_path) + >>> model.transform(df).take(1) == sameLocalModel.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + maxIter: int = 20, + seed: Optional[int] = None, + checkpointInterval: int = 10, + k: int = 10, + optimizer: str = "online", + learningOffset: float = 1024.0, + learningDecay: float = 0.51, + subsamplingRate: float = 0.05, + optimizeDocConcentration: bool = True, + docConcentration: Optional[List[float]] = None, + topicConcentration: Optional[float] = None, + topicDistributionCol: str = "topicDistribution", + keepLastCheckpoint: bool = True, + ): + """ + __init__(self, \\*, featuresCol="features", maxIter=20, seed=None, checkpointInterval=10,\ + k=10, optimizer="online", learningOffset=1024.0, learningDecay=0.51,\ + subsamplingRate=0.05, optimizeDocConcentration=True,\ + docConcentration=None, topicConcentration=None,\ + topicDistributionCol="topicDistribution", keepLastCheckpoint=True) + """ + super(LDA, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.clustering.LDA", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> LDAModel: + if self.getOptimizer() == "em": + return DistributedLDAModel(java_model) + else: + return LocalLDAModel(java_model) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + maxIter: int = 20, + seed: Optional[int] = None, + checkpointInterval: int = 10, + k: int = 10, + optimizer: str = "online", + learningOffset: float = 1024.0, + learningDecay: float = 0.51, + subsamplingRate: float = 0.05, + optimizeDocConcentration: bool = True, + docConcentration: Optional[List[float]] = None, + topicConcentration: Optional[float] = None, + topicDistributionCol: str = "topicDistribution", + keepLastCheckpoint: bool = True, + ) -> "LDA": + """ + setParams(self, \\*, featuresCol="features", maxIter=20, seed=None, checkpointInterval=10,\ + k=10, optimizer="online", learningOffset=1024.0, learningDecay=0.51,\ + subsamplingRate=0.05, optimizeDocConcentration=True,\ + docConcentration=None, topicConcentration=None,\ + topicDistributionCol="topicDistribution", keepLastCheckpoint=True) + + Sets params for LDA. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setCheckpointInterval(self, value: int) -> "LDA": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + @since("2.0.0") + def setSeed(self, value: int) -> "LDA": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("2.0.0") + def setK(self, value: int) -> "LDA": + """ + Sets the value of :py:attr:`k`. + + >>> algo = LDA().setK(10) + >>> algo.getK() + 10 + """ + return self._set(k=value) + + @since("2.0.0") + def setOptimizer(self, value: str) -> "LDA": + """ + Sets the value of :py:attr:`optimizer`. + Currently only support 'em' and 'online'. + + Examples + -------- + >>> algo = LDA().setOptimizer("em") + >>> algo.getOptimizer() + 'em' + """ + return self._set(optimizer=value) + + @since("2.0.0") + def setLearningOffset(self, value: float) -> "LDA": + """ + Sets the value of :py:attr:`learningOffset`. + + Examples + -------- + >>> algo = LDA().setLearningOffset(100) + >>> algo.getLearningOffset() + 100.0 + """ + return self._set(learningOffset=value) + + @since("2.0.0") + def setLearningDecay(self, value: float) -> "LDA": + """ + Sets the value of :py:attr:`learningDecay`. + + Examples + -------- + >>> algo = LDA().setLearningDecay(0.1) + >>> algo.getLearningDecay() + 0.1... + """ + return self._set(learningDecay=value) + + @since("2.0.0") + def setSubsamplingRate(self, value: float) -> "LDA": + """ + Sets the value of :py:attr:`subsamplingRate`. + + Examples + -------- + >>> algo = LDA().setSubsamplingRate(0.1) + >>> algo.getSubsamplingRate() + 0.1... + """ + return self._set(subsamplingRate=value) + + @since("2.0.0") + def setOptimizeDocConcentration(self, value: bool) -> "LDA": + """ + Sets the value of :py:attr:`optimizeDocConcentration`. + + Examples + -------- + >>> algo = LDA().setOptimizeDocConcentration(True) + >>> algo.getOptimizeDocConcentration() + True + """ + return self._set(optimizeDocConcentration=value) + + @since("2.0.0") + def setDocConcentration(self, value: List[float]) -> "LDA": + """ + Sets the value of :py:attr:`docConcentration`. + + Examples + -------- + >>> algo = LDA().setDocConcentration([0.1, 0.2]) + >>> algo.getDocConcentration() + [0.1..., 0.2...] + """ + return self._set(docConcentration=value) + + @since("2.0.0") + def setTopicConcentration(self, value: float) -> "LDA": + """ + Sets the value of :py:attr:`topicConcentration`. + + Examples + -------- + >>> algo = LDA().setTopicConcentration(0.5) + >>> algo.getTopicConcentration() + 0.5... + """ + return self._set(topicConcentration=value) + + @since("2.0.0") + def setTopicDistributionCol(self, value: str) -> "LDA": + """ + Sets the value of :py:attr:`topicDistributionCol`. + + Examples + -------- + >>> algo = LDA().setTopicDistributionCol("topicDistributionCol") + >>> algo.getTopicDistributionCol() + 'topicDistributionCol' + """ + return self._set(topicDistributionCol=value) + + @since("2.0.0") + def setKeepLastCheckpoint(self, value: bool) -> "LDA": + """ + Sets the value of :py:attr:`keepLastCheckpoint`. + + Examples + -------- + >>> algo = LDA().setKeepLastCheckpoint(False) + >>> algo.getKeepLastCheckpoint() + False + """ + return self._set(keepLastCheckpoint=value) + + @since("2.0.0") + def setMaxIter(self, value: int) -> "LDA": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.0.0") + def setFeaturesCol(self, value: str) -> "LDA": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + +@inherit_doc +class _PowerIterationClusteringParams(HasMaxIter, HasWeightCol): + """ + Params for :py:class:`PowerIterationClustering`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "The number of clusters to create. Must be > 1.", + typeConverter=TypeConverters.toInt, + ) + initMode: Param[str] = Param( + Params._dummy(), + "initMode", + "The initialization algorithm. This can be either " + + "'random' to use a random vector as vertex properties, or 'degree' to use " + + "a normalized sum of similarities with other vertices. Supported options: " + + "'random' and 'degree'.", + typeConverter=TypeConverters.toString, + ) + srcCol: Param[str] = Param( + Params._dummy(), + "srcCol", + "Name of the input column for source vertex IDs.", + typeConverter=TypeConverters.toString, + ) + dstCol: Param[str] = Param( + Params._dummy(), + "dstCol", + "Name of the input column for destination vertex IDs.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_PowerIterationClusteringParams, self).__init__(*args) + self._setDefault(k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst") + + @since("2.4.0") + def getK(self) -> int: + """ + Gets the value of :py:attr:`k` or its default value. + """ + return self.getOrDefault(self.k) + + @since("2.4.0") + def getInitMode(self) -> str: + """ + Gets the value of :py:attr:`initMode` or its default value. + """ + return self.getOrDefault(self.initMode) + + @since("2.4.0") + def getSrcCol(self) -> str: + """ + Gets the value of :py:attr:`srcCol` or its default value. + """ + return self.getOrDefault(self.srcCol) + + @since("2.4.0") + def getDstCol(self) -> str: + """ + Gets the value of :py:attr:`dstCol` or its default value. + """ + return self.getOrDefault(self.dstCol) + + +@inherit_doc +class PowerIterationClustering( + _PowerIterationClusteringParams, + JavaParams, + JavaMLReadable["PowerIterationClustering"], + JavaMLWritable, +): + """ + Power Iteration Clustering (PIC), a scalable graph clustering algorithm developed by + `Lin and Cohen `_. From the + abstract: PIC finds a very low-dimensional embedding of a dataset using truncated power + iteration on a normalized pair-wise similarity matrix of the data. + + This class is not yet an Estimator/Transformer, use :py:func:`assignClusters` method + to run the PowerIterationClustering algorithm. + + .. versionadded:: 2.4.0 + + Notes + ----- + See `Wikipedia on Spectral clustering `_ + + Examples + -------- + >>> data = [(1, 0, 0.5), + ... (2, 0, 0.5), (2, 1, 0.7), + ... (3, 0, 0.5), (3, 1, 0.7), (3, 2, 0.9), + ... (4, 0, 0.5), (4, 1, 0.7), (4, 2, 0.9), (4, 3, 1.1), + ... (5, 0, 0.5), (5, 1, 0.7), (5, 2, 0.9), (5, 3, 1.1), (5, 4, 1.3)] + >>> df = spark.createDataFrame(data).toDF("src", "dst", "weight").repartition(1) + >>> pic = PowerIterationClustering(k=2, weightCol="weight") + >>> pic.setMaxIter(40) + PowerIterationClustering... + >>> assignments = pic.assignClusters(df) + >>> assignments.sort(assignments.id).show(truncate=False) + +---+-------+ + |id |cluster| + +---+-------+ + |0 |0 | + |1 |0 | + |2 |0 | + |3 |0 | + |4 |0 | + |5 |1 | + +---+-------+ + ... + >>> pic_path = temp_path + "/pic" + >>> pic.save(pic_path) + >>> pic2 = PowerIterationClustering.load(pic_path) + >>> pic2.getK() + 2 + >>> pic2.getMaxIter() + 40 + >>> pic2.assignClusters(df).take(6) == assignments.take(6) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + k: int = 2, + maxIter: int = 20, + initMode: str = "random", + srcCol: str = "src", + dstCol: str = "dst", + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst",\ + weightCol=None) + """ + super(PowerIterationClustering, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.clustering.PowerIterationClustering", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.4.0") + def setParams( + self, + *, + k: int = 2, + maxIter: int = 20, + initMode: str = "random", + srcCol: str = "src", + dstCol: str = "dst", + weightCol: Optional[str] = None, + ) -> "PowerIterationClustering": + """ + setParams(self, \\*, k=2, maxIter=20, initMode="random", srcCol="src", dstCol="dst",\ + weightCol=None) + Sets params for PowerIterationClustering. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.4.0") + def setK(self, value: int) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + @since("2.4.0") + def setInitMode(self, value: str) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`initMode`. + """ + return self._set(initMode=value) + + @since("2.4.0") + def setSrcCol(self, value: str) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`srcCol`. + """ + return self._set(srcCol=value) + + @since("2.4.0") + def setDstCol(self, value: str) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`dstCol`. + """ + return self._set(dstCol=value) + + @since("2.4.0") + def setMaxIter(self, value: int) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.4.0") + def setWeightCol(self, value: str) -> "PowerIterationClustering": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("2.4.0") + def assignClusters(self, dataset: DataFrame) -> DataFrame: + """ + Run the PIC algorithm and returns a cluster assignment for each input vertex. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + A dataset with columns src, dst, weight representing the affinity matrix, + which is the matrix A in the PIC paper. Suppose the src column value is i, + the dst column value is j, the weight column value is similarity s,,ij,, + which must be nonnegative. This is a symmetric matrix and hence + s,,ij,, = s,,ji,,. For any (i, j) with nonzero similarity, there should be + either (i, j, s,,ij,,) or (j, i, s,,ji,,) in the input. Rows with i = j are + ignored, because we assume s,,ij,, = 0.0. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + A dataset that contains columns of vertex id and the corresponding cluster for + the id. The schema of it will be: + - id: Long + - cluster: Int + """ + self._transfer_params_to_java() + assert self._java_obj is not None + + jdf = self._java_obj.assignClusters(dataset._jdf) + return DataFrame(jdf, dataset.sparkSession) + + +if __name__ == "__main__": + import doctest + import numpy + import pyspark.ml.clustering + from pyspark.sql import SparkSession + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + globs = pyspark.ml.clustering.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.clustering tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/common.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/common.py new file mode 100644 index 0000000..dd6fee4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/common.py @@ -0,0 +1,146 @@ +# +# 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 Any, Callable, TYPE_CHECKING + +if TYPE_CHECKING: + from pyspark.ml._typing import C, JavaObjectOrPickleDump + +import py4j.protocol +from py4j.protocol import Py4JJavaError +from py4j.java_gateway import JavaObject +from py4j.java_collections import JavaArray, JavaList + +import pyspark.context +from pyspark import RDD, SparkContext +from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer +from pyspark.sql import DataFrame, SparkSession + +# Hack for support float('inf') in Py4j +_old_smart_decode = py4j.protocol.smart_decode + +_float_str_mapping = { + "nan": "NaN", + "inf": "Infinity", + "-inf": "-Infinity", +} + + +def _new_smart_decode(obj: Any) -> str: + if isinstance(obj, float): + s = str(obj) + return _float_str_mapping.get(s, s) + return _old_smart_decode(obj) + + +py4j.protocol.smart_decode = _new_smart_decode + + +_picklable_classes = [ + "SparseVector", + "DenseVector", + "SparseMatrix", + "DenseMatrix", +] + + +# this will call the ML version of pythonToJava() +def _to_java_object_rdd(rdd: RDD) -> JavaObject: + """Return an JavaRDD of Object by unpickling + + It will convert each Python object into Java object by Pickle, whenever the + RDD is serialized in batch or not. + """ + rdd = rdd._reserialize(AutoBatchedSerializer(CPickleSerializer())) + assert rdd.ctx._jvm is not None + return rdd.ctx._jvm.org.apache.spark.ml.python.MLSerDe.pythonToJava(rdd._jrdd, True) + + +def _py2java(sc: SparkContext, obj: Any) -> JavaObject: + """Convert Python object into Java""" + if isinstance(obj, RDD): + obj = _to_java_object_rdd(obj) + elif isinstance(obj, DataFrame): + obj = obj._jdf + elif isinstance(obj, SparkContext): + obj = obj._jsc + elif isinstance(obj, list): + obj = [_py2java(sc, x) for x in obj] + elif isinstance(obj, JavaObject): + pass + elif isinstance(obj, (int, float, bool, bytes, str)): + pass + else: + data = bytearray(CPickleSerializer().dumps(obj)) + assert sc._jvm is not None + obj = sc._jvm.org.apache.spark.ml.python.MLSerDe.loads(data) + return obj + + +def _java2py(sc: SparkContext, r: "JavaObjectOrPickleDump", encoding: str = "bytes") -> Any: + if isinstance(r, JavaObject): + clsName = r.getClass().getSimpleName() + # convert RDD into JavaRDD + if clsName != "JavaRDD" and clsName.endswith("RDD"): + r = r.toJavaRDD() + clsName = "JavaRDD" + + assert sc._jvm is not None + + if clsName == "JavaRDD": + jrdd = sc._jvm.org.apache.spark.ml.python.MLSerDe.javaToPython(r) + return RDD(jrdd, sc) + + if clsName == "Dataset": + return DataFrame(r, SparkSession._getActiveSessionOrCreate()) + + if clsName in _picklable_classes: + r = sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(r) + elif isinstance(r, (JavaArray, JavaList)): + try: + r = sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(r) + except Py4JJavaError: + pass # not picklable + + if isinstance(r, (bytearray, bytes)): + r = CPickleSerializer().loads(bytes(r), encoding=encoding) + return r + + +def callJavaFunc( + sc: pyspark.context.SparkContext, func: Callable[..., "JavaObjectOrPickleDump"], *args: Any +) -> "JavaObjectOrPickleDump": + """Call Java Function""" + java_args = [_py2java(sc, a) for a in args] + return _java2py(sc, func(*java_args)) + + +def inherit_doc(cls: "C") -> "C": + """ + A decorator that makes a class inherit documentation from its parents. + """ + for name, func in vars(cls).items(): + # only inherit docstring for public functions + if name.startswith("_"): + continue + if not func.__doc__: + for parent in cls.__bases__: + parent_func = getattr(parent, name, None) + if parent_func and getattr(parent_func, "__doc__", None): + func.__doc__ = parent_func.__doc__ + break + return cls diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/__init__.py new file mode 100644 index 0000000..2ee152f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/__init__.py @@ -0,0 +1,44 @@ +# +# 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. +# + +"""Spark Connect Python Client - ML module""" + +from pyspark.ml.connect.base import ( + Estimator, + Transformer, + Model, +) + +from pyspark.ml.connect import ( + feature, + evaluation, + tuning, +) + +from pyspark.ml.connect.pipeline import Pipeline, PipelineModel + +__all__ = [ + "Estimator", + "Transformer", + "Estimator", + "Model", + "feature", + "evaluation", + "Pipeline", + "PipelineModel", + "tuning", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/base.py new file mode 100644 index 0000000..f8ce0cb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/base.py @@ -0,0 +1,346 @@ +# +# 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 + +import pandas as pd + +from typing import ( + Any, + Generic, + List, + Optional, + TypeVar, + Union, + TYPE_CHECKING, + Tuple, + Callable, +) + +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, +) +from pyspark.ml.connect.util import transform_dataframe_column + +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 + """ + + @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 + + 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 + """ + + 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[["pd.Series"], Any]: + """ + Return a transformation function that accepts an instance of `pd.Series` as input 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 input pandas dataframe is appended with output + columns in place. + + 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]: + 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 + """ + + @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 + + 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 + """ + + pass + + +@inherit_doc +class _PredictorParams(HasLabelCol, HasFeaturesCol, HasPredictionCol): + """ + Params for :py:class:`Predictor` and :py:class:`PredictorModel`. + + .. versionadded:: 3.5.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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/classification.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/classification.py new file mode 100644 index 0000000..f8b525d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/classification.py @@ -0,0 +1,382 @@ +# +# 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 import keyword_only +from pyspark.ml.connect.base import _PredictorParams + +from pyspark.ml.param.shared import HasProbabilityCol + +from typing import Any, Dict, Union, List, Tuple, Callable, Optional +import numpy as np +import pandas as pd +import math + +from pyspark.sql import DataFrame +from pyspark.ml.common import inherit_doc +from pyspark.ml.torch.distributor import TorchDistributor +from pyspark.ml.param.shared import ( + HasMaxIter, + HasFitIntercept, + HasTol, + HasWeightCol, + HasSeed, + HasNumTrainWorkers, + HasBatchSize, + HasLearningRate, + HasMomentum, +) +from pyspark.ml.connect.base import Predictor, PredictionModel +from pyspark.ml.connect.io_utils import ParamsReadWrite, CoreModelReadWrite +from pyspark.sql.functions import lit, count, countDistinct + +import torch +import torch.nn as torch_nn + + +class _LogisticRegressionParams( + _PredictorParams, + HasMaxIter, + HasFitIntercept, + HasTol, + HasWeightCol, + HasNumTrainWorkers, + HasBatchSize, + HasLearningRate, + HasMomentum, + HasProbabilityCol, + HasSeed, +): + """ + Params for :py:class:`LogisticRegression` and :py:class:`LogisticRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + def __init__(self, *args: Any): + super(_LogisticRegressionParams, self).__init__(*args) + self._setDefault( + maxIter=100, + tol=1e-6, + batchSize=32, + learningRate=0.001, + momentum=0.9, + seed=0, + ) + + +def _train_logistic_regression_model_worker_fn( + num_samples_per_worker: int, + num_features: int, + batch_size: int, + max_iter: int, + num_classes: int, + learning_rate: float, + momentum: float, + fit_intercept: bool, + seed: int, +) -> Any: + from pyspark.ml.torch.distributor import _get_spark_partition_data_loader + from torch.nn.parallel import DistributedDataParallel as DDP + import torch.distributed + import torch.optim as optim + + # TODO: add a setting seed param. + torch.manual_seed(seed) + + # TODO: support training on GPU + # TODO: support L1 / L2 regularization + torch.distributed.init_process_group("gloo") + + linear_model = torch_nn.Linear( + num_features, num_classes, bias=fit_intercept, dtype=torch.float32 + ) + ddp_model = DDP(linear_model) + + loss_fn = torch_nn.CrossEntropyLoss() + + optimizer = optim.SGD(ddp_model.parameters(), lr=learning_rate, momentum=momentum) + data_loader = _get_spark_partition_data_loader( + num_samples_per_worker, + batch_size, + num_workers=0, + prefetch_factor=None, # type: ignore + ) + for i in range(max_iter): + ddp_model.train() + + step_count = 0 + + loss_sum = 0.0 + for x, target in data_loader: + optimizer.zero_grad() + output = ddp_model(x.to(torch.float32)) + loss = loss_fn(output, target.to(torch.long)) + loss.backward() + loss_sum += loss.detach().numpy() + optimizer.step() + step_count += 1 + + # TODO: early stopping + # When each epoch ends, computes loss on validation dataset and compare + # current epoch validation loss with last epoch validation loss, if + # less than provided `tol`, stop training. + + if torch.distributed.get_rank() == 0: + print(f"Progress: train epoch {i + 1} completes, train loss = {loss_sum / step_count}") + + if torch.distributed.get_rank() == 0: + return ddp_model.module.state_dict() + + return None + + +@inherit_doc +class LogisticRegression( + Predictor["LogisticRegressionModel"], _LogisticRegressionParams, ParamsReadWrite +): + """ + Logistic regression estimator. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.classification import LogisticRegression, LogisticRegressionModel + >>> lor = LogisticRegression(maxIter=20, learningRate=0.01) + >>> dataset = spark.createDataFrame([ + ... ([1.0, 2.0], 1), + ... ([2.0, -1.0], 1), + ... ([-3.0, -2.0], 0), + ... ([-1.0, -2.0], 0), + ... ], schema=['features', 'label']) + >>> lor_model = lor.fit(dataset) + >>> transformed_dataset = lor_model.transform(dataset) + >>> transformed_dataset.show() + +------------+-----+----------+--------------------+ + | features|label|prediction| probability| + +------------+-----+----------+--------------------+ + | [1.0, 2.0]| 1| 1|[0.02423273026943...| + | [2.0, -1.0]| 1| 1|[0.09334788471460...| + |[-3.0, -2.0]| 0| 0|[0.99808156490325...| + |[-1.0, -2.0]| 0| 0|[0.96210002899169...| + +------------+-----+----------+--------------------+ + >>> lor_model.saveToLocal("/tmp/lor_model") + >>> LogisticRegressionModel.loadFromLocal("/tmp/lor_model") + LogisticRegression_... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + maxIter: int = 100, + tol: float = 1e-6, + numTrainWorkers: int = 1, + batchSize: int = 32, + learningRate: float = 0.001, + momentum: float = 0.9, + seed: int = 0, + ): + """ + __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + probabilityCol: str = "probability", + maxIter: int = 100, + tol: float = 1e-6, + numTrainWorkers: int = 1, + batchSize: int = 32, + learningRate: float = 0.001, + momentum: float = 0.9, + seed: int = 0, + ) + """ + super(LogisticRegression, self).__init__() + kwargs = self._input_kwargs + self._set(**kwargs) + + def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "LogisticRegressionModel": + if isinstance(dataset, pd.DataFrame): + # TODO: support pandas dataframe fitting + raise NotImplementedError("Fitting pandas dataframe is not supported yet.") + + num_train_workers = self.getNumTrainWorkers() + batch_size = self.getBatchSize() + + # We don't need to persist the dataset because the shuffling result from the repartition + # has been cached. + dataset = dataset.select(self.getFeaturesCol(), self.getLabelCol()).repartition( + num_train_workers + ) + + # TODO: check label values are in range of [0, num_classes) + num_rows, num_classes = dataset.agg( + count(lit(1)), countDistinct(self.getLabelCol()) + ).head() # type: ignore[misc] + + num_batches_per_worker = math.ceil(num_rows / num_train_workers / batch_size) + num_samples_per_worker = num_batches_per_worker * batch_size + + num_features = len(dataset.select(self.getFeaturesCol()).head()[0]) # type: ignore[index] + + if num_classes < 2: + raise ValueError("Training dataset distinct labels must >= 2.") + + # TODO: support GPU. + distributor = TorchDistributor( + local_mode=False, use_gpu=False, num_processes=num_train_workers + ) + model_state_dict = distributor._train_on_dataframe( + _train_logistic_regression_model_worker_fn, + dataset, + num_samples_per_worker=num_samples_per_worker, + num_features=num_features, + batch_size=batch_size, + max_iter=self.getMaxIter(), + num_classes=num_classes, + learning_rate=self.getLearningRate(), + momentum=self.getMomentum(), + fit_intercept=self.getFitIntercept(), + seed=self.getSeed(), + ) + + dataset.unpersist() + + torch_model = torch_nn.Linear( + num_features, num_classes, bias=self.getFitIntercept(), dtype=torch.float32 + ) + torch_model.load_state_dict(model_state_dict) + + lor_model = LogisticRegressionModel( + torch_model, num_features=num_features, num_classes=num_classes + ) + lor_model._resetUid(self.uid) + return self._copyValues(lor_model) + + +@inherit_doc +class LogisticRegressionModel( + PredictionModel, _LogisticRegressionParams, ParamsReadWrite, CoreModelReadWrite +): + """ + Model fitted by LogisticRegression. + + .. versionadded:: 3.5.0 + """ + + def __init__( + self, + torch_model: Any = None, + num_features: Optional[int] = None, + num_classes: Optional[int] = None, + ): + super().__init__() + self.torch_model = torch_model + self.num_features = num_features + self.num_classes = num_classes + + @property + def numFeatures(self) -> int: + return self.num_features # type: ignore[return-value] + + @property + def numClasses(self) -> int: + return self.num_classes # type: ignore[return-value] + + def _input_columns(self) -> List[str]: + return [self.getOrDefault(self.featuresCol)] + + def _output_columns(self) -> List[Tuple[str, str]]: + output_cols = [(self.getOrDefault(self.predictionCol), "bigint")] + prob_col = self.getOrDefault(self.probabilityCol) + if prob_col: + output_cols += [(prob_col, "array")] + return output_cols + + def _get_transform_fn(self) -> Callable[["pd.Series"], Any]: + model_state_dict = self.torch_model.state_dict() + num_features = self.num_features + num_classes = self.num_classes + fit_intercept = self.getFitIntercept() + + def transform_fn(input_series: Any) -> Any: + torch_model = torch_nn.Linear( + num_features, # type: ignore[arg-type] + num_classes, # type: ignore[arg-type] + bias=fit_intercept, + dtype=torch.float32, + ) + # TODO: Use spark broadast for `model_state_dict`, + # it can improve performance when model is large. + torch_model.load_state_dict(model_state_dict) + + input_array = np.stack(input_series.values) + + with torch.inference_mode(): + result = torch_model(torch.tensor(input_array, dtype=torch.float32)) + predictions = torch.argmax(result, dim=1).numpy() + + if self.getProbabilityCol(): + probabilities = torch.softmax(result, dim=1).numpy() + + return pd.DataFrame( + { + self.getPredictionCol(): list(predictions), + self.getProbabilityCol(): list(probabilities), + }, + index=input_series.index.copy(), + ) + else: + return pd.Series(data=list(predictions), index=input_series.index.copy()) + + return transform_fn + + def _get_core_model_filename(self) -> str: + return self.__class__.__name__ + ".torch" + + def _save_core_model(self, path: str) -> None: + lor_torch_model = torch_nn.Sequential( + self.torch_model, + torch_nn.Softmax(dim=1), + ) + torch.save(lor_torch_model, path) + + def _load_core_model(self, path: str) -> None: + lor_torch_model = torch.load(path) + self.torch_model = lor_torch_model[0] + + def _get_extra_metadata(self) -> Dict[str, Any]: + return { + "num_features": self.num_features, + "num_classes": self.num_classes, + } + + def _load_extra_metadata(self, extra_metadata: Dict[str, Any]) -> None: + """ + Load extra metadata attribute from extra metadata json object. + """ + self.num_features = extra_metadata["num_features"] + self.num_classes = extra_metadata["num_classes"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/evaluation.py new file mode 100644 index 0000000..8124f10 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/evaluation.py @@ -0,0 +1,291 @@ +# +# 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 numpy as np +import pandas as pd +from typing import Any, Union, List, Tuple + +from pyspark import keyword_only +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import HasLabelCol, HasPredictionCol, HasProbabilityCol +from pyspark.ml.connect.base import Evaluator +from pyspark.ml.connect.io_utils import ParamsReadWrite +from pyspark.ml.connect.util import aggregate_dataframe +from pyspark.sql import DataFrame + + +class _TorchMetricEvaluator(Evaluator): + + metricName: Param[str] = Param( + Params._dummy(), + "metricName", + "metric name for the regression evaluator, valid values are 'mse' and 'r2'", + typeConverter=TypeConverters.toString, + ) + + def getMetricName(self) -> str: + """ + Gets the value of metricName or its default value. + + .. versionadded:: 3.5.0 + """ + return self.getOrDefault(self.metricName) + + def _get_torch_metric(self) -> Any: + raise NotImplementedError() + + def _get_input_cols(self) -> List[str]: + raise NotImplementedError() + + def _get_metric_update_inputs(self, dataset: "pd.DataFrame") -> Tuple[Any, Any]: + raise NotImplementedError() + + def _evaluate(self, dataset: Union["DataFrame", "pd.DataFrame"]) -> float: + torch_metric = self._get_torch_metric() + + def local_agg_fn(pandas_df: "pd.DataFrame") -> "pd.DataFrame": + torch_metric.update(*self._get_metric_update_inputs(pandas_df)) + return torch_metric + + def merge_agg_state(state1: Any, state2: Any) -> Any: + state1.merge_state([state2]) + return state1 + + def agg_state_to_result(state: Any) -> Any: + return state.compute().item() + + return aggregate_dataframe( + dataset, + self._get_input_cols(), + local_agg_fn, + merge_agg_state, + agg_state_to_result, + ) + + +def _get_rmse_torchmetric() -> Any: + import torch + import torcheval.metrics as torchmetrics + + class _RootMeanSquaredError(torchmetrics.MeanSquaredError): + def compute(self: Any) -> torch.Tensor: + return torch.sqrt(super().compute()) + + return _RootMeanSquaredError() + + +class RegressionEvaluator(_TorchMetricEvaluator, HasLabelCol, HasPredictionCol, ParamsReadWrite): + """ + Evaluator for Regression, which expects input columns prediction and label. + Supported metrics are 'rmse', 'mse' and 'r2'. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.evaluation import RegressionEvaluator + >>> eva = RegressionEvaluator(metricName='mse') + >>> dataset = spark.createDataFrame( + ... [(1.0, 2.0), (-1.0, -1.5)], schema=['label', 'prediction'] + ... ) + >>> eva.evaluate(dataset) + 0.625 + >>> eva.isLargerBetter() + False + """ + + @keyword_only + def __init__( + self, + *, + metricName: str = "rmse", + labelCol: str = "label", + predictionCol: str = "prediction", + ) -> None: + """ + __init__(self, *, metricName='rmse', labelCol='label', predictionCol='prediction') -> None: + """ + super().__init__() + self._set(metricName=metricName, labelCol=labelCol, predictionCol=predictionCol) + + def _get_torch_metric(self) -> Any: + import torcheval.metrics as torchmetrics + + metric_name = self.getOrDefault(self.metricName) + + if metric_name == "mse": + return torchmetrics.MeanSquaredError() + if metric_name == "r2": + return torchmetrics.R2Score() + if metric_name == "rmse": + return _get_rmse_torchmetric() + + raise ValueError(f"Unsupported regressor evaluator metric name: {metric_name}") + + def _get_input_cols(self) -> List[str]: + return [self.getPredictionCol(), self.getLabelCol()] + + def _get_metric_update_inputs(self, dataset: "pd.DataFrame") -> Tuple[Any, Any]: + import torch + + preds_tensor = torch.tensor(dataset[self.getPredictionCol()].values) + labels_tensor = torch.tensor(dataset[self.getLabelCol()].values) + return preds_tensor, labels_tensor + + def isLargerBetter(self) -> bool: + if self.getOrDefault(self.metricName) == "r2": + return True + + return False + + +class BinaryClassificationEvaluator( + _TorchMetricEvaluator, HasLabelCol, HasProbabilityCol, ParamsReadWrite +): + """ + Evaluator for binary classification, which expects input columns prediction and label. + Supported metrics are 'areaUnderROC' and 'areaUnderPR'. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.evaluation import BinaryClassificationEvaluator + >>> eva = BinaryClassificationEvaluator(metricName='areaUnderPR') + >>> dataset = spark.createDataFrame( + ... [(1, 0.6), (0, 0.55), (0, 0.1), (1, 0.6), (1, 0.4)], + ... schema=['label', 'probability'] + ... ) + >>> eva.evaluate(dataset) + 0.9166666865348816 + >>> eva.isLargerBetter() + True + """ + + @keyword_only + def __init__( + self, + *, + metricName: str = "areaUnderROC", + labelCol: str = "label", + probabilityCol: str = "probability", + ) -> None: + """ + __init__( + self, + *, + metricName='rmse', + labelCol='label', + probabilityCol='probability' + ) -> None: + """ + super().__init__() + self._set(metricName=metricName, labelCol=labelCol, probabilityCol=probabilityCol) + + def _get_torch_metric(self) -> Any: + import torcheval.metrics as torchmetrics + + metric_name = self.getOrDefault(self.metricName) + + if metric_name == "areaUnderROC": + return torchmetrics.BinaryAUROC() + if metric_name == "areaUnderPR": + return torchmetrics.BinaryAUPRC() + + raise ValueError(f"Unsupported binary classification evaluator metric name: {metric_name}") + + def _get_input_cols(self) -> List[str]: + return [self.getProbabilityCol(), self.getLabelCol()] + + def _get_metric_update_inputs(self, dataset: "pd.DataFrame") -> Tuple[Any, Any]: + import torch + + values = np.stack(dataset[self.getProbabilityCol()].values) # type: ignore[call-overload] + preds_tensor = torch.tensor(values) + if preds_tensor.dim() == 2: + preds_tensor = preds_tensor[:, 1] + labels_tensor = torch.tensor(dataset[self.getLabelCol()].values) + return preds_tensor, labels_tensor + + def isLargerBetter(self) -> bool: + return True + + +class MulticlassClassificationEvaluator( + _TorchMetricEvaluator, HasLabelCol, HasPredictionCol, ParamsReadWrite +): + """ + Evaluator for multiclass classification, which expects input columns prediction and label. + Supported metrics are 'accuracy'. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.evaluation import MulticlassClassificationEvaluator + >>> eva = MulticlassClassificationEvaluator(metricName='accuracy') + >>> dataset = spark.createDataFrame( + ... [(1, 1), (0, 0), (2, 2), (1, 0), (2, 1)], + ... schema=['label', 'prediction'] + ... ) + >>> eva.evaluate(dataset) + 0.6000000238418579 + >>> eva.isLargerBetter() + True + """ + + def __init__( + self, + metricName: str = "accuracy", + labelCol: str = "label", + predictionCol: str = "prediction", + ) -> None: + """ + __init__( + self, + *, + metricName='accuracy', + labelCol='label', + predictionCol='prediction' + ) -> None: + """ + super().__init__() + self._set(metricName=metricName, labelCol=labelCol, predictionCol=predictionCol) + + def _get_torch_metric(self) -> Any: + import torcheval.metrics as torchmetrics + + metric_name = self.getOrDefault(self.metricName) + + if metric_name == "accuracy": + return torchmetrics.MulticlassAccuracy() + + raise ValueError( + f"Unsupported multiclass classification evaluator metric name: {metric_name}" + ) + + def _get_input_cols(self) -> List[str]: + return [self.getPredictionCol(), self.getLabelCol()] + + def _get_metric_update_inputs(self, dataset: "pd.DataFrame") -> Tuple[Any, Any]: + import torch + + preds_tensor = torch.tensor(dataset[self.getPredictionCol()].values) + labels_tensor = torch.tensor(dataset[self.getLabelCol()].values) + return preds_tensor, labels_tensor + + def isLargerBetter(self) -> bool: + return True diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/feature.py new file mode 100644 index 0000000..42b4702 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/feature.py @@ -0,0 +1,258 @@ +# +# 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 numpy as np +import pandas as pd +import pickle +from typing import Any, Union, List, Tuple, Callable, Dict, Optional + +from pyspark import keyword_only +from pyspark.sql import DataFrame +from pyspark.ml.param.shared import HasInputCol, HasOutputCol +from pyspark.ml.connect.base import Estimator, Model +from pyspark.ml.connect.io_utils import ParamsReadWrite, CoreModelReadWrite +from pyspark.ml.connect.summarizer import summarize_dataframe + + +class MaxAbsScaler(Estimator, HasInputCol, HasOutputCol, ParamsReadWrite): + """ + Rescale each feature individually to range [-1, 1] by dividing through the largest maximum + absolute value in each feature. It does not shift/center the data, and thus does not destroy + any sparsity. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.feature import MaxAbsScaler + >>> scaler = MaxAbsScaler(inputCol='features', outputCol='scaled_features') + >>> dataset = spark.createDataFrame([ + ... ([1.0, 2.0],), + ... ([2.0, -1.0],), + ... ([-3.0, -2.0],), + ... ], schema=['features']) + >>> scaler_model = scaler.fit(dataset) + >>> transformed_dataset = scaler_model.transform(dataset) + >>> transformed_dataset.show(truncate=False) + +------------+--------------------------+ + |features |scaled_features | + +------------+--------------------------+ + |[1.0, 2.0] |[0.3333333333333333, 1.0] | + |[2.0, -1.0] |[0.6666666666666666, -0.5]| + |[-3.0, -2.0]|[-1.0, -1.0] | + +------------+--------------------------+ + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, inputCol: Optional[str] = None, outputCol: Optional[str] = None) -> None: + """ + __init__(self, \\*, inputCol=None, outputCol=None) + """ + super().__init__() + kwargs = self._input_kwargs + self._set(**kwargs) + + def _fit(self, dataset: Union["pd.DataFrame", "DataFrame"]) -> "MaxAbsScalerModel": + input_col = self.getInputCol() + + stat_res = summarize_dataframe(dataset, input_col, ["min", "max", "count"]) + min_values = stat_res["min"] + max_values = stat_res["max"] + n_samples_seen = stat_res["count"] + + max_abs_values = np.maximum(np.abs(min_values), np.abs(max_values)) + + model = MaxAbsScalerModel(max_abs_values, n_samples_seen) + model._resetUid(self.uid) + return self._copyValues(model) + + +class MaxAbsScalerModel(Model, HasInputCol, HasOutputCol, ParamsReadWrite, CoreModelReadWrite): + """ + Model fitted by MaxAbsScaler. + + .. versionadded:: 3.5.0 + """ + + def __init__( + self, max_abs_values: Optional["np.ndarray"] = None, n_samples_seen: Optional[int] = None + ) -> None: + super().__init__() + self.max_abs_values = max_abs_values + if max_abs_values is not None: + # if scale value is zero, replace it with 1.0 (for preventing division by zero) + self.scale_values = np.where(max_abs_values == 0.0, 1.0, max_abs_values) + self.n_samples_seen = n_samples_seen + + def _input_columns(self) -> List[str]: + return [self.getInputCol()] + + def _output_columns(self) -> List[Tuple[str, str]]: + return [(self.getOutputCol(), "array")] + + def _get_transform_fn(self) -> Callable[..., Any]: + scale_values = self.scale_values + + def transform_fn(series: Any) -> Any: + def map_value(x: "np.ndarray") -> "np.ndarray": + return x / scale_values + + return series.apply(map_value) + + return transform_fn + + def _get_core_model_filename(self) -> str: + return self.__class__.__name__ + ".sklearn.pkl" + + def _save_core_model(self, path: str) -> None: + from sklearn.preprocessing import MaxAbsScaler as sk_MaxAbsScaler + + sk_model = sk_MaxAbsScaler() + sk_model.scale_ = self.scale_values + sk_model.max_abs_ = self.max_abs_values + sk_model.n_features_in_ = len(self.max_abs_values) # type: ignore[arg-type] + sk_model.n_samples_seen_ = self.n_samples_seen + + with open(path, "wb") as fp: + pickle.dump(sk_model, fp) + + def _load_core_model(self, path: str) -> None: + with open(path, "rb") as fp: + sk_model = pickle.load(fp) + + self.max_abs_values = sk_model.max_abs_ + self.scale_values = sk_model.scale_ + self.n_samples_seen = sk_model.n_samples_seen_ + + +class StandardScaler(Estimator, HasInputCol, HasOutputCol, ParamsReadWrite): + """ + Standardizes features by removing the mean and scaling to unit variance using column summary + statistics on the samples in the training set. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.feature import StandardScaler + >>> scaler = StandardScaler(inputCol='features', outputCol='scaled_features') + >>> dataset = spark.createDataFrame([ + ... ([1.0, 2.0],), + ... ([2.0, -1.0],), + ... ([-3.0, -2.0],), + ... ], schema=['features']) + >>> scaler_model = scaler.fit(dataset) + >>> transformed_dataset = scaler_model.transform(dataset) + >>> transformed_dataset.show(truncate=False) + +------------+------------------------------------------+ + |features |scaled_features | + +------------+------------------------------------------+ + |[1.0, 2.0] |[0.3779644730092272, 1.1208970766356101] | + |[2.0, -1.0] |[0.7559289460184544, -0.3202563076101743] | + |[-3.0, -2.0]|[-1.1338934190276817, -0.8006407690254358]| + +------------+------------------------------------------+ + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, inputCol: Optional[str] = None, outputCol: Optional[str] = None) -> None: + """ + __init__(self, \\*, inputCol=None, outputCol=None) + """ + super().__init__() + kwargs = self._input_kwargs + self._set(**kwargs) + + def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "StandardScalerModel": + input_col = self.getInputCol() + + stat_result = summarize_dataframe(dataset, input_col, ["mean", "std", "count"]) + mean_values = stat_result["mean"] + std_values = stat_result["std"] + n_samples_seen = stat_result["count"] + + model = StandardScalerModel(mean_values, std_values, n_samples_seen) + model._resetUid(self.uid) + return self._copyValues(model) + + +class StandardScalerModel(Model, HasInputCol, HasOutputCol, ParamsReadWrite, CoreModelReadWrite): + """ + Model fitted by StandardScaler. + + .. versionadded:: 3.5.0 + """ + + def __init__( + self, + mean_values: Optional["np.ndarray"] = None, + std_values: Optional["np.ndarray"] = None, + n_samples_seen: Optional[int] = None, + ) -> None: + super().__init__() + self.mean_values = mean_values + self.std_values = std_values + if std_values is not None: + # if scale value is zero, replace it with 1.0 (for preventing division by zero) + self.scale_values = np.where(std_values == 0.0, 1.0, std_values) + self.n_samples_seen = n_samples_seen + + def _input_columns(self) -> List[str]: + return [self.getInputCol()] + + def _output_columns(self) -> List[Tuple[str, str]]: + return [(self.getOutputCol(), "array")] + + def _get_transform_fn(self) -> Callable[..., Any]: + mean_values = self.mean_values + scale_values = self.scale_values + + def transform_fn(series: Any) -> Any: + def map_value(x: "np.ndarray") -> "np.ndarray": + return (x - mean_values) / scale_values + + return series.apply(map_value) + + return transform_fn + + def _get_core_model_filename(self) -> str: + return self.__class__.__name__ + ".sklearn.pkl" + + def _save_core_model(self, path: str) -> None: + from sklearn.preprocessing import StandardScaler as sk_StandardScaler + + sk_model = sk_StandardScaler(with_mean=True, with_std=True) + sk_model.scale_ = self.scale_values + sk_model.var_ = self.std_values * self.std_values # type: ignore[operator] + sk_model.mean_ = self.mean_values + sk_model.n_features_in_ = len(self.std_values) # type: ignore[arg-type] + sk_model.n_samples_seen_ = self.n_samples_seen + + with open(path, "wb") as fp: + pickle.dump(sk_model, fp) + + def _load_core_model(self, path: str) -> None: + with open(path, "rb") as fp: + sk_model = pickle.load(fp) + + self.std_values = np.sqrt(sk_model.var_) + self.scale_values = sk_model.scale_ + self.mean_values = sk_model.mean_ + self.n_samples_seen = sk_model.n_samples_seen_ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/functions.py new file mode 100644 index 0000000..ab7e3ab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/functions.py @@ -0,0 +1,76 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from pyspark.ml import functions as PyMLFunctions + +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.functions import _invoke_function, _to_col, lit + + +def vector_to_array(col: Column, dtype: str = "float64") -> Column: + return _invoke_function("vector_to_array", _to_col(col), lit(dtype)) + + +vector_to_array.__doc__ = PyMLFunctions.vector_to_array.__doc__ + + +def array_to_vector(col: Column) -> Column: + return _invoke_function("array_to_vector", _to_col(col)) + + +array_to_vector.__doc__ = PyMLFunctions.array_to_vector.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.ml.connect.functions + + globs = pyspark.ml.connect.functions.__dict__.copy() + + # TODO: split vector_to_array doctest since it includes .mllib vectors + del pyspark.ml.connect.functions.vector_to_array.__doc__ + + # TODO: spark.createDataFrame should support UDT + del pyspark.ml.connect.functions.array_to_vector.__doc__ + + globs["spark"] = ( + PySparkSession.builder.appName("ml.connect.functions tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.ml.connect.functions, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/io_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/io_utils.py new file mode 100644 index 0000000..a09a244 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/io_utils.py @@ -0,0 +1,335 @@ +# +# 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 json +import shutil +import os +import tempfile +import time +from urllib.parse import urlparse +from typing import Any, Dict, List +from pyspark.ml.base import Params +from pyspark.sql import SparkSession +from pyspark.sql.utils import is_remote + + +from pyspark import __version__ as pyspark_version + + +_META_DATA_FILE_NAME = "metadata.json" + + +def _copy_file_from_local_to_fs(local_path: str, dest_path: str) -> None: + session = SparkSession.active() + if is_remote(): + session.copyFromLocalToFs(local_path, dest_path) + else: + jvm = session.sparkContext._gateway.jvm # type: ignore[union-attr] + jvm.org.apache.spark.ml.python.MLUtil.copyFileFromLocalToFs(local_path, dest_path) + + +def _copy_dir_from_local_to_fs(local_path: str, dest_path: str) -> None: + """ + Copy directory from local path to cloud storage path. + Limitation: Currently only one level directory is supported. + """ + assert os.path.isdir(local_path) + + file_list = os.listdir(local_path) + for file_name in file_list: + file_path = os.path.join(local_path, file_name) + dest_file_path = os.path.join(dest_path, file_name) + assert os.path.isfile(file_path) + _copy_file_from_local_to_fs(file_path, dest_file_path) + + +def _get_class(clazz: str) -> Any: + """ + Loads Python class from its name. + """ + parts = clazz.split(".") + module = ".".join(parts[:-1]) + m = __import__(module, fromlist=[parts[-1]]) + return getattr(m, parts[-1]) + + +class ParamsReadWrite(Params): + """ + The base interface Estimator / Transformer / Model / Evaluator needs to inherit + for supporting saving and loading. + """ + + def _get_extra_metadata(self) -> Any: + """ + Returns exta metadata of the instance + """ + return None + + def _get_skip_saving_params(self) -> List[str]: + """ + Returns params to be skipped when saving metadata. + """ + return [] + + def _get_metadata_to_save(self) -> Dict[str, Any]: + """ + Extract metadata of Estimator / Transformer / Model / Evaluator instance. + """ + extra_metadata = self._get_extra_metadata() + skipped_params = self._get_skip_saving_params() + + uid = self.uid + cls = self.__module__ + "." + self.__class__.__name__ + + # User-supplied param values + params = self._paramMap + json_params = {} + skipped_params = skipped_params or [] + for p in params: + if p.name not in skipped_params: + json_params[p.name] = params[p] + + # Default param values + json_default_params = {} + for p in self._defaultParamMap: + json_default_params[p.name] = self._defaultParamMap[p] + + metadata = { + "class": cls, + "timestamp": int(round(time.time() * 1000)), + "sparkVersion": pyspark_version, + "uid": uid, + "paramMap": json_params, + "defaultParamMap": json_default_params, + "type": "spark_connect", + } + if extra_metadata is not None: + assert isinstance(extra_metadata, dict) + metadata["extra"] = extra_metadata + + return metadata + + def _load_extra_metadata(self, metadata: Dict[str, Any]) -> None: + """ + Load extra metadata attribute from metadata json object. + """ + pass + + def _save_to_local(self, path: str) -> None: + metadata = self._save_to_node_path(path, []) + with open(os.path.join(path, _META_DATA_FILE_NAME), "w") as fp: + json.dump(metadata, fp) + + def saveToLocal(self, path: str, *, overwrite: bool = False) -> None: + """ + Save Estimator / Transformer / Model / Evaluator to provided local path. + + .. versionadded:: 3.5.0 + """ + if os.path.exists(path): + if overwrite: + if os.path.isdir(path): + shutil.rmtree(path) + else: + os.remove(path) + else: + raise ValueError(f"The path {path} already exists.") + + os.makedirs(path) + self._save_to_local(path) + + @classmethod + def _load_metadata(cls, metadata: Dict[str, Any]) -> "Params": + if "type" not in metadata or metadata["type"] != "spark_connect": + raise RuntimeError( + "The saved data is not saved by ML algorithm implemented in 'pyspark.ml.connect' " + "module." + ) + + class_name = metadata["class"] + instance = _get_class(class_name)() + instance._resetUid(metadata["uid"]) + + # Set user-supplied param values + for paramName in metadata["paramMap"]: + param = instance.getParam(paramName) + paramValue = metadata["paramMap"][paramName] + instance.set(param, paramValue) + + for paramName in metadata["defaultParamMap"]: + paramValue = metadata["defaultParamMap"][paramName] + instance._setDefault(**{paramName: paramValue}) + + if "extra" in metadata: + instance._load_extra_metadata(metadata["extra"]) + return instance + + @classmethod + def _load_instance_from_metadata(cls, metadata: Dict[str, Any], path: str) -> Any: + instance = cls._load_metadata(metadata) + + if isinstance(instance, CoreModelReadWrite): + core_model_path = metadata["core_model_path"] + instance._load_core_model(os.path.join(path, core_model_path)) + + if isinstance(instance, MetaAlgorithmReadWrite): + instance._load_meta_algorithm(path, metadata) + + return instance + + @classmethod + def _load_from_local(cls, path: str) -> "Params": + with open(os.path.join(path, _META_DATA_FILE_NAME), "r") as fp: + metadata = json.load(fp) + + return cls._load_instance_from_metadata(metadata, path) + + @classmethod + def loadFromLocal(cls, path: str) -> "Params": + """ + Load Estimator / Transformer / Model / Evaluator from provided local path. + + .. versionadded:: 3.5.0 + """ + return cls._load_from_local(path) + + def _save_to_node_path(self, root_path: str, node_path: List[str]) -> Any: + """ + Save the instance to provided node path, and return the node metadata. + """ + if isinstance(self, MetaAlgorithmReadWrite): + metadata = self._save_meta_algorithm(root_path, node_path) + else: + metadata = self._get_metadata_to_save() + if isinstance(self, CoreModelReadWrite): + core_model_path = ".".join(node_path + [self._get_core_model_filename()]) + self._save_core_model(os.path.join(root_path, core_model_path)) + metadata["core_model_path"] = core_model_path + + return metadata + + def save(self, path: str, *, overwrite: bool = False) -> None: + """ + Save Estimator / Transformer / Model / Evaluator to provided cloud storage path. + + .. versionadded:: 3.5.0 + """ + session = SparkSession.active() + path_exist = True + try: + session.read.format("binaryFile").load(path).head() + except Exception as e: + if "Path does not exist" in str(e): + path_exist = False + else: + # Unexpected error. + raise e + + if path_exist and not overwrite: + raise ValueError(f"The path {path} already exists.") + + tmp_local_dir = tempfile.mkdtemp(prefix="pyspark_ml_model_") + try: + self._save_to_local(tmp_local_dir) + _copy_dir_from_local_to_fs(tmp_local_dir, path) + finally: + shutil.rmtree(tmp_local_dir, ignore_errors=True) + + @classmethod + def load(cls, path: str) -> "Params": + """ + Load Estimator / Transformer / Model / Evaluator from provided cloud storage path. + + .. versionadded:: 3.5.0 + """ + session = SparkSession.active() + + tmp_local_dir = tempfile.mkdtemp(prefix="pyspark_ml_model_") + try: + file_data_df = session.read.format("binaryFile").load(path) + + for row in file_data_df.toLocalIterator(): + file_name = os.path.basename(urlparse(row.path).path) + file_content = bytes(row.content) + with open(os.path.join(tmp_local_dir, file_name), "wb") as f: + f.write(file_content) + + return cls._load_from_local(tmp_local_dir) + finally: + shutil.rmtree(tmp_local_dir, ignore_errors=True) + + +class CoreModelReadWrite: + def _get_core_model_filename(self) -> str: + """ + Returns the name of the file for saving the core model. + """ + raise NotImplementedError() + + def _save_core_model(self, path: str) -> None: + """ + Save the core model to provided local path. + Different pyspark models contain different type of core model, + e.g. for LogisticRegressionModel, its core model is a pytorch model. + """ + raise NotImplementedError() + + def _load_core_model(self, path: str) -> None: + """ + Load the core model from provided local path. + """ + raise NotImplementedError() + + +class MetaAlgorithmReadWrite(ParamsReadWrite): + """ + Meta-algorithm such as pipeline and cross validator must implement this interface. + """ + + def _get_child_stages(self) -> List[Any]: + raise NotImplementedError() + + def _save_meta_algorithm(self, root_path: str, node_path: List[str]) -> Dict[str, Any]: + raise NotImplementedError() + + def _load_meta_algorithm(self, root_path: str, node_metadata: Dict[str, Any]) -> None: + raise NotImplementedError() + + @staticmethod + def _get_all_nested_stages(instance: Any) -> List[Any]: + if isinstance(instance, MetaAlgorithmReadWrite): + child_stages = instance._get_child_stages() + else: + child_stages = [] + + nested_stages = [] + for stage in child_stages: + nested_stages.extend(MetaAlgorithmReadWrite._get_all_nested_stages(stage)) + + return [instance] + nested_stages + + @staticmethod + def get_uid_map(instance: Any) -> Dict[str, Any]: + all_nested_stages = MetaAlgorithmReadWrite._get_all_nested_stages(instance) + uid_map = {stage.uid: stage for stage in all_nested_stages} + if len(all_nested_stages) != len(uid_map): + raise RuntimeError( + f"{instance.__class__.__module__}.{instance.__class__.__name__}" + f"is a compound estimator with stages with duplicate " + f"UIDs. List of UIDs: {list(uid_map.keys())}." + ) + return uid_map diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/pipeline.py new file mode 100644 index 0000000..0f7b364 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/pipeline.py @@ -0,0 +1,262 @@ +# +# 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 pandas as pd +from typing import Any, Dict, List, Optional, Union, cast, TYPE_CHECKING + +from pyspark import keyword_only, since +from pyspark.ml.connect.base import Estimator, Model, Transformer +from pyspark.ml.connect.io_utils import ( + ParamsReadWrite, + MetaAlgorithmReadWrite, +) +from pyspark.ml.param import Param, Params +from pyspark.ml.common import inherit_doc +from pyspark.sql.dataframe import DataFrame + + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + + +class _PipelineReadWrite(MetaAlgorithmReadWrite): + def _get_child_stages(self) -> List[Any]: + if isinstance(self, Pipeline): + return list(self.getStages()) + elif isinstance(self, PipelineModel): + return list(self.stages) + else: + raise ValueError(f"Unknown type {self.__class__}") + + def _get_skip_saving_params(self) -> List[str]: + """ + Returns params to be skipped when saving metadata. + """ + return ["stages"] + + def _save_meta_algorithm(self, root_path: str, node_path: List[str]) -> Dict[str, Any]: + metadata = self._get_metadata_to_save() + metadata["stages"] = [] + + if isinstance(self, Pipeline): + stages = self.getStages() + elif isinstance(self, PipelineModel): + stages = self.stages + else: + raise ValueError(f"Unknown type {self.__class__}") + + for stage_index, stage in enumerate(stages): + stage_node_path = node_path + [f"pipeline_stage_{stage_index}"] + stage_metadata = stage._save_to_node_path( # type: ignore[attr-defined] + root_path, stage_node_path + ) + metadata["stages"].append(stage_metadata) + return metadata + + def _load_meta_algorithm(self, root_path: str, node_metadata: Dict[str, Any]) -> None: + stages = [] + for stage_meta in node_metadata["stages"]: + stage = ParamsReadWrite._load_instance_from_metadata(stage_meta, root_path) + stages.append(stage) + + if isinstance(self, Pipeline): + self.setStages(stages) + elif isinstance(self, PipelineModel): + self.stages = stages + else: + raise ValueError() + + +@inherit_doc +class Pipeline(Estimator["PipelineModel"], _PipelineReadWrite): + """ + A simple pipeline, which acts as an estimator. A Pipeline consists + of a sequence of stages, each of which is either an + :py:class:`Estimator` or a :py:class:`Transformer`. When + :py:meth:`Pipeline.fit` is called, the stages are executed in + order. If a stage is an :py:class:`Estimator`, its + :py:meth:`Estimator.fit` method will be called on the input + dataset to fit a model. Then the model, which is a transformer, + will be used to transform the dataset as the input to the next + stage. If a stage is a :py:class:`Transformer`, its + :py:meth:`Transformer.transform` method will be called to produce + the dataset for the next stage. The fitted model from a + :py:class:`Pipeline` is a :py:class:`PipelineModel`, which + consists of fitted models and transformers, corresponding to the + pipeline stages. If stages is an empty list, the pipeline acts as an + identity transformer. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect import Pipeline + >>> from pyspark.ml.connect.classification import LogisticRegression + >>> from pyspark.ml.connect.feature import StandardScaler + >>> scaler = StandardScaler(inputCol='features', outputCol='scaled_features') + >>> lor = LogisticRegression(maxIter=20, learningRate=0.01) + >>> pipeline=Pipeline(stages=[scaler, lor]) + >>> dataset = spark.createDataFrame([ + ... ([1.0, 2.0], 1), + ... ([2.0, -1.0], 1), + ... ([-3.0, -2.0], 0), + ... ([-1.0, -2.0], 0), + ... ], schema=['features', 'label']) + >>> pipeline_model = pipeline.fit(dataset) + >>> transformed_dataset = pipeline_model.transform(dataset) + >>> transformed_dataset.show() + +------------+-----+--------------------+----------+--------------------+ + | features|label| scaled_features|prediction| probability| + +------------+-----+--------------------+----------+--------------------+ + | [1.0, 2.0]| 1|[0.56373452100212...| 1|[0.02423273026943...| + | [2.0, -1.0]| 1|[1.01472213780381...| 1|[0.09334788471460...| + |[-3.0, -2.0]| 0|[-1.2402159462046...| 0|[0.99808156490325...| + |[-1.0, -2.0]| 0|[-0.3382407126012...| 0|[0.96210002899169...| + +------------+-----+--------------------+----------+--------------------+ + >>> pipeline_model.saveToLocal("/tmp/pipeline") + >>> loaded_pipeline_model = PipelineModel.loadFromLocal("/tmp/pipeline") + """ + + stages: Param[List[Params]] = Param( + Params._dummy(), "stages", "a list of pipeline stages" + ) # type: ignore[assignment] + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, stages: Optional[List[Params]] = None): + """ + __init__(self, \\*, stages=None) + """ + super(Pipeline, self).__init__() + kwargs = self._input_kwargs + self.setParams(**kwargs) + + def setStages(self, value: List[Params]) -> "Pipeline": + """ + Set pipeline stages. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + value : list + of :py:class:`pyspark.ml.connect.Transformer` + or :py:class:`pyspark.ml.connect.Estimator` + + Returns + ------- + :py:class:`Pipeline` + the pipeline instance + """ + return self._set(stages=value) + + @since("3.5.0") + def getStages(self) -> List[Params]: + """ + Get pipeline stages. + """ + return self.getOrDefault(self.stages) + + @keyword_only + @since("3.5.0") + def setParams(self, *, stages: Optional[List[Params]] = None) -> "Pipeline": + """ + setParams(self, \\*, stages=None) + Sets params for Pipeline. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "PipelineModel": + stages = self.getStages() + for stage in stages: + if not (isinstance(stage, Estimator) or isinstance(stage, Transformer)): + raise TypeError("Cannot recognize a pipeline stage of type %s." % type(stage)) + indexOfLastEstimator = -1 + for i, stage in enumerate(stages): + if isinstance(stage, Estimator): + indexOfLastEstimator = i + transformers: List[Transformer] = [] + for i, stage in enumerate(stages): + if i <= indexOfLastEstimator: + if isinstance(stage, Transformer): + transformers.append(stage) + dataset = stage.transform(dataset) + else: # must be an Estimator + model = stage.fit(dataset) # type: ignore[attr-defined] + transformers.append(model) + if i < indexOfLastEstimator: + dataset = model.transform(dataset) + else: + transformers.append(cast(Transformer, stage)) + pipeline_model = PipelineModel(transformers) # type: ignore[arg-type] + pipeline_model._resetUid(self.uid) + return pipeline_model + + def copy(self, extra: Optional["ParamMap"] = None) -> "Pipeline": + """ + Creates a copy of this instance. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + extra : dict, optional + extra parameters + + Returns + ------- + :py:class:`Pipeline` + new instance + """ + if extra is None: + extra = dict() + that = Params.copy(self, extra) + stages = [stage.copy(extra) for stage in that.getStages()] + return that.setStages(stages) + + +@inherit_doc +class PipelineModel(Model, _PipelineReadWrite): + """ + Represents a compiled pipeline with transformers and fitted models. + + .. versionadded:: 3.5.0 + """ + + def __init__(self, stages: Optional[List[Params]] = None): + super(PipelineModel, self).__init__() + self.stages = stages # type: ignore[assignment] + + def _transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]: + for t in self.stages: + dataset = t.transform(dataset) + return dataset + + def copy(self, extra: Optional["ParamMap"] = None) -> "PipelineModel": + """ + Creates a copy of this instance. + + .. versionadded:: 3.5.0 + + :param extra: extra parameters + :returns: new instance + """ + if extra is None: + extra = dict() + stages = [stage.copy(extra) for stage in self.stages] + return PipelineModel(stages) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/summarizer.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/summarizer.py new file mode 100644 index 0000000..dd42c52 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/summarizer.py @@ -0,0 +1,120 @@ +# +# 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 numpy as np +import pandas as pd +from typing import Any, Union, List, Dict + +from pyspark.sql import DataFrame +from pyspark.ml.connect.util import aggregate_dataframe + + +class SummarizerAggState: + def __init__(self, input_array: "np.ndarray") -> None: + self.min_values = input_array.copy() + self.max_values = input_array.copy() + self.count = 1 + self.sum_values = np.array(input_array.copy()) + self.square_sum_values = np.square(input_array.copy()) + + def update(self, input_array: "np.ndarray") -> None: + self.count += 1 + self.sum_values += input_array + self.square_sum_values += np.square(input_array) + self.min_values = np.minimum(self.min_values, input_array) + self.max_values = np.maximum(self.max_values, input_array) + + def merge(self, state: "SummarizerAggState") -> "SummarizerAggState": + self.count += state.count + self.sum_values += state.sum_values + self.square_sum_values += state.square_sum_values + self.min_values = np.minimum(self.min_values, state.min_values) + self.max_values = np.maximum(self.max_values, state.max_values) + return self + + def to_result(self, metrics: List[str]) -> Dict[str, Any]: + result = {} + + for metric in metrics: + if metric == "min": + result["min"] = self.min_values.copy() + if metric == "max": + result["max"] = self.max_values.copy() + if metric == "sum": + result["sum"] = self.sum_values.copy() + if metric == "mean": + result["mean"] = self.sum_values / self.count + if metric == "std": + if self.count <= 1: + raise ValueError( + "Standard deviation evaluation requires more than one row data." + ) + result["std"] = np.sqrt( + ( + (self.square_sum_values / self.count) + - np.square(self.sum_values / self.count) + ) + * (self.count / (self.count - 1)) + ) + if metric == "count": + result["count"] = self.count # type: ignore[assignment] + + return result + + +def summarize_dataframe( + dataframe: Union["DataFrame", "pd.DataFrame"], column: str, metrics: List[str] +) -> Dict[str, Any]: + """ + Summarize an array type column over a spark dataframe or a pandas dataframe + + Parameters + ---------- + dataframe : :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame` + input dataset, it can be either pandas dataframe or spark dataframe. + + column: + The name of the column to be summarized, it must be an array type column + and all values in the column must have the same length. + metrics: + The metrics to be summarized, available metrics are: + "min", "max", "sum", "mean", "count" + + Returns + ------- + Summary results as a dict, the keys in the dict are the metrics being summarized. + """ + + def local_agg_fn(pandas_df: "pd.DataFrame") -> Any: + state = None + for _, value_array in pandas_df[column].items(): + if state is None: + state = SummarizerAggState(value_array) + else: + state.update(value_array) + + return state + + def merge_agg_state(state1: Any, state2: Any) -> Any: + return state1.merge(state2) + + def agg_state_to_result(state: Any) -> Any: + return state.to_result(metrics) + + return aggregate_dataframe( + dataframe, [column], local_agg_fn, merge_agg_state, agg_state_to_result + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/tuning.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/tuning.py new file mode 100644 index 0000000..871e448 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/tuning.py @@ -0,0 +1,579 @@ +# +# 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 multiprocessing.pool import ThreadPool + +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + cast, + TYPE_CHECKING, +) + +import numpy as np +import pandas as pd + +from pyspark import keyword_only, since, inheritable_thread_target +from pyspark.ml.connect import Estimator, Model +from pyspark.ml.connect.base import Evaluator +from pyspark.ml.connect.io_utils import ( + MetaAlgorithmReadWrite, + ParamsReadWrite, +) +from pyspark.ml.param import Params, Param, TypeConverters +from pyspark.ml.param.shared import HasParallelism, HasSeed +from pyspark.sql.functions import col, lit, rand +from pyspark.sql.dataframe import DataFrame +from pyspark.sql import SparkSession + +from pyspark.sql.utils import is_remote + + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + + +class _ValidatorParams(HasSeed): + """ + Common params for TrainValidationSplit and CrossValidator. + """ + + estimator: Param[Estimator] = Param( + Params._dummy(), "estimator", "estimator to be cross-validated" + ) + estimatorParamMaps: Param[List["ParamMap"]] = Param( + Params._dummy(), "estimatorParamMaps", "estimator param maps" + ) + evaluator: Param[Evaluator] = Param( + Params._dummy(), + "evaluator", + "evaluator used to select hyper-parameters that maximize the validator metric", + ) + + @since("2.0.0") + def getEstimator(self) -> Estimator: + """ + Gets the value of estimator or its default value. + """ + return self.getOrDefault(self.estimator) + + @since("2.0.0") + def getEstimatorParamMaps(self) -> List["ParamMap"]: + """ + Gets the value of estimatorParamMaps or its default value. + """ + return self.getOrDefault(self.estimatorParamMaps) + + @since("2.0.0") + def getEvaluator(self) -> Evaluator: + """ + Gets the value of evaluator or its default value. + """ + return self.getOrDefault(self.evaluator) + + +class _CrossValidatorParams(_ValidatorParams): + """ + Params for :py:class:`CrossValidator` and :py:class:`CrossValidatorModel`. + + .. versionadded:: 3.5.0 + """ + + numFolds: Param[int] = Param( + Params._dummy(), + "numFolds", + "number of folds for cross validation", + typeConverter=TypeConverters.toInt, + ) + + foldCol: Param[str] = Param( + Params._dummy(), + "foldCol", + "Param for the column name of user " + + "specified fold number. Once this is specified, :py:class:`CrossValidator` " + + "won't do random k-fold split. Note that this column should be integer type " + + "with range [0, numFolds) and Spark will throw exception on out-of-range " + + "fold numbers.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_CrossValidatorParams, self).__init__(*args) + self._setDefault(numFolds=3, foldCol="") + + @since("1.4.0") + def getNumFolds(self) -> int: + """ + Gets the value of numFolds or its default value. + """ + return self.getOrDefault(self.numFolds) + + @since("3.1.0") + def getFoldCol(self) -> str: + """ + Gets the value of foldCol or its default value. + """ + return self.getOrDefault(self.foldCol) + + +def _parallelFitTasks( + estimator: Estimator, + train: DataFrame, + evaluator: Evaluator, + validation: DataFrame, + epm: Sequence["ParamMap"], +) -> List[Callable[[], Tuple[int, float]]]: + """ + Creates a list of callables which can be called from different threads to fit and evaluate + an estimator in parallel. Each callable returns an `(index, metric)` pair. + + Parameters + ---------- + est : :py:class:`pyspark.ml.baseEstimator` + he estimator to be fit. + train : :py:class:`pyspark.sql.DataFrame` + DataFrame, training data set, used for fitting. + eva : :py:class:`pyspark.ml.evaluation.Evaluator` + used to compute `metric` + validation : :py:class:`pyspark.sql.DataFrame` + DataFrame, validation data set, used for evaluation. + epm : :py:class:`collections.abc.Sequence` + Sequence of ParamMap, params maps to be used during fitting & evaluation. + collectSubModel : bool + Whether to collect sub model. + + Returns + ------- + tuple + (int, float), an index into `epm` and the associated metric value. + """ + + active_session = SparkSession.getActiveSession() + + if active_session is None: + raise RuntimeError( + "An active SparkSession is required for running cross valiator fit tasks." + ) + + def get_single_task(index: int, param_map: Any) -> Callable[[], Tuple[int, float]]: + def single_task() -> Tuple[int, float]: + if not is_remote(): + # Active session is thread-local variable, in background thread the active session + # is not set, the following line sets it as the main thread active session. + active_session._jvm.SparkSession.setActiveSession( # type: ignore[union-attr] + active_session._jsparkSession # type: ignore[union-attr] + ) + + model = estimator.fit(train, param_map) + metric = evaluator.evaluate( + model.transform(validation, param_map) # type: ignore[union-attr] + ) + return index, metric + + return single_task + + return [get_single_task(index, param_map) for index, param_map in enumerate(epm)] + + +class _CrossValidatorReadWrite(MetaAlgorithmReadWrite): + def _get_skip_saving_params(self) -> List[str]: + """ + Returns params to be skipped when saving metadata. + """ + return ["estimator", "estimatorParamMaps", "evaluator"] + + def _save_meta_algorithm(self, root_path: str, node_path: List[str]) -> Dict[str, Any]: + metadata = self._get_metadata_to_save() + metadata[ + "estimator" + ] = self.getEstimator()._save_to_node_path( # type: ignore[attr-defined] + root_path, node_path + ["crossvalidator_estimator"] + ) + metadata[ + "evaluator" + ] = self.getEvaluator()._save_to_node_path( # type: ignore[attr-defined] + root_path, node_path + ["crossvalidator_evaluator"] + ) + metadata["estimator_param_maps"] = [ + [ + {"parent": param.parent, "name": param.name, "value": value} + for param, value in param_map.items() + ] + for param_map in self.getEstimatorParamMaps() # type: ignore[attr-defined] + ] + + if isinstance(self, CrossValidatorModel): + metadata["avg_metrics"] = self.avgMetrics + metadata["std_metrics"] = self.stdMetrics + + metadata["best_model"] = self.bestModel._save_to_node_path( + root_path, node_path + ["crossvalidator_best_model"] + ) + return metadata + + def _load_meta_algorithm(self, root_path: str, node_metadata: Dict[str, Any]) -> None: + estimator = ParamsReadWrite._load_instance_from_metadata( + node_metadata["estimator"], root_path + ) + self.set(self.estimator, estimator) # type: ignore[attr-defined] + + evaluator = ParamsReadWrite._load_instance_from_metadata( + node_metadata["evaluator"], root_path + ) + self.set(self.evaluator, evaluator) # type: ignore[attr-defined] + + json_epm = node_metadata["estimator_param_maps"] + + uid_to_instances = MetaAlgorithmReadWrite.get_uid_map(estimator) + + epm = [] + for json_param_map in json_epm: + param_map = {} + for json_param in json_param_map: + est = uid_to_instances[json_param["parent"]] + param = getattr(est, json_param["name"]) + value = json_param["value"] + param_map[param] = value + epm.append(param_map) + + self.set(self.estimatorParamMaps, epm) # type: ignore[attr-defined] + + if isinstance(self, CrossValidatorModel): + self.avgMetrics = node_metadata["avg_metrics"] + self.stdMetrics = node_metadata["std_metrics"] + + self.bestModel = ParamsReadWrite._load_instance_from_metadata( + node_metadata["best_model"], root_path + ) + + +class CrossValidator( + Estimator["CrossValidatorModel"], + _CrossValidatorParams, + HasParallelism, + _CrossValidatorReadWrite, +): + """ + K-fold cross validation performs model selection by splitting the dataset into a set of + non-overlapping randomly partitioned folds which are used as separate training and test datasets + e.g., with k=3 folds, K-fold cross validation will generate 3 (training, test) dataset pairs, + each of which uses 2/3 of the data for training and 1/3 for testing. Each fold is used as the + test set exactly once. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> from pyspark.ml.connect.tuning import CrossValidator + >>> from pyspark.ml.connect.classification import LogisticRegression + >>> from pyspark.ml.connect.evaluation import BinaryClassificationEvaluator + >>> from pyspark.ml.tuning import ParamGridBuilder + >>> from sklearn.datasets import load_breast_cancer + >>> lor = LogisticRegression(maxIter=20, learningRate=0.01) + >>> ev = BinaryClassificationEvaluator() + >>> grid = ParamGridBuilder().addGrid(lor.maxIter, [2, 20]).build() + >>> cv = CrossValidator(estimator=lor, evaluator=ev, estimatorParamMaps=grid) + >>> sk_dataset = load_breast_cancer() + >>> train_dataset = spark.createDataFrame( + ... zip(sk_dataset.data.tolist(), [int(t) for t in sk_dataset.target]), + ... schema="features: array, label: long", + ... ) + >>> cv_model = cv.fit(train_dataset) + >>> transformed_dataset = cv_model.transform(train_dataset.limit(10)) + >>> cv_model.avgMetrics + [0.5527792527167658, 0.8348714668615984] + >>> cv_model.stdMetrics + [0.04902833489813031, 0.05247132866444953] + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + numFolds: int = 3, + seed: Optional[int] = None, + parallelism: int = 1, + foldCol: str = "", + ) -> None: + """ + __init__(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, numFolds=3,\ + seed=None, parallelism=1, foldCol="") + """ + super(CrossValidator, self).__init__() + self._setDefault(parallelism=1) + kwargs = self._input_kwargs + self._set(**kwargs) + + @keyword_only + @since("3.5.0") + def setParams( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + numFolds: int = 3, + seed: Optional[int] = None, + parallelism: int = 1, + foldCol: str = "", + ) -> "CrossValidator": + """ + setParams(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, numFolds=3,\ + seed=None, parallelism=1, collectSubModels=False, foldCol=""): + Sets params for cross validator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.5.0") + def setEstimator(self, value: Estimator) -> "CrossValidator": + """ + Sets the value of :py:attr:`estimator`. + """ + return self._set(estimator=value) + + @since("3.5.0") + def setEstimatorParamMaps(self, value: List["ParamMap"]) -> "CrossValidator": + """ + Sets the value of :py:attr:`estimatorParamMaps`. + """ + return self._set(estimatorParamMaps=value) + + @since("3.5.0") + def setEvaluator(self, value: Evaluator) -> "CrossValidator": + """ + Sets the value of :py:attr:`evaluator`. + """ + return self._set(evaluator=value) + + @since("3.5.0") + def setNumFolds(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`numFolds`. + """ + return self._set(numFolds=value) + + @since("3.5.0") + def setFoldCol(self, value: str) -> "CrossValidator": + """ + Sets the value of :py:attr:`foldCol`. + """ + return self._set(foldCol=value) + + def setSeed(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def setParallelism(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`parallelism`. + """ + return self._set(parallelism=value) + + def setCollectSubModels(self, value: bool) -> "CrossValidator": + """ + Sets the value of :py:attr:`collectSubModels`. + """ + return self._set(collectSubModels=value) + + @staticmethod + def _gen_avg_and_std_metrics( + metrics_all: List[List[float]], + ) -> Tuple[List[float], List[float]]: + avg_metrics = np.mean(metrics_all, axis=0) + std_metrics = np.std(metrics_all, axis=0) + return list(avg_metrics), list(std_metrics) + + def _fit(self, dataset: Union[pd.DataFrame, DataFrame]) -> "CrossValidatorModel": + if isinstance(dataset, pd.DataFrame): + # TODO: support pandas dataframe fitting + raise NotImplementedError("Fitting pandas dataframe is not supported yet.") + + est = self.getOrDefault(self.estimator) + epm = self.getOrDefault(self.estimatorParamMaps) + numModels = len(epm) + eva = self.getOrDefault(self.evaluator) + nFolds = self.getOrDefault(self.numFolds) + metrics_all = [[0.0] * numModels for i in range(nFolds)] + + pool = ThreadPool(processes=min(self.getParallelism(), numModels)) + + datasets = self._kFold(dataset) + for i in range(nFolds): + validation = datasets[i][1].cache() + train = datasets[i][0].cache() + + tasks = _parallelFitTasks(est, train, eva, validation, epm) + if not is_remote(): + tasks = list(map(inheritable_thread_target, tasks)) + + for j, metric in pool.imap_unordered(lambda f: f(), tasks): + metrics_all[i][j] = metric + + validation.unpersist() + train.unpersist() + + metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all) + + if eva.isLargerBetter(): + bestIndex = np.argmax(metrics) + else: + bestIndex = np.argmin(metrics) + bestModel = cast(Model, est.fit(dataset, epm[bestIndex])) + cv_model = self._copyValues( + CrossValidatorModel( + bestModel, + avgMetrics=metrics, + stdMetrics=std_metrics, + ) + ) + cv_model._resetUid(self.uid) + return cv_model + + def _kFold(self, dataset: DataFrame) -> List[Tuple[DataFrame, DataFrame]]: + nFolds = self.getOrDefault(self.numFolds) + foldCol = self.getOrDefault(self.foldCol) + + datasets = [] + if not foldCol: + # Do random k-fold split. + seed = self.getOrDefault(self.seed) + h = 1.0 / nFolds + randCol = self.uid + "_rand" + df = dataset.select("*", rand(seed).alias(randCol)) + for i in range(nFolds): + validateLB = i * h + validateUB = (i + 1) * h + condition = (df[randCol] >= validateLB) & (df[randCol] < validateUB) + validation = df.filter(condition) + train = df.filter(~condition) + datasets.append((train, validation)) + else: + # TODO: + # Add verification that foldCol column values are in range [0, nFolds) + for i in range(nFolds): + training = dataset.filter(col(foldCol) != lit(i)) + validation = dataset.filter(col(foldCol) == lit(i)) + if training.isEmpty(): + raise ValueError("The training data at fold %s is empty." % i) + if validation.isEmpty(): + raise ValueError("The validation data at fold %s is empty." % i) + datasets.append((training, validation)) + + return datasets + + def copy(self, extra: Optional["ParamMap"] = None) -> "CrossValidator": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies creates a deep copy of + the embedded paramMap, and copies the embedded and extra parameters over. + + + .. versionadded:: 3.5.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`CrossValidator` + Copy of this instance + """ + if extra is None: + extra = dict() + newCV = Params.copy(self, extra) + if self.isSet(self.estimator): + newCV.setEstimator(self.getEstimator().copy(extra)) + # estimatorParamMaps remain the same + if self.isSet(self.evaluator): + newCV.setEvaluator(self.getEvaluator().copy(extra)) + return newCV + + +class CrossValidatorModel(Model, _CrossValidatorParams, _CrossValidatorReadWrite): + """ + CrossValidatorModel contains the model with the highest average cross-validation + metric across folds and uses this model to transform input data. CrossValidatorModel + also tracks the metrics for each param map evaluated. + + .. versionadded:: 3.5.0 + """ + + def __init__( + self, + bestModel: Optional[Model] = None, + avgMetrics: Optional[List[float]] = None, + stdMetrics: Optional[List[float]] = None, + ) -> None: + super(CrossValidatorModel, self).__init__() + #: best model from cross validation + self.bestModel = bestModel + #: Average cross-validation metrics for each paramMap in + #: CrossValidator.estimatorParamMaps, in the corresponding order. + self.avgMetrics = avgMetrics or [] + #: standard deviation of metrics for each paramMap in + #: CrossValidator.estimatorParamMaps, in the corresponding order. + self.stdMetrics = stdMetrics or [] + + def _transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]: + return self.bestModel.transform(dataset) + + def copy(self, extra: Optional["ParamMap"] = None) -> "CrossValidatorModel": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies the underlying bestModel, + creates a deep copy of the embedded paramMap, and + copies the embedded and extra parameters over. + It does not copy the extra Params into the subModels. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`CrossValidatorModel` + Copy of this instance + """ + if extra is None: + extra = dict() + bestModel = self.bestModel.copy(extra) + avgMetrics = list(self.avgMetrics) + stdMetrics = list(self.stdMetrics) + + return self._copyValues( + CrossValidatorModel(bestModel, avgMetrics=avgMetrics, stdMetrics=stdMetrics), + extra=extra, + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/util.py new file mode 100644 index 0000000..a543a16 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/connect/util.py @@ -0,0 +1,173 @@ +# +# 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 pandas as pd +from typing import Any, Union, List, Tuple, Callable, Iterable + +from pyspark import cloudpickle +from pyspark.sql import DataFrame +from pyspark.sql.functions import col, pandas_udf + + +def aggregate_dataframe( + dataframe: Union["DataFrame", "pd.DataFrame"], + input_col_names: List[str], + local_agg_fn: Callable[["pd.DataFrame"], Any], + merge_agg_state: Callable[[Any, Any], Any], + agg_state_to_result: Callable[[Any], Any], +) -> Any: + """ + The function can be used to run arbitrary aggregation logic on a spark dataframe + or a pandas dataframe. + + Parameters + ---------- + dataframe : + A spark dataframe or a pandas dataframe + + input_col_names : + The name of columns that are used in aggregation + + local_agg_fn : + A user-defined function that converts a pandas dataframe to an object holding + aggregation state. The aggregation state object must be pickle-able by + `cloudpickle`. + + merge_agg_state : + A user-defined function that merges 2 aggregation state objects into one and + return the merged state. Either in-place modifying the first input state object + and returning it or creating a new state object are acceptable. + + agg_state_to_result : + A user-defined function that converts aggregation state object to final aggregation + result. + + Returns + ------- + Aggregation result. + """ + + if isinstance(dataframe, pd.DataFrame): + dataframe = dataframe[list(input_col_names)] + agg_state = local_agg_fn(dataframe) + return agg_state_to_result(agg_state) + + dataframe = dataframe.select(*input_col_names) + + def compute_state(iterator: Iterable["pd.DataFrame"]) -> Iterable["pd.DataFrame"]: + state = None + + for batch_pandas_df in iterator: + new_batch_state = local_agg_fn(batch_pandas_df) + if state is None: + state = new_batch_state + else: + state = merge_agg_state(state, new_batch_state) + + if state is None: + pickled_state = None + else: + pickled_state = cloudpickle.dumps(state) + yield pd.DataFrame({"state": [pickled_state]}) + + result_pdf = dataframe.mapInPandas(compute_state, schema="state binary").toPandas() + + merged_state = None + for state in result_pdf.state: + if state is None: + continue + state = cloudpickle.loads(state) + if merged_state is None: + merged_state = state + else: + merged_state = merge_agg_state(merged_state, state) + + return agg_state_to_result(merged_state) + + +def transform_dataframe_column( + dataframe: Union["DataFrame", "pd.DataFrame"], + input_cols: List[str], + transform_fn: Callable[..., Any], + output_cols: List[Tuple[str, str]], +) -> Union["DataFrame", "pd.DataFrame"]: + """ + Transform specified column of the input spark dataframe or pandas dataframe, + returns a new dataframe + + Parameters + ---------- + dataframe : + A spark dataframe or a pandas dataframe + + input_cols : + A list of names of input columns to be transformed + + transform_fn: + A transforming function with one arguments of `pandas.Series` type, + if the transform function output is only one column data, + return transformed result as a `pandas.Series` object, + otherwise return transformed result as a `pandas.DataFrame` object + with corresponding column names defined in `output_cols` argument. + The output pandas Series/DataFrame object must have the same index + with the input series. + + output_cols: + a list of output transformed columns, each elements in the list + is a tuple of (column_name, column_spark_type) + + Returns + ------- + 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 input pandas dataframe is appended with output + columns in place. + """ + + if len(output_cols) > 1: + output_col_name = "__spark_ml_transformer_output_tmp__" + spark_udf_return_type = ",".join( + [f"{col_name} {col_type}" for col_name, col_type in output_cols] + ) + else: + output_col_name, spark_udf_return_type = output_cols[0] + + if isinstance(dataframe, pd.DataFrame): + result_data = transform_fn(*[dataframe[col_name] for col_name in input_cols]) + if isinstance(result_data, pd.Series): + assert len(output_cols) == 1 + result_data = pd.DataFrame({output_col_name: result_data}) + else: + assert set(result_data.columns) == set(col_name for col_name, _ in output_cols) + result_data = result_data + + for col_name in result_data.columns: + dataframe.insert(len(dataframe.columns), col_name, result_data[col_name]) + return dataframe + + @pandas_udf(returnType=spark_udf_return_type) # type: ignore[call-overload] + def transform_fn_pandas_udf(s: "pd.Series") -> "pd.Series": + return transform_fn(s) + + result_spark_df = dataframe.withColumn(output_col_name, transform_fn_pandas_udf(*input_cols)) + + if len(output_cols) > 1: + return result_spark_df.withColumns( + {col_name: col(f"{output_col_name}.{col_name}") for col_name, _ in output_cols} + ).drop(output_col_name) + else: + return result_spark_df diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/__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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/deepspeed_distributor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/deepspeed_distributor.py new file mode 100644 index 0000000..7c2b8c4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/deepspeed_distributor.py @@ -0,0 +1,165 @@ +# +# 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 json +import sys +import tempfile +from typing import ( + Union, + Callable, + List, + Dict, + Optional, + Any, +) + +from pyspark.ml.torch.distributor import TorchDistributor + + +class DeepspeedTorchDistributor(TorchDistributor): + + _DEEPSPEED_SSL_CONF = "deepspeed.spark.distributor.ignoreSsl" + + def __init__( + self, + numGpus: int = 1, + nnodes: int = 1, + localMode: bool = True, + useGpu: bool = True, + deepspeedConfig: Optional[Union[str, Dict[str, Any]]] = None, + ): + """ + This class is used to run deepspeed training workloads with spark clusters. + The user has the option to specify the number of gpus per node + and the number of nodes (the same as if running from terminal), + as well as specify a deepspeed configuration file. + + Parameters + ---------- + numGpus: int + The number of GPUs to use per node (analagous to num_gpus in deepspeed command). + nnodes: int + The number of nodes that should be used for the run. + localMode: bool + Whether or not to run the training in a distributed fashion or just locally. + useGpu: bool + Boolean flag to determine whether to utilize gpus. + deepspeedConfig: Union[Dict[str,Any], str] or None: + The configuration file to be used for launching the deepspeed application. + If it's a dictionary containing the parameters, then we will create the file. + If None, deepspeed will fall back to default parameters. + + Examples + -------- + Run Deepspeed training function on a single node + + >>> def train(learning_rate): + ... import deepspeed + ... # rest of training function + ... return model + >>> distributor = DeepspeedTorchDistributor( + ... numGpus=4, + ... nnodes=1, + ... useGpu=True, + ... localMode=True, + ... deepspeedConfig="path/to/config.json") + >>> output = distributor.run(train, 0.01) + + Run Deepspeed training function on multiple nodes + + >>> distributor = DeepspeedTorchDistributor( + ... numGpus=4, + ... nnodes=3, + ... useGpu=True, + ... localMode=False, + ... deepspeedConfig="path/to/config.json") + >>> output = distributor.run(train, 0.01) + """ + num_processes = numGpus * nnodes + self.deepspeed_config = deepspeedConfig + super().__init__( + num_processes, + localMode, + useGpu, + _ssl_conf=DeepspeedTorchDistributor._DEEPSPEED_SSL_CONF, + ) + self.cleanup_deepspeed_conf = False + + @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: + json.dump(deepspeed_config, file) + return file.name + deepspeed_config_path = deepspeed_config + # Empty value means the deepspeed will fall back to default settings. + if deepspeed_config is None: + return "" + return deepspeed_config_path + + @staticmethod + def _create_torchrun_command( + input_params: Dict[str, Any], train_path: str, *args: Any + ) -> List[str]: + local_mode = input_params["local_mode"] + num_processes = input_params["num_processes"] + deepspeed_config = input_params["deepspeed_config"] + deepspeed_config_path = DeepspeedTorchDistributor._get_deepspeed_config_path( + deepspeed_config + ) + torchrun_args, processes_per_node = TorchDistributor._get_torchrun_args( + local_mode, num_processes + ) + args_string = list(map(str, args)) + command_to_run = [ + sys.executable, + "-m", + "torch.distributed.run", + *torchrun_args, + f"--nproc_per_node={processes_per_node}", + train_path, + *args_string, + "--deepspeed", + ] + + # Don't have the deepspeed_config argument if no path is provided or no parameters set + if deepspeed_config_path == "": + return command_to_run + return command_to_run + ["--deepspeed_config", deepspeed_config_path] + + @staticmethod + def _run_training_on_pytorch_file( + input_params: Dict[str, Any], train_path: str, *args: Any, **kwargs: Any + ) -> None: + if kwargs: + raise ValueError( + "DeepspeedTorchDistributor with pytorch file doesn't support keyword arguments" + ) + + log_streaming_client = input_params.get("log_streaming_client", None) + training_command = DeepspeedTorchDistributor._create_torchrun_command( + input_params, train_path, *args + ) + DeepspeedTorchDistributor._execute_command( + training_command, log_streaming_client=log_streaming_client + ) + + def run(self, train_object: Union[Callable, str], *args: Any, **kwargs: Any) -> Optional[Any]: + # If the "train_object" is a string, then we assume it's a filepath. + # Otherwise, we assume it's a function. + return self._run( + train_object, DeepspeedTorchDistributor._run_training_on_pytorch_file, *args, **kwargs + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/tests/test_deepspeed_distributor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/tests/test_deepspeed_distributor.py new file mode 100644 index 0000000..590e541 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/deepspeed/tests/test_deepspeed_distributor.py @@ -0,0 +1,306 @@ +# mypy: ignore-errors +# 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 contextlib import contextmanager +import os +import shutil +import sys +import textwrap +from typing import Any, Callable, Dict, Tuple +import unittest + +from pyspark import SparkConf, SparkContext +from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor +from pyspark.sql import SparkSession +from pyspark.ml.torch.tests.test_distributor import ( + get_local_mode_conf, + set_up_test_dirs, + get_distributed_mode_conf, +) + +have_deepspeed = True +try: + import deepspeed # noqa: F401 +except ImportError: + have_deepspeed = False + + +class DeepspeedTorchDistributorUnitTests(unittest.TestCase): + def _get_env_var(self, var_name: str, default_value: Any) -> Any: + value = os.getenv(var_name) + if value: + return value + os.environ[var_name] = str(default_value) + return default_value + + def _get_env_variables_distributed(self) -> Tuple[Any, Any, Any]: + master_addr = self._get_env_var("MASTER_ADDR", "127.0.0.1") + master_port = self._get_env_var("MASTER_PORT", 2000) + rank = self._get_env_var("RANK", 0) + return master_addr, master_port, rank + + def test_get_torchrun_args_local(self) -> None: + number_of_processes = 5 + expected_torchrun_args_local = ["--standalone", "--nnodes=1"] + expected_processes_per_node_local = number_of_processes + ( + get_local_mode_torchrun_args, + process_per_node, + ) = DeepspeedTorchDistributor._get_torchrun_args(True, number_of_processes) + self.assertEqual(get_local_mode_torchrun_args, expected_torchrun_args_local) + self.assertEqual(expected_processes_per_node_local, process_per_node) + + def test_get_torchrun_args_distributed(self) -> None: + number_of_processes = 5 + master_addr, master_port, rank = self._get_env_variables_distributed() + expected_torchrun_args_distributed = [ + f"--nnodes={number_of_processes}", + f"--node_rank={rank}", + f"--rdzv_endpoint={master_addr}:{master_port}", + "--rdzv_id=0", + ] + torchrun_args_distributed, process_per_node = DeepspeedTorchDistributor._get_torchrun_args( + False, number_of_processes + ) + self.assertEqual(torchrun_args_distributed, expected_torchrun_args_distributed) + self.assertEqual(process_per_node, 1) + + def test_create_torchrun_command_local(self) -> None: + deepspeed_conf = "path/to/deepspeed" + train_file_path = "path/to/exec" + num_procs = 10 + input_params: Dict[str, Any] = {} + input_params["local_mode"] = True + input_params["num_processes"] = num_procs + input_params["deepspeed_config"] = deepspeed_conf + + torchrun_local_args_expected = ["--standalone", "--nnodes=1"] + with self.subTest(msg="Testing local training with no extra args"): + local_cmd_no_args_expected = [ + sys.executable, + "-m", + "torch.distributed.run", + *torchrun_local_args_expected, + f"--nproc_per_node={num_procs}", + train_file_path, + "--deepspeed", + "--deepspeed_config", + deepspeed_conf, + ] + local_cmd = DeepspeedTorchDistributor._create_torchrun_command( + input_params, train_file_path + ) + self.assertEqual(local_cmd, local_cmd_no_args_expected) + with self.subTest(msg="Testing local training with extra args for the training script"): + local_mode_version_args = ["--arg1", "--arg2"] + local_cmd_args_expected = [ + sys.executable, + "-m", + "torch.distributed.run", + *torchrun_local_args_expected, + f"--nproc_per_node={num_procs}", + train_file_path, + *local_mode_version_args, + "--deepspeed", + "--deepspeed_config", + deepspeed_conf, + ] + + local_cmd_with_args = DeepspeedTorchDistributor._create_torchrun_command( + input_params, train_file_path, *local_mode_version_args + ) + self.assertEqual(local_cmd_with_args, local_cmd_args_expected) + + def test_create_torchrun_command_distributed(self) -> None: + deepspeed_conf = "path/to/deepspeed" + train_file_path = "path/to/exec" + num_procs = 10 + input_params: Dict[str, Any] = {} + input_params["local_mode"] = True + input_params["num_processes"] = num_procs + input_params["deepspeed_config"] = deepspeed_conf + ( + distributed_master_address, + distributed_master_port, + distributed_rank, + ) = self._get_env_variables_distributed() + distributed_torchrun_args = [ + f"--nnodes={num_procs}", + f"--node_rank={distributed_rank}", + f"--rdzv_endpoint={distributed_master_address}:{distributed_master_port}", + "--rdzv_id=0", + ] + with self.subTest(msg="Distributed training command verification with no extra args"): + distributed_cmd_no_args_expected = [ + sys.executable, + "-m", + "torch.distributed.run", + *distributed_torchrun_args, + "--nproc_per_node=1", + train_file_path, + "--deepspeed", + "--deepspeed_config", + deepspeed_conf, + ] + input_params["local_mode"] = False + distributed_command = DeepspeedTorchDistributor._create_torchrun_command( + input_params, train_file_path + ) + self.assertEqual(distributed_cmd_no_args_expected, distributed_command) + with self.subTest(msg="Distributed training command verification with extra arguments"): + distributed_extra_args = ["-args1", "--args2"] + distributed_cmd_args_expected = [ + sys.executable, + "-m", + "torch.distributed.run", + *distributed_torchrun_args, + "--nproc_per_node=1", + train_file_path, + *distributed_extra_args, + "--deepspeed", + "--deepspeed_config", + deepspeed_conf, + ] + distributed_command_with_args = DeepspeedTorchDistributor._create_torchrun_command( + input_params, train_file_path, *distributed_extra_args + ) + self.assertEqual(distributed_cmd_args_expected, distributed_command_with_args) + + +def _create_basic_function() -> Callable: + # TODO: swap out with better test function + # once Deepspeed better supports CPU + def pythagoras(leg1: float, leg2: float) -> float: + import deepspeed + + print(deepspeed.__version__) + return (leg1 * leg1 + leg2 * leg2) ** 0.5 + + return pythagoras + + +@contextmanager +def _create_pytorch_training_test_file(): + # Note: when Deepspeed CPU support becomes better, + # switch in more realistic training files using Deepspeed + # optimizations + constructs + str_to_write = textwrap.dedent( + """ + import sys + def pythagorean_thm(x : int, y: int): # type: ignore + import deepspeed # type: ignore + return (x*x + y*y)**0.5 # type: ignore + print(pythagorean_thm(int(sys.argv[1]), int(sys.argv[2])))""" + ) + cp_path = "/tmp/test_deepspeed_training_file.py" + with open(cp_path, "w") as f: + f.write(str_to_write) + yield cp_path + os.remove(cp_path) + + +# The program and function that we use in the end-to-end tests +# is very simple because in the Spark CI we only have access +# to CPUs and at this point in time, CPU support is limited +# in Deepspeed. Once Deepspeed better supports CPU training +# and inference, the hope is to switch out the training +# and file for the tests with more realistic testing +# that use Deepspeed constructs. +@unittest.skipIf(not have_deepspeed, "deepspeed is required for these tests") +class DeepspeedTorchDistributorDistributedEndToEnd(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() # noqa + # "loadDefaults" is set to False because if not, the SparkConf will + # use contain configurations from the LocalEndToEnd test, + # which causes the test to break. + conf = SparkConf(loadDefaults=False) + for k, v in get_distributed_mode_conf().items(): + conf = conf.set(k, v) + conf = conf.set( + "spark.worker.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + sc = SparkContext("local-cluster[2,2,512]", cls.__name__, conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + def test_simple_function_e2e(self) -> None: + train_fn = _create_basic_function() + # Arguments for the pythagoras function train_fn + x = 3 + y = 4 + dist = DeepspeedTorchDistributor(numGpus=2, useGpu=False, localMode=False) + output = dist.run(train_fn, x, y) + self.assertEqual(output, 5) + + def test_pytorch_file_e2e(self) -> None: + # TODO: change to better test script + # once Deepspeed CPU support is better + with _create_pytorch_training_test_file() as cp_path: + dist = DeepspeedTorchDistributor(numGpus=True, useGpu=False, localMode=False) + dist.run(cp_path, 2, 5) + + +@unittest.skipIf(not have_deepspeed, "deepspeed is required for these tests") +class DeepspeedDistributorLocalEndToEndTests(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.gpu_discovery_script_file_name, cls.mnist_dir_path = set_up_test_dirs() # noqa + conf = SparkConf() + for k, v in get_local_mode_conf().items(): + conf = conf.set(k, v) + conf = conf.set( + "spark.driver.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + sc = SparkContext("local-cluster[2,2,512]", cls.__name__, conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + def test_simple_function_e2e(self) -> None: + train_fn = _create_basic_function() + # Arguments for the pythagoras function train_fn + x = 3 + y = 4 + dist = DeepspeedTorchDistributor(numGpus=2, useGpu=False, localMode=True) + output = dist.run(train_fn, x, y) + self.assertEqual(output, 5) + + def test_pytorch_file_e2e(self) -> None: + with _create_pytorch_training_test_file() as path_to_train_file: + dist = DeepspeedTorchDistributor(numGpus=2, useGpu=False, localMode=True) + dist.run(path_to_train_file, 2, 5) + + +if __name__ == "__main__": + from pyspark.ml.deepspeed.tests.test_deepspeed_distributor import * # noqa: F401,F403 + + try: + import xmlrunner # type:ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/dl_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/dl_util.py new file mode 100644 index 0000000..8ead529 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/dl_util.py @@ -0,0 +1,150 @@ +# +# 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 os +import tempfile +import textwrap +from typing import Any, Callable + +from pyspark import cloudpickle + + +class FunctionPickler: + """ + This class provides a way to pickle a function and its arguments. + It also provides a way to create a script that can run a + function with arguments if they have them pickled to a file. + It also provides a way of extracting the conents of a pickle file. + """ + + @staticmethod + def pickle_fn_and_save( + fn: Callable, file_path: str, save_dir: str, *args: Any, **kwargs: Any + ) -> str: + """ + Given a function and args, this function will pickle them to a file. + + Parameters + ---------- + fn: Callable + The picklable function that will be pickled to a file. + file_path: str + The path where to save the pickled function, args, and kwargs. If it's the + empty string, the function will decide on a random name. + save_dir: str + The directory in which to save the file with the pickled function and arguments. + Does nothing if the path is specified. If both file_path and save_dir are empty, + the function will write the file to the current working directory with a random + name. + *args: Any + Arguments of fn that will be pickled. + **kwargs: Any + Key word arguments to fn that will be pickled. + + Returns + ------- + str + The path to the file where the function and arguments are pickled. + """ + if file_path != "": + with open(file_path, "wb") as f: + cloudpickle.dump((fn, args, kwargs), f) + return f.name + + if save_dir == "": + save_dir = os.getcwd() + + with tempfile.NamedTemporaryFile(dir=save_dir, delete=False) as f: + cloudpickle.dump((fn, args, kwargs), f) + return f.name + + @staticmethod + def create_fn_run_script( + pickled_fn_path: str, + fn_output_path: str, + script_path: str, + prefix_code: str = "", + suffix_code: str = "", + ) -> str: + """ + Given a file containing a pickled function and arguments, this function will create a + pytorch file that will execute the function and pickle the functions outputs. + + Parameters + ---------- + pickled_fn_path: str + This is the path of the file containing the pickled function, args, and kwargs. + fn_output_path: str + This is the location where the created file will save the pickled output of + the function. + script_path: str + This is the path which will be used for the created pytorch file. + prefix_code: str + This contains a string that the user can pass in which will be executed before + the code generated by this class to execute the function and save it. If + prefix_code is the empty string, nothing will be written before the auto- + generated code. + suffix_code: str + This contains a string of code that the user can pass in which will be executed + after the code generated by this class finishes executing. If suffix_code is + the empty string, nothing will be written after the auto-generated code. + + Returns + ------- + str + The path to the location of the newly created pytorch file. + """ + + code_snippet = textwrap.dedent( + f""" + from pyspark import cloudpickle + import os + + if __name__ == "__main__": + with open("{pickled_fn_path}", "rb") as f: + fn, args, kwargs = cloudpickle.load(f) + output = fn(*args, **kwargs) + with open("{fn_output_path}", "wb") as f: + cloudpickle.dump(output, f) + """ + ) + with open(script_path, "w") as f: + if prefix_code != "": + f.write(prefix_code) + f.write(code_snippet) + if suffix_code != "": + f.write(suffix_code) + + return script_path + + @staticmethod + def get_fn_output(fn_output_path: str) -> Any: + """ + Given a path to a file with pickled output, this function + will unpickle the output and return it to the user. + + Parameters + ---------- + fn_output_path: str + The path to the file containing the pickled output of a function. + + Returns + ------- + Any + The unpickled output stored in func_output_path + """ + with open(fn_output_path, "rb") as f: + return cloudpickle.load(f) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/evaluation.py new file mode 100644 index 0000000..19d123d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/evaluation.py @@ -0,0 +1,1166 @@ +# +# 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 sys +from abc import abstractmethod, ABCMeta + +from typing import Any, Dict, Optional, TYPE_CHECKING + +from pyspark import since, keyword_only +from pyspark.ml.wrapper import JavaParams +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import ( + HasLabelCol, + HasPredictionCol, + HasProbabilityCol, + HasRawPredictionCol, + HasFeaturesCol, + HasWeightCol, +) +from pyspark.ml.common import inherit_doc +from pyspark.ml.util import JavaMLReadable, JavaMLWritable +from pyspark.sql.dataframe import DataFrame + +if TYPE_CHECKING: + from pyspark.ml._typing import ( + ParamMap, + BinaryClassificationEvaluatorMetricType, + ClusteringEvaluatorDistanceMeasureType, + ClusteringEvaluatorMetricType, + MulticlassClassificationEvaluatorMetricType, + MultilabelClassificationEvaluatorMetricType, + RankingEvaluatorMetricType, + RegressionEvaluatorMetricType, + ) + + +__all__ = [ + "Evaluator", + "BinaryClassificationEvaluator", + "RegressionEvaluator", + "MulticlassClassificationEvaluator", + "MultilabelClassificationEvaluator", + "ClusteringEvaluator", + "RankingEvaluator", +] + + +@inherit_doc +class Evaluator(Params, metaclass=ABCMeta): + """ + Base class for evaluators that compute metrics from predictions. + + .. versionadded:: 1.4.0 + """ + + @abstractmethod + def _evaluate(self, dataset: 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:: 1.4.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. + """ + return True + + +@inherit_doc +class JavaEvaluator(JavaParams, Evaluator, metaclass=ABCMeta): + """ + Base class for :py:class:`Evaluator`s that wrap Java/Scala + implementations. + """ + + def _evaluate(self, dataset: DataFrame) -> float: + """ + Evaluates the output. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a dataset that contains labels/observations and predictions + + Returns + ------- + float + evaluation metric + """ + self._transfer_params_to_java() + assert self._java_obj is not None + return self._java_obj.evaluate(dataset._jdf) + + def isLargerBetter(self) -> bool: + self._transfer_params_to_java() + assert self._java_obj is not None + return self._java_obj.isLargerBetter() + + +@inherit_doc +class BinaryClassificationEvaluator( + JavaEvaluator, + HasLabelCol, + HasRawPredictionCol, + HasWeightCol, + JavaMLReadable["BinaryClassificationEvaluator"], + JavaMLWritable, +): + """ + Evaluator for binary classification, which expects input columns rawPrediction, label + and an optional weight column. + The rawPrediction column can be of type double (binary 0/1 prediction, or probability of label + 1) or of type vector (length-2 vector of raw predictions, scores, or label probabilities). + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> scoreAndLabels = map(lambda x: (Vectors.dense([1.0 - x[0], x[0]]), x[1]), + ... [(0.1, 0.0), (0.1, 1.0), (0.4, 0.0), (0.6, 0.0), (0.6, 1.0), (0.6, 1.0), (0.8, 1.0)]) + >>> dataset = spark.createDataFrame(scoreAndLabels, ["raw", "label"]) + ... + >>> evaluator = BinaryClassificationEvaluator() + >>> evaluator.setRawPredictionCol("raw") + BinaryClassificationEvaluator... + >>> evaluator.evaluate(dataset) + 0.70... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "areaUnderPR"}) + 0.83... + >>> bce_path = temp_path + "/bce" + >>> evaluator.save(bce_path) + >>> evaluator2 = BinaryClassificationEvaluator.load(bce_path) + >>> str(evaluator2.getRawPredictionCol()) + 'raw' + >>> scoreAndLabelsAndWeight = map(lambda x: (Vectors.dense([1.0 - x[0], x[0]]), x[1], x[2]), + ... [(0.1, 0.0, 1.0), (0.1, 1.0, 0.9), (0.4, 0.0, 0.7), (0.6, 0.0, 0.9), + ... (0.6, 1.0, 1.0), (0.6, 1.0, 0.3), (0.8, 1.0, 1.0)]) + >>> dataset = spark.createDataFrame(scoreAndLabelsAndWeight, ["raw", "label", "weight"]) + ... + >>> evaluator = BinaryClassificationEvaluator(rawPredictionCol="raw", weightCol="weight") + >>> evaluator.evaluate(dataset) + 0.70... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "areaUnderPR"}) + 0.82... + >>> evaluator.getNumBins() + 1000 + """ + + metricName: Param["BinaryClassificationEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + "metric name in evaluation (areaUnderROC|areaUnderPR)", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + + numBins: Param[int] = Param( + Params._dummy(), + "numBins", + "Number of bins to down-sample the curves " + "(ROC curve, PR curve) in area computation. If 0, no down-sampling will " + "occur. Must be >= 0.", + typeConverter=TypeConverters.toInt, + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + rawPredictionCol: str = "rawPrediction", + labelCol: str = "label", + metricName: "BinaryClassificationEvaluatorMetricType" = "areaUnderROC", + weightCol: Optional[str] = None, + numBins: int = 1000, + ): + """ + __init__(self, \\*, rawPredictionCol="rawPrediction", labelCol="label", \ + metricName="areaUnderROC", weightCol=None, numBins=1000) + """ + super(BinaryClassificationEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.BinaryClassificationEvaluator", self.uid + ) + self._setDefault(metricName="areaUnderROC", numBins=1000) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("1.4.0") + def setMetricName( + self, value: "BinaryClassificationEvaluatorMetricType" + ) -> "BinaryClassificationEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("1.4.0") + def getMetricName(self) -> str: + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("3.0.0") + def setNumBins(self, value: int) -> "BinaryClassificationEvaluator": + """ + Sets the value of :py:attr:`numBins`. + """ + return self._set(numBins=value) + + @since("3.0.0") + def getNumBins(self) -> int: + """ + Gets the value of numBins or its default value. + """ + return self.getOrDefault(self.numBins) + + def setLabelCol(self, value: str) -> "BinaryClassificationEvaluator": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def setRawPredictionCol(self, value: str) -> "BinaryClassificationEvaluator": + """ + Sets the value of :py:attr:`rawPredictionCol`. + """ + return self._set(rawPredictionCol=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "BinaryClassificationEvaluator": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + rawPredictionCol: str = "rawPrediction", + labelCol: str = "label", + metricName: "BinaryClassificationEvaluatorMetricType" = "areaUnderROC", + weightCol: Optional[str] = None, + numBins: int = 1000, + ) -> "BinaryClassificationEvaluator": + """ + setParams(self, \\*, rawPredictionCol="rawPrediction", labelCol="label", \ + metricName="areaUnderROC", weightCol=None, numBins=1000) + Sets params for binary classification evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +@inherit_doc +class RegressionEvaluator( + JavaEvaluator, + HasLabelCol, + HasPredictionCol, + HasWeightCol, + JavaMLReadable["RegressionEvaluator"], + JavaMLWritable, +): + """ + Evaluator for Regression, which expects input columns prediction, label + and an optional weight column. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> scoreAndLabels = [(-28.98343821, -27.0), (20.21491975, 21.5), + ... (-25.98418959, -22.0), (30.69731842, 33.0), (74.69283752, 71.0)] + >>> dataset = spark.createDataFrame(scoreAndLabels, ["raw", "label"]) + ... + >>> evaluator = RegressionEvaluator() + >>> evaluator.setPredictionCol("raw") + RegressionEvaluator... + >>> evaluator.evaluate(dataset) + 2.842... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "r2"}) + 0.993... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "mae"}) + 2.649... + >>> re_path = temp_path + "/re" + >>> evaluator.save(re_path) + >>> evaluator2 = RegressionEvaluator.load(re_path) + >>> str(evaluator2.getPredictionCol()) + 'raw' + >>> scoreAndLabelsAndWeight = [(-28.98343821, -27.0, 1.0), (20.21491975, 21.5, 0.8), + ... (-25.98418959, -22.0, 1.0), (30.69731842, 33.0, 0.6), (74.69283752, 71.0, 0.2)] + >>> dataset = spark.createDataFrame(scoreAndLabelsAndWeight, ["raw", "label", "weight"]) + ... + >>> evaluator = RegressionEvaluator(predictionCol="raw", weightCol="weight") + >>> evaluator.evaluate(dataset) + 2.740... + >>> evaluator.getThroughOrigin() + False + """ + + metricName: Param["RegressionEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + """metric name in evaluation - one of: + rmse - root mean squared error (default) + mse - mean squared error + r2 - r^2 metric + mae - mean absolute error + var - explained variance.""", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + + throughOrigin: Param[bool] = Param( + Params._dummy(), + "throughOrigin", + "whether the regression is through the origin.", + typeConverter=TypeConverters.toBoolean, + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "RegressionEvaluatorMetricType" = "rmse", + weightCol: Optional[str] = None, + throughOrigin: bool = False, + ): + """ + __init__(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="rmse", weightCol=None, throughOrigin=False) + """ + super(RegressionEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.RegressionEvaluator", self.uid + ) + self._setDefault(metricName="rmse", throughOrigin=False) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("1.4.0") + def setMetricName(self, value: "RegressionEvaluatorMetricType") -> "RegressionEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("1.4.0") + def getMetricName(self) -> "RegressionEvaluatorMetricType": + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("3.0.0") + def setThroughOrigin(self, value: bool) -> "RegressionEvaluator": + """ + Sets the value of :py:attr:`throughOrigin`. + """ + return self._set(throughOrigin=value) + + @since("3.0.0") + def getThroughOrigin(self) -> bool: + """ + Gets the value of throughOrigin or its default value. + """ + return self.getOrDefault(self.throughOrigin) + + def setLabelCol(self, value: str) -> "RegressionEvaluator": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def setPredictionCol(self, value: str) -> "RegressionEvaluator": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "RegressionEvaluator": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "RegressionEvaluatorMetricType" = "rmse", + weightCol: Optional[str] = None, + throughOrigin: bool = False, + ) -> "RegressionEvaluator": + """ + setParams(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="rmse", weightCol=None, throughOrigin=False) + Sets params for regression evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +@inherit_doc +class MulticlassClassificationEvaluator( + JavaEvaluator, + HasLabelCol, + HasPredictionCol, + HasWeightCol, + HasProbabilityCol, + JavaMLReadable["MulticlassClassificationEvaluator"], + JavaMLWritable, +): + """ + Evaluator for Multiclass Classification, which expects input + columns: prediction, label, weight (optional) and probabilityCol (only for logLoss). + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> scoreAndLabels = [(0.0, 0.0), (0.0, 1.0), (0.0, 0.0), + ... (1.0, 0.0), (1.0, 1.0), (1.0, 1.0), (1.0, 1.0), (2.0, 2.0), (2.0, 0.0)] + >>> dataset = spark.createDataFrame(scoreAndLabels, ["prediction", "label"]) + >>> evaluator = MulticlassClassificationEvaluator() + >>> evaluator.setPredictionCol("prediction") + MulticlassClassificationEvaluator... + >>> evaluator.evaluate(dataset) + 0.66... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "accuracy"}) + 0.66... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "truePositiveRateByLabel", + ... evaluator.metricLabel: 1.0}) + 0.75... + >>> evaluator.setMetricName("hammingLoss") + MulticlassClassificationEvaluator... + >>> evaluator.evaluate(dataset) + 0.33... + >>> mce_path = temp_path + "/mce" + >>> evaluator.save(mce_path) + >>> evaluator2 = MulticlassClassificationEvaluator.load(mce_path) + >>> str(evaluator2.getPredictionCol()) + 'prediction' + >>> scoreAndLabelsAndWeight = [(0.0, 0.0, 1.0), (0.0, 1.0, 1.0), (0.0, 0.0, 1.0), + ... (1.0, 0.0, 1.0), (1.0, 1.0, 1.0), (1.0, 1.0, 1.0), (1.0, 1.0, 1.0), + ... (2.0, 2.0, 1.0), (2.0, 0.0, 1.0)] + >>> dataset = spark.createDataFrame(scoreAndLabelsAndWeight, ["prediction", "label", "weight"]) + >>> evaluator = MulticlassClassificationEvaluator(predictionCol="prediction", + ... weightCol="weight") + >>> evaluator.evaluate(dataset) + 0.66... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "accuracy"}) + 0.66... + >>> predictionAndLabelsWithProbabilities = [ + ... (1.0, 1.0, 1.0, [0.1, 0.8, 0.1]), (0.0, 2.0, 1.0, [0.9, 0.05, 0.05]), + ... (0.0, 0.0, 1.0, [0.8, 0.2, 0.0]), (1.0, 1.0, 1.0, [0.3, 0.65, 0.05])] + >>> dataset = spark.createDataFrame(predictionAndLabelsWithProbabilities, ["prediction", + ... "label", "weight", "probability"]) + >>> evaluator = MulticlassClassificationEvaluator(predictionCol="prediction", + ... probabilityCol="probability") + >>> evaluator.setMetricName("logLoss") + MulticlassClassificationEvaluator... + >>> evaluator.evaluate(dataset) + 0.9682... + """ + + metricName: Param["MulticlassClassificationEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + "metric name in evaluation " + "(f1|accuracy|weightedPrecision|weightedRecall|weightedTruePositiveRate| " + "weightedFalsePositiveRate|weightedFMeasure|truePositiveRateByLabel| " + "falsePositiveRateByLabel|precisionByLabel|recallByLabel|fMeasureByLabel| " + "logLoss|hammingLoss)", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + metricLabel: Param[float] = Param( + Params._dummy(), + "metricLabel", + "The class whose metric will be computed in truePositiveRateByLabel|" + "falsePositiveRateByLabel|precisionByLabel|recallByLabel|fMeasureByLabel." + " Must be >= 0. The default value is 0.", + typeConverter=TypeConverters.toFloat, + ) + beta: Param[float] = Param( + Params._dummy(), + "beta", + "The beta value used in weightedFMeasure|fMeasureByLabel." + " Must be > 0. The default value is 1.", + typeConverter=TypeConverters.toFloat, + ) + eps: Param[float] = Param( + Params._dummy(), + "eps", + "log-loss is undefined for p=0 or p=1, so probabilities are clipped to " + "max(eps, min(1 - eps, p)). " + "Must be in range (0, 0.5). The default value is 1e-15.", + typeConverter=TypeConverters.toFloat, + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "MulticlassClassificationEvaluatorMetricType" = "f1", + weightCol: Optional[str] = None, + metricLabel: float = 0.0, + beta: float = 1.0, + probabilityCol: str = "probability", + eps: float = 1e-15, + ): + """ + __init__(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="f1", weightCol=None, metricLabel=0.0, beta=1.0, \ + probabilityCol="probability", eps=1e-15) + """ + super(MulticlassClassificationEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator", self.uid + ) + self._setDefault(metricName="f1", metricLabel=0.0, beta=1.0, eps=1e-15) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("1.5.0") + def setMetricName( + self, value: "MulticlassClassificationEvaluatorMetricType" + ) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("1.5.0") + def getMetricName(self) -> "MulticlassClassificationEvaluatorMetricType": + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("3.0.0") + def setMetricLabel(self, value: float) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`metricLabel`. + """ + return self._set(metricLabel=value) + + @since("3.0.0") + def getMetricLabel(self) -> float: + """ + Gets the value of metricLabel or its default value. + """ + return self.getOrDefault(self.metricLabel) + + @since("3.0.0") + def setBeta(self, value: float) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`beta`. + """ + return self._set(beta=value) + + @since("3.0.0") + def getBeta(self) -> float: + """ + Gets the value of beta or its default value. + """ + return self.getOrDefault(self.beta) + + @since("3.0.0") + def setEps(self, value: float) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`eps`. + """ + return self._set(eps=value) + + @since("3.0.0") + def getEps(self) -> float: + """ + Gets the value of eps or its default value. + """ + return self.getOrDefault(self.eps) + + def setLabelCol(self, value: str) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def setPredictionCol(self, value: str) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("3.0.0") + def setProbabilityCol(self, value: str) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`probabilityCol`. + """ + return self._set(probabilityCol=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "MulticlassClassificationEvaluator": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "MulticlassClassificationEvaluatorMetricType" = "f1", + weightCol: Optional[str] = None, + metricLabel: float = 0.0, + beta: float = 1.0, + probabilityCol: str = "probability", + eps: float = 1e-15, + ) -> "MulticlassClassificationEvaluator": + """ + setParams(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="f1", weightCol=None, metricLabel=0.0, beta=1.0, \ + probabilityCol="probability", eps=1e-15) + Sets params for multiclass classification evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +@inherit_doc +class MultilabelClassificationEvaluator( + JavaEvaluator, + HasLabelCol, + HasPredictionCol, + JavaMLReadable["MultilabelClassificationEvaluator"], + JavaMLWritable, +): + """ + Evaluator for Multilabel Classification, which expects two input + columns: prediction and label. + + .. versionadded:: 3.0.0 + + Notes + ----- + Experimental + + Examples + -------- + >>> scoreAndLabels = [([0.0, 1.0], [0.0, 2.0]), ([0.0, 2.0], [0.0, 1.0]), + ... ([], [0.0]), ([2.0], [2.0]), ([2.0, 0.0], [2.0, 0.0]), + ... ([0.0, 1.0, 2.0], [0.0, 1.0]), ([1.0], [1.0, 2.0])] + >>> dataset = spark.createDataFrame(scoreAndLabels, ["prediction", "label"]) + ... + >>> evaluator = MultilabelClassificationEvaluator() + >>> evaluator.setPredictionCol("prediction") + MultilabelClassificationEvaluator... + >>> evaluator.evaluate(dataset) + 0.63... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "accuracy"}) + 0.54... + >>> mlce_path = temp_path + "/mlce" + >>> evaluator.save(mlce_path) + >>> evaluator2 = MultilabelClassificationEvaluator.load(mlce_path) + >>> str(evaluator2.getPredictionCol()) + 'prediction' + """ + + metricName: Param["MultilabelClassificationEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + "metric name in evaluation " + "(subsetAccuracy|accuracy|hammingLoss|precision|recall|f1Measure|" + "precisionByLabel|recallByLabel|f1MeasureByLabel|microPrecision|" + "microRecall|microF1Measure)", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + metricLabel: Param[float] = Param( + Params._dummy(), + "metricLabel", + "The class whose metric will be computed in precisionByLabel|" + "recallByLabel|f1MeasureByLabel. " + "Must be >= 0. The default value is 0.", + typeConverter=TypeConverters.toFloat, + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "MultilabelClassificationEvaluatorMetricType" = "f1Measure", + metricLabel: float = 0.0, + ) -> None: + """ + __init__(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="f1Measure", metricLabel=0.0) + """ + super(MultilabelClassificationEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.MultilabelClassificationEvaluator", self.uid + ) + self._setDefault(metricName="f1Measure", metricLabel=0.0) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("3.0.0") + def setMetricName( + self, value: "MultilabelClassificationEvaluatorMetricType" + ) -> "MultilabelClassificationEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("3.0.0") + def getMetricName(self) -> "MultilabelClassificationEvaluatorMetricType": + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("3.0.0") + def setMetricLabel(self, value: float) -> "MultilabelClassificationEvaluator": + """ + Sets the value of :py:attr:`metricLabel`. + """ + return self._set(metricLabel=value) + + @since("3.0.0") + def getMetricLabel(self) -> float: + """ + Gets the value of metricLabel or its default value. + """ + return self.getOrDefault(self.metricLabel) + + @since("3.0.0") + def setLabelCol(self, value: str) -> "MultilabelClassificationEvaluator": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "MultilabelClassificationEvaluator": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @keyword_only + @since("3.0.0") + def setParams( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "MultilabelClassificationEvaluatorMetricType" = "f1Measure", + metricLabel: float = 0.0, + ) -> "MultilabelClassificationEvaluator": + """ + setParams(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="f1Measure", metricLabel=0.0) + Sets params for multilabel classification evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +@inherit_doc +class ClusteringEvaluator( + JavaEvaluator, + HasPredictionCol, + HasFeaturesCol, + HasWeightCol, + JavaMLReadable["ClusteringEvaluator"], + JavaMLWritable, +): + """ + Evaluator for Clustering results, which expects two input + columns: prediction and features. The metric computes the Silhouette + measure using the squared Euclidean distance. + + The Silhouette is a measure for the validation of the consistency + within clusters. It ranges between 1 and -1, where a value close to + 1 means that the points in a cluster are close to the other points + in the same cluster and far from the points of the other clusters. + + .. versionadded:: 2.3.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> featureAndPredictions = map(lambda x: (Vectors.dense(x[0]), x[1]), + ... [([0.0, 0.5], 0.0), ([0.5, 0.0], 0.0), ([10.0, 11.0], 1.0), + ... ([10.5, 11.5], 1.0), ([1.0, 1.0], 0.0), ([8.0, 6.0], 1.0)]) + >>> dataset = spark.createDataFrame(featureAndPredictions, ["features", "prediction"]) + ... + >>> evaluator = ClusteringEvaluator() + >>> evaluator.setPredictionCol("prediction") + ClusteringEvaluator... + >>> evaluator.evaluate(dataset) + 0.9079... + >>> featureAndPredictionsWithWeight = map(lambda x: (Vectors.dense(x[0]), x[1], x[2]), + ... [([0.0, 0.5], 0.0, 2.5), ([0.5, 0.0], 0.0, 2.5), ([10.0, 11.0], 1.0, 2.5), + ... ([10.5, 11.5], 1.0, 2.5), ([1.0, 1.0], 0.0, 2.5), ([8.0, 6.0], 1.0, 2.5)]) + >>> dataset = spark.createDataFrame( + ... featureAndPredictionsWithWeight, ["features", "prediction", "weight"]) + >>> evaluator = ClusteringEvaluator() + >>> evaluator.setPredictionCol("prediction") + ClusteringEvaluator... + >>> evaluator.setWeightCol("weight") + ClusteringEvaluator... + >>> evaluator.evaluate(dataset) + 0.9079... + >>> ce_path = temp_path + "/ce" + >>> evaluator.save(ce_path) + >>> evaluator2 = ClusteringEvaluator.load(ce_path) + >>> str(evaluator2.getPredictionCol()) + 'prediction' + """ + + metricName: Param["ClusteringEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + "metric name in evaluation (silhouette)", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + distanceMeasure: Param["ClusteringEvaluatorDistanceMeasureType"] = Param( + Params._dummy(), + "distanceMeasure", + "The distance measure. " + "Supported options: 'squaredEuclidean' and 'cosine'.", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + predictionCol: str = "prediction", + featuresCol: str = "features", + metricName: "ClusteringEvaluatorMetricType" = "silhouette", + distanceMeasure: str = "squaredEuclidean", + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, predictionCol="prediction", featuresCol="features", \ + metricName="silhouette", distanceMeasure="squaredEuclidean", weightCol=None) + """ + super(ClusteringEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.ClusteringEvaluator", self.uid + ) + self._setDefault(metricName="silhouette", distanceMeasure="squaredEuclidean") + kwargs = self._input_kwargs + self._set(**kwargs) + + @keyword_only + @since("2.3.0") + def setParams( + self, + *, + predictionCol: str = "prediction", + featuresCol: str = "features", + metricName: "ClusteringEvaluatorMetricType" = "silhouette", + distanceMeasure: str = "squaredEuclidean", + weightCol: Optional[str] = None, + ) -> "ClusteringEvaluator": + """ + setParams(self, \\*, predictionCol="prediction", featuresCol="features", \ + metricName="silhouette", distanceMeasure="squaredEuclidean", weightCol=None) + Sets params for clustering evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.3.0") + def setMetricName(self, value: "ClusteringEvaluatorMetricType") -> "ClusteringEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("2.3.0") + def getMetricName(self) -> "ClusteringEvaluatorMetricType": + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("2.4.0") + def setDistanceMeasure( + self, value: "ClusteringEvaluatorDistanceMeasureType" + ) -> "ClusteringEvaluator": + """ + Sets the value of :py:attr:`distanceMeasure`. + """ + return self._set(distanceMeasure=value) + + @since("2.4.0") + def getDistanceMeasure(self) -> "ClusteringEvaluatorDistanceMeasureType": + """ + Gets the value of `distanceMeasure` + """ + return self.getOrDefault(self.distanceMeasure) + + def setFeaturesCol(self, value: "str") -> "ClusteringEvaluator": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setPredictionCol(self, value: str) -> "ClusteringEvaluator": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("3.1.0") + def setWeightCol(self, value: str) -> "ClusteringEvaluator": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + +@inherit_doc +class RankingEvaluator( + JavaEvaluator, HasLabelCol, HasPredictionCol, JavaMLReadable["RankingEvaluator"], JavaMLWritable +): + """ + Evaluator for Ranking, which expects two input + columns: prediction and label. + + .. versionadded:: 3.0.0 + + Notes + ----- + Experimental + + Examples + -------- + >>> scoreAndLabels = [([1.0, 6.0, 2.0, 7.0, 8.0, 3.0, 9.0, 10.0, 4.0, 5.0], + ... [1.0, 2.0, 3.0, 4.0, 5.0]), + ... ([4.0, 1.0, 5.0, 6.0, 2.0, 7.0, 3.0, 8.0, 9.0, 10.0], [1.0, 2.0, 3.0]), + ... ([1.0, 2.0, 3.0, 4.0, 5.0], [])] + >>> dataset = spark.createDataFrame(scoreAndLabels, ["prediction", "label"]) + ... + >>> evaluator = RankingEvaluator() + >>> evaluator.setPredictionCol("prediction") + RankingEvaluator... + >>> evaluator.evaluate(dataset) + 0.35... + >>> evaluator.evaluate(dataset, {evaluator.metricName: "precisionAtK", evaluator.k: 2}) + 0.33... + >>> ranke_path = temp_path + "/ranke" + >>> evaluator.save(ranke_path) + >>> evaluator2 = RankingEvaluator.load(ranke_path) + >>> str(evaluator2.getPredictionCol()) + 'prediction' + """ + + metricName: Param["RankingEvaluatorMetricType"] = Param( + Params._dummy(), + "metricName", + "metric name in evaluation " + "(meanAveragePrecision|meanAveragePrecisionAtK|" + "precisionAtK|ndcgAtK|recallAtK)", + typeConverter=TypeConverters.toString, # type: ignore[arg-type] + ) + k: Param[int] = Param( + Params._dummy(), + "k", + "The ranking position value used in meanAveragePrecisionAtK|precisionAtK|" + "ndcgAtK|recallAtK. Must be > 0. The default value is 10.", + typeConverter=TypeConverters.toInt, + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "RankingEvaluatorMetricType" = "meanAveragePrecision", + k: int = 10, + ): + """ + __init__(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="meanAveragePrecision", k=10) + """ + super(RankingEvaluator, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.evaluation.RankingEvaluator", self.uid + ) + self._setDefault(metricName="meanAveragePrecision", k=10) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("3.0.0") + def setMetricName(self, value: "RankingEvaluatorMetricType") -> "RankingEvaluator": + """ + Sets the value of :py:attr:`metricName`. + """ + return self._set(metricName=value) + + @since("3.0.0") + def getMetricName(self) -> "RankingEvaluatorMetricType": + """ + Gets the value of metricName or its default value. + """ + return self.getOrDefault(self.metricName) + + @since("3.0.0") + def setK(self, value: int) -> "RankingEvaluator": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + @since("3.0.0") + def getK(self) -> int: + """ + Gets the value of k or its default value. + """ + return self.getOrDefault(self.k) + + @since("3.0.0") + def setLabelCol(self, value: str) -> "RankingEvaluator": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "RankingEvaluator": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @keyword_only + @since("3.0.0") + def setParams( + self, + *, + predictionCol: str = "prediction", + labelCol: str = "label", + metricName: "RankingEvaluatorMetricType" = "meanAveragePrecision", + k: int = 10, + ) -> "RankingEvaluator": + """ + setParams(self, \\*, predictionCol="prediction", labelCol="label", \ + metricName="meanAveragePrecision", k=10) + Sets params for ranking evaluator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +if __name__ == "__main__": + import doctest + import tempfile + import pyspark.ml.evaluation + from pyspark.sql import SparkSession + + globs = pyspark.ml.evaluation.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.evaluation tests").getOrCreate() + globs["spark"] = spark + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/feature.py new file mode 100644 index 0000000..349b509 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/feature.py @@ -0,0 +1,7474 @@ +# +# 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 ( + cast, + overload, + Any, + Dict, + Generic, + List, + Optional, + Tuple, + TypeVar, + Union, + TYPE_CHECKING, +) + +from pyspark import keyword_only, since, SparkContext +from pyspark.ml.linalg import _convert_to_vector, DenseMatrix, DenseVector, Vector +from pyspark.sql.dataframe import DataFrame +from pyspark.ml.param.shared import ( + HasThreshold, + HasThresholds, + HasInputCol, + HasOutputCol, + HasInputCols, + HasOutputCols, + HasHandleInvalid, + HasRelativeError, + HasFeaturesCol, + HasLabelCol, + HasSeed, + HasNumFeatures, + HasStepSize, + HasMaxIter, + TypeConverters, + Param, + Params, +) +from pyspark.ml.util import JavaMLReadable, JavaMLWritable +from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams, JavaTransformer, _jvm +from pyspark.ml.common import inherit_doc + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + +JM = TypeVar("JM", bound=JavaTransformer) +P = TypeVar("P", bound=Params) + +__all__ = [ + "Binarizer", + "BucketedRandomProjectionLSH", + "BucketedRandomProjectionLSHModel", + "Bucketizer", + "ChiSqSelector", + "ChiSqSelectorModel", + "CountVectorizer", + "CountVectorizerModel", + "DCT", + "ElementwiseProduct", + "FeatureHasher", + "HashingTF", + "IDF", + "IDFModel", + "Imputer", + "ImputerModel", + "IndexToString", + "Interaction", + "MaxAbsScaler", + "MaxAbsScalerModel", + "MinHashLSH", + "MinHashLSHModel", + "MinMaxScaler", + "MinMaxScalerModel", + "NGram", + "Normalizer", + "OneHotEncoder", + "OneHotEncoderModel", + "PCA", + "PCAModel", + "PolynomialExpansion", + "QuantileDiscretizer", + "RobustScaler", + "RobustScalerModel", + "RegexTokenizer", + "RFormula", + "RFormulaModel", + "SQLTransformer", + "StandardScaler", + "StandardScalerModel", + "StopWordsRemover", + "StringIndexer", + "StringIndexerModel", + "Tokenizer", + "UnivariateFeatureSelector", + "UnivariateFeatureSelectorModel", + "VarianceThresholdSelector", + "VarianceThresholdSelectorModel", + "VectorAssembler", + "VectorIndexer", + "VectorIndexerModel", + "VectorSizeHint", + "VectorSlicer", + "Word2Vec", + "Word2VecModel", +] + + +@inherit_doc +class Binarizer( + JavaTransformer, + HasThreshold, + HasThresholds, + HasInputCol, + HasOutputCol, + HasInputCols, + HasOutputCols, + JavaMLReadable["Binarizer"], + JavaMLWritable, +): + """ + Binarize a column of continuous features given a threshold. Since 3.0.0, + :py:class:`Binarize` can map multiple columns at once by setting the :py:attr:`inputCols` + parameter. Note that when both the :py:attr:`inputCol` and :py:attr:`inputCols` parameters + are set, an Exception will be thrown. The :py:attr:`threshold` parameter is used for + single column usage, and :py:attr:`thresholds` is for multiple columns. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> df = spark.createDataFrame([(0.5,)], ["values"]) + >>> binarizer = Binarizer(threshold=1.0, inputCol="values", outputCol="features") + >>> binarizer.setThreshold(1.0) + Binarizer... + >>> binarizer.setInputCol("values") + Binarizer... + >>> binarizer.setOutputCol("features") + Binarizer... + >>> binarizer.transform(df).head().features + 0.0 + >>> binarizer.setParams(outputCol="freqs").transform(df).head().freqs + 0.0 + >>> params = {binarizer.threshold: -0.5, binarizer.outputCol: "vector"} + >>> binarizer.transform(df, params).head().vector + 1.0 + >>> binarizerPath = temp_path + "/binarizer" + >>> binarizer.save(binarizerPath) + >>> loadedBinarizer = Binarizer.load(binarizerPath) + >>> loadedBinarizer.getThreshold() == binarizer.getThreshold() + True + >>> loadedBinarizer.transform(df).take(1) == binarizer.transform(df).take(1) + True + >>> df2 = spark.createDataFrame([(0.5, 0.3)], ["values1", "values2"]) + >>> binarizer2 = Binarizer(thresholds=[0.0, 1.0]) + >>> binarizer2.setInputCols(["values1", "values2"]).setOutputCols(["output1", "output2"]) + Binarizer... + >>> binarizer2.transform(df2).show() + +-------+-------+-------+-------+ + |values1|values2|output1|output2| + +-------+-------+-------+-------+ + | 0.5| 0.3| 1.0| 0.0| + +-------+-------+-------+-------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + threshold: Param[float] = Param( + Params._dummy(), + "threshold", + "Param for threshold used to binarize continuous features. " + + "The features greater than the threshold will be binarized to 1.0. " + + "The features equal to or less than the threshold will be binarized to 0.0", + typeConverter=TypeConverters.toFloat, + ) + thresholds: Param[List[float]] = Param( + Params._dummy(), + "thresholds", + "Param for array of threshold used to binarize continuous features. " + + "This is for multiple columns input. If transforming multiple columns " + + "and thresholds is not set, but threshold is set, then threshold will " + + "be applied across all columns.", + typeConverter=TypeConverters.toListFloat, + ) + + @overload + def __init__( + self, + *, + threshold: float = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + ): + ... + + @overload + def __init__( + self, + *, + thresholds: Optional[List[float]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + threshold: float = 0.0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + thresholds: Optional[List[float]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, threshold=0.0, inputCol=None, outputCol=None, thresholds=None, \ + inputCols=None, outputCols=None) + """ + super(Binarizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Binarizer", self.uid) + self._setDefault(threshold=0.0) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + threshold: float = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + ) -> "Binarizer": + ... + + @overload + def setParams( + self, + *, + thresholds: Optional[List[float]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ) -> "Binarizer": + ... + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + threshold: float = 0.0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + thresholds: Optional[List[float]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ) -> "Binarizer": + """ + setParams(self, \\*, threshold=0.0, inputCol=None, outputCol=None, thresholds=None, \ + inputCols=None, outputCols=None) + Sets params for this Binarizer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setThreshold(self, value: float) -> "Binarizer": + """ + Sets the value of :py:attr:`threshold`. + """ + return self._set(threshold=value) + + @since("3.0.0") + def setThresholds(self, value: List[float]) -> "Binarizer": + """ + Sets the value of :py:attr:`thresholds`. + """ + return self._set(thresholds=value) + + def setInputCol(self, value: str) -> "Binarizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "Binarizer": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "Binarizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "Binarizer": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + +class _LSHParams(HasInputCol, HasOutputCol): + """ + Mixin for Locality Sensitive Hashing (LSH) algorithm parameters. + """ + + numHashTables: Param[int] = Param( + Params._dummy(), + "numHashTables", + "number of hash tables, where " + + "increasing number of hash tables lowers the false negative rate, " + + "and decreasing it improves the running performance.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self, *args: Any): + super(_LSHParams, self).__init__(*args) + self._setDefault(numHashTables=1) + + def getNumHashTables(self) -> int: + """ + Gets the value of numHashTables or its default value. + """ + return self.getOrDefault(self.numHashTables) + + +class _LSH(JavaEstimator[JM], _LSHParams, JavaMLReadable, JavaMLWritable, Generic[JM]): + """ + Mixin for Locality Sensitive Hashing (LSH). + """ + + def setNumHashTables(self: P, value: int) -> P: + """ + Sets the value of :py:attr:`numHashTables`. + """ + return self._set(numHashTables=value) + + def setInputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +class _LSHModel(JavaModel, _LSHParams): + """ + Mixin for Locality Sensitive Hashing (LSH) models. + """ + + def setInputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def approxNearestNeighbors( + self, + dataset: DataFrame, + key: Vector, + numNearestNeighbors: int, + distCol: str = "distCol", + ) -> DataFrame: + """ + Given a large dataset and an item, approximately find at most k items which have the + closest distance to the item. If the :py:attr:`outputCol` is missing, the method will + transform the data; if the :py:attr:`outputCol` exists, it will use that. This allows + caching of the transformed data when necessary. + + Notes + ----- + This method is experimental and will likely change behavior in the next release. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + The dataset to search for nearest neighbors of the key. + key : :py:class:`pyspark.ml.linalg.Vector` + Feature vector representing the item to search for. + numNearestNeighbors : int + The maximum number of nearest neighbors. + distCol : str + Output column for storing the distance between each result row and the key. + Use "distCol" as default value if it's not specified. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + A dataset containing at most k items closest to the key. A column "distCol" is + added to show the distance between each row and the key. + """ + return self._call_java("approxNearestNeighbors", dataset, key, numNearestNeighbors, distCol) + + def approxSimilarityJoin( + self, + datasetA: DataFrame, + datasetB: DataFrame, + threshold: float, + distCol: str = "distCol", + ) -> DataFrame: + """ + Join two datasets to approximately find all pairs of rows whose distance are smaller than + the threshold. If the :py:attr:`outputCol` is missing, the method will transform the data; + if the :py:attr:`outputCol` exists, it will use that. This allows caching of the + transformed data when necessary. + + Parameters + ---------- + datasetA : :py:class:`pyspark.sql.DataFrame` + One of the datasets to join. + datasetB : :py:class:`pyspark.sql.DataFrame` + Another dataset to join. + threshold : float + The threshold for the distance of row pairs. + distCol : str, optional + Output column for storing the distance between each pair of rows. Use + "distCol" as default value if it's not specified. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + A joined dataset containing pairs of rows. The original rows are in columns + "datasetA" and "datasetB", and a column "distCol" is added to show the distance + between each pair. + """ + threshold = TypeConverters.toFloat(threshold) + return self._call_java("approxSimilarityJoin", datasetA, datasetB, threshold, distCol) + + +class _BucketedRandomProjectionLSHParams: + """ + Params for :py:class:`BucketedRandomProjectionLSH` and + :py:class:`BucketedRandomProjectionLSHModel`. + + .. versionadded:: 3.0.0 + """ + + bucketLength: Param[float] = Param( + Params._dummy(), + "bucketLength", + "the length of each hash bucket, " + "a larger bucket lowers the false negative rate.", + typeConverter=TypeConverters.toFloat, + ) + + @since("2.2.0") + def getBucketLength(self) -> float: + """ + Gets the value of bucketLength or its default value. + """ + return (cast(Params, self)).getOrDefault(self.bucketLength) + + +@inherit_doc +class BucketedRandomProjectionLSH( + _LSH["BucketedRandomProjectionLSHModel"], + _LSHParams, + _BucketedRandomProjectionLSHParams, + HasSeed, + JavaMLReadable["BucketedRandomProjectionLSH"], + JavaMLWritable, +): + """ + LSH class for Euclidean distance metrics. + The input is dense or sparse vectors, each of which represents a point in the Euclidean + distance space. The output will be vectors of configurable dimension. Hash values in the same + dimension are calculated by the same hash function. + + .. versionadded:: 2.2.0 + + Notes + ----- + + - `Stable Distributions in Wikipedia article on Locality-sensitive hashing \ + `_ + - `Hashing for Similarity Search: A Survey `_ + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.sql.functions import col + >>> data = [(0, Vectors.dense([-1.0, -1.0 ]),), + ... (1, Vectors.dense([-1.0, 1.0 ]),), + ... (2, Vectors.dense([1.0, -1.0 ]),), + ... (3, Vectors.dense([1.0, 1.0]),)] + >>> df = spark.createDataFrame(data, ["id", "features"]) + >>> brp = BucketedRandomProjectionLSH() + >>> brp.setInputCol("features") + BucketedRandomProjectionLSH... + >>> brp.setOutputCol("hashes") + BucketedRandomProjectionLSH... + >>> brp.setSeed(12345) + BucketedRandomProjectionLSH... + >>> brp.setBucketLength(1.0) + BucketedRandomProjectionLSH... + >>> model = brp.fit(df) + >>> model.getBucketLength() + 1.0 + >>> model.setOutputCol("hashes") + BucketedRandomProjectionLSHModel... + >>> model.transform(df).head() + Row(id=0, features=DenseVector([-1.0, -1.0]), hashes=[DenseVector([-1.0])]) + >>> data2 = [(4, Vectors.dense([2.0, 2.0 ]),), + ... (5, Vectors.dense([2.0, 3.0 ]),), + ... (6, Vectors.dense([3.0, 2.0 ]),), + ... (7, Vectors.dense([3.0, 3.0]),)] + >>> df2 = spark.createDataFrame(data2, ["id", "features"]) + >>> model.approxNearestNeighbors(df2, Vectors.dense([1.0, 2.0]), 1).collect() + [Row(id=4, features=DenseVector([2.0, 2.0]), hashes=[DenseVector([1.0])], distCol=1.0)] + >>> model.approxSimilarityJoin(df, df2, 3.0, distCol="EuclideanDistance").select( + ... col("datasetA.id").alias("idA"), + ... col("datasetB.id").alias("idB"), + ... col("EuclideanDistance")).show() + +---+---+-----------------+ + |idA|idB|EuclideanDistance| + +---+---+-----------------+ + | 3| 6| 2.23606797749979| + +---+---+-----------------+ + ... + >>> model.approxSimilarityJoin(df, df2, 3, distCol="EuclideanDistance").select( + ... col("datasetA.id").alias("idA"), + ... col("datasetB.id").alias("idB"), + ... col("EuclideanDistance")).show() + +---+---+-----------------+ + |idA|idB|EuclideanDistance| + +---+---+-----------------+ + | 3| 6| 2.23606797749979| + +---+---+-----------------+ + ... + >>> brpPath = temp_path + "/brp" + >>> brp.save(brpPath) + >>> brp2 = BucketedRandomProjectionLSH.load(brpPath) + >>> brp2.getBucketLength() == brp.getBucketLength() + True + >>> modelPath = temp_path + "/brp-model" + >>> model.save(modelPath) + >>> model2 = BucketedRandomProjectionLSHModel.load(modelPath) + >>> model.transform(df).head().hashes == model2.transform(df).head().hashes + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + seed: Optional[int] = None, + numHashTables: int = 1, + bucketLength: Optional[float] = None, + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, seed=None, numHashTables=1, \ + bucketLength=None) + """ + super(BucketedRandomProjectionLSH, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.BucketedRandomProjectionLSH", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.2.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + seed: Optional[int] = None, + numHashTables: int = 1, + bucketLength: Optional[float] = None, + ) -> "BucketedRandomProjectionLSH": + """ + setParams(self, \\*, inputCol=None, outputCol=None, seed=None, numHashTables=1, \ + bucketLength=None) + Sets params for this BucketedRandomProjectionLSH. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.2.0") + def setBucketLength(self, value: float) -> "BucketedRandomProjectionLSH": + """ + Sets the value of :py:attr:`bucketLength`. + """ + return self._set(bucketLength=value) + + def setSeed(self, value: int) -> "BucketedRandomProjectionLSH": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def _create_model(self, java_model: "JavaObject") -> "BucketedRandomProjectionLSHModel": + return BucketedRandomProjectionLSHModel(java_model) + + +class BucketedRandomProjectionLSHModel( + _LSHModel, + _BucketedRandomProjectionLSHParams, + JavaMLReadable["BucketedRandomProjectionLSHModel"], + JavaMLWritable, +): + r""" + Model fitted by :py:class:`BucketedRandomProjectionLSH`, where multiple random vectors are + stored. The vectors are normalized to be unit vectors and each vector is used in a hash + function: :math:`h_i(x) = floor(r_i \cdot x / bucketLength)` where :math:`r_i` is the + i-th random unit vector. The number of buckets will be `(max L2 norm of input vectors) / + bucketLength`. + + .. versionadded:: 2.2.0 + """ + + +@inherit_doc +class Bucketizer( + JavaTransformer, + HasInputCol, + HasOutputCol, + HasInputCols, + HasOutputCols, + HasHandleInvalid, + JavaMLReadable["Bucketizer"], + JavaMLWritable, +): + """ + Maps a column of continuous features to a column of feature buckets. Since 3.0.0, + :py:class:`Bucketizer` can map multiple columns at once by setting the :py:attr:`inputCols` + parameter. Note that when both the :py:attr:`inputCol` and :py:attr:`inputCols` parameters + are set, an Exception will be thrown. The :py:attr:`splits` parameter is only used for single + column usage, and :py:attr:`splitsArray` is for multiple columns. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> values = [(0.1, 0.0), (0.4, 1.0), (1.2, 1.3), (1.5, float("nan")), + ... (float("nan"), 1.0), (float("nan"), 0.0)] + >>> df = spark.createDataFrame(values, ["values1", "values2"]) + >>> bucketizer = Bucketizer() + >>> bucketizer.setSplits([-float("inf"), 0.5, 1.4, float("inf")]) + Bucketizer... + >>> bucketizer.setInputCol("values1") + Bucketizer... + >>> bucketizer.setOutputCol("buckets") + Bucketizer... + >>> bucketed = bucketizer.setHandleInvalid("keep").transform(df).collect() + >>> bucketed = bucketizer.setHandleInvalid("keep").transform(df.select("values1")) + >>> bucketed.show(truncate=False) + +-------+-------+ + |values1|buckets| + +-------+-------+ + |0.1 |0.0 | + |0.4 |0.0 | + |1.2 |1.0 | + |1.5 |2.0 | + |NaN |3.0 | + |NaN |3.0 | + +-------+-------+ + ... + >>> bucketizer.setParams(outputCol="b").transform(df).head().b + 0.0 + >>> bucketizerPath = temp_path + "/bucketizer" + >>> bucketizer.save(bucketizerPath) + >>> loadedBucketizer = Bucketizer.load(bucketizerPath) + >>> loadedBucketizer.getSplits() == bucketizer.getSplits() + True + >>> loadedBucketizer.transform(df).take(1) == bucketizer.transform(df).take(1) + True + >>> bucketed = bucketizer.setHandleInvalid("skip").transform(df).collect() + >>> len(bucketed) + 4 + >>> bucketizer2 = Bucketizer(splitsArray= + ... [[-float("inf"), 0.5, 1.4, float("inf")], [-float("inf"), 0.5, float("inf")]], + ... inputCols=["values1", "values2"], outputCols=["buckets1", "buckets2"]) + >>> bucketed2 = bucketizer2.setHandleInvalid("keep").transform(df) + >>> bucketed2.show(truncate=False) + +-------+-------+--------+--------+ + |values1|values2|buckets1|buckets2| + +-------+-------+--------+--------+ + |0.1 |0.0 |0.0 |0.0 | + |0.4 |1.0 |0.0 |1.0 | + |1.2 |1.3 |1.0 |1.0 | + |1.5 |NaN |2.0 |2.0 | + |NaN |1.0 |3.0 |1.0 | + |NaN |0.0 |3.0 |0.0 | + +-------+-------+--------+--------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + splits: Param[List[float]] = Param( + Params._dummy(), + "splits", + "Split points for mapping continuous features into buckets. With n+1 splits, " + + "there are n buckets. A bucket defined by splits x,y holds values in the " + + "range [x,y) except the last bucket, which also includes y. The splits " + + "should be of length >= 3 and strictly increasing. Values at -inf, inf must be " + + "explicitly provided to cover all Double values; otherwise, values outside the " + + "splits specified will be treated as errors.", + typeConverter=TypeConverters.toListFloat, + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "how to handle invalid entries " + "containing NaN values. Values outside the splits will always be treated " + "as errors. Options are 'skip' (filter out rows with invalid values), " + + "'error' (throw an error), or 'keep' (keep invalid values in a " + + "special additional bucket). Note that in the multiple column " + + "case, the invalid handling is applied to all columns. That said " + + "for 'error' it will throw an error if any invalids are found in " + + "any column, for 'skip' it will skip rows with any invalids in " + + "any columns, etc.", + typeConverter=TypeConverters.toString, + ) + + splitsArray: Param[List[List[float]]] = Param( + Params._dummy(), + "splitsArray", + "The array of split points for mapping " + + "continuous features into buckets for multiple columns. For each input " + + "column, with n+1 splits, there are n buckets. A bucket defined by " + + "splits x,y holds values in the range [x,y) except the last bucket, " + + "which also includes y. The splits should be of length >= 3 and " + + "strictly increasing. Values at -inf, inf must be explicitly provided " + + "to cover all Double values; otherwise, values outside the splits " + + "specified will be treated as errors.", + typeConverter=TypeConverters.toListListFloat, + ) + + @overload + def __init__( + self, + *, + splits: Optional[List[float]] = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + handleInvalid: str = ..., + ): + ... + + @overload + def __init__( + self, + *, + handleInvalid: str = ..., + splitsArray: Optional[List[List[float]]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + splits: Optional[List[float]] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + splitsArray: Optional[List[List[float]]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, splits=None, inputCol=None, outputCol=None, handleInvalid="error", \ + splitsArray=None, inputCols=None, outputCols=None) + """ + super(Bucketizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Bucketizer", self.uid) + self._setDefault(handleInvalid="error") + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + splits: Optional[List[float]] = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + handleInvalid: str = ..., + ) -> "Bucketizer": + ... + + @overload + def setParams( + self, + *, + handleInvalid: str = ..., + splitsArray: Optional[List[List[float]]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ) -> "Bucketizer": + ... + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + splits: Optional[List[float]] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + splitsArray: Optional[List[List[float]]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ) -> "Bucketizer": + """ + setParams(self, \\*, splits=None, inputCol=None, outputCol=None, handleInvalid="error", \ + splitsArray=None, inputCols=None, outputCols=None) + Sets params for this Bucketizer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setSplits(self, value: List[float]) -> "Bucketizer": + """ + Sets the value of :py:attr:`splits`. + """ + return self._set(splits=value) + + @since("1.4.0") + def getSplits(self) -> List[float]: + """ + Gets the value of threshold or its default value. + """ + return self.getOrDefault(self.splits) + + @since("3.0.0") + def setSplitsArray(self, value: List[List[float]]) -> "Bucketizer": + """ + Sets the value of :py:attr:`splitsArray`. + """ + return self._set(splitsArray=value) + + @since("3.0.0") + def getSplitsArray(self) -> List[List[float]]: + """ + Gets the array of split points or its default value. + """ + return self.getOrDefault(self.splitsArray) + + def setInputCol(self, value: str) -> "Bucketizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "Bucketizer": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "Bucketizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "Bucketizer": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + def setHandleInvalid(self, value: str) -> "Bucketizer": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + +class _CountVectorizerParams(JavaParams, HasInputCol, HasOutputCol): + """ + Params for :py:class:`CountVectorizer` and :py:class:`CountVectorizerModel`. + """ + + minTF: Param[float] = Param( + Params._dummy(), + "minTF", + "Filter to ignore rare words in" + + " a document. For each document, terms with frequency/count less than the given" + + " threshold are ignored. If this is an integer >= 1, then this specifies a count (of" + + " times the term must appear in the document); if this is a double in [0,1), then this " + + "specifies a fraction (out of the document's token count). Note that the parameter is " + + "only used in transform of CountVectorizerModel and does not affect fitting. Default 1.0", + typeConverter=TypeConverters.toFloat, + ) + minDF: Param[float] = Param( + Params._dummy(), + "minDF", + "Specifies the minimum number of" + + " different documents a term must appear in to be included in the vocabulary." + + " If this is an integer >= 1, this specifies the number of documents the term must" + + " appear in; if this is a double in [0,1), then this specifies the fraction of documents." + + " Default 1.0", + typeConverter=TypeConverters.toFloat, + ) + maxDF: Param[float] = Param( + Params._dummy(), + "maxDF", + "Specifies the maximum number of" + + " different documents a term could appear in to be included in the vocabulary." + + " A term that appears more than the threshold will be ignored. If this is an" + + " integer >= 1, this specifies the maximum number of documents the term could appear in;" + + " if this is a double in [0,1), then this specifies the maximum" + + " fraction of documents the term could appear in." + + " Default (2^63) - 1", + typeConverter=TypeConverters.toFloat, + ) + vocabSize: Param[int] = Param( + Params._dummy(), + "vocabSize", + "max size of the vocabulary. Default 1 << 18.", + typeConverter=TypeConverters.toInt, + ) + binary: Param[bool] = Param( + Params._dummy(), + "binary", + "Binary toggle to control the output vector values." + + " If True, all nonzero counts (after minTF filter applied) are set to 1. This is useful" + + " for discrete probabilistic models that model binary events rather than integer counts." + + " Default False", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self, *args: Any): + super(_CountVectorizerParams, self).__init__(*args) + self._setDefault(minTF=1.0, minDF=1.0, maxDF=2**63 - 1, vocabSize=1 << 18, binary=False) + + @since("1.6.0") + def getMinTF(self) -> float: + """ + Gets the value of minTF or its default value. + """ + return self.getOrDefault(self.minTF) + + @since("1.6.0") + def getMinDF(self) -> float: + """ + Gets the value of minDF or its default value. + """ + return self.getOrDefault(self.minDF) + + @since("2.4.0") + def getMaxDF(self) -> float: + """ + Gets the value of maxDF or its default value. + """ + return self.getOrDefault(self.maxDF) + + @since("1.6.0") + def getVocabSize(self) -> int: + """ + Gets the value of vocabSize or its default value. + """ + return self.getOrDefault(self.vocabSize) + + @since("2.0.0") + def getBinary(self) -> bool: + """ + Gets the value of binary or its default value. + """ + return self.getOrDefault(self.binary) + + +@inherit_doc +class CountVectorizer( + JavaEstimator["CountVectorizerModel"], + _CountVectorizerParams, + JavaMLReadable["CountVectorizer"], + JavaMLWritable, +): + """ + Extracts a vocabulary from document collections and generates a :py:attr:`CountVectorizerModel`. + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(0, ["a", "b", "c"]), (1, ["a", "b", "b", "c", "a"])], + ... ["label", "raw"]) + >>> cv = CountVectorizer() + >>> cv.setInputCol("raw") + CountVectorizer... + >>> cv.setOutputCol("vectors") + CountVectorizer... + >>> model = cv.fit(df) + >>> model.setInputCol("raw") + CountVectorizerModel... + >>> model.transform(df).show(truncate=False) + +-----+---------------+-------------------------+ + |label|raw |vectors | + +-----+---------------+-------------------------+ + |0 |[a, b, c] |(3,[0,1,2],[1.0,1.0,1.0])| + |1 |[a, b, b, c, a]|(3,[0,1,2],[2.0,2.0,1.0])| + +-----+---------------+-------------------------+ + ... + >>> sorted(model.vocabulary) == ['a', 'b', 'c'] + True + >>> countVectorizerPath = temp_path + "/count-vectorizer" + >>> cv.save(countVectorizerPath) + >>> loadedCv = CountVectorizer.load(countVectorizerPath) + >>> loadedCv.getMinDF() == cv.getMinDF() + True + >>> loadedCv.getMinTF() == cv.getMinTF() + True + >>> loadedCv.getVocabSize() == cv.getVocabSize() + True + >>> modelPath = temp_path + "/count-vectorizer-model" + >>> model.save(modelPath) + >>> loadedModel = CountVectorizerModel.load(modelPath) + >>> loadedModel.vocabulary == model.vocabulary + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + >>> fromVocabModel = CountVectorizerModel.from_vocabulary(["a", "b", "c"], + ... inputCol="raw", outputCol="vectors") + >>> fromVocabModel.transform(df).show(truncate=False) + +-----+---------------+-------------------------+ + |label|raw |vectors | + +-----+---------------+-------------------------+ + |0 |[a, b, c] |(3,[0,1,2],[1.0,1.0,1.0])| + |1 |[a, b, b, c, a]|(3,[0,1,2],[2.0,2.0,1.0])| + +-----+---------------+-------------------------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + minTF: float = 1.0, + minDF: float = 1.0, + maxDF: float = 2**63 - 1, + vocabSize: int = 1 << 18, + binary: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, minTF=1.0, minDF=1.0, maxDF=2 ** 63 - 1, vocabSize=1 << 18,\ + binary=False, inputCol=None,outputCol=None) + """ + super(CountVectorizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.CountVectorizer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + minTF: float = 1.0, + minDF: float = 1.0, + maxDF: float = 2**63 - 1, + vocabSize: int = 1 << 18, + binary: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "CountVectorizer": + """ + setParams(self, \\*, minTF=1.0, minDF=1.0, maxDF=2 ** 63 - 1, vocabSize=1 << 18,\ + binary=False, inputCol=None, outputCol=None) + Set the params for the CountVectorizer + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setMinTF(self, value: float) -> "CountVectorizer": + """ + Sets the value of :py:attr:`minTF`. + """ + return self._set(minTF=value) + + @since("1.6.0") + def setMinDF(self, value: float) -> "CountVectorizer": + """ + Sets the value of :py:attr:`minDF`. + """ + return self._set(minDF=value) + + @since("2.4.0") + def setMaxDF(self, value: float) -> "CountVectorizer": + """ + Sets the value of :py:attr:`maxDF`. + """ + return self._set(maxDF=value) + + @since("1.6.0") + def setVocabSize(self, value: int) -> "CountVectorizer": + """ + Sets the value of :py:attr:`vocabSize`. + """ + return self._set(vocabSize=value) + + @since("2.0.0") + def setBinary(self, value: bool) -> "CountVectorizer": + """ + Sets the value of :py:attr:`binary`. + """ + return self._set(binary=value) + + def setInputCol(self, value: str) -> "CountVectorizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "CountVectorizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "CountVectorizerModel": + return CountVectorizerModel(java_model) + + +@inherit_doc +class CountVectorizerModel( + JavaModel, _CountVectorizerParams, JavaMLReadable["CountVectorizerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`CountVectorizer`. + + .. versionadded:: 1.6.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "CountVectorizerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "CountVectorizerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @classmethod + @since("2.4.0") + def from_vocabulary( + cls, + vocabulary: List[str], + inputCol: str, + outputCol: Optional[str] = None, + minTF: Optional[float] = None, + binary: Optional[bool] = None, + ) -> "CountVectorizerModel": + """ + Construct the model directly from a vocabulary list of strings, + requires an active SparkContext. + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._gateway is not None + java_class = sc._gateway.jvm.java.lang.String + jvocab = CountVectorizerModel._new_java_array(vocabulary, java_class) + model = CountVectorizerModel._create_from_java_class( + "org.apache.spark.ml.feature.CountVectorizerModel", jvocab + ) + model.setInputCol(inputCol) + if outputCol is not None: + model.setOutputCol(outputCol) + if minTF is not None: + model.setMinTF(minTF) + if binary is not None: + model.setBinary(binary) + model._set(vocabSize=len(vocabulary)) + return model + + @property + @since("1.6.0") + def vocabulary(self) -> List[str]: + """ + An array of terms in the vocabulary. + """ + return self._call_java("vocabulary") + + @since("2.4.0") + def setMinTF(self, value: float) -> "CountVectorizerModel": + """ + Sets the value of :py:attr:`minTF`. + """ + return self._set(minTF=value) + + @since("2.4.0") + def setBinary(self, value: bool) -> "CountVectorizerModel": + """ + Sets the value of :py:attr:`binary`. + """ + return self._set(binary=value) + + +@inherit_doc +class DCT(JavaTransformer, HasInputCol, HasOutputCol, JavaMLReadable["DCT"], JavaMLWritable): + """ + A feature transformer that takes the 1D discrete cosine transform + of a real vector. No zero padding is performed on the input vector. + It returns a real vector of the same length representing the DCT. + The return vector is scaled such that the transform matrix is + unitary (aka scaled DCT-II). + + .. versionadded:: 1.6.0 + + Notes + ----- + `More information on Wikipedia \ + `_. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df1 = spark.createDataFrame([(Vectors.dense([5.0, 8.0, 6.0]),)], ["vec"]) + >>> dct = DCT( ) + >>> dct.setInverse(False) + DCT... + >>> dct.setInputCol("vec") + DCT... + >>> dct.setOutputCol("resultVec") + DCT... + >>> df2 = dct.transform(df1) + >>> df2.head().resultVec + DenseVector([10.969..., -0.707..., -2.041...]) + >>> df3 = DCT(inverse=True, inputCol="resultVec", outputCol="origVec").transform(df2) + >>> df3.head().origVec + DenseVector([5.0, 8.0, 6.0]) + >>> dctPath = temp_path + "/dct" + >>> dct.save(dctPath) + >>> loadedDtc = DCT.load(dctPath) + >>> loadedDtc.transform(df1).take(1) == dct.transform(df1).take(1) + True + >>> loadedDtc.getInverse() + False + """ + + _input_kwargs: Dict[str, Any] + + inverse: Param[bool] = Param( + Params._dummy(), + "inverse", + "Set transformer to perform inverse DCT, " + "default False.", + typeConverter=TypeConverters.toBoolean, + ) + + @keyword_only + def __init__( + self, + *, + inverse: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, inverse=False, inputCol=None, outputCol=None) + """ + super(DCT, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.DCT", self.uid) + self._setDefault(inverse=False) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + inverse: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "DCT": + """ + setParams(self, \\*, inverse=False, inputCol=None, outputCol=None) + Sets params for this DCT. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setInverse(self, value: bool) -> "DCT": + """ + Sets the value of :py:attr:`inverse`. + """ + return self._set(inverse=value) + + @since("1.6.0") + def getInverse(self) -> bool: + """ + Gets the value of inverse or its default value. + """ + return self.getOrDefault(self.inverse) + + def setInputCol(self, value: str) -> "DCT": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "DCT": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class ElementwiseProduct( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["ElementwiseProduct"], + JavaMLWritable, +): + """ + Outputs the Hadamard product (i.e., the element-wise product) of each input vector + with a provided "weight" vector. In other words, it scales each column of the dataset + by a scalar multiplier. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([2.0, 1.0, 3.0]),)], ["values"]) + >>> ep = ElementwiseProduct() + >>> ep.setScalingVec(Vectors.dense([1.0, 2.0, 3.0])) + ElementwiseProduct... + >>> ep.setInputCol("values") + ElementwiseProduct... + >>> ep.setOutputCol("eprod") + ElementwiseProduct... + >>> ep.transform(df).head().eprod + DenseVector([2.0, 2.0, 9.0]) + >>> ep.setParams(scalingVec=Vectors.dense([2.0, 3.0, 5.0])).transform(df).head().eprod + DenseVector([4.0, 3.0, 15.0]) + >>> elementwiseProductPath = temp_path + "/elementwise-product" + >>> ep.save(elementwiseProductPath) + >>> loadedEp = ElementwiseProduct.load(elementwiseProductPath) + >>> loadedEp.getScalingVec() == ep.getScalingVec() + True + >>> loadedEp.transform(df).take(1) == ep.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + scalingVec: Param[Vector] = Param( + Params._dummy(), + "scalingVec", + "Vector for hadamard product.", + typeConverter=TypeConverters.toVector, + ) + + @keyword_only + def __init__( + self, + *, + scalingVec: Optional[Vector] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, scalingVec=None, inputCol=None, outputCol=None) + """ + super(ElementwiseProduct, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.ElementwiseProduct", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + scalingVec: Optional[Vector] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "ElementwiseProduct": + """ + setParams(self, \\*, scalingVec=None, inputCol=None, outputCol=None) + Sets params for this ElementwiseProduct. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setScalingVec(self, value: Vector) -> "ElementwiseProduct": + """ + Sets the value of :py:attr:`scalingVec`. + """ + return self._set(scalingVec=value) + + @since("2.0.0") + def getScalingVec(self) -> Vector: + """ + Gets the value of scalingVec or its default value. + """ + return self.getOrDefault(self.scalingVec) + + def setInputCol(self, value: str) -> "ElementwiseProduct": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "ElementwiseProduct": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class FeatureHasher( + JavaTransformer, + HasInputCols, + HasOutputCol, + HasNumFeatures, + JavaMLReadable["FeatureHasher"], + JavaMLWritable, +): + """ + Feature hashing projects a set of categorical or numerical features into a feature vector of + specified dimension (typically substantially smaller than that of the original feature + space). This is done using the hashing trick (https://en.wikipedia.org/wiki/Feature_hashing) + to map features to indices in the feature vector. + + The FeatureHasher transformer operates on multiple columns. Each column may contain either + numeric or categorical features. Behavior and handling of column data types is as follows: + + * Numeric columns: + For numeric features, the hash value of the column name is used to map the + feature value to its index in the feature vector. By default, numeric features + are not treated as categorical (even when they are integers). To treat them + as categorical, specify the relevant columns in `categoricalCols`. + + * String columns: + For categorical features, the hash value of the string "column_name=value" + is used to map to the vector index, with an indicator value of `1.0`. + Thus, categorical features are "one-hot" encoded + (similarly to using :py:class:`OneHotEncoder` with `dropLast=false`). + + * Boolean columns: + Boolean values are treated in the same way as string columns. That is, + boolean features are represented as "column_name=true" or "column_name=false", + with an indicator value of `1.0`. + + Null (missing) values are ignored (implicitly zero in the resulting feature vector). + + Since a simple modulo is used to transform the hash function to a vector index, + it is advisable to use a power of two as the `numFeatures` parameter; + otherwise the features will not be mapped evenly to the vector indices. + + .. versionadded:: 2.3.0 + + Examples + -------- + >>> data = [(2.0, True, "1", "foo"), (3.0, False, "2", "bar")] + >>> cols = ["real", "bool", "stringNum", "string"] + >>> df = spark.createDataFrame(data, cols) + >>> hasher = FeatureHasher() + >>> hasher.setInputCols(cols) + FeatureHasher... + >>> hasher.setOutputCol("features") + FeatureHasher... + >>> hasher.transform(df).head().features + SparseVector(262144, {174475: 2.0, 247670: 1.0, 257907: 1.0, 262126: 1.0}) + >>> hasher.setCategoricalCols(["real"]).transform(df).head().features + SparseVector(262144, {171257: 1.0, 247670: 1.0, 257907: 1.0, 262126: 1.0}) + >>> hasherPath = temp_path + "/hasher" + >>> hasher.save(hasherPath) + >>> loadedHasher = FeatureHasher.load(hasherPath) + >>> loadedHasher.getNumFeatures() == hasher.getNumFeatures() + True + >>> loadedHasher.transform(df).head().features == hasher.transform(df).head().features + True + """ + + _input_kwargs: Dict[str, Any] + + categoricalCols: Param[List[str]] = Param( + Params._dummy(), + "categoricalCols", + "numeric columns to treat as categorical", + typeConverter=TypeConverters.toListString, + ) + + @keyword_only + def __init__( + self, + *, + numFeatures: int = 1 << 18, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + categoricalCols: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, numFeatures=1 << 18, inputCols=None, outputCol=None, \ + categoricalCols=None) + """ + super(FeatureHasher, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.FeatureHasher", self.uid) + self._setDefault(numFeatures=1 << 18) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.3.0") + def setParams( + self, + *, + numFeatures: int = 1 << 18, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + categoricalCols: Optional[List[str]] = None, + ) -> "FeatureHasher": + """ + setParams(self, \\*, numFeatures=1 << 18, inputCols=None, outputCol=None, \ + categoricalCols=None) + Sets params for this FeatureHasher. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.3.0") + def setCategoricalCols(self, value: List[str]) -> "FeatureHasher": + """ + Sets the value of :py:attr:`categoricalCols`. + """ + return self._set(categoricalCols=value) + + @since("2.3.0") + def getCategoricalCols(self) -> List[str]: + """ + Gets the value of binary or its default value. + """ + return self.getOrDefault(self.categoricalCols) + + def setInputCols(self, value: List[str]) -> "FeatureHasher": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "FeatureHasher": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setNumFeatures(self, value: int) -> "FeatureHasher": + """ + Sets the value of :py:attr:`numFeatures`. + """ + return self._set(numFeatures=value) + + +@inherit_doc +class HashingTF( + JavaTransformer, + HasInputCol, + HasOutputCol, + HasNumFeatures, + JavaMLReadable["HashingTF"], + JavaMLWritable, +): + """ + Maps a sequence of terms to their term frequencies using the hashing trick. + Currently we use Austin Appleby's MurmurHash 3 algorithm (MurmurHash3_x86_32) + to calculate the hash code value for the term object. + Since a simple modulo is used to transform the hash function to a column index, + it is advisable to use a power of two as the numFeatures parameter; + otherwise the features will not be mapped evenly to the columns. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],)], ["words"]) + >>> hashingTF = HashingTF(inputCol="words", outputCol="features") + >>> hashingTF.setNumFeatures(10) + HashingTF... + >>> hashingTF.transform(df).head().features + SparseVector(10, {5: 1.0, 7: 1.0, 8: 1.0}) + >>> hashingTF.setParams(outputCol="freqs").transform(df).head().freqs + SparseVector(10, {5: 1.0, 7: 1.0, 8: 1.0}) + >>> params = {hashingTF.numFeatures: 5, hashingTF.outputCol: "vector"} + >>> hashingTF.transform(df, params).head().vector + SparseVector(5, {0: 1.0, 2: 1.0, 3: 1.0}) + >>> hashingTFPath = temp_path + "/hashing-tf" + >>> hashingTF.save(hashingTFPath) + >>> loadedHashingTF = HashingTF.load(hashingTFPath) + >>> loadedHashingTF.getNumFeatures() == hashingTF.getNumFeatures() + True + >>> loadedHashingTF.transform(df).take(1) == hashingTF.transform(df).take(1) + True + >>> hashingTF.indexOf("b") + 5 + """ + + _input_kwargs: Dict[str, Any] + + binary: Param[bool] = Param( + Params._dummy(), + "binary", + "If True, all non zero counts are set to 1. " + + "This is useful for discrete probabilistic models that model binary events " + + "rather than integer counts. Default False.", + typeConverter=TypeConverters.toBoolean, + ) + + @keyword_only + def __init__( + self, + *, + numFeatures: int = 1 << 18, + binary: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, numFeatures=1 << 18, binary=False, inputCol=None, outputCol=None) + """ + super(HashingTF, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.HashingTF", self.uid) + self._setDefault(numFeatures=1 << 18, binary=False) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.3.0") + def setParams( + self, + *, + numFeatures: int = 1 << 18, + binary: bool = False, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "HashingTF": + """ + setParams(self, \\*, numFeatures=1 << 18, binary=False, inputCol=None, outputCol=None) + Sets params for this HashingTF. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setBinary(self, value: bool) -> "HashingTF": + """ + Sets the value of :py:attr:`binary`. + """ + return self._set(binary=value) + + @since("2.0.0") + def getBinary(self) -> bool: + """ + Gets the value of binary or its default value. + """ + return self.getOrDefault(self.binary) + + def setInputCol(self, value: str) -> "HashingTF": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "HashingTF": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setNumFeatures(self, value: int) -> "HashingTF": + """ + Sets the value of :py:attr:`numFeatures`. + """ + return self._set(numFeatures=value) + + @since("3.0.0") + def indexOf(self, term: Any) -> int: + """ + Returns the index of the input term. + """ + self._transfer_params_to_java() + assert self._java_obj is not None + return self._java_obj.indexOf(term) + + +class _IDFParams(HasInputCol, HasOutputCol): + """ + Params for :py:class:`IDF` and :py:class:`IDFModel`. + + .. versionadded:: 3.0.0 + """ + + minDocFreq: Param[int] = Param( + Params._dummy(), + "minDocFreq", + "minimum number of documents in which a term should appear for filtering", + typeConverter=TypeConverters.toInt, + ) + + @since("1.4.0") + def getMinDocFreq(self) -> int: + """ + Gets the value of minDocFreq or its default value. + """ + return self.getOrDefault(self.minDocFreq) + + def __init__(self, *args: Any): + super(_IDFParams, self).__init__(*args) + self._setDefault(minDocFreq=0) + + +@inherit_doc +class IDF(JavaEstimator["IDFModel"], _IDFParams, JavaMLReadable["IDF"], JavaMLWritable): + """ + Compute the Inverse Document Frequency (IDF) given a collection of documents. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import DenseVector + >>> df = spark.createDataFrame([(DenseVector([1.0, 2.0]),), + ... (DenseVector([0.0, 1.0]),), (DenseVector([3.0, 0.2]),)], ["tf"]) + >>> idf = IDF(minDocFreq=3) + >>> idf.setInputCol("tf") + IDF... + >>> idf.setOutputCol("idf") + IDF... + >>> model = idf.fit(df) + >>> model.setOutputCol("idf") + IDFModel... + >>> model.getMinDocFreq() + 3 + >>> model.idf + DenseVector([0.0, 0.0]) + >>> model.docFreq + [0, 3] + >>> model.numDocs == df.count() + True + >>> model.transform(df).head().idf + DenseVector([0.0, 0.0]) + >>> idf.setParams(outputCol="freqs").fit(df).transform(df).collect()[1].freqs + DenseVector([0.0, 0.0]) + >>> params = {idf.minDocFreq: 1, idf.outputCol: "vector"} + >>> idf.fit(df, params).transform(df).head().vector + DenseVector([0.2877, 0.0]) + >>> idfPath = temp_path + "/idf" + >>> idf.save(idfPath) + >>> loadedIdf = IDF.load(idfPath) + >>> loadedIdf.getMinDocFreq() == idf.getMinDocFreq() + True + >>> modelPath = temp_path + "/idf-model" + >>> model.save(modelPath) + >>> loadedModel = IDFModel.load(modelPath) + >>> loadedModel.transform(df).head().idf == model.transform(df).head().idf + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + minDocFreq: int = 0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, minDocFreq=0, inputCol=None, outputCol=None) + """ + super(IDF, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.IDF", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + minDocFreq: int = 0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "IDF": + """ + setParams(self, \\*, minDocFreq=0, inputCol=None, outputCol=None) + Sets params for this IDF. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setMinDocFreq(self, value: int) -> "IDF": + """ + Sets the value of :py:attr:`minDocFreq`. + """ + return self._set(minDocFreq=value) + + def setInputCol(self, value: str) -> "IDF": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "IDF": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "IDFModel": + return IDFModel(java_model) + + +class IDFModel(JavaModel, _IDFParams, JavaMLReadable["IDFModel"], JavaMLWritable): + """ + Model fitted by :py:class:`IDF`. + + .. versionadded:: 1.4.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "IDFModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "IDFModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.0.0") + def idf(self) -> Vector: + """ + Returns the IDF vector. + """ + return self._call_java("idf") + + @property + @since("3.0.0") + def docFreq(self) -> List[int]: + """ + Returns the document frequency. + """ + return self._call_java("docFreq") + + @property + @since("3.0.0") + def numDocs(self) -> int: + """ + Returns number of documents evaluated to compute idf + """ + return self._call_java("numDocs") + + +class _ImputerParams(HasInputCol, HasInputCols, HasOutputCol, HasOutputCols, HasRelativeError): + """ + Params for :py:class:`Imputer` and :py:class:`ImputerModel`. + + .. versionadded:: 3.0.0 + """ + + strategy: Param[str] = Param( + Params._dummy(), + "strategy", + "strategy for imputation. If mean, then replace missing values using the mean " + "value of the feature. If median, then replace missing values using the " + "median value of the feature. If mode, then replace missing using the most " + "frequent value of the feature.", + typeConverter=TypeConverters.toString, + ) + + missingValue: Param[float] = Param( + Params._dummy(), + "missingValue", + "The placeholder for the missing values. All occurrences of missingValue " + "will be imputed.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_ImputerParams, self).__init__(*args) + self._setDefault(strategy="mean", missingValue=float("nan"), relativeError=0.001) + + @since("2.2.0") + def getStrategy(self) -> str: + """ + Gets the value of :py:attr:`strategy` or its default value. + """ + return self.getOrDefault(self.strategy) + + @since("2.2.0") + def getMissingValue(self) -> float: + """ + Gets the value of :py:attr:`missingValue` or its default value. + """ + return self.getOrDefault(self.missingValue) + + +@inherit_doc +class Imputer( + JavaEstimator["ImputerModel"], _ImputerParams, JavaMLReadable["Imputer"], JavaMLWritable +): + """ + Imputation estimator for completing missing values, using the mean, median or mode + of the columns in which the missing values are located. The input columns should be of + numeric type. Currently Imputer does not support categorical features and + possibly creates incorrect values for a categorical feature. + + Note that the mean/median/mode value is computed after filtering out missing values. + All Null values in the input columns are treated as missing, and so are also imputed. For + computing median, :py:meth:`pyspark.sql.DataFrame.approxQuantile` is used with a + relative error of `0.001`. + + .. versionadded:: 2.2.0 + + Examples + -------- + >>> df = spark.createDataFrame([(1.0, float("nan")), (2.0, float("nan")), (float("nan"), 3.0), + ... (4.0, 4.0), (5.0, 5.0)], ["a", "b"]) + >>> imputer = Imputer() + >>> imputer.setInputCols(["a", "b"]) + Imputer... + >>> imputer.setOutputCols(["out_a", "out_b"]) + Imputer... + >>> imputer.getRelativeError() + 0.001 + >>> model = imputer.fit(df) + >>> model.setInputCols(["a", "b"]) + ImputerModel... + >>> model.getStrategy() + 'mean' + >>> model.surrogateDF.show() + +---+---+ + | a| b| + +---+---+ + |3.0|4.0| + +---+---+ + ... + >>> model.transform(df).show() + +---+---+-----+-----+ + | a| b|out_a|out_b| + +---+---+-----+-----+ + |1.0|NaN| 1.0| 4.0| + |2.0|NaN| 2.0| 4.0| + |NaN|3.0| 3.0| 3.0| + ... + >>> imputer.setStrategy("median").setMissingValue(1.0).fit(df).transform(df).show() + +---+---+-----+-----+ + | a| b|out_a|out_b| + +---+---+-----+-----+ + |1.0|NaN| 4.0| NaN| + ... + >>> df1 = spark.createDataFrame([(1.0,), (2.0,), (float("nan"),), (4.0,), (5.0,)], ["a"]) + >>> imputer1 = Imputer(inputCol="a", outputCol="out_a") + >>> model1 = imputer1.fit(df1) + >>> model1.surrogateDF.show() + +---+ + | a| + +---+ + |3.0| + +---+ + ... + >>> model1.transform(df1).show() + +---+-----+ + | a|out_a| + +---+-----+ + |1.0| 1.0| + |2.0| 2.0| + |NaN| 3.0| + ... + >>> imputer1.setStrategy("median").setMissingValue(1.0).fit(df1).transform(df1).show() + +---+-----+ + | a|out_a| + +---+-----+ + |1.0| 4.0| + ... + >>> df2 = spark.createDataFrame([(float("nan"),), (float("nan"),), (3.0,), (4.0,), (5.0,)], + ... ["b"]) + >>> imputer2 = Imputer(inputCol="b", outputCol="out_b") + >>> model2 = imputer2.fit(df2) + >>> model2.surrogateDF.show() + +---+ + | b| + +---+ + |4.0| + +---+ + ... + >>> model2.transform(df2).show() + +---+-----+ + | b|out_b| + +---+-----+ + |NaN| 4.0| + |NaN| 4.0| + |3.0| 3.0| + ... + >>> imputer2.setStrategy("median").setMissingValue(1.0).fit(df2).transform(df2).show() + +---+-----+ + | b|out_b| + +---+-----+ + |NaN| NaN| + ... + >>> imputerPath = temp_path + "/imputer" + >>> imputer.save(imputerPath) + >>> loadedImputer = Imputer.load(imputerPath) + >>> loadedImputer.getStrategy() == imputer.getStrategy() + True + >>> loadedImputer.getMissingValue() + 1.0 + >>> modelPath = temp_path + "/imputer-model" + >>> model.save(modelPath) + >>> loadedModel = ImputerModel.load(modelPath) + >>> loadedModel.transform(df).head().out_a == model.transform(df).head().out_a + True + """ + + _input_kwargs: Dict[str, Any] + + @overload + def __init__( + self, + *, + strategy: str = ..., + missingValue: float = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + relativeError: float = ..., + ): + ... + + @overload + def __init__( + self, + *, + strategy: str = ..., + missingValue: float = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + relativeError: float = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + strategy: str = "mean", + missingValue: float = float("nan"), + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + ): + """ + __init__(self, \\*, strategy="mean", missingValue=float("nan"), inputCols=None, \ + outputCols=None, inputCol=None, outputCol=None, relativeError=0.001): + """ + super(Imputer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Imputer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + strategy: str = ..., + missingValue: float = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + relativeError: float = ..., + ) -> "Imputer": + ... + + @overload + def setParams( + self, + *, + strategy: str = ..., + missingValue: float = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + relativeError: float = ..., + ) -> "Imputer": + ... + + @keyword_only + @since("2.2.0") + def setParams( + self, + *, + strategy: str = "mean", + missingValue: float = float("nan"), + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + ) -> "Imputer": + """ + setParams(self, \\*, strategy="mean", missingValue=float("nan"), inputCols=None, \ + outputCols=None, inputCol=None, outputCol=None, relativeError=0.001) + Sets params for this Imputer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.2.0") + def setStrategy(self, value: str) -> "Imputer": + """ + Sets the value of :py:attr:`strategy`. + """ + return self._set(strategy=value) + + @since("2.2.0") + def setMissingValue(self, value: float) -> "Imputer": + """ + Sets the value of :py:attr:`missingValue`. + """ + return self._set(missingValue=value) + + @since("2.2.0") + def setInputCols(self, value: List[str]) -> "Imputer": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("2.2.0") + def setOutputCols(self, value: List[str]) -> "Imputer": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @since("3.0.0") + def setInputCol(self, value: str) -> "Imputer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "Imputer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setRelativeError(self, value: float) -> "Imputer": + """ + Sets the value of :py:attr:`relativeError`. + """ + return self._set(relativeError=value) + + def _create_model(self, java_model: "JavaObject") -> "ImputerModel": + return ImputerModel(java_model) + + +class ImputerModel(JavaModel, _ImputerParams, JavaMLReadable["ImputerModel"], JavaMLWritable): + """ + Model fitted by :py:class:`Imputer`. + + .. versionadded:: 2.2.0 + """ + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "ImputerModel": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "ImputerModel": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @since("3.0.0") + def setInputCol(self, value: str) -> "ImputerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "ImputerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.2.0") + def surrogateDF(self) -> DataFrame: + """ + Returns a DataFrame containing inputCols and their corresponding surrogates, + which are used to replace the missing values in the input DataFrame. + """ + return self._call_java("surrogateDF") + + +@inherit_doc +class Interaction( + JavaTransformer, + HasInputCols, + HasOutputCol, + JavaMLReadable["Interaction"], + JavaMLWritable, +): + """ + Implements the feature interaction transform. This transformer takes in Double and Vector type + columns and outputs a flattened vector of their feature interactions. To handle interaction, + we first one-hot encode any nominal features. Then, a vector of the feature cross-products is + produced. + + For example, given the input feature values `Double(2)` and `Vector(3, 4)`, the output would be + `Vector(6, 8)` if all input features were numeric. If the first feature was instead nominal + with four categories, the output would then be `Vector(0, 0, 0, 0, 3, 4, 0, 0)`. + + .. versionadded:: 3.0.0 + + Examples + -------- + >>> df = spark.createDataFrame([(0.0, 1.0), (2.0, 3.0)], ["a", "b"]) + >>> interaction = Interaction() + >>> interaction.setInputCols(["a", "b"]) + Interaction... + >>> interaction.setOutputCol("ab") + Interaction... + >>> interaction.transform(df).show() + +---+---+-----+ + | a| b| ab| + +---+---+-----+ + |0.0|1.0|[0.0]| + |2.0|3.0|[6.0]| + +---+---+-----+ + ... + >>> interactionPath = temp_path + "/interaction" + >>> interaction.save(interactionPath) + >>> loadedInteraction = Interaction.load(interactionPath) + >>> loadedInteraction.transform(df).head().ab == interaction.transform(df).head().ab + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, inputCols: Optional[List[str]] = None, outputCol: Optional[str] = None): + """ + __init__(self, \\*, inputCols=None, outputCol=None): + """ + super(Interaction, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Interaction", self.uid) + self._setDefault() + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.0.0") + def setParams( + self, *, inputCols: Optional[List[str]] = None, outputCol: Optional[str] = None + ) -> "Interaction": + """ + setParams(self, \\*, inputCols=None, outputCol=None) + Sets params for this Interaction. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "Interaction": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "Interaction": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +class _MaxAbsScalerParams(HasInputCol, HasOutputCol): + """ + Params for :py:class:`MaxAbsScaler` and :py:class:`MaxAbsScalerModel`. + + .. versionadded:: 3.0.0 + """ + + pass + + +@inherit_doc +class MaxAbsScaler( + JavaEstimator["MaxAbsScalerModel"], + _MaxAbsScalerParams, + JavaMLReadable["MaxAbsScaler"], + JavaMLWritable, +): + """ + Rescale each feature individually to range [-1, 1] by dividing through the largest maximum + absolute value in each feature. It does not shift/center the data, and thus does not destroy + any sparsity. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([1.0]),), (Vectors.dense([2.0]),)], ["a"]) + >>> maScaler = MaxAbsScaler(outputCol="scaled") + >>> maScaler.setInputCol("a") + MaxAbsScaler... + >>> model = maScaler.fit(df) + >>> model.setOutputCol("scaledOutput") + MaxAbsScalerModel... + >>> model.transform(df).show() + +-----+------------+ + | a|scaledOutput| + +-----+------------+ + |[1.0]| [0.5]| + |[2.0]| [1.0]| + +-----+------------+ + ... + >>> scalerPath = temp_path + "/max-abs-scaler" + >>> maScaler.save(scalerPath) + >>> loadedMAScaler = MaxAbsScaler.load(scalerPath) + >>> loadedMAScaler.getInputCol() == maScaler.getInputCol() + True + >>> loadedMAScaler.getOutputCol() == maScaler.getOutputCol() + True + >>> modelPath = temp_path + "/max-abs-scaler-model" + >>> model.save(modelPath) + >>> loadedModel = MaxAbsScalerModel.load(modelPath) + >>> loadedModel.maxAbs == model.maxAbs + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, inputCol: Optional[str] = None, outputCol: Optional[str] = None): + """ + __init__(self, \\*, inputCol=None, outputCol=None) + """ + super(MaxAbsScaler, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.MaxAbsScaler", self.uid) + self._setDefault() + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.0.0") + def setParams( + self, *, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ) -> "MaxAbsScaler": + """ + setParams(self, \\*, inputCol=None, outputCol=None) + Sets params for this MaxAbsScaler. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def setInputCol(self, value: str) -> "MaxAbsScaler": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "MaxAbsScaler": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "MaxAbsScalerModel": + return MaxAbsScalerModel(java_model) + + +class MaxAbsScalerModel( + JavaModel, _MaxAbsScalerParams, JavaMLReadable["MaxAbsScalerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`MaxAbsScaler`. + + .. versionadded:: 2.0.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "MaxAbsScalerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "MaxAbsScalerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.0.0") + def maxAbs(self) -> Vector: + """ + Max Abs vector. + """ + return self._call_java("maxAbs") + + +@inherit_doc +class MinHashLSH( + _LSH["MinHashLSHModel"], + HasInputCol, + HasOutputCol, + HasSeed, + JavaMLReadable["MinHashLSH"], + JavaMLWritable, +): + + """ + LSH class for Jaccard distance. + The input can be dense or sparse vectors, but it is more efficient if it is sparse. + For example, `Vectors.sparse(10, [(2, 1.0), (3, 1.0), (5, 1.0)])` means there are 10 elements + in the space. This set contains elements 2, 3, and 5. Also, any input vector must have at + least 1 non-zero index, and all non-zero values are treated as binary "1" values. + + .. versionadded:: 2.2.0 + + Notes + ----- + See `Wikipedia on MinHash `_ + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.sql.functions import col + >>> data = [(0, Vectors.sparse(6, [0, 1, 2], [1.0, 1.0, 1.0]),), + ... (1, Vectors.sparse(6, [2, 3, 4], [1.0, 1.0, 1.0]),), + ... (2, Vectors.sparse(6, [0, 2, 4], [1.0, 1.0, 1.0]),)] + >>> df = spark.createDataFrame(data, ["id", "features"]) + >>> mh = MinHashLSH() + >>> mh.setInputCol("features") + MinHashLSH... + >>> mh.setOutputCol("hashes") + MinHashLSH... + >>> mh.setSeed(12345) + MinHashLSH... + >>> model = mh.fit(df) + >>> model.setInputCol("features") + MinHashLSHModel... + >>> model.transform(df).head() + Row(id=0, features=SparseVector(6, {0: 1.0, 1: 1.0, 2: 1.0}), hashes=[DenseVector([6179668... + >>> data2 = [(3, Vectors.sparse(6, [1, 3, 5], [1.0, 1.0, 1.0]),), + ... (4, Vectors.sparse(6, [2, 3, 5], [1.0, 1.0, 1.0]),), + ... (5, Vectors.sparse(6, [1, 2, 4], [1.0, 1.0, 1.0]),)] + >>> df2 = spark.createDataFrame(data2, ["id", "features"]) + >>> key = Vectors.sparse(6, [1, 2], [1.0, 1.0]) + >>> model.approxNearestNeighbors(df2, key, 1).collect() + [Row(id=5, features=SparseVector(6, {1: 1.0, 2: 1.0, 4: 1.0}), hashes=[DenseVector([6179668... + >>> model.approxSimilarityJoin(df, df2, 0.6, distCol="JaccardDistance").select( + ... col("datasetA.id").alias("idA"), + ... col("datasetB.id").alias("idB"), + ... col("JaccardDistance")).show() + +---+---+---------------+ + |idA|idB|JaccardDistance| + +---+---+---------------+ + | 0| 5| 0.5| + | 1| 4| 0.5| + +---+---+---------------+ + ... + >>> mhPath = temp_path + "/mh" + >>> mh.save(mhPath) + >>> mh2 = MinHashLSH.load(mhPath) + >>> mh2.getOutputCol() == mh.getOutputCol() + True + >>> modelPath = temp_path + "/mh-model" + >>> model.save(modelPath) + >>> model2 = MinHashLSHModel.load(modelPath) + >>> model.transform(df).head().hashes == model2.transform(df).head().hashes + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + seed: Optional[int] = None, + numHashTables: int = 1, + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, seed=None, numHashTables=1) + """ + super(MinHashLSH, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.MinHashLSH", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.2.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + seed: Optional[int] = None, + numHashTables: int = 1, + ) -> "MinHashLSH": + """ + setParams(self, \\*, inputCol=None, outputCol=None, seed=None, numHashTables=1) + Sets params for this MinHashLSH. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def setSeed(self, value: int) -> "MinHashLSH": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def _create_model(self, java_model: "JavaObject") -> "MinHashLSHModel": + return MinHashLSHModel(java_model) + + +class MinHashLSHModel(_LSHModel, JavaMLReadable, JavaMLWritable): + r""" + Model produced by :py:class:`MinHashLSH`, where where multiple hash functions are stored. Each + hash function is picked from the following family of hash functions, where :math:`a_i` and + :math:`b_i` are randomly chosen integers less than prime: + :math:`h_i(x) = ((x \cdot a_i + b_i) \mod prime)` This hash family is approximately min-wise + independent according to the reference. + + .. versionadded:: 2.2.0 + + Notes + ----- + See Tom Bohman, Colin Cooper, and Alan Frieze. "Min-wise independent linear permutations." + Electronic Journal of Combinatorics 7 (2000): R26. + """ + + +class _MinMaxScalerParams(HasInputCol, HasOutputCol): + """ + Params for :py:class:`MinMaxScaler` and :py:class:`MinMaxScalerModel`. + + .. versionadded:: 3.0.0 + """ + + min: Param[float] = Param( + Params._dummy(), + "min", + "Lower bound of the output feature range", + typeConverter=TypeConverters.toFloat, + ) + max: Param[float] = Param( + Params._dummy(), + "max", + "Upper bound of the output feature range", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_MinMaxScalerParams, self).__init__(*args) + self._setDefault(min=0.0, max=1.0) + + @since("1.6.0") + def getMin(self) -> float: + """ + Gets the value of min or its default value. + """ + return self.getOrDefault(self.min) + + @since("1.6.0") + def getMax(self) -> float: + """ + Gets the value of max or its default value. + """ + return self.getOrDefault(self.max) + + +@inherit_doc +class MinMaxScaler( + JavaEstimator["MinMaxScalerModel"], + _MinMaxScalerParams, + JavaMLReadable["MinMaxScaler"], + JavaMLWritable, +): + """ + Rescale each feature individually to a common range [min, max] linearly using column summary + statistics, which is also known as min-max normalization or Rescaling. The rescaled value for + feature E is calculated as, + + Rescaled(e_i) = (e_i - E_min) / (E_max - E_min) * (max - min) + min + + For the case E_max == E_min, Rescaled(e_i) = 0.5 * (max + min) + + .. versionadded:: 1.6.0 + + Notes + ----- + Since zero values will probably be transformed to non-zero values, output of the + transformer will be DenseVector even for sparse input. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([0.0]),), (Vectors.dense([2.0]),)], ["a"]) + >>> mmScaler = MinMaxScaler(outputCol="scaled") + >>> mmScaler.setInputCol("a") + MinMaxScaler... + >>> model = mmScaler.fit(df) + >>> model.setOutputCol("scaledOutput") + MinMaxScalerModel... + >>> model.originalMin + DenseVector([0.0]) + >>> model.originalMax + DenseVector([2.0]) + >>> model.transform(df).show() + +-----+------------+ + | a|scaledOutput| + +-----+------------+ + |[0.0]| [0.0]| + |[2.0]| [1.0]| + +-----+------------+ + ... + >>> minMaxScalerPath = temp_path + "/min-max-scaler" + >>> mmScaler.save(minMaxScalerPath) + >>> loadedMMScaler = MinMaxScaler.load(minMaxScalerPath) + >>> loadedMMScaler.getMin() == mmScaler.getMin() + True + >>> loadedMMScaler.getMax() == mmScaler.getMax() + True + >>> modelPath = temp_path + "/min-max-scaler-model" + >>> model.save(modelPath) + >>> loadedModel = MinMaxScalerModel.load(modelPath) + >>> loadedModel.originalMin == model.originalMin + True + >>> loadedModel.originalMax == model.originalMax + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + min: float = 0.0, + max: float = 1.0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, min=0.0, max=1.0, inputCol=None, outputCol=None) + """ + super(MinMaxScaler, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.MinMaxScaler", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + min: float = 0.0, + max: float = 1.0, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "MinMaxScaler": + """ + setParams(self, \\*, min=0.0, max=1.0, inputCol=None, outputCol=None) + Sets params for this MinMaxScaler. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setMin(self, value: float) -> "MinMaxScaler": + """ + Sets the value of :py:attr:`min`. + """ + return self._set(min=value) + + @since("1.6.0") + def setMax(self, value: float) -> "MinMaxScaler": + """ + Sets the value of :py:attr:`max`. + """ + return self._set(max=value) + + def setInputCol(self, value: str) -> "MinMaxScaler": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "MinMaxScaler": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "MinMaxScalerModel": + return MinMaxScalerModel(java_model) + + +class MinMaxScalerModel( + JavaModel, _MinMaxScalerParams, JavaMLReadable["MinMaxScalerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`MinMaxScaler`. + + .. versionadded:: 1.6.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "MinMaxScalerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "MinMaxScalerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setMin(self, value: float) -> "MinMaxScalerModel": + """ + Sets the value of :py:attr:`min`. + """ + return self._set(min=value) + + @since("3.0.0") + def setMax(self, value: float) -> "MinMaxScalerModel": + """ + Sets the value of :py:attr:`max`. + """ + return self._set(max=value) + + @property + @since("2.0.0") + def originalMin(self) -> Vector: + """ + Min value for each original column during fitting. + """ + return self._call_java("originalMin") + + @property + @since("2.0.0") + def originalMax(self) -> Vector: + """ + Max value for each original column during fitting. + """ + return self._call_java("originalMax") + + +@inherit_doc +class NGram(JavaTransformer, HasInputCol, HasOutputCol, JavaMLReadable["NGram"], JavaMLWritable): + """ + A feature transformer that converts the input array of strings into an array of n-grams. Null + values in the input array are ignored. + It returns an array of n-grams where each n-gram is represented by a space-separated string of + words. + When the input is empty, an empty array is returned. + When the input array length is less than n (number of elements per n-gram), no n-grams are + returned. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([Row(inputTokens=["a", "b", "c", "d", "e"])]) + >>> ngram = NGram(n=2) + >>> ngram.setInputCol("inputTokens") + NGram... + >>> ngram.setOutputCol("nGrams") + NGram... + >>> ngram.transform(df).head() + Row(inputTokens=['a', 'b', 'c', 'd', 'e'], nGrams=['a b', 'b c', 'c d', 'd e']) + >>> # Change n-gram length + >>> ngram.setParams(n=4).transform(df).head() + Row(inputTokens=['a', 'b', 'c', 'd', 'e'], nGrams=['a b c d', 'b c d e']) + >>> # Temporarily modify output column. + >>> ngram.transform(df, {ngram.outputCol: "output"}).head() + Row(inputTokens=['a', 'b', 'c', 'd', 'e'], output=['a b c d', 'b c d e']) + >>> ngram.transform(df).head() + Row(inputTokens=['a', 'b', 'c', 'd', 'e'], nGrams=['a b c d', 'b c d e']) + >>> # Must use keyword arguments to specify params. + >>> ngram.setParams("text") + Traceback (most recent call last): + ... + TypeError: Method setParams forces keyword arguments. + >>> ngramPath = temp_path + "/ngram" + >>> ngram.save(ngramPath) + >>> loadedNGram = NGram.load(ngramPath) + >>> loadedNGram.getN() == ngram.getN() + True + >>> loadedNGram.transform(df).take(1) == ngram.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + n: Param[int] = Param( + Params._dummy(), + "n", + "number of elements per n-gram (>=1)", + typeConverter=TypeConverters.toInt, + ) + + @keyword_only + def __init__( + self, *, n: int = 2, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ): + """ + __init__(self, \\*, n=2, inputCol=None, outputCol=None) + """ + super(NGram, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.NGram", self.uid) + self._setDefault(n=2) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.5.0") + def setParams( + self, *, n: int = 2, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ) -> "NGram": + """ + setParams(self, \\*, n=2, inputCol=None, outputCol=None) + Sets params for this NGram. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.5.0") + def setN(self, value: int) -> "NGram": + """ + Sets the value of :py:attr:`n`. + """ + return self._set(n=value) + + @since("1.5.0") + def getN(self) -> int: + """ + Gets the value of n or its default value. + """ + return self.getOrDefault(self.n) + + def setInputCol(self, value: str) -> "NGram": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "NGram": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class Normalizer( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["Normalizer"], + JavaMLWritable, +): + """ + Normalize a vector to have unit norm using the given p-norm. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> svec = Vectors.sparse(4, {1: 4.0, 3: 3.0}) + >>> df = spark.createDataFrame([(Vectors.dense([3.0, -4.0]), svec)], ["dense", "sparse"]) + >>> normalizer = Normalizer(p=2.0) + >>> normalizer.setInputCol("dense") + Normalizer... + >>> normalizer.setOutputCol("features") + Normalizer... + >>> normalizer.transform(df).head().features + DenseVector([0.6, -0.8]) + >>> normalizer.setParams(inputCol="sparse", outputCol="freqs").transform(df).head().freqs + SparseVector(4, {1: 0.8, 3: 0.6}) + >>> params = {normalizer.p: 1.0, normalizer.inputCol: "dense", normalizer.outputCol: "vector"} + >>> normalizer.transform(df, params).head().vector + DenseVector([0.4286, -0.5714]) + >>> normalizerPath = temp_path + "/normalizer" + >>> normalizer.save(normalizerPath) + >>> loadedNormalizer = Normalizer.load(normalizerPath) + >>> loadedNormalizer.getP() == normalizer.getP() + True + >>> loadedNormalizer.transform(df).take(1) == normalizer.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + p = Param(Params._dummy(), "p", "the p norm value.", typeConverter=TypeConverters.toFloat) + + @keyword_only + def __init__( + self, *, p: float = 2.0, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ): + """ + __init__(self, \\*, p=2.0, inputCol=None, outputCol=None) + """ + super(Normalizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Normalizer", self.uid) + self._setDefault(p=2.0) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, *, p: float = 2.0, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ) -> "Normalizer": + """ + setParams(self, \\*, p=2.0, inputCol=None, outputCol=None) + Sets params for this Normalizer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setP(self, value: float) -> "Normalizer": + """ + Sets the value of :py:attr:`p`. + """ + return self._set(p=value) + + @since("1.4.0") + def getP(self) -> float: + """ + Gets the value of p or its default value. + """ + return self.getOrDefault(self.p) + + def setInputCol(self, value: str) -> "Normalizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "Normalizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +class _OneHotEncoderParams( + HasInputCol, HasInputCols, HasOutputCol, HasOutputCols, HasHandleInvalid +): + """ + Params for :py:class:`OneHotEncoder` and :py:class:`OneHotEncoderModel`. + + .. versionadded:: 3.0.0 + """ + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "How to handle invalid data during " + + "transform(). Options are 'keep' (invalid data presented as an extra " + + "categorical feature) or error (throw an error). Note that this Param " + + "is only used during transform; during fitting, invalid data will " + + "result in an error.", + typeConverter=TypeConverters.toString, + ) + + dropLast: Param[bool] = Param( + Params._dummy(), + "dropLast", + "whether to drop the last category", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self, *args: Any): + super(_OneHotEncoderParams, self).__init__(*args) + self._setDefault(handleInvalid="error", dropLast=True) + + @since("2.3.0") + def getDropLast(self) -> bool: + """ + Gets the value of dropLast or its default value. + """ + return self.getOrDefault(self.dropLast) + + +@inherit_doc +class OneHotEncoder( + JavaEstimator["OneHotEncoderModel"], + _OneHotEncoderParams, + JavaMLReadable["OneHotEncoder"], + JavaMLWritable, +): + """ + A one-hot encoder that maps a column of category indices to a column of binary vectors, with + at most a single one-value per row that indicates the input category index. + For example with 5 categories, an input value of 2.0 would map to an output vector of + `[0.0, 0.0, 1.0, 0.0]`. + The last category is not included by default (configurable via :py:attr:`dropLast`), + because it makes the vector entries sum up to one, and hence linearly dependent. + So an input value of 4.0 maps to `[0.0, 0.0, 0.0, 0.0]`. + + When :py:attr:`handleInvalid` is configured to 'keep', an extra "category" indicating invalid + values is added as last category. So when :py:attr:`dropLast` is true, invalid values are + encoded as all-zeros vector. + + .. versionadded:: 2.3.0 + + Notes + ----- + This is different from scikit-learn's OneHotEncoder, which keeps all categories. + The output vectors are sparse. + + When encoding multi-column by using :py:attr:`inputCols` and + :py:attr:`outputCols` params, input/output cols come in pairs, specified by the order in + the arrays, and each pair is treated independently. + + See Also + -------- + StringIndexer : for converting categorical values into category indices + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(0.0,), (1.0,), (2.0,)], ["input"]) + >>> ohe = OneHotEncoder() + >>> ohe.setInputCols(["input"]) + OneHotEncoder... + >>> ohe.setOutputCols(["output"]) + OneHotEncoder... + >>> model = ohe.fit(df) + >>> model.setOutputCols(["output"]) + OneHotEncoderModel... + >>> model.getHandleInvalid() + 'error' + >>> model.transform(df).head().output + SparseVector(2, {0: 1.0}) + >>> single_col_ohe = OneHotEncoder(inputCol="input", outputCol="output") + >>> single_col_model = single_col_ohe.fit(df) + >>> single_col_model.transform(df).head().output + SparseVector(2, {0: 1.0}) + >>> ohePath = temp_path + "/ohe" + >>> ohe.save(ohePath) + >>> loadedOHE = OneHotEncoder.load(ohePath) + >>> loadedOHE.getInputCols() == ohe.getInputCols() + True + >>> modelPath = temp_path + "/ohe-model" + >>> model.save(modelPath) + >>> loadedModel = OneHotEncoderModel.load(modelPath) + >>> loadedModel.categorySizes == model.categorySizes + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @overload + def __init__( + self, + *, + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + handleInvalid: str = ..., + dropLast: bool = ..., + ): + ... + + @overload + def __init__( + self, + *, + handleInvalid: str = ..., + dropLast: bool = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + handleInvalid: str = "error", + dropLast: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, inputCols=None, outputCols=None, handleInvalid="error", dropLast=True, \ + inputCol=None, outputCol=None) + """ + super(OneHotEncoder, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.OneHotEncoder", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + handleInvalid: str = ..., + dropLast: bool = ..., + ) -> "OneHotEncoder": + ... + + @overload + def setParams( + self, + *, + handleInvalid: str = ..., + dropLast: bool = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + ) -> "OneHotEncoder": + ... + + @keyword_only + @since("2.3.0") + def setParams( + self, + *, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + handleInvalid: str = "error", + dropLast: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "OneHotEncoder": + """ + setParams(self, \\*, inputCols=None, outputCols=None, handleInvalid="error", \ + dropLast=True, inputCol=None, outputCol=None) + Sets params for this OneHotEncoder. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.3.0") + def setDropLast(self, value: bool) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`dropLast`. + """ + return self._set(dropLast=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @since("3.0.0") + def setHandleInvalid(self, value: str) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + @since("3.0.0") + def setInputCol(self, value: str) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "OneHotEncoder": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "OneHotEncoderModel": + return OneHotEncoderModel(java_model) + + +class OneHotEncoderModel( + JavaModel, _OneHotEncoderParams, JavaMLReadable["OneHotEncoderModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`OneHotEncoder`. + + .. versionadded:: 2.3.0 + """ + + @since("3.0.0") + def setDropLast(self, value: bool) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`dropLast`. + """ + return self._set(dropLast=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @since("3.0.0") + def setInputCol(self, value: str) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setHandleInvalid(self, value: str) -> "OneHotEncoderModel": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + @property + @since("2.3.0") + def categorySizes(self) -> List[int]: + """ + Original number of categories for each feature being encoded. + The array contains one value for each input column, in order. + """ + return self._call_java("categorySizes") + + +@inherit_doc +class PolynomialExpansion( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["PolynomialExpansion"], + JavaMLWritable, +): + """ + Perform feature expansion in a polynomial space. As said in `wikipedia of Polynomial Expansion + `_, "In mathematics, an + expansion of a product of sums expresses it as a sum of products by using the fact that + multiplication distributes over addition". Take a 2-variable feature vector as an example: + `(x, y)`, if we want to expand it with degree 2, then we get `(x, x * x, y, x * y, y * y)`. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([0.5, 2.0]),)], ["dense"]) + >>> px = PolynomialExpansion(degree=2) + >>> px.setInputCol("dense") + PolynomialExpansion... + >>> px.setOutputCol("expanded") + PolynomialExpansion... + >>> px.transform(df).head().expanded + DenseVector([0.5, 0.25, 2.0, 1.0, 4.0]) + >>> px.setParams(outputCol="test").transform(df).head().test + DenseVector([0.5, 0.25, 2.0, 1.0, 4.0]) + >>> polyExpansionPath = temp_path + "/poly-expansion" + >>> px.save(polyExpansionPath) + >>> loadedPx = PolynomialExpansion.load(polyExpansionPath) + >>> loadedPx.getDegree() == px.getDegree() + True + >>> loadedPx.transform(df).take(1) == px.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + degree: Param[int] = Param( + Params._dummy(), + "degree", + "the polynomial degree to expand (>= 1)", + typeConverter=TypeConverters.toInt, + ) + + @keyword_only + def __init__( + self, *, degree: int = 2, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ): + """ + __init__(self, \\*, degree=2, inputCol=None, outputCol=None) + """ + super(PolynomialExpansion, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.PolynomialExpansion", self.uid + ) + self._setDefault(degree=2) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, *, degree: int = 2, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ) -> "PolynomialExpansion": + """ + setParams(self, \\*, degree=2, inputCol=None, outputCol=None) + Sets params for this PolynomialExpansion. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setDegree(self, value: int) -> "PolynomialExpansion": + """ + Sets the value of :py:attr:`degree`. + """ + return self._set(degree=value) + + @since("1.4.0") + def getDegree(self) -> int: + """ + Gets the value of degree or its default value. + """ + return self.getOrDefault(self.degree) + + def setInputCol(self, value: str) -> "PolynomialExpansion": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "PolynomialExpansion": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class QuantileDiscretizer( + JavaEstimator, + HasInputCol, + HasOutputCol, + HasInputCols, + HasOutputCols, + HasHandleInvalid, + HasRelativeError, + JavaMLReadable["QuantileDiscretizer"], + JavaMLWritable, +): + """ + :py:class:`QuantileDiscretizer` takes a column with continuous features and outputs a column + with binned categorical features. The number of bins can be set using the :py:attr:`numBuckets` + parameter. It is possible that the number of buckets used will be less than this value, for + example, if there are too few distinct values of the input to create enough distinct quantiles. + Since 3.0.0, :py:class:`QuantileDiscretizer` can map multiple columns at once by setting the + :py:attr:`inputCols` parameter. If both of the :py:attr:`inputCol` and :py:attr:`inputCols` + parameters are set, an Exception will be thrown. To specify the number of buckets for each + column, the :py:attr:`numBucketsArray` parameter can be set, or if the number of buckets + should be the same across columns, :py:attr:`numBuckets` can be set as a convenience. + + .. versionadded:: 2.0.0 + + Notes + ----- + NaN handling: Note also that + :py:class:`QuantileDiscretizer` will raise an error when it finds NaN values in the dataset, + but the user can also choose to either keep or remove NaN values within the dataset by setting + :py:attr:`handleInvalid` parameter. If the user chooses to keep NaN values, they will be + handled specially and placed into their own bucket, for example, if 4 buckets are used, then + non-NaN data will be put into buckets[0-3], but NaNs will be counted in a special bucket[4]. + + Algorithm: The bin ranges are chosen using an approximate algorithm (see the documentation for + :py:meth:`pyspark.sql.DataFrameStatFunctions.approxQuantile` for a detailed description). + The precision of the approximation can be controlled with the + :py:attr:`relativeError` parameter. + The lower and upper bin bounds will be `-Infinity` and `+Infinity`, covering all real values. + + Examples + -------- + >>> values = [(0.1,), (0.4,), (1.2,), (1.5,), (float("nan"),), (float("nan"),)] + >>> df1 = spark.createDataFrame(values, ["values"]) + >>> qds1 = QuantileDiscretizer(inputCol="values", outputCol="buckets") + >>> qds1.setNumBuckets(2) + QuantileDiscretizer... + >>> qds1.setRelativeError(0.01) + QuantileDiscretizer... + >>> qds1.setHandleInvalid("error") + QuantileDiscretizer... + >>> qds1.getRelativeError() + 0.01 + >>> bucketizer = qds1.fit(df1) + >>> qds1.setHandleInvalid("keep").fit(df1).transform(df1).count() + 6 + >>> qds1.setHandleInvalid("skip").fit(df1).transform(df1).count() + 4 + >>> splits = bucketizer.getSplits() + >>> splits[0] + -inf + >>> print("%2.1f" % round(splits[1], 1)) + 0.4 + >>> bucketed = bucketizer.transform(df1).head() + >>> bucketed.buckets + 0.0 + >>> quantileDiscretizerPath = temp_path + "/quantile-discretizer" + >>> qds1.save(quantileDiscretizerPath) + >>> loadedQds = QuantileDiscretizer.load(quantileDiscretizerPath) + >>> loadedQds.getNumBuckets() == qds1.getNumBuckets() + True + >>> inputs = [(0.1, 0.0), (0.4, 1.0), (1.2, 1.3), (1.5, 1.5), + ... (float("nan"), float("nan")), (float("nan"), float("nan"))] + >>> df2 = spark.createDataFrame(inputs, ["input1", "input2"]) + >>> qds2 = QuantileDiscretizer(relativeError=0.01, handleInvalid="error", numBuckets=2, + ... inputCols=["input1", "input2"], outputCols=["output1", "output2"]) + >>> qds2.getRelativeError() + 0.01 + >>> qds2.setHandleInvalid("keep").fit(df2).transform(df2).show() + +------+------+-------+-------+ + |input1|input2|output1|output2| + +------+------+-------+-------+ + | 0.1| 0.0| 0.0| 0.0| + | 0.4| 1.0| 1.0| 1.0| + | 1.2| 1.3| 1.0| 1.0| + | 1.5| 1.5| 1.0| 1.0| + | NaN| NaN| 2.0| 2.0| + | NaN| NaN| 2.0| 2.0| + +------+------+-------+-------+ + ... + >>> qds3 = QuantileDiscretizer(relativeError=0.01, handleInvalid="error", + ... numBucketsArray=[5, 10], inputCols=["input1", "input2"], + ... outputCols=["output1", "output2"]) + >>> qds3.setHandleInvalid("skip").fit(df2).transform(df2).show() + +------+------+-------+-------+ + |input1|input2|output1|output2| + +------+------+-------+-------+ + | 0.1| 0.0| 1.0| 1.0| + | 0.4| 1.0| 2.0| 2.0| + | 1.2| 1.3| 3.0| 3.0| + | 1.5| 1.5| 4.0| 4.0| + +------+------+-------+-------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + numBuckets: Param[int] = Param( + Params._dummy(), + "numBuckets", + "Maximum number of buckets (quantiles, or " + + "categories) into which data points are grouped. Must be >= 2.", + typeConverter=TypeConverters.toInt, + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "how to handle invalid entries. " + + "Options are skip (filter out rows with invalid values), " + + "error (throw an error), or keep (keep invalid values in a special " + + "additional bucket). Note that in the multiple columns " + + "case, the invalid handling is applied to all columns. That said " + + "for 'error' it will throw an error if any invalids are found in " + + "any columns, for 'skip' it will skip rows with any invalids in " + + "any columns, etc.", + typeConverter=TypeConverters.toString, + ) + + numBucketsArray: Param[List[int]] = Param( + Params._dummy(), + "numBucketsArray", + "Array of number of buckets " + + "(quantiles, or categories) into which data points are grouped. " + + "This is for multiple columns input. If transforming multiple " + + "columns and numBucketsArray is not set, but numBuckets is set, " + + "then numBuckets will be applied across all columns.", + typeConverter=TypeConverters.toListInt, + ) + + @overload + def __init__( + self, + *, + numBuckets: int = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + relativeError: float = ..., + handleInvalid: str = ..., + ): + ... + + @overload + def __init__( + self, + *, + relativeError: float = ..., + handleInvalid: str = ..., + numBucketsArray: Optional[List[int]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + numBuckets: int = 2, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + handleInvalid: str = "error", + numBucketsArray: Optional[List[int]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, numBuckets=2, inputCol=None, outputCol=None, relativeError=0.001, \ + handleInvalid="error", numBucketsArray=None, inputCols=None, outputCols=None) + """ + super(QuantileDiscretizer, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.QuantileDiscretizer", self.uid + ) + self._setDefault(numBuckets=2, relativeError=0.001, handleInvalid="error") + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + numBuckets: int = ..., + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + relativeError: float = ..., + handleInvalid: str = ..., + ) -> "QuantileDiscretizer": + ... + + @overload + def setParams( + self, + *, + relativeError: float = ..., + handleInvalid: str = ..., + numBucketsArray: Optional[List[int]] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ) -> "QuantileDiscretizer": + ... + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + numBuckets: int = 2, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + handleInvalid: str = "error", + numBucketsArray: Optional[List[int]] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ) -> "QuantileDiscretizer": + """ + setParams(self, \\*, numBuckets=2, inputCol=None, outputCol=None, relativeError=0.001, \ + handleInvalid="error", numBucketsArray=None, inputCols=None, outputCols=None) + Set the params for the QuantileDiscretizer + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setNumBuckets(self, value: int) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`numBuckets`. + """ + return self._set(numBuckets=value) + + @since("2.0.0") + def getNumBuckets(self) -> int: + """ + Gets the value of numBuckets or its default value. + """ + return self.getOrDefault(self.numBuckets) + + @since("3.0.0") + def setNumBucketsArray(self, value: List[int]) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`numBucketsArray`. + """ + return self._set(numBucketsArray=value) + + @since("3.0.0") + def getNumBucketsArray(self) -> List[int]: + """ + Gets the value of numBucketsArray or its default value. + """ + return self.getOrDefault(self.numBucketsArray) + + @since("2.0.0") + def setRelativeError(self, value: float) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`relativeError`. + """ + return self._set(relativeError=value) + + def setInputCol(self, value: str) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + def setHandleInvalid(self, value: str) -> "QuantileDiscretizer": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + def _create_model(self, java_model: "JavaObject") -> Bucketizer: + """ + Private method to convert the java_model to a Python model. + """ + if self.isSet(self.inputCol): + return Bucketizer( + splits=list(java_model.getSplits()), + inputCol=self.getInputCol(), + outputCol=self.getOutputCol(), + handleInvalid=self.getHandleInvalid(), + ) + else: + splitsArrayList = [list(x) for x in list(java_model.getSplitsArray())] + return Bucketizer( + splitsArray=splitsArrayList, + inputCols=self.getInputCols(), + outputCols=self.getOutputCols(), + handleInvalid=self.getHandleInvalid(), + ) + + +class _RobustScalerParams(HasInputCol, HasOutputCol, HasRelativeError): + """ + Params for :py:class:`RobustScaler` and :py:class:`RobustScalerModel`. + + .. versionadded:: 3.0.0 + """ + + lower: Param[float] = Param( + Params._dummy(), + "lower", + "Lower quantile to calculate quantile range", + typeConverter=TypeConverters.toFloat, + ) + upper: Param[float] = Param( + Params._dummy(), + "upper", + "Upper quantile to calculate quantile range", + typeConverter=TypeConverters.toFloat, + ) + withCentering: Param[bool] = Param( + Params._dummy(), + "withCentering", + "Whether to center data with median", + typeConverter=TypeConverters.toBoolean, + ) + withScaling: Param[bool] = Param( + Params._dummy(), + "withScaling", + "Whether to scale the data to " "quantile range", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self, *args: Any): + super(_RobustScalerParams, self).__init__(*args) + self._setDefault( + lower=0.25, upper=0.75, withCentering=False, withScaling=True, relativeError=0.001 + ) + + @since("3.0.0") + def getLower(self) -> float: + """ + Gets the value of lower or its default value. + """ + return self.getOrDefault(self.lower) + + @since("3.0.0") + def getUpper(self) -> float: + """ + Gets the value of upper or its default value. + """ + return self.getOrDefault(self.upper) + + @since("3.0.0") + def getWithCentering(self) -> bool: + """ + Gets the value of withCentering or its default value. + """ + return self.getOrDefault(self.withCentering) + + @since("3.0.0") + def getWithScaling(self) -> bool: + """ + Gets the value of withScaling or its default value. + """ + return self.getOrDefault(self.withScaling) + + +@inherit_doc +class RobustScaler( + JavaEstimator, _RobustScalerParams, JavaMLReadable["RobustScaler"], JavaMLWritable +): + """ + RobustScaler removes the median and scales the data according to the quantile range. + The quantile range is by default IQR (Interquartile Range, quantile range between the + 1st quartile = 25th quantile and the 3rd quartile = 75th quantile) but can be configured. + Centering and scaling happen independently on each feature by computing the relevant + statistics on the samples in the training set. Median and quantile range are then + stored to be used on later data using the transform method. + Note that NaN values are ignored in the computation of medians and ranges. + + .. versionadded:: 3.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> data = [(0, Vectors.dense([0.0, 0.0]),), + ... (1, Vectors.dense([1.0, -1.0]),), + ... (2, Vectors.dense([2.0, -2.0]),), + ... (3, Vectors.dense([3.0, -3.0]),), + ... (4, Vectors.dense([4.0, -4.0]),),] + >>> df = spark.createDataFrame(data, ["id", "features"]) + >>> scaler = RobustScaler() + >>> scaler.setInputCol("features") + RobustScaler... + >>> scaler.setOutputCol("scaled") + RobustScaler... + >>> model = scaler.fit(df) + >>> model.setOutputCol("output") + RobustScalerModel... + >>> model.median + DenseVector([2.0, -2.0]) + >>> model.range + DenseVector([2.0, 2.0]) + >>> model.transform(df).collect()[1].output + DenseVector([0.5, -0.5]) + >>> scalerPath = temp_path + "/robust-scaler" + >>> scaler.save(scalerPath) + >>> loadedScaler = RobustScaler.load(scalerPath) + >>> loadedScaler.getWithCentering() == scaler.getWithCentering() + True + >>> loadedScaler.getWithScaling() == scaler.getWithScaling() + True + >>> modelPath = temp_path + "/robust-scaler-model" + >>> model.save(modelPath) + >>> loadedModel = RobustScalerModel.load(modelPath) + >>> loadedModel.median == model.median + True + >>> loadedModel.range == model.range + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + lower: float = 0.25, + upper: float = 0.75, + withCentering: bool = False, + withScaling: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + ): + """ + __init__(self, \\*, lower=0.25, upper=0.75, withCentering=False, withScaling=True, \ + inputCol=None, outputCol=None, relativeError=0.001) + """ + super(RobustScaler, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.RobustScaler", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.0.0") + def setParams( + self, + *, + lower: float = 0.25, + upper: float = 0.75, + withCentering: bool = False, + withScaling: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + relativeError: float = 0.001, + ) -> "RobustScaler": + """ + setParams(self, \\*, lower=0.25, upper=0.75, withCentering=False, withScaling=True, \ + inputCol=None, outputCol=None, relativeError=0.001) + Sets params for this RobustScaler. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.0.0") + def setLower(self, value: float) -> "RobustScaler": + """ + Sets the value of :py:attr:`lower`. + """ + return self._set(lower=value) + + @since("3.0.0") + def setUpper(self, value: float) -> "RobustScaler": + """ + Sets the value of :py:attr:`upper`. + """ + return self._set(upper=value) + + @since("3.0.0") + def setWithCentering(self, value: bool) -> "RobustScaler": + """ + Sets the value of :py:attr:`withCentering`. + """ + return self._set(withCentering=value) + + @since("3.0.0") + def setWithScaling(self, value: bool) -> "RobustScaler": + """ + Sets the value of :py:attr:`withScaling`. + """ + return self._set(withScaling=value) + + @since("3.0.0") + def setInputCol(self, value: str) -> "RobustScaler": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "RobustScaler": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setRelativeError(self, value: float) -> "RobustScaler": + """ + Sets the value of :py:attr:`relativeError`. + """ + return self._set(relativeError=value) + + def _create_model(self, java_model: "JavaObject") -> "RobustScalerModel": + return RobustScalerModel(java_model) + + +class RobustScalerModel( + JavaModel, _RobustScalerParams, JavaMLReadable["RobustScalerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`RobustScaler`. + + .. versionadded:: 3.0.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "RobustScalerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "RobustScalerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("3.0.0") + def median(self) -> Vector: + """ + Median of the RobustScalerModel. + """ + return self._call_java("median") + + @property + @since("3.0.0") + def range(self) -> Vector: + """ + Quantile range of the RobustScalerModel. + """ + return self._call_java("range") + + +@inherit_doc +class RegexTokenizer( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["RegexTokenizer"], + JavaMLWritable, +): + """ + A regex based tokenizer that extracts tokens either by using the + provided regex pattern (in Java dialect) to split the text + (default) or repeatedly matching the regex (if gaps is false). + Optional parameters also allow filtering tokens using a minimal + length. + It returns an array of strings that can be empty. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> df = spark.createDataFrame([("A B c",)], ["text"]) + >>> reTokenizer = RegexTokenizer() + >>> reTokenizer.setInputCol("text") + RegexTokenizer... + >>> reTokenizer.setOutputCol("words") + RegexTokenizer... + >>> reTokenizer.transform(df).head() + Row(text='A B c', words=['a', 'b', 'c']) + >>> # Change a parameter. + >>> reTokenizer.setParams(outputCol="tokens").transform(df).head() + Row(text='A B c', tokens=['a', 'b', 'c']) + >>> # Temporarily modify a parameter. + >>> reTokenizer.transform(df, {reTokenizer.outputCol: "words"}).head() + Row(text='A B c', words=['a', 'b', 'c']) + >>> reTokenizer.transform(df).head() + Row(text='A B c', tokens=['a', 'b', 'c']) + >>> # Must use keyword arguments to specify params. + >>> reTokenizer.setParams("text") + Traceback (most recent call last): + ... + TypeError: Method setParams forces keyword arguments. + >>> regexTokenizerPath = temp_path + "/regex-tokenizer" + >>> reTokenizer.save(regexTokenizerPath) + >>> loadedReTokenizer = RegexTokenizer.load(regexTokenizerPath) + >>> loadedReTokenizer.getMinTokenLength() == reTokenizer.getMinTokenLength() + True + >>> loadedReTokenizer.getGaps() == reTokenizer.getGaps() + True + >>> loadedReTokenizer.transform(df).take(1) == reTokenizer.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + minTokenLength: Param[int] = Param( + Params._dummy(), + "minTokenLength", + "minimum token length (>= 0)", + typeConverter=TypeConverters.toInt, + ) + gaps: Param[bool] = Param( + Params._dummy(), + "gaps", + "whether regex splits on gaps (True) or matches tokens " + "(False)", + ) + pattern: Param[str] = Param( + Params._dummy(), + "pattern", + "regex pattern (Java dialect) used for tokenizing", + typeConverter=TypeConverters.toString, + ) + toLowercase: Param[bool] = Param( + Params._dummy(), + "toLowercase", + "whether to convert all characters to " + "lowercase before tokenizing", + typeConverter=TypeConverters.toBoolean, + ) + + @keyword_only + def __init__( + self, + *, + minTokenLength: int = 1, + gaps: bool = True, + pattern: str = "\\s+", + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + toLowercase: bool = True, + ): + """ + __init__(self, \\*, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None, \ + outputCol=None, toLowercase=True) + """ + super(RegexTokenizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.RegexTokenizer", self.uid) + self._setDefault(minTokenLength=1, gaps=True, pattern="\\s+", toLowercase=True) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + minTokenLength: int = 1, + gaps: bool = True, + pattern: str = "\\s+", + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + toLowercase: bool = True, + ) -> "RegexTokenizer": + """ + setParams(self, \\*, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None, \ + outputCol=None, toLowercase=True) + Sets params for this RegexTokenizer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setMinTokenLength(self, value: int) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`minTokenLength`. + """ + return self._set(minTokenLength=value) + + @since("1.4.0") + def getMinTokenLength(self) -> int: + """ + Gets the value of minTokenLength or its default value. + """ + return self.getOrDefault(self.minTokenLength) + + @since("1.4.0") + def setGaps(self, value: bool) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`gaps`. + """ + return self._set(gaps=value) + + @since("1.4.0") + def getGaps(self) -> bool: + """ + Gets the value of gaps or its default value. + """ + return self.getOrDefault(self.gaps) + + @since("1.4.0") + def setPattern(self, value: str) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`pattern`. + """ + return self._set(pattern=value) + + @since("1.4.0") + def getPattern(self) -> str: + """ + Gets the value of pattern or its default value. + """ + return self.getOrDefault(self.pattern) + + @since("2.0.0") + def setToLowercase(self, value: bool) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`toLowercase`. + """ + return self._set(toLowercase=value) + + @since("2.0.0") + def getToLowercase(self) -> bool: + """ + Gets the value of toLowercase or its default value. + """ + return self.getOrDefault(self.toLowercase) + + def setInputCol(self, value: str) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "RegexTokenizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class SQLTransformer(JavaTransformer, JavaMLReadable["SQLTransformer"], JavaMLWritable): + """ + Implements the transforms which are defined by SQL statement. + Currently we only support SQL syntax like `SELECT ... FROM __THIS__` + where `__THIS__` represents the underlying table of the input dataset. + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> df = spark.createDataFrame([(0, 1.0, 3.0), (2, 2.0, 5.0)], ["id", "v1", "v2"]) + >>> sqlTrans = SQLTransformer( + ... statement="SELECT *, (v1 + v2) AS v3, (v1 * v2) AS v4 FROM __THIS__") + >>> sqlTrans.transform(df).head() + Row(id=0, v1=1.0, v2=3.0, v3=4.0, v4=3.0) + >>> sqlTransformerPath = temp_path + "/sql-transformer" + >>> sqlTrans.save(sqlTransformerPath) + >>> loadedSqlTrans = SQLTransformer.load(sqlTransformerPath) + >>> loadedSqlTrans.getStatement() == sqlTrans.getStatement() + True + >>> loadedSqlTrans.transform(df).take(1) == sqlTrans.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + statement = Param( + Params._dummy(), "statement", "SQL statement", typeConverter=TypeConverters.toString + ) + + @keyword_only + def __init__(self, *, statement: Optional[str] = None): + """ + __init__(self, \\*, statement=None) + """ + super(SQLTransformer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.SQLTransformer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams(self, *, statement: Optional[str] = None) -> "SQLTransformer": + """ + setParams(self, \\*, statement=None) + Sets params for this SQLTransformer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setStatement(self, value: str) -> "SQLTransformer": + """ + Sets the value of :py:attr:`statement`. + """ + return self._set(statement=value) + + @since("1.6.0") + def getStatement(self) -> str: + """ + Gets the value of statement or its default value. + """ + return self.getOrDefault(self.statement) + + +class _StandardScalerParams(HasInputCol, HasOutputCol): + """ + Params for :py:class:`StandardScaler` and :py:class:`StandardScalerModel`. + + .. versionadded:: 3.0.0 + """ + + withMean: Param[bool] = Param( + Params._dummy(), "withMean", "Center data with mean", typeConverter=TypeConverters.toBoolean + ) + withStd: Param[bool] = Param( + Params._dummy(), + "withStd", + "Scale to unit standard deviation", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self, *args: Any): + super(_StandardScalerParams, self).__init__(*args) + self._setDefault(withMean=False, withStd=True) + + @since("1.4.0") + def getWithMean(self) -> bool: + """ + Gets the value of withMean or its default value. + """ + return self.getOrDefault(self.withMean) + + @since("1.4.0") + def getWithStd(self) -> bool: + """ + Gets the value of withStd or its default value. + """ + return self.getOrDefault(self.withStd) + + +@inherit_doc +class StandardScaler( + JavaEstimator["StandardScalerModel"], + _StandardScalerParams, + JavaMLReadable["StandardScaler"], + JavaMLWritable, +): + """ + Standardizes features by removing the mean and scaling to unit variance using column summary + statistics on the samples in the training set. + + The "unit std" is computed using the `corrected sample standard deviation \ + `_, + which is computed as the square root of the unbiased sample variance. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([0.0]),), (Vectors.dense([2.0]),)], ["a"]) + >>> standardScaler = StandardScaler() + >>> standardScaler.setInputCol("a") + StandardScaler... + >>> standardScaler.setOutputCol("scaled") + StandardScaler... + >>> model = standardScaler.fit(df) + >>> model.getInputCol() + 'a' + >>> model.setOutputCol("output") + StandardScalerModel... + >>> model.mean + DenseVector([1.0]) + >>> model.std + DenseVector([1.4142]) + >>> model.transform(df).collect()[1].output + DenseVector([1.4142]) + >>> standardScalerPath = temp_path + "/standard-scaler" + >>> standardScaler.save(standardScalerPath) + >>> loadedStandardScaler = StandardScaler.load(standardScalerPath) + >>> loadedStandardScaler.getWithMean() == standardScaler.getWithMean() + True + >>> loadedStandardScaler.getWithStd() == standardScaler.getWithStd() + True + >>> modelPath = temp_path + "/standard-scaler-model" + >>> model.save(modelPath) + >>> loadedModel = StandardScalerModel.load(modelPath) + >>> loadedModel.std == model.std + True + >>> loadedModel.mean == model.mean + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + withMean: bool = False, + withStd: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, withMean=False, withStd=True, inputCol=None, outputCol=None) + """ + super(StandardScaler, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.StandardScaler", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + withMean: bool = False, + withStd: bool = True, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "StandardScaler": + """ + setParams(self, \\*, withMean=False, withStd=True, inputCol=None, outputCol=None) + Sets params for this StandardScaler. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setWithMean(self, value: bool) -> "StandardScaler": + """ + Sets the value of :py:attr:`withMean`. + """ + return self._set(withMean=value) + + @since("1.4.0") + def setWithStd(self, value: bool) -> "StandardScaler": + """ + Sets the value of :py:attr:`withStd`. + """ + return self._set(withStd=value) + + def setInputCol(self, value: str) -> "StandardScaler": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "StandardScaler": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "StandardScalerModel": + return StandardScalerModel(java_model) + + +class StandardScalerModel( + JavaModel, + _StandardScalerParams, + JavaMLReadable["StandardScalerModel"], + JavaMLWritable, +): + """ + Model fitted by :py:class:`StandardScaler`. + + .. versionadded:: 1.4.0 + """ + + def setInputCol(self, value: str) -> "StandardScalerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "StandardScalerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.0.0") + def std(self) -> Vector: + """ + Standard deviation of the StandardScalerModel. + """ + return self._call_java("std") + + @property + @since("2.0.0") + def mean(self) -> Vector: + """ + Mean of the StandardScalerModel. + """ + return self._call_java("mean") + + +class _StringIndexerParams( + JavaParams, HasHandleInvalid, HasInputCol, HasOutputCol, HasInputCols, HasOutputCols +): + """ + Params for :py:class:`StringIndexer` and :py:class:`StringIndexerModel`. + """ + + stringOrderType: Param[str] = Param( + Params._dummy(), + "stringOrderType", + "How to order labels of string column. The first label after " + + "ordering is assigned an index of 0. Supported options: " + + "frequencyDesc, frequencyAsc, alphabetDesc, alphabetAsc. " + + "Default is frequencyDesc. In case of equal frequency when " + + "under frequencyDesc/Asc, the strings are further sorted " + + "alphabetically", + typeConverter=TypeConverters.toString, + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "how to handle invalid data (unseen " + + "or NULL values) in features and label column of string type. " + + "Options are 'skip' (filter out rows with invalid data), " + + "error (throw an error), or 'keep' (put invalid data " + + "in a special additional bucket, at index numLabels).", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_StringIndexerParams, self).__init__(*args) + self._setDefault(handleInvalid="error", stringOrderType="frequencyDesc") + + @since("2.3.0") + def getStringOrderType(self) -> str: + """ + Gets the value of :py:attr:`stringOrderType` or its default value 'frequencyDesc'. + """ + return self.getOrDefault(self.stringOrderType) + + +@inherit_doc +class StringIndexer( + JavaEstimator["StringIndexerModel"], + _StringIndexerParams, + JavaMLReadable["StringIndexer"], + JavaMLWritable, +): + """ + A label indexer that maps a string column of labels to an ML column of label indices. + If the input column is numeric, we cast it to string and index the string values. + The indices are in [0, numLabels). By default, this is ordered by label frequencies + so the most frequent label gets index 0. The ordering behavior is controlled by + setting :py:attr:`stringOrderType`. Its default value is 'frequencyDesc'. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> stringIndexer = StringIndexer(inputCol="label", outputCol="indexed", + ... stringOrderType="frequencyDesc") + >>> stringIndexer.setHandleInvalid("error") + StringIndexer... + >>> model = stringIndexer.fit(stringIndDf) + >>> model.setHandleInvalid("error") + StringIndexerModel... + >>> td = model.transform(stringIndDf) + >>> sorted(set([(i[0], i[1]) for i in td.select(td.id, td.indexed).collect()]), + ... key=lambda x: x[0]) + [(0, 0.0), (1, 2.0), (2, 1.0), (3, 0.0), (4, 0.0), (5, 1.0)] + >>> inverter = IndexToString(inputCol="indexed", outputCol="label2", labels=model.labels) + >>> itd = inverter.transform(td) + >>> sorted(set([(i[0], str(i[1])) for i in itd.select(itd.id, itd.label2).collect()]), + ... key=lambda x: x[0]) + [(0, 'a'), (1, 'b'), (2, 'c'), (3, 'a'), (4, 'a'), (5, 'c')] + >>> stringIndexerPath = temp_path + "/string-indexer" + >>> stringIndexer.save(stringIndexerPath) + >>> loadedIndexer = StringIndexer.load(stringIndexerPath) + >>> loadedIndexer.getHandleInvalid() == stringIndexer.getHandleInvalid() + True + >>> modelPath = temp_path + "/string-indexer-model" + >>> model.save(modelPath) + >>> loadedModel = StringIndexerModel.load(modelPath) + >>> loadedModel.labels == model.labels + True + >>> indexToStringPath = temp_path + "/index-to-string" + >>> inverter.save(indexToStringPath) + >>> loadedInverter = IndexToString.load(indexToStringPath) + >>> loadedInverter.getLabels() == inverter.getLabels() + True + >>> loadedModel.transform(stringIndDf).take(1) == model.transform(stringIndDf).take(1) + True + >>> stringIndexer.getStringOrderType() + 'frequencyDesc' + >>> stringIndexer = StringIndexer(inputCol="label", outputCol="indexed", handleInvalid="error", + ... stringOrderType="alphabetDesc") + >>> model = stringIndexer.fit(stringIndDf) + >>> td = model.transform(stringIndDf) + >>> sorted(set([(i[0], i[1]) for i in td.select(td.id, td.indexed).collect()]), + ... key=lambda x: x[0]) + [(0, 2.0), (1, 1.0), (2, 0.0), (3, 2.0), (4, 2.0), (5, 0.0)] + >>> fromlabelsModel = StringIndexerModel.from_labels(["a", "b", "c"], + ... inputCol="label", outputCol="indexed", handleInvalid="error") + >>> result = fromlabelsModel.transform(stringIndDf) + >>> sorted(set([(i[0], i[1]) for i in result.select(result.id, result.indexed).collect()]), + ... key=lambda x: x[0]) + [(0, 0.0), (1, 1.0), (2, 2.0), (3, 0.0), (4, 0.0), (5, 2.0)] + >>> testData = sc.parallelize([Row(id=0, label1="a", label2="e"), + ... Row(id=1, label1="b", label2="f"), + ... Row(id=2, label1="c", label2="e"), + ... Row(id=3, label1="a", label2="f"), + ... Row(id=4, label1="a", label2="f"), + ... Row(id=5, label1="c", label2="f")], 3) + >>> multiRowDf = spark.createDataFrame(testData) + >>> inputs = ["label1", "label2"] + >>> outputs = ["index1", "index2"] + >>> stringIndexer = StringIndexer(inputCols=inputs, outputCols=outputs) + >>> model = stringIndexer.fit(multiRowDf) + >>> result = model.transform(multiRowDf) + >>> sorted(set([(i[0], i[1], i[2]) for i in result.select(result.id, result.index1, + ... result.index2).collect()]), key=lambda x: x[0]) + [(0, 0.0, 1.0), (1, 2.0, 0.0), (2, 1.0, 1.0), (3, 0.0, 0.0), (4, 0.0, 0.0), (5, 1.0, 0.0)] + >>> fromlabelsModel = StringIndexerModel.from_arrays_of_labels([["a", "b", "c"], ["e", "f"]], + ... inputCols=inputs, outputCols=outputs) + >>> result = fromlabelsModel.transform(multiRowDf) + >>> sorted(set([(i[0], i[1], i[2]) for i in result.select(result.id, result.index1, + ... result.index2).collect()]), key=lambda x: x[0]) + [(0, 0.0, 0.0), (1, 1.0, 1.0), (2, 2.0, 0.0), (3, 0.0, 1.0), (4, 0.0, 1.0), (5, 2.0, 1.0)] + """ + + _input_kwargs: Dict[str, Any] + + @overload + def __init__( + self, + *, + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + handleInvalid: str = ..., + stringOrderType: str = ..., + ): + ... + + @overload + def __init__( + self, + *, + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + handleInvalid: str = ..., + stringOrderType: str = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + handleInvalid: str = "error", + stringOrderType: str = "frequencyDesc", + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, inputCols=None, outputCols=None, \ + handleInvalid="error", stringOrderType="frequencyDesc") + """ + super(StringIndexer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.StringIndexer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + handleInvalid: str = ..., + stringOrderType: str = ..., + ) -> "StringIndexer": + ... + + @overload + def setParams( + self, + *, + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + handleInvalid: str = ..., + stringOrderType: str = ..., + ) -> "StringIndexer": + ... + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + handleInvalid: str = "error", + stringOrderType: str = "frequencyDesc", + ) -> "StringIndexer": + """ + setParams(self, \\*, inputCol=None, outputCol=None, inputCols=None, outputCols=None, \ + handleInvalid="error", stringOrderType="frequencyDesc") + Sets params for this StringIndexer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "StringIndexerModel": + return StringIndexerModel(java_model) + + @since("2.3.0") + def setStringOrderType(self, value: str) -> "StringIndexer": + """ + Sets the value of :py:attr:`stringOrderType`. + """ + return self._set(stringOrderType=value) + + def setInputCol(self, value: str) -> "StringIndexer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "StringIndexer": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "StringIndexer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "StringIndexer": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + def setHandleInvalid(self, value: str) -> "StringIndexer": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + +class StringIndexerModel( + JavaModel, _StringIndexerParams, JavaMLReadable["StringIndexerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`StringIndexer`. + + .. versionadded:: 1.4.0 + """ + + def setInputCol(self, value: str) -> "StringIndexerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "StringIndexerModel": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "StringIndexerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "StringIndexerModel": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @since("2.4.0") + def setHandleInvalid(self, value: str) -> "StringIndexerModel": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + @classmethod + @since("2.4.0") + def from_labels( + cls, + labels: List[str], + inputCol: str, + outputCol: Optional[str] = None, + handleInvalid: Optional[str] = None, + ) -> "StringIndexerModel": + """ + Construct the model directly from an array of label strings, + requires an active SparkContext. + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._gateway is not None + java_class = sc._gateway.jvm.java.lang.String + jlabels = StringIndexerModel._new_java_array(labels, java_class) + model = StringIndexerModel._create_from_java_class( + "org.apache.spark.ml.feature.StringIndexerModel", jlabels + ) + model.setInputCol(inputCol) + if outputCol is not None: + model.setOutputCol(outputCol) + if handleInvalid is not None: + model.setHandleInvalid(handleInvalid) + return model + + @classmethod + @since("3.0.0") + def from_arrays_of_labels( + cls, + arrayOfLabels: List[List[str]], + inputCols: List[str], + outputCols: Optional[List[str]] = None, + handleInvalid: Optional[str] = None, + ) -> "StringIndexerModel": + """ + Construct the model directly from an array of array of label strings, + requires an active SparkContext. + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._gateway is not None + java_class = sc._gateway.jvm.java.lang.String + jlabels = StringIndexerModel._new_java_array(arrayOfLabels, java_class) + model = StringIndexerModel._create_from_java_class( + "org.apache.spark.ml.feature.StringIndexerModel", jlabels + ) + model.setInputCols(inputCols) + if outputCols is not None: + model.setOutputCols(outputCols) + if handleInvalid is not None: + model.setHandleInvalid(handleInvalid) + return model + + @property + @since("1.5.0") + def labels(self) -> List[str]: + """ + Ordered list of labels, corresponding to indices to be assigned. + + .. deprecated:: 3.1.0 + It will be removed in future versions. Use `labelsArray` method instead. + """ + return self._call_java("labels") + + @property + @since("3.0.2") + def labelsArray(self) -> List[str]: + """ + Array of ordered list of labels, corresponding to indices to be assigned + for each input column. + """ + return self._call_java("labelsArray") + + +@inherit_doc +class IndexToString( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["IndexToString"], + JavaMLWritable, +): + """ + A :py:class:`pyspark.ml.base.Transformer` that maps a column of indices back to a new column of + corresponding string values. + The index-string mapping is either from the ML attributes of the input column, + or from user-supplied labels (which take precedence over ML attributes). + + .. versionadded:: 1.6.0 + + See Also + -------- + StringIndexer : for converting categorical values into category indices + """ + + _input_kwargs: Dict[str, Any] + + labels: Param[List[str]] = Param( + Params._dummy(), + "labels", + "Optional array of labels specifying index-string mapping." + + " If not provided or if empty, then metadata from inputCol is used instead.", + typeConverter=TypeConverters.toListString, + ) + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + labels: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, labels=None) + """ + super(IndexToString, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.IndexToString", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + labels: Optional[List[str]] = None, + ) -> "IndexToString": + """ + setParams(self, \\*, inputCol=None, outputCol=None, labels=None) + Sets params for this IndexToString. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setLabels(self, value: List[str]) -> "IndexToString": + """ + Sets the value of :py:attr:`labels`. + """ + return self._set(labels=value) + + @since("1.6.0") + def getLabels(self) -> List[str]: + """ + Gets the value of :py:attr:`labels` or its default value. + """ + return self.getOrDefault(self.labels) + + def setInputCol(self, value: str) -> "IndexToString": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "IndexToString": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +class StopWordsRemover( + JavaTransformer, + HasInputCol, + HasOutputCol, + HasInputCols, + HasOutputCols, + JavaMLReadable["StopWordsRemover"], + JavaMLWritable, +): + """ + A feature transformer that filters out stop words from input. + Since 3.0.0, :py:class:`StopWordsRemover` can filter out multiple columns at once by setting + the :py:attr:`inputCols` parameter. Note that when both the :py:attr:`inputCol` and + :py:attr:`inputCols` parameters are set, an Exception will be thrown. + + .. versionadded:: 1.6.0 + + Notes + ----- + null values from input array are preserved unless adding null to stopWords explicitly. + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],)], ["text"]) + >>> remover = StopWordsRemover(stopWords=["b"]) + >>> remover.setInputCol("text") + StopWordsRemover... + >>> remover.setOutputCol("words") + StopWordsRemover... + >>> remover.transform(df).head().words == ['a', 'c'] + True + >>> stopWordsRemoverPath = temp_path + "/stopwords-remover" + >>> remover.save(stopWordsRemoverPath) + >>> loadedRemover = StopWordsRemover.load(stopWordsRemoverPath) + >>> loadedRemover.getStopWords() == remover.getStopWords() + True + >>> loadedRemover.getCaseSensitive() == remover.getCaseSensitive() + True + >>> loadedRemover.transform(df).take(1) == remover.transform(df).take(1) + True + >>> df2 = spark.createDataFrame([(["a", "b", "c"], ["a", "b"])], ["text1", "text2"]) + >>> remover2 = StopWordsRemover(stopWords=["b"]) + >>> remover2.setInputCols(["text1", "text2"]).setOutputCols(["words1", "words2"]) + StopWordsRemover... + >>> remover2.transform(df2).show() + +---------+------+------+------+ + | text1| text2|words1|words2| + +---------+------+------+------+ + |[a, b, c]|[a, b]|[a, c]| [a]| + +---------+------+------+------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + stopWords: Param[List[str]] = Param( + Params._dummy(), + "stopWords", + "The words to be filtered out", + typeConverter=TypeConverters.toListString, + ) + caseSensitive: Param[bool] = Param( + Params._dummy(), + "caseSensitive", + "whether to do a case sensitive " + "comparison over the stop words", + typeConverter=TypeConverters.toBoolean, + ) + locale: Param[str] = Param( + Params._dummy(), + "locale", + "locale of the input. ignored when case sensitive " + "is true", + typeConverter=TypeConverters.toString, + ) + + @overload + def __init__( + self, + *, + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + stopWords: Optional[List[str]] = ..., + caseSensitive: bool = ..., + locale: Optional[str] = ..., + ): + ... + + @overload + def __init__( + self, + *, + stopWords: Optional[List[str]] = ..., + caseSensitive: bool = ..., + locale: Optional[str] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ): + ... + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + stopWords: Optional[List[str]] = None, + caseSensitive: bool = False, + locale: Optional[str] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, stopWords=None, caseSensitive=false, \ + locale=None, inputCols=None, outputCols=None) + """ + super(StopWordsRemover, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.StopWordsRemover", self.uid + ) + self._setDefault( + stopWords=StopWordsRemover.loadDefaultStopWords("english"), + caseSensitive=False, + locale=self._java_obj.getLocale(), + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @overload + def setParams( + self, + *, + inputCol: Optional[str] = ..., + outputCol: Optional[str] = ..., + stopWords: Optional[List[str]] = ..., + caseSensitive: bool = ..., + locale: Optional[str] = ..., + ) -> "StopWordsRemover": + ... + + @overload + def setParams( + self, + *, + stopWords: Optional[List[str]] = ..., + caseSensitive: bool = ..., + locale: Optional[str] = ..., + inputCols: Optional[List[str]] = ..., + outputCols: Optional[List[str]] = ..., + ) -> "StopWordsRemover": + ... + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + stopWords: Optional[List[str]] = None, + caseSensitive: bool = False, + locale: Optional[str] = None, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + ) -> "StopWordsRemover": + """ + setParams(self, \\*, inputCol=None, outputCol=None, stopWords=None, caseSensitive=false, \ + locale=None, inputCols=None, outputCols=None) + Sets params for this StopWordRemover. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setStopWords(self, value: List[str]) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`stopWords`. + """ + return self._set(stopWords=value) + + @since("1.6.0") + def getStopWords(self) -> List[str]: + """ + Gets the value of :py:attr:`stopWords` or its default value. + """ + return self.getOrDefault(self.stopWords) + + @since("1.6.0") + def setCaseSensitive(self, value: bool) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`caseSensitive`. + """ + return self._set(caseSensitive=value) + + @since("1.6.0") + def getCaseSensitive(self) -> bool: + """ + Gets the value of :py:attr:`caseSensitive` or its default value. + """ + return self.getOrDefault(self.caseSensitive) + + @since("2.4.0") + def setLocale(self, value: str) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`locale`. + """ + return self._set(locale=value) + + @since("2.4.0") + def getLocale(self) -> str: + """ + Gets the value of :py:attr:`locale`. + """ + return self.getOrDefault(self.locale) + + def setInputCol(self, value: str) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("3.0.0") + def setInputCols(self, value: List[str]) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + @since("3.0.0") + def setOutputCols(self, value: List[str]) -> "StopWordsRemover": + """ + Sets the value of :py:attr:`outputCols`. + """ + return self._set(outputCols=value) + + @staticmethod + @since("2.0.0") + def loadDefaultStopWords(language: str) -> List[str]: + """ + Loads the default stop words for the given language. + Supported languages: danish, dutch, english, finnish, french, german, hungarian, + italian, norwegian, portuguese, russian, spanish, swedish, turkish + """ + stopWordsObj = _jvm().org.apache.spark.ml.feature.StopWordsRemover + return list(stopWordsObj.loadDefaultStopWords(language)) + + +@inherit_doc +class Tokenizer( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["Tokenizer"], + JavaMLWritable, +): + """ + A tokenizer that converts the input string to lowercase and then + splits it by white spaces. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> df = spark.createDataFrame([("a b c",)], ["text"]) + >>> tokenizer = Tokenizer(outputCol="words") + >>> tokenizer.setInputCol("text") + Tokenizer... + >>> tokenizer.transform(df).head() + Row(text='a b c', words=['a', 'b', 'c']) + >>> # Change a parameter. + >>> tokenizer.setParams(outputCol="tokens").transform(df).head() + Row(text='a b c', tokens=['a', 'b', 'c']) + >>> # Temporarily modify a parameter. + >>> tokenizer.transform(df, {tokenizer.outputCol: "words"}).head() + Row(text='a b c', words=['a', 'b', 'c']) + >>> tokenizer.transform(df).head() + Row(text='a b c', tokens=['a', 'b', 'c']) + >>> # Must use keyword arguments to specify params. + >>> tokenizer.setParams("text") + Traceback (most recent call last): + ... + TypeError: Method setParams forces keyword arguments. + >>> tokenizerPath = temp_path + "/tokenizer" + >>> tokenizer.save(tokenizerPath) + >>> loadedTokenizer = Tokenizer.load(tokenizerPath) + >>> loadedTokenizer.transform(df).head().tokens == tokenizer.transform(df).head().tokens + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, inputCol: Optional[str] = None, outputCol: Optional[str] = None): + """ + __init__(self, \\*, inputCol=None, outputCol=None) + """ + super(Tokenizer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Tokenizer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.3.0") + def setParams( + self, *, inputCol: Optional[str] = None, outputCol: Optional[str] = None + ) -> "Tokenizer": + """ + setParams(self, \\*, inputCol=None, outputCol=None) + Sets params for this Tokenizer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def setInputCol(self, value: str) -> "Tokenizer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "Tokenizer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +@inherit_doc +class VectorAssembler( + JavaTransformer, + HasInputCols, + HasOutputCol, + HasHandleInvalid, + JavaMLReadable["VectorAssembler"], + JavaMLWritable, +): + """ + A feature transformer that merges multiple columns into a vector column. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> df = spark.createDataFrame([(1, 0, 3)], ["a", "b", "c"]) + >>> vecAssembler = VectorAssembler(outputCol="features") + >>> vecAssembler.setInputCols(["a", "b", "c"]) + VectorAssembler... + >>> vecAssembler.transform(df).head().features + DenseVector([1.0, 0.0, 3.0]) + >>> vecAssembler.setParams(outputCol="freqs").transform(df).head().freqs + DenseVector([1.0, 0.0, 3.0]) + >>> params = {vecAssembler.inputCols: ["b", "a"], vecAssembler.outputCol: "vector"} + >>> vecAssembler.transform(df, params).head().vector + DenseVector([0.0, 1.0]) + >>> vectorAssemblerPath = temp_path + "/vector-assembler" + >>> vecAssembler.save(vectorAssemblerPath) + >>> loadedAssembler = VectorAssembler.load(vectorAssemblerPath) + >>> loadedAssembler.transform(df).head().freqs == vecAssembler.transform(df).head().freqs + True + >>> dfWithNullsAndNaNs = spark.createDataFrame( + ... [(1.0, 2.0, None), (3.0, float("nan"), 4.0), (5.0, 6.0, 7.0)], ["a", "b", "c"]) + >>> vecAssembler2 = VectorAssembler(inputCols=["a", "b", "c"], outputCol="features", + ... handleInvalid="keep") + >>> vecAssembler2.transform(dfWithNullsAndNaNs).show() + +---+---+----+-------------+ + | a| b| c| features| + +---+---+----+-------------+ + |1.0|2.0|NULL|[1.0,2.0,NaN]| + |3.0|NaN| 4.0|[3.0,NaN,4.0]| + |5.0|6.0| 7.0|[5.0,6.0,7.0]| + +---+---+----+-------------+ + ... + >>> vecAssembler2.setParams(handleInvalid="skip").transform(dfWithNullsAndNaNs).show() + +---+---+---+-------------+ + | a| b| c| features| + +---+---+---+-------------+ + |5.0|6.0|7.0|[5.0,6.0,7.0]| + +---+---+---+-------------+ + ... + """ + + _input_kwargs: Dict[str, Any] + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "How to handle invalid data (NULL " + + "and NaN values). Options are 'skip' (filter out rows with invalid " + + "data), 'error' (throw an error), or 'keep' (return relevant number " + + "of NaN in the output). Column lengths are taken from the size of ML " + + "Attribute Group, which can be set using `VectorSizeHint` in a " + + "pipeline before `VectorAssembler`. Column lengths can also be " + + "inferred from first rows of the data since it is safe to do so but " + + "only in case of 'error' or 'skip').", + typeConverter=TypeConverters.toString, + ) + + @keyword_only + def __init__( + self, + *, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + ): + """ + __init__(self, \\*, inputCols=None, outputCol=None, handleInvalid="error") + """ + super(VectorAssembler, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.VectorAssembler", self.uid) + self._setDefault(handleInvalid="error") + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + inputCols: Optional[List[str]] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + ) -> "VectorAssembler": + """ + setParams(self, \\*, inputCols=None, outputCol=None, handleInvalid="error") + Sets params for this VectorAssembler. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def setInputCols(self, value: List[str]) -> "VectorAssembler": + """ + Sets the value of :py:attr:`inputCols`. + """ + return self._set(inputCols=value) + + def setOutputCol(self, value: str) -> "VectorAssembler": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setHandleInvalid(self, value: str) -> "VectorAssembler": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + +class _VectorIndexerParams(HasInputCol, HasOutputCol, HasHandleInvalid): + """ + Params for :py:class:`VectorIndexer` and :py:class:`VectorIndexerModel`. + + .. versionadded:: 3.0.0 + """ + + maxCategories: Param[int] = Param( + Params._dummy(), + "maxCategories", + "Threshold for the number of values a categorical feature can take " + + "(>= 2). If a feature is found to have > maxCategories values, then " + + "it is declared continuous.", + typeConverter=TypeConverters.toInt, + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "How to handle invalid data " + + "(unseen labels or NULL values). Options are 'skip' (filter out " + + "rows with invalid data), 'error' (throw an error), or 'keep' (put " + + "invalid data in a special additional bucket, at index of the number " + + "of categories of the feature).", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_VectorIndexerParams, self).__init__(*args) + self._setDefault(maxCategories=20, handleInvalid="error") + + @since("1.4.0") + def getMaxCategories(self) -> int: + """ + Gets the value of maxCategories or its default value. + """ + return self.getOrDefault(self.maxCategories) + + +@inherit_doc +class VectorIndexer( + JavaEstimator["VectorIndexerModel"], + _VectorIndexerParams, + HasHandleInvalid, + JavaMLReadable["VectorIndexer"], + JavaMLWritable, +): + """ + Class for indexing categorical feature columns in a dataset of `Vector`. + + This has 2 usage modes: + - Automatically identify categorical features (default behavior) + - This helps process a dataset of unknown vectors into a dataset with some continuous + features and some categorical features. The choice between continuous and categorical + is based upon a maxCategories parameter. + - Set maxCategories to the maximum number of categorical any categorical feature should + have. + - E.g.: Feature 0 has unique values {-1.0, 0.0}, and feature 1 values {1.0, 3.0, 5.0}. + If maxCategories = 2, then feature 0 will be declared categorical and use indices {0, 1}, + and feature 1 will be declared continuous. + - Index all features, if all features are categorical + - If maxCategories is set to be very large, then this will build an index of unique + values for all features. + - Warning: This can cause problems if features are continuous since this will collect ALL + unique values to the driver. + - E.g.: Feature 0 has unique values {-1.0, 0.0}, and feature 1 values {1.0, 3.0, 5.0}. + If maxCategories >= 3, then both features will be declared categorical. + + This returns a model which can transform categorical features to use 0-based indices. + + Index stability: + - This is not guaranteed to choose the same category index across multiple runs. + - If a categorical feature includes value 0, then this is guaranteed to map value 0 to + index 0. This maintains vector sparsity. + - More stability may be added in the future. + + TODO: Future extensions: The following functionality is planned for the future: + - Preserve metadata in transform; if a feature's metadata is already present, + do not recompute. + - Specify certain features to not index, either via a parameter or via existing metadata. + - Add warning if a categorical feature has only 1 category. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([(Vectors.dense([-1.0, 0.0]),), + ... (Vectors.dense([0.0, 1.0]),), (Vectors.dense([0.0, 2.0]),)], ["a"]) + >>> indexer = VectorIndexer(maxCategories=2, inputCol="a") + >>> indexer.setOutputCol("indexed") + VectorIndexer... + >>> model = indexer.fit(df) + >>> indexer.getHandleInvalid() + 'error' + >>> model.setOutputCol("output") + VectorIndexerModel... + >>> model.transform(df).head().output + DenseVector([1.0, 0.0]) + >>> model.numFeatures + 2 + >>> model.categoryMaps + {0: {0.0: 0, -1.0: 1}} + >>> indexer.setParams(outputCol="test").fit(df).transform(df).collect()[1].test + DenseVector([0.0, 1.0]) + >>> params = {indexer.maxCategories: 3, indexer.outputCol: "vector"} + >>> model2 = indexer.fit(df, params) + >>> model2.transform(df).head().vector + DenseVector([1.0, 0.0]) + >>> vectorIndexerPath = temp_path + "/vector-indexer" + >>> indexer.save(vectorIndexerPath) + >>> loadedIndexer = VectorIndexer.load(vectorIndexerPath) + >>> loadedIndexer.getMaxCategories() == indexer.getMaxCategories() + True + >>> modelPath = temp_path + "/vector-indexer-model" + >>> model.save(modelPath) + >>> loadedModel = VectorIndexerModel.load(modelPath) + >>> loadedModel.numFeatures == model.numFeatures + True + >>> loadedModel.categoryMaps == model.categoryMaps + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + >>> dfWithInvalid = spark.createDataFrame([(Vectors.dense([3.0, 1.0]),)], ["a"]) + >>> indexer.getHandleInvalid() + 'error' + >>> model3 = indexer.setHandleInvalid("skip").fit(df) + >>> model3.transform(dfWithInvalid).count() + 0 + >>> model4 = indexer.setParams(handleInvalid="keep", outputCol="indexed").fit(df) + >>> model4.transform(dfWithInvalid).head().indexed + DenseVector([2.0, 1.0]) + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + maxCategories: int = 20, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + ): + """ + __init__(self, \\*, maxCategories=20, inputCol=None, outputCol=None, handleInvalid="error") + """ + super(VectorIndexer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.VectorIndexer", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + maxCategories: int = 20, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + handleInvalid: str = "error", + ) -> "VectorIndexer": + """ + setParams(self, \\*, maxCategories=20, inputCol=None, outputCol=None, handleInvalid="error") + Sets params for this VectorIndexer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setMaxCategories(self, value: int) -> "VectorIndexer": + """ + Sets the value of :py:attr:`maxCategories`. + """ + return self._set(maxCategories=value) + + def setInputCol(self, value: str) -> "VectorIndexer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "VectorIndexer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setHandleInvalid(self, value: str) -> "VectorIndexer": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + def _create_model(self, java_model: "JavaObject") -> "VectorIndexerModel": + return VectorIndexerModel(java_model) + + +class VectorIndexerModel( + JavaModel, _VectorIndexerParams, JavaMLReadable["VectorIndexerModel"], JavaMLWritable +): + """ + Model fitted by :py:class:`VectorIndexer`. + + Transform categorical features to use 0-based indices instead of their original values. + - Categorical features are mapped to indices. + - Continuous features (columns) are left unchanged. + + This also appends metadata to the output column, marking features as Numeric (continuous), + Nominal (categorical), or Binary (either continuous or categorical). + Non-ML metadata is not carried over from the input to the output column. + + This maintains vector sparsity. + + .. versionadded:: 1.4.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "VectorIndexerModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "VectorIndexerModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("1.4.0") + def numFeatures(self) -> int: + """ + Number of features, i.e., length of Vectors which this transforms. + """ + return self._call_java("numFeatures") + + @property + @since("1.4.0") + def categoryMaps(self) -> Dict[int, Tuple[float, int]]: + """ + Feature value index. Keys are categorical feature indices (column indices). + Values are maps from original features values to 0-based category indices. + If a feature is not in this map, it is treated as continuous. + """ + return self._call_java("javaCategoryMaps") + + +@inherit_doc +class VectorSlicer( + JavaTransformer, + HasInputCol, + HasOutputCol, + JavaMLReadable["VectorSlicer"], + JavaMLWritable, +): + """ + This class takes a feature vector and outputs a new feature vector with a subarray + of the original features. + + The subset of features can be specified with either indices (`setIndices()`) + or names (`setNames()`). At least one feature must be selected. Duplicate features + are not allowed, so there can be no overlap between selected indices and names. + + The output vector will order features with the selected indices first (in the order given), + followed by the selected names (in the order given). + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (Vectors.dense([-2.0, 2.3, 0.0, 0.0, 1.0]),), + ... (Vectors.dense([0.0, 0.0, 0.0, 0.0, 0.0]),), + ... (Vectors.dense([0.6, -1.1, -3.0, 4.5, 3.3]),)], ["features"]) + >>> vs = VectorSlicer(outputCol="sliced", indices=[1, 4]) + >>> vs.setInputCol("features") + VectorSlicer... + >>> vs.transform(df).head().sliced + DenseVector([2.3, 1.0]) + >>> vectorSlicerPath = temp_path + "/vector-slicer" + >>> vs.save(vectorSlicerPath) + >>> loadedVs = VectorSlicer.load(vectorSlicerPath) + >>> loadedVs.getIndices() == vs.getIndices() + True + >>> loadedVs.getNames() == vs.getNames() + True + >>> loadedVs.transform(df).take(1) == vs.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + indices: Param[List[int]] = Param( + Params._dummy(), + "indices", + "An array of indices to select features from " + + "a vector column. There can be no overlap with names.", + typeConverter=TypeConverters.toListInt, + ) + names: Param[List[str]] = Param( + Params._dummy(), + "names", + "An array of feature names to select features from " + + "a vector column. These names must be specified by ML " + + "org.apache.spark.ml.attribute.Attribute. There can be no overlap with " + + "indices.", + typeConverter=TypeConverters.toListString, + ) + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + indices: Optional[List[int]] = None, + names: Optional[List[str]] = None, + ): + """ + __init__(self, \\*, inputCol=None, outputCol=None, indices=None, names=None) + """ + super(VectorSlicer, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.VectorSlicer", self.uid) + self._setDefault(indices=[], names=[]) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + indices: Optional[List[int]] = None, + names: Optional[List[str]] = None, + ) -> "VectorSlicer": + """ + setParams(self, \\*, inputCol=None, outputCol=None, indices=None, names=None): + Sets params for this VectorSlicer. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.6.0") + def setIndices(self, value: List[int]) -> "VectorSlicer": + """ + Sets the value of :py:attr:`indices`. + """ + return self._set(indices=value) + + @since("1.6.0") + def getIndices(self) -> List[int]: + """ + Gets the value of indices or its default value. + """ + return self.getOrDefault(self.indices) + + @since("1.6.0") + def setNames(self, value: List[str]) -> "VectorSlicer": + """ + Sets the value of :py:attr:`names`. + """ + return self._set(names=value) + + @since("1.6.0") + def getNames(self) -> List[str]: + """ + Gets the value of names or its default value. + """ + return self.getOrDefault(self.names) + + def setInputCol(self, value: str) -> "VectorSlicer": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "VectorSlicer": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + +class _Word2VecParams(HasStepSize, HasMaxIter, HasSeed, HasInputCol, HasOutputCol): + """ + Params for :py:class:`Word2Vec` and :py:class:`Word2VecModel`. + + .. versionadded:: 3.0.0 + """ + + vectorSize: Param[int] = Param( + Params._dummy(), + "vectorSize", + "the dimension of codes after transforming from words", + typeConverter=TypeConverters.toInt, + ) + numPartitions: Param[int] = Param( + Params._dummy(), + "numPartitions", + "number of partitions for sentences of words", + typeConverter=TypeConverters.toInt, + ) + minCount: Param[int] = Param( + Params._dummy(), + "minCount", + "the minimum number of times a token must appear to be included in the " + + "word2vec model's vocabulary", + typeConverter=TypeConverters.toInt, + ) + windowSize: Param[int] = Param( + Params._dummy(), + "windowSize", + "the window size (context words from [-window, window]). Default value is 5", + typeConverter=TypeConverters.toInt, + ) + maxSentenceLength: Param[int] = Param( + Params._dummy(), + "maxSentenceLength", + "Maximum length (in words) of each sentence in the input data. " + + "Any sentence longer than this threshold will " + + "be divided into chunks up to the size.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self, *args: Any): + super(_Word2VecParams, self).__init__(*args) + self._setDefault( + vectorSize=100, + minCount=5, + numPartitions=1, + stepSize=0.025, + maxIter=1, + windowSize=5, + maxSentenceLength=1000, + ) + + @since("1.4.0") + def getVectorSize(self) -> int: + """ + Gets the value of vectorSize or its default value. + """ + return self.getOrDefault(self.vectorSize) + + @since("1.4.0") + def getNumPartitions(self) -> int: + """ + Gets the value of numPartitions or its default value. + """ + return self.getOrDefault(self.numPartitions) + + @since("1.4.0") + def getMinCount(self) -> int: + """ + Gets the value of minCount or its default value. + """ + return self.getOrDefault(self.minCount) + + @since("2.0.0") + def getWindowSize(self) -> int: + """ + Gets the value of windowSize or its default value. + """ + return self.getOrDefault(self.windowSize) + + @since("2.0.0") + def getMaxSentenceLength(self) -> int: + """ + Gets the value of maxSentenceLength or its default value. + """ + return self.getOrDefault(self.maxSentenceLength) + + +@inherit_doc +class Word2Vec( + JavaEstimator["Word2VecModel"], + _Word2VecParams, + JavaMLReadable["Word2Vec"], + JavaMLWritable, +): + """ + Word2Vec trains a model of `Map(String, Vector)`, i.e. transforms a word into a code for further + natural language processing or machine learning process. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> sent = ("a b " * 100 + "a c " * 10).split(" ") + >>> doc = spark.createDataFrame([(sent,), (sent,)], ["sentence"]) + >>> word2Vec = Word2Vec(vectorSize=5, seed=42, inputCol="sentence", outputCol="model") + >>> word2Vec.setMaxIter(10) + Word2Vec... + >>> word2Vec.getMaxIter() + 10 + >>> word2Vec.clear(word2Vec.maxIter) + >>> model = word2Vec.fit(doc) + >>> model.getMinCount() + 5 + >>> model.setInputCol("sentence") + Word2VecModel... + >>> model.getVectors().show() + +----+--------------------+ + |word| vector| + +----+--------------------+ + | a|[0.0951... + | b|[-1.202... + | c|[0.3015... + +----+--------------------+ + ... + >>> model.findSynonymsArray("a", 2) + [('b', 0.015859...), ('c', -0.568079...)] + >>> from pyspark.sql.functions import format_number as fmt + >>> model.findSynonyms("a", 2).select("word", fmt("similarity", 5).alias("similarity")).show() + +----+----------+ + |word|similarity| + +----+----------+ + | b| 0.01586| + | c| -0.56808| + +----+----------+ + ... + >>> model.transform(doc).head().model + DenseVector([-0.4833, 0.1855, -0.273, -0.0509, -0.4769]) + >>> word2vecPath = temp_path + "/word2vec" + >>> word2Vec.save(word2vecPath) + >>> loadedWord2Vec = Word2Vec.load(word2vecPath) + >>> loadedWord2Vec.getVectorSize() == word2Vec.getVectorSize() + True + >>> loadedWord2Vec.getNumPartitions() == word2Vec.getNumPartitions() + True + >>> loadedWord2Vec.getMinCount() == word2Vec.getMinCount() + True + >>> modelPath = temp_path + "/word2vec-model" + >>> model.save(modelPath) + >>> loadedModel = Word2VecModel.load(modelPath) + >>> loadedModel.getVectors().first().word == model.getVectors().first().word + True + >>> loadedModel.getVectors().first().vector == model.getVectors().first().vector + True + >>> loadedModel.transform(doc).take(1) == model.transform(doc).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + vectorSize: int = 100, + minCount: int = 5, + numPartitions: int = 1, + stepSize: float = 0.025, + maxIter: int = 1, + seed: Optional[int] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + windowSize: int = 5, + maxSentenceLength: int = 1000, + ): + """ + __init__(self, \\*, vectorSize=100, minCount=5, numPartitions=1, stepSize=0.025, \ + maxIter=1, seed=None, inputCol=None, outputCol=None, windowSize=5, \ + maxSentenceLength=1000) + """ + super(Word2Vec, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.Word2Vec", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + vectorSize: int = 100, + minCount: int = 5, + numPartitions: int = 1, + stepSize: float = 0.025, + maxIter: int = 1, + seed: Optional[int] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + windowSize: int = 5, + maxSentenceLength: int = 1000, + ) -> "Word2Vec": + """ + setParams(self, \\*, minCount=5, numPartitions=1, stepSize=0.025, maxIter=1, \ + seed=None, inputCol=None, outputCol=None, windowSize=5, \ + maxSentenceLength=1000) + Sets params for this Word2Vec. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.4.0") + def setVectorSize(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`vectorSize`. + """ + return self._set(vectorSize=value) + + @since("1.4.0") + def setNumPartitions(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`numPartitions`. + """ + return self._set(numPartitions=value) + + @since("1.4.0") + def setMinCount(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`minCount`. + """ + return self._set(minCount=value) + + @since("2.0.0") + def setWindowSize(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`windowSize`. + """ + return self._set(windowSize=value) + + @since("2.0.0") + def setMaxSentenceLength(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`maxSentenceLength`. + """ + return self._set(maxSentenceLength=value) + + def setMaxIter(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + def setInputCol(self, value: str) -> "Word2Vec": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "Word2Vec": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setSeed(self, value: int) -> "Word2Vec": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("1.4.0") + def setStepSize(self, value: float) -> "Word2Vec": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + def _create_model(self, java_model: "JavaObject") -> "Word2VecModel": + return Word2VecModel(java_model) + + +class Word2VecModel(JavaModel, _Word2VecParams, JavaMLReadable["Word2VecModel"], JavaMLWritable): + """ + Model fitted by :py:class:`Word2Vec`. + + .. versionadded:: 1.4.0 + """ + + @since("1.5.0") + def getVectors(self) -> DataFrame: + """ + Returns the vector representation of the words as a dataframe + with two fields, word and vector. + """ + return self._call_java("getVectors") + + def setInputCol(self, value: str) -> "Word2VecModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "Word2VecModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @since("1.5.0") + def findSynonyms(self, word: Union[str, Vector], num: int) -> DataFrame: + """ + Find "num" number of words closest in similarity to "word". + word can be a string or vector representation. + Returns a dataframe with two fields word and similarity (which + gives the cosine similarity). + """ + if not isinstance(word, str): + word = _convert_to_vector(word) + return self._call_java("findSynonyms", word, num) + + @since("2.3.0") + def findSynonymsArray(self, word: Union[Vector, str], num: int) -> List[Tuple[str, float]]: + """ + Find "num" number of words closest in similarity to "word". + word can be a string or vector representation. + Returns an array with two fields word and similarity (which + gives the cosine similarity). + """ + if not isinstance(word, str): + word = _convert_to_vector(word) + assert self._java_obj is not None + tuples = self._java_obj.findSynonymsArray(word, num) + return list(map(lambda st: (st._1(), st._2()), list(tuples))) + + +class _PCAParams(HasInputCol, HasOutputCol): + """ + Params for :py:class:`PCA` and :py:class:`PCAModel`. + + .. versionadded:: 3.0.0 + """ + + k: Param[int] = Param( + Params._dummy(), + "k", + "the number of principal components", + typeConverter=TypeConverters.toInt, + ) + + @since("1.5.0") + def getK(self) -> int: + """ + Gets the value of k or its default value. + """ + return self.getOrDefault(self.k) + + +@inherit_doc +class PCA(JavaEstimator["PCAModel"], _PCAParams, JavaMLReadable["PCA"], JavaMLWritable): + """ + PCA trains a model to project vectors to a lower dimensional space of the + top :py:attr:`k` principal components. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> data = [(Vectors.sparse(5, [(1, 1.0), (3, 7.0)]),), + ... (Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),), + ... (Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0]),)] + >>> df = spark.createDataFrame(data,["features"]) + >>> pca = PCA(k=2, inputCol="features") + >>> pca.setOutputCol("pca_features") + PCA... + >>> model = pca.fit(df) + >>> model.getK() + 2 + >>> model.setOutputCol("output") + PCAModel... + >>> model.transform(df).collect()[0].output + DenseVector([1.648..., -4.013...]) + >>> model.explainedVariance + DenseVector([0.794..., 0.205...]) + >>> pcaPath = temp_path + "/pca" + >>> pca.save(pcaPath) + >>> loadedPca = PCA.load(pcaPath) + >>> loadedPca.getK() == pca.getK() + True + >>> modelPath = temp_path + "/pca-model" + >>> model.save(modelPath) + >>> loadedModel = PCAModel.load(modelPath) + >>> loadedModel.pc == model.pc + True + >>> loadedModel.explainedVariance == model.explainedVariance + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + k: Optional[int] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ): + """ + __init__(self, \\*, k=None, inputCol=None, outputCol=None) + """ + super(PCA, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.PCA", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + k: Optional[int] = None, + inputCol: Optional[str] = None, + outputCol: Optional[str] = None, + ) -> "PCA": + """ + setParams(self, \\*, k=None, inputCol=None, outputCol=None) + Set params for this PCA. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.5.0") + def setK(self, value: int) -> "PCA": + """ + Sets the value of :py:attr:`k`. + """ + return self._set(k=value) + + def setInputCol(self, value: str) -> "PCA": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setOutputCol(self, value: str) -> "PCA": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "PCAModel": + return PCAModel(java_model) + + +class PCAModel(JavaModel, _PCAParams, JavaMLReadable["PCAModel"], JavaMLWritable): + """ + Model fitted by :py:class:`PCA`. Transforms vectors to a lower dimensional space. + + .. versionadded:: 1.5.0 + """ + + @since("3.0.0") + def setInputCol(self, value: str) -> "PCAModel": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + @since("3.0.0") + def setOutputCol(self, value: str) -> "PCAModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.0.0") + def pc(self) -> DenseMatrix: + """ + Returns a principal components Matrix. + Each column is one principal component. + """ + return self._call_java("pc") + + @property + @since("2.0.0") + def explainedVariance(self) -> DenseVector: + """ + Returns a vector of proportions of variance + explained by each principal component. + """ + return self._call_java("explainedVariance") + + +class _RFormulaParams(HasFeaturesCol, HasLabelCol, HasHandleInvalid): + """ + Params for :py:class:`RFormula` and :py:class:`RFormula`. + + .. versionadded:: 3.0.0 + """ + + formula: Param[str] = Param( + Params._dummy(), "formula", "R model formula", typeConverter=TypeConverters.toString + ) + + forceIndexLabel: Param[bool] = Param( + Params._dummy(), + "forceIndexLabel", + "Force to index label whether it is numeric or string", + typeConverter=TypeConverters.toBoolean, + ) + + stringIndexerOrderType: Param[str] = Param( + Params._dummy(), + "stringIndexerOrderType", + "How to order categories of a string feature column used by " + + "StringIndexer. The last category after ordering is dropped " + + "when encoding strings. Supported options: frequencyDesc, " + + "frequencyAsc, alphabetDesc, alphabetAsc. The default value " + + "is frequencyDesc. When the ordering is set to alphabetDesc, " + + "RFormula drops the same category as R when encoding strings.", + typeConverter=TypeConverters.toString, + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "how to handle invalid entries. " + + "Options are 'skip' (filter out rows with invalid values), " + + "'error' (throw an error), or 'keep' (put invalid data in a special " + + "additional bucket, at index numLabels).", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_RFormulaParams, self).__init__(*args) + self._setDefault( + forceIndexLabel=False, stringIndexerOrderType="frequencyDesc", handleInvalid="error" + ) + + @since("1.5.0") + def getFormula(self) -> str: + """ + Gets the value of :py:attr:`formula`. + """ + return self.getOrDefault(self.formula) + + @since("2.1.0") + def getForceIndexLabel(self) -> bool: + """ + Gets the value of :py:attr:`forceIndexLabel`. + """ + return self.getOrDefault(self.forceIndexLabel) + + @since("2.3.0") + def getStringIndexerOrderType(self) -> str: + """ + Gets the value of :py:attr:`stringIndexerOrderType` or its default value 'frequencyDesc'. + """ + return self.getOrDefault(self.stringIndexerOrderType) + + +@inherit_doc +class RFormula( + JavaEstimator["RFormulaModel"], + _RFormulaParams, + JavaMLReadable["RFormula"], + JavaMLWritable, +): + """ + Implements the transforms required for fitting a dataset against an + R model formula. Currently we support a limited subset of the R + operators, including '~', '.', ':', '+', '-', '*', and '^'. + + .. versionadded:: 1.5.0 + + Notes + ----- + Also see the `R formula docs + `_. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (1.0, 1.0, "a"), + ... (0.0, 2.0, "b"), + ... (0.0, 0.0, "a") + ... ], ["y", "x", "s"]) + >>> rf = RFormula(formula="y ~ x + s") + >>> model = rf.fit(df) + >>> model.getLabelCol() + 'label' + >>> model.transform(df).show() + +---+---+---+---------+-----+ + | y| x| s| features|label| + +---+---+---+---------+-----+ + |1.0|1.0| a|[1.0,1.0]| 1.0| + |0.0|2.0| b|[2.0,0.0]| 0.0| + |0.0|0.0| a|[0.0,1.0]| 0.0| + +---+---+---+---------+-----+ + ... + >>> rf.fit(df, {rf.formula: "y ~ . - s"}).transform(df).show() + +---+---+---+--------+-----+ + | y| x| s|features|label| + +---+---+---+--------+-----+ + |1.0|1.0| a| [1.0]| 1.0| + |0.0|2.0| b| [2.0]| 0.0| + |0.0|0.0| a| [0.0]| 0.0| + +---+---+---+--------+-----+ + ... + >>> rFormulaPath = temp_path + "/rFormula" + >>> rf.save(rFormulaPath) + >>> loadedRF = RFormula.load(rFormulaPath) + >>> loadedRF.getFormula() == rf.getFormula() + True + >>> loadedRF.getFeaturesCol() == rf.getFeaturesCol() + True + >>> loadedRF.getLabelCol() == rf.getLabelCol() + True + >>> loadedRF.getHandleInvalid() == rf.getHandleInvalid() + True + >>> str(loadedRF) + 'RFormula(y ~ x + s) (uid=...)' + >>> modelPath = temp_path + "/rFormulaModel" + >>> model.save(modelPath) + >>> loadedModel = RFormulaModel.load(modelPath) + >>> loadedModel.uid == model.uid + True + >>> loadedModel.transform(df).show() + +---+---+---+---------+-----+ + | y| x| s| features|label| + +---+---+---+---------+-----+ + |1.0|1.0| a|[1.0,1.0]| 1.0| + |0.0|2.0| b|[2.0,0.0]| 0.0| + |0.0|0.0| a|[0.0,1.0]| 0.0| + +---+---+---+---------+-----+ + ... + >>> str(loadedModel) + 'RFormulaModel(ResolvedRFormula(label=y, terms=[x,s], hasIntercept=true)) (uid=...)' + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + formula: Optional[str] = None, + featuresCol: str = "features", + labelCol: str = "label", + forceIndexLabel: bool = False, + stringIndexerOrderType: str = "frequencyDesc", + handleInvalid: str = "error", + ): + """ + __init__(self, \\*, formula=None, featuresCol="features", labelCol="label", \ + forceIndexLabel=False, stringIndexerOrderType="frequencyDesc", \ + handleInvalid="error") + """ + super(RFormula, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.RFormula", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.5.0") + def setParams( + self, + *, + formula: Optional[str] = None, + featuresCol: str = "features", + labelCol: str = "label", + forceIndexLabel: bool = False, + stringIndexerOrderType: str = "frequencyDesc", + handleInvalid: str = "error", + ) -> "RFormula": + """ + setParams(self, \\*, formula=None, featuresCol="features", labelCol="label", \ + forceIndexLabel=False, stringIndexerOrderType="frequencyDesc", \ + handleInvalid="error") + Sets params for RFormula. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("1.5.0") + def setFormula(self, value: str) -> "RFormula": + """ + Sets the value of :py:attr:`formula`. + """ + return self._set(formula=value) + + @since("2.1.0") + def setForceIndexLabel(self, value: bool) -> "RFormula": + """ + Sets the value of :py:attr:`forceIndexLabel`. + """ + return self._set(forceIndexLabel=value) + + @since("2.3.0") + def setStringIndexerOrderType(self, value: str) -> "RFormula": + """ + Sets the value of :py:attr:`stringIndexerOrderType`. + """ + return self._set(stringIndexerOrderType=value) + + def setFeaturesCol(self, value: str) -> "RFormula": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setLabelCol(self, value: str) -> "RFormula": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def setHandleInvalid(self, value: str) -> "RFormula": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + def _create_model(self, java_model: "JavaObject") -> "RFormulaModel": + return RFormulaModel(java_model) + + def __str__(self) -> str: + formulaStr = self.getFormula() if self.isDefined(self.formula) else "" + return "RFormula(%s) (uid=%s)" % (formulaStr, self.uid) + + +class RFormulaModel(JavaModel, _RFormulaParams, JavaMLReadable["RFormulaModel"], JavaMLWritable): + """ + Model fitted by :py:class:`RFormula`. Fitting is required to determine the + factor levels of formula terms. + + .. versionadded:: 1.5.0 + """ + + def __str__(self) -> str: + resolvedFormula = self._call_java("resolvedFormula") + return "RFormulaModel(%s) (uid=%s)" % (resolvedFormula, self.uid) + + +class _SelectorParams(HasFeaturesCol, HasOutputCol, HasLabelCol): + """ + Params for :py:class:`Selector` and :py:class:`SelectorModel`. + + .. versionadded:: 3.1.0 + """ + + selectorType: Param[str] = Param( + Params._dummy(), + "selectorType", + "The selector type. " + + "Supported options: numTopFeatures (default), percentile, fpr, fdr, fwe.", + typeConverter=TypeConverters.toString, + ) + + numTopFeatures: Param[int] = Param( + Params._dummy(), + "numTopFeatures", + "Number of features that selector will select, ordered by ascending p-value. " + + "If the number of features is < numTopFeatures, then this will select " + + "all features.", + typeConverter=TypeConverters.toInt, + ) + + percentile: Param[float] = Param( + Params._dummy(), + "percentile", + "Percentile of features that selector " + "will select, ordered by ascending p-value.", + typeConverter=TypeConverters.toFloat, + ) + + fpr: Param[float] = Param( + Params._dummy(), + "fpr", + "The highest p-value for features to be kept.", + typeConverter=TypeConverters.toFloat, + ) + + fdr: Param[float] = Param( + Params._dummy(), + "fdr", + "The upper bound of the expected false discovery rate.", + typeConverter=TypeConverters.toFloat, + ) + + fwe: Param[float] = Param( + Params._dummy(), + "fwe", + "The upper bound of the expected family-wise error rate.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_SelectorParams, self).__init__(*args) + self._setDefault( + numTopFeatures=50, + selectorType="numTopFeatures", + percentile=0.1, + fpr=0.05, + fdr=0.05, + fwe=0.05, + ) + + @since("2.1.0") + def getSelectorType(self) -> str: + """ + Gets the value of selectorType or its default value. + """ + return self.getOrDefault(self.selectorType) + + @since("2.0.0") + def getNumTopFeatures(self) -> int: + """ + Gets the value of numTopFeatures or its default value. + """ + return self.getOrDefault(self.numTopFeatures) + + @since("2.1.0") + def getPercentile(self) -> float: + """ + Gets the value of percentile or its default value. + """ + return self.getOrDefault(self.percentile) + + @since("2.1.0") + def getFpr(self) -> float: + """ + Gets the value of fpr or its default value. + """ + return self.getOrDefault(self.fpr) + + @since("2.2.0") + def getFdr(self) -> float: + """ + Gets the value of fdr or its default value. + """ + return self.getOrDefault(self.fdr) + + @since("2.2.0") + def getFwe(self) -> float: + """ + Gets the value of fwe or its default value. + """ + return self.getOrDefault(self.fwe) + + +class _Selector(JavaEstimator[JM], _SelectorParams, JavaMLReadable, JavaMLWritable, Generic[JM]): + """ + Mixin for Selectors. + """ + + @since("2.1.0") + def setSelectorType(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`selectorType`. + """ + return self._set(selectorType=value) + + @since("2.0.0") + def setNumTopFeatures(self: P, value: int) -> P: + """ + Sets the value of :py:attr:`numTopFeatures`. + Only applicable when selectorType = "numTopFeatures". + """ + return self._set(numTopFeatures=value) + + @since("2.1.0") + def setPercentile(self: P, value: float) -> P: + """ + Sets the value of :py:attr:`percentile`. + Only applicable when selectorType = "percentile". + """ + return self._set(percentile=value) + + @since("2.1.0") + def setFpr(self: P, value: float) -> P: + """ + Sets the value of :py:attr:`fpr`. + Only applicable when selectorType = "fpr". + """ + return self._set(fpr=value) + + @since("2.2.0") + def setFdr(self: P, value: float) -> P: + """ + Sets the value of :py:attr:`fdr`. + Only applicable when selectorType = "fdr". + """ + return self._set(fdr=value) + + @since("2.2.0") + def setFwe(self: P, value: float) -> P: + """ + Sets the value of :py:attr:`fwe`. + Only applicable when selectorType = "fwe". + """ + return self._set(fwe=value) + + def setFeaturesCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setOutputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setLabelCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + +class _SelectorModel(JavaModel, _SelectorParams): + """ + Mixin for Selector models. + """ + + @since("3.0.0") + def setFeaturesCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setOutputCol(self: P, value: str) -> P: + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("2.0.0") + def selectedFeatures(self) -> List[int]: + """ + List of indices to select (filter). + """ + return self._call_java("selectedFeatures") + + +@inherit_doc +class ChiSqSelector( + _Selector["ChiSqSelectorModel"], + JavaMLReadable["ChiSqSelector"], + JavaMLWritable, +): + """ + Chi-Squared feature selection, which selects categorical features to use for predicting a + categorical label. + The selector supports different selection methods: `numTopFeatures`, `percentile`, `fpr`, + `fdr`, `fwe`. + + * `numTopFeatures` chooses a fixed number of top features according to a chi-squared test. + + * `percentile` is similar but chooses a fraction of all features + instead of a fixed number. + + * `fpr` chooses all features whose p-values are below a threshold, + thus controlling the false positive rate of selection. + + * `fdr` uses the `Benjamini-Hochberg procedure `_ + to choose all features whose false discovery rate is below a threshold. + + * `fwe` chooses all features whose p-values are below a threshold. The threshold is scaled by + 1/numFeatures, thus controlling the family-wise error rate of selection. + + By default, the selection method is `numTopFeatures`, with the default number of top features + set to 50. + + .. deprecated:: 3.1.0 + Use UnivariateFeatureSelector + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame( + ... [(Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0), + ... (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0), + ... (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0)], + ... ["features", "label"]) + >>> selector = ChiSqSelector(numTopFeatures=1, outputCol="selectedFeatures") + >>> model = selector.fit(df) + >>> model.getFeaturesCol() + 'features' + >>> model.setFeaturesCol("features") + ChiSqSelectorModel... + >>> model.transform(df).head().selectedFeatures + DenseVector([18.0]) + >>> model.selectedFeatures + [2] + >>> chiSqSelectorPath = temp_path + "/chi-sq-selector" + >>> selector.save(chiSqSelectorPath) + >>> loadedSelector = ChiSqSelector.load(chiSqSelectorPath) + >>> loadedSelector.getNumTopFeatures() == selector.getNumTopFeatures() + True + >>> modelPath = temp_path + "/chi-sq-selector-model" + >>> model.save(modelPath) + >>> loadedModel = ChiSqSelectorModel.load(modelPath) + >>> loadedModel.selectedFeatures == model.selectedFeatures + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + numTopFeatures: int = 50, + featuresCol: str = "features", + outputCol: Optional[str] = None, + labelCol: str = "label", + selectorType: str = "numTopFeatures", + percentile: float = 0.1, + fpr: float = 0.05, + fdr: float = 0.05, + fwe: float = 0.05, + ): + """ + __init__(self, \\*, numTopFeatures=50, featuresCol="features", outputCol=None, \ + labelCol="label", selectorType="numTopFeatures", percentile=0.1, fpr=0.05, \ + fdr=0.05, fwe=0.05) + """ + super(ChiSqSelector, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.ChiSqSelector", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + numTopFeatures: int = 50, + featuresCol: str = "features", + outputCol: Optional[str] = None, + labelCol: str = "label", + selectorType: str = "numTopFeatures", + percentile: float = 0.1, + fpr: float = 0.05, + fdr: float = 0.05, + fwe: float = 0.05, + ) -> "ChiSqSelector": + """ + setParams(self, \\*, numTopFeatures=50, featuresCol="features", outputCol=None, \ + labelCol="label", selectorType="numTopFeatures", percentile=0.1, fpr=0.05, \ + fdr=0.05, fwe=0.05) + Sets params for this ChiSqSelector. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "ChiSqSelectorModel": + return ChiSqSelectorModel(java_model) + + +class ChiSqSelectorModel(_SelectorModel, JavaMLReadable["ChiSqSelectorModel"], JavaMLWritable): + """ + Model fitted by :py:class:`ChiSqSelector`. + + .. versionadded:: 2.0.0 + """ + + +@inherit_doc +class VectorSizeHint( + JavaTransformer, + HasInputCol, + HasHandleInvalid, + JavaMLReadable["VectorSizeHint"], + JavaMLWritable, +): + """ + A feature transformer that adds size information to the metadata of a vector column. + VectorAssembler needs size information for its input columns and cannot be used on streaming + dataframes without this metadata. + + .. versionadded:: 2.3.0 + + Notes + ----- + VectorSizeHint modifies `inputCol` to include size metadata and does not have an outputCol. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml import Pipeline, PipelineModel + >>> data = [(Vectors.dense([1., 2., 3.]), 4.)] + >>> df = spark.createDataFrame(data, ["vector", "float"]) + >>> + >>> sizeHint = VectorSizeHint(inputCol="vector", size=3, handleInvalid="skip") + >>> vecAssembler = VectorAssembler(inputCols=["vector", "float"], outputCol="assembled") + >>> pipeline = Pipeline(stages=[sizeHint, vecAssembler]) + >>> + >>> pipelineModel = pipeline.fit(df) + >>> pipelineModel.transform(df).head().assembled + DenseVector([1.0, 2.0, 3.0, 4.0]) + >>> vectorSizeHintPath = temp_path + "/vector-size-hint-pipeline" + >>> pipelineModel.save(vectorSizeHintPath) + >>> loadedPipeline = PipelineModel.load(vectorSizeHintPath) + >>> loaded = loadedPipeline.transform(df).head().assembled + >>> expected = pipelineModel.transform(df).head().assembled + >>> loaded == expected + True + """ + + _input_kwargs: Dict[str, Any] + + size: Param[int] = Param( + Params._dummy(), "size", "Size of vectors in column.", typeConverter=TypeConverters.toInt + ) + + handleInvalid: Param[str] = Param( + Params._dummy(), + "handleInvalid", + "How to handle invalid vectors in inputCol. Invalid vectors include " + "nulls and vectors with the wrong size. The options are `skip` (filter " + "out rows with invalid vectors), `error` (throw an error) and " + "`optimistic` (do not check the vector size, and keep all rows). " + "`error` by default.", + TypeConverters.toString, + ) + + @keyword_only + def __init__( + self, + *, + inputCol: Optional[str] = None, + size: Optional[int] = None, + handleInvalid: str = "error", + ): + """ + __init__(self, \\*, inputCol=None, size=None, handleInvalid="error") + """ + super(VectorSizeHint, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.VectorSizeHint", self.uid) + self._setDefault(handleInvalid="error") + self.setParams(**self._input_kwargs) + + @keyword_only + @since("2.3.0") + def setParams( + self, + *, + inputCol: Optional[str] = None, + size: Optional[str] = None, + handleInvalid: str = "error", + ) -> "VectorSizeHint": + """ + setParams(self, \\*, inputCol=None, size=None, handleInvalid="error") + Sets params for this VectorSizeHint. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.3.0") + def getSize(self) -> int: + """Gets size param, the size of vectors in `inputCol`.""" + return self.getOrDefault(self.size) + + @since("2.3.0") + def setSize(self, value: int) -> "VectorSizeHint": + """Sets size param, the size of vectors in `inputCol`.""" + return self._set(size=value) + + def setInputCol(self, value: str) -> "VectorSizeHint": + """ + Sets the value of :py:attr:`inputCol`. + """ + return self._set(inputCol=value) + + def setHandleInvalid(self, value: str) -> "VectorSizeHint": + """ + Sets the value of :py:attr:`handleInvalid`. + """ + return self._set(handleInvalid=value) + + +class _VarianceThresholdSelectorParams(HasFeaturesCol, HasOutputCol): + """ + Params for :py:class:`VarianceThresholdSelector` and + :py:class:`VarianceThresholdSelectorModel`. + + .. versionadded:: 3.1.0 + """ + + varianceThreshold: Param[float] = Param( + Params._dummy(), + "varianceThreshold", + "Param for variance threshold. Features with a variance not " + + "greater than this threshold will be removed. The default value " + + "is 0.0.", + typeConverter=TypeConverters.toFloat, + ) + + @since("3.1.0") + def getVarianceThreshold(self) -> float: + """ + Gets the value of varianceThreshold or its default value. + """ + return self.getOrDefault(self.varianceThreshold) + + +@inherit_doc +class VarianceThresholdSelector( + JavaEstimator["VarianceThresholdSelectorModel"], + _VarianceThresholdSelectorParams, + JavaMLReadable["VarianceThresholdSelector"], + JavaMLWritable, +): + """ + Feature selector that removes all low-variance features. Features with a + (sample) variance not greater than the threshold will be removed. The default is to keep + all features with non-zero variance, i.e. remove the features that have the + same value in all samples. + + .. versionadded:: 3.1.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame( + ... [(Vectors.dense([6.0, 7.0, 0.0, 7.0, 6.0, 0.0]),), + ... (Vectors.dense([0.0, 9.0, 6.0, 0.0, 5.0, 9.0]),), + ... (Vectors.dense([0.0, 9.0, 3.0, 0.0, 5.0, 5.0]),), + ... (Vectors.dense([0.0, 9.0, 8.0, 5.0, 6.0, 4.0]),), + ... (Vectors.dense([8.0, 9.0, 6.0, 5.0, 4.0, 4.0]),), + ... (Vectors.dense([8.0, 9.0, 6.0, 0.0, 0.0, 0.0]),)], + ... ["features"]) + >>> selector = VarianceThresholdSelector(varianceThreshold=8.2, outputCol="selectedFeatures") + >>> model = selector.fit(df) + >>> model.getFeaturesCol() + 'features' + >>> model.setFeaturesCol("features") + VarianceThresholdSelectorModel... + >>> model.transform(df).head().selectedFeatures + DenseVector([6.0, 7.0, 0.0]) + >>> model.selectedFeatures + [0, 3, 5] + >>> varianceThresholdSelectorPath = temp_path + "/variance-threshold-selector" + >>> selector.save(varianceThresholdSelectorPath) + >>> loadedSelector = VarianceThresholdSelector.load(varianceThresholdSelectorPath) + >>> loadedSelector.getVarianceThreshold() == selector.getVarianceThreshold() + True + >>> modelPath = temp_path + "/variance-threshold-selector-model" + >>> model.save(modelPath) + >>> loadedModel = VarianceThresholdSelectorModel.load(modelPath) + >>> loadedModel.selectedFeatures == model.selectedFeatures + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + outputCol: Optional[str] = None, + varianceThreshold: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", outputCol=None, varianceThreshold=0.0) + """ + super(VarianceThresholdSelector, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.VarianceThresholdSelector", self.uid + ) + self._setDefault(varianceThreshold=0.0) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.1.0") + def setParams( + self, + *, + featuresCol: str = "features", + outputCol: Optional[str] = None, + varianceThreshold: float = 0.0, + ) -> "VarianceThresholdSelector": + """ + setParams(self, \\*, featuresCol="features", outputCol=None, varianceThreshold=0.0) + Sets params for this VarianceThresholdSelector. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.1.0") + def setVarianceThreshold(self, value: float) -> "VarianceThresholdSelector": + """ + Sets the value of :py:attr:`varianceThreshold`. + """ + return self._set(varianceThreshold=value) + + @since("3.1.0") + def setFeaturesCol(self, value: str) -> "VarianceThresholdSelector": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.1.0") + def setOutputCol(self, value: str) -> "VarianceThresholdSelector": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def _create_model(self, java_model: "JavaObject") -> "VarianceThresholdSelectorModel": + return VarianceThresholdSelectorModel(java_model) + + +class VarianceThresholdSelectorModel( + JavaModel, + _VarianceThresholdSelectorParams, + JavaMLReadable["VarianceThresholdSelectorModel"], + JavaMLWritable, +): + """ + Model fitted by :py:class:`VarianceThresholdSelector`. + + .. versionadded:: 3.1.0 + """ + + @since("3.1.0") + def setFeaturesCol(self, value: str) -> "VarianceThresholdSelectorModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.1.0") + def setOutputCol(self, value: str) -> "VarianceThresholdSelectorModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("3.1.0") + def selectedFeatures(self) -> List[int]: + """ + List of indices to select (filter). + """ + return self._call_java("selectedFeatures") + + +class _UnivariateFeatureSelectorParams(HasFeaturesCol, HasOutputCol, HasLabelCol): + """ + Params for :py:class:`UnivariateFeatureSelector` and + :py:class:`UnivariateFeatureSelectorModel`. + + .. versionadded:: 3.1.0 + """ + + featureType: Param[str] = Param( + Params._dummy(), + "featureType", + "The feature type. " + "Supported options: categorical, continuous.", + typeConverter=TypeConverters.toString, + ) + + labelType: Param[str] = Param( + Params._dummy(), + "labelType", + "The label type. " + "Supported options: categorical, continuous.", + typeConverter=TypeConverters.toString, + ) + + selectionMode: Param[str] = Param( + Params._dummy(), + "selectionMode", + "The selection mode. " + + "Supported options: numTopFeatures (default), percentile, fpr, " + + "fdr, fwe.", + typeConverter=TypeConverters.toString, + ) + + selectionThreshold: Param[float] = Param( + Params._dummy(), + "selectionThreshold", + "The upper bound of the " + "features that selector will select.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_UnivariateFeatureSelectorParams, self).__init__(*args) + self._setDefault(selectionMode="numTopFeatures") + + @since("3.1.1") + def getFeatureType(self) -> str: + """ + Gets the value of featureType or its default value. + """ + return self.getOrDefault(self.featureType) + + @since("3.1.1") + def getLabelType(self) -> str: + """ + Gets the value of labelType or its default value. + """ + return self.getOrDefault(self.labelType) + + @since("3.1.1") + def getSelectionMode(self) -> str: + """ + Gets the value of selectionMode or its default value. + """ + return self.getOrDefault(self.selectionMode) + + @since("3.1.1") + def getSelectionThreshold(self) -> float: + """ + Gets the value of selectionThreshold or its default value. + """ + return self.getOrDefault(self.selectionThreshold) + + +@inherit_doc +class UnivariateFeatureSelector( + JavaEstimator["UnivariateFeatureSelectorModel"], + _UnivariateFeatureSelectorParams, + JavaMLReadable["UnivariateFeatureSelector"], + JavaMLWritable, +): + """ + UnivariateFeatureSelector + Feature selector based on univariate statistical tests against labels. Currently, Spark + supports three Univariate Feature Selectors: chi-squared, ANOVA F-test and F-value. + User can choose Univariate Feature Selector by setting `featureType` and `labelType`, + and Spark will pick the score function based on the specified `featureType` and `labelType`. + + The following combination of `featureType` and `labelType` are supported: + + - `featureType` `categorical` and `labelType` `categorical`, Spark uses chi-squared, + i.e. chi2 in sklearn. + - `featureType` `continuous` and `labelType` `categorical`, Spark uses ANOVA F-test, + i.e. f_classif in sklearn. + - `featureType` `continuous` and `labelType` `continuous`, Spark uses F-value, + i.e. f_regression in sklearn. + + The `UnivariateFeatureSelector` supports different selection modes: `numTopFeatures`, + `percentile`, `fpr`, `fdr`, `fwe`. + + - `numTopFeatures` chooses a fixed number of top features according to a according to a + hypothesis. + - `percentile` is similar but chooses a fraction of all features + instead of a fixed number. + - `fpr` chooses all features whose p-values are below a threshold, + thus controlling the false positive rate of selection. + - `fdr` uses the `Benjamini-Hochberg procedure \ + `_ + to choose all features whose false discovery rate is below a threshold. + - `fwe` chooses all features whose p-values are below a threshold. The threshold is scaled by + 1 / `numFeatures`, thus controlling the family-wise error rate of selection. + + By default, the selection mode is `numTopFeatures`. + + .. versionadded:: 3.1.1 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame( + ... [(Vectors.dense([1.7, 4.4, 7.6, 5.8, 9.6, 2.3]), 3.0), + ... (Vectors.dense([8.8, 7.3, 5.7, 7.3, 2.2, 4.1]), 2.0), + ... (Vectors.dense([1.2, 9.5, 2.5, 3.1, 8.7, 2.5]), 1.0), + ... (Vectors.dense([3.7, 9.2, 6.1, 4.1, 7.5, 3.8]), 2.0), + ... (Vectors.dense([8.9, 5.2, 7.8, 8.3, 5.2, 3.0]), 4.0), + ... (Vectors.dense([7.9, 8.5, 9.2, 4.0, 9.4, 2.1]), 4.0)], + ... ["features", "label"]) + >>> selector = UnivariateFeatureSelector(outputCol="selectedFeatures") + >>> selector.setFeatureType("continuous").setLabelType("categorical").setSelectionThreshold(1) + UnivariateFeatureSelector... + >>> model = selector.fit(df) + >>> model.getFeaturesCol() + 'features' + >>> model.setFeaturesCol("features") + UnivariateFeatureSelectorModel... + >>> model.transform(df).head().selectedFeatures + DenseVector([7.6]) + >>> model.selectedFeatures + [2] + >>> selectorPath = temp_path + "/selector" + >>> selector.save(selectorPath) + >>> loadedSelector = UnivariateFeatureSelector.load(selectorPath) + >>> loadedSelector.getSelectionThreshold() == selector.getSelectionThreshold() + True + >>> modelPath = temp_path + "/selector-model" + >>> model.save(modelPath) + >>> loadedModel = UnivariateFeatureSelectorModel.load(modelPath) + >>> loadedModel.selectedFeatures == model.selectedFeatures + True + >>> loadedModel.transform(df).take(1) == model.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + outputCol: Optional[str] = None, + labelCol: str = "label", + selectionMode: str = "numTopFeatures", + ): + """ + __init__(self, \\*, featuresCol="features", outputCol=None, \ + labelCol="label", selectionMode="numTopFeatures") + """ + super(UnivariateFeatureSelector, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.feature.UnivariateFeatureSelector", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.1.1") + def setParams( + self, + *, + featuresCol: str = "features", + outputCol: Optional[str] = None, + labelCol: str = "label", + selectionMode: str = "numTopFeatures", + ) -> "UnivariateFeatureSelector": + """ + setParams(self, \\*, featuresCol="features", outputCol=None, \ + labelCol="label", selectionMode="numTopFeatures") + Sets params for this UnivariateFeatureSelector. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.1.1") + def setFeatureType(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`featureType`. + """ + return self._set(featureType=value) + + @since("3.1.1") + def setLabelType(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`labelType`. + """ + return self._set(labelType=value) + + @since("3.1.1") + def setSelectionMode(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`selectionMode`. + """ + return self._set(selectionMode=value) + + @since("3.1.1") + def setSelectionThreshold(self, value: float) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`selectionThreshold`. + """ + return self._set(selectionThreshold=value) + + def setFeaturesCol(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + def setOutputCol(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + def setLabelCol(self, value: str) -> "UnivariateFeatureSelector": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + def _create_model(self, java_model: "JavaObject") -> "UnivariateFeatureSelectorModel": + return UnivariateFeatureSelectorModel(java_model) + + +class UnivariateFeatureSelectorModel( + JavaModel, + _UnivariateFeatureSelectorParams, + JavaMLReadable["UnivariateFeatureSelectorModel"], + JavaMLWritable, +): + """ + Model fitted by :py:class:`UnivariateFeatureSelector`. + + .. versionadded:: 3.1.1 + """ + + @since("3.1.1") + def setFeaturesCol(self, value: str) -> "UnivariateFeatureSelectorModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.1.1") + def setOutputCol(self, value: str) -> "UnivariateFeatureSelectorModel": + """ + Sets the value of :py:attr:`outputCol`. + """ + return self._set(outputCol=value) + + @property + @since("3.1.1") + def selectedFeatures(self) -> List[int]: + """ + List of indices to select (filter). + """ + return self._call_java("selectedFeatures") + + +if __name__ == "__main__": + import doctest + import sys + import tempfile + + import pyspark.ml.feature + from pyspark.sql import Row, SparkSession + + globs = globals().copy() + features = pyspark.ml.feature.__dict__.copy() + globs.update(features) + + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.feature tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + testData = sc.parallelize( + [ + Row(id=0, label="a"), + Row(id=1, label="b"), + Row(id=2, label="c"), + Row(id=3, label="a"), + Row(id=4, label="a"), + Row(id=5, label="c"), + ], + 2, + ) + globs["stringIndDf"] = spark.createDataFrame(testData) + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/fpm.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/fpm.py new file mode 100644 index 0000000..cba4219 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/fpm.py @@ -0,0 +1,543 @@ +# +# 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 sys +from typing import Any, Dict, Optional, TYPE_CHECKING + +from pyspark import keyword_only, since +from pyspark.sql import DataFrame +from pyspark.ml.util import JavaMLWritable, JavaMLReadable +from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams +from pyspark.ml.param.shared import HasPredictionCol, Param, TypeConverters, Params + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + +__all__ = ["FPGrowth", "FPGrowthModel", "PrefixSpan"] + + +class _FPGrowthParams(HasPredictionCol): + """ + Params for :py:class:`FPGrowth` and :py:class:`FPGrowthModel`. + + .. versionadded:: 3.0.0 + """ + + itemsCol: Param[str] = Param( + Params._dummy(), "itemsCol", "items column name", typeConverter=TypeConverters.toString + ) + minSupport: Param[float] = Param( + Params._dummy(), + "minSupport", + "Minimal support level of the frequent pattern. [0.0, 1.0]. " + + "Any pattern that appears more than (minSupport * size-of-the-dataset) " + + "times will be output in the frequent itemsets.", + typeConverter=TypeConverters.toFloat, + ) + numPartitions: Param[int] = Param( + Params._dummy(), + "numPartitions", + "Number of partitions (at least 1) used by parallel FP-growth. " + + "By default the param is not set, " + + "and partition number of the input dataset is used.", + typeConverter=TypeConverters.toInt, + ) + minConfidence: Param[float] = Param( + Params._dummy(), + "minConfidence", + "Minimal confidence for generating Association Rule. [0.0, 1.0]. " + + "minConfidence will not affect the mining for frequent itemsets, " + + "but will affect the association rules generation.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_FPGrowthParams, self).__init__(*args) + self._setDefault( + minSupport=0.3, minConfidence=0.8, itemsCol="items", predictionCol="prediction" + ) + + def getItemsCol(self) -> str: + """ + Gets the value of itemsCol or its default value. + """ + return self.getOrDefault(self.itemsCol) + + def getMinSupport(self) -> float: + """ + Gets the value of minSupport or its default value. + """ + return self.getOrDefault(self.minSupport) + + def getNumPartitions(self) -> int: + """ + Gets the value of :py:attr:`numPartitions` or its default value. + """ + return self.getOrDefault(self.numPartitions) + + def getMinConfidence(self) -> float: + """ + Gets the value of minConfidence or its default value. + """ + return self.getOrDefault(self.minConfidence) + + +class FPGrowthModel(JavaModel, _FPGrowthParams, JavaMLWritable, JavaMLReadable["FPGrowthModel"]): + """ + Model fitted by FPGrowth. + + .. versionadded:: 2.2.0 + """ + + @since("3.0.0") + def setItemsCol(self, value: str) -> "FPGrowthModel": + """ + Sets the value of :py:attr:`itemsCol`. + """ + return self._set(itemsCol=value) + + @since("3.0.0") + def setMinConfidence(self, value: float) -> "FPGrowthModel": + """ + Sets the value of :py:attr:`minConfidence`. + """ + return self._set(minConfidence=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "FPGrowthModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @property + @since("2.2.0") + def freqItemsets(self) -> DataFrame: + """ + DataFrame with two columns: + * `items` - Itemset of the same type as the input column. + * `freq` - Frequency of the itemset (`LongType`). + """ + return self._call_java("freqItemsets") + + @property + @since("2.2.0") + def associationRules(self) -> DataFrame: + """ + DataFrame with four columns: + * `antecedent` - Array of the same type as the input column. + * `consequent` - Array of the same type as the input column. + * `confidence` - Confidence for the rule (`DoubleType`). + * `lift` - Lift for the rule (`DoubleType`). + """ + return self._call_java("associationRules") + + +class FPGrowth( + JavaEstimator[FPGrowthModel], _FPGrowthParams, JavaMLWritable, JavaMLReadable["FPGrowth"] +): + r""" + A parallel FP-growth algorithm to mine frequent itemsets. + + .. versionadded:: 2.2.0 + + Notes + ----- + + The algorithm is described in + Li et al., PFP: Parallel FP-Growth for Query Recommendation [1]_. + PFP distributes computation in such a way that each worker executes an + independent group of mining tasks. The FP-Growth algorithm is described in + Han et al., Mining frequent patterns without candidate generation [2]_ + + NULL values in the feature column are ignored during `fit()`. + + Internally `transform` `collects` and `broadcasts` association rules. + + + .. [1] Haoyuan Li, Yi Wang, Dong Zhang, Ming Zhang, and Edward Y. Chang. 2008. + Pfp: parallel fp-growth for query recommendation. + In Proceedings of the 2008 ACM conference on Recommender systems (RecSys '08). + Association for Computing Machinery, New York, NY, USA, 107-114. + DOI: https://doi.org/10.1145/1454008.1454027 + .. [2] Jiawei Han, Jian Pei, and Yiwen Yin. 2000. + Mining frequent patterns without candidate generation. + SIGMOD Rec. 29, 2 (June 2000), 1-12. + DOI: https://doi.org/10.1145/335191.335372 + + + Examples + -------- + >>> from pyspark.sql.functions import split + >>> data = (spark.read + ... .text("data/mllib/sample_fpgrowth.txt") + ... .select(split("value", "\s+").alias("items"))) + >>> data.show(truncate=False) + +------------------------+ + |items | + +------------------------+ + |[r, z, h, k, p] | + |[z, y, x, w, v, u, t, s]| + |[s, x, o, n, r] | + |[x, z, y, m, t, s, q, e]| + |[z] | + |[x, z, y, r, q, t, p] | + +------------------------+ + ... + >>> fp = FPGrowth(minSupport=0.2, minConfidence=0.7) + >>> fpm = fp.fit(data) + >>> fpm.setPredictionCol("newPrediction") + FPGrowthModel... + >>> fpm.freqItemsets.sort("items").show(5) + +---------+----+ + | items|freq| + +---------+----+ + | [p]| 2| + | [p, r]| 2| + |[p, r, z]| 2| + | [p, z]| 2| + | [q]| 2| + +---------+----+ + only showing top 5 rows + ... + >>> fpm.associationRules.sort("antecedent", "consequent").show(5) + +----------+----------+----------+----+------------------+ + |antecedent|consequent|confidence|lift| support| + +----------+----------+----------+----+------------------+ + | [p]| [r]| 1.0| 2.0|0.3333333333333333| + | [p]| [z]| 1.0| 1.2|0.3333333333333333| + | [p, r]| [z]| 1.0| 1.2|0.3333333333333333| + | [p, z]| [r]| 1.0| 2.0|0.3333333333333333| + | [q]| [t]| 1.0| 2.0|0.3333333333333333| + +----------+----------+----------+----+------------------+ + only showing top 5 rows + ... + >>> new_data = spark.createDataFrame([(["t", "s"], )], ["items"]) + >>> sorted(fpm.transform(new_data).first().newPrediction) + ['x', 'y', 'z'] + >>> model_path = temp_path + "/fpm_model" + >>> fpm.save(model_path) + >>> model2 = FPGrowthModel.load(model_path) + >>> fpm.transform(data).take(1) == model2.transform(data).take(1) + True + """ + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + minSupport: float = 0.3, + minConfidence: float = 0.8, + itemsCol: str = "items", + predictionCol: str = "prediction", + numPartitions: Optional[int] = None, + ): + """ + __init__(self, \\*, minSupport=0.3, minConfidence=0.8, itemsCol="items", \ + predictionCol="prediction", numPartitions=None) + """ + super(FPGrowth, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.fpm.FPGrowth", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.2.0") + def setParams( + self, + *, + minSupport: float = 0.3, + minConfidence: float = 0.8, + itemsCol: str = "items", + predictionCol: str = "prediction", + numPartitions: Optional[int] = None, + ) -> "FPGrowth": + """ + setParams(self, \\*, minSupport=0.3, minConfidence=0.8, itemsCol="items", \ + predictionCol="prediction", numPartitions=None) + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def setItemsCol(self, value: str) -> "FPGrowth": + """ + Sets the value of :py:attr:`itemsCol`. + """ + return self._set(itemsCol=value) + + def setMinSupport(self, value: float) -> "FPGrowth": + """ + Sets the value of :py:attr:`minSupport`. + """ + return self._set(minSupport=value) + + def setNumPartitions(self, value: int) -> "FPGrowth": + """ + Sets the value of :py:attr:`numPartitions`. + """ + return self._set(numPartitions=value) + + def setMinConfidence(self, value: float) -> "FPGrowth": + """ + Sets the value of :py:attr:`minConfidence`. + """ + return self._set(minConfidence=value) + + def setPredictionCol(self, value: str) -> "FPGrowth": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + def _create_model(self, java_model: "JavaObject") -> FPGrowthModel: + return FPGrowthModel(java_model) + + +class PrefixSpan(JavaParams): + """ + A parallel PrefixSpan algorithm to mine frequent sequential patterns. + The PrefixSpan algorithm is described in J. Pei, et al., PrefixSpan: Mining Sequential Patterns + Efficiently by Prefix-Projected Pattern Growth + (see `here `_). + This class is not yet an Estimator/Transformer, use :py:func:`findFrequentSequentialPatterns` + method to run the PrefixSpan algorithm. + + .. versionadded:: 2.4.0 + + Notes + ----- + See `Sequential Pattern Mining (Wikipedia) \ + `_ + + Examples + -------- + >>> from pyspark.ml.fpm import PrefixSpan + >>> from pyspark.sql import Row + >>> df = sc.parallelize([Row(sequence=[[1, 2], [3]]), + ... Row(sequence=[[1], [3, 2], [1, 2]]), + ... Row(sequence=[[1, 2], [5]]), + ... Row(sequence=[[6]])]).toDF() + >>> prefixSpan = PrefixSpan() + >>> prefixSpan.getMaxLocalProjDBSize() + 32000000 + >>> prefixSpan.getSequenceCol() + 'sequence' + >>> prefixSpan.setMinSupport(0.5) + PrefixSpan... + >>> prefixSpan.setMaxPatternLength(5) + PrefixSpan... + >>> prefixSpan.findFrequentSequentialPatterns(df).sort("sequence").show(truncate=False) + +----------+----+ + |sequence |freq| + +----------+----+ + |[[1]] |3 | + |[[1], [3]]|2 | + |[[2]] |3 | + |[[2, 1]] |3 | + |[[3]] |2 | + +----------+----+ + ... + """ + + _input_kwargs: Dict[str, Any] + + minSupport: Param[float] = Param( + Params._dummy(), + "minSupport", + "The minimal support level of the " + + "sequential pattern. Sequential pattern that appears more than " + + "(minSupport * size-of-the-dataset) times will be output. Must be >= 0.", + typeConverter=TypeConverters.toFloat, + ) + + maxPatternLength: Param[int] = Param( + Params._dummy(), + "maxPatternLength", + "The maximal length of the sequential pattern. Must be > 0.", + typeConverter=TypeConverters.toInt, + ) + + maxLocalProjDBSize: Param[int] = Param( + Params._dummy(), + "maxLocalProjDBSize", + "The maximum number of items (including delimiters used in the " + + "internal storage format) allowed in a projected database before " + + "local processing. If a projected database exceeds this size, " + + "another iteration of distributed prefix growth is run. " + + "Must be > 0.", + typeConverter=TypeConverters.toInt, + ) + + sequenceCol: Param[str] = Param( + Params._dummy(), + "sequenceCol", + "The name of the sequence column in " + + "dataset, rows with nulls in this column are ignored.", + typeConverter=TypeConverters.toString, + ) + + @keyword_only + def __init__( + self, + *, + minSupport: float = 0.1, + maxPatternLength: int = 10, + maxLocalProjDBSize: int = 32000000, + sequenceCol: str = "sequence", + ): + """ + __init__(self, \\*, minSupport=0.1, maxPatternLength=10, maxLocalProjDBSize=32000000, \ + sequenceCol="sequence") + """ + super(PrefixSpan, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.fpm.PrefixSpan", self.uid) + self._setDefault( + minSupport=0.1, maxPatternLength=10, maxLocalProjDBSize=32000000, sequenceCol="sequence" + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("2.4.0") + def setParams( + self, + *, + minSupport: float = 0.1, + maxPatternLength: int = 10, + maxLocalProjDBSize: int = 32000000, + sequenceCol: str = "sequence", + ) -> "PrefixSpan": + """ + setParams(self, \\*, minSupport=0.1, maxPatternLength=10, maxLocalProjDBSize=32000000, \ + sequenceCol="sequence") + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("3.0.0") + def setMinSupport(self, value: float) -> "PrefixSpan": + """ + Sets the value of :py:attr:`minSupport`. + """ + return self._set(minSupport=value) + + @since("3.0.0") + def getMinSupport(self) -> float: + """ + Gets the value of minSupport or its default value. + """ + return self.getOrDefault(self.minSupport) + + @since("3.0.0") + def setMaxPatternLength(self, value: int) -> "PrefixSpan": + """ + Sets the value of :py:attr:`maxPatternLength`. + """ + return self._set(maxPatternLength=value) + + @since("3.0.0") + def getMaxPatternLength(self) -> int: + """ + Gets the value of maxPatternLength or its default value. + """ + return self.getOrDefault(self.maxPatternLength) + + @since("3.0.0") + def setMaxLocalProjDBSize(self, value: int) -> "PrefixSpan": + """ + Sets the value of :py:attr:`maxLocalProjDBSize`. + """ + return self._set(maxLocalProjDBSize=value) + + @since("3.0.0") + def getMaxLocalProjDBSize(self) -> int: + """ + Gets the value of maxLocalProjDBSize or its default value. + """ + return self.getOrDefault(self.maxLocalProjDBSize) + + @since("3.0.0") + def setSequenceCol(self, value: str) -> "PrefixSpan": + """ + Sets the value of :py:attr:`sequenceCol`. + """ + return self._set(sequenceCol=value) + + @since("3.0.0") + def getSequenceCol(self) -> str: + """ + Gets the value of sequenceCol or its default value. + """ + return self.getOrDefault(self.sequenceCol) + + def findFrequentSequentialPatterns(self, dataset: DataFrame) -> DataFrame: + """ + Finds the complete set of frequent sequential patterns in the input sequences of itemsets. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + A dataframe containing a sequence column which is + `ArrayType(ArrayType(T))` type, T is the item type for the input dataset. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + A `DataFrame` that contains columns of sequence and corresponding frequency. + The schema of it will be: + + - `sequence: ArrayType(ArrayType(T))` (T is the item type) + - `freq: Long` + """ + + self._transfer_params_to_java() + assert self._java_obj is not None + jdf = self._java_obj.findFrequentSequentialPatterns(dataset._jdf) + return DataFrame(jdf, dataset.sparkSession) + + +if __name__ == "__main__": + import doctest + import pyspark.ml.fpm + from pyspark.sql import SparkSession + + globs = pyspark.ml.fpm.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.fpm tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/functions.py new file mode 100644 index 0000000..89b05b6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/functions.py @@ -0,0 +1,842 @@ +# +# 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 numpy as np +import pandas as pd +import uuid +from pyspark import SparkContext +from pyspark.sql.functions import pandas_udf +from pyspark.sql.column import Column, _to_java_column +from pyspark.sql.types import ( + ArrayType, + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, + StructType, +) +from pyspark.ml.util import try_remote_functions +from typing import Any, Callable, Iterator, List, Mapping, TYPE_CHECKING, Tuple, Union, Optional + +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)] + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._jvm is not None + return Column( + 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]))] + """ + sc = SparkContext._active_spark_context + assert sc is not None and sc._jvm is not None + return Column(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]) + 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 + + 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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/image.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/image.py new file mode 100644 index 0000000..6dc97ac --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/image.py @@ -0,0 +1,271 @@ +# +# 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. +# + +""" +.. attribute:: ImageSchema + + An attribute of this module that contains the instance of :class:`_ImageSchema`. + +.. autoclass:: _ImageSchema + :members: +""" + +import sys +from typing import Any, Dict, List, NoReturn, Optional, cast + +import numpy as np +from distutils.version import LooseVersion + +from pyspark import SparkContext +from pyspark.sql.types import Row, StructType, _create_row, _parse_datatype_json_string +from pyspark.sql import SparkSession + +__all__ = ["ImageSchema"] + + +class _ImageSchema: + """ + Internal class for `pyspark.ml.image.ImageSchema` attribute. Meant to be private and + not to be instantized. Use `pyspark.ml.image.ImageSchema` attribute to access the + APIs of this class. + """ + + def __init__(self) -> None: + self._imageSchema: Optional[StructType] = None + self._ocvTypes: Optional[Dict[str, int]] = None + self._columnSchema: Optional[StructType] = None + self._imageFields: Optional[List[str]] = None + self._undefinedImageType: Optional[str] = None + + @property + def imageSchema(self) -> StructType: + """ + Returns the image schema. + + Returns + ------- + :class:`StructType` + with a single column of images named "image" (nullable) + and having the same type returned by :meth:`columnSchema`. + + .. versionadded:: 2.3.0 + """ + + if self._imageSchema is None: + ctx = SparkContext._active_spark_context + assert ctx is not None and ctx._jvm is not None + jschema = ctx._jvm.org.apache.spark.ml.image.ImageSchema.imageSchema() + self._imageSchema = cast(StructType, _parse_datatype_json_string(jschema.json())) + return self._imageSchema + + @property + def ocvTypes(self) -> Dict[str, int]: + """ + Returns the OpenCV type mapping supported. + + Returns + ------- + dict + a dictionary containing the OpenCV type mapping supported. + + .. versionadded:: 2.3.0 + """ + + if self._ocvTypes is None: + ctx = SparkContext._active_spark_context + assert ctx is not None and ctx._jvm is not None + self._ocvTypes = dict(ctx._jvm.org.apache.spark.ml.image.ImageSchema.javaOcvTypes()) + return self._ocvTypes + + @property + def columnSchema(self) -> StructType: + """ + Returns the schema for the image column. + + Returns + ------- + :class:`StructType` + a schema for image column, + ``struct``. + + .. versionadded:: 2.4.0 + """ + + if self._columnSchema is None: + ctx = SparkContext._active_spark_context + assert ctx is not None and ctx._jvm is not None + jschema = ctx._jvm.org.apache.spark.ml.image.ImageSchema.columnSchema() + self._columnSchema = cast(StructType, _parse_datatype_json_string(jschema.json())) + return self._columnSchema + + @property + def imageFields(self) -> List[str]: + """ + Returns field names of image columns. + + Returns + ------- + list + a list of field names. + + .. versionadded:: 2.3.0 + """ + + if self._imageFields is None: + ctx = SparkContext._active_spark_context + assert ctx is not None and ctx._jvm is not None + self._imageFields = list(ctx._jvm.org.apache.spark.ml.image.ImageSchema.imageFields()) + return self._imageFields + + @property + def undefinedImageType(self) -> str: + """ + Returns the name of undefined image type for the invalid image. + + .. versionadded:: 2.3.0 + """ + + if self._undefinedImageType is None: + ctx = SparkContext._active_spark_context + assert ctx is not None and ctx._jvm is not None + self._undefinedImageType = ( + ctx._jvm.org.apache.spark.ml.image.ImageSchema.undefinedImageType() + ) + return self._undefinedImageType + + def toNDArray(self, image: Row) -> np.ndarray: + """ + Converts an image to an array with metadata. + + Parameters + ---------- + image : :class:`Row` + image: A row that contains the image to be converted. It should + have the attributes specified in `ImageSchema.imageSchema`. + + Returns + ------- + :class:`numpy.ndarray` + that is an image. + + .. versionadded:: 2.3.0 + """ + + if not isinstance(image, Row): + raise TypeError( + "image argument should be pyspark.sql.types.Row; however, " + "it got [%s]." % type(image) + ) + + if any(not hasattr(image, f) for f in self.imageFields): + raise ValueError( + "image argument should have attributes specified in " + "ImageSchema.imageSchema [%s]." % ", ".join(self.imageFields) + ) + + height = image.height + width = image.width + nChannels = image.nChannels + return np.ndarray( + shape=(height, width, nChannels), + dtype=np.uint8, + buffer=image.data, + strides=(width * nChannels, nChannels, 1), + ) + + def toImage(self, array: np.ndarray, origin: str = "") -> Row: + """ + Converts an array with metadata to a two-dimensional image. + + Parameters + ---------- + array : :class:`numpy.ndarray` + The array to convert to image. + origin : str + Path to the image, optional. + + Returns + ------- + :class:`Row` + that is a two dimensional image. + + .. versionadded:: 2.3.0 + """ + + if not isinstance(array, np.ndarray): + raise TypeError( + "array argument should be numpy.ndarray; however, it got [%s]." % type(array) + ) + + if array.ndim != 3: + raise ValueError("Invalid array shape") + + height, width, nChannels = array.shape + ocvTypes = ImageSchema.ocvTypes + if nChannels == 1: + mode = ocvTypes["CV_8UC1"] + elif nChannels == 3: + mode = ocvTypes["CV_8UC3"] + elif nChannels == 4: + mode = ocvTypes["CV_8UC4"] + else: + raise ValueError("Invalid number of channels") + + # Running `bytearray(numpy.array([1]))` fails in specific Python versions + # with a specific Numpy version, for example in Python 3.6.0 and NumPy 1.13.3. + # Here, it avoids it by converting it to bytes. + if LooseVersion(np.__version__) >= LooseVersion("1.9"): + data = bytearray(array.astype(dtype=np.uint8).ravel().tobytes()) + else: + # Numpy prior to 1.9 don't have `tobytes` method. + data = bytearray(array.astype(dtype=np.uint8).ravel()) + + # Creating new Row with _create_row(), because Row(name = value, ... ) + # orders fields by name, which conflicts with expected schema order + # when the new DataFrame is created by UDF + return _create_row(self.imageFields, [origin, height, width, nChannels, mode, data]) + + +ImageSchema = _ImageSchema() + + +# Monkey patch to disallow instantiation of this class. +def _disallow_instance(_: Any) -> NoReturn: + raise RuntimeError("Creating instance of _ImageSchema class is disallowed.") + + +_ImageSchema.__init__ = _disallow_instance # type: ignore[assignment] + + +def _test() -> None: + import doctest + import pyspark.ml.image + + globs = pyspark.ml.image.__dict__.copy() + spark = SparkSession.builder.master("local[2]").appName("ml.image tests").getOrCreate() + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.ml.image, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/linalg/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/linalg/__init__.py new file mode 100644 index 0000000..a5a849e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/linalg/__init__.py @@ -0,0 +1,1382 @@ +# +# 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. +# + +""" +MLlib utilities for linear algebra. For dense vectors, MLlib +uses the NumPy `array` type, so you can simply pass NumPy arrays +around. For sparse vectors, users can construct a :class:`SparseVector` +object from MLlib or pass SciPy `scipy.sparse` column vectors if +SciPy is available in their environment. +""" + +import sys +import array +import struct + +import numpy as np + +from pyspark.sql.types import ( + UserDefinedType, + StructField, + StructType, + ArrayType, + DoubleType, + IntegerType, + ByteType, + BooleanType, +) + +from typing import ( + Any, + Callable, + cast, + Dict, + Iterable, + List, + Optional, + overload, + Sequence, + Tuple, + Type, + TYPE_CHECKING, + Union, +) + + +__all__ = [ + "Vector", + "DenseVector", + "SparseVector", + "Vectors", + "Matrix", + "DenseMatrix", + "SparseMatrix", + "Matrices", +] + +if TYPE_CHECKING: + from pyspark.mllib._typing import NormType + from pyspark.ml._typing import VectorLike + from scipy.sparse import spmatrix + + +# Check whether we have SciPy. MLlib works without it too, but if we have it, some methods, +# such as _dot and _serialize_double_vector, start to support scipy.sparse matrices. + +try: + import scipy.sparse + + _have_scipy = True +except BaseException: + # No SciPy in environment, but that's okay + _have_scipy = False + + +def _convert_to_vector(d: Union["VectorLike", "spmatrix", range]) -> "Vector": + if isinstance(d, Vector): + return d + elif type(d) in (array.array, np.array, np.ndarray, list, tuple, range): + return DenseVector(d) + elif _have_scipy and scipy.sparse.issparse(d): + assert cast("spmatrix", d).shape[1] == 1, "Expected column vector" + # Make sure the converted csc_matrix has sorted indices. + csc = cast("spmatrix", d).tocsc() + if not csc.has_sorted_indices: + csc.sort_indices() + return SparseVector(cast("spmatrix", d).shape[0], csc.indices, csc.data) + else: + raise TypeError("Cannot convert type %s into Vector" % type(d)) + + +def _vector_size(v: Union["VectorLike", "spmatrix", range]) -> int: + """ + Returns the size of the vector. + + Examples + -------- + >>> _vector_size([1., 2., 3.]) + 3 + >>> _vector_size((1., 2., 3.)) + 3 + >>> _vector_size(array.array('d', [1., 2., 3.])) + 3 + >>> _vector_size(np.zeros(3)) + 3 + >>> _vector_size(np.zeros((3, 1))) + 3 + >>> _vector_size(np.zeros((1, 3))) + Traceback (most recent call last): + ... + ValueError: Cannot treat an ndarray of shape (1, 3) as a vector + """ + if isinstance(v, Vector): + return len(v) + elif type(v) in (array.array, list, tuple, range): + return len(v) + elif type(v) == np.ndarray: + if v.ndim == 1 or (v.ndim == 2 and v.shape[1] == 1): + return len(v) + else: + raise ValueError("Cannot treat an ndarray of shape %s as a vector" % str(v.shape)) + elif _have_scipy and scipy.sparse.issparse(v): + assert cast("spmatrix", v).shape[1] == 1, "Expected column vector" + return cast("spmatrix", v).shape[0] + else: + raise TypeError("Cannot treat type %s as a vector" % type(v)) + + +def _format_float(f: float, digits: int = 4) -> str: + s = str(round(f, digits)) + if "." in s: + s = s[: s.index(".") + 1 + digits] + return s + + +def _format_float_list(xs: Iterable[float]) -> List[str]: + return [_format_float(x) for x in xs] + + +def _double_to_long_bits(value: float) -> int: + if np.isnan(value): + value = float("nan") + # pack double into 64 bits, then unpack as long int + return struct.unpack("Q", struct.pack("d", value))[0] + + +class VectorUDT(UserDefinedType): + """ + SQL user-defined type (UDT) for Vector. + """ + + @classmethod + def sqlType(cls) -> StructType: + return StructType( + [ + StructField("type", ByteType(), False), + StructField("size", IntegerType(), True), + StructField("indices", ArrayType(IntegerType(), False), True), + StructField("values", ArrayType(DoubleType(), False), True), + ] + ) + + @classmethod + def module(cls) -> str: + return "pyspark.ml.linalg" + + @classmethod + def scalaUDT(cls) -> str: + return "org.apache.spark.ml.linalg.VectorUDT" + + def serialize( + self, obj: "Vector" + ) -> Tuple[int, Optional[int], Optional[List[int]], List[float]]: + if isinstance(obj, SparseVector): + indices = [int(i) for i in obj.indices] + values = [float(v) for v in obj.values] + return (0, obj.size, indices, values) + elif isinstance(obj, DenseVector): + values = [float(v) for v in obj] # type: ignore[attr-defined] + return (1, None, None, values) + else: + raise TypeError("cannot serialize %r of type %r" % (obj, type(obj))) + + def deserialize( + self, datum: Tuple[int, Optional[int], Optional[List[int]], List[float]] + ) -> "Vector": + assert ( + len(datum) == 4 + ), "VectorUDT.deserialize given row with length %d but requires 4" % len(datum) + tpe = datum[0] + if tpe == 0: + return SparseVector(cast(int, datum[1]), cast(List[int], datum[2]), datum[3]) + elif tpe == 1: + return DenseVector(datum[3]) + else: + raise ValueError("do not recognize type %r" % tpe) + + def simpleString(self) -> str: + return "vector" + + +class MatrixUDT(UserDefinedType): + """ + SQL user-defined type (UDT) for Matrix. + """ + + @classmethod + def sqlType(cls) -> StructType: + return StructType( + [ + StructField("type", ByteType(), False), + StructField("numRows", IntegerType(), False), + StructField("numCols", IntegerType(), False), + StructField("colPtrs", ArrayType(IntegerType(), False), True), + StructField("rowIndices", ArrayType(IntegerType(), False), True), + StructField("values", ArrayType(DoubleType(), False), True), + StructField("isTransposed", BooleanType(), False), + ] + ) + + @classmethod + def module(cls) -> str: + return "pyspark.ml.linalg" + + @classmethod + def scalaUDT(cls) -> str: + return "org.apache.spark.ml.linalg.MatrixUDT" + + def serialize( + self, obj: "Matrix" + ) -> Tuple[int, int, int, Optional[List[int]], Optional[List[int]], List[float], bool]: + if isinstance(obj, SparseMatrix): + colPtrs = [int(i) for i in obj.colPtrs] + rowIndices = [int(i) for i in obj.rowIndices] + values = [float(v) for v in obj.values] + return ( + 0, + obj.numRows, + obj.numCols, + colPtrs, + rowIndices, + values, + bool(obj.isTransposed), + ) + elif isinstance(obj, DenseMatrix): + values = [float(v) for v in obj.values] + return (1, obj.numRows, obj.numCols, None, None, values, bool(obj.isTransposed)) + else: + raise TypeError("cannot serialize type %r" % (type(obj))) + + def deserialize( + self, + datum: Tuple[int, int, int, Optional[List[int]], Optional[List[int]], List[float], bool], + ) -> "Matrix": + assert ( + len(datum) == 7 + ), "MatrixUDT.deserialize given row with length %d but requires 7" % len(datum) + tpe = datum[0] + if tpe == 0: + return SparseMatrix(*datum[1:]) # type: ignore[arg-type] + elif tpe == 1: + return DenseMatrix(datum[1], datum[2], datum[5], datum[6]) + else: + raise ValueError("do not recognize type %r" % tpe) + + def simpleString(self) -> str: + return "matrix" + + +class Vector: + + __UDT__ = VectorUDT() + + """ + Abstract class for DenseVector and SparseVector + """ + + def toArray(self) -> np.ndarray: + """ + Convert the vector into an numpy.ndarray + + :return: numpy.ndarray + """ + raise NotImplementedError + + def __len__(self) -> int: + raise NotImplementedError + + +class DenseVector(Vector): + """ + A dense vector represented by a value array. We use numpy array for + storage and arithmetics will be delegated to the underlying numpy + array. + + Examples + -------- + >>> v = Vectors.dense([1.0, 2.0]) + >>> u = Vectors.dense([3.0, 4.0]) + >>> v + u + DenseVector([4.0, 6.0]) + >>> 2 - v + DenseVector([1.0, 0.0]) + >>> v / 2 + DenseVector([0.5, 1.0]) + >>> v * u + DenseVector([3.0, 8.0]) + >>> u / v + DenseVector([3.0, 2.0]) + >>> u % 2 + DenseVector([1.0, 0.0]) + >>> -v + DenseVector([-1.0, -2.0]) + """ + + def __init__(self, ar: Union[bytes, np.ndarray, Iterable[float]]): + ar_: np.ndarray + if isinstance(ar, bytes): + ar_ = np.frombuffer(ar, dtype=np.float64) + elif not isinstance(ar, np.ndarray): + ar_ = np.array(ar, dtype=np.float64) + else: + ar_ = ar.astype(np.float64) if ar.dtype != np.float64 else ar + self.array = ar_ + + def __reduce__(self) -> Tuple[Type["DenseVector"], Tuple[bytes]]: + return DenseVector, (self.array.tobytes(),) + + def numNonzeros(self) -> int: + """ + Number of nonzero elements. This scans all active values and count non zeros + """ + return np.count_nonzero(self.array) + + def norm(self, p: "NormType") -> np.float64: + """ + Calculates the norm of a DenseVector. + + Examples + -------- + >>> a = DenseVector([0, -1, 2, -3]) + >>> a.norm(2) + 3.7... + >>> a.norm(1) + 6.0 + """ + return np.linalg.norm(self.array, p) + + def dot(self, other: Iterable[float]) -> np.float64: + """ + Compute the dot product of two Vectors. We support + (Numpy array, list, SparseVector, or SciPy sparse) + and a target NumPy array that is either 1- or 2-dimensional. + Equivalent to calling numpy.dot of the two vectors. + + Examples + -------- + >>> dense = DenseVector(array.array('d', [1., 2.])) + >>> dense.dot(dense) + 5.0 + >>> dense.dot(SparseVector(2, [0, 1], [2., 1.])) + 4.0 + >>> dense.dot(range(1, 3)) + 5.0 + >>> dense.dot(np.array(range(1, 3))) + 5.0 + >>> dense.dot([1.,]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> dense.dot(np.reshape([1., 2., 3., 4.], (2, 2), order='F')) + array([ 5., 11.]) + >>> dense.dot(np.reshape([1., 2., 3.], (3, 1), order='F')) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + if type(other) == np.ndarray: + if other.ndim > 1: + assert len(self) == other.shape[0], "dimension mismatch" + return np.dot(self.array, other) + elif _have_scipy and scipy.sparse.issparse(other): + assert len(self) == cast("spmatrix", other).shape[0], "dimension mismatch" + return cast("spmatrix", other).transpose().dot(self.toArray()) + else: + assert len(self) == _vector_size(other), "dimension mismatch" + if isinstance(other, SparseVector): + return other.dot(self) + elif isinstance(other, Vector): + return np.dot(self.toArray(), other.toArray()) + else: + return np.dot(self.toArray(), other) # type: ignore[call-overload] + + def squared_distance(self, other: Iterable[float]) -> np.float64: + """ + Squared distance of two Vectors. + + Examples + -------- + >>> dense1 = DenseVector(array.array('d', [1., 2.])) + >>> dense1.squared_distance(dense1) + 0.0 + >>> dense2 = np.array([2., 1.]) + >>> dense1.squared_distance(dense2) + 2.0 + >>> dense3 = [2., 1.] + >>> dense1.squared_distance(dense3) + 2.0 + >>> sparse1 = SparseVector(2, [0, 1], [2., 1.]) + >>> dense1.squared_distance(sparse1) + 2.0 + >>> dense1.squared_distance([1.,]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> dense1.squared_distance(SparseVector(1, [0,], [1.,])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + assert len(self) == _vector_size(other), "dimension mismatch" + if isinstance(other, SparseVector): + return other.squared_distance(self) + elif _have_scipy and scipy.sparse.issparse(other): + return _convert_to_vector(other).squared_distance(self) # type: ignore[attr-defined] + + if isinstance(other, Vector): + other = other.toArray() + elif not isinstance(other, np.ndarray): + other = np.array(other) + diff: np.ndarray = self.toArray() - other + return np.dot(diff, diff) + + def toArray(self) -> np.ndarray: + """ + Returns the underlying numpy.ndarray + """ + return self.array + + @property + def values(self) -> np.ndarray: + """ + Returns the underlying numpy.ndarray + """ + return self.array + + @overload + def __getitem__(self, item: int) -> np.float64: + ... + + @overload + def __getitem__(self, item: slice) -> np.ndarray: + ... + + def __getitem__(self, item: Union[int, slice]) -> Union[np.float64, np.ndarray]: + return self.array[item] + + def __len__(self) -> int: + return len(self.array) + + def __str__(self) -> str: + return "[" + ",".join([str(v) for v in self.array]) + "]" + + def __repr__(self) -> str: + return "DenseVector([%s])" % (", ".join(_format_float(i) for i in self.array)) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, DenseVector): + return np.array_equal(self.array, other.array) + elif isinstance(other, SparseVector): + if len(self) != other.size: + return False + return Vectors._equals(list(range(len(self))), self.array, other.indices, other.values) + return False + + def __ne__(self, other: Any) -> bool: + return not self == other + + def __hash__(self) -> int: + size = len(self) + result = 31 + size + nnz = 0 + i = 0 + while i < size and nnz < 128: + if self.array[i] != 0: + result = 31 * result + i + bits = _double_to_long_bits(self.array[i]) + result = 31 * result + (bits ^ (bits >> 32)) + nnz += 1 + i += 1 + return result + + def __getattr__(self, item: str) -> Any: + return getattr(self.array, item) + + def __neg__(self) -> "DenseVector": + return DenseVector(-self.array) + + def _delegate(op: str) -> Callable[["DenseVector", Any], "DenseVector"]: # type: ignore[misc] + def func(self: "DenseVector", other: Any) -> "DenseVector": + if isinstance(other, DenseVector): + other = other.array + return DenseVector(getattr(self.array, op)(other)) + + return func + + __add__ = _delegate("__add__") + __sub__ = _delegate("__sub__") + __mul__ = _delegate("__mul__") + __div__ = _delegate("__div__") + __truediv__ = _delegate("__truediv__") + __mod__ = _delegate("__mod__") + __radd__ = _delegate("__radd__") + __rsub__ = _delegate("__rsub__") + __rmul__ = _delegate("__rmul__") + __rdiv__ = _delegate("__rdiv__") + __rtruediv__ = _delegate("__rtruediv__") + __rmod__ = _delegate("__rmod__") + + +class SparseVector(Vector): + """ + A simple sparse vector class for passing data to MLlib. Users may + alternatively pass SciPy's {scipy.sparse} data types. + """ + + @overload + def __init__(self, size: int, __indices: bytes, __values: bytes): + ... + + @overload + def __init__(self, size: int, *args: Tuple[int, float]): + ... + + @overload + def __init__(self, size: int, __indices: Iterable[int], __values: Iterable[float]): + ... + + @overload + def __init__(self, size: int, __pairs: Iterable[Tuple[int, float]]): + ... + + @overload + def __init__(self, size: int, __map: Dict[int, float]): + ... + + def __init__( + self, + size: int, + *args: Union[ + bytes, Tuple[int, float], Iterable[float], Iterable[Tuple[int, float]], Dict[int, float] + ], + ): + """ + Create a sparse vector, using either a dictionary, a list of + (index, value) pairs, or two separate arrays of indices and + values (sorted by index). + + Examples + -------- + size : int + Size of the vector. + args + Active entries, as a dictionary {index: value, ...}, + a list of tuples [(index, value), ...], or a list of strictly + increasing indices and a list of corresponding values [index, ...], + [value, ...]. Inactive entries are treated as zeros. + + Examples + -------- + >>> SparseVector(4, {1: 1.0, 3: 5.5}) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> SparseVector(4, [(1, 1.0), (3, 5.5)]) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> SparseVector(4, [1, 3], [1.0, 5.5]) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> SparseVector(4, {1:1.0, 6:2.0}) + Traceback (most recent call last): + ... + AssertionError: Index 6 is out of the size of vector with size=4 + >>> SparseVector(4, {-1:1.0}) + Traceback (most recent call last): + ... + AssertionError: Contains negative index -1 + """ + self.size = int(size) + """ Size of the vector. """ + assert 1 <= len(args) <= 2, "must pass either 2 or 3 arguments" + if len(args) == 1: + pairs = args[0] + if type(pairs) == dict: + pairs = pairs.items() + pairs = cast(Iterable[Tuple[int, float]], sorted(pairs)) + self.indices = np.array([p[0] for p in pairs], dtype=np.int32) + """ A list of indices corresponding to active entries. """ + self.values = np.array([p[1] for p in pairs], dtype=np.float64) + """ A list of values corresponding to active entries. """ + else: + if isinstance(args[0], bytes): + assert isinstance(args[1], bytes), "values should be string too" + if args[0]: + self.indices = np.frombuffer(args[0], np.int32) + self.values = np.frombuffer(args[1], np.float64) + else: + # np.frombuffer() doesn't work well with empty string in older version + self.indices = np.array([], dtype=np.int32) + self.values = np.array([], dtype=np.float64) + else: + self.indices = np.array(args[0], dtype=np.int32) + self.values = np.array(args[1], dtype=np.float64) + assert len(self.indices) == len(self.values), "index and value arrays not same length" + for i in range(len(self.indices) - 1): + if self.indices[i] >= self.indices[i + 1]: + raise TypeError( + "Indices %s and %s are not strictly increasing" + % (self.indices[i], self.indices[i + 1]) + ) + + if self.indices.size > 0: + assert ( + np.max(self.indices) < self.size + ), "Index %d is out of the size of vector with size=%d" % ( + np.max(self.indices), + self.size, + ) + assert np.min(self.indices) >= 0, "Contains negative index %d" % (np.min(self.indices)) + + def numNonzeros(self) -> int: + """ + Number of nonzero elements. This scans all active values and count non zeros. + """ + return np.count_nonzero(self.values) + + def norm(self, p: "NormType") -> np.float64: + """ + Calculates the norm of a SparseVector. + + Examples + -------- + >>> a = SparseVector(4, [0, 1], [3., -4.]) + >>> a.norm(1) + 7.0 + >>> a.norm(2) + 5.0 + """ + return np.linalg.norm(self.values, p) + + def __reduce__(self) -> Tuple[Type["SparseVector"], Tuple[int, bytes, bytes]]: + return (SparseVector, (self.size, self.indices.tobytes(), self.values.tobytes())) + + def dot(self, other: Iterable[float]) -> np.float64: + """ + Dot product with a SparseVector or 1- or 2-dimensional Numpy array. + + Examples + -------- + >>> a = SparseVector(4, [1, 3], [3.0, 4.0]) + >>> a.dot(a) + 25.0 + >>> a.dot(array.array('d', [1., 2., 3., 4.])) + 22.0 + >>> b = SparseVector(4, [2], [1.0]) + >>> a.dot(b) + 0.0 + >>> a.dot(np.array([[1, 1], [2, 2], [3, 3], [4, 4]])) + array([ 22., 22.]) + >>> a.dot([1., 2., 3.]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(np.array([1., 2.])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(DenseVector([1., 2.])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(np.zeros((3, 2))) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + + if isinstance(other, np.ndarray): + if other.ndim not in [2, 1]: + raise ValueError("Cannot call dot with %d-dimensional array" % other.ndim) + assert len(self) == other.shape[0], "dimension mismatch" + return np.dot(self.values, other[self.indices]) + + assert len(self) == _vector_size(other), "dimension mismatch" + + if isinstance(other, DenseVector): + return np.dot(other.array[self.indices], self.values) + + elif isinstance(other, SparseVector): + # Find out common indices. + self_cmind = np.in1d(self.indices, other.indices, assume_unique=True) + self_values = self.values[self_cmind] + if self_values.size == 0: + return np.float64(0.0) + else: + other_cmind = np.in1d(other.indices, self.indices, assume_unique=True) + return np.dot(self_values, other.values[other_cmind]) + + else: + return self.dot(_convert_to_vector(other)) # type: ignore[arg-type] + + def squared_distance(self, other: Iterable[float]) -> np.float64: + """ + Squared distance from a SparseVector or 1-dimensional NumPy array. + + Examples + -------- + >>> a = SparseVector(4, [1, 3], [3.0, 4.0]) + >>> a.squared_distance(a) + 0.0 + >>> a.squared_distance(array.array('d', [1., 2., 3., 4.])) + 11.0 + >>> a.squared_distance(np.array([1., 2., 3., 4.])) + 11.0 + >>> b = SparseVector(4, [2], [1.0]) + >>> a.squared_distance(b) + 26.0 + >>> b.squared_distance(a) + 26.0 + >>> b.squared_distance([1., 2.]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> b.squared_distance(SparseVector(3, [1,], [1.0,])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + assert len(self) == _vector_size(other), "dimension mismatch" + + if isinstance(other, np.ndarray) or isinstance(other, DenseVector): + if isinstance(other, np.ndarray) and other.ndim != 1: + raise ValueError( + "Cannot call squared_distance with %d-dimensional array" % other.ndim + ) + if isinstance(other, DenseVector): + other = other.array + sparse_ind = np.zeros(other.size, dtype=bool) + sparse_ind[self.indices] = True + dist = other[sparse_ind] - self.values + result = np.dot(dist, dist) + + other_ind = other[~sparse_ind] + result += np.dot(other_ind, other_ind) + return result + + elif isinstance(other, SparseVector): + result = 0.0 + i, j = 0, 0 + while i < len(self.indices) and j < len(other.indices): + if self.indices[i] == other.indices[j]: + diff = self.values[i] - other.values[j] + result += diff * diff + i += 1 + j += 1 + elif self.indices[i] < other.indices[j]: + result += self.values[i] * self.values[i] + i += 1 + else: + result += other.values[j] * other.values[j] + j += 1 + while i < len(self.indices): + result += self.values[i] * self.values[i] + i += 1 + while j < len(other.indices): + result += other.values[j] * other.values[j] + j += 1 + return result + else: + return self.squared_distance(_convert_to_vector(other)) # type: ignore[arg-type] + + def toArray(self) -> np.ndarray: + """ + Returns a copy of this SparseVector as a 1-dimensional numpy.ndarray. + """ + arr = np.zeros((self.size,), dtype=np.float64) + arr[self.indices] = self.values + return arr + + def __len__(self) -> int: + return self.size + + def __str__(self) -> str: + inds = "[" + ",".join([str(i) for i in self.indices]) + "]" + vals = "[" + ",".join([str(v) for v in self.values]) + "]" + return "(" + ",".join((str(self.size), inds, vals)) + ")" + + def __repr__(self) -> str: + inds = self.indices + vals = self.values + entries = ", ".join( + ["{0}: {1}".format(inds[i], _format_float(vals[i])) for i in range(len(inds))] + ) + return "SparseVector({0}, {{{1}}})".format(self.size, entries) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, SparseVector): + return ( + other.size == self.size + and np.array_equal(other.indices, self.indices) + and np.array_equal(other.values, self.values) + ) + elif isinstance(other, DenseVector): + if self.size != len(other): + return False + return Vectors._equals(self.indices, self.values, list(range(len(other))), other.array) + return False + + def __getitem__(self, index: int) -> np.float64: + inds = self.indices + vals = self.values + if not isinstance(index, int): + raise TypeError("Indices must be of type integer, got type %s" % type(index)) + + if index >= self.size or index < -self.size: + raise IndexError("Index %d out of bounds." % index) + if index < 0: + index += self.size + + if (inds.size == 0) or (index > inds.item(-1)): + return np.float64(0.0) + + insert_index = np.searchsorted(inds, index) + row_ind = inds[insert_index] + if row_ind == index: + return vals[insert_index] + return np.float64(0.0) + + def __ne__(self, other: Any) -> bool: + return not self.__eq__(other) + + def __hash__(self) -> int: + result = 31 + self.size + nnz = 0 + i = 0 + while i < len(self.values) and nnz < 128: + if self.values[i] != 0: + result = 31 * result + int(self.indices[i]) + bits = _double_to_long_bits(self.values[i]) + result = 31 * result + (bits ^ (bits >> 32)) + nnz += 1 + i += 1 + return result + + +class Vectors: + + """ + Factory methods for working with vectors. + + Notes + ----- + Dense vectors are simply represented as NumPy array objects, + so there is no need to convert them for use in MLlib. For sparse vectors, + the factory methods in this class create an MLlib-compatible type, or users + can pass in SciPy's `scipy.sparse` column vectors. + """ + + @staticmethod + @overload + def sparse(size: int, __indices: bytes, __values: bytes) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, *args: Tuple[int, float]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __indices: Iterable[int], __values: Iterable[float]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __pairs: Iterable[Tuple[int, float]]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __map: Dict[int, float]) -> SparseVector: + ... + + @staticmethod + def sparse( + size: int, + *args: Union[ + bytes, Tuple[int, float], Iterable[float], Iterable[Tuple[int, float]], Dict[int, float] + ], + ) -> SparseVector: + """ + Create a sparse vector, using either a dictionary, a list of + (index, value) pairs, or two separate arrays of indices and + values (sorted by index). + + Parameters + ---------- + size : int + Size of the vector. + args + Non-zero entries, as a dictionary, list of tuples, + or two sorted lists containing indices and values. + + Examples + -------- + >>> Vectors.sparse(4, {1: 1.0, 3: 5.5}) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> Vectors.sparse(4, [(1, 1.0), (3, 5.5)]) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> Vectors.sparse(4, [1, 3], [1.0, 5.5]) + SparseVector(4, {1: 1.0, 3: 5.5}) + """ + return SparseVector(size, *args) # type: ignore[arg-type] + + @overload + @staticmethod + def dense(*elements: float) -> DenseVector: + ... + + @overload + @staticmethod + def dense(__arr: bytes) -> DenseVector: + ... + + @overload + @staticmethod + def dense(__arr: Iterable[float]) -> DenseVector: + ... + + @staticmethod + def dense(*elements: Union[float, bytes, np.ndarray, Iterable[float]]) -> DenseVector: + """ + Create a dense vector of 64-bit floats from a Python list or numbers. + + Examples + -------- + >>> Vectors.dense([1, 2, 3]) + DenseVector([1.0, 2.0, 3.0]) + >>> Vectors.dense(1.0, 2.0) + DenseVector([1.0, 2.0]) + """ + if len(elements) == 1 and not isinstance(elements[0], (float, int)): + # it's list, numpy.array or other iterable object. + elements = elements[0] # type: ignore[assignment] + return DenseVector(cast(Iterable[float], elements)) + + @staticmethod + def squared_distance(v1: Vector, v2: Vector) -> np.float64: + """ + Squared distance between two vectors. + a and b can be of type SparseVector, DenseVector, np.ndarray + or array.array. + + Examples + -------- + >>> a = Vectors.sparse(4, [(0, 1), (3, 4)]) + >>> b = Vectors.dense([2, 5, 4, 1]) + >>> a.squared_distance(b) + 51.0 + """ + v1, v2 = _convert_to_vector(v1), _convert_to_vector(v2) + return v1.squared_distance(v2) # type: ignore[attr-defined] + + @staticmethod + def norm(vector: Vector, p: "NormType") -> np.float64: + """ + Find norm of the given vector. + """ + return _convert_to_vector(vector).norm(p) # type: ignore[attr-defined] + + @staticmethod + def zeros(size: int) -> DenseVector: + return DenseVector(np.zeros(size)) + + @staticmethod + def _equals( + v1_indices: Union[Sequence[int], np.ndarray], + v1_values: Union[Sequence[float], np.ndarray], + v2_indices: Union[Sequence[int], np.ndarray], + v2_values: Union[Sequence[float], np.ndarray], + ) -> bool: + """ + Check equality between sparse/dense vectors, + v1_indices and v2_indices assume to be strictly increasing. + """ + v1_size = len(v1_values) + v2_size = len(v2_values) + k1 = 0 + k2 = 0 + all_equal = True + while all_equal: + while k1 < v1_size and v1_values[k1] == 0: + k1 += 1 + while k2 < v2_size and v2_values[k2] == 0: + k2 += 1 + + if k1 >= v1_size or k2 >= v2_size: + return k1 >= v1_size and k2 >= v2_size + + all_equal = v1_indices[k1] == v2_indices[k2] and v1_values[k1] == v2_values[k2] + k1 += 1 + k2 += 1 + return all_equal + + +class Matrix: + + __UDT__ = MatrixUDT() + + """ + Represents a local matrix. + """ + + def __init__(self, numRows: int, numCols: int, isTransposed: bool = False): + self.numRows = numRows + self.numCols = numCols + self.isTransposed = isTransposed + + def toArray(self) -> np.ndarray: + """ + Returns its elements in a numpy.ndarray. + """ + raise NotImplementedError + + @staticmethod + def _convert_to_array(array_like: Union[bytes, Iterable[float]], dtype: Any) -> np.ndarray: + """ + Convert Matrix attributes which are array-like or buffer to array. + """ + if isinstance(array_like, bytes): + return np.frombuffer(array_like, dtype=dtype) + return np.asarray(array_like, dtype=dtype) + + +class DenseMatrix(Matrix): + """ + Column-major dense matrix. + """ + + def __init__( + self, + numRows: int, + numCols: int, + values: Union[bytes, Iterable[float]], + isTransposed: bool = False, + ): + Matrix.__init__(self, numRows, numCols, isTransposed) + values = self._convert_to_array(values, np.float64) + assert len(values) == numRows * numCols + self.values = values + + def __reduce__(self) -> Tuple[Type["DenseMatrix"], Tuple[int, int, bytes, int]]: + return DenseMatrix, ( + self.numRows, + self.numCols, + self.values.tobytes(), + int(self.isTransposed), + ) + + def __str__(self) -> str: + """ + Pretty printing of a DenseMatrix + + Examples + -------- + >>> dm = DenseMatrix(2, 2, range(4)) + >>> print(dm) + DenseMatrix([[ 0., 2.], + [ 1., 3.]]) + >>> dm = DenseMatrix(2, 2, range(4), isTransposed=True) + >>> print(dm) + DenseMatrix([[ 0., 1.], + [ 2., 3.]]) + """ + # Inspired by __repr__ in scipy matrices. + array_lines = repr(self.toArray()).splitlines() + + # We need to adjust six spaces which is the difference in number + # of letters between "DenseMatrix" and "array" + x = "\n".join([(" " * 6 + line) for line in array_lines[1:]]) + return array_lines[0].replace("array", "DenseMatrix") + "\n" + x + + def __repr__(self) -> str: + """ + Representation of a DenseMatrix + + Examples + -------- + >>> dm = DenseMatrix(2, 2, range(4)) + >>> dm + DenseMatrix(2, 2, [0.0, 1.0, 2.0, 3.0], False) + """ + # If the number of values are less than seventeen then return as it is. + # Else return first eight values and last eight values. + if len(self.values) < 17: + entries = _format_float_list(self.values) + else: + entries = ( + _format_float_list(self.values[:8]) + ["..."] + _format_float_list(self.values[-8:]) + ) + + entries = ", ".join(entries) # type: ignore[assignment] + return "DenseMatrix({0}, {1}, [{2}], {3})".format( + self.numRows, self.numCols, entries, self.isTransposed + ) + + def toArray(self) -> np.ndarray: + """ + Return a :py:class:`numpy.ndarray` + + Examples + -------- + >>> m = DenseMatrix(2, 2, range(4)) + >>> m.toArray() + array([[ 0., 2.], + [ 1., 3.]]) + """ + if self.isTransposed: + return np.asfortranarray(self.values.reshape((self.numRows, self.numCols))) + else: + return self.values.reshape((self.numRows, self.numCols), order="F") + + def toSparse(self) -> "SparseMatrix": + """Convert to SparseMatrix""" + if self.isTransposed: + values = np.ravel(self.toArray(), order="F") + else: + values = self.values + indices = np.nonzero(values)[0] + colCounts = np.bincount(indices // self.numRows) + colPtrs = np.cumsum(np.hstack((0, colCounts, np.zeros(self.numCols - colCounts.size)))) + values = values[indices] + rowIndices = indices % self.numRows + + return SparseMatrix(self.numRows, self.numCols, colPtrs, rowIndices, values) + + def __getitem__(self, indices: Tuple[int, int]) -> np.float64: + i, j = indices + if i < 0 or i >= self.numRows: + raise IndexError("Row index %d is out of range [0, %d)" % (i, self.numRows)) + if j >= self.numCols or j < 0: + raise IndexError("Column index %d is out of range [0, %d)" % (j, self.numCols)) + + if self.isTransposed: + return self.values[i * self.numCols + j] + else: + return self.values[i + j * self.numRows] + + def __eq__(self, other: Any) -> bool: + if self.numRows != other.numRows or self.numCols != other.numCols: + return False + if isinstance(other, SparseMatrix): + return np.all(self.toArray() == other.toArray()).tolist() + + self_values = np.ravel(self.toArray(), order="F") + other_values = np.ravel(other.toArray(), order="F") + return np.all(self_values == other_values).tolist() + + +class SparseMatrix(Matrix): + """Sparse Matrix stored in CSC format.""" + + def __init__( + self, + numRows: int, + numCols: int, + colPtrs: Union[bytes, Iterable[int]], + rowIndices: Union[bytes, Iterable[int]], + values: Union[bytes, Iterable[float]], + isTransposed: bool = False, + ): + Matrix.__init__(self, numRows, numCols, isTransposed) + self.colPtrs = self._convert_to_array(colPtrs, np.int32) + self.rowIndices = self._convert_to_array(rowIndices, np.int32) + self.values = self._convert_to_array(values, np.float64) + + if self.isTransposed: + if self.colPtrs.size != numRows + 1: + raise ValueError( + "Expected colPtrs of size %d, got %d." % (numRows + 1, self.colPtrs.size) + ) + else: + if self.colPtrs.size != numCols + 1: + raise ValueError( + "Expected colPtrs of size %d, got %d." % (numCols + 1, self.colPtrs.size) + ) + if self.rowIndices.size != self.values.size: + raise ValueError( + "Expected rowIndices of length %d, got %d." + % (self.rowIndices.size, self.values.size) + ) + + def __str__(self) -> str: + """ + Pretty printing of a SparseMatrix + + Examples + -------- + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + >>> print(sm1) + 2 X 2 CSCMatrix + (0,0) 2.0 + (1,0) 3.0 + (1,1) 4.0 + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + >>> print(sm1) + 2 X 2 CSRMatrix + (0,0) 2.0 + (0,1) 3.0 + (1,1) 4.0 + """ + spstr = "{0} X {1} ".format(self.numRows, self.numCols) + if self.isTransposed: + spstr += "CSRMatrix\n" + else: + spstr += "CSCMatrix\n" + + cur_col = 0 + smlist = [] + + # Display first 16 values. + if len(self.values) <= 16: + zipindval = zip(self.rowIndices, self.values) + else: + zipindval = zip(self.rowIndices[:16], self.values[:16]) + for i, (rowInd, value) in enumerate(zipindval): + if self.colPtrs[cur_col + 1] <= i: + cur_col += 1 + if self.isTransposed: + smlist.append("({0},{1}) {2}".format(cur_col, rowInd, _format_float(value))) + else: + smlist.append("({0},{1}) {2}".format(rowInd, cur_col, _format_float(value))) + spstr += "\n".join(smlist) + + if len(self.values) > 16: + spstr += "\n.." * 2 + return spstr + + def __repr__(self) -> str: + """ + Representation of a SparseMatrix + + Examples + -------- + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + >>> sm1 + SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2.0, 3.0, 4.0], False) + """ + rowIndices = list(self.rowIndices) + colPtrs = list(self.colPtrs) + + if len(self.values) <= 16: + values = _format_float_list(self.values) + + else: + values = ( + _format_float_list(self.values[:8]) + ["..."] + _format_float_list(self.values[-8:]) + ) + rowIndices = rowIndices[:8] + ["..."] + rowIndices[-8:] + + if len(self.colPtrs) > 16: + colPtrs = colPtrs[:8] + ["..."] + colPtrs[-8:] + + values = ", ".join(values) # type: ignore[assignment] + rowIndices = ", ".join([str(ind) for ind in rowIndices]) # type: ignore[assignment] + colPtrs = ", ".join([str(ptr) for ptr in colPtrs]) # type: ignore[assignment] + return "SparseMatrix({0}, {1}, [{2}], [{3}], [{4}], {5})".format( + self.numRows, self.numCols, colPtrs, rowIndices, values, self.isTransposed + ) + + def __reduce__(self) -> Tuple[Type["SparseMatrix"], Tuple[int, int, bytes, bytes, bytes, int]]: + return SparseMatrix, ( + self.numRows, + self.numCols, + self.colPtrs.tobytes(), + self.rowIndices.tobytes(), + self.values.tobytes(), + int(self.isTransposed), + ) + + def __getitem__(self, indices: Tuple[int, int]) -> np.float64: + i, j = indices + if i < 0 or i >= self.numRows: + raise IndexError("Row index %d is out of range [0, %d)" % (i, self.numRows)) + if j < 0 or j >= self.numCols: + raise IndexError("Column index %d is out of range [0, %d)" % (j, self.numCols)) + + # If a CSR matrix is given, then the row index should be searched + # for in ColPtrs, and the column index should be searched for in the + # corresponding slice obtained from rowIndices. + if self.isTransposed: + j, i = i, j + + colStart = self.colPtrs[j] + colEnd = self.colPtrs[j + 1] + nz = self.rowIndices[colStart:colEnd] + ind = np.searchsorted(nz, i) + colStart + if ind < colEnd and self.rowIndices[ind] == i: + return self.values[ind] + else: + return np.float64(0.0) + + def toArray(self) -> np.ndarray: + """ + Return a numpy.ndarray + """ + A = np.zeros((self.numRows, self.numCols), dtype=np.float64, order="F") + for k in range(self.colPtrs.size - 1): + startptr = self.colPtrs[k] + endptr = self.colPtrs[k + 1] + if self.isTransposed: + A[k, self.rowIndices[startptr:endptr]] = self.values[startptr:endptr] + else: + A[self.rowIndices[startptr:endptr], k] = self.values[startptr:endptr] + return A + + def toDense(self) -> "DenseMatrix": + densevals = np.ravel(self.toArray(), order="F") + return DenseMatrix(self.numRows, self.numCols, densevals) + + # TODO: More efficient implementation: + def __eq__(self, other: Any) -> bool: + return np.all(self.toArray() == other.toArray()).tolist() + + +class Matrices: + @staticmethod + def dense(numRows: int, numCols: int, values: Union[bytes, Iterable[float]]) -> DenseMatrix: + """ + Create a DenseMatrix + """ + return DenseMatrix(numRows, numCols, values) + + @staticmethod + def sparse( + numRows: int, + numCols: int, + colPtrs: Union[bytes, Iterable[int]], + rowIndices: Union[bytes, Iterable[int]], + values: Union[bytes, Iterable[float]], + ) -> SparseMatrix: + """ + Create a SparseMatrix + """ + return SparseMatrix(numRows, numCols, colPtrs, rowIndices, values) + + +def _test() -> None: + import doctest + + try: + # Numpy 1.14+ changed it's string format. + np.set_printoptions(legacy="1.13") + except TypeError: + pass + (failure_count, test_count) = doctest.testmod(optionflags=doctest.ELLIPSIS) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/model_cache.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/model_cache.py new file mode 100644 index 0000000..f70133f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/model_cache.py @@ -0,0 +1,55 @@ +# +# 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 collections import OrderedDict +from threading import Lock +from typing import Callable, Optional +from uuid import UUID + + +class ModelCache: + """Cache for model prediction functions on executors. + + This requires the `spark.python.worker.reuse` configuration to be set to `true`, otherwise a + new python worker (with an empty cache) will be started for every task. + + If a python worker is idle for more than one minute (per the IDLE_WORKER_TIMEOUT_NS setting in + PythonWorkerFactory.scala), it will be killed, effectively clearing the cache until a new python + worker is started. + + Caching large models can lead to out-of-memory conditions, which may require adjusting spark + memory configurations, e.g. `spark.executor.memoryOverhead`. + """ + + _models: OrderedDict = OrderedDict() + _capacity: int = 3 # "reasonable" default size for now, make configurable later, if needed + _lock: Lock = Lock() + + @staticmethod + def add(uuid: UUID, predict_fn: Callable) -> None: + with ModelCache._lock: + ModelCache._models[uuid] = predict_fn + ModelCache._models.move_to_end(uuid) + if len(ModelCache._models) > ModelCache._capacity: + ModelCache._models.popitem(last=False) + + @staticmethod + def get(uuid: UUID) -> Optional[Callable]: + with ModelCache._lock: + predict_fn = ModelCache._models.get(uuid) + if predict_fn: + ModelCache._models.move_to_end(uuid) + return predict_fn diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/__init__.py new file mode 100644 index 0000000..74f7b0b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/__init__.py @@ -0,0 +1,602 @@ +# +# 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 array +from abc import ABCMeta +import copy +from typing import ( + Any, + Callable, + Generic, + List, + Optional, + overload, + TypeVar, + Union, + TYPE_CHECKING, +) + +import numpy as np +from py4j.java_gateway import JavaObject + +from pyspark.ml.linalg import DenseVector, Vector, Matrix +from pyspark.ml.util import Identifiable + + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + +__all__ = ["Param", "Params", "TypeConverters"] + +T = TypeVar("T") +P = TypeVar("P", bound="Params") + + +class Param(Generic[T]): + """ + A param with self-contained documentation. + + .. versionadded:: 1.3.0 + """ + + def __init__( + self, + parent: Identifiable, + name: str, + doc: str, + typeConverter: Optional[Callable[[Any], T]] = None, + ): + if not isinstance(parent, Identifiable): + raise TypeError("Parent must be an Identifiable but got type %s." % type(parent)) + self.parent = parent.uid + self.name = str(name) + self.doc = str(doc) + self.typeConverter = TypeConverters.identity if typeConverter is None else typeConverter + + def _copy_new_parent(self, parent: Any) -> "Param": + """Copy the current param to a new parent, must be a dummy param.""" + if self.parent == "undefined": + param = copy.copy(self) + param.parent = parent.uid + return param + else: + raise ValueError("Cannot copy from non-dummy parent %s." % parent) + + def __str__(self) -> str: + return str(self.parent) + "__" + self.name + + def __repr__(self) -> str: + return "Param(parent=%r, name=%r, doc=%r)" % (self.parent, self.name, self.doc) + + def __hash__(self) -> int: + return hash(str(self)) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, Param): + return self.parent == other.parent and self.name == other.name + else: + return False + + +class TypeConverters: + """ + Factory methods for common type conversion functions for `Param.typeConverter`. + + .. versionadded:: 2.0.0 + """ + + @staticmethod + def _is_numeric(value: Any) -> bool: + vtype = type(value) + return vtype in [int, float, np.float64, np.int64] or vtype.__name__ == "long" + + @staticmethod + def _is_integer(value: Any) -> bool: + return TypeConverters._is_numeric(value) and float(value).is_integer() + + @staticmethod + def _can_convert_to_list(value: Any) -> bool: + vtype = type(value) + return vtype in [list, np.ndarray, tuple, range, array.array] or isinstance(value, Vector) + + @staticmethod + def _can_convert_to_string(value: Any) -> bool: + vtype = type(value) + return isinstance(value, str) or vtype in [np.unicode_, np.string_, np.str_] + + @staticmethod + def identity(value: "T") -> "T": + """ + Dummy converter that just returns value. + """ + return value + + @staticmethod + def toList(value: Any) -> List: + """ + Convert a value to a list, if possible. + """ + if type(value) == list: + return value + elif type(value) in [np.ndarray, tuple, range, array.array]: + return list(value) + elif isinstance(value, Vector): + return list(value.toArray()) + else: + raise TypeError("Could not convert %s to list" % value) + + @staticmethod + def toListFloat(value: Any) -> List[float]: + """ + Convert a value to list of floats, if possible. + """ + if TypeConverters._can_convert_to_list(value): + value = TypeConverters.toList(value) + if all(map(lambda v: TypeConverters._is_numeric(v), value)): + return [float(v) for v in value] + raise TypeError("Could not convert %s to list of floats" % value) + + @staticmethod + def toListListFloat(value: Any) -> List[List[float]]: + """ + Convert a value to list of list of floats, if possible. + """ + if TypeConverters._can_convert_to_list(value): + value = TypeConverters.toList(value) + return [TypeConverters.toListFloat(v) for v in value] + raise TypeError("Could not convert %s to list of list of floats" % value) + + @staticmethod + def toListInt(value: Any) -> List[int]: + """ + Convert a value to list of ints, if possible. + """ + if TypeConverters._can_convert_to_list(value): + value = TypeConverters.toList(value) + if all(map(lambda v: TypeConverters._is_integer(v), value)): + return [int(v) for v in value] + raise TypeError("Could not convert %s to list of ints" % value) + + @staticmethod + def toListString(value: Any) -> List[str]: + """ + Convert a value to list of strings, if possible. + """ + if TypeConverters._can_convert_to_list(value): + value = TypeConverters.toList(value) + if all(map(lambda v: TypeConverters._can_convert_to_string(v), value)): + return [TypeConverters.toString(v) for v in value] + raise TypeError("Could not convert %s to list of strings" % value) + + @staticmethod + def toVector(value: Any) -> Vector: + """ + Convert a value to a MLlib Vector, if possible. + """ + if isinstance(value, Vector): + return value + elif TypeConverters._can_convert_to_list(value): + value = TypeConverters.toList(value) + if all(map(lambda v: TypeConverters._is_numeric(v), value)): + return DenseVector(value) + raise TypeError("Could not convert %s to vector" % value) + + @staticmethod + def toMatrix(value: Any) -> Matrix: + """ + Convert a value to a MLlib Matrix, if possible. + """ + if isinstance(value, Matrix): + return value + raise TypeError("Could not convert %s to matrix" % value) + + @staticmethod + def toFloat(value: Any) -> float: + """ + Convert a value to a float, if possible. + """ + if TypeConverters._is_numeric(value): + return float(value) + else: + raise TypeError("Could not convert %s to float" % value) + + @staticmethod + def toInt(value: Any) -> int: + """ + Convert a value to an int, if possible. + """ + if TypeConverters._is_integer(value): + return int(value) + else: + raise TypeError("Could not convert %s to int" % value) + + @staticmethod + def toString(value: Any) -> str: + """ + Convert a value to a string, if possible. + """ + if isinstance(value, str): + return value + elif type(value) in [np.string_, np.str_, np.unicode_]: + return str(value) + else: + raise TypeError("Could not convert %s to string type" % type(value)) + + @staticmethod + def toBoolean(value: Any) -> bool: + """ + Convert a value to a boolean, if possible. + """ + if type(value) == bool: + return value + else: + raise TypeError("Boolean Param requires value of type bool. Found %s." % type(value)) + + +class Params(Identifiable, metaclass=ABCMeta): + """ + Components that take parameters. This also provides an internal + param map to store parameter values attached to the instance. + + .. versionadded:: 1.3.0 + """ + + def __init__(self) -> None: + super(Params, self).__init__() + #: internal param map for user-supplied values param map + self._paramMap: "ParamMap" = {} + + #: internal param map for default values + self._defaultParamMap: "ParamMap" = {} + + #: value returned by :py:func:`params` + self._params: Optional[List[Param]] = None + + # Copy the params from the class to the object + self._copy_params() + + def _copy_params(self) -> None: + """ + Copy all params defined on the class to current object. + """ + cls = type(self) + src_name_attrs = [(x, getattr(cls, x)) for x in dir(cls)] + src_params = list(filter(lambda nameAttr: isinstance(nameAttr[1], Param), src_name_attrs)) + for name, param in src_params: + setattr(self, name, param._copy_new_parent(self)) + + @property + def params(self) -> List[Param]: + """ + Returns all params ordered by name. The default implementation + uses :py:func:`dir` to get all attributes of type + :py:class:`Param`. + """ + if self._params is None: + self._params = list( + filter( + lambda attr: isinstance(attr, Param), + [ + getattr(self, x) + for x in dir(self) + if x != "params" and not isinstance(getattr(type(self), x, None), property) + ], + ) + ) + return self._params + + def explainParam(self, param: Union[str, Param]) -> str: + """ + Explains a single param and returns its name, doc, and optional + default value and user-supplied value in a string. + """ + param = self._resolveParam(param) + values = [] + if self.isDefined(param): + if param in self._defaultParamMap: + values.append("default: %s" % self._defaultParamMap[param]) + if param in self._paramMap: + values.append("current: %s" % self._paramMap[param]) + else: + values.append("undefined") + valueStr = "(" + ", ".join(values) + ")" + return "%s: %s %s" % (param.name, param.doc, valueStr) + + def explainParams(self) -> str: + """ + Returns the documentation of all params with their optionally + default values and user-supplied values. + """ + return "\n".join([self.explainParam(param) for param in self.params]) + + def getParam(self, paramName: str) -> Param: + """ + Gets a param by its name. + """ + param = getattr(self, paramName) + if isinstance(param, Param): + return param + else: + raise ValueError("Cannot find param with name %s." % paramName) + + def isSet(self, param: Union[str, Param[Any]]) -> bool: + """ + Checks whether a param is explicitly set by user. + """ + param = self._resolveParam(param) + return param in self._paramMap + + def hasDefault(self, param: Union[str, Param[Any]]) -> bool: + """ + Checks whether a param has a default value. + """ + param = self._resolveParam(param) + return param in self._defaultParamMap + + def isDefined(self, param: Union[str, Param[Any]]) -> bool: + """ + Checks whether a param is explicitly set by user or has + a default value. + """ + return self.isSet(param) or self.hasDefault(param) + + def hasParam(self, paramName: str) -> bool: + """ + Tests whether this instance contains a param with a given + (string) name. + """ + if isinstance(paramName, str): + p = getattr(self, paramName, None) + return isinstance(p, Param) + else: + raise TypeError("hasParam(): paramName must be a string") + + @overload + def getOrDefault(self, param: str) -> Any: + ... + + @overload + def getOrDefault(self, param: Param[T]) -> T: + ... + + def getOrDefault(self, param: Union[str, Param[T]]) -> Union[Any, T]: + + """ + Gets the value of a param in the user-supplied param map or its + default value. Raises an error if neither is set. + """ + param = self._resolveParam(param) + if param in self._paramMap: + return self._paramMap[param] + else: + return self._defaultParamMap[param] + + def extractParamMap(self, extra: Optional["ParamMap"] = None) -> "ParamMap": + """ + Extracts the embedded default param values and user-supplied + values, and then merges them with extra values from input into + a flat param map, where the latter value is used if there exist + conflicts, i.e., with ordering: default param values < + user-supplied values < extra. + + Parameters + ---------- + extra : dict, optional + extra param values + + Returns + ------- + dict + merged param map + """ + if extra is None: + extra = dict() + paramMap = self._defaultParamMap.copy() + paramMap.update(self._paramMap) + paramMap.update(extra) + return paramMap + + def copy(self: P, extra: Optional["ParamMap"] = None) -> P: + """ + Creates a copy of this instance with the same uid and some + extra params. The default implementation creates a + shallow copy using :py:func:`copy.copy`, and then copies the + embedded and extra parameters over and returns the copy. + Subclasses should override this method if the default approach + is not sufficient. + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`Params` + Copy of this instance + """ + if extra is None: + extra = dict() + that = copy.copy(self) + that._paramMap = {} + that._defaultParamMap = {} + return self._copyValues(that, extra) + + def set(self, param: Param, value: Any) -> None: + """ + Sets a parameter in the embedded param map. + """ + self._shouldOwn(param) + try: + value = param.typeConverter(value) + except ValueError as e: + raise ValueError('Invalid param value given for param "%s". %s' % (param.name, e)) + self._paramMap[param] = value + + def _shouldOwn(self, param: Param) -> None: + """ + Validates that the input param belongs to this Params instance. + """ + if not (self.uid == param.parent and self.hasParam(param.name)): + raise ValueError("Param %r does not belong to %r." % (param, self)) + + def _resolveParam(self, param: Union[str, Param]) -> Param: + """ + Resolves a param and validates the ownership. + + Parameters + ---------- + param : str or :py:class:`Param` + param name or the param instance, which must + belong to this Params instance + + Returns + ------- + :py:class:`Param` + resolved param instance + """ + if isinstance(param, Param): + self._shouldOwn(param) + return param + elif isinstance(param, str): + return self.getParam(param) + else: + raise TypeError("Cannot resolve %r as a param." % param) + + def _testOwnParam(self, param_parent: str, param_name: str) -> bool: + """ + Test the ownership. Return True or False + """ + return self.uid == param_parent and self.hasParam(param_name) + + @staticmethod + def _dummy() -> "Params": + """ + Returns a dummy Params instance used as a placeholder to + generate docs. + """ + dummy = Params() + dummy.uid = "undefined" + return dummy + + def _set(self: P, **kwargs: Any) -> P: + """ + Sets user-supplied params. + """ + for param, value in kwargs.items(): + p = getattr(self, param) + if value is not None: + try: + value = p.typeConverter(value) + except TypeError as e: + raise TypeError('Invalid param value given for param "%s". %s' % (p.name, e)) + self._paramMap[p] = value + return self + + def clear(self, param: Param) -> None: + """ + Clears a param from the param map if it has been explicitly set. + """ + if self.isSet(param): + del self._paramMap[param] + + def _setDefault(self: P, **kwargs: Any) -> P: + """ + Sets default params. + """ + for param, value in kwargs.items(): + p = getattr(self, param) + if value is not None and not isinstance(value, JavaObject): + try: + value = p.typeConverter(value) + except TypeError as e: + raise TypeError( + 'Invalid default param value given for param "%s". %s' % (p.name, e) + ) + self._defaultParamMap[p] = value + return self + + def _copyValues(self, to: P, extra: Optional["ParamMap"] = None) -> P: + """ + Copies param values from this instance to another instance for + params shared by them. + + Parameters + ---------- + to : :py:class:`Params` + the target instance + extra : dict, optional + extra params to be copied + + Returns + ------- + :py:class:`Params` + the target instance with param values copied + """ + paramMap = self._paramMap.copy() + if isinstance(extra, dict): + for param, value in extra.items(): + if isinstance(param, Param): + paramMap[param] = value + else: + raise TypeError( + "Expecting a valid instance of Param, but received: {}".format(param) + ) + elif extra is not None: + raise TypeError( + "Expecting a dict, but received an object of type {}.".format(type(extra)) + ) + for param in self.params: + # copy default params + if param in self._defaultParamMap and to.hasParam(param.name): + to._defaultParamMap[to.getParam(param.name)] = self._defaultParamMap[param] + # copy explicitly set params + if param in paramMap and to.hasParam(param.name): + to._set(**{param.name: paramMap[param]}) + return to + + def _resetUid(self: P, newUid: Any) -> P: + """ + Changes the uid of this instance. This updates both + the stored uid and the parent uid of params and param maps. + This is used by persistence (loading). + + Parameters + ---------- + newUid + new uid to use, which is converted to unicode + + Returns + ------- + :py:class:`Params` + same instance, but with the uid and Param.parent values + updated, including within param maps + """ + newUid = str(newUid) + self.uid = newUid + newDefaultParamMap = dict() + newParamMap = dict() + for param in self.params: + newParam = copy.copy(param) + newParam.parent = newUid + if param in self._defaultParamMap: + newDefaultParamMap[newParam] = self._defaultParamMap[param] + if param in self._paramMap: + newParamMap[newParam] = self._paramMap[param] + param.parent = newUid + self._defaultParamMap = newDefaultParamMap + self._paramMap = newParamMap + return self diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/_shared_params_code_gen.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/_shared_params_code_gen.py new file mode 100644 index 0000000..2bec3a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/_shared_params_code_gen.py @@ -0,0 +1,368 @@ +# +# 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 Optional + +header = """# +# 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. +#""" + +# Code generator for shared params (shared.py). Run under this folder with: +# python _shared_params_code_gen.py > shared.py + +_type_for_type_converter = { + "TypeConverters.toBoolean": "bool", + "TypeConverters.toFloat": "float", + "TypeConverters.toInt": "int", + "TypeConverters.toListFloat": "List[float]", + "TypeConverters.toListString": "List[str]", + "TypeConverters.toString": "str", +} + + +def _gen_param_header( + name: str, doc: str, defaultValueStr: Optional[str], typeConverter: str, paramType: str +) -> str: + """ + Generates the header part for shared variables + + :param name: param name + :param doc: param doc + """ + Name = f"Has{name[0].upper()}{name[1:]}" + + template = f'''class {Name}(Params): + """ + Mixin for param {name}: {doc} + """ + + {name}: "Param[{paramType}]" = Param( + Params._dummy(), + "{name}", + "{doc}", + typeConverter={typeConverter}, + ) + + def __init__(self) -> None: + super({Name}, self).__init__()''' + + if defaultValueStr is not None: + template += f""" + self._setDefault({name}={defaultValueStr})""" + + return template + + +def _gen_param_code(name: str, paramType: str) -> str: + """ + Generates Python code for a shared param class. + + :param name: param name + :param doc: param doc + :param defaultValueStr: string representation of the default value + :return: code string + """ + # TODO: How to correctly inherit instance attributes? + return f''' + def get{name[0].upper()}{name[1:]}(self) -> {paramType}: + """ + Gets the value of {name} or its default value. + """ + return self.getOrDefault(self.{name})''' + + +if __name__ == "__main__": + print(header) + print("\n# DO NOT MODIFY THIS FILE! It was generated by _shared_params_code_gen.py.\n") + print("from typing import List\n") + print("from pyspark.ml.param import Param, Params, TypeConverters\n\n") + shared = [ + ( + "maxIter", + "max number of iterations (>= 0).", + None, + "TypeConverters.toInt", + ), + ( + "regParam", + "regularization parameter (>= 0).", + None, + "TypeConverters.toFloat", + ), + ( + "featuresCol", + "features column name.", + '"features"', + "TypeConverters.toString", + ), + ( + "labelCol", + "label column name.", + '"label"', + "TypeConverters.toString", + ), + ( + "predictionCol", + "prediction column name.", + '"prediction"', + "TypeConverters.toString", + ), + ( + "probabilityCol", + "Column name for predicted class conditional probabilities. " + + "Note: Not all models output well-calibrated probability estimates! " + + "These probabilities should be treated as confidences, not precise probabilities.", + '"probability"', + "TypeConverters.toString", + ), + ( + "rawPredictionCol", + "raw prediction (a.k.a. confidence) column name.", + '"rawPrediction"', + "TypeConverters.toString", + ), + ( + "inputCol", + "input column name.", + None, + "TypeConverters.toString", + ), + ( + "inputCols", + "input column names.", + None, + "TypeConverters.toListString", + ), + ( + "outputCol", + "output column name.", + 'self.uid + "__output"', + "TypeConverters.toString", + ), + ( + "outputCols", + "output column names.", + None, + "TypeConverters.toListString", + ), + ( + "numFeatures", + "Number of features. Should be greater than 0.", + "262144", + "TypeConverters.toInt", + ), + ( + "checkpointInterval", + "set checkpoint interval (>= 1) or disable checkpoint (-1). " + + "E.g. 10 means that the cache will get checkpointed every 10 iterations. Note: " + + "this setting will be ignored if the checkpoint directory is not set in " + + "the SparkContext.", + None, + "TypeConverters.toInt", + ), + ( + "seed", + "random seed.", + "hash(type(self).__name__)", + "TypeConverters.toInt", + ), + ( + "tol", + "the convergence tolerance for iterative algorithms (>= 0).", + None, + "TypeConverters.toFloat", + ), + ( + "relativeError", + "the relative target precision for the approximate quantile " + + "algorithm. Must be in the range [0, 1]", + "0.001", + "TypeConverters.toFloat", + ), + ( + "stepSize", + "Step size to be used for each iteration of optimization (>= 0).", + None, + "TypeConverters.toFloat", + ), + ( + "handleInvalid", + "how to handle invalid entries. Options are skip (which will filter " + + "out rows with bad values), or error (which will throw an error). " + + "More options may be added later.", + None, + "TypeConverters.toString", + ), + ( + "elasticNetParam", + "the ElasticNet mixing parameter, in range [0, 1]. For alpha = 0, " + + "the penalty is an L2 penalty. For alpha = 1, it is an L1 penalty.", + "0.0", + "TypeConverters.toFloat", + ), + ( + "fitIntercept", + "whether to fit an intercept term.", + "True", + "TypeConverters.toBoolean", + ), + ( + "standardization", + "whether to standardize the training features before fitting the " + "model.", + "True", + "TypeConverters.toBoolean", + ), + ( + "thresholds", + "Thresholds in multi-class classification to adjust the probability of " + + "predicting each class. Array must have length equal to the number of classes, with " + + "values > 0, excepting that at most one value may be 0. " + + "The class with largest value p/t is predicted, where p is the original " + + "probability of that class and t is the class's threshold.", + None, + "TypeConverters.toListFloat", + ), + ( + "threshold", + "threshold in binary classification prediction, in range [0, 1]", + "0.5", + "TypeConverters.toFloat", + ), + ( + "weightCol", + "weight column name. If this is not set or empty, we treat " + + "all instance weights as 1.0.", + None, + "TypeConverters.toString", + ), + ( + "solver", + "the solver algorithm for optimization. If this is not set or empty, " + + "default value is 'auto'.", + '"auto"', + "TypeConverters.toString", + ), + ( + "varianceCol", + "column name for the biased sample variance of prediction.", + None, + "TypeConverters.toString", + ), + ( + "aggregationDepth", + "suggested depth for treeAggregate (>= 2).", + "2", + "TypeConverters.toInt", + ), + ( + "parallelism", + "the number of threads to use when running parallel algorithms (>= 1).", + "1", + "TypeConverters.toInt", + ), + ( + "collectSubModels", + "Param for whether to collect a list of sub-models trained during " + + "tuning. If set to false, then only the single best sub-model will be available " + + "after fitting. If set to true, then all sub-models will be available. Warning: " + + "For large models, collecting all sub-models can cause OOMs on the Spark driver.", + "False", + "TypeConverters.toBoolean", + ), + ( + "loss", + "the loss function to be optimized.", + None, + "TypeConverters.toString", + ), + ( + "distanceMeasure", + "the distance measure. Supported options: 'euclidean' and 'cosine'.", + '"euclidean"', + "TypeConverters.toString", + ), + ( + "validationIndicatorCol", + "name of the column that indicates whether each row is for " + + "training or for validation. False indicates training; true indicates validation.", + None, + "TypeConverters.toString", + ), + ( + "blockSize", + "block size for stacking input data in matrices. Data is stacked within " + "partitions. If block size is more than remaining data in a partition then it is " + "adjusted to the size of this data.", + None, + "TypeConverters.toInt", + ), + ( + "maxBlockSizeInMB", + "maximum memory in MB for stacking input data into blocks. Data is " + + "stacked within partitions. If more than remaining data size in a partition then it " + + "is adjusted to the data size. Default 0.0 represents choosing optimal value, " + + "depends on specific algorithm. Must be >= 0.", + "0.0", + "TypeConverters.toFloat", + ), + ( + "numTrainWorkers", + "number of training workers", + "1", + "TypeConverters.toInt", + ), + ( + "batchSize", + "number of training batch size", + None, + "TypeConverters.toInt", + ), + ( + "learningRate", + "learning rate for training", + None, + "TypeConverters.toFloat", + ), + ( + "momentum", + "momentum for training optimizer", + None, + "TypeConverters.toFloat", + ), + ] + + code = [] + for name, doc, defaultValueStr, typeConverter in shared: + paramType = _type_for_type_converter.get(typeConverter, "None") + + param_code = _gen_param_header(name, doc, defaultValueStr, typeConverter, paramType) + code.append(param_code + "\n" + _gen_param_code(name, paramType)) + + print("\n\n\n".join(code)) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/shared.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/shared.py new file mode 100644 index 0000000..d61d206 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/param/shared.py @@ -0,0 +1,878 @@ +# +# 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. +# + +# DO NOT MODIFY THIS FILE! It was generated by _shared_params_code_gen.py. + +from typing import List + +from pyspark.ml.param import Param, Params, TypeConverters + + +class HasMaxIter(Params): + """ + Mixin for param maxIter: max number of iterations (>= 0). + """ + + maxIter: "Param[int]" = Param( + Params._dummy(), + "maxIter", + "max number of iterations (>= 0).", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasMaxIter, self).__init__() + + def getMaxIter(self) -> int: + """ + Gets the value of maxIter or its default value. + """ + return self.getOrDefault(self.maxIter) + + +class HasRegParam(Params): + """ + Mixin for param regParam: regularization parameter (>= 0). + """ + + regParam: "Param[float]" = Param( + Params._dummy(), + "regParam", + "regularization parameter (>= 0).", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasRegParam, self).__init__() + + def getRegParam(self) -> float: + """ + Gets the value of regParam or its default value. + """ + return self.getOrDefault(self.regParam) + + +class HasFeaturesCol(Params): + """ + Mixin for param featuresCol: features column name. + """ + + featuresCol: "Param[str]" = Param( + Params._dummy(), + "featuresCol", + "features column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasFeaturesCol, self).__init__() + self._setDefault(featuresCol="features") + + def getFeaturesCol(self) -> str: + """ + Gets the value of featuresCol or its default value. + """ + return self.getOrDefault(self.featuresCol) + + +class HasLabelCol(Params): + """ + Mixin for param labelCol: label column name. + """ + + labelCol: "Param[str]" = Param( + Params._dummy(), + "labelCol", + "label column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasLabelCol, self).__init__() + self._setDefault(labelCol="label") + + def getLabelCol(self) -> str: + """ + Gets the value of labelCol or its default value. + """ + return self.getOrDefault(self.labelCol) + + +class HasPredictionCol(Params): + """ + Mixin for param predictionCol: prediction column name. + """ + + predictionCol: "Param[str]" = Param( + Params._dummy(), + "predictionCol", + "prediction column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasPredictionCol, self).__init__() + self._setDefault(predictionCol="prediction") + + def getPredictionCol(self) -> str: + """ + Gets the value of predictionCol or its default value. + """ + return self.getOrDefault(self.predictionCol) + + +class HasProbabilityCol(Params): + """ + Mixin for param probabilityCol: Column name for predicted class conditional probabilities. Note: Not all models output well-calibrated probability estimates! These probabilities should be treated as confidences, not precise probabilities. + """ + + probabilityCol: "Param[str]" = Param( + Params._dummy(), + "probabilityCol", + "Column name for predicted class conditional probabilities. Note: Not all models output well-calibrated probability estimates! These probabilities should be treated as confidences, not precise probabilities.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasProbabilityCol, self).__init__() + self._setDefault(probabilityCol="probability") + + def getProbabilityCol(self) -> str: + """ + Gets the value of probabilityCol or its default value. + """ + return self.getOrDefault(self.probabilityCol) + + +class HasRawPredictionCol(Params): + """ + Mixin for param rawPredictionCol: raw prediction (a.k.a. confidence) column name. + """ + + rawPredictionCol: "Param[str]" = Param( + Params._dummy(), + "rawPredictionCol", + "raw prediction (a.k.a. confidence) column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasRawPredictionCol, self).__init__() + self._setDefault(rawPredictionCol="rawPrediction") + + def getRawPredictionCol(self) -> str: + """ + Gets the value of rawPredictionCol or its default value. + """ + return self.getOrDefault(self.rawPredictionCol) + + +class HasInputCol(Params): + """ + Mixin for param inputCol: input column name. + """ + + inputCol: "Param[str]" = Param( + Params._dummy(), + "inputCol", + "input column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasInputCol, self).__init__() + + def getInputCol(self) -> str: + """ + Gets the value of inputCol or its default value. + """ + return self.getOrDefault(self.inputCol) + + +class HasInputCols(Params): + """ + Mixin for param inputCols: input column names. + """ + + inputCols: "Param[List[str]]" = Param( + Params._dummy(), + "inputCols", + "input column names.", + typeConverter=TypeConverters.toListString, + ) + + def __init__(self) -> None: + super(HasInputCols, self).__init__() + + def getInputCols(self) -> List[str]: + """ + Gets the value of inputCols or its default value. + """ + return self.getOrDefault(self.inputCols) + + +class HasOutputCol(Params): + """ + Mixin for param outputCol: output column name. + """ + + outputCol: "Param[str]" = Param( + Params._dummy(), + "outputCol", + "output column name.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasOutputCol, self).__init__() + self._setDefault(outputCol=self.uid + "__output") + + def getOutputCol(self) -> str: + """ + Gets the value of outputCol or its default value. + """ + return self.getOrDefault(self.outputCol) + + +class HasOutputCols(Params): + """ + Mixin for param outputCols: output column names. + """ + + outputCols: "Param[List[str]]" = Param( + Params._dummy(), + "outputCols", + "output column names.", + typeConverter=TypeConverters.toListString, + ) + + def __init__(self) -> None: + super(HasOutputCols, self).__init__() + + def getOutputCols(self) -> List[str]: + """ + Gets the value of outputCols or its default value. + """ + return self.getOrDefault(self.outputCols) + + +class HasNumFeatures(Params): + """ + Mixin for param numFeatures: Number of features. Should be greater than 0. + """ + + numFeatures: "Param[int]" = Param( + Params._dummy(), + "numFeatures", + "Number of features. Should be greater than 0.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasNumFeatures, self).__init__() + self._setDefault(numFeatures=262144) + + def getNumFeatures(self) -> int: + """ + Gets the value of numFeatures or its default value. + """ + return self.getOrDefault(self.numFeatures) + + +class HasCheckpointInterval(Params): + """ + Mixin for param checkpointInterval: set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations. Note: this setting will be ignored if the checkpoint directory is not set in the SparkContext. + """ + + checkpointInterval: "Param[int]" = Param( + Params._dummy(), + "checkpointInterval", + "set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations. Note: this setting will be ignored if the checkpoint directory is not set in the SparkContext.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasCheckpointInterval, self).__init__() + + def getCheckpointInterval(self) -> int: + """ + Gets the value of checkpointInterval or its default value. + """ + return self.getOrDefault(self.checkpointInterval) + + +class HasSeed(Params): + """ + Mixin for param seed: random seed. + """ + + seed: "Param[int]" = Param( + Params._dummy(), + "seed", + "random seed.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasSeed, self).__init__() + self._setDefault(seed=hash(type(self).__name__)) + + def getSeed(self) -> int: + """ + Gets the value of seed or its default value. + """ + return self.getOrDefault(self.seed) + + +class HasTol(Params): + """ + Mixin for param tol: the convergence tolerance for iterative algorithms (>= 0). + """ + + tol: "Param[float]" = Param( + Params._dummy(), + "tol", + "the convergence tolerance for iterative algorithms (>= 0).", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasTol, self).__init__() + + def getTol(self) -> float: + """ + Gets the value of tol or its default value. + """ + return self.getOrDefault(self.tol) + + +class HasRelativeError(Params): + """ + Mixin for param relativeError: the relative target precision for the approximate quantile algorithm. Must be in the range [0, 1] + """ + + relativeError: "Param[float]" = Param( + Params._dummy(), + "relativeError", + "the relative target precision for the approximate quantile algorithm. Must be in the range [0, 1]", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasRelativeError, self).__init__() + self._setDefault(relativeError=0.001) + + def getRelativeError(self) -> float: + """ + Gets the value of relativeError or its default value. + """ + return self.getOrDefault(self.relativeError) + + +class HasStepSize(Params): + """ + Mixin for param stepSize: Step size to be used for each iteration of optimization (>= 0). + """ + + stepSize: "Param[float]" = Param( + Params._dummy(), + "stepSize", + "Step size to be used for each iteration of optimization (>= 0).", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasStepSize, self).__init__() + + def getStepSize(self) -> float: + """ + Gets the value of stepSize or its default value. + """ + return self.getOrDefault(self.stepSize) + + +class HasHandleInvalid(Params): + """ + Mixin for param handleInvalid: how to handle invalid entries. Options are skip (which will filter out rows with bad values), or error (which will throw an error). More options may be added later. + """ + + handleInvalid: "Param[str]" = Param( + Params._dummy(), + "handleInvalid", + "how to handle invalid entries. Options are skip (which will filter out rows with bad values), or error (which will throw an error). More options may be added later.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasHandleInvalid, self).__init__() + + def getHandleInvalid(self) -> str: + """ + Gets the value of handleInvalid or its default value. + """ + return self.getOrDefault(self.handleInvalid) + + +class HasElasticNetParam(Params): + """ + Mixin for param elasticNetParam: the ElasticNet mixing parameter, in range [0, 1]. For alpha = 0, the penalty is an L2 penalty. For alpha = 1, it is an L1 penalty. + """ + + elasticNetParam: "Param[float]" = Param( + Params._dummy(), + "elasticNetParam", + "the ElasticNet mixing parameter, in range [0, 1]. For alpha = 0, the penalty is an L2 penalty. For alpha = 1, it is an L1 penalty.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasElasticNetParam, self).__init__() + self._setDefault(elasticNetParam=0.0) + + def getElasticNetParam(self) -> float: + """ + Gets the value of elasticNetParam or its default value. + """ + return self.getOrDefault(self.elasticNetParam) + + +class HasFitIntercept(Params): + """ + Mixin for param fitIntercept: whether to fit an intercept term. + """ + + fitIntercept: "Param[bool]" = Param( + Params._dummy(), + "fitIntercept", + "whether to fit an intercept term.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super(HasFitIntercept, self).__init__() + self._setDefault(fitIntercept=True) + + def getFitIntercept(self) -> bool: + """ + Gets the value of fitIntercept or its default value. + """ + return self.getOrDefault(self.fitIntercept) + + +class HasStandardization(Params): + """ + Mixin for param standardization: whether to standardize the training features before fitting the model. + """ + + standardization: "Param[bool]" = Param( + Params._dummy(), + "standardization", + "whether to standardize the training features before fitting the model.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super(HasStandardization, self).__init__() + self._setDefault(standardization=True) + + def getStandardization(self) -> bool: + """ + Gets the value of standardization or its default value. + """ + return self.getOrDefault(self.standardization) + + +class HasThresholds(Params): + """ + Mixin for param thresholds: Thresholds in multi-class classification to adjust the probability of predicting each class. Array must have length equal to the number of classes, with values > 0, excepting that at most one value may be 0. The class with largest value p/t is predicted, where p is the original probability of that class and t is the class's threshold. + """ + + thresholds: "Param[List[float]]" = Param( + Params._dummy(), + "thresholds", + "Thresholds in multi-class classification to adjust the probability of predicting each class. Array must have length equal to the number of classes, with values > 0, excepting that at most one value may be 0. The class with largest value p/t is predicted, where p is the original probability of that class and t is the class's threshold.", + typeConverter=TypeConverters.toListFloat, + ) + + def __init__(self) -> None: + super(HasThresholds, self).__init__() + + def getThresholds(self) -> List[float]: + """ + Gets the value of thresholds or its default value. + """ + return self.getOrDefault(self.thresholds) + + +class HasThreshold(Params): + """ + Mixin for param threshold: threshold in binary classification prediction, in range [0, 1] + """ + + threshold: "Param[float]" = Param( + Params._dummy(), + "threshold", + "threshold in binary classification prediction, in range [0, 1]", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasThreshold, self).__init__() + self._setDefault(threshold=0.5) + + def getThreshold(self) -> float: + """ + Gets the value of threshold or its default value. + """ + return self.getOrDefault(self.threshold) + + +class HasWeightCol(Params): + """ + Mixin for param weightCol: weight column name. If this is not set or empty, we treat all instance weights as 1.0. + """ + + weightCol: "Param[str]" = Param( + Params._dummy(), + "weightCol", + "weight column name. If this is not set or empty, we treat all instance weights as 1.0.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasWeightCol, self).__init__() + + def getWeightCol(self) -> str: + """ + Gets the value of weightCol or its default value. + """ + return self.getOrDefault(self.weightCol) + + +class HasSolver(Params): + """ + Mixin for param solver: the solver algorithm for optimization. If this is not set or empty, default value is 'auto'. + """ + + solver: "Param[str]" = Param( + Params._dummy(), + "solver", + "the solver algorithm for optimization. If this is not set or empty, default value is 'auto'.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasSolver, self).__init__() + self._setDefault(solver="auto") + + def getSolver(self) -> str: + """ + Gets the value of solver or its default value. + """ + return self.getOrDefault(self.solver) + + +class HasVarianceCol(Params): + """ + Mixin for param varianceCol: column name for the biased sample variance of prediction. + """ + + varianceCol: "Param[str]" = Param( + Params._dummy(), + "varianceCol", + "column name for the biased sample variance of prediction.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasVarianceCol, self).__init__() + + def getVarianceCol(self) -> str: + """ + Gets the value of varianceCol or its default value. + """ + return self.getOrDefault(self.varianceCol) + + +class HasAggregationDepth(Params): + """ + Mixin for param aggregationDepth: suggested depth for treeAggregate (>= 2). + """ + + aggregationDepth: "Param[int]" = Param( + Params._dummy(), + "aggregationDepth", + "suggested depth for treeAggregate (>= 2).", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasAggregationDepth, self).__init__() + self._setDefault(aggregationDepth=2) + + def getAggregationDepth(self) -> int: + """ + Gets the value of aggregationDepth or its default value. + """ + return self.getOrDefault(self.aggregationDepth) + + +class HasParallelism(Params): + """ + Mixin for param parallelism: the number of threads to use when running parallel algorithms (>= 1). + """ + + parallelism: "Param[int]" = Param( + Params._dummy(), + "parallelism", + "the number of threads to use when running parallel algorithms (>= 1).", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasParallelism, self).__init__() + self._setDefault(parallelism=1) + + def getParallelism(self) -> int: + """ + Gets the value of parallelism or its default value. + """ + return self.getOrDefault(self.parallelism) + + +class HasCollectSubModels(Params): + """ + Mixin for param collectSubModels: Param for whether to collect a list of sub-models trained during tuning. If set to false, then only the single best sub-model will be available after fitting. If set to true, then all sub-models will be available. Warning: For large models, collecting all sub-models can cause OOMs on the Spark driver. + """ + + collectSubModels: "Param[bool]" = Param( + Params._dummy(), + "collectSubModels", + "Param for whether to collect a list of sub-models trained during tuning. If set to false, then only the single best sub-model will be available after fitting. If set to true, then all sub-models will be available. Warning: For large models, collecting all sub-models can cause OOMs on the Spark driver.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super(HasCollectSubModels, self).__init__() + self._setDefault(collectSubModels=False) + + def getCollectSubModels(self) -> bool: + """ + Gets the value of collectSubModels or its default value. + """ + return self.getOrDefault(self.collectSubModels) + + +class HasLoss(Params): + """ + Mixin for param loss: the loss function to be optimized. + """ + + loss: "Param[str]" = Param( + Params._dummy(), + "loss", + "the loss function to be optimized.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasLoss, self).__init__() + + def getLoss(self) -> str: + """ + Gets the value of loss or its default value. + """ + return self.getOrDefault(self.loss) + + +class HasDistanceMeasure(Params): + """ + Mixin for param distanceMeasure: the distance measure. Supported options: 'euclidean' and 'cosine'. + """ + + distanceMeasure: "Param[str]" = Param( + Params._dummy(), + "distanceMeasure", + "the distance measure. Supported options: 'euclidean' and 'cosine'.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasDistanceMeasure, self).__init__() + self._setDefault(distanceMeasure="euclidean") + + def getDistanceMeasure(self) -> str: + """ + Gets the value of distanceMeasure or its default value. + """ + return self.getOrDefault(self.distanceMeasure) + + +class HasValidationIndicatorCol(Params): + """ + Mixin for param validationIndicatorCol: name of the column that indicates whether each row is for training or for validation. False indicates training; true indicates validation. + """ + + validationIndicatorCol: "Param[str]" = Param( + Params._dummy(), + "validationIndicatorCol", + "name of the column that indicates whether each row is for training or for validation. False indicates training; true indicates validation.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(HasValidationIndicatorCol, self).__init__() + + def getValidationIndicatorCol(self) -> str: + """ + Gets the value of validationIndicatorCol or its default value. + """ + return self.getOrDefault(self.validationIndicatorCol) + + +class HasBlockSize(Params): + """ + Mixin for param blockSize: block size for stacking input data in matrices. Data is stacked within partitions. If block size is more than remaining data in a partition then it is adjusted to the size of this data. + """ + + blockSize: "Param[int]" = Param( + Params._dummy(), + "blockSize", + "block size for stacking input data in matrices. Data is stacked within partitions. If block size is more than remaining data in a partition then it is adjusted to the size of this data.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasBlockSize, self).__init__() + + def getBlockSize(self) -> int: + """ + Gets the value of blockSize or its default value. + """ + return self.getOrDefault(self.blockSize) + + +class HasMaxBlockSizeInMB(Params): + """ + Mixin for param maxBlockSizeInMB: maximum memory in MB for stacking input data into blocks. Data is stacked within partitions. If more than remaining data size in a partition then it is adjusted to the data size. Default 0.0 represents choosing optimal value, depends on specific algorithm. Must be >= 0. + """ + + maxBlockSizeInMB: "Param[float]" = Param( + Params._dummy(), + "maxBlockSizeInMB", + "maximum memory in MB for stacking input data into blocks. Data is stacked within partitions. If more than remaining data size in a partition then it is adjusted to the data size. Default 0.0 represents choosing optimal value, depends on specific algorithm. Must be >= 0.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasMaxBlockSizeInMB, self).__init__() + self._setDefault(maxBlockSizeInMB=0.0) + + def getMaxBlockSizeInMB(self) -> float: + """ + Gets the value of maxBlockSizeInMB or its default value. + """ + return self.getOrDefault(self.maxBlockSizeInMB) + + +class HasNumTrainWorkers(Params): + """ + Mixin for param numTrainWorkers: number of training workers + """ + + numTrainWorkers: "Param[int]" = Param( + Params._dummy(), + "numTrainWorkers", + "number of training workers", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasNumTrainWorkers, self).__init__() + self._setDefault(numTrainWorkers=1) + + def getNumTrainWorkers(self) -> int: + """ + Gets the value of numTrainWorkers or its default value. + """ + return self.getOrDefault(self.numTrainWorkers) + + +class HasBatchSize(Params): + """ + Mixin for param batchSize: number of training batch size + """ + + batchSize: "Param[int]" = Param( + Params._dummy(), + "batchSize", + "number of training batch size", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self) -> None: + super(HasBatchSize, self).__init__() + + def getBatchSize(self) -> int: + """ + Gets the value of batchSize or its default value. + """ + return self.getOrDefault(self.batchSize) + + +class HasLearningRate(Params): + """ + Mixin for param learningRate: learning rate for training + """ + + learningRate: "Param[float]" = Param( + Params._dummy(), + "learningRate", + "learning rate for training", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasLearningRate, self).__init__() + + def getLearningRate(self) -> float: + """ + Gets the value of learningRate or its default value. + """ + return self.getOrDefault(self.learningRate) + + +class HasMomentum(Params): + """ + Mixin for param momentum: momentum for training optimizer + """ + + momentum: "Param[float]" = Param( + Params._dummy(), + "momentum", + "momentum for training optimizer", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self) -> None: + super(HasMomentum, self).__init__() + + def getMomentum(self) -> float: + """ + Gets the value of momentum or its default value. + """ + return self.getOrDefault(self.momentum) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/pipeline.py new file mode 100644 index 0000000..24653d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/pipeline.py @@ -0,0 +1,451 @@ +# +# 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 os + +from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast, TYPE_CHECKING + +from pyspark import keyword_only, since, SparkContext +from pyspark.ml.base import Estimator, Model, Transformer +from pyspark.ml.param import Param, Params +from pyspark.ml.util import ( + MLReadable, + MLWritable, + JavaMLWriter, + JavaMLReader, + DefaultParamsReader, + DefaultParamsWriter, + MLWriter, + MLReader, + JavaMLReadable, + JavaMLWritable, +) +from pyspark.ml.wrapper import JavaParams +from pyspark.ml.common import inherit_doc +from pyspark.sql.dataframe import DataFrame + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap, PipelineStage + from py4j.java_gateway import JavaObject + + +@inherit_doc +class Pipeline(Estimator["PipelineModel"], MLReadable["Pipeline"], MLWritable): + """ + A simple pipeline, which acts as an estimator. A Pipeline consists + of a sequence of stages, each of which is either an + :py:class:`Estimator` or a :py:class:`Transformer`. When + :py:meth:`Pipeline.fit` is called, the stages are executed in + order. If a stage is an :py:class:`Estimator`, its + :py:meth:`Estimator.fit` method will be called on the input + dataset to fit a model. Then the model, which is a transformer, + will be used to transform the dataset as the input to the next + stage. If a stage is a :py:class:`Transformer`, its + :py:meth:`Transformer.transform` method will be called to produce + the dataset for the next stage. The fitted model from a + :py:class:`Pipeline` is a :py:class:`PipelineModel`, which + consists of fitted models and transformers, corresponding to the + pipeline stages. If stages is an empty list, the pipeline acts as an + identity transformer. + + .. versionadded:: 1.3.0 + """ + + stages: Param[List["PipelineStage"]] = Param( + Params._dummy(), "stages", "a list of pipeline stages" + ) + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__(self, *, stages: Optional[List["PipelineStage"]] = None): + """ + __init__(self, \\*, stages=None) + """ + super(Pipeline, self).__init__() + kwargs = self._input_kwargs + self.setParams(**kwargs) + + def setStages(self, value: List["PipelineStage"]) -> "Pipeline": + """ + Set pipeline stages. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + value : list + of :py:class:`pyspark.ml.Transformer` + or :py:class:`pyspark.ml.Estimator` + + Returns + ------- + :py:class:`Pipeline` + the pipeline instance + """ + return self._set(stages=value) + + @since("1.3.0") + def getStages(self) -> List["PipelineStage"]: + """ + Get pipeline stages. + """ + return self.getOrDefault(self.stages) + + @keyword_only + @since("1.3.0") + def setParams(self, *, stages: Optional[List["PipelineStage"]] = None) -> "Pipeline": + """ + setParams(self, \\*, stages=None) + Sets params for Pipeline. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _fit(self, dataset: DataFrame) -> "PipelineModel": + stages = self.getStages() + for stage in stages: + if not (isinstance(stage, Estimator) or isinstance(stage, Transformer)): + raise TypeError("Cannot recognize a pipeline stage of type %s." % type(stage)) + indexOfLastEstimator = -1 + for i, stage in enumerate(stages): + if isinstance(stage, Estimator): + indexOfLastEstimator = i + transformers: List[Transformer] = [] + for i, stage in enumerate(stages): + if i <= indexOfLastEstimator: + if isinstance(stage, Transformer): + transformers.append(stage) + dataset = stage.transform(dataset) + else: # must be an Estimator + model = stage.fit(dataset) + transformers.append(model) + if i < indexOfLastEstimator: + dataset = model.transform(dataset) + else: + transformers.append(cast(Transformer, stage)) + return PipelineModel(transformers) + + def copy(self, extra: Optional["ParamMap"] = None) -> "Pipeline": + """ + Creates a copy of this instance. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + extra : dict, optional + extra parameters + + Returns + ------- + :py:class:`Pipeline` + new instance + """ + if extra is None: + extra = dict() + that = Params.copy(self, extra) + stages = [stage.copy(extra) for stage in that.getStages()] + return that.setStages(stages) + + @since("2.0.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + allStagesAreJava = PipelineSharedReadWrite.checkStagesForJava(self.getStages()) + if allStagesAreJava: + return JavaMLWriter(self) # type: ignore[arg-type] + return PipelineWriter(self) + + @classmethod + @since("2.0.0") + def read(cls) -> "PipelineReader": + """Returns an MLReader instance for this class.""" + return PipelineReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "Pipeline": + """ + Given a Java Pipeline, create and return a Python wrapper of it. + Used for ML persistence. + """ + # Create a new instance of this stage. + py_stage = cls() + # Load information from java_stage to the instance. + py_stages: List["PipelineStage"] = [ + JavaParams._from_java(s) for s in java_stage.getStages() + ] + py_stage.setStages(py_stages) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java Pipeline. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + + gateway = SparkContext._gateway + assert gateway is not None and SparkContext._jvm is not None + + cls = SparkContext._jvm.org.apache.spark.ml.PipelineStage + java_stages = gateway.new_array(cls, len(self.getStages())) + for idx, stage in enumerate(self.getStages()): + java_stages[idx] = cast(JavaParams, stage)._to_java() + + _java_obj = JavaParams._new_java_obj("org.apache.spark.ml.Pipeline", self.uid) + _java_obj.setStages(java_stages) + + return _java_obj + + +@inherit_doc +class PipelineWriter(MLWriter): + """ + (Private) Specialization of :py:class:`MLWriter` for :py:class:`Pipeline` types + """ + + def __init__(self, instance: Pipeline): + super(PipelineWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + stages = self.instance.getStages() + PipelineSharedReadWrite.validateStages(stages) + PipelineSharedReadWrite.saveImpl(self.instance, stages, self.sc, path) + + +@inherit_doc +class PipelineReader(MLReader[Pipeline]): + """ + (Private) Specialization of :py:class:`MLReader` for :py:class:`Pipeline` types + """ + + def __init__(self, cls: Type[Pipeline]): + super(PipelineReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> Pipeline: + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if "language" not in metadata["paramMap"] or metadata["paramMap"]["language"] != "Python": + return JavaMLReader(cast(Type["JavaMLReadable[Pipeline]"], self.cls)).load(path) + else: + uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path) + return Pipeline(stages=stages)._resetUid(uid) + + +@inherit_doc +class PipelineModelWriter(MLWriter): + """ + (Private) Specialization of :py:class:`MLWriter` for :py:class:`PipelineModel` types + """ + + def __init__(self, instance: "PipelineModel"): + super(PipelineModelWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + stages = self.instance.stages + PipelineSharedReadWrite.validateStages(cast(List["PipelineStage"], stages)) + PipelineSharedReadWrite.saveImpl( + self.instance, cast(List["PipelineStage"], stages), self.sc, path + ) + + +@inherit_doc +class PipelineModelReader(MLReader["PipelineModel"]): + """ + (Private) Specialization of :py:class:`MLReader` for :py:class:`PipelineModel` types + """ + + def __init__(self, cls: Type["PipelineModel"]): + super(PipelineModelReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> "PipelineModel": + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if "language" not in metadata["paramMap"] or metadata["paramMap"]["language"] != "Python": + return JavaMLReader(cast(Type["JavaMLReadable[PipelineModel]"], self.cls)).load(path) + else: + uid, stages = PipelineSharedReadWrite.load(metadata, self.sc, path) + return PipelineModel(stages=cast(List[Transformer], stages))._resetUid(uid) + + +@inherit_doc +class PipelineModel(Model, MLReadable["PipelineModel"], MLWritable): + """ + Represents a compiled pipeline with transformers and fitted models. + + .. versionadded:: 1.3.0 + """ + + def __init__(self, stages: List[Transformer]): + super(PipelineModel, self).__init__() + self.stages = stages + + def _transform(self, dataset: DataFrame) -> DataFrame: + for t in self.stages: + dataset = t.transform(dataset) + return dataset + + def copy(self, extra: Optional["ParamMap"] = None) -> "PipelineModel": + """ + Creates a copy of this instance. + + .. versionadded:: 1.4.0 + + :param extra: extra parameters + :returns: new instance + """ + if extra is None: + extra = dict() + stages = [stage.copy(extra) for stage in self.stages] + return PipelineModel(stages) + + @since("2.0.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + allStagesAreJava = PipelineSharedReadWrite.checkStagesForJava( + cast(List["PipelineStage"], self.stages) + ) + if allStagesAreJava: + return JavaMLWriter(self) # type: ignore[arg-type] + return PipelineModelWriter(self) + + @classmethod + @since("2.0.0") + def read(cls) -> PipelineModelReader: + """Returns an MLReader instance for this class.""" + return PipelineModelReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "PipelineModel": + """ + Given a Java PipelineModel, create and return a Python wrapper of it. + Used for ML persistence. + """ + # Load information from java_stage to the instance. + py_stages: List[Transformer] = [JavaParams._from_java(s) for s in java_stage.stages()] + # Create a new instance of this stage. + py_stage = cls(py_stages) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java PipelineModel. Used for ML persistence. + + :return: Java object equivalent to this instance. + """ + + gateway = SparkContext._gateway + assert gateway is not None and SparkContext._jvm is not None + + cls = SparkContext._jvm.org.apache.spark.ml.Transformer + java_stages = gateway.new_array(cls, len(self.stages)) + for idx, stage in enumerate(self.stages): + java_stages[idx] = cast(JavaParams, stage)._to_java() + + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.PipelineModel", self.uid, java_stages + ) + + return _java_obj + + +@inherit_doc +class PipelineSharedReadWrite: + """ + Functions for :py:class:`MLReader` and :py:class:`MLWriter` shared between + :py:class:`Pipeline` and :py:class:`PipelineModel` + + .. versionadded:: 2.3.0 + """ + + @staticmethod + def checkStagesForJava(stages: List["PipelineStage"]) -> bool: + return all(isinstance(stage, JavaMLWritable) for stage in stages) + + @staticmethod + def validateStages(stages: List["PipelineStage"]) -> None: + """ + Check that all stages are Writable + """ + for stage in stages: + if not isinstance(stage, MLWritable): + raise ValueError( + "Pipeline write will fail on this pipeline " + + "because stage %s of type %s is not MLWritable", + stage.uid, + type(stage), + ) + + @staticmethod + def saveImpl( + instance: Union[Pipeline, PipelineModel], + stages: List["PipelineStage"], + sc: SparkContext, + path: str, + ) -> None: + """ + Save metadata and stages for a :py:class:`Pipeline` or :py:class:`PipelineModel` + - save metadata to path/metadata + - save stages to stages/IDX_UID + """ + stageUids = [stage.uid for stage in stages] + jsonParams = {"stageUids": stageUids, "language": "Python"} + DefaultParamsWriter.saveMetadata(instance, path, sc, paramMap=jsonParams) + stagesDir = os.path.join(path, "stages") + for index, stage in enumerate(stages): + cast(MLWritable, stage).write().save( + PipelineSharedReadWrite.getStagePath(stage.uid, index, len(stages), stagesDir) + ) + + @staticmethod + def load( + metadata: Dict[str, Any], sc: SparkContext, path: str + ) -> Tuple[str, List["PipelineStage"]]: + """ + Load metadata and stages for a :py:class:`Pipeline` or :py:class:`PipelineModel` + + Returns + ------- + tuple + (UID, list of stages) + """ + stagesDir = os.path.join(path, "stages") + stageUids = metadata["paramMap"]["stageUids"] + stages = [] + for index, stageUid in enumerate(stageUids): + stagePath = PipelineSharedReadWrite.getStagePath( + stageUid, index, len(stageUids), stagesDir + ) + stage: "PipelineStage" = DefaultParamsReader.loadParamsInstance(stagePath, sc) + stages.append(stage) + return (metadata["uid"], stages) + + @staticmethod + def getStagePath(stageUid: str, stageIdx: int, numStages: int, stagesDir: str) -> str: + """ + Get path for saving the given stage. + """ + stageIdxDigits = len(str(numStages)) + stageDir = str(stageIdx).zfill(stageIdxDigits) + "_" + stageUid + stagePath = os.path.join(stagesDir, stageDir) + return stagePath diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/recommendation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/recommendation.py new file mode 100644 index 0000000..873140e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/recommendation.py @@ -0,0 +1,748 @@ +# +# 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 sys +from typing import Any, Dict, Optional, TYPE_CHECKING + +from pyspark import since, keyword_only +from pyspark.ml.param.shared import ( + HasPredictionCol, + HasBlockSize, + HasMaxIter, + HasRegParam, + HasCheckpointInterval, + HasSeed, +) +from pyspark.ml.wrapper import JavaEstimator, JavaModel +from pyspark.ml.common import inherit_doc +from pyspark.ml.param import Params, TypeConverters, Param +from pyspark.ml.util import JavaMLWritable, JavaMLReadable +from pyspark.sql import DataFrame + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + + +__all__ = ["ALS", "ALSModel"] + + +@inherit_doc +class _ALSModelParams(HasPredictionCol, HasBlockSize): + """ + Params for :py:class:`ALS` and :py:class:`ALSModel`. + + .. versionadded:: 3.0.0 + """ + + userCol: Param[str] = Param( + Params._dummy(), + "userCol", + "column name for user ids. Ids must be within " + "the integer value range.", + typeConverter=TypeConverters.toString, + ) + itemCol: Param[str] = Param( + Params._dummy(), + "itemCol", + "column name for item ids. Ids must be within " + "the integer value range.", + typeConverter=TypeConverters.toString, + ) + coldStartStrategy: Param[str] = Param( + Params._dummy(), + "coldStartStrategy", + "strategy for dealing with " + + "unknown or new users/items at prediction time. This may be useful " + + "in cross-validation or production scenarios, for handling " + + "user/item ids the model has not seen in the training data. " + + "Supported values: 'nan', 'drop'.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_ALSModelParams, self).__init__(*args) + self._setDefault(blockSize=4096) + + @since("1.4.0") + def getUserCol(self) -> str: + """ + Gets the value of userCol or its default value. + """ + return self.getOrDefault(self.userCol) + + @since("1.4.0") + def getItemCol(self) -> str: + """ + Gets the value of itemCol or its default value. + """ + return self.getOrDefault(self.itemCol) + + @since("2.2.0") + def getColdStartStrategy(self) -> str: + """ + Gets the value of coldStartStrategy or its default value. + """ + return self.getOrDefault(self.coldStartStrategy) + + +@inherit_doc +class _ALSParams(_ALSModelParams, HasMaxIter, HasRegParam, HasCheckpointInterval, HasSeed): + """ + Params for :py:class:`ALS`. + + .. versionadded:: 3.0.0 + """ + + rank: Param[int] = Param( + Params._dummy(), "rank", "rank of the factorization", typeConverter=TypeConverters.toInt + ) + numUserBlocks: Param[int] = Param( + Params._dummy(), + "numUserBlocks", + "number of user blocks", + typeConverter=TypeConverters.toInt, + ) + numItemBlocks: Param[int] = Param( + Params._dummy(), + "numItemBlocks", + "number of item blocks", + typeConverter=TypeConverters.toInt, + ) + implicitPrefs: Param[bool] = Param( + Params._dummy(), + "implicitPrefs", + "whether to use implicit preference", + typeConverter=TypeConverters.toBoolean, + ) + alpha: Param[float] = Param( + Params._dummy(), + "alpha", + "alpha for implicit preference", + typeConverter=TypeConverters.toFloat, + ) + + ratingCol: Param[str] = Param( + Params._dummy(), + "ratingCol", + "column name for ratings", + typeConverter=TypeConverters.toString, + ) + nonnegative: Param[bool] = Param( + Params._dummy(), + "nonnegative", + "whether to use nonnegative constraint for least squares", + typeConverter=TypeConverters.toBoolean, + ) + intermediateStorageLevel: Param[str] = Param( + Params._dummy(), + "intermediateStorageLevel", + "StorageLevel for intermediate datasets. Cannot be 'NONE'.", + typeConverter=TypeConverters.toString, + ) + finalStorageLevel: Param[str] = Param( + Params._dummy(), + "finalStorageLevel", + "StorageLevel for ALS model factors.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_ALSParams, self).__init__(*args) + self._setDefault( + rank=10, + maxIter=10, + regParam=0.1, + numUserBlocks=10, + numItemBlocks=10, + implicitPrefs=False, + alpha=1.0, + userCol="user", + itemCol="item", + ratingCol="rating", + nonnegative=False, + checkpointInterval=10, + intermediateStorageLevel="MEMORY_AND_DISK", + finalStorageLevel="MEMORY_AND_DISK", + coldStartStrategy="nan", + ) + + @since("1.4.0") + def getRank(self) -> int: + """ + Gets the value of rank or its default value. + """ + return self.getOrDefault(self.rank) + + @since("1.4.0") + def getNumUserBlocks(self) -> int: + """ + Gets the value of numUserBlocks or its default value. + """ + return self.getOrDefault(self.numUserBlocks) + + @since("1.4.0") + def getNumItemBlocks(self) -> int: + """ + Gets the value of numItemBlocks or its default value. + """ + return self.getOrDefault(self.numItemBlocks) + + @since("1.4.0") + def getImplicitPrefs(self) -> bool: + """ + Gets the value of implicitPrefs or its default value. + """ + return self.getOrDefault(self.implicitPrefs) + + @since("1.4.0") + def getAlpha(self) -> float: + """ + Gets the value of alpha or its default value. + """ + return self.getOrDefault(self.alpha) + + @since("1.4.0") + def getRatingCol(self) -> str: + """ + Gets the value of ratingCol or its default value. + """ + return self.getOrDefault(self.ratingCol) + + @since("1.4.0") + def getNonnegative(self) -> bool: + """ + Gets the value of nonnegative or its default value. + """ + return self.getOrDefault(self.nonnegative) + + @since("2.0.0") + def getIntermediateStorageLevel(self) -> str: + """ + Gets the value of intermediateStorageLevel or its default value. + """ + return self.getOrDefault(self.intermediateStorageLevel) + + @since("2.0.0") + def getFinalStorageLevel(self) -> str: + """ + Gets the value of finalStorageLevel or its default value. + """ + return self.getOrDefault(self.finalStorageLevel) + + +@inherit_doc +class ALS(JavaEstimator["ALSModel"], _ALSParams, JavaMLWritable, JavaMLReadable["ALS"]): + """ + Alternating Least Squares (ALS) matrix factorization. + + ALS attempts to estimate the ratings matrix `R` as the product of + two lower-rank matrices, `X` and `Y`, i.e. `X * Yt = R`. Typically + these approximations are called 'factor' matrices. The general + approach is iterative. During each iteration, one of the factor + matrices is held constant, while the other is solved for using least + squares. The newly-solved factor matrix is then held constant while + solving for the other factor matrix. + + This is a blocked implementation of the ALS factorization algorithm + that groups the two sets of factors (referred to as "users" and + "products") into blocks and reduces communication by only sending + one copy of each user vector to each product block on each + iteration, and only for the product blocks that need that user's + feature vector. This is achieved by pre-computing some information + about the ratings matrix to determine the "out-links" of each user + (which blocks of products it will contribute to) and "in-link" + information for each product (which of the feature vectors it + receives from each user block it will depend on). This allows us to + send only an array of feature vectors between each user block and + product block, and have the product block find the users' ratings + and update the products based on these messages. + + For implicit preference data, the algorithm used is based on + `"Collaborative Filtering for Implicit Feedback Datasets", + `_, adapted for the blocked + approach used here. + + Essentially instead of finding the low-rank approximations to the + rating matrix `R`, this finds the approximations for a preference + matrix `P` where the elements of `P` are 1 if r > 0 and 0 if r <= 0. + The ratings then act as 'confidence' values related to strength of + indicated user preferences rather than explicit ratings given to + items. + + .. versionadded:: 1.4.0 + + Notes + ----- + The input rating dataframe to the ALS implementation should be deterministic. + Nondeterministic data can cause failure during fitting ALS model. + For example, an order-sensitive operation like sampling after a repartition makes + dataframe output nondeterministic, like `df.repartition(2).sample(False, 0.5, 1618)`. + Checkpointing sampled dataframe or adding a sort before sampling can help make the + dataframe deterministic. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(0, 0, 4.0), (0, 1, 2.0), (1, 1, 3.0), (1, 2, 4.0), (2, 1, 1.0), (2, 2, 5.0)], + ... ["user", "item", "rating"]) + >>> als = ALS(rank=10, seed=0) + >>> als.setMaxIter(5) + ALS... + >>> als.getMaxIter() + 5 + >>> als.setRegParam(0.1) + ALS... + >>> als.getRegParam() + 0.1 + >>> als.clear(als.regParam) + >>> model = als.fit(df) + >>> model.getBlockSize() + 4096 + >>> model.getUserCol() + 'user' + >>> model.setUserCol("user") + ALSModel... + >>> model.getItemCol() + 'item' + >>> model.setPredictionCol("newPrediction") + ALS... + >>> model.rank + 10 + >>> model.userFactors.orderBy("id").collect() + [Row(id=0, features=[...]), Row(id=1, ...), Row(id=2, ...)] + >>> test = spark.createDataFrame([(0, 2), (1, 0), (2, 0)], ["user", "item"]) + >>> predictions = sorted(model.transform(test).collect(), key=lambda r: r[0]) + >>> predictions[0] + Row(user=0, item=2, newPrediction=0.6929...) + >>> predictions[1] + Row(user=1, item=0, newPrediction=3.47356...) + >>> predictions[2] + Row(user=2, item=0, newPrediction=-0.899198...) + >>> user_recs = model.recommendForAllUsers(3) + >>> user_recs.where(user_recs.user == 0)\ + .select("recommendations.item", "recommendations.rating").collect() + [Row(item=[0, 1, 2], rating=[3.910..., 1.997..., 0.692...])] + >>> item_recs = model.recommendForAllItems(3) + >>> item_recs.where(item_recs.item == 2)\ + .select("recommendations.user", "recommendations.rating").collect() + [Row(user=[2, 1, 0], rating=[4.892..., 3.991..., 0.692...])] + >>> user_subset = df.where(df.user == 2) + >>> user_subset_recs = model.recommendForUserSubset(user_subset, 3) + >>> user_subset_recs.select("recommendations.item", "recommendations.rating").first() + Row(item=[2, 1, 0], rating=[4.892..., 1.076..., -0.899...]) + >>> item_subset = df.where(df.item == 0) + >>> item_subset_recs = model.recommendForItemSubset(item_subset, 3) + >>> item_subset_recs.select("recommendations.user", "recommendations.rating").first() + Row(user=[0, 1, 2], rating=[3.910..., 3.473..., -0.899...]) + >>> als_path = temp_path + "/als" + >>> als.save(als_path) + >>> als2 = ALS.load(als_path) + >>> als.getMaxIter() + 5 + >>> model_path = temp_path + "/als_model" + >>> model.save(model_path) + >>> model2 = ALSModel.load(model_path) + >>> model.rank == model2.rank + True + >>> sorted(model.userFactors.collect()) == sorted(model2.userFactors.collect()) + True + >>> sorted(model.itemFactors.collect()) == sorted(model2.itemFactors.collect()) + True + >>> model.transform(test).take(1) == model2.transform(test).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + rank: int = 10, + maxIter: int = 10, + regParam: float = 0.1, + numUserBlocks: int = 10, + numItemBlocks: int = 10, + implicitPrefs: bool = False, + alpha: float = 1.0, + userCol: str = "user", + itemCol: str = "item", + seed: Optional[int] = None, + ratingCol: str = "rating", + nonnegative: bool = False, + checkpointInterval: int = 10, + intermediateStorageLevel: str = "MEMORY_AND_DISK", + finalStorageLevel: str = "MEMORY_AND_DISK", + coldStartStrategy: str = "nan", + blockSize: int = 4096, + ): + """ + __init__(self, \\*, rank=10, maxIter=10, regParam=0.1, numUserBlocks=10, + numItemBlocks=10, implicitPrefs=False, alpha=1.0, userCol="user", itemCol="item", \ + seed=None, ratingCol="rating", nonnegative=False, checkpointInterval=10, \ + intermediateStorageLevel="MEMORY_AND_DISK", \ + finalStorageLevel="MEMORY_AND_DISK", coldStartStrategy="nan", blockSize=4096) + """ + super(ALS, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.recommendation.ALS", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + rank: int = 10, + maxIter: int = 10, + regParam: float = 0.1, + numUserBlocks: int = 10, + numItemBlocks: int = 10, + implicitPrefs: bool = False, + alpha: float = 1.0, + userCol: str = "user", + itemCol: str = "item", + seed: Optional[int] = None, + ratingCol: str = "rating", + nonnegative: bool = False, + checkpointInterval: int = 10, + intermediateStorageLevel: str = "MEMORY_AND_DISK", + finalStorageLevel: str = "MEMORY_AND_DISK", + coldStartStrategy: str = "nan", + blockSize: int = 4096, + ) -> "ALS": + """ + setParams(self, \\*, rank=10, maxIter=10, regParam=0.1, numUserBlocks=10, \ + numItemBlocks=10, implicitPrefs=False, alpha=1.0, userCol="user", itemCol="item", \ + seed=None, ratingCol="rating", nonnegative=False, checkpointInterval=10, \ + intermediateStorageLevel="MEMORY_AND_DISK", \ + finalStorageLevel="MEMORY_AND_DISK", coldStartStrategy="nan", blockSize=4096) + Sets params for ALS. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "ALSModel": + return ALSModel(java_model) + + @since("1.4.0") + def setRank(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`rank`. + """ + return self._set(rank=value) + + @since("1.4.0") + def setNumUserBlocks(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`numUserBlocks`. + """ + return self._set(numUserBlocks=value) + + @since("1.4.0") + def setNumItemBlocks(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`numItemBlocks`. + """ + return self._set(numItemBlocks=value) + + @since("1.4.0") + def setNumBlocks(self, value: int) -> "ALS": + """ + Sets both :py:attr:`numUserBlocks` and :py:attr:`numItemBlocks` to the specific value. + """ + self._set(numUserBlocks=value) + return self._set(numItemBlocks=value) + + @since("1.4.0") + def setImplicitPrefs(self, value: bool) -> "ALS": + """ + Sets the value of :py:attr:`implicitPrefs`. + """ + return self._set(implicitPrefs=value) + + @since("1.4.0") + def setAlpha(self, value: float) -> "ALS": + """ + Sets the value of :py:attr:`alpha`. + """ + return self._set(alpha=value) + + @since("1.4.0") + def setUserCol(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`userCol`. + """ + return self._set(userCol=value) + + @since("1.4.0") + def setItemCol(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`itemCol`. + """ + return self._set(itemCol=value) + + @since("1.4.0") + def setRatingCol(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`ratingCol`. + """ + return self._set(ratingCol=value) + + @since("1.4.0") + def setNonnegative(self, value: bool) -> "ALS": + """ + Sets the value of :py:attr:`nonnegative`. + """ + return self._set(nonnegative=value) + + @since("2.0.0") + def setIntermediateStorageLevel(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`intermediateStorageLevel`. + """ + return self._set(intermediateStorageLevel=value) + + @since("2.0.0") + def setFinalStorageLevel(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`finalStorageLevel`. + """ + return self._set(finalStorageLevel=value) + + @since("2.2.0") + def setColdStartStrategy(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`coldStartStrategy`. + """ + return self._set(coldStartStrategy=value) + + def setMaxIter(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + def setRegParam(self, value: float) -> "ALS": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + def setPredictionCol(self, value: str) -> "ALS": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + def setCheckpointInterval(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + def setSeed(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setBlockSize(self, value: int) -> "ALS": + """ + Sets the value of :py:attr:`blockSize`. + """ + return self._set(blockSize=value) + + +class ALSModel(JavaModel, _ALSModelParams, JavaMLWritable, JavaMLReadable["ALSModel"]): + """ + Model fitted by ALS. + + .. versionadded:: 1.4.0 + """ + + @since("3.0.0") + def setUserCol(self, value: str) -> "ALSModel": + """ + Sets the value of :py:attr:`userCol`. + """ + return self._set(userCol=value) + + @since("3.0.0") + def setItemCol(self, value: str) -> "ALSModel": + """ + Sets the value of :py:attr:`itemCol`. + """ + return self._set(itemCol=value) + + @since("3.0.0") + def setColdStartStrategy(self, value: str) -> "ALSModel": + """ + Sets the value of :py:attr:`coldStartStrategy`. + """ + return self._set(coldStartStrategy=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "ALSModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("3.0.0") + def setBlockSize(self, value: int) -> "ALSModel": + """ + Sets the value of :py:attr:`blockSize`. + """ + return self._set(blockSize=value) + + @property + @since("1.4.0") + def rank(self) -> int: + """rank of the matrix factorization model""" + return self._call_java("rank") + + @property + @since("1.4.0") + def userFactors(self) -> DataFrame: + """ + a DataFrame that stores user factors in two columns: `id` and + `features` + """ + return self._call_java("userFactors") + + @property + @since("1.4.0") + def itemFactors(self) -> DataFrame: + """ + a DataFrame that stores item factors in two columns: `id` and + `features` + """ + return self._call_java("itemFactors") + + def recommendForAllUsers(self, numItems: int) -> DataFrame: + """ + Returns top `numItems` items recommended for each user, for all users. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + numItems : int + max number of recommendations for each user + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + a DataFrame of (userCol, recommendations), where recommendations are + stored as an array of (itemCol, rating) Rows. + """ + return self._call_java("recommendForAllUsers", numItems) + + def recommendForAllItems(self, numUsers: int) -> DataFrame: + """ + Returns top `numUsers` users recommended for each item, for all items. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + numUsers : int + max number of recommendations for each item + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + a DataFrame of (itemCol, recommendations), where recommendations are + stored as an array of (userCol, rating) Rows. + """ + return self._call_java("recommendForAllItems", numUsers) + + def recommendForUserSubset(self, dataset: DataFrame, numItems: int) -> DataFrame: + """ + Returns top `numItems` items recommended for each user id in the input data set. Note that + if there are duplicate ids in the input dataset, only one set of recommendations per unique + id will be returned. + + .. versionadded:: 2.3.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a DataFrame containing a column of user ids. The column name must match `userCol`. + numItems : int + max number of recommendations for each user + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + a DataFrame of (userCol, recommendations), where recommendations are + stored as an array of (itemCol, rating) Rows. + """ + return self._call_java("recommendForUserSubset", dataset, numItems) + + def recommendForItemSubset(self, dataset: DataFrame, numUsers: int) -> DataFrame: + """ + Returns top `numUsers` users recommended for each item id in the input data set. Note that + if there are duplicate ids in the input dataset, only one set of recommendations per unique + id will be returned. + + .. versionadded:: 2.3.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a DataFrame containing a column of item ids. The column name must match `itemCol`. + numUsers : int + max number of recommendations for each item + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + a DataFrame of (itemCol, recommendations), where recommendations are + stored as an array of (userCol, rating) Rows. + """ + return self._call_java("recommendForItemSubset", dataset, numUsers) + + +if __name__ == "__main__": + import doctest + import pyspark.ml.recommendation + from pyspark.sql import SparkSession + + globs = pyspark.ml.recommendation.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.recommendation tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/regression.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/regression.py new file mode 100644 index 0000000..a4ce961 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/regression.py @@ -0,0 +1,3335 @@ +# +# 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 sys + +from typing import Any, Dict, Generic, List, Optional, TypeVar, TYPE_CHECKING + +from abc import ABCMeta + +from pyspark import keyword_only, since +from pyspark.ml import Predictor, PredictionModel +from pyspark.ml.base import _PredictorParams +from pyspark.ml.param.shared import ( + HasFeaturesCol, + HasLabelCol, + HasPredictionCol, + HasWeightCol, + Param, + Params, + TypeConverters, + HasMaxIter, + HasTol, + HasFitIntercept, + HasAggregationDepth, + HasMaxBlockSizeInMB, + HasRegParam, + HasSolver, + HasStepSize, + HasSeed, + HasElasticNetParam, + HasStandardization, + HasLoss, + HasVarianceCol, +) +from pyspark.ml.tree import ( + _DecisionTreeModel, + _DecisionTreeParams, + _TreeEnsembleModel, + _RandomForestParams, + _GBTParams, + _TreeRegressorParams, +) +from pyspark.ml.base import Transformer +from pyspark.ml.linalg import Vector, Matrix +from pyspark.ml.util import ( + JavaMLWritable, + JavaMLReadable, + HasTrainingSummary, + GeneralJavaMLWritable, +) +from pyspark.ml.wrapper import ( + JavaEstimator, + JavaModel, + JavaPredictor, + JavaPredictionModel, + JavaTransformer, + JavaWrapper, +) +from pyspark.ml.common import inherit_doc +from pyspark.sql import DataFrame + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + +T = TypeVar("T") +M = TypeVar("M", bound=Transformer) +JM = TypeVar("JM", bound=JavaTransformer) + + +__all__ = [ + "AFTSurvivalRegression", + "AFTSurvivalRegressionModel", + "DecisionTreeRegressor", + "DecisionTreeRegressionModel", + "GBTRegressor", + "GBTRegressionModel", + "GeneralizedLinearRegression", + "GeneralizedLinearRegressionModel", + "GeneralizedLinearRegressionSummary", + "GeneralizedLinearRegressionTrainingSummary", + "IsotonicRegression", + "IsotonicRegressionModel", + "LinearRegression", + "LinearRegressionModel", + "LinearRegressionSummary", + "LinearRegressionTrainingSummary", + "RandomForestRegressor", + "RandomForestRegressionModel", + "FMRegressor", + "FMRegressionModel", +] + + +class Regressor(Predictor[M], _PredictorParams, Generic[M], metaclass=ABCMeta): + """ + Regressor for regression tasks. + + .. versionadded:: 3.0.0 + """ + + pass + + +class RegressionModel(PredictionModel[T], _PredictorParams, metaclass=ABCMeta): + """ + Model produced by a ``Regressor``. + + .. versionadded:: 3.0.0 + """ + + pass + + +class _JavaRegressor(Regressor, JavaPredictor[JM], Generic[JM], metaclass=ABCMeta): + """ + Java Regressor for regression tasks. + + .. versionadded:: 3.0.0 + """ + + pass + + +class _JavaRegressionModel(RegressionModel, JavaPredictionModel[T], metaclass=ABCMeta): + """ + Java Model produced by a ``_JavaRegressor``. + To be mixed in with :class:`pyspark.ml.JavaModel` + + .. versionadded:: 3.0.0 + """ + + pass + + +class _LinearRegressionParams( + _PredictorParams, + HasRegParam, + HasElasticNetParam, + HasMaxIter, + HasTol, + HasFitIntercept, + HasStandardization, + HasWeightCol, + HasSolver, + HasAggregationDepth, + HasLoss, + HasMaxBlockSizeInMB, +): + """ + Params for :py:class:`LinearRegression` and :py:class:`LinearRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + solver: Param[str] = Param( + Params._dummy(), + "solver", + "The solver algorithm for optimization. Supported " + "options: auto, normal, l-bfgs.", + typeConverter=TypeConverters.toString, + ) + + loss: Param[str] = Param( + Params._dummy(), + "loss", + "The loss function to be optimized. Supported " + "options: squaredError, huber.", + typeConverter=TypeConverters.toString, + ) + + epsilon: Param[float] = Param( + Params._dummy(), + "epsilon", + "The shape parameter to control the amount of " + + "robustness. Must be > 1.0. Only valid when loss is huber", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_LinearRegressionParams, self).__init__(*args) + self._setDefault( + maxIter=100, + regParam=0.0, + tol=1e-6, + loss="squaredError", + epsilon=1.35, + maxBlockSizeInMB=0.0, + ) + + @since("2.3.0") + def getEpsilon(self) -> float: + """ + Gets the value of epsilon or its default value. + """ + return self.getOrDefault(self.epsilon) + + +@inherit_doc +class LinearRegression( + _JavaRegressor["LinearRegressionModel"], + _LinearRegressionParams, + JavaMLWritable, + JavaMLReadable["LinearRegression"], +): + """ + Linear regression. + + The learning objective is to minimize the specified loss function, with regularization. + This supports two kinds of loss: + + * squaredError (a.k.a squared loss) + * huber (a hybrid of squared error for relatively small errors and absolute error for \ + relatively large ones, and we estimate the scale parameter from training data) + + This supports multiple types of regularization: + + * none (a.k.a. ordinary least squares) + * L2 (ridge regression) + * L1 (Lasso) + * L2 + L1 (elastic net) + + .. versionadded:: 1.4.0 + + Notes + ----- + Fitting with huber loss only supports none and L2 regularization. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, 2.0, Vectors.dense(1.0)), + ... (0.0, 2.0, Vectors.sparse(1, [], []))], ["label", "weight", "features"]) + >>> lr = LinearRegression(regParam=0.0, solver="normal", weightCol="weight") + >>> lr.setMaxIter(5) + LinearRegression... + >>> lr.getMaxIter() + 5 + >>> lr.setRegParam(0.1) + LinearRegression... + >>> lr.getRegParam() + 0.1 + >>> lr.setRegParam(0.0) + LinearRegression... + >>> model = lr.fit(df) + >>> model.setFeaturesCol("features") + LinearRegressionModel... + >>> model.setPredictionCol("newPrediction") + LinearRegressionModel... + >>> model.getMaxIter() + 5 + >>> model.getMaxBlockSizeInMB() + 0.0 + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> abs(model.predict(test0.head().features) - (-1.0)) < 0.001 + True + >>> abs(model.transform(test0).head().newPrediction - (-1.0)) < 0.001 + True + >>> abs(model.coefficients[0] - 1.0) < 0.001 + True + >>> abs(model.intercept - 0.0) < 0.001 + True + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> abs(model.transform(test1).head().newPrediction - 1.0) < 0.001 + True + >>> lr.setParams(featuresCol="vector") + LinearRegression... + >>> lr_path = temp_path + "/lr" + >>> lr.save(lr_path) + >>> lr2 = LinearRegression.load(lr_path) + >>> lr2.getMaxIter() + 5 + >>> model_path = temp_path + "/lr_model" + >>> model.save(model_path) + >>> model2 = LinearRegressionModel.load(model_path) + >>> model.coefficients[0] == model2.coefficients[0] + True + >>> model.intercept == model2.intercept + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> model.numFeatures + 1 + >>> model.write().format("pmml").save(model_path + "_2") + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + elasticNetParam: float = 0.0, + tol: float = 1e-6, + fitIntercept: bool = True, + standardization: bool = True, + solver: str = "auto", + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + loss: str = "squaredError", + epsilon: float = 1.35, + maxBlockSizeInMB: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \ + standardization=True, solver="auto", weightCol=None, aggregationDepth=2, \ + loss="squaredError", epsilon=1.35, maxBlockSizeInMB=0.0) + """ + super(LinearRegression, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.LinearRegression", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxIter: int = 100, + regParam: float = 0.0, + elasticNetParam: float = 0.0, + tol: float = 1e-6, + fitIntercept: bool = True, + standardization: bool = True, + solver: str = "auto", + weightCol: Optional[str] = None, + aggregationDepth: int = 2, + loss: str = "squaredError", + epsilon: float = 1.35, + maxBlockSizeInMB: float = 0.0, + ) -> "LinearRegression": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \ + standardization=True, solver="auto", weightCol=None, aggregationDepth=2, \ + loss="squaredError", epsilon=1.35, maxBlockSizeInMB=0.0) + Sets params for linear regression. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "LinearRegressionModel": + return LinearRegressionModel(java_model) + + @since("2.3.0") + def setEpsilon(self, value: float) -> "LinearRegression": + """ + Sets the value of :py:attr:`epsilon`. + """ + return self._set(epsilon=value) + + def setMaxIter(self, value: int) -> "LinearRegression": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + def setRegParam(self, value: float) -> "LinearRegression": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + def setTol(self, value: float) -> "LinearRegression": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + def setElasticNetParam(self, value: float) -> "LinearRegression": + """ + Sets the value of :py:attr:`elasticNetParam`. + """ + return self._set(elasticNetParam=value) + + def setFitIntercept(self, value: bool) -> "LinearRegression": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + def setStandardization(self, value: bool) -> "LinearRegression": + """ + Sets the value of :py:attr:`standardization`. + """ + return self._set(standardization=value) + + def setWeightCol(self, value: str) -> "LinearRegression": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + def setSolver(self, value: str) -> "LinearRegression": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + def setAggregationDepth(self, value: int) -> "LinearRegression": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + def setLoss(self, value: str) -> "LinearRegression": + """ + Sets the value of :py:attr:`loss`. + """ + return self._set(lossType=value) + + @since("3.1.0") + def setMaxBlockSizeInMB(self, value: float) -> "LinearRegression": + """ + Sets the value of :py:attr:`maxBlockSizeInMB`. + """ + return self._set(maxBlockSizeInMB=value) + + +class LinearRegressionModel( + _JavaRegressionModel, + _LinearRegressionParams, + GeneralJavaMLWritable, + JavaMLReadable["LinearRegressionModel"], + HasTrainingSummary["LinearRegressionSummary"], +): + """ + Model fitted by :class:`LinearRegression`. + + .. versionadded:: 1.4.0 + """ + + @property + @since("2.0.0") + def coefficients(self) -> Vector: + """ + Model coefficients. + """ + return self._call_java("coefficients") + + @property + @since("1.4.0") + def intercept(self) -> float: + """ + Model intercept. + """ + return self._call_java("intercept") + + @property + @since("2.3.0") + def scale(self) -> float: + r""" + The value by which :math:`\|y - X'w\|` is scaled down when loss is "huber", otherwise 1.0. + """ + return self._call_java("scale") + + @property + @since("2.0.0") + def summary(self) -> "LinearRegressionTrainingSummary": + """ + Gets summary (residuals, MSE, r-squared ) of model on + training set. An exception is thrown if + `trainingSummary is None`. + """ + if self.hasSummary: + return LinearRegressionTrainingSummary(super(LinearRegressionModel, self).summary) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "LinearRegressionSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on, where dataset is an + instance of :py:class:`pyspark.sql.DataFrame` + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_lr_summary = self._call_java("evaluate", dataset) + return LinearRegressionSummary(java_lr_summary) + + +class LinearRegressionSummary(JavaWrapper): + """ + Linear regression results evaluated on a dataset. + + .. versionadded:: 2.0.0 + """ + + @property + @since("2.0.0") + def predictions(self) -> DataFrame: + """ + Dataframe outputted by the model's `transform` method. + """ + return self._call_java("predictions") + + @property + @since("2.0.0") + def predictionCol(self) -> str: + """ + Field in "predictions" which gives the predicted value of + the label at each instance. + """ + return self._call_java("predictionCol") + + @property + @since("2.0.0") + def labelCol(self) -> str: + """ + Field in "predictions" which gives the true label of each + instance. + """ + return self._call_java("labelCol") + + @property + @since("2.0.0") + def featuresCol(self) -> str: + """ + Field in "predictions" which gives the features of each instance + as a vector. + """ + return self._call_java("featuresCol") + + @property + @since("2.0.0") + def explainedVariance(self) -> float: + r""" + Returns the explained variance regression score. + explainedVariance = :math:`1 - \frac{variance(y - \hat{y})}{variance(y)}` + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark + versions. + + For additional information see + `Explained variation on Wikipedia \ + `_ + """ + return self._call_java("explainedVariance") + + @property + @since("2.0.0") + def meanAbsoluteError(self) -> float: + """ + Returns the mean absolute error, which is a risk function + corresponding to the expected value of the absolute error + loss or l1-norm loss. + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark + versions. + """ + return self._call_java("meanAbsoluteError") + + @property + @since("2.0.0") + def meanSquaredError(self) -> float: + """ + Returns the mean squared error, which is a risk function + corresponding to the expected value of the squared error + loss or quadratic loss. + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark + versions. + """ + return self._call_java("meanSquaredError") + + @property + @since("2.0.0") + def rootMeanSquaredError(self) -> float: + """ + Returns the root mean squared error, which is defined as the + square root of the mean squared error. + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark + versions. + """ + return self._call_java("rootMeanSquaredError") + + @property + @since("2.0.0") + def r2(self) -> float: + """ + Returns R^2, the coefficient of determination. + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark + versions. + + See also `Wikipedia coefficient of determination \ + `_ + """ + return self._call_java("r2") + + @property + @since("2.4.0") + def r2adj(self) -> float: + """ + Returns Adjusted R^2, the adjusted coefficient of determination. + + Notes + ----- + This ignores instance weights (setting all to 1.0) from + `LinearRegression.weightCol`. This will change in later Spark versions. + + `Wikipedia coefficient of determination, Adjusted R^2 \ + `_ + """ + return self._call_java("r2adj") + + @property + @since("2.0.0") + def residuals(self) -> DataFrame: + """ + Residuals (label - predicted value) + """ + return self._call_java("residuals") + + @property + @since("2.0.0") + def numInstances(self) -> int: + """ + Number of instances in DataFrame predictions + """ + return self._call_java("numInstances") + + @property + @since("2.2.0") + def degreesOfFreedom(self) -> int: + """ + Degrees of freedom. + """ + return self._call_java("degreesOfFreedom") + + @property + @since("2.0.0") + def devianceResiduals(self) -> List[float]: + """ + The weighted residuals, the usual residuals rescaled by the + square root of the instance weights. + """ + return self._call_java("devianceResiduals") + + @property + def coefficientStandardErrors(self) -> List[float]: + """ + Standard error of estimated coefficients and intercept. + This value is only available when using the "normal" solver. + + If :py:attr:`LinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + + .. versionadded:: 2.0.0 + + See Also + -------- + LinearRegression.solver + """ + return self._call_java("coefficientStandardErrors") + + @property + def tValues(self) -> List[float]: + """ + T-statistic of estimated coefficients and intercept. + This value is only available when using the "normal" solver. + + If :py:attr:`LinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + + .. versionadded:: 2.0.0 + + See Also + -------- + LinearRegression.solver + """ + return self._call_java("tValues") + + @property + def pValues(self) -> List[float]: + """ + Two-sided p-value of estimated coefficients and intercept. + This value is only available when using the "normal" solver. + + If :py:attr:`LinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + + .. versionadded:: 2.0.0 + + See Also + -------- + LinearRegression.solver + """ + return self._call_java("pValues") + + +@inherit_doc +class LinearRegressionTrainingSummary(LinearRegressionSummary): + """ + Linear regression training results. Currently, the training summary ignores the + training weights except for the objective trace. + + .. versionadded:: 2.0.0 + """ + + @property + def objectiveHistory(self) -> List[float]: + """ + Objective function (scaled loss + regularization) at each + iteration. + This value is only available when using the "l-bfgs" solver. + + .. versionadded:: 2.0.0 + + See Also + -------- + LinearRegression.solver + """ + return self._call_java("objectiveHistory") + + @property + def totalIterations(self) -> int: + """ + Number of training iterations until termination. + This value is only available when using the "l-bfgs" solver. + + .. versionadded:: 2.0.0 + + See Also + -------- + LinearRegression.solver + """ + return self._call_java("totalIterations") + + +class _IsotonicRegressionParams(HasFeaturesCol, HasLabelCol, HasPredictionCol, HasWeightCol): + """ + Params for :py:class:`IsotonicRegression` and :py:class:`IsotonicRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + isotonic: Param[bool] = Param( + Params._dummy(), + "isotonic", + "whether the output sequence should be isotonic/increasing (true) or" + + "antitonic/decreasing (false).", + typeConverter=TypeConverters.toBoolean, + ) + featureIndex: Param[int] = Param( + Params._dummy(), + "featureIndex", + "The index of the feature if featuresCol is a vector column, no effect otherwise.", + typeConverter=TypeConverters.toInt, + ) + + def __init__(self, *args: Any): + super(_IsotonicRegressionParams, self).__init__(*args) + self._setDefault(isotonic=True, featureIndex=0) + + def getIsotonic(self) -> bool: + """ + Gets the value of isotonic or its default value. + """ + return self.getOrDefault(self.isotonic) + + def getFeatureIndex(self) -> int: + """ + Gets the value of featureIndex or its default value. + """ + return self.getOrDefault(self.featureIndex) + + +@inherit_doc +class IsotonicRegression( + JavaEstimator, _IsotonicRegressionParams, HasWeightCol, JavaMLWritable, JavaMLReadable +): + """ + Currently implemented using parallelized pool adjacent violators algorithm. + Only univariate (single feature) algorithm supported. + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> ir = IsotonicRegression() + >>> model = ir.fit(df) + >>> model.setFeaturesCol("features") + IsotonicRegressionModel... + >>> model.numFeatures + 1 + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.transform(test0).head().prediction + 0.0 + >>> model.predict(test0.head().features[model.getFeatureIndex()]) + 0.0 + >>> model.boundaries + DenseVector([0.0, 1.0]) + >>> ir_path = temp_path + "/ir" + >>> ir.save(ir_path) + >>> ir2 = IsotonicRegression.load(ir_path) + >>> ir2.getIsotonic() + True + >>> model_path = temp_path + "/ir_model" + >>> model.save(model_path) + >>> model2 = IsotonicRegressionModel.load(model_path) + >>> model.boundaries == model2.boundaries + True + >>> model.predictions == model2.predictions + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + weightCol: Optional[str] = None, + isotonic: bool = True, + featureIndex: int = 0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + weightCol=None, isotonic=True, featureIndex=0): + """ + super(IsotonicRegression, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.IsotonicRegression", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + weightCol: Optional[str] = None, + isotonic: bool = True, + featureIndex: int = 0, + ) -> "IsotonicRegression": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + weightCol=None, isotonic=True, featureIndex=0): + Set the params for IsotonicRegression. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "IsotonicRegressionModel": + return IsotonicRegressionModel(java_model) + + def setIsotonic(self, value: bool) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`isotonic`. + """ + return self._set(isotonic=value) + + def setFeatureIndex(self, value: int) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`featureIndex`. + """ + return self._set(featureIndex=value) + + @since("1.6.0") + def setFeaturesCol(self, value: str) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("1.6.0") + def setPredictionCol(self, value: str) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @since("1.6.0") + def setLabelCol(self, value: str) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + @since("1.6.0") + def setWeightCol(self, value: str) -> "IsotonicRegression": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + +class IsotonicRegressionModel( + JavaModel, + _IsotonicRegressionParams, + JavaMLWritable, + JavaMLReadable["IsotonicRegressionModel"], +): + """ + Model fitted by :class:`IsotonicRegression`. + + .. versionadded:: 1.6.0 + """ + + @since("3.0.0") + def setFeaturesCol(self, value: str) -> "IsotonicRegressionModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.0.0") + def setPredictionCol(self, value: str) -> "IsotonicRegressionModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + def setFeatureIndex(self, value: int) -> "IsotonicRegressionModel": + """ + Sets the value of :py:attr:`featureIndex`. + """ + return self._set(featureIndex=value) + + @property + @since("1.6.0") + def boundaries(self) -> Vector: + """ + Boundaries in increasing order for which predictions are known. + """ + return self._call_java("boundaries") + + @property + @since("1.6.0") + def predictions(self) -> Vector: + """ + Predictions associated with the boundaries at the same index, monotone because of isotonic + regression. + """ + return self._call_java("predictions") + + @property + @since("3.0.0") + def numFeatures(self) -> int: + """ + Returns the number of features the model was trained on. If unknown, returns -1 + """ + return self._call_java("numFeatures") + + @since("3.0.0") + def predict(self, value: float) -> float: + """ + Predict label for the given features. + """ + return self._call_java("predict", value) + + +class _DecisionTreeRegressorParams(_DecisionTreeParams, _TreeRegressorParams, HasVarianceCol): + """ + Params for :py:class:`DecisionTreeRegressor` and :py:class:`DecisionTreeRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + def __init__(self, *args: Any): + super(_DecisionTreeRegressorParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + checkpointInterval=10, + impurity="variance", + leafCol="", + minWeightFractionPerNode=0.0, + ) + + +@inherit_doc +class DecisionTreeRegressor( + _JavaRegressor["DecisionTreeRegressionModel"], + _DecisionTreeRegressorParams, + JavaMLWritable, + JavaMLReadable["DecisionTreeRegressor"], +): + """ + `Decision tree `_ + learning algorithm for regression. + It supports both continuous and categorical features. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> dt = DecisionTreeRegressor(maxDepth=2) + >>> dt.setVarianceCol("variance") + DecisionTreeRegressor... + >>> model = dt.fit(df) + >>> model.getVarianceCol() + 'variance' + >>> model.setLeafCol("leafId") + DecisionTreeRegressionModel... + >>> model.depth + 1 + >>> model.numNodes + 3 + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> model.numFeatures + 1 + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> model.predictLeaf(test0.head().features) + 0.0 + >>> result.leafId + 0.0 + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 1.0 + >>> dtr_path = temp_path + "/dtr" + >>> dt.save(dtr_path) + >>> dt2 = DecisionTreeRegressor.load(dtr_path) + >>> dt2.getMaxDepth() + 2 + >>> model_path = temp_path + "/dtr_model" + >>> model.save(model_path) + >>> model2 = DecisionTreeRegressionModel.load(model_path) + >>> model.numNodes == model2.numNodes + True + >>> model.depth == model2.depth + True + >>> model.transform(test1).head().variance + 0.0 + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> df3 = spark.createDataFrame([ + ... (1.0, 0.2, Vectors.dense(1.0)), + ... (1.0, 0.8, Vectors.dense(1.0)), + ... (0.0, 1.0, Vectors.sparse(1, [], []))], ["label", "weight", "features"]) + >>> dt3 = DecisionTreeRegressor(maxDepth=2, weightCol="weight", varianceCol="variance") + >>> model3 = dt3.fit(df3) + >>> print(model3.toDebugString) + DecisionTreeRegressionModel...depth=1, numNodes=3... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "variance", + seed: Optional[int] = None, + varianceCol: Optional[str] = None, + weightCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + impurity="variance", seed=None, varianceCol=None, weightCol=None, \ + leafCol="", minWeightFractionPerNode=0.0) + """ + super(DecisionTreeRegressor, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.DecisionTreeRegressor", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "variance", + seed: Optional[int] = None, + varianceCol: Optional[str] = None, + weightCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + ) -> "DecisionTreeRegressor": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + impurity="variance", seed=None, varianceCol=None, weightCol=None, \ + leafCol="", minWeightFractionPerNode=0.0) + Sets params for the DecisionTreeRegressor. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "DecisionTreeRegressionModel": + return DecisionTreeRegressionModel(java_model) + + @since("1.4.0") + def setMaxDepth(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + @since("1.4.0") + def setMaxBins(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + @since("1.4.0") + def setMinInstancesPerNode(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + @since("1.4.0") + def setMinInfoGain(self, value: float) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + @since("1.4.0") + def setMaxMemoryInMB(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + @since("1.4.0") + def setCacheNodeIds(self, value: bool) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setCheckpointInterval(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + def setSeed(self, value: int) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("2.0.0") + def setVarianceCol(self, value: str) -> "DecisionTreeRegressor": + """ + Sets the value of :py:attr:`varianceCol`. + """ + return self._set(varianceCol=value) + + +@inherit_doc +class DecisionTreeRegressionModel( + _JavaRegressionModel, + _DecisionTreeModel, + _DecisionTreeRegressorParams, + JavaMLWritable, + JavaMLReadable["DecisionTreeRegressionModel"], +): + """ + Model fitted by :class:`DecisionTreeRegressor`. + + .. versionadded:: 1.4.0 + """ + + @since("3.0.0") + def setVarianceCol(self, value: str) -> "DecisionTreeRegressionModel": + """ + Sets the value of :py:attr:`varianceCol`. + """ + return self._set(varianceCol=value) + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + This generalizes the idea of "Gini" importance to other losses, + following the explanation of Gini importance from "Random Forests" documentation + by Leo Breiman and Adele Cutler, and following the implementation from scikit-learn. + + This feature importance is calculated as follows: + - importance(feature j) = sum (over nodes which split on feature j) of the gain, + where gain is scaled by the number of instances passing through node + - Normalize importances for tree to sum to 1. + + .. versionadded:: 2.0.0 + + Notes + ----- + Feature importance for single decision trees can have high variance due to + correlated predictor variables. Consider using a :py:class:`RandomForestRegressor` + to determine feature importance instead. + """ + return self._call_java("featureImportances") + + +class _RandomForestRegressorParams(_RandomForestParams, _TreeRegressorParams): + """ + Params for :py:class:`RandomForestRegressor` and :py:class:`RandomForestRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + def __init__(self, *args: Any): + super(_RandomForestRegressorParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + checkpointInterval=10, + impurity="variance", + subsamplingRate=1.0, + numTrees=20, + featureSubsetStrategy="auto", + leafCol="", + minWeightFractionPerNode=0.0, + bootstrap=True, + ) + + +@inherit_doc +class RandomForestRegressor( + _JavaRegressor["RandomForestRegressionModel"], + _RandomForestRegressorParams, + JavaMLWritable, + JavaMLReadable["RandomForestRegressor"], +): + """ + `Random Forest `_ + learning algorithm for regression. + It supports both continuous and categorical features. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from numpy import allclose + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> rf = RandomForestRegressor(numTrees=2, maxDepth=2) + >>> rf.getMinWeightFractionPerNode() + 0.0 + >>> rf.setSeed(42) + RandomForestRegressor... + >>> model = rf.fit(df) + >>> model.getBootstrap() + True + >>> model.getSeed() + 42 + >>> model.setLeafCol("leafId") + RandomForestRegressionModel... + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> allclose(model.treeWeights, [1.0, 1.0]) + True + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> model.predictLeaf(test0.head().features) + DenseVector([0.0, 0.0]) + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> result.leafId + DenseVector([0.0, 0.0]) + >>> model.numFeatures + 1 + >>> model.trees + [DecisionTreeRegressionModel...depth=..., DecisionTreeRegressionModel...] + >>> model.getNumTrees + 2 + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 0.5 + >>> rfr_path = temp_path + "/rfr" + >>> rf.save(rfr_path) + >>> rf2 = RandomForestRegressor.load(rfr_path) + >>> rf2.getNumTrees() + 2 + >>> model_path = temp_path + "/rfr_model" + >>> model.save(model_path) + >>> model2 = RandomForestRegressionModel.load(model_path) + >>> model.featureImportances == model2.featureImportances + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "variance", + subsamplingRate: float = 1.0, + seed: Optional[int] = None, + numTrees: int = 20, + featureSubsetStrategy: str = "auto", + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + bootstrap: Optional[bool] = True, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + impurity="variance", subsamplingRate=1.0, seed=None, numTrees=20, \ + featureSubsetStrategy="auto", leafCol=", minWeightFractionPerNode=0.0", \ + weightCol=None, bootstrap=True) + """ + super(RandomForestRegressor, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.RandomForestRegressor", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + checkpointInterval: int = 10, + impurity: str = "variance", + subsamplingRate: float = 1.0, + seed: Optional[int] = None, + numTrees: int = 20, + featureSubsetStrategy: str = "auto", + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + bootstrap: Optional[bool] = True, + ) -> "RandomForestRegressor": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10, \ + impurity="variance", subsamplingRate=1.0, seed=None, numTrees=20, \ + featureSubsetStrategy="auto", leafCol="", minWeightFractionPerNode=0.0, \ + weightCol=None, bootstrap=True) + Sets params for linear regression. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "RandomForestRegressionModel": + return RandomForestRegressionModel(java_model) + + def setMaxDepth(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + def setMaxBins(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + def setMinInstancesPerNode(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + def setMinInfoGain(self, value: float) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + def setMaxMemoryInMB(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + def setCacheNodeIds(self, value: bool) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setNumTrees(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`numTrees`. + """ + return self._set(numTrees=value) + + @since("3.0.0") + def setBootstrap(self, value: bool) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`bootstrap`. + """ + return self._set(bootstrap=value) + + @since("1.4.0") + def setSubsamplingRate(self, value: float) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`subsamplingRate`. + """ + return self._set(subsamplingRate=value) + + @since("2.4.0") + def setFeatureSubsetStrategy(self, value: str) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`featureSubsetStrategy`. + """ + return self._set(featureSubsetStrategy=value) + + def setCheckpointInterval(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + def setSeed(self, value: int) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + +class RandomForestRegressionModel( + _JavaRegressionModel[Vector], + _TreeEnsembleModel, + _RandomForestRegressorParams, + JavaMLWritable, + JavaMLReadable["RandomForestRegressionModel"], +): + """ + Model fitted by :class:`RandomForestRegressor`. + + .. versionadded:: 1.4.0 + """ + + @property + @since("2.0.0") + def trees(self) -> List[DecisionTreeRegressionModel]: + """Trees in this ensemble. Warning: These have null parent Estimators.""" + return [DecisionTreeRegressionModel(m) for m in list(self._call_java("trees"))] + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + Each feature's importance is the average of its importance across all trees in the ensemble + The importance vector is normalized to sum to 1. This method is suggested by Hastie et al. + (Hastie, Tibshirani, Friedman. "The Elements of Statistical Learning, 2nd Edition." 2001.) + and follows the implementation from scikit-learn. + + .. versionadded:: 2.0.0 + + Examples + -------- + DecisionTreeRegressionModel.featureImportances + """ + return self._call_java("featureImportances") + + +class _GBTRegressorParams(_GBTParams, _TreeRegressorParams): + """ + Params for :py:class:`GBTRegressor` and :py:class:`GBTRegressorModel`. + + .. versionadded:: 3.0.0 + """ + + supportedLossTypes: List[str] = ["squared", "absolute"] + + lossType: Param[str] = Param( + Params._dummy(), + "lossType", + "Loss function which GBT tries to minimize (case-insensitive). " + + "Supported options: " + + ", ".join(supportedLossTypes), + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_GBTRegressorParams, self).__init__(*args) + self._setDefault( + maxDepth=5, + maxBins=32, + minInstancesPerNode=1, + minInfoGain=0.0, + maxMemoryInMB=256, + cacheNodeIds=False, + subsamplingRate=1.0, + checkpointInterval=10, + lossType="squared", + maxIter=20, + stepSize=0.1, + impurity="variance", + featureSubsetStrategy="all", + validationTol=0.01, + leafCol="", + minWeightFractionPerNode=0.0, + ) + + @since("1.4.0") + def getLossType(self) -> str: + """ + Gets the value of lossType or its default value. + """ + return self.getOrDefault(self.lossType) + + +@inherit_doc +class GBTRegressor( + _JavaRegressor["GBTRegressionModel"], + _GBTRegressorParams, + JavaMLWritable, + JavaMLReadable["GBTRegressor"], +): + """ + `Gradient-Boosted Trees (GBTs) `_ + learning algorithm for regression. + It supports both continuous and categorical features. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from numpy import allclose + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> gbt = GBTRegressor(maxDepth=2, seed=42, leafCol="leafId") + >>> gbt.setMaxIter(5) + GBTRegressor... + >>> gbt.setMinWeightFractionPerNode(0.049) + GBTRegressor... + >>> gbt.getMaxIter() + 5 + >>> print(gbt.getImpurity()) + variance + >>> print(gbt.getFeatureSubsetStrategy()) + all + >>> model = gbt.fit(df) + >>> model.featureImportances + SparseVector(1, {0: 1.0}) + >>> model.numFeatures + 1 + >>> allclose(model.treeWeights, [1.0, 0.1, 0.1, 0.1, 0.1]) + True + >>> test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"]) + >>> model.predict(test0.head().features) + 0.0 + >>> model.predictLeaf(test0.head().features) + DenseVector([0.0, 0.0, 0.0, 0.0, 0.0]) + >>> result = model.transform(test0).head() + >>> result.prediction + 0.0 + >>> result.leafId + DenseVector([0.0, 0.0, 0.0, 0.0, 0.0]) + >>> test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"]) + >>> model.transform(test1).head().prediction + 1.0 + >>> gbtr_path = temp_path + "gbtr" + >>> gbt.save(gbtr_path) + >>> gbt2 = GBTRegressor.load(gbtr_path) + >>> gbt2.getMaxDepth() + 2 + >>> model_path = temp_path + "gbtr_model" + >>> model.save(model_path) + >>> model2 = GBTRegressionModel.load(model_path) + >>> model.featureImportances == model2.featureImportances + True + >>> model.treeWeights == model2.treeWeights + True + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + >>> model.trees + [DecisionTreeRegressionModel...depth=..., DecisionTreeRegressionModel...] + >>> validation = spark.createDataFrame([(0.0, Vectors.dense(-1.0))], + ... ["label", "features"]) + >>> model.evaluateEachIteration(validation, "squared") + [0.0, 0.0, 0.0, 0.0, 0.0] + >>> gbt = gbt.setValidationIndicatorCol("validationIndicator") + >>> gbt.getValidationIndicatorCol() + 'validationIndicator' + >>> gbt.getValidationTol() + 0.01 + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + subsamplingRate: float = 1.0, + checkpointInterval: int = 10, + lossType: str = "squared", + maxIter: int = 20, + stepSize: float = 0.1, + seed: Optional[int] = None, + impurity: str = "variance", + featureSubsetStrategy: str = "all", + validationTol: float = 0.1, + validationIndicatorCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, subsamplingRate=1.0, \ + checkpointInterval=10, lossType="squared", maxIter=20, stepSize=0.1, seed=None, \ + impurity="variance", featureSubsetStrategy="all", validationTol=0.01, \ + validationIndicatorCol=None, leafCol="", minWeightFractionPerNode=0.0, + weightCol=None) + """ + super(GBTRegressor, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.regression.GBTRegressor", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + maxMemoryInMB: int = 256, + cacheNodeIds: bool = False, + subsamplingRate: float = 1.0, + checkpointInterval: int = 10, + lossType: str = "squared", + maxIter: int = 20, + stepSize: float = 0.1, + seed: Optional[int] = None, + impurity: str = "variance", + featureSubsetStrategy: str = "all", + validationTol: float = 0.1, + validationIndicatorCol: Optional[str] = None, + leafCol: str = "", + minWeightFractionPerNode: float = 0.0, + weightCol: Optional[str] = None, + ) -> "GBTRegressor": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0, \ + maxMemoryInMB=256, cacheNodeIds=False, subsamplingRate=1.0, \ + checkpointInterval=10, lossType="squared", maxIter=20, stepSize=0.1, seed=None, \ + impurity="variance", featureSubsetStrategy="all", validationTol=0.01, \ + validationIndicatorCol=None, leafCol="", minWeightFractionPerNode=0.0, \ + weightCol=None) + Sets params for Gradient Boosted Tree Regression. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "GBTRegressionModel": + return GBTRegressionModel(java_model) + + @since("1.4.0") + def setMaxDepth(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`maxDepth`. + """ + return self._set(maxDepth=value) + + @since("1.4.0") + def setMaxBins(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`maxBins`. + """ + return self._set(maxBins=value) + + @since("1.4.0") + def setMinInstancesPerNode(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`minInstancesPerNode`. + """ + return self._set(minInstancesPerNode=value) + + @since("1.4.0") + def setMinInfoGain(self, value: float) -> "GBTRegressor": + """ + Sets the value of :py:attr:`minInfoGain`. + """ + return self._set(minInfoGain=value) + + @since("1.4.0") + def setMaxMemoryInMB(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`maxMemoryInMB`. + """ + return self._set(maxMemoryInMB=value) + + @since("1.4.0") + def setCacheNodeIds(self, value: bool) -> "GBTRegressor": + """ + Sets the value of :py:attr:`cacheNodeIds`. + """ + return self._set(cacheNodeIds=value) + + @since("1.4.0") + def setImpurity(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`impurity`. + """ + return self._set(impurity=value) + + @since("1.4.0") + def setLossType(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`lossType`. + """ + return self._set(lossType=value) + + @since("1.4.0") + def setSubsamplingRate(self, value: float) -> "GBTRegressor": + """ + Sets the value of :py:attr:`subsamplingRate`. + """ + return self._set(subsamplingRate=value) + + @since("2.4.0") + def setFeatureSubsetStrategy(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`featureSubsetStrategy`. + """ + return self._set(featureSubsetStrategy=value) + + @since("3.0.0") + def setValidationIndicatorCol(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`validationIndicatorCol`. + """ + return self._set(validationIndicatorCol=value) + + @since("1.4.0") + def setMaxIter(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("1.4.0") + def setCheckpointInterval(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`checkpointInterval`. + """ + return self._set(checkpointInterval=value) + + @since("1.4.0") + def setSeed(self, value: int) -> "GBTRegressor": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("1.4.0") + def setStepSize(self, value: float) -> "GBTRegressor": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + @since("3.0.0") + def setWeightCol(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("3.0.0") + def setMinWeightFractionPerNode(self, value: float) -> "GBTRegressor": + """ + Sets the value of :py:attr:`minWeightFractionPerNode`. + """ + return self._set(minWeightFractionPerNode=value) + + +class GBTRegressionModel( + _JavaRegressionModel[Vector], + _TreeEnsembleModel, + _GBTRegressorParams, + JavaMLWritable, + JavaMLReadable["GBTRegressionModel"], +): + """ + Model fitted by :class:`GBTRegressor`. + + .. versionadded:: 1.4.0 + """ + + @property + def featureImportances(self) -> Vector: + """ + Estimate of the importance of each feature. + + Each feature's importance is the average of its importance across all trees in the ensemble + The importance vector is normalized to sum to 1. This method is suggested by Hastie et al. + (Hastie, Tibshirani, Friedman. "The Elements of Statistical Learning, 2nd Edition." 2001.) + and follows the implementation from scikit-learn. + + .. versionadded:: 2.0.0 + + Examples + -------- + DecisionTreeRegressionModel.featureImportances + """ + return self._call_java("featureImportances") + + @property + @since("2.0.0") + def trees(self) -> List[DecisionTreeRegressionModel]: + """Trees in this ensemble. Warning: These have null parent Estimators.""" + return [DecisionTreeRegressionModel(m) for m in list(self._call_java("trees"))] + + def evaluateEachIteration(self, dataset: DataFrame, loss: str) -> List[float]: + """ + Method to compute error or loss for every iteration of gradient boosting. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on, where dataset is an + instance of :py:class:`pyspark.sql.DataFrame` + loss : str + The loss function used to compute error. + Supported options: squared, absolute + """ + return self._call_java("evaluateEachIteration", dataset, loss) + + +class _AFTSurvivalRegressionParams( + _PredictorParams, HasMaxIter, HasTol, HasFitIntercept, HasAggregationDepth, HasMaxBlockSizeInMB +): + """ + Params for :py:class:`AFTSurvivalRegression` and :py:class:`AFTSurvivalRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + censorCol: Param[str] = Param( + Params._dummy(), + "censorCol", + "censor column name. The value of this column could be 0 or 1. " + + "If the value is 1, it means the event has occurred i.e. " + + "uncensored; otherwise censored.", + typeConverter=TypeConverters.toString, + ) + quantileProbabilities: Param[List[float]] = Param( + Params._dummy(), + "quantileProbabilities", + "quantile probabilities array. Values of the quantile probabilities array " + + "should be in the range (0, 1) and the array should be non-empty.", + typeConverter=TypeConverters.toListFloat, + ) + quantilesCol: Param[str] = Param( + Params._dummy(), + "quantilesCol", + "quantiles column name. This column will output quantiles of " + + "corresponding quantileProbabilities if it is set.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_AFTSurvivalRegressionParams, self).__init__(*args) + self._setDefault( + censorCol="censor", + quantileProbabilities=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99], + maxIter=100, + tol=1e-6, + maxBlockSizeInMB=0.0, + ) + + @since("1.6.0") + def getCensorCol(self) -> str: + """ + Gets the value of censorCol or its default value. + """ + return self.getOrDefault(self.censorCol) + + @since("1.6.0") + def getQuantileProbabilities(self) -> List[float]: + """ + Gets the value of quantileProbabilities or its default value. + """ + return self.getOrDefault(self.quantileProbabilities) + + @since("1.6.0") + def getQuantilesCol(self) -> str: + """ + Gets the value of quantilesCol or its default value. + """ + return self.getOrDefault(self.quantilesCol) + + +@inherit_doc +class AFTSurvivalRegression( + _JavaRegressor["AFTSurvivalRegressionModel"], + _AFTSurvivalRegressionParams, + JavaMLWritable, + JavaMLReadable["AFTSurvivalRegression"], +): + """ + Accelerated Failure Time (AFT) Model Survival Regression + + Fit a parametric AFT survival regression model based on the Weibull distribution + of the survival time. + + Notes + ----- + For more information see Wikipedia page on + `AFT Model `_ + + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(1.0), 1.0), + ... (1e-40, Vectors.sparse(1, [], []), 0.0)], ["label", "features", "censor"]) + >>> aftsr = AFTSurvivalRegression() + >>> aftsr.setMaxIter(10) + AFTSurvivalRegression... + >>> aftsr.getMaxIter() + 10 + >>> aftsr.clear(aftsr.maxIter) + >>> model = aftsr.fit(df) + >>> model.getMaxBlockSizeInMB() + 0.0 + >>> model.setFeaturesCol("features") + AFTSurvivalRegressionModel... + >>> model.predict(Vectors.dense(6.3)) + 1.0 + >>> model.predictQuantiles(Vectors.dense(6.3)) + DenseVector([0.0101, 0.0513, 0.1054, 0.2877, 0.6931, 1.3863, 2.3026, 2.9957, 4.6052]) + >>> model.transform(df).show() + +-------+---------+------+----------+ + | label| features|censor|prediction| + +-------+---------+------+----------+ + | 1.0| [1.0]| 1.0| 1.0| + |1.0E-40|(1,[],[])| 0.0| 1.0| + +-------+---------+------+----------+ + ... + >>> aftsr_path = temp_path + "/aftsr" + >>> aftsr.save(aftsr_path) + >>> aftsr2 = AFTSurvivalRegression.load(aftsr_path) + >>> aftsr2.getMaxIter() + 100 + >>> model_path = temp_path + "/aftsr_model" + >>> model.save(model_path) + >>> model2 = AFTSurvivalRegressionModel.load(model_path) + >>> model.coefficients == model2.coefficients + True + >>> model.intercept == model2.intercept + True + >>> model.scale == model2.scale + True + >>> model.transform(df).take(1) == model2.transform(df).take(1) + True + + .. versionadded:: 1.6.0 + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + fitIntercept: bool = True, + maxIter: int = 100, + tol: float = 1e-6, + censorCol: str = "censor", + quantileProbabilities: List[float] = [ + 0.01, + 0.05, + 0.1, + 0.25, + 0.5, + 0.75, + 0.9, + 0.95, + 0.99, + ], # noqa: B005 + quantilesCol: Optional[str] = None, + aggregationDepth: int = 2, + maxBlockSizeInMB: float = 0.0, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + fitIntercept=True, maxIter=100, tol=1E-6, censorCol="censor", \ + quantileProbabilities=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99], \ + quantilesCol=None, aggregationDepth=2, maxBlockSizeInMB=0.0) + """ + super(AFTSurvivalRegression, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.AFTSurvivalRegression", self.uid + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("1.6.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + fitIntercept: bool = True, + maxIter: int = 100, + tol: float = 1e-6, + censorCol: str = "censor", + quantileProbabilities: List[float] = [ + 0.01, + 0.05, + 0.1, + 0.25, + 0.5, + 0.75, + 0.9, + 0.95, + 0.99, + ], # noqa: B005 + quantilesCol: Optional[str] = None, + aggregationDepth: int = 2, + maxBlockSizeInMB: float = 0.0, + ) -> "AFTSurvivalRegression": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + fitIntercept=True, maxIter=100, tol=1E-6, censorCol="censor", \ + quantileProbabilities=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99], \ + quantilesCol=None, aggregationDepth=2, maxBlockSizeInMB=0.0): + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "AFTSurvivalRegressionModel": + return AFTSurvivalRegressionModel(java_model) + + @since("1.6.0") + def setCensorCol(self, value: str) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`censorCol`. + """ + return self._set(censorCol=value) + + @since("1.6.0") + def setQuantileProbabilities(self, value: List[float]) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`quantileProbabilities`. + """ + return self._set(quantileProbabilities=value) + + @since("1.6.0") + def setQuantilesCol(self, value: str) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`quantilesCol`. + """ + return self._set(quantilesCol=value) + + @since("1.6.0") + def setMaxIter(self, value: int) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("1.6.0") + def setTol(self, value: float) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("1.6.0") + def setFitIntercept(self, value: bool) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + @since("2.1.0") + def setAggregationDepth(self, value: int) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + @since("3.1.0") + def setMaxBlockSizeInMB(self, value: int) -> "AFTSurvivalRegression": + """ + Sets the value of :py:attr:`maxBlockSizeInMB`. + """ + return self._set(maxBlockSizeInMB=value) + + +class AFTSurvivalRegressionModel( + _JavaRegressionModel[Vector], + _AFTSurvivalRegressionParams, + JavaMLWritable, + JavaMLReadable["AFTSurvivalRegressionModel"], +): + """ + Model fitted by :class:`AFTSurvivalRegression`. + + .. versionadded:: 1.6.0 + """ + + @since("3.0.0") + def setQuantileProbabilities(self, value: List[float]) -> "AFTSurvivalRegressionModel": + """ + Sets the value of :py:attr:`quantileProbabilities`. + """ + return self._set(quantileProbabilities=value) + + @since("3.0.0") + def setQuantilesCol(self, value: str) -> "AFTSurvivalRegressionModel": + """ + Sets the value of :py:attr:`quantilesCol`. + """ + return self._set(quantilesCol=value) + + @property + @since("2.0.0") + def coefficients(self) -> Vector: + """ + Model coefficients. + """ + return self._call_java("coefficients") + + @property + @since("1.6.0") + def intercept(self) -> float: + """ + Model intercept. + """ + return self._call_java("intercept") + + @property + @since("1.6.0") + def scale(self) -> float: + """ + Model scale parameter. + """ + return self._call_java("scale") + + @since("2.0.0") + def predictQuantiles(self, features: Vector) -> Vector: + """ + Predicted Quantiles + """ + return self._call_java("predictQuantiles", features) + + +class _GeneralizedLinearRegressionParams( + _PredictorParams, + HasFitIntercept, + HasMaxIter, + HasTol, + HasRegParam, + HasWeightCol, + HasSolver, + HasAggregationDepth, +): + """ + Params for :py:class:`GeneralizedLinearRegression` and + :py:class:`GeneralizedLinearRegressionModel`. + + .. versionadded:: 3.0.0 + """ + + family: Param[str] = Param( + Params._dummy(), + "family", + "The name of family which is a description of " + + "the error distribution to be used in the model. Supported options: " + + "gaussian (default), binomial, poisson, gamma and tweedie.", + typeConverter=TypeConverters.toString, + ) + link: Param[str] = Param( + Params._dummy(), + "link", + "The name of link function which provides the " + + "relationship between the linear predictor and the mean of the distribution " + + "function. Supported options: identity, log, inverse, logit, probit, cloglog " + + "and sqrt.", + typeConverter=TypeConverters.toString, + ) + linkPredictionCol: Param[str] = Param( + Params._dummy(), + "linkPredictionCol", + "link prediction (linear " + "predictor) column name", + typeConverter=TypeConverters.toString, + ) + variancePower: Param[float] = Param( + Params._dummy(), + "variancePower", + "The power in the variance function " + + "of the Tweedie distribution which characterizes the relationship " + + "between the variance and mean of the distribution. Only applicable " + + "for the Tweedie family. Supported values: 0 and [1, Inf).", + typeConverter=TypeConverters.toFloat, + ) + linkPower: Param[float] = Param( + Params._dummy(), + "linkPower", + "The index in the power link function. " + "Only applicable to the Tweedie family.", + typeConverter=TypeConverters.toFloat, + ) + solver: Param[str] = Param( + Params._dummy(), + "solver", + "The solver algorithm for optimization. Supported " + "options: irls.", + typeConverter=TypeConverters.toString, + ) + offsetCol: Param[str] = Param( + Params._dummy(), + "offsetCol", + "The offset column name. If this is not set " + + "or empty, we treat all instance offsets as 0.0", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_GeneralizedLinearRegressionParams, self).__init__(*args) + self._setDefault( + family="gaussian", + maxIter=25, + tol=1e-6, + regParam=0.0, + solver="irls", + variancePower=0.0, + aggregationDepth=2, + ) + + @since("2.0.0") + def getFamily(self) -> str: + """ + Gets the value of family or its default value. + """ + return self.getOrDefault(self.family) + + @since("2.0.0") + def getLinkPredictionCol(self) -> str: + """ + Gets the value of linkPredictionCol or its default value. + """ + return self.getOrDefault(self.linkPredictionCol) + + @since("2.0.0") + def getLink(self) -> str: + """ + Gets the value of link or its default value. + """ + return self.getOrDefault(self.link) + + @since("2.2.0") + def getVariancePower(self) -> float: + """ + Gets the value of variancePower or its default value. + """ + return self.getOrDefault(self.variancePower) + + @since("2.2.0") + def getLinkPower(self) -> float: + """ + Gets the value of linkPower or its default value. + """ + return self.getOrDefault(self.linkPower) + + @since("2.3.0") + def getOffsetCol(self) -> str: + """ + Gets the value of offsetCol or its default value. + """ + return self.getOrDefault(self.offsetCol) + + +@inherit_doc +class GeneralizedLinearRegression( + _JavaRegressor["GeneralizedLinearRegressionModel"], + _GeneralizedLinearRegressionParams, + JavaMLWritable, + JavaMLReadable["GeneralizedLinearRegression"], +): + """ + Generalized Linear Regression. + + Fit a Generalized Linear Model specified by giving a symbolic description of the linear + predictor (link function) and a description of the error distribution (family). It supports + "gaussian", "binomial", "poisson", "gamma" and "tweedie" as family. Valid link functions for + each family is listed below. The first link function of each family is the default one. + + * "gaussian" -> "identity", "log", "inverse" + + * "binomial" -> "logit", "probit", "cloglog" + + * "poisson" -> "log", "identity", "sqrt" + + * "gamma" -> "inverse", "identity", "log" + + * "tweedie" -> power link function specified through "linkPower". \ + The default link power in the tweedie family is 1 - variancePower. + + .. versionadded:: 2.0.0 + + Notes + ----- + For more information see Wikipedia page on + `GLM `_ + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> df = spark.createDataFrame([ + ... (1.0, Vectors.dense(0.0, 0.0)), + ... (1.0, Vectors.dense(1.0, 2.0)), + ... (2.0, Vectors.dense(0.0, 0.0)), + ... (2.0, Vectors.dense(1.0, 1.0)),], ["label", "features"]) + >>> glr = GeneralizedLinearRegression(family="gaussian", link="identity", linkPredictionCol="p") + >>> glr.setRegParam(0.1) + GeneralizedLinearRegression... + >>> glr.getRegParam() + 0.1 + >>> glr.clear(glr.regParam) + >>> glr.setMaxIter(10) + GeneralizedLinearRegression... + >>> glr.getMaxIter() + 10 + >>> glr.clear(glr.maxIter) + >>> model = glr.fit(df) + >>> model.setFeaturesCol("features") + GeneralizedLinearRegressionModel... + >>> model.getMaxIter() + 25 + >>> model.getAggregationDepth() + 2 + >>> transformed = model.transform(df) + >>> abs(transformed.head().prediction - 1.5) < 0.001 + True + >>> abs(transformed.head().p - 1.5) < 0.001 + True + >>> model.coefficients + DenseVector([1.5..., -1.0...]) + >>> model.numFeatures + 2 + >>> abs(model.intercept - 1.5) < 0.001 + True + >>> glr_path = temp_path + "/glr" + >>> glr.save(glr_path) + >>> glr2 = GeneralizedLinearRegression.load(glr_path) + >>> glr.getFamily() == glr2.getFamily() + True + >>> model_path = temp_path + "/glr_model" + >>> model.save(model_path) + >>> model2 = GeneralizedLinearRegressionModel.load(model_path) + >>> model.intercept == model2.intercept + True + >>> model.coefficients[0] == model2.coefficients[0] + True + >>> model.transform(df).take(1) == model2.transform(df).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + labelCol: str = "label", + featuresCol: str = "features", + predictionCol: str = "prediction", + family: str = "gaussian", + link: Optional[str] = None, + fitIntercept: bool = True, + maxIter: int = 25, + tol: float = 1e-6, + regParam: float = 0.0, + weightCol: Optional[str] = None, + solver: str = "irls", + linkPredictionCol: Optional[str] = None, + variancePower: float = 0.0, + linkPower: Optional[float] = None, + offsetCol: Optional[str] = None, + aggregationDepth: int = 2, + ): + """ + __init__(self, \\*, labelCol="label", featuresCol="features", predictionCol="prediction", \ + family="gaussian", link=None, fitIntercept=True, maxIter=25, tol=1e-6, \ + regParam=0.0, weightCol=None, solver="irls", linkPredictionCol=None, \ + variancePower=0.0, linkPower=None, offsetCol=None, aggregationDepth=2) + """ + super(GeneralizedLinearRegression, self).__init__() + self._java_obj = self._new_java_obj( + "org.apache.spark.ml.regression.GeneralizedLinearRegression", self.uid + ) + kwargs = self._input_kwargs + + self.setParams(**kwargs) + + @keyword_only + @since("2.0.0") + def setParams( + self, + *, + labelCol: str = "label", + featuresCol: str = "features", + predictionCol: str = "prediction", + family: str = "gaussian", + link: Optional[str] = None, + fitIntercept: bool = True, + maxIter: int = 25, + tol: float = 1e-6, + regParam: float = 0.0, + weightCol: Optional[str] = None, + solver: str = "irls", + linkPredictionCol: Optional[str] = None, + variancePower: float = 0.0, + linkPower: Optional[float] = None, + offsetCol: Optional[str] = None, + aggregationDepth: int = 2, + ) -> "GeneralizedLinearRegression": + """ + setParams(self, \\*, labelCol="label", featuresCol="features", predictionCol="prediction", \ + family="gaussian", link=None, fitIntercept=True, maxIter=25, tol=1e-6, \ + regParam=0.0, weightCol=None, solver="irls", linkPredictionCol=None, \ + variancePower=0.0, linkPower=None, offsetCol=None, aggregationDepth=2) + Sets params for generalized linear regression. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "GeneralizedLinearRegressionModel": + return GeneralizedLinearRegressionModel(java_model) + + @since("2.0.0") + def setFamily(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`family`. + """ + return self._set(family=value) + + @since("2.0.0") + def setLinkPredictionCol(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`linkPredictionCol`. + """ + return self._set(linkPredictionCol=value) + + @since("2.0.0") + def setLink(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`link`. + """ + return self._set(link=value) + + @since("2.2.0") + def setVariancePower(self, value: float) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`variancePower`. + """ + return self._set(variancePower=value) + + @since("2.2.0") + def setLinkPower(self, value: float) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`linkPower`. + """ + return self._set(linkPower=value) + + @since("2.3.0") + def setOffsetCol(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`offsetCol`. + """ + return self._set(offsetCol=value) + + @since("2.0.0") + def setMaxIter(self, value: int) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("2.0.0") + def setRegParam(self, value: float) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + @since("2.0.0") + def setTol(self, value: float) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("2.0.0") + def setFitIntercept(self, value: bool) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + @since("2.0.0") + def setWeightCol(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`weightCol`. + """ + return self._set(weightCol=value) + + @since("2.0.0") + def setSolver(self, value: str) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + @since("3.0.0") + def setAggregationDepth(self, value: int) -> "GeneralizedLinearRegression": + """ + Sets the value of :py:attr:`aggregationDepth`. + """ + return self._set(aggregationDepth=value) + + +class GeneralizedLinearRegressionModel( + _JavaRegressionModel[Vector], + _GeneralizedLinearRegressionParams, + JavaMLWritable, + JavaMLReadable["GeneralizedLinearRegressionModel"], + HasTrainingSummary["GeneralizedLinearRegressionTrainingSummary"], +): + """ + Model fitted by :class:`GeneralizedLinearRegression`. + + .. versionadded:: 2.0.0 + """ + + @since("3.0.0") + def setLinkPredictionCol(self, value: str) -> "GeneralizedLinearRegressionModel": + """ + Sets the value of :py:attr:`linkPredictionCol`. + """ + return self._set(linkPredictionCol=value) + + @property + @since("2.0.0") + def coefficients(self) -> Vector: + """ + Model coefficients. + """ + return self._call_java("coefficients") + + @property + @since("2.0.0") + def intercept(self) -> float: + """ + Model intercept. + """ + return self._call_java("intercept") + + @property + @since("2.0.0") + def summary(self) -> "GeneralizedLinearRegressionTrainingSummary": + """ + Gets summary (residuals, deviance, p-values) of model on + training set. An exception is thrown if + `trainingSummary is None`. + """ + if self.hasSummary: + return GeneralizedLinearRegressionTrainingSummary( + super(GeneralizedLinearRegressionModel, self).summary + ) + else: + raise RuntimeError( + "No training summary available for this %s" % self.__class__.__name__ + ) + + def evaluate(self, dataset: DataFrame) -> "GeneralizedLinearRegressionSummary": + """ + Evaluates the model on a test dataset. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + Test dataset to evaluate model on, where dataset is an + instance of :py:class:`pyspark.sql.DataFrame` + """ + if not isinstance(dataset, DataFrame): + raise TypeError("dataset must be a DataFrame but got %s." % type(dataset)) + java_glr_summary = self._call_java("evaluate", dataset) + return GeneralizedLinearRegressionSummary(java_glr_summary) + + +class GeneralizedLinearRegressionSummary(JavaWrapper): + """ + Generalized linear regression results evaluated on a dataset. + + .. versionadded:: 2.0.0 + """ + + @property + @since("2.0.0") + def predictions(self) -> DataFrame: + """ + Predictions output by the model's `transform` method. + """ + return self._call_java("predictions") + + @property + @since("2.0.0") + def predictionCol(self) -> str: + """ + Field in :py:attr:`predictions` which gives the predicted value of each instance. + This is set to a new column name if the original model's `predictionCol` is not set. + """ + return self._call_java("predictionCol") + + @property + @since("2.2.0") + def numInstances(self) -> int: + """ + Number of instances in DataFrame predictions. + """ + return self._call_java("numInstances") + + @property + @since("2.0.0") + def rank(self) -> int: + """ + The numeric rank of the fitted linear model. + """ + return self._call_java("rank") + + @property + @since("2.0.0") + def degreesOfFreedom(self) -> int: + """ + Degrees of freedom. + """ + return self._call_java("degreesOfFreedom") + + @property + @since("2.0.0") + def residualDegreeOfFreedom(self) -> int: + """ + The residual degrees of freedom. + """ + return self._call_java("residualDegreeOfFreedom") + + @property + @since("2.0.0") + def residualDegreeOfFreedomNull(self) -> int: + """ + The residual degrees of freedom for the null model. + """ + return self._call_java("residualDegreeOfFreedomNull") + + def residuals(self, residualsType: str = "deviance") -> DataFrame: + """ + Get the residuals of the fitted model by type. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + residualsType : str, optional + The type of residuals which should be returned. + Supported options: deviance (default), pearson, working, and response. + """ + return self._call_java("residuals", residualsType) + + @property + @since("2.0.0") + def nullDeviance(self) -> float: + """ + The deviance for the null model. + """ + return self._call_java("nullDeviance") + + @property + @since("2.0.0") + def deviance(self) -> float: + """ + The deviance for the fitted model. + """ + return self._call_java("deviance") + + @property + @since("2.0.0") + def dispersion(self) -> float: + """ + The dispersion of the fitted model. + It is taken as 1.0 for the "binomial" and "poisson" families, and otherwise + estimated by the residual Pearson's Chi-Squared statistic (which is defined as + sum of the squares of the Pearson residuals) divided by the residual degrees of freedom. + """ + return self._call_java("dispersion") + + @property + @since("2.0.0") + def aic(self) -> float: + """ + Akaike's "An Information Criterion"(AIC) for the fitted model. + """ + return self._call_java("aic") + + +@inherit_doc +class GeneralizedLinearRegressionTrainingSummary(GeneralizedLinearRegressionSummary): + """ + Generalized linear regression training results. + + .. versionadded:: 2.0.0 + """ + + @property + @since("2.0.0") + def numIterations(self) -> int: + """ + Number of training iterations. + """ + return self._call_java("numIterations") + + @property + @since("2.0.0") + def solver(self) -> str: + """ + The numeric solver used for training. + """ + return self._call_java("solver") + + @property + @since("2.0.0") + def coefficientStandardErrors(self) -> List[float]: + """ + Standard error of estimated coefficients and intercept. + + If :py:attr:`GeneralizedLinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + """ + return self._call_java("coefficientStandardErrors") + + @property + @since("2.0.0") + def tValues(self) -> List[float]: + """ + T-statistic of estimated coefficients and intercept. + + If :py:attr:`GeneralizedLinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + """ + return self._call_java("tValues") + + @property + @since("2.0.0") + def pValues(self) -> List[float]: + """ + Two-sided p-value of estimated coefficients and intercept. + + If :py:attr:`GeneralizedLinearRegression.fitIntercept` is set to True, + then the last element returned corresponds to the intercept. + """ + return self._call_java("pValues") + + def __repr__(self) -> str: + return self._call_java("toString") + + +class _FactorizationMachinesParams( + _PredictorParams, + HasMaxIter, + HasStepSize, + HasTol, + HasSolver, + HasSeed, + HasFitIntercept, + HasRegParam, + HasWeightCol, +): + """ + Params for :py:class:`FMRegressor`, :py:class:`FMRegressionModel`, :py:class:`FMClassifier` + and :py:class:`FMClassifierModel`. + + .. versionadded:: 3.0.0 + """ + + factorSize: Param[int] = Param( + Params._dummy(), + "factorSize", + "Dimensionality of the factor vectors, " + + "which are used to get pairwise interactions between variables", + typeConverter=TypeConverters.toInt, + ) + + fitLinear: Param[bool] = Param( + Params._dummy(), + "fitLinear", + "whether to fit linear term (aka 1-way term)", + typeConverter=TypeConverters.toBoolean, + ) + + miniBatchFraction: Param[float] = Param( + Params._dummy(), + "miniBatchFraction", + "fraction of the input data " + + "set that should be used for one iteration of gradient descent", + typeConverter=TypeConverters.toFloat, + ) + + initStd: Param[float] = Param( + Params._dummy(), + "initStd", + "standard deviation of initial coefficients", + typeConverter=TypeConverters.toFloat, + ) + + solver = Param( + Params._dummy(), + "solver", + "The solver algorithm for optimization. Supported " + "options: gd, adamW. (Default adamW)", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_FactorizationMachinesParams, self).__init__(*args) + self._setDefault( + factorSize=8, + fitIntercept=True, + fitLinear=True, + regParam=0.0, + miniBatchFraction=1.0, + initStd=0.01, + maxIter=100, + stepSize=1.0, + tol=1e-6, + solver="adamW", + ) + + @since("3.0.0") + def getFactorSize(self) -> int: + """ + Gets the value of factorSize or its default value. + """ + return self.getOrDefault(self.factorSize) + + @since("3.0.0") + def getFitLinear(self) -> bool: + """ + Gets the value of fitLinear or its default value. + """ + return self.getOrDefault(self.fitLinear) + + @since("3.0.0") + def getMiniBatchFraction(self) -> float: + """ + Gets the value of miniBatchFraction or its default value. + """ + return self.getOrDefault(self.miniBatchFraction) + + @since("3.0.0") + def getInitStd(self) -> float: + """ + Gets the value of initStd or its default value. + """ + return self.getOrDefault(self.initStd) + + +@inherit_doc +class FMRegressor( + _JavaRegressor["FMRegressionModel"], + _FactorizationMachinesParams, + JavaMLWritable, + JavaMLReadable["FMRegressor"], +): + """ + Factorization Machines learning algorithm for regression. + + solver Supports: + + * gd (normal mini-batch gradient descent) + * adamW (default) + + .. versionadded:: 3.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.regression import FMRegressor + >>> df = spark.createDataFrame([ + ... (2.0, Vectors.dense(2.0)), + ... (1.0, Vectors.dense(1.0)), + ... (0.0, Vectors.sparse(1, [], []))], ["label", "features"]) + >>> + >>> fm = FMRegressor(factorSize=2) + >>> fm.setSeed(16) + FMRegressor... + >>> model = fm.fit(df) + >>> model.getMaxIter() + 100 + >>> test0 = spark.createDataFrame([ + ... (Vectors.dense(-2.0),), + ... (Vectors.dense(0.5),), + ... (Vectors.dense(1.0),), + ... (Vectors.dense(4.0),)], ["features"]) + >>> model.transform(test0).show(10, False) + +--------+-------------------+ + |features|prediction | + +--------+-------------------+ + |[-2.0] |-1.9989237712341565| + |[0.5] |0.4956682219523814 | + |[1.0] |0.994586620589689 | + |[4.0] |3.9880970124135344 | + +--------+-------------------+ + ... + >>> model.intercept + -0.0032501766849261557 + >>> model.linear + DenseVector([0.9978]) + >>> model.factors + DenseMatrix(1, 2, [0.0173, 0.0021], 1) + >>> model_path = temp_path + "/fm_model" + >>> model.save(model_path) + >>> model2 = FMRegressionModel.load(model_path) + >>> model2.intercept + -0.0032501766849261557 + >>> model2.linear + DenseVector([0.9978]) + >>> model2.factors + DenseMatrix(1, 2, [0.0173, 0.0021], 1) + >>> model.transform(test0).take(1) == model2.transform(test0).take(1) + True + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + factorSize: int = 8, + fitIntercept: bool = True, + fitLinear: bool = True, + regParam: float = 0.0, + miniBatchFraction: float = 1.0, + initStd: float = 0.01, + maxIter: int = 100, + stepSize: float = 1.0, + tol: float = 1e-6, + solver: str = "adamW", + seed: Optional[int] = None, + ): + """ + __init__(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + factorSize=8, fitIntercept=True, fitLinear=True, regParam=0.0, \ + miniBatchFraction=1.0, initStd=0.01, maxIter=100, stepSize=1.0, \ + tol=1e-6, solver="adamW", seed=None) + """ + super(FMRegressor, self).__init__() + self._java_obj = self._new_java_obj("org.apache.spark.ml.regression.FMRegressor", self.uid) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + @since("3.0.0") + def setParams( + self, + *, + featuresCol: str = "features", + labelCol: str = "label", + predictionCol: str = "prediction", + factorSize: int = 8, + fitIntercept: bool = True, + fitLinear: bool = True, + regParam: float = 0.0, + miniBatchFraction: float = 1.0, + initStd: float = 0.01, + maxIter: int = 100, + stepSize: float = 1.0, + tol: float = 1e-6, + solver: str = "adamW", + seed: Optional[int] = None, + ) -> "FMRegressor": + """ + setParams(self, \\*, featuresCol="features", labelCol="label", predictionCol="prediction", \ + factorSize=8, fitIntercept=True, fitLinear=True, regParam=0.0, \ + miniBatchFraction=1.0, initStd=0.01, maxIter=100, stepSize=1.0, \ + tol=1e-6, solver="adamW", seed=None) + Sets Params for FMRegressor. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + def _create_model(self, java_model: "JavaObject") -> "FMRegressionModel": + return FMRegressionModel(java_model) + + @since("3.0.0") + def setFactorSize(self, value: int) -> "FMRegressor": + """ + Sets the value of :py:attr:`factorSize`. + """ + return self._set(factorSize=value) + + @since("3.0.0") + def setFitLinear(self, value: bool) -> "FMRegressor": + """ + Sets the value of :py:attr:`fitLinear`. + """ + return self._set(fitLinear=value) + + @since("3.0.0") + def setMiniBatchFraction(self, value: float) -> "FMRegressor": + """ + Sets the value of :py:attr:`miniBatchFraction`. + """ + return self._set(miniBatchFraction=value) + + @since("3.0.0") + def setInitStd(self, value: float) -> "FMRegressor": + """ + Sets the value of :py:attr:`initStd`. + """ + return self._set(initStd=value) + + @since("3.0.0") + def setMaxIter(self, value: int) -> "FMRegressor": + """ + Sets the value of :py:attr:`maxIter`. + """ + return self._set(maxIter=value) + + @since("3.0.0") + def setStepSize(self, value: float) -> "FMRegressor": + """ + Sets the value of :py:attr:`stepSize`. + """ + return self._set(stepSize=value) + + @since("3.0.0") + def setTol(self, value: float) -> "FMRegressor": + """ + Sets the value of :py:attr:`tol`. + """ + return self._set(tol=value) + + @since("3.0.0") + def setSolver(self, value: str) -> "FMRegressor": + """ + Sets the value of :py:attr:`solver`. + """ + return self._set(solver=value) + + @since("3.0.0") + def setSeed(self, value: int) -> "FMRegressor": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + @since("3.0.0") + def setFitIntercept(self, value: bool) -> "FMRegressor": + """ + Sets the value of :py:attr:`fitIntercept`. + """ + return self._set(fitIntercept=value) + + @since("3.0.0") + def setRegParam(self, value: float) -> "FMRegressor": + """ + Sets the value of :py:attr:`regParam`. + """ + return self._set(regParam=value) + + +class FMRegressionModel( + _JavaRegressionModel, + _FactorizationMachinesParams, + JavaMLWritable, + JavaMLReadable["FMRegressionModel"], +): + """ + Model fitted by :class:`FMRegressor`. + + .. versionadded:: 3.0.0 + """ + + @property + @since("3.0.0") + def intercept(self) -> float: + """ + Model intercept. + """ + return self._call_java("intercept") + + @property + @since("3.0.0") + def linear(self) -> Vector: + """ + Model linear term. + """ + return self._call_java("linear") + + @property + @since("3.0.0") + def factors(self) -> Matrix: + """ + Model factor term. + """ + return self._call_java("factors") + + +if __name__ == "__main__": + import doctest + import pyspark.ml.regression + from pyspark.sql import SparkSession + + globs = pyspark.ml.regression.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.regression tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/stat.py new file mode 100644 index 0000000..e0c4d05 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/stat.py @@ -0,0 +1,523 @@ +# +# 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 sys + +from typing import Optional, Tuple, TYPE_CHECKING + + +from pyspark import since, SparkContext +from pyspark.ml.common import _java2py, _py2java +from pyspark.ml.linalg import Matrix, Vector +from pyspark.ml.wrapper import JavaWrapper, _jvm +from pyspark.sql.column import Column, _to_seq +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.functions import lit + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + + +class ChiSquareTest: + """ + Conduct Pearson's independence test for every feature against the label. For each feature, + the (feature, label) pairs are converted into a contingency matrix for which the Chi-squared + statistic is computed. All label and feature values must be categorical. + + The null hypothesis is that the occurrence of the outcomes is statistically independent. + + .. versionadded:: 2.2.0 + + """ + + @staticmethod + def test( + dataset: DataFrame, featuresCol: str, labelCol: str, flatten: bool = False + ) -> DataFrame: + """ + Perform a Pearson's independence test using dataset. + + .. versionadded:: 2.2.0 + .. versionchanged:: 3.1.0 + Added optional ``flatten`` argument. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + DataFrame of categorical labels and categorical features. + Real-valued features will be treated as categorical for each distinct value. + featuresCol : str + Name of features column in dataset, of type `Vector` (`VectorUDT`). + labelCol : str + Name of label column in dataset, of any numerical type. + flatten : bool, optional + if True, flattens the returned dataframe. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + DataFrame containing the test result for every feature against the label. + If flatten is True, this DataFrame will contain one row per feature with the following + fields: + + - `featureIndex: int` + - `pValue: float` + - `degreesOfFreedom: int` + - `statistic: float` + + If flatten is False, this DataFrame will contain a single Row with the following fields: + + - `pValues: Vector` + - `degreesOfFreedom: Array[int]` + - `statistics: Vector` + + Each of these fields has one value per feature. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.stat import ChiSquareTest + >>> dataset = [[0, Vectors.dense([0, 0, 1])], + ... [0, Vectors.dense([1, 0, 1])], + ... [1, Vectors.dense([2, 1, 1])], + ... [1, Vectors.dense([3, 1, 1])]] + >>> dataset = spark.createDataFrame(dataset, ["label", "features"]) + >>> chiSqResult = ChiSquareTest.test(dataset, 'features', 'label') + >>> chiSqResult.select("degreesOfFreedom").collect()[0] + Row(degreesOfFreedom=[3, 1, 0]) + >>> chiSqResult = ChiSquareTest.test(dataset, 'features', 'label', True) + >>> row = chiSqResult.orderBy("featureIndex").collect() + >>> row[0].statistic + 4.0 + """ + sc = SparkContext._active_spark_context + assert sc is not None + + javaTestObj = _jvm().org.apache.spark.ml.stat.ChiSquareTest + args = [_py2java(sc, arg) for arg in (dataset, featuresCol, labelCol, flatten)] + return _java2py(sc, javaTestObj.test(*args)) + + +class Correlation: + """ + Compute the correlation matrix for the input dataset of Vectors using the specified method. + Methods currently supported: `pearson` (default), `spearman`. + + .. versionadded:: 2.2.0 + + Notes + ----- + For Spearman, a rank correlation, we need to create an RDD[Double] for each column + and sort it in order to retrieve the ranks and then join the columns back into an RDD[Vector], + which is fairly costly. Cache the input Dataset before calling corr with `method = 'spearman'` + to avoid recomputing the common lineage. + """ + + @staticmethod + def corr(dataset: DataFrame, column: str, method: str = "pearson") -> DataFrame: + """ + Compute the correlation matrix with specified method using dataset. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + A DataFrame. + column : str + The name of the column of vectors for which the correlation coefficient needs + to be computed. This must be a column of the dataset, and it must contain + Vector objects. + method : str, optional + String specifying the method to use for computing correlation. + Supported: `pearson` (default), `spearman`. + + Returns + ------- + A DataFrame that contains the correlation matrix of the column of vectors. This + DataFrame contains a single row and a single column of name `METHODNAME(COLUMN)`. + + Examples + -------- + >>> from pyspark.ml.linalg import DenseMatrix, Vectors + >>> from pyspark.ml.stat import Correlation + >>> dataset = [[Vectors.dense([1, 0, 0, -2])], + ... [Vectors.dense([4, 5, 0, 3])], + ... [Vectors.dense([6, 7, 0, 8])], + ... [Vectors.dense([9, 0, 0, 1])]] + >>> dataset = spark.createDataFrame(dataset, ['features']) + >>> pearsonCorr = Correlation.corr(dataset, 'features', 'pearson').collect()[0][0] + >>> print(str(pearsonCorr).replace('nan', 'NaN')) + DenseMatrix([[ 1. , 0.0556..., NaN, 0.4004...], + [ 0.0556..., 1. , NaN, 0.9135...], + [ NaN, NaN, 1. , NaN], + [ 0.4004..., 0.9135..., NaN, 1. ]]) + >>> spearmanCorr = Correlation.corr(dataset, 'features', method='spearman').collect()[0][0] + >>> print(str(spearmanCorr).replace('nan', 'NaN')) + DenseMatrix([[ 1. , 0.1054..., NaN, 0.4 ], + [ 0.1054..., 1. , NaN, 0.9486... ], + [ NaN, NaN, 1. , NaN], + [ 0.4 , 0.9486... , NaN, 1. ]]) + """ + sc = SparkContext._active_spark_context + assert sc is not None + + javaCorrObj = _jvm().org.apache.spark.ml.stat.Correlation + args = [_py2java(sc, arg) for arg in (dataset, column, method)] + return _java2py(sc, javaCorrObj.corr(*args)) + + +class KolmogorovSmirnovTest: + """ + Conduct the two-sided Kolmogorov Smirnov (KS) test for data sampled from a continuous + distribution. + + By comparing the largest difference between the empirical cumulative + distribution of the sample data and the theoretical distribution we can provide a test for the + the null hypothesis that the sample data comes from that theoretical distribution. + + .. versionadded:: 2.4.0 + + """ + + @staticmethod + def test(dataset: DataFrame, sampleCol: str, distName: str, *params: float) -> DataFrame: + """ + Conduct a one-sample, two-sided Kolmogorov-Smirnov test for probability distribution + equality. Currently supports the normal distribution, taking as parameters the mean and + standard deviation. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a Dataset or a DataFrame containing the sample of data to test. + sampleCol : str + Name of sample column in dataset, of any numerical type. + distName : str + a `string` name for a theoretical distribution, currently only support "norm". + params : float + a list of `float` values specifying the parameters to be used for the theoretical + distribution. For "norm" distribution, the parameters includes mean and variance. + + Returns + ------- + A DataFrame that contains the Kolmogorov-Smirnov test result for the input sampled data. + This DataFrame will contain a single Row with the following fields: + + - `pValue: Double` + - `statistic: Double` + + Examples + -------- + >>> from pyspark.ml.stat import KolmogorovSmirnovTest + >>> dataset = [[-1.0], [0.0], [1.0]] + >>> dataset = spark.createDataFrame(dataset, ['sample']) + >>> ksResult = KolmogorovSmirnovTest.test(dataset, 'sample', 'norm', 0.0, 1.0).first() + >>> round(ksResult.pValue, 3) + 1.0 + >>> round(ksResult.statistic, 3) + 0.175 + >>> dataset = [[2.0], [3.0], [4.0]] + >>> dataset = spark.createDataFrame(dataset, ['sample']) + >>> ksResult = KolmogorovSmirnovTest.test(dataset, 'sample', 'norm', 3.0, 1.0).first() + >>> round(ksResult.pValue, 3) + 1.0 + >>> round(ksResult.statistic, 3) + 0.175 + """ + sc = SparkContext._active_spark_context + assert sc is not None + + javaTestObj = _jvm().org.apache.spark.ml.stat.KolmogorovSmirnovTest + dataset = _py2java(sc, dataset) + params = [float(param) for param in params] # type: ignore[assignment] + return _java2py( + sc, javaTestObj.test(dataset, sampleCol, distName, _jvm().PythonUtils.toSeq(params)) + ) + + +class Summarizer: + """ + Tools for vectorized statistics on MLlib Vectors. + The methods in this package provide various statistics for Vectors contained inside DataFrames. + This class lets users pick the statistics they would like to extract for a given column. + + .. versionadded:: 2.4.0 + + Examples + -------- + >>> from pyspark.ml.stat import Summarizer + >>> from pyspark.sql import Row + >>> from pyspark.ml.linalg import Vectors + >>> summarizer = Summarizer.metrics("mean", "count") + >>> df = sc.parallelize([Row(weight=1.0, features=Vectors.dense(1.0, 1.0, 1.0)), + ... Row(weight=0.0, features=Vectors.dense(1.0, 2.0, 3.0))]).toDF() + >>> df.select(summarizer.summary(df.features, df.weight)).show(truncate=False) + +-----------------------------------+ + |aggregate_metrics(features, weight)| + +-----------------------------------+ + |{[1.0,1.0,1.0], 1} | + +-----------------------------------+ + >>> df.select(summarizer.summary(df.features)).show(truncate=False) + +--------------------------------+ + |aggregate_metrics(features, 1.0)| + +--------------------------------+ + |{[1.0,1.5,2.0], 2} | + +--------------------------------+ + >>> df.select(Summarizer.mean(df.features, df.weight)).show(truncate=False) + +--------------+ + |mean(features)| + +--------------+ + |[1.0,1.0,1.0] | + +--------------+ + >>> df.select(Summarizer.mean(df.features)).show(truncate=False) + +--------------+ + |mean(features)| + +--------------+ + |[1.0,1.5,2.0] | + +--------------+ + """ + + @staticmethod + @since("2.4.0") + def mean(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of mean summary + """ + return Summarizer._get_single_metric(col, weightCol, "mean") + + @staticmethod + @since("3.0.0") + def sum(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of sum summary + """ + return Summarizer._get_single_metric(col, weightCol, "sum") + + @staticmethod + @since("2.4.0") + def variance(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of variance summary + """ + return Summarizer._get_single_metric(col, weightCol, "variance") + + @staticmethod + @since("3.0.0") + def std(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of std summary + """ + return Summarizer._get_single_metric(col, weightCol, "std") + + @staticmethod + @since("2.4.0") + def count(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of count summary + """ + return Summarizer._get_single_metric(col, weightCol, "count") + + @staticmethod + @since("2.4.0") + def numNonZeros(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of numNonZero summary + """ + return Summarizer._get_single_metric(col, weightCol, "numNonZeros") + + @staticmethod + @since("2.4.0") + def max(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of max summary + """ + return Summarizer._get_single_metric(col, weightCol, "max") + + @staticmethod + @since("2.4.0") + def min(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of min summary + """ + return Summarizer._get_single_metric(col, weightCol, "min") + + @staticmethod + @since("2.4.0") + def normL1(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of normL1 summary + """ + return Summarizer._get_single_metric(col, weightCol, "normL1") + + @staticmethod + @since("2.4.0") + def normL2(col: Column, weightCol: Optional[Column] = None) -> Column: + """ + return a column of normL2 summary + """ + return Summarizer._get_single_metric(col, weightCol, "normL2") + + @staticmethod + def _check_param(featuresCol: Column, weightCol: Optional[Column]) -> Tuple[Column, Column]: + if weightCol is None: + weightCol = lit(1.0) + if not isinstance(featuresCol, Column) or not isinstance(weightCol, Column): + raise TypeError("featureCol and weightCol should be a Column") + return featuresCol, weightCol + + @staticmethod + def _get_single_metric(col: Column, weightCol: Optional[Column], metric: str) -> Column: + col, weightCol = Summarizer._check_param(col, weightCol) + return Column( + JavaWrapper._new_java_obj( + "org.apache.spark.ml.stat.Summarizer." + metric, col._jc, weightCol._jc + ) + ) + + @staticmethod + def metrics(*metrics: str) -> "SummaryBuilder": + """ + Given a list of metrics, provides a builder that it turns computes metrics from a column. + + See the documentation of :py:class:`Summarizer` for an example. + + The following metrics are accepted (case sensitive): + - mean: a vector that contains the coefficient-wise mean. + - sum: a vector that contains the coefficient-wise sum. + - variance: a vector that contains the coefficient-wise variance. + - std: a vector that contains the coefficient-wise standard deviation. + - count: the count of all vectors seen. + - numNonzeros: a vector with the number of non-zeros for each coefficients + - max: the maximum for each coefficient. + - min: the minimum for each coefficient. + - normL2: the Euclidean norm for each coefficient. + - normL1: the L1 norm of each coefficient (sum of the absolute values). + + .. versionadded:: 2.4.0 + + Notes + ----- + Currently, the performance of this interface is about 2x~3x slower than using the RDD + interface. + + Examples + -------- + metrics : str + metrics that can be provided. + + Returns + ------- + :py:class:`pyspark.ml.stat.SummaryBuilder` + """ + sc = SparkContext._active_spark_context + assert sc is not None + + js = JavaWrapper._new_java_obj( + "org.apache.spark.ml.stat.Summarizer.metrics", _to_seq(sc, metrics) + ) + return SummaryBuilder(js) + + +class SummaryBuilder(JavaWrapper): + """ + A builder object that provides summary statistics about a given column. + + Users should not directly create such builders, but instead use one of the methods in + :py:class:`pyspark.ml.stat.Summarizer` + + .. versionadded:: 2.4.0 + + """ + + def __init__(self, jSummaryBuilder: "JavaObject"): + super(SummaryBuilder, self).__init__(jSummaryBuilder) + + def summary(self, featuresCol: Column, weightCol: Optional[Column] = None) -> Column: + """ + Returns an aggregate object that contains the summary of the column with the requested + metrics. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + featuresCol : str + a column that contains features Vector object. + weightCol : str, optional + a column that contains weight value. Default weight is 1.0. + + Returns + ------- + :py:class:`pyspark.sql.Column` + an aggregate column that contains the statistics. The exact content of this + structure is determined during the creation of the builder. + """ + featuresCol, weightCol = Summarizer._check_param(featuresCol, weightCol) + assert self._java_obj is not None + + return Column(self._java_obj.summary(featuresCol._jc, weightCol._jc)) + + +class MultivariateGaussian: + """Represents a (mean, cov) tuple + + .. versionadded:: 3.0.0 + + Examples + -------- + >>> from pyspark.ml.linalg import DenseMatrix, Vectors + >>> from pyspark.ml.stat import MultivariateGaussian + >>> m = MultivariateGaussian(Vectors.dense([11,12]), DenseMatrix(2, 2, (1.0, 3.0, 5.0, 2.0))) + >>> (m.mean, m.cov.toArray()) + (DenseVector([11.0, 12.0]), array([[ 1., 5.], + [ 3., 2.]])) + """ + + def __init__(self, mean: Vector, cov: Matrix): + self.mean = mean + self.cov = cov + + +if __name__ == "__main__": + import doctest + import numpy + import pyspark.ml.stat + from pyspark.sql import SparkSession + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + + globs = pyspark.ml.stat.__dict__.copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.stat tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + + failure_count, test_count = doctest.testmod( + globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_classification.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_classification.py new file mode 100644 index 0000000..f3e621c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_classification.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_classification import ClassificationTestsMixin + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + + +@unittest.skipIf(not have_torch, "torch is required") +class ClassificationTestsOnConnect(ClassificationTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.connect.copyFromLocalToFs.allowDestLocal", "true") + .getOrCreate() + ) + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_classification import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_evaluation.py new file mode 100644 index 0000000..ce7cf03 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_evaluation.py @@ -0,0 +1,47 @@ +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_evaluation import EvaluationTestsMixin + +have_torcheval = True +try: + import torcheval # noqa: F401 +except ImportError: + have_torcheval = False + + +@unittest.skipIf(not have_torcheval, "torcheval is required") +class EvaluationTestsOnConnect(EvaluationTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.remote("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_evaluation import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_feature.py new file mode 100644 index 0000000..d7698c3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_feature.py @@ -0,0 +1,40 @@ +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_feature import FeatureTestsMixin + + +class FeatureTestsOnConnect(FeatureTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.remote("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_feature import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_function.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_function.py new file mode 100644 index 0000000..7da3d3f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_function.py @@ -0,0 +1,113 @@ +# +# 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 os +import unittest + +from pyspark.sql import SparkSession as PySparkSession +from pyspark.sql.dataframe import DataFrame as SDF +from pyspark.ml import functions as SF + +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.testing.connectutils import ( + should_test_connect, + ReusedConnectTestCase, +) +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + +if should_test_connect: + from pyspark.sql.connect.dataframe import DataFrame as CDF + from pyspark.ml.connect import functions as CF + + +class SparkConnectMLFunctionTests(ReusedConnectTestCase, PandasOnSparkTestUtils, SQLTestUtils): + """These test cases exercise the interface to the proto plan + generation but do not call Spark.""" + + @classmethod + def setUpClass(cls): + super(SparkConnectMLFunctionTests, cls).setUpClass() + # Disable the shared namespace so pyspark.sql.functions, etc point the regular + # PySpark libraries. + os.environ["PYSPARK_NO_NAMESPACE_SHARE"] = "1" + cls.connect = cls.spark # Switch Spark Connect session and regular PySpark sesion. + cls.spark = PySparkSession._instantiatedSession + assert cls.spark is not None + + @classmethod + def tearDownClass(cls): + cls.spark = cls.connect # Stopping Spark Connect closes the session in JVM at the server. + super(SparkConnectMLFunctionTests, cls).setUpClass() + del os.environ["PYSPARK_NO_NAMESPACE_SHARE"] + + def compare_by_show(self, df1, df2, n: int = 20, truncate: int = 20): + assert isinstance(df1, (SDF, CDF)) + if isinstance(df1, SDF): + str1 = df1._jdf.showString(n, truncate, False) + else: + str1 = df1._show_string(n, truncate, False) + + assert isinstance(df2, (SDF, CDF)) + if isinstance(df2, SDF): + str2 = df2._jdf.showString(n, truncate, False) + else: + str2 = df2._show_string(n, truncate, False) + + self.assertEqual(str1, str2) + + def test_array_vector_conversion(self): + query = """ + SELECT * FROM VALUES + (1, 4, ARRAY(1.0, 2.0, 3.0)), + (1, 2, ARRAY(-1.0, -2.0, -3.0)) + AS tab(a, b, c) + """ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.compare_by_show( + cdf.select(cdf.b, CF.array_to_vector(cdf.c)), + sdf.select(sdf.b, SF.array_to_vector(sdf.c)), + ) + + cdf1 = cdf.select("a", CF.array_to_vector(cdf.c).alias("d")) + sdf1 = sdf.select("a", SF.array_to_vector(sdf.c).alias("d")) + + self.compare_by_show( + cdf1.select(CF.vector_to_array(cdf1.d)), + sdf1.select(SF.vector_to_array(sdf1.d)), + ) + self.compare_by_show( + cdf1.select(CF.vector_to_array(cdf1.d, "float32")), + sdf1.select(SF.vector_to_array(sdf1.d, "float32")), + ) + self.compare_by_show( + cdf1.select(CF.vector_to_array(cdf1.d, "float64")), + sdf1.select(SF.vector_to_array(sdf1.d, "float64")), + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_function import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_pipeline.py new file mode 100644 index 0000000..e676c8b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_pipeline.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_pipeline import PipelineTestsMixin + + +class PipelineTestsOnConnect(PipelineTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.connect.copyFromLocalToFs.allowDestLocal", "true") + .getOrCreate() + ) + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_pipeline import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_summarizer.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_summarizer.py new file mode 100644 index 0000000..0b0537d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_summarizer.py @@ -0,0 +1,40 @@ +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_summarizer import SummarizerTestsMixin + + +class SummarizerTestsOnConnect(SummarizerTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.remote("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_summarizer import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_tuning.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_tuning.py new file mode 100644 index 0000000..18673d4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_connect_tuning.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest +from pyspark.sql import SparkSession +from pyspark.ml.tests.connect.test_legacy_mode_tuning import CrossValidatorTestsMixin + + +class CrossValidatorTestsOnConnect(CrossValidatorTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.connect.copyFromLocalToFs.allowDestLocal", "true") + .getOrCreate() + ) + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_connect_tuning import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py new file mode 100644 index 0000000..84d5829 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py @@ -0,0 +1,238 @@ +# -*- coding: utf-8 -*- +# +# 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 os +import tempfile +import unittest +import numpy as np +from pyspark.ml.connect.classification import ( + LogisticRegression as LORV2, + LogisticRegressionModel as LORV2Model, +) +from pyspark.sql import SparkSession + + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + + +class ClassificationTestsMixin: + @staticmethod + def _check_result(result_dataframe, expected_predictions, expected_probabilities=None): + np.testing.assert_array_equal(list(result_dataframe.prediction), expected_predictions) + if "probability" in result_dataframe.columns: + np.testing.assert_allclose( + list(result_dataframe.probability), + expected_probabilities, + rtol=1e-1, + ) + + def test_binary_classes_logistic_regression(self): + df1 = self.spark.createDataFrame( + [ + (1.0, [0.0, 5.0]), + (0.0, [1.0, 2.0]), + (1.0, [2.0, 1.0]), + (0.0, [3.0, 3.0]), + ] + * 100, + ["label", "features"], + ) + eval_df1 = self.spark.createDataFrame( + [ + ([0.0, 2.0],), + ([3.5, 3.0],), + ], + ["features"], + ) + + lorv2 = LORV2(maxIter=200, numTrainWorkers=2, learningRate=0.001) + assert lorv2.getMaxIter() == 200 + assert lorv2.getNumTrainWorkers() == 2 + assert lorv2.getOrDefault(lorv2.learningRate) == 0.001 + + model = lorv2.fit(df1) + assert model.uid == lorv2.uid + + expected_predictions = [1, 0] + expected_probabilities = [ + [0.217875, 0.782125], + [0.839615, 0.160385], + ] + + result = model.transform(eval_df1).toPandas() + self._check_result(result, expected_predictions, expected_probabilities) + local_transform_result = model.transform(eval_df1.toPandas()) + self._check_result(local_transform_result, expected_predictions, expected_probabilities) + + model.set(model.probabilityCol, "") + result_without_prob = model.transform(eval_df1).toPandas() + assert "probability" not in result_without_prob.columns + self._check_result(result_without_prob, expected_predictions, None) + + def test_multi_classes_logistic_regression(self): + df1 = self.spark.createDataFrame( + [ + (1.0, [1.0, 5.0]), + (2.0, [1.0, -2.0]), + (0.0, [-2.0, 1.5]), + ] + * 100, + ["label", "features"], + ) + eval_df1 = self.spark.createDataFrame( + [ + ([1.5, 5.0],), + ([1.0, -2.5],), + ([-2.0, 1.0],), + ], + ["features"], + ) + + lorv2 = LORV2(maxIter=200, numTrainWorkers=2, learningRate=0.001) + + model = lorv2.fit(df1) + + expected_predictions = [1, 2, 0] + expected_probabilities = [ + [5.526459e-03, 9.943553e-01, 1.183146e-04], + [4.629959e-03, 8.141352e-03, 9.872288e-01], + [9.624363e-01, 3.080821e-02, 6.755549e-03], + ] + + result = model.transform(eval_df1).toPandas() + self._check_result(result, expected_predictions, expected_probabilities) + local_transform_result = model.transform(eval_df1.toPandas()) + self._check_result(local_transform_result, expected_predictions, expected_probabilities) + + def test_save_load(self): + with tempfile.TemporaryDirectory() as tmp_dir: + estimator = LORV2(maxIter=2, numTrainWorkers=2, learningRate=0.001) + local_path = os.path.join(tmp_dir, "estimator") + estimator.saveToLocal(local_path) + loaded_estimator = LORV2.loadFromLocal(local_path) + assert loaded_estimator.uid == estimator.uid + assert loaded_estimator.getOrDefault(loaded_estimator.maxIter) == 2 + assert loaded_estimator.getOrDefault(loaded_estimator.numTrainWorkers) == 2 + assert loaded_estimator.getOrDefault(loaded_estimator.learningRate) == 0.001 + + # test overwriting + estimator2 = estimator.copy() + estimator2.set(estimator2.maxIter, 10) + estimator2.saveToLocal(local_path, overwrite=True) + loaded_estimator2 = LORV2.loadFromLocal(local_path) + assert loaded_estimator2.getOrDefault(loaded_estimator2.maxIter) == 10 + + fs_path = os.path.join(tmp_dir, "fs", "estimator") + estimator.save(fs_path) + loaded_estimator = LORV2.load(fs_path) + assert loaded_estimator.uid == estimator.uid + assert loaded_estimator.getOrDefault(loaded_estimator.maxIter) == 2 + assert loaded_estimator.getOrDefault(loaded_estimator.numTrainWorkers) == 2 + assert loaded_estimator.getOrDefault(loaded_estimator.learningRate) == 0.001 + + training_dataset = self.spark.createDataFrame( + [ + (1.0, [0.0, 5.0]), + (0.0, [1.0, 2.0]), + (1.0, [2.0, 1.0]), + (0.0, [3.0, 3.0]), + ] + * 100, + ["label", "features"], + ) + eval_df1 = self.spark.createDataFrame( + [ + ([0.0, 2.0],), + ([3.5, 3.0],), + ], + ["features"], + ) + + model = estimator.fit(training_dataset) + model_predictions = model.transform(eval_df1.toPandas()) + + assert model.uid == estimator.uid + + local_model_path = os.path.join(tmp_dir, "model") + model.saveToLocal(local_model_path) + + # test saved torch model can be loaded by pytorch solely + lor_torch_model = torch.load( + os.path.join(local_model_path, "LogisticRegressionModel.torch") + ) + + with torch.inference_mode(): + torch_infer_result = lor_torch_model( + torch.tensor(np.stack(list(eval_df1.toPandas().features)), dtype=torch.float32) + ).numpy() + + np.testing.assert_allclose( + np.stack(list(model_predictions.probability)), + torch_infer_result, + rtol=1e-4, + ) + + loaded_model = LORV2Model.loadFromLocal(local_model_path) + assert loaded_model.numFeatures == 2 + assert loaded_model.numClasses == 2 + assert loaded_model.getOrDefault(loaded_model.maxIter) == 2 + assert loaded_model.torch_model is not None + np.testing.assert_allclose( + loaded_model.torch_model.weight.detach().numpy(), + model.torch_model.weight.detach().numpy(), + ) + np.testing.assert_allclose( + loaded_model.torch_model.bias.detach().numpy(), + model.torch_model.bias.detach().numpy(), + ) + + # Test loaded model transformation. + loaded_model.transform(eval_df1.toPandas()) + + fs_model_path = os.path.join(tmp_dir, "fs", "model") + model.save(fs_model_path) + loaded_model = LORV2Model.load(fs_model_path) + assert loaded_model.numFeatures == 2 + assert loaded_model.numClasses == 2 + assert loaded_model.getOrDefault(loaded_model.maxIter) == 2 + assert loaded_model.torch_model is not None + # Test loaded model transformation works. + loaded_model.transform(eval_df1.toPandas()) + + +class ClassificationTests(ClassificationTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_classification import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py new file mode 100644 index 0000000..9ff26c1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py @@ -0,0 +1,194 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest +import numpy as np +import tempfile + +from pyspark.ml.connect.evaluation import ( + RegressionEvaluator, + BinaryClassificationEvaluator, + MulticlassClassificationEvaluator, +) +from pyspark.sql import SparkSession + + +have_torcheval = True +try: + import torcheval # noqa: F401 +except ImportError: + have_torcheval = False + + +class EvaluationTestsMixin: + def test_regressor_evaluator(self): + df1 = self.spark.createDataFrame( + [ + (0.5, 1.0), + (-0.5, -0.8), + (2.0, 3.0), + ], + schema=["label", "prediction"], + ) + + local_df1 = df1.toPandas() + + mse_evaluator = RegressionEvaluator( + metricName="mse", + labelCol="label", + predictionCol="prediction", + ) + + expected_mse = 0.4466666877269745 + mse = mse_evaluator.evaluate(df1) + mse_local = mse_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(mse, expected_mse) + np.testing.assert_almost_equal(mse_local, expected_mse) + + rmse_evaluator = RegressionEvaluator( + metricName="rmse", + labelCol="label", + predictionCol="prediction", + ) + + expected_rmse = 0.6683312709480042 + rmse = rmse_evaluator.evaluate(df1) + rmse_local = rmse_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(rmse, expected_rmse) + np.testing.assert_almost_equal(rmse_local, expected_rmse) + + r2_evaluator = RegressionEvaluator( + metricName="r2", + labelCol="label", + predictionCol="prediction", + ) + + expected_r2 = 0.5768420696258545 + r2 = r2_evaluator.evaluate(df1) + r2_local = r2_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(r2, expected_r2) + np.testing.assert_almost_equal(r2_local, expected_r2) + + # Test save / load + with tempfile.TemporaryDirectory() as tmp_dir: + r2_evaluator.saveToLocal(f"{tmp_dir}/ev") + loaded_evaluator = RegressionEvaluator.loadFromLocal(f"{tmp_dir}/ev") + assert loaded_evaluator.getMetricName() == "r2" + + def test_binary_classifier_evaluator(self): + df1 = self.spark.createDataFrame( + [ + (1, 0.2, [0.8, 0.2]), + (0, 0.6, [0.4, 0.6]), + (1, 0.8, [0.2, 0.8]), + (1, 0.7, [0.3, 0.7]), + (0, 0.4, [0.6, 0.4]), + (0, 0.3, [0.7, 0.3]), + ], + schema=["label", "prob", "prob2"], + ) + + local_df1 = df1.toPandas() + + for prob_col in ["prob", "prob2"]: + auroc_evaluator = BinaryClassificationEvaluator( + metricName="areaUnderROC", + labelCol="label", + probabilityCol=prob_col, + ) + + expected_auroc = 0.6667 + auroc = auroc_evaluator.evaluate(df1) + auroc_local = auroc_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(auroc, expected_auroc, decimal=2) + np.testing.assert_almost_equal(auroc_local, expected_auroc, decimal=2) + + auprc_evaluator = BinaryClassificationEvaluator( + metricName="areaUnderPR", + labelCol="label", + probabilityCol=prob_col, + ) + + expected_auprc = 0.8333 + auprc = auprc_evaluator.evaluate(df1) + auprc_local = auprc_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(auprc, expected_auprc, decimal=2) + np.testing.assert_almost_equal(auprc_local, expected_auprc, decimal=2) + + # Test save / load + with tempfile.TemporaryDirectory() as tmp_dir: + auprc_evaluator.saveToLocal(f"{tmp_dir}/ev") + loaded_evaluator = RegressionEvaluator.loadFromLocal(f"{tmp_dir}/ev") + assert loaded_evaluator.getMetricName() == "areaUnderPR" + + def test_multiclass_classifier_evaluator(self): + df1 = self.spark.createDataFrame( + [ + (1, 1), + (1, 1), + (2, 3), + (0, 0), + (0, 1), + (3, 1), + (3, 3), + (2, 2), + (1, 0), + (2, 2), + ], + schema=["label", "prediction"], + ) + + local_df1 = df1.toPandas() + + accuracy_evaluator = MulticlassClassificationEvaluator( + metricName="accuracy", + labelCol="label", + predictionCol="prediction", + ) + + expected_accuracy = 0.600 + accuracy = accuracy_evaluator.evaluate(df1) + accuracy_local = accuracy_evaluator.evaluate(local_df1) + np.testing.assert_almost_equal(accuracy, expected_accuracy, decimal=2) + np.testing.assert_almost_equal(accuracy_local, expected_accuracy, decimal=2) + + # Test save / load + with tempfile.TemporaryDirectory() as tmp_dir: + accuracy_evaluator.saveToLocal(f"{tmp_dir}/ev") + loaded_evaluator = RegressionEvaluator.loadFromLocal(f"{tmp_dir}/ev") + assert loaded_evaluator.getMetricName() == "accuracy" + + +@unittest.skipIf(not have_torcheval, "torcheval is required") +class EvaluationTests(EvaluationTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_evaluation import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py new file mode 100644 index 0000000..3aac4a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py @@ -0,0 +1,156 @@ +# -*- coding: utf-8 -*- +# +# 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 os +import pickle +import numpy as np +import tempfile +import unittest + +from pyspark.ml.connect.feature import ( + MaxAbsScaler, + MaxAbsScalerModel, + StandardScaler, + StandardScalerModel, +) +from pyspark.sql import SparkSession + + +class FeatureTestsMixin: + def test_max_abs_scaler(self): + df1 = self.spark.createDataFrame( + [ + ([2.0, 3.5, 1.5],), + ([-3.0, -0.5, -2.5],), + ], + schema=["features"], + ) + + scaler = MaxAbsScaler(inputCol="features", outputCol="scaled_features") + + model = scaler.fit(df1) + assert model.uid == scaler.uid + result = model.transform(df1).toPandas() + assert list(result.columns) == ["features", "scaled_features"] + + expected_result = [[2.0 / 3, 1.0, 0.6], [-1.0, -1.0 / 7, -1.0]] + + np.testing.assert_allclose(list(result.scaled_features), expected_result) + + local_df1 = df1.toPandas() + local_fit_model = scaler.fit(local_df1) + local_transform_result = local_fit_model.transform(local_df1) + assert id(local_transform_result) == id(local_df1) + assert list(local_transform_result.columns) == ["features", "scaled_features"] + + np.testing.assert_allclose(list(local_transform_result.scaled_features), expected_result) + + with tempfile.TemporaryDirectory() as tmp_dir: + estimator_path = os.path.join(tmp_dir, "estimator") + scaler.saveToLocal(estimator_path) + loaded_scaler = MaxAbsScaler.loadFromLocal(estimator_path) + assert loaded_scaler.getInputCol() == "features" + assert loaded_scaler.getOutputCol() == "scaled_features" + + model_path = os.path.join(tmp_dir, "model") + model.saveToLocal(model_path) + loaded_model = MaxAbsScalerModel.loadFromLocal(model_path) + + np.testing.assert_allclose(model.scale_values, loaded_model.scale_values) + np.testing.assert_allclose(model.max_abs_values, loaded_model.max_abs_values) + assert model.n_samples_seen == loaded_model.n_samples_seen + + # Test loading core model as scikit-learn model + with open(os.path.join(model_path, "MaxAbsScalerModel.sklearn.pkl"), "rb") as f: + sk_model = pickle.load(f) + sk_result = sk_model.transform(np.stack(list(local_df1.features))) + np.testing.assert_allclose(sk_result, expected_result) + + def test_standard_scaler(self): + df1 = self.spark.createDataFrame( + [ + ([2.0, 3.5, 1.5],), + ([-3.0, -0.5, -2.5],), + ([1.0, -1.5, 0.5],), + ], + schema=["features"], + ) + + scaler = StandardScaler(inputCol="features", outputCol="scaled_features") + model = scaler.fit(df1) + assert model.uid == scaler.uid + result = model.transform(df1).toPandas() + assert list(result.columns) == ["features", "scaled_features"] + + expected_result = [ + [0.7559289460184544, 1.1338934190276817, 0.8006407690254358], + [-1.1338934190276817, -0.3779644730092272, -1.1208970766356101], + [0.3779644730092272, -0.7559289460184544, 0.32025630761017426], + ] + + np.testing.assert_allclose(list(result.scaled_features), expected_result) + + local_df1 = df1.toPandas() + local_fit_model = scaler.fit(local_df1) + local_transform_result = local_fit_model.transform(local_df1) + assert id(local_transform_result) == id(local_df1) + assert list(local_transform_result.columns) == ["features", "scaled_features"] + + np.testing.assert_allclose(list(local_transform_result.scaled_features), expected_result) + + with tempfile.TemporaryDirectory() as tmp_dir: + estimator_path = os.path.join(tmp_dir, "estimator") + scaler.saveToLocal(estimator_path) + loaded_scaler = StandardScaler.loadFromLocal(estimator_path) + assert loaded_scaler.getInputCol() == "features" + assert loaded_scaler.getOutputCol() == "scaled_features" + + model_path = os.path.join(tmp_dir, "model") + model.saveToLocal(model_path) + loaded_model = StandardScalerModel.loadFromLocal(model_path) + + np.testing.assert_allclose(model.std_values, loaded_model.std_values) + np.testing.assert_allclose(model.mean_values, loaded_model.mean_values) + np.testing.assert_allclose(model.scale_values, loaded_model.scale_values) + assert model.n_samples_seen == loaded_model.n_samples_seen + + # Test loading core model as scikit-learn model + with open(os.path.join(model_path, "StandardScalerModel.sklearn.pkl"), "rb") as f: + sk_model = pickle.load(f) + sk_result = sk_model.transform(np.stack(list(local_df1.features))) + np.testing.assert_allclose(sk_result, expected_result) + + +class FeatureTests(FeatureTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_feature import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py new file mode 100644 index 0000000..5fd4f6f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py @@ -0,0 +1,184 @@ +# -*- coding: utf-8 -*- +# +# 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 os +import tempfile +import unittest +import numpy as np +from pyspark.ml.connect.feature import StandardScaler +from pyspark.ml.connect.classification import LogisticRegression as LORV2 +from pyspark.ml.connect.pipeline import Pipeline +from pyspark.sql import SparkSession + + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + + +class PipelineTestsMixin: + @staticmethod + def _check_result(result_dataframe, expected_predictions, expected_probabilities=None): + np.testing.assert_array_equal(list(result_dataframe.prediction), expected_predictions) + if "probability" in result_dataframe.columns: + np.testing.assert_allclose( + list(result_dataframe.probability), + expected_probabilities, + rtol=1e-1, + ) + + def test_pipeline(self): + train_dataset = self.spark.createDataFrame( + [ + (1.0, [0.0, 5.0]), + (0.0, [1.0, 2.0]), + (1.0, [2.0, 1.0]), + (0.0, [3.0, 3.0]), + ] + * 100, + ["label", "features"], + ) + eval_dataset = self.spark.createDataFrame( + [ + ([0.0, 2.0],), + ([3.5, 3.0],), + ], + ["features"], + ) + scaler = StandardScaler(inputCol="features", outputCol="scaled_features") + lorv2 = LORV2( + maxIter=200, numTrainWorkers=2, learningRate=0.001, featuresCol="scaled_features" + ) + + pipeline = Pipeline(stages=[scaler, lorv2]) + model = pipeline.fit(train_dataset) + assert model.uid == pipeline.uid + + expected_predictions = [1, 0] + expected_probabilities = [ + [0.117658, 0.882342], + [0.878738, 0.121262], + ] + + result = model.transform(eval_dataset).toPandas() + self._check_result(result, expected_predictions, expected_probabilities) + local_transform_result = model.transform(eval_dataset.toPandas()) + self._check_result(local_transform_result, expected_predictions, expected_probabilities) + + pipeline2 = Pipeline(stages=[pipeline]) + model2 = pipeline2.fit(train_dataset) + result2 = model2.transform(eval_dataset).toPandas() + self._check_result(result2, expected_predictions, expected_probabilities) + local_transform_result2 = model2.transform(eval_dataset.toPandas()) + self._check_result(local_transform_result2, expected_predictions, expected_probabilities) + + with tempfile.TemporaryDirectory() as tmp_dir: + pipeline_local_path = os.path.join(tmp_dir, "pipeline") + pipeline.saveToLocal(pipeline_local_path) + loaded_pipeline = Pipeline.loadFromLocal(pipeline_local_path) + + assert pipeline.uid == loaded_pipeline.uid + assert loaded_pipeline.getStages()[1].getMaxIter() == 200 + + pipeline_model_local_path = os.path.join(tmp_dir, "pipeline_model") + model.saveToLocal(pipeline_model_local_path) + loaded_model = Pipeline.loadFromLocal(pipeline_model_local_path) + + assert model.uid == loaded_model.uid + assert loaded_model.stages[1].getMaxIter() == 200 + + loaded_model_transform_result = loaded_model.transform(eval_dataset).toPandas() + self._check_result( + loaded_model_transform_result, expected_predictions, expected_probabilities + ) + + pipeline2_local_path = os.path.join(tmp_dir, "pipeline2") + pipeline2.saveToLocal(pipeline2_local_path) + loaded_pipeline2 = Pipeline.loadFromLocal(pipeline2_local_path) + + assert pipeline2.uid == loaded_pipeline2.uid + assert loaded_pipeline2.getStages()[0].getStages()[1].getMaxIter() == 200 + + pipeline2_model_local_path = os.path.join(tmp_dir, "pipeline2_model") + model2.saveToLocal(pipeline2_model_local_path) + loaded_model2 = Pipeline.loadFromLocal(pipeline2_model_local_path) + + assert model2.uid == loaded_model2.uid + assert loaded_model2.stages[0].stages[1].getMaxIter() == 200 + + loaded_model2_transform_result = loaded_model2.transform(eval_dataset).toPandas() + self._check_result( + loaded_model2_transform_result, expected_predictions, expected_probabilities + ) + + @staticmethod + def test_pipeline_copy(): + scaler = StandardScaler(inputCol="features", outputCol="scaled_features") + lorv2 = LORV2( + maxIter=200, numTrainWorkers=2, learningRate=0.001, featuresCol="scaled_features" + ) + + pipeline = Pipeline(stages=[scaler, lorv2]) + + copied_pipeline = pipeline.copy( + {scaler.inputCol: "f1", lorv2.maxIter: 10, lorv2.numTrainWorkers: 1} + ) + + stages = copied_pipeline.getStages() + + assert stages[0].getInputCol() == "f1" + assert stages[1].getOrDefault(stages[1].maxIter) == 10 + assert stages[1].getOrDefault(stages[1].numTrainWorkers) == 1 + assert stages[1].getOrDefault(stages[1].featuresCol) == "scaled_features" + + pipeline2 = Pipeline(stages=[pipeline]) + copied_pipeline2 = pipeline2.copy( + {scaler.inputCol: "f2", lorv2.maxIter: 20, lorv2.numTrainWorkers: 20} + ) + + stages = copied_pipeline2.getStages()[0].getStages() + + assert stages[0].getInputCol() == "f2" + assert stages[1].getOrDefault(stages[1].maxIter) == 20 + assert stages[1].getOrDefault(stages[1].numTrainWorkers) == 20 + assert stages[1].getOrDefault(stages[1].featuresCol) == "scaled_features" + + # test original stage instance params are not modified after pipeline copying. + assert scaler.getInputCol() == "features" + assert lorv2.getOrDefault(lorv2.maxIter) == 200 + + +class PipelineTests(PipelineTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_pipeline import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py new file mode 100644 index 0000000..49c092b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest +import numpy as np + +from pyspark.ml.connect.summarizer import summarize_dataframe +from pyspark.sql import SparkSession + + +class SummarizerTestsMixin: + def test_summarize_dataframe(self): + df1 = self.spark.createDataFrame( + [ + ([2.0, -1.5],), + ([-3.0, 0.5],), + ([1.0, 3.5],), + ], + schema=["features"], + ) + + df1_local = df1.toPandas() + + result = summarize_dataframe(df1, "features", ["min", "max", "sum", "mean", "std"]) + result_local = summarize_dataframe( + df1_local, "features", ["min", "max", "sum", "mean", "std"] + ) + expected_result = { + "min": [-3.0, -1.5], + "max": [2.0, 3.5], + "sum": [0.0, 2.5], + "mean": [0.0, 0.83333333], + "std": [2.64575131, 2.51661148], + } + + def assert_dict_allclose(dict1, dict2): + assert set(dict1.keys()) == set(dict2.keys()) + + for key in dict1: + np.testing.assert_allclose(dict1[key], dict2[key]) + + assert_dict_allclose(result, expected_result) + assert_dict_allclose(result_local, expected_result) + + +class SummarizerTests(SummarizerTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_summarizer import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py new file mode 100644 index 0000000..0ade227 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py @@ -0,0 +1,292 @@ +# -*- coding: utf-8 -*- +# +# 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 tempfile +import unittest +import numpy as np +import pandas as pd +from pyspark.ml.param import Param, Params +from pyspark.ml.connect import Model, Estimator +from pyspark.ml.connect.feature import StandardScaler +from pyspark.ml.connect.classification import LogisticRegression as LORV2 +from pyspark.ml.connect.pipeline import Pipeline +from pyspark.ml.connect.tuning import CrossValidator, CrossValidatorModel +from pyspark.ml.connect.evaluation import BinaryClassificationEvaluator, RegressionEvaluator +from pyspark.ml.tuning import ParamGridBuilder +from pyspark.sql import SparkSession +from pyspark.sql.functions import rand + +from sklearn.datasets import load_breast_cancer + + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + + +class HasInducedError(Params): + def __init__(self): + super(HasInducedError, self).__init__() + self.inducedError = Param( + self, "inducedError", "Uniformly-distributed error added to feature" + ) + + def getInducedError(self): + return self.getOrDefault(self.inducedError) + + +class InducedErrorModel(Model, HasInducedError): + def __init__(self): + super(InducedErrorModel, self).__init__() + + def _transform(self, dataset): + return dataset.withColumn( + "prediction", dataset.feature + (rand(0) * self.getInducedError()) + ) + + +class InducedErrorEstimator(Estimator, HasInducedError): + def __init__(self, inducedError=1.0): + super(InducedErrorEstimator, self).__init__() + self._set(inducedError=inducedError) + + def _fit(self, dataset): + model = InducedErrorModel() + self._copyValues(model) + return model + + +class CrossValidatorTestsMixin: + def test_gen_avg_and_std_metrics(self): + metrics_all = [ + [1.0, 3.0, 2.0, 4.0], + [3.0, 2.0, 2.0, 4.0], + [3.0, 2.5, 2.1, 8.0], + ] + avg_metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all) + assert np.allclose(avg_metrics, [2.33333333, 2.5, 2.03333333, 5.33333333]) + assert np.allclose(std_metrics, [0.94280904, 0.40824829, 0.04714045, 1.88561808]) + assert isinstance(avg_metrics, list) + assert isinstance(std_metrics, list) + + def test_copy(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="rmse") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator( + estimator=iee, + estimatorParamMaps=grid, + evaluator=evaluator, + numFolds=2, + ) + cvCopied = cv.copy() + for param in [ + lambda x: x.getEstimator().uid, + # SPARK-32092: CrossValidator.copy() needs to copy all existing params + lambda x: x.getNumFolds(), + lambda x: x.getFoldCol(), + lambda x: x.getParallelism(), + lambda x: x.getSeed(), + ]: + self.assertEqual(param(cv), param(cvCopied)) + + cvModel = cv.fit(dataset) + cvModelCopied = cvModel.copy() + for index in range(len(cvModel.avgMetrics)): + self.assertTrue( + abs(cvModel.avgMetrics[index] - cvModelCopied.avgMetrics[index]) < 0.0001 + ) + self.assertTrue(np.allclose(cvModel.stdMetrics, cvModelCopied.stdMetrics)) + # SPARK-32092: CrossValidatorModel.copy() needs to copy all existing params + for param in [lambda x: x.getNumFolds(), lambda x: x.getFoldCol(), lambda x: x.getSeed()]: + self.assertEqual(param(cvModel), param(cvModelCopied)) + + cvModel.avgMetrics[0] = "foo" + self.assertNotEqual( + cvModelCopied.avgMetrics[0], + "foo", + "Changing the original avgMetrics should not affect the copied model", + ) + cvModel.stdMetrics[0] = "foo" + self.assertNotEqual( + cvModelCopied.stdMetrics[0], + "foo", + "Changing the original stdMetrics should not affect the copied model", + ) + + def test_fit_minimize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="rmse") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + bestModel = cvModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(0.0, bestModelMetric, "Best model has RMSE of 0") + + def test_fit_maximize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="r2") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + bestModel = cvModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(1.0, bestModelMetric, "Best model has R-squared of 1") + + @staticmethod + def _check_result(result_dataframe, expected_predictions, expected_probabilities=None): + np.testing.assert_array_equal(list(result_dataframe.prediction), expected_predictions) + if "probability" in result_dataframe.columns: + np.testing.assert_allclose( + list(result_dataframe.probability), + expected_probabilities, + rtol=1e-1, + ) + + def test_crossvalidator_on_pipeline(self): + sk_dataset = load_breast_cancer() + + train_dataset = self.spark.createDataFrame( + zip(sk_dataset.data.tolist(), [int(t) for t in sk_dataset.target]), + schema="features: array, label: long", + ) + + scaler = StandardScaler(inputCol="features", outputCol="scaled_features") + lorv2 = LORV2(numTrainWorkers=2, featuresCol="scaled_features") + pipeline = Pipeline(stages=[scaler, lorv2]) + + grid2 = ParamGridBuilder().addGrid(lorv2.maxIter, [2, 200]).build() + cv = CrossValidator( + estimator=pipeline, + estimatorParamMaps=grid2, + parallelism=2, + evaluator=BinaryClassificationEvaluator(), + ) + cv_model = cv.fit(train_dataset) + transformed_result = ( + cv_model.transform(train_dataset).select("prediction", "probability").toPandas() + ) + expected_transformed_result = ( + cv_model.bestModel.transform(train_dataset) + .select("prediction", "probability") + .toPandas() + ) + pd.testing.assert_frame_equal(transformed_result, expected_transformed_result) + + assert cv_model.bestModel.stages[1].getMaxIter() == 200 + + # trial of index 2 should have better metric value + # because it sets higher `maxIter` param. + assert cv_model.avgMetrics[1] > cv_model.avgMetrics[0] + + def _verify_cv_saved_params(instance, loaded_instance): + assert instance.getEstimator().uid == loaded_instance.getEstimator().uid + assert instance.getEvaluator().uid == loaded_instance.getEvaluator().uid + assert instance.getEstimatorParamMaps() == loaded_instance.getEstimatorParamMaps() + + # Test save / load + with tempfile.TemporaryDirectory() as tmp_dir: + cv.saveToLocal(f"{tmp_dir}/cv") + loaded_cv = CrossValidator.loadFromLocal(f"{tmp_dir}/cv") + + _verify_cv_saved_params(cv, loaded_cv) + + cv_model.saveToLocal(f"{tmp_dir}/cv_model") + loaded_cv_model = CrossValidatorModel.loadFromLocal(f"{tmp_dir}/cv_model") + + _verify_cv_saved_params(cv_model, loaded_cv_model) + + assert cv_model.uid == loaded_cv_model.uid + assert cv_model.bestModel.uid == loaded_cv_model.bestModel.uid + assert cv_model.bestModel.stages[0].uid == loaded_cv_model.bestModel.stages[0].uid + assert cv_model.bestModel.stages[1].uid == loaded_cv_model.bestModel.stages[1].uid + assert loaded_cv_model.bestModel.stages[1].getMaxIter() == 200 + + np.testing.assert_allclose(cv_model.avgMetrics, loaded_cv_model.avgMetrics) + np.testing.assert_allclose(cv_model.stdMetrics, loaded_cv_model.stdMetrics) + + def test_crossvalidator_with_fold_col(self): + sk_dataset = load_breast_cancer() + + train_dataset = self.spark.createDataFrame( + zip( + sk_dataset.data.tolist(), + [int(t) for t in sk_dataset.target], + [int(i % 3) for i in range(len(sk_dataset.target))], + ), + schema="features: array, label: long, fold: long", + ) + + lorv2 = LORV2(numTrainWorkers=2) + + grid2 = ParamGridBuilder().addGrid(lorv2.maxIter, [2, 200]).build() + cv = CrossValidator( + estimator=lorv2, + estimatorParamMaps=grid2, + parallelism=2, + evaluator=BinaryClassificationEvaluator(), + foldCol="fold", + numFolds=3, + ) + cv.fit(train_dataset) + + +class CrossValidatorTests(CrossValidatorTestsMixin, unittest.TestCase): + def setUp(self) -> None: + self.spark = SparkSession.builder.master("local[2]").getOrCreate() + + def tearDown(self) -> None: + self.spark.stop() + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_legacy_mode_tuning import * # noqa: F401,F403 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py new file mode 100644 index 0000000..1855663 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py @@ -0,0 +1,49 @@ +# +# 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 unittest +from pyspark.sql import SparkSession + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + +from pyspark.ml.torch.tests.test_data_loader import TorchDistributorDataLoaderUnitTests + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorBaselineUnitTestsOnConnect(TorchDistributorDataLoaderUnitTests): + def setUp(self) -> None: + self.spark = ( + SparkSession.builder.remote("local[1]") + .config("spark.default.parallelism", "1") + .getOrCreate() + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_parity_torch_data_loader import * # noqa: F401,F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py new file mode 100644 index 0000000..b855332 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py @@ -0,0 +1,152 @@ +# +# 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 os +import shutil +import unittest + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + +from pyspark.sql import SparkSession + +from pyspark.ml.torch.tests.test_distributor import ( + TorchDistributorBaselineUnitTestsMixin, + TorchDistributorLocalUnitTestsMixin, + TorchDistributorDistributedUnitTestsMixin, + TorchWrapperUnitTestsMixin, + set_up_test_dirs, + get_local_mode_conf, + get_distributed_mode_conf, +) + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorBaselineUnitTestsOnConnect( + TorchDistributorBaselineUnitTestsMixin, unittest.TestCase +): + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.remote("local[4]").getOrCreate() + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorLocalUnitTestsOnConnect( + TorchDistributorLocalUnitTestsMixin, unittest.TestCase +): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + builder = SparkSession.builder.appName(cls.__name__) + for k, v in get_local_mode_conf().items(): + builder = builder.config(k, v) + builder = builder.config( + "spark.driver.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + cls.spark = builder.remote("local-cluster[2,2,512]").getOrCreate() + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + def _get_inputs_for_test_local_training_succeeds(self): + return [ + ("0,1,2", 1, True, "0,1,2"), + ("0,1,2", 3, True, "0,1,2"), + ("0,1,2", 2, False, "0,1,2"), + (None, 3, False, "NONE"), + ] + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorLocalUnitTestsIIOnConnect( + TorchDistributorLocalUnitTestsMixin, unittest.TestCase +): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + builder = SparkSession.builder.appName(cls.__name__) + for k, v in get_local_mode_conf().items(): + builder = builder.config(k, v) + + builder = builder.config( + "spark.driver.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + cls.spark = builder.remote("local[4]").getOrCreate() + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + def _get_inputs_for_test_local_training_succeeds(self): + return [ + ("0,1,2", 1, True, "0,1,2"), + ("0,1,2", 3, True, "0,1,2"), + ("0,1,2", 2, False, "0,1,2"), + (None, 3, False, "NONE"), + ] + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorDistributedUnitTestsOnConnect( + TorchDistributorDistributedUnitTestsMixin, unittest.TestCase +): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + builder = SparkSession.builder.appName(cls.__name__) + for k, v in get_distributed_mode_conf().items(): + builder = builder.config(k, v) + + builder = builder.config( + "spark.worker.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + cls.spark = builder.remote("local-cluster[2,2,512]").getOrCreate() + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchWrapperUnitTestsOnConnect(TorchWrapperUnitTestsMixin, unittest.TestCase): + pass + + +if __name__ == "__main__": + from pyspark.ml.tests.connect.test_parity_torch_distributor import * # noqa: F401,F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_algorithms.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_algorithms.py new file mode 100644 index 0000000..fb2507f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_algorithms.py @@ -0,0 +1,456 @@ +# +# 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 shutil import rmtree +import tempfile +import unittest + +import numpy as np + +from pyspark.ml.classification import ( + FMClassifier, + LogisticRegression, + MultilayerPerceptronClassifier, + OneVsRest, +) +from pyspark.ml.clustering import DistributedLDAModel, KMeans, LocalLDAModel, LDA, LDAModel +from pyspark.ml.fpm import FPGrowth +from pyspark.ml.linalg import Matrices, Vectors, DenseVector +from pyspark.ml.recommendation import ALS +from pyspark.ml.regression import GeneralizedLinearRegression, LinearRegression +from pyspark.sql import Row +from pyspark.testing.mlutils import SparkSessionTestCase + + +class LogisticRegressionTest(SparkSessionTestCase): + def test_binomial_logistic_regression_with_bound(self): + + df = self.spark.createDataFrame( + [ + (1.0, 1.0, Vectors.dense(0.0, 5.0)), + (0.0, 2.0, Vectors.dense(1.0, 2.0)), + (1.0, 3.0, Vectors.dense(2.0, 1.0)), + (0.0, 4.0, Vectors.dense(3.0, 3.0)), + ], + ["label", "weight", "features"], + ) + + lor = LogisticRegression( + regParam=0.01, + weightCol="weight", + lowerBoundsOnCoefficients=Matrices.dense(1, 2, [-1.0, -1.0]), + upperBoundsOnIntercepts=Vectors.dense(0.0), + ) + model = lor.fit(df) + self.assertTrue(np.allclose(model.coefficients.toArray(), [-0.2944, -0.0484], atol=1e-4)) + self.assertTrue(np.isclose(model.intercept, 0.0, atol=1e-4)) + + def test_multinomial_logistic_regression_with_bound(self): + + data_path = "data/mllib/sample_multiclass_classification_data.txt" + df = self.spark.read.format("libsvm").load(data_path) + + lor = LogisticRegression( + regParam=0.01, + lowerBoundsOnCoefficients=Matrices.dense(3, 4, range(12)), + upperBoundsOnIntercepts=Vectors.dense(0.0, 0.0, 0.0), + ) + model = lor.fit(df) + expected = [ + [4.593, 4.5516, 9.0099, 12.2904], + [1.0, 8.1093, 7.0, 10.0], + [3.041, 5.0, 8.0, 11.0], + ] + for i in range(0, len(expected)): + self.assertTrue( + np.allclose(model.coefficientMatrix.toArray()[i], expected[i], atol=1e-4) + ) + self.assertTrue( + np.allclose(model.interceptVector.toArray(), [-0.9057, -1.1392, -0.0033], atol=1e-4) + ) + + def test_logistic_regression_with_threshold(self): + + df = self.spark.createDataFrame( + [ + (1.0, 1.0, Vectors.dense(0.0, 5.0)), + (0.0, 2.0, Vectors.dense(1.0, 2.0)), + (1.0, 3.0, Vectors.dense(2.0, 1.0)), + (0.0, 4.0, Vectors.dense(3.0, 3.0)), + ], + ["label", "weight", "features"], + ) + + lor = LogisticRegression(weightCol="weight") + model = lor.fit(df) + + # status changes 1 + for t in [0.0, 0.1, 0.2, 0.5, 1.0]: + model.setThreshold(t).transform(df) + + # status changes 2 + [model.setThreshold(t).predict(Vectors.dense(0.0, 5.0)) for t in [0.0, 0.1, 0.2, 0.5, 1.0]] + + self.assertEqual( + [row.prediction for row in model.setThreshold(0.0).transform(df).collect()], + [1.0, 1.0, 1.0, 1.0], + ) + self.assertEqual( + [row.prediction for row in model.setThreshold(0.5).transform(df).collect()], + [0.0, 1.0, 1.0, 0.0], + ) + self.assertEqual( + [row.prediction for row in model.setThreshold(1.0).transform(df).collect()], + [0.0, 0.0, 0.0, 0.0], + ) + + +class MultilayerPerceptronClassifierTest(SparkSessionTestCase): + def test_raw_and_probability_prediction(self): + + data_path = "data/mllib/sample_multiclass_classification_data.txt" + df = self.spark.read.format("libsvm").load(data_path) + + mlp = MultilayerPerceptronClassifier( + maxIter=100, layers=[4, 5, 4, 3], blockSize=128, seed=123 + ) + model = mlp.fit(df) + test = self.sc.parallelize([Row(features=Vectors.dense(0.1, 0.1, 0.25, 0.25))]).toDF() + result = model.transform(test).head() + expected_prediction = 2.0 + expected_probability = [0.0, 0.0, 1.0] + expected_rawPrediction = [-11.6081922998, -8.15827998691, 22.17757045] + self.assertTrue(result.prediction, expected_prediction) + self.assertTrue(np.allclose(result.probability, expected_probability, atol=1e-4)) + # Use `assert_allclose` to show the value of `result.rawPrediction` in the assertion error + # message + np.testing.assert_allclose( + result.rawPrediction, + expected_rawPrediction, + rtol=0.3, + # Use the same default value as `np.allclose` + atol=1e-08, + ) + + +class OneVsRestTests(SparkSessionTestCase): + def test_copy(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense(1.0, 0.8)), + (1.0, Vectors.sparse(2, [], [])), + (2.0, Vectors.dense(0.5, 0.5)), + ], + ["label", "features"], + ) + lr = LogisticRegression(maxIter=5, regParam=0.01) + ovr = OneVsRest(classifier=lr) + ovr1 = ovr.copy({lr.maxIter: 10}) + self.assertEqual(ovr.getClassifier().getMaxIter(), 5) + self.assertEqual(ovr1.getClassifier().getMaxIter(), 10) + model = ovr.fit(df) + model1 = model.copy({model.predictionCol: "indexed"}) + self.assertEqual(model1.getPredictionCol(), "indexed") + + def test_output_columns(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense(1.0, 0.8)), + (1.0, Vectors.sparse(2, [], [])), + (2.0, Vectors.dense(0.5, 0.5)), + ], + ["label", "features"], + ) + lr = LogisticRegression(maxIter=5, regParam=0.01) + ovr = OneVsRest(classifier=lr, parallelism=1) + model = ovr.fit(df) + output = model.transform(df) + self.assertEqual(output.columns, ["label", "features", "rawPrediction", "prediction"]) + + def test_raw_prediction_column_is_of_vector_type(self): + # SPARK-35142: `OneVsRestModel` outputs raw prediction as a string column + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense(1.0, 0.8)), + (1.0, Vectors.sparse(2, [], [])), + (2.0, Vectors.dense(0.5, 0.5)), + ], + ["label", "features"], + ) + lr = LogisticRegression(maxIter=5, regParam=0.01) + ovr = OneVsRest(classifier=lr, parallelism=1) + model = ovr.fit(df) + row = model.transform(df).head() + self.assertIsInstance(row["rawPrediction"], DenseVector) + + def test_parallelism_does_not_change_output(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense(1.0, 0.8)), + (1.0, Vectors.sparse(2, [], [])), + (2.0, Vectors.dense(0.5, 0.5)), + ], + ["label", "features"], + ) + ovrPar1 = OneVsRest(classifier=LogisticRegression(maxIter=5, regParam=0.01), parallelism=1) + modelPar1 = ovrPar1.fit(df) + ovrPar2 = OneVsRest(classifier=LogisticRegression(maxIter=5, regParam=0.01), parallelism=2) + modelPar2 = ovrPar2.fit(df) + for i, model in enumerate(modelPar1.models): + self.assertTrue( + np.allclose( + model.coefficients.toArray(), + modelPar2.models[i].coefficients.toArray(), + atol=1e-4, + ) + ) + self.assertTrue(np.allclose(model.intercept, modelPar2.models[i].intercept, atol=1e-4)) + + def test_support_for_weightCol(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense(1.0, 0.8), 1.0), + (1.0, Vectors.sparse(2, [], []), 1.0), + (2.0, Vectors.dense(0.5, 0.5), 1.0), + ], + ["label", "features", "weight"], + ) + # classifier inherits hasWeightCol + lr = LogisticRegression(maxIter=5, regParam=0.01) + ovr = OneVsRest(classifier=lr, weightCol="weight") + self.assertIsNotNone(ovr.fit(df)) + # classifier doesn't inherit hasWeightCol + dt = FMClassifier() + ovr2 = OneVsRest(classifier=dt, weightCol="weight") + self.assertIsNotNone(ovr2.fit(df)) + + +class KMeansTests(SparkSessionTestCase): + def test_kmeans_cosine_distance(self): + data = [ + (Vectors.dense([1.0, 1.0]),), + (Vectors.dense([10.0, 10.0]),), + (Vectors.dense([1.0, 0.5]),), + (Vectors.dense([10.0, 4.4]),), + (Vectors.dense([-1.0, 1.0]),), + (Vectors.dense([-100.0, 90.0]),), + ] + df = self.spark.createDataFrame(data, ["features"]) + kmeans = KMeans(k=3, seed=1, distanceMeasure="cosine") + model = kmeans.fit(df) + result = model.transform(df).collect() + self.assertTrue(result[0].prediction == result[1].prediction) + self.assertTrue(result[2].prediction == result[3].prediction) + self.assertTrue(result[4].prediction == result[5].prediction) + + +class LDATest(SparkSessionTestCase): + def _compare(self, m1, m2): + """ + Temp method for comparing instances. + TODO: Replace with generic implementation once SPARK-14706 is merged. + """ + self.assertEqual(m1.uid, m2.uid) + self.assertEqual(type(m1), type(m2)) + self.assertEqual(len(m1.params), len(m2.params)) + for p in m1.params: + if m1.isDefined(p): + self.assertEqual(m1.getOrDefault(p), m2.getOrDefault(p)) + self.assertEqual(p.parent, m2.getParam(p.name).parent) + if isinstance(m1, LDAModel): + self.assertEqual(m1.vocabSize(), m2.vocabSize()) + self.assertEqual(m1.topicsMatrix(), m2.topicsMatrix()) + + def test_persistence(self): + # Test save/load for LDA, LocalLDAModel, DistributedLDAModel. + df = self.spark.createDataFrame( + [ + [1, Vectors.dense([0.0, 1.0])], + [2, Vectors.sparse(2, {0: 1.0})], + ], + ["id", "features"], + ) + # Fit model + lda = LDA(k=2, seed=1, optimizer="em") + distributedModel = lda.fit(df) + self.assertTrue(distributedModel.isDistributed()) + localModel = distributedModel.toLocal() + self.assertFalse(localModel.isDistributed()) + # Define paths + path = tempfile.mkdtemp() + lda_path = path + "/lda" + dist_model_path = path + "/distLDAModel" + local_model_path = path + "/localLDAModel" + # Test LDA + lda.save(lda_path) + lda2 = LDA.load(lda_path) + self._compare(lda, lda2) + # Test DistributedLDAModel + distributedModel.save(dist_model_path) + distributedModel2 = DistributedLDAModel.load(dist_model_path) + self._compare(distributedModel, distributedModel2) + # Test LocalLDAModel + localModel.save(local_model_path) + localModel2 = LocalLDAModel.load(local_model_path) + self._compare(localModel, localModel2) + # Clean up + try: + rmtree(path) + except OSError: + pass + + +class FPGrowthTests(SparkSessionTestCase): + def setUp(self): + super(FPGrowthTests, self).setUp() + self.data = self.spark.createDataFrame( + [([1, 2],), ([1, 2],), ([1, 2, 3],), ([1, 3],)], ["items"] + ) + + def test_association_rules(self): + fp = FPGrowth() + fpm = fp.fit(self.data) + + expected_association_rules = self.spark.createDataFrame( + [([3], [1], 1.0, 1.0, 0.5), ([2], [1], 1.0, 1.0, 0.75)], + ["antecedent", "consequent", "confidence", "lift", "support"], + ) + actual_association_rules = fpm.associationRules + + self.assertEqual(actual_association_rules.subtract(expected_association_rules).count(), 0) + self.assertEqual(expected_association_rules.subtract(actual_association_rules).count(), 0) + + def test_freq_itemsets(self): + fp = FPGrowth() + fpm = fp.fit(self.data) + + expected_freq_itemsets = self.spark.createDataFrame( + [([1], 4), ([2], 3), ([2, 1], 3), ([3], 2), ([3, 1], 2)], ["items", "freq"] + ) + actual_freq_itemsets = fpm.freqItemsets + + self.assertEqual(actual_freq_itemsets.subtract(expected_freq_itemsets).count(), 0) + self.assertEqual(expected_freq_itemsets.subtract(actual_freq_itemsets).count(), 0) + + def tearDown(self): + del self.data + + +class ALSTest(SparkSessionTestCase): + def test_storage_levels(self): + df = self.spark.createDataFrame( + [(0, 0, 4.0), (0, 1, 2.0), (1, 1, 3.0), (1, 2, 4.0), (2, 1, 1.0), (2, 2, 5.0)], + ["user", "item", "rating"], + ) + als = ALS().setMaxIter(1).setRank(1) + # test default params + als.fit(df) + self.assertEqual(als.getIntermediateStorageLevel(), "MEMORY_AND_DISK") + self.assertEqual(als._java_obj.getIntermediateStorageLevel(), "MEMORY_AND_DISK") + self.assertEqual(als.getFinalStorageLevel(), "MEMORY_AND_DISK") + self.assertEqual(als._java_obj.getFinalStorageLevel(), "MEMORY_AND_DISK") + # test non-default params + als.setIntermediateStorageLevel("MEMORY_ONLY_2") + als.setFinalStorageLevel("DISK_ONLY") + als.fit(df) + self.assertEqual(als.getIntermediateStorageLevel(), "MEMORY_ONLY_2") + self.assertEqual(als._java_obj.getIntermediateStorageLevel(), "MEMORY_ONLY_2") + self.assertEqual(als.getFinalStorageLevel(), "DISK_ONLY") + self.assertEqual(als._java_obj.getFinalStorageLevel(), "DISK_ONLY") + + +class GeneralizedLinearRegressionTest(SparkSessionTestCase): + def test_tweedie_distribution(self): + + df = self.spark.createDataFrame( + [ + (1.0, Vectors.dense(0.0, 0.0)), + (1.0, Vectors.dense(1.0, 2.0)), + (2.0, Vectors.dense(0.0, 0.0)), + (2.0, Vectors.dense(1.0, 1.0)), + ], + ["label", "features"], + ) + + glr = GeneralizedLinearRegression(family="tweedie", variancePower=1.6) + model = glr.fit(df) + self.assertTrue(np.allclose(model.coefficients.toArray(), [-0.4645, 0.3402], atol=1e-4)) + self.assertTrue(np.isclose(model.intercept, 0.7841, atol=1e-4)) + + model2 = glr.setLinkPower(-1.0).fit(df) + self.assertTrue(np.allclose(model2.coefficients.toArray(), [-0.6667, 0.5], atol=1e-4)) + self.assertTrue(np.isclose(model2.intercept, 0.6667, atol=1e-4)) + + def test_offset(self): + + df = self.spark.createDataFrame( + [ + (0.2, 1.0, 2.0, Vectors.dense(0.0, 5.0)), + (0.5, 2.1, 0.5, Vectors.dense(1.0, 2.0)), + (0.9, 0.4, 1.0, Vectors.dense(2.0, 1.0)), + (0.7, 0.7, 0.0, Vectors.dense(3.0, 3.0)), + ], + ["label", "weight", "offset", "features"], + ) + + glr = GeneralizedLinearRegression(family="poisson", weightCol="weight", offsetCol="offset") + model = glr.fit(df) + self.assertTrue( + np.allclose(model.coefficients.toArray(), [0.664647, -0.3192581], atol=1e-4) + ) + self.assertTrue(np.isclose(model.intercept, -1.561613, atol=1e-4)) + + +class LinearRegressionTest(SparkSessionTestCase): + def test_linear_regression_with_huber_loss(self): + + data_path = "data/mllib/sample_linear_regression_data.txt" + df = self.spark.read.format("libsvm").load(data_path) + + lir = LinearRegression(loss="huber", epsilon=2.0) + model = lir.fit(df) + + expectedCoefficients = [ + 0.136, + 0.7648, + -0.7761, + 2.4236, + 0.537, + 1.2612, + -0.333, + -0.5694, + -0.6311, + 0.6053, + ] + expectedIntercept = 0.1607 + expectedScale = 9.758 + + self.assertTrue(np.allclose(model.coefficients.toArray(), expectedCoefficients, atol=1e-3)) + self.assertTrue(np.isclose(model.intercept, expectedIntercept, atol=1e-3)) + self.assertTrue(np.isclose(model.scale, expectedScale, atol=1e-3)) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_algorithms import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_base.py new file mode 100644 index 0000000..6c3c51d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_base.py @@ -0,0 +1,96 @@ +# +# 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 unittest + +from pyspark.sql.types import DoubleType, IntegerType +from pyspark.testing.mlutils import ( + MockDataset, + MockEstimator, + MockUnaryTransformer, + MockTransformer, + SparkSessionTestCase, +) + + +class TransformerTests(unittest.TestCase): + def test_transform_invalid_type(self): + transformer = MockTransformer() + data = MockDataset() + self.assertRaises(TypeError, transformer.transform, data, "") + + +class UnaryTransformerTests(SparkSessionTestCase): + def test_unary_transformer_validate_input_type(self): + shiftVal = 3 + transformer = ( + MockUnaryTransformer(shiftVal=shiftVal).setInputCol("input").setOutputCol("output") + ) + + # should not raise any errors + transformer.validateInputType(DoubleType()) + + with self.assertRaises(TypeError): + # passing the wrong input type should raise an error + transformer.validateInputType(IntegerType()) + + def test_unary_transformer_transform(self): + shiftVal = 3 + transformer = ( + MockUnaryTransformer(shiftVal=shiftVal).setInputCol("input").setOutputCol("output") + ) + + df = self.spark.range(0, 10).toDF("input") + df = df.withColumn("input", df.input.cast(dataType="double")) + + transformed_df = transformer.transform(df) + results = transformed_df.select("input", "output").collect() + + for res in results: + self.assertEqual(res.input + shiftVal, res.output) + + +class EstimatorTest(unittest.TestCase): + def setUp(self): + self.estimator = MockEstimator() + self.data = MockDataset() + + def test_fit_invalid_params(self): + invalid_type_parms = "" + self.assertRaises(TypeError, self.estimator.fit, self.data, invalid_type_parms) + + def testDefaultFitMultiple(self): + N = 4 + params = [{self.estimator.fake: i} for i in range(N)] + modelIter = self.estimator.fitMultiple(self.data, params) + indexList = [] + for index, model in modelIter: + self.assertEqual(model.getFake(), index) + indexList.append(index) + self.assertEqual(sorted(indexList), list(range(N))) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_base import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_dl_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_dl_util.py new file mode 100644 index 0000000..abb1c6c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_dl_util.py @@ -0,0 +1,186 @@ +# +# 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 contextlib import contextmanager +import os +import textwrap +from typing import Any, BinaryIO, Callable, Iterator + +import unittest +from parameterized import parameterized +from pyspark import cloudpickle +from pyspark.ml.dl_util import FunctionPickler + + +class TestFunctionPickler(unittest.TestCase): + + # Function that will be used to test pickling. + @staticmethod + def _test_function(x: float, y: float) -> float: + return x**2 + y**2 + + def _check_if_test_function_pickled( + self, + file: BinaryIO, + desired_function: Callable, + output_value: Any, + *arguments, + **key_word_args, + ): + fn, args, kwargs = cloudpickle.load(file) + self.assertEqual(fn, desired_function) + self.assertEqual(args, arguments) + self.assertEqual(kwargs, key_word_args) + fn_output = fn(*args, **kwargs) + self.assertEqual(fn_output, output_value) + + @parameterized.expand( + [ + ("See if it pickles correctly with no path specified", "", ""), + ("See if it pickles correctly with path specified", "silly_bear", ""), + ( + "See if it pickles correctly with both path and save_dir specified", + "silly_bear", + "tmp_dir", + ), + ] + ) + def test_pickle_fn_and_save(self, _: str, file_path_to_save: str, save_dir: str): + x, y = 1, 3 # args of test_function + if save_dir != "": + os.makedirs(save_dir, exist_ok=True) + pickled_fn_path = FunctionPickler.pickle_fn_and_save( + TestFunctionPickler._test_function, file_path_to_save, save_dir, x, y + ) + if file_path_to_save != "": + self.assertEqual(file_path_to_save, pickled_fn_path) + + with open(pickled_fn_path, "rb") as f: + self._check_if_test_function_pickled(f, TestFunctionPickler._test_function, 10, x, y) + os.remove(pickled_fn_path) + + if save_dir != "": + os.rmdir(save_dir) + + def test_getting_output_from_pickle_file(self): + a, b = 2, 0 # arguments for _test_function + pickle_fn_file = FunctionPickler.pickle_fn_and_save( + TestFunctionPickler._test_function, "", "", a, b + ) + fn, args, kwargs = FunctionPickler.get_fn_output(pickle_fn_file) + self.assertEqual(fn, TestFunctionPickler._test_function) + self.assertEqual(len(args), 2) + self.assertEqual(len(kwargs), 0) + self.assertEqual(args[0], a) + self.assertEqual(args[1], b) + self.assertEqual(fn(*args, **kwargs), 4) + os.remove(pickle_fn_file) + + @contextmanager + def create_reference_file( + self, body: str, prefix: str = "", suffix: str = "", fname: str = "reference.py" + ) -> Iterator[None]: + try: + with open(fname, "w") as f: + if prefix != "": + f.write(prefix) + f.write(body) + if suffix != "": + f.write(suffix) + yield + finally: + os.remove(fname) + + def _create_code_snippet_body(self, pickled_fn_path: str, fn_output_save_path: str) -> str: + code_snippet = textwrap.dedent( + f""" + from pyspark import cloudpickle + import os + + if __name__ == "__main__": + with open("{pickled_fn_path}", "rb") as f: + fn, args, kwargs = cloudpickle.load(f) + output = fn(*args, **kwargs) + with open("{fn_output_save_path}", "wb") as f: + cloudpickle.dump(output, f) + """ + ) + return code_snippet + + def _are_two_files_identical(self, fpath1: str, fpath2: str) -> bool: + with open(fpath1, "rb") as f: + contents_one = f.read() + with open(fpath2, "rb") as f: + contents_two = f.read() + self.assertEqual(contents_one, contents_two) + return contents_one == contents_two + + @parameterized.expand( + [ + ("Check if it creates the correct file with no prefix nor suffix", "", ""), + ( + "Check if it creates the correct file with only prefix + body", + "print('hello before')\n", + "", + ), + ( + "Check if it creates the correct file with only suffix + boddy", + "", + "print('goodbye')", + ), + ( + "Check if it creates the correct file prefix, body, and suffix", + "print('hello')\n", + "print('goodbye')\n", + ), + ] + ) + def test_create_fn_run_script(self, mesg: str, prefix_test: str, suffix_test: str): + arg1, arg2 = 3, 4 + pickled_fn_path = FunctionPickler.pickle_fn_and_save( + TestFunctionPickler._test_function, "", "", arg1, arg2 + ) + fn_out_path = "output.pickled" + reference_path = "ref_result_file.py" + test_path = "test_result.py" + body_for_reference = self._create_code_snippet_body(pickled_fn_path, fn_out_path) + + with self.create_reference_file( + body_for_reference, prefix=prefix_test, suffix=suffix_test, fname=reference_path + ) as _: + executable_file_path = FunctionPickler.create_fn_run_script( + pickled_fn_path, + fn_out_path, + test_path, + prefix_code=prefix_test, + suffix_code=suffix_test, + ) + self.assertTrue(self._are_two_files_identical(reference_path, executable_file_path)) + os.remove(executable_file_path) + + os.remove(pickled_fn_path) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_dl_util import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_evaluation.py new file mode 100644 index 0000000..3c5ae3f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_evaluation.py @@ -0,0 +1,77 @@ +# +# 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 unittest + +import numpy as np + +from pyspark.ml.evaluation import ClusteringEvaluator, RegressionEvaluator +from pyspark.ml.linalg import Vectors +from pyspark.sql import Row +from pyspark.testing.mlutils import SparkSessionTestCase + + +class EvaluatorTests(SparkSessionTestCase): + def test_evaluate_invalid_type(self): + evaluator = RegressionEvaluator(metricName="r2") + df = self.spark.createDataFrame([Row(label=1.0, prediction=1.1)]) + invalid_type = "" + self.assertRaises(TypeError, evaluator.evaluate, df, invalid_type) + + def test_java_params(self): + """ + This tests a bug fixed by SPARK-18274 which causes multiple copies + of a Params instance in Python to be linked to the same Java instance. + """ + evaluator = RegressionEvaluator(metricName="r2") + df = self.spark.createDataFrame([Row(label=1.0, prediction=1.1)]) + evaluator.evaluate(df) + self.assertEqual(evaluator._java_obj.getMetricName(), "r2") + evaluatorCopy = evaluator.copy({evaluator.metricName: "mae"}) + evaluator.evaluate(df) + evaluatorCopy.evaluate(df) + self.assertEqual(evaluator._java_obj.getMetricName(), "r2") + self.assertEqual(evaluatorCopy._java_obj.getMetricName(), "mae") + + def test_clustering_evaluator_with_cosine_distance(self): + featureAndPredictions = map( + lambda x: (Vectors.dense(x[0]), x[1]), + [ + ([1.0, 1.0], 1.0), + ([10.0, 10.0], 1.0), + ([1.0, 0.5], 2.0), + ([10.0, 4.4], 2.0), + ([-1.0, 1.0], 3.0), + ([-100.0, 90.0], 3.0), + ], + ) + dataset = self.spark.createDataFrame(featureAndPredictions, ["features", "prediction"]) + evaluator = ClusteringEvaluator(predictionCol="prediction", distanceMeasure="cosine") + self.assertEqual(evaluator.getDistanceMeasure(), "cosine") + self.assertTrue(np.isclose(evaluator.evaluate(dataset), 0.992671213, atol=1e-5)) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_evaluation import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_feature.py new file mode 100644 index 0000000..0051d47 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_feature.py @@ -0,0 +1,401 @@ +# -*- coding: utf-8 -*- +# +# 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 unittest + +from pyspark.ml.feature import ( + Binarizer, + CountVectorizer, + CountVectorizerModel, + HashingTF, + IDF, + NGram, + RFormula, + StopWordsRemover, + StringIndexer, + StringIndexerModel, + VectorSizeHint, +) +from pyspark.ml.linalg import DenseVector, SparseVector, Vectors +from pyspark.sql import Row +from pyspark.testing.utils import QuietTest +from pyspark.testing.mlutils import check_params, SparkSessionTestCase + + +class FeatureTests(SparkSessionTestCase): + def test_binarizer(self): + b0 = Binarizer() + self.assertListEqual( + b0.params, + [b0.inputCol, b0.inputCols, b0.outputCol, b0.outputCols, b0.threshold, b0.thresholds], + ) + self.assertTrue(all([~b0.isSet(p) for p in b0.params])) + self.assertTrue(b0.hasDefault(b0.threshold)) + self.assertEqual(b0.getThreshold(), 0.0) + b0.setParams(inputCol="input", outputCol="output").setThreshold(1.0) + self.assertTrue(not all([b0.isSet(p) for p in b0.params])) + self.assertEqual(b0.getThreshold(), 1.0) + self.assertEqual(b0.getInputCol(), "input") + self.assertEqual(b0.getOutputCol(), "output") + + b0c = b0.copy({b0.threshold: 2.0}) + self.assertEqual(b0c.uid, b0.uid) + self.assertListEqual(b0c.params, b0.params) + self.assertEqual(b0c.getThreshold(), 2.0) + + b1 = Binarizer(threshold=2.0, inputCol="input", outputCol="output") + self.assertNotEqual(b1.uid, b0.uid) + self.assertEqual(b1.getThreshold(), 2.0) + self.assertEqual(b1.getInputCol(), "input") + self.assertEqual(b1.getOutputCol(), "output") + + def test_idf(self): + dataset = self.spark.createDataFrame( + [(DenseVector([1.0, 2.0]),), (DenseVector([0.0, 1.0]),), (DenseVector([3.0, 0.2]),)], + ["tf"], + ) + idf0 = IDF(inputCol="tf") + self.assertListEqual(idf0.params, [idf0.inputCol, idf0.minDocFreq, idf0.outputCol]) + idf0m = idf0.fit(dataset, {idf0.outputCol: "idf"}) + self.assertEqual( + idf0m.uid, idf0.uid, "Model should inherit the UID from its parent estimator." + ) + output = idf0m.transform(dataset) + self.assertIsNotNone(output.head().idf) + self.assertIsNotNone(idf0m.docFreq) + self.assertEqual(idf0m.numDocs, 3) + # Test that parameters transferred to Python Model + check_params(self, idf0m) + + def test_ngram(self): + dataset = self.spark.createDataFrame([Row(input=["a", "b", "c", "d", "e"])]) + ngram0 = NGram(n=4, inputCol="input", outputCol="output") + self.assertEqual(ngram0.getN(), 4) + self.assertEqual(ngram0.getInputCol(), "input") + self.assertEqual(ngram0.getOutputCol(), "output") + transformedDF = ngram0.transform(dataset) + self.assertEqual(transformedDF.head().output, ["a b c d", "b c d e"]) + + def test_stopwordsremover(self): + dataset = self.spark.createDataFrame([Row(input=["a", "panda"])]) + stopWordRemover = StopWordsRemover(inputCol="input", outputCol="output") + # Default + self.assertEqual(stopWordRemover.getInputCol(), "input") + transformedDF = stopWordRemover.transform(dataset) + self.assertEqual(transformedDF.head().output, ["panda"]) + self.assertEqual(type(stopWordRemover.getStopWords()), list) + self.assertTrue(isinstance(stopWordRemover.getStopWords()[0], str)) + # Custom + stopwords = ["panda"] + stopWordRemover.setStopWords(stopwords) + self.assertEqual(stopWordRemover.getInputCol(), "input") + self.assertEqual(stopWordRemover.getStopWords(), stopwords) + transformedDF = stopWordRemover.transform(dataset) + self.assertEqual(transformedDF.head().output, ["a"]) + # with language selection + stopwords = StopWordsRemover.loadDefaultStopWords("turkish") + dataset = self.spark.createDataFrame([Row(input=["acaba", "ama", "biri"])]) + stopWordRemover.setStopWords(stopwords) + self.assertEqual(stopWordRemover.getStopWords(), stopwords) + transformedDF = stopWordRemover.transform(dataset) + self.assertEqual(transformedDF.head().output, []) + # with locale + stopwords = ["BELKİ"] + dataset = self.spark.createDataFrame([Row(input=["belki"])]) + stopWordRemover.setStopWords(stopwords).setLocale("tr") + self.assertEqual(stopWordRemover.getStopWords(), stopwords) + transformedDF = stopWordRemover.transform(dataset) + self.assertEqual(transformedDF.head().output, []) + + def test_count_vectorizer_with_binary(self): + dataset = self.spark.createDataFrame( + [ + ( + 0, + "a a a b b c".split(" "), + SparseVector(3, {0: 1.0, 1: 1.0, 2: 1.0}), + ), + ( + 1, + "a a".split(" "), + SparseVector(3, {0: 1.0}), + ), + ( + 2, + "a b".split(" "), + SparseVector(3, {0: 1.0, 1: 1.0}), + ), + ( + 3, + "c".split(" "), + SparseVector(3, {2: 1.0}), + ), + ], + ["id", "words", "expected"], + ) + cv = CountVectorizer(binary=True, inputCol="words", outputCol="features") + model = cv.fit(dataset) + + transformedList = model.transform(dataset).select("features", "expected").collect() + + for r in transformedList: + feature, expected = r + self.assertEqual(feature, expected) + + def test_count_vectorizer_with_maxDF(self): + dataset = self.spark.createDataFrame( + [ + ( + 0, + "a b c d".split(" "), + SparseVector(3, {0: 1.0, 1: 1.0, 2: 1.0}), + ), + ( + 1, + "a b c".split(" "), + SparseVector(3, {0: 1.0, 1: 1.0}), + ), + ( + 2, + "a b".split(" "), + SparseVector(3, {0: 1.0}), + ), + ( + 3, + "a".split(" "), + SparseVector(3, {}), + ), + ], + ["id", "words", "expected"], + ) + cv = CountVectorizer(inputCol="words", outputCol="features") + model1 = cv.setMaxDF(3).fit(dataset) + self.assertEqual(model1.vocabulary, ["b", "c", "d"]) + + transformedList1 = model1.transform(dataset).select("features", "expected").collect() + + for r in transformedList1: + feature, expected = r + self.assertEqual(feature, expected) + + model2 = cv.setMaxDF(0.75).fit(dataset) + self.assertEqual(model2.vocabulary, ["b", "c", "d"]) + + transformedList2 = model2.transform(dataset).select("features", "expected").collect() + + for r in transformedList2: + feature, expected = r + self.assertEqual(feature, expected) + + def test_count_vectorizer_from_vocab(self): + model = CountVectorizerModel.from_vocabulary( + ["a", "b", "c"], inputCol="words", outputCol="features", minTF=2 + ) + self.assertEqual(model.vocabulary, ["a", "b", "c"]) + self.assertEqual(model.getMinTF(), 2) + + dataset = self.spark.createDataFrame( + [ + ( + 0, + "a a a b b c".split(" "), + SparseVector(3, {0: 3.0, 1: 2.0}), + ), + ( + 1, + "a a".split(" "), + SparseVector(3, {0: 2.0}), + ), + ( + 2, + "a b".split(" "), + SparseVector(3, {}), + ), + ], + ["id", "words", "expected"], + ) + + transformed_list = model.transform(dataset).select("features", "expected").collect() + + for r in transformed_list: + feature, expected = r + self.assertEqual(feature, expected) + + # Test an empty vocabulary + with QuietTest(self.sc): + with self.assertRaisesRegex(Exception, "vocabSize.*invalid.*0"): + CountVectorizerModel.from_vocabulary([], inputCol="words") + + # Test model with default settings can transform + model_default = CountVectorizerModel.from_vocabulary(["a", "b", "c"], inputCol="words") + transformed_list = ( + model_default.transform(dataset) + .select(model_default.getOrDefault(model_default.outputCol)) + .collect() + ) + self.assertEqual(len(transformed_list), 3) + + def test_rformula_force_index_label(self): + df = self.spark.createDataFrame( + [(1.0, 1.0, "a"), (0.0, 2.0, "b"), (1.0, 0.0, "a")], ["y", "x", "s"] + ) + # Does not index label by default since it's numeric type. + rf = RFormula(formula="y ~ x + s") + model = rf.fit(df) + transformedDF = model.transform(df) + self.assertEqual(transformedDF.head().label, 1.0) + # Force to index label. + rf2 = RFormula(formula="y ~ x + s").setForceIndexLabel(True) + model2 = rf2.fit(df) + transformedDF2 = model2.transform(df) + self.assertEqual(transformedDF2.head().label, 0.0) + + def test_rformula_string_indexer_order_type(self): + df = self.spark.createDataFrame( + [(1.0, 1.0, "a"), (0.0, 2.0, "b"), (1.0, 0.0, "a")], ["y", "x", "s"] + ) + rf = RFormula(formula="y ~ x + s", stringIndexerOrderType="alphabetDesc") + self.assertEqual(rf.getStringIndexerOrderType(), "alphabetDesc") + transformedDF = rf.fit(df).transform(df) + observed = transformedDF.select("features").collect() + expected = [[1.0, 0.0], [2.0, 1.0], [0.0, 0.0]] + for i in range(0, len(expected)): + self.assertTrue(all(observed[i]["features"].toArray() == expected[i])) + + def test_string_indexer_handle_invalid(self): + df = self.spark.createDataFrame([(0, "a"), (1, "d"), (2, None)], ["id", "label"]) + + si1 = StringIndexer( + inputCol="label", + outputCol="indexed", + handleInvalid="keep", + stringOrderType="alphabetAsc", + ) + model1 = si1.fit(df) + td1 = model1.transform(df) + actual1 = td1.select("id", "indexed").collect() + expected1 = [Row(id=0, indexed=0.0), Row(id=1, indexed=1.0), Row(id=2, indexed=2.0)] + self.assertEqual(actual1, expected1) + + si2 = si1.setHandleInvalid("skip") + model2 = si2.fit(df) + td2 = model2.transform(df) + actual2 = td2.select("id", "indexed").collect() + expected2 = [Row(id=0, indexed=0.0), Row(id=1, indexed=1.0)] + self.assertEqual(actual2, expected2) + + def test_string_indexer_from_labels(self): + model = StringIndexerModel.from_labels( + ["a", "b", "c"], inputCol="label", outputCol="indexed", handleInvalid="keep" + ) + self.assertEqual(model.labels, ["a", "b", "c"]) + self.assertEqual(model.labelsArray, [("a", "b", "c")]) + + df1 = self.spark.createDataFrame( + [(0, "a"), (1, "c"), (2, None), (3, "b"), (4, "b")], ["id", "label"] + ) + + result1 = model.transform(df1) + actual1 = result1.select("id", "indexed").collect() + expected1 = [ + Row(id=0, indexed=0.0), + Row(id=1, indexed=2.0), + Row(id=2, indexed=3.0), + Row(id=3, indexed=1.0), + Row(id=4, indexed=1.0), + ] + self.assertEqual(actual1, expected1) + + model_empty_labels = StringIndexerModel.from_labels( + [], inputCol="label", outputCol="indexed", handleInvalid="keep" + ) + actual2 = model_empty_labels.transform(df1).select("id", "indexed").collect() + expected2 = [ + Row(id=0, indexed=0.0), + Row(id=1, indexed=0.0), + Row(id=2, indexed=0.0), + Row(id=3, indexed=0.0), + Row(id=4, indexed=0.0), + ] + self.assertEqual(actual2, expected2) + + # Test model with default settings can transform + model_default = StringIndexerModel.from_labels(["a", "b", "c"], inputCol="label") + df2 = self.spark.createDataFrame( + [(0, "a"), (1, "c"), (2, "b"), (3, "b"), (4, "b")], ["id", "label"] + ) + transformed_list = ( + model_default.transform(df2) + .select(model_default.getOrDefault(model_default.outputCol)) + .collect() + ) + self.assertEqual(len(transformed_list), 5) + + def test_vector_size_hint(self): + df = self.spark.createDataFrame( + [ + (0, Vectors.dense([0.0, 10.0, 0.5])), + (1, Vectors.dense([1.0, 11.0, 0.5, 0.6])), + (2, Vectors.dense([2.0, 12.0])), + ], + ["id", "vector"], + ) + + sizeHint = VectorSizeHint(inputCol="vector", handleInvalid="skip") + sizeHint.setSize(3) + self.assertEqual(sizeHint.getSize(), 3) + + output = sizeHint.transform(df).head().vector + expected = DenseVector([0.0, 10.0, 0.5]) + self.assertEqual(output, expected) + + +class HashingTFTest(SparkSessionTestCase): + def test_apply_binary_term_freqs(self): + + df = self.spark.createDataFrame([(0, ["a", "a", "b", "c", "c", "c"])], ["id", "words"]) + n = 10 + hashingTF = HashingTF() + hashingTF.setInputCol("words").setOutputCol("features").setNumFeatures(n).setBinary(True) + output = hashingTF.transform(df) + features = output.select("features").first().features.toArray() + expected = Vectors.dense([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0]).toArray() + for i in range(0, n): + self.assertAlmostEqual( + features[i], + expected[i], + 14, + "Error at " + + str(i) + + ": expected " + + str(expected[i]) + + ", got " + + str(features[i]), + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_feature import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_functions.py new file mode 100644 index 0000000..894db2f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_functions.py @@ -0,0 +1,528 @@ +# +# 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 numpy as np +import pandas as pd +import unittest + +from pyspark.ml.functions import predict_batch_udf +from pyspark.sql.functions import array, struct, col +from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StructType, StructField, FloatType +from pyspark.testing.mlutils import SparkSessionTestCase + + +class PredictBatchUDFTests(SparkSessionTestCase): + def setUp(self): + super(PredictBatchUDFTests, self).setUp() + self.data = np.arange(0, 1000, dtype=np.float64).reshape(-1, 4) + + # 4 scalar columns + self.pdf = pd.DataFrame(self.data, columns=["a", "b", "c", "d"]) + self.df = self.spark.createDataFrame(self.pdf) + + # 1 tensor column of 4 doubles + self.pdf_tensor = pd.DataFrame() + self.pdf_tensor["t1"] = self.pdf.values.tolist() + self.df_tensor1 = self.spark.createDataFrame(self.pdf_tensor) + + # 2 tensor columns of 4 doubles and 3 doubles + self.pdf_tensor["t2"] = self.pdf.drop(columns="d").values.tolist() + self.df_tensor2 = self.spark.createDataFrame(self.pdf_tensor) + + # 4 scalar columns with 1 tensor column + self.pdf_scalar_tensor = self.pdf + self.pdf_scalar_tensor["t1"] = self.pdf.values.tolist() + self.df_scalar_tensor = self.spark.createDataFrame(self.pdf_scalar_tensor) + + def test_identity_single(self): + def make_predict_fn(): + def predict(inputs): + return inputs + + return predict + + identity = predict_batch_udf(make_predict_fn, return_type=DoubleType(), batch_size=5) + + # single column input => single column output (struct) + preds = self.df.withColumn("preds", identity(struct("a"))).toPandas() + self.assertTrue(preds["a"].equals(preds["preds"])) + + # single column input => single column output (col) + preds = self.df.withColumn("preds", identity(col("a"))).toPandas() + self.assertTrue(preds["a"].equals(preds["preds"])) + + # single column input => single column output (str) + preds = self.df.withColumn("preds", identity("a")).toPandas() + self.assertTrue(preds["a"].equals(preds["preds"])) + + # multiple column input, single input => ERROR + with self.assertRaisesRegex(Exception, "Multiple input columns found, but model expected"): + preds = self.df.withColumn("preds", identity("a", "b")).toPandas() + + # batch_size 1 + identity = predict_batch_udf(make_predict_fn, return_type=DoubleType(), batch_size=1) + preds = self.df.withColumn("preds", identity("a")).toPandas() + self.assertTrue(preds["a"].equals(preds["preds"])) + + def test_identity_multi(self): + # single input model + def make_predict_fn(): + def predict(inputs): + return {"a1": inputs[:, 0], "b1": inputs[:, 1]} + + return predict + + identity = predict_batch_udf( + make_predict_fn, + return_type=StructType( + [StructField("a1", DoubleType(), True), StructField("b1", DoubleType(), True)] + ), + batch_size=5, + ) + + # multiple columns using struct, single input => multiple column output + preds = ( + self.df.withColumn("preds", identity(struct("a", "b"))) + .select("a", "b", "preds.*") + .toPandas() + ) + self.assertTrue(preds["a"].equals(preds["a1"])) + self.assertTrue(preds["b"].equals(preds["b1"])) + + # multiple columns, single input => ERROR + with self.assertRaisesRegex(Exception, "Multiple input columns found, but model expected"): + preds = ( + self.df.withColumn("preds", identity("a", "b")) + .select("a", "b", "preds.*") + .toPandas() + ) + + # multiple input model + def predict_batch2_fn(): + def predict(in1, in2): + return {"a1": in1, "b1": in2} + + return predict + + identity2 = predict_batch_udf( + predict_batch2_fn, + return_type=StructType( + [StructField("a1", DoubleType(), True), StructField("b1", DoubleType(), True)] + ), + batch_size=5, + ) + + # multiple columns using struct, multiple inputs => multiple column output + preds = ( + self.df.withColumn("preds", identity2(struct("a", "b"))) + .select("a", "b", "preds.*") + .toPandas() + ) + self.assertTrue(preds["a"].equals(preds["a1"])) + self.assertTrue(preds["b"].equals(preds["b1"])) + + # multiple columns, multiple inputs => multiple column output + preds = ( + self.df.withColumn("preds", identity2(col("a"), col("b"))) + .select("a", "b", "preds.*") + .toPandas() + ) + self.assertTrue(preds["a"].equals(preds["a1"])) + self.assertTrue(preds["b"].equals(preds["b1"])) + + # multiple column input => multiple column output (str) + preds = ( + self.df.withColumn("preds", identity2("a", "b")).select("a", "b", "preds.*").toPandas() + ) + self.assertTrue(preds["a"].equals(preds["a1"])) + self.assertTrue(preds["b"].equals(preds["b1"])) + + def test_batching(self): + batch_size = 10 + + def make_predict_fn(): + def predict(inputs): + batch_size = len(inputs) + # just return the batch size as the "prediction" + outputs = [batch_size for i in inputs] + return np.array(outputs) + + return predict + + identity = predict_batch_udf( + make_predict_fn, return_type=IntegerType(), batch_size=batch_size + ) + + # struct + preds = self.df.withColumn("preds", identity(struct("a"))).toPandas() + batch_sizes = preds["preds"].to_numpy() + self.assertTrue(all(batch_sizes <= batch_size)) + + # col + preds = self.df.withColumn("preds", identity(col("a"))).toPandas() + batch_sizes = preds["preds"].to_numpy() + self.assertTrue(all(batch_sizes <= batch_size)) + + # struct + preds = self.df.withColumn("preds", identity("a")).toPandas() + batch_sizes = preds["preds"].to_numpy() + self.assertTrue(all(batch_sizes <= batch_size)) + + def test_caching(self): + def make_predict_fn(): + # emulate loading a model, this should only be invoked once (per worker process) + fake_output = np.random.random() + + def predict(inputs): + return np.array([fake_output for i in inputs]) + + return predict + + identity = predict_batch_udf(make_predict_fn, return_type=DoubleType(), batch_size=5) + + # results should be the same + df1 = self.df.withColumn("preds", identity(struct("a"))).toPandas() + df2 = self.df.withColumn("preds", identity(struct("a"))).toPandas() + self.assertTrue(df1.equals(df2)) + + identity = predict_batch_udf(make_predict_fn, return_type=DoubleType(), batch_size=5) + + # cache should now be invalidated and results should be different + df3 = self.df.withColumn("preds", identity(struct("a"))).toPandas() + self.assertFalse(df1.equals(df3)) + + def test_transform_scalar(self): + columns = self.df.columns + + # multiple scalar columns, single input, no input_tensor_shapes => single numpy array + def array_sum_fn(): + def predict(inputs): + return np.sum(inputs, axis=1) + + return predict + + sum_cols = predict_batch_udf(array_sum_fn, return_type=DoubleType(), batch_size=5) + preds = self.df.withColumn("preds", sum_cols(struct(*columns))).toPandas() + self.assertTrue(np.array_equal(np.sum(self.data, axis=1), preds["preds"].to_numpy())) + + with self.assertRaisesRegex(Exception, "Multiple input columns found, but model expected"): + preds = self.df.withColumn("preds", sum_cols(*[col(c) for c in columns])).toPandas() + + with self.assertRaisesRegex(Exception, "Multiple input columns found, but model expected"): + preds = self.df.withColumn("preds", sum_cols(*columns)).toPandas() + + # multiple scalar columns, multiple inputs, no input_tensor_shapes => list of numpy arrays + def list_sum_fn(): + def predict(a, b, c, d): + result = sum([a, b, c, d]) + return result + + return predict + + sum_cols = predict_batch_udf(list_sum_fn, return_type=DoubleType(), batch_size=5) + preds = self.df.withColumn("preds", sum_cols(*columns)).toPandas() + self.assertTrue(np.array_equal(np.sum(self.data, axis=1), preds["preds"].to_numpy())) + + # multiple scalar columns, mismatched inputs, no input_tensor_shapes => ERROR + def list_sum_fn(): + def predict(a, b, c): + result = sum([a, b, c]) + return result + + return predict + + sum_cols = predict_batch_udf(list_sum_fn, return_type=DoubleType(), batch_size=5) + with self.assertRaisesRegex(Exception, "Model expected 3 inputs, but received 4 columns"): + preds = self.df.withColumn("preds", sum_cols(*columns)).toPandas() + + # muliple scalar columns with one tensor_input_shape => single numpy array + sum_cols = predict_batch_udf( + array_sum_fn, return_type=DoubleType(), batch_size=5, input_tensor_shapes=[[4]] + ) + preds = self.df.withColumn("preds", sum_cols(struct(*columns))).toPandas() + self.assertTrue(np.array_equal(np.sum(self.data, axis=1), preds["preds"].to_numpy())) + + # muliple scalar columns with wrong tensor_input_shape => ERROR + sum_cols = predict_batch_udf( + array_sum_fn, return_type=DoubleType(), batch_size=5, input_tensor_shapes=[[3]] + ) + with self.assertRaisesRegex(Exception, "Input data does not match expected shape."): + self.df.withColumn("preds", sum_cols(struct(*columns))).toPandas() + + # scalar columns with multiple tensor_input_shapes => ERROR + sum_cols = predict_batch_udf( + array_sum_fn, + return_type=DoubleType(), + batch_size=5, + input_tensor_shapes=[[4], [4]], + ) + with self.assertRaisesRegex(Exception, "Multiple input_tensor_shapes found"): + self.df.withColumn("preds", sum_cols(struct(*columns))).toPandas() + + def test_transform_single_tensor(self): + columns1 = self.df_tensor1.columns + + def array_sum_fn(): + def predict(inputs): + return np.sum(inputs, axis=1) + + return predict + + # tensor column with no input_tensor_shapes => ERROR + sum_cols = predict_batch_udf(array_sum_fn, return_type=DoubleType(), batch_size=5) + with self.assertRaisesRegex(Exception, "Tensor columns require input_tensor_shapes"): + preds = self.df_tensor1.withColumn("preds", sum_cols(struct(*columns1))).toPandas() + + # tensor column with tensor_input_shapes => single numpy array + sum_cols = predict_batch_udf( + array_sum_fn, return_type=DoubleType(), batch_size=5, input_tensor_shapes=[[4]] + ) + preds = self.df_tensor1.withColumn("preds", sum_cols(struct(*columns1))).toPandas() + self.assertTrue(np.array_equal(np.sum(self.data, axis=1), preds["preds"].to_numpy())) + + # tensor column with multiple tensor_input_shapes => ERROR + sum_cols = predict_batch_udf( + array_sum_fn, + return_type=DoubleType(), + batch_size=5, + input_tensor_shapes=[[4], [3]], + ) + with self.assertRaisesRegex(Exception, "Multiple input_tensor_shapes found"): + preds = self.df_tensor1.withColumn("preds", sum_cols(struct(*columns1))).toPandas() + + def test_transform_multi_tensor(self): + def multi_sum_fn(): + def predict(t1, t2): + result = np.sum(t1, axis=1) + np.sum(t2, axis=1) + return result + + return predict + + # multiple tensor columns with tensor_input_shapes => list of numpy arrays + sum_cols = predict_batch_udf( + multi_sum_fn, + return_type=DoubleType(), + batch_size=5, + input_tensor_shapes=[[4], [3]], + ) + preds = self.df_tensor2.withColumn("preds", sum_cols("t1", "t2")).toPandas() + self.assertTrue( + np.array_equal( + np.sum(self.data, axis=1) + np.sum(self.data[:, 0:3], axis=1), + preds["preds"].to_numpy(), + ) + ) + + def test_mixed_input_shapes(self): + def mixed_sum_fn(): + # 4 scalars + 1 tensor + def predict(a, b, c, d, t1): + result = a + b + c + d + np.sum(t1, axis=1) + return result + + return predict + + # dense input_tensor_shapes + sum_cols = predict_batch_udf( + mixed_sum_fn, + return_type=DoubleType(), + batch_size=5, + input_tensor_shapes=[None, None, None, None, [4]], + ) + + preds = self.df_scalar_tensor.withColumn( + "preds", sum_cols("a", "b", "c", "d", "t1") + ).toPandas() + + self.assertTrue( + np.array_equal( + np.sum(self.data, axis=1) * 2, + preds["preds"].to_numpy(), + ) + ) + + # sparse input_tensor_shapes + sum_cols = predict_batch_udf( + mixed_sum_fn, + return_type=DoubleType(), + batch_size=5, + input_tensor_shapes={4: [4]}, + ) + + preds = self.df_scalar_tensor.withColumn( + "preds", sum_cols("a", "b", "c", "d", "t1") + ).toPandas() + + self.assertTrue( + np.array_equal( + np.sum(self.data, axis=1) * 2, + preds["preds"].to_numpy(), + ) + ) + + def test_return_multiple(self): + # columnar form (dictionary of numpy arrays) + def multiples_column_fn(): + def predict(inputs): + return {"x2": inputs * 2, "x3": inputs * 3} + + return predict + + multiples_col = predict_batch_udf( + multiples_column_fn, + return_type=StructType( + [StructField("x2", DoubleType(), True), StructField("x3", DoubleType(), True)] + ), + batch_size=5, + ) + preds = self.df.withColumn("preds", multiples_col("a")).select("a", "preds.*").toPandas() + + self.assertTrue(np.array_equal(self.data[:, 0] * 2, preds["x2"].to_numpy())) + self.assertTrue(np.array_equal(self.data[:, 0] * 3, preds["x3"].to_numpy())) + + # row form: list of dictionaries + def multiples_row_fn(): + def predict(inputs): + return [{"x2": x * 2, "x3": x * 3} for x in inputs] + + return predict + + multiples_row = predict_batch_udf( + multiples_row_fn, + return_type=StructType( + [StructField("x2", DoubleType(), True), StructField("x3", DoubleType(), True)] + ), + batch_size=5, + ) + preds = self.df.withColumn("preds", multiples_row("a")).select("a", "preds.*").toPandas() + + self.assertTrue(np.array_equal(self.data[:, 0] * 2, preds["x2"].to_numpy())) + self.assertTrue(np.array_equal(self.data[:, 0] * 3, preds["x3"].to_numpy())) + + def test_return_struct_with_array_field(self): + # column form + def multiples_with_array_fn(): + def predict(x, y): + return {"x2": x * 2, "y3": y * 3} + + return predict + + multiples_w_array = predict_batch_udf( + multiples_with_array_fn, + return_type=StructType( + [ + StructField("x2", DoubleType(), True), + StructField("y3", ArrayType(DoubleType()), True), + ] + ), + input_tensor_shapes=[[], [3]], + batch_size=5, + ) + preds = ( + self.df.withColumn("preds", multiples_w_array("a", array(["b", "c", "d"]))) + .select("a", "preds.*") + .toPandas() + ) + + self.assertTrue(np.array_equal(self.data[:, 0] * 2, np.array(preds["x2"]))) + self.assertTrue(np.array_equal(self.data[:, 1:4] * 3, np.vstack(preds["y3"]))) + + # row form: list of dictionaries + def multiples_row_array_fn(): + def predict(x, y): + return [{"x2": x * 2, "y3": y * 3} for x, y in zip(x, y)] + + return predict + + multiples_row_array = predict_batch_udf( + multiples_row_array_fn, + return_type=StructType( + [ + StructField("x2", DoubleType(), True), + StructField("y3", ArrayType(DoubleType()), True), + ] + ), + input_tensor_shapes=[[], [3]], + batch_size=5, + ) + + preds = ( + self.df.withColumn("preds", multiples_row_array("a", array(["b", "c", "d"]))) + .select("a", "preds.*") + .toPandas() + ) + + self.assertTrue(np.array_equal(self.data[:, 0] * 2, np.array(preds["x2"]))) + self.assertTrue(np.array_equal(self.data[:, 1:4] * 3, np.vstack(preds["y3"]))) + + # row form: list of dictionaries, malformed array + def multiples_row_array_fn(): + def predict(x, y): + return [{"x2": x * 2, "y3": np.reshape(y, (-1, 1)) * 3} for x, y in zip(x, y)] + + return predict + + multiples_row_array = predict_batch_udf( + multiples_row_array_fn, + return_type=StructType( + [ + StructField("x2", DoubleType(), True), + StructField("y3", ArrayType(DoubleType()), True), + ] + ), + input_tensor_shapes=[[], [3]], + batch_size=5, + ) + with self.assertRaisesRegex(Exception, "must be one-dimensional"): + preds = ( + self.df.withColumn("preds", multiples_row_array("a", array(["b", "c", "d"]))) + .select("a", "preds.*") + .toPandas() + ) + + def test_single_value_in_batch(self): + # SPARK-42250: batches consisting of single float value should work + df = self.spark.createDataFrame( + [[[0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0]]], schema=["t1", "t2"] + ) + + def make_multi_sum_fn(): + def predict(x1: np.ndarray, x2: np.ndarray) -> np.ndarray: + 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=1, + input_tensor_shapes=[[4], [3]], + ) + + [value] = df.select(multi_sum_udf("t1", "t2")).first() + self.assertEqual(value, 9.0) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_functions import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_image.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_image.py new file mode 100644 index 0000000..86fa46c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_image.py @@ -0,0 +1,82 @@ +# +# 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 unittest + +from pyspark.ml.image import ImageSchema +from pyspark.testing.mlutils import SparkSessionTestCase +from pyspark.sql import Row +from pyspark.testing.utils import QuietTest + + +class ImageFileFormatTest(SparkSessionTestCase): + def test_read_images(self): + data_path = "data/mllib/images/origin/kittens" + df = ( + self.spark.read.format("image") + .option("dropInvalid", True) + .option("recursiveFileLookup", True) + .load(data_path) + ) + self.assertEqual(df.count(), 4) + first_row = df.take(1)[0][0] + # compare `schema.simpleString()` instead of directly compare schema, + # because the df loaded from datasource may change schema column nullability. + self.assertEqual(df.schema.simpleString(), ImageSchema.imageSchema.simpleString()) + self.assertEqual( + df.schema["image"].dataType.simpleString(), ImageSchema.columnSchema.simpleString() + ) + array = ImageSchema.toNDArray(first_row) + self.assertEqual(len(array), first_row[1]) + self.assertEqual(ImageSchema.toImage(array, origin=first_row[0]), first_row) + expected = {"CV_8UC3": 16, "Undefined": -1, "CV_8U": 0, "CV_8UC1": 0, "CV_8UC4": 24} + self.assertEqual(ImageSchema.ocvTypes, expected) + expected = ["origin", "height", "width", "nChannels", "mode", "data"] + self.assertEqual(ImageSchema.imageFields, expected) + self.assertEqual(ImageSchema.undefinedImageType, "Undefined") + + with QuietTest(self.sc): + self.assertRaisesRegex( + TypeError, + "image argument should be pyspark.sql.types.Row; however", + lambda: ImageSchema.toNDArray("a"), + ) + + with QuietTest(self.sc): + self.assertRaisesRegex( + ValueError, + "image argument should have attributes specified in", + lambda: ImageSchema.toNDArray(Row(a=1)), + ) + + with QuietTest(self.sc): + self.assertRaisesRegex( + TypeError, + "array argument should be numpy.ndarray; however, it got", + lambda: ImageSchema.toImage("a"), + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_image import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_linalg.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_linalg.py new file mode 100644 index 0000000..6632d10 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_linalg.py @@ -0,0 +1,409 @@ +# +# 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 unittest +import array as pyarray + +from numpy import arange, array, array_equal, inf, ones, tile, zeros + +from pyspark.serializers import CPickleSerializer +from pyspark.ml.linalg import ( + DenseMatrix, + DenseVector, + MatrixUDT, + SparseMatrix, + SparseVector, + Vector, + VectorUDT, + Vectors, +) +from pyspark.testing.mllibutils import MLlibTestCase +from pyspark.sql import Row +from pyspark.sql.functions import unwrap_udt + + +class VectorTests(MLlibTestCase): + def _test_serialize(self, v): + ser = CPickleSerializer() + self.assertEqual(v, ser.loads(ser.dumps(v))) + jvec = self.sc._jvm.org.apache.spark.ml.python.MLSerDe.loads(bytearray(ser.dumps(v))) + nv = ser.loads(bytes(self.sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(jvec))) + self.assertEqual(v, nv) + vs = [v] * 100 + jvecs = self.sc._jvm.org.apache.spark.ml.python.MLSerDe.loads(bytearray(ser.dumps(vs))) + nvs = ser.loads(bytes(self.sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(jvecs))) + self.assertEqual(vs, nvs) + + def test_serialize(self): + self._test_serialize(DenseVector(range(10))) + self._test_serialize(DenseVector(array([1.0, 2.0, 3.0, 4.0]))) + self._test_serialize(DenseVector(pyarray.array("d", range(10)))) + self._test_serialize(SparseVector(4, {1: 1, 3: 2})) + self._test_serialize(SparseVector(3, {})) + self._test_serialize(DenseMatrix(2, 3, range(6))) + sm1 = SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0]) + self._test_serialize(sm1) + + def test_dot(self): + sv = SparseVector(4, {1: 1, 3: 2}) + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + lst = DenseVector([1, 2, 3, 4]) + mat = array( + [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]] + ) + arr = pyarray.array("d", [0, 1, 2, 3]) + self.assertEqual(10.0, sv.dot(dv)) + self.assertTrue(array_equal(array([3.0, 6.0, 9.0, 12.0]), sv.dot(mat))) + self.assertEqual(30.0, dv.dot(dv)) + self.assertTrue(array_equal(array([10.0, 20.0, 30.0, 40.0]), dv.dot(mat))) + self.assertEqual(30.0, lst.dot(dv)) + self.assertTrue(array_equal(array([10.0, 20.0, 30.0, 40.0]), lst.dot(mat))) + self.assertEqual(7.0, sv.dot(arr)) + + def test_squared_distance(self): + def squared_distance(a, b): + if isinstance(a, Vector): + return a.squared_distance(b) + else: + return b.squared_distance(a) + + sv = SparseVector(4, {1: 1, 3: 2}) + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + lst = DenseVector([4, 3, 2, 1]) + lst1 = [4, 3, 2, 1] + arr = pyarray.array("d", [0, 2, 1, 3]) + narr = array([0, 2, 1, 3]) + self.assertEqual(15.0, squared_distance(sv, dv)) + self.assertEqual(25.0, squared_distance(sv, lst)) + self.assertEqual(20.0, squared_distance(dv, lst)) + self.assertEqual(15.0, squared_distance(dv, sv)) + self.assertEqual(25.0, squared_distance(lst, sv)) + self.assertEqual(20.0, squared_distance(lst, dv)) + self.assertEqual(0.0, squared_distance(sv, sv)) + self.assertEqual(0.0, squared_distance(dv, dv)) + self.assertEqual(0.0, squared_distance(lst, lst)) + self.assertEqual(25.0, squared_distance(sv, lst1)) + self.assertEqual(3.0, squared_distance(sv, arr)) + self.assertEqual(3.0, squared_distance(sv, narr)) + + def test_hash(self): + v1 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v2 = SparseVector(4, [(1, 1.0), (3, 5.5)]) + v3 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v4 = SparseVector(4, [(1, 1.0), (3, 2.5)]) + self.assertEqual(hash(v1), hash(v2)) + self.assertEqual(hash(v1), hash(v3)) + self.assertEqual(hash(v2), hash(v3)) + self.assertFalse(hash(v1) == hash(v4)) + self.assertFalse(hash(v2) == hash(v4)) + + def test_eq(self): + v1 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v2 = SparseVector(4, [(1, 1.0), (3, 5.5)]) + v3 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v4 = SparseVector(6, [(1, 1.0), (3, 5.5)]) + v5 = DenseVector([0.0, 1.0, 0.0, 2.5]) + v6 = SparseVector(4, [(1, 1.0), (3, 2.5)]) + dm1 = DenseMatrix(2, 2, [2, 0, 0, 0]) + sm1 = SparseMatrix(2, 2, [0, 2, 3], [0], [2]) + self.assertEqual(v1, v2) + self.assertEqual(v1, v3) + self.assertFalse(v2 == v4) + self.assertFalse(v1 == v5) + self.assertFalse(v1 == v6) + # this is done as Dense and Sparse matrices can be semantically + # equal while still implementing a different __eq__ method + self.assertEqual(dm1, sm1) + self.assertEqual(sm1, dm1) + + def test_equals(self): + indices = [1, 2, 4] + values = [1.0, 3.0, 2.0] + self.assertTrue(Vectors._equals(indices, values, list(range(5)), [0.0, 1.0, 3.0, 0.0, 2.0])) + self.assertFalse( + Vectors._equals(indices, values, list(range(5)), [0.0, 3.0, 1.0, 0.0, 2.0]) + ) + self.assertFalse(Vectors._equals(indices, values, list(range(5)), [0.0, 3.0, 0.0, 2.0])) + self.assertFalse( + Vectors._equals(indices, values, list(range(5)), [0.0, 1.0, 3.0, 2.0, 2.0]) + ) + + def test_conversion(self): + # numpy arrays should be automatically upcast to float64 + # tests for fix of [SPARK-5089] + v = array([1, 2, 3, 4], dtype="float64") + dv = DenseVector(v) + self.assertTrue(dv.array.dtype == "float64") + v = array([1, 2, 3, 4], dtype="float32") + dv = DenseVector(v) + self.assertTrue(dv.array.dtype == "float64") + + def test_sparse_vector_indexing(self): + sv = SparseVector(5, {1: 1, 3: 2}) + self.assertEqual(sv[0], 0.0) + self.assertEqual(sv[3], 2.0) + self.assertEqual(sv[1], 1.0) + self.assertEqual(sv[2], 0.0) + self.assertEqual(sv[4], 0.0) + self.assertEqual(sv[-1], 0.0) + self.assertEqual(sv[-2], 2.0) + self.assertEqual(sv[-3], 0.0) + self.assertEqual(sv[-5], 0.0) + for ind in [5, -6]: + self.assertRaises(IndexError, sv.__getitem__, ind) + for ind in [7.8, "1"]: + self.assertRaises(TypeError, sv.__getitem__, ind) + + zeros = SparseVector(4, {}) + self.assertEqual(zeros[0], 0.0) + self.assertEqual(zeros[3], 0.0) + for ind in [4, -5]: + self.assertRaises(IndexError, zeros.__getitem__, ind) + + empty = SparseVector(0, {}) + for ind in [-1, 0, 1]: + self.assertRaises(IndexError, empty.__getitem__, ind) + + def test_sparse_vector_iteration(self): + self.assertListEqual(list(SparseVector(3, [], [])), [0.0, 0.0, 0.0]) + self.assertListEqual(list(SparseVector(5, [0, 3], [1.0, 2.0])), [1.0, 0.0, 0.0, 2.0, 0.0]) + + def test_matrix_indexing(self): + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10]) + expected = [[0, 6], [1, 8], [4, 10]] + for i in range(3): + for j in range(2): + self.assertEqual(mat[i, j], expected[i][j]) + + for i, j in [(-1, 0), (4, 1), (3, 4)]: + self.assertRaises(IndexError, mat.__getitem__, (i, j)) + + def test_repr_dense_matrix(self): + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10]) + self.assertTrue(repr(mat), "DenseMatrix(3, 2, [0.0, 1.0, 4.0, 6.0, 8.0, 10.0], False)") + + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10], True) + self.assertTrue(repr(mat), "DenseMatrix(3, 2, [0.0, 1.0, 4.0, 6.0, 8.0, 10.0], False)") + + mat = DenseMatrix(6, 3, zeros(18)) + self.assertTrue( + repr(mat), + "DenseMatrix(6, 3, [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], False)", + ) + + def test_repr_sparse_matrix(self): + sm1t = SparseMatrix( + 3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], isTransposed=True + ) + self.assertTrue( + repr(sm1t), + "SparseMatrix(3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], True)", + ) + + indices = tile(arange(6), 3) + values = ones(18) + sm = SparseMatrix(6, 3, [0, 6, 12, 18], indices, values) + self.assertTrue( + repr(sm), + "SparseMatrix(6, 3, [0, 6, 12, 18], \ + [0, 1, 2, 3, 4, 5, 0, 1, ..., 4, 5, 0, 1, 2, 3, 4, 5], \ + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, ..., \ + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], False)", + ) + + self.assertTrue( + str(sm), + "6 X 3 CSCMatrix\n\ + (0,0) 1.0\n(1,0) 1.0\n(2,0) 1.0\n(3,0) 1.0\n(4,0) 1.0\n(5,0) 1.0\n\ + (0,1) 1.0\n(1,1) 1.0\n(2,1) 1.0\n(3,1) 1.0\n(4,1) 1.0\n(5,1) 1.0\n\ + (0,2) 1.0\n(1,2) 1.0\n(2,2) 1.0\n(3,2) 1.0\n..\n..", + ) + + sm = SparseMatrix(1, 18, zeros(19), [], []) + self.assertTrue( + repr(sm), + "SparseMatrix(1, 18, \ + [0, 0, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0, 0, 0, 0, 0], [], [], False)", + ) + + def test_sparse_matrix(self): + # Test sparse matrix creation. + sm1 = SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0]) + self.assertEqual(sm1.numRows, 3) + self.assertEqual(sm1.numCols, 4) + self.assertEqual(sm1.colPtrs.tolist(), [0, 2, 2, 4, 4]) + self.assertEqual(sm1.rowIndices.tolist(), [1, 2, 1, 2]) + self.assertEqual(sm1.values.tolist(), [1.0, 2.0, 4.0, 5.0]) + self.assertTrue( + repr(sm1), + "SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0], False)", + ) + + # Test indexing + expected = [[0, 0, 0, 0], [1, 0, 4, 0], [2, 0, 5, 0]] + + for i in range(3): + for j in range(4): + self.assertEqual(expected[i][j], sm1[i, j]) + self.assertTrue(array_equal(sm1.toArray(), expected)) + + for i, j in [(-1, 1), (4, 3), (3, 5)]: + self.assertRaises(IndexError, sm1.__getitem__, (i, j)) + + # Test conversion to dense and sparse. + smnew = sm1.toDense().toSparse() + self.assertEqual(sm1.numRows, smnew.numRows) + self.assertEqual(sm1.numCols, smnew.numCols) + self.assertTrue(array_equal(sm1.colPtrs, smnew.colPtrs)) + self.assertTrue(array_equal(sm1.rowIndices, smnew.rowIndices)) + self.assertTrue(array_equal(sm1.values, smnew.values)) + + sm1t = SparseMatrix( + 3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], isTransposed=True + ) + self.assertEqual(sm1t.numRows, 3) + self.assertEqual(sm1t.numCols, 4) + self.assertEqual(sm1t.colPtrs.tolist(), [0, 2, 3, 5]) + self.assertEqual(sm1t.rowIndices.tolist(), [0, 1, 2, 0, 2]) + self.assertEqual(sm1t.values.tolist(), [3.0, 2.0, 4.0, 9.0, 8.0]) + + expected = [[3, 2, 0, 0], [0, 0, 4, 0], [9, 0, 8, 0]] + + for i in range(3): + for j in range(4): + self.assertEqual(expected[i][j], sm1t[i, j]) + self.assertTrue(array_equal(sm1t.toArray(), expected)) + + def test_dense_matrix_is_transposed(self): + mat1 = DenseMatrix(3, 2, [0, 4, 1, 6, 3, 9], isTransposed=True) + mat = DenseMatrix(3, 2, [0, 1, 3, 4, 6, 9]) + self.assertEqual(mat1, mat) + + expected = [[0, 4], [1, 6], [3, 9]] + for i in range(3): + for j in range(2): + self.assertEqual(mat1[i, j], expected[i][j]) + self.assertTrue(array_equal(mat1.toArray(), expected)) + + sm = mat1.toSparse() + self.assertTrue(array_equal(sm.rowIndices, [1, 2, 0, 1, 2])) + self.assertTrue(array_equal(sm.colPtrs, [0, 2, 5])) + self.assertTrue(array_equal(sm.values, [1, 3, 4, 6, 9])) + + def test_norms(self): + a = DenseVector([0, 2, 3, -1]) + self.assertAlmostEqual(a.norm(2), 3.742, 3) + self.assertTrue(a.norm(1), 6) + self.assertTrue(a.norm(inf), 3) + a = SparseVector(4, [0, 2], [3, -4]) + self.assertAlmostEqual(a.norm(2), 5) + self.assertTrue(a.norm(1), 7) + self.assertTrue(a.norm(inf), 4) + + tmp = SparseVector(4, [0, 2], [3, 0]) + self.assertEqual(tmp.numNonzeros(), 1) + + +class VectorUDTTests(MLlibTestCase): + + dv0 = DenseVector([]) + dv1 = DenseVector([1.0, 2.0]) + sv0 = SparseVector(2, [], []) + sv1 = SparseVector(2, [1], [2.0]) + udt = VectorUDT() + + def test_json_schema(self): + self.assertEqual(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + + def test_serialization(self): + for v in [self.dv0, self.dv1, self.sv0, self.sv1]: + self.assertEqual(v, self.udt.deserialize(self.udt.serialize(v))) + + def test_infer_schema(self): + rdd = self.sc.parallelize( + [Row(label=1.0, features=self.dv1), Row(label=0.0, features=self.sv1)] + ) + df = rdd.toDF() + schema = df.schema + field = [f for f in schema.fields if f.name == "features"][0] + self.assertEqual(field.dataType, self.udt) + vectors = df.rdd.map(lambda p: p.features).collect() + self.assertEqual(len(vectors), 2) + for v in vectors: + if isinstance(v, SparseVector): + self.assertEqual(v, self.sv1) + elif isinstance(v, DenseVector): + self.assertEqual(v, self.dv1) + else: + raise TypeError("expecting a vector but got %r of type %r" % (v, type(v))) + + def test_unwrap_udt(self): + df = self.spark.createDataFrame( + [(Vectors.dense(1.0, 2.0, 3.0),), (Vectors.sparse(3, {1: 1.0, 2: 5.5}),)], + ["vec"], + ) + results = df.select(unwrap_udt("vec").alias("v2")).collect() + unwrapped_vec = Row("type", "size", "indices", "values") + expected = [ + Row(v2=unwrapped_vec(1, None, None, [1.0, 2.0, 3.0])), + Row(v2=unwrapped_vec(0, 3, [1, 2], [1.0, 5.5])), + ] + self.assertEquals(results, expected) + + +class MatrixUDTTests(MLlibTestCase): + + dm1 = DenseMatrix(3, 2, [0, 1, 4, 5, 9, 10]) + dm2 = DenseMatrix(3, 2, [0, 1, 4, 5, 9, 10], isTransposed=True) + sm1 = SparseMatrix(1, 1, [0, 1], [0], [2.0]) + sm2 = SparseMatrix(2, 1, [0, 0, 1], [0], [5.0], isTransposed=True) + udt = MatrixUDT() + + def test_json_schema(self): + self.assertEqual(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + + def test_serialization(self): + for m in [self.dm1, self.dm2, self.sm1, self.sm2]: + self.assertEqual(m, self.udt.deserialize(self.udt.serialize(m))) + + def test_infer_schema(self): + rdd = self.sc.parallelize([("dense", self.dm1), ("sparse", self.sm1)]) + df = rdd.toDF() + schema = df.schema + self.assertTrue(schema.fields[1].dataType, self.udt) + matrices = df.rdd.map(lambda x: x._2).collect() + self.assertEqual(len(matrices), 2) + for m in matrices: + if isinstance(m, DenseMatrix): + self.assertTrue(m, self.dm1) + elif isinstance(m, SparseMatrix): + self.assertTrue(m, self.sm1) + else: + raise ValueError("Expected a matrix but got type %r" % type(m)) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_linalg import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_model_cache.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_model_cache.py new file mode 100644 index 0000000..f37c73e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_model_cache.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark.ml.model_cache import ModelCache +from pyspark.testing.mlutils import SparkSessionTestCase +from uuid import uuid4 + + +class ModelCacheTests(SparkSessionTestCase): + def setUp(self): + super(ModelCacheTests, self).setUp() + + def test_cache(self): + def predict_fn(inputs): + return inputs + + # add 10 items, expect last 3 items in cache + uuids = [uuid4() for i in range(10)] + for uuid in uuids: + ModelCache.add(uuid, predict_fn) + + self.assertTrue(len(ModelCache._models) == 3) + self.assertTrue(list(ModelCache._models.keys()) == uuids[7:10]) + + # get item, expect it to become most recently used + _ = ModelCache.get(uuids[8]) + expected_uuids = uuids[7:8] + uuids[9:10] + [uuids[8]] + self.assertTrue(list(ModelCache._models.keys()) == expected_uuids) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_model_cache import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_param.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_param.py new file mode 100644 index 0000000..8df50a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_param.py @@ -0,0 +1,441 @@ +# -*- coding: utf-8 -*- +# +# 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 inspect +import array as pyarray +import unittest + +import numpy as np + +from pyspark import keyword_only +from pyspark.ml.classification import LogisticRegression +from pyspark.ml.clustering import KMeans +from pyspark.ml.feature import ( + Binarizer, + Bucketizer, + ElementwiseProduct, + IndexToString, + MaxAbsScaler, + VectorSlicer, + Word2Vec, +) +from pyspark.ml.linalg import DenseVector, SparseVector, Vectors +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import HasInputCol, HasMaxIter, HasSeed +from pyspark.ml.regression import LinearRegressionModel, GeneralizedLinearRegressionModel +from pyspark.ml.wrapper import JavaParams +from pyspark.testing.mlutils import check_params, PySparkTestCase, SparkSessionTestCase + + +class ParamTypeConversionTests(PySparkTestCase): + """ + Test that param type conversion happens. + """ + + def test_int(self): + lr = LogisticRegression(maxIter=5.0) + self.assertEqual(lr.getMaxIter(), 5) + self.assertTrue(type(lr.getMaxIter()) == int) + self.assertRaises(TypeError, lambda: LogisticRegression(maxIter="notAnInt")) + self.assertRaises(TypeError, lambda: LogisticRegression(maxIter=5.1)) + + def test_float(self): + lr = LogisticRegression(tol=1) + self.assertEqual(lr.getTol(), 1.0) + self.assertTrue(type(lr.getTol()) == float) + self.assertRaises(TypeError, lambda: LogisticRegression(tol="notAFloat")) + + def test_vector(self): + ewp = ElementwiseProduct(scalingVec=[1, 3]) + self.assertEqual(ewp.getScalingVec(), DenseVector([1.0, 3.0])) + ewp = ElementwiseProduct(scalingVec=np.array([1.2, 3.4])) + self.assertEqual(ewp.getScalingVec(), DenseVector([1.2, 3.4])) + self.assertRaises(TypeError, lambda: ElementwiseProduct(scalingVec=["a", "b"])) + + def test_list(self): + lst = [0, 1] + for lst_like in [ + lst, + np.array(lst), + DenseVector(lst), + SparseVector(len(lst), range(len(lst)), lst), + pyarray.array("l", lst), + range(2), + tuple(lst), + ]: + converted = TypeConverters.toList(lst_like) + self.assertEqual(type(converted), list) + self.assertListEqual(converted, lst) + + def test_list_int(self): + for indices in [ + [1.0, 2.0], + np.array([1.0, 2.0]), + DenseVector([1.0, 2.0]), + SparseVector(2, {0: 1.0, 1: 2.0}), + range(1, 3), + (1.0, 2.0), + pyarray.array("d", [1.0, 2.0]), + ]: + vs = VectorSlicer(indices=indices) + self.assertListEqual(vs.getIndices(), [1, 2]) + self.assertTrue(all([type(v) == int for v in vs.getIndices()])) + self.assertRaises(TypeError, lambda: VectorSlicer(indices=["a", "b"])) + + def test_list_float(self): + b = Bucketizer(splits=[1, 4]) + self.assertEqual(b.getSplits(), [1.0, 4.0]) + self.assertTrue(all([type(v) == float for v in b.getSplits()])) + self.assertRaises(TypeError, lambda: Bucketizer(splits=["a", 1.0])) + + def test_list_list_float(self): + b = Bucketizer(splitsArray=[[-0.1, 0.5, 3], [-5, 1.5]]) + self.assertEqual(b.getSplitsArray(), [[-0.1, 0.5, 3.0], [-5.0, 1.5]]) + self.assertTrue(all([type(v) == list for v in b.getSplitsArray()])) + self.assertTrue(all([type(v) == float for v in b.getSplitsArray()[0]])) + self.assertTrue(all([type(v) == float for v in b.getSplitsArray()[1]])) + self.assertRaises(TypeError, lambda: Bucketizer(splitsArray=["a", 1.0])) + self.assertRaises(TypeError, lambda: Bucketizer(splitsArray=[[-5, 1.5], ["a", 1.0]])) + + def test_list_string(self): + for labels in [np.array(["a", "b"]), ["a", "b"], np.array(["a", "b"])]: + idx_to_string = IndexToString(labels=labels) + self.assertListEqual(idx_to_string.getLabels(), ["a", "b"]) + self.assertRaises(TypeError, lambda: IndexToString(labels=["a", 2])) + + def test_string(self): + lr = LogisticRegression() + for col in ["features", "features", np.str_("features")]: + lr.setFeaturesCol(col) + self.assertEqual(lr.getFeaturesCol(), "features") + self.assertRaises(TypeError, lambda: LogisticRegression(featuresCol=2.3)) + + def test_bool(self): + self.assertRaises(TypeError, lambda: LogisticRegression(fitIntercept=1)) + self.assertRaises(TypeError, lambda: LogisticRegression(fitIntercept="false")) + + +class TestParams(HasMaxIter, HasInputCol, HasSeed): + """ + A subclass of Params mixed with HasMaxIter, HasInputCol and HasSeed. + """ + + @keyword_only + def __init__(self, seed=None): + super(TestParams, self).__init__() + self._setDefault(maxIter=10) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + def setParams(self, seed=None): + """ + setParams(self, seed=None) + Sets params for this test. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +class OtherTestParams(HasMaxIter, HasInputCol, HasSeed): + """ + A subclass of Params mixed with HasMaxIter, HasInputCol and HasSeed. + """ + + @keyword_only + def __init__(self, seed=None): + super(OtherTestParams, self).__init__() + self._setDefault(maxIter=10) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + def setParams(self, seed=None): + """ + setParams(self, seed=None) + Sets params for this test. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + +class HasThrowableProperty(Params): + def __init__(self): + super(HasThrowableProperty, self).__init__() + self.p = Param(self, "none", "empty param") + + @property + def test_property(self): + raise RuntimeError("Test property to raise error when invoked") + + +class ParamTests(SparkSessionTestCase): + def test_copy_new_parent(self): + testParams = TestParams() + # Copying an instantiated param should fail + with self.assertRaises(ValueError): + testParams.maxIter._copy_new_parent(testParams) + # Copying a dummy param should succeed + TestParams.maxIter._copy_new_parent(testParams) + maxIter = testParams.maxIter + self.assertEqual(maxIter.name, "maxIter") + self.assertEqual(maxIter.doc, "max number of iterations (>= 0).") + self.assertTrue(maxIter.parent == testParams.uid) + + def test_param(self): + testParams = TestParams() + maxIter = testParams.maxIter + self.assertEqual(maxIter.name, "maxIter") + self.assertEqual(maxIter.doc, "max number of iterations (>= 0).") + self.assertTrue(maxIter.parent == testParams.uid) + + def test_hasparam(self): + testParams = TestParams() + self.assertTrue(all([testParams.hasParam(p.name) for p in testParams.params])) + self.assertFalse(testParams.hasParam("notAParameter")) + self.assertTrue(testParams.hasParam("maxIter")) + + def test_resolveparam(self): + testParams = TestParams() + self.assertEqual(testParams._resolveParam(testParams.maxIter), testParams.maxIter) + self.assertEqual(testParams._resolveParam("maxIter"), testParams.maxIter) + + self.assertEqual(testParams._resolveParam("maxIter"), testParams.maxIter) + self.assertRaises(AttributeError, lambda: testParams._resolveParam("아")) + + # Invalid type + invalid_type = 1 + self.assertRaises(TypeError, testParams._resolveParam, invalid_type) + + def test_params(self): + testParams = TestParams() + maxIter = testParams.maxIter + inputCol = testParams.inputCol + seed = testParams.seed + + params = testParams.params + self.assertEqual(params, [inputCol, maxIter, seed]) + + self.assertTrue(testParams.hasParam(maxIter.name)) + self.assertTrue(testParams.hasDefault(maxIter)) + self.assertFalse(testParams.isSet(maxIter)) + self.assertTrue(testParams.isDefined(maxIter)) + self.assertEqual(testParams.getMaxIter(), 10) + + self.assertTrue(testParams.hasParam(inputCol.name)) + self.assertFalse(testParams.hasDefault(inputCol)) + self.assertFalse(testParams.isSet(inputCol)) + self.assertFalse(testParams.isDefined(inputCol)) + with self.assertRaises(KeyError): + testParams.getInputCol() + + otherParam = Param( + Params._dummy(), + "otherParam", + "Parameter used to test that " + "set raises an error for a non-member parameter.", + typeConverter=TypeConverters.toString, + ) + with self.assertRaises(ValueError): + testParams.set(otherParam, "value") + + # Since the default is normally random, set it to a known number for debug str + testParams._setDefault(seed=41) + + self.assertEqual( + testParams.explainParams(), + "\n".join( + [ + "inputCol: input column name. (undefined)", + "maxIter: max number of iterations (>= 0). (default: 10)", + "seed: random seed. (default: 41)", + ] + ), + ) + + def test_clear_param(self): + df = self.spark.createDataFrame([(Vectors.dense([1.0]),), (Vectors.dense([2.0]),)], ["a"]) + maScaler = MaxAbsScaler(inputCol="a", outputCol="scaled") + model = maScaler.fit(df) + self.assertTrue(model.isSet(model.outputCol)) + self.assertEqual(model.getOutputCol(), "scaled") + model.clear(model.outputCol) + self.assertFalse(model.isSet(model.outputCol)) + self.assertEqual(model.getOutputCol()[:12], "MaxAbsScaler") + output = model.transform(df) + self.assertEqual(model.getOutputCol(), output.schema.names[1]) + + def test_kmeans_param(self): + algo = KMeans() + self.assertEqual(algo.getInitMode(), "k-means||") + algo.setK(10) + self.assertEqual(algo.getK(), 10) + algo.setInitSteps(10) + self.assertEqual(algo.getInitSteps(), 10) + self.assertEqual(algo.getDistanceMeasure(), "euclidean") + algo.setDistanceMeasure("cosine") + self.assertEqual(algo.getDistanceMeasure(), "cosine") + + def test_hasseed(self): + noSeedSpecd = TestParams() + withSeedSpecd = TestParams(seed=42) + other = OtherTestParams() + # Check that we no longer use 42 as the magic number + self.assertNotEqual(noSeedSpecd.getSeed(), 42) + origSeed = noSeedSpecd.getSeed() + # Check that we only compute the seed once + self.assertEqual(noSeedSpecd.getSeed(), origSeed) + # Check that a specified seed is honored + self.assertEqual(withSeedSpecd.getSeed(), 42) + # Check that a different class has a different seed + self.assertNotEqual(other.getSeed(), noSeedSpecd.getSeed()) + + def test_param_property_error(self): + param_store = HasThrowableProperty() + self.assertRaises(RuntimeError, lambda: param_store.test_property) + params = param_store.params # should not invoke the property 'test_property' + self.assertEqual(len(params), 1) + + def test_word2vec_param(self): + model = Word2Vec().setWindowSize(6) + # Check windowSize is set properly + self.assertEqual(model.getWindowSize(), 6) + + def test_copy_param_extras(self): + tp = TestParams(seed=42) + extra = {tp.getParam(TestParams.inputCol.name): "copy_input"} + tp_copy = tp.copy(extra=extra) + self.assertEqual(tp.uid, tp_copy.uid) + self.assertEqual(tp.params, tp_copy.params) + for k, v in extra.items(): + self.assertTrue(tp_copy.isDefined(k)) + self.assertEqual(tp_copy.getOrDefault(k), v) + copied_no_extra = {} + for k, v in tp_copy._paramMap.items(): + if k not in extra: + copied_no_extra[k] = v + self.assertEqual(tp._paramMap, copied_no_extra) + self.assertEqual(tp._defaultParamMap, tp_copy._defaultParamMap) + with self.assertRaises(TypeError): + tp.copy(extra={"unknown_parameter": None}) + with self.assertRaises(TypeError): + tp.copy(extra=["must be a dict"]) + + def test_logistic_regression_check_thresholds(self): + self.assertIsInstance( + LogisticRegression(threshold=0.5, thresholds=[0.5, 0.5]), LogisticRegression + ) + + self.assertRaisesRegex( + ValueError, + "Logistic Regression getThreshold found inconsistent.*$", + LogisticRegression, + threshold=0.42, + thresholds=[0.5, 0.5], + ) + + def test_preserve_set_state(self): + dataset = self.spark.createDataFrame([(0.5,)], ["data"]) + binarizer = Binarizer(inputCol="data") + self.assertFalse(binarizer.isSet("threshold")) + binarizer.transform(dataset) + binarizer._transfer_params_from_java() + self.assertFalse( + binarizer.isSet("threshold"), + "Params not explicitly set should remain unset after transform", + ) + + def test_default_params_transferred(self): + dataset = self.spark.createDataFrame([(0.5,)], ["data"]) + binarizer = Binarizer(inputCol="data") + # intentionally change the pyspark default, but don't set it + binarizer._defaultParamMap[binarizer.outputCol] = "my_default" + result = binarizer.transform(dataset).select("my_default").collect() + self.assertFalse(binarizer.isSet(binarizer.outputCol)) + self.assertEqual(result[0][0], 1.0) + + def test_lr_evaluate_invaild_type(self): + lr = LinearRegressionModel() + invalid_type = "" + self.assertRaises(TypeError, lr.evaluate, invalid_type) + + def test_glr_evaluate_invaild_type(self): + glr = GeneralizedLinearRegressionModel() + invalid_type = "" + self.assertRaises(TypeError, glr.evaluate, invalid_type) + + +class DefaultValuesTests(PySparkTestCase): + """ + Test :py:class:`JavaParams` classes to see if their default Param values match + those in their Scala counterparts. + """ + + def test_java_params(self): + import re + + import pyspark.ml.feature + import pyspark.ml.classification + import pyspark.ml.clustering + import pyspark.ml.evaluation + import pyspark.ml.pipeline + import pyspark.ml.recommendation + import pyspark.ml.regression + + modules = [ + pyspark.ml.feature, + pyspark.ml.classification, + pyspark.ml.clustering, + pyspark.ml.evaluation, + pyspark.ml.pipeline, + pyspark.ml.recommendation, + pyspark.ml.regression, + ] + for module in modules: + for name, cls in inspect.getmembers(module, inspect.isclass): + if ( + not name.endswith("Model") + and not name.endswith("Params") + and issubclass(cls, JavaParams) + and not inspect.isabstract(cls) + and not re.match("_?Java", name) + and name != "_LSH" + and name != "_Selector" + ): + check_params(self, cls(), check_params_exist=True) + + # Additional classes that need explicit construction + from pyspark.ml.feature import CountVectorizerModel, StringIndexerModel + + check_params( + self, CountVectorizerModel.from_vocabulary(["a"], "input"), check_params_exist=True + ) + check_params( + self, StringIndexerModel.from_labels(["a", "b"], "input"), check_params_exist=True + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_param import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_persistence.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_persistence.py new file mode 100644 index 0000000..406180d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_persistence.py @@ -0,0 +1,546 @@ +# +# 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 json +from shutil import rmtree +import tempfile +import unittest + +from pyspark.ml import Transformer +from pyspark.ml.classification import ( + DecisionTreeClassifier, + FMClassifier, + FMClassificationModel, + LogisticRegression, + MultilayerPerceptronClassifier, + MultilayerPerceptronClassificationModel, + OneVsRest, + OneVsRestModel, +) +from pyspark.ml.clustering import KMeans +from pyspark.ml.feature import Binarizer, Bucketizer, HashingTF, PCA +from pyspark.ml.linalg import Vectors +from pyspark.ml.param import Params +from pyspark.ml.pipeline import Pipeline, PipelineModel +from pyspark.ml.regression import ( + DecisionTreeRegressor, + GeneralizedLinearRegression, + GeneralizedLinearRegressionModel, + LinearRegression, +) +from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWriter +from pyspark.ml.wrapper import JavaParams +from pyspark.testing.mlutils import MockUnaryTransformer, SparkSessionTestCase + + +class TestDefaultSolver(SparkSessionTestCase): + def test_multilayer_load(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense([0.0, 0.0])), + (1.0, Vectors.dense([0.0, 1.0])), + (1.0, Vectors.dense([1.0, 0.0])), + (0.0, Vectors.dense([1.0, 1.0])), + ], + ["label", "features"], + ) + + mlp = MultilayerPerceptronClassifier(layers=[2, 2, 2], seed=123) + model = mlp.fit(df) + self.assertEqual(model.getSolver(), "l-bfgs") + transformed1 = model.transform(df) + path = tempfile.mkdtemp() + model_path = path + "/mlp" + model.save(model_path) + model2 = MultilayerPerceptronClassificationModel.load(model_path) + self.assertEqual(model2.getSolver(), "l-bfgs") + transformed2 = model2.transform(df) + self.assertEqual(transformed1.take(4), transformed2.take(4)) + + def test_fm_load(self): + df = self.spark.createDataFrame( + [(1.0, Vectors.dense(1.0)), (0.0, Vectors.sparse(1, [], []))], ["label", "features"] + ) + fm = FMClassifier(factorSize=2, maxIter=50, stepSize=2.0) + model = fm.fit(df) + self.assertEqual(model.getSolver(), "adamW") + transformed1 = model.transform(df) + path = tempfile.mkdtemp() + model_path = path + "/fm" + model.save(model_path) + model2 = FMClassificationModel.load(model_path) + self.assertEqual(model2.getSolver(), "adamW") + transformed2 = model2.transform(df) + self.assertEqual(transformed1.take(2), transformed2.take(2)) + + def test_glr_load(self): + df = self.spark.createDataFrame( + [ + (1.0, Vectors.dense(0.0, 0.0)), + (1.0, Vectors.dense(1.0, 2.0)), + (2.0, Vectors.dense(0.0, 0.0)), + (2.0, Vectors.dense(1.0, 1.0)), + ], + ["label", "features"], + ) + glr = GeneralizedLinearRegression(family="gaussian", link="identity", linkPredictionCol="p") + model = glr.fit(df) + self.assertEqual(model.getSolver(), "irls") + transformed1 = model.transform(df) + path = tempfile.mkdtemp() + model_path = path + "/glr" + model.save(model_path) + model2 = GeneralizedLinearRegressionModel.load(model_path) + self.assertEqual(model2.getSolver(), "irls") + transformed2 = model2.transform(df) + self.assertEqual(transformed1.take(4), transformed2.take(4)) + + +class PersistenceTest(SparkSessionTestCase): + def test_linear_regression(self): + lr = LinearRegression(maxIter=1) + path = tempfile.mkdtemp() + lr_path = path + "/lr" + lr.save(lr_path) + lr2 = LinearRegression.load(lr_path) + self.assertEqual(lr.uid, lr2.uid) + self.assertEqual(type(lr.uid), type(lr2.uid)) + self.assertEqual( + lr2.uid, + lr2.maxIter.parent, + "Loaded LinearRegression instance uid (%s) did not match Param's uid (%s)" + % (lr2.uid, lr2.maxIter.parent), + ) + self.assertEqual( + lr._defaultParamMap[lr.maxIter], + lr2._defaultParamMap[lr2.maxIter], + "Loaded LinearRegression instance default params did not match " + "original defaults", + ) + try: + rmtree(path) + except OSError: + pass + + def test_linear_regression_pmml_basic(self): + # Most of the validation is done in the Scala side, here we just check + # that we output text rather than parquet (e.g. that the format flag + # was respected). + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + lr = LinearRegression(maxIter=1) + model = lr.fit(df) + path = tempfile.mkdtemp() + lr_path = path + "/lr-pmml" + model.write().format("pmml").save(lr_path) + pmml_text_list = self.sc.textFile(lr_path).collect() + pmml_text = "\n".join(pmml_text_list) + self.assertIn("Apache Spark", pmml_text) + self.assertIn("PMML", pmml_text) + + def test_logistic_regression(self): + lr = LogisticRegression(maxIter=1) + path = tempfile.mkdtemp() + lr_path = path + "/logreg" + lr.save(lr_path) + lr2 = LogisticRegression.load(lr_path) + self.assertEqual( + lr2.uid, + lr2.maxIter.parent, + "Loaded LogisticRegression instance uid (%s) " + "did not match Param's uid (%s)" % (lr2.uid, lr2.maxIter.parent), + ) + self.assertEqual( + lr._defaultParamMap[lr.maxIter], + lr2._defaultParamMap[lr2.maxIter], + "Loaded LogisticRegression instance default params did not match " + + "original defaults", + ) + try: + rmtree(path) + except OSError: + pass + + def test_kmeans(self): + kmeans = KMeans(k=2, seed=1) + path = tempfile.mkdtemp() + km_path = path + "/km" + kmeans.save(km_path) + kmeans2 = KMeans.load(km_path) + self.assertEqual(kmeans.uid, kmeans2.uid) + self.assertEqual(type(kmeans.uid), type(kmeans2.uid)) + self.assertEqual( + kmeans2.uid, + kmeans2.k.parent, + "Loaded KMeans instance uid (%s) did not match Param's uid (%s)" + % (kmeans2.uid, kmeans2.k.parent), + ) + self.assertEqual( + kmeans._defaultParamMap[kmeans.k], + kmeans2._defaultParamMap[kmeans2.k], + "Loaded KMeans instance default params did not match " + "original defaults", + ) + try: + rmtree(path) + except OSError: + pass + + def test_kmean_pmml_basic(self): + # Most of the validation is done in the Scala side, here we just check + # that we output text rather than parquet (e.g. that the format flag + # was respected). + data = [ + (Vectors.dense([0.0, 0.0]),), + (Vectors.dense([1.0, 1.0]),), + (Vectors.dense([9.0, 8.0]),), + (Vectors.dense([8.0, 9.0]),), + ] + df = self.spark.createDataFrame(data, ["features"]) + kmeans = KMeans(k=2, seed=1) + model = kmeans.fit(df) + path = tempfile.mkdtemp() + km_path = path + "/km-pmml" + model.write().format("pmml").save(km_path) + pmml_text_list = self.sc.textFile(km_path).collect() + pmml_text = "\n".join(pmml_text_list) + self.assertIn("Apache Spark", pmml_text) + self.assertIn("PMML", pmml_text) + + def _compare_params(self, m1, m2, param): + """ + Compare 2 ML Params instances for the given param, and assert both have the same param value + and parent. The param must be a parameter of m1. + """ + # Prevent key not found error in case of some param in neither paramMap nor defaultParamMap. + if m1.isDefined(param): + paramValue1 = m1.getOrDefault(param) + paramValue2 = m2.getOrDefault(m2.getParam(param.name)) + if isinstance(paramValue1, Params): + self._compare_pipelines(paramValue1, paramValue2) + else: + self.assertEqual(paramValue1, paramValue2) # for general types param + # Assert parents are equal + self.assertEqual(param.parent, m2.getParam(param.name).parent) + else: + # If m1 is not defined param, then m2 should not, too. See SPARK-14931. + self.assertFalse(m2.isDefined(m2.getParam(param.name))) + + def _compare_pipelines(self, m1, m2): + """ + Compare 2 ML types, asserting that they are equivalent. + This currently supports: + - basic types + - Pipeline, PipelineModel + - OneVsRest, OneVsRestModel + This checks: + - uid + - type + - Param values and parents + """ + self.assertEqual(m1.uid, m2.uid) + self.assertEqual(type(m1), type(m2)) + if isinstance(m1, JavaParams) or isinstance(m1, Transformer): + self.assertEqual(len(m1.params), len(m2.params)) + for p in m1.params: + self._compare_params(m1, m2, p) + elif isinstance(m1, Pipeline): + self.assertEqual(len(m1.getStages()), len(m2.getStages())) + for s1, s2 in zip(m1.getStages(), m2.getStages()): + self._compare_pipelines(s1, s2) + elif isinstance(m1, PipelineModel): + self.assertEqual(len(m1.stages), len(m2.stages)) + for s1, s2 in zip(m1.stages, m2.stages): + self._compare_pipelines(s1, s2) + elif isinstance(m1, OneVsRest) or isinstance(m1, OneVsRestModel): + for p in m1.params: + self._compare_params(m1, m2, p) + if isinstance(m1, OneVsRestModel): + self.assertEqual(len(m1.models), len(m2.models)) + for x, y in zip(m1.models, m2.models): + self._compare_pipelines(x, y) + elif isinstance(m1, Params): + # Test on python backend Estimator/Transformer/Model/Evaluator + self.assertEqual(len(m1.params), len(m2.params)) + for p in m1.params: + self._compare_params(m1, m2, p) + else: + raise RuntimeError("_compare_pipelines does not yet support type: %s" % type(m1)) + + def test_pipeline_persistence(self): + """ + Pipeline[HashingTF, PCA] + """ + temp_path = tempfile.mkdtemp() + + try: + df = self.spark.createDataFrame([(["a", "b", "c"],), (["c", "d", "e"],)], ["words"]) + tf = HashingTF(numFeatures=10, inputCol="words", outputCol="features") + pca = PCA(k=2, inputCol="features", outputCol="pca_features") + pl = Pipeline(stages=[tf, pca]) + model = pl.fit(df) + + pipeline_path = temp_path + "/pipeline" + pl.save(pipeline_path) + loaded_pipeline = Pipeline.load(pipeline_path) + self._compare_pipelines(pl, loaded_pipeline) + + model_path = temp_path + "/pipeline-model" + model.save(model_path) + loaded_model = PipelineModel.load(model_path) + self._compare_pipelines(model, loaded_model) + finally: + try: + rmtree(temp_path) + except OSError: + pass + + def test_nested_pipeline_persistence(self): + """ + Pipeline[HashingTF, Pipeline[PCA]] + """ + temp_path = tempfile.mkdtemp() + + try: + df = self.spark.createDataFrame([(["a", "b", "c"],), (["c", "d", "e"],)], ["words"]) + tf = HashingTF(numFeatures=10, inputCol="words", outputCol="features") + pca = PCA(k=2, inputCol="features", outputCol="pca_features") + p0 = Pipeline(stages=[pca]) + pl = Pipeline(stages=[tf, p0]) + model = pl.fit(df) + + pipeline_path = temp_path + "/pipeline" + pl.save(pipeline_path) + loaded_pipeline = Pipeline.load(pipeline_path) + self._compare_pipelines(pl, loaded_pipeline) + + model_path = temp_path + "/pipeline-model" + model.save(model_path) + loaded_model = PipelineModel.load(model_path) + self._compare_pipelines(model, loaded_model) + finally: + try: + rmtree(temp_path) + except OSError: + pass + + def test_python_transformer_pipeline_persistence(self): + """ + Pipeline[MockUnaryTransformer, Binarizer] + """ + temp_path = tempfile.mkdtemp() + + try: + df = self.spark.range(0, 10).toDF("input") + tf = MockUnaryTransformer(shiftVal=2).setInputCol("input").setOutputCol("shiftedInput") + tf2 = Binarizer(threshold=6, inputCol="shiftedInput", outputCol="binarized") + pl = Pipeline(stages=[tf, tf2]) + model = pl.fit(df) + + pipeline_path = temp_path + "/pipeline" + pl.save(pipeline_path) + loaded_pipeline = Pipeline.load(pipeline_path) + self._compare_pipelines(pl, loaded_pipeline) + + model_path = temp_path + "/pipeline-model" + model.save(model_path) + loaded_model = PipelineModel.load(model_path) + self._compare_pipelines(model, loaded_model) + finally: + try: + rmtree(temp_path) + except OSError: + pass + + def _run_test_onevsrest(self, LogisticRegressionCls): + temp_path = tempfile.mkdtemp() + df = self.spark.createDataFrame( + [ + (0.0, 0.5, Vectors.dense(1.0, 0.8)), + (1.0, 0.5, Vectors.sparse(2, [], [])), + (2.0, 1.0, Vectors.dense(0.5, 0.5)), + ] + * 10, + ["label", "wt", "features"], + ) + + lr = LogisticRegressionCls(maxIter=5, regParam=0.01) + OneVsRest(classifier=lr) + + def reload_and_compare(ovr, suffix): + model = ovr.fit(df) + ovrPath = temp_path + "/{}".format(suffix) + ovr.save(ovrPath) + loadedOvr = OneVsRest.load(ovrPath) + self._compare_pipelines(ovr, loadedOvr) + modelPath = temp_path + "/{}Model".format(suffix) + model.save(modelPath) + loadedModel = OneVsRestModel.load(modelPath) + self._compare_pipelines(model, loadedModel) + + reload_and_compare(OneVsRest(classifier=lr), "ovr") + reload_and_compare(OneVsRest(classifier=lr).setWeightCol("wt"), "ovrw") + + def test_onevsrest(self): + from pyspark.testing.mlutils import DummyLogisticRegression + + self._run_test_onevsrest(LogisticRegression) + self._run_test_onevsrest(DummyLogisticRegression) + + def test_decisiontree_classifier(self): + dt = DecisionTreeClassifier(maxDepth=1) + path = tempfile.mkdtemp() + dtc_path = path + "/dtc" + dt.save(dtc_path) + dt2 = DecisionTreeClassifier.load(dtc_path) + self.assertEqual( + dt2.uid, + dt2.maxDepth.parent, + "Loaded DecisionTreeClassifier instance uid (%s) " + "did not match Param's uid (%s)" % (dt2.uid, dt2.maxDepth.parent), + ) + self.assertEqual( + dt._defaultParamMap[dt.maxDepth], + dt2._defaultParamMap[dt2.maxDepth], + "Loaded DecisionTreeClassifier instance default params did not match " + + "original defaults", + ) + try: + rmtree(path) + except OSError: + pass + + def test_decisiontree_regressor(self): + dt = DecisionTreeRegressor(maxDepth=1) + path = tempfile.mkdtemp() + dtr_path = path + "/dtr" + dt.save(dtr_path) + dt2 = DecisionTreeClassifier.load(dtr_path) + self.assertEqual( + dt2.uid, + dt2.maxDepth.parent, + "Loaded DecisionTreeRegressor instance uid (%s) " + "did not match Param's uid (%s)" % (dt2.uid, dt2.maxDepth.parent), + ) + self.assertEqual( + dt._defaultParamMap[dt.maxDepth], + dt2._defaultParamMap[dt2.maxDepth], + "Loaded DecisionTreeRegressor instance default params did not match " + + "original defaults", + ) + try: + rmtree(path) + except OSError: + pass + + def test_default_read_write(self): + temp_path = tempfile.mkdtemp() + + lr = LogisticRegression() + lr.setMaxIter(50) + lr.setThreshold(0.75) + writer = DefaultParamsWriter(lr) + + savePath = temp_path + "/lr" + writer.save(savePath) + + reader = DefaultParamsReadable.read() + lr2 = reader.load(savePath) + + self.assertEqual(lr.uid, lr2.uid) + self.assertEqual(lr.extractParamMap(), lr2.extractParamMap()) + + # test overwrite + lr.setThreshold(0.8) + writer.overwrite().save(savePath) + + reader = DefaultParamsReadable.read() + lr3 = reader.load(savePath) + + self.assertEqual(lr.uid, lr3.uid) + self.assertEqual(lr.extractParamMap(), lr3.extractParamMap()) + + def test_default_read_write_default_params(self): + lr = LogisticRegression() + self.assertFalse(lr.isSet(lr.getParam("threshold"))) + + lr.setMaxIter(50) + lr.setThreshold(0.75) + + # `threshold` is set by user, default param `predictionCol` is not set by user. + self.assertTrue(lr.isSet(lr.getParam("threshold"))) + self.assertFalse(lr.isSet(lr.getParam("predictionCol"))) + self.assertTrue(lr.hasDefault(lr.getParam("predictionCol"))) + + writer = DefaultParamsWriter(lr) + metadata = json.loads(writer._get_metadata_to_save(lr, self.sc)) + self.assertTrue("defaultParamMap" in metadata) + + reader = DefaultParamsReadable.read() + metadataStr = json.dumps(metadata, separators=[",", ":"]) + loadedMetadata = reader._parseMetaData( + metadataStr, + ) + reader.getAndSetParams(lr, loadedMetadata) + + self.assertTrue(lr.isSet(lr.getParam("threshold"))) + self.assertFalse(lr.isSet(lr.getParam("predictionCol"))) + self.assertTrue(lr.hasDefault(lr.getParam("predictionCol"))) + + # manually create metadata without `defaultParamMap` section. + del metadata["defaultParamMap"] + metadataStr = json.dumps(metadata, separators=[",", ":"]) + loadedMetadata = reader._parseMetaData( + metadataStr, + ) + with self.assertRaisesRegex(AssertionError, "`defaultParamMap` section not found"): + reader.getAndSetParams(lr, loadedMetadata) + + # Prior to 2.4.0, metadata doesn't have `defaultParamMap`. + metadata["sparkVersion"] = "2.3.0" + metadataStr = json.dumps(metadata, separators=[",", ":"]) + loadedMetadata = reader._parseMetaData( + metadataStr, + ) + reader.getAndSetParams(lr, loadedMetadata) + + # Test for SPARK-35542 fix. + def test_save_and_load_on_nested_list_params(self): + temp_path = tempfile.mkdtemp() + splitsArray = [ + [-float("inf"), 0.5, 1.4, float("inf")], + [-float("inf"), 0.1, 1.2, float("inf")], + ] + bucketizer = Bucketizer( + splitsArray=splitsArray, inputCols=["values", "values"], outputCols=["b1", "b2"] + ) + savePath = temp_path + "/bk" + bucketizer.write().overwrite().save(savePath) + loadedBucketizer = Bucketizer.load(savePath) + assert loadedBucketizer.getSplitsArray() == splitsArray + + +if __name__ == "__main__": + from pyspark.ml.tests.test_persistence import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_pipeline.py new file mode 100644 index 0000000..afc900c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_pipeline.py @@ -0,0 +1,71 @@ +# +# 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 unittest + +from pyspark.ml.pipeline import Pipeline +from pyspark.testing.mlutils import MockDataset, MockEstimator, MockTransformer, PySparkTestCase + + +class PipelineTests(PySparkTestCase): + def test_pipeline(self): + dataset = MockDataset() + estimator0 = MockEstimator() + transformer1 = MockTransformer() + estimator2 = MockEstimator() + transformer3 = MockTransformer() + pipeline = Pipeline(stages=[estimator0, transformer1, estimator2, transformer3]) + pipeline_model = pipeline.fit(dataset, {estimator0.fake: 0, transformer1.fake: 1}) + model0, transformer1, model2, transformer3 = pipeline_model.stages + self.assertEqual(0, model0.dataset_index) + self.assertEqual(0, model0.getFake()) + self.assertEqual(1, transformer1.dataset_index) + self.assertEqual(1, transformer1.getFake()) + self.assertEqual(2, dataset.index) + self.assertIsNone(model2.dataset_index, "The last model shouldn't be called in fit.") + self.assertIsNone( + transformer3.dataset_index, "The last transformer shouldn't be called in fit." + ) + dataset = pipeline_model.transform(dataset) + self.assertEqual(2, model0.dataset_index) + self.assertEqual(3, transformer1.dataset_index) + self.assertEqual(4, model2.dataset_index) + self.assertEqual(5, transformer3.dataset_index) + self.assertEqual(6, dataset.index) + + def test_identity_pipeline(self): + dataset = MockDataset() + + def doTransform(pipeline): + pipeline_model = pipeline.fit(dataset) + return pipeline_model.transform(dataset) + + # check that empty pipeline did not perform any transformation + self.assertEqual(dataset.index, doTransform(Pipeline(stages=[])).index) + # check that failure to set stages param will raise KeyError for missing param + self.assertRaises(KeyError, lambda: doTransform(Pipeline())) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_pipeline import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_stat.py new file mode 100644 index 0000000..6bab41b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_stat.py @@ -0,0 +1,52 @@ +# +# 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 unittest + +from pyspark.ml.linalg import Vectors +from pyspark.ml.stat import ChiSquareTest +from pyspark.sql import DataFrame +from pyspark.testing.mlutils import SparkSessionTestCase + + +class ChiSquareTestTests(SparkSessionTestCase): + def test_chisquaretest(self): + data = [ + [0, Vectors.dense([0, 1, 2])], + [1, Vectors.dense([1, 1, 1])], + [2, Vectors.dense([2, 1, 0])], + ] + df = self.spark.createDataFrame(data, ["label", "feat"]) + res = ChiSquareTest.test(df, "feat", "label") + # This line is hitting the collect bug described in #17218, commented for now. + # pValues = res.select("degreesOfFreedom").collect()) + self.assertIsInstance(res, DataFrame) + fieldNames = set(field.name for field in res.schema.fields) + expectedFields = ["pValues", "degreesOfFreedom", "statistics"] + self.assertTrue(all(field in fieldNames for field in expectedFields)) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_stat import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_training_summary.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_training_summary.py new file mode 100644 index 0000000..5704d71 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_training_summary.py @@ -0,0 +1,494 @@ +# +# 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 unittest + +from pyspark.ml.classification import ( + BinaryLogisticRegressionSummary, + BinaryRandomForestClassificationSummary, + FMClassifier, + FMClassificationSummary, + LinearSVC, + LinearSVCSummary, + LogisticRegression, + LogisticRegressionSummary, + MultilayerPerceptronClassifier, + MultilayerPerceptronClassificationSummary, + RandomForestClassificationSummary, + RandomForestClassifier, +) +from pyspark.ml.clustering import BisectingKMeans, GaussianMixture, KMeans +from pyspark.ml.linalg import Vectors +from pyspark.ml.regression import GeneralizedLinearRegression, LinearRegression +from pyspark.sql import DataFrame +from pyspark.testing.mlutils import SparkSessionTestCase + + +class TrainingSummaryTest(SparkSessionTestCase): + def test_linear_regression_summary(self): + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + lr = LinearRegression( + maxIter=5, regParam=0.0, solver="normal", weightCol="weight", fitIntercept=False + ) + model = lr.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertEqual(s.totalIterations, 0) + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.predictionCol, "prediction") + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.featuresCol, "features") + objHist = s.objectiveHistory + self.assertTrue(isinstance(objHist, list) and isinstance(objHist[0], float)) + self.assertAlmostEqual(s.explainedVariance, 0.25, 2) + self.assertAlmostEqual(s.meanAbsoluteError, 0.0) + self.assertAlmostEqual(s.meanSquaredError, 0.0) + self.assertAlmostEqual(s.rootMeanSquaredError, 0.0) + self.assertAlmostEqual(s.r2, 1.0, 2) + self.assertAlmostEqual(s.r2adj, 1.0, 2) + self.assertTrue(isinstance(s.residuals, DataFrame)) + self.assertEqual(s.numInstances, 2) + self.assertEqual(s.degreesOfFreedom, 1) + devResiduals = s.devianceResiduals + self.assertTrue(isinstance(devResiduals, list) and isinstance(devResiduals[0], float)) + coefStdErr = s.coefficientStandardErrors + self.assertTrue(isinstance(coefStdErr, list) and isinstance(coefStdErr[0], float)) + tValues = s.tValues + self.assertTrue(isinstance(tValues, list) and isinstance(tValues[0], float)) + pValues = s.pValues + self.assertTrue(isinstance(pValues, list) and isinstance(pValues[0], float)) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned + # The child class LinearRegressionTrainingSummary runs full test + sameSummary = model.evaluate(df) + self.assertAlmostEqual(sameSummary.explainedVariance, s.explainedVariance) + + def test_glr_summary(self): + from pyspark.ml.linalg import Vectors + + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + glr = GeneralizedLinearRegression( + family="gaussian", link="identity", weightCol="weight", fitIntercept=False + ) + model = glr.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertEqual(s.numIterations, 1) # this should default to a single iteration of WLS + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.predictionCol, "prediction") + self.assertEqual(s.numInstances, 2) + self.assertTrue(isinstance(s.residuals(), DataFrame)) + self.assertTrue(isinstance(s.residuals("pearson"), DataFrame)) + coefStdErr = s.coefficientStandardErrors + self.assertTrue(isinstance(coefStdErr, list) and isinstance(coefStdErr[0], float)) + tValues = s.tValues + self.assertTrue(isinstance(tValues, list) and isinstance(tValues[0], float)) + pValues = s.pValues + self.assertTrue(isinstance(pValues, list) and isinstance(pValues[0], float)) + self.assertEqual(s.degreesOfFreedom, 1) + self.assertEqual(s.residualDegreeOfFreedom, 1) + self.assertEqual(s.residualDegreeOfFreedomNull, 2) + self.assertEqual(s.rank, 1) + self.assertTrue(isinstance(s.solver, str)) + self.assertTrue(isinstance(s.aic, float)) + self.assertTrue(isinstance(s.deviance, float)) + self.assertTrue(isinstance(s.nullDeviance, float)) + self.assertTrue(isinstance(s.dispersion, float)) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned + # The child class GeneralizedLinearRegressionTrainingSummary runs full test + sameSummary = model.evaluate(df) + self.assertAlmostEqual(sameSummary.deviance, s.deviance) + + def test_binary_logistic_regression_summary(self): + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + lr = LogisticRegression(maxIter=5, regParam=0.01, weightCol="weight", fitIntercept=False) + model = lr.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.probabilityCol, "probability") + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.featuresCol, "features") + self.assertEqual(s.predictionCol, "prediction") + objHist = s.objectiveHistory + self.assertTrue(isinstance(objHist, list) and isinstance(objHist[0], float)) + self.assertGreater(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertTrue(isinstance(s.roc, DataFrame)) + self.assertAlmostEqual(s.areaUnderROC, 1.0, 2) + self.assertTrue(isinstance(s.pr, DataFrame)) + self.assertTrue(isinstance(s.fMeasureByThreshold, DataFrame)) + self.assertTrue(isinstance(s.precisionByThreshold, DataFrame)) + self.assertTrue(isinstance(s.recallByThreshold, DataFrame)) + self.assertAlmostEqual(s.accuracy, 1.0, 2) + self.assertAlmostEqual(s.weightedTruePositiveRate, 1.0, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.0, 2) + self.assertAlmostEqual(s.weightedRecall, 1.0, 2) + self.assertAlmostEqual(s.weightedPrecision, 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 1.0, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, BinaryLogisticRegressionSummary)) + self.assertAlmostEqual(sameSummary.areaUnderROC, s.areaUnderROC) + + def test_multiclass_logistic_regression_summary(self): + df = self.spark.createDataFrame( + [ + (1.0, 2.0, Vectors.dense(1.0)), + (0.0, 2.0, Vectors.sparse(1, [], [])), + (2.0, 2.0, Vectors.dense(2.0)), + (2.0, 2.0, Vectors.dense(1.9)), + ], + ["label", "weight", "features"], + ) + lr = LogisticRegression(maxIter=5, regParam=0.01, weightCol="weight", fitIntercept=False) + model = lr.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.probabilityCol, "probability") + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.featuresCol, "features") + self.assertEqual(s.predictionCol, "prediction") + objHist = s.objectiveHistory + self.assertTrue(isinstance(objHist, list) and isinstance(objHist[0], float)) + self.assertGreater(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertAlmostEqual(s.accuracy, 0.75, 2) + self.assertAlmostEqual(s.weightedTruePositiveRate, 0.75, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.25, 2) + self.assertAlmostEqual(s.weightedRecall, 0.75, 2) + self.assertAlmostEqual(s.weightedPrecision, 0.583, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 0.65, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 0.65, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, LogisticRegressionSummary)) + self.assertFalse(isinstance(sameSummary, BinaryLogisticRegressionSummary)) + self.assertAlmostEqual(sameSummary.accuracy, s.accuracy) + + def test_linear_svc_summary(self): + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0, 1.0, 1.0)), (0.0, 2.0, Vectors.dense(1.0, 2.0, 3.0))], + ["label", "weight", "features"], + ) + svc = LinearSVC(maxIter=5, weightCol="weight") + model = svc.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary() + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.scoreCol, "rawPrediction") + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.predictionCol, "prediction") + objHist = s.objectiveHistory + self.assertTrue(isinstance(objHist, list) and isinstance(objHist[0], float)) + self.assertGreater(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertTrue(isinstance(s.roc, DataFrame)) + self.assertAlmostEqual(s.areaUnderROC, 1.0, 2) + self.assertTrue(isinstance(s.pr, DataFrame)) + self.assertTrue(isinstance(s.fMeasureByThreshold, DataFrame)) + self.assertTrue(isinstance(s.precisionByThreshold, DataFrame)) + self.assertTrue(isinstance(s.recallByThreshold, DataFrame)) + print(s.weightedTruePositiveRate) + self.assertAlmostEqual(s.weightedTruePositiveRate, 1.0, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.0, 2) + self.assertAlmostEqual(s.weightedRecall, 1.0, 2) + self.assertAlmostEqual(s.weightedPrecision, 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 1.0, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, LinearSVCSummary)) + self.assertAlmostEqual(sameSummary.areaUnderROC, s.areaUnderROC) + + def test_binary_randomforest_classification_summary(self): + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + rf = RandomForestClassifier(weightCol="weight") + model = rf.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.predictionCol, "prediction") + self.assertEqual(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertTrue(isinstance(s.roc, DataFrame)) + self.assertAlmostEqual(s.areaUnderROC, 1.0, 2) + self.assertTrue(isinstance(s.pr, DataFrame)) + self.assertTrue(isinstance(s.fMeasureByThreshold, DataFrame)) + self.assertTrue(isinstance(s.precisionByThreshold, DataFrame)) + self.assertTrue(isinstance(s.recallByThreshold, DataFrame)) + self.assertAlmostEqual(s.accuracy, 1.0, 2) + self.assertAlmostEqual(s.weightedTruePositiveRate, 1.0, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.0, 2) + self.assertAlmostEqual(s.weightedRecall, 1.0, 2) + self.assertAlmostEqual(s.weightedPrecision, 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 1.0, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, BinaryRandomForestClassificationSummary)) + self.assertAlmostEqual(sameSummary.areaUnderROC, s.areaUnderROC) + + def test_multiclass_randomforest_classification_summary(self): + df = self.spark.createDataFrame( + [ + (1.0, 2.0, Vectors.dense(1.0)), + (0.0, 2.0, Vectors.sparse(1, [], [])), + (2.0, 2.0, Vectors.dense(2.0)), + (2.0, 2.0, Vectors.dense(1.9)), + ], + ["label", "weight", "features"], + ) + rf = RandomForestClassifier(weightCol="weight") + model = rf.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.predictionCol, "prediction") + self.assertEqual(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertAlmostEqual(s.accuracy, 1.0, 2) + self.assertAlmostEqual(s.weightedTruePositiveRate, 1.0, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.0, 2) + self.assertAlmostEqual(s.weightedRecall, 1.0, 2) + self.assertAlmostEqual(s.weightedPrecision, 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 1.0, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, RandomForestClassificationSummary)) + self.assertFalse(isinstance(sameSummary, BinaryRandomForestClassificationSummary)) + self.assertAlmostEqual(sameSummary.accuracy, s.accuracy) + + def test_fm_classification_summary(self): + df = self.spark.createDataFrame( + [ + (1.0, Vectors.dense(2.0)), + (0.0, Vectors.dense(2.0)), + (0.0, Vectors.dense(6.0)), + (1.0, Vectors.dense(3.0)), + ], + ["label", "features"], + ) + fm = FMClassifier(maxIter=5) + model = fm.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary() + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.scoreCol, "probability") + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.predictionCol, "prediction") + objHist = s.objectiveHistory + self.assertTrue(isinstance(objHist, list) and isinstance(objHist[0], float)) + self.assertGreater(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertTrue(isinstance(s.roc, DataFrame)) + self.assertAlmostEqual(s.areaUnderROC, 0.625, 2) + self.assertTrue(isinstance(s.pr, DataFrame)) + self.assertTrue(isinstance(s.fMeasureByThreshold, DataFrame)) + self.assertTrue(isinstance(s.precisionByThreshold, DataFrame)) + self.assertTrue(isinstance(s.recallByThreshold, DataFrame)) + self.assertAlmostEqual(s.weightedTruePositiveRate, 0.75, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.25, 2) + self.assertAlmostEqual(s.weightedRecall, 0.75, 2) + self.assertAlmostEqual(s.weightedPrecision, 0.8333333333333333, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 0.7333333333333334, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 0.7333333333333334, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, FMClassificationSummary)) + self.assertAlmostEqual(sameSummary.areaUnderROC, s.areaUnderROC) + + def test_mlp_classification_summary(self): + df = self.spark.createDataFrame( + [ + (0.0, Vectors.dense([0.0, 0.0])), + (1.0, Vectors.dense([0.0, 1.0])), + (1.0, Vectors.dense([1.0, 0.0])), + (0.0, Vectors.dense([1.0, 1.0])), + ], + ["label", "features"], + ) + mlp = MultilayerPerceptronClassifier(layers=[2, 2, 2], seed=123) + model = mlp.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary() + # test that api is callable and returns expected types + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.labelCol, "label") + self.assertEqual(s.predictionCol, "prediction") + self.assertGreater(s.totalIterations, 0) + self.assertTrue(isinstance(s.labels, list)) + self.assertTrue(isinstance(s.truePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.falsePositiveRateByLabel, list)) + self.assertTrue(isinstance(s.precisionByLabel, list)) + self.assertTrue(isinstance(s.recallByLabel, list)) + self.assertTrue(isinstance(s.fMeasureByLabel(), list)) + self.assertTrue(isinstance(s.fMeasureByLabel(1.0), list)) + self.assertAlmostEqual(s.accuracy, 1.0, 2) + self.assertAlmostEqual(s.weightedTruePositiveRate, 1.0, 2) + self.assertAlmostEqual(s.weightedFalsePositiveRate, 0.0, 2) + self.assertAlmostEqual(s.weightedRecall, 1.0, 2) + self.assertAlmostEqual(s.weightedPrecision, 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(), 1.0, 2) + self.assertAlmostEqual(s.weightedFMeasure(1.0), 1.0, 2) + # test evaluation (with training dataset) produces a summary with same values + # one check is enough to verify a summary is returned, Scala version runs full test + sameSummary = model.evaluate(df) + self.assertTrue(isinstance(sameSummary, MultilayerPerceptronClassificationSummary)) + self.assertAlmostEqual(sameSummary.accuracy, s.accuracy) + + def test_gaussian_mixture_summary(self): + data = [ + (Vectors.dense(1.0),), + (Vectors.dense(5.0),), + (Vectors.dense(10.0),), + (Vectors.sparse(1, [], []),), + ] + df = self.spark.createDataFrame(data, ["features"]) + gmm = GaussianMixture(k=2) + model = gmm.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.probabilityCol, "probability") + self.assertTrue(isinstance(s.probability, DataFrame)) + self.assertEqual(s.featuresCol, "features") + self.assertEqual(s.predictionCol, "prediction") + self.assertTrue(isinstance(s.cluster, DataFrame)) + self.assertEqual(len(s.clusterSizes), 2) + self.assertEqual(s.k, 2) + self.assertEqual(s.numIter, 3) + + def test_bisecting_kmeans_summary(self): + data = [ + (Vectors.dense(1.0),), + (Vectors.dense(5.0),), + (Vectors.dense(10.0),), + (Vectors.sparse(1, [], []),), + ] + df = self.spark.createDataFrame(data, ["features"]) + bkm = BisectingKMeans(k=2) + model = bkm.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.featuresCol, "features") + self.assertEqual(s.predictionCol, "prediction") + self.assertTrue(isinstance(s.cluster, DataFrame)) + self.assertEqual(len(s.clusterSizes), 2) + self.assertEqual(s.k, 2) + self.assertEqual(s.numIter, 20) + + def test_kmeans_summary(self): + data = [ + (Vectors.dense([0.0, 0.0]),), + (Vectors.dense([1.0, 1.0]),), + (Vectors.dense([9.0, 8.0]),), + (Vectors.dense([8.0, 9.0]),), + ] + df = self.spark.createDataFrame(data, ["features"]) + kmeans = KMeans(k=2, seed=1) + model = kmeans.fit(df) + self.assertTrue(model.hasSummary) + s = model.summary + self.assertTrue(isinstance(s.predictions, DataFrame)) + self.assertEqual(s.featuresCol, "features") + self.assertEqual(s.predictionCol, "prediction") + self.assertTrue(isinstance(s.cluster, DataFrame)) + self.assertEqual(len(s.clusterSizes), 2) + self.assertEqual(s.k, 2) + self.assertEqual(s.numIter, 1) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_training_summary import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_util.py new file mode 100644 index 0000000..55c9738 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_util.py @@ -0,0 +1,85 @@ +# +# 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 unittest + +from pyspark.ml import Pipeline +from pyspark.ml.classification import LogisticRegression, OneVsRest +from pyspark.ml.feature import VectorAssembler +from pyspark.ml.linalg import Vectors +from pyspark.ml.util import MetaAlgorithmReadWrite +from pyspark.testing.mlutils import SparkSessionTestCase + + +class MetaAlgorithmReadWriteTests(SparkSessionTestCase): + def test_getAllNestedStages(self): + def _check_uid_set_equal(stages, expected_stages): + uids = set(map(lambda x: x.uid, stages)) + expected_uids = set(map(lambda x: x.uid, expected_stages)) + self.assertEqual(uids, expected_uids) + + df1 = self.spark.createDataFrame( + [ + (Vectors.dense([1.0, 2.0]), 1.0), + (Vectors.dense([-1.0, -2.0]), 0.0), + ], + ["features", "label"], + ) + df2 = self.spark.createDataFrame( + [ + (1.0, 2.0, 1.0), + (1.0, 2.0, 0.0), + ], + ["a", "b", "label"], + ) + vs = VectorAssembler(inputCols=["a", "b"], outputCol="features") + lr = LogisticRegression() + pipeline = Pipeline(stages=[vs, lr]) + pipelineModel = pipeline.fit(df2) + ova = OneVsRest(classifier=lr) + ovaModel = ova.fit(df1) + + ova_pipeline = Pipeline(stages=[vs, ova]) + nested_pipeline = Pipeline(stages=[ova_pipeline]) + + _check_uid_set_equal( + MetaAlgorithmReadWrite.getAllNestedStages(pipeline), [pipeline, vs, lr] + ) + _check_uid_set_equal( + MetaAlgorithmReadWrite.getAllNestedStages(pipelineModel), + [pipelineModel] + pipelineModel.stages, + ) + _check_uid_set_equal(MetaAlgorithmReadWrite.getAllNestedStages(ova), [ova, lr]) + _check_uid_set_equal( + MetaAlgorithmReadWrite.getAllNestedStages(ovaModel), [ovaModel, lr] + ovaModel.models + ) + _check_uid_set_equal( + MetaAlgorithmReadWrite.getAllNestedStages(nested_pipeline), + [nested_pipeline, ova_pipeline, vs, ova, lr], + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_util import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_wrapper.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_wrapper.py new file mode 100644 index 0000000..33d93c0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/test_wrapper.py @@ -0,0 +1,138 @@ +# +# 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 unittest + +import py4j + +from pyspark.ml.linalg import DenseVector, Vectors +from pyspark.ml.regression import LinearRegression +from pyspark.ml.wrapper import ( + _java2py, + _py2java, + JavaParams, + JavaWrapper, +) +from pyspark.testing.mllibutils import MLlibTestCase +from pyspark.testing.mlutils import SparkSessionTestCase +from pyspark.testing.utils import eventually + + +class JavaWrapperMemoryTests(SparkSessionTestCase): + def test_java_object_gets_detached(self): + df = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + lr = LinearRegression( + maxIter=1, regParam=0.0, solver="normal", weightCol="weight", fitIntercept=False + ) + + model = lr.fit(df) + summary = model.summary + + self.assertIsInstance(model, JavaWrapper) + self.assertIsInstance(summary, JavaWrapper) + self.assertIsInstance(model, JavaParams) + self.assertNotIsInstance(summary, JavaParams) + + error_no_object = "Target Object ID does not exist for this gateway" + + self.assertIn("LinearRegression_", model._java_obj.toString()) + self.assertIn("LinearRegressionTrainingSummary", summary._java_obj.toString()) + + model.__del__() + + def condition(): + with self.assertRaisesRegex(py4j.protocol.Py4JError, error_no_object): + model._java_obj.toString() + self.assertIn("LinearRegressionTrainingSummary", summary._java_obj.toString()) + return True + + eventually(condition, timeout=10, catch_assertions=True) + + try: + summary.__del__() + except BaseException: + pass + + def condition(): + with self.assertRaisesRegex(py4j.protocol.Py4JError, error_no_object): + model._java_obj.toString() + with self.assertRaisesRegex(py4j.protocol.Py4JError, error_no_object): + summary._java_obj.toString() + return True + + eventually(condition, timeout=10, catch_assertions=True) + + +class WrapperTests(MLlibTestCase): + def test_new_java_array(self): + # test array of strings + str_list = ["a", "b", "c"] + java_class = self.sc._gateway.jvm.java.lang.String + java_array = JavaWrapper._new_java_array(str_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), str_list) + # test array of integers + int_list = [1, 2, 3] + java_class = self.sc._gateway.jvm.java.lang.Integer + java_array = JavaWrapper._new_java_array(int_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), int_list) + # test array of floats + float_list = [0.1, 0.2, 0.3] + java_class = self.sc._gateway.jvm.java.lang.Double + java_array = JavaWrapper._new_java_array(float_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), float_list) + # test array of bools + bool_list = [False, True, True] + java_class = self.sc._gateway.jvm.java.lang.Boolean + java_array = JavaWrapper._new_java_array(bool_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), bool_list) + # test array of Java DenseVectors + v1 = DenseVector([0.0, 1.0]) + v2 = DenseVector([1.0, 0.0]) + vec_java_list = [_py2java(self.sc, v1), _py2java(self.sc, v2)] + java_class = self.sc._gateway.jvm.org.apache.spark.ml.linalg.DenseVector + java_array = JavaWrapper._new_java_array(vec_java_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), [v1, v2]) + # test empty array + java_class = self.sc._gateway.jvm.java.lang.Integer + java_array = JavaWrapper._new_java_array([], java_class) + self.assertEqual(_java2py(self.sc, java_array), []) + # test array of array of strings + str_list = [["a", "b", "c"], ["d", "e"], ["f", "g", "h", "i"], []] + expected_str_list = [ + ("a", "b", "c", None), + ("d", "e", None, None), + ("f", "g", "h", "i"), + (None, None, None, None), + ] + java_class = self.sc._gateway.jvm.java.lang.String + java_array = JavaWrapper._new_java_array(str_list, java_class) + self.assertEqual(_java2py(self.sc, java_array), expected_str_list) + + +if __name__ == "__main__": + from pyspark.ml.tests.test_wrapper import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/__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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_basic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_basic.py new file mode 100644 index 0000000..20d63c6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_basic.py @@ -0,0 +1,151 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel +from pyspark.ml.evaluation import BinaryClassificationEvaluator +from pyspark.ml.linalg import Vectors +from pyspark.ml.tuning import ( + CrossValidator, + CrossValidatorModel, + ParamGridBuilder, +) +from pyspark.testing.mlutils import ( + DummyEvaluator, + DummyLogisticRegression, + DummyLogisticRegressionModel, + SparkSessionTestCase, +) +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class CrossValidatorIOBasicTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_trained_model(self, LogisticRegressionCls, LogisticRegressionModelCls): + # This tests saving and loading the trained model only. + # Save/load for CrossValidator will be added later: SPARK-13786 + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + lr = LogisticRegressionCls() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = BinaryClassificationEvaluator() + cv = CrossValidator( + estimator=lr, + estimatorParamMaps=grid, + evaluator=evaluator, + collectSubModels=True, + numFolds=4, + seed=42, + ) + cvModel = cv.fit(dataset) + lrModel = cvModel.bestModel + + lrModelPath = temp_path + "/lrModel" + lrModel.save(lrModelPath) + loadedLrModel = LogisticRegressionModelCls.load(lrModelPath) + self.assertEqual(loadedLrModel.uid, lrModel.uid) + self.assertEqual(loadedLrModel.intercept, lrModel.intercept) + + # SPARK-32092: Saving and then loading CrossValidatorModel should not change the params + cvModelPath = temp_path + "/cvModel" + cvModel.save(cvModelPath) + loadedCvModel = CrossValidatorModel.load(cvModelPath) + for param in [ + lambda x: x.getNumFolds(), + lambda x: x.getFoldCol(), + lambda x: x.getSeed(), + lambda x: len(x.subModels), + ]: + self.assertEqual(param(cvModel), param(loadedCvModel)) + + self.assertTrue(all(loadedCvModel.isSet(param) for param in loadedCvModel.params)) + + # mimic old version CrossValidatorModel (without stdMetrics attribute) + # test loading model backwards compatibility + cvModel2 = cvModel.copy() + cvModel2.stdMetrics = [] + cvModelPath2 = temp_path + "/cvModel2" + cvModel2.save(cvModelPath2) + loadedCvModel2 = CrossValidatorModel.load(cvModelPath2) + assert loadedCvModel2.stdMetrics == [] + + def test_save_load_trained_model(self): + self._run_test_save_load_trained_model(LogisticRegression, LogisticRegressionModel) + self._run_test_save_load_trained_model( + DummyLogisticRegression, DummyLogisticRegressionModel + ) + + def _run_test_save_load_simple_estimator(self, LogisticRegressionCls, evaluatorCls): + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + + lr = LogisticRegressionCls() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = evaluatorCls() + + # test save/load of CrossValidator + cv = CrossValidator(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + cvPath = temp_path + "/cv" + cv.save(cvPath) + loadedCV = CrossValidator.load(cvPath) + self.assertEqual(loadedCV.getEstimator().uid, cv.getEstimator().uid) + self.assertEqual(loadedCV.getEvaluator().uid, cv.getEvaluator().uid) + self.assert_param_maps_equal(loadedCV.getEstimatorParamMaps(), cv.getEstimatorParamMaps()) + + # test save/load of CrossValidatorModel + cvModelPath = temp_path + "/cvModel" + cvModel.save(cvModelPath) + loadedModel = CrossValidatorModel.load(cvModelPath) + self.assertEqual(loadedModel.bestModel.uid, cvModel.bestModel.uid) + + def test_save_load_simple_estimator(self): + self._run_test_save_load_simple_estimator(LogisticRegression, BinaryClassificationEvaluator) + self._run_test_save_load_simple_estimator(DummyLogisticRegression, DummyEvaluator) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_cv_io_basic import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_nested.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_nested.py new file mode 100644 index 0000000..3164588 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_nested.py @@ -0,0 +1,97 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.classification import LogisticRegression, OneVsRest +from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.linalg import Vectors +from pyspark.ml.tuning import ( + CrossValidator, + CrossValidatorModel, + ParamGridBuilder, +) +from pyspark.testing.mlutils import ( + DummyLogisticRegression, + SparkSessionTestCase, +) +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class CrossValidatorIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_nested_estimator(self, LogisticRegressionCls): + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + + ova = OneVsRest(classifier=LogisticRegressionCls()) + lr1 = LogisticRegressionCls().setMaxIter(100) + lr2 = LogisticRegressionCls().setMaxIter(150) + grid = ParamGridBuilder().addGrid(ova.classifier, [lr1, lr2]).build() + evaluator = MulticlassClassificationEvaluator() + + # test save/load of CrossValidator + cv = CrossValidator(estimator=ova, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + cvPath = temp_path + "/cv" + cv.save(cvPath) + loadedCV = CrossValidator.load(cvPath) + self.assert_param_maps_equal(loadedCV.getEstimatorParamMaps(), grid) + self.assertEqual(loadedCV.getEstimator().uid, cv.getEstimator().uid) + self.assertEqual(loadedCV.getEvaluator().uid, cv.getEvaluator().uid) + + originalParamMap = cv.getEstimatorParamMaps() + loadedParamMap = loadedCV.getEstimatorParamMaps() + for i, param in enumerate(loadedParamMap): + for p in param: + if p.name == "classifier": + self.assertEqual(param[p].uid, originalParamMap[i][p].uid) + else: + self.assertEqual(param[p], originalParamMap[i][p]) + + # test save/load of CrossValidatorModel + cvModelPath = temp_path + "/cvModel" + cvModel.save(cvModelPath) + loadedModel = CrossValidatorModel.load(cvModelPath) + self.assert_param_maps_equal(loadedModel.getEstimatorParamMaps(), grid) + self.assertEqual(loadedModel.bestModel.uid, cvModel.bestModel.uid) + + def test_save_load_nested_estimator(self): + self._run_test_save_load_nested_estimator(LogisticRegression) + self._run_test_save_load_nested_estimator(DummyLogisticRegression) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_cv_io_nested import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py new file mode 100644 index 0000000..acfe9c9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py @@ -0,0 +1,143 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.feature import HashingTF, Tokenizer +from pyspark.ml import Pipeline +from pyspark.ml.classification import LogisticRegression, OneVsRest +from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.tuning import ( + CrossValidator, + CrossValidatorModel, + ParamGridBuilder, +) +from pyspark.testing.mlutils import ( + DummyLogisticRegression, + SparkSessionTestCase, +) +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class CrossValidatorIOPipelineTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): + temp_path = tempfile.mkdtemp() + training = self.spark.createDataFrame( + [ + (0, "a b c d e spark", 1.0), + (1, "b d", 0.0), + (2, "spark f g h", 1.0), + (3, "hadoop mapreduce", 0.0), + (4, "b spark who", 1.0), + (5, "g d a y", 0.0), + (6, "spark fly", 1.0), + (7, "was mapreduce", 0.0), + ], + ["id", "text", "label"], + ) + + # Configure an ML pipeline, which consists of tree stages: tokenizer, hashingTF, and lr. + tokenizer = Tokenizer(inputCol="text", outputCol="words") + hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features") + + ova = OneVsRest(classifier=LogisticRegressionCls()) + lr1 = LogisticRegressionCls().setMaxIter(5) + lr2 = LogisticRegressionCls().setMaxIter(10) + + pipeline = Pipeline(stages=[tokenizer, hashingTF, ova]) + + paramGrid = ( + ParamGridBuilder() + .addGrid(hashingTF.numFeatures, [10, 100]) + .addGrid(ova.classifier, [lr1, lr2]) + .build() + ) + + crossval = CrossValidator( + estimator=pipeline, + estimatorParamMaps=paramGrid, + evaluator=MulticlassClassificationEvaluator(), + numFolds=2, + ) # use 3+ folds in practice + cvPath = temp_path + "/cv" + crossval.save(cvPath) + loadedCV = CrossValidator.load(cvPath) + self.assert_param_maps_equal(loadedCV.getEstimatorParamMaps(), paramGrid) + self.assertEqual(loadedCV.getEstimator().uid, crossval.getEstimator().uid) + + # Run cross-validation, and choose the best set of parameters. + cvModel = crossval.fit(training) + + # test save/load of CrossValidatorModel + cvModelPath = temp_path + "/cvModel" + cvModel.save(cvModelPath) + loadedModel = CrossValidatorModel.load(cvModelPath) + self.assertEqual(loadedModel.bestModel.uid, cvModel.bestModel.uid) + self.assertEqual(len(loadedModel.bestModel.stages), len(cvModel.bestModel.stages)) + for loadedStage, originalStage in zip( + loadedModel.bestModel.stages, cvModel.bestModel.stages + ): + self.assertEqual(loadedStage.uid, originalStage.uid) + + # Test nested pipeline + nested_pipeline = Pipeline(stages=[tokenizer, Pipeline(stages=[hashingTF, ova])]) + crossval2 = CrossValidator( + estimator=nested_pipeline, + estimatorParamMaps=paramGrid, + evaluator=MulticlassClassificationEvaluator(), + numFolds=2, + ) # use 3+ folds in practice + cv2Path = temp_path + "/cv2" + crossval2.save(cv2Path) + loadedCV2 = CrossValidator.load(cv2Path) + self.assert_param_maps_equal(loadedCV2.getEstimatorParamMaps(), paramGrid) + self.assertEqual(loadedCV2.getEstimator().uid, crossval2.getEstimator().uid) + + # Run cross-validation, and choose the best set of parameters. + cvModel2 = crossval2.fit(training) + # test save/load of CrossValidatorModel + cvModelPath2 = temp_path + "/cvModel2" + cvModel2.save(cvModelPath2) + loadedModel2 = CrossValidatorModel.load(cvModelPath2) + self.assertEqual(loadedModel2.bestModel.uid, cvModel2.bestModel.uid) + loaded_nested_pipeline_model = loadedModel2.bestModel.stages[1] + original_nested_pipeline_model = cvModel2.bestModel.stages[1] + self.assertEqual(loaded_nested_pipeline_model.uid, original_nested_pipeline_model.uid) + self.assertEqual( + len(loaded_nested_pipeline_model.stages), len(original_nested_pipeline_model.stages) + ) + for loadedStage, originalStage in zip( + loaded_nested_pipeline_model.stages, original_nested_pipeline_model.stages + ): + self.assertEqual(loadedStage.uid, originalStage.uid) + + def test_save_load_pipeline_estimator(self): + self._run_test_save_load_pipeline_estimator(LogisticRegression) + self._run_test_save_load_pipeline_estimator(DummyLogisticRegression) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_cv_io_pipeline import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tuning.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tuning.py new file mode 100644 index 0000000..ac0f846 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tuning.py @@ -0,0 +1,551 @@ +# +# 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 tempfile +import unittest + +import numpy as np +from pyspark.ml import Estimator, Model +from pyspark.ml.classification import LogisticRegression +from pyspark.ml.evaluation import ( + BinaryClassificationEvaluator, + RegressionEvaluator, +) +from pyspark.ml.linalg import Vectors +from pyspark.ml.param import Param, Params +from pyspark.ml.tuning import ( + CrossValidator, + CrossValidatorModel, + ParamGridBuilder, + TrainValidationSplit, + TrainValidationSplitModel, +) +from pyspark.sql.functions import rand +from pyspark.testing.mlutils import SparkSessionTestCase + + +class HasInducedError(Params): + def __init__(self): + super(HasInducedError, self).__init__() + self.inducedError = Param( + self, "inducedError", "Uniformly-distributed error added to feature" + ) + + def getInducedError(self): + return self.getOrDefault(self.inducedError) + + +class InducedErrorModel(Model, HasInducedError): + def __init__(self): + super(InducedErrorModel, self).__init__() + + def _transform(self, dataset): + return dataset.withColumn( + "prediction", dataset.feature + (rand(0) * self.getInducedError()) + ) + + +class InducedErrorEstimator(Estimator, HasInducedError): + def __init__(self, inducedError=1.0): + super(InducedErrorEstimator, self).__init__() + self._set(inducedError=inducedError) + + def _fit(self, dataset): + model = InducedErrorModel() + self._copyValues(model) + return model + + +class ParamGridBuilderTests(SparkSessionTestCase): + def test_addGrid(self): + with self.assertRaises(TypeError): + (ParamGridBuilder().addGrid("must be an instance of Param", ["not", "string"]).build()) + + +class ValidatorTestUtilsMixin: + def assert_param_maps_equal(self, paramMaps1, paramMaps2): + self.assertEqual(len(paramMaps1), len(paramMaps2)) + for paramMap1, paramMap2 in zip(paramMaps1, paramMaps2): + self.assertEqual(set(paramMap1.keys()), set(paramMap2.keys())) + for param in paramMap1.keys(): + v1 = paramMap1[param] + v2 = paramMap2[param] + if isinstance(v1, Params): + self.assertEqual(v1.uid, v2.uid) + else: + self.assertEqual(v1, v2) + + +class CrossValidatorTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def test_gen_avg_and_std_metrics(self): + metrics_all = [ + [1.0, 3.0, 2.0, 4.0], + [3.0, 2.0, 2.0, 4.0], + [3.0, 2.5, 2.1, 8.0], + ] + avg_metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all) + assert np.allclose(avg_metrics, [2.33333333, 2.5, 2.03333333, 5.33333333]) + assert np.allclose(std_metrics, [0.94280904, 0.40824829, 0.04714045, 1.88561808]) + assert isinstance(avg_metrics, list) + assert isinstance(std_metrics, list) + + def test_copy(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="rmse") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator( + estimator=iee, + estimatorParamMaps=grid, + evaluator=evaluator, + collectSubModels=True, + numFolds=2, + ) + cvCopied = cv.copy() + for param in [ + lambda x: x.getEstimator().uid, + # SPARK-32092: CrossValidator.copy() needs to copy all existing params + lambda x: x.getNumFolds(), + lambda x: x.getFoldCol(), + lambda x: x.getCollectSubModels(), + lambda x: x.getParallelism(), + lambda x: x.getSeed(), + ]: + self.assertEqual(param(cv), param(cvCopied)) + + cvModel = cv.fit(dataset) + cvModelCopied = cvModel.copy() + for index in range(len(cvModel.avgMetrics)): + self.assertTrue( + abs(cvModel.avgMetrics[index] - cvModelCopied.avgMetrics[index]) < 0.0001 + ) + self.assertTrue(np.allclose(cvModel.stdMetrics, cvModelCopied.stdMetrics)) + # SPARK-32092: CrossValidatorModel.copy() needs to copy all existing params + for param in [lambda x: x.getNumFolds(), lambda x: x.getFoldCol(), lambda x: x.getSeed()]: + self.assertEqual(param(cvModel), param(cvModelCopied)) + + cvModel.avgMetrics[0] = "foo" + self.assertNotEqual( + cvModelCopied.avgMetrics[0], + "foo", + "Changing the original avgMetrics should not affect the copied model", + ) + cvModel.stdMetrics[0] = "foo" + self.assertNotEqual( + cvModelCopied.stdMetrics[0], + "foo", + "Changing the original stdMetrics should not affect the copied model", + ) + cvModel.subModels[0][0].getInducedError = lambda: "foo" + self.assertNotEqual( + cvModelCopied.subModels[0][0].getInducedError(), + "foo", + "Changing the original subModels should not affect the copied model", + ) + + def test_fit_minimize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="rmse") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + bestModel = cvModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(0.0, bestModelMetric, "Best model has RMSE of 0") + + def test_fit_maximize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="r2") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + cv = CrossValidator(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + cvModel = cv.fit(dataset) + bestModel = cvModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(1.0, bestModelMetric, "Best model has R-squared of 1") + + def test_param_grid_type_coercion(self): + lr = LogisticRegression(maxIter=10) + paramGrid = ParamGridBuilder().addGrid(lr.regParam, [0.5, 1]).build() + for param in paramGrid: + for v in param.values(): + assert type(v) == float + + def test_parallel_evaluation(self): + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [5, 6]).build() + evaluator = BinaryClassificationEvaluator() + + # test save/load of CrossValidator + cv = CrossValidator(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator) + cv.setParallelism(1) + cvSerialModel = cv.fit(dataset) + cv.setParallelism(2) + cvParallelModel = cv.fit(dataset) + self.assertEqual(cvSerialModel.avgMetrics, cvParallelModel.avgMetrics) + self.assertEqual(cvSerialModel.stdMetrics, cvParallelModel.stdMetrics) + + def test_expose_sub_models(self): + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = BinaryClassificationEvaluator() + + numFolds = 3 + cv = CrossValidator( + estimator=lr, + estimatorParamMaps=grid, + evaluator=evaluator, + numFolds=numFolds, + collectSubModels=True, + ) + + def checkSubModels(subModels): + self.assertEqual(len(subModels), numFolds) + for i in range(numFolds): + self.assertEqual(len(subModels[i]), len(grid)) + + cvModel = cv.fit(dataset) + checkSubModels(cvModel.subModels) + + # Test the default value for option "persistSubModel" to be "true" + testSubPath = temp_path + "/testCrossValidatorSubModels" + savingPathWithSubModels = testSubPath + "cvModel3" + cvModel.save(savingPathWithSubModels) + cvModel3 = CrossValidatorModel.load(savingPathWithSubModels) + checkSubModels(cvModel3.subModels) + cvModel4 = cvModel3.copy() + checkSubModels(cvModel4.subModels) + + savingPathWithoutSubModels = testSubPath + "cvModel2" + cvModel.write().option("persistSubModels", "false").save(savingPathWithoutSubModels) + cvModel2 = CrossValidatorModel.load(savingPathWithoutSubModels) + self.assertEqual(cvModel2.subModels, None) + + for i in range(numFolds): + for j in range(len(grid)): + self.assertEqual(cvModel.subModels[i][j].uid, cvModel3.subModels[i][j].uid) + + def test_user_specified_folds(self): + from pyspark.sql import functions as F + + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ).repartition(2, "features") + + dataset_with_folds = ( + dataset.repartition(1) + .withColumn("random", rand(100)) + .withColumn( + "fold", + F.when(F.col("random") < 0.33, 0).when(F.col("random") < 0.66, 1).otherwise(2), + ) + .repartition(2, "features") + ) + + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [20]).build() + evaluator = BinaryClassificationEvaluator() + + cv = CrossValidator(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, numFolds=3) + cv_with_user_folds = CrossValidator( + estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, numFolds=3, foldCol="fold" + ) + + self.assertEqual(cv.getEstimator().uid, cv_with_user_folds.getEstimator().uid) + + cvModel1 = cv.fit(dataset) + cvModel2 = cv_with_user_folds.fit(dataset_with_folds) + for index in range(len(cvModel1.avgMetrics)): + print(abs(cvModel1.avgMetrics[index] - cvModel2.avgMetrics[index])) + self.assertTrue(abs(cvModel1.avgMetrics[index] - cvModel2.avgMetrics[index]) < 0.1) + + # test save/load of CrossValidator + temp_path = tempfile.mkdtemp() + cvPath = temp_path + "/cv" + cv_with_user_folds.save(cvPath) + loadedCV = CrossValidator.load(cvPath) + self.assertEqual(loadedCV.getFoldCol(), cv_with_user_folds.getFoldCol()) + + def test_invalid_user_specified_folds(self): + dataset_with_folds = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0, 0), + (Vectors.dense([0.4]), 1.0, 1), + (Vectors.dense([0.5]), 0.0, 2), + (Vectors.dense([0.6]), 1.0, 0), + (Vectors.dense([1.0]), 1.0, 1), + ] + * 10, + ["features", "label", "fold"], + ) + + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [20]).build() + evaluator = BinaryClassificationEvaluator() + + cv = CrossValidator( + estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, numFolds=2, foldCol="fold" + ) + with self.assertRaisesRegex(Exception, "Fold number must be in range"): + cv.fit(dataset_with_folds) + + cv = CrossValidator( + estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, numFolds=4, foldCol="fold" + ) + with self.assertRaisesRegex(Exception, "The validation data at fold 3 is empty"): + cv.fit(dataset_with_folds) + + +class TrainValidationSplitTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def test_fit_minimize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="rmse") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + tvs = TrainValidationSplit(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + tvsModel = tvs.fit(dataset) + bestModel = tvsModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + validationMetrics = tvsModel.validationMetrics + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(0.0, bestModelMetric, "Best model has RMSE of 0") + self.assertEqual( + len(grid), + len(validationMetrics), + "validationMetrics has the same size of grid parameter", + ) + self.assertEqual(0.0, min(validationMetrics)) + + def test_fit_maximize_metric(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="r2") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + tvs = TrainValidationSplit(estimator=iee, estimatorParamMaps=grid, evaluator=evaluator) + tvsModel = tvs.fit(dataset) + bestModel = tvsModel.bestModel + bestModelMetric = evaluator.evaluate(bestModel.transform(dataset)) + validationMetrics = tvsModel.validationMetrics + + self.assertEqual( + 0.0, bestModel.getOrDefault("inducedError"), "Best model should have zero induced error" + ) + self.assertEqual(1.0, bestModelMetric, "Best model has R-squared of 1") + self.assertEqual( + len(grid), + len(validationMetrics), + "validationMetrics has the same size of grid parameter", + ) + self.assertEqual(1.0, max(validationMetrics)) + + def test_parallel_evaluation(self): + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [5, 6]).build() + evaluator = BinaryClassificationEvaluator() + tvs = TrainValidationSplit(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator) + tvs.setParallelism(1) + tvsSerialModel = tvs.fit(dataset) + tvs.setParallelism(2) + tvsParallelModel = tvs.fit(dataset) + self.assertEqual(tvsSerialModel.validationMetrics, tvsParallelModel.validationMetrics) + + def test_expose_sub_models(self): + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + lr = LogisticRegression() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = BinaryClassificationEvaluator() + tvs = TrainValidationSplit( + estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, collectSubModels=True + ) + tvsModel = tvs.fit(dataset) + self.assertEqual(len(tvsModel.subModels), len(grid)) + + # Test the default value for option "persistSubModel" to be "true" + testSubPath = temp_path + "/testTrainValidationSplitSubModels" + savingPathWithSubModels = testSubPath + "cvModel3" + tvsModel.save(savingPathWithSubModels) + tvsModel3 = TrainValidationSplitModel.load(savingPathWithSubModels) + self.assertEqual(len(tvsModel3.subModels), len(grid)) + tvsModel4 = tvsModel3.copy() + self.assertEqual(len(tvsModel4.subModels), len(grid)) + + savingPathWithoutSubModels = testSubPath + "cvModel2" + tvsModel.write().option("persistSubModels", "false").save(savingPathWithoutSubModels) + tvsModel2 = TrainValidationSplitModel.load(savingPathWithoutSubModels) + self.assertEqual(tvsModel2.subModels, None) + + for i in range(len(grid)): + self.assertEqual(tvsModel.subModels[i].uid, tvsModel3.subModels[i].uid) + + def test_copy(self): + dataset = self.spark.createDataFrame( + [(10, 10.0), (50, 50.0), (100, 100.0), (500, 500.0)] * 10, ["feature", "label"] + ) + + iee = InducedErrorEstimator() + evaluator = RegressionEvaluator(metricName="r2") + + grid = ParamGridBuilder().addGrid(iee.inducedError, [100.0, 0.0, 10000.0]).build() + tvs = TrainValidationSplit( + estimator=iee, estimatorParamMaps=grid, evaluator=evaluator, collectSubModels=True + ) + tvsModel = tvs.fit(dataset) + tvsCopied = tvs.copy() + tvsModelCopied = tvsModel.copy() + + for param in [ + lambda x: x.getCollectSubModels(), + lambda x: x.getParallelism(), + lambda x: x.getSeed(), + lambda x: x.getTrainRatio(), + ]: + self.assertEqual(param(tvs), param(tvsCopied)) + + for param in [ + lambda x: x.getSeed(), + lambda x: x.getTrainRatio(), + ]: + self.assertEqual(param(tvsModel), param(tvsModelCopied)) + + self.assertEqual( + tvs.getEstimator().uid, + tvsCopied.getEstimator().uid, + "Copied TrainValidationSplit has the same uid of Estimator", + ) + + self.assertEqual(tvsModel.bestModel.uid, tvsModelCopied.bestModel.uid) + self.assertEqual( + len(tvsModel.validationMetrics), + len(tvsModelCopied.validationMetrics), + "Copied validationMetrics has the same size of the original", + ) + for index in range(len(tvsModel.validationMetrics)): + self.assertEqual( + tvsModel.validationMetrics[index], tvsModelCopied.validationMetrics[index] + ) + + tvsModel.validationMetrics[0] = "foo" + self.assertNotEqual( + tvsModelCopied.validationMetrics[0], + "foo", + "Changing the original validationMetrics should not affect the copied model", + ) + tvsModel.subModels[0].getInducedError = lambda: "foo" + self.assertNotEqual( + tvsModelCopied.subModels[0].getInducedError(), + "foo", + "Changing the original subModels should not affect the copied model", + ) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_tuning import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py new file mode 100644 index 0000000..991b837 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py @@ -0,0 +1,137 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel +from pyspark.ml.evaluation import BinaryClassificationEvaluator +from pyspark.ml.linalg import Vectors +from pyspark.ml.tuning import ( + ParamGridBuilder, + TrainValidationSplit, + TrainValidationSplitModel, +) +from pyspark.testing.mlutils import ( + DummyEvaluator, + DummyLogisticRegression, + DummyLogisticRegressionModel, + SparkSessionTestCase, +) +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class TrainValidationSplitIOBasicTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_trained_model(self, LogisticRegressionCls, LogisticRegressionModelCls): + # This tests saving and loading the trained model only. + # Save/load for TrainValidationSplit will be added later: SPARK-13786 + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + lr = LogisticRegressionCls() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = BinaryClassificationEvaluator() + tvs = TrainValidationSplit( + estimator=lr, + estimatorParamMaps=grid, + evaluator=evaluator, + collectSubModels=True, + seed=42, + ) + tvsModel = tvs.fit(dataset) + lrModel = tvsModel.bestModel + + lrModelPath = temp_path + "/lrModel" + lrModel.save(lrModelPath) + loadedLrModel = LogisticRegressionModelCls.load(lrModelPath) + self.assertEqual(loadedLrModel.uid, lrModel.uid) + self.assertEqual(loadedLrModel.intercept, lrModel.intercept) + + tvsModelPath = temp_path + "/tvsModel" + tvsModel.save(tvsModelPath) + loadedTvsModel = TrainValidationSplitModel.load(tvsModelPath) + for param in [ + lambda x: x.getSeed(), + lambda x: x.getTrainRatio(), + ]: + self.assertEqual(param(tvsModel), param(loadedTvsModel)) + + self.assertTrue(all(loadedTvsModel.isSet(param) for param in loadedTvsModel.params)) + + def test_save_load_trained_model(self): + self._run_test_save_load_trained_model(LogisticRegression, LogisticRegressionModel) + self._run_test_save_load_trained_model( + DummyLogisticRegression, DummyLogisticRegressionModel + ) + + def _run_test_save_load_simple_estimator(self, LogisticRegressionCls, evaluatorCls): + # This tests saving and loading the trained model only. + # Save/load for TrainValidationSplit will be added later: SPARK-13786 + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + lr = LogisticRegressionCls() + grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + evaluator = evaluatorCls() + tvs = TrainValidationSplit(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator) + tvsModel = tvs.fit(dataset) + + tvsPath = temp_path + "/tvs" + tvs.save(tvsPath) + loadedTvs = TrainValidationSplit.load(tvsPath) + self.assertEqual(loadedTvs.getEstimator().uid, tvs.getEstimator().uid) + self.assertEqual(loadedTvs.getEvaluator().uid, tvs.getEvaluator().uid) + self.assert_param_maps_equal(loadedTvs.getEstimatorParamMaps(), tvs.getEstimatorParamMaps()) + + tvsModelPath = temp_path + "/tvsModel" + tvsModel.save(tvsModelPath) + loadedModel = TrainValidationSplitModel.load(tvsModelPath) + self.assertEqual(loadedModel.bestModel.uid, tvsModel.bestModel.uid) + + def test_save_load_simple_estimator(self): + self._run_test_save_load_simple_estimator(LogisticRegression, BinaryClassificationEvaluator) + self._run_test_save_load_simple_estimator(DummyLogisticRegression, DummyEvaluator) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_tvs_io_basic import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py new file mode 100644 index 0000000..2abdf93 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py @@ -0,0 +1,96 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.classification import LogisticRegression, OneVsRest +from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.linalg import Vectors +from pyspark.ml.tuning import ( + ParamGridBuilder, + TrainValidationSplit, + TrainValidationSplitModel, +) +from pyspark.testing.mlutils import ( + DummyLogisticRegression, + SparkSessionTestCase, +) +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class TrainValidationSplitIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_nested_estimator(self, LogisticRegressionCls): + # This tests saving and loading the trained model only. + # Save/load for TrainValidationSplit will be added later: SPARK-13786 + temp_path = tempfile.mkdtemp() + dataset = self.spark.createDataFrame( + [ + (Vectors.dense([0.0]), 0.0), + (Vectors.dense([0.4]), 1.0), + (Vectors.dense([0.5]), 0.0), + (Vectors.dense([0.6]), 1.0), + (Vectors.dense([1.0]), 1.0), + ] + * 10, + ["features", "label"], + ) + ova = OneVsRest(classifier=LogisticRegressionCls()) + lr1 = LogisticRegressionCls().setMaxIter(100) + lr2 = LogisticRegressionCls().setMaxIter(150) + grid = ParamGridBuilder().addGrid(ova.classifier, [lr1, lr2]).build() + evaluator = MulticlassClassificationEvaluator() + + tvs = TrainValidationSplit(estimator=ova, estimatorParamMaps=grid, evaluator=evaluator) + tvsModel = tvs.fit(dataset) + tvsPath = temp_path + "/tvs" + tvs.save(tvsPath) + loadedTvs = TrainValidationSplit.load(tvsPath) + self.assert_param_maps_equal(loadedTvs.getEstimatorParamMaps(), grid) + self.assertEqual(loadedTvs.getEstimator().uid, tvs.getEstimator().uid) + self.assertEqual(loadedTvs.getEvaluator().uid, tvs.getEvaluator().uid) + + originalParamMap = tvs.getEstimatorParamMaps() + loadedParamMap = loadedTvs.getEstimatorParamMaps() + for i, param in enumerate(loadedParamMap): + for p in param: + if p.name == "classifier": + self.assertEqual(param[p].uid, originalParamMap[i][p].uid) + else: + self.assertEqual(param[p], originalParamMap[i][p]) + + tvsModelPath = temp_path + "/tvsModel" + tvsModel.save(tvsModelPath) + loadedModel = TrainValidationSplitModel.load(tvsModelPath) + self.assert_param_maps_equal(loadedModel.getEstimatorParamMaps(), grid) + self.assertEqual(loadedModel.bestModel.uid, tvsModel.bestModel.uid) + + def test_save_load_nested_estimator(self): + self._run_test_save_load_nested_estimator(LogisticRegression) + self._run_test_save_load_nested_estimator(DummyLogisticRegression) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_tvs_io_nested import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py new file mode 100644 index 0000000..3ad12ca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py @@ -0,0 +1,142 @@ +# +# 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 tempfile +import unittest + +from pyspark.ml.feature import HashingTF, Tokenizer +from pyspark.ml import Pipeline +from pyspark.ml.classification import LogisticRegression, OneVsRest +from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.tuning import ( + ParamGridBuilder, + TrainValidationSplit, + TrainValidationSplitModel, +) +from pyspark.testing.mlutils import ( + DummyLogisticRegression, + SparkSessionTestCase, +) + +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin + + +class TrainValidationSplitIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): + def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): + temp_path = tempfile.mkdtemp() + training = self.spark.createDataFrame( + [ + (0, "a b c d e spark", 1.0), + (1, "b d", 0.0), + (2, "spark f g h", 1.0), + (3, "hadoop mapreduce", 0.0), + (4, "b spark who", 1.0), + (5, "g d a y", 0.0), + (6, "spark fly", 1.0), + (7, "was mapreduce", 0.0), + ], + ["id", "text", "label"], + ) + + # Configure an ML pipeline, which consists of tree stages: tokenizer, hashingTF, and lr. + tokenizer = Tokenizer(inputCol="text", outputCol="words") + hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features") + + ova = OneVsRest(classifier=LogisticRegressionCls()) + lr1 = LogisticRegressionCls().setMaxIter(5) + lr2 = LogisticRegressionCls().setMaxIter(10) + + pipeline = Pipeline(stages=[tokenizer, hashingTF, ova]) + + paramGrid = ( + ParamGridBuilder() + .addGrid(hashingTF.numFeatures, [10, 100]) + .addGrid(ova.classifier, [lr1, lr2]) + .build() + ) + + tvs = TrainValidationSplit( + estimator=pipeline, + estimatorParamMaps=paramGrid, + evaluator=MulticlassClassificationEvaluator(), + ) + tvsPath = temp_path + "/tvs" + tvs.save(tvsPath) + loadedTvs = TrainValidationSplit.load(tvsPath) + self.assert_param_maps_equal(loadedTvs.getEstimatorParamMaps(), paramGrid) + self.assertEqual(loadedTvs.getEstimator().uid, tvs.getEstimator().uid) + + # Run train validation split, and choose the best set of parameters. + tvsModel = tvs.fit(training) + + # test save/load of CrossValidatorModel + tvsModelPath = temp_path + "/tvsModel" + tvsModel.save(tvsModelPath) + loadedModel = TrainValidationSplitModel.load(tvsModelPath) + self.assertEqual(loadedModel.bestModel.uid, tvsModel.bestModel.uid) + self.assertEqual(len(loadedModel.bestModel.stages), len(tvsModel.bestModel.stages)) + for loadedStage, originalStage in zip( + loadedModel.bestModel.stages, tvsModel.bestModel.stages + ): + self.assertEqual(loadedStage.uid, originalStage.uid) + + # Test nested pipeline + nested_pipeline = Pipeline(stages=[tokenizer, Pipeline(stages=[hashingTF, ova])]) + tvs2 = TrainValidationSplit( + estimator=nested_pipeline, + estimatorParamMaps=paramGrid, + evaluator=MulticlassClassificationEvaluator(), + ) + tvs2Path = temp_path + "/tvs2" + tvs2.save(tvs2Path) + loadedTvs2 = TrainValidationSplit.load(tvs2Path) + self.assert_param_maps_equal(loadedTvs2.getEstimatorParamMaps(), paramGrid) + self.assertEqual(loadedTvs2.getEstimator().uid, tvs2.getEstimator().uid) + + # Run train validation split, and choose the best set of parameters. + tvsModel2 = tvs2.fit(training) + # test save/load of CrossValidatorModel + tvsModelPath2 = temp_path + "/tvsModel2" + tvsModel2.save(tvsModelPath2) + loadedModel2 = TrainValidationSplitModel.load(tvsModelPath2) + self.assertEqual(loadedModel2.bestModel.uid, tvsModel2.bestModel.uid) + loaded_nested_pipeline_model = loadedModel2.bestModel.stages[1] + original_nested_pipeline_model = tvsModel2.bestModel.stages[1] + self.assertEqual(loaded_nested_pipeline_model.uid, original_nested_pipeline_model.uid) + self.assertEqual( + len(loaded_nested_pipeline_model.stages), len(original_nested_pipeline_model.stages) + ) + for loadedStage, originalStage in zip( + loaded_nested_pipeline_model.stages, original_nested_pipeline_model.stages + ): + self.assertEqual(loadedStage.uid, originalStage.uid) + + def test_save_load_pipeline_estimator(self): + self._run_test_save_load_pipeline_estimator(LogisticRegression) + self._run_test_save_load_pipeline_estimator(DummyLogisticRegression) + + +if __name__ == "__main__": + from pyspark.ml.tests.tuning.test_tvs_io_pipeline import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_classification.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_classification.yml new file mode 100644 index 0000000..45e5f1c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_classification.yml @@ -0,0 +1,38 @@ +# +# 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. +# + +- case: oneVsRest + main: | + from pyspark.ml.classification import ( + OneVsRest, OneVsRestModel, LogisticRegression, LogisticRegressionModel + ) + + # Should support + OneVsRest(classifier=LogisticRegression()) + OneVsRest(classifier=LogisticRegressionModel.load("/foo")) # E: Argument "classifier" to "OneVsRest" has incompatible type "LogisticRegressionModel"; expected "Optional[Classifier[]]" [arg-type] + OneVsRest(classifier="foo") # E: Argument "classifier" to "OneVsRest" has incompatible type "str"; expected "Optional[Classifier[]]" [arg-type] + + +- case: fitFMClassifier + main: | + from pyspark.sql import SparkSession + from pyspark.ml.classification import FMClassifier, FMClassificationModel + + spark = SparkSession.builder.getOrCreate() + fm_model: FMClassificationModel = FMClassifier().fit(spark.read.parquet("/foo")) + fm_model.linear.toArray() + fm_model.factors.numRows diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_clustering.yaml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_clustering.yaml new file mode 100644 index 0000000..bc50cf5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_clustering.yaml @@ -0,0 +1,33 @@ +# +# 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. +# + +- case: InheritedLDAMethods + main: | + from pyspark.ml.clustering import LDAModel, LocalLDAModel, DistributedLDAModel + + distributed_model = DistributedLDAModel.load("foo") + reveal_type(distributed_model) + reveal_type(distributed_model.setFeaturesCol("foo")) + + local_model = distributed_model.toLocal() + reveal_type(local_model) + reveal_type(local_model.setFeaturesCol("foo")) + out: | + main:4: note: Revealed type is "pyspark.ml.clustering.DistributedLDAModel" + main:5: note: Revealed type is "pyspark.ml.clustering.DistributedLDAModel" + main:8: note: Revealed type is "pyspark.ml.clustering.LocalLDAModel" + main:9: note: Revealed type is "pyspark.ml.clustering.LocalLDAModel" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_evaluation.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_evaluation.yml new file mode 100644 index 0000000..51e076c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_evaluation.yml @@ -0,0 +1,28 @@ +# +# 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. +# + +- case: BinaryClassificationEvaluator + main: | + from pyspark.ml.evaluation import BinaryClassificationEvaluator + + BinaryClassificationEvaluator().setMetricName("areaUnderROC") + BinaryClassificationEvaluator(metricName="areaUnderPR") + + BinaryClassificationEvaluator().setMetricName("foo") # E: Argument 1 to "setMetricName" of "BinaryClassificationEvaluator" has incompatible type "Literal['foo']"; expected "Literal['areaUnderROC', 'areaUnderPR']" [arg-type] + BinaryClassificationEvaluator(metricName="bar") # E: Argument "metricName" to "BinaryClassificationEvaluator" has incompatible type "Literal['bar']"; expected "Literal['areaUnderROC', 'areaUnderPR']" [arg-type] + + reveal_type(BinaryClassificationEvaluator.load("foo")) # N: Revealed type is "pyspark.ml.evaluation.BinaryClassificationEvaluator" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_feature.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_feature.yml new file mode 100644 index 0000000..0d1034a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_feature.yml @@ -0,0 +1,55 @@ +# +# 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. +# + + +- case: featureMethodChaining + main: | + from pyspark.ml.feature import NGram + + reveal_type(NGram().setInputCol("foo").setOutputCol("bar")) + + out: | + main:3: note: Revealed type is "pyspark.ml.feature.NGram" + + +- case: stringIndexerOverloads + main: | + from pyspark.ml.feature import StringIndexer + + # No arguments is OK + StringIndexer() + + StringIndexer(inputCol="foo") + StringIndexer(outputCol="bar") + StringIndexer(inputCol="foo", outputCol="bar") + + StringIndexer(inputCols=["foo"]) + StringIndexer(outputCols=["bar"]) + StringIndexer(inputCols=["foo"], outputCols=["bar"]) + + StringIndexer(inputCol="foo", outputCols=["bar"]) + StringIndexer(inputCols=["foo"], outputCol="bar") + + out: | + main:14: error: No overload variant of "StringIndexer" matches argument types "str", "List[str]" [call-overload] + main:14: note: Possible overload variants: + main:14: note: def StringIndexer(self, *, inputCol: Optional[str] = ..., outputCol: Optional[str] = ..., handleInvalid: str = ..., stringOrderType: str = ...) -> StringIndexer + main:14: note: def StringIndexer(self, *, inputCols: Optional[List[str]] = ..., outputCols: Optional[List[str]] = ..., handleInvalid: str = ..., stringOrderType: str = ...) -> StringIndexer + main:15: error: No overload variant of "StringIndexer" matches argument types "List[str]", "str" [call-overload] + main:15: note: Possible overload variants: + main:15: note: def StringIndexer(self, *, inputCol: Optional[str] = ..., outputCol: Optional[str] = ..., handleInvalid: str = ..., stringOrderType: str = ...) -> StringIndexer + main:15: note: def StringIndexer(self, *, inputCols: Optional[List[str]] = ..., outputCols: Optional[List[str]] = ..., handleInvalid: str = ..., stringOrderType: str = ...) -> StringIndexer diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_param.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_param.yml new file mode 100644 index 0000000..0b423f0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_param.yml @@ -0,0 +1,30 @@ +# +# 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. +# + +- case: paramGenric + main: | + from pyspark.ml.param import Param, Params, TypeConverters + + class Foo(Params): + foo = Param(Params(), "foo", "foo", TypeConverters.toInt) + def getFoo(self) -> int: + return self.getOrDefault(self.foo) + + class Bar(Params): + bar = Param(Params(), "bar", "bar", TypeConverters.toInt) + def getFoo(self) -> str: + return self.getOrDefault(self.bar) # E: Incompatible return value type (got "int", expected "str") [return-value] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_readable.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_readable.yml new file mode 100644 index 0000000..772133a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_readable.yml @@ -0,0 +1,28 @@ +# +# 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. +# + +- case: readLinearSVCModel + main: | + from pyspark.ml.classification import LinearSVCModel + + model1 = LinearSVCModel.load("dummy") + model1.coefficients.toArray() + model1.foo() # E: "LinearSVCModel" has no attribute "foo" [attr-defined] + + model2 = LinearSVCModel.read().load("dummy") + model2.coefficients.toArray() + model2.foo() # E: "LinearSVCModel" has no attribute "foo" [attr-defined] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_regression.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_regression.yml new file mode 100644 index 0000000..4a54a56 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tests/typing/test_regression.yml @@ -0,0 +1,53 @@ +# +# 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. +# + +- case: linearRegressionMethodChaining + main: | + from pyspark.ml.regression import LinearRegression, LinearRegressionModel + + lr = LinearRegression() + reveal_type(lr.setFeaturesCol("foo").setLabelCol("bar")) + + lrm = LinearRegressionModel.load("/foo") + reveal_type(lrm.setPredictionCol("baz")) + + out: | + main:4: note: Revealed type is "pyspark.ml.regression.LinearRegression" + main:7: note: Revealed type is "pyspark.ml.regression.LinearRegressionModel" + + +- case: loadFMRegressor + main: | + from pyspark.ml.regression import FMRegressor, FMRegressionModel + + fm = FMRegressor.load("/foo") + fm.setMiniBatchFraction(0.1) + + fm_model = FMRegressionModel.load("/bar") + fm_model.factors.numCols + + fm_model.foo() # E: "FMRegressionModel" has no attribute "foo" [attr-defined] + + +- case: loadLinearRegressor + main: | + from pyspark.ml.regression import LinearRegressionModel + + lr_model = LinearRegressionModel.load("/foo") + lr_model.getLabelCol().upper() + + lr_model.foo # E: "LinearRegressionModel" has no attribute "foo" [attr-defined] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/__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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/data.py new file mode 100644 index 0000000..0a5597f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/data.py @@ -0,0 +1,100 @@ +# +# 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 torch +import numpy as np +from typing import Any, Callable, Iterator +from pyspark.sql.types import StructType + + +class _SparkPartitionTorchDataset(torch.utils.data.IterableDataset): + def __init__(self, arrow_file_path: str, schema: "StructType", num_samples: int): + self.arrow_file_path = arrow_file_path + self.num_samples = num_samples + self.field_types = [field.dataType.simpleString() for field in schema] + self.field_converters = [ + _SparkPartitionTorchDataset._get_field_converter(field_type) + for field_type in self.field_types + ] + + @staticmethod + def _get_field_converter(field_type: str) -> Callable[[Any], Any]: + if field_type == "vector": + + def converter(value: Any) -> Any: + if value["type"] == 1: + # dense vector + return value["values"] + if value["type"] == 0: + # sparse vector + size = int(value["size"]) + sparse_array = np.zeros(size, dtype=np.float64) + sparse_array[value["indices"]] = value["values"] + return sparse_array + + elif field_type in [ + "float", + "double", + "int", + "bigint", + "smallint", + "array", + "array", + "array", + "array", + "array", + ]: + + def converter(value: Any) -> Any: + return value + + else: + raise ValueError( + "SparkPartitionTorchDataset does not support loading data from field of " + f"type {field_type}." + ) + return converter + + def __iter__(self) -> Iterator[Any]: + from pyspark.sql.pandas.serializers import ArrowStreamSerializer + + serializer = ArrowStreamSerializer() + + worker_info = torch.utils.data.get_worker_info() + if worker_info is not None and worker_info.num_workers > 1: + raise RuntimeError( + "SparkPartitionTorchDataset does not support multiple worker processes." + ) + + count = 0 + + while count < self.num_samples: + with open(self.arrow_file_path, "rb") as f: + batch_iter = serializer.load_stream(f) + for batch in batch_iter: + # TODO: we can optimize this further by directly extracting + # field data from arrow batch without converting it to + # pandas DataFrame. + batch_pdf = batch.to_pandas() + for row in batch_pdf.itertuples(index=False): + yield [ + field_converter(value) + for value, field_converter in zip(row, self.field_converters) + ] + count += 1 + if count == self.num_samples: + return diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/distributor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/distributor.py new file mode 100644 index 0000000..d0979f5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/distributor.py @@ -0,0 +1,1133 @@ +# +# 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 json +from contextlib import contextmanager +import collections +import logging +import math +import os +import random +import re +import shutil +import subprocess +import sys +import tempfile +import textwrap +import time +from typing import ( + Union, + Callable, + List, + Dict, + Optional, + Any, + Tuple, + Generator, + Iterator, +) + +from pyspark import cloudpickle +from pyspark.resource.information import ResourceInformation +from pyspark.sql import DataFrame, SparkSession +from pyspark.taskcontext import BarrierTaskContext +from pyspark.ml.torch.log_communication import ( # type: ignore + LogStreamingClient, + LogStreamingServer, +) + + +def _get_resources(session: SparkSession) -> Dict[str, ResourceInformation]: + resources: Dict[str, ResourceInformation] = {} + try: + resources = session.sparkContext.resources + except Exception: + resources = session._client._resources() # type: ignore[attr-defined] + return resources + + +def _get_conf(spark: SparkSession, key: str, default_value: str) -> str: + """Get the conf "key" from the given spark session, + or return the default value if the conf is not set. + + Parameters + ---------- + spark : :class:`SparkSession` + The :class:`SparkSession` for the distributor. + key : str + string for conf name + default_value : str + default value for the conf value for the given key + + Returns + ------- + str + Returns the string value that corresponds to the conf + """ + value = spark.conf.get(key, default_value) + assert value is not None + return value + + +# TODO(SPARK-41589): will move the functions and tests to an external file +# once we are in agreement about which functions should be in utils.py +def _get_conf_boolean(spark: SparkSession, key: str, default_value: str) -> bool: + value = _get_conf(spark=spark, key=key, default_value=default_value) + value = value.lower() + assert value in ["true", "false"] + return value == "true" + + +def _get_logger(name: str) -> logging.Logger: + """ + Gets a logger by name, or creates and configures it for the first time. + """ + logger = logging.getLogger(name) + logger.setLevel(logging.INFO) + # If the logger is configured, skip the configure + if not logger.handlers and not logging.getLogger().handlers: + handler = logging.StreamHandler(sys.stderr) + logger.addHandler(handler) + return logger + + +def _get_gpus_owned(context: Union[SparkSession, BarrierTaskContext]) -> List[str]: + """Gets the number of GPUs that Spark scheduled to the calling task. + + Parameters + ---------- + context : :class:`SparkSession` or :class:`BarrierTaskContext` + The :class:`SparkSession` or :class:`BarrierTaskContext` that has GPUs available. + + Returns + ------- + list + The correct mapping of addresses to workers. + + Raises + ------ + ValueError + Raised if the input addresses were not found. + """ + CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" + pattern = re.compile("^[1-9][0-9]*|0$") + if isinstance(context, BarrierTaskContext): + addresses = context.resources()["gpu"].addresses + else: + addresses = _get_resources(context)["gpu"].addresses + + if any(not pattern.match(address) for address in addresses): + raise ValueError( + f"Found GPU addresses {addresses} which " + "are not all in the correct format " + "for CUDA_VISIBLE_DEVICES, which requires " + "integers with no zero padding." + ) + if CUDA_VISIBLE_DEVICES in os.environ: + gpu_indices = list(map(int, addresses)) + gpu_list = os.environ[CUDA_VISIBLE_DEVICES].split(",") + gpu_owned = [gpu_list[i] for i in gpu_indices] + return gpu_owned + return addresses + + +SPARK_PARTITION_ARROW_DATA_FILE = "SPARK_PARTITION_ARROW_DATA_FILE" +SPARK_DATAFRAME_SCHEMA_FILE = "SPARK_DATAFRAME_SCHEMA_FILE" + + +class Distributor: + """ + The parent class for TorchDistributor. This class shouldn't be instantiated directly. + """ + + def __init__( + self, + num_processes: int = 1, + local_mode: bool = True, + use_gpu: bool = True, + ssl_conf: Optional[str] = None, + ): + from pyspark.sql.utils import is_remote + + self.is_remote = is_remote() + self.spark = SparkSession.active() + + # indicate whether the server side is local mode + self.is_spark_local_master = False + # Refer to 'org.apache.spark.util.Utils#isLocalMaster' + master = _get_conf(self.spark, "spark.master", "") + if master == "local" or master.startswith("local["): + self.is_spark_local_master = True + + self.logger = _get_logger(self.__class__.__name__) + self.num_processes = num_processes + self.local_mode = local_mode + self.use_gpu = use_gpu + self.num_tasks = self._get_num_tasks() + self.ssl_conf = ssl_conf + + def _create_input_params(self) -> Dict[str, Any]: + input_params = self.__dict__.copy() + for unneeded_param in [ + "spark", + "ssl_conf", + "logger", + "is_remote", + "is_spark_local_master", + ]: + del input_params[unneeded_param] + return input_params + + def _get_num_tasks(self) -> int: + """ + Returns the number of Spark tasks to use for distributed training + + Returns + ------- + int + The number of Spark tasks to use for distributed training + + Raises + ------ + RuntimeError + Raised when the SparkConf was misconfigured. + """ + if self.use_gpu: + if not self.local_mode: + key = "spark.task.resource.gpu.amount" + task_gpu_amount = int(_get_conf(self.spark, key, "0")) + if task_gpu_amount < 1: + raise RuntimeError(f"'{key}' was unset, so gpu usage is unavailable.") + # TODO(SPARK-41916): Address situation when spark.task.resource.gpu.amount > 1 + return math.ceil(self.num_processes / task_gpu_amount) + else: + key = "spark.driver.resource.gpu.amount" + if "gpu" not in _get_resources(self.spark): + raise RuntimeError("GPUs were unable to be found on the driver.") + num_available_gpus = int(_get_conf(self.spark, key, "0")) + if num_available_gpus == 0: + raise RuntimeError("GPU resources were not configured properly on the driver.") + if self.num_processes > num_available_gpus: + self.logger.warning( + "'num_processes' cannot be set to a value greater than the number of " + f"available GPUs on the driver, which is {num_available_gpus}. " + "'num_processes' was reset to be equal to the number of available GPUs.", + ) + self.num_processes = num_available_gpus + return self.num_processes + + def _validate_input_params(self) -> None: + if self.num_processes <= 0: + raise ValueError("num_proccesses has to be a positive integer") + + def _check_encryption(self) -> None: + """Checks to see if the user requires encrpytion of data. + If required, throw an exception since we don't support that. + + Raises + ------ + RuntimeError + Thrown when the user requires ssl encryption or when the user initializes + the Distributor parent class. + """ + if not hasattr(self, "ssl_conf"): + raise RuntimeError( + "Distributor doesn't have this functionality. Use TorchDistributor instead." + ) + is_ssl_enabled = _get_conf_boolean(self.spark, "spark.ssl.enabled", "false") + ignore_ssl = _get_conf_boolean(self.spark, self.ssl_conf, "false") # type: ignore + if is_ssl_enabled: + name = self.__class__.__name__ + if ignore_ssl: + self.logger.warning( + textwrap.dedent( + f""" + This cluster has TLS encryption enabled; + however, {name} does not + support data encryption in transit. + The Spark configuration + '{self.ssl_conf}' has been set to + 'true' to override this + configuration and use {name} anyway. Please + note this will cause model + parameters and possibly training data to + be sent between nodes unencrypted. + """, + ) + ) + return + raise RuntimeError( + textwrap.dedent( + f""" + This cluster has TLS encryption enabled; + however, {name} does not support + data encryption in transit. To override + this configuration and use {name} + anyway, you may set '{self.ssl_conf}' + to 'true' in the Spark configuration. Please note this + will cause model parameters and possibly training + data to be sent between nodes unencrypted. + """ + ) + ) + + +class TorchDistributor(Distributor): + """ + A class to support distributed training on PyTorch and PyTorch Lightning using PySpark. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + num_processes : int, optional + An integer that determines how many different concurrent + tasks are allowed. We expect spark.task.gpus = 1 for GPU-enabled training. Default + should be 1; we don't want to invoke multiple cores/gpus without explicit mention. + local_mode : bool, optional + A boolean that determines whether we are using the driver + node for training. Default should be false; we don't want to invoke executors without + explicit mention. + use_gpu : bool, optional + A boolean that indicates whether or not we are doing training + on the GPU. Note that there are differences in how GPU-enabled code looks like and + how CPU-specific code looks like. + + Examples + -------- + Run PyTorch Training locally on GPU (using a PyTorch native function) + + >>> def train(learning_rate): + ... import torch.distributed + ... torch.distributed.init_process_group(backend="nccl") + ... # ... + ... torch.destroy_process_group() + ... return model # or anything else + ... + >>> distributor = TorchDistributor( + ... num_processes=2, + ... local_mode=True, + ... use_gpu=True) + >>> model = distributor.run(train, 1e-3) + + Run PyTorch Training on GPU (using a file with PyTorch code) + + >>> distributor = TorchDistributor( + ... num_processes=2, + ... local_mode=False, + ... use_gpu=True) + >>> distributor.run("/path/to/train.py", "--learning-rate=1e-3") + + Run PyTorch Lightning Training on GPU + + >>> num_proc = 2 + >>> def train(): + ... from pytorch_lightning import Trainer + ... # ... + ... # required to set devices = 1 and num_nodes = num_processes for multi node + ... # required to set devices = num_processes and num_nodes = 1 for single node multi GPU + ... trainer = Trainer(accelerator="gpu", devices=1, num_nodes=num_proc, strategy="ddp") + ... trainer.fit() + ... # ... + ... return trainer + ... + >>> distributor = TorchDistributor( + ... num_processes=num_proc, + ... local_mode=True, + ... use_gpu=True) + >>> trainer = distributor.run(train) + """ + + _PICKLED_FUNC_FILE = "func.pickle" + _TRAIN_FILE = "train.py" + _PICKLED_OUTPUT_FILE = "output.pickle" + _TORCH_SSL_CONF = "pytorch.spark.distributor.ignoreSsl" + + def __init__( + self, + num_processes: int = 1, + local_mode: bool = True, + use_gpu: bool = True, + _ssl_conf: str = _TORCH_SSL_CONF, + ): + """Initializes the distributor. + + Parameters + ---------- + num_processes : int, optional + An integer that determines how many different concurrent + tasks are allowed. We expect spark.task.gpus = 1 for GPU-enabled training. Default + should be 1; we don't want to invoke multiple cores/gpus without explicit mention. + local_mode : bool, optional + A boolean that determines whether we are using the driver + node for training. Default should be false; we don't want to invoke executors without + explicit mention. + use_gpu : bool, optional + A boolean that indicates whether or not we are doing training + on the GPU. Note that there are differences in how GPU-enabled code looks like and + how CPU-specific code looks like. + + Raises + ------ + ValueError + If any of the parameters are incorrect. + RuntimeError + If an active SparkSession is unavailable. + """ + super().__init__(num_processes, local_mode, use_gpu, ssl_conf=_ssl_conf) + self._validate_input_params() + self.input_params = self._create_input_params() + + @staticmethod + def _get_torchrun_args(local_mode: bool, num_processes: int) -> Tuple[List[Any], int]: + """ + Given the mode and the number of processes, create the arguments to be given to for torch + + Parameters + --------- + local_mode: bool + Whether or not we are running training locally or in a distributed fashion + + num_processes: int + The number of processes that we are going to use + + Returns + ------ + Tuple[List[Any], int] + A tuple containing a list of arguments to pass as pytorch args, + as well as the number of processes per node + """ + if local_mode: + torchrun_args = ["--standalone", "--nnodes=1"] + processes_per_node = num_processes + return torchrun_args, processes_per_node + + master_addr = os.environ["MASTER_ADDR"] + master_port = os.environ["MASTER_PORT"] + node_rank = os.environ["RANK"] + torchrun_args = [ + f"--nnodes={num_processes}", + 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 + ] + processes_per_node = 1 + return torchrun_args, processes_per_node + + @staticmethod + def _create_torchrun_command( + input_params: Dict[str, Any], path_to_train_file: str, *args: Any + ) -> List[str]: + local_mode = input_params["local_mode"] + num_processes = input_params["num_processes"] + + torchrun_args, processes_per_node = TorchDistributor._get_torchrun_args( + local_mode=local_mode, num_processes=num_processes + ) + args_string = list(map(str, args)) # converting all args to strings + + return [ + sys.executable, + "-m", + "pyspark.ml.torch.torch_run_process_wrapper", + *torchrun_args, + f"--nproc_per_node={processes_per_node}", + path_to_train_file, + *args_string, + ] + + @staticmethod + def _execute_command( + cmd: List[str], + _prctl: bool = True, + redirect_to_stdout: bool = True, + log_streaming_client: Optional[LogStreamingClient] = None, + ) -> None: + _TAIL_LINES_TO_KEEP = 100 + + task = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + env=os.environ, + ) + task.stdin.close() # type: ignore + tail: collections.deque = collections.deque(maxlen=_TAIL_LINES_TO_KEEP) + try: + for line in task.stdout: # type: ignore + decoded = line.decode() + tail.append(decoded) + if redirect_to_stdout: + if ( + log_streaming_client + and not log_streaming_client.failed + and ( + log_streaming_client.sock.getsockname()[0] + == log_streaming_client.sock.getpeername()[0] + ) + ): + # If log_streaming_client and log_stream_server are in the same + # node (typical case is spark local mode), + # server side will redirect the log to STDOUT, + # to avoid STDOUT outputs duplication, skip redirecting + # logs to STDOUT in client side. + pass + else: + sys.stdout.write(decoded) + if log_streaming_client: + log_streaming_client.send(decoded.rstrip()) + task.wait() + finally: + if task.poll() is None: + try: + task.terminate() # SIGTERM + time.sleep(0.5) + if task.poll() is None: + task.kill() # SIGKILL + except OSError: + pass + if task.returncode != os.EX_OK: + if len(tail) == _TAIL_LINES_TO_KEEP: + last_n_msg = f"last {_TAIL_LINES_TO_KEEP} lines of the task output are" + else: + last_n_msg = "task output is" + task_output = "".join(tail) + raise RuntimeError( + f"Command {cmd} failed with return code {task.returncode}. " + f"The {last_n_msg} included below: {task_output}" + ) + + @staticmethod + def _get_output_from_framework_wrapper( + framework_wrapper: Optional[Callable], + input_params: Dict, + train_object: Union[Callable, str], + run_pytorch_file_fn: Optional[Callable], + *args: Any, + **kwargs: Any, + ) -> Optional[Any]: + """ + This function is meant to get the output from framework wrapper function by passing in the + correct arguments, depending on the type of train_object. + + Parameters + ---------- + framework_wrapper: Optional[Callable] + Function pointer that will be invoked. Can either be the function that runs distributed + training on files if train_object is a string. Otherwise, it will be the function that + runs distributed training for functions if the train_object is a Callable + input_params: Dict + A dictionary that maps parameter to arguments for the command to be created. + train_object: Union[Callable, str] + This input comes from the user. If the user inputs a string, then this means + it's a filepath. Otherwise, if the input is a function, then this means that + the user wants to run this function in a distributed manner. + run_pytorch_file_fn: Optional[Callable] + The function that will be used to run distributed training of a file; + mainly used for the distributed training using a function. + *args: Any + Extra arguments to be used by framework wrapper. + **kwargs: Any + Extra keyword args to be used. Not currently supported but kept for + future improvement. + + Returns + ------- + Optional[Any] + Returns the result of the framework_wrapper + """ + if not framework_wrapper: + raise RuntimeError("`framework_wrapper` is not set. ...") + # The object to train is a file path, so framework_wrapper is some + # run_training_on_pytorch_file function. + if type(train_object) is str: + return framework_wrapper(input_params, train_object, *args, **kwargs) + else: + # We are doing training with a function, will call run_training_on_pytorch_function + if not run_pytorch_file_fn: + run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file + return framework_wrapper( + input_params, train_object, run_pytorch_file_fn, *args, **kwargs + ) + + def _run_local_training( + self, + framework_wrapper_fn: Callable, + train_object: Union[Callable, str], + run_pytorch_file_fn: Optional[Callable], + *args: Any, + **kwargs: Any, + ) -> Optional[Any]: + CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" + cuda_state_was_set = CUDA_VISIBLE_DEVICES in os.environ + old_cuda_visible_devices = os.environ.get(CUDA_VISIBLE_DEVICES, "") + try: + # Only replace the GPUs with 'SparkContext.resources' in legacy mode. + # In connect mode, this replacement is skipped since only GPUs on the client side + # can be used. + if self.use_gpu and not self.is_remote: + gpus_owned = _get_gpus_owned(self.spark) + random.seed(hash(train_object)) + selected_gpus = [str(e) for e in random.sample(gpus_owned, self.num_processes)] + os.environ[CUDA_VISIBLE_DEVICES] = ",".join(selected_gpus) + + self.logger.info(f"Started local training with {self.num_processes} processes") + output = TorchDistributor._get_output_from_framework_wrapper( + framework_wrapper_fn, + self.input_params, + train_object, + run_pytorch_file_fn, + *args, + **kwargs, + ) + self.logger.info(f"Finished local training with {self.num_processes} processes") + + finally: + if cuda_state_was_set: + os.environ[CUDA_VISIBLE_DEVICES] = old_cuda_visible_devices + else: + if CUDA_VISIBLE_DEVICES in os.environ: + del os.environ[CUDA_VISIBLE_DEVICES] + + return output + + def _get_spark_task_function( + self, + framework_wrapper_fn: Optional[Callable], + train_object: Union[Callable, str], + run_pytorch_file_fn: Optional[Callable], + input_dataframe: Optional["DataFrame"], + *args: Any, + **kwargs: Any, + ) -> Callable: + """Creates a spark task function that is used inside `mapPartitions`. + + Parameters + ---------- + framework_wrapper_fn : Optional[Callable] + The function that determines whether we are running training + on a PyTorch file or a PyTorch function. + train_object : Union[Callable, str] + The actual train function/file. + + Returns + ------- + Callable + The wrapped function ready for use with `mapPartitions` + """ + num_processes = self.num_processes + use_gpu = self.use_gpu + input_params = self.input_params + driver_address = self.driver_address + log_streaming_server_port = self.log_streaming_server_port + is_spark_local_master = self.is_spark_local_master + driver_owned_gpus: List[str] = [] + if is_spark_local_master and use_gpu: + driver_owned_gpus = _get_gpus_owned(self.spark) + + if input_dataframe is not None: + schema_json = input_dataframe.schema.jsonValue() + else: + schema_json = None + + # Spark task program + def wrapped_train_fn(iterator): # type: ignore[no-untyped-def] + import os + import pandas as pd + import pyarrow + from pyspark import BarrierTaskContext + + CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" + + def get_free_port(address: str, context: "BarrierTaskContext") -> int: + port = "" + if context.partitionId() == 0: + try: + import socket + + sock = socket.socket() + sock.bind((address, 0)) + port = sock.getsockname()[1] + except socket.error: + pass + available_port = context.allGather(str(port))[0] + if not available_port: + raise RuntimeError("Failed to find free port for distributed training.") + return int(available_port) + + def set_torch_config(context: "BarrierTaskContext") -> None: + addrs = [e.address.split(":")[0] for e in context.getTaskInfos()] + + os.environ["MASTER_ADDR"] = str(addrs[0]) + os.environ["MASTER_PORT"] = str(get_free_port(addrs[0], context)) + os.environ["WORLD_SIZE"] = str(num_processes) + os.environ["NODE_RANK"] = str(context.partitionId()) + os.environ["RANK"] = str(context.partitionId()) + + if context.partitionId() >= num_processes: + raise ValueError( + "TorchDistributor._train_on_dataframe requires setting num_processes " + "equal to input spark dataframe partition number." + ) + + if is_spark_local_master: + # distributed training on a local mode spark cluster + def set_gpus(context: "BarrierTaskContext") -> None: + if CUDA_VISIBLE_DEVICES in os.environ: + return + + gpu_owned = driver_owned_gpus[context.partitionId()] + os.environ[CUDA_VISIBLE_DEVICES] = gpu_owned + + else: + + def set_gpus(context: "BarrierTaskContext") -> None: + if CUDA_VISIBLE_DEVICES in os.environ: + return + + gpus_owned = _get_gpus_owned(context) + os.environ[CUDA_VISIBLE_DEVICES] = ",".join(gpus_owned) + + context = BarrierTaskContext.get() + + if use_gpu: + set_gpus(context) + else: + os.environ[CUDA_VISIBLE_DEVICES] = "" + set_torch_config(context) + + log_streaming_client = LogStreamingClient(driver_address, log_streaming_server_port) + input_params["log_streaming_client"] = log_streaming_client + try: + with TorchDistributor._setup_spark_partition_data(iterator, schema_json): + output = TorchDistributor._get_output_from_framework_wrapper( + framework_wrapper_fn, + input_params, + train_object, + run_pytorch_file_fn, + *args, + **kwargs, + ) + finally: + try: + LogStreamingClient._destroy() + except BaseException: + pass + + if context.partitionId() == 0: + output_bytes = cloudpickle.dumps(output) + output_size = len(output_bytes) + + # In Spark Connect, DataFrame.collect stacks rows to size + # 'spark.connect.grpc.arrow.maxBatchSize' (default 4MiB), + # here use 4KiB for each chunk, which mean each arrow batch + # may contain about 1000 chunks. + chunks = [] + chunk_size = 4096 + index = 0 + while index < output_size: + chunks.append(output_bytes[index : index + chunk_size]) + index += chunk_size + + yield pyarrow.RecordBatch.from_pandas(pd.DataFrame(data={"chunk": chunks})) + + return wrapped_train_fn + + def _run_distributed_training( + self, + framework_wrapper_fn: Callable, + train_object: Union[Callable, str], + run_pytorch_file_fn: Optional[Callable], + spark_dataframe: Optional["DataFrame"], + *args: Any, + **kwargs: Any, + ) -> Optional[Any]: + if not framework_wrapper_fn: + raise RuntimeError("Unknown combination of parameters") + + log_streaming_server = LogStreamingServer() + self.driver_address = _get_conf(self.spark, "spark.driver.host", "") + assert self.driver_address != "" + try: + log_streaming_server.start(spark_host_address=self.driver_address) + time.sleep(1) # wait for the server to start + self.log_streaming_server_port = log_streaming_server.port + except Exception as e: + # If starting log streaming server failed, we don't need to break + # the distributor training but emit a warning instead. + self.log_streaming_server_port = -1 + self.logger.warning( + "Start torch distributor log streaming server failed, " + "You cannot receive logs sent from distributor workers, ", + f"error: {repr(e)}.", + ) + + try: + spark_task_function = self._get_spark_task_function( + framework_wrapper_fn, + train_object, + run_pytorch_file_fn, + spark_dataframe, + *args, + **kwargs, + ) + self._check_encryption() + self.logger.info( + f"Started distributed training with {self.num_processes} executor processes" + ) + if spark_dataframe is not None: + input_df = spark_dataframe + else: + input_df = self.spark.range( + start=0, end=self.num_tasks, step=1, numPartitions=self.num_tasks + ) + rows = input_df.mapInArrow( + func=spark_task_function, schema="chunk binary", barrier=True + ).collect() + output_bytes = b"".join([row.chunk for row in rows]) + result = cloudpickle.loads(output_bytes) + finally: + log_streaming_server.shutdown() + self.logger.info( + f"Finished distributed training with {self.num_processes} executor processes" + ) + return result + + @staticmethod + def _run_training_on_pytorch_file( + input_params: Dict[str, Any], train_path: str, *args: Any, **kwargs: Any + ) -> None: + if kwargs: + raise ValueError("Running pytorch file does not support key-word type arguments.") + log_streaming_client = input_params.get("log_streaming_client", None) + training_command = TorchDistributor._create_torchrun_command( + input_params, train_path, *args + ) + TorchDistributor._execute_command( + training_command, log_streaming_client=log_streaming_client + ) + + @staticmethod + @contextmanager + def _setup_files( + train_fn: Callable, *args: Any, **kwargs: Any + ) -> Generator[Tuple[str, str], None, None]: + save_dir = TorchDistributor._create_save_dir() + pickle_file_path = TorchDistributor._save_pickled_function( + save_dir, train_fn, *args, **kwargs + ) + output_file_path = os.path.join(save_dir, TorchDistributor._PICKLED_OUTPUT_FILE) + train_file_path = TorchDistributor._create_torchrun_train_file( + save_dir, pickle_file_path, output_file_path + ) + try: + yield (train_file_path, output_file_path) + finally: + TorchDistributor._cleanup_files(save_dir) + + @staticmethod + @contextmanager + def _setup_spark_partition_data( + partition_data_iterator: Iterator[Any], input_schema_json: Dict[str, Any] + ) -> Iterator[Any]: + from pyspark.sql.pandas.serializers import ArrowStreamSerializer + from pyspark.files import SparkFiles + import json + + if input_schema_json is None: + yield + return + + # We need to temporarily write partition data into a temp dir, + # partition data might be huge, so we need to write it under + # configured `SPARK_LOCAL_DIRS`. + save_dir = TorchDistributor._create_save_dir(root_dir=SparkFiles.getRootDirectory()) + + try: + serializer = ArrowStreamSerializer() + arrow_file_path = os.path.join(save_dir, "data.arrow") + with open(arrow_file_path, "wb") as f: + serializer.dump_stream(partition_data_iterator, f) + if f.tell() == 0: + # Nothing is written to file, this partition is empty + raise ValueError( + "Empty Spark partition is not allowed in " + "TorchDistributor.train_on_dataframe." + ) + + 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: + f.write(schema_json_string) + + os.environ[SPARK_PARTITION_ARROW_DATA_FILE] = arrow_file_path + os.environ[SPARK_DATAFRAME_SCHEMA_FILE] = schema_file_path + yield + finally: + os.environ.pop(SPARK_PARTITION_ARROW_DATA_FILE) + os.environ.pop(SPARK_DATAFRAME_SCHEMA_FILE) + TorchDistributor._cleanup_files(save_dir) + + @staticmethod + def _run_training_on_pytorch_function( + input_params: Dict[str, Any], + train_fn: Callable, + run_pytorch_file_fn: Optional[Callable], + *args: Any, + **kwargs: Any, + ) -> Any: + + if not run_pytorch_file_fn: + run_pytorch_file_fn = TorchDistributor._run_training_on_pytorch_file + + with TorchDistributor._setup_files(train_fn, *args, **kwargs) as ( + train_file_path, + output_file_path, + ): + run_pytorch_file_fn(input_params, train_file_path) + if not os.path.exists(output_file_path): + raise RuntimeError( + "TorchDistributor failed during training." + "View stdout logs for detailed error message." + ) + try: + output = TorchDistributor._get_pickled_output(output_file_path) + except Exception as e: + raise RuntimeError( + "TorchDistributor failed due to a pickling error. " + "View stdout logs for detailed error message." + ) from e + return output + + @staticmethod + def _create_save_dir(root_dir: Optional[str] = None) -> str: + # TODO: need to do this in a safe way to avoid issues during concurrent runs + return tempfile.mkdtemp(dir=root_dir) + + @staticmethod + def _cleanup_files(save_dir: str) -> None: + shutil.rmtree(save_dir, ignore_errors=True) + + @staticmethod + def _save_pickled_function( + save_dir: str, train_fn: Union[str, Callable], *args: Any, **kwargs: Any + ) -> str: + saved_pickle_path = os.path.join(save_dir, TorchDistributor._PICKLED_FUNC_FILE) + with open(saved_pickle_path, "wb") as f: + cloudpickle.dump((train_fn, args, kwargs), f) + return saved_pickle_path + + @staticmethod + def _create_torchrun_train_file( + save_dir_path: str, pickle_file_path: str, output_file_path: str + ) -> str: + code = textwrap.dedent( + f""" + from pyspark import cloudpickle + import os + + if __name__ == "__main__": + with open("{pickle_file_path}", "rb") as f: + train_fn, args, kwargs = cloudpickle.load(f) + output = train_fn(*args, **kwargs) + with open("{output_file_path}", "wb") as f: + cloudpickle.dump(output, f) + """ + ) + saved_file_path = os.path.join(save_dir_path, TorchDistributor._TRAIN_FILE) + with open(saved_file_path, "w") as f: + f.write(code) + return saved_file_path + + @staticmethod + def _get_pickled_output(output_file_path: str) -> Any: + with open(output_file_path, "rb") as f: + output = cloudpickle.load(f) + return output + + def run(self, train_object: Union[Callable, str], *args: Any, **kwargs: Any) -> Optional[Any]: + """Runs distributed training. + + Parameters + ---------- + train_object : callable object or str + Either a PyTorch function, PyTorch Lightning function, or the path to a python file + that launches distributed training. + args : + If train_object is a python function and not a path to a python file, args need + to be the input parameters to that function. It would look like + + >>> model = distributor.run(train, 1e-3, 64) + + where train is a function and 1e-3 and 64 are regular numeric inputs to the function. + + If train_object is a python file, then args would be the command-line arguments for + that python file which are all in the form of strings. An example would be + + >>> distributor.run("/path/to/train.py", "--learning-rate=1e-3", "--batch-size=64") + + where since the input is a path, all of the parameters are strings that can be + handled by argparse in that python file. + kwargs : + If train_object is a python function and not a path to a python file, kwargs need + to be the key-word input parameters to that function. It would look like + + >>> model = distributor.run(train, tol=1e-3, max_iter=64) + + where train is a function of 2 arguments `tol` and `max_iter`. + + If train_object is a python file, then you should not set kwargs arguments. + + Returns + ------- + Returns the output of train_object called with args inside spark rank 0 task if the + train_object is a Callable with an expected output. Returns None if train_object is + a file. + """ + return self._run( + train_object, TorchDistributor._run_training_on_pytorch_file, *args, **kwargs + ) + + def _run( + self, + train_object: Union[Callable, str], + run_pytorch_file_fn: Callable, + *args: Any, + **kwargs: Any, + ) -> Optional[Any]: + if isinstance(train_object, str): + framework_wrapper_fn = run_pytorch_file_fn + else: + framework_wrapper_fn = TorchDistributor._run_training_on_pytorch_function + if self.local_mode: + output = self._run_local_training( + framework_wrapper_fn, train_object, run_pytorch_file_fn, *args, **kwargs + ) + else: + output = self._run_distributed_training( + framework_wrapper_fn, train_object, run_pytorch_file_fn, None, *args, **kwargs + ) + return output + + def _train_on_dataframe( + self, + train_function: Callable, + spark_dataframe: "DataFrame", + *args: Any, + **kwargs: Any, + ) -> Any: + """ + Runs distributed training using provided Spark DataFrame as input data. + You should ensure the input Spark DataFrame have evenly distributed partitions, + and this method starts a barrier Spark job that each Spark task in the job + process one partition of the input Spark DataFrame. + + Parameters + ---------- + train_function : + Either a PyTorch function, PyTorch Lightning function that launches distributed + training. Note that inside the function, you can call + `pyspark.ml.torch.distributor.get_spark_partition_data_loader` API to get a torch + data loader, the data loader loads data from the corresponding partition of the + input Spark DataFrame. + spark_dataframe : + An input Spark DataFrame that can be used in PyTorch `train_function` function. + See `train_function` argument doc for details. + args : + `args` need to be the input parameters to `train_function` function. It would look like + + >>> model = distributor.run(train, 1e-3, 64) + + where train is a function and 1e-3 and 64 are regular numeric inputs to the function. + kwargs : + `kwargs` need to be the key-word input parameters to `train_function` function. + It would look like + + >>> model = distributor.run(train, tol=1e-3, max_iter=64) + + where train is a function of 2 arguments `tol` and `max_iter`. + + Returns + ------- + Returns the output of `train_function` called with args inside Spark rank 0 task. + """ + + if self.local_mode: + raise ValueError( + "TorchDistributor.train_on_dataframe requires setting " + "TorchDistributor.local_mode to False." + ) + + return self._run_distributed_training( + TorchDistributor._run_training_on_pytorch_function, + train_function, + TorchDistributor._run_training_on_pytorch_file, + spark_dataframe, + *args, + **kwargs, + ) + + +def _get_spark_partition_data_loader( + num_samples: int, batch_size: int, num_workers: int = 1, prefetch_factor: int = 2 +) -> Any: + """ + This function must be called inside the `train_function` where `train_function` + is the input argument of `TorchDistributor.train_on_dataframe`. + The function returns a pytorch data loader that loads data from + the corresponding spark partition data. + + Parameters + ---------- + num_samples : + Number of samples to generate per epoch. If `num_samples` is less than the number of + rows in the spark partition, it generate the first `num_samples` rows of + the spark partition, if `num_samples` is greater than the number of + rows in the spark partition, then after the iterator loaded all rows from the partition, + it wraps round back to the first row. + batch_size: + How many samples per batch to load. + num_workers: + How many subprocesses to use for data loading. + 0 means that the data will be loaded in the main process. + prefetch_factor: + Number of batches loaded in advance by each worker + """ + from pyspark.sql.types import StructType + from pyspark.ml.torch.data import _SparkPartitionTorchDataset + from torch.utils.data import DataLoader + + arrow_file = os.environ[SPARK_PARTITION_ARROW_DATA_FILE] + schema_file = os.environ[SPARK_DATAFRAME_SCHEMA_FILE] + + with open(schema_file, "r") as fp: + schema = StructType.fromJson(json.load(fp)) + + dataset = _SparkPartitionTorchDataset(arrow_file, schema, num_samples) + + if num_workers > 0: + return DataLoader( + dataset, batch_size, num_workers=num_workers, prefetch_factor=prefetch_factor + ) + else: + # if num_workers is zero, we cannot set `prefetch_factor` otherwise + # torch will raise error. + return DataLoader(dataset, batch_size, num_workers=num_workers) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/log_communication.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/log_communication.py new file mode 100644 index 0000000..8efa83e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/log_communication.py @@ -0,0 +1,198 @@ +# +# 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. +# +# type: ignore + +from contextlib import closing +import time +import socket +import socketserver +from struct import pack, unpack +import sys +import threading +import traceback +from typing import Generator +import warnings + +# Use b'\x00' as separator instead of b'\n', because the bytes are encoded in utf-8 +_SERVER_POLL_INTERVAL = 0.1 +_TRUNCATE_MSG_LEN = 4000 + +_log_print_lock = threading.Lock() # pylint: disable=invalid-name + + +def _get_log_print_lock() -> threading.Lock: + return _log_print_lock + + +class WriteLogToStdout(socketserver.StreamRequestHandler): + def _read_bline(self) -> Generator[bytes, None, None]: + while self.server.is_active: + packed_number_bytes = self.rfile.read(4) + if not packed_number_bytes: + time.sleep(_SERVER_POLL_INTERVAL) + continue + number_bytes = unpack(">i", packed_number_bytes)[0] + message = self.rfile.read(number_bytes) + yield message + + def handle(self) -> None: + self.request.setblocking(0) # non-blocking mode + for bline in self._read_bline(): + with _get_log_print_lock(): + sys.stderr.write(bline.decode("utf-8") + "\n") + sys.stderr.flush() + + +# What is run on the local driver +class LogStreamingServer: + def __init__(self) -> None: + self.server = None + self.serve_thread = None + self.port = None + + @staticmethod + def _get_free_port(spark_host_address: str = "") -> int: + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as tcp: + tcp.bind((spark_host_address, 0)) + _, port = tcp.getsockname() + return port + + def start(self, spark_host_address: str = "") -> None: + if self.server: + raise RuntimeError("Cannot start the server twice.") + + def serve_task(port: int) -> None: + with socketserver.ThreadingTCPServer(("0.0.0.0", port), WriteLogToStdout) as server: + self.server = server + server.is_active = True + server.serve_forever(poll_interval=_SERVER_POLL_INTERVAL) + + self.port = LogStreamingServer._get_free_port(spark_host_address) + self.serve_thread = threading.Thread(target=serve_task, args=(self.port,)) + self.serve_thread.setDaemon(True) + self.serve_thread.start() + + def shutdown(self) -> None: + if self.server: + # Sleep to ensure all log has been received and printed. + time.sleep(_SERVER_POLL_INTERVAL * 2) + # Before close we need flush to ensure all stdout buffer were printed. + sys.stdout.flush() + self.server.is_active = False + self.server.shutdown() + self.serve_thread.join() + self.server = None + self.serve_thread = None + + +class LogStreamingClientBase: + @staticmethod + def _maybe_truncate_msg(message: str) -> str: + if len(message) > _TRUNCATE_MSG_LEN: + message = message[:_TRUNCATE_MSG_LEN] + return message + "...(truncated)" + else: + return message + + def send(self, message: str) -> None: + pass + + def close(self) -> None: + pass + + +class LogStreamingClient(LogStreamingClientBase): + """ + A client that streams log messages to :class:`LogStreamingServer`. + In case of failures, the client will skip messages instead of raising an error. + """ + + _log_callback_client = None + _server_address = None + _singleton_lock = threading.Lock() + + @staticmethod + def _init(address: str, port: int) -> None: + LogStreamingClient._server_address = (address, port) + + @staticmethod + def _destroy() -> None: + LogStreamingClient._server_address = None + if LogStreamingClient._log_callback_client is not None: + LogStreamingClient._log_callback_client.close() + + def __init__(self, address: str, port: int, timeout: int = 10): + """ + Creates a connection to the logging server and authenticates.This client is best effort, + if authentication or sending a message fails, the client will be marked as not alive and + stop trying to send message. + + :param address: Address where the service is running. + :param port: Port where the service is listening for new connections. + """ + self.address = address + self.port = port + self.timeout = timeout + self.sock = None + self.failed = True + self._lock = threading.RLock() + + def _fail(self, error_msg: str) -> None: + self.failed = True + warnings.warn(f"{error_msg}: {traceback.format_exc()}\n") + + def _connect(self) -> None: + if self.port == -1: + self._fail("Log streaming server is not available.") + return + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(self.timeout) + sock.connect((self.address, self.port)) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self.sock = sock + self.failed = False + except (OSError, IOError): # pylint: disable=broad-except + self._fail("Error connecting log streaming server") + + def send(self, message: str) -> None: + """ + Sends a message. + """ + with self._lock: + if self.sock is None: + self._connect() + if not self.failed: + try: + message = LogStreamingClientBase._maybe_truncate_msg(message) + # TODO: + # 1) addressing issue: idle TCP connection might get disconnected by + # cloud provider + # 2) sendall may block when server is busy handling data. + binary_message = message.encode("utf-8") + packed_number_bytes = pack(">i", len(binary_message)) + self.sock.sendall(packed_number_bytes + binary_message) + except Exception: # pylint: disable=broad-except + self._fail("Error sending logs to driver, stopping log streaming") + + def close(self) -> None: + """ + Closes the connection. + """ + if self.sock: + self.sock.close() + self.sock = None diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_data_loader.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_data_loader.py new file mode 100644 index 0000000..67ab6e3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_data_loader.py @@ -0,0 +1,136 @@ +# +# 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 numpy as np +import unittest + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + +from pyspark.ml.torch.distributor import ( + TorchDistributor, + _get_spark_partition_data_loader, +) +from pyspark.sql import SparkSession +from pyspark.ml.linalg import Vectors + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorDataLoaderUnitTests(unittest.TestCase): + def setUp(self) -> None: + self.spark = ( + SparkSession.builder.master("local[1]") + .config("spark.default.parallelism", "1") + .getOrCreate() + ) + + def tearDown(self) -> None: + self.spark.stop() + + def _check_data_loader_result_correctness(self, result, expected): + assert len(result) == len(expected) + + for res_row, exp_row in zip(result, expected): + assert len(res_row) == len(exp_row) + for res_field, exp_field in zip(res_row, exp_row): + np.testing.assert_almost_equal(res_field.numpy(), exp_field) + + def test_data_loader(self): + spark_df = self.spark.createDataFrame( + [ + (Vectors.dense([1.0, 2.0, 3.5]), 0, 10.5), + (Vectors.sparse(3, [1, 2], [4.5, 5.5]), 3, 12.5), + (Vectors.dense([6.0, 7.0, 8.5]), 1, 1.5), + (Vectors.sparse(3, [0, 2], [-2.5, -6.5]), 2, 9.5), + ], + schema=["features", "label", "weight"], + ) + + torch_distributor = TorchDistributor(local_mode=False, use_gpu=False) + + def train_function(num_samples, batch_size): + data_loader = _get_spark_partition_data_loader(num_samples, batch_size) + return list(data_loader) + + result = torch_distributor._train_on_dataframe( + train_function, + spark_df, + num_samples=4, + batch_size=2, + ) + self._check_data_loader_result_correctness( + result, + [ + [[[1.0, 2.0, 3.5], [0.0, 4.5, 5.5]], [0, 3], [10.5, 12.5]], + [[[6.0, 7.0, 8.5], [-2.5, 0.0, -6.5]], [1, 2], [1.5, 9.5]], + ], + ) + + result = torch_distributor._train_on_dataframe( + train_function, + spark_df, + num_samples=4, + batch_size=3, + ) + self._check_data_loader_result_correctness( + result, + [ + [ + [[1.0, 2.0, 3.5], [0.0, 4.5, 5.5], [6.0, 7.0, 8.5]], + [0, 3, 1], + [10.5, 12.5, 1.5], + ], + [[[-2.5, 0.0, -6.5]], [2], [9.5]], + ], + ) + + result = torch_distributor._train_on_dataframe( + train_function, + spark_df, + num_samples=6, + batch_size=3, + ) + self._check_data_loader_result_correctness( + result, + [ + [ + [[1.0, 2.0, 3.5], [0.0, 4.5, 5.5], [6.0, 7.0, 8.5]], + [0, 3, 1], + [10.5, 12.5, 1.5], + ], + [ + [[-2.5, 0.0, -6.5], [1.0, 2.0, 3.5], [0.0, 4.5, 5.5]], + [2, 0, 3], + [9.5, 10.5, 12.5], + ], + ], + ) + + +if __name__ == "__main__": + from pyspark.ml.torch.tests.test_data_loader import * # noqa: F401,F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_distributor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_distributor.py new file mode 100644 index 0000000..364ed83 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_distributor.py @@ -0,0 +1,561 @@ +# +# 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 contextlib +import os +import shutil +from six import StringIO +import stat +import subprocess +import sys +import time +import tempfile +import threading +from typing import Callable, Dict, Any +import unittest +from unittest.mock import patch + +have_torch = True +try: + import torch # noqa: F401 +except ImportError: + have_torch = False + +from pyspark import SparkConf, SparkContext +from pyspark.ml.torch.distributor import TorchDistributor, _get_gpus_owned +from pyspark.ml.torch.torch_run_process_wrapper import clean_and_terminate, check_parent_alive +from pyspark.sql import SparkSession +from pyspark.testing.utils import SPARK_HOME + + +@contextlib.contextmanager +def patch_stdout() -> StringIO: + """patch stdout and give an output""" + sys_stdout = sys.stdout + io_out = StringIO() + sys.stdout = io_out + try: + yield io_out + finally: + sys.stdout = sys_stdout + + +def create_training_function(mnist_dir_path: str) -> Callable: + import torch.nn as nn + import torch.nn.functional as F + from torchvision import transforms, datasets + + batch_size = 100 + num_epochs = 1 + momentum = 0.5 + + train_dataset = datasets.MNIST( + mnist_dir_path, + train=True, + download=True, + transform=transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ), + ) + + class Net(nn.Module): + def __init__(self) -> None: + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x: Any) -> Any: + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x) + + def train_fn(learning_rate: float) -> Any: + import torch + import torch.optim as optim + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel as DDP + from torch.utils.data.distributed import DistributedSampler + + dist.init_process_group("gloo") + + train_sampler = DistributedSampler(dataset=train_dataset) + data_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, sampler=train_sampler + ) + + model = Net() + ddp_model = DDP(model) + optimizer = optim.SGD(ddp_model.parameters(), lr=learning_rate, momentum=momentum) + for epoch in range(1, num_epochs + 1): + ddp_model.train() + for _, (data, target) in enumerate(data_loader): + optimizer.zero_grad() + output = ddp_model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + print(f"epoch {epoch} finished.") + + return "success" * 4096 + + return train_fn + + +def set_up_test_dirs(): + gpu_discovery_script_file = tempfile.NamedTemporaryFile(delete=False) + gpu_discovery_script_file_name = gpu_discovery_script_file.name + try: + gpu_discovery_script_file.write( + b'echo {\\"name\\": \\"gpu\\", \\"addresses\\": [\\"0\\",\\"1\\",\\"2\\"]}' + ) + finally: + gpu_discovery_script_file.close() + + # create temporary directory for Worker resources coordination + tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(tempdir.name) + os.chmod( + gpu_discovery_script_file_name, + stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH, + ) + mnist_dir_path = tempfile.mkdtemp() + + return (gpu_discovery_script_file_name, mnist_dir_path) + + +def get_local_mode_conf(): + return { + "spark.test.home": SPARK_HOME, + "spark.driver.resource.gpu.amount": "3", + "spark.driver.memory": "512M", + "spark.executor.memory": "512M", + } + + +def get_distributed_mode_conf(): + return { + "spark.test.home": SPARK_HOME, + "spark.worker.resource.gpu.amount": "3", + "spark.task.cpus": "2", + "spark.task.resource.gpu.amount": "1", + "spark.executor.resource.gpu.amount": "1", + "spark.driver.memory": "512M", + "spark.executor.memory": "512M", + } + + +class TorchDistributorBaselineUnitTestsMixin: + def setup_env_vars(self, input_map: Dict[str, str]) -> None: + for key, value in input_map.items(): + os.environ[key] = value + + def delete_env_vars(self, input_map: Dict[str, str]) -> None: + for key in input_map.keys(): + del os.environ[key] + + def test_validate_correct_inputs(self) -> None: + inputs = [ + (1, True, False), + (100, True, False), + (1, False, False), + (100, False, False), + ] + for num_processes, local_mode, use_gpu in inputs: + with self.subTest(): + expected_params = { + "num_processes": num_processes, + "local_mode": local_mode, + "use_gpu": use_gpu, + "num_tasks": num_processes, + } + dist = TorchDistributor(num_processes, local_mode, use_gpu) + self.assertEqual(expected_params, dist.input_params) + + def test_validate_incorrect_inputs(self) -> None: + inputs = [ + (0, False, False, ValueError, "positive"), + ] + for num_processes, local_mode, use_gpu, error, message in inputs: + with self.subTest(): + with self.assertRaisesRegex(error, message): + TorchDistributor(num_processes, local_mode, use_gpu) + + def test_encryption_passes(self) -> None: + inputs = [ + ("spark.ssl.enabled", "false", "pytorch.spark.distributor.ignoreSsl", "true"), + ("spark.ssl.enabled", "false", "pytorch.spark.distributor.ignoreSsl", "false"), + ("spark.ssl.enabled", "true", "pytorch.spark.distributor.ignoreSsl", "true"), + ] + for ssl_conf_key, ssl_conf_value, pytorch_conf_key, pytorch_conf_value in inputs: + with self.subTest(): + self.spark.conf.set(ssl_conf_key, ssl_conf_value) + self.spark.conf.set(pytorch_conf_key, pytorch_conf_value) + distributor = TorchDistributor(1, True, False) + distributor._check_encryption() + + def test_encryption_fails(self) -> None: + # this is the only combination that should fail + inputs = [("spark.ssl.enabled", "true", "pytorch.spark.distributor.ignoreSsl", "false")] + for ssl_conf_key, ssl_conf_value, pytorch_conf_key, pytorch_conf_value in inputs: + with self.subTest(): + with self.assertRaisesRegex(Exception, "encryption"): + self.spark.conf.set(ssl_conf_key, ssl_conf_value) + self.spark.conf.set(pytorch_conf_key, pytorch_conf_value) + distributor = TorchDistributor(1, True, False) + distributor._check_encryption() + + def test_get_num_tasks_fails(self) -> None: + inputs = [1, 5, 4] + + # This is when the conf isn't set and we request GPUs + for num_processes in inputs: + with self.subTest(): + with self.assertRaisesRegex(RuntimeError, "driver"): + TorchDistributor(num_processes, True, True) + with self.assertRaisesRegex(RuntimeError, "unset"): + TorchDistributor(num_processes, False, True) + + def test_execute_command(self) -> None: + """Test that run command runs the process and logs are written correctly""" + + with patch_stdout() as output: + stdout_command = ["echo", "hello_stdout"] + TorchDistributor._execute_command(stdout_command) + self.assertIn( + "hello_stdout", output.getvalue().strip(), "hello_stdout should print to stdout" + ) + + with patch_stdout() as output: + stderr_command = ["bash", "-c", "echo hello_stderr >&2"] + TorchDistributor._execute_command(stderr_command) + self.assertIn( + "hello_stderr", output.getvalue().strip(), "hello_stderr should print to stdout" + ) + + # include command in the exception message + with self.assertRaisesRegex(RuntimeError, "exit 1"): + error_command = ["bash", "-c", "exit 1"] + TorchDistributor._execute_command(error_command) + + with self.assertRaisesRegex(RuntimeError, "abcdef"): + error_command = ["bash", "-c", "'abc''def'"] + TorchDistributor._execute_command(error_command) + + def test_create_torchrun_command(self) -> None: + train_path = "train.py" + args_string = ["1", "3"] + local_mode_input_params = {"num_processes": 4, "local_mode": True} + + expected_local_mode_output = [ + sys.executable, + "-m", + "pyspark.ml.torch.torch_run_process_wrapper", + "--standalone", + "--nnodes=1", + "--nproc_per_node=4", + "train.py", + "1", + "3", + ] + self.assertEqual( + TorchDistributor._create_torchrun_command( + local_mode_input_params, train_path, *args_string + ), + expected_local_mode_output, + ) + + distributed_mode_input_params = {"num_processes": 4, "local_mode": False} + input_env_vars = {"MASTER_ADDR": "localhost", "MASTER_PORT": "9350", "RANK": "3"} + + args_number = [1, 3] # testing conversion to strings + self.setup_env_vars(input_env_vars) + expected_distributed_mode_output = [ + sys.executable, + "-m", + "pyspark.ml.torch.torch_run_process_wrapper", + "--nnodes=4", + "--node_rank=3", + "--rdzv_endpoint=localhost:9350", + "--rdzv_id=0", + "--nproc_per_node=1", + "train.py", + "1", + "3", + ] + self.assertEqual( + TorchDistributor._create_torchrun_command( + distributed_mode_input_params, train_path, *args_number + ), + expected_distributed_mode_output, + ) + self.delete_env_vars(input_env_vars) + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorBaselineUnitTests(TorchDistributorBaselineUnitTestsMixin, unittest.TestCase): + @classmethod + def setUpClass(cls): + conf = SparkConf() + sc = SparkContext("local[4]", conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + +class TorchDistributorLocalUnitTestsMixin: + def setup_env_vars(self, input_map: Dict[str, str]) -> None: + for key, value in input_map.items(): + os.environ[key] = value + + def delete_env_vars(self, input_map: Dict[str, str]) -> None: + for key in input_map.keys(): + del os.environ[key] + + def test_get_num_tasks_locally(self) -> None: + succeeds = [1, 2] + fails = [4, 8] + for num_processes in succeeds: + with self.subTest(): + expected_output = num_processes + distributor = TorchDistributor(num_processes, True, True) + self.assertEqual(distributor._get_num_tasks(), expected_output) + + for num_processes in fails: + with self.subTest(): + with self.assertLogs("TorchDistributor", level="WARNING") as log: + distributor = TorchDistributor(num_processes, True, True) + self.assertEqual(len(log.records), 1) + self.assertEqual(distributor.num_processes, 3) + + def test_get_gpus_owned_local(self) -> None: + addresses = ["0", "1", "2"] + self.assertEqual(_get_gpus_owned(self.spark), addresses) + + env_vars = {"CUDA_VISIBLE_DEVICES": "3,4,5"} + self.setup_env_vars(env_vars) + self.assertEqual(_get_gpus_owned(self.spark), ["3", "4", "5"]) + self.delete_env_vars(env_vars) + + def _get_inputs_for_test_local_training_succeeds(self): + return [ + ("0,1,2", 1, True, "1"), + ("0,1,2", 3, True, "1,2,0"), + ("0,1,2", 2, False, "0,1,2"), + (None, 3, False, "NONE"), + ] + + def test_local_training_succeeds(self) -> None: + CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" + inputs = self._get_inputs_for_test_local_training_succeeds() + + for i, (cuda_env_var, num_processes, use_gpu, expected) in enumerate(inputs): + with self.subTest(f"subtest: {i + 1}"): + # setup + if cuda_env_var: + self.setup_env_vars({CUDA_VISIBLE_DEVICES: cuda_env_var}) + + dist = TorchDistributor(num_processes, True, use_gpu) + dist._run_training_on_pytorch_file = lambda *args: os.environ.get( + CUDA_VISIBLE_DEVICES, "NONE" + ) + self.assertEqual( + expected, + dist._run_local_training(dist._run_training_on_pytorch_file, "train.py", None), + ) + # cleanup + if cuda_env_var: + self.delete_env_vars({CUDA_VISIBLE_DEVICES: cuda_env_var}) + + def test_local_file_with_pytorch(self) -> None: + test_file_path = "python/test_support/test_pytorch_training_file.py" + learning_rate_str = "0.01" + TorchDistributor(num_processes=2, local_mode=True, use_gpu=False).run( + test_file_path, learning_rate_str + ) + + def test_end_to_end_run_locally(self) -> None: + train_fn = create_training_function(self.mnist_dir_path) + output = TorchDistributor(num_processes=2, local_mode=True, use_gpu=False).run( + train_fn, 0.001 + ) + self.assertEqual(output, "success" * 4096) + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorLocalUnitTests(TorchDistributorLocalUnitTestsMixin, unittest.TestCase): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + conf = SparkConf() + for k, v in get_local_mode_conf().items(): + conf = conf.set(k, v) + conf = conf.set( + "spark.driver.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + + sc = SparkContext("local-cluster[2,2,512]", cls.__name__, conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorLocalUnitTestsII(TorchDistributorLocalUnitTestsMixin, unittest.TestCase): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + conf = SparkConf() + for k, v in get_local_mode_conf().items(): + conf = conf.set(k, v) + conf = conf.set( + "spark.driver.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + + sc = SparkContext("local[4]", cls.__name__, conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + +class TorchDistributorDistributedUnitTestsMixin: + def test_dist_training_succeeds(self) -> None: + CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" + inputs = [ + ("0,1,2", 2, True, "0"), + ] + + for i, (_, num_processes, use_gpu, expected) in enumerate(inputs): + with self.subTest(f"subtest: {i + 1}"): + dist = TorchDistributor(num_processes, False, use_gpu) + dist._run_training_on_pytorch_file = lambda *args: os.environ.get( + CUDA_VISIBLE_DEVICES, "NONE" + ) + self.assertEqual( + expected, + dist._run_distributed_training( + dist._run_training_on_pytorch_file, + "...", + TorchDistributor._run_training_on_pytorch_file, + None, + ), + ) + + def test_get_num_tasks_distributed(self) -> None: + inputs = [(1, 8, 8), (2, 8, 4), (3, 8, 3)] + + for spark_conf_value, num_processes, expected_output in inputs: + with self.subTest(): + self.spark.conf.set("spark.task.resource.gpu.amount", str(spark_conf_value)) + distributor = TorchDistributor(num_processes, False, True) + self.assertEqual(distributor._get_num_tasks(), expected_output) + + self.spark.conf.set("spark.task.resource.gpu.amount", "1") + + def test_distributed_file_with_pytorch(self) -> None: + test_file_path = "python/test_support/test_pytorch_training_file.py" + learning_rate_str = "0.01" + TorchDistributor(num_processes=2, local_mode=False, use_gpu=False).run( + test_file_path, learning_rate_str + ) + + def test_end_to_end_run_distributedly(self) -> None: + train_fn = create_training_function(self.mnist_dir_path) + output = TorchDistributor(num_processes=2, local_mode=False, use_gpu=False).run( + train_fn, 0.001 + ) + self.assertEqual(output, "success" * 4096) + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchDistributorDistributedUnitTests( + TorchDistributorDistributedUnitTestsMixin, unittest.TestCase +): + @classmethod + def setUpClass(cls): + (cls.gpu_discovery_script_file_name, cls.mnist_dir_path) = set_up_test_dirs() + conf = SparkConf() + for k, v in get_distributed_mode_conf().items(): + conf = conf.set(k, v) + conf = conf.set( + "spark.worker.resource.gpu.discoveryScript", cls.gpu_discovery_script_file_name + ) + + sc = SparkContext("local-cluster[2,2,512]", cls.__name__, conf=conf) + cls.spark = SparkSession(sc) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.mnist_dir_path) + os.unlink(cls.gpu_discovery_script_file_name) + cls.spark.stop() + + +class TorchWrapperUnitTestsMixin: + def test_clean_and_terminate(self) -> None: + def kill_task(task: "subprocess.Popen") -> None: + time.sleep(1) + clean_and_terminate(task) + + command = [sys.executable, "-c", '"import time; time.sleep(20)"'] + task = subprocess.Popen(command) + t = threading.Thread(target=kill_task, args=(task,)) + t.start() + time.sleep(2) + self.assertEqual(task.poll(), 0) # implies task ended + + @patch("pyspark.ml.torch.torch_run_process_wrapper.clean_and_terminate") + def test_check_parent_alive(self, mock_clean_and_terminate: Callable) -> None: + command = [sys.executable, "-c", '"import time; time.sleep(2)"'] + task = subprocess.Popen(command) + t = threading.Thread(target=check_parent_alive, args=(task,), daemon=True) + t.start() + time.sleep(2) + self.assertEqual(mock_clean_and_terminate.call_count, 0) + + +@unittest.skipIf(not have_torch, "torch is required") +class TorchWrapperUnitTests(TorchWrapperUnitTestsMixin, unittest.TestCase): + pass + + +if __name__ == "__main__": + from pyspark.ml.torch.tests.test_distributor import * # noqa: F401,F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_log_communication.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_log_communication.py new file mode 100644 index 0000000..164c755 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/tests/test_log_communication.py @@ -0,0 +1,172 @@ +# +# 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 absolute_import, division, print_function + +import contextlib +from six import StringIO +import sys +import time +from typing import Any, Callable +import unittest + +import pyspark.ml.torch.log_communication +from pyspark.ml.torch.log_communication import ( + LogStreamingServer, + LogStreamingClient, + LogStreamingClientBase, + _SERVER_POLL_INTERVAL, +) + + +@contextlib.contextmanager +def patch_stderr() -> StringIO: + """patch stdout and give an output""" + sys_stderr = sys.stderr + io_out = StringIO() + sys.stderr = io_out + try: + yield io_out + finally: + sys.stderr = sys_stderr + + +class LogStreamingServiceTestCase(unittest.TestCase): + def setUp(self) -> None: + self.default_truncate_msg_len = pyspark.ml.torch.log_communication._TRUNCATE_MSG_LEN + pyspark.ml.torch.log_communication._TRUNCATE_MSG_LEN = 10 + + def tearDown(self) -> None: + pyspark.ml.torch.log_communication._TRUNCATE_MSG_LEN = self.default_truncate_msg_len + + def basic_test(self) -> None: + server = LogStreamingServer() + server.start() + time.sleep(1) + client = LogStreamingClient("localhost", server.port) + with patch_stderr() as output: + client.send("msg 001") + client.send("msg 002") + time.sleep(_SERVER_POLL_INTERVAL + 1) + output = output.getvalue() + self.assertIn("msg 001\nmsg 002\n", output) + client.close() + server.shutdown() + + def test_truncate_message(self) -> None: + msg1 = "abc" + assert LogStreamingClientBase._maybe_truncate_msg(msg1) == msg1 + msg2 = "abcdefghijkl" + assert LogStreamingClientBase._maybe_truncate_msg(msg2) == "abcdefghij...(truncated)" + + def test_multiple_clients(self) -> None: + server = LogStreamingServer() + server.start() + time.sleep(1) + client1 = LogStreamingClient("localhost", server.port) + client2 = LogStreamingClient("localhost", server.port) + with patch_stderr() as output: + client1.send("c1 msg1") + time.sleep(_SERVER_POLL_INTERVAL + 1) + client2.send("c2 msg1") + time.sleep(_SERVER_POLL_INTERVAL + 1) + client1.send("c1 msg2") + time.sleep(_SERVER_POLL_INTERVAL + 1) + client2.send("c2 msg2") + time.sleep(_SERVER_POLL_INTERVAL + 1) + output = output.getvalue() + self.assertIn("c1 msg1\nc2 msg1\nc1 msg2\nc2 msg2\n", output) + client1.close() + client2.close() + server.shutdown() + + def test_client_should_fail_gracefully(self) -> None: + server = LogStreamingServer() + server.start() + time.sleep(1) + client = LogStreamingClient("localhost", server.port) + client.send("msg 001") + server.shutdown() + for i in range(5): + client.send("msg 002") + time.sleep(_SERVER_POLL_INTERVAL + 1) + self.assertTrue(client.failed) + client.close() + + def test_client_send_intermittently(self) -> None: + server = LogStreamingServer() + server.start() + time.sleep(1) + client = LogStreamingClient("localhost", server.port) + with patch_stderr() as output: + client._connect() + # test client send half message first + client.send("msg part1") + time.sleep(_SERVER_POLL_INTERVAL + 1) + # test client send another half message + client.send(" msg part2") + time.sleep(_SERVER_POLL_INTERVAL + 1) + output = output.getvalue() + self.assertIn("msg part1\n msg part2\n", output) + client.close() + server.shutdown() + + @staticmethod + def test_server_shutdown() -> None: + def run_test(client_ops: Callable) -> None: + server = LogStreamingServer() + server.start() + time.sleep(1) + client = LogStreamingClient("localhost", server.port) + client_ops(client) + server.shutdown() + client.close() + + def client_ops_close(client: Any) -> None: + client.close() + + def client_ops_send_half_msg(client: Any) -> None: + # Test server only recv incomplete message from client can exit. + client._connect() + client.sock.sendall(b"msg part1 ") + time.sleep(_SERVER_POLL_INTERVAL + 1) + + def client_ops_send_a_msg(client: Any) -> None: + client.send("msg1") + time.sleep(_SERVER_POLL_INTERVAL + 1) + + def client_ops_send_a_msg_and_close(client: Any) -> None: + client.send("msg1") + client.close() + time.sleep(_SERVER_POLL_INTERVAL + 1) + + run_test(client_ops_close) + run_test(client_ops_send_half_msg) + run_test(client_ops_send_a_msg) + run_test(client_ops_send_a_msg_and_close) + + +if __name__ == "__main__": + from pyspark.ml.torch.tests.test_log_communication import * # noqa: F401,F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/torch_run_process_wrapper.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/torch_run_process_wrapper.py new file mode 100644 index 0000000..67ec492 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/torch/torch_run_process_wrapper.py @@ -0,0 +1,83 @@ +# +# 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 os +import signal +import subprocess +import sys +import threading +import time +from typing import Any + + +def clean_and_terminate(task: "subprocess.Popen") -> None: + task.terminate() + time.sleep(0.5) + if task.poll() is None: + task.kill() + # TODO(SPARK-41775): Cleanup temp files + + +def check_parent_alive(task: "subprocess.Popen") -> None: + orig_parent_id = os.getppid() + while True: + if os.getppid() != orig_parent_id: + clean_and_terminate(task) + break + time.sleep(0.5) + + +if __name__ == "__main__": + """ + This is a wrapper around torch.distributed.run and it kills the child process + if the parent process fails, crashes, or exits. + """ + + args = sys.argv[1:] + + cmd = [sys.executable, "-m", "torch.distributed.run", *args] + task = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + env=os.environ, + ) + t = threading.Thread(target=check_parent_alive, args=(task,), daemon=True) + + def sigterm_handler(*args: Any) -> None: + clean_and_terminate(task) + os._exit(0) + + signal.signal(signal.SIGTERM, sigterm_handler) + + t.start() + task.stdin.close() # type: ignore[union-attr] + try: + for line in task.stdout: # type: ignore[union-attr] + decoded = line.decode() + print(decoded.rstrip()) + task.wait() + finally: + if task.poll() is None: + try: + task.terminate() + time.sleep(0.5) + if task.poll() is None: + task.kill() + except OSError: + pass diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tree.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tree.py new file mode 100644 index 0000000..5143c32 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tree.py @@ -0,0 +1,434 @@ +# +# 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 List, Sequence, TypeVar, TYPE_CHECKING + +from pyspark import since +from pyspark.ml.linalg import Vector +from pyspark.ml.param import Params +from pyspark.ml.param.shared import ( + HasCheckpointInterval, + HasSeed, + HasWeightCol, + Param, + TypeConverters, + HasMaxIter, + HasStepSize, + HasValidationIndicatorCol, +) +from pyspark.ml.wrapper import JavaPredictionModel +from pyspark.ml.common import inherit_doc + +if TYPE_CHECKING: + from pyspark.ml._typing import P + +T = TypeVar("T") + + +@inherit_doc +class _DecisionTreeModel(JavaPredictionModel[T]): + """ + Abstraction for Decision Tree models. + + .. versionadded:: 1.5.0 + """ + + @property + @since("1.5.0") + def numNodes(self) -> int: + """Return number of nodes of the decision tree.""" + return self._call_java("numNodes") + + @property + @since("1.5.0") + def depth(self) -> int: + """Return depth of the decision tree.""" + return self._call_java("depth") + + @property + @since("2.0.0") + def toDebugString(self) -> str: + """Full description of model.""" + return self._call_java("toDebugString") + + @since("3.0.0") + def predictLeaf(self, value: Vector) -> float: + """ + Predict the indices of the leaves corresponding to the feature vector. + """ + return self._call_java("predictLeaf", value) + + +class _DecisionTreeParams(HasCheckpointInterval, HasSeed, HasWeightCol): + """ + Mixin for Decision Tree parameters. + """ + + leafCol: Param[str] = Param( + Params._dummy(), + "leafCol", + "Leaf indices column name. Predicted leaf " + + "index of each instance in each tree by preorder.", + typeConverter=TypeConverters.toString, + ) + + maxDepth: Param[int] = Param( + Params._dummy(), + "maxDepth", + "Maximum depth of the tree. (>= 0) E.g., " + + "depth 0 means 1 leaf node; depth 1 means 1 internal node + 2 leaf nodes. " + + "Must be in range [0, 30].", + typeConverter=TypeConverters.toInt, + ) + + maxBins: Param[int] = Param( + Params._dummy(), + "maxBins", + "Max number of bins for discretizing continuous " + + "features. Must be >=2 and >= number of categories for any categorical " + + "feature.", + typeConverter=TypeConverters.toInt, + ) + + minInstancesPerNode: Param[int] = Param( + Params._dummy(), + "minInstancesPerNode", + "Minimum number of " + + "instances each child must have after split. If a split causes " + + "the left or right child to have fewer than " + + "minInstancesPerNode, the split will be discarded as invalid. " + + "Should be >= 1.", + typeConverter=TypeConverters.toInt, + ) + + minWeightFractionPerNode: Param[float] = Param( + Params._dummy(), + "minWeightFractionPerNode", + "Minimum " + "fraction of the weighted sample count that each child " + "must have after split. If a split causes the fraction " + "of the total weight in the left or right child to be " + "less than minWeightFractionPerNode, the split will be " + "discarded as invalid. Should be in interval [0.0, 0.5).", + typeConverter=TypeConverters.toFloat, + ) + + minInfoGain: Param[float] = Param( + Params._dummy(), + "minInfoGain", + "Minimum information gain for a split " + "to be considered at a tree node.", + typeConverter=TypeConverters.toFloat, + ) + + maxMemoryInMB: Param[int] = Param( + Params._dummy(), + "maxMemoryInMB", + "Maximum memory in MB allocated to " + + "histogram aggregation. If too small, then 1 node will be split per " + + "iteration, and its aggregates may exceed this size.", + typeConverter=TypeConverters.toInt, + ) + + cacheNodeIds: Param[bool] = Param( + Params._dummy(), + "cacheNodeIds", + "If false, the algorithm will pass " + + "trees to executors to match instances with nodes. If true, the " + + "algorithm will cache node IDs for each instance. Caching can speed " + + "up training of deeper trees. Users can set how often should the cache " + + "be checkpointed or disable it by setting checkpointInterval.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super(_DecisionTreeParams, self).__init__() + + def setLeafCol(self: "P", value: str) -> "P": + """ + Sets the value of :py:attr:`leafCol`. + """ + return self._set(leafCol=value) + + def getLeafCol(self) -> str: + """ + Gets the value of leafCol or its default value. + """ + return self.getOrDefault(self.leafCol) + + def getMaxDepth(self) -> int: + """ + Gets the value of maxDepth or its default value. + """ + return self.getOrDefault(self.maxDepth) + + def getMaxBins(self) -> int: + """ + Gets the value of maxBins or its default value. + """ + return self.getOrDefault(self.maxBins) + + def getMinInstancesPerNode(self) -> int: + """ + Gets the value of minInstancesPerNode or its default value. + """ + return self.getOrDefault(self.minInstancesPerNode) + + def getMinWeightFractionPerNode(self) -> float: + """ + Gets the value of minWeightFractionPerNode or its default value. + """ + return self.getOrDefault(self.minWeightFractionPerNode) + + def getMinInfoGain(self) -> float: + """ + Gets the value of minInfoGain or its default value. + """ + return self.getOrDefault(self.minInfoGain) + + def getMaxMemoryInMB(self) -> int: + """ + Gets the value of maxMemoryInMB or its default value. + """ + return self.getOrDefault(self.maxMemoryInMB) + + def getCacheNodeIds(self) -> bool: + """ + Gets the value of cacheNodeIds or its default value. + """ + return self.getOrDefault(self.cacheNodeIds) + + +@inherit_doc +class _TreeEnsembleModel(JavaPredictionModel[T]): + """ + (private abstraction) + Represents a tree ensemble model. + """ + + @property + @since("2.0.0") + def trees(self) -> Sequence["_DecisionTreeModel"]: + """Trees in this ensemble. Warning: These have null parent Estimators.""" + return [_DecisionTreeModel(m) for m in list(self._call_java("trees"))] + + @property + @since("2.0.0") + def getNumTrees(self) -> int: + """Number of trees in ensemble.""" + return self._call_java("getNumTrees") + + @property + @since("1.5.0") + def treeWeights(self) -> List[float]: + """Return the weights for each tree""" + return list(self._call_java("javaTreeWeights")) + + @property + @since("2.0.0") + def totalNumNodes(self) -> int: + """Total number of nodes, summed over all trees in the ensemble.""" + return self._call_java("totalNumNodes") + + @property + @since("2.0.0") + def toDebugString(self) -> str: + """Full description of model.""" + return self._call_java("toDebugString") + + @since("3.0.0") + def predictLeaf(self, value: Vector) -> float: + """ + Predict the indices of the leaves corresponding to the feature vector. + """ + return self._call_java("predictLeaf", value) + + +class _TreeEnsembleParams(_DecisionTreeParams): + """ + Mixin for Decision Tree-based ensemble algorithms parameters. + """ + + subsamplingRate: Param[float] = Param( + Params._dummy(), + "subsamplingRate", + "Fraction of the training data " + "used for learning each decision tree, in range (0, 1].", + typeConverter=TypeConverters.toFloat, + ) + + supportedFeatureSubsetStrategies: List[str] = ["auto", "all", "onethird", "sqrt", "log2"] + + featureSubsetStrategy: Param[str] = Param( + Params._dummy(), + "featureSubsetStrategy", + "The number of features to consider for splits at each tree node. Supported " + + "options: 'auto' (choose automatically for task: If numTrees == 1, set to " + + "'all'. If numTrees > 1 (forest), set to 'sqrt' for classification and to " + + "'onethird' for regression), 'all' (use all features), 'onethird' (use " + + "1/3 of the features), 'sqrt' (use sqrt(number of features)), 'log2' (use " + + "log2(number of features)), 'n' (when n is in the range (0, 1.0], use " + + "n * number of features. When n is in the range (1, number of features), use" + + " n features). default = 'auto'", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(_TreeEnsembleParams, self).__init__() + + @since("1.4.0") + def getSubsamplingRate(self) -> float: + """ + Gets the value of subsamplingRate or its default value. + """ + return self.getOrDefault(self.subsamplingRate) + + @since("1.4.0") + def getFeatureSubsetStrategy(self) -> str: + """ + Gets the value of featureSubsetStrategy or its default value. + """ + return self.getOrDefault(self.featureSubsetStrategy) + + +class _RandomForestParams(_TreeEnsembleParams): + """ + Private class to track supported random forest parameters. + """ + + numTrees: Param[int] = Param( + Params._dummy(), + "numTrees", + "Number of trees to train (>= 1).", + typeConverter=TypeConverters.toInt, + ) + + bootstrap: Param[bool] = Param( + Params._dummy(), + "bootstrap", + "Whether bootstrap samples are used " "when building trees.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super(_RandomForestParams, self).__init__() + + @since("1.4.0") + def getNumTrees(self) -> int: + """ + Gets the value of numTrees or its default value. + """ + return self.getOrDefault(self.numTrees) + + @since("3.0.0") + def getBootstrap(self) -> bool: + """ + Gets the value of bootstrap or its default value. + """ + return self.getOrDefault(self.bootstrap) + + +class _GBTParams(_TreeEnsembleParams, HasMaxIter, HasStepSize, HasValidationIndicatorCol): + """ + Private class to track supported GBT params. + """ + + stepSize: Param[float] = Param( + Params._dummy(), + "stepSize", + "Step size (a.k.a. learning rate) in interval (0, 1] for shrinking " + + "the contribution of each estimator.", + typeConverter=TypeConverters.toFloat, + ) + + validationTol: Param[float] = Param( + Params._dummy(), + "validationTol", + "Threshold for stopping early when fit with validation is used. " + + "If the error rate on the validation input changes by less than the " + + "validationTol, then learning will stop early (before `maxIter`). " + + "This parameter is ignored when fit without validation is used.", + typeConverter=TypeConverters.toFloat, + ) + + @since("3.0.0") + def getValidationTol(self) -> float: + """ + Gets the value of validationTol or its default value. + """ + return self.getOrDefault(self.validationTol) + + +class _HasVarianceImpurity(Params): + """ + Private class to track supported impurity measures. + """ + + supportedImpurities: List[str] = ["variance"] + + impurity: Param[str] = Param( + Params._dummy(), + "impurity", + "Criterion used for information gain calculation (case-insensitive). " + + "Supported options: " + + ", ".join(supportedImpurities), + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(_HasVarianceImpurity, self).__init__() + + @since("1.4.0") + def getImpurity(self) -> str: + """ + Gets the value of impurity or its default value. + """ + return self.getOrDefault(self.impurity) + + +class _TreeClassifierParams(Params): + """ + Private class to track supported impurity measures. + + .. versionadded:: 1.4.0 + """ + + supportedImpurities: List[str] = ["entropy", "gini"] + + impurity: Param[str] = Param( + Params._dummy(), + "impurity", + "Criterion used for information gain calculation (case-insensitive). " + + "Supported options: " + + ", ".join(supportedImpurities), + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super(_TreeClassifierParams, self).__init__() + + @since("1.6.0") + def getImpurity(self) -> str: + """ + Gets the value of impurity or its default value. + """ + return self.getOrDefault(self.impurity) + + +class _TreeRegressorParams(_HasVarianceImpurity): + """ + Private class to track supported impurity measures. + """ + + pass diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tuning.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tuning.py new file mode 100644 index 0000000..63f5122 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/tuning.py @@ -0,0 +1,1741 @@ +# +# 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 os +import sys +import itertools +from multiprocessing.pool import ThreadPool + +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, + overload, + TYPE_CHECKING, +) + +import numpy as np + +from pyspark import keyword_only, since, SparkContext, inheritable_thread_target +from pyspark.ml import Estimator, Transformer, Model +from pyspark.ml.common import inherit_doc, _py2java, _java2py +from pyspark.ml.evaluation import Evaluator, JavaEvaluator +from pyspark.ml.param import Params, Param, TypeConverters +from pyspark.ml.param.shared import HasCollectSubModels, HasParallelism, HasSeed +from pyspark.ml.util import ( + DefaultParamsReader, + DefaultParamsWriter, + MetaAlgorithmReadWrite, + MLReadable, + MLReader, + MLWritable, + MLWriter, + JavaMLReader, + JavaMLWriter, +) +from pyspark.ml.wrapper import JavaParams, JavaEstimator, JavaWrapper +from pyspark.sql.functions import col, lit, rand, UserDefinedFunction +from pyspark.sql.types import BooleanType + +from pyspark.sql.dataframe import DataFrame + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + from py4j.java_gateway import JavaObject + from py4j.java_collections import JavaArray + +__all__ = [ + "ParamGridBuilder", + "CrossValidator", + "CrossValidatorModel", + "TrainValidationSplit", + "TrainValidationSplitModel", +] + + +def _parallelFitTasks( + est: Estimator, + train: DataFrame, + eva: Evaluator, + validation: DataFrame, + epm: Sequence["ParamMap"], + collectSubModel: bool, +) -> List[Callable[[], Tuple[int, float, Transformer]]]: + """ + Creates a list of callables which can be called from different threads to fit and evaluate + an estimator in parallel. Each callable returns an `(index, metric)` pair. + + Parameters + ---------- + est : :py:class:`pyspark.ml.baseEstimator` + he estimator to be fit. + train : :py:class:`pyspark.sql.DataFrame` + DataFrame, training data set, used for fitting. + eva : :py:class:`pyspark.ml.evaluation.Evaluator` + used to compute `metric` + validation : :py:class:`pyspark.sql.DataFrame` + DataFrame, validation data set, used for evaluation. + epm : :py:class:`collections.abc.Sequence` + Sequence of ParamMap, params maps to be used during fitting & evaluation. + collectSubModel : bool + Whether to collect sub model. + + Returns + ------- + tuple + (int, float, subModel), an index into `epm` and the associated metric value. + """ + modelIter = est.fitMultiple(train, epm) + + def singleTask() -> Tuple[int, float, Transformer]: + index, model = next(modelIter) + # TODO: duplicate evaluator to take extra params from input + # Note: Supporting tuning params in evaluator need update method + # `MetaAlgorithmReadWrite.getAllNestedStages`, make it return + # all nested stages and evaluators + metric = eva.evaluate(model.transform(validation, epm[index])) + return index, metric, model if collectSubModel else None + + return [singleTask] * len(epm) + + +class ParamGridBuilder: + r""" + Builder for a param grid used in grid search-based model selection. + + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.classification import LogisticRegression + >>> lr = LogisticRegression() + >>> output = ParamGridBuilder() \ + ... .baseOn({lr.labelCol: 'l'}) \ + ... .baseOn([lr.predictionCol, 'p']) \ + ... .addGrid(lr.regParam, [1.0, 2.0]) \ + ... .addGrid(lr.maxIter, [1, 5]) \ + ... .build() + >>> expected = [ + ... {lr.regParam: 1.0, lr.maxIter: 1, lr.labelCol: 'l', lr.predictionCol: 'p'}, + ... {lr.regParam: 2.0, lr.maxIter: 1, lr.labelCol: 'l', lr.predictionCol: 'p'}, + ... {lr.regParam: 1.0, lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}, + ... {lr.regParam: 2.0, lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}] + >>> len(output) == len(expected) + True + >>> all([m in expected for m in output]) + True + """ + + def __init__(self) -> None: + self._param_grid: "ParamMap" = {} + + @since("1.4.0") + def addGrid(self, param: Param[Any], values: List[Any]) -> "ParamGridBuilder": + """ + Sets the given parameters in this grid to fixed values. + + param must be an instance of Param associated with an instance of Params + (such as Estimator or Transformer). + """ + if isinstance(param, Param): + self._param_grid[param] = values + else: + raise TypeError("param must be an instance of Param") + + return self + + @overload + def baseOn(self, __args: "ParamMap") -> "ParamGridBuilder": + ... + + @overload + def baseOn(self, *args: Tuple[Param, Any]) -> "ParamGridBuilder": + ... + + @since("1.4.0") + def baseOn(self, *args: Union["ParamMap", Tuple[Param, Any]]) -> "ParamGridBuilder": + """ + Sets the given parameters in this grid to fixed values. + Accepts either a parameter dictionary or a list of (parameter, value) pairs. + """ + if isinstance(args[0], dict): + self.baseOn(*args[0].items()) + else: + for (param, value) in args: + self.addGrid(param, [value]) + + return self + + @since("1.4.0") + def build(self) -> List["ParamMap"]: + """ + Builds and returns all combinations of parameters specified + by the param grid. + """ + keys = self._param_grid.keys() + grid_values = self._param_grid.values() + + def to_key_value_pairs( + keys: Iterable[Param], values: Iterable[Any] + ) -> Sequence[Tuple[Param, Any]]: + return [(key, key.typeConverter(value)) for key, value in zip(keys, values)] + + return [dict(to_key_value_pairs(keys, prod)) for prod in itertools.product(*grid_values)] + + +class _ValidatorParams(HasSeed): + """ + Common params for TrainValidationSplit and CrossValidator. + """ + + estimator: Param[Estimator] = Param( + Params._dummy(), "estimator", "estimator to be cross-validated" + ) + estimatorParamMaps: Param[List["ParamMap"]] = Param( + Params._dummy(), "estimatorParamMaps", "estimator param maps" + ) + evaluator: Param[Evaluator] = Param( + Params._dummy(), + "evaluator", + "evaluator used to select hyper-parameters that maximize the validator metric", + ) + + @since("2.0.0") + def getEstimator(self) -> Estimator: + """ + Gets the value of estimator or its default value. + """ + return self.getOrDefault(self.estimator) + + @since("2.0.0") + def getEstimatorParamMaps(self) -> List["ParamMap"]: + """ + Gets the value of estimatorParamMaps or its default value. + """ + return self.getOrDefault(self.estimatorParamMaps) + + @since("2.0.0") + def getEvaluator(self) -> Evaluator: + """ + Gets the value of evaluator or its default value. + """ + return self.getOrDefault(self.evaluator) + + @classmethod + def _from_java_impl( + cls, java_stage: "JavaObject" + ) -> Tuple[Estimator, List["ParamMap"], Evaluator]: + """ + Return Python estimator, estimatorParamMaps, and evaluator from a Java ValidatorParams. + """ + + # Load information from java_stage to the instance. + estimator: Estimator = JavaParams._from_java(java_stage.getEstimator()) + evaluator: Evaluator = JavaParams._from_java(java_stage.getEvaluator()) + if isinstance(estimator, JavaEstimator): + epms = [ + estimator._transfer_param_map_from_java(epm) + for epm in java_stage.getEstimatorParamMaps() + ] + elif MetaAlgorithmReadWrite.isMetaEstimator(estimator): + # Meta estimator such as Pipeline, OneVsRest + epms = _ValidatorSharedReadWrite.meta_estimator_transfer_param_maps_from_java( + estimator, java_stage.getEstimatorParamMaps() + ) + else: + raise ValueError("Unsupported estimator used in tuning: " + str(estimator)) + + return estimator, epms, evaluator + + def _to_java_impl(self) -> Tuple["JavaObject", "JavaObject", "JavaObject"]: + """ + Return Java estimator, estimatorParamMaps, and evaluator from this Python instance. + """ + + gateway = SparkContext._gateway + assert gateway is not None and SparkContext._jvm is not None + + cls = SparkContext._jvm.org.apache.spark.ml.param.ParamMap + + estimator = self.getEstimator() + if isinstance(estimator, JavaEstimator): + java_epms = gateway.new_array(cls, len(self.getEstimatorParamMaps())) + for idx, epm in enumerate(self.getEstimatorParamMaps()): + java_epms[idx] = estimator._transfer_param_map_to_java(epm) + elif MetaAlgorithmReadWrite.isMetaEstimator(estimator): + # Meta estimator such as Pipeline, OneVsRest + java_epms = _ValidatorSharedReadWrite.meta_estimator_transfer_param_maps_to_java( + estimator, self.getEstimatorParamMaps() + ) + else: + raise ValueError("Unsupported estimator used in tuning: " + str(estimator)) + + java_estimator = cast(JavaEstimator, self.getEstimator())._to_java() + java_evaluator = cast(JavaEvaluator, self.getEvaluator())._to_java() + return java_estimator, java_epms, java_evaluator + + +class _ValidatorSharedReadWrite: + @staticmethod + def meta_estimator_transfer_param_maps_to_java( + pyEstimator: Estimator, pyParamMaps: Sequence["ParamMap"] + ) -> "JavaArray": + pyStages = MetaAlgorithmReadWrite.getAllNestedStages(pyEstimator) + stagePairs = list(map(lambda stage: (stage, cast(JavaParams, stage)._to_java()), pyStages)) + sc = SparkContext._active_spark_context + + assert ( + sc is not None and SparkContext._jvm is not None and SparkContext._gateway is not None + ) + + paramMapCls = SparkContext._jvm.org.apache.spark.ml.param.ParamMap + javaParamMaps = SparkContext._gateway.new_array(paramMapCls, len(pyParamMaps)) + + for idx, pyParamMap in enumerate(pyParamMaps): + javaParamMap = JavaWrapper._new_java_obj("org.apache.spark.ml.param.ParamMap") + for pyParam, pyValue in pyParamMap.items(): + javaParam = None + for pyStage, javaStage in stagePairs: + if pyStage._testOwnParam(pyParam.parent, pyParam.name): + javaParam = javaStage.getParam(pyParam.name) + break + if javaParam is None: + raise ValueError("Resolve param in estimatorParamMaps failed: " + str(pyParam)) + if isinstance(pyValue, Params) and hasattr(pyValue, "_to_java"): + javaValue = cast(JavaParams, pyValue)._to_java() + else: + javaValue = _py2java(sc, pyValue) + pair = javaParam.w(javaValue) + javaParamMap.put([pair]) + javaParamMaps[idx] = javaParamMap + return javaParamMaps + + @staticmethod + def meta_estimator_transfer_param_maps_from_java( + pyEstimator: Estimator, javaParamMaps: "JavaArray" + ) -> List["ParamMap"]: + pyStages = MetaAlgorithmReadWrite.getAllNestedStages(pyEstimator) + stagePairs = list(map(lambda stage: (stage, cast(JavaParams, stage)._to_java()), pyStages)) + sc = SparkContext._active_spark_context + + assert sc is not None and sc._jvm is not None + + pyParamMaps = [] + for javaParamMap in javaParamMaps: + pyParamMap = dict() + for javaPair in javaParamMap.toList(): + javaParam = javaPair.param() + pyParam = None + for pyStage, javaStage in stagePairs: + if pyStage._testOwnParam(javaParam.parent(), javaParam.name()): + pyParam = pyStage.getParam(javaParam.name()) + if pyParam is None: + raise ValueError( + "Resolve param in estimatorParamMaps failed: " + + javaParam.parent() + + "." + + javaParam.name() + ) + javaValue = javaPair.value() + pyValue: Any + if sc._jvm.Class.forName( + "org.apache.spark.ml.util.DefaultParamsWritable" + ).isInstance(javaValue): + pyValue = JavaParams._from_java(javaValue) + else: + pyValue = _java2py(sc, javaValue) + pyParamMap[pyParam] = pyValue + pyParamMaps.append(pyParamMap) + return pyParamMaps + + @staticmethod + def is_java_convertible(instance: _ValidatorParams) -> bool: + allNestedStages = MetaAlgorithmReadWrite.getAllNestedStages(instance.getEstimator()) + evaluator_convertible = isinstance(instance.getEvaluator(), JavaParams) + estimator_convertible = all(map(lambda stage: hasattr(stage, "_to_java"), allNestedStages)) + return estimator_convertible and evaluator_convertible + + @staticmethod + def saveImpl( + path: str, + instance: _ValidatorParams, + sc: SparkContext, + extraMetadata: Optional[Dict[str, Any]] = None, + ) -> None: + numParamsNotJson = 0 + jsonEstimatorParamMaps = [] + for paramMap in instance.getEstimatorParamMaps(): + jsonParamMap = [] + for p, v in paramMap.items(): + jsonParam: Dict[str, Any] = {"parent": p.parent, "name": p.name} + if ( + (isinstance(v, Estimator) and not MetaAlgorithmReadWrite.isMetaEstimator(v)) + or isinstance(v, Transformer) + or isinstance(v, Evaluator) + ): + relative_path = f"epm_{p.name}{numParamsNotJson}" + param_path = os.path.join(path, relative_path) + numParamsNotJson += 1 + cast(MLWritable, v).save(param_path) + jsonParam["value"] = relative_path + jsonParam["isJson"] = False + elif isinstance(v, MLWritable): + raise RuntimeError( + "ValidatorSharedReadWrite.saveImpl does not handle parameters of type: " + "MLWritable that are not Estimator/Evaluator/Transformer, and if parameter " + "is estimator, it cannot be meta estimator such as Validator or OneVsRest" + ) + else: + jsonParam["value"] = v + jsonParam["isJson"] = True + jsonParamMap.append(jsonParam) + jsonEstimatorParamMaps.append(jsonParamMap) + + skipParams = ["estimator", "evaluator", "estimatorParamMaps"] + jsonParams = DefaultParamsWriter.extractJsonParams(instance, skipParams) + jsonParams["estimatorParamMaps"] = jsonEstimatorParamMaps + + DefaultParamsWriter.saveMetadata(instance, path, sc, extraMetadata, jsonParams) + evaluatorPath = os.path.join(path, "evaluator") + cast(MLWritable, instance.getEvaluator()).save(evaluatorPath) + estimatorPath = os.path.join(path, "estimator") + cast(MLWritable, instance.getEstimator()).save(estimatorPath) + + @staticmethod + def load( + path: str, sc: SparkContext, metadata: Dict[str, Any] + ) -> Tuple[Dict[str, Any], Estimator, Evaluator, List["ParamMap"]]: + evaluatorPath = os.path.join(path, "evaluator") + evaluator: Evaluator = DefaultParamsReader.loadParamsInstance(evaluatorPath, sc) + estimatorPath = os.path.join(path, "estimator") + estimator: Estimator = DefaultParamsReader.loadParamsInstance(estimatorPath, sc) + + uidToParams = MetaAlgorithmReadWrite.getUidMap(estimator) + uidToParams[evaluator.uid] = evaluator + + jsonEstimatorParamMaps = metadata["paramMap"]["estimatorParamMaps"] + + estimatorParamMaps = [] + for jsonParamMap in jsonEstimatorParamMaps: + paramMap = {} + for jsonParam in jsonParamMap: + est = uidToParams[jsonParam["parent"]] + param = getattr(est, jsonParam["name"]) + if "isJson" not in jsonParam or ("isJson" in jsonParam and jsonParam["isJson"]): + value = jsonParam["value"] + else: + relativePath = jsonParam["value"] + valueSavedPath = os.path.join(path, relativePath) + value = DefaultParamsReader.loadParamsInstance(valueSavedPath, sc) + paramMap[param] = value + estimatorParamMaps.append(paramMap) + + return metadata, estimator, evaluator, estimatorParamMaps + + @staticmethod + def validateParams(instance: _ValidatorParams) -> None: + estiamtor = instance.getEstimator() + evaluator = instance.getEvaluator() + uidMap = MetaAlgorithmReadWrite.getUidMap(estiamtor) + + for elem in [evaluator] + list(uidMap.values()): + if not isinstance(elem, MLWritable): + raise ValueError( + f"Validator write will fail because it contains {elem.uid} " + f"which is not writable." + ) + + estimatorParamMaps = instance.getEstimatorParamMaps() + paramErr = ( + "Validator save requires all Params in estimatorParamMaps to apply to " + "its Estimator, An extraneous Param was found: " + ) + for paramMap in estimatorParamMaps: + for param in paramMap: + if param.parent not in uidMap: + raise ValueError(paramErr + repr(param)) + + @staticmethod + def getValidatorModelWriterPersistSubModelsParam(writer: MLWriter) -> bool: + if "persistsubmodels" in writer.optionMap: + persistSubModelsParam = writer.optionMap["persistsubmodels"].lower() + if persistSubModelsParam == "true": + return True + elif persistSubModelsParam == "false": + return False + else: + raise ValueError( + f"persistSubModels option value {persistSubModelsParam} is invalid, " + f"the possible values are True, 'True' or False, 'False'" + ) + else: + return writer.instance.subModels is not None # type: ignore[attr-defined] + + +_save_with_persist_submodels_no_submodels_found_err: str = ( + "When persisting tuning models, you can only set persistSubModels to true if the tuning " + "was done with collectSubModels set to true. To save the sub-models, try rerunning fitting " + "with collectSubModels set to true." +) + + +@inherit_doc +class CrossValidatorReader(MLReader["CrossValidator"]): + def __init__(self, cls: Type["CrossValidator"]): + super(CrossValidatorReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> "CrossValidator": + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load( + path, self.sc, metadata + ) + cv = CrossValidator( + estimator=estimator, estimatorParamMaps=estimatorParamMaps, evaluator=evaluator + ) + cv = cv._resetUid(metadata["uid"]) + DefaultParamsReader.getAndSetParams(cv, metadata, skipParams=["estimatorParamMaps"]) + return cv + + +@inherit_doc +class CrossValidatorWriter(MLWriter): + def __init__(self, instance: "CrossValidator"): + super(CrossValidatorWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _ValidatorSharedReadWrite.validateParams(self.instance) + _ValidatorSharedReadWrite.saveImpl(path, self.instance, self.sc) + + +@inherit_doc +class CrossValidatorModelReader(MLReader["CrossValidatorModel"]): + def __init__(self, cls: Type["CrossValidatorModel"]): + super(CrossValidatorModelReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> "CrossValidatorModel": + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load( + path, self.sc, metadata + ) + numFolds = metadata["paramMap"]["numFolds"] + bestModelPath = os.path.join(path, "bestModel") + bestModel: Model = DefaultParamsReader.loadParamsInstance(bestModelPath, self.sc) + avgMetrics = metadata["avgMetrics"] + if "stdMetrics" in metadata: + stdMetrics = metadata["stdMetrics"] + else: + stdMetrics = None + persistSubModels = ("persistSubModels" in metadata) and metadata["persistSubModels"] + + if persistSubModels: + subModels = [[None] * len(estimatorParamMaps)] * numFolds + for splitIndex in range(numFolds): + for paramIndex in range(len(estimatorParamMaps)): + modelPath = os.path.join( + path, "subModels", f"fold{splitIndex}", f"{paramIndex}" + ) + subModels[splitIndex][paramIndex] = DefaultParamsReader.loadParamsInstance( + modelPath, self.sc + ) + else: + subModels = None + + cvModel = CrossValidatorModel( + bestModel, + avgMetrics=avgMetrics, + subModels=cast(List[List[Model]], subModels), + stdMetrics=stdMetrics, + ) + cvModel = cvModel._resetUid(metadata["uid"]) + cvModel.set(cvModel.estimator, estimator) + cvModel.set(cvModel.estimatorParamMaps, estimatorParamMaps) + cvModel.set(cvModel.evaluator, evaluator) + DefaultParamsReader.getAndSetParams( + cvModel, metadata, skipParams=["estimatorParamMaps"] + ) + return cvModel + + +@inherit_doc +class CrossValidatorModelWriter(MLWriter): + def __init__(self, instance: "CrossValidatorModel"): + super(CrossValidatorModelWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _ValidatorSharedReadWrite.validateParams(self.instance) + instance = self.instance + persistSubModels = _ValidatorSharedReadWrite.getValidatorModelWriterPersistSubModelsParam( + self + ) + extraMetadata = {"avgMetrics": instance.avgMetrics, "persistSubModels": persistSubModels} + if instance.stdMetrics: + extraMetadata["stdMetrics"] = instance.stdMetrics + + _ValidatorSharedReadWrite.saveImpl(path, instance, self.sc, extraMetadata=extraMetadata) + bestModelPath = os.path.join(path, "bestModel") + cast(MLWritable, instance.bestModel).save(bestModelPath) + if persistSubModels: + if instance.subModels is None: + raise ValueError(_save_with_persist_submodels_no_submodels_found_err) + subModelsPath = os.path.join(path, "subModels") + for splitIndex in range(instance.getNumFolds()): + splitPath = os.path.join(subModelsPath, f"fold{splitIndex}") + for paramIndex in range(len(instance.getEstimatorParamMaps())): + modelPath = os.path.join(splitPath, f"{paramIndex}") + cast(MLWritable, instance.subModels[splitIndex][paramIndex]).save(modelPath) + + +class _CrossValidatorParams(_ValidatorParams): + """ + Params for :py:class:`CrossValidator` and :py:class:`CrossValidatorModel`. + + .. versionadded:: 3.0.0 + """ + + numFolds: Param[int] = Param( + Params._dummy(), + "numFolds", + "number of folds for cross validation", + typeConverter=TypeConverters.toInt, + ) + + foldCol: Param[str] = Param( + Params._dummy(), + "foldCol", + "Param for the column name of user " + + "specified fold number. Once this is specified, :py:class:`CrossValidator` " + + "won't do random k-fold split. Note that this column should be integer type " + + "with range [0, numFolds) and Spark will throw exception on out-of-range " + + "fold numbers.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self, *args: Any): + super(_CrossValidatorParams, self).__init__(*args) + self._setDefault(numFolds=3, foldCol="") + + @since("1.4.0") + def getNumFolds(self) -> int: + """ + Gets the value of numFolds or its default value. + """ + return self.getOrDefault(self.numFolds) + + @since("3.1.0") + def getFoldCol(self) -> str: + """ + Gets the value of foldCol or its default value. + """ + return self.getOrDefault(self.foldCol) + + +class CrossValidator( + Estimator["CrossValidatorModel"], + _CrossValidatorParams, + HasParallelism, + HasCollectSubModels, + MLReadable["CrossValidator"], + MLWritable, +): + """ + + K-fold cross validation performs model selection by splitting the dataset into a set of + non-overlapping randomly partitioned folds which are used as separate training and test datasets + e.g., with k=3 folds, K-fold cross validation will generate 3 (training, test) dataset pairs, + each of which uses 2/3 of the data for training and 1/3 for testing. Each fold is used as the + test set exactly once. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.ml.classification import LogisticRegression + >>> from pyspark.ml.evaluation import BinaryClassificationEvaluator + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.tuning import CrossValidator, ParamGridBuilder, CrossValidatorModel + >>> import tempfile + >>> dataset = spark.createDataFrame( + ... [(Vectors.dense([0.0]), 0.0), + ... (Vectors.dense([0.4]), 1.0), + ... (Vectors.dense([0.5]), 0.0), + ... (Vectors.dense([0.6]), 1.0), + ... (Vectors.dense([1.0]), 1.0)] * 10, + ... ["features", "label"]) + >>> lr = LogisticRegression() + >>> grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + >>> evaluator = BinaryClassificationEvaluator() + >>> cv = CrossValidator(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, + ... parallelism=2) + >>> cvModel = cv.fit(dataset) + >>> cvModel.getNumFolds() + 3 + >>> cvModel.avgMetrics[0] + 0.5 + >>> path = tempfile.mkdtemp() + >>> model_path = path + "/model" + >>> cvModel.write().save(model_path) + >>> cvModelRead = CrossValidatorModel.read().load(model_path) + >>> cvModelRead.avgMetrics + [0.5, ... + >>> evaluator.evaluate(cvModel.transform(dataset)) + 0.8333... + >>> evaluator.evaluate(cvModelRead.transform(dataset)) + 0.8333... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + numFolds: int = 3, + seed: Optional[int] = None, + parallelism: int = 1, + collectSubModels: bool = False, + foldCol: str = "", + ) -> None: + """ + __init__(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, numFolds=3,\ + seed=None, parallelism=1, collectSubModels=False, foldCol="") + """ + super(CrossValidator, self).__init__() + self._setDefault(parallelism=1) + kwargs = self._input_kwargs + self._set(**kwargs) + + @keyword_only + @since("1.4.0") + def setParams( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + numFolds: int = 3, + seed: Optional[int] = None, + parallelism: int = 1, + collectSubModels: bool = False, + foldCol: str = "", + ) -> "CrossValidator": + """ + setParams(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, numFolds=3,\ + seed=None, parallelism=1, collectSubModels=False, foldCol=""): + Sets params for cross validator. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setEstimator(self, value: Estimator) -> "CrossValidator": + """ + Sets the value of :py:attr:`estimator`. + """ + return self._set(estimator=value) + + @since("2.0.0") + def setEstimatorParamMaps(self, value: List["ParamMap"]) -> "CrossValidator": + """ + Sets the value of :py:attr:`estimatorParamMaps`. + """ + return self._set(estimatorParamMaps=value) + + @since("2.0.0") + def setEvaluator(self, value: Evaluator) -> "CrossValidator": + """ + Sets the value of :py:attr:`evaluator`. + """ + return self._set(evaluator=value) + + @since("1.4.0") + def setNumFolds(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`numFolds`. + """ + return self._set(numFolds=value) + + @since("3.1.0") + def setFoldCol(self, value: str) -> "CrossValidator": + """ + Sets the value of :py:attr:`foldCol`. + """ + return self._set(foldCol=value) + + def setSeed(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def setParallelism(self, value: int) -> "CrossValidator": + """ + Sets the value of :py:attr:`parallelism`. + """ + return self._set(parallelism=value) + + def setCollectSubModels(self, value: bool) -> "CrossValidator": + """ + Sets the value of :py:attr:`collectSubModels`. + """ + return self._set(collectSubModels=value) + + @staticmethod + def _gen_avg_and_std_metrics(metrics_all: List[List[float]]) -> Tuple[List[float], List[float]]: + avg_metrics = np.mean(metrics_all, axis=0) + std_metrics = np.std(metrics_all, axis=0) + return list(avg_metrics), list(std_metrics) + + def _fit(self, dataset: DataFrame) -> "CrossValidatorModel": + est = self.getOrDefault(self.estimator) + epm = self.getOrDefault(self.estimatorParamMaps) + numModels = len(epm) + eva = self.getOrDefault(self.evaluator) + nFolds = self.getOrDefault(self.numFolds) + metrics_all = [[0.0] * numModels for i in range(nFolds)] + + pool = ThreadPool(processes=min(self.getParallelism(), numModels)) + subModels = None + collectSubModelsParam = self.getCollectSubModels() + if collectSubModelsParam: + subModels = [[None for j in range(numModels)] for i in range(nFolds)] + + datasets = self._kFold(dataset) + for i in range(nFolds): + validation = datasets[i][1].cache() + train = datasets[i][0].cache() + + tasks = map( + inheritable_thread_target, + _parallelFitTasks(est, train, eva, validation, epm, collectSubModelsParam), + ) + for j, metric, subModel in pool.imap_unordered(lambda f: f(), tasks): + metrics_all[i][j] = metric + if collectSubModelsParam: + assert subModels is not None + subModels[i][j] = subModel + + validation.unpersist() + train.unpersist() + + metrics, std_metrics = CrossValidator._gen_avg_and_std_metrics(metrics_all) + + if eva.isLargerBetter(): + bestIndex = np.argmax(metrics) + else: + bestIndex = np.argmin(metrics) + bestModel = est.fit(dataset, epm[bestIndex]) + return self._copyValues( + CrossValidatorModel(bestModel, metrics, cast(List[List[Model]], subModels), std_metrics) + ) + + def _kFold(self, dataset: DataFrame) -> List[Tuple[DataFrame, DataFrame]]: + nFolds = self.getOrDefault(self.numFolds) + foldCol = self.getOrDefault(self.foldCol) + + datasets = [] + if not foldCol: + # Do random k-fold split. + seed = self.getOrDefault(self.seed) + h = 1.0 / nFolds + randCol = self.uid + "_rand" + df = dataset.select("*", rand(seed).alias(randCol)) + for i in range(nFolds): + validateLB = i * h + validateUB = (i + 1) * h + condition = (df[randCol] >= validateLB) & (df[randCol] < validateUB) + validation = df.filter(condition) + train = df.filter(~condition) + datasets.append((train, validation)) + else: + # Use user-specified fold numbers. + def checker(foldNum: int) -> bool: + if foldNum < 0 or foldNum >= nFolds: + raise ValueError( + "Fold number must be in range [0, %s), but got %s." % (nFolds, foldNum) + ) + return True + + checker_udf = UserDefinedFunction(checker, BooleanType()) + for i in range(nFolds): + training = dataset.filter(checker_udf(dataset[foldCol]) & (col(foldCol) != lit(i))) + validation = dataset.filter( + checker_udf(dataset[foldCol]) & (col(foldCol) == lit(i)) + ) + if training.rdd.getNumPartitions() == 0 or len(training.take(1)) == 0: + raise ValueError("The training data at fold %s is empty." % i) + if validation.rdd.getNumPartitions() == 0 or len(validation.take(1)) == 0: + raise ValueError("The validation data at fold %s is empty." % i) + datasets.append((training, validation)) + + return datasets + + def copy(self, extra: Optional["ParamMap"] = None) -> "CrossValidator": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies creates a deep copy of + the embedded paramMap, and copies the embedded and extra parameters over. + + + .. versionadded:: 1.4.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`CrossValidator` + Copy of this instance + """ + if extra is None: + extra = dict() + newCV = Params.copy(self, extra) + if self.isSet(self.estimator): + newCV.setEstimator(self.getEstimator().copy(extra)) + # estimatorParamMaps remain the same + if self.isSet(self.evaluator): + newCV.setEvaluator(self.getEvaluator().copy(extra)) + return newCV + + @since("2.3.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + if _ValidatorSharedReadWrite.is_java_convertible(self): + return JavaMLWriter(self) # type: ignore[arg-type] + return CrossValidatorWriter(self) + + @classmethod + @since("2.3.0") + def read(cls) -> CrossValidatorReader: + """Returns an MLReader instance for this class.""" + return CrossValidatorReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "CrossValidator": + """ + Given a Java CrossValidator, create and return a Python wrapper of it. + Used for ML persistence. + """ + + estimator, epms, evaluator = super(CrossValidator, cls)._from_java_impl(java_stage) + numFolds = java_stage.getNumFolds() + seed = java_stage.getSeed() + parallelism = java_stage.getParallelism() + collectSubModels = java_stage.getCollectSubModels() + foldCol = java_stage.getFoldCol() + # Create a new instance of this stage. + py_stage = cls( + estimator=estimator, + estimatorParamMaps=epms, + evaluator=evaluator, + numFolds=numFolds, + seed=seed, + parallelism=parallelism, + collectSubModels=collectSubModels, + foldCol=foldCol, + ) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java CrossValidator. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + + estimator, epms, evaluator = super(CrossValidator, self)._to_java_impl() + + _java_obj = JavaParams._new_java_obj("org.apache.spark.ml.tuning.CrossValidator", self.uid) + _java_obj.setEstimatorParamMaps(epms) + _java_obj.setEvaluator(evaluator) + _java_obj.setEstimator(estimator) + _java_obj.setSeed(self.getSeed()) + _java_obj.setNumFolds(self.getNumFolds()) + _java_obj.setParallelism(self.getParallelism()) + _java_obj.setCollectSubModels(self.getCollectSubModels()) + _java_obj.setFoldCol(self.getFoldCol()) + + return _java_obj + + +class CrossValidatorModel( + Model, _CrossValidatorParams, MLReadable["CrossValidatorModel"], MLWritable +): + """ + CrossValidatorModel contains the model with the highest average cross-validation + metric across folds and uses this model to transform input data. CrossValidatorModel + also tracks the metrics for each param map evaluated. + + .. versionadded:: 1.4.0 + + Notes + ----- + Since version 3.3.0, CrossValidatorModel contains a new attribute "stdMetrics", + which represent standard deviation of metrics for each paramMap in + CrossValidator.estimatorParamMaps. + """ + + def __init__( + self, + bestModel: Model, + avgMetrics: Optional[List[float]] = None, + subModels: Optional[List[List[Model]]] = None, + stdMetrics: Optional[List[float]] = None, + ): + super(CrossValidatorModel, self).__init__() + #: best model from cross validation + self.bestModel = bestModel + #: Average cross-validation metrics for each paramMap in + #: CrossValidator.estimatorParamMaps, in the corresponding order. + self.avgMetrics = avgMetrics or [] + #: sub model list from cross validation + self.subModels = subModels + #: standard deviation of metrics for each paramMap in + #: CrossValidator.estimatorParamMaps, in the corresponding order. + self.stdMetrics = stdMetrics or [] + + def _transform(self, dataset: DataFrame) -> DataFrame: + return self.bestModel.transform(dataset) + + def copy(self, extra: Optional["ParamMap"] = None) -> "CrossValidatorModel": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies the underlying bestModel, + creates a deep copy of the embedded paramMap, and + copies the embedded and extra parameters over. + It does not copy the extra Params into the subModels. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`CrossValidatorModel` + Copy of this instance + """ + if extra is None: + extra = dict() + bestModel = self.bestModel.copy(extra) + avgMetrics = list(self.avgMetrics) + assert self.subModels is not None + subModels = [ + [sub_model.copy() for sub_model in fold_sub_models] + for fold_sub_models in self.subModels + ] + stdMetrics = list(self.stdMetrics) + return self._copyValues( + CrossValidatorModel(bestModel, avgMetrics, subModels, stdMetrics), extra=extra + ) + + @since("2.3.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + if _ValidatorSharedReadWrite.is_java_convertible(self): + return JavaMLWriter(self) # type: ignore[arg-type] + return CrossValidatorModelWriter(self) + + @classmethod + @since("2.3.0") + def read(cls) -> CrossValidatorModelReader: + """Returns an MLReader instance for this class.""" + return CrossValidatorModelReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "CrossValidatorModel": + """ + Given a Java CrossValidatorModel, create and return a Python wrapper of it. + Used for ML persistence. + """ + sc = SparkContext._active_spark_context + assert sc is not None + + bestModel: Model = JavaParams._from_java(java_stage.bestModel()) + avgMetrics = _java2py(sc, java_stage.avgMetrics()) + estimator, epms, evaluator = super(CrossValidatorModel, cls)._from_java_impl(java_stage) + + py_stage = cls(bestModel=bestModel, avgMetrics=avgMetrics) + params = { + "evaluator": evaluator, + "estimator": estimator, + "estimatorParamMaps": epms, + "numFolds": java_stage.getNumFolds(), + "foldCol": java_stage.getFoldCol(), + "seed": java_stage.getSeed(), + } + for param_name, param_val in params.items(): + py_stage = py_stage._set(**{param_name: param_val}) + + if java_stage.hasSubModels(): + py_stage.subModels = [ + [JavaParams._from_java(sub_model) for sub_model in fold_sub_models] + for fold_sub_models in java_stage.subModels() + ] + + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java CrossValidatorModel. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + + sc = SparkContext._active_spark_context + assert sc is not None + + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.tuning.CrossValidatorModel", + self.uid, + cast(JavaParams, self.bestModel)._to_java(), + _py2java(sc, self.avgMetrics), + ) + estimator, epms, evaluator = super(CrossValidatorModel, self)._to_java_impl() + + params = { + "evaluator": evaluator, + "estimator": estimator, + "estimatorParamMaps": epms, + "numFolds": self.getNumFolds(), + "foldCol": self.getFoldCol(), + "seed": self.getSeed(), + } + for param_name, param_val in params.items(): + java_param = _java_obj.getParam(param_name) + pair = java_param.w(param_val) + _java_obj.set(pair) + + if self.subModels is not None: + java_sub_models = [ + [cast(JavaParams, sub_model)._to_java() for sub_model in fold_sub_models] + for fold_sub_models in self.subModels + ] + _java_obj.setSubModels(java_sub_models) + return _java_obj + + +@inherit_doc +class TrainValidationSplitReader(MLReader["TrainValidationSplit"]): + def __init__(self, cls: Type["TrainValidationSplit"]): + super(TrainValidationSplitReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> "TrainValidationSplit": + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load( + path, self.sc, metadata + ) + tvs = TrainValidationSplit( + estimator=estimator, estimatorParamMaps=estimatorParamMaps, evaluator=evaluator + ) + tvs = tvs._resetUid(metadata["uid"]) + DefaultParamsReader.getAndSetParams(tvs, metadata, skipParams=["estimatorParamMaps"]) + return tvs + + +@inherit_doc +class TrainValidationSplitWriter(MLWriter): + def __init__(self, instance: "TrainValidationSplit"): + super(TrainValidationSplitWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _ValidatorSharedReadWrite.validateParams(self.instance) + _ValidatorSharedReadWrite.saveImpl(path, self.instance, self.sc) + + +@inherit_doc +class TrainValidationSplitModelReader(MLReader["TrainValidationSplitModel"]): + def __init__(self, cls: Type["TrainValidationSplitModel"]): + super(TrainValidationSplitModelReader, self).__init__() + self.cls = cls + + def load(self, path: str) -> "TrainValidationSplitModel": + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + if not DefaultParamsReader.isPythonParamsInstance(metadata): + return JavaMLReader(self.cls).load(path) # type: ignore[arg-type] + else: + metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load( + path, self.sc, metadata + ) + bestModelPath = os.path.join(path, "bestModel") + bestModel: Model = DefaultParamsReader.loadParamsInstance(bestModelPath, self.sc) + validationMetrics = metadata["validationMetrics"] + persistSubModels = ("persistSubModels" in metadata) and metadata["persistSubModels"] + + if persistSubModels: + subModels = [None] * len(estimatorParamMaps) + for paramIndex in range(len(estimatorParamMaps)): + modelPath = os.path.join(path, "subModels", f"{paramIndex}") + subModels[paramIndex] = DefaultParamsReader.loadParamsInstance( + modelPath, self.sc + ) + else: + subModels = None + + tvsModel = TrainValidationSplitModel( + bestModel, + validationMetrics=validationMetrics, + subModels=cast(Optional[List[Model]], subModels), + ) + tvsModel = tvsModel._resetUid(metadata["uid"]) + tvsModel.set(tvsModel.estimator, estimator) + tvsModel.set(tvsModel.estimatorParamMaps, estimatorParamMaps) + tvsModel.set(tvsModel.evaluator, evaluator) + DefaultParamsReader.getAndSetParams( + tvsModel, metadata, skipParams=["estimatorParamMaps"] + ) + return tvsModel + + +@inherit_doc +class TrainValidationSplitModelWriter(MLWriter): + def __init__(self, instance: "TrainValidationSplitModel"): + super(TrainValidationSplitModelWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + _ValidatorSharedReadWrite.validateParams(self.instance) + instance = self.instance + persistSubModels = _ValidatorSharedReadWrite.getValidatorModelWriterPersistSubModelsParam( + self + ) + + extraMetadata = { + "validationMetrics": instance.validationMetrics, + "persistSubModels": persistSubModels, + } + _ValidatorSharedReadWrite.saveImpl(path, instance, self.sc, extraMetadata=extraMetadata) + bestModelPath = os.path.join(path, "bestModel") + cast(MLWritable, instance.bestModel).save(bestModelPath) + if persistSubModels: + if instance.subModels is None: + raise ValueError(_save_with_persist_submodels_no_submodels_found_err) + subModelsPath = os.path.join(path, "subModels") + for paramIndex in range(len(instance.getEstimatorParamMaps())): + modelPath = os.path.join(subModelsPath, f"{paramIndex}") + cast(MLWritable, instance.subModels[paramIndex]).save(modelPath) + + +class _TrainValidationSplitParams(_ValidatorParams): + """ + Params for :py:class:`TrainValidationSplit` and :py:class:`TrainValidationSplitModel`. + + .. versionadded:: 3.0.0 + """ + + trainRatio: Param[float] = Param( + Params._dummy(), + "trainRatio", + "Param for ratio between train and\ + validation data. Must be between 0 and 1.", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, *args: Any): + super(_TrainValidationSplitParams, self).__init__(*args) + self._setDefault(trainRatio=0.75) + + @since("2.0.0") + def getTrainRatio(self) -> float: + """ + Gets the value of trainRatio or its default value. + """ + return self.getOrDefault(self.trainRatio) + + +class TrainValidationSplit( + Estimator["TrainValidationSplitModel"], + _TrainValidationSplitParams, + HasParallelism, + HasCollectSubModels, + MLReadable["TrainValidationSplit"], + MLWritable, +): + """ + Validation for hyper-parameter tuning. Randomly splits the input dataset into train and + validation sets, and uses evaluation metric on the validation set to select the best model. + Similar to :class:`CrossValidator`, but only splits the set once. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> from pyspark.ml.classification import LogisticRegression + >>> from pyspark.ml.evaluation import BinaryClassificationEvaluator + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.tuning import TrainValidationSplit, ParamGridBuilder + >>> from pyspark.ml.tuning import TrainValidationSplitModel + >>> import tempfile + >>> dataset = spark.createDataFrame( + ... [(Vectors.dense([0.0]), 0.0), + ... (Vectors.dense([0.4]), 1.0), + ... (Vectors.dense([0.5]), 0.0), + ... (Vectors.dense([0.6]), 1.0), + ... (Vectors.dense([1.0]), 1.0)] * 10, + ... ["features", "label"]).repartition(1) + >>> lr = LogisticRegression() + >>> grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build() + >>> evaluator = BinaryClassificationEvaluator() + >>> tvs = TrainValidationSplit(estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, + ... parallelism=1, seed=42) + >>> tvsModel = tvs.fit(dataset) + >>> tvsModel.getTrainRatio() + 0.75 + >>> tvsModel.validationMetrics + [0.5, ... + >>> path = tempfile.mkdtemp() + >>> model_path = path + "/model" + >>> tvsModel.write().save(model_path) + >>> tvsModelRead = TrainValidationSplitModel.read().load(model_path) + >>> tvsModelRead.validationMetrics + [0.5, ... + >>> evaluator.evaluate(tvsModel.transform(dataset)) + 0.833... + >>> evaluator.evaluate(tvsModelRead.transform(dataset)) + 0.833... + """ + + _input_kwargs: Dict[str, Any] + + @keyword_only + def __init__( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + trainRatio: float = 0.75, + parallelism: int = 1, + collectSubModels: bool = False, + seed: Optional[int] = None, + ) -> None: + """ + __init__(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, \ + trainRatio=0.75, parallelism=1, collectSubModels=False, seed=None) + """ + super(TrainValidationSplit, self).__init__() + self._setDefault(parallelism=1) + kwargs = self._input_kwargs + self._set(**kwargs) + + @since("2.0.0") + @keyword_only + def setParams( + self, + *, + estimator: Optional[Estimator] = None, + estimatorParamMaps: Optional[List["ParamMap"]] = None, + evaluator: Optional[Evaluator] = None, + trainRatio: float = 0.75, + parallelism: int = 1, + collectSubModels: bool = False, + seed: Optional[int] = None, + ) -> "TrainValidationSplit": + """ + setParams(self, \\*, estimator=None, estimatorParamMaps=None, evaluator=None, \ + trainRatio=0.75, parallelism=1, collectSubModels=False, seed=None): + Sets params for the train validation split. + """ + kwargs = self._input_kwargs + return self._set(**kwargs) + + @since("2.0.0") + def setEstimator(self, value: Estimator) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`estimator`. + """ + return self._set(estimator=value) + + @since("2.0.0") + def setEstimatorParamMaps(self, value: List["ParamMap"]) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`estimatorParamMaps`. + """ + return self._set(estimatorParamMaps=value) + + @since("2.0.0") + def setEvaluator(self, value: Evaluator) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`evaluator`. + """ + return self._set(evaluator=value) + + @since("2.0.0") + def setTrainRatio(self, value: float) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`trainRatio`. + """ + return self._set(trainRatio=value) + + def setSeed(self, value: int) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`seed`. + """ + return self._set(seed=value) + + def setParallelism(self, value: int) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`parallelism`. + """ + return self._set(parallelism=value) + + def setCollectSubModels(self, value: bool) -> "TrainValidationSplit": + """ + Sets the value of :py:attr:`collectSubModels`. + """ + return self._set(collectSubModels=value) + + def _fit(self, dataset: DataFrame) -> "TrainValidationSplitModel": + est = self.getOrDefault(self.estimator) + epm = self.getOrDefault(self.estimatorParamMaps) + numModels = len(epm) + eva = self.getOrDefault(self.evaluator) + tRatio = self.getOrDefault(self.trainRatio) + seed = self.getOrDefault(self.seed) + randCol = self.uid + "_rand" + df = dataset.select("*", rand(seed).alias(randCol)) + condition = df[randCol] >= tRatio + validation = df.filter(condition).cache() + train = df.filter(~condition).cache() + + subModels = None + collectSubModelsParam = self.getCollectSubModels() + if collectSubModelsParam: + subModels = [None for i in range(numModels)] + + tasks = map( + inheritable_thread_target, + _parallelFitTasks(est, train, eva, validation, epm, collectSubModelsParam), + ) + pool = ThreadPool(processes=min(self.getParallelism(), numModels)) + metrics = [None] * numModels + for j, metric, subModel in pool.imap_unordered(lambda f: f(), tasks): + metrics[j] = metric + if collectSubModelsParam: + assert subModels is not None + subModels[j] = subModel + + train.unpersist() + validation.unpersist() + + if eva.isLargerBetter(): + bestIndex = np.argmax(cast(List[float], metrics)) + else: + bestIndex = np.argmin(cast(List[float], metrics)) + bestModel = est.fit(dataset, epm[bestIndex]) + return self._copyValues( + TrainValidationSplitModel( + bestModel, + cast(List[float], metrics), + subModels, # type: ignore[arg-type] + ) + ) + + def copy(self, extra: Optional["ParamMap"] = None) -> "TrainValidationSplit": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies creates a deep copy of + the embedded paramMap, and copies the embedded and extra parameters over. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`TrainValidationSplit` + Copy of this instance + """ + if extra is None: + extra = dict() + newTVS = Params.copy(self, extra) + if self.isSet(self.estimator): + newTVS.setEstimator(self.getEstimator().copy(extra)) + # estimatorParamMaps remain the same + if self.isSet(self.evaluator): + newTVS.setEvaluator(self.getEvaluator().copy(extra)) + return newTVS + + @since("2.3.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + if _ValidatorSharedReadWrite.is_java_convertible(self): + return JavaMLWriter(self) # type: ignore[arg-type] + return TrainValidationSplitWriter(self) + + @classmethod + @since("2.3.0") + def read(cls) -> TrainValidationSplitReader: + """Returns an MLReader instance for this class.""" + return TrainValidationSplitReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "TrainValidationSplit": + """ + Given a Java TrainValidationSplit, create and return a Python wrapper of it. + Used for ML persistence. + """ + + estimator, epms, evaluator = super(TrainValidationSplit, cls)._from_java_impl(java_stage) + trainRatio = java_stage.getTrainRatio() + seed = java_stage.getSeed() + parallelism = java_stage.getParallelism() + collectSubModels = java_stage.getCollectSubModels() + # Create a new instance of this stage. + py_stage = cls( + estimator=estimator, + estimatorParamMaps=epms, + evaluator=evaluator, + trainRatio=trainRatio, + seed=seed, + parallelism=parallelism, + collectSubModels=collectSubModels, + ) + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java TrainValidationSplit. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + + estimator, epms, evaluator = super(TrainValidationSplit, self)._to_java_impl() + + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.tuning.TrainValidationSplit", self.uid + ) + _java_obj.setEstimatorParamMaps(epms) + _java_obj.setEvaluator(evaluator) + _java_obj.setEstimator(estimator) + _java_obj.setTrainRatio(self.getTrainRatio()) + _java_obj.setSeed(self.getSeed()) + _java_obj.setParallelism(self.getParallelism()) + _java_obj.setCollectSubModels(self.getCollectSubModels()) + return _java_obj + + +class TrainValidationSplitModel( + Model, _TrainValidationSplitParams, MLReadable["TrainValidationSplitModel"], MLWritable +): + """ + Model from train validation split. + + .. versionadded:: 2.0.0 + """ + + def __init__( + self, + bestModel: Model, + validationMetrics: Optional[List[float]] = None, + subModels: Optional[List[Model]] = None, + ): + super(TrainValidationSplitModel, self).__init__() + #: best model from train validation split + self.bestModel = bestModel + #: evaluated validation metrics + self.validationMetrics = validationMetrics or [] + #: sub models from train validation split + self.subModels = subModels + + def _transform(self, dataset: DataFrame) -> DataFrame: + return self.bestModel.transform(dataset) + + def copy(self, extra: Optional["ParamMap"] = None) -> "TrainValidationSplitModel": + """ + Creates a copy of this instance with a randomly generated uid + and some extra params. This copies the underlying bestModel, + creates a deep copy of the embedded paramMap, and + copies the embedded and extra parameters over. + And, this creates a shallow copy of the validationMetrics. + It does not copy the extra Params into the subModels. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`TrainValidationSplitModel` + Copy of this instance + """ + if extra is None: + extra = dict() + bestModel = self.bestModel.copy(extra) + validationMetrics = list(self.validationMetrics) + assert self.subModels is not None + subModels = [model.copy() for model in self.subModels] + return self._copyValues( + TrainValidationSplitModel(bestModel, validationMetrics, subModels), extra=extra + ) + + @since("2.3.0") + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + if _ValidatorSharedReadWrite.is_java_convertible(self): + return JavaMLWriter(self) # type: ignore[arg-type] + return TrainValidationSplitModelWriter(self) + + @classmethod + @since("2.3.0") + def read(cls) -> TrainValidationSplitModelReader: + """Returns an MLReader instance for this class.""" + return TrainValidationSplitModelReader(cls) + + @classmethod + def _from_java(cls, java_stage: "JavaObject") -> "TrainValidationSplitModel": + """ + Given a Java TrainValidationSplitModel, create and return a Python wrapper of it. + Used for ML persistence. + """ + + # Load information from java_stage to the instance. + sc = SparkContext._active_spark_context + assert sc is not None + + bestModel: Model = JavaParams._from_java(java_stage.bestModel()) + validationMetrics = _java2py(sc, java_stage.validationMetrics()) + estimator, epms, evaluator = super(TrainValidationSplitModel, cls)._from_java_impl( + java_stage + ) + # Create a new instance of this stage. + py_stage = cls(bestModel=bestModel, validationMetrics=validationMetrics) + params = { + "evaluator": evaluator, + "estimator": estimator, + "estimatorParamMaps": epms, + "trainRatio": java_stage.getTrainRatio(), + "seed": java_stage.getSeed(), + } + for param_name, param_val in params.items(): + py_stage = py_stage._set(**{param_name: param_val}) + + if java_stage.hasSubModels(): + py_stage.subModels = [ + JavaParams._from_java(sub_model) for sub_model in java_stage.subModels() + ] + + py_stage._resetUid(java_stage.uid()) + return py_stage + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance to a Java TrainValidationSplitModel. Used for ML persistence. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + + sc = SparkContext._active_spark_context + assert sc is not None + + _java_obj = JavaParams._new_java_obj( + "org.apache.spark.ml.tuning.TrainValidationSplitModel", + self.uid, + cast(JavaParams, self.bestModel)._to_java(), + _py2java(sc, self.validationMetrics), + ) + estimator, epms, evaluator = super(TrainValidationSplitModel, self)._to_java_impl() + + params = { + "evaluator": evaluator, + "estimator": estimator, + "estimatorParamMaps": epms, + "trainRatio": self.getTrainRatio(), + "seed": self.getSeed(), + } + for param_name, param_val in params.items(): + java_param = _java_obj.getParam(param_name) + pair = java_param.w(param_val) + _java_obj.set(pair) + + if self.subModels is not None: + java_sub_models = [ + cast(JavaParams, sub_model)._to_java() for sub_model in self.subModels + ] + _java_obj.setSubModels(java_sub_models) + + return _java_obj + + +if __name__ == "__main__": + import doctest + + from pyspark.sql import SparkSession + + globs = globals().copy() + + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("ml.tuning tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/util.py new file mode 100644 index 0000000..6467694 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/util.py @@ -0,0 +1,749 @@ +# +# 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 json +import os +import time +import uuid +import functools + +from typing import ( + Any, + Callable, + Dict, + Generic, + List, + Optional, + Sequence, + Type, + TypeVar, + cast, + TYPE_CHECKING, +) + + +from pyspark import SparkContext, since +from pyspark.ml.common import inherit_doc +from pyspark.sql import SparkSession +from pyspark.sql.utils import is_remote +from pyspark.util import VersionUtils + +if TYPE_CHECKING: + from py4j.java_gateway import JavaGateway, JavaObject + from pyspark.ml._typing import PipelineStage + from pyspark.ml.base import Params + from pyspark.ml.wrapper import JavaWrapper + +T = TypeVar("T") +RW = TypeVar("RW", bound="BaseReadWrite") +W = TypeVar("W", bound="MLWriter") +JW = TypeVar("JW", bound="JavaMLWriter") +RL = TypeVar("RL", bound="MLReadable") +JR = TypeVar("JR", bound="JavaMLReader") + +FuncT = TypeVar("FuncT", bound=Callable[..., Any]) + + +def _jvm() -> "JavaGateway": + """ + Returns the JVM view associated with SparkContext. Must be called + after SparkContext is initialized. + """ + jvm = SparkContext._jvm + if jvm: + return jvm + else: + raise AttributeError("Cannot load _jvm from SparkContext. Is SparkContext initialized?") + + +class Identifiable: + """ + Object with a unique ID. + """ + + def __init__(self) -> None: + #: A unique id for the object. + self.uid = self._randomUID() + + def __repr__(self) -> str: + return self.uid + + @classmethod + def _randomUID(cls) -> str: + """ + Generate a unique string id for the object. The default implementation + concatenates the class name, "_", and 12 random hex chars. + """ + return str(cls.__name__ + "_" + uuid.uuid4().hex[-12:]) + + +@inherit_doc +class BaseReadWrite: + """ + Base class for MLWriter and MLReader. Stores information about the SparkContext + and SparkSession. + + .. versionadded:: 2.3.0 + """ + + def __init__(self) -> None: + self._sparkSession: Optional[SparkSession] = None + + def session(self: RW, sparkSession: SparkSession) -> RW: + """ + Sets the Spark Session to use for saving/loading. + """ + self._sparkSession = sparkSession + return self + + @property + def sparkSession(self) -> SparkSession: + """ + Returns the user-specified Spark Session or the default. + """ + if self._sparkSession is None: + self._sparkSession = SparkSession._getActiveSessionOrCreate() + assert self._sparkSession is not None + return self._sparkSession + + @property + def sc(self) -> SparkContext: + """ + Returns the underlying `SparkContext`. + """ + assert self.sparkSession is not None + return self.sparkSession.sparkContext + + +@inherit_doc +class MLWriter(BaseReadWrite): + """ + Utility class that can save ML instances. + + .. versionadded:: 2.0.0 + """ + + def __init__(self) -> None: + super(MLWriter, self).__init__() + self.shouldOverwrite: bool = False + self.optionMap: Dict[str, Any] = {} + + def _handleOverwrite(self, path: str) -> None: + from pyspark.ml.wrapper import JavaWrapper + + _java_obj = JavaWrapper._new_java_obj("org.apache.spark.ml.util.FileSystemOverwrite") + wrapper = JavaWrapper(_java_obj) + wrapper._call_java("handleOverwrite", path, True, self.sparkSession._jsparkSession) + + def save(self, path: str) -> None: + """Save the ML instance to the input path.""" + if self.shouldOverwrite: + self._handleOverwrite(path) + self.saveImpl(path) + + def saveImpl(self, path: str) -> None: + """ + save() handles overwriting and then calls this method. Subclasses should override this + method to implement the actual saving of the instance. + """ + raise NotImplementedError("MLWriter is not yet implemented for type: %s" % type(self)) + + def overwrite(self) -> "MLWriter": + """Overwrites if the output path already exists.""" + self.shouldOverwrite = True + return self + + def option(self, key: str, value: Any) -> "MLWriter": + """ + Adds an option to the underlying MLWriter. See the documentation for the specific model's + writer for possible options. The option name (key) is case-insensitive. + """ + self.optionMap[key.lower()] = str(value) + return self + + +@inherit_doc +class GeneralMLWriter(MLWriter): + """ + Utility class that can save ML instances in different formats. + + .. versionadded:: 2.4.0 + """ + + def format(self, source: str) -> "GeneralMLWriter": + """ + Specifies the format of ML export ("pmml", "internal", or the fully qualified class + name for export). + """ + self.source = source + return self + + +@inherit_doc +class JavaMLWriter(MLWriter): + """ + (Private) Specialization of :py:class:`MLWriter` for :py:class:`JavaParams` types + """ + + _jwrite: "JavaObject" + + def __init__(self, instance: "JavaMLWritable"): + super(JavaMLWriter, self).__init__() + _java_obj = instance._to_java() # type: ignore[attr-defined] + self._jwrite = _java_obj.write() + + def save(self, path: str) -> None: + """Save the ML instance to the input path.""" + if not isinstance(path, str): + raise TypeError("path should be a string, got type %s" % type(path)) + self._jwrite.save(path) + + def overwrite(self) -> "JavaMLWriter": + """Overwrites if the output path already exists.""" + self._jwrite.overwrite() + return self + + def option(self, key: str, value: str) -> "JavaMLWriter": + self._jwrite.option(key, value) + return self + + def session(self, sparkSession: SparkSession) -> "JavaMLWriter": + """Sets the Spark Session to use for saving.""" + self._jwrite.session(sparkSession._jsparkSession) + return self + + +@inherit_doc +class GeneralJavaMLWriter(JavaMLWriter): + """ + (Private) Specialization of :py:class:`GeneralMLWriter` for :py:class:`JavaParams` types + """ + + def __init__(self, instance: "JavaMLWritable"): + super(GeneralJavaMLWriter, self).__init__(instance) + + def format(self, source: str) -> "GeneralJavaMLWriter": + """ + Specifies the format of ML export ("pmml", "internal", or the fully qualified class + name for export). + """ + self._jwrite.format(source) + return self + + +@inherit_doc +class MLWritable: + """ + Mixin for ML instances that provide :py:class:`MLWriter`. + + .. versionadded:: 2.0.0 + """ + + def write(self) -> MLWriter: + """Returns an MLWriter instance for this ML instance.""" + raise NotImplementedError("MLWritable is not yet implemented for type: %r" % type(self)) + + def save(self, path: str) -> None: + """Save this ML instance to the given path, a shortcut of 'write().save(path)'.""" + self.write().save(path) + + +@inherit_doc +class JavaMLWritable(MLWritable): + """ + (Private) Mixin for ML instances that provide :py:class:`JavaMLWriter`. + """ + + def write(self) -> JavaMLWriter: + """Returns an MLWriter instance for this ML instance.""" + return JavaMLWriter(self) + + +@inherit_doc +class GeneralJavaMLWritable(JavaMLWritable): + """ + (Private) Mixin for ML instances that provide :py:class:`GeneralJavaMLWriter`. + """ + + def write(self) -> GeneralJavaMLWriter: + """Returns an GeneralMLWriter instance for this ML instance.""" + return GeneralJavaMLWriter(self) + + +@inherit_doc +class MLReader(BaseReadWrite, Generic[RL]): + """ + Utility class that can load ML instances. + + .. versionadded:: 2.0.0 + """ + + def __init__(self) -> None: + super(MLReader, self).__init__() + + def load(self, path: str) -> RL: + """Load the ML instance from the input path.""" + raise NotImplementedError("MLReader is not yet implemented for type: %s" % type(self)) + + +@inherit_doc +class JavaMLReader(MLReader[RL]): + """ + (Private) Specialization of :py:class:`MLReader` for :py:class:`JavaParams` types + """ + + def __init__(self, clazz: Type["JavaMLReadable[RL]"]) -> None: + super(JavaMLReader, self).__init__() + self._clazz = clazz + self._jread = self._load_java_obj(clazz).read() + + def load(self, path: str) -> RL: + """Load the ML instance from the input path.""" + if not isinstance(path, str): + raise TypeError("path should be a string, got type %s" % type(path)) + java_obj = self._jread.load(path) + if not hasattr(self._clazz, "_from_java"): + raise NotImplementedError( + "This Java ML type cannot be loaded into Python currently: %r" % self._clazz + ) + return self._clazz._from_java(java_obj) # type: ignore[attr-defined] + + def session(self: JR, sparkSession: SparkSession) -> JR: + """Sets the Spark Session to use for loading.""" + self._jread.session(sparkSession._jsparkSession) + return self + + @classmethod + def _java_loader_class(cls, clazz: Type["JavaMLReadable[RL]"]) -> str: + """ + Returns the full class name of the Java ML instance. The default + implementation replaces "pyspark" by "org.apache.spark" in + the Python full class name. + """ + java_package = clazz.__module__.replace("pyspark", "org.apache.spark") + if clazz.__name__ in ("Pipeline", "PipelineModel"): + # Remove the last package name "pipeline" for Pipeline and PipelineModel. + java_package = ".".join(java_package.split(".")[0:-1]) + return java_package + "." + clazz.__name__ + + @classmethod + def _load_java_obj(cls, clazz: Type["JavaMLReadable[RL]"]) -> "JavaObject": + """Load the peer Java object of the ML instance.""" + java_class = cls._java_loader_class(clazz) + java_obj = _jvm() + for name in java_class.split("."): + java_obj = getattr(java_obj, name) + return java_obj + + +@inherit_doc +class MLReadable(Generic[RL]): + """ + Mixin for instances that provide :py:class:`MLReader`. + + .. versionadded:: 2.0.0 + """ + + @classmethod + def read(cls) -> MLReader[RL]: + """Returns an MLReader instance for this class.""" + raise NotImplementedError("MLReadable.read() not implemented for type: %r" % cls) + + @classmethod + def load(cls, path: str) -> RL: + """Reads an ML instance from the input path, a shortcut of `read().load(path)`.""" + return cls.read().load(path) + + +@inherit_doc +class JavaMLReadable(MLReadable[RL]): + """ + (Private) Mixin for instances that provide JavaMLReader. + """ + + @classmethod + def read(cls) -> JavaMLReader[RL]: + """Returns an MLReader instance for this class.""" + return JavaMLReader(cls) + + +@inherit_doc +class DefaultParamsWritable(MLWritable): + """ + Helper trait for making simple :py:class:`Params` types writable. If a :py:class:`Params` + class stores all data as :py:class:`Param` values, then extending this trait will provide + a default implementation of writing saved instances of the class. + This only handles simple :py:class:`Param` types; e.g., it will not handle + :py:class:`pyspark.sql.DataFrame`. See :py:class:`DefaultParamsReadable`, the counterpart + to this class. + + .. versionadded:: 2.3.0 + """ + + def write(self) -> MLWriter: + """Returns a DefaultParamsWriter instance for this class.""" + from pyspark.ml.param import Params + + if isinstance(self, Params): + return DefaultParamsWriter(self) + else: + raise TypeError( + "Cannot use DefaultParamsWritable with type %s because it does not " + + " extend Params.", + type(self), + ) + + +@inherit_doc +class DefaultParamsWriter(MLWriter): + """ + Specialization of :py:class:`MLWriter` for :py:class:`Params` types + + Class for writing Estimators and Transformers whose parameters are JSON-serializable. + + .. versionadded:: 2.3.0 + """ + + def __init__(self, instance: "Params"): + super(DefaultParamsWriter, self).__init__() + self.instance = instance + + def saveImpl(self, path: str) -> None: + DefaultParamsWriter.saveMetadata(self.instance, path, self.sc) + + @staticmethod + def extractJsonParams(instance: "Params", skipParams: Sequence[str]) -> Dict[str, Any]: + paramMap = instance.extractParamMap() + jsonParams = { + param.name: value for param, value in paramMap.items() if param.name not in skipParams + } + return jsonParams + + @staticmethod + def saveMetadata( + instance: "Params", + path: str, + sc: SparkContext, + extraMetadata: Optional[Dict[str, Any]] = None, + paramMap: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Saves metadata + Params to: path + "/metadata" + + - class + - timestamp + - sparkVersion + - uid + - paramMap + - defaultParamMap (since 2.4.0) + - (optionally, extra metadata) + + Parameters + ---------- + extraMetadata : dict, optional + Extra metadata to be saved at same level as uid, paramMap, etc. + paramMap : dict, optional + If given, this is saved in the "paramMap" field. + """ + metadataPath = os.path.join(path, "metadata") + metadataJson = DefaultParamsWriter._get_metadata_to_save( + instance, sc, extraMetadata, paramMap + ) + sc.parallelize([metadataJson], 1).saveAsTextFile(metadataPath) + + @staticmethod + def _get_metadata_to_save( + instance: "Params", + sc: SparkContext, + extraMetadata: Optional[Dict[str, Any]] = None, + paramMap: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Helper for :py:meth:`DefaultParamsWriter.saveMetadata` which extracts the JSON to save. + This is useful for ensemble models which need to save metadata for many sub-models. + + Notes + ----- + See :py:meth:`DefaultParamsWriter.saveMetadata` for details on what this includes. + """ + uid = instance.uid + cls = instance.__module__ + "." + instance.__class__.__name__ + + # User-supplied param values + params = instance._paramMap + jsonParams = {} + if paramMap is not None: + jsonParams = paramMap + else: + for p in params: + jsonParams[p.name] = params[p] + + # Default param values + jsonDefaultParams = {} + for p in instance._defaultParamMap: + jsonDefaultParams[p.name] = instance._defaultParamMap[p] + + basicMetadata = { + "class": cls, + "timestamp": int(round(time.time() * 1000)), + "sparkVersion": sc.version, + "uid": uid, + "paramMap": jsonParams, + "defaultParamMap": jsonDefaultParams, + } + if extraMetadata is not None: + basicMetadata.update(extraMetadata) + return json.dumps(basicMetadata, separators=(",", ":")) + + +@inherit_doc +class DefaultParamsReadable(MLReadable[RL]): + """ + Helper trait for making simple :py:class:`Params` types readable. + If a :py:class:`Params` class stores all data as :py:class:`Param` values, + then extending this trait will provide a default implementation of reading saved + instances of the class. This only handles simple :py:class:`Param` types; + e.g., it will not handle :py:class:`pyspark.sql.DataFrame`. See + :py:class:`DefaultParamsWritable`, the counterpart to this class. + + .. versionadded:: 2.3.0 + """ + + @classmethod + def read(cls) -> "DefaultParamsReader[RL]": + """Returns a DefaultParamsReader instance for this class.""" + return DefaultParamsReader(cls) + + +@inherit_doc +class DefaultParamsReader(MLReader[RL]): + """ + Specialization of :py:class:`MLReader` for :py:class:`Params` types + + Default :py:class:`MLReader` implementation for transformers and estimators that + contain basic (json-serializable) params and no data. This will not handle + more complex params or types with data (e.g., models with coefficients). + + .. versionadded:: 2.3.0 + """ + + def __init__(self, cls: Type[DefaultParamsReadable[RL]]): + super(DefaultParamsReader, self).__init__() + self.cls = cls + + @staticmethod + def __get_class(clazz: str) -> Type[RL]: + """ + Loads Python class from its name. + """ + parts = clazz.split(".") + module = ".".join(parts[:-1]) + m = __import__(module, fromlist=[parts[-1]]) + return getattr(m, parts[-1]) + + def load(self, path: str) -> RL: + metadata = DefaultParamsReader.loadMetadata(path, self.sc) + py_type: Type[RL] = DefaultParamsReader.__get_class(metadata["class"]) + instance = py_type() + cast("Params", instance)._resetUid(metadata["uid"]) + DefaultParamsReader.getAndSetParams(instance, metadata) + return instance + + @staticmethod + def loadMetadata(path: str, sc: SparkContext, expectedClassName: str = "") -> Dict[str, Any]: + """ + Load metadata saved using :py:meth:`DefaultParamsWriter.saveMetadata` + + Parameters + ---------- + path : str + sc : :py:class:`pyspark.SparkContext` + expectedClassName : str, optional + If non empty, this is checked against the loaded metadata. + """ + metadataPath = os.path.join(path, "metadata") + metadataStr = sc.textFile(metadataPath, 1).first() + loadedVals = DefaultParamsReader._parseMetaData(metadataStr, expectedClassName) + return loadedVals + + @staticmethod + def _parseMetaData(metadataStr: str, expectedClassName: str = "") -> Dict[str, Any]: + """ + Parse metadata JSON string produced by :py:meth`DefaultParamsWriter._get_metadata_to_save`. + This is a helper function for :py:meth:`DefaultParamsReader.loadMetadata`. + + Parameters + ---------- + metadataStr : str + JSON string of metadata + expectedClassName : str, optional + If non empty, this is checked against the loaded metadata. + """ + metadata = json.loads(metadataStr) + className = metadata["class"] + if len(expectedClassName) > 0: + assert className == expectedClassName, ( + "Error loading metadata: Expected " + + "class name {} but found class name {}".format(expectedClassName, className) + ) + return metadata + + @staticmethod + def getAndSetParams( + instance: RL, metadata: Dict[str, Any], skipParams: Optional[List[str]] = None + ) -> None: + """ + Extract Params from metadata, and set them in the instance. + """ + # Set user-supplied param values + for paramName in metadata["paramMap"]: + param = cast("Params", instance).getParam(paramName) + if skipParams is None or paramName not in skipParams: + paramValue = metadata["paramMap"][paramName] + cast("Params", instance).set(param, paramValue) + + # Set default param values + majorAndMinorVersions = VersionUtils.majorMinorVersion(metadata["sparkVersion"]) + major = majorAndMinorVersions[0] + minor = majorAndMinorVersions[1] + + # For metadata file prior to Spark 2.4, there is no default section. + if major > 2 or (major == 2 and minor >= 4): + assert "defaultParamMap" in metadata, ( + "Error loading metadata: Expected " + "`defaultParamMap` section not found" + ) + + for paramName in metadata["defaultParamMap"]: + paramValue = metadata["defaultParamMap"][paramName] + cast("Params", instance)._setDefault(**{paramName: paramValue}) + + @staticmethod + def isPythonParamsInstance(metadata: Dict[str, Any]) -> bool: + return metadata["class"].startswith("pyspark.ml.") + + @staticmethod + def loadParamsInstance(path: str, sc: SparkContext) -> RL: + """ + Load a :py:class:`Params` instance from the given path, and return it. + This assumes the instance inherits from :py:class:`MLReadable`. + """ + metadata = DefaultParamsReader.loadMetadata(path, sc) + if DefaultParamsReader.isPythonParamsInstance(metadata): + pythonClassName = metadata["class"] + else: + pythonClassName = metadata["class"].replace("org.apache.spark", "pyspark") + py_type: Type[RL] = DefaultParamsReader.__get_class(pythonClassName) + instance = py_type.load(path) + return instance + + +@inherit_doc +class HasTrainingSummary(Generic[T]): + """ + Base class for models that provides Training summary. + + .. versionadded:: 3.0.0 + """ + + @property + @since("2.1.0") + def hasSummary(self) -> bool: + """ + Indicates whether a training summary exists for this model + instance. + """ + return cast("JavaWrapper", self)._call_java("hasSummary") + + @property + @since("2.1.0") + def summary(self) -> T: + """ + Gets summary of the model trained on the training set. An exception is thrown if + no summary exists. + """ + return cast("JavaWrapper", self)._call_java("summary") + + +class MetaAlgorithmReadWrite: + @staticmethod + def isMetaEstimator(pyInstance: Any) -> bool: + from pyspark.ml import Estimator, Pipeline + from pyspark.ml.tuning import _ValidatorParams + from pyspark.ml.classification import OneVsRest + + return ( + isinstance(pyInstance, Pipeline) + or isinstance(pyInstance, OneVsRest) + or (isinstance(pyInstance, Estimator) and isinstance(pyInstance, _ValidatorParams)) + ) + + @staticmethod + def getAllNestedStages(pyInstance: Any) -> List["Params"]: + from pyspark.ml import Pipeline, PipelineModel + from pyspark.ml.tuning import _ValidatorParams + from pyspark.ml.classification import OneVsRest, OneVsRestModel + + # TODO: We need to handle `RFormulaModel.pipelineModel` here after Pyspark RFormulaModel + # support pipelineModel property. + pySubStages: Sequence["Params"] + + if isinstance(pyInstance, Pipeline): + pySubStages = pyInstance.getStages() + elif isinstance(pyInstance, PipelineModel): + pySubStages = cast(List["PipelineStage"], pyInstance.stages) + elif isinstance(pyInstance, _ValidatorParams): + raise ValueError("PySpark does not support nested validator.") + elif isinstance(pyInstance, OneVsRest): + pySubStages = [pyInstance.getClassifier()] + elif isinstance(pyInstance, OneVsRestModel): + pySubStages = [pyInstance.getClassifier()] + pyInstance.models # type: ignore[operator] + else: + pySubStages = [] + + nestedStages = [] + for pySubStage in pySubStages: + nestedStages.extend(MetaAlgorithmReadWrite.getAllNestedStages(pySubStage)) + + return [pyInstance] + nestedStages + + @staticmethod + def getUidMap(instance: Any) -> Dict[str, "Params"]: + nestedStages = MetaAlgorithmReadWrite.getAllNestedStages(instance) + uidMap = {stage.uid: stage for stage in nestedStages} + if len(nestedStages) != len(uidMap): + raise RuntimeError( + f"{instance.__class__.__module__}.{instance.__class__.__name__}" + f".load found a compound estimator with stages with duplicate " + f"UIDs. List of UIDs: {list(uidMap.keys())}." + ) + return uidMap + + +def try_remote_functions(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.ml.connect import functions + + return getattr(functions, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/ml/wrapper.py b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/wrapper.py new file mode 100644 index 0000000..5eee3ee --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/ml/wrapper.py @@ -0,0 +1,465 @@ +# +# 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, Optional, List, Type, TypeVar, TYPE_CHECKING + +from pyspark import since +from pyspark import SparkContext +from pyspark.sql import DataFrame +from pyspark.ml import Estimator, Predictor, PredictionModel, Transformer, Model +from pyspark.ml.base import _PredictorParams +from pyspark.ml.param import Param, Params +from pyspark.ml.util import _jvm +from pyspark.ml.common import inherit_doc, _java2py, _py2java + + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + from py4j.java_gateway import JavaObject, JavaClass + + +T = TypeVar("T") +JW = TypeVar("JW", bound="JavaWrapper") +JM = TypeVar("JM", bound="JavaTransformer") +JP = TypeVar("JP", bound="JavaParams") + + +class JavaWrapper: + """ + Wrapper class for a Java companion object + """ + + def __init__(self, java_obj: Optional["JavaObject"] = None): + super(JavaWrapper, self).__init__() + self._java_obj = java_obj + + def __del__(self) -> None: + if SparkContext._active_spark_context and self._java_obj is not None: + SparkContext._active_spark_context._gateway.detach( # type: ignore[union-attr] + self._java_obj + ) + + @classmethod + def _create_from_java_class(cls: Type[JW], java_class: str, *args: Any) -> JW: + """ + Construct this object from given Java classname and arguments + """ + java_obj = JavaWrapper._new_java_obj(java_class, *args) + return cls(java_obj) + + def _call_java(self, name: str, *args: Any) -> Any: + m = getattr(self._java_obj, name) + sc = SparkContext._active_spark_context + assert sc is not None + + java_args = [_py2java(sc, arg) for arg in args] + return _java2py(sc, m(*java_args)) + + @staticmethod + def _new_java_obj(java_class: str, *args: Any) -> "JavaObject": + """ + Returns a new Java object. + """ + sc = SparkContext._active_spark_context + assert sc is not None + + java_obj = _jvm() + for name in java_class.split("."): + java_obj = getattr(java_obj, name) + java_args = [_py2java(sc, arg) for arg in args] + return java_obj(*java_args) + + @staticmethod + def _new_java_array(pylist: List[Any], java_class: "JavaClass") -> "JavaObject": + """ + Create a Java array of given java_class type. Useful for + calling a method with a Scala Array from Python with Py4J. + If the param pylist is a 2D array, then a 2D java array will be returned. + The returned 2D java array is a square, non-jagged 2D array that is big + enough for all elements. The empty slots in the inner Java arrays will + be filled with null to make the non-jagged 2D array. + + Parameters + ---------- + pylist : list + Python list to convert to a Java Array. + java_class : :py:class:`py4j.java_gateway.JavaClass` + Java class to specify the type of Array. Should be in the + form of sc._gateway.jvm.* (sc is a valid Spark Context). + + Example primitive Java classes: + + - basestring -> sc._gateway.jvm.java.lang.String + - int -> sc._gateway.jvm.java.lang.Integer + - float -> sc._gateway.jvm.java.lang.Double + - bool -> sc._gateway.jvm.java.lang.Boolean + + Returns + ------- + :py:class:`py4j.java_collections.JavaArray` + Java Array of converted pylist. + """ + sc = SparkContext._active_spark_context + assert sc is not None + assert sc._gateway is not None + + java_array = None + if len(pylist) > 0 and isinstance(pylist[0], list): + # If pylist is a 2D array, then a 2D java array will be created. + # The 2D array is a square, non-jagged 2D array that is big enough for all elements. + inner_array_length = 0 + for i in range(len(pylist)): + inner_array_length = max(inner_array_length, len(pylist[i])) + java_array = sc._gateway.new_array(java_class, len(pylist), inner_array_length) + for i in range(len(pylist)): + for j in range(len(pylist[i])): + java_array[i][j] = pylist[i][j] + else: + java_array = sc._gateway.new_array(java_class, len(pylist)) + for i in range(len(pylist)): + java_array[i] = pylist[i] + return java_array + + +@inherit_doc +class JavaParams(JavaWrapper, Params, metaclass=ABCMeta): + """ + Utility class to help create wrapper classes from Java/Scala + implementations of pipeline components. + """ + + #: The param values in the Java object should be + #: synced with the Python wrapper in fit/transform/evaluate/copy. + + def _make_java_param_pair(self, param: Param[T], value: T) -> "JavaObject": + """ + Makes a Java param pair. + """ + sc = SparkContext._active_spark_context + assert sc is not None and self._java_obj is not None + + param = self._resolveParam(param) + java_param = self._java_obj.getParam(param.name) + java_value = _py2java(sc, value) + return java_param.w(java_value) + + def _transfer_params_to_java(self) -> None: + """ + Transforms the embedded params to the companion Java object. + """ + assert self._java_obj is not None + + pair_defaults = [] + for param in self.params: + if self.isSet(param): + pair = self._make_java_param_pair(param, self._paramMap[param]) + self._java_obj.set(pair) + if self.hasDefault(param): + pair = self._make_java_param_pair(param, self._defaultParamMap[param]) + pair_defaults.append(pair) + if len(pair_defaults) > 0: + sc = SparkContext._active_spark_context + assert sc is not None and sc._jvm is not None + + pair_defaults_seq = sc._jvm.PythonUtils.toSeq(pair_defaults) + self._java_obj.setDefault(pair_defaults_seq) + + def _transfer_param_map_to_java(self, pyParamMap: "ParamMap") -> "JavaObject": + """ + Transforms a Python ParamMap into a Java ParamMap. + """ + paramMap = JavaWrapper._new_java_obj("org.apache.spark.ml.param.ParamMap") + for param in self.params: + if param in pyParamMap: + pair = self._make_java_param_pair(param, pyParamMap[param]) + paramMap.put([pair]) + return paramMap + + def _create_params_from_java(self) -> None: + """ + SPARK-10931: Temporary fix to create params that are defined in the Java obj but not here + """ + assert self._java_obj is not None + + java_params = list(self._java_obj.params()) + from pyspark.ml.param import Param + + for java_param in java_params: + java_param_name = java_param.name() + if not hasattr(self, java_param_name): + param: Param[Any] = Param(self, java_param_name, java_param.doc()) + setattr(param, "created_from_java_param", True) + setattr(self, java_param_name, param) + self._params = None # need to reset so self.params will discover new params + + def _transfer_params_from_java(self) -> None: + """ + Transforms the embedded params from the companion Java object. + """ + sc = SparkContext._active_spark_context + assert sc is not None and self._java_obj is not None + + for param in self.params: + if self._java_obj.hasParam(param.name): + java_param = self._java_obj.getParam(param.name) + # SPARK-14931: Only check set params back to avoid default params mismatch. + if self._java_obj.isSet(java_param): + java_value = self._java_obj.getOrDefault(java_param) + if param.typeConverter.__name__.startswith("toList"): + value = [_java2py(sc, x) for x in list(java_value)] + else: + value = _java2py(sc, java_value) + self._set(**{param.name: value}) + # SPARK-10931: Temporary fix for params that have a default in Java + if self._java_obj.hasDefault(java_param) and not self.isDefined(param): + value = _java2py(sc, self._java_obj.getDefault(java_param)).get() + self._setDefault(**{param.name: value}) + + def _transfer_param_map_from_java(self, javaParamMap: "JavaObject") -> "ParamMap": + """ + Transforms a Java ParamMap into a Python ParamMap. + """ + sc = SparkContext._active_spark_context + assert sc is not None + + paramMap = dict() + for pair in javaParamMap.toList(): + param = pair.param() + if self.hasParam(str(param.name())): + paramMap[self.getParam(param.name())] = _java2py(sc, pair.value()) + return paramMap + + @staticmethod + def _empty_java_param_map() -> "JavaObject": + """ + Returns an empty Java ParamMap reference. + """ + return _jvm().org.apache.spark.ml.param.ParamMap() + + def _to_java(self) -> "JavaObject": + """ + Transfer this instance's Params to the wrapped Java object, and return the Java object. + Used for ML persistence. + + Meta-algorithms such as Pipeline should override this method. + + Returns + ------- + py4j.java_gateway.JavaObject + Java object equivalent to this instance. + """ + self._transfer_params_to_java() + return self._java_obj + + @staticmethod + def _from_java(java_stage: "JavaObject") -> "JP": # type: ignore + """ + Given a Java object, create and return a Python wrapper of it. + Used for ML persistence. + + Meta-algorithms such as Pipeline should override this method as a classmethod. + """ + + def __get_class(clazz: str) -> Type[JP]: + """ + Loads Python class from its name. + """ + parts = clazz.split(".") + module = ".".join(parts[:-1]) + m = __import__(module, fromlist=[parts[-1]]) + return getattr(m, parts[-1]) + + stage_name = java_stage.getClass().getName().replace("org.apache.spark", "pyspark") + # Generate a default new instance from the stage_name class. + py_type = __get_class(stage_name) + if issubclass(py_type, JavaParams): + # Load information from java_stage to the instance. + py_stage = py_type() + py_stage._java_obj = java_stage + + # SPARK-10931: Temporary fix so that persisted models would own params from Estimator + if issubclass(py_type, JavaModel): + py_stage._create_params_from_java() + + py_stage._resetUid(java_stage.uid()) + py_stage._transfer_params_from_java() + elif hasattr(py_type, "_from_java"): + py_stage = py_type._from_java(java_stage) + else: + raise NotImplementedError( + "This Java stage cannot be loaded into Python currently: %r" % stage_name + ) + return py_stage + + def copy(self: "JP", extra: Optional["ParamMap"] = None) -> "JP": + """ + Creates a copy of this instance with the same uid and some + extra params. This implementation first calls Params.copy and + then make a copy of the companion Java pipeline component with + extra params. So both the Python wrapper and the Java pipeline + component get copied. + + Parameters + ---------- + extra : dict, optional + Extra parameters to copy to the new instance + + Returns + ------- + :py:class:`JavaParams` + Copy of this instance + """ + if extra is None: + extra = dict() + that = super(JavaParams, self).copy(extra) + if self._java_obj is not None: + that._java_obj = self._java_obj.copy(self._empty_java_param_map()) + that._transfer_params_to_java() + return that + + def clear(self, param: Param) -> None: + """ + Clears a param from the param map if it has been explicitly set. + """ + assert self._java_obj is not None + + super(JavaParams, self).clear(param) + java_param = self._java_obj.getParam(param.name) + self._java_obj.clear(java_param) + + +@inherit_doc +class JavaEstimator(JavaParams, Estimator[JM], metaclass=ABCMeta): + """ + Base class for :py:class:`Estimator`s that wrap Java/Scala + implementations. + """ + + @abstractmethod + def _create_model(self, java_model: "JavaObject") -> JM: + """ + Creates a model from the input Java model reference. + """ + raise NotImplementedError() + + def _fit_java(self, dataset: DataFrame) -> "JavaObject": + """ + Fits a Java model to the input dataset. + + Examples + -------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + + Returns + ------- + py4j.java_gateway.JavaObject + fitted Java model + """ + assert self._java_obj is not None + + self._transfer_params_to_java() + return self._java_obj.fit(dataset._jdf) + + def _fit(self, dataset: DataFrame) -> JM: + java_model = self._fit_java(dataset) + model = self._create_model(java_model) + return self._copyValues(model) + + +@inherit_doc +class JavaTransformer(JavaParams, Transformer, metaclass=ABCMeta): + """ + Base class for :py:class:`Transformer`s that wrap Java/Scala + implementations. Subclasses should ensure they have the transformer Java object + available as _java_obj. + """ + + def _transform(self, dataset: DataFrame) -> DataFrame: + assert self._java_obj is not None + + self._transfer_params_to_java() + return DataFrame(self._java_obj.transform(dataset._jdf), dataset.sparkSession) + + +@inherit_doc +class JavaModel(JavaTransformer, Model, metaclass=ABCMeta): + """ + Base class for :py:class:`Model`s that wrap Java/Scala + implementations. Subclasses should inherit this class before + param mix-ins, because this sets the UID from the Java model. + """ + + def __init__(self, java_model: Optional["JavaObject"] = None): + """ + Initialize this instance with a Java model object. + Subclasses should call this constructor, initialize params, + and then call _transfer_params_from_java. + + This instance can be instantiated without specifying java_model, + it will be assigned after that, but this scenario only used by + :py:class:`JavaMLReader` to load models. This is a bit of a + hack, but it is easiest since a proper fix would require + MLReader (in pyspark.ml.util) to depend on these wrappers, but + these wrappers depend on pyspark.ml.util (both directly and via + other ML classes). + """ + super(JavaModel, self).__init__(java_model) + if java_model is not None: + + # SPARK-10931: This is a temporary fix to allow models to own params + # from estimators. Eventually, these params should be in models through + # using common base classes between estimators and models. + self._create_params_from_java() + + self._resetUid(java_model.uid()) + + def __repr__(self) -> str: + return self._call_java("toString") + + +@inherit_doc +class JavaPredictor(Predictor, JavaEstimator[JM], _PredictorParams, Generic[JM], metaclass=ABCMeta): + """ + (Private) Java Estimator for prediction tasks (regression and classification). + """ + + pass + + +@inherit_doc +class JavaPredictionModel(PredictionModel[T], JavaModel, _PredictorParams): + """ + (Private) Java Model for prediction tasks (regression and classification). + """ + + @property + @since("2.1.0") + def numFeatures(self) -> int: + """ + Returns the number of features the model was trained on. If unknown, returns -1 + """ + return self._call_java("numFeatures") + + @since("3.0.0") + def predict(self, value: T) -> float: + """ + Predict label for the given features. + """ + return self._call_java("predict", value) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/__init__.py new file mode 100644 index 0000000..392a693 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/__init__.py @@ -0,0 +1,44 @@ +# +# 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. +# + +""" +RDD-based machine learning APIs for Python (in maintenance mode). + +The `pyspark.mllib` package is in maintenance mode as of the Spark 2.0.0 release to encourage +migration to the DataFrame-based APIs under the `pyspark.ml` package. +""" +# MLlib currently needs NumPy 1.4+, so complain if lower + +import numpy + +ver = [int(x) for x in numpy.version.version.split(".")[:2]] +if ver < [1, 4]: + raise RuntimeError("MLlib requires NumPy 1.4+") + +__all__ = [ + "classification", + "clustering", + "feature", + "fpm", + "linalg", + "random", + "recommendation", + "regression", + "stat", + "tree", + "util", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/_typing.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/_typing.pyi new file mode 100644 index 0000000..f196bd5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/_typing.pyi @@ -0,0 +1,33 @@ +# +# 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 List, Tuple, TypeVar, Union + +from typing_extensions import Literal + +from pyspark.mllib.linalg import Vector +from numpy import ndarray # noqa: F401 +from py4j.java_gateway import JavaObject + +VectorLike = Union[ndarray, Vector, List[float], Tuple[float, ...]] +C = TypeVar("C", bound=type) +JavaObjectOrPickleDump = Union[JavaObject, bytearray, bytes] + +CorrMethodType = Union[Literal["spearman"], Literal["pearson"]] +KolmogorovSmirnovTestDistNameType = Literal["norm"] +NormType = Union[None, float, Literal["fro"], Literal["nuc"]] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/classification.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/classification.py new file mode 100644 index 0000000..1e1795d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/classification.py @@ -0,0 +1,989 @@ +# +# 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 math import exp +import sys +import warnings +from typing import Any, Iterable, Optional, Union, overload, TYPE_CHECKING + +import numpy + +from pyspark import RDD, SparkContext, since +from pyspark.streaming.dstream import DStream +from pyspark.mllib.common import callMLlibFunc, _py2java, _java2py +from pyspark.mllib.linalg import _convert_to_vector +from pyspark.mllib.regression import ( + LabeledPoint, + LinearModel, + _regression_train_wrapper, + StreamingLinearAlgorithm, +) +from pyspark.mllib.util import Saveable, Loader, inherit_doc +from pyspark.mllib.linalg import Vector +from pyspark.mllib.regression import LabeledPoint + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike + + +__all__ = [ + "LogisticRegressionModel", + "LogisticRegressionWithSGD", + "LogisticRegressionWithLBFGS", + "SVMModel", + "SVMWithSGD", + "NaiveBayesModel", + "NaiveBayes", + "StreamingLogisticRegressionWithSGD", +] + + +class LinearClassificationModel(LinearModel): + """ + A private abstract class representing a multiclass classification + model. The categories are represented by int values: 0, 1, 2, etc. + """ + + def __init__(self, weights: Vector, intercept: float) -> None: + super(LinearClassificationModel, self).__init__(weights, intercept) + self._threshold: Optional[float] = None + + @since("1.4.0") + def setThreshold(self, value: float) -> None: + """ + Sets the threshold that separates positive predictions from + negative predictions. An example with prediction score greater + than or equal to this threshold is identified as a positive, + and negative otherwise. It is used for binary classification + only. + """ + self._threshold = value + + @property + @since("1.4.0") + def threshold(self) -> Optional[float]: + """ + Returns the threshold (if any) used for converting raw + prediction scores into 0/1 predictions. It is used for + binary classification only. + """ + return self._threshold + + @since("1.4.0") + def clearThreshold(self) -> None: + """ + Clears the threshold so that `predict` will output raw + prediction scores. It is used for binary classification only. + """ + self._threshold = None + + @overload + def predict(self, test: "VectorLike") -> Union[int, float]: + ... + + @overload + def predict(self, test: RDD["VectorLike"]) -> RDD[Union[int, float]]: + ... + + def predict( + self, test: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[RDD[Union[int, float]], Union[int, float]]: + """ + Predict values for a single data point or an RDD of points + using the model trained. + + .. versionadded:: 1.4.0 + """ + raise NotImplementedError + + +class LogisticRegressionModel(LinearClassificationModel): + + """ + Classification model trained using Multinomial/Binary Logistic + Regression. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + weights : :py:class:`pyspark.mllib.linalg.Vector` + Weights computed for every feature. + intercept : float + Intercept computed for this model. (Only used in Binary Logistic + Regression. In Multinomial Logistic Regression, the intercepts will + not be a single value, so the intercepts will be part of the + weights.) + numFeatures : int + The dimension of the features. + numClasses : int + The number of possible outcomes for k classes classification problem + in Multinomial Logistic Regression. By default, it is binary + logistic regression so numClasses will be set to 2. + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> data = [ + ... LabeledPoint(0.0, [0.0, 1.0]), + ... LabeledPoint(1.0, [1.0, 0.0]), + ... ] + >>> lrm = LogisticRegressionWithSGD.train(sc.parallelize(data), iterations=10) + >>> lrm.predict([1.0, 0.0]) + 1 + >>> lrm.predict([0.0, 1.0]) + 0 + >>> lrm.predict(sc.parallelize([[1.0, 0.0], [0.0, 1.0]])).collect() + [1, 0] + >>> lrm.clearThreshold() + >>> lrm.predict([0.0, 1.0]) + 0.279... + + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) + ... ] + >>> lrm = LogisticRegressionWithSGD.train(sc.parallelize(sparse_data), iterations=10) + >>> lrm.predict(numpy.array([0.0, 1.0])) + 1 + >>> lrm.predict(numpy.array([1.0, 0.0])) + 0 + >>> lrm.predict(SparseVector(2, {1: 1.0})) + 1 + >>> lrm.predict(SparseVector(2, {0: 1.0})) + 0 + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> lrm.save(sc, path) + >>> sameModel = LogisticRegressionModel.load(sc, path) + >>> sameModel.predict(numpy.array([0.0, 1.0])) + 1 + >>> sameModel.predict(SparseVector(2, {0: 1.0})) + 0 + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except BaseException: + ... pass + >>> multi_class_data = [ + ... LabeledPoint(0.0, [0.0, 1.0, 0.0]), + ... LabeledPoint(1.0, [1.0, 0.0, 0.0]), + ... LabeledPoint(2.0, [0.0, 0.0, 1.0]) + ... ] + >>> data = sc.parallelize(multi_class_data) + >>> mcm = LogisticRegressionWithLBFGS.train(data, iterations=10, numClasses=3) + >>> mcm.predict([0.0, 0.5, 0.0]) + 0 + >>> mcm.predict([0.8, 0.0, 0.0]) + 1 + >>> mcm.predict([0.0, 0.0, 0.3]) + 2 + """ + + def __init__( + self, weights: Vector, intercept: float, numFeatures: int, numClasses: int + ) -> None: + super(LogisticRegressionModel, self).__init__(weights, intercept) + self._numFeatures = int(numFeatures) + self._numClasses = int(numClasses) + self._threshold = 0.5 + if self._numClasses == 2: + self._dataWithBiasSize = None + self._weightsMatrix = None + else: + self._dataWithBiasSize = self._coeff.size // ( # type: ignore[attr-defined] + self._numClasses - 1 + ) + self._weightsMatrix = self._coeff.toArray().reshape( + self._numClasses - 1, self._dataWithBiasSize + ) + + @property + @since("1.4.0") + def numFeatures(self) -> int: + """ + Dimension of the features. + """ + return self._numFeatures + + @property + @since("1.4.0") + def numClasses(self) -> int: + """ + Number of possible outcomes for k classes classification problem + in Multinomial Logistic Regression. + """ + return self._numClasses + + @overload + def predict(self, x: "VectorLike") -> Union[int, float]: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[Union[int, float]]: + ... + + def predict( + self, x: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[RDD[Union[int, float]], Union[int, float]]: + """ + Predict values for a single data point or an RDD of points + using the model trained. + + .. versionadded:: 0.9.0 + """ + if isinstance(x, RDD): + return x.map(lambda v: self.predict(v)) + + x = _convert_to_vector(x) + if self.numClasses == 2: + margin = self.weights.dot(x) + self._intercept # type: ignore[attr-defined] + if margin > 0: + prob = 1 / (1 + exp(-margin)) + else: + exp_margin = exp(margin) + prob = exp_margin / (1 + exp_margin) + if self._threshold is None: + return prob + else: + return 1 if prob > self._threshold else 0 + else: + assert self._weightsMatrix is not None + + best_class = 0 + max_margin = 0.0 + if x.size + 1 == self._dataWithBiasSize: # type: ignore[attr-defined] + for i in range(0, self._numClasses - 1): + margin = ( + x.dot(self._weightsMatrix[i][0 : x.size]) # type: ignore[attr-defined] + + self._weightsMatrix[i][x.size] # type: ignore[attr-defined] + ) + if margin > max_margin: + max_margin = margin + best_class = i + 1 + else: + for i in range(0, self._numClasses - 1): + margin = x.dot(self._weightsMatrix[i]) # type: ignore[attr-defined] + if margin > max_margin: + max_margin = margin + best_class = i + 1 + return best_class + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """ + Save this model to the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.classification.LogisticRegressionModel( + _py2java(sc, self._coeff), self.intercept, self.numFeatures, self.numClasses + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "LogisticRegressionModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.classification.LogisticRegressionModel.load( + sc._jsc.sc(), path + ) + weights = _java2py(sc, java_model.weights()) + intercept = java_model.intercept() + numFeatures = java_model.numFeatures() + numClasses = java_model.numClasses() + threshold = java_model.getThreshold().get() + model = LogisticRegressionModel(weights, intercept, numFeatures, numClasses) + model.setThreshold(threshold) + return model + + def __repr__(self) -> str: + return ( + "pyspark.mllib.LogisticRegressionModel: intercept = {}, " + "numFeatures = {}, numClasses = {}, threshold = {}" + ).format(self._intercept, self._numFeatures, self._numClasses, self._threshold) + + +class LogisticRegressionWithSGD: + """ + Train a classification model for Binary Logistic Regression using Stochastic Gradient Descent. + + .. versionadded:: 0.9.0 + .. deprecated:: 2.0.0 + Use ml.classification.LogisticRegression or LogisticRegressionWithLBFGS. + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + step: float = 1.0, + miniBatchFraction: float = 1.0, + initialWeights: Optional["VectorLike"] = None, + regParam: float = 0.01, + regType: str = "l2", + intercept: bool = False, + validateData: bool = True, + convergenceTol: float = 0.001, + ) -> LogisticRegressionModel: + """ + Train a logistic regression model on the given data. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of :py:class:`pyspark.mllib.regression.LabeledPoint`. + iterations : int, optional + The number of iterations. + (default: 100) + step : float, optional + The step parameter used in SGD. + (default: 1.0) + miniBatchFraction : float, optional + Fraction of data to be used for each SGD iteration. + (default: 1.0) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + regParam : float, optional + The regularizer parameter. + (default: 0.01) + regType : str, optional + The type of regularizer used for training our model. + Supported values: + + - "l1" for using L1 regularization + - "l2" for using L2 regularization (default) + - None for no regularization + + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e., whether bias + features are activated or not). + (default: False) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + convergenceTol : float, optional + A condition which decides iteration termination. + (default: 0.001) + """ + warnings.warn( + "Deprecated in 2.0.0. Use ml.classification.LogisticRegression or " + "LogisticRegressionWithLBFGS.", + FutureWarning, + ) + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainLogisticRegressionModelWithSGD", + rdd, + int(iterations), + float(step), + float(miniBatchFraction), + i, + float(regParam), + regType, + bool(intercept), + bool(validateData), + float(convergenceTol), + ) + + return _regression_train_wrapper(train, LogisticRegressionModel, data, initialWeights) + + +class LogisticRegressionWithLBFGS: + """ + Train a classification model for Multinomial/Binary Logistic Regression + using Limited-memory BFGS. + + Standard feature scaling and L2 regularization are used by default. + .. versionadded:: 1.2.0 + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + initialWeights: Optional["VectorLike"] = None, + regParam: float = 0.0, + regType: str = "l2", + intercept: bool = False, + corrections: int = 10, + tolerance: float = 1e-6, + validateData: bool = True, + numClasses: int = 2, + ) -> LogisticRegressionModel: + """ + Train a logistic regression model on the given data. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of :py:class:`pyspark.mllib.regression.LabeledPoint`. + iterations : int, optional + The number of iterations. + (default: 100) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + regParam : float, optional + The regularizer parameter. + (default: 0.01) + regType : str, optional + The type of regularizer used for training our model. + Supported values: + + - "l1" for using L1 regularization + - "l2" for using L2 regularization (default) + - None for no regularization + + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e., whether bias + features are activated or not). + (default: False) + corrections : int, optional + The number of corrections used in the LBFGS update. + If a known updater is used for binary classification, + it calls the ml implementation and this parameter will + have no effect. (default: 10) + tolerance : float, optional + The convergence tolerance of iterations for L-BFGS. + (default: 1e-6) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + numClasses : int, optional + The number of classes (i.e., outcomes) a label can take in + Multinomial Logistic Regression. + (default: 2) + + Examples + -------- + >>> data = [ + ... LabeledPoint(0.0, [0.0, 1.0]), + ... LabeledPoint(1.0, [1.0, 0.0]), + ... ] + >>> lrm = LogisticRegressionWithLBFGS.train(sc.parallelize(data), iterations=10) + >>> lrm.predict([1.0, 0.0]) + 1 + >>> lrm.predict([0.0, 1.0]) + 0 + """ + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainLogisticRegressionModelWithLBFGS", + rdd, + int(iterations), + i, + float(regParam), + regType, + bool(intercept), + int(corrections), + float(tolerance), + bool(validateData), + int(numClasses), + ) + + if initialWeights is None: + if numClasses == 2: + initialWeights = [0.0] * len(data.first().features) + else: + if intercept: + initialWeights = [0.0] * (len(data.first().features) + 1) * (numClasses - 1) + else: + initialWeights = [0.0] * len(data.first().features) * (numClasses - 1) + return _regression_train_wrapper(train, LogisticRegressionModel, data, initialWeights) + + +class SVMModel(LinearClassificationModel): + + """ + Model for Support Vector Machines (SVMs). + + .. versionadded:: 0.9.0 + + Parameters + ---------- + weights : :py:class:`pyspark.mllib.linalg.Vector` + Weights computed for every feature. + intercept : float + Intercept computed for this model. + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(1.0, [1.0]), + ... LabeledPoint(1.0, [2.0]), + ... LabeledPoint(1.0, [3.0]) + ... ] + >>> svm = SVMWithSGD.train(sc.parallelize(data), iterations=10) + >>> svm.predict([1.0]) + 1 + >>> svm.predict(sc.parallelize([[1.0]])).collect() + [1] + >>> svm.clearThreshold() + >>> svm.predict(numpy.array([1.0])) + 1.44... + + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {0: -1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) + ... ] + >>> svm = SVMWithSGD.train(sc.parallelize(sparse_data), iterations=10) + >>> svm.predict(SparseVector(2, {1: 1.0})) + 1 + >>> svm.predict(SparseVector(2, {0: -1.0})) + 0 + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> svm.save(sc, path) + >>> sameModel = SVMModel.load(sc, path) + >>> sameModel.predict(SparseVector(2, {1: 1.0})) + 1 + >>> sameModel.predict(SparseVector(2, {0: -1.0})) + 0 + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except BaseException: + ... pass + """ + + def __init__(self, weights: Vector, intercept: float) -> None: + super(SVMModel, self).__init__(weights, intercept) + self._threshold = 0.0 + + @overload + def predict(self, x: "VectorLike") -> Union[int, float]: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[Union[int, float]]: + ... + + def predict( + self, x: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[RDD[Union[int, float]], Union[int, float]]: + """ + Predict values for a single data point or an RDD of points + using the model trained. + + .. versionadded:: 0.9.0 + """ + if isinstance(x, RDD): + return x.map(lambda v: self.predict(v)) + + x = _convert_to_vector(x) + margin = self.weights.dot(x) + self.intercept # type: ignore[attr-defined] + if self._threshold is None: + return margin + else: + return 1 if margin > self._threshold else 0 + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """ + Save this model to the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.classification.SVMModel( + _py2java(sc, self._coeff), self.intercept + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "SVMModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.classification.SVMModel.load(sc._jsc.sc(), path) + weights = _java2py(sc, java_model.weights()) + intercept = java_model.intercept() + threshold = java_model.getThreshold().get() + model = SVMModel(weights, intercept) + model.setThreshold(threshold) + return model + + +class SVMWithSGD: + """ + Train a Support Vector Machine (SVM) using Stochastic Gradient Descent. + + .. versionadded:: 0.9.0 + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + step: float = 1.0, + regParam: float = 0.01, + miniBatchFraction: float = 1.0, + initialWeights: Optional["VectorLike"] = None, + regType: str = "l2", + intercept: bool = False, + validateData: bool = True, + convergenceTol: float = 0.001, + ) -> SVMModel: + """ + Train a support vector machine on the given data. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of :py:class:`pyspark.mllib.regression.LabeledPoint`. + iterations : int, optional + The number of iterations. + (default: 100) + step : float, optional + The step parameter used in SGD. + (default: 1.0) + regParam : float, optional + The regularizer parameter. + (default: 0.01) + miniBatchFraction : float, optional + Fraction of data to be used for each SGD iteration. + (default: 1.0) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + regType : str, optional + The type of regularizer used for training our model. + Allowed values: + + - "l1" for using L1 regularization + - "l2" for using L2 regularization (default) + - None for no regularization + + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e. whether bias + features are activated or not). + (default: False) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + convergenceTol : float, optional + A condition which decides iteration termination. + (default: 0.001) + """ + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainSVMModelWithSGD", + rdd, + int(iterations), + float(step), + float(regParam), + float(miniBatchFraction), + i, + regType, + bool(intercept), + bool(validateData), + float(convergenceTol), + ) + + return _regression_train_wrapper(train, SVMModel, data, initialWeights) + + +@inherit_doc +class NaiveBayesModel(Saveable, Loader["NaiveBayesModel"]): + + """ + Model for Naive Bayes classifiers. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + labels : :py:class:`numpy.ndarray` + List of labels. + pi : :py:class:`numpy.ndarray` + Log of class priors, whose dimension is C, number of labels. + theta : :py:class:`numpy.ndarray` + Log of class conditional probabilities, whose dimension is C-by-D, + where D is number of features. + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> data = [ + ... LabeledPoint(0.0, [0.0, 0.0]), + ... LabeledPoint(0.0, [0.0, 1.0]), + ... LabeledPoint(1.0, [1.0, 0.0]), + ... ] + >>> model = NaiveBayes.train(sc.parallelize(data)) + >>> model.predict(numpy.array([0.0, 1.0])) + 0.0 + >>> model.predict(numpy.array([1.0, 0.0])) + 1.0 + >>> model.predict(sc.parallelize([[1.0, 0.0]])).collect() + [1.0] + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {1: 0.0})), + ... LabeledPoint(0.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {0: 1.0})) + ... ] + >>> model = NaiveBayes.train(sc.parallelize(sparse_data)) + >>> model.predict(SparseVector(2, {1: 1.0})) + 0.0 + >>> model.predict(SparseVector(2, {0: 1.0})) + 1.0 + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = NaiveBayesModel.load(sc, path) + >>> sameModel.predict(SparseVector(2, {0: 1.0})) == model.predict(SparseVector(2, {0: 1.0})) + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + def __init__(self, labels: numpy.ndarray, pi: numpy.ndarray, theta: numpy.ndarray) -> None: + self.labels = labels + self.pi = pi + self.theta = theta + + @overload + def predict(self, x: "VectorLike") -> numpy.float64: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[numpy.float64]: + ... + + @since("0.9.0") + def predict( + self, x: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[numpy.float64, RDD[numpy.float64]]: + """ + Return the most likely class for a data vector + or an RDD of vectors + """ + if isinstance(x, RDD): + return x.map(lambda v: self.predict(v)) + x = _convert_to_vector(x) + return self.labels[ + numpy.argmax(self.pi + x.dot(self.theta.transpose())) # type: ignore[attr-defined] + ] + + def save(self, sc: SparkContext, path: str) -> None: + """ + Save this model to the given path. + """ + assert sc._jvm is not None + + java_labels = _py2java(sc, self.labels.tolist()) + java_pi = _py2java(sc, self.pi.tolist()) + java_theta = _py2java(sc, self.theta.tolist()) + java_model = sc._jvm.org.apache.spark.mllib.classification.NaiveBayesModel( + java_labels, java_pi, java_theta + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "NaiveBayesModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.classification.NaiveBayesModel.load( + sc._jsc.sc(), path + ) + # Can not unpickle array.array from Pickle in Python3 with "bytes" + py_labels = _java2py(sc, java_model.labels(), "latin1") + py_pi = _java2py(sc, java_model.pi(), "latin1") + py_theta = _java2py(sc, java_model.theta(), "latin1") + return NaiveBayesModel(py_labels, py_pi, numpy.array(py_theta)) + + +class NaiveBayes: + """ + Train a Multinomial Naive Bayes model. + + .. versionadded:: 0.9.0 + """ + + @classmethod + def train(cls, data: RDD[LabeledPoint], lambda_: float = 1.0) -> NaiveBayesModel: + """ + Train a Naive Bayes model given an RDD of (label, features) + vectors. + + This is the `Multinomial NB `_ which + can handle all kinds of discrete data. For example, by + converting documents into TF-IDF vectors, it can be used for + document classification. By making every vector a 0-1 vector, + it can also be used as `Bernoulli NB `_. + The input feature values must be nonnegative. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of :py:class:`pyspark.mllib.regression.LabeledPoint`. + lambda\\_ : float, optional + The smoothing parameter. + (default: 1.0) + """ + first = data.first() + if not isinstance(first, LabeledPoint): + raise ValueError("`data` should be an RDD of LabeledPoint") + labels, pi, theta = callMLlibFunc("trainNaiveBayesModel", data, lambda_) + return NaiveBayesModel(labels.toArray(), pi.toArray(), numpy.array(theta)) + + +@inherit_doc +class StreamingLogisticRegressionWithSGD(StreamingLinearAlgorithm): + """ + Train or predict a logistic regression model on streaming data. + Training uses Stochastic Gradient Descent to update the model based on + each new batch of incoming data from a DStream. + + Each batch of data is assumed to be an RDD of LabeledPoints. + The number of data points per batch can vary, but the number + of features must be constant. An initial weight + vector must be provided. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + stepSize : float, optional + Step size for each iteration of gradient descent. + (default: 0.1) + numIterations : int, optional + Number of iterations run for each batch of data. + (default: 50) + miniBatchFraction : float, optional + Fraction of each batch of data to use for updates. + (default: 1.0) + regParam : float, optional + L2 Regularization parameter. + (default: 0.0) + convergenceTol : float, optional + Value used to determine when to terminate iterations. + (default: 0.001) + """ + + def __init__( + self, + stepSize: float = 0.1, + numIterations: int = 50, + miniBatchFraction: float = 1.0, + regParam: float = 0.0, + convergenceTol: float = 0.001, + ) -> None: + self.stepSize = stepSize + self.numIterations = numIterations + self.regParam = regParam + self.miniBatchFraction = miniBatchFraction + self.convergenceTol = convergenceTol + self._model: Optional[LogisticRegressionModel] = None + super(StreamingLogisticRegressionWithSGD, self).__init__(model=self._model) + + @since("1.5.0") + def setInitialWeights( + self, initialWeights: "VectorLike" + ) -> "StreamingLogisticRegressionWithSGD": + """ + Set the initial value of weights. + + This must be set before running trainOn and predictOn. + """ + initialWeights = _convert_to_vector(initialWeights) + + # LogisticRegressionWithSGD does only binary classification. + self._model = LogisticRegressionModel( + initialWeights, 0, initialWeights.size, 2 # type: ignore[attr-defined] + ) + return self + + @since("1.5.0") + def trainOn(self, dstream: "DStream[LabeledPoint]") -> None: + """Train the model on the incoming dstream.""" + self._validate(dstream) + + def update(rdd: RDD[LabeledPoint]) -> None: + # LogisticRegressionWithSGD.train raises an error for an empty RDD. + if not rdd.isEmpty(): + self._model = LogisticRegressionWithSGD.train( + rdd, + self.numIterations, + self.stepSize, + self.miniBatchFraction, + self._model.weights, # type: ignore[union-attr] + regParam=self.regParam, + convergenceTol=self.convergenceTol, + ) + + dstream.foreachRDD(update) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.mllib.classification + + globs = pyspark.mllib.classification.__dict__.copy() + spark = ( + SparkSession.builder.master("local[4]").appName("mllib.classification tests").getOrCreate() + ) + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/clustering.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/clustering.py new file mode 100644 index 0000000..4595268 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/clustering.py @@ -0,0 +1,1318 @@ +# +# 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 sys +import array as pyarray +from math import exp, log +from collections import namedtuple +from typing import Any, List, Optional, Tuple, TypeVar, Union, overload, TYPE_CHECKING + +import numpy as np +from numpy import array, random, tile + +from pyspark import SparkContext, since +from pyspark.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc, callJavaFunc, _py2java, _java2py +from pyspark.mllib.linalg import SparseVector, _convert_to_vector, DenseVector # noqa: F401 +from pyspark.mllib.stat.distribution import MultivariateGaussian +from pyspark.mllib.util import Saveable, Loader, inherit_doc, JavaLoader, JavaSaveable +from pyspark.streaming import DStream + +if TYPE_CHECKING: + from py4j.java_gateway import JavaObject + from pyspark.mllib._typing import VectorLike + +T = TypeVar("T") + +__all__ = [ + "BisectingKMeansModel", + "BisectingKMeans", + "KMeansModel", + "KMeans", + "GaussianMixtureModel", + "GaussianMixture", + "PowerIterationClusteringModel", + "PowerIterationClustering", + "StreamingKMeans", + "StreamingKMeansModel", + "LDA", + "LDAModel", +] + + +@inherit_doc +class BisectingKMeansModel(JavaModelWrapper): + """ + A clustering model derived from the bisecting k-means method. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> data = array([0.0,0.0, 1.0,1.0, 9.0,8.0, 8.0,9.0]).reshape(4, 2) + >>> bskm = BisectingKMeans() + >>> model = bskm.train(sc.parallelize(data, 2), k=4) + >>> p = array([0.0, 0.0]) + >>> model.predict(p) + 0 + >>> model.k + 4 + >>> model.computeCost(p) + 0.0 + """ + + def __init__(self, java_model: "JavaObject"): + super(BisectingKMeansModel, self).__init__(java_model) + self.centers = [c.toArray() for c in self.call("clusterCenters")] + + @property + @since("2.0.0") + def clusterCenters(self) -> List[np.ndarray]: + """Get the cluster centers, represented as a list of NumPy + arrays.""" + return self.centers + + @property + @since("2.0.0") + def k(self) -> int: + """Get the number of clusters""" + return self.call("k") + + @overload + def predict(self, x: "VectorLike") -> int: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[int]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[int, RDD[int]]: + """ + Find the cluster that each of the points belongs to in this + model. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + A data point (or RDD of points) to determine cluster index. + :py:class:`pyspark.mllib.linalg.Vector` can be replaced with equivalent + objects (list, tuple, numpy.ndarray). + + Returns + ------- + int or :py:class:`pyspark.RDD` of int + Predicted cluster index or an RDD of predicted cluster indices + if the input is an RDD. + """ + if isinstance(x, RDD): + vecs = x.map(_convert_to_vector) + return self.call("predict", vecs) + + x = _convert_to_vector(x) + return self.call("predict", x) + + def computeCost(self, x: Union["VectorLike", RDD["VectorLike"]]) -> float: + """ + Return the Bisecting K-means cost (sum of squared distances of + points to their nearest center) for this model on the given + data. If provided with an RDD of points returns the sum. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + point : :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). + """ + if isinstance(x, RDD): + vecs = x.map(_convert_to_vector) + return self.call("computeCost", vecs) + + return self.call("computeCost", _convert_to_vector(x)) + + +class BisectingKMeans: + """ + A bisecting k-means algorithm based on the paper "A comparison of + document clustering techniques" by Steinbach, Karypis, and Kumar, + with modification to fit Spark. + The algorithm starts from a single cluster that contains all points. + Iteratively it finds divisible clusters on the bottom level and + bisects each of them using k-means, until there are `k` leaf + clusters in total or no leaf clusters are divisible. + The bisecting steps of clusters on the same level are grouped + together to increase parallelism. If bisecting all divisible + clusters on the bottom level would result more than `k` leaf + clusters, larger clusters get higher priority. + + .. versionadded:: 2.0.0 + + Notes + ----- + See the original paper [1]_ + + .. [1] Steinbach, M. et al. "A Comparison of Document Clustering Techniques." (2000). + KDD Workshop on Text Mining, 2000 + http://glaros.dtc.umn.edu/gkhome/fetch/papers/docclusterKDDTMW00.pdf + """ + + @classmethod + def train( + cls, + rdd: RDD["VectorLike"], + k: int = 4, + maxIterations: int = 20, + minDivisibleClusterSize: float = 1.0, + seed: int = -1888008604, + ) -> BisectingKMeansModel: + """ + Runs the bisecting k-means algorithm return the model. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + rdd : :py:class:`pyspark.RDD` + Training points as an `RDD` of `Vector` or convertible + sequence types. + k : int, optional + The desired number of leaf clusters. The actual number could + be smaller if there are no divisible leaf clusters. + (default: 4) + maxIterations : int, optional + Maximum number of iterations allowed to split clusters. + (default: 20) + minDivisibleClusterSize : float, optional + Minimum number of points (if >= 1.0) or the minimum proportion + of points (if < 1.0) of a divisible cluster. + (default: 1) + seed : int, optional + Random seed value for cluster initialization. + (default: -1888008604 from classOf[BisectingKMeans].getName.##) + """ + java_model = callMLlibFunc( + "trainBisectingKMeans", + rdd.map(_convert_to_vector), + k, + maxIterations, + minDivisibleClusterSize, + seed, + ) + return BisectingKMeansModel(java_model) + + +@inherit_doc +class KMeansModel(Saveable, Loader["KMeansModel"]): + + """A clustering model derived from the k-means method. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> data = array([0.0,0.0, 1.0,1.0, 9.0,8.0, 8.0,9.0]).reshape(4, 2) + >>> model = KMeans.train( + ... sc.parallelize(data), 2, maxIterations=10, initializationMode="random", + ... seed=50, initializationSteps=5, epsilon=1e-4) + >>> model.predict(array([0.0, 0.0])) == model.predict(array([1.0, 1.0])) + True + >>> model.predict(array([8.0, 9.0])) == model.predict(array([9.0, 8.0])) + True + >>> model.k + 2 + >>> model.computeCost(sc.parallelize(data)) + 2.0 + >>> model = KMeans.train(sc.parallelize(data), 2) + >>> sparse_data = [ + ... SparseVector(3, {1: 1.0}), + ... SparseVector(3, {1: 1.1}), + ... SparseVector(3, {2: 1.0}), + ... SparseVector(3, {2: 1.1}) + ... ] + >>> model = KMeans.train(sc.parallelize(sparse_data), 2, initializationMode="k-means||", + ... seed=50, initializationSteps=5, epsilon=1e-4) + >>> model.predict(array([0., 1., 0.])) == model.predict(array([0, 1.1, 0.])) + True + >>> model.predict(array([0., 0., 1.])) == model.predict(array([0, 0, 1.1])) + True + >>> model.predict(sparse_data[0]) == model.predict(sparse_data[1]) + True + >>> model.predict(sparse_data[2]) == model.predict(sparse_data[3]) + True + >>> isinstance(model.clusterCenters, list) + True + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = KMeansModel.load(sc, path) + >>> sameModel.predict(sparse_data[0]) == model.predict(sparse_data[0]) + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + + >>> data = array([-383.1,-382.9, 28.7,31.2, 366.2,367.3]).reshape(3, 2) + >>> model = KMeans.train(sc.parallelize(data), 3, maxIterations=0, + ... initialModel = KMeansModel([(-1000.0,-1000.0),(5.0,5.0),(1000.0,1000.0)])) + >>> model.clusterCenters + [array([-1000., -1000.]), array([ 5., 5.]), array([ 1000., 1000.])] + """ + + def __init__(self, centers: List["VectorLike"]): + self.centers = centers + + @property + @since("1.0.0") + def clusterCenters(self) -> List["VectorLike"]: + """Get the cluster centers, represented as a list of NumPy arrays.""" + return self.centers + + @property + @since("1.4.0") + def k(self) -> int: + """Total number of clusters.""" + return len(self.centers) + + @overload + def predict(self, x: "VectorLike") -> int: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[int]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[int, RDD[int]]: + """ + Find the cluster that each of the points belongs to in this + model. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + A data point (or RDD of points) to determine cluster index. + :py:class:`pyspark.mllib.linalg.Vector` can be replaced with equivalent + objects (list, tuple, numpy.ndarray). + + Returns + ------- + int or :py:class:`pyspark.RDD` of int + Predicted cluster index or an RDD of predicted cluster indices + if the input is an RDD. + """ + best = 0 + best_distance = float("inf") + if isinstance(x, RDD): + return x.map(self.predict) + + x = _convert_to_vector(x) + for i in range(len(self.centers)): + distance = x.squared_distance(self.centers[i]) # type: ignore[attr-defined] + if distance < best_distance: + best = i + best_distance = distance + return best + + def computeCost(self, rdd: RDD["VectorLike"]) -> float: + """ + Return the K-means cost (sum of squared distances of points to + their nearest center) for this model on the given + data. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + rdd : ::py:class:`pyspark.RDD` + The RDD of points to compute the cost on. + """ + cost = callMLlibFunc( + "computeCostKmeansModel", + rdd.map(_convert_to_vector), + [_convert_to_vector(c) for c in self.centers], + ) + return cost + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """ + Save this model to the given path. + """ + assert sc._jvm is not None + + java_centers = _py2java(sc, [_convert_to_vector(c) for c in self.centers]) + java_model = sc._jvm.org.apache.spark.mllib.clustering.KMeansModel(java_centers) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "KMeansModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.clustering.KMeansModel.load(sc._jsc.sc(), path) + return KMeansModel(_java2py(sc, java_model.clusterCenters())) + + +class KMeans: + """ + K-means clustering. + + .. versionadded:: 0.9.0 + """ + + @classmethod + def train( + cls, + rdd: RDD["VectorLike"], + k: int, + maxIterations: int = 100, + initializationMode: str = "k-means||", + seed: Optional[int] = None, + initializationSteps: int = 2, + epsilon: float = 1e-4, + initialModel: Optional[KMeansModel] = None, + distanceMeasure: str = "euclidean", + ) -> "KMeansModel": + """ + Train a k-means clustering model. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + rdd : ::py:class:`pyspark.RDD` + Training points as an `RDD` of :py:class:`pyspark.mllib.linalg.Vector` + or convertible sequence types. + k : int + Number of clusters to create. + maxIterations : int, optional + Maximum number of iterations allowed. + (default: 100) + initializationMode : str, optional + The initialization algorithm. This can be either "random" or + "k-means||". + (default: "k-means||") + seed : int, optional + Random seed value for cluster initialization. Set as None to + generate seed based on system time. + (default: None) + initializationSteps : + Number of steps for the k-means|| initialization mode. + This is an advanced setting -- the default of 2 is almost + always enough. + (default: 2) + epsilon : float, optional + Distance threshold within which a center will be considered to + have converged. If all centers move less than this Euclidean + distance, iterations are stopped. + (default: 1e-4) + initialModel : :py:class:`KMeansModel`, optional + Initial cluster centers can be provided as a KMeansModel object + rather than using the random or k-means|| initializationModel. + (default: None) + distanceMeasure : str, optional + The distance measure used by the k-means algorithm. + (default: "euclidean") + """ + clusterInitialModel = [] + if initialModel is not None: + if not isinstance(initialModel, KMeansModel): + raise TypeError( + "initialModel is of " + str(type(initialModel)) + ". It needs " + "to be of " + ) + clusterInitialModel = [_convert_to_vector(c) for c in initialModel.clusterCenters] + model = callMLlibFunc( + "trainKMeansModel", + rdd.map(_convert_to_vector), + k, + maxIterations, + initializationMode, + seed, + initializationSteps, + epsilon, + clusterInitialModel, + distanceMeasure, + ) + centers = callJavaFunc(rdd.context, model.clusterCenters) + return KMeansModel([c.toArray() for c in centers]) + + +@inherit_doc +class GaussianMixtureModel(JavaModelWrapper, JavaSaveable, JavaLoader["GaussianMixtureModel"]): + + """ + A clustering model derived from the Gaussian Mixture Model method. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import Vectors, DenseMatrix + >>> from numpy.testing import assert_equal + >>> from shutil import rmtree + >>> import os, tempfile + + >>> clusterdata_1 = sc.parallelize(array([-0.1,-0.05,-0.01,-0.1, + ... 0.9,0.8,0.75,0.935, + ... -0.83,-0.68,-0.91,-0.76 ]).reshape(6, 2), 2) + >>> model = GaussianMixture.train(clusterdata_1, 3, convergenceTol=0.0001, + ... maxIterations=50, seed=10) + >>> labels = model.predict(clusterdata_1).collect() + >>> labels[0]==labels[1] + False + >>> labels[1]==labels[2] + False + >>> labels[4]==labels[5] + True + >>> model.predict([-0.1,-0.05]) + 0 + >>> softPredicted = model.predictSoft([-0.1,-0.05]) + >>> abs(softPredicted[0] - 1.0) < 0.03 + True + >>> abs(softPredicted[1] - 0.0) < 0.03 + True + >>> abs(softPredicted[2] - 0.0) < 0.03 + True + + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = GaussianMixtureModel.load(sc, path) + >>> assert_equal(model.weights, sameModel.weights) + >>> mus, sigmas = list( + ... zip(*[(g.mu, g.sigma) for g in model.gaussians])) + >>> sameMus, sameSigmas = list( + ... zip(*[(g.mu, g.sigma) for g in sameModel.gaussians])) + >>> mus == sameMus + True + >>> sigmas == sameSigmas + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + + >>> data = array([-5.1971, -2.5359, -3.8220, + ... -5.2211, -5.0602, 4.7118, + ... 6.8989, 3.4592, 4.6322, + ... 5.7048, 4.6567, 5.5026, + ... 4.5605, 5.2043, 6.2734]) + >>> clusterdata_2 = sc.parallelize(data.reshape(5,3)) + >>> model = GaussianMixture.train(clusterdata_2, 2, convergenceTol=0.0001, + ... maxIterations=150, seed=4) + >>> labels = model.predict(clusterdata_2).collect() + >>> labels[0]==labels[1] + True + >>> labels[2]==labels[3]==labels[4] + True + """ + + @property + @since("1.4.0") + def weights(self) -> np.ndarray: + """ + Weights for each Gaussian distribution in the mixture, where weights[i] is + the weight for Gaussian i, and weights.sum == 1. + """ + return array(self.call("weights")) + + @property + @since("1.4.0") + def gaussians(self) -> List[MultivariateGaussian]: + """ + Array of MultivariateGaussian where gaussians[i] represents + the Multivariate Gaussian (Normal) Distribution for Gaussian i. + """ + return [ + MultivariateGaussian(gaussian[0], gaussian[1]) for gaussian in self.call("gaussians") + ] + + @property + @since("1.4.0") + def k(self) -> int: + """Number of gaussians in mixture.""" + return len(self.weights) + + @overload + def predict(self, x: "VectorLike") -> np.int64: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[int]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[np.int64, RDD[int]]: + """ + Find the cluster to which the point 'x' or each point in RDD 'x' + has maximum membership in this model. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + A feature vector or an RDD of vectors representing data points. + + Returns + ------- + numpy.float64 or :py:class:`pyspark.RDD` of int + Predicted cluster label or an RDD of predicted cluster labels + if the input is an RDD. + """ + if isinstance(x, RDD): + cluster_labels = self.predictSoft(x).map(lambda z: z.index(max(z))) + return cluster_labels + else: + z = self.predictSoft(x) + return z.argmax() + + @overload + def predictSoft(self, x: "VectorLike") -> np.ndarray: + ... + + @overload + def predictSoft(self, x: RDD["VectorLike"]) -> RDD[pyarray.array]: + ... + + def predictSoft( + self, x: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[np.ndarray, RDD[pyarray.array]]: + """ + Find the membership of point 'x' or each point in RDD 'x' to all mixture components. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + A feature vector or an RDD of vectors representing data points. + + Returns + ------- + numpy.ndarray or :py:class:`pyspark.RDD` + The membership value to all mixture components for vector 'x' + or each vector in RDD 'x'. + """ + if isinstance(x, RDD): + means, sigmas = zip(*[(g.mu, g.sigma) for g in self.gaussians]) + membership_matrix = callMLlibFunc( + "predictSoftGMM", + x.map(_convert_to_vector), + _convert_to_vector(self.weights), + means, + sigmas, + ) + return membership_matrix.map(lambda x: pyarray.array("d", x)) + else: + return self.call("predictSoft", _convert_to_vector(x)).toArray() + + @classmethod + def load(cls, sc: SparkContext, path: str) -> "GaussianMixtureModel": + """Load the GaussianMixtureModel from disk. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + sc : :py:class:`SparkContext` + path : str + Path to where the model is stored. + """ + assert sc._jvm is not None + + model = cls._load_java(sc, path) + wrapper = sc._jvm.org.apache.spark.mllib.api.python.GaussianMixtureModelWrapper(model) + return cls(wrapper) + + +class GaussianMixture: + """ + Learning algorithm for Gaussian Mixtures using the expectation-maximization algorithm. + + .. versionadded:: 1.3.0 + """ + + @classmethod + def train( + cls, + rdd: RDD["VectorLike"], + k: int, + convergenceTol: float = 1e-3, + maxIterations: int = 100, + seed: Optional[int] = None, + initialModel: Optional[GaussianMixtureModel] = None, + ) -> GaussianMixtureModel: + """ + Train a Gaussian Mixture clustering model. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + rdd : ::py:class:`pyspark.RDD` + Training points as an `RDD` of :py:class:`pyspark.mllib.linalg.Vector` + or convertible sequence types. + k : int + Number of independent Gaussians in the mixture model. + convergenceTol : float, optional + Maximum change in log-likelihood at which convergence is + considered to have occurred. + (default: 1e-3) + maxIterations : int, optional + Maximum number of iterations allowed. + (default: 100) + seed : int, optional + Random seed for initial Gaussian distribution. Set as None to + generate seed based on system time. + (default: None) + initialModel : GaussianMixtureModel, optional + Initial GMM starting point, bypassing the random + initialization. + (default: None) + """ + initialModelWeights = None + initialModelMu = None + initialModelSigma = None + if initialModel is not None: + if initialModel.k != k: + raise ValueError( + "Mismatched cluster count, initialModel.k = %s, however k = %s" + % (initialModel.k, k) + ) + initialModelWeights = list(initialModel.weights) + initialModelMu = [initialModel.gaussians[i].mu for i in range(initialModel.k)] + initialModelSigma = [initialModel.gaussians[i].sigma for i in range(initialModel.k)] + java_model = callMLlibFunc( + "trainGaussianMixtureModel", + rdd.map(_convert_to_vector), + k, + convergenceTol, + maxIterations, + seed, + initialModelWeights, + initialModelMu, + initialModelSigma, + ) + return GaussianMixtureModel(java_model) + + +class PowerIterationClusteringModel( + JavaModelWrapper, JavaSaveable, JavaLoader["PowerIterationClusteringModel"] +): + + """ + Model produced by :py:class:`PowerIterationClustering`. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> import math + >>> def genCircle(r, n): + ... points = [] + ... for i in range(0, n): + ... theta = 2.0 * math.pi * i / n + ... points.append((r * math.cos(theta), r * math.sin(theta))) + ... return points + ... + >>> def sim(x, y): + ... dist2 = (x[0] - y[0]) * (x[0] - y[0]) + (x[1] - y[1]) * (x[1] - y[1]) + ... return math.exp(-dist2 / 2.0) + ... + >>> r1 = 1.0 + >>> n1 = 10 + >>> r2 = 4.0 + >>> n2 = 40 + >>> n = n1 + n2 + >>> points = genCircle(r1, n1) + genCircle(r2, n2) + >>> similarities = [(i, j, sim(points[i], points[j])) for i in range(1, n) for j in range(0, i)] + >>> rdd = sc.parallelize(similarities, 2) + >>> model = PowerIterationClustering.train(rdd, 2, 40) + >>> model.k + 2 + >>> result = sorted(model.assignments().collect(), key=lambda x: x.id) + >>> result[0].cluster == result[1].cluster == result[2].cluster == result[3].cluster + True + >>> result[4].cluster == result[5].cluster == result[6].cluster == result[7].cluster + True + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = PowerIterationClusteringModel.load(sc, path) + >>> sameModel.k + 2 + >>> result = sorted(model.assignments().collect(), key=lambda x: x.id) + >>> result[0].cluster == result[1].cluster == result[2].cluster == result[3].cluster + True + >>> result[4].cluster == result[5].cluster == result[6].cluster == result[7].cluster + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + @property + @since("1.5.0") + def k(self) -> int: + """ + Returns the number of clusters. + """ + return self.call("k") + + @since("1.5.0") + def assignments(self) -> RDD["PowerIterationClustering.Assignment"]: + """ + Returns the cluster assignments of this model. + """ + return self.call("getAssignments").map(lambda x: (PowerIterationClustering.Assignment(*x))) + + @classmethod + @since("1.5.0") + def load(cls, sc: SparkContext, path: str) -> "PowerIterationClusteringModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + model = cls._load_java(sc, path) + wrapper = sc._jvm.org.apache.spark.mllib.api.python.PowerIterationClusteringModelWrapper( + model + ) + return PowerIterationClusteringModel(wrapper) + + +class PowerIterationClustering: + """ + Power Iteration Clustering (PIC), a scalable graph clustering algorithm. + + + Developed by Lin and Cohen [1]_. From the abstract: + + "PIC finds a very low-dimensional embedding of a + dataset using truncated power iteration on a normalized pair-wise + similarity matrix of the data." + + .. versionadded:: 1.5.0 + + .. [1] Lin, Frank & Cohen, William. (2010). Power Iteration Clustering. + http://www.cs.cmu.edu/~frank/papers/icml2010-pic-final.pdf + """ + + @classmethod + def train( + cls, + rdd: RDD[Tuple[int, int, float]], + k: int, + maxIterations: int = 100, + initMode: str = "random", + ) -> PowerIterationClusteringModel: + r""" + Train PowerIterationClusteringModel + + .. versionadded:: 1.5.0 + + Parameters + ---------- + rdd : :py:class:`pyspark.RDD` + An RDD of (i, j, s\ :sub:`ij`\) tuples representing the + affinity matrix, which is the matrix A in the PIC paper. The + similarity s\ :sub:`ij`\ must be nonnegative. This is a symmetric + matrix and hence s\ :sub:`ij`\ = s\ :sub:`ji`\ For any (i, j) with + nonzero similarity, there should be either (i, j, s\ :sub:`ij`\) or + (j, i, s\ :sub:`ji`\) in the input. Tuples with i = j are ignored, + because it is assumed s\ :sub:`ij`\ = 0.0. + k : int + Number of clusters. + maxIterations : int, optional + Maximum number of iterations of the PIC algorithm. + (default: 100) + initMode : str, optional + Initialization mode. This can be either "random" to use + a random vector as vertex properties, or "degree" to use + normalized sum similarities. + (default: "random") + """ + model = callMLlibFunc( + "trainPowerIterationClusteringModel", + rdd.map(_convert_to_vector), + int(k), + int(maxIterations), + initMode, + ) + return PowerIterationClusteringModel(model) + + class Assignment(namedtuple("Assignment", ["id", "cluster"])): + """ + Represents an (id, cluster) tuple. + + .. versionadded:: 1.5.0 + """ + + +class StreamingKMeansModel(KMeansModel): + """ + Clustering model which can perform an online update of the centroids. + + The update formula for each centroid is given by + + - c_t+1 = ((c_t * n_t * a) + (x_t * m_t)) / (n_t + m_t) + - n_t+1 = n_t * a + m_t + + where + + - c_t: Centroid at the n_th iteration. + - n_t: Number of samples (or) weights associated with the centroid + at the n_th iteration. + - x_t: Centroid of the new data closest to c_t. + - m_t: Number of samples (or) weights of the new data closest to c_t + - c_t+1: New centroid. + - n_t+1: New number of weights. + - a: Decay Factor, which gives the forgetfulness. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + clusterCenters : list of :py:class:`pyspark.mllib.linalg.Vector` or covertible + Initial cluster centers. + clusterWeights : :py:class:`pyspark.mllib.linalg.Vector` or covertible + List of weights assigned to each cluster. + + Notes + ----- + If a is set to 1, it is the weighted mean of the previous + and new data. If it set to zero, the old centroids are completely + forgotten. + + Examples + -------- + >>> initCenters = [[0.0, 0.0], [1.0, 1.0]] + >>> initWeights = [1.0, 1.0] + >>> stkm = StreamingKMeansModel(initCenters, initWeights) + >>> data = sc.parallelize([[-0.1, -0.1], [0.1, 0.1], + ... [0.9, 0.9], [1.1, 1.1]]) + >>> stkm = stkm.update(data, 1.0, "batches") + >>> stkm.centers + array([[ 0., 0.], + [ 1., 1.]]) + >>> stkm.predict([-0.1, -0.1]) + 0 + >>> stkm.predict([0.9, 0.9]) + 1 + >>> stkm.clusterWeights + [3.0, 3.0] + >>> decayFactor = 0.0 + >>> data = sc.parallelize([DenseVector([1.5, 1.5]), DenseVector([0.2, 0.2])]) + >>> stkm = stkm.update(data, 0.0, "batches") + >>> stkm.centers + array([[ 0.2, 0.2], + [ 1.5, 1.5]]) + >>> stkm.clusterWeights + [1.0, 1.0] + >>> stkm.predict([0.2, 0.2]) + 0 + >>> stkm.predict([1.5, 1.5]) + 1 + """ + + def __init__(self, clusterCenters: List["VectorLike"], clusterWeights: "VectorLike"): + super(StreamingKMeansModel, self).__init__(centers=clusterCenters) + self._clusterWeights = list(clusterWeights) # type: ignore[arg-type] + + @property + @since("1.5.0") + def clusterWeights(self) -> List[np.float64]: + """Return the cluster weights.""" + return self._clusterWeights + + @since("1.5.0") + def update( + self, data: RDD["VectorLike"], decayFactor: float, timeUnit: str + ) -> "StreamingKMeansModel": + """Update the centroids, according to data + + .. versionadded:: 1.5.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + RDD with new data for the model update. + decayFactor : float + Forgetfulness of the previous centroids. + timeUnit : str + Can be "batches" or "points". If points, then the decay factor + is raised to the power of number of new points and if batches, + then decay factor will be used as is. + """ + if not isinstance(data, RDD): + raise TypeError("Data should be of an RDD, got %s." % type(data)) + data = data.map(_convert_to_vector) + decayFactor = float(decayFactor) + if timeUnit not in ["batches", "points"]: + raise ValueError("timeUnit should be 'batches' or 'points', got %s." % timeUnit) + vectorCenters = [_convert_to_vector(center) for center in self.centers] + updatedModel = callMLlibFunc( + "updateStreamingKMeansModel", + vectorCenters, + self._clusterWeights, + data, + decayFactor, + timeUnit, + ) + self.centers = array(updatedModel[0]) # type: ignore[assignment] + self._clusterWeights = list(updatedModel[1]) + return self + + +class StreamingKMeans: + """ + Provides methods to set k, decayFactor, timeUnit to configure the + KMeans algorithm for fitting and predicting on incoming dstreams. + More details on how the centroids are updated are provided under the + docs of StreamingKMeansModel. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + k : int, optional + Number of clusters. + (default: 2) + decayFactor : float, optional + Forgetfulness of the previous centroids. + (default: 1.0) + timeUnit : str, optional + Can be "batches" or "points". If points, then the decay factor is + raised to the power of number of new points and if batches, then + decay factor will be used as is. + (default: "batches") + """ + + def __init__(self, k: int = 2, decayFactor: float = 1.0, timeUnit: str = "batches"): + self._k = k + self._decayFactor = decayFactor + if timeUnit not in ["batches", "points"]: + raise ValueError("timeUnit should be 'batches' or 'points', got %s." % timeUnit) + self._timeUnit = timeUnit + self._model: Optional[StreamingKMeansModel] = None + + @since("1.5.0") + def latestModel(self) -> Optional[StreamingKMeansModel]: + """Return the latest model""" + return self._model + + def _validate(self, dstream: Any) -> None: + if self._model is None: + raise ValueError( + "Initial centers should be set either by setInitialCenters " "or setRandomCenters." + ) + if not isinstance(dstream, DStream): + raise TypeError( + "Expected dstream to be of type DStream, " "got type %s" % type(dstream) + ) + + @since("1.5.0") + def setK(self, k: int) -> "StreamingKMeans": + """Set number of clusters.""" + self._k = k + return self + + @since("1.5.0") + def setDecayFactor(self, decayFactor: float) -> "StreamingKMeans": + """Set decay factor.""" + self._decayFactor = decayFactor + return self + + @since("1.5.0") + def setHalfLife(self, halfLife: float, timeUnit: str) -> "StreamingKMeans": + """ + Set number of batches after which the centroids of that + particular batch has half the weightage. + """ + self._timeUnit = timeUnit + self._decayFactor = exp(log(0.5) / halfLife) + return self + + @since("1.5.0") + def setInitialCenters( + self, centers: List["VectorLike"], weights: List[float] + ) -> "StreamingKMeans": + """ + Set initial centers. Should be set before calling trainOn. + """ + self._model = StreamingKMeansModel(centers, weights) + return self + + @since("1.5.0") + def setRandomCenters(self, dim: int, weight: float, seed: int) -> "StreamingKMeans": + """ + Set the initial centers to be random samples from + a gaussian population with constant weights. + """ + rng = random.RandomState(seed) + clusterCenters = rng.randn(self._k, dim) + clusterWeights = tile(weight, self._k) + self._model = StreamingKMeansModel(clusterCenters, clusterWeights) # type: ignore[arg-type] + return self + + @since("1.5.0") + def trainOn(self, dstream: "DStream[VectorLike]") -> None: + """Train the model on the incoming dstream.""" + self._validate(dstream) + + def update(rdd: RDD["VectorLike"]) -> None: + self._model.update(rdd, self._decayFactor, self._timeUnit) # type: ignore[union-attr] + + dstream.foreachRDD(update) + + @since("1.5.0") + def predictOn(self, dstream: "DStream[VectorLike]") -> "DStream[int]": + """ + Make predictions on a dstream. + Returns a transformed dstream object + """ + self._validate(dstream) + return dstream.map(lambda x: self._model.predict(x)) # type: ignore[union-attr] + + @since("1.5.0") + def predictOnValues(self, dstream: "DStream[Tuple[T, VectorLike]]") -> "DStream[Tuple[T, int]]": + """ + Make predictions on a keyed dstream. + Returns a transformed dstream object. + """ + self._validate(dstream) + return dstream.mapValues(lambda x: self._model.predict(x)) # type: ignore[union-attr] + + +class LDAModel(JavaModelWrapper, JavaSaveable, Loader["LDAModel"]): + + """A clustering model derived from the LDA method. + + Latent Dirichlet Allocation (LDA), a topic model designed for text documents. + Terminology + + - "word" = "term": an element of the vocabulary + - "token": instance of a term appearing in a document + - "topic": multinomial distribution over words representing some concept + + .. versionadded:: 1.5.0 + + Notes + ----- + See the original LDA paper (journal version) [1]_ + + .. [1] Blei, D. et al. "Latent Dirichlet Allocation." + J. Mach. Learn. Res. 3 (2003): 993-1022. + https://www.jmlr.org/papers/v3/blei03a + + Examples + -------- + >>> from pyspark.mllib.linalg import Vectors + >>> from numpy.testing import assert_almost_equal, assert_equal + >>> data = [ + ... [1, Vectors.dense([0.0, 1.0])], + ... [2, SparseVector(2, {0: 1.0})], + ... ] + >>> rdd = sc.parallelize(data) + >>> model = LDA.train(rdd, k=2, seed=1) + >>> model.vocabSize() + 2 + >>> model.describeTopics() + [([1, 0], [0.5..., 0.49...]), ([0, 1], [0.5..., 0.49...])] + >>> model.describeTopics(1) + [([1], [0.5...]), ([0], [0.5...])] + + >>> topics = model.topicsMatrix() + >>> topics_expect = array([[0.5, 0.5], [0.5, 0.5]]) + >>> assert_almost_equal(topics, topics_expect, 1) + + >>> import os, tempfile + >>> from shutil import rmtree + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = LDAModel.load(sc, path) + >>> assert_equal(sameModel.topicsMatrix(), model.topicsMatrix()) + >>> sameModel.vocabSize() == model.vocabSize() + True + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + @since("1.5.0") + def topicsMatrix(self) -> np.ndarray: + """Inferred topics, where each topic is represented by a distribution over terms.""" + return self.call("topicsMatrix").toArray() + + @since("1.5.0") + def vocabSize(self) -> int: + """Vocabulary size (number of terms or terms in the vocabulary)""" + return self.call("vocabSize") + + def describeTopics( + self, maxTermsPerTopic: Optional[int] = None + ) -> List[Tuple[List[int], List[float]]]: + """Return the topics described by weighted terms. + + .. versionadded:: 1.6.0 + .. warning:: If vocabSize and k are large, this can return a large object! + + Parameters + ---------- + maxTermsPerTopic : int, optional + Maximum number of terms to collect for each topic. + (default: vocabulary size) + + Returns + ------- + list + Array over topics. Each topic is represented as a pair of + matching arrays: (term indices, term weights in topic). + Each topic's terms are sorted in order of decreasing weight. + """ + if maxTermsPerTopic is None: + topics = self.call("describeTopics") + else: + topics = self.call("describeTopics", maxTermsPerTopic) + return topics + + @classmethod + def load(cls, sc: SparkContext, path: str) -> "LDAModel": + """Load the LDAModel from disk. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + path : str + Path to where the model is stored. + """ + if not isinstance(sc, SparkContext): + raise TypeError("sc should be a SparkContext, got type %s" % type(sc)) + if not isinstance(path, str): + raise TypeError("path should be a string, got type %s" % type(path)) + model = callMLlibFunc("loadLDAModel", sc, path) + return LDAModel(model) + + +class LDA: + """ + Train Latent Dirichlet Allocation (LDA) model. + + .. versionadded:: 1.5.0 + """ + + @classmethod + def train( + cls, + rdd: RDD[Tuple[int, "VectorLike"]], + k: int = 10, + maxIterations: int = 20, + docConcentration: float = -1.0, + topicConcentration: float = -1.0, + seed: Optional[int] = None, + checkpointInterval: int = 10, + optimizer: str = "em", + ) -> LDAModel: + """Train a LDA model. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + rdd : :py:class:`pyspark.RDD` + RDD of documents, which are tuples of document IDs and term + (word) count vectors. The term count vectors are "bags of + words" with a fixed-size vocabulary (where the vocabulary size + is the length of the vector). Document IDs must be unique + and >= 0. + k : int, optional + Number of topics to infer, i.e., the number of soft cluster + centers. + (default: 10) + maxIterations : int, optional + Maximum number of iterations allowed. + (default: 20) + docConcentration : float, optional + Concentration parameter (commonly named "alpha") for the prior + placed on documents' distributions over topics ("theta"). + (default: -1.0) + topicConcentration : float, optional + Concentration parameter (commonly named "beta" or "eta") for + the prior placed on topics' distributions over terms. + (default: -1.0) + seed : int, optional + Random seed for cluster initialization. Set as None to generate + seed based on system time. + (default: None) + checkpointInterval : int, optional + Period (in iterations) between checkpoints. + (default: 10) + optimizer : str, optional + LDAOptimizer used to perform the actual calculation. Currently + "em", "online" are supported. + (default: "em") + """ + model = callMLlibFunc( + "trainLDAModel", + rdd, + k, + maxIterations, + docConcentration, + topicConcentration, + seed, + checkpointInterval, + optimizer, + ) + return LDAModel(model) + + +def _test() -> None: + import doctest + import numpy + import pyspark.mllib.clustering + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + globs = pyspark.mllib.clustering.__dict__.copy() + globs["sc"] = SparkContext("local[4]", "PythonTest", batchSize=2) + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/common.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/common.py new file mode 100644 index 0000000..c5e1a7e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/common.py @@ -0,0 +1,174 @@ +# +# 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 Any, Callable, TYPE_CHECKING + +if TYPE_CHECKING: + from pyspark.mllib._typing import C, JavaObjectOrPickleDump + +import py4j.protocol +from py4j.protocol import Py4JJavaError +from py4j.java_gateway import JavaObject +from py4j.java_collections import JavaArray, JavaList + +import pyspark.context +from pyspark import RDD, SparkContext +from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer +from pyspark.sql import DataFrame, SparkSession + +# Hack for support float('inf') in Py4j +_old_smart_decode = py4j.protocol.smart_decode + +_float_str_mapping = { + "nan": "NaN", + "inf": "Infinity", + "-inf": "-Infinity", +} + + +def _new_smart_decode(obj: Any) -> str: + if isinstance(obj, float): + s = str(obj) + return _float_str_mapping.get(s, s) + return _old_smart_decode(obj) + + +py4j.protocol.smart_decode = _new_smart_decode + + +_picklable_classes = [ + "LinkedList", + "SparseVector", + "DenseVector", + "DenseMatrix", + "Rating", + "LabeledPoint", +] + + +# this will call the MLlib version of pythonToJava() +def _to_java_object_rdd(rdd: RDD) -> JavaObject: + """Return a JavaRDD of Object by unpickling + + It will convert each Python object into Java object by Pickle, whenever the + RDD is serialized in batch or not. + """ + rdd = rdd._reserialize(AutoBatchedSerializer(CPickleSerializer())) + assert rdd.ctx._jvm is not None + return rdd.ctx._jvm.org.apache.spark.mllib.api.python.SerDe.pythonToJava(rdd._jrdd, True) + + +def _py2java(sc: SparkContext, obj: Any) -> JavaObject: + """Convert Python object into Java""" + if isinstance(obj, RDD): + obj = _to_java_object_rdd(obj) + elif isinstance(obj, DataFrame): + obj = obj._jdf + elif isinstance(obj, SparkContext): + obj = obj._jsc + elif isinstance(obj, list): + obj = [_py2java(sc, x) for x in obj] + elif isinstance(obj, JavaObject): + pass + elif isinstance(obj, (int, float, bool, bytes, str)): + pass + else: + data = bytearray(CPickleSerializer().dumps(obj)) + assert sc._jvm is not None + obj = sc._jvm.org.apache.spark.mllib.api.python.SerDe.loads(data) + return obj + + +def _java2py(sc: SparkContext, r: "JavaObjectOrPickleDump", encoding: str = "bytes") -> Any: + if isinstance(r, JavaObject): + clsName = r.getClass().getSimpleName() + # convert RDD into JavaRDD + if clsName != "JavaRDD" and clsName.endswith("RDD"): + r = r.toJavaRDD() + clsName = "JavaRDD" + + assert sc._jvm is not None + + if clsName == "JavaRDD": + jrdd = sc._jvm.org.apache.spark.mllib.api.python.SerDe.javaToPython(r) + return RDD(jrdd, sc) + + if clsName == "Dataset": + return DataFrame(r, SparkSession._getActiveSessionOrCreate()) + + if clsName in _picklable_classes: + r = sc._jvm.org.apache.spark.mllib.api.python.SerDe.dumps(r) + elif isinstance(r, (JavaArray, JavaList)): + try: + r = sc._jvm.org.apache.spark.mllib.api.python.SerDe.dumps(r) + except Py4JJavaError: + pass # not pickable + + if isinstance(r, (bytearray, bytes)): + r = CPickleSerializer().loads(bytes(r), encoding=encoding) + return r + + +def callJavaFunc( + sc: pyspark.context.SparkContext, func: Callable[..., "JavaObjectOrPickleDump"], *args: Any +) -> Any: + """Call Java Function""" + java_args = [_py2java(sc, a) for a in args] + return _java2py(sc, func(*java_args)) + + +def callMLlibFunc(name: str, *args: Any) -> Any: + """Call API in PythonMLLibAPI""" + sc = SparkContext.getOrCreate() + assert sc._jvm is not None + api = getattr(sc._jvm.PythonMLLibAPI(), name) + return callJavaFunc(sc, api, *args) + + +class JavaModelWrapper: + """ + Wrapper for the model in JVM + """ + + def __init__(self, java_model: JavaObject): + self._sc = SparkContext.getOrCreate() + self._java_model = java_model + + def __del__(self) -> None: + assert self._sc._gateway is not None + self._sc._gateway.detach(self._java_model) + + def call(self, name: str, *a: Any) -> Any: + """Call method of java_model""" + return callJavaFunc(self._sc, getattr(self._java_model, name), *a) + + +def inherit_doc(cls: "C") -> "C": + """ + A decorator that makes a class inherit documentation from its parents. + """ + for name, func in vars(cls).items(): + # only inherit docstring for public functions + if name.startswith("_"): + continue + if not func.__doc__: + for parent in cls.__bases__: + parent_func = getattr(parent, name, None) + if parent_func and getattr(parent_func, "__doc__", None): + func.__doc__ = parent_func.__doc__ + break + return cls diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/evaluation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/evaluation.py new file mode 100644 index 0000000..2a8991d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/evaluation.py @@ -0,0 +1,691 @@ +# +# 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 Generic, List, Optional, Tuple, TypeVar, Union + +import sys + +from pyspark import since +from pyspark.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc +from pyspark.mllib.linalg import Matrix +from pyspark.sql import SQLContext +from pyspark.sql.types import ArrayType, DoubleType, StructField, StructType + +__all__ = [ + "BinaryClassificationMetrics", + "RegressionMetrics", + "MulticlassMetrics", + "RankingMetrics", +] + +T = TypeVar("T") + + +class BinaryClassificationMetrics(JavaModelWrapper): + """ + Evaluator for binary classification. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + scoreAndLabels : :py:class:`pyspark.RDD` + an RDD of score, label and optional weight. + + Examples + -------- + >>> scoreAndLabels = sc.parallelize([ + ... (0.1, 0.0), (0.1, 1.0), (0.4, 0.0), (0.6, 0.0), (0.6, 1.0), (0.6, 1.0), (0.8, 1.0)], 2) + >>> metrics = BinaryClassificationMetrics(scoreAndLabels) + >>> metrics.areaUnderROC + 0.70... + >>> metrics.areaUnderPR + 0.83... + >>> metrics.unpersist() + >>> scoreAndLabelsWithOptWeight = sc.parallelize([ + ... (0.1, 0.0, 1.0), (0.1, 1.0, 0.4), (0.4, 0.0, 0.2), (0.6, 0.0, 0.6), (0.6, 1.0, 0.9), + ... (0.6, 1.0, 0.5), (0.8, 1.0, 0.7)], 2) + >>> metrics = BinaryClassificationMetrics(scoreAndLabelsWithOptWeight) + >>> metrics.areaUnderROC + 0.79... + >>> metrics.areaUnderPR + 0.88... + """ + + def __init__(self, scoreAndLabels: RDD[Tuple[float, float]]): + sc = scoreAndLabels.ctx + sql_ctx = SQLContext.getOrCreate(sc) + numCol = len(scoreAndLabels.first()) + schema = StructType( + [ + StructField("score", DoubleType(), nullable=False), + StructField("label", DoubleType(), nullable=False), + ] + ) + if numCol == 3: + schema.add("weight", DoubleType(), False) + df = sql_ctx.createDataFrame(scoreAndLabels, schema=schema) + assert sc._jvm is not None + java_class = sc._jvm.org.apache.spark.mllib.evaluation.BinaryClassificationMetrics + java_model = java_class(df._jdf) + super(BinaryClassificationMetrics, self).__init__(java_model) + + @property + @since("1.4.0") + def areaUnderROC(self) -> float: + """ + Computes the area under the receiver operating characteristic + (ROC) curve. + """ + return self.call("areaUnderROC") + + @property + @since("1.4.0") + def areaUnderPR(self) -> float: + """ + Computes the area under the precision-recall curve. + """ + return self.call("areaUnderPR") + + @since("1.4.0") + def unpersist(self) -> None: + """ + Unpersists intermediate RDDs used in the computation. + """ + self.call("unpersist") + + +class RegressionMetrics(JavaModelWrapper): + """ + Evaluator for regression. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + predictionAndObservations : :py:class:`pyspark.RDD` + an RDD of prediction, observation and optional weight. + + Examples + -------- + >>> predictionAndObservations = sc.parallelize([ + ... (2.5, 3.0), (0.0, -0.5), (2.0, 2.0), (8.0, 7.0)]) + >>> metrics = RegressionMetrics(predictionAndObservations) + >>> metrics.explainedVariance + 8.859... + >>> metrics.meanAbsoluteError + 0.5... + >>> metrics.meanSquaredError + 0.37... + >>> metrics.rootMeanSquaredError + 0.61... + >>> metrics.r2 + 0.94... + >>> predictionAndObservationsWithOptWeight = sc.parallelize([ + ... (2.5, 3.0, 0.5), (0.0, -0.5, 1.0), (2.0, 2.0, 0.3), (8.0, 7.0, 0.9)]) + >>> metrics = RegressionMetrics(predictionAndObservationsWithOptWeight) + >>> metrics.rootMeanSquaredError + 0.68... + """ + + def __init__(self, predictionAndObservations: RDD[Tuple[float, float]]): + sc = predictionAndObservations.ctx + sql_ctx = SQLContext.getOrCreate(sc) + numCol = len(predictionAndObservations.first()) + schema = StructType( + [ + StructField("prediction", DoubleType(), nullable=False), + StructField("observation", DoubleType(), nullable=False), + ] + ) + if numCol == 3: + schema.add("weight", DoubleType(), False) + df = sql_ctx.createDataFrame(predictionAndObservations, schema=schema) + assert sc._jvm is not None + java_class = sc._jvm.org.apache.spark.mllib.evaluation.RegressionMetrics + java_model = java_class(df._jdf) + super(RegressionMetrics, self).__init__(java_model) + + @property + @since("1.4.0") + def explainedVariance(self) -> float: + r""" + Returns the explained variance regression score. + explainedVariance = :math:`1 - \frac{variance(y - \hat{y})}{variance(y)}` + """ + return self.call("explainedVariance") + + @property + @since("1.4.0") + def meanAbsoluteError(self) -> float: + """ + Returns the mean absolute error, which is a risk function corresponding to the + expected value of the absolute error loss or l1-norm loss. + """ + return self.call("meanAbsoluteError") + + @property + @since("1.4.0") + def meanSquaredError(self) -> float: + """ + Returns the mean squared error, which is a risk function corresponding to the + expected value of the squared error loss or quadratic loss. + """ + return self.call("meanSquaredError") + + @property + @since("1.4.0") + def rootMeanSquaredError(self) -> float: + """ + Returns the root mean squared error, which is defined as the square root of + the mean squared error. + """ + return self.call("rootMeanSquaredError") + + @property + @since("1.4.0") + def r2(self) -> float: + """ + Returns R^2^, the coefficient of determination. + """ + return self.call("r2") + + +class MulticlassMetrics(JavaModelWrapper): + """ + Evaluator for multiclass classification. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + predictionAndLabels : :py:class:`pyspark.RDD` + an RDD of prediction, label, optional weight and optional probability. + + Examples + -------- + >>> predictionAndLabels = sc.parallelize([(0.0, 0.0), (0.0, 1.0), (0.0, 0.0), + ... (1.0, 0.0), (1.0, 1.0), (1.0, 1.0), (1.0, 1.0), (2.0, 2.0), (2.0, 0.0)]) + >>> metrics = MulticlassMetrics(predictionAndLabels) + >>> metrics.confusionMatrix().toArray() + array([[ 2., 1., 1.], + [ 1., 3., 0.], + [ 0., 0., 1.]]) + >>> metrics.falsePositiveRate(0.0) + 0.2... + >>> metrics.precision(1.0) + 0.75... + >>> metrics.recall(2.0) + 1.0... + >>> metrics.fMeasure(0.0, 2.0) + 0.52... + >>> metrics.accuracy + 0.66... + >>> metrics.weightedFalsePositiveRate + 0.19... + >>> metrics.weightedPrecision + 0.68... + >>> metrics.weightedRecall + 0.66... + >>> metrics.weightedFMeasure() + 0.66... + >>> metrics.weightedFMeasure(2.0) + 0.65... + >>> predAndLabelsWithOptWeight = sc.parallelize([(0.0, 0.0, 1.0), (0.0, 1.0, 1.0), + ... (0.0, 0.0, 1.0), (1.0, 0.0, 1.0), (1.0, 1.0, 1.0), (1.0, 1.0, 1.0), (1.0, 1.0, 1.0), + ... (2.0, 2.0, 1.0), (2.0, 0.0, 1.0)]) + >>> metrics = MulticlassMetrics(predAndLabelsWithOptWeight) + >>> metrics.confusionMatrix().toArray() + array([[ 2., 1., 1.], + [ 1., 3., 0.], + [ 0., 0., 1.]]) + >>> metrics.falsePositiveRate(0.0) + 0.2... + >>> metrics.precision(1.0) + 0.75... + >>> metrics.recall(2.0) + 1.0... + >>> metrics.fMeasure(0.0, 2.0) + 0.52... + >>> metrics.accuracy + 0.66... + >>> metrics.weightedFalsePositiveRate + 0.19... + >>> metrics.weightedPrecision + 0.68... + >>> metrics.weightedRecall + 0.66... + >>> metrics.weightedFMeasure() + 0.66... + >>> metrics.weightedFMeasure(2.0) + 0.65... + >>> predictionAndLabelsWithProbabilities = sc.parallelize([ + ... (1.0, 1.0, 1.0, [0.1, 0.8, 0.1]), (0.0, 2.0, 1.0, [0.9, 0.05, 0.05]), + ... (0.0, 0.0, 1.0, [0.8, 0.2, 0.0]), (1.0, 1.0, 1.0, [0.3, 0.65, 0.05])]) + >>> metrics = MulticlassMetrics(predictionAndLabelsWithProbabilities) + >>> metrics.logLoss() + 0.9682... + """ + + def __init__(self, predictionAndLabels: RDD[Tuple[float, float]]): + sc = predictionAndLabels.ctx + sql_ctx = SQLContext.getOrCreate(sc) + numCol = len(predictionAndLabels.first()) + schema = StructType( + [ + StructField("prediction", DoubleType(), nullable=False), + StructField("label", DoubleType(), nullable=False), + ] + ) + if numCol >= 3: + schema.add("weight", DoubleType(), False) + if numCol == 4: + schema.add("probability", ArrayType(DoubleType(), False), False) + df = sql_ctx.createDataFrame(predictionAndLabels, schema) + assert sc._jvm is not None + java_class = sc._jvm.org.apache.spark.mllib.evaluation.MulticlassMetrics + java_model = java_class(df._jdf) + super(MulticlassMetrics, self).__init__(java_model) + + @since("1.4.0") + def confusionMatrix(self) -> Matrix: + """ + Returns confusion matrix: predicted classes are in columns, + they are ordered by class label ascending, as in "labels". + """ + return self.call("confusionMatrix") + + @since("1.4.0") + def truePositiveRate(self, label: float) -> float: + """ + Returns true positive rate for a given label (category). + """ + return self.call("truePositiveRate", label) + + @since("1.4.0") + def falsePositiveRate(self, label: float) -> float: + """ + Returns false positive rate for a given label (category). + """ + return self.call("falsePositiveRate", label) + + @since("1.4.0") + def precision(self, label: float) -> float: + """ + Returns precision. + """ + return self.call("precision", float(label)) + + @since("1.4.0") + def recall(self, label: float) -> float: + """ + Returns recall. + """ + return self.call("recall", float(label)) + + @since("1.4.0") + def fMeasure(self, label: float, beta: Optional[float] = None) -> float: + """ + Returns f-measure. + """ + if beta is None: + return self.call("fMeasure", label) + else: + return self.call("fMeasure", label, beta) + + @property + @since("2.0.0") + def accuracy(self) -> float: + """ + Returns accuracy (equals to the total number of correctly classified instances + out of the total number of instances). + """ + return self.call("accuracy") + + @property + @since("1.4.0") + def weightedTruePositiveRate(self) -> float: + """ + Returns weighted true positive rate. + (equals to precision, recall and f-measure) + """ + return self.call("weightedTruePositiveRate") + + @property + @since("1.4.0") + def weightedFalsePositiveRate(self) -> float: + """ + Returns weighted false positive rate. + """ + return self.call("weightedFalsePositiveRate") + + @property + @since("1.4.0") + def weightedRecall(self) -> float: + """ + Returns weighted averaged recall. + (equals to precision, recall and f-measure) + """ + return self.call("weightedRecall") + + @property + @since("1.4.0") + def weightedPrecision(self) -> float: + """ + Returns weighted averaged precision. + """ + return self.call("weightedPrecision") + + @since("1.4.0") + def weightedFMeasure(self, beta: Optional[float] = None) -> float: + """ + Returns weighted averaged f-measure. + """ + if beta is None: + return self.call("weightedFMeasure") + else: + return self.call("weightedFMeasure", beta) + + @since("3.0.0") + def logLoss(self, eps: float = 1e-15) -> float: + """ + Returns weighted logLoss. + """ + return self.call("logLoss", eps) + + +class RankingMetrics(JavaModelWrapper, Generic[T]): + """ + Evaluator for ranking algorithms. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + predictionAndLabels : :py:class:`pyspark.RDD` + an RDD of (predicted ranking, ground truth set) pairs + or (predicted ranking, ground truth set, + relevance value of ground truth set). + Since 3.4.0, it supports ndcg evaluation with relevance value. + + Examples + -------- + >>> predictionAndLabels = sc.parallelize([ + ... ([1, 6, 2, 7, 8, 3, 9, 10, 4, 5], [1, 2, 3, 4, 5]), + ... ([4, 1, 5, 6, 2, 7, 3, 8, 9, 10], [1, 2, 3]), + ... ([1, 2, 3, 4, 5], [])]) + >>> metrics = RankingMetrics(predictionAndLabels) + >>> metrics.precisionAt(1) + 0.33... + >>> metrics.precisionAt(5) + 0.26... + >>> metrics.precisionAt(15) + 0.17... + >>> metrics.meanAveragePrecision + 0.35... + >>> metrics.meanAveragePrecisionAt(1) + 0.3333333333333333... + >>> metrics.meanAveragePrecisionAt(2) + 0.25... + >>> metrics.ndcgAt(3) + 0.33... + >>> metrics.ndcgAt(10) + 0.48... + >>> metrics.recallAt(1) + 0.06... + >>> metrics.recallAt(5) + 0.35... + >>> metrics.recallAt(15) + 0.66... + """ + + def __init__( + self, + predictionAndLabels: Union[ + RDD[Tuple[List[T], List[T]]], RDD[Tuple[List[T], List[T], List[float]]] + ], + ): + sc = predictionAndLabels.ctx + sql_ctx = SQLContext.getOrCreate(sc) + df = sql_ctx.createDataFrame( + predictionAndLabels, schema=sql_ctx._inferSchema(predictionAndLabels) + ) + java_model = callMLlibFunc("newRankingMetrics", df._jdf) + super(RankingMetrics, self).__init__(java_model) + + @since("1.4.0") + def precisionAt(self, k: int) -> float: + """ + Compute the average precision of all the queries, truncated at ranking position k. + + If for a query, the ranking algorithm returns n (n < k) results, the precision value + will be computed as #(relevant items retrieved) / k. This formula also applies when + the size of the ground truth set is less than k. + + If a query has an empty ground truth set, zero will be used as precision together + with a log warning. + """ + return self.call("precisionAt", int(k)) + + @property + @since("1.4.0") + def meanAveragePrecision(self) -> float: + """ + Returns the mean average precision (MAP) of all the queries. + If a query has an empty ground truth set, the average precision will be zero and + a log warning is generated. + """ + return self.call("meanAveragePrecision") + + @since("3.0.0") + def meanAveragePrecisionAt(self, k: int) -> float: + """ + Returns the mean average precision (MAP) at first k ranking of all the queries. + If a query has an empty ground truth set, the average precision will be zero and + a log warning is generated. + """ + return self.call("meanAveragePrecisionAt", int(k)) + + @since("1.4.0") + def ndcgAt(self, k: int) -> float: + """ + Compute the average NDCG value of all the queries, truncated at ranking position k. + The discounted cumulative gain at position k is computed as: + sum,,i=1,,^k^ (2^{relevance of ''i''th item}^ - 1) / log(i + 1), + and the NDCG is obtained by dividing the DCG value on the ground truth set. + In the current implementation, the relevance value is binary. + If a query has an empty ground truth set, zero will be used as NDCG together with + a log warning. + """ + return self.call("ndcgAt", int(k)) + + @since("3.0.0") + def recallAt(self, k: int) -> float: + """ + Compute the average recall of all the queries, truncated at ranking position k. + + If for a query, the ranking algorithm returns n results, the recall value + will be computed as #(relevant items retrieved) / #(ground truth set). + This formula also applies when the size of the ground truth set is less than k. + + If a query has an empty ground truth set, zero will be used as recall together + with a log warning. + """ + return self.call("recallAt", int(k)) + + +class MultilabelMetrics(JavaModelWrapper): + """ + Evaluator for multilabel classification. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + predictionAndLabels : :py:class:`pyspark.RDD` + an RDD of (predictions, labels) pairs, + both are non-null Arrays, each with unique elements. + + Examples + -------- + >>> predictionAndLabels = sc.parallelize([([0.0, 1.0], [0.0, 2.0]), ([0.0, 2.0], [0.0, 1.0]), + ... ([], [0.0]), ([2.0], [2.0]), ([2.0, 0.0], [2.0, 0.0]), + ... ([0.0, 1.0, 2.0], [0.0, 1.0]), ([1.0], [1.0, 2.0])]) + >>> metrics = MultilabelMetrics(predictionAndLabels) + >>> metrics.precision(0.0) + 1.0 + >>> metrics.recall(1.0) + 0.66... + >>> metrics.f1Measure(2.0) + 0.5 + >>> metrics.precision() + 0.66... + >>> metrics.recall() + 0.64... + >>> metrics.f1Measure() + 0.63... + >>> metrics.microPrecision + 0.72... + >>> metrics.microRecall + 0.66... + >>> metrics.microF1Measure + 0.69... + >>> metrics.hammingLoss + 0.33... + >>> metrics.subsetAccuracy + 0.28... + >>> metrics.accuracy + 0.54... + """ + + def __init__(self, predictionAndLabels: RDD[Tuple[List[float], List[float]]]): + sc = predictionAndLabels.ctx + sql_ctx = SQLContext.getOrCreate(sc) + df = sql_ctx.createDataFrame( + predictionAndLabels, schema=sql_ctx._inferSchema(predictionAndLabels) + ) + assert sc._jvm is not None + java_class = sc._jvm.org.apache.spark.mllib.evaluation.MultilabelMetrics + java_model = java_class(df._jdf) + super(MultilabelMetrics, self).__init__(java_model) + + @since("1.4.0") + def precision(self, label: Optional[float] = None) -> float: + """ + Returns precision or precision for a given label (category) if specified. + """ + if label is None: + return self.call("precision") + else: + return self.call("precision", float(label)) + + @since("1.4.0") + def recall(self, label: Optional[float] = None) -> float: + """ + Returns recall or recall for a given label (category) if specified. + """ + if label is None: + return self.call("recall") + else: + return self.call("recall", float(label)) + + @since("1.4.0") + def f1Measure(self, label: Optional[float] = None) -> float: + """ + Returns f1Measure or f1Measure for a given label (category) if specified. + """ + if label is None: + return self.call("f1Measure") + else: + return self.call("f1Measure", float(label)) + + @property + @since("1.4.0") + def microPrecision(self) -> float: + """ + Returns micro-averaged label-based precision. + (equals to micro-averaged document-based precision) + """ + return self.call("microPrecision") + + @property + @since("1.4.0") + def microRecall(self) -> float: + """ + Returns micro-averaged label-based recall. + (equals to micro-averaged document-based recall) + """ + return self.call("microRecall") + + @property + @since("1.4.0") + def microF1Measure(self) -> float: + """ + Returns micro-averaged label-based f1-measure. + (equals to micro-averaged document-based f1-measure) + """ + return self.call("microF1Measure") + + @property + @since("1.4.0") + def hammingLoss(self) -> float: + """ + Returns Hamming-loss. + """ + return self.call("hammingLoss") + + @property + @since("1.4.0") + def subsetAccuracy(self) -> float: + """ + Returns subset accuracy. + (for equal sets of labels) + """ + return self.call("subsetAccuracy") + + @property + @since("1.4.0") + def accuracy(self) -> float: + """ + Returns accuracy. + """ + return self.call("accuracy") + + +def _test() -> None: + import doctest + import numpy + from pyspark.sql import SparkSession + import pyspark.mllib.evaluation + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + globs = pyspark.mllib.evaluation.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("mllib.evaluation tests").getOrCreate() + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/feature.py new file mode 100644 index 0000000..2a8cb7d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/feature.py @@ -0,0 +1,1085 @@ +# +# 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. +# + +""" +Python package for feature in MLlib. +""" +import sys +import warnings +from typing import Dict, Hashable, Iterable, List, Optional, Tuple, Union, overload, TYPE_CHECKING + +from py4j.protocol import Py4JJavaError + +from pyspark import since +from pyspark.rdd import RDD +from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper +from pyspark.mllib.linalg import Vectors, _convert_to_vector +from pyspark.mllib.util import JavaLoader, JavaSaveable + +from pyspark.context import SparkContext +from pyspark.mllib.linalg import Vector +from pyspark.mllib.regression import LabeledPoint +from py4j.java_collections import JavaMap + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike + from py4j.java_collections import JavaMap + +__all__ = [ + "Normalizer", + "StandardScalerModel", + "StandardScaler", + "HashingTF", + "IDFModel", + "IDF", + "Word2Vec", + "Word2VecModel", + "ChiSqSelector", + "ChiSqSelectorModel", + "ElementwiseProduct", +] + + +class VectorTransformer: + """ + Base class for transformation of a vector or RDD of vector + """ + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Applies transformation on a vector. + + Parameters + ---------- + vector : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + vector or convertible or RDD to be transformed. + """ + raise NotImplementedError + + +class Normalizer(VectorTransformer): + r""" + Normalizes samples individually to unit L\ :sup:`p`\ norm + + For any 1 <= `p` < float('inf'), normalizes samples using + sum(abs(vector) :sup:`p`) :sup:`(1/p)` as norm. + + For `p` = float('inf'), max(abs(vector)) will be used as norm for + normalization. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + p : float, optional + Normalization in L^p^ space, p = 2 by default. + + Examples + -------- + >>> from pyspark.mllib.linalg import Vectors + >>> v = Vectors.dense(range(3)) + >>> nor = Normalizer(1) + >>> nor.transform(v) + DenseVector([0.0, 0.3333, 0.6667]) + + >>> rdd = sc.parallelize([v]) + >>> nor.transform(rdd).collect() + [DenseVector([0.0, 0.3333, 0.6667])] + + >>> nor2 = Normalizer(float("inf")) + >>> nor2.transform(v) + DenseVector([0.0, 0.5, 1.0]) + """ + + def __init__(self, p: float = 2.0): + assert p >= 1.0, "p should be greater than 1.0" + self.p = float(p) + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Applies unit length normalization on a vector. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + vector : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + vector or RDD of vector to be normalized. + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + normalized vector(s). If the norm of the input is zero, it + will return the input vector. + """ + if isinstance(vector, RDD): + vector = vector.map(_convert_to_vector) + else: + vector = _convert_to_vector(vector) + return callMLlibFunc("normalizeVector", self.p, vector) + + +class JavaVectorTransformer(JavaModelWrapper, VectorTransformer): + """ + Wrapper for the model in JVM + """ + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Applies transformation on a vector or an RDD[Vector]. + + Parameters + ---------- + vector : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Input vector(s) to be transformed. + + Notes + ----- + In Python, transform cannot currently be used within + an RDD transformation or action. + Call transform directly on the RDD instead. + """ + if isinstance(vector, RDD): + vector = vector.map(_convert_to_vector) + else: + vector = _convert_to_vector(vector) + return self.call("transform", vector) + + +class StandardScalerModel(JavaVectorTransformer): + """ + Represents a StandardScaler model that can transform vectors. + + .. versionadded:: 1.2.0 + """ + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Applies standardization transformation on a vector. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + vector : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Input vector(s) to be standardized. + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Standardized vector(s). If the variance of a column is + zero, it will return default `0.0` for the column with + zero variance. + + Notes + ----- + In Python, transform cannot currently be used within + an RDD transformation or action. + Call transform directly on the RDD instead. + """ + return JavaVectorTransformer.transform(self, vector) + + @since("1.4.0") + def setWithMean(self, withMean: bool) -> "StandardScalerModel": + """ + Setter of the boolean which decides + whether it uses mean or not + """ + self.call("setWithMean", withMean) + return self + + @since("1.4.0") + def setWithStd(self, withStd: bool) -> "StandardScalerModel": + """ + Setter of the boolean which decides + whether it uses std or not + """ + self.call("setWithStd", withStd) + return self + + @property + @since("2.0.0") + def withStd(self) -> bool: + """ + Returns if the model scales the data to unit standard deviation. + """ + return self.call("withStd") + + @property + @since("2.0.0") + def withMean(self) -> bool: + """ + Returns if the model centers the data before scaling. + """ + return self.call("withMean") + + @property + @since("2.0.0") + def std(self) -> Vector: + """ + Return the column standard deviation values. + """ + return self.call("std") + + @property + @since("2.0.0") + def mean(self) -> Vector: + """ + Return the column mean values. + """ + return self.call("mean") + + +class StandardScaler: + """ + Standardizes features by removing the mean and scaling to unit + variance using column summary statistics on the samples in the + training set. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + withMean : bool, optional + False by default. Centers the data with mean + before scaling. It will build a dense output, so take + care when applying to sparse input. + withStd : bool, optional + True by default. Scales the data to unit + standard deviation. + + Examples + -------- + >>> vs = [Vectors.dense([-2.0, 2.3, 0]), Vectors.dense([3.8, 0.0, 1.9])] + >>> dataset = sc.parallelize(vs) + >>> standardizer = StandardScaler(True, True) + >>> model = standardizer.fit(dataset) + >>> result = model.transform(dataset) + >>> for r in result.collect(): r + DenseVector([-0.7071, 0.7071, -0.7071]) + DenseVector([0.7071, -0.7071, 0.7071]) + >>> int(model.std[0]) + 4 + >>> int(model.mean[0]*10) + 9 + >>> model.withStd + True + >>> model.withMean + True + """ + + def __init__(self, withMean: bool = False, withStd: bool = True): + if not (withMean or withStd): + warnings.warn("Both withMean and withStd are false. The model does nothing.") + self.withMean = withMean + self.withStd = withStd + + def fit(self, dataset: RDD["VectorLike"]) -> "StandardScalerModel": + """ + Computes the mean and variance and stores as a model to be used + for later scaling. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.RDD` + The data used to compute the mean and variance + to build the transformation model. + + Returns + ------- + :py:class:`StandardScalerModel` + """ + dataset = dataset.map(_convert_to_vector) + jmodel = callMLlibFunc("fitStandardScaler", self.withMean, self.withStd, dataset) + return StandardScalerModel(jmodel) + + +class ChiSqSelectorModel(JavaVectorTransformer): + """ + Represents a Chi Squared selector model. + + .. versionadded:: 1.4.0 + """ + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Applies transformation on a vector. + + .. versionadded:: 1.4.0 + + Examples + -------- + vector : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Input vector(s) to be transformed. + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + transformed vector(s). + """ + return JavaVectorTransformer.transform(self, vector) + + +class ChiSqSelector: + """ + Creates a ChiSquared feature selector. + The selector supports different selection methods: `numTopFeatures`, `percentile`, `fpr`, + `fdr`, `fwe`. + + * `numTopFeatures` chooses a fixed number of top features according to a chi-squared test. + + * `percentile` is similar but chooses a fraction of all features + instead of a fixed number. + + * `fpr` chooses all features whose p-values are below a threshold, + thus controlling the false positive rate of selection. + + * `fdr` uses the `Benjamini-Hochberg procedure `_ + to choose all features whose false discovery rate is below a threshold. + + * `fwe` chooses all features whose p-values are below a threshold. The threshold is scaled by + 1/numFeatures, thus controlling the family-wise error rate of selection. + + By default, the selection method is `numTopFeatures`, with the default number of top features + set to 50. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector, DenseVector + >>> from pyspark.mllib.regression import LabeledPoint + >>> data = sc.parallelize([ + ... LabeledPoint(0.0, SparseVector(3, {0: 8.0, 1: 7.0})), + ... LabeledPoint(1.0, SparseVector(3, {1: 9.0, 2: 6.0})), + ... LabeledPoint(1.0, [0.0, 9.0, 8.0]), + ... LabeledPoint(2.0, [7.0, 9.0, 5.0]), + ... LabeledPoint(2.0, [8.0, 7.0, 3.0]) + ... ]) + >>> model = ChiSqSelector(numTopFeatures=1).fit(data) + >>> model.transform(SparseVector(3, {1: 9.0, 2: 6.0})) + SparseVector(1, {}) + >>> model.transform(DenseVector([7.0, 9.0, 5.0])) + DenseVector([7.0]) + >>> model = ChiSqSelector(selectorType="fpr", fpr=0.2).fit(data) + >>> model.transform(SparseVector(3, {1: 9.0, 2: 6.0})) + SparseVector(1, {}) + >>> model.transform(DenseVector([7.0, 9.0, 5.0])) + DenseVector([7.0]) + >>> model = ChiSqSelector(selectorType="percentile", percentile=0.34).fit(data) + >>> model.transform(DenseVector([7.0, 9.0, 5.0])) + DenseVector([7.0]) + """ + + def __init__( + self, + numTopFeatures: int = 50, + selectorType: str = "numTopFeatures", + percentile: float = 0.1, + fpr: float = 0.05, + fdr: float = 0.05, + fwe: float = 0.05, + ): + self.numTopFeatures = numTopFeatures + self.selectorType = selectorType + self.percentile = percentile + self.fpr = fpr + self.fdr = fdr + self.fwe = fwe + + @since("2.1.0") + def setNumTopFeatures(self, numTopFeatures: int) -> "ChiSqSelector": + """ + set numTopFeature for feature selection by number of top features. + Only applicable when selectorType = "numTopFeatures". + """ + self.numTopFeatures = int(numTopFeatures) + return self + + @since("2.1.0") + def setPercentile(self, percentile: float) -> "ChiSqSelector": + """ + set percentile [0.0, 1.0] for feature selection by percentile. + Only applicable when selectorType = "percentile". + """ + self.percentile = float(percentile) + return self + + @since("2.1.0") + def setFpr(self, fpr: float) -> "ChiSqSelector": + """ + set FPR [0.0, 1.0] for feature selection by FPR. + Only applicable when selectorType = "fpr". + """ + self.fpr = float(fpr) + return self + + @since("2.2.0") + def setFdr(self, fdr: float) -> "ChiSqSelector": + """ + set FDR [0.0, 1.0] for feature selection by FDR. + Only applicable when selectorType = "fdr". + """ + self.fdr = float(fdr) + return self + + @since("2.2.0") + def setFwe(self, fwe: float) -> "ChiSqSelector": + """ + set FWE [0.0, 1.0] for feature selection by FWE. + Only applicable when selectorType = "fwe". + """ + self.fwe = float(fwe) + return self + + @since("2.1.0") + def setSelectorType(self, selectorType: str) -> "ChiSqSelector": + """ + set the selector type of the ChisqSelector. + Supported options: "numTopFeatures" (default), "percentile", "fpr", "fdr", "fwe". + """ + self.selectorType = str(selectorType) + return self + + def fit(self, data: RDD[LabeledPoint]) -> "ChiSqSelectorModel": + """ + Returns a ChiSquared feature selector. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` of :py:class:`pyspark.mllib.regression.LabeledPoint` + containing the labeled dataset with categorical features. + Real-valued features will be treated as categorical for each + distinct value. Apply feature discretizer before using this function. + """ + jmodel = callMLlibFunc( + "fitChiSqSelector", + self.selectorType, + self.numTopFeatures, + self.percentile, + self.fpr, + self.fdr, + self.fwe, + data, + ) + return ChiSqSelectorModel(jmodel) + + +class PCAModel(JavaVectorTransformer): + """ + Model fitted by [[PCA]] that can project vectors to a low-dimensional space using PCA. + + .. versionadded:: 1.5.0 + """ + + +class PCA: + """ + A feature transformer that projects vectors to a low-dimensional space using PCA. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> data = [Vectors.sparse(5, [(1, 1.0), (3, 7.0)]), + ... Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]), + ... Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0])] + >>> model = PCA(2).fit(sc.parallelize(data)) + >>> pcArray = model.transform(Vectors.sparse(5, [(1, 1.0), (3, 7.0)])).toArray() + >>> pcArray[0] + 1.648... + >>> pcArray[1] + -4.013... + """ + + def __init__(self, k: int): + """ + Parameters + ---------- + k : int + number of principal components. + """ + self.k = int(k) + + def fit(self, data: RDD["VectorLike"]) -> PCAModel: + """ + Computes a [[PCAModel]] that contains the principal components of the input vectors. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + source vectors + """ + jmodel = callMLlibFunc("fitPCA", self.k, data) + return PCAModel(jmodel) + + +class HashingTF: + """ + Maps a sequence of terms to their term frequencies using the hashing + trick. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + numFeatures : int, optional + number of features (default: 2^20) + + Notes + ----- + The terms must be hashable (can not be dict/set/list...). + + Examples + -------- + >>> htf = HashingTF(100) + >>> doc = "a a b b c d".split(" ") + >>> htf.transform(doc) + SparseVector(100, {...}) + """ + + def __init__(self, numFeatures: int = 1 << 20): + self.numFeatures = numFeatures + self.binary = False + + @since("2.0.0") + def setBinary(self, value: bool) -> "HashingTF": + """ + If True, term frequency vector will be binary such that non-zero + term counts will be set to 1 + (default: False) + """ + self.binary = value + return self + + @since("1.2.0") + def indexOf(self, term: Hashable) -> int: + """Returns the index of the input term.""" + return hash(term) % self.numFeatures + + @overload + def transform(self, document: Iterable[Hashable]) -> Vector: + ... + + @overload + def transform(self, document: RDD[Iterable[Hashable]]) -> RDD[Vector]: + ... + + @since("1.2.0") + def transform( + self, document: Union[Iterable[Hashable], RDD[Iterable[Hashable]]] + ) -> Union[Vector, RDD[Vector]]: + """ + Transforms the input document (list of terms) to term frequency + vectors, or transform the RDD of document to RDD of term + frequency vectors. + """ + if isinstance(document, RDD): + return document.map(self.transform) + + freq: Dict[int, float] = {} + for term in document: + i = self.indexOf(term) + freq[i] = 1.0 if self.binary else freq.get(i, 0) + 1.0 + return Vectors.sparse(self.numFeatures, freq.items()) + + +class IDFModel(JavaVectorTransformer): + """ + Represents an IDF model that can transform term frequency vectors. + + .. versionadded:: 1.2.0 + """ + + @overload + def transform(self, x: "VectorLike") -> Vector: + ... + + @overload + def transform(self, x: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[Vector, RDD[Vector]]: + """ + Transforms term frequency (TF) vectors to TF-IDF vectors. + + If `minDocFreq` was set for the IDF calculation, + the terms which occur in fewer than `minDocFreq` + documents will have an entry of 0. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + an RDD of term frequency vectors or a term frequency + vector + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + an RDD of TF-IDF vectors or a TF-IDF vector + + Notes + ----- + In Python, transform cannot currently be used within + an RDD transformation or action. + Call transform directly on the RDD instead. + """ + return JavaVectorTransformer.transform(self, x) + + @since("1.4.0") + def idf(self) -> Vector: + """ + Returns the current IDF vector. + """ + return self.call("idf") + + @since("3.0.0") + def docFreq(self) -> List[int]: + """ + Returns the document frequency. + """ + return self.call("docFreq") + + @since("3.0.0") + def numDocs(self) -> int: + """ + Returns number of documents evaluated to compute idf + """ + return self.call("numDocs") + + +class IDF: + """ + Inverse document frequency (IDF). + + The standard formulation is used: `idf = log((m + 1) / (d(t) + 1))`, + where `m` is the total number of documents and `d(t)` is the number + of documents that contain term `t`. + + This implementation supports filtering out terms which do not appear + in a minimum number of documents (controlled by the variable + `minDocFreq`). For terms that are not in at least `minDocFreq` + documents, the IDF is found as 0, resulting in TF-IDFs of 0. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + minDocFreq : int + minimum of documents in which a term should appear for filtering + + Examples + -------- + >>> n = 4 + >>> freqs = [Vectors.sparse(n, (1, 3), (1.0, 2.0)), + ... Vectors.dense([0.0, 1.0, 2.0, 3.0]), + ... Vectors.sparse(n, [1], [1.0])] + >>> data = sc.parallelize(freqs) + >>> idf = IDF() + >>> model = idf.fit(data) + >>> tfidf = model.transform(data) + >>> for r in tfidf.collect(): r + SparseVector(4, {1: 0.0, 3: 0.5754}) + DenseVector([0.0, 0.0, 1.3863, 0.863]) + SparseVector(4, {1: 0.0}) + >>> model.transform(Vectors.dense([0.0, 1.0, 2.0, 3.0])) + DenseVector([0.0, 0.0, 1.3863, 0.863]) + >>> model.transform([0.0, 1.0, 2.0, 3.0]) + DenseVector([0.0, 0.0, 1.3863, 0.863]) + >>> model.transform(Vectors.sparse(n, (1, 3), (1.0, 2.0))) + SparseVector(4, {1: 0.0, 3: 0.5754}) + """ + + def __init__(self, minDocFreq: int = 0): + self.minDocFreq = minDocFreq + + def fit(self, dataset: RDD["VectorLike"]) -> IDFModel: + """ + Computes the inverse document frequency. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.RDD` + an RDD of term frequency vectors + """ + if not isinstance(dataset, RDD): + raise TypeError("dataset should be an RDD of term frequency vectors") + jmodel = callMLlibFunc("fitIDF", self.minDocFreq, dataset.map(_convert_to_vector)) + return IDFModel(jmodel) + + +class Word2VecModel(JavaVectorTransformer, JavaSaveable, JavaLoader["Word2VecModel"]): + """ + class for Word2Vec model + """ + + def transform(self, word: str) -> Vector: # type: ignore[override] + """ + Transforms a word to its vector representation + + .. versionadded:: 1.2.0 + + Parameters + ---------- + word : str + a word + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` + vector representation of word(s) + + Notes + ----- + Local use only + """ + try: + return self.call("transform", word) + except Py4JJavaError: + raise ValueError("%s not found" % word) + + def findSynonyms(self, word: Union[str, "VectorLike"], num: int) -> Iterable[Tuple[str, float]]: + """ + Find synonyms of a word + + .. versionadded:: 1.2.0 + + Parameters + ---------- + + word : str or :py:class:`pyspark.mllib.linalg.Vector` + a word or a vector representation of word + num : int + number of synonyms to find + + Returns + ------- + :py:class:`collections.abc.Iterable` + array of (word, cosineSimilarity) + + Notes + ----- + Local use only + """ + if not isinstance(word, str): + word = _convert_to_vector(word) + words, similarity = self.call("findSynonyms", word, num) + return zip(words, similarity) + + @since("1.4.0") + def getVectors(self) -> "JavaMap": + """ + Returns a map of words to their vector representations. + """ + return self.call("getVectors") + + @classmethod + @since("1.5.0") + def load(cls, sc: SparkContext, path: str) -> "Word2VecModel": + """ + Load a model from the given path. + """ + assert sc._jvm is not None + + jmodel = sc._jvm.org.apache.spark.mllib.feature.Word2VecModel.load(sc._jsc.sc(), path) + model = sc._jvm.org.apache.spark.mllib.api.python.Word2VecModelWrapper(jmodel) + return Word2VecModel(model) + + +class Word2Vec: + """Word2Vec creates vector representation of words in a text corpus. + The algorithm first constructs a vocabulary from the corpus + and then learns vector representation of words in the vocabulary. + The vector representation can be used as features in + natural language processing and machine learning algorithms. + + We used skip-gram model in our implementation and hierarchical + softmax method to train the model. The variable names in the + implementation matches the original C implementation. + + For original C implementation, + see https://code.google.com/p/word2vec/ + For research papers, see + Efficient Estimation of Word Representations in Vector Space + and Distributed Representations of Words and Phrases and their + Compositionality. + + .. versionadded:: 1.2.0 + + Examples + -------- + >>> sentence = "a b " * 100 + "a c " * 10 + >>> localDoc = [sentence, sentence] + >>> doc = sc.parallelize(localDoc).map(lambda line: line.split(" ")) + >>> model = Word2Vec().setVectorSize(10).setSeed(42).fit(doc) + + Querying for synonyms of a word will not return that word: + + >>> syms = model.findSynonyms("a", 2) + >>> [s[0] for s in syms] + ['b', 'c'] + + But querying for synonyms of a vector may return the word whose + representation is that vector: + + >>> vec = model.transform("a") + >>> syms = model.findSynonyms(vec, 2) + >>> [s[0] for s in syms] + ['a', 'b'] + + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = Word2VecModel.load(sc, path) + >>> model.transform("a") == sameModel.transform("a") + True + >>> syms = sameModel.findSynonyms("a", 2) + >>> [s[0] for s in syms] + ['b', 'c'] + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + def __init__(self) -> None: + """ + Construct Word2Vec instance + """ + self.vectorSize = 100 + self.learningRate = 0.025 + self.numPartitions = 1 + self.numIterations = 1 + self.seed: Optional[int] = None + self.minCount = 5 + self.windowSize = 5 + + @since("1.2.0") + def setVectorSize(self, vectorSize: int) -> "Word2Vec": + """ + Sets vector size (default: 100). + """ + self.vectorSize = vectorSize + return self + + @since("1.2.0") + def setLearningRate(self, learningRate: float) -> "Word2Vec": + """ + Sets initial learning rate (default: 0.025). + """ + self.learningRate = learningRate + return self + + @since("1.2.0") + def setNumPartitions(self, numPartitions: int) -> "Word2Vec": + """ + Sets number of partitions (default: 1). Use a small number for + accuracy. + """ + self.numPartitions = numPartitions + return self + + @since("1.2.0") + def setNumIterations(self, numIterations: int) -> "Word2Vec": + """ + Sets number of iterations (default: 1), which should be smaller + than or equal to number of partitions. + """ + self.numIterations = numIterations + return self + + @since("1.2.0") + def setSeed(self, seed: int) -> "Word2Vec": + """ + Sets random seed. + """ + self.seed = seed + return self + + @since("1.4.0") + def setMinCount(self, minCount: int) -> "Word2Vec": + """ + Sets minCount, the minimum number of times a token must appear + to be included in the word2vec model's vocabulary (default: 5). + """ + self.minCount = minCount + return self + + @since("2.0.0") + def setWindowSize(self, windowSize: int) -> "Word2Vec": + """ + Sets window size (default: 5). + """ + self.windowSize = windowSize + return self + + def fit(self, data: RDD[List[str]]) -> "Word2VecModel": + """ + Computes the vector representation of each word in vocabulary. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + training data. RDD of list of string + + Returns + ------- + :py:class:`Word2VecModel` + """ + if not isinstance(data, RDD): + raise TypeError("data should be an RDD of list of string") + jmodel = callMLlibFunc( + "trainWord2VecModel", + data, + int(self.vectorSize), + float(self.learningRate), + int(self.numPartitions), + int(self.numIterations), + self.seed, + int(self.minCount), + int(self.windowSize), + ) + return Word2VecModel(jmodel) + + +class ElementwiseProduct(VectorTransformer): + """ + Scales each column of the vector, with the supplied weight vector. + i.e the elementwise product. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> weight = Vectors.dense([1.0, 2.0, 3.0]) + >>> eprod = ElementwiseProduct(weight) + >>> a = Vectors.dense([2.0, 1.0, 3.0]) + >>> eprod.transform(a) + DenseVector([2.0, 2.0, 9.0]) + >>> b = Vectors.dense([9.0, 3.0, 4.0]) + >>> rdd = sc.parallelize([a, b]) + >>> eprod.transform(rdd).collect() + [DenseVector([2.0, 2.0, 9.0]), DenseVector([9.0, 6.0, 12.0])] + """ + + def __init__(self, scalingVector: Vector) -> None: + self.scalingVector = _convert_to_vector(scalingVector) + + @overload + def transform(self, vector: "VectorLike") -> Vector: + ... + + @overload + def transform(self, vector: RDD["VectorLike"]) -> RDD[Vector]: + ... + + def transform( + self, vector: Union["VectorLike", RDD["VectorLike"]] + ) -> Union[Vector, RDD[Vector]]: + """ + Computes the Hadamard product of the vector. + + .. versionadded:: 1.5.0 + """ + if isinstance(vector, RDD): + vector = vector.map(_convert_to_vector) + + else: + vector = _convert_to_vector(vector) + return callMLlibFunc("elementwiseProductVector", self.scalingVector, vector) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + + globs = globals().copy() + spark = SparkSession.builder.master("local[4]").appName("mllib.feature tests").getOrCreate() + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + sys.path.pop(0) + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/fpm.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/fpm.py new file mode 100644 index 0000000..4b26ca6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/fpm.py @@ -0,0 +1,233 @@ +# +# 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 sys + +from typing import Any, Generic, List, NamedTuple, TypeVar + +from pyspark import since, SparkContext +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc +from pyspark.mllib.util import JavaSaveable, JavaLoader, inherit_doc +from pyspark.rdd import RDD + +__all__ = ["FPGrowth", "FPGrowthModel", "PrefixSpan", "PrefixSpanModel"] + +T = TypeVar("T") + + +@inherit_doc +class FPGrowthModel(JavaModelWrapper, JavaSaveable, JavaLoader["FPGrowthModel"]): + """ + A FP-Growth model for mining frequent itemsets + using the Parallel FP-Growth algorithm. + + .. versionadded:: 1.4.0 + + Examples + -------- + >>> data = [["a", "b", "c"], ["a", "b", "d", "e"], ["a", "c", "e"], ["a", "c", "f"]] + >>> rdd = sc.parallelize(data, 2) + >>> model = FPGrowth.train(rdd, 0.6, 2) + >>> sorted(model.freqItemsets().collect()) + [FreqItemset(items=['a'], freq=4), FreqItemset(items=['c'], freq=3), ... + >>> model_path = temp_path + "/fpm" + >>> model.save(sc, model_path) + >>> sameModel = FPGrowthModel.load(sc, model_path) + >>> sorted(model.freqItemsets().collect()) == sorted(sameModel.freqItemsets().collect()) + True + """ + + @since("1.4.0") + def freqItemsets(self) -> RDD["FPGrowth.FreqItemset"]: + """ + Returns the frequent itemsets of this model. + """ + return self.call("getFreqItemsets").map(lambda x: (FPGrowth.FreqItemset(x[0], x[1]))) + + @classmethod + @since("2.0.0") + def load(cls, sc: SparkContext, path: str) -> "FPGrowthModel": + """ + Load a model from the given path. + """ + model = cls._load_java(sc, path) + assert sc._jvm is not None + wrapper = sc._jvm.org.apache.spark.mllib.api.python.FPGrowthModelWrapper(model) + return FPGrowthModel(wrapper) + + +class FPGrowth: + """ + A Parallel FP-growth algorithm to mine frequent itemsets. + + .. versionadded:: 1.4.0 + """ + + @classmethod + def train( + cls, data: RDD[List[T]], minSupport: float = 0.3, numPartitions: int = -1 + ) -> "FPGrowthModel": + """ + Computes an FP-Growth model that contains frequent itemsets. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The input data set, each element contains a transaction. + minSupport : float, optional + The minimal support level. + (default: 0.3) + numPartitions : int, optional + The number of partitions used by parallel FP-growth. A value + of -1 will use the same number as input data. + (default: -1) + """ + model = callMLlibFunc("trainFPGrowthModel", data, float(minSupport), int(numPartitions)) + return FPGrowthModel(model) + + class FreqItemset(NamedTuple): + """ + Represents an (items, freq) tuple. + + .. versionadded:: 1.4.0 + """ + + items: List[Any] + freq: int + + +@inherit_doc +class PrefixSpanModel(JavaModelWrapper, Generic[T]): + """ + Model fitted by PrefixSpan + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> data = [ + ... [["a", "b"], ["c"]], + ... [["a"], ["c", "b"], ["a", "b"]], + ... [["a", "b"], ["e"]], + ... [["f"]]] + >>> rdd = sc.parallelize(data, 2) + >>> model = PrefixSpan.train(rdd) + >>> sorted(model.freqSequences().collect()) + [FreqSequence(sequence=[['a']], freq=3), FreqSequence(sequence=[['a'], ['a']], freq=1), ... + """ + + @since("1.6.0") + def freqSequences(self) -> RDD["PrefixSpan.FreqSequence"]: + """Gets frequent sequences""" + return self.call("getFreqSequences").map(lambda x: PrefixSpan.FreqSequence(x[0], x[1])) + + +class PrefixSpan: + """ + A parallel PrefixSpan algorithm to mine frequent sequential patterns. + The PrefixSpan algorithm is described in Jian Pei et al (2001) [1]_ + + .. versionadded:: 1.6.0 + + .. [1] Jian Pei et al., + "PrefixSpan,: mining sequential patterns efficiently by prefix-projected pattern growth," + Proceedings 17th International Conference on Data Engineering, Heidelberg, + Germany, 2001, pp. 215-224, + doi: https://doi.org/10.1109/ICDE.2001.914830 + """ + + @classmethod + def train( + cls, + data: RDD[List[List[T]]], + minSupport: float = 0.1, + maxPatternLength: int = 10, + maxLocalProjDBSize: int = 32000000, + ) -> PrefixSpanModel[T]: + """ + Finds the complete set of frequent sequential patterns in the + input sequences of itemsets. + + .. versionadded:: 1.6.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The input data set, each element contains a sequence of + itemsets. + minSupport : float, optional + The minimal support level of the sequential pattern, any + pattern that appears more than (minSupport * + size-of-the-dataset) times will be output. + (default: 0.1) + maxPatternLength : int, optional + The maximal length of the sequential pattern, any pattern + that appears less than maxPatternLength will be output. + (default: 10) + maxLocalProjDBSize : int, optional + The maximum number of items (including delimiters used in the + internal storage format) allowed in a projected database before + local processing. If a projected database exceeds this size, + another iteration of distributed prefix growth is run. + (default: 32000000) + """ + model = callMLlibFunc( + "trainPrefixSpanModel", data, minSupport, maxPatternLength, maxLocalProjDBSize + ) + return PrefixSpanModel(model) + + class FreqSequence(NamedTuple): + """ + Represents a (sequence, freq) tuple. + + .. versionadded:: 1.6.0 + """ + + sequence: List[List[Any]] + freq: int + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.mllib.fpm + + globs = pyspark.mllib.fpm.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("mllib.fpm tests").getOrCreate() + globs["sc"] = spark.sparkContext + import tempfile + + temp_path = tempfile.mkdtemp() + globs["temp_path"] = temp_path + try: + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + finally: + from shutil import rmtree + + try: + rmtree(temp_path) + except OSError: + pass + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/__init__.py new file mode 100644 index 0000000..f752f54 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/__init__.py @@ -0,0 +1,1653 @@ +# +# 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. +# + +""" +MLlib utilities for linear algebra. For dense vectors, MLlib +uses the NumPy `array` type, so you can simply pass NumPy arrays +around. For sparse vectors, users can construct a :class:`SparseVector` +object from MLlib or pass SciPy `scipy.sparse` column vectors if +SciPy is available in their environment. +""" + +import sys +import array +import struct + +import numpy as np + +from pyspark import since +from pyspark.ml import linalg as newlinalg +from pyspark.sql.types import ( + UserDefinedType, + StructField, + StructType, + ArrayType, + DoubleType, + IntegerType, + ByteType, + BooleanType, +) + +from typing import ( + Any, + Callable, + cast, + Dict, + Generic, + Iterable, + List, + Optional, + overload, + Sequence, + Tuple, + Type, + TypeVar, + TYPE_CHECKING, + Union, +) + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike, NormType + from scipy.sparse import spmatrix + from numpy.typing import ArrayLike + + +QT = TypeVar("QT") +RT = TypeVar("RT") + + +__all__ = [ + "Vector", + "DenseVector", + "SparseVector", + "Vectors", + "Matrix", + "DenseMatrix", + "SparseMatrix", + "Matrices", + "QRDecomposition", +] + + +# Check whether we have SciPy. MLlib works without it too, but if we have it, some methods, +# such as _dot and _serialize_double_vector, start to support scipy.sparse matrices. + +try: + import scipy.sparse + + _have_scipy = True +except BaseException: + # No SciPy in environment, but that's okay + _have_scipy = False + + +def _convert_to_vector(d: Union["VectorLike", "spmatrix", range]) -> "Vector": + if isinstance(d, Vector): + return d + elif type(d) in (array.array, np.array, np.ndarray, list, tuple, range): + return DenseVector(d) + elif _have_scipy and scipy.sparse.issparse(d): + assert cast("spmatrix", d).shape[1] == 1, "Expected column vector" + # Make sure the converted csc_matrix has sorted indices. + csc = cast("spmatrix", d).tocsc() + if not csc.has_sorted_indices: + csc.sort_indices() + return SparseVector(cast("spmatrix", d).shape[0], csc.indices, csc.data) + else: + raise TypeError("Cannot convert type %s into Vector" % type(d)) + + +def _vector_size(v: Union["VectorLike", "spmatrix", range]) -> int: + """ + Returns the size of the vector. + + Examples + -------- + >>> _vector_size([1., 2., 3.]) + 3 + >>> _vector_size((1., 2., 3.)) + 3 + >>> _vector_size(array.array('d', [1., 2., 3.])) + 3 + >>> _vector_size(np.zeros(3)) + 3 + >>> _vector_size(np.zeros((3, 1))) + 3 + >>> _vector_size(np.zeros((1, 3))) + Traceback (most recent call last): + ... + ValueError: Cannot treat an ndarray of shape (1, 3) as a vector + """ + if isinstance(v, Vector): + return len(v) + elif type(v) in (array.array, list, tuple, range): + return len(v) + elif type(v) == np.ndarray: + if v.ndim == 1 or (v.ndim == 2 and v.shape[1] == 1): + return len(v) + else: + raise ValueError("Cannot treat an ndarray of shape %s as a vector" % str(v.shape)) + elif _have_scipy and scipy.sparse.issparse(v): + assert cast("spmatrix", v).shape[1] == 1, "Expected column vector" + return cast("spmatrix", v).shape[0] + else: + raise TypeError("Cannot treat type %s as a vector" % type(v)) + + +def _format_float(f: float, digits: int = 4) -> str: + s = str(round(f, digits)) + if "." in s: + s = s[: s.index(".") + 1 + digits] + return s + + +def _format_float_list(xs: Iterable[float]) -> List[str]: + return [_format_float(x) for x in xs] + + +def _double_to_long_bits(value: float) -> int: + if np.isnan(value): + value = float("nan") + # pack double into 64 bits, then unpack as long int + return struct.unpack("Q", struct.pack("d", value))[0] + + +class VectorUDT(UserDefinedType): + """ + SQL user-defined type (UDT) for Vector. + """ + + @classmethod + def sqlType(cls) -> StructType: + return StructType( + [ + StructField("type", ByteType(), False), + StructField("size", IntegerType(), True), + StructField("indices", ArrayType(IntegerType(), False), True), + StructField("values", ArrayType(DoubleType(), False), True), + ] + ) + + @classmethod + def module(cls) -> str: + return "pyspark.mllib.linalg" + + @classmethod + def scalaUDT(cls) -> str: + return "org.apache.spark.mllib.linalg.VectorUDT" + + def serialize( + self, obj: "Vector" + ) -> Tuple[int, Optional[int], Optional[List[int]], List[float]]: + if isinstance(obj, SparseVector): + indices = [int(i) for i in obj.indices] + values = [float(v) for v in obj.values] + return (0, obj.size, indices, values) + elif isinstance(obj, DenseVector): + values = [float(v) for v in obj] # type: ignore[attr-defined] + return (1, None, None, values) + else: + raise TypeError("cannot serialize %r of type %r" % (obj, type(obj))) + + def deserialize( + self, datum: Tuple[int, Optional[int], Optional[List[int]], List[float]] + ) -> "Vector": + assert ( + len(datum) == 4 + ), "VectorUDT.deserialize given row with length %d but requires 4" % len(datum) + tpe = datum[0] + if tpe == 0: + return SparseVector(cast(int, datum[1]), cast(List[int], datum[2]), datum[3]) + elif tpe == 1: + return DenseVector(datum[3]) + else: + raise ValueError("do not recognize type %r" % tpe) + + def simpleString(self) -> str: + return "vector" + + +class MatrixUDT(UserDefinedType): + """ + SQL user-defined type (UDT) for Matrix. + """ + + @classmethod + def sqlType(cls) -> StructType: + return StructType( + [ + StructField("type", ByteType(), False), + StructField("numRows", IntegerType(), False), + StructField("numCols", IntegerType(), False), + StructField("colPtrs", ArrayType(IntegerType(), False), True), + StructField("rowIndices", ArrayType(IntegerType(), False), True), + StructField("values", ArrayType(DoubleType(), False), True), + StructField("isTransposed", BooleanType(), False), + ] + ) + + @classmethod + def module(cls) -> str: + return "pyspark.mllib.linalg" + + @classmethod + def scalaUDT(cls) -> str: + return "org.apache.spark.mllib.linalg.MatrixUDT" + + def serialize( + self, obj: "Matrix" + ) -> Tuple[int, int, int, Optional[List[int]], Optional[List[int]], List[float], bool]: + if isinstance(obj, SparseMatrix): + colPtrs = [int(i) for i in obj.colPtrs] + rowIndices = [int(i) for i in obj.rowIndices] + values = [float(v) for v in obj.values] + return ( + 0, + obj.numRows, + obj.numCols, + colPtrs, + rowIndices, + values, + bool(obj.isTransposed), + ) + elif isinstance(obj, DenseMatrix): + values = [float(v) for v in obj.values] + return (1, obj.numRows, obj.numCols, None, None, values, bool(obj.isTransposed)) + else: + raise TypeError("cannot serialize type %r" % (type(obj))) + + def deserialize( + self, + datum: Tuple[int, int, int, Optional[List[int]], Optional[List[int]], List[float], bool], + ) -> "Matrix": + assert ( + len(datum) == 7 + ), "MatrixUDT.deserialize given row with length %d but requires 7" % len(datum) + tpe = datum[0] + if tpe == 0: + return SparseMatrix( + datum[1], + datum[2], + cast(List[int], datum[3]), + cast(List[int], datum[4]), + datum[5], + datum[6], + ) + elif tpe == 1: + return DenseMatrix(datum[1], datum[2], datum[5], datum[6]) + else: + raise ValueError("do not recognize type %r" % tpe) + + def simpleString(self) -> str: + return "matrix" + + +class Vector: + + __UDT__ = VectorUDT() + + """ + Abstract class for DenseVector and SparseVector + """ + + def toArray(self) -> np.ndarray: + """ + Convert the vector into an numpy.ndarray + + Returns + ------- + :py:class:`numpy.ndarray` + """ + raise NotImplementedError + + def asML(self) -> newlinalg.Vector: + """ + Convert this vector to the new mllib-local representation. + This does NOT copy the data; it copies references. + + Returns + ------- + :py:class:`pyspark.ml.linalg.Vector` + """ + raise NotImplementedError + + def __len__(self) -> int: + raise NotImplementedError + + +class DenseVector(Vector): + """ + A dense vector represented by a value array. We use numpy array for + storage and arithmetics will be delegated to the underlying numpy + array. + + Examples + -------- + >>> v = Vectors.dense([1.0, 2.0]) + >>> u = Vectors.dense([3.0, 4.0]) + >>> v + u + DenseVector([4.0, 6.0]) + >>> 2 - v + DenseVector([1.0, 0.0]) + >>> v / 2 + DenseVector([0.5, 1.0]) + >>> v * u + DenseVector([3.0, 8.0]) + >>> u / v + DenseVector([3.0, 2.0]) + >>> u % 2 + DenseVector([1.0, 0.0]) + >>> -v + DenseVector([-1.0, -2.0]) + """ + + def __init__(self, ar: Union[bytes, np.ndarray, Iterable[float]]): + ar_: np.ndarray + if isinstance(ar, bytes): + ar_ = np.frombuffer(ar, dtype=np.float64) + elif not isinstance(ar, np.ndarray): + ar_ = np.array(ar, dtype=np.float64) + else: + ar_ = ar.astype(np.float64) if ar.dtype != np.float64 else ar + self.array = ar_ + + @staticmethod + def parse(s: str) -> "DenseVector": + """ + Parse string representation back into the DenseVector. + + Examples + -------- + >>> DenseVector.parse(' [ 0.0,1.0,2.0, 3.0]') + DenseVector([0.0, 1.0, 2.0, 3.0]) + """ + start = s.find("[") + if start == -1: + raise ValueError("Array should start with '['.") + end = s.find("]") + if end == -1: + raise ValueError("Array should end with ']'.") + s = s[start + 1 : end] + + try: + values = [float(val) for val in s.split(",") if val] + except ValueError: + raise ValueError("Unable to parse values from %s" % s) + return DenseVector(values) + + def __reduce__(self) -> Tuple[Type["DenseVector"], Tuple[bytes]]: + return DenseVector, (self.array.tobytes(),) + + def numNonzeros(self) -> int: + """ + Number of nonzero elements. This scans all active values and count non zeros + """ + return np.count_nonzero(self.array) + + def norm(self, p: "NormType") -> np.float64: + """ + Calculates the norm of a DenseVector. + + Examples + -------- + >>> a = DenseVector([0, -1, 2, -3]) + >>> a.norm(2) + 3.7... + >>> a.norm(1) + 6.0 + """ + return np.linalg.norm(self.array, p) + + def dot(self, other: Iterable[float]) -> np.float64: + """ + Compute the dot product of two Vectors. We support + (Numpy array, list, SparseVector, or SciPy sparse) + and a target NumPy array that is either 1- or 2-dimensional. + Equivalent to calling numpy.dot of the two vectors. + + Examples + -------- + >>> dense = DenseVector(array.array('d', [1., 2.])) + >>> dense.dot(dense) + 5.0 + >>> dense.dot(SparseVector(2, [0, 1], [2., 1.])) + 4.0 + >>> dense.dot(range(1, 3)) + 5.0 + >>> dense.dot(np.array(range(1, 3))) + 5.0 + >>> dense.dot([1.,]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> dense.dot(np.reshape([1., 2., 3., 4.], (2, 2), order='F')) + array([ 5., 11.]) + >>> dense.dot(np.reshape([1., 2., 3.], (3, 1), order='F')) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + if type(other) == np.ndarray: + if other.ndim > 1: + assert len(self) == other.shape[0], "dimension mismatch" + return np.dot(self.array, other) + elif _have_scipy and scipy.sparse.issparse(other): + assert len(self) == cast("spmatrix", other).shape[0], "dimension mismatch" + return cast("spmatrix", other).transpose().dot(self.toArray()) + else: + assert len(self) == _vector_size(other), "dimension mismatch" + if isinstance(other, SparseVector): + return other.dot(self) + elif isinstance(other, Vector): + return np.dot(self.toArray(), other.toArray()) + else: + return np.dot(self.toArray(), cast("ArrayLike", other)) + + def squared_distance(self, other: Iterable[float]) -> np.float64: + """ + Squared distance of two Vectors. + + Examples + -------- + >>> dense1 = DenseVector(array.array('d', [1., 2.])) + >>> dense1.squared_distance(dense1) + 0.0 + >>> dense2 = np.array([2., 1.]) + >>> dense1.squared_distance(dense2) + 2.0 + >>> dense3 = [2., 1.] + >>> dense1.squared_distance(dense3) + 2.0 + >>> sparse1 = SparseVector(2, [0, 1], [2., 1.]) + >>> dense1.squared_distance(sparse1) + 2.0 + >>> dense1.squared_distance([1.,]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> dense1.squared_distance(SparseVector(1, [0,], [1.,])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + assert len(self) == _vector_size(other), "dimension mismatch" + if isinstance(other, SparseVector): + return other.squared_distance(self) + elif _have_scipy and scipy.sparse.issparse(other): + return _convert_to_vector(other).squared_distance(self) # type: ignore[attr-defined] + + if isinstance(other, Vector): + other = other.toArray() + elif not isinstance(other, np.ndarray): + other = np.array(other) + diff: np.ndarray = self.toArray() - other + return np.dot(diff, diff) + + def toArray(self) -> np.ndarray: + """ + Returns an numpy.ndarray + """ + return self.array + + def asML(self) -> newlinalg.DenseVector: + """ + Convert this vector to the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Returns + ------- + :py:class:`pyspark.ml.linalg.DenseVector` + """ + return newlinalg.DenseVector(self.array) + + @property + def values(self) -> np.ndarray: + """ + Returns a list of values + """ + return self.array + + @overload + def __getitem__(self, item: int) -> np.float64: + ... + + @overload + def __getitem__(self, item: slice) -> np.ndarray: + ... + + def __getitem__(self, item: Union[int, slice]) -> Union[np.float64, np.ndarray]: + return self.array[item] + + def __len__(self) -> int: + return len(self.array) + + def __str__(self) -> str: + return "[" + ",".join([str(v) for v in self.array]) + "]" + + def __repr__(self) -> str: + return "DenseVector([%s])" % (", ".join(_format_float(i) for i in self.array)) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, DenseVector): + return np.array_equal(self.array, other.array) + elif isinstance(other, SparseVector): + if len(self) != other.size: + return False + return Vectors._equals(list(range(len(self))), self.array, other.indices, other.values) + return False + + def __ne__(self, other: Any) -> bool: + return not self == other + + def __hash__(self) -> int: + size = len(self) + result = 31 + size + nnz = 0 + i = 0 + while i < size and nnz < 128: + if self.array[i] != 0: + result = 31 * result + i + bits = _double_to_long_bits(self.array[i]) + result = 31 * result + (bits ^ (bits >> 32)) + nnz += 1 + i += 1 + return result + + def __getattr__(self, item: str) -> Any: + return getattr(self.array, item) + + def __neg__(self) -> "DenseVector": + return DenseVector(-self.array) + + def _delegate(op: str) -> Callable[["DenseVector", Any], "DenseVector"]: # type: ignore[misc] + def func(self: "DenseVector", other: Any) -> "DenseVector": + if isinstance(other, DenseVector): + other = other.array + return DenseVector(getattr(self.array, op)(other)) + + return func + + __add__ = _delegate("__add__") + __sub__ = _delegate("__sub__") + __mul__ = _delegate("__mul__") + __div__ = _delegate("__div__") + __truediv__ = _delegate("__truediv__") + __mod__ = _delegate("__mod__") + __radd__ = _delegate("__radd__") + __rsub__ = _delegate("__rsub__") + __rmul__ = _delegate("__rmul__") + __rdiv__ = _delegate("__rdiv__") + __rtruediv__ = _delegate("__rtruediv__") + __rmod__ = _delegate("__rmod__") + + +class SparseVector(Vector): + """ + A simple sparse vector class for passing data to MLlib. Users may + alternatively pass SciPy's {scipy.sparse} data types. + """ + + @overload + def __init__(self, size: int, __indices: bytes, __values: bytes): + ... + + @overload + def __init__(self, size: int, *args: Tuple[int, float]): + ... + + @overload + def __init__(self, size: int, __indices: Iterable[int], __values: Iterable[float]): + ... + + @overload + def __init__(self, size: int, __pairs: Iterable[Tuple[int, float]]): + ... + + @overload + def __init__(self, size: int, __map: Dict[int, float]): + ... + + def __init__( + self, + size: int, + *args: Union[ + bytes, Tuple[int, float], Iterable[float], Iterable[Tuple[int, float]], Dict[int, float] + ], + ): + """ + Create a sparse vector, using either a dictionary, a list of + (index, value) pairs, or two separate arrays of indices and + values (sorted by index). + + Parameters + ---------- + size : int + Size of the vector. + args + Active entries, as a dictionary {index: value, ...}, + a list of tuples [(index, value), ...], or a list of strictly + increasing indices and a list of corresponding values [index, ...], + [value, ...]. Inactive entries are treated as zeros. + + Examples + -------- + >>> SparseVector(4, {1: 1.0, 3: 5.5}) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> SparseVector(4, [(1, 1.0), (3, 5.5)]) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> SparseVector(4, [1, 3], [1.0, 5.5]) + SparseVector(4, {1: 1.0, 3: 5.5}) + """ + self.size = int(size) + """ Size of the vector. """ + assert 1 <= len(args) <= 2, "must pass either 2 or 3 arguments" + if len(args) == 1: + pairs = args[0] + if type(pairs) == dict: + pairs = pairs.items() + pairs = cast(Iterable[Tuple[int, float]], sorted(pairs)) + self.indices = np.array([p[0] for p in pairs], dtype=np.int32) + """ A list of indices corresponding to active entries. """ + self.values = np.array([p[1] for p in pairs], dtype=np.float64) + """ A list of values corresponding to active entries. """ + else: + if isinstance(args[0], bytes): + assert isinstance(args[1], bytes), "values should be string too" + if args[0]: + self.indices = np.frombuffer(args[0], np.int32) + self.values = np.frombuffer(args[1], np.float64) + else: + # np.frombuffer() doesn't work well with empty string in older version + self.indices = np.array([], dtype=np.int32) + self.values = np.array([], dtype=np.float64) + else: + self.indices = np.array(args[0], dtype=np.int32) + self.values = np.array(args[1], dtype=np.float64) + assert len(self.indices) == len(self.values), "index and value arrays not same length" + for i in range(len(self.indices) - 1): + if self.indices[i] >= self.indices[i + 1]: + raise TypeError( + "Indices %s and %s are not strictly increasing" + % (self.indices[i], self.indices[i + 1]) + ) + + def numNonzeros(self) -> int: + """ + Number of nonzero elements. This scans all active values and count non zeros. + """ + return np.count_nonzero(self.values) + + def norm(self, p: "NormType") -> np.float64: + """ + Calculates the norm of a SparseVector. + + Examples + -------- + >>> a = SparseVector(4, [0, 1], [3., -4.]) + >>> a.norm(1) + 7.0 + >>> a.norm(2) + 5.0 + """ + return np.linalg.norm(self.values, p) + + def __reduce__(self) -> Tuple[Type["SparseVector"], Tuple[int, bytes, bytes]]: + return ( + SparseVector, + ( + self.size, + self.indices.tobytes(), + self.values.tobytes(), + ), + ) + + @staticmethod + def parse(s: str) -> "SparseVector": + """ + Parse string representation back into the SparseVector. + + Examples + -------- + >>> SparseVector.parse(' (4, [0,1 ],[ 4.0,5.0] )') + SparseVector(4, {0: 4.0, 1: 5.0}) + """ + start = s.find("(") + if start == -1: + raise ValueError("Tuple should start with '('") + end = s.find(")") + if end == -1: + raise ValueError("Tuple should end with ')'") + s = s[start + 1 : end].strip() + + size = s[: s.find(",")] + try: + size = int(size) # type: ignore[assignment] + except ValueError: + raise ValueError("Cannot parse size %s." % size) + + ind_start = s.find("[") + if ind_start == -1: + raise ValueError("Indices array should start with '['.") + ind_end = s.find("]") + if ind_end == -1: + raise ValueError("Indices array should end with ']'") + new_s = s[ind_start + 1 : ind_end] + ind_list = new_s.split(",") + try: + indices = [int(ind) for ind in ind_list if ind] + except ValueError: + raise ValueError("Unable to parse indices from %s." % new_s) + s = s[ind_end + 1 :].strip() + + val_start = s.find("[") + if val_start == -1: + raise ValueError("Values array should start with '['.") + val_end = s.find("]") + if val_end == -1: + raise ValueError("Values array should end with ']'.") + val_list = s[val_start + 1 : val_end].split(",") + try: + values = [float(val) for val in val_list if val] + except ValueError: + raise ValueError("Unable to parse values from %s." % s) + return SparseVector(cast(int, size), indices, values) + + def dot(self, other: Iterable[float]) -> np.float64: + """ + Dot product with a SparseVector or 1- or 2-dimensional Numpy array. + + Examples + -------- + >>> a = SparseVector(4, [1, 3], [3.0, 4.0]) + >>> a.dot(a) + 25.0 + >>> a.dot(array.array('d', [1., 2., 3., 4.])) + 22.0 + >>> b = SparseVector(4, [2], [1.0]) + >>> a.dot(b) + 0.0 + >>> a.dot(np.array([[1, 1], [2, 2], [3, 3], [4, 4]])) + array([ 22., 22.]) + >>> a.dot([1., 2., 3.]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(np.array([1., 2.])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(DenseVector([1., 2.])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> a.dot(np.zeros((3, 2))) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + + if isinstance(other, np.ndarray): + if other.ndim not in [2, 1]: + raise ValueError("Cannot call dot with %d-dimensional array" % other.ndim) + assert len(self) == other.shape[0], "dimension mismatch" + return np.dot(self.values, other[self.indices]) + + assert len(self) == _vector_size(other), "dimension mismatch" + + if isinstance(other, DenseVector): + return np.dot(other.array[self.indices], self.values) + + elif isinstance(other, SparseVector): + # Find out common indices. + self_cmind = np.in1d(self.indices, other.indices, assume_unique=True) + self_values = self.values[self_cmind] + if self_values.size == 0: + return np.float64(0.0) + else: + other_cmind = np.in1d(other.indices, self.indices, assume_unique=True) + return np.dot(self_values, other.values[other_cmind]) + + else: + return self.dot(_convert_to_vector(other)) # type: ignore[arg-type] + + def squared_distance(self, other: Iterable[float]) -> np.float64: + """ + Squared distance from a SparseVector or 1-dimensional NumPy array. + + Examples + -------- + >>> a = SparseVector(4, [1, 3], [3.0, 4.0]) + >>> a.squared_distance(a) + 0.0 + >>> a.squared_distance(array.array('d', [1., 2., 3., 4.])) + 11.0 + >>> a.squared_distance(np.array([1., 2., 3., 4.])) + 11.0 + >>> b = SparseVector(4, [2], [1.0]) + >>> a.squared_distance(b) + 26.0 + >>> b.squared_distance(a) + 26.0 + >>> b.squared_distance([1., 2.]) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + >>> b.squared_distance(SparseVector(3, [1,], [1.0,])) + Traceback (most recent call last): + ... + AssertionError: dimension mismatch + """ + assert len(self) == _vector_size(other), "dimension mismatch" + + if isinstance(other, np.ndarray) or isinstance(other, DenseVector): + if isinstance(other, np.ndarray) and other.ndim != 1: + raise ValueError( + "Cannot call squared_distance with %d-dimensional array" % other.ndim + ) + if isinstance(other, DenseVector): + other = other.array + sparse_ind = np.zeros(other.size, dtype=bool) + sparse_ind[self.indices] = True + dist = other[sparse_ind] - self.values + result = np.dot(dist, dist) + + other_ind = other[~sparse_ind] + result += np.dot(other_ind, other_ind) + return result + + elif isinstance(other, SparseVector): + result = 0.0 + i, j = 0, 0 + while i < len(self.indices) and j < len(other.indices): + if self.indices[i] == other.indices[j]: + diff = self.values[i] - other.values[j] + result += diff * diff + i += 1 + j += 1 + elif self.indices[i] < other.indices[j]: + result += self.values[i] * self.values[i] + i += 1 + else: + result += other.values[j] * other.values[j] + j += 1 + while i < len(self.indices): + result += self.values[i] * self.values[i] + i += 1 + while j < len(other.indices): + result += other.values[j] * other.values[j] + j += 1 + return result + else: + return self.squared_distance(_convert_to_vector(other)) # type: ignore[arg-type] + + def toArray(self) -> np.ndarray: + """ + Returns a copy of this SparseVector as a 1-dimensional NumPy array. + """ + arr = np.zeros((self.size,), dtype=np.float64) + arr[self.indices] = self.values + return arr + + def asML(self) -> newlinalg.SparseVector: + """ + Convert this vector to the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Returns + ------- + :py:class:`pyspark.ml.linalg.SparseVector` + """ + return newlinalg.SparseVector(self.size, self.indices, self.values) + + def __len__(self) -> int: + return self.size + + def __str__(self) -> str: + inds = "[" + ",".join([str(i) for i in self.indices]) + "]" + vals = "[" + ",".join([str(v) for v in self.values]) + "]" + return "(" + ",".join((str(self.size), inds, vals)) + ")" + + def __repr__(self) -> str: + inds = self.indices + vals = self.values + entries = ", ".join( + ["{0}: {1}".format(inds[i], _format_float(vals[i])) for i in range(len(inds))] + ) + return "SparseVector({0}, {{{1}}})".format(self.size, entries) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, SparseVector): + return ( + other.size == self.size + and np.array_equal(other.indices, self.indices) + and np.array_equal(other.values, self.values) + ) + elif isinstance(other, DenseVector): + if self.size != len(other): + return False + return Vectors._equals(self.indices, self.values, list(range(len(other))), other.array) + return False + + def __getitem__(self, index: int) -> np.float64: + inds = self.indices + vals = self.values + if not isinstance(index, int): + raise TypeError("Indices must be of type integer, got type %s" % type(index)) + + if index >= self.size or index < -self.size: + raise IndexError("Index %d out of bounds." % index) + if index < 0: + index += self.size + + if (inds.size == 0) or (index > inds.item(-1)): + return np.float64(0.0) + + insert_index = np.searchsorted(inds, index) + row_ind = inds[insert_index] + if row_ind == index: + return vals[insert_index] + return np.float64(0.0) + + def __ne__(self, other: Any) -> bool: + return not self.__eq__(other) + + def __hash__(self) -> int: + result = 31 + self.size + nnz = 0 + i = 0 + while i < len(self.values) and nnz < 128: + if self.values[i] != 0: + result = 31 * result + int(self.indices[i]) + bits = _double_to_long_bits(self.values[i]) + result = 31 * result + (bits ^ (bits >> 32)) + nnz += 1 + i += 1 + return result + + +class Vectors: + + """ + Factory methods for working with vectors. + + Notes + ----- + Dense vectors are simply represented as NumPy array objects, + so there is no need to convert them for use in MLlib. For sparse vectors, + the factory methods in this class create an MLlib-compatible type, or users + can pass in SciPy's `scipy.sparse` column vectors. + """ + + @staticmethod + @overload + def sparse(size: int, __indices: bytes, __values: bytes) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, *args: Tuple[int, float]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __indices: Iterable[int], __values: Iterable[float]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __pairs: Iterable[Tuple[int, float]]) -> SparseVector: + ... + + @staticmethod + @overload + def sparse(size: int, __map: Dict[int, float]) -> SparseVector: + ... + + @staticmethod + def sparse( + size: int, + *args: Union[ + bytes, Tuple[int, float], Iterable[float], Iterable[Tuple[int, float]], Dict[int, float] + ], + ) -> SparseVector: + """ + Create a sparse vector, using either a dictionary, a list of + (index, value) pairs, or two separate arrays of indices and + values (sorted by index). + + Parameters + ---------- + size : int + Size of the vector. + args + Non-zero entries, as a dictionary, list of tuples, + or two sorted lists containing indices and values. + + Examples + -------- + >>> Vectors.sparse(4, {1: 1.0, 3: 5.5}) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> Vectors.sparse(4, [(1, 1.0), (3, 5.5)]) + SparseVector(4, {1: 1.0, 3: 5.5}) + >>> Vectors.sparse(4, [1, 3], [1.0, 5.5]) + SparseVector(4, {1: 1.0, 3: 5.5}) + """ + return SparseVector(size, *args) # type: ignore[arg-type] + + @overload + @staticmethod + def dense(*elements: float) -> DenseVector: + ... + + @overload + @staticmethod + def dense(__arr: bytes) -> DenseVector: + ... + + @overload + @staticmethod + def dense(__arr: Iterable[float]) -> DenseVector: + ... + + @staticmethod + def dense(*elements: Union[float, bytes, np.ndarray, Iterable[float]]) -> DenseVector: + """ + Create a dense vector of 64-bit floats from a Python list or numbers. + + Examples + -------- + >>> Vectors.dense([1, 2, 3]) + DenseVector([1.0, 2.0, 3.0]) + >>> Vectors.dense(1.0, 2.0) + DenseVector([1.0, 2.0]) + """ + if len(elements) == 1 and not isinstance(elements[0], (float, int)): + # it's list, numpy.array or other iterable object. + elements = elements[0] # type: ignore[assignment] + return DenseVector(cast(Iterable[float], elements)) + + @staticmethod + def fromML(vec: newlinalg.DenseVector) -> DenseVector: + """ + Convert a vector from the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + vec : :py:class:`pyspark.ml.linalg.Vector` + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Vector` + """ + if isinstance(vec, newlinalg.DenseVector): + return DenseVector(vec.array) + elif isinstance(vec, newlinalg.SparseVector): + return SparseVector(vec.size, vec.indices, vec.values) + else: + raise TypeError("Unsupported vector type %s" % type(vec)) + + @staticmethod + def stringify(vector: Vector) -> str: + """ + Converts a vector into a string, which can be recognized by + Vectors.parse(). + + Examples + -------- + >>> Vectors.stringify(Vectors.sparse(2, [1], [1.0])) + '(2,[1],[1.0])' + >>> Vectors.stringify(Vectors.dense([0.0, 1.0])) + '[0.0,1.0]' + """ + return str(vector) + + @staticmethod + def squared_distance(v1: Vector, v2: Vector) -> np.float64: + """ + Squared distance between two vectors. + a and b can be of type SparseVector, DenseVector, np.ndarray + or array.array. + + Examples + -------- + >>> a = Vectors.sparse(4, [(0, 1), (3, 4)]) + >>> b = Vectors.dense([2, 5, 4, 1]) + >>> a.squared_distance(b) + 51.0 + """ + v1, v2 = _convert_to_vector(v1), _convert_to_vector(v2) + return v1.squared_distance(v2) # type: ignore[attr-defined] + + @staticmethod + def norm(vector: Vector, p: "NormType") -> np.float64: + """ + Find norm of the given vector. + """ + return _convert_to_vector(vector).norm(p) # type: ignore[attr-defined] + + @staticmethod + def parse(s: str) -> Vector: + """Parse a string representation back into the Vector. + + Examples + -------- + >>> Vectors.parse('[2,1,2 ]') + DenseVector([2.0, 1.0, 2.0]) + >>> Vectors.parse(' ( 100, [0], [2])') + SparseVector(100, {0: 2.0}) + """ + if s.find("(") == -1 and s.find("[") != -1: + return DenseVector.parse(s) + elif s.find("(") != -1: + return SparseVector.parse(s) + else: + raise ValueError("Cannot find tokens '[' or '(' from the input string.") + + @staticmethod + def zeros(size: int) -> DenseVector: + return DenseVector(np.zeros(size)) + + @staticmethod + def _equals( + v1_indices: Union[Sequence[int], np.ndarray], + v1_values: Union[Sequence[float], np.ndarray], + v2_indices: Union[Sequence[int], np.ndarray], + v2_values: Union[Sequence[float], np.ndarray], + ) -> bool: + """ + Check equality between sparse/dense vectors, + v1_indices and v2_indices assume to be strictly increasing. + """ + v1_size = len(v1_values) + v2_size = len(v2_values) + k1 = 0 + k2 = 0 + all_equal = True + while all_equal: + while k1 < v1_size and v1_values[k1] == 0: + k1 += 1 + while k2 < v2_size and v2_values[k2] == 0: + k2 += 1 + + if k1 >= v1_size or k2 >= v2_size: + return k1 >= v1_size and k2 >= v2_size + + all_equal = v1_indices[k1] == v2_indices[k2] and v1_values[k1] == v2_values[k2] + k1 += 1 + k2 += 1 + return all_equal + + +class Matrix: + + __UDT__ = MatrixUDT() + + """ + Represents a local matrix. + """ + + def __init__(self, numRows: int, numCols: int, isTransposed: bool = False) -> None: + self.numRows = numRows + self.numCols = numCols + self.isTransposed = isTransposed + + def toArray(self) -> np.ndarray: + """ + Returns its elements in a NumPy ndarray. + """ + raise NotImplementedError + + def asML(self) -> newlinalg.Matrix: + """ + Convert this matrix to the new mllib-local representation. + This does NOT copy the data; it copies references. + """ + raise NotImplementedError + + @staticmethod + def _convert_to_array(array_like: Union[bytes, Iterable[float]], dtype: Any) -> np.ndarray: + """ + Convert Matrix attributes which are array-like or buffer to array. + """ + if isinstance(array_like, bytes): + return np.frombuffer(array_like, dtype=dtype) + return np.asarray(array_like, dtype=dtype) + + +class DenseMatrix(Matrix): + """ + Column-major dense matrix. + """ + + def __init__( + self, + numRows: int, + numCols: int, + values: Union[bytes, Iterable[float]], + isTransposed: bool = False, + ): + Matrix.__init__(self, numRows, numCols, isTransposed) + values = self._convert_to_array(values, np.float64) + assert len(values) == numRows * numCols + self.values = values + + def __reduce__(self) -> Tuple[Type["DenseMatrix"], Tuple[int, int, bytes, int]]: + return DenseMatrix, ( + self.numRows, + self.numCols, + self.values.tobytes(), + int(self.isTransposed), + ) + + def __str__(self) -> str: + """ + Pretty printing of a DenseMatrix + + Examples + -------- + >>> dm = DenseMatrix(2, 2, range(4)) + >>> print(dm) + DenseMatrix([[ 0., 2.], + [ 1., 3.]]) + >>> dm = DenseMatrix(2, 2, range(4), isTransposed=True) + >>> print(dm) + DenseMatrix([[ 0., 1.], + [ 2., 3.]]) + """ + # Inspired by __repr__ in scipy matrices. + array_lines = repr(self.toArray()).splitlines() + + # We need to adjust six spaces which is the difference in number + # of letters between "DenseMatrix" and "array" + x = "\n".join([(" " * 6 + line) for line in array_lines[1:]]) + return array_lines[0].replace("array", "DenseMatrix") + "\n" + x + + def __repr__(self) -> str: + """ + Representation of a DenseMatrix + + Examples + -------- + >>> dm = DenseMatrix(2, 2, range(4)) + >>> dm + DenseMatrix(2, 2, [0.0, 1.0, 2.0, 3.0], False) + """ + # If the number of values are less than seventeen then return as it is. + # Else return first eight values and last eight values. + if len(self.values) < 17: + entries = _format_float_list(self.values) + else: + entries = ( + _format_float_list(self.values[:8]) + ["..."] + _format_float_list(self.values[-8:]) + ) + + return "DenseMatrix({0}, {1}, [{2}], {3})".format( + self.numRows, self.numCols, ", ".join(entries), self.isTransposed + ) + + def toArray(self) -> np.ndarray: + """ + Return an numpy.ndarray + + Examples + -------- + >>> m = DenseMatrix(2, 2, range(4)) + >>> m.toArray() + array([[ 0., 2.], + [ 1., 3.]]) + """ + if self.isTransposed: + return np.asfortranarray(self.values.reshape((self.numRows, self.numCols))) + else: + return self.values.reshape((self.numRows, self.numCols), order="F") + + def toSparse(self) -> "SparseMatrix": + """Convert to SparseMatrix""" + if self.isTransposed: + values = np.ravel(self.toArray(), order="F") + else: + values = self.values + indices = np.nonzero(values)[0] + colCounts = np.bincount(indices // self.numRows) + colPtrs = np.cumsum(np.hstack((0, colCounts, np.zeros(self.numCols - colCounts.size)))) + values = values[indices] + rowIndices = indices % self.numRows + + return SparseMatrix(self.numRows, self.numCols, colPtrs, rowIndices, values) + + def asML(self) -> newlinalg.DenseMatrix: + """ + Convert this matrix to the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Returns + ------- + :py:class:`pyspark.ml.linalg.DenseMatrix` + """ + return newlinalg.DenseMatrix(self.numRows, self.numCols, self.values, self.isTransposed) + + def __getitem__(self, indices: Tuple[int, int]) -> np.float64: + i, j = indices + if i < 0 or i >= self.numRows: + raise IndexError("Row index %d is out of range [0, %d)" % (i, self.numRows)) + if j >= self.numCols or j < 0: + raise IndexError("Column index %d is out of range [0, %d)" % (j, self.numCols)) + + if self.isTransposed: + return self.values[i * self.numCols + j] + else: + return self.values[i + j * self.numRows] + + def __eq__(self, other: Any) -> bool: + if self.numRows != other.numRows or self.numCols != other.numCols: + return False + if isinstance(other, SparseMatrix): + return np.all(self.toArray() == other.toArray()).tolist() + + self_values = np.ravel(self.toArray(), order="F") + other_values = np.ravel(other.toArray(), order="F") + return np.all(self_values == other_values).tolist() + + +class SparseMatrix(Matrix): + """Sparse Matrix stored in CSC format.""" + + def __init__( + self, + numRows: int, + numCols: int, + colPtrs: Union[bytes, Iterable[int]], + rowIndices: Union[bytes, Iterable[int]], + values: Union[bytes, Iterable[float]], + isTransposed: bool = False, + ) -> None: + Matrix.__init__(self, numRows, numCols, isTransposed) + self.colPtrs = self._convert_to_array(colPtrs, np.int32) + self.rowIndices = self._convert_to_array(rowIndices, np.int32) + self.values = self._convert_to_array(values, np.float64) + + if self.isTransposed: + if self.colPtrs.size != numRows + 1: + raise ValueError( + "Expected colPtrs of size %d, got %d." % (numRows + 1, self.colPtrs.size) + ) + else: + if self.colPtrs.size != numCols + 1: + raise ValueError( + "Expected colPtrs of size %d, got %d." % (numCols + 1, self.colPtrs.size) + ) + if self.rowIndices.size != self.values.size: + raise ValueError( + "Expected rowIndices of length %d, got %d." + % (self.rowIndices.size, self.values.size) + ) + + def __str__(self) -> str: + """ + Pretty printing of a SparseMatrix + + Examples + -------- + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + >>> print(sm1) + 2 X 2 CSCMatrix + (0,0) 2.0 + (1,0) 3.0 + (1,1) 4.0 + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + >>> print(sm1) + 2 X 2 CSRMatrix + (0,0) 2.0 + (0,1) 3.0 + (1,1) 4.0 + """ + spstr = "{0} X {1} ".format(self.numRows, self.numCols) + if self.isTransposed: + spstr += "CSRMatrix\n" + else: + spstr += "CSCMatrix\n" + + cur_col = 0 + smlist = [] + + # Display first 16 values. + if len(self.values) <= 16: + zipindval = zip(self.rowIndices, self.values) + else: + zipindval = zip(self.rowIndices[:16], self.values[:16]) + for i, (rowInd, value) in enumerate(zipindval): + if self.colPtrs[cur_col + 1] <= i: + cur_col += 1 + if self.isTransposed: + smlist.append("({0},{1}) {2}".format(cur_col, rowInd, _format_float(value))) + else: + smlist.append("({0},{1}) {2}".format(rowInd, cur_col, _format_float(value))) + spstr += "\n".join(smlist) + + if len(self.values) > 16: + spstr += "\n.." * 2 + return spstr + + def __repr__(self) -> str: + """ + Representation of a SparseMatrix + + Examples + -------- + >>> sm1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + >>> sm1 + SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2.0, 3.0, 4.0], False) + """ + rowIndices = list(self.rowIndices) + colPtrs = list(self.colPtrs) + + if len(self.values) <= 16: + values = _format_float_list(self.values) + + else: + values = ( + _format_float_list(self.values[:8]) + ["..."] + _format_float_list(self.values[-8:]) + ) + rowIndices = rowIndices[:8] + ["..."] + rowIndices[-8:] + + if len(self.colPtrs) > 16: + colPtrs = colPtrs[:8] + ["..."] + colPtrs[-8:] + + return "SparseMatrix({0}, {1}, [{2}], [{3}], [{4}], {5})".format( + self.numRows, + self.numCols, + ", ".join([str(ptr) for ptr in colPtrs]), + ", ".join([str(ind) for ind in rowIndices]), + ", ".join(values), + self.isTransposed, + ) + + def __reduce__(self) -> Tuple[Type["SparseMatrix"], Tuple[int, int, bytes, bytes, bytes, int]]: + return SparseMatrix, ( + self.numRows, + self.numCols, + self.colPtrs.tobytes(), + self.rowIndices.tobytes(), + self.values.tobytes(), + int(self.isTransposed), + ) + + def __getitem__(self, indices: Tuple[int, int]) -> np.float64: + i, j = indices + if i < 0 or i >= self.numRows: + raise IndexError("Row index %d is out of range [0, %d)" % (i, self.numRows)) + if j < 0 or j >= self.numCols: + raise IndexError("Column index %d is out of range [0, %d)" % (j, self.numCols)) + + # If a CSR matrix is given, then the row index should be searched + # for in ColPtrs, and the column index should be searched for in the + # corresponding slice obtained from rowIndices. + if self.isTransposed: + j, i = i, j + + colStart = self.colPtrs[j] + colEnd = self.colPtrs[j + 1] + nz = self.rowIndices[colStart:colEnd] + ind = np.searchsorted(nz, i) + colStart + if ind < colEnd and self.rowIndices[ind] == i: + return self.values[ind] + else: + return np.float64(0.0) + + def toArray(self) -> np.ndarray: + """ + Return an numpy.ndarray + """ + A = np.zeros((self.numRows, self.numCols), dtype=np.float64, order="F") + for k in range(self.colPtrs.size - 1): + startptr = self.colPtrs[k] + endptr = self.colPtrs[k + 1] + if self.isTransposed: + A[k, self.rowIndices[startptr:endptr]] = self.values[startptr:endptr] + else: + A[self.rowIndices[startptr:endptr], k] = self.values[startptr:endptr] + return A + + def toDense(self) -> "DenseMatrix": + densevals = np.ravel(self.toArray(), order="F") + return DenseMatrix(self.numRows, self.numCols, densevals) + + def asML(self) -> newlinalg.SparseMatrix: + """ + Convert this matrix to the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Returns + ------- + :py:class:`pyspark.ml.linalg.SparseMatrix` + """ + return newlinalg.SparseMatrix( + self.numRows, + self.numCols, + self.colPtrs, + self.rowIndices, + self.values, + self.isTransposed, + ) + + # TODO: More efficient implementation: + def __eq__(self, other: Any) -> bool: + assert isinstance(other, Matrix) + return np.all(self.toArray() == other.toArray()).tolist() + + +class Matrices: + @staticmethod + def dense(numRows: int, numCols: int, values: Union[bytes, Iterable[float]]) -> DenseMatrix: + """ + Create a DenseMatrix + """ + return DenseMatrix(numRows, numCols, values) + + @staticmethod + def sparse( + numRows: int, + numCols: int, + colPtrs: Union[bytes, Iterable[int]], + rowIndices: Union[bytes, Iterable[int]], + values: Union[bytes, Iterable[float]], + ) -> SparseMatrix: + """ + Create a SparseMatrix + """ + return SparseMatrix(numRows, numCols, colPtrs, rowIndices, values) + + @staticmethod + def fromML(mat: newlinalg.Matrix) -> Matrix: + """ + Convert a matrix from the new mllib-local representation. + This does NOT copy the data; it copies references. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + mat : :py:class:`pyspark.ml.linalg.Matrix` + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Matrix` + """ + if isinstance(mat, newlinalg.DenseMatrix): + return DenseMatrix(mat.numRows, mat.numCols, mat.values, mat.isTransposed) + elif isinstance(mat, newlinalg.SparseMatrix): + return SparseMatrix( + mat.numRows, mat.numCols, mat.colPtrs, mat.rowIndices, mat.values, mat.isTransposed + ) + else: + raise TypeError("Unsupported matrix type %s" % type(mat)) + + +class QRDecomposition(Generic[QT, RT]): + """ + Represents QR factors. + """ + + def __init__(self, Q: QT, R: RT) -> None: + self._Q = Q + self._R = R + + @property + @since("2.0.0") + def Q(self) -> QT: + """ + An orthogonal matrix Q in a QR decomposition. + May be null if not computed. + """ + return self._Q + + @property + @since("2.0.0") + def R(self) -> RT: + """ + An upper triangular matrix R in a QR decomposition. + """ + return self._R + + +def _test() -> None: + import doctest + import numpy + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + (failure_count, test_count) = doctest.testmod(optionflags=doctest.ELLIPSIS) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/distributed.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/distributed.py new file mode 100644 index 0000000..ecdb4e7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/linalg/distributed.py @@ -0,0 +1,1662 @@ +# +# 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. +# + +""" +Package for distributed linear algebra. +""" + +import sys +from typing import Any, Generic, Optional, Tuple, TypeVar, Union, TYPE_CHECKING + +from py4j.java_gateway import JavaObject + +from pyspark import RDD, since +from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper +from pyspark.mllib.linalg import _convert_to_vector, DenseMatrix, Matrix, QRDecomposition, Vector +from pyspark.mllib.stat import MultivariateStatisticalSummary +from pyspark.sql import DataFrame +from pyspark.storagelevel import StorageLevel + +UT = TypeVar("UT", bound="DistributedMatrix") +VT = TypeVar("VT", bound="Matrix") + +if TYPE_CHECKING: + from pyspark.ml._typing import VectorLike + +__all__ = [ + "BlockMatrix", + "CoordinateMatrix", + "DistributedMatrix", + "IndexedRow", + "IndexedRowMatrix", + "MatrixEntry", + "RowMatrix", + "SingularValueDecomposition", +] + + +class DistributedMatrix: + """ + Represents a distributively stored matrix backed by one or + more RDDs. + + """ + + def numRows(self) -> int: + """Get or compute the number of rows.""" + raise NotImplementedError + + def numCols(self) -> int: + """Get or compute the number of cols.""" + raise NotImplementedError + + +class RowMatrix(DistributedMatrix): + """ + Represents a row-oriented distributed Matrix with no meaningful + row indices. + + + Parameters + ---------- + rows : :py:class:`pyspark.RDD` or :py:class:`pyspark.sql.DataFrame` + An RDD or DataFrame of vectors. If a DataFrame is provided, it must have a single + vector typed column. + numRows : int, optional + Number of rows in the matrix. A non-positive + value means unknown, at which point the number + of rows will be determined by the number of + records in the `rows` RDD. + numCols : int, optional + Number of columns in the matrix. A non-positive + value means unknown, at which point the number + of columns will be determined by the size of + the first row. + """ + + def __init__( + self, + rows: Union[RDD[Vector], DataFrame], + numRows: int = 0, + numCols: int = 0, + ): + """ + Note: This docstring is not shown publicly. + + Create a wrapper over a Java RowMatrix. + + Publicly, we require that `rows` be an RDD or DataFrame. However, for + internal usage, `rows` can also be a Java RowMatrix + object, in which case we can wrap it directly. This + assists in clean matrix conversions. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [4, 5, 6]]) + >>> mat = RowMatrix(rows) + + >>> mat_diff = RowMatrix(rows) + >>> (mat_diff._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + False + + >>> mat_same = RowMatrix(mat._java_matrix_wrapper._java_model) + >>> (mat_same._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + True + """ + if isinstance(rows, RDD): + rows = rows.map(_convert_to_vector) + java_matrix = callMLlibFunc("createRowMatrix", rows, int(numRows), int(numCols)) + elif isinstance(rows, DataFrame): + java_matrix = callMLlibFunc("createRowMatrix", rows, int(numRows), int(numCols)) + elif isinstance(rows, JavaObject) and rows.getClass().getSimpleName() == "RowMatrix": + java_matrix = rows + else: + raise TypeError("rows should be an RDD of vectors, got %s" % type(rows)) + + self._java_matrix_wrapper = JavaModelWrapper(java_matrix) + + @property + def rows(self) -> RDD[Vector]: + """ + Rows of the RowMatrix stored as an RDD of vectors. + + Examples + -------- + >>> mat = RowMatrix(sc.parallelize([[1, 2, 3], [4, 5, 6]])) + >>> rows = mat.rows + >>> rows.first() + DenseVector([1.0, 2.0, 3.0]) + """ + return self._java_matrix_wrapper.call("rows") + + def numRows(self) -> int: + """ + Get or compute the number of rows. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [4, 5, 6], + ... [7, 8, 9], [10, 11, 12]]) + + >>> mat = RowMatrix(rows) + >>> print(mat.numRows()) + 4 + + >>> mat = RowMatrix(rows, 7, 6) + >>> print(mat.numRows()) + 7 + """ + return self._java_matrix_wrapper.call("numRows") + + def numCols(self) -> int: + """ + Get or compute the number of cols. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [4, 5, 6], + ... [7, 8, 9], [10, 11, 12]]) + + >>> mat = RowMatrix(rows) + >>> print(mat.numCols()) + 3 + + >>> mat = RowMatrix(rows, 7, 6) + >>> print(mat.numCols()) + 6 + """ + return self._java_matrix_wrapper.call("numCols") + + def computeColumnSummaryStatistics(self) -> MultivariateStatisticalSummary: + """ + Computes column-wise summary statistics. + + .. versionadded:: 2.0.0 + + Returns + ------- + :py:class:`MultivariateStatisticalSummary` + object containing column-wise summary statistics. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [4, 5, 6]]) + >>> mat = RowMatrix(rows) + + >>> colStats = mat.computeColumnSummaryStatistics() + >>> colStats.mean() + array([ 2.5, 3.5, 4.5]) + """ + java_col_stats = self._java_matrix_wrapper.call("computeColumnSummaryStatistics") + return MultivariateStatisticalSummary(java_col_stats) + + def computeCovariance(self) -> Matrix: + """ + Computes the covariance matrix, treating each row as an + observation. + + .. versionadded:: 2.0.0 + + Notes + ----- + This cannot be computed on matrices with more than 65535 columns. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2], [2, 1]]) + >>> mat = RowMatrix(rows) + + >>> mat.computeCovariance() + DenseMatrix(2, 2, [0.5, -0.5, -0.5, 0.5], 0) + """ + return self._java_matrix_wrapper.call("computeCovariance") + + def computeGramianMatrix(self) -> Matrix: + """ + Computes the Gramian matrix `A^T A`. + + .. versionadded:: 2.0.0 + + Notes + ----- + This cannot be computed on matrices with more than 65535 columns. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [4, 5, 6]]) + >>> mat = RowMatrix(rows) + + >>> mat.computeGramianMatrix() + DenseMatrix(3, 3, [17.0, 22.0, 27.0, 22.0, 29.0, 36.0, 27.0, 36.0, 45.0], 0) + """ + return self._java_matrix_wrapper.call("computeGramianMatrix") + + @since("2.0.0") + def columnSimilarities(self, threshold: float = 0.0) -> "CoordinateMatrix": + """ + Compute similarities between columns of this matrix. + + The threshold parameter is a trade-off knob between estimate + quality and computational cost. + + The default threshold setting of 0 guarantees deterministically + correct results, but uses the brute-force approach of computing + normalized dot products. + + Setting the threshold to positive values uses a sampling + approach and incurs strictly less computational cost than the + brute-force approach. However the similarities computed will + be estimates. + + The sampling guarantees relative-error correctness for those + pairs of columns that have similarity greater than the given + similarity threshold. + + To describe the guarantee, we set some notation: + + - Let A be the smallest in magnitude non-zero element of + this matrix. + - Let B be the largest in magnitude non-zero element of + this matrix. + - Let L be the maximum number of non-zeros per row. + + For example, for {0,1} matrices: A=B=1. + Another example, for the Netflix matrix: A=1, B=5 + + For those column pairs that are above the threshold, the + computed similarity is correct to within 20% relative error + with probability at least 1 - (0.981)^10/B^ + + The shuffle size is bounded by the *smaller* of the following + two expressions: + + - O(n log(n) L / (threshold * A)) + - O(m L^2^) + + The latter is the cost of the brute-force approach, so for + non-zero thresholds, the cost is always cheaper than the + brute-force approach. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + threshold : float, optional + Set to 0 for deterministic guaranteed + correctness. Similarities above this + threshold are estimated with the cost vs + estimate quality trade-off described above. + + Returns + ------- + :py:class:`CoordinateMatrix` + An n x n sparse upper-triangular CoordinateMatrix of + cosine similarities between columns of this matrix. + + Examples + -------- + >>> rows = sc.parallelize([[1, 2], [1, 5]]) + >>> mat = RowMatrix(rows) + + >>> sims = mat.columnSimilarities() + >>> sims.entries.first().value + 0.91914503... + """ + java_sims_mat = self._java_matrix_wrapper.call("columnSimilarities", float(threshold)) + return CoordinateMatrix(java_sims_mat) + + def tallSkinnyQR( + self, computeQ: bool = False + ) -> QRDecomposition[Optional["RowMatrix"], Matrix]: + """ + Compute the QR decomposition of this RowMatrix. + + The implementation is designed to optimize the QR decomposition + (factorization) for the RowMatrix of a tall and skinny shape [1]_. + + .. [1] Paul G. Constantine, David F. Gleich. "Tall and skinny QR + factorizations in MapReduce architectures" + https://doi.org/10.1145/1996092.1996103 + + .. versionadded:: 2.0.0 + + Parameters + ---------- + computeQ : bool, optional + whether to computeQ + + Returns + ------- + :py:class:`pyspark.mllib.linalg.QRDecomposition` + QRDecomposition(Q: RowMatrix, R: Matrix), where + Q = None if computeQ = false. + + Examples + -------- + >>> rows = sc.parallelize([[3, -6], [4, -8], [0, 1]]) + >>> mat = RowMatrix(rows) + >>> decomp = mat.tallSkinnyQR(True) + >>> Q = decomp.Q + >>> R = decomp.R + + >>> # Test with absolute values + >>> absQRows = Q.rows.map(lambda row: abs(row.toArray()).tolist()) + >>> absQRows.collect() + [[0.6..., 0.0], [0.8..., 0.0], [0.0, 1.0]] + + >>> # Test with absolute values + >>> abs(R.toArray()).tolist() + [[5.0, 10.0], [0.0, 1.0]] + """ + decomp = JavaModelWrapper(self._java_matrix_wrapper.call("tallSkinnyQR", computeQ)) + if computeQ: + java_Q = decomp.call("Q") + Q = RowMatrix(java_Q) + else: + Q = None + R = decomp.call("R") + return QRDecomposition(Q, R) + + def computeSVD( + self, k: int, computeU: bool = False, rCond: float = 1e-9 + ) -> "SingularValueDecomposition[RowMatrix, Matrix]": + """ + Computes the singular value decomposition of the RowMatrix. + + The given row matrix A of dimension (m X n) is decomposed into + U * s * V'T where + + - U: (m X k) (left singular vectors) is a RowMatrix whose + columns are the eigenvectors of (A X A') + - s: DenseVector consisting of square root of the eigenvalues + (singular values) in descending order. + - v: (n X k) (right singular vectors) is a Matrix whose columns + are the eigenvectors of (A' X A) + + For more specific details on implementation, please refer + the Scala documentation. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + k : int + Number of leading singular values to keep (`0 < k <= n`). + It might return less than k if there are numerically zero singular values + or there are not enough Ritz values converged before the maximum number of + Arnoldi update iterations is reached (in case that matrix A is ill-conditioned). + computeU : bool, optional + Whether or not to compute U. If set to be + True, then U is computed by A * V * s^-1 + rCond : float, optional + Reciprocal condition number. All singular values + smaller than rCond * s[0] are treated as zero + where s[0] is the largest singular value. + + Returns + ------- + :py:class:`SingularValueDecomposition` + + Examples + -------- + >>> rows = sc.parallelize([[3, 1, 1], [-1, 3, 1]]) + >>> rm = RowMatrix(rows) + + >>> svd_model = rm.computeSVD(2, True) + >>> svd_model.U.rows.collect() + [DenseVector([-0.7071, 0.7071]), DenseVector([-0.7071, -0.7071])] + >>> svd_model.s + DenseVector([3.4641, 3.1623]) + >>> svd_model.V + DenseMatrix(3, 2, [-0.4082, -0.8165, -0.4082, 0.8944, -0.4472, ...0.0], 0) + """ + j_model = self._java_matrix_wrapper.call("computeSVD", int(k), bool(computeU), float(rCond)) + return SingularValueDecomposition(j_model) + + def computePrincipalComponents(self, k: int) -> Matrix: + """ + Computes the k principal components of the given row matrix + + .. versionadded:: 2.2.0 + + Notes + ----- + This cannot be computed on matrices with more than 65535 columns. + + Parameters + ---------- + k : int + Number of principal components to keep. + + Returns + ------- + :py:class:`pyspark.mllib.linalg.DenseMatrix` + + Examples + -------- + >>> rows = sc.parallelize([[1, 2, 3], [2, 4, 5], [3, 6, 1]]) + >>> rm = RowMatrix(rows) + + >>> # Returns the two principal components of rm + >>> pca = rm.computePrincipalComponents(2) + >>> pca + DenseMatrix(3, 2, [-0.349, -0.6981, 0.6252, -0.2796, -0.5592, -0.7805], 0) + + >>> # Transform into new dimensions with the greatest variance. + >>> rm.multiply(pca).rows.collect() # doctest: +NORMALIZE_WHITESPACE + [DenseVector([0.1305, -3.7394]), DenseVector([-0.3642, -6.6983]), \ + DenseVector([-4.6102, -4.9745])] + """ + return self._java_matrix_wrapper.call("computePrincipalComponents", k) + + def multiply(self, matrix: Matrix) -> "RowMatrix": + """ + Multiply this matrix by a local dense matrix on the right. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + matrix : :py:class:`pyspark.mllib.linalg.Matrix` + a local dense matrix whose number of rows must match the number of columns + of this matrix + + Returns + ------- + :py:class:`RowMatrix` + + Examples + -------- + >>> rm = RowMatrix(sc.parallelize([[0, 1], [2, 3]])) + >>> rm.multiply(DenseMatrix(2, 2, [0, 2, 1, 3])).rows.collect() + [DenseVector([2.0, 3.0]), DenseVector([6.0, 11.0])] + """ + if not isinstance(matrix, DenseMatrix): + raise TypeError("Only multiplication with DenseMatrix is supported.") + j_model = self._java_matrix_wrapper.call("multiply", matrix) + return RowMatrix(j_model) + + +class SingularValueDecomposition(JavaModelWrapper, Generic[UT, VT]): + """ + Represents singular value decomposition (SVD) factors. + + .. versionadded:: 2.2.0 + """ + + @property + @since("2.2.0") + def U(self) -> Optional[UT]: # type: ignore[return] + """ + Returns a distributed matrix whose columns are the left + singular vectors of the SingularValueDecomposition if computeU was set to be True. + """ + u = self.call("U") + if u is not None: + mat_name = u.getClass().getSimpleName() + if mat_name == "RowMatrix": + return RowMatrix(u) # type: ignore[return-value] + elif mat_name == "IndexedRowMatrix": + return IndexedRowMatrix(u) # type: ignore[return-value] + else: + raise TypeError("Expected RowMatrix/IndexedRowMatrix got %s" % mat_name) + + @property + @since("2.2.0") + def s(self) -> Vector: + """ + Returns a DenseVector with singular values in descending order. + """ + return self.call("s") + + @property + @since("2.2.0") + def V(self) -> VT: + """ + Returns a DenseMatrix whose columns are the right singular + vectors of the SingularValueDecomposition. + """ + return self.call("V") + + +class IndexedRow: + """ + Represents a row of an IndexedRowMatrix. + + Just a wrapper over a (int, vector) tuple. + + Parameters + ---------- + index : int + The index for the given row. + vector : :py:class:`pyspark.mllib.linalg.Vector` or convertible + The row in the matrix at the given index. + """ + + def __init__(self, index: int, vector: "VectorLike") -> None: + self.index = int(index) + self.vector = _convert_to_vector(vector) + + def __repr__(self) -> str: + return "IndexedRow(%s, %s)" % (self.index, self.vector) + + +def _convert_to_indexed_row(row: Any) -> IndexedRow: + if isinstance(row, IndexedRow): + return row + elif isinstance(row, tuple) and len(row) == 2: + return IndexedRow(*row) + else: + raise TypeError("Cannot convert type %s into IndexedRow" % type(row)) + + +class IndexedRowMatrix(DistributedMatrix): + """ + Represents a row-oriented distributed Matrix with indexed rows. + + Parameters + ---------- + rows : :py:class:`pyspark.RDD` + An RDD of IndexedRows or (int, vector) tuples or a DataFrame consisting of a + int typed column of indices and a vector typed column. + numRows : int, optional + Number of rows in the matrix. A non-positive + value means unknown, at which point the number + of rows will be determined by the max row + index plus one. + numCols : int, optional + Number of columns in the matrix. A non-positive + value means unknown, at which point the number + of columns will be determined by the size of + the first row. + """ + + def __init__( + self, + rows: RDD[Union[Tuple[int, "VectorLike"], IndexedRow]], + numRows: int = 0, + numCols: int = 0, + ): + """ + Note: This docstring is not shown publicly. + + Create a wrapper over a Java IndexedRowMatrix. + + Publicly, we require that `rows` be an RDD or DataFrame. However, for + internal usage, `rows` can also be a Java IndexedRowMatrix + object, in which case we can wrap it directly. This + assists in clean matrix conversions. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(1, [4, 5, 6])]) + >>> mat = IndexedRowMatrix(rows) + + >>> mat_diff = IndexedRowMatrix(rows) + >>> (mat_diff._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + False + + >>> mat_same = IndexedRowMatrix(mat._java_matrix_wrapper._java_model) + >>> (mat_same._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + True + """ + if isinstance(rows, RDD): + rows = rows.map(_convert_to_indexed_row) + # We use DataFrames for serialization of IndexedRows from + # Python, so first convert the RDD to a DataFrame on this + # side. This will convert each IndexedRow to a Row + # containing the 'index' and 'vector' values, which can + # both be easily serialized. We will convert back to + # IndexedRows on the Scala side. + java_matrix = callMLlibFunc( + "createIndexedRowMatrix", rows.toDF(), int(numRows), int(numCols) + ) + elif isinstance(rows, DataFrame): + java_matrix = callMLlibFunc("createIndexedRowMatrix", rows, int(numRows), int(numCols)) + elif isinstance(rows, JavaObject) and rows.getClass().getSimpleName() == "IndexedRowMatrix": + java_matrix = rows + else: + raise TypeError( + "rows should be an RDD of IndexedRows or (int, vector) tuples, " + "got %s" % type(rows) + ) + + self._java_matrix_wrapper = JavaModelWrapper(java_matrix) + + @property + def rows(self) -> RDD[IndexedRow]: + """ + Rows of the IndexedRowMatrix stored as an RDD of IndexedRows. + + Examples + -------- + >>> mat = IndexedRowMatrix(sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(1, [4, 5, 6])])) + >>> rows = mat.rows + >>> rows.first() + IndexedRow(0, [1.0,2.0,3.0]) + """ + # We use DataFrames for serialization of IndexedRows from + # Java, so we first convert the RDD of rows to a DataFrame + # on the Scala/Java side. Then we map each Row in the + # DataFrame back to an IndexedRow on this side. + rows_df = callMLlibFunc("getIndexedRows", self._java_matrix_wrapper._java_model) + rows = rows_df.rdd.map(lambda row: IndexedRow(row[0], row[1])) + return rows + + def numRows(self) -> int: + """ + Get or compute the number of rows. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(1, [4, 5, 6]), + ... IndexedRow(2, [7, 8, 9]), + ... IndexedRow(3, [10, 11, 12])]) + + >>> mat = IndexedRowMatrix(rows) + >>> print(mat.numRows()) + 4 + + >>> mat = IndexedRowMatrix(rows, 7, 6) + >>> print(mat.numRows()) + 7 + """ + return self._java_matrix_wrapper.call("numRows") + + def numCols(self) -> int: + """ + Get or compute the number of cols. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(1, [4, 5, 6]), + ... IndexedRow(2, [7, 8, 9]), + ... IndexedRow(3, [10, 11, 12])]) + + >>> mat = IndexedRowMatrix(rows) + >>> print(mat.numCols()) + 3 + + >>> mat = IndexedRowMatrix(rows, 7, 6) + >>> print(mat.numCols()) + 6 + """ + return self._java_matrix_wrapper.call("numCols") + + def columnSimilarities(self) -> "CoordinateMatrix": + """ + Compute all cosine similarities between columns. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(6, [4, 5, 6])]) + >>> mat = IndexedRowMatrix(rows) + >>> cs = mat.columnSimilarities() + >>> print(cs.numCols()) + 3 + """ + java_coordinate_matrix = self._java_matrix_wrapper.call("columnSimilarities") + return CoordinateMatrix(java_coordinate_matrix) + + def computeGramianMatrix(self) -> Matrix: + """ + Computes the Gramian matrix `A^T A`. + + .. versionadded:: 2.0.0 + + Notes + ----- + This cannot be computed on matrices with more than 65535 columns. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(1, [4, 5, 6])]) + >>> mat = IndexedRowMatrix(rows) + + >>> mat.computeGramianMatrix() + DenseMatrix(3, 3, [17.0, 22.0, 27.0, 22.0, 29.0, 36.0, 27.0, 36.0, 45.0], 0) + """ + return self._java_matrix_wrapper.call("computeGramianMatrix") + + def toRowMatrix(self) -> RowMatrix: + """ + Convert this matrix to a RowMatrix. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(6, [4, 5, 6])]) + >>> mat = IndexedRowMatrix(rows).toRowMatrix() + >>> mat.rows.collect() + [DenseVector([1.0, 2.0, 3.0]), DenseVector([4.0, 5.0, 6.0])] + """ + java_row_matrix = self._java_matrix_wrapper.call("toRowMatrix") + return RowMatrix(java_row_matrix) + + def toCoordinateMatrix(self) -> "CoordinateMatrix": + """ + Convert this matrix to a CoordinateMatrix. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 0]), + ... IndexedRow(6, [0, 5])]) + >>> mat = IndexedRowMatrix(rows).toCoordinateMatrix() + >>> mat.entries.take(3) + [MatrixEntry(0, 0, 1.0), MatrixEntry(0, 1, 0.0), MatrixEntry(6, 0, 0.0)] + """ + java_coordinate_matrix = self._java_matrix_wrapper.call("toCoordinateMatrix") + return CoordinateMatrix(java_coordinate_matrix) + + def toBlockMatrix(self, rowsPerBlock: int = 1024, colsPerBlock: int = 1024) -> "BlockMatrix": + """ + Convert this matrix to a BlockMatrix. + + Parameters + ---------- + rowsPerBlock : int, optional + Number of rows that make up each block. + The blocks forming the final rows are not + required to have the given number of rows. + colsPerBlock : int, optional + Number of columns that make up each block. + The blocks forming the final columns are not + required to have the given number of columns. + + Examples + -------- + >>> rows = sc.parallelize([IndexedRow(0, [1, 2, 3]), + ... IndexedRow(6, [4, 5, 6])]) + >>> mat = IndexedRowMatrix(rows).toBlockMatrix() + + >>> # This IndexedRowMatrix will have 7 effective rows, due to + >>> # the highest row index being 6, and the ensuing + >>> # BlockMatrix will have 7 rows as well. + >>> print(mat.numRows()) + 7 + + >>> print(mat.numCols()) + 3 + """ + java_block_matrix = self._java_matrix_wrapper.call( + "toBlockMatrix", rowsPerBlock, colsPerBlock + ) + return BlockMatrix(java_block_matrix, rowsPerBlock, colsPerBlock) + + def computeSVD( + self, k: int, computeU: bool = False, rCond: float = 1e-9 + ) -> SingularValueDecomposition["IndexedRowMatrix", Matrix]: + """ + Computes the singular value decomposition of the IndexedRowMatrix. + + The given row matrix A of dimension (m X n) is decomposed into + U * s * V'T where + + * U: (m X k) (left singular vectors) is a IndexedRowMatrix + whose columns are the eigenvectors of (A X A') + * s: DenseVector consisting of square root of the eigenvalues + (singular values) in descending order. + * v: (n X k) (right singular vectors) is a Matrix whose columns + are the eigenvectors of (A' X A) + + For more specific details on implementation, please refer + the scala documentation. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + k : int + Number of leading singular values to keep (`0 < k <= n`). + It might return less than k if there are numerically zero singular values + or there are not enough Ritz values converged before the maximum number of + Arnoldi update iterations is reached (in case that matrix A is ill-conditioned). + computeU : bool, optional + Whether or not to compute U. If set to be + True, then U is computed by A * V * s^-1 + rCond : float, optional + Reciprocal condition number. All singular values + smaller than rCond * s[0] are treated as zero + where s[0] is the largest singular value. + + Returns + ------- + :py:class:`SingularValueDecomposition` + + Examples + -------- + >>> rows = [(0, (3, 1, 1)), (1, (-1, 3, 1))] + >>> irm = IndexedRowMatrix(sc.parallelize(rows)) + >>> svd_model = irm.computeSVD(2, True) + >>> svd_model.U.rows.collect() # doctest: +NORMALIZE_WHITESPACE + [IndexedRow(0, [-0.707106781187,0.707106781187]),\ + IndexedRow(1, [-0.707106781187,-0.707106781187])] + >>> svd_model.s + DenseVector([3.4641, 3.1623]) + >>> svd_model.V + DenseMatrix(3, 2, [-0.4082, -0.8165, -0.4082, 0.8944, -0.4472, ...0.0], 0) + """ + j_model = self._java_matrix_wrapper.call("computeSVD", int(k), bool(computeU), float(rCond)) + return SingularValueDecomposition(j_model) + + def multiply(self, matrix: Matrix) -> "IndexedRowMatrix": + """ + Multiply this matrix by a local dense matrix on the right. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + matrix : :py:class:`pyspark.mllib.linalg.Matrix` + a local dense matrix whose number of rows must match the number of columns + of this matrix + + Returns + ------- + :py:class:`IndexedRowMatrix` + + Examples + -------- + >>> mat = IndexedRowMatrix(sc.parallelize([(0, (0, 1)), (1, (2, 3))])) + >>> mat.multiply(DenseMatrix(2, 2, [0, 2, 1, 3])).rows.collect() + [IndexedRow(0, [2.0,3.0]), IndexedRow(1, [6.0,11.0])] + """ + if not isinstance(matrix, DenseMatrix): + raise TypeError("Only multiplication with DenseMatrix is supported.") + return IndexedRowMatrix(self._java_matrix_wrapper.call("multiply", matrix)) + + +class MatrixEntry: + """ + Represents an entry of a CoordinateMatrix. + + Just a wrapper over a (int, int, float) tuple. + + Parameters + ---------- + i : int + The row index of the matrix. + j : int + The column index of the matrix. + value : float + The (i, j)th entry of the matrix, as a float. + """ + + def __init__(self, i: int, j: int, value: float) -> None: + self.i = int(i) + self.j = int(j) + self.value = float(value) + + def __repr__(self) -> str: + return "MatrixEntry(%s, %s, %s)" % (self.i, self.j, self.value) + + +def _convert_to_matrix_entry(entry: Any) -> MatrixEntry: + if isinstance(entry, MatrixEntry): + return entry + elif isinstance(entry, tuple) and len(entry) == 3: + return MatrixEntry(*entry) + else: + raise TypeError("Cannot convert type %s into MatrixEntry" % type(entry)) + + +class CoordinateMatrix(DistributedMatrix): + """ + Represents a matrix in coordinate format. + + Parameters + ---------- + entries : :py:class:`pyspark.RDD` + An RDD of MatrixEntry inputs or + (int, int, float) tuples. + numRows : int, optional + Number of rows in the matrix. A non-positive + value means unknown, at which point the number + of rows will be determined by the max row + index plus one. + numCols : int, optional + Number of columns in the matrix. A non-positive + value means unknown, at which point the number + of columns will be determined by the max row + index plus one. + """ + + def __init__( + self, + entries: RDD[Union[Tuple[int, int, float], MatrixEntry]], + numRows: int = 0, + numCols: int = 0, + ): + """ + Note: This docstring is not shown publicly. + + Create a wrapper over a Java CoordinateMatrix. + + Publicly, we require that `rows` be an RDD. However, for + internal usage, `rows` can also be a Java CoordinateMatrix + object, in which case we can wrap it directly. This + assists in clean matrix conversions. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(6, 4, 2.1)]) + >>> mat = CoordinateMatrix(entries) + + >>> mat_diff = CoordinateMatrix(entries) + >>> (mat_diff._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + False + + >>> mat_same = CoordinateMatrix(mat._java_matrix_wrapper._java_model) + >>> (mat_same._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + True + """ + if isinstance(entries, RDD): + entries = entries.map(_convert_to_matrix_entry) + # We use DataFrames for serialization of MatrixEntry entries + # from Python, so first convert the RDD to a DataFrame on + # this side. This will convert each MatrixEntry to a Row + # containing the 'i', 'j', and 'value' values, which can + # each be easily serialized. We will convert back to + # MatrixEntry inputs on the Scala side. + java_matrix = callMLlibFunc( + "createCoordinateMatrix", entries.toDF(), int(numRows), int(numCols) + ) + elif ( + isinstance(entries, JavaObject) + and entries.getClass().getSimpleName() == "CoordinateMatrix" + ): + java_matrix = entries + else: + raise TypeError( + "entries should be an RDD of MatrixEntry entries or " + "(int, int, float) tuples, got %s" % type(entries) + ) + + self._java_matrix_wrapper = JavaModelWrapper(java_matrix) + + @property + def entries(self) -> RDD[MatrixEntry]: + """ + Entries of the CoordinateMatrix stored as an RDD of + MatrixEntries. + + Examples + -------- + >>> mat = CoordinateMatrix(sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(6, 4, 2.1)])) + >>> entries = mat.entries + >>> entries.first() + MatrixEntry(0, 0, 1.2) + """ + # We use DataFrames for serialization of MatrixEntry entries + # from Java, so we first convert the RDD of entries to a + # DataFrame on the Scala/Java side. Then we map each Row in + # the DataFrame back to a MatrixEntry on this side. + entries_df = callMLlibFunc("getMatrixEntries", self._java_matrix_wrapper._java_model) + entries = entries_df.rdd.map(lambda row: MatrixEntry(row[0], row[1], row[2])) + return entries + + def numRows(self) -> int: + """ + Get or compute the number of rows. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(1, 0, 2), + ... MatrixEntry(2, 1, 3.7)]) + + >>> mat = CoordinateMatrix(entries) + >>> print(mat.numRows()) + 3 + + >>> mat = CoordinateMatrix(entries, 7, 6) + >>> print(mat.numRows()) + 7 + """ + return self._java_matrix_wrapper.call("numRows") + + def numCols(self) -> int: + """ + Get or compute the number of cols. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(1, 0, 2), + ... MatrixEntry(2, 1, 3.7)]) + + >>> mat = CoordinateMatrix(entries) + >>> print(mat.numCols()) + 2 + + >>> mat = CoordinateMatrix(entries, 7, 6) + >>> print(mat.numCols()) + 6 + """ + return self._java_matrix_wrapper.call("numCols") + + def transpose(self) -> "CoordinateMatrix": + """ + Transpose this CoordinateMatrix. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(1, 0, 2), + ... MatrixEntry(2, 1, 3.7)]) + >>> mat = CoordinateMatrix(entries) + >>> mat_transposed = mat.transpose() + + >>> print(mat_transposed.numRows()) + 2 + + >>> print(mat_transposed.numCols()) + 3 + """ + java_transposed_matrix = self._java_matrix_wrapper.call("transpose") + return CoordinateMatrix(java_transposed_matrix) + + def toRowMatrix(self) -> RowMatrix: + """ + Convert this matrix to a RowMatrix. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(6, 4, 2.1)]) + >>> mat = CoordinateMatrix(entries).toRowMatrix() + + >>> # This CoordinateMatrix will have 7 effective rows, due to + >>> # the highest row index being 6, but the ensuing RowMatrix + >>> # will only have 2 rows since there are only entries on 2 + >>> # unique rows. + >>> print(mat.numRows()) + 2 + + >>> # This CoordinateMatrix will have 5 columns, due to the + >>> # highest column index being 4, and the ensuing RowMatrix + >>> # will have 5 columns as well. + >>> print(mat.numCols()) + 5 + """ + java_row_matrix = self._java_matrix_wrapper.call("toRowMatrix") + return RowMatrix(java_row_matrix) + + def toIndexedRowMatrix(self) -> IndexedRowMatrix: + """ + Convert this matrix to an IndexedRowMatrix. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(6, 4, 2.1)]) + >>> mat = CoordinateMatrix(entries).toIndexedRowMatrix() + + >>> # This CoordinateMatrix will have 7 effective rows, due to + >>> # the highest row index being 6, and the ensuing + >>> # IndexedRowMatrix will have 7 rows as well. + >>> print(mat.numRows()) + 7 + + >>> # This CoordinateMatrix will have 5 columns, due to the + >>> # highest column index being 4, and the ensuing + >>> # IndexedRowMatrix will have 5 columns as well. + >>> print(mat.numCols()) + 5 + """ + java_indexed_row_matrix = self._java_matrix_wrapper.call("toIndexedRowMatrix") + return IndexedRowMatrix(java_indexed_row_matrix) + + def toBlockMatrix(self, rowsPerBlock: int = 1024, colsPerBlock: int = 1024) -> "BlockMatrix": + """ + Convert this matrix to a BlockMatrix. + + Parameters + ---------- + rowsPerBlock : int, optional + Number of rows that make up each block. + The blocks forming the final rows are not + required to have the given number of rows. + colsPerBlock : int, optional + Number of columns that make up each block. + The blocks forming the final columns are not + required to have the given number of columns. + + Examples + -------- + >>> entries = sc.parallelize([MatrixEntry(0, 0, 1.2), + ... MatrixEntry(6, 4, 2.1)]) + >>> mat = CoordinateMatrix(entries).toBlockMatrix() + + >>> # This CoordinateMatrix will have 7 effective rows, due to + >>> # the highest row index being 6, and the ensuing + >>> # BlockMatrix will have 7 rows as well. + >>> print(mat.numRows()) + 7 + + >>> # This CoordinateMatrix will have 5 columns, due to the + >>> # highest column index being 4, and the ensuing + >>> # BlockMatrix will have 5 columns as well. + >>> print(mat.numCols()) + 5 + """ + java_block_matrix = self._java_matrix_wrapper.call( + "toBlockMatrix", rowsPerBlock, colsPerBlock + ) + return BlockMatrix(java_block_matrix, rowsPerBlock, colsPerBlock) + + +def _convert_to_matrix_block_tuple(block: Any) -> Tuple[Tuple[int, int], Matrix]: + if ( + isinstance(block, tuple) + and len(block) == 2 + and isinstance(block[0], tuple) + and len(block[0]) == 2 + and isinstance(block[1], Matrix) + ): + blockRowIndex = int(block[0][0]) + blockColIndex = int(block[0][1]) + subMatrix = block[1] + return ((blockRowIndex, blockColIndex), subMatrix) + else: + raise TypeError("Cannot convert type %s into a sub-matrix block tuple" % type(block)) + + +class BlockMatrix(DistributedMatrix): + """ + Represents a distributed matrix in blocks of local matrices. + + Parameters + ---------- + blocks : :py:class:`pyspark.RDD` + An RDD of sub-matrix blocks + ((blockRowIndex, blockColIndex), sub-matrix) that + form this distributed matrix. If multiple blocks + with the same index exist, the results for + operations like add and multiply will be + unpredictable. + rowsPerBlock : int + Number of rows that make up each block. + The blocks forming the final rows are not + required to have the given number of rows. + colsPerBlock : int + Number of columns that make up each block. + The blocks forming the final columns are not + required to have the given number of columns. + numRows : int, optional + Number of rows of this matrix. If the supplied + value is less than or equal to zero, the number + of rows will be calculated when `numRows` is + invoked. + numCols : int, optional + Number of columns of this matrix. If the supplied + value is less than or equal to zero, the number + of columns will be calculated when `numCols` is + invoked. + """ + + def __init__( + self, + blocks: RDD[Tuple[Tuple[int, int], Matrix]], + rowsPerBlock: int, + colsPerBlock: int, + numRows: int = 0, + numCols: int = 0, + ): + """ + Note: This docstring is not shown publicly. + + Create a wrapper over a Java BlockMatrix. + + Publicly, we require that `blocks` be an RDD. However, for + internal usage, `blocks` can also be a Java BlockMatrix + object, in which case we can wrap it directly. This + assists in clean matrix conversions. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + + >>> mat_diff = BlockMatrix(blocks, 3, 2) + >>> (mat_diff._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + False + + >>> mat_same = BlockMatrix(mat._java_matrix_wrapper._java_model, 3, 2) + >>> (mat_same._java_matrix_wrapper._java_model == + ... mat._java_matrix_wrapper._java_model) + True + """ + if isinstance(blocks, RDD): + blocks = blocks.map(_convert_to_matrix_block_tuple) + # We use DataFrames for serialization of sub-matrix blocks + # from Python, so first convert the RDD to a DataFrame on + # this side. This will convert each sub-matrix block + # tuple to a Row containing the 'blockRowIndex', + # 'blockColIndex', and 'subMatrix' values, which can + # each be easily serialized. We will convert back to + # ((blockRowIndex, blockColIndex), sub-matrix) tuples on + # the Scala side. + java_matrix = callMLlibFunc( + "createBlockMatrix", + blocks.toDF(), + int(rowsPerBlock), + int(colsPerBlock), + int(numRows), + int(numCols), + ) + elif isinstance(blocks, JavaObject) and blocks.getClass().getSimpleName() == "BlockMatrix": + java_matrix = blocks + else: + raise TypeError( + "blocks should be an RDD of sub-matrix blocks as " + "((int, int), matrix) tuples, got %s" % type(blocks) + ) + + self._java_matrix_wrapper = JavaModelWrapper(java_matrix) + + @property + def blocks(self) -> RDD[Tuple[Tuple[int, int], Matrix]]: + """ + The RDD of sub-matrix blocks + ((blockRowIndex, blockColIndex), sub-matrix) that form this + distributed matrix. + + Examples + -------- + >>> mat = BlockMatrix( + ... sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]), 3, 2) + >>> blocks = mat.blocks + >>> blocks.first() + ((0, 0), DenseMatrix(3, 2, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 0)) + + """ + # We use DataFrames for serialization of sub-matrix blocks + # from Java, so we first convert the RDD of blocks to a + # DataFrame on the Scala/Java side. Then we map each Row in + # the DataFrame back to a sub-matrix block on this side. + blocks_df = callMLlibFunc("getMatrixBlocks", self._java_matrix_wrapper._java_model) + blocks = blocks_df.rdd.map(lambda row: ((row[0][0], row[0][1]), row[1])) + return blocks + + @property + def rowsPerBlock(self) -> int: + """ + Number of rows that make up each block. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + >>> mat.rowsPerBlock + 3 + """ + return self._java_matrix_wrapper.call("rowsPerBlock") + + @property + def colsPerBlock(self) -> int: + """ + Number of columns that make up each block. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + >>> mat.colsPerBlock + 2 + """ + return self._java_matrix_wrapper.call("colsPerBlock") + + @property + def numRowBlocks(self) -> int: + """ + Number of rows of blocks in the BlockMatrix. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + >>> mat.numRowBlocks + 2 + """ + return self._java_matrix_wrapper.call("numRowBlocks") + + @property + def numColBlocks(self) -> int: + """ + Number of columns of blocks in the BlockMatrix. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + >>> mat.numColBlocks + 1 + """ + return self._java_matrix_wrapper.call("numColBlocks") + + def numRows(self) -> int: + """ + Get or compute the number of rows. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + + >>> mat = BlockMatrix(blocks, 3, 2) + >>> print(mat.numRows()) + 6 + + >>> mat = BlockMatrix(blocks, 3, 2, 7, 6) + >>> print(mat.numRows()) + 7 + """ + return self._java_matrix_wrapper.call("numRows") + + def numCols(self) -> int: + """ + Get or compute the number of cols. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + + >>> mat = BlockMatrix(blocks, 3, 2) + >>> print(mat.numCols()) + 2 + + >>> mat = BlockMatrix(blocks, 3, 2, 7, 6) + >>> print(mat.numCols()) + 6 + """ + return self._java_matrix_wrapper.call("numCols") + + @since("2.0.0") + def cache(self) -> "BlockMatrix": + """ + Caches the underlying RDD. + """ + self._java_matrix_wrapper.call("cache") + return self + + @since("2.0.0") + def persist(self, storageLevel: StorageLevel) -> "BlockMatrix": + """ + Persists the underlying RDD with the specified storage level. + """ + if not isinstance(storageLevel, StorageLevel): + raise TypeError("`storageLevel` should be a StorageLevel, got %s" % type(storageLevel)) + javaStorageLevel = self._java_matrix_wrapper._sc._getJavaStorageLevel(storageLevel) + self._java_matrix_wrapper.call("persist", javaStorageLevel) + return self + + @since("2.0.0") + def validate(self) -> None: + """ + Validates the block matrix info against the matrix data (`blocks`) + and throws an exception if any error is found. + """ + self._java_matrix_wrapper.call("validate") + + def add(self, other: "BlockMatrix") -> "BlockMatrix": + """ + Adds two block matrices together. The matrices must have the + same size and matching `rowsPerBlock` and `colsPerBlock` values. + If one of the sub matrix blocks that are being added is a + SparseMatrix, the resulting sub matrix block will also be a + SparseMatrix, even if it is being added to a DenseMatrix. If + two dense sub matrix blocks are added, the output block will + also be a DenseMatrix. + + Examples + -------- + >>> dm1 = Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6]) + >>> dm2 = Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]) + >>> sm = Matrices.sparse(3, 2, [0, 1, 3], [0, 1, 2], [7, 11, 12]) + >>> blocks1 = sc.parallelize([((0, 0), dm1), ((1, 0), dm2)]) + >>> blocks2 = sc.parallelize([((0, 0), dm1), ((1, 0), dm2)]) + >>> blocks3 = sc.parallelize([((0, 0), sm), ((1, 0), dm2)]) + >>> mat1 = BlockMatrix(blocks1, 3, 2) + >>> mat2 = BlockMatrix(blocks2, 3, 2) + >>> mat3 = BlockMatrix(blocks3, 3, 2) + + >>> mat1.add(mat2).toLocalMatrix() + DenseMatrix(6, 2, [2.0, 4.0, 6.0, 14.0, 16.0, 18.0, 8.0, 10.0, 12.0, 20.0, 22.0, 24.0], 0) + + >>> mat1.add(mat3).toLocalMatrix() + DenseMatrix(6, 2, [8.0, 2.0, 3.0, 14.0, 16.0, 18.0, 4.0, 16.0, 18.0, 20.0, 22.0, 24.0], 0) + """ + if not isinstance(other, BlockMatrix): + raise TypeError("Other should be a BlockMatrix, got %s" % type(other)) + + other_java_block_matrix = other._java_matrix_wrapper._java_model + java_block_matrix = self._java_matrix_wrapper.call("add", other_java_block_matrix) + return BlockMatrix(java_block_matrix, self.rowsPerBlock, self.colsPerBlock) + + def subtract(self, other: "BlockMatrix") -> "BlockMatrix": + """ + Subtracts the given block matrix `other` from this block matrix: + `this - other`. The matrices must have the same size and + matching `rowsPerBlock` and `colsPerBlock` values. If one of + the sub matrix blocks that are being subtracted is a + SparseMatrix, the resulting sub matrix block will also be a + SparseMatrix, even if it is being subtracted from a DenseMatrix. + If two dense sub matrix blocks are subtracted, the output block + will also be a DenseMatrix. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> dm1 = Matrices.dense(3, 2, [3, 1, 5, 4, 6, 2]) + >>> dm2 = Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]) + >>> sm = Matrices.sparse(3, 2, [0, 1, 3], [0, 1, 2], [1, 2, 3]) + >>> blocks1 = sc.parallelize([((0, 0), dm1), ((1, 0), dm2)]) + >>> blocks2 = sc.parallelize([((0, 0), dm2), ((1, 0), dm1)]) + >>> blocks3 = sc.parallelize([((0, 0), sm), ((1, 0), dm2)]) + >>> mat1 = BlockMatrix(blocks1, 3, 2) + >>> mat2 = BlockMatrix(blocks2, 3, 2) + >>> mat3 = BlockMatrix(blocks3, 3, 2) + + >>> mat1.subtract(mat2).toLocalMatrix() + DenseMatrix(6, 2, [-4.0, -7.0, -4.0, 4.0, 7.0, 4.0, -6.0, -5.0, -10.0, 6.0, 5.0, 10.0], 0) + + >>> mat2.subtract(mat3).toLocalMatrix() + DenseMatrix(6, 2, [6.0, 8.0, 9.0, -4.0, -7.0, -4.0, 10.0, 9.0, 9.0, -6.0, -5.0, -10.0], 0) + """ + if not isinstance(other, BlockMatrix): + raise TypeError("Other should be a BlockMatrix, got %s" % type(other)) + + other_java_block_matrix = other._java_matrix_wrapper._java_model + java_block_matrix = self._java_matrix_wrapper.call("subtract", other_java_block_matrix) + return BlockMatrix(java_block_matrix, self.rowsPerBlock, self.colsPerBlock) + + def multiply(self, other: "BlockMatrix") -> "BlockMatrix": + """ + Left multiplies this BlockMatrix by `other`, another + BlockMatrix. The `colsPerBlock` of this matrix must equal the + `rowsPerBlock` of `other`. If `other` contains any SparseMatrix + blocks, they will have to be converted to DenseMatrix blocks. + The output BlockMatrix will only consist of DenseMatrix blocks. + This may cause some performance issues until support for + multiplying two sparse matrices is added. + + Examples + -------- + >>> dm1 = Matrices.dense(2, 3, [1, 2, 3, 4, 5, 6]) + >>> dm2 = Matrices.dense(2, 3, [7, 8, 9, 10, 11, 12]) + >>> dm3 = Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6]) + >>> dm4 = Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]) + >>> sm = Matrices.sparse(3, 2, [0, 1, 3], [0, 1, 2], [7, 11, 12]) + >>> blocks1 = sc.parallelize([((0, 0), dm1), ((0, 1), dm2)]) + >>> blocks2 = sc.parallelize([((0, 0), dm3), ((1, 0), dm4)]) + >>> blocks3 = sc.parallelize([((0, 0), sm), ((1, 0), dm4)]) + >>> mat1 = BlockMatrix(blocks1, 2, 3) + >>> mat2 = BlockMatrix(blocks2, 3, 2) + >>> mat3 = BlockMatrix(blocks3, 3, 2) + + >>> mat1.multiply(mat2).toLocalMatrix() + DenseMatrix(2, 2, [242.0, 272.0, 350.0, 398.0], 0) + + >>> mat1.multiply(mat3).toLocalMatrix() + DenseMatrix(2, 2, [227.0, 258.0, 394.0, 450.0], 0) + """ + if not isinstance(other, BlockMatrix): + raise TypeError("Other should be a BlockMatrix, got %s" % type(other)) + + other_java_block_matrix = other._java_matrix_wrapper._java_model + java_block_matrix = self._java_matrix_wrapper.call("multiply", other_java_block_matrix) + return BlockMatrix(java_block_matrix, self.rowsPerBlock, self.colsPerBlock) + + def transpose(self) -> "BlockMatrix": + """ + Transpose this BlockMatrix. Returns a new BlockMatrix + instance sharing the same underlying data. Is a lazy operation. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2) + + >>> mat_transposed = mat.transpose() + >>> mat_transposed.toLocalMatrix() + DenseMatrix(2, 6, [1.0, 4.0, 2.0, 5.0, 3.0, 6.0, 7.0, 10.0, 8.0, 11.0, 9.0, 12.0], 0) + """ + java_transposed_matrix = self._java_matrix_wrapper.call("transpose") + return BlockMatrix(java_transposed_matrix, self.colsPerBlock, self.rowsPerBlock) + + def toLocalMatrix(self) -> Matrix: + """ + Collect the distributed matrix on the driver as a DenseMatrix. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2).toLocalMatrix() + + >>> # This BlockMatrix will have 6 effective rows, due to + >>> # having two sub-matrix blocks stacked, each with 3 rows. + >>> # The ensuing DenseMatrix will also have 6 rows. + >>> print(mat.numRows) + 6 + + >>> # This BlockMatrix will have 2 effective columns, due to + >>> # having two sub-matrix blocks stacked, each with 2 + >>> # columns. The ensuing DenseMatrix will also have 2 columns. + >>> print(mat.numCols) + 2 + """ + return self._java_matrix_wrapper.call("toLocalMatrix") + + def toIndexedRowMatrix(self) -> IndexedRowMatrix: + """ + Convert this matrix to an IndexedRowMatrix. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(3, 2, [1, 2, 3, 4, 5, 6])), + ... ((1, 0), Matrices.dense(3, 2, [7, 8, 9, 10, 11, 12]))]) + >>> mat = BlockMatrix(blocks, 3, 2).toIndexedRowMatrix() + + >>> # This BlockMatrix will have 6 effective rows, due to + >>> # having two sub-matrix blocks stacked, each with 3 rows. + >>> # The ensuing IndexedRowMatrix will also have 6 rows. + >>> print(mat.numRows()) + 6 + + >>> # This BlockMatrix will have 2 effective columns, due to + >>> # having two sub-matrix blocks stacked, each with 2 columns. + >>> # The ensuing IndexedRowMatrix will also have 2 columns. + >>> print(mat.numCols()) + 2 + """ + java_indexed_row_matrix = self._java_matrix_wrapper.call("toIndexedRowMatrix") + return IndexedRowMatrix(java_indexed_row_matrix) + + def toCoordinateMatrix(self) -> CoordinateMatrix: + """ + Convert this matrix to a CoordinateMatrix. + + Examples + -------- + >>> blocks = sc.parallelize([((0, 0), Matrices.dense(1, 2, [1, 2])), + ... ((1, 0), Matrices.dense(1, 2, [7, 8]))]) + >>> mat = BlockMatrix(blocks, 1, 2).toCoordinateMatrix() + >>> mat.entries.take(3) + [MatrixEntry(0, 0, 1.0), MatrixEntry(0, 1, 2.0), MatrixEntry(1, 0, 7.0)] + """ + java_coordinate_matrix = self._java_matrix_wrapper.call("toCoordinateMatrix") + return CoordinateMatrix(java_coordinate_matrix) + + +def _test() -> None: + import doctest + import numpy + from pyspark.sql import SparkSession + from pyspark.mllib.linalg import Matrices + import pyspark.mllib.linalg.distributed + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + globs = pyspark.mllib.linalg.distributed.__dict__.copy() + spark = ( + SparkSession.builder.master("local[2]") + .appName("mllib.linalg.distributed tests") + .getOrCreate() + ) + globs["sc"] = spark.sparkContext + globs["Matrices"] = Matrices + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/random.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/random.py new file mode 100644 index 0000000..1342148 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/random.py @@ -0,0 +1,698 @@ +# +# 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. +# + +""" +Python package for random data generation. +""" + +import sys +from functools import wraps +from typing import Any, Callable, Optional + +import numpy as np + +from pyspark.mllib.common import callMLlibFunc +from pyspark.context import SparkContext +from pyspark.rdd import RDD +from pyspark.mllib.linalg import Vector + + +__all__ = [ + "RandomRDDs", +] + + +def toArray(f: Callable[..., RDD[Vector]]) -> Callable[..., RDD[np.ndarray]]: + @wraps(f) + def func(sc: SparkContext, *a: Any, **kw: Any) -> RDD[np.ndarray]: + rdd = f(sc, *a, **kw) + return rdd.map(lambda vec: vec.toArray()) + + return func + + +class RandomRDDs: + """ + Generator methods for creating RDDs comprised of i.i.d samples from + some distribution. + + .. versionadded:: 1.1.0 + """ + + @staticmethod + def uniformRDD( + sc: SparkContext, size: int, numPartitions: Optional[int] = None, seed: Optional[int] = None + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the + uniform distribution U(0.0, 1.0). + + To transform the distribution in the generated RDD from U(0.0, 1.0) + to U(a, b), use + ``RandomRDDs.uniformRDD(sc, n, p, seed).map(lambda v: a + (b - a) * v)`` + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + used to create the RDD. + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of float comprised of i.i.d. samples ~ `U(0.0, 1.0)`. + + Examples + -------- + >>> x = RandomRDDs.uniformRDD(sc, 100).collect() + >>> len(x) + 100 + >>> max(x) <= 1.0 and min(x) >= 0.0 + True + >>> RandomRDDs.uniformRDD(sc, 100, 4).getNumPartitions() + 4 + >>> parts = RandomRDDs.uniformRDD(sc, 100, seed=4).getNumPartitions() + >>> parts == sc.defaultParallelism + True + """ + return callMLlibFunc("uniformRDD", sc._jsc, size, numPartitions, seed) + + @staticmethod + def normalRDD( + sc: SparkContext, size: int, numPartitions: Optional[int] = None, seed: Optional[int] = None + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the standard normal + distribution. + + To transform the distribution in the generated RDD from standard normal + to some other normal N(mean, sigma^2), use + ``RandomRDDs.normal(sc, n, p, seed).map(lambda v: mean + sigma * v)`` + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + used to create the RDD. + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of float comprised of i.i.d. samples ~ N(0.0, 1.0). + + Examples + -------- + >>> x = RandomRDDs.normalRDD(sc, 1000, seed=1) + >>> stats = x.stats() + >>> stats.count() + 1000 + >>> abs(stats.mean() - 0.0) < 0.1 + True + >>> abs(stats.stdev() - 1.0) < 0.1 + True + """ + return callMLlibFunc("normalRDD", sc._jsc, size, numPartitions, seed) + + @staticmethod + def logNormalRDD( + sc: SparkContext, + mean: float, + std: float, + size: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the log normal + distribution with the input mean and standard distribution. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + used to create the RDD. + mean : float + mean for the log Normal distribution + std : float + std for the log Normal distribution + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + RDD of float comprised of i.i.d. samples ~ log N(mean, std). + + Examples + -------- + >>> from math import sqrt, exp + >>> mean = 0.0 + >>> std = 1.0 + >>> expMean = exp(mean + 0.5 * std * std) + >>> expStd = sqrt((exp(std * std) - 1.0) * exp(2.0 * mean + std * std)) + >>> x = RandomRDDs.logNormalRDD(sc, mean, std, 1000, seed=2) + >>> stats = x.stats() + >>> stats.count() + 1000 + >>> abs(stats.mean() - expMean) < 0.5 + True + >>> from math import sqrt + >>> abs(stats.stdev() - expStd) < 0.5 + True + """ + return callMLlibFunc( + "logNormalRDD", sc._jsc, float(mean), float(std), size, numPartitions, seed + ) + + @staticmethod + def poissonRDD( + sc: SparkContext, + mean: float, + size: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the Poisson + distribution with the input mean. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + mean : float + Mean, or lambda, for the Poisson distribution. + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of float comprised of i.i.d. samples ~ Pois(mean). + + Examples + -------- + >>> mean = 100.0 + >>> x = RandomRDDs.poissonRDD(sc, mean, 1000, seed=2) + >>> stats = x.stats() + >>> stats.count() + 1000 + >>> abs(stats.mean() - mean) < 0.5 + True + >>> from math import sqrt + >>> abs(stats.stdev() - sqrt(mean)) < 0.5 + True + """ + return callMLlibFunc("poissonRDD", sc._jsc, float(mean), size, numPartitions, seed) + + @staticmethod + def exponentialRDD( + sc: SparkContext, + mean: float, + size: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the Exponential + distribution with the input mean. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + mean : float + Mean, or 1 / lambda, for the Exponential distribution. + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of float comprised of i.i.d. samples ~ Exp(mean). + + Examples + -------- + >>> mean = 2.0 + >>> x = RandomRDDs.exponentialRDD(sc, mean, 1000, seed=2) + >>> stats = x.stats() + >>> stats.count() + 1000 + >>> abs(stats.mean() - mean) < 0.5 + True + >>> from math import sqrt + >>> abs(stats.stdev() - sqrt(mean)) < 0.5 + True + """ + return callMLlibFunc("exponentialRDD", sc._jsc, float(mean), size, numPartitions, seed) + + @staticmethod + def gammaRDD( + sc: SparkContext, + shape: float, + scale: float, + size: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[float]: + """ + Generates an RDD comprised of i.i.d. samples from the Gamma + distribution with the input shape and scale. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + shape : float + shape (> 0) parameter for the Gamma distribution + scale : float + scale (> 0) parameter for the Gamma distribution + size : int + Size of the RDD. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of float comprised of i.i.d. samples ~ Gamma(shape, scale). + + Examples + -------- + >>> from math import sqrt + >>> shape = 1.0 + >>> scale = 2.0 + >>> expMean = shape * scale + >>> expStd = sqrt(shape * scale * scale) + >>> x = RandomRDDs.gammaRDD(sc, shape, scale, 1000, seed=2) + >>> stats = x.stats() + >>> stats.count() + 1000 + >>> abs(stats.mean() - expMean) < 0.5 + True + >>> abs(stats.stdev() - expStd) < 0.5 + True + """ + return callMLlibFunc( + "gammaRDD", sc._jsc, float(shape), float(scale), size, numPartitions, seed + ) + + @staticmethod + @toArray + def uniformVectorRDD( + sc: SparkContext, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the uniform distribution U(0.0, 1.0). + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + numRows : int + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD. + seed : int, optional + Seed for the RNG that generates the seed for the generator in each partition. + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d samples ~ `U(0.0, 1.0)`. + + Examples + -------- + >>> import numpy as np + >>> mat = np.matrix(RandomRDDs.uniformVectorRDD(sc, 10, 10).collect()) + >>> mat.shape + (10, 10) + >>> mat.max() <= 1.0 and mat.min() >= 0.0 + True + >>> RandomRDDs.uniformVectorRDD(sc, 10, 10, 4).getNumPartitions() + 4 + """ + return callMLlibFunc("uniformVectorRDD", sc._jsc, numRows, numCols, numPartitions, seed) + + @staticmethod + @toArray + def normalVectorRDD( + sc: SparkContext, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the standard normal distribution. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + numRows : int + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d. samples ~ `N(0.0, 1.0)`. + + Examples + -------- + >>> import numpy as np + >>> mat = np.matrix(RandomRDDs.normalVectorRDD(sc, 100, 100, seed=1).collect()) + >>> mat.shape + (100, 100) + >>> abs(mat.mean() - 0.0) < 0.1 + True + >>> abs(mat.std() - 1.0) < 0.1 + True + """ + return callMLlibFunc("normalVectorRDD", sc._jsc, numRows, numCols, numPartitions, seed) + + @staticmethod + @toArray + def logNormalVectorRDD( + sc: SparkContext, + mean: float, + std: float, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the log normal distribution. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + mean : float + Mean of the log normal distribution + std : float + Standard Deviation of the log normal distribution + numRows : int + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d. samples ~ log `N(mean, std)`. + + Examples + -------- + >>> import numpy as np + >>> from math import sqrt, exp + >>> mean = 0.0 + >>> std = 1.0 + >>> expMean = exp(mean + 0.5 * std * std) + >>> expStd = sqrt((exp(std * std) - 1.0) * exp(2.0 * mean + std * std)) + >>> m = RandomRDDs.logNormalVectorRDD(sc, mean, std, 100, 100, seed=1).collect() + >>> mat = np.matrix(m) + >>> mat.shape + (100, 100) + >>> abs(mat.mean() - expMean) < 0.1 + True + >>> abs(mat.std() - expStd) < 0.1 + True + """ + return callMLlibFunc( + "logNormalVectorRDD", + sc._jsc, + float(mean), + float(std), + numRows, + numCols, + numPartitions, + seed, + ) + + @staticmethod + @toArray + def poissonVectorRDD( + sc: SparkContext, + mean: float, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the Poisson distribution with the input mean. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + mean : float + Mean, or lambda, for the Poisson distribution. + numRows : float + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`) + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d. samples ~ Pois(mean). + + Examples + -------- + >>> import numpy as np + >>> mean = 100.0 + >>> rdd = RandomRDDs.poissonVectorRDD(sc, mean, 100, 100, seed=1) + >>> mat = np.mat(rdd.collect()) + >>> mat.shape + (100, 100) + >>> abs(mat.mean() - mean) < 0.5 + True + >>> from math import sqrt + >>> abs(mat.std() - sqrt(mean)) < 0.5 + True + """ + return callMLlibFunc( + "poissonVectorRDD", sc._jsc, float(mean), numRows, numCols, numPartitions, seed + ) + + @staticmethod + @toArray + def exponentialVectorRDD( + sc: SparkContext, + mean: float, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the Exponential distribution with the input mean. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + mean : float + Mean, or 1 / lambda, for the Exponential distribution. + numRows : int + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`) + seed : int, optional + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d. samples ~ Exp(mean). + + Examples + -------- + >>> import numpy as np + >>> mean = 0.5 + >>> rdd = RandomRDDs.exponentialVectorRDD(sc, mean, 100, 100, seed=1) + >>> mat = np.mat(rdd.collect()) + >>> mat.shape + (100, 100) + >>> abs(mat.mean() - mean) < 0.5 + True + >>> from math import sqrt + >>> abs(mat.std() - sqrt(mean)) < 0.5 + True + """ + return callMLlibFunc( + "exponentialVectorRDD", sc._jsc, float(mean), numRows, numCols, numPartitions, seed + ) + + @staticmethod + @toArray + def gammaVectorRDD( + sc: SparkContext, + shape: float, + scale: float, + numRows: int, + numCols: int, + numPartitions: Optional[int] = None, + seed: Optional[int] = None, + ) -> RDD[Vector]: + """ + Generates an RDD comprised of vectors containing i.i.d. samples drawn + from the Gamma distribution. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + SparkContext used to create the RDD. + shape : float + Shape (> 0) of the Gamma distribution + scale : float + Scale (> 0) of the Gamma distribution + numRows : int + Number of Vectors in the RDD. + numCols : int + Number of elements in each Vector. + numPartitions : int, optional + Number of partitions in the RDD (default: `sc.defaultParallelism`). + seed : int, optional, + Random seed (default: a random long integer). + + Returns + ------- + :py:class:`pyspark.RDD` + RDD of Vector with vectors containing i.i.d. samples ~ Gamma(shape, scale). + + Examples + -------- + >>> import numpy as np + >>> from math import sqrt + >>> shape = 1.0 + >>> scale = 2.0 + >>> expMean = shape * scale + >>> expStd = sqrt(shape * scale * scale) + >>> mat = np.matrix(RandomRDDs.gammaVectorRDD(sc, shape, scale, 100, 100, seed=1).collect()) + >>> mat.shape + (100, 100) + >>> abs(mat.mean() - expMean) < 0.1 + True + >>> abs(mat.std() - expStd) < 0.1 + True + """ + return callMLlibFunc( + "gammaVectorRDD", + sc._jsc, + float(shape), + float(scale), + numRows, + numCols, + numPartitions, + seed, + ) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + + globs = globals().copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("mllib.random tests").getOrCreate() + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/recommendation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/recommendation.py new file mode 100644 index 0000000..7ff8fdd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/recommendation.py @@ -0,0 +1,389 @@ +# +# 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 array +import sys +from typing import Any, List, NamedTuple, Optional, Tuple, Type, Union + +from pyspark import SparkContext, since +from pyspark.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc, inherit_doc +from pyspark.mllib.util import JavaLoader, JavaSaveable +from pyspark.sql import DataFrame + +__all__ = ["MatrixFactorizationModel", "ALS", "Rating"] + + +class Rating(NamedTuple): + """ + Represents a (user, product, rating) tuple. + + .. versionadded:: 1.2.0 + + Examples + -------- + >>> r = Rating(1, 2, 5.0) + >>> (r.user, r.product, r.rating) + (1, 2, 5.0) + >>> (r[0], r[1], r[2]) + (1, 2, 5.0) + """ + + user: int + product: int + rating: float + + def __reduce__(self) -> Tuple[Type["Rating"], Tuple[int, int, float]]: + return Rating, (int(self.user), int(self.product), float(self.rating)) + + +@inherit_doc +class MatrixFactorizationModel( + JavaModelWrapper, JavaSaveable, JavaLoader["MatrixFactorizationModel"] +): + + """A matrix factorisation model trained by regularized alternating + least-squares. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> r1 = (1, 1, 1.0) + >>> r2 = (1, 2, 2.0) + >>> r3 = (2, 1, 2.0) + >>> ratings = sc.parallelize([r1, r2, r3]) + >>> model = ALS.trainImplicit(ratings, 1, seed=10) + >>> model.predict(2, 2) + 0.4... + + >>> testset = sc.parallelize([(1, 2), (1, 1)]) + >>> model = ALS.train(ratings, 2, seed=0) + >>> model.predictAll(testset).collect() + [Rating(user=1, product=1, rating=1.0...), Rating(user=1, product=2, rating=1.9...)] + + >>> model = ALS.train(ratings, 4, seed=10) + >>> model.userFeatures().collect() + [(1, array('d', [...])), (2, array('d', [...]))] + + >>> model.recommendUsers(1, 2) + [Rating(user=2, product=1, rating=1.9...), Rating(user=1, product=1, rating=1.0...)] + >>> model.recommendProducts(1, 2) + [Rating(user=1, product=2, rating=1.9...), Rating(user=1, product=1, rating=1.0...)] + >>> model.rank + 4 + + >>> first_user = model.userFeatures().take(1)[0] + >>> latents = first_user[1] + >>> len(latents) + 4 + + >>> model.productFeatures().collect() + [(1, array('d', [...])), (2, array('d', [...]))] + + >>> first_product = model.productFeatures().take(1)[0] + >>> latents = first_product[1] + >>> len(latents) + 4 + + >>> products_for_users = model.recommendProductsForUsers(1).collect() + >>> len(products_for_users) + 2 + >>> products_for_users[0] + (1, (Rating(user=1, product=2, rating=...),)) + + >>> users_for_products = model.recommendUsersForProducts(1).collect() + >>> len(users_for_products) + 2 + >>> users_for_products[0] + (1, (Rating(user=2, product=1, rating=...),)) + + >>> model = ALS.train(ratings, 1, nonnegative=True, seed=123456789) + >>> model.predict(2, 2) + 3.73... + + >>> df = sqlContext.createDataFrame([Rating(1, 1, 1.0), Rating(1, 2, 2.0), Rating(2, 1, 2.0)]) + >>> model = ALS.train(df, 1, nonnegative=True, seed=123456789) + >>> model.predict(2, 2) + 3.73... + + >>> model = ALS.trainImplicit(ratings, 1, nonnegative=True, seed=123456789) + >>> model.predict(2, 2) + 0.4... + + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> model.save(sc, path) + >>> sameModel = MatrixFactorizationModel.load(sc, path) + >>> sameModel.predict(2, 2) + 0.4... + >>> sameModel.predictAll(testset).collect() + [Rating(... + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + @since("0.9.0") + def predict(self, user: int, product: int) -> float: + """ + Predicts rating for the given user and product. + """ + return self._java_model.predict(int(user), int(product)) + + @since("0.9.0") + def predictAll(self, user_product: RDD[Tuple[int, int]]) -> RDD[Rating]: + """ + Returns a list of predicted ratings for input user and product + pairs. + """ + assert isinstance(user_product, RDD), "user_product should be RDD of (user, product)" + first = user_product.first() + assert len(first) == 2, "user_product should be RDD of (user, product)" + user_product = user_product.map(lambda u_p: (int(u_p[0]), int(u_p[1]))) + return self.call("predict", user_product) + + @since("1.2.0") + def userFeatures(self) -> RDD[Tuple[int, array.array]]: + """ + Returns a paired RDD, where the first element is the user and the + second is an array of features corresponding to that user. + """ + return self.call("getUserFeatures").mapValues(lambda v: array.array("d", v)) + + @since("1.2.0") + def productFeatures(self) -> RDD[Tuple[int, array.array]]: + """ + Returns a paired RDD, where the first element is the product and the + second is an array of features corresponding to that product. + """ + return self.call("getProductFeatures").mapValues(lambda v: array.array("d", v)) + + @since("1.4.0") + def recommendUsers(self, product: int, num: int) -> List[Rating]: + """ + Recommends the top "num" number of users for a given product and + returns a list of Rating objects sorted by the predicted rating in + descending order. + """ + return list(self.call("recommendUsers", product, num)) + + @since("1.4.0") + def recommendProducts(self, user: int, num: int) -> List[Rating]: + """ + Recommends the top "num" number of products for a given user and + returns a list of Rating objects sorted by the predicted rating in + descending order. + """ + return list(self.call("recommendProducts", user, num)) + + def recommendProductsForUsers(self, num: int) -> RDD[Tuple[int, Tuple[Rating, ...]]]: + """ + Recommends the top "num" number of products for all users. The + number of recommendations returned per user may be less than "num". + """ + return self.call("wrappedRecommendProductsForUsers", num) + + def recommendUsersForProducts(self, num: int) -> RDD[Tuple[int, Tuple[Rating, ...]]]: + """ + Recommends the top "num" number of users for all products. The + number of recommendations returned per product may be less than + "num". + """ + return self.call("wrappedRecommendUsersForProducts", num) + + @property + @since("1.4.0") + def rank(self) -> int: + """Rank for the features in this model""" + return self.call("rank") + + @classmethod + @since("1.3.1") + def load(cls, sc: SparkContext, path: str) -> "MatrixFactorizationModel": + """Load a model from the given path""" + model = cls._load_java(sc, path) + assert sc._jvm is not None + wrapper = sc._jvm.org.apache.spark.mllib.api.python.MatrixFactorizationModelWrapper(model) + return MatrixFactorizationModel(wrapper) + + +class ALS: + """Alternating Least Squares matrix factorization + + .. versionadded:: 0.9.0 + """ + + @classmethod + def _prepare(cls, ratings: Any) -> RDD[Rating]: + if isinstance(ratings, RDD): + pass + elif isinstance(ratings, DataFrame): + ratings = ratings.rdd + else: + raise TypeError( + "Ratings should be represented by either an RDD or a DataFrame, " + "but got %s." % type(ratings) + ) + first = ratings.first() + if isinstance(first, Rating): + pass + elif isinstance(first, (tuple, list)): + ratings = ratings.map(lambda x: Rating(*x)) + else: + raise TypeError("Expect a Rating or a tuple/list, but got %s." % type(first)) + return ratings + + @classmethod + def train( + cls, + ratings: Union[RDD[Rating], RDD[Tuple[int, int, float]]], + rank: int, + iterations: int = 5, + lambda_: float = 0.01, + blocks: int = -1, + nonnegative: bool = False, + seed: Optional[int] = None, + ) -> MatrixFactorizationModel: + """ + Train a matrix factorization model given an RDD of ratings by users + for a subset of products. The ratings matrix is approximated as the + product of two lower-rank matrices of a given rank (number of + features). To solve for these features, ALS is run iteratively with + a configurable level of parallelism. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + ratings : :py:class:`pyspark.RDD` + RDD of `Rating` or (userID, productID, rating) tuple. + rank : int + Number of features to use (also referred to as the number of latent factors). + iterations : int, optional + Number of iterations of ALS. + (default: 5) + lambda\\_ : float, optional + Regularization parameter. + (default: 0.01) + blocks : int, optional + Number of blocks used to parallelize the computation. A value + of -1 will use an auto-configured number of blocks. + (default: -1) + nonnegative : bool, optional + A value of True will solve least-squares with nonnegativity + constraints. + (default: False) + seed : bool, optional + Random seed for initial matrix factorization model. A value + of None will use system time as the seed. + (default: None) + """ + model = callMLlibFunc( + "trainALSModel", + cls._prepare(ratings), + rank, + iterations, + lambda_, + blocks, + nonnegative, + seed, + ) + return MatrixFactorizationModel(model) + + @classmethod + def trainImplicit( + cls, + ratings: Union[RDD[Rating], RDD[Tuple[int, int, float]]], + rank: int, + iterations: int = 5, + lambda_: float = 0.01, + blocks: int = -1, + alpha: float = 0.01, + nonnegative: bool = False, + seed: Optional[int] = None, + ) -> MatrixFactorizationModel: + """ + Train a matrix factorization model given an RDD of 'implicit + preferences' of users for a subset of products. The ratings matrix + is approximated as the product of two lower-rank matrices of a + given rank (number of features). To solve for these features, ALS + is run iteratively with a configurable level of parallelism. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + ratings : :py:class:`pyspark.RDD` + RDD of `Rating` or (userID, productID, rating) tuple. + rank : int + Number of features to use (also referred to as the number of latent factors). + iterations : int, optional + Number of iterations of ALS. + (default: 5) + lambda\\_ : float, optional + Regularization parameter. + (default: 0.01) + blocks : int, optional + Number of blocks used to parallelize the computation. A value + of -1 will use an auto-configured number of blocks. + (default: -1) + alpha : float, optional + A constant used in computing confidence. + (default: 0.01) + nonnegative : bool, optional + A value of True will solve least-squares with nonnegativity + constraints. + (default: False) + seed : int, optional + Random seed for initial matrix factorization model. A value + of None will use system time as the seed. + (default: None) + """ + model = callMLlibFunc( + "trainImplicitALSModel", + cls._prepare(ratings), + rank, + iterations, + lambda_, + blocks, + alpha, + nonnegative, + seed, + ) + return MatrixFactorizationModel(model) + + +def _test() -> None: + import doctest + import pyspark.mllib.recommendation + from pyspark.sql import SQLContext + + globs = pyspark.mllib.recommendation.__dict__.copy() + sc = SparkContext("local[4]", "PythonTest") + globs["sc"] = sc + globs["sqlContext"] = SQLContext(sc) + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/regression.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/regression.py new file mode 100644 index 0000000..cac3294 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/regression.py @@ -0,0 +1,1067 @@ +# +# 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 sys +import warnings +from typing import ( + Any, + Callable, + Iterable, + Optional, + Tuple, + Type, + TypeVar, + Union, + overload, + TYPE_CHECKING, +) + +import numpy as np + +from pyspark import RDD, since +from pyspark.streaming.dstream import DStream +from pyspark.mllib.common import callMLlibFunc, _py2java, _java2py, inherit_doc +from pyspark.mllib.linalg import _convert_to_vector +from pyspark.mllib.util import Saveable, Loader +from pyspark.rdd import RDD +from pyspark.context import SparkContext +from pyspark.mllib.linalg import Vector + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike + + +LM = TypeVar("LM") +K = TypeVar("K") + +__all__ = [ + "LabeledPoint", + "LinearModel", + "LinearRegressionModel", + "LinearRegressionWithSGD", + "RidgeRegressionModel", + "RidgeRegressionWithSGD", + "LassoModel", + "LassoWithSGD", + "IsotonicRegressionModel", + "IsotonicRegression", + "StreamingLinearAlgorithm", + "StreamingLinearRegressionWithSGD", +] + + +class LabeledPoint: + + """ + Class that represents the features and labels of a data point. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + label : int + Label for this data point. + features : :py:class:`pyspark.mllib.linalg.Vector` or convertible + Vector of features for this point (NumPy array, list, + pyspark.mllib.linalg.SparseVector, or scipy.sparse column matrix). + + Notes + ----- + 'label' and 'features' are accessible as class attributes. + """ + + def __init__(self, label: float, features: Iterable[float]): + self.label = float(label) + self.features = _convert_to_vector(features) + + def __reduce__(self) -> Tuple[Type["LabeledPoint"], Tuple[float, Vector]]: + return (LabeledPoint, (self.label, self.features)) + + def __str__(self) -> str: + return "(" + ",".join((str(self.label), str(self.features))) + ")" + + def __repr__(self) -> str: + return "LabeledPoint(%s, %s)" % (self.label, self.features) + + +class LinearModel: + + """ + A linear model that has a vector of coefficients and an intercept. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + weights : :py:class:`pyspark.mllib.linalg.Vector` + Weights computed for every feature. + intercept : float + Intercept computed for this model. + """ + + def __init__(self, weights: Vector, intercept: float): + self._coeff = _convert_to_vector(weights) + self._intercept = float(intercept) + + @property + @since("1.0.0") + def weights(self) -> Vector: + """Weights computed for every feature.""" + return self._coeff + + @property + @since("1.0.0") + def intercept(self) -> float: + """Intercept computed for this model.""" + return self._intercept + + def __repr__(self) -> str: + return "(weights=%s, intercept=%r)" % (self._coeff, self._intercept) + + +@inherit_doc +class LinearRegressionModelBase(LinearModel): + + """A linear regression model. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> lrmb = LinearRegressionModelBase(np.array([1.0, 2.0]), 0.1) + >>> abs(lrmb.predict(np.array([-1.03, 7.777])) - 14.624) < 1e-6 + True + >>> abs(lrmb.predict(SparseVector(2, {0: -1.03, 1: 7.777})) - 14.624) < 1e-6 + True + """ + + @overload + def predict(self, x: "VectorLike") -> float: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[float]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]: + """ + Predict the value of the dependent variable given a vector or + an RDD of vectors containing values for the independent variables. + + .. versionadded:: 0.9.0 + """ + if isinstance(x, RDD): + return x.map(self.predict) + x = _convert_to_vector(x) + return self.weights.dot(x) + self.intercept # type: ignore[attr-defined] + + +@inherit_doc +class LinearRegressionModel(LinearRegressionModelBase): + + """A linear regression model derived from a least-squares fit. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> from pyspark.mllib.regression import LabeledPoint + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(1.0, [1.0]), + ... LabeledPoint(3.0, [2.0]), + ... LabeledPoint(2.0, [3.0]) + ... ] + >>> lrm = LinearRegressionWithSGD.train(sc.parallelize(data), iterations=10, + ... initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> abs(lrm.predict(sc.parallelize([[1.0]])).collect()[0] - 1) < 0.5 + True + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> lrm.save(sc, path) + >>> sameModel = LinearRegressionModel.load(sc, path) + >>> abs(sameModel.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(sameModel.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except BaseException: + ... pass + >>> data = [ + ... LabeledPoint(0.0, SparseVector(1, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(1, {0: 1.0})), + ... LabeledPoint(3.0, SparseVector(1, {0: 2.0})), + ... LabeledPoint(2.0, SparseVector(1, {0: 3.0})) + ... ] + >>> lrm = LinearRegressionWithSGD.train(sc.parallelize(data), iterations=10, + ... initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> lrm = LinearRegressionWithSGD.train(sc.parallelize(data), iterations=10, step=1.0, + ... miniBatchFraction=1.0, initialWeights=np.array([1.0]), regParam=0.1, regType="l2", + ... intercept=True, validateData=True) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + """ + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """Save a LinearRegressionModel.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.LinearRegressionModel( + _py2java(sc, self._coeff), self.intercept + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "LinearRegressionModel": + """Load a LinearRegressionModel.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.LinearRegressionModel.load( + sc._jsc.sc(), path + ) + weights = _java2py(sc, java_model.weights()) + intercept = java_model.intercept() + model = LinearRegressionModel(weights, intercept) + return model + + +# train_func should take two parameters, namely data and initial_weights, and +# return the result of a call to the appropriate JVM stub. +# _regression_train_wrapper is responsible for setup and error checking. +def _regression_train_wrapper( + train_func: Callable[[RDD[LabeledPoint], Vector], Iterable[Any]], + modelClass: Type[LM], + data: RDD[LabeledPoint], + initial_weights: Optional["VectorLike"], +) -> LM: + from pyspark.mllib.classification import LogisticRegressionModel + + first = data.first() + if not isinstance(first, LabeledPoint): + raise TypeError("data should be an RDD of LabeledPoint, but got %s" % type(first)) + if initial_weights is None: + initial_weights = [0.0] * len(data.first().features) + if modelClass == LogisticRegressionModel: + weights, intercept, numFeatures, numClasses = train_func( + data, _convert_to_vector(initial_weights) + ) + return modelClass(weights, intercept, numFeatures, numClasses) # type: ignore[call-arg] + else: + weights, intercept = train_func(data, _convert_to_vector(initial_weights)) + return modelClass(weights, intercept) # type: ignore[call-arg] + + +class LinearRegressionWithSGD: + """ + Train a linear regression model with no regularization using Stochastic Gradient Descent. + + .. versionadded:: 0.9.0 + .. deprecated:: 2.0.0 + Use :py:class:`pyspark.ml.regression.LinearRegression`. + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + step: float = 1.0, + miniBatchFraction: float = 1.0, + initialWeights: Optional["VectorLike"] = None, + regParam: float = 0.0, + regType: Optional[str] = None, + intercept: bool = False, + validateData: bool = True, + convergenceTol: float = 0.001, + ) -> LinearRegressionModel: + """ + Train a linear regression model using Stochastic Gradient + Descent (SGD). This solves the least squares regression + formulation + + f(weights) = 1/(2n) ||A weights - y||^2 + + which is the mean squared error. Here the data matrix has n rows, + and the input RDD holds the set of rows of A, each with its + corresponding right hand side label y. + See also the documentation for the precise formulation. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of LabeledPoint. + iterations : int, optional + The number of iterations. + (default: 100) + step : float, optional + The step parameter used in SGD. + (default: 1.0) + miniBatchFraction : float, optional + Fraction of data to be used for each SGD iteration. + (default: 1.0) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + regParam : float, optional + The regularizer parameter. + (default: 0.0) + regType : str, optional + The type of regularizer used for training our model. + Supported values: + + - "l1" for using L1 regularization + - "l2" for using L2 regularization + - None for no regularization (default) + + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e., whether bias + features are activated or not). + (default: False) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + convergenceTol : float, optional + A condition which decides iteration termination. + (default: 0.001) + """ + warnings.warn("Deprecated in 2.0.0. Use ml.regression.LinearRegression.", FutureWarning) + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainLinearRegressionModelWithSGD", + rdd, + int(iterations), + float(step), + float(miniBatchFraction), + i, + float(regParam), + regType, + bool(intercept), + bool(validateData), + float(convergenceTol), + ) + + return _regression_train_wrapper(train, LinearRegressionModel, data, initialWeights) + + +@inherit_doc +class LassoModel(LinearRegressionModelBase): + + """A linear regression model derived from a least-squares fit with + an l_1 penalty term. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> from pyspark.mllib.regression import LabeledPoint + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(1.0, [1.0]), + ... LabeledPoint(3.0, [2.0]), + ... LabeledPoint(2.0, [3.0]) + ... ] + >>> lrm = LassoWithSGD.train( + ... sc.parallelize(data), iterations=10, initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> abs(lrm.predict(sc.parallelize([[1.0]])).collect()[0] - 1) < 0.5 + True + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> lrm.save(sc, path) + >>> sameModel = LassoModel.load(sc, path) + >>> abs(sameModel.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(sameModel.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except BaseException: + ... pass + >>> data = [ + ... LabeledPoint(0.0, SparseVector(1, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(1, {0: 1.0})), + ... LabeledPoint(3.0, SparseVector(1, {0: 2.0})), + ... LabeledPoint(2.0, SparseVector(1, {0: 3.0})) + ... ] + >>> lrm = LinearRegressionWithSGD.train(sc.parallelize(data), iterations=10, + ... initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> lrm = LassoWithSGD.train(sc.parallelize(data), iterations=10, step=1.0, + ... regParam=0.01, miniBatchFraction=1.0, initialWeights=np.array([1.0]), intercept=True, + ... validateData=True) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + """ + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """Save a LassoModel.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.LassoModel( + _py2java(sc, self._coeff), self.intercept + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "LassoModel": + """Load a LassoModel.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.LassoModel.load(sc._jsc.sc(), path) + weights = _java2py(sc, java_model.weights()) + intercept = java_model.intercept() + model = LassoModel(weights, intercept) + return model + + +class LassoWithSGD: + """ + Train a regression model with L1-regularization using Stochastic Gradient Descent. + + .. versionadded:: 0.9.0 + .. deprecated:: 2.0.0 + Use :py:class:`pyspark.ml.regression.LinearRegression` with elasticNetParam = 1.0. + Note the default regParam is 0.01 for LassoWithSGD, but is 0.0 for LinearRegression. + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + step: float = 1.0, + regParam: float = 0.01, + miniBatchFraction: float = 1.0, + initialWeights: Optional["VectorLike"] = None, + intercept: bool = False, + validateData: bool = True, + convergenceTol: float = 0.001, + ) -> LassoModel: + """ + Train a regression model with L1-regularization using Stochastic + Gradient Descent. This solves the l1-regularized least squares + regression formulation + + f(weights) = 1/(2n) ||A weights - y||^2 + regParam ||weights||_1 + + Here the data matrix has n rows, and the input RDD holds the set + of rows of A, each with its corresponding right hand side label y. + See also the documentation for the precise formulation. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of LabeledPoint. + iterations : int, optional + The number of iterations. + (default: 100) + step : float, optional + The step parameter used in SGD. + (default: 1.0) + regParam : float, optional + The regularizer parameter. + (default: 0.01) + miniBatchFraction : float, optional + Fraction of data to be used for each SGD iteration. + (default: 1.0) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e. whether bias + features are activated or not). + (default: False) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + convergenceTol : float, optional + A condition which decides iteration termination. + (default: 0.001) + """ + warnings.warn( + "Deprecated in 2.0.0. Use ml.regression.LinearRegression with elasticNetParam = 1.0. " + "Note the default regParam is 0.01 for LassoWithSGD, but is 0.0 for LinearRegression.", + FutureWarning, + ) + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainLassoModelWithSGD", + rdd, + int(iterations), + float(step), + float(regParam), + float(miniBatchFraction), + i, + bool(intercept), + bool(validateData), + float(convergenceTol), + ) + + return _regression_train_wrapper(train, LassoModel, data, initialWeights) + + +@inherit_doc +class RidgeRegressionModel(LinearRegressionModelBase): + + """A linear regression model derived from a least-squares fit with + an l_2 penalty term. + + .. versionadded:: 0.9.0 + + Examples + -------- + >>> from pyspark.mllib.linalg import SparseVector + >>> from pyspark.mllib.regression import LabeledPoint + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(1.0, [1.0]), + ... LabeledPoint(3.0, [2.0]), + ... LabeledPoint(2.0, [3.0]) + ... ] + >>> lrm = RidgeRegressionWithSGD.train(sc.parallelize(data), iterations=10, + ... initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> abs(lrm.predict(sc.parallelize([[1.0]])).collect()[0] - 1) < 0.5 + True + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> lrm.save(sc, path) + >>> sameModel = RidgeRegressionModel.load(sc, path) + >>> abs(sameModel.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(sameModel.predict(np.array([1.0])) - 1) < 0.5 + True + >>> abs(sameModel.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except BaseException: + ... pass + >>> data = [ + ... LabeledPoint(0.0, SparseVector(1, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(1, {0: 1.0})), + ... LabeledPoint(3.0, SparseVector(1, {0: 2.0})), + ... LabeledPoint(2.0, SparseVector(1, {0: 3.0})) + ... ] + >>> lrm = LinearRegressionWithSGD.train(sc.parallelize(data), iterations=10, + ... initialWeights=np.array([1.0])) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + >>> lrm = RidgeRegressionWithSGD.train(sc.parallelize(data), iterations=10, step=1.0, + ... regParam=0.01, miniBatchFraction=1.0, initialWeights=np.array([1.0]), intercept=True, + ... validateData=True) + >>> abs(lrm.predict(np.array([0.0])) - 0) < 0.5 + True + >>> abs(lrm.predict(SparseVector(1, {0: 1.0})) - 1) < 0.5 + True + """ + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """Save a RidgeRegressionMode.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.RidgeRegressionModel( + _py2java(sc, self._coeff), self.intercept + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "RidgeRegressionModel": + """Load a RidgeRegressionMode.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.RidgeRegressionModel.load( + sc._jsc.sc(), path + ) + weights = _java2py(sc, java_model.weights()) + intercept = java_model.intercept() + model = RidgeRegressionModel(weights, intercept) + return model + + +class RidgeRegressionWithSGD: + """ + Train a regression model with L2-regularization using Stochastic Gradient Descent. + + .. versionadded:: 0.9.0 + .. deprecated:: 2.0.0 + Use :py:class:`pyspark.ml.regression.LinearRegression` with elasticNetParam = 0.0. + Note the default regParam is 0.01 for RidgeRegressionWithSGD, but is 0.0 for + LinearRegression. + """ + + @classmethod + def train( + cls, + data: RDD[LabeledPoint], + iterations: int = 100, + step: float = 1.0, + regParam: float = 0.01, + miniBatchFraction: float = 1.0, + initialWeights: Optional["VectorLike"] = None, + intercept: bool = False, + validateData: bool = True, + convergenceTol: float = 0.001, + ) -> RidgeRegressionModel: + """ + Train a regression model with L2-regularization using Stochastic + Gradient Descent. This solves the l2-regularized least squares + regression formulation + + f(weights) = 1/(2n) ||A weights - y||^2 + regParam/2 ||weights||^2 + + Here the data matrix has n rows, and the input RDD holds the set + of rows of A, each with its corresponding right hand side label y. + See also the documentation for the precise formulation. + + .. versionadded:: 0.9.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + The training data, an RDD of LabeledPoint. + iterations : int, optional + The number of iterations. + (default: 100) + step : float, optional + The step parameter used in SGD. + (default: 1.0) + regParam : float, optional + The regularizer parameter. + (default: 0.01) + miniBatchFraction : float, optional + Fraction of data to be used for each SGD iteration. + (default: 1.0) + initialWeights : :py:class:`pyspark.mllib.linalg.Vector` or convertible, optional + The initial weights. + (default: None) + intercept : bool, optional + Boolean parameter which indicates the use or not of the + augmented representation for training data (i.e. whether bias + features are activated or not). + (default: False) + validateData : bool, optional + Boolean parameter which indicates if the algorithm should + validate data before training. + (default: True) + convergenceTol : float, optional + A condition which decides iteration termination. + (default: 0.001) + """ + warnings.warn( + "Deprecated in 2.0.0. Use ml.regression.LinearRegression with elasticNetParam = 0.0. " + "Note the default regParam is 0.01 for RidgeRegressionWithSGD, but is 0.0 for " + "LinearRegression.", + FutureWarning, + ) + + def train(rdd: RDD[LabeledPoint], i: Vector) -> Iterable[Any]: + return callMLlibFunc( + "trainRidgeModelWithSGD", + rdd, + int(iterations), + float(step), + float(regParam), + float(miniBatchFraction), + i, + bool(intercept), + bool(validateData), + float(convergenceTol), + ) + + return _regression_train_wrapper(train, RidgeRegressionModel, data, initialWeights) + + +class IsotonicRegressionModel(Saveable, Loader["IsotonicRegressionModel"]): + + """ + Regression model for isotonic regression. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + boundaries : ndarray + Array of boundaries for which predictions are known. Boundaries + must be sorted in increasing order. + predictions : ndarray + Array of predictions associated to the boundaries at the same + index. Results of isotonic regression and therefore monotone. + isotonic : true + Indicates whether this is isotonic or antitonic. + + Examples + -------- + >>> data = [(1, 0, 1), (2, 1, 1), (3, 2, 1), (1, 3, 1), (6, 4, 1), (17, 5, 1), (16, 6, 1)] + >>> irm = IsotonicRegression.train(sc.parallelize(data)) + >>> irm.predict(3) + 2.0 + >>> irm.predict(5) + 16.5 + >>> irm.predict(sc.parallelize([3, 5])).collect() + [2.0, 16.5] + >>> import os, tempfile + >>> path = tempfile.mkdtemp() + >>> irm.save(sc, path) + >>> sameModel = IsotonicRegressionModel.load(sc, path) + >>> sameModel.predict(3) + 2.0 + >>> sameModel.predict(5) + 16.5 + >>> from shutil import rmtree + >>> try: + ... rmtree(path) + ... except OSError: + ... pass + """ + + def __init__(self, boundaries: np.ndarray, predictions: np.ndarray, isotonic: bool): + self.boundaries = boundaries + self.predictions = predictions + self.isotonic = isotonic + + @overload + def predict(self, x: float) -> np.float64: + ... + + @overload + def predict(self, x: "VectorLike") -> np.ndarray: + ... + + @overload + def predict(self, x: RDD[float]) -> RDD[np.float64]: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[np.ndarray]: + ... + + def predict( + self, x: Union[float, "VectorLike", RDD[float], RDD["VectorLike"]] + ) -> Union[np.float64, np.ndarray, RDD[np.float64], RDD[np.ndarray]]: + """ + Predict labels for provided features. + Using a piecewise linear function. + 1) If x exactly matches a boundary then associated prediction + is returned. In case there are multiple predictions with the + same boundary then one of them is returned. Which one is + undefined (same as java.util.Arrays.binarySearch). + 2) If x is lower or higher than all boundaries then first or + last prediction is returned respectively. In case there are + multiple predictions with the same boundary then the lowest + or highest is returned respectively. + 3) If x falls between two values in boundary array then + prediction is treated as piecewise linear function and + interpolated value is returned. In case there are multiple + values with the same boundary then the same rules as in 2) + are used. + + + .. versionadded:: 1.4.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Feature or RDD of Features to be labeled. + """ + if isinstance(x, RDD): + return x.map(lambda v: self.predict(v)) + return np.interp(x, self.boundaries, self.predictions) # type: ignore[arg-type] + + @since("1.4.0") + def save(self, sc: SparkContext, path: str) -> None: + """Save an IsotonicRegressionModel.""" + java_boundaries = _py2java(sc, self.boundaries.tolist()) + java_predictions = _py2java(sc, self.predictions.tolist()) + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.IsotonicRegressionModel( + java_boundaries, java_predictions, self.isotonic + ) + java_model.save(sc._jsc.sc(), path) + + @classmethod + @since("1.4.0") + def load(cls, sc: SparkContext, path: str) -> "IsotonicRegressionModel": + """Load an IsotonicRegressionModel.""" + assert sc._jvm is not None + + java_model = sc._jvm.org.apache.spark.mllib.regression.IsotonicRegressionModel.load( + sc._jsc.sc(), path + ) + py_boundaries = _java2py(sc, java_model.boundaryVector()).toArray() + py_predictions = _java2py(sc, java_model.predictionVector()).toArray() + return IsotonicRegressionModel(py_boundaries, py_predictions, java_model.isotonic) + + +class IsotonicRegression: + """ + Isotonic regression. + Currently implemented using parallelized pool adjacent violators + algorithm. Only univariate (single feature) algorithm supported. + + .. versionadded:: 1.4.0 + + Notes + ----- + Sequential PAV implementation based on + Tibshirani, Ryan J., Holger Hoefling, and Robert Tibshirani (2011) [1]_ + + Sequential PAV parallelization based on + Kearsley, Anthony J., Richard A. Tapia, and Michael W. Trosset (1996) [2]_ + + See also + `Isotonic regression (Wikipedia) `_. + + .. [1] Tibshirani, Ryan J., Holger Hoefling, and Robert Tibshirani. + "Nearly-isotonic regression." Technometrics 53.1 (2011): 54-61. + Available from http://www.stat.cmu.edu/~ryantibs/papers/neariso.pdf + .. [2] Kearsley, Anthony J., Richard A. Tapia, and Michael W. Trosset + "An approach to parallelizing isotonic regression." + Applied Mathematics and Parallel Computing. Physica-Verlag HD, 1996. 141-147. + Available from http://softlib.rice.edu/pub/CRPC-TRs/reports/CRPC-TR96640.pdf + """ + + @classmethod + def train(cls, data: RDD["VectorLike"], isotonic: bool = True) -> IsotonicRegressionModel: + """ + Train an isotonic regression model on the given data. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + RDD of (label, feature, weight) tuples. + isotonic : bool, optional + Whether this is isotonic (which is default) or antitonic. + (default: True) + """ + boundaries, predictions = callMLlibFunc( + "trainIsotonicRegressionModel", data.map(_convert_to_vector), bool(isotonic) + ) + return IsotonicRegressionModel(boundaries.toArray(), predictions.toArray(), isotonic) + + +class StreamingLinearAlgorithm: + """ + Base class that has to be inherited by any StreamingLinearAlgorithm. + + Prevents reimplementation of methods predictOn and predictOnValues. + + .. versionadded:: 1.5.0 + """ + + def __init__(self, model: Optional[LinearModel]): + self._model = model + + @since("1.5.0") + def latestModel(self) -> Optional[LinearModel]: + """ + Returns the latest model. + """ + return self._model + + def _validate(self, dstream: Any) -> None: + if not isinstance(dstream, DStream): + raise TypeError("dstream should be a DStream object, got %s" % type(dstream)) + if not self._model: + raise ValueError("Model must be initialized using setInitialWeights") + + def predictOn(self, dstream: "DStream[VectorLike]") -> "DStream[float]": + """ + Use the model to make predictions on batches of data from a + DStream. + + .. versionadded:: 1.5.0 + + Returns + ------- + :py:class:`pyspark.streaming.DStream` + DStream containing predictions. + """ + self._validate(dstream) + return dstream.map(lambda x: self._model.predict(x)) # type: ignore[union-attr] + + def predictOnValues( + self, dstream: "DStream[Tuple[K, VectorLike]]" + ) -> "DStream[Tuple[K, float]]": + """ + Use the model to make predictions on the values of a DStream and + carry over its keys. + + .. versionadded:: 1.5.0 + + Returns + ------- + :py:class:`pyspark.streaming.DStream` + DStream containing predictions. + """ + self._validate(dstream) + return dstream.mapValues(lambda x: self._model.predict(x)) # type: ignore[union-attr] + + +@inherit_doc +class StreamingLinearRegressionWithSGD(StreamingLinearAlgorithm): + """ + Train or predict a linear regression model on streaming data. + Training uses Stochastic Gradient Descent to update the model + based on each new batch of incoming data from a DStream + (see `LinearRegressionWithSGD` for model equation). + + Each batch of data is assumed to be an RDD of LabeledPoints. + The number of data points per batch can vary, but the number + of features must be constant. An initial weight vector must + be provided. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + stepSize : float, optional + Step size for each iteration of gradient descent. + (default: 0.1) + numIterations : int, optional + Number of iterations run for each batch of data. + (default: 50) + miniBatchFraction : float, optional + Fraction of each batch of data to use for updates. + (default: 1.0) + convergenceTol : float, optional + Value used to determine when to terminate iterations. + (default: 0.001) + """ + + def __init__( + self, + stepSize: float = 0.1, + numIterations: int = 50, + miniBatchFraction: float = 1.0, + convergenceTol: float = 0.001, + ): + self.stepSize = stepSize + self.numIterations = numIterations + self.miniBatchFraction = miniBatchFraction + self.convergenceTol = convergenceTol + self._model: Optional[LinearModel] = None + super(StreamingLinearRegressionWithSGD, self).__init__(model=self._model) + + @since("1.5.0") + def setInitialWeights(self, initialWeights: "VectorLike") -> "StreamingLinearRegressionWithSGD": + """ + Set the initial value of weights. + + This must be set before running trainOn and predictOn + """ + initialWeights = _convert_to_vector(initialWeights) + self._model = LinearRegressionModel(initialWeights, 0) + return self + + @since("1.5.0") + def trainOn(self, dstream: "DStream[LabeledPoint]") -> None: + """Train the model on the incoming dstream.""" + self._validate(dstream) + + def update(rdd: RDD[LabeledPoint]) -> None: + # LinearRegressionWithSGD.train raises an error for an empty RDD. + if not rdd.isEmpty(): + assert self._model is not None + self._model = LinearRegressionWithSGD.train( + rdd, + self.numIterations, + self.stepSize, + self.miniBatchFraction, + self._model.weights, + intercept=self._model.intercept, # type: ignore[arg-type] + convergenceTol=self.convergenceTol, + ) + + dstream.foreachRDD(update) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.mllib.regression + + globs = pyspark.mllib.regression.__dict__.copy() + spark = SparkSession.builder.master("local[2]").appName("mllib.regression tests").getOrCreate() + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/KernelDensity.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/KernelDensity.py new file mode 100644 index 0000000..febf4fd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/KernelDensity.py @@ -0,0 +1,59 @@ +# +# 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 Iterable, Optional + +import numpy as np +from numpy import ndarray + +from pyspark.mllib.common import callMLlibFunc +from pyspark.rdd import RDD + + +class KernelDensity: + """ + Estimate probability density at required points given an RDD of samples + from the population. + + Examples + -------- + >>> kd = KernelDensity() + >>> sample = sc.parallelize([0.0, 1.0]) + >>> kd.setSample(sample) + >>> kd.estimate([0.0, 1.0]) + array([ 0.12938758, 0.12938758]) + """ + + def __init__(self) -> None: + self._bandwidth: float = 1.0 + self._sample: Optional[RDD[float]] = None + + def setBandwidth(self, bandwidth: float) -> None: + """Set bandwidth of each sample. Defaults to 1.0""" + self._bandwidth = bandwidth + + def setSample(self, sample: RDD[float]) -> None: + """Set sample points from the population. Should be a RDD""" + if not isinstance(sample, RDD): + raise TypeError("samples should be a RDD, received %s" % type(sample)) + self._sample = sample + + def estimate(self, points: Iterable[float]) -> ndarray: + """Estimate the probability density at points""" + points = list(points) + densities = callMLlibFunc("estimateKernelDensity", self._sample, self._bandwidth, points) + return np.asarray(densities) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/__init__.py new file mode 100644 index 0000000..dc7a6da --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/__init__.py @@ -0,0 +1,34 @@ +# +# 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. +# + +""" +Python package for statistical functions in MLlib. +""" + +from pyspark.mllib.stat._statistics import Statistics, MultivariateStatisticalSummary +from pyspark.mllib.stat.distribution import MultivariateGaussian +from pyspark.mllib.stat.test import ChiSqTestResult, KolmogorovSmirnovTestResult +from pyspark.mllib.stat.KernelDensity import KernelDensity + +__all__ = [ + "Statistics", + "MultivariateStatisticalSummary", + "ChiSqTestResult", + "KolmogorovSmirnovTestResult", + "MultivariateGaussian", + "KernelDensity", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/_statistics.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/_statistics.py new file mode 100644 index 0000000..a784e0e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/_statistics.py @@ -0,0 +1,409 @@ +# +# 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 sys +from typing import cast, overload, List, Optional, TYPE_CHECKING, Union + +from numpy import ndarray +from py4j.java_gateway import JavaObject + +from pyspark.rdd import RDD +from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper +from pyspark.mllib.linalg import Matrix, Vector, _convert_to_vector +from pyspark.mllib.regression import LabeledPoint +from pyspark.mllib.stat.test import ChiSqTestResult, KolmogorovSmirnovTestResult + +if TYPE_CHECKING: + from pyspark.mllib._typing import CorrMethodType, KolmogorovSmirnovTestDistNameType + +__all__ = ["MultivariateStatisticalSummary", "Statistics"] + + +class MultivariateStatisticalSummary(JavaModelWrapper): + + """ + Trait for multivariate statistical summary of a data matrix. + """ + + def mean(self) -> ndarray: + return cast(JavaObject, self.call("mean")).toArray() + + def variance(self) -> ndarray: + return cast(JavaObject, self.call("variance")).toArray() + + def count(self) -> int: + return int(self.call("count")) + + def numNonzeros(self) -> ndarray: + return cast(JavaObject, self.call("numNonzeros")).toArray() + + def max(self) -> ndarray: + return cast(JavaObject, self.call("max")).toArray() + + def min(self) -> ndarray: + return cast(JavaObject, self.call("min")).toArray() + + def normL1(self) -> ndarray: + return cast(JavaObject, self.call("normL1")).toArray() + + def normL2(self) -> ndarray: + return cast(JavaObject, self.call("normL2")).toArray() + + +class Statistics: + @staticmethod + def colStats(rdd: RDD[Vector]) -> MultivariateStatisticalSummary: + """ + Computes column-wise summary statistics for the input RDD[Vector]. + + Parameters + ---------- + rdd : :py:class:`pyspark.RDD` + an RDD[Vector] for which column-wise summary statistics + are to be computed. + + Returns + ------- + :class:`MultivariateStatisticalSummary` + object containing column-wise summary statistics. + + Examples + -------- + >>> from pyspark.mllib.linalg import Vectors + >>> rdd = sc.parallelize([Vectors.dense([2, 0, 0, -2]), + ... Vectors.dense([4, 5, 0, 3]), + ... Vectors.dense([6, 7, 0, 8])]) + >>> cStats = Statistics.colStats(rdd) + >>> cStats.mean() + array([ 4., 4., 0., 3.]) + >>> cStats.variance() + array([ 4., 13., 0., 25.]) + >>> cStats.count() + 3 + >>> cStats.numNonzeros() + array([ 3., 2., 0., 3.]) + >>> cStats.max() + array([ 6., 7., 0., 8.]) + >>> cStats.min() + array([ 2., 0., 0., -2.]) + """ + cStats = callMLlibFunc("colStats", rdd.map(_convert_to_vector)) + return MultivariateStatisticalSummary(cStats) + + @overload + @staticmethod + def corr(x: RDD[Vector], *, method: Optional["CorrMethodType"] = ...) -> Matrix: + ... + + @overload + @staticmethod + def corr(x: RDD[float], y: RDD[float], method: Optional["CorrMethodType"] = ...) -> float: + ... + + @staticmethod + def corr( + x: Union[RDD[Vector], RDD[float]], + y: Optional[RDD[float]] = None, + method: Optional["CorrMethodType"] = None, + ) -> Union[float, Matrix]: + """ + Compute the correlation (matrix) for the input RDD(s) using the + specified method. + Methods currently supported: `pearson (default), spearman`. + + If a single RDD of Vectors is passed in, a correlation matrix + comparing the columns in the input RDD is returned. Use `method` + to specify the method to be used for single RDD inout. + If two RDDs of floats are passed in, a single float is returned. + + Parameters + ---------- + x : :py:class:`pyspark.RDD` + an RDD of vector for which the correlation matrix is to be computed, + or an RDD of float of the same cardinality as y when y is specified. + y : :py:class:`pyspark.RDD`, optional + an RDD of float of the same cardinality as x. + method : str, optional + String specifying the method to use for computing correlation. + Supported: `pearson` (default), `spearman` + + Returns + ------- + :py:class:`pyspark.mllib.linalg.Matrix` + Correlation matrix comparing columns in x. + + Examples + -------- + >>> x = sc.parallelize([1.0, 0.0, -2.0], 2) + >>> y = sc.parallelize([4.0, 5.0, 3.0], 2) + >>> zeros = sc.parallelize([0.0, 0.0, 0.0], 2) + >>> abs(Statistics.corr(x, y) - 0.6546537) < 1e-7 + True + >>> Statistics.corr(x, y) == Statistics.corr(x, y, "pearson") + True + >>> Statistics.corr(x, y, "spearman") + 0.5 + >>> from math import isnan + >>> isnan(Statistics.corr(x, zeros)) + True + >>> from pyspark.mllib.linalg import Vectors + >>> rdd = sc.parallelize([Vectors.dense([1, 0, 0, -2]), Vectors.dense([4, 5, 0, 3]), + ... Vectors.dense([6, 7, 0, 8]), Vectors.dense([9, 0, 0, 1])]) + >>> pearsonCorr = Statistics.corr(rdd) + >>> print(str(pearsonCorr).replace('nan', 'NaN')) + [[ 1. 0.05564149 NaN 0.40047142] + [ 0.05564149 1. NaN 0.91359586] + [ NaN NaN 1. NaN] + [ 0.40047142 0.91359586 NaN 1. ]] + >>> spearmanCorr = Statistics.corr(rdd, method="spearman") + >>> print(str(spearmanCorr).replace('nan', 'NaN')) + [[ 1. 0.10540926 NaN 0.4 ] + [ 0.10540926 1. NaN 0.9486833 ] + [ NaN NaN 1. NaN] + [ 0.4 0.9486833 NaN 1. ]] + >>> try: + ... Statistics.corr(rdd, "spearman") + ... print("Method name as second argument without 'method=' shouldn't be allowed.") + ... except TypeError: + ... pass + """ + # Check inputs to determine whether a single value or a matrix is needed for output. + # Since it's legal for users to use the method name as the second argument, we need to + # check if y is used to specify the method name instead. + if type(y) == str: + raise TypeError("Use 'method=' to specify method name.") + + if not y: + return cast( + JavaObject, callMLlibFunc("corr", x.map(_convert_to_vector), method) + ).toArray() + else: + return cast( + float, + callMLlibFunc("corr", cast(RDD[float], x).map(float), y.map(float), method), + ) + + @overload + @staticmethod + def chiSqTest(observed: Matrix) -> ChiSqTestResult: + ... + + @overload + @staticmethod + def chiSqTest(observed: Vector, expected: Optional[Vector] = ...) -> ChiSqTestResult: + ... + + @overload + @staticmethod + def chiSqTest(observed: RDD[LabeledPoint]) -> List[ChiSqTestResult]: + ... + + @staticmethod + def chiSqTest( + observed: Union[Matrix, RDD[LabeledPoint], Vector], expected: Optional[Vector] = None + ) -> Union[ChiSqTestResult, List[ChiSqTestResult]]: + """ + If `observed` is Vector, conduct Pearson's chi-squared goodness + of fit test of the observed data against the expected distribution, + or against the uniform distribution (by default), with each category + having an expected frequency of `1 / len(observed)`. + + If `observed` is matrix, conduct Pearson's independence test on the + input contingency matrix, which cannot contain negative entries or + columns or rows that sum up to 0. + + If `observed` is an RDD of LabeledPoint, conduct Pearson's independence + test for every feature against the label across the input RDD. + For each feature, the (feature, label) pairs are converted into a + contingency matrix for which the chi-squared statistic is computed. + All label and feature values must be categorical. + + Parameters + ---------- + observed : :py:class:`pyspark.mllib.linalg.Vector` or \ + :py:class:`pyspark.mllib.linalg.Matrix` + it could be a vector containing the observed categorical + counts/relative frequencies, or the contingency matrix + (containing either counts or relative frequencies), + or an RDD of LabeledPoint containing the labeled dataset + with categorical features. Real-valued features will be + treated as categorical for each distinct value. + expected : :py:class:`pyspark.mllib.linalg.Vector` + Vector containing the expected categorical counts/relative + frequencies. `expected` is rescaled if the `expected` sum + differs from the `observed` sum. + + Returns + ------- + :py:class:`pyspark.mllib.stat.ChiSqTestResult` + object containing the test statistic, degrees + of freedom, p-value, the method used, and the null hypothesis. + + Notes + ----- + `observed` cannot contain negative values + + Examples + -------- + >>> from pyspark.mllib.linalg import Vectors, Matrices + >>> observed = Vectors.dense([4, 6, 5]) + >>> pearson = Statistics.chiSqTest(observed) + >>> print(pearson.statistic) + 0.4 + >>> pearson.degreesOfFreedom + 2 + >>> print(round(pearson.pValue, 4)) + 0.8187 + >>> pearson.method + 'pearson' + >>> pearson.nullHypothesis + 'observed follows the same distribution as expected.' + + >>> observed = Vectors.dense([21, 38, 43, 80]) + >>> expected = Vectors.dense([3, 5, 7, 20]) + >>> pearson = Statistics.chiSqTest(observed, expected) + >>> print(round(pearson.pValue, 4)) + 0.0027 + + >>> data = [40.0, 24.0, 29.0, 56.0, 32.0, 42.0, 31.0, 10.0, 0.0, 30.0, 15.0, 12.0] + >>> chi = Statistics.chiSqTest(Matrices.dense(3, 4, data)) + >>> print(round(chi.statistic, 4)) + 21.9958 + + >>> data = [LabeledPoint(0.0, Vectors.dense([0.5, 10.0])), + ... LabeledPoint(0.0, Vectors.dense([1.5, 20.0])), + ... LabeledPoint(1.0, Vectors.dense([1.5, 30.0])), + ... LabeledPoint(0.0, Vectors.dense([3.5, 30.0])), + ... LabeledPoint(0.0, Vectors.dense([3.5, 40.0])), + ... LabeledPoint(1.0, Vectors.dense([3.5, 40.0])),] + >>> rdd = sc.parallelize(data, 4) + >>> chi = Statistics.chiSqTest(rdd) + >>> print(chi[0].statistic) + 0.75 + >>> print(chi[1].statistic) + 1.5 + """ + if isinstance(observed, RDD): + if not isinstance(observed.first(), LabeledPoint): + raise ValueError("observed should be an RDD of LabeledPoint") + jmodels = callMLlibFunc("chiSqTest", observed) + return [ChiSqTestResult(m) for m in jmodels] + + if isinstance(observed, Matrix): + jmodel = callMLlibFunc("chiSqTest", observed) + else: + if expected and len(expected) != len(observed): + raise ValueError("`expected` should have same length with `observed`") + jmodel = callMLlibFunc("chiSqTest", _convert_to_vector(observed), expected) + return ChiSqTestResult(jmodel) + + @staticmethod + def kolmogorovSmirnovTest( + data: RDD[float], distName: "KolmogorovSmirnovTestDistNameType" = "norm", *params: float + ) -> KolmogorovSmirnovTestResult: + """ + Performs the Kolmogorov-Smirnov (KS) test for data sampled from + a continuous distribution. It tests the null hypothesis that + the data is generated from a particular distribution. + + The given data is sorted and the Empirical Cumulative + Distribution Function (ECDF) is calculated + which for a given point is the number of points having a CDF + value lesser than it divided by the total number of points. + + Since the data is sorted, this is a step function + that rises by (1 / length of data) for every ordered point. + + The KS statistic gives us the maximum distance between the + ECDF and the CDF. Intuitively if this statistic is large, the + probability that the null hypothesis is true becomes small. + For specific details of the implementation, please have a look + at the Scala documentation. + + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + RDD, samples from the data + distName : str, optional + string, currently only "norm" is supported. + (Normal distribution) to calculate the + theoretical distribution of the data. + params + additional values which need to be provided for + a certain distribution. + If not provided, the default values are used. + + Returns + ------- + :py:class:`pyspark.mllib.stat.KolmogorovSmirnovTestResult` + object containing the test statistic, degrees of freedom, p-value, + the method used, and the null hypothesis. + + Examples + -------- + >>> kstest = Statistics.kolmogorovSmirnovTest + >>> data = sc.parallelize([-1.0, 0.0, 1.0]) + >>> ksmodel = kstest(data, "norm") + >>> print(round(ksmodel.pValue, 3)) + 1.0 + >>> print(round(ksmodel.statistic, 3)) + 0.175 + >>> ksmodel.nullHypothesis + 'Sample follows theoretical distribution' + + >>> data = sc.parallelize([2.0, 3.0, 4.0]) + >>> ksmodel = kstest(data, "norm", 3.0, 1.0) + >>> print(round(ksmodel.pValue, 3)) + 1.0 + >>> print(round(ksmodel.statistic, 3)) + 0.175 + """ + if not isinstance(data, RDD): + raise TypeError("data should be an RDD, got %s." % type(data)) + if not isinstance(distName, str): + raise TypeError("distName should be a string, got %s." % type(distName)) + + param_list = [float(param) for param in params] + return KolmogorovSmirnovTestResult( + callMLlibFunc("kolmogorovSmirnovTest", data, distName, param_list) + ) + + +def _test() -> None: + import doctest + import numpy + from pyspark.sql import SparkSession + + try: + # Numpy 1.14+ changed it's string format. + numpy.set_printoptions(legacy="1.13") + except TypeError: + pass + globs = globals().copy() + spark = ( + SparkSession.builder.master("local[4]").appName("mllib.stat.statistics tests").getOrCreate() + ) + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/distribution.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/distribution.py new file mode 100644 index 0000000..95e94ce --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/distribution.py @@ -0,0 +1,39 @@ +# +# 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. +# + +__all__ = ["MultivariateGaussian"] + +from typing import NamedTuple + +from pyspark.mllib.linalg import Matrix, Vector + + +class MultivariateGaussian(NamedTuple): + + """Represents a (mu, sigma) tuple + + Examples + -------- + >>> m = MultivariateGaussian(Vectors.dense([11,12]),DenseMatrix(2, 2, (1.0, 3.0, 5.0, 2.0))) + >>> (m.mu, m.sigma.toArray()) + (DenseVector([11.0, 12.0]), array([[ 1., 5.],[ 3., 2.]])) + >>> (m[0], m[1]) + (DenseVector([11.0, 12.0]), array([[ 1., 5.],[ 3., 2.]])) + """ + + mu: Vector + sigma: Matrix diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/test.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/test.py new file mode 100644 index 0000000..b84ecfc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/stat/test.py @@ -0,0 +1,86 @@ +# +# 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 Generic, Tuple, TypeVar + +from pyspark.mllib.common import inherit_doc, JavaModelWrapper + + +__all__ = ["ChiSqTestResult", "KolmogorovSmirnovTestResult"] + +DF = TypeVar("DF", int, float, Tuple[int, ...], Tuple[float, ...]) + + +class TestResult(JavaModelWrapper, Generic[DF]): + """ + Base class for all test results. + """ + + @property + def pValue(self) -> float: + """ + The probability of obtaining a test statistic result at least as + extreme as the one that was actually observed, assuming that the + null hypothesis is true. + """ + return self._java_model.pValue() + + @property + def degreesOfFreedom(self) -> DF: + """ + Returns the degree(s) of freedom of the hypothesis test. + Return type should be Number(e.g. Int, Double) or tuples of Numbers. + """ + return self._java_model.degreesOfFreedom() + + @property + def statistic(self) -> float: + """ + Test statistic. + """ + return self._java_model.statistic() + + @property + def nullHypothesis(self) -> str: + """ + Null hypothesis of the test. + """ + return self._java_model.nullHypothesis() + + def __str__(self) -> str: + return self._java_model.toString() + + +@inherit_doc +class ChiSqTestResult(TestResult[int]): + """ + Contains test results for the chi-squared hypothesis test. + """ + + @property + def method(self) -> str: + """ + Name of the test method + """ + return self._java_model.method() + + +@inherit_doc +class KolmogorovSmirnovTestResult(TestResult[int]): + """ + Contains test results for the Kolmogorov-Smirnov test. + """ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_algorithms.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_algorithms.py new file mode 100644 index 0000000..dc48c2c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_algorithms.py @@ -0,0 +1,353 @@ +# +# 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 os +import tempfile +from shutil import rmtree +import unittest + +from numpy import array, array_equal +from py4j.protocol import Py4JJavaError + +from pyspark.mllib.fpm import FPGrowth +from pyspark.mllib.recommendation import Rating +from pyspark.mllib.regression import LabeledPoint +from pyspark.serializers import CPickleSerializer +from pyspark.testing.mllibutils import MLlibTestCase +from pyspark.testing.utils import eventually + + +class ListTests(MLlibTestCase): + + """ + Test MLlib algorithms on plain lists, to make sure they're passed through + as NumPy arrays. + """ + + def test_bisecting_kmeans(self): + from pyspark.mllib.clustering import BisectingKMeans + + data = array([0.0, 0.0, 1.0, 1.0, 9.0, 8.0, 8.0, 9.0]).reshape(4, 2) + bskm = BisectingKMeans() + model = bskm.train(self.sc.parallelize(data, 2), k=4) + p = array([0.0, 0.0]) + rdd_p = self.sc.parallelize([p]) + self.assertEqual(model.predict(p), model.predict(rdd_p).first()) + self.assertEqual(model.computeCost(p), model.computeCost(rdd_p)) + self.assertEqual(model.k, len(model.clusterCenters)) + + def test_kmeans(self): + from pyspark.mllib.clustering import KMeans + + data = [ + [0, 1.1], + [0, 1.2], + [1.1, 0], + [1.2, 0], + ] + clusters = KMeans.train( + self.sc.parallelize(data), + 2, + initializationMode="k-means||", + initializationSteps=7, + epsilon=1e-4, + ) + self.assertEqual(clusters.predict(data[0]), clusters.predict(data[1])) + self.assertEqual(clusters.predict(data[2]), clusters.predict(data[3])) + + def test_kmeans_deterministic(self): + from pyspark.mllib.clustering import KMeans + + X = range(0, 100, 10) + Y = range(0, 100, 10) + data = [[x, y] for x, y in zip(X, Y)] + clusters1 = KMeans.train( + self.sc.parallelize(data), + 3, + initializationMode="k-means||", + seed=42, + initializationSteps=7, + epsilon=1e-4, + ) + clusters2 = KMeans.train( + self.sc.parallelize(data), + 3, + initializationMode="k-means||", + seed=42, + initializationSteps=7, + epsilon=1e-4, + ) + centers1 = clusters1.centers + centers2 = clusters2.centers + for c1, c2 in zip(centers1, centers2): + # TODO: Allow small numeric difference. + self.assertTrue(array_equal(c1, c2)) + + def test_gmm(self): + from pyspark.mllib.clustering import GaussianMixture + + def condition(): + data = self.sc.parallelize( + [ + [1, 2], + [8, 9], + [-4, -3], + [-6, -7], + ] + ) + clusters = GaussianMixture.train( + data, 2, convergenceTol=0.001, maxIterations=10, seed=1 + ) + labels = clusters.predict(data).collect() + self.assertEqual(labels[0], labels[1]) + self.assertEqual(labels[2], labels[3]) + return True + + eventually(condition, timeout=60, catch_assertions=True) + + def test_gmm_deterministic(self): + from pyspark.mllib.clustering import GaussianMixture + + x = range(0, 100, 10) + y = range(0, 100, 10) + data = self.sc.parallelize([[a, b] for a, b in zip(x, y)]) + clusters1 = GaussianMixture.train(data, 5, convergenceTol=0.001, maxIterations=10, seed=63) + clusters2 = GaussianMixture.train(data, 5, convergenceTol=0.001, maxIterations=10, seed=63) + for c1, c2 in zip(clusters1.weights, clusters2.weights): + self.assertEqual(round(c1, 7), round(c2, 7)) + + def test_gmm_with_initial_model(self): + from pyspark.mllib.clustering import GaussianMixture + + data = self.sc.parallelize([(-10, -5), (-9, -4), (10, 5), (9, 4)]) + + gmm1 = GaussianMixture.train(data, 2, convergenceTol=0.001, maxIterations=10, seed=63) + gmm2 = GaussianMixture.train( + data, 2, convergenceTol=0.001, maxIterations=10, seed=63, initialModel=gmm1 + ) + self.assertAlmostEqual((gmm1.weights - gmm2.weights).sum(), 0.0) + + def test_classification(self): + from pyspark.mllib.classification import LogisticRegressionWithSGD, SVMWithSGD, NaiveBayes + from pyspark.mllib.tree import ( + DecisionTree, + DecisionTreeModel, + RandomForest, + RandomForestModel, + GradientBoostedTrees, + GradientBoostedTreesModel, + ) + + data = [ + LabeledPoint(0.0, [1, 0, 0]), + LabeledPoint(1.0, [0, 1, 1]), + LabeledPoint(0.0, [2, 0, 0]), + LabeledPoint(1.0, [0, 2, 1]), + ] + rdd = self.sc.parallelize(data) + features = [p.features.tolist() for p in data] + + temp_dir = tempfile.mkdtemp() + + lr_model = LogisticRegressionWithSGD.train(rdd, iterations=10) + self.assertTrue(lr_model.predict(features[0]) <= 0) + self.assertTrue(lr_model.predict(features[1]) > 0) + self.assertTrue(lr_model.predict(features[2]) <= 0) + self.assertTrue(lr_model.predict(features[3]) > 0) + + svm_model = SVMWithSGD.train(rdd, iterations=10) + self.assertTrue(svm_model.predict(features[0]) <= 0) + self.assertTrue(svm_model.predict(features[1]) > 0) + self.assertTrue(svm_model.predict(features[2]) <= 0) + self.assertTrue(svm_model.predict(features[3]) > 0) + + nb_model = NaiveBayes.train(rdd) + self.assertTrue(nb_model.predict(features[0]) <= 0) + self.assertTrue(nb_model.predict(features[1]) > 0) + self.assertTrue(nb_model.predict(features[2]) <= 0) + self.assertTrue(nb_model.predict(features[3]) > 0) + + categoricalFeaturesInfo = {0: 3} # feature 0 has 3 categories + dt_model = DecisionTree.trainClassifier( + rdd, numClasses=2, categoricalFeaturesInfo=categoricalFeaturesInfo, maxBins=4 + ) + self.assertTrue(dt_model.predict(features[0]) <= 0) + self.assertTrue(dt_model.predict(features[1]) > 0) + self.assertTrue(dt_model.predict(features[2]) <= 0) + self.assertTrue(dt_model.predict(features[3]) > 0) + + dt_model_dir = os.path.join(temp_dir, "dt") + dt_model.save(self.sc, dt_model_dir) + same_dt_model = DecisionTreeModel.load(self.sc, dt_model_dir) + self.assertEqual(same_dt_model.toDebugString(), dt_model.toDebugString()) + + rf_model = RandomForest.trainClassifier( + rdd, + numClasses=2, + categoricalFeaturesInfo=categoricalFeaturesInfo, + numTrees=10, + maxBins=4, + seed=1, + ) + self.assertTrue(rf_model.predict(features[0]) <= 0) + self.assertTrue(rf_model.predict(features[1]) > 0) + self.assertTrue(rf_model.predict(features[2]) <= 0) + self.assertTrue(rf_model.predict(features[3]) > 0) + + rf_model_dir = os.path.join(temp_dir, "rf") + rf_model.save(self.sc, rf_model_dir) + same_rf_model = RandomForestModel.load(self.sc, rf_model_dir) + self.assertEqual(same_rf_model.toDebugString(), rf_model.toDebugString()) + + gbt_model = GradientBoostedTrees.trainClassifier( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, numIterations=4 + ) + self.assertTrue(gbt_model.predict(features[0]) <= 0) + self.assertTrue(gbt_model.predict(features[1]) > 0) + self.assertTrue(gbt_model.predict(features[2]) <= 0) + self.assertTrue(gbt_model.predict(features[3]) > 0) + + gbt_model_dir = os.path.join(temp_dir, "gbt") + gbt_model.save(self.sc, gbt_model_dir) + same_gbt_model = GradientBoostedTreesModel.load(self.sc, gbt_model_dir) + self.assertEqual(same_gbt_model.toDebugString(), gbt_model.toDebugString()) + + try: + rmtree(temp_dir) + except OSError: + pass + + def test_regression(self): + from pyspark.mllib.regression import ( + LinearRegressionWithSGD, + LassoWithSGD, + RidgeRegressionWithSGD, + ) + from pyspark.mllib.tree import DecisionTree, RandomForest, GradientBoostedTrees + + data = [ + LabeledPoint(-1.0, [0, -1]), + LabeledPoint(1.0, [0, 1]), + LabeledPoint(-1.0, [0, -2]), + LabeledPoint(1.0, [0, 2]), + ] + rdd = self.sc.parallelize(data) + features = [p.features.tolist() for p in data] + + lr_model = LinearRegressionWithSGD.train(rdd, iterations=10) + self.assertTrue(lr_model.predict(features[0]) <= 0) + self.assertTrue(lr_model.predict(features[1]) > 0) + self.assertTrue(lr_model.predict(features[2]) <= 0) + self.assertTrue(lr_model.predict(features[3]) > 0) + + lasso_model = LassoWithSGD.train(rdd, iterations=10) + self.assertTrue(lasso_model.predict(features[0]) <= 0) + self.assertTrue(lasso_model.predict(features[1]) > 0) + self.assertTrue(lasso_model.predict(features[2]) <= 0) + self.assertTrue(lasso_model.predict(features[3]) > 0) + + rr_model = RidgeRegressionWithSGD.train(rdd, iterations=10) + self.assertTrue(rr_model.predict(features[0]) <= 0) + self.assertTrue(rr_model.predict(features[1]) > 0) + self.assertTrue(rr_model.predict(features[2]) <= 0) + self.assertTrue(rr_model.predict(features[3]) > 0) + + categoricalFeaturesInfo = {0: 2} # feature 0 has 2 categories + dt_model = DecisionTree.trainRegressor( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, maxBins=4 + ) + self.assertTrue(dt_model.predict(features[0]) <= 0) + self.assertTrue(dt_model.predict(features[1]) > 0) + self.assertTrue(dt_model.predict(features[2]) <= 0) + self.assertTrue(dt_model.predict(features[3]) > 0) + + rf_model = RandomForest.trainRegressor( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, numTrees=10, maxBins=4, seed=1 + ) + self.assertTrue(rf_model.predict(features[0]) <= 0) + self.assertTrue(rf_model.predict(features[1]) > 0) + self.assertTrue(rf_model.predict(features[2]) <= 0) + self.assertTrue(rf_model.predict(features[3]) > 0) + + gbt_model = GradientBoostedTrees.trainRegressor( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, numIterations=4 + ) + self.assertTrue(gbt_model.predict(features[0]) <= 0) + self.assertTrue(gbt_model.predict(features[1]) > 0) + self.assertTrue(gbt_model.predict(features[2]) <= 0) + self.assertTrue(gbt_model.predict(features[3]) > 0) + + try: + LinearRegressionWithSGD.train(rdd, initialWeights=array([1.0, 1.0]), iterations=10) + LassoWithSGD.train(rdd, initialWeights=array([1.0, 1.0]), iterations=10) + RidgeRegressionWithSGD.train(rdd, initialWeights=array([1.0, 1.0]), iterations=10) + except ValueError: + self.fail() + + # Verify that maxBins is being passed through + GradientBoostedTrees.trainRegressor( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, numIterations=4, maxBins=32 + ) + with self.assertRaises(Exception): + GradientBoostedTrees.trainRegressor( + rdd, categoricalFeaturesInfo=categoricalFeaturesInfo, numIterations=4, maxBins=1 + ) + + +class ALSTests(MLlibTestCase): + def test_als_ratings_serialize(self): + ser = CPickleSerializer() + r = Rating(7, 1123, 3.14) + jr = self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.loads(bytearray(ser.dumps(r))) + nr = ser.loads(bytes(self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.dumps(jr))) + self.assertEqual(r.user, nr.user) + self.assertEqual(r.product, nr.product) + self.assertAlmostEqual(r.rating, nr.rating, 2) + + def test_als_ratings_id_long_error(self): + ser = CPickleSerializer() + r = Rating(1205640308657491975, 50233468418, 1.0) + # rating user id exceeds max int value, should fail when pickled + self.assertRaises( + Py4JJavaError, + self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.loads, + bytearray(ser.dumps(r)), + ) + + +class FPGrowthTest(MLlibTestCase): + def test_fpgrowth(self): + data = [["a", "b", "c"], ["a", "b", "d", "e"], ["a", "c", "e"], ["a", "c", "f"]] + rdd = self.sc.parallelize(data, 2) + model1 = FPGrowth.train(rdd, 0.6, 2) + # use default data partition number when numPartitions is not specified + model2 = FPGrowth.train(rdd, 0.6) + self.assertEqual( + sorted(model1.freqItemsets().collect()), sorted(model2.freqItemsets().collect()) + ) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_algorithms import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_feature.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_feature.py new file mode 100644 index 0000000..ca06f39 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_feature.py @@ -0,0 +1,192 @@ +# +# 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 math import sqrt +import unittest + +from numpy import array, abs, tile + +from pyspark.mllib.linalg import SparseVector, DenseVector, Vectors +from pyspark.mllib.linalg.distributed import RowMatrix +from pyspark.mllib.feature import HashingTF, IDF, StandardScaler, ElementwiseProduct, Word2Vec +from pyspark.testing.mllibutils import MLlibTestCase + + +class FeatureTest(MLlibTestCase): + def test_idf_model(self): + data = [ + Vectors.dense([1, 2, 6, 0, 2, 3, 1, 1, 0, 0, 3]), + Vectors.dense([1, 3, 0, 1, 3, 0, 0, 2, 0, 0, 1]), + Vectors.dense([1, 4, 1, 0, 0, 4, 9, 0, 1, 2, 0]), + Vectors.dense([2, 1, 0, 3, 0, 0, 5, 0, 2, 3, 9]), + ] + model = IDF().fit(self.sc.parallelize(data, 2)) + idf = model.idf() + self.assertEqual(len(idf), 11) + + +class Word2VecTests(MLlibTestCase): + def test_word2vec_setters(self): + model = ( + Word2Vec() + .setVectorSize(2) + .setLearningRate(0.01) + .setNumPartitions(2) + .setNumIterations(10) + .setSeed(1024) + .setMinCount(3) + .setWindowSize(6) + ) + self.assertEqual(model.vectorSize, 2) + self.assertTrue(model.learningRate < 0.02) + self.assertEqual(model.numPartitions, 2) + self.assertEqual(model.numIterations, 10) + self.assertEqual(model.seed, 1024) + self.assertEqual(model.minCount, 3) + self.assertEqual(model.windowSize, 6) + + def test_word2vec_get_vectors(self): + data = [ + ["a", "b", "c", "d", "e", "f", "g"], + ["a", "b", "c", "d", "e", "f"], + ["a", "b", "c", "d", "e"], + ["a", "b", "c", "d"], + ["a", "b", "c"], + ["a", "b"], + ["a"], + ] + model = Word2Vec().fit(self.sc.parallelize(data)) + self.assertEqual(len(model.getVectors()), 3) + + +class StandardScalerTests(MLlibTestCase): + def test_model_setters(self): + data = [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0]] + model = StandardScaler().fit(self.sc.parallelize(data)) + self.assertIsNotNone(model.setWithMean(True)) + self.assertIsNotNone(model.setWithStd(True)) + self.assertEqual(model.transform([1.0, 2.0, 3.0]), DenseVector([-1.0, -1.0, -1.0])) + + def test_model_transform(self): + data = [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0]] + model = StandardScaler().fit(self.sc.parallelize(data)) + self.assertEqual(model.transform([1.0, 2.0, 3.0]), DenseVector([1.0, 2.0, 3.0])) + + +class ElementwiseProductTests(MLlibTestCase): + def test_model_transform(self): + weight = Vectors.dense([3, 2, 1]) + + densevec = Vectors.dense([4, 5, 6]) + sparsevec = Vectors.sparse(3, [0], [1]) + eprod = ElementwiseProduct(weight) + self.assertEqual(eprod.transform(densevec), DenseVector([12, 10, 6])) + self.assertEqual(eprod.transform(sparsevec), SparseVector(3, [0], [3])) + + +class HashingTFTest(MLlibTestCase): + def test_binary_term_freqs(self): + hashingTF = HashingTF(100).setBinary(True) + doc = "a a b c c c".split(" ") + n = hashingTF.numFeatures + output = hashingTF.transform(doc).toArray() + expected = Vectors.sparse( + n, + {hashingTF.indexOf("a"): 1.0, hashingTF.indexOf("b"): 1.0, hashingTF.indexOf("c"): 1.0}, + ).toArray() + for i in range(0, n): + self.assertAlmostEqual( + output[i], + expected[i], + 14, + "Error at " + str(i) + ": expected " + str(expected[i]) + ", got " + str(output[i]), + ) + + +class DimensionalityReductionTests(MLlibTestCase): + + denseData = [ + Vectors.dense([0.0, 1.0, 2.0]), + Vectors.dense([3.0, 4.0, 5.0]), + Vectors.dense([6.0, 7.0, 8.0]), + Vectors.dense([9.0, 0.0, 1.0]), + ] + sparseData = [ + Vectors.sparse(3, [(1, 1.0), (2, 2.0)]), + Vectors.sparse(3, [(0, 3.0), (1, 4.0), (2, 5.0)]), + Vectors.sparse(3, [(0, 6.0), (1, 7.0), (2, 8.0)]), + Vectors.sparse(3, [(0, 9.0), (2, 1.0)]), + ] + + def assertEqualUpToSign(self, vecA, vecB): + eq1 = vecA - vecB + eq2 = vecA + vecB + self.assertTrue(sum(abs(eq1)) < 1e-6 or sum(abs(eq2)) < 1e-6) + + def test_svd(self): + denseMat = RowMatrix(self.sc.parallelize(self.denseData)) + sparseMat = RowMatrix(self.sc.parallelize(self.sparseData)) + m = 4 + n = 3 + for mat in [denseMat, sparseMat]: + for k in range(1, 4): + rm = mat.computeSVD(k, computeU=True) + self.assertEqual(rm.s.size, k) + self.assertEqual(rm.U.numRows(), m) + self.assertEqual(rm.U.numCols(), k) + self.assertEqual(rm.V.numRows, n) + self.assertEqual(rm.V.numCols, k) + + # Test that U returned is None if computeU is set to False. + self.assertEqual(mat.computeSVD(1).U, None) + + # Test that low rank matrices cannot have number of singular values + # greater than a limit. + rm = RowMatrix(self.sc.parallelize(tile([1, 2, 3], (3, 1)))) + self.assertEqual(rm.computeSVD(3, False, 1e-6).s.size, 1) + + def test_pca(self): + expected_pcs = array( + [ + [0.0, 1.0, 0.0], + [sqrt(2.0) / 2.0, 0.0, sqrt(2.0) / 2.0], + [sqrt(2.0) / 2.0, 0.0, -sqrt(2.0) / 2.0], + ] + ) + n = 3 + denseMat = RowMatrix(self.sc.parallelize(self.denseData)) + sparseMat = RowMatrix(self.sc.parallelize(self.sparseData)) + for mat in [denseMat, sparseMat]: + for k in range(1, 4): + pcs = mat.computePrincipalComponents(k) + self.assertEqual(pcs.numRows, n) + self.assertEqual(pcs.numCols, k) + + # We can just test the updated principal component for equality. + self.assertEqualUpToSign(pcs.toArray()[:, k - 1], expected_pcs[:, k - 1]) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_feature import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_linalg.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_linalg.py new file mode 100644 index 0000000..6d8dfcd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_linalg.py @@ -0,0 +1,680 @@ +# +# 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 array as pyarray +import unittest + +from numpy import array, array_equal, zeros, arange, tile, ones, inf + +import pyspark.ml.linalg as newlinalg +from pyspark.serializers import CPickleSerializer +from pyspark.mllib.linalg import ( + Vector, + SparseVector, + DenseVector, + VectorUDT, + _convert_to_vector, + DenseMatrix, + SparseMatrix, + Vectors, + Matrices, + MatrixUDT, +) +from pyspark.mllib.linalg.distributed import RowMatrix, IndexedRowMatrix, IndexedRow +from pyspark.mllib.regression import LabeledPoint +from pyspark.sql import Row +from pyspark.testing.mllibutils import MLlibTestCase +from pyspark.testing.utils import have_scipy + + +class VectorTests(MLlibTestCase): + def _test_serialize(self, v): + ser = CPickleSerializer() + self.assertEqual(v, ser.loads(ser.dumps(v))) + jvec = self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.loads(bytearray(ser.dumps(v))) + nv = ser.loads(bytes(self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.dumps(jvec))) + self.assertEqual(v, nv) + vs = [v] * 100 + jvecs = self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.loads(bytearray(ser.dumps(vs))) + nvs = ser.loads(bytes(self.sc._jvm.org.apache.spark.mllib.api.python.SerDe.dumps(jvecs))) + self.assertEqual(vs, nvs) + + def test_serialize(self): + self._test_serialize(DenseVector(range(10))) + self._test_serialize(DenseVector(array([1.0, 2.0, 3.0, 4.0]))) + self._test_serialize(DenseVector(pyarray.array("d", range(10)))) + self._test_serialize(SparseVector(4, {1: 1, 3: 2})) + self._test_serialize(SparseVector(3, {})) + self._test_serialize(DenseMatrix(2, 3, range(6))) + sm1 = SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0]) + self._test_serialize(sm1) + + def test_dot(self): + sv = SparseVector(4, {1: 1, 3: 2}) + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + lst = DenseVector([1, 2, 3, 4]) + mat = array( + [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]] + ) + arr = pyarray.array("d", [0, 1, 2, 3]) + self.assertEqual(10.0, sv.dot(dv)) + self.assertTrue(array_equal(array([3.0, 6.0, 9.0, 12.0]), sv.dot(mat))) + self.assertEqual(30.0, dv.dot(dv)) + self.assertTrue(array_equal(array([10.0, 20.0, 30.0, 40.0]), dv.dot(mat))) + self.assertEqual(30.0, lst.dot(dv)) + self.assertTrue(array_equal(array([10.0, 20.0, 30.0, 40.0]), lst.dot(mat))) + self.assertEqual(7.0, sv.dot(arr)) + + def test_squared_distance(self): + def squared_distance(a, b): + if isinstance(a, Vector): + return a.squared_distance(b) + else: + return b.squared_distance(a) + + sv = SparseVector(4, {1: 1, 3: 2}) + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + lst = DenseVector([4, 3, 2, 1]) + lst1 = [4, 3, 2, 1] + arr = pyarray.array("d", [0, 2, 1, 3]) + narr = array([0, 2, 1, 3]) + self.assertEqual(15.0, squared_distance(sv, dv)) + self.assertEqual(25.0, squared_distance(sv, lst)) + self.assertEqual(20.0, squared_distance(dv, lst)) + self.assertEqual(15.0, squared_distance(dv, sv)) + self.assertEqual(25.0, squared_distance(lst, sv)) + self.assertEqual(20.0, squared_distance(lst, dv)) + self.assertEqual(0.0, squared_distance(sv, sv)) + self.assertEqual(0.0, squared_distance(dv, dv)) + self.assertEqual(0.0, squared_distance(lst, lst)) + self.assertEqual(25.0, squared_distance(sv, lst1)) + self.assertEqual(3.0, squared_distance(sv, arr)) + self.assertEqual(3.0, squared_distance(sv, narr)) + + def test_hash(self): + v1 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v2 = SparseVector(4, [(1, 1.0), (3, 5.5)]) + v3 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v4 = SparseVector(4, [(1, 1.0), (3, 2.5)]) + self.assertEqual(hash(v1), hash(v2)) + self.assertEqual(hash(v1), hash(v3)) + self.assertEqual(hash(v2), hash(v3)) + self.assertFalse(hash(v1) == hash(v4)) + self.assertFalse(hash(v2) == hash(v4)) + + def test_eq(self): + v1 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v2 = SparseVector(4, [(1, 1.0), (3, 5.5)]) + v3 = DenseVector([0.0, 1.0, 0.0, 5.5]) + v4 = SparseVector(6, [(1, 1.0), (3, 5.5)]) + v5 = DenseVector([0.0, 1.0, 0.0, 2.5]) + v6 = SparseVector(4, [(1, 1.0), (3, 2.5)]) + dm1 = DenseMatrix(2, 2, [2, 0, 0, 0]) + sm1 = SparseMatrix(2, 2, [0, 2, 3], [0], [2]) + self.assertEqual(v1, v2) + self.assertEqual(v1, v3) + self.assertFalse(v2 == v4) + self.assertFalse(v1 == v5) + self.assertFalse(v1 == v6) + # this is done as Dense and Sparse matrices can be semantically + # equal while still implementing a different __eq__ method + self.assertEqual(dm1, sm1) + self.assertEqual(sm1, dm1) + + def test_equals(self): + indices = [1, 2, 4] + values = [1.0, 3.0, 2.0] + self.assertTrue(Vectors._equals(indices, values, list(range(5)), [0.0, 1.0, 3.0, 0.0, 2.0])) + self.assertFalse( + Vectors._equals(indices, values, list(range(5)), [0.0, 3.0, 1.0, 0.0, 2.0]) + ) + self.assertFalse(Vectors._equals(indices, values, list(range(5)), [0.0, 3.0, 0.0, 2.0])) + self.assertFalse( + Vectors._equals(indices, values, list(range(5)), [0.0, 1.0, 3.0, 2.0, 2.0]) + ) + + def test_conversion(self): + # numpy arrays should be automatically upcast to float64 + # tests for fix of [SPARK-5089] + v = array([1, 2, 3, 4], dtype="float64") + dv = DenseVector(v) + self.assertTrue(dv.array.dtype == "float64") + v = array([1, 2, 3, 4], dtype="float32") + dv = DenseVector(v) + self.assertTrue(dv.array.dtype == "float64") + + def test_sparse_vector_indexing(self): + sv = SparseVector(5, {1: 1, 3: 2}) + self.assertEqual(sv[0], 0.0) + self.assertEqual(sv[3], 2.0) + self.assertEqual(sv[1], 1.0) + self.assertEqual(sv[2], 0.0) + self.assertEqual(sv[4], 0.0) + self.assertEqual(sv[-1], 0.0) + self.assertEqual(sv[-2], 2.0) + self.assertEqual(sv[-3], 0.0) + self.assertEqual(sv[-5], 0.0) + for ind in [5, -6]: + self.assertRaises(IndexError, sv.__getitem__, ind) + for ind in [7.8, "1"]: + self.assertRaises(TypeError, sv.__getitem__, ind) + + zeros = SparseVector(4, {}) + self.assertEqual(zeros[0], 0.0) + self.assertEqual(zeros[3], 0.0) + for ind in [4, -5]: + self.assertRaises(IndexError, zeros.__getitem__, ind) + + empty = SparseVector(0, {}) + for ind in [-1, 0, 1]: + self.assertRaises(IndexError, empty.__getitem__, ind) + + def test_sparse_vector_iteration(self): + self.assertListEqual(list(SparseVector(3, [], [])), [0.0, 0.0, 0.0]) + self.assertListEqual(list(SparseVector(5, [0, 3], [1.0, 2.0])), [1.0, 0.0, 0.0, 2.0, 0.0]) + + def test_matrix_indexing(self): + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10]) + expected = [[0, 6], [1, 8], [4, 10]] + for i in range(3): + for j in range(2): + self.assertEqual(mat[i, j], expected[i][j]) + + for i, j in [(-1, 0), (4, 1), (3, 4)]: + self.assertRaises(IndexError, mat.__getitem__, (i, j)) + + def test_repr_dense_matrix(self): + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10]) + self.assertTrue(repr(mat), "DenseMatrix(3, 2, [0.0, 1.0, 4.0, 6.0, 8.0, 10.0], False)") + + mat = DenseMatrix(3, 2, [0, 1, 4, 6, 8, 10], True) + self.assertTrue(repr(mat), "DenseMatrix(3, 2, [0.0, 1.0, 4.0, 6.0, 8.0, 10.0], False)") + + mat = DenseMatrix(6, 3, zeros(18)) + self.assertTrue( + repr(mat), + "DenseMatrix(6, 3, [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], False)", + ) + + def test_repr_sparse_matrix(self): + sm1t = SparseMatrix( + 3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], isTransposed=True + ) + self.assertTrue( + repr(sm1t), + "SparseMatrix(3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], True)", + ) + + indices = tile(arange(6), 3) + values = ones(18) + sm = SparseMatrix(6, 3, [0, 6, 12, 18], indices, values) + self.assertTrue( + repr(sm), + "SparseMatrix(6, 3, [0, 6, 12, 18], \ + [0, 1, 2, 3, 4, 5, 0, 1, ..., 4, 5, 0, 1, 2, 3, 4, 5], \ + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, ..., \ + 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], False)", + ) + + self.assertTrue( + str(sm), + "6 X 3 CSCMatrix\n\ + (0,0) 1.0\n(1,0) 1.0\n(2,0) 1.0\n(3,0) 1.0\n(4,0) 1.0\n(5,0) 1.0\n\ + (0,1) 1.0\n(1,1) 1.0\n(2,1) 1.0\n(3,1) 1.0\n(4,1) 1.0\n(5,1) 1.0\n\ + (0,2) 1.0\n(1,2) 1.0\n(2,2) 1.0\n(3,2) 1.0\n..\n..", + ) + + sm = SparseMatrix(1, 18, zeros(19), [], []) + self.assertTrue( + repr(sm), + "SparseMatrix(1, 18, \ + [0, 0, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0, 0, 0, 0, 0], [], [], False)", + ) + + def test_sparse_matrix(self): + # Test sparse matrix creation. + sm1 = SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0]) + self.assertEqual(sm1.numRows, 3) + self.assertEqual(sm1.numCols, 4) + self.assertEqual(sm1.colPtrs.tolist(), [0, 2, 2, 4, 4]) + self.assertEqual(sm1.rowIndices.tolist(), [1, 2, 1, 2]) + self.assertEqual(sm1.values.tolist(), [1.0, 2.0, 4.0, 5.0]) + self.assertTrue( + repr(sm1), + "SparseMatrix(3, 4, [0, 2, 2, 4, 4], [1, 2, 1, 2], [1.0, 2.0, 4.0, 5.0], False)", + ) + + # Test indexing + expected = [[0, 0, 0, 0], [1, 0, 4, 0], [2, 0, 5, 0]] + + for i in range(3): + for j in range(4): + self.assertEqual(expected[i][j], sm1[i, j]) + self.assertTrue(array_equal(sm1.toArray(), expected)) + + for i, j in [(-1, 1), (4, 3), (3, 5)]: + self.assertRaises(IndexError, sm1.__getitem__, (i, j)) + + # Test conversion to dense and sparse. + smnew = sm1.toDense().toSparse() + self.assertEqual(sm1.numRows, smnew.numRows) + self.assertEqual(sm1.numCols, smnew.numCols) + self.assertTrue(array_equal(sm1.colPtrs, smnew.colPtrs)) + self.assertTrue(array_equal(sm1.rowIndices, smnew.rowIndices)) + self.assertTrue(array_equal(sm1.values, smnew.values)) + + sm1t = SparseMatrix( + 3, 4, [0, 2, 3, 5], [0, 1, 2, 0, 2], [3.0, 2.0, 4.0, 9.0, 8.0], isTransposed=True + ) + self.assertEqual(sm1t.numRows, 3) + self.assertEqual(sm1t.numCols, 4) + self.assertEqual(sm1t.colPtrs.tolist(), [0, 2, 3, 5]) + self.assertEqual(sm1t.rowIndices.tolist(), [0, 1, 2, 0, 2]) + self.assertEqual(sm1t.values.tolist(), [3.0, 2.0, 4.0, 9.0, 8.0]) + + expected = [[3, 2, 0, 0], [0, 0, 4, 0], [9, 0, 8, 0]] + + for i in range(3): + for j in range(4): + self.assertEqual(expected[i][j], sm1t[i, j]) + self.assertTrue(array_equal(sm1t.toArray(), expected)) + + def test_dense_matrix_is_transposed(self): + mat1 = DenseMatrix(3, 2, [0, 4, 1, 6, 3, 9], isTransposed=True) + mat = DenseMatrix(3, 2, [0, 1, 3, 4, 6, 9]) + self.assertEqual(mat1, mat) + + expected = [[0, 4], [1, 6], [3, 9]] + for i in range(3): + for j in range(2): + self.assertEqual(mat1[i, j], expected[i][j]) + self.assertTrue(array_equal(mat1.toArray(), expected)) + + sm = mat1.toSparse() + self.assertTrue(array_equal(sm.rowIndices, [1, 2, 0, 1, 2])) + self.assertTrue(array_equal(sm.colPtrs, [0, 2, 5])) + self.assertTrue(array_equal(sm.values, [1, 3, 4, 6, 9])) + + def test_parse_vector(self): + a = DenseVector([]) + self.assertEqual(str(a), "[]") + self.assertEqual(Vectors.parse(str(a)), a) + a = DenseVector([3, 4, 6, 7]) + self.assertEqual(str(a), "[3.0,4.0,6.0,7.0]") + self.assertEqual(Vectors.parse(str(a)), a) + a = SparseVector(4, [], []) + self.assertEqual(str(a), "(4,[],[])") + self.assertEqual(SparseVector.parse(str(a)), a) + a = SparseVector(4, [0, 2], [3, 4]) + self.assertEqual(str(a), "(4,[0,2],[3.0,4.0])") + self.assertEqual(Vectors.parse(str(a)), a) + a = SparseVector(10, [0, 1], [4, 5]) + self.assertEqual(SparseVector.parse(" (10, [0,1 ],[ 4.0,5.0] )"), a) + + def test_norms(self): + a = DenseVector([0, 2, 3, -1]) + self.assertAlmostEqual(a.norm(2), 3.742, 3) + self.assertTrue(a.norm(1), 6) + self.assertTrue(a.norm(inf), 3) + a = SparseVector(4, [0, 2], [3, -4]) + self.assertAlmostEqual(a.norm(2), 5) + self.assertTrue(a.norm(1), 7) + self.assertTrue(a.norm(inf), 4) + + tmp = SparseVector(4, [0, 2], [3, 0]) + self.assertEqual(tmp.numNonzeros(), 1) + + def test_ml_mllib_vector_conversion(self): + # to ml + # dense + mllibDV = Vectors.dense([1, 2, 3]) + mlDV1 = newlinalg.Vectors.dense([1, 2, 3]) + mlDV2 = mllibDV.asML() + self.assertEqual(mlDV2, mlDV1) + # sparse + mllibSV = Vectors.sparse(4, {1: 1.0, 3: 5.5}) + mlSV1 = newlinalg.Vectors.sparse(4, {1: 1.0, 3: 5.5}) + mlSV2 = mllibSV.asML() + self.assertEqual(mlSV2, mlSV1) + # from ml + # dense + mllibDV1 = Vectors.dense([1, 2, 3]) + mlDV = newlinalg.Vectors.dense([1, 2, 3]) + mllibDV2 = Vectors.fromML(mlDV) + self.assertEqual(mllibDV1, mllibDV2) + # sparse + mllibSV1 = Vectors.sparse(4, {1: 1.0, 3: 5.5}) + mlSV = newlinalg.Vectors.sparse(4, {1: 1.0, 3: 5.5}) + mllibSV2 = Vectors.fromML(mlSV) + self.assertEqual(mllibSV1, mllibSV2) + + def test_ml_mllib_matrix_conversion(self): + # to ml + # dense + mllibDM = Matrices.dense(2, 2, [0, 1, 2, 3]) + mlDM1 = newlinalg.Matrices.dense(2, 2, [0, 1, 2, 3]) + mlDM2 = mllibDM.asML() + self.assertEqual(mlDM2, mlDM1) + # transposed + mllibDMt = DenseMatrix(2, 2, [0, 1, 2, 3], True) + mlDMt1 = newlinalg.DenseMatrix(2, 2, [0, 1, 2, 3], True) + mlDMt2 = mllibDMt.asML() + self.assertEqual(mlDMt2, mlDMt1) + # sparse + mllibSM = Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + mlSM1 = newlinalg.Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + mlSM2 = mllibSM.asML() + self.assertEqual(mlSM2, mlSM1) + # transposed + mllibSMt = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + mlSMt1 = newlinalg.SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + mlSMt2 = mllibSMt.asML() + self.assertEqual(mlSMt2, mlSMt1) + # from ml + # dense + mllibDM1 = Matrices.dense(2, 2, [1, 2, 3, 4]) + mlDM = newlinalg.Matrices.dense(2, 2, [1, 2, 3, 4]) + mllibDM2 = Matrices.fromML(mlDM) + self.assertEqual(mllibDM1, mllibDM2) + # transposed + mllibDMt1 = DenseMatrix(2, 2, [1, 2, 3, 4], True) + mlDMt = newlinalg.DenseMatrix(2, 2, [1, 2, 3, 4], True) + mllibDMt2 = Matrices.fromML(mlDMt) + self.assertEqual(mllibDMt1, mllibDMt2) + # sparse + mllibSM1 = Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + mlSM = newlinalg.Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]) + mllibSM2 = Matrices.fromML(mlSM) + self.assertEqual(mllibSM1, mllibSM2) + # transposed + mllibSMt1 = SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + mlSMt = newlinalg.SparseMatrix(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4], True) + mllibSMt2 = Matrices.fromML(mlSMt) + self.assertEqual(mllibSMt1, mllibSMt2) + + +class VectorUDTTests(MLlibTestCase): + + dv0 = DenseVector([]) + dv1 = DenseVector([1.0, 2.0]) + sv0 = SparseVector(2, [], []) + sv1 = SparseVector(2, [1], [2.0]) + udt = VectorUDT() + + def test_json_schema(self): + self.assertEqual(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + + def test_serialization(self): + for v in [self.dv0, self.dv1, self.sv0, self.sv1]: + self.assertEqual(v, self.udt.deserialize(self.udt.serialize(v))) + + def test_infer_schema(self): + rdd = self.sc.parallelize([LabeledPoint(1.0, self.dv1), LabeledPoint(0.0, self.sv1)]) + df = rdd.toDF() + schema = df.schema + field = [f for f in schema.fields if f.name == "features"][0] + self.assertEqual(field.dataType, self.udt) + vectors = df.rdd.map(lambda p: p.features).collect() + self.assertEqual(len(vectors), 2) + for v in vectors: + if isinstance(v, SparseVector): + self.assertEqual(v, self.sv1) + elif isinstance(v, DenseVector): + self.assertEqual(v, self.dv1) + else: + raise TypeError("expecting a vector but got %r of type %r" % (v, type(v))) + + def test_row_matrix_from_dataframe(self): + from pyspark.errors import IllegalArgumentException + + df = self.spark.createDataFrame([Row(Vectors.dense(1))]) + row_matrix = RowMatrix(df) + self.assertEqual(row_matrix.numRows(), 1) + self.assertEqual(row_matrix.numCols(), 1) + with self.assertRaises(IllegalArgumentException): + RowMatrix(df.selectExpr("'monkey'")) + + def test_indexed_row_matrix_from_dataframe(self): + from pyspark.errors import IllegalArgumentException + + df = self.spark.createDataFrame([Row(int(0), Vectors.dense(1))]) + matrix = IndexedRowMatrix(df) + self.assertEqual(matrix.numRows(), 1) + self.assertEqual(matrix.numCols(), 1) + with self.assertRaises(IllegalArgumentException): + IndexedRowMatrix(df.drop("_1")) + + def test_row_matrix_invalid_type(self): + rows = self.sc.parallelize([[1, 2, 3], [4, 5, 6]]) + invalid_type = "" + matrix = RowMatrix(rows) + self.assertRaises(TypeError, matrix.multiply, invalid_type) + + irows = self.sc.parallelize([IndexedRow(0, [1, 2, 3]), IndexedRow(1, [4, 5, 6])]) + imatrix = IndexedRowMatrix(irows) + self.assertRaises(TypeError, imatrix.multiply, invalid_type) + + +class MatrixUDTTests(MLlibTestCase): + + dm1 = DenseMatrix(3, 2, [0, 1, 4, 5, 9, 10]) + dm2 = DenseMatrix(3, 2, [0, 1, 4, 5, 9, 10], isTransposed=True) + sm1 = SparseMatrix(1, 1, [0, 1], [0], [2.0]) + sm2 = SparseMatrix(2, 1, [0, 0, 1], [0], [5.0], isTransposed=True) + udt = MatrixUDT() + + def test_json_schema(self): + self.assertEqual(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + + def test_serialization(self): + for m in [self.dm1, self.dm2, self.sm1, self.sm2]: + self.assertEqual(m, self.udt.deserialize(self.udt.serialize(m))) + + def test_infer_schema(self): + rdd = self.sc.parallelize([("dense", self.dm1), ("sparse", self.sm1)]) + df = rdd.toDF() + schema = df.schema + self.assertTrue(schema.fields[1].dataType, self.udt) + matrices = df.rdd.map(lambda x: x._2).collect() + self.assertEqual(len(matrices), 2) + for m in matrices: + if isinstance(m, DenseMatrix): + self.assertTrue(m, self.dm1) + elif isinstance(m, SparseMatrix): + self.assertTrue(m, self.sm1) + else: + raise ValueError("Expected a matrix but got type %r" % type(m)) + + +@unittest.skipIf(not have_scipy, "SciPy not installed") +class SciPyTests(MLlibTestCase): + + """ + Test both vector operations and MLlib algorithms with SciPy sparse matrices, + if SciPy is available. + """ + + def test_serialize(self): + from scipy.sparse import lil_matrix + + ser = CPickleSerializer() + lil = lil_matrix((4, 1)) + lil[1, 0] = 1 + lil[3, 0] = 2 + sv = SparseVector(4, {1: 1, 3: 2}) + self.assertEqual(sv, _convert_to_vector(lil)) + self.assertEqual(sv, _convert_to_vector(lil.tocsc())) + self.assertEqual(sv, _convert_to_vector(lil.tocoo())) + self.assertEqual(sv, _convert_to_vector(lil.tocsr())) + self.assertEqual(sv, _convert_to_vector(lil.todok())) + + def serialize(d): + return ser.loads(ser.dumps(_convert_to_vector(d))) + + self.assertEqual(sv, serialize(lil)) + self.assertEqual(sv, serialize(lil.tocsc())) + self.assertEqual(sv, serialize(lil.tocsr())) + self.assertEqual(sv, serialize(lil.todok())) + + def test_convert_to_vector(self): + from scipy.sparse import csc_matrix + + # Create a CSC matrix with non-sorted indices + indptr = array([0, 2]) + indices = array([3, 1]) + data = array([2.0, 1.0]) + csc = csc_matrix((data, indices, indptr)) + self.assertFalse(csc.has_sorted_indices) + sv = SparseVector(4, {1: 1, 3: 2}) + self.assertEqual(sv, _convert_to_vector(csc)) + + def test_dot(self): + from scipy.sparse import lil_matrix + + lil = lil_matrix((4, 1)) + lil[1, 0] = 1 + lil[3, 0] = 2 + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + self.assertEqual(10.0, dv.dot(lil)) + + def test_squared_distance(self): + from scipy.sparse import lil_matrix + + lil = lil_matrix((4, 1)) + lil[1, 0] = 3 + lil[3, 0] = 2 + dv = DenseVector(array([1.0, 2.0, 3.0, 4.0])) + sv = SparseVector(4, {0: 1, 1: 2, 2: 3, 3: 4}) + self.assertEqual(15.0, dv.squared_distance(lil)) + self.assertEqual(15.0, sv.squared_distance(lil)) + + def scipy_matrix(self, size, values): + """Create a column SciPy matrix from a dictionary of values""" + from scipy.sparse import lil_matrix + + lil = lil_matrix((size, 1)) + for key, value in values.items(): + lil[key, 0] = value + return lil + + def test_clustering(self): + from pyspark.mllib.clustering import KMeans + + data = [ + self.scipy_matrix(3, {1: 1.0}), + self.scipy_matrix(3, {1: 1.1}), + self.scipy_matrix(3, {2: 1.0}), + self.scipy_matrix(3, {2: 1.1}), + ] + clusters = KMeans.train(self.sc.parallelize(data), 2, initializationMode="k-means||") + self.assertEqual(clusters.predict(data[0]), clusters.predict(data[1])) + self.assertEqual(clusters.predict(data[2]), clusters.predict(data[3])) + + def test_classification(self): + from pyspark.mllib.classification import LogisticRegressionWithSGD, SVMWithSGD, NaiveBayes + from pyspark.mllib.tree import DecisionTree + + data = [ + LabeledPoint(0.0, self.scipy_matrix(2, {0: 1.0})), + LabeledPoint(1.0, self.scipy_matrix(2, {1: 1.0})), + LabeledPoint(0.0, self.scipy_matrix(2, {0: 2.0})), + LabeledPoint(1.0, self.scipy_matrix(2, {1: 2.0})), + ] + rdd = self.sc.parallelize(data) + features = [p.features for p in data] + + lr_model = LogisticRegressionWithSGD.train(rdd) + self.assertTrue(lr_model.predict(features[0]) <= 0) + self.assertTrue(lr_model.predict(features[1]) > 0) + self.assertTrue(lr_model.predict(features[2]) <= 0) + self.assertTrue(lr_model.predict(features[3]) > 0) + + svm_model = SVMWithSGD.train(rdd) + self.assertTrue(svm_model.predict(features[0]) <= 0) + self.assertTrue(svm_model.predict(features[1]) > 0) + self.assertTrue(svm_model.predict(features[2]) <= 0) + self.assertTrue(svm_model.predict(features[3]) > 0) + + nb_model = NaiveBayes.train(rdd) + self.assertTrue(nb_model.predict(features[0]) <= 0) + self.assertTrue(nb_model.predict(features[1]) > 0) + self.assertTrue(nb_model.predict(features[2]) <= 0) + self.assertTrue(nb_model.predict(features[3]) > 0) + + categoricalFeaturesInfo = {0: 3} # feature 0 has 3 categories + dt_model = DecisionTree.trainClassifier( + rdd, numClasses=2, categoricalFeaturesInfo=categoricalFeaturesInfo + ) + self.assertTrue(dt_model.predict(features[0]) <= 0) + self.assertTrue(dt_model.predict(features[1]) > 0) + self.assertTrue(dt_model.predict(features[2]) <= 0) + self.assertTrue(dt_model.predict(features[3]) > 0) + + def test_regression(self): + from pyspark.mllib.regression import ( + LinearRegressionWithSGD, + LassoWithSGD, + RidgeRegressionWithSGD, + ) + from pyspark.mllib.tree import DecisionTree + + data = [ + LabeledPoint(-1.0, self.scipy_matrix(2, {1: -1.0})), + LabeledPoint(1.0, self.scipy_matrix(2, {1: 1.0})), + LabeledPoint(-1.0, self.scipy_matrix(2, {1: -2.0})), + LabeledPoint(1.0, self.scipy_matrix(2, {1: 2.0})), + ] + rdd = self.sc.parallelize(data) + features = [p.features for p in data] + + lr_model = LinearRegressionWithSGD.train(rdd) + self.assertTrue(lr_model.predict(features[0]) <= 0) + self.assertTrue(lr_model.predict(features[1]) > 0) + self.assertTrue(lr_model.predict(features[2]) <= 0) + self.assertTrue(lr_model.predict(features[3]) > 0) + + lasso_model = LassoWithSGD.train(rdd) + self.assertTrue(lasso_model.predict(features[0]) <= 0) + self.assertTrue(lasso_model.predict(features[1]) > 0) + self.assertTrue(lasso_model.predict(features[2]) <= 0) + self.assertTrue(lasso_model.predict(features[3]) > 0) + + rr_model = RidgeRegressionWithSGD.train(rdd) + self.assertTrue(rr_model.predict(features[0]) <= 0) + self.assertTrue(rr_model.predict(features[1]) > 0) + self.assertTrue(rr_model.predict(features[2]) <= 0) + self.assertTrue(rr_model.predict(features[3]) > 0) + + categoricalFeaturesInfo = {0: 2} # feature 0 has 2 categories + dt_model = DecisionTree.trainRegressor(rdd, categoricalFeaturesInfo=categoricalFeaturesInfo) + self.assertTrue(dt_model.predict(features[0]) <= 0) + self.assertTrue(dt_model.predict(features[1]) > 0) + self.assertTrue(dt_model.predict(features[2]) <= 0) + self.assertTrue(dt_model.predict(features[3]) > 0) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_linalg import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_stat.py new file mode 100644 index 0000000..4fcd49f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_stat.py @@ -0,0 +1,206 @@ +# +# 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 array as pyarray +import unittest + +from numpy import array + +from pyspark.mllib.linalg import Vectors, Matrices +from pyspark.mllib.random import RandomRDDs +from pyspark.mllib.regression import LabeledPoint +from pyspark.mllib.stat import Statistics +from pyspark.errors import IllegalArgumentException +from pyspark.testing.mllibutils import MLlibTestCase + + +class StatTests(MLlibTestCase): + # SPARK-4023 + def test_col_with_different_rdds(self): + # numpy + data = RandomRDDs.normalVectorRDD(self.sc, 1000, 10, 10) + summary = Statistics.colStats(data) + self.assertEqual(1000, summary.count()) + # array + data = self.sc.parallelize([range(10)] * 10) + summary = Statistics.colStats(data) + self.assertEqual(10, summary.count()) + # array + data = self.sc.parallelize([pyarray.array("d", range(10))] * 10) + summary = Statistics.colStats(data) + self.assertEqual(10, summary.count()) + + def test_col_norms(self): + data = RandomRDDs.normalVectorRDD(self.sc, 1000, 10, 10) + summary = Statistics.colStats(data) + self.assertEqual(10, len(summary.normL1())) + self.assertEqual(10, len(summary.normL2())) + + data2 = self.sc.parallelize(range(10)).map(lambda x: Vectors.dense(x)) + summary2 = Statistics.colStats(data2) + self.assertEqual(array([45.0]), summary2.normL1()) + import math + + expectedNormL2 = math.sqrt(sum(map(lambda x: x * x, range(10)))) + self.assertTrue(math.fabs(summary2.normL2()[0] - expectedNormL2) < 1e-14) + + +class ChiSqTestTests(MLlibTestCase): + def test_goodness_of_fit(self): + from numpy import inf + + observed = Vectors.dense([4, 6, 5]) + pearson = Statistics.chiSqTest(observed) + + # Validated against the R command `chisq.test(c(4, 6, 5), p=c(1/3, 1/3, 1/3))` + self.assertEqual(pearson.statistic, 0.4) + self.assertEqual(pearson.degreesOfFreedom, 2) + self.assertAlmostEqual(pearson.pValue, 0.8187, 4) + + # Different expected and observed sum + observed1 = Vectors.dense([21, 38, 43, 80]) + expected1 = Vectors.dense([3, 5, 7, 20]) + pearson1 = Statistics.chiSqTest(observed1, expected1) + + # Results validated against the R command + # `chisq.test(c(21, 38, 43, 80), p=c(3/35, 1/7, 1/5, 4/7))` + self.assertAlmostEqual(pearson1.statistic, 14.1429, 4) + self.assertEqual(pearson1.degreesOfFreedom, 3) + self.assertAlmostEqual(pearson1.pValue, 0.002717, 4) + + # Vectors with different sizes + observed3 = Vectors.dense([1.0, 2.0, 3.0]) + expected3 = Vectors.dense([1.0, 2.0, 3.0, 4.0]) + self.assertRaises(ValueError, Statistics.chiSqTest, observed3, expected3) + + # Negative counts in observed + neg_obs = Vectors.dense([1.0, 2.0, 3.0, -4.0]) + self.assertRaises(IllegalArgumentException, Statistics.chiSqTest, neg_obs, expected1) + + # Count = 0.0 in expected but not observed + zero_expected = Vectors.dense([1.0, 0.0, 3.0]) + pearson_inf = Statistics.chiSqTest(observed, zero_expected) + self.assertEqual(pearson_inf.statistic, inf) + self.assertEqual(pearson_inf.degreesOfFreedom, 2) + self.assertEqual(pearson_inf.pValue, 0.0) + + # 0.0 in expected and observed simultaneously + zero_observed = Vectors.dense([2.0, 0.0, 1.0]) + self.assertRaises( + IllegalArgumentException, Statistics.chiSqTest, zero_observed, zero_expected + ) + + def test_matrix_independence(self): + data = [40.0, 24.0, 29.0, 56.0, 32.0, 42.0, 31.0, 10.0, 0.0, 30.0, 15.0, 12.0] + chi = Statistics.chiSqTest(Matrices.dense(3, 4, data)) + + # Results validated against R command + # `chisq.test(rbind(c(40, 56, 31, 30),c(24, 32, 10, 15), c(29, 42, 0, 12)))` + self.assertAlmostEqual(chi.statistic, 21.9958, 4) + self.assertEqual(chi.degreesOfFreedom, 6) + self.assertAlmostEqual(chi.pValue, 0.001213, 4) + + # Negative counts + neg_counts = Matrices.dense(2, 2, [4.0, 5.0, 3.0, -3.0]) + self.assertRaises(IllegalArgumentException, Statistics.chiSqTest, neg_counts) + + # Row sum = 0.0 + row_zero = Matrices.dense(2, 2, [0.0, 1.0, 0.0, 2.0]) + self.assertRaises(IllegalArgumentException, Statistics.chiSqTest, row_zero) + + # Column sum = 0.0 + col_zero = Matrices.dense(2, 2, [0.0, 0.0, 2.0, 2.0]) + self.assertRaises(IllegalArgumentException, Statistics.chiSqTest, col_zero) + + def test_chi_sq_pearson(self): + data = [ + LabeledPoint(0.0, Vectors.dense([0.5, 10.0])), + LabeledPoint(0.0, Vectors.dense([1.5, 20.0])), + LabeledPoint(1.0, Vectors.dense([1.5, 30.0])), + LabeledPoint(0.0, Vectors.dense([3.5, 30.0])), + LabeledPoint(0.0, Vectors.dense([3.5, 40.0])), + LabeledPoint(1.0, Vectors.dense([3.5, 40.0])), + ] + + for numParts in [2, 4, 6, 8]: + chi = Statistics.chiSqTest(self.sc.parallelize(data, numParts)) + feature1 = chi[0] + self.assertEqual(feature1.statistic, 0.75) + self.assertEqual(feature1.degreesOfFreedom, 2) + self.assertAlmostEqual(feature1.pValue, 0.6873, 4) + + feature2 = chi[1] + self.assertEqual(feature2.statistic, 1.5) + self.assertEqual(feature2.degreesOfFreedom, 3) + self.assertAlmostEqual(feature2.pValue, 0.6823, 4) + + def test_right_number_of_results(self): + num_cols = 1001 + sparse_data = [ + LabeledPoint(0.0, Vectors.sparse(num_cols, [(100, 2.0)])), + LabeledPoint(0.1, Vectors.sparse(num_cols, [(200, 1.0)])), + ] + chi = Statistics.chiSqTest(self.sc.parallelize(sparse_data)) + self.assertEqual(len(chi), num_cols) + self.assertIsNotNone(chi[1000]) + + +class KolmogorovSmirnovTest(MLlibTestCase): + def test_R_implementation_equivalence(self): + data = self.sc.parallelize( + [ + 1.1626852897838, + -0.585924465893051, + 1.78546500331661, + -1.33259371048501, + -0.446566766553219, + 0.569606122374976, + -2.88971761441412, + -0.869018343326555, + -0.461702683149641, + -0.555540910137444, + -0.0201353678515895, + -0.150382224136063, + -0.628126755843964, + 1.32322085193283, + -1.52135057001199, + -0.437427868856691, + 0.970577579543399, + 0.0282226444247749, + -0.0857821886527593, + 0.389214404984942, + ] + ) + model = Statistics.kolmogorovSmirnovTest(data, "norm") + self.assertAlmostEqual(model.statistic, 0.189, 3) + self.assertAlmostEqual(model.pValue, 0.422, 3) + + model = Statistics.kolmogorovSmirnovTest(data, "norm", 0, 1) + self.assertAlmostEqual(model.statistic, 0.189, 3) + self.assertAlmostEqual(model.pValue, 0.422, 3) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_stat import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_streaming_algorithms.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_streaming_algorithms.py new file mode 100644 index 0000000..5a06742 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_streaming_algorithms.py @@ -0,0 +1,471 @@ +# +# 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 os +import unittest + +from numpy import array, random, exp, dot, all, mean, abs +from numpy import sum as array_sum + +from pyspark import SparkContext +from pyspark.mllib.clustering import StreamingKMeans, StreamingKMeansModel +from pyspark.mllib.classification import StreamingLogisticRegressionWithSGD +from pyspark.mllib.linalg import Vectors +from pyspark.mllib.regression import LabeledPoint, StreamingLinearRegressionWithSGD +from pyspark.mllib.util import LinearDataGenerator +from pyspark.streaming import StreamingContext +from pyspark.testing.utils import eventually + + +class MLLibStreamingTestCase(unittest.TestCase): + def setUp(self): + self.sc = SparkContext("local[4]", "MLlib tests") + self.ssc = StreamingContext(self.sc, 1.0) + + def tearDown(self): + self.ssc.stop(False) + self.sc.stop() + + +class StreamingKMeansTest(MLLibStreamingTestCase): + def test_model_params(self): + """Test that the model params are set correctly""" + stkm = StreamingKMeans() + stkm.setK(5).setDecayFactor(0.0) + self.assertEqual(stkm._k, 5) + self.assertEqual(stkm._decayFactor, 0.0) + + # Model not set yet. + self.assertIsNone(stkm.latestModel()) + self.assertRaises(ValueError, stkm.trainOn, [0.0, 1.0]) + + stkm.setInitialCenters(centers=[[0.0, 0.0], [1.0, 1.0]], weights=[1.0, 1.0]) + self.assertEqual(stkm.latestModel().centers, [[0.0, 0.0], [1.0, 1.0]]) + self.assertEqual(stkm.latestModel().clusterWeights, [1.0, 1.0]) + + def test_accuracy_for_single_center(self): + """Test that parameters obtained are correct for a single center.""" + centers, batches = self.streamingKMeansDataGenerator( + batches=5, numPoints=5, k=1, d=5, r=0.1, seed=0 + ) + stkm = StreamingKMeans(1) + stkm.setInitialCenters([[0.0, 0.0, 0.0, 0.0, 0.0]], [0.0]) + input_stream = self.ssc.queueStream([self.sc.parallelize(batch, 1) for batch in batches]) + stkm.trainOn(input_stream) + + self.ssc.start() + + def condition(): + self.assertEqual(stkm.latestModel().clusterWeights, [25.0]) + return True + + eventually(condition, catch_assertions=True) + + realCenters = array_sum(array(centers), axis=0) + for i in range(5): + modelCenters = stkm.latestModel().centers[0][i] + self.assertAlmostEqual(centers[0][i], modelCenters, 1) + self.assertAlmostEqual(realCenters[i], modelCenters, 1) + + def streamingKMeansDataGenerator(self, batches, numPoints, k, d, r, seed, centers=None): + rng = random.RandomState(seed) + + # Generate centers. + centers = [rng.randn(d) for i in range(k)] + + return centers, [ + [Vectors.dense(centers[j % k] + r * rng.randn(d)) for j in range(numPoints)] + for i in range(batches) + ] + + def test_trainOn_model(self): + """Test the model on toy data with four clusters.""" + stkm = StreamingKMeans() + initCenters = [[1.0, 1.0], [-1.0, 1.0], [-1.0, -1.0], [1.0, -1.0]] + stkm.setInitialCenters(centers=initCenters, weights=[1.0, 1.0, 1.0, 1.0]) + + # Create a toy dataset by setting a tiny offset for each point. + offsets = [[0, 0.1], [0, -0.1], [0.1, 0], [-0.1, 0]] + batches = [] + for offset in offsets: + batches.append( + [[offset[0] + center[0], offset[1] + center[1]] for center in initCenters] + ) + + batches = [self.sc.parallelize(batch, 1) for batch in batches] + input_stream = self.ssc.queueStream(batches) + stkm.trainOn(input_stream) + self.ssc.start() + + # Give enough time to train the model. + def condition(): + finalModel = stkm.latestModel() + self.assertTrue(all(finalModel.centers == array(initCenters))) + self.assertEqual(finalModel.clusterWeights, [5.0, 5.0, 5.0, 5.0]) + return True + + eventually(condition, 90, catch_assertions=True) + + def test_predictOn_model(self): + """Test that the model predicts correctly on toy data.""" + stkm = StreamingKMeans() + stkm._model = StreamingKMeansModel( + clusterCenters=[[1.0, 1.0], [-1.0, 1.0], [-1.0, -1.0], [1.0, -1.0]], + clusterWeights=[1.0, 1.0, 1.0, 1.0], + ) + + predict_data = [[[1.5, 1.5]], [[-1.5, 1.5]], [[-1.5, -1.5]], [[1.5, -1.5]]] + predict_data = [self.sc.parallelize(batch, 1) for batch in predict_data] + predict_stream = self.ssc.queueStream(predict_data) + predict_val = stkm.predictOn(predict_stream) + + result = [] + + def update(rdd): + rdd_collect = rdd.collect() + if rdd_collect: + result.append(rdd_collect) + + predict_val.foreachRDD(update) + self.ssc.start() + + def condition(): + self.assertEqual(result, [[0], [1], [2], [3]]) + return True + + eventually(condition, catch_assertions=True) + + @unittest.skip("SPARK-10086: Flaky StreamingKMeans test in PySpark") + def test_trainOn_predictOn(self): + """Test that prediction happens on the updated model.""" + stkm = StreamingKMeans(decayFactor=0.0, k=2) + stkm.setInitialCenters([[0.0], [1.0]], [1.0, 1.0]) + + # Since decay factor is set to zero, once the first batch + # is passed the clusterCenters are updated to [-0.5, 0.7] + # which causes 0.2 & 0.3 to be classified as 1, even though the + # classification based in the initial model would have been 0 + # proving that the model is updated. + batches = [[[-0.5], [0.6], [0.8]], [[0.2], [-0.1], [0.3]]] + batches = [self.sc.parallelize(batch) for batch in batches] + input_stream = self.ssc.queueStream(batches) + predict_results = [] + + def collect(rdd): + rdd_collect = rdd.collect() + if rdd_collect: + predict_results.append(rdd_collect) + + stkm.trainOn(input_stream) + predict_stream = stkm.predictOn(input_stream) + predict_stream.foreachRDD(collect) + + self.ssc.start() + + def condition(): + self.assertEqual(predict_results, [[0, 1, 1], [1, 0, 1]]) + return True + + eventually(condition, catch_assertions=True) + + +class StreamingLogisticRegressionWithSGDTests(MLLibStreamingTestCase): + @staticmethod + def generateLogisticInput(offset, scale, nPoints, seed): + """ + Generate 1 / (1 + exp(-x * scale + offset)) + + where, + x is randomly distributed and the threshold + and labels for each sample in x is obtained from a random uniform + distribution. + """ + rng = random.RandomState(seed) + x = rng.randn(nPoints) + sigmoid = 1.0 / (1 + exp(-(dot(x, scale) + offset))) + y_p = rng.rand(nPoints) + cut_off = y_p <= sigmoid + y_p[cut_off] = 1.0 + y_p[~cut_off] = 0.0 + return [LabeledPoint(y_p[i], Vectors.dense([x[i]])) for i in range(nPoints)] + + def test_parameter_accuracy(self): + """ + Test that the final value of weights is close to the desired value. + """ + input_batches = [ + self.sc.parallelize(self.generateLogisticInput(0, 1.5, 100, 42 + i)) for i in range(20) + ] + input_stream = self.ssc.queueStream(input_batches) + + slr = StreamingLogisticRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([0.0]) + slr.trainOn(input_stream) + + self.ssc.start() + + def condition(): + rel = (1.5 - slr.latestModel().weights.array[0]) / 1.5 + self.assertAlmostEqual(rel, 0.1, 1) + return True + + eventually(condition, timeout=120.0, catch_assertions=True) + + def test_convergence(self): + """ + Test that weights converge to the required value on toy data. + """ + input_batches = [ + self.sc.parallelize(self.generateLogisticInput(0, 1.5, 100, 42 + i)) for i in range(20) + ] + input_stream = self.ssc.queueStream(input_batches) + models = [] + + slr = StreamingLogisticRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([0.0]) + slr.trainOn(input_stream) + input_stream.foreachRDD(lambda x: models.append(slr.latestModel().weights[0])) + + self.ssc.start() + + def condition(): + self.assertEqual(len(models), len(input_batches)) + return True + + # We want all batches to finish for this test. + eventually(condition, 120, catch_assertions=True) + + t_models = array(models) + diff = t_models[1:] - t_models[:-1] + # Test that weights improve with a small tolerance + self.assertTrue(all(diff >= -0.1)) + self.assertTrue(array_sum(diff > 0) > 1) + + @staticmethod + def calculate_accuracy_error(true, predicted): + return sum(abs(array(true) - array(predicted))) / len(true) + + def test_predictions(self): + """Test predicted values on a toy model.""" + input_batches = [] + for i in range(20): + batch = self.sc.parallelize(self.generateLogisticInput(0, 1.5, 100, 42 + i)) + input_batches.append(batch.map(lambda x: (x.label, x.features))) + input_stream = self.ssc.queueStream(input_batches) + + slr = StreamingLogisticRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([1.5]) + predict_stream = slr.predictOnValues(input_stream) + true_predicted = [] + predict_stream.foreachRDD(lambda x: true_predicted.append(x.collect())) + self.ssc.start() + + def condition(): + self.assertEqual(len(true_predicted), len(input_batches)) + return True + + eventually(condition, catch_assertions=True) + + # Test that the accuracy error is no more than 0.4 on each batch. + for batch in true_predicted: + true, predicted = zip(*batch) + self.assertTrue(self.calculate_accuracy_error(true, predicted) < 0.4) + + @unittest.skipIf( + "COVERAGE_PROCESS_START" in os.environ, "Flaky with coverage enabled, skipping for now." + ) + def test_training_and_prediction(self): + """Test that the model improves on toy data with no. of batches""" + input_batches = [ + self.sc.parallelize(self.generateLogisticInput(0, 1.5, 100, 42 + i)) for i in range(40) + ] + predict_batches = [b.map(lambda lp: (lp.label, lp.features)) for b in input_batches] + + slr = StreamingLogisticRegressionWithSGD(stepSize=0.01, numIterations=25) + slr.setInitialWeights([-0.1]) + errors = [] + + def collect_errors(rdd): + true, predicted = zip(*rdd.collect()) + errors.append(self.calculate_accuracy_error(true, predicted)) + + input_stream = self.ssc.queueStream(input_batches) + predict_stream = self.ssc.queueStream(predict_batches) + slr.trainOn(input_stream) + ps = slr.predictOnValues(predict_stream) + ps.foreachRDD(lambda x: collect_errors(x)) + + self.ssc.start() + + def condition(): + # Test that the improvement in error is > 0.3 + if len(errors) == len(predict_batches): + self.assertGreater(errors[1] - errors[-1], 0.3) + if len(errors) >= 3 and errors[1] - errors[-1] > 0.3: + return True + return "Latest errors: " + ", ".join(map(lambda x: str(x), errors)) + + eventually(condition, timeout=180.0) + + +class StreamingLinearRegressionWithTests(MLLibStreamingTestCase): + def assertArrayAlmostEqual(self, array1, array2, dec): + for i, j in array1, array2: + self.assertAlmostEqual(i, j, dec) + + def test_parameter_accuracy(self): + """Test that coefs are predicted accurately by fitting on toy data.""" + + # Test that fitting (10*X1 + 10*X2), (X1, X2) gives coefficients + # (10, 10) + slr = StreamingLinearRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([0.0, 0.0]) + xMean = [0.0, 0.0] + xVariance = [1.0 / 3.0, 1.0 / 3.0] + + # Create ten batches with 100 sample points in each. + batches = [] + for i in range(10): + batch = LinearDataGenerator.generateLinearInput( + 0.0, [10.0, 10.0], xMean, xVariance, 100, 42 + i, 0.1 + ) + batches.append(self.sc.parallelize(batch)) + + input_stream = self.ssc.queueStream(batches) + slr.trainOn(input_stream) + self.ssc.start() + + def condition(): + self.assertArrayAlmostEqual(slr.latestModel().weights.array, [10.0, 10.0], 1) + self.assertAlmostEqual(slr.latestModel().intercept, 0.0, 1) + return True + + eventually(condition, catch_assertions=True) + + def test_parameter_convergence(self): + """Test that the model parameters improve with streaming data.""" + slr = StreamingLinearRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([0.0]) + + # Create ten batches with 100 sample points in each. + batches = [] + for i in range(10): + batch = LinearDataGenerator.generateLinearInput( + 0.0, [10.0], [0.0], [1.0 / 3.0], 100, 42 + i, 0.1 + ) + batches.append(self.sc.parallelize(batch)) + + model_weights = [] + input_stream = self.ssc.queueStream(batches) + input_stream.foreachRDD(lambda x: model_weights.append(slr.latestModel().weights[0])) + slr.trainOn(input_stream) + self.ssc.start() + + def condition(): + self.assertEqual(len(model_weights), len(batches)) + return True + + # We want all batches to finish for this test. + eventually(condition, 90, catch_assertions=True) + + w = array(model_weights) + diff = w[1:] - w[:-1] + self.assertTrue(all(diff >= -0.1)) + + def test_prediction(self): + """Test prediction on a model with weights already set.""" + # Create a model with initial Weights equal to coefs + slr = StreamingLinearRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([10.0, 10.0]) + + # Create ten batches with 100 sample points in each. + batches = [] + for i in range(10): + batch = LinearDataGenerator.generateLinearInput( + 0.0, [10.0, 10.0], [0.0, 0.0], [1.0 / 3.0, 1.0 / 3.0], 100, 42 + i, 0.1 + ) + batches.append(self.sc.parallelize(batch).map(lambda lp: (lp.label, lp.features))) + + input_stream = self.ssc.queueStream(batches) + output_stream = slr.predictOnValues(input_stream) + samples = [] + output_stream.foreachRDD(lambda x: samples.append(x.collect())) + + self.ssc.start() + + def condition(): + self.assertEqual(len(samples), len(batches)) + return True + + # We want all batches to finish for this test. + eventually(condition, catch_assertions=True) + + # Test that mean absolute error on each batch is less than 0.1 + for batch in samples: + true, predicted = zip(*batch) + self.assertTrue(mean(abs(array(true) - array(predicted))) < 0.1) + + @unittest.skipIf( + "COVERAGE_PROCESS_START" in os.environ, "Flaky with coverage enabled, skipping for now." + ) + def test_train_prediction(self): + """Test that error on test data improves as model is trained.""" + slr = StreamingLinearRegressionWithSGD(stepSize=0.2, numIterations=25) + slr.setInitialWeights([0.0]) + + # Create fifteen batches with 100 sample points in each. + batches = [] + for i in range(15): + batch = LinearDataGenerator.generateLinearInput( + 0.0, [10.0], [0.0], [1.0 / 3.0], 100, 42 + i, 0.1 + ) + batches.append(self.sc.parallelize(batch)) + + predict_batches = [b.map(lambda lp: (lp.label, lp.features)) for b in batches] + errors = [] + + def func(rdd): + true, predicted = zip(*rdd.collect()) + errors.append(mean(abs(true) - abs(predicted))) + + input_stream = self.ssc.queueStream(batches) + output_stream = self.ssc.queueStream(predict_batches) + slr.trainOn(input_stream) + output_stream = slr.predictOnValues(output_stream) + output_stream.foreachRDD(func) + self.ssc.start() + + def condition(): + if len(errors) == len(predict_batches): + self.assertGreater(errors[1] - errors[-1], 2) + if len(errors) >= 3 and errors[1] - errors[-1] > 2: + return True + return "Latest errors: " + ", ".join(map(lambda x: str(x), errors)) + + eventually(condition, timeout=180.0) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_streaming_algorithms import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_util.py new file mode 100644 index 0000000..28a53af --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tests/test_util.py @@ -0,0 +1,108 @@ +# +# 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 os +import tempfile +import unittest + +from pyspark.mllib.common import _to_java_object_rdd +from pyspark.mllib.util import LinearDataGenerator +from pyspark.mllib.util import MLUtils +from pyspark.mllib.linalg import SparseVector, DenseVector, Vectors +from pyspark.mllib.random import RandomRDDs +from pyspark.testing.mllibutils import MLlibTestCase + + +class MLUtilsTests(MLlibTestCase): + def test_append_bias(self): + data = [2.0, 2.0, 2.0] + ret = MLUtils.appendBias(data) + self.assertEqual(ret[3], 1.0) + self.assertEqual(type(ret), DenseVector) + + def test_append_bias_with_vector(self): + data = Vectors.dense([2.0, 2.0, 2.0]) + ret = MLUtils.appendBias(data) + self.assertEqual(ret[3], 1.0) + self.assertEqual(type(ret), DenseVector) + + def test_append_bias_with_sp_vector(self): + data = Vectors.sparse(3, {0: 2.0, 2: 2.0}) + expected = Vectors.sparse(4, {0: 2.0, 2: 2.0, 3: 1.0}) + # Returned value must be SparseVector + ret = MLUtils.appendBias(data) + self.assertEqual(ret, expected) + self.assertEqual(type(ret), SparseVector) + + def test_load_vectors(self): + import shutil + + data = [[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]] + temp_dir = tempfile.mkdtemp() + load_vectors_path = os.path.join(temp_dir, "test_load_vectors") + try: + self.sc.parallelize(data).saveAsTextFile(load_vectors_path) + ret_rdd = MLUtils.loadVectors(self.sc, load_vectors_path) + ret = ret_rdd.collect() + self.assertEqual(len(ret), 2) + self.assertEqual(ret[0], DenseVector([1.0, 2.0, 3.0])) + self.assertEqual(ret[1], DenseVector([1.0, 2.0, 3.0])) + except BaseException: + self.fail() + finally: + shutil.rmtree(load_vectors_path) + + +class LinearDataGeneratorTests(MLlibTestCase): + def test_dim(self): + linear_data = LinearDataGenerator.generateLinearInput( + intercept=0.0, + weights=[0.0, 0.0, 0.0], + xMean=[0.0, 0.0, 0.0], + xVariance=[0.33, 0.33, 0.33], + nPoints=4, + seed=0, + eps=0.1, + ) + self.assertEqual(len(linear_data), 4) + for point in linear_data: + self.assertEqual(len(point.features), 3) + + linear_data = LinearDataGenerator.generateLinearRDD( + sc=self.sc, nexamples=6, nfeatures=2, eps=0.1, nParts=2, intercept=0.0 + ).collect() + self.assertEqual(len(linear_data), 6) + for point in linear_data: + self.assertEqual(len(point.features), 2) + + +class SerDeTest(MLlibTestCase): + def test_to_java_object_rdd(self): # SPARK-6660 + data = RandomRDDs.uniformRDD(self.sc, 10, 5, seed=0) + self.assertEqual(_to_java_object_rdd(data).count(), 10) + + +if __name__ == "__main__": + from pyspark.mllib.tests.test_util import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tree.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tree.py new file mode 100644 index 0000000..8a5c25d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/tree.py @@ -0,0 +1,888 @@ +# +# 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 sys +import random + +from pyspark import RDD, since +from pyspark.mllib.common import callMLlibFunc, inherit_doc, JavaModelWrapper +from pyspark.mllib.linalg import _convert_to_vector +from pyspark.mllib.regression import LabeledPoint +from pyspark.mllib.util import JavaLoader, JavaSaveable +from typing import Dict, Optional, Tuple, Union, overload, TYPE_CHECKING +from pyspark.rdd import RDD + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike + + +__all__ = [ + "DecisionTreeModel", + "DecisionTree", + "RandomForestModel", + "RandomForest", + "GradientBoostedTreesModel", + "GradientBoostedTrees", +] + + +class TreeEnsembleModel(JavaModelWrapper, JavaSaveable): + """TreeEnsembleModel + + .. versionadded:: 1.3.0 + """ + + @overload + def predict(self, x: "VectorLike") -> float: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[float]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]: + """ + Predict values for a single data point or an RDD of points using + the model trained. + + .. versionadded:: 1.3.0 + + Notes + ----- + In Python, predict cannot currently be used within an RDD + transformation or action. + Call predict directly on the RDD instead. + """ + if isinstance(x, RDD): + return self.call("predict", x.map(_convert_to_vector)) + + else: + return self.call("predict", _convert_to_vector(x)) + + @since("1.3.0") + def numTrees(self) -> int: + """ + Get number of trees in ensemble. + """ + return self.call("numTrees") + + @since("1.3.0") + def totalNumNodes(self) -> int: + """ + Get total number of nodes, summed over all trees in the ensemble. + """ + return self.call("totalNumNodes") + + def __repr__(self) -> str: + """Summary of model""" + return self._java_model.toString() + + @since("1.3.0") + def toDebugString(self) -> str: + """Full model""" + return self._java_model.toDebugString() + + +class DecisionTreeModel(JavaModelWrapper, JavaSaveable, JavaLoader["DecisionTreeModel"]): + """ + A decision tree model for classification or regression. + + .. versionadded:: 1.1.0 + """ + + @overload + def predict(self, x: "VectorLike") -> float: + ... + + @overload + def predict(self, x: RDD["VectorLike"]) -> RDD[float]: + ... + + def predict(self, x: Union["VectorLike", RDD["VectorLike"]]) -> Union[float, RDD[float]]: + """ + Predict the label of one or more examples. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + Data point (feature vector), or an RDD of data points (feature + vectors). + + Notes + ----- + In Python, predict cannot currently be used within an RDD + transformation or action. + Call predict directly on the RDD instead. + """ + if isinstance(x, RDD): + return self.call("predict", x.map(_convert_to_vector)) + + else: + return self.call("predict", _convert_to_vector(x)) + + @since("1.1.0") + def numNodes(self) -> int: + """Get number of nodes in tree, including leaf nodes.""" + return self._java_model.numNodes() + + @since("1.1.0") + def depth(self) -> int: + """ + Get depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + """ + return self._java_model.depth() + + def __repr__(self) -> str: + """summary of model.""" + return self._java_model.toString() + + @since("1.2.0") + def toDebugString(self) -> str: + """full model.""" + return self._java_model.toDebugString() + + @classmethod + def _java_loader_class(cls) -> str: + return "org.apache.spark.mllib.tree.model.DecisionTreeModel" + + +class DecisionTree: + """ + Learning algorithm for a decision tree model for classification or + regression. + + .. versionadded:: 1.1.0 + """ + + @classmethod + def _train( + cls, + data: RDD[LabeledPoint], + type: str, + numClasses: int, + features: Dict[int, int], + impurity: str = "gini", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + ) -> DecisionTreeModel: + first = data.first() + assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" + model = callMLlibFunc( + "trainDecisionTreeModel", + data, + type, + numClasses, + features, + impurity, + maxDepth, + maxBins, + minInstancesPerNode, + minInfoGain, + ) + return DecisionTreeModel(model) + + @classmethod + def trainClassifier( + cls, + data: RDD[LabeledPoint], + numClasses: int, + categoricalFeaturesInfo: Dict[int, int], + impurity: str = "gini", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + ) -> DecisionTreeModel: + """ + Train a decision tree model for classification. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + Training data: RDD of LabeledPoint. Labels should take values + {0, 1, ..., numClasses-1}. + numClasses : int + Number of classes for classification. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + impurity : str, optional + Criterion used for information gain calculation. + Supported values: "gini" or "entropy". + (default: "gini") + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 5) + maxBins : int, optional + Number of bins used for finding splits at each node. + (default: 32) + minInstancesPerNode : int, optional + Minimum number of instances required at child nodes to create + the parent split. + (default: 1) + minInfoGain : float, optional + Minimum info gain required to create a split. + (default: 0.0) + + Returns + ------- + :py:class:`DecisionTreeModel` + + Examples + -------- + >>> from numpy import array + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import DecisionTree + >>> + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(1.0, [1.0]), + ... LabeledPoint(1.0, [2.0]), + ... LabeledPoint(1.0, [3.0]) + ... ] + >>> model = DecisionTree.trainClassifier(sc.parallelize(data), 2, {}) + >>> print(model) + DecisionTreeModel classifier of depth 1 with 3 nodes + + >>> print(model.toDebugString()) + DecisionTreeModel classifier of depth 1 with 3 nodes + If (feature 0 <= 0.5) + Predict: 0.0 + Else (feature 0 > 0.5) + Predict: 1.0 + >>> model.predict(array([1.0])) + 1.0 + >>> model.predict(array([0.0])) + 0.0 + >>> rdd = sc.parallelize([[1.0], [0.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.0] + """ + return cls._train( + data, + "classification", + numClasses, + categoricalFeaturesInfo, + impurity, + maxDepth, + maxBins, + minInstancesPerNode, + minInfoGain, + ) + + @classmethod + @since("1.1.0") + def trainRegressor( + cls, + data: RDD[LabeledPoint], + categoricalFeaturesInfo: Dict[int, int], + impurity: str = "variance", + maxDepth: int = 5, + maxBins: int = 32, + minInstancesPerNode: int = 1, + minInfoGain: float = 0.0, + ) -> DecisionTreeModel: + """ + Train a decision tree model for regression. + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + Training data: RDD of LabeledPoint. Labels are real numbers. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + impurity : str, optional + Criterion used for information gain calculation. + The only supported value for regression is "variance". + (default: "variance") + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 5) + maxBins : int, optional + Number of bins used for finding splits at each node. + (default: 32) + minInstancesPerNode : int, optional + Minimum number of instances required at child nodes to create + the parent split. + (default: 1) + minInfoGain : float, optional + Minimum info gain required to create a split. + (default: 0.0) + + Returns + ------- + :py:class:`DecisionTreeModel` + + Examples + -------- + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import DecisionTree + >>> from pyspark.mllib.linalg import SparseVector + >>> + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(0.0, SparseVector(2, {0: 0.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) + ... ] + >>> + >>> model = DecisionTree.trainRegressor(sc.parallelize(sparse_data), {}) + >>> model.predict(SparseVector(2, {1: 1.0})) + 1.0 + >>> model.predict(SparseVector(2, {1: 0.0})) + 0.0 + >>> rdd = sc.parallelize([[0.0, 1.0], [0.0, 0.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.0] + """ + return cls._train( + data, + "regression", + 0, + categoricalFeaturesInfo, + impurity, + maxDepth, + maxBins, + minInstancesPerNode, + minInfoGain, + ) + + +@inherit_doc +class RandomForestModel(TreeEnsembleModel, JavaLoader["RandomForestModel"]): + """ + Represents a random forest model. + + .. versionadded:: 1.2.0 + """ + + @classmethod + def _java_loader_class(cls) -> str: + return "org.apache.spark.mllib.tree.model.RandomForestModel" + + +class RandomForest: + """ + Learning algorithm for a random forest model for classification or + regression. + + .. versionadded:: 1.2.0 + """ + + supportedFeatureSubsetStrategies: Tuple[str, ...] = ("auto", "all", "sqrt", "log2", "onethird") + + @classmethod + def _train( + cls, + data: RDD[LabeledPoint], + algo: str, + numClasses: int, + categoricalFeaturesInfo: Dict[int, int], + numTrees: int, + featureSubsetStrategy: str, + impurity: str, + maxDepth: int, + maxBins: int, + seed: Optional[int], + ) -> RandomForestModel: + first = data.first() + assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" + if featureSubsetStrategy not in cls.supportedFeatureSubsetStrategies: + raise ValueError("unsupported featureSubsetStrategy: %s" % featureSubsetStrategy) + if seed is None: + seed = random.randint(0, 1 << 30) + model = callMLlibFunc( + "trainRandomForestModel", + data, + algo, + numClasses, + categoricalFeaturesInfo, + numTrees, + featureSubsetStrategy, + impurity, + maxDepth, + maxBins, + seed, + ) + return RandomForestModel(model) + + @classmethod + def trainClassifier( + cls, + data: RDD[LabeledPoint], + numClasses: int, + categoricalFeaturesInfo: Dict[int, int], + numTrees: int, + featureSubsetStrategy: str = "auto", + impurity: str = "gini", + maxDepth: int = 4, + maxBins: int = 32, + seed: Optional[int] = None, + ) -> RandomForestModel: + """ + Train a random forest model for binary or multiclass + classification. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + Training dataset: RDD of LabeledPoint. Labels should take values + {0, 1, ..., numClasses-1}. + numClasses : int + Number of classes for classification. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + numTrees : int + Number of trees in the random forest. + featureSubsetStrategy : str, optional + Number of features to consider for splits at each node. + Supported values: "auto", "all", "sqrt", "log2", "onethird". + If "auto" is set, this parameter is set based on numTrees: + if numTrees == 1, set to "all"; + if numTrees > 1 (forest) set to "sqrt". + (default: "auto") + impurity : str, optional + Criterion used for information gain calculation. + Supported values: "gini" or "entropy". + (default: "gini") + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 4) + maxBins : int, optional + Maximum number of bins used for splitting features. + (default: 32) + seed : int, Optional + Random seed for bootstrapping and choosing feature subsets. + Set as None to generate seed based on system time. + (default: None) + + Returns + ------- + :py:class:`RandomForestModel` + that can be used for prediction. + + Examples + -------- + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import RandomForest + >>> + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(0.0, [1.0]), + ... LabeledPoint(1.0, [2.0]), + ... LabeledPoint(1.0, [3.0]) + ... ] + >>> model = RandomForest.trainClassifier(sc.parallelize(data), 2, {}, 3, seed=42) + >>> model.numTrees() + 3 + >>> model.totalNumNodes() + 7 + >>> print(model) + TreeEnsembleModel classifier with 3 trees + >>> print(model.toDebugString()) + TreeEnsembleModel classifier with 3 trees + Tree 0: + Predict: 1.0 + Tree 1: + If (feature 0 <= 1.5) + Predict: 0.0 + Else (feature 0 > 1.5) + Predict: 1.0 + Tree 2: + If (feature 0 <= 1.5) + Predict: 0.0 + Else (feature 0 > 1.5) + Predict: 1.0 + >>> model.predict([2.0]) + 1.0 + >>> model.predict([0.0]) + 0.0 + >>> rdd = sc.parallelize([[3.0], [1.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.0] + """ + return cls._train( + data, + "classification", + numClasses, + categoricalFeaturesInfo, + numTrees, + featureSubsetStrategy, + impurity, + maxDepth, + maxBins, + seed, + ) + + @classmethod + def trainRegressor( + cls, + data: RDD[LabeledPoint], + categoricalFeaturesInfo: Dict[int, int], + numTrees: int, + featureSubsetStrategy: str = "auto", + impurity: str = "variance", + maxDepth: int = 4, + maxBins: int = 32, + seed: Optional[int] = None, + ) -> RandomForestModel: + """ + Train a random forest model for regression. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + Training dataset: RDD of LabeledPoint. Labels are real numbers. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + numTrees : int + Number of trees in the random forest. + featureSubsetStrategy : str, optional + Number of features to consider for splits at each node. + Supported values: "auto", "all", "sqrt", "log2", "onethird". + If "auto" is set, this parameter is set based on numTrees: + + - if numTrees == 1, set to "all"; + - if numTrees > 1 (forest) set to "onethird" for regression. + + (default: "auto") + impurity : str, optional + Criterion used for information gain calculation. + The only supported value for regression is "variance". + (default: "variance") + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 4) + maxBins : int, optional + Maximum number of bins used for splitting features. + (default: 32) + seed : int, optional + Random seed for bootstrapping and choosing feature subsets. + Set as None to generate seed based on system time. + (default: None) + + Returns + ------- + :py:class:`RandomForestModel` + that can be used for prediction. + + Examples + -------- + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import RandomForest + >>> from pyspark.mllib.linalg import SparseVector + >>> + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) + ... ] + >>> + >>> model = RandomForest.trainRegressor(sc.parallelize(sparse_data), {}, 2, seed=42) + >>> model.numTrees() + 2 + >>> model.totalNumNodes() + 4 + >>> model.predict(SparseVector(2, {1: 1.0})) + 1.0 + >>> model.predict(SparseVector(2, {0: 1.0})) + 0.5 + >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.5] + """ + return cls._train( + data, + "regression", + 0, + categoricalFeaturesInfo, + numTrees, + featureSubsetStrategy, + impurity, + maxDepth, + maxBins, + seed, + ) + + +@inherit_doc +class GradientBoostedTreesModel(TreeEnsembleModel, JavaLoader["GradientBoostedTreesModel"]): + """ + Represents a gradient-boosted tree model. + + .. versionadded:: 1.3.0 + """ + + @classmethod + def _java_loader_class(cls) -> str: + return "org.apache.spark.mllib.tree.model.GradientBoostedTreesModel" + + +class GradientBoostedTrees: + """ + Learning algorithm for a gradient boosted trees model for + classification or regression. + + .. versionadded:: 1.3.0 + """ + + @classmethod + def _train( + cls, + data: RDD[LabeledPoint], + algo: str, + categoricalFeaturesInfo: Dict[int, int], + loss: str, + numIterations: int, + learningRate: float, + maxDepth: int, + maxBins: int, + ) -> GradientBoostedTreesModel: + first = data.first() + assert isinstance(first, LabeledPoint), "the data should be RDD of LabeledPoint" + model = callMLlibFunc( + "trainGradientBoostedTreesModel", + data, + algo, + categoricalFeaturesInfo, + loss, + numIterations, + learningRate, + maxDepth, + maxBins, + ) + return GradientBoostedTreesModel(model) + + @classmethod + def trainClassifier( + cls, + data: RDD[LabeledPoint], + categoricalFeaturesInfo: Dict[int, int], + loss: str = "logLoss", + numIterations: int = 100, + learningRate: float = 0.1, + maxDepth: int = 3, + maxBins: int = 32, + ) -> GradientBoostedTreesModel: + """ + Train a gradient-boosted trees model for classification. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + Training dataset: RDD of LabeledPoint. Labels should take values + {0, 1}. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + loss : str, optional + Loss function used for minimization during gradient boosting. + Supported values: "logLoss", "leastSquaresError", + "leastAbsoluteError". + (default: "logLoss") + numIterations : int, optional + Number of iterations of boosting. + (default: 100) + learningRate : float, optional + Learning rate for shrinking the contribution of each estimator. + The learning rate should be between in the interval (0, 1]. + (default: 0.1) + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 3) + maxBins : int, optional + Maximum number of bins used for splitting features. DecisionTree + requires maxBins >= max categories. + (default: 32) + + Returns + ------- + :py:class:`GradientBoostedTreesModel` + that can be used for prediction. + + Examples + -------- + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import GradientBoostedTrees + >>> + >>> data = [ + ... LabeledPoint(0.0, [0.0]), + ... LabeledPoint(0.0, [1.0]), + ... LabeledPoint(1.0, [2.0]), + ... LabeledPoint(1.0, [3.0]) + ... ] + >>> + >>> model = GradientBoostedTrees.trainClassifier(sc.parallelize(data), {}, numIterations=10) + >>> model.numTrees() + 10 + >>> model.totalNumNodes() + 30 + >>> print(model) # it already has newline + TreeEnsembleModel classifier with 10 trees + >>> model.predict([2.0]) + 1.0 + >>> model.predict([0.0]) + 0.0 + >>> rdd = sc.parallelize([[2.0], [0.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.0] + """ + return cls._train( + data, + "classification", + categoricalFeaturesInfo, + loss, + numIterations, + learningRate, + maxDepth, + maxBins, + ) + + @classmethod + def trainRegressor( + cls, + data: RDD[LabeledPoint], + categoricalFeaturesInfo: Dict[int, int], + loss: str = "leastSquaresError", + numIterations: int = 100, + learningRate: float = 0.1, + maxDepth: int = 3, + maxBins: int = 32, + ) -> GradientBoostedTreesModel: + """ + Train a gradient-boosted trees model for regression. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + data : + Training dataset: RDD of LabeledPoint. Labels are real numbers. + categoricalFeaturesInfo : dict + Map storing arity of categorical features. An entry (n -> k) + indicates that feature n is categorical with k categories + indexed from 0: {0, 1, ..., k-1}. + loss : str, optional + Loss function used for minimization during gradient boosting. + Supported values: "logLoss", "leastSquaresError", + "leastAbsoluteError". + (default: "leastSquaresError") + numIterations : int, optional + Number of iterations of boosting. + (default: 100) + learningRate : float, optional + Learning rate for shrinking the contribution of each estimator. + The learning rate should be between in the interval (0, 1]. + (default: 0.1) + maxDepth : int, optional + Maximum depth of tree (e.g. depth 0 means 1 leaf node, depth 1 + means 1 internal node + 2 leaf nodes). + (default: 3) + maxBins : int, optional + Maximum number of bins used for splitting features. DecisionTree + requires maxBins >= max categories. + (default: 32) + + Returns + ------- + :py:class:`GradientBoostedTreesModel` + that can be used for prediction. + + Examples + -------- + >>> from pyspark.mllib.regression import LabeledPoint + >>> from pyspark.mllib.tree import GradientBoostedTrees + >>> from pyspark.mllib.linalg import SparseVector + >>> + >>> sparse_data = [ + ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 1.0})), + ... LabeledPoint(0.0, SparseVector(2, {0: 1.0})), + ... LabeledPoint(1.0, SparseVector(2, {1: 2.0})) + ... ] + >>> + >>> data = sc.parallelize(sparse_data) + >>> model = GradientBoostedTrees.trainRegressor(data, {}, numIterations=10) + >>> model.numTrees() + 10 + >>> model.totalNumNodes() + 12 + >>> model.predict(SparseVector(2, {1: 1.0})) + 1.0 + >>> model.predict(SparseVector(2, {0: 1.0})) + 0.0 + >>> rdd = sc.parallelize([[0.0, 1.0], [1.0, 0.0]]) + >>> model.predict(rdd).collect() + [1.0, 0.0] + """ + return cls._train( + data, + "regression", + categoricalFeaturesInfo, + loss, + numIterations, + learningRate, + maxDepth, + maxBins, + ) + + +def _test() -> None: + import doctest + + globs = globals().copy() + from pyspark.sql import SparkSession + + spark = SparkSession.builder.master("local[4]").appName("mllib.tree tests").getOrCreate() + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod( + globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/util.py new file mode 100644 index 0000000..8f28e2c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/mllib/util.py @@ -0,0 +1,659 @@ +# +# 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 sys +from functools import reduce +import numpy as np + +from pyspark import SparkContext, since +from pyspark.mllib.common import callMLlibFunc, inherit_doc +from pyspark.mllib.linalg import Vectors, SparseVector, _convert_to_vector +from pyspark.sql import DataFrame +from typing import Generic, Iterable, List, Optional, Tuple, Type, TypeVar, cast, TYPE_CHECKING +from pyspark.context import SparkContext +from pyspark.mllib.linalg import Vector +from pyspark.rdd import RDD +from pyspark.sql.dataframe import DataFrame + +T = TypeVar("T") +L = TypeVar("L", bound="Loader") +JL = TypeVar("JL", bound="JavaLoader") + +if TYPE_CHECKING: + from pyspark.mllib._typing import VectorLike + from py4j.java_gateway import JavaObject + from pyspark.mllib.regression import LabeledPoint + + +class MLUtils: + + """ + Helper methods to load, save and pre-process data used in MLlib. + + .. versionadded:: 1.0.0 + """ + + @staticmethod + def _parse_libsvm_line(line: str) -> Tuple[float, np.ndarray, np.ndarray]: + """ + Parses a line in LIBSVM format into (label, indices, values). + """ + items = line.split(None) + label = float(items[0]) + nnz = len(items) - 1 + indices = np.zeros(nnz, dtype=np.int32) + values = np.zeros(nnz) + for i in range(nnz): + index, value = items[1 + i].split(":") + indices[i] = int(index) - 1 + values[i] = float(value) + return label, indices, values + + @staticmethod + def _convert_labeled_point_to_libsvm(p: "LabeledPoint") -> str: + """Converts a LabeledPoint to a string in LIBSVM format.""" + from pyspark.mllib.regression import LabeledPoint + + assert isinstance(p, LabeledPoint) + items = [str(p.label)] + v = _convert_to_vector(p.features) + if isinstance(v, SparseVector): + nnz = len(v.indices) + for i in range(nnz): + items.append(str(v.indices[i] + 1) + ":" + str(v.values[i])) + else: + for i in range(len(v)): + items.append(str(i + 1) + ":" + str(v[i])) # type: ignore[index] + return " ".join(items) + + @staticmethod + def loadLibSVMFile( + sc: SparkContext, path: str, numFeatures: int = -1, minPartitions: Optional[int] = None + ) -> RDD["LabeledPoint"]: + """ + Loads labeled data in the LIBSVM format into an RDD of + LabeledPoint. The LIBSVM format is a text-based format used by + LIBSVM and LIBLINEAR. Each line represents a labeled sparse + feature vector using the following format: + + label index1:value1 index2:value2 ... + + where the indices are one-based and in ascending order. This + method parses each line into a LabeledPoint, where the feature + indices are converted to zero-based. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + Spark context + path : str + file or directory path in any Hadoop-supported file system URI + numFeatures : int, optional + number of features, which will be determined + from the input data if a nonpositive value + is given. This is useful when the dataset is + already split into multiple files and you + want to load them separately, because some + features may not present in certain files, + which leads to inconsistent feature + dimensions. + minPartitions : int, optional + min number of partitions + + Returns + ------- + :py:class:`pyspark.RDD` + labeled data stored as an RDD of LabeledPoint + + Examples + -------- + >>> from tempfile import NamedTemporaryFile + >>> from pyspark.mllib.util import MLUtils + >>> from pyspark.mllib.regression import LabeledPoint + >>> tempFile = NamedTemporaryFile(delete=True) + >>> _ = tempFile.write(b"+1 1:1.0 3:2.0 5:3.0\\n-1\\n-1 2:4.0 4:5.0 6:6.0") + >>> tempFile.flush() + >>> examples = MLUtils.loadLibSVMFile(sc, tempFile.name).collect() + >>> tempFile.close() + >>> examples[0] + LabeledPoint(1.0, (6,[0,2,4],[1.0,2.0,3.0])) + >>> examples[1] + LabeledPoint(-1.0, (6,[],[])) + >>> examples[2] + LabeledPoint(-1.0, (6,[1,3,5],[4.0,5.0,6.0])) + """ + from pyspark.mllib.regression import LabeledPoint + + lines = sc.textFile(path, minPartitions) + parsed = lines.map(lambda l: MLUtils._parse_libsvm_line(l)) + if numFeatures <= 0: + parsed.cache() + numFeatures = parsed.map(lambda x: -1 if x[1].size == 0 else x[1][-1]).reduce(max) + 1 + return parsed.map( + lambda x: LabeledPoint( + x[0], Vectors.sparse(numFeatures, x[1], x[2]) # type: ignore[arg-type] + ) + ) + + @staticmethod + def saveAsLibSVMFile(data: RDD["LabeledPoint"], dir: str) -> None: + """ + Save labeled data in LIBSVM format. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + data : :py:class:`pyspark.RDD` + an RDD of LabeledPoint to be saved + dir : str + directory to save the data + + Examples + -------- + >>> from tempfile import NamedTemporaryFile + >>> from fileinput import input + >>> from pyspark.mllib.regression import LabeledPoint + >>> from glob import glob + >>> from pyspark.mllib.util import MLUtils + >>> examples = [LabeledPoint(1.1, Vectors.sparse(3, [(0, 1.23), (2, 4.56)])), + ... LabeledPoint(0.0, Vectors.dense([1.01, 2.02, 3.03]))] + >>> tempFile = NamedTemporaryFile(delete=True) + >>> tempFile.close() + >>> MLUtils.saveAsLibSVMFile(sc.parallelize(examples), tempFile.name) + >>> ''.join(sorted(input(glob(tempFile.name + "/part-0000*")))) + '0.0 1:1.01 2:2.02 3:3.03\\n1.1 1:1.23 3:4.56\\n' + """ + lines = data.map(lambda p: MLUtils._convert_labeled_point_to_libsvm(p)) + lines.saveAsTextFile(dir) + + @staticmethod + def loadLabeledPoints( + sc: SparkContext, path: str, minPartitions: Optional[int] = None + ) -> RDD["LabeledPoint"]: + """ + Load labeled points saved using RDD.saveAsTextFile. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + Spark context + path : str + file or directory path in any Hadoop-supported file system URI + minPartitions : int, optional + min number of partitions + + Returns + ------- + :py:class:`pyspark.RDD` + labeled data stored as an RDD of LabeledPoint + + Examples + -------- + >>> from tempfile import NamedTemporaryFile + >>> from pyspark.mllib.util import MLUtils + >>> from pyspark.mllib.regression import LabeledPoint + >>> examples = [LabeledPoint(1.1, Vectors.sparse(3, [(0, -1.23), (2, 4.56e-7)])), + ... LabeledPoint(0.0, Vectors.dense([1.01, 2.02, 3.03]))] + >>> tempFile = NamedTemporaryFile(delete=True) + >>> tempFile.close() + >>> sc.parallelize(examples, 1).saveAsTextFile(tempFile.name) + >>> MLUtils.loadLabeledPoints(sc, tempFile.name).collect() + [LabeledPoint(1.1, (3,[0,2],[-1.23,4.56e-07])), LabeledPoint(0.0, [1.01,2.02,3.03])] + """ + minPartitions = minPartitions or min(sc.defaultParallelism, 2) + return callMLlibFunc("loadLabeledPoints", sc, path, minPartitions) + + @staticmethod + @since("1.5.0") + def appendBias(data: Vector) -> Vector: + """ + Returns a new vector with `1.0` (bias) appended to + the end of the input vector. + """ + vec = _convert_to_vector(data) + if isinstance(vec, SparseVector): + newIndices = np.append(vec.indices, len(vec)) + newValues = np.append(vec.values, 1.0) + return SparseVector(len(vec) + 1, newIndices, newValues) + else: + return _convert_to_vector(np.append(vec.toArray(), 1.0)) + + @staticmethod + @since("1.5.0") + def loadVectors(sc: SparkContext, path: str) -> RDD[Vector]: + """ + Loads vectors saved using `RDD[Vector].saveAsTextFile` + with the default number of partitions. + """ + return callMLlibFunc("loadVectors", sc, path) + + @staticmethod + def convertVectorColumnsToML(dataset: DataFrame, *cols: str) -> DataFrame: + """ + Converts vector columns in an input DataFrame from the + :py:class:`pyspark.mllib.linalg.Vector` type to the new + :py:class:`pyspark.ml.linalg.Vector` type under the `spark.ml` + package. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + \\*cols : str + Vector columns to be converted. + + New vector columns will be ignored. If unspecified, all old + vector columns will be converted excepted nested ones. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + the input dataset with old vector columns converted to the + new vector type + + Examples + -------- + >>> import pyspark + >>> from pyspark.mllib.linalg import Vectors + >>> from pyspark.mllib.util import MLUtils + >>> df = spark.createDataFrame( + ... [(0, Vectors.sparse(2, [1], [1.0]), Vectors.dense(2.0, 3.0))], + ... ["id", "x", "y"]) + >>> r1 = MLUtils.convertVectorColumnsToML(df).first() + >>> isinstance(r1.x, pyspark.ml.linalg.SparseVector) + True + >>> isinstance(r1.y, pyspark.ml.linalg.DenseVector) + True + >>> r2 = MLUtils.convertVectorColumnsToML(df, "x").first() + >>> isinstance(r2.x, pyspark.ml.linalg.SparseVector) + True + >>> isinstance(r2.y, pyspark.mllib.linalg.DenseVector) + True + """ + if not isinstance(dataset, DataFrame): + raise TypeError("Input dataset must be a DataFrame but got {}.".format(type(dataset))) + return callMLlibFunc("convertVectorColumnsToML", dataset, list(cols)) + + @staticmethod + def convertVectorColumnsFromML(dataset: DataFrame, *cols: str) -> DataFrame: + """ + Converts vector columns in an input DataFrame to the + :py:class:`pyspark.mllib.linalg.Vector` type from the new + :py:class:`pyspark.ml.linalg.Vector` type under the `spark.ml` + package. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + \\*cols : str + Vector columns to be converted. + + Old vector columns will be ignored. If unspecified, all new + vector columns will be converted except nested ones. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + the input dataset with new vector columns converted to the + old vector type + + Examples + -------- + >>> import pyspark + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.mllib.util import MLUtils + >>> df = spark.createDataFrame( + ... [(0, Vectors.sparse(2, [1], [1.0]), Vectors.dense(2.0, 3.0))], + ... ["id", "x", "y"]) + >>> r1 = MLUtils.convertVectorColumnsFromML(df).first() + >>> isinstance(r1.x, pyspark.mllib.linalg.SparseVector) + True + >>> isinstance(r1.y, pyspark.mllib.linalg.DenseVector) + True + >>> r2 = MLUtils.convertVectorColumnsFromML(df, "x").first() + >>> isinstance(r2.x, pyspark.mllib.linalg.SparseVector) + True + >>> isinstance(r2.y, pyspark.ml.linalg.DenseVector) + True + """ + if not isinstance(dataset, DataFrame): + raise TypeError("Input dataset must be a DataFrame but got {}.".format(type(dataset))) + return callMLlibFunc("convertVectorColumnsFromML", dataset, list(cols)) + + @staticmethod + def convertMatrixColumnsToML(dataset: DataFrame, *cols: str) -> DataFrame: + """ + Converts matrix columns in an input DataFrame from the + :py:class:`pyspark.mllib.linalg.Matrix` type to the new + :py:class:`pyspark.ml.linalg.Matrix` type under the `spark.ml` + package. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + \\*cols : str + Matrix columns to be converted. + + New matrix columns will be ignored. If unspecified, all old + matrix columns will be converted excepted nested ones. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + the input dataset with old matrix columns converted to the + new matrix type + + Examples + -------- + >>> import pyspark + >>> from pyspark.mllib.linalg import Matrices + >>> from pyspark.mllib.util import MLUtils + >>> df = spark.createDataFrame( + ... [(0, Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]), + ... Matrices.dense(2, 2, range(4)))], ["id", "x", "y"]) + >>> r1 = MLUtils.convertMatrixColumnsToML(df).first() + >>> isinstance(r1.x, pyspark.ml.linalg.SparseMatrix) + True + >>> isinstance(r1.y, pyspark.ml.linalg.DenseMatrix) + True + >>> r2 = MLUtils.convertMatrixColumnsToML(df, "x").first() + >>> isinstance(r2.x, pyspark.ml.linalg.SparseMatrix) + True + >>> isinstance(r2.y, pyspark.mllib.linalg.DenseMatrix) + True + """ + if not isinstance(dataset, DataFrame): + raise TypeError("Input dataset must be a DataFrame but got {}.".format(type(dataset))) + return callMLlibFunc("convertMatrixColumnsToML", dataset, list(cols)) + + @staticmethod + def convertMatrixColumnsFromML(dataset: DataFrame, *cols: str) -> DataFrame: + """ + Converts matrix columns in an input DataFrame to the + :py:class:`pyspark.mllib.linalg.Matrix` type from the new + :py:class:`pyspark.ml.linalg.Matrix` type under the `spark.ml` + package. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + \\*cols : str + Matrix columns to be converted. + + Old matrix columns will be ignored. If unspecified, all new + matrix columns will be converted except nested ones. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` + the input dataset with new matrix columns converted to the + old matrix type + + Examples + -------- + >>> import pyspark + >>> from pyspark.ml.linalg import Matrices + >>> from pyspark.mllib.util import MLUtils + >>> df = spark.createDataFrame( + ... [(0, Matrices.sparse(2, 2, [0, 2, 3], [0, 1, 1], [2, 3, 4]), + ... Matrices.dense(2, 2, range(4)))], ["id", "x", "y"]) + >>> r1 = MLUtils.convertMatrixColumnsFromML(df).first() + >>> isinstance(r1.x, pyspark.mllib.linalg.SparseMatrix) + True + >>> isinstance(r1.y, pyspark.mllib.linalg.DenseMatrix) + True + >>> r2 = MLUtils.convertMatrixColumnsFromML(df, "x").first() + >>> isinstance(r2.x, pyspark.mllib.linalg.SparseMatrix) + True + >>> isinstance(r2.y, pyspark.ml.linalg.DenseMatrix) + True + """ + if not isinstance(dataset, DataFrame): + raise TypeError("Input dataset must be a DataFrame but got {}.".format(type(dataset))) + return callMLlibFunc("convertMatrixColumnsFromML", dataset, list(cols)) + + +class Saveable: + """ + Mixin for models and transformers which may be saved as files. + + .. versionadded:: 1.3.0 + """ + + def save(self, sc: SparkContext, path: str) -> None: + """ + Save this model to the given path. + + This saves: + * human-readable (JSON) model metadata to path/metadata/ + * Parquet formatted data to path/data/ + + The model may be loaded using :py:meth:`Loader.load`. + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + Spark context used to save model data. + path : str + Path specifying the directory in which to save + this model. If the directory already exists, + this method throws an exception. + """ + raise NotImplementedError + + +@inherit_doc +class JavaSaveable(Saveable): + """ + Mixin for models that provide save() through their Scala + implementation. + + .. versionadded:: 1.3.0 + """ + + _java_model: "JavaObject" + + @since("1.3.0") + def save(self, sc: SparkContext, path: str) -> None: + """Save this model to the given path.""" + if not isinstance(sc, SparkContext): + raise TypeError("sc should be a SparkContext, got type %s" % type(sc)) + if not isinstance(path, str): + raise TypeError("path should be a string, got type %s" % type(path)) + self._java_model.save(sc._jsc.sc(), path) + + +class Loader(Generic[T]): + """ + Mixin for classes which can load saved models from files. + + .. versionadded:: 1.3.0 + """ + + @classmethod + def load(cls: Type[L], sc: SparkContext, path: str) -> L: + """ + Load a model from the given path. The model should have been + saved using :py:meth:`Saveable.save`. + + Parameters + ---------- + sc : :py:class:`pyspark.SparkContext` + Spark context used for loading model files. + path : str + Path specifying the directory to which the model was saved. + + Returns + ------- + object + model instance + """ + raise NotImplementedError + + +@inherit_doc +class JavaLoader(Loader[T]): + """ + Mixin for classes which can load saved models using its Scala + implementation. + + .. versionadded:: 1.3.0 + """ + + @classmethod + def _java_loader_class(cls) -> str: + """ + Returns the full class name of the Java loader. The default + implementation replaces "pyspark" by "org.apache.spark" in + the Python full class name. + """ + java_package = cls.__module__.replace("pyspark", "org.apache.spark") + return ".".join([java_package, cls.__name__]) + + @classmethod + def _load_java(cls, sc: SparkContext, path: str) -> "JavaObject": + """ + Load a Java model from the given path. + """ + java_class = cls._java_loader_class() + java_obj: "JavaObject" = reduce(getattr, java_class.split("."), sc._jvm) + return java_obj.load(sc._jsc.sc(), path) + + @classmethod + @since("1.3.0") + def load(cls: Type[JL], sc: SparkContext, path: str) -> JL: + """Load a model from the given path.""" + java_model = cls._load_java(sc, path) + return cls(java_model) # type: ignore[call-arg] + + +class LinearDataGenerator: + """Utils for generating linear data. + + .. versionadded:: 1.5.0 + """ + + @staticmethod + def generateLinearInput( + intercept: float, + weights: "VectorLike", + xMean: "VectorLike", + xVariance: "VectorLike", + nPoints: int, + seed: int, + eps: float, + ) -> List["LabeledPoint"]: + """ + .. versionadded:: 1.5.0 + + Parameters + ---------- + intercept : float + bias factor, the term c in X'w + c + weights : :py:class:`pyspark.mllib.linalg.Vector` or convertible + feature vector, the term w in X'w + c + xMean : :py:class:`pyspark.mllib.linalg.Vector` or convertible + Point around which the data X is centered. + xVariance : :py:class:`pyspark.mllib.linalg.Vector` or convertible + Variance of the given data + nPoints : int + Number of points to be generated + seed : int + Random Seed + eps : float + Used to scale the noise. If eps is set high, + the amount of gaussian noise added is more. + + Returns + ------- + list + of :py:class:`pyspark.mllib.regression.LabeledPoints` of length nPoints + """ + weights = [float(weight) for weight in cast(Iterable[float], weights)] + xMean = [float(mean) for mean in cast(Iterable[float], xMean)] + xVariance = [float(var) for var in cast(Iterable[float], xVariance)] + return list( + callMLlibFunc( + "generateLinearInputWrapper", + float(intercept), + weights, + xMean, + xVariance, + int(nPoints), + int(seed), + float(eps), + ) + ) + + @staticmethod + @since("1.5.0") + def generateLinearRDD( + sc: SparkContext, + nexamples: int, + nfeatures: int, + eps: float, + nParts: int = 2, + intercept: float = 0.0, + ) -> RDD["LabeledPoint"]: + """ + Generate an RDD of LabeledPoints. + """ + return callMLlibFunc( + "generateLinearRDDWrapper", + sc, + int(nexamples), + int(nfeatures), + float(eps), + int(nParts), + float(intercept), + ) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + + globs = globals().copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + spark = SparkSession.builder.master("local[2]").appName("mllib.util tests").getOrCreate() + globs["spark"] = spark + globs["sc"] = spark.sparkContext + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/__init__.py new file mode 100644 index 0000000..980aeab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/__init__.py @@ -0,0 +1,165 @@ +# +# 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. +# + +""" +.. versionadded:: 3.2.0 + pandas API on Spark +""" + +import os +import sys +import warnings +from distutils.version import LooseVersion +from typing import Any + +from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions +from pyspark.pandas.missing.scalars import MissingPandasLikeScalars +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version + +try: + require_minimum_pandas_version() + require_minimum_pyarrow_version() +except ImportError as e: + if os.environ.get("SPARK_TESTING"): + warnings.warn(str(e)) + sys.exit() + else: + raise + + +import pyarrow + +if ( + LooseVersion(pyarrow.__version__) >= LooseVersion("2.0.0") + and "PYARROW_IGNORE_TIMEZONE" not in os.environ +): + warnings.warn( + "'PYARROW_IGNORE_TIMEZONE' environment variable was not set. It is required to " + "set this environment variable to '1' in both driver and executor sides if you use " + "pyarrow>=2.0.0. " + "pandas-on-Spark will set it for you but it does not work if there is a Spark context " + "already launched." + ) + os.environ["PYARROW_IGNORE_TIMEZONE"] = "1" + +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.indexes.category import CategoricalIndex +from pyspark.pandas.indexes.datetimes import DatetimeIndex +from pyspark.pandas.indexes.multi import MultiIndex +from pyspark.pandas.indexes.numeric import Float64Index, Int64Index +from pyspark.pandas.indexes.timedelta import TimedeltaIndex +from pyspark.pandas.series import Series +from pyspark.pandas.groupby import NamedAgg + +__all__ = [ # noqa: F405 + "read_csv", + "read_parquet", + "to_datetime", + "date_range", + "from_pandas", + "get_dummies", + "DataFrame", + "Series", + "Index", + "MultiIndex", + "Int64Index", + "Float64Index", + "CategoricalIndex", + "DatetimeIndex", + "TimedeltaIndex", + "sql", + "range", + "concat", + "melt", + "get_option", + "set_option", + "reset_option", + "read_sql_table", + "read_sql_query", + "read_sql", + "options", + "option_context", + "NamedAgg", +] + + +def _auto_patch_spark() -> None: + import os + import logging + + # Attach a usage logger. 'KOALAS_USAGE_LOGGER' is legacy, and it's for compatibility. + logger_module = os.getenv("PYSPARK_PANDAS_USAGE_LOGGER", os.getenv("KOALAS_USAGE_LOGGER", "")) + if logger_module != "": + try: + from pyspark.pandas import usage_logging + + usage_logging.attach(logger_module) + except Exception as e: + logger = logging.getLogger("pyspark.pandas.usage_logger") + logger.warning( + "Tried to attach usage logger `{}`, but an exception was raised: {}".format( + logger_module, str(e) + ) + ) + + +_frame_has_class_getitem = False +_series_has_class_getitem = False + + +def _auto_patch_pandas() -> None: + import pandas as pd + + # In order to use it in test cases. + global _frame_has_class_getitem + global _series_has_class_getitem + + _frame_has_class_getitem = hasattr(pd.DataFrame, "__class_getitem__") + _series_has_class_getitem = hasattr(pd.Series, "__class_getitem__") + + if sys.version_info >= (3, 7): + # Just in case pandas implements '__class_getitem__' later. + if not _frame_has_class_getitem: + pd.DataFrame.__class_getitem__ = ( # type: ignore[attr-defined] + lambda params: DataFrame.__class_getitem__(params) + ) + + if not _series_has_class_getitem: + pd.Series.__class_getitem__ = ( # type: ignore[attr-defined] + lambda params: Series.__class_getitem__(params) + ) + + +_auto_patch_spark() +_auto_patch_pandas() + +# Import after the usage logger is attached. +from pyspark.pandas.config import get_option, options, option_context, reset_option, set_option +from pyspark.pandas.namespace import * # noqa: F403 +from pyspark.pandas.sql_formatter import sql + + +def __getattr__(key: str) -> Any: + if key.startswith("__"): + raise AttributeError(key) + if hasattr(MissingPandasLikeScalars, key): + raise getattr(MissingPandasLikeScalars, key) + if hasattr(MissingPandasLikeGeneralFunctions, key): + return getattr(MissingPandasLikeGeneralFunctions, key) + else: + raise AttributeError("module 'pyspark.pandas' has no attribute '%s'" % (key)) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/_typing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/_typing.py new file mode 100644 index 0000000..0e5b86a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/_typing.py @@ -0,0 +1,52 @@ +# +# 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 datetime +import decimal +from typing import Any, Tuple, TypeVar, Union, TYPE_CHECKING + +import numpy as np +from pandas.api.extensions import ExtensionDtype + + +if TYPE_CHECKING: + from pyspark.pandas.base import IndexOpsMixin + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.generic import Frame + from pyspark.pandas.indexes.base import Index + from pyspark.pandas.series import Series + + +# TypeVars +T = TypeVar("T") + +FrameLike = TypeVar("FrameLike", bound="Frame") +IndexOpsLike = TypeVar("IndexOpsLike", bound="IndexOpsMixin") + +# Type aliases +Scalar = Union[ + int, float, bool, str, bytes, decimal.Decimal, datetime.date, datetime.datetime, None +] + +# TODO: use the actual type parameters. +Label = Tuple[Any, ...] +Name = Union[Any, Label] + +Axis = Union[int, str] +Dtype = Union[np.dtype, ExtensionDtype] + +DataFrameOrSeries = Union["DataFrame", "Series"] +SeriesOrIndex = Union["Series", "Index"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/accessors.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/accessors.py new file mode 100644 index 0000000..4e96f4d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/accessors.py @@ -0,0 +1,989 @@ +# +# 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. +# +""" +pandas-on-Spark specific features. +""" +import inspect +from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING, cast, List +from types import FunctionType + +import numpy as np # noqa: F401 +import pandas as pd + +from pyspark.sql import functions as F +from pyspark.sql.functions import pandas_udf +from pyspark.sql.types import DataType, LongType, StructField, StructType + +from pyspark.pandas._typing import DataFrameOrSeries, Name +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + SPARK_INDEX_NAME_FORMAT, + SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_PATTERN, +) +from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType +from pyspark.pandas.utils import ( + is_name_like_value, + is_name_like_tuple, + name_like_string, + scol_for, + verify_temp_column_name, + log_advice, +) + +if TYPE_CHECKING: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series + from pyspark.sql._typing import UserDefinedFunctionLike + + +class PandasOnSparkFrameMethods: + """pandas-on-Spark specific features for DataFrame.""" + + def __init__(self, frame: "DataFrame"): + self._psdf = frame + + def attach_id_column(self, id_type: str, column: Name) -> "DataFrame": + """ + Attach a column to be used as an identifier of rows similar to the default index. + + See also `Default Index type + `_. + + Parameters + ---------- + id_type : string + The id type. + + - 'sequence' : a sequence that increases one by one. + + .. note:: this uses Spark's Window without specifying partition specification. + This leads to moving all data into a single partition in a single machine and + could cause serious performance degradation. + Avoid this method with very large datasets. + + - 'distributed-sequence' : a sequence that increases one by one, + by group-by and group-map approach in a distributed manner. + - 'distributed' : a monotonically increasing sequence simply by using PySpark’s + monotonically_increasing_id function in a fully distributed manner. + + column : string or tuple of string + The column name. + + Returns + ------- + DataFrame + The DataFrame attached the column. + + Examples + -------- + >>> df = ps.DataFrame({"x": ['a', 'b', 'c']}) + >>> df.pandas_on_spark.attach_id_column(id_type="sequence", column="id") + x id + 0 a 0 + 1 b 1 + 2 c 2 + + >>> df.pandas_on_spark.attach_id_column(id_type="distributed-sequence", column=0) + x 0 + 0 a 0 + 1 b 1 + 2 c 2 + + >>> df.pandas_on_spark.attach_id_column(id_type="distributed", column=0.0) + ... # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE + x 0.0 + 0 a ... + 1 b ... + 2 c ... + + For multi-index columns: + + >>> df = ps.DataFrame({("x", "y"): ['a', 'b', 'c']}) + >>> df.pandas_on_spark.attach_id_column(id_type="sequence", column=("id-x", "id-y")) + x id-x + y id-y + 0 a 0 + 1 b 1 + 2 c 2 + + >>> df.pandas_on_spark.attach_id_column(id_type="distributed-sequence", column=(0, 1.0)) + x 0 + y 1.0 + 0 a 0 + 1 b 1 + 2 c 2 + """ + from pyspark.pandas.frame import DataFrame + + if id_type == "sequence": + attach_func = InternalFrame.attach_sequence_column + elif id_type == "distributed-sequence": + attach_func = InternalFrame.attach_distributed_sequence_column + elif id_type == "distributed": + attach_func = InternalFrame.attach_distributed_column + else: + raise ValueError( + "id_type should be one of 'sequence', 'distributed-sequence' and 'distributed'" + ) + + assert is_name_like_value(column, allow_none=False), column + if not is_name_like_tuple(column): + column = (column,) + + internal = self._psdf._internal + + if len(column) != internal.column_labels_level: + raise ValueError( + "The given column `{}` must be the same length as the existing columns.".format( + column + ) + ) + elif column in internal.column_labels: + raise ValueError( + "The given column `{}` already exists.".format(name_like_string(column)) + ) + + # Make sure the underlying Spark column names are the form of + # `name_like_string(column_label)`. + sdf = internal.spark_frame.select( + [ + scol.alias(SPARK_INDEX_NAME_FORMAT(i)) + for i, scol in enumerate(internal.index_spark_columns) + ] + + [ + scol.alias(name_like_string(label)) + for scol, label in zip(internal.data_spark_columns, internal.column_labels) + ] + ) + sdf = attach_func(sdf, name_like_string(column)) + + return DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, SPARK_INDEX_NAME_FORMAT(i)) for i in range(internal.index_level) + ], + index_names=internal.index_names, + index_fields=internal.index_fields, + column_labels=internal.column_labels + [column], + data_spark_columns=( + [scol_for(sdf, name_like_string(label)) for label in internal.column_labels] + + [scol_for(sdf, name_like_string(column))] + ), + data_fields=internal.data_fields + + [ + InternalField.from_struct_field( + StructField(name_like_string(column), LongType(), nullable=False) + ) + ], + column_label_names=internal.column_label_names, + ).resolved_copy + ) + + def apply_batch( + self, func: Callable[..., pd.DataFrame], args: Tuple = (), **kwds: Any + ) -> "DataFrame": + """ + Apply a function that takes pandas DataFrame and outputs pandas DataFrame. The pandas + DataFrame given to the function is of a batch used internally. + + See also `Transform and apply a function + `_. + + .. note:: the `func` is unable to access the whole input frame. pandas-on-Spark + internally splits the input series into multiple batches and calls `func` with each + batch multiple times. Therefore, operations such as global aggregations are impossible. + See the example below. + + >>> # This case does not return the length of whole frame but of the batch internally + ... # used. + ... def length(pdf) -> ps.DataFrame[int, [int]]: + ... return pd.DataFrame([len(pdf)]) + ... + >>> df = ps.DataFrame({'A': range(1000)}) + >>> df.pandas_on_spark.apply_batch(length) # doctest: +SKIP + c0 + 0 83 + 1 83 + 2 83 + ... + 10 83 + 11 83 + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]: + ... return x + 1 + + If the return type is specified, the output column names become + `c0, c1, c2 ... cn`. These names are positionally mapped to the returned + DataFrame in ``func``. + + To specify the column names, you can assign them in a NumPy compound type style + as below: + + >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]: + ... return x + 1 + + >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]}) + >>> def plus_one(x) -> ps.DataFrame[ + ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]: + ... return x + 1 + + Parameters + ---------- + func : function + Function to apply to each pandas frame. + args : tuple + Positional arguments to pass to `func` in addition to the + array/series. + **kwds + Additional keyword arguments to pass as keywords arguments to + `func`. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.apply: For row/columnwise operations. + DataFrame.applymap: For elementwise operations. + DataFrame.aggregate: Only perform aggregating type operations. + DataFrame.transform: Only perform transforming type operations. + Series.pandas_on_spark.transform_batch: transform the search as each pandas chunks. + + Examples + -------- + >>> df = ps.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B']) + >>> df + A B + 0 1 2 + 1 3 4 + 2 5 6 + + >>> def query_func(pdf) -> ps.DataFrame[int, [int, int]]: + ... return pdf.query('A == 1') + >>> df.pandas_on_spark.apply_batch(query_func) + c0 c1 + 0 1 2 + + >>> def query_func(pdf) -> ps.DataFrame[("idx", int), [("A", int), ("B", int)]]: + ... return pdf.query('A == 1') + >>> df.pandas_on_spark.apply_batch(query_func) # doctest: +NORMALIZE_WHITESPACE + A B + idx + 0 1 2 + + You can also omit the type hints so pandas-on-Spark infers the return schema as below: + + >>> df.pandas_on_spark.apply_batch(lambda pdf: pdf.query('A == 1')) + A B + 0 1 2 + + You can also specify extra arguments. + + >>> def calculation(pdf, y, z) -> ps.DataFrame[int, [int, int]]: + ... return pdf ** y + z + >>> df.pandas_on_spark.apply_batch(calculation, args=(10,), z=20) + c0 c1 + 0 21 1044 + 1 59069 1048596 + 2 9765645 60466196 + + You can also use ``np.ufunc`` and built-in functions as input. + + >>> df.pandas_on_spark.apply_batch(np.add, args=(10,)) + A B + 0 11 12 + 1 13 14 + 2 15 16 + + >>> (df * -1).pandas_on_spark.apply_batch(abs) + A B + 0 1 2 + 1 3 4 + 2 5 6 + + """ + # TODO: codes here partially duplicate `DataFrame.apply`. Can we deduplicate? + + from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.frame import DataFrame + from pyspark import pandas as ps + + if not isinstance(func, FunctionType): + assert callable(func), "the first argument should be a callable function." + f = func + # Note that the return type hint specified here affects actual return + # type in Spark (e.g., infer_return_type). And, MyPy does not allow + # redefinition of a function. + func = lambda *args, **kwargs: f(*args, **kwargs) # noqa: E731 + + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + + original_func = func + + def new_func(o: Any) -> pd.DataFrame: + return original_func(o, *args, **kwds) + + self_applied: DataFrame = DataFrame(self._psdf._internal.resolved_copy) + + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + log_advice( + "If the type hints is not specified for `apply_batch`, " + "it is expensive to infer the data type internally." + ) + limit = ps.get_option("compute.shortcut_limit") + pdf = self_applied.head(limit + 1)._to_internal_pandas() + applied = new_func(pdf) + if not isinstance(applied, pd.DataFrame): + raise ValueError( + "The given function should return a frame; however, " + "the return type was %s." % type(applied) + ) + psdf: DataFrame = DataFrame(applied) + if len(pdf) <= limit: + return psdf + + index_fields = [field.normalize_spark_type() for field in psdf._internal.index_fields] + data_fields = [field.normalize_spark_type() for field in psdf._internal.data_fields] + + return_schema = StructType([field.struct_field for field in index_fields + data_fields]) + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, new_func, return_schema, retain_index=True + ) + sdf = self_applied._internal.spark_frame.mapInPandas( + lambda iterator: map(output_func, iterator), schema=return_schema + ) + + # If schema is inferred, we can restore indexes too. + internal = psdf._internal.with_new_sdf( + spark_frame=sdf, index_fields=index_fields, data_fields=data_fields + ) + else: + return_type = infer_return_type(original_func) + is_return_dataframe = isinstance(return_type, DataFrameType) + if not is_return_dataframe: + raise TypeError( + "The given function should specify a frame as its type " + "hints; however, the return type was %s." % return_sig + ) + index_fields = cast(DataFrameType, return_type).index_fields + should_retain_index = len(index_fields) > 0 + return_schema = cast(DataFrameType, return_type).spark_type + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, new_func, return_schema, retain_index=should_retain_index + ) + sdf = self_applied._internal.to_internal_spark_frame.mapInPandas( + lambda iterator: map(output_func, iterator), schema=return_schema + ) + + index_spark_columns = None + index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None + + if should_retain_index: + index_spark_columns = [ + scol_for(sdf, index_field.struct_field.name) for index_field in index_fields + ] + + if not any( + [ + SPARK_INDEX_NAME_PATTERN.match(index_field.struct_field.name) + for index_field in index_fields + ] + ): + index_names = [(index_field.struct_field.name,) for index_field in index_fields] + internal = InternalFrame( + spark_frame=sdf, + index_names=index_names, + index_spark_columns=index_spark_columns, + index_fields=index_fields, + data_fields=cast(DataFrameType, return_type).data_fields, + ) + return DataFrame(internal) + + def transform_batch( + self, func: Callable[..., Union[pd.DataFrame, pd.Series]], *args: Any, **kwargs: Any + ) -> DataFrameOrSeries: + """ + Transform chunks with a function that takes pandas DataFrame and outputs pandas DataFrame. + The pandas DataFrame given to the function is of a batch used internally. The length of + each input and output should be the same. + + See also `Transform and apply a function + `_. + + .. note:: the `func` is unable to access the whole input frame. pandas-on-Spark + internally splits the input series into multiple batches and calls `func` with each + batch multiple times. Therefore, operations such as global aggregations are impossible. + See the example below. + + >>> # This case does not return the length of whole frame but of the batch internally + ... # used. + ... def length(pdf) -> ps.DataFrame[int]: + ... return pd.DataFrame([len(pdf)] * len(pdf)) + ... + >>> df = ps.DataFrame({'A': range(1000)}) + >>> df.pandas_on_spark.transform_batch(length) # doctest: +SKIP + c0 + 0 83 + 1 83 + 2 83 + ... + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]: + ... return x + 1 + + If the return type is specified, the output column names become + `c0, c1, c2 ... cn`. These names are positionally mapped to the returned + DataFrame in ``func``. + + To specify the column names, you can assign them in a NumPy compound type style + as below: + + >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]: + ... return x + 1 + + >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]}) + >>> def plus_one(x) -> ps.DataFrame[ + ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]: + ... return x + 1 + + Parameters + ---------- + func : function + Function to transform each pandas frame. + *args + Positional arguments to pass to func. + **kwargs + Keyword arguments to pass to func. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.pandas_on_spark.apply_batch: For row/columnwise operations. + Series.pandas_on_spark.transform_batch: transform the search as each pandas chunks. + + Examples + -------- + >>> df = ps.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B']) + >>> df + A B + 0 1 2 + 1 3 4 + 2 5 6 + + >>> def plus_one_func(pdf) -> ps.DataFrame[int, [int, int]]: + ... return pdf + 1 + >>> df.pandas_on_spark.transform_batch(plus_one_func) + c0 c1 + 0 2 3 + 1 4 5 + 2 6 7 + + >>> def plus_one_func(pdf) -> ps.DataFrame[("index", int), [('A', int), ('B', int)]]: + ... return pdf + 1 + >>> df.pandas_on_spark.transform_batch(plus_one_func) # doctest: +NORMALIZE_WHITESPACE + A B + index + 0 2 3 + 1 4 5 + 2 6 7 + + >>> def plus_one_func(pdf) -> ps.Series[int]: + ... return pdf.B + 1 + >>> df.pandas_on_spark.transform_batch(plus_one_func) + 0 3 + 1 5 + 2 7 + dtype: int64 + + You can also omit the type hints so pandas-on-Spark infers the return schema as below: + + >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf + 1) + A B + 0 2 3 + 1 4 5 + 2 6 7 + + >>> (df * -1).pandas_on_spark.transform_batch(abs) + A B + 0 1 2 + 1 3 4 + 2 5 6 + + Note that you should not transform the index. The index information will not change. + + >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf.B + 1) + 0 3 + 1 5 + 2 7 + Name: B, dtype: int64 + + You can also specify extra arguments as below. + + >>> df.pandas_on_spark.transform_batch(lambda pdf, a, b, c: pdf.B + a + b + c, 1, 2, c=3) + 0 8 + 1 10 + 2 12 + Name: B, dtype: int64 + """ + from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import first_series + from pyspark import pandas as ps + + assert callable(func), "the first argument should be a callable function." + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + should_retain_index = should_infer_schema + original_func = func + + def new_func(o: Any) -> Union[pd.DataFrame, pd.Series]: + return original_func(o, *args, **kwargs) + + def apply_func(pdf: pd.DataFrame) -> pd.DataFrame: + return new_func(pdf).to_frame() + + def pandas_series_func( + f: Callable[[pd.DataFrame], pd.DataFrame], return_type: DataType + ) -> "UserDefinedFunctionLike": + ff = f + + @pandas_udf(returnType=return_type) # type: ignore[call-overload] + def udf(pdf: pd.DataFrame) -> pd.Series: + return first_series(ff(pdf)) + + return udf + + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + log_advice( + "If the type hints is not specified for `transform_batch`, " + "it is expensive to infer the data type internally." + ) + limit = ps.get_option("compute.shortcut_limit") + pdf = self._psdf.head(limit + 1)._to_internal_pandas() + transformed = new_func(pdf) + if not isinstance(transformed, (pd.DataFrame, pd.Series)): + raise ValueError( + "The given function should return a frame; however, " + "the return type was %s." % type(transformed) + ) + if len(transformed) != len(pdf): + raise ValueError("transform_batch cannot produce aggregated results") + psdf_or_psser = ps.from_pandas(transformed) + + if isinstance(psdf_or_psser, ps.Series): + psser = psdf_or_psser + + field = psser._internal.data_fields[0].normalize_spark_type() + + return_schema = StructType([field.struct_field]) + output_func = GroupBy._make_pandas_df_builder_func( + self._psdf, apply_func, return_schema, retain_index=False + ) + + pudf = pandas_series_func(output_func, return_type=field.spark_type) + columns = self._psdf._internal.spark_columns + # TODO: Index will be lost in this case. + internal = self._psdf._internal.copy( + column_labels=psser._internal.column_labels, + data_spark_columns=[pudf(F.struct(*columns)).alias(field.name)], + data_fields=[field], + column_label_names=psser._internal.column_label_names, + ) + return first_series(DataFrame(internal)) + else: + psdf = cast(DataFrame, psdf_or_psser) + if len(pdf) <= limit: + # only do the short cut when it returns a frame to avoid + # operations on different dataframes in case of series. + return psdf + + index_fields = [ + field.normalize_spark_type() for field in psdf._internal.index_fields + ] + data_fields = [field.normalize_spark_type() for field in psdf._internal.data_fields] + + return_schema = StructType( + [field.struct_field for field in index_fields + data_fields] + ) + + self_applied: DataFrame = DataFrame(self._psdf._internal.resolved_copy) + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, + new_func, # type: ignore[arg-type] + return_schema, + retain_index=True, + ) + columns = self_applied._internal.spark_columns + + pudf = pandas_udf( # type: ignore[call-overload] + output_func, returnType=return_schema + ) + temp_struct_column = verify_temp_column_name( + self_applied._internal.spark_frame, "__temp_struct__" + ) + applied = pudf(F.struct(*columns)).alias(temp_struct_column) + sdf = self_applied._internal.spark_frame.select(applied) + sdf = sdf.selectExpr("%s.*" % temp_struct_column) + + return DataFrame( + psdf._internal.with_new_sdf( + spark_frame=sdf, index_fields=index_fields, data_fields=data_fields + ) + ) + else: + return_type = infer_return_type(original_func) + is_return_series = isinstance(return_type, SeriesType) + is_return_dataframe = isinstance(return_type, DataFrameType) + if not is_return_dataframe and not is_return_series: + raise TypeError( + "The given function should specify a frame or series as its type " + "hints; however, the return type was %s." % return_sig + ) + if is_return_series: + field = InternalField( + dtype=cast(SeriesType, return_type).dtype, + struct_field=StructField( + name=SPARK_DEFAULT_SERIES_NAME, + dataType=cast(SeriesType, return_type).spark_type, + ), + ).normalize_spark_type() + + return_schema = StructType([field.struct_field]) + output_func = GroupBy._make_pandas_df_builder_func( + self._psdf, apply_func, return_schema, retain_index=False + ) + + pudf = pandas_series_func(output_func, return_type=field.spark_type) + columns = self._psdf._internal.spark_columns + internal = self._psdf._internal.copy( + column_labels=[None], + data_spark_columns=[pudf(F.struct(*columns)).alias(field.name)], + data_fields=[field], + column_label_names=None, + ) + return first_series(DataFrame(internal)) + else: + index_fields = cast(DataFrameType, return_type).index_fields + index_fields = [index_field.normalize_spark_type() for index_field in index_fields] + data_fields = [ + field.normalize_spark_type() + for field in cast(DataFrameType, return_type).data_fields + ] + normalized_fields = index_fields + data_fields + return_schema = StructType([field.struct_field for field in normalized_fields]) + should_retain_index = len(index_fields) > 0 + + self_applied = DataFrame(self._psdf._internal.resolved_copy) + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, + new_func, # type: ignore[arg-type] + return_schema, + retain_index=should_retain_index, + ) + columns = self_applied._internal.spark_columns + + pudf = pandas_udf( # type: ignore[call-overload] + output_func, returnType=return_schema + ) + temp_struct_column = verify_temp_column_name( + self_applied._internal.spark_frame, "__temp_struct__" + ) + applied = pudf(F.struct(*columns)).alias(temp_struct_column) + sdf = self_applied._internal.spark_frame.select(applied) + sdf = sdf.selectExpr("%s.*" % temp_struct_column) + + index_spark_columns = None + index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None + + if should_retain_index: + index_spark_columns = [ + scol_for(sdf, index_field.struct_field.name) for index_field in index_fields + ] + + if not any( + [ + SPARK_INDEX_NAME_PATTERN.match(index_field.struct_field.name) + for index_field in index_fields + ] + ): + index_names = [ + (index_field.struct_field.name,) for index_field in index_fields + ] + internal = InternalFrame( + spark_frame=sdf, + index_names=index_names, + index_spark_columns=index_spark_columns, + index_fields=index_fields, + data_fields=data_fields, + ) + return DataFrame(internal) + + +class PandasOnSparkSeriesMethods: + """pandas-on-Spark specific features for Series.""" + + def __init__(self, series: "Series"): + self._psser = series + + def transform_batch( + self, func: Callable[..., pd.Series], *args: Any, **kwargs: Any + ) -> "Series": + """ + Transform the data with the function that takes pandas Series and outputs pandas Series. + The pandas Series given to the function is of a batch used internally. + + See also `Transform and apply a function + `_. + + .. note:: the `func` is unable to access the whole input series. pandas-on-Spark + internally splits the input series into multiple batches and calls `func` with each + batch multiple times. Therefore, operations such as global aggregations are impossible. + See the example below. + + >>> # This case does not return the length of whole frame but of the batch internally + ... # used. + ... def length(pser) -> ps.Series[int]: + ... return pd.Series([len(pser)] * len(pser)) + ... + >>> df = ps.DataFrame({'A': range(1000)}) + >>> df.A.pandas_on_spark.transform_batch(length) # doctest: +SKIP + c0 + 0 83 + 1 83 + 2 83 + ... + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def plus_one(x) -> ps.Series[int]: + ... return x + 1 + + Parameters + ---------- + func : function + Function to apply to each pandas frame. + *args + Positional arguments to pass to func. + **kwargs + Keyword arguments to pass to func. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.pandas_on_spark.apply_batch : Similar but it takes pandas DataFrame as its + internal batch. + + Examples + -------- + >>> df = ps.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B']) + >>> df + A B + 0 1 2 + 1 3 4 + 2 5 6 + + >>> def plus_one_func(pser) -> ps.Series[np.int64]: + ... return pser + 1 + >>> df.A.pandas_on_spark.transform_batch(plus_one_func) + 0 2 + 1 4 + 2 6 + Name: A, dtype: int64 + + You can also omit the type hints so pandas-on-Spark infers the return schema as below: + + >>> df.A.pandas_on_spark.transform_batch(lambda pser: pser + 1) + 0 2 + 1 4 + 2 6 + Name: A, dtype: int64 + + You can also specify extra arguments. + + >>> def plus_one_func(pser, a, b, c=3) -> ps.Series[np.int64]: + ... return pser + a + b + c + >>> df.A.pandas_on_spark.transform_batch(plus_one_func, 1, b=2) + 0 7 + 1 9 + 2 11 + Name: A, dtype: int64 + + You can also use ``np.ufunc`` and built-in functions as input. + + >>> df.A.pandas_on_spark.transform_batch(np.add, 10) + 0 11 + 1 13 + 2 15 + Name: A, dtype: int64 + + >>> (df * -1).A.pandas_on_spark.transform_batch(abs) + 0 1 + 1 3 + 2 5 + Name: A, dtype: int64 + """ + assert callable(func), "the first argument should be a callable function." + + return_sig = None + try: + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + except TypeError: + # Falls back to schema inference if it fails to get signature. + pass + + return_type = None + if return_sig is not None: + # Extract the signature arguments from this function. + sig_return = infer_return_type(func) + if not isinstance(sig_return, SeriesType): + raise ValueError( + "Expected the return type of this function to be of type column," + " but found type {}".format(sig_return) + ) + return_type = sig_return + + return self._transform_batch(lambda c: func(c, *args, **kwargs), return_type) + + def _transform_batch( + self, func: Callable[..., pd.Series], return_type: Optional[Union[SeriesType, ScalarType]] + ) -> "Series": + from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.series import Series, first_series + from pyspark import pandas as ps + + if not isinstance(func, FunctionType): + f = func + # Note that the return type hint specified here affects actual return + # type in Spark (e.g., infer_return_type). And, MyPy does not allow + # redefinition of a function. + func = lambda *args, **kwargs: f(*args, **kwargs) # noqa: E731 + + if return_type is None: + # TODO: In this case, it avoids the shortcut for now (but only infers schema) + # because it returns a series from a different DataFrame and it has a different + # anchor. We should fix this to allow the shortcut or only allow to infer + # schema. + limit = ps.get_option("compute.shortcut_limit") + pser = self._psser.head(limit + 1)._to_internal_pandas() + transformed = pser.transform(func) + psser: Series = Series(transformed) + + field = psser._internal.data_fields[0].normalize_spark_type() + else: + spark_return_type = return_type.spark_type + dtype = return_type.dtype + field = InternalField( + dtype=dtype, + struct_field=StructField( + name=self._psser._internal.data_spark_column_names[0], + dataType=spark_return_type, + ), + ) + + psdf = self._psser.to_frame() + columns = psdf._internal.spark_column_names + + def pandas_concat(*series: pd.Series) -> pd.DataFrame: + # The input can only be a DataFrame for struct from Spark 3.0. + # This works around makeing the input as a frame. See SPARK-27240 + pdf = pd.concat(series, axis=1) + pdf.columns = columns + return pdf + + def apply_func(pdf: pd.DataFrame) -> pd.DataFrame: + return func(first_series(pdf)).to_frame() + + return_schema = StructType([StructField(SPARK_DEFAULT_SERIES_NAME, field.spark_type)]) + output_func = GroupBy._make_pandas_df_builder_func( + psdf, apply_func, return_schema, retain_index=False + ) + + @pandas_udf(returnType=field.spark_type) # type: ignore[call-overload] + def pudf(*series: pd.Series) -> pd.Series: + return first_series(output_func(pandas_concat(*series))) + + return self._psser._with_new_scol( + scol=pudf(*psdf._internal.spark_columns).alias(field.name), field=field + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.accessors + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.accessors.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.accessors tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.accessors, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/base.py new file mode 100644 index 0000000..e005fd1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/base.py @@ -0,0 +1,1804 @@ +# +# 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. +# + +""" +Base and utility classes for pandas-on-Spark objects. +""" +import warnings +from abc import ABCMeta, abstractmethod +from functools import wraps, partial +from itertools import chain +from typing import Any, Callable, Optional, Sequence, Tuple, Union, cast, TYPE_CHECKING + +import numpy as np +import pandas as pd +from pandas.api.types import is_list_like, CategoricalDtype # type: ignore[attr-defined] +from pyspark.sql import functions as F, Column, Window +from pyspark.sql.types import LongType, BooleanType, NumericType + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import Axis, Dtype, IndexOpsLike, Label, SeriesOrIndex +from pyspark.pandas.config import get_option, option_context +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + NATURAL_ORDER_COLUMN_NAME, + SPARK_DEFAULT_INDEX_NAME, +) +from pyspark.pandas.spark.accessors import SparkIndexOpsMethods +from pyspark.pandas.typedef import extension_dtypes +from pyspark.pandas.utils import ( + combine_frames, + same_anchor, + scol_for, + validate_axis, + ERROR_MESSAGE_CANNOT_COMBINE, +) +from pyspark.pandas.frame import DataFrame + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName + + from pyspark.pandas.data_type_ops.base import DataTypeOps + from pyspark.pandas.series import Series + + +def should_alignment_for_column_op(self: SeriesOrIndex, other: SeriesOrIndex) -> bool: + from pyspark.pandas.series import Series + + if isinstance(self, Series) and isinstance(other, Series): + return not same_anchor(self, other) + else: + return self._internal.spark_frame is not other._internal.spark_frame + + +def align_diff_index_ops( + func: Callable[..., Column], this_index_ops: SeriesOrIndex, *args: Any +) -> SeriesOrIndex: + """ + Align the `IndexOpsMixin` objects and apply the function. + + Parameters + ---------- + func : The function to apply + this_index_ops : IndexOpsMixin + A base `IndexOpsMixin` object + args : list of other arguments including other `IndexOpsMixin` objects + + Returns + ------- + `Index` if all `this_index_ops` and arguments are `Index`; otherwise `Series` + """ + from pyspark.pandas.indexes import Index + from pyspark.pandas.series import Series, first_series + + cols = [arg for arg in args if isinstance(arg, IndexOpsMixin)] + + if isinstance(this_index_ops, Series) and all(isinstance(col, Series) for col in cols): + combined = combine_frames( + this_index_ops.to_frame(), + *[cast(Series, col).rename(i) for i, col in enumerate(cols)], + how="full", + ) + + return column_op(func)( + combined["this"]._psser_for(combined["this"]._internal.column_labels[0]), + *[ + combined["that"]._psser_for(label) + for label in combined["that"]._internal.column_labels + ], + ).rename(this_index_ops.name) + else: + # This could cause as many counts, reset_index calls, joins for combining + # as the number of `Index`s in `args`. So far it's fine since we can assume the ops + # only work between at most two `Index`s. We might need to fix it in the future. + + self_len = len(this_index_ops) + if any(len(col) != self_len for col in args if isinstance(col, IndexOpsMixin)): + raise ValueError("operands could not be broadcast together with shapes") + + with option_context("compute.default_index_type", "distributed-sequence"): + if isinstance(this_index_ops, Index) and all(isinstance(col, Index) for col in cols): + return Index( + column_op(func)( + this_index_ops.to_series().reset_index(drop=True), + *[ + arg.to_series().reset_index(drop=True) + if isinstance(arg, Index) + else arg + for arg in args + ], + ).sort_index(), + name=this_index_ops.name, + ) + elif isinstance(this_index_ops, Series): + this = cast(DataFrame, this_index_ops.reset_index()) + that = [ + cast(Series, col.to_series() if isinstance(col, Index) else col) + .rename(i) + .reset_index(drop=True) + for i, col in enumerate(cols) + ] + + combined = combine_frames(this, *that, how="full").sort_index() + combined = combined.set_index( + combined._internal.column_labels[: this_index_ops._internal.index_level] + ) + combined.index.names = this_index_ops._internal.index_names + + return column_op(func)( + first_series(combined["this"]), + *[ + combined["that"]._psser_for(label) + for label in combined["that"]._internal.column_labels + ], + ).rename(this_index_ops.name) + else: + this = this_index_ops.to_frame().reset_index(drop=True) + + that_series = next(col for col in cols if isinstance(col, Series)) + that_frame = that_series._psdf[ + [ + cast(Series, col.to_series() if isinstance(col, Index) else col).rename(i) + for i, col in enumerate(cols) + ] + ] + + combined = combine_frames(this, that_frame.reset_index()).sort_index() + + self_index = ( + combined["this"].set_index(combined["this"]._internal.column_labels).index + ) + + other = combined["that"].set_index( + combined["that"]._internal.column_labels[: that_series._internal.index_level] + ) + other.index.names = that_series._internal.index_names + + return column_op(func)( + self_index, + *[ + other._psser_for(label) + for label, col in zip(other._internal.column_labels, cols) + ], + ).rename(that_series.name) + + +def booleanize_null(scol: Column, f: Callable[..., Column]) -> Column: + """ + Booleanize Null in Spark Column + """ + comp_ops = [ + getattr(Column, "__{}__".format(comp_op)) + for comp_op in ["eq", "ne", "lt", "le", "ge", "gt"] + ] + + if f in comp_ops: + # if `f` is "!=", fill null with True otherwise False + filler = f == Column.__ne__ + scol = F.when(scol.isNull(), filler).otherwise(scol) + + return scol + + +def column_op(f: Callable[..., Column]) -> Callable[..., SeriesOrIndex]: + """ + A decorator that wraps APIs taking/returning Spark Column so that pandas-on-Spark Series can be + supported too. If this decorator is used for the `f` function that takes Spark Column and + returns Spark Column, decorated `f` takes pandas-on-Spark Series as well and returns + pandas-on-Spark Series. + + :param f: a function that takes Spark Column and returns Spark Column. + :param self: pandas-on-Spark Series + :param args: arguments that the function `f` takes. + """ + + @wraps(f) + def wrapper(self: SeriesOrIndex, *args: Any) -> SeriesOrIndex: + from pyspark.pandas.indexes.base import Index + from pyspark.pandas.series import Series + + # It is possible for the function `f` to take other arguments than Spark Column. + # To cover this case, explicitly check if the argument is pandas-on-Spark Series and + # extract Spark Column. For other arguments, they are used as are. + cols = [arg for arg in args if isinstance(arg, (Series, Index))] + + if all(not should_alignment_for_column_op(self, col) for col in cols): + # Same DataFrame anchors + scol = f( + self.spark.column, + *[arg.spark.column if isinstance(arg, IndexOpsMixin) else arg for arg in args], + ) + + field = InternalField.from_struct_field( + self._internal.spark_frame.select(scol).schema[0], + use_extension_dtypes=any( + isinstance(col.dtype, extension_dtypes) for col in [self] + cols + ), + ) + + if not field.is_extension_dtype: + scol = booleanize_null(scol, f).alias(field.name) + + if isinstance(self, Series) or not any(isinstance(col, Series) for col in cols): + index_ops = self._with_new_scol(scol, field=field) + else: + psser = next(col for col in cols if isinstance(col, Series)) + index_ops = psser._with_new_scol(scol, field=field) + elif get_option("compute.ops_on_diff_frames"): + index_ops = align_diff_index_ops(f, self, *args) + else: + raise ValueError(ERROR_MESSAGE_CANNOT_COMBINE) + + if not all(self.name == col.name for col in cols): + index_ops = index_ops.rename(None) + + return index_ops + + return wrapper + + +def numpy_column_op(f: Callable[..., Column]) -> Callable[..., SeriesOrIndex]: + @wraps(f) + def wrapper(self: SeriesOrIndex, *args: Any) -> SeriesOrIndex: + # PySpark does not support NumPy type out of the box. For now, we convert NumPy types + # into some primitive types understandable in PySpark. + new_args = [] + for arg in args: + # TODO: This is a quick hack to support NumPy type. We should revisit this. + if isinstance(self.spark.data_type, LongType) and isinstance(arg, np.timedelta64): + new_args.append(float(arg / np.timedelta64(1, "s"))) + else: + new_args.append(arg) + return column_op(f)(self, *new_args) + + return wrapper + + +class IndexOpsMixin(object, metaclass=ABCMeta): + """common ops mixin to support a unified interface / docs for Series / Index + + Assuming there are following attributes or properties and functions. + """ + + @property + @abstractmethod + def _internal(self) -> InternalFrame: + pass + + @property + @abstractmethod + def _psdf(self) -> DataFrame: + pass + + @abstractmethod + def _with_new_scol( + self: IndexOpsLike, scol: Column, *, field: Optional[InternalField] = None + ) -> IndexOpsLike: + pass + + @property + @abstractmethod + def _column_label(self) -> Optional[Label]: + pass + + @property + @abstractmethod + def spark(self: IndexOpsLike) -> SparkIndexOpsMethods[IndexOpsLike]: + pass + + @property + def _dtype_op(self) -> "DataTypeOps": + from pyspark.pandas.data_type_ops.base import DataTypeOps + + return DataTypeOps(self.dtype, self.spark.data_type) + + @abstractmethod + def copy(self: IndexOpsLike) -> IndexOpsLike: + pass + + # arithmetic operators + def __neg__(self: IndexOpsLike) -> IndexOpsLike: + return self._dtype_op.neg(self) + + def __add__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.add(self, other) + + def __sub__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.sub(self, other) + + def __mul__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.mul(self, other) + + def __truediv__(self, other: Any) -> SeriesOrIndex: + """ + __truediv__ has different behaviour between pandas and PySpark for several cases. + 1. When dividing np.inf by zero, PySpark returns null whereas pandas returns np.inf + 2. When dividing a positive number by zero, PySpark returns null + whereas pandas returns np.inf + 3. When divide -np.inf by zero, PySpark returns null whereas pandas returns -np.inf + 4. When divide negative number by zero, PySpark returns null whereas pandas returns -np.inf + + +-------------------------------------------+ + | dividend (divisor: 0) | PySpark | pandas | + |-----------------------|---------|---------| + | np.inf | null | np.inf | + | -np.inf | null | -np.inf | + | 10 | null | np.inf | + | -10 | null | -np.inf | + +-----------------------|---------|---------+ + """ + return self._dtype_op.truediv(self, other) + + def __mod__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.mod(self, other) + + def __radd__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.radd(self, other) + + def __rsub__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rsub(self, other) + + def __rmul__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rmul(self, other) + + def __rtruediv__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rtruediv(self, other) + + def __floordiv__(self, other: Any) -> SeriesOrIndex: + """ + __floordiv__ has different behaviour between pandas and PySpark for several cases. + 1. When dividing np.inf by zero, PySpark returns null whereas pandas returns np.inf + 2. When dividing a positive number by zero, PySpark returns null + whereas pandas returns np.inf + 3. When divide -np.inf by zero, PySpark returns null whereas pandas returns -np.inf + 4. When divide negative number by zero, PySpark returns null whereas pandas returns -np.inf + + +-------------------------------------------+ + | dividend (divisor: 0) | PySpark | pandas | + |-----------------------|---------|---------| + | np.inf | null | np.inf | + | -np.inf | null | -np.inf | + | 10 | null | np.inf | + | -10 | null | -np.inf | + +-----------------------|---------|---------+ + """ + return self._dtype_op.floordiv(self, other) + + def __rfloordiv__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rfloordiv(self, other) + + def __rmod__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rmod(self, other) + + def __pow__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.pow(self, other) + + def __rpow__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rpow(self, other) + + def __abs__(self: IndexOpsLike) -> IndexOpsLike: + return self._dtype_op.abs(self) + + # comparison operators + def __eq__(self, other: Any) -> SeriesOrIndex: # type: ignore[override] + # pandas always returns False for all items with dict and set. + if isinstance(other, (dict, set)): + return self != self + else: + return self._dtype_op.eq(self, other) + + def __ne__(self, other: Any) -> SeriesOrIndex: # type: ignore[override] + return self._dtype_op.ne(self, other) + + def __lt__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.lt(self, other) + + def __le__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.le(self, other) + + def __ge__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.ge(self, other) + + def __gt__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.gt(self, other) + + def __invert__(self: IndexOpsLike) -> IndexOpsLike: + return self._dtype_op.invert(self) + + # `and`, `or`, `not` cannot be overloaded in Python, + # so use bitwise operators as boolean operators + def __and__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.__and__(self, other) + + def __or__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.__or__(self, other) + + def __rand__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rand(self, other) + + def __ror__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.ror(self, other) + + def __xor__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.xor(self, other) + + def __rxor__(self, other: Any) -> SeriesOrIndex: + return self._dtype_op.rxor(self, other) + + def __len__(self) -> int: + return len(self._psdf) + + # NDArray Compat + def __array_ufunc__( + self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any + ) -> SeriesOrIndex: + from pyspark.pandas import numpy_compat + + # Try dunder methods first. + result = numpy_compat.maybe_dispatch_ufunc_to_dunder_op( + self, ufunc, method, *inputs, **kwargs + ) + + # After that, we try with PySpark APIs. + if result is NotImplemented: + result = numpy_compat.maybe_dispatch_ufunc_to_spark_func( + self, ufunc, method, *inputs, **kwargs + ) + + if result is not NotImplemented: + return cast(SeriesOrIndex, result) + else: + # TODO: support more APIs? + raise NotImplementedError( + "pandas-on-Spark objects currently do not support %s." % ufunc + ) + + @property + def dtype(self) -> Dtype: + """Return the dtype object of the underlying data. + + Examples + -------- + >>> s = ps.Series([1, 2, 3]) + >>> s.dtype + dtype('int64') + + >>> s = ps.Series(list('abc')) + >>> s.dtype + dtype('O') + + >>> s = ps.Series(pd.date_range('20130101', periods=3)) + >>> s.dtype + dtype('>> s.rename("a").to_frame().set_index("a").index.dtype + dtype(' bool: + """ + Returns true if the current object is empty. Otherwise, it returns false. + + >>> ps.range(10).id.empty + False + + >>> ps.range(0).id.empty + True + + >>> ps.DataFrame({}, index=list('abc')).index.empty + False + """ + return self._internal.resolved_copy.spark_frame.rdd.isEmpty() + + @property + def hasnans(self) -> bool: + """ + Return True if it has any missing values. Otherwise, it returns False. + + >>> ps.DataFrame({}, index=list('abc')).index.hasnans + False + + >>> ps.Series(['a', None]).hasnans + True + + >>> ps.Series([1.0, 2.0, np.nan]).hasnans + True + + >>> ps.Series([1, 2, 3]).hasnans + False + + >>> (ps.Series([1.0, 2.0, np.nan]) + 1).hasnans + True + + >>> ps.Series([1, 2, 3]).rename("a").to_frame().set_index("a").index.hasnans + False + """ + return self.isnull().any() + + @property + def is_monotonic(self) -> bool: + """ + Return boolean if values in the object are monotonically increasing. + + .. note:: the current implementation of is_monotonic requires to shuffle + and aggregate multiple times to check the order locally and globally, + which is potentially expensive. In case of multi-index, all data is + transferred to a single node which can easily cause out-of-memory errors. + + .. note:: Disable the Spark config `spark.sql.optimizer.nestedSchemaPruning.enabled` + for multi-index if you're using pandas-on-Spark < 1.7.0 with PySpark 3.1.1. + + .. deprecated:: 3.4.0 + + Returns + ------- + is_monotonic : bool + + Examples + -------- + >>> ser = ps.Series(['1/1/2018', '3/1/2018', '4/1/2018']) + >>> ser.is_monotonic + True + + >>> df = ps.DataFrame({'dates': [None, '1/1/2018', '2/1/2018', '3/1/2018']}) + >>> df.dates.is_monotonic + False + + >>> df.index.is_monotonic + True + + >>> ser = ps.Series([1]) + >>> ser.is_monotonic + True + + >>> ser = ps.Series([]) + >>> ser.is_monotonic + True + + >>> ser.rename("a").to_frame().set_index("a").index.is_monotonic + True + + >>> ser = ps.Series([5, 4, 3, 2, 1], index=[1, 2, 3, 4, 5]) + >>> ser.is_monotonic + False + + >>> ser.index.is_monotonic + True + + Support for MultiIndex + + >>> midx = ps.MultiIndex.from_tuples( + ... [('x', 'a'), ('x', 'b'), ('y', 'c'), ('y', 'd'), ('z', 'e')]) + >>> midx # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('y', 'c'), + ('y', 'd'), + ('z', 'e')], + ) + >>> midx.is_monotonic + True + + >>> midx = ps.MultiIndex.from_tuples( + ... [('z', 'a'), ('z', 'b'), ('y', 'c'), ('y', 'd'), ('x', 'e')]) + >>> midx # doctest: +SKIP + MultiIndex([('z', 'a'), + ('z', 'b'), + ('y', 'c'), + ('y', 'd'), + ('x', 'e')], + ) + >>> midx.is_monotonic + False + """ + warnings.warn( + "is_monotonic is deprecated and will be removed in a future version. " + "Use is_monotonic_increasing instead.", + FutureWarning, + ) + return self._is_monotonic("increasing") + + @property + def is_monotonic_increasing(self) -> bool: + """ + Return boolean if values in the object are monotonically increasing. + + .. note:: the current implementation of is_monotonic_increasing requires to shuffle + and aggregate multiple times to check the order locally and globally, + which is potentially expensive. In case of multi-index, all data is + transferred to a single node which can easily cause out-of-memory errors. + + .. note:: Disable the Spark config `spark.sql.optimizer.nestedSchemaPruning.enabled` + for multi-index if you're using pandas-on-Spark < 1.7.0 with PySpark 3.1.1. + + Returns + ------- + is_monotonic : bool + + Examples + -------- + >>> ser = ps.Series(['1/1/2018', '3/1/2018', '4/1/2018']) + >>> ser.is_monotonic_increasing + True + + >>> df = ps.DataFrame({'dates': [None, '1/1/2018', '2/1/2018', '3/1/2018']}) + >>> df.dates.is_monotonic_increasing + False + + >>> df.index.is_monotonic_increasing + True + + >>> ser = ps.Series([1]) + >>> ser.is_monotonic_increasing + True + + >>> ser = ps.Series([]) + >>> ser.is_monotonic_increasing + True + + >>> ser.rename("a").to_frame().set_index("a").index.is_monotonic_increasing + True + + >>> ser = ps.Series([5, 4, 3, 2, 1], index=[1, 2, 3, 4, 5]) + >>> ser.is_monotonic_increasing + False + + >>> ser.index.is_monotonic_increasing + True + + Support for MultiIndex + + >>> midx = ps.MultiIndex.from_tuples( + ... [('x', 'a'), ('x', 'b'), ('y', 'c'), ('y', 'd'), ('z', 'e')]) + >>> midx # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('y', 'c'), + ('y', 'd'), + ('z', 'e')], + ) + >>> midx.is_monotonic_increasing + True + + >>> midx = ps.MultiIndex.from_tuples( + ... [('z', 'a'), ('z', 'b'), ('y', 'c'), ('y', 'd'), ('x', 'e')]) + >>> midx # doctest: +SKIP + MultiIndex([('z', 'a'), + ('z', 'b'), + ('y', 'c'), + ('y', 'd'), + ('x', 'e')], + ) + >>> midx.is_monotonic_increasing + False + """ + return self._is_monotonic("increasing") + + @property + def is_monotonic_decreasing(self) -> bool: + """ + Return boolean if values in the object are monotonically decreasing. + + .. note:: the current implementation of is_monotonic_decreasing requires to shuffle + and aggregate multiple times to check the order locally and globally, + which is potentially expensive. In case of multi-index, all data is transferred + to a single node which can easily cause out-of-memory errors. + + .. note:: Disable the Spark config `spark.sql.optimizer.nestedSchemaPruning.enabled` + for multi-index if you're using pandas-on-Spark < 1.7.0 with PySpark 3.1.1. + + Returns + ------- + is_monotonic : bool + + Examples + -------- + >>> ser = ps.Series(['4/1/2018', '3/1/2018', '1/1/2018']) + >>> ser.is_monotonic_decreasing + True + + >>> df = ps.DataFrame({'dates': [None, '3/1/2018', '2/1/2018', '1/1/2018']}) + >>> df.dates.is_monotonic_decreasing + False + + >>> df.index.is_monotonic_decreasing + False + + >>> ser = ps.Series([1]) + >>> ser.is_monotonic_decreasing + True + + >>> ser = ps.Series([]) + >>> ser.is_monotonic_decreasing + True + + >>> ser.rename("a").to_frame().set_index("a").index.is_monotonic_decreasing + True + + >>> ser = ps.Series([5, 4, 3, 2, 1], index=[1, 2, 3, 4, 5]) + >>> ser.is_monotonic_decreasing + True + + >>> ser.index.is_monotonic_decreasing + False + + Support for MultiIndex + + >>> midx = ps.MultiIndex.from_tuples( + ... [('x', 'a'), ('x', 'b'), ('y', 'c'), ('y', 'd'), ('z', 'e')]) + >>> midx # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('y', 'c'), + ('y', 'd'), + ('z', 'e')], + ) + >>> midx.is_monotonic_decreasing + False + + >>> midx = ps.MultiIndex.from_tuples( + ... [('z', 'e'), ('z', 'd'), ('y', 'c'), ('y', 'b'), ('x', 'a')]) + >>> midx # doctest: +SKIP + MultiIndex([('z', 'a'), + ('z', 'b'), + ('y', 'c'), + ('y', 'd'), + ('x', 'e')], + ) + >>> midx.is_monotonic_decreasing + True + """ + return self._is_monotonic("decreasing") + + def _is_locally_monotonic_spark_column(self, order: str) -> Column: + window = ( + Window.partitionBy(F.col("__partition_id")) + .orderBy(NATURAL_ORDER_COLUMN_NAME) + .rowsBetween(-1, -1) + ) + + if order == "increasing": + return (F.col("__origin") >= F.lag(F.col("__origin"), 1).over(window)) & F.col( + "__origin" + ).isNotNull() + else: + return (F.col("__origin") <= F.lag(F.col("__origin"), 1).over(window)) & F.col( + "__origin" + ).isNotNull() + + def _is_monotonic(self, order: str) -> bool: + assert order in ("increasing", "decreasing") + + sdf = self._internal.spark_frame + + sdf = ( + sdf.select( + F.spark_partition_id().alias( + "__partition_id" + ), # Make sure we use the same partition id in the whole job. + F.col(NATURAL_ORDER_COLUMN_NAME), + self.spark.column.alias("__origin"), + ) + .select( + F.col("__partition_id"), + F.col("__origin"), + self._is_locally_monotonic_spark_column(order).alias( + "__comparison_within_partition" + ), + ) + .groupby(F.col("__partition_id")) + .agg( + F.min(F.col("__origin")).alias("__partition_min"), + F.max(F.col("__origin")).alias("__partition_max"), + F.min(F.coalesce(F.col("__comparison_within_partition"), F.lit(True))).alias( + "__comparison_within_partition" + ), + ) + ) + + # Now we're windowing the aggregation results without partition specification. + # The number of rows here will be the same as partitions, which is expected + # to be small. + window = Window.orderBy(F.col("__partition_id")).rowsBetween(-1, -1) + if order == "increasing": + comparison_col = F.col("__partition_min") >= F.lag(F.col("__partition_max"), 1).over( + window + ) + else: + comparison_col = F.col("__partition_min") <= F.lag(F.col("__partition_max"), 1).over( + window + ) + + sdf = sdf.select( + comparison_col.alias("__comparison_between_partitions"), + F.col("__comparison_within_partition"), + ) + + ret = sdf.select( + F.min(F.coalesce(F.col("__comparison_between_partitions"), F.lit(True))) + & F.min(F.coalesce(F.col("__comparison_within_partition"), F.lit(True))) + ).collect()[0][0] + if ret is None: + return True + else: + return ret + + @property + def ndim(self) -> int: + """ + Return an int representing the number of array dimensions. + + Return 1 for Series / Index / MultiIndex. + + Examples + -------- + + For Series + + >>> s = ps.Series([None, 1, 2, 3, 4], index=[4, 5, 2, 1, 8]) + >>> s.ndim + 1 + + For Index + + >>> s.index.ndim + 1 + + For MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [1, 1, 1, 1, 1, 2, 1, 2, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + >>> s.index.ndim + 1 + """ + return 1 + + def astype(self: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + """ + Cast a pandas-on-Spark object to a specified dtype ``dtype``. + + Parameters + ---------- + dtype : data type + Use a numpy.dtype or Python type to cast entire pandas object to + the same type. + + Returns + ------- + casted : same type as caller + + See Also + -------- + to_datetime : Convert argument to datetime. + + Examples + -------- + >>> ser = ps.Series([1, 2], dtype='int32') + >>> ser + 0 1 + 1 2 + dtype: int32 + + >>> ser.astype('int64') + 0 1 + 1 2 + dtype: int64 + + >>> ser.rename("a").to_frame().set_index("a").index.astype('int64') # doctest: +SKIP + Int64Index([1, 2], dtype='int64', name='a') + """ + return self._dtype_op.astype(self, dtype) + + def isin(self: IndexOpsLike, values: Sequence[Any]) -> IndexOpsLike: + """ + Check whether `values` are contained in Series or Index. + + Return a boolean Series or Index showing whether each element in the Series + matches an element in the passed sequence of `values` exactly. + + Parameters + ---------- + values : set or list-like + The sequence of values to test. + + Returns + ------- + isin : Series (bool dtype) or Index (bool dtype) + + Examples + -------- + >>> s = ps.Series(['lama', 'cow', 'lama', 'beetle', 'lama', + ... 'hippo'], name='animal') + >>> s.isin(['cow', 'lama']) + 0 True + 1 True + 2 True + 3 False + 4 True + 5 False + Name: animal, dtype: bool + + Passing a single string as ``s.isin('lama')`` will raise an error. Use + a list of one element instead: + + >>> s.isin(['lama']) + 0 True + 1 False + 2 True + 3 False + 4 True + 5 False + Name: animal, dtype: bool + + >>> s.rename("a").to_frame().set_index("a").index.isin(['lama']) # doctest: +SKIP + Index([True, False, True, False, True, False], dtype='bool', name='a') + """ + if not is_list_like(values): + raise TypeError( + "only list-like objects are allowed to be passed" + " to isin(), you passed a [{values_type}]".format(values_type=type(values).__name__) + ) + + values = ( + cast(np.ndarray, values).tolist() if isinstance(values, np.ndarray) else list(values) + ) + + other = [F.lit(v) for v in values] + scol = self.spark.column.isin(other) + field = self._internal.data_fields[0].copy( + dtype=np.dtype("bool"), spark_type=BooleanType(), nullable=False + ) + return self._with_new_scol(scol=F.coalesce(scol, F.lit(False)), field=field) + + def isnull(self: IndexOpsLike) -> IndexOpsLike: + """ + Detect existing (non-missing) values. + + Return a boolean same-sized object indicating if the values are NA. + NA values, such as None or numpy.NaN, get mapped to True values. + Everything else gets mapped to False values. Characters such as empty strings '' or + numpy.inf are not considered NA values + (unless you set pandas.options.mode.use_inf_as_na = True). + + Returns + ------- + Series or Index : Mask of bool values for each element in Series + that indicates whether an element is not an NA value. + + Examples + -------- + >>> ser = ps.Series([5, 6, np.NaN]) + >>> ser.isna() # doctest: +NORMALIZE_WHITESPACE + 0 False + 1 False + 2 True + dtype: bool + + >>> ser.rename("a").to_frame().set_index("a").index.isna() # doctest: +SKIP + Index([False, False, True], dtype='bool', name='a') + """ + from pyspark.pandas.indexes import MultiIndex + + if isinstance(self, MultiIndex): + raise NotImplementedError("isna is not defined for MultiIndex") + + return self._dtype_op.isnull(self) + + isna = isnull + + def notnull(self: IndexOpsLike) -> IndexOpsLike: + """ + Detect existing (non-missing) values. + Return a boolean same-sized object indicating if the values are not NA. + Non-missing values get mapped to True. + Characters such as empty strings '' or numpy.inf are not considered NA values + (unless you set pandas.options.mode.use_inf_as_na = True). + NA values, such as None or numpy.NaN, get mapped to False values. + + Returns + ------- + Series or Index : Mask of bool values for each element in Series + that indicates whether an element is not an NA value. + + Examples + -------- + Show which entries in a Series are not NA. + + >>> ser = ps.Series([5, 6, np.NaN]) + >>> ser + 0 5.0 + 1 6.0 + 2 NaN + dtype: float64 + + >>> ser.notna() + 0 True + 1 True + 2 False + dtype: bool + + >>> ser.rename("a").to_frame().set_index("a").index.notna() # doctest: +SKIP + Index([True, True, False], dtype='bool', name='a') + """ + from pyspark.pandas.indexes import MultiIndex + + if isinstance(self, MultiIndex): + raise NotImplementedError("notna is not defined for MultiIndex") + return (~self.isnull()).rename(self.name) # type: ignore[attr-defined] + + notna = notnull + + # TODO: axis and many arguments should be implemented. + def all(self, axis: Axis = 0, skipna: bool = True) -> bool: + """ + Return whether all elements are True. + + Returns True unless there at least one element within a series that is + False or equivalent (e.g. zero or empty) + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Indicate which axis or axes should be reduced. + + * 0 / 'index' : reduce the index, return a Series whose index is the + original column labels. + + skipna : boolean, default True + Exclude NA values, such as None or numpy.NaN. + If an entire row/column is NA values and `skipna` is True, + then the result will be True, as for an empty row/column. + If `skipna` is False, numpy.NaNs are treated as True because these are + not equal to zero, Nones are treated as False. + + Examples + -------- + >>> ps.Series([True, True]).all() + True + + >>> ps.Series([True, False]).all() + False + + >>> ps.Series([0, 1]).all() + False + + >>> ps.Series([1, 2, 3]).all() + True + + >>> ps.Series([True, True, None]).all() + True + + >>> ps.Series([True, True, None]).all(skipna=False) + False + + >>> ps.Series([True, False, None]).all() + False + + >>> ps.Series([]).all() + True + + >>> ps.Series([np.nan]).all() + True + + >>> ps.Series([np.nan]).all(skipna=False) + True + + >>> ps.Series([None]).all() + True + + >>> ps.Series([None]).all(skipna=False) + False + + >>> df = ps.Series([True, False, None]).rename("a").to_frame() + >>> df.set_index("a").index.all() + False + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + sdf = self._internal.spark_frame.select(self.spark.column) + col = scol_for(sdf, sdf.columns[0]) + + # `any` and `every` was added as of Spark 3.0. + # ret = sdf.select(F.expr("every(CAST(`%s` AS BOOLEAN))" % sdf.columns[0])).collect()[0][0] + # We use min as its alternative as below. + if isinstance(self.spark.data_type, NumericType) or skipna: + # np.nan takes no effect to the result; None takes no effect if `skipna` + ret = sdf.select(F.min(F.coalesce(col.cast("boolean"), F.lit(True)))).collect()[0][0] + else: + # Take None as False when not `skipna` + ret = sdf.select( + F.min(F.when(col.isNull(), F.lit(False)).otherwise(col.cast("boolean"))) + ).collect()[0][0] + + if ret is None: + return True + else: + return ret + + # TODO: axis, skipna, and many arguments should be implemented. + def any(self, axis: Axis = 0) -> bool: + """ + Return whether any element is True. + + Returns False unless there is at least one element within a series that is + True or equivalent (e.g. non-zero or non-empty). + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Indicate which axis or axes should be reduced. + + * 0 / 'index' : reduce the index, return a Series whose index is the + original column labels. + + Examples + -------- + >>> ps.Series([False, False]).any() + False + + >>> ps.Series([True, False]).any() + True + + >>> ps.Series([0, 0]).any() + False + + >>> ps.Series([0, 1, 2]).any() + True + + >>> ps.Series([False, False, None]).any() + False + + >>> ps.Series([True, False, None]).any() + True + + >>> ps.Series([]).any() + False + + >>> ps.Series([np.nan]).any() + False + + >>> df = ps.Series([True, False, None]).rename("a").to_frame() + >>> df.set_index("a").index.any() + True + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + sdf = self._internal.spark_frame.select(self.spark.column) + col = scol_for(sdf, sdf.columns[0]) + + # Note that we're ignoring `None`s here for now. + # any and every was added as of Spark 3.0 + # ret = sdf.select(F.expr("any(CAST(`%s` AS BOOLEAN))" % sdf.columns[0])).collect()[0][0] + # Here we use max as its alternative: + ret = sdf.select(F.max(F.coalesce(col.cast("boolean"), F.lit(False)))).collect()[0][0] + if ret is None: + return False + else: + return ret + + # TODO: add frep and axis parameter + def shift( + self: IndexOpsLike, periods: int = 1, fill_value: Optional[Any] = None + ) -> IndexOpsLike: + """ + Shift Series/Index by desired number of periods. + + .. note:: the current implementation of shift uses Spark's Window without + specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int + Number of periods to shift. Can be positive or negative. + fill_value : object, optional + The scalar value to use for newly introduced missing values. + The default depends on the dtype of self. For numeric data, np.nan is used. + + Returns + ------- + Copy of input Series/Index, shifted. + + Examples + -------- + >>> df = ps.DataFrame({'Col1': [10, 20, 15, 30, 45], + ... 'Col2': [13, 23, 18, 33, 48], + ... 'Col3': [17, 27, 22, 37, 52]}, + ... columns=['Col1', 'Col2', 'Col3']) + + >>> df.Col1.shift(periods=3) + 0 NaN + 1 NaN + 2 NaN + 3 10.0 + 4 20.0 + Name: Col1, dtype: float64 + + >>> df.Col2.shift(periods=3, fill_value=0) + 0 0 + 1 0 + 2 0 + 3 13 + 4 23 + Name: Col2, dtype: int64 + + >>> df.index.shift(periods=3, fill_value=0) # doctest: +SKIP + Int64Index([0, 0, 0, 0, 1], dtype='int64') + """ + return self._shift(periods, fill_value).spark.analyzed + + def _shift( + self: IndexOpsLike, + periods: int, + fill_value: Any, + *, + part_cols: Sequence["ColumnOrName"] = (), + ) -> IndexOpsLike: + if not isinstance(periods, int): + raise TypeError("periods should be an int; however, got [%s]" % type(periods).__name__) + + if periods == 0: + return self.copy() + + col = self.spark.column + window = ( + Window.partitionBy(*part_cols) + .orderBy(NATURAL_ORDER_COLUMN_NAME) + .rowsBetween(-periods, -periods) + ) + lag_col = F.lag(col, periods).over(window) + col = F.when(lag_col.isNull() | F.isnan(lag_col), fill_value).otherwise(lag_col) + return self._with_new_scol(col, field=self._internal.data_fields[0].copy(nullable=True)) + + # TODO: Update Documentation for Bins Parameter when its supported + def value_counts( + self, + normalize: bool = False, + sort: bool = True, + ascending: bool = False, + bins: None = None, + dropna: bool = True, + ) -> "Series": + """ + Return a Series containing counts of unique values. + The resulting object will be in descending order so that the + first element is the most frequently-occurring element. + Excludes NA values by default. + + Parameters + ---------- + normalize : boolean, default False + If True then the object returned will contain the relative + frequencies of the unique values. + sort : boolean, default True + Sort by values. + ascending : boolean, default False + Sort in ascending order. + bins : Not Yet Supported + dropna : boolean, default True + Don't include counts of NaN. + + Returns + ------- + counts : Series + + See Also + -------- + Series.count: Number of non-NA elements in a Series. + + Examples + -------- + For Series + + >>> df = ps.DataFrame({'x':[0, 0, 1, 1, 1, np.nan]}) + >>> df.x.value_counts() # doctest: +NORMALIZE_WHITESPACE + 1.0 3 + 0.0 2 + Name: x, dtype: int64 + + With `normalize` set to `True`, returns the relative frequency by + dividing all values by the sum of values. + + >>> df.x.value_counts(normalize=True) # doctest: +NORMALIZE_WHITESPACE + 1.0 0.6 + 0.0 0.4 + Name: x, dtype: float64 + + **dropna** + With `dropna` set to `False` we can also see NaN index values. + + >>> df.x.value_counts(dropna=False) # doctest: +NORMALIZE_WHITESPACE + 1.0 3 + 0.0 2 + NaN 1 + Name: x, dtype: int64 + + For Index + + >>> idx = ps.Index([3, 1, 2, 3, 4, np.nan]) + >>> idx # doctest: +SKIP + Float64Index([3.0, 1.0, 2.0, 3.0, 4.0, nan], dtype='float64') + + >>> idx.value_counts().sort_index() + 1.0 1 + 2.0 1 + 3.0 2 + 4.0 1 + dtype: int64 + + **sort** + + With `sort` set to `False`, the result wouldn't be sorted by number of count. + + >>> idx.value_counts(sort=True).sort_index() + 1.0 1 + 2.0 1 + 3.0 2 + 4.0 1 + dtype: int64 + + **normalize** + + With `normalize` set to `True`, returns the relative frequency by + dividing all values by the sum of values. + + >>> idx.value_counts(normalize=True).sort_index() + 1.0 0.2 + 2.0 0.2 + 3.0 0.4 + 4.0 0.2 + dtype: float64 + + **dropna** + + With `dropna` set to `False` we can also see NaN index values. + + >>> idx.value_counts(dropna=False).sort_index() # doctest: +SKIP + 1.0 1 + 2.0 1 + 3.0 2 + 4.0 1 + NaN 1 + dtype: int64 + + For MultiIndex. + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [1, 1, 1, 1, 1, 2, 1, 2, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + >>> s.index # doctest: +SKIP + MultiIndex([( 'lama', 'weight'), + ( 'lama', 'weight'), + ( 'lama', 'weight'), + ( 'cow', 'weight'), + ( 'cow', 'weight'), + ( 'cow', 'length'), + ('falcon', 'weight'), + ('falcon', 'length'), + ('falcon', 'length')], + ) + + >>> s.index.value_counts().sort_index() + (cow, length) 1 + (cow, weight) 2 + (falcon, length) 2 + (falcon, weight) 1 + (lama, weight) 3 + dtype: int64 + + >>> s.index.value_counts(normalize=True).sort_index() + (cow, length) 0.111111 + (cow, weight) 0.222222 + (falcon, length) 0.222222 + (falcon, weight) 0.111111 + (lama, weight) 0.333333 + dtype: float64 + + If Index has name, keep the name up. + + >>> idx = ps.Index([0, 0, 0, 1, 1, 2, 3], name='pandas-on-Spark') + >>> idx.value_counts().sort_index() + 0 3 + 1 2 + 2 1 + 3 1 + Name: pandas-on-Spark, dtype: int64 + """ + from pyspark.pandas.series import first_series, Series + + if isinstance(self, Series): + warnings.warn( + "The resulting Series will have a fixed name of 'count' from 4.0.0.", + FutureWarning, + ) + + if bins is not None: + raise NotImplementedError("value_counts currently does not support bins") + + if dropna: + sdf_dropna = self._internal.spark_frame.select(self.spark.column).dropna() + else: + sdf_dropna = self._internal.spark_frame.select(self.spark.column) + index_name = SPARK_DEFAULT_INDEX_NAME + column_name = self._internal.data_spark_column_names[0] + sdf = sdf_dropna.groupby(scol_for(sdf_dropna, column_name).alias(index_name)).count() + if sort: + if ascending: + sdf = sdf.orderBy(F.col("count")) + else: + sdf = sdf.orderBy(F.col("count").desc()) + + if normalize: + drop_sum = sdf_dropna.count() + sdf = sdf.withColumn("count", F.col("count") / F.lit(drop_sum)) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, index_name)], + column_labels=self._internal.column_labels, + data_spark_columns=[scol_for(sdf, "count")], + column_label_names=self._internal.column_label_names, + ) + + return first_series(DataFrame(internal)) + + def nunique(self, dropna: bool = True, approx: bool = False, rsd: float = 0.05) -> int: + """ + Return number of unique elements in the object. + Excludes NA values by default. + + Parameters + ---------- + dropna : bool, default True + Don’t include NaN in the count. + approx: bool, default False + If False, will use the exact algorithm and return the exact number of unique. + If True, it uses the HyperLogLog approximate algorithm, which is significantly faster + for large amount of data. + Note: This parameter is specific to pandas-on-Spark and is not found in pandas. + rsd: float, default 0.05 + Maximum estimation error allowed in the HyperLogLog algorithm. + Note: Just like ``approx`` this parameter is specific to pandas-on-Spark. + + Returns + ------- + int + + See Also + -------- + DataFrame.nunique: Method nunique for DataFrame. + Series.count: Count non-NA/null observations in the Series. + + Examples + -------- + >>> ps.Series([1, 2, 3, np.nan]).nunique() + 3 + + >>> ps.Series([1, 2, 3, np.nan]).nunique(dropna=False) + 4 + + On big data, we recommend using the approximate algorithm to speed up this function. + The result will be very close to the exact unique count. + + >>> ps.Series([1, 2, 3, np.nan]).nunique(approx=True) + 3 + + >>> idx = ps.Index([1, 1, 2, None]) + >>> idx # doctest: +SKIP + Float64Index([1.0, 1.0, 2.0, nan], dtype='float64') + + >>> idx.nunique() + 2 + + >>> idx.nunique(dropna=False) + 3 + """ + res = self._internal.spark_frame.select([self._nunique(dropna, approx, rsd)]) + return res.collect()[0][0] + + def _nunique(self, dropna: bool = True, approx: bool = False, rsd: float = 0.05) -> Column: + colname = self._internal.data_spark_column_names[0] + count_fn = cast( + Callable[[Column], Column], + partial(F.approx_count_distinct, rsd=rsd) if approx else F.countDistinct, + ) + if dropna: + return count_fn(self.spark.column).alias(colname) + else: + return ( + count_fn(self.spark.column) + + F.when( + F.count(F.when(self.spark.column.isNull(), 1).otherwise(None)) >= 1, 1 + ).otherwise(0) + ).alias(colname) + + def take(self: IndexOpsLike, indices: Sequence[int]) -> IndexOpsLike: + """ + Return the elements in the given *positional* indices along an axis. + + This means that we are not indexing according to actual values in + the index attribute of the object. We are indexing according to the + actual position of the element in the object. + + Parameters + ---------- + indices : array-like + An array of ints indicating which positions to take. + + Returns + ------- + taken : same type as caller + An array-like containing the elements taken from the object. + + See Also + -------- + DataFrame.loc : Select a subset of a DataFrame by labels. + DataFrame.iloc : Select a subset of a DataFrame by positions. + numpy.take : Take elements from an array along an axis. + + Examples + -------- + + Series + + >>> psser = ps.Series([100, 200, 300, 400, 500]) + >>> psser + 0 100 + 1 200 + 2 300 + 3 400 + 4 500 + dtype: int64 + + >>> psser.take([0, 2, 4]).sort_index() + 0 100 + 2 300 + 4 500 + dtype: int64 + + Index + + >>> psidx = ps.Index([100, 200, 300, 400, 500]) + >>> psidx # doctest: +SKIP + Int64Index([100, 200, 300, 400, 500], dtype='int64') + + >>> psidx.take([0, 2, 4]).sort_values() # doctest: +SKIP + Int64Index([100, 300, 500], dtype='int64') + + MultiIndex + + >>> psmidx = ps.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("x", "c")]) + >>> psmidx # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('x', 'c')], + ) + + >>> psmidx.take([0, 2]) # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'c')], + ) + """ + if not is_list_like(indices) or isinstance(indices, (dict, set)): + raise TypeError("`indices` must be a list-like except dict or set") + if isinstance(self, ps.Series): + return cast(IndexOpsLike, self.iloc[indices]) + else: + return cast(IndexOpsLike, self._psdf.iloc[indices].index) + + def factorize( + self: IndexOpsLike, sort: bool = True, na_sentinel: Optional[int] = -1 + ) -> Tuple[IndexOpsLike, pd.Index]: + """ + Encode the object as an enumerated type or categorical variable. + + This method is useful for obtaining a numeric representation of an + array when all that matters is identifying distinct values. + + Parameters + ---------- + sort : bool, default True + na_sentinel : int or None, default -1 + Value to mark "not found". If None, will not drop the NaN + from the uniques of the values. + + .. deprecated:: 3.4.0 + + Returns + ------- + codes : Series or Index + A Series or Index that's an indexer into `uniques`. + ``uniques.take(codes)`` will have the same values as `values`. + uniques : pd.Index + The unique valid values. + + .. note :: + + Even if there's a missing value in `values`, `uniques` will + *not* contain an entry for it. + + Examples + -------- + >>> psser = ps.Series(['b', None, 'a', 'c', 'b']) + >>> codes, uniques = psser.factorize() + >>> codes + 0 1 + 1 -1 + 2 0 + 3 2 + 4 1 + dtype: int32 + >>> uniques + Index(['a', 'b', 'c'], dtype='object') + + >>> codes, uniques = psser.factorize(na_sentinel=None) + >>> codes + 0 1 + 1 3 + 2 0 + 3 2 + 4 1 + dtype: int32 + >>> uniques + Index(['a', 'b', 'c', None], dtype='object') + + >>> codes, uniques = psser.factorize(na_sentinel=-2) + >>> codes + 0 1 + 1 -2 + 2 0 + 3 2 + 4 1 + dtype: int32 + >>> uniques + Index(['a', 'b', 'c'], dtype='object') + + For Index: + + >>> psidx = ps.Index(['b', None, 'a', 'c', 'b']) + >>> codes, uniques = psidx.factorize() + >>> codes # doctest: +SKIP + Int64Index([1, -1, 0, 2, 1], dtype='int64') + >>> uniques + Index(['a', 'b', 'c'], dtype='object') + """ + from pyspark.pandas.series import first_series + + assert (na_sentinel is None) or isinstance(na_sentinel, int) + assert sort is True + + warnings.warn( + "Argument `na_sentinel` will be removed in 4.0.0.", + FutureWarning, + ) + + if isinstance(self.dtype, CategoricalDtype): + categories = self.dtype.categories + if len(categories) == 0: + scol = F.lit(None) + else: + kvs = list( + chain( + *[ + (F.lit(code), F.lit(category)) + for code, category in enumerate(categories) + ] + ) + ) + map_scol = F.create_map(*kvs) + scol = map_scol[self.spark.column] + codes, uniques = self._with_new_scol( + scol.alias(self._internal.data_spark_column_names[0]) + ).factorize(na_sentinel=na_sentinel) + return codes, uniques.astype(self.dtype) + + uniq_sdf = self._internal.spark_frame.select(self.spark.column).distinct() + + # Check number of uniques and constructs sorted `uniques_list` + max_compute_count = get_option("compute.max_rows") + if max_compute_count is not None: + uniq_pdf = uniq_sdf.limit(max_compute_count + 1).toPandas() + if len(uniq_pdf) > max_compute_count: + raise ValueError( + "Current Series has more then {0} unique values. " + "Please set 'compute.max_rows' by using 'pyspark.pandas.config.set_option' " + "to more than {0} rows. Note that, before changing the " + "'compute.max_rows', this operation is considerably expensive.".format( + max_compute_count + ) + ) + else: + uniq_pdf = uniq_sdf.toPandas() + # pandas takes both NaN and null in Spark to np.nan, so de-duplication is required + uniq_series = first_series(uniq_pdf).drop_duplicates() + uniques_list = uniq_series.tolist() + uniques_list = sorted(uniques_list, key=lambda x: (pd.isna(x), x)) + + # Constructs `unique_to_code` mapping non-na unique to code + unique_to_code = {} + if na_sentinel is not None: + na_sentinel_code = na_sentinel + code = 0 + for unique in uniques_list: + if pd.isna(unique): + if na_sentinel is None: + na_sentinel_code = code + else: + unique_to_code[unique] = code + code += 1 + + kvs = list( + chain(*([(F.lit(unique), F.lit(code)) for unique, code in unique_to_code.items()])) + ) + + if len(kvs) == 0: # uniques are all missing values + new_scol = F.lit(na_sentinel_code) + else: + map_scol = F.create_map(*kvs) + null_scol = F.when(self.isnull().spark.column, F.lit(na_sentinel_code)) + new_scol = null_scol.otherwise(map_scol[self.spark.column]) + + codes = self._with_new_scol(new_scol.alias(self._internal.data_spark_column_names[0])) + + if na_sentinel is not None: + # Drops the NaN from the uniques of the values + uniques_list = [x for x in uniques_list if not pd.isna(x)] + + uniques = pd.Index(uniques_list) + + return codes, uniques + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.base + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.base.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.base tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.base, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/categorical.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/categorical.py new file mode 100644 index 0000000..36b11ca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/categorical.py @@ -0,0 +1,822 @@ +# +# 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 Any, Callable, List, Optional, Union, TYPE_CHECKING, cast +import warnings + +import pandas as pd +from pandas.api.types import ( # type: ignore[attr-defined] + CategoricalDtype, + is_dict_like, + is_list_like, +) + +from pyspark.pandas.internal import InternalField +from pyspark.pandas.data_type_ops.categorical_ops import _to_cat +from pyspark.sql import functions as F +from pyspark.sql.types import StructField + +if TYPE_CHECKING: + import pyspark.pandas as ps + + +class CategoricalAccessor: + """ + Accessor object for categorical properties of the Series values. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.categories + Index(['a', 'b', 'c'], dtype='object') + + >>> s.cat.codes + 0 0 + 1 1 + 2 1 + 3 2 + 4 2 + 5 2 + dtype: int8 + """ + + def __init__(self, series: "ps.Series"): + if not isinstance(series.dtype, CategoricalDtype): + raise ValueError("Cannot call CategoricalAccessor on type {}".format(series.dtype)) + self._data = series + + @property + def _dtype(self) -> CategoricalDtype: + return cast(CategoricalDtype, self._data.dtype) + + @property + def categories(self) -> pd.Index: + """ + The categories of this categorical. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.categories + Index(['a', 'b', 'c'], dtype='object') + """ + return self._dtype.categories + + @categories.setter + def categories(self, categories: Union[pd.Index, List]) -> None: + dtype = CategoricalDtype(categories, ordered=self.ordered) + + if len(self.categories) != len(dtype.categories): + raise ValueError( + "new categories need to have the same number of items as the old categories!" + ) + + internal = self._data._psdf._internal.with_new_spark_column( + self._data._column_label, + self._data.spark.column, + field=self._data._internal.data_fields[0].copy(dtype=dtype), + ) + self._data._psdf._update_internal_frame(internal) + + @property + def ordered(self) -> bool: + """ + Whether the categories have an ordered relationship. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.ordered + False + """ + return self._dtype.ordered + + @property + def codes(self) -> "ps.Series": + """ + Return Series of codes as well as the index. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.codes + 0 0 + 1 1 + 2 1 + 3 2 + 4 2 + 5 2 + dtype: int8 + """ + return self._data._with_new_scol( + self._data.spark.column, + field=InternalField.from_struct_field( + StructField( + name=self._data._internal.data_spark_column_names[0], + dataType=self._data.spark.data_type, + nullable=self._data.spark.nullable, + ) + ), + ).rename() + + def add_categories(self, new_categories: Union[pd.Index, Any, List]) -> Optional["ps.Series"]: + """ + Add new categories. + + `new_categories` will be included at the last/highest place in the + categories and will be unused directly after this call. + + Parameters + ---------- + new_categories : category or list-like of category + The new categories to be included. + + Returns + ------- + Series or None + Categorical with new categories added or None if ``inplace=True``. + + Raises + ------ + ValueError + If the new categories include old categories or do not validate as + categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.add_categories('x') # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (4, object): ['a', 'b', 'c', 'x'] + """ + from pyspark.pandas.frame import DataFrame + + categories: List[Any] + if is_list_like(new_categories): + categories = list(new_categories) + else: + categories = [new_categories] + + if any(cat in self.categories for cat in categories): + raise ValueError( + "new categories must not include old categories: {{{cats}}}".format( + cats=", ".join(set(str(cat) for cat in categories if cat in self.categories)) + ) + ) + + internal = self._data._psdf._internal.with_new_spark_column( + self._data._column_label, + self._data.spark.column, + field=self._data._internal.data_fields[0].copy( + dtype=CategoricalDtype(list(self.categories) + categories, ordered=self.ordered) + ), + ) + return DataFrame(internal)._psser_for(self._data._column_label).copy() + + def _set_ordered(self, *, ordered: bool, inplace: bool) -> Optional["ps.Series"]: + from pyspark.pandas.frame import DataFrame + + if self.ordered == ordered: + if inplace: + return None + else: + return self._data.copy() + else: + internal = self._data._psdf._internal.with_new_spark_column( + self._data._column_label, + self._data.spark.column, + field=self._data._internal.data_fields[0].copy( + dtype=CategoricalDtype(categories=self.categories, ordered=ordered) + ), + ) + if inplace: + self._data._psdf._update_internal_frame(internal) + return None + else: + return DataFrame(internal)._psser_for(self._data._column_label).copy() + + def as_ordered(self, inplace: bool = False) -> Optional["ps.Series"]: + """ + Set the Categorical to be ordered. + + Parameters + ---------- + inplace : bool, default False + Whether or not to set the ordered attribute in-place or return + a copy of this categorical with ordered set to True. + + .. deprecated:: 3.4.0 + + Returns + ------- + Series or None + Ordered Categorical or None if ``inplace=True``. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.as_ordered() # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a' < 'b' < 'c'] + """ + if inplace: + warnings.warn( + "The `inplace` parameter in as_ordered is deprecated " + "and will be removed in a future version.", + FutureWarning, + ) + return self._set_ordered(ordered=True, inplace=inplace) + + def as_unordered(self, inplace: bool = False) -> Optional["ps.Series"]: + """ + Set the Categorical to be unordered. + + Parameters + ---------- + inplace : bool, default False + Whether or not to set the ordered attribute in-place or return + a copy of this categorical with ordered set to False. + + .. deprecated:: 3.4.0 + + Returns + ------- + Series or None + Unordered Categorical or None if ``inplace=True``. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category").cat.as_ordered() + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a' < 'b' < 'c'] + + >>> s.cat.as_unordered() # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + """ + if inplace: + warnings.warn( + "The `inplace` parameter in as_unordered is deprecated " + "and will be removed in a future version.", + FutureWarning, + ) + return self._set_ordered(ordered=False, inplace=inplace) + + def remove_categories(self, removals: Union[pd.Index, Any, List]) -> Optional["ps.Series"]: + """ + Remove the specified categories. + + `removals` must be included in the old categories. Values which were in + the removed categories will be set to NaN + + Parameters + ---------- + removals : category or list of categories + The categories which should be removed. + + Returns + ------- + Series or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If the removals are not contained in the categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.remove_categories('b') # doctest: +SKIP + 0 a + 1 NaN + 2 NaN + 3 c + 4 c + 5 c + dtype: category + Categories (2, object): ['a', 'c'] + """ + categories: List[Any] + if is_list_like(removals): + categories = [cat for cat in removals if cat is not None] + elif removals is None: + categories = [] + else: + categories = [removals] + + if any(cat not in self.categories for cat in categories): + raise ValueError( + "removals must all be in old categories: {{{cats}}}".format( + cats=", ".join( + set(str(cat) for cat in categories if cat not in self.categories) + ) + ) + ) + + if len(categories) == 0: + return self._data.copy() + else: + dtype = CategoricalDtype( + [cat for cat in self.categories if cat not in categories], ordered=self.ordered + ) + return self._data.astype(dtype) + + def remove_unused_categories(self) -> Optional["ps.Series"]: + """ + Remove categories which are not used. + + Returns + ------- + cat : Series or None + Categorical with unused categories dropped or None if ``inplace=True``. + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> s = ps.Series(pd.Categorical(list("abbccc"), categories=['a', 'b', 'c', 'd'])) + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (4, object): ['a', 'b', 'c', 'd'] + + >>> s.cat.remove_unused_categories() # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + """ + categories = set(self._data.drop_duplicates()._to_pandas()) + removals = [cat for cat in self.categories if cat not in categories] + return self.remove_categories(removals=removals) + + def rename_categories( + self, new_categories: Union[list, dict, Callable] + ) -> Optional["ps.Series"]: + """ + Rename categories. + + Parameters + ---------- + new_categories : list-like, dict-like or callable + + New categories which will replace old categories. + + * list-like: all items must be unique and the number of items in + the new categories must match the existing number of categories. + + * dict-like: specifies a mapping from + old categories to new. Categories not contained in the mapping + are passed through and extra categories in the mapping are + ignored. + + * callable : a callable that is called on all items in the old + categories and whose return values comprise the new categories. + + Returns + ------- + cat : Series or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If new categories are list-like and do not have the same number of + items than the current categories or do not validate as categories + + See Also + -------- + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> s = ps.Series(["a", "a", "b"], dtype="category") + >>> s.cat.rename_categories([0, 1]) # doctest: +SKIP + 0 0 + 1 0 + 2 1 + dtype: category + Categories (2, int64): [0, 1] + + For dict-like ``new_categories``, extra keys are ignored and + categories not in the dictionary are passed through + + >>> s.cat.rename_categories({'a': 'A', 'c': 'C'}) # doctest: +SKIP + 0 A + 1 A + 2 b + dtype: category + Categories (2, object): ['A', 'b'] + + You may also provide a callable to create the new categories + + >>> s.cat.rename_categories(lambda x: x.upper()) # doctest: +SKIP + 0 A + 1 A + 2 B + dtype: category + Categories (2, object): ['A', 'B'] + """ + from pyspark.pandas.frame import DataFrame + + if is_dict_like(new_categories): + categories = [cast(dict, new_categories).get(item, item) for item in self.categories] + elif callable(new_categories): + categories = [new_categories(item) for item in self.categories] + elif is_list_like(new_categories): + if len(self.categories) != len(new_categories): + raise ValueError( + "new categories need to have the same number of items as the old categories!" + ) + categories = cast(list, new_categories) + else: + raise TypeError("new_categories must be list-like, dict-like or callable.") + + internal = self._data._psdf._internal.with_new_spark_column( + self._data._column_label, + self._data.spark.column, + field=self._data._internal.data_fields[0].copy( + dtype=CategoricalDtype(categories=categories, ordered=self.ordered) + ), + ) + + return DataFrame(internal)._psser_for(self._data._column_label).copy() + + def reorder_categories( + self, + new_categories: Union[pd.Index, List], + ordered: Optional[bool] = None, + ) -> Optional["ps.Series"]: + """ + Reorder categories as specified in new_categories. + + `new_categories` needs to include all old categories and no new category + items. + + Parameters + ---------- + new_categories : Index-like + The categories in new order. + ordered : bool, optional + Whether or not the categorical is treated as an ordered categorical. + If not given, do not change the ordered information. + + Returns + ------- + cat : Series or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If the new categories do not contain all old category items or any + new ones + + See Also + -------- + rename_categories : Rename categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.reorder_categories(['c', 'b', 'a'], ordered=True) # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['c' < 'b' < 'a'] + """ + if not is_list_like(new_categories): + raise TypeError( + "Parameter 'new_categories' must be list-like, was '{}'".format(new_categories) + ) + elif len(set(new_categories)) != len(set(self.categories)) or any( + cat not in self.categories for cat in new_categories + ): + raise ValueError("items in new_categories are not the same as in old categories") + + if ordered is None: + ordered = self.ordered + + if new_categories == list(self.categories) and ordered == self.ordered: + return self._data.copy() + else: + dtype = CategoricalDtype(categories=new_categories, ordered=ordered) + return _to_cat(self._data).astype(dtype) + + def set_categories( + self, + new_categories: Union[pd.Index, List], + ordered: Optional[bool] = None, + rename: bool = False, + ) -> Optional["ps.Series"]: + """ + Set the categories to the specified new_categories. + + `new_categories` can include new categories (which will result in + unused categories) or remove old categories (which results in values + set to NaN). If `rename==True`, the categories will simply be renamed + (less or more items than in old categories will result in values set to + NaN or in unused categories respectively). + + This method can be used to perform more than one action of adding, + removing, and reordering simultaneously and is therefore faster than + performing the individual steps via the more specialised methods. + + On the other hand this methods does not do checks (e.g., whether the + old categories are included in the new categories on a reorder), which + can result in surprising changes, for example when using special string + dtypes, which does not consider a S1 string equal to a single char + python string. + + Parameters + ---------- + new_categories : Index-like + The categories in new order. + ordered : bool, default False + Whether or not the categorical is treated as an ordered categorical. + If not given, do not change the ordered information. + rename : bool, default False + Whether or not the new_categories should be considered as a rename + of the old categories or as reordered categories. + + Returns + ------- + Series with reordered categories or None if inplace. + + Raises + ------ + ValueError + If new_categories does not validate as categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + + Examples + -------- + >>> s = ps.Series(list("abbccc"), dtype="category") + >>> s # doctest: +SKIP + 0 a + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (3, object): ['a', 'b', 'c'] + + >>> s.cat.set_categories(['b', 'c']) # doctest: +SKIP + 0 NaN + 1 b + 2 b + 3 c + 4 c + 5 c + dtype: category + Categories (2, object): ['b', 'c'] + + >>> s.cat.set_categories([1, 2, 3], rename=True) # doctest: +SKIP + 0 1 + 1 2 + 2 2 + 3 3 + 4 3 + 5 3 + dtype: category + Categories (3, int64): [1, 2, 3] + + >>> s.cat.set_categories([1, 2, 3], rename=True, ordered=True) # doctest: +SKIP + 0 1 + 1 2 + 2 2 + 3 3 + 4 3 + 5 3 + dtype: category + Categories (3, int64): [1 < 2 < 3] + """ + from pyspark.pandas.frame import DataFrame + + if not is_list_like(new_categories): + raise TypeError( + "Parameter 'new_categories' must be list-like, was '{}'".format(new_categories) + ) + + if ordered is None: + ordered = self.ordered + + new_dtype = CategoricalDtype(new_categories, ordered=ordered) + scol = self._data.spark.column + + if rename: + new_scol = ( + F.when(scol >= len(new_categories), F.lit(-1).cast(self._data.spark.data_type)) + .otherwise(scol) + .alias(self._data._internal.data_spark_column_names[0]) + ) + + internal = self._data._psdf._internal.with_new_spark_column( + self._data._column_label, + new_scol, + field=self._data._internal.data_fields[0].copy(dtype=new_dtype), + ) + + return DataFrame(internal)._psser_for(self._data._column_label).copy() + else: + return self._data.astype(new_dtype) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.categorical + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.categorical.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.categorical tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.categorical, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/config.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/config.py new file mode 100644 index 0000000..79cb859 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/config.py @@ -0,0 +1,539 @@ +# +# 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. +# + +""" +Infrastructure of options for pandas-on-Spark. +""" +from contextlib import contextmanager +import json +from typing import Any, Callable, Dict, Iterator, List, Tuple, Union + +from pyspark._globals import _NoValue, _NoValueType + +from pyspark.pandas.utils import default_session + + +__all__ = ["get_option", "set_option", "reset_option", "options", "option_context"] + + +class Option: + """ + Option class that defines an option with related properties. + + This class holds all information relevant to the one option. Also, + Its instance can validate if the given value is acceptable or not. + + It is currently for internal usage only. + + Parameters + ---------- + key: str, keyword-only argument + the option name to use. + doc: str, keyword-only argument + the documentation for the current option. + default: Any, keyword-only argument + default value for this option. + types: Union[Tuple[type, ...], type], keyword-only argument + default is str. It defines the expected types for this option. It is + used with `isinstance` to validate the given value to this option. + check_func: Tuple[Callable[[Any], bool], str], keyword-only argument + default is a function that always returns `True` with an empty string. + It defines: + - a function to check the given value to this option + - the error message to show when this check is failed + When new value is set to this option, this function is called to check + if the given value is valid. + + Examples + -------- + >>> option = Option( + ... key='option.name', + ... doc="this is a test option", + ... default="default", + ... types=(float, int), + ... check_func=(lambda v: v > 0, "should be a positive float")) + + >>> option.validate('abc') # doctest: +NORMALIZE_WHITESPACE + Traceback (most recent call last): + ... + TypeError: The value for option 'option.name' was ; + however, expected types are [(, )]. + + >>> option.validate(-1.1) + Traceback (most recent call last): + ... + ValueError: should be a positive float + + >>> option.validate(1.1) + """ + + def __init__( + self, + *, + key: str, + doc: str, + default: Any, + types: Union[Tuple[type, ...], type] = str, + check_func: Tuple[Callable[[Any], bool], str] = (lambda v: True, ""), + ): + self.key = key + self.doc = doc + self.default = default + self.types = types + self.check_func = check_func + + def validate(self, v: Any) -> None: + """ + Validate the given value and throw an exception with related information such as key. + """ + if not isinstance(v, self.types): + raise TypeError( + "The value for option '%s' was %s; however, expected types are " + "[%s]." % (self.key, type(v), str(self.types)) + ) + if not self.check_func[0](v): + raise ValueError(self.check_func[1]) + + +# Available options. +# +# NOTE: if you are fixing or adding an option here, make sure you execute `show_options()` and +# copy & paste the results into show_options +# 'docs/source/user_guide/pandas_on_spark/options.rst' as well. +# See the examples below: +# >>> from pyspark.pandas.config import show_options +# >>> show_options() +_options: List[Option] = [ + Option( + key="display.max_rows", + doc=( + "This sets the maximum number of rows pandas-on-Spark should output when printing out " + "various output. For example, this value determines the number of rows to be " + "shown at the repr() in a dataframe. Set `None` to unlimit the input length. " + "Default is 1000." + ), + default=1000, + types=(int, type(None)), + check_func=( + lambda v: v is None or v >= 0, + "'display.max_rows' should be greater than or equal to 0.", + ), + ), + Option( + key="compute.max_rows", + doc=( + "'compute.max_rows' sets the limit of the current pandas-on-Spark DataFrame. " + "Set `None` to unlimit the input length. When the limit is set, it is executed " + "by the shortcut by collecting the data into the driver, and then using the pandas " + "API. If the limit is unset, the operation is executed by PySpark. Default is 1000." + ), + default=1000, + types=(int, type(None)), + check_func=( + lambda v: v is None or v >= 0, + "'compute.max_rows' should be greater than or equal to 0.", + ), + ), + Option( + key="compute.shortcut_limit", + doc=( + "'compute.shortcut_limit' sets the limit for a shortcut. " + "It computes the specified number of rows and uses its schema. When the dataframe " + "length is larger than this limit, pandas-on-Spark uses PySpark to compute." + ), + default=1000, + types=int, + check_func=( + lambda v: v >= 0, + "'compute.shortcut_limit' should be greater than or equal to 0.", + ), + ), + Option( + key="compute.ops_on_diff_frames", + doc=( + "This determines whether or not to operate between two different dataframes. " + "For example, 'combine_frames' function internally performs a join operation which " + "can be expensive in general. So, if `compute.ops_on_diff_frames` variable is not " + "True, that method throws an exception." + ), + default=False, + types=bool, + ), + Option( + key="compute.default_index_type", + doc=("This sets the default index type: sequence, distributed and distributed-sequence."), + default="distributed-sequence", + types=str, + check_func=( + lambda v: v in ("sequence", "distributed", "distributed-sequence"), + "Index type should be one of 'sequence', 'distributed', 'distributed-sequence'.", + ), + ), + Option( + key="compute.default_index_cache", + doc=( + "This sets the default storage level for temporary RDDs cached in " + "distributed-sequence indexing: 'NONE', 'DISK_ONLY', 'DISK_ONLY_2', " + "'DISK_ONLY_3', 'MEMORY_ONLY', 'MEMORY_ONLY_2', 'MEMORY_ONLY_SER', " + "'MEMORY_ONLY_SER_2', 'MEMORY_AND_DISK', 'MEMORY_AND_DISK_2', " + "'MEMORY_AND_DISK_SER', 'MEMORY_AND_DISK_SER_2', 'OFF_HEAP', " + "'LOCAL_CHECKPOINT'." + ), + default="MEMORY_AND_DISK_SER", + types=str, + check_func=( + lambda v: v + in ( + "NONE", + "DISK_ONLY", + "DISK_ONLY_2", + "DISK_ONLY_3", + "MEMORY_ONLY", + "MEMORY_ONLY_2", + "MEMORY_ONLY_SER", + "MEMORY_ONLY_SER_2", + "MEMORY_AND_DISK", + "MEMORY_AND_DISK_2", + "MEMORY_AND_DISK_SER", + "MEMORY_AND_DISK_SER_2", + "OFF_HEAP", + "LOCAL_CHECKPOINT", + ), + "Index type should be one of 'NONE', 'DISK_ONLY', 'DISK_ONLY_2', " + "'DISK_ONLY_3', 'MEMORY_ONLY', 'MEMORY_ONLY_2', 'MEMORY_ONLY_SER', " + "'MEMORY_ONLY_SER_2', 'MEMORY_AND_DISK', 'MEMORY_AND_DISK_2', " + "'MEMORY_AND_DISK_SER', 'MEMORY_AND_DISK_SER_2', 'OFF_HEAP', " + "'LOCAL_CHECKPOINT'.", + ), + ), + Option( + key="compute.ordered_head", + doc=( + "'compute.ordered_head' sets whether or not to operate head with natural ordering. " + "pandas-on-Spark does not guarantee the row ordering so `head` could return some " + "rows from distributed partitions. If 'compute.ordered_head' is set to True, " + "pandas-on-Spark performs natural ordering beforehand, but it will cause a " + "performance overhead." + ), + default=False, + types=bool, + ), + Option( + key="compute.eager_check", + doc=( + "'compute.eager_check' sets whether or not to launch some Spark jobs just for the sake " + "of validation. If 'compute.eager_check' is set to True, pandas-on-Spark performs the " + "validation beforehand, but it will cause a performance overhead. Otherwise, " + "pandas-on-Spark skip the validation and will be slightly different from pandas. " + "Affected APIs: `Series.dot`, `Series.asof`, `Series.compare`, " + "`FractionalExtensionOps.astype`, `IntegralExtensionOps.astype`, " + "`FractionalOps.astype`, `DecimalOps.astype`, `skipna of statistical functions`." + ), + default=True, + types=bool, + ), + Option( + key="compute.isin_limit", + doc=( + "'compute.isin_limit' sets the limit for filtering by 'Column.isin(list)'. " + "If the length of the ‘list’ is above the limit, broadcast join is used instead " + "for better performance." + ), + default=80, + types=int, + check_func=( + lambda v: v >= 0, + "'compute.isin_limit' should be greater than or equal to 0.", + ), + ), + Option( + key="plotting.max_rows", + doc=( + "'plotting.max_rows' sets the visual limit on top-n-based plots such as `plot.bar` " + "and `plot.pie`. If it is set to 1000, the first 1000 data points will be used " + "for plotting. Default is 1000." + ), + default=1000, + types=int, + check_func=( + lambda v: v >= 0, + "'plotting.max_rows' should be greater than or equal to 0.", + ), + ), + Option( + key="plotting.sample_ratio", + doc=( + "'plotting.sample_ratio' sets the proportion of data that will be plotted for sample-" + "based plots such as `plot.line` and `plot.area`. " + "This option defaults to 'plotting.max_rows' option." + ), + default=None, + types=(float, type(None)), + check_func=( + lambda v: v is None or 1 >= v >= 0, + "'plotting.sample_ratio' should be 1.0 >= value >= 0.0.", + ), + ), + Option( + key="plotting.backend", + doc=( + "Backend to use for plotting. Default is plotly. " + "Supports any package that has a top-level `.plot` method. " + "Known options are: [matplotlib, plotly]." + ), + default="plotly", + types=str, + ), +] + +_options_dict: Dict[str, Option] = dict(zip((option.key for option in _options), _options)) + +_key_format = "pandas_on_Spark.{}".format + + +class OptionError(AttributeError, KeyError): + pass + + +def show_options() -> None: + """ + Make a pretty table that can be copied and pasted into public documentation. + This is currently for an internal purpose. + + Examples + -------- + >>> show_options() # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + ================... =======... =====================... + Option Default Description + ================... =======... =====================... + display.max_rows 1000 This sets the maximum... + ... + ================... =======... =====================... + """ + + import textwrap + + header = ["Option", "Default", "Description"] + row_format = "{:<31} {:<23} {:<53}" + + print(row_format.format("=" * 31, "=" * 23, "=" * 53)) + print(row_format.format(*header)) + print(row_format.format("=" * 31, "=" * 23, "=" * 53)) + + for option in _options: + doc = textwrap.fill(option.doc, 53) + formatted = "".join([line + "\n" + (" " * 56) for line in doc.split("\n")]).rstrip() + print(row_format.format(option.key, repr(option.default), formatted)) + + print(row_format.format("=" * 31, "=" * 23, "=" * 53)) + + +def get_option(key: str, default: Union[Any, _NoValueType] = _NoValue) -> Any: + """ + Retrieves the value of the specified option. + + Parameters + ---------- + key : str + The key which should match a single option. + default : object + The default value if the option is not set yet. The value should be JSON serializable. + + Returns + ------- + result : the value of the option + + Raises + ------ + OptionError : if no such option exists and the default is not provided + """ + _check_option(key) + if default is _NoValue: + default = _options_dict[key].default + _options_dict[key].validate(default) + spark_session = default_session() + + return json.loads(spark_session.conf.get(_key_format(key), default=json.dumps(default))) + + +def set_option(key: str, value: Any) -> None: + """ + Sets the value of the specified option. + + Parameters + ---------- + key : str + The key which should match a single option. + value : object + New value of option. The value should be JSON serializable. + + Returns + ------- + None + """ + _check_option(key) + _options_dict[key].validate(value) + spark_session = default_session() + + spark_session.conf.set(_key_format(key), json.dumps(value)) + + +def reset_option(key: str) -> None: + """ + Reset one option to their default value. + + Pass "all" as an argument to reset all options. + + Parameters + ---------- + key : str + If specified only option will be reset. + + Returns + ------- + None + """ + _check_option(key) + default_session().conf.unset(_key_format(key)) + + +@contextmanager +def option_context(*args: Any) -> Iterator[None]: + """ + Context manager to temporarily set options in the `with` statement context. + + You need to invoke ``option_context(pat, val, [(pat, val), ...])``. + + Examples + -------- + >>> with option_context('display.max_rows', 10, 'compute.max_rows', 5): + ... print(get_option('display.max_rows'), get_option('compute.max_rows')) + 10 5 + >>> print(get_option('display.max_rows'), get_option('compute.max_rows')) + 1000 1000 + """ + if len(args) == 0 or len(args) % 2 != 0: + raise ValueError("Need to invoke as option_context(pat, val, [(pat, val), ...]).") + opts = dict(zip(args[::2], args[1::2])) + orig_opts = {key: get_option(key) for key in opts} + try: + for key, value in opts.items(): + set_option(key, value) + yield + finally: + for key, value in orig_opts.items(): + set_option(key, value) + + +def _check_option(key: str) -> None: + if key not in _options_dict: + raise OptionError( + "No such option: '{}'. Available options are [{}]".format( + key, ", ".join(list(_options_dict.keys())) + ) + ) + + +class DictWrapper: + """provide attribute-style access to a nested dict""" + + def __init__(self, d: Dict[str, Option], prefix: str = ""): + object.__setattr__(self, "d", d) + object.__setattr__(self, "prefix", prefix) + + def __setattr__(self, key: str, val: Any) -> None: + prefix = object.__getattribute__(self, "prefix") + d = object.__getattribute__(self, "d") + if prefix: + prefix += "." + canonical_key = prefix + key + + candidates = [ + k for k in d.keys() if all(x in k.split(".") for x in canonical_key.split(".")) + ] + if len(candidates) == 1 and candidates[0] == canonical_key: + set_option(canonical_key, val) + else: + raise OptionError( + "No such option: '{}'. Available options are [{}]".format( + key, ", ".join(list(_options_dict.keys())) + ) + ) + + def __getattr__(self, key: str) -> Union["DictWrapper", Any]: + prefix = object.__getattribute__(self, "prefix") + d = object.__getattribute__(self, "d") + if prefix: + prefix += "." + canonical_key = prefix + key + + candidates = [ + k for k in d.keys() if all(x in k.split(".") for x in canonical_key.split(".")) + ] + if len(candidates) == 1 and candidates[0] == canonical_key: + return get_option(canonical_key) + elif len(candidates) == 0: + raise OptionError( + "No such option: '{}'. Available options are [{}]".format( + key, ", ".join(list(_options_dict.keys())) + ) + ) + else: + return DictWrapper(d, canonical_key) + + def __dir__(self) -> List[str]: + prefix = object.__getattribute__(self, "prefix") + d = object.__getattribute__(self, "d") + + if prefix == "": + candidates = d.keys() + offset = 0 + else: + candidates = [k for k in d.keys() if all(x in k.split(".") for x in prefix.split("."))] + offset = len(prefix) + 1 # prefix (e.g. "compute.") to trim. + return [c[offset:] for c in candidates] + + +options = DictWrapper(_options_dict) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.config + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.config.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.config tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.config, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/correlation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/correlation.py new file mode 100644 index 0000000..75d3a85 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/correlation.py @@ -0,0 +1,262 @@ +# +# 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 List + +from pyspark.sql import DataFrame as SparkDataFrame, functions as F +from pyspark.sql.window import Window + +from pyspark.pandas.utils import verify_temp_column_name + + +CORRELATION_VALUE_1_COLUMN = "__correlation_value_1_input__" +CORRELATION_VALUE_2_COLUMN = "__correlation_value_2_input__" +CORRELATION_CORR_OUTPUT_COLUMN = "__correlation_corr_output__" +CORRELATION_COUNT_OUTPUT_COLUMN = "__correlation_count_output__" + + +def compute(sdf: SparkDataFrame, groupKeys: List[str], method: str) -> SparkDataFrame: + """ + Compute correlation per group, excluding NA/null values. + + Input PySpark Dataframe should contain column `CORRELATION_VALUE_1_COLUMN` and + column `CORRELATION_VALUE_2_COLUMN`, as well as the group columns. + + The returned PySpark Dataframe will contain the correlation column + `CORRELATION_CORR_OUTPUT_COLUMN` and the non-null count column + `CORRELATION_COUNT_OUTPUT_COLUMN`, as well as the group columns. + """ + assert len(groupKeys) > 0 + assert method in ["pearson", "spearman", "kendall"] + + sdf = sdf.select( + *[F.col(key) for key in groupKeys], + *[ + # assign both columns nulls, if some of them are null + F.when( + F.isnull(CORRELATION_VALUE_1_COLUMN) | F.isnull(CORRELATION_VALUE_2_COLUMN), + F.lit(None), + ) + .otherwise(F.col(CORRELATION_VALUE_1_COLUMN)) + .alias(CORRELATION_VALUE_1_COLUMN), + F.when( + F.isnull(CORRELATION_VALUE_1_COLUMN) | F.isnull(CORRELATION_VALUE_2_COLUMN), + F.lit(None), + ) + .otherwise(F.col(CORRELATION_VALUE_2_COLUMN)) + .alias(CORRELATION_VALUE_2_COLUMN), + ], + ) + + if method in ["pearson", "spearman"]: + # convert values to avg ranks for spearman correlation + if method == "spearman": + ROW_NUMBER_COLUMN = verify_temp_column_name( + sdf, "__correlation_spearman_row_number_temp_column__" + ) + DENSE_RANK_COLUMN = verify_temp_column_name( + sdf, "__correlation_spearman_dense_rank_temp_column__" + ) + window = Window.partitionBy(groupKeys) + + # CORRELATION_VALUE_1_COLUMN: value -> avg rank + # for example: + # values: 3, 4, 5, 7, 7, 7, 9, 9, 10 + # avg ranks: 1.0, 2.0, 3.0, 5.0, 5.0, 5.0, 7.5, 7.5, 9.0 + sdf = ( + sdf.withColumn( + ROW_NUMBER_COLUMN, + F.row_number().over( + window.orderBy(F.asc_nulls_last(CORRELATION_VALUE_1_COLUMN)) + ), + ) + # drop nulls but make sure each group contains at least one row + .where(~F.isnull(CORRELATION_VALUE_1_COLUMN) | (F.col(ROW_NUMBER_COLUMN) == 1)) + .withColumn( + DENSE_RANK_COLUMN, + F.dense_rank().over( + window.orderBy(F.asc_nulls_last(CORRELATION_VALUE_1_COLUMN)) + ), + ) + .withColumn( + CORRELATION_VALUE_1_COLUMN, + F.when(F.isnull(CORRELATION_VALUE_1_COLUMN), F.lit(None)).otherwise( + F.avg(ROW_NUMBER_COLUMN).over( + window.orderBy(F.asc(DENSE_RANK_COLUMN)).rangeBetween(0, 0) + ) + ), + ) + ) + + # CORRELATION_VALUE_2_COLUMN: value -> avg rank + sdf = ( + sdf.withColumn( + ROW_NUMBER_COLUMN, + F.row_number().over( + window.orderBy(F.asc_nulls_last(CORRELATION_VALUE_2_COLUMN)) + ), + ) + .withColumn( + DENSE_RANK_COLUMN, + F.dense_rank().over( + window.orderBy(F.asc_nulls_last(CORRELATION_VALUE_2_COLUMN)) + ), + ) + .withColumn( + CORRELATION_VALUE_2_COLUMN, + F.when(F.isnull(CORRELATION_VALUE_2_COLUMN), F.lit(None)).otherwise( + F.avg(ROW_NUMBER_COLUMN).over( + window.orderBy(F.asc(DENSE_RANK_COLUMN)).rangeBetween(0, 0) + ) + ), + ) + ) + + sdf = sdf.groupby(groupKeys).agg( + F.corr(CORRELATION_VALUE_1_COLUMN, CORRELATION_VALUE_2_COLUMN).alias( + CORRELATION_CORR_OUTPUT_COLUMN + ), + F.count( + F.when( + ~F.isnull(CORRELATION_VALUE_1_COLUMN), + 1, + ) + ).alias(CORRELATION_COUNT_OUTPUT_COLUMN), + ) + + return sdf + + else: + # kendall correlation + ROW_NUMBER_1_2_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_row_number_1_2_temp_column__" + ) + sdf = sdf.withColumn( + ROW_NUMBER_1_2_COLUMN, + F.row_number().over( + Window.partitionBy(groupKeys).orderBy( + F.asc_nulls_last(CORRELATION_VALUE_1_COLUMN), + F.asc_nulls_last(CORRELATION_VALUE_2_COLUMN), + ) + ), + ) + + # drop nulls but make sure each group contains at least one row + sdf = sdf.where(~F.isnull(CORRELATION_VALUE_1_COLUMN) | (F.col(ROW_NUMBER_1_2_COLUMN) == 1)) + + CORRELATION_VALUE_X_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_value_x_temp_column__" + ) + CORRELATION_VALUE_Y_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_value_y_temp_column__" + ) + ROW_NUMBER_X_Y_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_row_number_x_y_temp_column__" + ) + sdf2 = sdf.select( + *[F.col(key) for key in groupKeys], + *[ + F.col(CORRELATION_VALUE_1_COLUMN).alias(CORRELATION_VALUE_X_COLUMN), + F.col(CORRELATION_VALUE_2_COLUMN).alias(CORRELATION_VALUE_Y_COLUMN), + F.col(ROW_NUMBER_1_2_COLUMN).alias(ROW_NUMBER_X_Y_COLUMN), + ], + ) + + sdf = sdf.join(sdf2, groupKeys, "inner").where( + F.col(ROW_NUMBER_1_2_COLUMN) <= F.col(ROW_NUMBER_X_Y_COLUMN) + ) + + # compute P, Q, T, U in tau_b = (P - Q) / sqrt((P + Q + T) * (P + Q + U)) + # see https://github.com/scipy/scipy/blob/v1.9.1/scipy/stats/_stats_py.py#L5015-L5222 + CORRELATION_KENDALL_P_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_tau_b_p_temp_column__" + ) + CORRELATION_KENDALL_Q_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_tau_b_q_temp_column__" + ) + CORRELATION_KENDALL_T_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_tau_b_t_temp_column__" + ) + CORRELATION_KENDALL_U_COLUMN = verify_temp_column_name( + sdf, "__correlation_kendall_tau_b_u_temp_column__" + ) + + pair_cond = ~F.isnull(CORRELATION_VALUE_1_COLUMN) & ( + F.col(ROW_NUMBER_1_2_COLUMN) < F.col(ROW_NUMBER_X_Y_COLUMN) + ) + + p_cond = ( + (F.col(CORRELATION_VALUE_1_COLUMN) < F.col(CORRELATION_VALUE_X_COLUMN)) + & (F.col(CORRELATION_VALUE_2_COLUMN) < F.col(CORRELATION_VALUE_Y_COLUMN)) + ) | ( + (F.col(CORRELATION_VALUE_1_COLUMN) > F.col(CORRELATION_VALUE_X_COLUMN)) + & (F.col(CORRELATION_VALUE_2_COLUMN) > F.col(CORRELATION_VALUE_Y_COLUMN)) + ) + q_cond = ( + (F.col(CORRELATION_VALUE_1_COLUMN) < F.col(CORRELATION_VALUE_X_COLUMN)) + & (F.col(CORRELATION_VALUE_2_COLUMN) > F.col(CORRELATION_VALUE_Y_COLUMN)) + ) | ( + (F.col(CORRELATION_VALUE_1_COLUMN) > F.col(CORRELATION_VALUE_X_COLUMN)) + & (F.col(CORRELATION_VALUE_2_COLUMN) < F.col(CORRELATION_VALUE_Y_COLUMN)) + ) + t_cond = (F.col(CORRELATION_VALUE_1_COLUMN) == F.col(CORRELATION_VALUE_X_COLUMN)) & ( + F.col(CORRELATION_VALUE_2_COLUMN) != F.col(CORRELATION_VALUE_Y_COLUMN) + ) + u_cond = (F.col(CORRELATION_VALUE_1_COLUMN) != F.col(CORRELATION_VALUE_X_COLUMN)) & ( + F.col(CORRELATION_VALUE_2_COLUMN) == F.col(CORRELATION_VALUE_Y_COLUMN) + ) + + sdf = ( + sdf.groupby(groupKeys) + .agg( + F.count(F.when(pair_cond & p_cond, 1)).alias(CORRELATION_KENDALL_P_COLUMN), + F.count(F.when(pair_cond & q_cond, 1)).alias(CORRELATION_KENDALL_Q_COLUMN), + F.count(F.when(pair_cond & t_cond, 1)).alias(CORRELATION_KENDALL_T_COLUMN), + F.count(F.when(pair_cond & u_cond, 1)).alias(CORRELATION_KENDALL_U_COLUMN), + F.max( + F.when( + ~F.isnull(CORRELATION_VALUE_1_COLUMN), F.col(ROW_NUMBER_X_Y_COLUMN) + ).otherwise(F.lit(0)) + ).alias(CORRELATION_COUNT_OUTPUT_COLUMN), + ) + .withColumn( + CORRELATION_CORR_OUTPUT_COLUMN, + (F.col(CORRELATION_KENDALL_P_COLUMN) - F.col(CORRELATION_KENDALL_Q_COLUMN)) + / F.sqrt( + ( + ( + F.col(CORRELATION_KENDALL_P_COLUMN) + + F.col(CORRELATION_KENDALL_Q_COLUMN) + + (F.col(CORRELATION_KENDALL_T_COLUMN)) + ) + ) + * ( + ( + F.col(CORRELATION_KENDALL_P_COLUMN) + + F.col(CORRELATION_KENDALL_Q_COLUMN) + + (F.col(CORRELATION_KENDALL_U_COLUMN)) + ) + ) + ), + ) + ) + + sdf = sdf.select( + *[F.col(key) for key in groupKeys], + *[CORRELATION_CORR_OUTPUT_COLUMN, CORRELATION_COUNT_OUTPUT_COLUMN], + ) + return sdf diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/base.py new file mode 100644 index 0000000..5d497a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/base.py @@ -0,0 +1,519 @@ +# +# 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 numbers +from abc import ABCMeta +from itertools import chain +from typing import Any, Optional, Union + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark.sql import functions as F +from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + DataType, + DateType, + DayTimeIntervalType, + DecimalType, + FractionalType, + IntegralType, + MapType, + NullType, + NumericType, + StringType, + StructType, + TimestampType, + TimestampNTZType, + UserDefinedType, +) +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.typedef import extension_dtypes +from pyspark.pandas.typedef.typehints import ( + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, + spark_type_to_pandas_dtype, +) + +# For supporting Spark Connect +from pyspark.sql.utils import get_column_class + +if extension_dtypes_available: + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + +if extension_float_dtypes_available: + from pandas import Float32Dtype, Float64Dtype + +if extension_object_dtypes_available: + from pandas import BooleanDtype, StringDtype + + +def is_valid_operand_for_numeric_arithmetic(operand: Any, *, allow_bool: bool = True) -> bool: + """Check whether the `operand` is valid for arithmetic operations against numerics.""" + from pyspark.pandas.base import IndexOpsMixin + + if isinstance(operand, numbers.Number): + return not isinstance(operand, bool) or allow_bool + elif isinstance(operand, IndexOpsMixin): + if isinstance(operand.dtype, CategoricalDtype): + return False + else: + return isinstance(operand.spark.data_type, NumericType) or ( + allow_bool and isinstance(operand.spark.data_type, BooleanType) + ) + else: + return False + + +def transform_boolean_operand_to_numeric( + operand: Any, *, spark_type: Optional[DataType] = None +) -> Any: + """Transform boolean operand to numeric. + + If the `operand` is: + - a boolean IndexOpsMixin, transform the `operand` to the `spark_type`. + - a boolean literal, transform to the int value. + Otherwise, return the operand as it is. + """ + from pyspark.pandas.base import IndexOpsMixin + + if isinstance(operand, IndexOpsMixin) and isinstance(operand.spark.data_type, BooleanType): + assert spark_type, "spark_type must be provided if the operand is a boolean IndexOpsMixin" + assert isinstance(spark_type, NumericType), "spark_type must be NumericType" + dtype = spark_type_to_pandas_dtype( + spark_type, use_extension_dtypes=operand._internal.data_fields[0].is_extension_dtype + ) + return operand._with_new_scol( + operand.spark.column.cast(spark_type), + field=operand._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type), + ) + elif isinstance(operand, bool): + return int(operand) + else: + return operand + + +def _as_categorical_type( + index_ops: IndexOpsLike, dtype: CategoricalDtype, spark_type: DataType +) -> IndexOpsLike: + """Cast `index_ops` to categorical dtype, given `dtype` and `spark_type`.""" + assert isinstance(dtype, CategoricalDtype) + if dtype.categories is None: + codes, uniques = index_ops.factorize() + categories = uniques.astype(index_ops.dtype) + return codes._with_new_scol( + codes.spark.column, + field=codes._internal.data_fields[0].copy( + dtype=CategoricalDtype(categories=categories) + ), + ) + else: + categories = dtype.categories + if len(categories) == 0: + scol = F.lit(-1) + else: + kvs = chain( + *[(F.lit(category), F.lit(code)) for code, category in enumerate(categories)] + ) + map_scol = F.create_map(*kvs) + + scol = F.coalesce(map_scol[index_ops.spark.column], F.lit(-1)) + return index_ops._with_new_scol( + scol.cast(spark_type), + field=index_ops._internal.data_fields[0].copy( + dtype=dtype, spark_type=spark_type, nullable=False + ), + ) + + +def _as_bool_type(index_ops: IndexOpsLike, dtype: Dtype) -> IndexOpsLike: + """Cast `index_ops` to BooleanType Spark type, given `dtype`.""" + spark_type = BooleanType() + if isinstance(dtype, extension_dtypes): + scol = index_ops.spark.column.cast(spark_type) + else: + scol = F.when(index_ops.spark.column.isNull(), F.lit(False)).otherwise( + index_ops.spark.column.cast(spark_type) + ) + return index_ops._with_new_scol( + scol, field=index_ops._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type) + ) + + +def _as_string_type( + index_ops: IndexOpsLike, dtype: Dtype, *, null_str: str = str(None) +) -> IndexOpsLike: + """Cast `index_ops` to StringType Spark type, given `dtype` and `null_str`, + representing null Spark column. Note that `null_str` is for non-extension dtypes only. + """ + spark_type = StringType() + if isinstance(dtype, extension_dtypes): + scol = index_ops.spark.column.cast(spark_type) + else: + casted = index_ops.spark.column.cast(spark_type) + scol = F.when(index_ops.spark.column.isNull(), null_str).otherwise(casted) + return index_ops._with_new_scol( + scol, field=index_ops._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type) + ) + + +def _as_other_type(index_ops: IndexOpsLike, dtype: Dtype, spark_type: DataType) -> IndexOpsLike: + """Cast `index_ops` to a `dtype` (`spark_type`) that needs no pre-processing. + + Destination types that need pre-processing: CategoricalDtype, BooleanType, and StringType. + """ + from pyspark.pandas.internal import InternalField + + need_pre_process = ( + isinstance(dtype, CategoricalDtype) + or isinstance(spark_type, BooleanType) + or isinstance(spark_type, StringType) + ) + assert not need_pre_process, "Pre-processing is needed before the type casting." + + scol = index_ops.spark.column.cast(spark_type) + return index_ops._with_new_scol(scol, field=InternalField(dtype=dtype)) + + +def _sanitize_list_like(operand: Any) -> None: + """Raise TypeError if operand is list-like.""" + if isinstance(operand, (list, tuple, dict, set)): + raise TypeError("The operation can not be applied to %s." % type(operand).__name__) + + +def _is_valid_for_logical_operator(right: Any) -> bool: + from pyspark.pandas.base import IndexOpsMixin + + return isinstance(right, (int, bool)) or ( + isinstance(right, IndexOpsMixin) + and ( + isinstance(right.spark.data_type, BooleanType) + or isinstance(right.spark.data_type, IntegralType) + ) + ) + + +def _is_boolean_type(right: Any) -> bool: + from pyspark.pandas.base import IndexOpsMixin + + return isinstance(right, bool) or ( + isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, BooleanType) + ) + + +def _is_extension_dtypes(object: Any) -> bool: + """ + Check whether the type of given object is extension dtype or not. + Extention dtype includes Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, BooleanDtype, + StringDtype, Float32Dtype and Float64Dtype. + """ + return isinstance(getattr(object, "dtype", None), extension_dtypes) + + +class DataTypeOps(object, metaclass=ABCMeta): + """The base class for binary operations of pandas-on-Spark objects (of different data types).""" + + def __new__(cls, dtype: Dtype, spark_type: DataType) -> "DataTypeOps": + from pyspark.pandas.data_type_ops.binary_ops import BinaryOps + from pyspark.pandas.data_type_ops.boolean_ops import BooleanOps, BooleanExtensionOps + from pyspark.pandas.data_type_ops.categorical_ops import CategoricalOps + from pyspark.pandas.data_type_ops.complex_ops import ArrayOps, MapOps, StructOps + from pyspark.pandas.data_type_ops.date_ops import DateOps + from pyspark.pandas.data_type_ops.datetime_ops import DatetimeOps, DatetimeNTZOps + from pyspark.pandas.data_type_ops.null_ops import NullOps + from pyspark.pandas.data_type_ops.num_ops import ( + DecimalOps, + FractionalExtensionOps, + FractionalOps, + IntegralExtensionOps, + IntegralOps, + ) + from pyspark.pandas.data_type_ops.string_ops import StringOps, StringExtensionOps + from pyspark.pandas.data_type_ops.timedelta_ops import TimedeltaOps + from pyspark.pandas.data_type_ops.udt_ops import UDTOps + + if isinstance(dtype, CategoricalDtype): + return object.__new__(CategoricalOps) + elif isinstance(spark_type, DecimalType): + return object.__new__(DecimalOps) + elif isinstance(spark_type, FractionalType): + if extension_float_dtypes_available and type(dtype) in [Float32Dtype, Float64Dtype]: + return object.__new__(FractionalExtensionOps) + else: + return object.__new__(FractionalOps) + elif isinstance(spark_type, IntegralType): + if extension_dtypes_available and type(dtype) in [ + Int8Dtype, + Int16Dtype, + Int32Dtype, + Int64Dtype, + ]: + return object.__new__(IntegralExtensionOps) + else: + return object.__new__(IntegralOps) + elif isinstance(spark_type, StringType): + if extension_object_dtypes_available and isinstance(dtype, StringDtype): + return object.__new__(StringExtensionOps) + else: + return object.__new__(StringOps) + elif isinstance(spark_type, BooleanType): + if extension_object_dtypes_available and isinstance(dtype, BooleanDtype): + return object.__new__(BooleanExtensionOps) + else: + return object.__new__(BooleanOps) + elif isinstance(spark_type, TimestampType): + return object.__new__(DatetimeOps) + elif isinstance(spark_type, TimestampNTZType): + return object.__new__(DatetimeNTZOps) + elif isinstance(spark_type, DateType): + return object.__new__(DateOps) + elif isinstance(spark_type, DayTimeIntervalType): + return object.__new__(TimedeltaOps) + elif isinstance(spark_type, BinaryType): + return object.__new__(BinaryOps) + elif isinstance(spark_type, ArrayType): + return object.__new__(ArrayOps) + elif isinstance(spark_type, MapType): + return object.__new__(MapOps) + elif isinstance(spark_type, StructType): + return object.__new__(StructOps) + elif isinstance(spark_type, NullType): + return object.__new__(NullOps) + elif isinstance(spark_type, UserDefinedType): + return object.__new__(UDTOps) + else: + raise TypeError("Type %s was not understood." % dtype) + + def __init__(self, dtype: Dtype, spark_type: DataType): + self.dtype = dtype + self.spark_type = spark_type + + @property + def pretty_name(self) -> str: + raise NotImplementedError() + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Addition can not be applied to %s." % self.pretty_name) + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Subtraction can not be applied to %s." % self.pretty_name) + + def mul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Multiplication can not be applied to %s." % self.pretty_name) + + def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("True division can not be applied to %s." % self.pretty_name) + + def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Floor division can not be applied to %s." % self.pretty_name) + + def mod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Modulo can not be applied to %s." % self.pretty_name) + + def pow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Exponentiation can not be applied to %s." % self.pretty_name) + + def radd(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Addition can not be applied to %s." % self.pretty_name) + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Subtraction can not be applied to %s." % self.pretty_name) + + def rmul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Multiplication can not be applied to %s." % self.pretty_name) + + def rtruediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("True division can not be applied to %s." % self.pretty_name) + + def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Floor division can not be applied to %s." % self.pretty_name) + + def rmod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Modulo can not be applied to %s." % self.pretty_name) + + def rpow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Exponentiation can not be applied to %s." % self.pretty_name) + + def __and__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Bitwise and can not be applied to %s." % self.pretty_name) + + def xor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Bitwise xor can not be applied to %s." % self.pretty_name) + + def __or__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("Bitwise or can not be applied to %s." % self.pretty_name) + + def rand(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return left.__and__(right) + + def rxor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return left ^ right + + def ror(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return left.__or__(right) + + def neg(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("Unary - can not be applied to %s." % self.pretty_name) + + def abs(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("abs() can not be applied to %s." % self.pretty_name) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("< can not be applied to %s." % self.pretty_name) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("<= can not be applied to %s." % self.pretty_name) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("> can not be applied to %s." % self.pretty_name) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError(">= can not be applied to %s." % self.pretty_name) + + def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + if isinstance(right, (list, tuple)): + from pyspark.pandas.series import first_series, scol_for + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.internal import NATURAL_ORDER_COLUMN_NAME, InternalField + + if len(left) != len(right): + raise ValueError("Lengths must be equal") + + sdf = left._internal.spark_frame + structed_scol = F.struct( + sdf[NATURAL_ORDER_COLUMN_NAME], + *left._internal.index_spark_columns, + left.spark.column, + ) + # The size of the list is expected to be small. + collected_structed_scol = F.collect_list(structed_scol) + # Sort the array by NATURAL_ORDER_COLUMN so that we can guarantee the order. + collected_structed_scol = F.array_sort(collected_structed_scol) + right_values_scol = F.array(*(F.lit(x) for x in right)) + index_scol_names = left._internal.index_spark_column_names + scol_name = left._internal.spark_column_name_for(left._internal.column_labels[0]) + # Compare the values of left and right by using zip_with function. + cond = F.zip_with( + collected_structed_scol, + right_values_scol, + lambda x, y: F.struct( + *[ + x[index_scol_name].alias(index_scol_name) + for index_scol_name in index_scol_names + ], + F.when(x[scol_name].isNull() | y.isNull(), False) + .otherwise( + x[scol_name] == y, + ) + .alias(scol_name), + ), + ).alias(scol_name) + # 1. `sdf_new` here looks like the below (the first field of each set is Index): + # +----------------------------------------------------------+ + # |0 | + # +----------------------------------------------------------+ + # |[{0, false}, {1, true}, {2, false}, {3, true}, {4, false}]| + # +----------------------------------------------------------+ + sdf_new = sdf.select(cond) + # 2. `sdf_new` after the explode looks like the below: + # +----------+ + # | col| + # +----------+ + # |{0, false}| + # | {1, true}| + # |{2, false}| + # | {3, true}| + # |{4, false}| + # +----------+ + sdf_new = sdf_new.select(F.explode(scol_name)) + # 3. Here, the final `sdf_new` looks like the below: + # +-----------------+-----+ + # |__index_level_0__| 0| + # +-----------------+-----+ + # | 0|false| + # | 1| true| + # | 2|false| + # | 3| true| + # | 4|false| + # +-----------------+-----+ + sdf_new = sdf_new.select("col.*") + + index_spark_columns = [ + scol_for(sdf_new, index_scol_name) for index_scol_name in index_scol_names + ] + data_spark_columns = [scol_for(sdf_new, scol_name)] + + internal = left._internal.copy( + spark_frame=sdf_new, + index_spark_columns=index_spark_columns, + data_spark_columns=data_spark_columns, + index_fields=[ + InternalField.from_struct_field(index_field) + for index_field in sdf_new.select(index_spark_columns).schema.fields + ], + data_fields=[ + InternalField.from_struct_field( + sdf_new.select(data_spark_columns).schema.fields[0] + ) + ], + ) + return first_series(DataFrame(internal)) + else: + from pyspark.pandas.base import column_op + + Column = get_column_class() + return column_op(Column.__eq__)(left, right) + + def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + + Column = get_column_class() + return column_op(Column.__ne__)(left, right) + + def invert(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("Unary ~ can not be applied to %s." % self.pretty_name) + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return col + + def prepare(self, col: pd.Series) -> pd.Series: + """Prepare column when from_pandas.""" + return col.replace({np.nan: None}) + + def isnull(self, index_ops: IndexOpsLike) -> IndexOpsLike: + return index_ops._with_new_scol( + index_ops.spark.column.isNull(), + field=index_ops._internal.data_fields[0].copy( + dtype=np.dtype("bool"), spark_type=BooleanType(), nullable=False + ), + ) + + def nan_to_null(self, index_ops: IndexOpsLike) -> IndexOpsLike: + return index_ops.copy() + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + raise TypeError("astype can not be applied to %s." % self.pretty_name) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/binary_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/binary_ops.py new file mode 100644 index 0000000..f528d3e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/binary_ops.py @@ -0,0 +1,98 @@ +# +# 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 Any, Union, cast + +from pandas.api.types import CategoricalDtype + +from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql import functions as F +from pyspark.sql.types import BinaryType, BooleanType, StringType +from pyspark.sql.utils import pyspark_column_op + + +class BinaryOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with BinaryType. + """ + + @property + def pretty_name(self) -> str: + return "binaries" + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, BinaryType): + return column_op(F.concat)(left, right) + elif isinstance(right, bytes): + return column_op(F.concat)(left, F.lit(right)) + else: + raise TypeError( + "Concatenation can not be applied to %s and the given type." % self.pretty_name + ) + + def radd(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if isinstance(right, bytes): + return cast( + SeriesOrIndex, left._with_new_scol(F.concat(F.lit(right), left.spark.column)) + ) + else: + raise TypeError( + "Concatenation can not be applied to %s and the given type." % self.pretty_name + ) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + # Cannot cast binary to boolean in Spark. + # We should cast binary to str first, and cast it to boolean + return index_ops.astype(str).astype(bool) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype) + else: + return _as_other_type(index_ops, dtype, spark_type) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/boolean_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/boolean_ops.py new file mode 100644 index 0000000..11f376d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/boolean_ops.py @@ -0,0 +1,426 @@ +# +# 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 numbers +from typing import Any, Union + +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + is_valid_operand_for_numeric_arithmetic, + transform_boolean_operand_to_numeric, + _as_bool_type, + _as_categorical_type, + _as_other_type, + _sanitize_list_like, + _is_valid_for_logical_operator, + _is_boolean_type, +) +from pyspark.pandas.typedef.typehints import as_spark_type, extension_dtypes, pandas_on_spark_type +from pyspark.sql import functions as F +from pyspark.sql.column import Column as PySparkColumn +from pyspark.sql.types import BooleanType, StringType +from pyspark.sql.utils import get_column_class +from pyspark.errors import PySparkValueError + + +class BooleanOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: BooleanType. + """ + + @property + def pretty_name(self) -> str: + return "bools" + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError( + "Addition can not be applied to %s and the given type." % self.pretty_name + ) + + if isinstance(right, bool): + return left.__or__(right) + elif isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left + right + else: + assert isinstance(right, IndexOpsMixin) + if isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, BooleanType): + return left.__or__(right) + else: + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left + right + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right, allow_bool=False): + raise TypeError( + "Subtraction can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left - right + else: + assert isinstance(right, IndexOpsMixin) + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left - right + + def mul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError( + "Multiplication can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, bool): + return left.__and__(right) + elif isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left * right + else: + assert isinstance(right, IndexOpsMixin) + if isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, BooleanType): + return left.__and__(right) + else: + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left * right + + def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right, allow_bool=False): + raise TypeError( + "True division can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left / right + else: + assert isinstance(right, IndexOpsMixin) + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left / right + + def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right, allow_bool=False): + raise TypeError( + "Floor division can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left // right + else: + assert isinstance(right, IndexOpsMixin) + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left // right + + def mod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right, allow_bool=False): + raise TypeError( + "Modulo can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left % right + else: + assert isinstance(right, IndexOpsMixin) + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left % right + + def pow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right, allow_bool=False): + raise TypeError( + "Exponentiation can not be applied to %s and the given type." % self.pretty_name + ) + if isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return left**right + else: + assert isinstance(right, IndexOpsMixin) + left = transform_boolean_operand_to_numeric(left, spark_type=right.spark.data_type) + return left**right + + def radd(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, bool): + return left.__or__(right) + elif isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right + left + else: + raise TypeError( + "Addition can not be applied to %s and the given type." % self.pretty_name + ) + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, numbers.Number) and not isinstance(right, bool): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right - left + else: + raise TypeError( + "Subtraction can not be applied to %s and the given type." % self.pretty_name + ) + + def rmul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, bool): + return left.__and__(right) + elif isinstance(right, numbers.Number): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right * left + else: + raise TypeError( + "Multiplication can not be applied to %s and the given type." % self.pretty_name + ) + + def rtruediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, numbers.Number) and not isinstance(right, bool): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right / left + else: + raise TypeError( + "True division can not be applied to %s and the given type." % self.pretty_name + ) + + def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, numbers.Number) and not isinstance(right, bool): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right // left + else: + raise TypeError( + "Floor division can not be applied to %s and the given type." % self.pretty_name + ) + + def rpow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, numbers.Number) and not isinstance(right, bool): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right**left + else: + raise TypeError( + "Exponentiation can not be applied to %s and the given type." % self.pretty_name + ) + + def rmod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, numbers.Number) and not isinstance(right, bool): + left = transform_boolean_operand_to_numeric(left, spark_type=as_spark_type(type(right))) + return right % left + else: + raise TypeError( + "Modulo can not be applied to %s and the given type." % self.pretty_name + ) + + def __and__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, IndexOpsMixin) and isinstance(right.dtype, extension_dtypes): + return right.__and__(left) + else: + + def and_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + scol = left & right + return F.when(scol.isNull(), False).otherwise(scol) + + return column_op(and_func)(left, right) + + def xor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, IndexOpsMixin) and isinstance(right.dtype, extension_dtypes): + return right ^ left + elif _is_valid_for_logical_operator(right): + + def xor_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + scol = left.cast("integer").bitwiseXOR(right.cast("integer")).cast("boolean") + return F.when(scol.isNull(), False).otherwise(scol) + + return column_op(xor_func)(left, right) + else: + raise TypeError("XOR can not be applied to given types.") + + def __or__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, IndexOpsMixin) and isinstance(right.dtype, extension_dtypes): + return right.__or__(left) + else: + + def or_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + if is_null: + return F.lit(False) + else: + scol = left | F.lit(right) + return F.when(left.isNull() | scol.isNull(), False).otherwise(scol) + + return column_op(or_func)(left, right) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + return _as_bool_type(index_ops, dtype) + elif isinstance(spark_type, StringType): + if isinstance(dtype, extension_dtypes): + scol = F.when( + index_ops.spark.column.isNotNull(), + F.when(index_ops.spark.column, "True").otherwise("False"), + ) + nullable = index_ops.spark.nullable + else: + null_str = str(pd.NA) if isinstance(self, BooleanExtensionOps) else str(None) + casted = F.when(index_ops.spark.column, "True").otherwise("False") + scol = F.when(index_ops.spark.column.isNull(), null_str).otherwise(casted) + nullable = False + return index_ops._with_new_scol( + scol, + field=index_ops._internal.data_fields[0].copy( + dtype=dtype, spark_type=spark_type, nullable=nullable + ), + ) + else: + return _as_other_type(index_ops, dtype, spark_type) + + def neg(self, operand: IndexOpsLike) -> IndexOpsLike: + return ~operand + + def abs(self, operand: IndexOpsLike) -> IndexOpsLike: + return operand + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__lt__)(left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__le__)(left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__ge__)(left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__gt__)(left, right) + + def invert(self, operand: IndexOpsLike) -> IndexOpsLike: + return operand._with_new_scol(~operand.spark.column, field=operand._internal.data_fields[0]) + + +class BooleanExtensionOps(BooleanOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type BooleanType, + and dtype BooleanDtype. + """ + + @property + def pretty_name(self) -> str: + return "booleans" + + def __and__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + def and_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + return left & right + + return column_op(and_func)(left, right) + + def __or__(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + def or_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + return left | right + + return column_op(or_func)(left, right) + + def xor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if _is_boolean_type(right): + + def xor_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + return left.cast("integer").bitwiseXOR(right.cast("integer")).cast("boolean") + + return column_op(xor_func)(left, right) + else: + raise TypeError("XOR can not be applied to given types.") + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return col.astype(self.dtype) + + def neg(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("Unary - can not be applied to %s." % self.pretty_name) + + def invert(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("Unary ~ can not be applied to %s." % self.pretty_name) + + def abs(self, operand: IndexOpsLike) -> IndexOpsLike: + raise TypeError("abs() can not be applied to %s." % self.pretty_name) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/categorical_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/categorical_ops.py new file mode 100644 index 0000000..824666b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/categorical_ops.py @@ -0,0 +1,141 @@ +# +# 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 itertools import chain +from typing import cast, Any, Union + +import pandas as pd +import numpy as np +from pandas.api.types import is_list_like, CategoricalDtype # type: ignore[attr-defined] + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.data_type_ops.base import _sanitize_list_like, DataTypeOps +from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql import functions as F +from pyspark.sql.utils import pyspark_column_op + + +class CategoricalOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with categorical types. + """ + + @property + def pretty_name(self) -> str: + return "categoricals" + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return pd.Series( + pd.Categorical.from_codes( + col.replace(np.nan, -1).astype(int), + categories=cast(CategoricalDtype, self.dtype).categories, + ordered=cast(CategoricalDtype, self.dtype).ordered, + ) + ) + + def prepare(self, col: pd.Series) -> pd.Series: + """Prepare column when from_pandas.""" + return col.cat.codes + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, _ = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype) and ( + (dtype.categories is None) or (index_ops.dtype == dtype) + ): + return index_ops.copy() + + return _to_cat(index_ops).astype(dtype) + + def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__eq__", is_equality_comparison=True) + + def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__ne__", is_equality_comparison=True) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__lt__") + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__le__") + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__gt__") + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return _compare(left, right, "__ge__") + + +def _compare( + left: IndexOpsLike, + right: Any, + func_name: str, + *, + is_equality_comparison: bool = False, +) -> SeriesOrIndex: + """ + Compare a Categorical operand `left` to `right` with the given Spark Column function. + + Parameters + ---------- + left: A Categorical operand + right: The other operand to compare with + func_name: The Spark Column function name to apply + is_equality_comparison: True if it is equality comparison, ie. == or !=. False by default. + + Returns + ------- + SeriesOrIndex + """ + if isinstance(right, IndexOpsMixin) and isinstance(right.dtype, CategoricalDtype): + if not is_equality_comparison: + if not cast(CategoricalDtype, left.dtype).ordered: + raise TypeError("Unordered Categoricals can only compare equality or not.") + # Check if categoricals have the same dtype, same categories, and same ordered + if hash(left.dtype) != hash(right.dtype): + raise TypeError("Categoricals can only be compared if 'categories' are the same.") + if cast(CategoricalDtype, left.dtype).ordered: + return pyspark_column_op(func_name, left, right) + else: + return pyspark_column_op(func_name, _to_cat(left), _to_cat(right)) + elif not is_list_like(right): + categories = cast(CategoricalDtype, left.dtype).categories + if right not in categories: + raise TypeError("Cannot compare a Categorical with a scalar, which is not a category.") + right_code = categories.get_loc(right) + return pyspark_column_op(func_name, left, right_code) + else: + raise TypeError("Cannot compare a Categorical with the given type.") + + +def _to_cat(index_ops: IndexOpsLike) -> IndexOpsLike: + categories = cast(CategoricalDtype, index_ops.dtype).categories + if len(categories) == 0: + scol = F.lit(None) + else: + kvs = chain(*[(F.lit(code), F.lit(category)) for code, category in enumerate(categories)]) + map_scol = F.create_map(*kvs) + scol = map_scol[index_ops.spark.column] + return index_ops._with_new_scol(scol) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/complex_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/complex_ops.py new file mode 100644 index 0000000..415301e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/complex_ops.py @@ -0,0 +1,145 @@ +# +# 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 Any, Union, cast + +from pandas.api.types import CategoricalDtype + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_bool_type, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql import functions as F, Column +from pyspark.sql.types import ArrayType, BooleanType, NumericType, StringType + + +class ArrayOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with ArrayType. + """ + + @property + def pretty_name(self) -> str: + return "arrays" + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, IndexOpsMixin) or ( + isinstance(right, IndexOpsMixin) and not isinstance(right.spark.data_type, ArrayType) + ): + raise TypeError( + "Concatenation can not be applied to %s and the given type." % self.pretty_name + ) + + left_type = cast(ArrayType, left.spark.data_type).elementType + right_type = right.spark.data_type.elementType + + if left_type != right_type and not ( + isinstance(left_type, NumericType) and isinstance(right_type, NumericType) + ): + raise TypeError( + "Concatenation can only be applied to %s of the same type" % self.pretty_name + ) + + return column_op(F.concat)(left, right) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__lt__)(left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__le__)(left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__ge__)(left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__gt__)(left, right) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + return _as_bool_type(index_ops, dtype) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype) + else: + return _as_other_type(index_ops, dtype, spark_type) + + +class MapOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with MapType. + """ + + @property + def pretty_name(self) -> str: + return "maps" + + +class StructOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with StructType. + """ + + @property + def pretty_name(self) -> str: + return "structs" + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__lt__)(left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__le__)(left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__ge__)(left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + return column_op(Column.__gt__)(left, right) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/date_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/date_ops.py new file mode 100644 index 0000000..51d1018 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/date_ops.py @@ -0,0 +1,127 @@ +# +# 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 datetime +import warnings +from typing import Any, Union + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark.sql import functions as F +from pyspark.sql.types import BooleanType, DateType, StringType +from pyspark.sql.utils import get_column_class + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.typedef import pandas_on_spark_type + + +class DateOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: DateType. + """ + + @property + def pretty_name(self) -> str: + return "dates" + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + # Note that date subtraction casts arguments to integer. This is to mimic pandas's + # behaviors. pandas returns 'timedelta64[ns]' in days from date's subtraction. + msg = ( + "Note that there is a behavior difference of date subtraction. " + "The date subtraction returns an integer in days, " + "whereas pandas returns 'timedelta64[ns]'." + ) + if isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, DateType): + warnings.warn(msg, UserWarning) + return column_op(F.datediff)(left, right).astype("long") + elif isinstance(right, datetime.date) and not isinstance(right, datetime.datetime): + warnings.warn(msg, UserWarning) + return column_op(F.datediff)(left, F.lit(right)).astype("long") + else: + raise TypeError("Date subtraction can only be applied to date series.") + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + # Note that date subtraction casts arguments to integer. This is to mimic pandas's + # behaviors. pandas returns 'timedelta64[ns]' in days from date's subtraction. + msg = ( + "Note that there is a behavior difference of date subtraction. " + "The date subtraction returns an integer in days, " + "whereas pandas returns 'timedelta64[ns]'." + ) + if isinstance(right, datetime.date) and not isinstance(right, datetime.datetime): + warnings.warn(msg, UserWarning) + return -column_op(F.datediff)(left, F.lit(right)).astype("long") + else: + raise TypeError("Date subtraction can only be applied to date series.") + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__lt__)(left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__le__)(left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__ge__)(left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + from pyspark.pandas.base import column_op + + _sanitize_list_like(right) + Column = get_column_class() + return column_op(Column.__gt__)(left, right) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + return index_ops._with_new_scol( + index_ops.spark.column.isNotNull(), + field=index_ops._internal.data_fields[0].copy( + dtype=np.dtype(bool), spark_type=spark_type, nullable=False + ), + ) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str=str(pd.NaT)) + else: + return _as_other_type(index_ops, dtype, spark_type) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/datetime_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/datetime_ops.py new file mode 100644 index 0000000..ea9b994 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/datetime_ops.py @@ -0,0 +1,171 @@ +# +# 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 datetime +import warnings +from typing import Any, Union, cast + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark import SparkContext +from pyspark.sql import Column, functions as F +from pyspark.sql.types import ( + BooleanType, + LongType, + StringType, + TimestampType, + TimestampNTZType, + NumericType, +) +from pyspark.sql.utils import pyspark_column_op + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.typedef import pandas_on_spark_type + + +class DatetimeOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: TimestampType. + """ + + @property + def pretty_name(self) -> str: + return "datetimes" + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + # Note that timestamp subtraction casts arguments to integer. This is to mimic pandas's + # behaviors. pandas returns 'timedelta64[ns]' from 'datetime64[ns]'s subtraction. + msg = ( + "Note that there is a behavior difference of timestamp subtraction. " + "The timestamp subtraction returns an integer in seconds, " + "whereas pandas returns 'timedelta64[ns]'." + ) + if isinstance(right, IndexOpsMixin) and isinstance( + right.spark.data_type, (TimestampType, TimestampNTZType) + ): + warnings.warn(msg, UserWarning) + return left.astype("long") - right.astype("long") + elif isinstance(right, datetime.datetime): + warnings.warn(msg, UserWarning) + return cast( + SeriesOrIndex, + left._with_new_scol( + left.astype("long").spark.column + - self._cast_spark_column_timestamp_to_long(F.lit(right)), + field=left._internal.data_fields[0].copy( + dtype=np.dtype("int64"), spark_type=LongType() + ), + ), + ) + else: + raise TypeError("Datetime subtraction can only be applied to datetime series.") + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + # Note that timestamp subtraction casts arguments to integer. This is to mimic pandas's + # behaviors. pandas returns 'timedelta64[ns]' from 'datetime64[ns]'s subtraction. + msg = ( + "Note that there is a behavior difference of timestamp subtraction. " + "The timestamp subtraction returns an integer in seconds, " + "whereas pandas returns 'timedelta64[ns]'." + ) + if isinstance(right, datetime.datetime): + warnings.warn(msg, UserWarning) + return cast( + SeriesOrIndex, + left._with_new_scol( + self._cast_spark_column_timestamp_to_long(F.lit(right)) + - left.astype("long").spark.column, + field=left._internal.data_fields[0].copy( + dtype=np.dtype("int64"), spark_type=LongType() + ), + ), + ) + else: + raise TypeError("Datetime subtraction can only be applied to datetime series.") + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right) + + def prepare(self, col: pd.Series) -> pd.Series: + """Prepare column when from_pandas.""" + return col + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + raise TypeError("cannot astype a %s to [bool]" % self.pretty_name) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str=str(pd.NaT)) + else: + return _as_other_type(index_ops, dtype, spark_type) + + def _cast_spark_column_timestamp_to_long(self, scol: Column) -> Column: + return scol.cast(LongType()) + + +class DatetimeNTZOps(DatetimeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: + TimestampNTZType. + """ + + def _cast_spark_column_timestamp_to_long(self, scol: Column) -> Column: + jvm = SparkContext._active_spark_context._jvm + return Column(jvm.PythonSQLUtils.castTimestampNTZToLong(scol._jc)) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, NumericType): + from pyspark.pandas.internal import InternalField + + scol = self._cast_spark_column_timestamp_to_long(index_ops.spark.column).cast( + spark_type + ) + return index_ops._with_new_scol(scol, field=InternalField(dtype=dtype)) + else: + return super(DatetimeNTZOps, self).astype(index_ops, dtype) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/null_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/null_ops.py new file mode 100644 index 0000000..329a379 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/null_ops.py @@ -0,0 +1,83 @@ +# +# 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 Any, Union + +from pandas.api.types import CategoricalDtype, is_list_like # type: ignore[attr-defined] + +from pyspark.pandas._typing import Dtype, IndexOpsLike +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_bool_type, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas._typing import SeriesOrIndex +from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql.types import BooleanType, StringType +from pyspark.sql.utils import pyspark_column_op +from pyspark.pandas.base import IndexOpsMixin + + +class NullOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with Spark type: NullType. + """ + + @property + def pretty_name(self) -> str: + return "nulls" + + def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + # We can directly use `super().eq` when given object is list, tuple, dict or set. + if not isinstance(right, IndexOpsMixin) and is_list_like(right): + return super().eq(left, right) + return pyspark_column_op("__eq__", left, right, fillna=False) + + def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ne__", left, right, fillna=True) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right, fillna=False) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right, fillna=False) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right, fillna=False) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right, fillna=False) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + return _as_bool_type(index_ops, dtype) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype) + else: + return _as_other_type(index_ops, dtype, spark_type) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/num_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/num_ops.py new file mode 100644 index 0000000..af5e387 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/num_ops.py @@ -0,0 +1,588 @@ +# +# 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 numbers +from typing import Any, Union + +import numpy as np +import pandas as pd +from pandas.api.types import ( # type: ignore[attr-defined] + is_bool_dtype, + is_integer_dtype, + CategoricalDtype, + is_list_like, +) + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import column_op, IndexOpsMixin, numpy_column_op +from pyspark.pandas.config import get_option +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + is_valid_operand_for_numeric_arithmetic, + transform_boolean_operand_to_numeric, + _as_bool_type, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, + _is_valid_for_logical_operator, + _is_boolean_type, +) +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.typedef.typehints import extension_dtypes, pandas_on_spark_type +from pyspark.sql import functions as F +from pyspark.sql import Column as PySparkColumn +from pyspark.sql.types import ( + BooleanType, + DataType, + StringType, +) +from pyspark.errors import PySparkValueError + +# For Supporting Spark Connect +from pyspark.sql.utils import pyspark_column_op, get_column_class + + +def _non_fractional_astype( + index_ops: IndexOpsLike, dtype: Dtype, spark_type: DataType +) -> IndexOpsLike: + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + return _as_bool_type(index_ops, dtype) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str="NaN") + else: + return _as_other_type(index_ops, dtype, spark_type) + + +class NumericOps(DataTypeOps): + """The class for binary operations of numeric pandas-on-Spark objects.""" + + @property + def pretty_name(self) -> str: + return "numerics" + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Addition can not be applied to given types.") + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + Column = get_column_class() + return column_op(Column.__add__)(left, right) + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Subtraction can not be applied to given types.") + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + Column = get_column_class() + return column_op(Column.__sub__)(left, right) + + def mod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Modulo can not be applied to given types.") + + def mod(left: PySparkColumn, right: Any) -> PySparkColumn: + return ((left % right) + right) % right + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return column_op(mod)(left, right) + + def pow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Exponentiation can not be applied to given types.") + + Column = get_column_class() + + def pow_func(left: Column, right: Any) -> Column: # type: ignore[valid-type] + return ( + F.when(left == 1, left) # type: ignore + .when(F.lit(right) == 0, 1) + .otherwise(Column.__pow__(left, right)) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return column_op(pow_func)(left, right) + + def radd(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Addition can not be applied to given types.") + right = transform_boolean_operand_to_numeric(right) + Column = get_column_class() + return column_op(Column.__radd__)(left, right) + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Subtraction can not be applied to given types.") + right = transform_boolean_operand_to_numeric(right) + Column = get_column_class() + return column_op(Column.__rsub__)(left, right) + + def rmul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Multiplication can not be applied to given types.") + right = transform_boolean_operand_to_numeric(right) + Column = get_column_class() + return column_op(Column.__rmul__)(left, right) + + def rpow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Exponentiation can not be applied to given types.") + + Column = get_column_class() + + def rpow_func(left: Column, right: Any) -> Column: # type: ignore[valid-type] + return F.when(F.lit(right == 1), right).otherwise(Column.__rpow__(left, right)) + + right = transform_boolean_operand_to_numeric(right) + return column_op(rpow_func)(left, right) + + def rmod(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Modulo can not be applied to given types.") + + def rmod(left: PySparkColumn, right: Any) -> PySparkColumn: + return ((right % left) + left) % left + + right = transform_boolean_operand_to_numeric(right) + return column_op(rmod)(left, right) + + def neg(self, operand: IndexOpsLike) -> IndexOpsLike: + return operand._with_new_scol(-operand.spark.column, field=operand._internal.data_fields[0]) + + def abs(self, operand: IndexOpsLike) -> IndexOpsLike: + return operand._with_new_scol( + F.abs(operand.spark.column), field=operand._internal.data_fields[0] + ) + + def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + # We can directly use `super().eq` when given object is list, tuple, dict or set. + if not isinstance(right, IndexOpsMixin) and is_list_like(right): + return super().eq(left, right) + return pyspark_column_op("__eq__", left, right, fillna=False) + + def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ne__", left, right, fillna=True) + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right, fillna=False) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right, fillna=False) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right, fillna=False) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right, fillna=False) + + +class IntegralOps(NumericOps): + """ + The class for binary operations of pandas-on-Spark objects with spark types: + LongType, IntegerType, ByteType and ShortType. + """ + + def xor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if isinstance(right, IndexOpsMixin) and isinstance(right.dtype, extension_dtypes): + return right ^ left + elif _is_valid_for_logical_operator(right): + right_is_boolean = _is_boolean_type(right) + + def xor_func(left: PySparkColumn, right: Any) -> PySparkColumn: + try: + is_null = pd.isna(right) + except PySparkValueError: + # Complaining `PySparkValueError` means that `right` is a Column. + is_null = False + + right = F.lit(None) if is_null else F.lit(right) + return ( + left.bitwiseXOR(right.cast("integer")).cast("boolean") + if right_is_boolean + else left.bitwiseXOR(right) + ) + + return column_op(xor_func)(left, right) + else: + raise TypeError("XOR can not be applied to given types.") + + @property + def pretty_name(self) -> str: + return "integrals" + + def mul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, StringType): + return column_op(SF.repeat)(right, left) + + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Multiplication can not be applied to given types.") + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + Column = get_column_class() + return column_op(Column.__mul__)(left, right) + + def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("True division can not be applied to given types.") + + def truediv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when( + F.lit(right != 0) | F.lit(right).isNull(), + left.__div__(right), + ).otherwise(F.lit(np.inf).__div__(left)) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(truediv)(left, right) + + def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Floor division can not be applied to given types.") + + def floordiv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(F.lit(right is np.nan), np.nan).otherwise( + F.when( + F.lit(right != 0) | F.lit(right).isNull(), + F.floor(left.__div__(right)), + ).otherwise(F.lit(np.inf).__div__(left)) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(floordiv)(left, right) + + def rtruediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("True division can not be applied to given types.") + + def rtruediv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(left == 0, F.lit(np.inf).__div__(right)).otherwise( + F.lit(right).__truediv__(left) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(rtruediv)(left, right) + + def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Floor division can not be applied to given types.") + + def rfloordiv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(F.lit(left == 0), F.lit(np.inf).__div__(right)).otherwise( + F.floor(F.lit(right).__div__(left)) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(rfloordiv)(left, right) + + def invert(self, operand: IndexOpsLike) -> IndexOpsLike: + return operand._with_new_scol( + F.bitwise_not(operand.spark.column), field=operand._internal.data_fields[0] + ) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + return _non_fractional_astype(index_ops, dtype, spark_type) + + +class FractionalOps(NumericOps): + """ + The class for binary operations of pandas-on-Spark objects with spark types: + FloatType, DoubleType. + """ + + @property + def pretty_name(self) -> str: + return "fractions" + + def mul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Multiplication can not be applied to given types.") + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + Column = get_column_class() + return column_op(Column.__mul__)(left, right) + + def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("True division can not be applied to given types.") + + def truediv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when( + F.lit(right != 0) | F.lit(right).isNull(), + left.__div__(right), + ).otherwise( + F.when(F.lit(left == np.inf) | F.lit(left == -np.inf), left).otherwise( + F.lit(np.inf).__div__(left) + ) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(truediv)(left, right) + + def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not is_valid_operand_for_numeric_arithmetic(right): + raise TypeError("Floor division can not be applied to given types.") + + def floordiv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(F.lit(right is np.nan), np.nan).otherwise( + F.when( + F.lit(right != 0) | F.lit(right).isNull(), + F.floor(left.__div__(right)), + ).otherwise( + F.when(F.lit(left == np.inf) | F.lit(left == -np.inf), left).otherwise( + F.lit(np.inf).__div__(left) + ) + ) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(floordiv)(left, right) + + def rtruediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("True division can not be applied to given types.") + + def rtruediv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(left == 0, F.lit(np.inf).__div__(right)).otherwise( + F.lit(right).__truediv__(left) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(rtruediv)(left, right) + + def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Floor division can not be applied to given types.") + + def rfloordiv(left: PySparkColumn, right: Any) -> PySparkColumn: + return F.when(F.lit(left == 0), F.lit(np.inf).__div__(right)).otherwise( + F.when(F.lit(left) == np.nan, np.nan).otherwise(F.floor(F.lit(right).__div__(left))) + ) + + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) + return numpy_column_op(rfloordiv)(left, right) + + def isnull(self, index_ops: IndexOpsLike) -> IndexOpsLike: + return index_ops._with_new_scol( + index_ops.spark.column.isNull() | F.isnan(index_ops.spark.column), + field=index_ops._internal.data_fields[0].copy( + dtype=np.dtype("bool"), spark_type=BooleanType(), nullable=False + ), + ) + + def nan_to_null(self, index_ops: IndexOpsLike) -> IndexOpsLike: + # Special handle floating point types because Spark's count treats nan as a valid value, + # whereas pandas count doesn't include nan. + return index_ops._with_new_scol( + F.nanvl(index_ops.spark.column, F.lit(None)), + field=index_ops._internal.data_fields[0].copy(nullable=True), + ) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if is_integer_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if get_option("compute.eager_check") and index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to integer" % self.pretty_name + ) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + if isinstance(dtype, extension_dtypes): + scol = index_ops.spark.column.cast(spark_type) + else: + scol = F.when( + index_ops.spark.column.isNull() | F.isnan(index_ops.spark.column), + F.lit(True), + ).otherwise(index_ops.spark.column.cast(spark_type)) + return index_ops._with_new_scol( + scol.alias(index_ops._internal.data_spark_column_names[0]), + field=index_ops._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type), + ) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str=str(np.nan)) + else: + return _as_other_type(index_ops, dtype, spark_type) + + +class DecimalOps(FractionalOps): + """ + The class for decimal operations of pandas-on-Spark objects with spark type: + DecimalType. + """ + + @property + def pretty_name(self) -> str: + return "decimal" + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("< can not be applied to %s." % self.pretty_name) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("<= can not be applied to %s." % self.pretty_name) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError("> can not be applied to %s." % self.pretty_name) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + raise TypeError(">= can not be applied to %s." % self.pretty_name) + + def isnull(self, index_ops: IndexOpsLike) -> IndexOpsLike: + return index_ops._with_new_scol( + index_ops.spark.column.isNull(), + field=index_ops._internal.data_fields[0].copy( + dtype=np.dtype("bool"), spark_type=BooleanType(), nullable=False + ), + ) + + def nan_to_null(self, index_ops: IndexOpsLike) -> IndexOpsLike: + return index_ops.copy() + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + if is_integer_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if get_option("compute.eager_check") and index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to integer" % self.pretty_name + ) + return _non_fractional_astype(index_ops, dtype, spark_type) + + def rpow(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if not isinstance(right, numbers.Number): + raise TypeError("Exponentiation can not be applied to given types.") + + Column = get_column_class() + + def rpow_func(left: Column, right: Any) -> Column: # type: ignore[valid-type] + return ( + F.when(left.isNull(), np.nan) # type: ignore + .when(F.lit(right == 1), right) + .otherwise(Column.__rpow__(left, right)) + ) + + right = transform_boolean_operand_to_numeric(right) + return column_op(rpow_func)(left, right) + + +class IntegralExtensionOps(IntegralOps): + """ + The class for binary operations of pandas-on-Spark objects with one of the + - spark types: + LongType, IntegerType, ByteType and ShortType + - dtypes: + Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + """ + + def xor(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + raise TypeError("XOR can not be applied to given types.") + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return col.astype(self.dtype) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + if get_option("compute.eager_check"): + if is_integer_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to integer" % self.pretty_name + ) + elif is_bool_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to bool" % self.pretty_name + ) + return _non_fractional_astype(index_ops, dtype, spark_type) + + +class FractionalExtensionOps(FractionalOps): + """ + The class for binary operations of pandas-on-Spark objects with one of the + - spark types: + FloatType, DoubleType and DecimalType + - dtypes: + Float32Dtype, Float64Dtype + """ + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return col.astype(self.dtype) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + if get_option("compute.eager_check"): + if is_integer_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to integer" % self.pretty_name + ) + elif is_bool_dtype(dtype) and not isinstance(dtype, extension_dtypes): + if index_ops.hasnans: + raise ValueError( + "Cannot convert %s with missing values to bool" % self.pretty_name + ) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + if isinstance(dtype, extension_dtypes): + scol = index_ops.spark.column.cast(spark_type) + else: + scol = F.when( + index_ops.spark.column.isNull() | F.isnan(index_ops.spark.column), + F.lit(True), + ).otherwise(index_ops.spark.column.cast(spark_type)) + return index_ops._with_new_scol( + scol.alias(index_ops._internal.data_spark_column_names[0]), + field=index_ops._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type), + ) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str=str(np.nan)) + else: + return _as_other_type(index_ops, dtype, spark_type) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/string_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/string_ops.py new file mode 100644 index 0000000..1c282f2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/string_ops.py @@ -0,0 +1,154 @@ +# +# 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 Any, Union, cast + +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark.sql import functions as F +from pyspark.sql.types import IntegralType, StringType +from pyspark.sql.utils import pyspark_column_op + +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.typedef import extension_dtypes, pandas_on_spark_type +from pyspark.sql.types import BooleanType + + +class StringOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: StringType. + """ + + @property + def pretty_name(self) -> str: + return "strings" + + def add(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, str): + return cast( + SeriesOrIndex, + left._with_new_scol( + F.concat(left.spark.column, F.lit(right)), field=left._internal.data_fields[0] + ), + ) + elif isinstance(right, IndexOpsMixin) and isinstance(right.spark.data_type, StringType): + return column_op(F.concat)(left, right) + else: + raise TypeError("Addition can not be applied to given types.") + + def mul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, int): + return cast( + SeriesOrIndex, + left._with_new_scol( + SF.repeat(left.spark.column, right), field=left._internal.data_fields[0] + ), + ) + elif ( + isinstance(right, IndexOpsMixin) + and isinstance(right.spark.data_type, IntegralType) + and not isinstance(right.dtype, CategoricalDtype) + ): + return column_op(SF.repeat)(left, right) + else: + raise TypeError("Multiplication can not be applied to given types.") + + def radd(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, str): + return cast( + SeriesOrIndex, + left._with_new_scol( + F.concat(F.lit(right), left.spark.column), field=left._internal.data_fields[0] + ), + ) + else: + raise TypeError("Addition can not be applied to given types.") + + def rmul(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + if isinstance(right, int): + return cast( + SeriesOrIndex, + left._with_new_scol( + SF.repeat(left.spark.column, right), field=left._internal.data_fields[0] + ), + ) + else: + raise TypeError("Multiplication can not be applied to given types.") + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right) + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + + if isinstance(spark_type, BooleanType): + if isinstance(dtype, extension_dtypes): + scol = index_ops.spark.column.cast(spark_type) + else: + scol = F.when(index_ops.spark.column.isNull(), F.lit(False)).otherwise( + F.length(index_ops.spark.column) > 0 + ) + return index_ops._with_new_scol( + scol, + field=index_ops._internal.data_fields[0].copy(dtype=dtype, spark_type=spark_type), + ) + elif isinstance(spark_type, StringType): + null_str = str(pd.NA) if isinstance(self, StringExtensionOps) else str(None) + return _as_string_type(index_ops, dtype, null_str=null_str) + else: + return _as_other_type(index_ops, dtype, spark_type) + + +class StringExtensionOps(StringOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type StringType, + and dtype StringDtype. + """ + + def restore(self, col: pd.Series) -> pd.Series: + """Restore column when to_pandas.""" + return col.astype(self.dtype) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/timedelta_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/timedelta_ops.py new file mode 100644 index 0000000..7a9da85 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/timedelta_ops.py @@ -0,0 +1,101 @@ +# +# 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 datetime import timedelta +from typing import Any, Union + +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark.sql.types import ( + BooleanType, + DayTimeIntervalType, + StringType, +) +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.data_type_ops.base import ( + DataTypeOps, + _as_categorical_type, + _as_other_type, + _as_string_type, + _sanitize_list_like, +) +from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql.utils import pyspark_column_op + + +class TimedeltaOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with spark type: DayTimeIntervalType. + """ + + @property + def pretty_name(self) -> str: + return "timedelta" + + def astype(self, index_ops: IndexOpsLike, dtype: Union[str, type, Dtype]) -> IndexOpsLike: + dtype, spark_type = pandas_on_spark_type(dtype) + + if isinstance(dtype, CategoricalDtype): + return _as_categorical_type(index_ops, dtype, spark_type) + elif isinstance(spark_type, BooleanType): + raise TypeError("cannot astype a %s to [bool]" % self.pretty_name) + elif isinstance(spark_type, StringType): + return _as_string_type(index_ops, dtype, null_str=str(pd.NaT)) + else: + return _as_other_type(index_ops, dtype, spark_type) + + def prepare(self, col: pd.Series) -> pd.Series: + """Prepare column when from_pandas.""" + return col + + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if ( + isinstance(right, IndexOpsMixin) + and isinstance(right.spark.data_type, DayTimeIntervalType) + or isinstance(right, timedelta) + ): + return pyspark_column_op("__sub__", left, right) + else: + raise TypeError("Timedelta subtraction can only be applied to timedelta series.") + + def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + + if isinstance(right, timedelta): + return pyspark_column_op("__rsub__", left, right) + else: + raise TypeError("Timedelta subtraction can only be applied to timedelta series.") + + def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__lt__", left, right) + + def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__le__", left, right) + + def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__ge__", left, right) + + def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: + _sanitize_list_like(right) + return pyspark_column_op("__gt__", left, right) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/udt_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/udt_ops.py new file mode 100644 index 0000000..ba08162 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/data_type_ops/udt_ops.py @@ -0,0 +1,29 @@ +# +# 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.pandas.data_type_ops.base import DataTypeOps + + +class UDTOps(DataTypeOps): + """ + The class for binary operations of pandas-on-Spark objects with Spark type: + UserDefinedType or its subclasses. + """ + + @property + def pretty_name(self) -> str: + return "user defined types" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/datetimes.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/datetimes.py new file mode 100644 index 0000000..752f6f4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/datetimes.py @@ -0,0 +1,891 @@ +# +# 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. +# + +""" +Date/Time related functions on pandas-on-Spark Series +""" +import warnings +from typing import Any, Optional, Union, no_type_check + +import numpy as np +import pandas as pd # noqa: F401 +from pandas.tseries.offsets import DateOffset + +import pyspark.pandas as ps +import pyspark.sql.functions as F +from pyspark.sql.types import DateType, TimestampType, TimestampNTZType, LongType + + +class DatetimeMethods: + """Date/Time methods for pandas-on-Spark Series""" + + def __init__(self, series: "ps.Series"): + if not isinstance(series.spark.data_type, (DateType, TimestampType, TimestampNTZType)): + raise ValueError( + "Cannot call DatetimeMethods on type {}".format(series.spark.data_type) + ) + self._data = series + + # Properties + @property + def date(self) -> "ps.Series": + """ + Returns a Series of python datetime.date objects (namely, the date + part of Timestamps without timezone information). + """ + # TODO: Hit a weird exception + # syntax error in attribute name: `to_date(`start_date`)` with alias + return self._data.spark.transform(F.to_date) + + @property + def time(self) -> "ps.Series": + raise NotImplementedError() + + @property + def timetz(self) -> "ps.Series": + raise NotImplementedError() + + @property + def year(self) -> "ps.Series": + """ + The year of the datetime. + """ + return self._data.spark.transform(lambda c: F.year(c).cast(LongType())) + + @property + def month(self) -> "ps.Series": + """ + The month of the timestamp as January = 1 December = 12. + """ + return self._data.spark.transform(lambda c: F.month(c).cast(LongType())) + + @property + def day(self) -> "ps.Series": + """ + The days of the datetime. + """ + return self._data.spark.transform(lambda c: F.dayofmonth(c).cast(LongType())) + + @property + def hour(self) -> "ps.Series": + """ + The hours of the datetime. + """ + return self._data.spark.transform(lambda c: F.hour(c).cast(LongType())) + + @property + def minute(self) -> "ps.Series": + """ + The minutes of the datetime. + """ + return self._data.spark.transform(lambda c: F.minute(c).cast(LongType())) + + @property + def second(self) -> "ps.Series": + """ + The seconds of the datetime. + """ + return self._data.spark.transform(lambda c: F.second(c).cast(LongType())) + + @property + def microsecond(self) -> "ps.Series": + """ + The microseconds of the datetime. + """ + + def pandas_microsecond(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.dt.microsecond + + return self._data.pandas_on_spark.transform_batch(pandas_microsecond) + + @property + def nanosecond(self) -> "ps.Series": + raise NotImplementedError() + + # TODO(SPARK-42617): Support isocalendar.week and replace it. + # See also https://github.com/pandas-dev/pandas/pull/33595. + @property + def week(self) -> "ps.Series": + """ + The week ordinal of the year. + + .. deprecated:: 3.4.0 + """ + warnings.warn( + "weekofyear and week have been deprecated.", + FutureWarning, + ) + return self._data.spark.transform(lambda c: F.weekofyear(c).cast(LongType())) + + @property + def weekofyear(self) -> "ps.Series": + return self.week + + weekofyear.__doc__ = week.__doc__ + + @property + def dayofweek(self) -> "ps.Series": + """ + The day of the week with Monday=0, Sunday=6. + + Return the day of the week. It is assumed the week starts on + Monday, which is denoted by 0 and ends on Sunday which is denoted + by 6. This method is available on both Series with datetime + values (using the `dt` accessor). + + Returns + ------- + Series + Containing integers indicating the day number. + + See Also + -------- + Series.dt.dayofweek : Alias. + Series.dt.weekday : Alias. + Series.dt.day_name : Returns the name of the day of the week. + + Examples + -------- + >>> s = ps.from_pandas(pd.date_range('2016-12-31', '2017-01-08', freq='D').to_series()) + >>> s.dt.dayofweek + 2016-12-31 5 + 2017-01-01 6 + 2017-01-02 0 + 2017-01-03 1 + 2017-01-04 2 + 2017-01-05 3 + 2017-01-06 4 + 2017-01-07 5 + 2017-01-08 6 + dtype: int64 + """ + + def pandas_dayofweek(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.dt.dayofweek + + return self._data.pandas_on_spark.transform_batch(pandas_dayofweek) + + @property + def weekday(self) -> "ps.Series": + return self.dayofweek + + weekday.__doc__ = dayofweek.__doc__ + + @property + def dayofyear(self) -> "ps.Series": + """ + The ordinal day of the year. + """ + + def pandas_dayofyear(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.dt.dayofyear + + return self._data.pandas_on_spark.transform_batch(pandas_dayofyear) + + @property + def quarter(self) -> "ps.Series": + """ + The quarter of the date. + """ + + def pandas_quarter(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.dt.quarter + + return self._data.pandas_on_spark.transform_batch(pandas_quarter) + + @property + def is_month_start(self) -> "ps.Series": + """ + Indicates whether the date is the first day of the month. + + Returns + ------- + Series + For Series, returns a Series with boolean values. + + See Also + -------- + is_month_end : Return a boolean indicating whether the date + is the last day of the month. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> s = ps.Series(pd.date_range("2018-02-27", periods=3)) + >>> s + 0 2018-02-27 + 1 2018-02-28 + 2 2018-03-01 + dtype: datetime64[ns] + + >>> s.dt.is_month_start + 0 False + 1 False + 2 True + dtype: bool + """ + + def pandas_is_month_start(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_month_start + + return self._data.pandas_on_spark.transform_batch(pandas_is_month_start) + + @property + def is_month_end(self) -> "ps.Series": + """ + Indicates whether the date is the last day of the month. + + Returns + ------- + Series + For Series, returns a Series with boolean values. + + See Also + -------- + is_month_start : Return a boolean indicating whether the date + is the first day of the month. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> s = ps.Series(pd.date_range("2018-02-27", periods=3)) + >>> s + 0 2018-02-27 + 1 2018-02-28 + 2 2018-03-01 + dtype: datetime64[ns] + + >>> s.dt.is_month_end + 0 False + 1 True + 2 False + dtype: bool + """ + + def pandas_is_month_end(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_month_end + + return self._data.pandas_on_spark.transform_batch(pandas_is_month_end) + + @property + def is_quarter_start(self) -> "ps.Series": + """ + Indicator for whether the date is the first day of a quarter. + + Returns + ------- + is_quarter_start : Series + The same type as the original data with boolean values. Series will + have the same name and index. + + See Also + -------- + quarter : Return the quarter of the date. + is_quarter_end : Similar property for indicating the quarter start. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> df = ps.DataFrame({'dates': pd.date_range("2017-03-30", + ... periods=4)}) + >>> df + dates + 0 2017-03-30 + 1 2017-03-31 + 2 2017-04-01 + 3 2017-04-02 + + >>> df.dates.dt.quarter + 0 1 + 1 1 + 2 2 + 3 2 + Name: dates, dtype: int64 + + >>> df.dates.dt.is_quarter_start + 0 False + 1 False + 2 True + 3 False + Name: dates, dtype: bool + """ + + def pandas_is_quarter_start(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_quarter_start + + return self._data.pandas_on_spark.transform_batch(pandas_is_quarter_start) + + @property + def is_quarter_end(self) -> "ps.Series": + """ + Indicator for whether the date is the last day of a quarter. + + Returns + ------- + is_quarter_end : Series + The same type as the original data with boolean values. Series will + have the same name and index. + + See Also + -------- + quarter : Return the quarter of the date. + is_quarter_start : Similar property indicating the quarter start. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> df = ps.DataFrame({'dates': pd.date_range("2017-03-30", + ... periods=4)}) + >>> df + dates + 0 2017-03-30 + 1 2017-03-31 + 2 2017-04-01 + 3 2017-04-02 + + >>> df.dates.dt.quarter + 0 1 + 1 1 + 2 2 + 3 2 + Name: dates, dtype: int64 + + >>> df.dates.dt.is_quarter_start + 0 False + 1 False + 2 True + 3 False + Name: dates, dtype: bool + """ + + def pandas_is_quarter_end(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_quarter_end + + return self._data.pandas_on_spark.transform_batch(pandas_is_quarter_end) + + @property + def is_year_start(self) -> "ps.Series": + """ + Indicate whether the date is the first day of a year. + + Returns + ------- + Series + The same type as the original data with boolean values. Series will + have the same name and index. + + See Also + -------- + is_year_end : Similar property indicating the last day of the year. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> dates = ps.Series(pd.date_range("2017-12-30", periods=3)) + >>> dates + 0 2017-12-30 + 1 2017-12-31 + 2 2018-01-01 + dtype: datetime64[ns] + + >>> dates.dt.is_year_start + 0 False + 1 False + 2 True + dtype: bool + """ + + def pandas_is_year_start(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_year_start + + return self._data.pandas_on_spark.transform_batch(pandas_is_year_start) + + @property + def is_year_end(self) -> "ps.Series": + """ + Indicate whether the date is the last day of the year. + + Returns + ------- + Series + The same type as the original data with boolean values. Series will + have the same name and index. + + See Also + -------- + is_year_start : Similar property indicating the start of the year. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> dates = ps.Series(pd.date_range("2017-12-30", periods=3)) + >>> dates + 0 2017-12-30 + 1 2017-12-31 + 2 2018-01-01 + dtype: datetime64[ns] + + >>> dates.dt.is_year_end + 0 False + 1 True + 2 False + dtype: bool + """ + + def pandas_is_year_end(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_year_end + + return self._data.pandas_on_spark.transform_batch(pandas_is_year_end) + + @property + def is_leap_year(self) -> "ps.Series": + """ + Boolean indicator if the date belongs to a leap year. + + A leap year is a year, which has 366 days (instead of 365) including + 29th of February as an intercalary day. + Leap years are years which are multiples of four with the exception + of years divisible by 100 but not by 400. + + Returns + ------- + Series + Booleans indicating if dates belong to a leap year. + + Examples + -------- + This method is available on Series with datetime values under + the ``.dt`` accessor. + + >>> dates_series = ps.Series(pd.date_range("2012-01-01", "2015-01-01", freq="Y")) + >>> dates_series + 0 2012-12-31 + 1 2013-12-31 + 2 2014-12-31 + dtype: datetime64[ns] + + >>> dates_series.dt.is_leap_year + 0 True + 1 False + 2 False + dtype: bool + """ + + def pandas_is_leap_year(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.dt.is_leap_year + + return self._data.pandas_on_spark.transform_batch(pandas_is_leap_year) + + @property + def daysinmonth(self) -> "ps.Series": + """ + The number of days in the month. + """ + + def pandas_daysinmonth(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.dt.daysinmonth + + return self._data.pandas_on_spark.transform_batch(pandas_daysinmonth) + + @property + def days_in_month(self) -> "ps.Series": + return self.daysinmonth + + days_in_month.__doc__ = daysinmonth.__doc__ + + # Methods + + @no_type_check + def tz_localize(self, tz) -> "ps.Series": + """ + Localize tz-naive Datetime column to tz-aware Datetime column. + """ + # Neither tz-naive or tz-aware datetime exists in Spark + raise NotImplementedError() + + @no_type_check + def tz_convert(self, tz) -> "ps.Series": + """ + Convert tz-aware Datetime column from one time zone to another. + """ + # tz-aware datetime doesn't exist in Spark + raise NotImplementedError() + + def normalize(self) -> "ps.Series": + """ + Convert times to midnight. + + The time component of the date-time is converted to midnight i.e. + 00:00:00. This is useful in cases, when the time does not matter. + Length is unaltered. The time zones are unaffected. + + This method is available on Series with datetime values under + the ``.dt`` accessor, and directly on Datetime Array. + + Returns + ------- + Series + The same type as the original data. Series will have the same + name and index. + + See Also + -------- + floor : Floor the series to the specified freq. + ceil : Ceil the series to the specified freq. + round : Round the series to the specified freq. + + Examples + -------- + >>> series = ps.Series(pd.Series(pd.date_range('2012-1-1 12:45:31', periods=3, freq='M'))) + >>> series.dt.normalize() + 0 2012-01-31 + 1 2012-02-29 + 2 2012-03-31 + dtype: datetime64[ns] + """ + + def pandas_normalize(s) -> ps.Series[np.datetime64]: # type: ignore[no-untyped-def] + return s.dt.normalize() + + return self._data.pandas_on_spark.transform_batch(pandas_normalize) + + def strftime(self, date_format: str) -> "ps.Series": + """ + Convert to a string Series using specified date_format. + + Return an series of formatted strings specified by date_format, which + supports the same string format as the python standard library. Details + of the string format can be found in the python string format + doc. + + Parameters + ---------- + date_format : str + Date format string (example: "%%Y-%%m-%%d"). + + Returns + ------- + Series + Series of formatted strings. + + See Also + -------- + to_datetime : Convert the given argument to datetime. + normalize : Return series with times to midnight. + round : Round the series to the specified freq. + floor : Floor the series to the specified freq. + + Examples + -------- + >>> series = ps.Series(pd.date_range(pd.Timestamp("2018-03-10 09:00"), + ... periods=3, freq='s')) + >>> series + 0 2018-03-10 09:00:00 + 1 2018-03-10 09:00:01 + 2 2018-03-10 09:00:02 + dtype: datetime64[ns] + + >>> series.dt.strftime('%B %d, %Y, %r') + 0 March 10, 2018, 09:00:00 AM + 1 March 10, 2018, 09:00:01 AM + 2 March 10, 2018, 09:00:02 AM + dtype: object + """ + + def pandas_strftime(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.dt.strftime(date_format) + + return self._data.pandas_on_spark.transform_batch(pandas_strftime) + + def round(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps.Series": + """ + Perform round operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to round the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + nonexistent : 'shift_forward', 'shift_backward, 'NaT', timedelta, default 'raise' + A nonexistent time does not exist in a particular timezone + where clocks moved forward due to DST. + + - 'shift_forward' will shift the nonexistent time forward to the + closest existing time + - 'shift_backward' will shift the nonexistent time backward to the + closest existing time + - 'NaT' will return NaT where there are nonexistent times + - timedelta objects will shift nonexistent times by the timedelta + - 'raise' will raise an NonExistentTimeError if there are + nonexistent times + + .. note:: this option only works with pandas 0.24.0+ + + Returns + ------- + Series + a Series with the same index for a Series. + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> series = ps.Series(pd.date_range('1/1/2018 11:59:00', periods=3, freq='min')) + >>> series + 0 2018-01-01 11:59:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 12:01:00 + dtype: datetime64[ns] + + >>> series.dt.round("H") + 0 2018-01-01 12:00:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 12:00:00 + dtype: datetime64[ns] + """ + + def pandas_round(s) -> ps.Series[np.datetime64]: # type: ignore[no-untyped-def] + return s.dt.round(freq, *args, **kwargs) + + return self._data.pandas_on_spark.transform_batch(pandas_round) + + def floor(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps.Series": + """ + Perform floor operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to floor the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + nonexistent : 'shift_forward', 'shift_backward, 'NaT', timedelta, default 'raise' + A nonexistent time does not exist in a particular timezone + where clocks moved forward due to DST. + + - 'shift_forward' will shift the nonexistent time forward to the + closest existing time + - 'shift_backward' will shift the nonexistent time backward to the + closest existing time + - 'NaT' will return NaT where there are nonexistent times + - timedelta objects will shift nonexistent times by the timedelta + - 'raise' will raise an NonExistentTimeError if there are + nonexistent times + + .. note:: this option only works with pandas 0.24.0+ + + Returns + ------- + Series + a Series with the same index for a Series. + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> series = ps.Series(pd.date_range('1/1/2018 11:59:00', periods=3, freq='min')) + >>> series + 0 2018-01-01 11:59:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 12:01:00 + dtype: datetime64[ns] + + >>> series.dt.floor("H") + 0 2018-01-01 11:00:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 12:00:00 + dtype: datetime64[ns] + """ + + def pandas_floor(s) -> ps.Series[np.datetime64]: # type: ignore[no-untyped-def] + return s.dt.floor(freq, *args, **kwargs) + + return self._data.pandas_on_spark.transform_batch(pandas_floor) + + def ceil(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps.Series": + """ + Perform ceil operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to round the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + nonexistent : 'shift_forward', 'shift_backward, 'NaT', timedelta, default 'raise' + A nonexistent time does not exist in a particular timezone + where clocks moved forward due to DST. + + - 'shift_forward' will shift the nonexistent time forward to the + closest existing time + - 'shift_backward' will shift the nonexistent time backward to the + closest existing time + - 'NaT' will return NaT where there are nonexistent times + - timedelta objects will shift nonexistent times by the timedelta + - 'raise' will raise an NonExistentTimeError if there are + nonexistent times + + .. note:: this option only works with pandas 0.24.0+ + + Returns + ------- + Series + a Series with the same index for a Series. + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> series = ps.Series(pd.date_range('1/1/2018 11:59:00', periods=3, freq='min')) + >>> series + 0 2018-01-01 11:59:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 12:01:00 + dtype: datetime64[ns] + + >>> series.dt.ceil("H") + 0 2018-01-01 12:00:00 + 1 2018-01-01 12:00:00 + 2 2018-01-01 13:00:00 + dtype: datetime64[ns] + """ + + def pandas_ceil(s) -> ps.Series[np.datetime64]: # type: ignore[no-untyped-def] + return s.dt.ceil(freq, *args, **kwargs) + + return self._data.pandas_on_spark.transform_batch(pandas_ceil) + + def month_name(self, locale: Optional[str] = None) -> "ps.Series": + """ + Return the month names of the series with specified locale. + + Parameters + ---------- + locale : str, optional + Locale determining the language in which to return the month name. + Default is English locale. + + Returns + ------- + Series + Series of month names. + + Examples + -------- + >>> series = ps.Series(pd.date_range(start='2018-01', freq='M', periods=3)) + >>> series + 0 2018-01-31 + 1 2018-02-28 + 2 2018-03-31 + dtype: datetime64[ns] + + >>> series.dt.month_name() + 0 January + 1 February + 2 March + dtype: object + """ + + def pandas_month_name(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.dt.month_name(locale=locale) + + return self._data.pandas_on_spark.transform_batch(pandas_month_name) + + def day_name(self, locale: Optional[str] = None) -> "ps.Series": + """ + Return the day names of the series with specified locale. + + Parameters + ---------- + locale : str, optional + Locale determining the language in which to return the day name. + Default is English locale. + + Returns + ------- + Series + Series of day names. + + Examples + -------- + >>> series = ps.Series(pd.date_range(start='2018-01-01', freq='D', periods=3)) + >>> series + 0 2018-01-01 + 1 2018-01-02 + 2 2018-01-03 + dtype: datetime64[ns] + + >>> series.dt.day_name() + 0 Monday + 1 Tuesday + 2 Wednesday + dtype: object + """ + + def pandas_day_name(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.dt.day_name(locale=locale) + + return self._data.pandas_on_spark.transform_batch(pandas_day_name) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.datetimes + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.datetimes.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.datetimes tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.datetimes, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/exceptions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/exceptions.py new file mode 100644 index 0000000..d93f0bf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/exceptions.py @@ -0,0 +1,150 @@ +# +# 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. +# + +""" +Exceptions/Errors used in pandas-on-Spark. +""" +from typing import Optional + + +class DataError(Exception): + pass + + +class SparkPandasIndexingError(Exception): + pass + + +def code_change_hint(pandas_function: Optional[str], spark_target_function: Optional[str]) -> str: + if pandas_function is not None and spark_target_function is not None: + return "You are trying to use pandas function {}, use spark function {}".format( + pandas_function, spark_target_function + ) + elif pandas_function is not None and spark_target_function is None: + return ( + "You are trying to use pandas function {}, checkout the spark " + "user guide to find a relevant function" + ).format(pandas_function) + elif pandas_function is None and spark_target_function is not None: + return "Use spark function {}".format(spark_target_function) + else: # both none + return "Checkout the spark user guide to find a relevant function" + + +class SparkPandasNotImplementedError(NotImplementedError): + def __init__( + self, + pandas_function: Optional[str] = None, + spark_target_function: Optional[str] = None, + description: str = "", + ): + self.pandas_source = pandas_function + self.spark_target = spark_target_function + hint = code_change_hint(pandas_function, spark_target_function) + if len(description) > 0: + description += " " + hint + else: + description = hint + super().__init__(description) + + +class PandasNotImplementedError(NotImplementedError): + def __init__( + self, + class_name: str, + method_name: Optional[str] = None, + arg_name: Optional[str] = None, + property_name: Optional[str] = None, + scalar_name: Optional[str] = None, + deprecated: bool = False, + reason: str = "", + ): + assert [method_name is not None, property_name is not None, scalar_name is not None].count( + True + ) == 1 + self.class_name = class_name + self.method_name = method_name + self.arg_name = arg_name + if method_name is not None: + if arg_name is not None: + msg = "The method `{0}.{1}()` does not support `{2}` parameter. {3}".format( + class_name, method_name, arg_name, reason + ) + else: + if deprecated: + msg = ( + "The method `{0}.{1}()` is deprecated in pandas and will therefore " + + "not be supported in pandas-on-Spark. {2}" + ).format(class_name, method_name, reason) + else: + if reason == "": + reason = " yet." + else: + reason = ". " + reason + msg = "The method `{0}.{1}()` is not implemented{2}".format( + class_name, method_name, reason + ) + elif scalar_name is not None: + msg = ( + "The scalar `{0}.{1}` is not reimplemented in pyspark.pandas;" + " use `pd.{1}`.".format(class_name, scalar_name) + ) + else: + if deprecated: + msg = ( + "The property `{0}.{1}()` is deprecated in pandas and will therefore " + + "not be supported in pandas-on-Spark. {2}" + ).format(class_name, property_name, reason) + else: + if reason == "": + reason = " yet." + else: + reason = ". " + reason + msg = "The property `{0}.{1}()` is not implemented{2}".format( + class_name, property_name, reason + ) + super().__init__(msg) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.exceptions + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.exceptions.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.exceptions tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.exceptions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/extensions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/extensions.py new file mode 100644 index 0000000..af4d160 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/extensions.py @@ -0,0 +1,388 @@ +# +# 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 Callable, Generic, Optional, Type, Union, TYPE_CHECKING +import warnings + +from pyspark.pandas._typing import T + +if TYPE_CHECKING: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.indexes import Index + from pyspark.pandas.series import Series + + +class CachedAccessor(Generic[T]): + """ + Custom property-like object. + + A descriptor for caching accessors: + + Parameters + ---------- + name : str + Namespace that accessor methods, properties, etc will be accessed under, e.g. "foo" for a + dataframe accessor yields the accessor ``df.foo`` + accessor: cls + Class with the extension methods. + + Notes + ----- + For accessor, the class's __init__ method assumes that you are registering an accessor for one + of ``Series``, ``DataFrame``, or ``Index``. + + This object is not meant to be instantiated directly. Instead, use register_dataframe_accessor, + register_series_accessor, or register_index_accessor. + + The pandas-on-Spark accessor is modified based on pandas.core.accessor. + """ + + def __init__(self, name: str, accessor: Type[T]) -> None: + self._name = name + self._accessor = accessor + + def __get__( + self, obj: Optional[Union["DataFrame", "Series", "Index"]], cls: Type[T] + ) -> Union[T, Type[T]]: + if obj is None: + return self._accessor + accessor_obj = self._accessor(obj) # type: ignore[call-arg] + object.__setattr__(obj, self._name, accessor_obj) + return accessor_obj + + +def _register_accessor( + name: str, cls: Union[Type["DataFrame"], Type["Series"], Type["Index"]] +) -> Callable[[Type[T]], Type[T]]: + """ + Register a custom accessor on {klass} objects. + + Parameters + ---------- + name : str + Name under which the accessor should be registered. A warning is issued if this name + conflicts with a preexisting attribute. + + Returns + ------- + callable + A class decorator. + + See Also + -------- + register_dataframe_accessor: Register a custom accessor on DataFrame objects + register_series_accessor: Register a custom accessor on Series objects + register_index_accessor: Register a custom accessor on Index objects + + Notes + ----- + When accessed, your accessor will be initialized with the pandas-on-Spark object the user + is interacting with. The code signature must be: + + .. code-block:: python + + def __init__(self, pandas_on_spark_obj): + # constructor logic + ... + + In the pandas API, if data passed to your accessor has an incorrect dtype, it's recommended to + raise an ``AttributeError`` for consistency purposes. In pandas-on-Spark, ``ValueError`` is more + frequently used to annotate when a value's datatype is unexpected for a given method/function. + + Ultimately, you can structure this however you like, but pandas-on-Spark would likely do + something like this: + + >>> ps.Series(['a', 'b']).dt + ... + Traceback (most recent call last): + ... + ValueError: Cannot call DatetimeMethods on type StringType() + + Note: This function is not meant to be used directly - instead, use register_dataframe_accessor, + register_series_accessor, or register_index_accessor. + """ + + def decorator(accessor: Type[T]) -> Type[T]: + if hasattr(cls, name): + msg = ( + "registration of accessor {0} under name '{1}' for type {2} is overriding " + "a preexisting attribute with the same name.".format(accessor, name, cls.__name__) + ) + + warnings.warn( + msg, + UserWarning, + stacklevel=2, + ) + setattr(cls, name, CachedAccessor(name, accessor)) + return accessor + + return decorator + + +def register_dataframe_accessor(name: str) -> Callable[[Type[T]], Type[T]]: + """ + Register a custom accessor with a DataFrame + + Parameters + ---------- + name : str + name used when calling the accessor after its registered + + Returns + ------- + callable + A class decorator. + + See Also + -------- + register_series_accessor: Register a custom accessor on Series objects + register_index_accessor: Register a custom accessor on Index objects + + Notes + ----- + When accessed, your accessor will be initialized with the pandas-on-Spark object the user + is interacting with. The accessor's init method should always ingest the object being accessed. + See the examples for the init signature. + + In the pandas API, if data passed to your accessor has an incorrect dtype, it's recommended to + raise an ``AttributeError`` for consistency purposes. In pandas-on-Spark, ``ValueError`` is more + frequently used to annotate when a value's datatype is unexpected for a given method/function. + + Ultimately, you can structure this however you like, but pandas-on-Spark would likely do + something like this: + + >>> ps.Series(['a', 'b']).dt + ... + Traceback (most recent call last): + ... + ValueError: Cannot call DatetimeMethods on type StringType() + + Examples + -------- + In your library code:: + + from pyspark.pandas.extensions import register_dataframe_accessor + + @register_dataframe_accessor("geo") + class GeoAccessor: + + def __init__(self, pandas_on_spark_obj): + self._obj = pandas_on_spark_obj + # other constructor logic + + @property + def center(self): + # return the geographic center point of this DataFrame + lat = self._obj.latitude + lon = self._obj.longitude + return (float(lon.mean()), float(lat.mean())) + + def plot(self): + # plot this array's data on a map + pass + + Then, in an ipython session:: + + >>> ## Import if the accessor is in the other file. + >>> # from my_ext_lib import GeoAccessor + >>> psdf = ps.DataFrame({"longitude": np.linspace(0,10), + ... "latitude": np.linspace(0, 20)}) + >>> psdf.geo.center # doctest: +SKIP + (5.0, 10.0) + + >>> psdf.geo.plot() # doctest: +SKIP + """ + from pyspark.pandas import DataFrame + + return _register_accessor(name, DataFrame) + + +def register_series_accessor(name: str) -> Callable[[Type[T]], Type[T]]: + """ + Register a custom accessor with a Series object + + Parameters + ---------- + name : str + name used when calling the accessor after its registered + + Returns + ------- + callable + A class decorator. + + See Also + -------- + register_dataframe_accessor: Register a custom accessor on DataFrame objects + register_index_accessor: Register a custom accessor on Index objects + + Notes + ----- + When accessed, your accessor will be initialized with the pandas-on-Spark object the user is + interacting with. The code signature must be:: + + def __init__(self, pandas_on_spark_obj): + # constructor logic + ... + + In the pandas API, if data passed to your accessor has an incorrect dtype, it's recommended to + raise an ``AttributeError`` for consistency purposes. In pandas-on-Spark, ``ValueError`` is more + frequently used to annotate when a value's datatype is unexpected for a given method/function. + + Ultimately, you can structure this however you like, but pandas-on-Spark would likely do + something like this: + + >>> ps.Series(['a', 'b']).dt + ... + Traceback (most recent call last): + ... + ValueError: Cannot call DatetimeMethods on type StringType() + + Examples + -------- + In your library code:: + + from pyspark.pandas.extensions import register_series_accessor + + @register_series_accessor("geo") + class GeoAccessor: + + def __init__(self, pandas_on_spark_obj): + self._obj = pandas_on_spark_obj + + @property + def is_valid(self): + # boolean check to see if series contains valid geometry + return True + + Then, in an ipython session:: + + >>> ## Import if the accessor is in the other file. + >>> # from my_ext_lib import GeoAccessor + >>> psdf = ps.DataFrame({"longitude": np.linspace(0,10), + ... "latitude": np.linspace(0, 20)}) + >>> psdf.longitude.geo.is_valid # doctest: +SKIP + True + """ + from pyspark.pandas import Series + + return _register_accessor(name, Series) + + +def register_index_accessor(name: str) -> Callable[[Type[T]], Type[T]]: + """ + Register a custom accessor with an Index + + Parameters + ---------- + name : str + name used when calling the accessor after its registered + + Returns + ------- + callable + A class decorator. + + See Also + -------- + register_dataframe_accessor: Register a custom accessor on DataFrame objects + register_series_accessor: Register a custom accessor on Series objects + + Notes + ----- + When accessed, your accessor will be initialized with the pandas-on-Spark object the user is + interacting with. The code signature must be:: + + def __init__(self, pandas_on_spark_obj): + # constructor logic + ... + + In the pandas API, if data passed to your accessor has an incorrect dtype, it's recommended to + raise an ``AttributeError`` for consistency purposes. In pandas-on-Spark, ``ValueError`` is more + frequently used to annotate when a value's datatype is unexpected for a given method/function. + + Ultimately, you can structure this however you like, but pandas-on-Spark would likely do + something like this: + + >>> ps.Series(['a', 'b']).dt + ... + Traceback (most recent call last): + ... + ValueError: Cannot call DatetimeMethods on type StringType() + + Examples + -------- + In your library code:: + + from pyspark.pandas.extensions import register_index_accessor + + @register_index_accessor("foo") + class CustomAccessor: + + def __init__(self, pandas_on_spark_obj): + self._obj = pandas_on_spark_obj + self.item = "baz" + + @property + def bar(self): + # return item value + return self.item + + Then, in an ipython session:: + + >>> ## Import if the accessor is in the other file. + >>> # from my_ext_lib import CustomAccessor + >>> psdf = ps.DataFrame({"longitude": np.linspace(0,10), + ... "latitude": np.linspace(0, 20)}) + >>> psdf.index.foo.bar # doctest: +SKIP + 'baz' + """ + from pyspark.pandas import Index + + return _register_accessor(name, Index) + + +def _test() -> None: + import os + import doctest + import sys + import numpy + from pyspark.sql import SparkSession + import pyspark.pandas.extensions + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.extensions.__dict__.copy() + globs["np"] = numpy + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.extensions tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.extensions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/frame.py new file mode 100644 index 0000000..2956ca7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/frame.py @@ -0,0 +1,13738 @@ +# +# 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. +# + +""" +A wrapper class for Spark DataFrame to behave like pandas DataFrame. +""" +from collections import defaultdict, namedtuple +from collections.abc import Mapping +import re +import warnings +import inspect +import json +import types +from functools import partial, reduce +import sys +from itertools import zip_longest, chain +from types import TracebackType +from typing import ( + Any, + Callable, + Dict, + Generic, + IO, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, + no_type_check, + TYPE_CHECKING, +) +import datetime + +import numpy as np +import pandas as pd +from pandas.api.types import ( # type: ignore[attr-defined] + is_bool_dtype, + is_list_like, + is_dict_like, + is_scalar, +) +from pandas.tseries.frequencies import DateOffset, to_offset + +if TYPE_CHECKING: + from pandas.io.formats.style import Styler + +from pandas.core.dtypes.common import infer_dtype_from_object +from pandas.core.accessor import CachedAccessor +from pandas.core.dtypes.inference import is_sequence +from pyspark import StorageLevel +from pyspark.sql import Column as PySparkColumn, DataFrame as PySparkDataFrame, functions as F +from pyspark.sql.functions import pandas_udf +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DoubleType, + NumericType, + Row, + StringType, + StructField, + StructType, + DecimalType, + TimestampType, + TimestampNTZType, +) +from pyspark.sql.window import Window + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import ( + Axis, + DataFrameOrSeries, + Dtype, + Label, + Name, + Scalar, + T, +) +from pyspark.pandas.accessors import PandasOnSparkFrameMethods +from pyspark.pandas.config import option_context, get_option +from pyspark.pandas.correlation import ( + compute, + CORRELATION_VALUE_1_COLUMN, + CORRELATION_VALUE_2_COLUMN, + CORRELATION_CORR_OUTPUT_COLUMN, + CORRELATION_COUNT_OUTPUT_COLUMN, +) +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.spark.accessors import SparkFrameMethods, CachedSparkFrameMethods +from pyspark.pandas.utils import ( + align_diff_frames, + column_labels_level, + combine_frames, + default_session, + is_name_like_tuple, + is_name_like_value, + is_testing, + name_like_string, + same_anchor, + scol_for, + validate_arguments_and_invoke_function, + validate_axis, + validate_bool_kwarg, + validate_how, + validate_mode, + verify_temp_column_name, + log_advice, +) +from pyspark.pandas.generic import Frame +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + HIDDEN_COLUMNS, + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, + SPARK_DEFAULT_INDEX_NAME, + SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_PATTERN, +) +from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame +from pyspark.pandas.typedef.typehints import ( + as_spark_type, + infer_return_type, + pandas_on_spark_type, + spark_type_to_pandas_dtype, + DataFrameType, + SeriesType, + ScalarType, + create_tuple_for_frame_type, +) +from pyspark.pandas.plot import PandasOnSparkPlotAccessor + +# For supporting Spark Connect +from pyspark.sql.utils import get_column_class, get_dataframe_class + +if TYPE_CHECKING: + from pyspark.sql._typing import OptionalPrimitiveType + + from pyspark.pandas.groupby import DataFrameGroupBy + from pyspark.pandas.resample import DataFrameResampler + from pyspark.pandas.indexes import Index + from pyspark.pandas.series import Series + + +# These regular expression patterns are compiled and defined here to avoid compiling the same +# pattern every time it is used in _repr_ and _repr_html_ in DataFrame. +# Two patterns basically seek the footer string from Pandas' +REPR_PATTERN = re.compile(r"\n\n\[(?P[0-9]+) rows x (?P[0-9]+) columns\]$") +REPR_HTML_PATTERN = re.compile( + r"\n\(?P[0-9]+) rows × (?P[0-9]+) columns\<\/p\>\n\<\/div\>$" +) + + +_flex_doc_FRAME = """ +Get {desc} of dataframe and other, element-wise (binary operator `{op_name}`). + +Equivalent to ``{equiv}``. With the reverse version, `{reverse}`. + +Among flexible wrappers (`add`, `sub`, `mul`, `div`) to +arithmetic operators: `+`, `-`, `*`, `/`, `//`. + +Parameters +---------- +other : scalar + Any single data + +Returns +------- +DataFrame + Result of the arithmetic operation. + +Examples +-------- +>>> df = ps.DataFrame({{'angles': [0, 3, 4], +... 'degrees': [360, 180, 360]}}, +... index=['circle', 'triangle', 'rectangle'], +... columns=['angles', 'degrees']) +>>> df + angles degrees +circle 0 360 +triangle 3 180 +rectangle 4 360 + +Add a scalar with operator version which returns the same +results. Also, the reverse version. + +>>> df + 1 + angles degrees +circle 1 361 +triangle 4 181 +rectangle 5 361 + +>>> df.add(1) + angles degrees +circle 1 361 +triangle 4 181 +rectangle 5 361 + +>>> df.add(df) + angles degrees +circle 0 720 +triangle 6 360 +rectangle 8 720 + +>>> df + df + df + angles degrees +circle 0 1080 +triangle 9 540 +rectangle 12 1080 + +>>> df.radd(1) + angles degrees +circle 1 361 +triangle 4 181 +rectangle 5 361 + +Divide and true divide by constant with reverse version. + +>>> df / 10 + angles degrees +circle 0.0 36.0 +triangle 0.3 18.0 +rectangle 0.4 36.0 + +>>> df.div(10) + angles degrees +circle 0.0 36.0 +triangle 0.3 18.0 +rectangle 0.4 36.0 + +>>> df.rdiv(10) + angles degrees +circle inf 0.027778 +triangle 3.333333 0.055556 +rectangle 2.500000 0.027778 + +>>> df.truediv(10) + angles degrees +circle 0.0 36.0 +triangle 0.3 18.0 +rectangle 0.4 36.0 + +>>> df.rtruediv(10) + angles degrees +circle inf 0.027778 +triangle 3.333333 0.055556 +rectangle 2.500000 0.027778 + +Subtract by constant with reverse version. + +>>> df - 1 + angles degrees +circle -1 359 +triangle 2 179 +rectangle 3 359 + +>>> df.sub(1) + angles degrees +circle -1 359 +triangle 2 179 +rectangle 3 359 + +>>> df.rsub(1) + angles degrees +circle 1 -359 +triangle -2 -179 +rectangle -3 -359 + +Multiply by constant with the reverse version. + +>>> df * 1 + angles degrees +circle 0 360 +triangle 3 180 +rectangle 4 360 + +>>> df.mul(1) + angles degrees +circle 0 360 +triangle 3 180 +rectangle 4 360 + +>>> df.rmul(1) + angles degrees +circle 0 360 +triangle 3 180 +rectangle 4 360 + +Floor Divide by constant with reverse version. + +>>> df // 10 + angles degrees +circle 0.0 36.0 +triangle 0.0 18.0 +rectangle 0.0 36.0 + +>>> df.floordiv(10) + angles degrees +circle 0.0 36.0 +triangle 0.0 18.0 +rectangle 0.0 36.0 + +>>> df.rfloordiv(10) # doctest: +SKIP + angles degrees +circle inf 0.0 +triangle 3.0 0.0 +rectangle 2.0 0.0 + +Mod by constant with reverse version. + +>>> df % 2 + angles degrees +circle 0 0 +triangle 1 0 +rectangle 0 0 + +>>> df.mod(2) + angles degrees +circle 0 0 +triangle 1 0 +rectangle 0 0 + +>>> df.rmod(2) + angles degrees +circle NaN 2 +triangle 2.0 2 +rectangle 2.0 2 + +Power by constant with reverse version. + +>>> df ** 2 + angles degrees +circle 0.0 129600.0 +triangle 9.0 32400.0 +rectangle 16.0 129600.0 + +>>> df.pow(2) + angles degrees +circle 0.0 129600.0 +triangle 9.0 32400.0 +rectangle 16.0 129600.0 + +>>> df.rpow(2) + angles degrees +circle 1.0 2.348543e+108 +triangle 8.0 1.532496e+54 +rectangle 16.0 2.348543e+108 +""" + + +class DataFrame(Frame, Generic[T]): + """ + pandas-on-Spark DataFrame that corresponds to pandas DataFrame logically. This holds Spark + DataFrame internally. + + :ivar _internal: an internal immutable Frame to manage metadata. + :type _internal: InternalFrame + + Parameters + ---------- + data : numpy ndarray (structured or homogeneous), dict, pandas DataFrame, + Spark DataFrame, pandas-on-Spark DataFrame or pandas-on-Spark Series. + Dict can contain Series, arrays, constants, or list-like objects + index : Index or array-like + Index to use for the resulting frame. Will default to RangeIndex if + no indexing information part of input data and no index provided + columns : Index or array-like + Column labels to use for the resulting frame. Will default to + RangeIndex (0, 1, 2, ..., n) if no column labels are provided + dtype : dtype, default None + Data type to force. Only a single dtype is allowed. If None, infer + copy : boolean, default False + Copy data from inputs. Only affects DataFrame / 2d ndarray input + + .. versionchanged:: 3.4.0 + Since 3.4.0, it deals with `data` and `index` in this approach: + 1, when `data` is a distributed dataset (Internal DataFrame/Spark DataFrame/ + pandas-on-Spark DataFrame/pandas-on-Spark Series), it will first parallelize + the `index` if necessary, and then try to combine the `data` and `index`; + Note that if `data` and `index` doesn't have the same anchor, then + `compute.ops_on_diff_frames` should be turned on; + 2, when `data` is a local dataset (Pandas DataFrame/numpy ndarray/list/etc), + it will first collect the `index` to driver if necessary, and then apply + the `pandas.DataFrame(...)` creation internally; + + Examples + -------- + Constructing DataFrame from a dictionary. + + >>> d = {'col1': [1, 2], 'col2': [3, 4]} + >>> df = ps.DataFrame(data=d, columns=['col1', 'col2']) + >>> df + col1 col2 + 0 1 3 + 1 2 4 + + Constructing DataFrame from pandas DataFrame + + >>> df = ps.DataFrame(pd.DataFrame(data=d, columns=['col1', 'col2'])) + >>> df + col1 col2 + 0 1 3 + 1 2 4 + + Notice that the inferred dtype is int64. + + >>> df.dtypes + col1 int64 + col2 int64 + dtype: object + + To enforce a single dtype: + + >>> df = ps.DataFrame(data=d, dtype=np.int8) + >>> df.dtypes + col1 int8 + col2 int8 + dtype: object + + Constructing DataFrame from numpy ndarray: + + >>> import numpy as np + >>> ps.DataFrame(data=np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]), + ... columns=['a', 'b', 'c', 'd', 'e']) + a b c d e + 0 1 2 3 4 5 + 1 6 7 8 9 0 + + Constructing DataFrame from numpy ndarray with Pandas index: + + >>> import numpy as np + >>> import pandas as pd + + >>> ps.DataFrame(data=np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]), + ... index=pd.Index([1, 4]), columns=['a', 'b', 'c', 'd', 'e']) + a b c d e + 1 1 2 3 4 5 + 4 6 7 8 9 0 + + Constructing DataFrame from numpy ndarray with pandas-on-Spark index: + + >>> import numpy as np + >>> import pandas as pd + >>> ps.DataFrame(data=np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]), + ... index=ps.Index([1, 4]), columns=['a', 'b', 'c', 'd', 'e']) + a b c d e + 1 1 2 3 4 5 + 4 6 7 8 9 0 + + Constructing DataFrame from Pandas DataFrame with Pandas index: + + >>> import numpy as np + >>> import pandas as pd + >>> pdf = pd.DataFrame(data=np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]), + ... columns=['a', 'b', 'c', 'd', 'e']) + >>> ps.DataFrame(data=pdf, index=pd.Index([1, 4])) + a b c d e + 1 6.0 7.0 8.0 9.0 0.0 + 4 NaN NaN NaN NaN NaN + + Constructing DataFrame from Pandas DataFrame with pandas-on-Spark index: + + >>> import numpy as np + >>> import pandas as pd + >>> pdf = pd.DataFrame(data=np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]), + ... columns=['a', 'b', 'c', 'd', 'e']) + >>> ps.DataFrame(data=pdf, index=ps.Index([1, 4])) + a b c d e + 1 6.0 7.0 8.0 9.0 0.0 + 4 NaN NaN NaN NaN NaN + + Constructing DataFrame from Spark DataFrame with Pandas index: + + >>> import pandas as pd + >>> sdf = spark.createDataFrame([("Data", 1), ("Bricks", 2)], ["x", "y"]) + >>> ps.DataFrame(data=sdf, index=pd.Index([0, 1, 2])) + Traceback (most recent call last): + ... + ValueError: Cannot combine the series or dataframe...'compute.ops_on_diff_frames' option. + + Enable 'compute.ops_on_diff_frames' to combine SparkDataFrame and Pandas index + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... ps.DataFrame(data=sdf, index=pd.Index([0, 1, 2])) + x y + 0 Data 1.0 + 1 Bricks 2.0 + 2 None NaN + + Constructing DataFrame from Spark DataFrame with pandas-on-Spark index: + + >>> import pandas as pd + >>> sdf = spark.createDataFrame([("Data", 1), ("Bricks", 2)], ["x", "y"]) + >>> ps.DataFrame(data=sdf, index=ps.Index([0, 1, 2])) + Traceback (most recent call last): + ... + ValueError: Cannot combine the series or dataframe...'compute.ops_on_diff_frames' option. + + Enable 'compute.ops_on_diff_frames' to combine Spark DataFrame and pandas-on-Spark index + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... ps.DataFrame(data=sdf, index=ps.Index([0, 1, 2])) + x y + 0 Data 1.0 + 1 Bricks 2.0 + 2 None NaN + """ + + def __init__( # type: ignore[no-untyped-def] + self, data=None, index=None, columns=None, dtype=None, copy=False + ): + SparkDataFrame = get_dataframe_class() + index_assigned = False + if isinstance(data, InternalFrame): + assert columns is None + assert dtype is None + assert not copy + if index is None: + internal = data + elif isinstance(data, SparkDataFrame): + assert columns is None + assert dtype is None + assert not copy + if index is None: + internal = InternalFrame(spark_frame=data, index_spark_columns=None) + elif isinstance(data, ps.DataFrame): + assert columns is None + assert dtype is None + assert not copy + if index is None: + internal = data._internal + elif isinstance(data, ps.Series): + assert dtype is None + assert not copy + # For pandas compatibility when `columns` contains only one valid column. + if columns is not None: + assert isinstance(columns, (dict, list, tuple)) + assert len(columns) == 1 + columns = list(columns.keys()) if isinstance(columns, dict) else columns + assert columns[0] == data._internal.data_spark_column_names[0] + if index is None: + internal = data.to_frame()._internal + else: + from pyspark.pandas.indexes.base import Index + + if index is not None and isinstance(index, Index): + # with local data, collect ps.Index to driver + # to avoid mismatched results between + # ps.DataFrame([1, 2], index=ps.Index([1, 2])) + # and + # pd.DataFrame([1, 2], index=pd.Index([1, 2])) + index = index._to_pandas() + + pdf = pd.DataFrame(data=data, index=index, columns=columns, dtype=dtype, copy=copy) + internal = InternalFrame.from_pandas(pdf) + index_assigned = True + + if index is not None and not index_assigned: + # TODO(SPARK-40226): Support MultiIndex + if isinstance(index, (ps.MultiIndex, pd.MultiIndex)): + raise ValueError("Cannot combine a Distributed Dataset with a MultiIndex") + + data_df = ps.DataFrame(data=data, index=None, columns=columns, dtype=dtype, copy=copy) + index_ps = ps.Index(index) + index_df = index_ps.to_frame() + + if same_anchor(data_df, index_df): + data_labels = data_df._internal.column_labels + data_pssers = [data_df._psser_for(label) for label in data_labels] + index_labels = index_df._internal.column_labels + index_pssers = [index_df._psser_for(label) for label in index_labels] + internal = data_df._internal.with_new_columns(data_pssers + index_pssers) + + combined = ps.DataFrame(internal).set_index(index_labels) + combined.index.name = index_ps.name + else: + # drop un-matched rows in `data` + # note that `combine_frames` cannot work with a MultiIndex for now + combined = combine_frames(data_df, index_df, how="right") + combined_labels = combined._internal.column_labels + index_labels = [label for label in combined_labels if label[0] == "that"] + combined = combined.set_index(index_labels) + + combined._internal._column_labels = data_df._internal.column_labels + combined._internal._column_label_names = data_df._internal._column_label_names + combined._internal._index_names = index_df._internal.column_labels + combined.index.name = index_ps.name + + internal = combined._internal + + object.__setattr__(self, "_internal_frame", internal) + + @property + def _pssers(self) -> Dict[Label, "Series"]: + """Return a dict of column label -> Series which anchors `self`.""" + from pyspark.pandas.series import Series + + if not hasattr(self, "_psseries"): + object.__setattr__( + self, + "_psseries", + {label: Series(data=self, index=label) for label in self._internal.column_labels}, + ) + else: + psseries = cast(Dict[Label, Series], self._psseries) # type: ignore[has-type] + assert len(self._internal.column_labels) == len(psseries), ( + len(self._internal.column_labels), + len(psseries), + ) + if any(self is not psser._psdf for psser in psseries.values()): + # Refresh the dict to contain only Series anchoring `self`. + self._psseries = { + label: ( + psseries[label] + if self is psseries[label]._psdf + else Series(data=self, index=label) + ) + for label in self._internal.column_labels + } + return self._psseries + + @property + def _internal(self) -> InternalFrame: + return cast(InternalFrame, self._internal_frame) # type: ignore[has-type] + + def _update_internal_frame( + self, + internal: InternalFrame, + check_same_anchor: bool = True, + anchor_force_disconnect: bool = False, + ) -> None: + """ + Update InternalFrame with the given one. + + If the column_label is changed or the new InternalFrame is not the same `anchor` or the + `anchor_force_disconnect` flag is set to True, disconnect the original anchor and create + a new one. + + If `check_same_anchor` is `False`, checking whether the same anchor is ignored + and force to update the InternalFrame, e.g., replacing the internal with the resolved_copy, + updating the underlying Spark DataFrame which need to combine a different Spark DataFrame. + + Parameters + ---------- + internal : InternalFrame + The new InternalFrame + check_same_anchor : bool + Whether checking the same anchor + anchor_force_disconnect : bool + Force to disconnect the original anchor and create a new one + """ + from pyspark.pandas.series import Series + + if hasattr(self, "_psseries"): + psseries = {} + + for old_label, new_label in zip_longest( + self._internal.column_labels, internal.column_labels + ): + if old_label is not None: + psser = self._pssers[old_label] + + renamed = old_label != new_label + not_same_anchor = check_same_anchor and not same_anchor(internal, psser) + + if renamed or not_same_anchor or anchor_force_disconnect: + psdf: DataFrame = DataFrame(self._internal.select_column(old_label)) + psser._update_anchor(psdf) + psser = None + else: + psser = None + if new_label is not None: + if psser is None: + psser = Series(data=self, index=new_label) + psseries[new_label] = psser + + self._psseries = psseries + + self._internal_frame = internal + + if hasattr(self, "_repr_pandas_cache"): + del self._repr_pandas_cache + + @property + def ndim(self) -> int: + """ + Return an int representing the number of array dimensions. + + return 2 for DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=['cobra', 'viper', None], + ... columns=['max_speed', 'shield']) + >>> df # doctest: +SKIP + max_speed shield + cobra 1 2 + viper 4 5 + None 7 8 + >>> df.ndim + 2 + """ + return 2 + + @property + def axes(self) -> List: + """ + Return a list representing the axes of the DataFrame. + + It has the row axis labels and column axis labels as the only members. + They are returned in that order. + + Examples + -------- + + >>> df = ps.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.axes # doctest: +SKIP + [Int64Index([0, 1], dtype='int64'), Index(['col1', 'col2'], dtype='object')] + """ + return [self.index, self.columns] + + def _reduce_for_stat_function( + self, + sfun: Callable[["Series"], PySparkColumn], + name: str, + axis: Optional[Axis] = None, + numeric_only: bool = True, + skipna: bool = True, + **kwargs: Any, + ) -> "Series": + """ + Applies sfun to each column and returns a pd.Series where the number of rows equals the + number of columns. + + Parameters + ---------- + sfun : either an 1-arg function that takes a Column and returns a Column, or + a 2-arg function that takes a Column and its DataType and returns a Column. + axis: used only for sanity check because the series only supports index axis. + name : original pandas API name. + axis : axis to apply. 0 or 1, or 'index' or 'columns. + numeric_only : bool, default True + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. Only 'DataFrame.count' uses this parameter + currently. + skipna : bool, default True + Exclude NA/null values when computing the result. + """ + from pyspark.pandas.series import Series, first_series + + axis = validate_axis(axis) + if axis == 0: + min_count = kwargs.get("min_count", 0) + + exprs = [F.lit(None).cast(StringType()).alias(SPARK_DEFAULT_INDEX_NAME)] + new_column_labels = [] + for label in self._internal.column_labels: + psser = self._psser_for(label) + + is_numeric_or_boolean = isinstance( + psser.spark.data_type, (NumericType, BooleanType) + ) + keep_column = not numeric_only or is_numeric_or_boolean + + if keep_column: + if not skipna and get_option("compute.eager_check") and psser.hasnans: + scol = F.first(F.lit(np.nan)) + else: + scol = sfun(psser) + + if min_count > 0: + scol = F.when(Frame._count_expr(psser) >= min_count, scol) + + exprs.append(scol.alias(name_like_string(label))) + new_column_labels.append(label) + + if len(exprs) == 1: + return Series([]) + + sdf = self._internal.spark_frame.select(*exprs) + + # The data is expected to be small so it's fine to transpose/use the default index. + with ps.option_context("compute.max_rows", 1): + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + column_labels=new_column_labels, + column_label_names=self._internal.column_label_names, + ) + return first_series(DataFrame(internal).transpose()) + + else: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + limit = get_option("compute.shortcut_limit") + pdf = self.head(limit + 1)._to_internal_pandas() + pser = getattr(pdf, name)(axis=axis, numeric_only=numeric_only, **kwargs) + if len(pdf) <= limit: + return Series(pser) + + @pandas_udf(returnType=as_spark_type(pser.dtype.type)) # type: ignore[call-overload] + def calculate_columns_axis(*cols: pd.Series) -> pd.Series: + return getattr(pd.concat(cols, axis=1), name)( + axis=axis, numeric_only=numeric_only, **kwargs + ) + + column_name = verify_temp_column_name( + self._internal.spark_frame.select(self._internal.index_spark_columns), + "__calculate_columns_axis__", + ) + sdf = self._internal.spark_frame.select( + self._internal.index_spark_columns + + [calculate_columns_axis(*self._internal.data_spark_columns).alias(column_name)] + ) + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + return first_series(DataFrame(internal)).rename(pser.name) + + def _psser_for(self, label: Label) -> "Series": + """ + Create Series with a proper column label. + + The given label must be verified to exist in `InternalFrame.column_labels`. + + For example, in some method, self is like: + + >>> self = ps.range(3) + + `self._psser_for(label)` can be used with `InternalFrame.column_labels`: + + >>> self._psser_for(self._internal.column_labels[0]) + 0 0 + 1 1 + 2 2 + Name: id, dtype: int64 + + `self._psser_for(label)` must not be used directly with user inputs. + In that case, `self[label]` should be used instead, which checks the label exists or not: + + >>> self['id'] + 0 0 + 1 1 + 2 2 + Name: id, dtype: int64 + """ + return self._pssers[label] + + def _apply_series_op( + self, + op: Callable[["Series"], Union["Series", PySparkColumn]], + should_resolve: bool = False, + ) -> "DataFrame": + applied = [] + for label in self._internal.column_labels: + applied.append(op(self._psser_for(label))) + internal = self._internal.with_new_columns(applied) + if should_resolve: + internal = internal.resolved_copy + return DataFrame(internal) + + # Arithmetic Operators + def _map_series_op(self, op: str, other: Any) -> "DataFrame": + from pyspark.pandas.base import IndexOpsMixin + + if not isinstance(other, DataFrame) and ( + isinstance(other, IndexOpsMixin) or is_sequence(other) + ): + raise TypeError( + "%s with a sequence is currently not supported; " + "however, got %s." % (op, type(other).__name__) + ) + + if isinstance(other, DataFrame): + if self._internal.column_labels_level != other._internal.column_labels_level: + raise ValueError("cannot join with no overlapping index names") + + if not same_anchor(self, other): + # Different DataFrames + def apply_op( + psdf: DataFrame, + this_column_labels: List[Label], + that_column_labels: List[Label], + ) -> Iterator[Tuple["Series", Label]]: + for this_label, that_label in zip(this_column_labels, that_column_labels): + yield ( + getattr(psdf._psser_for(this_label), op)( + psdf._psser_for(that_label) + ).rename(this_label), + this_label, + ) + + return align_diff_frames(apply_op, self, other, fillna=True, how="full") + else: + applied = [] + column_labels = [] + for label in self._internal.column_labels: + if label in other._internal.column_labels: + applied.append(getattr(self._psser_for(label), op)(other._psser_for(label))) + else: + applied.append( + F.lit(None) + .cast(self._internal.spark_type_for(label)) + .alias(name_like_string(label)) + ) + column_labels.append(label) + for label in other._internal.column_labels: + if label not in column_labels: + applied.append( + F.lit(None) + .cast(other._internal.spark_type_for(label)) + .alias(name_like_string(label)) + ) + column_labels.append(label) + internal = self._internal.with_new_columns(applied, column_labels=column_labels) + return DataFrame(internal) + else: + return self._apply_series_op(lambda psser: getattr(psser, op)(other)) + + def __add__(self, other: Any) -> "DataFrame": + return self._map_series_op("add", other) + + def __radd__(self, other: Any) -> "DataFrame": + return self._map_series_op("radd", other) + + def __truediv__(self, other: Any) -> "DataFrame": + return self._map_series_op("truediv", other) + + def __rtruediv__(self, other: Any) -> "DataFrame": + return self._map_series_op("rtruediv", other) + + def __mul__(self, other: Any) -> "DataFrame": + return self._map_series_op("mul", other) + + def __rmul__(self, other: Any) -> "DataFrame": + return self._map_series_op("rmul", other) + + def __sub__(self, other: Any) -> "DataFrame": + return self._map_series_op("sub", other) + + def __rsub__(self, other: Any) -> "DataFrame": + return self._map_series_op("rsub", other) + + def __pow__(self, other: Any) -> "DataFrame": + return self._map_series_op("pow", other) + + def __rpow__(self, other: Any) -> "DataFrame": + return self._map_series_op("rpow", other) + + def __mod__(self, other: Any) -> "DataFrame": + return self._map_series_op("mod", other) + + def __rmod__(self, other: Any) -> "DataFrame": + return self._map_series_op("rmod", other) + + def __floordiv__(self, other: Any) -> "DataFrame": + return self._map_series_op("floordiv", other) + + def __rfloordiv__(self, other: Any) -> "DataFrame": + return self._map_series_op("rfloordiv", other) + + def __abs__(self) -> "DataFrame": + return self._apply_series_op(lambda psser: abs(psser)) + + def __neg__(self) -> "DataFrame": + return self._apply_series_op(lambda psser: -psser) + + def add(self, other: Any) -> "DataFrame": + return self + other + + # create accessor for plot + plot = CachedAccessor("plot", PandasOnSparkPlotAccessor) + + # create accessor for Spark related methods. + spark = CachedAccessor("spark", SparkFrameMethods) + + # create accessor for pandas-on-Spark specific methods. + pandas_on_spark = CachedAccessor("pandas_on_spark", PandasOnSparkFrameMethods) + + # keep the name "koalas" for backward compatibility. + koalas = CachedAccessor("koalas", PandasOnSparkFrameMethods) + + @no_type_check + def hist(self, bins=10, **kwds): + return self.plot.hist(bins, **kwds) + + hist.__doc__ = PandasOnSparkPlotAccessor.hist.__doc__ + + @no_type_check + def boxplot(self, **kwds): + return self.plot.box(**kwds) + + boxplot.__doc__ = PandasOnSparkPlotAccessor.box.__doc__ + + @no_type_check + def kde(self, bw_method=None, ind=None, **kwds): + return self.plot.kde(bw_method, ind, **kwds) + + kde.__doc__ = PandasOnSparkPlotAccessor.kde.__doc__ + + add.__doc__ = _flex_doc_FRAME.format( + desc="Addition", op_name="+", equiv="dataframe + other", reverse="radd" + ) + + def radd(self, other: Any) -> "DataFrame": + return other + self + + radd.__doc__ = _flex_doc_FRAME.format( + desc="Addition", op_name="+", equiv="other + dataframe", reverse="add" + ) + + def div(self, other: Any) -> "DataFrame": + return self / other + + div.__doc__ = _flex_doc_FRAME.format( + desc="Floating division", op_name="/", equiv="dataframe / other", reverse="rdiv" + ) + + divide = div + + def rdiv(self, other: Any) -> "DataFrame": + return other / self + + rdiv.__doc__ = _flex_doc_FRAME.format( + desc="Floating division", op_name="/", equiv="other / dataframe", reverse="div" + ) + + def truediv(self, other: Any) -> "DataFrame": + return self / other + + truediv.__doc__ = _flex_doc_FRAME.format( + desc="Floating division", op_name="/", equiv="dataframe / other", reverse="rtruediv" + ) + + def rtruediv(self, other: Any) -> "DataFrame": + return other / self + + rtruediv.__doc__ = _flex_doc_FRAME.format( + desc="Floating division", op_name="/", equiv="other / dataframe", reverse="truediv" + ) + + def mul(self, other: Any) -> "DataFrame": + return self * other + + mul.__doc__ = _flex_doc_FRAME.format( + desc="Multiplication", op_name="*", equiv="dataframe * other", reverse="rmul" + ) + + multiply = mul + + def rmul(self, other: Any) -> "DataFrame": + return other * self + + rmul.__doc__ = _flex_doc_FRAME.format( + desc="Multiplication", op_name="*", equiv="other * dataframe", reverse="mul" + ) + + def sub(self, other: Any) -> "DataFrame": + return self - other + + sub.__doc__ = _flex_doc_FRAME.format( + desc="Subtraction", op_name="-", equiv="dataframe - other", reverse="rsub" + ) + + subtract = sub + + def rsub(self, other: Any) -> "DataFrame": + return other - self + + rsub.__doc__ = _flex_doc_FRAME.format( + desc="Subtraction", op_name="-", equiv="other - dataframe", reverse="sub" + ) + + def mod(self, other: Any) -> "DataFrame": + return self % other + + mod.__doc__ = _flex_doc_FRAME.format( + desc="Modulo", op_name="%", equiv="dataframe % other", reverse="rmod" + ) + + def rmod(self, other: Any) -> "DataFrame": + return other % self + + rmod.__doc__ = _flex_doc_FRAME.format( + desc="Modulo", op_name="%", equiv="other % dataframe", reverse="mod" + ) + + def pow(self, other: Any) -> "DataFrame": + return self**other + + pow.__doc__ = _flex_doc_FRAME.format( + desc="Exponential power of series", op_name="**", equiv="dataframe ** other", reverse="rpow" + ) + + def rpow(self, other: Any) -> "DataFrame": + return other**self + + rpow.__doc__ = _flex_doc_FRAME.format( + desc="Exponential power", op_name="**", equiv="other ** dataframe", reverse="pow" + ) + + def floordiv(self, other: Any) -> "DataFrame": + return self // other + + floordiv.__doc__ = _flex_doc_FRAME.format( + desc="Integer division", op_name="//", equiv="dataframe // other", reverse="rfloordiv" + ) + + def rfloordiv(self, other: Any) -> "DataFrame": + return other // self + + rfloordiv.__doc__ = _flex_doc_FRAME.format( + desc="Integer division", op_name="//", equiv="other // dataframe", reverse="floordiv" + ) + + # Comparison Operators + def __eq__(self, other: Any) -> "DataFrame": # type: ignore[override] + return self._map_series_op("eq", other) + + def __ne__(self, other: Any) -> "DataFrame": # type: ignore[override] + return self._map_series_op("ne", other) + + def __lt__(self, other: Any) -> "DataFrame": + return self._map_series_op("lt", other) + + def __le__(self, other: Any) -> "DataFrame": + return self._map_series_op("le", other) + + def __ge__(self, other: Any) -> "DataFrame": + return self._map_series_op("ge", other) + + def __gt__(self, other: Any) -> "DataFrame": + return self._map_series_op("gt", other) + + def eq(self, other: Any) -> "DataFrame": + """ + Compare if the current value is equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.eq(1) + a b + a True True + b False False + c False True + d False False + """ + return self == other + + equals = eq + + def gt(self, other: Any) -> "DataFrame": + """ + Compare if the current value is greater than the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.gt(2) + a b + a False False + b False False + c True False + d True False + """ + return self > other + + def ge(self, other: Any) -> "DataFrame": + """ + Compare if the current value is greater than or equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.ge(1) + a b + a True True + b True False + c True True + d True False + """ + return self >= other + + def lt(self, other: Any) -> "DataFrame": + """ + Compare if the current value is less than the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.lt(1) + a b + a False False + b False False + c False False + d False False + """ + return self < other + + def le(self, other: Any) -> "DataFrame": + """ + Compare if the current value is less than or equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.le(2) + a b + a True True + b True False + c False True + d False False + """ + return self <= other + + def ne(self, other: Any) -> "DataFrame": + """ + Compare if the current value is not equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.ne(1) + a b + a False False + b True True + c True False + d True True + """ + return self != other + + def applymap(self, func: Callable[[Any], Any]) -> "DataFrame": + """ + Apply a function to a Dataframe elementwise. + + This method applies a function that accepts and returns a scalar + to every element of a DataFrame. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def square(x) -> np.int32: + ... return x ** 2 + + pandas-on-Spark uses return type hints and does not try to infer the type. + + Parameters + ---------- + func : callable + Python function returns a single value from a single value. + + Returns + ------- + DataFrame + Transformed DataFrame. + + Examples + -------- + >>> df = ps.DataFrame([[1, 2.12], [3.356, 4.567]]) + >>> df + 0 1 + 0 1.000 2.120 + 1 3.356 4.567 + + >>> def str_len(x) -> int: + ... return len(str(x)) + >>> df.applymap(str_len) + 0 1 + 0 3 4 + 1 5 5 + + >>> def power(x) -> float: + ... return x ** 2 + >>> df.applymap(power) + 0 1 + 0 1.000000 4.494400 + 1 11.262736 20.857489 + + You can omit type hints and let pandas-on-Spark infer its type. + + >>> df.applymap(lambda x: x ** 2) + 0 1 + 0 1.000000 4.494400 + 1 11.262736 20.857489 + """ + + # TODO: We can implement shortcut theoretically since it creates new DataFrame + # anyway and we don't have to worry about operations on different DataFrames. + return self._apply_series_op(lambda psser: psser.apply(func)) + + # TODO: not all arguments are implemented comparing to pandas' for now. + def aggregate(self, func: Union[List[str], Dict[Name, List[str]]]) -> "DataFrame": + """Aggregate using one or more operations over the specified axis. + + Parameters + ---------- + func : dict or a list + a dict mapping from column name (string) to + aggregate functions (list of strings). + If a list is given, the aggregation is performed against + all columns. + + Returns + ------- + DataFrame + + Notes + ----- + `agg` is an alias for `aggregate`. Use the alias. + + See Also + -------- + DataFrame.apply : Invoke function on DataFrame. + DataFrame.transform : Only perform transforming type operations. + DataFrame.groupby : Perform operations over groups. + Series.aggregate : The equivalent function for Series. + + Examples + -------- + >>> df = ps.DataFrame([[1, 2, 3], + ... [4, 5, 6], + ... [7, 8, 9], + ... [np.nan, np.nan, np.nan]], + ... columns=['A', 'B', 'C']) + + >>> df + A B C + 0 1.0 2.0 3.0 + 1 4.0 5.0 6.0 + 2 7.0 8.0 9.0 + 3 NaN NaN NaN + + Aggregate these functions over the rows. + + >>> df.agg(['sum', 'min'])[['A', 'B', 'C']].sort_index() + A B C + min 1.0 2.0 3.0 + sum 12.0 15.0 18.0 + + Different aggregations per column. + + >>> df.agg({'A' : ['sum', 'min'], 'B' : ['min', 'max']})[['A', 'B']].sort_index() + A B + max NaN 8.0 + min 1.0 2.0 + sum 12.0 NaN + + For multi-index columns: + + >>> df.columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + >>> df.agg(['sum', 'min'])[[("X", "A"), ("X", "B"), ("Y", "C")]].sort_index() + X Y + A B C + min 1.0 2.0 3.0 + sum 12.0 15.0 18.0 + + >>> aggregated = df.agg({("X", "A") : ['sum', 'min'], ("X", "B") : ['min', 'max']}) + >>> aggregated[[("X", "A"), ("X", "B")]].sort_index() # doctest: +NORMALIZE_WHITESPACE + X + A B + max NaN 8.0 + min 1.0 2.0 + sum 12.0 NaN + """ + from pyspark.pandas.groupby import GroupBy + + if isinstance(func, list): + if all((isinstance(f, str) for f in func)): + func = dict([(column, func) for column in self.columns]) + else: + raise ValueError( + "If the given function is a list, it " + "should only contains function names as strings." + ) + + if not isinstance(func, dict) or not all( + is_name_like_value(key) + and ( + isinstance(value, str) + or (isinstance(value, list) and all(isinstance(v, str) for v in value)) + ) + for key, value in func.items() + ): + raise ValueError( + "aggs must be a dict mapping from column name to aggregate " + "functions (string or list of strings)." + ) + + with option_context("compute.default_index_type", "distributed"): + psdf: DataFrame = DataFrame(GroupBy._spark_groupby(self, func)) + + # The codes below basically convert: + # + # A B + # sum min min max + # 0 12.0 1.0 2.0 8.0 + # + # to: + # A B + # max NaN 8.0 + # min 1.0 2.0 + # sum 12.0 NaN + # + # Aggregated output is usually pretty much small. + + return psdf.stack().droplevel(0)[list(func.keys())] + + agg = aggregate + + def corr(self, method: str = "pearson", min_periods: Optional[int] = None) -> "DataFrame": + """ + Compute pairwise correlation of columns, excluding NA/null values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + method : {'pearson', 'spearman', 'kendall'} + * pearson : standard correlation coefficient + * spearman : Spearman rank correlation + * kendall : Kendall Tau correlation coefficient + + .. versionchanged:: 3.4.0 + support 'kendall' for method parameter + min_periods : int, optional + Minimum number of observations required per pair of columns + to have a valid result. + + .. versionadded:: 3.4.0 + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.corrwith + Series.corr + + Notes + ----- + 1. Pearson, Kendall and Spearman correlation are currently computed using pairwise + complete observations. + + 2. The complexity of Kendall correlation is O(#row * #row), if the dataset is too + large, sampling ahead of correlation computation is recommended. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df.corr('pearson') + dogs cats + dogs 1.000000 -0.851064 + cats -0.851064 1.000000 + + >>> df.corr('spearman') + dogs cats + dogs 1.000000 -0.948683 + cats -0.948683 1.000000 + + >>> df.corr('kendall') + dogs cats + dogs 1.000000 -0.912871 + cats -0.912871 1.000000 + """ + if method not in ["pearson", "spearman", "kendall"]: + raise ValueError(f"Invalid method {method}") + if min_periods is not None and not isinstance(min_periods, int): + raise TypeError(f"Invalid min_periods type {type(min_periods).__name__}") + + min_periods = 1 if min_periods is None else min_periods + internal = self._internal.resolved_copy + numeric_labels = [ + label + for label in internal.column_labels + if isinstance(internal.spark_type_for(label), (NumericType, BooleanType)) + ] + numeric_scols = [ + 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, "__corr_index_1_temp_column__") + index_2_col_name = verify_temp_column_name(sdf, "__corr_index_2_temp_column__") + + # simple dataset + # +---+---+----+ + # | A| B| C| + # +---+---+----+ + # | 1| 2| 3.0| + # | 4| 1|NULL| + # +---+---+----+ + + pair_scols = [] + 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(CORRELATION_VALUE_1_COLUMN), + numeric_scols[j].alias(CORRELATION_VALUE_2_COLUMN), + ) + ) + + # +-------------------+-------------------+-------------------+-------------------+ + # |__tmp_index_1_col__|__tmp_index_2_col__|__tmp_value_1_col__|__tmp_value_2_col__| + # +-------------------+-------------------+-------------------+-------------------+ + # | 0| 0| 1.0| 1.0| + # | 0| 1| 1.0| 2.0| + # | 0| 2| 1.0| 3.0| + # | 1| 1| 2.0| 2.0| + # | 1| 2| 2.0| 3.0| + # | 2| 2| 3.0| 3.0| + # | 0| 0| 4.0| 4.0| + # | 0| 1| 4.0| 1.0| + # | 0| 2| NULL| NULL| + # | 1| 1| 1.0| 1.0| + # | 1| 2| NULL| NULL| + # | 2| 2| NULL| NULL| + # +-------------------+-------------------+-------------------+-------------------+ + sdf = sdf.select(F.inline(F.array(*pair_scols))) + + sdf = compute(sdf=sdf, groupKeys=[index_1_col_name, index_2_col_name], method=method) + if method == "kendall": + sdf = sdf.withColumn( + CORRELATION_CORR_OUTPUT_COLUMN, + F.when(F.col(index_1_col_name) == F.col(index_2_col_name), F.lit(1.0)).otherwise( + F.col(CORRELATION_CORR_OUTPUT_COLUMN) + ), + ) + + sdf = sdf.withColumn( + CORRELATION_CORR_OUTPUT_COLUMN, + F.when(F.col(CORRELATION_COUNT_OUTPUT_COLUMN) < min_periods, F.lit(None)).otherwise( + F.col(CORRELATION_CORR_OUTPUT_COLUMN) + ), + ) + + # +-------------------+-------------------+----------------+ + # |__tmp_index_1_col__|__tmp_index_2_col__|__tmp_corr_col__| + # +-------------------+-------------------+----------------+ + # | 2| 2| NULL| + # | 1| 2| NULL| + # | 2| 1| NULL| + # | 1| 1| 1.0| + # | 0| 0| 1.0| + # | 0| 1| -1.0| + # | 1| 0| -1.0| + # | 0| 2| NULL| + # | 2| 0| NULL| + # +-------------------+-------------------+----------------+ + + auxiliary_col_name = verify_temp_column_name(sdf, "__corr_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.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(CORRELATION_CORR_OUTPUT_COLUMN), + ) + + # +-------------------+--------------------+ + # |__tmp_index_1_col__| __tmp_array_col__| + # +-------------------+--------------------+ + # | 0|[{0, 1.0}, {1, -1...| + # | 1|[{0, -1.0}, {1, 1...| + # | 2|[{0, null}, {1, n...| + # +-------------------+--------------------+ + array_col_name = verify_temp_column_name(sdf, "__corr_array_temp_column__") + sdf = ( + sdf.groupby(index_1_col_name) + .agg( + F.array_sort( + F.collect_list( + F.struct(F.col(index_2_col_name), F.col(CORRELATION_CORR_OUTPUT_COLUMN)) + ) + ).alias(array_col_name) + ) + .orderBy(index_1_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}.{CORRELATION_CORR_OUTPUT_COLUMN}"), + ) + + index_col_names: List[str] = [] + if internal.column_labels_level > 1: + for level in range(0, internal.column_labels_level): + index_col_name = SPARK_INDEX_NAME_FORMAT(level) + indices = [label[level] for label in numeric_labels] + sdf = sdf.withColumn(index_col_name, F.get(F.lit(indices), F.col(index_1_col_name))) + index_col_names.append(index_col_name) + else: + sdf = sdf.withColumn( + SPARK_DEFAULT_INDEX_NAME, + F.get(F.lit(numeric_col_names), F.col(index_1_col_name)), + ) + index_col_names = [SPARK_DEFAULT_INDEX_NAME] + + sdf = sdf.select(*index_col_names, *numeric_col_names) + + return DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, index_col_name) for index_col_name in index_col_names + ], + column_labels=numeric_labels, + column_label_names=internal.column_label_names, + ) + ) + + def corrwith( + self, other: DataFrameOrSeries, axis: Axis = 0, drop: bool = False, method: str = "pearson" + ) -> "Series": + """ + Compute pairwise correlation. + + Pairwise correlation is computed between rows or columns of + DataFrame with rows or columns of Series or DataFrame. DataFrames + are first aligned along both axes before computing the + correlations. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + other : DataFrame, Series + Object with which to compute correlations. + axis : int, default 0 or 'index' + Can only be set to 0 now. + drop : bool, default False + Drop missing indices from result. + method : {'pearson', 'spearman', 'kendall'} + * pearson : standard correlation coefficient + * spearman : Spearman rank correlation + * kendall : Kendall Tau correlation coefficient + + Returns + ------- + Series + Pairwise correlations. + + See Also + -------- + DataFrame.corr : Compute pairwise correlation of columns. + + Examples + -------- + >>> df1 = ps.DataFrame({ + ... "A":[1, 5, 7, 8], + ... "X":[5, 8, 4, 3], + ... "C":[10, 4, 9, 3]}) + >>> df1.corrwith(df1[["X", "C"]]).sort_index() + A NaN + C 1.0 + X 1.0 + dtype: float64 + + >>> df2 = ps.DataFrame({ + ... "A":[5, 3, 6, 4], + ... "B":[11, 2, 4, 3], + ... "C":[4, 3, 8, 5]}) + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... df1.corrwith(df2).sort_index() + A -0.041703 + B NaN + C 0.395437 + X NaN + dtype: float64 + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... df1.corrwith(df2, method="kendall").sort_index() + A 0.0 + B NaN + C 0.0 + X NaN + dtype: float64 + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... df1.corrwith(df2.B, method="spearman").sort_index() + A -0.4 + C 0.8 + X -0.2 + dtype: float64 + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... df2.corrwith(df1.X).sort_index() + A -0.597614 + B -0.151186 + C -0.642857 + dtype: float64 + """ + from pyspark.pandas.series import Series, first_series + + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError("corrwith currently only works for axis=0") + if method not in ["pearson", "spearman", "kendall"]: + raise ValueError(f"Invalid method {method}") + if not isinstance(other, (DataFrame, Series)): + raise TypeError("unsupported type: {}".format(type(other).__name__)) + + right_is_series = isinstance(other, Series) + + if same_anchor(self, other): + combined = self + this = self + that = other + else: + combined = combine_frames(self, other, how="inner") + this = combined["this"] + that = combined["that"] + + sdf = combined._internal.spark_frame + index_col_name = verify_temp_column_name(sdf, "__corrwith_index_temp_column__") + + this_numeric_column_labels: List[Label] = [] + for column_label in this._internal.column_labels: + if isinstance(this._internal.spark_type_for(column_label), (NumericType, BooleanType)): + this_numeric_column_labels.append(column_label) + + that_numeric_column_labels: List[Label] = [] + for column_label in that._internal.column_labels: + if isinstance(that._internal.spark_type_for(column_label), (NumericType, BooleanType)): + that_numeric_column_labels.append(column_label) + + intersect_numeric_column_labels: List[Label] = [] + diff_numeric_column_labels: List[Label] = [] + pair_scols = [] + if right_is_series: + intersect_numeric_column_labels = this_numeric_column_labels + that_scol = that._internal.spark_column_for(that_numeric_column_labels[0]).cast( + "double" + ) + for numeric_column_label in intersect_numeric_column_labels: + this_scol = this._internal.spark_column_for(numeric_column_label).cast("double") + pair_scols.append( + F.struct( + F.lit(name_like_string(numeric_column_label)).alias(index_col_name), + this_scol.alias(CORRELATION_VALUE_1_COLUMN), + that_scol.alias(CORRELATION_VALUE_2_COLUMN), + ) + ) + else: + for numeric_column_label in this_numeric_column_labels: + if numeric_column_label in that_numeric_column_labels: + intersect_numeric_column_labels.append(numeric_column_label) + else: + diff_numeric_column_labels.append(numeric_column_label) + for numeric_column_label in that_numeric_column_labels: + if numeric_column_label not in this_numeric_column_labels: + diff_numeric_column_labels.append(numeric_column_label) + for numeric_column_label in intersect_numeric_column_labels: + this_scol = this._internal.spark_column_for(numeric_column_label).cast("double") + that_scol = that._internal.spark_column_for(numeric_column_label).cast("double") + pair_scols.append( + F.struct( + F.lit(name_like_string(numeric_column_label)).alias(index_col_name), + this_scol.alias(CORRELATION_VALUE_1_COLUMN), + that_scol.alias(CORRELATION_VALUE_2_COLUMN), + ) + ) + + if len(pair_scols) > 0: + sdf = sdf.select(F.inline(F.array(*pair_scols))) + + sdf = compute(sdf=sdf, groupKeys=[index_col_name], method=method).select( + index_col_name, CORRELATION_CORR_OUTPUT_COLUMN + ) + + else: + sdf = self._internal.spark_frame.select( + F.lit(None).cast("string").alias(index_col_name), + F.lit(None).cast("double").alias(CORRELATION_CORR_OUTPUT_COLUMN), + ).limit(0) + + if not drop and len(diff_numeric_column_labels) > 0: + sdf2 = ( + self._internal.spark_frame.select( + F.lit([name_like_string(label) for label in diff_numeric_column_labels]).alias( + index_col_name + ) + ) + .limit(1) + .select(F.explode(index_col_name).alias(index_col_name)) + ) + sdf = sdf.unionByName(sdf2, allowMissingColumns=True) + + sdf = sdf.withColumn( + NATURAL_ORDER_COLUMN_NAME, + F.monotonically_increasing_id(), + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, index_col_name)], + column_labels=[(CORRELATION_CORR_OUTPUT_COLUMN,)], + column_label_names=self._internal.column_label_names, + ) + sser = first_series(DataFrame(internal)) + sser.name = None + return sser + + def items(self) -> Iterator[Tuple[Name, "Series"]]: + """ + Iterator over (column name, Series) pairs. + + Iterates over the DataFrame columns, returning a tuple with + the column name and the content as a Series. + + Returns + ------- + label : object + The column names for the DataFrame being iterated over. + content : Series + The column entries belonging to each label, as a Series. + + Examples + -------- + >>> df = ps.DataFrame({'species': ['bear', 'bear', 'marsupial'], + ... 'population': [1864, 22000, 80000]}, + ... index=['panda', 'polar', 'koala'], + ... columns=['species', 'population']) + >>> df + species population + panda bear 1864 + polar bear 22000 + koala marsupial 80000 + + >>> for label, content in df.iteritems(): + ... print('label:', label) + ... print('content:', content.to_string()) + ... + ... # doctest: +SKIP + label: species + content: panda bear + polar bear + koala marsupial + label: population + content: panda 1864 + polar 22000 + koala 80000 + """ + return ( + (label if len(label) > 1 else label[0], self._psser_for(label)) + for label in self._internal.column_labels + ) + + def iterrows(self) -> Iterator[Tuple[Name, pd.Series]]: + """ + Iterate over DataFrame rows as (index, Series) pairs. + + Yields + ------ + index : label or tuple of label + The index of the row. A tuple for a `MultiIndex`. + data : pandas.Series + The data of the row as a Series. + + it : generator + A generator that iterates over the rows of the frame. + + Notes + ----- + + 1. Because ``iterrows`` returns a Series for each row, + it does **not** preserve dtypes across the rows (dtypes are + preserved across columns for DataFrames). For example, + + >>> df = ps.DataFrame([[1, 1.5]], columns=['int', 'float']) + >>> row = next(df.iterrows())[1] + >>> row + int 1.0 + float 1.5 + Name: 0, dtype: float64 + >>> print(row['int'].dtype) + float64 + >>> print(df['int'].dtype) + int64 + + To preserve dtypes while iterating over the rows, it is better + to use :meth:`itertuples` which returns namedtuples of the values + and which is generally faster than ``iterrows``. + + 2. You should **never modify** something you are iterating over. + This is not guaranteed to work in all cases. Depending on the + data types, the iterator returns a copy and not a view, and writing + to it will have no effect. + """ + + columns = self.columns + internal_index_columns = self._internal.index_spark_column_names + internal_data_columns = self._internal.data_spark_column_names + + def extract_kv_from_spark_row(row: Row) -> Tuple[Name, Any]: + k = ( + row[internal_index_columns[0]] + if len(internal_index_columns) == 1 + else tuple(row[c] for c in internal_index_columns) + ) + v = [row[c] for c in internal_data_columns] + return k, v + + for k, v in map( + extract_kv_from_spark_row, self._internal.resolved_copy.spark_frame.toLocalIterator() + ): + s = pd.Series(v, index=columns, name=k) + yield k, s + + def itertuples( + self, index: bool = True, name: Optional[str] = "PandasOnSpark" + ) -> Iterator[Tuple]: + """ + Iterate over DataFrame rows as namedtuples. + + Parameters + ---------- + index : bool, default True + If True, return the index as the first element of the tuple. + name : str or None, default "PandasOnSpark" + The name of the returned namedtuples or None to return regular + tuples. + + Returns + ------- + iterator + An object to iterate over namedtuples for each row in the + DataFrame with the first field possibly being the index and + following fields being the column values. + + See Also + -------- + DataFrame.iterrows : Iterate over DataFrame rows as (index, Series) + pairs. + DataFrame.items : Iterate over (column name, Series) pairs. + + Notes + ----- + The column names will be renamed to positional names if they are + invalid Python identifiers, repeated, or start with an underscore. + + Examples + -------- + >>> df = ps.DataFrame({'num_legs': [4, 2], 'num_wings': [0, 2]}, + ... index=['dog', 'hawk']) + >>> df + num_legs num_wings + dog 4 0 + hawk 2 2 + + >>> for row in df.itertuples(): + ... print(row) + ... + PandasOnSpark(Index='dog', num_legs=4, num_wings=0) + PandasOnSpark(Index='hawk', num_legs=2, num_wings=2) + + By setting the `index` parameter to False we can remove the index + as the first element of the tuple: + + >>> for row in df.itertuples(index=False): + ... print(row) + ... + PandasOnSpark(num_legs=4, num_wings=0) + PandasOnSpark(num_legs=2, num_wings=2) + + With the `name` parameter set we set a custom name for the yielded + namedtuples: + + >>> for row in df.itertuples(name='Animal'): + ... print(row) + ... + Animal(Index='dog', num_legs=4, num_wings=0) + Animal(Index='hawk', num_legs=2, num_wings=2) + """ + fields = list(self.columns) + if index: + fields.insert(0, "Index") + + index_spark_column_names = self._internal.index_spark_column_names + data_spark_column_names = self._internal.data_spark_column_names + + def extract_kv_from_spark_row(row: Row) -> Tuple[Name, Any]: + k = ( + row[index_spark_column_names[0]] + if len(index_spark_column_names) == 1 + else tuple(row[c] for c in index_spark_column_names) + ) + v = [row[c] for c in data_spark_column_names] + return k, v + + can_return_named_tuples = sys.version_info >= (3, 7) or len(self.columns) + index < 255 + + if name is not None and can_return_named_tuples: + itertuple = namedtuple(name, fields, rename=True) # type: ignore[misc] + for k, v in map( + extract_kv_from_spark_row, + self._internal.resolved_copy.spark_frame.toLocalIterator(), + ): + yield itertuple._make(([k] if index else []) + list(v)) + else: + for k, v in map( + extract_kv_from_spark_row, + self._internal.resolved_copy.spark_frame.toLocalIterator(), + ): + yield tuple(([k] if index else []) + list(v)) + + def iteritems(self) -> Iterator[Tuple[Name, "Series"]]: + """ + This is an alias of ``items``. + + .. deprecated:: 3.4.0 + iteritems is deprecated and will be removed in a future version. + Use .items instead. + """ + warnings.warn( + "Deprecated in 3.4.0, and will be removed in 4.0.0. Use DataFrame.items instead.", + FutureWarning, + ) + return self.items() + + def to_clipboard(self, excel: bool = True, sep: Optional[str] = None, **kwargs: Any) -> None: + """ + Copy object to the system clipboard. + + Write a text representation of object to the system clipboard. + This can be pasted into Excel, for example. + + .. note:: This method should only be used if the resulting DataFrame is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + excel : bool, default True + - True, use the provided separator, writing in a csv format for + allowing easy pasting into excel. + - False, write a string representation of the object to the + clipboard. + + sep : str, default ``'\\t'`` + Field delimiter. + **kwargs + These parameters will be passed to DataFrame.to_csv. + + Notes + ----- + Requirements for your platform. + + - Linux : `xclip`, or `xsel` (with `gtk` or `PyQt4` modules) + - Windows : none + - OS X : none + + See Also + -------- + read_clipboard : Read text from clipboard. + + Examples + -------- + Copy the contents of a DataFrame to the clipboard. + + >>> df = ps.DataFrame([[1, 2, 3], [4, 5, 6]], columns=['A', 'B', 'C']) # doctest: +SKIP + >>> df.to_clipboard(sep=',') # doctest: +SKIP + ... # Wrote the following to the system clipboard: + ... # ,A,B,C + ... # 0,1,2,3 + ... # 1,4,5,6 + + We can omit the index by passing the keyword `index` and setting + it to false. + + >>> df.to_clipboard(sep=',', index=False) # doctest: +SKIP + ... # Wrote the following to the system clipboard: + ... # A,B,C + ... # 1,2,3 + ... # 4,5,6 + + This function also works for Series: + + >>> df = ps.Series([1, 2, 3, 4, 5, 6, 7], name='x') # doctest: +SKIP + >>> df.to_clipboard(sep=',') # doctest: +SKIP + ... # Wrote the following to the system clipboard: + ... # 0, 1 + ... # 1, 2 + ... # 2, 3 + ... # 3, 4 + ... # 4, 5 + ... # 5, 6 + ... # 6, 7 + """ + + args = locals() + psdf = self + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_clipboard, pd.DataFrame.to_clipboard, args + ) + + def to_html( + self, + buf: Optional[IO[str]] = None, + columns: Optional[Sequence[Name]] = None, + col_space: Optional[Union[str, int, Dict[Name, Union[str, int]]]] = None, + header: bool = True, + index: bool = True, + na_rep: str = "NaN", + formatters: Optional[ + Union[List[Callable[[Any], str]], Dict[Name, Callable[[Any], str]]] + ] = None, + float_format: Optional[Callable[[float], str]] = None, + sparsify: Optional[bool] = None, + index_names: bool = True, + justify: Optional[str] = None, + max_rows: Optional[int] = None, + max_cols: Optional[int] = None, + show_dimensions: bool = False, + decimal: str = ".", + bold_rows: bool = True, + classes: Optional[Union[str, list, tuple]] = None, + escape: bool = True, + notebook: bool = False, + border: Optional[int] = None, + table_id: Optional[str] = None, + render_links: bool = False, + ) -> Optional[str]: + """ + Render a DataFrame as an HTML table. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. If the input + is large, set max_rows parameter. + + Parameters + ---------- + buf : StringIO-like, optional + Buffer to write to. + columns : sequence, optional, default None + The subset of columns to write. Writes all columns by default. + col_space : int, optional + The minimum width of each column. + header : bool, optional + Write out the column names. If a list of strings is given, it + is assumed to be aliases for the column names + index : bool, optional, default True + Whether to print index (row) labels. + na_rep : str, optional, default 'NaN' + String representation of NAN to use. + formatters : list or dict of one-param. functions, optional + Formatter functions to apply to columns' elements by position or + name. + The result of each function must be a Unicode string. + List must be of length equal to the number of columns. + float_format : one-parameter function, optional, default None + Formatter function to apply to columns' elements if they are + floats. The result of this function must be a Unicode string. + sparsify : bool, optional, default True + Set to False for a DataFrame with a hierarchical index to print + every multiindex key at each row. + index_names : bool, optional, default True + Prints the names of the indexes. + justify : str, default None + How to justify the column labels. If None uses the option from + the print configuration (controlled by set_option), 'right' out + of the box. Valid values are + + * left + * right + * center + * justify + * justify-all + * start + * end + * inherit + * match-parent + * initial + * unset. + max_rows : int, optional + Maximum number of rows to display in the console. + max_cols : int, optional + Maximum number of columns to display in the console. + show_dimensions : bool, default False + Display DataFrame dimensions (number of rows by number of columns). + decimal : str, default '.' + Character recognized as decimal separator, e.g. ',' in Europe. + bold_rows : bool, default True + Make the row labels bold in the output. + classes : str or list or tuple, default None + CSS class(es) to apply to the resulting html table. + escape : bool, default True + Convert the characters <, >, and & to HTML-safe sequences. + notebook : {True, False}, default False + Whether the generated HTML is for IPython Notebook. + border : int + A ``border=border`` attribute is included in the opening + `` tag. By default ``pd.options.html.border``. + table_id : str, optional + A css id is included in the opening `
    ` tag if specified. + render_links : bool, default False + Convert URLs to HTML links (only works with pandas 0.24+). + + Returns + ------- + str (or Unicode, depending on data and options) + String representation of the dataframe. + + See Also + -------- + to_string : Convert DataFrame to a string. + """ + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + if max_rows is not None: + psdf = self.head(max_rows) + else: + psdf = self + + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_html, pd.DataFrame.to_html, args + ) + + def to_string( + self, + buf: Optional[IO[str]] = None, + columns: Optional[Sequence[Name]] = None, + col_space: Optional[Union[str, int, Dict[Name, Union[str, int]]]] = None, + header: bool = True, + index: bool = True, + na_rep: str = "NaN", + formatters: Optional[ + Union[List[Callable[[Any], str]], Dict[Name, Callable[[Any], str]]] + ] = None, + float_format: Optional[Callable[[float], str]] = None, + sparsify: Optional[bool] = None, + index_names: bool = True, + justify: Optional[str] = None, + max_rows: Optional[int] = None, + max_cols: Optional[int] = None, + show_dimensions: bool = False, + decimal: str = ".", + line_width: Optional[int] = None, + ) -> Optional[str]: + """ + Render a DataFrame to a console-friendly tabular output. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. If the input + is large, set max_rows parameter. + + Parameters + ---------- + buf : StringIO-like, optional + Buffer to write to. + columns : sequence, optional, default None + The subset of columns to write. Writes all columns by default. + col_space : int, optional + The minimum width of each column. + header : bool, optional + Write out the column names. If a list of strings is given, it + is assumed to be aliases for the column names + index : bool, optional, default True + Whether to print index (row) labels. + na_rep : str, optional, default 'NaN' + String representation of NAN to use. + formatters : list or dict of one-param. functions, optional + Formatter functions to apply to columns' elements by position or + name. + The result of each function must be a Unicode string. + List must be of length equal to the number of columns. + float_format : one-parameter function, optional, default None + Formatter function to apply to columns' elements if they are + floats. The result of this function must be a Unicode string. + sparsify : bool, optional, default True + Set to False for a DataFrame with a hierarchical index to print + every multiindex key at each row. + index_names : bool, optional, default True + Prints the names of the indexes. + justify : str, default None + How to justify the column labels. If None uses the option from + the print configuration (controlled by set_option), 'right' out + of the box. Valid values are + + * left + * right + * center + * justify + * justify-all + * start + * end + * inherit + * match-parent + * initial + * unset. + max_rows : int, optional + Maximum number of rows to display in the console. + max_cols : int, optional + Maximum number of columns to display in the console. + show_dimensions : bool, default False + Display DataFrame dimensions (number of rows by number of columns). + decimal : str, default '.' + Character recognized as decimal separator, e.g. ',' in Europe. + line_width : int, optional + Width to wrap a line in characters. + + Returns + ------- + str (or Unicode, depending on data and options) + String representation of the dataframe. + + See Also + -------- + to_html : Convert DataFrame to HTML. + + Examples + -------- + >>> df = ps.DataFrame({'col1': [1, 2, 3], 'col2': [4, 5, 6]}, columns=['col1', 'col2']) + >>> print(df.to_string()) + col1 col2 + 0 1 4 + 1 2 5 + 2 3 6 + + >>> print(df.to_string(max_rows=2)) + col1 col2 + 0 1 4 + 1 2 5 + """ + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + if max_rows is not None: + psdf = self.head(max_rows) + else: + psdf = self + + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_string, pd.DataFrame.to_string, args + ) + + def to_dict(self, orient: str = "dict", into: Type = dict) -> Union[List, Mapping]: + """ + Convert the DataFrame to a dictionary. + + The type of the key-value pairs can be customized with the parameters + (see below). + + .. note:: This method should only be used if the resulting pandas DataFrame is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + orient : str {'dict', 'list', 'series', 'split', 'records', 'index'} + Determines the type of the values of the dictionary. + + - 'dict' (default) : dict like {column -> {index -> value}} + - 'list' : dict like {column -> [values]} + - 'series' : dict like {column -> Series(values)} + - 'split' : dict like + {'index' -> [index], 'columns' -> [columns], 'data' -> [values]} + - 'records' : list like + [{column -> value}, ... , {column -> value}] + - 'index' : dict like {index -> {column -> value}} + + Abbreviations are allowed. `s` indicates `series` and `sp` + indicates `split`. + + into : class, default dict + The collections.abc.Mapping subclass used for all Mappings + in the return value. Can be the actual class or an empty + instance of the mapping type you want. If you want a + collections.defaultdict, you must pass it initialized. + + Returns + ------- + dict, list or collections.abc.Mapping + Return a collections.abc.Mapping object representing the DataFrame. + The resulting transformation depends on the `orient` parameter. + + Examples + -------- + >>> df = ps.DataFrame({'col1': [1, 2], + ... 'col2': [0.5, 0.75]}, + ... index=['row1', 'row2'], + ... columns=['col1', 'col2']) + >>> df + col1 col2 + row1 1 0.50 + row2 2 0.75 + + >>> df_dict = df.to_dict() + >>> sorted([(key, sorted(values.items())) for key, values in df_dict.items()]) + [('col1', [('row1', 1), ('row2', 2)]), ('col2', [('row1', 0.5), ('row2', 0.75)])] + + You can specify the return orientation. + + >>> df_dict = df.to_dict('series') + >>> sorted(df_dict.items()) + [('col1', row1 1 + row2 2 + Name: col1, dtype: int64), ('col2', row1 0.50 + row2 0.75 + Name: col2, dtype: float64)] + + >>> df_dict = df.to_dict('split') + >>> sorted(df_dict.items()) # doctest: +ELLIPSIS + [('columns', ['col1', 'col2']), ('data', [[1..., 0.75]]), ('index', ['row1', 'row2'])] + + >>> df_dict = df.to_dict('records') + >>> [sorted(values.items()) for values in df_dict] # doctest: +ELLIPSIS + [[('col1', 1...), ('col2', 0.5)], [('col1', 2...), ('col2', 0.75)]] + + >>> df_dict = df.to_dict('index') + >>> sorted([(key, sorted(values.items())) for key, values in df_dict.items()]) + [('row1', [('col1', 1), ('col2', 0.5)]), ('row2', [('col1', 2), ('col2', 0.75)])] + + You can also specify the mapping type. + + >>> from collections import OrderedDict, defaultdict + >>> df.to_dict(into=OrderedDict) + OrderedDict([('col1', OrderedDict([('row1', 1), ('row2', 2)])), \ +('col2', OrderedDict([('row1', 0.5), ('row2', 0.75)]))]) + + If you want a `defaultdict`, you need to initialize it: + + >>> dd = defaultdict(list) + >>> df.to_dict('records', into=dd) # doctest: +ELLIPSIS + [defaultdict(, {'col..., 'col...}), \ +defaultdict(, {'col..., 'col...})] + """ + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + psdf = self + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_dict, pd.DataFrame.to_dict, args + ) + + def to_latex( + self, + buf: Optional[IO[str]] = None, + columns: Optional[List[Name]] = None, + col_space: Optional[int] = None, + header: bool = True, + index: bool = True, + na_rep: str = "NaN", + formatters: Optional[ + Union[List[Callable[[Any], str]], Dict[Name, Callable[[Any], str]]] + ] = None, + float_format: Optional[Callable[[float], str]] = None, + sparsify: Optional[bool] = None, + index_names: bool = True, + bold_rows: bool = False, + column_format: Optional[str] = None, + longtable: Optional[bool] = None, + escape: Optional[bool] = None, + encoding: Optional[str] = None, + decimal: str = ".", + multicolumn: Optional[bool] = None, + multicolumn_format: Optional[str] = None, + multirow: Optional[bool] = None, + ) -> Optional[str]: + r""" + Render an object to a LaTeX tabular environment table. + + Render an object to a tabular environment table. You can splice this into a LaTeX + document. Requires usepackage{booktabs}. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. If the input + is large, consider alternative formats. + + Parameters + ---------- + buf : file descriptor or None + Buffer to write to. If None, the output is returned as a string. + columns : list of label, optional + The subset of columns to write. Writes all columns by default. + col_space : int, optional + The minimum width of each column. + + .. deprecated:: 3.4.0 + + header : bool or list of str, default True + Write out the column names. If a list of strings is given, it is assumed to be aliases + for the column names. + index : bool, default True + Write row names (index). + na_rep : str, default ‘NaN’ + Missing data representation. + formatters : list of functions or dict of {str: function}, optional + Formatter functions to apply to columns’ elements by position or name. The result of + each function must be a Unicode string. List must be of length equal to the number of + columns. + float_format : str, optional + Format string for floating point numbers. + sparsify : bool, optional + Set to False for a DataFrame with a hierarchical index to print every multiindex key at + each row. By default the value will be read from the config module. + index_names : bool, default True + Prints the names of the indexes. + bold_rows : bool, default False + Make the row labels bold in the output. + column_format : str, optional + The columns format as specified in LaTeX table format e.g. ‘rcl’ for 3 columns. By + default, ‘l’ will be used for all columns except columns of numbers, which default + to ‘r’. + longtable : bool, optional + By default the value will be read from the pandas config module. Use a longtable + environment instead of tabular. Requires adding a usepackage{longtable} to your LaTeX + preamble. + escape : bool, optional + 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. + decimal : str, default ‘.’ + Character recognized as decimal separator, e.g. ‘,’ in Europe. + multicolumn : bool, default True + Use multicolumn to enhance MultiIndex columns. The default will be read from the config + module. + multicolumn_format : str, default ‘l’ + The alignment for multicolumns, similar to column_format The default will be read from + the config module. + multirow : bool, default False + Use multirow to enhance MultiIndex rows. Requires adding a usepackage{multirow} to your + LaTeX preamble. Will print centered labels (instead of top-aligned) across the contained + rows, separating groups via clines. The default will be read from the pandas config + module. + + Returns + ------- + str or None + If buf is None, returns the resulting LateX format as a string. Otherwise returns None. + + See Also + -------- + DataFrame.to_string : Render a DataFrame to a console-friendly + tabular output. + DataFrame.to_html : Render a DataFrame as an HTML table. + + + Examples + -------- + >>> df = ps.DataFrame({'name': ['Raphael', 'Donatello'], + ... 'mask': ['red', 'purple'], + ... 'weapon': ['sai', 'bo staff']}, + ... columns=['name', 'mask', 'weapon']) + >>> print(df.to_latex(index=False)) # doctest: +NORMALIZE_WHITESPACE + \begin{tabular}{lll} + \toprule + name & mask & weapon \\ + \midrule + Raphael & red & sai \\ + Donatello & purple & bo staff \\ + \bottomrule + \end{tabular} + """ + warnings.warn( + "Argument `col_space` will be removed in 4.0.0.", + FutureWarning, + ) + + args = locals() + psdf = self + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_latex, pd.DataFrame.to_latex, args + ) + + # TODO: enable doctests once we drop Spark 2.3.x (due to type coercion logic + # when creating arrays) + def transpose(self) -> "DataFrame": + """ + Transpose index and columns. + + Reflect the DataFrame over its main diagonal by writing rows as columns + and vice-versa. The property :attr:`.T` is an accessor to the method + :meth:`transpose`. + + .. note:: This method is based on an expensive operation due to the nature + of big data. Internally it needs to generate each row for each value, and + then group twice - it is a huge operation. To prevent misuse, this method + has the 'compute.max_rows' default limit of input length and raises a ValueError. + + >>> from pyspark.pandas.config import option_context + >>> with option_context('compute.max_rows', 1000): # doctest: +NORMALIZE_WHITESPACE + ... ps.DataFrame({'a': range(1001)}).transpose() + Traceback (most recent call last): + ... + ValueError: Current DataFrame's length exceeds the given limit of 1000 rows. + Please set 'compute.max_rows' by using 'pyspark.pandas.config.set_option' + to retrieve more than 1000 rows. Note that, before changing the + 'compute.max_rows', this operation is considerably expensive. + + Returns + ------- + DataFrame + The transposed DataFrame. + + Notes + ----- + Transposing a DataFrame with mixed dtypes will result in a homogeneous + DataFrame with the coerced dtype. For instance, if int and float have + to be placed in same column, it becomes float. If type coercion is not + possible, it fails. + + Also, note that the values in index should be unique because they become + unique column names. + + In addition, if Spark 2.3 is used, the types should always be exactly same. + + Examples + -------- + **Square DataFrame with homogeneous dtype** + + >>> d1 = {'col1': [1, 2], 'col2': [3, 4]} + >>> df1 = ps.DataFrame(data=d1, columns=['col1', 'col2']) + >>> df1 + col1 col2 + 0 1 3 + 1 2 4 + + >>> df1_transposed = df1.T.sort_index() # doctest: +SKIP + >>> df1_transposed # doctest: +SKIP + 0 1 + col1 1 2 + col2 3 4 + + When the dtype is homogeneous in the original DataFrame, we get a + transposed DataFrame with the same dtype: + + >>> df1.dtypes + col1 int64 + col2 int64 + dtype: object + >>> df1_transposed.dtypes # doctest: +SKIP + 0 int64 + 1 int64 + dtype: object + + **Non-square DataFrame with mixed dtypes** + + >>> d2 = {'score': [9.5, 8], + ... 'kids': [0, 0], + ... 'age': [12, 22]} + >>> df2 = ps.DataFrame(data=d2, columns=['score', 'kids', 'age']) + >>> df2 + score kids age + 0 9.5 0 12 + 1 8.0 0 22 + + >>> df2_transposed = df2.T.sort_index() # doctest: +SKIP + >>> df2_transposed # doctest: +SKIP + 0 1 + age 12.0 22.0 + kids 0.0 0.0 + score 9.5 8.0 + + When the DataFrame has mixed dtypes, we get a transposed DataFrame with + the coerced dtype: + + >>> df2.dtypes + score float64 + kids int64 + age int64 + dtype: object + + >>> df2_transposed.dtypes # doctest: +SKIP + 0 float64 + 1 float64 + dtype: object + """ + max_compute_count = get_option("compute.max_rows") + if max_compute_count is not None: + pdf = self.head(max_compute_count + 1)._to_internal_pandas() + if len(pdf) > max_compute_count: + raise ValueError( + "Current DataFrame's length exceeds the given limit of {0} rows. " + "Please set 'compute.max_rows' by using 'pyspark.pandas.config.set_option' " + "to retrieve more than {0} rows. Note that, before changing the " + "'compute.max_rows', this operation is considerably expensive.".format( + max_compute_count + ) + ) + return DataFrame(pdf.transpose()) + + # Explode the data to be pairs. + # + # For instance, if the current input DataFrame is as below: + # + # +------+------+------+------+------+ + # |index1|index2|(a,x1)|(a,x2)|(b,x3)| + # +------+------+------+------+------+ + # | y1| z1| 1| 0| 0| + # | y2| z2| 0| 50| 0| + # | y3| z3| 3| 2| 1| + # +------+------+------+------+------+ + # + # Output of `exploded_df` becomes as below: + # + # +-----------------+-----------------+-----------------+-----+ + # | index|__index_level_0__|__index_level_1__|value| + # +-----------------+-----------------+-----------------+-----+ + # |{"a":["y1","z1"]}| a| x1| 1| + # |{"a":["y1","z1"]}| a| x2| 0| + # |{"a":["y1","z1"]}| b| x3| 0| + # |{"a":["y2","z2"]}| a| x1| 0| + # |{"a":["y2","z2"]}| a| x2| 50| + # |{"a":["y2","z2"]}| b| x3| 0| + # |{"a":["y3","z3"]}| a| x1| 3| + # |{"a":["y3","z3"]}| a| x2| 2| + # |{"a":["y3","z3"]}| b| x3| 1| + # +-----------------+-----------------+-----------------+-----+ + pairs = F.explode( + F.array( + *[ + F.struct( + *[ + F.lit(col).alias(SPARK_INDEX_NAME_FORMAT(i)) + for i, col in enumerate(label) + ], + *[self._internal.spark_column_for(label).alias("value")], + ) + for label in self._internal.column_labels + ] + ) + ) + + exploded_df = self._internal.spark_frame.withColumn("pairs", pairs).select( + [ + F.to_json( + F.struct( + F.array(*[scol for scol in self._internal.index_spark_columns]).alias("a") + ) + ).alias("index"), + F.col("pairs.*"), + ] + ) + + # After that, executes pivot with key and its index column. + # Note that index column should contain unique values since column names + # should be unique. + internal_index_columns = [ + SPARK_INDEX_NAME_FORMAT(i) for i in range(self._internal.column_labels_level) + ] + pivoted_df = exploded_df.groupBy(internal_index_columns).pivot("index") + + transposed_df = pivoted_df.agg(F.first(F.col("value"))) + + new_data_columns = list( + filter(lambda x: x not in internal_index_columns, transposed_df.columns) + ) + + column_labels = [ + None if len(label) == 1 and label[0] is None else label + for label in (tuple(json.loads(col)["a"]) for col in new_data_columns) + ] + + internal = InternalFrame( + spark_frame=transposed_df, + index_spark_columns=[scol_for(transposed_df, col) for col in internal_index_columns], + index_names=self._internal.column_label_names, + column_labels=column_labels, + data_spark_columns=[scol_for(transposed_df, col) for col in new_data_columns], + column_label_names=self._internal.index_names, + ) + + return DataFrame(internal) + + T = property(transpose) + + def apply( + self, func: Callable, axis: Axis = 0, args: Sequence[Any] = (), **kwds: Any + ) -> Union["Series", "DataFrame", "Index"]: + """ + Apply a function along an axis of the DataFrame. + + Objects passed to the function are Series objects whose index is + either the DataFrame's index (``axis=0``) or the DataFrame's columns + (``axis=1``). + + See also `Transform and apply a function + `_. + + .. note:: when `axis` is 0 or 'index', the `func` is unable to access + to the whole input series. pandas-on-Spark internally splits the input series into + multiple batches and calls `func` with each batch multiple times. Therefore, operations + such as global aggregations are impossible. See the example below. + + >>> # This case does not return the length of whole series but of the batch internally + ... # used. + ... def length(s) -> int: + ... return len(s) + ... + >>> df = ps.DataFrame({'A': range(1000)}) + >>> df.apply(length, axis=0) # doctest: +SKIP + 0 83 + 1 83 + 2 83 + ... + 10 83 + 11 83 + dtype: int32 + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify the return type as `Series` or scalar value in ``func``, + for instance, as below: + + >>> def square(s) -> ps.Series[np.int32]: + ... return s ** 2 + + pandas-on-Spark uses return type hints and does not try to infer the type. + + In case when axis is 1, it requires to specify `DataFrame` or scalar value + with type hints as below: + + >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]: + ... return x + 1 + + If the return type is specified as `DataFrame`, the output column names become + `c0, c1, c2 ... cn`. These names are positionally mapped to the returned + DataFrame in ``func``. + + To specify the column names, you can assign them in a pandas style as below: + + >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]: + ... return x + 1 + + >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]}) + >>> def plus_one(x) -> ps.DataFrame[ + ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]: + ... return x + 1 + + Parameters + ---------- + func : function + Function to apply to each column or row. + axis : {0 or 'index', 1 or 'columns'}, default 0 + Axis along which the function is applied: + + * 0 or 'index': apply function to each column. + * 1 or 'columns': apply function to each row. + args : tuple + Positional arguments to pass to `func` in addition to the + array/series. + **kwds + Additional keyword arguments to pass as keywords arguments to + `func`. + + Returns + ------- + Series or DataFrame + Result of applying ``func`` along the given axis of the + DataFrame. + + See Also + -------- + DataFrame.applymap : For elementwise operations. + DataFrame.aggregate : Only perform aggregating type operations. + DataFrame.transform : Only perform transforming type operations. + Series.apply : The equivalent function for Series. + + Examples + -------- + >>> df = ps.DataFrame([[4, 9]] * 3, columns=['A', 'B']) + >>> df + A B + 0 4 9 + 1 4 9 + 2 4 9 + + Using a numpy universal function (in this case the same as + ``np.sqrt(df)``): + + >>> def sqrt(x) -> ps.Series[float]: + ... return np.sqrt(x) + ... + >>> df.apply(sqrt, axis=0) + A B + 0 2.0 3.0 + 1 2.0 3.0 + 2 2.0 3.0 + + You can omit type hints and let pandas-on-Spark infer its type. + + >>> df.apply(np.sqrt, axis=0) + A B + 0 2.0 3.0 + 1 2.0 3.0 + 2 2.0 3.0 + + When `axis` is 1 or 'columns', it applies the function for each row. + + >>> def summation(x) -> np.int64: + ... return np.sum(x) + ... + >>> df.apply(summation, axis=1) + 0 13 + 1 13 + 2 13 + dtype: int64 + + You can omit type hints and let pandas-on-Spark infer its type. + + >>> df.apply(np.sum, axis=1) + 0 13 + 1 13 + 2 13 + dtype: int64 + + >>> df.apply(max, axis=1) + 0 9 + 1 9 + 2 9 + dtype: int64 + + Returning a list-like will result in a Series + + >>> df.apply(lambda x: [1, 2], axis=1) + 0 [1, 2] + 1 [1, 2] + 2 [1, 2] + dtype: object + + To specify the types when `axis` is '1', it should use DataFrame[...] + annotation. In this case, the column names are automatically generated. + + >>> def identify(x) -> ps.DataFrame[('index', int), [('A', np.int64), ('B', np.int64)]]: + ... return x + ... + >>> df.apply(identify, axis=1) # doctest: +NORMALIZE_WHITESPACE + A B + index + 0 4 9 + 1 4 9 + 2 4 9 + + You can also specify extra arguments. + + >>> def plus_two(a, b, c) -> ps.DataFrame[np.int64, [np.int64, np.int64]]: + ... return a + b + c + ... + >>> df.apply(plus_two, axis=1, args=(1,), c=3) + c0 c1 + 0 8 13 + 1 8 13 + 2 8 13 + """ + from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.series import first_series + + if not isinstance(func, types.FunctionType): + assert callable(func), "the first argument should be a callable function." + f = func + # Note that the return type hints specified here affects actual return + # type in Spark (e.g., infer_return_type). And MyPy does not allow + # redefinition of a function. + func = lambda *args, **kwargs: f(*args, **kwargs) # noqa: E731 + + axis = validate_axis(axis) + should_return_series = False + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + should_retain_index = should_infer_schema + + def apply_func(pdf: pd.DataFrame) -> pd.DataFrame: + pdf_or_pser = pdf.apply(func, axis=axis, args=args, **kwds) # type: ignore[arg-type] + if isinstance(pdf_or_pser, pd.Series): + return pdf_or_pser.to_frame() + else: + return pdf_or_pser + + self_applied: DataFrame = DataFrame(self._internal.resolved_copy) + + column_labels: Optional[List[Label]] = None + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + log_advice( + "If the type hints is not specified for `apply`, " + "it is expensive to infer the data type internally." + ) + limit = get_option("compute.shortcut_limit") + pdf = self_applied.head(limit + 1)._to_internal_pandas() + applied = pdf.apply(func, axis=axis, args=args, **kwds) # type: ignore[arg-type] + psser_or_psdf = ps.from_pandas(applied) + if len(pdf) <= limit: + return psser_or_psdf + + psdf = psser_or_psdf + if isinstance(psser_or_psdf, ps.Series): + should_return_series = True + psdf = psser_or_psdf._psdf + + index_fields = [field.normalize_spark_type() for field in psdf._internal.index_fields] + data_fields = [field.normalize_spark_type() for field in psdf._internal.data_fields] + + return_schema = StructType([field.struct_field for field in index_fields + data_fields]) + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, apply_func, return_schema, retain_index=should_retain_index + ) + sdf = self_applied._internal.to_internal_spark_frame.mapInPandas( + lambda iterator: map(output_func, iterator), schema=return_schema + ) + + # If schema is inferred, we can restore indexes too. + internal = psdf._internal.with_new_sdf( + spark_frame=sdf, index_fields=index_fields, data_fields=data_fields + ) + else: + return_type = infer_return_type(func) + require_index_axis = isinstance(return_type, SeriesType) + require_column_axis = isinstance(return_type, DataFrameType) + index_fields = None + + if require_index_axis: + if axis != 0: + raise TypeError( + "The given function should specify a scalar or a series as its type " + "hints when axis is 0 or 'index'; however, the return type " + "was %s" % return_sig + ) + dtype = cast(SeriesType, return_type).dtype + spark_type = cast(SeriesType, return_type).spark_type + data_fields = [ + InternalField( + dtype=dtype, struct_field=StructField(name=name, dataType=spark_type) + ) + for name in self_applied.columns + ] + return_schema = StructType([field.struct_field for field in data_fields]) + elif require_column_axis: + if axis != 1: + raise TypeError( + "The given function should specify a scalar or a frame as its type " + "hints when axis is 1 or 'column'; however, the return type " + "was %s" % return_sig + ) + index_fields = cast(DataFrameType, return_type).index_fields + should_retain_index = len(index_fields) > 0 + data_fields = cast(DataFrameType, return_type).data_fields + return_schema = cast(DataFrameType, return_type).spark_type + else: + # any axis is fine. + should_return_series = True + spark_type = cast(ScalarType, return_type).spark_type + dtype = cast(ScalarType, return_type).dtype + data_fields = [ + InternalField( + dtype=dtype, + struct_field=StructField( + name=SPARK_DEFAULT_SERIES_NAME, dataType=spark_type + ), + ) + ] + return_schema = StructType([field.struct_field for field in data_fields]) + column_labels = [None] + + output_func = GroupBy._make_pandas_df_builder_func( + self_applied, apply_func, return_schema, retain_index=should_retain_index + ) + sdf = self_applied._internal.to_internal_spark_frame.mapInPandas( + lambda iterator: map(output_func, iterator), schema=return_schema + ) + + index_spark_columns = None + index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None + + if should_retain_index: + index_spark_columns = [ + scol_for(sdf, index_field.struct_field.name) for index_field in index_fields + ] + + if not any( + [ + SPARK_INDEX_NAME_PATTERN.match(index_field.struct_field.name) + for index_field in index_fields + ] + ): + index_names = [(index_field.struct_field.name,) for index_field in index_fields] + internal = InternalFrame( + spark_frame=sdf, + index_names=index_names, + index_spark_columns=index_spark_columns, + index_fields=index_fields, + data_fields=data_fields, + column_labels=column_labels, + ) + + result: DataFrame = DataFrame(internal) + if should_return_series: + return first_series(result) + else: + return result + + def transform( + self, func: Callable[..., "Series"], axis: Axis = 0, *args: Any, **kwargs: Any + ) -> "DataFrame": + """ + Call ``func`` on self producing a Series with transformed values + and that has the same length as its input. + + See also `Transform and apply a function + `_. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def square(x) -> ps.Series[np.int32]: + ... return x ** 2 + + pandas-on-Spark uses return type hints and does not try to infer the type. + + .. note:: the series within ``func`` is actually multiple pandas series as the + segments of the whole pandas-on-Spark series; therefore, the length of each series + is not guaranteed. As an example, an aggregation against each series + does work as a global aggregation but an aggregation of each segment. See + below: + + >>> def func(x) -> ps.Series[np.int32]: + ... return x + sum(x) + + Parameters + ---------- + func : function + Function to use for transforming the data. It must work when pandas Series + is passed. + axis : int, default 0 or 'index' + Can only be set to 0 now. + *args + Positional arguments to pass to func. + **kwargs + Keyword arguments to pass to func. + + Returns + ------- + DataFrame + A DataFrame that must have the same length as self. + + Raises + ------ + Exception : If the returned DataFrame has a different length than self. + + See Also + -------- + DataFrame.aggregate : Only perform aggregating type operations. + DataFrame.apply : Invoke function on DataFrame. + Series.transform : The equivalent function for Series. + + Examples + -------- + >>> df = ps.DataFrame({'A': range(3), 'B': range(1, 4)}, columns=['A', 'B']) + >>> df + A B + 0 0 1 + 1 1 2 + 2 2 3 + + >>> def square(x) -> ps.Series[np.int32]: + ... return x ** 2 + >>> df.transform(square) + A B + 0 0 1 + 1 1 4 + 2 4 9 + + You can omit type hints and let pandas-on-Spark infer its type. + + >>> df.transform(lambda x: x ** 2) + A B + 0 0 1 + 1 1 4 + 2 4 9 + + For multi-index columns: + + >>> df.columns = [('X', 'A'), ('X', 'B')] + >>> df.transform(square) # doctest: +NORMALIZE_WHITESPACE + X + A B + 0 0 1 + 1 1 4 + 2 4 9 + + >>> (df * -1).transform(abs) # doctest: +NORMALIZE_WHITESPACE + X + A B + 0 0 1 + 1 1 2 + 2 2 3 + + You can also specify extra arguments. + + >>> def calculation(x, y, z) -> ps.Series[int]: + ... return x ** y + z + >>> df.transform(calculation, y=10, z=20) # doctest: +NORMALIZE_WHITESPACE + X + A B + 0 20 21 + 1 21 1044 + 2 1044 59069 + """ + if not isinstance(func, types.FunctionType): + assert callable(func), "the first argument should be a callable function." + f = func + # Note that the return type hints specified here affects actual return + # type in Spark (e.g., infer_return_type). And, MyPy does not allow + # redefinition of a function. + func = lambda *args, **kwargs: f(*args, **kwargs) # noqa: E731 + + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + log_advice( + "If the type hints is not specified for `transform`, " + "it is expensive to infer the data type internally." + ) + limit = get_option("compute.shortcut_limit") + pdf = self.head(limit + 1)._to_internal_pandas() + transformed = pdf.transform(func, axis, *args, **kwargs) # type: ignore[arg-type] + psdf: DataFrame = DataFrame(transformed) + if len(pdf) <= limit: + return psdf + + applied = [] + data_fields = [] + for input_label, output_label in zip( + self._internal.column_labels, psdf._internal.column_labels + ): + psser = self._psser_for(input_label) + + field = psdf._internal.field_for(output_label).normalize_spark_type() + data_fields.append(field) + + return_schema = field.spark_type + applied.append( + psser.pandas_on_spark._transform_batch( + func=lambda c: func(c, *args, **kwargs), + return_type=SeriesType(field.dtype, return_schema), + ) + ) + + internal = self._internal.with_new_columns(applied, data_fields=data_fields) + return DataFrame(internal) + else: + return self._apply_series_op( + lambda psser: psser.pandas_on_spark.transform_batch(func, *args, **kwargs) + ) + + def pop(self, item: Name) -> "DataFrame": + """ + Return item and drop from frame. Raise KeyError if not found. + + Parameters + ---------- + item : str + Label of column to be popped. + + Returns + ------- + Series + + Examples + -------- + >>> df = ps.DataFrame([('falcon', 'bird', 389.0), + ... ('parrot', 'bird', 24.0), + ... ('lion', 'mammal', 80.5), + ... ('monkey','mammal', np.nan)], + ... columns=('name', 'class', 'max_speed')) + + >>> df + name class max_speed + 0 falcon bird 389.0 + 1 parrot bird 24.0 + 2 lion mammal 80.5 + 3 monkey mammal NaN + + >>> df.pop('class') + 0 bird + 1 bird + 2 mammal + 3 mammal + Name: class, dtype: object + + >>> df + name max_speed + 0 falcon 389.0 + 1 parrot 24.0 + 2 lion 80.5 + 3 monkey NaN + + Also support for MultiIndex + + >>> df = ps.DataFrame([('falcon', 'bird', 389.0), + ... ('parrot', 'bird', 24.0), + ... ('lion', 'mammal', 80.5), + ... ('monkey','mammal', np.nan)], + ... columns=('name', 'class', 'max_speed')) + >>> columns = [('a', 'name'), ('a', 'class'), ('b', 'max_speed')] + >>> df.columns = pd.MultiIndex.from_tuples(columns) + >>> df + a b + name class max_speed + 0 falcon bird 389.0 + 1 parrot bird 24.0 + 2 lion mammal 80.5 + 3 monkey mammal NaN + + >>> df.pop('a') + name class + 0 falcon bird + 1 parrot bird + 2 lion mammal + 3 monkey mammal + + >>> df + b + max_speed + 0 389.0 + 1 24.0 + 2 80.5 + 3 NaN + """ + result = self[item] + self._update_internal_frame(self.drop(columns=item)._internal) + return result + + # TODO: add axis parameter can work when '1' or 'columns' + def xs(self, key: Name, axis: Axis = 0, level: Optional[int] = None) -> DataFrameOrSeries: + """ + Return cross-section from the DataFrame. + + This method takes a `key` argument to select data at a particular + level of a MultiIndex. + + Parameters + ---------- + key : label or tuple of label + Label contained in the index, or partially in a MultiIndex. + axis : 0 or 'index', 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. + + Returns + ------- + DataFrame or Series + Cross-section from the original DataFrame + corresponding to the selected index levels. + + See Also + -------- + DataFrame.loc : Access a group of rows and columns + by label(s) or a boolean array. + DataFrame.iloc : Purely integer-location based indexing + for selection by position. + + Examples + -------- + >>> d = {'num_legs': [4, 4, 2, 2], + ... 'num_wings': [0, 0, 2, 2], + ... 'class': ['mammal', 'mammal', 'mammal', 'bird'], + ... 'animal': ['cat', 'dog', 'bat', 'penguin'], + ... 'locomotion': ['walks', 'walks', 'flies', 'walks']} + >>> df = ps.DataFrame(data=d) + >>> df = df.set_index(['class', 'animal', 'locomotion']) + >>> df # doctest: +NORMALIZE_WHITESPACE + num_legs num_wings + class animal locomotion + mammal cat walks 4 0 + dog walks 4 0 + bat flies 2 2 + bird penguin walks 2 2 + + Get values at specified index + + >>> df.xs('mammal') # doctest: +NORMALIZE_WHITESPACE + num_legs num_wings + animal locomotion + cat walks 4 0 + dog walks 4 0 + bat flies 2 2 + + Get values at several indexes + + >>> df.xs(('mammal', 'dog')) # doctest: +NORMALIZE_WHITESPACE + num_legs num_wings + locomotion + walks 4 0 + + >>> df.xs(('mammal', 'dog', 'walks')) # doctest: +NORMALIZE_WHITESPACE + num_legs 4 + num_wings 0 + Name: (mammal, dog, walks), dtype: int64 + + Get values at specified index and level + + >>> df.xs('cat', level=1) # doctest: +NORMALIZE_WHITESPACE + num_legs num_wings + class locomotion + mammal walks 4 0 + """ + from pyspark.pandas.series import first_series + + if not is_name_like_value(key): + raise TypeError("'key' should be a scalar value or tuple that contains scalar values") + + if level is not None and is_name_like_tuple(key): + raise KeyError(key) + + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + if not is_name_like_tuple(key): + key = (key,) + if len(key) > self._internal.index_level: + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + len(key), self._internal.index_level + ) + ) + if level is None: + level = 0 + + rows = [ + self._internal.index_spark_columns[lvl] == index for lvl, index in enumerate(key, level) + ] + internal = self._internal.with_filter(reduce(lambda x, y: x & y, rows)) + + if len(key) == self._internal.index_level: + psdf: DataFrame = DataFrame(internal) + pdf = psdf.head(2)._to_internal_pandas() + if len(pdf) == 0: + raise KeyError(key) + elif len(pdf) > 1: + return psdf + else: + return first_series(DataFrame(pdf.transpose())) + else: + index_spark_columns = ( + internal.index_spark_columns[:level] + + internal.index_spark_columns[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, + index_names=index_names, + index_fields=index_fields, + ).resolved_copy + return DataFrame(internal) + + # TODO(SPARK-42620): Add `inclusive` parameter and replace `include_start` & `include_end`. + # See https://github.com/pandas-dev/pandas/issues/43248 + def between_time( + self, + start_time: Union[datetime.time, str], + end_time: Union[datetime.time, str], + include_start: bool = True, + include_end: bool = True, + axis: Axis = 0, + ) -> "DataFrame": + """ + Select values between particular times of the day (example: 9:00-9:30 AM). + + By setting ``start_time`` to be later than ``end_time``, + you can get the times that are *not* between the two times. + + Parameters + ---------- + start_time : datetime.time or str + Initial time as a time filter limit. + end_time : datetime.time or str + End time as a time filter limit. + include_start : bool, default True + Whether the start time needs to be included in the result. + + .. deprecated:: 3.4.0 + + include_end : bool, default True + Whether the end time needs to be included in the result. + + .. deprecated:: 3.4.0 + + axis : {0 or 'index', 1 or 'columns'}, default 0 + Determine range time on index or columns value. + + Returns + ------- + DataFrame + Data from the original object filtered to the specified dates range. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + See Also + -------- + at_time : Select values at a particular time of the day. + first : Select initial periods of time series based on a date offset. + last : Select final periods of time series based on a date offset. + DatetimeIndex.indexer_between_time : Get just the index locations for + values between particular times of the day. + + Examples + -------- + >>> idx = pd.date_range('2018-04-09', periods=4, freq='1D20min') + >>> psdf = ps.DataFrame({'A': [1, 2, 3, 4]}, index=idx) + >>> psdf + A + 2018-04-09 00:00:00 1 + 2018-04-10 00:20:00 2 + 2018-04-11 00:40:00 3 + 2018-04-12 01:00:00 4 + + >>> psdf.between_time('0:15', '0:45') # doctest: +SKIP + A + 2018-04-10 00:20:00 2 + 2018-04-11 00:40:00 3 + + You get the times that are *not* between two times by setting + ``start_time`` later than ``end_time``: + + >>> psdf.between_time('0:45', '0:15') # doctest: +SKIP + A + 2018-04-09 00:00:00 1 + 2018-04-12 01:00:00 4 + """ + axis = validate_axis(axis) + + if axis != 0: + raise NotImplementedError("between_time currently only works for axis=0") + + if not isinstance(self.index, ps.DatetimeIndex): + raise TypeError("Index must be DatetimeIndex") + + psdf = self.copy() + psdf.index.name = verify_temp_column_name(psdf, "__index_name__") + return_types = [psdf.index.dtype] + list(psdf.dtypes) + + def pandas_between_time( # type: ignore[no-untyped-def] + pdf, + ) -> ps.DataFrame[return_types]: # type: ignore[valid-type] + return pdf.between_time(start_time, end_time, include_start, include_end).reset_index() + + # apply_batch will remove the index of the pandas-on-Spark DataFrame and attach a + # default index, which will never be used. Use "distributed" index as a dummy to + # avoid overhead. + with option_context("compute.default_index_type", "distributed"): + psdf = psdf.pandas_on_spark.apply_batch(pandas_between_time) + + return DataFrame( + self._internal.copy( + spark_frame=psdf._internal.spark_frame, + index_spark_columns=psdf._internal.data_spark_columns[:1], + index_fields=psdf._internal.data_fields[:1], + data_spark_columns=psdf._internal.data_spark_columns[1:], + data_fields=psdf._internal.data_fields[1:], + ) + ) + + # TODO: implement axis=1 + def at_time( + self, time: Union[datetime.time, str], asof: bool = False, axis: Axis = 0 + ) -> "DataFrame": + """ + Select values at particular time of day (example: 9:30AM). + + Parameters + ---------- + time : datetime.time or str + axis : {0 or 'index', 1 or 'columns'}, default 0 + + Returns + ------- + DataFrame + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + See Also + -------- + between_time : Select values between particular times of the day. + DatetimeIndex.indexer_at_time : Get just the index locations for + values at particular time of the day. + + Examples + -------- + >>> idx = pd.date_range('2018-04-09', periods=4, freq='12H') + >>> psdf = ps.DataFrame({'A': [1, 2, 3, 4]}, index=idx) + >>> psdf + A + 2018-04-09 00:00:00 1 + 2018-04-09 12:00:00 2 + 2018-04-10 00:00:00 3 + 2018-04-10 12:00:00 4 + + >>> psdf.at_time('12:00') + A + 2018-04-09 12:00:00 2 + 2018-04-10 12:00:00 4 + """ + if asof: + raise NotImplementedError("'asof' argument is not supported") + + axis = validate_axis(axis) + + if axis != 0: + raise NotImplementedError("at_time currently only works for axis=0") + + if not isinstance(self.index, ps.DatetimeIndex): + raise TypeError("Index must be DatetimeIndex") + + psdf = self.copy() + psdf.index.name = verify_temp_column_name(psdf, "__index_name__") + return_types = [psdf.index.dtype] + list(psdf.dtypes) + + def pandas_at_time( # type: ignore[no-untyped-def] + pdf, + ) -> ps.DataFrame[return_types]: # type: ignore[valid-type] + return pdf.at_time(time, asof, axis).reset_index() + + # apply_batch will remove the index of the pandas-on-Spark DataFrame and attach + # a default index, which will never be used. Use "distributed" index as a dummy + # to avoid overhead. + with option_context("compute.default_index_type", "distributed"): + psdf = psdf.pandas_on_spark.apply_batch(pandas_at_time) + + return DataFrame( + self._internal.copy( + spark_frame=psdf._internal.spark_frame, + index_spark_columns=psdf._internal.data_spark_columns[:1], + index_fields=psdf._internal.data_fields[:1], + data_spark_columns=psdf._internal.data_spark_columns[1:], + data_fields=psdf._internal.data_fields[1:], + ) + ) + + def where( + self, + cond: DataFrameOrSeries, + other: Union[DataFrameOrSeries, Any] = np.nan, + axis: Axis = None, + ) -> "DataFrame": + """ + Replace values where the condition is False. + + Parameters + ---------- + cond : boolean DataFrame + Where cond is True, keep the original value. Where False, + replace with corresponding value from other. + other : scalar, DataFrame + Entries where cond is False are replaced with corresponding value from other. + axis : int, default None + Can only be set to 0 now for compatibility with pandas. + + Returns + ------- + DataFrame + + Examples + -------- + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> df1 = ps.DataFrame({'A': [0, 1, 2, 3, 4], 'B':[100, 200, 300, 400, 500]}) + >>> df2 = ps.DataFrame({'A': [0, -1, -2, -3, -4], 'B':[-100, -200, -300, -400, -500]}) + >>> df1 + A B + 0 0 100 + 1 1 200 + 2 2 300 + 3 3 400 + 4 4 500 + >>> df2 + A B + 0 0 -100 + 1 -1 -200 + 2 -2 -300 + 3 -3 -400 + 4 -4 -500 + + >>> df1.where(df1 > 0).sort_index() + A B + 0 NaN 100.0 + 1 1.0 200.0 + 2 2.0 300.0 + 3 3.0 400.0 + 4 4.0 500.0 + + >>> df1.where(df1 > 1, 10).sort_index() + A B + 0 10 100 + 1 10 200 + 2 2 300 + 3 3 400 + 4 4 500 + + >>> df1.where(df1 > 1, df1 + 100).sort_index() + A B + 0 100 100 + 1 101 200 + 2 2 300 + 3 3 400 + 4 4 500 + + >>> df1.where(df1 > 1, df2).sort_index() + A B + 0 0 100 + 1 -1 200 + 2 2 300 + 3 3 400 + 4 4 500 + + When the column name of cond is different from self, it treats all values are False + + >>> cond = ps.DataFrame({'C': [0, -1, -2, -3, -4], 'D':[4, 3, 2, 1, 0]}) % 3 == 0 + >>> cond + C D + 0 True False + 1 False True + 2 False False + 3 True False + 4 False True + + >>> df1.where(cond).sort_index() + A B + 0 NaN NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 NaN NaN + + When the type of cond is Series, it just check boolean regardless of column name + + >>> cond = ps.Series([1, 2]) > 1 + >>> cond + 0 False + 1 True + dtype: bool + + >>> df1.where(cond).sort_index() + A B + 0 NaN NaN + 1 1.0 200.0 + 2 NaN NaN + 3 NaN NaN + 4 NaN NaN + + >>> reset_option("compute.ops_on_diff_frames") + """ + from pyspark.pandas.series import Series + + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + tmp_cond_col_name = "__tmp_cond_col_{}__".format + tmp_other_col_name = "__tmp_other_col_{}__".format + + psdf = self.copy() + + tmp_cond_col_names = [ + tmp_cond_col_name(name_like_string(label)) for label in self._internal.column_labels + ] + if isinstance(cond, DataFrame): + cond = cond[ + [ + ( + cond._internal.spark_column_for(label) + if label in cond._internal.column_labels + else F.lit(False) + ).alias(name) + for label, name in zip(self._internal.column_labels, tmp_cond_col_names) + ] + ] + psdf[tmp_cond_col_names] = cond + elif isinstance(cond, Series): + cond = cond.to_frame() + cond = cond[ + [cond._internal.data_spark_columns[0].alias(name) for name in tmp_cond_col_names] + ] + psdf[tmp_cond_col_names] = cond + else: + raise TypeError("type of cond must be a DataFrame or Series") + + tmp_other_col_names = [ + tmp_other_col_name(name_like_string(label)) for label in self._internal.column_labels + ] + if isinstance(other, DataFrame): + other = other[ + [ + ( + other._internal.spark_column_for(label) + if label in other._internal.column_labels + else F.lit(np.nan) + ).alias(name) + for label, name in zip(self._internal.column_labels, tmp_other_col_names) + ] + ] + psdf[tmp_other_col_names] = other + elif isinstance(other, Series): + other = other.to_frame() + other = other[ + [other._internal.data_spark_columns[0].alias(name) for name in tmp_other_col_names] + ] + psdf[tmp_other_col_names] = other + else: + for label in self._internal.column_labels: + psdf[tmp_other_col_name(name_like_string(label))] = other + + # above logic make spark dataframe looks like below: + # +-----------------+---+---+------------------+-------------------+------------------+--... + # |__index_level_0__| A| B|__tmp_cond_col_A__|__tmp_other_col_A__|__tmp_cond_col_B__|__... + # +-----------------+---+---+------------------+-------------------+------------------+--... + # | 0| 0|100| true| 0| false| ... + # | 1| 1|200| false| -1| false| ... + # | 3| 3|400| true| -3| false| ... + # | 2| 2|300| false| -2| true| ... + # | 4| 4|500| false| -4| false| ... + # +-----------------+---+---+------------------+-------------------+------------------+--... + + data_spark_columns = [] + for label in self._internal.column_labels: + data_spark_columns.append( + F.when( + psdf[tmp_cond_col_name(name_like_string(label))].spark.column, + psdf._internal.spark_column_for(label), + ) + .otherwise(psdf[tmp_other_col_name(name_like_string(label))].spark.column) + .alias(psdf._internal.spark_column_name_for(label)) + ) + + return DataFrame( + psdf._internal.with_new_columns( + data_spark_columns, column_labels=self._internal.column_labels # TODO: dtypes? + ) + ) + + def mask( + self, cond: DataFrameOrSeries, other: Union[DataFrameOrSeries, Any] = np.nan + ) -> "DataFrame": + """ + Replace values where the condition is True. + + Parameters + ---------- + cond : boolean DataFrame + Where cond is False, keep the original value. Where True, + replace with corresponding value from other. + other : scalar, DataFrame + Entries where cond is True are replaced with corresponding value from other. + + Returns + ------- + DataFrame + + Examples + -------- + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> df1 = ps.DataFrame({'A': [0, 1, 2, 3, 4], 'B':[100, 200, 300, 400, 500]}) + >>> df2 = ps.DataFrame({'A': [0, -1, -2, -3, -4], 'B':[-100, -200, -300, -400, -500]}) + >>> df1 + A B + 0 0 100 + 1 1 200 + 2 2 300 + 3 3 400 + 4 4 500 + >>> df2 + A B + 0 0 -100 + 1 -1 -200 + 2 -2 -300 + 3 -3 -400 + 4 -4 -500 + + >>> df1.mask(df1 > 0).sort_index() + A B + 0 0.0 NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 NaN NaN + + >>> df1.mask(df1 > 1, 10).sort_index() + A B + 0 0 10 + 1 1 10 + 2 10 10 + 3 10 10 + 4 10 10 + + >>> df1.mask(df1 > 1, df1 + 100).sort_index() + A B + 0 0 200 + 1 1 300 + 2 102 400 + 3 103 500 + 4 104 600 + + >>> df1.mask(df1 > 1, df2).sort_index() + A B + 0 0 -100 + 1 1 -200 + 2 -2 -300 + 3 -3 -400 + 4 -4 -500 + + >>> reset_option("compute.ops_on_diff_frames") + """ + from pyspark.pandas.series import Series + + if not isinstance(cond, (DataFrame, Series)): + raise TypeError("type of cond must be a DataFrame or Series") + + cond_inversed = cond._apply_series_op(lambda psser: ~psser) + return self.where(cond_inversed, other) + + @property + def index(self) -> "Index": + """The index (row labels) Column of the DataFrame. + + Currently not supported when the DataFrame has no index. + + See Also + -------- + Index + """ + from pyspark.pandas.indexes.base import Index + + return Index._new_instance(self) + + @property + def empty(self) -> bool: + """ + Returns true if the current DataFrame is empty. Otherwise, returns false. + + Examples + -------- + >>> ps.range(10).empty + False + + >>> ps.range(0).empty + True + + >>> ps.DataFrame({}, index=list('abc')).empty + True + """ + return ( + len(self._internal.column_labels) == 0 + or self._internal.resolved_copy.spark_frame.isEmpty() + ) + + @property + def style(self) -> "Styler": + """ + Property returning a Styler object containing methods for + building a styled HTML representation for the DataFrame. + + Examples + -------- + >>> ps.range(1001).style # doctest: +SKIP + + """ + max_results = get_option("compute.max_rows") + if max_results is not None: + pdf = self.head(max_results + 1)._to_internal_pandas() + if len(pdf) > max_results: + warnings.warn( + "'style' property will only use top %s rows." % max_results, UserWarning + ) + return pdf.head(max_results).style + else: + return self._to_internal_pandas().style + + def set_index( + self, + keys: Union[Name, List[Name]], + drop: bool = True, + append: bool = False, + inplace: bool = False, + ) -> Optional["DataFrame"]: + """Set the DataFrame index (row labels) using one or more existing columns. + + Set the DataFrame index (row labels) using one or more existing + columns or arrays (of the correct length). The index can replace the + existing index or expand on it. + + Parameters + ---------- + keys : label or array-like or list of labels/arrays + This parameter can be either a single column key, a single array of + the same length as the calling DataFrame, or a list containing an + arbitrary combination of column keys and arrays. Here, "array" + encompasses :class:`Series`, :class:`Index` and ``np.ndarray``. + drop : bool, default True + Delete columns to be used as the new index. + append : bool, default False + Whether to append columns to existing index. + inplace : bool, default False + Modify the DataFrame in place (do not create a new object). + + Returns + ------- + DataFrame + Changed row labels. + + See Also + -------- + DataFrame.reset_index : Opposite of set_index. + + Examples + -------- + >>> df = ps.DataFrame({'month': [1, 4, 7, 10], + ... 'year': [2012, 2014, 2013, 2014], + ... 'sale': [55, 40, 84, 31]}, + ... columns=['month', 'year', 'sale']) + >>> df + month year sale + 0 1 2012 55 + 1 4 2014 40 + 2 7 2013 84 + 3 10 2014 31 + + Set the index to become the 'month' column: + + >>> df.set_index('month') # doctest: +NORMALIZE_WHITESPACE + year sale + month + 1 2012 55 + 4 2014 40 + 7 2013 84 + 10 2014 31 + + Create a MultiIndex using columns 'year' and 'month': + + >>> df.set_index(['year', 'month']) # doctest: +NORMALIZE_WHITESPACE + sale + year month + 2012 1 55 + 2014 4 40 + 2013 7 84 + 2014 10 31 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + key_list: List[Label] + if is_name_like_tuple(keys): + key_list = [cast(Label, keys)] + elif is_name_like_value(keys): + key_list = [(keys,)] + else: + key_list = [key if is_name_like_tuple(key) else (key,) for key in keys] + columns = set(self._internal.column_labels) + for key in key_list: + if key not in columns: + raise KeyError(name_like_string(key)) + + if drop: + column_labels = [ + label for label in self._internal.column_labels if label not in key_list + ] + else: + column_labels = self._internal.column_labels + if append: + index_spark_columns = self._internal.index_spark_columns + [ + self._internal.spark_column_for(label) for label in key_list + ] + index_names = self._internal.index_names + key_list + index_fields = self._internal.index_fields + [ + self._internal.field_for(label) for label in key_list + ] + else: + index_spark_columns = [self._internal.spark_column_for(label) for label in key_list] + index_names = key_list + index_fields = [self._internal.field_for(label) for label in key_list] + + internal = self._internal.copy( + index_spark_columns=index_spark_columns, + index_names=index_names, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=[self._internal.spark_column_for(label) for label in column_labels], + data_fields=[self._internal.field_for(label) for label in column_labels], + ) + + if inplace: + self._update_internal_frame(internal) + return None + else: + return DataFrame(internal) + + def reset_index( + self, + level: Optional[Union[int, Name, Sequence[Union[int, Name]]]] = None, + drop: bool = False, + inplace: bool = False, + col_level: int = 0, + col_fill: str = "", + ) -> Optional["DataFrame"]: + """Reset the index, or a level of it. + + For DataFrame with multi-level index, return new DataFrame with labeling information in + the columns under the index names, defaulting to 'level_0', 'level_1', etc. if any are None. + For a standard index, the index name will be used (if set), otherwise a default 'index' or + 'level_0' (if 'index' is already taken) will be used. + + Parameters + ---------- + level : int, str, tuple, or list, default None + Only remove the given levels from the index. Removes all levels by + default. + drop : bool, default False + Do not try to insert index into dataframe columns. This reset + the index to the default integer index. + inplace : bool, default False + Modify the DataFrame in place (do not create a new object). + col_level : int or str, default 0 + If the columns have multiple levels, determines which level the + labels are inserted into. By default it is inserted into the first + level. + col_fill : object, default '' + If the columns have multiple levels, determines how the other + levels are named. If None then the index name is repeated. + + Returns + ------- + DataFrame + DataFrame with the new index. + + See Also + -------- + DataFrame.set_index : Opposite of reset_index. + + Examples + -------- + >>> df = ps.DataFrame([('bird', 389.0), + ... ('bird', 24.0), + ... ('mammal', 80.5), + ... ('mammal', np.nan)], + ... index=['falcon', 'parrot', 'lion', 'monkey'], + ... columns=('class', 'max_speed')) + >>> df + class max_speed + falcon bird 389.0 + parrot bird 24.0 + lion mammal 80.5 + monkey mammal NaN + + When we reset the index, the old index is added as a column. Unlike pandas, pandas-on-Spark + does not automatically add a sequential index. The following 0, 1, 2, 3 are only + there when we display the DataFrame. + + >>> df.reset_index() + index class max_speed + 0 falcon bird 389.0 + 1 parrot bird 24.0 + 2 lion mammal 80.5 + 3 monkey mammal NaN + + We can use the `drop` parameter to avoid the old index being added as + a column: + + >>> df.reset_index(drop=True) + class max_speed + 0 bird 389.0 + 1 bird 24.0 + 2 mammal 80.5 + 3 mammal NaN + + You can also use `reset_index` with `MultiIndex`. + + >>> index = pd.MultiIndex.from_tuples([('bird', 'falcon'), + ... ('bird', 'parrot'), + ... ('mammal', 'lion'), + ... ('mammal', 'monkey')], + ... names=['class', 'name']) + >>> columns = pd.MultiIndex.from_tuples([('speed', 'max'), + ... ('species', 'type')]) + >>> df = ps.DataFrame([(389.0, 'fly'), + ... ( 24.0, 'fly'), + ... ( 80.5, 'run'), + ... (np.nan, 'jump')], + ... index=index, + ... columns=columns) + >>> df # doctest: +NORMALIZE_WHITESPACE + speed species + max type + class name + bird falcon 389.0 fly + parrot 24.0 fly + mammal lion 80.5 run + monkey NaN jump + + If the index has multiple levels, we can reset a subset of them: + + >>> df.reset_index(level='class') # doctest: +NORMALIZE_WHITESPACE + class speed species + max type + name + falcon bird 389.0 fly + parrot bird 24.0 fly + lion mammal 80.5 run + monkey mammal NaN jump + + If we are not dropping the index, by default, it is placed in the top + level. We can place it in another level: + + >>> df.reset_index(level='class', col_level=1) # doctest: +NORMALIZE_WHITESPACE + speed species + class max type + name + falcon bird 389.0 fly + parrot bird 24.0 fly + lion mammal 80.5 run + monkey mammal NaN jump + + When the index is inserted under another level, we can specify under + which one with the parameter `col_fill`: + + >>> df.reset_index(level='class', col_level=1, + ... col_fill='species') # doctest: +NORMALIZE_WHITESPACE + species speed species + class max type + name + falcon bird 389.0 fly + parrot bird 24.0 fly + lion mammal 80.5 run + monkey mammal NaN jump + + If we specify a nonexistent level for `col_fill`, it is created: + + >>> df.reset_index(level='class', col_level=1, + ... col_fill='genus') # doctest: +NORMALIZE_WHITESPACE + genus speed species + class max type + name + falcon bird 389.0 fly + parrot bird 24.0 fly + lion mammal 80.5 run + monkey mammal NaN jump + """ + inplace = validate_bool_kwarg(inplace, "inplace") + multi_index = self._internal.index_level > 1 + + def rename(index: int) -> Label: + if multi_index: + return ("level_{}".format(index),) + else: + if ("index",) not in self._internal.column_labels: + return ("index",) + else: + return ("level_{}".format(index),) + + if level is None: + new_column_labels = [ + name if name is not None else rename(i) + for i, name in enumerate(self._internal.index_names) + ] + new_data_spark_columns = [ + scol.alias(name_like_string(label)) + for scol, label in zip(self._internal.index_spark_columns, new_column_labels) + ] + new_data_fields = self._internal.index_fields + + index_spark_columns = [] + index_names = [] + index_fields = [] + else: + if is_list_like(level): + level = list(cast(Sequence[Union[int, Name]], level)) + if isinstance(level, int) or is_name_like_tuple(level): + level_list = [cast(Union[int, Label], level)] + elif is_name_like_value(level): + level_list = [(level,)] + else: + level_list = [ + lvl if isinstance(lvl, int) or is_name_like_tuple(lvl) else (lvl,) + for lvl in level + ] + + if all(isinstance(lvl, int) for lvl in level_list): + int_level_list = cast(List[int], level_list) + for lev in int_level_list: + if lev >= self._internal.index_level: + raise IndexError( + "Too many levels: Index has only {} level, not {}".format( + self._internal.index_level, lev + 1 + ) + ) + idx = int_level_list + elif all(is_name_like_tuple(lev) for lev in level_list): + idx = [] + for label in cast(List[Label], level_list): + try: + i = self._internal.index_names.index(label) + idx.append(i) + except ValueError: + if multi_index: + raise KeyError("Level unknown not found") + else: + raise KeyError( + "Level unknown must be same as name ({})".format( + name_like_string(self._internal.index_names[0]) + ) + ) + else: + raise ValueError("Level should be all int or all string.") + idx.sort() + + new_column_labels = [] + new_data_spark_columns = [] + new_data_fields = [] + + index_spark_columns = self._internal.index_spark_columns.copy() + index_names = self._internal.index_names.copy() + index_fields = self._internal.index_fields.copy() + + for i in idx[::-1]: + name = index_names.pop(i) + new_column_labels.insert(0, name if name is not None else rename(i)) + + scol = index_spark_columns.pop(i) + new_data_spark_columns.insert(0, scol.alias(name_like_string(name))) + + new_data_fields.insert(0, index_fields.pop(i).copy(name=name_like_string(name))) + + if drop: + new_data_spark_columns = [] + new_column_labels = [] + new_data_fields = [] + + for label in new_column_labels: + if label in self._internal.column_labels: + raise ValueError("cannot insert {}, already exists".format(name_like_string(label))) + + if self._internal.column_labels_level > 1: + column_depth = len(self._internal.column_labels[0]) + if col_level >= column_depth: + raise IndexError( + "Too many levels: Index has only {} levels, not {}".format( + column_depth, col_level + 1 + ) + ) + if any(col_level + len(label) > column_depth for label in new_column_labels): + raise ValueError("Item must have length equal to number of levels.") + new_column_labels = [ + tuple( + ([col_fill] * col_level) + + list(label) + + ([col_fill] * (column_depth - (len(label) + col_level))) + ) + for label in new_column_labels + ] + + internal = self._internal.copy( + index_spark_columns=index_spark_columns, + index_names=index_names, + index_fields=index_fields, + column_labels=new_column_labels + self._internal.column_labels, + data_spark_columns=new_data_spark_columns + self._internal.data_spark_columns, + data_fields=new_data_fields + self._internal.data_fields, + ) + + if inplace: + self._update_internal_frame(internal) + return None + else: + return DataFrame(internal) + + def isnull(self) -> "DataFrame": + """ + Detects missing values for items in the current Dataframe. + + Return a boolean same-sized Dataframe indicating if the values are NA. + NA values, such as None or numpy.NaN, gets mapped to True values. + Everything else gets mapped to False values. + + See Also + -------- + DataFrame.notnull + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, None), (.6, None), (.2, .1)]) + >>> df.isnull() + 0 1 + 0 False False + 1 False True + 2 False True + 3 False False + + >>> df = ps.DataFrame([[None, 'bee', None], ['dog', None, 'fly']]) + >>> df.isnull() + 0 1 2 + 0 True False True + 1 False True False + """ + return self._apply_series_op(lambda psser: psser.isnull()) + + isna = isnull + + def notnull(self) -> "DataFrame": + """ + Detects non-missing values for items in the current Dataframe. + + This function takes a dataframe and indicates whether it's + values are valid (not missing, which is ``NaN`` in numeric + datatypes, ``None`` or ``NaN`` in objects and ``NaT`` in datetimelike). + + See Also + -------- + DataFrame.isnull + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, None), (.6, None), (.2, .1)]) + >>> df.notnull() + 0 1 + 0 True True + 1 True False + 2 True False + 3 True True + + >>> df = ps.DataFrame([['ant', 'bee', 'cat'], ['dog', None, 'fly']]) + >>> df.notnull() + 0 1 2 + 0 True True True + 1 True False True + """ + return self._apply_series_op(lambda psser: psser.notnull()) + + notna = notnull + + def insert( + self, + loc: int, + column: Name, + value: Union[Scalar, "Series", Iterable], + allow_duplicates: bool = False, + ) -> None: + """ + Insert column into DataFrame at specified location. + + Raises a ValueError if `column` is already contained in the DataFrame, + unless `allow_duplicates` is set to True. + + Parameters + ---------- + loc : int + Insertion index. Must verify 0 <= loc <= len(columns). + column : str, number, or hashable object + Label of the inserted column. + value : int, Series, or array-like + allow_duplicates : bool, optional + + Examples + -------- + >>> psdf = ps.DataFrame([1, 2, 3]) + >>> psdf.sort_index() + 0 + 0 1 + 1 2 + 2 3 + >>> psdf.insert(0, 'x', 4) + >>> psdf.sort_index() + x 0 + 0 4 1 + 1 4 2 + 2 4 3 + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + + >>> psdf.insert(1, 'y', [5, 6, 7]) + >>> psdf.sort_index() + x y 0 + 0 4 5 1 + 1 4 6 2 + 2 4 7 3 + + >>> psdf.insert(2, 'z', ps.Series([8, 9, 10])) + >>> psdf.sort_index() + x y z 0 + 0 4 5 8 1 + 1 4 6 9 2 + 2 4 7 10 3 + + >>> reset_option("compute.ops_on_diff_frames") + """ + if not isinstance(loc, int): + raise TypeError("loc must be int") + + assert 0 <= loc <= len(self.columns) + assert allow_duplicates is False + + if not is_name_like_value(column): + raise TypeError( + '"column" should be a scalar value or tuple that contains scalar values' + ) + + # TODO(SPARK-37723): Support tuple for non-MultiIndex column name. + if is_name_like_tuple(column): + if self._internal.column_labels_level > 1: + if len(column) != len(self.columns.levels): # type: ignore[attr-defined] + # To be consistent with pandas + raise ValueError('"column" must have length equal to number of column levels.') + else: + raise NotImplementedError( + "Assigning column name as tuple is only supported for MultiIndex columns " + "for now." + ) + + if column in self.columns: + raise ValueError("cannot insert %s, already exists" % str(column)) + + psdf = self.copy() + psdf[column] = value + columns = psdf.columns[:-1].insert(loc, psdf.columns[-1]) + psdf = psdf[columns] + self._update_internal_frame(psdf._internal) + + # TODO: add frep and axis parameter + def shift(self, periods: int = 1, fill_value: Optional[Any] = None) -> "DataFrame": + """ + Shift DataFrame by desired number of periods. + + .. note:: the current implementation of shift uses Spark's Window without + specifying partition specification. This leads to moving all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int + Number of periods to shift. Can be positive or negative. + fill_value : object, optional + The scalar value to use for newly introduced missing values. + The default depends on the dtype of self. For numeric data, np.nan is used. + + Returns + ------- + Copy of input DataFrame, shifted. + + Examples + -------- + >>> df = ps.DataFrame({'Col1': [10, 20, 15, 30, 45], + ... 'Col2': [13, 23, 18, 33, 48], + ... 'Col3': [17, 27, 22, 37, 52]}, + ... columns=['Col1', 'Col2', 'Col3']) + + >>> df.shift(periods=3) + Col1 Col2 Col3 + 0 NaN NaN NaN + 1 NaN NaN NaN + 2 NaN NaN NaN + 3 10.0 13.0 17.0 + 4 20.0 23.0 27.0 + + >>> df.shift(periods=3, fill_value=0) + Col1 Col2 Col3 + 0 0 0 0 + 1 0 0 0 + 2 0 0 0 + 3 10 13 17 + 4 20 23 27 + + """ + return self._apply_series_op( + lambda psser: psser._shift(periods, fill_value), should_resolve=True + ) + + # TODO: axis should support 1 or 'columns' either at this moment + def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame": + """ + First discrete difference of element. + + Calculates the difference of a DataFrame element compared with another element in the + DataFrame (default is the element in the same column of the previous row). + + .. note:: the current implementation of diff uses Spark's Window without + specifying partition specification. This leads to moving all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int, default 1 + Periods to shift for calculating difference, accepts negative values. + axis : int, default 0 or 'index' + Can only be set to 0 now. + + Returns + ------- + diffed : DataFrame + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2, 3, 4, 5, 6], + ... 'b': [1, 1, 2, 3, 5, 8], + ... 'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c']) + >>> df + a b c + 0 1 1 1 + 1 2 1 4 + 2 3 2 9 + 3 4 3 16 + 4 5 5 25 + 5 6 8 36 + + >>> df.diff() + a b c + 0 NaN NaN NaN + 1 1.0 0.0 3.0 + 2 1.0 1.0 5.0 + 3 1.0 1.0 7.0 + 4 1.0 2.0 9.0 + 5 1.0 3.0 11.0 + + Difference with previous column + + >>> df.diff(periods=3) + a b c + 0 NaN NaN NaN + 1 NaN NaN NaN + 2 NaN NaN NaN + 3 3.0 2.0 15.0 + 4 3.0 4.0 21.0 + 5 3.0 6.0 27.0 + + Difference with following row + + >>> df.diff(periods=-1) + a b c + 0 -1.0 0.0 -3.0 + 1 -1.0 -1.0 -5.0 + 2 -1.0 -1.0 -7.0 + 3 -1.0 -2.0 -9.0 + 4 -1.0 -3.0 -11.0 + 5 NaN NaN NaN + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + return self._apply_series_op(lambda psser: psser._diff(periods), should_resolve=True) + + # TODO: axis should support 1 or 'columns' either at this moment + def nunique( + self, + axis: Axis = 0, + dropna: bool = True, + approx: bool = False, + rsd: float = 0.05, + ) -> "Series": + """ + Return number of unique elements in the object. + + Excludes NA values by default. + + Parameters + ---------- + axis : int, default 0 or 'index' + Can only be set to 0 now. + dropna : bool, default True + Don’t include NaN in the count. + approx: bool, default False + If False, will use the exact algorithm and return the exact number of unique. + If True, it uses the HyperLogLog approximate algorithm, which is significantly faster + for large amounts of data. + Note: This parameter is specific to pandas-on-Spark and is not found in pandas. + rsd: float, default 0.05 + Maximum estimation error allowed in the HyperLogLog algorithm. + Note: Just like ``approx`` this parameter is specific to pandas-on-Spark. + + Returns + ------- + The number of unique values per column as a pandas-on-Spark Series. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 3], 'B': [np.nan, 3, np.nan]}) + >>> df.nunique() + A 3 + B 1 + dtype: int64 + + >>> df.nunique(dropna=False) + A 3 + B 2 + dtype: int64 + + On big data, we recommend using the approximate algorithm to speed up this function. + The result will be very close to the exact unique count. + + >>> df.nunique(approx=True) + A 3 + B 1 + dtype: int64 + """ + from pyspark.pandas.series import first_series + + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + sdf = self._internal.spark_frame.select( + [F.lit(None).cast(StringType()).alias(SPARK_DEFAULT_INDEX_NAME)] + + [ + self._psser_for(label)._nunique(dropna, approx, rsd) + for label in self._internal.column_labels + ] + ) + + # The data is expected to be small so it's fine to transpose/use the default index. + with ps.option_context("compute.max_rows", 1): + internal = self._internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + index_names=[None], + index_fields=[None], + data_spark_columns=[ + scol_for(sdf, col) for col in self._internal.data_spark_column_names + ], + data_fields=None, + ) + return first_series(DataFrame(internal).transpose()) + + def round(self, decimals: Union[int, Dict[Name, int], "Series"] = 0) -> "DataFrame": + """ + Round a DataFrame to a variable number of decimal places. + + Parameters + ---------- + decimals : int, dict, Series + Number of decimal places to round each column to. If an int is + given, round each column to the same number of places. + Otherwise dict and Series round to variable numbers of places. + Column names should be in the keys if `decimals` is a + dict-like, or in the index if `decimals` is a Series. Any + columns not included in `decimals` will be left as is. Elements + of `decimals` which are not columns of the input will be + ignored. + + .. note:: If `decimals` is a Series, it is expected to be small, + as all the data is loaded into the driver's memory. + + Returns + ------- + DataFrame + + See Also + -------- + Series.round + + Examples + -------- + >>> df = ps.DataFrame({'A':[0.028208, 0.038683, 0.877076], + ... 'B':[0.992815, 0.645646, 0.149370], + ... 'C':[0.173891, 0.577595, 0.491027]}, + ... columns=['A', 'B', 'C'], + ... index=['first', 'second', 'third']) + >>> df + A B C + first 0.028208 0.992815 0.173891 + second 0.038683 0.645646 0.577595 + third 0.877076 0.149370 0.491027 + + >>> df.round(2) + A B C + first 0.03 0.99 0.17 + second 0.04 0.65 0.58 + third 0.88 0.15 0.49 + + >>> df.round({'A': 1, 'C': 2}) + A B C + first 0.0 0.992815 0.17 + second 0.0 0.645646 0.58 + third 0.9 0.149370 0.49 + + >>> decimals = ps.Series([1, 0, 2], index=['A', 'B', 'C']) + >>> df.round(decimals) + A B C + first 0.0 1.0 0.17 + second 0.0 1.0 0.58 + third 0.9 0.0 0.49 + """ + if isinstance(decimals, ps.Series): + decimals_dict = { + k if isinstance(k, tuple) else (k,): v + for k, v in decimals._to_internal_pandas().items() + } + elif isinstance(decimals, dict): + decimals_dict = {k if is_name_like_tuple(k) else (k,): v for k, v in decimals.items()} + elif isinstance(decimals, int): + decimals_dict = {k: decimals for k in self._internal.column_labels} + else: + raise TypeError("decimals must be an integer, a dict-like or a Series") + + def op(psser: ps.Series) -> Union[ps.Series, PySparkColumn]: + label = psser._column_label + if label in decimals_dict: + return F.round(psser.spark.column, decimals_dict[label]) + else: + return psser + + return self._apply_series_op(op) + + def _mark_duplicates( + self, + subset: Optional[Union[Name, List[Name]]] = None, + keep: Union[bool, str] = "first", + ) -> Tuple[PySparkDataFrame, str]: + if subset is None: + subset_list = self._internal.column_labels + else: + if is_name_like_tuple(subset): + subset_list = [cast(Label, subset)] + elif is_name_like_value(subset): + subset_list = [(subset,)] + else: + subset_list = [sub if is_name_like_tuple(sub) else (sub,) for sub in subset] + diff = set(subset_list).difference(set(self._internal.column_labels)) + if len(diff) > 0: + raise KeyError(", ".join([name_like_string(d) for d in diff])) + group_cols = [self._internal.spark_column_name_for(label) for label in subset_list] + + sdf = self._internal.resolved_copy.spark_frame + + column = verify_temp_column_name(sdf, "__duplicated__") + + if keep == "first" or keep == "last": + if keep == "first": + ord_func = F.asc + else: + ord_func = F.desc + window = ( + Window.partitionBy(*group_cols) + .orderBy(ord_func(NATURAL_ORDER_COLUMN_NAME)) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + sdf = sdf.withColumn(column, F.row_number().over(window) > 1) + elif not keep: + window = Window.partitionBy(*group_cols).rowsBetween( + Window.unboundedPreceding, Window.unboundedFollowing + ) + sdf = sdf.withColumn(column, F.count("*").over(window) > 1) + else: + raise ValueError("'keep' only supports 'first', 'last' and False") + return sdf, column + + def duplicated( + self, + subset: Optional[Union[Name, List[Name]]] = None, + keep: Union[bool, str] = "first", + ) -> "Series": + """ + Return boolean Series denoting duplicate rows, optionally only considering certain columns. + + Parameters + ---------- + subset : column label or sequence of labels, optional + Only consider certain columns for identifying duplicates, + default use all of the columns + keep : {'first', 'last', False}, default 'first' + - ``first`` : Mark duplicates as ``True`` except for the first occurrence. + - ``last`` : Mark duplicates as ``True`` except for the last occurrence. + - False : Mark all duplicates as ``True``. + + Returns + ------- + duplicated : Series + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 3], 'b': [1, 1, 1, 4], 'c': [1, 1, 1, 5]}, + ... columns = ['a', 'b', 'c']) + >>> df + a b c + 0 1 1 1 + 1 1 1 1 + 2 1 1 1 + 3 3 4 5 + + >>> df.duplicated().sort_index() + 0 False + 1 True + 2 True + 3 False + dtype: bool + + Mark duplicates as ``True`` except for the last occurrence. + + >>> df.duplicated(keep='last').sort_index() + 0 True + 1 True + 2 False + 3 False + dtype: bool + + Mark all duplicates as ``True``. + + >>> df.duplicated(keep=False).sort_index() + 0 True + 1 True + 2 True + 3 False + dtype: bool + """ + from pyspark.pandas.series import first_series + + sdf, column = self._mark_duplicates(subset, keep) + + sdf = sdf.select( + self._internal.index_spark_columns + + [scol_for(sdf, column).alias(SPARK_DEFAULT_SERIES_NAME)] + ) + return first_series( + DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + column_labels=[None], + data_spark_columns=[scol_for(sdf, SPARK_DEFAULT_SERIES_NAME)], + ) + ) + ) + + # TODO: support other as DataFrame or array-like + def dot(self, other: "Series") -> "Series": + """ + Compute the matrix multiplication between the DataFrame and others. + + This method computes the matrix product between the DataFrame and the + values of an other Series + + It can also be called using ``self @ other`` in Python >= 3.5. + + .. note:: This method is based on an expensive operation due to the nature + of big data. Internally it needs to generate each row for each value, and + then group twice - it is a huge operation. To prevent misuse, this method + has the 'compute.max_rows' default limit of input length and raises a ValueError. + + >>> from pyspark.pandas.config import option_context + >>> with option_context( + ... 'compute.max_rows', 1000, "compute.ops_on_diff_frames", True + ... ): # doctest: +NORMALIZE_WHITESPACE + ... psdf = ps.DataFrame({'a': range(1001)}) + ... psser = ps.Series([2], index=['a']) + ... psdf.dot(psser) + Traceback (most recent call last): + ... + ValueError: Current DataFrame's length exceeds the given limit of 1000 rows. + Please set 'compute.max_rows' by using 'pyspark.pandas.config.set_option' + to retrieve more than 1000 rows. Note that, before changing the + 'compute.max_rows', this operation is considerably expensive. + + Parameters + ---------- + other : Series + The other object to compute the matrix product with. + + Returns + ------- + Series + Return the matrix product between self and other as a Series. + + See Also + -------- + Series.dot: Similar method for Series. + + Notes + ----- + The dimensions of DataFrame and other must be compatible to + compute the matrix multiplication. In addition, the column names of + DataFrame and the index of other must contain the same values, as they + will be aligned prior to the multiplication. + + The dot method for Series computes the inner product, instead of the + matrix product here. + + Examples + -------- + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> psdf = ps.DataFrame([[0, 1, -2, -1], [1, 1, 1, 1]]) + >>> psser = ps.Series([1, 1, 2, 1]) + >>> psdf.dot(psser) + 0 -4 + 1 5 + dtype: int64 + + Note how shuffling of the objects does not change the result. + + >>> psser2 = psser.reindex([1, 0, 2, 3]) + >>> psdf.dot(psser2) + 0 -4 + 1 5 + dtype: int64 + >>> psdf @ psser2 + 0 -4 + 1 5 + dtype: int64 + >>> reset_option("compute.ops_on_diff_frames") + """ + if not isinstance(other, ps.Series): + raise TypeError("Unsupported type {}".format(type(other).__name__)) + else: + return cast(ps.Series, other.dot(self.transpose())).rename(None) + + def __matmul__(self, other: "Series") -> "Series": + """ + Matrix multiplication using binary `@` operator in Python>=3.5. + """ + return self.dot(other) + + def to_table( + self, + name: str, + format: Optional[str] = None, + mode: str = "w", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, + ) -> None: + if index_col is None: + log_advice( + "If `index_col` is not specified for `to_table`, " + "the existing index is lost when converting to table." + ) + mode = validate_mode(mode) + return self.spark.to_table(name, format, mode, partition_cols, index_col, **options) + + to_table.__doc__ = SparkFrameMethods.to_table.__doc__ + + def to_delta( + self, + path: str, + mode: str = "w", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """ + Write the DataFrame out as a Delta Lake table. + + Parameters + ---------- + path : str, required + Path to write to. + mode : str + Python write mode, default 'w'. + + .. note:: mode can accept the strings for Spark writing mode. + Such as 'append', 'overwrite', 'ignore', 'error', 'errorifexists'. + + - 'append' (equivalent to 'a'): Append the new data to existing data. + - 'overwrite' (equivalent to 'w'): Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols : str or list of str, optional, default None + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default the index is always lost. + options : dict + All other options passed directly into Delta Lake. + + See Also + -------- + read_delta + DataFrame.to_parquet + DataFrame.to_table + DataFrame.to_spark_io + + Examples + -------- + + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df + date country code + 0 2012-01-31 12:00:00 KR 1 + 1 2012-02-29 12:00:00 US 2 + 2 2012-03-31 12:00:00 JP 3 + + Create a new Delta Lake table, partitioned by one column: + + >>> df.to_delta('%s/to_delta/foo' % path, partition_cols='date') # doctest: +SKIP + + Partitioned by two columns: + + >>> df.to_delta('%s/to_delta/bar' % path, + ... partition_cols=['date', 'country']) # doctest: +SKIP + + Overwrite an existing table's partitions, using the 'replaceWhere' capability in Delta: + + >>> df.to_delta('%s/to_delta/bar' % path, + ... mode='overwrite', replaceWhere='date >= "2012-01-01"') # doctest: +SKIP + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `to_delta`, " + "the existing index is lost when converting to Delta." + ) + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") # type: ignore[assignment] + + mode = validate_mode(mode) + self.spark.to_spark_io( + path=path, + mode=mode, + format="delta", + partition_cols=partition_cols, + index_col=index_col, + **options, + ) + + def to_parquet( + self, + path: str, + mode: str = "w", + partition_cols: Optional[Union[str, List[str]]] = None, + compression: Optional[str] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, + ) -> None: + """ + Write the DataFrame out as a Parquet file or directory. + + Parameters + ---------- + path : str, required + Path to write to. + mode : str + Python write mode, default 'w'. + + .. note:: mode can accept the strings for Spark writing mode. + Such as 'append', 'overwrite', 'ignore', 'error', 'errorifexists'. + + - 'append' (equivalent to 'a'): Append the new data to existing data. + - 'overwrite' (equivalent to 'w'): Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols : str or list of str, optional, default None + Names of partitioning columns + compression : str {'none', 'uncompressed', 'snappy', 'gzip', 'lzo', 'brotli', 'lz4', 'zstd'} + Compression codec to use when saving to file. If None is set, it uses the + value specified in `spark.sql.parquet.compression.codec`. + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default the index is always lost. + options : dict + All other options passed directly into Spark's data source. + + See Also + -------- + read_parquet + DataFrame.to_delta + DataFrame.to_table + DataFrame.to_spark_io + + Examples + -------- + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df + date country code + 0 2012-01-31 12:00:00 KR 1 + 1 2012-02-29 12:00:00 US 2 + 2 2012-03-31 12:00:00 JP 3 + + >>> df.to_parquet('%s/to_parquet/foo.parquet' % path, partition_cols='date') + + >>> df.to_parquet( + ... '%s/to_parquet/foo.parquet' % path, + ... mode = 'overwrite', + ... partition_cols=['date', 'country']) + + Notes + ----- + pandas API on Spark writes Parquet files into the directory, `path`, and writes + multiple part files in the directory unlike pandas. + pandas API on Spark respects HDFS's property such as 'fs.default.name'. + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `to_parquet`, " + "the existing index is lost when converting to Parquet." + ) + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + mode = validate_mode(mode) + builder = self.to_spark(index_col=index_col).write.mode(mode) + if partition_cols is not None: + builder.partitionBy(partition_cols) + if compression is not None: + builder.option("compression", compression) + builder.options(**options).format("parquet").save(path) + + def to_orc( + self, + path: str, + mode: str = "w", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """ + Write a DataFrame to the ORC format. + + Parameters + ---------- + path : str + Path to write to. + mode : str + Python write mode, default 'w'. + + .. note:: mode can accept the strings for Spark writing mode. + Such as 'append', 'overwrite', 'ignore', 'error', 'errorifexists'. + + - 'append' (equivalent to 'a'): Append the new data to existing data. + - 'overwrite' (equivalent to 'w'): Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols : str or list of str, optional, default None + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default the index is always lost. + options : dict + All other options passed directly into Spark's data source. + + See Also + -------- + read_orc + DataFrame.to_delta + DataFrame.to_parquet + DataFrame.to_table + DataFrame.to_spark_io + + Examples + -------- + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df + date country code + 0 2012-01-31 12:00:00 KR 1 + 1 2012-02-29 12:00:00 US 2 + 2 2012-03-31 12:00:00 JP 3 + + >>> df.to_orc('%s/to_orc/foo.orc' % path, partition_cols='date') + + >>> df.to_orc( + ... '%s/to_orc/foo.orc' % path, + ... mode = 'overwrite', + ... partition_cols=['date', 'country']) + + Notes + ----- + pandas API on Spark writes ORC files into the directory, `path`, and writes + multiple part files in the directory unlike pandas. + pandas API on Spark respects HDFS's property such as 'fs.default.name'. + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `to_orc`, " + "the existing index is lost when converting to ORC." + ) + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") # type: ignore[assignment] + + mode = validate_mode(mode) + self.spark.to_spark_io( + path=path, + mode=mode, + format="orc", + partition_cols=partition_cols, + index_col=index_col, + **options, + ) + + def to_spark_io( + self, + path: Optional[str] = None, + format: Optional[str] = None, + mode: str = "overwrite", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """An alias for :func:`DataFrame.spark.to_spark_io`. + See :meth:`pyspark.pandas.spark.accessors.SparkFrameMethods.to_spark_io`. + + .. deprecated:: 3.2.0 + Use :func:`DataFrame.spark.to_spark_io` instead. + """ + warnings.warn("Deprecated in 3.2, Use DataFrame.spark.to_spark_io instead.", FutureWarning) + return self.spark.to_spark_io(path, format, mode, partition_cols, index_col, **options) + + to_spark_io.__doc__ = SparkFrameMethods.to_spark_io.__doc__ + + def to_spark(self, index_col: Optional[Union[str, List[str]]] = None) -> PySparkDataFrame: + if index_col is None: + log_advice( + "If `index_col` is not specified for `to_spark`, " + "the existing index is lost when converting to Spark DataFrame." + ) + return self._to_spark(index_col) + + to_spark.__doc__ = SparkFrameMethods.__doc__ + + def _to_spark(self, index_col: Optional[Union[str, List[str]]] = None) -> PySparkDataFrame: + """ + Same as `to_spark()`, without issuing the advice log when `index_col` is not specified + for internal usage. + """ + return self.spark.frame(index_col) + + def to_pandas(self) -> pd.DataFrame: + """ + Return a pandas DataFrame. + + .. note:: This method should only be used if the resulting pandas DataFrame is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df.to_pandas() + dogs cats + 0 0.2 0.3 + 1 0.0 0.6 + 2 0.6 0.0 + 3 0.2 0.1 + """ + log_advice( + "`to_pandas` loads all data into the driver's memory. " + "It should only be used if the resulting pandas DataFrame is expected to be small." + ) + return self._to_pandas() + + def _to_pandas(self) -> pd.DataFrame: + """ + Same as `to_pandas()`, without issuing the advice log for internal usage. + """ + return self._internal.to_pandas_frame.copy() + + def assign(self, **kwargs: Any) -> "DataFrame": + """ + Assign new columns to a DataFrame. + + Returns a new object with all original columns in addition to new ones. + Existing columns that are re-assigned will be overwritten. + + Parameters + ---------- + **kwargs : dict of {str: callable, Series or Index} + The column names are keywords. If the values are + callable, they are computed on the DataFrame and + assigned to the new columns. The callable must not + change input DataFrame (though pandas-on-Spark doesn't check it). + If the values are not callable, (e.g. a Series or a literal), + they are simply assigned. + + Returns + ------- + DataFrame + A new DataFrame with the new columns in addition to + all the existing columns. + + Examples + -------- + >>> df = ps.DataFrame({'temp_c': [17.0, 25.0]}, + ... index=['Portland', 'Berkeley']) + >>> df + temp_c + Portland 17.0 + Berkeley 25.0 + + Where the value is a callable, evaluated on `df`: + + >>> df.assign(temp_f=lambda x: x.temp_c * 9 / 5 + 32) + temp_c temp_f + Portland 17.0 62.6 + Berkeley 25.0 77.0 + + Alternatively, the same behavior can be achieved by directly + referencing an existing Series or sequence and you can also + create multiple columns within the same assign. + + >>> assigned = df.assign(temp_f=df['temp_c'] * 9 / 5 + 32, + ... temp_k=df['temp_c'] + 273.15, + ... temp_idx=df.index) + >>> assigned[['temp_c', 'temp_f', 'temp_k', 'temp_idx']] + temp_c temp_f temp_k temp_idx + Portland 17.0 62.6 290.15 Portland + Berkeley 25.0 77.0 298.15 Berkeley + + Notes + ----- + Assigning multiple columns within the same ``assign`` is possible + but you cannot refer to newly created or modified columns. This + feature is supported in pandas for Python 3.6 and later but not in + pandas-on-Spark. In pandas-on-Spark, all items are computed first, + and then assigned. + """ + return self._assign(kwargs) + + def _assign(self, kwargs: Any) -> "DataFrame": + assert isinstance(kwargs, dict) + from pyspark.pandas.indexes import MultiIndex + from pyspark.pandas.series import IndexOpsMixin + + Column = get_column_class() + for k, v in kwargs.items(): + is_invalid_assignee = ( + not (isinstance(v, (IndexOpsMixin, Column)) or callable(v) or is_scalar(v)) + ) or isinstance(v, MultiIndex) + if is_invalid_assignee: + raise TypeError( + "Column assignment doesn't support type " "{0}".format(type(v).__name__) + ) + if callable(v): + kwargs[k] = v(self) + + pairs = { + (k if is_name_like_tuple(k) else (k,)): ( + (v.spark.column, v._internal.data_fields[0]) + if isinstance(v, IndexOpsMixin) and not isinstance(v, MultiIndex) + else (v, None) + if isinstance(v, Column) + else (F.lit(v), None) + ) + for k, v in kwargs.items() + } + + scols = [] + data_fields = [] + for label in self._internal.column_labels: + for i in range(len(label)): + if label[: len(label) - i] in pairs: + scol, field = pairs[label[: len(label) - i]] + + name = self._internal.spark_column_name_for(label) + scol = scol.alias(name) + if field is not None: + field = field.copy(name=name) + break + else: + scol = self._internal.spark_column_for(label) + field = self._internal.field_for(label) + scols.append(scol) + data_fields.append(field) + + column_labels = self._internal.column_labels.copy() + for label, (scol, field) in pairs.items(): + if label not in set(i[: len(label)] for i in self._internal.column_labels): + name = name_like_string(label) + scols.append(scol.alias(name)) + if field is not None: + field = field.copy(name=name) + data_fields.append(field) + + column_labels.append(label) + + level = self._internal.column_labels_level + column_labels = [ + tuple(list(label) + ([""] * (level - len(label)))) for label in column_labels + ] + + internal = self._internal.with_new_columns( + cast(Sequence[Union[PySparkColumn, "Series"]], scols), + column_labels=column_labels, + data_fields=data_fields, + ) + return DataFrame(internal) + + @staticmethod + def from_records( + data: Union[np.ndarray, List[tuple], dict, pd.DataFrame], + index: Union[str, list, np.ndarray] = None, + exclude: list = None, + columns: list = None, + coerce_float: bool = False, + nrows: int = None, + ) -> "DataFrame": + """ + Convert structured or recorded ndarray to DataFrame. + + Parameters + ---------- + data : ndarray (structured dtype), list of tuples, dict, or DataFrame + index : string, list of fields, array-like + Field of array to use as the index, alternately a specific set of input labels to use + exclude : sequence, default None + Columns or fields to exclude + columns : sequence, default None + Column names to use. If the passed data do not have names associated with them, this + argument provides names for the columns. Otherwise this argument indicates the order of + the columns in the result (any names not found in the data will become all-NA columns) + coerce_float : boolean, default False + Attempt to convert values of non-string, non-numeric objects (like decimal.Decimal) to + floating point, useful for SQL result sets + nrows : int, default None + Number of rows to read if data is an iterator + + Returns + ------- + df : DataFrame + + Examples + -------- + Use dict as input + + >>> ps.DataFrame.from_records({'A': [1, 2, 3]}) + A + 0 1 + 1 2 + 2 3 + + Use list of tuples as input + + >>> ps.DataFrame.from_records([(1, 2), (3, 4)]) + 0 1 + 0 1 2 + 1 3 4 + + Use NumPy array as input + + >>> ps.DataFrame.from_records(np.eye(3)) + 0 1 2 + 0 1.0 0.0 0.0 + 1 0.0 1.0 0.0 + 2 0.0 0.0 1.0 + """ + return DataFrame( + pd.DataFrame.from_records(data, index, exclude, columns, coerce_float, nrows) + ) + + def to_records( + self, + index: bool = True, + column_dtypes: Optional[Union[str, Dtype, Dict[Name, Union[str, Dtype]]]] = None, + index_dtypes: Optional[Union[str, Dtype, Dict[Name, Union[str, Dtype]]]] = None, + ) -> np.recarray: + """ + Convert DataFrame to a NumPy record array. + + Index will be included as the first field of the record array if + requested. + + .. note:: This method should only be used if the resulting NumPy ndarray is + expected to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + index : bool, default True + Include index in resulting record array, stored in 'index' + field or using the index label, if set. + column_dtypes : str, type, dict, default None + If a string or type, the data type to store all columns. If + a dictionary, a mapping of column names and indices (zero-indexed) + to specific data types. + index_dtypes : str, type, dict, default None + If a string or type, the data type to store all index levels. If + a dictionary, a mapping of index level names and indices + (zero-indexed) to specific data types. + This mapping is applied only if `index=True`. + + Returns + ------- + numpy.recarray + NumPy ndarray with the DataFrame labels as fields and each row + of the DataFrame as entries. + + See Also + -------- + DataFrame.from_records: Convert structured or record ndarray + to DataFrame. + numpy.recarray: An ndarray that allows field access using + attributes, analogous to typed columns in a + spreadsheet. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2], 'B': [0.5, 0.75]}, + ... index=['a', 'b']) + >>> df + A B + a 1 0.50 + b 2 0.75 + + >>> df.to_records() # doctest: +SKIP + rec.array([('a', 1, 0.5 ), ('b', 2, 0.75)], + dtype=[('index', 'O'), ('A', '>> df.to_records(index=False) # doctest: +SKIP + rec.array([(1, 0.5 ), (2, 0.75)], + dtype=[('A', '>> df.to_records(column_dtypes={"A": "int32"}) # doctest: +SKIP + rec.array([('a', 1, 0.5 ), ('b', 2, 0.75)], + dtype=[('index', 'O'), ('A', '>> df.to_records(index_dtypes=" "DataFrame": + """ + Make a copy of this object's indices and data. + + Parameters + ---------- + deep : bool, default True + this parameter is not supported but just dummy parameter to match pandas. + + Returns + ------- + copy : DataFrame + + Examples + -------- + >>> df = ps.DataFrame({'x': [1, 2], 'y': [3, 4], 'z': [5, 6], 'w': [7, 8]}, + ... columns=['x', 'y', 'z', 'w']) + >>> df + x y z w + 0 1 3 5 7 + 1 2 4 6 8 + >>> df_copy = df.copy() + >>> df_copy + x y z w + 0 1 3 5 7 + 1 2 4 6 8 + """ + return DataFrame(self._internal) + + def dropna( + self, + axis: Axis = 0, + how: str = "any", + thresh: Optional[int] = None, + subset: Optional[Union[Name, List[Name]]] = None, + inplace: bool = False, + ) -> Optional["DataFrame"]: + """ + Remove missing values. + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Determine if rows or columns which contain missing values are + removed. + + * 0, or 'index' : Drop rows which contain missing values. + how : {'any', 'all'}, default 'any' + Determine if row or column is removed from DataFrame, when we have + at least one NA or all NA. + + * 'any' : If any NA values are present, drop that row or column. + * 'all' : If all values are NA, drop that row or column. + + thresh : int, optional + Require that many non-NA values. + subset : array-like, optional + Labels along other axis to consider, e.g. if you are dropping rows + these would be a list of columns to include. + inplace : bool, default False + If True, do operation inplace and return None. + + Returns + ------- + DataFrame + DataFrame with NA entries dropped from it. + + See Also + -------- + DataFrame.drop : Drop specified labels from columns. + DataFrame.isnull: Indicate missing values. + DataFrame.notnull : Indicate existing (non-missing) values. + + Examples + -------- + >>> df = ps.DataFrame({"name": ['Alfred', 'Batman', 'Catwoman'], + ... "toy": [None, 'Batmobile', 'Bullwhip'], + ... "born": [None, "1940-04-25", None]}, + ... columns=['name', 'toy', 'born']) + >>> df + name toy born + 0 Alfred None None + 1 Batman Batmobile 1940-04-25 + 2 Catwoman Bullwhip None + + Drop the rows where at least one element is missing. + + >>> df.dropna() + name toy born + 1 Batman Batmobile 1940-04-25 + + Drop the columns where at least one element is missing. + + >>> df.dropna(axis='columns') + name + 0 Alfred + 1 Batman + 2 Catwoman + + Drop the rows where all elements are missing. + + >>> df.dropna(how='all') + name toy born + 0 Alfred None None + 1 Batman Batmobile 1940-04-25 + 2 Catwoman Bullwhip None + + Keep only the rows with at least 2 non-NA values. + + >>> df.dropna(thresh=2) + name toy born + 1 Batman Batmobile 1940-04-25 + 2 Catwoman Bullwhip None + + Define in which columns to look for missing values. + + >>> df.dropna(subset=['name', 'born']) + name toy born + 1 Batman Batmobile 1940-04-25 + + Keep the DataFrame with valid entries in the same variable. + + >>> df.dropna(inplace=True) + >>> df + name toy born + 1 Batman Batmobile 1940-04-25 + """ + axis = validate_axis(axis) + inplace = validate_bool_kwarg(inplace, "inplace") + + if thresh is None: + if how is None: + raise TypeError("must specify how or thresh") + elif how not in ("any", "all"): + raise ValueError("invalid how option: {h}".format(h=how)) + + labels: Optional[List[Label]] + if subset is not None: + if isinstance(subset, str): + labels = [(subset,)] + elif isinstance(subset, tuple): + labels = [subset] + else: + labels = [sub if isinstance(sub, tuple) else (sub,) for sub in subset] + else: + labels = None + + if axis == 0: + if labels is not None: + invalids = [label for label in labels if label not in self._internal.column_labels] + if len(invalids) > 0: + raise KeyError(invalids) + else: + labels = self._internal.column_labels + + cnt = reduce( + lambda x, y: x + y, + [ + F.when(self._psser_for(label).notna().spark.column, 1).otherwise(0) + for label in labels + ], + F.lit(0), + ) + if thresh is not None: + pred = cnt >= F.lit(int(thresh)) + elif how == "any": + pred = cnt == F.lit(len(labels)) + elif how == "all": + pred = cnt > F.lit(0) + + internal = self._internal.with_filter(pred) + if inplace: + self._update_internal_frame(internal) + return None + else: + return DataFrame(internal) + else: + assert axis == 1 + + internal = self._internal.resolved_copy + + if labels is not None: + if any(len(lbl) != internal.index_level for lbl in labels): + raise ValueError( + "The length of each subset must be the same as the index size." + ) + + cond = reduce( + lambda x, y: x | y, + [ + reduce( + lambda x, y: x & y, + [ + scol == F.lit(part) + for part, scol in zip(lbl, internal.index_spark_columns) + ], + ) + for lbl in labels + ], + ) + + internal = internal.with_filter(cond) + + psdf: DataFrame = DataFrame(internal) + + null_counts = [] + for label in internal.column_labels: + psser = psdf._psser_for(label) + cond = psser.isnull().spark.column + null_counts.append( + F.sum(F.when(~cond, 1).otherwise(0)).alias(name_like_string(label)) + ) + + counts = internal.spark_frame.select(null_counts + [F.count("*")]).head() + + if thresh is not None: + column_labels = [ + label + for label, cnt in zip(internal.column_labels, counts) + if (cnt or 0) >= int(thresh) + ] + elif how == "any": + column_labels = [ + label + for label, cnt in zip(internal.column_labels, counts) + if (cnt or 0) == counts[-1] + ] + elif how == "all": + column_labels = [ + label for label, cnt in zip(internal.column_labels, counts) if (cnt or 0) > 0 + ] + + psdf = self[column_labels] + if inplace: + self._update_internal_frame(psdf._internal) + return None + else: + return psdf + + # TODO: add 'limit' when value parameter exists + def fillna( + self, + value: Optional[Union[Any, Dict[Name, Any]]] = None, + method: Optional[str] = None, + axis: Optional[Axis] = None, + inplace: bool = False, + limit: Optional[int] = None, + ) -> Optional["DataFrame"]: + """Fill NA/NaN values. + + .. note:: the current implementation of 'method' parameter in fillna uses Spark's Window + without specifying partition specification. This leads to moving all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + value : scalar, dict, Series + Value to use to fill holes. alternately a dict/Series of values + specifying which value to use for each column. + DataFrame is not supported. + method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None + Method to use for filling holes in reindexed Series pad / ffill: propagate last valid + observation forward to next valid backfill / bfill: + use NEXT valid observation to fill gap + axis : {0 or `index`} + 1 and `columns` are not supported. + inplace : boolean, default False + Fill in place (do not create a new object) + limit : int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame + DataFrame with NA entries filled. + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'A': [None, 3, None, None], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> df + A B C D + 0 NaN 2.0 NaN 0 + 1 3.0 4.0 NaN 1 + 2 NaN NaN NaN 5 + 3 NaN 3.0 1.0 4 + + Replace all NaN elements with 0s. + + >>> df.fillna(0) + A B C D + 0 0.0 2.0 0.0 0 + 1 3.0 4.0 0.0 1 + 2 0.0 0.0 0.0 5 + 3 0.0 3.0 1.0 4 + + We can also propagate non-null values forward or backward. + + >>> df.fillna(method='ffill') + A B C D + 0 NaN 2.0 NaN 0 + 1 3.0 4.0 NaN 1 + 2 3.0 4.0 NaN 5 + 3 3.0 3.0 1.0 4 + + Replace all NaN elements in column 'A', 'B', 'C', and 'D', with 0, 1, + 2, and 3 respectively. + + >>> values = {'A': 0, 'B': 1, 'C': 2, 'D': 3} + >>> df.fillna(value=values) + A B C D + 0 0.0 2.0 2.0 0 + 1 3.0 4.0 2.0 1 + 2 0.0 1.0 2.0 5 + 3 0.0 3.0 1.0 4 + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError("fillna currently only works for axis=0 or axis='index'") + + if value is not None: + if not isinstance(value, (float, int, str, bool, dict, pd.Series)): + raise TypeError("Unsupported type %s" % type(value).__name__) + if limit is not None: + raise ValueError("limit parameter for value is not support now") + if isinstance(value, pd.Series): + value = value.to_dict() + if isinstance(value, dict): + for v in value.values(): + if not isinstance(v, (float, int, str, bool)): + raise TypeError("Unsupported type %s" % type(v).__name__) + value = {k if is_name_like_tuple(k) else (k,): v for k, v in value.items()} + + def op(psser: ps.Series) -> ps.Series: + label = psser._column_label + for k, v in value.items(): + if k == label[: len(k)]: + return psser._fillna( + value=value[k], method=method, axis=axis, limit=limit + ) + else: + return psser + + else: + + def op(psser: ps.Series) -> ps.Series: + return psser._fillna(value=value, method=method, axis=axis, limit=limit) + + elif method is not None: + + def op(psser: ps.Series) -> ps.Series: + return psser._fillna(value=value, method=method, axis=axis, limit=limit) + + else: + raise ValueError("Must specify a fillna 'value' or 'method' parameter.") + + psdf = self._apply_series_op(op, should_resolve=(method is not None)) + + inplace = validate_bool_kwarg(inplace, "inplace") + if inplace: + self._update_internal_frame(psdf._internal, check_same_anchor=False) + return None + else: + return psdf + + def interpolate( + self, + method: str = "linear", + limit: Optional[int] = None, + limit_direction: Optional[str] = None, + limit_area: Optional[str] = None, + ) -> "DataFrame": + if method not in ["linear"]: + raise NotImplementedError("interpolate currently works only for method='linear'") + if (limit is not None) and (not limit > 0): + raise ValueError("limit must be > 0.") + if (limit_direction is not None) and ( + limit_direction not in ["forward", "backward", "both"] + ): + raise ValueError("invalid limit_direction: '{}'".format(limit_direction)) + if (limit_area is not None) and (limit_area not in ["inside", "outside"]): + raise ValueError("invalid limit_area: '{}'".format(limit_area)) + + numeric_col_names = [] + for label in self._internal.column_labels: + psser = self._psser_for(label) + if isinstance(psser.spark.data_type, (NumericType, BooleanType)): + numeric_col_names.append(psser.name) + + psdf = self[numeric_col_names] + return psdf._apply_series_op( + lambda psser: psser._interpolate( + method=method, limit=limit, limit_direction=limit_direction, limit_area=limit_area + ), + should_resolve=True, + ) + + def replace( + self, + to_replace: Optional[Union[Any, List, Tuple, Dict]] = None, + value: Optional[Any] = None, + inplace: bool = False, + limit: Optional[int] = None, + regex: bool = False, + method: str = "pad", + ) -> Optional["DataFrame"]: + """ + Returns a new DataFrame replacing a value with another value. + + Parameters + ---------- + to_replace : int, float, string, list, tuple or dict + Value to be replaced. + value : int, float, string, list or tuple + Value to use to replace holes. The replacement value must be an int, float, + or string. + If value is a list or tuple, value should be of the same length with to_replace. + inplace : boolean, default False + Fill in place (do not create a new object) + + Returns + ------- + DataFrame + Object after replacement. + + Examples + -------- + >>> df = ps.DataFrame({"name": ['Ironman', 'Captain America', 'Thor', 'Hulk'], + ... "weapon": ['Mark-45', 'Shield', 'Mjolnir', 'Smash']}, + ... columns=['name', 'weapon']) + >>> df + name weapon + 0 Ironman Mark-45 + 1 Captain America Shield + 2 Thor Mjolnir + 3 Hulk Smash + + Scalar `to_replace` and `value` + + >>> df.replace('Ironman', 'War-Machine') + name weapon + 0 War-Machine Mark-45 + 1 Captain America Shield + 2 Thor Mjolnir + 3 Hulk Smash + + List like `to_replace` and `value` + + >>> df.replace(['Ironman', 'Captain America'], ['Rescue', 'Hawkeye'], inplace=True) + >>> df + name weapon + 0 Rescue Mark-45 + 1 Hawkeye Shield + 2 Thor Mjolnir + 3 Hulk Smash + + Dicts can be used to specify different replacement values for different existing values + To use a dict in this way the value parameter should be None + + >>> df.replace({'Mjolnir': 'Stormbuster'}) + name weapon + 0 Rescue Mark-45 + 1 Hawkeye Shield + 2 Thor Stormbuster + 3 Hulk Smash + + Dict can specify that different values should be replaced in different columns + The value parameter should not be None in this case + + >>> df.replace({'weapon': 'Mjolnir'}, 'Stormbuster') + name weapon + 0 Rescue Mark-45 + 1 Hawkeye Shield + 2 Thor Stormbuster + 3 Hulk Smash + + Nested dictionaries + The value parameter should be None to use a nested dict in this way + + >>> df.replace({'weapon': {'Mjolnir': 'Stormbuster'}}) + name weapon + 0 Rescue Mark-45 + 1 Hawkeye Shield + 2 Thor Stormbuster + 3 Hulk Smash + """ + if method != "pad": + raise NotImplementedError("replace currently works only for method='pad") + if limit is not None: + raise NotImplementedError("replace currently works only when limit=None") + if regex is not False: + raise NotImplementedError("replace currently doesn't supports regex") + inplace = validate_bool_kwarg(inplace, "inplace") + + if value is not None and not isinstance(value, (int, float, str, list, tuple, dict)): + raise TypeError("Unsupported type {}".format(type(value).__name__)) + if to_replace is not None and not isinstance( + to_replace, (int, float, str, list, tuple, dict) + ): + raise TypeError("Unsupported type {}".format(type(to_replace).__name__)) + + if isinstance(value, (list, tuple)) and isinstance(to_replace, (list, tuple)): + if len(value) != len(to_replace): + raise ValueError("Length of to_replace and value must be same") + + if isinstance(to_replace, dict) and ( + value is not None or all(isinstance(i, dict) for i in to_replace.values()) + ): + to_replace_dict = to_replace + + def op(psser: ps.Series) -> ps.Series: + if psser.name in to_replace_dict: + return psser.replace( + to_replace=to_replace_dict[psser.name], value=value, regex=regex + ) + else: + return psser + + else: + + def op(psser: ps.Series) -> ps.Series: + return psser.replace(to_replace=to_replace, value=value, regex=regex) + + psdf = self._apply_series_op(op) + if inplace: + self._update_internal_frame(psdf._internal) + return None + else: + return psdf + + def clip(self, lower: Union[float, int] = None, upper: Union[float, int] = None) -> "DataFrame": + """ + Trim values at input threshold(s). + + Assigns values outside boundary-to-boundary values. + + Parameters + ---------- + lower : float or int, default None + Minimum threshold value. All values below this threshold will be set to it. + upper : float or int, default None + Maximum threshold value. All values above this threshold will be set to it. + + Returns + ------- + DataFrame + DataFrame with the values outside the clip boundaries replaced. + + Examples + -------- + >>> ps.DataFrame({'A': [0, 2, 4]}).clip(1, 3) + A + 0 1 + 1 2 + 2 3 + + Notes + ----- + One difference between this implementation and pandas is that running + pd.DataFrame({'A': ['a', 'b']}).clip(0, 1) will crash with "TypeError: '<=' not supported + between instances of 'str' and 'int'" while ps.DataFrame({'A': ['a', 'b']}).clip(0, 1) + will output the original DataFrame, simply ignoring the incompatible types. + """ + if is_list_like(lower) or is_list_like(upper): + raise TypeError( + "List-like value are not supported for 'lower' and 'upper' at the " + "moment" + ) + + if lower is None and upper is None: + return self + + return self._apply_series_op(lambda psser: psser.clip(lower=lower, upper=upper)) + + def head(self, n: int = 5) -> "DataFrame": + """ + Return the first `n` rows. + + This function returns the first `n` rows for the object based + on position. It is useful for quickly testing if your object + has the right type of data in it. + + Parameters + ---------- + n : int, default 5 + Number of rows to select. + + Returns + ------- + obj_head : same type as caller + The first `n` rows of the caller object. + + Examples + -------- + >>> df = ps.DataFrame({'animal':['alligator', 'bee', 'falcon', 'lion', + ... 'monkey', 'parrot', 'shark', 'whale', 'zebra']}) + >>> df + animal + 0 alligator + 1 bee + 2 falcon + 3 lion + 4 monkey + 5 parrot + 6 shark + 7 whale + 8 zebra + + Viewing the first 5 lines + + >>> df.head() + animal + 0 alligator + 1 bee + 2 falcon + 3 lion + 4 monkey + + Viewing the first `n` lines (three in this case) + + >>> df.head(3) + animal + 0 alligator + 1 bee + 2 falcon + """ + if n < 0: + n = len(self) + n + if n <= 0: + return DataFrame(self._internal.with_filter(F.lit(False))) + else: + sdf = self._internal.resolved_copy.spark_frame + if get_option("compute.ordered_head"): + sdf = sdf.orderBy(NATURAL_ORDER_COLUMN_NAME) + return DataFrame(self._internal.with_new_sdf(sdf.limit(n))) + + def last(self, offset: Union[str, DateOffset]) -> "DataFrame": + """ + Select final periods of time series data based on a date offset. + + When having a DataFrame with dates as index, this function can + select the last few rows based on a date offset. + + Parameters + ---------- + offset : str or DateOffset + The offset length of the data that will be selected. For instance, + '3D' will display all the rows having their index within the last 3 days. + + Returns + ------- + DataFrame + A subset of the caller. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + Examples + -------- + + >>> index = pd.date_range('2018-04-09', periods=4, freq='2D') + >>> psdf = ps.DataFrame({'A': [1, 2, 3, 4]}, index=index) + >>> psdf + A + 2018-04-09 1 + 2018-04-11 2 + 2018-04-13 3 + 2018-04-15 4 + + Get the rows for the last 3 days: + + >>> psdf.last('3D') + A + 2018-04-13 3 + 2018-04-15 4 + + Notice the data for 3 last calendar days were returned, not the last + 3 observed days in the dataset, and therefore data for 2018-04-11 was + not returned. + """ + # Check index type should be format DateTime + if not isinstance(self.index, ps.DatetimeIndex): + raise TypeError("'last' only supports a DatetimeIndex") + + from_date = cast( + int, + cast(datetime.datetime, self.index.max()) - cast(datetime.timedelta, to_offset(offset)), + ) + + return cast(DataFrame, self.loc[from_date:]) + + def first(self, offset: Union[str, DateOffset]) -> "DataFrame": + """ + Select first periods of time series data based on a date offset. + + When having a DataFrame with dates as index, this function can + select the first few rows based on a date offset. + + Parameters + ---------- + offset : str or DateOffset + The offset length of the data that will be selected. For instance, + '3D' will display all the rows having their index within the first 3 days. + + Returns + ------- + DataFrame + A subset of the caller. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + Examples + -------- + + >>> index = pd.date_range('2018-04-09', periods=4, freq='2D') + >>> psdf = ps.DataFrame({'A': [1, 2, 3, 4]}, index=index) + >>> psdf + A + 2018-04-09 1 + 2018-04-11 2 + 2018-04-13 3 + 2018-04-15 4 + + Get the rows for the last 3 days: + + >>> psdf.first('3D') + A + 2018-04-09 1 + 2018-04-11 2 + + Notice the data for 3 first calendar days were returned, not the first + 3 observed days in the dataset, and therefore data for 2018-04-13 was + not returned. + """ + # Check index type should be format DatetimeIndex + if not isinstance(self.index, ps.DatetimeIndex): + raise TypeError("'first' only supports a DatetimeIndex") + + to_date = cast( + int, + cast(datetime.datetime, self.index.min()) + cast(datetime.timedelta, to_offset(offset)), + ) + + return cast(DataFrame, self.loc[:to_date]) + + def pivot_table( + self, + values: Optional[Union[Name, List[Name]]] = None, + index: Optional[List[Name]] = None, + columns: Optional[Name] = None, + aggfunc: Union[str, Dict[Name, str]] = "mean", + fill_value: Optional[Any] = None, + ) -> "DataFrame": + """ + Create a spreadsheet-style pivot table as a DataFrame. The levels in + the pivot table will be stored in MultiIndex objects (hierarchical + indexes) on the index and columns of the result DataFrame. + + Parameters + ---------- + values : column to aggregate. + They should be either a list less than three or a string. + index : column (string) or list of columns + If an array is passed, it must be the same length as the data. + The list should contain string. + columns : column + Columns used in the pivot operation. Only one column is supported and + it should be a string. + aggfunc : function (string), dict, default mean + If dict is passed, the key is column to aggregate and value + is function or list of functions. + fill_value : scalar, default None + Value to replace missing values with. + + Returns + ------- + table : DataFrame + + Examples + -------- + >>> df = ps.DataFrame({"A": ["foo", "foo", "foo", "foo", "foo", + ... "bar", "bar", "bar", "bar"], + ... "B": ["one", "one", "one", "two", "two", + ... "one", "one", "two", "two"], + ... "C": ["small", "large", "large", "small", + ... "small", "large", "small", "small", + ... "large"], + ... "D": [1, 2, 2, 3, 3, 4, 5, 6, 7], + ... "E": [2, 4, 5, 5, 6, 6, 8, 9, 9]}, + ... columns=['A', 'B', 'C', 'D', 'E']) + >>> df + A B C D E + 0 foo one small 1 2 + 1 foo one large 2 4 + 2 foo one large 2 5 + 3 foo two small 3 5 + 4 foo two small 3 6 + 5 bar one large 4 6 + 6 bar one small 5 8 + 7 bar two small 6 9 + 8 bar two large 7 9 + + This first example aggregates values by taking the sum. + + >>> table = df.pivot_table(values='D', index=['A', 'B'], + ... columns='C', aggfunc='sum') + >>> table.sort_index() # doctest: +NORMALIZE_WHITESPACE + C large small + A B + bar one 4.0 5 + two 7.0 6 + foo one 4.0 1 + two NaN 6 + + We can also fill missing values using the `fill_value` parameter. + + >>> table = df.pivot_table(values='D', index=['A', 'B'], + ... columns='C', aggfunc='sum', fill_value=0) + >>> table.sort_index() # doctest: +NORMALIZE_WHITESPACE + C large small + A B + bar one 4 5 + two 7 6 + foo one 4 1 + two 0 6 + + We can also calculate multiple types of aggregations for any given + value column. + + >>> table = df.pivot_table(values=['D'], index =['C'], + ... columns="A", aggfunc={'D': 'mean'}) + >>> table.sort_index() # doctest: +NORMALIZE_WHITESPACE + D + A bar foo + C + large 5.5 2.000000 + small 5.5 2.333333 + + The next example aggregates on multiple values. + + >>> table = df.pivot_table(index=['C'], columns="A", values=['D', 'E'], + ... aggfunc={'D': 'mean', 'E': 'sum'}) + >>> table.sort_index() # doctest: +NORMALIZE_WHITESPACE + D E + A bar foo bar foo + C + large 5.5 2.000000 15 9 + small 5.5 2.333333 17 13 + """ + if not is_name_like_value(columns): + raise TypeError("columns should be one column name.") + + if not is_name_like_value(values) and not ( + isinstance(values, list) and all(is_name_like_value(v) for v in values) + ): + raise TypeError("values should be one column or list of columns.") + + if not isinstance(aggfunc, str) and ( + not isinstance(aggfunc, dict) + or not all( + is_name_like_value(key) and isinstance(value, str) for key, value in aggfunc.items() + ) + ): + raise TypeError( + "aggfunc must be a dict mapping from column name " + "to aggregate functions (string)." + ) + + if isinstance(aggfunc, dict) and index is None: + raise NotImplementedError( + "pivot_table doesn't support aggfunc" " as dict and without index." + ) + if isinstance(values, list) and index is None: + raise NotImplementedError("values can't be a list without index.") + + if columns not in self.columns: + raise ValueError("Wrong columns {}.".format(name_like_string(columns))) + if not is_name_like_tuple(columns): + columns = (columns,) + + if isinstance(values, list): + values = [col if is_name_like_tuple(col) else (col,) for col in values] + if not all( + isinstance(self._internal.spark_type_for(col), NumericType) for col in values + ): + raise TypeError("values should be a numeric type.") + else: + values = values if is_name_like_tuple(values) else (values,) + if not isinstance(self._internal.spark_type_for(values), NumericType): + raise TypeError("values should be a numeric type.") + + if isinstance(aggfunc, str): + if isinstance(values, list): + agg_cols = [ + F.expr( + "{1}(`{0}`) as `{0}`".format( + self._internal.spark_column_name_for(value), aggfunc + ) + ) + for value in values + ] + else: + agg_cols = [ + F.expr( + "{1}(`{0}`) as `{0}`".format( + self._internal.spark_column_name_for(values), aggfunc + ) + ) + ] + elif isinstance(aggfunc, dict): + aggfunc = { + key if is_name_like_tuple(key) else (key,): value for key, value in aggfunc.items() + } + agg_cols = [ + F.expr( + "{1}(`{0}`) as `{0}`".format(self._internal.spark_column_name_for(key), value) + ) + for key, value in aggfunc.items() + ] + agg_columns = [key for key, _ in aggfunc.items()] + + if set(agg_columns) != set(values): + raise ValueError("Columns in aggfunc must be the same as values.") + + sdf = self._internal.resolved_copy.spark_frame + if index is None: + sdf = ( + sdf.groupBy() + .pivot(pivot_col=self._internal.spark_column_name_for(columns)) + .agg(*agg_cols) + ) + + elif isinstance(index, list): + index = [label if is_name_like_tuple(label) else (label,) for label in index] + sdf = ( + sdf.groupBy([self._internal.spark_column_name_for(label) for label in index]) + .pivot(pivot_col=self._internal.spark_column_name_for(columns)) + .agg(*agg_cols) + ) + else: + raise TypeError("index should be a None or a list of columns.") + + if fill_value is not None and isinstance(fill_value, (int, float)): + sdf = sdf.fillna(fill_value) + + psdf: DataFrame + if index is not None: + index_columns = [self._internal.spark_column_name_for(label) for label in index] + index_fields = [self._internal.field_for(label) for label in index] + + if isinstance(values, list): + data_columns = [column for column in sdf.columns if column not in index_columns] + + if len(values) > 1: + # If we have two values, Spark will return column's name + # in this format: column_values, where column contains + # their values in the DataFrame and values is + # the column list passed to the pivot_table(). + # E.g. if column is b and values is ['b','e'], + # then ['2_b', '2_e', '3_b', '3_e']. + + # We sort the columns of Spark DataFrame by values. + data_columns.sort(key=lambda x: x.split("_", 1)[1]) + sdf = sdf.select(index_columns + data_columns) + + column_name_to_index = dict( + zip(self._internal.data_spark_column_names, self._internal.column_labels) + ) + column_labels = [ + tuple(list(column_name_to_index[name.split("_")[1]]) + [name.split("_")[0]]) + for name in data_columns + ] + column_label_names = ( + [cast(Optional[Name], None)] * column_labels_level(values) + ) + [columns] + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_columns], + index_names=index, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + column_label_names=column_label_names, + ) + psdf = DataFrame(internal) + else: + column_labels = [tuple(list(values[0]) + [column]) for column in data_columns] + column_label_names = ([cast(Optional[Name], None)] * len(values[0])) + [columns] + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_columns], + index_names=index, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + column_label_names=column_label_names, + ) + psdf = DataFrame(internal) + else: + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_columns], + index_names=index, + index_fields=index_fields, + column_label_names=[columns], + ) + psdf = DataFrame(internal) + else: + index_values = values + index_map: Dict[str, Optional[Label]] = {} + for i, index_value in enumerate(index_values): + colname = SPARK_INDEX_NAME_FORMAT(i) + sdf = sdf.withColumn(colname, F.lit(index_value)) + index_map[colname] = None + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_map.keys()], + index_names=list(index_map.values()), + column_label_names=[columns], + ) + psdf = DataFrame(internal) + + psdf_columns = psdf.columns + if isinstance(psdf_columns, pd.MultiIndex): + psdf.columns = psdf_columns.set_levels( + psdf_columns.levels[-1].astype( # type: ignore[index] + spark_type_to_pandas_dtype(self._psser_for(columns).spark.data_type) + ), + level=-1, + ) + else: + psdf.columns = psdf_columns.astype( + spark_type_to_pandas_dtype(self._psser_for(columns).spark.data_type) + ) + + return psdf + + def pivot( + self, + index: Optional[Name] = None, + columns: Optional[Name] = None, + values: Optional[Name] = None, + ) -> "DataFrame": + """ + Return reshaped DataFrame organized by given index / column values. + + Reshape data (produce a "pivot" table) based on column values. Uses + unique values from specified `index` / `columns` to form axes of the + resulting DataFrame. This function does not support data + aggregation. + + Parameters + ---------- + index : string, optional + Column to use to make new frame's index. If None, uses + existing index. + columns : string + Column to use to make new frame's columns. + values : string, object or a list of the previous + Column(s) to use for populating new frame's values. + + Returns + ------- + DataFrame + Returns reshaped DataFrame. + + See Also + -------- + DataFrame.pivot_table : Generalization of pivot that can handle + duplicate values for one index/column pair. + + Examples + -------- + >>> df = ps.DataFrame({'foo': ['one', 'one', 'one', 'two', 'two', + ... 'two'], + ... 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], + ... 'baz': [1, 2, 3, 4, 5, 6], + ... 'zoo': ['x', 'y', 'z', 'q', 'w', 't']}, + ... columns=['foo', 'bar', 'baz', 'zoo']) + >>> df + foo bar baz zoo + 0 one A 1 x + 1 one B 2 y + 2 one C 3 z + 3 two A 4 q + 4 two B 5 w + 5 two C 6 t + + >>> df.pivot(index='foo', columns='bar', values='baz').sort_index() + ... # doctest: +NORMALIZE_WHITESPACE + bar A B C + foo + one 1 2 3 + two 4 5 6 + + >>> df.pivot(columns='bar', values='baz').sort_index() # doctest: +NORMALIZE_WHITESPACE + bar A B C + 0 1.0 NaN NaN + 1 NaN 2.0 NaN + 2 NaN NaN 3.0 + 3 4.0 NaN NaN + 4 NaN 5.0 NaN + 5 NaN NaN 6.0 + + Notice that, unlike pandas raises an ValueError when duplicated values are found. + Pandas-on-Spark's pivot still works with its first value it meets during operation because + pivot is an expensive operation, and it is preferred to permissively execute over failing + fast when processing large data. + + >>> df = ps.DataFrame({"foo": ['one', 'one', 'two', 'two'], + ... "bar": ['A', 'A', 'B', 'C'], + ... "baz": [1, 2, 3, 4]}, columns=['foo', 'bar', 'baz']) + >>> df + foo bar baz + 0 one A 1 + 1 one A 2 + 2 two B 3 + 3 two C 4 + + >>> df.pivot(index='foo', columns='bar', values='baz').sort_index() + ... # doctest: +NORMALIZE_WHITESPACE + bar A B C + foo + one 1.0 NaN NaN + two NaN 3.0 4.0 + + It also supports multi-index and multi-index column. + >>> df.columns = pd.MultiIndex.from_tuples([('a', 'foo'), ('a', 'bar'), ('b', 'baz')]) + + >>> df = df.set_index(('a', 'bar'), append=True) + >>> df # doctest: +NORMALIZE_WHITESPACE + a b + foo baz + (a, bar) + 0 A one 1 + 1 A one 2 + 2 B two 3 + 3 C two 4 + + >>> df.pivot(columns=('a', 'foo'), values=('b', 'baz')).sort_index() + ... # doctest: +NORMALIZE_WHITESPACE + ('a', 'foo') one two + (a, bar) + 0 A 1.0 NaN + 1 A 2.0 NaN + 2 B NaN 3.0 + 3 C NaN 4.0 + + """ + if columns is None: + raise ValueError("columns should be set.") + + if values is None: + raise ValueError("values should be set.") + + should_use_existing_index = index is not None + if should_use_existing_index: + df = self + index_labels = [index] + else: + # The index after `reset_index()` will never be used, so use "distributed" index + # as a dummy to avoid overhead. + with option_context("compute.default_index_type", "distributed"): + df = self.reset_index() + index_labels = df._internal.column_labels[: self._internal.index_level] + + df = df.pivot_table(index=index_labels, columns=columns, values=values, aggfunc="first") + + if should_use_existing_index: + return df + else: + internal = df._internal.copy(index_names=self._internal.index_names) + return DataFrame(internal) + + @property + def columns(self) -> pd.Index: + """The column labels of the DataFrame.""" + names = [ + name if name is None or len(name) > 1 else name[0] + for name in self._internal.column_label_names + ] + if self._internal.column_labels_level > 1: + columns = pd.MultiIndex.from_tuples(self._internal.column_labels, names=names) + else: + columns = pd.Index([label[0] for label in self._internal.column_labels], name=names[0]) + return columns + + @columns.setter + def columns(self, columns: Union[pd.Index, List[Name]]) -> None: + if isinstance(columns, pd.MultiIndex): + column_labels = columns.tolist() + else: + column_labels = [ + col if is_name_like_tuple(col, allow_none=False) else (col,) for col in columns + ] + + if len(self._internal.column_labels) != len(column_labels): + raise ValueError( + "Length mismatch: Expected axis has {} elements, " + "new values have {} elements".format( + len(self._internal.column_labels), len(column_labels) + ) + ) + + column_label_names: Optional[List] + if isinstance(columns, pd.Index): + column_label_names = [ + name if is_name_like_tuple(name) else (name,) for name in columns.names + ] + else: + column_label_names = None + + pssers = [ + self._psser_for(label).rename(name) + for label, name in zip(self._internal.column_labels, column_labels) + ] + self._update_internal_frame( + self._internal.with_new_columns(pssers, column_label_names=column_label_names) + ) + + @property + def dtypes(self) -> pd.Series: + """Return the dtypes in the DataFrame. + + This returns a Series with the data type of each column. The result's index is the original + DataFrame's columns. Columns with mixed types are stored with the object dtype. + + Returns + ------- + pd.Series + The data type of each column. + + Examples + -------- + >>> df = ps.DataFrame({'a': list('abc'), + ... 'b': list(range(1, 4)), + ... 'c': np.arange(3, 6).astype('i1'), + ... 'd': np.arange(4.0, 7.0, dtype='float64'), + ... 'e': [True, False, True], + ... 'f': pd.date_range('20130101', periods=3)}, + ... columns=['a', 'b', 'c', 'd', 'e', 'f']) + >>> df.dtypes + a object + b int64 + c int8 + d float64 + e bool + f datetime64[ns] + dtype: object + """ + return pd.Series( + [self._psser_for(label).dtype for label in self._internal.column_labels], + index=pd.Index( + [label if len(label) > 1 else label[0] for label in self._internal.column_labels] + ), + ) + + def select_dtypes( + self, + include: Optional[Union[str, List[str]]] = None, + exclude: Optional[Union[str, List[str]]] = None, + ) -> "DataFrame": + """ + Return a subset of the DataFrame's columns based on the column dtypes. + + Parameters + ---------- + include, exclude : scalar or list-like + A selection of dtypes or strings to be included/excluded. At least + one of these parameters must be supplied. It also takes Spark SQL + DDL type strings, for instance, 'string' and 'date'. + + Returns + ------- + DataFrame + The subset of the frame including the dtypes in ``include`` and + excluding the dtypes in ``exclude``. + + Raises + ------ + ValueError + * If both of ``include`` and ``exclude`` are empty + + >>> df = ps.DataFrame({'a': [1, 2] * 3, + ... 'b': [True, False] * 3, + ... 'c': [1.0, 2.0] * 3}) + >>> df.select_dtypes() + Traceback (most recent call last): + ... + ValueError: at least one of include or exclude must be nonempty + + * If ``include`` and ``exclude`` have overlapping elements + + >>> df = ps.DataFrame({'a': [1, 2] * 3, + ... 'b': [True, False] * 3, + ... 'c': [1.0, 2.0] * 3}) + >>> df.select_dtypes(include='a', exclude='a') + Traceback (most recent call last): + ... + ValueError: include and exclude overlap on {'a'} + + Notes + ----- + * To select datetimes, use ``np.datetime64``, ``'datetime'`` or + ``'datetime64'`` + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2] * 3, + ... 'b': [True, False] * 3, + ... 'c': [1.0, 2.0] * 3, + ... 'd': ['a', 'b'] * 3}, columns=['a', 'b', 'c', 'd']) + >>> df + a b c d + 0 1 True 1.0 a + 1 2 False 2.0 b + 2 1 True 1.0 a + 3 2 False 2.0 b + 4 1 True 1.0 a + 5 2 False 2.0 b + + >>> df.select_dtypes(include='bool') + b + 0 True + 1 False + 2 True + 3 False + 4 True + 5 False + + >>> df.select_dtypes(include=['float64'], exclude=['int']) + c + 0 1.0 + 1 2.0 + 2 1.0 + 3 2.0 + 4 1.0 + 5 2.0 + + >>> df.select_dtypes(include=['int'], exclude=['float64']) + a + 0 1 + 1 2 + 2 1 + 3 2 + 4 1 + 5 2 + + >>> df.select_dtypes(exclude=['int']) + b c d + 0 True 1.0 a + 1 False 2.0 b + 2 True 1.0 a + 3 False 2.0 b + 4 True 1.0 a + 5 False 2.0 b + + Spark SQL DDL type strings can be used as well. + + >>> df.select_dtypes(exclude=['string']) + a b c + 0 1 True 1.0 + 1 2 False 2.0 + 2 1 True 1.0 + 3 2 False 2.0 + 4 1 True 1.0 + 5 2 False 2.0 + """ + from pyspark.sql.types import _parse_datatype_string + + include_list: List[str] + if not is_list_like(include): + include_list = [cast(str, include)] if include is not None else [] + else: + include_list = list(include) + exclude_list: List[str] + if not is_list_like(exclude): + exclude_list = [cast(str, exclude)] if exclude is not None else [] + else: + exclude_list = list(exclude) + + if not any((include_list, exclude_list)): + raise ValueError("at least one of include or exclude must be " "nonempty") + + # can't both include AND exclude! + if set(include_list).intersection(set(exclude_list)): + raise ValueError( + "include and exclude overlap on {inc_ex}".format( + inc_ex=set(include_list).intersection(set(exclude_list)) + ) + ) + + # Handle Spark types + include_spark_type = [] + for inc in include_list: + try: + include_spark_type.append(_parse_datatype_string(inc)) + except BaseException: + pass + + exclude_spark_type = [] + for exc in exclude_list: + try: + exclude_spark_type.append(_parse_datatype_string(exc)) + except BaseException: + pass + + # Handle pandas types + include_numpy_type = [] + for inc in include_list: + try: + include_numpy_type.append(infer_dtype_from_object(inc)) + except BaseException: + pass + + exclude_numpy_type = [] + for exc in exclude_list: + try: + exclude_numpy_type.append(infer_dtype_from_object(exc)) + except BaseException: + pass + + column_labels = [] + for label in self._internal.column_labels: + if len(include_list) > 0: + should_include = ( + infer_dtype_from_object(self._psser_for(label).dtype.name) in include_numpy_type + or self._internal.spark_type_for(label) in include_spark_type + ) + else: + should_include = not ( + infer_dtype_from_object(self._psser_for(label).dtype.name) in exclude_numpy_type + or self._internal.spark_type_for(label) in exclude_spark_type + ) + + if should_include: + column_labels.append(label) + + return DataFrame( + self._internal.with_new_columns([self._psser_for(label) for label in column_labels]) + ) + + def droplevel( + self, level: Union[int, Name, List[Union[int, Name]]], axis: Axis = 0 + ) -> "DataFrame": + """ + Return DataFrame with requested index / column level(s) removed. + + Parameters + ---------- + level: int, str, or list-like + If a string is given, must be the name of a level If list-like, elements must + be names or positional indexes of levels. + + axis: {0 or ‘index’, 1 or ‘columns’}, default 0 + + Returns + ------- + DataFrame with requested index / column level(s) removed. + + Examples + -------- + >>> df = ps.DataFrame( + ... [[3, 4], [7, 8], [11, 12]], + ... index=pd.MultiIndex.from_tuples([(1, 2), (5, 6), (9, 10)], names=["a", "b"]), + ... ) + + >>> df.columns = pd.MultiIndex.from_tuples([ + ... ('c', 'e'), ('d', 'f') + ... ], names=['level_1', 'level_2']) + + >>> df # doctest: +NORMALIZE_WHITESPACE + level_1 c d + level_2 e f + a b + 1 2 3 4 + 5 6 7 8 + 9 10 11 12 + + >>> df.droplevel('a') # doctest: +NORMALIZE_WHITESPACE + level_1 c d + level_2 e f + b + 2 3 4 + 6 7 8 + 10 11 12 + + >>> df.droplevel('level_2', axis=1) # doctest: +NORMALIZE_WHITESPACE + level_1 c d + a b + 1 2 3 4 + 5 6 7 8 + 9 10 11 12 + """ + axis = validate_axis(axis) + if axis == 0: + if not isinstance(level, (tuple, list)): # huh? + level = [level] + + names = self.index.names + nlevels = self._internal.index_level + + int_level = set() + for n in level: + if isinstance(n, int): + if n < 0: + n = n + nlevels + if n < 0: + raise IndexError( + "Too many levels: Index has only {} levels, " + "{} is not a valid level number".format(nlevels, (n - nlevels)) + ) + if n >= nlevels: + raise IndexError( + "Too many levels: Index has only {} levels, not {}".format( + nlevels, (n + 1) + ) + ) + else: + if n not in names: + raise KeyError("Level {} not found".format(n)) + n = names.index(n) + int_level.add(n) + + if len(level) >= nlevels: + raise ValueError( + "Cannot remove {} levels from an index with {} levels: " + "at least one level must be left.".format(len(level), nlevels) + ) + + index_spark_columns, index_names, index_fields = zip( + *[ + item + for i, item in enumerate( + zip( + self._internal.index_spark_columns, + self._internal.index_names, + self._internal.index_fields, + ) + ) + if i not in int_level + ] + ) + + internal = self._internal.copy( + index_spark_columns=list(index_spark_columns), + index_names=list(index_names), + index_fields=list(index_fields), + ) + return DataFrame(internal) + else: + psdf = self.copy() + psdf.columns = psdf.columns.droplevel(level) # type: ignore[arg-type] + return psdf + + def drop( + self, + labels: Optional[Union[Name, List[Name]]] = None, + axis: Optional[Axis] = 0, + index: Union[Name, List[Name]] = None, + columns: Union[Name, List[Name]] = None, + ) -> "DataFrame": + """ + Drop specified labels from columns. + + Remove rows and/or columns by specifying label names and corresponding axis, + or by specifying directly index and/or column names. + Drop rows of a MultiIndex DataFrame is not supported yet. + + Parameters + ---------- + labels : single label or list-like + Column labels to drop. + axis : {0 or 'index', 1 or 'columns'}, default 0 + + .. versionchanged:: 3.3 + Set dropping by index is default. + index : single label or list-like + Alternative to specifying axis (``labels, axis=0`` + is equivalent to ``index=columns``). + + .. versionchanged:: 3.3 + Added dropping rows by 'index'. + columns : single label or list-like + Alternative to specifying axis (``labels, axis=1`` + is equivalent to ``columns=labels``). + + Returns + ------- + dropped : DataFrame + + See Also + -------- + Series.dropna + + Examples + -------- + >>> df = ps.DataFrame(np.arange(12).reshape(3, 4), columns=['A', 'B', 'C', 'D']) + >>> df + A B C D + 0 0 1 2 3 + 1 4 5 6 7 + 2 8 9 10 11 + + Drop columns + + >>> df.drop(['B', 'C'], axis=1) + A D + 0 0 3 + 1 4 7 + 2 8 11 + + >>> df.drop(columns=['B', 'C']) + A D + 0 0 3 + 1 4 7 + 2 8 11 + + Drop a row by index + + >>> df.drop([0, 1]) + A B C D + 2 8 9 10 11 + + >>> df.drop(index=[0, 1], columns='A') + B C D + 2 9 10 11 + + Also support dropping columns for MultiIndex + + >>> df = ps.DataFrame({'x': [1, 2], 'y': [3, 4], 'z': [5, 6], 'w': [7, 8]}, + ... columns=['x', 'y', 'z', 'w']) + >>> columns = [('a', 'x'), ('a', 'y'), ('b', 'z'), ('b', 'w')] + >>> df.columns = pd.MultiIndex.from_tuples(columns) + >>> df # doctest: +NORMALIZE_WHITESPACE + a b + x y z w + 0 1 3 5 7 + 1 2 4 6 8 + >>> df.drop(labels='a', axis=1) # doctest: +NORMALIZE_WHITESPACE + b + z w + 0 5 7 + 1 6 8 + + Notes + ----- + Currently, dropping rows of a MultiIndex DataFrame is not supported yet. + """ + if labels is not None: + if index is not None or columns is not None: + raise ValueError("Cannot specify both 'labels' and 'index'/'columns'") + axis = validate_axis(axis) + if axis == 1: + return self.drop(index=index, columns=labels) + else: + return self.drop(index=labels, columns=columns) + else: + if index is None and columns is None: + raise ValueError("Need to specify at least one of 'labels' or 'columns' or 'index'") + + internal = self._internal + if index is not None: + if is_name_like_tuple(index) or is_name_like_value(index): + index = [index] + + if len(index) > 0: + if internal.index_level == 1: + internal = internal.resolved_copy + + if len(index) <= ps.get_option("compute.isin_limit"): + self_index_type = self.index.spark.data_type + cond = ~internal.index_spark_columns[0].isin( + [F.lit(label).cast(self_index_type) for label in index] + ) + internal = internal.with_filter(cond) + else: + index_sdf_col = "__index" + index_sdf = default_session().createDataFrame( + pd.DataFrame({index_sdf_col: index}) + ) + joined_sdf = internal.spark_frame.join( + other=F.broadcast(index_sdf), + on=( + internal.index_spark_columns[0] + == scol_for(index_sdf, index_sdf_col) + ), + how="anti", + ) + internal = internal.with_new_sdf(joined_sdf) + else: + raise NotImplementedError( + "Drop rows of MultiIndex DataFrame is not supported yet" + ) + if columns is not None: + if is_name_like_tuple(columns): + columns = [columns] + elif is_name_like_value(columns): + columns = [(columns,)] + else: + columns = [col if is_name_like_tuple(col) else (col,) for col in columns] + + if len(columns) > 0: + drop_column_labels = set( + label + for label in internal.column_labels + for col in columns + if label[: len(col)] == col + ) + if len(drop_column_labels) == 0: + raise KeyError(columns) + + keep_columns_and_labels = [ + (column, label) + for column, label in zip( + self._internal.data_spark_column_names, self._internal.column_labels + ) + if label not in drop_column_labels + ] + + cols, labels = ( + zip(*keep_columns_and_labels) + if len(keep_columns_and_labels) > 0 + else ([], []) + ) + internal = internal.with_new_columns( + [self._psser_for(label) for label in labels] + ) + return DataFrame(internal) + + def _prepare_sort_by_scols(self, by: Union[Name, List[Name]]) -> List[PySparkColumn]: + if is_name_like_value(by): + by = [by] + else: + assert is_list_like(by), type(by) + new_by = [] + for colname in by: + ser = self[colname] + if not isinstance(ser, ps.Series): + raise ValueError( + "The column %s is not unique. For a multi-index, the label must be a tuple " + "with elements corresponding to each level." % name_like_string(colname) + ) + new_by.append(ser.spark.column) + return new_by + + def _sort( + self, + by: Sequence[PySparkColumn], + ascending: Union[bool, List[bool]], + na_position: str, + keep: str = "first", + ) -> "DataFrame": + if isinstance(ascending, bool): + ascending = [ascending] * len(by) + if len(ascending) != len(by): + raise ValueError( + "Length of ascending ({}) != length of by ({})".format(len(ascending), len(by)) + ) + if na_position not in ("first", "last"): + raise ValueError("invalid na_position: '{}'".format(na_position)) + + Column = get_column_class() + # Mapper: Get a spark colum + # n function for (ascending, na_position) combination + mapper = { + (True, "first"): Column.asc_nulls_first, + (True, "last"): Column.asc_nulls_last, + (False, "first"): Column.desc_nulls_first, + (False, "last"): Column.desc_nulls_last, + } + by = [mapper[(asc, na_position)](scol) for scol, asc in zip(by, ascending)] + + natural_order_scol = F.col(NATURAL_ORDER_COLUMN_NAME) + + if keep == "last": + natural_order_scol = Column.desc(natural_order_scol) + elif keep == "all": + raise NotImplementedError("`keep`=all is not implemented yet.") + elif keep != "first": + raise ValueError('keep must be either "first", "last" or "all".') + sdf = self._internal.resolved_copy.spark_frame.sort(*by, natural_order_scol) + return DataFrame(self._internal.with_new_sdf(sdf)) + + def sort_values( + self, + by: Union[Name, List[Name]], + ascending: Union[bool, List[bool]] = True, + inplace: bool = False, + na_position: str = "last", + ignore_index: bool = False, + ) -> Optional["DataFrame"]: + """ + Sort by the values along either axis. + + Parameters + ---------- + by : str or list of str + ascending : bool or list of bool, default True + Sort ascending vs. descending. Specify list for multiple sort + orders. If this is a list of bools, must match the length of + the by. + inplace : bool, default False + if True, perform operation in-place + na_position : {'first', 'last'}, default 'last' + `first` puts NaNs at the beginning, `last` puts NaNs at the end + ignore_index : bool, default False + If True, the resulting axis will be labeled 0, 1, …, n - 1. + + Returns + ------- + sorted_obj : DataFrame + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'col1': ['A', 'B', None, 'D', 'C'], + ... 'col2': [2, 9, 8, 7, 4], + ... 'col3': [0, 9, 4, 2, 3], + ... }, + ... columns=['col1', 'col2', 'col3'], + ... index=['a', 'b', 'c', 'd', 'e']) + >>> df + col1 col2 col3 + a A 2 0 + b B 9 9 + c None 8 4 + d D 7 2 + e C 4 3 + + Sort by col1 + + >>> df.sort_values(by=['col1']) + col1 col2 col3 + a A 2 0 + b B 9 9 + e C 4 3 + d D 7 2 + c None 8 4 + + Ignore index for the resulting axis + + >>> df.sort_values(by=['col1'], ignore_index=True) + col1 col2 col3 + 0 A 2 0 + 1 B 9 9 + 2 C 4 3 + 3 D 7 2 + 4 None 8 4 + + Sort Descending + + >>> df.sort_values(by='col1', ascending=False) + col1 col2 col3 + d D 7 2 + e C 4 3 + b B 9 9 + a A 2 0 + c None 8 4 + + Sort by multiple columns + + >>> df = ps.DataFrame({ + ... 'col1': ['A', 'A', 'B', None, 'D', 'C'], + ... 'col2': [2, 1, 9, 8, 7, 4], + ... 'col3': [0, 1, 9, 4, 2, 3], + ... }, + ... columns=['col1', 'col2', 'col3']) + >>> df.sort_values(by=['col1', 'col2']) + col1 col2 col3 + 1 A 1 1 + 0 A 2 0 + 2 B 9 9 + 5 C 4 3 + 4 D 7 2 + 3 None 8 4 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + new_by = self._prepare_sort_by_scols(by) + + psdf = self._sort(by=new_by, ascending=ascending, na_position=na_position) + + if inplace: + if ignore_index: + psdf.reset_index(drop=True, inplace=inplace) + self._update_internal_frame(psdf._internal) + return None + else: + return psdf.reset_index(drop=True) if ignore_index else psdf + + def sort_index( + self, + axis: Axis = 0, + level: Optional[Union[int, List[int]]] = None, + ascending: bool = True, + inplace: bool = False, + kind: str = None, + na_position: str = "last", + ignore_index: bool = False, + ) -> Optional["DataFrame"]: + """ + Sort object by labels (along an axis) + + Parameters + ---------- + axis : index, columns to direct sorting. Currently, only axis = 0 is supported. + level : int or level name or list of ints or list of level names + if not None, sort on values in specified index level(s) + ascending : boolean, default True + Sort ascending vs. descending + inplace : bool, default False + if True, perform operation in-place + kind : str, default None + pandas-on-Spark does not allow specifying the sorting algorithm now, + default None + na_position : {‘first’, ‘last’}, default ‘last’ + first puts NaNs at the beginning, last puts NaNs at the end. Not implemented for + MultiIndex. + ignore_index : bool, default False + If True, the resulting axis will be labeled 0, 1, …, n - 1. + + .. versionadded:: 3.4.0 + + Returns + ------- + sorted_obj : DataFrame + + Examples + -------- + >>> df = ps.DataFrame({'A': [2, 1, np.nan]}, index=['b', 'a', np.nan]) + + >>> df.sort_index() # doctest: +SKIP + A + a 1.0 + b 2.0 + None NaN + + >>> df.sort_index(ascending=False) # doctest: +SKIP + A + b 2.0 + a 1.0 + None NaN + + >>> df.sort_index(na_position='first') # doctest: +SKIP + A + None NaN + a 1.0 + b 2.0 + + >>> df.sort_index(ignore_index=True) + A + 0 1.0 + 1 2.0 + 2 NaN + + >>> df.sort_index(inplace=True) + >>> df # doctest: +SKIP + A + a 1.0 + b 2.0 + None NaN + + >>> df = ps.DataFrame({'A': range(4), 'B': range(4)[::-1]}, + ... index=[['b', 'b', 'a', 'a'], [1, 0, 1, 0]], + ... columns=['A', 'B']) + + >>> df.sort_index() + A B + a 0 3 0 + 1 2 1 + b 0 1 2 + 1 0 3 + + >>> df.sort_index(level=1) + A B + b 0 1 2 + a 0 3 0 + b 1 0 3 + a 1 2 1 + + >>> df.sort_index(level=[1, 0]) + A B + a 0 3 0 + b 0 1 2 + a 1 2 1 + b 1 0 3 + + >>> df.sort_index(ignore_index=True) + A B + 0 3 0 + 1 2 1 + 2 1 2 + 3 0 3 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError("No other axis than 0 are supported now") + if kind is not None: + raise NotImplementedError("Specifying the sorting algorithm is not supported now.") + + if level is None or (is_list_like(level) and len(level) == 0): # type: ignore[arg-type] + by = self._internal.index_spark_columns + elif is_list_like(level): + by = [ + self._internal.index_spark_columns[lvl] for lvl in level # type: ignore[union-attr] + ] + else: + by = [self._internal.index_spark_columns[level]] # type: ignore[index] + + psdf = self._sort(by=by, ascending=ascending, na_position=na_position) + if inplace: + if ignore_index: + psdf.reset_index(drop=True, inplace=inplace) + self._update_internal_frame(psdf._internal) + return None + else: + return psdf.reset_index(drop=True) if ignore_index else psdf + + def swaplevel( + self, i: Union[int, Name] = -2, j: Union[int, Name] = -1, axis: Axis = 0 + ) -> "DataFrame": + """ + Swap levels i and j in a MultiIndex on a particular axis. + + Parameters + ---------- + i, j : int or str + Levels of the indices to be swapped. Can pass level name as string. + axis : {0 or 'index', 1 or 'columns'}, default 0 + The axis to swap levels on. 0 or 'index' for row-wise, 1 or + 'columns' for column-wise. + + Returns + ------- + DataFrame + DataFrame with levels swapped in MultiIndex. + + Examples + -------- + >>> midx = pd.MultiIndex.from_arrays( + ... [['red', 'blue'], [1, 2], ['s', 'm']], names = ['color', 'number', 'size']) + >>> midx # doctest: +SKIP + MultiIndex([( 'red', 1, 's'), + ('blue', 2, 'm')], + names=['color', 'number', 'size']) + + Swap levels in a MultiIndex on index. + + >>> psdf = ps.DataFrame({'x': [5, 6], 'y':[5, 6]}, index=midx) + >>> psdf # doctest: +NORMALIZE_WHITESPACE + x y + color number size + red 1 s 5 5 + blue 2 m 6 6 + + >>> psdf.swaplevel() # doctest: +NORMALIZE_WHITESPACE + x y + color size number + red s 1 5 5 + blue m 2 6 6 + + >>> psdf.swaplevel(0, 1) # doctest: +NORMALIZE_WHITESPACE + x y + number color size + 1 red s 5 5 + 2 blue m 6 6 + + >>> psdf.swaplevel('number', 'size') # doctest: +NORMALIZE_WHITESPACE + x y + color size number + red s 1 5 5 + blue m 2 6 6 + + Swap levels in a MultiIndex on columns. + + >>> psdf = ps.DataFrame({'x': [5, 6], 'y':[5, 6]}) + >>> psdf.columns = midx + >>> psdf + color red blue + number 1 2 + size s m + 0 5 5 + 1 6 6 + + >>> psdf.swaplevel(axis=1) + color red blue + size s m + number 1 2 + 0 5 5 + 1 6 6 + + >>> psdf.swaplevel(axis=1) + color red blue + size s m + number 1 2 + 0 5 5 + 1 6 6 + + >>> psdf.swaplevel(0, 1, axis=1) + number 1 2 + color red blue + size s m + 0 5 5 + 1 6 6 + + >>> psdf.swaplevel('number', 'color', axis=1) + number 1 2 + color red blue + size s m + 0 5 5 + 1 6 6 + """ + axis = validate_axis(axis) + if axis == 0: + internal = self._swaplevel_index(i, j) + else: + assert axis == 1 + internal = self._swaplevel_columns(i, j) + + return DataFrame(internal) + + def swapaxes(self, i: Axis, j: Axis, copy: bool = True) -> "DataFrame": + """ + Interchange axes and swap values axes appropriately. + + .. note:: This method is based on an expensive operation due to the nature + of big data. Internally it needs to generate each row for each value, and + then group twice - it is a huge operation. To prevent misuse, this method + has the 'compute.max_rows' default limit of input length and raises a ValueError. + + >>> from pyspark.pandas.config import option_context + >>> with option_context('compute.max_rows', 1000): # doctest: +NORMALIZE_WHITESPACE + ... ps.DataFrame({'a': range(1001)}).swapaxes(i=0, j=1) + Traceback (most recent call last): + ... + ValueError: Current DataFrame's length exceeds the given limit of 1000 rows. + Please set 'compute.max_rows' by using 'pyspark.pandas.config.set_option' + to retrieve more than 1000 rows. Note that, before changing the + 'compute.max_rows', this operation is considerably expensive. + + Parameters + ---------- + i: {0 or 'index', 1 or 'columns'}. The axis to swap. + j: {0 or 'index', 1 or 'columns'}. The axis to swap. + copy : bool, default True. + + Returns + ------- + DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame( + ... [[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=['x', 'y', 'z'], columns=['a', 'b', 'c'] + ... ) + >>> psdf + a b c + x 1 2 3 + y 4 5 6 + z 7 8 9 + >>> psdf.swapaxes(i=1, j=0) + x y z + a 1 4 7 + b 2 5 8 + c 3 6 9 + >>> psdf.swapaxes(i=1, j=1) + a b c + x 1 2 3 + y 4 5 6 + z 7 8 9 + """ + assert copy is True + + i = validate_axis(i) + j = validate_axis(j) + + return self.copy() if i == j else self.transpose() + + def _swaplevel_columns(self, i: Union[int, Name], j: Union[int, Name]) -> InternalFrame: + assert isinstance(self.columns, pd.MultiIndex) + for index in (i, j): + if not isinstance(index, int) and index not in self.columns.names: + raise KeyError("Level %s not found" % index) + + i = i if isinstance(i, int) else self.columns.names.index(i) + j = j if isinstance(j, int) else self.columns.names.index(j) + for index in (i, j): + if index >= len(self.columns) or index < -len(self.columns): + raise IndexError( + "Too many levels: Columns have only %s levels, " + "%s is not a valid level number" % (self._internal.index_level, index) + ) + + column_label_names = self._internal.column_label_names.copy() + column_label_names[i], column_label_names[j], = ( + column_label_names[j], + column_label_names[i], + ) + column_labels = self._internal._column_labels + column_label_list = [list(label) for label in column_labels] + for label_list in column_label_list: + label_list[i], label_list[j] = label_list[j], label_list[i] + column_labels = [tuple(x) for x in column_label_list] + internal = self._internal.copy( + column_label_names=list(column_label_names), column_labels=list(column_labels) + ) + return internal + + def _swaplevel_index(self, i: Union[int, Name], j: Union[int, Name]) -> InternalFrame: + assert isinstance(self.index, ps.MultiIndex) + for index in (i, j): + if not isinstance(index, int) and index not in self.index.names: + raise KeyError("Level %s not found" % index) + + i = i if isinstance(i, int) else self.index.names.index(i) + j = j if isinstance(j, int) else self.index.names.index(j) + for index in (i, j): + if index >= self._internal.index_level or index < -self._internal.index_level: + raise IndexError( + "Too many levels: Index has only %s levels, " + "%s is not a valid level number" % (self._internal.index_level, index) + ) + + index_map = list( + zip( + self._internal.index_spark_columns, + self._internal.index_names, + self._internal.index_fields, + ) + ) + index_map[i], index_map[j] = index_map[j], index_map[i] + index_spark_columns, index_names, index_fields = zip(*index_map) + internal = self._internal.copy( + index_spark_columns=list(index_spark_columns), + index_names=list(index_names), + index_fields=list(index_fields), + ) + return internal + + def nlargest( + self, n: int, columns: Union[Name, List[Name]], keep: str = "first" + ) -> "DataFrame": + """ + Return the first `n` rows ordered by `columns` in descending order. + + Return the first `n` rows with the largest values in `columns`, in + descending order. The columns that are not specified are returned as + well, but not used for ordering. + + This method is equivalent to + ``df.sort_values(columns, ascending=False).head(n)``, but more + performant in pandas. + In pandas-on-Spark, thanks to Spark's lazy execution and query optimizer, + the two would have same performance. + + Parameters + ---------- + n : int + Number of rows to return. + columns : label or list of labels + Column label(s) to order by. + keep : {'first', 'last'}, default 'first'. 'all' is not implemented yet. + Determines which duplicates (if any) to keep. + - ``first`` : Keep the first occurrence. + - ``last`` : Keep the last occurrence. + + Returns + ------- + DataFrame + The first `n` rows ordered by the given columns in descending + order. + + See Also + -------- + DataFrame.nsmallest : Return the first `n` rows ordered by `columns` in + ascending order. + DataFrame.sort_values : Sort DataFrame by the values. + DataFrame.head : Return the first `n` rows without re-ordering. + + Notes + ----- + + This function cannot be used with all column types. For example, when + specifying columns with `object` or `category` dtypes, ``TypeError`` is + raised. + + Examples + -------- + >>> df = ps.DataFrame({'X': [1, 2, 3, 5, 6, 7, np.nan], + ... 'Y': [6, 7, 8, 9, 10, 11, 12]}) + >>> df + X Y + 0 1.0 6 + 1 2.0 7 + 2 3.0 8 + 3 5.0 9 + 4 6.0 10 + 5 7.0 11 + 6 NaN 12 + + In the following example, we will use ``nlargest`` to select the three + rows having the largest values in column "X". + + >>> df.nlargest(n=3, columns='X') + X Y + 5 7.0 11 + 4 6.0 10 + 3 5.0 9 + + To order by the largest values in column "Y" and then "X", we can + specify multiple columns like in the next example. + + >>> df.nlargest(n=3, columns=['Y', 'X']) + X Y + 6 NaN 12 + 5 7.0 11 + 4 6.0 10 + + The examples below show how ties are resolved, which is decided by `keep`. + + >>> tied_df = ps.DataFrame({'X': [1, 2, 2, 3, 3]}, index=['a', 'b', 'c', 'd', 'e']) + >>> tied_df + X + a 1 + b 2 + c 2 + d 3 + e 3 + + When using keep='first' (default), ties are resolved in order: + + >>> tied_df.nlargest(3, 'X') + X + d 3 + e 3 + b 2 + + >>> tied_df.nlargest(3, 'X', keep='first') + X + d 3 + e 3 + b 2 + + When using keep='last', ties are resolved in reverse order: + + >>> tied_df.nlargest(3, 'X', keep='last') + X + e 3 + d 3 + c 2 + """ + by_scols = self._prepare_sort_by_scols(columns) + return self._sort(by=by_scols, ascending=False, na_position="last", keep=keep).head(n=n) + + def nsmallest( + self, n: int, columns: Union[Name, List[Name]], keep: str = "first" + ) -> "DataFrame": + """ + Return the first `n` rows ordered by `columns` in ascending order. + + Return the first `n` rows with the smallest values in `columns`, in + ascending order. The columns that are not specified are returned as + well, but not used for ordering. + + This method is equivalent to ``df.sort_values(columns, ascending=True).head(n)``, + but more performant. In pandas-on-Spark, thanks to Spark's lazy execution and query + optimizer, the two would have same performance. + + Parameters + ---------- + n : int + Number of items to retrieve. + columns : list or str + Column name or names to order by. + keep : {'first', 'last'}, default 'first'. 'all' is not implemented yet. + Determines which duplicates (if any) to keep. + - ``first`` : Keep the first occurrence. + - ``last`` : Keep the last occurrence. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.nlargest : Return the first `n` rows ordered by `columns` in + descending order. + DataFrame.sort_values : Sort DataFrame by the values. + DataFrame.head : Return the first `n` rows without re-ordering. + + Examples + -------- + >>> df = ps.DataFrame({'X': [1, 2, 3, 5, 6, 7, np.nan], + ... 'Y': [6, 7, 8, 9, 10, 11, 12]}) + >>> df + X Y + 0 1.0 6 + 1 2.0 7 + 2 3.0 8 + 3 5.0 9 + 4 6.0 10 + 5 7.0 11 + 6 NaN 12 + + In the following example, we will use ``nsmallest`` to select the + three rows having the smallest values in column "X". + + >>> df.nsmallest(n=3, columns='X') # doctest: +NORMALIZE_WHITESPACE + X Y + 0 1.0 6 + 1 2.0 7 + 2 3.0 8 + + To order by the smallest values in column "Y" and then "X", we can + specify multiple columns like in the next example. + + >>> df.nsmallest(n=3, columns=['Y', 'X']) # doctest: +NORMALIZE_WHITESPACE + X Y + 0 1.0 6 + 1 2.0 7 + 2 3.0 8 + + The examples below show how ties are resolved, which is decided by `keep`. + + >>> tied_df = ps.DataFrame({'X': [1, 1, 2, 2, 3]}, index=['a', 'b', 'c', 'd', 'e']) + >>> tied_df + X + a 1 + b 1 + c 2 + d 2 + e 3 + + When using keep='first' (default), ties are resolved in order: + + >>> tied_df.nsmallest(3, 'X') + X + a 1 + b 1 + c 2 + + >>> tied_df.nsmallest(3, 'X', keep='first') + X + a 1 + b 1 + c 2 + + When using keep='last', ties are resolved in reverse order: + + >>> tied_df.nsmallest(3, 'X', keep='last') + X + b 1 + a 1 + d 2 + """ + by_scols = self._prepare_sort_by_scols(columns) + return self._sort(by=by_scols, ascending=True, na_position="last", keep=keep).head(n=n) + + def isin(self, values: Union[List, Dict]) -> "DataFrame": + """ + Whether each element in the DataFrame is contained in values. + + Parameters + ---------- + values : iterable or dict + The sequence of values to test. If values are a dict, + the keys must be the column names, which must match. + Series and DataFrame are not supported. + + Returns + ------- + DataFrame + DataFrame of booleans showing whether each element in the DataFrame + is contained in values. + + Examples + -------- + >>> df = ps.DataFrame({'num_legs': [2, 4], 'num_wings': [2, 0]}, + ... index=['falcon', 'dog'], + ... columns=['num_legs', 'num_wings']) + >>> df + num_legs num_wings + falcon 2 2 + dog 4 0 + + When ``values`` is a list check whether every value in the DataFrame + is present in the list (which animals have 0 or 2 legs or wings) + + >>> df.isin([0, 2]) + num_legs num_wings + falcon True True + dog False True + + When ``values`` is a dict, we can pass values to check for each + column separately: + + >>> df.isin({'num_wings': [0, 3]}) + num_legs num_wings + falcon False False + dog False True + """ + if isinstance(values, (pd.DataFrame, pd.Series)): + raise NotImplementedError("DataFrame and Series are not supported") + if isinstance(values, dict) and not set(values.keys()).issubset(self.columns): + raise AttributeError( + "'DataFrame' object has no attribute %s" + % (set(values.keys()).difference(self.columns)) + ) + + data_spark_columns = [] + if isinstance(values, dict): + for i, col in enumerate(self.columns): + if col in values: + item = values[col] + item = item.tolist() if isinstance(item, np.ndarray) else list(item) + + scol = self._internal.spark_column_for(self._internal.column_labels[i]).isin( + [F.lit(v) for v in item] + ) + scol = F.coalesce(scol, F.lit(False)) + else: + scol = F.lit(False) + data_spark_columns.append(scol.alias(self._internal.data_spark_column_names[i])) + elif is_list_like(values): + values = ( + cast(np.ndarray, values).tolist() + if isinstance(values, np.ndarray) + else list(values) + ) + + for label in self._internal.column_labels: + scol = self._internal.spark_column_for(label).isin([F.lit(v) for v in values]) + scol = F.coalesce(scol, F.lit(False)) + data_spark_columns.append(scol.alias(self._internal.spark_column_name_for(label))) + else: + raise TypeError("Values should be iterable, Series, DataFrame or dict.") + + return DataFrame( + self._internal.with_new_columns( + data_spark_columns, + data_fields=[ + field.copy(dtype=np.dtype("bool"), spark_type=BooleanType(), nullable=False) + for field in self._internal.data_fields + ], + ) + ) + + @property + def shape(self) -> Tuple[int, int]: + """ + Return a tuple representing the dimensionality of the DataFrame. + + Examples + -------- + >>> df = ps.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) + >>> df.shape + (2, 2) + + >>> df = ps.DataFrame({'col1': [1, 2], 'col2': [3, 4], + ... 'col3': [5, 6]}) + >>> df.shape + (2, 3) + """ + return len(self), len(self.columns) + + def merge( + self, + right: "DataFrame", + how: str = "inner", + on: Optional[Union[Name, List[Name]]] = None, + left_on: Optional[Union[Name, List[Name]]] = None, + right_on: Optional[Union[Name, List[Name]]] = None, + left_index: bool = False, + right_index: bool = False, + suffixes: Tuple[str, str] = ("_x", "_y"), + ) -> "DataFrame": + """ + Merge DataFrame objects with a database-style join. + + The index of the resulting DataFrame will be one of the following: + - 0...n if no index is used for merging + - Index of the left DataFrame if merged only on the index of the right DataFrame + - Index of the right DataFrame if merged only on the index of the left DataFrame + - All involved indices if merged using the indices of both DataFrames + e.g. if `left` with indices (a, x) and `right` with indices (b, x), the result will + be an index (x, a, b) + + Parameters + ---------- + right: Object to merge with. + how: Type of merge to be performed. + {'left', 'right', 'outer', 'inner'}, default 'inner' + + left: use only keys from left frame, like a SQL left outer join; not preserve + key order unlike pandas. + right: use only keys from right frame, like a SQL right outer join; not preserve + key order unlike pandas. + outer: use union of keys from both frames, like a SQL full outer join; sort keys + lexicographically. + inner: use intersection of keys from both frames, like a SQL inner join; + not preserve the order of the left keys unlike pandas. + on: Column or index level names to join on. These must be found in both DataFrames. If on + is None and not merging on indexes then this defaults to the intersection of the + columns in both DataFrames. + left_on: Column or index level names to join on in the left DataFrame. Can also + be an array or list of arrays of the length of the left DataFrame. + These arrays are treated as if they are columns. + right_on: Column or index level names to join on in the right DataFrame. Can also + be an array or list of arrays of the length of the right DataFrame. + These arrays are treated as if they are columns. + left_index: Use the index from the left DataFrame as the join key(s). If it is a + MultiIndex, the number of keys in the other DataFrame (either the index or a number of + columns) must match the number of levels. + right_index: Use the index from the right DataFrame as the join key. Same caveats as + left_index. + suffixes: Suffix to apply to overlapping column names in the left and right side, + respectively. + + Returns + ------- + DataFrame + A DataFrame of the two merged objects. + + See Also + -------- + DataFrame.join : Join columns of another DataFrame. + DataFrame.update : Modify in place using non-NA values from another DataFrame. + DataFrame.hint : Specifies some hint on the current DataFrame. + broadcast : Marks a DataFrame as small enough for use in broadcast joins. + + Examples + -------- + >>> df1 = ps.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [1, 2, 3, 5]}, + ... columns=['lkey', 'value']) + >>> df2 = ps.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [5, 6, 7, 8]}, + ... columns=['rkey', 'value']) + >>> df1 + lkey value + 0 foo 1 + 1 bar 2 + 2 baz 3 + 3 foo 5 + >>> df2 + rkey value + 0 foo 5 + 1 bar 6 + 2 baz 7 + 3 foo 8 + + Merge df1 and df2 on the lkey and rkey columns. The value columns have + the default suffixes, _x and _y, appended. + + >>> merged = df1.merge(df2, left_on='lkey', right_on='rkey') + >>> merged.sort_values(by=['lkey', 'value_x', 'rkey', 'value_y']) # doctest: +ELLIPSIS + lkey value_x rkey value_y + ...bar 2 bar 6 + ...baz 3 baz 7 + ...foo 1 foo 5 + ...foo 1 foo 8 + ...foo 5 foo 5 + ...foo 5 foo 8 + + >>> left_psdf = ps.DataFrame({'A': [1, 2]}) + >>> right_psdf = ps.DataFrame({'B': ['x', 'y']}, index=[1, 2]) + + >>> left_psdf.merge(right_psdf, left_index=True, right_index=True).sort_index() + A B + 1 2 x + + >>> left_psdf.merge(right_psdf, left_index=True, right_index=True, how='left').sort_index() + A B + 0 1 None + 1 2 x + + >>> left_psdf.merge(right_psdf, left_index=True, right_index=True, how='right').sort_index() + A B + 1 2.0 x + 2 NaN y + + >>> left_psdf.merge(right_psdf, left_index=True, right_index=True, how='outer').sort_index() + A B + 0 1.0 None + 1 2.0 x + 2 NaN y + + Notes + ----- + As described in #263, joining string columns currently returns None for missing values + instead of NaN. + """ + + def to_list(os: Optional[Union[Name, List[Name]]]) -> List[Label]: + if os is None: + return [] + elif is_name_like_tuple(os): + return [cast(Label, os)] + elif is_name_like_value(os): + return [(os,)] + else: + return [o if is_name_like_tuple(o) else (o,) for o in os] + + if isinstance(right, ps.Series): + right = right.to_frame() + + if on: + if left_on or right_on: + raise ValueError( + 'Can only pass argument "on" OR "left_on" and "right_on", ' + "not a combination of both." + ) + left_key_names = list(map(self._internal.spark_column_name_for, to_list(on))) + right_key_names = list(map(right._internal.spark_column_name_for, to_list(on))) + else: + # TODO: need special handling for multi-index. + if left_index: + left_key_names = self._internal.index_spark_column_names + else: + left_key_names = list(map(self._internal.spark_column_name_for, to_list(left_on))) + if right_index: + right_key_names = right._internal.index_spark_column_names + else: + right_key_names = list( + map(right._internal.spark_column_name_for, to_list(right_on)) + ) + + if left_key_names and not right_key_names: + raise ValueError("Must pass right_on or right_index=True") + if right_key_names and not left_key_names: + raise ValueError("Must pass left_on or left_index=True") + if not left_key_names and not right_key_names: + common = list(self.columns.intersection(right.columns)) + if len(common) == 0: + raise ValueError( + "No common columns to perform merge on. Merge options: " + "left_on=None, right_on=None, left_index=False, right_index=False" + ) + left_key_names = list(map(self._internal.spark_column_name_for, to_list(common))) + right_key_names = list(map(right._internal.spark_column_name_for, to_list(common))) + if len(left_key_names) != len(right_key_names): + raise ValueError("len(left_keys) must equal len(right_keys)") + + # We should distinguish the name to avoid ambiguous column name after merging. + right_prefix = "__right_" + right_key_names = [right_prefix + right_key_name for right_key_name in right_key_names] + + how = validate_how(how) + + def resolve(internal: InternalFrame, side: str) -> InternalFrame: + def rename(col: str) -> str: + return "__{}_{}".format(side, col) + + internal = internal.resolved_copy + sdf = internal.spark_frame + sdf = sdf.select( + *[ + scol_for(sdf, col).alias(rename(col)) + for col in sdf.columns + if col not in HIDDEN_COLUMNS + ], + *HIDDEN_COLUMNS, + ) + return internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.index_spark_column_names + ], + index_fields=[ + field.copy(name=rename(field.name)) for field in internal.index_fields + ], + data_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.data_spark_column_names + ], + data_fields=[field.copy(name=rename(field.name)) for field in internal.data_fields], + ) + + left_internal = self._internal.resolved_copy + right_internal = resolve(right._internal, "right") + + left_table = left_internal.spark_frame.alias("left_table") + right_table = right_internal.spark_frame.alias("right_table") + + left_key_columns = [scol_for(left_table, label) for label in left_key_names] + right_key_columns = [scol_for(right_table, label) for label in right_key_names] + + join_condition = reduce( + lambda x, y: x & y, + [lkey == rkey for lkey, rkey in zip(left_key_columns, right_key_columns)], + ) + + joined_table = left_table.join(right_table, join_condition, how=how) + + # Unpack suffixes tuple for convenience + left_suffix = suffixes[0] + right_suffix = suffixes[1] + + # Append suffixes to columns with the same name to avoid conflicts later + duplicate_columns = set(left_internal.column_labels) & set(right_internal.column_labels) + + exprs = [] + data_columns = [] + column_labels = [] + + def left_scol_for(label: Label) -> PySparkColumn: + return scol_for(left_table, left_internal.spark_column_name_for(label)) + + def right_scol_for(label: Label) -> PySparkColumn: + return scol_for(right_table, right_internal.spark_column_name_for(label)) + + for label in left_internal.column_labels: + col = left_internal.spark_column_name_for(label) + scol = left_scol_for(label) + if label in duplicate_columns: + spark_column_name = left_internal.spark_column_name_for(label) + if ( + spark_column_name in left_key_names + and (right_prefix + spark_column_name) in right_key_names + ): + right_scol = right_scol_for(label) + if how == "right": + scol = right_scol.alias(col) + elif how == "full": + scol = F.when(scol.isNotNull(), scol).otherwise(right_scol).alias(col) + else: + pass + else: + col = col + left_suffix + scol = scol.alias(col) + label = tuple([str(label[0]) + left_suffix] + list(label[1:])) + exprs.append(scol) + data_columns.append(col) + column_labels.append(label) + for label in right_internal.column_labels: + # recover `right_prefix` here. + col = right_internal.spark_column_name_for(label)[len(right_prefix) :] + scol = right_scol_for(label).alias(col) + if label in duplicate_columns: + spark_column_name = left_internal.spark_column_name_for(label) + if ( + spark_column_name in left_key_names + and (right_prefix + spark_column_name) in right_key_names + ): + continue + else: + col = col + right_suffix + scol = scol.alias(col) + label = tuple([str(label[0]) + right_suffix] + list(label[1:])) + exprs.append(scol) + data_columns.append(col) + column_labels.append(label) + + left_index_scols = left_internal.index_spark_columns + right_index_scols = right_internal.index_spark_columns + + # Retain indices if they are used for joining + if left_index: + if right_index: + if how in ("inner", "left"): + exprs.extend(left_index_scols) + index_spark_column_names = left_internal.index_spark_column_names + index_names = left_internal.index_names + elif how == "right": + exprs.extend(right_index_scols) + index_spark_column_names = right_internal.index_spark_column_names + index_names = right_internal.index_names + else: + index_spark_column_names = left_internal.index_spark_column_names + index_names = left_internal.index_names + for col, left_scol, right_scol in zip( + index_spark_column_names, left_index_scols, right_index_scols + ): + scol = F.when(left_scol.isNotNull(), left_scol).otherwise(right_scol) + exprs.append(scol.alias(col)) + else: + exprs.extend(right_index_scols) + index_spark_column_names = right_internal.index_spark_column_names + index_names = right_internal.index_names + elif right_index: + exprs.extend(left_index_scols) + index_spark_column_names = left_internal.index_spark_column_names + index_names = left_internal.index_names + else: + index_spark_column_names = [] + index_names = [] + + selected_columns = joined_table.select(*exprs) + + internal = InternalFrame( + spark_frame=selected_columns, + index_spark_columns=[ + scol_for(selected_columns, col) for col in index_spark_column_names + ], + index_names=index_names, + column_labels=column_labels, + data_spark_columns=[scol_for(selected_columns, col) for col in data_columns], + ) + return DataFrame(internal) + + def join( + self, + right: "DataFrame", + on: Optional[Union[Name, List[Name]]] = None, + how: str = "left", + lsuffix: str = "", + rsuffix: str = "", + ) -> "DataFrame": + """ + Join columns of another DataFrame. + + Join columns with `right` DataFrame either on index or on a key column. Efficiently join + multiple DataFrame objects by index at once by passing a list. + + Parameters + ---------- + right: DataFrame, Series + on: str, list of str, or array-like, optional + Column or index level name(s) in the caller to join on the index in `right`, otherwise + joins index-on-index. If multiple values given, the `right` DataFrame must have a + MultiIndex. Can pass an array as the join key if it is not already contained in the + calling DataFrame. Like an Excel VLOOKUP operation. + how: {'left', 'right', 'outer', 'inner'}, default 'left' + How to handle the operation of the two objects. + + * left: use `left` frame’s index (or column if on is specified). + * right: use `right`’s index. + * outer: form union of `left` frame’s index (or column if on is specified) with + right’s index, and sort it. lexicographically. + * inner: form intersection of `left` frame’s index (or column if on is specified) + with `right`’s index, preserving the order of the `left`’s one. + lsuffix : str, default '' + Suffix to use from left frame's overlapping columns. + rsuffix : str, default '' + Suffix to use from `right` frame's overlapping columns. + + Returns + ------- + DataFrame + A dataframe containing columns from both the `left` and `right`. + + See Also + -------- + DataFrame.merge: For column(s)-on-columns(s) operations. + DataFrame.update : Modify in place using non-NA values from another DataFrame. + DataFrame.hint : Specifies some hint on the current DataFrame. + broadcast : Marks a DataFrame as small enough for use in broadcast joins. + + Notes + ----- + Parameters on, lsuffix, and rsuffix are not supported when passing a list of DataFrame + objects. + + Examples + -------- + >>> psdf1 = ps.DataFrame({'key': ['K0', 'K1', 'K2', 'K3'], + ... 'A': ['A0', 'A1', 'A2', 'A3']}, + ... columns=['key', 'A']) + >>> psdf2 = ps.DataFrame({'key': ['K0', 'K1', 'K2'], + ... 'B': ['B0', 'B1', 'B2']}, + ... columns=['key', 'B']) + >>> psdf1 + key A + 0 K0 A0 + 1 K1 A1 + 2 K2 A2 + 3 K3 A3 + >>> psdf2 + key B + 0 K0 B0 + 1 K1 B1 + 2 K2 B2 + + Join DataFrames using their indexes. + + >>> join_psdf = psdf1.join(psdf2, lsuffix='_left', rsuffix='_right') + >>> join_psdf.sort_values(by=join_psdf.columns) + key_left A key_right B + 0 K0 A0 K0 B0 + 1 K1 A1 K1 B1 + 2 K2 A2 K2 B2 + 3 K3 A3 None None + + If we want to join using the key columns, we need to set key to be the index in both df and + right. The joined DataFrame will have key as its index. + + >>> join_psdf = psdf1.set_index('key').join(psdf2.set_index('key')) + >>> join_psdf.sort_values(by=join_psdf.columns) # doctest: +NORMALIZE_WHITESPACE + A B + key + K0 A0 B0 + K1 A1 B1 + K2 A2 B2 + K3 A3 None + + Another option to join using the key columns is to use the on parameter. DataFrame.join + always uses right’s index but we can use any column in df. This method does not preserve + the original DataFrame’s index in the result unlike pandas. + + >>> join_psdf = psdf1.join(psdf2.set_index('key'), on='key') + >>> join_psdf.index # doctest: +SKIP + Int64Index([0, 1, 2, 3], dtype='int64') + """ + if isinstance(right, ps.Series): + common = list(self.columns.intersection([right.name])) + else: + common = list(self.columns.intersection(right.columns)) + if len(common) > 0 and not lsuffix and not rsuffix: + raise ValueError( + "columns overlap but no suffix specified: " "{rename}".format(rename=common) + ) + + need_set_index = False + if on: + if not is_list_like(on): + on = [on] + if len(on) != right._internal.index_level: + raise ValueError( + 'len(left_on) must equal the number of levels in the index of "right"' + ) + + need_set_index = len(set(on) & set(self.index.names)) == 0 + if need_set_index: + self = self.set_index(on) + join_psdf = self.merge( + right, left_index=True, right_index=True, how=how, suffixes=(lsuffix, rsuffix) + ) + return join_psdf.reset_index() if need_set_index else join_psdf + + def combine_first(self, other: "DataFrame") -> "DataFrame": + """ + Update null elements with value in the same location in `other`. + + Combine two DataFrame objects by filling null values in one DataFrame + with non-null values from other DataFrame. The row and column indexes + of the resulting DataFrame will be the union of the two. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + other : DataFrame + Provided DataFrame to use to fill null values. + + Returns + ------- + DataFrame + + Examples + -------- + >>> ps.set_option("compute.ops_on_diff_frames", True) + >>> df1 = ps.DataFrame({'A': [None, 0], 'B': [None, 4]}) + >>> df2 = ps.DataFrame({'A': [1, 1], 'B': [3, 3]}) + + >>> df1.combine_first(df2).sort_index() + A B + 0 1.0 3.0 + 1 0.0 4.0 + + Null values persist if the location of that null value does not exist in other + + >>> df1 = ps.DataFrame({'A': [None, 0], 'B': [4, None]}) + >>> df2 = ps.DataFrame({'B': [3, 3], 'C': [1, 1]}, index=[1, 2]) + + >>> df1.combine_first(df2).sort_index() + A B C + 0 NaN 4.0 NaN + 1 0.0 3.0 1.0 + 2 NaN 3.0 1.0 + >>> ps.reset_option("compute.ops_on_diff_frames") + """ + if not isinstance(other, DataFrame): + raise TypeError("`combine_first` only allows `DataFrame` for parameter `other`") + if same_anchor(self, other): + combined = self + this = self + that = other + else: + combined = combine_frames(self, other) + this = combined["this"] + that = combined["that"] + + intersect_column_labels = set(self._internal.column_labels).intersection( + set(other._internal.column_labels) + ) + + column_labels, data_spark_columns = [], [] + for column_label in this._internal.column_labels: + this_scol = this._internal.spark_column_for(column_label) + if column_label in intersect_column_labels: + that_scol = that._internal.spark_column_for(column_label) + this_scol_name = this._internal.spark_column_name_for(column_label) + combined_scol = ( + F.when(this_scol.isNull(), that_scol).otherwise(this_scol).alias(this_scol_name) + ) + data_spark_columns.append(combined_scol) + else: + data_spark_columns.append(this_scol) + column_labels.append(column_label) + + for column_label in that._internal.column_labels: + if column_label not in intersect_column_labels: + that_scol = that._internal.spark_column_for(column_label) + data_spark_columns.append(that_scol) + column_labels.append(column_label) + + internal = combined._internal.copy( + column_labels=column_labels, + data_spark_columns=data_spark_columns, + data_fields=None, # TODO: dtype? + column_label_names=self._internal.column_label_names, + ) + return DataFrame(internal) + + def append( + self, + other: "DataFrame", + ignore_index: bool = False, + verify_integrity: bool = False, + sort: bool = False, + ) -> "DataFrame": + """ + Append rows of other to the end of caller, returning a new object. + + Columns in other that are not in the caller are added as new columns. + + .. deprecated:: 3.4.0 + + Parameters + ---------- + other : DataFrame or Series/dict-like object, or list of these + The data to append. + + ignore_index : boolean, default False + If True, do not use the index labels. + + verify_integrity : boolean, default False + If True, raise ValueError on creating index with duplicates. + + sort : boolean, default False + Currently not supported. + + Returns + ------- + appended : DataFrame + + Examples + -------- + >>> df = ps.DataFrame([[1, 2], [3, 4]], columns=list('AB')) + + >>> df.append(df) + A B + 0 1 2 + 1 3 4 + 0 1 2 + 1 3 4 + + >>> df.append(df, ignore_index=True) + A B + 0 1 2 + 1 3 4 + 2 1 2 + 3 3 4 + """ + warnings.warn( + "The DataFrame.append method is deprecated " + "and will be removed in 4.0.0. " + "Use pyspark.pandas.concat instead.", + FutureWarning, + ) + if isinstance(other, ps.Series): + raise TypeError("DataFrames.append() does not support appending Series to DataFrames") + if sort: + raise NotImplementedError("The 'sort' parameter is currently not supported") + + if not ignore_index: + index_scols = self._internal.index_spark_columns + if len(index_scols) != other._internal.index_level: + raise ValueError("Both DataFrames have to have the same number of index levels") + + if ( + verify_integrity + and len(index_scols) > 0 + and ( + self._internal.spark_frame.select(index_scols) + .intersect( + other._internal.spark_frame.select(other._internal.index_spark_columns) + ) + .count() + ) + > 0 + ): + raise ValueError("Indices have overlapping values") + + # Lazy import to avoid circular dependency issues + from pyspark.pandas.namespace import concat + + return cast(DataFrame, concat([self, other], ignore_index=ignore_index)) + + # TODO: add 'filter_func' and 'errors' parameter + def update(self, other: "DataFrame", join: str = "left", overwrite: bool = True) -> None: + """ + Modify in place using non-NA values from another DataFrame. + Aligns on indices. There is no return value. + + Parameters + ---------- + other : DataFrame, or Series + join : 'left', default 'left' + Only left join is implemented, keeping the index and columns of the original object. + overwrite : bool, default True + How to handle non-NA values for overlapping keys: + + * True: overwrite original DataFrame's values with values from `other`. + * False: only update values that are NA in the original DataFrame. + + Returns + ------- + None : method directly changes calling object + + See Also + -------- + DataFrame.merge : For column(s)-on-columns(s) operations. + DataFrame.join : Join columns of another DataFrame. + DataFrame.hint : Specifies some hint on the current DataFrame. + broadcast : Marks a DataFrame as small enough for use in broadcast joins. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 3], 'B': [400, 500, 600]}, columns=['A', 'B']) + >>> new_df = ps.DataFrame({'B': [4, 5, 6], 'C': [7, 8, 9]}, columns=['B', 'C']) + >>> df.update(new_df) + >>> df.sort_index() + A B + 0 1 4 + 1 2 5 + 2 3 6 + + The DataFrame's length does not increase because of the update, + only values at matching index/column labels are updated. + + >>> df = ps.DataFrame({'A': ['a', 'b', 'c'], 'B': ['x', 'y', 'z']}, columns=['A', 'B']) + >>> new_df = ps.DataFrame({'B': ['d', 'e', 'f', 'g', 'h', 'i']}, columns=['B']) + >>> df.update(new_df) + >>> df.sort_index() + A B + 0 a d + 1 b e + 2 c f + + For Series, its name attribute must be set. + + >>> df = ps.DataFrame({'A': ['a', 'b', 'c'], 'B': ['x', 'y', 'z']}, columns=['A', 'B']) + >>> new_column = ps.Series(['d', 'e'], name='B', index=[0, 2]) + >>> df.update(new_column) + >>> df.sort_index() + A B + 0 a d + 1 b y + 2 c e + + If `other` contains None the corresponding values are not updated in the original dataframe. + + >>> df = ps.DataFrame({'A': [1, 2, 3], 'B': [400, 500, 600]}, columns=['A', 'B']) + >>> new_df = ps.DataFrame({'B': [4, None, 6]}, columns=['B']) + >>> df.update(new_df) + >>> df.sort_index() + A B + 0 1 4.0 + 1 2 500.0 + 2 3 6.0 + """ + if join != "left": + raise NotImplementedError("Only left join is supported") + + if isinstance(other, ps.Series): + other = other.to_frame() + + update_columns = list( + set(self._internal.column_labels).intersection(set(other._internal.column_labels)) + ) + update_sdf = self.join( + other[update_columns], rsuffix="_new" + )._internal.resolved_copy.spark_frame + + data_fields = self._internal.data_fields.copy() + for column_labels in update_columns: + column_name = self._internal.spark_column_name_for(column_labels) + old_col = scol_for(update_sdf, column_name) + new_col = scol_for( + update_sdf, other._internal.spark_column_name_for(column_labels) + "_new" + ) + if overwrite: + update_sdf = update_sdf.withColumn( + column_name, F.when(new_col.isNull(), old_col).otherwise(new_col) + ) + else: + update_sdf = update_sdf.withColumn( + column_name, F.when(old_col.isNull(), new_col).otherwise(old_col) + ) + data_fields[self._internal.column_labels.index(column_labels)] = None + sdf = update_sdf.select( + *[scol_for(update_sdf, col) for col in self._internal.spark_column_names], + *HIDDEN_COLUMNS, + ) + internal = self._internal.with_new_sdf(sdf, data_fields=data_fields) + self._update_internal_frame(internal, check_same_anchor=False) + + def cov(self, min_periods: Optional[int] = None, ddof: int = 1) -> "DataFrame": + """ + Compute pairwise covariance of columns, excluding NA/null values. + + Compute the pairwise covariance among the series of a DataFrame. + The returned data frame is the `covariance matrix + `__ of the columns + of the DataFrame. + + Both NA and null values are automatically excluded from the + calculation. (See the note below about bias from missing values.) + 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``. + + This method is generally used for the analysis of time series data to + understand the relationship between different measures across time. + + .. versionadded:: 3.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. + + .. versionadded:: 3.4.0 + + Returns + ------- + DataFrame + The covariance matrix of the series of the DataFrame. + + See Also + -------- + Series.cov : Compute covariance with another Series. + + Examples + -------- + >>> df = ps.DataFrame([(1, 2), (0, 3), (2, 0), (1, 1)], + ... columns=['dogs', 'cats']) + >>> df.cov() + dogs cats + dogs 0.666667 -1.000000 + cats -1.000000 1.666667 + + >>> np.random.seed(42) + >>> df = ps.DataFrame(np.random.randn(1000, 5), + ... columns=['a', 'b', 'c', 'd', 'e']) + >>> df.cov() + a b c d e + a 0.998438 -0.020161 0.059277 -0.008943 0.014144 + b -0.020161 1.059352 -0.008543 -0.024738 0.009826 + c 0.059277 -0.008543 1.010670 -0.001486 -0.000271 + d -0.008943 -0.024738 -0.001486 0.921297 -0.013692 + e 0.014144 0.009826 -0.000271 -0.013692 0.977795 + >>> df.cov(ddof=2) + a b c d e + a 0.999439 -0.020181 0.059336 -0.008952 0.014159 + b -0.020181 1.060413 -0.008551 -0.024762 0.009836 + c 0.059336 -0.008551 1.011683 -0.001487 -0.000271 + d -0.008952 -0.024762 -0.001487 0.922220 -0.013705 + e 0.014159 0.009836 -0.000271 -0.013705 0.978775 + >>> df.cov(ddof=-1) + a b c d e + a 0.996444 -0.020121 0.059158 -0.008926 0.014116 + b -0.020121 1.057235 -0.008526 -0.024688 0.009807 + c 0.059158 -0.008526 1.008650 -0.001483 -0.000270 + d -0.008926 -0.024688 -0.001483 0.919456 -0.013664 + e 0.014116 0.009807 -0.000270 -0.013664 0.975842 + + **Minimum number of periods** + + This method also supports an optional ``min_periods`` keyword + that specifies the required minimum number of non-NA observations for + each column pair to have a valid result: + + >>> np.random.seed(42) + >>> df = pd.DataFrame(np.random.randn(20, 3), + ... columns=['a', 'b', 'c']) + >>> df.loc[df.index[:5], 'a'] = np.nan + >>> df.loc[df.index[5:10], 'b'] = np.nan + >>> sdf = ps.from_pandas(df) + >>> sdf.cov(min_periods=12) + a b c + a 0.316741 NaN -0.150812 + b NaN 1.248003 0.191417 + c -0.150812 0.191417 0.895202 + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + min_periods = 1 if min_periods is None else min_periods + + # Only compute covariance for Boolean and Numeric except Decimal + psdf = self[ + [ + col + for col in self.columns + if isinstance(self[col].spark.data_type, BooleanType) + or ( + isinstance(self[col].spark.data_type, NumericType) + and not isinstance(self[col].spark.data_type, DecimalType) + ) + ] + ] + + num_cols = len(psdf.columns) + cov = np.zeros([num_cols, num_cols]) + + if num_cols == 0: + return DataFrame() + + if len(psdf) < min_periods: + cov.fill(np.nan) + return DataFrame(cov, columns=psdf.columns, index=psdf.columns) + + data_cols = psdf._internal.data_spark_column_names + cov_scols = [] + count_not_null_scols = [] + + # Count number of null row between two columns + # Example: + # a b c + # 0 1 1 1 + # 1 NaN 2 2 + # 2 3 NaN 3 + # 3 4 4 4 + # + # a b c + # a count(a, a) count(a, b) count(a, c) + # b count(b, b) count(b, c) + # c count(c, c) + # + # count_not_null_scols = + # [F.count(a, a), F.count(a, b), F.count(a, c), F.count(b, b), F.count(b, c), F.count(c, c)] + for r in range(0, num_cols): + for c in range(r, num_cols): + count_not_null_scols.append( + F.count( + F.when(F.col(data_cols[r]).isNotNull() & F.col(data_cols[c]).isNotNull(), 1) + ) + ) + + count_not_null = ( + psdf._internal.spark_frame.replace(float("nan"), None) + .select(*count_not_null_scols) + .head(1)[0] + ) + + # Calculate covariance between two columns + # Example: + # with min_periods = 3 + # a b c + # 0 1 1 1 + # 1 NaN 2 2 + # 2 3 NaN 3 + # 3 4 4 4 + # + # a b c + # a cov(a, a) None cov(a, c) + # b cov(b, b) cov(b, c) + # c cov(c, c) + # + # cov_scols = [F.cov(a, a), None, F.cov(a, c), F.cov(b, b), F.cov(b, c), F.cov(c, c)] + step = 0 + for r in range(0, num_cols): + step += r + for c in range(r, num_cols): + cov_scols.append( + SF.covar( + F.col(data_cols[r]).cast("double"), F.col(data_cols[c]).cast("double"), ddof + ) + if count_not_null[r * num_cols + c - step] >= min_periods + else F.lit(None) + ) + + pair_cov = psdf._internal.spark_frame.select(*cov_scols).head(1)[0] + + # Convert from row to 2D array + # Example: + # pair_cov = [cov(a, a), None, cov(a, c), cov(b, b), cov(b, c), cov(c, c)] + # + # cov = + # + # a b c + # a cov(a, a) None cov(a, c) + # b cov(b, b) cov(b, c) + # c cov(c, c) + step = 0 + for r in range(0, num_cols): + step += r + for c in range(r, num_cols): + cov[r][c] = pair_cov[r * num_cols + c - step] + + # Copy values + # Example: + # cov = + # a b c + # a cov(a, a) None cov(a, c) + # b None cov(b, b) cov(b, c) + # c cov(a, c) cov(b, c) cov(c, c) + cov = cov + cov.T - np.diag(np.diag(cov)) + return DataFrame(cov, columns=psdf.columns, index=psdf.columns) + + def sample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + replace: bool = False, + random_state: Optional[int] = None, + ignore_index: bool = False, + ) -> "DataFrame": + """ + Return a random sample of items from an axis of object. + + Please call this function using named argument by specifying the ``frac`` argument. + + You can use `random_state` for reproducibility. However, note that different from pandas, + specifying a seed in pandas-on-Spark/Spark does not guarantee the sampled rows will + be fixed. The result set depends on not only the seed, but also how the data is distributed + across machines and to some extent network randomness when shuffle operations are involved. + Even in the simplest case, the result set will depend on the system's CPU core count. + + Parameters + ---------- + n : int, optional + Number of items to return. This is currently NOT supported. Use frac instead. + frac : float, optional + Fraction of axis items to return. + replace : bool, default False + Sample with or without replacement. + random_state : int, optional + Seed for the random number generator (if int). + ignore_index : bool, default False + If True, the resulting index will be labeled 0, 1, …, n - 1. + + .. versionadded:: 3.4.0 + + Returns + ------- + Series or DataFrame + A new object of same type as caller containing the sampled items. + + Examples + -------- + >>> df = ps.DataFrame({'num_legs': [2, 4, 8, 0], + ... 'num_wings': [2, 0, 0, 0], + ... 'num_specimen_seen': [10, 2, 1, 8]}, + ... index=['falcon', 'dog', 'spider', 'fish'], + ... columns=['num_legs', 'num_wings', 'num_specimen_seen']) + >>> df # doctest: +SKIP + num_legs num_wings num_specimen_seen + falcon 2 2 10 + dog 4 0 2 + spider 8 0 1 + fish 0 0 8 + + A random 25% sample of the ``DataFrame``. + Note that we use `random_state` to ensure the reproducibility of + the examples. + + >>> df.sample(frac=0.25, random_state=1) # doctest: +SKIP + num_legs num_wings num_specimen_seen + falcon 2 2 10 + fish 0 0 8 + + A random 50% sample of the ``DataFrame``, while ignoring the index. + + >>> df.sample(frac=0.5, random_state=1, ignore_index=True) # doctest: +SKIP + num_legs num_wings num_specimen_seen + 0 4 0 2 + 1 8 0 1 + 2 0 0 8 + + Extract 25% random elements from the ``Series`` ``df['num_legs']`` with replacement + so, the same items could appear more than once. + + >>> df['num_legs'].sample(frac=0.4, replace=True, random_state=1) # doctest: +SKIP + falcon 2 + spider 8 + spider 8 + Name: num_legs, dtype: int64 + + Specifying the exact number of items to return is not supported now. + + >>> df.sample(n=5) # doctest: +ELLIPSIS + Traceback (most recent call last): + ... + NotImplementedError: Function sample currently does not support specifying ... + """ + # Note: we don't run any of the doctests because the result can change depending on the + # system's core count. + if n is not None: + raise NotImplementedError( + "Function sample currently does not support specifying " + "exact number of items to return. Use frac instead." + ) + + if frac is None: + raise ValueError("frac must be specified.") + + sdf = self._internal.resolved_copy.spark_frame.sample( + withReplacement=replace, fraction=frac, seed=random_state + ) + if ignore_index: + return DataFrame(sdf.drop(*self._internal.index_spark_column_names)) + else: + return DataFrame(self._internal.with_new_sdf(sdf)) + + def astype(self, dtype: Union[str, Dtype, Dict[Name, Union[str, Dtype]]]) -> "DataFrame": + """ + Cast a pandas-on-Spark object to a specified dtype ``dtype``. + + Parameters + ---------- + dtype : data type, or dict of column name -> data type + Use a numpy.dtype or Python type to cast entire pandas-on-Spark object to + the same type. Alternatively, use {col: dtype, ...}, where col is a + column label and dtype is a numpy.dtype or Python type to cast one + or more of the DataFrame's columns to column-specific types. + + Returns + ------- + casted : same type as caller + + See Also + -------- + to_datetime : Convert argument to datetime. + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2, 3], 'b': [1, 2, 3]}, dtype='int64') + >>> df + a b + 0 1 1 + 1 2 2 + 2 3 3 + + Convert to float type: + + >>> df.astype('float') + a b + 0 1.0 1.0 + 1 2.0 2.0 + 2 3.0 3.0 + + Convert to int64 type back: + + >>> df.astype('int64') + a b + 0 1 1 + 1 2 2 + 2 3 3 + + Convert column a to float type: + + >>> df.astype({'a': float}) + a b + 0 1.0 1 + 1 2.0 2 + 2 3.0 3 + + """ + applied = [] + if is_dict_like(dtype): + dtype_dict = cast(Dict[Name, Union[str, Dtype]], dtype) + for col_name in dtype_dict.keys(): + if col_name not in self.columns: + raise KeyError( + "Only a column name can be used for the " + "key in a dtype mappings argument." + ) + for col_name, col in self.items(): + if col_name in dtype_dict: + applied.append(col.astype(dtype=dtype_dict[col_name])) + else: + applied.append(col) + else: + for col_name, col in self.items(): + applied.append(col.astype(dtype=cast(Union[str, Dtype], dtype))) + return DataFrame(self._internal.with_new_columns(applied)) + + def add_prefix(self, prefix: str) -> "DataFrame": + """ + Prefix labels with string `prefix`. + + For Series, the row labels are prefixed. + For DataFrame, the column labels are prefixed. + + Parameters + ---------- + prefix : str + The string to add before each label. + + Returns + ------- + DataFrame + New DataFrame with updated labels. + + See Also + -------- + Series.add_prefix: Prefix row labels with string `prefix`. + Series.add_suffix: Suffix row labels with string `suffix`. + DataFrame.add_suffix: Suffix column labels with string `suffix`. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 3, 4], 'B': [3, 4, 5, 6]}, columns=['A', 'B']) + >>> df + A B + 0 1 3 + 1 2 4 + 2 3 5 + 3 4 6 + + >>> df.add_prefix('col_') + col_A col_B + 0 1 3 + 1 2 4 + 2 3 5 + 3 4 6 + """ + assert isinstance(prefix, str) + return self._apply_series_op( + lambda psser: psser.rename(tuple([prefix + i for i in psser._column_label])) + ) + + def add_suffix(self, suffix: str) -> "DataFrame": + """ + Suffix labels with string `suffix`. + + For Series, the row labels are suffixed. + For DataFrame, the column labels are suffixed. + + Parameters + ---------- + suffix : str + The string to add before each label. + + Returns + ------- + DataFrame + New DataFrame with updated labels. + + See Also + -------- + Series.add_prefix: Prefix row labels with string `prefix`. + Series.add_suffix: Suffix row labels with string `suffix`. + DataFrame.add_prefix: Prefix column labels with string `prefix`. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 3, 4], 'B': [3, 4, 5, 6]}, columns=['A', 'B']) + >>> df + A B + 0 1 3 + 1 2 4 + 2 3 5 + 3 4 6 + + >>> df.add_suffix('_col') + A_col B_col + 0 1 3 + 1 2 4 + 2 3 5 + 3 4 6 + """ + assert isinstance(suffix, str) + return self._apply_series_op( + lambda psser: psser.rename(tuple([i + suffix for i in psser._column_label])) + ) + + # TODO: include and exclude should be implemented. + def describe(self, percentiles: Optional[List[float]] = None) -> "DataFrame": + """ + Generate descriptive statistics that summarize the central tendency, + dispersion and shape of a dataset's distribution, excluding + ``NaN`` values. + + Analyzes both numeric and object series, as well + as ``DataFrame`` column sets of mixed data types. The output + will vary depending on what is provided. Refer to the notes + below for more detail. + + Parameters + ---------- + percentiles : list of ``float`` in range [0.0, 1.0], default [0.25, 0.5, 0.75] + A list of percentiles to be computed. + + Returns + ------- + DataFrame + Summary statistics of the Dataframe provided. + + See Also + -------- + DataFrame.count: Count number of non-NA/null observations. + DataFrame.max: Maximum of the values in the object. + DataFrame.min: Minimum of the values in the object. + DataFrame.mean: Mean of the values. + DataFrame.std: Standard deviation of the observations. + + Notes + ----- + For numeric data, the result's index will include ``count``, + ``mean``, ``std``, ``min``, ``25%``, ``50%``, ``75%``, ``max``. + + For object data (e.g. strings or timestamps), the result’s index will include + ``count``, ``unique``, ``top``, and ``freq``. + The ``top`` is the most common value. The ``freq`` is the most common value’s frequency. + Timestamps also include the ``first`` and ``last`` items. + + Examples + -------- + Describing a numeric ``Series``. + + >>> s = ps.Series([1, 2, 3]) + >>> s.describe() + count 3.0 + mean 2.0 + std 1.0 + min 1.0 + 25% 1.0 + 50% 2.0 + 75% 3.0 + max 3.0 + dtype: float64 + + Describing a ``DataFrame``. Only numeric fields are returned. + + >>> df = ps.DataFrame({'numeric1': [1, 2, 3], + ... 'numeric2': [4.0, 5.0, 6.0], + ... 'object': ['a', 'b', 'c'] + ... }, + ... columns=['numeric1', 'numeric2', 'object']) + >>> df.describe() + numeric1 numeric2 + count 3.0 3.0 + mean 2.0 5.0 + std 1.0 1.0 + min 1.0 4.0 + 25% 1.0 4.0 + 50% 2.0 5.0 + 75% 3.0 6.0 + max 3.0 6.0 + + For multi-index columns: + + >>> df.columns = [('num', 'a'), ('num', 'b'), ('obj', 'c')] + >>> df.describe() # doctest: +NORMALIZE_WHITESPACE + num + a b + count 3.0 3.0 + mean 2.0 5.0 + std 1.0 1.0 + min 1.0 4.0 + 25% 1.0 4.0 + 50% 2.0 5.0 + 75% 3.0 6.0 + max 3.0 6.0 + + >>> df[('num', 'b')].describe() + count 3.0 + mean 5.0 + std 1.0 + min 4.0 + 25% 4.0 + 50% 5.0 + 75% 6.0 + max 6.0 + Name: (num, b), dtype: float64 + + Describing a ``DataFrame`` and selecting custom percentiles. + + >>> df = ps.DataFrame({'numeric1': [1, 2, 3], + ... 'numeric2': [4.0, 5.0, 6.0] + ... }, + ... columns=['numeric1', 'numeric2']) + >>> df.describe(percentiles = [0.85, 0.15]) + numeric1 numeric2 + count 3.0 3.0 + mean 2.0 5.0 + std 1.0 1.0 + min 1.0 4.0 + 15% 1.0 4.0 + 50% 2.0 5.0 + 85% 3.0 6.0 + max 3.0 6.0 + + Describing a column from a ``DataFrame`` by accessing it as + an attribute. + + >>> df.numeric1.describe() + count 3.0 + mean 2.0 + std 1.0 + min 1.0 + 25% 1.0 + 50% 2.0 + 75% 3.0 + max 3.0 + Name: numeric1, dtype: float64 + + Describing a column from a ``DataFrame`` by accessing it as + an attribute and selecting custom percentiles. + + >>> df.numeric1.describe(percentiles = [0.85, 0.15]) + count 3.0 + mean 2.0 + std 1.0 + min 1.0 + 15% 1.0 + 50% 2.0 + 85% 3.0 + max 3.0 + Name: numeric1, dtype: float64 + """ + psser_numeric: List[Series] = [] + psser_string: List[Series] = [] + psser_timestamp: List[Series] = [] + spark_data_types: List[DataType] = [] + column_labels: Optional[List[Label]] = [] + column_names: List[str] = [] + for label in self._internal.column_labels: + psser = self._psser_for(label) + spark_data_type = psser.spark.data_type + if isinstance(spark_data_type, NumericType): + psser_numeric.append(psser) + column_labels.append(label) + spark_data_types.append(spark_data_type) + elif isinstance(spark_data_type, (TimestampType, TimestampNTZType)): + psser_timestamp.append(psser) + column_labels.append(label) + spark_data_types.append(spark_data_type) + else: + psser_string.append(psser) + column_names.append(self._internal.spark_column_name_for(label)) + + if percentiles is not None: + if any((p < 0.0) or (p > 1.0) for p in percentiles): + raise ValueError("Percentiles should all be in the interval [0, 1]") + # appending 50% if not in percentiles already + percentiles = (percentiles + [0.5]) if 0.5 not in percentiles else percentiles + else: + percentiles = [0.25, 0.5, 0.75] + + # Identify the cases + is_all_string_type = ( + len(psser_numeric) == 0 and len(psser_timestamp) == 0 and len(psser_string) > 0 + ) + is_all_numeric_type = len(psser_numeric) > 0 and len(psser_timestamp) == 0 + has_timestamp_type = len(psser_timestamp) > 0 + has_numeric_type = len(psser_numeric) > 0 + + if is_all_string_type: + # Handling string type columns + # We will retrive the `count`, `unique`, `top` and `freq`. + internal = self._internal.resolved_copy + exprs_string = [ + internal.spark_column_for(psser._column_label) for psser in psser_string + ] + sdf = internal.spark_frame.select(*exprs_string) + + # Get `count` & `unique` for each columns + counts, uniques = map(lambda x: x[1:], sdf.summary("count", "count_distinct").take(2)) + # Handling Empty DataFrame + if len(counts) == 0 or counts[0] == "0": + data = dict() + for psser in psser_string: + data[psser.name] = [0, 0, np.nan, np.nan] + return DataFrame(data, index=["count", "unique", "top", "freq"]) + + # Get `top` & `freq` for each columns + tops = [] + freqs = [] + # TODO(SPARK-37711): We should do it in single pass since invoking Spark job + # for every columns is too expensive. + for column in exprs_string: + top, freq = sdf.groupby(column).count().sort("count", ascending=False).first() + tops.append(str(top)) + freqs.append(str(freq)) + + stats = [counts, uniques, tops, freqs] + stats_names = ["count", "unique", "top", "freq"] + + result: DataFrame = DataFrame( + data=stats, + index=stats_names, + columns=column_names, + ) + elif is_all_numeric_type: + # Handling numeric columns + exprs_numeric = [ + psser._dtype_op.nan_to_null(psser).spark.column for psser in psser_numeric + ] + formatted_perc = ["{:.0%}".format(p) for p in sorted(percentiles)] + stats = ["count", "mean", "stddev", "min", *formatted_perc, "max"] + + # In this case, we can simply use `summary` to calculate the stats. + sdf = self._internal.spark_frame.select(*exprs_numeric).summary(*stats) + sdf = sdf.replace("stddev", "std", subset=["summary"]) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, "summary")], + column_labels=column_labels, + data_spark_columns=[ + scol_for(sdf, self._internal.spark_column_name_for(label)) + for label in column_labels + ], + ) + result = DataFrame(internal).astype("float64") + elif has_timestamp_type: + internal = self._internal.resolved_copy + column_names = [ + internal.spark_column_name_for(column_label) for column_label in column_labels + ] + column_length = len(column_labels) + + # Apply stat functions for each column. + count_exprs = map(F.count, column_names) + min_exprs = map(F.min, column_names) + # Here we try to flat the multiple maps into single list that contains each calculated + # percentile using `chain`. + # e.g. flat the `[, ]` + # to `[Column<'percentile_approx(A, 0.2, 10000)'>, + # Column<'percentile_approx(B, 0.2, 10000)'>, + # Column<'percentile_approx(A, 0.5, 10000)'>, + # Column<'percentile_approx(B, 0.5, 10000)'>]` + perc_exprs = chain( + *[ + map(F.percentile_approx, column_names, [percentile] * column_length) + for percentile in percentiles + ] + ) + max_exprs = map(F.max, column_names) + mean_exprs = [] + for column_name, spark_data_type in zip(column_names, spark_data_types): + mean_exprs.append(F.mean(column_name).astype(spark_data_type)) + exprs = [*count_exprs, *mean_exprs, *min_exprs, *perc_exprs, *max_exprs] + + formatted_perc = ["{:.0%}".format(p) for p in sorted(percentiles)] + stats_names = ["count", "mean", "min", *formatted_perc, "max"] + + # If not all columns are timestamp type, + # we also need to calculate the `std` for numeric columns + if has_numeric_type: + std_exprs = [] + for label, spark_data_type in zip(column_labels, spark_data_types): + column_name = label[0] + if isinstance(spark_data_type, (TimestampType, TimestampNTZType)): + std_exprs.append(F.lit(None).alias("stddev_samp({})".format(column_name))) + else: + std_exprs.append(F.stddev(column_name)) + exprs.extend(std_exprs) + stats_names.append("std") + + # Select stats for all columns at once. + sdf = internal.spark_frame.select(exprs) + stat_values = sdf.first() + + num_stats = int(len(exprs) / column_length) + # `column_name_stats_kv` is key-value store that has column name as key, and + # the stats as values e.g. {"A": [{count_value}, {min_value}, ...], + # "B": [{count_value}, {min_value} ...]} + column_name_stats_kv: Dict[str, List[str]] = defaultdict(list) + for i, column_name in enumerate(column_names): + for first_stat_idx in range(num_stats): + column_name_stats_kv[column_name].append( + stat_values[(first_stat_idx * column_length) + i] + ) + + # For timestamp type columns, we should cast the column type to string. + for key, spark_data_type in zip(column_name_stats_kv, spark_data_types): + if isinstance(spark_data_type, (TimestampType, TimestampNTZType)): + column_name_stats_kv[key] = [str(value) for value in column_name_stats_kv[key]] + + result: DataFrame = DataFrame( # type: ignore[no-redef] + data=column_name_stats_kv, + index=stats_names, + columns=column_names, + ) + else: + # Empty DataFrame without column + raise ValueError("Cannot describe a DataFrame without columns") + + return result + + def drop_duplicates( + self, + subset: Optional[Union[Name, List[Name]]] = None, + keep: Union[bool, str] = "first", + inplace: bool = False, + ignore_index: bool = False, + ) -> Optional["DataFrame"]: + """ + Return DataFrame with duplicate rows removed, optionally only + considering certain columns. + + Parameters + ---------- + subset : column label or sequence of labels, optional + Only consider certain columns for identifying duplicates, by + default use all the columns. + keep : {'first', 'last', False}, default 'first' + Determines which duplicates (if any) to keep. + - ``first`` : Drop duplicates except for the first occurrence. + - ``last`` : Drop duplicates except for the last occurrence. + - False : Drop all duplicates. + inplace : boolean, default False + Whether to drop duplicates in place or to return a copy. + ignore_index : boolean, default False + If True, the resulting axis will be labeled 0, 1, …, n - 1. + + Returns + ------- + DataFrame + DataFrame with duplicates removed or None if ``inplace=True``. + + >>> df = ps.DataFrame( + ... {'a': [1, 2, 2, 2, 3], 'b': ['a', 'a', 'a', 'c', 'd']}, columns = ['a', 'b']) + >>> df + a b + 0 1 a + 1 2 a + 2 2 a + 3 2 c + 4 3 d + + >>> df.drop_duplicates().sort_index() + a b + 0 1 a + 1 2 a + 3 2 c + 4 3 d + + >>> df.drop_duplicates(ignore_index=True).sort_index() + a b + 0 1 a + 1 2 a + 2 2 c + 3 3 d + + >>> df.drop_duplicates('a').sort_index() + a b + 0 1 a + 1 2 a + 4 3 d + + >>> df.drop_duplicates(['a', 'b']).sort_index() + a b + 0 1 a + 1 2 a + 3 2 c + 4 3 d + + >>> df.drop_duplicates(keep='last').sort_index() + a b + 0 1 a + 2 2 a + 3 2 c + 4 3 d + + >>> df.drop_duplicates(keep=False).sort_index() + a b + 0 1 a + 3 2 c + 4 3 d + """ + inplace = validate_bool_kwarg(inplace, "inplace") + + sdf, column = self._mark_duplicates(subset, keep) + + sdf = sdf.where(~scol_for(sdf, column)).drop(column) + internal = self._internal.with_new_sdf(sdf) + psdf: DataFrame = DataFrame(internal) + + if inplace: + if ignore_index: + psdf.reset_index(drop=True, inplace=inplace) + self._update_internal_frame(psdf._internal) + return None + else: + return psdf.reset_index(drop=True) if ignore_index else psdf + + def reindex( + self, + labels: Optional[Sequence[Any]] = None, + index: Optional[Union["Index", Sequence[Any]]] = None, + columns: Optional[Union[pd.Index, Sequence[Any]]] = None, + axis: Optional[Axis] = None, + copy: Optional[bool] = True, + fill_value: Optional[Any] = None, + ) -> "DataFrame": + """ + Conform DataFrame to new index with optional filling logic, placing + NA/NaN in locations having no value in the previous index. A new object + is produced unless the new index is equivalent to the current one and + ``copy=False``. + + Parameters + ---------- + labels: array-like, optional + New labels / index to conform the axis specified by ‘axis’ to. + index, columns: array-like, optional + New labels / index to conform to, should be specified using keywords. + Preferably an Index object to avoid duplicating data + axis: int or str, optional + Axis to target. Can be either the axis name (‘index’, ‘columns’) or + number (0, 1). + copy : bool, default True + Return a new object, even if the passed indexes are the same. + fill_value : scalar, default np.NaN + Value to use for missing values. Defaults to NaN, but can be any + "compatible" value. + + Returns + ------- + DataFrame with changed index. + + See Also + -------- + DataFrame.set_index : Set row labels. + DataFrame.reset_index : Remove row labels or move them to new columns. + + Examples + -------- + + ``DataFrame.reindex`` supports two calling conventions + + * ``(index=index_labels, columns=column_labels, ...)`` + * ``(labels, axis={'index', 'columns'}, ...)`` + + We *highly* recommend using keyword arguments to clarify your + intent. + + Create a dataframe with some fictional data. + + >>> index = ['Firefox', 'Chrome', 'Safari', 'IE10', 'Konqueror'] + >>> df = ps.DataFrame({ + ... 'http_status': [200, 200, 404, 404, 301], + ... 'response_time': [0.04, 0.02, 0.07, 0.08, 1.0]}, + ... index=index, + ... columns=['http_status', 'response_time']) + >>> df + http_status response_time + Firefox 200 0.04 + Chrome 200 0.02 + Safari 404 0.07 + IE10 404 0.08 + Konqueror 301 1.00 + + Create a new index and reindex the dataframe. By default + values in the new index that do not have corresponding + records in the dataframe are assigned ``NaN``. + + >>> new_index= ['Safari', 'Iceweasel', 'Comodo Dragon', 'IE10', + ... 'Chrome'] + >>> df.reindex(new_index).sort_index() + http_status response_time + Chrome 200.0 0.02 + Comodo Dragon NaN NaN + IE10 404.0 0.08 + Iceweasel NaN NaN + Safari 404.0 0.07 + + We can fill in the missing values by passing a value to + the keyword ``fill_value``. + + >>> df.reindex(new_index, fill_value=0, copy=False).sort_index() + http_status response_time + Chrome 200 0.02 + Comodo Dragon 0 0.00 + IE10 404 0.08 + Iceweasel 0 0.00 + Safari 404 0.07 + + We can also reindex the columns. + + >>> df.reindex(columns=['http_status', 'user_agent']).sort_index() + http_status user_agent + Chrome 200 NaN + Firefox 200 NaN + IE10 404 NaN + Konqueror 301 NaN + Safari 404 NaN + + Or we can use "axis-style" keyword arguments + + >>> df.reindex(['http_status', 'user_agent'], axis="columns").sort_index() + http_status user_agent + Chrome 200 NaN + Firefox 200 NaN + IE10 404 NaN + Konqueror 301 NaN + Safari 404 NaN + + To further illustrate the filling functionality in + ``reindex``, we will create a dataframe with a + monotonically increasing index (for example, a sequence + of dates). + + >>> date_index = pd.date_range('1/1/2010', periods=6, freq='D') + >>> df2 = ps.DataFrame({"prices": [100, 101, np.nan, 100, 89, 88]}, + ... index=date_index) + >>> df2.sort_index() + prices + 2010-01-01 100.0 + 2010-01-02 101.0 + 2010-01-03 NaN + 2010-01-04 100.0 + 2010-01-05 89.0 + 2010-01-06 88.0 + + Suppose we decide to expand the dataframe to cover a wider + date range. + + >>> date_index2 = pd.date_range('12/29/2009', periods=10, freq='D') + >>> df2.reindex(date_index2).sort_index() + prices + 2009-12-29 NaN + 2009-12-30 NaN + 2009-12-31 NaN + 2010-01-01 100.0 + 2010-01-02 101.0 + 2010-01-03 NaN + 2010-01-04 100.0 + 2010-01-05 89.0 + 2010-01-06 88.0 + 2010-01-07 NaN + """ + if axis is not None and (index is not None or columns is not None): + raise TypeError("Cannot specify both 'axis' and any of 'index' or 'columns'.") + + if labels is not None: + axis = validate_axis(axis) + if axis == 0: + index = labels + elif axis == 1: + columns = labels + + if index is not None and not is_list_like(index): + raise TypeError( + "Index must be called with a collection of some kind, " + "%s was passed" % type(index) + ) + + if columns is not None and not is_list_like(columns): + raise TypeError( + "Columns must be called with a collection of some kind, " + "%s was passed" % type(columns) + ) + + df = self + + if index is not None: + df = df._reindex_index(index, fill_value) + + if columns is not None: + df = df._reindex_columns(columns, fill_value) + + # Copy + if copy and df is self: + return df.copy() + else: + return df + + def _reindex_index( + self, index: Optional[Union["Index", Sequence[Any]]], fill_value: Optional[Any] + ) -> "DataFrame": + # When axis is index, we can mimic pandas by a right outer join. + nlevels = self._internal.index_level + assert nlevels <= 1 or ( + isinstance(index, ps.MultiIndex) and nlevels == index.nlevels + ), "MultiIndex DataFrame can only be reindexed with a similar pandas-on-Spark MultiIndex." + + index_columns = self._internal.index_spark_column_names + frame = self._internal.resolved_copy.spark_frame.drop(NATURAL_ORDER_COLUMN_NAME) + + if isinstance(index, ps.Index): + if nlevels != index.nlevels: + return DataFrame(index._internal.with_new_columns([])).reindex( + columns=self.columns, fill_value=fill_value + ) + + index_names = index._internal.index_names + scols = index._internal.index_spark_columns + labels = index._internal.spark_frame.select( + [scol.alias(index_column) for scol, index_column in zip(scols, index_columns)] + ) + else: + index = ps.Index(list(index)) + labels = index._internal.spark_frame.select(index.spark.column.alias(index_columns[0])) + index_names = self._internal.index_names + + if fill_value is not None: + frame_index_columns = [ + verify_temp_column_name(frame, "__frame_index_column_{}__".format(i)) + for i in range(nlevels) + ] + index_scols = [ + scol_for(frame, index_col).alias(frame_index_col) + for index_col, frame_index_col in zip(index_columns, frame_index_columns) + ] + scols = self._internal.resolved_copy.data_spark_columns + frame = frame.select(index_scols + scols) + + temp_fill_value = verify_temp_column_name(frame, "__fill_value__") + labels = labels.withColumn(temp_fill_value, F.lit(fill_value)) + + frame_index_scols = [scol_for(frame, col) for col in frame_index_columns] + labels_index_scols = [scol_for(labels, col) for col in index_columns] + + joined_df = frame.join( + labels, + on=[fcol == lcol for fcol, lcol in zip(frame_index_scols, labels_index_scols)], + how="right", + ) + + joined_df = joined_df.select( + *labels_index_scols, + *[ + F.when( + reduce( + lambda c1, c2: c1 & c2, + [ + fcol.isNull() & lcol.isNotNull() + for fcol, lcol in zip(frame_index_scols, labels_index_scols) + ], + ), + scol_for(joined_df, temp_fill_value), + ) + .otherwise(scol_for(joined_df, col)) + .alias(col) + for col in self._internal.data_spark_column_names + ], + ) + data_fields = None + else: + joined_df = frame.join(labels, on=index_columns, how="right") + data_fields = [field.copy(nullable=True) for field in self._internal.data_fields] + + sdf = joined_df.drop(NATURAL_ORDER_COLUMN_NAME) + internal = self._internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=index_names, + index_fields=[ + field.copy(name=name) + for field, name in zip( + index._internal.index_fields, self._internal.index_spark_column_names + ) + ], + data_spark_columns=[ + scol_for(sdf, col) for col in self._internal.data_spark_column_names + ], + data_fields=data_fields, + ) + return DataFrame(internal) + + def _reindex_columns( + self, columns: Optional[Union[pd.Index, Sequence[Any]]], fill_value: Optional[Any] + ) -> "DataFrame": + level = self._internal.column_labels_level + if level > 1: + label_columns = list(columns) + for col in label_columns: + if not isinstance(col, tuple): + raise TypeError("Expected tuple, got {}".format(type(col).__name__)) + else: + label_columns = [(col,) for col in columns] + for col in label_columns: + if len(col) != level: + raise ValueError( + "shape (1,{}) doesn't match the shape (1,{})".format(len(col), level) + ) + fill_value = np.nan if fill_value is None else fill_value + scols_or_pssers: List[Union[PySparkColumn, "Series"]] = [] + labels = [] + for label in label_columns: + if label in self._internal.column_labels: + scols_or_pssers.append(self._psser_for(label)) + else: + scols_or_pssers.append(F.lit(fill_value).alias(name_like_string(label))) + labels.append(label) + + if isinstance(columns, pd.Index): + column_label_names = [ + name if is_name_like_tuple(name) else (name,) for name in columns.names + ] + internal = self._internal.with_new_columns( + scols_or_pssers, column_labels=labels, column_label_names=column_label_names + ) + else: + internal = self._internal.with_new_columns(scols_or_pssers, column_labels=labels) + + return DataFrame(internal) + + def reindex_like(self, other: "DataFrame", copy: bool = True) -> "DataFrame": + """ + Return a DataFrame with matching indices as other object. + + Conform the object to the same index on all axes. Places NA/NaN in locations + having no value in the previous index. A new object is produced unless the + new index is equivalent to the current one and copy=False. + + Parameters + ---------- + other : DataFrame + Its row and column indices are used to define the new indices + of this object. + copy : bool, default True + Return a new object, even if the passed indexes are the same. + + Returns + ------- + DataFrame + DataFrame with changed indices on each axis. + + See Also + -------- + DataFrame.set_index : Set row labels. + DataFrame.reset_index : Remove row labels or move them to new columns. + DataFrame.reindex : Change to new indices or expand indices. + + Notes + ----- + Same as calling + ``.reindex(index=other.index, columns=other.columns,...)``. + + Examples + -------- + + >>> df1 = ps.DataFrame([[24.3, 75.7, 'high'], + ... [31, 87.8, 'high'], + ... [22, 71.6, 'medium'], + ... [35, 95, 'medium']], + ... columns=['temp_celsius', 'temp_fahrenheit', + ... 'windspeed'], + ... index=pd.date_range(start='2014-02-12', + ... end='2014-02-15', freq='D')) + >>> df1 + temp_celsius temp_fahrenheit windspeed + 2014-02-12 24.3 75.7 high + 2014-02-13 31.0 87.8 high + 2014-02-14 22.0 71.6 medium + 2014-02-15 35.0 95.0 medium + + >>> df2 = ps.DataFrame([[28, 'low'], + ... [30, 'low'], + ... [35.1, 'medium']], + ... columns=['temp_celsius', 'windspeed'], + ... index=pd.DatetimeIndex(['2014-02-12', '2014-02-13', + ... '2014-02-15'])) + >>> df2 + temp_celsius windspeed + 2014-02-12 28.0 low + 2014-02-13 30.0 low + 2014-02-15 35.1 medium + + >>> df2.reindex_like(df1).sort_index() # doctest: +NORMALIZE_WHITESPACE + temp_celsius temp_fahrenheit windspeed + 2014-02-12 28.0 NaN low + 2014-02-13 30.0 NaN low + 2014-02-14 NaN NaN None + 2014-02-15 35.1 NaN medium + """ + + if isinstance(other, DataFrame): + return self.reindex(index=other.index, columns=other.columns, copy=copy) + else: + raise TypeError("other must be a pandas-on-Spark DataFrame") + + def melt( + self, + id_vars: Optional[Union[Name, List[Name]]] = None, + value_vars: Optional[Union[Name, List[Name]]] = None, + var_name: Optional[Union[str, List[str]]] = None, + value_name: str = "value", + ) -> "DataFrame": + """ + Unpivot a DataFrame from wide format to long format, optionally + leaving identifier variables set. + + This function is useful to massage a DataFrame into a format where one + or more columns are identifier variables (`id_vars`), while all other + columns, considered measured variables (`value_vars`), are "unpivoted" to + the row axis, leaving just two non-identifier columns, 'variable' and + 'value'. + + Parameters + ---------- + frame : DataFrame + id_vars : tuple, list, or ndarray, optional + Column(s) to use as identifier variables. + value_vars : tuple, list, or ndarray, optional + Column(s) to unpivot. If not specified, uses all columns that + are not set as `id_vars`. + var_name : scalar, default 'variable' + Name to use for the 'variable' column. If None it uses `frame.columns.name` or + ‘variable’. + value_name : scalar, default 'value' + Name to use for the 'value' column. + + Returns + ------- + DataFrame + Unpivoted DataFrame. + + Examples + -------- + >>> df = ps.DataFrame({'A': {0: 'a', 1: 'b', 2: 'c'}, + ... 'B': {0: 1, 1: 3, 2: 5}, + ... 'C': {0: 2, 1: 4, 2: 6}}, + ... columns=['A', 'B', 'C']) + >>> df + A B C + 0 a 1 2 + 1 b 3 4 + 2 c 5 6 + + >>> ps.melt(df) + variable value + 0 A a + 1 B 1 + 2 C 2 + 3 A b + 4 B 3 + 5 C 4 + 6 A c + 7 B 5 + 8 C 6 + + >>> df.melt(id_vars='A') + A variable value + 0 a B 1 + 1 a C 2 + 2 b B 3 + 3 b C 4 + 4 c B 5 + 5 c C 6 + + >>> df.melt(value_vars='A') + variable value + 0 A a + 1 A b + 2 A c + + >>> ps.melt(df, id_vars=['A', 'B']) + A B variable value + 0 a 1 C 2 + 1 b 3 C 4 + 2 c 5 C 6 + + >>> df.melt(id_vars=['A'], value_vars=['C']) + A variable value + 0 a C 2 + 1 b C 4 + 2 c C 6 + + The names of 'variable' and 'value' columns can be customized: + + >>> ps.melt(df, id_vars=['A'], value_vars=['B'], + ... var_name='myVarname', value_name='myValname') + A myVarname myValname + 0 a B 1 + 1 b B 3 + 2 c B 5 + """ + column_labels = self._internal.column_labels + + if id_vars is None: + id_vars = [] + else: + if isinstance(id_vars, tuple): + if self._internal.column_labels_level == 1: + id_vars = [idv if is_name_like_tuple(idv) else (idv,) for idv in id_vars] + else: + raise ValueError( + "id_vars must be a list of tuples" " when columns are a MultiIndex" + ) + elif is_name_like_value(id_vars): + id_vars = [(id_vars,)] + else: + id_vars = [idv if is_name_like_tuple(idv) else (idv,) for idv in id_vars] + + non_existence_col = [idv for idv in id_vars if idv not in column_labels] + if len(non_existence_col) != 0: + raveled_column_labels: np.ndarray[Any, np.dtype[Any]] = np.ravel(column_labels) + missing = [ + nec for nec in np.ravel(non_existence_col) if nec not in raveled_column_labels + ] + if len(missing) != 0: + raise KeyError( + "The following 'id_vars' are not present" + " in the DataFrame: {}".format(missing) + ) + else: + raise KeyError( + "None of {} are in the {}".format(non_existence_col, column_labels) + ) + + if value_vars is None: + value_vars = [] + else: + if isinstance(value_vars, tuple): + if self._internal.column_labels_level == 1: + value_vars = [ + valv if is_name_like_tuple(valv) else (valv,) for valv in value_vars + ] + else: + raise ValueError( + "value_vars must be a list of tuples" " when columns are a MultiIndex" + ) + elif is_name_like_value(value_vars): + value_vars = [(value_vars,)] + else: + value_vars = [valv if is_name_like_tuple(valv) else (valv,) for valv in value_vars] + + non_existence_col = [valv for valv in value_vars if valv not in column_labels] + if len(non_existence_col) != 0: + raveled_column_labels = np.ravel(column_labels) + missing = [ + nec for nec in np.ravel(non_existence_col) if nec not in raveled_column_labels + ] + if len(missing) != 0: + raise KeyError( + "The following 'value_vars' are not present" + " in the DataFrame: {}".format(missing) + ) + else: + raise KeyError( + "None of {} are in the {}".format(non_existence_col, column_labels) + ) + + if len(value_vars) == 0: + value_vars = column_labels + + column_labels = [label for label in column_labels if label not in id_vars] + + sdf = self._internal.spark_frame + + if var_name is None: + if ( + self._internal.column_labels_level == 1 + and self._internal.column_label_names[0] is None + ): + var_name = ["variable"] + else: + var_name = [ + name_like_string(name) if name is not None else "variable_{}".format(i) + for i, name in enumerate(self._internal.column_label_names) + ] + elif isinstance(var_name, str): + var_name = [var_name] + + pairs = F.explode( + F.array( + *[ + F.struct( + *[F.lit(c).alias(name) for c, name in zip(label, var_name)], + *[self._internal.spark_column_for(label).alias(value_name)], + ) + for label in column_labels + if label in value_vars + ] + ) + ) + + columns = ( + [ + self._internal.spark_column_for(label).alias(name_like_string(label)) + for label in id_vars + ] + + [F.col("pairs.`%s`" % name) for name in var_name] + + [F.col("pairs.`%s`" % value_name)] + ) + exploded_df = sdf.withColumn("pairs", pairs).select(columns) + + return DataFrame( + InternalFrame( + spark_frame=exploded_df, + index_spark_columns=None, + column_labels=( + [label if len(label) == 1 else (name_like_string(label),) for label in id_vars] + + [(name,) for name in var_name] + + [(value_name,)] + ), + ) + ) + + def stack(self) -> DataFrameOrSeries: + """ + Stack the prescribed level(s) from columns to index. + + Return a reshaped DataFrame or Series having a multi-level + index with one or more new inner-most levels compared to the current + DataFrame. The new inner-most levels are created by pivoting the + columns of the current dataframe: + + - if the columns have a single level, the output is a Series + - if the columns have multiple levels, the new index + level(s) is (are) taken from the prescribed level(s) and + the output is a DataFrame. + + The new index levels are sorted. + + Returns + ------- + DataFrame or Series + Stacked dataframe or series. + + See Also + -------- + DataFrame.unstack : Unstack prescribed level(s) from index axis + onto column axis. + DataFrame.pivot : Reshape dataframe from long format to wide + format. + DataFrame.pivot_table : Create a spreadsheet-style pivot table + as a DataFrame. + + Notes + ----- + The function is named by analogy with a collection of books + being reorganized from being side by side on a horizontal + position (the columns of the dataframe) to being stacked + vertically on top of each other (in the index of the + dataframe). + + Examples + -------- + **Single level columns** + + >>> df_single_level_cols = ps.DataFrame([[0, 1], [2, 3]], + ... index=['cat', 'dog'], + ... columns=['weight', 'height']) + + Stacking a dataframe with a single level column axis returns a Series: + + >>> df_single_level_cols + weight height + cat 0 1 + dog 2 3 + >>> df_single_level_cols.stack().sort_index() + cat height 1 + weight 0 + dog height 3 + weight 2 + dtype: int64 + + **Multi level columns: simple case** + + >>> multicol1 = pd.MultiIndex.from_tuples([('weight', 'kg'), + ... ('weight', 'pounds')]) + >>> df_multi_level_cols1 = ps.DataFrame([[1, 2], [2, 4]], + ... index=['cat', 'dog'], + ... columns=multicol1) + + Stacking a dataframe with a multi-level column axis: + + >>> df_multi_level_cols1 # doctest: +NORMALIZE_WHITESPACE + weight + kg pounds + cat 1 2 + dog 2 4 + >>> df_multi_level_cols1.stack().sort_index() + weight + cat kg 1 + pounds 2 + dog kg 2 + pounds 4 + + **Missing values** + + >>> multicol2 = pd.MultiIndex.from_tuples([('weight', 'kg'), + ... ('height', 'm')]) + >>> df_multi_level_cols2 = ps.DataFrame([[1.0, 2.0], [3.0, 4.0]], + ... index=['cat', 'dog'], + ... columns=multicol2) + + It is common to have missing values when stacking a dataframe + with multi-level columns, as the stacked dataframe typically + has more values than the original dataframe. Missing values + are filled with NaNs: + + >>> df_multi_level_cols2 + weight height + kg m + cat 1.0 2.0 + dog 3.0 4.0 + >>> df_multi_level_cols2.stack().sort_index() # doctest: +SKIP + height weight + cat kg NaN 1.0 + m 2.0 NaN + dog kg NaN 3.0 + m 4.0 NaN + """ + from pyspark.pandas.series import first_series + + if len(self._internal.column_labels) == 0: + return DataFrame( + self._internal.copy( + column_label_names=self._internal.column_label_names[:-1] + ).with_filter(F.lit(False)) + ) + + column_labels: Dict[Label, Dict[Any, PySparkColumn]] = defaultdict(dict) + index_values = set() + should_returns_series = False + for label in self._internal.column_labels: + new_label = label[:-1] + if len(new_label) == 0: + new_label = None + should_returns_series = True + value = label[-1] + + scol = self._internal.spark_column_for(label) + column_labels[new_label][value] = scol + + index_values.add(value) + + column_labels = dict(sorted(column_labels.items(), key=lambda x: x[0])) + + index_name = self._internal.column_label_names[-1] + column_label_names = self._internal.column_label_names[:-1] + if len(column_label_names) == 0: + column_label_names = [None] + + index_column = SPARK_INDEX_NAME_FORMAT(self._internal.index_level) + data_columns = [name_like_string(label) for label in column_labels] + + structs = [ + F.struct( + *[F.lit(value).alias(index_column)], + *[ + ( + column_labels[label][value] + if value in column_labels[label] + else F.lit(None) + ).alias(name) + for label, name in zip(column_labels, data_columns) + ], + ).alias(value) + for value in index_values + ] + + pairs = F.explode(F.array(*structs)) + + sdf = self._internal.spark_frame.withColumn("pairs", pairs) + sdf = sdf.select( + self._internal.index_spark_columns + + [sdf["pairs"][index_column].alias(index_column)] + + [sdf["pairs"][name].alias(name) for name in data_columns] + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) + for col in (self._internal.index_spark_column_names + [index_column]) + ], + index_names=self._internal.index_names + [index_name], + index_fields=self._internal.index_fields + [None], + column_labels=list(column_labels), + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + column_label_names=column_label_names, + ) + psdf: DataFrame = DataFrame(internal) + + if should_returns_series: + return first_series(psdf) + else: + return psdf + + def unstack(self) -> DataFrameOrSeries: + """ + Pivot the (necessarily hierarchical) index labels. + + Returns a DataFrame having a new level of column labels whose inner-most level + consists of the pivoted index labels. + + If the index is not a MultiIndex, the output will be a Series. + + .. note:: If the index is a MultiIndex, the output DataFrame could be very wide, and + it could cause a serious performance degradation since Spark partitions its row based. + + Returns + ------- + Series or DataFrame + + See Also + -------- + DataFrame.pivot : Pivot a table based on column values. + DataFrame.stack : Pivot a level of the column labels (inverse operation from unstack). + + Examples + -------- + >>> df = ps.DataFrame({"A": {"0": "a", "1": "b", "2": "c"}, + ... "B": {"0": "1", "1": "3", "2": "5"}, + ... "C": {"0": "2", "1": "4", "2": "6"}}, + ... columns=["A", "B", "C"]) + >>> df + A B C + 0 a 1 2 + 1 b 3 4 + 2 c 5 6 + + >>> df.unstack().sort_index() + A 0 a + 1 b + 2 c + B 0 1 + 1 3 + 2 5 + C 0 2 + 1 4 + 2 6 + dtype: object + + >>> df.columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C')]) + >>> df.unstack().sort_index() + X A 0 a + 1 b + 2 c + B 0 1 + 1 3 + 2 5 + Y C 0 2 + 1 4 + 2 6 + dtype: object + + For MultiIndex case: + + >>> df = ps.DataFrame({"A": ["a", "b", "c"], + ... "B": [1, 3, 5], + ... "C": [2, 4, 6]}, + ... columns=["A", "B", "C"]) + >>> df = df.set_index('A', append=True) + >>> df # doctest: +NORMALIZE_WHITESPACE + B C + A + 0 a 1 2 + 1 b 3 4 + 2 c 5 6 + >>> df.unstack().sort_index() # doctest: +NORMALIZE_WHITESPACE + B C + A a b c a b c + 0 1.0 NaN NaN 2.0 NaN NaN + 1 NaN 3.0 NaN NaN 4.0 NaN + 2 NaN NaN 5.0 NaN NaN 6.0 + """ + from pyspark.pandas.series import first_series + + if self._internal.index_level > 1: + # The index after `reset_index()` will never be used, so use "distributed" index + # as a dummy to avoid overhead. + with option_context("compute.default_index_type", "distributed"): + df = self.reset_index() + index = df._internal.column_labels[: self._internal.index_level - 1] + columns = df.columns[self._internal.index_level - 1] + df = df.pivot_table( + index=index, columns=columns, values=self._internal.column_labels, aggfunc="first" + ) + internal = df._internal.copy( + index_names=self._internal.index_names[:-1], + index_fields=df._internal.index_fields[: self._internal.index_level - 1], + column_label_names=( + df._internal.column_label_names[:-1] + + [ + None + if self._internal.index_names[-1] is None + else df._internal.column_label_names[-1] + ] + ), + ) + return DataFrame(internal) + + # TODO: Codes here are similar with melt. Should we deduplicate? + column_labels = self._internal.column_labels + ser_name = SPARK_DEFAULT_SERIES_NAME + sdf = self._internal.spark_frame + new_index_columns = [ + SPARK_INDEX_NAME_FORMAT(i) for i in range(self._internal.column_labels_level) + ] + + new_index_map = list(zip_longest(new_index_columns, self._internal.column_label_names, [])) + + pairs = F.explode( + F.array( + *[ + F.struct( + *[F.lit(c).alias(name) for c, name in zip(idx, new_index_columns)], + *[self._internal.spark_column_for(idx).alias(ser_name)], + ) + for idx in column_labels + ] + ) + ) + + columns = [ + F.col("pairs.%s" % name) + for name in new_index_columns[: self._internal.column_labels_level] + ] + [F.col("pairs.%s" % ser_name)] + + new_index_len = len(new_index_columns) + existing_index_columns = [] + for i, (index_name, index_field) in enumerate( + zip(self._internal.index_names, self._internal.index_fields) + ): + name = SPARK_INDEX_NAME_FORMAT(i + new_index_len) + new_index_map.append((name, index_name, index_field.copy(name=name))) + existing_index_columns.append(self._internal.index_spark_columns[i].alias(name)) + + exploded_df = sdf.withColumn("pairs", pairs).select(existing_index_columns + columns) + + index_spark_column_names, index_names, index_fields = zip(*new_index_map) + return first_series( + DataFrame( + InternalFrame( + exploded_df, + index_spark_columns=[ + scol_for(exploded_df, col) for col in index_spark_column_names + ], + index_names=list(index_names), + index_fields=list(index_fields), + column_labels=[None], + ) + ) + ) + + # TODO: axis, level and **kwargs should be implemented. + def all( + self, axis: Axis = 0, bool_only: Optional[bool] = None, skipna: bool = True + ) -> "Series": + """ + Return whether all elements are True. + + Returns True unless there is at least one element within a series that is + False or equivalent (e.g. zero or empty) + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Indicate which axis or axes should be reduced. + + * 0 / 'index' : reduce the index, return a Series whose index is the + original column labels. + + bool_only : bool, default None + Include only boolean columns. If None, will attempt to use everything, + then use only boolean data. + + skipna : boolean, default True + Exclude NA values, such as None or numpy.NaN. + If an entire row/column is NA values and `skipna` is True, + then the result will be True, as for an empty row/column. + If `skipna` is False, numpy.NaNs are treated as True because these are + not equal to zero, Nones are treated as False. + + Returns + ------- + Series + + Examples + -------- + Create a dataframe from a dictionary. + + >>> df = ps.DataFrame({ + ... 'col1': [True, True, True], + ... 'col2': [True, False, False], + ... 'col3': [0, 0, 0], + ... 'col4': [1, 2, 3], + ... 'col5': [True, True, None], + ... 'col6': [True, False, None]}, + ... columns=['col1', 'col2', 'col3', 'col4', 'col5', 'col6']) + + Default behavior checks if column-wise values all return True. + + >>> df.all() + col1 True + col2 False + col3 False + col4 True + col5 True + col6 False + dtype: bool + + Include NA values when set `skipna=False`. + + >>> df[['col5', 'col6']].all(skipna=False) + col5 False + col6 False + dtype: bool + + Include only boolean columns when set `bool_only=True`. + + >>> df.all(bool_only=True) + col1 True + col2 False + dtype: bool + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + column_labels = self._internal.column_labels + if bool_only: + column_labels = self._bool_column_labels(column_labels) + if len(column_labels) == 0: + return ps.Series([], dtype=bool) + + applied: List[PySparkColumn] = [] + for label in column_labels: + scol = self._internal.spark_column_for(label) + + if isinstance(self._internal.spark_type_for(label), NumericType) or skipna: + # np.nan takes no effect to the result; None takes no effect if `skipna` + all_col = F.min(F.coalesce(scol.cast("boolean"), F.lit(True))) + else: + # Take None as False when not `skipna` + all_col = F.min(F.when(scol.isNull(), F.lit(False)).otherwise(scol.cast("boolean"))) + applied.append(F.when(all_col.isNull(), True).otherwise(all_col)) + + return self._result_aggregated(column_labels, applied) + + # TODO: axis, skipna, level and **kwargs should be implemented. + def any(self, axis: Axis = 0, bool_only: Optional[bool] = None) -> "Series": + """ + Return whether any element is True. + + Returns False unless there is at least one element within a series that is + True or equivalent (e.g. non-zero or non-empty). + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Indicate which axis or axes should be reduced. + + * 0 / 'index' : reduce the index, return a Series whose index is the + original column labels. + + bool_only : bool, default None + Include only boolean columns. If None, will attempt to use everything, + then use only boolean data. + + Returns + ------- + Series + + Examples + -------- + Create a dataframe from a dictionary. + + >>> df = ps.DataFrame({ + ... 'col1': [False, False, False], + ... 'col2': [True, False, False], + ... 'col3': [0, 0, 1], + ... 'col4': [0, 1, 2], + ... 'col5': [False, False, None], + ... 'col6': [True, False, None]}, + ... columns=['col1', 'col2', 'col3', 'col4', 'col5', 'col6']) + + Default behavior checks if column-wise values all return True. + + >>> df.any() + col1 False + col2 True + col3 True + col4 True + col5 False + col6 True + dtype: bool + + Include only boolean columns when set `bool_only=True`. + + >>> df.any(bool_only=True) + col1 False + col2 True + dtype: bool + + Returns empty Series when the DataFrame is empty. + >>> df[[]].any() + Series([], dtype: bool) + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + column_labels = self._internal.column_labels + if bool_only: + column_labels = self._bool_column_labels(column_labels) + if len(column_labels) == 0: + return ps.Series([], dtype=bool) + + applied: List[PySparkColumn] = [] + for label in column_labels: + scol = self._internal.spark_column_for(label) + any_col = F.max(F.coalesce(scol.cast("boolean"), F.lit(False))) + applied.append(F.when(any_col.isNull(), False).otherwise(any_col)) + + return self._result_aggregated(column_labels, applied) + + def _bool_column_labels(self, column_labels: List[Label]) -> List[Label]: + """ + Filter column labels of boolean columns (without None). + """ + bool_column_labels = [] + for label in column_labels: + psser = self._psser_for(label) + if is_bool_dtype(psser): + # Rely on dtype rather than spark type because + # columns that consist of bools and Nones should be excluded + # if bool_only is True + bool_column_labels.append(label) + return bool_column_labels + + def _result_aggregated( + self, column_labels: List[Label], scols: Sequence[PySparkColumn] + ) -> "Series": + """ + Given aggregated Spark columns and respective column labels from the original + pandas-on-Spark DataFrame, construct the result Series. + """ + from pyspark.pandas.series import first_series + + cols = [] + result_scol_name = "value" + for label, applied_col in zip(column_labels, scols): + cols.append( + F.struct( + *[F.lit(col).alias(SPARK_INDEX_NAME_FORMAT(i)) for i, col in enumerate(label)], + *[applied_col.alias(result_scol_name)], + ) + ) + # Statements under this comment implement spark frame transformations as below: + # From: + # +-------------------------------------------------------------------------------------+ + # |arrays | + # +-------------------------------------------------------------------------------------+ + # |[{col1, true}, {col2, true}, {col3, false}, {col4, true}]| + # +-------------------------------------------------------------------------------------+ + # To: + # +-------------+ + # |col | + # +-------------+ + # |{col1, true} | + # |{col2, true} | + # |{col3, false}| + # |{col4, true} | + # +-------------+ + # To: + # +-----------------+-----+ + # |__index_level_0__|value| + # +-----------------+-----+ + # |col1 |true | + # |col2 |true | + # |col3 |false| + # |col4 |true | + # +-----------------+-----+ + sdf = self._internal.spark_frame.select(F.array(*cols).alias("arrays")).select( + F.explode(F.col("arrays")) + ) + sdf = sdf.selectExpr("col.*") + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, SPARK_INDEX_NAME_FORMAT(i)) + for i in range(self._internal.column_labels_level) + ], + index_names=self._internal.column_label_names, + column_labels=[None], + data_spark_columns=[scol_for(sdf, result_scol_name)], + ) + + # (cont.) The result Series should look as below: + # col1 False + # col2 True + # col3 True + # col4 True + # dtype: bool + return first_series(DataFrame(internal)) + + # TODO: add axis, pct, na_option parameter + def rank( + self, method: str = "average", ascending: bool = True, numeric_only: Optional[bool] = None + ) -> "DataFrame": + """ + Compute numerical data ranks (1 through n) along axis. Equal values are + assigned a rank that is the average of the ranks of those values. + + .. note:: the current implementation of rank uses Spark's Window without + specifying partition specification. This leads to moving all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + method : {'average', 'min', 'max', 'first', 'dense'} + * average: average rank of group + * min: lowest rank in group + * max: highest rank in group + * first: ranks assigned in order they appear in the array + * dense: like 'min', but rank always increases by 1 between groups + ascending : boolean, default True + False for ranks by high (1) to low (N) + numeric_only : bool, optional + For DataFrame objects, rank only numeric columns if set to True. + + Returns + ------- + ranks : same type as caller + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 2, 3], 'B': [4, 3, 2, 1]}, columns=['A', 'B']) + >>> df + A B + 0 1 4 + 1 2 3 + 2 2 2 + 3 3 1 + + >>> df.rank().sort_index() + A B + 0 1.0 4.0 + 1 2.5 3.0 + 2 2.5 2.0 + 3 4.0 1.0 + + If method is set to 'min', it uses lowest rank in group. + + >>> df.rank(method='min').sort_index() + A B + 0 1.0 4.0 + 1 2.0 3.0 + 2 2.0 2.0 + 3 4.0 1.0 + + If method is set to 'max', it uses highest rank in group. + + >>> df.rank(method='max').sort_index() + A B + 0 1.0 4.0 + 1 3.0 3.0 + 2 3.0 2.0 + 3 4.0 1.0 + + If method is set to 'dense', it leaves no gaps in group. + + >>> df.rank(method='dense').sort_index() + A B + 0 1.0 4.0 + 1 2.0 3.0 + 2 2.0 2.0 + 3 3.0 1.0 + + If numeric_only is set to 'True', rank only numeric columns. + + >>> df = ps.DataFrame({'A': [1, 2, 2, 3], 'B': ['a', 'b', 'd', 'c']}, columns= ['A', 'B']) + >>> df + A B + 0 1 a + 1 2 b + 2 2 d + 3 3 c + >>> df.rank(numeric_only=True) + A + 0 1.0 + 1 2.5 + 2 2.5 + 3 4.0 + """ + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `None` in 4.0.0.", + FutureWarning, + ) + if numeric_only: + numeric_col_names = [] + for label in self._internal.column_labels: + psser = self._psser_for(label) + if isinstance(psser.spark.data_type, (NumericType, BooleanType)): + numeric_col_names.append(psser.name) + + psdf = self[numeric_col_names] if numeric_only else self + return psdf._apply_series_op( + lambda psser: psser._rank(method=method, ascending=ascending), should_resolve=True + ) + + def filter( + self, + items: Optional[Sequence[Any]] = None, + like: Optional[str] = None, + regex: Optional[str] = None, + axis: Optional[Axis] = None, + ) -> "DataFrame": + """ + Subset rows or columns of dataframe according to labels in + the specified index. + + Note that this routine does not filter a dataframe on its + contents. The filter is applied to the labels of the index. + + Parameters + ---------- + items : list-like + Keep labels from axis which are in items. + like : string + Keep labels from axis for which "like in label == True". + regex : string (regular expression) + Keep labels from axis for which re.search(regex, label) == True. + axis : int or string axis name + The axis to filter on. By default this is the info axis, + 'index' for Series, 'columns' for DataFrame. + + Returns + ------- + same type as input object + + See Also + -------- + DataFrame.loc + + Notes + ----- + The ``items``, ``like``, and ``regex`` parameters are + enforced to be mutually exclusive. + + ``axis`` defaults to the info axis that is used when indexing + with ``[]``. + + Examples + -------- + >>> df = ps.DataFrame(np.array(([1, 2, 3], [4, 5, 6])), + ... index=['mouse', 'rabbit'], + ... columns=['one', 'two', 'three']) + + >>> # select columns by name + >>> df.filter(items=['one', 'three']) + one three + mouse 1 3 + rabbit 4 6 + + >>> # select columns by regular expression + >>> df.filter(regex='e$', axis=1) + one three + mouse 1 3 + rabbit 4 6 + + >>> # select rows containing 'bbi' + >>> df.filter(like='bbi', axis=0) + one two three + rabbit 4 5 6 + + For a Series, + + >>> # select rows by name + >>> df.one.filter(items=['rabbit']) + rabbit 4 + Name: one, dtype: int64 + + >>> # select rows by regular expression + >>> df.one.filter(regex='e$') + mouse 1 + Name: one, dtype: int64 + + >>> # select rows containing 'bbi' + >>> df.one.filter(like='bbi') + rabbit 4 + Name: one, dtype: int64 + """ + if sum(x is not None for x in (items, like, regex)) > 1: + raise TypeError( + "Keyword arguments `items`, `like`, or `regex` " "are mutually exclusive" + ) + + axis = validate_axis(axis, none_axis=1) + + index_scols = self._internal.index_spark_columns + + if items is not None: + if is_list_like(items): + items = list(items) + else: + raise ValueError("items should be a list-like object.") + if axis == 0: + if len(index_scols) == 1: + if len(items) <= ps.get_option("compute.isin_limit"): + col = index_scols[0].isin([F.lit(item) for item in items]) + return DataFrame(self._internal.with_filter(col)) + else: + item_sdf_col = verify_temp_column_name( + self._internal.spark_frame, "__item__" + ) + item_sdf = default_session().createDataFrame( + pd.DataFrame({item_sdf_col: items}) + ) + joined_sdf = self._internal.spark_frame.join( + other=F.broadcast(item_sdf), + on=(index_scols[0] == scol_for(item_sdf, item_sdf_col)), + how="semi", + ) + + return DataFrame(self._internal.with_new_sdf(joined_sdf)) + + else: + # for multi-index + col = None + for item in items: + if not isinstance(item, tuple): + raise TypeError("Unsupported type {}".format(type(item).__name__)) + if not item: + raise ValueError("The item should not be empty.") + midx_col = None + for i, element in enumerate(item): + if midx_col is None: + midx_col = index_scols[i] == F.lit(element) + else: + midx_col = midx_col & (index_scols[i] == F.lit(element)) + if col is None: + col = midx_col + else: + col = col | midx_col + return DataFrame(self._internal.with_filter(col)) + else: + return self[items] + elif like is not None: + if axis == 0: + col = None + for index_scol in index_scols: + if col is None: + col = index_scol.contains(like) + else: + col = col | index_scol.contains(like) + return DataFrame(self._internal.with_filter(col)) + else: + column_labels = self._internal.column_labels + output_labels = [label for label in column_labels if any(like in i for i in label)] + return self[output_labels] + elif regex is not None: + if axis == 0: + col = None + for index_scol in index_scols: + if col is None: + col = index_scol.rlike(regex) + else: + col = col | index_scol.rlike(regex) + return DataFrame(self._internal.with_filter(col)) + else: + column_labels = self._internal.column_labels + matcher = re.compile(regex) + output_labels = [ + label + for label in column_labels + if any(matcher.search(i) is not None for i in label) + ] + return self[output_labels] + else: + raise TypeError("Must pass either `items`, `like`, or `regex`") + + def rename( + self, + mapper: Optional[Union[Dict, Callable[[Any], Any]]] = None, + index: Optional[Union[Dict, Callable[[Any], Any]]] = None, + columns: Optional[Union[Dict, Callable[[Any], Any]]] = None, + axis: Axis = "index", + inplace: bool = False, + level: Optional[int] = None, + errors: str = "ignore", + ) -> Optional["DataFrame"]: + + """ + Alter axes labels. + Function / dict values must be unique (1-to-1). Labels not contained in a dict / Series + will be left as-is. Extra labels listed don’t throw an error. + + Parameters + ---------- + mapper : dict-like or function + Dict-like or functions transformations to apply to that axis’ values. + Use either `mapper` and `axis` to specify the axis to target with `mapper`, or `index` + and `columns`. + index : dict-like or function + Alternative to specifying axis ("mapper, axis=0" is equivalent to "index=mapper"). + columns : dict-like or function + Alternative to specifying axis ("mapper, axis=1" is equivalent to "columns=mapper"). + axis : int or str, default 'index' + Axis to target with mapper. Can be either the axis name ('index', 'columns') or + number (0, 1). + inplace : bool, default False + Whether to return a new DataFrame. + level : int or level name, default None + In case of a MultiIndex, only rename labels in the specified level. + errors : {'ignore', 'raise'}, default 'ignore' + If 'raise', raise a `KeyError` when a dict-like `mapper`, `index`, or `columns` + contains labels that are not present in the Index being transformed. If 'ignore', + existing keys will be renamed, and extra keys will be ignored. + + Returns + ------- + DataFrame with the renamed axis labels. + + Raises + ------ + `KeyError` + If any of the labels is not found in the selected axis and "errors='raise'". + + Examples + -------- + >>> psdf1 = ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) + >>> psdf1.rename(columns={"A": "a", "B": "c"}) # doctest: +NORMALIZE_WHITESPACE + a c + 0 1 4 + 1 2 5 + 2 3 6 + + >>> psdf1.rename(index={1: 10, 2: 20}) # doctest: +NORMALIZE_WHITESPACE + A B + 0 1 4 + 10 2 5 + 20 3 6 + + >>> psdf1.rename(columns={"A": "a", "C": "c"}, errors="raise") + Traceback (most recent call last): + ... + KeyError: 'Index include value which is not in the `mapper`' + + >>> def str_lower(s) -> str: + ... return str.lower(s) + >>> psdf1.rename(str_lower, axis='columns') # doctest: +NORMALIZE_WHITESPACE + a b + 0 1 4 + 1 2 5 + 2 3 6 + + >>> def mul10(x) -> int: + ... return x * 10 + >>> psdf1.rename(mul10, axis='index') # doctest: +NORMALIZE_WHITESPACE + A B + 0 1 4 + 10 2 5 + 20 3 6 + + >>> idx = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C'), ('Y', 'D')]) + >>> psdf2 = ps.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], columns=idx) + >>> psdf2.rename(columns=str_lower, level=0) # doctest: +NORMALIZE_WHITESPACE + x y + A B C D + 0 1 2 3 4 + 1 5 6 7 8 + + >>> psdf3 = ps.DataFrame([[1, 2], [3, 4], [5, 6], [7, 8]], index=idx, columns=list('ab')) + >>> psdf3.rename(index=str_lower) # doctest: +NORMALIZE_WHITESPACE + a b + x a 1 2 + b 3 4 + y c 5 6 + d 7 8 + """ + + def gen_mapper_fn( + mapper: Union[Dict, Callable[[Any], Any]], skip_return_type: bool = False + ) -> Tuple[Callable[[Any], Any], Dtype, DataType]: + if isinstance(mapper, dict): + mapper_dict = mapper + + type_set = set(map(lambda x: type(x), mapper_dict.values())) + if len(type_set) > 1: + raise ValueError("Mapper dict should have the same value type.") + dtype, spark_return_type = pandas_on_spark_type(list(type_set)[0]) + + def mapper_fn(x: Any) -> Any: + if x in mapper_dict: + return mapper_dict[x] + else: + if errors == "raise": + raise KeyError("Index include value which is not in the `mapper`") + return x + + return mapper_fn, dtype, spark_return_type + elif callable(mapper): + mapper_callable = cast(Callable, mapper) + + def mapper_fn(x: Any) -> Any: + return mapper_callable(x) + + if skip_return_type: + return mapper_fn, None, None + else: + return_type = cast(ScalarType, infer_return_type(mapper)) + dtype = return_type.dtype + spark_return_type = return_type.spark_type + return mapper_fn, dtype, spark_return_type + else: + raise ValueError( + "`mapper` or `index` or `columns` should be " + "either dict-like or function type." + ) + + index_mapper_fn = None + index_mapper_ret_stype = None + columns_mapper_fn = None + + inplace = validate_bool_kwarg(inplace, "inplace") + if mapper: + axis = validate_axis(axis) + if axis == 0: + index_mapper_fn, index_mapper_ret_dtype, index_mapper_ret_stype = gen_mapper_fn( + mapper + ) + elif axis == 1: + columns_mapper_fn, _, _ = gen_mapper_fn(mapper) + else: + if index: + index_mapper_fn, index_mapper_ret_dtype, index_mapper_ret_stype = gen_mapper_fn( + index + ) + if columns: + columns_mapper_fn, _, _ = gen_mapper_fn(columns, skip_return_type=True) + + if not index and not columns: + raise ValueError("Either `index` or `columns` should be provided.") + + psdf = self.copy() + if index_mapper_fn: + # rename index labels, if `level` is None, rename all index columns, otherwise only + # rename the corresponding level index. + # implement this by transform the underlying spark dataframe, + # Example: + # suppose the psdf index column in underlying spark dataframe is "index_0", "index_1", + # if rename level 0 index labels, will do: + # ``psdf._sdf.withColumn("index_0", mapper_fn_udf(col("index_0"))`` + # if rename all index labels (`level` is None), then will do: + # ``` + # psdf._sdf.withColumn("index_0", mapper_fn_udf(col("index_0")) + # .withColumn("index_1", mapper_fn_udf(col("index_1")) + # ``` + + index_columns = psdf._internal.index_spark_column_names + num_indices = len(index_columns) + if level is not None and (level < 0 or level >= num_indices): + raise ValueError("level should be an integer between [0, %s)" % num_indices) + + @pandas_udf(returnType=index_mapper_ret_stype) # type: ignore[call-overload] + def index_mapper_udf(s: pd.Series) -> pd.Series: + return s.map(index_mapper_fn) + + index_spark_columns = psdf._internal.index_spark_columns.copy() + index_fields = psdf._internal.index_fields.copy() + if level is None: + for i in range(num_indices): + index_spark_columns[i] = index_mapper_udf(index_spark_columns[i]).alias( + index_columns[i] + ) + index_fields[i] = index_fields[i].copy( + dtype=index_mapper_ret_dtype, + spark_type=index_mapper_ret_stype, + nullable=True, + ) + else: + index_spark_columns[level] = index_mapper_udf(index_spark_columns[level]).alias( + index_columns[level] + ) + index_fields[level] = index_fields[level].copy( + dtype=index_mapper_ret_dtype, + spark_type=index_mapper_ret_stype, + nullable=True, + ) + psdf = DataFrame( + psdf._internal.copy( + index_spark_columns=index_spark_columns, index_fields=index_fields + ) + ) + if columns_mapper_fn: + # rename column name. + # Will modify the `_internal._column_labels` and transform underlying spark dataframe + # to the same column name with `_internal._column_labels`. + if level: + if level < 0 or level >= psdf._internal.column_labels_level: + raise ValueError("level should be an integer between [0, column_labels_level)") + + def gen_new_column_labels_entry(column_labels_entry: Label) -> Label: + if level is None: + # rename all level columns + return tuple(map(columns_mapper_fn, column_labels_entry)) + else: + # only rename specified level column + entry_list = list(column_labels_entry) + entry_list[level] = columns_mapper_fn(entry_list[level]) + return tuple(entry_list) + + new_column_labels = list(map(gen_new_column_labels_entry, psdf._internal.column_labels)) + + new_data_pssers = [ + psdf._psser_for(old_label).rename(new_label) + for old_label, new_label in zip(psdf._internal.column_labels, new_column_labels) + ] + psdf = DataFrame(psdf._internal.with_new_columns(new_data_pssers)) + if inplace: + self._update_internal_frame(psdf._internal) + return None + else: + return psdf + + def rename_axis( + self, + mapper: Union[Any, Sequence[Any], Dict[Name, Any], Callable[[Name], Any]] = None, + index: Union[Any, Sequence[Any], Dict[Name, Any], Callable[[Name], Any]] = None, + columns: Union[Any, Sequence[Any], Dict[Name, Any], Callable[[Name], Any]] = None, + axis: Optional[Axis] = 0, + inplace: Optional[bool] = False, + ) -> Optional["DataFrame"]: + """ + Set the name of the axis for the index or columns. + + Parameters + ---------- + mapper : scalar, list-like, optional + A scalar, list-like, dict-like or functions transformations to + apply to the axis name attribute. + index, columns : scalar, list-like, dict-like or function, optional + A scalar, list-like, dict-like or functions transformations to + apply to that axis' values. + + Use either ``mapper`` and ``axis`` to + specify the axis to target with ``mapper``, or ``index`` + and/or ``columns``. + axis : {0 or 'index', 1 or 'columns'}, default 0 + The axis to rename. + inplace : bool, default False + Modifies the object directly, instead of creating a new DataFrame. + + Returns + ------- + DataFrame, or None if `inplace` is True. + + See Also + -------- + Series.rename : Alter Series index labels or name. + DataFrame.rename : Alter DataFrame index labels or name. + Index.rename : Set new names on index. + + Notes + ----- + ``DataFrame.rename_axis`` supports two calling conventions + + * ``(index=index_mapper, columns=columns_mapper, ...)`` + * ``(mapper, axis={'index', 'columns'}, ...)`` + + The first calling convention will only modify the names of + the index and/or the names of the Index object that is the columns. + + The second calling convention will modify the names of the + corresponding index specified by axis. + + We *highly* recommend using keyword arguments to clarify your + intent. + + Examples + -------- + >>> df = ps.DataFrame({"num_legs": [4, 4, 2], + ... "num_arms": [0, 0, 2]}, + ... index=["dog", "cat", "monkey"], + ... columns=["num_legs", "num_arms"]) + >>> df + num_legs num_arms + dog 4 0 + cat 4 0 + monkey 2 2 + + >>> df = df.rename_axis("animal").sort_index() + >>> df # doctest: +NORMALIZE_WHITESPACE + num_legs num_arms + animal + cat 4 0 + dog 4 0 + monkey 2 2 + + >>> df = df.rename_axis("limbs", axis="columns").sort_index() + >>> df # doctest: +NORMALIZE_WHITESPACE + limbs num_legs num_arms + animal + cat 4 0 + dog 4 0 + monkey 2 2 + + **MultiIndex** + + >>> index = pd.MultiIndex.from_product([['mammal'], + ... ['dog', 'cat', 'monkey']], + ... names=['type', 'name']) + >>> df = ps.DataFrame({"num_legs": [4, 4, 2], + ... "num_arms": [0, 0, 2]}, + ... index=index, + ... columns=["num_legs", "num_arms"]) + >>> df # doctest: +NORMALIZE_WHITESPACE + num_legs num_arms + type name + mammal dog 4 0 + cat 4 0 + monkey 2 2 + + >>> df.rename_axis(index={'type': 'class'}).sort_index() # doctest: +NORMALIZE_WHITESPACE + num_legs num_arms + class name + mammal cat 4 0 + dog 4 0 + monkey 2 2 + + >>> df.rename_axis(index=str.upper).sort_index() # doctest: +NORMALIZE_WHITESPACE + num_legs num_arms + TYPE NAME + mammal cat 4 0 + dog 4 0 + monkey 2 2 + """ + + def gen_names( + v: Union[Any, Sequence[Any], Dict[Name, Any], Callable[[Name], Any]], + curnames: List[Name], + ) -> List[Label]: + newnames: List[Name] + if is_scalar(v): + newnames = [cast(Name, v)] + elif is_list_like(v) and not is_dict_like(v): + newnames = list(cast(Sequence[Name], v)) + elif is_dict_like(v): + v_dict = cast(Dict[Name, Name], v) + newnames = [v_dict[name] if name in v_dict else name for name in curnames] + elif callable(v): + v_callable = cast(Callable[[Name], Name], v) + newnames = [v_callable(name) for name in curnames] + else: + raise ValueError( + "`mapper` or `index` or `columns` should be " + "either dict-like or function type." + ) + + if len(newnames) != len(curnames): + raise ValueError( + "Length of new names must be {}, got {}".format(len(curnames), len(newnames)) + ) + + return [name if is_name_like_tuple(name) else (name,) for name in newnames] + + if mapper is not None and (index is not None or columns is not None): + raise TypeError("Cannot specify both 'mapper' and any of 'index' or 'columns'.") + + if mapper is not None: + axis = validate_axis(axis) + if axis == 0: + index = mapper + elif axis == 1: + columns = mapper + + column_label_names = ( + gen_names(columns, self.columns.names) + if columns is not None + else self._internal.column_label_names + ) + index_names = ( + gen_names(index, self.index.names) if index is not None else self._internal.index_names + ) + + internal = self._internal.copy( + index_names=index_names, column_label_names=column_label_names + ) + if inplace: + self._update_internal_frame(internal) + return None + else: + return DataFrame(internal) + + def keys(self) -> pd.Index: + """ + Return alias for columns. + + Returns + ------- + Index + Columns of the DataFrame. + + Examples + -------- + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=['cobra', 'viper', 'sidewinder'], + ... columns=['max_speed', 'shield']) + >>> df + max_speed shield + cobra 1 2 + viper 4 5 + sidewinder 7 8 + + >>> df.keys() + Index(['max_speed', 'shield'], dtype='object') + """ + return self.columns + + def pct_change(self, periods: int = 1) -> "DataFrame": + """ + Percentage change between the current and a prior element. + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to moving all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int, default 1 + Periods to shift for forming percent change. + + Returns + ------- + DataFrame + + Examples + -------- + Percentage change in French franc, Deutsche Mark, and Italian lira + from 1980-01-01 to 1980-03-01. + + >>> df = ps.DataFrame({ + ... 'FR': [4.0405, 4.0963, 4.3149], + ... 'GR': [1.7246, 1.7482, 1.8519], + ... 'IT': [804.74, 810.01, 860.13]}, + ... index=['1980-01-01', '1980-02-01', '1980-03-01']) + >>> df + FR GR IT + 1980-01-01 4.0405 1.7246 804.74 + 1980-02-01 4.0963 1.7482 810.01 + 1980-03-01 4.3149 1.8519 860.13 + + >>> df.pct_change() + FR GR IT + 1980-01-01 NaN NaN NaN + 1980-02-01 0.013810 0.013684 0.006549 + 1980-03-01 0.053365 0.059318 0.061876 + + You can set periods to shift for forming percent change + + >>> df.pct_change(2) + FR GR IT + 1980-01-01 NaN NaN NaN + 1980-02-01 NaN NaN NaN + 1980-03-01 0.067912 0.073814 0.06883 + """ + window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-periods, -periods) + + def op(psser: ps.Series) -> PySparkColumn: + prev_row = F.lag(psser.spark.column, periods).over(window) + return ((psser.spark.column - prev_row) / prev_row).alias( + psser._internal.data_spark_column_names[0] + ) + + return self._apply_series_op(op, should_resolve=True) + + # TODO: axis = 1 + def idxmax(self, axis: Axis = 0) -> "Series": + """ + Return index of first occurrence of maximum over requested axis. + NA/null values are excluded. + + .. note:: This API collect all rows with maximum value using `to_pandas()` + because we suppose the number of rows with max values are usually small in general. + + Parameters + ---------- + axis : 0 or 'index' + Can only be set to 0 now. + + Returns + ------- + Series + + See Also + -------- + Series.idxmax + + Examples + -------- + >>> psdf = ps.DataFrame({'a': [1, 2, 3, 2], + ... 'b': [4.0, 2.0, 3.0, 1.0], + ... 'c': [300, 200, 400, 200]}) + >>> psdf + a b c + 0 1 4.0 300 + 1 2 2.0 200 + 2 3 3.0 400 + 3 2 1.0 200 + + >>> psdf.idxmax() + a 2 + b 0 + c 2 + dtype: int64 + + For Multi-column Index + + >>> psdf = ps.DataFrame({'a': [1, 2, 3, 2], + ... 'b': [4.0, 2.0, 3.0, 1.0], + ... 'c': [300, 200, 400, 200]}) + >>> psdf.columns = pd.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> psdf + a b c + x y z + 0 1 4.0 300 + 1 2 2.0 200 + 2 3 3.0 400 + 3 2 1.0 200 + + >>> psdf.idxmax() + a x 2 + b y 0 + c z 2 + dtype: int64 + """ + max_cols = map(lambda scol: F.max(scol), self._internal.data_spark_columns) + sdf_max = self._internal.spark_frame.select(*max_cols).head() + # `sdf_max` looks like below + # +------+------+------+ + # |(a, x)|(b, y)|(c, z)| + # +------+------+------+ + # | 3| 4.0| 400| + # +------+------+------+ + + conds = ( + scol == max_val for scol, max_val in zip(self._internal.data_spark_columns, sdf_max) + ) + cond = reduce(lambda x, y: x | y, conds) + + psdf: DataFrame = DataFrame(self._internal.with_filter(cond)) + + return cast(ps.Series, ps.from_pandas(psdf._to_internal_pandas().idxmax())) + + # TODO: axis = 1 + def idxmin(self, axis: Axis = 0) -> "Series": + """ + Return index of first occurrence of minimum over requested axis. + NA/null values are excluded. + + .. note:: This API collect all rows with minimum value using `to_pandas()` + because we suppose the number of rows with min values are usually small in general. + + Parameters + ---------- + axis : 0 or 'index' + Can only be set to 0 now. + + Returns + ------- + Series + + See Also + -------- + Series.idxmin + + Examples + -------- + >>> psdf = ps.DataFrame({'a': [1, 2, 3, 2], + ... 'b': [4.0, 2.0, 3.0, 1.0], + ... 'c': [300, 200, 400, 200]}) + >>> psdf + a b c + 0 1 4.0 300 + 1 2 2.0 200 + 2 3 3.0 400 + 3 2 1.0 200 + + >>> psdf.idxmin() + a 0 + b 3 + c 1 + dtype: int64 + + For Multi-column Index + + >>> psdf = ps.DataFrame({'a': [1, 2, 3, 2], + ... 'b': [4.0, 2.0, 3.0, 1.0], + ... 'c': [300, 200, 400, 200]}) + >>> psdf.columns = pd.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> psdf + a b c + x y z + 0 1 4.0 300 + 1 2 2.0 200 + 2 3 3.0 400 + 3 2 1.0 200 + + >>> psdf.idxmin() + a x 0 + b y 3 + c z 1 + dtype: int64 + """ + min_cols = map(lambda scol: F.min(scol), self._internal.data_spark_columns) + sdf_min = self._internal.spark_frame.select(*min_cols).head() + + conds = ( + scol == min_val for scol, min_val in zip(self._internal.data_spark_columns, sdf_min) + ) + cond = reduce(lambda x, y: x | y, conds) + + psdf: DataFrame = DataFrame(self._internal.with_filter(cond)) + + return cast(ps.Series, ps.from_pandas(psdf._to_internal_pandas().idxmin())) + + # TODO(SPARK-41619): Add `show_counts` parameter and replace with `null_counts`. + def info( + self, + verbose: Optional[bool] = None, + buf: Optional[IO[str]] = None, + max_cols: Optional[int] = None, + ) -> None: + """ + Print a concise summary of a DataFrame. + + This method prints information about a DataFrame including + the index dtype and column dtypes, non-null values and memory usage. + + Parameters + ---------- + verbose : bool, optional + Whether to print the full summary. + buf : writable buffer, defaults to sys.stdout + Where to send the output. By default the output is printed to + sys.stdout. Pass a writable buffer if you need to further process + the output. + max_cols : int, optional + When to switch from the verbose to the truncated output. If the + DataFrame has more than `max_cols` columns, the truncated output + is used. + null_counts : bool, optional + Whether to show the non-null counts. + + .. deprecated:: 3.4.0 + + Returns + ------- + None + This method prints a summary of a DataFrame and returns None. + + See Also + -------- + DataFrame.describe: Generate descriptive statistics of DataFrame + columns. + + Examples + -------- + >>> int_values = [1, 2, 3, 4, 5] + >>> text_values = ['alpha', 'beta', 'gamma', 'delta', 'epsilon'] + >>> float_values = [0.0, 0.25, 0.5, 0.75, 1.0] + >>> df = ps.DataFrame( + ... {"int_col": int_values, "text_col": text_values, "float_col": float_values}, + ... columns=['int_col', 'text_col', 'float_col']) + >>> df + int_col text_col float_col + 0 1 alpha 0.00 + 1 2 beta 0.25 + 2 3 gamma 0.50 + 3 4 delta 0.75 + 4 5 epsilon 1.00 + + Prints information of all columns: + + >>> df.info(verbose=True) # doctest: +SKIP + + Index: 5 entries, 0 to 4 + Data columns (total 3 columns): + # Column Non-Null Count Dtype + --- ------ -------------- ----- + 0 int_col 5 non-null int64 + 1 text_col 5 non-null object + 2 float_col 5 non-null float64 + dtypes: float64(1), int64(1), object(1) + + Prints a summary of columns count and its dtypes but not per column + information: + + >>> df.info(verbose=False) # doctest: +SKIP + + Index: 5 entries, 0 to 4 + Columns: 3 entries, int_col to float_col + dtypes: float64(1), int64(1), object(1) + + Pipe output of DataFrame.info to buffer instead of sys.stdout, get + buffer content and writes to a text file: + + >>> import io + >>> buffer = io.StringIO() + >>> df.info(buf=buffer) + >>> s = buffer.getvalue() + >>> with open('%s/info.txt' % path, "w", + ... encoding="utf-8") as f: + ... _ = f.write(s) + >>> with open('%s/info.txt' % path) as f: + ... f.readlines() # doctest: +SKIP + ["\\n", + 'Index: 5 entries, 0 to 4\\n', + 'Data columns (total 3 columns):\\n', + ' # Column Non-Null Count Dtype \\n', + '--- ------ -------------- ----- \\n', + ' 0 int_col 5 non-null int64 \\n', + ' 1 text_col 5 non-null object \\n', + ' 2 float_col 5 non-null float64\\n', + 'dtypes: float64(1), int64(1), object(1)'] + """ + # To avoid pandas' existing config affects pandas-on-Spark. + # TODO: should we have corresponding pandas-on-Spark configs? + with pd.option_context( + "display.max_info_columns", sys.maxsize, "display.max_info_rows", sys.maxsize + ): + try: + # hack to use pandas' info as is. + object.__setattr__(self, "_data", self) + count_func = self.count + self.count = ( # type: ignore[assignment] + lambda: count_func()._to_pandas() # type: ignore[assignment, misc, union-attr] + ) + return pd.DataFrame.info( + self, # type: ignore[arg-type] + verbose=verbose, + buf=buf, + max_cols=max_cols, + memory_usage=False, + ) + finally: + del self._data + self.count = count_func # type: ignore[assignment] + + # TODO: fix parameter 'axis' and 'numeric_only' to work same as pandas' + def quantile( + self, + q: Union[float, Iterable[float]] = 0.5, + axis: Axis = 0, + numeric_only: bool = True, + accuracy: int = 10000, + ) -> DataFrameOrSeries: + """ + Return value at the given quantile. + + .. note:: Unlike pandas', the quantile in pandas-on-Spark is an approximated quantile + based upon approximate percentile computation because computing quantile across a + large dataset is extremely expensive. + + Parameters + ---------- + q : float or array-like, default 0.5 (50% quantile) + 0 <= q <= 1, the quantile(s) to compute. + axis : int or str, default 0 or 'index' + Can only be set to 0 now. + numeric_only : bool, default True + If False, the quantile of datetime and time delta data will be computed as well. + Can only be set to True now. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + + Returns + ------- + Series or DataFrame + If q is an array, a DataFrame will be returned where the + index is q, the columns are the columns of self, and the values are the quantiles. + If q is a float, a Series will be returned where the + index is the columns of self and the values are the quantiles. + + Examples + -------- + >>> psdf = ps.DataFrame({'a': [1, 2, 3, 4, 5], 'b': [6, 7, 8, 9, 0]}) + >>> psdf + a b + 0 1 6 + 1 2 7 + 2 3 8 + 3 4 9 + 4 5 0 + + >>> psdf.quantile(.5) + a 3.0 + b 7.0 + Name: 0.5, dtype: float64 + + >>> psdf.quantile([.25, .5, .75]) + a b + 0.25 2.0 6.0 + 0.50 3.0 7.0 + 0.75 4.0 8.0 + """ + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + if not isinstance(accuracy, int): + raise TypeError( + "accuracy must be an integer; however, got [%s]" % type(accuracy).__name__ + ) + + qq: Union[float, List[float]] = list(q) if isinstance(q, Iterable) else q + + for v in qq if isinstance(qq, list) else [qq]: + if not isinstance(v, float): + raise TypeError( + "q must be a float or an array of floats; however, [%s] found." % type(v) + ) + if v < 0.0 or v > 1.0: + raise ValueError("percentiles should all be in the interval [0, 1].") + + def quantile(psser: "Series") -> PySparkColumn: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, (BooleanType, NumericType)): + return F.percentile_approx(spark_column.cast(DoubleType()), qq, accuracy) + else: + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + if isinstance(qq, list): + # First calculate the percentiles from all columns and map it to each `quantiles` + # by creating each entry as a struct. So, it becomes an array of structs as below: + # + # +-----------------------------------------+ + # | arrays| + # +-----------------------------------------+ + # |[[0.25, 2, 6], [0.5, 3, 7], [0.75, 4, 8]]| + # +-----------------------------------------+ + + percentile_cols: List[PySparkColumn] = [] + percentile_col_names: List[str] = [] + column_labels: List[Label] = [] + for label, column in zip( + self._internal.column_labels, self._internal.data_spark_column_names + ): + psser = self._psser_for(label) + + is_numeric_or_boolean = isinstance( + psser.spark.data_type, (NumericType, BooleanType) + ) + keep_column = not numeric_only or is_numeric_or_boolean + + if keep_column: + percentile_col = quantile(psser) + percentile_cols.append(percentile_col.alias(column)) + percentile_col_names.append(column) + column_labels.append(label) + + if len(percentile_cols) == 0: + return DataFrame(index=qq) + + sdf = self._internal.spark_frame.select(percentile_cols) + # Here, after select percentile cols, a spark_frame looks like below: + # +---------+---------+ + # | a| b| + # +---------+---------+ + # |[2, 3, 4]|[6, 7, 8]| + # +---------+---------+ + + cols_dict: Dict[str, List[PySparkColumn]] = {} + for column in percentile_col_names: + cols_dict[column] = list() + for i in range(len(qq)): + cols_dict[column].append(scol_for(sdf, column)[i].alias(column)) + + internal_index_column = SPARK_DEFAULT_INDEX_NAME + cols = [] + for i, col in enumerate(zip(*cols_dict.values())): + cols.append(F.struct(F.lit(qq[i]).alias(internal_index_column), *col)) + sdf = sdf.select(F.array(*cols).alias("arrays")) + + # And then, explode it and manually set the index. + # +-----------------+---+---+ + # |__index_level_0__| a| b| + # +-----------------+---+---+ + # | 0.25| 2| 6| + # | 0.5| 3| 7| + # | 0.75| 4| 8| + # +-----------------+---+---+ + sdf = sdf.select(F.explode(F.col("arrays"))).selectExpr("col.*") + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, internal_index_column)], + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in percentile_col_names], + ) + return DataFrame(internal) + else: + return self._reduce_for_stat_function( + quantile, name="quantile", numeric_only=numeric_only + ).rename(qq) + + def query(self, expr: str, inplace: bool = False) -> Optional["DataFrame"]: + """ + Query the columns of a DataFrame with a boolean expression. + + .. note:: Internal columns that starting with a '__' prefix are able to access, however, + they are not supposed to be accessed. + + .. note:: This API delegates to Spark SQL so the syntax follows Spark SQL. Therefore, the + pandas specific syntax such as `@` is not supported. If you want the pandas syntax, + you can work around with :meth:`DataFrame.pandas_on_spark.apply_batch`, but you should + be aware that `query_func` will be executed at different nodes in a distributed manner. + So, for example to use `@` syntax, make sure the variable is serialized by + putting it within the closure as below. + + >>> df = ps.DataFrame({'A': range(2000), 'B': range(2000)}) + >>> def query_func(pdf): + ... num = 1995 + ... return pdf.query('A > @num') + >>> df.pandas_on_spark.apply_batch(query_func) + A B + 1996 1996 1996 + 1997 1997 1997 + 1998 1998 1998 + 1999 1999 1999 + + Parameters + ---------- + expr : str + The query string to evaluate. + + You can refer to column names that contain spaces by surrounding + them in backticks. + + For example, if one of your columns is called ``a a`` and you want + to sum it with ``b``, your query should be ```a a` + b``. + + inplace : bool + Whether the query should modify the data in place or return + a modified copy. + + Returns + ------- + DataFrame + DataFrame resulting from the provided query expression. + + Examples + -------- + >>> df = ps.DataFrame({'A': range(1, 6), + ... 'B': range(10, 0, -2), + ... 'C C': range(10, 5, -1)}) + >>> df + A B C C + 0 1 10 10 + 1 2 8 9 + 2 3 6 8 + 3 4 4 7 + 4 5 2 6 + + >>> df.query('A > B') + A B C C + 4 5 2 6 + + The previous expression is equivalent to + + >>> df[df.A > df.B] + A B C C + 4 5 2 6 + + For columns with spaces in their name, you can use backtick quoting. + + >>> df.query('B == `C C`') + A B C C + 0 1 10 10 + + The previous expression is equivalent to + + >>> df[df.B == df['C C']] + A B C C + 0 1 10 10 + """ + if isinstance(self.columns, pd.MultiIndex): + raise TypeError("Doesn't support for MultiIndex columns") + if not isinstance(expr, str): + raise TypeError( + "expr must be a string to be evaluated, {} given".format(type(expr).__name__) + ) + inplace = validate_bool_kwarg(inplace, "inplace") + + data_columns = [label[0] for label in self._internal.column_labels] + sdf = self._internal.spark_frame.select( + self._internal.index_spark_columns + + [ + scol.alias(col) + for scol, col in zip(self._internal.data_spark_columns, data_columns) + ] + ).filter(expr) + internal = self._internal.with_new_sdf(sdf, data_columns=data_columns) + + if inplace: + self._update_internal_frame(internal) + return None + else: + return DataFrame(internal) + + def take(self, indices: List[int], axis: Axis = 0, **kwargs: Any) -> "DataFrame": + """ + Return the elements in the given *positional* indices along an axis. + + This means that we are not indexing according to actual values in + the index attribute of the object. We are indexing according to the + actual position of the element in the object. + + Parameters + ---------- + indices : array-like + An array of ints indicating which positions to take. + axis : {0 or 'index', 1 or 'columns', None}, default 0 + The axis on which to select elements. ``0`` means that we are + selecting rows, ``1`` means that we are selecting columns. + **kwargs + For compatibility with :meth:`numpy.take`. Has no effect on the + output. + + Returns + ------- + taken : same type as caller + An array-like containing the elements taken from the object. + + See Also + -------- + DataFrame.loc : Select a subset of a DataFrame by labels. + DataFrame.iloc : Select a subset of a DataFrame by positions. + numpy.take : Take elements from an array along an axis. + + Examples + -------- + >>> df = ps.DataFrame([('falcon', 'bird', 389.0), + ... ('parrot', 'bird', 24.0), + ... ('lion', 'mammal', 80.5), + ... ('monkey', 'mammal', np.nan)], + ... columns=['name', 'class', 'max_speed'], + ... index=[0, 2, 3, 1]) + >>> df + name class max_speed + 0 falcon bird 389.0 + 2 parrot bird 24.0 + 3 lion mammal 80.5 + 1 monkey mammal NaN + + Take elements at positions 0 and 3 along the axis 0 (default). + + Note how the actual indices selected (0 and 1) do not correspond to + our selected indices 0 and 3. That's because we are selecting the 0th + and 3rd rows, not rows whose indices equal 0 and 3. + + >>> df.take([0, 3]).sort_index() + name class max_speed + 0 falcon bird 389.0 + 1 monkey mammal NaN + + Take elements at indices 1 and 2 along the axis 1 (column selection). + + >>> df.take([1, 2], axis=1) + class max_speed + 0 bird 389.0 + 2 bird 24.0 + 3 mammal 80.5 + 1 mammal NaN + + We may take elements using negative integers for positive indices, + starting from the end of the object, just like with Python lists. + + >>> df.take([-1, -2]).sort_index() + name class max_speed + 1 monkey mammal NaN + 3 lion mammal 80.5 + """ + axis = validate_axis(axis) + if not is_list_like(indices) or isinstance(indices, (dict, set)): + raise TypeError("`indices` must be a list-like except dict or set") + if axis == 0: + return cast(DataFrame, self.iloc[indices, :]) + else: + return cast(DataFrame, self.iloc[:, indices]) + + def eval(self, expr: str, inplace: bool = False) -> Optional[DataFrameOrSeries]: + """ + Evaluate a string describing operations on DataFrame columns. + + Operates on columns only, not specific rows or elements. This allows + `eval` to run arbitrary code, which can make you vulnerable to code + injection if you pass user input to this function. + + Parameters + ---------- + expr : str + The expression string to evaluate. + inplace : bool, default False + If the expression contains an assignment, whether to perform the + operation inplace and mutate the existing DataFrame. Otherwise, + a new DataFrame is returned. + + Returns + ------- + The result of the evaluation. + + See Also + -------- + DataFrame.query : Evaluates a boolean expression to query the columns + of a frame. + DataFrame.assign : Can evaluate an expression or function to create new + values for a column. + eval : Evaluate a Python expression as a string using various + backends. + + Examples + -------- + >>> df = ps.DataFrame({'A': range(1, 6), 'B': range(10, 0, -2)}) + >>> df + A B + 0 1 10 + 1 2 8 + 2 3 6 + 3 4 4 + 4 5 2 + >>> df.eval('A + B') + 0 11 + 1 10 + 2 9 + 3 8 + 4 7 + dtype: int64 + + Assignment is allowed though by default the original DataFrame is not + modified. + + >>> df.eval('C = A + B') + A B C + 0 1 10 11 + 1 2 8 10 + 2 3 6 9 + 3 4 4 8 + 4 5 2 7 + >>> df + A B + 0 1 10 + 1 2 8 + 2 3 6 + 3 4 4 + 4 5 2 + + Use ``inplace=True`` to modify the original DataFrame. + + >>> df.eval('C = A + B', inplace=True) + >>> df + A B C + 0 1 10 11 + 1 2 8 10 + 2 3 6 9 + 3 4 4 8 + 4 5 2 7 + """ + from pyspark.pandas.series import first_series + + if isinstance(self.columns, pd.MultiIndex): + raise TypeError("`eval` is not supported for multi-index columns") + inplace = validate_bool_kwarg(inplace, "inplace") + should_return_series = False + series_name = None + should_return_scalar = False + + # Since `eval_func` doesn't have a type hint, inferring the schema is always preformed + # in the `apply_batch`. Hence, the variables `should_return_series`, `series_name`, + # and `should_return_scalar` can be updated. + def eval_func(pdf): # type: ignore[no-untyped-def] + nonlocal should_return_series + nonlocal series_name + nonlocal should_return_scalar + result_inner = pdf.eval(expr, inplace=inplace) + if inplace: + result_inner = pdf + if isinstance(result_inner, pd.Series): + should_return_series = True + series_name = result_inner.name + result_inner = result_inner.to_frame() + elif is_scalar(result_inner): + should_return_scalar = True + result_inner = pd.Series(result_inner).to_frame() + return result_inner + + result = self.pandas_on_spark.apply_batch(eval_func) + if inplace: + # Here, the result is always a frame because the error is thrown during schema inference + # from pandas. + self._update_internal_frame(result._internal, check_same_anchor=False) + return None + elif should_return_series: + return first_series(result).rename(series_name) + elif should_return_scalar: + return first_series(result)[0] + else: + # Returns a frame + return result + + def explode(self, column: Name, ignore_index: bool = False) -> "DataFrame": + """ + Transform each element of a list-like to a row, replicating index values. + + Parameters + ---------- + column : str or tuple + Column to explode. + ignore_index : bool, default False + If True, the resulting index will be labeled 0, 1, …, n - 1. + + Returns + ------- + DataFrame + Exploded lists to rows of the subset columns; + index will be duplicated for these rows. + + See Also + -------- + DataFrame.unstack : Pivot a level of the (necessarily hierarchical) + index labels. + DataFrame.melt : Unpivot a DataFrame from wide format to long format. + + Examples + -------- + >>> df = ps.DataFrame({'A': [[1, 2, 3], [], [3, 4]], 'B': 1}) + >>> df + A B + 0 [1, 2, 3] 1 + 1 [] 1 + 2 [3, 4] 1 + + >>> df.explode('A') + A B + 0 1.0 1 + 0 2.0 1 + 0 3.0 1 + 1 NaN 1 + 2 3.0 1 + 2 4.0 1 + + >>> df.explode('A', ignore_index=True) + A B + 0 1.0 1 + 1 2.0 1 + 2 3.0 1 + 3 NaN 1 + 4 3.0 1 + 5 4.0 1 + """ + from pyspark.pandas.series import Series + + if not is_name_like_value(column): + raise TypeError("column must be a scalar") + + psdf: DataFrame = DataFrame(self._internal.resolved_copy) + psser = psdf[column] + if not isinstance(psser, Series): + raise ValueError( + "The column %s is not unique. For a multi-index, the label must be a tuple " + "with elements corresponding to each level." % name_like_string(column) + ) + if not isinstance(psser.spark.data_type, ArrayType): + return self.copy() + + sdf = psdf._internal.spark_frame.withColumn( + psser._internal.data_spark_column_names[0], F.explode_outer(psser.spark.column) + ) + + data_fields = psdf._internal.data_fields.copy() + idx = psdf._internal.column_labels.index(psser._column_label) + field = data_fields[idx] + spark_type = cast(ArrayType, field.spark_type).elementType + dtype = spark_type_to_pandas_dtype(spark_type) + data_fields[idx] = field.copy(dtype=dtype, spark_type=spark_type, nullable=True) + + internal = psdf._internal.with_new_sdf(sdf, data_fields=data_fields) + result_df: DataFrame = DataFrame(internal) + return result_df.reset_index(drop=True) if ignore_index else result_df + + def mad(self, axis: Axis = 0) -> "Series": + """ + Return the mean absolute deviation of values. + + .. deprecated:: 3.4.0 + + Parameters + ---------- + axis : {index (0), columns (1)} + Axis for the function to be applied on. + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + >>> df.mad() + a 0.666667 + b 0.066667 + dtype: float64 + + >>> df.mad(axis=1) # doctest: +SKIP + 0 0.45 + 1 0.90 + 2 1.35 + 3 NaN + dtype: float64 + """ + warnings.warn( + "The 'mad' method is deprecated and will be removed in 4.0.0. " + "To compute the same result, you may do `(df - df.mean()).abs().mean()`.", + FutureWarning, + ) + from pyspark.pandas.series import first_series + + axis = validate_axis(axis) + + if axis == 0: + + def get_spark_column(psdf: DataFrame, label: Label) -> PySparkColumn: + scol = psdf._internal.spark_column_for(label) + col_type = psdf._internal.spark_type_for(label) + + if isinstance(col_type, BooleanType): + scol = scol.cast("integer") + + return scol + + new_column_labels: List[Label] = [] + for label in self._internal.column_labels: + # Filtering out only columns of numeric and boolean type column. + dtype = self._psser_for(label).spark.data_type + if isinstance(dtype, (NumericType, BooleanType)): + new_column_labels.append(label) + + new_columns = [ + F.avg(get_spark_column(self, label)).alias(name_like_string(label)) + for label in new_column_labels + ] + + mean_data = self._internal.spark_frame.select(*new_columns).first() + + new_columns = [ + F.avg( + F.abs(get_spark_column(self, label) - mean_data[name_like_string(label)]) + ).alias(name_like_string(label)) + for label in new_column_labels + ] + + sdf = self._internal.spark_frame.select( + *[F.lit(None).cast(StringType()).alias(SPARK_DEFAULT_INDEX_NAME)], *new_columns + ) + + # The data is expected to be small so it's fine to transpose/use the default index. + with ps.option_context("compute.max_rows", 1): + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + column_labels=new_column_labels, + column_label_names=self._internal.column_label_names, + ) + return first_series(DataFrame(internal).transpose()) + + else: + + @pandas_udf(returnType=DoubleType()) # type: ignore[call-overload] + def calculate_columns_axis(*cols: pd.Series) -> pd.Series: + return pd.concat(cols, axis=1).mad(axis=1) + + internal = self._internal.copy( + column_labels=[None], + data_spark_columns=[ + calculate_columns_axis(*self._internal.data_spark_columns).alias( + SPARK_DEFAULT_SERIES_NAME + ) + ], + data_fields=[None], + column_label_names=None, + ) + return first_series(DataFrame(internal)) + + def mode(self, axis: Axis = 0, numeric_only: bool = False, dropna: bool = True) -> "DataFrame": + """ + Get the mode(s) of each element along the selected axis. + + The mode of a set of values is the value that appears most often. + It can be multiple values. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + Axis for the function to be applied on. + numeric_only : bool, default False + If True, only apply to numeric columns. + dropna : bool, default True + Don't consider counts of NaN/NaT. + + Returns + ------- + DataFrame + The modes of each column or row. + + See Also + -------- + Series.mode : Return the highest frequency value in a Series. + Series.value_counts : Return the counts of values in a Series. + + Examples + -------- + >>> df = ps.DataFrame([('bird', 2, 2), + ... ('mammal', 4, np.nan), + ... ('arthropod', 8, 0), + ... ('bird', 2, np.nan)], + ... index=('falcon', 'horse', 'spider', 'ostrich'), + ... columns=('species', 'legs', 'wings')) + >>> df + species legs wings + falcon bird 2 2.0 + horse mammal 4 NaN + spider arthropod 8 0.0 + ostrich bird 2 NaN + + By default missing values are not considered, and the mode of wings + are both 0 and 2. Because the resulting DataFrame has two rows, + the second row of ``species`` and ``legs`` contains ``NaN``. + + >>> df.mode() + species legs wings + 0 bird 2.0 0.0 + 1 None NaN 2.0 + + Setting ``dropna=False`` ``NaN`` values are considered and they can be + the mode (like for wings). + + >>> df.mode(dropna=False) + species legs wings + 0 bird 2 NaN + + Setting ``numeric_only=True``, only the mode of numeric columns is + computed, and columns of other types are ignored. + + >>> df.mode(numeric_only=True) + legs wings + 0 2.0 0.0 + 1 NaN 2.0 + """ + axis = validate_axis(axis, none_axis=0) + if axis != 0: + raise ValueError('axis should be either 0 or "index" currently.') + if numeric_only is None and axis == 0: + numeric_only = True + + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + + mode_scols: List[PySparkColumn] = [] + mode_col_names: List[str] = [] + mode_labels: List[Label] = [] + for label, col_name in zip( + self._internal.column_labels, self._internal.data_spark_column_names + ): + psser = self._psser_for(label) + is_numeric = isinstance(psser.spark.data_type, (NumericType, BooleanType)) + + if not numeric_only or is_numeric: + scol = psser.spark.column + mode_scol = SF.mode(scol, dropna).alias(col_name) + mode_scols.append(mode_scol) + mode_col_names.append(col_name) + mode_labels.append(label) + + # Here, after aggregation, a spark_frame looks like below: + # +-------+----+----------+ + # |species|legs| wings| + # +-------+----+----------+ + # | [bird]| [2]|[0.0, 2.0]| + # +-------+----+----------+ + sdf = self._internal.spark_frame.select(mode_scols) + sdf = sdf.select(*[F.array_sort(F.col(name)).alias(name) for name in mode_col_names]) + + zip_col_name = verify_temp_column_name(sdf, "__mode_zip_tmp_col__") + explode_col_name = verify_temp_column_name(sdf, "__mode_explode_tmp_col__") + + # After this transformation, sdf turns out to be: + # +-------+----+-----+ + # |species|legs|wings| + # +-------+----+-----+ + # | bird| 2| 0.0| + # | NULL|NULL| 2.0| + # +-------+----+-----+ + sdf = ( + sdf.select(F.arrays_zip(*[F.col(name) for name in mode_col_names]).alias(zip_col_name)) + .select(F.explode(F.col(zip_col_name)).alias(explode_col_name)) + .select( + *[ + F.col("{0}.{1}".format(explode_col_name, name)).alias(name) + for name in mode_col_names + ] + ) + ) + + sdf = sdf.withColumn(SPARK_DEFAULT_INDEX_NAME, F.monotonically_increasing_id()) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + column_labels=mode_labels, + data_spark_columns=[scol_for(sdf, col) for col in mode_col_names], + ) + return DataFrame(internal) + + def tail(self, n: int = 5) -> "DataFrame": + """ + Return the last `n` rows. + + This function returns last `n` rows from the object based on + position. It is useful for quickly verifying data, for example, + after sorting or appending rows. + + For negative values of `n`, this function returns all rows except + the first `n` rows, equivalent to ``df[n:]``. + + Parameters + ---------- + n : int, default 5 + Number of rows to select. + + Returns + ------- + type of caller + The last `n` rows of the caller object. + + See Also + -------- + DataFrame.head : The first `n` rows of the caller object. + + Examples + -------- + >>> df = ps.DataFrame({'animal': ['alligator', 'bee', 'falcon', 'lion', + ... 'monkey', 'parrot', 'shark', 'whale', 'zebra']}) + >>> df + animal + 0 alligator + 1 bee + 2 falcon + 3 lion + 4 monkey + 5 parrot + 6 shark + 7 whale + 8 zebra + + Viewing the last 5 lines + + >>> df.tail() # doctest: +SKIP + animal + 4 monkey + 5 parrot + 6 shark + 7 whale + 8 zebra + + Viewing the last `n` lines (three in this case) + + >>> df.tail(3) # doctest: +SKIP + animal + 6 shark + 7 whale + 8 zebra + + For negative values of `n` + + >>> df.tail(-3) # doctest: +SKIP + animal + 3 lion + 4 monkey + 5 parrot + 6 shark + 7 whale + 8 zebra + """ + if not isinstance(n, int): + raise TypeError("bad operand type for unary -: '{}'".format(type(n).__name__)) + if n < 0: + n = len(self) + n + if n <= 0: + return ps.DataFrame(self._internal.with_filter(F.lit(False))) + # Should use `resolved_copy` here for the case like `(psdf + 1).tail()` + sdf = self._internal.resolved_copy.spark_frame + rows = sdf.tail(n) + new_sdf = default_session().createDataFrame(rows, sdf.schema) + + return DataFrame(self._internal.with_new_sdf(new_sdf)) + + def align( + self, + other: DataFrameOrSeries, + join: str = "outer", + axis: Optional[Axis] = None, + copy: bool = True, + ) -> Tuple["DataFrame", DataFrameOrSeries]: + """ + Align two objects on their axes with the specified join method. + + Join method is specified for each axis Index. + + Parameters + ---------- + other : DataFrame or Series + join : {{'outer', 'inner', 'left', 'right'}}, default 'outer' + axis : allowed axis of the other object, default None + Align on index (0), columns (1), or both (None). + copy : bool, default True + Always returns new objects. If copy=False and no reindexing is + required then original objects are returned. + + Returns + ------- + (left, right) : (DataFrame, type of other) + Aligned objects. + + Examples + -------- + >>> ps.set_option("compute.ops_on_diff_frames", True) + >>> df1 = ps.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30]) + >>> df2 = ps.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12]) + + Align both axis: + + >>> aligned_l, aligned_r = df1.align(df2) + >>> aligned_l.sort_index() + a b c + 10 1.0 a NaN + 11 NaN None NaN + 12 NaN None NaN + 20 2.0 b NaN + 30 3.0 c NaN + >>> aligned_r.sort_index() + a b c + 10 4.0 NaN d + 11 5.0 NaN e + 12 6.0 NaN f + 20 NaN NaN None + 30 NaN NaN None + + Align only axis=0 (index): + + >>> aligned_l, aligned_r = df1.align(df2, axis=0) + >>> aligned_l.sort_index() + a b + 10 1.0 a + 11 NaN None + 12 NaN None + 20 2.0 b + 30 3.0 c + >>> aligned_r.sort_index() + a c + 10 4.0 d + 11 5.0 e + 12 6.0 f + 20 NaN None + 30 NaN None + + Align only axis=1 (column): + + >>> aligned_l, aligned_r = df1.align(df2, axis=1) + >>> aligned_l.sort_index() + a b c + 10 1 a NaN + 20 2 b NaN + 30 3 c NaN + >>> aligned_r.sort_index() + a b c + 10 4 NaN d + 11 5 NaN e + 12 6 NaN f + + Align with the join type "inner": + + >>> aligned_l, aligned_r = df1.align(df2, join="inner") + >>> aligned_l.sort_index() + a + 10 1 + >>> aligned_r.sort_index() + a + 10 4 + + Align with a Series: + + >>> s = ps.Series([7, 8, 9], index=[10, 11, 12]) + >>> aligned_l, aligned_r = df1.align(s, axis=0) + >>> aligned_l.sort_index() + a b + 10 1.0 a + 11 NaN None + 12 NaN None + 20 2.0 b + 30 3.0 c + >>> aligned_r.sort_index() + 10 7.0 + 11 8.0 + 12 9.0 + 20 NaN + 30 NaN + dtype: float64 + + >>> ps.reset_option("compute.ops_on_diff_frames") + """ + from pyspark.pandas.series import Series, first_series + + if not isinstance(other, (DataFrame, Series)): + raise TypeError("unsupported type: {}".format(type(other).__name__)) + + how = validate_how(join) + axis = validate_axis(axis, None) + + right_is_series = isinstance(other, Series) + if right_is_series: + if axis is None: + raise ValueError("Must specify axis=0 or 1") + elif axis != 0: + raise NotImplementedError( + "align currently only works for axis=0 when right is Series" + ) + + left = self + right = other + + if (axis is None or axis == 0) and not same_anchor(left, right): + combined = combine_frames(left, right, how=how) + left = combined["this"] + right = combined["that"] + + if right_is_series: + right = first_series(cast(DataFrame[Any], right)).rename(other.name) + + if ( + axis is None or axis == 1 + ) and left._internal.column_labels != right._internal.column_labels: + + if left._internal.column_labels_level != right._internal.column_labels_level: + raise ValueError("cannot join with no overlapping index names") + + left = left.copy() + right = right.copy() + + if how == "full": + column_labels = sorted( + list(set(left._internal.column_labels) | set(right._internal.column_labels)) + ) + elif how == "inner": + column_labels = sorted( + list(set(left._internal.column_labels) & set(right._internal.column_labels)) + ) + elif how == "left": + column_labels = left._internal.column_labels + else: + column_labels = right._internal.column_labels + + for label in column_labels: + if label not in left._internal.column_labels: + left[label] = F.lit(None).cast(DoubleType()) + left = left[column_labels] + for label in column_labels: + if label not in right._internal.column_labels: + right[label] = F.lit(None).cast(DoubleType()) + right = right[column_labels] + + return (left.copy(), right.copy()) if copy else (left, right) + + @staticmethod + def from_dict( + data: Dict[Name, Sequence[Any]], + orient: str = "columns", + dtype: Union[str, Dtype] = None, + columns: Optional[List[Name]] = None, + ) -> "DataFrame": + """ + Construct DataFrame from dict of array-like or dicts. + + Creates DataFrame object from dictionary by columns or by index + allowing dtype specification. + + Parameters + ---------- + data : dict + Of the form {field : array-like} or {field : dict}. + orient : {'columns', 'index'}, default 'columns' + The "orientation" of the data. If the keys of the passed dict + should be the columns of the resulting DataFrame, pass 'columns' + (default). Otherwise, if the keys should be rows, pass 'index'. + dtype : dtype, default None + Data type to force, otherwise infer. + columns : list, default None + Column labels to use when ``orient='index'``. Raises a ValueError + if used with ``orient='columns'``. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.from_records : DataFrame from structured ndarray, sequence + of tuples or dicts, or DataFrame. + DataFrame : DataFrame object creation using constructor. + + Examples + -------- + By default the keys of the dict become the DataFrame columns: + + >>> data = {'col_1': [3, 2, 1, 0], 'col_2': [10, 20, 30, 40]} + >>> ps.DataFrame.from_dict(data) + col_1 col_2 + 0 3 10 + 1 2 20 + 2 1 30 + 3 0 40 + + Specify ``orient='index'`` to create the DataFrame using dictionary + keys as rows: + + >>> data = {'row_1': [3, 2, 1, 0], 'row_2': [10, 20, 30, 40]} + >>> ps.DataFrame.from_dict(data, orient='index').sort_index() + 0 1 2 3 + row_1 3 2 1 0 + row_2 10 20 30 40 + + When using the 'index' orientation, the column names can be + specified manually: + + >>> ps.DataFrame.from_dict(data, orient='index', + ... columns=['A', 'B', 'C', 'D']).sort_index() + A B C D + row_1 3 2 1 0 + row_2 10 20 30 40 + """ + return DataFrame( + pd.DataFrame.from_dict( + data, orient=orient, dtype=dtype, columns=columns # type: ignore[arg-type] + ) + ) + + # Override the `groupby` to specify the actual return type annotation. + def groupby( + self, + by: Union[Name, "Series", List[Union[Name, "Series"]]], + axis: Axis = 0, + as_index: bool = True, + dropna: bool = True, + ) -> "DataFrameGroupBy": + return cast( + "DataFrameGroupBy", super().groupby(by=by, axis=axis, as_index=as_index, dropna=dropna) + ) + + groupby.__doc__ = Frame.groupby.__doc__ + + def _build_groupby( + self, by: List[Union["Series", Label]], as_index: bool, dropna: bool + ) -> "DataFrameGroupBy": + from pyspark.pandas.groupby import DataFrameGroupBy + + return DataFrameGroupBy._build(self, by, as_index=as_index, dropna=dropna) + + def resample( + self, + rule: str, + closed: Optional[str] = None, + label: Optional[str] = None, + on: Optional["Series"] = None, + ) -> "DataFrameResampler": + """ + Resample time-series data. + + Convenience method for frequency conversion and resampling of time series. + The object must have a datetime-like index (only support `DatetimeIndex` for now), + or the caller must pass the label of a datetime-like + series/index to the ``on`` keyword parameter. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + rule : str + The offset string or object representing target conversion. + Currently, supported units are {'Y', 'A', 'M', 'D', 'H', + 'T', 'MIN', 'S'}. + closed : {{'right', 'left'}}, default None + Which side of bin interval is closed. The default is 'left' + for all frequency offsets except for 'A', 'Y' and 'M' which all + have a default of 'right'. + label : {{'right', 'left'}}, default None + Which bin edge label to label bucket with. The default is 'left' + for all frequency offsets except for 'A', 'Y' and 'M' which all + have a default of 'right'. + on : Series, optional + For a DataFrame, column to use instead of index for resampling. + Column must be datetime-like. + + Returns + ------- + DataFrameResampler + + See Also + -------- + Series.resample : Resample a Series. + groupby : Group by mapping, function, label, or list of labels. + """ + from pyspark.pandas.indexes import DatetimeIndex + from pyspark.pandas.resample import DataFrameResampler + + if on is None and not isinstance(self.index, DatetimeIndex): + raise NotImplementedError("resample currently works only for DatetimeIndex") + if on is not None and not isinstance( + as_spark_type(on.dtype), (TimestampType, TimestampNTZType) + ): + raise NotImplementedError("`on` currently works only for TimestampType") + + agg_columns: List[ps.Series] = [] + for column_label in self._internal.column_labels: + if isinstance(self._internal.spark_type_for(column_label), (NumericType, BooleanType)): + agg_columns.append(self._psser_for(column_label)) + + if len(agg_columns) == 0: + raise ValueError("No available aggregation columns!") + + return DataFrameResampler( + psdf=self, + resamplekey=on, + rule=rule, + closed=closed, + label=label, + agg_columns=agg_columns, + ) + + def _to_internal_pandas(self) -> pd.DataFrame: + """ + Return a pandas DataFrame directly from _internal to avoid overhead of copy. + + This method is for internal use only. + """ + return self._internal.to_pandas_frame + + def _get_or_create_repr_pandas_cache(self, n: int) -> Union[pd.DataFrame, pd.Series]: + if not hasattr(self, "_repr_pandas_cache") or n not in self._repr_pandas_cache: + object.__setattr__( + self, "_repr_pandas_cache", {n: self.head(n + 1)._to_internal_pandas()} + ) + return self._repr_pandas_cache[n] + + def __repr__(self) -> str: + max_display_count = get_option("display.max_rows") + if max_display_count is None: + return self._to_internal_pandas().to_string() + + pdf = cast("DataFrame", self._get_or_create_repr_pandas_cache(max_display_count)) + pdf_length = len(pdf) + pdf = cast("DataFrame", pdf.iloc[:max_display_count]) + if pdf_length > max_display_count: + repr_string = pdf.to_string(show_dimensions=True) + match = REPR_PATTERN.search(repr_string) + if match is not None: + nrows = match.group("rows") + ncols = match.group("columns") + footer = "\n\n[Showing only the first {nrows} rows x {ncols} columns]".format( + nrows=nrows, ncols=ncols + ) + return REPR_PATTERN.sub(footer, repr_string) + return pdf.to_string() + + def _repr_html_(self) -> str: + max_display_count = get_option("display.max_rows") + if max_display_count is None: + return self._to_internal_pandas().to_html(notebook=True) + + pdf = self._get_or_create_repr_pandas_cache(max_display_count) + pdf_length = len(pdf) + pdf = pdf.iloc[:max_display_count] + if pdf_length > max_display_count: + repr_html = pdf.to_html(show_dimensions=True, notebook=True) + match = REPR_HTML_PATTERN.search(repr_html) + if match is not None: + nrows = match.group("rows") + ncols = match.group("columns") + by = chr(215) + footer = ( + "\n

    Showing only the first {rows} rows " + "{by} {cols} columns

    \n".format(rows=nrows, by=by, cols=ncols) + ) + return REPR_HTML_PATTERN.sub(footer, repr_html) + return pdf.to_html(notebook=True) + + def __getitem__(self, key: Any) -> Any: + from pyspark.pandas.series import Series + + if key is None: + raise KeyError("none key") + elif isinstance(key, Series): + return self.loc[key.astype(bool)] + elif isinstance(key, slice): + if any(type(n) == int or None for n in [key.start, key.stop]): + # Seems like pandas Frame always uses int as positional search when slicing + # with ints. + return self.iloc[key] + return self.loc[key] + elif is_name_like_value(key): + return self.loc[:, key] + elif is_list_like(key): + return self.loc[:, list(key)] + + def __setitem__(self, key: Any, value: Any) -> None: + from pyspark.pandas.series import Series + + if isinstance(value, (DataFrame, Series)) and not same_anchor(value, self): + # Different Series or DataFrames + level = self._internal.column_labels_level + key = DataFrame._index_normalized_label(level, key) + value = DataFrame._index_normalized_frame(level, value) + + def assign_columns( + psdf: DataFrame, this_column_labels: List[Label], that_column_labels: List[Label] + ) -> Iterator[Tuple["Series", Label]]: + assert len(key) == len(that_column_labels) + # Note that here intentionally uses `zip_longest` that combine + # that_columns. + for k, this_label, that_label in zip_longest( + key, this_column_labels, that_column_labels + ): + yield (psdf._psser_for(that_label), tuple(["that", *k])) + if this_label is not None and this_label[1:] != k: + yield (psdf._psser_for(this_label), this_label) + + psdf = align_diff_frames(assign_columns, self, value, fillna=False, how="left") + elif isinstance(value, list): + if len(self) != len(value): + raise ValueError("Length of values does not match length of index") + + # TODO: avoid using default index? + with option_context( + "compute.default_index_type", + "distributed-sequence", + "compute.ops_on_diff_frames", + True, + ): + psdf = self.reset_index() + psdf[key] = ps.DataFrame(value) + psdf = psdf.set_index(psdf.columns[: self._internal.index_level]) + psdf.index.names = self.index.names + + elif isinstance(key, list): + assert isinstance(value, DataFrame) + # Same DataFrames. + field_names = value.columns + psdf = self._assign({k: value[c] for k, c in zip(key, field_names)}) + else: + # Same Series. + psdf = self._assign({key: value}) + + # Since Spark 3.4, df.__setitem__ generates a new dataframe instead of operating + # in-place to follow pandas v1.4 behavior, see also SPARK-38946. + self._update_internal_frame(psdf._internal, anchor_force_disconnect=True) + + @staticmethod + def _index_normalized_label(level: int, labels: Union[Name, Sequence[Name]]) -> List[Label]: + """ + Returns a label that is normalized against the current column index level. + For example, the key "abc" can be ("abc", "", "") if the current Frame has + a multi-index for its column + """ + if is_name_like_tuple(labels): + labels = [labels] + elif is_name_like_value(labels): + labels = [(labels,)] + else: + labels = [k if is_name_like_tuple(k) else (k,) for k in labels] + + if any(len(label) > level for label in labels): + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + max(len(label) for label in labels), level + ) + ) + return [tuple(list(label) + ([""] * (level - len(label)))) for label in labels] + + @staticmethod + def _index_normalized_frame(level: int, psser_or_psdf: DataFrameOrSeries) -> "DataFrame": + """ + Returns a frame that is normalized against the current column index level. + For example, the name in `pd.Series([...], name="abc")` can be can be + ("abc", "", "") if the current DataFrame has a multi-index for its column + """ + from pyspark.pandas.series import Series + + if isinstance(psser_or_psdf, Series): + psdf = psser_or_psdf.to_frame() + else: + assert isinstance(psser_or_psdf, DataFrame), type(psser_or_psdf) + psdf = psser_or_psdf.copy() + + psdf.columns = pd.MultiIndex.from_tuples( + [ + tuple([name_like_string(label)] + ([""] * (level - 1))) + for label in psdf._internal.column_labels + ], + ) + + return psdf + + def __getattr__(self, key: str) -> Any: + if key.startswith("__"): + raise AttributeError(key) + if hasattr(MissingPandasLikeDataFrame, key): + property_or_func = getattr(MissingPandasLikeDataFrame, key) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + + try: + return self.loc[:, key] + except KeyError: + raise AttributeError( + "'%s' object has no attribute '%s'" % (self.__class__.__name__, key) + ) + + def __setattr__(self, key: str, value: Any) -> None: + try: + object.__getattribute__(self, key) + return object.__setattr__(self, key, value) + except AttributeError: + pass + + if (key,) in self._internal.column_labels: + self[key] = value + else: + msg = "pandas-on-Spark doesn't allow columns to be created via a new attribute name" + if is_testing(): + raise AssertionError(msg) + else: + warnings.warn(msg, UserWarning) + + def __len__(self) -> int: + return self._internal.resolved_copy.spark_frame.count() + + def __dir__(self) -> Iterable[str]: + fields = [ + f for f in self._internal.resolved_copy.spark_frame.schema.fieldNames() if " " not in f + ] + return list(super().__dir__()) + fields + + def __iter__(self) -> Iterator[Name]: + return iter(self.columns) + + # NDArray Compat + def __array_ufunc__( + self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any + ) -> "DataFrame": + # TODO: is it possible to deduplicate it with '_map_series_op'? + if all(isinstance(inp, DataFrame) for inp in inputs) and any( + not same_anchor(inp, inputs[0]) for inp in inputs + ): + # binary only + assert len(inputs) == 2 + this = inputs[0] + that = inputs[1] + if this._internal.column_labels_level != that._internal.column_labels_level: + raise ValueError("cannot join with no overlapping index names") + + # Different DataFrames + def apply_op( + psdf: DataFrame, this_column_labels: List[Label], that_column_labels: List[Label] + ) -> Iterator[Tuple["Series", Label]]: + for this_label, that_label in zip(this_column_labels, that_column_labels): + yield ( + ufunc( + psdf._psser_for(this_label), psdf._psser_for(that_label), **kwargs + ).rename(this_label), + this_label, + ) + + return align_diff_frames(apply_op, this, that, fillna=True, how="full") + else: + # DataFrame and Series + applied = [] + this = inputs[0] + assert all(inp is this for inp in inputs if isinstance(inp, DataFrame)) + + for label in this._internal.column_labels: + arguments = [] + for inp in inputs: + arguments.append(inp[label] if isinstance(inp, DataFrame) else inp) + # both binary and unary. + applied.append(ufunc(*arguments, **kwargs).rename(label)) + + internal = this._internal.with_new_columns(applied) + return DataFrame(internal) + + def __class_getitem__(cls, params: Any) -> object: + # See https://github.com/python/typing/issues/193 + # we always wraps the given type hints by a tuple to mimic the variadic generic. + return create_tuple_for_frame_type(params) + + +def _reduce_spark_multi(sdf: PySparkDataFrame, aggs: List[PySparkColumn]) -> Any: + """ + Performs a reduction on a spark DataFrame, the functions being known SQL aggregate functions. + """ + SparkDataFrame = get_dataframe_class() + assert isinstance(sdf, SparkDataFrame) + sdf0 = sdf.agg(*aggs) + lst = sdf0.limit(2).toPandas() + assert len(lst) == 1, (sdf, lst) + row = lst.iloc[0] + lst2 = list(row) + assert len(lst2) == len(aggs), (row, lst2) + return lst2 + + +class CachedDataFrame(DataFrame): + """ + Cached pandas-on-Spark DataFrame, which corresponds to pandas DataFrame logically, but + internally it caches the corresponding Spark DataFrame. + """ + + def __init__(self, internal: InternalFrame, storage_level: Optional[StorageLevel] = None): + if storage_level is None: + object.__setattr__(self, "_cached", internal.spark_frame.cache()) + elif isinstance(storage_level, StorageLevel): + object.__setattr__(self, "_cached", internal.spark_frame.persist(storage_level)) + else: + raise TypeError( + "Only a valid pyspark.StorageLevel type is acceptable for the `storage_level`" + ) + super().__init__(internal) + + def __enter__(self) -> "CachedDataFrame": + return self + + def __exit__( + self, + exception_type: Optional[Type[BaseException]], + exception_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + self.spark.unpersist() + return None + + # create accessor for Spark related methods. + spark = CachedAccessor("spark", CachedSparkFrameMethods) + + +def _test() -> None: + import os + import doctest + import shutil + import sys + import tempfile + import uuid + from pyspark.sql import SparkSession + import pyspark.pandas.frame + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.frame.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.frame tests").getOrCreate() + ) + globs["spark"] = spark + + db_name = "db%s" % str(uuid.uuid4()).replace("-", "") + spark.sql("CREATE DATABASE %s" % db_name) + globs["db"] = db_name + + path = tempfile.mkdtemp() + globs["path"] = path + + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.frame, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + + shutil.rmtree(path, ignore_errors=True) + spark.sql("DROP DATABASE IF EXISTS %s CASCADE" % db_name) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/generic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/generic.py new file mode 100644 index 0000000..c346889 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/generic.py @@ -0,0 +1,3560 @@ +# +# 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. +# + +""" +A base class of DataFrame/Column to behave like pandas DataFrame/Series. +""" +from abc import ABCMeta, abstractmethod +from collections import Counter +from functools import reduce +from typing import ( + Any, + Callable, + Dict, + Iterable, + IO, + List, + Optional, + NoReturn, + Tuple, + Union, + TYPE_CHECKING, + cast, +) +import warnings + +import numpy as np +import pandas as pd +from pandas.api.types import is_list_like # type: ignore[attr-defined] + +from pyspark.sql import Column, functions as F +from pyspark.sql.types import ( + BooleanType, + DoubleType, + LongType, + NumericType, +) + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import ( + Axis, + DataFrameOrSeries, + Dtype, + FrameLike, + Label, + Name, + Scalar, +) +from pyspark.pandas.indexing import AtIndexer, iAtIndexer, iLocIndexer, LocIndexer +from pyspark.pandas.internal import InternalFrame +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.typedef import spark_type_to_pandas_dtype +from pyspark.pandas.utils import ( + is_name_like_tuple, + is_name_like_value, + name_like_string, + scol_for, + sql_conf, + validate_arguments_and_invoke_function, + validate_axis, + validate_mode, + SPARK_CONF_ARROW_ENABLED, + log_advice, +) + +if TYPE_CHECKING: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.indexes.base import Index + from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.series import Series + from pyspark.pandas.window import Rolling, Expanding, ExponentialMoving + + +bool_type = bool + + +class Frame(object, metaclass=ABCMeta): + """ + The base class for both DataFrame and Series. + """ + + @abstractmethod + def __getitem__(self, key: Any) -> Any: + pass + + @property + @abstractmethod + def _internal(self) -> InternalFrame: + pass + + @abstractmethod + def _apply_series_op( + self: FrameLike, + op: Callable[["Series"], Union["Series", Column]], + should_resolve: bool = False, + ) -> FrameLike: + pass + + @abstractmethod + def _reduce_for_stat_function( + self, + sfun: Callable[["Series"], Column], + name: str, + axis: Optional[Axis] = None, + numeric_only: bool = True, + skipna: bool = True, + **kwargs: Any, + ) -> Union["Series", Scalar]: + pass + + @property + @abstractmethod + def dtypes(self) -> Union[pd.Series, Dtype]: + pass + + @abstractmethod + def to_pandas(self) -> Union[pd.DataFrame, pd.Series]: + pass + + @abstractmethod + def _to_pandas(self) -> Union[pd.DataFrame, pd.Series]: + pass + + @property + @abstractmethod + def index(self) -> "Index": + pass + + @abstractmethod + def copy(self: FrameLike) -> FrameLike: + pass + + @abstractmethod + def _to_internal_pandas(self) -> Union[pd.DataFrame, pd.Series]: + pass + + @abstractmethod + def head(self: FrameLike, n: int = 5) -> FrameLike: + pass + + # TODO: add 'axis' parameter + def cummin(self: FrameLike, skipna: bool = True) -> FrameLike: + """ + Return cumulative minimum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative minimum. + + .. note:: the current implementation of cummin uses Spark's Window without + specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + skipna: boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.min: Return the minimum over DataFrame axis. + DataFrame.cummax: Return cumulative maximum over DataFrame axis. + DataFrame.cummin: Return cumulative minimum over DataFrame axis. + DataFrame.cumsum: Return cumulative sum over DataFrame axis. + Series.min: Return the minimum over Series axis. + Series.cummax: Return cumulative maximum over Series axis. + Series.cummin: Return cumulative minimum over Series axis. + Series.cumsum: Return cumulative sum over Series axis. + Series.cumprod: Return cumulative product over Series axis. + + Examples + -------- + >>> df = ps.DataFrame([[2.0, 1.0], [3.0, None], [1.0, 0.0]], columns=list('AB')) + >>> df + A B + 0 2.0 1.0 + 1 3.0 NaN + 2 1.0 0.0 + + By default, iterates over rows and finds the minimum in each column. + + >>> df.cummin() + A B + 0 2.0 1.0 + 1 2.0 NaN + 2 1.0 0.0 + + It works identically in Series. + + >>> df.A.cummin() + 0 2.0 + 1 2.0 + 2 1.0 + Name: A, dtype: float64 + """ + return self._apply_series_op(lambda psser: psser._cum(F.min, skipna), should_resolve=True) + + # TODO: add 'axis' parameter + def cummax(self: FrameLike, skipna: bool = True) -> FrameLike: + """ + Return cumulative maximum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative maximum. + + .. note:: the current implementation of cummax uses Spark's Window without + specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + skipna: boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.max: Return the maximum over DataFrame axis. + DataFrame.cummax: Return cumulative maximum over DataFrame axis. + DataFrame.cummin: Return cumulative minimum over DataFrame axis. + DataFrame.cumsum: Return cumulative sum over DataFrame axis. + DataFrame.cumprod: Return cumulative product over DataFrame axis. + Series.max: Return the maximum over Series axis. + Series.cummax: Return cumulative maximum over Series axis. + Series.cummin: Return cumulative minimum over Series axis. + Series.cumsum: Return cumulative sum over Series axis. + Series.cumprod: Return cumulative product over Series axis. + + Examples + -------- + >>> df = ps.DataFrame([[2.0, 1.0], [3.0, None], [1.0, 0.0]], columns=list('AB')) + >>> df + A B + 0 2.0 1.0 + 1 3.0 NaN + 2 1.0 0.0 + + By default, iterates over rows and finds the maximum in each column. + + >>> df.cummax() + A B + 0 2.0 1.0 + 1 3.0 NaN + 2 3.0 1.0 + + It works identically in Series. + + >>> df.B.cummax() + 0 1.0 + 1 NaN + 2 1.0 + Name: B, dtype: float64 + """ + return self._apply_series_op(lambda psser: psser._cum(F.max, skipna), should_resolve=True) + + # TODO: add 'axis' parameter + def cumsum(self: FrameLike, skipna: bool = True) -> FrameLike: + """ + Return cumulative sum over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative sum. + + .. note:: the current implementation of cumsum uses Spark's Window without + specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + skipna: boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.sum: Return the sum over DataFrame axis. + DataFrame.cummax: Return cumulative maximum over DataFrame axis. + DataFrame.cummin: Return cumulative minimum over DataFrame axis. + DataFrame.cumsum: Return cumulative sum over DataFrame axis. + DataFrame.cumprod: Return cumulative product over DataFrame axis. + Series.sum: Return the sum over Series axis. + Series.cummax: Return cumulative maximum over Series axis. + Series.cummin: Return cumulative minimum over Series axis. + Series.cumsum: Return cumulative sum over Series axis. + Series.cumprod: Return cumulative product over Series axis. + + Examples + -------- + >>> df = ps.DataFrame([[2.0, 1.0], [3.0, None], [1.0, 0.0]], columns=list('AB')) + >>> df + A B + 0 2.0 1.0 + 1 3.0 NaN + 2 1.0 0.0 + + By default, iterates over rows and finds the sum in each column. + + >>> df.cumsum() + A B + 0 2.0 1.0 + 1 5.0 NaN + 2 6.0 1.0 + + It works identically in Series. + + >>> df.A.cumsum() + 0 2.0 + 1 5.0 + 2 6.0 + Name: A, dtype: float64 + """ + return self._apply_series_op(lambda psser: psser._cumsum(skipna), should_resolve=True) + + # TODO: add 'axis' parameter + # TODO: use pandas_udf to support negative values and other options later + # other window except unbounded ones is supported as of Spark 3.0. + def cumprod(self: FrameLike, skipna: bool = True) -> FrameLike: + """ + Return cumulative product over a DataFrame or Series axis. + + Returns a DataFrame or Series of the same size containing the cumulative product. + + .. note:: the current implementation of cumprod uses Spark's Window without + specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + .. note:: unlike pandas', pandas-on-Spark's emulates cumulative product by + ``exp(sum(log(...)))`` trick. Therefore, it only works for positive numbers. + + Parameters + ---------- + skipna: boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.cummax: Return cumulative maximum over DataFrame axis. + DataFrame.cummin: Return cumulative minimum over DataFrame axis. + DataFrame.cumsum: Return cumulative sum over DataFrame axis. + DataFrame.cumprod: Return cumulative product over DataFrame axis. + Series.cummax: Return cumulative maximum over Series axis. + Series.cummin: Return cumulative minimum over Series axis. + Series.cumsum: Return cumulative sum over Series axis. + Series.cumprod: Return cumulative product over Series axis. + + Raises + ------ + Exception: If the values is equal to or lower than 0. + + Examples + -------- + >>> df = ps.DataFrame([[2.0, 1.0], [3.0, None], [4.0, 10.0]], columns=list('AB')) + >>> df + A B + 0 2.0 1.0 + 1 3.0 NaN + 2 4.0 10.0 + + By default, iterates over rows and finds the sum in each column. + + >>> df.cumprod() + A B + 0 2.0 1.0 + 1 6.0 NaN + 2 24.0 10.0 + + It works identically in Series. + + >>> df.A.cumprod() + 0 2.0 + 1 6.0 + 2 24.0 + Name: A, dtype: float64 + """ + return self._apply_series_op(lambda psser: psser._cumprod(skipna), should_resolve=True) + + # TODO: Although this has removed pandas >= 1.0.0, but we're keeping this as deprecated + # since we're using this for `DataFrame.info` internally. + # We can drop it once our minimal pandas version becomes 1.0.0. + def get_dtype_counts(self) -> pd.Series: + """ + Return counts of unique dtypes in this object. + + .. deprecated:: 0.14.0 + + Returns + ------- + dtype: pd.Series + Series with the count of columns with each dtype. + + See Also + -------- + dtypes: Return the dtypes in this object. + + Examples + -------- + >>> a = [['a', 1, 1], ['b', 2, 2], ['c', 3, 3]] + >>> df = ps.DataFrame(a, columns=['str', 'int1', 'int2']) + >>> df + str int1 int2 + 0 a 1 1 + 1 b 2 2 + 2 c 3 3 + + >>> df.get_dtype_counts().sort_values() + object 1 + int64 2 + dtype: int64 + + >>> df.str.get_dtype_counts().sort_values() + object 1 + dtype: int64 + """ + warnings.warn( + "`get_dtype_counts` has been deprecated and will be " + "removed in a future version. For DataFrames use " + "`.dtypes.value_counts()", + FutureWarning, + ) + if not isinstance(self.dtypes, Iterable): + dtypes = [self.dtypes] + else: + dtypes = list(self.dtypes) + return pd.Series(dict(Counter([d.name for d in dtypes]))) + + def pipe(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + r""" + Apply func(self, \*args, \*\*kwargs). + + Parameters + ---------- + func: function + function to apply to the DataFrame. + ``args``, and ``kwargs`` are passed into ``func``. + Alternatively a ``(callable, data_keyword)`` tuple where + ``data_keyword`` is a string indicating the keyword of + ``callable`` that expects the DataFrames. + args: iterable, optional + positional arguments passed into ``func``. + kwargs: mapping, optional + a dictionary of keyword arguments passed into ``func``. + + Returns + ------- + object: the return type of ``func``. + + Notes + ----- + Use ``.pipe`` when chaining together functions that expect + Series, DataFrames or GroupBy objects. For example, given + + >>> df = ps.DataFrame({'category': ['A', 'A', 'B'], + ... 'col1': [1, 2, 3], + ... 'col2': [4, 5, 6]}, + ... columns=['category', 'col1', 'col2']) + >>> def keep_category_a(df): + ... return df[df['category'] == 'A'] + >>> def add_one(df, column): + ... return df.assign(col3=df[column] + 1) + >>> def multiply(df, column1, column2): + ... return df.assign(col4=df[column1] * df[column2]) + + + instead of writing + + >>> multiply(add_one(keep_category_a(df), column="col1"), column1="col2", column2="col3") + category col1 col2 col3 col4 + 0 A 1 4 2 8 + 1 A 2 5 3 15 + + + You can write + + >>> (df.pipe(keep_category_a) + ... .pipe(add_one, column="col1") + ... .pipe(multiply, column1="col2", column2="col3") + ... ) + category col1 col2 col3 col4 + 0 A 1 4 2 8 + 1 A 2 5 3 15 + + + If you have a function that takes the data as the second + argument, pass a tuple indicating which keyword expects the + data. For example, suppose ``f`` takes its data as ``df``: + + >>> def multiply_2(column1, df, column2): + ... return df.assign(col4=df[column1] * df[column2]) + + + Then you can write + + >>> (df.pipe(keep_category_a) + ... .pipe(add_one, column="col1") + ... .pipe((multiply_2, 'df'), column1="col2", column2="col3") + ... ) + category col1 col2 col3 col4 + 0 A 1 4 2 8 + 1 A 2 5 3 15 + + You can use lambda as well + + >>> ps.Series([1, 2, 3]).pipe(lambda x: (x + 1).rename("value")) + 0 2 + 1 3 + 2 4 + Name: value, dtype: int64 + """ + + if isinstance(func, tuple): + func, target = func + if target in kwargs: + raise ValueError("%s is both the pipe target and a keyword " "argument" % target) + kwargs[target] = self + return func(*args, **kwargs) + else: + return func(self, *args, **kwargs) + + def to_numpy(self) -> np.ndarray: + """ + A NumPy ndarray representing the values in this DataFrame or Series. + + .. note:: This method should only be used if the resulting NumPy ndarray is expected + to be small, as all the data is loaded into the driver's memory. + + Returns + ------- + numpy.ndarray + + Examples + -------- + >>> ps.DataFrame({"A": [1, 2], "B": [3, 4]}).to_numpy() + array([[1, 3], + [2, 4]]) + + With heterogeneous data, the lowest common type will have to be used. + + >>> ps.DataFrame({"A": [1, 2], "B": [3.0, 4.5]}).to_numpy() + array([[1. , 3. ], + [2. , 4.5]]) + + For a mix of numeric and non-numeric types, the output array will have object dtype. + + >>> df = ps.DataFrame({"A": [1, 2], "B": [3.0, 4.5], "C": pd.date_range('2000', periods=2)}) + >>> df.to_numpy() + array([[1, 3.0, Timestamp('2000-01-01 00:00:00')], + [2, 4.5, Timestamp('2000-01-02 00:00:00')]], dtype=object) + + For Series, + + >>> ps.Series(['a', 'b', 'a']).to_numpy() + array(['a', 'b', 'a'], dtype=object) + """ + log_advice( + "`to_numpy` loads all data into the driver's memory. " + "It should only be used if the resulting NumPy ndarray is expected to be small." + ) + return cast(np.ndarray, self._to_pandas().values) + + @property + def values(self) -> np.ndarray: + """ + Return a Numpy representation of the DataFrame or the Series. + + .. warning:: We recommend using `DataFrame.to_numpy()` or `Series.to_numpy()` instead. + + .. note:: This method should only be used if the resulting NumPy ndarray is expected + to be small, as all the data is loaded into the driver's memory. + + Returns + ------- + numpy.ndarray + + Examples + -------- + A DataFrame where all columns are the same type (e.g., int64) results in an array of + the same type. + + >>> df = ps.DataFrame({'age': [ 3, 29], + ... 'height': [94, 170], + ... 'weight': [31, 115]}) + >>> df + age height weight + 0 3 94 31 + 1 29 170 115 + >>> df.dtypes + age int64 + height int64 + weight int64 + dtype: object + >>> df.values + array([[ 3, 94, 31], + [ 29, 170, 115]]) + + A DataFrame with mixed type columns(e.g., str/object, int64, float32) results in an ndarray + of the broadest type that accommodates these mixed types (e.g., object). + + >>> df2 = ps.DataFrame([('parrot', 24.0, 'second'), + ... ('lion', 80.5, 'first'), + ... ('monkey', np.nan, None)], + ... columns=('name', 'max_speed', 'rank')) + >>> df2.dtypes + name object + max_speed float64 + rank object + dtype: object + >>> df2.values + array([['parrot', 24.0, 'second'], + ['lion', 80.5, 'first'], + ['monkey', nan, None]], dtype=object) + + For Series, + + >>> ps.Series([1, 2, 3]).values + array([1, 2, 3]) + + >>> ps.Series(list('aabc')).values + array(['a', 'a', 'b', 'c'], dtype=object) + """ + warnings.warn("We recommend using `{}.to_numpy()` instead.".format(type(self).__name__)) + return self.to_numpy() + + def to_csv( + self, + path: Optional[str] = None, + sep: str = ",", + na_rep: str = "", + columns: Optional[List[Name]] = None, + header: bool = True, + quotechar: str = '"', + date_format: Optional[str] = None, + escapechar: Optional[str] = None, + num_files: Optional[int] = None, + mode: str = "w", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, + ) -> Optional[str]: + r""" + Write object to a comma-separated values (csv) file. + + .. note:: pandas-on-Spark `to_csv` writes files to a path or URI. Unlike pandas', + pandas-on-Spark respects HDFS's property such as 'fs.default.name'. + + .. note:: pandas-on-Spark writes CSV files into the directory, `path`, and writes + multiple `part-...` files in the directory when `path` is specified. + This behavior was inherited from Apache Spark. The number of partitions can + be controlled by `num_files`. This is deprecated. + Use `DataFrame.spark.repartition` instead. + + Parameters + ---------- + path: str, default None + File path. If None is provided the result is returned as a string. + sep: str, default ',' + String of length 1. Field delimiter for the output file. + na_rep: str, default '' + Missing data representation. + columns: sequence, optional + Columns to write. + header: bool or list of str, default True + Write out the column names. If a list of strings is given it is + assumed to be aliases for the column names. + quotechar: str, default '\"' + String of length 1. Character used to quote fields. + date_format: str, default None + Format string for datetime objects. + escapechar: str, default None + String of length 1. Character used to escape `sep` and `quotechar` + when appropriate. + num_files: the number of partitions to be written in `path` directory when + this is a path. This is deprecated. Use `DataFrame.spark.repartition` instead. + mode: str + Python write mode, default 'w'. + + .. note:: mode can accept the strings for Spark writing mode. + Such as 'append', 'overwrite', 'ignore', 'error', 'errorifexists'. + + - 'append' (equivalent to 'a'): Append the new data to existing data. + - 'overwrite' (equivalent to 'w'): Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols: str or list of str, optional, default None + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + options: keyword arguments for additional options specific to PySpark. + These kwargs are specific to PySpark's CSV options to pass. Check + the options in PySpark's API documentation for spark.write.csv(...). + It has higher priority and overwrites all other options. + This parameter only works when `path` is specified. + + Returns + ------- + str or None + + See Also + -------- + read_csv + DataFrame.to_delta + DataFrame.to_table + DataFrame.to_parquet + DataFrame.to_spark_io + + Examples + -------- + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df.sort_values(by="date") # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + date country code + ... 2012-01-31 12:00:00 KR 1 + ... 2012-02-29 12:00:00 US 2 + ... 2012-03-31 12:00:00 JP 3 + + >>> print(df.to_csv()) # doctest: +NORMALIZE_WHITESPACE + date,country,code + 2012-01-31 12:00:00,KR,1 + 2012-02-29 12:00:00,US,2 + 2012-03-31 12:00:00,JP,3 + + >>> df.cummax().to_csv(path=r'%s/to_csv/foo.csv' % path, num_files=1) + >>> ps.read_csv( + ... path=r'%s/to_csv/foo.csv' % path + ... ).sort_values(by="date") # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + date country code + ... 2012-01-31 12:00:00 KR 1 + ... 2012-02-29 12:00:00 US 2 + ... 2012-03-31 12:00:00 US 3 + + In case of Series, + + >>> print(df.date.to_csv()) # doctest: +NORMALIZE_WHITESPACE + date + 2012-01-31 12:00:00 + 2012-02-29 12:00:00 + 2012-03-31 12:00:00 + + >>> df.date.to_csv(path=r'%s/to_csv/foo.csv' % path, num_files=1) + >>> ps.read_csv( + ... path=r'%s/to_csv/foo.csv' % path + ... ).sort_values(by="date") # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + date + ... 2012-01-31 12:00:00 + ... 2012-02-29 12:00:00 + ... 2012-03-31 12:00:00 + + You can preserve the index in the roundtrip as below. + + >>> df.set_index("country", append=True, inplace=True) + >>> df.date.to_csv( + ... path=r'%s/to_csv/bar.csv' % path, + ... num_files=1, + ... index_col=["index1", "index2"]) + >>> ps.read_csv( + ... path=r'%s/to_csv/bar.csv' % path, index_col=["index1", "index2"] + ... ).sort_values(by="date") # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + date + index1 index2 + ... ... 2012-01-31 12:00:00 + ... ... 2012-02-29 12:00:00 + ... ... 2012-03-31 12:00:00 + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + if path is None: + # If path is none, just collect and use pandas's to_csv. + return self._to_pandas().to_csv( + None, + sep=sep, + na_rep=na_rep, + columns=columns, + header=header, + quotechar=quotechar, + date_format=date_format, + escapechar=escapechar, + index=False, + ) + + if isinstance(self, ps.DataFrame): + psdf = self + else: + assert isinstance(self, ps.Series) + psdf = self.to_frame() + + if columns is None: + column_labels = psdf._internal.column_labels + else: + column_labels = [] + for col in columns: + if is_name_like_tuple(col): + label = cast(Label, col) + else: + label = cast(Label, (col,)) + if label not in psdf._internal.column_labels: + raise KeyError(name_like_string(label)) + column_labels.append(label) + + if isinstance(index_col, str): + index_cols = [index_col] + elif index_col is None: + index_cols = [] + else: + index_cols = index_col + + if header is True and psdf._internal.column_labels_level > 1: + raise ValueError("to_csv only support one-level index column now") + elif isinstance(header, list): + sdf = psdf.to_spark(index_col) + sdf = sdf.select( + [scol_for(sdf, name_like_string(label)) for label in index_cols] + + [ + scol_for(sdf, str(i) if label is None else name_like_string(label)).alias( + new_name + ) + for i, (label, new_name) in enumerate(zip(column_labels, header)) + ] + ) + header = True + else: + sdf = psdf.to_spark(index_col) + sdf = sdf.select( + [scol_for(sdf, name_like_string(label)) for label in index_cols] + + [ + scol_for(sdf, str(i) if label is None else name_like_string(label)) + for i, label in enumerate(column_labels) + ] + ) + + if num_files is not None: + warnings.warn( + "`num_files` has been deprecated and might be removed in a future version. " + "Use `DataFrame.spark.repartition` instead.", + FutureWarning, + ) + sdf = sdf.repartition(num_files) + + mode = validate_mode(mode) + builder = sdf.write.mode(mode) + if partition_cols is not None: + builder.partitionBy(partition_cols) + builder._set_opts( + sep=sep, + nullValue=na_rep, + header=header, + quote=quotechar, + dateFormat=date_format, + charToEscapeQuoteEscaping=escapechar, + ) + builder.options(**options).format("csv").save(path) + return None + + def to_json( + self, + path: Optional[str] = None, + compression: str = "uncompressed", + num_files: Optional[int] = None, + mode: str = "w", + orient: str = "records", + lines: bool = True, + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, + ) -> Optional[str]: + """ + Convert the object to a JSON string. + + .. note:: pandas-on-Spark `to_json` writes files to a path or URI. Unlike pandas', + pandas-on-Spark respects HDFS's property such as 'fs.default.name'. + + .. note:: pandas-on-Spark writes JSON files into the directory, `path`, and writes + multiple `part-...` files in the directory when `path` is specified. + This behavior was inherited from Apache Spark. The number of partitions can + be controlled by `num_files`. This is deprecated. + Use `DataFrame.spark.repartition` instead. + + .. note:: output JSON format is different from pandas'. It always uses `orient='records'` + for its output. This behavior might have to change soon. + + .. note:: Set `ignoreNullFields` keyword argument to `True` to omit `None` or `NaN` values + when writing JSON objects. It works only when `path` is provided. + + Note NaN's and None will be converted to null and datetime objects + will be converted to UNIX timestamps. + + Parameters + ---------- + path: string, optional + File path. If not specified, the result is returned as + a string. + lines: bool, default True + If ‘orient’ is ‘records’ write out line delimited JSON format. + Will throw ValueError if incorrect ‘orient’ since others are not + list like. It should be always True for now. + orient: str, default 'records' + It should be always 'records' for now. + compression: {'gzip', 'bz2', 'xz', None} + A string representing the compression to use in the output file, + only used when the first argument is a filename. By default, the + compression is inferred from the filename. + num_files: the number of partitions to be written in `path` directory when + this is a path. This is deprecated. Use `DataFrame.spark.repartition` instead. + mode: str + Python write mode, default 'w'. + + .. note:: mode can accept the strings for Spark writing mode. + Such as 'append', 'overwrite', 'ignore', 'error', 'errorifexists'. + + - 'append' (equivalent to 'a'): Append the new data to existing data. + - 'overwrite' (equivalent to 'w'): Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols: str or list of str, optional, default None + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + options: keyword arguments for additional options specific to PySpark. + It is specific to PySpark's JSON options to pass. Check + the options in PySpark's API documentation for `spark.write.json(...)`. + It has a higher priority and overwrites all other options. + This parameter only works when `path` is specified. + + Returns + ------- + str or None + + Examples + -------- + >>> df = ps.DataFrame([['a', 'b'], ['c', 'd']], + ... columns=['col 1', 'col 2']) + >>> df.to_json() + '[{"col 1":"a","col 2":"b"},{"col 1":"c","col 2":"d"}]' + + >>> df['col 1'].to_json() + '[{"col 1":"a"},{"col 1":"c"}]' + + >>> df.to_json(path=r'%s/to_json/foo.json' % path, num_files=1) + >>> ps.read_json( + ... path=r'%s/to_json/foo.json' % path + ... ).sort_values(by="col 1") + col 1 col 2 + 0 a b + 1 c d + + >>> df['col 1'].to_json(path=r'%s/to_json/foo.json' % path, num_files=1, index_col="index") + >>> ps.read_json( + ... path=r'%s/to_json/foo.json' % path, index_col="index" + ... ).sort_values(by="col 1") # doctest: +NORMALIZE_WHITESPACE + col 1 + index + 0 a + 1 c + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + default_options: Dict[str, Any] = {"ignoreNullFields": False} + options = {**default_options, **options} + + if not lines: + raise NotImplementedError("lines=False is not implemented yet.") + + if orient != "records": + raise NotImplementedError("orient='records' is supported only for now.") + + if path is None: + # If path is none, just collect and use pandas's to_json. + psdf_or_ser = self + pdf = psdf_or_ser._to_pandas() + if isinstance(self, ps.Series): + pdf = pdf.to_frame() + # To make the format consistent and readable by `read_json`, convert it to pandas' and + # use 'records' orient for now. + return pdf.to_json(orient="records") + + if isinstance(self, ps.DataFrame): + psdf = self + else: + assert isinstance(self, ps.Series) + psdf = self.to_frame() + sdf = psdf.to_spark(index_col=index_col) + + if num_files is not None: + warnings.warn( + "`num_files` has been deprecated and might be removed in a future version. " + "Use `DataFrame.spark.repartition` instead.", + FutureWarning, + ) + sdf = sdf.repartition(num_files) + + mode = validate_mode(mode) + builder = sdf.write.mode(mode) + if partition_cols is not None: + builder.partitionBy(partition_cols) + builder._set_opts(compression=compression) + builder.options(**options).format("json").save(path) + return None + + def to_excel( + self, + excel_writer: Union[str, pd.ExcelWriter], + sheet_name: str = "Sheet1", + na_rep: str = "", + float_format: Optional[str] = None, + columns: Optional[Union[str, List[str]]] = None, + header: bool = True, + index: bool = True, + index_label: Optional[Union[str, List[str]]] = None, + startrow: int = 0, + startcol: int = 0, + engine: Optional[str] = None, + merge_cells: bool = True, + encoding: Optional[str] = None, + inf_rep: str = "inf", + verbose: bool = True, + freeze_panes: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Write object to an Excel sheet. + + .. note:: This method should only be used if the resulting DataFrame is expected + to be small, as all the data is loaded into the driver's memory. + + To write a single object to an Excel .xlsx file it is only necessary to + specify a target file name. To write to multiple sheets it is necessary to + create an `ExcelWriter` object with a target file name, and specify a sheet + in the file to write to. + + Multiple sheets may be written to by specifying unique `sheet_name`. + With all data written to the file it is necessary to save the changes. + Note that creating an `ExcelWriter` object with a file name that already + exists will result in the contents of the existing file being erased. + + Parameters + ---------- + excel_writer: str or ExcelWriter object + File path or existing ExcelWriter. + sheet_name: str, default 'Sheet1' + Name of sheet which will contain DataFrame. + na_rep: str, default '' + Missing data representation. + float_format: str, optional + Format string for floating point numbers. For example + ``float_format="%%.2f"`` will format 0.1234 to 0.12. + columns: sequence or list of str, optional + Columns to write. + header: bool or list of str, default True + Write out the column names. If a list of string is given it is + assumed to be aliases for the column names. + index: bool, default True + Write row names (index). + index_label: str or sequence, optional + Column label for index column(s) if desired. If not specified, and + `header` and `index` are True, then the index names are used. A + sequence should be given if the DataFrame uses MultiIndex. + startrow: int, default 0 + Upper left cell row to dump data frame. + startcol: int, default 0 + Upper left cell column to dump data frame. + engine: str, optional + Write engine to use, 'openpyxl' or 'xlsxwriter'. You can also set this + via the options ``io.excel.xlsx.writer``, ``io.excel.xls.writer``, and + ``io.excel.xlsm.writer``. + merge_cells: bool, default True + Write MultiIndex and Hierarchical Rows as merged cells. + encoding: str, optional + Encoding of the resulting excel file. Only necessary for xlwt, + other writers support unicode natively. + + .. deprecated:: 3.4.0 + + inf_rep: str, default 'inf' + Representation for infinity (there is no native representation for + infinity in Excel). + verbose: bool, default True + Display more information in the error logs. + + .. deprecated:: 3.4.0 + + freeze_panes: tuple of int (length 2), optional + Specifies the one-based bottommost row and rightmost column that + is to be frozen. + + Notes + ----- + Once a workbook has been saved it is not possible write further data + without rewriting the whole workbook. + + See Also + -------- + read_excel: Read Excel file. + + Examples + -------- + Create, write to, and save a workbook: + + >>> df1 = ps.DataFrame([['a', 'b'], ['c', 'd']], + ... index=['row 1', 'row 2'], + ... columns=['col 1', 'col 2']) + >>> df1.to_excel("output.xlsx") # doctest: +SKIP + + To specify the sheet name: + + >>> df1.to_excel("output.xlsx") # doctest: +SKIP + >>> df1.to_excel("output.xlsx", + ... sheet_name='Sheet_name_1') # doctest: +SKIP + + If you wish to write to more than one sheet in the workbook, it is + necessary to specify an ExcelWriter object: + + >>> with pd.ExcelWriter('output.xlsx') as writer: # doctest: +SKIP + ... df1.to_excel(writer, sheet_name='Sheet_name_1') + ... df2.to_excel(writer, sheet_name='Sheet_name_2') + + To set the library that is used to write the Excel file, + you can pass the `engine` keyword (the default engine is + automatically chosen depending on the file extension): + + >>> df1.to_excel('output1.xlsx', engine='xlsxwriter') # doctest: +SKIP + """ + log_advice( + "`to_excel` loads all data into the driver's memory. " + "It should only be used if the resulting DataFrame is expected to be small." + ) + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + psdf = self + + if isinstance(self, ps.DataFrame): + f = pd.DataFrame.to_excel + elif isinstance(self, ps.Series): + f = pd.Series.to_excel + else: + raise TypeError( + "Constructor expects DataFrame or Series; however, " "got [%s]" % (self,) + ) + return validate_arguments_and_invoke_function( + psdf._to_internal_pandas(), self.to_excel, f, args + ) + + def mean( + self, axis: Optional[Axis] = None, skipna: bool = True, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return the mean of the values. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + mean: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.mean() + a 2.0 + b 0.2 + dtype: float64 + + >>> df.mean(axis=1) + 0 0.55 + 1 1.10 + 2 1.65 + 3 NaN + dtype: float64 + + On a Series: + + >>> df['a'].mean() + 2.0 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def mean(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + return F.mean(spark_column) + + return self._reduce_for_stat_function( + mean, + name="mean", + axis=axis, + numeric_only=numeric_only, + skipna=skipna, + ) + + def sum( + self, + axis: Optional[Axis] = None, + skipna: bool = True, + numeric_only: bool = None, + min_count: int = 0, + ) -> Union[Scalar, "Series"]: + """ + Return the sum of the values. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Added *skipna* to exclude. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + min_count: int, default 0 + The required number of valid values to perform the operation. If fewer than + ``min_count`` non-NA values are present the result will be NA. + + Returns + ------- + sum: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, np.nan, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.sum() + a 6.0 + b 0.4 + dtype: float64 + + >>> df.sum(axis=1) + 0 1.1 + 1 2.0 + 2 3.3 + 3 0.0 + dtype: float64 + + >>> df.sum(min_count=3) + a 6.0 + b NaN + dtype: float64 + + >>> df.sum(axis=1, min_count=1) + 0 1.1 + 1 2.0 + 2 3.3 + 3 NaN + dtype: float64 + + On a Series: + + >>> df['a'].sum() + 6.0 + + >>> df['a'].sum(min_count=3) + 6.0 + >>> df['b'].sum(min_count=3) + nan + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + elif numeric_only is True and axis == 1: + numeric_only = None + + def sum(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + return F.coalesce(F.sum(spark_column), F.lit(0)) + + return self._reduce_for_stat_function( + sum, + name="sum", + axis=axis, + numeric_only=numeric_only, + min_count=min_count, + skipna=skipna, + ) + + def product( + self, + axis: Optional[Axis] = None, + skipna: bool = True, + numeric_only: bool = None, + min_count: int = 0, + ) -> Union[Scalar, "Series"]: + """ + Return the product of the values. + + .. note:: unlike pandas', pandas-on-Spark's emulates product by ``exp(sum(log(...)))`` + trick. Therefore, it only works for positive numbers. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + min_count: int, default 0 + The required number of valid values to perform the operation. If fewer than + ``min_count`` non-NA values are present the result will be NA. + + Examples + -------- + On a DataFrame: + + Non-numeric type column is not included to the result. + + >>> psdf = ps.DataFrame({'A': [1, 2, 3, 4, 5], + ... 'B': [10, 20, 30, 40, 50], + ... 'C': ['a', 'b', 'c', 'd', 'e']}) + >>> psdf + A B C + 0 1 10 a + 1 2 20 b + 2 3 30 c + 3 4 40 d + 4 5 50 e + + >>> psdf.prod() + A 120 + B 12000000 + dtype: int64 + + If there is no numeric type columns, returns empty Series. + + >>> ps.DataFrame({"key": ['a', 'b', 'c'], "val": ['x', 'y', 'z']}).prod() # doctest: +SKIP + Series([], dtype: float64) + + On a Series: + + >>> ps.Series([1, 2, 3, 4, 5]).prod() + 120 + + By default, the product of an empty or all-NA Series is ``1`` + + >>> ps.Series([]).prod() # doctest: +SKIP + 1.0 + + This can be controlled with the ``min_count`` parameter + + >>> ps.Series([]).prod(min_count=1) # doctest: +SKIP + nan + """ + axis = validate_axis(axis) + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `None` in 4.0.0.", + FutureWarning, + ) + + if numeric_only is None and axis == 0: + numeric_only = True + elif numeric_only is True and axis == 1: + numeric_only = None + + def prod(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + return SF.product(spark_column, skipna) + + return self._reduce_for_stat_function( + prod, + name="prod", + axis=axis, + numeric_only=numeric_only, + min_count=min_count, + skipna=skipna, + ) + + prod = product + + def skew( + self, axis: Optional[Axis] = None, skipna: bool = True, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return unbiased skew normalized by N-1. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + skew: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.skew() + a 0.0 + b 0.0 + dtype: float64 + + On a Series: + + >>> df['a'].skew() + 0.0 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def skew(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + return SF.skew(spark_column) + + return self._reduce_for_stat_function( + skew, + name="skew", + axis=axis, + numeric_only=numeric_only, + skipna=skipna, + ) + + def kurtosis( + self, axis: Optional[Axis] = None, skipna: bool = True, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return unbiased kurtosis using Fisher’s definition of kurtosis (kurtosis of normal == 0.0). + Normalized by N-1. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + kurt: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan, 6], 'b': [0.1, 0.2, 0.3, np.nan, 0.8]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.kurtosis() + a 1.500000 + b 2.703924 + dtype: float64 + + On a Series: + + >>> df['a'].kurtosis() + 1.5 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def kurtosis(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + return SF.kurt(spark_column) + + return self._reduce_for_stat_function( + kurtosis, + name="kurtosis", + axis=axis, + numeric_only=numeric_only, + skipna=skipna, + ) + + kurt = kurtosis + + def min( + self, axis: Optional[Axis] = None, skipna: bool = True, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return the minimum of the values. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + If True, include only float, int, boolean columns. This parameter is mainly for + pandas compatibility. False is supported; however, the columns should + be all numeric or all non-numeric. + + Returns + ------- + min: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.min() + a 1.0 + b 0.1 + dtype: float64 + + >>> df.min(axis=1) + 0 0.1 + 1 0.2 + 2 0.3 + 3 NaN + dtype: float64 + + On a Series: + + >>> df['a'].min() + 1.0 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + elif numeric_only is True and axis == 1: + numeric_only = None + + return self._reduce_for_stat_function( + lambda psser: F.min(psser.spark.column), + name="min", + axis=axis, + numeric_only=numeric_only, + skipna=skipna, + ) + + def max( + self, axis: Optional[Axis] = None, skipna: bool = True, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return the maximum of the values. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + If True, include only float, int, boolean columns. This parameter is mainly for + pandas compatibility. False is supported; however, the columns should + be all numeric or all non-numeric. + + Returns + ------- + max: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.max() + a 3.0 + b 0.3 + dtype: float64 + + >>> df.max(axis=1) + 0 1.0 + 1 2.0 + 2 3.0 + 3 NaN + dtype: float64 + + On a Series: + + >>> df['a'].max() + 3.0 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + elif numeric_only is True and axis == 1: + numeric_only = None + + return self._reduce_for_stat_function( + lambda psser: F.max(psser.spark.column), + name="max", + axis=axis, + numeric_only=numeric_only, + skipna=skipna, + ) + + def count( + self, axis: Optional[Axis] = None, numeric_only: bool = False + ) -> Union[Scalar, "Series"]: + """ + Count non-NA cells for each column. + + The values `None`, `NaN` are considered NA. + + Parameters + ---------- + axis: {0 or ‘index’, 1 or ‘columns’}, default 0 + If 0 or ‘index’ counts are generated for each column. If 1 or ‘columns’ counts are + generated for each row. + numeric_only: bool, default False + If True, include only float, int, boolean columns. This parameter is mainly for + pandas compatibility. + + Returns + ------- + max: scalar for a Series, and a Series for a DataFrame. + + See Also + -------- + DataFrame.shape: Number of DataFrame rows and columns (including NA + elements). + DataFrame.isna: Boolean same-sized DataFrame showing places of NA + elements. + + Examples + -------- + Constructing DataFrame from a dictionary: + + >>> df = ps.DataFrame({"Person": + ... ["John", "Myla", "Lewis", "John", "Myla"], + ... "Age": [24., np.nan, 21., 33, 26], + ... "Single": [False, True, True, True, False]}, + ... columns=["Person", "Age", "Single"]) + >>> df + Person Age Single + 0 John 24.0 False + 1 Myla NaN True + 2 Lewis 21.0 True + 3 John 33.0 True + 4 Myla 26.0 False + + Notice the uncounted NA values: + + >>> df.count() + Person 5 + Age 4 + Single 5 + dtype: int64 + + >>> df.count(axis=1) + 0 3 + 1 2 + 2 3 + 3 3 + 4 3 + dtype: int64 + + On a Series: + + >>> df['Person'].count() + 5 + + >>> df['Age'].count() + 4 + """ + + return self._reduce_for_stat_function( + Frame._count_expr, name="count", axis=axis, numeric_only=numeric_only + ) + + def std( + self, + axis: Optional[Axis] = None, + skipna: bool = True, + ddof: int = 1, + numeric_only: bool = None, + ) -> Union[Scalar, "Series"]: + """ + Return sample standard deviation. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + ddof: int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + .. versionchanged:: 3.4.0 + Supported including arbitary integers. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + std: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.std() + a 1.0 + b 0.1 + dtype: float64 + + >>> df.std(ddof=2) + a 1.414214 + b 0.141421 + dtype: float64 + + >>> df.std(axis=1) + 0 0.636396 + 1 1.272792 + 2 1.909188 + 3 NaN + dtype: float64 + + >>> df.std(ddof=0) + a 0.816497 + b 0.081650 + dtype: float64 + + On a Series: + + >>> df['a'].std() + 1.0 + + >>> df['a'].std(ddof=0) + 0.816496580927726 + + >>> df['a'].std(ddof=-1) + 0.707106... + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def std(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + return SF.stddev(spark_column, ddof) + + return self._reduce_for_stat_function( + std, name="std", axis=axis, numeric_only=numeric_only, ddof=ddof, skipna=skipna + ) + + def var( + self, axis: Optional[Axis] = None, ddof: int = 1, numeric_only: bool = None + ) -> Union[Scalar, "Series"]: + """ + Return unbiased variance. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + ddof: int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + .. versionchanged:: 3.4.0 + Supported including arbitary integers. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + var: scalar for a Series, and a Series for a DataFrame. + + Examples + -------- + + >>> df = ps.DataFrame({'a': [1, 2, 3, np.nan], 'b': [0.1, 0.2, 0.3, np.nan]}, + ... columns=['a', 'b']) + + On a DataFrame: + + >>> df.var() + a 1.00 + b 0.01 + dtype: float64 + + >>> df.var(ddof=2) + a 2.00 + b 0.02 + dtype: float64 + + >>> df.var(axis=1) + 0 0.405 + 1 1.620 + 2 3.645 + 3 NaN + dtype: float64 + + >>> df.var(ddof=0) + a 0.666667 + b 0.006667 + dtype: float64 + + On a Series: + + >>> df['a'].var() + 1.0 + + >>> df['a'].var(ddof=0) + 0.6666666666666666 + + >>> df['a'].var(ddof=-2) + 0.4 + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def var(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + return SF.var(spark_column, ddof) + + return self._reduce_for_stat_function( + var, name="var", axis=axis, numeric_only=numeric_only, ddof=ddof + ) + + def median( + self, + axis: Optional[Axis] = None, + skipna: bool = True, + numeric_only: bool = None, + accuracy: int = 10000, + ) -> Union[Scalar, "Series"]: + """ + Return the median of the values for the requested axis. + + .. note:: Unlike pandas', the median in pandas-on-Spark is an approximated median based upon + approximate percentile computation because computing median across a large dataset + is extremely expensive. + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + accuracy: int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + + Returns + ------- + median: scalar or Series + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'a': [24., 21., 25., 33., 26.], 'b': [1, 2, 3, 4, 5]}, columns=['a', 'b']) + >>> df + a b + 0 24.0 1 + 1 21.0 2 + 2 25.0 3 + 3 33.0 4 + 4 26.0 5 + + On a DataFrame: + + >>> df.median() + a 25.0 + b 3.0 + dtype: float64 + + On a Series: + + >>> df['a'].median() + 25.0 + >>> (df['b'] + 100).median() + 103.0 + + For multi-index columns, + + >>> df.columns = pd.MultiIndex.from_tuples([('x', 'a'), ('y', 'b')]) + >>> df + x y + a b + 0 24.0 1 + 1 21.0 2 + 2 25.0 3 + 3 33.0 4 + 4 26.0 5 + + On a DataFrame: + + >>> df.median() + x a 25.0 + y b 3.0 + dtype: float64 + + >>> df.median(axis=1) + 0 12.5 + 1 11.5 + 2 14.0 + 3 18.5 + 4 15.5 + dtype: float64 + + On a Series: + + >>> df[('x', 'a')].median() + 25.0 + >>> (df[('y', 'b')] + 100).median() + 103.0 + """ + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + if not isinstance(accuracy, int): + raise TypeError( + "accuracy must be an integer; however, got [%s]" % type(accuracy).__name__ + ) + + def median(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, (BooleanType, NumericType)): + return F.percentile_approx(spark_column.cast(DoubleType()), 0.5, accuracy) + else: + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + return self._reduce_for_stat_function( + median, + name="median", + numeric_only=numeric_only, + axis=axis, + skipna=skipna, + ) + + def sem( + self, + axis: Optional[Axis] = None, + skipna: bool = True, + ddof: int = 1, + numeric_only: bool = None, + ) -> Union[Scalar, "Series"]: + """ + Return unbiased standard error of the mean over requested axis. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + axis: {index (0), columns (1)} + Axis for the function to be applied on. + skipna: bool, default True + Exclude NA/null values when computing the result. + + .. versionchanged:: 3.4.0 + Supported including NA/null values. + ddof: int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + .. versionchanged:: 3.4.0 + Supported including arbitary integers. + numeric_only: bool, default None + Include only float, int, boolean columns. False is not supported. This parameter + is mainly for pandas compatibility. + + Returns + ------- + scalar(for Series) or Series(for DataFrame) + + Examples + -------- + >>> psdf = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + >>> psdf + a b + 0 1 4 + 1 2 5 + 2 3 6 + + >>> psdf.sem() + a 0.57735 + b 0.57735 + dtype: float64 + + >>> psdf.sem(ddof=0) + a 0.471405 + b 0.471405 + dtype: float64 + + >>> psdf.sem(ddof=2) + a 0.816497 + b 0.816497 + dtype: float64 + + >>> psdf.sem(axis=1) + 0 1.5 + 1 1.5 + 2 1.5 + dtype: float64 + + Support for Series + + >>> psser = psdf.a + >>> psser + 0 1 + 1 2 + 2 3 + Name: a, dtype: int64 + + >>> psser.sem() + 0.5773502691896258 + + >>> psser.sem(ddof=0) + 0.47140452079103173 + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + axis = validate_axis(axis) + + if numeric_only is None and axis == 0: + numeric_only = True + + def std(psser: "Series") -> Column: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, BooleanType): + spark_column = spark_column.cast(LongType()) + elif not isinstance(spark_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + return SF.stddev(spark_column, ddof) + + def sem(psser: "Series") -> Column: + return std(psser) / F.sqrt(Frame._count_expr(psser)) + + return self._reduce_for_stat_function( + sem, + name="sem", + numeric_only=numeric_only, + axis=axis, + ddof=ddof, + skipna=skipna, + ) + + @property + def size(self) -> int: + """ + Return an int representing the number of elements in this object. + + Return the number of rows if Series. Otherwise return the number of + rows times number of columns if DataFrame. + + Examples + -------- + >>> s = ps.Series({'a': 1, 'b': 2, 'c': None}) + >>> s.size + 3 + + >>> df = ps.DataFrame({'col1': [1, 2, None], 'col2': [3, 4, None]}) + >>> df.size + 6 + + >>> df = ps.DataFrame(index=[1, 2, None]) + >>> df.size + 0 + """ + num_columns = len(self._internal.data_spark_columns) + if num_columns == 0: + return 0 + else: + return len(self) * num_columns # type: ignore[arg-type] + + def abs(self: FrameLike) -> FrameLike: + """ + Return a Series/DataFrame with absolute numeric value of each element. + + Returns + ------- + abs: Series/DataFrame containing the absolute value of each element. + + Examples + -------- + + Absolute numeric values in a Series. + + >>> s = ps.Series([-1.10, 2, -3.33, 4]) + >>> s.abs() + 0 1.10 + 1 2.00 + 2 3.33 + 3 4.00 + dtype: float64 + + Absolute numeric values in a DataFrame. + + >>> df = ps.DataFrame({ + ... 'a': [4, 5, 6, 7], + ... 'b': [10, 20, 30, 40], + ... 'c': [100, 50, -30, -50] + ... }, + ... columns=['a', 'b', 'c']) + >>> df.abs() + a b c + 0 4 10 100 + 1 5 20 50 + 2 6 30 30 + 3 7 40 50 + """ + + def abs(psser: "Series") -> Union["Series", Column]: + if isinstance(psser.spark.data_type, BooleanType): + return psser + elif isinstance(psser.spark.data_type, NumericType): + return psser._with_new_scol( + F.abs(psser.spark.column), field=psser._internal.data_fields[0] + ) + else: + raise TypeError( + "bad operand type for abs(): {} ({})".format( + spark_type_to_pandas_dtype(psser.spark.data_type), + psser.spark.data_type.simpleString(), + ) + ) + + return self._apply_series_op(abs) + + # TODO: by argument only support the grouping name and as_index only for now. Documentation + # should be updated when it's supported. + def groupby( + self: FrameLike, + by: Union[Name, "Series", List[Union[Name, "Series"]]], + axis: Axis = 0, + as_index: bool = True, + dropna: bool = True, + ) -> "GroupBy[FrameLike]": + """ + Group DataFrame or Series using one or more columns. + + A groupby operation involves some combination of splitting the + object, applying a function, and combining the results. This can be + used to group large amounts of data and compute operations on these + groups. + + Parameters + ---------- + by: Series, label, or list of labels + Used to determine the groups for the groupby. + If Series is passed, the Series or dict VALUES + will be used to determine the groups. A label or list of + labels may be passed to group by the columns in ``self``. + axis: int, default 0 or 'index' + Can only be set to 0 now. + as_index: bool, default True + For aggregated output, return object with group labels as the + index. Only relevant for DataFrame input. as_index=False is + effectively "SQL-style" grouped output. + dropna: bool, default True + If True, and if group keys contain NA values, + NA values together with row/column will be dropped. + If False, NA values will also be treated as the key in groups. + + Returns + ------- + DataFrameGroupBy or SeriesGroupBy + Depends on the calling object and returns groupby object that + contains information about the groups. + + See Also + -------- + pyspark.pandas.groupby.GroupBy + + Examples + -------- + >>> df = ps.DataFrame({'Animal': ['Falcon', 'Falcon', + ... 'Parrot', 'Parrot'], + ... 'Max Speed': [380., 370., 24., 26.]}, + ... columns=['Animal', 'Max Speed']) + >>> df + Animal Max Speed + 0 Falcon 380.0 + 1 Falcon 370.0 + 2 Parrot 24.0 + 3 Parrot 26.0 + + >>> df.groupby(['Animal']).mean().sort_index() # doctest: +NORMALIZE_WHITESPACE + Max Speed + Animal + Falcon 375.0 + Parrot 25.0 + + >>> df.groupby(['Animal'], as_index=False).mean().sort_values('Animal') + ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE + Animal Max Speed + ...Falcon 375.0 + ...Parrot 25.0 + + We can also choose to include NA in group keys or not by setting dropna parameter, + the default setting is True: + + >>> l = [[1, 2, 3], [1, None, 4], [2, 1, 3], [1, 2, 2]] + >>> df = ps.DataFrame(l, columns=["a", "b", "c"]) + >>> df.groupby(by=["b"]).sum().sort_index() # doctest: +NORMALIZE_WHITESPACE + a c + b + 1.0 2 3 + 2.0 2 5 + + >>> df.groupby(by=["b"], dropna=False).sum().sort_index() # doctest: +NORMALIZE_WHITESPACE + a c + b + 1.0 2 3 + 2.0 2 5 + NaN 1 4 + """ + new_by: List[Union[Label, ps.Series]] + if isinstance(by, ps.DataFrame): + raise ValueError("Grouper for '{}' not 1-dimensional".format(type(by).__name__)) + elif isinstance(by, ps.Series): + new_by = [by] + elif is_name_like_tuple(by): + if isinstance(self, ps.Series): + raise KeyError(by) + new_by = [cast(Label, by)] + elif is_name_like_value(by): + if isinstance(self, ps.Series): + raise KeyError(by) + new_by = [cast(Label, (by,))] + elif is_list_like(by): + new_by = [] + for key in by: + if isinstance(key, ps.DataFrame): + raise ValueError( + "Grouper for '{}' not 1-dimensional".format(type(key).__name__) + ) + elif isinstance(key, ps.Series): + new_by.append(key) + elif is_name_like_tuple(key): + if isinstance(self, ps.Series): + raise KeyError(key) + new_by.append(cast(Label, key)) + elif is_name_like_value(key): + if isinstance(self, ps.Series): + raise KeyError(key) + new_by.append(cast(Label, (key,))) + else: + raise ValueError( + "Grouper for '{}' not 1-dimensional".format(type(key).__name__) + ) + else: + raise ValueError("Grouper for '{}' not 1-dimensional".format(type(by).__name__)) + if not len(new_by): + raise ValueError("No group keys passed!") + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + return self._build_groupby(by=new_by, as_index=as_index, dropna=dropna) + + @abstractmethod + def _build_groupby( + self: FrameLike, by: List[Union["Series", Label]], as_index: bool, dropna: bool + ) -> "GroupBy[FrameLike]": + pass + + def bool(self) -> bool: + """ + Return the bool of a single element in the current object. + + This must be a boolean scalar value, either True or False. Raise a ValueError if + the object does not have exactly 1 element, or that element is not boolean + + Returns + ------- + bool + + Examples + -------- + >>> ps.DataFrame({'a': [True]}).bool() + True + + >>> ps.Series([False]).bool() + False + + If there are non-boolean or multiple values exist, it raises an exception in all + cases as below. + + >>> ps.DataFrame({'a': ['a']}).bool() + Traceback (most recent call last): + ... + ValueError: bool cannot act on a non-boolean single element DataFrame + + >>> ps.DataFrame({'a': [True], 'b': [False]}).bool() # doctest: +NORMALIZE_WHITESPACE + Traceback (most recent call last): + ... + ValueError: The truth value of a DataFrame is ambiguous. Use a.empty, a.bool(), + a.item(), a.any() or a.all(). + + >>> ps.Series([1]).bool() + Traceback (most recent call last): + ... + ValueError: bool cannot act on a non-boolean single element DataFrame + """ + if isinstance(self, ps.DataFrame): + df = self + elif isinstance(self, ps.Series): + df = self.to_dataframe() + return df.head(2)._to_internal_pandas().bool() + + def first_valid_index(self) -> Optional[Union[Scalar, Tuple[Scalar, ...]]]: + """ + Retrieves the index of the first valid value. + + Returns + ------- + scalar, tuple, or None + + Examples + -------- + + Support for DataFrame + + >>> psdf = ps.DataFrame({'a': [None, 2, 3, 2], + ... 'b': [None, 2.0, 3.0, 1.0], + ... 'c': [None, 200, 400, 200]}, + ... index=['Q', 'W', 'E', 'R']) + >>> psdf + a b c + Q NaN NaN NaN + W 2.0 2.0 200.0 + E 3.0 3.0 400.0 + R 2.0 1.0 200.0 + + >>> psdf.first_valid_index() + 'W' + + Support for MultiIndex columns + + >>> psdf.columns = pd.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> psdf + a b c + x y z + Q NaN NaN NaN + W 2.0 2.0 200.0 + E 3.0 3.0 400.0 + R 2.0 1.0 200.0 + + >>> psdf.first_valid_index() + 'W' + + Support for Series. + + >>> s = ps.Series([None, None, 3, 4, 5], index=[100, 200, 300, 400, 500]) + >>> s + 100 NaN + 200 NaN + 300 3.0 + 400 4.0 + 500 5.0 + dtype: float64 + + >>> s.first_valid_index() + 300 + + Support for MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([None, None, None, None, 250, 1.5, 320, 1, 0.3], index=midx) + >>> s + lama speed NaN + weight NaN + length NaN + cow speed NaN + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.first_valid_index() + ('cow', 'weight') + """ + data_spark_columns = self._internal.data_spark_columns + + if len(data_spark_columns) == 0: + return None + + cond = reduce(lambda x, y: x & y, map(lambda x: x.isNotNull(), data_spark_columns)) + + with sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + # Disable Arrow to keep row ordering. + first_valid_row = ( + self._internal.spark_frame.filter(cond) + .select(self._internal.index_spark_columns) + .limit(1) + .toPandas() + ) + + # For Empty Series or DataFrame, returns None. + if len(first_valid_row) == 0: + return None + + first_valid_row = first_valid_row.iloc[0] + if len(first_valid_row) == 1: + return first_valid_row.iloc[0] + else: + return tuple(first_valid_row) + + def last_valid_index(self) -> Optional[Union[Scalar, Tuple[Scalar, ...]]]: + """ + Return index for last non-NA/null value. + + Returns + ------- + scalar, tuple, or None + + Notes + ----- + This API only works with PySpark >= 3.0. + + Examples + -------- + + Support for DataFrame + + >>> psdf = ps.DataFrame({'a': [1, 2, 3, None], + ... 'b': [1.0, 2.0, 3.0, None], + ... 'c': [100, 200, 400, None]}, + ... index=['Q', 'W', 'E', 'R']) + >>> psdf + a b c + Q 1.0 1.0 100.0 + W 2.0 2.0 200.0 + E 3.0 3.0 400.0 + R NaN NaN NaN + + >>> psdf.last_valid_index() # doctest: +SKIP + 'E' + + Support for MultiIndex columns + + >>> psdf.columns = pd.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> psdf + a b c + x y z + Q 1.0 1.0 100.0 + W 2.0 2.0 200.0 + E 3.0 3.0 400.0 + R NaN NaN NaN + + >>> psdf.last_valid_index() # doctest: +SKIP + 'E' + + Support for Series. + + >>> s = ps.Series([1, 2, 3, None, None], index=[100, 200, 300, 400, 500]) + >>> s + 100 1.0 + 200 2.0 + 300 3.0 + 400 NaN + 500 NaN + dtype: float64 + + >>> s.last_valid_index() # doctest: +SKIP + 300 + + Support for MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([250, 1.5, 320, 1, 0.3, None, None, None, None], index=midx) + >>> s + lama speed 250.0 + weight 1.5 + length 320.0 + cow speed 1.0 + weight 0.3 + length NaN + falcon speed NaN + weight NaN + length NaN + dtype: float64 + + >>> s.last_valid_index() # doctest: +SKIP + ('cow', 'weight') + """ + data_spark_columns = self._internal.data_spark_columns + + if len(data_spark_columns) == 0: + return None + + cond = reduce(lambda x, y: x & y, map(lambda x: x.isNotNull(), data_spark_columns)) + + last_valid_rows = ( + self._internal.spark_frame.filter(cond) + .select(self._internal.index_spark_columns) + .tail(1) + ) + + # For Empty Series or DataFrame, returns None. + if len(last_valid_rows) == 0: + return None + + last_valid_row = last_valid_rows[0] + + if len(last_valid_row) == 1: + return last_valid_row[0] + else: + return tuple(last_valid_row) + + # TODO: 'center', 'win_type', 'on', 'axis' parameter should be implemented. + def rolling( + self: FrameLike, window: int, min_periods: Optional[int] = None + ) -> "Rolling[FrameLike]": + """ + Provide rolling transformations. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + Parameters + ---------- + window: int, or offset + Size of the moving window. + This is the number of observations used for calculating the statistic. + Each window will be a fixed size. + + min_periods: int, default None + Minimum number of observations in window required to have a value + (otherwise result is NA). + For a window that is specified by an offset, min_periods will default to 1. + Otherwise, min_periods will default to the size of the window. + + Returns + ------- + a Window sub-classed for the operation + """ + from pyspark.pandas.window import Rolling + + return Rolling(self, window=window, min_periods=min_periods) + + # TODO: 'center' and 'axis' parameter should be implemented. + # 'axis' implementation, refer https://github.com/databricks/koalas/pull/607 + def expanding(self: FrameLike, min_periods: int = 1) -> "Expanding[FrameLike]": + """ + Provide expanding transformations. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + Parameters + ---------- + min_periods: int, default 1 + Minimum number of observations in window required to have a value + (otherwise result is NA). + + Returns + ------- + a Window sub-classed for the operation + """ + from pyspark.pandas.window import Expanding + + return Expanding(self, min_periods=min_periods) + + # TODO: 'adjust', 'axis', 'method' parameter should be implemented. + def ewm( + self: FrameLike, + com: Optional[float] = None, + span: Optional[float] = None, + halflife: Optional[float] = None, + alpha: Optional[float] = None, + min_periods: Optional[int] = None, + ignore_na: bool_type = False, + ) -> "ExponentialMoving[FrameLike]": + """ + Provide exponentially weighted window transformations. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + com: float, optional + Specify decay in terms of center of mass. + alpha = 1 / (1 + com), for com >= 0. + + span: float, optional + Specify decay in terms of span. + alpha = 2 / (span + 1), for span >= 1. + + halflife: float, optional + Specify decay in terms of half-life. + alpha = 1 - exp(-ln(2) / halflife), for halflife > 0. + + alpha: float, optional + Specify smoothing factor alpha directly. + 0 < alpha <= 1. + + min_periods: int, default None + Minimum number of observations in window required to have a value + (otherwise result is NA). + + ignore_na: bool, default False + Ignore missing values when calculating weights. + + - When ``ignore_na=False`` (default), weights are based on absolute positions. + For example, the weights of :math:`x_0` and :math:`x_2` used in calculating + the final weighted average of [:math:`x_0`, None, :math:`x_2`] are + :math:`(1-\alpha)^2` and :math:`1` if ``adjust=True``, and + :math:`(1-\alpha)^2` and :math:`\alpha` if ``adjust=False``. + + - When ``ignore_na=True``, weights are based + on relative positions. For example, the weights of :math:`x_0` and :math:`x_2` + used in calculating the final weighted average of + [:math:`x_0`, None, :math:`x_2`] are :math:`1-\alpha` and :math:`1` if + ``adjust=True``, and :math:`1-\alpha` and :math:`\alpha` if ``adjust=False``. + + Returns + ------- + a Window sub-classed for the operation + """ + from pyspark.pandas.window import ExponentialMoving + + return ExponentialMoving( + self, + com=com, + span=span, + halflife=halflife, + alpha=alpha, + min_periods=min_periods, + ignore_na=ignore_na, + ) + + def get(self, key: Any, default: Optional[Any] = None) -> Any: + """ + Get item from object for given key (DataFrame column, Panel slice, + etc.). Returns default value if not found. + + Parameters + ---------- + key: object + + Returns + ------- + value: same type as items contained in object + + Examples + -------- + >>> df = ps.DataFrame({'x':range(3), 'y':['a','b','b'], 'z':['a','b','b']}, + ... columns=['x', 'y', 'z'], index=[10, 20, 20]) + >>> df + x y z + 10 0 a a + 20 1 b b + 20 2 b b + + >>> df.get('x') + 10 0 + 20 1 + 20 2 + Name: x, dtype: int64 + + >>> df.get(['x', 'y']) + x y + 10 0 a + 20 1 b + 20 2 b + + >>> df.x.get(10) + 0 + + >>> df.x.get(20) + 20 1 + 20 2 + Name: x, dtype: int64 + + >>> df.x.get(15, -1) + -1 + """ + try: + return self[key] + except (KeyError, ValueError, IndexError): + return default + + def squeeze(self, axis: Optional[Axis] = None) -> Union[Scalar, "DataFrame", "Series"]: + """ + Squeeze 1 dimensional axis objects into scalars. + + Series or DataFrames with a single element are squeezed to a scalar. + DataFrames with a single column or a single row are squeezed to a + Series. Otherwise the object is unchanged. + + This method is most useful when you don't know if your + object is a Series or DataFrame, but you do know it has just a single + column. In that case you can safely call `squeeze` to ensure you have a + Series. + + Parameters + ---------- + axis: {0 or 'index', 1 or 'columns', None}, default None + A specific axis to squeeze. By default, all length-1 axes are + squeezed. + + Returns + ------- + DataFrame, Series, or scalar + The projection after squeezing `axis` or all the axes. + + See Also + -------- + Series.iloc: Integer-location based indexing for selecting scalars. + DataFrame.iloc: Integer-location based indexing for selecting Series. + Series.to_frame: Inverse of DataFrame.squeeze for a + single-column DataFrame. + + Examples + -------- + >>> primes = ps.Series([2, 3, 5, 7]) + + Slicing might produce a Series with a single value: + + >>> even_primes = primes[primes % 2 == 0] + >>> even_primes + 0 2 + dtype: int64 + + >>> even_primes.squeeze() + 2 + + Squeezing objects with more than one value in every axis does nothing: + + >>> odd_primes = primes[primes % 2 == 1] + >>> odd_primes + 1 3 + 2 5 + 3 7 + dtype: int64 + + >>> odd_primes.squeeze() + 1 3 + 2 5 + 3 7 + dtype: int64 + + Squeezing is even more effective when used with DataFrames. + + >>> df = ps.DataFrame([[1, 2], [3, 4]], columns=['a', 'b']) + >>> df + a b + 0 1 2 + 1 3 4 + + Slicing a single column will produce a DataFrame with the columns + having only one value: + + >>> df_a = df[['a']] + >>> df_a + a + 0 1 + 1 3 + + The columns can be squeezed down, resulting in a Series: + + >>> df_a.squeeze('columns') + 0 1 + 1 3 + Name: a, dtype: int64 + + Slicing a single row from a single column will produce a single + scalar DataFrame: + + >>> df_1a = df.loc[[1], ['a']] + >>> df_1a + a + 1 3 + + Squeezing the rows produces a single scalar Series: + + >>> df_1a.squeeze('rows') + a 3 + Name: 1, dtype: int64 + + Squeezing all axes will project directly into a scalar: + + >>> df_1a.squeeze() + 3 + """ + if axis is not None: + axis = "index" if axis == "rows" else axis + axis = validate_axis(axis) + + if isinstance(self, ps.DataFrame): + from pyspark.pandas.series import first_series + + is_squeezable = len(self.columns[:2]) == 1 + # If DataFrame has multiple columns, there is no change. + if not is_squeezable: + return self + series_from_column = first_series(self) + has_single_value = len(series_from_column.head(2)) == 1 + # If DataFrame has only a single value, use pandas API directly. + if has_single_value: + result = self._to_internal_pandas().squeeze(axis) + return ps.Series(result) if isinstance(result, pd.Series) else result + elif axis == 0: + return self + else: + return series_from_column + else: + # The case of Series is simple. + # If Series has only a single value, just return it as a scalar. + # Otherwise, there is no change. + self_top_two = cast("Series", self).head(2) + has_single_value = len(self_top_two) == 1 + return cast(Union[Scalar, ps.Series], self_top_two[0] if has_single_value else self) + + def truncate( + self, + before: Optional[Any] = None, + after: Optional[Any] = None, + axis: Optional[Axis] = None, + copy: bool_type = True, + ) -> DataFrameOrSeries: + """ + Truncate a Series or DataFrame before and after some index value. + + This is a useful shorthand for boolean indexing based on index + values above or below certain thresholds. + + .. note:: This API is dependent on :meth:`Index.is_monotonic_increasing` + which can be expensive. + + Parameters + ---------- + before: date, str, int + Truncate all rows before this index value. + after: date, str, int + Truncate all rows after this index value. + axis: {0 or 'index', 1 or 'columns'}, optional + Axis to truncate. Truncates the index (rows) by default. + copy: bool, default is True, + Return a copy of the truncated section. + + Returns + ------- + type of caller + The truncated Series or DataFrame. + + See Also + -------- + DataFrame.loc: Select a subset of a DataFrame by label. + DataFrame.iloc: Select a subset of a DataFrame by position. + + Examples + -------- + >>> df = ps.DataFrame({'A': ['a', 'b', 'c', 'd', 'e'], + ... 'B': ['f', 'g', 'h', 'i', 'j'], + ... 'C': ['k', 'l', 'm', 'n', 'o']}, + ... index=[1, 2, 3, 4, 5]) + >>> df + A B C + 1 a f k + 2 b g l + 3 c h m + 4 d i n + 5 e j o + + >>> df.truncate(before=2, after=4) + A B C + 2 b g l + 3 c h m + 4 d i n + + The columns of a DataFrame can be truncated. + + >>> df.truncate(before="A", after="B", axis="columns") + A B + 1 a f + 2 b g + 3 c h + 4 d i + 5 e j + + For Series, only rows can be truncated. + + >>> df['A'].truncate(before=2, after=4) + 2 b + 3 c + 4 d + Name: A, dtype: object + + A Series has index that sorted integers. + + >>> s = ps.Series([10, 20, 30, 40, 50, 60, 70], + ... index=[1, 2, 3, 4, 5, 6, 7]) + >>> s + 1 10 + 2 20 + 3 30 + 4 40 + 5 50 + 6 60 + 7 70 + dtype: int64 + + >>> s.truncate(2, 5) + 2 20 + 3 30 + 4 40 + 5 50 + dtype: int64 + + A Series has index that sorted strings. + + >>> s = ps.Series([10, 20, 30, 40, 50, 60, 70], + ... index=['a', 'b', 'c', 'd', 'e', 'f', 'g']) + >>> s + a 10 + b 20 + c 30 + d 40 + e 50 + f 60 + g 70 + dtype: int64 + + >>> s.truncate('b', 'e') + b 20 + c 30 + d 40 + e 50 + dtype: int64 + """ + from pyspark.pandas.series import first_series + + axis = validate_axis(axis) + indexes = self.index + indexes_increasing = indexes.is_monotonic_increasing + if not indexes_increasing and not indexes.is_monotonic_decreasing: + raise ValueError("truncate requires a sorted index") + if (before is None) and (after is None): + return cast(Union[ps.DataFrame, ps.Series], self.copy() if copy else self) + if (before is not None and after is not None) and before > after: + raise ValueError("Truncate: %s must be after %s" % (after, before)) + + if isinstance(self, ps.Series): + if indexes_increasing: + result = first_series( + self.to_frame().loc[before:after] # type: ignore[arg-type] + ).rename(self.name) + else: + result = first_series( + self.to_frame().loc[after:before] # type: ignore[arg-type] + ).rename(self.name) + elif isinstance(self, ps.DataFrame): + if axis == 0: + if indexes_increasing: + result = self.loc[before:after] # type: ignore[assignment] + else: + result = self.loc[after:before] # type: ignore[assignment] + elif axis == 1: + result = self.loc[:, before:after] # type: ignore[assignment] + + return cast(DataFrameOrSeries, result.copy() if copy else result) + + def to_markdown( + self, buf: Optional[Union[IO[str], str]] = None, mode: Optional[str] = None + ) -> str: + """ + Print Series or DataFrame in Markdown-friendly format. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + buf: writable buffer, defaults to sys.stdout + Where to send the output. By default, the output is printed to + sys.stdout. Pass a writable buffer if you need to further process + the output. + mode: str, optional + Mode in which file is opened. + **kwargs + These parameters will be passed to `tabulate`. + + Returns + ------- + str + Series or DataFrame in Markdown-friendly format. + + Notes + ----- + Requires the `tabulate `_ package. + + Examples + -------- + >>> psser = ps.Series(["elk", "pig", "dog", "quetzal"], name="animal") + >>> print(psser.to_markdown()) # doctest: +SKIP + | | animal | + |---:|:---------| + | 0 | elk | + | 1 | pig | + | 2 | dog | + | 3 | quetzal | + + >>> psdf = ps.DataFrame( + ... data={"animal_1": ["elk", "pig"], "animal_2": ["dog", "quetzal"]} + ... ) + >>> print(psdf.to_markdown()) # doctest: +SKIP + | | animal_1 | animal_2 | + |---:|:-----------|:-----------| + | 0 | elk | dog | + | 1 | pig | quetzal | + """ + log_advice( + "`to_markdown` loads all data into the driver's memory. " + "It should only be used if the resulting pandas object is expected to be small." + ) + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + psser_or_psdf = self + internal_pandas = psser_or_psdf._to_internal_pandas() + return validate_arguments_and_invoke_function( + internal_pandas, self.to_markdown, type(internal_pandas).to_markdown, args + ) + + @abstractmethod + def fillna( + self: FrameLike, + value: Optional[Any] = None, + method: Optional[str] = None, + axis: Optional[Axis] = None, + inplace: bool_type = False, + limit: Optional[int] = None, + ) -> FrameLike: + pass + + # TODO: add 'downcast' when value parameter exists + def bfill( + self: FrameLike, + axis: Optional[Axis] = None, + inplace: bool_type = False, + limit: Optional[int] = None, + ) -> FrameLike: + """ + Synonym for `DataFrame.fillna()` or `Series.fillna()` with ``method=`bfill```. + + .. note:: the current implementation of 'bfill' uses Spark's Window + without specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + axis: {0 or `index`} + 1 and `columns` are not supported. + inplace: boolean, default False + Fill in place (do not create a new object) + limit: int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame or Series + DataFrame or Series with NA entries filled. + + Examples + -------- + >>> psdf = ps.DataFrame({ + ... 'A': [None, 3, None, None], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> psdf + A B C D + 0 NaN 2.0 NaN 0 + 1 3.0 4.0 NaN 1 + 2 NaN NaN NaN 5 + 3 NaN 3.0 1.0 4 + + Propagate non-null values backward. + + >>> psdf.bfill() + A B C D + 0 3.0 2.0 1.0 0 + 1 3.0 4.0 1.0 1 + 2 NaN 3.0 1.0 5 + 3 NaN 3.0 1.0 4 + + For Series + + >>> psser = ps.Series([None, None, None, 1]) + >>> psser + 0 NaN + 1 NaN + 2 NaN + 3 1.0 + dtype: float64 + + >>> psser.bfill() + 0 1.0 + 1 1.0 + 2 1.0 + 3 1.0 + dtype: float64 + """ + return self.fillna(method="bfill", axis=axis, inplace=inplace, limit=limit) + + backfill = bfill + + # TODO: add 'downcast' when value parameter exists + def ffill( + self: FrameLike, + axis: Optional[Axis] = None, + inplace: bool_type = False, + limit: Optional[int] = None, + ) -> FrameLike: + """ + Synonym for `DataFrame.fillna()` or `Series.fillna()` with ``method=`ffill```. + + .. note:: the current implementation of 'ffill' uses Spark's Window + without specifying partition specification. This leads to moveing all data into a + single a partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + axis: {0 or `index`} + 1 and `columns` are not supported. + inplace: boolean, default False + Fill in place (do not create a new object) + limit: int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame or Series + DataFrame or Series with NA entries filled. + + Examples + -------- + >>> psdf = ps.DataFrame({ + ... 'A': [None, 3, None, None], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> psdf + A B C D + 0 NaN 2.0 NaN 0 + 1 3.0 4.0 NaN 1 + 2 NaN NaN NaN 5 + 3 NaN 3.0 1.0 4 + + Propagate non-null values forward. + + >>> psdf.ffill() + A B C D + 0 NaN 2.0 NaN 0 + 1 3.0 4.0 NaN 1 + 2 3.0 4.0 NaN 5 + 3 3.0 3.0 1.0 4 + + For Series + + >>> psser = ps.Series([2, 4, None, 3]) + >>> psser + 0 2.0 + 1 4.0 + 2 NaN + 3 3.0 + dtype: float64 + + >>> psser.ffill() + 0 2.0 + 1 4.0 + 2 4.0 + 3 3.0 + dtype: float64 + """ + return self.fillna(method="ffill", axis=axis, inplace=inplace, limit=limit) + + pad = ffill + + # TODO: add 'axis', 'inplace', 'downcast' + def interpolate( + self: FrameLike, + method: str = "linear", + limit: Optional[int] = None, + limit_direction: Optional[str] = None, + limit_area: Optional[str] = None, + ) -> FrameLike: + """ + Fill NaN values using an interpolation method. + + .. note:: the current implementation of interpolate uses Spark's Window without + specifying partition specification. This leads to moveing all data into a + single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + method: str, default 'linear' + Interpolation technique to use. One of: + + * 'linear': Ignore the index and treat the values as equally + spaced. + + limit: int, optional + Maximum number of consecutive NaNs to fill. Must be greater than + 0. + + limit_direction: str, default None + Consecutive NaNs will be filled in this direction. + One of {{'forward', 'backward', 'both'}}. + + limit_area: str, default None + If limit is specified, consecutive NaNs will be filled with this restriction. One of: + + * None: No fill restriction. + * 'inside': Only fill NaNs surrounded by valid values (interpolate). + * 'outside': Only fill NaNs outside valid values (extrapolate). + + Returns + ------- + Series or DataFrame or None + Returns the same object type as the caller, interpolated at + some or all NA values. + + See Also + -------- + fillna: Fill missing values using different methods. + + Examples + -------- + Filling in NA via linear interpolation. + + >>> s = ps.Series([0, 1, np.nan, 3]) + >>> s + 0 0.0 + 1 1.0 + 2 NaN + 3 3.0 + dtype: float64 + >>> s.interpolate() + 0 0.0 + 1 1.0 + 2 2.0 + 3 3.0 + dtype: float64 + + Fill the DataFrame forward (that is, going down) along each column + using linear interpolation. + + Note how the last entry in column 'a' is interpolated differently, + because there is no entry after it to use for interpolation. + Note how the first entry in column 'b' remains NA, because there + is no entry before it to use for interpolation. + + >>> df = ps.DataFrame([(0.0, np.nan, -1.0, 1.0), + ... (np.nan, 2.0, np.nan, np.nan), + ... (2.0, 3.0, np.nan, 9.0), + ... (np.nan, 4.0, -4.0, 16.0)], + ... columns=list('abcd')) + >>> df + a b c d + 0 0.0 NaN -1.0 1.0 + 1 NaN 2.0 NaN NaN + 2 2.0 3.0 NaN 9.0 + 3 NaN 4.0 -4.0 16.0 + >>> df.interpolate(method='linear') + a b c d + 0 0.0 NaN -1.0 1.0 + 1 1.0 2.0 -2.0 5.0 + 2 2.0 3.0 -3.0 9.0 + 3 2.0 4.0 -4.0 16.0 + """ + return self.interpolate( + method=method, limit=limit, limit_direction=limit_direction, limit_area=limit_area + ) + + @property + def at(self) -> AtIndexer: + return AtIndexer(self) + + at.__doc__ = AtIndexer.__doc__ + + @property + def iat(self) -> iAtIndexer: + return iAtIndexer(self) + + iat.__doc__ = iAtIndexer.__doc__ + + @property + def iloc(self) -> iLocIndexer: + return iLocIndexer(self) + + iloc.__doc__ = iLocIndexer.__doc__ + + @property + def loc(self) -> LocIndexer: + return LocIndexer(self) + + loc.__doc__ = LocIndexer.__doc__ + + def __bool__(self) -> NoReturn: + raise ValueError( + "The truth value of a {0} is ambiguous. " + "Use a.empty, a.bool(), a.item(), a.any() or a.all().".format(self.__class__.__name__) + ) + + @staticmethod + def _count_expr(psser: "Series") -> Column: + return F.count(psser._dtype_op.nan_to_null(psser).spark.column) + + +def _test() -> None: + import os + import doctest + import shutil + import sys + import tempfile + from pyspark.sql import SparkSession + import pyspark.pandas.generic + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.generic.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.generic tests") + .getOrCreate() + ) + + path = tempfile.mkdtemp() + globs["path"] = path + + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.generic, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + + shutil.rmtree(path, ignore_errors=True) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/groupby.py new file mode 100644 index 0000000..663a635 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/groupby.py @@ -0,0 +1,4448 @@ +# +# 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. +# + +""" +A wrapper for GroupedData to behave like pandas GroupBy. +""" +from abc import ABCMeta, abstractmethod +import inspect +from collections import defaultdict, namedtuple +from distutils.version import LooseVersion +from functools import partial +from itertools import product +from typing import ( + Any, + Callable, + Dict, + Generic, + Iterator, + Mapping, + List, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, + cast, + TYPE_CHECKING, +) +import warnings + +import pandas as pd +from pandas.api.types import is_number, is_hashable, is_list_like # type: ignore[attr-defined] + +if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + from pandas.core.common import _builtin_table # type: ignore[attr-defined] +else: + from pandas.core.base import SelectionMixin + + _builtin_table = SelectionMixin._builtin_table # type: ignore[attr-defined] + +from pyspark.sql import Column, DataFrame as SparkDataFrame, Window, functions as F +from pyspark.sql.types import ( + BooleanType, + DataType, + DoubleType, + NumericType, + StructField, + StructType, + StringType, +) + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import Axis, FrameLike, Label, Name +from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + HIDDEN_COLUMNS, + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, + SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_PATTERN, +) +from pyspark.pandas.missing.groupby import ( + MissingPandasLikeDataFrameGroupBy, + MissingPandasLikeSeriesGroupBy, +) +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.config import get_option +from pyspark.pandas.utils import ( + align_diff_frames, + is_name_like_tuple, + is_name_like_value, + name_like_string, + same_anchor, + scol_for, + verify_temp_column_name, + log_advice, +) +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale +from pyspark.pandas.exceptions import DataError + +if TYPE_CHECKING: + from pyspark.pandas.window import RollingGroupby, ExpandingGroupby, ExponentialMovingGroupby + + +# to keep it the same as pandas +NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"]) + + +class GroupBy(Generic[FrameLike], metaclass=ABCMeta): + """ + :ivar _psdf: The parent dataframe that is used to perform the groupby + :type _psdf: DataFrame + :ivar _groupkeys: The list of keys that will be used to perform the grouping + :type _groupkeys: List[Series] + """ + + def __init__( + self, + psdf: DataFrame, + groupkeys: List[Series], + as_index: bool, + dropna: bool, + column_labels_to_exclude: Set[Label], + agg_columns_selected: bool, + agg_columns: List[Series], + ): + self._psdf = psdf + self._groupkeys = groupkeys + self._as_index = as_index + self._dropna = dropna + self._column_labels_to_exclude = column_labels_to_exclude + self._agg_columns_selected = agg_columns_selected + self._agg_columns = agg_columns + + @property + def _groupkeys_scols(self) -> List[Column]: + return [s.spark.column for s in self._groupkeys] + + @property + def _agg_columns_scols(self) -> List[Column]: + return [s.spark.column for s in self._agg_columns] + + @abstractmethod + def _apply_series_op( + self, + op: Callable[["SeriesGroupBy"], Series], + should_resolve: bool = False, + numeric_only: bool = False, + ) -> FrameLike: + pass + + @abstractmethod + def _handle_output(self, psdf: DataFrame) -> FrameLike: + pass + + # TODO: Series support is not implemented yet. + # TODO: not all arguments are implemented comparing to pandas' for now. + def aggregate( + self, + func_or_funcs: Optional[Union[str, List[str], Dict[Name, Union[str, List[str]]]]] = None, + *args: Any, + **kwargs: Any, + ) -> DataFrame: + """Aggregate using one or more operations over the specified axis. + + Parameters + ---------- + func_or_funcs : dict, str or list + a dict mapping from column name (string) to + aggregate functions (string or list of strings). + + Returns + ------- + Series or DataFrame + + The return can be: + + * Series : when DataFrame.agg is called with a single function + * DataFrame : when DataFrame.agg is called with several functions + + Return Series or DataFrame. + + Notes + ----- + `agg` is an alias for `aggregate`. Use the alias. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 2], + ... 'B': [1, 2, 3, 4], + ... 'C': [0.362, 0.227, 1.267, -0.562]}, + ... columns=['A', 'B', 'C']) + + >>> df + A B C + 0 1 1 0.362 + 1 1 2 0.227 + 2 2 3 1.267 + 3 2 4 -0.562 + + Different aggregations per column + + >>> aggregated = df.groupby('A').agg({'B': 'min', 'C': 'sum'}) + >>> aggregated[['B', 'C']].sort_index() # doctest: +NORMALIZE_WHITESPACE + B C + A + 1 1 0.589 + 2 3 0.705 + + >>> aggregated = df.groupby('A').agg({'B': ['min', 'max']}) + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + B + min max + A + 1 1 2 + 2 3 4 + + >>> aggregated = df.groupby('A').agg('min') + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + B C + A + 1 1 0.227 + 2 3 -0.562 + + >>> aggregated = df.groupby('A').agg(['min', 'max']) + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + B C + min max min max + A + 1 1 2 0.227 0.362 + 2 3 4 -0.562 1.267 + + To control the output names with different aggregations per column, pandas-on-Spark + also supports 'named aggregation' or nested renaming in .agg. It can also be + used when applying multiple aggregation functions to specific columns. + + >>> aggregated = df.groupby('A').agg(b_max=ps.NamedAgg(column='B', aggfunc='max')) + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + b_max + A + 1 2 + 2 4 + + >>> aggregated = df.groupby('A').agg(b_max=('B', 'max'), b_min=('B', 'min')) + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + b_max b_min + A + 1 2 1 + 2 4 3 + + >>> aggregated = df.groupby('A').agg(b_max=('B', 'max'), c_min=('C', 'min')) + >>> aggregated.sort_index() # doctest: +NORMALIZE_WHITESPACE + b_max c_min + A + 1 2 0.227 + 2 4 -0.562 + """ + # I think current implementation of func and arguments in pandas-on-Spark for aggregate + # is different than pandas, later once arguments are added, this could be removed. + if func_or_funcs is None and kwargs is None: + raise ValueError("No aggregation argument or function specified.") + + relabeling = func_or_funcs is None and is_multi_agg_with_relabel(**kwargs) + if relabeling: + ( + func_or_funcs, + columns, + order, + ) = normalize_keyword_aggregation( # type: ignore[assignment] + kwargs + ) + + if not isinstance(func_or_funcs, (str, list)): + if not isinstance(func_or_funcs, dict) or not all( + is_name_like_value(key) + and ( + isinstance(value, str) + or isinstance(value, list) + and all(isinstance(v, str) for v in value) + ) + for key, value in func_or_funcs.items() + ): + raise ValueError( + "aggs must be a dict mapping from column name " + "to aggregate functions (string or list of strings)." + ) + + else: + agg_cols = [col.name for col in self._agg_columns] + func_or_funcs = {col: func_or_funcs for col in agg_cols} + + psdf: DataFrame = DataFrame( + GroupBy._spark_groupby(self._psdf, func_or_funcs, self._groupkeys) + ) + + if self._dropna: + psdf = DataFrame( + psdf._internal.with_new_sdf( + psdf._internal.spark_frame.dropna( + subset=psdf._internal.index_spark_column_names + ) + ) + ) + + if not self._as_index: + should_drop_index = set( + i for i, gkey in enumerate(self._groupkeys) if gkey._psdf is not self._psdf + ) + if len(should_drop_index) > 0: + psdf = psdf.reset_index(level=should_drop_index, drop=True) + if len(should_drop_index) < len(self._groupkeys): + psdf = psdf.reset_index() + + if relabeling: + psdf = psdf[order] + psdf.columns = columns # type: ignore[assignment] + return psdf + + agg = aggregate + + @staticmethod + def _spark_groupby( + psdf: DataFrame, + func: Mapping[Name, Union[str, List[str]]], + groupkeys: Sequence[Series] = (), + ) -> InternalFrame: + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(groupkeys))] + groupkey_scols = [s.spark.column.alias(name) for s, name in zip(groupkeys, groupkey_names)] + + multi_aggs = any(isinstance(v, list) for v in func.values()) + reordered = [] + data_columns = [] + column_labels = [] + for key, value in func.items(): + label = key if is_name_like_tuple(key) else (key,) + if len(label) != psdf._internal.column_labels_level: + raise TypeError("The length of the key must be the same as the column label level.") + for aggfunc in [value] if isinstance(value, str) else value: + column_label = tuple(list(label) + [aggfunc]) if multi_aggs else label + column_labels.append(column_label) + + data_col = name_like_string(column_label) + data_columns.append(data_col) + + col_name = psdf._internal.spark_column_name_for(label) + if aggfunc == "nunique": + reordered.append( + F.expr("count(DISTINCT `{0}`) as `{1}`".format(col_name, data_col)) + ) + + # Implement "quartiles" aggregate function for ``describe``. + elif aggfunc == "quartiles": + reordered.append( + F.expr( + "percentile_approx(`{0}`, array(0.25, 0.5, 0.75)) as `{1}`".format( + col_name, data_col + ) + ) + ) + + else: + reordered.append( + F.expr("{1}(`{0}`) as `{2}`".format(col_name, aggfunc, data_col)) + ) + + sdf = psdf._internal.spark_frame.select(groupkey_scols + psdf._internal.data_spark_columns) + sdf = sdf.groupby(*groupkey_names).agg(*reordered) + + return InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(groupkeys, groupkey_names) + ], + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + ) + + def count(self) -> FrameLike: + """ + Compute count of group, excluding missing values. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 1, 2], + ... 'B': [np.nan, 2, 3, 4, 5], + ... 'C': [1, 2, 1, 1, 2]}, columns=['A', 'B', 'C']) + >>> df.groupby('A').count().sort_index() # doctest: +NORMALIZE_WHITESPACE + B C + A + 1 2 3 + 2 2 2 + """ + return self._reduce_for_stat_function(F.count) + + def first(self, numeric_only: Optional[bool] = False, min_count: int = -1) -> FrameLike: + """ + Compute first of group values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + min_count : int, default -1 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 3, 4, 4], "D": ["a", "b", "a", "a"]}) + >>> df + A B C D + 0 1 True 3 a + 1 2 False 3 b + 2 1 False 4 a + 3 2 True 4 a + + >>> df.groupby("A").first().sort_index() + B C D + A + 1 True 3 a + 2 False 3 b + + Include only float, int, boolean columns when set numeric_only True. + + >>> df.groupby("A").first(numeric_only=True).sort_index() + B C + A + 1 True 3 + 2 False 3 + + >>> df.groupby("D").first().sort_index() + A B C + D + a 1 True 3 + b 2 False 3 + + >>> df.groupby("D").first(min_count=3).sort_index() + A B C + D + a 1.0 True 3.0 + b NaN None NaN + """ + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + return self._reduce_for_stat_function( + lambda col: F.first(col, ignorenulls=True), + accepted_spark_types=(NumericType, BooleanType) if numeric_only else None, + min_count=min_count, + ) + + def last(self, numeric_only: Optional[bool] = False, min_count: int = -1) -> FrameLike: + """ + Compute last of group values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + min_count : int, default -1 + The required number of valid values to perform the operation. If fewer + than ``min_count`` non-NA values are present the result will be NA. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 3, 4, 4], "D": ["a", "a", "b", "a"]}) + >>> df + A B C D + 0 1 True 3 a + 1 2 False 3 a + 2 1 False 4 b + 3 2 True 4 a + + >>> df.groupby("A").last().sort_index() + B C D + A + 1 False 4 b + 2 True 4 a + + Include only float, int, boolean columns when set numeric_only True. + + >>> df.groupby("A").last(numeric_only=True).sort_index() + B C + A + 1 False 4 + 2 True 4 + + >>> df.groupby("D").last().sort_index() + A B C + D + a 2 True 4 + b 1 False 4 + + >>> df.groupby("D").last(min_count=3).sort_index() + A B C + D + a 2.0 True 4.0 + b NaN None NaN + """ + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + return self._reduce_for_stat_function( + lambda col: F.last(col, ignorenulls=True), + accepted_spark_types=(NumericType, BooleanType) if numeric_only else None, + min_count=min_count, + ) + + def max(self, numeric_only: Optional[bool] = False, min_count: int = -1) -> FrameLike: + """ + Compute max of group values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + min_count : bool, default -1 + The required number of valid values to perform the operation. If fewer + than min_count non-NA values are present the result will be NA. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "a", "b", "a"]}) + + >>> df.groupby("A").max().sort_index() + B C D + A + 1 True 3 b + 2 True 4 a + + Include only float, int, boolean columns when set numeric_only True. + + >>> df.groupby("A").max(numeric_only=True).sort_index() + B C + A + 1 True 3 + 2 True 4 + + >>> df.groupby("D").max().sort_index() + A B C + D + a 2 True 4 + b 1 False 3 + + >>> df.groupby("D").max(min_count=3).sort_index() + A B C + D + a 2.0 True 4.0 + b NaN None NaN + """ + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + return self._reduce_for_stat_function( + F.max, + accepted_spark_types=(NumericType, BooleanType) if numeric_only else None, + min_count=min_count, + ) + + def mean(self, numeric_only: Optional[bool] = True) -> FrameLike: + """ + Compute mean of groups, excluding missing values. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + + Returns + ------- + pyspark.pandas.Series or pyspark.pandas.DataFrame + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 1, 2], + ... 'B': [np.nan, 2, 3, 4, 5], + ... 'C': [1, 2, 1, 1, 2], + ... 'D': [True, False, True, False, True]}) + + Groupby one column and return the mean of the remaining columns in + each group. + + >>> df.groupby('A').mean().sort_index() # doctest: +NORMALIZE_WHITESPACE + B C D + A + 1 3.0 1.333333 0.333333 + 2 4.0 1.500000 1.000000 + """ + self._validate_agg_columns(numeric_only=numeric_only, function_name="median") + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + + return self._reduce_for_stat_function( + F.mean, accepted_spark_types=(NumericType,), bool_to_numeric=True + ) + + # TODO: 'q' accepts list like type + def quantile(self, q: float = 0.5, accuracy: int = 10000) -> FrameLike: + """ + Return group values at the given quantile. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + q : float, default 0.5 (50% quantile) + Value between 0 and 1 providing the quantile to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + This is a panda-on-Spark specific parameter. + + Returns + ------- + pyspark.pandas.Series or pyspark.pandas.DataFrame + Return type determined by caller of GroupBy object. + + Notes + ----- + `quantile` in pandas-on-Spark are using distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas, also + `interpolation` parameter is not supported yet. + + See Also + -------- + pyspark.pandas.Series.quantile + pyspark.pandas.DataFrame.quantile + pyspark.sql.functions.percentile_approx + + Examples + -------- + >>> df = ps.DataFrame([ + ... ['a', 1], ['a', 2], ['a', 3], + ... ['b', 1], ['b', 3], ['b', 5] + ... ], columns=['key', 'val']) + + Groupby one column and return the quantile of the remaining columns in + each group. + + >>> df.groupby('key').quantile() + val + key + a 2.0 + b 3.0 + """ + if is_list_like(q): + raise NotImplementedError("q doesn't support for list like type for now") + if not is_number(q): + raise TypeError("must be real number, not %s" % type(q).__name__) + if not 0 <= q <= 1: + raise ValueError("'q' must be between 0 and 1. Got '%s' instead" % q) + return self._reduce_for_stat_function( + lambda col: F.percentile_approx(col.cast(DoubleType()), q, accuracy), + accepted_spark_types=(NumericType, BooleanType), + bool_to_numeric=True, + ) + + def min(self, numeric_only: Optional[bool] = False, min_count: int = -1) -> FrameLike: + """ + Compute min of group values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + min_count : bool, default -1 + The required number of valid values to perform the operation. If fewer + than min_count non-NA values are present the result will be NA. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "a", "b", "a"]}) + >>> df.groupby("A").min().sort_index() + B C D + A + 1 False 3 a + 2 False 4 a + + Include only float, int, boolean columns when set numeric_only True. + + >>> df.groupby("A").min(numeric_only=True).sort_index() + B C + A + 1 False 3 + 2 False 4 + + >>> df.groupby("D").min().sort_index() + A B C + D + a 1 False 3 + b 1 False 3 + + + >>> df.groupby("D").min(min_count=3).sort_index() + A B C + D + a 1.0 False 3.0 + b NaN None NaN + """ + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + return self._reduce_for_stat_function( + F.min, + accepted_spark_types=(NumericType, BooleanType) if numeric_only else None, + min_count=min_count, + ) + + # TODO: sync the doc. + def std(self, ddof: int = 1) -> FrameLike: + """ + Compute standard deviation of groups, excluding missing values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + ddof : int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + .. versionchanged:: 3.4.0 + Supported including arbitary integers. + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "b", "b", "a"]}) + + >>> df.groupby("A").std() + B C + A + 1 0.707107 0.0 + 2 0.707107 0.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + # Raise the TypeError when all aggregation columns are of unaccepted data types + any_accepted = any( + isinstance(_agg_col.spark.data_type, (NumericType, BooleanType)) + for _agg_col in self._agg_columns + ) + if not any_accepted: + raise TypeError( + "Unaccepted data types of aggregation columns; numeric or bool expected." + ) + + def std(col: Column) -> Column: + return SF.stddev(col, ddof) + + return self._reduce_for_stat_function( + std, + accepted_spark_types=(NumericType,), + bool_to_numeric=True, + ) + + def sum(self, numeric_only: Optional[bool] = True, min_count: int = 0) -> FrameLike: + """ + Compute sum of group values + + .. versionadded:: 3.3.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + It takes no effect since only numeric columns can be support here. + + .. versionadded:: 3.4.0 + min_count : int, default 0 + The required number of valid values to perform the operation. + If fewer than min_count non-NA values are present the result will be NA. + + .. versionadded:: 3.4.0 + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "a", "b", "a"]}) + + >>> df.groupby("A").sum().sort_index() + B C + A + 1 1 6 + 2 1 8 + + >>> df.groupby("D").sum().sort_index() + A B C + D + a 5 2 11 + b 1 0 3 + + >>> df.groupby("D").sum(min_count=3).sort_index() + A B C + D + a 5.0 2.0 11.0 + b NaN NaN NaN + + Notes + ----- + There is a behavior difference between pandas-on-Spark and pandas: + + * when there is a non-numeric aggregation column, it will be ignored + even if `numeric_only` is False. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + if numeric_only is not None and not isinstance(numeric_only, bool): + raise TypeError("numeric_only must be None or bool") + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + if numeric_only is not None and not numeric_only: + unsupported = [ + col.name + for col in self._agg_columns + if not isinstance(col.spark.data_type, (NumericType, BooleanType)) + ] + if len(unsupported) > 0: + log_advice( + "GroupBy.sum() can only support numeric and bool columns even if" + f"numeric_only=False, skip unsupported columns: {unsupported}" + ) + + return self._reduce_for_stat_function( + F.sum, + accepted_spark_types=(NumericType, BooleanType), + bool_to_numeric=True, + min_count=min_count, + ) + + # TODO: sync the doc. + def var(self, ddof: int = 1) -> FrameLike: + """ + Compute variance of groups, excluding missing values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + ddof : int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + .. versionchanged:: 3.4.0 + Supported including arbitary integers. + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 2], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "b", "b", "a"]}) + + >>> df.groupby("A").var() + B C + A + 1 0.5 0.0 + 2 0.5 0.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + def var(col: Column) -> Column: + return SF.var(col, ddof) + + return self._reduce_for_stat_function( + var, + accepted_spark_types=(NumericType,), + bool_to_numeric=True, + ) + + def skew(self) -> FrameLike: + """ + Compute skewness of groups, excluding missing values. + + .. versionadded:: 3.4.0 + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 1], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "b", "b", "a"]}) + + >>> df.groupby("A").skew() + B C + A + 1 -1.732051 1.732051 + 2 NaN NaN + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + return self._reduce_for_stat_function( + SF.skew, + accepted_spark_types=(NumericType,), + bool_to_numeric=True, + ) + + # TODO: 'axis', 'skipna', 'level' parameter should be implemented. + def mad(self) -> FrameLike: + """ + Compute mean absolute deviation of groups, excluding missing values. + + .. versionadded:: 3.4.0 + + .. deprecated:: 3.4.0 + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 1], "B": [True, False, False, True], + ... "C": [3, 4, 3, 4], "D": ["a", "b", "b", "a"]}) + + >>> df.groupby("A").mad() + B C + A + 1 0.444444 0.444444 + 2 0.000000 0.000000 + + >>> df.B.groupby(df.A).mad() + A + 1 0.444444 + 2 0.000000 + Name: B, dtype: float64 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + warnings.warn( + "The 'mad' method is deprecated and will be removed in a future version. " + "To compute the same result, you may do `(group_df - group_df.mean()).abs().mean()`.", + FutureWarning, + ) + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + internal, agg_columns, sdf = self._prepare_reduce( + groupkey_names=groupkey_names, + accepted_spark_types=(NumericType, BooleanType), + bool_to_numeric=False, + ) + psdf: DataFrame = DataFrame(internal) + + if len(psdf._internal.column_labels) > 0: + window = Window.partitionBy(groupkey_names).rowsBetween( + Window.unboundedPreceding, Window.unboundedFollowing + ) + new_agg_scols = {} + new_stat_scols = [] + for agg_column in agg_columns: + # it is not able to directly use 'self._reduce_for_stat_function', due to + # 'it is not allowed to use a window function inside an aggregate function'. + # so we need to create temporary columns to compute the 'abs(x - avg(x))' here. + agg_column_name = agg_column._internal.data_spark_column_names[0] + new_agg_column_name = verify_temp_column_name( + psdf._internal.spark_frame, "__tmp_agg_col_{}__".format(agg_column_name) + ) + casted_agg_scol = F.col(agg_column_name).cast("double") + new_agg_scols[new_agg_column_name] = F.abs( + casted_agg_scol - F.avg(casted_agg_scol).over(window) + ) + new_stat_scols.append(F.avg(F.col(new_agg_column_name)).alias(agg_column_name)) + + sdf = ( + psdf._internal.spark_frame.withColumns(new_agg_scols) + .groupby(groupkey_names) + .agg(*new_stat_scols) + ) + else: + sdf = sdf.select(*groupkey_names).distinct() + + internal = internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + data_spark_columns=[scol_for(sdf, col) for col in internal.data_spark_column_names], + data_fields=None, + ) + + return self._prepare_return(DataFrame(internal)) + + def sem(self, ddof: int = 1) -> FrameLike: + """ + Compute standard error of the mean of groups, excluding missing values. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + ddof : int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + Examples + -------- + >>> df = ps.DataFrame({"A": [1, 2, 1, 1], "B": [True, False, False, True], + ... "C": [3, None, 3, 4], "D": ["a", "b", "b", "a"]}) + + >>> df.groupby("A").sem() + B C + A + 1 0.333333 0.333333 + 2 NaN NaN + + >>> df.groupby("D").sem(ddof=1) + A B C + D + a 0.0 0.0 0.5 + b 0.5 0.0 NaN + + >>> df.B.groupby(df.A).sem() + A + 1 0.333333 + 2 NaN + Name: B, dtype: float64 + + See Also + -------- + pyspark.pandas.Series.sem + pyspark.pandas.DataFrame.sem + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + # Raise the TypeError when all aggregation columns are of unaccepted data types + any_accepted = any( + isinstance(_agg_col.spark.data_type, (NumericType, BooleanType)) + for _agg_col in self._agg_columns + ) + if not any_accepted: + raise TypeError( + "Unaccepted data types of aggregation columns; numeric or bool expected." + ) + + def sem(col: Column) -> Column: + return SF.stddev(col, ddof) / F.sqrt(F.count(col)) + + return self._reduce_for_stat_function( + sem, + accepted_spark_types=(NumericType, BooleanType), + bool_to_numeric=True, + ) + + # TODO: 1, 'n' accepts list and slice; 2, implement 'dropna' parameter + def nth(self, n: int) -> FrameLike: + """ + Take the nth row from each group. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + n : int + A single nth value for the row + + Returns + ------- + Series or DataFrame + + Notes + ----- + 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__`. + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 1, 2], + ... 'B': [np.nan, 2, 3, 4, 5]}, columns=['A', 'B']) + >>> g = df.groupby('A') + >>> g.nth(0) + B + A + 1 NaN + 2 3.0 + >>> g.nth(1) + B + A + 1 2.0 + 2 5.0 + >>> g.nth(-1) + B + A + 1 4.0 + 2 5.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + """ + if isinstance(n, slice) or is_list_like(n): + raise NotImplementedError("n doesn't support slice or list for now") + if not isinstance(n, int): + raise TypeError("Invalid index %s" % type(n).__name__) + + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + internal, agg_columns, sdf = self._prepare_reduce( + groupkey_names=groupkey_names, + accepted_spark_types=None, + bool_to_numeric=False, + ) + psdf: DataFrame = DataFrame(internal) + + if len(psdf._internal.column_labels) > 0: + window1 = Window.partitionBy(*groupkey_names).orderBy(NATURAL_ORDER_COLUMN_NAME) + tmp_row_number_col = verify_temp_column_name(sdf, "__tmp_row_number_col__") + if n >= 0: + sdf = ( + psdf._internal.spark_frame.withColumn( + tmp_row_number_col, F.row_number().over(window1) + ) + .where(F.col(tmp_row_number_col) == n + 1) + .drop(tmp_row_number_col) + ) + else: + window2 = Window.partitionBy(*groupkey_names).rowsBetween( + Window.unboundedPreceding, Window.unboundedFollowing + ) + tmp_group_size_col = verify_temp_column_name(sdf, "__tmp_group_size_col__") + sdf = ( + psdf._internal.spark_frame.withColumn( + tmp_group_size_col, F.count(F.lit(0)).over(window2) + ) + .withColumn(tmp_row_number_col, F.row_number().over(window1)) + .where(F.col(tmp_row_number_col) == F.col(tmp_group_size_col) + 1 + n) + .drop(tmp_group_size_col, tmp_row_number_col) + ) + else: + sdf = sdf.select(*groupkey_names).distinct() + + internal = internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + data_spark_columns=[scol_for(sdf, col) for col in internal.data_spark_column_names], + data_fields=None, + ) + + return self._prepare_return(DataFrame(internal)) + + def prod(self, numeric_only: Optional[bool] = True, min_count: int = 0) -> FrameLike: + """ + Compute prod of groups. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + min_count : int, default 0 + The required number of valid values to perform the operation. + If fewer than min_count non-NA values are present the result will be NA. + + Returns + ------- + Series or DataFrame + Computed prod of values within each group. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> df = ps.DataFrame( + ... { + ... "A": [1, 1, 2, 1, 2], + ... "B": [np.nan, 2, 3, 4, 5], + ... "C": [1, 2, 1, 1, 2], + ... "D": [True, False, True, False, True], + ... } + ... ) + + Groupby one column and return the prod of the remaining columns in + each group. + + >>> df.groupby('A').prod().sort_index() + B C D + A + 1 8.0 2 0 + 2 15.0 2 1 + + >>> df.groupby('A').prod(min_count=3).sort_index() + B C D + A + 1 NaN 2.0 0.0 + 2 NaN NaN NaN + """ + if not isinstance(min_count, int): + raise TypeError("min_count must be integer") + + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + + self._validate_agg_columns(numeric_only=numeric_only, function_name="prod") + + return self._reduce_for_stat_function( + lambda col: SF.product(col, True), + accepted_spark_types=(NumericType, BooleanType), + bool_to_numeric=True, + min_count=min_count, + ) + + def all(self, skipna: bool = True) -> FrameLike: + """ + Returns True if all values in the group are truthful, else False. + + Parameters + ---------- + skipna : bool, default True + Flag to ignore NA(nan/null) values during truth testing. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + ... 'B': [True, True, True, False, False, + ... False, None, True, None, False]}, + ... columns=['A', 'B']) + >>> df + A B + 0 1 True + 1 1 True + 2 2 True + 3 2 False + 4 3 False + 5 3 False + 6 4 None + 7 4 True + 8 5 None + 9 5 False + + >>> df.groupby('A').all().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 1 True + 2 False + 3 False + 4 True + 5 False + + >>> df.groupby('A').all(skipna=False).sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 1 True + 2 False + 3 False + 4 False + 5 False + """ + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + internal, _, sdf = self._prepare_reduce(groupkey_names) + psdf: DataFrame = DataFrame(internal) + + def sfun(scol: Column, scol_type: DataType) -> Column: + if isinstance(scol_type, NumericType) or skipna: + # np.nan takes no effect to the result; None takes no effect if `skipna` + all_col = F.min(F.coalesce(scol.cast("boolean"), F.lit(True))) + else: + # Take None as False when not `skipna` + all_col = F.min(F.when(scol.isNull(), F.lit(False)).otherwise(scol.cast("boolean"))) + return all_col + + if len(psdf._internal.column_labels) > 0: + stat_exprs = [] + for label in psdf._internal.column_labels: + psser = psdf._psser_for(label) + stat_exprs.append( + sfun( + psser._dtype_op.nan_to_null(psser).spark.column, psser.spark.data_type + ).alias(psser._internal.data_spark_column_names[0]) + ) + sdf = sdf.groupby(*groupkey_names).agg(*stat_exprs) + else: + sdf = sdf.select(*groupkey_names).distinct() + + internal = internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + data_spark_columns=[scol_for(sdf, col) for col in internal.data_spark_column_names], + data_fields=None, + ) + + return self._prepare_return(DataFrame(internal)) + + # TODO: skipna should be implemented. + def any(self) -> FrameLike: + """ + Returns True if any value in the group is truthful, else False. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + ... 'B': [True, True, True, False, False, + ... False, None, True, None, False]}, + ... columns=['A', 'B']) + >>> df + A B + 0 1 True + 1 1 True + 2 2 True + 3 2 False + 4 3 False + 5 3 False + 6 4 None + 7 4 True + 8 5 None + 9 5 False + + >>> df.groupby('A').any().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 1 True + 2 True + 3 False + 4 True + 5 False + """ + return self._reduce_for_stat_function( + lambda col: F.max(F.coalesce(col.cast("boolean"), F.lit(False))) + ) + + # TODO: groupby multiply columns should be implemented. + def size(self) -> Series: + """ + Compute group sizes. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 2, 3, 3, 3], + ... 'B': [1, 1, 2, 3, 3, 3]}, + ... columns=['A', 'B']) + >>> df + A B + 0 1 1 + 1 2 1 + 2 2 2 + 3 3 3 + 4 3 3 + 5 3 3 + + >>> df.groupby('A').size().sort_index() + A + 1 1 + 2 2 + 3 3 + dtype: int64 + + >>> df.groupby(['A', 'B']).size().sort_index() + A B + 1 1 1 + 2 1 1 + 2 1 + 3 3 3 + dtype: int64 + + For Series, + + >>> df.B.groupby(df.A).size().sort_index() + A + 1 1 + 2 2 + 3 3 + Name: B, dtype: int64 + + >>> df.groupby(df.A).B.size().sort_index() + A + 1 1 + 2 2 + 3 3 + Name: B, dtype: int64 + """ + groupkeys = self._groupkeys + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(groupkeys))] + groupkey_scols = [s.spark.column.alias(name) for s, name in zip(groupkeys, groupkey_names)] + sdf = self._psdf._internal.spark_frame.select( + groupkey_scols + self._psdf._internal.data_spark_columns + ) + sdf = sdf.groupby(*groupkey_names).count() + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(groupkeys, groupkey_names) + ], + column_labels=[None], + data_spark_columns=[scol_for(sdf, "count")], + ) + return first_series(DataFrame(internal)) + + def diff(self, periods: int = 1) -> FrameLike: + """ + First discrete difference of element. + + Calculates the difference of a DataFrame element compared with another element in the + DataFrame group (default is the element in the same column of the previous row). + + Parameters + ---------- + periods : int, default 1 + Periods to shift for calculating difference, accepts negative values. + + Returns + ------- + diffed : DataFrame or Series + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2, 3, 4, 5, 6], + ... 'b': [1, 1, 2, 3, 5, 8], + ... 'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c']) + >>> df + a b c + 0 1 1 1 + 1 2 1 4 + 2 3 2 9 + 3 4 3 16 + 4 5 5 25 + 5 6 8 36 + + >>> df.groupby(['b']).diff().sort_index() + a c + 0 NaN NaN + 1 1.0 3.0 + 2 NaN NaN + 3 NaN NaN + 4 NaN NaN + 5 NaN NaN + + Difference with previous column in a group. + + >>> df.groupby(['b'])['a'].diff().sort_index() + 0 NaN + 1 1.0 + 2 NaN + 3 NaN + 4 NaN + 5 NaN + Name: a, dtype: float64 + """ + return self._apply_series_op( + lambda sg: sg._psser._diff(periods, part_cols=sg._groupkeys_scols), should_resolve=True + ) + + def cumcount(self, ascending: bool = True) -> Series: + """ + Number each item in each group from 0 to the length of that group - 1. + + Essentially this is equivalent to + + .. code-block:: python + + self.apply(lambda x: pd.Series(np.arange(len(x)), x.index)) + + Parameters + ---------- + ascending : bool, default True + If False, number in reverse, from length of group - 1 to 0. + + Returns + ------- + Series + Sequence number of each element within each group. + + Examples + -------- + + >>> df = ps.DataFrame([['a'], ['a'], ['a'], ['b'], ['b'], ['a']], + ... columns=['A']) + >>> df + A + 0 a + 1 a + 2 a + 3 b + 4 b + 5 a + >>> df.groupby('A').cumcount().sort_index() + 0 0 + 1 1 + 2 2 + 3 0 + 4 1 + 5 3 + dtype: int64 + >>> df.groupby('A').cumcount(ascending=False).sort_index() + 0 3 + 1 2 + 2 1 + 3 1 + 4 0 + 5 0 + dtype: int64 + """ + ret = ( + self._groupkeys[0] + .rename() + .spark.transform(lambda _: F.lit(0)) + ._cum(F.count, True, part_cols=self._groupkeys_scols, ascending=ascending) + - 1 + ) + internal = ret._internal.resolved_copy + return first_series(DataFrame(internal)) + + def cummax(self) -> FrameLike: + """ + Cumulative max for each group. + + Returns + ------- + Series or DataFrame + + See Also + -------- + Series.cummax + DataFrame.cummax + + Examples + -------- + >>> df = ps.DataFrame( + ... [[1, None, 4], [1, 0.1, 3], [1, 20.0, 2], [4, 10.0, 1]], + ... columns=list('ABC')) + >>> df + A B C + 0 1 NaN 4 + 1 1 0.1 3 + 2 1 20.0 2 + 3 4 10.0 1 + + By default, iterates over rows and finds the sum in each column. + + >>> df.groupby("A").cummax().sort_index() + B C + 0 NaN 4 + 1 0.1 4 + 2 20.0 4 + 3 10.0 1 + + It works as below in Series. + + >>> df.C.groupby(df.A).cummax().sort_index() + 0 4 + 1 4 + 2 4 + 3 1 + Name: C, dtype: int64 + """ + return self._apply_series_op( + lambda sg: sg._psser._cum(F.max, True, part_cols=sg._groupkeys_scols), + should_resolve=True, + numeric_only=True, + ) + + def cummin(self) -> FrameLike: + """ + Cumulative min for each group. + + Returns + ------- + Series or DataFrame + + See Also + -------- + Series.cummin + DataFrame.cummin + + Examples + -------- + >>> df = ps.DataFrame( + ... [[1, None, 4], [1, 0.1, 3], [1, 20.0, 2], [4, 10.0, 1]], + ... columns=list('ABC')) + >>> df + A B C + 0 1 NaN 4 + 1 1 0.1 3 + 2 1 20.0 2 + 3 4 10.0 1 + + By default, iterates over rows and finds the sum in each column. + + >>> df.groupby("A").cummin().sort_index() + B C + 0 NaN 4 + 1 0.1 3 + 2 0.1 2 + 3 10.0 1 + + It works as below in Series. + + >>> df.B.groupby(df.A).cummin().sort_index() + 0 NaN + 1 0.1 + 2 0.1 + 3 10.0 + Name: B, dtype: float64 + """ + return self._apply_series_op( + lambda sg: sg._psser._cum(F.min, True, part_cols=sg._groupkeys_scols), + should_resolve=True, + numeric_only=True, + ) + + def cumprod(self) -> FrameLike: + """ + Cumulative product for each group. + + Returns + ------- + Series or DataFrame + + See Also + -------- + Series.cumprod + DataFrame.cumprod + + Examples + -------- + >>> df = ps.DataFrame( + ... [[1, None, 4], [1, 0.1, 3], [1, 20.0, 2], [4, 10.0, 1]], + ... columns=list('ABC')) + >>> df + A B C + 0 1 NaN 4 + 1 1 0.1 3 + 2 1 20.0 2 + 3 4 10.0 1 + + By default, iterates over rows and finds the sum in each column. + + >>> df.groupby("A").cumprod().sort_index() + B C + 0 NaN 4 + 1 0.1 12 + 2 2.0 24 + 3 10.0 1 + + It works as below in Series. + + >>> df.B.groupby(df.A).cumprod().sort_index() + 0 NaN + 1 0.1 + 2 2.0 + 3 10.0 + Name: B, dtype: float64 + """ + return self._apply_series_op( + lambda sg: sg._psser._cumprod(True, part_cols=sg._groupkeys_scols), + should_resolve=True, + numeric_only=True, + ) + + def cumsum(self) -> FrameLike: + """ + Cumulative sum for each group. + + Returns + ------- + Series or DataFrame + + See Also + -------- + Series.cumsum + DataFrame.cumsum + + Examples + -------- + >>> df = ps.DataFrame( + ... [[1, None, 4], [1, 0.1, 3], [1, 20.0, 2], [4, 10.0, 1]], + ... columns=list('ABC')) + >>> df + A B C + 0 1 NaN 4 + 1 1 0.1 3 + 2 1 20.0 2 + 3 4 10.0 1 + + By default, iterates over rows and finds the sum in each column. + + >>> df.groupby("A").cumsum().sort_index() + B C + 0 NaN 4 + 1 0.1 7 + 2 20.1 9 + 3 10.0 1 + + It works as below in Series. + + >>> df.B.groupby(df.A).cumsum().sort_index() + 0 NaN + 1 0.1 + 2 20.1 + 3 10.0 + Name: B, dtype: float64 + """ + return self._apply_series_op( + lambda sg: sg._psser._cumsum(True, part_cols=sg._groupkeys_scols), + should_resolve=True, + numeric_only=True, + ) + + def apply(self, func: Callable, *args: Any, **kwargs: Any) -> Union[DataFrame, Series]: + """ + Apply function `func` group-wise and combine the results together. + + The function passed to `apply` must take a DataFrame as its first + argument and return a DataFrame. `apply` will + then take care of combining the results back together into a single + dataframe. `apply` is therefore a highly flexible + grouping method. + + While `apply` is a very flexible method, its downside is that + using it can be quite a bit slower than using more specific methods + like `agg` or `transform`. pandas-on-Spark offers a wide range of method that will + be much faster than using `apply` for their specific purposes, so try to + use them before reaching for `apply`. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def pandas_div(x) -> ps.DataFrame[int, [float, float]]: + ... return x[['B', 'C']] / x[['B', 'C']] + + If the return type is specified, the output column names become + `c0, c1, c2 ... cn`. These names are positionally mapped to the returned + DataFrame in ``func``. + + To specify the column names, you can assign them in a NumPy compound type style + as below: + + >>> def pandas_div(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]: + ... return x[['B', 'C']] / x[['B', 'C']] + + >>> pdf = pd.DataFrame({'B': [1.], 'C': [3.]}) + >>> def plus_one(x) -> ps.DataFrame[ + ... (pdf.index.name, pdf.index.dtype), zip(pdf.columns, pdf.dtypes)]: + ... return x[['B', 'C']] / x[['B', 'C']] + + .. note:: the dataframe within ``func`` is actually a pandas dataframe. Therefore, + any pandas API within this function is allowed. + + Parameters + ---------- + func : callable + A callable that takes a DataFrame as its first argument, and + returns a dataframe. + *args + Positional arguments to pass to func. + **kwargs + Keyword arguments to pass to func. + + Returns + ------- + applied : DataFrame or Series + + See Also + -------- + aggregate : Apply aggregate function to the GroupBy object. + DataFrame.apply : Apply a function to a DataFrame. + Series.apply : Apply a function to a Series. + + Examples + -------- + >>> df = ps.DataFrame({'A': 'a a b'.split(), + ... 'B': [1, 2, 3], + ... 'C': [4, 6, 5]}, columns=['A', 'B', 'C']) + >>> g = df.groupby('A') + + Notice that ``g`` has two groups, ``a`` and ``b``. + Calling `apply` in various ways, we can get different grouping results: + + Below the functions passed to `apply` takes a DataFrame as + its argument and returns a DataFrame. `apply` combines the result for + each group together into a new DataFrame: + + >>> def plus_min(x): + ... return x + x.min() + >>> g.apply(plus_min).sort_index() # doctest: +SKIP + A B C + 0 aa 2 8 + 1 aa 3 10 + 2 bb 6 10 + + >>> g.apply(sum).sort_index() # doctest: +NORMALIZE_WHITESPACE + A B C + A + a aa 3 10 + b b 3 5 + + >>> g.apply(len).sort_index() # doctest: +NORMALIZE_WHITESPACE + A + a 2 + b 1 + dtype: int64 + + You can specify the type hint and prevent schema inference for better performance. + + >>> def pandas_div(x) -> ps.DataFrame[int, [float, float]]: + ... return x[['B', 'C']] / x[['B', 'C']] + >>> g.apply(pandas_div).sort_index() # doctest: +SKIP + c0 c1 + 0 1.0 1.0 + 1 1.0 1.0 + 2 1.0 1.0 + + >>> def pandas_div(x) -> ps.DataFrame[("index", int), [("f1", float), ("f2", float)]]: + ... return x[['B', 'C']] / x[['B', 'C']] + >>> g.apply(pandas_div).sort_index() # doctest: +SKIP + f1 f2 + index + 0 1.0 1.0 + 1 1.0 1.0 + 2 1.0 1.0 + + In case of Series, it works as below. + + >>> def plus_max(x) -> ps.Series[int]: + ... return x + x.max() + >>> df.B.groupby(df.A).apply(plus_max).sort_index() # doctest: +SKIP + 0 6 + 1 3 + 2 4 + Name: B, dtype: int64 + + >>> def plus_min(x): + ... return x + x.min() + >>> df.B.groupby(df.A).apply(plus_min).sort_index() # doctest: +SKIP + 0 2 + 1 3 + 2 6 + Name: B, dtype: int64 + + You can also return a scalar value as an aggregated value of the group: + + >>> def plus_length(x) -> int: + ... return len(x) + >>> df.B.groupby(df.A).apply(plus_length).sort_index() # doctest: +SKIP + 0 1 + 1 2 + Name: B, dtype: int64 + + The extra arguments to the function can be passed as below. + + >>> def calculation(x, y, z) -> int: + ... return len(x) + y * z + >>> df.B.groupby(df.A).apply(calculation, 5, z=10).sort_index() # doctest: +SKIP + 0 51 + 1 52 + Name: B, dtype: int64 + """ + if not callable(func): + raise TypeError("%s object is not callable" % type(func).__name__) + + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + should_retain_index = should_infer_schema + + is_series_groupby = isinstance(self, SeriesGroupBy) + + psdf = self._psdf + + if self._agg_columns_selected: + agg_columns = self._agg_columns + else: + agg_columns = [ + psdf._psser_for(label) + for label in psdf._internal.column_labels + if label not in self._column_labels_to_exclude + ] + + psdf, groupkey_labels, groupkey_names = GroupBy._prepare_group_map_apply( + psdf, self._groupkeys, agg_columns + ) + + if is_series_groupby: + name = psdf.columns[-1] + pandas_apply = _builtin_table.get(func, func) + else: + f = _builtin_table.get(func, func) + + def pandas_apply(pdf: pd.DataFrame, *a: Any, **k: Any) -> Any: + return f(pdf.drop(groupkey_names, axis=1), *a, **k) + + should_return_series = False + + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + log_advice( + "If the type hints is not specified for `groupby.apply`, " + "it is expensive to infer the data type internally." + ) + limit = get_option("compute.shortcut_limit") + # Ensure sampling rows >= 2 to make sure apply's infer schema is accurate + # See related: https://github.com/pandas-dev/pandas/issues/46893 + sample_limit = limit + 1 if limit else 2 + pdf = psdf.head(sample_limit)._to_internal_pandas() + groupkeys = [ + pdf[groupkey_name].rename(psser.name) + for groupkey_name, psser in zip(groupkey_names, self._groupkeys) + ] + grouped = pdf.groupby(groupkeys) + if is_series_groupby: + pser_or_pdf = grouped[name].apply(pandas_apply, *args, **kwargs) + else: + pser_or_pdf = grouped.apply(pandas_apply, *args, **kwargs) + psser_or_psdf = ps.from_pandas(pser_or_pdf.infer_objects()) + + if len(pdf) <= limit: + if isinstance(psser_or_psdf, ps.Series) and is_series_groupby: + psser_or_psdf = psser_or_psdf.rename(cast(SeriesGroupBy, self)._psser.name) + return cast(Union[Series, DataFrame], psser_or_psdf) + + if len(grouped) <= 1: + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn( + "The amount of data for return type inference might not be large enough. " + "Consider increasing an option `compute.shortcut_limit`." + ) + + if isinstance(psser_or_psdf, Series): + should_return_series = True + psdf_from_pandas = psser_or_psdf._psdf + else: + psdf_from_pandas = cast(DataFrame, psser_or_psdf) + + index_fields = [ + field.normalize_spark_type() for field in psdf_from_pandas._internal.index_fields + ] + data_fields = [ + field.normalize_spark_type() for field in psdf_from_pandas._internal.data_fields + ] + return_schema = StructType([field.struct_field for field in index_fields + data_fields]) + else: + return_type = infer_return_type(func) + if not is_series_groupby and isinstance(return_type, SeriesType): + raise TypeError( + "Series as a return type hint at frame groupby is not supported " + "currently; however got [%s]. Use DataFrame type hint instead." % return_sig + ) + + if isinstance(return_type, DataFrameType): + data_fields = return_type.data_fields + return_schema = return_type.spark_type + index_fields = return_type.index_fields + should_retain_index = len(index_fields) > 0 + psdf_from_pandas = None + else: + should_return_series = True + dtype = cast(Union[SeriesType, ScalarType], return_type).dtype + spark_type = cast(Union[SeriesType, ScalarType], return_type).spark_type + if is_series_groupby: + data_fields = [ + InternalField( + dtype=dtype, struct_field=StructField(name=name, dataType=spark_type) + ) + ] + else: + data_fields = [ + InternalField( + dtype=dtype, + struct_field=StructField( + name=SPARK_DEFAULT_SERIES_NAME, dataType=spark_type + ), + ) + ] + return_schema = StructType([field.struct_field for field in data_fields]) + + def pandas_groupby_apply(pdf: pd.DataFrame) -> pd.DataFrame: + + if is_series_groupby: + pdf_or_ser = pdf.groupby(groupkey_names)[name].apply(pandas_apply, *args, **kwargs) + else: + pdf_or_ser = pdf.groupby(groupkey_names).apply(pandas_apply, *args, **kwargs) + if should_return_series and isinstance(pdf_or_ser, pd.DataFrame): + pdf_or_ser = pdf_or_ser.stack() + + if not isinstance(pdf_or_ser, pd.DataFrame): + return pd.DataFrame(pdf_or_ser) + else: + return pdf_or_ser + + sdf = GroupBy._spark_group_map_apply( + psdf, + pandas_groupby_apply, + [psdf._internal.spark_column_for(label) for label in groupkey_labels], + return_schema, + retain_index=should_retain_index, + ) + + if should_retain_index: + # If schema is inferred, we can restore indexes too. + if psdf_from_pandas is not None: + internal = psdf_from_pandas._internal.with_new_sdf( + spark_frame=sdf, index_fields=index_fields, data_fields=data_fields + ) + else: + index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None + + index_spark_columns = [ + scol_for(sdf, index_field.struct_field.name) for index_field in index_fields + ] + + if not any( + [ + SPARK_INDEX_NAME_PATTERN.match(index_field.struct_field.name) + for index_field in index_fields + ] + ): + index_names = [(index_field.struct_field.name,) for index_field in index_fields] + internal = InternalFrame( + spark_frame=sdf, + index_names=index_names, + index_spark_columns=index_spark_columns, + index_fields=index_fields, + data_fields=data_fields, + ) + else: + # Otherwise, it loses index. + internal = InternalFrame( + spark_frame=sdf, index_spark_columns=None, data_fields=data_fields + ) + + if should_return_series: + psser = first_series(DataFrame(internal)) + if is_series_groupby: + psser = psser.rename(cast(SeriesGroupBy, self)._psser.name) + return psser + else: + return DataFrame(internal) + + # TODO: implement 'dropna' parameter + def filter(self, func: Callable[[FrameLike], FrameLike]) -> FrameLike: + """ + Return a copy of a DataFrame excluding elements from groups that + do not satisfy the boolean criterion specified by func. + + Parameters + ---------- + f : function + Function to apply to each subframe. Should return True or False. + dropna : Drop groups that do not pass the filter. True by default; + if False, groups that evaluate False are filled with NaNs. + + Returns + ------- + filtered : DataFrame or Series + + Notes + ----- + Each subframe is endowed the attribute 'name' in case you need to know + which group you are working on. + + Examples + -------- + >>> df = ps.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar', + ... 'foo', 'bar'], + ... 'B' : [1, 2, 3, 4, 5, 6], + ... 'C' : [2.0, 5., 8., 1., 2., 9.]}, columns=['A', 'B', 'C']) + >>> grouped = df.groupby('A') + >>> grouped.filter(lambda x: x['B'].mean() > 3.) + A B C + 1 bar 2 5.0 + 3 bar 4 1.0 + 5 bar 6 9.0 + + >>> df.B.groupby(df.A).filter(lambda x: x.mean() > 3.) + 1 2 + 3 4 + 5 6 + Name: B, dtype: int64 + """ + if not callable(func): + raise TypeError("%s object is not callable" % type(func).__name__) + + is_series_groupby = isinstance(self, SeriesGroupBy) + + psdf = self._psdf + + if self._agg_columns_selected: + agg_columns = self._agg_columns + else: + agg_columns = [ + psdf._psser_for(label) + for label in psdf._internal.column_labels + if label not in self._column_labels_to_exclude + ] + + data_schema = ( + psdf[agg_columns]._internal.resolved_copy.spark_frame.drop(*HIDDEN_COLUMNS).schema + ) + + psdf, groupkey_labels, groupkey_names = GroupBy._prepare_group_map_apply( + psdf, self._groupkeys, agg_columns + ) + + if is_series_groupby: + + def pandas_filter(pdf: pd.DataFrame) -> pd.DataFrame: + return pd.DataFrame(pdf.groupby(groupkey_names)[pdf.columns[-1]].filter(func)) + + else: + f = _builtin_table.get(func, func) + + def wrapped_func(pdf: pd.DataFrame) -> pd.DataFrame: + return f(pdf.drop(groupkey_names, axis=1)) + + def pandas_filter(pdf: pd.DataFrame) -> pd.DataFrame: + return pdf.groupby(groupkey_names).filter(wrapped_func).drop(groupkey_names, axis=1) + + sdf = GroupBy._spark_group_map_apply( + psdf, + pandas_filter, + [psdf._internal.spark_column_for(label) for label in groupkey_labels], + data_schema, + retain_index=True, + ) + + psdf = DataFrame(self._psdf[agg_columns]._internal.with_new_sdf(sdf)) + if is_series_groupby: + return cast(FrameLike, first_series(psdf)) + else: + return cast(FrameLike, psdf) + + @staticmethod + def _prepare_group_map_apply( + psdf: DataFrame, groupkeys: List[Series], agg_columns: List[Series] + ) -> Tuple[DataFrame, List[Label], List[str]]: + groupkey_labels: List[Label] = [ + verify_temp_column_name(psdf, "__groupkey_{}__".format(i)) + for i in range(len(groupkeys)) + ] + psdf = psdf[[s.rename(label) for s, label in zip(groupkeys, groupkey_labels)] + agg_columns] + groupkey_names = [label if len(label) > 1 else label[0] for label in groupkey_labels] + return DataFrame(psdf._internal.resolved_copy), groupkey_labels, groupkey_names + + @staticmethod + def _spark_group_map_apply( + psdf: DataFrame, + func: Callable[[pd.DataFrame], pd.DataFrame], + groupkeys_scols: List[Column], + return_schema: StructType, + retain_index: bool, + ) -> SparkDataFrame: + output_func = GroupBy._make_pandas_df_builder_func(psdf, func, return_schema, retain_index) + sdf = psdf._internal.spark_frame.drop(*HIDDEN_COLUMNS) + return sdf.groupby(*groupkeys_scols).applyInPandas(output_func, return_schema) + + @staticmethod + def _make_pandas_df_builder_func( + psdf: DataFrame, + func: Callable[[pd.DataFrame], pd.DataFrame], + return_schema: StructType, + retain_index: bool, + ) -> Callable[[pd.DataFrame], pd.DataFrame]: + """ + Creates a function that can be used inside the pandas UDF. This function can construct + the same pandas DataFrame as if the pandas-on-Spark DataFrame is collected to driver side. + The index, column labels, etc. are re-constructed within the function. + """ + from pyspark.sql.utils import is_timestamp_ntz_preferred + + arguments_for_restore_index = psdf._internal.arguments_for_restore_index + prefer_timestamp_ntz = is_timestamp_ntz_preferred() + + def rename_output(pdf: pd.DataFrame) -> pd.DataFrame: + pdf = InternalFrame.restore_index(pdf.copy(), **arguments_for_restore_index) + + pdf = func(pdf) + + # If schema should be inferred, we don't restore the index. pandas seems to restore + # the index in some cases. + # When Spark output type is specified, without executing it, we don't know + # if we should restore the index or not. For instance, see the example in + # https://github.com/databricks/koalas/issues/628. + pdf, _, _, _, _ = InternalFrame.prepare_pandas_frame( + pdf, retain_index=retain_index, prefer_timestamp_ntz=prefer_timestamp_ntz + ) + + # Just positionally map the column names to given schema's. + pdf.columns = return_schema.names + + return pdf + + return rename_output + + def rank(self, method: str = "average", ascending: bool = True) -> FrameLike: + """ + Provide the rank of values within each group. + + Parameters + ---------- + method : {'average', 'min', 'max', 'first', 'dense'}, default 'average' + * average: average rank of group + * min: lowest rank in group + * max: highest rank in group + * first: ranks assigned in order they appear in the array + * dense: like 'min', but rank always increases by 1 between groups + ascending : boolean, default True + False for ranks by high (1) to low (N) + + Returns + ------- + DataFrame with ranking of values within each group + + Examples + -------- + + >>> df = ps.DataFrame({ + ... 'a': [1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b']) + >>> df + a b + 0 1 1 + 1 1 2 + 2 1 2 + 3 2 2 + 4 2 3 + 5 2 3 + 6 3 3 + 7 3 4 + 8 3 4 + + >>> df.groupby("a").rank().sort_index() + b + 0 1.0 + 1 2.5 + 2 2.5 + 3 1.0 + 4 2.5 + 5 2.5 + 6 1.0 + 7 2.5 + 8 2.5 + + >>> df.b.groupby(df.a).rank(method='max').sort_index() + 0 1.0 + 1 3.0 + 2 3.0 + 3 1.0 + 4 3.0 + 5 3.0 + 6 1.0 + 7 3.0 + 8 3.0 + Name: b, dtype: float64 + + """ + return self._apply_series_op( + lambda sg: sg._psser._rank(method, ascending, part_cols=sg._groupkeys_scols), + should_resolve=True, + ) + + # TODO: add axis parameter + def idxmax(self, skipna: bool = True) -> FrameLike: + """ + Return index of first occurrence of maximum over requested axis in group. + NA/null values are excluded. + + Parameters + ---------- + skipna : boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + See Also + -------- + Series.idxmax + DataFrame.idxmax + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 2, 2, 3], + ... 'b': [1, 2, 3, 4, 5], + ... 'c': [5, 4, 3, 2, 1]}, columns=['a', 'b', 'c']) + + >>> df.groupby(['a'])['b'].idxmax().sort_index() # doctest: +NORMALIZE_WHITESPACE + a + 1 1 + 2 3 + 3 4 + Name: b, dtype: int64 + + >>> df.groupby(['a']).idxmax().sort_index() # doctest: +NORMALIZE_WHITESPACE + b c + a + 1 1 0 + 2 3 2 + 3 4 4 + """ + if self._psdf._internal.index_level != 1: + raise ValueError("idxmax only support one-level index now") + + 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) + index = self._psdf._internal.index_spark_column_names[0] + + stat_exprs = [] + for psser, scol in zip(self._agg_columns, self._agg_columns_scols): + name = psser._internal.data_spark_column_names[0] + + if skipna: + order_column = scol.desc_nulls_last() + else: + order_column = scol.desc_nulls_first() + + window = Window.partitionBy(*groupkey_names).orderBy( + order_column, NATURAL_ORDER_COLUMN_NAME + ) + sdf = sdf.withColumn( + name, F.when(F.row_number().over(window) == 1, scol_for(sdf, index)).otherwise(None) + ) + stat_exprs.append(F.max(scol_for(sdf, name)).alias(name)) + + sdf = sdf.groupby(*groupkey_names).agg(*stat_exprs) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in self._groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(self._groupkeys, groupkey_names) + ], + column_labels=[psser._column_label for psser in self._agg_columns], + data_spark_columns=[ + scol_for(sdf, psser._internal.data_spark_column_names[0]) + for psser in self._agg_columns + ], + ) + return self._handle_output(DataFrame(internal)) + + # TODO: add axis parameter + def idxmin(self, skipna: bool = True) -> FrameLike: + """ + Return index of first occurrence of minimum over requested axis in group. + NA/null values are excluded. + + Parameters + ---------- + skipna : boolean, default True + Exclude NA/null values. If an entire row/column is NA, the result will be NA. + + See Also + -------- + Series.idxmin + DataFrame.idxmin + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 2, 2, 3], + ... 'b': [1, 2, 3, 4, 5], + ... 'c': [5, 4, 3, 2, 1]}, columns=['a', 'b', 'c']) + + >>> df.groupby(['a'])['b'].idxmin().sort_index() # doctest: +NORMALIZE_WHITESPACE + a + 1 0 + 2 2 + 3 4 + Name: b, dtype: int64 + + >>> df.groupby(['a']).idxmin().sort_index() # doctest: +NORMALIZE_WHITESPACE + b c + a + 1 0 1 + 2 2 3 + 3 4 4 + """ + if self._psdf._internal.index_level != 1: + raise ValueError("idxmin only support one-level index now") + + 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) + index = self._psdf._internal.index_spark_column_names[0] + + stat_exprs = [] + for psser, scol in zip(self._agg_columns, self._agg_columns_scols): + name = psser._internal.data_spark_column_names[0] + + if skipna: + order_column = scol.asc_nulls_last() + else: + order_column = scol.asc_nulls_first() + + window = Window.partitionBy(*groupkey_names).orderBy( + order_column, NATURAL_ORDER_COLUMN_NAME + ) + sdf = sdf.withColumn( + name, F.when(F.row_number().over(window) == 1, scol_for(sdf, index)).otherwise(None) + ) + stat_exprs.append(F.max(scol_for(sdf, name)).alias(name)) + + sdf = sdf.groupby(*groupkey_names).agg(*stat_exprs) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in self._groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(self._groupkeys, groupkey_names) + ], + column_labels=[psser._column_label for psser in self._agg_columns], + data_spark_columns=[ + scol_for(sdf, psser._internal.data_spark_column_names[0]) + for psser in self._agg_columns + ], + ) + return self._handle_output(DataFrame(internal)) + + def fillna( + self, + value: Optional[Any] = None, + method: Optional[str] = None, + axis: Optional[Axis] = None, + inplace: bool = False, + limit: Optional[int] = None, + ) -> FrameLike: + """Fill NA/NaN values in group. + + Parameters + ---------- + value : scalar, dict, Series + Value to use to fill holes. alternately a dict/Series of values + specifying which value to use for each column. + DataFrame is not supported. + method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None + Method to use for filling holes in reindexed Series pad / ffill: propagate last valid + observation forward to next valid backfill / bfill: + use NEXT valid observation to fill gap + axis : {0 or `index`} + 1 and `columns` are not supported. + inplace : boolean, default False + Fill in place (do not create a new object) + limit : int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame + DataFrame with NA entries filled. + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'A': [1, 1, 2, 2], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> df + A B C D + 0 1 2.0 NaN 0 + 1 1 4.0 NaN 1 + 2 2 NaN NaN 5 + 3 2 3.0 1.0 4 + + We can also propagate non-null values forward or backward in group. + + >>> df.groupby(['A'])['B'].fillna(method='ffill').sort_index() + 0 2.0 + 1 4.0 + 2 NaN + 3 3.0 + Name: B, dtype: float64 + + >>> df.groupby(['A']).fillna(method='bfill').sort_index() + B C D + 0 2.0 NaN 0 + 1 4.0 NaN 1 + 2 3.0 1.0 5 + 3 3.0 1.0 4 + """ + return self._apply_series_op( + lambda sg: sg._psser._fillna( + value=value, method=method, axis=axis, limit=limit, part_cols=sg._groupkeys_scols + ), + should_resolve=(method is not None), + ) + + def bfill(self, limit: Optional[int] = None) -> FrameLike: + """ + Synonym for `DataFrame.fillna()` with ``method=`bfill```. + + Parameters + ---------- + axis : {0 or `index`} + 1 and `columns` are not supported. + inplace : boolean, default False + Fill in place (do not create a new object) + limit : int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame + DataFrame with NA entries filled. + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'A': [1, 1, 2, 2], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> df + A B C D + 0 1 2.0 NaN 0 + 1 1 4.0 NaN 1 + 2 2 NaN NaN 5 + 3 2 3.0 1.0 4 + + Propagate non-null values backward. + + >>> df.groupby(['A']).bfill().sort_index() + B C D + 0 2.0 NaN 0 + 1 4.0 NaN 1 + 2 3.0 1.0 5 + 3 3.0 1.0 4 + """ + return self.fillna(method="bfill", limit=limit) + + def backfill(self, limit: Optional[int] = None) -> FrameLike: + """ + Alias for bfill. + + .. deprecated:: 3.4.0 + """ + warnings.warn( + "The GroupBy.backfill method is deprecated " + "and will be removed in a future version. " + "Use GroupBy.bfill instead.", + FutureWarning, + ) + return self.bfill(limit=limit) + + def ffill(self, limit: Optional[int] = None) -> FrameLike: + """ + Synonym for `DataFrame.fillna()` with ``method=`ffill```. + + Parameters + ---------- + axis : {0 or `index`} + 1 and `columns` are not supported. + inplace : boolean, default False + Fill in place (do not create a new object) + limit : int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + DataFrame + DataFrame with NA entries filled. + + Examples + -------- + >>> df = ps.DataFrame({ + ... 'A': [1, 1, 2, 2], + ... 'B': [2, 4, None, 3], + ... 'C': [None, None, None, 1], + ... 'D': [0, 1, 5, 4] + ... }, + ... columns=['A', 'B', 'C', 'D']) + >>> df + A B C D + 0 1 2.0 NaN 0 + 1 1 4.0 NaN 1 + 2 2 NaN NaN 5 + 3 2 3.0 1.0 4 + + Propagate non-null values forward. + + >>> df.groupby(['A']).ffill().sort_index() + B C D + 0 2.0 NaN 0 + 1 4.0 NaN 1 + 2 NaN NaN 5 + 3 3.0 1.0 4 + """ + return self.fillna(method="ffill", limit=limit) + + def pad(self, limit: Optional[int] = None) -> FrameLike: + """ + Alias for ffill. + + .. deprecated:: 3.4.0 + """ + warnings.warn( + "The GroupBy.pad method is deprecated " + "and will be removed in a future version. " + "Use GroupBy.ffill instead.", + FutureWarning, + ) + return self.ffill(limit=limit) + + def _limit(self, n: int, asc: bool) -> FrameLike: + """ + Private function for tail and head. + """ + psdf = self._psdf + + if self._agg_columns_selected: + agg_columns = self._agg_columns + else: + agg_columns = [ + psdf._psser_for(label) + for label in psdf._internal.column_labels + if label not in self._column_labels_to_exclude + ] + + psdf, groupkey_labels, _ = GroupBy._prepare_group_map_apply( + psdf, + self._groupkeys, + agg_columns, + ) + + groupkey_scols = [psdf._internal.spark_column_for(label) for label in groupkey_labels] + + sdf = psdf._internal.spark_frame + + window = Window.partitionBy(*groupkey_scols) + # This part is handled differently depending on whether it is a tail or a head. + ordered_window = ( + window.orderBy(F.col(NATURAL_ORDER_COLUMN_NAME).asc()) + if asc + else window.orderBy(F.col(NATURAL_ORDER_COLUMN_NAME).desc()) + ) + + if n >= 0 or LooseVersion(pd.__version__) < LooseVersion("1.4.0"): + tmp_row_num_col = verify_temp_column_name(sdf, "__row_number__") + sdf = ( + sdf.withColumn(tmp_row_num_col, F.row_number().over(ordered_window)) + .filter(F.col(tmp_row_num_col) <= n) + .drop(tmp_row_num_col) + ) + else: + # Pandas supports Groupby positional indexing since v1.4.0 + # https://pandas.pydata.org/docs/whatsnew/v1.4.0.html#groupby-positional-indexing + # + # To support groupby positional indexing, we need add a `__tmp_lag__` column to help + # us filtering rows before the specified offset row. + # + # For example for the dataframe: + # >>> df = ps.DataFrame([["g", "g0"], + # ... ["g", "g1"], + # ... ["g", "g2"], + # ... ["g", "g3"], + # ... ["h", "h0"], + # ... ["h", "h1"]], columns=["A", "B"]) + # >>> df.groupby("A").head(-1) + # + # Below is a result to show the `__tmp_lag__` column for above df, the limit n is + # `-1`, the `__tmp_lag__` will be set to `0` in rows[:-1], and left will be set to + # `null`: + # + # >>> sdf.withColumn(tmp_lag_col, F.lag(F.lit(0), -1).over(ordered_window)) + # +-----------------+--------------+---+---+-----------------+-----------+ + # |__index_level_0__|__groupkey_0__| A| B|__natural_order__|__tmp_lag__| + # +-----------------+--------------+---+---+-----------------+-----------+ + # | 0| g| g| g0| 0| 0| + # | 1| g| g| g1| 8589934592| 0| + # | 2| g| g| g2| 17179869184| 0| + # | 3| g| g| g3| 25769803776| null| + # | 4| h| h| h0| 34359738368| 0| + # | 5| h| h| h1| 42949672960| null| + # +-----------------+--------------+---+---+-----------------+-----------+ + # + tmp_lag_col = verify_temp_column_name(sdf, "__tmp_lag__") + sdf = ( + sdf.withColumn(tmp_lag_col, F.lag(F.lit(0), n).over(ordered_window)) + .where(~F.isnull(F.col(tmp_lag_col))) + .drop(tmp_lag_col) + ) + + internal = psdf._internal.with_new_sdf(sdf) + return self._handle_output(DataFrame(internal).drop(groupkey_labels, axis=1)) + + def head(self, n: int = 5) -> FrameLike: + """ + Return first n rows of each group. + + Returns + ------- + DataFrame or Series + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [2, 3, 1, 4, 6, 9, 8, 10, 7, 5], + ... 'c': [3, 5, 2, 5, 1, 2, 6, 4, 3, 6]}, + ... columns=['a', 'b', 'c'], + ... index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6]) + >>> df + a b c + 7 1 2 3 + 2 1 3 5 + 4 1 1 2 + 1 1 4 5 + 3 2 6 1 + 4 2 9 2 + 9 2 8 6 + 10 3 10 4 + 5 3 7 3 + 6 3 5 6 + + >>> df.groupby('a').head(2).sort_index() + a b c + 2 1 3 5 + 3 2 6 1 + 4 2 9 2 + 5 3 7 3 + 7 1 2 3 + 10 3 10 4 + + >>> df.groupby('a')['b'].head(2).sort_index() + 2 3 + 3 6 + 4 9 + 5 7 + 7 2 + 10 10 + Name: b, dtype: int64 + + Supports Groupby positional indexing Since pandas on Spark 3.4 (with pandas 1.4+): + + >>> df = ps.DataFrame([["g", "g0"], + ... ["g", "g1"], + ... ["g", "g2"], + ... ["g", "g3"], + ... ["h", "h0"], + ... ["h", "h1"]], columns=["A", "B"]) + >>> df.groupby("A").head(-1) # doctest: +SKIP + A B + 0 g g0 + 1 g g1 + 2 g g2 + 4 h h0 + """ + return self._limit(n, asc=True) + + def tail(self, n: int = 5) -> FrameLike: + """ + Return last n rows of each group. + + Similar to `.apply(lambda x: x.tail(n))`, but it returns a subset of rows from + the original DataFrame with original index and order preserved (`as_index` flag is ignored). + + Does not work for negative values of n. + + Returns + ------- + DataFrame or Series + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [2, 3, 1, 4, 6, 9, 8, 10, 7, 5], + ... 'c': [3, 5, 2, 5, 1, 2, 6, 4, 3, 6]}, + ... columns=['a', 'b', 'c'], + ... index=[7, 2, 3, 1, 3, 4, 9, 10, 5, 6]) + >>> df + a b c + 7 1 2 3 + 2 1 3 5 + 3 1 1 2 + 1 1 4 5 + 3 2 6 1 + 4 2 9 2 + 9 2 8 6 + 10 3 10 4 + 5 3 7 3 + 6 3 5 6 + + >>> df.groupby('a').tail(2).sort_index() + a b c + 1 1 4 5 + 3 1 1 2 + 4 2 9 2 + 5 3 7 3 + 6 3 5 6 + 9 2 8 6 + + >>> df.groupby('a')['b'].tail(2).sort_index() + 1 4 + 3 1 + 4 9 + 5 7 + 6 5 + 9 8 + Name: b, dtype: int64 + + Supports Groupby positional indexing Since pandas on Spark 3.4 (with pandas 1.4+): + + >>> df = ps.DataFrame([["g", "g0"], + ... ["g", "g1"], + ... ["g", "g2"], + ... ["g", "g3"], + ... ["h", "h0"], + ... ["h", "h1"]], columns=["A", "B"]) + >>> df.groupby("A").tail(-1) # doctest: +SKIP + A B + 3 g g3 + 2 g g2 + 1 g g1 + 5 h h1 + """ + return self._limit(n, asc=False) + + def shift(self, periods: int = 1, fill_value: Optional[Any] = None) -> FrameLike: + """ + Shift each group by periods observations. + + Parameters + ---------- + periods : integer, default 1 + number of periods to shift + fill_value : optional + + Returns + ------- + Series or DataFrame + Object shifted within each group. + + Examples + -------- + + >>> df = ps.DataFrame({ + ... 'a': [1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b']) + >>> df + a b + 0 1 1 + 1 1 2 + 2 1 2 + 3 2 2 + 4 2 3 + 5 2 3 + 6 3 3 + 7 3 4 + 8 3 4 + + >>> df.groupby('a').shift().sort_index() # doctest: +SKIP + b + 0 NaN + 1 1.0 + 2 2.0 + 3 NaN + 4 2.0 + 5 3.0 + 6 NaN + 7 3.0 + 8 4.0 + + >>> df.groupby('a').shift(periods=-1, fill_value=0).sort_index() # doctest: +SKIP + b + 0 2 + 1 2 + 2 0 + 3 3 + 4 3 + 5 0 + 6 4 + 7 4 + 8 0 + """ + return self._apply_series_op( + lambda sg: sg._psser._shift(periods, fill_value, part_cols=sg._groupkeys_scols), + should_resolve=True, + ) + + def transform(self, func: Callable[..., pd.Series], *args: Any, **kwargs: Any) -> FrameLike: + """ + Apply function column-by-column to the GroupBy object. + + The function passed to `transform` must take a Series as its first + argument and return a Series. The given function is executed for + each series in each grouped data. + + While `transform` is a very flexible method, its downside is that + using it can be quite a bit slower than using more specific methods + like `agg` or `transform`. pandas-on-Spark offers a wide range of method that will + be much faster than using `transform` for their specific purposes, so try to + use them before reaching for `transform`. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def convert_to_string(x) -> ps.Series[str]: + ... return x.apply("a string {}".format) + + When the given function has the return type annotated, the original index of the + GroupBy object will be lost, and a default index will be attached to the result. + Please be careful about configuring the default index. See also `Default Index Type + `_. + + .. note:: the series within ``func`` is actually a pandas series. Therefore, + any pandas API within this function is allowed. + + + Parameters + ---------- + func : callable + A callable that takes a Series as its first argument, and + returns a Series. + *args + Positional arguments to pass to func. + **kwargs + Keyword arguments to pass to func. + + Returns + ------- + applied : DataFrame + + See Also + -------- + aggregate : Apply aggregate function to the GroupBy object. + Series.apply : Apply a function to a Series. + + Examples + -------- + + >>> df = ps.DataFrame({'A': [0, 0, 1], + ... 'B': [1, 2, 3], + ... 'C': [4, 6, 5]}, columns=['A', 'B', 'C']) + + >>> g = df.groupby('A') + + Notice that ``g`` has two groups, ``0`` and ``1``. + Calling `transform` in various ways, we can get different grouping results: + Below the functions passed to `transform` takes a Series as + its argument and returns a Series. `transform` applies the function on each series + in each grouped data, and combine them into a new DataFrame: + + >>> def convert_to_string(x) -> ps.Series[str]: + ... return x.apply("a string {}".format) + >>> g.transform(convert_to_string) # doctest: +NORMALIZE_WHITESPACE + B C + 0 a string 1 a string 4 + 1 a string 2 a string 6 + 2 a string 3 a string 5 + + >>> def plus_max(x) -> ps.Series[int]: + ... return x + x.max() + >>> g.transform(plus_max) # doctest: +NORMALIZE_WHITESPACE + B C + 0 3 10 + 1 4 12 + 2 6 10 + + You can omit the type hint and let pandas-on-Spark infer its type. + + >>> def plus_min(x): + ... return x + x.min() + >>> g.transform(plus_min) # doctest: +NORMALIZE_WHITESPACE + B C + 0 2 8 + 1 3 10 + 2 6 10 + + In case of Series, it works as below. + + >>> df.B.groupby(df.A).transform(plus_max) + 0 3 + 1 4 + 2 6 + Name: B, dtype: int64 + + >>> (df * -1).B.groupby(df.A).transform(abs) + 0 1 + 1 2 + 2 3 + Name: B, dtype: int64 + + You can also specify extra arguments to pass to the function. + + >>> def calculation(x, y, z) -> ps.Series[int]: + ... return x + x.min() + y + z + >>> g.transform(calculation, 5, z=20) # doctest: +NORMALIZE_WHITESPACE + B C + 0 27 33 + 1 28 35 + 2 31 35 + """ + if not callable(func): + raise TypeError("%s object is not callable" % type(func).__name__) + + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + + psdf, groupkey_labels, groupkey_names = GroupBy._prepare_group_map_apply( + self._psdf, self._groupkeys, agg_columns=self._agg_columns + ) + + def pandas_transform(pdf: pd.DataFrame) -> pd.DataFrame: + return pdf.groupby(groupkey_names).transform(func, *args, **kwargs) + + should_infer_schema = return_sig is None + + if should_infer_schema: + # Here we execute with the first 1000 to get the return type. + # If the records were less than 1000, it uses pandas API directly for a shortcut. + log_advice( + "If the type hints is not specified for `groupby.transform`, " + "it is expensive to infer the data type internally." + ) + limit = get_option("compute.shortcut_limit") + pdf = psdf.head(limit + 1)._to_internal_pandas() + pdf = pdf.groupby(groupkey_names).transform(func, *args, **kwargs) + psdf_from_pandas: DataFrame = DataFrame(pdf) + return_schema = force_decimal_precision_scale( + as_nullable_spark_type( + psdf_from_pandas._internal.spark_frame.drop(*HIDDEN_COLUMNS).schema + ) + ) + if len(pdf) <= limit: + return self._handle_output(psdf_from_pandas) + + sdf = GroupBy._spark_group_map_apply( + psdf, + pandas_transform, + [psdf._internal.spark_column_for(label) for label in groupkey_labels], + return_schema, + retain_index=True, + ) + # If schema is inferred, we can restore indexes too. + internal = psdf_from_pandas._internal.with_new_sdf( + sdf, + index_fields=[ + field.copy(nullable=True) for field in psdf_from_pandas._internal.index_fields + ], + data_fields=[ + field.copy(nullable=True) for field in psdf_from_pandas._internal.data_fields + ], + ) + else: + return_type = infer_return_type(func) + if not isinstance(return_type, SeriesType): + raise TypeError( + "Expected the return type of this function to be of Series type, " + "but found type {}".format(return_type) + ) + + dtype = return_type.dtype + spark_type = return_type.spark_type + + data_fields = [ + InternalField(dtype=dtype, struct_field=StructField(name=c, dataType=spark_type)) + for c in psdf._internal.data_spark_column_names + if c not in groupkey_names + ] + + return_schema = StructType([field.struct_field for field in data_fields]) + + sdf = GroupBy._spark_group_map_apply( + psdf, + pandas_transform, + [psdf._internal.spark_column_for(label) for label in groupkey_labels], + return_schema, + retain_index=False, + ) + # Otherwise, it loses index. + internal = InternalFrame( + spark_frame=sdf, index_spark_columns=None, data_fields=data_fields + ) + + return self._handle_output(DataFrame(internal)) + + def nunique(self, dropna: bool = True) -> FrameLike: + """ + Return DataFrame with number of distinct observations per group for each column. + + Parameters + ---------- + dropna : boolean, default True + Don’t include NaN in the counts. + + Returns + ------- + nunique : DataFrame or Series + + Examples + -------- + + >>> df = ps.DataFrame({'id': ['spam', 'egg', 'egg', 'spam', + ... 'ham', 'ham'], + ... 'value1': [1, 5, 5, 2, 5, 5], + ... 'value2': list('abbaxy')}, columns=['id', 'value1', 'value2']) + >>> df + id value1 value2 + 0 spam 1 a + 1 egg 5 b + 2 egg 5 b + 3 spam 2 a + 4 ham 5 x + 5 ham 5 y + + >>> df.groupby('id').nunique().sort_index() # doctest: +SKIP + value1 value2 + id + egg 1 1 + ham 1 2 + spam 2 1 + + >>> df.groupby('id')['value1'].nunique().sort_index() # doctest: +NORMALIZE_WHITESPACE + id + egg 1 + ham 1 + spam 2 + Name: value1, dtype: int64 + """ + if dropna: + + def stat_function(col: Column) -> Column: + return F.countDistinct(col) + + else: + + def stat_function(col: Column) -> Column: + return F.countDistinct(col) + F.when( + F.count(F.when(col.isNull(), 1).otherwise(None)) >= 1, 1 + ).otherwise(0) + + return self._reduce_for_stat_function(stat_function) + + def rolling( + self, window: int, min_periods: Optional[int] = None + ) -> "RollingGroupby[FrameLike]": + """ + Return an rolling grouper, providing rolling + functionality per group. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + Parameters + ---------- + window : int, or offset + Size of the moving window. + This is the number of observations used for calculating the statistic. + Each window will be a fixed size. + + min_periods : int, default 1 + Minimum number of observations in window required to have a value + (otherwise result is NA). + + See Also + -------- + Series.groupby + DataFrame.groupby + """ + from pyspark.pandas.window import RollingGroupby + + return RollingGroupby(self, window, min_periods=min_periods) + + def expanding(self, min_periods: int = 1) -> "ExpandingGroupby[FrameLike]": + """ + Return an expanding grouper, providing expanding + functionality per group. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + Parameters + ---------- + min_periods : int, default 1 + Minimum number of observations in window required to have a value + (otherwise result is NA). + + See Also + -------- + Series.groupby + DataFrame.groupby + """ + from pyspark.pandas.window import ExpandingGroupby + + return ExpandingGroupby(self, min_periods=min_periods) + + # TODO: 'adjust', 'axis', 'method' parameter should be implemented. + def ewm( + self, + com: Optional[float] = None, + span: Optional[float] = None, + halflife: Optional[float] = None, + alpha: Optional[float] = None, + min_periods: Optional[int] = None, + ignore_na: bool = False, + ) -> "ExponentialMovingGroupby[FrameLike]": + """ + Return an ewm grouper, providing ewm functionality per group. + + .. note:: 'min_periods' in pandas-on-Spark works as a fixed window size unlike pandas. + Unlike pandas, NA is also counted as the period. This might be changed + soon. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + com : float, optional + Specify decay in terms of center of mass. + alpha = 1 / (1 + com), for com >= 0. + + span : float, optional + Specify decay in terms of span. + alpha = 2 / (span + 1), for span >= 1. + + halflife : float, optional + Specify decay in terms of half-life. + alpha = 1 - exp(-ln(2) / halflife), for halflife > 0. + + alpha : float, optional + Specify smoothing factor alpha directly. + 0 < alpha <= 1. + + min_periods : int, default None + Minimum number of observations in window required to have a value + (otherwise result is NA). + + ignore_na : bool, default False + Ignore missing values when calculating weights. + + - When ``ignore_na=False`` (default), weights are based on absolute positions. + For example, the weights of :math:`x_0` and :math:`x_2` used in calculating + the final weighted average of [:math:`x_0`, None, :math:`x_2`] are + :math:`(1-\alpha)^2` and :math:`1` if ``adjust=True``, and + :math:`(1-\alpha)^2` and :math:`\alpha` if ``adjust=False``. + + - When ``ignore_na=True``, weights are based + on relative positions. For example, the weights of :math:`x_0` and :math:`x_2` + used in calculating the final weighted average of + [:math:`x_0`, None, :math:`x_2`] are :math:`1-\alpha` and :math:`1` if + ``adjust=True``, and :math:`1-\alpha` and :math:`\alpha` if ``adjust=False``. + """ + from pyspark.pandas.window import ExponentialMovingGroupby + + return ExponentialMovingGroupby( + self, + com=com, + span=span, + halflife=halflife, + alpha=alpha, + min_periods=min_periods, + ignore_na=ignore_na, + ) + + def get_group(self, name: Union[Name, List[Name]]) -> FrameLike: + """ + Construct DataFrame from group with provided name. + + Parameters + ---------- + name : object + The name of the group to get as a DataFrame. + + Returns + ------- + group : same type as obj + + Examples + -------- + >>> psdf = ps.DataFrame([('falcon', 'bird', 389.0), + ... ('parrot', 'bird', 24.0), + ... ('lion', 'mammal', 80.5), + ... ('monkey', 'mammal', np.nan)], + ... columns=['name', 'class', 'max_speed'], + ... index=[0, 2, 3, 1]) + >>> psdf + name class max_speed + 0 falcon bird 389.0 + 2 parrot bird 24.0 + 3 lion mammal 80.5 + 1 monkey mammal NaN + + >>> psdf.groupby("class").get_group("bird").sort_index() + name class max_speed + 0 falcon bird 389.0 + 2 parrot bird 24.0 + + >>> psdf.groupby("class").get_group("mammal").sort_index() + name class max_speed + 1 monkey mammal NaN + 3 lion mammal 80.5 + """ + groupkeys = self._groupkeys + if not is_hashable(name): + raise TypeError("unhashable type: '{}'".format(type(name).__name__)) + elif len(groupkeys) > 1: + if not isinstance(name, tuple): + raise ValueError("must supply a tuple to get_group with multiple grouping keys") + if len(groupkeys) != len(name): + raise ValueError( + "must supply a same-length tuple to get_group with multiple grouping keys" + ) + if not is_list_like(name): + name = [name] + cond = F.lit(True) + for groupkey, item in zip(groupkeys, name): + scol = groupkey.spark.column + cond = cond & (scol == item) + if self._agg_columns_selected: + internal = self._psdf._internal + spark_frame = internal.spark_frame.select( + internal.index_spark_columns + self._agg_columns_scols + ).filter(cond) + + internal = internal.copy( + spark_frame=spark_frame, + index_spark_columns=[ + scol_for(spark_frame, col) for col in internal.index_spark_column_names + ], + column_labels=[s._column_label for s in self._agg_columns], + data_spark_columns=[ + scol_for(spark_frame, s._internal.data_spark_column_names[0]) + for s in self._agg_columns + ], + data_fields=[s._internal.data_fields[0] for s in self._agg_columns], + ) + else: + internal = self._psdf._internal.with_filter(cond) + if internal.spark_frame.head() is None: + raise KeyError(name) + + return self._handle_output(DataFrame(internal)) + + def median(self, numeric_only: Optional[bool] = True, accuracy: int = 10000) -> FrameLike: + """ + Compute median of groups, excluding missing values. + + For multiple groupings, the result index will be a MultiIndex + + .. note:: Unlike pandas', the median in pandas-on-Spark is an approximated median based upon + approximate percentile computation because computing median across a large dataset + is extremely expensive. + + Parameters + ---------- + numeric_only : bool, default False + Include only float, int, boolean columns. If None, will attempt to use + everything, then use only numeric data. + + .. versionadded:: 3.4.0 + + Returns + ------- + Series or DataFrame + Median of values within each group. + + Examples + -------- + >>> psdf = ps.DataFrame({'a': [1., 1., 1., 1., 2., 2., 2., 3., 3., 3.], + ... 'b': [2., 3., 1., 4., 6., 9., 8., 10., 7., 5.], + ... 'c': [3., 5., 2., 5., 1., 2., 6., 4., 3., 6.]}, + ... columns=['a', 'b', 'c'], + ... index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6]) + >>> psdf + a b c + 7 1.0 2.0 3.0 + 2 1.0 3.0 5.0 + 4 1.0 1.0 2.0 + 1 1.0 4.0 5.0 + 3 2.0 6.0 1.0 + 4 2.0 9.0 2.0 + 9 2.0 8.0 6.0 + 10 3.0 10.0 4.0 + 5 3.0 7.0 3.0 + 6 3.0 5.0 6.0 + + DataFrameGroupBy + + >>> psdf.groupby('a').median().sort_index() # doctest: +NORMALIZE_WHITESPACE + b c + a + 1.0 2.0 3.0 + 2.0 8.0 2.0 + 3.0 7.0 4.0 + + SeriesGroupBy + + >>> psdf.groupby('a')['b'].median().sort_index() + a + 1.0 2.0 + 2.0 8.0 + 3.0 7.0 + Name: b, dtype: float64 + """ + if not isinstance(accuracy, int): + raise TypeError( + "accuracy must be an integer; however, got [%s]" % type(accuracy).__name__ + ) + + self._validate_agg_columns(numeric_only=numeric_only, function_name="median") + + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `True` in 4.0.0.", + FutureWarning, + ) + + def stat_function(col: Column) -> Column: + return F.percentile_approx(col, 0.5, accuracy) + + return self._reduce_for_stat_function( + stat_function, + accepted_spark_types=(NumericType,), + bool_to_numeric=True, + ) + + def _validate_agg_columns(self, numeric_only: Optional[bool], function_name: str) -> None: + """Validate aggregation columns and raise an error or a warning following pandas.""" + has_non_numeric = False + for _agg_col in self._agg_columns: + if not isinstance(_agg_col.spark.data_type, (NumericType, BooleanType)): + has_non_numeric = True + break + if has_non_numeric: + if isinstance(self, SeriesGroupBy): + raise TypeError("Only numeric aggregation column is accepted.") + + if not numeric_only and has_non_numeric: + warnings.warn( + "Dropping invalid columns in DataFrameGroupBy.%s is deprecated. " + "In a future version, a TypeError will be raised. " + "Before calling .%s, select only columns which should be " + "valid for the function." % (function_name, function_name), + FutureWarning, + ) + + def _reduce_for_stat_function( + self, + sfun: Callable[[Column], Column], + accepted_spark_types: Optional[Tuple[Type[DataType], ...]] = None, + bool_to_numeric: bool = False, + **kwargs: Any, + ) -> FrameLike: + """Apply an aggregate function `sfun` per column and reduce to a FrameLike. + + Parameters + ---------- + sfun : The aggregate function to apply per column. + accepted_spark_types: Accepted spark types of columns to be aggregated; + default None means all spark types are accepted. + bool_to_numeric: If True, boolean columns are converted to numeric columns, which + are accepted for all statistical functions regardless of + `accepted_spark_types`. + """ + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + internal, _, sdf = self._prepare_reduce( + groupkey_names, accepted_spark_types, bool_to_numeric + ) + psdf: DataFrame = DataFrame(internal) + + if len(psdf._internal.column_labels) > 0: + min_count = kwargs.get("min_count", 0) + stat_exprs = [] + for label in psdf._internal.column_labels: + psser = psdf._psser_for(label) + input_scol = psser._dtype_op.nan_to_null(psser).spark.column + output_scol = sfun(input_scol) + + if min_count > 0: + output_scol = F.when( + F.count(F.when(~F.isnull(input_scol), F.lit(0))) >= min_count, output_scol + ) + + stat_exprs.append(output_scol.alias(psser._internal.data_spark_column_names[0])) + sdf = sdf.groupby(*groupkey_names).agg(*stat_exprs) + else: + sdf = sdf.select(*groupkey_names).distinct() + + internal = internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + data_spark_columns=[scol_for(sdf, col) for col in internal.data_spark_column_names], + data_fields=None, + ) + psdf = DataFrame(internal) + + return self._prepare_return(psdf) + + def _prepare_return(self, psdf: DataFrame) -> FrameLike: + if self._dropna: + psdf = DataFrame( + psdf._internal.with_new_sdf( + psdf._internal.spark_frame.dropna( + subset=psdf._internal.index_spark_column_names + ) + ) + ) + if not self._as_index: + should_drop_index = set( + i for i, gkey in enumerate(self._groupkeys) if gkey._psdf is not self._psdf + ) + if len(should_drop_index) > 0: + psdf = psdf.reset_index(level=should_drop_index, drop=True) + if len(should_drop_index) < len(self._groupkeys): + psdf = psdf.reset_index() + return self._handle_output(psdf) + + def _prepare_reduce( + self, + groupkey_names: List, + accepted_spark_types: Optional[Tuple[Type[DataType], ...]] = None, + bool_to_numeric: bool = False, + ) -> Tuple[InternalFrame, List[Series], SparkDataFrame]: + groupkey_scols = [s.alias(name) for s, name in zip(self._groupkeys_scols, groupkey_names)] + agg_columns = [] + for psser in self._agg_columns: + if bool_to_numeric and isinstance(psser.spark.data_type, BooleanType): + agg_columns.append(psser.astype(int)) + elif (accepted_spark_types is None) or isinstance( + psser.spark.data_type, accepted_spark_types + ): + agg_columns.append(psser) + sdf = self._psdf._internal.spark_frame.select( + *groupkey_scols, *[psser.spark.column for psser in agg_columns] + ) + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in self._groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(self._groupkeys, groupkey_names) + ], + data_spark_columns=[ + scol_for(sdf, psser._internal.data_spark_column_names[0]) for psser in agg_columns + ], + column_labels=[psser._column_label for psser in agg_columns], + data_fields=[psser._internal.data_fields[0] for psser in agg_columns], + column_label_names=self._psdf._internal.column_label_names, + ) + return internal, agg_columns, sdf + + @staticmethod + def _resolve_grouping_from_diff_dataframes( + psdf: DataFrame, by: List[Union[Series, Label]] + ) -> Tuple[DataFrame, List[Series], Set[Label]]: + column_labels_level = psdf._internal.column_labels_level + + column_labels = [] + additional_pssers = [] + additional_column_labels = [] + tmp_column_labels = set() + for i, col_or_s in enumerate(by): + if isinstance(col_or_s, Series): + if col_or_s._psdf is psdf: + column_labels.append(col_or_s._column_label) + elif same_anchor(col_or_s, psdf): + temp_label = verify_temp_column_name(psdf, "__tmp_groupkey_{}__".format(i)) + column_labels.append(temp_label) + additional_pssers.append(col_or_s.rename(temp_label)) + additional_column_labels.append(temp_label) + else: + temp_label = verify_temp_column_name( + psdf, + tuple( + ([""] * (column_labels_level - 1)) + ["__tmp_groupkey_{}__".format(i)] + ), + ) + column_labels.append(temp_label) + tmp_column_labels.add(temp_label) + elif isinstance(col_or_s, tuple): + psser = psdf[col_or_s] + if not isinstance(psser, Series): + raise ValueError(name_like_string(col_or_s)) + column_labels.append(col_or_s) + else: + raise ValueError(col_or_s) + + psdf = DataFrame( + psdf._internal.with_new_columns( + [psdf._psser_for(label) for label in psdf._internal.column_labels] + + additional_pssers + ) + ) + + def assign_columns( + psdf: DataFrame, this_column_labels: List[Label], that_column_labels: List[Label] + ) -> Iterator[Tuple[Series, Label]]: + raise NotImplementedError( + "Duplicated labels with groupby() and " + "'compute.ops_on_diff_frames' option is not supported currently " + "Please use unique labels in series and frames." + ) + + for col_or_s, label in zip(by, column_labels): + if label in tmp_column_labels: + psser = col_or_s + psdf = align_diff_frames( + assign_columns, + psdf, + psser.rename(label), + fillna=False, + how="inner", + preserve_order_column=True, + ) + + tmp_column_labels |= set(additional_column_labels) + + new_by_series = [] + for col_or_s, label in zip(by, column_labels): + if label in tmp_column_labels: + psser = col_or_s + new_by_series.append(psdf._psser_for(label).rename(psser.name)) + else: + new_by_series.append(psdf._psser_for(label)) + + return psdf, new_by_series, tmp_column_labels + + @staticmethod + def _resolve_grouping(psdf: DataFrame, by: List[Union[Series, Label]]) -> List[Series]: + new_by_series = [] + for col_or_s in by: + if isinstance(col_or_s, Series): + new_by_series.append(col_or_s) + elif isinstance(col_or_s, tuple): + psser = psdf[col_or_s] + if not isinstance(psser, Series): + raise ValueError(name_like_string(col_or_s)) + new_by_series.append(psser) + else: + raise ValueError(col_or_s) + return new_by_series + + +class DataFrameGroupBy(GroupBy[DataFrame]): + @staticmethod + def _build( + psdf: DataFrame, by: List[Union[Series, Label]], as_index: bool, dropna: bool + ) -> "DataFrameGroupBy": + if any(isinstance(col_or_s, Series) and not same_anchor(psdf, col_or_s) for col_or_s in by): + ( + psdf, + new_by_series, + column_labels_to_exclude, + ) = GroupBy._resolve_grouping_from_diff_dataframes(psdf, by) + else: + new_by_series = GroupBy._resolve_grouping(psdf, by) + column_labels_to_exclude = set() + return DataFrameGroupBy( + psdf, + new_by_series, + as_index=as_index, + dropna=dropna, + column_labels_to_exclude=column_labels_to_exclude, + ) + + def __init__( + self, + psdf: DataFrame, + by: List[Series], + as_index: bool, + dropna: bool, + column_labels_to_exclude: Set[Label], + agg_columns: List[Label] = None, + ): + agg_columns_selected = agg_columns is not None + if agg_columns_selected: + for label in agg_columns: + if label in column_labels_to_exclude: + raise KeyError(label) + else: + agg_columns = [ + label + for label in psdf._internal.column_labels + if not any(label == key._column_label and key._psdf is psdf for key in by) + and label not in column_labels_to_exclude + ] + + super().__init__( + psdf=psdf, + groupkeys=by, + as_index=as_index, + dropna=dropna, + column_labels_to_exclude=column_labels_to_exclude, + agg_columns_selected=agg_columns_selected, + agg_columns=[psdf[label] for label in agg_columns], + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeDataFrameGroupBy, item): + property_or_func = getattr(MissingPandasLikeDataFrameGroupBy, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + return self.__getitem__(item) + + def __getitem__(self, item: Any) -> GroupBy: + if self._as_index and is_name_like_value(item): + return SeriesGroupBy( + self._psdf._psser_for(item if is_name_like_tuple(item) else (item,)), + self._groupkeys, + dropna=self._dropna, + ) + else: + if is_name_like_tuple(item): + item = [item] + elif is_name_like_value(item): + item = [(item,)] + else: + item = [i if is_name_like_tuple(i) else (i,) for i in item] + if not self._as_index: + groupkey_names = set(key._column_label for key in self._groupkeys) + for name in item: + if name in groupkey_names: + raise ValueError( + "cannot insert {}, already exists".format(name_like_string(name)) + ) + return DataFrameGroupBy( + self._psdf, + self._groupkeys, + as_index=self._as_index, + dropna=self._dropna, + column_labels_to_exclude=self._column_labels_to_exclude, + agg_columns=item, + ) + + def _apply_series_op( + self, + op: Callable[["SeriesGroupBy"], Series], + should_resolve: bool = False, + numeric_only: bool = False, + ) -> DataFrame: + applied = [] + for column in self._agg_columns: + applied.append(op(column.groupby(self._groupkeys))) + if numeric_only: + applied = [col for col in applied if isinstance(col.spark.data_type, NumericType)] + if not applied: + raise DataError("No numeric types to aggregate") + internal = self._psdf._internal.with_new_columns(applied, keep_order=False) + if should_resolve: + internal = internal.resolved_copy + return DataFrame(internal) + + def _handle_output(self, psdf: DataFrame) -> DataFrame: + return psdf + + # TODO: Implement 'percentiles', 'include', and 'exclude' arguments. + # TODO: Add ``DataFrame.select_dtypes`` to See Also when 'include' + # and 'exclude' arguments are implemented. + def describe(self) -> DataFrame: + """ + Generate descriptive statistics that summarize the central tendency, + dispersion and shape of a dataset's distribution, excluding + ``NaN`` values. + + Analyzes both numeric and object series, as well + as ``DataFrame`` column sets of mixed data types. The output + will vary depending on what is provided. Refer to the notes + below for more detail. + + .. 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 of the DataFrame provided. + + See Also + -------- + DataFrame.count + DataFrame.max + DataFrame.min + DataFrame.mean + DataFrame.std + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> df + a b c + 0 1 4 7 + 1 1 5 8 + 2 3 6 9 + + Describing a ``DataFrame``. By default only numeric fields + are returned. + + >>> described = df.groupby('a').describe() + >>> described.sort_index() # doctest: +NORMALIZE_WHITESPACE + b c + count mean std min 25% 50% 75% max 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 2.0 7.5 0.707107 7.0 7.0 7.0 8.0 8.0 + 3 1.0 6.0 NaN 6.0 6.0 6.0 6.0 6.0 1.0 9.0 NaN 9.0 9.0 9.0 9.0 9.0 + + """ + for col in self._agg_columns: + if isinstance(col.spark.data_type, StringType): + raise NotImplementedError( + "DataFrameGroupBy.describe() doesn't support for string type for now" + ) + + psdf = self.aggregate(["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) + + # Reorder columns lexicographically by agg column followed by stats. + stats = ["count", "mean", "std", "min"] + formatted_percentiles + ["max"] + column_labels = [tuple(list(label) + [s]) for label, s in product(agg_column_labels, stats)] + data_columns = map(name_like_string, column_labels) + + # Reindex the DataFrame to reflect initial grouping and agg columns. + internal = psdf._internal.copy( + spark_frame=sdf, + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + data_fields=None, + ) + + # Cast columns to ``"float64"`` to match `pandas.DataFrame.groupby`. + return DataFrame(internal).astype("float64") + + +class SeriesGroupBy(GroupBy[Series]): + @staticmethod + def _build( + psser: Series, by: List[Union[Series, Label]], as_index: bool, dropna: bool + ) -> "SeriesGroupBy": + if any( + isinstance(col_or_s, Series) and not same_anchor(psser, col_or_s) for col_or_s in by + ): + psdf, new_by_series, _ = GroupBy._resolve_grouping_from_diff_dataframes( + psser.to_frame(), by + ) + return SeriesGroupBy( + first_series(psdf).rename(psser.name), + new_by_series, + as_index=as_index, + dropna=dropna, + ) + else: + new_by_series = GroupBy._resolve_grouping(psser._psdf, by) + return SeriesGroupBy(psser, new_by_series, as_index=as_index, dropna=dropna) + + def __init__(self, psser: Series, by: List[Series], as_index: bool = True, dropna: bool = True): + if not as_index: + raise TypeError("as_index=False only valid with DataFrame") + super().__init__( + psdf=psser._psdf, + groupkeys=by, + as_index=True, + dropna=dropna, + column_labels_to_exclude=set(), + agg_columns_selected=True, + agg_columns=[psser], + ) + self._psser = psser + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeSeriesGroupBy, item): + property_or_func = getattr(MissingPandasLikeSeriesGroupBy, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + def _apply_series_op( + self, + op: Callable[["SeriesGroupBy"], Series], + should_resolve: bool = False, + numeric_only: bool = False, + ) -> Series: + if numeric_only and not isinstance(self._agg_columns[0].spark.data_type, NumericType): + raise DataError("No numeric types to aggregate") + psser = op(self) + if should_resolve: + internal = psser._internal.resolved_copy + return first_series(DataFrame(internal)) + else: + return psser.copy() + + def _handle_output(self, psdf: DataFrame) -> Series: + return first_series(psdf).rename(self._psser.name) + + def agg(self, *args: Any, **kwargs: Any) -> None: + return MissingPandasLikeSeriesGroupBy.agg(self, *args, **kwargs) + + def aggregate(self, *args: Any, **kwargs: Any) -> None: + return MissingPandasLikeSeriesGroupBy.aggregate(self, *args, **kwargs) + + def size(self) -> Series: + return super().size().rename(self._psser.name) + + size.__doc__ = GroupBy.size.__doc__ + + # TODO: add keep parameter + def nsmallest(self, n: int = 5) -> Series: + """ + Return the smallest `n` elements. + + Parameters + ---------- + n : int + Number of items to retrieve. + + See Also + -------- + pyspark.pandas.Series.nsmallest + pyspark.pandas.DataFrame.nsmallest + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b']) + + >>> df.groupby(['a'])['b'].nsmallest(1).sort_index() # doctest: +NORMALIZE_WHITESPACE + a + 1 0 1 + 2 3 2 + 3 6 3 + Name: b, dtype: int64 + """ + if self._psser._internal.index_level > 1: + raise ValueError("nsmallest do not support multi-index now") + + groupkey_col_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + sdf = self._psser._internal.spark_frame.select( + *[scol.alias(name) for scol, name in zip(self._groupkeys_scols, groupkey_col_names)], + *[ + scol.alias(SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i, scol in enumerate(self._psser._internal.index_spark_columns) + ], + self._psser.spark.column, + NATURAL_ORDER_COLUMN_NAME, + ) + + window = Window.partitionBy(*groupkey_col_names).orderBy( + scol_for(sdf, self._psser._internal.data_spark_column_names[0]).asc(), + NATURAL_ORDER_COLUMN_NAME, + ) + + temp_rank_column = verify_temp_column_name(sdf, "__rank__") + sdf = ( + sdf.withColumn(temp_rank_column, F.row_number().over(window)) + .filter(F.col(temp_rank_column) <= n) + .drop(temp_rank_column) + ).drop(NATURAL_ORDER_COLUMN_NAME) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=( + [scol_for(sdf, col) for col in groupkey_col_names] + + [ + scol_for(sdf, SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i in range(self._psdf._internal.index_level) + ] + ), + index_names=( + [psser._column_label for psser in self._groupkeys] + + self._psdf._internal.index_names + ), + index_fields=( + [ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(self._groupkeys, groupkey_col_names) + ] + + [ + field.copy(name=SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i, field in enumerate(self._psdf._internal.index_fields) + ] + ), + column_labels=[self._psser._column_label], + data_spark_columns=[scol_for(sdf, self._psser._internal.data_spark_column_names[0])], + data_fields=[self._psser._internal.data_fields[0]], + ) + return first_series(DataFrame(internal)) + + # TODO: add keep parameter + def nlargest(self, n: int = 5) -> Series: + """ + Return the first n rows ordered by columns in descending order in group. + + Return the first n rows with the smallest values in columns, in descending order. + The columns that are not specified are returned as well, but not used for ordering. + + Parameters + ---------- + n : int + Number of items to retrieve. + + See Also + -------- + pyspark.pandas.Series.nlargest + pyspark.pandas.DataFrame.nlargest + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b']) + + >>> df.groupby(['a'])['b'].nlargest(1).sort_index() # doctest: +NORMALIZE_WHITESPACE + a + 1 1 2 + 2 4 3 + 3 7 4 + Name: b, dtype: int64 + """ + if self._psser._internal.index_level > 1: + raise ValueError("nlargest do not support multi-index now") + + groupkey_col_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(self._groupkeys))] + sdf = self._psser._internal.spark_frame.select( + *[scol.alias(name) for scol, name in zip(self._groupkeys_scols, groupkey_col_names)], + *[ + scol.alias(SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i, scol in enumerate(self._psser._internal.index_spark_columns) + ], + self._psser.spark.column, + NATURAL_ORDER_COLUMN_NAME, + ) + + window = Window.partitionBy(*groupkey_col_names).orderBy( + scol_for(sdf, self._psser._internal.data_spark_column_names[0]).desc(), + NATURAL_ORDER_COLUMN_NAME, + ) + + temp_rank_column = verify_temp_column_name(sdf, "__rank__") + sdf = ( + sdf.withColumn(temp_rank_column, F.row_number().over(window)) + .filter(F.col(temp_rank_column) <= n) + .drop(temp_rank_column) + ).drop(NATURAL_ORDER_COLUMN_NAME) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=( + [scol_for(sdf, col) for col in groupkey_col_names] + + [ + scol_for(sdf, SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i in range(self._psdf._internal.index_level) + ] + ), + index_names=( + [psser._column_label for psser in self._groupkeys] + + self._psdf._internal.index_names + ), + index_fields=( + [ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(self._groupkeys, groupkey_col_names) + ] + + [ + field.copy(name=SPARK_INDEX_NAME_FORMAT(i + len(self._groupkeys))) + for i, field in enumerate(self._psdf._internal.index_fields) + ] + ), + column_labels=[self._psser._column_label], + data_spark_columns=[scol_for(sdf, self._psser._internal.data_spark_column_names[0])], + data_fields=[self._psser._internal.data_fields[0]], + ) + return first_series(DataFrame(internal)) + + # TODO: add bins, normalize parameter + def value_counts( + self, sort: Optional[bool] = None, ascending: Optional[bool] = None, dropna: bool = True + ) -> Series: + """ + Compute group sizes. + + Parameters + ---------- + sort : boolean, default None + Sort by frequencies. + ascending : boolean, default False + Sort in ascending order. + dropna : boolean, default True + Don't include counts of NaN. + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> df = ps.DataFrame({'A': [1, 2, 2, 3, 3, 3], + ... 'B': [1, 1, 2, 3, 3, np.nan]}, + ... columns=['A', 'B']) + >>> df + A B + 0 1 1.0 + 1 2 1.0 + 2 2 2.0 + 3 3 3.0 + 4 3 3.0 + 5 3 NaN + + >>> df.groupby('A')['B'].value_counts().sort_index() # doctest: +NORMALIZE_WHITESPACE + A B + 1 1.0 1 + 2 1.0 1 + 2.0 1 + 3 3.0 2 + Name: B, dtype: int64 + + Don't include counts of NaN when dropna is False. + + >>> df.groupby('A')['B'].value_counts( + ... dropna=False).sort_index() # doctest: +NORMALIZE_WHITESPACE + A B + 1 1.0 1 + 2 1.0 1 + 2.0 1 + 3 3.0 2 + NaN 1 + Name: B, dtype: int64 + """ + warnings.warn( + "The resulting Series will have a fixed name of 'count' from 4.0.0.", + FutureWarning, + ) + groupkeys = self._groupkeys + self._agg_columns + groupkey_names = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(groupkeys))] + groupkey_cols = [s.spark.column.alias(name) for s, name in zip(groupkeys, groupkey_names)] + + sdf = self._psdf._internal.spark_frame + + agg_column = self._agg_columns[0]._internal.data_spark_column_names[0] + sdf = sdf.groupby(*groupkey_cols).count().withColumnRenamed("count", agg_column) + + if self._dropna: + _groupkey_column_names = groupkey_names[: len(self._groupkeys)] + sdf = sdf.dropna(subset=_groupkey_column_names) + + if dropna: + _agg_columns_names = groupkey_names[len(self._groupkeys) :] + sdf = sdf.dropna(subset=_agg_columns_names) + + if sort: + if ascending: + sdf = sdf.orderBy(scol_for(sdf, agg_column).asc()) + else: + sdf = sdf.orderBy(scol_for(sdf, agg_column).desc()) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in groupkey_names], + index_names=[psser._column_label for psser in groupkeys], + index_fields=[ + psser._internal.data_fields[0].copy(name=name) + for psser, name in zip(groupkeys, groupkey_names) + ], + column_labels=[self._agg_columns[0]._column_label], + data_spark_columns=[scol_for(sdf, agg_column)], + ) + return first_series(DataFrame(internal)) + + def unique(self) -> Series: + """ + Return unique values in group. + + Unique is returned in order of unknown. It does NOT sort. + + See Also + -------- + pyspark.pandas.Series.unique + pyspark.pandas.Index.unique + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 1, 2, 2, 2, 3, 3, 3], + ... 'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b']) + + >>> df.groupby(['a'])['b'].unique().sort_index() # doctest: +SKIP + a + 1 [1, 2] + 2 [2, 3] + 3 [3, 4] + Name: b, dtype: object + """ + return self._reduce_for_stat_function(F.collect_set) + + +def is_multi_agg_with_relabel(**kwargs: Any) -> bool: + """ + Check whether the kwargs pass to .agg look like multi-agg with relabling. + + Parameters + ---------- + **kwargs : dict + + Returns + ------- + bool + + Examples + -------- + >>> is_multi_agg_with_relabel(a='max') + False + >>> is_multi_agg_with_relabel(a_max=('a', 'max'), + ... a_min=('a', 'min')) + True + >>> is_multi_agg_with_relabel() + False + """ + if not kwargs: + return False + return all(isinstance(v, tuple) and len(v) == 2 for v in kwargs.values()) + + +def normalize_keyword_aggregation( + kwargs: Dict[str, Tuple[Name, str]], +) -> Tuple[Dict[Name, List[str]], List[str], List[Tuple]]: + """ + Normalize user-provided kwargs. + + Transforms from the new ``Dict[str, NamedAgg]`` style kwargs + to the old defaultdict[str, List[scalar]]. + + Parameters + ---------- + kwargs : dict + + Returns + ------- + aggspec : dict + The transformed kwargs. + columns : List[str] + The user-provided keys. + order : List[Tuple[str, str]] + Pairs of the input and output column names. + + Examples + -------- + >>> normalize_keyword_aggregation({'output': ('input', 'sum')}) + (defaultdict(, {'input': ['sum']}), ['output'], [('input', 'sum')]) + """ + aggspec: Dict[Union[Any, Tuple], List[str]] = defaultdict(list) + order: List[Tuple] = [] + columns, pairs = zip(*kwargs.items()) + + for column, aggfunc in pairs: + if column in aggspec: + aggspec[column].append(aggfunc) + else: + aggspec[column] = [aggfunc] + + order.append((column, aggfunc)) + # For MultiIndex, we need to flatten the tuple, e.g. (('y', 'A'), 'max') needs to be + # flattened to ('y', 'A', 'max'), it won't do anything on normal Index. + if isinstance(order[0][0], tuple): + order = [(*levs, method) for levs, method in order] + return aggspec, list(columns), order + + +def _test() -> None: + import os + import doctest + import sys + import numpy + from pyspark.sql import SparkSession + import pyspark.pandas.groupby + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.groupby.__dict__.copy() + globs["np"] = numpy + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.groupby tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.groupby, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/__init__.py new file mode 100644 index 0000000..7fde6ff --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/__init__.py @@ -0,0 +1,21 @@ +# +# 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.pandas.indexes.base import Index # noqa: F401 +from pyspark.pandas.indexes.datetimes import DatetimeIndex # noqa: F401 +from pyspark.pandas.indexes.multi import MultiIndex # noqa: F401 +from pyspark.pandas.indexes.numeric import Float64Index, Int64Index # noqa: F401 +from pyspark.pandas.indexes.timedelta import TimedeltaIndex # noqa: F401 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/base.py new file mode 100644 index 0000000..a8fd07a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/base.py @@ -0,0 +1,2783 @@ +# +# 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 functools import partial +from typing import ( + Any, + Callable, + Iterator, + List, + Optional, + Tuple, + Union, + cast, + no_type_check, + TYPE_CHECKING, +) +import warnings + +import pandas as pd +import numpy as np +from pandas.api.types import ( # type: ignore[attr-defined] + is_list_like, + is_interval_dtype, + is_bool_dtype, + is_categorical_dtype, + is_integer_dtype, + is_float_dtype, + is_numeric_dtype, + is_object_dtype, +) +from pandas.core.accessor import CachedAccessor +from pandas.io.formats.printing import pprint_thing +from pandas.api.types import CategoricalDtype, is_hashable # type: ignore[attr-defined] +from pandas._libs import lib + +from pyspark.sql.column import Column +from pyspark.sql import functions as F +from pyspark.sql.types import ( + DayTimeIntervalType, + FractionalType, + IntegralType, + TimestampType, + TimestampNTZType, +) + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import Dtype, Label, Name, Scalar +from pyspark.pandas.config import get_option, option_context +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.missing.indexes import MissingPandasLikeIndex +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.spark.accessors import SparkIndexMethods +from pyspark.pandas.utils import ( + is_name_like_tuple, + is_name_like_value, + name_like_string, + same_anchor, + scol_for, + verify_temp_column_name, + validate_bool_kwarg, + validate_index_loc, + ERROR_MESSAGE_CANNOT_COMBINE, + log_advice, +) +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + DEFAULT_SERIES_NAME, + SPARK_DEFAULT_INDEX_NAME, + SPARK_INDEX_NAME_FORMAT, +) + +if TYPE_CHECKING: + from pyspark.pandas.spark.accessors import SparkIndexOpsMethods + + +class Index(IndexOpsMixin): + """ + pandas-on-Spark Index that corresponds to pandas Index logically. This might hold Spark Column + internally. + + Parameters + ---------- + data : array-like (1-dimensional) + dtype : dtype, default None + If dtype is None, we find the dtype that best fits the data. + If an actual dtype is provided, we coerce to that dtype if it's safe. + Otherwise, an error will be raised. + copy : bool + Make a copy of input ndarray. + name : object + Name to be stored in the index. + tupleize_cols : bool (default: True) + When True, attempt to create a MultiIndex if possible. + + See Also + -------- + MultiIndex : A multi-level, or hierarchical, Index. + DatetimeIndex : Index of datetime64 data. + Int64Index : A special case of :class:`Index` with purely integer labels. + Float64Index : A special case of :class:`Index` with purely float labels. + + Examples + -------- + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=[1, 2, 3]).index # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + >>> ps.DataFrame({'a': [1, 2, 3]}, index=list('abc')).index # doctest: +SKIP + Index(['a', 'b', 'c'], dtype='object') + + >>> ps.Index([1, 2, 3]) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + >>> ps.Index(list('abc')) + Index(['a', 'b', 'c'], dtype='object') + + From a Series: + + >>> s = ps.Series([1, 2, 3], index=[10, 20, 30]) + >>> ps.Index(s) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + From an Index: + + >>> idx = ps.Index([1, 2, 3]) + >>> ps.Index(idx) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + """ + + def __new__( + cls, + data: Optional[Any] = None, + dtype: Optional[Union[str, Dtype]] = None, + copy: bool = False, + name: Optional[Name] = None, + tupleize_cols: bool = True, + **kwargs: Any, + ) -> "Index": + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, Series): + if dtype is not None: + data = data.astype(dtype) + if name is not None: + data = data.rename(name) + + internal = InternalFrame( + spark_frame=data._internal.spark_frame, + index_spark_columns=data._internal.data_spark_columns, + index_names=data._internal.column_labels, + index_fields=data._internal.data_fields, + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return DataFrame(internal).index + elif isinstance(data, Index): + if copy: + data = data.copy() + if dtype is not None: + data = data.astype(dtype) + if name is not None: + data = data.rename(name) + return data + + return cast( + Index, + ps.from_pandas( + pd.Index( + data=data, + dtype=dtype, + copy=copy, + name=name, + tupleize_cols=tupleize_cols, + **kwargs, + ) + ), + ) + + @staticmethod + def _new_instance(anchor: DataFrame) -> "Index": + from pyspark.pandas.indexes.category import CategoricalIndex + from pyspark.pandas.indexes.datetimes import DatetimeIndex + from pyspark.pandas.indexes.multi import MultiIndex + from pyspark.pandas.indexes.numeric import Float64Index, Int64Index + from pyspark.pandas.indexes.timedelta import TimedeltaIndex + + instance: Index + if anchor._internal.index_level > 1: + instance = object.__new__(MultiIndex) + elif isinstance(anchor._internal.index_fields[0].dtype, CategoricalDtype): + instance = object.__new__(CategoricalIndex) + elif isinstance( + anchor._internal.spark_type_for(anchor._internal.index_spark_columns[0]), IntegralType + ): + instance = object.__new__(Int64Index) + elif isinstance( + anchor._internal.spark_type_for(anchor._internal.index_spark_columns[0]), FractionalType + ): + instance = object.__new__(Float64Index) + elif isinstance( + anchor._internal.spark_type_for(anchor._internal.index_spark_columns[0]), + (TimestampType, TimestampNTZType), + ): + instance = object.__new__(DatetimeIndex) + elif isinstance( + anchor._internal.spark_type_for(anchor._internal.index_spark_columns[0]), + DayTimeIntervalType, + ): + instance = object.__new__(TimedeltaIndex) + else: + instance = object.__new__(Index) + + instance._anchor = anchor # type: ignore[attr-defined] + return instance + + @property + def _psdf(self) -> DataFrame: + return self._anchor + + @property + def _internal(self) -> InternalFrame: + internal = self._psdf._internal + return internal.copy( + column_labels=internal.index_names, + data_spark_columns=internal.index_spark_columns, + data_fields=internal.index_fields, + column_label_names=None, + ) + + @property + def _column_label(self) -> Optional[Label]: + return self._psdf._internal.index_names[0] + + def _with_new_scol(self, scol: Column, *, field: Optional[InternalField] = None) -> "Index": + """ + Copy pandas-on-Spark Index with the new Spark Column. + + :param scol: the new Spark Column + :return: the copied Index + """ + internal = self._internal.copy( + index_spark_columns=[scol.alias(SPARK_DEFAULT_INDEX_NAME)], + index_fields=[ + field + if field is None or field.struct_field is None + else field.copy(name=SPARK_DEFAULT_INDEX_NAME) + ], + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return DataFrame(internal).index + + spark: "SparkIndexOpsMethods" = CachedAccessor( # type: ignore[assignment] + "spark", SparkIndexMethods + ) + + # This method is used via `DataFrame.info` API internally. + def _summary(self, name: Optional[str] = None) -> str: + """ + Return a summarized representation. + + Parameters + ---------- + name : str + name to use in the summary representation + + Returns + ------- + String with a summarized representation of the index + """ + head, tail, total_count = tuple( + self._internal.spark_frame.select( + F.first(self.spark.column), F.last(self.spark.column), F.count(F.expr("*")) + ) + .toPandas() + .iloc[0] + ) + + if total_count > 0: + index_summary = ", %s to %s" % (pprint_thing(head), pprint_thing(tail)) + else: + index_summary = "" + + if name is None: + name = type(self).__name__ + return "%s: %s entries%s" % (name, total_count, index_summary) + + @property + def size(self) -> int: + """ + Return an int representing the number of elements in this object. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats'], + ... index=list('abcd')) + >>> df.index.size + 4 + + >>> df.set_index('dogs', append=True).index.size + 4 + """ + return len(self) + + @property + def shape(self) -> tuple: + """ + Return a tuple of the shape of the underlying data. + + Examples + -------- + >>> idx = ps.Index(['a', 'b', 'c']) + >>> idx + Index(['a', 'b', 'c'], dtype='object') + >>> idx.shape + (3,) + + >>> midx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> midx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z')], + ) + >>> midx.shape + (3,) + """ + return (len(self._psdf),) + + def identical(self, other: "Index") -> bool: + """ + Similar to equals, but check that other comparable attributes are + also equal. + + Returns + ------- + bool + If two Index objects have equal elements and same type True, + otherwise False. + + Examples + -------- + + >>> from pyspark.pandas.config import option_context + >>> idx = ps.Index(['a', 'b', 'c']) + >>> midx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + + For Index + + >>> idx.identical(idx) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... idx.identical(ps.Index(['a', 'b', 'c'])) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... idx.identical(ps.Index(['b', 'b', 'a'])) + False + >>> idx.identical(midx) + False + + For MultiIndex + + >>> midx.identical(midx) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... midx.identical(ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')])) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... midx.identical(ps.MultiIndex.from_tuples([('c', 'z'), ('b', 'y'), ('a', 'x')])) + False + >>> midx.identical(idx) + False + """ + from pyspark.pandas.indexes.multi import MultiIndex + + self_name = self.names if isinstance(self, MultiIndex) else self.name + other_name = other.names if isinstance(other, MultiIndex) else other.name + + return ( + self_name == other_name # to support non-index comparison by short-circuiting. + and self.equals(other) + ) + + def equals(self, other: "Index") -> bool: + """ + Determine if two Index objects contain the same elements. + + Returns + ------- + bool + True if "other" is an Index and it has the same elements as calling + index; False otherwise. + + Examples + -------- + + >>> from pyspark.pandas.config import option_context + >>> idx = ps.Index(['a', 'b', 'c']) + >>> idx.name = "name" + >>> midx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> midx.names = ("nameA", "nameB") + + For Index + + >>> idx.equals(idx) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... idx.equals(ps.Index(['a', 'b', 'c'])) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... idx.equals(ps.Index(['b', 'b', 'a'])) + False + >>> idx.equals(midx) + False + + For MultiIndex + + >>> midx.equals(midx) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... midx.equals(ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')])) + True + >>> with option_context('compute.ops_on_diff_frames', True): + ... midx.equals(ps.MultiIndex.from_tuples([('c', 'z'), ('b', 'y'), ('a', 'x')])) + False + >>> midx.equals(idx) + False + """ + if same_anchor(self, other): + return True + elif type(self) == type(other): + if get_option("compute.ops_on_diff_frames"): + # TODO: avoid using default index? + with option_context("compute.default_index_type", "distributed-sequence"): + # Directly using Series from both self and other seems causing + # some exceptions when 'compute.ops_on_diff_frames' is enabled. + # Working around for now via using frames. + return ( + cast(Series, self.to_series("self").reset_index(drop=True)) + == cast(Series, other.to_series("other").reset_index(drop=True)) + ).all() + else: + raise ValueError(ERROR_MESSAGE_CANNOT_COMBINE) + else: + return False + + def transpose(self) -> "Index": + """ + Return the transpose, For index, It will be index itself. + + Examples + -------- + >>> idx = ps.Index(['a', 'b', 'c']) + >>> idx + Index(['a', 'b', 'c'], dtype='object') + + >>> idx.transpose() + Index(['a', 'b', 'c'], dtype='object') + + For MultiIndex + + >>> midx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> midx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z')], + ) + + >>> midx.transpose() # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z')], + ) + """ + return self + + T = property(transpose) + + def _to_internal_pandas(self) -> pd.Index: + """ + Return a pandas Index directly from _internal to avoid overhead of copy. + + This method is for internal use only. + """ + return self._psdf._internal.to_pandas_frame.index + + def to_pandas(self) -> pd.Index: + """ + Return a pandas Index. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats'], + ... index=list('abcd')) + >>> df['dogs'].index.to_pandas() + Index(['a', 'b', 'c', 'd'], dtype='object') + """ + log_advice( + "`to_pandas` loads all data into the driver's memory. " + "It should only be used if the resulting pandas Index is expected to be small." + ) + return self._to_pandas() + + def _to_pandas(self) -> pd.Index: + """ + Same as `to_pandas()`, without issuing the advice log for internal usage. + """ + return self._to_internal_pandas().copy() + + def to_numpy(self, dtype: Optional[Union[str, Dtype]] = None, copy: bool = False) -> np.ndarray: + """ + A NumPy ndarray representing the values in this Index or MultiIndex. + + .. note:: This method should only be used if the resulting NumPy ndarray is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + dtype : str or numpy.dtype, optional + The dtype to pass to :meth:`numpy.asarray` + copy : bool, default False + Whether to ensure that the returned value is not a view on + another array. Note that ``copy=False`` does not *ensure* that + ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensures that + a copy is made, even if not strictly necessary. + + Returns + ------- + numpy.ndarray + + Examples + -------- + >>> ps.Series([1, 2, 3, 4]).index.to_numpy() + array([0, 1, 2, 3]) + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=[[1, 2, 3], [4, 5, 6]]).index.to_numpy() + array([(1, 4), (2, 5), (3, 6)], dtype=object) + """ + log_advice( + "`to_numpy` loads all data into the driver's memory. " + "It should only be used if the resulting NumPy ndarray is expected to be small." + ) + result = np.asarray( + self._to_internal_pandas()._values, dtype=dtype # type: ignore[arg-type,attr-defined] + ) + if copy: + result = result.copy() + return result + + def map( + self, mapper: Union[dict, Callable[[Any], Any], pd.Series], na_action: Optional[str] = None + ) -> "Index": + """ + Map values using input correspondence (a dict, Series, or function). + + Parameters + ---------- + mapper : function, dict, or pd.Series + Mapping correspondence. + na_action : {None, 'ignore'} + If ‘ignore’, propagate NA values, without passing them to the mapping correspondence. + + Returns + ------- + applied : Index, inferred + The output of the mapping function applied to the index. + + Examples + -------- + >>> psidx = ps.Index([1, 2, 3]) + + >>> psidx.map({1: "one", 2: "two", 3: "three"}) + Index(['one', 'two', 'three'], dtype='object') + + >>> psidx.map(lambda id: "{id} + 1".format(id=id)) + Index(['1 + 1', '2 + 1', '3 + 1'], dtype='object') + + >>> pser = pd.Series(["one", "two", "three"], index=[1, 2, 3]) + >>> psidx.map(pser) + Index(['one', 'two', 'three'], dtype='object') + """ + if isinstance(mapper, dict): + if len(set(type(k) for k in mapper.values())) > 1: + raise TypeError( + "If the mapper is a dictionary, its values must be of the same type" + ) + + return Index( + self.to_series().pandas_on_spark.transform_batch( + lambda pser: pser.map(mapper, na_action) + ) + ).rename(self.name) + + @property + def values(self) -> np.ndarray: + """ + Return an array representing the data in the Index. + + .. warning:: We recommend using `Index.to_numpy()` instead. + + .. note:: This method should only be used if the resulting NumPy ndarray is expected + to be small, as all the data is loaded into the driver's memory. + + Returns + ------- + numpy.ndarray + + Examples + -------- + >>> ps.Series([1, 2, 3, 4]).index.values + array([0, 1, 2, 3]) + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=[[1, 2, 3], [4, 5, 6]]).index.values + array([(1, 4), (2, 5), (3, 6)], dtype=object) + """ + warnings.warn("We recommend using `{}.to_numpy()` instead.".format(type(self).__name__)) + return self.to_numpy() + + @property + def asi8(self) -> np.ndarray: + """ + Integer representation of the values. + + .. warning:: We recommend using `Index.to_numpy()` instead. + + .. note:: This method should only be used if the resulting NumPy ndarray is expected + to be small, as all the data is loaded into the driver's memory. + + .. deprecated:: 3.4.0 + + Returns + ------- + numpy.ndarray + An ndarray with int64 dtype. + + Examples + -------- + >>> ps.Index([1, 2, 3]).asi8 + array([1, 2, 3]) + + Returns None for non-int64 dtype + + >>> ps.Index(['a', 'b', 'c']).asi8 is None + True + """ + warnings.warn( + "Index.asi8 is deprecated and will be removed in 4.0.0. " "Use Index.astype instead.", + FutureWarning, + ) + if isinstance(self.spark.data_type, IntegralType): + return self.to_numpy() + elif isinstance(self.spark.data_type, (TimestampType, TimestampNTZType)): + return np.array(list(map(lambda x: x.astype(np.int64), self.to_numpy()))) + else: + return None + + @property + def has_duplicates(self) -> bool: + """ + If index has duplicates, return True, otherwise False. + + Examples + -------- + >>> idx = ps.Index([1, 5, 7, 7]) + >>> idx.has_duplicates + True + + >>> idx = ps.Index([1, 5, 7]) + >>> idx.has_duplicates + False + + >>> idx = ps.Index(["Watermelon", "Orange", "Apple", + ... "Watermelon"]) + >>> idx.has_duplicates + True + + >>> idx = ps.Index(["Orange", "Apple", + ... "Watermelon"]) + >>> idx.has_duplicates + False + """ + sdf = self._internal.spark_frame.select(self.spark.column) + scol = scol_for(sdf, sdf.columns[0]) + + return sdf.select(F.count(scol) != F.countDistinct(scol)).first()[0] + + @property + def is_unique(self) -> bool: + """ + Return if the index has unique values. + + Examples + -------- + >>> idx = ps.Index([1, 5, 7, 7]) + >>> idx.is_unique + False + + >>> idx = ps.Index([1, 5, 7]) + >>> idx.is_unique + True + + >>> idx = ps.Index(["Watermelon", "Orange", "Apple", + ... "Watermelon"]) + >>> idx.is_unique + False + + >>> idx = ps.Index(["Orange", "Apple", + ... "Watermelon"]) + >>> idx.is_unique + True + """ + return not self.has_duplicates + + @property + def name(self) -> Name: + """Return name of the Index.""" + return self.names[0] + + @name.setter + def name(self, name: Name) -> None: + self.names = [name] + + @property + def names(self) -> List[Name]: + """Return names of the Index.""" + return [ + name if name is None or len(name) > 1 else name[0] + for name in self._internal.index_names + ] + + @names.setter + def names(self, names: List[Name]) -> None: + if not is_list_like(names): + raise ValueError("Names must be a list-like") + if self._internal.index_level != len(names): + raise ValueError( + "Length of new names must be {}, got {}".format( + self._internal.index_level, len(names) + ) + ) + if self._internal.index_level == 1: + self.rename(names[0], inplace=True) + else: + self.rename(names, inplace=True) + + @property + def nlevels(self) -> int: + """ + Number of levels in Index & MultiIndex. + + Examples + -------- + >>> psdf = ps.DataFrame({"a": [1, 2, 3]}, index=pd.Index(['a', 'b', 'c'], name="idx")) + >>> psdf.index.nlevels + 1 + + >>> psdf = ps.DataFrame({'a': [1, 2, 3]}, index=[list('abc'), list('def')]) + >>> psdf.index.nlevels + 2 + """ + return self._internal.index_level + + def rename(self, name: Union[Name, List[Name]], inplace: bool = False) -> Optional["Index"]: + """ + Alter Index or MultiIndex name. + Able to set new names without level. Defaults to returning a new index. + + Parameters + ---------- + name : label or list of labels + Name(s) to set. + inplace : boolean, default False + Modifies the object directly, instead of creating a new Index or MultiIndex. + + Returns + ------- + Index or MultiIndex + The same type as the caller or None if inplace is True. + + Examples + -------- + >>> df = ps.DataFrame({'a': ['A', 'C'], 'b': ['A', 'B']}, columns=['a', 'b']) + >>> df.index.rename("c") # doctest: +SKIP + Int64Index([0, 1], dtype='int64', name='c') + + >>> df.set_index("a", inplace=True) + >>> df.index.rename("d") + Index(['A', 'C'], dtype='object', name='d') + + You can also change the index name in place. + + >>> df.index.rename("e", inplace=True) + >>> df.index + Index(['A', 'C'], dtype='object', name='e') + + >>> df # doctest: +NORMALIZE_WHITESPACE + b + e + A A + C B + + Support for MultiIndex + + >>> psidx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y')]) + >>> psidx.names = ['hello', 'pandas-on-Spark'] + >>> psidx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + names=['hello', 'pandas-on-Spark']) + + >>> psidx.rename(['aloha', 'databricks']) # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + names=['aloha', 'databricks']) + """ + names = self._verify_for_rename(name) + + internal = self._psdf._internal.copy(index_names=names) + + if inplace: + self._psdf._update_internal_frame(internal) + return None + else: + return DataFrame(internal).index + + def _verify_for_rename(self, name: Name) -> List[Label]: + if is_hashable(name): + if is_name_like_tuple(name): + return [name] + elif is_name_like_value(name): + return [(name,)] + raise TypeError("Index.name must be a hashable type") + + # TODO: add downcast parameter for fillna function + def fillna(self, value: Scalar) -> "Index": + """ + Fill NA/NaN values with the specified value. + + Parameters + ---------- + value : scalar + Scalar value to use to fill holes (example: 0). This value cannot be a list-likes. + + Returns + ------- + Index : + filled with value + + Examples + -------- + >>> idx = ps.Index([1, 2, None]) + >>> idx # doctest: +SKIP + Float64Index([1.0, 2.0, nan], dtype='float64') + + >>> idx.fillna(0) # doctest: +SKIP + Float64Index([1.0, 2.0, 0.0], dtype='float64') + """ + if not isinstance(value, (float, int, str, bool)): + raise TypeError("Unsupported type %s" % type(value).__name__) + sdf = self._internal.spark_frame.fillna(value) + + internal = InternalFrame( # TODO: dtypes? + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + ) + return DataFrame(internal).index + + def drop_duplicates(self, keep: Union[bool, str] = "first") -> "Index": + """ + Return Index with duplicate values removed. + + Parameters + ---------- + keep : {'first', 'last', ``False``}, default 'first' + Method to handle dropping duplicates: + - 'first' : Drop duplicates except for the first occurrence. + - 'last' : Drop duplicates except for the last occurrence. + - ``False`` : Drop all duplicates. + + Returns + ------- + deduplicated : Index + + See Also + -------- + Series.drop_duplicates : Equivalent method on Series. + DataFrame.drop_duplicates : Equivalent method on DataFrame. + + Examples + -------- + Generate an Index with duplicate values. + + >>> idx = ps.Index(['lama', 'cow', 'lama', 'beetle', 'lama', 'hippo']) + + >>> idx.drop_duplicates().sort_values() + Index(['beetle', 'cow', 'hippo', 'lama'], dtype='object') + """ + with ps.option_context("compute.default_index_type", "distributed"): + # The attached index caused by `reset_index` below is used for sorting only, + # and it will be dropped soon, + # so we enforce “distributed” default index type + psser = self.to_series().reset_index(drop=True) + return Index(psser.drop_duplicates(keep=keep).sort_index()) + + def to_series(self, name: Optional[Name] = None) -> Series: + """ + Create a Series with both index and values equal to the index keys + useful with map for returning an indexer based on an index. + + Parameters + ---------- + name : string, optional + name of resulting Series. If None, defaults to name of original + index + + Returns + ------- + Series : dtype will be based on the type of the Index values. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats'], + ... index=list('abcd')) + >>> df['dogs'].index.to_series() + a a + b b + c c + d d + dtype: object + """ + if not is_hashable(name): + raise TypeError("Series.name must be a hashable type") + scol = self.spark.column + field = self._internal.data_fields[0] + if name is not None: + scol = scol.alias(name_like_string(name)) + field = field.copy(name=name_like_string(name)) + elif self._internal.index_level == 1: + name = self.name + column_labels: List[Optional[Label]] = [name if is_name_like_tuple(name) else (name,)] + internal = self._internal.copy( + column_labels=column_labels, + data_spark_columns=[scol], + data_fields=[field], + column_label_names=None, + ) + return first_series(DataFrame(internal)) + + def to_frame(self, index: bool = True, name: Optional[Name] = None) -> DataFrame: + """ + Create a DataFrame with a column containing the Index. + + Parameters + ---------- + index : boolean, default True + Set the index of the returned DataFrame as the original Index. + name : object, default None + The passed name should substitute for the index name (if it has + one). + + Returns + ------- + DataFrame + DataFrame containing the original Index data. + + See Also + -------- + Index.to_series : Convert an Index to a Series. + Series.to_frame : Convert Series to DataFrame. + + Examples + -------- + >>> idx = ps.Index(['Ant', 'Bear', 'Cow'], name='animal') + >>> idx.to_frame() # doctest: +NORMALIZE_WHITESPACE + animal + animal + Ant Ant + Bear Bear + Cow Cow + + By default, the original Index is reused. To enforce a new Index: + + >>> idx.to_frame(index=False) + animal + 0 Ant + 1 Bear + 2 Cow + + To override the name of the resulting column, specify `name`: + + >>> idx.to_frame(name='zoo') # doctest: +NORMALIZE_WHITESPACE + zoo + animal + Ant Ant + Bear Bear + Cow Cow + """ + if name is None: + if self._internal.index_names[0] is None: + name = (DEFAULT_SERIES_NAME,) + else: + name = self._internal.index_names[0] + elif not is_name_like_tuple(name): + if is_name_like_value(name): + name = (name,) + else: + raise TypeError("unhashable type: '{}'".format(type(name).__name__)) + + return self._to_frame(index=index, names=[name]) + + def _to_frame(self, index: bool, names: List[Label]) -> DataFrame: + if index: + index_spark_columns = self._internal.index_spark_columns + index_names = self._internal.index_names + index_fields = self._internal.index_fields + else: + index_spark_columns = [] + index_names = [] + index_fields = [] + + internal = InternalFrame( + spark_frame=self._internal.spark_frame, + index_spark_columns=index_spark_columns, + index_names=index_names, + index_fields=index_fields, + column_labels=names, + data_spark_columns=self._internal.index_spark_columns, + data_fields=self._internal.index_fields, + ) + return DataFrame(internal) + + def is_boolean(self) -> bool: + """ + Return if the current index type is a boolean type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[True]).index.is_boolean() + True + """ + return is_bool_dtype(self.dtype) + + def is_categorical(self) -> bool: + """ + Return if the current index type is a categorical type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[1]).index.is_categorical() + False + """ + return is_categorical_dtype(self.dtype) + + def is_floating(self) -> bool: + """ + Return if the current index type is a floating type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[1]).index.is_floating() + False + """ + return is_float_dtype(self.dtype) + + def is_integer(self) -> bool: + """ + Return if the current index type is an integer type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[1]).index.is_integer() + True + """ + return is_integer_dtype(self.dtype) + + def is_interval(self) -> bool: + """ + Return if the current index type is an interval type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[1]).index.is_interval() + False + """ + return is_interval_dtype(self.dtype) + + def is_numeric(self) -> bool: + """ + Return if the current index type is a numeric type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=[1]).index.is_numeric() + True + """ + return is_numeric_dtype(self.dtype) + + def is_object(self) -> bool: + """ + Return if the current index type is an object type. + + Examples + -------- + >>> ps.DataFrame({'a': [1]}, index=["a"]).index.is_object() + True + """ + return is_object_dtype(self.dtype) + + def is_type_compatible(self, kind: str) -> bool: + """ + Whether the index type is compatible with the provided type. + + .. deprecated:: 3.4.0 + + Examples + -------- + >>> psidx = ps.Index([1, 2, 3]) + >>> psidx.is_type_compatible('integer') + True + + >>> psidx = ps.Index([1.0, 2.0, 3.0]) + >>> psidx.is_type_compatible('integer') + False + >>> psidx.is_type_compatible('floating') + True + """ + warnings.warn( + "Index.is_type_compatible is deprecated and will be removed in 4.0.0. " + "Use Index.isin instead.", + FutureWarning, + ) + return kind == self.inferred_type + + def dropna(self, how: str = "any") -> "Index": + """ + Return Index or MultiIndex without NA/NaN values + + Parameters + ---------- + how : {'any', 'all'}, default 'any' + If the Index is a MultiIndex, drop the value when any or all levels + are NaN. + + Returns + ------- + Index or MultiIndex + + Examples + -------- + + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=['cobra', 'viper', None], + ... columns=['max_speed', 'shield']) + >>> df # doctest: +SKIP + max_speed shield + cobra 1 2 + viper 4 5 + None 7 8 + + >>> df.index.dropna() + Index(['cobra', 'viper'], dtype='object') + + Also support for MultiIndex + + + >>> tuples = [(np.nan, 1.0), (2.0, 2.0), (np.nan, np.nan), (3.0, np.nan)] + >>> midx = ps.MultiIndex.from_tuples(tuples) + >>> midx # doctest: +SKIP + MultiIndex([(nan, 1.0), + (2.0, 2.0), + (nan, nan), + (3.0, nan)], + ) + + >>> midx.dropna() # doctest: +SKIP + MultiIndex([(2.0, 2.0)], + ) + + >>> midx.dropna(how="all") # doctest: +SKIP + MultiIndex([(nan, 1.0), + (2.0, 2.0), + (3.0, nan)], + ) + """ + if how not in ("any", "all"): + raise ValueError("invalid how option: %s" % how) + + sdf = self._internal.spark_frame.select(self._internal.index_spark_columns).dropna(how=how) + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + return DataFrame(internal).index + + def unique(self, level: Optional[Union[int, Name]] = None) -> "Index": + """ + Return unique values in the index. + + Be aware the order of unique values might be different than pandas.Index.unique + + Parameters + ---------- + level : int or str, optional, default is None + + Returns + ------- + Index without duplicates + + See Also + -------- + Series.unique + groupby.SeriesGroupBy.unique + + Examples + -------- + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=[1, 1, 3]).index.unique().sort_values() + ... # doctest: +SKIP + Int64Index([1, 3], dtype='int64') + + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=['d', 'e', 'e']).index.unique().sort_values() + Index(['d', 'e'], dtype='object') + + MultiIndex + + >>> ps.MultiIndex.from_tuples([("A", "X"), ("A", "Y"), ("A", "X")]).unique() + ... # doctest: +SKIP + MultiIndex([('A', 'X'), + ('A', 'Y')], + ) + """ + if level is not None: + self._validate_index_level(level) + scols = self._internal.index_spark_columns + sdf = self._psdf._internal.spark_frame.select(scols).distinct() + return DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + ).index + + # TODO: add error parameter + def drop(self, labels: List[Any]) -> "Index": + """ + Make new Index with passed list of labels deleted. + + Parameters + ---------- + labels : array-like + + Returns + ------- + dropped : Index + + Examples + -------- + >>> index = ps.Index([1, 2, 3]) + >>> index # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + >>> index.drop([1]) # doctest: +SKIP + Int64Index([2, 3], dtype='int64') + """ + internal = self._internal.resolved_copy + sdf = internal.spark_frame[~internal.index_spark_columns[0].isin(labels)] + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return DataFrame(internal).index + + def _validate_index_level(self, level: Union[int, Name]) -> None: + """ + Validate index level. + For single-level Index getting level number is a no-op, but some + verification must be done like in MultiIndex. + """ + if isinstance(level, int): + if level < 0 and level != -1: + raise IndexError( + "Too many levels: Index has only 1 level," + " %d is not a valid level number" % (level,) + ) + elif level > 0: + raise IndexError("Too many levels:" " Index has only 1 level, not %d" % (level + 1)) + elif level != self.name: + raise KeyError( + "Requested level ({}) does not match index name ({})".format(level, self.name) + ) + + def get_level_values(self, level: Union[int, Name]) -> "Index": + """ + Return Index if a valid level is given. + + Examples + -------- + >>> psidx = ps.Index(['a', 'b', 'c'], name='ks') + >>> psidx.get_level_values(0) + Index(['a', 'b', 'c'], dtype='object', name='ks') + + >>> psidx.get_level_values('ks') + Index(['a', 'b', 'c'], dtype='object', name='ks') + """ + self._validate_index_level(level) + return self + + def copy(self, name: Optional[Name] = None, deep: Optional[bool] = None) -> "Index": + """ + Make a copy of this object. name sets those attributes on the new object. + + Parameters + ---------- + name : string, optional + to set name of index + deep : None + this parameter is not supported but just dummy parameter to match pandas. + + Examples + -------- + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=['cobra', 'viper', 'sidewinder'], + ... columns=['max_speed', 'shield']) + >>> df + max_speed shield + cobra 1 2 + viper 4 5 + sidewinder 7 8 + >>> df.index + Index(['cobra', 'viper', 'sidewinder'], dtype='object') + + Copy index + + >>> df.index.copy() + Index(['cobra', 'viper', 'sidewinder'], dtype='object') + + Copy index with name + + >>> df.index.copy(name='snake') + Index(['cobra', 'viper', 'sidewinder'], dtype='object', name='snake') + """ + result = self._psdf[[]].index + if name: + result.name = name + return result + + def droplevel(self, level: Union[int, Name, List[Union[int, Name]]]) -> "Index": + """ + Return index with requested level(s) removed. + If resulting index has only 1 level left, the result will be + of Index type, not MultiIndex. + + Parameters + ---------- + level : int, str, tuple, or list-like, default 0 + If a string is given, must be the name of a level + If list-like, elements must be names or indexes of levels. + + Returns + ------- + Index or MultiIndex + + Examples + -------- + >>> midx = ps.DataFrame({'a': ['a', 'b']}, index=[['a', 'x'], ['b', 'y'], [1, 2]]).index + >>> midx # doctest: +SKIP + MultiIndex([('a', 'b', 1), + ('x', 'y', 2)], + ) + >>> midx.droplevel([0, 1]) # doctest: +SKIP + Int64Index([1, 2], dtype='int64') + >>> midx.droplevel(0) # doctest: +SKIP + MultiIndex([('b', 1), + ('y', 2)], + ) + >>> midx.names = [("a", "b"), "b", "c"] + >>> midx.droplevel([('a', 'b')]) # doctest: +SKIP + MultiIndex([('b', 1), + ('y', 2)], + names=['b', 'c']) + """ + names = self.names + nlevels = self.nlevels + if not is_list_like(level): + levels = [cast(Union[int, Name], level)] + else: + levels = cast(List[Union[int, Name]], level) + + int_level = set() + for n in levels: + if isinstance(n, int): + if n < 0: + n = n + nlevels + if n < 0: + raise IndexError( + "Too many levels: Index has only {} levels, " + "{} is not a valid level number".format(nlevels, (n - nlevels)) + ) + if n >= nlevels: + raise IndexError( + "Too many levels: Index has only {} levels, not {}".format(nlevels, n + 1) + ) + else: + if n not in names: + raise KeyError("Level {} not found".format(n)) + n = names.index(n) + int_level.add(n) + + if len(levels) >= nlevels: + raise ValueError( + "Cannot remove {} levels from an index with {} " + "levels: at least one level must be " + "left.".format(len(levels), nlevels) + ) + + index_spark_columns, index_names, index_fields = zip( + *[ + item + for i, item in enumerate( + zip( + self._internal.index_spark_columns, + self._internal.index_names, + self._internal.index_fields, + ) + ) + if i not in int_level + ] + ) + + internal = self._internal.copy( + index_spark_columns=list(index_spark_columns), + index_names=list(index_names), + index_fields=list(index_fields), + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return DataFrame(internal).index + + def symmetric_difference( + self, + other: "Index", + result_name: Optional[Name] = None, + sort: Optional[bool] = None, + ) -> "Index": + """ + Compute the symmetric difference of two Index objects. + + Parameters + ---------- + other : Index or array-like + result_name : str + sort : True or None, default None + Whether to sort the resulting index. + * True : Attempt to sort the result. + * None : Do not sort the result. + + Returns + ------- + symmetric_difference : Index + + Notes + ----- + ``symmetric_difference`` contains elements that appear in either + ``idx1`` or ``idx2`` but not both. Equivalent to the Index created by + ``idx1.difference(idx2) | idx2.difference(idx1)`` with duplicates + dropped. + + Examples + -------- + >>> s1 = ps.Series([1, 2, 3, 4], index=[1, 2, 3, 4]) + >>> s2 = ps.Series([1, 2, 3, 4], index=[2, 3, 4, 5]) + + >>> s1.index.symmetric_difference(s2.index) # doctest: +SKIP + Int64Index([5, 1], dtype='int64') + + You can set name of result Index. + + >>> s1.index.symmetric_difference(s2.index, result_name='pandas-on-Spark') # doctest: +SKIP + Int64Index([5, 1], dtype='int64', name='pandas-on-Spark') + + You can set sort to `True`, if you want to sort the resulting index. + + >>> s1.index.symmetric_difference(s2.index, sort=True) # doctest: +SKIP + Int64Index([1, 5], dtype='int64') + + You can also use the ``^`` operator: + + >>> s1.index ^ s2.index # doctest: +SKIP + Int64Index([5, 1], dtype='int64') + """ + if type(self) != type(other): + raise NotImplementedError( + "Doesn't support symmetric_difference between Index & MultiIndex for now" + ) + + sdf_self = self._psdf._internal.spark_frame.select(self._internal.index_spark_columns) + sdf_other = other._psdf._internal.spark_frame.select(other._internal.index_spark_columns) + + sdf_symdiff = sdf_self.union(sdf_other).subtract(sdf_self.intersect(sdf_other)) + + if sort: + sdf_symdiff = sdf_symdiff.sort(*self._internal.index_spark_column_names) + + internal = InternalFrame( + spark_frame=sdf_symdiff, + index_spark_columns=[ + scol_for(sdf_symdiff, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + result = DataFrame(internal).index + + if result_name: + result.name = result_name + + return result + + def sort_values( + self, return_indexer: bool = False, ascending: bool = True + ) -> Union["Index", Tuple["Index", "Index"]]: + """ + Return a sorted copy of the index, and optionally return the indices that + sorted the index itself. + + .. note:: This method is not supported for pandas when index has NaN value. + pandas raises unexpected TypeError, but we support treating NaN + as the smallest value. + This method returns indexer as a pandas-on-Spark index while + pandas returns it as a list. That's because indexer in pandas-on-Spark + may not fit in memory. + + Parameters + ---------- + return_indexer : bool, default False + Should the indices that would sort the index be returned. + ascending : bool, default True + Should the index values be sorted in an ascending order. + + Returns + ------- + sorted_index : ps.Index or ps.MultiIndex + Sorted copy of the index. + indexer : ps.Index + The indices that the index itself was sorted by. + + See Also + -------- + Series.sort_values : Sort values of a Series. + DataFrame.sort_values : Sort values in a DataFrame. + + Examples + -------- + >>> idx = ps.Index([10, 100, 1, 1000]) + >>> idx # doctest: +SKIP + Int64Index([10, 100, 1, 1000], dtype='int64') + + Sort values in ascending order (default behavior). + + >>> idx.sort_values() # doctest: +SKIP + Int64Index([1, 10, 100, 1000], dtype='int64') + + Sort values in descending order. + + >>> idx.sort_values(ascending=False) # doctest: +SKIP + Int64Index([1000, 100, 10, 1], dtype='int64') + + Sort values in descending order, and also get the indices idx was sorted by. + + >>> idx.sort_values(ascending=False, return_indexer=True) # doctest: +SKIP + (Int64Index([1000, 100, 10, 1], dtype='int64'), Int64Index([3, 1, 0, 2], dtype='int64')) + + Support for MultiIndex. + + >>> psidx = ps.MultiIndex.from_tuples([('a', 'x', 1), ('c', 'y', 2), ('b', 'z', 3)]) + >>> psidx # doctest: +SKIP + MultiIndex([('a', 'x', 1), + ('c', 'y', 2), + ('b', 'z', 3)], + ) + + >>> psidx.sort_values() # doctest: +SKIP + MultiIndex([('a', 'x', 1), + ('b', 'z', 3), + ('c', 'y', 2)], + ) + + >>> psidx.sort_values(ascending=False) # doctest: +SKIP + MultiIndex([('c', 'y', 2), + ('b', 'z', 3), + ('a', 'x', 1)], + ) + + >>> psidx.sort_values(ascending=False, return_indexer=True) # doctest: +SKIP + (MultiIndex([('c', 'y', 2), + ('b', 'z', 3), + ('a', 'x', 1)], + ), Int64Index([1, 2, 0], dtype='int64')) + """ + sdf = self._internal.spark_frame + if return_indexer: + sequence_col = verify_temp_column_name(sdf, "__distributed_sequence_column__") + sdf = InternalFrame.attach_distributed_sequence_column(sdf, column_name=sequence_col) + + ordered_sdf = sdf.orderBy(*self._internal.index_spark_columns, ascending=ascending) + sdf = ordered_sdf.select(self._internal.index_spark_columns) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + sorted_index = DataFrame(internal).index + + if return_indexer: + alias_sequence_scol = scol_for(ordered_sdf, sequence_col).alias( + SPARK_DEFAULT_INDEX_NAME + ) + indexer_sdf = ordered_sdf.select(alias_sequence_scol) + indexer_internal = InternalFrame( + spark_frame=indexer_sdf, + index_spark_columns=[scol_for(indexer_sdf, SPARK_DEFAULT_INDEX_NAME)], + ) + indexer = DataFrame(indexer_internal).index + return sorted_index, indexer + else: + return sorted_index + + @no_type_check + def sort(self, *args, **kwargs) -> None: + """ + Use sort_values instead. + """ + raise TypeError("cannot sort an Index object in-place, use sort_values instead") + + def min(self) -> Union[Scalar, Tuple[Scalar, ...]]: + """ + Return the minimum value of the Index. + + Returns + ------- + scalar + Minimum value. + + See Also + -------- + Index.max : Return the maximum value of the object. + Series.min : Return the minimum value in a Series. + DataFrame.min : Return the minimum values in a DataFrame. + + Examples + -------- + >>> idx = ps.Index([3, 2, 1]) + >>> idx.min() + 1 + + >>> idx = ps.Index(['c', 'b', 'a']) + >>> idx.min() + 'a' + + For a MultiIndex, the maximum is determined lexicographically. + + >>> idx = ps.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2)]) + >>> idx.min() + ('a', 'x', 1) + """ + sdf = self._internal.spark_frame + min_row = ( + sdf.select(F.min(F.struct(*self._internal.index_spark_columns)).alias("min_row")) + .select("min_row.*") + .toPandas() + ) + result = tuple(min_row.iloc[0]) + + return result if len(result) > 1 else result[0] + + def max(self) -> Union[Scalar, Tuple[Scalar, ...]]: + """ + Return the maximum value of the Index. + + Returns + ------- + scalar + Maximum value. + + See Also + -------- + Index.min : Return the minimum value in an Index. + Series.max : Return the maximum value in a Series. + DataFrame.max : Return the maximum values in a DataFrame. + + Examples + -------- + >>> idx = ps.Index([3, 2, 1]) + >>> idx.max() + 3 + + >>> idx = ps.Index(['c', 'b', 'a']) + >>> idx.max() + 'c' + + For a MultiIndex, the maximum is determined lexicographically. + + >>> idx = ps.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2)]) + >>> idx.max() + ('b', 'y', 2) + """ + sdf = self._internal.spark_frame + max_row = ( + sdf.select(F.max(F.struct(*self._internal.index_spark_columns)).alias("max_row")) + .select("max_row.*") + .toPandas() + ) + result = tuple(max_row.iloc[0]) + + return result if len(result) > 1 else result[0] + + def delete(self, loc: Union[int, List[int]]) -> "Index": + """ + Make new Index with passed location(-s) deleted. + + .. note:: this API can be pretty expensive since it is based on + a global sequence internally. + + Returns + ------- + new_index : Index + + Examples + -------- + >>> psidx = ps.Index([10, 10, 9, 8, 4, 2, 4, 4, 2, 2, 10, 10]) + >>> psidx # doctest: +SKIP + Int64Index([10, 10, 9, 8, 4, 2, 4, 4, 2, 2, 10, 10], dtype='int64') + + >>> psidx.delete(0).sort_values() # doctest: +SKIP + Int64Index([2, 2, 2, 4, 4, 4, 8, 9, 10, 10, 10], dtype='int64') + + >>> psidx.delete([0, 1, 2, 3, 10, 11]).sort_values() # doctest: +SKIP + Int64Index([2, 2, 2, 4, 4, 4], dtype='int64') + + MultiIndex + + >>> psidx = ps.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2), ('c', 'z', 3)]) + >>> psidx # doctest: +SKIP + MultiIndex([('a', 'x', 1), + ('b', 'y', 2), + ('c', 'z', 3)], + ) + + >>> psidx.delete([0, 2]).sort_values() # doctest: +SKIP + MultiIndex([('b', 'y', 2)], + ) + """ + length = len(self) + + def is_len_exceeded(index: int) -> bool: + """Check if the given index is exceeded the length or not""" + return index >= length if index >= 0 else abs(index) > length + + if not is_list_like(loc): + if is_len_exceeded(cast(int, loc)): + raise IndexError( + "index {} is out of bounds for axis 0 with size {}".format(loc, length) + ) + locs = [cast(int, loc)] + else: + for index in cast(List[int], loc): + if is_len_exceeded(index): + raise IndexError( + "index {} is out of bounds for axis 0 with size {}".format(index, length) + ) + locs = cast(List[int], loc) + + locs = [int(item) for item in locs] + locs = [item if item >= 0 else length + item for item in locs] + + # we need a temporary column such as '__index_value_0__' + # since 'InternalFrame.attach_default_index' will be failed + # when self._scol has name of '__index_level_0__' + index_value_column_format = "__index_value_{}__" + + sdf = self._internal._sdf + index_value_column_names = [ + verify_temp_column_name(sdf, index_value_column_format.format(i)) + for i in range(self._internal.index_level) + ] + index_value_columns = [ + index_scol.alias(index_vcol_name) + for index_scol, index_vcol_name in zip( + self._internal.index_spark_columns, index_value_column_names + ) + ] + sdf = sdf.select(index_value_columns) + + sdf = InternalFrame.attach_default_index(sdf, default_index_type="distributed-sequence") + # sdf here looks as below + # +-----------------+-----------------+-----------------+-----------------+ + # |__index_level_0__|__index_value_0__|__index_value_1__|__index_value_2__| + # +-----------------+-----------------+-----------------+-----------------+ + # | 0| a| x| 1| + # | 1| b| y| 2| + # | 2| c| z| 3| + # +-----------------+-----------------+-----------------+-----------------+ + + # delete rows which are matched with given `loc` + sdf = sdf.where(~F.col(SPARK_INDEX_NAME_FORMAT(0)).isin(locs)) + sdf = sdf.select(index_value_column_names) + # sdf here looks as below, we should alias them back to origin spark column names + # +-----------------+-----------------+-----------------+ + # |__index_value_0__|__index_value_1__|__index_value_2__| + # +-----------------+-----------------+-----------------+ + # | c| z| 3| + # +-----------------+-----------------+-----------------+ + index_origin_columns = [ + F.col(index_vcol_name).alias(index_scol_name) + for index_vcol_name, index_scol_name in zip( + index_value_column_names, self._internal.index_spark_column_names + ) + ] + sdf = sdf.select(index_origin_columns) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + + return DataFrame(internal).index + + def append(self, other: "Index") -> "Index": + """ + Append a collection of Index options together. + + Parameters + ---------- + other : Index + + Returns + ------- + appended : Index + + Examples + -------- + >>> psidx = ps.Index([10, 5, 0, 5, 10, 5, 0, 10]) + >>> psidx # doctest: +SKIP + Int64Index([10, 5, 0, 5, 10, 5, 0, 10], dtype='int64') + + >>> psidx.append(psidx) # doctest: +SKIP + Int64Index([10, 5, 0, 5, 10, 5, 0, 10, 10, 5, 0, 5, 10, 5, 0, 10], dtype='int64') + + Support for MiltiIndex + + >>> psidx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y')]) + >>> psidx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + ) + + >>> psidx.append(psidx) # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('a', 'x'), + ('b', 'y')], + ) + """ + from pyspark.pandas.indexes.multi import MultiIndex + from pyspark.pandas.indexes.category import CategoricalIndex + + if isinstance(self, MultiIndex) != isinstance(other, MultiIndex): + raise NotImplementedError( + "append() between Index & MultiIndex is currently not supported" + ) + if self._internal.index_level != other._internal.index_level: + raise NotImplementedError( + "append() between MultiIndexs with different levels is currently not supported" + ) + + index_fields = self._index_fields_for_union_like(other, func_name="append") + # Since pandas 1.5.0, the order of category matters. + if isinstance(other, CategoricalIndex): + other = other.reorder_categories(self.categories.to_list()) + + sdf_self = self._internal.spark_frame.select(self._internal.index_spark_columns) + sdf_other = other._internal.spark_frame.select(other._internal.index_spark_columns) + sdf_appended = sdf_self.union(sdf_other) + + # names should be kept when MultiIndex, but Index wouldn't keep its name. + if isinstance(self, MultiIndex): + index_names = self._internal.index_names + else: + index_names = None + + internal = InternalFrame( + spark_frame=sdf_appended, + index_spark_columns=[ + scol_for(sdf_appended, col) for col in self._internal.index_spark_column_names + ], + index_names=index_names, + index_fields=index_fields, + ) + + return DataFrame(internal).index + + def argmax(self) -> int: + """ + Return a maximum argument indexer. + + Parameters + ---------- + skipna : bool, default True + + Returns + ------- + maximum argument indexer + + Examples + -------- + >>> psidx = ps.Index([10, 9, 8, 7, 100, 5, 4, 3, 100, 3]) + >>> psidx # doctest: +SKIP + Int64Index([10, 9, 8, 7, 100, 5, 4, 3, 100, 3], dtype='int64') + + >>> psidx.argmax() + 4 + """ + sdf = self._internal.spark_frame.select(self.spark.column) + sequence_col = verify_temp_column_name(sdf, "__distributed_sequence_column__") + sdf = InternalFrame.attach_distributed_sequence_column(sdf, column_name=sequence_col) + # spark_frame here looks like below + # +-----------------+---------------+ + # |__index_level_0__|__index_value__| + # +-----------------+---------------+ + # | 0| 10| + # | 4| 100| + # | 2| 8| + # | 3| 7| + # | 6| 4| + # | 5| 5| + # | 7| 3| + # | 8| 100| + # | 1| 9| + # +-----------------+---------------+ + + return ( + sdf.orderBy( + scol_for(sdf, self._internal.data_spark_column_names[0]).desc(), + F.col(sequence_col).asc(), + ) + .select(sequence_col) + .first()[0] + ) + + def argmin(self) -> int: + """ + Return a minimum argument indexer. + + Parameters + ---------- + skipna : bool, default True + + Returns + ------- + minimum argument indexer + + Examples + -------- + >>> psidx = ps.Index([10, 9, 8, 7, 100, 5, 4, 3, 100, 3]) + >>> psidx # doctest: +SKIP + Int64Index([10, 9, 8, 7, 100, 5, 4, 3, 100, 3], dtype='int64') + + >>> psidx.argmin() + 7 + """ + sdf = self._internal.spark_frame.select(self.spark.column) + sequence_col = verify_temp_column_name(sdf, "__distributed_sequence_column__") + sdf = InternalFrame.attach_distributed_sequence_column(sdf, column_name=sequence_col) + + return ( + sdf.orderBy( + scol_for(sdf, self._internal.data_spark_column_names[0]).asc(), + F.col(sequence_col).asc(), + ) + .select(sequence_col) + .first()[0] + ) + + def set_names( + self, + names: Union[Name, List[Name]], + level: Optional[Union[int, Name, List[Union[int, Name]]]] = None, + inplace: bool = False, + ) -> Optional["Index"]: + """ + Set Index or MultiIndex name. + Able to set new names partially and by level. + + Parameters + ---------- + names : label or list of label + Name(s) to set. + level : int, label or list of int or label, optional + If the index is a MultiIndex, level(s) to set (None for all + levels). Otherwise level must be None. + inplace : bool, default False + Modifies the object directly, instead of creating a new Index or + MultiIndex. + + Returns + ------- + Index + The same type as the caller or None if inplace is True. + + See Also + -------- + Index.rename : Able to set new names without level. + + Examples + -------- + >>> idx = ps.Index([1, 2, 3, 4]) + >>> idx # doctest: +SKIP + Int64Index([1, 2, 3, 4], dtype='int64') + + >>> idx.set_names('quarter') # doctest: +SKIP + Int64Index([1, 2, 3, 4], dtype='int64', name='quarter') + + For MultiIndex + + >>> idx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y')]) + >>> idx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + ) + + >>> idx.set_names(['kind', 'year'], inplace=True) + >>> idx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + names=['kind', 'year']) + + >>> idx.set_names('species', level=0) # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y')], + names=['species', 'year']) + """ + from pyspark.pandas.indexes.multi import MultiIndex + + if isinstance(self, MultiIndex) and level is not None: + self_names = self.names + self_names[level] = names # type: ignore[index] + names = self_names + return self.rename(name=names, inplace=inplace) + + def difference(self, other: "Index", sort: Optional[bool] = None) -> "Index": + """ + Return a new Index with elements from the index that are not in + `other`. + + This is the set difference of two Index objects. + + Parameters + ---------- + other : Index or array-like + sort : True or None, default None + Whether to sort the resulting index. + * True : Attempt to sort the result. + * None : Do not sort the result. + + Returns + ------- + difference : Index + + Examples + -------- + + >>> idx1 = ps.Index([2, 1, 3, 4]) + >>> idx2 = ps.Index([3, 4, 5, 6]) + >>> idx1.difference(idx2, sort=True) # doctest: +SKIP + Int64Index([1, 2], dtype='int64') + + MultiIndex + + >>> midx1 = ps.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2), ('c', 'z', 3)]) + >>> midx2 = ps.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'z', 2), ('k', 'z', 3)]) + >>> midx1.difference(midx2) # doctest: +SKIP + MultiIndex([('b', 'y', 2), + ('c', 'z', 3)], + ) + """ + from pyspark.pandas.indexes.multi import MultiIndex + + # Check if the `self` and `other` have different index types. + # 1. `self` is Index, `other` is MultiIndex + # 2. `self` is MultiIndex, `other` is Index + is_index_types_different = isinstance(other, Index) and not isinstance(self, type(other)) + if is_index_types_different: + if isinstance(self, MultiIndex): + # In case `self` is MultiIndex and `other` is Index, + # return MultiIndex without its names. + return self.rename([None] * len(self)) + elif isinstance(self, Index): + # In case `self` is Index and `other` is MultiIndex, + # return Index without its name. + return self.rename(None) + + if not isinstance(other, (Index, Series, tuple, list, set, dict)): + raise TypeError("Input must be Index or array-like") + if not isinstance(sort, (type(None), type(True))): + raise ValueError( + "The 'sort' keyword only takes the values of None or True; {} was passed.".format( + sort + ) + ) + # Handling MultiIndex when `other` is not MultiIndex. + if isinstance(self, MultiIndex) and not isinstance(other, MultiIndex): + is_other_list_of_tuples = isinstance(other, (list, set, dict)) and all( + [isinstance(item, tuple) for item in other] + ) + if is_other_list_of_tuples: + other = MultiIndex.from_tuples(other) # type: ignore[arg-type] + else: + raise TypeError("other must be a MultiIndex or a list of tuples") + + if not isinstance(other, Index): + other = Index(other) + + sdf_self = self._internal.spark_frame + sdf_other = other._internal.spark_frame + idx_self = self._internal.index_spark_columns + idx_other = other._internal.index_spark_columns + sdf_diff = sdf_self.select(idx_self).subtract(sdf_other.select(idx_other)) + internal = InternalFrame( + spark_frame=sdf_diff, + index_spark_columns=[ + scol_for(sdf_diff, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + result = DataFrame(internal).index + # Name(s) will be kept when only name(s) of (Multi)Index are the same. + if isinstance(self, type(other)) and isinstance(self, MultiIndex): + if self.names == other.names: + result.names = self.names + elif isinstance(self, type(other)) and not isinstance(self, MultiIndex): + if self.name == other.name: + result.name = self.name + return result if sort is None else cast(Index, result.sort_values()) + + @property + def is_all_dates(self) -> bool: + """ + Return if all data types of the index are datetime. + remember that since pandas-on-Spark does not support multiple data types in an index, + so it returns True if any type of data is datetime. + + .. deprecated:: 3.4.0 + + Examples + -------- + >>> from datetime import datetime + + >>> idx = ps.Index([datetime(2019, 1, 1, 0, 0, 0), datetime(2019, 2, 3, 0, 0, 0)]) + >>> idx + DatetimeIndex(['2019-01-01', '2019-02-03'], dtype='datetime64[ns]', freq=None) + + >>> idx.is_all_dates + True + + >>> idx = ps.Index([datetime(2019, 1, 1, 0, 0, 0), None]) + >>> idx + DatetimeIndex(['2019-01-01', 'NaT'], dtype='datetime64[ns]', freq=None) + + >>> idx.is_all_dates + True + + >>> idx = ps.Index([0, 1, 2]) + >>> idx # doctest: +SKIP + Int64Index([0, 1, 2], dtype='int64') + + >>> idx.is_all_dates + False + """ + warnings.warn( + "Index.is_all_dates is deprecated, will be removed in a future version. " + "check index.inferred_type instead", + FutureWarning, + ) + return isinstance(self.spark.data_type, (TimestampType, TimestampNTZType)) + + def repeat(self, repeats: int) -> "Index": + """ + Repeat elements of a Index/MultiIndex. + + Returns a new Index/MultiIndex where each element of the current Index/MultiIndex + is repeated consecutively a given number of times. + + Parameters + ---------- + repeats : int + The number of repetitions for each element. This should be a + non-negative integer. Repeating 0 times will return an empty + Index. + + Returns + ------- + repeated_index : Index/MultiIndex + Newly created Index/MultiIndex with repeated elements. + + See Also + -------- + Series.repeat : Equivalent function for Series. + + Examples + -------- + >>> idx = ps.Index(['a', 'b', 'c']) + >>> idx + Index(['a', 'b', 'c'], dtype='object') + >>> idx.repeat(2) + Index(['a', 'b', 'c', 'a', 'b', 'c'], dtype='object') + + For MultiIndex, + + >>> midx = ps.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')]) + >>> midx # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('y', 'c')], + ) + >>> midx.repeat(2) # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('y', 'c'), + ('x', 'a'), + ('x', 'b'), + ('y', 'c')], + ) + >>> midx.repeat(0) # doctest: +SKIP + MultiIndex([], ) + """ + if not isinstance(repeats, int): + raise TypeError( + "`repeats` argument must be integer, but got {}".format(type(repeats).__name__) + ) + elif repeats < 0: + raise ValueError("negative dimensions are not allowed") + + psdf: DataFrame = DataFrame(self._internal.resolved_copy) + if repeats == 0: + return DataFrame(psdf._internal.with_filter(F.lit(False))).index + else: + return ps.concat([psdf] * repeats).index + + def asof(self, label: Any) -> Scalar: + """ + Return the label from the index, or, if not present, the previous one. + + Assuming that the index is sorted, return the passed index label if it + is in the index, or return the previous index label if the passed one + is not in the index. + + .. note:: This API is dependent on :meth:`Index.is_monotonic_increasing` + which can be expensive. + + Parameters + ---------- + label : object + The label up to which the method returns the latest index label. + + Returns + ------- + object + The passed label if it is in the index. The previous label if the + passed label is not in the sorted index or `NaN` if there is no + such label. + + Examples + -------- + `Index.asof` returns the latest index label up to the passed label. + + >>> idx = ps.Index(['2013-12-31', '2014-01-02', '2014-01-03']) + >>> idx.asof('2014-01-01') + '2013-12-31' + + If the label is in the index, the method returns the passed label. + + >>> idx.asof('2014-01-02') + '2014-01-02' + + If all of the labels in the index are later than the passed label, + NaN is returned. + + >>> idx.asof('1999-01-02') + nan + """ + sdf = self._internal.spark_frame + if self.is_monotonic_increasing: + sdf = sdf.where(self.spark.column <= F.lit(label).cast(self.spark.data_type)).select( + F.max(self.spark.column) + ) + elif self.is_monotonic_decreasing: + sdf = sdf.where(self.spark.column >= F.lit(label).cast(self.spark.data_type)).select( + F.min(self.spark.column) + ) + else: + raise ValueError("index must be monotonic increasing or decreasing") + + result = sdf.toPandas().iloc[0, 0] + return result if result is not None else np.nan + + def _index_fields_for_union_like( + self: "Index", other: "Index", func_name: str + ) -> Optional[List[InternalField]]: + if self._internal.index_fields == other._internal.index_fields: + return self._internal.index_fields + elif all( + left.dtype == right.dtype + and (isinstance(left.dtype, CategoricalDtype) or left.spark_type == right.spark_type) + for left, right in zip(self._internal.index_fields, other._internal.index_fields) + ): + return [ + left.copy(nullable=left.nullable or right.nullable) + if left.spark_type == right.spark_type + else InternalField(dtype=left.dtype) + for left, right in zip(self._internal.index_fields, other._internal.index_fields) + ] + elif any( + isinstance(field.dtype, CategoricalDtype) + for field in self._internal.index_fields + other._internal.index_fields + ): + # TODO: non-categorical or categorical with different categories + raise NotImplementedError( + "{}() between CategoricalIndex and non-categorical or " + "categorical with different categories is currently not supported".format(func_name) + ) + else: + return None + + def union( + self, other: Union[DataFrame, Series, "Index", List], sort: Optional[bool] = None + ) -> "Index": + """ + Form the union of two Index objects. + + Parameters + ---------- + other : Index or array-like + sort : bool or None, default None + Whether to sort the resulting Index. + + Returns + ------- + union : Index + + Examples + -------- + + Index + + >>> idx1 = ps.Index([1, 2, 3, 4]) + >>> idx2 = ps.Index([3, 4, 5, 6]) + >>> idx1.union(idx2).sort_values() # doctest: +SKIP + Int64Index([1, 2, 3, 4, 5, 6], dtype='int64') + + MultiIndex + + >>> midx1 = ps.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("x", "c"), ("x", "d")]) + >>> midx2 = ps.MultiIndex.from_tuples([("x", "c"), ("x", "d"), ("x", "e"), ("x", "f")]) + >>> midx1.union(midx2).sort_values() # doctest: +SKIP + MultiIndex([('x', 'a'), + ('x', 'b'), + ('x', 'c'), + ('x', 'd'), + ('x', 'e'), + ('x', 'f')], + ) + """ + from pyspark.pandas.indexes.multi import MultiIndex + + sort = True if sort is None else sort + sort = validate_bool_kwarg(sort, "sort") + other_idx: Index + if isinstance(self, MultiIndex): + if isinstance(other, MultiIndex): + other_idx = other + elif isinstance(other, list) and all(isinstance(item, tuple) for item in other): + other_idx = MultiIndex.from_tuples(other) + else: + raise TypeError("other must be a MultiIndex or a list of tuples") + else: + if isinstance(other, MultiIndex): + # TODO: We can't support different type of values in a single column for now. + raise NotImplementedError("Union between Index and MultiIndex is not yet supported") + elif isinstance(other, DataFrame): + raise ValueError("Index data must be 1-dimensional") + else: + other_idx = Index(other) + + index_fields = self._index_fields_for_union_like(other_idx, func_name="union") + + sdf_self = self._internal.spark_frame.select(self._internal.index_spark_columns) + sdf_other = other_idx._internal.spark_frame.select(other_idx._internal.index_spark_columns) + sdf = sdf_self.unionAll(sdf_other).exceptAll(sdf_self.intersectAll(sdf_other)) + if sort: + sdf = sdf.sort(*self._internal.index_spark_column_names) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=index_fields, + ) + + return DataFrame(internal).index + + def holds_integer(self) -> bool: + """ + Whether the type is an integer type. + Always return False for MultiIndex. + + Notes + ----- + When Index contains null values the result can be different with pandas + since pandas-on-Spark cast integer to float when Index contains null values. + + >>> ps.Index([1, 2, 3, None]) # doctest: +SKIP + Float64Index([1.0, 2.0, 3.0, nan], dtype='float64') + + Examples + -------- + >>> psidx = ps.Index([1, 2, 3, 4]) + >>> psidx.holds_integer() + True + + Returns False for string type. + + >>> psidx = ps.Index(["A", "B", "C", "D"]) + >>> psidx.holds_integer() + False + + Returns False for float type. + + >>> psidx = ps.Index([1.1, 2.2, 3.3, 4.4]) + >>> psidx.holds_integer() + False + """ + return isinstance(self.spark.data_type, IntegralType) + + def intersection(self, other: Union[DataFrame, Series, "Index", List]) -> "Index": + """ + Form the intersection of two Index objects. + + This returns a new Index with elements common to the index and `other`. + + Parameters + ---------- + other : Index or array-like + + Returns + ------- + intersection : Index + + Examples + -------- + >>> idx1 = ps.Index([1, 2, 3, 4]) + >>> idx2 = ps.Index([3, 4, 5, 6]) + >>> idx1.intersection(idx2).sort_values() # doctest: +SKIP + Int64Index([3, 4], dtype='int64') + """ + from pyspark.pandas.indexes.multi import MultiIndex + + other_idx: Index + if isinstance(other, DataFrame): + raise ValueError("Index data must be 1-dimensional") + elif isinstance(other, MultiIndex): + # Always returns a no-named empty Index if `other` is MultiIndex. + return self._psdf.head(0).index.rename(None) + elif isinstance(other, Index): + other_idx = other + spark_frame_other = other_idx.to_frame()._to_spark() + keep_name = self.name == other_idx.name + elif isinstance(other, Series): + other_idx = Index(other) + spark_frame_other = other_idx.to_frame()._to_spark() + keep_name = True + elif is_list_like(other): + other_idx = Index(other) + if isinstance(other_idx, MultiIndex): + raise ValueError("Names should be list-like for a MultiIndex") + spark_frame_other = other_idx.to_frame()._to_spark() + keep_name = True + else: + raise TypeError("Input must be Index or array-like") + + index_fields = self._index_fields_for_union_like(other_idx, func_name="intersection") + + spark_frame_self = self.to_frame(name=SPARK_DEFAULT_INDEX_NAME)._to_spark() + spark_frame_intersected = spark_frame_self.intersect(spark_frame_other) + if keep_name: + index_names = self._internal.index_names + else: + index_names = None + + internal = InternalFrame( + spark_frame=spark_frame_intersected, + index_spark_columns=[scol_for(spark_frame_intersected, SPARK_DEFAULT_INDEX_NAME)], + index_names=index_names, + index_fields=index_fields, + ) + + return DataFrame(internal).index + + def item(self) -> Union[Scalar, Tuple[Scalar, ...]]: + """ + Return the first element of the underlying data as a python scalar. + + Returns + ------- + scalar + The first element of Index. + + Raises + ------ + ValueError + If the data is not length-1. + + Examples + -------- + >>> psidx = ps.Index([10]) + >>> psidx.item() + 10 + """ + return self.to_series().item() + + def insert(self, loc: int, item: Any) -> "Index": + """ + Make new Index inserting new item at location. + + Follows Python list.append semantics for negative values. + + .. versionchanged:: 3.4.0 + Raise IndexError when loc is out of bounds to follow Pandas 1.4+ behavior + + Parameters + ---------- + loc : int + item : object + + Returns + ------- + new_index : Index + + Examples + -------- + >>> psidx = ps.Index([1, 2, 3, 4, 5]) + >>> psidx.insert(3, 100) # doctest: +SKIP + Int64Index([1, 2, 3, 100, 4, 5], dtype='int64') + + For negative values + + >>> psidx = ps.Index([1, 2, 3, 4, 5]) + >>> psidx.insert(-3, 100) # doctest: +SKIP + Int64Index([1, 2, 100, 3, 4, 5], dtype='int64') + """ + validate_index_loc(self, loc) + loc = loc + len(self) if loc < 0 else loc + + index_name = self._internal.index_spark_column_names[0] + sdf_before = self.to_frame(name=index_name)[:loc]._to_spark() + sdf_middle = Index([item], dtype=self.dtype).to_frame(name=index_name)._to_spark() + sdf_after = self.to_frame(name=index_name)[loc:]._to_spark() + sdf = sdf_before.union(sdf_middle).union(sdf_after) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=[InternalField(field.dtype) for field in self._internal.index_fields], + ) + return DataFrame(internal).index + + def view(self) -> "Index": + """ + this is defined as a copy with the same identity + """ + return self.copy() + + def to_list(self) -> List: + """ + Return a list of the values. + + These are each a scalar type, which is a Python scalar + (for str, int, float) or a pandas scalar + (for Timestamp/Timedelta/Interval/Period) + + .. note:: This method should only be used if the resulting list is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + Index + + >>> idx = ps.Index([1, 2, 3, 4, 5]) + >>> idx.to_list() + [1, 2, 3, 4, 5] + + MultiIndex + + >>> tuples = [(1, 'red'), (1, 'blue'), (2, 'red'), (2, 'green')] + >>> midx = ps.MultiIndex.from_tuples(tuples) + >>> midx.to_list() + [(1, 'red'), (1, 'blue'), (2, 'red'), (2, 'green')] + """ + log_advice( + "`to_list` loads all data into the driver's memory. " + "It should only be used if the resulting list is expected to be small." + ) + return self._to_internal_pandas().tolist() + + tolist = to_list + + @property + def inferred_type(self) -> str: + """ + Return a string of the type inferred from the values. + + Examples + -------- + >>> from datetime import datetime + >>> ps.Index([1, 2, 3]).inferred_type + 'integer' + + >>> ps.Index([1.0, 2.0, 3.0]).inferred_type + 'floating' + + >>> ps.Index(['a', 'b', 'c']).inferred_type + 'string' + + >>> ps.Index([True, False, True, False]).inferred_type + 'boolean' + """ + return lib.infer_dtype([self.to_series().head(1).item()]) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeIndex, item): + property_or_func = getattr(MissingPandasLikeIndex, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item)) + + def __repr__(self) -> str: + max_display_count = get_option("display.max_rows") + if max_display_count is None: + return repr(self._to_internal_pandas()) + + pindex = self._psdf._get_or_create_repr_pandas_cache(max_display_count).index + + pindex_length = len(pindex) + repr_string = repr(pindex[:max_display_count]) + + if pindex_length > max_display_count: + footer = "\nShowing only the first {}".format(max_display_count) + return repr_string + footer + return repr_string + + def __iter__(self) -> Iterator: + return MissingPandasLikeIndex.__iter__(self) + + def __and__(self, other: "Index") -> "Index": + warnings.warn( + "Index.__and__ operating as a set operation is deprecated, " + "in the future this will be a logical operation matching Series.__and__. " + "Use index.intersection(other) instead", + FutureWarning, + ) + return self.intersection(other) + + def __or__(self, other: "Index") -> "Index": + warnings.warn( + "Index.__or__ operating as a set operation is deprecated, " + "in the future this will be a logical operation matching Series.__or__. " + "Use index.union(other) instead", + FutureWarning, + ) + return self.union(other) + + def __xor__(self, other: "Index") -> "Index": + warnings.warn( + "Index.__xor__ operating as a set operation is deprecated, " + "in the future this will be a logical operation matching Series.__xor__. " + "Use index.symmetric_difference(other) instead", + FutureWarning, + ) + return self.symmetric_difference(other) + + def __rxor__(self, other: Any) -> "Index": + return NotImplemented + + def __bool__(self) -> bool: + raise ValueError( + "The truth value of a {0} is ambiguous. " + "Use a.empty, a.bool(), a.item(), a.any() or a.all().".format(self.__class__.__name__) + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.indexes.base + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexes.base.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexes.base tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexes.base, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/category.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/category.py new file mode 100644 index 0000000..7bc8780 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/category.py @@ -0,0 +1,773 @@ +# +# 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 warnings +from typing import Any, Callable, List, Optional, Union, cast, no_type_check + +import pandas as pd +from pandas.api.types import is_hashable, CategoricalDtype # type: ignore[attr-defined] + +from pyspark import pandas as ps +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.internal import InternalField +from pyspark.pandas.series import Series +from pyspark.sql.types import StructField + + +class CategoricalIndex(Index): + """ + Index based on an underlying `Categorical`. + + CategoricalIndex can only take on a limited, + and usually fixed, number of possible values (`categories`). Also, + it might have an order, but numerical operations + (additions, divisions, ...) are not possible. + + Parameters + ---------- + data : array-like (1-dimensional) + The values of the categorical. If `categories` are given, values not in + `categories` will be replaced with NaN. + categories : index-like, optional + The categories for the categorical. Items need to be unique. + If the categories are not given here (and also not in `dtype`), they + will be inferred from the `data`. + ordered : bool, optional + Whether or not this categorical is treated as an ordered + categorical. If not given here or in `dtype`, the resulting + categorical will be unordered. + dtype : CategoricalDtype or "category", optional + If :class:`CategoricalDtype`, cannot be used together with + `categories` or `ordered`. + copy : bool, default False + Make a copy of input ndarray. + name : object, optional + Name to be stored in the index. + + See Also + -------- + Index : The base pandas-on-Spark Index type. + + Examples + -------- + >>> ps.CategoricalIndex(["a", "b", "c", "a", "b", "c"]) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c', 'a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + ``CategoricalIndex`` can also be instantiated from a ``Categorical``: + + >>> c = pd.Categorical(["a", "b", "c", "a", "b", "c"]) + >>> ps.CategoricalIndex(c) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c', 'a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + Ordered ``CategoricalIndex`` can have a min and max value. + + >>> ci = ps.CategoricalIndex( + ... ["a", "b", "c", "a", "b", "c"], ordered=True, categories=["c", "b", "a"] + ... ) + >>> ci # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c', 'a', 'b', 'c'], + categories=['c', 'b', 'a'], ordered=True, dtype='category') + + From a Series: + + >>> s = ps.Series(["a", "b", "c", "a", "b", "c"], index=[10, 20, 30, 40, 50, 60]) + >>> ps.CategoricalIndex(s) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c', 'a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + From an Index: + + >>> idx = ps.Index(["a", "b", "c", "a", "b", "c"]) + >>> ps.CategoricalIndex(idx) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c', 'a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + """ + + @no_type_check + def __new__(cls, data=None, categories=None, ordered=None, dtype=None, copy=False, name=None): + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, (Series, Index)): + if dtype is None: + dtype = "category" + return Index(data, dtype=dtype, copy=copy, name=name) + + return ps.from_pandas( + pd.CategoricalIndex( + data=data, categories=categories, ordered=ordered, dtype=dtype, name=name + ) + ) + + @property + def dtype(self) -> CategoricalDtype: + return cast(CategoricalDtype, super().dtype) + + @property + def codes(self) -> Index: + """ + The category codes of this categorical. + + Codes are an Index of integers which are the positions of the actual + values in the categories Index. + + There is no setter, use the other categorical methods and the normal item + setter to change values in the categorical. + + Returns + ------- + Index + A non-writable view of the `codes` Index. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.codes # doctest: +SKIP + Int64Index([0, 1, 1, 2, 2, 2], dtype='int64') + """ + return self._with_new_scol( + self.spark.column, + field=InternalField.from_struct_field( + StructField( + name=self._internal.index_spark_column_names[0], + dataType=self.spark.data_type, + nullable=self.spark.nullable, + ) + ), + ).rename(None) + + @property + def categories(self) -> pd.Index: + """ + The categories of this categorical. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.categories + Index(['a', 'b', 'c'], dtype='object') + """ + return self.dtype.categories + + @categories.setter + def categories(self, categories: Union[pd.Index, List]) -> None: + dtype = CategoricalDtype(categories, ordered=self.ordered) + + if len(self.categories) != len(dtype.categories): + raise ValueError( + "new categories need to have the same number of items as the old categories!" + ) + + internal = self._psdf._internal.copy( + index_fields=[self._internal.index_fields[0].copy(dtype=dtype)] + ) + self._psdf._update_internal_frame(internal) + + @property + def ordered(self) -> bool: + """ + Whether the categories have an ordered relationship. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.ordered + False + """ + return self.dtype.ordered + + def add_categories( + self, new_categories: Union[pd.Index, Any, List], inplace: bool = False + ) -> Optional["CategoricalIndex"]: + """ + Add new categories. + + `new_categories` will be included at the last/highest place in the + categories and will be unused directly after this call. + + Parameters + ---------- + new_categories : category or list-like of category + The new categories to be included. + inplace : bool, default False + Whether or not to add the categories inplace or return a copy of + this categorical with added categories. + + .. deprecated:: 3.2.0 + + Returns + ------- + CategoricalIndex or None + Categorical with new categories added or None if ``inplace=True``. + + Raises + ------ + ValueError + If the new categories include old categories or do not validate as + categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.add_categories('x') # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c', 'x'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex( + self.to_series().cat.add_categories(new_categories=new_categories) + ).rename(self.name) + + def as_ordered(self, inplace: bool = False) -> Optional["CategoricalIndex"]: + """ + Set the Categorical to be ordered. + + Parameters + ---------- + inplace : bool, default False + Whether or not to set the ordered attribute in-place or return + a copy of this categorical with ordered set to True. + + Returns + ------- + CategoricalIndex or None + Ordered Categorical or None if ``inplace=True``. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.as_ordered() # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=True, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex(self.to_series().cat.as_ordered()).rename(self.name) + + def as_unordered(self, inplace: bool = False) -> Optional["CategoricalIndex"]: + """ + Set the Categorical to be unordered. + + Parameters + ---------- + inplace : bool, default False + Whether or not to set the ordered attribute in-place or return + a copy of this categorical with ordered set to False. + + Returns + ------- + CategoricalIndex or None + Unordered Categorical or None if ``inplace=True``. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")).as_ordered() + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=True, dtype='category') + + >>> idx.as_unordered() # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex(self.to_series().cat.as_unordered()).rename(self.name) + + def remove_categories( + self, removals: Union[pd.Index, Any, List], inplace: bool = False + ) -> Optional["CategoricalIndex"]: + """ + Remove the specified categories. + + `removals` must be included in the old categories. Values which were in + the removed categories will be set to NaN + + Parameters + ---------- + removals : category or list of categories + The categories which should be removed. + inplace : bool, default False + Whether or not to remove the categories inplace or return a copy of + this categorical with removed categories. + + .. deprecated:: 3.2.0 + + Returns + ------- + CategoricalIndex or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If the removals are not contained in the categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.remove_categories('b') # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', nan, nan, 'c', 'c', 'c'], + categories=['a', 'c'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex(self.to_series().cat.remove_categories(removals)).rename(self.name) + + def remove_unused_categories(self, inplace: bool = False) -> Optional["CategoricalIndex"]: + """ + Remove categories which are not used. + + Parameters + ---------- + inplace : bool, default False + Whether or not to drop unused categories inplace or return a copy of + this categorical with unused categories dropped. + + .. deprecated:: 3.2.0 + + Returns + ------- + cat : CategoricalIndex or None + Categorical with unused categories dropped or None if ``inplace=True``. + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc"), categories=['a', 'b', 'c', 'd']) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c', 'd'], ordered=False, dtype='category') + + >>> idx.remove_unused_categories() # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex(self.to_series().cat.remove_unused_categories()).rename(self.name) + + def rename_categories( + self, new_categories: Union[list, dict, Callable], inplace: bool = False + ) -> Optional["CategoricalIndex"]: + """ + Rename categories. + + Parameters + ---------- + new_categories : list-like, dict-like or callable + + New categories which will replace old categories. + + * list-like: all items must be unique and the number of items in + the new categories must match the existing number of categories. + + * dict-like: specifies a mapping from + old categories to new. Categories not contained in the mapping + are passed through and extra categories in the mapping are + ignored. + + * callable : a callable that is called on all items in the old + categories and whose return values comprise the new categories. + + inplace : bool, default False + Whether or not to rename the categories inplace or return a copy of + this categorical with renamed categories. + + .. deprecated:: 3.2.0 + + Returns + ------- + cat : CategoricalIndex or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If new categories are list-like and do not have the same number of + items than the current categories or do not validate as categories + + See Also + -------- + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> idx = ps.CategoricalIndex(["a", "a", "b"]) + >>> idx.rename_categories([0, 1]) + CategoricalIndex([0, 0, 1], categories=[0, 1], ordered=False, dtype='category') + + For dict-like ``new_categories``, extra keys are ignored and + categories not in the dictionary are passed through + + >>> idx.rename_categories({'a': 'A', 'c': 'C'}) + CategoricalIndex(['A', 'A', 'b'], categories=['A', 'b'], ordered=False, dtype='category') + + You may also provide a callable to create the new categories + + >>> idx.rename_categories(lambda x: x.upper()) + CategoricalIndex(['A', 'A', 'B'], categories=['A', 'B'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex(self.to_series().cat.rename_categories(new_categories)).rename( + self.name + ) + + def reorder_categories( + self, + new_categories: Union[pd.Index, Any, List], + ordered: Optional[bool] = None, + inplace: bool = False, + ) -> Optional["CategoricalIndex"]: + """ + Reorder categories as specified in new_categories. + + `new_categories` needs to include all old categories and no new category + items. + + Parameters + ---------- + new_categories : Index-like + The categories in new order. + ordered : bool, optional + Whether or not the categorical is treated as an ordered categorical. + If not given, do not change the ordered information. + inplace : bool, default False + Whether or not to reorder the categories inplace or return a copy of + this categorical with reordered categories. + + .. deprecated:: 3.2.0 + + Returns + ------- + cat : CategoricalIndex or None + Categorical with removed categories or None if ``inplace=True``. + + Raises + ------ + ValueError + If the new categories do not contain all old category items or any + new ones + + See Also + -------- + rename_categories : Rename categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + set_categories : Set the categories to the specified ones. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.reorder_categories(['c', 'b', 'a']) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['c', 'b', 'a'], ordered=False, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex( + self.to_series().cat.reorder_categories(new_categories=new_categories, ordered=ordered) + ).rename(self.name) + + def set_categories( + self, + new_categories: Union[pd.Index, List], + ordered: Optional[bool] = None, + rename: bool = False, + inplace: bool = False, + ) -> Optional["CategoricalIndex"]: + """ + Set the categories to the specified new_categories. + + `new_categories` can include new categories (which will result in + unused categories) or remove old categories (which results in values + set to NaN). If `rename==True`, the categories will simply be renamed + (less or more items than in old categories will result in values set to + NaN or in unused categories respectively). + + This method can be used to perform more than one action of adding, + removing, and reordering simultaneously and is therefore faster than + performing the individual steps via the more specialised methods. + + On the other hand this methods does not do checks (e.g., whether the + old categories are included in the new categories on a reorder), which + can result in surprising changes, for example when using special string + dtypes, which does not consider a S1 string equal to a single char + python string. + + Parameters + ---------- + new_categories : Index-like + The categories in new order. + ordered : bool, default False + Whether or not the categorical is treated as an ordered categorical. + If not given, do not change the ordered information. + rename : bool, default False + Whether or not the new_categories should be considered as a rename + of the old categories or as reordered categories. + inplace : bool, default False + Whether or not to reorder the categories in-place or return a copy + of this categorical with reordered categories. + + .. deprecated:: 3.2.0 + + Returns + ------- + CategoricalIndex with reordered categories or None if inplace. + + Raises + ------ + ValueError + If new_categories does not validate as categories + + See Also + -------- + rename_categories : Rename categories. + reorder_categories : Reorder categories. + add_categories : Add new categories. + remove_categories : Remove the specified categories. + remove_unused_categories : Remove categories which are not used. + + Examples + -------- + >>> idx = ps.CategoricalIndex(list("abbccc")) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'b', 'c', 'c', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.set_categories(['b', 'c']) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex([nan, 'b', 'b', 'c', 'c', 'c'], + categories=['b', 'c'], ordered=False, dtype='category') + + >>> idx.set_categories([1, 2, 3], rename=True) + CategoricalIndex([1, 2, 2, 3, 3, 3], categories=[1, 2, 3], ordered=False, dtype='category') + + >>> idx.set_categories([1, 2, 3], rename=True, ordered=True) + CategoricalIndex([1, 2, 2, 3, 3, 3], categories=[1, 2, 3], ordered=True, dtype='category') + """ + if inplace: + warnings.warn( + "Argument `inplace` will be removed in 4.0.0.", + FutureWarning, + ) + raise ValueError("cannot use inplace with CategoricalIndex") + + return CategoricalIndex( + self.to_series().cat.set_categories(new_categories, ordered=ordered, rename=rename) + ).rename(self.name) + + def map( # type: ignore[override] + self, mapper: Union[dict, Callable[[Any], Any], pd.Series] + ) -> "Index": + """ + Map values using input correspondence (a dict, Series, or function). + + Maps the values (their categories, not the codes) of the index to new + categories. If the mapping correspondence is one-to-one the result is a + `CategoricalIndex` which has the same order property as the original, + otherwise an `Index` is returned. + + If a `dict` or `Series` is used any unmapped category is mapped to missing values. + Note that if this happens an `Index` will be returned. + + Parameters + ---------- + mapper : function, dict, or Series + Mapping correspondence. + + Returns + ------- + CategoricalIndex or Index + Mapped index. + + See Also + -------- + Index.map : Apply a mapping correspondence on an `Index`. + Series.map : Apply a mapping correspondence on a `Series` + Series.apply : Apply more complex functions on a `Series` + + Examples + -------- + >>> idx = ps.CategoricalIndex(['a', 'b', 'c']) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=False, dtype='category') + + >>> idx.map(lambda x: x.upper()) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['A', 'B', 'C'], + categories=['A', 'B', 'C'], ordered=False, dtype='category') + + >>> pser = pd.Series([1, 2, 3], index=pd.CategoricalIndex(['a', 'b', 'c'], ordered=True)) + >>> idx.map(pser) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex([1, 2, 3], + categories=[1, 2, 3], ordered=False, dtype='category') + + >>> idx.map({'a': 'first', 'b': 'second', 'c': 'third'}) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['first', 'second', 'third'], + categories=['first', 'second', 'third'], ordered=False, dtype='category') + + If the mapping is one-to-one the ordering of the categories is preserved: + + >>> idx = ps.CategoricalIndex(['a', 'b', 'c'], ordered=True) + >>> idx # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex(['a', 'b', 'c'], + categories=['a', 'b', 'c'], ordered=True, dtype='category') + + >>> idx.map({'a': 3, 'b': 2, 'c': 1}) # doctest: +NORMALIZE_WHITESPACE + CategoricalIndex([3, 2, 1], + categories=[3, 2, 1], ordered=True, dtype='category') + + If the mapping is not one-to-one an `Index` is returned: + + >>> idx.map({'a': 'first', 'b': 'second', 'c': 'first'}) + Index(['first', 'second', 'first'], dtype='object') + + If a `dict` is used, all unmapped categories are mapped to None and + the result is an `Index`: + + >>> idx.map({'a': 'first', 'b': 'second'}) + Index(['first', 'second', None], dtype='object') + """ + return super().map(mapper) + + @no_type_check + def all(self, *args, **kwargs) -> None: + raise TypeError("Cannot perform 'all' with this index type: %s" % type(self).__name__) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.indexes.category + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexes.category.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexes.category tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexes.category, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/datetimes.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/datetimes.py new file mode 100644 index 0000000..9adef61 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/datetimes.py @@ -0,0 +1,843 @@ +# +# 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 datetime +import warnings +from functools import partial +from typing import Any, Optional, Union, cast, no_type_check + +import pandas as pd +from pandas.api.types import is_hashable # type: ignore[attr-defined] +from pandas.tseries.offsets import DateOffset +from pyspark._globals import _NoValue + +from pyspark import pandas as ps +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.missing.indexes import MissingPandasLikeDatetimeIndex +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.utils import verify_temp_column_name + + +class DatetimeIndex(Index): + """ + Immutable ndarray-like of datetime64 data. + + Parameters + ---------- + data : array-like (1-dimensional), optional + Optional datetime-like data to construct index with. + freq : str or pandas offset object, optional + One of pandas date offset strings or corresponding objects. The string + 'infer' can be passed in order to set the frequency of the index as the + inferred frequency upon creation. + normalize : bool, default False + Normalize start/end dates to midnight before generating date range. + closed : {'left', 'right'}, optional + Set whether to include `start` and `end` that are on the + boundary. The default includes boundary points on either end. + ambiguous : 'infer', bool-ndarray, 'NaT', default 'raise' + When clocks moved backward due to DST, ambiguous times may arise. + For example in Central European Time (UTC+01), when going from 03:00 + DST to 02:00 non-DST, 02:30:00 local time occurs both at 00:30:00 UTC + and at 01:30:00 UTC. In such a situation, the `ambiguous` parameter + dictates how ambiguous times should be handled. + + - 'infer' will attempt to infer fall dst-transition hours based on + order + - bool-ndarray where True signifies a DST time, False signifies a + non-DST time (note that this flag is only applicable for ambiguous + times) + - 'NaT' will return NaT where there are ambiguous times + - 'raise' will raise an AmbiguousTimeError if there are ambiguous times. + dayfirst : bool, default False + If True, parse dates in `data` with the day first order. + yearfirst : bool, default False + If True parse dates in `data` with the year first order. + dtype : numpy.dtype or str, default None + Note that the only NumPy dtype allowed is ‘datetime64[ns]’. + copy : bool, default False + Make a copy of input ndarray. + name : label, default None + Name to be stored in the index. + + See Also + -------- + Index : The base pandas Index type. + to_datetime : Convert argument to datetime. + + Examples + -------- + >>> ps.DatetimeIndex(['1970-01-01', '1970-01-01', '1970-01-01']) + DatetimeIndex(['1970-01-01', '1970-01-01', '1970-01-01'], dtype='datetime64[ns]', freq=None) + + From a Series: + + >>> from datetime import datetime + >>> s = ps.Series([datetime(2021, 3, 1), datetime(2021, 3, 2)], index=[10, 20]) + >>> ps.DatetimeIndex(s) + DatetimeIndex(['2021-03-01', '2021-03-02'], dtype='datetime64[ns]', freq=None) + + From an Index: + + >>> idx = ps.DatetimeIndex(['1970-01-01', '1970-01-01', '1970-01-01']) + >>> ps.DatetimeIndex(idx) + DatetimeIndex(['1970-01-01', '1970-01-01', '1970-01-01'], dtype='datetime64[ns]', freq=None) + """ + + @no_type_check + def __new__( + cls, + data=None, + freq=_NoValue, + normalize=False, + closed=None, + ambiguous="raise", + dayfirst=False, + yearfirst=False, + dtype=None, + copy=False, + name=None, + ) -> "DatetimeIndex": + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, (Series, Index)): + if dtype is None: + dtype = "datetime64[ns]" + return cast(DatetimeIndex, Index(data, dtype=dtype, copy=copy, name=name)) + + kwargs = dict( + data=data, + normalize=normalize, + closed=closed, + ambiguous=ambiguous, + dayfirst=dayfirst, + yearfirst=yearfirst, + dtype=dtype, + copy=copy, + name=name, + ) + if freq is not _NoValue: + kwargs["freq"] = freq + + return cast(DatetimeIndex, ps.from_pandas(pd.DatetimeIndex(**kwargs))) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeDatetimeIndex, item): + property_or_func = getattr(MissingPandasLikeDatetimeIndex, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError("'DatetimeIndex' object has no attribute '{}'".format(item)) + + # Properties + @property + def year(self) -> Index: + """ + The year of the datetime. + """ + return Index(self.to_series().dt.year) + + @property + def month(self) -> Index: + """ + The month of the timestamp as January = 1 December = 12. + """ + return Index(self.to_series().dt.month) + + @property + def day(self) -> Index: + """ + The days of the datetime. + """ + warnings.warn( + "`day` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.day) + + @property + def hour(self) -> Index: + """ + The hours of the datetime. + """ + warnings.warn( + "`hour` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.hour) + + @property + def minute(self) -> Index: + """ + The minutes of the datetime. + """ + warnings.warn( + "`minute` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.minute) + + @property + def second(self) -> Index: + """ + The seconds of the datetime. + """ + warnings.warn( + "`second` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.second) + + @property + def microsecond(self) -> Index: + """ + The microseconds of the datetime. + """ + warnings.warn( + "`microsecond` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.microsecond) + + @property + def week(self) -> Index: + """ + The week ordinal of the year. + + .. deprecated:: 3.5.0 + """ + warnings.warn( + "`week` is deprecated in 3.5.0 and will be removed in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.week) + + @property + def weekofyear(self) -> Index: + warnings.warn( + "`weekofyear` is deprecated in 3.5.0 and will be removed in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.weekofyear) + + weekofyear.__doc__ = week.__doc__ + + @property + def dayofweek(self) -> Index: + """ + The day of the week with Monday=0, Sunday=6. + Return the day of the week. It is assumed the week starts on + Monday, which is denoted by 0 and ends on Sunday which is denoted + by 6. This method is available on both Series with datetime + values (using the `dt` accessor) or DatetimeIndex. + + Returns + ------- + Series or Index + Containing integers indicating the day number. + + See Also + -------- + Series.dt.dayofweek : Alias. + Series.dt.weekday : Alias. + Series.dt.day_name : Returns the name of the day of the week. + + Examples + -------- + >>> idx = ps.date_range('2016-12-31', '2017-01-08', freq='D') # doctest: +SKIP + >>> idx.dayofweek # doctest: +SKIP + Int64Index([5, 6, 0, 1, 2, 3, 4, 5, 6], dtype='int64') + """ + warnings.warn( + "`dayofweek` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.dayofweek) + + @property + def day_of_week(self) -> Index: + warnings.warn( + "`day_of_week` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return self.dayofweek + + day_of_week.__doc__ = dayofweek.__doc__ + + @property + def weekday(self) -> Index: + warnings.warn( + "`weekday` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.weekday) + + weekday.__doc__ = dayofweek.__doc__ + + @property + def dayofyear(self) -> Index: + """ + The ordinal day of the year. + """ + warnings.warn( + "`dayofyear` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.dayofyear) + + @property + def day_of_year(self) -> Index: + warnings.warn( + "`day_of_year` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return self.dayofyear + + day_of_year.__doc__ = dayofyear.__doc__ + + @property + def quarter(self) -> Index: + """ + The quarter of the date. + """ + warnings.warn( + "`quarter` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.quarter) + + @property + def is_month_start(self) -> Index: + """ + Indicates whether the date is the first day of the month. + + Returns + ------- + Index + Returns a Index with boolean values + + See Also + -------- + is_month_end : Return a boolean indicating whether the date + is the last day of the month. + + Examples + -------- + >>> idx = ps.date_range("2018-02-27", periods=3) # doctest: +SKIP + >>> idx.is_month_start # doctest: +SKIP + Index([False, False, True], dtype='bool') + """ + return Index(self.to_series().dt.is_month_start) + + @property + def is_month_end(self) -> Index: + """ + Indicates whether the date is the last day of the month. + + Returns + ------- + Index + Returns an Index with boolean values. + + See Also + -------- + is_month_start : Return a boolean indicating whether the date + is the first day of the month. + + Examples + -------- + >>> idx = ps.date_range("2018-02-27", periods=3) # doctest: +SKIP + >>> idx.is_month_end # doctest: +SKIP + Index([False, True, False], dtype='bool') + """ + return Index(self.to_series().dt.is_month_end) + + @property + def is_quarter_start(self) -> Index: + """ + Indicator for whether the date is the first day of a quarter. + + Returns + ------- + is_quarter_start : Index + Returns an Index with boolean values. + + See Also + -------- + quarter : Return the quarter of the date. + is_quarter_end : Similar property for indicating the quarter start. + + Examples + -------- + >>> idx = ps.date_range('2017-03-30', periods=4) # doctest: +SKIP + >>> idx.is_quarter_start # doctest: +SKIP + Index([False, False, True, False], dtype='bool') + """ + return Index(self.to_series().dt.is_quarter_start) + + @property + def is_quarter_end(self) -> Index: + """ + Indicator for whether the date is the last day of a quarter. + + Returns + ------- + is_quarter_end : Index + Returns an Index with boolean values. + + See Also + -------- + quarter : Return the quarter of the date. + is_quarter_start : Similar property indicating the quarter start. + + Examples + -------- + >>> idx = ps.date_range('2017-03-30', periods=4) # doctest: +SKIP + >>> idx.is_quarter_end # doctest: +SKIP + Index([False, True, False, False], dtype='bool') + """ + return Index(self.to_series().dt.is_quarter_end) + + @property + def is_year_start(self) -> Index: + """ + Indicate whether the date is the first day of a year. + + Returns + ------- + Index + Returns an Index with boolean values. + + See Also + -------- + is_year_end : Similar property indicating the last day of the year. + + Examples + -------- + >>> idx = ps.date_range("2017-12-30", periods=3) # doctest: +SKIP + >>> idx.is_year_start # doctest: +SKIP + Index([False, False, True], dtype='bool') + """ + return Index(self.to_series().dt.is_year_start) + + @property + def is_year_end(self) -> Index: + """ + Indicate whether the date is the last day of the year. + + Returns + ------- + Index + Returns an Index with boolean values. + + See Also + -------- + is_year_start : Similar property indicating the start of the year. + + Examples + -------- + >>> idx = ps.date_range("2017-12-30", periods=3) # doctest: +SKIP + >>> idx.is_year_end # doctest: +SKIP + Index([False, True, False], dtype='bool') + """ + return Index(self.to_series().dt.is_year_end) + + @property + def is_leap_year(self) -> Index: + """ + Boolean indicator if the date belongs to a leap year. + + A leap year is a year, which has 366 days (instead of 365) including + 29th of February as an intercalary day. + Leap years are years which are multiples of four with the exception + of years divisible by 100 but not by 400. + + Returns + ------- + Index + Booleans indicating if dates belong to a leap year. + + Examples + -------- + >>> idx = ps.date_range("2012-01-01", "2015-01-01", freq="Y") # doctest: +SKIP + >>> idx.is_leap_year # doctest: +SKIP + Index([True, False, False], dtype='bool') + """ + return Index(self.to_series().dt.is_leap_year) + + @property + def daysinmonth(self) -> Index: + """ + The number of days in the month. + """ + warnings.warn( + "`daysinmonth` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.daysinmonth) + + @property + def days_in_month(self) -> Index: + warnings.warn( + "`days_in_month` will return int32 index instead of int 64 index in 4.0.0.", + FutureWarning, + ) + return Index(self.to_series().dt.days_in_month) + + days_in_month.__doc__ = daysinmonth.__doc__ + + # Methods + def ceil(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "DatetimeIndex": + """ + Perform ceil operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to ceil the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + Returns + ------- + DatetimeIndex + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> rng = ps.date_range('1/1/2018 11:59:00', periods=3, freq='min') # doctest: +SKIP + >>> rng.ceil('H') # doctest: +SKIP + DatetimeIndex(['2018-01-01 12:00:00', '2018-01-01 12:00:00', + '2018-01-01 13:00:00'], + dtype='datetime64[ns]', freq=None) + """ + disallow_nanoseconds(freq) + + return DatetimeIndex(self.to_series().dt.ceil(freq, *args, **kwargs)) + + def floor(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "DatetimeIndex": + """ + Perform floor operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to floor the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + Returns + ------- + DatetimeIndex + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> rng = ps.date_range('1/1/2018 11:59:00', periods=3, freq='min') # doctest: +SKIP + >>> rng.floor("H") # doctest: +SKIP + DatetimeIndex(['2018-01-01 11:00:00', '2018-01-01 12:00:00', + '2018-01-01 12:00:00'], + dtype='datetime64[ns]', freq=None) + """ + disallow_nanoseconds(freq) + + return DatetimeIndex(self.to_series().dt.floor(freq, *args, **kwargs)) + + def round(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "DatetimeIndex": + """ + Perform round operation on the data to the specified freq. + + Parameters + ---------- + freq : str or Offset + The frequency level to round the index to. Must be a fixed + frequency like 'S' (second) not 'ME' (month end). + + Returns + ------- + DatetimeIndex + + Raises + ------ + ValueError if the `freq` cannot be converted. + + Examples + -------- + >>> rng = ps.date_range('1/1/2018 11:59:00', periods=3, freq='min') # doctest: +SKIP + >>> rng.round("H") # doctest: +SKIP + DatetimeIndex(['2018-01-01 12:00:00', '2018-01-01 12:00:00', + '2018-01-01 12:00:00'], + dtype='datetime64[ns]', freq=None) + """ + disallow_nanoseconds(freq) + + return DatetimeIndex(self.to_series().dt.round(freq, *args, **kwargs)) + + def month_name(self, locale: Optional[str] = None) -> Index: + """ + Return the month names of the DatetimeIndex with specified locale. + + Parameters + ---------- + locale : str, optional + Locale determining the language in which to return the month name. + Default is English locale. + + Returns + ------- + Index + Index of month names. + + Examples + -------- + >>> idx = ps.date_range(start='2018-01', freq='M', periods=3) # doctest: +SKIP + >>> idx.month_name() # doctest: +SKIP + Index(['January', 'February', 'March'], dtype='object') + """ + return Index(self.to_series().dt.month_name(locale)) + + def day_name(self, locale: Optional[str] = None) -> Index: + """ + Return the day names of the series with specified locale. + + Parameters + ---------- + locale : str, optional + Locale determining the language in which to return the day name. + Default is English locale. + + Returns + ------- + Index + Index of day names. + + Examples + -------- + >>> idx = ps.date_range(start='2018-01-01', freq='D', periods=3) # doctest: +SKIP + >>> idx.day_name() # doctest: +SKIP + Index(['Monday', 'Tuesday', 'Wednesday'], dtype='object') + """ + return Index(self.to_series().dt.day_name(locale)) + + def normalize(self) -> "DatetimeIndex": + """ + Convert times to midnight. + + The time component of the date-time is converted to midnight i.e. + 00:00:00. This is useful in cases, when the time does not matter. + Length is unaltered. The time zones are unaffected. + + This method is available on Series with datetime values under + the ``.dt`` accessor. + + Returns + ------- + DatetimeIndex + The same type as the original data. + + See Also + -------- + floor : Floor the series to the specified freq. + ceil : Ceil the series to the specified freq. + round : Round the series to the specified freq. + + Examples + -------- + >>> idx = ps.date_range(start='2014-08-01 10:00', freq='H', periods=3) # doctest: +SKIP + >>> idx.normalize() # doctest: +SKIP + DatetimeIndex(['2014-08-01', '2014-08-01', '2014-08-01'], dtype='datetime64[ns]', freq=None) + """ + return DatetimeIndex(self.to_series().dt.normalize()) + + def strftime(self, date_format: str) -> Index: + """ + Convert to a string Index using specified date_format. + + Return an Index of formatted strings specified by date_format, which + supports the same string format as the python standard library. Details + of the string format can be found in the python string format + doc. + + Parameters + ---------- + date_format : str + Date format string (example: "%%Y-%%m-%%d"). + + Returns + ------- + Index + Index of formatted strings. + + See Also + -------- + normalize : Return series with times to midnight. + round : Round the series to the specified freq. + floor : Floor the series to the specified freq. + + Examples + -------- + >>> idx = ps.date_range(pd.Timestamp("2018-03-10 09:00"), periods=3, freq='s') + ... # doctest: +SKIP + >>> idx.strftime('%B %d, %Y, %r') # doctest: +SKIP + Index(['March 10, 2018, 09:00:00 AM', 'March 10, 2018, 09:00:01 AM', + 'March 10, 2018, 09:00:02 AM'], + dtype='object') + """ + return Index(self.to_series().dt.strftime(date_format)) + + def indexer_between_time( + self, + start_time: Union[datetime.time, str], + end_time: Union[datetime.time, str], + include_start: bool = True, + include_end: bool = True, + ) -> Index: + """ + Return index locations of values between particular times of day + (example: 9:00-9:30AM). + + Parameters + ---------- + start_time, end_time : datetime.time, str + Time passed either as object (datetime.time) or as string in + appropriate format ("%H:%M", "%H%M", "%I:%M%p", "%I%M%p", + "%H:%M:%S", "%H%M%S", "%I:%M:%S%p","%I%M%S%p"). + include_start : bool, default True + include_end : bool, default True + + Returns + ------- + values_between_time : Index of integers + + Examples + -------- + >>> psidx = ps.date_range("2000-01-01", periods=3, freq="T") # doctest: +SKIP + >>> psidx # doctest: +SKIP + DatetimeIndex(['2000-01-01 00:00:00', '2000-01-01 00:01:00', + '2000-01-01 00:02:00'], + dtype='datetime64[ns]', freq=None) + + >>> psidx.indexer_between_time("00:01", "00:02").sort_values() # doctest: +SKIP + Int64Index([1, 2], dtype='int64') + + >>> psidx.indexer_between_time("00:01", "00:02", include_end=False) # doctest: +SKIP + Int64Index([1], dtype='int64') + + >>> psidx.indexer_between_time("00:01", "00:02", include_start=False) # doctest: +SKIP + Int64Index([2], dtype='int64') + """ + + def pandas_between_time(pdf) -> ps.DataFrame[int]: # type: ignore[no-untyped-def] + return pdf.between_time(start_time, end_time, include_start, include_end) + + psdf = self.to_frame()[[]] + id_column_name = verify_temp_column_name(psdf, "__id_column__") + psdf = psdf.pandas_on_spark.attach_id_column("distributed-sequence", id_column_name) + with ps.option_context("compute.default_index_type", "distributed"): + # The attached index in the statement below will be dropped soon, + # so we enforce “distributed” default index type + psdf = psdf.pandas_on_spark.apply_batch(pandas_between_time) + return ps.Index(first_series(psdf).rename(self.name)) + + def indexer_at_time(self, time: Union[datetime.time, str], asof: bool = False) -> Index: + """ + Return index locations of values at particular time of day + (example: 9:30AM). + + Parameters + ---------- + time : datetime.time or str + Time passed in either as object (datetime.time) or as string in + appropriate format ("%H:%M", "%H%M", "%I:%M%p", "%I%M%p", + "%H:%M:%S", "%H%M%S", "%I:%M:%S%p", "%I%M%S%p"). + + Returns + ------- + values_at_time : Index of integers + + Examples + -------- + >>> psidx = ps.date_range("2000-01-01", periods=3, freq="T") # doctest: +SKIP + >>> psidx # doctest: +SKIP + DatetimeIndex(['2000-01-01 00:00:00', '2000-01-01 00:01:00', + '2000-01-01 00:02:00'], + dtype='datetime64[ns]', freq=None) + + >>> psidx.indexer_at_time("00:00") # doctest: +SKIP + Int64Index([0], dtype='int64') + + >>> psidx.indexer_at_time("00:01") # doctest: +SKIP + Int64Index([1], dtype='int64') + """ + if asof: + raise NotImplementedError("'asof' argument is not supported") + + def pandas_at_time(pdf) -> ps.DataFrame[int]: # type: ignore[no-untyped-def] + return pdf.at_time(time, asof) + + psdf = self.to_frame()[[]] + id_column_name = verify_temp_column_name(psdf, "__id_column__") + psdf = psdf.pandas_on_spark.attach_id_column("distributed-sequence", id_column_name) + with ps.option_context("compute.default_index_type", "distributed"): + # The attached index in the statement below will be dropped soon, + # so we enforce “distributed” default index type + psdf = psdf.pandas_on_spark.apply_batch(pandas_at_time) + return ps.Index(first_series(psdf).rename(self.name)) + + @no_type_check + def all(self, *args, **kwargs) -> None: + raise TypeError("Cannot perform 'all' with this index type: %s" % type(self).__name__) + + +def disallow_nanoseconds(freq: Union[str, DateOffset]) -> None: + if freq in ["N", "ns"]: + raise ValueError("nanoseconds is not supported") + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.indexes.datetimes + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexes.datetimes.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexes.datetimes tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexes.datetimes, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/multi.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/multi.py new file mode 100644 index 0000000..dd93e31 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/multi.py @@ -0,0 +1,1323 @@ +# +# 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 functools import partial, reduce +from typing import Any, Callable, Iterator, List, Optional, Tuple, Union, cast, no_type_check + +import pandas as pd +from pandas.api.types import is_hashable, is_list_like # type: ignore[attr-defined] + +from pyspark.sql import functions as F, Column as PySparkColumn, Window +from pyspark.sql.types import DataType +from pyspark.sql.utils import get_column_class + +# For running doctests and reference resolution in PyCharm. +from pyspark import pandas as ps +from pyspark.pandas._typing import Label, Name, Scalar +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.missing.indexes import MissingPandasLikeMultiIndex +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.utils import ( + compare_disallow_null, + is_name_like_tuple, + name_like_string, + scol_for, + verify_temp_column_name, + validate_index_loc, +) +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, +) + + +class MultiIndex(Index): + """ + pandas-on-Spark MultiIndex that corresponds to pandas MultiIndex logically. This might hold + Spark Column internally. + + Parameters + ---------- + levels : sequence of arrays + The unique labels for each level. + codes : sequence of arrays + Integers for each level designating which label at each location. + sortorder : optional int + Level of sortedness (must be lexicographically sorted by that + level). + names : optional sequence of objects + Names for each of the index levels. (name is accepted for compat). + copy : bool, default False + Copy the meta-data. + verify_integrity : bool, default True + Check that the levels/codes are consistent and valid. + + See Also + -------- + MultiIndex.from_arrays : Convert list of arrays to MultiIndex. + MultiIndex.from_product : Create a MultiIndex from the cartesian product + of iterables. + MultiIndex.from_tuples : Convert list of tuples to a MultiIndex. + MultiIndex.from_frame : Make a MultiIndex from a DataFrame. + Index : A single-level Index. + + Examples + -------- + >>> ps.DataFrame({'a': ['a', 'b', 'c']}, index=[[1, 2, 3], [4, 5, 6]]).index # doctest: +SKIP + MultiIndex([(1, 4), + (2, 5), + (3, 6)], + ) + + >>> ps.DataFrame({'a': [1, 2, 3]}, index=[list('abc'), list('def')]).index # doctest: +SKIP + MultiIndex([('a', 'd'), + ('b', 'e'), + ('c', 'f')], + ) + """ + + @no_type_check + def __new__( + cls, + levels=None, + codes=None, + sortorder=None, + names=None, + dtype=None, + copy=False, + name=None, + verify_integrity: bool = True, + ) -> "MultiIndex": + pidx = pd.MultiIndex( + levels=levels, + codes=codes, + sortorder=sortorder, + names=names, + dtype=dtype, + copy=copy, + name=name, + verify_integrity=verify_integrity, + ) + return ps.from_pandas(pidx) + + @property + def _internal(self) -> InternalFrame: + internal = self._psdf._internal + scol = F.struct(*internal.index_spark_columns) + return internal.copy( + column_labels=[None], + data_spark_columns=[scol], + data_fields=[None], + column_label_names=None, + ) + + @property + def _column_label(self) -> Optional[Label]: + return None + + def __abs__(self) -> "MultiIndex": + raise TypeError("TypeError: cannot perform __abs__ with this index type: MultiIndex") + + def _with_new_scol( + self, scol: PySparkColumn, *, field: Optional[InternalField] = None + ) -> "MultiIndex": + raise NotImplementedError("Not supported for type MultiIndex") + + @no_type_check + def any(self, *args, **kwargs) -> None: + raise TypeError("cannot perform any with this index type: MultiIndex") + + @no_type_check + def all(self, *args, **kwargs) -> None: + raise TypeError("cannot perform all with this index type: MultiIndex") + + @staticmethod + def from_tuples( + tuples: List[Tuple], + sortorder: Optional[int] = None, + names: Optional[List[Name]] = None, + ) -> "MultiIndex": + """ + Convert list of tuples to MultiIndex. + + Parameters + ---------- + tuples : list / sequence of tuple-likes + Each tuple is the index of one row/column. + sortorder : int or None + Level of sortedness (must be lexicographically sorted by that level). + names : list / sequence of str, optional + Names for the levels in the index. + + Returns + ------- + index : MultiIndex + + Examples + -------- + + >>> tuples = [(1, 'red'), (1, 'blue'), + ... (2, 'red'), (2, 'blue')] + >>> ps.MultiIndex.from_tuples(tuples, names=('number', 'color')) # doctest: +SKIP + MultiIndex([(1, 'red'), + (1, 'blue'), + (2, 'red'), + (2, 'blue')], + names=['number', 'color']) + """ + return cast( + MultiIndex, + ps.from_pandas( + pd.MultiIndex.from_tuples(tuples=tuples, sortorder=sortorder, names=names) + ), + ) + + @staticmethod + def from_arrays( + arrays: List[List], + sortorder: Optional[int] = None, + names: Optional[List[Name]] = None, + ) -> "MultiIndex": + """ + Convert arrays to MultiIndex. + + Parameters + ---------- + arrays: list / sequence of array-likes + Each array-like gives one level’s value for each data point. len(arrays) + is the number of levels. + sortorder: int or None + Level of sortedness (must be lexicographically sorted by that level). + names: list / sequence of str, optional + Names for the levels in the index. + + Returns + ------- + index: MultiIndex + + Examples + -------- + + >>> arrays = [[1, 1, 2, 2], ['red', 'blue', 'red', 'blue']] + >>> ps.MultiIndex.from_arrays(arrays, names=('number', 'color')) # doctest: +SKIP + MultiIndex([(1, 'red'), + (1, 'blue'), + (2, 'red'), + (2, 'blue')], + names=['number', 'color']) + """ + return cast( + MultiIndex, + ps.from_pandas( + pd.MultiIndex.from_arrays(arrays=arrays, sortorder=sortorder, names=names) + ), + ) + + @staticmethod + def from_product( + iterables: List[List], + sortorder: Optional[int] = None, + names: Optional[List[Name]] = None, + ) -> "MultiIndex": + """ + Make a MultiIndex from the cartesian product of multiple iterables. + + Parameters + ---------- + iterables : list / sequence of iterables + Each iterable has unique labels for each level of the index. + sortorder : int or None + Level of sortedness (must be lexicographically sorted by that + level). + names : list / sequence of str, optional + Names for the levels in the index. + + Returns + ------- + index : MultiIndex + + See Also + -------- + MultiIndex.from_arrays : Convert list of arrays to MultiIndex. + MultiIndex.from_tuples : Convert list of tuples to MultiIndex. + + Examples + -------- + >>> numbers = [0, 1, 2] + >>> colors = ['green', 'purple'] + >>> ps.MultiIndex.from_product([numbers, colors], + ... names=['number', 'color']) # doctest: +SKIP + MultiIndex([(0, 'green'), + (0, 'purple'), + (1, 'green'), + (1, 'purple'), + (2, 'green'), + (2, 'purple')], + names=['number', 'color']) + """ + return cast( + MultiIndex, + ps.from_pandas( + pd.MultiIndex.from_product(iterables=iterables, sortorder=sortorder, names=names) + ), + ) + + @staticmethod + def from_frame(df: DataFrame, names: Optional[List[Name]] = None) -> "MultiIndex": + """ + Make a MultiIndex from a DataFrame. + + Parameters + ---------- + df : DataFrame + DataFrame to be converted to MultiIndex. + names : list-like, optional + If no names are provided, use the column names, or tuple of column + names if the column is a MultiIndex. If a sequence, overwrite + names with the given sequence. + + Returns + ------- + MultiIndex + The MultiIndex representation of the given DataFrame. + + See Also + -------- + MultiIndex.from_arrays : Convert list of arrays to MultiIndex. + MultiIndex.from_tuples : Convert list of tuples to MultiIndex. + MultiIndex.from_product : Make a MultiIndex from cartesian product + of iterables. + + Examples + -------- + >>> df = ps.DataFrame([['HI', 'Temp'], ['HI', 'Precip'], + ... ['NJ', 'Temp'], ['NJ', 'Precip']], + ... columns=['a', 'b']) + >>> df # doctest: +SKIP + a b + 0 HI Temp + 1 HI Precip + 2 NJ Temp + 3 NJ Precip + + >>> ps.MultiIndex.from_frame(df) # doctest: +SKIP + MultiIndex([('HI', 'Temp'), + ('HI', 'Precip'), + ('NJ', 'Temp'), + ('NJ', 'Precip')], + names=['a', 'b']) + + Using explicit names, instead of the column names + + >>> ps.MultiIndex.from_frame(df, names=['state', 'observation']) # doctest: +SKIP + MultiIndex([('HI', 'Temp'), + ('HI', 'Precip'), + ('NJ', 'Temp'), + ('NJ', 'Precip')], + names=['state', 'observation']) + """ + if not isinstance(df, DataFrame): + raise TypeError("Input must be a DataFrame") + sdf = df._to_spark() + + if names is None: + names = df._internal.column_labels + elif not is_list_like(names): + raise TypeError("Names should be list-like for a MultiIndex") + else: + names = [name if is_name_like_tuple(name) else (name,) for name in names] + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in sdf.columns], + index_names=names, + ) + return cast(MultiIndex, DataFrame(internal).index) + + @property + def name(self) -> Name: + raise PandasNotImplementedError(class_name="pd.MultiIndex", property_name="name") + + @name.setter + def name(self, name: Name) -> None: + raise PandasNotImplementedError(class_name="pd.MultiIndex", property_name="name") + + @property + def dtypes(self) -> pd.Series: + """Return the dtypes as a Series for the underlying MultiIndex. + + .. versionadded:: 3.3.0 + + Returns + ------- + pd.Series + The data type of each level. + + Examples + -------- + >>> psmidx = ps.MultiIndex.from_arrays( + ... [[0, 1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8, 9]], + ... names=("zero", "one"), + ... ) + >>> psmidx.dtypes + zero int64 + one int64 + dtype: object + """ + return pd.Series( + [field.dtype for field in self._internal.index_fields], + index=pd.Index( + [name if len(name) > 1 else name[0] for name in self._internal.index_names] + ), + ) + + def _verify_for_rename(self, name: List[Name]) -> List[Label]: # type: ignore[override] + if is_list_like(name): + if self._internal.index_level != len(name): + raise ValueError( + "Length of new names must be {}, got {}".format( + self._internal.index_level, len(name) + ) + ) + if any(not is_hashable(n) for n in name): + raise TypeError("MultiIndex.name must be a hashable type") + return [n if is_name_like_tuple(n) else (n,) for n in name] + else: + raise TypeError("Must pass list-like as `names`.") + + def swaplevel(self, i: int = -2, j: int = -1) -> "MultiIndex": + """ + Swap level i with level j. + Calling this method does not change the ordering of the values. + + Parameters + ---------- + i : int, str, default -2 + First level of index to be swapped. Can pass level name as string. + Parameter types can be mixed. + j : int, str, default -1 + Second level of index to be swapped. Can pass level name as string. + Parameter types can be mixed. + + Returns + ------- + MultiIndex + A new MultiIndex. + + Examples + -------- + >>> midx = ps.MultiIndex.from_arrays([['a', 'b'], [1, 2]], names = ['word', 'number']) + >>> midx # doctest: +SKIP + MultiIndex([('a', 1), + ('b', 2)], + names=['word', 'number']) + + >>> midx.swaplevel(0, 1) # doctest: +SKIP + MultiIndex([(1, 'a'), + (2, 'b')], + names=['number', 'word']) + + >>> midx.swaplevel('number', 'word') # doctest: +SKIP + MultiIndex([(1, 'a'), + (2, 'b')], + names=['number', 'word']) + """ + for index in (i, j): + if not isinstance(index, int) and index not in self.names: + raise KeyError("Level %s not found" % index) + + i = i if isinstance(i, int) else self.names.index(i) + j = j if isinstance(j, int) else self.names.index(j) + + for index in (i, j): + if index >= len(self.names) or index < -len(self.names): + raise IndexError( + "Too many levels: Index has only %s levels, " + "%s is not a valid level number" % (len(self.names), index) + ) + + index_map = list( + zip( + self._internal.index_spark_columns, + self._internal.index_names, + self._internal.index_fields, + ) + ) + index_map[i], index_map[j] = index_map[j], index_map[i] + index_spark_columns, index_names, index_fields = zip(*index_map) + internal = self._internal.copy( + index_spark_columns=list(index_spark_columns), + index_names=list(index_names), + index_fields=list(index_fields), + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return cast(MultiIndex, DataFrame(internal).index) + + @property + def levshape(self) -> Tuple[int, ...]: + """ + A tuple with the length of each level. + + Examples + -------- + >>> midx = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> midx # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z')], + ) + + >>> midx.levshape + (3, 3) + """ + result = self._internal.spark_frame.agg( + *(F.countDistinct(c) for c in self._internal.index_spark_columns) + ).collect()[0] + return tuple(result) + + @staticmethod + def _comparator_for_monotonic_increasing( + data_type: DataType, + ) -> Callable[ + [PySparkColumn, PySparkColumn, Callable[[PySparkColumn, PySparkColumn], PySparkColumn]], + PySparkColumn, + ]: + return compare_disallow_null + + def _is_monotonic(self, order: str) -> bool: + if order == "increasing": + return self._is_monotonic_increasing().all() + else: + return self._is_monotonic_decreasing().all() + + def _is_monotonic_increasing(self) -> Series: + window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-1, -1) + + cond = F.lit(True) + has_not_null = F.lit(True) + Column = get_column_class() + for scol in self._internal.index_spark_columns[::-1]: + data_type = self._internal.spark_type_for(scol) + prev = F.lag(scol, 1).over(window) + compare = MultiIndex._comparator_for_monotonic_increasing(data_type) + # Since pandas 1.1.4, null value is not allowed at any levels of MultiIndex. + # Therefore, we should check `has_not_null` over all levels. + has_not_null = has_not_null & scol.isNotNull() + cond = F.when(scol.eqNullSafe(prev), cond).otherwise(compare(scol, prev, Column.__gt__)) + + cond = has_not_null & (prev.isNull() | cond) + + cond_name = verify_temp_column_name( + self._internal.spark_frame.select(self._internal.index_spark_columns), + "__is_monotonic_increasing_cond__", + ) + + sdf = self._internal.spark_frame.select( + self._internal.index_spark_columns + [cond.alias(cond_name)] + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + + return first_series(DataFrame(internal)) + + @staticmethod + def _comparator_for_monotonic_decreasing( + data_type: DataType, + ) -> Callable[ + [PySparkColumn, PySparkColumn, Callable[[PySparkColumn, PySparkColumn], PySparkColumn]], + PySparkColumn, + ]: + return compare_disallow_null + + def _is_monotonic_decreasing(self) -> Series: + window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-1, -1) + + cond = F.lit(True) + has_not_null = F.lit(True) + Column = get_column_class() + for scol in self._internal.index_spark_columns[::-1]: + data_type = self._internal.spark_type_for(scol) + prev = F.lag(scol, 1).over(window) + compare = MultiIndex._comparator_for_monotonic_increasing(data_type) + # Since pandas 1.1.4, null value is not allowed at any levels of MultiIndex. + # Therefore, we should check `has_not_null` over all levels. + has_not_null = has_not_null & scol.isNotNull() + cond = F.when(scol.eqNullSafe(prev), cond).otherwise(compare(scol, prev, Column.__lt__)) + + cond = has_not_null & (prev.isNull() | cond) + + cond_name = verify_temp_column_name( + self._internal.spark_frame.select(self._internal.index_spark_columns), + "__is_monotonic_decreasing_cond__", + ) + + sdf = self._internal.spark_frame.select( + self._internal.index_spark_columns + [cond.alias(cond_name)] + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + + return first_series(DataFrame(internal)) + + def to_frame( # type: ignore[override] + self, index: bool = True, name: Optional[List[Name]] = None + ) -> DataFrame: + """ + Create a DataFrame with the levels of the MultiIndex as columns. + Column ordering is determined by the DataFrame constructor with data as + a dict. + + Parameters + ---------- + index : boolean, default True + Set the index of the returned DataFrame as the original MultiIndex. + name : list / sequence of strings, optional + The passed names should substitute index level names. + + Returns + ------- + DataFrame : a DataFrame containing the original MultiIndex data. + + See Also + -------- + DataFrame + + Examples + -------- + >>> tuples = [(1, 'red'), (1, 'blue'), + ... (2, 'red'), (2, 'blue')] + >>> idx = ps.MultiIndex.from_tuples(tuples, names=('number', 'color')) + >>> idx # doctest: +SKIP + MultiIndex([(1, 'red'), + (1, 'blue'), + (2, 'red'), + (2, 'blue')], + names=['number', 'color']) + >>> idx.to_frame() # doctest: +NORMALIZE_WHITESPACE + number color + number color + 1 red 1 red + blue 1 blue + 2 red 2 red + blue 2 blue + + By default, the original Index is reused. To enforce a new Index: + + >>> idx.to_frame(index=False) + number color + 0 1 red + 1 1 blue + 2 2 red + 3 2 blue + + To override the name of the resulting column, specify `name`: + + >>> idx.to_frame(name=['n', 'c']) # doctest: +NORMALIZE_WHITESPACE + n c + number color + 1 red 1 red + blue 1 blue + 2 red 2 red + blue 2 blue + """ + if name is None: + name = [ + name if name is not None else (i,) + for i, name in enumerate(self._internal.index_names) + ] + elif is_list_like(name): + if len(name) != self._internal.index_level: + raise ValueError("'name' should have same length as number of levels on index.") + name = [n if is_name_like_tuple(n) else (n,) for n in name] + else: + raise TypeError("'name' must be a list / sequence of column names.") + + return self._to_frame(index=index, names=name) + + def to_pandas(self) -> pd.MultiIndex: + """ + Return a pandas MultiIndex. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats'], + ... index=[list('abcd'), list('efgh')]) + >>> df['dogs'].index.to_pandas() # doctest: +SKIP + MultiIndex([('a', 'e'), + ('b', 'f'), + ('c', 'g'), + ('d', 'h')], + ) + """ + # TODO: We might need to handle internal state change. + # So far, we don't have any functions to change the internal state of MultiIndex except for + # series-like operations. In that case, it creates a new Index object instead of MultiIndex. + return cast(pd.MultiIndex, super().to_pandas()) + + def _to_pandas(self) -> pd.MultiIndex: + """ + Same as `to_pandas()`, without issuing the advice log for internal usage. + """ + return cast(pd.MultiIndex, super()._to_pandas()) + + def nunique(self, dropna: bool = True, approx: bool = False, rsd: float = 0.05) -> int: + raise NotImplementedError("nunique is not defined for MultiIndex") + + # TODO: add 'name' parameter after pd.MultiIndex.name is implemented + def copy(self, deep: Optional[bool] = None) -> "MultiIndex": # type: ignore[override] + """ + Make a copy of this object. + + Parameters + ---------- + deep : None + this parameter is not supported but just dummy parameter to match pandas. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats'], + ... index=[list('abcd'), list('efgh')]) + >>> df['dogs'].index # doctest: +SKIP + MultiIndex([('a', 'e'), + ('b', 'f'), + ('c', 'g'), + ('d', 'h')], + ) + + Copy index + + >>> df.index.copy() # doctest: +SKIP + MultiIndex([('a', 'e'), + ('b', 'f'), + ('c', 'g'), + ('d', 'h')], + ) + """ + return cast(MultiIndex, super().copy(deep=deep)) + + def symmetric_difference( # type: ignore[override] + self, + other: Index, + result_name: Optional[List[Name]] = None, + sort: Optional[bool] = None, + ) -> "MultiIndex": + """ + Compute the symmetric difference of two MultiIndex objects. + + Parameters + ---------- + other : Index or array-like + result_name : list + sort : True or None, default None + Whether to sort the resulting index. + * True : Attempt to sort the result. + * None : Do not sort the result. + + Returns + ------- + symmetric_difference : MultiIndex + + Notes + ----- + ``symmetric_difference`` contains elements that appear in either + ``idx1`` or ``idx2`` but not both. Equivalent to the Index created by + ``idx1.difference(idx2) | idx2.difference(idx1)`` with duplicates + dropped. + + Examples + -------- + >>> midx1 = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 0, 0, 0, 1, 2, 0, 1, 2]]) + >>> midx2 = pd.MultiIndex([['pandas-on-Spark', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 0, 0, 0, 1, 2, 0, 1, 2]]) + >>> s1 = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx1) + >>> s2 = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx2) + + >>> s1.index.symmetric_difference(s2.index) # doctest: +SKIP + MultiIndex([('pandas-on-Spark', 'speed'), + ( 'lama', 'speed')], + ) + + You can set names of the result Index. + + >>> s1.index.symmetric_difference(s2.index, result_name=['a', 'b']) # doctest: +SKIP + MultiIndex([('pandas-on-Spark', 'speed'), + ( 'lama', 'speed')], + names=['a', 'b']) + + You can set sort to `True`, if you want to sort the resulting index. + + >>> s1.index.symmetric_difference(s2.index, sort=True) # doctest: +SKIP + MultiIndex([('pandas-on-Spark', 'speed'), + ( 'lama', 'speed')], + ) + + You can also use the ``^`` operator: + + >>> s1.index ^ s2.index # doctest: +SKIP + MultiIndex([('pandas-on-Spark', 'speed'), + ( 'lama', 'speed')], + ) + """ + if type(self) != type(other): + raise NotImplementedError( + "Doesn't support symmetric_difference between Index & MultiIndex for now" + ) + + sdf_self = self._psdf._internal.spark_frame.select(self._internal.index_spark_columns) + sdf_other = other._psdf._internal.spark_frame.select(other._internal.index_spark_columns) + + sdf_symdiff = sdf_self.union(sdf_other).subtract(sdf_self.intersect(sdf_other)) + + if sort: + sdf_symdiff = sdf_symdiff.sort(*self._internal.index_spark_columns) + + internal = InternalFrame( + spark_frame=sdf_symdiff, + index_spark_columns=[ + scol_for(sdf_symdiff, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=self._internal.index_fields, + ) + result = cast(MultiIndex, DataFrame(internal).index) + + if result_name: + result.names = result_name + + return result + + # TODO: ADD error parameter + def drop(self, codes: List[Any], level: Optional[Union[int, Name]] = None) -> "MultiIndex": + """ + Make new MultiIndex with passed list of labels deleted + + Parameters + ---------- + codes : array-like + Must be a list of tuples + level : int or level name, default None + + Returns + ------- + dropped : MultiIndex + + Examples + -------- + >>> index = ps.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')]) + >>> index # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z')], + ) + + >>> index.drop(['a']) # doctest: +SKIP + MultiIndex([('b', 'y'), + ('c', 'z')], + ) + + >>> index.drop(['x', 'y'], level=1) # doctest: +SKIP + MultiIndex([('c', 'z')], + ) + """ + internal = self._internal.resolved_copy + sdf = internal.spark_frame + index_scols = internal.index_spark_columns + if level is None: + scol = index_scols[0] + elif isinstance(level, int): + scol = index_scols[level] + else: + scol = None + for index_spark_column, index_name in zip( + internal.index_spark_columns, internal.index_names + ): + if not isinstance(level, tuple): + level = (level,) + if level == index_name: + if scol is not None: + raise ValueError( + "The name {} occurs multiple times, use a level number".format( + name_like_string(level) + ) + ) + scol = index_spark_column + if scol is None: + raise KeyError("Level {} not found".format(name_like_string(level))) + sdf = sdf[~scol.isin(codes)] + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in internal.index_spark_column_names], + index_names=internal.index_names, + index_fields=internal.index_fields, + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return cast(MultiIndex, DataFrame(internal).index) + + def drop_duplicates(self, keep: Union[bool, str] = "first") -> "MultiIndex": + """ + Return MultiIndex with duplicate values removed. + + Parameters + ---------- + keep : {'first', 'last', ``False``}, default 'first' + Method to handle dropping duplicates: + - 'first' : Drop duplicates except for the first occurrence. + - 'last' : Drop duplicates except for the last occurrence. + - ``False`` : Drop all duplicates. + + Returns + ------- + deduplicated : MultiIndex + + See Also + -------- + Series.drop_duplicates : Equivalent method on Series. + DataFrame.drop_duplicates : Equivalent method on DataFrame. + + Examples + -------- + Generate a MultiIndex with duplicate values. + + >>> arrays = [[1, 2, 3, 1, 2], ["red", "blue", "black", "red", "blue"]] + >>> midx = ps.MultiIndex.from_arrays(arrays, names=("number", "color")) + >>> midx + MultiIndex([(1, 'red'), + (2, 'blue'), + (3, 'black'), + (1, 'red'), + (2, 'blue')], + names=['number', 'color']) + + >>> midx.drop_duplicates() + MultiIndex([(1, 'red'), + (2, 'blue'), + (3, 'black')], + names=['number', 'color']) + + >>> midx.drop_duplicates(keep='first') + MultiIndex([(1, 'red'), + (2, 'blue'), + (3, 'black')], + names=['number', 'color']) + + >>> midx.drop_duplicates(keep='last') + MultiIndex([(3, 'black'), + (1, 'red'), + (2, 'blue')], + names=['number', 'color']) + + >>> midx.drop_duplicates(keep=False) + MultiIndex([(3, 'black')], + names=['number', 'color']) + """ + with ps.option_context("compute.default_index_type", "distributed"): + # The attached index caused by `reset_index` below is used for sorting only, + # and it will be dropped soon, + # so we enforce “distributed” default index type + psdf = self.to_frame().reset_index(drop=True) + return ps.MultiIndex.from_frame(psdf.drop_duplicates(keep=keep).sort_index()) + + def argmax(self) -> None: + raise TypeError("reduction operation 'argmax' not allowed for this dtype") + + def argmin(self) -> None: + raise TypeError("reduction operation 'argmin' not allowed for this dtype") + + def asof(self, label: Any) -> None: + raise NotImplementedError( + "only the default get_loc method is currently supported for MultiIndex" + ) + + @property + def is_all_dates(self) -> bool: + """ + is_all_dates always returns False for MultiIndex + + Examples + -------- + >>> from datetime import datetime + + >>> idx = ps.MultiIndex.from_tuples( + ... [(datetime(2019, 1, 1, 0, 0, 0), datetime(2019, 1, 1, 0, 0, 0)), + ... (datetime(2019, 1, 1, 0, 0, 0), datetime(2019, 1, 1, 0, 0, 0))]) + >>> idx # doctest: +SKIP + MultiIndex([('2019-01-01', '2019-01-01'), + ('2019-01-01', '2019-01-01')], + ) + + >>> idx.is_all_dates + False + """ + return False + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeMultiIndex, item): + property_or_func = getattr(MissingPandasLikeMultiIndex, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError("'MultiIndex' object has no attribute '{}'".format(item)) + + def _get_level_number(self, level: Union[int, Name]) -> int: + """ + Return the level number if a valid level is given. + """ + count = self.names.count(level) + if (count > 1) and not isinstance(level, int): + raise ValueError("The name %s occurs multiple times, use a level number" % level) + if level in self.names: + level = self.names.index(level) + elif isinstance(level, int): + nlevels = self.nlevels + if level >= nlevels: + raise IndexError( + "Too many levels: Index has only %d " + "levels, %d is not a valid level number" % (nlevels, level) + ) + if level < 0: + if (level + nlevels) < 0: + raise IndexError( + "Too many levels: Index has only %d levels, " + "not %d" % (nlevels, level + 1) + ) + level = level + nlevels + else: + raise KeyError("Level %s not found" % str(level)) + + return level + + def get_level_values(self, level: Union[int, Name]) -> Index: + """ + Return vector of label values for requested level, + equal to the length of the index. + + Parameters + ---------- + level : int or str + ``level`` is either the integer position of the level in the + MultiIndex, or the name of the level. + + Returns + ------- + values : Index + Values is a level of this MultiIndex converted to + a single :class:`Index` (or subclass thereof). + + Examples + -------- + + Create a MultiIndex: + + >>> mi = ps.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'a')]) + >>> mi.names = ['level_1', 'level_2'] + + Get level values by supplying level as either integer or name: + + >>> mi.get_level_values(0) + Index(['x', 'x', 'y'], dtype='object', name='level_1') + + >>> mi.get_level_values('level_2') + Index(['a', 'b', 'a'], dtype='object', name='level_2') + """ + level = self._get_level_number(level) + index_scol = self._internal.index_spark_columns[level] + index_name = self._internal.index_names[level] + index_field = self._internal.index_fields[level] + internal = self._internal.copy( + index_spark_columns=[index_scol], + index_names=[index_name], + index_fields=[index_field], + column_labels=[], + data_spark_columns=[], + data_fields=[], + ) + return DataFrame(internal).index + + def insert(self, loc: int, item: Any) -> Index: + """ + Make new MultiIndex inserting new item at location. + + Follows Python list.append semantics for negative values. + + .. versionchanged:: 3.4.0 + Raise IndexError when loc is out of bounds to follow Pandas 1.4+ behavior + + Parameters + ---------- + loc : int + item : object + + Returns + ------- + new_index : MultiIndex + + Examples + -------- + >>> psmidx = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + >>> psmidx.insert(3, ("h", "j")) # doctest: +SKIP + MultiIndex([('a', 'x'), + ('b', 'y'), + ('c', 'z'), + ('h', 'j')], + ) + + For negative values + + >>> psmidx.insert(-2, ("h", "j")) # doctest: +SKIP + MultiIndex([('a', 'x'), + ('h', 'j'), + ('b', 'y'), + ('c', 'z')], + ) + """ + validate_index_loc(self, loc) + loc = loc + len(self) if loc < 0 else loc + + index_name: List[Label] = [(name,) for name in self._internal.index_spark_column_names] + sdf_before = self.to_frame(name=index_name)[:loc]._to_spark() + sdf_middle = Index([item]).to_frame(name=index_name)._to_spark() + sdf_after = self.to_frame(name=index_name)[loc:]._to_spark() + sdf = sdf_before.union(sdf_middle).union(sdf_after) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=[InternalField(field.dtype) for field in self._internal.index_fields], + ) + return DataFrame(internal).index + + def item(self) -> Tuple[Scalar, ...]: + """ + Return the first element of the underlying data as a python tuple. + + Returns + ------- + tuple + The first element of MultiIndex. + + Raises + ------ + ValueError + If the data is not length-1. + + Examples + -------- + >>> psmidx = ps.MultiIndex.from_tuples([('a', 'x')]) + >>> psmidx.item() + ('a', 'x') + """ + return self._psdf.head(2)._to_internal_pandas().index.item() + + def intersection(self, other: Union[DataFrame, Series, Index, List]) -> "MultiIndex": + """ + Form the intersection of two Index objects. + + This returns a new Index with elements common to the index and `other`. + + Parameters + ---------- + other : Index or array-like + + Returns + ------- + intersection : MultiIndex + + Examples + -------- + >>> midx1 = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + >>> midx2 = ps.MultiIndex.from_tuples([("c", "z"), ("d", "w")]) + >>> midx1.intersection(midx2).sort_values() # doctest: +SKIP + MultiIndex([('c', 'z')], + ) + """ + if isinstance(other, Series) or not is_list_like(other): + raise TypeError("other must be a MultiIndex or a list of tuples") + elif isinstance(other, DataFrame): + raise ValueError("Index data must be 1-dimensional") + elif isinstance(other, MultiIndex): + spark_frame_other = other.to_frame()._to_spark() + keep_name = self.names == other.names + elif isinstance(other, Index): + # Always returns an empty MultiIndex if `other` is Index. + return cast(MultiIndex, self.to_frame().head(0).index) + elif not all(isinstance(item, tuple) for item in other): + raise TypeError("other must be a MultiIndex or a list of tuples") + else: + other = MultiIndex.from_tuples(list(other)) + spark_frame_other = cast(MultiIndex, other).to_frame()._to_spark() + keep_name = True + + index_fields = self._index_fields_for_union_like(other, func_name="intersection") + + default_name: List[Name] = [SPARK_INDEX_NAME_FORMAT(i) for i in range(self.nlevels)] + spark_frame_self = self.to_frame(name=default_name)._to_spark() + spark_frame_intersected = spark_frame_self.intersect(spark_frame_other) + if keep_name: + index_names = self._internal.index_names + else: + index_names = None + + internal = InternalFrame( + spark_frame=spark_frame_intersected, + index_spark_columns=[ + scol_for(spark_frame_intersected, cast(str, col)) for col in default_name + ], + index_names=index_names, + index_fields=index_fields, + ) + return cast(MultiIndex, DataFrame(internal).index) + + def equal_levels(self, other: "MultiIndex") -> bool: + """ + Return True if the levels of both MultiIndex objects are the same + + .. versionadded:: 3.3.0 + + Examples + -------- + >>> psmidx1 = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + >>> psmidx2 = ps.MultiIndex.from_tuples([("b", "y"), ("a", "x"), ("c", "z")]) + >>> psmidx1.equal_levels(psmidx2) + True + + >>> psmidx2 = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "j")]) + >>> psmidx1.equal_levels(psmidx2) + False + """ + nlevels = self.nlevels + if nlevels != other.nlevels: + return False + + self_sdf = self._internal.spark_frame + other_sdf = other._internal.spark_frame + subtract_list = [] + for nlevel in range(nlevels): + self_index_scol = self._internal.index_spark_columns[nlevel] + other_index_scol = other._internal.index_spark_columns[nlevel] + self_subtract_other = self_sdf.select(self_index_scol).subtract( + other_sdf.select(other_index_scol) + ) + subtract_list.append(self_subtract_other) + + unioned_subtracts = reduce(lambda x, y: x.union(y), subtract_list) + return len(unioned_subtracts.head(1)) == 0 + + @property + def hasnans(self) -> bool: + raise NotImplementedError("hasnans is not defined for MultiIndex") + + @property + def inferred_type(self) -> str: + """ + Return a string of the type inferred from the values. + """ + # Always returns "mixed" for MultiIndex + return "mixed" + + @property + def asi8(self) -> None: + """ + Integer representation of the values. + """ + # Always returns None for MultiIndex + return None + + def factorize( + self, sort: bool = True, na_sentinel: Optional[int] = -1 + ) -> Tuple["MultiIndex", pd.Index]: + return MissingPandasLikeMultiIndex.factorize(self, sort=sort, na_sentinel=na_sentinel) + + def __iter__(self) -> Iterator: + return MissingPandasLikeMultiIndex.__iter__(self) + + def map( + self, + mapper: Union[dict, Callable[[Any], Any], pd.Series] = None, + na_action: Optional[str] = None, + ) -> "Index": + return MissingPandasLikeMultiIndex.map(self, mapper, na_action) + + +def _test() -> None: + import os + import doctest + import sys + import numpy + from pyspark.sql import SparkSession + import pyspark.pandas.indexes.multi + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexes.multi.__dict__.copy() + globs["np"] = numpy + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexes.multi tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexes.multi, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/numeric.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/numeric.py new file mode 100644 index 0000000..d0b5bc5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/numeric.py @@ -0,0 +1,210 @@ +# +# 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 warnings +from typing import Any, Optional, Union, cast + +import pandas as pd +from pandas.api.types import is_hashable # type: ignore[attr-defined] + +from pyspark import pandas as ps +from pyspark.pandas._typing import Dtype, Name +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.series import Series + + +class NumericIndex(Index): + """ + Provide numeric type operations. + This is an abstract class. + """ + + pass + + +class IntegerIndex(NumericIndex): + """ + This is an abstract class for Int64Index. + """ + + pass + + +class Int64Index(IntegerIndex): + """ + Immutable sequence used for indexing and alignment. The basic object + storing axis labels for all pandas objects. Int64Index is a special case + of `Index` with purely integer labels. + + .. deprecated:: 3.4.0 + + Parameters + ---------- + data : array-like (1-dimensional) + dtype : NumPy dtype (default: int64) + copy : bool + Make a copy of input ndarray. + name : object + Name to be stored in the index. + + See Also + -------- + Index : The base pandas-on-Spark Index type. + Float64Index : A special case of :class:`Index` with purely float labels. + + Notes + ----- + An Index instance can **only** contain hashable objects. + + Examples + -------- + >>> ps.Int64Index([1, 2, 3]) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + From a Series: + + >>> s = ps.Series([1, 2, 3], index=[10, 20, 30]) + >>> ps.Int64Index(s) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + From an Index: + + >>> idx = ps.Index([1, 2, 3]) + >>> ps.Int64Index(idx) # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + """ + + def __new__( + cls, + data: Optional[Any] = None, + dtype: Optional[Union[str, Dtype]] = None, + copy: bool = False, + name: Optional[Name] = None, + ) -> "Int64Index": + warnings.warn( + "Int64Index is deprecated in 3.4.0, and will be removed in 4.0.0. Use Index instead.", + FutureWarning, + ) + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, (Series, Index)): + if dtype is None: + dtype = "int64" + return cast(Int64Index, Index(data, dtype=dtype, copy=copy, name=name)) + + return cast( + Int64Index, ps.from_pandas(pd.Int64Index(data=data, dtype=dtype, copy=copy, name=name)) + ) + + +class Float64Index(NumericIndex): + """ + Immutable sequence used for indexing and alignment. The basic object + storing axis labels for all pandas objects. Float64Index is a special case + of `Index` with purely float labels. + + .. deprecated:: 3.4.0 + + Parameters + ---------- + data : array-like (1-dimensional) + dtype : NumPy dtype (default: float64) + copy : bool + Make a copy of input ndarray. + name : object + Name to be stored in the index. + + See Also + -------- + Index : The base pandas-on-Spark Index type. + Int64Index : A special case of :class:`Index` with purely integer labels. + + Notes + ----- + An Index instance can **only** contain hashable objects. + + Examples + -------- + >>> ps.Float64Index([1.0, 2.0, 3.0]) # doctest: +SKIP + Float64Index([1.0, 2.0, 3.0], dtype='float64') + + From a Series: + + >>> s = ps.Series([1, 2, 3], index=[10, 20, 30]) + >>> ps.Float64Index(s) # doctest: +SKIP + Float64Index([1.0, 2.0, 3.0], dtype='float64') + + From an Index: + + >>> idx = ps.Index([1, 2, 3]) + >>> ps.Float64Index(idx) # doctest: +SKIP + Float64Index([1.0, 2.0, 3.0], dtype='float64') + """ + + def __new__( + cls, + data: Optional[Any] = None, + dtype: Optional[Union[str, Dtype]] = None, + copy: bool = False, + name: Optional[Name] = None, + ) -> "Float64Index": + warnings.warn( + "Float64Index is deprecated in 3.4.0, and will be removed in 4.0.0. Use Index instead.", + FutureWarning, + ) + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, (Series, Index)): + if dtype is None: + dtype = "float64" + return cast(Float64Index, Index(data, dtype=dtype, copy=copy, name=name)) + + return cast( + Float64Index, + ps.from_pandas(pd.Float64Index(data=data, dtype=dtype, copy=copy, name=name)), + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.indexes.numeric + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexes.numeric.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexes.numeric tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexes.numeric, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/timedelta.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/timedelta.py new file mode 100644 index 0000000..36cbb7c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexes/timedelta.py @@ -0,0 +1,197 @@ +# +# 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 cast, no_type_check, Any +from functools import partial + +import pandas as pd +from pandas.api.types import is_hashable # type: ignore[attr-defined] + +from pyspark import pandas as ps +from pyspark._globals import _NoValue +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.missing.indexes import MissingPandasLikeTimedeltaIndex +from pyspark.pandas.series import Series +from pyspark.sql import functions as F + + +HOURS_PER_DAY = 24 +MINUTES_PER_HOUR = 60 +SECONDS_PER_MINUTE = 60 +MILLIS_PER_SECOND = 1000 +MICROS_PER_MILLIS = 1000 + +SECONDS_PER_HOUR = MINUTES_PER_HOUR * SECONDS_PER_MINUTE +SECONDS_PER_DAY = HOURS_PER_DAY * SECONDS_PER_HOUR +MICROS_PER_SECOND = MILLIS_PER_SECOND * MICROS_PER_MILLIS + + +class TimedeltaIndex(Index): + """ + Immutable ndarray-like of timedelta64 data, represented internally as int64, and + which can be boxed to timedelta objects. + + Parameters + ---------- + data : array-like (1-dimensional), optional + Optional timedelta-like data to construct index with. + unit : unit of the arg (D,h,m,s,ms,us,ns) denote the unit, optional + Which is an integer/float number. + freq : str or pandas offset object, optional + One of pandas date offset strings or corresponding objects. The string + 'infer' can be passed in order to set the frequency of the index as the + inferred frequency upon creation. + copy : bool + Make a copy of input ndarray. + name : object + Name to be stored in the index. + + See Also + -------- + Index : The base pandas Index type. + + Examples + -------- + >>> from datetime import timedelta + >>> ps.TimedeltaIndex([timedelta(1), timedelta(microseconds=2)]) + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days 00:00:00', '0 days 00:00:00.000002'], + dtype='timedelta64[ns]', freq=None) + + From an Series: + + >>> s = ps.Series([timedelta(1), timedelta(microseconds=2)], index=[10, 20]) + >>> ps.TimedeltaIndex(s) + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days 00:00:00', '0 days 00:00:00.000002'], + dtype='timedelta64[ns]', freq=None) + + From an Index: + + >>> idx = ps.TimedeltaIndex([timedelta(1), timedelta(microseconds=2)]) + >>> ps.TimedeltaIndex(idx) + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days 00:00:00', '0 days 00:00:00.000002'], + dtype='timedelta64[ns]', freq=None) + """ + + @no_type_check + def __new__( + cls, + data=None, + unit=None, + freq=_NoValue, + closed=None, + dtype=None, + copy=False, + name=None, + ) -> "TimedeltaIndex": + if not is_hashable(name): + raise TypeError("Index.name must be a hashable type") + + if isinstance(data, (Series, Index)): + if dtype is None: + dtype = "timedelta64[ns]" + return cast(TimedeltaIndex, Index(data, dtype=dtype, copy=copy, name=name)) + + kwargs = dict( + data=data, + unit=unit, + closed=closed, + dtype=dtype, + copy=copy, + name=name, + ) + if freq is not _NoValue: + kwargs["freq"] = freq + + return cast(TimedeltaIndex, ps.from_pandas(pd.TimedeltaIndex(**kwargs))) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeTimedeltaIndex, item): + property_or_func = getattr(MissingPandasLikeTimedeltaIndex, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + + raise AttributeError("'TimedeltaIndex' object has no attribute '{}'".format(item)) + + @property + def days(self) -> Index: + """ + Number of days for each element. + """ + + def pandas_days(x) -> int: # type: ignore[no-untyped-def] + return x.days + + return Index(self.to_series().transform(pandas_days)) + + @property + def seconds(self) -> Index: + """ + Number of seconds (>= 0 and less than 1 day) for each element. + """ + + @no_type_check + def get_seconds(scol): + hour_scol = F.date_part(F.lit("HOUR"), scol) + minute_scol = F.date_part(F.lit("MINUTE"), scol) + second_scol = F.date_part(F.lit("SECOND"), scol) + return ( + F.when( + hour_scol < 0, + SECONDS_PER_DAY + hour_scol * SECONDS_PER_HOUR, + ).otherwise(hour_scol * SECONDS_PER_HOUR) + + F.when( + minute_scol < 0, + SECONDS_PER_DAY + minute_scol * SECONDS_PER_MINUTE, + ).otherwise(minute_scol * SECONDS_PER_MINUTE) + + F.when( + second_scol < 0, + SECONDS_PER_DAY + second_scol, + ).otherwise(second_scol) + ).cast("int") + + return Index(self.to_series().spark.transform(get_seconds)) + + @property + def microseconds(self) -> Index: + """ + Number of microseconds (>= 0 and less than 1 second) for each element. + """ + + @no_type_check + def get_microseconds(scol): + second_scol = F.date_part(F.lit("SECOND"), scol) + return ( + ( + F.when( + (second_scol >= 0) & (second_scol < 1), + second_scol, + ) + .when(second_scol < 0, 1 + second_scol) + .otherwise(0) + ) + * MICROS_PER_SECOND + ).cast("int") + + return Index(self.to_series().spark.transform(get_microseconds)) + + @no_type_check + def all(self, *args, **kwargs) -> None: + raise TypeError("Cannot perform 'all' with this index type: %s" % type(self).__name__) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexing.py new file mode 100644 index 0000000..cf8a2c0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/indexing.py @@ -0,0 +1,1862 @@ +# +# 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. +# + +""" +A loc indexer for pandas-on-Spark DataFrame/Series. +""" +from abc import ABCMeta, abstractmethod +from collections.abc import Iterable +from functools import reduce +from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union, cast, Sized + +import pandas as pd +from pandas.api.types import is_list_like # type: ignore[attr-defined] +from pyspark.sql import functions as F, Column as PySparkColumn +from pyspark.sql.types import BooleanType, LongType, DataType +from pyspark.errors import AnalysisException +import numpy as np + +from pyspark import pandas as ps # noqa: F401 +from pyspark.pandas._typing import Label, Name, Scalar +from pyspark.pandas.internal import ( + DEFAULT_SERIES_NAME, + InternalField, + InternalFrame, + NATURAL_ORDER_COLUMN_NAME, + SPARK_DEFAULT_SERIES_NAME, +) +from pyspark.pandas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError +from pyspark.pandas.utils import ( + is_name_like_tuple, + is_name_like_value, + lazy_property, + name_like_string, + same_anchor, + scol_for, + spark_column_equals, + verify_temp_column_name, +) + +# For Supporting Spark Connect +from pyspark.sql.utils import get_column_class + +if TYPE_CHECKING: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.generic import Frame + from pyspark.pandas.series import Series + + +class IndexerLike: + def __init__(self, psdf_or_psser: "Frame"): + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series + + assert isinstance( + psdf_or_psser, (DataFrame, Series) + ), "unexpected argument type: {}".format(type(psdf_or_psser)) + self._psdf_or_psser = psdf_or_psser + + @property + def _is_df(self) -> bool: + from pyspark.pandas.frame import DataFrame + + return isinstance(self._psdf_or_psser, DataFrame) + + @property + def _is_series(self) -> bool: + from pyspark.pandas.series import Series + + return isinstance(self._psdf_or_psser, Series) + + @property + def _psdf(self) -> "DataFrame": + if self._is_df: + return cast("DataFrame", self._psdf_or_psser) + else: + assert self._is_series + return self._psdf_or_psser._psdf + + @property + def _internal(self) -> InternalFrame: + return self._psdf._internal + + +class AtIndexer(IndexerLike): + """ + Access a single value for a row/column label pair. + If the index is not unique, all matching pairs are returned as an array. + Like ``loc``, in that both provide label-based lookups. Use ``at`` if you only need to + get a single value in a DataFrame or Series. + + .. note:: Unlike pandas, pandas-on-Spark only allows using ``at`` to get values but not to + set them. + + .. note:: Warning: If ``row_index`` matches a lot of rows, large amounts of data will be + fetched, potentially causing your machine to run out of memory. + + Raises + ------ + KeyError + When label does not exist in DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame([[0, 2, 3], [0, 4, 1], [10, 20, 30]], + ... index=[4, 5, 5], columns=['A', 'B', 'C']) + >>> psdf + A B C + 4 0 2 3 + 5 0 4 1 + 5 10 20 30 + + Get value at specified row/column pair + + >>> psdf.at[4, 'B'] + 2 + + Get array if an index occurs multiple times + + >>> psdf.at[5, 'B'] + array([ 4, 20]) + """ + + def __getitem__(self, key: Any) -> Union["Series", "DataFrame", Scalar]: + if self._is_df: + if not isinstance(key, tuple) or len(key) != 2: + raise TypeError("Use DataFrame.at like .at[row_index, column_name]") + row_sel, col_sel = key + else: + assert self._is_series, type(self._psdf_or_psser) + if isinstance(key, tuple) and len(key) != 1: + raise TypeError("Use Series.at like .at[row_index]") + row_sel = key + col_sel = self._psdf_or_psser._column_label + + if self._internal.index_level == 1: + if not is_name_like_value(row_sel, allow_none=False, allow_tuple=False): + raise ValueError("At based indexing on a single index can only have a single value") + row_sel = (row_sel,) + else: + if not is_name_like_tuple(row_sel, allow_none=False): + raise ValueError("At based indexing on multi-index can only have tuple values") + + if col_sel is not None: + if not is_name_like_value(col_sel, allow_none=False): + raise ValueError("At based indexing on multi-index can only have tuple values") + if not is_name_like_tuple(col_sel): + col_sel = (col_sel,) + + cond = reduce( + lambda x, y: x & y, + [scol == row for scol, row in zip(self._internal.index_spark_columns, row_sel)], + ) + pdf = ( + self._internal.spark_frame.drop(NATURAL_ORDER_COLUMN_NAME) + .filter(cond) + .select(self._internal.spark_column_for(col_sel)) + .toPandas() + ) + + if len(pdf) < 1: + raise KeyError(name_like_string(row_sel)) + + values = pdf.iloc[:, 0].values + return ( + values if (len(row_sel) < self._internal.index_level or len(values) > 1) else values[0] + ) + + +class iAtIndexer(IndexerLike): + """ + Access a single value for a row/column pair by integer position. + + Like ``iloc``, in that both provide integer-based lookups. Use + ``iat`` if you only need to get or set a single value in a DataFrame + or Series. + + Raises + ------ + KeyError + When label does not exist in DataFrame + + Examples + -------- + >>> df = ps.DataFrame([[0, 2, 3], [0, 4, 1], [10, 20, 30]], + ... columns=['A', 'B', 'C']) + >>> df + A B C + 0 0 2 3 + 1 0 4 1 + 2 10 20 30 + + Get value at specified row/column pair + + >>> df.iat[1, 2] + 1 + + Get value within a series + + >>> psser = ps.Series([1, 2, 3], index=[10, 20, 30]) + >>> psser + 10 1 + 20 2 + 30 3 + dtype: int64 + + >>> psser.iat[1] + 2 + """ + + def __getitem__(self, key: Any) -> Union["Series", "DataFrame", Scalar]: + if self._is_df: + if not isinstance(key, tuple) or len(key) != 2: + raise TypeError( + "Use DataFrame.iat like .iat[row_integer_position, column_integer_position]" + ) + row_sel, col_sel = key + if not isinstance(row_sel, int) or not isinstance(col_sel, int): + raise ValueError("iAt based indexing can only have integer indexers") + return self._psdf_or_psser.iloc[row_sel, col_sel] + else: + assert self._is_series, type(self._psdf_or_psser) + if not isinstance(key, int) and len(key) != 1: + raise TypeError("Use Series.iat like .iat[row_integer_position]") + if not isinstance(key, int): + raise ValueError("iAt based indexing can only have integer indexers") + return self._psdf_or_psser.iloc[key] + + +class LocIndexerLike(IndexerLike, metaclass=ABCMeta): + def _select_rows( + self, rows_sel: Any + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """ + Dispatch the logic for select rows to more specific methods by `rows_sel` argument types. + + Parameters + ---------- + rows_sel : the key specified to select rows. + + Returns + ------- + Tuple of Spark column, int, int: + + * The Spark column for the condition to filter the rows. + * The number of rows when the selection can be simplified by limit. + * The remaining index rows if the result index size is shrunk. + """ + from pyspark.pandas.series import Series + + Column = get_column_class() + if rows_sel is None: + return None, None, None + elif isinstance(rows_sel, Series): + return self._select_rows_by_series(rows_sel) + elif isinstance(rows_sel, Column): + return self._select_rows_by_spark_column(rows_sel) + elif isinstance(rows_sel, slice): + if rows_sel == slice(None): + # If slice is None - select everything, so nothing to do + return None, None, None + return self._select_rows_by_slice(rows_sel) + elif isinstance(rows_sel, tuple): + return self._select_rows_else(rows_sel) + elif is_list_like(rows_sel): + return self._select_rows_by_iterable(rows_sel) + else: + return self._select_rows_else(rows_sel) + + def _select_cols( + self, cols_sel: Any, missing_keys: Optional[List[Name]] = None + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """ + Dispatch the logic for select columns to more specific methods by `cols_sel` argument types. + + Parameters + ---------- + cols_sel : the key specified to select columns. + + Returns + ------- + Tuple of list of column label, list of Spark columns, list of dtypes, bool: + + * The column labels selected. + * The Spark columns selected. + * The field metadata selected. + * The boolean value whether Series should be returned or not. + * The Series name if needed. + """ + from pyspark.pandas.series import Series + + Column = get_column_class() + if cols_sel is None: + column_labels = self._internal.column_labels + data_spark_columns = self._internal.data_spark_columns + data_fields = self._internal.data_fields + return column_labels, data_spark_columns, data_fields, False, None + elif isinstance(cols_sel, Series): + return self._select_cols_by_series(cols_sel, missing_keys) + elif isinstance(cols_sel, Column): + return self._select_cols_by_spark_column(cols_sel, missing_keys) + elif isinstance(cols_sel, slice): + if cols_sel == slice(None): + # If slice is None - select everything, so nothing to do + column_labels = self._internal.column_labels + data_spark_columns = self._internal.data_spark_columns + data_fields = self._internal.data_fields + return column_labels, data_spark_columns, data_fields, False, None + return self._select_cols_by_slice(cols_sel, missing_keys) + elif isinstance(cols_sel, tuple): + return self._select_cols_else(cols_sel, missing_keys) + elif is_list_like(cols_sel): + return self._select_cols_by_iterable(cols_sel, missing_keys) + else: + return self._select_cols_else(cols_sel, missing_keys) + + # Methods for row selection + + @abstractmethod + def _select_rows_by_series( + self, rows_sel: "Series" + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """Select rows by `Series` type key.""" + pass + + @abstractmethod + def _select_rows_by_spark_column( + self, rows_sel: PySparkColumn + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """Select rows by Spark `Column` type key.""" + pass + + @abstractmethod + def _select_rows_by_slice( + self, rows_sel: slice + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """Select rows by `slice` type key.""" + pass + + @abstractmethod + def _select_rows_by_iterable( + self, rows_sel: Iterable + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """Select rows by `Iterable` type key.""" + pass + + @abstractmethod + def _select_rows_else( + self, rows_sel: Any + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + """Select rows by other type key.""" + pass + + # Methods for col selection + + @abstractmethod + def _select_cols_by_series( + self, cols_sel: "Series", missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """Select columns by `Series` type key.""" + pass + + @abstractmethod + def _select_cols_by_spark_column( + self, cols_sel: PySparkColumn, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """Select columns by Spark `Column` type key.""" + pass + + @abstractmethod + def _select_cols_by_slice( + self, cols_sel: slice, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """Select columns by `slice` type key.""" + pass + + @abstractmethod + def _select_cols_by_iterable( + self, cols_sel: Iterable, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """Select columns by `Iterable` type key.""" + pass + + @abstractmethod + def _select_cols_else( + self, cols_sel: Any, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + """Select columns by other type key.""" + pass + + def __getitem__(self, key: Any) -> Union["Series", "DataFrame"]: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series, first_series + + if self._is_series: + if isinstance(key, Series) and not same_anchor(key, self._psdf_or_psser): + name = self._psdf_or_psser.name or DEFAULT_SERIES_NAME + psdf = self._psdf_or_psser.to_frame(name) + temp_col = verify_temp_column_name(psdf, "__temp_col__") + + psdf[temp_col] = key + return type(self)(psdf[name].rename(self._psdf_or_psser.name))[psdf[temp_col]] + + cond, limit, remaining_index = self._select_rows(key) + if cond is None and limit is None: + return self._psdf_or_psser + + column_label = self._psdf_or_psser._column_label + column_labels = [column_label] + data_spark_columns = [self._internal.spark_column_for(column_label)] + data_fields = [self._internal.field_for(column_label)] + returns_series = True + series_name = self._psdf_or_psser.name + else: + assert self._is_df + if isinstance(key, tuple): + if len(key) != 2: + raise SparkPandasIndexingError("Only accepts pairs of candidates") + rows_sel, cols_sel = key + else: + rows_sel = key + cols_sel = None + + if isinstance(rows_sel, Series) and not same_anchor(rows_sel, self._psdf_or_psser): + psdf = self._psdf_or_psser.copy() + temp_col = verify_temp_column_name(cast("DataFrame", psdf), "__temp_col__") + + psdf[temp_col] = rows_sel + return type(self)(psdf)[psdf[temp_col], cols_sel][list(self._psdf_or_psser.columns)] + + cond, limit, remaining_index = self._select_rows(rows_sel) + ( + column_labels, + data_spark_columns, + data_fields, + returns_series, + series_name, + ) = self._select_cols(cols_sel) + + if cond is None and limit is None and returns_series: + psser = self._psdf_or_psser._psser_for(column_labels[0]) + if series_name is not None and series_name != psser.name: + psser = psser.rename(series_name) + return psser + + if remaining_index is not None: + index_spark_columns = self._internal.index_spark_columns[-remaining_index:] + index_names = self._internal.index_names[-remaining_index:] + index_fields = self._internal.index_fields[-remaining_index:] + else: + index_spark_columns = self._internal.index_spark_columns + index_names = self._internal.index_names + index_fields = self._internal.index_fields + + if len(column_labels) > 0: + column_labels = column_labels.copy() + column_labels_level = max( + len(label) if label is not None else 1 for label in column_labels + ) + none_column = 0 + for i, label in enumerate(column_labels): + if label is None: + label = (none_column,) + none_column += 1 + if len(label) < column_labels_level: + label = tuple(list(label) + ([""]) * (column_labels_level - len(label))) + column_labels[i] = label + + if i == 0 and none_column == 1: + column_labels = [None] + + column_label_names = self._internal.column_label_names[-column_labels_level:] + else: + column_label_names = self._internal.column_label_names + + try: + sdf = self._internal.spark_frame + + if cond is not None: + index_columns = sdf.select(index_spark_columns).columns + data_columns = sdf.select(data_spark_columns).columns + sdf = sdf.filter(cond).select(index_spark_columns + data_spark_columns) + index_spark_columns = [scol_for(sdf, col) for col in index_columns] + data_spark_columns = [scol_for(sdf, col) for col in data_columns] + + if limit is not None: + if limit >= 0: + sdf = sdf.limit(limit) + else: + sdf = sdf.limit(sdf.count() + limit) + sdf = sdf.drop(NATURAL_ORDER_COLUMN_NAME) + except AnalysisException: + raise KeyError( + "[{}] don't exist in columns".format( + [col._jc.toString() for col in data_spark_columns] + ) + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=index_spark_columns, + index_names=index_names, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=data_spark_columns, + data_fields=data_fields, + column_label_names=column_label_names, + ) + psdf = DataFrame(internal) + + psdf_or_psser: Union[DataFrame, Series] + if returns_series: + psdf_or_psser = first_series(psdf) + if series_name is not None and series_name != psdf_or_psser.name: + psdf_or_psser = psdf_or_psser.rename(series_name) + else: + psdf_or_psser = psdf + + if remaining_index is not None and remaining_index == 0: + pdf_or_pser = psdf_or_psser.head(2)._to_pandas() + length = len(pdf_or_pser) + if length == 0: + raise KeyError(name_like_string(key)) + elif length == 1: + return pdf_or_pser.iloc[0] + else: + return psdf_or_psser + else: + return psdf_or_psser + + def __setitem__(self, key: Any, value: Any) -> None: + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series, first_series + + Column = get_column_class() + if self._is_series: + if ( + isinstance(key, Series) + and (isinstance(self, iLocIndexer) or not same_anchor(key, self._psdf_or_psser)) + ) or ( + isinstance(value, Series) + and (isinstance(self, iLocIndexer) or not same_anchor(value, self._psdf_or_psser)) + ): + if self._psdf_or_psser.name is None: + psdf = self._psdf_or_psser.to_frame() + column_label = psdf._internal.column_labels[0] + else: + psdf = self._psdf_or_psser._psdf.copy() + column_label = self._psdf_or_psser._column_label + temp_natural_order = verify_temp_column_name(psdf, "__temp_natural_order__") + temp_key_col = verify_temp_column_name(psdf, "__temp_key_col__") + temp_value_col = verify_temp_column_name(psdf, "__temp_value_col__") + + psdf[temp_natural_order] = F.monotonically_increasing_id() + if isinstance(key, Series): + psdf[temp_key_col] = key + if isinstance(value, Series): + psdf[temp_value_col] = value + psdf = psdf.sort_values(temp_natural_order).drop(columns=temp_natural_order) + + psser = psdf._psser_for(column_label) + if isinstance(key, Series): + key = F.col( + "`{}`".format(psdf[temp_key_col]._internal.data_spark_column_names[0]) + ) + if isinstance(value, Series): + value = F.col( + "`{}`".format(psdf[temp_value_col]._internal.data_spark_column_names[0]) + ) + + type(self)(psser)[key] = value + + if self._psdf_or_psser.name is None: + psser = psser.rename() + + self._psdf_or_psser._psdf._update_internal_frame( + psser._psdf[ + self._psdf_or_psser._psdf._internal.column_labels + ]._internal.resolved_copy, + check_same_anchor=False, + ) + return + + if isinstance(value, DataFrame): + raise ValueError("Incompatible indexer with DataFrame") + + cond, limit, remaining_index = self._select_rows(key) + if cond is None: + cond = F.lit(True) + if limit is not None: + cond = cond & ( + self._internal.spark_frame[cast(iLocIndexer, self)._sequence_col] < F.lit(limit) + ) + + if isinstance(value, (Series, Column)): + if remaining_index is not None and remaining_index == 0: + raise ValueError( + "No axis named {} for object type {}".format(key, type(value).__name__) + ) + if isinstance(value, Series): + value = value.spark.column + else: + value = F.lit(value) + scol = ( + F.when(cond, value) + .otherwise(self._internal.spark_column_for(self._psdf_or_psser._column_label)) + .alias(name_like_string(self._psdf_or_psser.name or SPARK_DEFAULT_SERIES_NAME)) + ) + + internal = self._internal.with_new_spark_column( + self._psdf_or_psser._column_label, scol # TODO: dtype? + ) + self._psdf_or_psser._psdf._update_internal_frame(internal, check_same_anchor=False) + else: + assert self._is_df + + if isinstance(key, tuple): + if len(key) != 2: + raise SparkPandasIndexingError("Only accepts pairs of candidates") + rows_sel, cols_sel = key + else: + rows_sel = key + cols_sel = None + + if isinstance(value, DataFrame): + if len(value.columns) == 1: + value = first_series(value) + else: + raise ValueError("Only a dataframe with one column can be assigned") + + if ( + isinstance(rows_sel, Series) + and ( + isinstance(self, iLocIndexer) or not same_anchor(rows_sel, self._psdf_or_psser) + ) + ) or ( + isinstance(value, Series) + and (isinstance(self, iLocIndexer) or not same_anchor(value, self._psdf_or_psser)) + ): + psdf = cast(DataFrame, self._psdf_or_psser.copy()) + temp_natural_order = verify_temp_column_name(psdf, "__temp_natural_order__") + temp_key_col = verify_temp_column_name(psdf, "__temp_key_col__") + temp_value_col = verify_temp_column_name(psdf, "__temp_value_col__") + + psdf[temp_natural_order] = F.monotonically_increasing_id() + if isinstance(rows_sel, Series): + psdf[temp_key_col] = rows_sel + if isinstance(value, Series): + psdf[temp_value_col] = value + psdf = psdf.sort_values(temp_natural_order).drop(columns=temp_natural_order) + + if isinstance(rows_sel, Series): + rows_sel = F.col( + "`{}`".format(psdf[temp_key_col]._internal.data_spark_column_names[0]) + ) + if isinstance(value, Series): + value = F.col( + "`{}`".format(psdf[temp_value_col]._internal.data_spark_column_names[0]) + ) + + type(self)(psdf)[rows_sel, cols_sel] = value + + self._psdf_or_psser._update_internal_frame( + psdf[list(self._psdf_or_psser.columns)]._internal.resolved_copy, + check_same_anchor=False, + ) + return + + cond, limit, remaining_index = self._select_rows(rows_sel) + missing_keys: List[Name] = [] + _, data_spark_columns, _, _, _ = self._select_cols(cols_sel, missing_keys=missing_keys) + + if cond is None: + cond = F.lit(True) + if limit is not None: + cond = cond & ( + self._internal.spark_frame[cast(iLocIndexer, self)._sequence_col] < F.lit(limit) + ) + + if isinstance(value, (Series, Column)): + if remaining_index is not None and remaining_index == 0: + raise ValueError("Incompatible indexer with Series") + if len(data_spark_columns) > 1: + raise ValueError("shape mismatch") + if isinstance(value, Series): + value = value.spark.column + else: + value = F.lit(value) + + new_data_spark_columns = [] + new_fields = [] + for new_scol, spark_column_name, new_field in zip( + self._internal.data_spark_columns, + self._internal.data_spark_column_names, + self._internal.data_fields, + ): + for scol in data_spark_columns: + if spark_column_equals(new_scol, scol): + new_scol = F.when(cond, value).otherwise(scol).alias(spark_column_name) + new_field = InternalField.from_struct_field( + self._internal.spark_frame.select(new_scol).schema[0], + use_extension_dtypes=new_field.is_extension_dtype, + ) + break + new_data_spark_columns.append(new_scol) + new_fields.append(new_field) + + column_labels = self._internal.column_labels.copy() + for missing in missing_keys: + if is_name_like_tuple(missing): + label = cast(Label, missing) + else: + label = cast(Label, (missing,)) + if len(label) < self._internal.column_labels_level: + label = tuple( + list(label) + ([""] * (self._internal.column_labels_level - len(label))) + ) + elif len(label) > self._internal.column_labels_level: + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + len(label), self._internal.column_labels_level + ) + ) + column_labels.append(label) + new_data_spark_columns.append(F.when(cond, value).alias(name_like_string(label))) + new_fields.append(None) + + internal = self._internal.with_new_columns( + new_data_spark_columns, column_labels=column_labels, data_fields=new_fields + ) + self._psdf_or_psser._update_internal_frame(internal, check_same_anchor=False) + + +class LocIndexer(LocIndexerLike): + """ + Access a group of rows and columns by label(s) or a boolean Series. + + ``.loc[]`` is primarily label based, but may also be used with a + conditional boolean Series derived from the DataFrame or Series. + + Allowed inputs are: + + - A single label, e.g. ``5`` or ``'a'``, (note that ``5`` is + interpreted as a *label* of the index, and **never** as an + integer position along the index) for column selection. + + - A list or array of labels, e.g. ``['a', 'b', 'c']``. + + - A slice object with labels, e.g. ``'a':'f'``. + + - A conditional boolean Series derived from the DataFrame or Series + + - A boolean array of the same length as the column axis being sliced, + e.g. ``[True, False, True]``. + + - An alignable boolean pandas Series to the column axis being sliced. + The index of the key will be aligned before masking. + + Not allowed inputs which pandas allows are: + + - A boolean array of the same length as the row axis being sliced, + e.g. ``[True, False, True]``. + - A ``callable`` function with one argument (the calling Series, DataFrame + or Panel) and that returns valid output for indexing (one of the above) + + .. note:: MultiIndex is not supported yet. + + .. note:: Note that contrary to usual python slices, **both** the + start and the stop are included, and the step of the slice is not allowed. + + .. note:: With a list or array of labels for row selection, + pandas-on-Spark behaves as a filter without reordering by the labels. + + See Also + -------- + Series.loc : Access group of values using labels. + + Examples + -------- + **Getting values** + + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=['cobra', 'viper', 'sidewinder'], + ... columns=['max_speed', 'shield']) + >>> df + max_speed shield + cobra 1 2 + viper 4 5 + sidewinder 7 8 + + Single label. Note this returns the row as a Series. + + >>> df.loc['viper'] + max_speed 4 + shield 5 + Name: viper, dtype: int64 + + List of labels. Note using ``[[]]`` returns a DataFrame. + Also note that pandas-on-Spark behaves just a filter without reordering by the labels. + + >>> df.loc[['viper', 'sidewinder']] + max_speed shield + viper 4 5 + sidewinder 7 8 + + >>> df.loc[['sidewinder', 'viper']] + max_speed shield + viper 4 5 + sidewinder 7 8 + + Single label for column. + + >>> df.loc['cobra', 'shield'] + 2 + + List of labels for row. + + >>> df.loc[['cobra'], 'shield'] + cobra 2 + Name: shield, dtype: int64 + + List of labels for column. + + >>> df.loc['cobra', ['shield']] + shield 2 + Name: cobra, dtype: int64 + + List of labels for both row and column. + + >>> df.loc[['cobra'], ['shield']] + shield + cobra 2 + + Slice with labels for row and single label for column. + Note that both the start and stop of the slice are included. + + >>> df.loc['cobra':'viper', 'max_speed'] + cobra 1 + viper 4 + Name: max_speed, dtype: int64 + + Conditional that returns a boolean Series + + >>> df.loc[df['shield'] > 6] + max_speed shield + sidewinder 7 8 + + Conditional that returns a boolean Series with column labels specified + + >>> df.loc[df['shield'] > 6, ['max_speed']] + max_speed + sidewinder 7 + + A boolean array of the same length as the column axis being sliced. + + >>> df.loc[:, [False, True]] + shield + cobra 2 + viper 5 + sidewinder 8 + + An alignable boolean Series to the column axis being sliced. + + >>> df.loc[:, pd.Series([False, True], index=['max_speed', 'shield'])] + shield + cobra 2 + viper 5 + sidewinder 8 + + **Setting values** + + Setting value for all items matching the list of labels. + + >>> df.loc[['viper', 'sidewinder'], ['shield']] = 50 + >>> df + max_speed shield + cobra 1 2 + viper 4 50 + sidewinder 7 50 + + Setting value for an entire row + + >>> df.loc['cobra'] = 10 + >>> df + max_speed shield + cobra 10 10 + viper 4 50 + sidewinder 7 50 + + Set value for an entire column + + >>> df.loc[:, 'max_speed'] = 30 + >>> df + max_speed shield + cobra 30 10 + viper 30 50 + sidewinder 30 50 + + Set value for an entire list of columns + + >>> df.loc[:, ['max_speed', 'shield']] = 100 + >>> df + max_speed shield + cobra 100 100 + viper 100 100 + sidewinder 100 100 + + Set value with Series + + >>> df.loc[:, 'shield'] = df['shield'] * 2 + >>> df + max_speed shield + cobra 100 200 + viper 100 200 + sidewinder 100 200 + + **Getting values on a DataFrame with an index that has integer labels** + + Another example using integers for the index + + >>> df = ps.DataFrame([[1, 2], [4, 5], [7, 8]], + ... index=[7, 8, 9], + ... columns=['max_speed', 'shield']) + >>> df + max_speed shield + 7 1 2 + 8 4 5 + 9 7 8 + + Slice with integer labels for rows. Note that both + the start and stop of the slice are included. + + >>> df.loc[7:9] + max_speed shield + 7 1 2 + 8 4 5 + 9 7 8 + """ + + @staticmethod + def _NotImplemented(description: str) -> SparkPandasNotImplementedError: + return SparkPandasNotImplementedError( + description=description, + pandas_function=".loc[..., ...]", + spark_target_function="select, where", + ) + + def _select_rows_by_series( + self, rows_sel: "Series" + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + assert isinstance(rows_sel.spark.data_type, BooleanType), rows_sel.spark.data_type + return rows_sel.spark.column, None, None + + def _select_rows_by_spark_column( + self, rows_sel: PySparkColumn + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + spark_type = self._internal.spark_frame.select(rows_sel).schema[0].dataType + assert isinstance(spark_type, BooleanType), spark_type + return rows_sel, None, None + + def _select_rows_by_slice( + self, rows_sel: slice + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + from pyspark.pandas.indexes import MultiIndex + + if rows_sel.step is not None: + raise LocIndexer._NotImplemented("Cannot use step with Spark.") + elif self._internal.index_level == 1: + sdf = self._internal.spark_frame + index = self._psdf_or_psser.index + index_column = index.to_series() + index_data_type = index_column.spark.data_type + start = rows_sel.start + stop = rows_sel.stop + + # get natural order from '__natural_order__' from start to stop + # to keep natural order. + start_and_stop = ( + sdf.select(index_column.spark.column, NATURAL_ORDER_COLUMN_NAME) + .where( + (index_column.spark.column == F.lit(start).cast(index_data_type)) + | (index_column.spark.column == F.lit(stop).cast(index_data_type)) + ) + .collect() + ) + + start = [row[1] for row in start_and_stop if row[0] == start] + start = start[0] if len(start) > 0 else None + + stop = [row[1] for row in start_and_stop if row[0] == stop] + stop = stop[-1] if len(stop) > 0 else None + + conds: List[PySparkColumn] = [] + if start is not None: + conds.append(F.col(NATURAL_ORDER_COLUMN_NAME) >= F.lit(start).cast(LongType())) + if stop is not None: + conds.append(F.col(NATURAL_ORDER_COLUMN_NAME) <= F.lit(stop).cast(LongType())) + + # if index order is not monotonic increasing or decreasing + # and specified values don't exist in index, raise KeyError + if (start is None and rows_sel.start is not None) or ( + stop is None and rows_sel.stop is not None + ): + + inc = index_column.is_monotonic_increasing + if inc is False: + dec = index_column.is_monotonic_decreasing + + if start is None and rows_sel.start is not None: + start = rows_sel.start + if inc is not False: + conds.append( + index_column.spark.column >= F.lit(start).cast(index_data_type) + ) + elif dec is not False: + conds.append( + index_column.spark.column <= F.lit(start).cast(index_data_type) + ) + else: + raise KeyError(rows_sel.start) + if stop is None and rows_sel.stop is not None: + stop = rows_sel.stop + if inc is not False: + conds.append(index_column.spark.column <= F.lit(stop).cast(index_data_type)) + elif dec is not False: + conds.append(index_column.spark.column >= F.lit(stop).cast(index_data_type)) + else: + raise KeyError(rows_sel.stop) + + return reduce(lambda x, y: x & y, conds), None, None + else: + from pyspark.sql.types import StructType + + index = self._psdf_or_psser.index + index_data_type = [ # type: ignore[assignment] + f.dataType for f in cast(StructType, index.to_series().spark.data_type) + ] + + start = rows_sel.start + if start is not None: + if not isinstance(start, tuple): + start = (start,) + if len(start) == 0: + start = None + stop = rows_sel.stop + if stop is not None: + if not isinstance(stop, tuple): + stop = (stop,) + if len(stop) == 0: + stop = None + + depth = max( + len(start) if start is not None else 0, len(stop) if stop is not None else 0 + ) + if depth == 0: + return None, None, None + elif ( + depth > self._internal.index_level + or not index.droplevel(list(range(self._internal.index_level)[depth:])).is_monotonic + ): + raise KeyError( + "Key length ({}) was greater than MultiIndex sort depth".format(depth) + ) + + conds = [] + if start is not None: + cond = F.lit(True) + for scol, value, dt in list( + zip( + self._internal.index_spark_columns, + cast(Tuple[int, ...], start), + cast(List[DataType], index_data_type), + ) + )[::-1]: + compare = MultiIndex._comparator_for_monotonic_increasing(dt) + Column = get_column_class() + cond = F.when(scol.eqNullSafe(F.lit(value).cast(dt)), cond).otherwise( + compare(scol, F.lit(value).cast(dt), Column.__gt__) + ) + conds.append(cond) + if stop is not None: + cond = F.lit(True) + for scol, value, dt in list( + zip( + self._internal.index_spark_columns, + cast(Tuple[int, ...], stop), + cast(List[DataType], index_data_type), + ) + )[::-1]: + compare = MultiIndex._comparator_for_monotonic_increasing(dt) + Column = get_column_class() + cond = F.when(scol.eqNullSafe(F.lit(value).cast(dt)), cond).otherwise( + compare(scol, F.lit(value).cast(dt), Column.__lt__) + ) + conds.append(cond) + + return reduce(lambda x, y: x & y, conds), None, None + + def _select_rows_by_iterable( + self, rows_sel: Iterable + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + rows_sel = list(rows_sel) + if len(rows_sel) == 0: + return F.lit(False), None, None + elif self._internal.index_level == 1: + index_column = self._psdf_or_psser.index.to_series() + index_data_type = index_column.spark.data_type + if len(rows_sel) == 1: + return ( + index_column.spark.column == F.lit(rows_sel[0]).cast(index_data_type), + None, + None, + ) + else: + return ( + index_column.spark.column.isin( + [F.lit(r).cast(index_data_type) for r in rows_sel] + ), + None, + None, + ) + else: + raise LocIndexer._NotImplemented("Cannot select with MultiIndex with Spark.") + + def _select_rows_else( + self, rows_sel: Any + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + if not isinstance(rows_sel, tuple): + rows_sel = (rows_sel,) + if len(rows_sel) > self._internal.index_level: + raise SparkPandasIndexingError("Too many indexers") + + rows = [scol == value for scol, value in zip(self._internal.index_spark_columns, rows_sel)] + return ( + reduce(lambda x, y: x & y, rows), + None, + self._internal.index_level - len(rows_sel), + ) + + def _get_from_multiindex_column( + self, + key: Optional[Label], + missing_keys: Optional[List[Name]], + labels: Optional[List[Tuple[Label, Label]]] = None, + recursed: int = 0, + ) -> Tuple[ + List[Label], Optional[List[PySparkColumn]], List[InternalField], bool, Optional[Name] + ]: + """Select columns from multi-index columns.""" + assert isinstance(key, tuple) + if labels is None: + labels = [(label, label) for label in self._internal.column_labels] + for k in key: + labels = [ + (label, None if lbl is None else lbl[1:]) + for label, lbl in labels + if (lbl is None and k is None) or (lbl is not None and lbl[0] == k) + ] + if len(labels) == 0: + if missing_keys is None: + raise KeyError(k) + else: + missing_keys.append(key) + return [], [], [], False, None + + if all(lbl is not None and len(lbl) > 0 and lbl[0] == "" for _, lbl in labels): + # If the head is '', drill down recursively. + labels = [(label, tuple([str(key), *lbl[1:]])) for i, (label, lbl) in enumerate(labels)] + return self._get_from_multiindex_column((str(key),), missing_keys, labels, recursed + 1) + else: + returns_series = all(lbl is None or len(lbl) == 0 for _, lbl in labels) + series_name: Optional[Name] + if returns_series: + label_set = set(label for label, _ in labels) + assert len(label_set) == 1 + label = list(label_set)[0] + column_labels = [label] + data_spark_columns = [self._internal.spark_column_for(label)] + data_fields = [self._internal.field_for(label)] + if label is None: + series_name = None + else: + if recursed > 0: + label = label[:-recursed] + series_name = label if len(label) > 1 else label[0] + else: + column_labels = [ + None if lbl is None or lbl == (None,) else lbl for _, lbl in labels + ] + data_spark_columns = [self._internal.spark_column_for(label) for label, _ in labels] + data_fields = [self._internal.field_for(label) for label, _ in labels] + series_name = None + + return column_labels, data_spark_columns, data_fields, returns_series, series_name + + def _select_cols_by_series( + self, cols_sel: "Series", missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + column_labels = cols_sel._internal.column_labels + data_spark_columns = cols_sel._internal.data_spark_columns + data_fields = cols_sel._internal.data_fields + return column_labels, data_spark_columns, data_fields, True, None + + def _select_cols_by_spark_column( + self, cols_sel: PySparkColumn, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + column_labels: List[Label] = [(self._internal.spark_frame.select(cols_sel).columns[0],)] + data_spark_columns = [cols_sel] + return column_labels, data_spark_columns, None, True, None + + def _select_cols_by_slice( + self, cols_sel: slice, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + start, stop = self._psdf_or_psser.columns.slice_locs( + start=cols_sel.start, end=cols_sel.stop + ) + column_labels = self._internal.column_labels[start:stop] + data_spark_columns = self._internal.data_spark_columns[start:stop] + data_fields = self._internal.data_fields[start:stop] + return column_labels, data_spark_columns, data_fields, False, None + + def _select_cols_by_iterable( + self, cols_sel: Iterable, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + from pyspark.pandas.series import Series + + Column = get_column_class() + if all(isinstance(key, Series) for key in cols_sel): + column_labels = [key._column_label for key in cols_sel] + data_spark_columns = [key.spark.column for key in cols_sel] + data_fields = [key._internal.data_fields[0] for key in cols_sel] + elif all(isinstance(key, Column) for key in cols_sel): + column_labels = [ + (self._internal.spark_frame.select(col).columns[0],) for col in cols_sel + ] + data_spark_columns = list(cols_sel) + data_fields = None + elif all(isinstance(key, bool) for key in cols_sel) or all( + isinstance(key, np.bool_) for key in cols_sel + ): + if len(cast(Sized, cols_sel)) != len(self._internal.column_labels): + raise IndexError( + "Boolean index has wrong length: %s instead of %s" + % (len(cast(Sized, cols_sel)), len(self._internal.column_labels)) + ) + if isinstance(cols_sel, pd.Series): + if not cols_sel.index.sort_values().equals(self._psdf.columns.sort_values()): + raise SparkPandasIndexingError( + "Unalignable boolean Series provided as indexer " + "(index of the boolean Series and of the indexed object do not match)" + ) + else: + column_labels = [ + column_label + for column_label in self._internal.column_labels + if cols_sel[column_label if len(column_label) > 1 else column_label[0]] + ] + data_spark_columns = [ + self._internal.spark_column_for(column_label) + for column_label in column_labels + ] + data_fields = [ + self._internal.field_for(column_label) for column_label in column_labels + ] + else: + column_labels = [ + self._internal.column_labels[i] for i, col in enumerate(cols_sel) if col + ] + data_spark_columns = [ + self._internal.data_spark_columns[i] for i, col in enumerate(cols_sel) if col + ] + data_fields = [ + self._internal.data_fields[i] for i, col in enumerate(cols_sel) if col + ] + elif any(isinstance(key, tuple) for key in cols_sel) and any( + not is_name_like_tuple(key) for key in cols_sel + ): + raise TypeError( + "Expected tuple, got {}".format( + type(set(key for key in cols_sel if not is_name_like_tuple(key)).pop()) + ) + ) + else: + if missing_keys is None and all(isinstance(key, tuple) for key in cols_sel): + level = self._internal.column_labels_level + if any(len(key) != level for key in cols_sel): + raise ValueError("All the key level should be the same as column index level.") + + column_labels = [] + data_spark_columns = [] + data_fields = [] + for key in cols_sel: + found = False + for label in self._internal.column_labels: + if label == key or label[0] == key: + column_labels.append(label) + data_spark_columns.append(self._internal.spark_column_for(label)) + data_fields.append(self._internal.field_for(label)) + found = True + if not found: + if missing_keys is None: + raise KeyError("['{}'] not in index".format(name_like_string(key))) + else: + missing_keys.append(key) + + return column_labels, data_spark_columns, data_fields, False, None + + def _select_cols_else( + self, cols_sel: Any, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + if not is_name_like_tuple(cols_sel): + cols_sel = (cols_sel,) + return self._get_from_multiindex_column(cols_sel, missing_keys) + + +class iLocIndexer(LocIndexerLike): + """ + Purely integer-location based indexing for selection by position. + + ``.iloc[]`` is primarily integer position based (from ``0`` to + ``length-1`` of the axis), but may also be used with a conditional boolean Series. + + Allowed inputs are: + + - An integer for column selection, e.g. ``5``. + - A list or array of integers for row selection with distinct index values, + e.g. ``[3, 4, 0]`` + - A list or array of integers for column selection, e.g. ``[4, 3, 0]``. + - A boolean array for column selection. + - A slice object with ints for row and column selection, e.g. ``1:7``. + + Not allowed inputs which pandas allows are: + + - A list or array of integers for row selection with duplicated indexes, + e.g. ``[4, 4, 0]``. + - A boolean array for row selection. + - A ``callable`` function with one argument (the calling Series, DataFrame + or Panel) and that returns valid output for indexing (one of the above). + This is useful in method chains when you don't have a reference to the + calling object but would like to base your selection on some value. + + ``.iloc`` will raise ``IndexError`` if a requested indexer is + out-of-bounds, except *slice* indexers which allow out-of-bounds + indexing (this conforms with python/numpy *slice* semantics). + + See Also + -------- + DataFrame.loc : Purely label-location based indexer for selection by label. + Series.iloc : Purely integer-location based indexing for + selection by position. + + Examples + -------- + + >>> mydict = [{'a': 1, 'b': 2, 'c': 3, 'd': 4}, + ... {'a': 100, 'b': 200, 'c': 300, 'd': 400}, + ... {'a': 1000, 'b': 2000, 'c': 3000, 'd': 4000 }] + >>> df = ps.DataFrame(mydict, columns=['a', 'b', 'c', 'd']) + >>> df + a b c d + 0 1 2 3 4 + 1 100 200 300 400 + 2 1000 2000 3000 4000 + + **Indexing just the rows** + + A scalar integer for row selection. + + >>> df.iloc[1] + a 100 + b 200 + c 300 + d 400 + Name: 1, dtype: int64 + + >>> df.iloc[[0]] + a b c d + 0 1 2 3 4 + + With a `slice` object. + + >>> df.iloc[:3] + a b c d + 0 1 2 3 4 + 1 100 200 300 400 + 2 1000 2000 3000 4000 + + **Indexing both axes** + + You can mix the indexer types for the index and columns. Use ``:`` to + select the entire axis. + + With scalar integers. + + >>> df.iloc[:1, 1] + 0 2 + Name: b, dtype: int64 + + With lists of integers. + + >>> df.iloc[:2, [1, 3]] + b d + 0 2 4 + 1 200 400 + + With `slice` objects. + + >>> df.iloc[:2, 0:3] + a b c + 0 1 2 3 + 1 100 200 300 + + With a boolean array whose length matches the columns. + + >>> df.iloc[:, [True, False, True, False]] + a c + 0 1 3 + 1 100 300 + 2 1000 3000 + + **Setting values** + + Setting value for all items matching the list of labels. + + >>> df.iloc[[1, 2], [1]] = 50 + >>> df + a b c d + 0 1 2 3 4 + 1 100 50 300 400 + 2 1000 50 3000 4000 + + Setting value for an entire row + + >>> df.iloc[0] = 10 + >>> df + a b c d + 0 10 10 10 10 + 1 100 50 300 400 + 2 1000 50 3000 4000 + + Set value for an entire column + + >>> df.iloc[:, 2] = 30 + >>> df + a b c d + 0 10 10 30 10 + 1 100 50 30 400 + 2 1000 50 30 4000 + + Set value for an entire list of columns + + >>> df.iloc[:, [2, 3]] = 100 + >>> df + a b c d + 0 10 10 100 100 + 1 100 50 100 100 + 2 1000 50 100 100 + + Set value with Series + + >>> df.iloc[:, 3] = df.iloc[:, 3] * 2 + >>> df + a b c d + 0 10 10 100 200 + 1 100 50 100 200 + 2 1000 50 100 200 + """ + + @staticmethod + def _NotImplemented(description: str) -> SparkPandasNotImplementedError: + return SparkPandasNotImplementedError( + description=description, + pandas_function=".iloc[..., ...]", + spark_target_function="select, where", + ) + + @lazy_property + def _internal(self) -> "InternalFrame": + # Use resolved_copy to fix the natural order. + internal = super()._internal.resolved_copy + sdf = InternalFrame.attach_distributed_sequence_column( + internal.spark_frame, column_name=self._sequence_col + ) + return internal.with_new_sdf(spark_frame=sdf.orderBy(NATURAL_ORDER_COLUMN_NAME)) + + @lazy_property + def _sequence_col(self) -> str: + # Use resolved_copy to fix the natural order. + internal = super()._internal.resolved_copy + return verify_temp_column_name(internal.spark_frame, "__distributed_sequence_column__") + + def _select_rows_by_series( + self, rows_sel: "Series" + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + raise iLocIndexer._NotImplemented( + ".iloc requires numeric slice, conditional " + "boolean Index or a sequence of positions as int, " + "got {}".format(type(rows_sel)) + ) + + def _select_rows_by_spark_column( + self, rows_sel: PySparkColumn + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + raise iLocIndexer._NotImplemented( + ".iloc requires numeric slice, conditional " + "boolean Index or a sequence of positions as int, " + "got {}".format(type(rows_sel)) + ) + + def _select_rows_by_slice( + self, rows_sel: slice + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + def verify_type(i: int) -> None: + if not isinstance(i, int): + raise TypeError( + "cannot do slice indexing with these indexers [{}] of {}".format(i, type(i)) + ) + + has_negative = False + start = rows_sel.start + if start is not None: + verify_type(start) + if start == 0: + start = None + elif start < 0: + has_negative = True + stop = rows_sel.stop + if stop is not None: + verify_type(stop) + if stop < 0: + has_negative = True + + step = rows_sel.step + if step is not None: + verify_type(step) + if step == 0: + raise ValueError("slice step cannot be zero") + else: + step = 1 + + if start is None and step == 1: + return None, stop, None + + sdf = self._internal.spark_frame + sequence_scol = sdf[self._sequence_col] + + if has_negative or (step < 0 and start is None): + cnt = sdf.count() + + cond = [] + if start is not None: + if start < 0: + start = start + cnt + if step >= 0: + cond.append(sequence_scol >= F.lit(start).cast(LongType())) + else: + cond.append(sequence_scol <= F.lit(start).cast(LongType())) + if stop is not None: + if stop < 0: + stop = stop + cnt + if step >= 0: + cond.append(sequence_scol < F.lit(stop).cast(LongType())) + else: + cond.append(sequence_scol > F.lit(stop).cast(LongType())) + if step != 1: + if step > 0: + start = start or 0 + else: + start = start or (cnt - 1) + cond.append(((sequence_scol - start) % F.lit(step).cast(LongType())) == F.lit(0)) + + return reduce(lambda x, y: x & y, cond), None, None + + def _select_rows_by_iterable( + self, rows_sel: Iterable + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + sdf = self._internal.spark_frame + + if any(isinstance(key, (int, np.int64, np.int32)) and key < 0 for key in rows_sel): + offset = sdf.count() + else: + offset = 0 + + new_rows_sel = [] + for key in list(rows_sel): + if not isinstance(key, (int, np.int64, np.int32)): + raise TypeError( + "cannot do positional indexing with these indexers [{}] of {}".format( + key, type(key) + ) + ) + if key < 0: + key = key + offset + new_rows_sel.append(key) + + if len(new_rows_sel) != len(set(new_rows_sel)): + raise NotImplementedError( + "Duplicated row selection is not currently supported; " + "however, normalized index was [%s]" % new_rows_sel + ) + + if len(new_rows_sel) == 0: + cond = F.lit(False) + else: + cond = sdf[self._sequence_col].isin( + [F.lit(int(key)).cast(LongType()) for key in new_rows_sel] + ) + return cond, None, None + + def _select_rows_else( + self, rows_sel: Any + ) -> Tuple[Optional[PySparkColumn], Optional[int], Optional[int]]: + if isinstance(rows_sel, int): + sdf = self._internal.spark_frame + return (sdf[self._sequence_col] == rows_sel), None, 0 + elif isinstance(rows_sel, tuple): + raise SparkPandasIndexingError("Too many indexers") + else: + raise iLocIndexer._NotImplemented( + ".iloc requires numeric slice, conditional " + "boolean Index or a sequence of positions as int, " + "got {}".format(type(rows_sel)) + ) + + def _select_cols_by_series( + self, cols_sel: "Series", missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + raise ValueError( + "Location based indexing can only have [integer, integer slice, " + "listlike of integers, boolean array] types, got {}".format(cols_sel) + ) + + def _select_cols_by_spark_column( + self, cols_sel: PySparkColumn, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + raise ValueError( + "Location based indexing can only have [integer, integer slice, " + "listlike of integers, boolean array] types, got {}".format(cols_sel) + ) + + def _select_cols_by_slice( + self, cols_sel: slice, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + if all( + s is None or isinstance(s, int) for s in (cols_sel.start, cols_sel.stop, cols_sel.step) + ): + column_labels = self._internal.column_labels[cols_sel] + data_spark_columns = self._internal.data_spark_columns[cols_sel] + data_fields = self._internal.data_fields[cols_sel] + return column_labels, data_spark_columns, data_fields, False, None + else: + not_none = ( + cols_sel.start + if cols_sel.start is not None + else cols_sel.stop + if cols_sel.stop is not None + else cols_sel.step + ) + raise TypeError( + "cannot do slice indexing with these indexers {} of {}".format( + not_none, type(not_none) + ) + ) + + def _select_cols_by_iterable( + self, cols_sel: Iterable, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + if all(isinstance(s, bool) for s in cols_sel): + cols_sel = [i for i, s in enumerate(cols_sel) if s] + if all(isinstance(s, int) for s in cols_sel): + column_labels = [self._internal.column_labels[s] for s in cols_sel] + data_spark_columns = [self._internal.data_spark_columns[s] for s in cols_sel] + data_fields = [self._internal.data_fields[s] for s in cols_sel] + return column_labels, data_spark_columns, data_fields, False, None + else: + raise TypeError("cannot perform reduce with flexible type") + + def _select_cols_else( + self, cols_sel: Any, missing_keys: Optional[List[Name]] + ) -> Tuple[ + List[Label], + Optional[List[PySparkColumn]], + Optional[List[InternalField]], + bool, + Optional[Name], + ]: + if isinstance(cols_sel, int): + if cols_sel > len(self._internal.column_labels): + raise KeyError(cols_sel) + column_labels = [self._internal.column_labels[cols_sel]] + data_spark_columns = [self._internal.data_spark_columns[cols_sel]] + data_fields = [self._internal.data_fields[cols_sel]] + return column_labels, data_spark_columns, data_fields, True, None + else: + raise ValueError( + "Location based indexing can only have [integer, integer slice, " + "listlike of integers, boolean array] types, got {}".format(cols_sel) + ) + + def __setitem__(self, key: Any, value: Any) -> None: + Column = get_column_class() + if not isinstance(value, Column) and is_list_like(value): + iloc_item = self[key] + if not is_list_like(key) or not is_list_like(iloc_item): + raise ValueError("setting an array element with a sequence.") + else: + shape_iloc_item = iloc_item.shape + len_iloc_item = shape_iloc_item[0] + len_value = len(value) + if len_iloc_item != len_value: + if self._is_series: + raise ValueError( + "cannot set using a list-like indexer with a different length than " + "the value" + ) + else: + raise ValueError( + "shape mismatch: value array of shape ({},) could not be broadcast " + "to indexing result of shape {}".format(len_value, shape_iloc_item) + ) + super().__setitem__(key, value) + # Update again with resolved_copy to drop extra columns. + self._psdf._update_internal_frame( + self._psdf._internal.resolved_copy, check_same_anchor=False + ) + + # Clean up implicitly cached properties to be able to reuse the indexer. + del self._internal + del self._sequence_col + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.indexing + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.indexing.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.indexing tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.indexing, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/internal.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/internal.py new file mode 100644 index 0000000..e025d91 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/internal.py @@ -0,0 +1,1680 @@ +# +# 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. +# + +""" +An internal immutable DataFrame with some metadata to manage indexes. +""" +import re +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, TYPE_CHECKING, cast + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype # noqa: F401 +from pyspark._globals import _NoValue, _NoValueType +from pyspark.sql import ( + functions as F, + Column as PySparkColumn, + DataFrame as PySparkDataFrame, + Window, +) +from pyspark.sql.types import ( # noqa: F401 + BooleanType, + DataType, + LongType, + StructField, + StructType, + StringType, +) +from pyspark.sql.utils import is_timestamp_ntz_preferred + +# For supporting Spark Connect +from pyspark.sql.utils import is_remote, get_column_class, get_dataframe_class + +# For running doctests and reference resolution in PyCharm. +from pyspark import pandas as ps +from pyspark.pandas._typing import Label + +if TYPE_CHECKING: + # This is required in old Python 3.5 to prevent circular reference. + from pyspark.pandas.series import Series +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale +from pyspark.pandas.data_type_ops.base import DataTypeOps +from pyspark.pandas.typedef import ( + Dtype, + as_spark_type, + extension_dtypes, + infer_pd_series_spark_type, + spark_type_to_pandas_dtype, +) +from pyspark.pandas.utils import ( + column_labels_level, + default_session, + is_name_like_tuple, + is_testing, + lazy_property, + name_like_string, + scol_for, + spark_column_equals, +) + + +# A function to turn given numbers to Spark columns that represent pandas-on-Spark index. +SPARK_INDEX_NAME_FORMAT = "__index_level_{}__".format +SPARK_DEFAULT_INDEX_NAME = SPARK_INDEX_NAME_FORMAT(0) +# A pattern to check if the name of a Spark column is a pandas-on-Spark index name or not. +SPARK_INDEX_NAME_PATTERN = re.compile(r"__index_level_[0-9]+__") + +NATURAL_ORDER_COLUMN_NAME = "__natural_order__" + +HIDDEN_COLUMNS = {NATURAL_ORDER_COLUMN_NAME} + +DEFAULT_SERIES_NAME = 0 +SPARK_DEFAULT_SERIES_NAME = str(DEFAULT_SERIES_NAME) + + +class InternalField: + """ + The internal field to store the dtype as well as the Spark's StructField optionally. + + Parameters + ---------- + dtype : numpy.dtype or pandas' ExtensionDtype + The dtype for the field + struct_field : StructField, optional + The `StructField` for the field. If None, InternalFrame will properly set. + """ + + def __init__(self, dtype: Dtype, struct_field: Optional[StructField] = None): + self._dtype = dtype + self._struct_field = struct_field + + @staticmethod + def from_struct_field( + struct_field: StructField, *, use_extension_dtypes: bool = False + ) -> "InternalField": + """ + Returns a new InternalField object created from the given StructField. + + The dtype will be inferred from the data type of the given StructField. + + Parameters + ---------- + struct_field : StructField + The StructField used to create a new InternalField object. + use_extension_dtypes : bool + If True, try to use the extension dtypes. + + Returns + ------- + InternalField + """ + return InternalField( + dtype=spark_type_to_pandas_dtype( + struct_field.dataType, use_extension_dtypes=use_extension_dtypes + ), + struct_field=struct_field, + ) + + @property + def dtype(self) -> Dtype: + """Return the dtype for the field.""" + return self._dtype + + @property + def struct_field(self) -> Optional[StructField]: + """Return the StructField for the field.""" + return self._struct_field + + @property + def name(self) -> str: + """Return the field name if the StructField exists.""" + assert self.struct_field is not None + return self.struct_field.name + + @property + def spark_type(self) -> DataType: + """Return the spark data type for the field if the StructField exists.""" + assert self.struct_field is not None + return self.struct_field.dataType + + @property + def nullable(self) -> bool: + """Return the nullability for the field if the StructField exists.""" + assert self.struct_field is not None + return self.struct_field.nullable + + @property + def metadata(self) -> Dict[str, Any]: + """Return the metadata for the field if the StructField exists.""" + assert self.struct_field is not None + return self.struct_field.metadata + + @property + def is_extension_dtype(self) -> bool: + """Return whether the dtype for the field is an extension type or not.""" + return isinstance(self.dtype, extension_dtypes) + + def normalize_spark_type(self) -> "InternalField": + """Return a new InternalField object with normalized Spark data type.""" + assert self.struct_field is not None + return self.copy( + spark_type=force_decimal_precision_scale(as_nullable_spark_type(self.spark_type)), + nullable=True, + ) + + def copy( + self, + *, + name: Union[str, _NoValueType] = _NoValue, + dtype: Union[Dtype, _NoValueType] = _NoValue, + spark_type: Union[DataType, _NoValueType] = _NoValue, + nullable: Union[bool, _NoValueType] = _NoValue, + metadata: Union[Optional[Dict[str, Any]], _NoValueType] = _NoValue, + ) -> "InternalField": + """Copy the InternalField object.""" + if name is _NoValue: + name = self.name + if dtype is _NoValue: + dtype = self.dtype + if spark_type is _NoValue: + spark_type = self.spark_type + if nullable is _NoValue: + nullable = self.nullable + if metadata is _NoValue: + metadata = self.metadata + return InternalField( + dtype=cast(Dtype, dtype), + struct_field=StructField( + name=cast(str, name), + dataType=cast(DataType, spark_type), + nullable=cast(bool, nullable), + metadata=cast(Optional[Dict[str, Any]], metadata), + ), + ) + + def __eq__(self, other: Any) -> bool: + return ( + isinstance(other, InternalField) + and self.dtype == other.dtype + and self.struct_field == other.struct_field + ) + + def __repr__(self) -> str: + return "InternalField(dtype={dtype}, struct_field={struct_field})".format( + dtype=self.dtype, struct_field=self.struct_field + ) + + +class InternalFrame: + """ + The internal immutable DataFrame which manages Spark DataFrame and column names and index + information. + + .. note:: this is an internal class. It is not supposed to be exposed to users and users + should not directly access to it. + + The internal immutable DataFrame represents the index information for a DataFrame it belongs to. + For instance, if we have a pandas-on-Spark DataFrame as below, pandas DataFrame does not + store the index as columns. + + >>> psdf = ps.DataFrame({ + ... 'A': [1, 2, 3, 4], + ... 'B': [5, 6, 7, 8], + ... 'C': [9, 10, 11, 12], + ... 'D': [13, 14, 15, 16], + ... 'E': [17, 18, 19, 20]}, columns = ['A', 'B', 'C', 'D', 'E']) + >>> psdf # doctest: +NORMALIZE_WHITESPACE + A B C D E + 0 1 5 9 13 17 + 1 2 6 10 14 18 + 2 3 7 11 15 19 + 3 4 8 12 16 20 + + However, all columns including index column are also stored in Spark DataFrame internally + as below. + + >>> psdf._internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +-----------------+---+---+---+---+---+ + |__index_level_0__| A| B| C| D| E| + +-----------------+---+---+---+---+---+ + | 0| 1| 5| 9| 13| 17| + | 1| 2| 6| 10| 14| 18| + | 2| 3| 7| 11| 15| 19| + | 3| 4| 8| 12| 16| 20| + +-----------------+---+---+---+---+---+ + + To fill this gap, the current metadata is used by mapping Spark's internal column + to pandas-on-Spark's index. See the method below: + + * `spark_frame` represents the internal Spark DataFrame + + * `data_spark_column_names` represents non-indexing Spark column names + + * `data_spark_columns` represents non-indexing Spark columns + + * `data_fields` represents non-indexing InternalFields + + * `index_spark_column_names` represents internal index Spark column names + + * `index_spark_columns` represents internal index Spark columns + + * `index_fields` represents index InternalFields + + * `spark_column_names` represents all columns + + * `index_names` represents the external index name as a label + + * `to_internal_spark_frame` represents Spark DataFrame derived by the metadata. Includes index. + + * `to_pandas_frame` represents pandas DataFrame derived by the metadata + + >>> internal = psdf._internal + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+---+---+---+---+---+-----------------+ + |__index_level_0__| A| B| C| D| E|__natural_order__| + +-----------------+---+---+---+---+---+-----------------+ + | 0| 1| 5| 9| 13| 17| ...| + | 1| 2| 6| 10| 14| 18| ...| + | 2| 3| 7| 11| 15| 19| ...| + | 3| 4| 8| 12| 16| 20| ...| + +-----------------+---+---+---+---+---+-----------------+ + >>> internal.data_spark_column_names + ['A', 'B', 'C', 'D', 'E'] + >>> internal.index_spark_column_names + ['__index_level_0__'] + >>> internal.spark_column_names + ['__index_level_0__', 'A', 'B', 'C', 'D', 'E'] + >>> internal.index_names + [None] + >>> internal.data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('A', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('B', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('C', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('D', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('E', LongType(), False))] + >>> internal.index_fields + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', LongType(), False))] + >>> internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +-----------------+---+---+---+---+---+ + |__index_level_0__| A| B| C| D| E| + +-----------------+---+---+---+---+---+ + | 0| 1| 5| 9| 13| 17| + | 1| 2| 6| 10| 14| 18| + | 2| 3| 7| 11| 15| 19| + | 3| 4| 8| 12| 16| 20| + +-----------------+---+---+---+---+---+ + >>> internal.to_pandas_frame + A B C D E + 0 1 5 9 13 17 + 1 2 6 10 14 18 + 2 3 7 11 15 19 + 3 4 8 12 16 20 + + In case that index is set to one of the existing columns as below: + + >>> psdf1 = psdf.set_index("A") + >>> psdf1 # doctest: +NORMALIZE_WHITESPACE + B C D E + A + 1 5 9 13 17 + 2 6 10 14 18 + 3 7 11 15 19 + 4 8 12 16 20 + + >>> psdf1._internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +---+---+---+---+---+ + | A| B| C| D| E| + +---+---+---+---+---+ + | 1| 5| 9| 13| 17| + | 2| 6| 10| 14| 18| + | 3| 7| 11| 15| 19| + | 4| 8| 12| 16| 20| + +---+---+---+---+---+ + + >>> internal = psdf1._internal + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+---+---+---+---+---+-----------------+ + |__index_level_0__| A| B| C| D| E|__natural_order__| + +-----------------+---+---+---+---+---+-----------------+ + | 0| 1| 5| 9| 13| 17| ...| + | 1| 2| 6| 10| 14| 18| ...| + | 2| 3| 7| 11| 15| 19| ...| + | 3| 4| 8| 12| 16| 20| ...| + +-----------------+---+---+---+---+---+-----------------+ + >>> internal.data_spark_column_names + ['B', 'C', 'D', 'E'] + >>> internal.index_spark_column_names + ['A'] + >>> internal.spark_column_names + ['A', 'B', 'C', 'D', 'E'] + >>> internal.index_names + [('A',)] + >>> internal.data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('B', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('C', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('D', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('E', LongType(), False))] + >>> internal.index_fields + [InternalField(dtype=int64, struct_field=StructField('A', LongType(), False))] + >>> internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +---+---+---+---+---+ + | A| B| C| D| E| + +---+---+---+---+---+ + | 1| 5| 9| 13| 17| + | 2| 6| 10| 14| 18| + | 3| 7| 11| 15| 19| + | 4| 8| 12| 16| 20| + +---+---+---+---+---+ + >>> internal.to_pandas_frame # doctest: +NORMALIZE_WHITESPACE + B C D E + A + 1 5 9 13 17 + 2 6 10 14 18 + 3 7 11 15 19 + 4 8 12 16 20 + + In case that index becomes a multi index as below: + + >>> psdf2 = psdf.set_index("A", append=True) + >>> psdf2 # doctest: +NORMALIZE_WHITESPACE + B C D E + A + 0 1 5 9 13 17 + 1 2 6 10 14 18 + 2 3 7 11 15 19 + 3 4 8 12 16 20 + + >>> psdf2._internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +-----------------+---+---+---+---+---+ + |__index_level_0__| A| B| C| D| E| + +-----------------+---+---+---+---+---+ + | 0| 1| 5| 9| 13| 17| + | 1| 2| 6| 10| 14| 18| + | 2| 3| 7| 11| 15| 19| + | 3| 4| 8| 12| 16| 20| + +-----------------+---+---+---+---+---+ + + >>> internal = psdf2._internal + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+---+---+---+---+---+-----------------+ + |__index_level_0__| A| B| C| D| E|__natural_order__| + +-----------------+---+---+---+---+---+-----------------+ + | 0| 1| 5| 9| 13| 17| ...| + | 1| 2| 6| 10| 14| 18| ...| + | 2| 3| 7| 11| 15| 19| ...| + | 3| 4| 8| 12| 16| 20| ...| + +-----------------+---+---+---+---+---+-----------------+ + >>> internal.data_spark_column_names + ['B', 'C', 'D', 'E'] + >>> internal.index_spark_column_names + ['__index_level_0__', 'A'] + >>> internal.spark_column_names + ['__index_level_0__', 'A', 'B', 'C', 'D', 'E'] + >>> internal.index_names + [None, ('A',)] + >>> internal.data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('B', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('C', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('D', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('E', LongType(), False))] + >>> internal.index_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('A', LongType(), False))] + >>> internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +-----------------+---+---+---+---+---+ + |__index_level_0__| A| B| C| D| E| + +-----------------+---+---+---+---+---+ + | 0| 1| 5| 9| 13| 17| + | 1| 2| 6| 10| 14| 18| + | 2| 3| 7| 11| 15| 19| + | 3| 4| 8| 12| 16| 20| + +-----------------+---+---+---+---+---+ + >>> internal.to_pandas_frame # doctest: +NORMALIZE_WHITESPACE + B C D E + A + 0 1 5 9 13 17 + 1 2 6 10 14 18 + 2 3 7 11 15 19 + 3 4 8 12 16 20 + + For multi-level columns, it also holds column_labels + + >>> columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), + ... ('Y', 'C'), ('Y', 'D')]) + >>> psdf3 = ps.DataFrame([ + ... [1, 2, 3, 4], + ... [5, 6, 7, 8], + ... [9, 10, 11, 12], + ... [13, 14, 15, 16], + ... [17, 18, 19, 20]], columns = columns) + >>> psdf3 # doctest: +NORMALIZE_WHITESPACE + X Y + A B C D + 0 1 2 3 4 + 1 5 6 7 8 + 2 9 10 11 12 + 3 13 14 15 16 + 4 17 18 19 20 + + >>> internal = psdf3._internal + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+------+------+------+------+-----------------+ + |__index_level_0__|(X, A)|(X, B)|(Y, C)|(Y, D)|__natural_order__| + +-----------------+------+------+------+------+-----------------+ + | 0| 1| 2| 3| 4| ...| + | 1| 5| 6| 7| 8| ...| + | 2| 9| 10| 11| 12| ...| + | 3| 13| 14| 15| 16| ...| + | 4| 17| 18| 19| 20| ...| + +-----------------+------+------+------+------+-----------------+ + >>> internal.data_spark_column_names + ['(X, A)', '(X, B)', '(Y, C)', '(Y, D)'] + >>> internal.column_labels + [('X', 'A'), ('X', 'B'), ('Y', 'C'), ('Y', 'D')] + + For Series, it also holds scol to represent the column. + + >>> psseries = psdf1.B + >>> psseries + A + 1 5 + 2 6 + 3 7 + 4 8 + Name: B, dtype: int64 + + >>> internal = psseries._internal + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+---+---+---+---+---+-----------------+ + |__index_level_0__| A| B| C| D| E|__natural_order__| + +-----------------+---+---+---+---+---+-----------------+ + | 0| 1| 5| 9| 13| 17| ...| + | 1| 2| 6| 10| 14| 18| ...| + | 2| 3| 7| 11| 15| 19| ...| + | 3| 4| 8| 12| 16| 20| ...| + +-----------------+---+---+---+---+---+-----------------+ + >>> internal.data_spark_column_names + ['B'] + >>> internal.index_spark_column_names + ['A'] + >>> internal.spark_column_names + ['A', 'B'] + >>> internal.index_names + [('A',)] + >>> internal.data_fields + [InternalField(dtype=int64, struct_field=StructField('B', LongType(), False))] + >>> internal.index_fields + [InternalField(dtype=int64, struct_field=StructField('A', LongType(), False))] + >>> internal.to_internal_spark_frame.show() # doctest: +NORMALIZE_WHITESPACE + +---+---+ + | A| B| + +---+---+ + | 1| 5| + | 2| 6| + | 3| 7| + | 4| 8| + +---+---+ + >>> internal.to_pandas_frame # doctest: +NORMALIZE_WHITESPACE + B + A + 1 5 + 2 6 + 3 7 + 4 8 + """ + + def __init__( + self, + spark_frame: PySparkDataFrame, + index_spark_columns: Optional[List[PySparkColumn]], + index_names: Optional[List[Optional[Label]]] = None, + index_fields: Optional[List[InternalField]] = None, + column_labels: Optional[List[Label]] = None, + data_spark_columns: Optional[List[PySparkColumn]] = None, + data_fields: Optional[List[InternalField]] = None, + column_label_names: Optional[List[Optional[Label]]] = None, + ): + """ + Create a new internal immutable DataFrame to manage Spark DataFrame, column fields and + index fields and names. + + :param spark_frame: Spark DataFrame to be managed. + :param index_spark_columns: list of Spark Column + Spark Columns for the index. + :param index_names: list of tuples + the index names. + :param index_fields: list of InternalField + the InternalFields for the index columns + :param column_labels: list of tuples with the same length + The multi-level values in the tuples. + :param data_spark_columns: list of Spark Column + Spark Columns to appear as columns. If this is None, calculated + from spark_frame. + :param data_fields: list of InternalField + the InternalFields for the data columns + :param column_label_names: Names for each of the column index levels. + + See the examples below to refer what each parameter means. + + >>> column_labels = pd.MultiIndex.from_tuples( + ... [('a', 'x'), ('a', 'y'), ('b', 'z')], names=["column_labels_a", "column_labels_b"]) + >>> row_index = pd.MultiIndex.from_tuples( + ... [('foo', 'bar'), ('foo', 'bar'), ('zoo', 'bar')], + ... names=["row_index_a", "row_index_b"]) + >>> psdf = ps.DataFrame( + ... [[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=row_index, columns=column_labels) + >>> psdf.set_index(('a', 'x'), append=True, inplace=True) + >>> psdf # doctest: +NORMALIZE_WHITESPACE + column_labels_a a b + column_labels_b y z + row_index_a row_index_b (a, x) + foo bar 1 2 3 + 4 5 6 + zoo bar 7 8 9 + + >>> internal = psdf._internal + + >>> internal.spark_frame.show() # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS + +-----------------+-----------------+------+------+------+... + |__index_level_0__|__index_level_1__|(a, x)|(a, y)|(b, z)|... + +-----------------+-----------------+------+------+------+... + | foo| bar| 1| 2| 3|... + | foo| bar| 4| 5| 6|... + | zoo| bar| 7| 8| 9|... + +-----------------+-----------------+------+------+------+... + + >>> internal.index_spark_columns # doctest: +SKIP + [Column<'__index_level_0__'>, Column<'__index_level_1__'>, Column<'(a, x)'>] + + >>> internal.index_names + [('row_index_a',), ('row_index_b',), ('a', 'x')] + + >>> internal.index_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=object, + struct_field=StructField('__index_level_0__', StringType(), False)), + InternalField(dtype=object, + struct_field=StructField('__index_level_1__', StringType(), False)), + InternalField(dtype=int64, + struct_field=StructField('(a, x)', LongType(), False))] + + >>> internal.column_labels + [('a', 'y'), ('b', 'z')] + + >>> internal.data_spark_columns # doctest: +SKIP + [Column<'(a, y)'>, Column<'(b, z)'>] + + >>> internal.data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('(a, y)', LongType(), False)), + InternalField(dtype=int64, struct_field=StructField('(b, z)', LongType(), False))] + + >>> internal.column_label_names + [('column_labels_a',), ('column_labels_b',)] + """ + SparkDataFrame = get_dataframe_class() + assert isinstance(spark_frame, SparkDataFrame) + assert not spark_frame.isStreaming, "pandas-on-Spark does not support Structured Streaming." + + if not index_spark_columns: + if data_spark_columns is not None: + if column_labels is not None: + data_spark_columns = [ + scol.alias(name_like_string(label)) + for scol, label in zip(data_spark_columns, column_labels) + ] + spark_frame = spark_frame.select(data_spark_columns) + + assert not any(SPARK_INDEX_NAME_PATTERN.match(name) for name in spark_frame.columns), ( + "Index columns should not appear in columns of the Spark DataFrame. Avoid " + "index column names [%s]." % SPARK_INDEX_NAME_PATTERN + ) + + # Create default index. + spark_frame = InternalFrame.attach_default_index(spark_frame) + index_spark_columns = [scol_for(spark_frame, SPARK_DEFAULT_INDEX_NAME)] + + index_fields = [ + InternalField.from_struct_field( + StructField(SPARK_DEFAULT_INDEX_NAME, LongType(), nullable=False) + ) + ] + + if data_spark_columns is not None: + data_struct_fields = [ + field + for field in spark_frame.schema.fields + if field.name != SPARK_DEFAULT_INDEX_NAME + ] + data_spark_columns = [ + scol_for(spark_frame, field.name) for field in data_struct_fields + ] + if data_fields is not None: + data_fields = [ + field.copy( + name=name_like_string(struct_field.name), + ) + for field, struct_field in zip(data_fields, data_struct_fields) + ] + + if NATURAL_ORDER_COLUMN_NAME not in spark_frame.columns: + spark_frame = spark_frame.withColumn( + NATURAL_ORDER_COLUMN_NAME, F.monotonically_increasing_id() + ) + + self._sdf = spark_frame + + # index_spark_columns + Column = get_column_class() + assert all( + isinstance(index_scol, Column) for index_scol in index_spark_columns + ), index_spark_columns + + self._index_spark_columns: List[Column] = index_spark_columns # type: ignore[valid-type] + + # data_spark_columns + if data_spark_columns is None: + data_spark_columns = [ + scol_for(spark_frame, col) + for col in spark_frame.columns + if all( + not spark_column_equals(scol_for(spark_frame, col), index_scol) + for index_scol in index_spark_columns + ) + and col not in HIDDEN_COLUMNS + ] + else: + assert all(isinstance(scol, Column) for scol in data_spark_columns) + + self._data_spark_columns: List[Column] = data_spark_columns # type: ignore[valid-type] + + # fields + if index_fields is None: + index_fields = [None] * len(index_spark_columns) + if data_fields is None: + data_fields = [None] * len(data_spark_columns) + + assert len(index_spark_columns) == len(index_fields), ( + len(index_spark_columns), + len(index_fields), + ) + assert len(data_spark_columns) == len(data_fields), ( + len(data_spark_columns), + len(data_fields), + ) + + if any(field is None or field.struct_field is None for field in index_fields) and any( + field is None or field.struct_field is None for field in data_fields + ): + schema = spark_frame.select(index_spark_columns + data_spark_columns).schema + fields = [ + InternalField.from_struct_field(struct_field) + if field is None + else InternalField(field.dtype, struct_field) + if field.struct_field is None + else field + for field, struct_field in zip(index_fields + data_fields, schema.fields) + ] + index_fields = fields[: len(index_spark_columns)] + data_fields = fields[len(index_spark_columns) :] + elif any(field is None or field.struct_field is None for field in index_fields): + schema = spark_frame.select(index_spark_columns).schema + index_fields = [ + InternalField.from_struct_field(struct_field) + if field is None + else InternalField(field.dtype, struct_field) + if field.struct_field is None + else field + for field, struct_field in zip(index_fields, schema.fields) + ] + elif any(field is None or field.struct_field is None for field in data_fields): + schema = spark_frame.select(data_spark_columns).schema + data_fields = [ + InternalField.from_struct_field(struct_field) + if field is None + else InternalField(field.dtype, struct_field) + if field.struct_field is None + else field + for field, struct_field in zip(data_fields, schema.fields) + ] + + assert all( + isinstance(ops.dtype, Dtype.__args__) # type: ignore[attr-defined] + and ( + ops.dtype == np.dtype("object") + or as_spark_type(ops.dtype, raise_error=False) is not None + ) + for ops in index_fields + ), index_fields + + if is_testing(): + struct_fields = spark_frame.select(index_spark_columns).schema.fields + if is_remote(): + # TODO(SPARK-42965): For some reason, the metadata of StructField is different + # in a few tests when using Spark Connect. However, the function works properly. + # Therefore, we temporarily perform Spark Connect tests by excluding metadata + # until the issue is resolved. + def remove_metadata(struct_field: StructField) -> StructField: + new_struct_field = StructField( + struct_field.name, struct_field.dataType, struct_field.nullable + ) + return new_struct_field + + assert all( + remove_metadata(index_field.struct_field) == remove_metadata(struct_field) + for index_field, struct_field in zip(index_fields, struct_fields) + ), (index_fields, struct_fields) + else: + assert all( + index_field.struct_field == struct_field + for index_field, struct_field in zip(index_fields, struct_fields) + ), (index_fields, struct_fields) + + self._index_fields: List[InternalField] = index_fields + + assert all( + isinstance(ops.dtype, Dtype.__args__) # type: ignore[attr-defined] + and ( + ops.dtype == np.dtype("object") + or as_spark_type(ops.dtype, raise_error=False) is not None + ) + for ops in data_fields + ), data_fields + + if is_testing(): + struct_fields = spark_frame.select(data_spark_columns).schema.fields + if is_remote(): + # TODO(SPARK-42965): For some reason, the metadata of StructField is different + # in a few tests when using Spark Connect. However, the function works properly. + # Therefore, we temporarily perform Spark Connect tests by excluding metadata + # until the issue is resolved. + def remove_metadata(struct_field: StructField) -> StructField: + new_struct_field = StructField( + struct_field.name, struct_field.dataType, struct_field.nullable + ) + return new_struct_field + + assert all( + remove_metadata(data_field.struct_field) == remove_metadata(struct_field) + for data_field, struct_field in zip(data_fields, struct_fields) + ), (data_fields, struct_fields) + else: + assert all( + data_field.struct_field == struct_field + for data_field, struct_field in zip(data_fields, struct_fields) + ), (data_fields, struct_fields) + + self._data_fields: List[InternalField] = data_fields + + # index_names + if not index_names: + index_names = [None] * len(index_spark_columns) + + assert len(index_spark_columns) == len(index_names), ( + len(index_spark_columns), + len(index_names), + ) + assert all( + is_name_like_tuple(index_name, check_type=True) for index_name in index_names + ), index_names + + self._index_names: List[Optional[Label]] = index_names + + # column_labels + if column_labels is None: + column_labels = [(col,) for col in spark_frame.select(self._data_spark_columns).columns] + else: + assert len(column_labels) == len(self._data_spark_columns), ( + len(column_labels), + len(self._data_spark_columns), + ) + if len(column_labels) == 1: + column_label = column_labels[0] + assert is_name_like_tuple(column_label, check_type=True), column_label + else: + assert all( + is_name_like_tuple(column_label, check_type=True) + for column_label in column_labels + ), column_labels + assert len(set(len(label) for label in column_labels)) <= 1, column_labels + + self._column_labels: List[Label] = column_labels + + # column_label_names + if column_label_names is None: + column_label_names = [None] * column_labels_level(self._column_labels) + else: + if len(self._column_labels) > 0: + assert len(column_label_names) == column_labels_level(self._column_labels), ( + len(column_label_names), + column_labels_level(self._column_labels), + ) + else: + assert len(column_label_names) > 0, len(column_label_names) + assert all( + is_name_like_tuple(column_label_name, check_type=True) + for column_label_name in column_label_names + ), column_label_names + + self._column_label_names: List[Optional[Label]] = column_label_names + + @staticmethod + def attach_default_index( + sdf: PySparkDataFrame, default_index_type: Optional[str] = None + ) -> PySparkDataFrame: + """ + This method attaches a default index to Spark DataFrame. Spark does not have the index + notion so corresponding column should be generated. + There are several types of default index can be configured by `compute.default_index_type`. + + >>> spark_frame = ps.range(10).to_spark() + >>> spark_frame + DataFrame[id: bigint] + + It adds the default index column '__index_level_0__'. + + >>> spark_frame = InternalFrame.attach_default_index(spark_frame) + >>> spark_frame + DataFrame[__index_level_0__: bigint, id: bigint] + + It throws an exception if the given column name already exists. + + >>> InternalFrame.attach_default_index(spark_frame) + ... # doctest: +ELLIPSIS + Traceback (most recent call last): + ... + AssertionError: '__index_level_0__' already exists... + """ + index_column = SPARK_DEFAULT_INDEX_NAME + assert ( + index_column not in sdf.columns + ), "'%s' already exists in the Spark column names '%s'" % (index_column, sdf.columns) + + if default_index_type is None: + default_index_type = ps.get_option("compute.default_index_type") + + if default_index_type == "sequence": + return InternalFrame.attach_sequence_column(sdf, column_name=index_column) + elif default_index_type == "distributed-sequence": + return InternalFrame.attach_distributed_sequence_column(sdf, column_name=index_column) + elif default_index_type == "distributed": + return InternalFrame.attach_distributed_column(sdf, column_name=index_column) + else: + raise ValueError( + "'compute.default_index_type' should be one of 'sequence'," + " 'distributed-sequence' and 'distributed'" + ) + + @staticmethod + def attach_sequence_column(sdf: PySparkDataFrame, column_name: str) -> PySparkDataFrame: + scols = [scol_for(sdf, column) for column in sdf.columns] + sequential_index = ( + F.row_number().over(Window.orderBy(F.monotonically_increasing_id())).cast("long") - 1 + ) + return sdf.select(sequential_index.alias(column_name), *scols) + + @staticmethod + def attach_distributed_column(sdf: PySparkDataFrame, column_name: str) -> PySparkDataFrame: + scols = [scol_for(sdf, column) for column in sdf.columns] + # Does not add an alias to avoid having some changes in protobuf definition for now. + # The alias is more for query strings in DataFrame.explain, and they are cosmetic changes. + if is_remote(): + return sdf.select(F.monotonically_increasing_id().alias(column_name), *scols) + jvm = sdf.sparkSession._jvm + tag = jvm.org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FUNC_ALIAS() + jexpr = F.monotonically_increasing_id()._jc.expr() + jexpr.setTagValue(tag, "distributed_index") + return sdf.select(PySparkColumn(jvm.Column(jexpr)).alias(column_name), *scols) + + @staticmethod + def attach_distributed_sequence_column( + sdf: PySparkDataFrame, column_name: str + ) -> PySparkDataFrame: + """ + This method attaches a Spark column that has a sequence in a distributed manner. + This is equivalent to the column assigned when default index type 'distributed-sequence'. + + >>> sdf = ps.DataFrame(['a', 'b', 'c']).to_spark() + >>> sdf = InternalFrame.attach_distributed_sequence_column(sdf, column_name="sequence") + >>> sdf.show() # doctest: +NORMALIZE_WHITESPACE + +--------+---+ + |sequence| 0| + +--------+---+ + | 0| a| + | 1| b| + | 2| c| + +--------+---+ + """ + if len(sdf.columns) > 0: + if is_remote(): + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.expressions import DistributedSequenceID + + return sdf.select( + ConnectColumn(DistributedSequenceID()).alias(column_name), + "*", # type: ignore[call-overload] + ) + else: + return PySparkDataFrame( + sdf._jdf.toDF().withSequenceColumn(column_name), + sdf.sparkSession, + ) + else: + cnt = sdf.count() + if cnt > 0: + return default_session().range(cnt).toDF(column_name) + else: + return default_session().createDataFrame( + [], schema=StructType().add(column_name, data_type=LongType(), nullable=False) + ) + + def spark_column_for(self, label: Label) -> PySparkColumn: + """Return Spark Column for the given column label.""" + column_labels_to_scol = dict(zip(self.column_labels, self.data_spark_columns)) + if label in column_labels_to_scol: + return column_labels_to_scol[label] + else: + raise KeyError(name_like_string(label)) + + def spark_column_name_for(self, label_or_scol: Union[Label, PySparkColumn]) -> str: + """Return the actual Spark column name for the given column label.""" + Column = get_column_class() + if isinstance(label_or_scol, Column): + return self.spark_frame.select(label_or_scol).columns[0] + else: + return self.field_for(label_or_scol).name # type: ignore[arg-type] + + def spark_type_for(self, label_or_scol: Union[Label, PySparkColumn]) -> DataType: + """Return DataType for the given column label.""" + Column = get_column_class() + if isinstance(label_or_scol, Column): + return self.spark_frame.select(label_or_scol).schema[0].dataType + else: + return self.field_for(label_or_scol).spark_type # type: ignore[arg-type] + + def spark_column_nullable_for(self, label_or_scol: Union[Label, PySparkColumn]) -> bool: + """Return nullability for the given column label.""" + Column = get_column_class() + if isinstance(label_or_scol, Column): + return self.spark_frame.select(label_or_scol).schema[0].nullable + else: + return self.field_for(label_or_scol).nullable # type: ignore[arg-type] + + def field_for(self, label: Label) -> InternalField: + """Return InternalField for the given column label.""" + column_labels_to_fields = dict(zip(self.column_labels, self.data_fields)) + if label in column_labels_to_fields: + return column_labels_to_fields[label] + else: + raise KeyError(name_like_string(label)) + + @property + def spark_frame(self) -> PySparkDataFrame: + """Return the managed Spark DataFrame.""" + return self._sdf + + @lazy_property + def data_spark_column_names(self) -> List[str]: + """Return the managed column field names.""" + return [field.name for field in self.data_fields] + + @property + def data_spark_columns(self) -> List[PySparkColumn]: + """Return Spark Columns for the managed data columns.""" + return self._data_spark_columns + + @property + def index_spark_column_names(self) -> List[str]: + """Return the managed index field names.""" + return [field.name for field in self.index_fields] + + @property + def index_spark_columns(self) -> List[PySparkColumn]: + """Return Spark Columns for the managed index columns.""" + return self._index_spark_columns + + @lazy_property + def spark_column_names(self) -> List[str]: + """Return all the field names including index field names.""" + return self.spark_frame.select(self.spark_columns).columns + + @lazy_property + def spark_columns(self) -> List[PySparkColumn]: + """Return Spark Columns for the managed columns including index columns.""" + index_spark_columns = self.index_spark_columns + return index_spark_columns + [ + spark_column + for spark_column in self.data_spark_columns + if all(not spark_column_equals(spark_column, scol) for scol in index_spark_columns) + ] + + @property + def index_names(self) -> List[Optional[Label]]: + """Return the managed index names.""" + return self._index_names + + @lazy_property + def index_level(self) -> int: + """Return the level of the index.""" + return len(self._index_names) + + @property + def column_labels(self) -> List[Label]: + """Return the managed column index.""" + return self._column_labels + + @lazy_property + def column_labels_level(self) -> int: + """Return the level of the column index.""" + return len(self._column_label_names) + + @property + def column_label_names(self) -> List[Optional[Label]]: + """Return names of the index levels.""" + return self._column_label_names + + @property + def index_fields(self) -> List[InternalField]: + """Return InternalFields for the managed index columns.""" + return self._index_fields + + @property + def data_fields(self) -> List[InternalField]: + """Return InternalFields for the managed columns.""" + return self._data_fields + + @lazy_property + def to_internal_spark_frame(self) -> PySparkDataFrame: + """ + Return as Spark DataFrame. This contains index columns as well + and should be only used for internal purposes. + """ + index_spark_columns = self.index_spark_columns + data_columns = [] + for spark_column in self.data_spark_columns: + if all(not spark_column_equals(spark_column, scol) for scol in index_spark_columns): + data_columns.append(spark_column) + return self.spark_frame.select(index_spark_columns + data_columns) + + @lazy_property + def to_pandas_frame(self) -> pd.DataFrame: + """Return as pandas DataFrame.""" + sdf = self.to_internal_spark_frame + pdf = sdf.toPandas() + if len(pdf) == 0 and len(sdf.schema) > 0: + pdf = pdf.astype( + {field.name: spark_type_to_pandas_dtype(field.dataType) for field in sdf.schema} + ) + + return InternalFrame.restore_index(pdf, **self.arguments_for_restore_index) + + @lazy_property + def arguments_for_restore_index(self) -> Dict: + """Create arguments for `restore_index`.""" + column_names = [] + fields = self.index_fields.copy() + + for spark_column, column_name, field in zip( + self.data_spark_columns, self.data_spark_column_names, self.data_fields + ): + for index_spark_column_name, index_spark_column in zip( + self.index_spark_column_names, self.index_spark_columns + ): + if spark_column_equals(spark_column, index_spark_column): + column_names.append(index_spark_column_name) + break + else: + column_names.append(column_name) + fields.append(field) + + return dict( + index_columns=self.index_spark_column_names, + index_names=self.index_names, + data_columns=column_names, + column_labels=self.column_labels, + column_label_names=self.column_label_names, + fields=fields, + ) + + @staticmethod + def restore_index( + pdf: pd.DataFrame, + *, + index_columns: List[str], + index_names: List[Label], + data_columns: List[str], + column_labels: List[Label], + column_label_names: List[Label], + fields: List[InternalField] = None, + ) -> pd.DataFrame: + """ + Restore pandas DataFrame indices using the metadata. + + :param pdf: the pandas DataFrame to be processed. + :param index_columns: the original column names for index columns. + :param index_names: the index names after restored. + :param data_columns: the original column names for data columns. + :param column_labels: the column labels after restored. + :param column_label_names: the column label names after restored. + :param fields: the fields after restored. + :return: the restored pandas DataFrame + + >>> from numpy import dtype + >>> pdf = pd.DataFrame({"index": [10, 20, 30], "a": ['a', 'b', 'c'], "b": [0, 2, 1]}) + >>> InternalFrame.restore_index( + ... pdf, + ... index_columns=["index"], + ... index_names=[("idx",)], + ... data_columns=["a", "b", "index"], + ... column_labels=[("x",), ("y",), ("z",)], + ... column_label_names=[("lv1",)], + ... fields=[ + ... InternalField( + ... dtype=dtype('int64'), + ... struct_field=StructField(name='index', dataType=LongType(), nullable=False), + ... ), + ... InternalField( + ... dtype=dtype('object'), + ... struct_field=StructField(name='a', dataType=StringType(), nullable=False), + ... ), + ... InternalField( + ... dtype=CategoricalDtype(categories=["i", "j", "k"]), + ... struct_field=StructField(name='b', dataType=LongType(), nullable=False), + ... ), + ... ], + ... ) # doctest: +NORMALIZE_WHITESPACE + lv1 x y z + idx + 10 a i 10 + 20 b k 20 + 30 c j 30 + """ + for col, field in zip(pdf.columns, fields): + pdf[col] = DataTypeOps(field.dtype, field.spark_type).restore(pdf[col]) + + append = False + for index_field in index_columns: + drop = index_field not in data_columns + pdf = pdf.set_index(index_field, drop=drop, append=append) + append = True + pdf = pdf[data_columns] + + pdf.index.names = [ + name if name is None or len(name) > 1 else name[0] for name in index_names + ] + + names = [name if name is None or len(name) > 1 else name[0] for name in column_label_names] + if len(column_label_names) > 1: + pdf.columns = pd.MultiIndex.from_tuples(column_labels, names=names) + else: + pdf.columns = pd.Index( + [None if label is None else label[0] for label in column_labels], + name=names[0], + ) + + return pdf + + @lazy_property + def resolved_copy(self) -> "InternalFrame": + """Copy the immutable InternalFrame with the updates resolved.""" + sdf = self.spark_frame.select(self.spark_columns + list(HIDDEN_COLUMNS)) + return self.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in self.index_spark_column_names], + data_spark_columns=[scol_for(sdf, col) for col in self.data_spark_column_names], + ) + + def with_new_sdf( + self, + spark_frame: PySparkDataFrame, + *, + index_fields: Optional[List[InternalField]] = None, + data_columns: Optional[List[str]] = None, + data_fields: Optional[List[InternalField]] = None, + ) -> "InternalFrame": + """Copy the immutable InternalFrame with the updates by the specified Spark DataFrame. + + :param spark_frame: the new Spark DataFrame + :param index_fields: the new InternalFields for the index columns. + If None, the original dtyeps are used. + :param data_columns: the new column names. If None, the original one is used. + :param data_fields: the new InternalFields for the data columns. + If None, the original dtyeps are used. + :return: the copied InternalFrame. + """ + if index_fields is None: + index_fields = self.index_fields + else: + assert len(index_fields) == len(self.index_fields), ( + len(index_fields), + len(self.index_fields), + ) + + if data_columns is None: + data_columns = self.data_spark_column_names + else: + assert len(data_columns) == len(self.column_labels), ( + len(data_columns), + len(self.column_labels), + ) + + if data_fields is None: + data_fields = self.data_fields + else: + assert len(data_fields) == len(self.column_labels), ( + len(data_fields), + len(self.column_labels), + ) + + sdf = spark_frame.drop(NATURAL_ORDER_COLUMN_NAME) + return self.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in self.index_spark_column_names], + index_fields=index_fields, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + data_fields=data_fields, + ) + + def with_new_columns( + self, + scols_or_pssers: Sequence[Union[PySparkColumn, "Series"]], + *, + column_labels: Optional[List[Label]] = None, + data_fields: Optional[List[InternalField]] = None, + column_label_names: Union[Optional[List[Optional[Label]]], _NoValueType] = _NoValue, + keep_order: bool = True, + ) -> "InternalFrame": + """ + Copy the immutable InternalFrame with the updates by the specified Spark Columns or Series. + + :param scols_or_pssers: the new Spark Columns or Series. + :param column_labels: the new column index. + If None, the column_labels of the corresponding `scols_or_pssers` is used if it is + Series; otherwise the original one is used. + :param data_fields: the new InternalFields for the data columns. + If None, the dtypes of the corresponding `scols_or_pssers` is used if it is Series; + otherwise the dtypes will be inferred from the corresponding `scols_or_pssers`. + :param column_label_names: the new names of the column index levels. + :return: the copied InternalFrame. + """ + from pyspark.pandas.series import Series + + if column_labels is None: + if all(isinstance(scol_or_psser, Series) for scol_or_psser in scols_or_pssers): + column_labels = [cast(Series, psser)._column_label for psser in scols_or_pssers] + else: + assert len(scols_or_pssers) == len(self.column_labels), ( + len(scols_or_pssers), + len(self.column_labels), + ) + column_labels = [] + for scol_or_psser, label in zip(scols_or_pssers, self.column_labels): + if isinstance(scol_or_psser, Series): + column_labels.append(scol_or_psser._column_label) + else: + column_labels.append(label) + else: + assert len(scols_or_pssers) == len(column_labels), ( + len(scols_or_pssers), + len(column_labels), + ) + + data_spark_columns = [] + for scol_or_psser in scols_or_pssers: + if isinstance(scol_or_psser, Series): + scol = scol_or_psser.spark.column + else: + scol = scol_or_psser + data_spark_columns.append(scol) + + if data_fields is None: + data_fields = [] + for scol_or_psser in scols_or_pssers: + if isinstance(scol_or_psser, Series): + data_fields.append(scol_or_psser._internal.data_fields[0]) + else: + data_fields.append(None) + else: + assert len(scols_or_pssers) == len(data_fields), ( + len(scols_or_pssers), + len(data_fields), + ) + + sdf = self.spark_frame + if not keep_order: + sdf = self.spark_frame.select(self.index_spark_columns + data_spark_columns) + index_spark_columns = [scol_for(sdf, col) for col in self.index_spark_column_names] + data_spark_columns = [ + scol_for(sdf, col) for col in self.spark_frame.select(data_spark_columns).columns + ] + else: + index_spark_columns = self.index_spark_columns + + if column_label_names is _NoValue: + column_label_names = self._column_label_names + + return self.copy( + spark_frame=sdf, + index_spark_columns=index_spark_columns, + column_labels=column_labels, + data_spark_columns=data_spark_columns, + data_fields=data_fields, + column_label_names=column_label_names, + ) + + def with_filter(self, pred: Union[PySparkColumn, "Series"]) -> "InternalFrame": + """ + Copy the immutable InternalFrame with the updates by the predicate. + + :param pred: the predicate to filter. + :return: the copied InternalFrame. + """ + from pyspark.pandas.series import Series + + if isinstance(pred, Series): + assert isinstance(pred.spark.data_type, BooleanType), pred.spark.data_type + condition = pred.spark.column + else: + condition = pred + spark_type = self.spark_frame.select(condition).schema[0].dataType + assert isinstance(spark_type, BooleanType), spark_type + + return self.with_new_sdf(self.spark_frame.filter(condition).select(self.spark_columns)) + + def with_new_spark_column( + self, + column_label: Label, + scol: PySparkColumn, + *, + field: Optional[InternalField] = None, + keep_order: bool = True, + ) -> "InternalFrame": + """ + Copy the immutable InternalFrame with the updates by the specified Spark Column. + + :param column_label: the column label to be updated. + :param scol: the new Spark Column + :param field: the new InternalField for the data column. + If not specified, the InternalField will be inferred from the spark Column. + :return: the copied InternalFrame. + """ + assert column_label in self.column_labels, column_label + + idx = self.column_labels.index(column_label) + data_spark_columns = self.data_spark_columns.copy() + data_spark_columns[idx] = scol + data_fields = self.data_fields.copy() + data_fields[idx] = field + return self.with_new_columns( + data_spark_columns, data_fields=data_fields, keep_order=keep_order + ) + + def select_column(self, column_label: Label) -> "InternalFrame": + """ + Copy the immutable InternalFrame with the specified column. + + :param column_label: the column label to use. + :return: the copied InternalFrame. + """ + assert column_label in self.column_labels, column_label + + return self.copy( + column_labels=[column_label], + data_spark_columns=[self.spark_column_for(column_label)], + data_fields=[self.field_for(column_label)], + column_label_names=None, + ) + + def copy( + self, + *, + spark_frame: Union[PySparkDataFrame, _NoValueType] = _NoValue, + index_spark_columns: Union[List[PySparkColumn], _NoValueType] = _NoValue, + index_names: Union[Optional[List[Optional[Label]]], _NoValueType] = _NoValue, + index_fields: Union[Optional[List[InternalField]], _NoValueType] = _NoValue, + column_labels: Union[Optional[List[Label]], _NoValueType] = _NoValue, + data_spark_columns: Union[Optional[List[PySparkColumn]], _NoValueType] = _NoValue, + data_fields: Union[Optional[List[InternalField]], _NoValueType] = _NoValue, + column_label_names: Union[Optional[List[Optional[Label]]], _NoValueType] = _NoValue, + ) -> "InternalFrame": + """ + Copy the immutable InternalFrame. + + :param spark_frame: the new Spark DataFrame. If not specified, the original one is used. + :param index_spark_columns: the list of Spark Column. + If not specified, the original ones are used. + :param index_names: the index names. If not specified, the original ones are used. + :param index_fields: the new InternalFields for the index columns. + If not specified, the original metadata are used. + :param column_labels: the new column labels. If not specified, the original ones are used. + :param data_spark_columns: the new Spark Columns. + If not specified, the original ones are used. + :param data_fields: the new InternalFields for the data columns. + If not specified, the original metadata are used. + :param column_label_names: the new names of the column index levels. + If not specified, the original ones are used. + :return: the copied immutable InternalFrame. + """ + if spark_frame is _NoValue: + spark_frame = self.spark_frame + if index_spark_columns is _NoValue: + index_spark_columns = self.index_spark_columns + if index_names is _NoValue: + index_names = self.index_names + if index_fields is _NoValue: + index_fields = self.index_fields + if column_labels is _NoValue: + column_labels = self.column_labels + if data_spark_columns is _NoValue: + data_spark_columns = self.data_spark_columns + if data_fields is _NoValue: + data_fields = self.data_fields + if column_label_names is _NoValue: + column_label_names = self.column_label_names + return InternalFrame( + spark_frame=cast(PySparkDataFrame, spark_frame), + index_spark_columns=cast(List[PySparkColumn], index_spark_columns), + index_names=cast(Optional[List[Optional[Label]]], index_names), + index_fields=cast(Optional[List[InternalField]], index_fields), + column_labels=cast(Optional[List[Label]], column_labels), + data_spark_columns=cast(Optional[List[PySparkColumn]], data_spark_columns), + data_fields=cast(Optional[List[InternalField]], data_fields), + column_label_names=cast(Optional[List[Optional[Label]]], column_label_names), + ) + + @staticmethod + def from_pandas(pdf: pd.DataFrame) -> "InternalFrame": + """Create an immutable DataFrame from pandas DataFrame. + + :param pdf: :class:`pd.DataFrame` + :return: the created immutable DataFrame + """ + + index_names: List[Optional[Label]] = [ + name if name is None or isinstance(name, tuple) else (name,) for name in pdf.index.names + ] + + columns = pdf.columns + column_labels: List[Label] + if isinstance(columns, pd.MultiIndex): + column_labels = columns.tolist() + else: + column_labels = [(col,) for col in columns] + + column_label_names: List[Optional[Label]] = [ + name if name is None or isinstance(name, tuple) else (name,) for name in columns.names + ] + + prefer_timestamp_ntz = is_timestamp_ntz_preferred() + + ( + pdf, + index_columns, + index_fields, + data_columns, + data_fields, + ) = InternalFrame.prepare_pandas_frame(pdf, prefer_timestamp_ntz=prefer_timestamp_ntz) + + schema = StructType([field.struct_field for field in index_fields + data_fields]) + + sdf = default_session().createDataFrame(pdf, schema=schema) + return InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_columns], + index_names=index_names, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + data_fields=data_fields, + column_label_names=column_label_names, + ) + + @staticmethod + def prepare_pandas_frame( + pdf: pd.DataFrame, *, retain_index: bool = True, prefer_timestamp_ntz: bool = False + ) -> Tuple[pd.DataFrame, List[str], List[InternalField], List[str], List[InternalField]]: + """ + Prepare pandas DataFrame for creating Spark DataFrame. + + :param pdf: the pandas DataFrame to be prepared. + :param retain_index: whether the indices should be retained. + :return: the tuple of + - the prepared pandas dataFrame + - index column names for Spark DataFrame + - the InternalFields for the index columns of the given pandas DataFrame + - data column names for Spark DataFrame + - the InternalFields for the data columns of the given pandas DataFrame + + >>> pdf = pd.DataFrame( + ... {("x", "a"): ['a', 'b', 'c'], + ... ("y", "b"): pd.Categorical(["i", "k", "j"], categories=["i", "j", "k"])}, + ... index=[10, 20, 30]) + >>> prepared, index_columns, index_fields, data_columns, data_fields = ( + ... InternalFrame.prepare_pandas_frame(pdf) + ... ) + >>> prepared + __index_level_0__ (x, a) (y, b) + 0 10 a 0 + 1 20 b 2 + 2 30 c 1 + >>> index_columns + ['__index_level_0__'] + >>> index_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', + LongType(), False))] + >>> data_columns + ['(x, a)', '(y, b)'] + >>> data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=object, struct_field=StructField('(x, a)', StringType(), False)), + InternalField(dtype=category, struct_field=StructField('(y, b)', ByteType(), False))] + + >>> import datetime + >>> pdf = pd.DataFrame({ + ... "dt": [datetime.datetime(1970, 1, 1)], "dt_obj": [datetime.datetime(1970, 1, 1)] + ... }) + >>> pdf.dt_obj = pdf.dt_obj.astype("object") + >>> _, _, _, _, data_fields = ( + ... InternalFrame.prepare_pandas_frame(pdf, prefer_timestamp_ntz=True) + ... ) + >>> data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=datetime64[ns], + struct_field=StructField('dt', TimestampNTZType(), False)), + InternalField(dtype=object, + struct_field=StructField('dt_obj', TimestampNTZType(), False))] + + >>> pdf = pd.DataFrame({ + ... "td": [datetime.timedelta(0)], "td_obj": [datetime.timedelta(0)] + ... }) + >>> pdf.td_obj = pdf.td_obj.astype("object") + >>> _, _, _, _, data_fields = ( + ... InternalFrame.prepare_pandas_frame(pdf) + ... ) + >>> data_fields # doctest: +NORMALIZE_WHITESPACE + [InternalField(dtype=timedelta64[ns], + struct_field=StructField('td', DayTimeIntervalType(0, 3), False)), + InternalField(dtype=object, + struct_field=StructField('td_obj', DayTimeIntervalType(0, 3), False))] + """ + pdf = pdf.copy() + + data_columns = [name_like_string(col) for col in pdf.columns] + pdf.columns = data_columns + + if retain_index: + index_nlevels = pdf.index.nlevels + index_columns = [SPARK_INDEX_NAME_FORMAT(i) for i in range(index_nlevels)] + pdf.index.names = index_columns + reset_index = pdf.reset_index() + else: + index_nlevels = 0 + index_columns = [] + reset_index = pdf + + index_dtypes = list(reset_index.dtypes)[:index_nlevels] + data_dtypes = list(reset_index.dtypes)[index_nlevels:] + + for col, dtype in zip(reset_index.columns, reset_index.dtypes): + spark_type = infer_pd_series_spark_type(reset_index[col], dtype, prefer_timestamp_ntz) + reset_index[col] = DataTypeOps(dtype, spark_type).prepare(reset_index[col]) + + fields = [ + InternalField( + dtype=dtype, + struct_field=StructField( + name=str(name), + dataType=infer_pd_series_spark_type(col, dtype, prefer_timestamp_ntz), + nullable=bool(col.isnull().any()), + ), + ) + for (name, col), dtype in zip(reset_index.items(), index_dtypes + data_dtypes) + ] + + return ( + reset_index, + index_columns, + fields[:index_nlevels], + data_columns, + fields[index_nlevels:], + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.internal + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.internal.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.internal tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.internal, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/__init__.py new file mode 100644 index 0000000..bfe1597 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/__init__.py @@ -0,0 +1,48 @@ +# +# 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.pandas.exceptions import PandasNotImplementedError + + +def unsupported_function(class_name, method_name, deprecated=False, reason=""): + def unsupported_function(*args, **kwargs): + raise PandasNotImplementedError( + class_name=class_name, method_name=method_name, reason=reason + ) + + def deprecated_function(*args, **kwargs): + raise PandasNotImplementedError( + class_name=class_name, method_name=method_name, deprecated=deprecated, reason=reason + ) + + return deprecated_function if deprecated else unsupported_function + + +def unsupported_property(class_name, property_name, deprecated=False, reason=""): + @property + def unsupported_property(self): + raise PandasNotImplementedError( + class_name=class_name, property_name=property_name, reason=reason + ) + + @property + def deprecated_property(self): + raise PandasNotImplementedError( + class_name=class_name, property_name=property_name, deprecated=deprecated, reason=reason + ) + + return deprecated_property if deprecated else unsupported_property diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/common.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/common.py new file mode 100644 index 0000000..e6530a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/common.py @@ -0,0 +1,76 @@ +# +# 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. +# + + +def memory_usage(f): + return f( + "memory_usage", + reason="Unlike pandas, most DataFrames are not materialized in memory in Spark " + "(and pandas-on-Spark), and as a result memory_usage() does not do what you intend it " + "to do. Use Spark's web UI to monitor disk and memory usage of your application.", + ) + + +def array(f): + return f( + "array", + reason="If you want to collect your data as an NumPy array, use 'to_numpy()' instead.", + ) + + +def to_pickle(f): + return f( + "to_pickle", + reason="For storage, we encourage you to use Delta or Parquet, instead of Python pickle " + "format.", + ) + + +def to_xarray(f): + return f( + "to_xarray", + reason="If you want to collect your data as an NumPy array, use 'to_numpy()' instead.", + ) + + +def to_list(f): + return f( + "to_list", + reason="If you want to collect your data as an NumPy array, use 'to_numpy()' instead.", + ) + + +def tolist(f): + return f( + "tolist", + reason="If you want to collect your data as an NumPy array, use 'to_numpy()' instead.", + ) + + +def __iter__(f): + return f( + "__iter__", + reason="If you want to collect your data as an NumPy array, use 'to_numpy()' instead.", + ) + + +def duplicated(f): + return f( + "duplicated", + reason="'duplicated' API returns np.ndarray and the data size is too large." + "You can just use DataFrame.deduplicated instead", + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/frame.py new file mode 100644 index 0000000..5ba81c8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/frame.py @@ -0,0 +1,63 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property, common + + +def _unsupported_function(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pd.DataFrame", method_name=method_name, deprecated=deprecated, reason=reason + ) + + +def _unsupported_property(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pd.DataFrame", property_name=property_name, deprecated=deprecated, reason=reason + ) + + +class MissingPandasLikeDataFrame: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Functions + asfreq = _unsupported_function("asfreq") + asof = _unsupported_function("asof") + combine = _unsupported_function("combine") + compare = _unsupported_function("compare") + convert_dtypes = _unsupported_function("convert_dtypes") + infer_objects = _unsupported_function("infer_objects") + reorder_levels = _unsupported_function("reorder_levels") + set_axis = _unsupported_function("set_axis") + to_feather = _unsupported_function("to_feather") + to_gbq = _unsupported_function("to_gbq") + to_hdf = _unsupported_function("to_hdf") + to_period = _unsupported_function("to_period") + to_sql = _unsupported_function("to_sql") + to_stata = _unsupported_function("to_stata") + to_timestamp = _unsupported_function("to_timestamp") + tz_convert = _unsupported_function("tz_convert") + tz_localize = _unsupported_function("tz_localize") + + # Deprecated functions + lookup = _unsupported_function( + "lookup", deprecated=True, reason="Use DataFrame.melt and DataFrame.loc instead." + ) + + # Functions we won't support. + to_pickle = common.to_pickle(_unsupported_function) + memory_usage = common.memory_usage(_unsupported_function) + to_xarray = common.to_xarray(_unsupported_function) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/general_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/general_functions.py new file mode 100644 index 0000000..4713331 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/general_functions.py @@ -0,0 +1,43 @@ +# +# 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.pandas.missing import unsupported_function + + +def _unsupported_function(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pd", method_name=method_name, deprecated=deprecated, reason=reason + ) + + +class MissingPandasLikeGeneralFunctions: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + pivot = _unsupported_function("pivot") + pivot_table = _unsupported_function("pivot_table") + crosstab = _unsupported_function("crosstab") + cut = _unsupported_function("cut") + qcut = _unsupported_function("qcut") + merge_ordered = _unsupported_function("merge_ordered") + factorize = _unsupported_function("factorize") + unique = _unsupported_function("unique") + wide_to_long = _unsupported_function("wide_to_long") + bdate_range = _unsupported_function("bdate_range") + period_range = _unsupported_function("period_range") + infer_freq = _unsupported_function("infer_freq") + interval_range = _unsupported_function("interval_range") + eval = _unsupported_function("eval") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/groupby.py new file mode 100644 index 0000000..55a4a1d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/groupby.py @@ -0,0 +1,93 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property + + +def _unsupported_function(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pd.groupby.GroupBy", + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pd.groupby.GroupBy", + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +class MissingPandasLikeDataFrameGroupBy: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + corr = _unsupported_property("corr") + corrwith = _unsupported_property("corrwith") + cov = _unsupported_property("cov") + dtypes = _unsupported_property("dtypes") + groups = _unsupported_property("groups") + hist = _unsupported_property("hist") + indices = _unsupported_property("indices") + ngroups = _unsupported_property("ngroups") + plot = _unsupported_property("plot") + + # Deprecated properties + take = _unsupported_property("take", deprecated=True) + + # Functions + boxplot = _unsupported_function("boxplot") + ngroup = _unsupported_function("ngroup") + ohlc = _unsupported_function("ohlc") + pct_change = _unsupported_function("pct_change") + pipe = _unsupported_function("pipe") + resample = _unsupported_function("resample") + + +class MissingPandasLikeSeriesGroupBy: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + corr = _unsupported_property("corr") + cov = _unsupported_property("cov") + dtype = _unsupported_property("dtype") + groups = _unsupported_property("groups") + hist = _unsupported_property("hist") + indices = _unsupported_property("indices") + is_monotonic_decreasing = _unsupported_property("is_monotonic_decreasing") + is_monotonic_increasing = _unsupported_property("is_monotonic_increasing") + ngroups = _unsupported_property("ngroups") + plot = _unsupported_property("plot") + + # Deprecated properties + take = _unsupported_property("take", deprecated=True) + + # 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") + pipe = _unsupported_function("pipe") + resample = _unsupported_function("resample") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/indexes.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/indexes.py new file mode 100644 index 0000000..2419908 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/indexes.py @@ -0,0 +1,184 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property, common + + +def _unsupported_function(method_name, deprecated=False, reason="", cls="Index"): + return unsupported_function( + class_name="pd.{}".format(cls), + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property(property_name, deprecated=False, reason="", cls="Index"): + return unsupported_property( + class_name="pd.{}".format(cls), + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +class MissingPandasLikeIndex: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + nbytes = _unsupported_property("nbytes") + + # Functions + argsort = _unsupported_function("argsort") + asof_locs = _unsupported_function("asof_locs") + format = _unsupported_function("format") + get_indexer = _unsupported_function("get_indexer") + get_indexer_for = _unsupported_function("get_indexer_for") + get_indexer_non_unique = _unsupported_function("get_indexer_non_unique") + get_loc = _unsupported_function("get_loc") + get_slice_bound = _unsupported_function("get_slice_bound") + groupby = _unsupported_function("groupby") + is_ = _unsupported_function("is_") + join = _unsupported_function("join") + putmask = _unsupported_function("putmask") + ravel = _unsupported_function("ravel") + reindex = _unsupported_function("reindex") + searchsorted = _unsupported_function("searchsorted") + slice_indexer = _unsupported_function("slice_indexer") + slice_locs = _unsupported_function("slice_locs") + sortlevel = _unsupported_function("sortlevel") + to_flat_index = _unsupported_function("to_flat_index") + where = _unsupported_function("where") + is_mixed = _unsupported_function("is_mixed") + + # Deprecated functions + to_native_types = _unsupported_function("to_native_types", deprecated=True) + + # Properties we won't support. + array = common.array(_unsupported_property) + duplicated = common.duplicated(_unsupported_property) + + # Functions we won't support. + memory_usage = common.memory_usage(_unsupported_function) + __iter__ = common.__iter__(_unsupported_function) + + +class MissingPandasLikeDatetimeIndex(MissingPandasLikeIndex): + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + nanosecond = _unsupported_property("nanosecond", cls="DatetimeIndex") + date = _unsupported_property("date", cls="DatetimeIndex") + time = _unsupported_property("time", cls="DatetimeIndex") + timetz = _unsupported_property("timetz", cls="DatetimeIndex") + tz = _unsupported_property("tz", cls="DatetimeIndex") + freq = _unsupported_property("freq", cls="DatetimeIndex") + freqstr = _unsupported_property("freqstr", cls="DatetimeIndex") + inferred_freq = _unsupported_property("inferred_freq", cls="DatetimeIndex") + + # Functions + snap = _unsupported_function("snap", cls="DatetimeIndex") + tz_convert = _unsupported_function("tz_convert", cls="DatetimeIndex") + tz_localize = _unsupported_function("tz_localize", cls="DatetimeIndex") + to_period = _unsupported_function("to_period", cls="DatetimeIndex") + to_perioddelta = _unsupported_function("to_perioddelta", cls="DatetimeIndex") + to_pydatetime = _unsupported_function("to_pydatetime", cls="DatetimeIndex") + mean = _unsupported_function("mean", cls="DatetimeIndex") + std = _unsupported_function("std", cls="DatetimeIndex") + + +class MissingPandasLikeTimedeltaIndex(MissingPandasLikeIndex): + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + nanoseconds = _unsupported_property("nanoseconds", cls="TimedeltaIndex") + components = _unsupported_property("components", cls="TimedeltaIndex") + inferred_freq = _unsupported_property("inferred_freq", cls="TimedeltaIndex") + + # Functions + to_pytimedelta = _unsupported_function("to_pytimedelta", cls="TimedeltaIndex") + round = _unsupported_function("round", cls="TimedeltaIndex") + floor = _unsupported_function("floor", cls="TimedeltaIndex") + ceil = _unsupported_function("ceil", cls="TimedeltaIndex") + mean = _unsupported_function("mean", cls="TimedeltaIndex") + + +class MissingPandasLikeMultiIndex: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Functions + argsort = _unsupported_function("argsort") + asof_locs = _unsupported_function("asof_locs") + factorize = _unsupported_function("factorize") + format = _unsupported_function("format") + get_indexer = _unsupported_function("get_indexer") + get_indexer_for = _unsupported_function("get_indexer_for") + get_indexer_non_unique = _unsupported_function("get_indexer_non_unique") + get_loc = _unsupported_function("get_loc") + get_loc_level = _unsupported_function("get_loc_level") + get_locs = _unsupported_function("get_locs") + get_slice_bound = _unsupported_function("get_slice_bound") + get_value = _unsupported_function("get_value") + groupby = _unsupported_function("groupby") + is_ = _unsupported_function("is_") + is_lexsorted = _unsupported_function("is_lexsorted") + join = _unsupported_function("join") + map = _unsupported_function("map") + putmask = _unsupported_function("putmask") + ravel = _unsupported_function("ravel") + reindex = _unsupported_function("reindex") + remove_unused_levels = _unsupported_function("remove_unused_levels") + reorder_levels = _unsupported_function("reorder_levels") + searchsorted = _unsupported_function("searchsorted") + set_codes = _unsupported_function("set_codes") + set_levels = _unsupported_function("set_levels") + slice_indexer = _unsupported_function("slice_indexer") + slice_locs = _unsupported_function("slice_locs") + sortlevel = _unsupported_function("sortlevel") + to_flat_index = _unsupported_function("to_flat_index") + truncate = _unsupported_function("truncate") + where = _unsupported_function("where") + + # Deprecated functions + is_mixed = _unsupported_function( + "is_mixed", deprecated=True, reason="Check index.inferred_type directly instead." + ) + set_value = _unsupported_function("set_value", deprecated=True) + to_native_types = _unsupported_function("to_native_types", deprecated=True) + + # Functions we won't support. + array = common.array(_unsupported_property) + duplicated = common.duplicated(_unsupported_property) + codes = _unsupported_property( + "codes", + reason="'codes' requires to collect all data into the driver which is against the " + "design principle of pandas-on-Spark. Alternatively, you could call 'to_pandas()' and" + " use 'codes' property in pandas.", + ) + levels = _unsupported_property( + "levels", + reason="'levels' requires to collect all data into the driver which is against the " + "design principle of pandas-on-Spark. Alternatively, you could call 'to_pandas()' and" + " use 'levels' property in pandas.", + ) + __iter__ = common.__iter__(_unsupported_function) + + # Properties we won't support. + memory_usage = common.memory_usage(_unsupported_function) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/resample.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/resample.py new file mode 100644 index 0000000..549e3c5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/resample.py @@ -0,0 +1,101 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property + + +def _unsupported_function(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pd.resample.Resampler", + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pd.resample.Resampler", + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +class MissingPandasLikeDataFrameResampler: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + groups = _unsupported_property("groups") + indices = _unsupported_property("indices") + + # Functions + get_group = _unsupported_property("get_group") + apply = _unsupported_function("apply") + aggregate = _unsupported_function("aggregate") + transform = _unsupported_function("transform") + pipe = _unsupported_function("pipe") + ffill = _unsupported_function("ffill") + bfill = _unsupported_function("bfill") + nearest = _unsupported_function("nearest") + fillna = _unsupported_function("fillna") + asfreq = _unsupported_function("asfreq") + interpolate = _unsupported_function("interpolate") + count = _unsupported_function("count") + nunique = _unsupported_function("nunique") + first = _unsupported_function("first") + last = _unsupported_function("last") + median = _unsupported_function("median") + ohlc = _unsupported_function("ohlc") + prod = _unsupported_function("prod") + size = _unsupported_function("size") + sem = _unsupported_function("sem") + quantile = _unsupported_function("quantile") + + +class MissingPandasLikeSeriesResampler: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Properties + groups = _unsupported_property("groups") + indices = _unsupported_property("indices") + + # Functions + get_group = _unsupported_property("get_group") + apply = _unsupported_function("apply") + aggregate = _unsupported_function("aggregate") + transform = _unsupported_function("transform") + pipe = _unsupported_function("pipe") + ffill = _unsupported_function("ffill") + bfill = _unsupported_function("bfill") + nearest = _unsupported_function("nearest") + fillna = _unsupported_function("fillna") + asfreq = _unsupported_function("asfreq") + interpolate = _unsupported_function("interpolate") + count = _unsupported_function("count") + nunique = _unsupported_function("nunique") + first = _unsupported_function("first") + last = _unsupported_function("last") + median = _unsupported_function("median") + ohlc = _unsupported_function("ohlc") + prod = _unsupported_function("prod") + size = _unsupported_function("size") + sem = _unsupported_function("sem") + quantile = _unsupported_function("quantile") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/scalars.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/scalars.py new file mode 100644 index 0000000..e4b2831 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/scalars.py @@ -0,0 +1,29 @@ +# +# 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.pandas.exceptions import PandasNotImplementedError + + +def _unsupported_scalar(scalar_name): + return PandasNotImplementedError(class_name="ps", scalar_name=scalar_name) + + +class MissingPandasLikeScalars: + Timestamp = _unsupported_scalar("Timestamp") + Timedelta = _unsupported_scalar("Timedelta") + Period = _unsupported_scalar("Period") + Interval = _unsupported_scalar("Interval") + Categorical = _unsupported_scalar("Categorical") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/series.py new file mode 100644 index 0000000..4ee860d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/series.py @@ -0,0 +1,69 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property, common + + +def _unsupported_function(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pd.Series", method_name=method_name, deprecated=deprecated, reason=reason + ) + + +def _unsupported_property(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pd.Series", property_name=property_name, deprecated=deprecated, reason=reason + ) + + +class MissingPandasLikeSeries: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + # Functions + asfreq = _unsupported_function("asfreq") + combine = _unsupported_function("combine") + convert_dtypes = _unsupported_function("convert_dtypes") + infer_objects = _unsupported_function("infer_objects") + reorder_levels = _unsupported_function("reorder_levels") + set_axis = _unsupported_function("set_axis") + to_hdf = _unsupported_function("to_hdf") + to_period = _unsupported_function("to_period") + to_sql = _unsupported_function("to_sql") + to_timestamp = _unsupported_function("to_timestamp") + tz_convert = _unsupported_function("tz_convert") + tz_localize = _unsupported_function("tz_localize") + view = _unsupported_function("view") + + # Properties we won't support. + array = common.array(_unsupported_property) + nbytes = _unsupported_property( + "nbytes", + reason="'nbytes' requires to compute whole dataset. You can calculate manually it, " + "with its 'itemsize', by explicitly executing its count. Use Spark's web UI " + "to monitor disk and memory usage of your application in general.", + ) + + # Functions we won't support. + memory_usage = common.memory_usage(_unsupported_function) + to_pickle = common.to_pickle(_unsupported_function) + to_xarray = common.to_xarray(_unsupported_function) + __iter__ = common.__iter__(_unsupported_function) + ravel = _unsupported_function( + "ravel", + reason="If you want to collect your flattened underlying data as an NumPy array, " + "use 'to_numpy().ravel()' instead.", + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/window.py new file mode 100644 index 0000000..a6d423d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/missing/window.py @@ -0,0 +1,168 @@ +# +# 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.pandas.missing import unsupported_function, unsupported_property + + +def _unsupported_function_expanding(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pandas.core.window.Expanding", + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property_expanding(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pandas.core.window.Expanding", + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_function_rolling(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pandas.core.window.Rolling", + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property_rolling(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pandas.core.window.Rolling", + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_function_exponential_moving(method_name, deprecated=False, reason=""): + return unsupported_function( + class_name="pandas.core.window.ExponentialMovingWindow", + method_name=method_name, + deprecated=deprecated, + reason=reason, + ) + + +def _unsupported_property_exponential_moving(property_name, deprecated=False, reason=""): + return unsupported_property( + class_name="pandas.core.window.ExponentialMovingWindow", + property_name=property_name, + deprecated=deprecated, + reason=reason, + ) + + +class MissingPandasLikeExpanding: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + agg = _unsupported_function_expanding("agg") + aggregate = _unsupported_function_expanding("aggregate") + 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") + is_datetimelike = _unsupported_property_expanding("is_datetimelike") + is_freq_type = _unsupported_property_expanding("is_freq_type") + ndim = _unsupported_property_expanding("ndim") + + +class MissingPandasLikeRolling: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + agg = _unsupported_function_rolling("agg") + aggregate = _unsupported_function_rolling("aggregate") + 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") + is_datetimelike = _unsupported_property_rolling("is_datetimelike") + is_freq_type = _unsupported_property_rolling("is_freq_type") + ndim = _unsupported_property_rolling("ndim") + + +class MissingPandasLikeExpandingGroupby: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + agg = _unsupported_function_expanding("agg") + aggregate = _unsupported_function_expanding("aggregate") + 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") + is_datetimelike = _unsupported_property_expanding("is_datetimelike") + is_freq_type = _unsupported_property_expanding("is_freq_type") + ndim = _unsupported_property_expanding("ndim") + + +class MissingPandasLikeRollingGroupby: + # NOTE: Please update the pandas-on-Spark reference document when implementing the new API. + # Documentation path: `python/docs/source/reference/pyspark.pandas/`. + + agg = _unsupported_function_rolling("agg") + aggregate = _unsupported_function_rolling("aggregate") + 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") + is_datetimelike = _unsupported_property_rolling("is_datetimelike") + is_freq_type = _unsupported_property_rolling("is_freq_type") + ndim = _unsupported_property_rolling("ndim") + + +class MissingPandasLikeExponentialMoving: + sum = _unsupported_function_exponential_moving("sum") + var = _unsupported_function_exponential_moving("var") + std = _unsupported_function_exponential_moving("std") + cov = _unsupported_function_exponential_moving("cov") + corr = _unsupported_function_exponential_moving("corr") + + adjust = _unsupported_property_exponential_moving("adjust") + axis = _unsupported_property_exponential_moving("axis") + method = _unsupported_property_exponential_moving("method") + + +class MissingPandasLikeExponentialMovingGroupby: + sum = _unsupported_function_exponential_moving("sum") + var = _unsupported_function_exponential_moving("var") + std = _unsupported_function_exponential_moving("std") + cov = _unsupported_function_exponential_moving("cov") + corr = _unsupported_function_exponential_moving("corr") + + adjust = _unsupported_property_exponential_moving("adjust") + axis = _unsupported_property_exponential_moving("axis") + method = _unsupported_property_exponential_moving("method") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/mlflow.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/mlflow.py new file mode 100644 index 0000000..b78ae93 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/mlflow.py @@ -0,0 +1,238 @@ +# +# 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. +# + +""" +MLflow-related functions to load models and apply them to pandas-on-Spark dataframes. +""" +from typing import List, Union + +from pyspark.sql.types import DataType +import pandas as pd +import numpy as np +from typing import Any + +from pyspark.pandas._typing import Label, Dtype +from pyspark.pandas.utils import lazy_property, default_session +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.typedef import as_spark_type + +__all__ = ["PythonModelWrapper", "load_model"] + + +class PythonModelWrapper: + """ + A wrapper around MLflow's Python object model. + + This wrapper acts as a predictor on pandas-on-Spark + + """ + + def __init__(self, model_uri: str, return_type_hint: Union[str, type, Dtype]): + self._model_uri = model_uri + self._return_type_hint = return_type_hint + + @lazy_property + def _return_type(self) -> DataType: + hint = self._return_type_hint + # The logic is simple for now, because it corresponds to the default + # case: continuous predictions + # TODO: do something smarter, for example when there is a sklearn.Classifier (it should + # return an integer or a categorical) + # We can do the same for pytorch/tensorflow/keras models by looking at the output types. + # However, this is probably better done in mlflow than here. + if hint == "infer" or not hint: + hint = np.float64 + return as_spark_type(hint) + + @lazy_property + def _model(self) -> Any: + """ + The return object has to follow the API of mlflow.pyfunc.PythonModel. + """ + from mlflow import pyfunc + + return pyfunc.load_model(model_uri=self._model_uri) + + @lazy_property + def _model_udf(self) -> Any: + from mlflow import pyfunc + + spark = default_session() + return pyfunc.spark_udf(spark, model_uri=self._model_uri, result_type=self._return_type) + + def __str__(self) -> str: + return "PythonModelWrapper({})".format(str(self._model)) + + def __repr__(self) -> str: + return "PythonModelWrapper({})".format(repr(self._model)) + + def predict(self, data: Union[DataFrame, pd.DataFrame]) -> Union[Series, pd.Series]: + """ + Returns a prediction on the data. + + If the data is a pandas-on-Spark DataFrame, the return is a pandas-on-Spark Series. + + If the data is a pandas Dataframe, the return is the expected output of the underlying + pyfunc object (typically a pandas Series or a numpy array). + """ + if isinstance(data, pd.DataFrame): + return self._model.predict(data) + elif isinstance(data, DataFrame): + return_col = self._model_udf(*data._internal.data_spark_columns) + # TODO: the columns should be named according to the mlflow spec + # However, this is only possible with spark >= 3.0 + # s = F.struct(*data.columns) + # return_col = self._model_udf(s) + column_labels: List[Label] = [ + (col,) for col in data._internal.spark_frame.select(return_col).columns + ] + internal = data._internal.copy( + column_labels=column_labels, data_spark_columns=[return_col], data_fields=None + ) + return first_series(DataFrame(internal)) + else: + raise ValueError("unknown data type: {}".format(type(data).__name__)) + + +def load_model( + model_uri: str, predict_type: Union[str, type, Dtype] = "infer" +) -> PythonModelWrapper: + """ + Loads an MLflow model into a wrapper that can be used both for pandas and pandas-on-Spark + DataFrame. + + Parameters + ---------- + model_uri : str + URI pointing to the model. See MLflow documentation for more details. + predict_type : a python basic type, a numpy basic type, a Spark type or 'infer'. + This is the return type that is expected when calling the predict function of the model. + If 'infer' is specified, the wrapper will attempt to automatically determine the return type + based on the model type. + + Returns + ------- + PythonModelWrapper + A wrapper around MLflow PythonModel objects. This wrapper is expected to adhere to the + interface of mlflow.pyfunc.PythonModel. + + Examples + -------- + Here is a full example that creates a model with scikit-learn and saves the model with + MLflow. The model is then loaded as a predictor that can be applied on a pandas-on-Spark + Dataframe. + + We first initialize our MLflow environment: + + >>> from mlflow.tracking import MlflowClient, set_tracking_uri + >>> import mlflow.sklearn + >>> from tempfile import mkdtemp + >>> d = mkdtemp("pandas_on_spark_mlflow") + >>> set_tracking_uri("file:%s"%d) + >>> client = MlflowClient() + >>> exp_id = mlflow.create_experiment("my_experiment") + >>> exp = mlflow.set_experiment("my_experiment") + + We aim at learning this numerical function using a simple linear regressor. + + >>> from sklearn.linear_model import LinearRegression + >>> train = pd.DataFrame({"x1": np.arange(8), "x2": np.arange(8)**2, + ... "y": np.log(2 + np.arange(8))}) + >>> train_x = train[["x1", "x2"]] + >>> train_y = train[["y"]] + >>> with mlflow.start_run(): + ... lr = LinearRegression() + ... lr.fit(train_x, train_y) + ... mlflow.sklearn.log_model(lr, "model") + LinearRegression... + + Now that our model is logged using MLflow, we load it back and apply it on a pandas-on-Spark + dataframe: + + >>> from pyspark.pandas.mlflow import load_model + >>> run_info = client.search_runs(exp_id)[-1].info + >>> model = load_model("runs:/{run_id}/model".format(run_id=run_info.run_id)) + >>> prediction_df = ps.DataFrame({"x1": [2.0], "x2": [4.0]}) + >>> prediction_df["prediction"] = model.predict(prediction_df) + >>> prediction_df + x1 x2 prediction + 0 2.0 4.0 1.355551 + + The model also works on pandas DataFrames as expected: + + >>> model.predict(prediction_df[["x1", "x2"]].to_pandas()) + array([[1.35555142]]) + + Notes + ----- + Currently, the model prediction can only be merged back with the existing dataframe. + Other columns must be manually joined. + For example, this code will not work: + + >>> df = ps.DataFrame({"x1": [2.0], "x2": [3.0], "z": [-1]}) + >>> features = df[["x1", "x2"]] + >>> y = model.predict(features) + >>> # Works: + >>> features["y"] = y # doctest: +SKIP + >>> # Will fail with a message about dataframes not aligned. + >>> df["y"] = y # doctest: +SKIP + + A current workaround is to use the .merge() function, using the feature values + as merging keys. + + >>> features['y'] = y + >>> everything = df.merge(features, on=['x1', 'x2']) + >>> everything + x1 x2 z y + 0 2.0 3.0 -1 1.376932 + """ + return PythonModelWrapper(model_uri, predict_type) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.mlflow + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.mlflow.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.mlflow tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.mlflow, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + try: + import mlflow # noqa: F401 + import sklearn # noqa: F401 + + _test() + except ImportError: + pass diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/namespace.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/namespace.py new file mode 100644 index 0000000..5689471 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/namespace.py @@ -0,0 +1,3807 @@ +# +# 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. +# + +""" +Wrappers around spark that correspond to common pandas functions. +""" +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Set, + Sized, + Tuple, + Type, + Union, + cast, + no_type_check, +) +from collections.abc import Iterable +from datetime import tzinfo +from functools import reduce +from io import BytesIO +import json +import warnings + +import numpy as np +import pandas as pd +from pandas.api.types import ( # type: ignore[attr-defined] + is_datetime64_dtype, + is_datetime64tz_dtype, + is_list_like, +) +from pandas.tseries.offsets import DateOffset +import pyarrow as pa +import pyarrow.parquet as pq +from pyspark.sql import functions as F, Column as PySparkColumn +from pyspark.sql.functions import pandas_udf +from pyspark.sql.types import ( + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + BooleanType, + TimestampType, + TimestampNTZType, + DecimalType, + StringType, + DateType, + StructType, + DataType, +) +from pyspark.sql.dataframe import DataFrame as PySparkDataFrame + +from pyspark import pandas as ps +from pyspark.pandas._typing import Axis, Dtype, Label, Name +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.utils import ( + align_diff_frames, + default_session, + is_name_like_tuple, + is_name_like_value, + name_like_string, + same_anchor, + scol_for, + validate_axis, + log_advice, +) +from pyspark.pandas.frame import DataFrame, _reduce_spark_multi +from pyspark.pandas.internal import ( + InternalFrame, + DEFAULT_SERIES_NAME, + HIDDEN_COLUMNS, + SPARK_INDEX_NAME_FORMAT, +) +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale +from pyspark.pandas.indexes import Index, DatetimeIndex, TimedeltaIndex +from pyspark.pandas.indexes.multi import MultiIndex + +# For Supporting Spark Connect +from pyspark.sql.utils import get_column_class + +__all__ = [ + "from_pandas", + "range", + "read_csv", + "read_delta", + "read_table", + "read_spark_io", + "read_parquet", + "read_clipboard", + "read_excel", + "read_html", + "to_datetime", + "date_range", + "to_timedelta", + "timedelta_range", + "get_dummies", + "concat", + "melt", + "isna", + "isnull", + "notna", + "notnull", + "read_sql_table", + "read_sql_query", + "read_sql", + "read_json", + "merge", + "merge_asof", + "to_numeric", + "broadcast", + "read_orc", +] + + +def from_pandas(pobj: Union[pd.DataFrame, pd.Series, pd.Index]) -> Union[Series, DataFrame, Index]: + """Create a pandas-on-Spark DataFrame, Series or Index from a pandas DataFrame, Series or Index. + + This is similar to Spark's `SparkSession.createDataFrame()` with pandas DataFrame, + but this also works with pandas Series and picks the index. + + Parameters + ---------- + pobj : pandas.DataFrame or pandas.Series + pandas DataFrame or Series to read. + + Returns + ------- + Series or DataFrame + If a pandas Series is passed in, this function returns a pandas-on-Spark Series. + If a pandas DataFrame is passed in, this function returns a pandas-on-Spark DataFrame. + """ + if isinstance(pobj, pd.Series): + return Series(pobj) + elif isinstance(pobj, pd.DataFrame): + return DataFrame(pobj) + elif isinstance(pobj, pd.Index): + return DataFrame(pd.DataFrame(index=pobj)).index + else: + raise TypeError("Unknown data type: {}".format(type(pobj).__name__)) + + +_range = range # built-in range + + +def range( + start: int, end: Optional[int] = None, step: int = 1, num_partitions: Optional[int] = None +) -> DataFrame: + """ + Create a DataFrame with some range of numbers. + + The resulting DataFrame has a single int64 column named `id`, containing elements in a range + from ``start`` to ``end`` (exclusive) with step value ``step``. If only the first parameter + (i.e. start) is specified, we treat it as the end value with the start value being 0. + + This is like the range function in SparkSession and is used primarily for testing. + + Parameters + ---------- + start : int + the start value (inclusive) + end : int, optional + the end value (exclusive) + step : int, optional, default 1 + the incremental step + num_partitions : int, optional + the number of partitions of the DataFrame + + Returns + ------- + DataFrame + + Examples + -------- + When the first parameter is specified, we generate a range of values up till that number. + + >>> ps.range(5) + id + 0 0 + 1 1 + 2 2 + 3 3 + 4 4 + + When start, end, and step are specified: + + >>> ps.range(start = 100, end = 200, step = 20) + id + 0 100 + 1 120 + 2 140 + 3 160 + 4 180 + """ + sdf = default_session().range(start=start, end=end, step=step, numPartitions=num_partitions) + return DataFrame(sdf) + + +def read_csv( + path: Union[str, List[str]], + sep: str = ",", + header: Union[str, int, None] = "infer", + names: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + usecols: Optional[Union[List[int], List[str], Callable[[str], bool]]] = None, + squeeze: bool = False, + mangle_dupe_cols: bool = True, + dtype: Optional[Union[str, Dtype, Dict[str, Union[str, Dtype]]]] = None, + nrows: Optional[int] = None, + parse_dates: bool = False, + quotechar: Optional[str] = None, + escapechar: Optional[str] = None, + comment: Optional[str] = None, + encoding: Optional[str] = None, + **options: Any, +) -> Union[DataFrame, Series]: + """Read CSV (comma-separated) file into DataFrame or Series. + + Parameters + ---------- + path : str or list + Path(s) of the CSV file(s) to be read. + sep : str, default ‘,’ + Delimiter to use. Non empty string. + header : int, default ‘infer’ + Whether to use the column names, and the start of the data. + Default behavior is to infer the column names: if no names are passed + the behavior is identical to `header=0` and column names are inferred from + the first line of the file, if column names are passed explicitly then + the behavior is identical to `header=None`. Explicitly pass `header=0` to be + able to replace existing names + names : str or array-like, optional + List of column names to use. If file contains no header row, then you should + explicitly pass `header=None`. Duplicates in this list will cause an error to be issued. + If a string is given, it should be a DDL-formatted string in Spark SQL, which is + preferred to avoid schema inference for better performance. + index_col: str or list of str, optional, default: None + Index column of table in Spark. + usecols : list-like or callable, optional + Return a subset of the columns. If list-like, all elements must either be + positional (i.e. integer indices into the document columns) or strings that + correspond to column names provided either by the user in names or inferred + from the document header row(s). + If callable, the callable function will be evaluated against the column names, + returning names where the callable function evaluates to `True`. + squeeze : bool, default False + If the parsed data only contains one column then return a Series. + + .. deprecated:: 3.4.0 + + mangle_dupe_cols : bool, default True + Duplicate columns will be specified as 'X0', 'X1', ... 'XN', rather + than 'X' ... 'X'. Passing in False will cause data to be overwritten if + there are duplicate names in the columns. + Currently only `True` is allowed. + + .. deprecated:: 3.4.0 + + dtype : Type name or dict of column -> type, default None + Data type for data or columns. E.g. {‘a’: np.float64, ‘b’: np.int32} Use str or object + together with suitable na_values settings to preserve and not interpret dtype. + nrows : int, default None + Number of rows to read from the CSV file. + parse_dates : boolean or list of ints or names or list of lists or dict, default `False`. + Currently only `False` is allowed. + quotechar : str (length 1), optional + The character used to denote the start and end of a quoted item. Quoted items can include + the delimiter and it will be ignored. + escapechar : str (length 1), default None + One-character string used to escape other characters. + comment: str, optional + Indicates the line should not be parsed. + encoding: str, optional + Indicates the encoding to read file + options : dict + All other options passed directly into Spark's data source. + + Returns + ------- + DataFrame or Series + + See Also + -------- + DataFrame.to_csv : Write DataFrame to a comma-separated values (csv) file. + + Examples + -------- + >>> ps.read_csv('data.csv') # doctest: +SKIP + + Load multiple CSV files as a single DataFrame: + + >>> ps.read_csv(['data-01.csv', 'data-02.csv']) # doctest: +SKIP + """ + # For latin-1 encoding is same as iso-8859-1, that's why its mapped to iso-8859-1. + encoding_mapping = {"latin-1": "iso-8859-1"} + + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + if mangle_dupe_cols is not True: + raise ValueError("mangle_dupe_cols can only be `True`: %s" % mangle_dupe_cols) + if parse_dates is not False: + raise ValueError("parse_dates can only be `False`: %s" % parse_dates) + + if usecols is not None and not callable(usecols): + usecols = list(usecols) # type: ignore[assignment] + + if usecols is None or callable(usecols) or len(usecols) > 0: + reader = default_session().read + reader.option("inferSchema", True) + reader.option("sep", sep) + + if header == "infer": + header = 0 if names is None else None + if header == 0: + reader.option("header", True) + elif header is None: + reader.option("header", False) + else: + raise ValueError("Unknown header argument {}".format(header)) + + if quotechar is not None: + reader.option("quote", quotechar) + if escapechar is not None: + reader.option("escape", escapechar) + + if comment is not None: + if not isinstance(comment, str) or len(comment) != 1: + raise ValueError("Only length-1 comment characters supported") + reader.option("comment", comment) + + reader.options(**options) + + if encoding is not None: + reader.option("encoding", encoding_mapping.get(encoding, encoding)) + + column_labels: Dict[Any, str] + if isinstance(names, str): + sdf = reader.schema(names).csv(path) + column_labels = {col: col for col in sdf.columns} + else: + sdf = reader.csv(path) + if is_list_like(names): + names = list(names) + if len(set(names)) != len(names): + raise ValueError("Found non-unique column index") + if len(names) != len(sdf.columns): + raise ValueError( + "The number of names [%s] does not match the number " + "of columns [%d]. Try names by a Spark SQL DDL-formatted " + "string." % (len(sdf.schema), len(names)) + ) + column_labels = dict(zip(names, sdf.columns)) + elif header is None: + column_labels = dict(enumerate(sdf.columns)) + else: + column_labels = {col: col for col in sdf.columns} + + if usecols is not None: + missing: List[Union[int, str]] + if callable(usecols): + column_labels = { + label: col for label, col in column_labels.items() if usecols(label) + } + missing = [] + elif all(isinstance(col, int) for col in usecols): + usecols_ints = cast(List[int], usecols) + new_column_labels = { + label: col + for i, (label, col) in enumerate(column_labels.items()) + if i in usecols_ints + } + missing = [ + col + for col in usecols_ints + if ( + col >= len(column_labels) + or list(column_labels)[col] not in new_column_labels + ) + ] + column_labels = new_column_labels + elif all(isinstance(col, str) for col in usecols): + new_column_labels = { + label: col for label, col in column_labels.items() if label in usecols + } + missing = [col for col in usecols if col not in new_column_labels] + column_labels = new_column_labels + else: + raise ValueError( + "'usecols' must either be list-like of all strings, " + "all unicode, all integers or a callable." + ) + if len(missing) > 0: + raise ValueError( + "Usecols do not match columns, columns expected but not " "found: %s" % missing + ) + + if len(column_labels) > 0: + sdf = sdf.select([scol_for(sdf, col) for col in column_labels.values()]) + else: + sdf = default_session().createDataFrame([], schema=StructType()) + else: + sdf = default_session().createDataFrame([], schema=StructType()) + column_labels = {} + + if nrows is not None: + sdf = sdf.limit(nrows) + + index_spark_column_names: List[str] + index_names: List[Label] + if index_col is not None: + if isinstance(index_col, (str, int)): + index_col = [index_col] + for col in index_col: + if col not in column_labels: + raise KeyError(col) + index_spark_column_names = [column_labels[col] for col in index_col] + index_names = [(col,) for col in index_col] + column_labels = { + label: col for label, col in column_labels.items() if label not in index_col + } + else: + log_advice( + "If `index_col` is not specified for `read_csv`, " + "the default index is attached which can cause additional overhead." + ) + index_spark_column_names = [] + index_names = [] + + psdf: DataFrame = DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_spark_column_names], + index_names=index_names, + column_labels=[ + label if is_name_like_tuple(label) else (label,) for label in column_labels + ], + data_spark_columns=[scol_for(sdf, col) for col in column_labels.values()], + ) + ) + + if dtype is not None: + if isinstance(dtype, dict): + for col, tpe in dtype.items(): + psdf[col] = psdf[col].astype(tpe) + else: + for col in psdf.columns: + psdf[col] = psdf[col].astype(dtype) + + if squeeze and len(psdf.columns) == 1: + return first_series(psdf) + else: + return psdf + + +def read_json( + path: str, lines: bool = True, index_col: Optional[Union[str, List[str]]] = None, **options: Any +) -> DataFrame: + """ + Convert a JSON string to DataFrame. + + Parameters + ---------- + path : string + File path + lines : bool, default True + Read the file as a JSON object per line. It should be always True for now. + index_col : str or list of str, optional, default: None + Index column of table in Spark. + options : dict + All other options passed directly into Spark's data source. + + Examples + -------- + >>> df = ps.DataFrame([['a', 'b'], ['c', 'd']], + ... columns=['col 1', 'col 2']) + + >>> df.to_json(path=r'%s/read_json/foo.json' % path, num_files=1) + >>> ps.read_json( + ... path=r'%s/read_json/foo.json' % path + ... ).sort_values(by="col 1") + col 1 col 2 + 0 a b + 1 c d + + >>> df.to_json(path=r'%s/read_json/foo.json' % path, num_files=1, lineSep='___') + >>> ps.read_json( + ... path=r'%s/read_json/foo.json' % path, lineSep='___' + ... ).sort_values(by="col 1") + col 1 col 2 + 0 a b + 1 c d + + You can preserve the index in the roundtrip as below. + + >>> df.to_json(path=r'%s/read_json/bar.json' % path, num_files=1, index_col="index") + >>> ps.read_json( + ... path=r'%s/read_json/bar.json' % path, index_col="index" + ... ).sort_values(by="col 1") # doctest: +NORMALIZE_WHITESPACE + col 1 col 2 + index + 0 a b + 1 c d + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `read_json`, " + "the default index is attached which can cause additional overhead." + ) + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + if not lines: + raise NotImplementedError("lines=False is not implemented yet.") + + return read_spark_io(path, format="json", index_col=index_col, **options) + + +def read_delta( + path: str, + version: Optional[str] = None, + timestamp: Optional[str] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, +) -> DataFrame: + """ + Read a Delta Lake table on some file system and return a DataFrame. + + If the Delta Lake table is already stored in the catalog (aka the metastore), use 'read_table'. + + Parameters + ---------- + path : string + Path to the Delta Lake table. + version : string, optional + Specifies the table version (based on Delta's internal transaction version) to read from, + using Delta's time travel feature. This sets Delta's 'versionAsOf' option. Note that + this parameter and `timestamp` parameter cannot be used together, otherwise it will raise a + `ValueError`. + timestamp : string, optional + Specifies the table version (based on timestamp) to read from, + using Delta's time travel feature. This must be a valid date or timestamp string in Spark, + and sets Delta's 'timestampAsOf' option. Note that this parameter and `version` parameter + cannot be used together, otherwise it will raise a `ValueError`. + index_col : str or list of str, optional, default: None + Index column of table in Spark. + options + Additional options that can be passed onto Delta. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.to_delta + read_table + read_spark_io + read_parquet + + Examples + -------- + >>> ps.range(1).to_delta('%s/read_delta/foo' % path) # doctest: +SKIP + >>> ps.read_delta('%s/read_delta/foo' % path) # doctest: +SKIP + id + 0 0 + + >>> ps.range(10, 15, num_partitions=1).to_delta('%s/read_delta/foo' % path, + ... mode='overwrite') # doctest: +SKIP + >>> ps.read_delta('%s/read_delta/foo' % path) # doctest: +SKIP + id + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + + >>> ps.read_delta('%s/read_delta/foo' % path, version=0) # doctest: +SKIP + id + 0 0 + + You can preserve the index in the roundtrip as below. + + >>> ps.range(10, 15, num_partitions=1).to_delta( + ... '%s/read_delta/bar' % path, index_col="index") # doctest: +SKIP + >>> ps.read_delta('%s/read_delta/bar' % path, index_col="index") # doctest: +SKIP + id + index + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `read_delta`, " + "the default index is attached which can cause additional overhead." + ) + if version is not None and timestamp is not None: + raise ValueError("version and timestamp cannot be used together.") + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + if version is not None: + options["versionAsOf"] = version + if timestamp is not None: + options["timestampAsOf"] = timestamp + return read_spark_io(path, format="delta", index_col=index_col, **options) + + +def read_table(name: str, index_col: Optional[Union[str, List[str]]] = None) -> DataFrame: + """ + Read a Spark table and return a DataFrame. + + Parameters + ---------- + name : string + Table name in Spark. + + index_col : str or list of str, optional, default: None + Index column of table in Spark. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.to_table + read_delta + read_parquet + read_spark_io + + Examples + -------- + >>> ps.range(1).to_table('%s.my_table' % db) + >>> ps.read_table('%s.my_table' % db) + id + 0 0 + + >>> ps.range(1).to_table('%s.my_table' % db, index_col="index") + >>> ps.read_table('%s.my_table' % db, index_col="index") # doctest: +NORMALIZE_WHITESPACE + id + index + 0 0 + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `read_table`, " + "the default index is attached which can cause additional overhead." + ) + sdf = default_session().read.table(name) + index_spark_columns, index_names = _get_index_map(sdf, index_col) + + return DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + + +def read_spark_io( + path: Optional[str] = None, + format: Optional[str] = None, + schema: Union[str, "StructType"] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, +) -> DataFrame: + """Load a DataFrame from a Spark data source. + + Parameters + ---------- + path : string, optional + Path to the data source. + format : string, optional + Specifies the output data source format. Some common ones are: + + - 'delta' + - 'parquet' + - 'orc' + - 'json' + - 'csv' + schema : string or StructType, optional + Input schema. If none, Spark tries to infer the schema automatically. + The schema can either be a Spark StructType, or a DDL-formatted string like + `col0 INT, col1 DOUBLE`. + index_col : str or list of str, optional, default: None + Index column of table in Spark. + options : dict + All other options passed directly into Spark's data source. + + See Also + -------- + DataFrame.to_spark_io + DataFrame.read_table + DataFrame.read_delta + DataFrame.read_parquet + + Examples + -------- + >>> ps.range(1).to_spark_io('%s/read_spark_io/data.parquet' % path) + >>> ps.read_spark_io( + ... '%s/read_spark_io/data.parquet' % path, format='parquet', schema='id long') + id + 0 0 + + >>> ps.range(10, 15, num_partitions=1).to_spark_io('%s/read_spark_io/data.json' % path, + ... format='json', lineSep='__') + >>> ps.read_spark_io( + ... '%s/read_spark_io/data.json' % path, format='json', schema='id long', lineSep='__') + id + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + + You can preserve the index in the roundtrip as below. + + >>> ps.range(10, 15, num_partitions=1).to_spark_io('%s/read_spark_io/data.orc' % path, + ... format='orc', index_col="index") + >>> ps.read_spark_io( + ... path=r'%s/read_spark_io/data.orc' % path, format="orc", index_col="index") + ... # doctest: +NORMALIZE_WHITESPACE + id + index + 0 10 + 1 11 + 2 12 + 3 13 + 4 14 + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + sdf = default_session().read.load(path=path, format=format, schema=schema, **options) + index_spark_columns, index_names = _get_index_map(sdf, index_col) + + return DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + + +def read_parquet( + path: str, + columns: Optional[List[str]] = None, + index_col: Optional[List[str]] = None, + pandas_metadata: bool = False, + **options: Any, +) -> DataFrame: + """Load a parquet object from the file path, returning a DataFrame. + + Parameters + ---------- + path : string + File path + columns : list, default=None + If not None, only these columns will be read from the file. + index_col : str or list of str, optional, default: None + Index column of table in Spark. + pandas_metadata : bool, default: False + If True, try to respect the metadata if the Parquet file is written from pandas. + options : dict + All other options passed directly into Spark's data source. + + Returns + ------- + DataFrame + + See Also + -------- + DataFrame.to_parquet + DataFrame.read_table + DataFrame.read_delta + DataFrame.read_spark_io + + Examples + -------- + >>> ps.range(1).to_parquet('%s/read_spark_io/data.parquet' % path) + >>> ps.read_parquet('%s/read_spark_io/data.parquet' % path, columns=['id']) + id + 0 0 + + You can preserve the index in the roundtrip as below. + + >>> ps.range(1).to_parquet('%s/read_spark_io/data.parquet' % path, index_col="index") + >>> ps.read_parquet('%s/read_spark_io/data.parquet' % path, columns=['id'], index_col="index") + ... # doctest: +NORMALIZE_WHITESPACE + id + index + 0 0 + """ + if index_col is None: + log_advice( + "If `index_col` is not specified for `read_parquet`, " + "the default index is attached which can cause additional overhead." + ) + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + if columns is not None: + columns = list(columns) + + index_names = None + + if index_col is None and pandas_metadata: + # Try to read pandas metadata + + @pandas_udf( # type: ignore[call-overload] + "index_col array, index_names array" + ) + def read_index_metadata(pser: pd.Series) -> pd.DataFrame: + binary = pser.iloc[0] + metadata = pq.ParquetFile(pa.BufferReader(binary)).metadata.metadata + if b"pandas" in metadata: + pandas_metadata = json.loads(metadata[b"pandas"].decode("utf8")) + if all(isinstance(col, str) for col in pandas_metadata["index_columns"]): + index_col = [] + index_names = [] + for col in pandas_metadata["index_columns"]: + index_col.append(col) + for column in pandas_metadata["columns"]: + if column["field_name"] == col: + index_names.append(column["name"]) + break + else: + index_names.append(None) + return pd.DataFrame({"index_col": [index_col], "index_names": [index_names]}) + return pd.DataFrame({"index_col": [None], "index_names": [None]}) + + index_col, index_names = ( + default_session() + .read.format("binaryFile") + .load(path) + .limit(1) + .select(read_index_metadata("content").alias("index_metadata")) + .select("index_metadata.*") + .head() + ) + + psdf = read_spark_io(path=path, format="parquet", options=options, index_col=index_col) + + if columns is not None: + new_columns = [c for c in columns if c in psdf.columns] + if len(new_columns) > 0: + psdf = psdf[new_columns] + else: + sdf = default_session().createDataFrame([], schema=StructType()) + index_spark_columns, index_names = _get_index_map(sdf, index_col) + psdf = DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=index_spark_columns, + index_names=index_names, + ) + ) + + if index_names is not None: + psdf.index.names = index_names + + return psdf + + +def read_clipboard(sep: str = r"\s+", **kwargs: Any) -> DataFrame: + r""" + Read text from clipboard and pass to read_csv. See read_csv for the + full argument list + + Parameters + ---------- + sep : str, default '\s+' + A string or regex delimiter. The default of '\s+' denotes + one or more whitespace characters. + + See Also + -------- + DataFrame.to_clipboard : Write text out to clipboard. + + Returns + ------- + parsed : DataFrame + """ + return cast(DataFrame, from_pandas(pd.read_clipboard(sep, **kwargs))) + + +def read_excel( + io: Union[str, Any], + sheet_name: Union[str, int, List[Union[str, int]], None] = 0, + header: Union[int, List[int]] = 0, + names: Optional[List] = None, + index_col: Optional[List[int]] = None, + usecols: Optional[Union[int, str, List[Union[int, str]], Callable[[str], bool]]] = None, + squeeze: bool = False, + dtype: Optional[Dict[str, Union[str, Dtype]]] = None, + engine: Optional[str] = None, + converters: Optional[Dict] = None, + true_values: Optional[Any] = None, + false_values: Optional[Any] = None, + skiprows: Optional[Union[int, List[int]]] = None, + nrows: Optional[int] = None, + na_values: Optional[Any] = None, + keep_default_na: bool = True, + verbose: bool = False, + parse_dates: Union[bool, List, Dict] = False, + date_parser: Optional[Callable] = None, + thousands: Optional[str] = None, + comment: Optional[str] = None, + skipfooter: int = 0, + convert_float: bool = True, + mangle_dupe_cols: bool = True, + **kwds: Any, +) -> Union[DataFrame, Series, Dict[str, Union[DataFrame, Series]]]: + """ + Read an Excel file into a pandas-on-Spark DataFrame or Series. + + Support both `xls` and `xlsx` file extensions from a local filesystem or URL. + Support an option to read a single sheet or a list of sheets. + + Parameters + ---------- + io : str, file descriptor, pathlib.Path, ExcelFile or xlrd.Book + The string could be a URL. The value URL must be available in Spark's DataFrameReader. + + .. note:: + If the underlying Spark is below 3.0, the parameter as a string is not supported. + You can use `ps.from_pandas(pd.read_excel(...))` as a workaround. + + sheet_name : str, int, list, or None, default 0 + Strings are used for sheet names. Integers are used in zero-indexed + sheet positions. Lists of strings/integers are used to request + multiple sheets. Specify None to get all sheets. + + Available cases: + + * Defaults to ``0``: 1st sheet as a `DataFrame` + * ``1``: 2nd sheet as a `DataFrame` + * ``"Sheet1"``: Load sheet with name "Sheet1" + * ``[0, 1, "Sheet5"]``: Load first, second and sheet named "Sheet5" + as a dict of `DataFrame` + * None: All sheets. + + header : int, list of int, default 0 + Row (0-indexed) to use for the column labels of the parsed + DataFrame. If a list of integers is passed those row positions will + be combined into a ``MultiIndex``. Use None if there is no header. + names : array-like, default None + List of column names to use. If file contains no header row, + then you should explicitly pass header=None. + index_col : int, list of int, default None + Column (0-indexed) to use as the row labels of the DataFrame. + Pass None if there is no such column. If a list is passed, + those columns will be combined into a ``MultiIndex``. If a + subset of data is selected with ``usecols``, index_col + is based on the subset. + usecols : int, str, list-like, or callable default None + Return a subset of the columns. + + * If None, then parse all columns. + * If str, then indicates comma separated list of Excel column letters + and column ranges (e.g. "A:E" or "A,C,E:F"). Ranges are inclusive of + both sides. + * If list of int, then indicates list of column numbers to be parsed. + * If list of string, then indicates list of column names to be parsed. + * If callable, then evaluate each column name against it and parse the + column if the callable returns ``True``. + squeeze : bool, default False + If the parsed data only contains one column then return a Series. + + .. deprecated:: 3.4.0 + + dtype : Type name or dict of column -> type, default None + Data type for data or columns. E.g. {'a': np.float64, 'b': np.int32} + Use `object` to preserve data as stored in Excel and not interpret dtype. + If converters are specified, they will be applied INSTEAD + of dtype conversion. + engine : str, default None + If io is not a buffer or path, this must be set to identify io. + Acceptable values are None or xlrd. + converters : dict, default None + Dict of functions for converting values in certain columns. Keys can + either be integers or column labels, values are functions that take one + input argument, the Excel cell content, and return the transformed + content. + true_values : list, default None + Values to consider as True. + false_values : list, default None + Values to consider as False. + skiprows : list-like + Rows to skip at the beginning (0-indexed). + nrows : int, default None + Number of rows to parse. + na_values : scalar, str, list-like, or dict, default None + Additional strings to recognize as NA/NaN. If dict passed, specific + per-column NA values. By default the following values are interpreted + as NaN. + keep_default_na : bool, default True + If na_values are specified and keep_default_na is False the default NaN + values are overridden, otherwise they're appended to. + verbose : bool, default False + Indicate number of NA values placed in non-numeric columns. + parse_dates : bool, list-like, or dict, default False + The behavior is as follows: + + * bool. If True -> try parsing the index. + * list of int or names. e.g. If [1, 2, 3] -> try parsing columns 1, 2, 3 + each as a separate date column. + * list of lists. e.g. If [[1, 3]] -> combine columns 1 and 3 and parse as + a single date column. + * dict, e.g. {{'foo' : [1, 3]}} -> parse columns 1, 3 as date and call + result 'foo' + + If a column or index contains an unparseable date, the entire column or + index will be returned unaltered as an object data type. For non-standard + datetime parsing, use ``pd.to_datetime`` after ``pd.read_csv`` + + Note: A fast-path exists for iso8601-formatted dates. + date_parser : function, optional + Function to use for converting a sequence of string columns to an array of + datetime instances. The default uses ``dateutil.parser.parser`` to do the + conversion. pandas-on-Spark will try to call `date_parser` in three different ways, + advancing to the next if an exception occurs: 1) Pass one or more arrays + (as defined by `parse_dates`) as arguments; 2) concatenate (row-wise) the + string values from the columns defined by `parse_dates` into a single array + and pass that; and 3) call `date_parser` once for each row using one or + more strings (corresponding to the columns defined by `parse_dates`) as + arguments. + thousands : str, default None + Thousands separator for parsing string columns to numeric. Note that + this parameter is only necessary for columns stored as TEXT in Excel, + any numeric columns will automatically be parsed, regardless of display + format. + comment : str, default None + Comments out remainder of line. Pass a character or characters to this + argument to indicate comments in the input file. Any data between the + comment string and the end of the current line is ignored. + skipfooter : int, default 0 + Rows at the end to skip (0-indexed). + convert_float : bool, default True + Convert integral floats to int (i.e., 1.0 --> 1). If False, all numeric + data will be read in as floats: Excel stores all numbers as floats + internally. + + .. deprecated:: 3.4.0 + + mangle_dupe_cols : bool, default True + Duplicate columns will be specified as 'X', 'X.1', ...'X.N', rather than + 'X'...'X'. Passing in False will cause data to be overwritten if there + are duplicate names in the columns. + + .. deprecated:: 3.4.0 + + **kwds : optional + Optional keyword arguments can be passed to ``TextFileReader``. + + Returns + ------- + DataFrame or dict of DataFrames + DataFrame from the passed in Excel file. See notes in sheet_name + argument for more information on when a dict of DataFrames is returned. + + See Also + -------- + DataFrame.to_excel : Write DataFrame to an Excel file. + DataFrame.to_csv : Write DataFrame to a comma-separated values (csv) file. + read_csv : Read a comma-separated values (csv) file into DataFrame. + + Examples + -------- + The file can be read using the file name as string or an open file object: + + >>> ps.read_excel('tmp.xlsx', index_col=0) # doctest: +SKIP + Name Value + 0 string1 1 + 1 string2 2 + 2 #Comment 3 + + >>> ps.read_excel(open('tmp.xlsx', 'rb'), + ... sheet_name='Sheet3') # doctest: +SKIP + Unnamed: 0 Name Value + 0 0 string1 1 + 1 1 string2 2 + 2 2 #Comment 3 + + Index and header can be specified via the `index_col` and `header` arguments + + >>> ps.read_excel('tmp.xlsx', index_col=None, header=None) # doctest: +SKIP + 0 1 2 + 0 NaN Name Value + 1 0.0 string1 1 + 2 1.0 string2 2 + 3 2.0 #Comment 3 + + Column types are inferred but can be explicitly specified + + >>> ps.read_excel('tmp.xlsx', index_col=0, + ... dtype={'Name': str, 'Value': float}) # doctest: +SKIP + Name Value + 0 string1 1.0 + 1 string2 2.0 + 2 #Comment 3.0 + + True, False, and NA values, and thousands separators have defaults, + but can be explicitly specified, too. Supply the values you would like + as strings or lists of strings! + + >>> ps.read_excel('tmp.xlsx', index_col=0, + ... na_values=['string1', 'string2']) # doctest: +SKIP + Name Value + 0 None 1 + 1 None 2 + 2 #Comment 3 + + Comment lines in the excel input file can be skipped using the `comment` kwarg + + >>> ps.read_excel('tmp.xlsx', index_col=0, comment='#') # doctest: +SKIP + Name Value + 0 string1 1.0 + 1 string2 2.0 + 2 None NaN + """ + + def pd_read_excel( + io_or_bin: Any, sn: Union[str, int, List[Union[str, int]], None], sq: bool + ) -> pd.DataFrame: + return pd.read_excel( + io=BytesIO(io_or_bin) if isinstance(io_or_bin, (bytes, bytearray)) else io_or_bin, + sheet_name=sn, + header=header, + names=names, + index_col=index_col, + usecols=usecols, + squeeze=sq, + dtype=dtype, + engine=engine, + converters=converters, + true_values=true_values, + false_values=false_values, + skiprows=skiprows, + nrows=nrows, + na_values=na_values, + keep_default_na=keep_default_na, + verbose=verbose, + parse_dates=parse_dates, # type: ignore[arg-type] + date_parser=date_parser, + thousands=thousands, + comment=comment, + skipfooter=skipfooter, + convert_float=convert_float, + mangle_dupe_cols=mangle_dupe_cols, + **kwds, + ) + + if isinstance(io, str): + # 'binaryFile' format is available since Spark 3.0.0. + binaries = default_session().read.format("binaryFile").load(io).select("content").head(2) + io_or_bin = binaries[0][0] + single_file = len(binaries) == 1 + else: + io_or_bin = io + single_file = True + + pdf_or_psers = pd_read_excel(io_or_bin, sn=sheet_name, sq=squeeze) + + if single_file: + if isinstance(pdf_or_psers, dict): + return { + sn: cast(Union[DataFrame, Series], from_pandas(pdf_or_pser)) + for sn, pdf_or_pser in pdf_or_psers.items() + } + else: + return cast(Union[DataFrame, Series], from_pandas(pdf_or_psers)) + else: + + def read_excel_on_spark( + pdf_or_pser: Union[pd.DataFrame, pd.Series], + sn: Union[str, int, List[Union[str, int]], None], + ) -> Union[DataFrame, Series]: + if isinstance(pdf_or_pser, pd.Series): + pdf = pdf_or_pser.to_frame() + else: + pdf = pdf_or_pser + + psdf = cast(DataFrame, from_pandas(pdf)) + return_schema = force_decimal_precision_scale( + as_nullable_spark_type(psdf._internal.spark_frame.drop(*HIDDEN_COLUMNS).schema) + ) + + def output_func(pdf: pd.DataFrame) -> pd.DataFrame: + pdf = pd.concat( + [pd_read_excel(bin, sn=sn, sq=False) for bin in pdf[pdf.columns[0]]] + ) + + reset_index = pdf.reset_index() + for name, col in reset_index.items(): + dt = col.dtype + if is_datetime64_dtype(dt) or is_datetime64tz_dtype(dt): + continue + reset_index[name] = col.replace({np.nan: None}) + pdf = reset_index + + # Just positionally map the column names to given schema's. + return pdf.rename(columns=dict(zip(pdf.columns, return_schema.names))) + + sdf = ( + default_session() + .read.format("binaryFile") + .load(io) + .select("content") + .mapInPandas(lambda iterator: map(output_func, iterator), schema=return_schema) + ) + + psdf = DataFrame(psdf._internal.with_new_sdf(sdf)) + if squeeze and len(psdf.columns) == 1: + return first_series(psdf) + else: + return psdf + + if isinstance(pdf_or_psers, dict): + return { + sn: read_excel_on_spark(pdf_or_pser, sn) for sn, pdf_or_pser in pdf_or_psers.items() + } + else: + return read_excel_on_spark(pdf_or_psers, sheet_name) + + +def read_html( + io: Union[str, Any], + match: str = ".+", + flavor: Optional[str] = None, + header: Optional[Union[int, List[int]]] = None, + index_col: Optional[Union[int, List[int]]] = None, + skiprows: Optional[Union[int, List[int], slice]] = None, + attrs: Optional[Dict[str, str]] = None, + parse_dates: bool = False, + thousands: str = ",", + encoding: Optional[str] = None, + decimal: str = ".", + converters: Optional[Dict] = None, + na_values: Optional[Any] = None, + keep_default_na: bool = True, + displayed_only: bool = True, +) -> List[DataFrame]: + r"""Read HTML tables into a ``list`` of ``DataFrame`` objects. + + Parameters + ---------- + io : str or file-like + A URL, a file-like object, or a raw string containing HTML. Note that + lxml only accepts the http, FTP and file URL protocols. If you have a + URL that starts with ``'https'`` you might try removing the ``'s'``. + + match : str or compiled regular expression, optional + The set of tables containing text matching this regex or string will be + returned. Unless the HTML is extremely simple you will probably need to + pass a non-empty string here. Defaults to '.+' (match any non-empty + string). The default value will return all tables contained on a page. + This value is converted to a regular expression so that there is + consistent behavior between Beautiful Soup and lxml. + + flavor : str or None, container of strings + The parsing engine to use. 'bs4' and 'html5lib' are synonymous with + each other, they are both there for backwards compatibility. The + default of ``None`` tries to use ``lxml`` to parse and if that fails it + falls back on ``bs4`` + ``html5lib``. + + header : int or list-like or None, optional + The row (or list of rows for a :class:`~ps.MultiIndex`) to use to + make the columns headers. + + index_col : int or list-like or None, optional + The column (or list of columns) to use to create the index. + + skiprows : int or list-like or slice or None, optional + 0-based. Number of rows to skip after parsing the column integer. If a + sequence of integers or a slice is given, will skip the rows indexed by + that sequence. Note that a single element sequence means 'skip the nth + row' whereas an integer means 'skip n rows'. + + attrs : dict or None, optional + This is a dictionary of attributes that you can pass to use to identify + the table in the HTML. These are not checked for validity before being + passed to lxml or Beautiful Soup. However, these attributes must be + valid HTML table attributes to work correctly. For example, :: + + attrs = {'id': 'table'} + + is a valid attribute dictionary because the 'id' HTML tag attribute is + a valid HTML attribute for *any* HTML tag as per `this document + `__. :: + + attrs = {'asdf': 'table'} + + is *not* a valid attribute dictionary because 'asdf' is not a valid + HTML attribute even if it is a valid XML attribute. Valid HTML 4.01 + table attributes can be found `here + `__. A + working draft of the HTML 5 spec can be found `here + `__. It contains the + latest information on table attributes for the modern web. + + parse_dates : bool, optional + See :func:`~ps.read_csv` for more details. + + thousands : str, optional + Separator to use to parse thousands. Defaults to ``','``. + + encoding : str or None, optional + The encoding used to decode the web page. Defaults to ``None``.``None`` + preserves the previous encoding behavior, which depends on the + underlying parser library (e.g., the parser library will try to use + the encoding provided by the document). + + decimal : str, default '.' + Character to recognize as decimal point (example: use ',' for European + data). + + converters : dict, default None + Dict of functions for converting values in certain columns. Keys can + either be integers or column labels, values are functions that take one + input argument, the cell (not column) content, and return the + transformed content. + + na_values : iterable, default None + Custom NA values + + keep_default_na : bool, default True + If na_values are specified and keep_default_na is False the default NaN + values are overridden, otherwise they're appended to + + displayed_only : bool, default True + Whether elements with "display: none" should be parsed + + Returns + ------- + dfs : list of DataFrames + + See Also + -------- + read_csv + DataFrame.to_html + """ + pdfs = pd.read_html( + io=io, + match=match, + flavor=flavor, + header=header, + index_col=index_col, + skiprows=skiprows, + attrs=attrs, + parse_dates=parse_dates, + thousands=thousands, + encoding=encoding, + decimal=decimal, + converters=converters, + na_values=na_values, + keep_default_na=keep_default_na, + displayed_only=displayed_only, + ) + return cast(List[DataFrame], [from_pandas(pdf) for pdf in pdfs]) + + +# TODO: add `coerce_float` and 'parse_dates' parameters +def read_sql_table( + table_name: str, + con: str, + schema: Optional[str] = None, + index_col: Optional[Union[str, List[str]]] = None, + columns: Optional[Union[str, List[str]]] = None, + **options: Any, +) -> DataFrame: + """ + Read SQL database table into a DataFrame. + + Given a table name and a JDBC URI, returns a DataFrame. + + Parameters + ---------- + table_name : str + Name of SQL table in database. + con : str + A JDBC URI could be provided as str. + + .. note:: The URI must be JDBC URI instead of Python's database URI. + + schema : str, default None + Name of SQL schema in database to query (if database flavor + supports this). Uses default schema if None (default). + index_col : str or list of str, optional, default: None + Column(s) to set as index(MultiIndex). + columns : list, default None + List of column names to select from SQL table. + options : dict + All other options passed directly into Spark's JDBC data source. + + Returns + ------- + DataFrame + A SQL table is returned as two-dimensional data structure with labeled + axes. + + See Also + -------- + read_sql_query : Read SQL query into a DataFrame. + read_sql : Read SQL query or database table into a DataFrame. + + Examples + -------- + >>> ps.read_sql_table('table_name', 'jdbc:postgresql:db_name') # doctest: +SKIP + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + reader = default_session().read + reader.option("dbtable", table_name) + reader.option("url", con) + if schema is not None: + reader.schema(schema) + reader.options(**options) + sdf = reader.format("jdbc").load() + index_spark_columns, index_names = _get_index_map(sdf, index_col) + psdf: DataFrame = DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + if columns is not None: + if isinstance(columns, str): + columns = [columns] + psdf = psdf[columns] + return psdf + + +# TODO: add `coerce_float`, `params`, and 'parse_dates' parameters +def read_sql_query( + sql: str, con: str, index_col: Optional[Union[str, List[str]]] = None, **options: Any +) -> DataFrame: + """Read SQL query into a DataFrame. + + Returns a DataFrame corresponding to the result set of the query + string. Optionally provide an `index_col` parameter to use one of the + columns as the index, otherwise default index will be used. + + .. note:: Some database might hit the issue of Spark: SPARK-27596 + + Parameters + ---------- + sql : string SQL query + SQL query to be executed. + con : str + A JDBC URI could be provided as str. + + .. note:: The URI must be JDBC URI instead of Python's database URI. + + index_col : string or list of strings, optional, default: None + Column(s) to set as index(MultiIndex). + options : dict + All other options passed directly into Spark's JDBC data source. + + Returns + ------- + DataFrame + + See Also + -------- + read_sql_table : Read SQL database table into a DataFrame. + read_sql + + Examples + -------- + >>> ps.read_sql_query('SELECT * FROM table_name', 'jdbc:postgresql:db_name') # doctest: +SKIP + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + reader = default_session().read + reader.option("query", sql) + reader.option("url", con) + reader.options(**options) + sdf = reader.format("jdbc").load() + index_spark_columns, index_names = _get_index_map(sdf, index_col) + return DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + + +# TODO: add `coerce_float`, `params`, and 'parse_dates' parameters +def read_sql( + sql: str, + con: str, + index_col: Optional[Union[str, List[str]]] = None, + columns: Optional[Union[str, List[str]]] = None, + **options: Any, +) -> DataFrame: + """ + Read SQL query or database table into a DataFrame. + + This function is a convenience wrapper around ``read_sql_table`` and + ``read_sql_query`` (for backward compatibility). It will delegate + to the specific function depending on the provided input. A SQL query + will be routed to ``read_sql_query``, while a database table name will + be routed to ``read_sql_table``. Note that the delegated function might + have more specific notes about their functionality not listed here. + + .. note:: Some database might hit the issue of Spark: SPARK-27596 + + Parameters + ---------- + sql : string + SQL query to be executed or a table name. + con : str + A JDBC URI could be provided as str. + + .. note:: The URI must be JDBC URI instead of Python's database URI. + + index_col : string or list of strings, optional, default: None + Column(s) to set as index(MultiIndex). + columns : list, default: None + List of column names to select from SQL table (only used when reading + a table). + options : dict + All other options passed directly into Spark's JDBC data source. + + Returns + ------- + DataFrame + + See Also + -------- + read_sql_table : Read SQL database table into a DataFrame. + read_sql_query : Read SQL query into a DataFrame. + + Examples + -------- + >>> ps.read_sql('table_name', 'jdbc:postgresql:db_name') # doctest: +SKIP + >>> ps.read_sql('SELECT * FROM table_name', 'jdbc:postgresql:db_name') # doctest: +SKIP + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + striped = sql.strip() + if " " not in striped: # TODO: identify the table name or not more precisely. + return read_sql_table(sql, con, index_col=index_col, columns=columns, **options) + else: + return read_sql_query(sql, con, index_col=index_col, **options) + + +@no_type_check +def to_datetime( + arg, + errors: str = "raise", + format: Optional[str] = None, + unit: Optional[str] = None, + infer_datetime_format: bool = False, + origin: str = "unix", +): + """ + Convert argument to datetime. + + Parameters + ---------- + arg : integer, float, string, datetime, list, tuple, 1-d array, Series + or DataFrame/dict-like + + errors : {'ignore', 'raise', 'coerce'}, default 'raise' + + - If 'raise', then invalid parsing will raise an exception + - If 'coerce', then invalid parsing will be set as NaT + - If 'ignore', then invalid parsing will return the input + format : string, default None + strftime to parse time, eg "%d/%m/%Y", note that "%f" will parse + all the way up to nanoseconds. + unit : string, default None + unit of the arg (D,s,ms,us,ns) denote the unit, which is an + integer or float number. This will be based off the origin. + Example, with unit='ms' and origin='unix' (the default), this + would calculate the number of milliseconds to the unix epoch start. + infer_datetime_format : boolean, default False + If True and no `format` is given, attempt to infer the format of the + datetime strings, and if it can be inferred, switch to a faster + method of parsing them. In some cases this can increase the parsing + speed by ~5-10x. + origin : scalar, default 'unix' + Define the reference date. The numeric values would be parsed as number + of units (defined by `unit`) since this reference date. + + - If 'unix' (or POSIX) time; origin is set to 1970-01-01. + - If 'julian', unit must be 'D', and origin is set to beginning of + Julian Calendar. Julian day number 0 is assigned to the day starting + at noon on January 1, 4713 BC. + - If Timestamp convertible, origin is set to Timestamp identified by + origin. + + Returns + ------- + ret : datetime if parsing succeeded. + Return type depends on input: + + - list-like: DatetimeIndex + - Series: Series of datetime64 dtype + - scalar: Timestamp + + In case when it is not possible to return designated types (e.g. when + any element of input is before Timestamp.min or after Timestamp.max) + return will have datetime.datetime type (or corresponding + array/Series). + + Examples + -------- + Assembling a datetime from multiple columns of a DataFrame. The keys can be + common abbreviations like ['year', 'month', 'day', 'minute', 'second', + 'ms', 'us', 'ns']) or plurals of the same + + >>> df = ps.DataFrame({'year': [2015, 2016], + ... 'month': [2, 3], + ... 'day': [4, 5]}) + >>> ps.to_datetime(df) + 0 2015-02-04 + 1 2016-03-05 + dtype: datetime64[ns] + + If a date does not meet the `timestamp limitations + `_, passing errors='ignore' + will return the original input instead of raising any exception. + + Passing errors='coerce' will force an out-of-bounds date to NaT, + in addition to forcing non-dates (or non-parseable dates) to NaT. + + >>> ps.to_datetime('13000101', format='%Y%m%d', errors='ignore') # doctest: +SKIP + datetime.datetime(1300, 1, 1, 0, 0) + >>> ps.to_datetime('13000101', format='%Y%m%d', errors='coerce') + NaT + + Passing infer_datetime_format=True can often-times speedup a parsing + if its not an ISO8601 format exactly, but in a regular format. + + >>> s = ps.Series(['3/11/2000', '3/12/2000', '3/13/2000'] * 1000) + >>> s.head() + 0 3/11/2000 + 1 3/12/2000 + 2 3/13/2000 + 3 3/11/2000 + 4 3/12/2000 + dtype: object + + >>> import timeit + >>> timeit.timeit( + ... lambda: repr(ps.to_datetime(s, infer_datetime_format=True)), + ... number = 1) # doctest: +SKIP + 0.35832712500000063 + + >>> timeit.timeit( + ... lambda: repr(ps.to_datetime(s, infer_datetime_format=False)), + ... number = 1) # doctest: +SKIP + 0.8895321660000004 + + Using a unix epoch time + + >>> ps.to_datetime(1490195805, unit='s') + Timestamp('2017-03-22 15:16:45') + >>> ps.to_datetime(1490195805433502912, unit='ns') + Timestamp('2017-03-22 15:16:45.433502912') + + Using a non-unix epoch origin + + >>> ps.to_datetime([1, 2, 3], unit='D', origin=pd.Timestamp('1960-01-01')) + DatetimeIndex(['1960-01-02', '1960-01-03', '1960-01-04'], dtype='datetime64[ns]', freq=None) + """ + + # mappings for assembling units + # From pandas: pandas.core.tools.datetimes + _unit_map = { + "year": "year", + "years": "year", + "month": "month", + "months": "month", + "day": "day", + "days": "day", + "hour": "h", + "hours": "h", + "minute": "m", + "minutes": "m", + "second": "s", + "seconds": "s", + "ms": "ms", + "millisecond": "ms", + "milliseconds": "ms", + "us": "us", + "microsecond": "us", + "microseconds": "us", + } + + def pandas_to_datetime( + pser_or_pdf: Union[pd.DataFrame, pd.Series], cols: Optional[List[str]] = None + ) -> Series[np.datetime64]: + if isinstance(pser_or_pdf, pd.DataFrame): + pser_or_pdf = pser_or_pdf[cols] + return pd.to_datetime( + pser_or_pdf, + errors=errors, + format=format, + unit=unit, + infer_datetime_format=infer_datetime_format, + origin=origin, + ) + + if isinstance(arg, Series): + return arg.pandas_on_spark.transform_batch(pandas_to_datetime) + if isinstance(arg, DataFrame): + unit = {k: _unit_map[k.lower()] for k in arg.keys() if k.lower() in _unit_map} + unit_rev = {v: k for k, v in unit.items()} + list_cols = [unit_rev["year"], unit_rev["month"], unit_rev["day"]] + for u in ["h", "m", "s", "ms", "us"]: + value = unit_rev.get(u) + if value is not None and value in arg: + list_cols.append(value) + + psdf = arg[list_cols] + return psdf.pandas_on_spark.transform_batch(pandas_to_datetime, list_cols) + return pd.to_datetime( + arg, + errors=errors, + format=format, + unit=unit, + infer_datetime_format=infer_datetime_format, + origin=origin, + ) + + +# TODO(SPARK-42621): Add `inclusive` parameter and replace `closed`. +# See https://github.com/pandas-dev/pandas/issues/40245 +def date_range( + start: Union[str, Any] = None, + end: Union[str, Any] = None, + periods: Optional[int] = None, + freq: Optional[Union[str, DateOffset]] = None, + tz: Optional[Union[str, tzinfo]] = None, + normalize: bool = False, + name: Optional[str] = None, + closed: Optional[str] = None, + **kwargs: Any, +) -> DatetimeIndex: + """ + Return a fixed frequency DatetimeIndex. + + Parameters + ---------- + start : str or datetime-like, optional + Left bound for generating dates. + end : str or datetime-like, optional + Right bound for generating dates. + periods : int, optional + Number of periods to generate. + freq : str or DateOffset, default 'D' + Frequency strings can have multiples, e.g. '5H'. + tz : str or tzinfo, optional + Time zone name for returning localized DatetimeIndex, for example + 'Asia/Hong_Kong'. By default, the resulting DatetimeIndex is + time zone naive. + normalize : bool, default False + Normalize start/end dates to midnight before generating date range. + name : str, default None + Name of the resulting DatetimeIndex. + closed : {None, 'left', 'right'}, optional + Make the interval closed with respect to the given frequency to + the 'left', 'right', or both sides (None, the default). + + .. deprecated:: 3.4.0 + + **kwargs + For compatibility. Has no effect on the result. + + Returns + ------- + rng : DatetimeIndex + + See Also + -------- + DatetimeIndex : An immutable container for datetimes. + + Notes + ----- + Of the four parameters ``start``, ``end``, ``periods``, and ``freq``, + exactly three must be specified. If ``freq`` is omitted, the resulting + ``DatetimeIndex`` will have ``periods`` linearly spaced elements between + ``start`` and ``end`` (closed on both sides). + + To learn more about the frequency strings, please see `this link + `__. + + Examples + -------- + **Specifying the values** + + The next four examples generate the same `DatetimeIndex`, but vary + the combination of `start`, `end` and `periods`. + + Specify `start` and `end`, with the default daily frequency. + + >>> ps.date_range(start='1/1/2018', end='1/08/2018') # doctest: +SKIP + DatetimeIndex(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-04', + '2018-01-05', '2018-01-06', '2018-01-07', '2018-01-08'], + dtype='datetime64[ns]', freq=None) + + Specify `start` and `periods`, the number of periods (days). + + >>> ps.date_range(start='1/1/2018', periods=8) # doctest: +SKIP + DatetimeIndex(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-04', + '2018-01-05', '2018-01-06', '2018-01-07', '2018-01-08'], + dtype='datetime64[ns]', freq=None) + + Specify `end` and `periods`, the number of periods (days). + + >>> ps.date_range(end='1/1/2018', periods=8) # doctest: +SKIP + DatetimeIndex(['2017-12-25', '2017-12-26', '2017-12-27', '2017-12-28', + '2017-12-29', '2017-12-30', '2017-12-31', '2018-01-01'], + dtype='datetime64[ns]', freq=None) + + Specify `start`, `end`, and `periods`; the frequency is generated + automatically (linearly spaced). + + >>> ps.date_range( + ... start='2018-04-24', end='2018-04-27', periods=3 + ... ) # doctest: +SKIP + DatetimeIndex(['2018-04-24 00:00:00', '2018-04-25 12:00:00', + '2018-04-27 00:00:00'], + dtype='datetime64[ns]', freq=None) + + **Other Parameters** + + Changed the `freq` (frequency) to ``'M'`` (month end frequency). + + >>> ps.date_range(start='1/1/2018', periods=5, freq='M') # doctest: +SKIP + DatetimeIndex(['2018-01-31', '2018-02-28', '2018-03-31', '2018-04-30', + '2018-05-31'], + dtype='datetime64[ns]', freq=None) + + Multiples are allowed + + >>> ps.date_range(start='1/1/2018', periods=5, freq='3M') # doctest: +SKIP + DatetimeIndex(['2018-01-31', '2018-04-30', '2018-07-31', '2018-10-31', + '2019-01-31'], + dtype='datetime64[ns]', freq=None) + + `freq` can also be specified as an Offset object. + + >>> ps.date_range( + ... start='1/1/2018', periods=5, freq=pd.offsets.MonthEnd(3) + ... ) # doctest: +SKIP + DatetimeIndex(['2018-01-31', '2018-04-30', '2018-07-31', '2018-10-31', + '2019-01-31'], + dtype='datetime64[ns]', freq=None) + + `closed` controls whether to include `start` and `end` that are on the + boundary. The default includes boundary points on either end. + + >>> ps.date_range( + ... start='2017-01-01', end='2017-01-04', closed=None + ... ) # doctest: +SKIP + DatetimeIndex(['2017-01-01', '2017-01-02', '2017-01-03', '2017-01-04'], + dtype='datetime64[ns]', freq=None) + + Use ``closed='left'`` to exclude `end` if it falls on the boundary. + + >>> ps.date_range( + ... start='2017-01-01', end='2017-01-04', closed='left' + ... ) # doctest: +SKIP + DatetimeIndex(['2017-01-01', '2017-01-02', '2017-01-03'], dtype='datetime64[ns]', freq=None) + + Use ``closed='right'`` to exclude `start` if it falls on the boundary. + + >>> ps.date_range( + ... start='2017-01-01', end='2017-01-04', closed='right' + ... ) # doctest: +SKIP + DatetimeIndex(['2017-01-02', '2017-01-03', '2017-01-04'], dtype='datetime64[ns]', freq=None) + """ + assert freq not in ["N", "ns"], "nanoseconds is not supported" + assert tz is None, "Localized DatetimeIndex is not supported" + if closed is not None: + warnings.warn( + "Argument `closed` is deprecated in 3.4.0 and will be removed in 4.0.0.", + FutureWarning, + ) + + return cast( + DatetimeIndex, + ps.from_pandas( + pd.date_range( + start=start, + end=end, + periods=periods, + freq=freq, + tz=tz, + normalize=normalize, + name=name, + closed=closed, + **kwargs, + ) + ), + ) + + +@no_type_check +def to_timedelta( + arg, + unit: Optional[str] = None, + errors: str = "raise", +): + """ + Convert argument to timedelta. + + Parameters + ---------- + arg : str, timedelta, list-like or Series + The data to be converted to timedelta. + unit : str, optional + Denotes the unit of the arg for numeric `arg`. Defaults to ``"ns"``. + + Possible values: + * 'W' + * 'D' / 'days' / 'day' + * 'hours' / 'hour' / 'hr' / 'h' + * 'm' / 'minute' / 'min' / 'minutes' / 'T' + * 'S' / 'seconds' / 'sec' / 'second' + * 'ms' / 'milliseconds' / 'millisecond' / 'milli' / 'millis' / 'L' + * 'us' / 'microseconds' / 'microsecond' / 'micro' / 'micros' / 'U' + * 'ns' / 'nanoseconds' / 'nano' / 'nanos' / 'nanosecond' / 'N' + + Must not be specified when `arg` context strings and ``errors="raise"``. + errors : {'ignore', 'raise', 'coerce'}, default 'raise' + - If 'raise', then invalid parsing will raise an exception. + - If 'coerce', then invalid parsing will be set as NaT. + - If 'ignore', then invalid parsing will return the input. + + Returns + ------- + ret : timedelta64, TimedeltaIndex or Series of timedelta64 if parsing succeeded. + + See Also + -------- + DataFrame.astype : Cast argument to a specified dtype. + to_datetime : Convert argument to datetime. + + Notes + ----- + If the precision is higher than nanoseconds, the precision of the duration is + truncated to nanoseconds for string inputs. + + Examples + -------- + Parsing a single string to a Timedelta: + + >>> ps.to_timedelta('1 days 06:05:01.00003') + Timedelta('1 days 06:05:01.000030') + >>> ps.to_timedelta('15.5us') # doctest: +SKIP + Timedelta('0 days 00:00:00.000015500') + + Parsing a list or array of strings: + + >>> ps.to_timedelta(['1 days 06:05:01.00003', '15.5us', 'nan']) # doctest: +SKIP + TimedeltaIndex(['1 days 06:05:01.000030', '0 days 00:00:00.000015500', NaT], + dtype='timedelta64[ns]', freq=None) + + Converting numbers by specifying the `unit` keyword argument: + + >>> ps.to_timedelta(np.arange(5), unit='s') # doctest: +SKIP + TimedeltaIndex(['0 days 00:00:00', '0 days 00:00:01', '0 days 00:00:02', + '0 days 00:00:03', '0 days 00:00:04'], + dtype='timedelta64[ns]', freq=None) + >>> ps.to_timedelta(np.arange(5), unit='d') # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['0 days', '1 days', '2 days', '3 days', '4 days'], + dtype='timedelta64[ns]', freq=None) + """ + + def pandas_to_timedelta(pser: pd.Series) -> np.timedelta64: + return pd.to_timedelta( + arg=pser, + unit=unit, + errors=errors, + ) + + if isinstance(arg, Series): + return arg.transform(pandas_to_timedelta) + + else: + return pd.to_timedelta( + arg=arg, + unit=unit, + errors=errors, + ) + + +def timedelta_range( + start: Union[str, Any] = None, + end: Union[str, Any] = None, + periods: Optional[int] = None, + freq: Optional[Union[str, DateOffset]] = None, + name: Optional[str] = None, + closed: Optional[str] = None, +) -> TimedeltaIndex: + """ + Return a fixed frequency TimedeltaIndex, with day as the default frequency. + + Parameters + ---------- + start : str or timedelta-like, optional + Left bound for generating timedeltas. + end : str or timedelta-like, optional + Right bound for generating timedeltas. + periods : int, optional + Number of periods to generate. + freq : str or DateOffset, default 'D' + Frequency strings can have multiples, e.g. '5H'. + name : str, default None + Name of the resulting TimedeltaIndex. + closed : {None, 'left', 'right'}, optional + Make the interval closed with respect to the given frequency to + the 'left', 'right', or both sides (None, the default). + + Returns + ------- + TimedeltaIndex + + Notes + ----- + Of the four parameters ``start``, ``end``, ``periods``, and ``freq``, + exactly three must be specified. If ``freq`` is omitted, the resulting + ``TimedeltaIndex`` will have ``periods`` linearly spaced elements between + ``start`` and ``end`` (closed on both sides). + + To learn more about the frequency strings, please see `this link + `__. + + Examples + -------- + >>> ps.timedelta_range(start='1 day', periods=4) # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days', '2 days', '3 days', '4 days'], dtype='timedelta64[ns]', freq=None) + + The closed parameter specifies which endpoint is included. + The default behavior is to include both endpoints. + + >>> ps.timedelta_range(start='1 day', periods=4, closed='right') + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['2 days', '3 days', '4 days'], dtype='timedelta64[ns]', freq=None) + + The freq parameter specifies the frequency of the TimedeltaIndex. + Only fixed frequencies can be passed, non-fixed frequencies such as ‘M’ (month end) will raise. + + >>> ps.timedelta_range(start='1 day', end='2 days', freq='6H') + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days 00:00:00', '1 days 06:00:00', '1 days 12:00:00', + '1 days 18:00:00', '2 days 00:00:00'], + dtype='timedelta64[ns]', freq=None) + + Specify start, end, and periods; the frequency is generated automatically (linearly spaced). + + >>> ps.timedelta_range(start='1 day', end='5 days', periods=4) + ... # doctest: +NORMALIZE_WHITESPACE + TimedeltaIndex(['1 days 00:00:00', '2 days 08:00:00', '3 days 16:00:00', + '5 days 00:00:00'], + dtype='timedelta64[ns]', freq=None) + """ + assert freq not in ["N", "ns"], "nanoseconds is not supported" + + return cast( + TimedeltaIndex, + ps.from_pandas( + pd.timedelta_range( + start=start, + end=end, + periods=periods, + freq=freq, + name=name, + closed=closed, + ) + ), + ) + + +def get_dummies( + data: Union[DataFrame, Series], + prefix: Optional[Union[str, List[str], Dict[str, str]]] = None, + prefix_sep: str = "_", + dummy_na: bool = False, + columns: Optional[Union[Name, List[Name]]] = None, + sparse: bool = False, + drop_first: bool = False, + dtype: Optional[Union[str, Dtype]] = None, +) -> DataFrame: + """ + Convert categorical variable into dummy/indicator variables, also + known as one hot encoding. + + Parameters + ---------- + data : array-like, Series, or DataFrame + prefix : string, list of strings, or dict of strings, default None + String to append DataFrame column names. + Pass a list with length equal to the number of columns + when calling get_dummies on a DataFrame. Alternatively, `prefix` + can be a dictionary mapping column names to prefixes. + prefix_sep : string, default '_' + If appending prefix, separator/delimiter to use. Or pass a + list or dictionary as with `prefix.` + dummy_na : bool, default False + Add a column to indicate NaNs, if False NaNs are ignored. + columns : list-like, default None + Column names in the DataFrame to be encoded. + If `columns` is None then all the columns with + `object` or `category` dtype will be converted. + sparse : bool, default False + Whether the dummy-encoded columns should be be backed by + a :class:`SparseArray` (True) or a regular NumPy array (False). + In pandas-on-Spark, this value must be "False". + drop_first : bool, default False + Whether to get k-1 dummies out of k categorical levels by removing the + first level. + dtype : dtype, default np.uint8 + Data type for new columns. Only a single dtype is allowed. + + Returns + ------- + dummies : DataFrame + + See Also + -------- + Series.str.get_dummies + + Examples + -------- + >>> s = ps.Series(list('abca')) + + >>> ps.get_dummies(s) + a b c + 0 1 0 0 + 1 0 1 0 + 2 0 0 1 + 3 1 0 0 + + >>> df = ps.DataFrame({'A': ['a', 'b', 'a'], 'B': ['b', 'a', 'c'], + ... 'C': [1, 2, 3]}, + ... columns=['A', 'B', 'C']) + + >>> ps.get_dummies(df, prefix=['col1', 'col2']) + C col1_a col1_b col2_a col2_b col2_c + 0 1 1 0 0 1 0 + 1 2 0 1 1 0 0 + 2 3 1 0 0 0 1 + + >>> ps.get_dummies(ps.Series(list('abcaa'))) + a b c + 0 1 0 0 + 1 0 1 0 + 2 0 0 1 + 3 1 0 0 + 4 1 0 0 + + >>> ps.get_dummies(ps.Series(list('abcaa')), drop_first=True) + b c + 0 0 0 + 1 1 0 + 2 0 1 + 3 0 0 + 4 0 0 + + >>> ps.get_dummies(ps.Series(list('abc')), dtype=float) + a b c + 0 1.0 0.0 0.0 + 1 0.0 1.0 0.0 + 2 0.0 0.0 1.0 + """ + if sparse is not False: + raise NotImplementedError("get_dummies currently does not support sparse") + + if columns is not None and not is_list_like(columns): + raise TypeError("Input must be a list-like for parameter `columns`") + + if dtype is None: + dtype = "byte" + + if isinstance(data, Series): + if prefix is not None: + prefix = [str(prefix)] + psdf = data.to_frame() + column_labels = psdf._internal.column_labels + remaining_columns = [] + else: + if isinstance(prefix, str): + raise NotImplementedError( + "get_dummies currently does not support prefix as string types" + ) + psdf = data.copy() + + if columns is None: + column_labels = [ + label + for label in psdf._internal.column_labels + if isinstance( + psdf._internal.spark_type_for(label), _get_dummies_default_accept_types + ) + ] + else: + if is_name_like_tuple(columns): + column_labels = [ + label + for label in psdf._internal.column_labels + if label[: len(columns)] == columns + ] + if len(column_labels) == 0: + raise KeyError(name_like_string(columns)) + if prefix is None: + prefix = [ + str(label[len(columns) :]) + if len(label) > len(columns) + 1 + else label[len(columns)] + if len(label) == len(columns) + 1 + else "" + for label in column_labels + ] + elif any(isinstance(col, tuple) for col in columns) and any( + not is_name_like_tuple(col) for col in columns + ): + raise ValueError( + "Expected tuple, got {}".format( + type(set(col for col in columns if not is_name_like_tuple(col)).pop()) + ) + ) + else: + column_labels = [ + label + for key in columns + for label in psdf._internal.column_labels + if label == key or label[0] == key + ] + if len(column_labels) == 0: + if columns is None: + return psdf + raise KeyError("{} not in index".format(columns)) + + if prefix is None: + prefix = [str(label) if len(label) > 1 else label[0] for label in column_labels] + + column_labels_set = set(column_labels) + remaining_columns = [ + ( + psdf[label] + if psdf._internal.column_labels_level == 1 + else psdf[label].rename(name_like_string(label)) + ) + for label in psdf._internal.column_labels + if label not in column_labels_set + ] + + if any( + not isinstance(psdf._internal.spark_type_for(label), _get_dummies_acceptable_types) + for label in column_labels + ): + raise NotImplementedError( + "get_dummies currently only accept {} values".format( + ", ".join( + [cast(Type[DataType], t).typeName() for t in _get_dummies_acceptable_types] + ) + ) + ) + + if prefix is not None and len(column_labels) != len(prefix): + raise ValueError( + "Length of 'prefix' ({}) did not match the length of " + "the columns being encoded ({}).".format(len(prefix), len(column_labels)) + ) + elif isinstance(prefix, dict): + prefix = [prefix[column_label[0]] for column_label in column_labels] + + all_values = _reduce_spark_multi( + psdf._internal.spark_frame, + [F.collect_set(psdf._internal.spark_column_for(label)) for label in column_labels], + ) + for i, label in enumerate(column_labels): + values = all_values[i] + if isinstance(values, np.ndarray): + values = values.tolist() + values = sorted(values) + if drop_first: + values = values[1:] + + def column_name(v: Any) -> Name: + if prefix is None or cast(List[str], prefix)[i] == "": + return v + else: + return "{}{}{}".format(cast(List[str], prefix)[i], prefix_sep, v) + + for value in values: + remaining_columns.append( + (psdf[label].notnull() & (psdf[label] == value)) + .astype(dtype) + .rename(column_name(value)) + ) + if dummy_na: + remaining_columns.append(psdf[label].isnull().astype(dtype).rename(column_name(np.nan))) + + return psdf[remaining_columns] + + +# TODO: there are many parameters to implement and support. See pandas's pd.concat. +def concat( + objs: List[Union[DataFrame, Series]], + axis: Axis = 0, + join: str = "outer", + ignore_index: bool = False, + sort: bool = False, +) -> Union[Series, DataFrame]: + """ + Concatenate pandas-on-Spark objects along a particular axis with optional set logic + along the other axes. + + Parameters + ---------- + objs : a sequence of Series or DataFrame + Any None objects will be dropped silently unless + they are all None in which case a ValueError will be raised + axis : {0/'index', 1/'columns'}, default 0 + The axis to concatenate along. + join : {'inner', 'outer'}, default 'outer' + How to handle indexes on other axis (or axes). + ignore_index : bool, default False + If True, do not use the index values along the concatenation axis. The + resulting axis will be labeled 0, ..., n - 1. This is useful if you are + concatenating objects where the concatenation axis does not have + meaningful indexing information. Note the index values on the other + axes are still respected in the join. + sort : bool, default False + Sort non-concatenation axis if it is not already aligned. + + Returns + ------- + object, type of objs + When concatenating all ``Series`` along the index (axis=0), a + ``Series`` is returned. When ``objs`` contains at least one + ``DataFrame``, a ``DataFrame`` is returned. When concatenating along + the columns (axis=1), a ``DataFrame`` is returned. + + See Also + -------- + Series.append : Concatenate Series. + DataFrame.join : Join DataFrames using indexes. + DataFrame.merge : Merge DataFrames by indexes or columns. + + Examples + -------- + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + + Combine two ``Series``. + + >>> s1 = ps.Series(['a', 'b']) + >>> s2 = ps.Series(['c', 'd']) + >>> ps.concat([s1, s2]) + 0 a + 1 b + 0 c + 1 d + dtype: object + + Clear the existing index and reset it in the result + by setting the ``ignore_index`` option to ``True``. + + >>> ps.concat([s1, s2], ignore_index=True) + 0 a + 1 b + 2 c + 3 d + dtype: object + + Combine two ``DataFrame`` objects with identical columns. + + >>> df1 = ps.DataFrame([['a', 1], ['b', 2]], + ... columns=['letter', 'number']) + >>> df1 + letter number + 0 a 1 + 1 b 2 + >>> df2 = ps.DataFrame([['c', 3], ['d', 4]], + ... columns=['letter', 'number']) + >>> df2 + letter number + 0 c 3 + 1 d 4 + + >>> ps.concat([df1, df2]) + letter number + 0 a 1 + 1 b 2 + 0 c 3 + 1 d 4 + + Combine ``DataFrame`` and ``Series`` objects with different columns. + + >>> ps.concat([df2, s1]) + letter number 0 + 0 c 3.0 None + 1 d 4.0 None + 0 None NaN a + 1 None NaN b + + Combine ``DataFrame`` objects with overlapping columns + and return everything. Columns outside the intersection will + be filled with ``None`` values. + + >>> df3 = ps.DataFrame([['c', 3, 'cat'], ['d', 4, 'dog']], + ... columns=['letter', 'number', 'animal']) + >>> df3 + letter number animal + 0 c 3 cat + 1 d 4 dog + + >>> ps.concat([df1, df3]) + letter number animal + 0 a 1 None + 1 b 2 None + 0 c 3 cat + 1 d 4 dog + + Sort the columns. + + >>> ps.concat([df1, df3], sort=True) + animal letter number + 0 None a 1 + 1 None b 2 + 0 cat c 3 + 1 dog d 4 + + Combine ``DataFrame`` objects with overlapping columns + and return only those that are shared by passing ``inner`` to + the ``join`` keyword argument. + + >>> ps.concat([df1, df3], join="inner") + letter number + 0 a 1 + 1 b 2 + 0 c 3 + 1 d 4 + + >>> df4 = ps.DataFrame([['bird', 'polly'], ['monkey', 'george']], + ... columns=['animal', 'name']) + + Combine with column axis. + + >>> ps.concat([df1, df4], axis=1) + letter number animal name + 0 a 1 bird polly + 1 b 2 monkey george + + >>> reset_option("compute.ops_on_diff_frames") + """ + if isinstance(objs, (DataFrame, IndexOpsMixin)) or not isinstance( + objs, Iterable + ): # TODO: support dict + raise TypeError( + "first argument must be an iterable of pandas-on-Spark " + "objects, you passed an object of type " + '"{name}"'.format(name=type(objs).__name__) + ) + + if len(cast(Sized, objs)) == 0: + raise ValueError("No objects to concatenate") + objs = list(filter(lambda obj: obj is not None, objs)) + if len(objs) == 0: + raise ValueError("All objects passed were None") + + for obj in objs: + if not isinstance(obj, (Series, DataFrame)): + raise TypeError( + "cannot concatenate object of type " + "'{name}" + "; only ps.Series " + "and ps.DataFrame are valid".format(name=type(objs).__name__) + ) + + if join not in ["inner", "outer"]: + raise ValueError("Only can inner (intersect) or outer (union) join the other axis.") + + axis = validate_axis(axis) + psdf: DataFrame + if axis == 1: + psdfs: List[DataFrame] = [ + obj.to_frame() if isinstance(obj, Series) else obj for obj in objs + ] + + level: int = min(psdf._internal.column_labels_level for psdf in psdfs) + psdfs = [ + DataFrame._index_normalized_frame(level, psdf) + if psdf._internal.column_labels_level > level + else psdf + for psdf in psdfs + ] + + concat_psdf = psdfs[0] + column_labels: List[Label] = concat_psdf._internal.column_labels.copy() + + psdfs_not_same_anchor = [] + for psdf in psdfs[1:]: + duplicated = [label for label in psdf._internal.column_labels if label in column_labels] + if len(duplicated) > 0: + pretty_names = [name_like_string(label) for label in duplicated] + raise ValueError( + "Labels have to be unique; however, got duplicated labels %s." % pretty_names + ) + column_labels.extend(psdf._internal.column_labels) + + if same_anchor(concat_psdf, psdf): + concat_psdf = DataFrame( + concat_psdf._internal.with_new_columns( + [ + concat_psdf._psser_for(label) + for label in concat_psdf._internal.column_labels + ] + + [psdf._psser_for(label) for label in psdf._internal.column_labels] + ) + ) + else: + psdfs_not_same_anchor.append(psdf) + + if len(psdfs_not_same_anchor) > 0: + + @no_type_check + def resolve_func(psdf, this_column_labels, that_column_labels): + raise AssertionError("This should not happen.") + + for psdf in psdfs_not_same_anchor: + if join == "inner": + concat_psdf = align_diff_frames( + resolve_func, + concat_psdf, + psdf, + fillna=False, + how="inner", + ) + elif join == "outer": + concat_psdf = align_diff_frames( + resolve_func, + concat_psdf, + psdf, + fillna=False, + how="full", + ) + + concat_psdf = concat_psdf[column_labels] + + if ignore_index: + concat_psdf.columns = list( # type: ignore[assignment] + map(str, _range(len(concat_psdf.columns))) + ) + + if sort: + concat_psdf = concat_psdf.sort_index() + + return concat_psdf + + # Series, Series ... + # We should return Series if objects are all Series. + should_return_series = all(map(lambda obj: isinstance(obj, Series), objs)) + + # DataFrame, Series ... & Series, Series ... + # In this case, we should return DataFrame. + new_objs: List[DataFrame] = [] + num_series = 0 + series_names = set() + for obj in objs: + if isinstance(obj, Series): + num_series += 1 + series_names.add(obj.name) + new_objs.append(obj.to_frame(DEFAULT_SERIES_NAME)) + else: + assert isinstance(obj, DataFrame) + new_objs.append(obj) + + column_labels_levels: Set[int] = set(obj._internal.column_labels_level for obj in new_objs) + if len(column_labels_levels) != 1: + raise ValueError("MultiIndex columns should have the same levels") + + # DataFrame, DataFrame, ... + # All Series are converted into DataFrame and then compute concat. + if not ignore_index: + indices_of_psdfs = [psdf.index for psdf in new_objs] + index_of_first_psdf = indices_of_psdfs[0] + for index_of_psdf in indices_of_psdfs: + if index_of_first_psdf.names != index_of_psdf.names: + raise ValueError( + "Index type and names should be same in the objects to concatenate. " + "You passed different indices " + "{index_of_first_psdf} and {index_of_psdf}".format( + index_of_first_psdf=index_of_first_psdf.names, + index_of_psdf=index_of_psdf.names, + ) + ) + + column_labels_of_psdfs = [psdf._internal.column_labels for psdf in new_objs] + index_names_of_psdfs: List[List[Optional[Label]]] + if ignore_index: + index_names_of_psdfs = [[] for _ in new_objs] + else: + index_names_of_psdfs = [psdf._internal.index_names for psdf in new_objs] + + if all(name == index_names_of_psdfs[0] for name in index_names_of_psdfs) and all( + idx == column_labels_of_psdfs[0] for idx in column_labels_of_psdfs + ): + # If all columns are in the same order and values, use it. + psdfs = new_objs + else: + if join == "inner": + interested_columns = set.intersection(*map(lambda x: set(x), column_labels_of_psdfs)) + # Keep the column order with its firsts DataFrame. + merged_columns = [ + label for label in column_labels_of_psdfs[0] if label in interested_columns + ] + + # If sort is True, sort to follow pandas 1.4+ behavior. + if sort: + # FIXME: better ordering + merged_columns = sorted(merged_columns, key=name_like_string) + + psdfs = [psdf[merged_columns] for psdf in new_objs] + elif join == "outer": + merged_columns = [] + for labels in column_labels_of_psdfs: + merged_columns.extend(label for label in labels if label not in merged_columns) + + assert len(merged_columns) > 0 + + # If sort is True, always sort + if sort: + # FIXME: better ordering + merged_columns = sorted(merged_columns, key=name_like_string) + + psdfs = [] + for psdf in new_objs: + columns_to_add = list(set(merged_columns) - set(psdf._internal.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)) + + data_columns = psdf._internal.data_spark_column_names + [ + name_like_string(label) for label in columns_to_add + ] + psdf = DataFrame( + psdf._internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in psdf._internal.index_spark_column_names + ], + column_labels=(psdf._internal.column_labels + columns_to_add), + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + data_fields=(psdf._internal.data_fields + ([None] * len(columns_to_add))), + ) + ) + + psdfs.append(psdf[merged_columns]) + + if ignore_index: + sdfs = [ + psdf._internal.spark_frame.select(psdf._internal.data_spark_columns) for psdf in psdfs + ] + else: + sdfs = [ + psdf._internal.spark_frame.select( + psdf._internal.index_spark_columns + psdf._internal.data_spark_columns + ) + for psdf in psdfs + ] + concatenated = reduce(lambda x, y: x.union(y), sdfs) + + if ignore_index: + index_spark_column_names = [] + index_names = [] + index_fields = [] + else: + index_spark_column_names = psdfs[0]._internal.index_spark_column_names + index_names = psdfs[0]._internal.index_names + index_fields = psdfs[0]._internal.index_fields + + result_psdf: DataFrame = DataFrame( + psdfs[0]._internal.copy( + spark_frame=concatenated, + index_spark_columns=[scol_for(concatenated, col) for col in index_spark_column_names], + index_names=index_names, + index_fields=index_fields, + data_spark_columns=[ + scol_for(concatenated, col) for col in psdfs[0]._internal.data_spark_column_names + ], + data_fields=None, # TODO: dtypes? + ) + ) + + if should_return_series: + # If all input were Series, we should return Series. + if len(series_names) == 1: + name = series_names.pop() + else: + name = None + return first_series(result_psdf).rename(name) + else: + return result_psdf + + +def melt( + frame: DataFrame, + id_vars: Optional[Union[Name, List[Name]]] = None, + value_vars: Optional[Union[Name, List[Name]]] = None, + var_name: Optional[Union[str, List[str]]] = None, + value_name: str = "value", +) -> DataFrame: + return DataFrame.melt(frame, id_vars, value_vars, var_name, value_name) + + +melt.__doc__ = DataFrame.melt.__doc__ + + +@no_type_check +def isna(obj): + """ + Detect missing values for an array-like object. + + This function takes a scalar or array-like object and indicates + whether values are missing (``NaN`` in numeric arrays, ``None`` or ``NaN`` + in object arrays). + + Parameters + ---------- + obj : scalar or array-like + Object to check for null or missing values. + + Returns + ------- + bool or array-like of bool + For scalar input, returns a scalar boolean. + For array input, returns an array of boolean indicating whether each + corresponding element is missing. + + See Also + -------- + Series.isna : Detect missing values in a Series. + Series.isnull : Detect missing values in a Series. + DataFrame.isna : Detect missing values in a DataFrame. + DataFrame.isnull : Detect missing values in a DataFrame. + Index.isna : Detect missing values in an Index. + Index.isnull : Detect missing values in an Index. + + Examples + -------- + Scalar arguments (including strings) result in a scalar boolean. + + >>> ps.isna('dog') + False + + >>> ps.isna(np.nan) + True + + ndarrays result in an ndarray of booleans. + + >>> array = np.array([[1, np.nan, 3], [4, 5, np.nan]]) + >>> array + array([[ 1., nan, 3.], + [ 4., 5., nan]]) + >>> ps.isna(array) + array([[False, True, False], + [False, False, True]]) + + For Series and DataFrame, the same type is returned, containing booleans. + + >>> df = ps.DataFrame({'a': ['ant', 'bee', 'cat'], 'b': ['dog', None, 'fly']}) + >>> df + a b + 0 ant dog + 1 bee None + 2 cat fly + + >>> ps.isna(df) + a b + 0 False False + 1 False True + 2 False False + + >>> ps.isnull(df.b) + 0 False + 1 True + 2 False + Name: b, dtype: bool + """ + # TODO: Add back: + # notnull : Boolean inverse of pandas.isnull. + # into the See Also in the docstring. It does not find the method in the latest numpydoc. + if isinstance(obj, (DataFrame, Series)): + return obj.isnull() + else: + return pd.isnull(obj) + + +isnull = isna + + +@no_type_check +def notna(obj): + """ + Detect existing (non-missing) values. + + Return a boolean same-sized object indicating if the values are not NA. + Non-missing values get mapped to True. NA values, such as None or + :attr:`numpy.NaN`, get mapped to False values. + + Returns + ------- + bool or array-like of bool + Mask of bool values for each element that + indicates whether an element is not an NA value. + + See Also + -------- + isna : Detect missing values for an array-like object. + Series.notna : Boolean inverse of Series.isna. + DataFrame.notnull : Boolean inverse of DataFrame.isnull. + Index.notna : Boolean inverse of Index.isna. + Index.notnull : Boolean inverse of Index.isnull. + + Examples + -------- + Show which entries in a DataFrame are not NA. + + >>> df = ps.DataFrame({'age': [5, 6, np.NaN], + ... 'born': [pd.NaT, pd.Timestamp('1939-05-27'), + ... pd.Timestamp('1940-04-25')], + ... 'name': ['Alfred', 'Batman', ''], + ... 'toy': [None, 'Batmobile', 'Joker']}) + >>> df + age born name toy + 0 5.0 NaT Alfred None + 1 6.0 1939-05-27 Batman Batmobile + 2 NaN 1940-04-25 Joker + + >>> df.notnull() + age born name toy + 0 True False True False + 1 True True True True + 2 False True True True + + Show which entries in a Series are not NA. + + >>> ser = ps.Series([5, 6, np.NaN]) + >>> ser + 0 5.0 + 1 6.0 + 2 NaN + dtype: float64 + + >>> ps.notna(ser) + 0 True + 1 True + 2 False + dtype: bool + + >>> ps.notna(ser.index) + True + """ + # TODO: Add back: + # Series.notnull :Boolean inverse of Series.isnull. + # DataFrame.notna :Boolean inverse of DataFrame.isna. + # into the See Also in the docstring. It does not find the method in the latest numpydoc. + if isinstance(obj, (DataFrame, Series)): + return obj.notna() + else: + return pd.notna(obj) + + +notnull = notna + + +def merge( + obj: DataFrame, + right: DataFrame, + how: str = "inner", + on: Optional[Union[Name, List[Name]]] = None, + left_on: Optional[Union[Name, List[Name]]] = None, + right_on: Optional[Union[Name, List[Name]]] = None, + left_index: bool = False, + right_index: bool = False, + suffixes: Tuple[str, str] = ("_x", "_y"), +) -> "DataFrame": + """ + Merge DataFrame objects with a database-style join. + + The index of the resulting DataFrame will be one of the following: + - 0...n if no index is used for merging + - Index of the left DataFrame if merged only on the index of the right DataFrame + - Index of the right DataFrame if merged only on the index of the left DataFrame + - All involved indices if merged using the indices of both DataFrames + e.g. if `left` with indices (a, x) and `right` with indices (b, x), the result will + be an index (x, a, b) + + Parameters + ---------- + right: Object to merge with. + how: Type of merge to be performed. + {'left', 'right', 'outer', 'inner'}, default 'inner' + + left: use only keys from left frame, like a SQL left outer join; preserve key + order. + right: use only keys from right frame, like a SQL right outer join; preserve key + order. + outer: use union of keys from both frames, like a SQL full outer join; sort keys + lexicographically. + inner: use intersection of keys from both frames, like a SQL inner join; + preserve the order of the left keys. + on: Column or index level names to join on. These must be found in both DataFrames. If on + is None and not merging on indexes then this defaults to the intersection of the + columns in both DataFrames. + left_on: Column or index level names to join on in the left DataFrame. Can also + be an array or list of arrays of the length of the left DataFrame. + These arrays are treated as if they are columns. + right_on: Column or index level names to join on in the right DataFrame. Can also + be an array or list of arrays of the length of the right DataFrame. + These arrays are treated as if they are columns. + left_index: Use the index from the left DataFrame as the join key(s). If it is a + MultiIndex, the number of keys in the other DataFrame (either the index or a number of + columns) must match the number of levels. + right_index: Use the index from the right DataFrame as the join key. Same caveats as + left_index. + suffixes: Suffix to apply to overlapping column names in the left and right side, + respectively. + + Returns + ------- + DataFrame + A DataFrame of the two merged objects. + + Examples + -------- + + >>> df1 = ps.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [1, 2, 3, 5]}, + ... columns=['lkey', 'value']) + >>> df2 = ps.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [5, 6, 7, 8]}, + ... columns=['rkey', 'value']) + >>> df1 + lkey value + 0 foo 1 + 1 bar 2 + 2 baz 3 + 3 foo 5 + >>> df2 + rkey value + 0 foo 5 + 1 bar 6 + 2 baz 7 + 3 foo 8 + + Merge df1 and df2 on the lkey and rkey columns. The value columns have + the default suffixes, _x and _y, appended. + + >>> merged = ps.merge(df1, df2, left_on='lkey', right_on='rkey') + >>> merged.sort_values(by=['lkey', 'value_x', 'rkey', 'value_y']) # doctest: +ELLIPSIS + lkey value_x rkey value_y + ...bar 2 bar 6 + ...baz 3 baz 7 + ...foo 1 foo 5 + ...foo 1 foo 8 + ...foo 5 foo 5 + ...foo 5 foo 8 + + >>> left_psdf = ps.DataFrame({'A': [1, 2]}) + >>> right_psdf = ps.DataFrame({'B': ['x', 'y']}, index=[1, 2]) + + >>> ps.merge(left_psdf, right_psdf, left_index=True, right_index=True).sort_index() + A B + 1 2 x + + >>> ps.merge(left_psdf, right_psdf, left_index=True, right_index=True, how='left').sort_index() + A B + 0 1 None + 1 2 x + + >>> ps.merge(left_psdf, right_psdf, left_index=True, right_index=True, how='right').sort_index() + A B + 1 2.0 x + 2 NaN y + + >>> ps.merge(left_psdf, right_psdf, left_index=True, right_index=True, how='outer').sort_index() + A B + 0 1.0 None + 1 2.0 x + 2 NaN y + + Notes + ----- + As described in #263, joining string columns currently returns None for missing values + instead of NaN. + """ + return obj.merge( + right, + how=how, + on=on, + left_on=left_on, + right_on=right_on, + left_index=left_index, + right_index=right_index, + suffixes=suffixes, + ) + + +def merge_asof( + left: Union[DataFrame, Series], + right: Union[DataFrame, Series], + on: Optional[Name] = None, + left_on: Optional[Name] = None, + right_on: Optional[Name] = None, + left_index: bool = False, + right_index: bool = False, + by: Optional[Union[Name, List[Name]]] = None, + left_by: Optional[Union[Name, List[Name]]] = None, + right_by: Optional[Union[Name, List[Name]]] = None, + suffixes: Tuple[str, str] = ("_x", "_y"), + tolerance: Optional[Any] = None, + allow_exact_matches: bool = True, + direction: str = "backward", +) -> DataFrame: + """ + Perform an asof merge. + + This is like a left-join except that we match on nearest + key rather than equal keys. + + For each row in the left DataFrame: + + - A "backward" search selects the last row in the right DataFrame whose + 'on' key is less than or equal to the left's key. + + - A "forward" search selects the first row in the right DataFrame whose + 'on' key is greater than or equal to the left's key. + + - A "nearest" search selects the row in the right DataFrame who's 'on' + key is closest in absolute distance to the left's key. + + Optionally match on equivalent keys with 'by' before searching with 'on'. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + left : DataFrame or named Series + right : DataFrame or named Series + on : label + Field name to join on. Must be found in both DataFrames. + The data MUST be ordered. This must be a numeric column, + such as datetimelike, integer, or float. On or left_on/right_on + must be given. + left_on : label + Field name to join on in left DataFrame. + right_on : label + Field name to join on in right DataFrame. + left_index : bool + Use the index of the left DataFrame as the join key. + right_index : bool + Use the index of the right DataFrame as the join key. + by : column name or list of column names + Match on these columns before performing merge operation. + left_by : column name + Field names to match on in the left DataFrame. + right_by : column name + Field names to match on in the right DataFrame. + suffixes : 2-length sequence (tuple, list, ...) + Suffix to apply to overlapping column names in the left and right + side, respectively. + tolerance : int or Timedelta, optional, default None + Select asof tolerance within this range; must be compatible + with the merge index. + allow_exact_matches : bool, default True + + - If True, allow matching with the same 'on' value + (i.e. less-than-or-equal-to / greater-than-or-equal-to) + - If False, don't match the same 'on' value + (i.e., strictly less-than / strictly greater-than). + + direction : 'backward' (default), 'forward', or 'nearest' + Whether to search for prior, subsequent, or closest matches. + + Returns + ------- + merged : DataFrame + + See Also + -------- + merge : Merge with a database-style join. + merge_ordered : Merge with optional filling/interpolation. + + Examples + -------- + >>> left = ps.DataFrame({"a": [1, 5, 10], "left_val": ["a", "b", "c"]}) + >>> left + a left_val + 0 1 a + 1 5 b + 2 10 c + + >>> right = ps.DataFrame({"a": [1, 2, 3, 6, 7], "right_val": [1, 2, 3, 6, 7]}) + >>> right + a right_val + 0 1 1 + 1 2 2 + 2 3 3 + 3 6 6 + 4 7 7 + + >>> ps.merge_asof(left, right, on="a").sort_values("a").reset_index(drop=True) + a left_val right_val + 0 1 a 1 + 1 5 b 3 + 2 10 c 7 + + >>> ps.merge_asof( + ... left, + ... right, + ... on="a", + ... allow_exact_matches=False + ... ).sort_values("a").reset_index(drop=True) + a left_val right_val + 0 1 a NaN + 1 5 b 3.0 + 2 10 c 7.0 + + >>> ps.merge_asof( + ... left, + ... right, + ... on="a", + ... direction="forward" + ... ).sort_values("a").reset_index(drop=True) + a left_val right_val + 0 1 a 1.0 + 1 5 b 6.0 + 2 10 c NaN + + >>> ps.merge_asof( + ... left, + ... right, + ... on="a", + ... direction="nearest" + ... ).sort_values("a").reset_index(drop=True) + a left_val right_val + 0 1 a 1 + 1 5 b 6 + 2 10 c 7 + + We can use indexed DataFrames as well. + + >>> left = ps.DataFrame({"left_val": ["a", "b", "c"]}, index=[1, 5, 10]) + >>> left + left_val + 1 a + 5 b + 10 c + + >>> right = ps.DataFrame({"right_val": [1, 2, 3, 6, 7]}, index=[1, 2, 3, 6, 7]) + >>> right + right_val + 1 1 + 2 2 + 3 3 + 6 6 + 7 7 + + >>> ps.merge_asof(left, right, left_index=True, right_index=True).sort_index() + left_val right_val + 1 a 1 + 5 b 3 + 10 c 7 + + Here is a real-world times-series example + + >>> quotes = ps.DataFrame( + ... { + ... "time": [ + ... pd.Timestamp("2016-05-25 13:30:00.023"), + ... pd.Timestamp("2016-05-25 13:30:00.023"), + ... pd.Timestamp("2016-05-25 13:30:00.030"), + ... pd.Timestamp("2016-05-25 13:30:00.041"), + ... pd.Timestamp("2016-05-25 13:30:00.048"), + ... pd.Timestamp("2016-05-25 13:30:00.049"), + ... pd.Timestamp("2016-05-25 13:30:00.072"), + ... pd.Timestamp("2016-05-25 13:30:00.075") + ... ], + ... "ticker": [ + ... "GOOG", + ... "MSFT", + ... "MSFT", + ... "MSFT", + ... "GOOG", + ... "AAPL", + ... "GOOG", + ... "MSFT" + ... ], + ... "bid": [720.50, 51.95, 51.97, 51.99, 720.50, 97.99, 720.50, 52.01], + ... "ask": [720.93, 51.96, 51.98, 52.00, 720.93, 98.01, 720.88, 52.03] + ... } + ... ) + >>> quotes + time ticker bid ask + 0 2016-05-25 13:30:00.023 GOOG 720.50 720.93 + 1 2016-05-25 13:30:00.023 MSFT 51.95 51.96 + 2 2016-05-25 13:30:00.030 MSFT 51.97 51.98 + 3 2016-05-25 13:30:00.041 MSFT 51.99 52.00 + 4 2016-05-25 13:30:00.048 GOOG 720.50 720.93 + 5 2016-05-25 13:30:00.049 AAPL 97.99 98.01 + 6 2016-05-25 13:30:00.072 GOOG 720.50 720.88 + 7 2016-05-25 13:30:00.075 MSFT 52.01 52.03 + + >>> trades = ps.DataFrame( + ... { + ... "time": [ + ... pd.Timestamp("2016-05-25 13:30:00.023"), + ... pd.Timestamp("2016-05-25 13:30:00.038"), + ... pd.Timestamp("2016-05-25 13:30:00.048"), + ... pd.Timestamp("2016-05-25 13:30:00.048"), + ... pd.Timestamp("2016-05-25 13:30:00.048") + ... ], + ... "ticker": ["MSFT", "MSFT", "GOOG", "GOOG", "AAPL"], + ... "price": [51.95, 51.95, 720.77, 720.92, 98.0], + ... "quantity": [75, 155, 100, 100, 100] + ... } + ... ) + >>> trades + time ticker price quantity + 0 2016-05-25 13:30:00.023 MSFT 51.95 75 + 1 2016-05-25 13:30:00.038 MSFT 51.95 155 + 2 2016-05-25 13:30:00.048 GOOG 720.77 100 + 3 2016-05-25 13:30:00.048 GOOG 720.92 100 + 4 2016-05-25 13:30:00.048 AAPL 98.00 100 + + By default we are taking the asof of the quotes + + >>> ps.merge_asof( + ... trades, quotes, on="time", by="ticker" + ... ).sort_values(["time", "ticker", "price"]).reset_index(drop=True) + time ticker price quantity bid ask + 0 2016-05-25 13:30:00.023 MSFT 51.95 75 51.95 51.96 + 1 2016-05-25 13:30:00.038 MSFT 51.95 155 51.97 51.98 + 2 2016-05-25 13:30:00.048 AAPL 98.00 100 NaN NaN + 3 2016-05-25 13:30:00.048 GOOG 720.77 100 720.50 720.93 + 4 2016-05-25 13:30:00.048 GOOG 720.92 100 720.50 720.93 + + We only asof within 2ms between the quote time and the trade time + + >>> ps.merge_asof( + ... trades, + ... quotes, + ... on="time", + ... by="ticker", + ... tolerance=sf.expr("INTERVAL 2 MILLISECONDS") # pd.Timedelta("2ms") + ... ).sort_values(["time", "ticker", "price"]).reset_index(drop=True) + time ticker price quantity bid ask + 0 2016-05-25 13:30:00.023 MSFT 51.95 75 51.95 51.96 + 1 2016-05-25 13:30:00.038 MSFT 51.95 155 NaN NaN + 2 2016-05-25 13:30:00.048 AAPL 98.00 100 NaN NaN + 3 2016-05-25 13:30:00.048 GOOG 720.77 100 720.50 720.93 + 4 2016-05-25 13:30:00.048 GOOG 720.92 100 720.50 720.93 + + We only asof within 10ms between the quote time and the trade time + and we exclude exact matches on time. However *prior* data will + propagate forward + + >>> ps.merge_asof( + ... trades, + ... quotes, + ... on="time", + ... by="ticker", + ... tolerance=sf.expr("INTERVAL 10 MILLISECONDS"), # pd.Timedelta("10ms") + ... allow_exact_matches=False + ... ).sort_values(["time", "ticker", "price"]).reset_index(drop=True) + time ticker price quantity bid ask + 0 2016-05-25 13:30:00.023 MSFT 51.95 75 NaN NaN + 1 2016-05-25 13:30:00.038 MSFT 51.95 155 51.97 51.98 + 2 2016-05-25 13:30:00.048 AAPL 98.00 100 NaN NaN + 3 2016-05-25 13:30:00.048 GOOG 720.77 100 NaN NaN + 4 2016-05-25 13:30:00.048 GOOG 720.92 100 NaN NaN + """ + + def to_list(os: Optional[Union[Name, List[Name]]]) -> List[Label]: + if os is None: + return [] + elif is_name_like_tuple(os): + return [cast(Label, os)] + elif is_name_like_value(os): + return [(os,)] + else: + return [o if is_name_like_tuple(o) else (o,) for o in os] + + if isinstance(left, Series): + left = left.to_frame() + if isinstance(right, Series): + right = right.to_frame() + + if on: + if left_on or right_on: + raise ValueError( + 'Can only pass argument "on" OR "left_on" and "right_on", ' + "not a combination of both." + ) + left_as_of_names = list(map(left._internal.spark_column_name_for, to_list(on))) + right_as_of_names = list(map(right._internal.spark_column_name_for, to_list(on))) + else: + if left_index: + if isinstance(left.index, MultiIndex): + raise ValueError("left can only have one index") + left_as_of_names = left._internal.index_spark_column_names + else: + left_as_of_names = list(map(left._internal.spark_column_name_for, to_list(left_on))) + if right_index: + if isinstance(right.index, MultiIndex): + raise ValueError("right can only have one index") + right_as_of_names = right._internal.index_spark_column_names + else: + right_as_of_names = list(map(right._internal.spark_column_name_for, to_list(right_on))) + + if left_as_of_names and not right_as_of_names: + raise ValueError("Must pass right_on or right_index=True") + if right_as_of_names and not left_as_of_names: + raise ValueError("Must pass left_on or left_index=True") + if not left_as_of_names and not right_as_of_names: + common = list(left.columns.intersection(right.columns)) + if len(common) == 0: + raise ValueError( + "No common columns to perform merge on. Merge options: " + "left_on=None, right_on=None, left_index=False, right_index=False" + ) + left_as_of_names = list(map(left._internal.spark_column_name_for, to_list(common))) + right_as_of_names = list(map(right._internal.spark_column_name_for, to_list(common))) + + if len(left_as_of_names) != 1: + raise ValueError("can only asof on a key for left") + if len(right_as_of_names) != 1: + raise ValueError("can only asof on a key for right") + + if by: + if left_by or right_by: + raise ValueError('Can only pass argument "by" OR "left_by" and "right_by".') + left_join_on_names = list(map(left._internal.spark_column_name_for, to_list(by))) + right_join_on_names = list(map(right._internal.spark_column_name_for, to_list(by))) + else: + left_join_on_names = list(map(left._internal.spark_column_name_for, to_list(left_by))) + right_join_on_names = list(map(right._internal.spark_column_name_for, to_list(right_by))) + + if left_join_on_names and not right_join_on_names: + raise ValueError("missing right_by") + if right_join_on_names and not left_join_on_names: + raise ValueError("missing left_by") + if len(left_join_on_names) != len(right_join_on_names): + raise ValueError("left_by and right_by must be same length") + + # We should distinguish the name to avoid ambiguous column name after merging. + right_prefix = "__right_" + right_as_of_names = [right_prefix + right_as_of_name for right_as_of_name in right_as_of_names] + right_join_on_names = [ + right_prefix + right_join_on_name for right_join_on_name in right_join_on_names + ] + + left_as_of_name = left_as_of_names[0] + right_as_of_name = right_as_of_names[0] + + def resolve(internal: InternalFrame, side: str) -> InternalFrame: + def rename(col: str) -> str: + return "__{}_{}".format(side, col) + + internal = internal.resolved_copy + sdf = internal.spark_frame + sdf = sdf.select( + *[ + scol_for(sdf, col).alias(rename(col)) + for col in sdf.columns + if col not in HIDDEN_COLUMNS + ], + *HIDDEN_COLUMNS, + ) + return internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.index_spark_column_names + ], + index_fields=[field.copy(name=rename(field.name)) for field in internal.index_fields], + data_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.data_spark_column_names + ], + data_fields=[field.copy(name=rename(field.name)) for field in internal.data_fields], + ) + + left_internal = left._internal.resolved_copy + right_internal = resolve(right._internal, "right") + + left_table = left_internal.spark_frame.alias("left_table") + right_table = right_internal.spark_frame.alias("right_table") + + left_as_of_column = scol_for(left_table, left_as_of_name) + right_as_of_column = scol_for(right_table, right_as_of_name) + + if left_join_on_names: + left_join_on_columns = [scol_for(left_table, label) for label in left_join_on_names] + right_join_on_columns = [scol_for(right_table, label) for label in right_join_on_names] + on = reduce( + lambda lft, rgt: lft & rgt, + [lft == rgt for lft, rgt in zip(left_join_on_columns, right_join_on_columns)], + ) + else: + on = None + + Column = get_column_class() + if tolerance is not None and not isinstance(tolerance, Column): + tolerance = F.lit(tolerance) + + as_of_joined_table = left_table._joinAsOf( + right_table, + leftAsOfColumn=left_as_of_column, + rightAsOfColumn=right_as_of_column, + on=on, + how="left", + tolerance=tolerance, + allowExactMatches=allow_exact_matches, + direction=direction, + ) + + # Unpack suffixes tuple for convenience + left_suffix = suffixes[0] + right_suffix = suffixes[1] + + # Append suffixes to columns with the same name to avoid conflicts later + duplicate_columns = set(left_internal.column_labels) & set(right_internal.column_labels) + + exprs = [] + data_columns = [] + column_labels = [] + + def left_scol_for(label: Label) -> Column: # type: ignore[valid-type] + return scol_for(as_of_joined_table, left_internal.spark_column_name_for(label)) + + def right_scol_for(label: Label) -> Column: # type: ignore[valid-type] + return scol_for(as_of_joined_table, right_internal.spark_column_name_for(label)) + + for label in left_internal.column_labels: + col = left_internal.spark_column_name_for(label) + scol = left_scol_for(label) + if label in duplicate_columns: + spark_column_name = left_internal.spark_column_name_for(label) + if spark_column_name in (left_as_of_names + left_join_on_names) and ( + (right_prefix + spark_column_name) in (right_as_of_names + right_join_on_names) + ): + pass + else: + col = col + left_suffix + scol = scol.alias(col) # type: ignore[attr-defined] + label = tuple([str(label[0]) + left_suffix] + list(label[1:])) + exprs.append(scol) + data_columns.append(col) + column_labels.append(label) + for label in right_internal.column_labels: + # recover `right_prefix` here. + col = right_internal.spark_column_name_for(label)[len(right_prefix) :] + scol = right_scol_for(label).alias(col) # type: ignore[attr-defined] + if label in duplicate_columns: + spark_column_name = left_internal.spark_column_name_for(label) + if spark_column_name in left_as_of_names + left_join_on_names and ( + (right_prefix + spark_column_name) in right_as_of_names + right_join_on_names + ): + continue + else: + col = col + right_suffix + scol = scol.alias(col) # type: ignore[attr-defined] + label = tuple([str(label[0]) + right_suffix] + list(label[1:])) + exprs.append(scol) + data_columns.append(col) + column_labels.append(label) + + # Retain indices if they are used for joining + if left_index or right_index: + index_spark_column_names = [ + SPARK_INDEX_NAME_FORMAT(i) for i in range(len(left_internal.index_spark_column_names)) + ] + left_index_scols = [ + scol.alias(name) + for scol, name in zip(left_internal.index_spark_columns, index_spark_column_names) + ] + exprs.extend(left_index_scols) + index_names = left_internal.index_names + else: + index_spark_column_names = [] + index_names = [] + + selected_columns = as_of_joined_table.select(*exprs) + + internal = InternalFrame( + spark_frame=selected_columns, + index_spark_columns=[scol_for(selected_columns, col) for col in index_spark_column_names], + index_names=index_names, + column_labels=column_labels, + data_spark_columns=[scol_for(selected_columns, col) for col in data_columns], + ) + return DataFrame(internal) + + +@no_type_check +def to_numeric(arg, errors="raise"): + """ + Convert argument to a numeric type. + + Parameters + ---------- + arg : scalar, list, tuple, 1-d array, or Series + Argument to be converted. + errors : {'raise', 'coerce'}, default 'raise' + * If 'coerce', then invalid parsing will be set as NaN. + * If 'raise', then invalid parsing will raise an exception. + * If 'ignore', then invalid parsing will return the input. + + .. note:: 'ignore' doesn't work yet when `arg` is pandas-on-Spark Series. + + Returns + ------- + ret : numeric if parsing succeeded. + + See Also + -------- + DataFrame.astype : Cast argument to a specified dtype. + to_datetime : Convert argument to datetime. + to_timedelta : Convert argument to timedelta. + numpy.ndarray.astype : Cast a numpy array to a specified type. + + Examples + -------- + + >>> psser = ps.Series(['1.0', '2', '-3']) + >>> psser + 0 1.0 + 1 2 + 2 -3 + dtype: object + + >>> ps.to_numeric(psser) + 0 1.0 + 1 2.0 + 2 -3.0 + dtype: float32 + + If given Series contains invalid value to cast float, just cast it to `np.nan` + when `errors` is set to "coerce". + + >>> psser = ps.Series(['apple', '1.0', '2', '-3']) + >>> psser + 0 apple + 1 1.0 + 2 2 + 3 -3 + dtype: object + + >>> ps.to_numeric(psser, errors="coerce") + 0 NaN + 1 1.0 + 2 2.0 + 3 -3.0 + dtype: float32 + + Also support for list, tuple, np.array, or a scalar + + >>> ps.to_numeric(['1.0', '2', '-3']) + array([ 1., 2., -3.]) + + >>> ps.to_numeric(('1.0', '2', '-3')) + array([ 1., 2., -3.]) + + >>> ps.to_numeric(np.array(['1.0', '2', '-3'])) + array([ 1., 2., -3.]) + + >>> ps.to_numeric('1.0') + 1.0 + """ + if isinstance(arg, Series): + if errors == "coerce": + return arg._with_new_scol(arg.spark.column.cast("float")) + elif errors == "raise": + scol = arg.spark.column + scol_casted = scol.cast("float") + cond = F.when( + F.assert_true(scol.isNull() | scol_casted.isNotNull()).isNull(), scol_casted + ) + return arg._with_new_scol(cond) + elif errors == "ignore": + raise NotImplementedError("'ignore' is not implemented yet, when the `arg` is Series.") + else: + raise ValueError("invalid error value specified") + else: + return pd.to_numeric(arg, errors=errors) + + +def broadcast(obj: DataFrame) -> DataFrame: + """ + Marks a DataFrame as small enough for use in broadcast joins. + + .. deprecated:: 3.2.0 + Use :func:`DataFrame.spark.hint` instead. + + Parameters + ---------- + obj : DataFrame + + Returns + ------- + ret : DataFrame with broadcast hint. + + See Also + -------- + DataFrame.merge : Merge DataFrame objects with a database-style join. + DataFrame.join : Join columns of another DataFrame. + DataFrame.update : Modify in place using non-NA values from another DataFrame. + DataFrame.hint : Specifies some hint on the current DataFrame. + + Examples + -------- + >>> df1 = ps.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [1, 2, 3, 5]}, + ... columns=['lkey', 'value']).set_index('lkey') + >>> df2 = ps.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [5, 6, 7, 8]}, + ... columns=['rkey', 'value']).set_index('rkey') + >>> merged = df1.merge(ps.broadcast(df2), left_index=True, right_index=True) + >>> merged.spark.explain() # doctest: +ELLIPSIS + == Physical Plan == + ... + ...BroadcastHashJoin... + ... + """ + warnings.warn( + "`broadcast` has been deprecated and might be removed in a future version. " + "Use `DataFrame.spark.hint` with 'broadcast' for `name` parameter instead.", + FutureWarning, + ) + if not isinstance(obj, DataFrame): + raise TypeError("Invalid type : expected DataFrame got {}".format(type(obj).__name__)) + return DataFrame( + obj._internal.with_new_sdf(F.broadcast(obj._internal.resolved_copy.spark_frame)) + ) + + +def read_orc( + path: str, + columns: Optional[List[str]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: Any, +) -> "DataFrame": + """ + Load an ORC object from the file path, returning a DataFrame. + + Parameters + ---------- + path : str + The path string storing the ORC file to be read. + columns : list, default None + If not None, only these columns will be read from the file. + index_col : str or list of str, optional, default: None + Index column of table in Spark. + options : dict + All other options passed directly into Spark's data source. + + Returns + ------- + DataFrame + + Examples + -------- + >>> ps.range(1).to_orc('%s/read_spark_io/data.orc' % path) + >>> ps.read_orc('%s/read_spark_io/data.orc' % path, columns=['id']) + id + 0 0 + + You can preserve the index in the roundtrip as below. + + >>> ps.range(1).to_orc('%s/read_spark_io/data.orc' % path, index_col="index") + >>> ps.read_orc('%s/read_spark_io/data.orc' % path, columns=['id'], index_col="index") + ... # doctest: +NORMALIZE_WHITESPACE + id + index + 0 0 + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") + + psdf = read_spark_io(path, format="orc", index_col=index_col, **options) + + if columns is not None: + psdf_columns = psdf.columns + new_columns = list() + for column in list(columns): + if column in psdf_columns: + new_columns.append(column) + else: + raise ValueError("Unknown column name '{}'".format(column)) + psdf = psdf[new_columns] + + return psdf + + +def _get_index_map( + sdf: PySparkDataFrame, index_col: Optional[Union[str, List[str]]] = None +) -> Tuple[Optional[List[PySparkColumn]], Optional[List[Label]]]: + index_spark_columns: Optional[List[PySparkColumn]] + index_names: Optional[List[Label]] + if index_col is not None: + if isinstance(index_col, str): + index_col = [index_col] + sdf_columns = set(sdf.columns) + for col in index_col: + if col not in sdf_columns: + raise KeyError(col) + index_spark_columns = [scol_for(sdf, col) for col in index_col] + index_names = [(col,) for col in index_col] + else: + index_spark_columns = None + index_names = None + + return index_spark_columns, index_names + + +_get_dummies_default_accept_types = (DecimalType, StringType, DateType) +_get_dummies_acceptable_types = _get_dummies_default_accept_types + ( + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + BooleanType, + TimestampType, + TimestampNTZType, +) + + +def _test() -> None: + import os + import doctest + import shutil + import sys + import tempfile + import uuid + from pyspark.sql import SparkSession + import pyspark.pandas.namespace + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.namespace.__dict__.copy() + globs["ps"] = pyspark.pandas + globs["sf"] = F + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.namespace tests") + .getOrCreate() + ) + + db_name = "db%s" % str(uuid.uuid4()).replace("-", "") + spark.sql("CREATE DATABASE %s" % db_name) + globs["db"] = db_name + + path = tempfile.mkdtemp() + globs["path"] = path + + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.namespace, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + + shutil.rmtree(path, ignore_errors=True) + spark.sql("DROP DATABASE IF EXISTS %s CASCADE" % db_name) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/numpy_compat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/numpy_compat.py new file mode 100644 index 0000000..96f2ab1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/numpy_compat.py @@ -0,0 +1,260 @@ +# +# 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 Any, Callable, no_type_check + +import numpy as np +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.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] + "arcsin": F.asin, + "arcsinh": pandas_udf(lambda s: np.arcsinh(s), DoubleType()), # type: ignore[call-overload] + "arctan": F.atan, + "arctanh": pandas_udf(lambda s: np.arctanh(s), DoubleType()), # type: ignore[call-overload] + "bitwise_not": F.bitwiseNOT, + "cbrt": F.cbrt, + "ceil": F.ceil, + # It requires complex type which pandas-on-Spark does not support yet + "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] + "degrees": F.degrees, + "exp": F.exp, + "exp2": pandas_udf(lambda s: np.exp2(s), DoubleType()), # type: ignore[call-overload] + "expm1": F.expm1, + "fabs": pandas_udf(lambda s: np.fabs(s), DoubleType()), # type: ignore[call-overload] + "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"), + "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] + "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), + "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] + "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] + "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] +} + +binary_np_spark_mappings = { + "arctan2": F.atan2, + "bitwise_and": lambda c1, c2: c1.bitwiseAND(c2), + "bitwise_or": lambda c1, c2: c1.bitwiseOR(c2), + "bitwise_xor": lambda c1, c2: c1.bitwiseXOR(c2), + "copysign": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.copysign(s1, s2), DoubleType() + ), + "float_power": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.float_power(s1, s2), DoubleType() + ), + "floor_divide": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.floor_divide(s1, s2), DoubleType() + ), + "fmax": pandas_udf(lambda s1, s2: np.fmax(s1, s2), DoubleType()), # type: ignore[call-overload] + "fmin": pandas_udf(lambda s1, s2: np.fmin(s1, s2), DoubleType()), # type: ignore[call-overload] + "fmod": pandas_udf(lambda s1, s2: np.fmod(s1, s2), DoubleType()), # type: ignore[call-overload] + "gcd": pandas_udf(lambda s1, s2: np.gcd(s1, s2), DoubleType()), # type: ignore[call-overload] + "heaviside": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.heaviside(s1, s2), DoubleType() + ), + "hypot": F.hypot, + "lcm": pandas_udf(lambda s1, s2: np.lcm(s1, s2), DoubleType()), # type: ignore[call-overload] + "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() + ), + "logaddexp": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.logaddexp(s1, s2), DoubleType() + ), + "logaddexp2": pandas_udf( # type: ignore[call-overload] + lambda s1, s2: np.logaddexp2(s1, s2), DoubleType() + ), + "logical_and": lambda c1, c2: c1.cast(BooleanType()) & c2.cast(BooleanType()), + "logical_or": lambda c1, c2: c1.cast(BooleanType()) | c2.cast(BooleanType()), + "logical_xor": lambda c1, c2: ( + # mimics xor by logical operators. + (c1.cast(BooleanType()) | c2.cast(BooleanType())) + & (~(c1.cast(BooleanType())) | ~(c2.cast(BooleanType()))) + ), + "maximum": F.greatest, + "minimum": F.least, + "modf": pandas_udf(lambda s1, s2: np.modf(s1, s2), DoubleType()), # type: ignore[call-overload] + "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() + ), +} + + +# Copied from pandas. +# See also https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses +def maybe_dispatch_ufunc_to_dunder_op( + ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any +) -> IndexOpsMixin: + special = { + "add", + "sub", + "mul", + "pow", + "mod", + "floordiv", + "truediv", + "divmod", + "eq", + "ne", + "lt", + "gt", + "le", + "ge", + "remainder", + "matmul", + } + aliases = { + "absolute": "abs", + "multiply": "mul", + "floor_divide": "floordiv", + "true_divide": "truediv", + "power": "pow", + "remainder": "mod", + "divide": "truediv", + "equal": "eq", + "not_equal": "ne", + "less": "lt", + "less_equal": "le", + "greater": "gt", + "greater_equal": "ge", + } + + # For op(., Array) -> Array.__r{op}__ + flipped = { + "lt": "__gt__", + "le": "__ge__", + "gt": "__lt__", + "ge": "__le__", + "eq": "__eq__", + "ne": "__ne__", + } + + op_name = ufunc.__name__ + op_name = aliases.get(op_name, op_name) + + @no_type_check + def not_implemented(*args, **kwargs): + return NotImplemented + + if method == "__call__" and op_name in special and kwargs.get("out") is None: + if isinstance(inputs[0], type(ser_or_index)): + name = "__{}__".format(op_name) + return getattr(ser_or_index, name, not_implemented)(inputs[1]) + else: + name = flipped.get(op_name, "__r{}__".format(op_name)) + return getattr(ser_or_index, name, not_implemented)(inputs[0]) + else: + return NotImplemented + + +# See also https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses +def maybe_dispatch_ufunc_to_spark_func( + ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any +) -> IndexOpsMixin: + from pyspark.pandas.base import column_op + + op_name = ufunc.__name__ + + if ( + method == "__call__" + and (op_name in unary_np_spark_mappings or op_name in binary_np_spark_mappings) + and kwargs.get("out") is None + ): + + np_spark_map_func = unary_np_spark_mappings.get(op_name) or binary_np_spark_mappings.get( + op_name + ) + + @no_type_check + def convert_arguments(*args): + args = [F.lit(inp) for inp in args] + return np_spark_map_func(*args) + + return column_op(convert_arguments)(*inputs) + else: + return NotImplemented + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.numpy_compat + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.numpy_compat.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.numpy_compat tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.numpy_compat, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/__init__.py new file mode 100644 index 0000000..d00e002 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/__init__.py @@ -0,0 +1,17 @@ +# +# 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.pandas.plot.core import * # noqa: F401,F403 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/core.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/core.py new file mode 100644 index 0000000..ccae96a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/core.py @@ -0,0 +1,1213 @@ +# +# 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 importlib + +import pandas as pd +import numpy as np +from pyspark.ml.feature import Bucketizer +from pyspark.mllib.stat import KernelDensity +from pyspark.sql import functions as F +from pandas.core.base import PandasObject +from pandas.core.dtypes.inference import is_integer + +from pyspark.pandas.missing import unsupported_function +from pyspark.pandas.config import get_option +from pyspark.pandas.utils import name_like_string + + +class TopNPlotBase: + def get_top_n(self, data): + from pyspark.pandas import DataFrame, Series + + max_rows = get_option("plotting.max_rows") + # Simply use the first 1k elements and make it into a pandas dataframe + # For categorical variables, it is likely called from df.x.value_counts().plot.xxx(). + if isinstance(data, (Series, DataFrame)): + data = data.head(max_rows + 1)._to_pandas() + else: + raise TypeError("Only DataFrame and Series are supported for plotting.") + + self.partial = False + if len(data) > max_rows: + self.partial = True + data = data.iloc[:max_rows] + return data + + def set_result_text(self, ax): + max_rows = get_option("plotting.max_rows") + assert hasattr(self, "partial") + + if self.partial: + ax.text( + 1, + 1, + "showing top {} elements only".format(max_rows), + size=6, + ha="right", + va="bottom", + transform=ax.transAxes, + ) + + +class SampledPlotBase: + def get_sampled(self, data): + from pyspark.pandas import DataFrame, Series + + fraction = get_option("plotting.sample_ratio") + if fraction is None: + fraction = 1 / (len(data) / get_option("plotting.max_rows")) + fraction = min(1.0, fraction) + self.fraction = fraction + + if isinstance(data, (DataFrame, Series)): + if isinstance(data, Series): + data = data.to_frame() + sampled = data._internal.resolved_copy.spark_frame.sample(fraction=self.fraction) + return DataFrame(data._internal.with_new_sdf(sampled))._to_pandas() + else: + raise TypeError("Only DataFrame and Series are supported for plotting.") + + def set_result_text(self, ax): + assert hasattr(self, "fraction") + + if self.fraction < 1: + ax.text( + 1, + 1, + "showing the sampled result by fraction %s" % self.fraction, + size=6, + ha="right", + va="bottom", + transform=ax.transAxes, + ) + + +class NumericPlotBase: + @staticmethod + def prepare_numeric_data(data): + from pyspark.pandas.series import Series + + if isinstance(data, Series): + data = data.to_frame() + + numeric_data = data.select_dtypes( + include=["byte", "decimal", "integer", "float", "long", "double", np.datetime64] + ) + + # no empty frames or series allowed + if len(numeric_data.columns) == 0: + raise TypeError( + "Empty {0!r}: no numeric data to " "plot".format(numeric_data.__class__.__name__) + ) + + return data, numeric_data + + +class HistogramPlotBase(NumericPlotBase): + @staticmethod + def prepare_hist_data(data, bins): + data, numeric_data = NumericPlotBase.prepare_numeric_data(data) + if is_integer(bins): + # computes boundaries for the column + bins = HistogramPlotBase.get_bins(data._to_spark(), bins) + + return numeric_data, bins + + @staticmethod + def get_bins(sdf, bins): + # 'data' is a Spark DataFrame that selects all columns. + if len(sdf.columns) > 1: + min_col = F.least(*map(F.min, sdf)) + max_col = F.greatest(*map(F.max, sdf)) + else: + min_col = F.min(sdf.columns[-1]) + max_col = F.max(sdf.columns[-1]) + boundaries = sdf.select(min_col, max_col).first() + + # divides the boundaries into bins + if boundaries[0] == boundaries[1]: + boundaries = (boundaries[0] - 0.5, boundaries[1] + 0.5) + + return np.linspace(boundaries[0], boundaries[1], bins + 1) + + @staticmethod + def compute_hist(psdf, bins): + # 'data' is a Spark DataFrame that selects one column. + assert isinstance(bins, (np.ndarray, np.generic)) + + sdf = psdf._internal.spark_frame + scols = [] + input_column_names = [] + for label in psdf._internal.column_labels: + input_column_name = name_like_string(label) + input_column_names.append(input_column_name) + scols.append(psdf._internal.spark_column_for(label).alias(input_column_name)) + sdf = sdf.select(*scols) + + # 1. Make the bucket output flat to: + # +----------+-------+ + # |__group_id|buckets| + # +----------+-------+ + # |0 |0.0 | + # |0 |0.0 | + # |0 |1.0 | + # |0 |2.0 | + # |0 |3.0 | + # |0 |3.0 | + # |1 |0.0 | + # |1 |1.0 | + # |1 |1.0 | + # |1 |2.0 | + # |1 |1.0 | + # |1 |0.0 | + # +----------+-------+ + colnames = sdf.columns + bucket_names = ["__{}_bucket".format(colname) for colname in colnames] + + output_df = None + for group_id, (colname, bucket_name) in enumerate(zip(colnames, bucket_names)): + # creates a Bucketizer to get corresponding bin of each value + bucketizer = Bucketizer( + splits=bins, inputCol=colname, outputCol=bucket_name, handleInvalid="skip" + ) + + bucket_df = bucketizer.transform(sdf) + + if output_df is None: + output_df = bucket_df.select( + F.lit(group_id).alias("__group_id"), F.col(bucket_name).alias("__bucket") + ) + else: + output_df = output_df.union( + bucket_df.select( + F.lit(group_id).alias("__group_id"), F.col(bucket_name).alias("__bucket") + ) + ) + + # 2. Calculate the count based on each group and bucket. + # +----------+-------+------+ + # |__group_id|buckets| count| + # +----------+-------+------+ + # |0 |0.0 |2 | + # |0 |1.0 |1 | + # |0 |2.0 |1 | + # |0 |3.0 |2 | + # |1 |0.0 |2 | + # |1 |1.0 |3 | + # |1 |2.0 |1 | + # +----------+-------+------+ + result = ( + output_df.groupby("__group_id", "__bucket") + .agg(F.count("*").alias("count")) + .toPandas() + .sort_values(by=["__group_id", "__bucket"]) + ) + + # 3. Fill empty bins and calculate based on each group id. From: + # +----------+--------+------+ + # |__group_id|__bucket| count| + # +----------+--------+------+ + # |0 |0.0 |2 | + # |0 |1.0 |1 | + # |0 |2.0 |1 | + # |0 |3.0 |2 | + # +----------+--------+------+ + # +----------+--------+------+ + # |__group_id|__bucket| count| + # +----------+--------+------+ + # |1 |0.0 |2 | + # |1 |1.0 |3 | + # |1 |2.0 |1 | + # +----------+--------+------+ + # + # to: + # +-----------------+ + # |__values1__bucket| + # +-----------------+ + # |2 | + # |1 | + # |1 | + # |2 | + # |0 | + # +-----------------+ + # +-----------------+ + # |__values2__bucket| + # +-----------------+ + # |2 | + # |3 | + # |1 | + # |0 | + # |0 | + # +-----------------+ + output_series = [] + for i, (input_column_name, bucket_name) in enumerate(zip(input_column_names, bucket_names)): + current_bucket_result = result[result["__group_id"] == i] + # generates a pandas DF with one row for each bin + # we need this as some of the bins may be empty + indexes = pd.DataFrame({"__bucket": np.arange(0, len(bins) - 1)}) + # merges the bins with counts on it and fills remaining ones with zeros + pdf = indexes.merge(current_bucket_result, how="left", on=["__bucket"]).fillna(0)[ + ["count"] + ] + pdf.columns = [input_column_name] + output_series.append(pdf[input_column_name]) + + return output_series + + +class BoxPlotBase: + @staticmethod + def compute_multicol_stats(data, colnames, whis, precision): + # Computes mean, median, Q1 and Q3 with approx_percentile and precision + scol = [] + for colname in colnames: + scol.append( + F.percentile_approx( + "`%s`" % colname, [0.25, 0.50, 0.75], int(1.0 / precision) + ).alias("{}_percentiles%".format(colname)) + ) + scol.append(F.mean("`%s`" % colname).alias("{}_mean".format(colname))) + + # a_percentiles a_mean b_percentiles b_mean + # 0 [3.0, 3.2, 3.2] 3.18 [5.1, 5.9, 6.4] 5.86 + pdf = data._internal.resolved_copy.spark_frame.select(*scol).toPandas() + + i = 0 + multicol_stats = {} + for colname in colnames: + q1, med, q3 = pdf.iloc[0, i] + iqr = q3 - q1 + lfence = q1 - whis * iqr + ufence = q3 + whis * iqr + i += 1 + + mean = pdf.iloc[0, i] + i += 1 + + multicol_stats[colname] = { + "mean": mean, + "med": med, + "q1": q1, + "q3": q3, + "lfence": lfence, + "ufence": ufence, + } + + return multicol_stats + + @staticmethod + def compute_stats(data, colname, whis, precision): + # Computes mean, median, Q1 and Q3 with approx_percentile and precision + pdf = data._psdf._internal.resolved_copy.spark_frame.agg( + *[ + F.expr( + "approx_percentile(`{}`, {}, {})".format(colname, q, int(1.0 / precision)) + ).alias("{}_{}%".format(colname, int(q * 100))) + for q in [0.25, 0.50, 0.75] + ], + F.mean("`%s`" % colname).alias("{}_mean".format(colname)), + ).toPandas() + + # Computes IQR and Tukey's fences + iqr = "{}_iqr".format(colname) + p75 = "{}_75%".format(colname) + p25 = "{}_25%".format(colname) + pdf.loc[:, iqr] = pdf.loc[:, p75] - pdf.loc[:, p25] + pdf.loc[:, "{}_lfence".format(colname)] = pdf.loc[:, p25] - whis * pdf.loc[:, iqr] + pdf.loc[:, "{}_ufence".format(colname)] = pdf.loc[:, p75] + whis * pdf.loc[:, iqr] + + qnames = ["25%", "50%", "75%", "mean", "lfence", "ufence"] + col_summ = pdf[["{}_{}".format(colname, q) for q in qnames]] + col_summ.columns = qnames + lfence, ufence = col_summ["lfence"], col_summ["ufence"] + + stats = { + "mean": col_summ["mean"].values[0], + "med": col_summ["50%"].values[0], + "q1": col_summ["25%"].values[0], + "q3": col_summ["75%"].values[0], + } + + return stats, (lfence.values[0], ufence.values[0]) + + @staticmethod + def multicol_outliers(data, multicol_stats): + scols = {} + for colname, stats in multicol_stats.items(): + scols["__{}_outlier".format(colname)] = ~F.col("`%s`" % colname).between( + stats["lfence"], stats["ufence"] + ) + return data._internal.resolved_copy.spark_frame.withColumns(scols) + + @staticmethod + def outliers(data, colname, lfence, ufence): + # Builds expression to identify outliers + expression = F.col("`%s`" % colname).between(lfence, ufence) + # Creates a column to flag rows as outliers or not + return data._psdf._internal.resolved_copy.spark_frame.withColumn( + "__{}_outlier".format(colname), ~expression + ) + + @staticmethod + def calc_multicol_whiskers(colnames, multicol_outliers): + # Computes min and max values of non-outliers - the whiskers + scols = [] + for colname in colnames: + outlier_colname = "__{}_outlier".format(colname) + scols.append( + F.min(F.when(~F.col(outlier_colname), F.col(colname)).otherwise(F.lit(None))).alias( + "__{}_min".format(colname) + ) + ) + scols.append( + F.max(F.when(~F.col(outlier_colname), F.col(colname)).otherwise(F.lit(None))).alias( + "__{}_max".format(colname) + ) + ) + + pdf = multicol_outliers.select(*scols).toPandas() + + i = 0 + whiskers = {} + for colname in colnames: + min = pdf.iloc[0, i] + i += 1 + max = pdf.iloc[0, i] + i += 1 + whiskers[colname] = { + "min": min, + "max": max, + } + + return whiskers + + @staticmethod + def calc_whiskers(colname, outliers): + # Computes min and max values of non-outliers - the whiskers + minmax = ( + outliers.filter("not `__{}_outlier`".format(colname)) + .agg(F.min("`%s`" % colname).alias("min"), F.max(colname).alias("max")) + .toPandas() + ) + return minmax.iloc[0][["min", "max"]].values + + @staticmethod + def get_fliers(colname, outliers, min_val): + # Filters only the outliers, should "showfliers" be True + fliers_df = outliers.filter("`__{}_outlier`".format(colname)) + + # If it shows fliers, take the top 1k with highest absolute values + # Here we normalize the values by subtracting the minimum value from + # each, and use absolute values. + order_col = F.abs(F.col("`{}`".format(colname)) - min_val.item()) + fliers = ( + fliers_df.select(F.col("`{}`".format(colname))) + .orderBy(order_col) + .limit(1001) + .toPandas()[colname] + .values + ) + + return fliers + + +class KdePlotBase(NumericPlotBase): + @staticmethod + def prepare_kde_data(data): + _, numeric_data = NumericPlotBase.prepare_numeric_data(data) + return numeric_data + + @staticmethod + def get_ind(sdf, ind): + def calc_min_max(): + if len(sdf.columns) > 1: + min_col = F.least(*map(F.min, sdf)) + max_col = F.greatest(*map(F.max, sdf)) + else: + min_col = F.min(sdf.columns[-1]) + max_col = F.max(sdf.columns[-1]) + return sdf.select(min_col, max_col).first() + + if ind is None: + min_val, max_val = calc_min_max() + sample_range = max_val - min_val + ind = np.linspace( + min_val - 0.5 * sample_range, + max_val + 0.5 * sample_range, + 1000, + ) + elif is_integer(ind): + min_val, max_val = calc_min_max() + sample_range = max_val - min_val + ind = np.linspace( + min_val - 0.5 * sample_range, + max_val + 0.5 * sample_range, + ind, + ) + return ind + + @staticmethod + def compute_kde(sdf, bw_method=None, ind=None): + # 'sdf' is a Spark DataFrame that selects one column. + + # Using RDD is slow so we might have to change it to Dataset based implementation + # once Spark has that implementation. + sample = sdf.rdd.map(lambda x: float(x[0])) + kd = KernelDensity() + kd.setSample(sample) + + assert isinstance(bw_method, (int, float)), "'bw_method' must be set as a scalar number." + + if bw_method is not None: + # Match the bandwidth with Spark. + kd.setBandwidth(float(bw_method)) + return kd.estimate(list(map(float, ind))) + + +class PandasOnSparkPlotAccessor(PandasObject): + """ + Series/Frames plotting accessor and method. + + Uses the backend specified by the + option ``plotting.backend``. By default, plotly is used. + + Plotting methods can also be accessed by calling the accessor as a method + with the ``kind`` argument: + ``s.plot(kind='hist')`` is equivalent to ``s.plot.hist()`` + """ + + pandas_plot_data_map = { + "pie": TopNPlotBase().get_top_n, + "bar": TopNPlotBase().get_top_n, + "barh": TopNPlotBase().get_top_n, + "scatter": TopNPlotBase().get_top_n, + "area": SampledPlotBase().get_sampled, + "line": SampledPlotBase().get_sampled, + } + _backends = {} # type: ignore[var-annotated] + + def __init__(self, data): + self.data = data + + @staticmethod + def _find_backend(backend): + """ + Find a pandas-on-Spark plotting backend + """ + try: + return PandasOnSparkPlotAccessor._backends[backend] + except KeyError: + try: + module = importlib.import_module(backend) + except ImportError: + # We re-raise later on. + pass + else: + if hasattr(module, "plot") or hasattr(module, "plot_pandas_on_spark"): + # Validate that the interface is implemented when the option + # is set, rather than at plot time. + PandasOnSparkPlotAccessor._backends[backend] = module + return module + + raise ValueError( + "Could not find plotting backend '{backend}'. Ensure that you've installed " + "the package providing the '{backend}' entrypoint, or that the package has a " + "top-level `.plot` method.".format(backend=backend) + ) + + @staticmethod + def _get_plot_backend(backend=None): + backend = backend or get_option("plotting.backend") + # Shortcut + if backend in PandasOnSparkPlotAccessor._backends: + return PandasOnSparkPlotAccessor._backends[backend] + + if backend == "matplotlib": + # Because matplotlib is an optional dependency, + # we need to attempt an import here to raise an ImportError if needed. + try: + # test if matplotlib can be imported + import matplotlib # noqa: F401 + from pyspark.pandas.plot import matplotlib as module + except ImportError: + raise ImportError( + "matplotlib is required for plotting when the " + "default backend 'matplotlib' is selected." + ) from None + + PandasOnSparkPlotAccessor._backends["matplotlib"] = module + elif backend == "plotly": + try: + # test if plotly can be imported + import plotly # noqa: F401 + from pyspark.pandas.plot import plotly as module + except ImportError: + raise ImportError( + "plotly is required for plotting when the " + "default backend 'plotly' is selected." + ) from None + + PandasOnSparkPlotAccessor._backends["plotly"] = module + else: + module = PandasOnSparkPlotAccessor._find_backend(backend) + PandasOnSparkPlotAccessor._backends[backend] = module + return module + + def __call__(self, kind="line", backend=None, **kwargs): + plot_backend = PandasOnSparkPlotAccessor._get_plot_backend(backend) + plot_data = self.data + + kind = {"density": "kde"}.get(kind, kind) + if hasattr(plot_backend, "plot_pandas_on_spark"): + # use if there's pandas-on-Spark specific method. + return plot_backend.plot_pandas_on_spark(plot_data, kind=kind, **kwargs) + else: + # fallback to use pandas' + if not PandasOnSparkPlotAccessor.pandas_plot_data_map[kind]: + raise NotImplementedError( + "'%s' plot is not supported with '%s' plot " + "backend yet." % (kind, plot_backend.__name__) + ) + plot_data = PandasOnSparkPlotAccessor.pandas_plot_data_map[kind](plot_data) + return plot_backend.plot(plot_data, kind=kind, **kwargs) + + def line(self, x=None, y=None, **kwargs): + """ + Plot DataFrame/Series as lines. + + This function is useful to plot lines using Series's values + as coordinates. + + Parameters + ---------- + x : int or str, optional + Columns to use for the horizontal axis. + Either the location or the label of the columns to be used. + By default, it will use the DataFrame indices. + y : int, str, or list of them, optional + The values to be plotted. + Either the location or the label of the columns to be used. + By default, it will use the remaining DataFrame numeric columns. + **kwds + Keyword arguments to pass on to :meth:`Series.plot` or :meth:`DataFrame.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + See Also + -------- + plotly.express.line : Plot y versus x as lines and/or markers (plotly). + matplotlib.pyplot.plot : Plot y versus x as lines and/or markers (matplotlib). + + Examples + -------- + Basic plot. + + For Series: + + .. plotly:: + + >>> s = ps.Series([1, 3, 2]) + >>> s.plot.line() # doctest: +SKIP + + For DataFrame: + + .. plotly:: + + The following example shows the populations for some animals + over the years. + + >>> df = ps.DataFrame({'pig': [20, 18, 489, 675, 1776], + ... 'horse': [4, 25, 281, 600, 1900]}, + ... index=[1990, 1997, 2003, 2009, 2014]) + >>> df.plot.line() # doctest: +SKIP + + .. plotly:: + + The following example shows the relationship between both + populations. + + >>> df = ps.DataFrame({'pig': [20, 18, 489, 675, 1776], + ... 'horse': [4, 25, 281, 600, 1900]}, + ... index=[1990, 1997, 2003, 2009, 2014]) + >>> df.plot.line(x='pig', y='horse') # doctest: +SKIP + """ + return self(kind="line", x=x, y=y, **kwargs) + + def bar(self, x=None, y=None, **kwds): + """ + Vertical bar plot. + + Parameters + ---------- + x : label or position, optional + Allows plotting of one column versus another. + If not specified, the index of the DataFrame is used. + y : label or position, optional + Allows plotting of one column versus another. + If not specified, all numerical columns are used. + **kwds : optional + Additional keyword arguments are documented in + :meth:`pyspark.pandas.Series.plot` or + :meth:`pyspark.pandas.DataFrame.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Examples + -------- + Basic plot. + + For Series: + + .. plotly:: + + >>> s = ps.Series([1, 3, 2]) + >>> s.plot.bar() # doctest: +SKIP + + For DataFrame: + + .. plotly:: + + >>> df = ps.DataFrame({'lab': ['A', 'B', 'C'], 'val': [10, 30, 20]}) + >>> df.plot.bar(x='lab', y='val') # doctest: +SKIP + + Plot a whole dataframe to a bar plot. Each column is stacked with a + distinct color along the horizontal axis. + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.bar() # doctest: +SKIP + + Instead of stacking, the figure can be split by column with plotly + APIs. + + .. plotly:: + + >>> from plotly.subplots import make_subplots + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> fig = (make_subplots(rows=2, cols=1) + ... .add_trace(df.plot.bar(y='speed').data[0], row=1, col=1) + ... .add_trace(df.plot.bar(y='speed').data[0], row=1, col=1) + ... .add_trace(df.plot.bar(y='lifespan').data[0], row=2, col=1)) + >>> fig # doctest: +SKIP + + Plot a single column. + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.bar(y='speed') # doctest: +SKIP + + Plot only selected categories for the DataFrame. + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.bar(x='lifespan') # doctest: +SKIP + """ + from pyspark.pandas import DataFrame, Series + + if isinstance(self.data, Series): + return self(kind="bar", **kwds) + elif isinstance(self.data, DataFrame): + return self(kind="bar", x=x, y=y, **kwds) + + def barh(self, x=None, y=None, **kwargs): + """ + Make a horizontal bar plot. + + A horizontal bar plot is a plot that presents quantitative data with + rectangular bars with lengths proportional to the values that they + represent. A bar plot shows comparisons among discrete categories. One + axis of the plot shows the specific categories being compared, and the + other axis represents a measured value. + + Parameters + ---------- + x : label or position, default DataFrame.index + Column to be used for categories. + y : label or position, default All numeric columns in dataframe + Columns to be plotted from the DataFrame. + **kwds + Keyword arguments to pass on to + :meth:`pyspark.pandas.DataFrame.plot` or :meth:`pyspark.pandas.Series.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + See Also + -------- + plotly.express.bar : Plot a vertical bar plot using plotly. + matplotlib.axes.Axes.bar : Plot a vertical bar plot using matplotlib. + + Examples + -------- + For Series: + + .. plotly:: + + >>> df = ps.DataFrame({'lab': ['A', 'B', 'C'], 'val': [10, 30, 20]}) + >>> df.val.plot.barh() # doctest: +SKIP + + For DataFrame: + + .. plotly:: + + >>> df = ps.DataFrame({'lab': ['A', 'B', 'C'], 'val': [10, 30, 20]}) + >>> df.plot.barh(x='lab', y='val') # doctest: +SKIP + + Plot a whole DataFrame to a horizontal bar plot + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.barh() # doctest: +SKIP + + Plot a column of the DataFrame to a horizontal bar plot + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.barh(y='speed') # doctest: +SKIP + + Plot DataFrame versus the desired column + + .. plotly:: + + >>> speed = [0.1, 17.5, 40, 48, 52, 69, 88] + >>> lifespan = [2, 8, 70, 1.5, 25, 12, 28] + >>> index = ['snail', 'pig', 'elephant', + ... 'rabbit', 'giraffe', 'coyote', 'horse'] + >>> df = ps.DataFrame({'speed': speed, + ... 'lifespan': lifespan}, index=index) + >>> df.plot.barh(x='lifespan') # doctest: +SKIP + """ + from pyspark.pandas import DataFrame, Series + + if isinstance(self.data, Series): + return self(kind="barh", **kwargs) + elif isinstance(self.data, DataFrame): + return self(kind="barh", x=x, y=y, **kwargs) + + def box(self, **kwds): + """ + Make a box plot of the Series columns. + + Parameters + ---------- + **kwds : optional + Additional keyword arguments are documented in + :meth:`pyspark.pandas.Series.plot`. + + precision: scalar, default = 0.01 + This argument is used by pandas-on-Spark to compute approximate statistics + for building a boxplot. Use *smaller* values to get more precise + statistics (matplotlib-only). + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Notes + ----- + There are behavior differences between pandas-on-Spark and pandas. + + * pandas-on-Spark computes approximate statistics - expect differences between + pandas and pandas-on-Spark boxplots, especially regarding 1st and 3rd quartiles. + * The `whis` argument is only supported as a single number. + * pandas-on-Spark doesn't support the following argument(s) (matplotlib-only). + + * `bootstrap` argument is not supported + * `autorange` argument is not supported + + Examples + -------- + Draw a box plot from a DataFrame with four columns of randomly + generated data. + + For Series: + + .. plotly:: + + >>> data = np.random.randn(25, 4) + >>> df = ps.DataFrame(data, columns=list('ABCD')) + >>> df['A'].plot.box() # doctest: +SKIP + + This is an unsupported function for DataFrame type + """ + from pyspark.pandas import DataFrame, Series + + if isinstance(self.data, (Series, DataFrame)): + return self(kind="box", **kwds) + + def hist(self, bins=10, **kwds): + """ + Draw one histogram of the DataFrame’s columns. + A `histogram`_ is a representation of the distribution of data. + This function calls :meth:`plotting.backend.plot`, + on each series in the DataFrame, resulting in one histogram per column. + + .. _histogram: https://en.wikipedia.org/wiki/Histogram + + Parameters + ---------- + bins : integer or sequence, default 10 + Number of histogram bins to be used. If an integer is given, bins + 1 + bin edges are calculated and returned. If bins is a sequence, it gives + bin edges, including left edge of first bin and right edge of last + bin. In this case, bins are returned unmodified. + **kwds + All other plotting keyword arguments to be passed to + plotting backend. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Examples + -------- + Basic plot. + + For Series: + + .. plotly:: + + >>> s = ps.Series([1, 3, 2]) + >>> s.plot.hist() # doctest: +SKIP + + For DataFrame: + + .. plotly:: + + >>> df = pd.DataFrame( + ... np.random.randint(1, 7, 6000), + ... columns=['one']) + >>> df['two'] = df['one'] + np.random.randint(1, 7, 6000) + >>> df = ps.from_pandas(df) + >>> df.plot.hist(bins=12, alpha=0.5) # doctest: +SKIP + """ + return self(kind="hist", bins=bins, **kwds) + + def kde(self, bw_method=None, ind=None, **kwargs): + """ + Generate Kernel Density Estimate plot using Gaussian kernels. + + Parameters + ---------- + bw_method : scalar + The method used to calculate the estimator bandwidth. + See KernelDensity in PySpark for more information. + ind : NumPy array or integer, optional + Evaluation points for the estimated PDF. If None (default), + 1000 equally spaced points are used. If `ind` is a NumPy array, the + KDE is evaluated at the points passed. If `ind` is an integer, + `ind` number of equally spaced points are used. + **kwargs : optional + Keyword arguments to pass on to :meth:`pandas-on-Spark.Series.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Examples + -------- + A scalar bandwidth should be specified. Using a small bandwidth value can + lead to over-fitting, while using a large bandwidth value may result + in under-fitting: + + .. plotly:: + + >>> s = ps.Series([1, 2, 2.5, 3, 3.5, 4, 5]) + >>> s.plot.kde(bw_method=0.3) # doctest: +SKIP + + .. plotly:: + + >>> s = ps.Series([1, 2, 2.5, 3, 3.5, 4, 5]) + >>> s.plot.kde(bw_method=3) # doctest: +SKIP + + The `ind` parameter determines the evaluation points for the + plot of the estimated KDF: + + .. plotly:: + + >>> s = ps.Series([1, 2, 2.5, 3, 3.5, 4, 5]) + >>> s.plot.kde(ind=[1, 2, 3, 4, 5], bw_method=0.3) # doctest: +SKIP + + For DataFrame, it works in the same way as Series: + + .. plotly:: + + >>> df = ps.DataFrame({ + ... 'x': [1, 2, 2.5, 3, 3.5, 4, 5], + ... 'y': [4, 4, 4.5, 5, 5.5, 6, 6], + ... }) + >>> df.plot.kde(bw_method=0.3) # doctest: +SKIP + + .. plotly:: + + >>> df = ps.DataFrame({ + ... 'x': [1, 2, 2.5, 3, 3.5, 4, 5], + ... 'y': [4, 4, 4.5, 5, 5.5, 6, 6], + ... }) + >>> df.plot.kde(bw_method=3) # doctest: +SKIP + + .. plotly:: + + >>> df = ps.DataFrame({ + ... 'x': [1, 2, 2.5, 3, 3.5, 4, 5], + ... 'y': [4, 4, 4.5, 5, 5.5, 6, 6], + ... }) + >>> df.plot.kde(ind=[1, 2, 3, 4, 5, 6], bw_method=0.3) # doctest: +SKIP + """ + return self(kind="kde", bw_method=bw_method, ind=ind, **kwargs) + + density = kde + + def area(self, x=None, y=None, **kwds): + """ + Draw a stacked area plot. + + An area plot displays quantitative data visually. + This function wraps the plotly area function. + + Parameters + ---------- + x : label or position, optional + Coordinates for the X axis. By default it uses the index. + y : label or position, optional + Column to plot. By default it uses all columns. + stacked : bool, default True + Area plots are stacked by default. Set to False to create an + unstacked plot (matplotlib-only). + **kwds : optional + Additional keyword arguments are documented in + :meth:`DataFrame.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Examples + -------- + + For Series + + .. plotly:: + + >>> df = ps.DataFrame({ + ... 'sales': [3, 2, 3, 9, 10, 6], + ... 'signups': [5, 5, 6, 12, 14, 13], + ... 'visits': [20, 42, 28, 62, 81, 50], + ... }, index=pd.date_range(start='2018/01/01', end='2018/07/01', + ... freq='M')) + >>> df.sales.plot.area() # doctest: +SKIP + + For DataFrame + + .. plotly:: + + >>> df = ps.DataFrame({ + ... 'sales': [3, 2, 3, 9, 10, 6], + ... 'signups': [5, 5, 6, 12, 14, 13], + ... 'visits': [20, 42, 28, 62, 81, 50], + ... }, index=pd.date_range(start='2018/01/01', end='2018/07/01', + ... freq='M')) + >>> df.plot.area() # doctest: +SKIP + """ + from pyspark.pandas import DataFrame, Series + + if isinstance(self.data, Series): + return self(kind="area", **kwds) + elif isinstance(self.data, DataFrame): + return self(kind="area", x=x, y=y, **kwds) + + def pie(self, **kwds): + """ + Generate a pie plot. + + A pie plot is a proportional representation of the numerical data in a + column. This function wraps :meth:`plotly.express.pie` for the + specified column. + + Parameters + ---------- + y : int or label, optional + Label or position of the column to plot. + If not provided, ``subplots=True`` argument must be passed (matplotlib-only). + **kwds + Keyword arguments to pass on to :meth:`pandas-on-Spark.Series.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + Examples + -------- + + For Series: + + .. plotly:: + + >>> df = ps.DataFrame({'mass': [0.330, 4.87, 5.97], + ... 'radius': [2439.7, 6051.8, 6378.1]}, + ... index=['Mercury', 'Venus', 'Earth']) + >>> df.mass.plot.pie() # doctest: +SKIP + + + For DataFrame: + + .. plotly:: + + >>> df = ps.DataFrame({'mass': [0.330, 4.87, 5.97], + ... 'radius': [2439.7, 6051.8, 6378.1]}, + ... index=['Mercury', 'Venus', 'Earth']) + >>> df.plot.pie(y='mass') # doctest: +SKIP + """ + from pyspark.pandas import DataFrame, Series + + if isinstance(self.data, Series): + return self(kind="pie", **kwds) + else: + # pandas will raise an error if y is None and subplots if not True + if ( + isinstance(self.data, DataFrame) + and kwds.get("y", None) is None + and not kwds.get("subplots", False) + ): + raise ValueError( + "pie requires either y column or 'subplots=True' (matplotlib-only)" + ) + return self(kind="pie", **kwds) + + def scatter(self, x, y, **kwds): + """ + Create a scatter plot with varying marker point size and color. + + The coordinates of each point are defined by two dataframe columns and + filled circles are used to represent each point. This kind of plot is + useful to see complex correlations between two variables. Points could + be for instance natural 2D coordinates like longitude and latitude in + a map or, in general, any pair of metrics that can be plotted against + each other. + + Parameters + ---------- + x : int or str + The column name or column position to be used as horizontal + coordinates for each point. + y : int or str + The column name or column position to be used as vertical + coordinates for each point. + s : scalar or array_like, optional + (matplotlib-only). + c : str, int or array_like, optional + (matplotlib-only). + + **kwds: Optional + Keyword arguments to pass on to :meth:`pyspark.pandas.DataFrame.plot`. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + Return an custom object when ``backend!=plotly``. + Return an ndarray when ``subplots=True`` (matplotlib-only). + + See Also + -------- + plotly.express.scatter : Scatter plot using multiple input data + formats (plotly). + matplotlib.pyplot.scatter : Scatter plot using multiple input data + formats (matplotlib). + + Examples + -------- + Let's see how to draw a scatter plot using coordinates from the values + in a DataFrame's columns. + + .. plotly:: + + >>> df = ps.DataFrame([[5.1, 3.5, 0], [4.9, 3.0, 0], [7.0, 3.2, 1], + ... [6.4, 3.2, 1], [5.9, 3.0, 2]], + ... columns=['length', 'width', 'species']) + >>> df.plot.scatter(x='length', y='width') # doctest: +SKIP + + And now with dark scheme: + + .. plotly:: + + >>> df = ps.DataFrame([[5.1, 3.5, 0], [4.9, 3.0, 0], [7.0, 3.2, 1], + ... [6.4, 3.2, 1], [5.9, 3.0, 2]], + ... columns=['length', 'width', 'species']) + >>> fig = df.plot.scatter(x='length', y='width') + >>> fig.update_layout(template="plotly_dark") # doctest: +SKIP + """ + return self(kind="scatter", x=x, y=y, **kwds) + + def hexbin(self, **kwds): + return unsupported_function(class_name="pd.DataFrame", method_name="hexbin")() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/matplotlib.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/matplotlib.py new file mode 100644 index 0000000..39e862b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/matplotlib.py @@ -0,0 +1,928 @@ +# +# 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 warnings +from distutils.version import LooseVersion + +import matplotlib as mat +import numpy as np +from matplotlib.axes._base import _process_plot_format +from pandas.core.dtypes.inference import is_list_like +from pandas.io.formats.printing import pprint_thing + +from pandas.plotting._matplotlib import ( # type: ignore[attr-defined] + BarPlot as PandasBarPlot, + BoxPlot as PandasBoxPlot, + HistPlot as PandasHistPlot, + PiePlot as PandasPiePlot, + AreaPlot as PandasAreaPlot, + LinePlot as PandasLinePlot, + BarhPlot as PandasBarhPlot, + ScatterPlot as PandasScatterPlot, + KdePlot as PandasKdePlot, +) +from pandas.plotting._core import PlotAccessor +from pandas.plotting._matplotlib.core import MPLPlot as PandasMPLPlot + +from pyspark.pandas.plot import ( + TopNPlotBase, + SampledPlotBase, + HistogramPlotBase, + BoxPlotBase, + unsupported_function, + KdePlotBase, +) + +_all_kinds = PlotAccessor._all_kinds # type: ignore[attr-defined] + + +class PandasOnSparkBarPlot(PandasBarPlot, TopNPlotBase): + _kind = "bar" + + def __init__(self, data, **kwargs): + super().__init__(self.get_top_n(data), **kwargs) + + def _plot(self, ax, x, y, w, start=0, log=False, **kwds): + self.set_result_text(ax) + return ax.bar(x, y, w, bottom=start, log=log, **kwds) + + +class PandasOnSparkBoxPlot(PandasBoxPlot, BoxPlotBase): + _kind = "box" + + def boxplot( + self, + ax, + bxpstats, + notch=None, + sym=None, + vert=None, + whis=None, + positions=None, + widths=None, + patch_artist=None, + bootstrap=None, + usermedians=None, + conf_intervals=None, + meanline=None, + showmeans=None, + showcaps=None, + showbox=None, + showfliers=None, + boxprops=None, + labels=None, + flierprops=None, + medianprops=None, + meanprops=None, + capprops=None, + whiskerprops=None, + manage_ticks=None, + # manage_xticks is for compatibility of matplotlib < 3.1.0. + # Remove this when minimum version is 3.0.0 + manage_xticks=None, + autorange=False, + zorder=None, + precision=None, + ): + def update_dict(dictionary, rc_name, properties): + """Loads properties in the dictionary from rc file if not already + in the dictionary""" + rc_str = "boxplot.{0}.{1}" + if dictionary is None: + dictionary = dict() + for prop_dict in properties: + dictionary.setdefault(prop_dict, mat.rcParams[rc_str.format(rc_name, prop_dict)]) + return dictionary + + # Common property dictionaries loading from rc + flier_props = [ + "color", + "marker", + "markerfacecolor", + "markeredgecolor", + "markersize", + "linestyle", + "linewidth", + ] + default_props = ["color", "linewidth", "linestyle"] + + boxprops = update_dict(boxprops, "boxprops", default_props) + whiskerprops = update_dict(whiskerprops, "whiskerprops", default_props) + capprops = update_dict(capprops, "capprops", default_props) + medianprops = update_dict(medianprops, "medianprops", default_props) + meanprops = update_dict(meanprops, "meanprops", default_props) + flierprops = update_dict(flierprops, "flierprops", flier_props) + + if patch_artist: + boxprops["linestyle"] = "solid" + boxprops["edgecolor"] = boxprops.pop("color") + + # if non-default sym value, put it into the flier dictionary + # the logic for providing the default symbol ('b+') now lives + # in bxp in the initial value of final_flierprops + # handle all of the `sym` related logic here so we only have to pass + # on the flierprops dict. + if sym is not None: + # no-flier case, which should really be done with + # 'showfliers=False' but none-the-less deal with it to keep back + # compatibility + if sym == "": + # blow away existing dict and make one for invisible markers + flierprops = dict(linestyle="none", marker="", color="none") + # turn the fliers off just to be safe + showfliers = False + # now process the symbol string + else: + # process the symbol string + # discarded linestyle + _, marker, color = _process_plot_format(sym) + # if we have a marker, use it + if marker is not None: + flierprops["marker"] = marker + # if we have a color, use it + if color is not None: + # assume that if color is passed in the user want + # filled symbol, if the users want more control use + # flierprops + flierprops["color"] = color + flierprops["markerfacecolor"] = color + flierprops["markeredgecolor"] = color + + # replace medians if necessary: + if usermedians is not None: + if len(np.ravel(usermedians)) != len(bxpstats) or np.shape(usermedians)[0] != len( + bxpstats + ): + raise ValueError("usermedians length not compatible with x") + else: + # reassign medians as necessary + for stats, med in zip(bxpstats, usermedians): + if med is not None: + stats["med"] = med + + if conf_intervals is not None: + if np.shape(conf_intervals)[0] != len(bxpstats): + err_mess = "conf_intervals length not compatible with x" + raise ValueError(err_mess) + else: + for stats, ci in zip(bxpstats, conf_intervals): + if ci is not None: + if len(ci) != 2: + raise ValueError("each confidence interval must " "have two values") + else: + if ci[0] is not None: + stats["cilo"] = ci[0] + if ci[1] is not None: + stats["cihi"] = ci[1] + + should_manage_ticks = True + if manage_xticks is not None: + should_manage_ticks = manage_xticks + if manage_ticks is not None: + should_manage_ticks = manage_ticks + + if LooseVersion(mat.__version__) < LooseVersion("3.1.0"): + extra_args = {"manage_xticks": should_manage_ticks} + else: + extra_args = {"manage_ticks": should_manage_ticks} + + artists = ax.bxp( + bxpstats, + positions=positions, + widths=widths, + vert=vert, + patch_artist=patch_artist, + shownotches=notch, + showmeans=showmeans, + showcaps=showcaps, + showbox=showbox, + boxprops=boxprops, + flierprops=flierprops, + medianprops=medianprops, + meanprops=meanprops, + meanline=meanline, + showfliers=showfliers, + capprops=capprops, + whiskerprops=whiskerprops, + zorder=zorder, + **extra_args, + ) + return artists + + def _plot(self, ax, bxpstats, column_num=None, return_type="axes", **kwds): + bp = self.boxplot(ax, bxpstats, **kwds) + + if return_type == "dict": + return bp, bp + elif return_type == "both": + return self.BP(ax=ax, lines=bp), bp + else: + return ax, bp + + def _compute_plot_data(self): + colname = self.data.name + spark_column_name = self.data._internal.spark_column_name_for(self.data._column_label) + data = self.data + + # Updates all props with the rc defaults from matplotlib + self.kwds.update(PandasOnSparkBoxPlot.rc_defaults(**self.kwds)) + + # Gets some important kwds + showfliers = self.kwds.get("showfliers", False) + whis = self.kwds.get("whis", 1.5) + labels = self.kwds.get("labels", [colname]) + + # This one is pandas-on-Spark specific to control precision for approx_percentile + precision = self.kwds.get("precision", 0.01) + + # # Computes mean, median, Q1 and Q3 with approx_percentile and precision + col_stats, col_fences = BoxPlotBase.compute_stats(data, spark_column_name, whis, precision) + + # # Creates a column to flag rows as outliers or not + outliers = BoxPlotBase.outliers(data, spark_column_name, *col_fences) + + # # Computes min and max values of non-outliers - the whiskers + whiskers = BoxPlotBase.calc_whiskers(spark_column_name, outliers) + + if showfliers: + fliers = BoxPlotBase.get_fliers(spark_column_name, outliers, whiskers[0]) + else: + fliers = [] + + # Builds bxpstats dict + stats = [] + item = { + "mean": col_stats["mean"], + "med": col_stats["med"], + "q1": col_stats["q1"], + "q3": col_stats["q3"], + "whislo": whiskers[0], + "whishi": whiskers[1], + "fliers": fliers, + "label": labels[0], + } + stats.append(item) + + self.data = {labels[0]: stats} + + def _make_plot(self): + bxpstats = list(self.data.values())[0] + ax = self._get_ax(0) + kwds = self.kwds.copy() + + for stats in bxpstats: + if len(stats["fliers"]) > 1000: + stats["fliers"] = stats["fliers"][:1000] + ax.text( + 1, + 1, + "showing top 1,000 fliers only", + size=6, + ha="right", + va="bottom", + transform=ax.transAxes, + ) + + ret, bp = self._plot(ax, bxpstats, column_num=0, return_type=self.return_type, **kwds) + self.maybe_color_bp(bp) + self._return_obj = ret + + labels = [lbl for lbl, _ in self.data.items()] + labels = [pprint_thing(lbl) for lbl in labels] + if not self.use_index: + labels = [pprint_thing(key) for key in range(len(labels))] + self._set_ticklabels(ax, labels) + + @staticmethod + def rc_defaults( + notch=None, + vert=None, + whis=None, + patch_artist=None, + bootstrap=None, + meanline=None, + showmeans=None, + showcaps=None, + showbox=None, + showfliers=None, + **kwargs, + ): + # Missing arguments default to rcParams. + if whis is None: + whis = mat.rcParams["boxplot.whiskers"] + if bootstrap is None: + bootstrap = mat.rcParams["boxplot.bootstrap"] + + if notch is None: + notch = mat.rcParams["boxplot.notch"] + if vert is None: + vert = mat.rcParams["boxplot.vertical"] + if patch_artist is None: + patch_artist = mat.rcParams["boxplot.patchartist"] + if meanline is None: + meanline = mat.rcParams["boxplot.meanline"] + if showmeans is None: + showmeans = mat.rcParams["boxplot.showmeans"] + if showcaps is None: + showcaps = mat.rcParams["boxplot.showcaps"] + if showbox is None: + showbox = mat.rcParams["boxplot.showbox"] + if showfliers is None: + showfliers = mat.rcParams["boxplot.showfliers"] + + return dict( + whis=whis, + bootstrap=bootstrap, + notch=notch, + vert=vert, + patch_artist=patch_artist, + meanline=meanline, + showmeans=showmeans, + showcaps=showcaps, + showbox=showbox, + showfliers=showfliers, + ) + + +class PandasOnSparkHistPlot(PandasHistPlot, HistogramPlotBase): + _kind = "hist" + + def _args_adjust(self): + if is_list_like(self.bottom): + self.bottom = np.array(self.bottom) + + def _compute_plot_data(self): + self.data, self.bins = HistogramPlotBase.prepare_hist_data(self.data, self.bins) + + def _make_plot(self): + # TODO: this logic is similar to KdePlot. Might have to deduplicate it. + # 'num_colors' requires to calculate `shape` which has to count all. + # Use 1 for now to save the computation. + colors = self._get_colors(num_colors=1) + stacking_id = self._get_stacking_id() + output_series = HistogramPlotBase.compute_hist(self.data, self.bins) + + for (i, label), y in zip(enumerate(self.data._internal.column_labels), output_series): + ax = self._get_ax(i) + + kwds = self.kwds.copy() + + label = pprint_thing(label if len(label) > 1 else label[0]) + # `if hasattr(...)` makes plotting compatible with pandas < 1.3, + # see pandas-dev/pandas#40078. + label = ( + self._mark_right_label(label, index=i) + if hasattr(self, "_mark_right_label") + else label + ) + kwds["label"] = label + + style, kwds = self._apply_style_colors(colors, kwds, i, label) + if style is not None: + kwds["style"] = style + + kwds = self._make_plot_keywords(kwds, y) + artists = self._plot(ax, y, column_num=i, stacking_id=stacking_id, **kwds) + # `if hasattr(...)` makes plotting compatible with pandas < 1.3, + # see pandas-dev/pandas#40078. + self._append_legend_handles_labels(artists[0], label) if hasattr( + self, "_append_legend_handles_labels" + ) else self._add_legend_handle(artists[0], label, index=i) + + @classmethod + def _plot(cls, ax, y, style=None, bins=None, bottom=0, column_num=0, stacking_id=None, **kwds): + if column_num == 0: + cls._initialize_stacker(ax, stacking_id, len(bins) - 1) + + base = np.zeros(len(bins) - 1) + bottom = bottom + cls._get_stacked_values(ax, stacking_id, base, kwds["label"]) + + # Since the counts were computed already, we use them as weights and just generate + # one entry for each bin + n, bins, patches = ax.hist(bins[:-1], bins=bins, bottom=bottom, weights=y, **kwds) + + cls._update_stacker(ax, stacking_id, n) + return patches + + +class PandasOnSparkPiePlot(PandasPiePlot, TopNPlotBase): + _kind = "pie" + + def __init__(self, data, **kwargs): + super().__init__(self.get_top_n(data), **kwargs) + + def _make_plot(self): + self.set_result_text(self._get_ax(0)) + super()._make_plot() + + +class PandasOnSparkAreaPlot(PandasAreaPlot, SampledPlotBase): + _kind = "area" + + def __init__(self, data, **kwargs): + super().__init__(self.get_sampled(data), **kwargs) + + def _make_plot(self): + self.set_result_text(self._get_ax(0)) + super()._make_plot() + + +class PandasOnSparkLinePlot(PandasLinePlot, SampledPlotBase): + _kind = "line" + + def __init__(self, data, **kwargs): + super().__init__(self.get_sampled(data), **kwargs) + + def _make_plot(self): + self.set_result_text(self._get_ax(0)) + super()._make_plot() + + +class PandasOnSparkBarhPlot(PandasBarhPlot, TopNPlotBase): + _kind = "barh" + + def __init__(self, data, **kwargs): + super().__init__(self.get_top_n(data), **kwargs) + + def _make_plot(self): + self.set_result_text(self._get_ax(0)) + super()._make_plot() + + +class PandasOnSparkScatterPlot(PandasScatterPlot, TopNPlotBase): + _kind = "scatter" + + def __init__(self, data, x, y, **kwargs): + super().__init__(self.get_top_n(data), x, y, **kwargs) + + def _make_plot(self): + self.set_result_text(self._get_ax(0)) + super()._make_plot() + + +class PandasOnSparkKdePlot(PandasKdePlot, KdePlotBase): + _kind = "kde" + + def _compute_plot_data(self): + self.data = KdePlotBase.prepare_kde_data(self.data) + + def _make_plot(self): + # 'num_colors' requires to calculate `shape` which has to count all. + # Use 1 for now to save the computation. + colors = self._get_colors(num_colors=1) + stacking_id = self._get_stacking_id() + + sdf = self.data._internal.spark_frame + + for i, label in enumerate(self.data._internal.column_labels): + # 'y' is a Spark DataFrame that selects one column. + y = sdf.select(self.data._internal.spark_column_for(label)) + ax = self._get_ax(i) + + kwds = self.kwds.copy() + + label = pprint_thing(label if len(label) > 1 else label[0]) + # `if hasattr(...)` makes plotting compatible with pandas < 1.3, + # see pandas-dev/pandas#40078. + label = ( + self._mark_right_label(label, index=i) + if hasattr(self, "_mark_right_label") + else label + ) + kwds["label"] = label + + style, kwds = self._apply_style_colors(colors, kwds, i, label) + if style is not None: + kwds["style"] = style + + kwds = self._make_plot_keywords(kwds, y) + artists = self._plot(ax, y, column_num=i, stacking_id=stacking_id, **kwds) + # `if hasattr(...)` makes plotting compatible with pandas < 1.3, + # see pandas-dev/pandas#40078. + self._append_legend_handles_labels(artists[0], label) if hasattr( + self, "_append_legend_handles_labels" + ) else self._add_legend_handle(artists[0], label, index=i) + + def _get_ind(self, y): + return KdePlotBase.get_ind(y, self.ind) + + @classmethod + def _plot( + cls, ax, y, style=None, bw_method=None, ind=None, column_num=None, stacking_id=None, **kwds + ): + y = KdePlotBase.compute_kde(y, bw_method=bw_method, ind=ind) + lines = PandasMPLPlot._plot(ax, ind, y, style=style, **kwds) + return lines + + +_klasses = [ + PandasOnSparkHistPlot, + PandasOnSparkBarPlot, + PandasOnSparkBoxPlot, + PandasOnSparkPiePlot, + PandasOnSparkAreaPlot, + PandasOnSparkLinePlot, + PandasOnSparkBarhPlot, + PandasOnSparkScatterPlot, + PandasOnSparkKdePlot, +] +_plot_klass = {getattr(klass, "_kind"): klass for klass in _klasses} +_common_kinds = {"area", "bar", "barh", "box", "hist", "kde", "line", "pie"} +_series_kinds = _common_kinds.union(set()) +_dataframe_kinds = _common_kinds.union({"scatter", "hexbin"}) +_pandas_on_spark_all_kinds = _common_kinds.union(_series_kinds).union(_dataframe_kinds) + + +def plot_pandas_on_spark(data, kind, **kwargs): + if kind not in _pandas_on_spark_all_kinds: + raise ValueError("{} is not a valid plot kind".format(kind)) + + from pyspark.pandas import DataFrame, Series + + if isinstance(data, Series): + if kind not in _series_kinds: + return unsupported_function(class_name="pd.Series", method_name=kind)() + return plot_series(data=data, kind=kind, **kwargs) + elif isinstance(data, DataFrame): + if kind not in _dataframe_kinds: + return unsupported_function(class_name="pd.DataFrame", method_name=kind)() + return plot_frame(data=data, kind=kind, **kwargs) + + +def plot_series( + data, + kind="line", + ax=None, # Series unique + figsize=None, + use_index=True, + title=None, + grid=None, + legend=False, + style=None, + logx=False, + logy=False, + loglog=False, + xticks=None, + yticks=None, + xlim=None, + ylim=None, + rot=None, + fontsize=None, + colormap=None, + table=False, + yerr=None, + xerr=None, + label=None, + secondary_y=False, # Series unique + **kwds, +): + """ + Make plots of Series using matplotlib / pylab. + + Each plot kind has a corresponding method on the + ``Series.plot`` accessor: + ``s.plot(kind='line')`` is equivalent to + ``s.plot.line()``. + + Parameters + ---------- + data : Series + + kind : str + - 'line' : line plot (default) + - 'bar' : vertical bar plot + - 'barh' : horizontal bar plot + - 'hist' : histogram + - 'box' : boxplot + - 'kde' : Kernel Density Estimation plot + - 'density' : same as 'kde' + - 'area' : area plot + - 'pie' : pie plot + + ax : matplotlib axes object + If not passed, uses gca() + figsize : a tuple (width, height) in inches + use_index : boolean, default True + Use index as ticks for x axis + title : string or list + Title to use for the plot. If a string is passed, print the string at + the top of the figure. If a list is passed and `subplots` is True, + print each item in the list above the corresponding subplot. + grid : boolean, default None (matlab style default) + Axis grid lines + legend : False/True/'reverse' + Place legend on axis subplots + style : list or dict + matplotlib line style per column + logx : boolean, default False + Use log scaling on x axis + logy : boolean, default False + Use log scaling on y axis + loglog : boolean, default False + Use log scaling on both x and y axes + xticks : sequence + Values to use for the xticks + yticks : sequence + Values to use for the yticks + xlim : 2-tuple/list + ylim : 2-tuple/list + rot : int, default None + Rotation for ticks (xticks for vertical, yticks for horizontal plots) + fontsize : int, default None + Font size for xticks and yticks + colormap : str or matplotlib colormap object, default None + Colormap to select colors from. If string, load colormap with that name + from matplotlib. + colorbar : boolean, optional + If True, plot colorbar (only relevant for 'scatter' and 'hexbin' plots) + position : float + Specify relative alignments for bar plot layout. + From 0 (left/bottom-end) to 1 (right/top-end). Default is 0.5 (center) + table : boolean, Series or DataFrame, default False + If True, draw a table using the data in the DataFrame and the data will + be transposed to meet matplotlib's default layout. + If a Series or DataFrame is passed, use passed data to draw a table. + yerr : DataFrame, Series, array-like, dict and str + See :ref:`Plotting with Error Bars ` for + detail. + xerr : same types as yerr. + label : label argument to provide to plot + secondary_y : boolean or sequence of ints, default False + If True then y-axis will be on the right + mark_right : boolean, default True + When using a secondary_y axis, automatically mark the column + labels with "(right)" in the legend + **kwds : keywords + Options to pass to matplotlib plotting method + + Returns + ------- + axes : :class:`matplotlib.axes.Axes` or numpy.ndarray of them + + Notes + ----- + + - See matplotlib documentation online for more on this subject + - If `kind` = 'bar' or 'barh', you can specify relative alignments + for bar plot layout by `position` keyword. + From 0 (left/bottom-end) to 1 (right/top-end). Default is 0.5 (center) + """ + + # function copied from pandas.plotting._core + # so it calls modified _plot below + + import matplotlib.pyplot as plt + + if ax is None and len(plt.get_fignums()) > 0: + with plt.rc_context(): + ax = plt.gca() + ax = PandasMPLPlot._get_ax_layer(ax) + return _plot( + data, + kind=kind, + ax=ax, + figsize=figsize, + use_index=use_index, + title=title, + grid=grid, + legend=legend, + style=style, + logx=logx, + logy=logy, + loglog=loglog, + xticks=xticks, + yticks=yticks, + xlim=xlim, + ylim=ylim, + rot=rot, + fontsize=fontsize, + colormap=colormap, + table=table, + yerr=yerr, + xerr=xerr, + label=label, + secondary_y=secondary_y, + **kwds, + ) + + +def plot_frame( + data, + x=None, + y=None, + kind="line", + ax=None, + subplots=False, + sharex=None, + sharey=False, + layout=None, + figsize=None, + use_index=True, + title=None, + grid=None, + legend=True, + style=None, + logx=False, + logy=False, + loglog=False, + xticks=None, + yticks=None, + xlim=None, + ylim=None, + rot=None, + fontsize=None, + colormap=None, + table=False, + yerr=None, + xerr=None, + secondary_y=False, + sort_columns=False, + **kwds, +): + """ + Make plots of DataFrames using matplotlib / pylab. + + Each plot kind has a corresponding method on the + ``DataFrame.plot`` accessor: + ``psdf.plot(kind='line')`` is equivalent to + ``psdf.plot.line()``. + + Parameters + ---------- + data : DataFrame + + kind : str + - 'line' : line plot (default) + - 'bar' : vertical bar plot + - 'barh' : horizontal bar plot + - 'hist' : histogram + - 'box' : boxplot + - 'kde' : Kernel Density Estimation plot + - 'density' : same as 'kde' + - 'area' : area plot + - 'pie' : pie plot + - 'scatter' : scatter plot + ax : matplotlib axes object + If not passed, uses gca() + x : label or position, default None + y : label, position or list of label, positions, default None + Allows plotting of one column versus another. + figsize : a tuple (width, height) in inches + use_index : boolean, default True + Use index as ticks for x axis + title : string or list + Title to use for the plot. If a string is passed, print the string at + the top of the figure. If a list is passed and `subplots` is True, + print each item in the list above the corresponding subplot. + grid : boolean, default None (matlab style default) + Axis grid lines + legend : False/True/'reverse' + Place legend on axis subplots + style : list or dict + matplotlib line style per column + logx : boolean, default False + Use log scaling on x axis + logy : boolean, default False + Use log scaling on y axis + loglog : boolean, default False + Use log scaling on both x and y axes + xticks : sequence + Values to use for the xticks + yticks : sequence + Values to use for the yticks + xlim : 2-tuple/list + ylim : 2-tuple/list + sharex: bool or None, default is None + Whether to share x axis or not. + sharey: bool, default is False + Whether to share y axis or not. + rot : int, default None + Rotation for ticks (xticks for vertical, yticks for horizontal plots) + fontsize : int, default None + Font size for xticks and yticks + colormap : str or matplotlib colormap object, default None + Colormap to select colors from. If string, load colormap with that name + from matplotlib. + colorbar : boolean, optional + If True, plot colorbar (only relevant for 'scatter' and 'hexbin' plots) + position : float + Specify relative alignments for bar plot layout. + From 0 (left/bottom-end) to 1 (right/top-end). Default is 0.5 (center) + table : boolean, Series or DataFrame, default False + If True, draw a table using the data in the DataFrame and the data will + be transposed to meet matplotlib's default layout. + If a Series or DataFrame is passed, use passed data to draw a table. + yerr : DataFrame, Series, array-like, dict and str + See :ref:`Plotting with Error Bars ` for + detail. + xerr : same types as yerr. + label : label argument to provide to plot + secondary_y : boolean or sequence of ints, default False + If True then y-axis will be on the right + mark_right : boolean, default True + When using a secondary_y axis, automatically mark the column + labels with "(right)" in the legend + sort_columns: bool, default is False + When True, will sort values on plots. + + .. deprecated:: 3.4.0 + + **kwds : keywords + Options to pass to matplotlib plotting method + + Returns + ------- + axes : :class:`matplotlib.axes.Axes` or numpy.ndarray of them + + Notes + ----- + + - See matplotlib documentation online for more on this subject + - If `kind` = 'bar' or 'barh', you can specify relative alignments + for bar plot layout by `position` keyword. + From 0 (left/bottom-end) to 1 (right/top-end). Default is 0.5 (center) + """ + warnings.warn( + "Argument `sort_columns` will be removed in 4.0.0.", + FutureWarning, + ) + + return _plot( + data, + kind=kind, + x=x, + y=y, + ax=ax, + figsize=figsize, + use_index=use_index, + title=title, + grid=grid, + legend=legend, + subplots=subplots, + style=style, + logx=logx, + logy=logy, + loglog=loglog, + xticks=xticks, + yticks=yticks, + xlim=xlim, + ylim=ylim, + rot=rot, + fontsize=fontsize, + colormap=colormap, + table=table, + yerr=yerr, + xerr=xerr, + sharex=sharex, + sharey=sharey, + secondary_y=secondary_y, + layout=layout, + sort_columns=sort_columns, + **kwds, + ) + + +def _plot(data, x=None, y=None, subplots=False, ax=None, kind="line", **kwds): + from pyspark.pandas import DataFrame + + # function copied from pandas.plotting._core + # and adapted to handle pandas-on-Spark DataFrame and Series + + kind = kind.lower().strip() + kind = {"density": "kde"}.get(kind, kind) + if kind in _all_kinds: + klass = _plot_klass[kind] + else: + raise ValueError("%r is not a valid plot kind" % kind) + + # scatter and hexbin are inherited from PlanePlot which require x and y + if kind in ("scatter", "hexbin"): + plot_obj = klass(data, x, y, subplots=subplots, ax=ax, kind=kind, **kwds) + else: + + # check data type and do preprocess before applying plot + if isinstance(data, DataFrame): + if x is not None: + data = data.set_index(x) + # TODO: check if value of y is plottable + if y is not None: + data = data[y] + + plot_obj = klass(data, subplots=subplots, ax=ax, kind=kind, **kwds) + plot_obj.generate() + plot_obj.draw() + return plot_obj.result diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/plotly.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/plotly.py new file mode 100644 index 0000000..d54166a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/plot/plotly.py @@ -0,0 +1,261 @@ +# +# 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 inspect +from typing import TYPE_CHECKING, Union + +import pandas as pd + +from pyspark.pandas.plot import ( + HistogramPlotBase, + name_like_string, + PandasOnSparkPlotAccessor, + BoxPlotBase, + KdePlotBase, +) + +if TYPE_CHECKING: + import pyspark.pandas as ps + + +def plot_pandas_on_spark(data: Union["ps.DataFrame", "ps.Series"], kind: str, **kwargs): + import plotly + + # pandas-on-Spark specific plots + if kind == "pie": + return plot_pie(data, **kwargs) + if kind == "hist": + return plot_histogram(data, **kwargs) + if kind == "box": + return plot_box(data, **kwargs) + if kind == "kde" or kind == "density": + return plot_kde(data, **kwargs) + + # Other plots. + return plotly.plot(PandasOnSparkPlotAccessor.pandas_plot_data_map[kind](data), kind, **kwargs) + + +def plot_pie(data: Union["ps.DataFrame", "ps.Series"], **kwargs): + from plotly import express + + data = PandasOnSparkPlotAccessor.pandas_plot_data_map["pie"](data) + + if isinstance(data, pd.Series): + pdf = data.to_frame() + return express.pie(pdf, values=pdf.columns[0], names=pdf.index, **kwargs) + elif isinstance(data, pd.DataFrame): + values = kwargs.pop("y", None) + default_names = None + if values is not None: + default_names = data.index + + return express.pie( + data, + values=kwargs.pop("values", values), + names=kwargs.pop("names", default_names), + **kwargs, + ) + else: + raise RuntimeError("Unexpected type: [%s]" % type(data)) + + +def plot_histogram(data: Union["ps.DataFrame", "ps.Series"], **kwargs): + import plotly.graph_objs as go + import pyspark.pandas as ps + + bins = kwargs.get("bins", 10) + y = kwargs.get("y") + if y and isinstance(data, ps.DataFrame): + # Note that the results here are matched with matplotlib. x and y + # handling is different from pandas' plotly output. + data = data[y] + psdf, bins = HistogramPlotBase.prepare_hist_data(data, bins) + assert len(bins) > 2, "the number of buckets must be higher than 2." + output_series = HistogramPlotBase.compute_hist(psdf, bins) + prev = float("%.9f" % bins[0]) # to make it prettier, truncate. + text_bins = [] + for b in bins[1:]: + norm_b = float("%.9f" % b) + text_bins.append("[%s, %s)" % (prev, norm_b)) + prev = norm_b + text_bins[-1] = text_bins[-1][:-1] + "]" # replace ) to ] for the last bucket. + + bins = 0.5 * (bins[:-1] + bins[1:]) + + output_series = list(output_series) + bars = [] + for series in output_series: + bars.append( + go.Bar( + x=bins, + y=series, + name=name_like_string(series.name), + text=text_bins, + hovertemplate=( + "variable=" + name_like_string(series.name) + "
    value=%{text}
    count=%{y}" + ), + ) + ) + + layout_keys = inspect.signature(go.Layout).parameters.keys() + layout_kwargs = {k: v for k, v in kwargs.items() if k in layout_keys} + + fig = go.Figure(data=bars, layout=go.Layout(**layout_kwargs)) + fig["layout"]["barmode"] = "stack" + fig["layout"]["xaxis"]["title"] = "value" + fig["layout"]["yaxis"]["title"] = "count" + return fig + + +def plot_box(data: Union["ps.DataFrame", "ps.Series"], **kwargs): + import plotly.graph_objs as go + import pyspark.pandas as ps + from pyspark.sql.types import NumericType + + # 'whis' isn't actually an argument in plotly (but in matplotlib). But seems like + # plotly doesn't expose the reach of the whiskers to the beyond the first and + # third quartiles (?). Looks they use default 1.5. + whis = kwargs.pop("whis", 1.5) + # 'precision' is pandas-on-Spark specific to control precision for approx_percentile + precision = kwargs.pop("precision", 0.01) + + # Plotly options + boxpoints = kwargs.pop("boxpoints", "suspectedoutliers") + notched = kwargs.pop("notched", False) + if boxpoints not in ["suspectedoutliers", False]: + raise ValueError( + "plotly plotting backend does not support 'boxpoints' set to '%s'. " + "Set to 'suspectedoutliers' or False." % boxpoints + ) + if notched: + raise ValueError( + "plotly plotting backend does not support 'notched' set to '%s'. " + "Set to False." % notched + ) + + fig = go.Figure() + if isinstance(data, ps.Series): + colname = name_like_string(data.name) + spark_column_name = data._internal.spark_column_name_for(data._column_label) + + # Computes mean, median, Q1 and Q3 with approx_percentile and precision + col_stats, col_fences = BoxPlotBase.compute_stats(data, spark_column_name, whis, precision) + + # Creates a column to flag rows as outliers or not + outliers = BoxPlotBase.outliers(data, spark_column_name, *col_fences) + + # Computes min and max values of non-outliers - the whiskers + whiskers = BoxPlotBase.calc_whiskers(spark_column_name, outliers) + + fliers = None + if boxpoints: + fliers = BoxPlotBase.get_fliers(spark_column_name, outliers, whiskers[0]) + fliers = [fliers] if len(fliers) > 0 else None + + fig.add_trace( + go.Box( + name=colname, + q1=[col_stats["q1"]], + median=[col_stats["med"]], + q3=[col_stats["q3"]], + mean=[col_stats["mean"]], + lowerfence=[whiskers[0]], + upperfence=[whiskers[1]], + y=fliers, + boxpoints=boxpoints, + notched=notched, + **kwargs, # this is for workarounds. Box takes different options from express.box. + ) + ) + fig["layout"]["xaxis"]["title"] = colname + + else: + numeric_column_names = [] + for column_label in data._internal.column_labels: + if isinstance(data._internal.spark_type_for(column_label), NumericType): + numeric_column_names.append(name_like_string(column_label)) + + # Computes mean, median, Q1 and Q3 with approx_percentile and precision + multicol_stats = BoxPlotBase.compute_multicol_stats( + data, numeric_column_names, whis, precision + ) + + # Creates a column to flag rows as outliers or not + outliers = BoxPlotBase.multicol_outliers(data, multicol_stats) + + # Computes min and max values of non-outliers - the whiskers + whiskers = BoxPlotBase.calc_multicol_whiskers(numeric_column_names, outliers) + + i = 0 + for colname in numeric_column_names: + col_stats = multicol_stats[colname] + col_whiskers = whiskers[colname] + + fig.add_trace( + go.Box( + x=[i], + name=colname, + q1=[col_stats["q1"]], + median=[col_stats["med"]], + q3=[col_stats["q3"]], + mean=[col_stats["mean"]], + lowerfence=[col_whiskers["min"]], + upperfence=[col_whiskers["max"]], + y=None, # todo: support y=fliers + boxpoints=boxpoints, + notched=notched, + **kwargs, + ) + ) + i += 1 + + fig["layout"]["yaxis"]["title"] = "value" + return fig + + +def plot_kde(data: Union["ps.DataFrame", "ps.Series"], **kwargs): + from plotly import express + import pyspark.pandas as ps + + if isinstance(data, ps.DataFrame) and "color" not in kwargs: + kwargs["color"] = "names" + + psdf = KdePlotBase.prepare_kde_data(data) + sdf = psdf._internal.spark_frame + data_columns = psdf._internal.data_spark_columns + ind = KdePlotBase.get_ind(sdf.select(*data_columns), kwargs.pop("ind", None)) + bw_method = kwargs.pop("bw_method", None) + + pdfs = [] + for label in psdf._internal.column_labels: + pdfs.append( + pd.DataFrame( + { + "Density": KdePlotBase.compute_kde( + sdf.select(psdf._internal.spark_column_for(label)), + ind=ind, + bw_method=bw_method, + ), + "names": name_like_string(label), + "index": ind, + } + ) + ) + pdf = pd.concat(pdfs) + + fig = express.line(pdf, x="index", y="Density", **kwargs) + fig["layout"]["xaxis"]["title"] = None + return fig diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/resample.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/resample.py new file mode 100644 index 0000000..30f8c9d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/resample.py @@ -0,0 +1,816 @@ +# +# 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. +# + +""" +A wrapper for ResampledData to behave like pandas Resampler. +""" +from abc import ABCMeta, abstractmethod +from distutils.version import LooseVersion +from functools import partial +from typing import ( + Any, + Generic, + List, + Optional, + Union, +) + +import numpy as np + +import pandas as pd +from pandas.tseries.frequencies import to_offset + +if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + from pandas.core.common import _builtin_table # type: ignore[attr-defined] +else: + from pandas.core.base import SelectionMixin + + _builtin_table = SelectionMixin._builtin_table # type: ignore[attr-defined] + +from pyspark import SparkContext +from pyspark.sql import Column, functions as F +from pyspark.sql.types import ( + NumericType, + StructField, + TimestampNTZType, + DataType, +) + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import FrameLike +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + SPARK_DEFAULT_INDEX_NAME, +) +from pyspark.pandas.missing.resample import ( + MissingPandasLikeDataFrameResampler, + MissingPandasLikeSeriesResampler, +) +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.utils import ( + scol_for, + verify_temp_column_name, +) +from pyspark.sql.utils import is_remote +from pyspark.pandas.spark.functions import timestampdiff + + +class Resampler(Generic[FrameLike], metaclass=ABCMeta): + """ + Class for resampling datetimelike data, a groupby-like operation. + + It's easiest to use obj.resample(...) to use Resampler. + + Parameters + ---------- + psdf : DataFrame + + Returns + ------- + a Resampler of the appropriate type + + Notes + ----- + After resampling, see aggregate, apply, and transform functions. + """ + + def __init__( + self, + psdf: DataFrame, + resamplekey: Optional[Series], + rule: str, + closed: Optional[str] = None, + label: Optional[str] = None, + agg_columns: List[Series] = [], + ): + self._psdf = psdf + self._resamplekey = resamplekey + + self._offset = to_offset(rule) + if self._offset.rule_code not in ["A-DEC", "M", "D", "H", "T", "S"]: + raise ValueError("rule code {} is not supported".format(self._offset.rule_code)) + if not getattr(self._offset, "n") > 0: + raise ValueError("rule offset must be positive") + + if closed is None: + self._closed = "right" if self._offset.rule_code in ["A-DEC", "M"] else "left" + elif closed in ["left", "right"]: + self._closed = closed + else: + raise ValueError("invalid closed: '{}'".format(closed)) + + if label is None: + self._label = "right" if self._offset.rule_code in ["A-DEC", "M"] else "left" + elif label in ["left", "right"]: + self._label = label + else: + raise ValueError("invalid label: '{}'".format(label)) + + self._agg_columns = agg_columns + + @property + def _resamplekey_scol(self) -> Column: + if self._resamplekey is None: + return self._psdf.index.spark.column + else: + return self._resamplekey.spark.column + + @property + def _resamplekey_type(self) -> DataType: + if self._resamplekey is None: + return self._psdf.index.spark.data_type + else: + return self._resamplekey.spark.data_type + + @property + def _agg_columns_scols(self) -> List[Column]: + return [s.spark.column for s in self._agg_columns] + + def get_make_interval( # type: ignore[return] + self, unit: str, col: Union[Column, int, float] + ) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import lit, make_interval + + col = col if not isinstance(col, (int, float)) else lit(col) # type: ignore[assignment] + if unit == "MONTH": + return make_interval(months=col) # type: ignore + if unit == "HOUR": + return make_interval(hours=col) # type: ignore + if unit == "MINUTE": + return make_interval(mins=col) # type: ignore + if unit == "SECOND": + return make_interval(secs=col) # type: ignore + else: + sql_utils = SparkContext._active_spark_context._jvm.PythonSQLUtils + col = col._jc if isinstance(col, Column) else F.lit(col)._jc + return sql_utils.makeInterval(unit, col) + + def _bin_timestamp(self, origin: pd.Timestamp, ts_scol: Column) -> Column: + key_type = self._resamplekey_type + origin_scol = F.lit(origin) + (rule_code, n) = (self._offset.rule_code, getattr(self._offset, "n")) + left_closed, right_closed = (self._closed == "left", self._closed == "right") + left_labeled, right_labeled = (self._label == "left", self._label == "right") + + if rule_code == "A-DEC": + assert ( + origin.month == 12 + and origin.day == 31 + and origin.hour == 0 + and origin.minute == 0 + and origin.second == 0 + ) + + diff = F.year(ts_scol) - F.year(origin_scol) + mod = F.lit(0) if n == 1 else (diff % n) + edge_cond = (mod == 0) & (F.month(ts_scol) == 12) & (F.dayofmonth(ts_scol) == 31) + + edge_label = F.year(ts_scol) + if left_closed and right_labeled: + edge_label += n + elif right_closed and left_labeled: + edge_label -= n + + if left_labeled: + non_edge_label = F.when(mod == 0, F.year(ts_scol) - n).otherwise( + F.year(ts_scol) - mod + ) + else: + non_edge_label = F.when(mod == 0, F.year(ts_scol)).otherwise( + F.year(ts_scol) - (mod - n) + ) + + ret = F.to_timestamp( + F.make_date( + F.when(edge_cond, edge_label).otherwise(non_edge_label), F.lit(12), F.lit(31) + ) + ) + + elif rule_code == "M": + assert ( + origin.is_month_end + and origin.hour == 0 + and origin.minute == 0 + and origin.second == 0 + ) + + diff = ( + (F.year(ts_scol) - F.year(origin_scol)) * 12 + + F.month(ts_scol) + - F.month(origin_scol) + ) + mod = F.lit(0) if n == 1 else (diff % n) + edge_cond = (mod == 0) & (F.dayofmonth(ts_scol) == F.dayofmonth(F.last_day(ts_scol))) + + truncated_ts_scol = F.date_trunc("MONTH", ts_scol) + edge_label = truncated_ts_scol + if left_closed and right_labeled: + edge_label += self.get_make_interval("MONTH", n) + elif right_closed and left_labeled: + edge_label -= self.get_make_interval("MONTH", n) + + if left_labeled: + non_edge_label = F.when( + mod == 0, + truncated_ts_scol - self.get_make_interval("MONTH", n), + ).otherwise(truncated_ts_scol - self.get_make_interval("MONTH", mod)) + else: + non_edge_label = F.when(mod == 0, truncated_ts_scol).otherwise( + truncated_ts_scol - self.get_make_interval("MONTH", mod - n) + ) + + ret = F.to_timestamp( + F.last_day(F.when(edge_cond, edge_label).otherwise(non_edge_label)) + ) + + elif rule_code == "D": + assert origin.hour == 0 and origin.minute == 0 and origin.second == 0 + + if n == 1: + # NOTE: the logic to process '1D' is different from the cases with n>1, + # since hour/minute/second parts are taken into account to determine edges! + edge_cond = ( + (F.hour(ts_scol) == 0) & (F.minute(ts_scol) == 0) & (F.second(ts_scol) == 0) + ) + + if left_closed and left_labeled: + ret = F.date_trunc("DAY", ts_scol) + elif left_closed and right_labeled: + ret = F.date_trunc("DAY", F.date_add(ts_scol, 1)) + elif right_closed and left_labeled: + ret = F.when(edge_cond, F.date_trunc("DAY", F.date_sub(ts_scol, 1))).otherwise( + F.date_trunc("DAY", ts_scol) + ) + else: + ret = F.when(edge_cond, F.date_trunc("DAY", ts_scol)).otherwise( + F.date_trunc("DAY", F.date_add(ts_scol, 1)) + ) + + else: + diff = F.datediff(end=ts_scol, start=origin_scol) + mod = diff % n + + edge_cond = mod == 0 + + truncated_ts_scol = F.date_trunc("DAY", ts_scol) + edge_label = truncated_ts_scol + if left_closed and right_labeled: + edge_label = F.date_add(truncated_ts_scol, n) + elif right_closed and left_labeled: + edge_label = F.date_sub(truncated_ts_scol, n) + + if left_labeled: + non_edge_label = F.date_sub(truncated_ts_scol, mod) + else: + non_edge_label = F.date_sub(truncated_ts_scol, mod - n) + + ret = F.when(edge_cond, edge_label).otherwise(non_edge_label) + + elif rule_code in ["H", "T", "S"]: + unit_mapping = {"H": "HOUR", "T": "MINUTE", "S": "SECOND"} + unit_str = unit_mapping[rule_code] + + truncated_ts_scol = F.date_trunc(unit_str, ts_scol) + if isinstance(key_type, TimestampNTZType): + truncated_ts_scol = F.to_timestamp_ntz(truncated_ts_scol) + diff = timestampdiff(unit_str, origin_scol, truncated_ts_scol) + mod = F.lit(0) if n == 1 else (diff % F.lit(n)) + + if rule_code == "H": + assert origin.minute == 0 and origin.second == 0 + edge_cond = (mod == 0) & (F.minute(ts_scol) == 0) & (F.second(ts_scol) == 0) + elif rule_code == "T": + assert origin.second == 0 + edge_cond = (mod == 0) & (F.second(ts_scol) == 0) + else: + edge_cond = mod == 0 + + edge_label = truncated_ts_scol + if left_closed and right_labeled: + edge_label += self.get_make_interval(unit_str, n) + elif right_closed and left_labeled: + edge_label -= self.get_make_interval(unit_str, n) + + if left_labeled: + non_edge_label = F.when(mod == 0, truncated_ts_scol).otherwise( + truncated_ts_scol - self.get_make_interval(unit_str, mod) + ) + else: + non_edge_label = F.when( + mod == 0, + truncated_ts_scol + self.get_make_interval(unit_str, n), + ).otherwise(truncated_ts_scol - self.get_make_interval(unit_str, mod - n)) + + ret = F.when(edge_cond, edge_label).otherwise(non_edge_label) + + else: + raise ValueError("Got the unexpected unit {}".format(rule_code)) + + if isinstance(key_type, TimestampNTZType): + return F.to_timestamp_ntz(ret) + else: + return ret + + def _downsample(self, f: str) -> DataFrame: + """ + Downsample the defined function. + + Parameters + ---------- + how : string / mapped function + **kwargs : kw args passed to how function + """ + + # a simple example to illustrate the computation: + # dates = [ + # datetime(2012, 1, 2), + # datetime(2012, 5, 3), + # datetime(2022, 5, 3), + # ] + # index = pd.DatetimeIndex(dates) + # pdf = pd.DataFrame(np.array([1,2,3]), index=index, columns=['A']) + # pdf.resample('3Y').max() + # A + # 2012-12-31 2.0 + # 2015-12-31 NaN + # 2018-12-31 NaN + # 2021-12-31 NaN + # 2024-12-31 3.0 + # + # in this case: + # 1, obtain one origin point to bin all timestamps, we can get one (2009-12-31) + # from the minimum timestamp (2012-01-02); + # 2, the default intervals for 'Y' are right-closed, so intervals are: + # (2009-12-31, 2012-12-31], (2012-12-31, 2015-12-31], (2015-12-31, 2018-12-31], ... + # 3, bin all timestamps, for example, 2022-05-03 belongs to interval + # (2021-12-31, 2024-12-31], since the default label is 'right', label it with the right + # edge 2024-12-31; + # 4, some intervals maybe too large for this down sampling, so we need to pad the dataframe + # to avoid missing some results, like: 2015-12-31, 2018-12-31 and 2021-12-31; + # 5, union the binned dataframe and padded dataframe, and apply aggregation 'max' to get + # the final results; + + # one action to obtain the range, in the future we may cache it in the index. + ts_min, ts_max = ( + self._psdf._internal.spark_frame.select( + F.min(self._resamplekey_scol), F.max(self._resamplekey_scol) + ) + .toPandas() + .iloc[0] + ) + + # the logic to obtain an origin point to bin the timestamps is too complex to follow, + # here just use Pandas' resample on a 1-length series to get it. + ts_origin = ( + pd.Series([0], index=[ts_min]) + .resample(rule=self._offset.freqstr, closed=self._closed, label="left") + .sum() + .index[0] + ) + assert ts_origin <= ts_min + + bin_col_name = "__tmp_resample_bin_col__" + bin_col_label = verify_temp_column_name(self._psdf, bin_col_name) + bin_col_field = InternalField( + dtype=np.dtype("datetime64[ns]"), + struct_field=StructField(bin_col_name, self._resamplekey_type, True), + ) + bin_scol = self._bin_timestamp(ts_origin, self._resamplekey_scol) + + agg_columns = [ + psser for psser in self._agg_columns if (isinstance(psser.spark.data_type, NumericType)) + ] + assert len(agg_columns) > 0 + + # in the binning side, label the timestamps according to the origin and the freq(rule) + bin_sdf = self._psdf._internal.spark_frame.select( + F.col(SPARK_DEFAULT_INDEX_NAME), + bin_scol.alias(bin_col_name), + *[psser.spark.column for psser in agg_columns], + ) + + # in the padding side, insert necessary points + # again, directly apply Pandas' resample on a 2-length series to obtain the indices + pad_sdf = ( + ps.from_pandas( + pd.Series([0, 0], index=[ts_min, ts_max]) + .resample(rule=self._offset.freqstr, closed=self._closed, label=self._label) + .sum() + .index + ) + ._internal.spark_frame.select(F.col(SPARK_DEFAULT_INDEX_NAME).alias(bin_col_name)) + .where((ts_min <= F.col(bin_col_name)) & (F.col(bin_col_name) <= ts_max)) + ) + + # union the above two spark dataframes. + sdf = bin_sdf.unionByName(pad_sdf, allowMissingColumns=True).where( + ~F.isnull(F.col(bin_col_name)) + ) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + data_spark_columns=[F.col(bin_col_name)] + + [scol_for(sdf, psser._internal.data_spark_column_names[0]) for psser in agg_columns], + column_labels=[bin_col_label] + [psser._column_label for psser in agg_columns], + data_fields=[bin_col_field] + + [psser._internal.data_fields[0].copy(nullable=True) for psser in agg_columns], + column_label_names=self._psdf._internal.column_label_names, + ) + psdf: DataFrame = DataFrame(internal) + + groupby = psdf.groupby(psdf._psser_for(bin_col_label), dropna=False) + downsampled = getattr(groupby, f)() + downsampled.index.name = None + + return downsampled + + @abstractmethod + def _handle_output(self, psdf: DataFrame) -> FrameLike: + pass + + def min(self) -> FrameLike: + """ + Compute min of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").min().sort_index() + A B + 2022-05-01 0.171162 0.338864 + 2022-05-04 0.010527 0.561204 + 2022-05-07 NaN NaN + 2022-05-10 0.813726 0.745100 + """ + return self._handle_output(self._downsample("min")) + + def max(self) -> FrameLike: + """ + Compute max of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").max().sort_index() + A B + 2022-05-01 0.420538 0.859182 + 2022-05-04 0.270533 0.691041 + 2022-05-07 NaN NaN + 2022-05-10 0.813726 0.745100 + """ + return self._handle_output(self._downsample("max")) + + def sum(self) -> FrameLike: + """ + Compute sum of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").sum().sort_index() + A B + 2022-05-01 0.800160 1.679727 + 2022-05-04 0.281060 1.252245 + 2022-05-07 0.000000 0.000000 + 2022-05-10 0.813726 0.745100 + """ + return self._handle_output(self._downsample("sum").fillna(0.0)) + + def mean(self) -> FrameLike: + """ + Compute mean of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").mean().sort_index() + A B + 2022-05-01 0.266720 0.559909 + 2022-05-04 0.140530 0.626123 + 2022-05-07 NaN NaN + 2022-05-10 0.813726 0.745100 + """ + return self._handle_output(self._downsample("mean")) + + def std(self) -> FrameLike: + """ + Compute std of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").std().sort_index() + A B + 2022-05-01 0.134509 0.268835 + 2022-05-04 0.183852 0.091809 + 2022-05-07 NaN NaN + 2022-05-10 NaN NaN + """ + return self._handle_output(self._downsample("std")) + + def var(self) -> FrameLike: + """ + Compute var of resampled values. + + .. versionadded:: 3.4.0 + + See Also + -------- + pyspark.pandas.Series.groupby + pyspark.pandas.DataFrame.groupby + + Examples + -------- + >>> import numpy as np + >>> from datetime import datetime + >>> np.random.seed(22) + >>> dates = [ + ... datetime(2022, 5, 1, 4, 5, 6), + ... datetime(2022, 5, 3), + ... datetime(2022, 5, 3, 23, 59, 59), + ... datetime(2022, 5, 4), + ... pd.NaT, + ... datetime(2022, 5, 4, 0, 0, 1), + ... datetime(2022, 5, 11), + ... ] + >>> df = ps.DataFrame( + ... np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=["A", "B"] + ... ) + >>> df + A B + 2022-05-01 04:05:06 0.208461 0.481681 + 2022-05-03 00:00:00 0.420538 0.859182 + 2022-05-03 23:59:59 0.171162 0.338864 + 2022-05-04 00:00:00 0.270533 0.691041 + NaT 0.220405 0.811951 + 2022-05-04 00:00:01 0.010527 0.561204 + 2022-05-11 00:00:00 0.813726 0.745100 + >>> df.resample("3D").var().sort_index() + A B + 2022-05-01 0.018093 0.072272 + 2022-05-04 0.033802 0.008429 + 2022-05-07 NaN NaN + 2022-05-10 NaN NaN + """ + return self._handle_output(self._downsample("var")) + + +class DataFrameResampler(Resampler[DataFrame]): + def __init__( + self, + psdf: DataFrame, + resamplekey: Optional[Series], + rule: str, + closed: Optional[str] = None, + label: Optional[str] = None, + agg_columns: List[Series] = [], + ): + super().__init__( + psdf=psdf, + resamplekey=resamplekey, + rule=rule, + closed=closed, + label=label, + agg_columns=agg_columns, + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeDataFrameResampler, item): + property_or_func = getattr(MissingPandasLikeDataFrameResampler, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + + def _handle_output(self, psdf: DataFrame) -> DataFrame: + return psdf + + +class SeriesResampler(Resampler[Series]): + def __init__( + self, + psser: Series, + resamplekey: Optional[Series], + rule: str, + closed: Optional[str] = None, + label: Optional[str] = None, + agg_columns: List[Series] = [], + ): + super().__init__( + psdf=psser._psdf, + resamplekey=resamplekey, + rule=rule, + closed=closed, + label=label, + agg_columns=agg_columns, + ) + self._psser = psser + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeSeriesResampler, item): + property_or_func = getattr(MissingPandasLikeSeriesResampler, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + + def _handle_output(self, psdf: DataFrame) -> Series: + return first_series(psdf).rename(self._psser.name) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.resample + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.resample.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.resample tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.resample, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/series.py new file mode 100644 index 0000000..95ca92e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/series.py @@ -0,0 +1,7440 @@ +# +# 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. +# + +""" +A wrapper class for Spark Column to behave like pandas Series. +""" +import datetime +import re +import inspect +import warnings +from collections.abc import Mapping +from functools import partial, reduce +from typing import ( + Any, + Callable, + Dict, + Generic, + IO, + Iterable, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, + no_type_check, + overload, + TYPE_CHECKING, +) + +import numpy as np +import pandas as pd +from pandas.core.accessor import CachedAccessor +from pandas.io.formats.printing import pprint_thing +from pandas.api.types import ( # type: ignore[attr-defined] + is_list_like, + is_hashable, + CategoricalDtype, +) +from pandas.tseries.frequencies import DateOffset +from pyspark.sql import functions as F, Column as PySparkColumn, DataFrame as SparkDataFrame +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + IntegralType, + LongType, + NumericType, + Row, + StructType, + TimestampType, +) +from pyspark.sql.window import Window +from pyspark.sql.utils import get_column_class, get_window_class + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas._typing import Axis, Dtype, Label, Name, Scalar, T +from pyspark.pandas.accessors import PandasOnSparkSeriesMethods +from pyspark.pandas.categorical import CategoricalAccessor +from pyspark.pandas.config import get_option +from pyspark.pandas.correlation import ( + compute, + CORRELATION_VALUE_1_COLUMN, + CORRELATION_VALUE_2_COLUMN, + CORRELATION_CORR_OUTPUT_COLUMN, + CORRELATION_COUNT_OUTPUT_COLUMN, +) +from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.exceptions import SparkPandasIndexingError +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.generic import Frame +from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + DEFAULT_SERIES_NAME, + NATURAL_ORDER_COLUMN_NAME, + SPARK_DEFAULT_INDEX_NAME, + SPARK_DEFAULT_SERIES_NAME, +) +from pyspark.pandas.missing.series import MissingPandasLikeSeries +from pyspark.pandas.plot import PandasOnSparkPlotAccessor +from pyspark.pandas.utils import ( + combine_frames, + is_name_like_tuple, + is_name_like_value, + name_like_string, + same_anchor, + scol_for, + sql_conf, + validate_arguments_and_invoke_function, + validate_axis, + validate_bool_kwarg, + verify_temp_column_name, + SPARK_CONF_ARROW_ENABLED, + log_advice, +) +from pyspark.pandas.datetimes import DatetimeMethods +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.spark.accessors import SparkSeriesMethods +from pyspark.pandas.strings import StringMethods +from pyspark.pandas.typedef import ( + infer_return_type, + spark_type_to_pandas_dtype, + ScalarType, + SeriesType, + create_type_for_series_type, +) +from pyspark.pandas.typedef.typehints import as_spark_type + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName + + from pyspark.pandas.groupby import SeriesGroupBy + from pyspark.pandas.resample import SeriesResampler + from pyspark.pandas.indexes import Index + from pyspark.pandas.spark.accessors import SparkIndexOpsMethods + +# This regular expression pattern is compiled and defined here to avoid to compile the same +# pattern every time it is used in _repr_ in Series. +# This pattern basically seeks the footer string from pandas' +REPR_PATTERN = re.compile(r"Length: (?P[0-9]+)") + +_flex_doc_SERIES = """ +Return {desc} of series and other, element-wise (binary operator `{op_name}`). + +Equivalent to ``{equiv}`` + +Parameters +---------- +other : Series or scalar value +fill_value : Scalar value, default None + +Returns +------- +Series + The result of the operation. + +See Also +-------- +Series.{reverse} + +{series_examples} +""" + +_add_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.add(df.b) +a 4.0 +b NaN +c 6.0 +d NaN +dtype: float64 + +>>> df.a.radd(df.b) +a 4.0 +b NaN +c 6.0 +d NaN +dtype: float64 +""" + +_sub_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.subtract(df.b) +a 0.0 +b NaN +c 2.0 +d NaN +dtype: float64 + +>>> df.a.rsub(df.b) +a 0.0 +b NaN +c -2.0 +d NaN +dtype: float64 +""" + +_mul_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.multiply(df.b) +a 4.0 +b NaN +c 8.0 +d NaN +dtype: float64 + +>>> df.a.rmul(df.b) +a 4.0 +b NaN +c 8.0 +d NaN +dtype: float64 +""" + +_div_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.divide(df.b) +a 1.0 +b NaN +c 2.0 +d NaN +dtype: float64 + +>>> df.a.rdiv(df.b) +a 1.0 +b NaN +c 0.5 +d NaN +dtype: float64 +""" + +_pow_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.pow(df.b) +a 4.0 +b NaN +c 16.0 +d NaN +dtype: float64 + +>>> df.a.rpow(df.b) +a 4.0 +b NaN +c 16.0 +d NaN +dtype: float64 +""" + +_mod_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.mod(df.b) +a 0.0 +b NaN +c 0.0 +d NaN +dtype: float64 + +>>> df.a.rmod(df.b) +a 0.0 +b NaN +c 2.0 +d NaN +dtype: float64 +""" + +_floordiv_example_SERIES = """ +Examples +-------- +>>> df = ps.DataFrame({'a': [2, 2, 4, np.nan], +... 'b': [2, np.nan, 2, np.nan]}, +... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) +>>> df + a b +a 2.0 2.0 +b 2.0 NaN +c 4.0 2.0 +d NaN NaN + +>>> df.a.floordiv(df.b) +a 1.0 +b NaN +c 2.0 +d NaN +dtype: float64 + +>>> df.a.rfloordiv(df.b) +a 1.0 +b NaN +c 0.0 +d NaN +dtype: float64 +""" + +# Needed to disambiguate Series.str and str type +str_type = str + + +class Series(Frame, IndexOpsMixin, Generic[T]): + """ + pandas-on-Spark Series that corresponds to pandas Series logically. This holds Spark Column + internally. + + :ivar _internal: an internal immutable Frame to manage metadata. + :type _internal: InternalFrame + :ivar _psdf: Parent's pandas-on-Spark DataFrame + :type _psdf: ps.DataFrame + + Parameters + ---------- + data : array-like, dict, or scalar value, pandas Series + Contains data stored in Series + Note that if `data` is a pandas Series, other arguments should not be used. + index : array-like or Index (1d) + Values must be hashable and have the same length as `data`. + Non-unique index values are allowed. Will default to + RangeIndex (0, 1, 2, ..., n) if not provided. If both a dict and index + sequence is used, the index will override the keys found in the + dict. + dtype : numpy.dtype or None + If None, dtype will be inferred + copy : boolean, default False + Copy input data + """ + + def __init__( # type: ignore[no-untyped-def] + self, data=None, index=None, dtype=None, name=None, copy=False, fastpath=False + ): + assert data is not None + + self._anchor: DataFrame + self._col_label: Label + if isinstance(data, DataFrame): + assert dtype is None + assert name is None + assert not copy + assert not fastpath + + self._anchor = data + self._col_label = index + else: + if isinstance(data, pd.Series): + assert index is None + assert dtype is None + assert name is None + assert not copy + assert not fastpath + s = data + else: + from pyspark.pandas.indexes.base import Index + + if isinstance(index, Index): + raise TypeError( + "The given index cannot be a pandas-on-Spark index. " + "Try pandas index or array-like." + ) + + s = pd.Series( + data=data, index=index, dtype=dtype, name=name, copy=copy, fastpath=fastpath + ) + internal = InternalFrame.from_pandas(pd.DataFrame(s)) + if s.name is None: + internal = internal.copy(column_labels=[None]) + anchor = DataFrame(internal) + + self._anchor = anchor + self._col_label = anchor._internal.column_labels[0] + object.__setattr__(anchor, "_psseries", {self._column_label: self}) + + @property + def _psdf(self) -> DataFrame: + return self._anchor + + @property + def _internal(self) -> InternalFrame: + return self._psdf._internal.select_column(self._column_label) + + @property + def _column_label(self) -> Optional[Label]: + return self._col_label + + def _update_anchor(self, psdf: DataFrame) -> None: + assert psdf._internal.column_labels == [self._column_label], ( + psdf._internal.column_labels, + [self._column_label], + ) + self._anchor = psdf + object.__setattr__(psdf, "_psseries", {self._column_label: self}) + + def _with_new_scol( + self, scol: PySparkColumn, *, field: Optional[InternalField] = None + ) -> "Series": + """ + Copy pandas-on-Spark Series with the new Spark Column. + + :param scol: the new Spark Column + :return: the copied Series + """ + name = name_like_string(self._column_label) + internal = self._internal.copy( + data_spark_columns=[scol.alias(name)], + data_fields=[ + field if field is None or field.struct_field is None else field.copy(name=name) + ], + ) + return first_series(DataFrame(internal)) + + spark: "SparkIndexOpsMethods" = CachedAccessor( # type: ignore[assignment] + "spark", SparkSeriesMethods + ) + + @property + def dtypes(self) -> Dtype: + """Return the dtype object of the underlying data. + + >>> s = ps.Series(list('abc')) + >>> s.dtype == s.dtypes + True + """ + return self.dtype + + @property + def axes(self) -> List["Index"]: + """ + Return a list of the row axis labels. + + Examples + -------- + + >>> psser = ps.Series([1, 2, 3]) + >>> psser.axes # doctest: +SKIP + [Int64Index([0, 1, 2], dtype='int64')] + """ + return [self.index] + + # Arithmetic Operators + def add(self, other: Any, fill_value: Union[int, str, float] = None) -> "Series": + if fill_value is not None: + if isinstance(other, (int, str, float)): + scol = self.spark.column + scol = F.when(scol.isNull() | F.isnan(scol), fill_value).otherwise(scol) + self = self._with_new_scol(scol) + else: + raise NotImplementedError( + "`fill_value` currently only works when type of `other` is in (int, str, float)" + ) + + return self + other + + add.__doc__ = _flex_doc_SERIES.format( + desc="Addition", + op_name="+", + equiv="series + other", + reverse="radd", + series_examples=_add_example_SERIES, + ) + + def radd(self, other: Any, fill_value: Union[int, str, float] = None) -> "Series": + if fill_value is not None: + if isinstance(other, (int, str, float)): + scol = self.spark.column + scol = F.when(scol.isNull() | F.isnan(scol), fill_value).otherwise(scol) + self = self._with_new_scol(scol) + else: + raise NotImplementedError( + "`fill_value` currently only works when type of `other` is in (int, str, float)" + ) + + return other + self + + radd.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Addition", + op_name="+", + equiv="other + series", + reverse="add", + series_examples=_add_example_SERIES, + ) + + def div(self, other: Any) -> "Series": + return self / other + + div.__doc__ = _flex_doc_SERIES.format( + desc="Floating division", + op_name="/", + equiv="series / other", + reverse="rdiv", + series_examples=_div_example_SERIES, + ) + + divide = div + + def rdiv(self, other: Any) -> "Series": + return other / self + + rdiv.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Floating division", + op_name="/", + equiv="other / series", + reverse="div", + series_examples=_div_example_SERIES, + ) + + def truediv(self, other: Any) -> "Series": + return self / other + + truediv.__doc__ = _flex_doc_SERIES.format( + desc="Floating division", + op_name="/", + equiv="series / other", + reverse="rtruediv", + series_examples=_div_example_SERIES, + ) + + def rtruediv(self, other: Any) -> "Series": + return other / self + + rtruediv.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Floating division", + op_name="/", + equiv="other / series", + reverse="truediv", + series_examples=_div_example_SERIES, + ) + + def mul(self, other: Any) -> "Series": + return self * other + + mul.__doc__ = _flex_doc_SERIES.format( + desc="Multiplication", + op_name="*", + equiv="series * other", + reverse="rmul", + series_examples=_mul_example_SERIES, + ) + + multiply = mul + + def rmul(self, other: Any) -> "Series": + return other * self + + rmul.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Multiplication", + op_name="*", + equiv="other * series", + reverse="mul", + series_examples=_mul_example_SERIES, + ) + + def sub(self, other: Any) -> "Series": + return self - other + + sub.__doc__ = _flex_doc_SERIES.format( + desc="Subtraction", + op_name="-", + equiv="series - other", + reverse="rsub", + series_examples=_sub_example_SERIES, + ) + + subtract = sub + + def rsub(self, other: Any) -> "Series": + return other - self + + rsub.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Subtraction", + op_name="-", + equiv="other - series", + reverse="sub", + series_examples=_sub_example_SERIES, + ) + + def mod(self, other: Any) -> "Series": + return self % other + + mod.__doc__ = _flex_doc_SERIES.format( + desc="Modulo", + op_name="%", + equiv="series % other", + reverse="rmod", + series_examples=_mod_example_SERIES, + ) + + def rmod(self, other: Any) -> "Series": + return other % self + + rmod.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Modulo", + op_name="%", + equiv="other % series", + reverse="mod", + series_examples=_mod_example_SERIES, + ) + + def pow(self, other: Any) -> "Series": + return self**other + + pow.__doc__ = _flex_doc_SERIES.format( + desc="Exponential power of series", + op_name="**", + equiv="series ** other", + reverse="rpow", + series_examples=_pow_example_SERIES, + ) + + def rpow(self, other: Any) -> "Series": + return other**self + + rpow.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Exponential power", + op_name="**", + equiv="other ** series", + reverse="pow", + series_examples=_pow_example_SERIES, + ) + + def floordiv(self, other: Any) -> "Series": + return self // other + + floordiv.__doc__ = _flex_doc_SERIES.format( + desc="Integer division", + op_name="//", + equiv="series // other", + reverse="rfloordiv", + series_examples=_floordiv_example_SERIES, + ) + + def rfloordiv(self, other: Any) -> "Series": + return other // self + + rfloordiv.__doc__ = _flex_doc_SERIES.format( + desc="Reverse Integer division", + op_name="//", + equiv="other // series", + reverse="floordiv", + series_examples=_floordiv_example_SERIES, + ) + + # create accessor for pandas-on-Spark specific methods. + pandas_on_spark = CachedAccessor("pandas_on_spark", PandasOnSparkSeriesMethods) + + # keep the name "koalas" for backward compatibility. + koalas = CachedAccessor("koalas", PandasOnSparkSeriesMethods) + + # Comparison Operators + def eq(self, other: Any) -> "Series": + """ + Compare if the current value is equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a == 1 + a True + b False + c False + d False + Name: a, dtype: bool + + >>> df.b.eq(1) + a True + b False + c True + d False + Name: b, dtype: bool + """ + return self == other + + equals = eq + + def gt(self, other: Any) -> "Series": + """ + Compare if the current value is greater than the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a > 1 + a False + b True + c True + d True + Name: a, dtype: bool + + >>> df.b.gt(1) + a False + b False + c False + d False + Name: b, dtype: bool + """ + return self > other + + def ge(self, other: Any) -> "Series": + """ + Compare if the current value is greater than or equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a >= 2 + a False + b True + c True + d True + Name: a, dtype: bool + + >>> df.b.ge(2) + a False + b False + c False + d False + Name: b, dtype: bool + """ + return self >= other + + def lt(self, other: Any) -> "Series": + """ + Compare if the current value is less than the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a < 1 + a False + b False + c False + d False + Name: a, dtype: bool + + >>> df.b.lt(2) + a True + b False + c True + d False + Name: b, dtype: bool + """ + return self < other + + def le(self, other: Any) -> "Series": + """ + Compare if the current value is less than or equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a <= 2 + a True + b True + c False + d False + Name: a, dtype: bool + + >>> df.b.le(2) + a True + b False + c True + d False + Name: b, dtype: bool + """ + return self <= other + + def ne(self, other: Any) -> "Series": + """ + Compare if the current value is not equal to the other. + + >>> df = ps.DataFrame({'a': [1, 2, 3, 4], + ... 'b': [1, np.nan, 1, np.nan]}, + ... index=['a', 'b', 'c', 'd'], columns=['a', 'b']) + + >>> df.a != 1 + a False + b True + c True + d True + Name: a, dtype: bool + + >>> df.b.ne(1) + a False + b True + c False + d True + Name: b, dtype: bool + """ + return self != other + + def divmod(self, other: Any) -> Tuple["Series", "Series"]: + """ + Return Integer division and modulo of series and other, element-wise + (binary operator `divmod`). + + Parameters + ---------- + other : Series or scalar value + + Returns + ------- + 2-Tuple of Series + The result of the operation. + + See Also + -------- + Series.rdivmod + """ + return self.floordiv(other), self.mod(other) + + def rdivmod(self, other: Any) -> Tuple["Series", "Series"]: + """ + Return Integer division and modulo of series and other, element-wise + (binary operator `rdivmod`). + + Parameters + ---------- + other : Series or scalar value + + Returns + ------- + 2-Tuple of Series + The result of the operation. + + See Also + -------- + Series.divmod + """ + return self.rfloordiv(other), self.rmod(other) + + def between(self, left: Any, right: Any, inclusive: Union[bool, str] = "both") -> "Series": + """ + Return boolean Series equivalent to left <= series <= right. + This function returns a boolean vector containing `True` wherever the + corresponding Series element is between the boundary values `left` and + `right`. NA values are treated as `False`. + + Parameters + ---------- + left : scalar or list-like + Left boundary. + right : scalar or list-like + Right boundary. + inclusive : {"both", "neither", "left", "right"} or boolean. "both" by default. + Include boundaries. Whether to set each bound as closed or open. + Booleans are deprecated in favour of `both` or `neither`. + + Returns + ------- + Series + Series representing whether each element is between left and + right (inclusive). + + See Also + -------- + Series.gt : Greater than of series and other. + Series.lt : Less than of series and other. + + Notes + ----- + This function is equivalent to ``(left <= ser) & (ser <= right)`` + + Examples + -------- + >>> s = ps.Series([2, 0, 4, 8, np.nan]) + + Boundary values are included by default: + + >>> s.between(0, 4) + 0 True + 1 True + 2 True + 3 False + 4 False + dtype: bool + + With `inclusive` set to "neither" boundary values are excluded: + + >>> s.between(0, 4, inclusive="neither") + 0 True + 1 False + 2 False + 3 False + 4 False + dtype: bool + + With `inclusive` set to "right" only right boundary value is included: + + >>> s.between(0, 4, inclusive="right") + 0 True + 1 False + 2 True + 3 False + 4 False + dtype: bool + + With `inclusive` set to "left" only left boundary value is included: + + >>> s.between(0, 4, inclusive="left") + 0 True + 1 True + 2 False + 3 False + 4 False + dtype: bool + + `left` and `right` can be any scalar value: + + >>> s = ps.Series(['Alice', 'Bob', 'Carol', 'Eve']) + >>> s.between('Anna', 'Daniel') + 0 False + 1 True + 2 True + 3 False + dtype: bool + """ + if inclusive is True or inclusive is False: + warnings.warn( + "Boolean inputs to the `inclusive` argument are deprecated in " + "favour of `both` or `neither`.", + FutureWarning, + ) + if inclusive: + inclusive = "both" + else: + inclusive = "neither" + + if inclusive == "both": + lmask = self >= left + rmask = self <= right + elif inclusive == "left": + lmask = self >= left + rmask = self < right + elif inclusive == "right": + lmask = self > left + rmask = self <= right + elif inclusive == "neither": + lmask = self > left + rmask = self < right + else: + raise ValueError( + "Inclusive has to be either string of 'both'," "'left', 'right', or 'neither'." + ) + + return lmask & rmask + + def cov(self, other: "Series", min_periods: Optional[int] = None, ddof: int = 1) -> float: + """ + Compute covariance with Series, excluding missing values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + other : Series + Series with which to compute the covariance. + min_periods : int, optional + Minimum number of observations needed 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. + + .. versionadded:: 3.4.0 + + Returns + ------- + float + Covariance between Series and other + + Examples + -------- + >>> from pyspark.pandas.config import set_option, reset_option + >>> s1 = ps.Series([0.90010907, 0.13484424, 0.62036035]) + >>> s2 = ps.Series([0.12528585, 0.26962463, 0.51111198]) + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.cov(s2) + -0.016857... + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.cov(s2, ddof=2) + -0.033715... + """ + if not isinstance(other, Series): + raise TypeError("unsupported type: %s" % type(other)) + if not np.issubdtype(self.dtype, np.number): # type: ignore[arg-type] + raise TypeError("unsupported dtype: %s" % self.dtype) + if not np.issubdtype(other.dtype, np.number): # type: ignore[arg-type] + raise TypeError("unsupported dtype: %s" % other.dtype) + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + min_periods = 1 if min_periods is None else min_periods + + if same_anchor(self, other): + sdf = self._internal.spark_frame.select(self.spark.column, other.spark.column) + else: + combined = combine_frames(self.to_frame(), other.to_frame()) + sdf = combined._internal.spark_frame.select(*combined._internal.data_spark_columns) + + sdf = sdf.dropna() + + if len(sdf.head(min_periods)) < min_periods: + return np.nan + else: + sdf = sdf.select(SF.covar(F.col(sdf.columns[0]), F.col(sdf.columns[1]), ddof)) + return sdf.head(1)[0][0] + + # TODO: NaN and None when ``arg`` is an empty dict + # TODO: Support ps.Series ``arg`` + def map( + self, arg: Union[Dict, Callable[[Any], Any], pd.Series], na_action: Optional[str] = None + ) -> "Series": + """ + Map values of Series according to input correspondence. + + Used for substituting each value in a Series with another value, + that may be derived from a function, a ``dict``. + + .. note:: make sure the size of the dictionary is not huge because it could + downgrade the performance or throw OutOfMemoryError due to a huge + expression within Spark. Consider the input as a function as an + alternative instead in this case. + + Parameters + ---------- + arg : function, dict or pd.Series + Mapping correspondence. + na_action : + If `ignore`, propagate NA values, without passing them to the mapping correspondence. + + Returns + ------- + Series + Same index as caller. + + See Also + -------- + Series.apply : For applying more complex functions on a Series. + DataFrame.applymap : Apply a function element-wise on a whole DataFrame. + + Notes + ----- + When ``arg`` is a dictionary, values in Series that are not in the + dictionary (as keys) is converted to ``None``. However, if the + dictionary is a ``dict`` subclass that defines ``__missing__`` (i.e. + provides a method for default values), then this default is used + rather than ``None``. + + Examples + -------- + >>> s = ps.Series(['cat', 'dog', None, 'rabbit']) + >>> s + 0 cat + 1 dog + 2 None + 3 rabbit + dtype: object + + ``map`` accepts a ``dict``. Values that are not found + in the ``dict`` are converted to ``None``, unless the dict has a default + value (e.g. ``defaultdict``): + + >>> s.map({'cat': 'kitten', 'dog': 'puppy'}) + 0 kitten + 1 puppy + 2 None + 3 None + dtype: object + + It also accepts a pandas Series: + + >>> pser = pd.Series(['kitten', 'puppy'], index=['cat', 'dog']) + >>> s.map(pser) + 0 kitten + 1 puppy + 2 None + 3 None + dtype: object + + It also accepts a function: + + >>> def format(x) -> str: + ... return 'I am a {}'.format(x) + + >>> s.map(format) + 0 I am a cat + 1 I am a dog + 2 I am a None + 3 I am a rabbit + dtype: object + + To avoid applying the function to missing values (and keep them as NaN) + na_action='ignore' can be used: + + >>> s.map('I am a {}'.format, na_action='ignore') + 0 I am a cat + 1 I am a dog + 2 None + 3 I am a rabbit + dtype: object + """ + if isinstance(arg, (dict, pd.Series)): + is_start = True + # In case dictionary is empty. + current = F.when(F.lit(False), F.lit(None).cast(self.spark.data_type)) + + for to_replace, value in arg.items(): + if is_start: + current = F.when(self.spark.column == F.lit(to_replace), value) + is_start = False + else: + current = current.when(self.spark.column == F.lit(to_replace), value) + + if hasattr(arg, "__missing__"): + tmp_val = arg[np._NoValue] # type: ignore[attr-defined] + # Remove in case it's set in defaultdict. + del arg[np._NoValue] # type: ignore[attr-defined] + current = current.otherwise(F.lit(tmp_val)) + else: + current = current.otherwise(F.lit(None).cast(self.spark.data_type)) + return self._with_new_scol(current) + else: + return self.pandas_on_spark.transform_batch(lambda pser: pser.map(arg, na_action)) + + @property + def shape(self) -> Tuple[int]: + """Return a tuple of the shape of the underlying data.""" + return (len(self),) + + @property + def name(self) -> Name: + """Return name of the Series.""" + name = self._column_label + if name is not None and len(name) == 1: + return name[0] + else: + return name + + @name.setter + def name(self, name: Name) -> None: + self.rename(name, inplace=True) + + # TODO: Currently, changing index labels taking dictionary/Series is not supported. + def rename( + self, index: Optional[Union[Name, Callable[[Any], Any]]] = None, **kwargs: Any + ) -> "Series": + """ + Alter Series index labels or name. + + Parameters + ---------- + index : scalar or function, optional + Functions are transformations to apply to the index. + Scalar will alter the Series.name attribute. + + inplace : bool, default False + Whether to return a new Series. If True then value of copy is + ignored. + + Returns + ------- + Series + Series with index labels or name altered. + + Examples + -------- + + >>> s = ps.Series([1, 2, 3]) + >>> s + 0 1 + 1 2 + 2 3 + dtype: int64 + + >>> s.rename("my_name") # scalar, changes Series.name + 0 1 + 1 2 + 2 3 + Name: my_name, dtype: int64 + + >>> s.rename(lambda x: x ** 2) # function, changes labels + 0 1 + 1 2 + 4 3 + dtype: int64 + """ + if index is None: + pass + if callable(index): + if kwargs.get("inplace", False): + raise ValueError("inplace True is not supported yet for a function 'index'") + frame = self.to_frame() + new_index_name = verify_temp_column_name(frame, "__index_name__") + frame[new_index_name] = self.index.map(index) + frame.set_index(new_index_name, inplace=True) + frame.index.name = self.index.name + return first_series(frame).rename(self.name) + elif isinstance(index, (pd.Series, dict)): + raise ValueError("'index' of %s type is not supported yet" % type(index).__name__) + elif not is_hashable(index): + raise TypeError("Series.name must be a hashable type") + elif not isinstance(index, tuple): + index = (index,) + name = name_like_string(index) + scol = self.spark.column.alias(name) + field = self._internal.data_fields[0].copy(name=name) + + internal = self._internal.copy( + column_labels=[index], + data_spark_columns=[scol], + data_fields=[field], + column_label_names=None, + ) + psdf: DataFrame = DataFrame(internal) + + if kwargs.get("inplace", False): + self._col_label = index + self._update_anchor(psdf) + return self + else: + return first_series(psdf) + + def rename_axis( + self, mapper: Optional[Any] = None, index: Optional[Any] = None, inplace: bool = False + ) -> Optional["Series"]: + """ + Set the name of the axis for the index or columns. + + Parameters + ---------- + mapper, index : scalar, list-like, dict-like or function, optional + A scalar, list-like, dict-like or functions transformations to + apply to the index values. + inplace : bool, default False + Modifies the object directly, instead of creating a new Series. + + Returns + ------- + Series, or None if `inplace` is True. + + See Also + -------- + Series.rename : Alter Series index labels or name. + DataFrame.rename : Alter DataFrame index labels or name. + Index.rename : Set new names on index. + + Examples + -------- + >>> s = ps.Series(["dog", "cat", "monkey"], name="animal") + >>> s # doctest: +NORMALIZE_WHITESPACE + 0 dog + 1 cat + 2 monkey + Name: animal, dtype: object + >>> s.rename_axis("index").sort_index() # doctest: +NORMALIZE_WHITESPACE + index + 0 dog + 1 cat + 2 monkey + Name: animal, dtype: object + + **MultiIndex** + + >>> index = pd.MultiIndex.from_product([['mammal'], + ... ['dog', 'cat', 'monkey']], + ... names=['type', 'name']) + >>> s = ps.Series([4, 4, 2], index=index, name='num_legs') + >>> s # doctest: +NORMALIZE_WHITESPACE + type name + mammal dog 4 + cat 4 + monkey 2 + Name: num_legs, dtype: int64 + >>> s.rename_axis(index={'type': 'class'}).sort_index() # doctest: +NORMALIZE_WHITESPACE + class name + mammal cat 4 + dog 4 + monkey 2 + Name: num_legs, dtype: int64 + >>> s.rename_axis(index=str.upper).sort_index() # doctest: +NORMALIZE_WHITESPACE + TYPE NAME + mammal cat 4 + dog 4 + monkey 2 + Name: num_legs, dtype: int64 + """ + psdf = self.to_frame().rename_axis(mapper=mapper, index=index, inplace=False) + if inplace: + self._update_anchor(psdf) + return None + else: + return first_series(psdf) + + @property + def index(self) -> "ps.Index": + """The index (axis labels) Column of the Series. + + See Also + -------- + Index + """ + return self._psdf.index + + @property + def is_unique(self) -> bool: + """ + Return boolean if values in the object are unique + + Returns + ------- + is_unique : boolean + + >>> ps.Series([1, 2, 3]).is_unique + True + >>> ps.Series([1, 2, 2]).is_unique + False + >>> ps.Series([1, 2, 3, None]).is_unique + True + """ + scol = self.spark.column + + # Here we check: + # 1. the distinct count without nulls and count without nulls for non-null values + # 2. count null values and see if null is a distinct value. + # + # This workaround is to calculate the distinct count including nulls in + # single pass. Note that COUNT(DISTINCT expr) in Spark is designed to ignore nulls. + return self._internal.spark_frame.select( + (F.count(scol) == F.countDistinct(scol)) + & (F.count(F.when(scol.isNull(), 1).otherwise(None)) <= 1) + ).collect()[0][0] + + def reset_index( + self, + level: Optional[Union[int, Name, Sequence[Union[int, Name]]]] = None, + drop: bool = False, + name: Optional[Name] = None, + inplace: bool = False, + ) -> Optional[Union["Series", DataFrame]]: + """ + Generate a new DataFrame or Series with the index reset. + + This is useful when the index needs to be treated as a column, + or when the index is meaningless and needs to be reset + to the default before another operation. + + Parameters + ---------- + level : int, str, tuple, or list, default optional + For a Series with a MultiIndex, only remove the specified levels from the index. + Removes all levels by default. + drop : bool, default False + Just reset the index, without inserting it as a column in the new DataFrame. + name : object, optional + The name to use for the column containing the original Series values. + Uses self.name by default. This argument is ignored when drop is True. + inplace : bool, default False + Modify the Series in place (do not create a new object). + + Returns + ------- + Series or DataFrame + When `drop` is False (the default), a DataFrame is returned. + The newly created columns will come first in the DataFrame, + followed by the original Series values. + When `drop` is True, a `Series` is returned. + In either case, if ``inplace=True``, no value is returned. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4], index=pd.Index(['a', 'b', 'c', 'd'], name='idx')) + + Generate a DataFrame with default index. + + >>> s.reset_index() + idx 0 + 0 a 1 + 1 b 2 + 2 c 3 + 3 d 4 + + To specify the name of the new column use `name`. + + >>> s.reset_index(name='values') + idx values + 0 a 1 + 1 b 2 + 2 c 3 + 3 d 4 + + To generate a new Series with the default set `drop` to True. + + >>> s.reset_index(drop=True) + 0 1 + 1 2 + 2 3 + 3 4 + dtype: int64 + + To update the Series in place, without generating a new one + set `inplace` to True. Note that it also requires ``drop=True``. + + >>> s.reset_index(inplace=True, drop=True) + >>> s + 0 1 + 1 2 + 2 3 + 3 4 + dtype: int64 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + if inplace and not drop: + raise TypeError("Cannot reset_index inplace on a Series to create a DataFrame") + + if drop: + psdf = self._psdf[[self.name]] + else: + psser = self + if name is not None: + psser = psser.rename(name) + psdf = psser.to_frame() + psdf = psdf.reset_index(level=level, drop=drop) + if drop: + if inplace: + self._update_anchor(psdf) + return None + else: + return first_series(psdf) + else: + return psdf + + def to_frame(self, name: Optional[Name] = None) -> DataFrame: + """ + Convert Series to DataFrame. + + Parameters + ---------- + name : object, default None + The passed name should substitute for the series name (if it has + one). + + Returns + ------- + DataFrame + DataFrame representation of Series. + + Examples + -------- + >>> s = ps.Series(["a", "b", "c"]) + >>> s.to_frame() + 0 + 0 a + 1 b + 2 c + + >>> s = ps.Series(["a", "b", "c"], name="vals") + >>> s.to_frame() + vals + 0 a + 1 b + 2 c + """ + if name is not None: + renamed = self.rename(name) + elif self._column_label is None: + renamed = self.rename(DEFAULT_SERIES_NAME) + else: + renamed = self + return DataFrame(renamed._internal) + + to_dataframe = to_frame + + def to_string( + self, + buf: Optional[IO[str]] = None, + na_rep: str = "NaN", + float_format: Optional[Callable[[float], str]] = None, + header: bool = True, + index: bool = True, + length: bool = False, + dtype: bool = False, + name: bool = False, + max_rows: Optional[int] = None, + ) -> Optional[str]: + """ + Render a string representation of the Series. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. If the input + is large, set max_rows parameter. + + Parameters + ---------- + buf : StringIO-like, optional + buffer to write to + na_rep : string, optional + string representation of NAN to use, default 'NaN' + float_format : one-parameter function, optional + formatter function to apply to columns' elements if they are floats + default None + header : boolean, default True + Add the Series header (index name) + index : bool, optional + Add index (row) labels, default True + length : boolean, default False + Add the Series length + dtype : boolean, default False + Add the Series dtype + name : boolean, default False + Add the Series name if not None + max_rows : int, optional + Maximum number of rows to show before truncating. If None, show + all. + + Returns + ------- + formatted : string (if not buffer passed) + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], columns=['dogs', 'cats']) + >>> print(df['dogs'].to_string()) + 0 0.2 + 1 0.0 + 2 0.6 + 3 0.2 + + >>> print(df['dogs'].to_string(max_rows=2)) + 0 0.2 + 1 0.0 + """ + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + if max_rows is not None: + psseries = self.head(max_rows) + else: + psseries = self + + return validate_arguments_and_invoke_function( + psseries._to_internal_pandas(), self.to_string, pd.Series.to_string, args + ) + + def to_clipboard(self, excel: bool = True, sep: Optional[str] = None, **kwargs: Any) -> None: + # Docstring defined below by reusing DataFrame.to_clipboard's. + args = locals() + psseries = self + + return validate_arguments_and_invoke_function( + psseries._to_internal_pandas(), self.to_clipboard, pd.Series.to_clipboard, args + ) + + to_clipboard.__doc__ = DataFrame.to_clipboard.__doc__ + + def to_dict(self, into: Type = dict) -> Mapping: + """ + Convert Series to {label -> value} dict or dict-like object. + + .. note:: This method should only be used if the resulting pandas DataFrame is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + into : class, default dict + The collections.abc.Mapping subclass to use as the return + object. Can be the actual class or an empty + instance of the mapping type you want. If you want a + collections.defaultdict, you must pass it initialized. + + Returns + ------- + collections.abc.Mapping + Key-value representation of Series. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4]) + >>> s_dict = s.to_dict() + >>> sorted(s_dict.items()) + [(0, 1), (1, 2), (2, 3), (3, 4)] + + >>> from collections import OrderedDict, defaultdict + >>> s.to_dict(OrderedDict) + OrderedDict([(0, 1), (1, 2), (2, 3), (3, 4)]) + + >>> dd = defaultdict(list) + >>> s.to_dict(dd) # doctest: +ELLIPSIS + defaultdict(, {...}) + """ + # Make sure locals() call is at the top of the function so we don't capture local variables. + args = locals() + psseries = self + return validate_arguments_and_invoke_function( + psseries._to_internal_pandas(), self.to_dict, pd.Series.to_dict, args + ) + + def to_latex( + self, + buf: Optional[IO[str]] = None, + columns: Optional[List[Name]] = None, + col_space: Optional[int] = None, + header: bool = True, + index: bool = True, + na_rep: str = "NaN", + formatters: Optional[ + Union[List[Callable[[Any], str]], Dict[Name, Callable[[Any], str]]] + ] = None, + float_format: Optional[Callable[[float], str]] = None, + sparsify: Optional[bool] = None, + index_names: bool = True, + bold_rows: bool = False, + column_format: Optional[str] = None, + longtable: Optional[bool] = None, + escape: Optional[bool] = None, + encoding: Optional[str] = None, + decimal: str = ".", + multicolumn: Optional[bool] = None, + multicolumn_format: Optional[str] = None, + multirow: Optional[bool] = None, + ) -> Optional[str]: + warnings.warn( + "Argument `col_space` will be removed in 4.0.0.", + FutureWarning, + ) + + args = locals() + psseries = self + return validate_arguments_and_invoke_function( + psseries._to_internal_pandas(), self.to_latex, pd.Series.to_latex, args + ) + + to_latex.__doc__ = DataFrame.to_latex.__doc__ + + def to_pandas(self) -> pd.Series: + """ + Return a pandas Series. + + .. note:: This method should only be used if the resulting pandas object is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], columns=['dogs', 'cats']) + >>> df['dogs'].to_pandas() + 0 0.2 + 1 0.0 + 2 0.6 + 3 0.2 + Name: dogs, dtype: float64 + """ + log_advice( + "`to_pandas` loads all data into the driver's memory. " + "It should only be used if the resulting pandas Series is expected to be small." + ) + return self._to_pandas() + + def _to_pandas(self) -> pd.Series: + """ + Same as `to_pandas()`, without issuing the advice log for internal usage. + """ + return self._to_internal_pandas().copy() + + def to_list(self) -> List: + """ + Return a list of the values. + + These are each a scalar type, which is a Python scalar + (for str, int, float) or a pandas scalar + (for Timestamp/Timedelta/Interval/Period) + + .. note:: This method should only be used if the resulting list is expected + to be small, as all the data is loaded into the driver's memory. + + """ + log_advice( + "`to_list` loads all data into the driver's memory. " + "It should only be used if the resulting list is expected to be small." + ) + return self._to_internal_pandas().tolist() + + tolist = to_list + + def duplicated(self, keep: Union[bool, str] = "first") -> "Series": + """ + Indicate duplicate Series values. + + Duplicated values are indicated as ``True`` values in the resulting + Series. Either all duplicates, all except the first or all except the + last occurrence of duplicates can be indicated. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + keep : {'first', 'last', False}, default 'first' + Method to handle marking duplicates: + - 'first' : Mark duplicates as ``True`` except for the first occurrence. + - 'last' : Mark duplicates as ``True`` except for the last occurrence. + - ``False`` : Mark all duplicates as ``True``. + + Returns + ------- + Series + Series indicating whether each value has occurred in the + preceding values + + See Also + -------- + Index.drop_duplicates : Remove duplicate values from Index. + DataFrame.duplicated : Equivalent method on DataFrame. + Series.drop_duplicates : Remove duplicate values from Series. + + Examples + -------- + By default, for each set of duplicated values, the first occurrence is + set on False and all others on True: + + >>> animals = ps.Series(['lama', 'cow', 'lama', 'beetle', 'lama']) + >>> animals.duplicated().sort_index() + 0 False + 1 False + 2 True + 3 False + 4 True + dtype: bool + + which is equivalent to + + >>> animals.duplicated(keep='first').sort_index() + 0 False + 1 False + 2 True + 3 False + 4 True + dtype: bool + + By using 'last', the last occurrence of each set of duplicated values + is set on False and all others on True: + + >>> animals.duplicated(keep='last').sort_index() + 0 True + 1 False + 2 True + 3 False + 4 False + dtype: bool + + By setting keep on ``False``, all duplicates are True: + + >>> animals.duplicated(keep=False).sort_index() + 0 True + 1 False + 2 True + 3 False + 4 True + dtype: bool + """ + return self._psdf[[self.name]].duplicated(keep=keep).rename(self.name) + + def drop_duplicates( + self, keep: Union[bool, str] = "first", inplace: bool = False + ) -> Optional["Series"]: + """ + Return Series with duplicate values removed. + + Parameters + ---------- + keep : {'first', 'last', ``False``}, default 'first' + Method to handle dropping duplicates: + - 'first' : Drop duplicates except for the first occurrence. + - 'last' : Drop duplicates except for the last occurrence. + - ``False`` : Drop all duplicates. + inplace : bool, default ``False`` + If ``True``, performs operation inplace and returns None. + + Returns + ------- + Series + Series with duplicates dropped. + + Examples + -------- + Generate a Series with duplicated entries. + + >>> s = ps.Series(['lama', 'cow', 'lama', 'beetle', 'lama', 'hippo'], + ... name='animal') + >>> s.sort_index() + 0 lama + 1 cow + 2 lama + 3 beetle + 4 lama + 5 hippo + Name: animal, dtype: object + + With the 'keep' parameter, the selection behavior of duplicated values + can be changed. The value 'first' keeps the first occurrence for each + set of duplicated entries. The default value of keep is 'first'. + + >>> s.drop_duplicates().sort_index() + 0 lama + 1 cow + 3 beetle + 5 hippo + Name: animal, dtype: object + + The value 'last' for parameter 'keep' keeps the last occurrence for + each set of duplicated entries. + + >>> s.drop_duplicates(keep='last').sort_index() + 1 cow + 3 beetle + 4 lama + 5 hippo + Name: animal, dtype: object + + The value ``False`` for parameter 'keep' discards all sets of + duplicated entries. Setting the value of 'inplace' to ``True`` performs + the operation inplace and returns ``None``. + + >>> s.drop_duplicates(keep=False, inplace=True) + >>> s.sort_index() + 1 cow + 3 beetle + 5 hippo + Name: animal, dtype: object + """ + inplace = validate_bool_kwarg(inplace, "inplace") + psdf = self._psdf[[self.name]].drop_duplicates(keep=keep) + + if inplace: + self._update_anchor(psdf) + return None + else: + return first_series(psdf) + + def reindex(self, index: Optional[Any] = None, fill_value: Optional[Any] = None) -> "Series": + """ + Conform Series to new index with optional filling logic, placing + NA/NaN in locations having no value in the previous index. A new object + is produced. + + Parameters + ---------- + index: array-like, optional + New labels / index to conform to, should be specified using keywords. + Preferably an Index object to avoid duplicating data + fill_value : scalar, default np.NaN + Value to use for missing values. Defaults to NaN, but can be any + "compatible" value. + + Returns + ------- + Series with changed index. + + See Also + -------- + Series.reset_index : Remove row labels or move them to new columns. + + Examples + -------- + + Create a series with some fictional data. + + >>> index = ['Firefox', 'Chrome', 'Safari', 'IE10', 'Konqueror'] + >>> ser = ps.Series([200, 200, 404, 404, 301], + ... index=index, name='http_status') + >>> ser + Firefox 200 + Chrome 200 + Safari 404 + IE10 404 + Konqueror 301 + Name: http_status, dtype: int64 + + Create a new index and reindex the Series. By default + values in the new index that do not have corresponding + records in the Series are assigned ``NaN``. + + >>> new_index= ['Safari', 'Iceweasel', 'Comodo Dragon', 'IE10', + ... 'Chrome'] + >>> ser.reindex(new_index).sort_index() + Chrome 200.0 + Comodo Dragon NaN + IE10 404.0 + Iceweasel NaN + Safari 404.0 + Name: http_status, dtype: float64 + + We can fill in the missing values by passing a value to + the keyword ``fill_value``. + + >>> ser.reindex(new_index, fill_value=0).sort_index() + Chrome 200 + Comodo Dragon 0 + IE10 404 + Iceweasel 0 + Safari 404 + Name: http_status, dtype: int64 + + To further illustrate the filling functionality in + ``reindex``, we will create a Series with a + monotonically increasing index (for example, a sequence + of dates). + + >>> date_index = pd.date_range('1/1/2010', periods=6, freq='D') + >>> ser2 = ps.Series([100, 101, np.nan, 100, 89, 88], + ... name='prices', index=date_index) + >>> ser2.sort_index() + 2010-01-01 100.0 + 2010-01-02 101.0 + 2010-01-03 NaN + 2010-01-04 100.0 + 2010-01-05 89.0 + 2010-01-06 88.0 + Name: prices, dtype: float64 + + Suppose we decide to expand the series to cover a wider + date range. + + >>> date_index2 = pd.date_range('12/29/2009', periods=10, freq='D') + >>> ser2.reindex(date_index2).sort_index() + 2009-12-29 NaN + 2009-12-30 NaN + 2009-12-31 NaN + 2010-01-01 100.0 + 2010-01-02 101.0 + 2010-01-03 NaN + 2010-01-04 100.0 + 2010-01-05 89.0 + 2010-01-06 88.0 + 2010-01-07 NaN + Name: prices, dtype: float64 + """ + + return first_series(self.to_frame().reindex(index=index, fill_value=fill_value)).rename( + self.name + ) + + def reindex_like(self, other: Union["Series", "DataFrame"]) -> "Series": + """ + Return a Series with matching indices as other object. + + Conform the object to the same index on all axes. Places NA/NaN in locations + having no value in the previous index. + + Parameters + ---------- + other : Series or DataFrame + Its row and column indices are used to define the new indices + of this object. + + Returns + ------- + Series + Series with changed indices on each axis. + + See Also + -------- + DataFrame.set_index : Set row labels. + DataFrame.reset_index : Remove row labels or move them to new columns. + DataFrame.reindex : Change to new indices or expand indices. + + Notes + ----- + Same as calling + ``.reindex(index=other.index, ...)``. + + Examples + -------- + + >>> s1 = ps.Series([24.3, 31.0, 22.0, 35.0], + ... index=pd.date_range(start='2014-02-12', + ... end='2014-02-15', freq='D'), + ... name="temp_celsius") + >>> s1 + 2014-02-12 24.3 + 2014-02-13 31.0 + 2014-02-14 22.0 + 2014-02-15 35.0 + Name: temp_celsius, dtype: float64 + + >>> s2 = ps.Series(["low", "low", "medium"], + ... index=pd.DatetimeIndex(['2014-02-12', '2014-02-13', + ... '2014-02-15']), + ... name="winspeed") + >>> s2 + 2014-02-12 low + 2014-02-13 low + 2014-02-15 medium + Name: winspeed, dtype: object + + >>> s2.reindex_like(s1).sort_index() + 2014-02-12 low + 2014-02-13 low + 2014-02-14 None + 2014-02-15 medium + Name: winspeed, dtype: object + """ + if isinstance(other, (Series, DataFrame)): + return self.reindex(index=other.index) + else: + raise TypeError("other must be a pandas-on-Spark Series or DataFrame") + + def fillna( + self, + value: Optional[Any] = None, + method: Optional[str] = None, + axis: Optional[Axis] = None, + inplace: bool = False, + limit: Optional[int] = None, + ) -> Optional["Series"]: + """Fill NA/NaN values. + + .. note:: the current implementation of 'method' parameter in fillna uses Spark's Window + without specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + value : scalar, dict, Series + Value to use to fill holes. alternately a dict/Series of values + specifying which value to use for each column. + DataFrame is not supported. + method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None + Method to use for filling holes in reindexed Series pad / ffill: propagate last valid + observation forward to next valid backfill / bfill: + use NEXT valid observation to fill gap + axis : {0 or `index`} + 1 and `columns` are not supported. + inplace : boolean, default False + Fill in place (do not create a new object) + limit : int, default None + If method is specified, this is the maximum number of consecutive NaN values to + forward/backward fill. In other words, if there is a gap with more than this number of + consecutive NaNs, it will only be partially filled. If method is not specified, + this is the maximum number of entries along the entire axis where NaNs will be filled. + Must be greater than 0 if not None + + Returns + ------- + Series + Series with NA entries filled. + + Examples + -------- + >>> s = ps.Series([np.nan, 2, 3, 4, np.nan, 6], name='x') + >>> s + 0 NaN + 1 2.0 + 2 3.0 + 3 4.0 + 4 NaN + 5 6.0 + Name: x, dtype: float64 + + Replace all NaN elements with 0s. + + >>> s.fillna(0) + 0 0.0 + 1 2.0 + 2 3.0 + 3 4.0 + 4 0.0 + 5 6.0 + Name: x, dtype: float64 + + We can also propagate non-null values forward or backward. + + >>> s.fillna(method='ffill') + 0 NaN + 1 2.0 + 2 3.0 + 3 4.0 + 4 4.0 + 5 6.0 + Name: x, dtype: float64 + + >>> s = ps.Series([np.nan, 'a', 'b', 'c', np.nan], name='x') + >>> s.fillna(method='ffill') + 0 None + 1 a + 2 b + 3 c + 4 c + Name: x, dtype: object + """ + psser = self._fillna(value=value, method=method, axis=axis, limit=limit) + + if method is not None: + psser = DataFrame(psser._psdf._internal.resolved_copy)._psser_for(self._column_label) + + inplace = validate_bool_kwarg(inplace, "inplace") + if inplace: + self._psdf._update_internal_frame(psser._psdf._internal, check_same_anchor=False) + return None + else: + return psser.copy() + + def _fillna( + self, + value: Optional[Any] = None, + method: Optional[str] = None, + axis: Optional[Axis] = None, + limit: Optional[int] = None, + part_cols: Sequence["ColumnOrName"] = (), + ) -> "Series": + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError("fillna currently only works for axis=0 or axis='index'") + if (value is None) and (method is None): + raise ValueError("Must specify a fillna 'value' or 'method' parameter.") + if (method is not None) and (method not in ["ffill", "pad", "backfill", "bfill"]): + raise ValueError("Expecting 'pad', 'ffill', 'backfill' or 'bfill'.") + + scol = self.spark.column + + if not self.spark.nullable and not isinstance( + self.spark.data_type, (FloatType, DoubleType) + ): + return self._psdf.copy()._psser_for(self._column_label) + + cond = self.isnull().spark.column + + if value is not None: + if not isinstance(value, (float, int, str, bool)): + raise TypeError("Unsupported type %s" % type(value).__name__) + if limit is not None: + raise NotImplementedError("limit parameter for value is not support now") + scol = F.when(cond, value).otherwise(scol) + else: + if method in ["ffill", "pad"]: + func = F.last + end = Window.currentRow - 1 + if limit is not None: + begin = Window.currentRow - limit + else: + begin = Window.unboundedPreceding + elif method in ["bfill", "backfill"]: + func = F.first + begin = Window.currentRow + 1 + if limit is not None: + end = Window.currentRow + limit + else: + end = Window.unboundedFollowing + + window = ( + Window.partitionBy(*part_cols) + .orderBy(NATURAL_ORDER_COLUMN_NAME) + .rowsBetween(begin, end) + ) + scol = F.when(cond, func(scol, True).over(window)).otherwise(scol) + + return DataFrame( + self._psdf._internal.with_new_spark_column( + self._column_label, scol.alias(name_like_string(self.name)) # TODO: dtype? + ) + )._psser_for(self._column_label) + + def interpolate( + self, + method: str = "linear", + limit: Optional[int] = None, + limit_direction: Optional[str] = None, + limit_area: Optional[str] = None, + ) -> "Series": + return self._interpolate( + method=method, limit=limit, limit_direction=limit_direction, limit_area=limit_area + ) + + def _interpolate( + self, + method: str = "linear", + limit: Optional[int] = None, + limit_direction: Optional[str] = None, + limit_area: Optional[str] = None, + ) -> "Series": + if method not in ["linear"]: + raise NotImplementedError("interpolate currently works only for method='linear'") + if (limit is not None) and (not limit > 0): + raise ValueError("limit must be > 0.") + if (limit_direction is not None) and ( + limit_direction not in ["forward", "backward", "both"] + ): + raise ValueError("invalid limit_direction: '{}'".format(limit_direction)) + if (limit_area is not None) and (limit_area not in ["inside", "outside"]): + raise ValueError("invalid limit_area: '{}'".format(limit_area)) + + if not self.spark.nullable and not isinstance( + self.spark.data_type, (FloatType, DoubleType) + ): + return self._psdf.copy()._psser_for(self._column_label) + + scol = self.spark.column + last_non_null = SF.last_non_null(scol) + null_index = SF.null_index(scol) + + Window = get_window_class() + window_forward = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + last_non_null_forward = last_non_null.over(window_forward) + null_index_forward = null_index.over(window_forward) + + window_backward = Window.orderBy(F.desc(NATURAL_ORDER_COLUMN_NAME)).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + last_non_null_backward = last_non_null.over(window_backward) + null_index_backward = null_index.over(window_backward) + + fill = (last_non_null_backward - last_non_null_forward) / ( + null_index_backward + null_index_forward + ) * null_index_forward + last_non_null_forward + + fill_cond = ~F.isnull(last_non_null_backward) & ~F.isnull(last_non_null_forward) + + pad_head = F.lit(None) + pad_head_cond = F.lit(False) + pad_tail = F.lit(None) + pad_tail_cond = F.lit(False) + + # inputs -> NaN, NaN, 1.0, NaN, NaN, NaN, 5.0, NaN, NaN + if limit_direction is None or limit_direction == "forward": + # outputs -> NaN, NaN, 1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 5.0 + pad_tail = last_non_null_forward + pad_tail_cond = F.isnull(last_non_null_backward) & ~F.isnull(last_non_null_forward) + if limit is not None: + # outputs (limit=1) -> NaN, NaN, 1.0, 2.0, NaN, NaN, 5.0, 5.0, NaN + fill_cond = fill_cond & (null_index_forward <= F.lit(limit)) + pad_tail_cond = pad_tail_cond & (null_index_forward <= F.lit(limit)) + + elif limit_direction == "backward": + # outputs -> 1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, NaN, NaN + pad_head = last_non_null_backward + pad_head_cond = ~F.isnull(last_non_null_backward) & F.isnull(last_non_null_forward) + if limit is not None: + # outputs (limit=1) -> NaN, 1.0, 1.0, NaN, NaN, 4.0, 5.0, NaN, NaN + fill_cond = fill_cond & (null_index_backward <= F.lit(limit)) + pad_head_cond = pad_head_cond & (null_index_backward <= F.lit(limit)) + + else: + # outputs -> 1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 5.0 + pad_head = last_non_null_backward + pad_head_cond = ~F.isnull(last_non_null_backward) & F.isnull(last_non_null_forward) + pad_tail = last_non_null_forward + pad_tail_cond = F.isnull(last_non_null_backward) & ~F.isnull(last_non_null_forward) + if limit is not None: + # outputs (limit=1) -> NaN, 1.0, 1.0, 2.0, NaN, 4.0, 5.0, 5.0, NaN + fill_cond = fill_cond & ( + (null_index_forward <= F.lit(limit)) | (null_index_backward <= F.lit(limit)) + ) + pad_head_cond = pad_head_cond & (null_index_backward <= F.lit(limit)) + pad_tail_cond = pad_tail_cond & (null_index_forward <= F.lit(limit)) + + if limit_area == "inside": + pad_head_cond = F.lit(False) + pad_tail_cond = F.lit(False) + elif limit_area == "outside": + fill_cond = F.lit(False) + + cond = self.isnull().spark.column + scol = ( + F.when(cond & fill_cond, fill) + .when(cond & pad_head_cond, pad_head) + .when(cond & pad_tail_cond, pad_tail) + .otherwise(scol) + ) + + return DataFrame( + self._psdf._internal.with_new_spark_column(self._column_label, scol) # TODO: dtype? + )._psser_for(self._column_label) + + def dropna(self, axis: Axis = 0, inplace: bool = False, **kwargs: Any) -> Optional["Series"]: + """ + Return a new Series with missing values removed. + + Parameters + ---------- + axis : {0 or 'index'}, default 0 + There is only one axis to drop values from. + inplace : bool, default False + If True, do operation inplace and return None. + **kwargs + Not in use. + + Returns + ------- + Series + Series with NA entries dropped from it. + + Examples + -------- + >>> ser = ps.Series([1., 2., np.nan]) + >>> ser + 0 1.0 + 1 2.0 + 2 NaN + dtype: float64 + + Drop NA values from a Series. + + >>> ser.dropna() + 0 1.0 + 1 2.0 + dtype: float64 + + Keep the Series with valid entries in the same variable. + + >>> ser.dropna(inplace=True) + >>> ser + 0 1.0 + 1 2.0 + dtype: float64 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + # TODO: last two examples from pandas produce different results. + psdf = self._psdf[[self.name]].dropna(axis=axis, inplace=False) + if inplace: + self._update_anchor(psdf) + return None + else: + return first_series(psdf) + + def clip( + self, + lower: Union[float, int] = None, + upper: Union[float, int] = None, + inplace: bool = False, + ) -> "Series": + """ + Trim values at input threshold(s). + + Assigns values outside boundary-to-boundary values. + + Parameters + ---------- + lower : float or int, default None + Minimum threshold value. All values below this threshold will be set to it. + upper : float or int, default None + Maximum threshold value. All values above this threshold will be set to it. + inplace : bool, default False + if True, perform operation in-place + + Returns + ------- + Series + Series with the values outside the clip boundaries replaced + + Examples + -------- + >>> psser = ps.Series([0, 2, 4]) + >>> psser + 0 0 + 1 2 + 2 4 + dtype: int64 + + >>> psser.clip(1, 3) + 0 1 + 1 2 + 2 3 + dtype: int64 + + Clip can be performed in-place. + + >>> psser.clip(2, 3, inplace=True) + >>> psser + 0 2 + 1 2 + 2 3 + dtype: int64 + + Notes + ----- + One difference between this implementation and pandas is that running + `pd.Series(['a', 'b']).clip(0, 1)` will crash with "TypeError: '<=' not supported between + instances of 'str' and 'int'" while `ps.Series(['a', 'b']).clip(0, 1)` will output the + original Series, simply ignoring the incompatible types. + """ + if is_list_like(lower) or is_list_like(upper): + raise TypeError( + "List-like value are not supported for 'lower' and 'upper' at the " + "moment" + ) + + if lower is None and upper is None: + return self + + if isinstance(self.spark.data_type, NumericType): + scol = self.spark.column + if lower is not None: + scol = F.when(scol < lower, lower).otherwise(scol) + if upper is not None: + scol = F.when(scol > upper, upper).otherwise(scol) + if inplace: + internal = self._internal.copy( + data_spark_columns=[scol.alias(self._internal.data_spark_column_names[0])], + data_fields=[self._internal.data_fields[0]], + ) + self._psdf._update_internal_frame(internal, check_same_anchor=False) + return None + else: + return self._with_new_scol( + scol.alias(self._internal.data_spark_column_names[0]), + field=self._internal.data_fields[0], + ) + else: + return self + + def drop( + self, + labels: Optional[Union[Name, List[Name]]] = None, + index: Optional[Union[Name, List[Name]]] = None, + columns: Optional[Union[Name, List[Name]]] = None, + level: Optional[int] = None, + inplace: bool = False, + ) -> "Series": + """ + Return Series with specified index labels removed. + + Remove elements of a Series based on specifying the index labels. + When using a multi-index, labels on different levels can be removed by specifying the level. + + Parameters + ---------- + labels : single label or list-like + Index labels to drop. + index : single label or list-like + Redundant for application on Series, but index can be used instead of labels. + columns : single label or list-like + No change is made to the Series; use ‘index’ or ‘labels’ instead. + + .. versionadded:: 3.4.0 + level : int or level name, optional + For MultiIndex, level for which the labels will be removed. + inplace: bool, default False + If True, do operation inplace and return None + + .. versionadded:: 3.4.0 + + Returns + ------- + Series + Series with specified index labels removed. + + See Also + -------- + Series.dropna + + Examples + -------- + >>> s = ps.Series(data=np.arange(3), index=['A', 'B', 'C']) + >>> s + A 0 + B 1 + C 2 + dtype: int64 + + Drop single label A + + >>> s.drop('A') + B 1 + C 2 + dtype: int64 + + Drop labels B and C + + >>> s.drop(labels=['B', 'C']) + A 0 + dtype: int64 + + With 'index' rather than 'labels' returns exactly same result. + + >>> s.drop(index='A') + B 1 + C 2 + dtype: int64 + + >>> s.drop(index=['B', 'C']) + A 0 + dtype: int64 + + With 'columns', no change is made to the Series. + + >>> s.drop(columns=['A']) + A 0 + B 1 + C 2 + dtype: int64 + + With 'inplace=True', do operation inplace and return None. + + >>> s.drop(index=['B', 'C'], inplace=True) + >>> s + A 0 + dtype: int64 + + Also support for MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx) + >>> s + lama speed 45.0 + weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.drop(labels='weight', level=1) + lama speed 45.0 + length 1.2 + cow speed 30.0 + length 1.5 + falcon speed 320.0 + length 0.3 + dtype: float64 + + >>> s.drop(('lama', 'weight')) + lama speed 45.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.drop([('lama', 'speed'), ('falcon', 'weight')]) + lama weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + length 0.3 + dtype: float64 + """ + dropped = self._drop( + labels=labels, index=index, level=level, inplace=inplace, columns=columns + ) + return None if dropped is None else first_series(dropped) + + def _drop( + self, + labels: Optional[Union[Name, List[Name]]] = None, + index: Optional[Union[Name, List[Name]]] = None, + level: Optional[int] = None, + inplace: bool = False, + columns: Optional[Union[Name, List[Name]]] = None, + ) -> Optional[DataFrame]: + if labels is not None: + if columns is not None or index is not None: + raise ValueError("Cannot specify both 'labels' and 'index'/'columns'") + return self._drop(index=labels, level=level, inplace=inplace, columns=columns) + if index is not None: + internal = self._internal + if level is None: + level = 0 + if level >= internal.index_level: + raise ValueError("'level' should be less than the number of indexes") + + if is_name_like_tuple(index): + index_list = [cast(Label, index)] + elif is_name_like_value(index): + index_list = [(index,)] + elif all(is_name_like_value(idxes, allow_tuple=False) for idxes in index): + index_list = [(idex,) for idex in index] + elif not all(is_name_like_tuple(idxes) for idxes in index): + raise ValueError( + "If the given index is a list, it " + "should only contains names as all tuples or all non tuples " + "that contain index names" + ) + else: + index_list = cast(List[Label], index) + + drop_index_scols = [] + for idxes in index_list: + try: + index_scols = [ + internal.index_spark_columns[lvl] == idx + for lvl, idx in enumerate(idxes, level) + ] + except IndexError: + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + internal.index_level, len(idxes) + ) + ) + drop_index_scols.append(reduce(lambda x, y: x & y, index_scols)) + + cond = ~reduce(lambda x, y: x | y, drop_index_scols) + dropped_internal = internal.with_filter(cond) + if inplace: + self._update_anchor(DataFrame(dropped_internal)) + return None + else: + return DataFrame(dropped_internal) + elif columns is not None: + return self._psdf + else: + raise ValueError("Need to specify at least one of 'labels', 'index' or 'columns'") + + def head(self, n: int = 5) -> "Series": + """ + Return the first n rows. + + This function returns the first n rows for the object based on position. + It is useful for quickly testing if your object has the right type of data in it. + + Parameters + ---------- + n : Integer, default = 5 + + Returns + ------- + The first n rows of the caller object. + + Examples + -------- + >>> df = ps.DataFrame({'animal':['alligator', 'bee', 'falcon', 'lion']}) + >>> df.animal.head(2) # doctest: +NORMALIZE_WHITESPACE + 0 alligator + 1 bee + Name: animal, dtype: object + """ + return first_series(self.to_frame().head(n)).rename(self.name) + + def last(self, offset: Union[str, DateOffset]) -> "Series": + """ + Select final periods of time series data based on a date offset. + + When having a Series with dates as index, this function can + select the last few elements based on a date offset. + + Parameters + ---------- + offset : str or DateOffset + The offset length of the data that will be selected. For instance, + '3D' will display all the rows having their index within the last 3 days. + + Returns + ------- + Series + A subset of the caller. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + Examples + -------- + >>> index = pd.date_range('2018-04-09', periods=4, freq='2D') + >>> psser = ps.Series([1, 2, 3, 4], index=index) + >>> psser + 2018-04-09 1 + 2018-04-11 2 + 2018-04-13 3 + 2018-04-15 4 + dtype: int64 + + Get the rows for the last 3 days: + + >>> psser.last('3D') + 2018-04-13 3 + 2018-04-15 4 + dtype: int64 + + Notice the data for 3 last calendar days were returned, not the last + 3 observed days in the dataset, and therefore data for 2018-04-11 was + not returned. + """ + return first_series(self.to_frame().last(offset)).rename(self.name) + + def first(self, offset: Union[str, DateOffset]) -> "Series": + """ + Select first periods of time series data based on a date offset. + + When having a Series with dates as index, this function can + select the first few elements based on a date offset. + + Parameters + ---------- + offset : str or DateOffset + The offset length of the data that will be selected. For instance, + '3D' will display all the rows having their index within the first 3 days. + + Returns + ------- + Series + A subset of the caller. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + Examples + -------- + >>> index = pd.date_range('2018-04-09', periods=4, freq='2D') + >>> psser = ps.Series([1, 2, 3, 4], index=index) + >>> psser + 2018-04-09 1 + 2018-04-11 2 + 2018-04-13 3 + 2018-04-15 4 + dtype: int64 + + Get the rows for the first 3 days: + + >>> psser.first('3D') + 2018-04-09 1 + 2018-04-11 2 + dtype: int64 + + Notice the data for 3 first calendar days were returned, not the first + 3 observed days in the dataset, and therefore data for 2018-04-13 was + not returned. + """ + return first_series(self.to_frame().first(offset)).rename(self.name) + + # TODO: Categorical type isn't supported (due to PySpark's limitation) and + # some doctests related with timestamps were not added. + def unique(self) -> "Series": + """ + Return unique values of Series object. + + Uniques are returned in order of appearance. Hash table-based unique, + therefore does NOT sort. + + .. note:: This method returns newly created Series whereas pandas returns + the unique values as a NumPy array. + + Returns + ------- + Returns the unique values as a Series. + + See Also + -------- + Index.unique + groupby.SeriesGroupBy.unique + + Examples + -------- + >>> psser = ps.Series([2, 1, 3, 3], name='A') + >>> psser.unique().sort_values() + 1 1 + 0 2 + 2 3 + Name: A, dtype: int64 + + >>> ps.Series([pd.Timestamp('2016-01-01') for _ in range(3)]).unique() + 0 2016-01-01 + dtype: datetime64[ns] + + >>> psser.name = ('x', 'a') + >>> psser.unique().sort_values() + 1 1 + 0 2 + 2 3 + Name: (x, a), dtype: int64 + """ + sdf = self._internal.spark_frame.select(self.spark.column).distinct() + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=None, + column_labels=[self._column_label], + data_spark_columns=[scol_for(sdf, self._internal.data_spark_column_names[0])], + data_fields=[self._internal.data_fields[0]], + column_label_names=self._internal.column_label_names, + ) + return first_series(DataFrame(internal)) + + def sort_values( + self, + ascending: bool = True, + inplace: bool = False, + na_position: str = "last", + ignore_index: bool = False, + ) -> Optional["Series"]: + """ + Sort by the values. + + Sort a Series in ascending or descending order by some criterion. + + Parameters + ---------- + ascending : bool or list of bool, default True + Sort ascending vs. descending. Specify list for multiple sort + orders. If this is a list of bools, must match the length of + the by. + inplace : bool, default False + if True, perform operation in-place + na_position : {'first', 'last'}, default 'last' + `first` puts NaNs at the beginning, `last` puts NaNs at the end + ignore_index : bool, default False + If True, the resulting axis will be labeled 0, 1, …, n - 1. + + .. versionadded:: 3.4.0 + + Returns + ------- + sorted_obj : Series ordered by values. + + Examples + -------- + >>> s = ps.Series([np.nan, 1, 3, 10, 5]) + >>> s + 0 NaN + 1 1.0 + 2 3.0 + 3 10.0 + 4 5.0 + dtype: float64 + + Sort values ascending order (default behaviour) + + >>> s.sort_values(ascending=True) + 1 1.0 + 2 3.0 + 4 5.0 + 3 10.0 + 0 NaN + dtype: float64 + + Sort values descending order + + >>> s.sort_values(ascending=False) + 3 10.0 + 4 5.0 + 2 3.0 + 1 1.0 + 0 NaN + dtype: float64 + + Sort values descending order and ignoring index + + >>> s.sort_values(ascending=False, ignore_index=True) + 0 10.0 + 1 5.0 + 2 3.0 + 3 1.0 + 4 NaN + dtype: float64 + + Sort values inplace + + >>> s.sort_values(ascending=False, inplace=True) + >>> s + 3 10.0 + 4 5.0 + 2 3.0 + 1 1.0 + 0 NaN + dtype: float64 + + Sort values putting NAs first + + >>> s.sort_values(na_position='first') + 0 NaN + 1 1.0 + 2 3.0 + 4 5.0 + 3 10.0 + dtype: float64 + + Sort a series of strings + + >>> s = ps.Series(['z', 'b', 'd', 'a', 'c']) + >>> s + 0 z + 1 b + 2 d + 3 a + 4 c + dtype: object + + >>> s.sort_values() + 3 a + 1 b + 4 c + 2 d + 0 z + dtype: object + """ + inplace = validate_bool_kwarg(inplace, "inplace") + psdf = self._psdf[[self.name]]._sort( + by=[self.spark.column], ascending=ascending, na_position=na_position + ) + + if inplace: + if ignore_index: + psdf.reset_index(drop=True, inplace=inplace) + self._update_anchor(psdf) + return None + else: + return first_series(psdf.reset_index(drop=True)) if ignore_index else first_series(psdf) + + def sort_index( + self, + axis: Axis = 0, + level: Optional[Union[int, List[int]]] = None, + ascending: bool = True, + inplace: bool = False, + kind: str = None, + na_position: str = "last", + ignore_index: bool = False, + ) -> Optional["Series"]: + """ + Sort object by labels (along an axis) + + Parameters + ---------- + axis : index, columns to direct sorting. Currently, only axis = 0 is supported. + level : int or level name or list of ints or list of level names + if not None, sort on values in specified index level(s) + ascending : boolean, default True + Sort ascending vs. descending + inplace : bool, default False + if True, perform operation in-place + kind : str, default None + pandas-on-Spark does not allow specifying the sorting algorithm now, + default None + na_position : {‘first’, ‘last’}, default ‘last’ + first puts NaNs at the beginning, last puts NaNs at the end. Not implemented for + MultiIndex. + ignore_index : bool, default False + If True, the resulting axis will be labeled 0, 1, …, n - 1. + + .. versionadded:: 3.4.0 + + Returns + ------- + sorted_obj : Series + + Examples + -------- + >>> s = ps.Series([2, 1, np.nan], index=['b', 'a', np.nan]) + + >>> s.sort_index() # doctest: +SKIP + a 1.0 + b 2.0 + None NaN + dtype: float64 + + >>> s.sort_index(ignore_index=True) + 0 1.0 + 1 2.0 + 2 NaN + dtype: float64 + + >>> s.sort_index(ascending=False) # doctest: +SKIP + b 2.0 + a 1.0 + None NaN + dtype: float64 + + >>> s.sort_index(na_position='first') # doctest: +SKIP + None NaN + a 1.0 + b 2.0 + dtype: float64 + + >>> s.sort_index(inplace=True) + >>> s # doctest: +SKIP + a 1.0 + b 2.0 + None NaN + dtype: float64 + + Multi-index series. + + >>> s = ps.Series(range(4), index=[['b', 'b', 'a', 'a'], [1, 0, 1, 0]], name='0') + + >>> s.sort_index() + a 0 3 + 1 2 + b 0 1 + 1 0 + Name: 0, dtype: int64 + + >>> s.sort_index(level=1) # doctest: +SKIP + a 0 3 + b 0 1 + a 1 2 + b 1 0 + Name: 0, dtype: int64 + + >>> s.sort_index(level=[1, 0]) + a 0 3 + b 0 1 + a 1 2 + b 1 0 + Name: 0, dtype: int64 + """ + inplace = validate_bool_kwarg(inplace, "inplace") + psdf = self._psdf[[self.name]].sort_index( + axis=axis, level=level, ascending=ascending, kind=kind, na_position=na_position + ) + + if inplace: + if ignore_index: + psdf.reset_index(drop=True, inplace=inplace) + self._update_anchor(psdf) + return None + else: + return first_series(psdf.reset_index(drop=True)) if ignore_index else first_series(psdf) + + def swaplevel( + self, i: Union[int, Name] = -2, j: Union[int, Name] = -1, copy: bool = True + ) -> "Series": + """ + Swap levels i and j in a MultiIndex. + Default is to swap the two innermost levels of the index. + + Parameters + ---------- + i, j : int, str + Level of the indices to be swapped. Can pass level name as string. + copy : bool, default True + Whether to copy underlying data. Must be True. + + Returns + ------- + Series + Series with levels swapped in MultiIndex. + + Examples + -------- + >>> midx = pd.MultiIndex.from_arrays([['a', 'b'], [1, 2]], names = ['word', 'number']) + >>> midx # doctest: +SKIP + MultiIndex([('a', 1), + ('b', 2)], + names=['word', 'number']) + >>> psser = ps.Series(['x', 'y'], index=midx) + >>> psser + word number + a 1 x + b 2 y + dtype: object + >>> psser.swaplevel() + number word + 1 a x + 2 b y + dtype: object + >>> psser.swaplevel(0, 1) + number word + 1 a x + 2 b y + dtype: object + >>> psser.swaplevel('number', 'word') + number word + 1 a x + 2 b y + dtype: object + """ + assert copy is True + + return first_series(self.to_frame().swaplevel(i, j, axis=0)).rename(self.name) + + def swapaxes(self, i: Axis, j: Axis, copy: bool = True) -> "Series": + """ + Interchange axes and swap values axes appropriately. + + Parameters + ---------- + i: {0 or 'index', 1 or 'columns'}. The axis to swap. + j: {0 or 'index', 1 or 'columns'}. The axis to swap. + copy : bool, default True. + + Returns + ------- + Series + + Examples + -------- + >>> psser = ps.Series([1, 2, 3], index=["x", "y", "z"]) + >>> psser + x 1 + y 2 + z 3 + dtype: int64 + >>> + >>> psser.swapaxes(0, 0) + x 1 + y 2 + z 3 + dtype: int64 + """ + assert copy is True + + i = validate_axis(i) + j = validate_axis(j) + if not i == j == 0: + raise ValueError("Axis must be 0 for Series") + + return self.copy() + + def add_prefix(self, prefix: str) -> "Series": + """ + Prefix labels with string `prefix`. + + For Series, the row labels are prefixed. + For DataFrame, the column labels are prefixed. + + Parameters + ---------- + prefix : str + The string to add before each label. + + Returns + ------- + Series + New Series with updated labels. + + See Also + -------- + Series.add_suffix: Suffix column labels with string `suffix`. + DataFrame.add_suffix: Suffix column labels with string `suffix`. + DataFrame.add_prefix: Prefix column labels with string `prefix`. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4]) + >>> s + 0 1 + 1 2 + 2 3 + 3 4 + dtype: int64 + + >>> s.add_prefix('item_') + item_0 1 + item_1 2 + item_2 3 + item_3 4 + dtype: int64 + """ + assert isinstance(prefix, str) + internal = self._internal.resolved_copy + sdf = internal.spark_frame.select( + [ + F.concat(F.lit(prefix), index_spark_column).alias(index_spark_column_name) + for index_spark_column, index_spark_column_name in zip( + internal.index_spark_columns, internal.index_spark_column_names + ) + ] + + internal.data_spark_columns + ) + return first_series( + DataFrame(internal.with_new_sdf(sdf, index_fields=([None] * internal.index_level))) + ) + + def add_suffix(self, suffix: str) -> "Series": + """ + Suffix labels with string suffix. + + For Series, the row labels are suffixed. + For DataFrame, the column labels are suffixed. + + Parameters + ---------- + suffix : str + The string to add after each label. + + Returns + ------- + Series + New Series with updated labels. + + See Also + -------- + Series.add_prefix: Prefix row labels with string `prefix`. + DataFrame.add_prefix: Prefix column labels with string `prefix`. + DataFrame.add_suffix: Suffix column labels with string `suffix`. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4]) + >>> s + 0 1 + 1 2 + 2 3 + 3 4 + dtype: int64 + + >>> s.add_suffix('_item') + 0_item 1 + 1_item 2 + 2_item 3 + 3_item 4 + dtype: int64 + """ + assert isinstance(suffix, str) + internal = self._internal.resolved_copy + sdf = internal.spark_frame.select( + [ + F.concat(index_spark_column, F.lit(suffix)).alias(index_spark_column_name) + for index_spark_column, index_spark_column_name in zip( + internal.index_spark_columns, internal.index_spark_column_names + ) + ] + + internal.data_spark_columns + ) + return first_series( + DataFrame(internal.with_new_sdf(sdf, index_fields=([None] * internal.index_level))) + ) + + def autocorr(self, lag: int = 1) -> float: + """ + Compute the lag-N autocorrelation. + + This method computes the Pearson correlation between + the Series and its shifted self. + + .. note:: the current implementation of rank uses Spark's Window without + specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + lag : int, default 1 + Number of lags to apply before performing autocorrelation. + + Returns + ------- + float + The Pearson correlation between self and self.shift(lag). + + See Also + -------- + Series.corr : Compute the correlation between two Series. + Series.shift : Shift index by desired number of periods. + DataFrame.corr : Compute pairwise correlation of columns. + + Notes + ----- + If the Pearson correlation is not well defined return 'NaN'. + + Examples + -------- + >>> s = ps.Series([.2, .0, .6, .2, np.nan, .5, .6]) + >>> s.autocorr() # doctest: +ELLIPSIS + -0.141219... + >>> s.autocorr(0) # doctest: +ELLIPSIS + 1.0... + >>> s.autocorr(2) # doctest: +ELLIPSIS + 0.970725... + >>> s.autocorr(-3) # doctest: +ELLIPSIS + 0.277350... + >>> s.autocorr(5) # doctest: +ELLIPSIS + -1.000000... + >>> s.autocorr(6) # doctest: +ELLIPSIS + nan + + If the Pearson correlation is not well defined, then 'NaN' is returned. + + >>> s = ps.Series([1, 0, 0, 0]) + >>> s.autocorr() + nan + """ + # This implementation is suboptimal because it moves all data to a single partition, + # global sort should be used instead of window, but it should be a start + if not isinstance(lag, int): + raise TypeError("lag should be an int; however, got [%s]" % type(lag).__name__) + + sdf = self._internal.spark_frame + scol = self.spark.column + if lag == 0: + corr = sdf.select(F.corr(scol, scol)).head()[0] + else: + lag_scol = F.lag(scol, lag).over(Window.orderBy(NATURAL_ORDER_COLUMN_NAME)) + lag_col_name = verify_temp_column_name(sdf, "__autocorr_lag_tmp_col__") + corr = ( + sdf.withColumn(lag_col_name, lag_scol) + .select(F.corr(scol, F.col(lag_col_name))) + .head()[0] + ) + return np.nan if corr is None else corr + + def corr( + self, other: "Series", method: str = "pearson", min_periods: Optional[int] = None + ) -> float: + """ + Compute correlation with `other` Series, excluding missing values. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + other : Series + method : {'pearson', 'spearman', 'kendall'} + * pearson : standard correlation coefficient + * spearman : Spearman rank correlation + * kendall : Kendall Tau correlation coefficient + + .. versionchanged:: 3.4.0 + support 'kendall' for method parameter + min_periods : int, optional + Minimum number of observations needed to have a valid result. + + .. versionadded:: 3.4.0 + + Returns + ------- + correlation : float + + Notes + ----- + The complexity of Kendall correlation is O(#row * #row), if the dataset is too + large, sampling ahead of correlation computation is recommended. + + Examples + -------- + >>> df = ps.DataFrame({'s1': [.2, .0, .6, .2], + ... 's2': [.3, .6, .0, .1]}) + >>> s1 = df.s1 + >>> s2 = df.s2 + >>> s1.corr(s2, method='pearson') + -0.85106... + + >>> s1.corr(s2, method='spearman') + -0.94868... + + >>> s1.corr(s2, method='kendall') + -0.91287... + + >>> s1 = ps.Series([1, np.nan, 2, 1, 1, 2, 3]) + >>> s2 = ps.Series([3, 4, 1, 1, 5]) + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.corr(s2, method="pearson") + -0.52223... + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.corr(s2, method="spearman") + -0.54433... + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.corr(s2, method="kendall") + -0.51639... + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.corr(s2, method="kendall", min_periods=5) + nan + """ + if method not in ["pearson", "spearman", "kendall"]: + raise ValueError(f"Invalid method {method}") + if not isinstance(other, Series): + raise TypeError("'other' must be a Series") + if min_periods is not None and not isinstance(min_periods, int): + raise TypeError(f"Invalid min_periods type {type(min_periods).__name__}") + + min_periods = 1 if min_periods is None else min_periods + + if same_anchor(self, other): + combined = self + this = self + that = other + else: + combined = combine_frames(self._psdf, other._psdf) # type: ignore[assignment] + this = combined["this"] + that = combined["that"] + + sdf = combined._internal.spark_frame + index_col_name = verify_temp_column_name(sdf, "__ser_corr_index_temp_column__") + this_scol = this._internal.spark_column_for(this._internal.column_labels[0]) + that_scol = that._internal.spark_column_for(that._internal.column_labels[0]) + + sdf = sdf.select( + F.lit(0).alias(index_col_name), + this_scol.cast("double").alias(CORRELATION_VALUE_1_COLUMN), + that_scol.cast("double").alias(CORRELATION_VALUE_2_COLUMN), + ) + + sdf = compute(sdf=sdf, groupKeys=[index_col_name], method=method).select( + F.when( + F.col(CORRELATION_COUNT_OUTPUT_COLUMN) < min_periods, F.lit(None).cast("double") + ).otherwise(F.col(CORRELATION_CORR_OUTPUT_COLUMN)) + ) + + results = sdf.take(1) + if len(results) == 0: + raise ValueError("attempt to get corr of an empty sequence") + else: + return np.nan if results[0][0] is None else results[0][0] + + def nsmallest(self, n: int = 5) -> "Series": + """ + Return the smallest `n` elements. + + Parameters + ---------- + n : int, default 5 + Return this many ascending sorted values. + + Returns + ------- + Series + The `n` smallest values in the Series, sorted in increasing order. + + See Also + -------- + Series.nlargest: Get the `n` largest elements. + Series.sort_values: Sort Series by values. + Series.head: Return the first `n` rows. + + Notes + ----- + Faster than ``.sort_values().head(n)`` for small `n` relative to + the size of the ``Series`` object. + In pandas-on-Spark, thanks to Spark's lazy execution and query optimizer, + the two would have same performance. + + Examples + -------- + >>> data = [1, 2, 3, 4, np.nan ,6, 7, 8] + >>> s = ps.Series(data) + >>> s + 0 1.0 + 1 2.0 + 2 3.0 + 3 4.0 + 4 NaN + 5 6.0 + 6 7.0 + 7 8.0 + dtype: float64 + + The `n` largest elements where ``n=5`` by default. + + >>> s.nsmallest() + 0 1.0 + 1 2.0 + 2 3.0 + 3 4.0 + 5 6.0 + dtype: float64 + + >>> s.nsmallest(3) + 0 1.0 + 1 2.0 + 2 3.0 + dtype: float64 + """ + return self.sort_values(ascending=True).head(n) + + def nlargest(self, n: int = 5) -> "Series": + """ + Return the largest `n` elements. + + Parameters + ---------- + n : int, default 5 + + Returns + ------- + Series + The `n` largest values in the Series, sorted in decreasing order. + + See Also + -------- + Series.nsmallest: Get the `n` smallest elements. + Series.sort_values: Sort Series by values. + Series.head: Return the first `n` rows. + + Notes + ----- + Faster than ``.sort_values(ascending=False).head(n)`` for small `n` + relative to the size of the ``Series`` object. + + In pandas-on-Spark, thanks to Spark's lazy execution and query optimizer, + the two would have same performance. + + Examples + -------- + >>> data = [1, 2, 3, 4, np.nan ,6, 7, 8] + >>> s = ps.Series(data) + >>> s + 0 1.0 + 1 2.0 + 2 3.0 + 3 4.0 + 4 NaN + 5 6.0 + 6 7.0 + 7 8.0 + dtype: float64 + + The `n` largest elements where ``n=5`` by default. + + >>> s.nlargest() + 7 8.0 + 6 7.0 + 5 6.0 + 3 4.0 + 2 3.0 + dtype: float64 + + >>> s.nlargest(n=3) + 7 8.0 + 6 7.0 + 5 6.0 + dtype: float64 + + + """ + return self.sort_values(ascending=False).head(n) + + def append( + self, to_append: "Series", ignore_index: bool = False, verify_integrity: bool = False + ) -> "Series": + """ + Concatenate two or more Series. + + .. deprecated:: 3.4.0 + + Parameters + ---------- + to_append : Series or list/tuple of Series + ignore_index : boolean, default False + If True, do not use the index labels. + verify_integrity : boolean, default False + If True, raise Exception on creating index with duplicates + + Returns + ------- + appended : Series + + Examples + -------- + >>> s1 = ps.Series([1, 2, 3]) + >>> s2 = ps.Series([4, 5, 6]) + >>> s3 = ps.Series([4, 5, 6], index=[3,4,5]) + + >>> s1.append(s2) # doctest: +SKIP + 0 1 + 1 2 + 2 3 + 0 4 + 1 5 + 2 6 + dtype: int64 + + >>> s1.append(s3) # doctest: +SKIP + 0 1 + 1 2 + 2 3 + 3 4 + 4 5 + 5 6 + dtype: int64 + + With ignore_index set to True: + + >>> s1.append(s2, ignore_index=True) # doctest: +SKIP + 0 1 + 1 2 + 2 3 + 3 4 + 4 5 + 5 6 + dtype: int64 + """ + warnings.warn( + "The Series.append method is deprecated " + "and will be removed in 4.0.0. " + "Use pyspark.pandas.concat instead.", + FutureWarning, + ) + return first_series( + self.to_frame().append(to_append.to_frame(), ignore_index, verify_integrity) + ).rename(self.name) + + def sample( + self, + n: Optional[int] = None, + frac: Optional[float] = None, + replace: bool = False, + random_state: Optional[int] = None, + ignore_index: bool = False, + ) -> "Series": + return first_series( + self.to_frame().sample( + n=n, + frac=frac, + replace=replace, + random_state=random_state, + ignore_index=ignore_index, + ) + ).rename(self.name) + + sample.__doc__ = DataFrame.sample.__doc__ + + @no_type_check + def hist(self, bins=10, **kwds): + return self.plot.hist(bins, **kwds) + + hist.__doc__ = PandasOnSparkPlotAccessor.hist.__doc__ + + def apply(self, func: Callable, args: Sequence[Any] = (), **kwds: Any) -> "Series": + """ + Invoke function on values of Series. + + Can be a Python function that only works on the Series. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def square(x) -> np.int32: + ... return x ** 2 + + pandas-on-Spark uses return type hint and does not try to infer the type. + + Parameters + ---------- + func : function + Python function to apply. Note that type hint for return type is required. + args : tuple + Positional arguments passed to func after the series value. + **kwds + Additional keyword arguments passed to func. + + Returns + ------- + Series + + See Also + -------- + Series.aggregate : Only perform aggregating type operations. + Series.transform : Only perform transforming type operations. + DataFrame.apply : The equivalent function for DataFrame. + + Examples + -------- + Create a Series with typical summer temperatures for each city. + + >>> s = ps.Series([20, 21, 12], + ... index=['London', 'New York', 'Helsinki']) + >>> s + London 20 + New York 21 + Helsinki 12 + dtype: int64 + + + Square the values by defining a function and passing it as an + argument to ``apply()``. + + >>> def square(x) -> np.int64: + ... return x ** 2 + >>> s.apply(square) + London 400 + New York 441 + Helsinki 144 + dtype: int64 + + + Define a custom function that needs additional positional + arguments and pass these additional arguments using the + ``args`` keyword + + >>> def subtract_custom_value(x, custom_value) -> np.int64: + ... return x - custom_value + + >>> s.apply(subtract_custom_value, args=(5,)) + London 15 + New York 16 + Helsinki 7 + dtype: int64 + + + Define a custom function that takes keyword arguments + and pass these arguments to ``apply`` + + >>> def add_custom_values(x, **kwargs) -> np.int64: + ... for month in kwargs: + ... x += kwargs[month] + ... return x + + >>> s.apply(add_custom_values, june=30, july=20, august=25) + London 95 + New York 96 + Helsinki 87 + dtype: int64 + + + Use a function from the Numpy library + + >>> def numpy_log(col) -> np.float64: + ... return np.log(col) + >>> s.apply(numpy_log) + London 2.995732 + New York 3.044522 + Helsinki 2.484907 + dtype: float64 + + + You can omit the type hint and let pandas-on-Spark infer its type. + + >>> s.apply(np.log) + London 2.995732 + New York 3.044522 + Helsinki 2.484907 + dtype: float64 + + """ + assert callable(func), "the first argument should be a callable function." + try: + spec = inspect.getfullargspec(func) + return_sig = spec.annotations.get("return", None) + should_infer_schema = return_sig is None + except TypeError: + # Falls back to schema inference if it fails to get signature. + should_infer_schema = True + + def apply_each(s: Any) -> pd.Series: + return s.apply(func, args=args, **kwds) + + if should_infer_schema: + return self.pandas_on_spark._transform_batch(apply_each, None) + else: + sig_return = infer_return_type(func) + if not isinstance(sig_return, ScalarType): + raise ValueError( + "Expected the return type of this function to be of scalar type, " + "but found type {}".format(sig_return) + ) + return_type = sig_return + return self.pandas_on_spark._transform_batch(apply_each, return_type) + + # TODO: not all arguments are implemented comparing to pandas' for now. + def aggregate(self, func: Union[str, List[str]]) -> Union[Scalar, "Series"]: + """Aggregate using one or more operations over the specified axis. + + Parameters + ---------- + func : str or a list of str + function name(s) as string apply to series. + + Returns + ------- + scalar, Series + The return can be: + - scalar : when Series.agg is called with single function + - Series : when Series.agg is called with several functions + + Notes + ----- + `agg` is an alias for `aggregate`. Use the alias. + + See Also + -------- + Series.apply : Invoke function on a Series. + Series.transform : Only perform transforming type operations. + Series.groupby : Perform operations over groups. + DataFrame.aggregate : The equivalent function for DataFrame. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4]) + >>> s.agg('min') + 1 + + >>> s.agg(['min', 'max']).sort_index() + max 4 + min 1 + dtype: int64 + """ + if isinstance(func, list): + return first_series(self.to_frame().aggregate(func)).rename(self.name) + elif isinstance(func, str): + return getattr(self, func)() + else: + raise TypeError("func must be a string or list of strings") + + agg = aggregate + + def transpose(self, *args: Any, **kwargs: Any) -> "Series": + """ + Return the transpose, which is self. + + Examples + -------- + It returns the same object as the transpose of the given series object, which is by + definition self. + + >>> s = ps.Series([1, 2, 3]) + >>> s + 0 1 + 1 2 + 2 3 + dtype: int64 + + >>> s.transpose() + 0 1 + 1 2 + 2 3 + dtype: int64 + """ + return self.copy() + + T = property(transpose) + + def transform( + self, func: Union[Callable, List[Callable]], axis: Axis = 0, *args: Any, **kwargs: Any + ) -> Union["Series", DataFrame]: + """ + Call ``func`` producing the same type as `self` with transformed values + and that has the same axis length as input. + + .. note:: this API executes the function once to infer the type which is + potentially expensive, for instance, when the dataset is created after + aggregations or sorting. + + To avoid this, specify return type in ``func``, for instance, as below: + + >>> def square(x) -> np.int32: + ... return x ** 2 + + pandas-on-Spark uses return type hint and does not try to infer the type. + + Parameters + ---------- + func : function or list + A function or a list of functions to use for transforming the data. + axis : int, default 0 or 'index' + Can only be set to 0 now. + *args + Positional arguments to pass to `func`. + **kwargs + Keyword arguments to pass to `func`. + + Returns + ------- + An instance of the same type with `self` that must have the same length as input. + + See Also + -------- + Series.aggregate : Only perform aggregating type operations. + Series.apply : Invoke function on Series. + DataFrame.transform : The equivalent function for DataFrame. + + Examples + -------- + + >>> s = ps.Series(range(3)) + >>> s + 0 0 + 1 1 + 2 2 + dtype: int64 + + >>> def sqrt(x) -> float: + ... return np.sqrt(x) + >>> s.transform(sqrt) + 0 0.000000 + 1 1.000000 + 2 1.414214 + dtype: float64 + + Even though the resulting instance must have the same length as the + input, it is possible to provide several input functions: + + >>> def exp(x) -> float: + ... return np.exp(x) + >>> s.transform([sqrt, exp]) + sqrt exp + 0 0.000000 1.000000 + 1 1.000000 2.718282 + 2 1.414214 7.389056 + + You can omit the type hint and let pandas-on-Spark infer its type. + + >>> s.transform([np.sqrt, np.exp]) + sqrt exp + 0 0.000000 1.000000 + 1 1.000000 2.718282 + 2 1.414214 7.389056 + """ + axis = validate_axis(axis) + if axis != 0: + raise NotImplementedError('axis should be either 0 or "index" currently.') + + if isinstance(func, list): + applied = [] + for f in func: + applied.append(self.apply(f, args=args, **kwargs).rename(f.__name__)) + + internal = self._internal.with_new_columns(applied) + return DataFrame(internal) + else: + return self.apply(func, args=args, **kwargs) + + def round(self, decimals: int = 0) -> "Series": + """ + Round each value in a Series to the given number of decimals. + + Parameters + ---------- + decimals : int + Number of decimal places to round to (default: 0). + If decimals are negative, it specifies the number of + positions to the left of the decimal point. + + Returns + ------- + Series object + + See Also + -------- + DataFrame.round + + Examples + -------- + >>> df = ps.Series([0.028208, 0.038683, 0.877076], name='x') + >>> df + 0 0.028208 + 1 0.038683 + 2 0.877076 + Name: x, dtype: float64 + + >>> df.round(2) + 0 0.03 + 1 0.04 + 2 0.88 + Name: x, dtype: float64 + """ + if not isinstance(decimals, int): + raise TypeError("decimals must be an integer") + scol = F.round(self.spark.column, decimals) + return self._with_new_scol( + scol, + field=( + self._internal.data_fields[0].copy(nullable=True) + if not isinstance(self.spark.data_type, DecimalType) + else None + ), + ) + + # TODO: add 'interpolation' parameter. + def quantile( + self, q: Union[float, Iterable[float]] = 0.5, accuracy: int = 10000 + ) -> Union[Scalar, "Series"]: + """ + Return value at the given quantile. + + .. note:: Unlike pandas', the quantile in pandas-on-Spark is an approximated quantile + based upon approximate percentile computation because computing quantile across + a large dataset is extremely expensive. + + Parameters + ---------- + q : float or array-like, default 0.5 (50% quantile) + 0 <= q <= 1, the quantile(s) to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + + Returns + ------- + float or Series + If the current object is a Series and ``q`` is an array, a Series will be + returned where the index is ``q`` and the values are the quantiles, otherwise + a float will be returned. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4, 5]) + >>> s.quantile(.5) + 3.0 + + >>> (s + 1).quantile(.5) + 4.0 + + >>> s.quantile([.25, .5, .75]) + 0.25 2.0 + 0.50 3.0 + 0.75 4.0 + dtype: float64 + + >>> (s + 1).quantile([.25, .5, .75]) + 0.25 3.0 + 0.50 4.0 + 0.75 5.0 + dtype: float64 + """ + if isinstance(q, Iterable): + return first_series( + cast( + "ps.DataFrame", + self.to_frame().quantile(q=q, axis=0, numeric_only=False, accuracy=accuracy), + ) + ).rename(self.name) + else: + if not isinstance(accuracy, int): + raise TypeError( + "accuracy must be an integer; however, got [%s]" % type(accuracy).__name__ + ) + + if not isinstance(q, float): + raise TypeError( + "q must be a float or an array of floats; however, [%s] found." % type(q) + ) + q_float = q + if q_float < 0.0 or q_float > 1.0: + raise ValueError("percentiles should all be in the interval [0, 1].") + + def quantile(psser: Series) -> PySparkColumn: + spark_type = psser.spark.data_type + spark_column = psser.spark.column + if isinstance(spark_type, (BooleanType, NumericType)): + return F.percentile_approx(spark_column.cast(DoubleType()), q_float, accuracy) + else: + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(spark_type), spark_type.simpleString() + ) + ) + + return self._reduce_for_stat_function(quantile, name="quantile") + + # TODO: add axis, pct, na_option parameter + def rank( + self, method: str = "average", ascending: bool = True, numeric_only: Optional[bool] = None + ) -> "Series": + """ + Compute numerical data ranks (1 through n) along axis. Equal values are + assigned a rank that is the average of the ranks of those values. + + .. note:: the current implementation of rank uses Spark's Window without + specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + method : {'average', 'min', 'max', 'first', 'dense'} + * average: average rank of group + * min: lowest rank in group + * max: highest rank in group + * first: ranks assigned in order they appear in the array + * dense: like 'min', but rank always increases by 1 between groups + ascending : boolean, default True + False for ranks by high (1) to low (N) + numeric_only : bool, optional + If set to True, rank numeric Series, or return an empty Series for non-numeric Series + + Returns + ------- + ranks : same type as caller + + Examples + -------- + >>> s = ps.Series([1, 2, 2, 3], name='A') + >>> s + 0 1 + 1 2 + 2 2 + 3 3 + Name: A, dtype: int64 + + >>> s.rank() + 0 1.0 + 1 2.5 + 2 2.5 + 3 4.0 + Name: A, dtype: float64 + + If method is set to 'min', it uses lowest rank in group. + + >>> s.rank(method='min') + 0 1.0 + 1 2.0 + 2 2.0 + 3 4.0 + Name: A, dtype: float64 + + If method is set to 'max', it uses highest rank in group. + + >>> s.rank(method='max') + 0 1.0 + 1 3.0 + 2 3.0 + 3 4.0 + Name: A, dtype: float64 + + If method is set to 'first', it is assigned rank in order without groups. + + >>> s.rank(method='first') + 0 1.0 + 1 2.0 + 2 3.0 + 3 4.0 + Name: A, dtype: float64 + + If method is set to 'dense', it leaves no gaps in group. + + >>> s.rank(method='dense') + 0 1.0 + 1 2.0 + 2 2.0 + 3 3.0 + Name: A, dtype: float64 + + If numeric_only is set to 'True', rank only numeric Series, + return an empty Series otherwise. + + >>> s = ps.Series(['a', 'b', 'c'], name='A', index=['x', 'y', 'z']) + >>> s + x a + y b + z c + Name: A, dtype: object + + >>> s.rank(numeric_only=True) + Series([], Name: A, dtype: float64) + """ + warnings.warn( + "Default value of `numeric_only` will be changed to `False` " + "instead of `None` in 4.0.0.", + FutureWarning, + ) + is_numeric = isinstance(self.spark.data_type, (NumericType, BooleanType)) + if numeric_only and not is_numeric: + return ps.Series([], dtype="float64", name=self.name) + else: + return self._rank(method, ascending).spark.analyzed + + def _rank( + self, + method: str = "average", + ascending: bool = True, + *, + part_cols: Sequence["ColumnOrName"] = (), + ) -> "Series": + if method not in ["average", "min", "max", "first", "dense"]: + msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'" + raise ValueError(msg) + + if self._internal.index_level > 1: + raise NotImplementedError("rank do not support MultiIndex now") + + Column = get_column_class() + if ascending: + asc_func = Column.asc + else: + asc_func = Column.desc + + if method == "first": + window = ( + Window.orderBy( + asc_func(self.spark.column), + asc_func(F.col(NATURAL_ORDER_COLUMN_NAME)), + ) + .partitionBy(*part_cols) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + scol = F.row_number().over(window) + elif method == "dense": + window = ( + Window.orderBy(asc_func(self.spark.column)) + .partitionBy(*part_cols) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + scol = F.dense_rank().over(window) + else: + if method == "average": + stat_func = F.mean + elif method == "min": + stat_func = F.min + elif method == "max": + stat_func = F.max + window1 = ( + Window.orderBy(asc_func(self.spark.column)) + .partitionBy(*part_cols) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + + window2 = Window.partitionBy( + cast("List[ColumnOrName]", [self.spark.column]) + list(part_cols) + ).rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + scol = stat_func(F.row_number().over(window1)).over(window2) + return self._with_new_scol(scol.cast(DoubleType())) + + def filter( + self, + items: Optional[Sequence[Any]] = None, + like: Optional[str] = None, + regex: Optional[str] = None, + axis: Optional[Axis] = None, + ) -> "Series": + axis = validate_axis(axis) + if axis == 1: + raise ValueError("Series does not support columns axis.") + return first_series( + self.to_frame().filter(items=items, like=like, regex=regex, axis=axis), + ).rename(self.name) + + filter.__doc__ = DataFrame.filter.__doc__ + + def describe(self, percentiles: Optional[List[float]] = None) -> "Series": + return first_series(self.to_frame().describe(percentiles)).rename(self.name) + + describe.__doc__ = DataFrame.describe.__doc__ + + def diff(self, periods: int = 1) -> "Series": + """ + First discrete difference of element. + + Calculates the difference of a Series element compared with another element in the + DataFrame (default is the element in the same column of the previous row). + + .. note:: the current implementation of diff uses Spark's Window without + specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int, default 1 + Periods to shift for calculating difference, accepts negative values. + + Returns + ------- + diffed : Series + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 2, 3, 4, 5, 6], + ... 'b': [1, 1, 2, 3, 5, 8], + ... 'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c']) + >>> df + a b c + 0 1 1 1 + 1 2 1 4 + 2 3 2 9 + 3 4 3 16 + 4 5 5 25 + 5 6 8 36 + + >>> df.b.diff() + 0 NaN + 1 0.0 + 2 1.0 + 3 1.0 + 4 2.0 + 5 3.0 + Name: b, dtype: float64 + + Difference with previous value + + >>> df.c.diff(periods=3) + 0 NaN + 1 NaN + 2 NaN + 3 15.0 + 4 21.0 + 5 27.0 + Name: c, dtype: float64 + + Difference with following value + + >>> df.c.diff(periods=-1) + 0 -3.0 + 1 -5.0 + 2 -7.0 + 3 -9.0 + 4 -11.0 + 5 NaN + Name: c, dtype: float64 + """ + return self._diff(periods).spark.analyzed + + def _diff(self, periods: int, *, part_cols: Sequence["ColumnOrName"] = ()) -> "Series": + if not isinstance(periods, int): + raise TypeError("periods should be an int; however, got [%s]" % type(periods).__name__) + window = ( + Window.partitionBy(*part_cols) + .orderBy(NATURAL_ORDER_COLUMN_NAME) + .rowsBetween(-periods, -periods) + ) + scol = self.spark.column - F.lag(self.spark.column, periods).over(window) + return self._with_new_scol(scol, field=self._internal.data_fields[0].copy(nullable=True)) + + def idxmax(self, skipna: bool = True) -> Union[Tuple, Any]: + """ + Return the row label of the maximum value. + + If multiple values equal the maximum, the first row label with that + value is returned. + + Parameters + ---------- + skipna : bool, default True + Exclude NA/null values. If the entire Series is NA, the result + will be NA. + + Returns + ------- + Index + Label of the maximum value. + + Raises + ------ + ValueError + If the Series is empty. + + See Also + -------- + Series.idxmin : Return index *label* of the first occurrence + of minimum of values. + + Examples + -------- + >>> s = ps.Series(data=[1, None, 4, 3, 5], + ... index=['A', 'B', 'C', 'D', 'E']) + >>> s + A 1.0 + B NaN + C 4.0 + D 3.0 + E 5.0 + dtype: float64 + + >>> s.idxmax() + 'E' + + If `skipna` is False and there is an NA value in the data, + the function returns ``nan``. + + >>> s.idxmax(skipna=False) + nan + + In case of multi-index, you get a tuple: + + >>> index = pd.MultiIndex.from_arrays([ + ... ['a', 'a', 'b', 'b'], ['c', 'd', 'e', 'f']], names=('first', 'second')) + >>> s = ps.Series(data=[1, None, 4, 5], index=index) + >>> s + first second + a c 1.0 + d NaN + b e 4.0 + f 5.0 + dtype: float64 + + >>> s.idxmax() + ('b', 'f') + + If multiple values equal the maximum, the first row label with that + value is returned. + + >>> s = ps.Series([1, 100, 1, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) + >>> s + 10 1 + 3 100 + 5 1 + 2 100 + 1 1 + 8 100 + dtype: int64 + + >>> s.idxmax() + 3 + """ + sdf = self._internal.spark_frame + scol = self.spark.column + index_scols = self._internal.index_spark_columns + + if skipna: + sdf = sdf.orderBy(scol.desc_nulls_last(), NATURAL_ORDER_COLUMN_NAME) + else: + sdf = sdf.orderBy(scol.desc_nulls_first(), NATURAL_ORDER_COLUMN_NAME) + + results = sdf.select([scol] + index_scols).take(1) + if len(results) == 0: + raise ValueError("attempt to get idxmin of an empty sequence") + if results[0][0] is None: + # This will only happen when skipna is False because we will + # place nulls first. + return np.nan + values = list(results[0][1:]) + if len(values) == 1: + return values[0] + else: + return tuple(values) + + def idxmin(self, skipna: bool = True) -> Union[Tuple, Any]: + """ + Return the row label of the minimum value. + + If multiple values equal the minimum, the first row label with that + value is returned. + + Parameters + ---------- + skipna : bool, default True + Exclude NA/null values. If the entire Series is NA, the result + will be NA. + + Returns + ------- + Index + Label of the minimum value. + + Raises + ------ + ValueError + If the Series is empty. + + See Also + -------- + Series.idxmax : Return index *label* of the first occurrence + of maximum of values. + + Notes + ----- + This method is the Series version of ``ndarray.argmin``. This method + returns the label of the minimum, while ``ndarray.argmin`` returns + the position. To get the position, use ``series.values.argmin()``. + + Examples + -------- + >>> s = ps.Series(data=[1, None, 4, 0], + ... index=['A', 'B', 'C', 'D']) + >>> s + A 1.0 + B NaN + C 4.0 + D 0.0 + dtype: float64 + + >>> s.idxmin() + 'D' + + If `skipna` is False and there is an NA value in the data, + the function returns ``nan``. + + >>> s.idxmin(skipna=False) + nan + + In case of multi-index, you get a tuple: + + >>> index = pd.MultiIndex.from_arrays([ + ... ['a', 'a', 'b', 'b'], ['c', 'd', 'e', 'f']], names=('first', 'second')) + >>> s = ps.Series(data=[1, None, 4, 0], index=index) + >>> s + first second + a c 1.0 + d NaN + b e 4.0 + f 0.0 + dtype: float64 + + >>> s.idxmin() + ('b', 'f') + + If multiple values equal the minimum, the first row label with that + value is returned. + + >>> s = ps.Series([1, 100, 1, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) + >>> s + 10 1 + 3 100 + 5 1 + 2 100 + 1 1 + 8 100 + dtype: int64 + + >>> s.idxmin() + 10 + """ + sdf = self._internal.spark_frame + scol = self.spark.column + index_scols = self._internal.index_spark_columns + + if skipna: + sdf = sdf.orderBy(scol.asc_nulls_last(), NATURAL_ORDER_COLUMN_NAME) + else: + sdf = sdf.orderBy(scol.asc_nulls_first(), NATURAL_ORDER_COLUMN_NAME) + + results = sdf.select([scol] + index_scols).take(1) + if len(results) == 0: + raise ValueError("attempt to get idxmin of an empty sequence") + if results[0][0] is None: + # This will only happen when skipna is False because we will + # place nulls first. + return np.nan + values = list(results[0][1:]) + if len(values) == 1: + return values[0] + else: + return tuple(values) + + def pop(self, item: Name) -> Union["Series", Scalar]: + """ + Return item and drop from series. + + Parameters + ---------- + item : label + Label of index to be popped. + + Returns + ------- + Value that is popped from series. + + Examples + -------- + >>> s = ps.Series(data=np.arange(3), index=['A', 'B', 'C']) + >>> s + A 0 + B 1 + C 2 + dtype: int64 + + >>> s.pop('A') + 0 + + >>> s + B 1 + C 2 + dtype: int64 + + >>> s = ps.Series(data=np.arange(3), index=['A', 'A', 'C']) + >>> s + A 0 + A 1 + C 2 + dtype: int64 + + >>> s.pop('A') + A 0 + A 1 + dtype: int64 + + >>> s + C 2 + dtype: int64 + + Also support for MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx) + >>> s + lama speed 45.0 + weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.pop('lama') + speed 45.0 + weight 200.0 + length 1.2 + dtype: float64 + + >>> s + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + Also support for MultiIndex with several indexes. + + >>> midx = pd.MultiIndex([['a', 'b', 'c'], + ... ['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 0, 0, 0, 1, 1, 1], + ... [0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 0, 2]] + ... ) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx) + >>> s + a lama speed 45.0 + weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + b falcon speed 320.0 + speed 1.0 + length 0.3 + dtype: float64 + + >>> s.pop(('a', 'lama')) + speed 45.0 + weight 200.0 + length 1.2 + dtype: float64 + + >>> s + a cow speed 30.0 + weight 250.0 + length 1.5 + b falcon speed 320.0 + speed 1.0 + length 0.3 + dtype: float64 + + >>> s.pop(('b', 'falcon', 'speed')) + (b, falcon, speed) 320.0 + (b, falcon, speed) 1.0 + dtype: float64 + """ + if not is_name_like_value(item): + raise TypeError("'key' should be string or tuple that contains strings") + if not is_name_like_tuple(item): + item = (item,) + if self._internal.index_level < len(item): + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + len(item), self._internal.index_level + ) + ) + + internal = self._internal + scols = internal.index_spark_columns[len(item) :] + [self.spark.column] + rows = [internal.spark_columns[level] == index for level, index in enumerate(item)] + sdf = internal.spark_frame.filter(reduce(lambda x, y: x & y, rows)).select(scols) + + psdf = self._drop(item) + self._update_anchor(psdf) + + if self._internal.index_level == len(item): + # if spark_frame has one column and one data, return data only without frame + pdf = sdf.limit(2).toPandas() + length = len(pdf) + if length == 1: + val = pdf[internal.data_spark_column_names[0]].iloc[0] + if isinstance(self.dtype, CategoricalDtype): + return self.dtype.categories[val] + else: + return val + + item_string = name_like_string(item) + sdf = sdf.withColumn(SPARK_DEFAULT_INDEX_NAME, F.lit(str(item_string))) + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, SPARK_DEFAULT_INDEX_NAME)], + column_labels=[self._column_label], + data_fields=[self._internal.data_fields[0]], + ) + return first_series(DataFrame(internal)) + else: + internal = internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in internal.index_spark_column_names[len(item) :] + ], + index_fields=internal.index_fields[len(item) :], + index_names=self._internal.index_names[len(item) :], + data_spark_columns=[scol_for(sdf, internal.data_spark_column_names[0])], + ) + return first_series(DataFrame(internal)) + + def copy(self, deep: bool = True) -> "Series": + """ + Make a copy of this object's indices and data. + + Parameters + ---------- + deep : bool, default True + this parameter is not supported but just dummy parameter to match pandas. + + Returns + ------- + copy : Series + + Examples + -------- + >>> s = ps.Series([1, 2], index=["a", "b"]) + >>> s + a 1 + b 2 + dtype: int64 + >>> s_copy = s.copy() + >>> s_copy + a 1 + b 2 + dtype: int64 + """ + return first_series(DataFrame(self._internal)) + + def mode(self, dropna: bool = True) -> "Series": + """ + Return the mode(s) of the dataset. + + Always returns Series even if only one value is returned. + + .. versionchanged:: 3.4.0 + Series name is preserved to follow pandas 1.4+ behavior. + + Parameters + ---------- + dropna : bool, default True + Don't consider counts of NaN/NaT. + + Returns + ------- + Series + Modes of the Series. + + Examples + -------- + >>> s = ps.Series([0, 0, 1, 1, 1, np.nan, np.nan, np.nan]) + >>> s + 0 0.0 + 1 0.0 + 2 1.0 + 3 1.0 + 4 1.0 + 5 NaN + 6 NaN + 7 NaN + dtype: float64 + + >>> s.mode() + 0 1.0 + dtype: float64 + + If there are several same modes, all items are shown + + >>> s = ps.Series([0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, + ... np.nan, np.nan, np.nan]) + >>> s + 0 0.0 + 1 0.0 + 2 1.0 + 3 1.0 + 4 1.0 + 5 2.0 + 6 2.0 + 7 2.0 + 8 3.0 + 9 3.0 + 10 3.0 + 11 NaN + 12 NaN + 13 NaN + dtype: float64 + + >>> s.mode().sort_values() + 0 1.0 + 1 2.0 + 2 3.0 + dtype: float64 + + With 'dropna' set to 'False', we can also see NaN in the result + + >>> s.mode(False).sort_values() + 0 1.0 + 1 2.0 + 2 3.0 + 3 NaN + dtype: float64 + """ + scol = self.spark.column + name = self._internal.data_spark_column_names[0] + sdf = ( + self._internal.spark_frame.select(SF.mode(scol, dropna).alias(name)) + .select(F.array_sort(F.col(name)).alias(name)) + .select(F.explode(F.col(name)).alias(name)) + ) + internal = InternalFrame(spark_frame=sdf, index_spark_columns=None, column_labels=[None]) + ser_mode = first_series(DataFrame(internal)) + ser_mode.name = self.name + return ser_mode + + def keys(self) -> "ps.Index": + """ + Return alias for index. + + Returns + ------- + Index + Index of the Series. + + Examples + -------- + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> psser = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + + >>> psser.keys() # doctest: +SKIP + MultiIndex([( 'lama', 'speed'), + ( 'lama', 'weight'), + ( 'lama', 'length'), + ( 'cow', 'speed'), + ( 'cow', 'weight'), + ( 'cow', 'length'), + ('falcon', 'speed'), + ('falcon', 'weight'), + ('falcon', 'length')], + ) + """ + return self.index + + # TODO: introduce 'method', 'limit', 'in_place'; fully support 'regex' + def replace( + self, + to_replace: Optional[Union[Any, List, Tuple, Dict]] = None, + value: Optional[Union[List, Tuple]] = None, + regex: Union[str, bool] = False, + ) -> "Series": + """ + Replace values given in to_replace with value. + Values of the Series are replaced with other values dynamically. + + .. note:: For partial pattern matching, the replacement is against the whole string, + which is different from pandas. That's by the nature of underlying Spark API. + + Parameters + ---------- + to_replace : str, list, tuple, dict, Series, int, float, or None + How to find the values that will be replaced. + * numeric, str: + + - numeric: numeric values equal to to_replace will be replaced with value + - str: string exactly matching to_replace will be replaced with value + + * list of str or numeric: + + - if to_replace and value are both lists or tuples, they must be the same length. + - str and numeric rules apply as above. + + * dict: + + - Dicts can be used to specify different replacement values for different + existing values. + For example, {'a': 'b', 'y': 'z'} replaces the value ‘a’ with ‘b’ and ‘y’ + with ‘z’. To use a dict in this way the value parameter should be None. + - For a DataFrame a dict can specify that different values should be replaced + in different columns. For example, {'a': 1, 'b': 'z'} looks for the value 1 + in column ‘a’ and the value ‘z’ in column ‘b’ and replaces these values with + whatever is specified in value. + The value parameter should not be None in this case. + You can treat this as a special case of passing two lists except that you are + specifying the column to search in. + + See the examples section for examples of each of these. + + value : scalar, dict, list, tuple, str default None + Value to replace any values matching to_replace with. + For a DataFrame a dict of values can be used to specify which value to use + for each column (columns not in the dict will not be filled). + Regular expressions, strings and lists or dicts of such objects are also allowed. + + regex: bool or str, default False + Whether to interpret to_replace and/or value as regular expressions. + If this is True then to_replace must be a string. + Alternatively, this could be a regular expression in which case to_replace must be None. + + + Returns + ------- + Series + Object after replacement. + + Examples + -------- + + Scalar `to_replace` and `value` + + >>> s = ps.Series([0, 1, 2, 3, 4]) + >>> s + 0 0 + 1 1 + 2 2 + 3 3 + 4 4 + dtype: int64 + + >>> s.replace(0, 5) + 0 5 + 1 1 + 2 2 + 3 3 + 4 4 + dtype: int64 + + List-like `to_replace` + + >>> s.replace([0, 4], 5000) + 0 5000 + 1 1 + 2 2 + 3 3 + 4 5000 + dtype: int64 + + >>> s.replace([1, 2, 3], [10, 20, 30]) + 0 0 + 1 10 + 2 20 + 3 30 + 4 4 + dtype: int64 + + Dict-like `to_replace` + + >>> s.replace({1: 1000, 2: 2000, 3: 3000, 4: 4000}) + 0 0 + 1 1000 + 2 2000 + 3 3000 + 4 4000 + dtype: int64 + + Also support for MultiIndex + + >>> midx = pd.MultiIndex([['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx) + >>> s + lama speed 45.0 + weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.replace(45, 450) + lama speed 450.0 + weight 200.0 + length 1.2 + cow speed 30.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.replace([45, 30, 320], 500) + lama speed 500.0 + weight 200.0 + length 1.2 + cow speed 500.0 + weight 250.0 + length 1.5 + falcon speed 500.0 + weight 1.0 + length 0.3 + dtype: float64 + + >>> s.replace({45: 450, 30: 300}) + lama speed 450.0 + weight 200.0 + length 1.2 + cow speed 300.0 + weight 250.0 + length 1.5 + falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + Regular expression `to_replace` + + >>> psser = ps.Series(['bat', 'foo', 'bait', 'abc', 'bar', 'zoo']) + >>> psser.replace(to_replace=r'^ba.$', value='new', regex=True) + 0 new + 1 foo + 2 bait + 3 abc + 4 new + 5 zoo + dtype: object + + >>> psser.replace(value='new', regex=r'^.oo$') + 0 bat + 1 new + 2 bait + 3 abc + 4 bar + 5 new + dtype: object + + For partial pattern matching, the replacement is against the whole string + + >>> psser.replace('ba', 'xx', regex=True) + 0 xx + 1 foo + 2 xx + 3 abc + 4 xx + 5 zoo + dtype: object + """ + if isinstance(regex, str): + if to_replace is not None: + raise ValueError("'to_replace' must be 'None' if 'regex' is not a bool") + to_replace = regex + regex = True + elif not isinstance(regex, bool): + raise NotImplementedError("'regex' of %s type is not supported" % type(regex).__name__) + elif regex is True: + assert isinstance( + to_replace, str + ), "If 'regex' is True then 'to_replace' must be a string" + + if to_replace is None: + return self.fillna(method="ffill") + if not isinstance(to_replace, (str, list, tuple, dict, int, float)): + raise TypeError("'to_replace' should be one of str, list, tuple, dict, int, float") + + to_replace = list(to_replace) if isinstance(to_replace, tuple) else to_replace + value = list(value) if isinstance(value, tuple) else value + if isinstance(to_replace, list) and isinstance(value, list): + if not len(to_replace) == len(value): + raise ValueError( + "Replacement lists must match in length. Expecting {} got {}".format( + len(to_replace), len(value) + ) + ) + to_replace = {k: v for k, v in zip(to_replace, value)} + if isinstance(to_replace, dict): + is_start = True + if len(to_replace) == 0: + current = self.spark.column + else: + for to_replace_, value in to_replace.items(): + cond = ( + (F.isnan(self.spark.column) | self.spark.column.isNull()) + if pd.isna(to_replace_) + else (self.spark.column == F.lit(to_replace_)) + ) + if is_start: + current = F.when(cond, value) + is_start = False + else: + current = current.when(cond, value) + current = current.otherwise(self.spark.column) + else: + if regex: + # to_replace must be a string + cond = self.spark.column.rlike(cast(str, to_replace)) + else: + cond = self.spark.column.isin(to_replace) + # to_replace may be a scalar + if np.array(pd.isna(to_replace)).any(): + cond = cond | F.isnan(self.spark.column) | self.spark.column.isNull() + current = F.when(cond, value).otherwise(self.spark.column) + + return self._with_new_scol(current) # TODO: dtype? + + def update(self, other: "Series") -> None: + """ + Modify Series in place using non-NA values from passed Series. Aligns on index. + + Parameters + ---------- + other : Series + + Examples + -------- + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> s = ps.Series([1, 2, 3]) + >>> s.update(ps.Series([4, 5, 6])) + >>> s.sort_index() + 0 4 + 1 5 + 2 6 + dtype: int64 + + >>> s = ps.Series(['a', 'b', 'c']) + >>> s.update(ps.Series(['d', 'e'], index=[0, 2])) + >>> s.sort_index() + 0 d + 1 b + 2 e + dtype: object + + >>> s = ps.Series([1, 2, 3]) + >>> s.update(ps.Series([4, 5, 6, 7, 8])) + >>> s.sort_index() + 0 4 + 1 5 + 2 6 + dtype: int64 + + >>> s = ps.Series([1, 2, 3], index=[10, 11, 12]) + >>> s + 10 1 + 11 2 + 12 3 + dtype: int64 + + >>> s.update(ps.Series([4, 5, 6])) + >>> s.sort_index() + 10 1 + 11 2 + 12 3 + dtype: int64 + + >>> s.update(ps.Series([4, 5, 6], index=[11, 12, 13])) + >>> s.sort_index() + 10 1 + 11 4 + 12 5 + dtype: int64 + + If ``other`` contains NaNs the corresponding values are not updated + in the original Series. + + >>> s = ps.Series([1, 2, 3]) + >>> s.update(ps.Series([4, np.nan, 6])) + >>> s.sort_index() + 0 4.0 + 1 2.0 + 2 6.0 + dtype: float64 + + >>> reset_option("compute.ops_on_diff_frames") + """ + if not isinstance(other, Series): + raise TypeError("'other' must be a Series") + + if same_anchor(self, other): + scol = ( + F.when(other.spark.column.isNotNull(), other.spark.column) + .otherwise(self.spark.column) + .alias(self._psdf._internal.spark_column_name_for(self._column_label)) + ) + internal = self._psdf._internal.with_new_spark_column( + self._column_label, scol # TODO: dtype? + ) + self._psdf._update_internal_frame(internal) + else: + combined = combine_frames(self._psdf, other._psdf, how="leftouter") + + this_scol = combined["this"]._internal.spark_column_for(self._column_label) + that_scol = combined["that"]._internal.spark_column_for(other._column_label) + + scol = ( + F.when(that_scol.isNotNull(), that_scol) + .otherwise(this_scol) + .alias(self._psdf._internal.spark_column_name_for(self._column_label)) + ) + + internal = combined["this"]._internal.with_new_spark_column( + self._column_label, scol # TODO: dtype? + ) + + self._psdf._update_internal_frame(internal.resolved_copy, check_same_anchor=False) + + def where(self, cond: "Series", other: Any = np.nan) -> "Series": + """ + Replace values where the condition is False. + + Parameters + ---------- + cond : boolean Series + Where cond is True, keep the original value. Where False, + replace with corresponding value from other. + other : scalar, Series + Entries where cond is False are replaced with corresponding value from other. + + Returns + ------- + Series + + Examples + -------- + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> s1 = ps.Series([0, 1, 2, 3, 4]) + >>> s2 = ps.Series([100, 200, 300, 400, 500]) + >>> s1.where(s1 > 0).sort_index() + 0 NaN + 1 1.0 + 2 2.0 + 3 3.0 + 4 4.0 + dtype: float64 + + >>> s1.where(s1 > 1, 10).sort_index() + 0 10 + 1 10 + 2 2 + 3 3 + 4 4 + dtype: int64 + + >>> s1.where(s1 > 1, s1 + 100).sort_index() + 0 100 + 1 101 + 2 2 + 3 3 + 4 4 + dtype: int64 + + >>> s1.where(s1 > 1, s2).sort_index() + 0 100 + 1 200 + 2 2 + 3 3 + 4 4 + dtype: int64 + + >>> reset_option("compute.ops_on_diff_frames") + """ + assert isinstance(cond, Series) + + # We should check the DataFrame from both `cond` and `other`. + should_try_ops_on_diff_frame = not same_anchor(cond, self) or ( + isinstance(other, Series) and not same_anchor(other, self) + ) + + if should_try_ops_on_diff_frame: + # Try to perform it with 'compute.ops_on_diff_frame' option. + psdf = self.to_frame() + tmp_cond_col = verify_temp_column_name(psdf, "__tmp_cond_col__") + tmp_other_col = verify_temp_column_name(psdf, "__tmp_other_col__") + + psdf[tmp_cond_col] = cond + psdf[tmp_other_col] = other + + # above logic makes a Spark DataFrame looks like below: + # +-----------------+---+----------------+-----------------+ + # |__index_level_0__| 0|__tmp_cond_col__|__tmp_other_col__| + # +-----------------+---+----------------+-----------------+ + # | 0| 0| false| 100| + # | 1| 1| false| 200| + # | 3| 3| true| 400| + # | 2| 2| true| 300| + # | 4| 4| true| 500| + # +-----------------+---+----------------+-----------------+ + condition = ( + F.when( + psdf[tmp_cond_col].spark.column, + psdf._psser_for(psdf._internal.column_labels[0]).spark.column, + ) + .otherwise(psdf[tmp_other_col].spark.column) + .alias(psdf._internal.data_spark_column_names[0]) + ) + + internal = psdf._internal.with_new_columns( + [condition], column_labels=self._internal.column_labels + ) + return first_series(DataFrame(internal)) + else: + if isinstance(other, Series): + other = other.spark.column + condition = ( + F.when(cond.spark.column, self.spark.column) + .otherwise(other) + .alias(self._internal.data_spark_column_names[0]) + ) + return self._with_new_scol(condition) + + def mask(self, cond: "Series", other: Any = np.nan) -> "Series": + """ + Replace values where the condition is True. + + Parameters + ---------- + cond : boolean Series + Where cond is False, keep the original value. Where True, + replace with corresponding value from other. + other : scalar, Series + Entries where cond is True are replaced with corresponding value from other. + + Returns + ------- + Series + + Examples + -------- + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> s1 = ps.Series([0, 1, 2, 3, 4]) + >>> s2 = ps.Series([100, 200, 300, 400, 500]) + >>> s1.mask(s1 > 0).sort_index() + 0 0.0 + 1 NaN + 2 NaN + 3 NaN + 4 NaN + dtype: float64 + + >>> s1.mask(s1 > 1, 10).sort_index() + 0 0 + 1 1 + 2 10 + 3 10 + 4 10 + dtype: int64 + + >>> s1.mask(s1 > 1, s1 + 100).sort_index() + 0 0 + 1 1 + 2 102 + 3 103 + 4 104 + dtype: int64 + + >>> s1.mask(s1 > 1, s2).sort_index() + 0 0 + 1 1 + 2 300 + 3 400 + 4 500 + dtype: int64 + + >>> reset_option("compute.ops_on_diff_frames") + """ + return self.where(~cond, other) + + def xs(self, key: Name, level: Optional[int] = None) -> "Series": + """ + Return cross-section from the Series. + + This method takes a `key` argument to select data at a particular + level of a MultiIndex. + + Parameters + ---------- + key : label or tuple of label + Label contained in the index, or partially in a MultiIndex. + 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. + + Returns + ------- + Series + Cross-section from the original Series + corresponding to the selected index levels. + + Examples + -------- + >>> midx = pd.MultiIndex([['a', 'b', 'c'], + ... ['lama', 'cow', 'falcon'], + ... ['speed', 'weight', 'length']], + ... [[0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 0, 0, 1, 1, 1, 2, 2, 2], + ... [0, 1, 2, 0, 1, 2, 0, 1, 2]]) + >>> s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], + ... index=midx) + >>> s + a lama speed 45.0 + weight 200.0 + length 1.2 + b cow speed 30.0 + weight 250.0 + length 1.5 + c falcon speed 320.0 + weight 1.0 + length 0.3 + dtype: float64 + + Get values at specified index + + >>> s.xs('a') + lama speed 45.0 + weight 200.0 + length 1.2 + dtype: float64 + + Get values at several indexes + + >>> s.xs(('a', 'lama')) + speed 45.0 + weight 200.0 + length 1.2 + dtype: float64 + + Get values at specified index and level + + >>> s.xs('lama', level=1) + a speed 45.0 + weight 200.0 + length 1.2 + dtype: float64 + """ + if not isinstance(key, tuple): + key = (key,) + if level is None: + level = 0 + + internal = self._internal + scols = ( + internal.index_spark_columns[:level] + + internal.index_spark_columns[level + len(key) :] + + [self.spark.column] + ) + rows = [internal.spark_columns[lvl] == index for lvl, index in enumerate(key, level)] + sdf = internal.spark_frame.filter(reduce(lambda x, y: x & y, rows)).select(scols) + + if internal.index_level == len(key): + # if spark_frame has one column and one data, return data only without frame + pdf = sdf.limit(2).toPandas() + length = len(pdf) + if length == 1: + return pdf[self._internal.data_spark_column_names[0]].iloc[0] + + index_spark_column_names = ( + internal.index_spark_column_names[:level] + + internal.index_spark_column_names[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( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_spark_column_names], + index_names=index_names, + index_fields=index_fields, + data_spark_columns=[scol_for(sdf, internal.data_spark_column_names[0])], + ) + return first_series(DataFrame(internal)) + + def pct_change(self, periods: int = 1) -> "Series": + """ + Percentage change between the current and a prior element. + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to moveing all data into + a single partition in a single machine and could cause serious + performance degradation. Avoid this method with very large datasets. + + Parameters + ---------- + periods : int, default 1 + Periods to shift for forming percent change. + + Returns + ------- + Series + + Examples + -------- + + >>> psser = ps.Series([90, 91, 85], index=[2, 4, 1]) + >>> psser + 2 90 + 4 91 + 1 85 + dtype: int64 + + >>> psser.pct_change() + 2 NaN + 4 0.011111 + 1 -0.065934 + dtype: float64 + + >>> psser.sort_index().pct_change() + 1 NaN + 2 0.058824 + 4 0.011111 + dtype: float64 + + >>> psser.pct_change(periods=2) + 2 NaN + 4 NaN + 1 -0.055556 + dtype: float64 + """ + scol = self.spark.column + + window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween(-periods, -periods) + prev_row = F.lag(scol, periods).over(window) + + return self._with_new_scol((scol - prev_row) / prev_row).spark.analyzed + + def combine_first(self, other: "Series") -> "Series": + """ + Combine Series values, choosing the calling Series's values first. + + Parameters + ---------- + other : Series + The value(s) to be combined with the `Series`. + + Returns + ------- + Series + The result of combining the Series with the other object. + + See Also + -------- + Series.combine : Perform element-wise operation on two Series + using a given function. + + Notes + ----- + Result index will be the union of the two indexes. + + Examples + -------- + >>> s1 = ps.Series([1, np.nan]) + >>> s2 = ps.Series([3, 4]) + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s1.combine_first(s2) + 0 1.0 + 1 4.0 + dtype: float64 + """ + if not isinstance(other, ps.Series): + raise TypeError("`combine_first` only allows `Series` for parameter `other`") + if same_anchor(self, other): + this = self.spark.column + that = other.spark.column + combined = self._psdf + else: + combined = combine_frames(self._psdf, other._psdf) + this = combined["this"]._internal.spark_column_for(self._column_label) + that = combined["that"]._internal.spark_column_for(other._column_label) + # If `self` has missing value, use value of `other` + cond = F.when(this.isNull(), that).otherwise(this) + # If `self` and `other` come from same frame, the anchor should be kept + if same_anchor(self, other): + return self._with_new_scol(cond) # TODO: dtype? + index_scols = combined._internal.index_spark_columns + sdf = combined._internal.spark_frame.select( + *index_scols, cond.alias(self._internal.data_spark_column_names[0]) + ).distinct() + internal = self._internal.with_new_sdf( + sdf, index_fields=combined._internal.index_fields, data_fields=[None] # TODO: dtype? + ) + return first_series(DataFrame(internal)) + + def dot(self, other: Union["Series", DataFrame]) -> Union[Scalar, "Series"]: + """ + Compute the dot product between the Series and the columns of other. + + This method computes the dot product between the Series and another + one, or the Series and each columns of a DataFrame. + + It can also be called using `self @ other` in Python >= 3.5. + + .. note:: This API is slightly different from pandas when indexes from both Series + are not aligned and config 'compute.eager_check' is False. pandas raise an exception; + however, pandas-on-Spark just proceeds and performs by ignoring mismatches with NaN + permissively. + + >>> pdf1 = pd.Series([1, 2, 3], index=[0, 1, 2]) + >>> pdf2 = pd.Series([1, 2, 3], index=[0, 1, 3]) + >>> pdf1.dot(pdf2) # doctest: +SKIP + ... + ValueError: matrices are not aligned + + >>> psdf1 = ps.Series([1, 2, 3], index=[0, 1, 2]) + >>> psdf2 = ps.Series([1, 2, 3], index=[0, 1, 3]) + >>> with ps.option_context("compute.eager_check", False): + ... psdf1.dot(psdf2) # doctest: +SKIP + ... + 5 + + Parameters + ---------- + other : Series, DataFrame. + The other object to compute the dot product with its columns. + + Returns + ------- + scalar, Series + Return the dot product of the Series and other if other is a + Series, the Series of the dot product of Series and each row of + other if other is a DataFrame. + + Notes + ----- + The Series and other must share the same index if other are a Series + or a DataFrame. + + Examples + -------- + >>> s = ps.Series([0, 1, 2, 3]) + + >>> s.dot(s) + 14 + + >>> s @ s + 14 + + >>> psdf = ps.DataFrame({'x': [0, 1, 2, 3], 'y': [0, -1, -2, -3]}) + >>> psdf + x y + 0 0 0 + 1 1 -1 + 2 2 -2 + 3 3 -3 + + >>> with ps.option_context("compute.ops_on_diff_frames", True): + ... s.dot(psdf) + ... + x 14 + y -14 + dtype: int64 + """ + if not same_anchor(self, other): + if get_option("compute.eager_check") and not cast( + ps.Index, self.index.sort_values() + ).equals(cast(ps.Index, other.index.sort_values())): + raise ValueError("matrices are not aligned") + elif len(self.index) != len(other.index): + raise ValueError("matrices are not aligned") + + if isinstance(other, DataFrame): + other_copy: DataFrame = other.copy() + column_labels = other_copy._internal.column_labels + + self_column_label = verify_temp_column_name(other_copy, "__self_column__") + other_copy[self_column_label] = self + self_psser = other_copy._psser_for(self_column_label) + + product_pssers = [ + cast(Series, other_copy._psser_for(label) * self_psser) for label in column_labels + ] + + dot_product_psser = DataFrame( + other_copy._internal.with_new_columns(product_pssers, column_labels=column_labels) + ).sum() + + return cast(Series, dot_product_psser).rename(self.name) + + else: + assert isinstance(other, Series) + return (self * other).sum() + + def __matmul__(self, other: Union["Series", DataFrame]) -> Union[Scalar, "Series"]: + """ + Matrix multiplication using binary `@` operator in Python>=3.5. + """ + return self.dot(other) + + def repeat(self, repeats: Union[int, "Series"]) -> "Series": + """ + Repeat elements of a Series. + + Returns a new Series where each element of the current Series + is repeated consecutively a given number of times. + + Parameters + ---------- + repeats : int or Series + The number of repetitions for each element. This should be a + non-negative integer. Repeating 0 times will return an empty + Series. + + Returns + ------- + Series + Newly created Series with repeated elements. + + See Also + -------- + Index.repeat : Equivalent function for Index. + + Examples + -------- + >>> s = ps.Series(['a', 'b', 'c']) + >>> s + 0 a + 1 b + 2 c + dtype: object + >>> s.repeat(2) + 0 a + 1 b + 2 c + 0 a + 1 b + 2 c + dtype: object + >>> ps.Series([1, 2, 3]).repeat(0) + Series([], dtype: int64) + """ + if not isinstance(repeats, (int, Series)): + raise TypeError( + "`repeats` argument must be integer or Series, but got {}".format(type(repeats)) + ) + + if isinstance(repeats, Series): + if not same_anchor(self, repeats): + psdf = self.to_frame() + temp_repeats = verify_temp_column_name(psdf, "__temp_repeats__") + psdf[temp_repeats] = repeats + return ( + psdf._psser_for(psdf._internal.column_labels[0]) + .repeat(psdf[temp_repeats]) + .rename(self.name) + ) + else: + scol = F.explode( + F.array_repeat(self.spark.column, repeats.astype("int32").spark.column) + ).alias(name_like_string(self.name)) + sdf = self._internal.spark_frame.select(self._internal.index_spark_columns + [scol]) + internal = self._internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + data_spark_columns=[scol_for(sdf, name_like_string(self.name))], + ) + return first_series(DataFrame(internal)) + else: + if repeats < 0: + raise ValueError("negative dimensions are not allowed") + + psdf = self._psdf[[self.name]] + if repeats == 0: + return first_series(DataFrame(psdf._internal.with_filter(F.lit(False)))) + else: + return first_series(cast("ps.DataFrame", ps.concat([psdf] * repeats))) + + def asof(self, where: Union[Any, List]) -> Union[Scalar, "Series"]: + """ + Return the last row(s) without any NaNs before `where`. + + The last row (for each element in `where`, if list) without any + NaN is taken. + + If there is no good value, NaN is returned. + + .. note:: This API is dependent on :meth:`Index.is_monotonic_increasing` + which is expensive. + + Parameters + ---------- + where : index or array-like of indices + + Returns + ------- + scalar or Series + + The return can be: + + * scalar : when `self` is a Series and `where` is a scalar + * Series: when `self` is a Series and `where` is an array-like + + Return scalar or Series + + Notes + ----- + 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 + + Examples + -------- + >>> s = ps.Series([1, 2, np.nan, 4], index=[10, 20, 30, 40]) + >>> s + 10 1.0 + 20 2.0 + 30 NaN + 40 4.0 + dtype: float64 + + A scalar `where`. + + >>> s.asof(20) + 2.0 + + For a sequence `where`, a Series is returned. The first value is + NaN, because the first element of `where` is before the first + index value. + + >>> s.asof([5, 20]).sort_index() + 5 NaN + 20 2.0 + dtype: float64 + + Missing values are not considered. The following is ``2.0``, not + NaN, even though NaN is at the index location for ``30``. + + >>> s.asof(30) + 2.0 + + >>> s = ps.Series([1, 2, np.nan, 4], index=[10, 30, 20, 40]) + >>> with ps.option_context("compute.eager_check", False): + ... s.asof(20) + ... + 1.0 + """ + should_return_series = True + if isinstance(self.index, ps.MultiIndex): + raise ValueError("asof is not supported for a MultiIndex") + if isinstance(where, (ps.Index, ps.Series, DataFrame)): + raise ValueError("where cannot be an Index, Series or a DataFrame") + if get_option("compute.eager_check") and not self.index.is_monotonic_increasing: + raise ValueError("asof requires a sorted index") + if not is_list_like(where): + should_return_series = False + where = [where] + internal = self._internal.resolved_copy + index_scol = internal.index_spark_columns[0] + index_type = internal.spark_type_for(index_scol) + spark_column = internal.data_spark_columns[0] + monotonically_increasing_id_column = verify_temp_column_name( + internal.spark_frame, "__monotonically_increasing_id__" + ) + cond = [ + F.max_by( + spark_column, + F.when( + (index_scol <= F.lit(index).cast(index_type)) & spark_column.isNotNull() + if pd.notna(index) + # If index is nan and the value of the col is not null + # then return monotonically_increasing_id. This will let max by + # to return last index value, which is the behaviour of pandas + else spark_column.isNotNull(), + monotonically_increasing_id_column, + ), + ) + for index in where + ] + + sdf = internal.spark_frame.withColumn( + monotonically_increasing_id_column, F.monotonically_increasing_id() + ).select(cond) + + if not should_return_series: + with sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + # Disable Arrow to keep row ordering. + result = sdf.limit(1).toPandas().iloc[0, 0] + return result if result is not None else np.nan + + # The data is expected to be small so it's fine to transpose/use default index. + with ps.option_context("compute.default_index_type", "distributed", "compute.max_rows", 1): + if len(where) == len(set(where)) and not isinstance(index_type, TimestampType): + psdf: DataFrame = DataFrame(sdf) + psdf.columns = pd.Index(where) + return first_series(psdf.transpose()).rename(self.name) + else: + # If `where` has duplicate items, leverage the pandas directly + # since pandas API on Spark doesn't support the duplicate column name. + pdf: pd.DataFrame = sdf.limit(1).toPandas() + pdf.columns = pd.Index(where) + return first_series(DataFrame(pdf.transpose())).rename(self.name) + + def mad(self) -> float: + """ + Return the mean absolute deviation of values. + + .. deprecated:: 3.4.0 + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4]) + >>> s + 0 1 + 1 2 + 2 3 + 3 4 + dtype: int64 + + >>> s.mad() + 1.0 + """ + warnings.warn( + "The 'mad' method is deprecated and will be removed in 4.0.0. " + "To compute the same result, you may do `(series - series.mean()).abs().mean()`.", + FutureWarning, + ) + sdf = self._internal.spark_frame + spark_column = self.spark.column + avg = unpack_scalar(sdf.select(F.avg(spark_column))) + mad = unpack_scalar(sdf.select(F.avg(F.abs(spark_column - avg)))) + + return mad + + def unstack(self, level: int = -1) -> DataFrame: + """ + Unstack, a.k.a. pivot, Series with MultiIndex to produce DataFrame. + The level involved will automatically get sorted. + + Notes + ----- + Unlike pandas, pandas-on-Spark doesn't check whether an index is duplicated or not + because the checking of duplicated index requires scanning whole data which + can be quite expensive. + + Parameters + ---------- + level : int, str, or list of these, default last level + Level(s) to unstack, can pass level name. + + Returns + ------- + DataFrame + Unstacked Series. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4], + ... index=pd.MultiIndex.from_product([['one', 'two'], + ... ['a', 'b']])) + >>> s + one a 1 + b 2 + two a 3 + b 4 + dtype: int64 + + >>> s.unstack(level=-1).sort_index() + a b + one 1 2 + two 3 4 + + >>> s.unstack(level=0).sort_index() + one two + a 1 3 + b 2 4 + """ + if not isinstance(self.index, ps.MultiIndex): + raise ValueError("Series.unstack only support for a MultiIndex") + index_nlevels = self.index.nlevels + if level > 0 and (level > index_nlevels - 1): + raise IndexError( + "Too many levels: Index has only {} levels, not {}".format(index_nlevels, level + 1) + ) + elif level < 0 and (level < -index_nlevels): + raise IndexError( + "Too many levels: Index has only {} levels, {} is not a valid level number".format( + index_nlevels, level + ) + ) + + internal = self._internal.resolved_copy + + index_map = list( + zip(internal.index_spark_column_names, internal.index_names, internal.index_fields) + ) + pivot_col, column_label_names, _ = index_map.pop(level) + index_scol_names, index_names, index_fields = zip(*index_map) + col = internal.data_spark_column_names[0] + + sdf = internal.spark_frame + sdf = sdf.groupby(list(index_scol_names)).pivot(pivot_col).agg(F.first(scol_for(sdf, col))) + + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in index_scol_names], + index_names=list(index_names), + index_fields=list(index_fields), + column_label_names=[column_label_names], + ) + internal = internal.copy( + data_fields=[ + field.copy(dtype=self._internal.data_fields[0].dtype) + for field in internal.data_fields + ] + ) + return DataFrame(internal) + + def item(self) -> Scalar: + """ + Return the first element of the underlying data as a Python scalar. + + Returns + ------- + scalar + The first element of Series. + + Raises + ------ + ValueError + If the data is not length-1. + + Examples + -------- + >>> psser = ps.Series([10]) + >>> psser.item() + 10 + """ + return self.head(2)._to_internal_pandas().item() + + def items(self) -> Iterable[Tuple[Name, Any]]: + """ + Lazily iterate over (index, value) tuples. + + This method returns an iterable tuple (index, value). This is + convenient if you want to create a lazy iterator. + + .. note:: Unlike pandas', the iteritems in pandas-on-Spark returns generator rather + zip object + + Returns + ------- + iterable + Iterable of tuples containing the (index, value) pairs from a + Series. + + See Also + -------- + DataFrame.items : Iterate over (column name, Series) pairs. + DataFrame.iterrows : Iterate over DataFrame rows as (index, Series) pairs. + + Examples + -------- + >>> s = ps.Series(['A', 'B', 'C']) + >>> for index, value in s.items(): + ... print("Index : {}, Value : {}".format(index, value)) + Index : 0, Value : A + Index : 1, Value : B + Index : 2, Value : C + """ + internal_index_columns = self._internal.index_spark_column_names + internal_data_column = self._internal.data_spark_column_names[0] + + def extract_kv_from_spark_row(row: Row) -> Tuple[Name, Any]: + k = ( + row[internal_index_columns[0]] + if len(internal_index_columns) == 1 + else tuple(row[c] for c in internal_index_columns) + ) + v = row[internal_data_column] + return k, v + + for k, v in map( + extract_kv_from_spark_row, self._internal.resolved_copy.spark_frame.toLocalIterator() + ): + yield k, v + + def iteritems(self) -> Iterable[Tuple[Name, Any]]: + """ + This is an alias of ``items``. + + .. deprecated:: 3.4.0 + iteritems is deprecated and will be removed in a future version. + Use .items instead. + """ + warnings.warn( + "Deprecated in 3.4, and will be removed in 4.0.0. Use Series.items instead.", + FutureWarning, + ) + return self.items() + + def droplevel(self, level: Union[int, Name, List[Union[int, Name]]]) -> "Series": + """ + Return Series with requested index level(s) removed. + + Parameters + ---------- + level : int, str, or list-like + If a string is given, must be the name of a level + If list-like, elements must be names or positional indexes + of levels. + + Returns + ------- + Series + Series with requested index level(s) removed. + + Examples + -------- + >>> psser = ps.Series( + ... [1, 2, 3], + ... index=pd.MultiIndex.from_tuples( + ... [("x", "a"), ("x", "b"), ("y", "c")], names=["level_1", "level_2"] + ... ), + ... ) + >>> psser + level_1 level_2 + x a 1 + b 2 + y c 3 + dtype: int64 + + Removing specific index level by level + + >>> psser.droplevel(0) + level_2 + a 1 + b 2 + c 3 + dtype: int64 + + Removing specific index level by name + + >>> psser.droplevel("level_2") + level_1 + x 1 + x 2 + y 3 + dtype: int64 + """ + return first_series(self.to_frame().droplevel(level=level, axis=0)).rename(self.name) + + def tail(self, n: int = 5) -> "Series": + """ + Return the last `n` rows. + + This function returns last `n` rows from the object based on + position. It is useful for quickly verifying data, for example, + after sorting or appending rows. + + For negative values of `n`, this function returns all rows except + the first `n` rows, equivalent to ``df[n:]``. + + Parameters + ---------- + n : int, default 5 + Number of rows to select. + + Returns + ------- + type of caller + The last `n` rows of the caller object. + + See Also + -------- + DataFrame.head : The first `n` rows of the caller object. + + Examples + -------- + >>> psser = ps.Series([1, 2, 3, 4, 5]) + >>> psser + 0 1 + 1 2 + 2 3 + 3 4 + 4 5 + dtype: int64 + + >>> psser.tail(3) # doctest: +SKIP + 2 3 + 3 4 + 4 5 + dtype: int64 + """ + return first_series(self.to_frame().tail(n=n)).rename(self.name) + + def explode(self) -> "Series": + """ + Transform each element of a list-like to a row. + + Returns + ------- + Series + Exploded lists to rows; index will be duplicated for these rows. + + See Also + -------- + Series.str.split : Split string values on specified separator. + Series.unstack : Unstack, a.k.a. pivot, Series with MultiIndex + to produce DataFrame. + DataFrame.melt : Unpivot a DataFrame from wide format to long format. + DataFrame.explode : Explode a DataFrame from list-like + columns to long format. + + Examples + -------- + >>> psser = ps.Series([[1, 2, 3], [], [3, 4]]) + >>> psser + 0 [1, 2, 3] + 1 [] + 2 [3, 4] + dtype: object + + >>> psser.explode() # doctest: +SKIP + 0 1.0 + 0 2.0 + 0 3.0 + 1 NaN + 2 3.0 + 2 4.0 + dtype: float64 + """ + if not isinstance(self.spark.data_type, ArrayType): + return self.copy() + + scol = F.explode_outer(self.spark.column).alias(name_like_string(self._column_label)) + + internal = self._internal.with_new_columns([scol], keep_order=False) + return first_series(DataFrame(internal)) + + def argsort(self) -> "Series": + """ + Return the integer indices that would sort the Series values. + Unlike pandas, the index order is not preserved in the result. + + Returns + ------- + Series + Positions of values within the sort order with -1 indicating + nan values. + + Examples + -------- + >>> psser = ps.Series([3, 3, 4, 1, 6, 2, 3, 7, 8, 7, 10]) + >>> psser + 0 3 + 1 3 + 2 4 + 3 1 + 4 6 + 5 2 + 6 3 + 7 7 + 8 8 + 9 7 + 10 10 + dtype: int64 + + >>> psser.argsort().sort_index() + 0 3 + 1 5 + 2 0 + 3 1 + 4 6 + 5 2 + 6 4 + 7 7 + 8 9 + 9 8 + 10 10 + dtype: int64 + """ + notnull = self.loc[self.notnull()] + + sdf_for_index = notnull._internal.spark_frame.select(notnull._internal.index_spark_columns) + + tmp_join_key = verify_temp_column_name(sdf_for_index, "__tmp_join_key__") + sdf_for_index = InternalFrame.attach_distributed_sequence_column( + sdf_for_index, tmp_join_key + ) + # sdf_for_index: + # +----------------+-----------------+ + # |__tmp_join_key__|__index_level_0__| + # +----------------+-----------------+ + # | 0| 0| + # | 1| 1| + # | 2| 2| + # | 3| 3| + # | 4| 4| + # +----------------+-----------------+ + + sdf_for_data = notnull._internal.spark_frame.select( + notnull.spark.column.alias("values"), NATURAL_ORDER_COLUMN_NAME + ) + sdf_for_data = InternalFrame.attach_distributed_sequence_column( + sdf_for_data, SPARK_DEFAULT_SERIES_NAME + ) + # sdf_for_data: + # +---+------+-----------------+ + # | 0|values|__natural_order__| + # +---+------+-----------------+ + # | 0| 3| 25769803776| + # | 1| 3| 51539607552| + # | 2| 4| 77309411328| + # | 3| 1| 103079215104| + # | 4| 2| 128849018880| + # +---+------+-----------------+ + + sdf_for_data = sdf_for_data.sort( + scol_for(sdf_for_data, "values"), NATURAL_ORDER_COLUMN_NAME + ).drop("values", NATURAL_ORDER_COLUMN_NAME) + + tmp_join_key = verify_temp_column_name(sdf_for_data, "__tmp_join_key__") + sdf_for_data = InternalFrame.attach_distributed_sequence_column(sdf_for_data, tmp_join_key) + # sdf_for_index: sdf_for_data: + # +----------------+-----------------+ +----------------+---+ + # |__tmp_join_key__|__index_level_0__| |__tmp_join_key__| 0| + # +----------------+-----------------+ +----------------+---+ + # | 0| 0| | 0| 3| + # | 1| 1| | 1| 4| + # | 2| 2| | 2| 0| + # | 3| 3| | 3| 1| + # | 4| 4| | 4| 2| + # +----------------+-----------------+ +----------------+---+ + + sdf = sdf_for_index.join(sdf_for_data, on=tmp_join_key).drop(tmp_join_key) + + internal = self._internal.with_new_sdf( + spark_frame=sdf, + data_columns=[SPARK_DEFAULT_SERIES_NAME], + index_fields=[ + InternalField(dtype=field.dtype) for field in self._internal.index_fields + ], + data_fields=[None], + ) + psser = first_series(DataFrame(internal)) + + return cast( + Series, + ps.concat([psser, self.loc[self.isnull()].spark.transform(lambda _: F.lit(-1))]), + ) + + def argmax(self, axis: Axis = None, skipna: bool = True) -> int: + """ + Return int position of the largest value in the Series. + + If the maximum is achieved in multiple locations, + the first row position is returned. + + Parameters + ---------- + axis : None + Dummy argument for consistency with Series. + skipna : bool, default True + Exclude NA/null values. + + Returns + ------- + int + Row position of the maximum value. + + Examples + -------- + Consider dataset containing cereal calories + + >>> s = ps.Series({'Corn Flakes': 100.0, 'Almond Delight': 110.0, 'Unknown': np.nan, + ... 'Cinnamon Toast Crunch': 120.0, 'Cocoa Puff': 110.0}) + >>> s + Corn Flakes 100.0 + Almond Delight 110.0 + Unknown NaN + Cinnamon Toast Crunch 120.0 + Cocoa Puff 110.0 + dtype: float64 + + >>> s.argmax() + 3 + + >>> s.argmax(skipna=False) + -1 + """ + axis = validate_axis(axis, none_axis=0) + if axis == 1: + raise ValueError("axis can only be 0 or 'index'") + sdf = self._internal.spark_frame.select(self.spark.column, NATURAL_ORDER_COLUMN_NAME) + seq_col_name = verify_temp_column_name(sdf, "__distributed_sequence_column__") + sdf = InternalFrame.attach_distributed_sequence_column( + sdf, + seq_col_name, + ) + scol = scol_for(sdf, self._internal.data_spark_column_names[0]) + + if skipna: + sdf = sdf.orderBy(scol.desc_nulls_last(), NATURAL_ORDER_COLUMN_NAME, seq_col_name) + else: + sdf = sdf.orderBy(scol.desc_nulls_first(), NATURAL_ORDER_COLUMN_NAME, seq_col_name) + + results = sdf.select(scol, seq_col_name).take(1) + + if len(results) == 0: + raise ValueError("attempt to get argmax of an empty sequence") + else: + max_value = results[0] + # If the maximum is achieved in multiple locations, the first row position is returned. + return -1 if max_value[0] is None else max_value[1] + + def argmin(self, axis: Axis = None, skipna: bool = True) -> int: + """ + Return int position of the smallest value in the Series. + + If the minimum is achieved in multiple locations, + the first row position is returned. + + Parameters + ---------- + axis : None + Dummy argument for consistency with Series. + skipna : bool, default True + Exclude NA/null values. + + Returns + ------- + int + Row position of the minimum value. + + Examples + -------- + Consider dataset containing cereal calories + + >>> s = ps.Series({'Corn Flakes': 100.0, 'Almond Delight': 110.0, + ... 'Cinnamon Toast Crunch': 120.0, 'Cocoa Puff': 110.0}) + >>> s # doctest: +SKIP + Corn Flakes 100.0 + Almond Delight 110.0 + Cinnamon Toast Crunch 120.0 + Cocoa Puff 110.0 + dtype: float64 + + >>> s.argmin() # doctest: +SKIP + 0 + """ + axis = validate_axis(axis, none_axis=0) + if axis == 1: + raise ValueError("axis can only be 0 or 'index'") + sdf = self._internal.spark_frame.select(self.spark.column, NATURAL_ORDER_COLUMN_NAME) + seq_col_name = verify_temp_column_name(sdf, "__distributed_sequence_column__") + sdf = InternalFrame.attach_distributed_sequence_column( + sdf, + seq_col_name, + ) + scol = scol_for(sdf, self._internal.data_spark_column_names[0]) + + if skipna: + sdf = sdf.orderBy(scol.asc_nulls_last(), NATURAL_ORDER_COLUMN_NAME, seq_col_name) + else: + sdf = sdf.orderBy(scol.asc_nulls_first(), NATURAL_ORDER_COLUMN_NAME, seq_col_name) + + results = sdf.select(scol, seq_col_name).take(1) + + if len(results) == 0: + raise ValueError("attempt to get argmin of an empty sequence") + else: + min_value = results[0] + # If the maximum is achieved in multiple locations, the first row position is returned. + return -1 if min_value[0] is None else min_value[1] + + def compare( + self, other: "Series", keep_shape: bool = False, keep_equal: bool = False + ) -> DataFrame: + """ + Compare to another Series and show the differences. + + .. note:: This API is slightly different from pandas when indexes from both Series + are not identical and config 'compute.eager_check' is False. pandas raise an exception; + however, pandas-on-Spark just proceeds and performs by ignoring mismatches. + + >>> psser1 = ps.Series([1, 2, 3, 4, 5], index=pd.Index([1, 2, 3, 4, 5])) + >>> psser2 = ps.Series([1, 2, 3, 4, 5], index=pd.Index([1, 2, 4, 3, 6])) + >>> psser1.compare(psser2) # doctest: +SKIP + ... + ValueError: Can only compare identically-labeled Series objects + + >>> with ps.option_context("compute.eager_check", False): + ... psser1.compare(psser2) # doctest: +SKIP + ... + self other + 3 3.0 4.0 + 4 4.0 3.0 + 5 5.0 NaN + 6 NaN 5.0 + + Parameters + ---------- + other : Series + Object to compare with. + keep_shape : bool, default False + If true, all rows and columns are kept. + Otherwise, only the ones with different values are kept. + keep_equal : bool, default False + If true, the result keeps values that are equal. + Otherwise, equal values are shown as NaNs. + + Returns + ------- + DataFrame + + Notes + ----- + Matching NaNs will not appear as a difference. + + Examples + -------- + + >>> from pyspark.pandas.config import set_option, reset_option + >>> set_option("compute.ops_on_diff_frames", True) + >>> s1 = ps.Series(["a", "b", "c", "d", "e"]) + >>> s2 = ps.Series(["a", "a", "c", "b", "e"]) + + Align the differences on columns + + >>> s1.compare(s2).sort_index() + self other + 1 b a + 3 d b + + Keep all original rows + + >>> s1.compare(s2, keep_shape=True).sort_index() + self other + 0 None None + 1 b a + 2 None None + 3 d b + 4 None None + + Keep all original rows and all original values + + >>> s1.compare(s2, keep_shape=True, keep_equal=True).sort_index() + self other + 0 a a + 1 b a + 2 c c + 3 d b + 4 e e + + >>> reset_option("compute.ops_on_diff_frames") + """ + combined: DataFrame + if same_anchor(self, other): + self_column_label = verify_temp_column_name(other.to_frame(), "__self_column__") + other_column_label = verify_temp_column_name(self.to_frame(), "__other_column__") + combined = DataFrame( + self._internal.with_new_columns( + [self.rename(self_column_label), other.rename(other_column_label)] + ) + ) + else: + if get_option("compute.eager_check") and not self.index.equals(other.index): + raise ValueError("Can only compare identically-labeled Series objects") + + combined = combine_frames(self.to_frame(), other.to_frame()) + + this_column_label = "self" + that_column_label = "other" + if keep_equal and keep_shape: + combined.columns = pd.Index([this_column_label, that_column_label]) + return combined + + this_data_scol = combined._internal.data_spark_columns[0] + that_data_scol = combined._internal.data_spark_columns[1] + index_scols = combined._internal.index_spark_columns + sdf = combined._internal.spark_frame + if keep_shape: + this_scol = ( + F.when(this_data_scol == that_data_scol, None) + .otherwise(this_data_scol) + .alias(this_column_label) + ) + this_field = combined._internal.data_fields[0].copy( + name=this_column_label, nullable=True + ) + + that_scol = ( + F.when(this_data_scol == that_data_scol, None) + .otherwise(that_data_scol) + .alias(that_column_label) + ) + that_field = combined._internal.data_fields[1].copy( + name=that_column_label, nullable=True + ) + else: + sdf = sdf.filter(~this_data_scol.eqNullSafe(that_data_scol)) + + this_scol = this_data_scol.alias(this_column_label) + this_field = combined._internal.data_fields[0].copy(name=this_column_label) + + that_scol = that_data_scol.alias(that_column_label) + that_field = combined._internal.data_fields[1].copy(name=that_column_label) + + sdf = sdf.select(*index_scols, this_scol, that_scol, NATURAL_ORDER_COLUMN_NAME) + internal = InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, col) for col in self._internal.index_spark_column_names + ], + index_names=self._internal.index_names, + index_fields=combined._internal.index_fields, + column_labels=[(this_column_label,), (that_column_label,)], + data_spark_columns=[scol_for(sdf, this_column_label), scol_for(sdf, that_column_label)], + data_fields=[this_field, that_field], + column_label_names=[None], + ) + return DataFrame(internal) + + # TODO(SPARK-40553): 1, support array-like 'value'; 2, add parameter 'sorter' + def searchsorted(self, value: Any, side: str = "left") -> int: + """ + Find indices where elements should be inserted to maintain order. + + Find the indices into a sorted Series self such that, if the corresponding elements + in value were inserted before the indices, the order of self would be preserved. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + value : scalar + Values to insert into self. + side : {‘left’, ‘right’}, optional + If ‘left’, the index of the first suitable location found is given. + If ‘right’, return the last such index. If there is no suitable index, + return either 0 or N (where N is the length of self). + + Returns + ------- + int + insertion point + + Notes + ----- + The Series must be monotonically sorted, otherwise wrong locations will likely be returned. + + Examples + -------- + >>> ser = ps.Series([1, 2, 2, 3]) + >>> ser.searchsorted(0) + 0 + >>> ser.searchsorted(1) + 0 + >>> ser.searchsorted(2) + 1 + >>> ser.searchsorted(5) + 4 + >>> ser.searchsorted(0, side="right") + 0 + >>> ser.searchsorted(1, side="right") + 1 + >>> ser.searchsorted(2, side="right") + 3 + >>> ser.searchsorted(5, side="right") + 4 + """ + if side not in ["left", "right"]: + raise ValueError(f"Invalid side {side}") + + sdf = self._internal.spark_frame + index_col_name = verify_temp_column_name(sdf, "__search_sorted_index_col__") + value_col_name = verify_temp_column_name(sdf, "__search_sorted_value_col__") + sdf = InternalFrame.attach_distributed_sequence_column( + sdf.select(self.spark.column.alias(value_col_name)), index_col_name + ) + + if side == "left": + results = sdf.select( + F.min(F.when(F.lit(value) <= F.col(value_col_name), F.col(index_col_name))), + F.count(F.lit(0)), + ).take(1) + else: + results = sdf.select( + F.min(F.when(F.lit(value) < F.col(value_col_name), F.col(index_col_name))), + F.count(F.lit(0)), + ).take(1) + + if len(results) == 0: + return 0 + else: + return results[0][1] if results[0][0] is None else results[0][0] + + def align( + self, + other: Union[DataFrame, "Series"], + join: str = "outer", + axis: Optional[Axis] = None, + copy: bool = True, + ) -> Tuple["Series", Union[DataFrame, "Series"]]: + """ + Align two objects on their axes with the specified join method. + + Join method is specified for each axis Index. + + Parameters + ---------- + other : DataFrame or Series + join : {{'outer', 'inner', 'left', 'right'}}, default 'outer' + axis : allowed axis of the other object, default None + Align on index (0), columns (1), or both (None). + copy : bool, default True + Always returns new objects. If copy=False and no reindexing is + required then original objects are returned. + + Returns + ------- + (left, right) : (Series, type of other) + Aligned objects. + + Examples + -------- + >>> ps.set_option("compute.ops_on_diff_frames", True) + >>> s1 = ps.Series([7, 8, 9], index=[10, 11, 12]) + >>> s2 = ps.Series(["g", "h", "i"], index=[10, 20, 30]) + + >>> aligned_l, aligned_r = s1.align(s2) + >>> aligned_l.sort_index() + 10 7.0 + 11 8.0 + 12 9.0 + 20 NaN + 30 NaN + dtype: float64 + >>> aligned_r.sort_index() + 10 g + 11 None + 12 None + 20 h + 30 i + dtype: object + + Align with the join type "inner": + + >>> aligned_l, aligned_r = s1.align(s2, join="inner") + >>> aligned_l.sort_index() + 10 7 + dtype: int64 + >>> aligned_r.sort_index() + 10 g + dtype: object + + Align with a DataFrame: + + >>> df = ps.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30]) + >>> aligned_l, aligned_r = s1.align(df) + >>> aligned_l.sort_index() + 10 7.0 + 11 8.0 + 12 9.0 + 20 NaN + 30 NaN + dtype: float64 + >>> aligned_r.sort_index() + a b + 10 1.0 a + 11 NaN None + 12 NaN None + 20 2.0 b + 30 3.0 c + + >>> ps.reset_option("compute.ops_on_diff_frames") + """ + axis = validate_axis(axis) + if axis == 1: + raise ValueError("Series does not support columns axis.") + + self_df = self.to_frame() + left, right = self_df.align(other, join=join, axis=axis, copy=False) + + if left is self_df: + left_ser = self + else: + left_ser = first_series(left).rename(self.name) + + return (left_ser.copy(), right.copy()) if copy else (left_ser, right) + + # TODO(SPARK-42620): Add `inclusive` parameter and replace `include_start` & `include_end`. + # See https://github.com/pandas-dev/pandas/issues/43248 + def between_time( + self, + start_time: Union[datetime.time, str], + end_time: Union[datetime.time, str], + include_start: bool = True, + include_end: bool = True, + axis: Axis = 0, + ) -> "Series": + """ + Select values between particular times of the day (example: 9:00-9:30 AM). + + By setting ``start_time`` to be later than ``end_time``, + you can get the times that are *not* between the two times. + + Parameters + ---------- + start_time : datetime.time or str + Initial time as a time filter limit. + end_time : datetime.time or str + End time as a time filter limit. + include_start : bool, default True + Whether the start time needs to be included in the result. + + .. deprecated:: 3.4.0 + + include_end : bool, default True + Whether the end time needs to be included in the result. + + .. deprecated:: 3.4.0 + + axis : {0 or 'index', 1 or 'columns'}, default 0 + Determine range time on index or columns value. + + Returns + ------- + Series + Data from the original object filtered to the specified dates range. + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + See Also + -------- + at_time : Select values at a particular time of the day. + last : Select final periods of time series based on a date offset. + DatetimeIndex.indexer_between_time : Get just the index locations for + values between particular times of the day. + + Examples + -------- + >>> idx = pd.date_range('2018-04-09', periods=4, freq='1D20min') + >>> psser = ps.Series([1, 2, 3, 4], index=idx) + >>> psser + 2018-04-09 00:00:00 1 + 2018-04-10 00:20:00 2 + 2018-04-11 00:40:00 3 + 2018-04-12 01:00:00 4 + dtype: int64 + + >>> psser.between_time('0:15', '0:45') # doctest: +SKIP + 2018-04-10 00:20:00 2 + 2018-04-11 00:40:00 3 + dtype: int64 + """ + return first_series( + self.to_frame().between_time(start_time, end_time, include_start, include_end, axis) + ).rename(self.name) + + def at_time( + self, time: Union[datetime.time, str], asof: bool = False, axis: Axis = 0 + ) -> "Series": + """ + Select values at particular time of day (example: 9:30AM). + + Parameters + ---------- + time : datetime.time or str + axis : {0 or 'index', 1 or 'columns'}, default 0 + + Returns + ------- + Series + + Raises + ------ + TypeError + If the index is not a :class:`DatetimeIndex` + + See Also + -------- + between_time : Select values between particular times of the day. + DatetimeIndex.indexer_at_time : Get just the index locations for + values at particular time of the day. + + Examples + -------- + >>> idx = pd.date_range('2018-04-09', periods=4, freq='12H') + >>> psser = ps.Series([1, 2, 3, 4], index=idx) + >>> psser + 2018-04-09 00:00:00 1 + 2018-04-09 12:00:00 2 + 2018-04-10 00:00:00 3 + 2018-04-10 12:00:00 4 + dtype: int64 + + >>> psser.at_time('12:00') + 2018-04-09 12:00:00 2 + 2018-04-10 12:00:00 4 + dtype: int64 + """ + return first_series(self.to_frame().at_time(time, asof, axis)).rename(self.name) + + def _cum( + self, + func: Callable[[PySparkColumn], PySparkColumn], + skipna: bool, + part_cols: Sequence["ColumnOrName"] = (), + ascending: bool = True, + ) -> "Series": + # This is used to cummin, cummax, cumsum, etc. + + if ascending: + window = ( + Window.orderBy(F.asc(NATURAL_ORDER_COLUMN_NAME)) + .partitionBy(*part_cols) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + else: + window = ( + Window.orderBy(F.desc(NATURAL_ORDER_COLUMN_NAME)) + .partitionBy(*part_cols) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + + if skipna: + # There is a behavior difference between pandas and PySpark. In case of cummax, + # + # Input: + # A B + # 0 2.0 1.0 + # 1 5.0 NaN + # 2 1.0 0.0 + # 3 2.0 4.0 + # 4 4.0 9.0 + # + # pandas: + # A B + # 0 2.0 1.0 + # 1 5.0 NaN + # 2 5.0 1.0 + # 3 5.0 4.0 + # 4 5.0 9.0 + # + # PySpark: + # A B + # 0 2.0 1.0 + # 1 5.0 1.0 + # 2 5.0 1.0 + # 3 5.0 4.0 + # 4 5.0 9.0 + + scol = F.when( + # Manually sets nulls given the column defined above. + self.spark.column.isNull(), + F.lit(None), + ).otherwise(func(self.spark.column).over(window)) + else: + # Here, we use two Windows. + # One for real data. + # The other one for setting nulls after the first null it meets. + # + # There is a behavior difference between pandas and PySpark. In case of cummax, + # + # Input: + # A B + # 0 2.0 1.0 + # 1 5.0 NaN + # 2 1.0 0.0 + # 3 2.0 4.0 + # 4 4.0 9.0 + # + # pandas: + # A B + # 0 2.0 1.0 + # 1 5.0 NaN + # 2 5.0 NaN + # 3 5.0 NaN + # 4 5.0 NaN + # + # PySpark: + # A B + # 0 2.0 1.0 + # 1 5.0 1.0 + # 2 5.0 1.0 + # 3 5.0 4.0 + # 4 5.0 9.0 + scol = F.when( + # By going through with max, it sets True after the first time it meets null. + F.max(self.spark.column.isNull()).over(window), + # Manually sets nulls given the column defined above. + F.lit(None), + ).otherwise(func(self.spark.column).over(window)) + + return self._with_new_scol(scol) + + def _cumsum(self, skipna: bool, part_cols: Sequence["ColumnOrName"] = ()) -> "Series": + psser = self + if isinstance(psser.spark.data_type, BooleanType): + psser = psser.spark.transform(lambda scol: scol.cast(LongType())) + elif not isinstance(psser.spark.data_type, NumericType): + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(psser.spark.data_type), + psser.spark.data_type.simpleString(), + ) + ) + return psser._cum(F.sum, skipna, part_cols) + + def _cumprod(self, skipna: bool, part_cols: Sequence["ColumnOrName"] = ()) -> "Series": + if isinstance(self.spark.data_type, BooleanType): + scol = self._cum( + lambda scol: F.min(F.coalesce(scol, F.lit(True))), skipna, part_cols + ).spark.column.cast(LongType()) + elif isinstance(self.spark.data_type, NumericType): + num_zeros = self._cum( + lambda scol: F.sum(F.when(scol == 0, 1).otherwise(0)), skipna, part_cols + ).spark.column + num_negatives = self._cum( + lambda scol: F.sum(F.when(scol < 0, 1).otherwise(0)), skipna, part_cols + ).spark.column + sign = F.when(num_negatives % 2 == 0, 1).otherwise(-1) + + abs_prod = F.exp( + self._cum(lambda scol: F.sum(F.log(F.abs(scol))), skipna, part_cols).spark.column + ) + + scol = F.when(num_zeros > 0, 0).otherwise(sign * abs_prod) + + if isinstance(self.spark.data_type, IntegralType): + scol = F.round(scol).cast(LongType()) + else: + raise TypeError( + "Could not convert {} ({}) to numeric".format( + spark_type_to_pandas_dtype(self.spark.data_type), + self.spark.data_type.simpleString(), + ) + ) + + return self._with_new_scol(scol) + + # ---------------------------------------------------------------------- + # Accessor Methods + # ---------------------------------------------------------------------- + dt = CachedAccessor("dt", DatetimeMethods) + str = CachedAccessor("str", StringMethods) + cat = CachedAccessor("cat", CategoricalAccessor) + plot = CachedAccessor("plot", PandasOnSparkPlotAccessor) + + # ---------------------------------------------------------------------- + + def _apply_series_op( + self, op: Callable[["Series"], Union["Series", PySparkColumn]], should_resolve: bool = False + ) -> "Series": + psser_or_scol = op(self) + if isinstance(psser_or_scol, Series): + psser = psser_or_scol + else: + psser = self._with_new_scol(psser_or_scol) + if should_resolve: + internal = psser._internal.resolved_copy + return first_series(DataFrame(internal)) + else: + return psser.copy() + + def _reduce_for_stat_function( + self, + sfun: Callable[["Series"], PySparkColumn], + name: str_type, + axis: Optional[Axis] = None, + numeric_only: bool = True, + skipna: bool = True, + **kwargs: Any, + ) -> Scalar: + """ + Applies sfun to the column and returns a scalar + + Parameters + ---------- + sfun : the stats function to be used for aggregation + name : original pandas API name. + axis : used only for sanity check because series only support index axis. + numeric_only : not used by this implementation, but passed down by stats functions. + skipna: exclude NA/null values when computing the result. + """ + axis = validate_axis(axis) + if axis == 1: + raise NotImplementedError("Series does not support columns axis.") + + if not skipna and get_option("compute.eager_check") and self.hasnans: + scol = F.first(F.lit(np.nan)) + else: + scol = sfun(self) + + min_count = kwargs.get("min_count", 0) + if min_count > 0: + scol = F.when(Frame._count_expr(self) >= min_count, scol) + + result = unpack_scalar(self._internal.spark_frame.select(scol)) + return result if result is not None else np.nan + + # Override the `groupby` to specify the actual return type annotation. + def groupby( + self, + by: Union[Name, "Series", List[Union[Name, "Series"]]], + axis: Axis = 0, + as_index: bool = True, + dropna: bool = True, + ) -> "SeriesGroupBy": + return cast( + "SeriesGroupBy", super().groupby(by=by, axis=axis, as_index=as_index, dropna=dropna) + ) + + groupby.__doc__ = Frame.groupby.__doc__ + + def _build_groupby( + self, by: List[Union["Series", Label]], as_index: bool, dropna: bool + ) -> "SeriesGroupBy": + from pyspark.pandas.groupby import SeriesGroupBy + + return SeriesGroupBy._build(self, by, as_index=as_index, dropna=dropna) + + def resample( + self, + rule: str_type, + closed: Optional[str_type] = None, + label: Optional[str_type] = None, + on: Optional["Series"] = None, + ) -> "SeriesResampler": + """ + Resample time-series data. + + Convenience method for frequency conversion and resampling of time series. + The object must have a datetime-like index (only support `DatetimeIndex` for now), + or the caller must pass the label of a datetime-like + series/index to the ``on`` keyword parameter. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + rule : str + The offset string or object representing target conversion. + Currently, supported units are {'Y', 'A', 'M', 'D', 'H', + 'T', 'MIN', 'S'}. + closed : {{'right', 'left'}}, default None + Which side of bin interval is closed. The default is 'left' + for all frequency offsets except for 'A', 'Y' and 'M' which all + have a default of 'right'. + label : {{'right', 'left'}}, default None + Which bin edge label to label bucket with. The default is 'left' + for all frequency offsets except for 'A', 'Y' and 'M' which all + have a default of 'right'. + on : Series, optional + For a DataFrame, column to use instead of index for resampling. + Column must be datetime-like. + + Returns + ------- + SeriesResampler + + + Examples + -------- + Start by creating a series with 9 one minute timestamps. + + >>> index = pd.date_range('1/1/2000', periods=9, freq='T') + >>> series = ps.Series(range(9), index=index, name='V') + >>> series + 2000-01-01 00:00:00 0 + 2000-01-01 00:01:00 1 + 2000-01-01 00:02:00 2 + 2000-01-01 00:03:00 3 + 2000-01-01 00:04:00 4 + 2000-01-01 00:05:00 5 + 2000-01-01 00:06:00 6 + 2000-01-01 00:07:00 7 + 2000-01-01 00:08:00 8 + Name: V, dtype: int64 + + Downsample the series into 3 minute bins and sum the values + of the timestamps falling into a bin. + + >>> series.resample('3T').sum().sort_index() + 2000-01-01 00:00:00 3.0 + 2000-01-01 00:03:00 12.0 + 2000-01-01 00:06:00 21.0 + Name: V, dtype: float64 + + Downsample the series into 3 minute bins as above, but label each + bin using the right edge instead of the left. Please note that the + value in the bucket used as the label is not included in the bucket, + which it labels. For example, in the original series the + bucket ``2000-01-01 00:03:00`` contains the value 3, but the summed + value in the resampled bucket with the label ``2000-01-01 00:03:00`` + does not include 3 (if it did, the summed value would be 6, not 3). + To include this value, close the right side of the bin interval as + illustrated in the example below this one. + + >>> series.resample('3T', label='right').sum().sort_index() + 2000-01-01 00:03:00 3.0 + 2000-01-01 00:06:00 12.0 + 2000-01-01 00:09:00 21.0 + Name: V, dtype: float64 + + Downsample the series into 3 minute bins as above, but close the right + side of the bin interval. + + >>> series.resample('3T', label='right', closed='right').sum().sort_index() + 2000-01-01 00:00:00 0.0 + 2000-01-01 00:03:00 6.0 + 2000-01-01 00:06:00 15.0 + 2000-01-01 00:09:00 15.0 + Name: V, dtype: float64 + + Upsample the series into 30 second bins. + + >>> series.resample('30S').sum().sort_index()[0:5] # Select first 5 rows + 2000-01-01 00:00:00 0.0 + 2000-01-01 00:00:30 0.0 + 2000-01-01 00:01:00 1.0 + 2000-01-01 00:01:30 0.0 + 2000-01-01 00:02:00 2.0 + Name: V, dtype: float64 + + See Also + -------- + DataFrame.resample : Resample a DataFrame. + groupby : Group by mapping, function, label, or list of labels. + """ + from pyspark.pandas.indexes import DatetimeIndex + from pyspark.pandas.resample import SeriesResampler + + if on is None and not isinstance(self.index, DatetimeIndex): + raise NotImplementedError("resample currently works only for DatetimeIndex") + if on is not None and not isinstance(as_spark_type(on.dtype), TimestampType): + raise NotImplementedError("`on` currently works only for TimestampType") + + agg_columns: List[ps.Series] = [] + column_label = self._internal.column_labels[0] + if isinstance(self._internal.spark_type_for(column_label), (NumericType, BooleanType)): + agg_columns.append(self) + + if len(agg_columns) == 0: + raise ValueError("No available aggregation columns!") + + return SeriesResampler( + psser=self, + resamplekey=on, + rule=rule, + closed=closed, + label=label, + agg_columns=agg_columns, + ) + + def __getitem__(self, key: Any) -> Any: + try: + if (isinstance(key, slice) and any(type(n) == int for n in [key.start, key.stop])) or ( + type(key) == int + and not isinstance(self.index.spark.data_type, (IntegerType, LongType)) + ): + # Seems like pandas Series always uses int as positional search when slicing + # with ints, searches based on index values when the value is int. + return self.iloc[key] + return self.loc[key] + except SparkPandasIndexingError: + raise KeyError( + "Key length ({}) exceeds index depth ({})".format( + len(key), self._internal.index_level + ) + ) + + def __getattr__(self, item: str_type) -> Any: + if item.startswith("__"): + raise AttributeError(item) + if hasattr(MissingPandasLikeSeries, item): + property_or_func = getattr(MissingPandasLikeSeries, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError("'Series' object has no attribute '{}'".format(item)) + + def _to_internal_pandas(self) -> pd.Series: + """ + Return a pandas Series directly from _internal to avoid overhead of copy. + + This method is for internal use only. + """ + return self._psdf._internal.to_pandas_frame[self.name] + + def __repr__(self) -> str_type: + max_display_count = get_option("display.max_rows") + if max_display_count is None: + return self._to_internal_pandas().to_string( + name=bool(self.name), dtype=bool(self.dtype) + ) + + pser = self._psdf._get_or_create_repr_pandas_cache(max_display_count)[self.name] + pser_length = len(pser) + pser = pser.iloc[:max_display_count] + if pser_length > max_display_count: + repr_string = pser.to_string(length=True) + rest, prev_footer = repr_string.rsplit("\n", 1) + match = REPR_PATTERN.search(prev_footer) + if match is not None: + length = match.group("length") + dtype_name = str(self.dtype.name) + if self.name is None: + footer = "\ndtype: {dtype}\nShowing only the first {length}".format( + length=length, dtype=pprint_thing(dtype_name) + ) + else: + footer = ( + "\nName: {name}, dtype: {dtype}" + "\nShowing only the first {length}".format( + length=length, name=self.name, dtype=pprint_thing(dtype_name) + ) + ) + return rest + footer + return pser.to_string(name=self.name, dtype=self.dtype) + + def __dir__(self) -> Iterable[str_type]: + if not isinstance(self.spark.data_type, StructType): + fields = [] + else: + fields = [f for f in self.spark.data_type.fieldNames() if " " not in f] + return list(super().__dir__()) + fields + + def __iter__(self) -> None: + return MissingPandasLikeSeries.__iter__(self) + + # In order to support the type hints such as Series[...]. See DataFrame.__class_getitem__. + def __class_getitem__(cls, params: Any) -> Type[SeriesType]: + return create_type_for_series_type(params) + + +def unpack_scalar(sdf: SparkDataFrame) -> Any: + """ + Takes a dataframe that is supposed to contain a single row with a single scalar value, + and returns this value. + """ + lst = sdf.limit(2).toPandas() + assert len(lst) == 1, (sdf, lst) + row = lst.iloc[0] + lst2 = list(row) + assert len(lst2) == 1, (row, lst2) + return lst2[0] + + +@overload +def first_series(df: DataFrame) -> Series: + ... + + +@overload +def first_series(df: pd.DataFrame) -> pd.Series: + ... + + +def first_series(df: Union[DataFrame, pd.DataFrame]) -> Union[Series, pd.Series]: + """ + Takes a DataFrame and returns the first column of the DataFrame as a Series + """ + assert isinstance(df, (DataFrame, pd.DataFrame)), type(df) + if isinstance(df, DataFrame): + return df._psser_for(df._internal.column_labels[0]) + else: + return df[df.columns[0]] + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.series + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.series.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.series tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.series, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/accessors.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/accessors.py new file mode 100644 index 0000000..f55f70e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/accessors.py @@ -0,0 +1,1281 @@ +# +# 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. +# + +""" +Spark related features. Usually, the features here are missing in pandas +but Spark has it. +""" +from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING, Callable, Generic, List, Optional, Union + +from pyspark import StorageLevel +from pyspark.sql import Column as PySparkColumn, DataFrame as PySparkDataFrame +from pyspark.sql.types import DataType, StructType + +from pyspark.pandas._typing import IndexOpsLike +from pyspark.pandas.internal import InternalField + +# For Supporting Spark Connect +from pyspark.sql.utils import get_column_class, get_dataframe_class + +if TYPE_CHECKING: + from pyspark.sql._typing import OptionalPrimitiveType + from pyspark._typing import PrimitiveType + + import pyspark.pandas as ps + from pyspark.pandas.frame import CachedDataFrame + + +class SparkIndexOpsMethods(Generic[IndexOpsLike], metaclass=ABCMeta): + """Spark related features. Usually, the features here are missing in pandas + but Spark has it.""" + + def __init__(self, data: IndexOpsLike): + self._data = data + + @property + def data_type(self) -> DataType: + """Returns the data type as defined by Spark, as a Spark DataType object.""" + return self._data._internal.spark_type_for(self._data._column_label) + + @property + def nullable(self) -> bool: + """Returns the nullability as defined by Spark.""" + return self._data._internal.spark_column_nullable_for(self._data._column_label) + + @property + def column(self) -> PySparkColumn: + """ + Spark Column object representing the Series/Index. + + .. note:: This Spark Column object is strictly stick to its base DataFrame the Series/Index + was derived from. + """ + return self._data._internal.spark_column_for(self._data._column_label) + + def transform(self, func: Callable[[PySparkColumn], PySparkColumn]) -> IndexOpsLike: + """ + Applies a function that takes and returns a Spark column. It allows natively + applying a Spark function and column APIs with the Spark column internally used + in Series or Index. The output length of the Spark column should be the same as input's. + + .. note:: It requires to have the same input and output length; therefore, + the aggregate Spark functions such as count does not work. + + Parameters + ---------- + func : function + Function to use for transforming the data by using Spark columns. + + Returns + ------- + Series or Index + + Raises + ------ + ValueError : If the output from the function is not a Spark column. + + Examples + -------- + >>> from pyspark.sql.functions import log + >>> df = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, columns=["a", "b"]) + >>> df + a b + 0 1 4 + 1 2 5 + 2 3 6 + + >>> df.a.spark.transform(lambda c: log(c)) + 0 0.000000 + 1 0.693147 + 2 1.098612 + Name: a, dtype: float64 + + >>> df.index.spark.transform(lambda c: c + 10) # doctest: +SKIP + Int64Index([10, 11, 12], dtype='int64') + + >>> df.a.spark.transform(lambda c: c + df.b.spark.column) + 0 5 + 1 7 + 2 9 + Name: a, dtype: int64 + """ + from pyspark.pandas import MultiIndex + + if isinstance(self._data, MultiIndex): + raise NotImplementedError("MultiIndex does not support spark.transform yet.") + output = func(self._data.spark.column) + Column = get_column_class() + if not isinstance(output, Column): + raise ValueError( + "The output of the function [%s] should be of a " + "pyspark.sql.Column; however, got [%s]." % (func, type(output)) + ) + # Trigger the resolution so it throws an exception if anything does wrong + # within the function, for example, + # `df1.a.spark.transform(lambda _: F.col("non-existent"))`. + field = InternalField.from_struct_field( + self._data._internal.spark_frame.select(output).schema.fields[0] + ) + return self._data._with_new_scol(scol=output, field=field) + + @property + @abstractmethod + def analyzed(self) -> IndexOpsLike: + pass + + +class SparkSeriesMethods(SparkIndexOpsMethods["ps.Series"]): + def apply(self, func: Callable[[PySparkColumn], PySparkColumn]) -> "ps.Series": + """ + Applies a function that takes and returns a Spark column. It allows to natively + apply a Spark function and column APIs with the Spark column internally used + in Series or Index. + + .. note:: It forces to lose the index and end up using the default index. It is + preferred to use :meth:`Series.spark.transform` or `:meth:`DataFrame.spark.apply` + with specifying the `index_col`. + + .. note:: It does not require to have the same length of the input and output. + However, it requires to create a new DataFrame internally which will require + to set `compute.ops_on_diff_frames` to compute even with the same origin + DataFrame is expensive, whereas :meth:`Series.spark.transform` does not + require it. + + Parameters + ---------- + func : function + Function to apply the function against the data by using Spark columns. + + Returns + ------- + Series + + Raises + ------ + ValueError : If the output from the function is not a Spark column. + + Examples + -------- + >>> from pyspark import pandas as ps + >>> from pyspark.sql.functions import count, lit + >>> df = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, columns=["a", "b"]) + >>> df + a b + 0 1 4 + 1 2 5 + 2 3 6 + + >>> df.a.spark.apply(lambda c: count(c)) + 0 3 + Name: a, dtype: int64 + + >>> df.a.spark.apply(lambda c: c + df.b.spark.column) + 0 5 + 1 7 + 2 9 + Name: a, dtype: int64 + """ + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series, first_series + from pyspark.pandas.internal import HIDDEN_COLUMNS + + output = func(self._data.spark.column) + Column = get_column_class() + if not isinstance(output, Column): + raise ValueError( + "The output of the function [%s] should be of a " + "pyspark.sql.Column; however, got [%s]." % (func, type(output)) + ) + assert isinstance(self._data, Series) + + sdf = self._data._internal.spark_frame.drop(*HIDDEN_COLUMNS).select(output) + # Lose index. + return first_series(DataFrame(sdf)).rename(self._data.name) + + @property + def analyzed(self) -> "ps.Series": + """ + Returns a new Series with the analyzed Spark DataFrame. + + After multiple operations, the underlying Spark plan could grow huge + and make the Spark planner take a long time to finish the planning. + + This function is for the workaround to avoid it. + + .. note:: After analyzing, operations between the analyzed Series and the original one + will **NOT** work without setting a config `compute.ops_on_diff_frames` to `True`. + + Returns + ------- + Series + + Examples + -------- + >>> ser = ps.Series([1, 2, 3]) + >>> ser + 0 1 + 1 2 + 2 3 + dtype: int64 + + The analyzed one should return the same value. + + >>> ser.spark.analyzed + 0 1 + 1 2 + 2 3 + dtype: int64 + + However, it won't work with the same anchor Series. + + >>> ser + ser.spark.analyzed + Traceback (most recent call last): + ... + ValueError: ... enable 'compute.ops_on_diff_frames' option. + + >>> with ps.option_context('compute.ops_on_diff_frames', True): + ... (ser + ser.spark.analyzed).sort_index() + 0 2 + 1 4 + 2 6 + dtype: int64 + """ + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import first_series + + return first_series(DataFrame(self._data._internal.resolved_copy)) + + +class SparkIndexMethods(SparkIndexOpsMethods["ps.Index"]): + @property + def analyzed(self) -> "ps.Index": + """ + Returns a new Index with the analyzed Spark DataFrame. + + After multiple operations, the underlying Spark plan could grow huge + and make the Spark planner take a long time to finish the planning. + + This function is for the workaround to avoid it. + + .. note:: After analyzing, operations between the analyzed Series and the original one + will **NOT** work without setting a config `compute.ops_on_diff_frames` to `True`. + + Returns + ------- + Index + + Examples + -------- + >>> import pyspark.pandas as ps + >>> idx = ps.Index([1, 2, 3]) + >>> idx # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + The analyzed one should return the same value. + + >>> idx.spark.analyzed # doctest: +SKIP + Int64Index([1, 2, 3], dtype='int64') + + However, it won't work with the same anchor Index. + + >>> idx + idx.spark.analyzed + Traceback (most recent call last): + ... + ValueError: ... enable 'compute.ops_on_diff_frames' option. + + >>> with ps.option_context('compute.ops_on_diff_frames', True): + ... (idx + idx.spark.analyzed).sort_values() # doctest: +SKIP + Int64Index([2, 4, 6], dtype='int64') + """ + from pyspark.pandas.frame import DataFrame + + return DataFrame(self._data._internal.resolved_copy).index + + +class SparkFrameMethods: + """Spark related features. Usually, the features here are missing in pandas + but Spark has it.""" + + def __init__(self, frame: "ps.DataFrame"): + self._psdf = frame + + def schema(self, index_col: Optional[Union[str, List[str]]] = None) -> StructType: + """ + Returns the underlying Spark schema. + + Returns + ------- + pyspark.sql.types.StructType + The underlying Spark schema. + + Parameters + ---------- + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + + Examples + -------- + >>> df = ps.DataFrame({'a': list('abc'), + ... 'b': list(range(1, 4)), + ... 'c': np.arange(3, 6).astype('i1'), + ... 'd': np.arange(4.0, 7.0, dtype='float64'), + ... 'e': [True, False, True], + ... 'f': pd.date_range('20130101', periods=3)}, + ... columns=['a', 'b', 'c', 'd', 'e', 'f']) + >>> df.spark.schema().simpleString() + 'struct' + >>> df.spark.schema(index_col='index').simpleString() + 'struct' + """ + return self.frame(index_col).schema + + def print_schema(self, index_col: Optional[Union[str, List[str]]] = None) -> None: + """ + Prints out the underlying Spark schema in the tree format. + + Parameters + ---------- + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + + Returns + ------- + None + + Examples + -------- + >>> df = ps.DataFrame({'a': list('abc'), + ... 'b': list(range(1, 4)), + ... 'c': np.arange(3, 6).astype('i1'), + ... 'd': np.arange(4.0, 7.0, dtype='float64'), + ... 'e': [True, False, True], + ... 'f': pd.date_range('20130101', periods=3)}, + ... columns=['a', 'b', 'c', 'd', 'e', 'f']) + >>> df.spark.print_schema() # doctest: +NORMALIZE_WHITESPACE + root + |-- a: string (nullable = false) + |-- b: long (nullable = false) + |-- c: byte (nullable = false) + |-- d: double (nullable = false) + |-- e: boolean (nullable = false) + |-- f: timestamp (nullable = false) + >>> df.spark.print_schema(index_col='index') # doctest: +NORMALIZE_WHITESPACE + root + |-- index: long (nullable = false) + |-- a: string (nullable = false) + |-- b: long (nullable = false) + |-- c: byte (nullable = false) + |-- d: double (nullable = false) + |-- e: boolean (nullable = false) + |-- f: timestamp (nullable = false) + """ + self.frame(index_col).printSchema() + + def frame(self, index_col: Optional[Union[str, List[str]]] = None) -> PySparkDataFrame: + """ + Return the current DataFrame as a Spark DataFrame. :meth:`DataFrame.spark.frame` is an + alias of :meth:`DataFrame.to_spark`. + + Parameters + ---------- + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + + See Also + -------- + DataFrame.to_spark + DataFrame.pandas_api + DataFrame.spark.frame + + Examples + -------- + By default, this method loses the index as below. + + >>> df = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> df.to_spark().show() # doctest: +NORMALIZE_WHITESPACE + +---+---+---+ + | a| b| c| + +---+---+---+ + | 1| 4| 7| + | 2| 5| 8| + | 3| 6| 9| + +---+---+---+ + + >>> df = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> df.spark.frame().show() # doctest: +NORMALIZE_WHITESPACE + +---+---+---+ + | a| b| c| + +---+---+---+ + | 1| 4| 7| + | 2| 5| 8| + | 3| 6| 9| + +---+---+---+ + + If `index_col` is set, it keeps the index column as specified. + + >>> df.to_spark(index_col="index").show() # doctest: +NORMALIZE_WHITESPACE + +-----+---+---+---+ + |index| a| b| c| + +-----+---+---+---+ + | 0| 1| 4| 7| + | 1| 2| 5| 8| + | 2| 3| 6| 9| + +-----+---+---+---+ + + Keeping an index column is useful when you want to call some Spark APIs and + convert it back to pandas-on-Spark DataFrame without creating a default index, which + can affect performance. + + >>> spark_df = df.to_spark(index_col="index") + >>> spark_df = spark_df.filter("a == 2") + >>> spark_df.pandas_api(index_col="index") # doctest: +NORMALIZE_WHITESPACE + a b c + index + 1 2 5 8 + + In case of multi-index, specify a list to `index_col`. + + >>> new_df = df.set_index("a", append=True) + >>> new_spark_df = new_df.to_spark(index_col=["index_1", "index_2"]) + >>> new_spark_df.show() # doctest: +NORMALIZE_WHITESPACE + +-------+-------+---+---+ + |index_1|index_2| b| c| + +-------+-------+---+---+ + | 0| 1| 4| 7| + | 1| 2| 5| 8| + | 2| 3| 6| 9| + +-------+-------+---+---+ + + Can be converted back to pandas-on-Spark DataFrame. + + >>> new_spark_df.pandas_api( + ... index_col=["index_1", "index_2"]) # doctest: +NORMALIZE_WHITESPACE + b c + index_1 index_2 + 0 1 4 7 + 1 2 5 8 + 2 3 6 9 + """ + from pyspark.pandas.utils import name_like_string + + psdf = self._psdf + + data_column_names = [] + data_columns = [] + for i, (label, spark_column, column_name) in enumerate( + zip( + psdf._internal.column_labels, + psdf._internal.data_spark_columns, + psdf._internal.data_spark_column_names, + ) + ): + name = str(i) if label is None else name_like_string(label) + data_column_names.append(name) + if column_name != name: + spark_column = spark_column.alias(name) + data_columns.append(spark_column) + + if index_col is None: + return psdf._internal.spark_frame.select(data_columns) + else: + if isinstance(index_col, str): + index_col = [index_col] + + old_index_scols = psdf._internal.index_spark_columns + + if len(index_col) != len(old_index_scols): + raise ValueError( + "length of index columns is %s; however, the length of the given " + "'index_col' is %s." % (len(old_index_scols), len(index_col)) + ) + + if any(col in data_column_names for col in index_col): + raise ValueError("'index_col' cannot be overlapped with other columns.") + + new_index_scols = [ + index_scol.alias(col) for index_scol, col in zip(old_index_scols, index_col) + ] + return psdf._internal.spark_frame.select(new_index_scols + data_columns) + + def cache(self) -> "CachedDataFrame": + """ + Yields and caches the current DataFrame. + + The pandas-on-Spark DataFrame is yielded as a protected resource and its corresponding + data is cached which gets uncached after execution goes off the context. + + If you want to specify the StorageLevel manually, use :meth:`DataFrame.spark.persist` + + See Also + -------- + DataFrame.spark.persist + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df + dogs cats + 0 0.2 0.3 + 1 0.0 0.6 + 2 0.6 0.0 + 3 0.2 0.1 + + >>> with df.spark.cache() as cached_df: + ... print(cached_df.count()) + ... + dogs 4 + cats 4 + dtype: int64 + + >>> df = df.spark.cache() + >>> df.to_pandas().mean(axis=1) + 0 0.25 + 1 0.30 + 2 0.30 + 3 0.15 + dtype: float64 + + To uncache the dataframe, use `unpersist` function + + >>> df.spark.unpersist() + """ + from pyspark.pandas.frame import CachedDataFrame + + self._psdf._update_internal_frame( + self._psdf._internal.resolved_copy, check_same_anchor=False + ) + return CachedDataFrame(self._psdf._internal) + + def persist( + self, storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK + ) -> "CachedDataFrame": + """ + Yields and caches the current DataFrame with a specific StorageLevel. + If a StorageLevel is not given, the `MEMORY_AND_DISK` level is used by default like PySpark. + + The pandas-on-Spark DataFrame is yielded as a protected resource and its corresponding + data is cached which gets uncached after execution goes off the context. + + See Also + -------- + DataFrame.spark.cache + + Examples + -------- + >>> import pyspark + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df + dogs cats + 0 0.2 0.3 + 1 0.0 0.6 + 2 0.6 0.0 + 3 0.2 0.1 + + Set the StorageLevel to `MEMORY_ONLY`. + + >>> with df.spark.persist(pyspark.StorageLevel.MEMORY_ONLY) as cached_df: + ... print(cached_df.spark.storage_level) + ... print(cached_df.count()) + ... + Memory Serialized 1x Replicated + dogs 4 + cats 4 + dtype: int64 + + Set the StorageLevel to `DISK_ONLY`. + + >>> with df.spark.persist(pyspark.StorageLevel.DISK_ONLY) as cached_df: + ... print(cached_df.spark.storage_level) + ... print(cached_df.count()) + ... + Disk Serialized 1x Replicated + dogs 4 + cats 4 + dtype: int64 + + If a StorageLevel is not given, it uses `MEMORY_AND_DISK` by default. + + >>> with df.spark.persist() as cached_df: + ... print(cached_df.spark.storage_level) + ... print(cached_df.count()) + ... + Disk Memory Serialized 1x Replicated + dogs 4 + cats 4 + dtype: int64 + + >>> df = df.spark.persist() + >>> df.to_pandas().mean(axis=1) + 0 0.25 + 1 0.30 + 2 0.30 + 3 0.15 + dtype: float64 + + To uncache the dataframe, use `unpersist` function + + >>> df.spark.unpersist() + """ + from pyspark.pandas.frame import CachedDataFrame + + self._psdf._update_internal_frame( + self._psdf._internal.resolved_copy, check_same_anchor=False + ) + return CachedDataFrame(self._psdf._internal, storage_level=storage_level) + + def hint(self, name: str, *parameters: "PrimitiveType") -> "ps.DataFrame": + """ + Specifies some hint on the current DataFrame. + + Parameters + ---------- + name : A name of the hint. + parameters : Optional parameters. + + Returns + ------- + ret : DataFrame with the hint. + + See Also + -------- + broadcast : Marks a DataFrame as small enough for use in broadcast joins. + + Examples + -------- + >>> df1 = ps.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [1, 2, 3, 5]}, + ... columns=['lkey', 'value']).set_index('lkey') + >>> df2 = ps.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'], + ... 'value': [5, 6, 7, 8]}, + ... columns=['rkey', 'value']).set_index('rkey') + >>> merged = df1.merge(df2.spark.hint("broadcast"), left_index=True, right_index=True) + >>> merged.spark.explain() # doctest: +ELLIPSIS + == Physical Plan == + ... + ...BroadcastHashJoin... + ... + """ + from pyspark.pandas.frame import DataFrame + + internal = self._psdf._internal.resolved_copy + return DataFrame(internal.with_new_sdf(internal.spark_frame.hint(name, *parameters))) + + def to_table( + self, + name: str, + format: Optional[str] = None, + mode: str = "overwrite", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """ + Write the DataFrame into a Spark table. :meth:`DataFrame.spark.to_table` + is an alias of :meth:`DataFrame.to_table`. + + Parameters + ---------- + name : str, required + Table name in Spark. + format : string, optional + Specifies the output data source format. Some common ones are: + + - 'delta' + - 'parquet' + - 'orc' + - 'json' + - 'csv' + + mode : str {'append', 'overwrite', 'ignore', 'error', 'errorifexists'}, default + 'overwrite'. Specifies the behavior of the save operation when the table exists + already. + + - 'append': Append the new data to existing data. + - 'overwrite': Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + + partition_cols : str or list of str, optional, default None + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + options + Additional options passed directly to Spark. + + Returns + ------- + None + + See Also + -------- + read_table + DataFrame.to_spark_io + DataFrame.spark.to_spark_io + DataFrame.to_parquet + + Examples + -------- + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df + date country code + 0 2012-01-31 12:00:00 KR 1 + 1 2012-02-29 12:00:00 US 2 + 2 2012-03-31 12:00:00 JP 3 + + >>> df.to_table('%s.my_table' % db, partition_cols='date') + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") # type: ignore[assignment] + + self._psdf.spark.frame(index_col=index_col).write.saveAsTable( + name=name, format=format, mode=mode, partitionBy=partition_cols, **options + ) + + def to_spark_io( + self, + path: Optional[str] = None, + format: Optional[str] = None, + mode: str = "overwrite", + partition_cols: Optional[Union[str, List[str]]] = None, + index_col: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """Write the DataFrame out to a Spark data source. :meth:`DataFrame.spark.to_spark_io` + is an alias of :meth:`DataFrame.to_spark_io`. + + Parameters + ---------- + path : string, optional + Path to the data source. + format : string, optional + Specifies the output data source format. Some common ones are: + + - 'delta' + - 'parquet' + - 'orc' + - 'json' + - 'csv' + mode : str {'append', 'overwrite', 'ignore', 'error', 'errorifexists'}, default + 'overwrite'. Specifies the behavior of the save operation when data already exists. + + - 'append': Append the new data to existing data. + - 'overwrite': Overwrite existing data. + - 'ignore': Silently ignore this operation if data already exists. + - 'error' or 'errorifexists': Throw an exception if data already exists. + partition_cols : str or list of str, optional + Names of partitioning columns + index_col: str or list of str, optional, default: None + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + options : dict + All other options passed directly into Spark's data source. + + Returns + ------- + None + + See Also + -------- + read_spark_io + DataFrame.to_delta + DataFrame.to_parquet + DataFrame.to_table + DataFrame.to_spark_io + DataFrame.spark.to_spark_io + + Examples + -------- + >>> df = ps.DataFrame(dict( + ... date=list(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M')), + ... country=['KR', 'US', 'JP'], + ... code=[1, 2 ,3]), columns=['date', 'country', 'code']) + >>> df + date country code + 0 2012-01-31 12:00:00 KR 1 + 1 2012-02-29 12:00:00 US 2 + 2 2012-03-31 12:00:00 JP 3 + + >>> df.to_spark_io(path='%s/to_spark_io/foo.json' % path, format='json') + """ + if "options" in options and isinstance(options.get("options"), dict) and len(options) == 1: + options = options.get("options") # type: ignore[assignment] + + self._psdf.spark.frame(index_col=index_col).write.save( + path=path, format=format, mode=mode, partitionBy=partition_cols, **options + ) + + def explain(self, extended: Optional[bool] = None, mode: Optional[str] = None) -> None: + """ + Prints the underlying (logical and physical) Spark plans to the console for debugging + purpose. + + Parameters + ---------- + extended : boolean, default ``False``. + If ``False``, prints only the physical plan. + mode : string, default ``None``. + The expected output format of plans. + + Returns + ------- + None + + Examples + -------- + >>> df = ps.DataFrame({'id': range(10)}) + >>> df.spark.explain() # doctest: +ELLIPSIS + == Physical Plan == + ... + + >>> df.spark.explain(True) # doctest: +ELLIPSIS + == Parsed Logical Plan == + ... + == Analyzed Logical Plan == + ... + == Optimized Logical Plan == + ... + == Physical Plan == + ... + + >>> df.spark.explain("extended") # doctest: +ELLIPSIS + == Parsed Logical Plan == + ... + == Analyzed Logical Plan == + ... + == Optimized Logical Plan == + ... + == Physical Plan == + ... + + >>> df.spark.explain(mode="extended") # doctest: +ELLIPSIS + == Parsed Logical Plan == + ... + == Analyzed Logical Plan == + ... + == Optimized Logical Plan == + ... + == Physical Plan == + ... + """ + self._psdf._internal.to_internal_spark_frame.explain(extended, mode) + + def apply( + self, + func: Callable[[PySparkDataFrame], PySparkDataFrame], + index_col: Optional[Union[str, List[str]]] = None, + ) -> "ps.DataFrame": + """ + Applies a function that takes and returns a Spark DataFrame. It allows natively + apply a Spark function and column APIs with the Spark column internally used + in Series or Index. + + .. note:: set `index_col` and keep the column named as so in the output Spark + DataFrame to avoid using the default index to prevent performance penalty. + If you omit `index_col`, it will use default index which is potentially + expensive in general. + + .. note:: it will lose column labels. This is a synonym of + ``func(psdf.to_spark(index_col)).pandas_api(index_col)``. + + Parameters + ---------- + func : function + Function to apply the function against the data by using Spark DataFrame. + + Returns + ------- + DataFrame + + Raises + ------ + ValueError : If the output from the function is not a Spark DataFrame. + + Examples + -------- + >>> psdf = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, columns=["a", "b"]) + >>> psdf + a b + 0 1 4 + 1 2 5 + 2 3 6 + + >>> psdf.spark.apply( + ... lambda sdf: sdf.selectExpr("a + b as c", "index"), index_col="index") + ... # doctest: +NORMALIZE_WHITESPACE + c + index + 0 5 + 1 7 + 2 9 + + The case below ends up with using the default index, which should be avoided + if possible. + + >>> psdf.spark.apply(lambda sdf: sdf.groupby("a").count().sort("a")) + a count + 0 1 1 + 1 2 1 + 2 3 1 + """ + output = func(self.frame(index_col)) + SparkDataFrame = get_dataframe_class() + if not isinstance(output, SparkDataFrame): + raise ValueError( + "The output of the function [%s] should be of a " + "pyspark.sql.DataFrame; however, got [%s]." % (func, type(output)) + ) + return output.pandas_api(index_col) + + def repartition(self, num_partitions: int) -> "ps.DataFrame": + """ + Returns a new DataFrame partitioned by the given partitioning expressions. The + resulting DataFrame is hash partitioned. + + Parameters + ---------- + num_partitions : int + The target number of partitions. + + Returns + ------- + DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame({"age": [5, 5, 2, 2], + ... "name": ["Bob", "Bob", "Alice", "Alice"]}).set_index("age") + >>> psdf.sort_index() # doctest: +NORMALIZE_WHITESPACE + name + age + 2 Alice + 2 Alice + 5 Bob + 5 Bob + >>> new_psdf = psdf.spark.repartition(7) + >>> new_psdf.to_spark().rdd.getNumPartitions() + 7 + >>> new_psdf.sort_index() # doctest: +NORMALIZE_WHITESPACE + name + age + 2 Alice + 2 Alice + 5 Bob + 5 Bob + """ + from pyspark.pandas.frame import DataFrame + + internal = self._psdf._internal.resolved_copy + repartitioned_sdf = internal.spark_frame.repartition(num_partitions) + return DataFrame(internal.with_new_sdf(repartitioned_sdf)) + + def coalesce(self, num_partitions: int) -> "ps.DataFrame": + """ + Returns a new DataFrame that has exactly `num_partitions` partitions. + + .. note:: This operation results in a narrow dependency, e.g. if you go from 1000 + partitions to 100 partitions, there will not be a shuffle, instead each of the 100 new + partitions will claim 10 of the current partitions. If a larger number of partitions is + requested, it will stay at the current number of partitions. However, if you're doing a + drastic coalesce, e.g. to num_partitions = 1, this may result in your computation taking + place on fewer nodes than you like (e.g. one node in the case of num_partitions = 1). To + avoid this, you can call repartition(). This will add a shuffle step, but means the + current upstream partitions will be executed in parallel (per whatever the current + partitioning is). + + Parameters + ---------- + num_partitions : int + The target number of partitions. + + Returns + ------- + DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame({"age": [5, 5, 2, 2], + ... "name": ["Bob", "Bob", "Alice", "Alice"]}).set_index("age") + >>> psdf.sort_index() # doctest: +NORMALIZE_WHITESPACE + name + age + 2 Alice + 2 Alice + 5 Bob + 5 Bob + >>> new_psdf = psdf.spark.coalesce(1) + >>> new_psdf.to_spark().rdd.getNumPartitions() + 1 + >>> new_psdf.sort_index() # doctest: +NORMALIZE_WHITESPACE + name + age + 2 Alice + 2 Alice + 5 Bob + 5 Bob + """ + from pyspark.pandas.frame import DataFrame + + internal = self._psdf._internal.resolved_copy + coalesced_sdf = internal.spark_frame.coalesce(num_partitions) + return DataFrame(internal.with_new_sdf(coalesced_sdf)) + + def checkpoint(self, eager: bool = True) -> "ps.DataFrame": + """Returns a checkpointed version of this DataFrame. + + Checkpointing can be used to truncate the logical plan of this DataFrame, which is + especially useful in iterative algorithms where the plan may grow exponentially. It will be + saved to files inside the checkpoint directory set with `SparkContext.setCheckpointDir`. + + Parameters + ---------- + eager : bool + Whether to checkpoint this DataFrame immediately + + Returns + ------- + DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame({"a": ["a", "b", "c"]}) + >>> psdf + a + 0 a + 1 b + 2 c + >>> new_psdf = psdf.spark.checkpoint() # doctest: +SKIP + >>> new_psdf # doctest: +SKIP + a + 0 a + 1 b + 2 c + """ + from pyspark.pandas.frame import DataFrame + + internal = self._psdf._internal.resolved_copy + checkpointed_sdf = internal.spark_frame.checkpoint(eager) + return DataFrame(internal.with_new_sdf(checkpointed_sdf)) + + def local_checkpoint(self, eager: bool = True) -> "ps.DataFrame": + """Returns a locally checkpointed version of this DataFrame. + + Checkpointing can be used to truncate the logical plan of this DataFrame, which is + especially useful in iterative algorithms where the plan may grow exponentially. Local + checkpoints are stored in the executors using the caching subsystem and therefore they are + not reliable. + + Parameters + ---------- + eager : bool + Whether to locally checkpoint this DataFrame immediately + + Returns + ------- + DataFrame + + Examples + -------- + >>> psdf = ps.DataFrame({"a": ["a", "b", "c"]}) + >>> psdf + a + 0 a + 1 b + 2 c + >>> new_psdf = psdf.spark.local_checkpoint() + >>> new_psdf + a + 0 a + 1 b + 2 c + """ + from pyspark.pandas.frame import DataFrame + + internal = self._psdf._internal.resolved_copy + checkpointed_sdf = internal.spark_frame.localCheckpoint(eager) + return DataFrame(internal.with_new_sdf(checkpointed_sdf)) + + @property + def analyzed(self) -> "ps.DataFrame": + """ + Returns a new DataFrame with the analyzed Spark DataFrame. + + After multiple operations, the underlying Spark plan could grow huge + and make the Spark planner take a long time to finish the planning. + + This function is for the workaround to avoid it. + + .. note:: After analysis, operations between the analyzed DataFrame and the original one + will **NOT** work without setting a config `compute.ops_on_diff_frames` to `True`. + + Returns + ------- + DataFrame + + Examples + -------- + >>> df = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, columns=["a", "b"]) + >>> df + a b + 0 1 4 + 1 2 5 + 2 3 6 + + The analyzed one should return the same value. + + >>> df.spark.analyzed + a b + 0 1 4 + 1 2 5 + 2 3 6 + + However, it won't work with the same anchor Series. + + >>> df + df.spark.analyzed + Traceback (most recent call last): + ... + ValueError: ... enable 'compute.ops_on_diff_frames' option. + + >>> with ps.option_context('compute.ops_on_diff_frames', True): + ... (df + df.spark.analyzed).sort_index() + a b + 0 2 8 + 1 4 10 + 2 6 12 + """ + from pyspark.pandas.frame import DataFrame + + return DataFrame(self._psdf._internal.resolved_copy) + + +class CachedSparkFrameMethods(SparkFrameMethods): + """Spark related features for cached DataFrame. This is usually created via + `df.spark.cache()`.""" + + def __init__(self, frame: "CachedDataFrame"): + super().__init__(frame) + + @property + def storage_level(self) -> StorageLevel: + """ + Return the storage level of this cache. + + Examples + -------- + >>> import pyspark.pandas as ps + >>> import pyspark + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df + dogs cats + 0 0.2 0.3 + 1 0.0 0.6 + 2 0.6 0.0 + 3 0.2 0.1 + + >>> with df.spark.cache() as cached_df: + ... print(cached_df.spark.storage_level) + ... + Disk Memory Deserialized 1x Replicated + + Set the StorageLevel to `MEMORY_ONLY`. + + >>> with df.spark.persist(pyspark.StorageLevel.MEMORY_ONLY) as cached_df: + ... print(cached_df.spark.storage_level) + ... + Memory Serialized 1x Replicated + """ + return self._psdf._cached.storageLevel + + def unpersist(self) -> None: + """ + The `unpersist` function is used to uncache the pandas-on-Spark DataFrame when it + is not used with the `with` statement. + + Returns + ------- + None + + Examples + -------- + >>> df = ps.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)], + ... columns=['dogs', 'cats']) + >>> df = df.spark.cache() + + To uncache the dataframe, use `unpersist` function + + >>> df.spark.unpersist() + """ + if self._psdf._cached.is_cached: + self._psdf._cached.unpersist() + + +def _test() -> None: + import os + import doctest + import shutil + import sys + import tempfile + import uuid + import numpy + import pandas + from pyspark.sql import SparkSession + import pyspark.pandas.spark.accessors + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.spark.accessors.__dict__.copy() + globs["np"] = numpy + globs["pd"] = pandas + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.spark.accessors tests") + .getOrCreate() + ) + + db_name = "db%s" % str(uuid.uuid4()).replace("-", "") + spark.sql("CREATE DATABASE %s" % db_name) + globs["db"] = db_name + + path = tempfile.mkdtemp() + globs["path"] = path + + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.spark.accessors, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + + shutil.rmtree(path, ignore_errors=True) + spark.sql("DROP DATABASE IF EXISTS %s CASCADE" % db_name) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/functions.py new file mode 100644 index 0000000..d6f6c6f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/functions.py @@ -0,0 +1,203 @@ +# +# 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. +# +""" +Additional Spark functions used in pandas-on-Spark. +""" +from typing import Union + +from pyspark import SparkContext +import pyspark.sql.functions as F +from pyspark.sql.column import Column + +# For supporting Spark Connect +from pyspark.sql.utils import is_remote + + +def product(col: Column, dropna: bool) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_product", + col, # type: ignore[arg-type] + lit(dropna), + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasProduct(col._jc, dropna)) + + +def stddev(col: Column, ddof: int) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_stddev", + col, # type: ignore[arg-type] + lit(ddof), + ) + + else: + + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasStddev(col._jc, ddof)) + + +def var(col: Column, ddof: int) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_var", + col, # type: ignore[arg-type] + lit(ddof), + ) + + else: + + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasVariance(col._jc, ddof)) + + +def skew(col: Column) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_skew", + col, # type: ignore[arg-type] + ) + + else: + + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasSkewness(col._jc)) + + +def kurt(col: Column) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_kurt", + col, # type: ignore[arg-type] + ) + + else: + + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasKurtosis(col._jc)) + + +def mode(col: Column, dropna: bool) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_mode", + col, # type: ignore[arg-type] + lit(dropna), + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasMode(col._jc, dropna)) + + +def covar(col1: Column, col2: Column, ddof: int) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "pandas_covar", + col1, # type: ignore[arg-type] + col2, # type: ignore[arg-type] + lit(ddof), + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.pandasCovar(col1._jc, col2._jc, ddof)) + + +def repeat(col: Column, n: Union[int, Column]) -> Column: + """ + Repeats a string column n times, and returns it as a new string column. + """ + _n = F.lit(n) if isinstance(n, int) else n + return F.call_udf("repeat", col, _n) + + +def ewm(col: Column, alpha: float, ignore_na: bool) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "ewm", + col, # type: ignore[arg-type] + lit(alpha), + lit(ignore_na), + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.ewm(col._jc, alpha, ignore_na)) + + +def last_non_null(col: Column) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns + + return _invoke_function_over_columns( # type: ignore[return-value] + "last_non_null", + col, # type: ignore[arg-type] + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.lastNonNull(col._jc)) + + +def null_index(col: Column) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns + + return _invoke_function_over_columns( # type: ignore[return-value] + "null_index", + col, # type: ignore[arg-type] + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.nullIndex(col._jc)) + + +def timestampdiff(unit: str, start: Column, end: Column) -> Column: + if is_remote(): + from pyspark.sql.connect.functions import _invoke_function_over_columns, lit + + return _invoke_function_over_columns( # type: ignore[return-value] + "timestampdiff", + lit(unit), + start, # type: ignore[arg-type] + end, # type: ignore[arg-type] + ) + + else: + sc = SparkContext._active_spark_context + return Column(sc._jvm.PythonSQLUtils.timestampDiff(unit, start._jc, end._jc)) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/utils.py new file mode 100644 index 0000000..9b8b5bb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/spark/utils.py @@ -0,0 +1,202 @@ +# +# 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. +# +""" +Helpers and utilities to deal with PySpark instances +""" +from typing import overload + +from pyspark.sql.types import DecimalType, StructType, MapType, ArrayType, StructField, DataType + + +@overload +def as_nullable_spark_type(dt: StructType) -> StructType: + ... + + +@overload +def as_nullable_spark_type(dt: ArrayType) -> ArrayType: + ... + + +@overload +def as_nullable_spark_type(dt: MapType) -> MapType: + ... + + +@overload +def as_nullable_spark_type(dt: DataType) -> DataType: + ... + + +def as_nullable_spark_type(dt: DataType) -> DataType: + """ + Returns a nullable schema or data types. + + Examples + -------- + >>> from pyspark.sql.types import * + >>> as_nullable_spark_type(StructType([ + ... StructField("A", IntegerType(), True), + ... StructField("B", FloatType(), False)])) # doctest: +NORMALIZE_WHITESPACE + StructType([StructField('A', IntegerType(), True), StructField('B', FloatType(), True)]) + + >>> as_nullable_spark_type(StructType([ + ... StructField("A", + ... StructType([ + ... StructField('a', + ... MapType(IntegerType(), + ... ArrayType(IntegerType(), False), False), False), + ... StructField('b', StringType(), True)])), + ... StructField("B", FloatType(), False)])) # doctest: +NORMALIZE_WHITESPACE + StructType([StructField('A', + StructType([StructField('a', + MapType(IntegerType(), + ArrayType(IntegerType(), True), True), True), + StructField('b', StringType(), True)]), True), + StructField('B', FloatType(), True)]) + """ + if isinstance(dt, StructType): + new_fields = [] + for field in dt.fields: + new_fields.append( + StructField( + field.name, + as_nullable_spark_type(field.dataType), + nullable=True, + metadata=field.metadata, + ) + ) + return StructType(new_fields) + elif isinstance(dt, ArrayType): + return ArrayType(as_nullable_spark_type(dt.elementType), containsNull=True) + elif isinstance(dt, MapType): + return MapType( + as_nullable_spark_type(dt.keyType), + as_nullable_spark_type(dt.valueType), + valueContainsNull=True, + ) + else: + return dt + + +@overload +def force_decimal_precision_scale( + dt: StructType, *, precision: int = ..., scale: int = ... +) -> StructType: + ... + + +@overload +def force_decimal_precision_scale( + dt: ArrayType, *, precision: int = ..., scale: int = ... +) -> ArrayType: + ... + + +@overload +def force_decimal_precision_scale( + dt: MapType, *, precision: int = ..., scale: int = ... +) -> MapType: + ... + + +@overload +def force_decimal_precision_scale( + dt: DataType, *, precision: int = ..., scale: int = ... +) -> DataType: + ... + + +def force_decimal_precision_scale( + dt: DataType, *, precision: int = 38, scale: int = 18 +) -> DataType: + """ + Returns a data type with a fixed decimal type. + + The precision and scale of the decimal type are fixed with the given values. + + Examples + -------- + >>> from pyspark.sql.types import * + >>> force_decimal_precision_scale(StructType([ + ... StructField("A", DecimalType(10, 0), True), + ... StructField("B", DecimalType(14, 7), False)])) # doctest: +NORMALIZE_WHITESPACE + StructType([StructField('A', DecimalType(38,18), True), + StructField('B', DecimalType(38,18), False)]) + + >>> force_decimal_precision_scale(StructType([ + ... StructField("A", + ... StructType([ + ... StructField('a', + ... MapType(DecimalType(5, 0), + ... ArrayType(DecimalType(20, 0), False), False), False), + ... StructField('b', StringType(), True)])), + ... StructField("B", DecimalType(30, 15), False)]), + ... precision=30, scale=15) # doctest: +NORMALIZE_WHITESPACE + StructType([StructField('A', + StructType([StructField('a', + MapType(DecimalType(30,15), + ArrayType(DecimalType(30,15), False), False), False), + StructField('b', StringType(), True)]), True), + StructField('B', DecimalType(30,15), False)]) + """ + if isinstance(dt, StructType): + new_fields = [] + for field in dt.fields: + new_fields.append( + StructField( + field.name, + force_decimal_precision_scale(field.dataType, precision=precision, scale=scale), + nullable=field.nullable, + metadata=field.metadata, + ) + ) + return StructType(new_fields) + elif isinstance(dt, ArrayType): + return ArrayType( + force_decimal_precision_scale(dt.elementType, precision=precision, scale=scale), + containsNull=dt.containsNull, + ) + elif isinstance(dt, MapType): + return MapType( + force_decimal_precision_scale(dt.keyType, precision=precision, scale=scale), + force_decimal_precision_scale(dt.valueType, precision=precision, scale=scale), + valueContainsNull=dt.valueContainsNull, + ) + elif isinstance(dt, DecimalType): + return DecimalType(precision=precision, scale=scale) + else: + return dt + + +def _test() -> None: + import doctest + import sys + import pyspark.pandas.spark.utils + + globs = pyspark.pandas.spark.utils.__dict__.copy() + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.spark.utils, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_formatter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_formatter.py new file mode 100644 index 0000000..8593703 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_formatter.py @@ -0,0 +1,308 @@ +# +# 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 os +import string +from typing import Any, Dict, Optional, Union, List, Sequence, Mapping, Tuple +import uuid +import warnings + +import pandas as pd + +from pyspark.pandas.internal import InternalFrame +from pyspark.pandas.namespace import _get_index_map +from pyspark import pandas as ps +from pyspark.sql import SparkSession +from pyspark.pandas.utils import default_session +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.series import Series + + +__all__ = ["sql"] + + +# This is not used in this file. It's for legacy sql_processor. +_CAPTURE_SCOPES = 3 + + +def sql( + query: str, + index_col: Optional[Union[str, List[str]]] = None, + args: Optional[Union[Dict[str, Any], List]] = None, + **kwargs: Any, +) -> DataFrame: + """ + Execute a SQL query and return the result as a pandas-on-Spark DataFrame. + + This function acts as a standard Python string formatter with understanding + the following variable types: + + * pandas-on-Spark DataFrame + * pandas-on-Spark Series + * pandas DataFrame + * pandas Series + * string + + Also the method can bind named parameters to SQL literals from `args`. + + Parameters + ---------- + query : str + the SQL query + index_col : str or list of str, optional + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + + .. note:: If you want to preserve the index, explicitly use :func:`DataFrame.reset_index`, + and pass it to the SQL statement with `index_col` parameter. + + For example, + + >>> psdf = ps.DataFrame({"A": [1, 2, 3], "B":[4, 5, 6]}, index=['a', 'b', 'c']) + >>> new_psdf = psdf.reset_index() + >>> ps.sql("SELECT * FROM {new_psdf}", index_col="index", new_psdf=new_psdf) + ... # doctest: +NORMALIZE_WHITESPACE + A B + index + a 1 4 + b 2 5 + c 3 6 + + For MultiIndex, + + >>> psdf = ps.DataFrame( + ... {"A": [1, 2, 3], "B": [4, 5, 6]}, + ... index=pd.MultiIndex.from_tuples( + ... [("a", "b"), ("c", "d"), ("e", "f")], names=["index1", "index2"] + ... ), + ... ) + >>> new_psdf = psdf.reset_index() + >>> ps.sql( + ... "SELECT * FROM {new_psdf}", index_col=["index1", "index2"], new_psdf=new_psdf) + ... # doctest: +NORMALIZE_WHITESPACE + A B + index1 index2 + a b 1 4 + c d 2 5 + e f 3 6 + + Also note that the index name(s) should be matched to the existing name. + args : dict or list + A dictionary of parameter names to Python objects or a list of Python objects + that can be converted to SQL literal expressions. See + + Supported Data Types for supported value types in Python. + For example, dictionary keys: "rank", "name", "birthdate"; + dictionary values: 1, "Steven", datetime.date(2023, 4, 2). + A value can be also a `Column` of literal expression, in that case it is taken as is. + + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Added positional parameters. + + kwargs + other variables that the user want to set that can be referenced in the query + + Returns + ------- + pandas-on-Spark DataFrame + + Examples + -------- + + Calling a built-in SQL function. + + >>> ps.sql("SELECT * FROM range(10) where id > 7") + id + 0 8 + 1 9 + + >>> ps.sql("SELECT * FROM range(10) WHERE id > {bound1} AND id < {bound2}", bound1=7, bound2=9) + id + 0 8 + + >>> mydf = ps.range(10) + >>> x = tuple(range(4)) + >>> ps.sql("SELECT {ser} FROM {mydf} WHERE id IN {x}", ser=mydf.id, mydf=mydf, x=x) + id + 0 0 + 1 1 + 2 2 + 3 3 + + Mixing pandas-on-Spark and pandas DataFrames in a join operation. Note that the index is + dropped. + + >>> ps.sql(''' + ... SELECT m1.a, m2.b + ... FROM {table1} m1 INNER JOIN {table2} m2 + ... ON m1.key = m2.key + ... ORDER BY m1.a, m2.b''', + ... table1=ps.DataFrame({"a": [1,2], "key": ["a", "b"]}), + ... table2=pd.DataFrame({"b": [3,4,5], "key": ["a", "b", "b"]})) + a b + 0 1 3 + 1 2 4 + 2 2 5 + + Also, it is possible to query using Series. + + >>> psdf = ps.DataFrame({"A": [1, 2, 3], "B":[4, 5, 6]}, index=['a', 'b', 'c']) + >>> ps.sql("SELECT {mydf.A} FROM {mydf}", mydf=psdf) + A + 0 1 + 1 2 + 2 3 + + And substitude named parameters with the `:` prefix by SQL literals. + + >>> ps.sql("SELECT * FROM range(10) WHERE id > :bound1", args={"bound1":7}) + id + 0 8 + 1 9 + + Or positional parameters marked by `?` in the SQL query by SQL literals. + + >>> ps.sql("SELECT * FROM range(10) WHERE id > ?", args=[7]) + id + 0 8 + 1 9 + """ + if os.environ.get("PYSPARK_PANDAS_SQL_LEGACY") == "1": + from pyspark.pandas import sql_processor + + warnings.warn( + "Deprecated in 3.3.0, and the legacy behavior " + "will be removed in the future releases.", + FutureWarning, + ) + return sql_processor.sql(query, index_col=index_col, **kwargs) + + session = default_session() + formatter = PandasSQLStringFormatter(session) + try: + sdf = session.sql(formatter.format(query, **kwargs), args) + finally: + formatter.clear() + + index_spark_columns, index_names = _get_index_map(sdf, index_col) + + return DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + + +class PandasSQLStringFormatter(string.Formatter): + """ + A standard ``string.Formatter`` in Python that can understand pandas-on-Spark instances + with basic Python objects. This object must be clear after the use for single SQL + query; cannot be reused across multiple SQL queries without cleaning. + """ + + def __init__(self, session: SparkSession) -> None: + self._session: SparkSession = session + self._temp_views: List[Tuple[DataFrame, str]] = [] + self._ref_sers: List[Tuple[Series, str]] = [] + + def vformat(self, format_string: str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> str: + ret = super(PandasSQLStringFormatter, self).vformat(format_string, args, kwargs) + + for ref, n in self._ref_sers: + if not any((ref is v for v in df._pssers.values()) for df, _ in self._temp_views): + # If referred DataFrame does not hold the given Series, raise an error. + raise ValueError("The series in {%s} does not refer any dataframe specified." % n) + return ret + + def get_field(self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> Any: + obj, first = super(PandasSQLStringFormatter, self).get_field(field_name, args, kwargs) + return self._convert_value(obj, field_name), first + + def _convert_value(self, val: Any, name: str) -> Optional[str]: + """ + Converts the given value into a SQL string. + """ + if isinstance(val, pd.Series): + # Return the column name from pandas Series directly. + return ps.from_pandas(val).to_frame()._to_spark().columns[0] + elif isinstance(val, Series): + # Return the column name of pandas-on-Spark Series iff its DataFrame was + # referred. The check will be done in `vformat` after we parse all. + self._ref_sers.append((val, name)) + return val.to_frame()._to_spark().columns[0] + elif isinstance(val, (DataFrame, pd.DataFrame)): + df_name = "_pandas_api_%s" % str(uuid.uuid4()).replace("-", "") + + if isinstance(val, pd.DataFrame): + # Don't store temp view for plain pandas instances + # because it is unable to know which pandas DataFrame + # holds which Series. + val = ps.from_pandas(val) + else: + for df, n in self._temp_views: + if df is val: + return n + self._temp_views.append((val, df_name)) + + val._to_spark().createOrReplaceTempView(df_name) + return df_name + elif isinstance(val, str): + # This is matched to behavior from JVM implementation. + # See `sql` definition from `sql/catalyst/src/main/scala/org/apache/spark/ + # sql/catalyst/expressions/literals.scala` + return "'" + val.replace("\\", "\\\\").replace("'", "\\'") + "'" + else: + return val + + def clear(self) -> None: + for _, n in self._temp_views: + self._session.catalog.dropTempView(n) + self._temp_views = [] + self._ref_sers = [] + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.sql_formatter + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.sql_formatter.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.sql_formatter tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.sql_formatter, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_processor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_processor.py new file mode 100644 index 0000000..0142e8d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/sql_processor.py @@ -0,0 +1,394 @@ +# +# 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 _string # type: ignore[import] +from typing import Any, Dict, Optional, Union, List +import inspect +import pandas as pd + +from pyspark.sql import SparkSession, DataFrame as SDataFrame + +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.pandas.utils import default_session +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.series import Series +from pyspark.pandas.internal import InternalFrame +from pyspark.pandas.namespace import _get_index_map + + +__all__ = ["sql"] + +from builtins import globals as builtin_globals +from builtins import locals as builtin_locals + + +def sql( + query: str, + index_col: Optional[Union[str, List[str]]] = None, + globals: Optional[Dict[str, Any]] = None, + locals: Optional[Dict[str, Any]] = None, + **kwargs: Any, +) -> DataFrame: + """ + Execute a SQL query and return the result as a pandas-on-Spark DataFrame. + + This function also supports embedding Python variables (locals, globals, and parameters) + in the SQL statement by wrapping them in curly braces. See examples section for details. + + In addition to the locals, globals and parameters, the function will also attempt + to determine if the program currently runs in an IPython (or Jupyter) environment + and to import the variables from this environment. The variables have the same + precedence as globals. + + The following variable types are supported: + + * string + * int + * float + * list, tuple, range of above types + * pandas-on-Spark DataFrame + * pandas-on-Spark Series + * pandas DataFrame + + Parameters + ---------- + query : str + the SQL query + index_col : str or list of str, optional + Column names to be used in Spark to represent pandas-on-Spark's index. The index name + in pandas-on-Spark is ignored. By default, the index is always lost. + + .. note:: If you want to preserve the index, explicitly use :func:`DataFrame.reset_index`, + and pass it to the SQL statement with `index_col` parameter. + + For example, + + >>> from pyspark.pandas import sql_processor + >>> # we will call 'sql_processor' directly in doctests so decrease one level. + >>> sql_processor._CAPTURE_SCOPES = 2 + >>> sql = sql_processor.sql + >>> psdf = ps.DataFrame({"A": [1, 2, 3], "B":[4, 5, 6]}, index=['a', 'b', 'c']) + >>> psdf_reset_index = psdf.reset_index() + >>> sql("SELECT * FROM {psdf_reset_index}", index_col="index") + ... # doctest: +NORMALIZE_WHITESPACE + A B + index + a 1 4 + b 2 5 + c 3 6 + + For MultiIndex, + + >>> psdf = ps.DataFrame( + ... {"A": [1, 2, 3], "B": [4, 5, 6]}, + ... index=pd.MultiIndex.from_tuples( + ... [("a", "b"), ("c", "d"), ("e", "f")], names=["index1", "index2"] + ... ), + ... ) + >>> psdf_reset_index = psdf.reset_index() + >>> sql("SELECT * FROM {psdf_reset_index}", index_col=["index1", "index2"]) + ... # doctest: +NORMALIZE_WHITESPACE + A B + index1 index2 + a b 1 4 + c d 2 5 + e f 3 6 + + Also note that the index name(s) should be matched to the existing name. + + globals : dict, optional + the dictionary of global variables, if explicitly set by the user + locals : dict, optional + the dictionary of local variables, if explicitly set by the user + kwargs + other variables that the user may want to set manually that can be referenced in the query + + Returns + ------- + pandas-on-Spark DataFrame + + Examples + -------- + + Calling a built-in SQL function. + + >>> sql("select * from range(10) where id > 7") + id + 0 8 + 1 9 + + A query can also reference a local variable or parameter by wrapping them in curly braces: + + >>> bound1 = 7 + >>> sql("select * from range(10) where id > {bound1} and id < {bound2}", bound2=9) + id + 0 8 + + You can also wrap a DataFrame with curly braces to query it directly. Note that when you do + that, the indexes, if any, automatically become top level columns. + + >>> mydf = ps.range(10) + >>> x = range(4) + >>> sql("SELECT * from {mydf} WHERE id IN {x}") + id + 0 0 + 1 1 + 2 2 + 3 3 + + Queries can also be arbitrarily nested in functions: + + >>> def statement(): + ... mydf2 = ps.DataFrame({"x": range(2)}) + ... return sql("SELECT * from {mydf2}") + >>> statement() + x + 0 0 + 1 1 + + Mixing pandas-on-Spark and pandas DataFrames in a join operation. Note that the index is + dropped. + + >>> sql(''' + ... SELECT m1.a, m2.b + ... FROM {table1} m1 INNER JOIN {table2} m2 + ... ON m1.key = m2.key + ... ORDER BY m1.a, m2.b''', + ... table1=ps.DataFrame({"a": [1,2], "key": ["a", "b"]}), + ... table2=pd.DataFrame({"b": [3,4,5], "key": ["a", "b", "b"]})) + a b + 0 1 3 + 1 2 4 + 2 2 5 + + Also, it is possible to query using Series. + + >>> myser = ps.Series({'a': [1.0, 2.0, 3.0], 'b': [15.0, 30.0, 45.0]}) + >>> sql("SELECT * from {myser}") + 0 + 0 [1.0, 2.0, 3.0] + 1 [15.0, 30.0, 45.0] + """ + if globals is None: + globals = _get_ipython_scope() + _globals = builtin_globals() if globals is None else dict(globals) + _locals = builtin_locals() if locals is None else dict(locals) + # The default choice is the globals + _dict = dict(_globals) + # The vars: + _scope = _get_local_scope() + _dict.update(_scope) + # Then the locals + _dict.update(_locals) + # Highest order of precedence is the locals + _dict.update(kwargs) + return SQLProcessor(_dict, query, default_session()).execute(index_col) + + +_CAPTURE_SCOPES = 3 + + +def _get_local_scope() -> Dict[str, Any]: + # Get 2 scopes above (_get_local_scope -> sql -> ...) to capture the vars there. + try: + return inspect.stack()[_CAPTURE_SCOPES][0].f_locals + except Exception: + # TODO (rxin, thunterdb): use a narrower scope exception. + # See https://github.com/databricks/koalas/pull/448 + return {} + + +def _get_ipython_scope() -> Dict[str, Any]: + """ + Tries to extract the dictionary of variables if the program is running + in an IPython notebook environment. + """ + try: + from IPython import get_ipython + + shell = get_ipython() + return shell.user_ns + except Exception: + # TODO (rxin, thunterdb): use a narrower scope exception. + # See https://github.com/databricks/koalas/pull/448 + return None + + +# Originally from pymysql package +_escape_table = [chr(x) for x in range(128)] +_escape_table[0] = "\\0" +_escape_table[ord("\\")] = "\\\\" +_escape_table[ord("\n")] = "\\n" +_escape_table[ord("\r")] = "\\r" +_escape_table[ord("\032")] = "\\Z" +_escape_table[ord('"')] = '\\"' +_escape_table[ord("'")] = "\\'" + + +def escape_sql_string(value: str) -> str: + """Escapes value without adding quotes. + + >>> escape_sql_string("foo\\nbar") + 'foo\\\\nbar' + + >>> escape_sql_string("'abc'de") + "\\\\'abc\\\\'de" + + >>> escape_sql_string('"abc"de') + '\\\\"abc\\\\"de' + """ + return value.translate(_escape_table) + + +class SQLProcessor: + def __init__(self, scope: Dict[str, Any], statement: str, session: SparkSession): + self._scope = scope + self._statement = statement + # All the temporary views created when executing this statement + # The key is the name of the variable in {} + # The value is the cached Spark Dataframe. + self._temp_views: Dict[str, SDataFrame] = {} + # All the other variables, converted to a normalized form. + # The normalized form is typically a string + self._cached_vars: Dict[str, Any] = {} + # The SQL statement after: + # - all the dataframes have been registered as temporary views + # - all the values have been converted normalized to equivalent SQL representations + self._normalized_statement: Optional[str] = None + self._session = session + + def execute(self, index_col: Optional[Union[str, List[str]]]) -> DataFrame: + """ + Returns a DataFrame for which the SQL statement has been executed by + the underlying SQL engine. + + >>> from pyspark.pandas import sql_processor + >>> # we will call 'sql_processor' directly in doctests so decrease one level. + >>> sql_processor._CAPTURE_SCOPES = 2 + >>> sql = sql_processor.sql + >>> str0 = 'abc' + >>> sql("select {str0}") + abc + 0 abc + + >>> str1 = 'abc"abc' + >>> str2 = "abc'abc" + >>> sql("select {str0}, {str1}, {str2}") + abc abc"abc abc'abc + 0 abc abc"abc abc'abc + + >>> strs = ['a', 'b'] + >>> sql("select 'a' in {strs} as cond1, 'c' in {strs} as cond2") + cond1 cond2 + 0 True False + """ + blocks = _string.formatter_parser(self._statement) + # TODO: use a string builder + res = "" + try: + for (pre, inner, _, _) in blocks: + var_next = "" if inner is None else self._convert(inner) + res = res + pre + var_next + self._normalized_statement = res + + sdf = self._session.sql(self._normalized_statement) + finally: + for v in self._temp_views: + self._session.catalog.dropTempView(v) + + index_spark_columns, index_names = _get_index_map(sdf, index_col) + + return DataFrame( + InternalFrame( + spark_frame=sdf, index_spark_columns=index_spark_columns, index_names=index_names + ) + ) + + def _convert(self, key: str) -> Any: + """ + Given a {} key, returns an equivalent SQL representation. + This conversion performs all the necessary escaping so that the string + returned can be directly injected into the SQL statement. + """ + # Already cached? + if key in self._cached_vars: + return self._cached_vars[key] + # Analyze: + if key not in self._scope: + raise ValueError( + "The key {} in the SQL statement was not found in global," + " local or parameters variables".format(key) + ) + var = self._scope[key] + fillin = self._convert_var(var) + self._cached_vars[key] = fillin + return fillin + + def _convert_var(self, var: Any) -> Any: + """ + Converts a python object into a string that is legal SQL. + """ + if isinstance(var, (int, float)): + return str(var) + if isinstance(var, Series): + return self._convert_var(var.to_dataframe()) + if isinstance(var, pd.DataFrame): + return self._convert_var(ps.DataFrame(var)) + if isinstance(var, DataFrame): + df_id = "pandas_on_spark_" + str(id(var)) + if df_id not in self._temp_views: + sdf = var._to_spark() + sdf.createOrReplaceTempView(df_id) + self._temp_views[df_id] = sdf + return df_id + if isinstance(var, str): + return '"' + escape_sql_string(var) + '"' + if isinstance(var, list): + return "(" + ", ".join([self._convert_var(v) for v in var]) + ")" + if isinstance(var, (tuple, range)): + return self._convert_var(list(var)) + raise ValueError("Unsupported variable type {}: {}".format(type(var).__name__, str(var))) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.sql_processor + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.sql_processor.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.sql_processor tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.sql_processor, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/strings.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/strings.py new file mode 100644 index 0000000..7c489de --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/strings.py @@ -0,0 +1,2371 @@ +# +# 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. +# + +""" +String functions on pandas-on-Spark Series +""" +import warnings +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Union, + cast, + no_type_check, +) + +import numpy as np + +import pandas as pd +from pyspark.sql.types import StringType, BinaryType, ArrayType, LongType, MapType +from pyspark.sql import functions as F +from pyspark.sql.functions import pandas_udf + +import pyspark.pandas as ps +from pyspark.pandas.spark import functions as SF + + +class StringMethods: + """String methods for pandas-on-Spark Series""" + + def __init__(self, series: "ps.Series"): + if not isinstance(series.spark.data_type, (StringType, BinaryType, ArrayType)): + raise ValueError("Cannot call StringMethods on type {}".format(series.spark.data_type)) + self._data = series + + # Methods + def capitalize(self) -> "ps.Series": + """ + Convert Strings in the series to be capitalized. + + Examples + -------- + >>> s = ps.Series(['lower', 'CAPITALS', 'this is a sentence', 'SwApCaSe']) + >>> s + 0 lower + 1 CAPITALS + 2 this is a sentence + 3 SwApCaSe + dtype: object + + >>> s.str.capitalize() + 0 Lower + 1 Capitals + 2 This is a sentence + 3 Swapcase + dtype: object + """ + + def pandas_capitalize(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.capitalize() + + return self._data.pandas_on_spark.transform_batch(pandas_capitalize) + + def title(self) -> "ps.Series": + """ + Convert Strings in the series to be title case. + + Examples + -------- + >>> s = ps.Series(['lower', 'CAPITALS', 'this is a sentence', 'SwApCaSe']) + >>> s + 0 lower + 1 CAPITALS + 2 this is a sentence + 3 SwApCaSe + dtype: object + + >>> s.str.title() + 0 Lower + 1 Capitals + 2 This Is A Sentence + 3 Swapcase + dtype: object + """ + + def pandas_title(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.title() + + return self._data.pandas_on_spark.transform_batch(pandas_title) + + def lower(self) -> "ps.Series": + """ + Convert strings in the Series/Index to all lowercase. + + Examples + -------- + >>> s = ps.Series(['lower', 'CAPITALS', 'this is a sentence', 'SwApCaSe']) + >>> s + 0 lower + 1 CAPITALS + 2 this is a sentence + 3 SwApCaSe + dtype: object + + >>> s.str.lower() + 0 lower + 1 capitals + 2 this is a sentence + 3 swapcase + dtype: object + """ + return self._data.spark.transform(F.lower) + + def upper(self) -> "ps.Series": + """ + Convert strings in the Series/Index to all uppercase. + + Examples + -------- + >>> s = ps.Series(['lower', 'CAPITALS', 'this is a sentence', 'SwApCaSe']) + >>> s + 0 lower + 1 CAPITALS + 2 this is a sentence + 3 SwApCaSe + dtype: object + + >>> s.str.upper() + 0 LOWER + 1 CAPITALS + 2 THIS IS A SENTENCE + 3 SWAPCASE + dtype: object + """ + return self._data.spark.transform(F.upper) + + def swapcase(self) -> "ps.Series": + """ + Convert strings in the Series/Index to be swap cased. + + Examples + -------- + >>> s = ps.Series(['lower', 'CAPITALS', 'this is a sentence', 'SwApCaSe']) + >>> s + 0 lower + 1 CAPITALS + 2 this is a sentence + 3 SwApCaSe + dtype: object + + >>> s.str.swapcase() + 0 LOWER + 1 capitals + 2 THIS IS A SENTENCE + 3 sWaPcAsE + dtype: object + """ + + def pandas_swapcase(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.swapcase() + + return self._data.pandas_on_spark.transform_batch(pandas_swapcase) + + def startswith(self, pattern: str, na: Optional[Any] = None) -> "ps.Series": + """ + Test if the start of each string element matches a pattern. + + Equivalent to :func:`str.startswith`. + + Parameters + ---------- + pattern : str + Character sequence. Regular expressions are not accepted. + na : object, default None + Object shown if element is not a string. NaN converted to None. + + Returns + ------- + Series of bool or object + pandas-on-Spark Series of booleans indicating whether the given pattern + matches the start of each string element. + + Examples + -------- + >>> s = ps.Series(['bat', 'Bear', 'cat', np.nan]) + >>> s + 0 bat + 1 Bear + 2 cat + 3 None + dtype: object + + >>> s.str.startswith('b') + 0 True + 1 False + 2 False + 3 None + dtype: object + + Specifying na to be False instead of None. + + >>> s.str.startswith('b', na=False) + 0 True + 1 False + 2 False + 3 False + dtype: bool + """ + + def pandas_startswith(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.startswith(pattern, na) + + return self._data.pandas_on_spark.transform_batch(pandas_startswith) + + def endswith(self, pattern: str, na: Optional[Any] = None) -> "ps.Series": + """ + Test if the end of each string element matches a pattern. + + Equivalent to :func:`str.endswith`. + + Parameters + ---------- + pattern : str + Character sequence. Regular expressions are not accepted. + na : object, default None + Object shown if element is not a string. NaN converted to None. + + Returns + ------- + Series of bool or object + pandas-on-Spark Series of booleans indicating whether the given pattern + matches the end of each string element. + + Examples + -------- + >>> s = ps.Series(['bat', 'Bear', 'cat', np.nan]) + >>> s + 0 bat + 1 Bear + 2 cat + 3 None + dtype: object + + >>> s.str.endswith('t') + 0 True + 1 False + 2 True + 3 None + dtype: object + + Specifying na to be False instead of None. + + >>> s.str.endswith('t', na=False) + 0 True + 1 False + 2 True + 3 False + dtype: bool + """ + + def pandas_endswith(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.endswith(pattern, na) + + return self._data.pandas_on_spark.transform_batch(pandas_endswith) + + def strip(self, to_strip: Optional[str] = None) -> "ps.Series": + """ + Remove leading and trailing characters. + + Strip whitespaces (including newlines) or a set of specified + characters from each string in the Series/Index from left and + right sides. Equivalent to :func:`str.strip`. + + Parameters + ---------- + to_strip : str + Specifying the set of characters to be removed. All combinations + of this set of characters will be stripped. If None then + whitespaces are removed. + + Returns + ------- + Series of objects + + Examples + -------- + >>> s = ps.Series(['1. Ant.', '2. Bee!\\t', None]) + >>> s + 0 1. Ant. + 1 2. Bee!\\t + 2 None + dtype: object + + >>> s.str.strip() + 0 1. Ant. + 1 2. Bee! + 2 None + dtype: object + + >>> s.str.strip('12.') + 0 Ant + 1 Bee!\\t + 2 None + dtype: object + + >>> s.str.strip('.!\\t') + 0 1. Ant + 1 2. Bee + 2 None + dtype: object + """ + + def pandas_strip(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.strip(to_strip) + + return self._data.pandas_on_spark.transform_batch(pandas_strip) + + def lstrip(self, to_strip: Optional[str] = None) -> "ps.Series": + """ + Remove leading characters. + + Strip whitespaces (including newlines) or a set of specified + characters from each string in the Series/Index from left side. + Equivalent to :func:`str.lstrip`. + + Parameters + ---------- + to_strip : str + Specifying the set of characters to be removed. All combinations + of this set of characters will be stripped. If None then + whitespaces are removed. + + Returns + ------- + Series of object + + Examples + -------- + >>> s = ps.Series(['1. Ant.', '2. Bee!\\t', None]) + >>> s + 0 1. Ant. + 1 2. Bee!\\t + 2 None + dtype: object + + >>> s.str.lstrip('12.') + 0 Ant. + 1 Bee!\\t + 2 None + dtype: object + """ + + def pandas_lstrip(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.lstrip(to_strip) + + return self._data.pandas_on_spark.transform_batch(pandas_lstrip) + + def rstrip(self, to_strip: Optional[str] = None) -> "ps.Series": + """ + Remove trailing characters. + + Strip whitespaces (including newlines) or a set of specified + characters from each string in the Series/Index from right side. + Equivalent to :func:`str.rstrip`. + + Parameters + ---------- + to_strip : str + Specifying the set of characters to be removed. All combinations + of this set of characters will be stripped. If None then + whitespaces are removed. + + Returns + ------- + Series of object + + Examples + -------- + >>> s = ps.Series(['1. Ant.', '2. Bee!\\t', None]) + >>> s + 0 1. Ant. + 1 2. Bee!\\t + 2 None + dtype: object + + >>> s.str.rstrip('.!\\t') + 0 1. Ant + 1 2. Bee + 2 None + dtype: object + """ + + def pandas_rstrip(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.rstrip(to_strip) + + return self._data.pandas_on_spark.transform_batch(pandas_rstrip) + + def get(self, i: int) -> "ps.Series": + """ + Extract element from each string or string list/tuple in the Series + at the specified position. + + Parameters + ---------- + i : int + Position of element to extract. + + Returns + ------- + Series of objects + + Examples + -------- + >>> s1 = ps.Series(["String", "123"]) + >>> s1 + 0 String + 1 123 + dtype: object + + >>> s1.str.get(1) + 0 t + 1 2 + dtype: object + + >>> s1.str.get(-1) + 0 g + 1 3 + dtype: object + + >>> s2 = ps.Series([["a", "b", "c"], ["x", "y"]]) + >>> s2 + 0 [a, b, c] + 1 [x, y] + dtype: object + + >>> s2.str.get(0) + 0 a + 1 x + dtype: object + + >>> s2.str.get(2) + 0 c + 1 None + dtype: object + """ + + def pandas_get(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.get(i) + + return self._data.pandas_on_spark.transform_batch(pandas_get) + + def isalnum(self) -> "ps.Series": + """ + Check whether all characters in each string are alphanumeric. + + This is equivalent to running the Python string method + :func:`str.isalnum` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s1 = ps.Series(['one', 'one1', '1', '']) + + >>> s1.str.isalnum() + 0 True + 1 True + 2 True + 3 False + dtype: bool + + Note that checks against characters mixed with any additional + punctuation or whitespace will evaluate too false for an alphanumeric + check. + + >>> s2 = ps.Series(['A B', '1.5', '3,000']) + >>> s2.str.isalnum() + 0 False + 1 False + 2 False + dtype: bool + """ + + def pandas_isalnum(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isalnum() + + return self._data.pandas_on_spark.transform_batch(pandas_isalnum) + + def isalpha(self) -> "ps.Series": + """ + Check whether all characters in each string are alphabetic. + + This is equivalent to running the Python string method + :func:`str.isalpha` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s1 = ps.Series(['one', 'one1', '1', '']) + + >>> s1.str.isalpha() + 0 True + 1 False + 2 False + 3 False + dtype: bool + """ + + def pandas_isalpha(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isalpha() + + return self._data.pandas_on_spark.transform_batch(pandas_isalpha) + + def isdigit(self) -> "ps.Series": + """ + Check whether all characters in each string are digits. + + This is equivalent to running the Python string method + :func:`str.isdigit` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series(['23', '³', '⅕', '']) + + The s.str.isdecimal method checks for characters used to form numbers + in base 10. + + >>> s.str.isdecimal() + 0 True + 1 False + 2 False + 3 False + dtype: bool + + The s.str.isdigit method is the same as s.str.isdecimal but also + includes special digits, like superscripted and subscripted digits in + Unicode. + + >>> s.str.isdigit() + 0 True + 1 True + 2 False + 3 False + dtype: bool + + The s.str.isnumeric method is the same as s.str.isdigit but also + includes other characters that can represent quantities such as unicode + fractions. + + >>> s.str.isnumeric() + 0 True + 1 True + 2 True + 3 False + dtype: bool + """ + + def pandas_isdigit(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isdigit() + + return self._data.pandas_on_spark.transform_batch(pandas_isdigit) + + def isspace(self) -> "ps.Series": + """ + Check whether all characters in each string are whitespaces. + + This is equivalent to running the Python string method + :func:`str.isspace` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series([' ', '\\t\\r\\n ', '']) + >>> s.str.isspace() + 0 True + 1 True + 2 False + dtype: bool + """ + + def pandas_isspace(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isspace() + + return self._data.pandas_on_spark.transform_batch(pandas_isspace) + + def islower(self) -> "ps.Series": + """ + Check whether all characters in each string are lowercase. + + This is equivalent to running the Python string method + :func:`str.islower` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series(['leopard', 'Golden Eagle', 'SNAKE', '']) + >>> s.str.islower() + 0 True + 1 False + 2 False + 3 False + dtype: bool + """ + + def pandas_isspace(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.islower() + + return self._data.pandas_on_spark.transform_batch(pandas_isspace) + + def isupper(self) -> "ps.Series": + """ + Check whether all characters in each string are uppercase. + + This is equivalent to running the Python string method + :func:`str.isupper` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series(['leopard', 'Golden Eagle', 'SNAKE', '']) + >>> s.str.isupper() + 0 False + 1 False + 2 True + 3 False + dtype: bool + """ + + def pandas_isspace(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isupper() + + return self._data.pandas_on_spark.transform_batch(pandas_isspace) + + def istitle(self) -> "ps.Series": + """ + Check whether all characters in each string are title case. + + This is equivalent to running the Python string method + :func:`str.istitle` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series(['leopard', 'Golden Eagle', 'SNAKE', '']) + + The s.str.istitle method checks for whether all words are in title + case (whether only the first letter of each word is capitalized). + Words are assumed to be as any sequence of non-numeric characters + separated by whitespace characters. + + >>> s.str.istitle() + 0 False + 1 True + 2 False + 3 False + dtype: bool + """ + + def pandas_istitle(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.istitle() + + return self._data.pandas_on_spark.transform_batch(pandas_istitle) + + def isnumeric(self) -> "ps.Series": + """ + Check whether all characters in each string are numeric. + + This is equivalent to running the Python string method + :func:`str.isnumeric` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s1 = ps.Series(['one', 'one1', '1', '']) + >>> s1.str.isnumeric() + 0 False + 1 False + 2 True + 3 False + dtype: bool + + >>> s2 = ps.Series(['23', '³', '⅕', '']) + + The s2.str.isdecimal method checks for characters used to form numbers + in base 10. + + >>> s2.str.isdecimal() + 0 True + 1 False + 2 False + 3 False + dtype: bool + + The s2.str.isdigit method is the same as s2.str.isdecimal but also + includes special digits, like superscripted and subscripted digits in + Unicode. + + >>> s2.str.isdigit() + 0 True + 1 True + 2 False + 3 False + dtype: bool + + The s2.str.isnumeric method is the same as s2.str.isdigit but also + includes other characters that can represent quantities such as unicode + fractions. + + >>> s2.str.isnumeric() + 0 True + 1 True + 2 True + 3 False + dtype: bool + """ + + def pandas_isnumeric(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isnumeric() + + return self._data.pandas_on_spark.transform_batch(pandas_isnumeric) + + def isdecimal(self) -> "ps.Series": + """ + Check whether all characters in each string are decimals. + + This is equivalent to running the Python string method + :func:`str.isdecimal` for each element of the Series/Index. + If a string has zero characters, False is returned for that check. + + Examples + -------- + >>> s = ps.Series(['23', '³', '⅕', '']) + + The s.str.isdecimal method checks for characters used to form numbers + in base 10. + + >>> s.str.isdecimal() + 0 True + 1 False + 2 False + 3 False + dtype: bool + + The s.str.isdigit method is the same as s.str.isdecimal but also + includes special digits, like superscripted and subscripted digits in + Unicode. + + >>> s.str.isdigit() + 0 True + 1 True + 2 False + 3 False + dtype: bool + + The s.str.isnumeric method is the same as s.str.isdigit but also + includes other characters that can represent quantities such as unicode + fractions. + + >>> s.str.isnumeric() + 0 True + 1 True + 2 True + 3 False + dtype: bool + """ + + def pandas_isdecimal(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.isdecimal() + + return self._data.pandas_on_spark.transform_batch(pandas_isdecimal) + + @no_type_check + def cat(self, others=None, sep=None, na_rep=None, join=None) -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + def center(self, width: int, fillchar: str = " ") -> "ps.Series": + """ + Filling left and right side of strings in the Series/Index with an + additional character. Equivalent to :func:`str.center`. + + Parameters + ---------- + width : int + Minimum width of resulting string; additional characters will be + filled with fillchar. + fillchar : str + Additional character for filling, default is whitespace. + + Returns + ------- + Series of objects + + Examples + -------- + >>> s = ps.Series(["caribou", "tiger"]) + >>> s + 0 caribou + 1 tiger + dtype: object + + >>> s.str.center(width=10, fillchar='-') + 0 -caribou-- + 1 --tiger--- + dtype: object + """ + + def pandas_center(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.center(width, fillchar) + + return self._data.pandas_on_spark.transform_batch(pandas_center) + + def contains( + self, pat: str, case: bool = True, flags: int = 0, na: Any = None, regex: bool = True + ) -> "ps.Series": + """ + Test if pattern or regex is contained within a string of a Series. + + Return boolean Series based on whether a given pattern or regex is + contained within a string of a Series. + + Analogous to :func:`match`, but less strict, relying on + :func:`re.search` instead of :func:`re.match`. + + Parameters + ---------- + pat : str + Character sequence or regular expression. + case : bool, default True + If True, case sensitive. + flags : int, default 0 (no flags) + Flags to pass through to the re module, e.g. re.IGNORECASE. + na : default None + Fill value for missing values. NaN converted to None. + regex : bool, default True + If True, assumes the pat is a regular expression. + If False, treats the pat as a literal string. + + + Returns + ------- + Series of boolean values or object + A Series of boolean values indicating whether the given pattern is + contained within the string of each element of the Series. + + Examples + -------- + Returning a Series of booleans using only a literal pattern. + + >>> s1 = ps.Series(['Mouse', 'dog', 'house and parrot', '23', np.NaN]) + >>> s1.str.contains('og', regex=False) + 0 False + 1 True + 2 False + 3 False + 4 None + dtype: object + + Specifying case sensitivity using case. + + >>> s1.str.contains('oG', case=True, regex=True) + 0 False + 1 False + 2 False + 3 False + 4 None + dtype: object + + Specifying na to be False instead of NaN replaces NaN values with + False. If Series does not contain NaN values the resultant dtype will + be bool, otherwise, an object dtype. + + >>> s1.str.contains('og', na=False, regex=True) + 0 False + 1 True + 2 False + 3 False + 4 False + dtype: bool + + Returning ‘house’ or ‘dog’ when either expression occurs in a string. + + >>> s1.str.contains('house|dog', regex=True) + 0 False + 1 True + 2 True + 3 False + 4 None + dtype: object + + Ignoring case sensitivity using flags with regex. + + >>> import re + >>> s1.str.contains('PARROT', flags=re.IGNORECASE, regex=True) + 0 False + 1 False + 2 True + 3 False + 4 None + dtype: object + + Returning any digit using regular expression. + + >>> s1.str.contains('[0-9]', regex=True) + 0 False + 1 False + 2 False + 3 True + 4 None + dtype: object + + Ensure pat is a not a literal pattern when regex is set to True. + Note in the following example one might expect only s2[1] and s2[3] + to return True. However, ‘.0’ as a regex matches any character followed + by a 0. + + >>> s2 = ps.Series(['40','40.0','41','41.0','35']) + >>> s2.str.contains('.0', regex=True) + 0 True + 1 True + 2 False + 3 True + 4 False + dtype: bool + """ + + def pandas_contains(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.contains(pat, case, flags, na, regex) + + return self._data.pandas_on_spark.transform_batch(pandas_contains) + + def count(self, pat: str, flags: int = 0) -> "ps.Series": + """ + Count occurrences of pattern in each string of the Series. + + This function is used to count the number of times a particular regex + pattern is repeated in each of the string elements of the Series. + + Parameters + ---------- + pat : str + Valid regular expression. + flags : int, default 0 (no flags) + Flags for the re module. + + Returns + ------- + Series of int + A Series containing the integer counts of pattern matches. + + Examples + -------- + >>> s = ps.Series(['A', 'B', 'Aaba', 'Baca', np.NaN, 'CABA', 'cat']) + >>> s.str.count('a') + 0 0.0 + 1 0.0 + 2 2.0 + 3 2.0 + 4 NaN + 5 0.0 + 6 1.0 + dtype: float64 + + Escape '$' to find the literal dollar sign. + + >>> s = ps.Series(['$', 'B', 'Aab$', '$$ca', 'C$B$', 'cat']) + >>> s.str.count('\\$') + 0 1 + 1 0 + 2 1 + 3 2 + 4 2 + 5 0 + dtype: int64 + """ + + def pandas_count(s) -> ps.Series[int]: # type: ignore[no-untyped-def] + return s.str.count(pat, flags) + + return self._data.pandas_on_spark.transform_batch(pandas_count) + + @no_type_check + def decode(self, encoding, errors="strict") -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + @no_type_check + def encode(self, encoding, errors="strict") -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + @no_type_check + def extract(self, pat, flags=0, expand=True) -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + @no_type_check + def extractall(self, pat, flags=0) -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + def find(self, sub: str, start: int = 0, end: Optional[int] = None) -> "ps.Series": + """ + Return lowest indexes in each string in the Series where the + substring is fully contained between [start:end]. + + Return -1 on failure. Equivalent to standard :func:`str.find`. + + Parameters + ---------- + sub : str + Substring being searched. + start : int + Left edge index. + end : int + Right edge index. + + Returns + ------- + Series of int + Series of lowest matching indexes. + + Examples + -------- + >>> s = ps.Series(['apple', 'oranges', 'bananas']) + + >>> s.str.find('a') + 0 0 + 1 2 + 2 1 + dtype: int64 + + >>> s.str.find('a', start=2) + 0 -1 + 1 2 + 2 3 + dtype: int64 + + >>> s.str.find('a', end=1) + 0 0 + 1 -1 + 2 -1 + dtype: int64 + + >>> s.str.find('a', start=2, end=2) + 0 -1 + 1 -1 + 2 -1 + dtype: int64 + """ + + def pandas_find(s) -> ps.Series[int]: # type: ignore[no-untyped-def] + return s.str.find(sub, start, end) + + return self._data.pandas_on_spark.transform_batch(pandas_find) + + def findall(self, pat: str, flags: int = 0) -> "ps.Series": + """ + Find all occurrences of pattern or regular expression in the Series. + + Equivalent to applying :func:`re.findall` to all the elements in + the Series. + + Parameters + ---------- + pat : str + Pattern or regular expression. + flags : int, default 0 (no flags) + `re` module flags, e.g. `re.IGNORECASE`. + + Returns + ------- + Series of object + All non-overlapping matches of pattern or regular expression in + each string of this Series. + + Examples + -------- + >>> s = ps.Series(['Lion', 'Monkey', 'Rabbit']) + + The search for the pattern ‘Monkey’ returns one match: + + >>> s.str.findall('Monkey') + 0 [] + 1 [Monkey] + 2 [] + dtype: object + + On the other hand, the search for the pattern ‘MONKEY’ doesn’t return + any match: + + >>> s.str.findall('MONKEY') + 0 [] + 1 [] + 2 [] + dtype: object + + Flags can be added to the pattern or regular expression. For instance, + to find the pattern ‘MONKEY’ ignoring the case: + + >>> import re + >>> s.str.findall('MONKEY', flags=re.IGNORECASE) + 0 [] + 1 [Monkey] + 2 [] + dtype: object + + When the pattern matches more than one string in the Series, all + matches are returned: + + >>> s.str.findall('on') + 0 [on] + 1 [on] + 2 [] + dtype: object + + Regular expressions are supported too. For instance, the search for all + the strings ending with the word ‘on’ is shown next: + + >>> s.str.findall('on$') + 0 [on] + 1 [] + 2 [] + dtype: object + + If the pattern is found more than once in the same string, then a list + of multiple strings is returned: + + >>> s.str.findall('b') + 0 [] + 1 [] + 2 [b, b] + dtype: object + """ + # type hint does not support to specify array type yet. + @pandas_udf( # type: ignore[call-overload] + returnType=ArrayType(StringType(), containsNull=True) + ) + def pudf(s: pd.Series) -> pd.Series: + return s.str.findall(pat, flags) + + return self._data._with_new_scol(scol=pudf(self._data.spark.column)) + + def index(self, sub: str, start: int = 0, end: Optional[int] = None) -> "ps.Series": + """ + Return lowest indexes in each string where the substring is fully + contained between [start:end]. + + This is the same as :func:`str.find` except instead of returning -1, + it raises a ValueError when the substring is not found. Equivalent to + standard :func:`str.index`. + + Parameters + ---------- + sub : str + Substring being searched. + start : int + Left edge index. + end : int + Right edge index. + + Returns + ------- + Series of int + Series of lowest matching indexes. + + Examples + -------- + >>> s = ps.Series(['apple', 'oranges', 'bananas']) + + >>> s.str.index('a') + 0 0 + 1 2 + 2 1 + dtype: int64 + + The following expression throws an exception: + + >>> s.str.index('a', start=2) # doctest: +SKIP + """ + + def pandas_index(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.str.index(sub, start, end) + + return self._data.pandas_on_spark.transform_batch(pandas_index) + + def join(self, sep: str) -> "ps.Series": + """ + Join lists contained as elements in the Series with passed delimiter. + + If the elements of a Series are lists themselves, join the content of + these lists using the delimiter passed to the function. This function + is an equivalent to calling :func:`str.join` on the lists. + + Parameters + ---------- + sep : str + Delimiter to use between list entries. + + Returns + ------- + Series of object + Series with list entries concatenated by intervening occurrences of + the delimiter. + + See Also + -------- + str.split : Split strings around given separator/delimiter. + str.rsplit : Splits string around given separator/delimiter, + starting from the right. + + Examples + -------- + Example with a list that contains a None element. + + >>> s = ps.Series([['lion', 'elephant', 'zebra'], + ... ['cat', None, 'dog']]) + >>> s + 0 [lion, elephant, zebra] + 1 [cat, None, dog] + dtype: object + + Join all lists using a ‘-‘. The list containing None will produce None. + + >>> s.str.join('-') + 0 lion-elephant-zebra + 1 None + dtype: object + """ + + def pandas_join(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.join(sep) + + return self._data.pandas_on_spark.transform_batch(pandas_join) + + def len(self) -> "ps.Series": + """ + Computes the length of each element in the Series. + + The element may be a sequence (such as a string, tuple or list). + + Returns + ------- + Series of int + A Series of integer values indicating the length of each element in + the Series. + + Examples + -------- + Returns the length (number of characters) in a string. Returns the + number of entries for lists or tuples. + + >>> s1 = ps.Series(['dog', 'monkey']) + >>> s1.str.len() + 0 3 + 1 6 + dtype: int64 + + >>> s2 = ps.Series([["a", "b", "c"], []]) + >>> s2.str.len() + 0 3 + 1 0 + dtype: int64 + """ + if isinstance(self._data.spark.data_type, (ArrayType, MapType)): + return self._data.spark.transform(lambda c: F.size(c).cast(LongType())) + else: + return self._data.spark.transform(lambda c: F.length(c).cast(LongType())) + + def ljust(self, width: int, fillchar: str = " ") -> "ps.Series": + """ + Filling right side of strings in the Series with an additional + character. Equivalent to :func:`str.ljust`. + + Parameters + ---------- + width : int + Minimum width of resulting string; additional characters will be + filled with `fillchar`. + fillchar : str + Additional character for filling, default is whitespace. + + Returns + ------- + Series of object + + Examples + -------- + >>> s = ps.Series(["caribou", "tiger"]) + >>> s + 0 caribou + 1 tiger + dtype: object + + >>> s.str.ljust(width=10, fillchar='-') + 0 caribou--- + 1 tiger----- + dtype: object + """ + + def pandas_ljust(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.ljust(width, fillchar) + + return self._data.pandas_on_spark.transform_batch(pandas_ljust) + + def match(self, pat: str, case: bool = True, flags: int = 0, na: Any = np.NaN) -> "ps.Series": + """ + Determine if each string matches a regular expression. + + Analogous to :func:`contains`, but more strict, relying on + :func:`re.match` instead of :func:`re.search`. + + Parameters + ---------- + pat : str + Character sequence or regular expression. + case : bool, default True + If True, case sensitive. + flags : int, default 0 (no flags) + Flags to pass through to the re module, e.g. re.IGNORECASE. + na : default NaN + Fill value for missing values. + + Returns + ------- + Series of boolean values or object + A Series of boolean values indicating whether the given pattern can + be matched in the string of each element of the Series. + + Examples + -------- + >>> s = ps.Series(['Mouse', 'dog', 'house and parrot', '23', np.NaN]) + >>> s.str.match('dog') + 0 False + 1 True + 2 False + 3 False + 4 None + dtype: object + + >>> s.str.match('mouse|dog', case=False) + 0 True + 1 True + 2 False + 3 False + 4 None + dtype: object + + >>> s.str.match('.+and.+', na=True) + 0 False + 1 False + 2 True + 3 False + 4 True + dtype: bool + + >>> import re + >>> s.str.match('MOUSE', flags=re.IGNORECASE) + 0 True + 1 False + 2 False + 3 False + 4 None + dtype: object + """ + + def pandas_match(s) -> ps.Series[bool]: # type: ignore[no-untyped-def] + return s.str.match(pat, case, flags, na) + + return self._data.pandas_on_spark.transform_batch(pandas_match) + + def normalize(self, form: str) -> "ps.Series": + """ + Return the Unicode normal form for the strings in the Series. + + For more information on the forms, see the + :func:`unicodedata.normalize`. + + Parameters + ---------- + form : {‘NFC’, ‘NFKC’, ‘NFD’, ‘NFKD’} + Unicode form. + + Returns + ------- + Series of objects + A Series of normalized strings. + """ + + def pandas_normalize(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.normalize(form) + + return self._data.pandas_on_spark.transform_batch(pandas_normalize) + + def pad(self, width: int, side: str = "left", fillchar: str = " ") -> "ps.Series": + """ + Pad strings in the Series up to width. + + Parameters + ---------- + width : int + Minimum width of resulting string; additional characters will be + filled with character defined in `fillchar`. + side : {‘left’, ‘right’, ‘both’}, default ‘left’ + Side from which to fill resulting string. + fillchar : str, default ' ' + Additional character for filling, default is whitespace. + + Returns + ------- + Series of object + Returns Series with minimum number of char in object. + + Examples + -------- + >>> s = ps.Series(["caribou", "tiger"]) + >>> s + 0 caribou + 1 tiger + dtype: object + + >>> s.str.pad(width=10) + 0 caribou + 1 tiger + dtype: object + + >>> s.str.pad(width=10, side='right', fillchar='-') + 0 caribou--- + 1 tiger----- + dtype: object + + >>> s.str.pad(width=10, side='both', fillchar='-') + 0 -caribou-- + 1 --tiger--- + dtype: object + """ + + def pandas_pad(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.pad(width, side, fillchar) + + return self._data.pandas_on_spark.transform_batch(pandas_pad) + + def partition(self, sep: str = " ", expand: bool = True) -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + def repeat(self, repeats: int) -> "ps.Series": + """ + Duplicate each string in the Series. + + Parameters + ---------- + repeats : int + Repeat the string given number of times (int). Sequence of int + is not supported. + + Returns + ------- + Series of object + Series or Index of repeated string objects specified by input + parameter repeats. + + Examples + -------- + >>> s = ps.Series(['a', 'b', 'c']) + >>> s + 0 a + 1 b + 2 c + dtype: object + + Single int repeats string in Series + + >>> s.str.repeat(repeats=2) + 0 aa + 1 bb + 2 cc + dtype: object + """ + if not isinstance(repeats, int): + raise TypeError("repeats expects an int parameter") + return self._data.spark.transform(lambda c: SF.repeat(col=c, n=repeats)) + + def replace( + self, + pat: str, + repl: Union[str, Callable[[str], str]], + n: int = -1, + case: Optional[bool] = None, + flags: int = 0, + regex: bool = True, + ) -> "ps.Series": + """ + Replace occurrences of pattern/regex in the Series with some other + string. Equivalent to :func:`str.replace` or :func:`re.sub`. + + Parameters + ---------- + pat : str or compiled regex + String can be a character sequence or regular expression. + repl : str or callable + Replacement string or a callable. The callable is passed the regex + match object and must return a replacement string to be used. See + :func:`re.sub`. + n : int, default -1 (all) + Number of replacements to make from start. + case : boolean, default None + If True, case sensitive (the default if pat is a string). + Set to False for case insensitive. + Cannot be set if pat is a compiled regex. + flags: int, default 0 (no flags) + re module flags, e.g. re.IGNORECASE. + Cannot be set if pat is a compiled regex. + regex : boolean, default True + If True, assumes the passed-in pattern is a regular expression. + If False, treats the pattern as a literal string. + Cannot be set to False if pat is a compile regex or repl is a + callable. + + Returns + ------- + Series of object + A copy of the string with all matching occurrences of pat replaced + by repl. + + Examples + -------- + When pat is a string and regex is True (the default), the given pat is + compiled as a regex. When repl is a string, it replaces matching regex + patterns as with :func:`re.sub`. NaN value(s) in the Series are changed + to None: + + >>> ps.Series(['foo', 'fuz', np.nan]).str.replace('f.', 'ba', regex=True) + 0 bao + 1 baz + 2 None + dtype: object + + When pat is a string and regex is False, every pat is replaced with + repl as with :func:`str.replace`: + + >>> ps.Series(['f.o', 'fuz', np.nan]).str.replace('f.', 'ba', regex=False) + 0 bao + 1 fuz + 2 None + dtype: object + + When repl is a callable, it is called on every pat using + :func:`re.sub`. The callable should expect one positional argument (a + regex object) and return a string. + + Reverse every lowercase alphabetic word: + + >>> repl = lambda m: m.group(0)[::-1] + >>> ps.Series(['foo 123', 'bar baz', np.nan]).str.replace(r'[a-z]+', repl) + 0 oof 123 + 1 rab zab + 2 None + dtype: object + + Using regex groups (extract second group and swap case): + + >>> pat = r"(?P\\w+) (?P\\w+) (?P\\w+)" + >>> repl = lambda m: m.group('two').swapcase() + >>> ps.Series(['One Two Three', 'Foo Bar Baz']).str.replace(pat, repl) + 0 tWO + 1 bAR + dtype: object + + Using a compiled regex with flags: + + >>> import re + >>> regex_pat = re.compile(r'FUZ', flags=re.IGNORECASE) + >>> ps.Series(['foo', 'fuz', np.nan]).str.replace(regex_pat, 'bar') + 0 foo + 1 bar + 2 None + dtype: object + """ + warnings.warn( + "Default value of `regex` will be changed to `False` instead of `True` in 4.0.0.", + FutureWarning, + ) + + def pandas_replace(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.replace(pat, repl, n=n, case=case, flags=flags, regex=regex) + + return self._data.pandas_on_spark.transform_batch(pandas_replace) + + def rfind(self, sub: str, start: int = 0, end: Optional[int] = None) -> "ps.Series": + """ + Return highest indexes in each string in the Series where the + substring is fully contained between [start:end]. + + Return -1 on failure. Equivalent to standard :func:`str.rfind`. + + Parameters + ---------- + sub : str + Substring being searched. + start : int + Left edge index. + end : int + Right edge index. + + Returns + ------- + Series of int + Series of highest matching indexes. + + Examples + -------- + >>> s = ps.Series(['apple', 'oranges', 'bananas']) + + >>> s.str.rfind('a') + 0 0 + 1 2 + 2 5 + dtype: int64 + + >>> s.str.rfind('a', start=2) + 0 -1 + 1 2 + 2 5 + dtype: int64 + + >>> s.str.rfind('a', end=1) + 0 0 + 1 -1 + 2 -1 + dtype: int64 + + >>> s.str.rfind('a', start=2, end=2) + 0 -1 + 1 -1 + 2 -1 + dtype: int64 + """ + + def pandas_rfind(s) -> ps.Series[int]: # type: ignore[no-untyped-def] + return s.str.rfind(sub, start, end) + + return self._data.pandas_on_spark.transform_batch(pandas_rfind) + + def rindex(self, sub: str, start: int = 0, end: Optional[int] = None) -> "ps.Series": + """ + Return highest indexes in each string where the substring is fully + contained between [start:end]. + + This is the same as :func:`str.rfind` except instead of returning -1, + it raises a ValueError when the substring is not found. Equivalent to + standard :func:`str.rindex`. + + Parameters + ---------- + sub : str + Substring being searched. + start : int + Left edge index. + end : int + Right edge index. + + Returns + ------- + Series of int + Series of highest matching indexes. + + Examples + -------- + >>> s = ps.Series(['apple', 'oranges', 'bananas']) + + >>> s.str.rindex('a') + 0 0 + 1 2 + 2 5 + dtype: int64 + + The following expression throws an exception: + + >>> s.str.rindex('a', start=2) # doctest: +SKIP + """ + + def pandas_rindex(s) -> ps.Series[np.int64]: # type: ignore[no-untyped-def] + return s.str.rindex(sub, start, end) + + return self._data.pandas_on_spark.transform_batch(pandas_rindex) + + def rjust(self, width: int, fillchar: str = " ") -> "ps.Series": + """ + Filling left side of strings in the Series with an additional + character. Equivalent to :func:`str.rjust`. + + Parameters + ---------- + width : int + Minimum width of resulting string; additional characters will be + filled with `fillchar`. + fillchar : str + Additional character for filling, default is whitespace. + + Returns + ------- + Series of object + + Examples + -------- + >>> s = ps.Series(["caribou", "tiger"]) + >>> s + 0 caribou + 1 tiger + dtype: object + + >>> s.str.rjust(width=10) + 0 caribou + 1 tiger + dtype: object + + >>> s.str.rjust(width=10, fillchar='-') + 0 ---caribou + 1 -----tiger + dtype: object + """ + + def pandas_rjust(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.rjust(width, fillchar) + + return self._data.pandas_on_spark.transform_batch(pandas_rjust) + + def rpartition(self, sep: str = " ", expand: bool = True) -> "ps.Series": + """ + Not supported. + """ + raise NotImplementedError() + + def slice( + self, start: Optional[int] = None, stop: Optional[int] = None, step: Optional[int] = None + ) -> "ps.Series": + """ + Slice substrings from each element in the Series. + + Parameters + ---------- + start : int, optional + Start position for slice operation. + stop : int, optional + Stop position for slice operation. + step : int, optional + Step size for slice operation. + + Returns + ------- + Series of object + Series from sliced substrings from original string objects. + + Examples + -------- + >>> s = ps.Series(["koala", "fox", "chameleon"]) + >>> s + 0 koala + 1 fox + 2 chameleon + dtype: object + + >>> s.str.slice(start=1) + 0 oala + 1 ox + 2 hameleon + dtype: object + + >>> s.str.slice(stop=2) + 0 ko + 1 fo + 2 ch + dtype: object + + >>> s.str.slice(step=2) + 0 kaa + 1 fx + 2 caeen + dtype: object + + >>> s.str.slice(start=0, stop=5, step=3) + 0 kl + 1 f + 2 cm + dtype: object + """ + + def pandas_slice(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.slice(start, stop, step) + + return self._data.pandas_on_spark.transform_batch(pandas_slice) + + def slice_replace( + self, start: Optional[int] = None, stop: Optional[int] = None, repl: Optional[str] = None + ) -> "ps.Series": + """ + Slice substrings from each element in the Series. + + Parameters + ---------- + start : int, optional + Start position for slice operation. If not specified (None), the + slice is unbounded on the left, i.e. slice from the start of the + string. + stop : int, optional + Stop position for slice operation. If not specified (None), the + slice is unbounded on the right, i.e. slice until the end of the + string. + repl : str, optional + String for replacement. If not specified (None), the sliced region + is replaced with an empty string. + + Returns + ------- + Series of object + Series from sliced substrings from original string objects. + + Examples + -------- + >>> s = ps.Series(['a', 'ab', 'abc', 'abdc', 'abcde']) + >>> s + 0 a + 1 ab + 2 abc + 3 abdc + 4 abcde + dtype: object + + Specify just start, meaning replace start until the end of the string + with repl. + + >>> s.str.slice_replace(1, repl='X') + 0 aX + 1 aX + 2 aX + 3 aX + 4 aX + dtype: object + + Specify just stop, meaning the start of the string to stop is replaced + with repl, and the rest of the string is included. + + >>> s.str.slice_replace(stop=2, repl='X') + 0 X + 1 X + 2 Xc + 3 Xdc + 4 Xcde + dtype: object + + Specify start and stop, meaning the slice from start to stop is + replaced with repl. Everything before or after start and stop is + included as is. + + >>> s.str.slice_replace(start=1, stop=3, repl='X') + 0 aX + 1 aX + 2 aX + 3 aXc + 4 aXde + dtype: object + """ + + def pandas_slice_replace(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.slice_replace(start, stop, repl) + + return self._data.pandas_on_spark.transform_batch(pandas_slice_replace) + + def split( + self, pat: Optional[str] = None, n: int = -1, expand: bool = False + ) -> Union["ps.Series", "ps.DataFrame"]: + """ + Split strings around given separator/delimiter. + + Splits the string in the Series from the beginning, at the specified + delimiter string. Equivalent to :func:`str.split`. + + Parameters + ---------- + pat : str, optional + String or regular expression to split on. If not specified, split + on whitespace. + n : int, default -1 (all) + Limit number of splits in output. None, 0 and -1 will be + interpreted as return all splits. + expand : bool, default False + Expand the split strings into separate columns. + + * If ``True``, `n` must be a positive integer, and return DataFrame expanding + dimensionality. + * If ``False``, return Series, containing lists of strings. + + Returns + ------- + Series, DataFrame + Type matches caller unless `expand=True` (see Notes). + + See Also + -------- + str.rsplit : Splits string around given separator/delimiter, + starting from the right. + str.join : Join lists contained as elements in the Series/Index + with passed delimiter. + + Notes + ----- + The handling of the `n` keyword depends on the number of found splits: + + - If found splits > `n`, make first `n` splits only + - If found splits <= `n`, make all splits + - If for a certain row the number of found splits < `n`, + append `None` for padding up to `n` if ``expand=True`` + + If using ``expand=True``, Series callers return DataFrame objects with `n + 1` columns. + + .. note:: Even if `n` is much larger than found splits, the number of columns does NOT + shrink unlike pandas. + + Examples + -------- + >>> s = ps.Series(["this is a regular sentence", + ... "https://docs.python.org/3/tutorial/index.html", + ... np.nan]) + + In the default setting, the string is split by whitespace. + + >>> s.str.split() # doctest: +SKIP + 0 [this, is, a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + Without the n parameter, the outputs of rsplit and split are identical. + + >>> s.str.rsplit() # doctest: +SKIP + 0 [this, is, a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + The n parameter can be used to limit the number of splits on the + delimiter. The outputs of split and rsplit are different. + + >>> s.str.split(n=2) # doctest: +SKIP + 0 [this, is, a regular sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + >>> s.str.rsplit(n=2) # doctest: +SKIP + 0 [this is a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + The pat parameter can be used to split by other characters. + + >>> s.str.split(pat = "/") # doctest: +SKIP + 0 [this is a regular sentence] + 1 [https:, , docs.python.org, 3, tutorial, index... + 2 None + dtype: object + + When using ``expand=True``, the split elements will expand out into + separate columns. If NaN is present, it is propagated throughout + the columns during the split. + + >>> s.str.split(n=4, expand=True) # doctest: +SKIP + 0 1 2 3 4 + 0 this is a regular sentence + 1 https://docs.python.org/3/tutorial/index.html None None None None + 2 None None None None None + + For slightly more complex use cases like splitting the html document name + from a url, a combination of parameter settings can be used. + + >>> s.str.rsplit("/", n=1, expand=True) # doctest: +SKIP + 0 1 + 0 this is a regular sentence None + 1 https://docs.python.org/3/tutorial index.html + 2 None None + + Remember to escape special characters when explicitly using regular + expressions. + + >>> s = ps.Series(["1+1=2"]) + >>> s.str.split(r"\\+|=", n=2, expand=True) # doctest: +SKIP + 0 1 2 + 0 1 1 2 + """ + from pyspark.pandas.frame import DataFrame + + if expand and n <= 0: + raise NotImplementedError("expand=True is currently only supported with n > 0.") + + # type hint does not support to specify array type yet. + return_type = ArrayType(StringType(), containsNull=True) + + @pandas_udf(returnType=return_type) # type: ignore[call-overload] + def pudf(s: pd.Series) -> pd.Series: + return s.str.split(pat, n) + + psser = self._data._with_new_scol( + pudf(self._data.spark.column).alias(self._data._internal.data_spark_column_names[0]), + field=self._data._internal.data_fields[0].copy(spark_type=return_type, nullable=True), + ) + + if expand: + psdf = psser.to_frame() + scol = psdf._internal.data_spark_columns[0] + spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)] + column_labels = [(i,) for i in range(n + 1)] + internal = psdf._internal.with_new_columns( + spark_columns, + column_labels=cast(Optional[List], column_labels), + data_fields=[ + self._data._internal.data_fields[0].copy(name=str(i), nullable=True) + for i in range(n + 1) + ], + ) + return DataFrame(internal) + else: + return psser + + def rsplit( + self, pat: Optional[str] = None, n: int = -1, expand: bool = False + ) -> Union["ps.Series", "ps.DataFrame"]: + """ + Split strings around given separator/delimiter. + + Splits the string in the Series from the end, at the specified + delimiter string. Equivalent to :func:`str.rsplit`. + + Parameters + ---------- + pat : str, optional + String or regular expression to split on. If not specified, split + on whitespace. + n : int, default -1 (all) + Limit number of splits in output. None, 0 and -1 will be + interpreted as return all splits. + expand : bool, default False + Expand the split strings into separate columns. + + * If ``True``, `n` must be a positive integer, and return DataFrame expanding + dimensionality. + * If ``False``, return Series, containing lists of strings. + + Returns + ------- + Series, DataFrame + Type matches caller unless `expand=True` (see Notes). + + See Also + -------- + str.split : Split strings around given separator/delimiter. + str.join : Join lists contained as elements in the Series/Index + with passed delimiter. + + Notes + ----- + The handling of the `n` keyword depends on the number of found splits: + + - If found splits > `n`, make first `n` splits only + - If found splits <= `n`, make all splits + - If for a certain row the number of found splits < `n`, + append `None` for padding up to `n` if ``expand=True`` + + If using ``expand=True``, Series callers return DataFrame objects with `n + 1` columns. + + .. note:: Even if `n` is much larger than found splits, the number of columns does NOT + shrink unlike pandas. + + Examples + -------- + >>> s = ps.Series(["this is a regular sentence", + ... "https://docs.python.org/3/tutorial/index.html", + ... np.nan]) + + In the default setting, the string is split by whitespace. + + >>> s.str.split() # doctest: +SKIP + 0 [this, is, a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + Without the n parameter, the outputs of rsplit and split are identical. + + >>> s.str.rsplit() # doctest: +SKIP + 0 [this, is, a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + The n parameter can be used to limit the number of splits on the + delimiter. The outputs of split and rsplit are different. + + >>> s.str.split(n=2) # doctest: +SKIP + 0 [this, is, a regular sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + >>> s.str.rsplit(n=2) # doctest: +SKIP + 0 [this is a, regular, sentence] + 1 [https://docs.python.org/3/tutorial/index.html] + 2 None + dtype: object + + When using ``expand=True``, the split elements will expand out into + separate columns. If NaN is present, it is propagated throughout + the columns during the split. + + >>> s.str.split(n=4, expand=True) # doctest: +SKIP + 0 1 2 3 4 + 0 this is a regular sentence + 1 https://docs.python.org/3/tutorial/index.html None None None None + 2 None None None None None + + For slightly more complex use cases like splitting the html document name + from a url, a combination of parameter settings can be used. + + >>> s.str.rsplit("/", n=1, expand=True) # doctest: +SKIP + 0 1 + 0 this is a regular sentence None + 1 https://docs.python.org/3/tutorial index.html + 2 None None + + Remember to escape special characters when explicitly using regular + expressions. + + >>> s = ps.Series(["1+1=2"]) + >>> s.str.split(r"\\+|=", n=2, expand=True) # doctest: +SKIP + 0 1 2 + 0 1 1 2 + """ + from pyspark.pandas.frame import DataFrame + + if expand and n <= 0: + raise NotImplementedError("expand=True is currently only supported with n > 0.") + + # type hint does not support to specify array type yet. + return_type = ArrayType(StringType(), containsNull=True) + + @pandas_udf(returnType=return_type) # type: ignore[call-overload] + def pudf(s: pd.Series) -> pd.Series: + return s.str.rsplit(pat, n) + + psser = self._data._with_new_scol( + pudf(self._data.spark.column).alias(self._data._internal.data_spark_column_names[0]), + field=self._data._internal.data_fields[0].copy(spark_type=return_type, nullable=True), + ) + + if expand: + psdf = psser.to_frame() + scol = psdf._internal.data_spark_columns[0] + spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)] + column_labels = [(i,) for i in range(n + 1)] + internal = psdf._internal.with_new_columns( + spark_columns, + column_labels=cast(Optional[List], column_labels), + data_fields=[ + self._data._internal.data_fields[0].copy(name=str(i), nullable=True) + for i in range(n + 1) + ], + ) + return DataFrame(internal) + else: + return psser + + def translate(self, table: Dict) -> "ps.Series": + """ + Map all characters in the string through the given mapping table. + Equivalent to standard :func:`str.translate`. + + Parameters + ---------- + table : dict + Table is a mapping of Unicode ordinals to Unicode ordinals, + strings, or None. Unmapped characters are left untouched. + Characters mapped to None are deleted. :func:`str.maketrans` is a + helper function for making translation tables. + + Returns + ------- + Series of object + Series with translated strings. + + Examples + -------- + >>> s = ps.Series(["dog", "cat", "bird"]) + >>> m = str.maketrans({'a': 'X', 'i': 'Y', 'o': None}) + >>> s.str.translate(m) + 0 dg + 1 cXt + 2 bYrd + dtype: object + """ + + def pandas_translate(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.translate(table) + + return self._data.pandas_on_spark.transform_batch(pandas_translate) + + def wrap(self, width: int, **kwargs: bool) -> "ps.Series": + """ + Wrap long strings in the Series to be formatted in paragraphs with + length less than a given width. + + This method has the same keyword parameters and defaults as + :class:`textwrap.TextWrapper`. + + Parameters + ---------- + width : int + Maximum line-width. Lines separated with newline char. + expand_tabs : bool, optional + If true, tab characters will be expanded to spaces (default: True). + replace_whitespace : bool, optional + If true, each whitespace character remaining after tab expansion + will be replaced by a single space (default: True). + drop_whitespace : bool, optional + If true, whitespace that, after wrapping, happens to end up at the + beginning or end of a line is dropped (default: True). + break_long_words : bool, optional + If true, then words longer than width will be broken to + ensure that no lines are longer than width. If it is false, long + words will not be broken, and some lines may be longer than width + (default: True). + break_on_hyphens : bool, optional + If true, wrapping will occur preferably on whitespace and right + after hyphens in compound words, as it is customary in English. + If false, only whitespaces will be considered as potentially good + places for line breaks, but you need to set break_long_words to + false if you want truly insecable words (default: True). + + Returns + ------- + Series of object + Series with wrapped strings. + + Examples + -------- + >>> s = ps.Series(['line to be wrapped', 'another line to be wrapped']) + >>> s.str.wrap(12) + 0 line to be\\nwrapped + 1 another line\\nto be\\nwrapped + dtype: object + """ + + def pandas_wrap(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.wrap(width, **kwargs) + + return self._data.pandas_on_spark.transform_batch(pandas_wrap) + + def zfill(self, width: int) -> "ps.Series": + """ + Pad strings in the Series by prepending ‘0’ characters. + + Strings in the Series are padded with ‘0’ characters on the left of the + string to reach a total string length width. Strings in the Series with + length greater or equal to width are unchanged. + + Differs from :func:`str.zfill` which has special handling for ‘+’/’-‘ + in the string. + + Parameters + ---------- + width : int + Minimum length of resulting string; strings with length less than + width be prepended with ‘0’ characters. + + Returns + ------- + Series of object + Series with '0' left-padded strings. + + Examples + -------- + >>> s = ps.Series(['-1', '1', '1000', np.nan]) + >>> s + 0 -1 + 1 1 + 2 1000 + 3 None + dtype: object + + Note that NaN is not a string, therefore it is converted to NaN. The + minus sign in '-1' is treated as a regular character and the zero is + added to the left of it (:func:`str.zfill` would have moved it to the + left). 1000 remains unchanged as it is longer than width. + + >>> s.str.zfill(3) # doctest: +SKIP + 0 -01 + 1 001 + 2 1000 + 3 None + dtype: object + """ + + def pandas_zfill(s) -> ps.Series[str]: # type: ignore[no-untyped-def] + return s.str.zfill(width) + + return self._data.pandas_on_spark.transform_batch(pandas_zfill) + + @no_type_check + def get_dummies(self, sep: str = "|") -> "ps.DataFrame": + """ + Not supported. + """ + raise NotImplementedError() + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.strings + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.strings.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]") + .appName("pyspark.pandas.strings tests") + .getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.strings, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/supported_api_gen.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/supported_api_gen.py new file mode 100644 index 0000000..06591c5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/supported_api_gen.py @@ -0,0 +1,390 @@ +# +# 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. +# + +""" +Generate 'Supported pandas APIs' documentation file +""" +import warnings +from distutils.version import LooseVersion +from enum import Enum, unique +from inspect import getmembers, isclass, isfunction, signature +from typing import Any, Callable, Dict, List, NamedTuple, Set, TextIO, Tuple + +import pyspark.pandas as ps +import pyspark.pandas.groupby as psg +import pyspark.pandas.window as psw +from pyspark.pandas.exceptions import PandasNotImplementedError + +import pandas as pd +import pandas.core.groupby as pdg +import pandas.core.window as pdw + +MAX_MISSING_PARAMS_SIZE = 5 +COMMON_PARAMETER_SET = { + "kwargs", + "args", + "cls", +} # These are not counted as missing parameters. +MODULE_GROUP_MATCH = [(pd, ps), (pdw, psw), (pdg, psg)] + +RST_HEADER = """ +===================== +Supported pandas API +===================== + +.. currentmodule:: pyspark.pandas + +The following table shows the pandas APIs that implemented or non-implemented from pandas API on +Spark. Some pandas API do not implement full parameters, so the third column shows missing +parameters for each API. + +* 'Y' in the second column means it's implemented including its whole parameter. +* 'N' means it's not implemented yet. +* 'P' means it's partially implemented with the missing of some parameters. + +All API in the list below computes the data with distributed execution except the ones that require +the local execution by design. For example, `DataFrame.to_numpy() `__ +requires to collect the data to the driver side. + +If there is non-implemented pandas API or parameter you want, you can create an `Apache Spark +JIRA `__ to request or to contribute by +your own. + +The API list is updated based on the `latest pandas official API reference +`__. + +""" + + +@unique +class Implemented(Enum): + IMPLEMENTED = "Y" + NOT_IMPLEMENTED = "N" + PARTIALLY_IMPLEMENTED = "P" + + +class SupportedStatus(NamedTuple): + """ + Defines a supported status for specific pandas API + """ + + implemented: str + missing: str + + +def generate_supported_api(output_rst_file_path: str) -> None: + """ + Generate supported APIs status dictionary. + + Parameters + ---------- + output_rst_file_path : str + The path to the document file in RST format. + + Write supported APIs documentation. + """ + pandas_latest_version = "2.0.3" + if LooseVersion(pd.__version__) != LooseVersion(pandas_latest_version): + msg = ( + "Warning: Latest version of pandas (%s) is required to generate the documentation; " + "however, your version was %s" % (pandas_latest_version, pd.__version__) + ) + warnings.warn(msg, UserWarning) + raise ImportError(msg) + + all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]] = {} + for pd_module_group, ps_module_group in MODULE_GROUP_MATCH: + pd_modules = _get_pd_modules(pd_module_group) + _update_all_supported_status( + all_supported_status, pd_modules, pd_module_group, ps_module_group + ) + _write_rst(output_rst_file_path, all_supported_status) + + +def _create_supported_by_module( + module_name: str, pd_module_group: Any, ps_module_group: Any +) -> Dict[str, SupportedStatus]: + """ + Retrieves supported status of pandas module + + Parameters + ---------- + module_name : str + Class name that exists in the path of the module. + pd_module_group : Any + Specific path of importable pandas module. + ps_module_group: Any + Specific path of importable pyspark.pandas module. + """ + pd_module = getattr(pd_module_group, module_name) if module_name else pd_module_group + try: + ps_module = getattr(ps_module_group, module_name) if module_name else ps_module_group + except (AttributeError, PandasNotImplementedError): + # module not implemented + return {} + + pd_funcs = dict( + [ + m + for m in getmembers(pd_module, isfunction) + if not m[0].startswith("_") and m[0] in pd_module.__dict__ + ] + ) + if not pd_funcs: + return {} + + ps_funcs = dict( + [ + m + for m in getmembers(ps_module, isfunction) + if not m[0].startswith("_") and m[0] in ps_module.__dict__ + ] + ) + + return _organize_by_implementation_status( + module_name, pd_funcs, ps_funcs, pd_module_group, ps_module_group + ) + + +def _organize_by_implementation_status( + module_name: str, + pd_funcs: Dict[str, Callable], + ps_funcs: Dict[str, Callable], + pd_module_group: Any, + ps_module_group: Any, +) -> Dict[str, SupportedStatus]: + """ + Check the implementation status and parameters of both modules. + + Parameters + ---------- + module_name : str + Class name that exists in the path of the module. + pd_funcs: Dict[str, Callable] + function name and function object mapping of pandas module. + ps_funcs: Dict[str, Callable] + function name and function object mapping of pyspark.pandas module. + pd_module_group : Any + Specific path of importable pandas module. + ps_module_group: Any + Specific path of importable pyspark.pandas module. + """ + pd_dict = {} + for pd_func_name, pd_func in pd_funcs.items(): + ps_func = ps_funcs.get(pd_func_name) + if ps_func: + missing_set = ( + set(signature(pd_func).parameters) + - set(signature(ps_func).parameters) + - COMMON_PARAMETER_SET + ) + if missing_set: + # partially implemented + pd_dict[pd_func_name] = SupportedStatus( + implemented=Implemented.PARTIALLY_IMPLEMENTED.value, + missing=_transform_missing( + module_name, + pd_func_name, + missing_set, + pd_module_group.__name__, + ps_module_group.__name__, + ), + ) + else: + # implemented including it's whole parameter + pd_dict[pd_func_name] = SupportedStatus( + implemented=Implemented.IMPLEMENTED.value, missing="" + ) + else: + # not implemented yet + pd_dict[pd_func_name] = SupportedStatus( + implemented=Implemented.NOT_IMPLEMENTED.value, missing="" + ) + return pd_dict + + +def _transform_missing( + module_name: str, + pd_func_name: str, + missing_set: Set[str], + pd_module_path: str, + ps_module_path: str, +) -> str: + """ + Transform missing parameters into table information string. + + Parameters + ---------- + module_name : str + Class name that exists in the path of the module. + pd_func_name : str + Name of pandas API. + missing_set : Set[str] + A set of parameters not yet implemented. + pd_module_path : str + Path string of pandas module. + ps_module_path : str + Path string of pyspark.pandas module. + + Examples + -------- + >>> _transform_missing("DataFrame", "add", {"axis", "fill_value", "level"}, + ... "pandas.DataFrame", "pyspark.pandas.DataFrame") + '``axis`` , ``fill_value`` , ``level``' + """ + missing_str = " , ".join("``%s``" % x for x in sorted(missing_set)[:MAX_MISSING_PARAMS_SIZE]) + if len(missing_set) > MAX_MISSING_PARAMS_SIZE: + module_dot_func = "%s.%s" % (module_name, pd_func_name) if module_name else pd_func_name + additional_str = ( + " and more. See the " + + "`%s.%s " % (pd_module_path, module_dot_func) + + "`__ and " % (pd_module_path, module_dot_func) + + "`%s.%s " % (ps_module_path, module_dot_func) + + "`__ for detail." % (ps_module_path, module_dot_func) + ) + missing_str += additional_str + return missing_str + + +def _get_pd_modules(pd_module_group: Any) -> List[str]: + """ + Returns sorted pandas member list from pandas module path. + + Parameters + ---------- + pd_module_group : Any + Specific path of importable pandas module. + """ + return sorted([m[0] for m in getmembers(pd_module_group, isclass) if not m[0].startswith("_")]) + + +def _update_all_supported_status( + all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]], + pd_modules: List[str], + pd_module_group: Any, + ps_module_group: Any, +) -> None: + """ + Updates supported status across multiple module paths. + + Parameters + ---------- + all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]] + Data that stores the supported status across multiple module paths. + pd_modules: List[str] + Name list of pandas modules. + pd_module_group : Any + Specific path of importable pandas module. + ps_module_group: Any + Specific path of importable pyspark.pandas module. + """ + pd_modules += [""] # for General Function APIs + for module_name in pd_modules: + supported_status = _create_supported_by_module( + module_name, pd_module_group, ps_module_group + ) + if supported_status: + all_supported_status[(module_name, ps_module_group.__name__)] = supported_status + + +def _write_table( + module_name: str, + module_path: str, + supported_status: Dict[str, SupportedStatus], + w_fd: TextIO, +) -> None: + """ + Write table by using Sphinx list-table directive. + """ + lines = [] + if module_name: + lines.append(module_name) + else: + lines.append("General Function") + lines.append(" API\n") + lines.append("-" * 100) + lines.append("\n") + lines.append(".. currentmodule:: %s" % module_path) + if module_name: + lines.append(".%s\n" % module_name) + else: + lines.append("\n") + lines.append("\n") + lines.append(".. list-table::\n") + lines.append(" :header-rows: 1\n") + lines.append("\n") + lines.append(" * - API\n") + lines.append(" - Implemented\n") + lines.append(" - Missing parameters\n") + for func_str, status in supported_status.items(): + func_str = _escape_func_str(func_str) + if status.implemented == Implemented.NOT_IMPLEMENTED.value: + lines.append(" * - %s\n" % func_str) + else: + lines.append(" * - :func:`%s`\n" % func_str) + lines.append(" - %s\n" % status.implemented) + lines.append(" - \n") if not status.missing else lines.append( + " - %s\n" % status.missing + ) + w_fd.writelines(lines) + + +def _escape_func_str(func_str: str) -> str: + """ + Transforms which affecting rst data format. + """ + # TODO: Take into account that this function can create links incorrectly + # We can create alias links or links to parent methods + if func_str.endswith("_"): + return func_str[:-1] + "\_" # noqa: W605 + else: + return func_str + + +def _write_rst( + output_rst_file_path: str, + all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]], +) -> None: + """ + Writes the documentation to the target file path. + """ + with open(output_rst_file_path, "w") as w_fd: + w_fd.write(RST_HEADER) + for module_info, supported_status in all_supported_status.items(): + module, module_path = module_info + if supported_status: + _write_table(module, module_path, supported_status, w_fd) + w_fd.write("\n") + + +def _test() -> None: + import doctest + import sys + + import pyspark.pandas.supported_api_gen + + globs = pyspark.pandas.supported_api_gen.__dict__.copy() + (failure_count, test_count) = doctest.testmod(pyspark.pandas.supported_api_gen, globs=globs) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_any_all.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_any_all.py new file mode 100644 index 0000000..3574254 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_any_all.py @@ -0,0 +1,177 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameAnyAllMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43812): Enable DataFrameTests.test_all for pandas 2.0.0.", + ) + def test_all(self): + pdf = pd.DataFrame( + { + "col1": [False, False, False], + "col2": [True, False, False], + "col3": [0, 0, 1], + "col4": [0, 1, 2], + "col5": [False, False, None], + "col6": [True, False, None], + }, + index=np.random.rand(3), + ) + pdf.name = "x" + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.all(), pdf.all()) + self.assert_eq(psdf.all(bool_only=True), pdf.all(bool_only=True)) + self.assert_eq(psdf.all(bool_only=False), pdf.all(bool_only=False)) + self.assert_eq(psdf[["col5"]].all(bool_only=True), pdf[["col5"]].all(bool_only=True)) + self.assert_eq(psdf[["col5"]].all(bool_only=False), pdf[["col5"]].all(bool_only=False)) + + columns = pd.MultiIndex.from_tuples( + [ + ("a", "col1"), + ("a", "col2"), + ("a", "col3"), + ("b", "col4"), + ("b", "col5"), + ("c", "col6"), + ] + ) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.all(), pdf.all()) + self.assert_eq(psdf.all(bool_only=True), pdf.all(bool_only=True)) + self.assert_eq(psdf.all(bool_only=False), pdf.all(bool_only=False)) + + columns.names = ["X", "Y"] + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.all(), pdf.all()) + self.assert_eq(psdf.all(bool_only=True), pdf.all(bool_only=True)) + self.assert_eq(psdf.all(bool_only=False), pdf.all(bool_only=False)) + + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psdf.all(axis=1) + + # Test skipna + pdf = pd.DataFrame({"A": [True, True], "B": [1, np.nan], "C": [True, None]}) + pdf.name = "x" + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf[["A", "B"]].all(skipna=False), pdf[["A", "B"]].all(skipna=False)) + self.assert_eq(psdf[["A", "C"]].all(skipna=False), pdf[["A", "C"]].all(skipna=False)) + self.assert_eq(psdf[["B", "C"]].all(skipna=False), pdf[["B", "C"]].all(skipna=False)) + self.assert_eq(psdf.all(skipna=False), pdf.all(skipna=False)) + self.assert_eq(psdf.all(skipna=True), pdf.all(skipna=True)) + self.assert_eq(psdf.all(), pdf.all()) + self.assert_eq( + ps.DataFrame([np.nan]).all(skipna=False), pd.DataFrame([np.nan]).all(skipna=False) + ) + self.assert_eq(ps.DataFrame([None]).all(skipna=True), pd.DataFrame([None]).all(skipna=True)) + + def test_any(self): + pdf = pd.DataFrame( + { + "col1": [False, False, False], + "col2": [True, False, False], + "col3": [0, 0, 1], + "col4": [0, 1, 2], + "col5": [False, False, None], + "col6": [True, False, None], + }, + index=np.random.rand(3), + ) + pdf.name = "x" + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.any(), pdf.any()) + self.assert_eq(psdf.any(bool_only=True), pdf.any(bool_only=True)) + self.assert_eq(psdf.any(bool_only=False), pdf.any(bool_only=False)) + self.assert_eq(psdf[["col5"]].all(bool_only=True), pdf[["col5"]].all(bool_only=True)) + self.assert_eq(psdf[["col5"]].all(bool_only=False), pdf[["col5"]].all(bool_only=False)) + + columns = pd.MultiIndex.from_tuples( + [ + ("a", "col1"), + ("a", "col2"), + ("a", "col3"), + ("b", "col4"), + ("b", "col5"), + ("c", "col6"), + ] + ) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.any(), pdf.any()) + self.assert_eq(psdf.any(bool_only=True), pdf.any(bool_only=True)) + self.assert_eq(psdf.any(bool_only=False), pdf.any(bool_only=False)) + + columns.names = ["X", "Y"] + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.any(), pdf.any()) + self.assert_eq(psdf.any(bool_only=True), pdf.any(bool_only=True)) + self.assert_eq(psdf.any(bool_only=False), pdf.any(bool_only=False)) + + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psdf.any(axis=1) + + +class FrameAnyAllTests(FrameAnyAllMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_any_all import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_apply_func.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_apply_func.py new file mode 100644 index 0000000..37cc4a4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_apply_func.py @@ -0,0 +1,575 @@ +# +# 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 datetime import datetime +from distutils.version import LooseVersion +import sys +import unittest +from typing import List + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Function application, GroupBy & Window' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#function-application-groupby-window +# as well as 'apply_batch*' and 'transform_batch*'. +class FrameApplyFunctionMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_apply(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 100, + "b": [1.0, 1.0, 2.0, 3.0, 5.0, 8.0] * 100, + "c": [1, 4, 9, 16, 25, 36] * 100, + }, + columns=["a", "b", "c"], + index=np.random.rand(600), + ) + psdf = ps.DataFrame(pdf) + + self.assert_eq( + psdf.apply(lambda x: x + 1).sort_index(), pdf.apply(lambda x: x + 1).sort_index() + ) + self.assert_eq( + psdf.apply(lambda x, b: x + b, args=(1,)).sort_index(), + pdf.apply(lambda x, b: x + b, args=(1,)).sort_index(), + ) + self.assert_eq( + psdf.apply(lambda x, b: x + b, b=1).sort_index(), + pdf.apply(lambda x, b: x + b, b=1).sort_index(), + ) + + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.apply(lambda x: x + 1).sort_index(), pdf.apply(lambda x: x + 1).sort_index() + ) + self.assert_eq( + psdf.apply(lambda x, b: x + b, args=(1,)).sort_index(), + pdf.apply(lambda x, b: x + b, args=(1,)).sort_index(), + ) + self.assert_eq( + psdf.apply(lambda x, b: x + b, b=1).sort_index(), + pdf.apply(lambda x, b: x + b, b=1).sort_index(), + ) + + # returning a Series + self.assert_eq( + psdf.apply(lambda x: len(x), axis=1).sort_index(), + pdf.apply(lambda x: len(x), axis=1).sort_index(), + ) + self.assert_eq( + psdf.apply(lambda x, c: len(x) + c, axis=1, c=100).sort_index(), + pdf.apply(lambda x, c: len(x) + c, axis=1, c=100).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.apply(lambda x: len(x), axis=1).sort_index(), + pdf.apply(lambda x: len(x), axis=1).sort_index(), + ) + self.assert_eq( + psdf.apply(lambda x, c: len(x) + c, axis=1, c=100).sort_index(), + pdf.apply(lambda x, c: len(x) + c, axis=1, c=100).sort_index(), + ) + + with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"): + psdf.apply(1) + + with self.assertRaisesRegex(TypeError, "The given function.*1 or 'column'; however"): + + def f1(_) -> ps.DataFrame[int]: + pass + + psdf.apply(f1, axis=0) + + with self.assertRaisesRegex(TypeError, "The given function.*0 or 'index'; however"): + + def f2(_) -> ps.Series[int]: + pass + + psdf.apply(f2, axis=1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.apply(lambda x: x + 1).sort_index(), pdf.apply(lambda x: x + 1).sort_index() + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.apply(lambda x: x + 1).sort_index(), pdf.apply(lambda x: x + 1).sort_index() + ) + + # returning a Series + self.assert_eq( + psdf.apply(lambda x: len(x), axis=1).sort_index(), + pdf.apply(lambda x: len(x), axis=1).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.apply(lambda x: len(x), axis=1).sort_index(), + pdf.apply(lambda x: len(x), axis=1).sort_index(), + ) + + def test_apply_with_type(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + + def identify1(x) -> ps.DataFrame[int, int]: + return x + + # Type hints set the default column names, and we use default index for + # pandas API on Spark. Here we ignore both diff. + actual = psdf.apply(identify1, axis=1) + expected = pdf.apply(identify1, axis=1) + self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy())) + + def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405 + return x + + actual = psdf.apply(identify2, axis=1) + expected = pdf.apply(identify2, axis=1) + self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy())) + + def test_apply_batch(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 100, + "b": [1.0, 1.0, 2.0, 3.0, 5.0, 8.0] * 100, + "c": [1, 4, 9, 16, 25, 36] * 100, + }, + columns=["a", "b", "c"], + index=np.random.rand(600), + ) + psdf = ps.DataFrame(pdf) + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda pdf, a: pdf + a, args=(1,)).sort_index(), + (pdf + 1).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda pdf: pdf + 1).sort_index(), + (pdf + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda pdf, b: pdf + b, b=1).sort_index(), + (pdf + 1).sort_index(), + ) + + with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"): + psdf.pandas_on_spark.apply_batch(1) + + with self.assertRaisesRegex(TypeError, "The given function.*frame as its type hints"): + + def f2(_) -> ps.Series[int]: + pass + + psdf.pandas_on_spark.apply_batch(f2) + + with self.assertRaisesRegex(ValueError, "The given function should return a frame"): + psdf.pandas_on_spark.apply_batch(lambda pdf: 1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda x: x + 1).sort_index(), (pdf + 1).sort_index() + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda x: x + 1).sort_index(), + (pdf + 1).sort_index(), + ) + + def test_apply_batch_with_type(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + + def identify1(x) -> ps.DataFrame[int, int]: + return x + + # Type hints set the default column names, and we use default index for + # pandas API on Spark. Here we ignore both diff. + actual = psdf.pandas_on_spark.apply_batch(identify1) + expected = pdf + self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy())) + + def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405 + return x + + actual = psdf.pandas_on_spark.apply_batch(identify2) + expected = pdf + self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy())) + + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [[e] for e in [4, 5, 6, 3, 2, 1, 0, 0, 0]]}, + index=np.random.rand(9), + ) + psdf = ps.from_pandas(pdf) + + def identify3(x) -> ps.DataFrame[float, [int, List[int]]]: + return x + + actual = psdf.pandas_on_spark.apply_batch(identify3) + actual.columns = ["a", "b"] + self.assert_eq(actual, pdf) + + # For NumPy typing, NumPy version should be 1.21+ and Python version should be 3.8+ + if sys.version_info >= (3, 8) and LooseVersion(np.__version__) >= LooseVersion("1.21"): + import numpy.typing as ntp + + psdf = ps.from_pandas(pdf) + + def identify4( + x, + ) -> ps.DataFrame[float, [int, ntp.NDArray[int]]]: + return x + + actual = psdf.pandas_on_spark.apply_batch(identify4) + actual.columns = ["a", "b"] + self.assert_eq(actual, pdf) + + arrays = [[1, 2, 3, 4, 5, 6, 7, 8, 9], ["a", "b", "c", "d", "e", "f", "g", "h", "i"]] + idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [[e] for e in [4, 5, 6, 3, 2, 1, 0, 0, 0]]}, + index=idx, + ) + psdf = ps.from_pandas(pdf) + + def identify4(x) -> ps.DataFrame[[int, str], [int, List[int]]]: + return x + + actual = psdf.pandas_on_spark.apply_batch(identify4) + actual.index.names = ["number", "color"] + actual.columns = ["a", "b"] + self.assert_eq(actual, pdf) + + def identify5( + x, + ) -> ps.DataFrame[ + [("number", int), ("color", str)], [("a", int), ("b", List[int])] # noqa: F405 + ]: + return x + + actual = psdf.pandas_on_spark.apply_batch(identify5) + self.assert_eq(actual, pdf) + + def test_transform(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 100, + "b": [1.0, 1.0, 2.0, 3.0, 5.0, 8.0] * 100, + "c": [1, 4, 9, 16, 25, 36] * 100, + }, + columns=["a", "b", "c"], + index=np.random.rand(600), + ) + psdf = ps.DataFrame(pdf) + self.assert_eq( + psdf.transform(lambda x: x + 1).sort_index(), + pdf.transform(lambda x: x + 1).sort_index(), + ) + self.assert_eq( + psdf.transform(lambda x, y: x + y, y=2).sort_index(), + pdf.transform(lambda x, y: x + y, y=2).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.transform(lambda x: x + 1).sort_index(), + pdf.transform(lambda x: x + 1).sort_index(), + ) + self.assert_eq( + psdf.transform(lambda x, y: x + y, y=1).sort_index(), + pdf.transform(lambda x, y: x + y, y=1).sort_index(), + ) + + with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"): + psdf.transform(1) + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psdf.transform(lambda x: x + 1, axis=1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.transform(lambda x: x + 1).sort_index(), + pdf.transform(lambda x: x + 1).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.transform(lambda x: x + 1).sort_index(), + pdf.transform(lambda x: x + 1).sort_index(), + ) + + def test_transform_batch(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 100, + "b": [1.0, 1.0, 2.0, 3.0, 5.0, 8.0] * 100, + "c": [1, 4, 9, 16, 25, 36] * 100, + }, + columns=["a", "b", "c"], + index=np.random.rand(600), + ) + psdf = ps.DataFrame(pdf) + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.c + 1).sort_index(), + (pdf.c + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf, a: pdf + a, 1).sort_index(), + (pdf + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf, a: pdf.c + a, a=1).sort_index(), + (pdf.c + 1).sort_index(), + ) + + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf + 1).sort_index(), + (pdf + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.b + 1).sort_index(), + (pdf.b + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf, a: pdf + a, 1).sort_index(), + (pdf + 1).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf, a: pdf.c + a, a=1).sort_index(), + (pdf.c + 1).sort_index(), + ) + + with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"): + psdf.pandas_on_spark.transform_batch(1) + + with self.assertRaisesRegex(ValueError, "The given function should return a frame"): + psdf.pandas_on_spark.transform_batch(lambda pdf: 1) + + with self.assertRaisesRegex( + ValueError, "transform_batch cannot produce aggregated results" + ): + psdf.pandas_on_spark.transform_batch(lambda pdf: pd.Series(1)) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda x: x + 1).sort_index(), + (pdf + 1).sort_index(), + ) + with option_context("compute.shortcut_limit", 500): + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda x: x + 1).sort_index(), + (pdf + 1).sort_index(), + ) + + def test_transform_batch_with_type(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + + def identify1(x) -> ps.DataFrame[int, int]: + return x + + # Type hints set the default column names, and we use default index for + # pandas API on Spark. Here we ignore both diff. + actual = psdf.pandas_on_spark.transform_batch(identify1) + expected = pdf + self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy())) + + def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405 + return x + + actual = psdf.pandas_on_spark.transform_batch(identify2) + expected = pdf + self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy())) + + def test_transform_batch_same_anchor(self): + psdf = ps.range(10) + psdf["d"] = psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.id + 1) + self.assert_eq( + psdf, + pd.DataFrame({"id": list(range(10)), "d": list(range(1, 11))}, columns=["id", "d"]), + ) + + psdf = ps.range(10) + + def plus_one(pdf) -> ps.Series[np.int64]: + return pdf.id + 1 + + psdf["d"] = psdf.pandas_on_spark.transform_batch(plus_one) + self.assert_eq( + psdf, + pd.DataFrame({"id": list(range(10)), "d": list(range(1, 11))}, columns=["id", "d"]), + ) + + psdf = ps.range(10) + + def plus_one(ser) -> ps.Series[np.int64]: + return ser + 1 + + psdf["d"] = psdf.id.pandas_on_spark.transform_batch(plus_one) + self.assert_eq( + psdf, + pd.DataFrame({"id": list(range(10)), "d": list(range(1, 11))}, columns=["id", "d"]), + ) + + def test_pipe(self): + psdf = ps.DataFrame( + {"category": ["A", "A", "B"], "col1": [1, 2, 3], "col2": [4, 5, 6]}, + columns=["category", "col1", "col2"], + ) + + self.assertRaisesRegex( + ValueError, + "arg is both the pipe target and a keyword argument", + lambda: psdf.pipe((lambda x: x, "arg"), arg="1"), + ) + + def test_aggregate(self): + pdf = pd.DataFrame( + [[1, 2, 3], [4, 5, 6], [7, 8, 9], [np.nan, np.nan, np.nan]], columns=["A", "B", "C"] + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.agg(["sum", "min"])[["A", "B", "C"]].sort_index(), # TODO?: fix column order + pdf.agg(["sum", "min"])[["A", "B", "C"]].sort_index(), + ) + self.assert_eq( + psdf.agg({"A": ["sum", "min"], "B": ["min", "max"]})[["A", "B"]].sort_index(), + pdf.agg({"A": ["sum", "min"], "B": ["min", "max"]})[["A", "B"]].sort_index(), + ) + + self.assertRaises(KeyError, lambda: psdf.agg({"A": ["sum", "min"], "X": ["min", "max"]})) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.agg(["sum", "min"])[[("X", "A"), ("X", "B"), ("Y", "C")]].sort_index(), + pdf.agg(["sum", "min"])[[("X", "A"), ("X", "B"), ("Y", "C")]].sort_index(), + ) + self.assert_eq( + psdf.agg({("X", "A"): ["sum", "min"], ("X", "B"): ["min", "max"]})[ + [("X", "A"), ("X", "B")] + ].sort_index(), + pdf.agg({("X", "A"): ["sum", "min"], ("X", "B"): ["min", "max"]})[ + [("X", "A"), ("X", "B")] + ].sort_index(), + ) + + self.assertRaises(TypeError, lambda: psdf.agg({"X": ["sum", "min"], "Y": ["min", "max"]})) + + # non-string names + pdf = pd.DataFrame( + [[1, 2, 3], [4, 5, 6], [7, 8, 9], [np.nan, np.nan, np.nan]], columns=[10, 20, 30] + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.agg(["sum", "min"])[[10, 20, 30]].sort_index(), + pdf.agg(["sum", "min"])[[10, 20, 30]].sort_index(), + ) + self.assert_eq( + psdf.agg({10: ["sum", "min"], 20: ["min", "max"]})[[10, 20]].sort_index(), + pdf.agg({10: ["sum", "min"], 20: ["min", "max"]})[[10, 20]].sort_index(), + ) + + columns = pd.MultiIndex.from_tuples([("X", 10), ("X", 20), ("Y", 30)]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.agg(["sum", "min"])[[("X", 10), ("X", 20), ("Y", 30)]].sort_index(), + pdf.agg(["sum", "min"])[[("X", 10), ("X", 20), ("Y", 30)]].sort_index(), + ) + self.assert_eq( + psdf.agg({("X", 10): ["sum", "min"], ("X", 20): ["min", "max"]})[ + [("X", 10), ("X", 20)] + ].sort_index(), + pdf.agg({("X", 10): ["sum", "min"], ("X", 20): ["min", "max"]})[ + [("X", 10), ("X", 20)] + ].sort_index(), + ) + + pdf = pd.DataFrame( + [datetime(2019, 2, 2, 0, 0, 0, 0), datetime(2019, 2, 3, 0, 0, 0, 0)], + columns=["timestamp"], + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.timestamp.min(), pdf.timestamp.min()) + self.assert_eq(psdf.timestamp.max(), pdf.timestamp.max()) + + self.assertRaises(ValueError, lambda: psdf.agg(("sum", "min"))) + + +class FrameApplyFunctionTests(FrameApplyFunctionMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_apply_func import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_binary_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_binary_ops.py new file mode 100644 index 0000000..e382b3e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_binary_ops.py @@ -0,0 +1,235 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Binary operator functions' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#binary-operator-functions +class FrameBinaryOpsMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_binary_operators(self): + pdf = pd.DataFrame( + {"A": [0, 2, 4], "B": [4, 2, 0], "X": [-1, 10, 0]}, index=np.random.rand(3) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf + psdf.copy(), pdf + pdf.copy()) + self.assert_eq(psdf + psdf.loc[:, ["A", "B"]], pdf + pdf.loc[:, ["A", "B"]]) + self.assert_eq(psdf.loc[:, ["A", "B"]] + psdf, pdf.loc[:, ["A", "B"]] + pdf) + + self.assertRaisesRegex( + ValueError, + "it comes from a different dataframe", + lambda: ps.range(10).add(ps.range(10)), + ) + + self.assertRaisesRegex( + TypeError, + "add with a sequence is currently not supported", + lambda: ps.range(10).add(ps.range(10).id), + ) + + psdf_other = psdf.copy() + psdf_other.columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "C")]) + self.assertRaisesRegex( + ValueError, + "cannot join with no overlapping index names", + lambda: psdf.add(psdf_other), + ) + + def test_binary_operator_add(self): + # Positive + pdf = pd.DataFrame({"a": ["x"], "b": ["y"], "c": [1], "d": [2]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["a"] + psdf["b"], pdf["a"] + pdf["b"]) + self.assert_eq(psdf["c"] + psdf["d"], pdf["c"] + pdf["d"]) + + # Negative + ks_err_msg = "Addition can not be applied to given types" + + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] + psdf["c"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["c"] + psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["c"] + "literal") + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" + psdf["c"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 + psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] + 1) + + def test_binary_operator_sub(self): + # Positive + pdf = pd.DataFrame({"a": [2], "b": [1]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["a"] - psdf["b"], pdf["a"] - pdf["b"]) + + # Negative + psdf = ps.DataFrame({"a": ["x"], "b": [1]}) + ks_err_msg = "Subtraction can not be applied to given types" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] - psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] - "literal") + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" - psdf["b"]) + + ks_err_msg = "Subtraction can not be applied to strings" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] - psdf["b"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 - psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] - 1) + + psdf = ps.DataFrame({"a": ["x"], "b": ["y"]}) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] - psdf["b"]) + + def test_binary_operator_truediv(self): + # Positive + pdf = pd.DataFrame({"a": [3], "b": [2]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["a"] / psdf["b"], pdf["a"] / pdf["b"]) + + # Negative + psdf = ps.DataFrame({"a": ["x"], "b": [1]}) + + ks_err_msg = "True division can not be applied to given types" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] / psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] / "literal") + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" / psdf["b"]) + + ks_err_msg = "True division can not be applied to strings" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] / psdf["b"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 / psdf["a"]) + + def test_binary_operator_floordiv(self): + psdf = ps.DataFrame({"a": ["x"], "b": [1]}) + + ks_err_msg = "Floor division can not be applied to strings" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] // psdf["b"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 // psdf["a"]) + + ks_err_msg = "Floor division can not be applied to given types" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] // psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] // "literal") + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" // psdf["b"]) + + def test_binary_operator_mod(self): + # Positive + pdf = pd.DataFrame({"a": [3], "b": [2]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["a"] % psdf["b"], pdf["a"] % pdf["b"]) + + # Negative + psdf = ps.DataFrame({"a": ["x"], "b": [1]}) + ks_err_msg = "Modulo can not be applied to given types" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] % psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] % "literal") + + ks_err_msg = "Modulo can not be applied to strings" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] % psdf["b"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 1 % psdf["a"]) + + def test_binary_operator_multiply(self): + # Positive + pdf = pd.DataFrame({"a": ["x", "y"], "b": [1, 2], "c": [3, 4]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["b"] * psdf["c"], pdf["b"] * pdf["c"]) + self.assert_eq(psdf["c"] * psdf["b"], pdf["c"] * pdf["b"]) + self.assert_eq(psdf["a"] * psdf["b"], pdf["a"] * pdf["b"]) + self.assert_eq(psdf["b"] * psdf["a"], pdf["b"] * pdf["a"]) + self.assert_eq(psdf["a"] * 2, pdf["a"] * 2) + self.assert_eq(psdf["b"] * 2, pdf["b"] * 2) + self.assert_eq(2 * psdf["a"], 2 * pdf["a"]) + self.assert_eq(2 * psdf["b"], 2 * pdf["b"]) + + # Negative + psdf = ps.DataFrame({"a": ["x"], "b": [2]}) + ks_err_msg = "Multiplication can not be applied to given types" + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["b"] * "literal") + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" * psdf["b"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] * "literal") + + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] * psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] * 0.1) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: 0.1 * psdf["a"]) + self.assertRaisesRegex(TypeError, ks_err_msg, lambda: "literal" * psdf["a"]) + + def test_combine_first(self): + pdf = pd.DataFrame( + {("X", "A"): [None, 0], ("X", "B"): [4, None], ("Y", "C"): [3, 3], ("Y", "B"): [1, 1]} + ) + pdf1, pdf2 = pdf["X"], pdf["Y"] + psdf = ps.from_pandas(pdf) + psdf1, psdf2 = psdf["X"], psdf["Y"] + + if LooseVersion(pd.__version__) >= LooseVersion("1.2.0"): + self.assert_eq(pdf1.combine_first(pdf2), psdf1.combine_first(psdf2)) + else: + # pandas < 1.2.0 returns unexpected dtypes, + # please refer to https://github.com/pandas-dev/pandas/issues/28481 for details + expected_pdf = pd.DataFrame({"A": [None, 0], "B": [4.0, 1.0], "C": [3, 3]}) + self.assert_eq(expected_pdf, psdf1.combine_first(psdf2)) + + def test_dot(self): + psdf = self.psdf + + with self.assertRaisesRegex(TypeError, "Unsupported type DataFrame"): + psdf.dot(psdf) + + def test_rfloordiv(self): + pdf = pd.DataFrame( + {"angles": [0, 3, 4], "degrees": [360, 180, 360]}, + index=["circle", "triangle", "rectangle"], + columns=["angles", "degrees"], + ) + psdf = ps.from_pandas(pdf) + + expected_result = pdf.rfloordiv(10) + self.assert_eq(psdf.rfloordiv(10), expected_result) + + +class FrameBinaryOpsTests(FrameBinaryOpsMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_binary_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_combine.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_combine.py new file mode 100644 index 0000000..dd55c0f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_combine.py @@ -0,0 +1,653 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Combining / joining / merging' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#combining-joining-merging +class FrameCombineMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43562): Enable DataFrameTests.test_append for pandas 2.0.0.", + ) + def test_append(self): + pdf = pd.DataFrame([[1, 2], [3, 4]], columns=list("AB")) + psdf = ps.from_pandas(pdf) + other_pdf = pd.DataFrame([[3, 4], [5, 6]], columns=list("BC"), index=[2, 3]) + other_psdf = ps.from_pandas(other_pdf) + + self.assert_eq(psdf.append(psdf), pdf.append(pdf)) + self.assert_eq(psdf.append(psdf, ignore_index=True), pdf.append(pdf, ignore_index=True)) + + # Assert DataFrames with non-matching columns + self.assert_eq(psdf.append(other_psdf), pdf.append(other_pdf)) + + # Assert appending a Series fails + msg = "DataFrames.append() does not support appending Series to DataFrames" + with self.assertRaises(TypeError, msg=msg): + psdf.append(psdf["A"]) + + # Assert using the sort parameter raises an exception + msg = "The 'sort' parameter is currently not supported" + with self.assertRaises(NotImplementedError, msg=msg): + psdf.append(psdf, sort=True) + + # Assert using 'verify_integrity' only raises an exception for overlapping indices + self.assert_eq( + psdf.append(other_psdf, verify_integrity=True), + pdf.append(other_pdf, verify_integrity=True), + ) + msg = "Indices have overlapping values" + with self.assertRaises(ValueError, msg=msg): + psdf.append(psdf, verify_integrity=True) + + # Skip integrity verification when ignore_index=True + self.assert_eq( + psdf.append(psdf, ignore_index=True, verify_integrity=True), + pdf.append(pdf, ignore_index=True, verify_integrity=True), + ) + + # Assert appending multi-index DataFrames + multi_index_pdf = pd.DataFrame([[1, 2], [3, 4]], columns=list("AB"), index=[[2, 3], [4, 5]]) + multi_index_psdf = ps.from_pandas(multi_index_pdf) + other_multi_index_pdf = pd.DataFrame( + [[5, 6], [7, 8]], columns=list("AB"), index=[[2, 3], [6, 7]] + ) + other_multi_index_psdf = ps.from_pandas(other_multi_index_pdf) + + self.assert_eq( + multi_index_psdf.append(multi_index_psdf), multi_index_pdf.append(multi_index_pdf) + ) + + # Assert DataFrames with non-matching columns + self.assert_eq( + multi_index_psdf.append(other_multi_index_psdf), + multi_index_pdf.append(other_multi_index_pdf), + ) + + # Assert using 'verify_integrity' only raises an exception for overlapping indices + self.assert_eq( + multi_index_psdf.append(other_multi_index_psdf, verify_integrity=True), + multi_index_pdf.append(other_multi_index_pdf, verify_integrity=True), + ) + with self.assertRaises(ValueError, msg=msg): + multi_index_psdf.append(multi_index_psdf, verify_integrity=True) + + # Skip integrity verification when ignore_index=True + self.assert_eq( + multi_index_psdf.append(multi_index_psdf, ignore_index=True, verify_integrity=True), + multi_index_pdf.append(multi_index_pdf, ignore_index=True, verify_integrity=True), + ) + + # Assert trying to append DataFrames with different index levels + msg = "Both DataFrames have to have the same number of index levels" + with self.assertRaises(ValueError, msg=msg): + psdf.append(multi_index_psdf) + + # Skip index level check when ignore_index=True + self.assert_eq( + psdf.append(multi_index_psdf, ignore_index=True), + pdf.append(multi_index_pdf, ignore_index=True), + ) + + columns = pd.MultiIndex.from_tuples([("A", "X"), ("A", "Y")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.append(psdf), pdf.append(pdf)) + + def test_merge(self): + left_pdf = pd.DataFrame( + { + "lkey": ["foo", "bar", "baz", "foo", "bar", "l"], + "value": [1, 2, 3, 5, 6, 7], + "x": list("abcdef"), + }, + columns=["lkey", "value", "x"], + ) + right_pdf = pd.DataFrame( + { + "rkey": ["baz", "foo", "bar", "baz", "foo", "r"], + "value": [4, 5, 6, 7, 8, 9], + "y": list("efghij"), + }, + columns=["rkey", "value", "y"], + ) + right_pser = pd.Series(list("defghi"), name="x", index=[5, 6, 7, 8, 9, 10]) + + left_psdf = ps.from_pandas(left_pdf) + right_psdf = ps.from_pandas(right_pdf) + right_psser = ps.from_pandas(right_pser) + + def check(op, right_psdf=right_psdf, right_pdf=right_pdf): + ps_res = op(left_psdf, right_psdf) + ps_res = ps_res._to_pandas() + ps_res = ps_res.sort_values(by=list(ps_res.columns)) + ps_res = ps_res.reset_index(drop=True) + p_res = op(left_pdf, right_pdf) + p_res = p_res.sort_values(by=list(p_res.columns)) + p_res = p_res.reset_index(drop=True) + self.assert_eq(ps_res, p_res) + + check(lambda left, right: left.merge(right)) + check(lambda left, right: left.merge(right, on="value")) + check(lambda left, right: left.merge(right, on=("value",))) + check(lambda left, right: left.merge(right, left_on="lkey", right_on="rkey")) + check(lambda left, right: left.set_index("lkey").merge(right.set_index("rkey"))) + check( + lambda left, right: left.set_index("lkey").merge( + right, left_index=True, right_on="rkey" + ) + ) + check( + lambda left, right: left.merge( + right.set_index("rkey"), left_on="lkey", right_index=True + ) + ) + check( + lambda left, right: left.set_index("lkey").merge( + right.set_index("rkey"), left_index=True, right_index=True + ) + ) + + # MultiIndex + check( + lambda left, right: left.merge( + right, left_on=["lkey", "value"], right_on=["rkey", "value"] + ) + ) + check( + lambda left, right: left.set_index(["lkey", "value"]).merge( + right, left_index=True, right_on=["rkey", "value"] + ) + ) + check( + lambda left, right: left.merge( + right.set_index(["rkey", "value"]), left_on=["lkey", "value"], right_index=True + ) + ) + # TODO: when both left_index=True and right_index=True with multi-index + # check(lambda left, right: left.set_index(['lkey', 'value']).merge( + # right.set_index(['rkey', 'value']), left_index=True, right_index=True)) + + # join types + for how in ["inner", "left", "right", "outer"]: + check(lambda left, right: left.merge(right, on="value", how=how)) + check(lambda left, right: left.merge(right, left_on="lkey", right_on="rkey", how=how)) + + # suffix + check( + lambda left, right: left.merge( + right, left_on="lkey", right_on="rkey", suffixes=["_left", "_right"] + ) + ) + + # Test Series on the right + check(lambda left, right: left.merge(right), right_psser, right_pser) + check( + lambda left, right: left.merge(right, left_on="x", right_on="x"), + right_psser, + right_pser, + ) + check( + lambda left, right: left.set_index("x").merge(right, left_index=True, right_on="x"), + right_psser, + right_pser, + ) + + # Test join types with Series + for how in ["inner", "left", "right", "outer"]: + check(lambda left, right: left.merge(right, how=how), right_psser, right_pser) + check( + lambda left, right: left.merge(right, left_on="x", right_on="x", how=how), + right_psser, + right_pser, + ) + + # suffix with Series + check( + lambda left, right: left.merge( + right, + suffixes=["_left", "_right"], + how="outer", + left_index=True, + right_index=True, + ), + right_psser, + right_pser, + ) + + # multi-index columns + left_columns = pd.MultiIndex.from_tuples([(10, "lkey"), (10, "value"), (20, "x")]) + left_pdf.columns = left_columns + left_psdf.columns = left_columns + + right_columns = pd.MultiIndex.from_tuples([(10, "rkey"), (10, "value"), (30, "y")]) + right_pdf.columns = right_columns + right_psdf.columns = right_columns + + check(lambda left, right: left.merge(right)) + check(lambda left, right: left.merge(right, on=[(10, "value")])) + check( + lambda left, right: (left.set_index((10, "lkey")).merge(right.set_index((10, "rkey")))) + ) + check( + lambda left, right: ( + left.set_index((10, "lkey")).merge( + right.set_index((10, "rkey")), left_index=True, right_index=True + ) + ) + ) + # TODO: when both left_index=True and right_index=True with multi-index columns + # check(lambda left, right: left.merge(right, + # left_on=[('a', 'lkey')], right_on=[('a', 'rkey')])) + # check(lambda left, right: (left.set_index(('a', 'lkey')) + # .merge(right, left_index=True, right_on=[('a', 'rkey')]))) + + # non-string names + left_pdf.columns = [10, 100, 1000] + left_psdf.columns = [10, 100, 1000] + + right_pdf.columns = [20, 100, 2000] + right_psdf.columns = [20, 100, 2000] + + check(lambda left, right: left.merge(right)) + check(lambda left, right: left.merge(right, on=[100])) + check(lambda left, right: (left.set_index(10).merge(right.set_index(20)))) + check( + lambda left, right: ( + left.set_index(10).merge(right.set_index(20), left_index=True, right_index=True) + ) + ) + + def test_merge_same_anchor(self): + pdf = pd.DataFrame( + { + "lkey": ["foo", "bar", "baz", "foo", "bar", "l"], + "rkey": ["baz", "foo", "bar", "baz", "foo", "r"], + "value": [1, 1, 3, 5, 6, 7], + "x": list("abcdef"), + "y": list("efghij"), + }, + columns=["lkey", "rkey", "value", "x", "y"], + ) + psdf = ps.from_pandas(pdf) + + left_pdf = pdf[["lkey", "value", "x"]] + right_pdf = pdf[["rkey", "value", "y"]] + left_psdf = psdf[["lkey", "value", "x"]] + right_psdf = psdf[["rkey", "value", "y"]] + + def check(op, right_psdf=right_psdf, right_pdf=right_pdf): + k_res = op(left_psdf, right_psdf) + k_res = k_res._to_pandas() + k_res = k_res.sort_values(by=list(k_res.columns)) + k_res = k_res.reset_index(drop=True) + p_res = op(left_pdf, right_pdf) + p_res = p_res.sort_values(by=list(p_res.columns)) + p_res = p_res.reset_index(drop=True) + self.assert_eq(k_res, p_res) + + check(lambda left, right: left.merge(right)) + check(lambda left, right: left.merge(right, on="value")) + check(lambda left, right: left.merge(right, left_on="lkey", right_on="rkey")) + check(lambda left, right: left.set_index("lkey").merge(right.set_index("rkey"))) + check( + lambda left, right: left.set_index("lkey").merge( + right, left_index=True, right_on="rkey" + ) + ) + check( + lambda left, right: left.merge( + right.set_index("rkey"), left_on="lkey", right_index=True + ) + ) + check( + lambda left, right: left.set_index("lkey").merge( + right.set_index("rkey"), left_index=True, right_index=True + ) + ) + + def test_merge_retains_indices(self): + left_pdf = pd.DataFrame({"A": [0, 1]}) + right_pdf = pd.DataFrame({"B": [1, 2]}, index=[1, 2]) + left_psdf = ps.from_pandas(left_pdf) + right_psdf = ps.from_pandas(right_pdf) + + self.assert_eq( + left_psdf.merge(right_psdf, left_index=True, right_index=True), + left_pdf.merge(right_pdf, left_index=True, right_index=True), + ) + self.assert_eq( + left_psdf.merge(right_psdf, left_on="A", right_index=True), + left_pdf.merge(right_pdf, left_on="A", right_index=True), + ) + self.assert_eq( + left_psdf.merge(right_psdf, left_index=True, right_on="B"), + left_pdf.merge(right_pdf, left_index=True, right_on="B"), + ) + self.assert_eq( + left_psdf.merge(right_psdf, left_on="A", right_on="B"), + left_pdf.merge(right_pdf, left_on="A", right_on="B"), + ) + + def test_merge_how_parameter(self): + left_pdf = pd.DataFrame({"A": [1, 2]}) + right_pdf = pd.DataFrame({"B": ["x", "y"]}, index=[1, 2]) + left_psdf = ps.from_pandas(left_pdf) + right_psdf = ps.from_pandas(right_pdf) + + psdf = left_psdf.merge(right_psdf, left_index=True, right_index=True) + pdf = left_pdf.merge(right_pdf, left_index=True, right_index=True) + self.assert_eq( + psdf.sort_values(by=list(psdf.columns)).reset_index(drop=True), + pdf.sort_values(by=list(pdf.columns)).reset_index(drop=True), + ) + + psdf = left_psdf.merge(right_psdf, left_index=True, right_index=True, how="left") + pdf = left_pdf.merge(right_pdf, left_index=True, right_index=True, how="left") + self.assert_eq( + psdf.sort_values(by=list(psdf.columns)).reset_index(drop=True), + pdf.sort_values(by=list(pdf.columns)).reset_index(drop=True), + ) + + psdf = left_psdf.merge(right_psdf, left_index=True, right_index=True, how="right") + pdf = left_pdf.merge(right_pdf, left_index=True, right_index=True, how="right") + self.assert_eq( + psdf.sort_values(by=list(psdf.columns)).reset_index(drop=True), + pdf.sort_values(by=list(pdf.columns)).reset_index(drop=True), + ) + + psdf = left_psdf.merge(right_psdf, left_index=True, right_index=True, how="outer") + pdf = left_pdf.merge(right_pdf, left_index=True, right_index=True, how="outer") + self.assert_eq( + psdf.sort_values(by=list(psdf.columns)).reset_index(drop=True), + pdf.sort_values(by=list(pdf.columns)).reset_index(drop=True), + ) + + def test_merge_raises(self): + left = ps.DataFrame( + {"value": [1, 2, 3, 5, 6], "x": list("abcde")}, + columns=["value", "x"], + index=["foo", "bar", "baz", "foo", "bar"], + ) + right = ps.DataFrame( + {"value": [4, 5, 6, 7, 8], "y": list("fghij")}, + columns=["value", "y"], + index=["baz", "foo", "bar", "baz", "foo"], + ) + + with self.assertRaisesRegex(ValueError, "No common columns to perform merge on"): + left[["x"]].merge(right[["y"]]) + + with self.assertRaisesRegex(ValueError, "not a combination of both"): + left.merge(right, on="value", left_on="x") + + with self.assertRaisesRegex(ValueError, "Must pass right_on or right_index=True"): + left.merge(right, left_on="x") + + with self.assertRaisesRegex(ValueError, "Must pass right_on or right_index=True"): + left.merge(right, left_index=True) + + with self.assertRaisesRegex(ValueError, "Must pass left_on or left_index=True"): + left.merge(right, right_on="y") + + with self.assertRaisesRegex(ValueError, "Must pass left_on or left_index=True"): + left.merge(right, right_index=True) + + with self.assertRaisesRegex( + ValueError, "len\\(left_keys\\) must equal len\\(right_keys\\)" + ): + left.merge(right, left_on="value", right_on=["value", "y"]) + + with self.assertRaisesRegex( + ValueError, "len\\(left_keys\\) must equal len\\(right_keys\\)" + ): + left.merge(right, left_on=["value", "x"], right_on="value") + + with self.assertRaisesRegex(ValueError, "['inner', 'left', 'right', 'full', 'outer']"): + left.merge(right, left_index=True, right_index=True, how="foo") + + with self.assertRaisesRegex(KeyError, "id"): + left.merge(right, on="id") + + def test_join(self): + # check basic function + pdf1 = pd.DataFrame( + {"key": ["K0", "K1", "K2", "K3"], "A": ["A0", "A1", "A2", "A3"]}, columns=["key", "A"] + ) + pdf2 = pd.DataFrame( + {"key": ["K0", "K1", "K2"], "B": ["B0", "B1", "B2"]}, columns=["key", "B"] + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + join_pdf = pdf1.join(pdf2, lsuffix="_left", rsuffix="_right") + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.join(psdf2, lsuffix="_left", rsuffix="_right") + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + + self.assert_eq(join_pdf, join_psdf) + + # join with duplicated columns in Series + with self.assertRaisesRegex(ValueError, "columns overlap but no suffix specified"): + ks1 = ps.Series(["A1", "A5"], index=[1, 2], name="A") + psdf1.join(ks1, how="outer") + # join with duplicated columns in DataFrame + with self.assertRaisesRegex(ValueError, "columns overlap but no suffix specified"): + psdf1.join(psdf2, how="outer") + + # check `on` parameter + join_pdf = pdf1.join(pdf2.set_index("key"), on="key", lsuffix="_left", rsuffix="_right") + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.join(psdf2.set_index("key"), on="key", lsuffix="_left", rsuffix="_right") + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + self.assert_eq(join_pdf.reset_index(drop=True), join_psdf.reset_index(drop=True)) + + join_pdf = pdf1.set_index("key").join( + pdf2.set_index("key"), on="key", lsuffix="_left", rsuffix="_right" + ) + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.set_index("key").join( + psdf2.set_index("key"), on="key", lsuffix="_left", rsuffix="_right" + ) + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + self.assert_eq(join_pdf.reset_index(drop=True), join_psdf.reset_index(drop=True)) + + # multi-index columns + columns1 = pd.MultiIndex.from_tuples([("x", "key"), ("Y", "A")]) + columns2 = pd.MultiIndex.from_tuples([("x", "key"), ("Y", "B")]) + pdf1.columns = columns1 + pdf2.columns = columns2 + psdf1.columns = columns1 + psdf2.columns = columns2 + + join_pdf = pdf1.join(pdf2, lsuffix="_left", rsuffix="_right") + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.join(psdf2, lsuffix="_left", rsuffix="_right") + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + + self.assert_eq(join_pdf, join_psdf) + + # check `on` parameter + join_pdf = pdf1.join( + pdf2.set_index(("x", "key")), on=[("x", "key")], lsuffix="_left", rsuffix="_right" + ) + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.join( + psdf2.set_index(("x", "key")), on=[("x", "key")], lsuffix="_left", rsuffix="_right" + ) + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + + self.assert_eq(join_pdf.reset_index(drop=True), join_psdf.reset_index(drop=True)) + + join_pdf = pdf1.set_index(("x", "key")).join( + pdf2.set_index(("x", "key")), on=[("x", "key")], lsuffix="_left", rsuffix="_right" + ) + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.set_index(("x", "key")).join( + psdf2.set_index(("x", "key")), on=[("x", "key")], lsuffix="_left", rsuffix="_right" + ) + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + + self.assert_eq(join_pdf.reset_index(drop=True), join_psdf.reset_index(drop=True)) + + # multi-index + midx1 = pd.MultiIndex.from_tuples( + [("w", "a"), ("x", "b"), ("y", "c"), ("z", "d")], names=["index1", "index2"] + ) + midx2 = pd.MultiIndex.from_tuples( + [("w", "a"), ("x", "b"), ("y", "c")], names=["index1", "index2"] + ) + pdf1.index = midx1 + pdf2.index = midx2 + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + join_pdf = pdf1.join(pdf2, on=["index1", "index2"], rsuffix="_right") + join_pdf.sort_values(by=list(join_pdf.columns), inplace=True) + + join_psdf = psdf1.join(psdf2, on=["index1", "index2"], rsuffix="_right") + join_psdf.sort_values(by=list(join_psdf.columns), inplace=True) + + self.assert_eq(join_pdf, join_psdf) + + with self.assertRaisesRegex( + ValueError, r'len\(left_on\) must equal the number of levels in the index of "right"' + ): + psdf1.join(psdf2, on=["index1"], rsuffix="_right") + + def test_update(self): + # check base function + def get_data(left_columns=None, right_columns=None): + left_pdf = pd.DataFrame( + {"A": ["1", "2", "3", "4"], "B": ["100", "200", np.nan, np.nan]}, columns=["A", "B"] + ) + right_pdf = pd.DataFrame( + {"B": ["x", np.nan, "y", np.nan], "C": ["100", "200", "300", "400"]}, + columns=["B", "C"], + ) + + left_psdf = ps.DataFrame( + {"A": ["1", "2", "3", "4"], "B": ["100", "200", None, None]}, columns=["A", "B"] + ) + right_psdf = ps.DataFrame( + {"B": ["x", None, "y", None], "C": ["100", "200", "300", "400"]}, columns=["B", "C"] + ) + if left_columns is not None: + left_pdf.columns = left_columns + left_psdf.columns = left_columns + if right_columns is not None: + right_pdf.columns = right_columns + right_psdf.columns = right_columns + return left_psdf, left_pdf, right_psdf, right_pdf + + left_psdf, left_pdf, right_psdf, right_pdf = get_data() + pser = left_pdf.B + psser = left_psdf.B + left_pdf.update(right_pdf) + left_psdf.update(right_psdf) + self.assert_eq(left_pdf.sort_values(by=["A", "B"]), left_psdf.sort_values(by=["A", "B"])) + # Skip due to pandas bug: https://github.com/pandas-dev/pandas/issues/47188 + if not (LooseVersion("1.4.0") <= LooseVersion(pd.__version__) <= LooseVersion("1.4.2")): + self.assert_eq(psser.sort_index(), pser.sort_index()) + + left_psdf, left_pdf, right_psdf, right_pdf = get_data() + left_pdf.update(right_pdf, overwrite=False) + left_psdf.update(right_psdf, overwrite=False) + self.assert_eq(left_pdf.sort_values(by=["A", "B"]), left_psdf.sort_values(by=["A", "B"])) + + with self.assertRaises(NotImplementedError): + left_psdf.update(right_psdf, join="right") + + # multi-index columns + left_columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) + right_columns = pd.MultiIndex.from_tuples([("X", "B"), ("Y", "C")]) + + left_psdf, left_pdf, right_psdf, right_pdf = get_data( + left_columns=left_columns, right_columns=right_columns + ) + left_pdf.update(right_pdf) + left_psdf.update(right_psdf) + self.assert_eq( + left_pdf.sort_values(by=[("X", "A"), ("X", "B")]), + left_psdf.sort_values(by=[("X", "A"), ("X", "B")]), + ) + + left_psdf, left_pdf, right_psdf, right_pdf = get_data( + left_columns=left_columns, right_columns=right_columns + ) + left_pdf.update(right_pdf, overwrite=False) + left_psdf.update(right_psdf, overwrite=False) + self.assert_eq( + left_pdf.sort_values(by=[("X", "A"), ("X", "B")]), + left_psdf.sort_values(by=[("X", "A"), ("X", "B")]), + ) + + right_columns = pd.MultiIndex.from_tuples([("Y", "B"), ("Y", "C")]) + left_psdf, left_pdf, right_psdf, right_pdf = get_data( + left_columns=left_columns, right_columns=right_columns + ) + left_pdf.update(right_pdf) + left_psdf.update(right_psdf) + self.assert_eq( + left_pdf.sort_values(by=[("X", "A"), ("X", "B")]), + left_psdf.sort_values(by=[("X", "A"), ("X", "B")]), + ) + + +class FrameCombineTests(FrameCombineMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_combine import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_compute.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_compute.py new file mode 100644 index 0000000..5ce273c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_compute.py @@ -0,0 +1,463 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Computations / Descriptive Stats' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#computations-descriptive-stats +class FrameComputeMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_abs(self): + pdf = pd.DataFrame({"a": [-2, -1, 0, 1]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(abs(psdf), abs(pdf)) + self.assert_eq(np.abs(psdf), np.abs(pdf)) + + def test_clip(self): + pdf = pd.DataFrame( + {"A": [0, 2, 4], "B": [4, 2, 0], "X": [-1, 10, 0]}, index=np.random.rand(3) + ) + psdf = ps.from_pandas(pdf) + + # Assert list-like values are not accepted for 'lower' and 'upper' + msg = "List-like value are not supported for 'lower' and 'upper' at the moment" + with self.assertRaises(TypeError, msg=msg): + psdf.clip(lower=[1]) + with self.assertRaises(TypeError, msg=msg): + psdf.clip(upper=[1]) + + # Assert no lower or upper + self.assert_eq(psdf.clip(), pdf.clip()) + # Assert lower only + self.assert_eq(psdf.clip(1), pdf.clip(1)) + # Assert upper only + self.assert_eq(psdf.clip(upper=3), pdf.clip(upper=3)) + # Assert lower and upper + self.assert_eq(psdf.clip(1, 3), pdf.clip(1, 3)) + + pdf["clip"] = pdf.A.clip(lower=1, upper=3) + psdf["clip"] = psdf.A.clip(lower=1, upper=3) + self.assert_eq(psdf, pdf) + + # Assert behavior on string values + str_psdf = ps.DataFrame({"A": ["a", "b", "c"]}, index=np.random.rand(3)) + self.assert_eq(str_psdf.clip(1, 3), str_psdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43560): Enable DataFrameSlowTests.test_mad for pandas 2.0.0.", + ) + def test_mad(self): + pdf = pd.DataFrame( + { + "A": [1, 2, None, 4, np.nan], + "B": [-0.1, 0.2, -0.3, np.nan, 0.5], + "C": ["a", "b", "c", "d", "e"], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.mad(), pdf.mad()) + self.assert_eq(psdf.mad(axis=1), pdf.mad(axis=1)) + + with self.assertRaises(ValueError): + psdf.mad(axis=2) + + # MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "X"), ("A", "Y"), ("A", "Z")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.mad(), pdf.mad()) + self.assert_eq(psdf.mad(axis=1), pdf.mad(axis=1)) + + pdf = pd.DataFrame({"A": [True, True, False, False], "B": [True, False, False, True]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.mad(), pdf.mad()) + self.assert_eq(psdf.mad(axis=1), pdf.mad(axis=1)) + + def test_mode(self): + pdf = pd.DataFrame( + { + "A": [1, 2, None, 4, 5, 4, 2], + "B": [-0.1, 0.2, -0.3, np.nan, 0.5, -0.1, -0.1], + "C": ["d", "b", "c", "c", "e", "a", "a"], + "D": [np.nan, np.nan, np.nan, np.nan, 0.1, -0.1, -0.1], + "E": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.mode(), pdf.mode()) + self.assert_eq(psdf.mode(numeric_only=True), pdf.mode(numeric_only=True)) + self.assert_eq(psdf.mode(dropna=False), pdf.mode(dropna=False)) + + # dataframe with single column + for c in ["A", "B", "C", "D", "E"]: + self.assert_eq(psdf[[c]].mode(), pdf[[c]].mode()) + + with self.assertRaises(ValueError): + psdf.mode(axis=2) + + def f(index, iterator): + return ["3", "3", "3", "3", "4"] if index == 3 else ["0", "1", "2", "3", "4"] + + rdd = self.spark.sparkContext.parallelize( + [ + 1, + ], + 4, + ).mapPartitionsWithIndex(f) + df = self.spark.createDataFrame(rdd, schema="string") + psdf = df.pandas_api() + self.assert_eq(psdf.mode(), psdf._to_pandas().mode()) + + def test_round(self): + pdf = pd.DataFrame( + { + "A": [0.028208, 0.038683, 0.877076], + "B": [0.992815, 0.645646, 0.149370], + "C": [0.173891, 0.577595, 0.491027], + }, + columns=["A", "B", "C"], + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + + pser = pd.Series([1, 0, 2], index=["A", "B", "C"]) + psser = ps.Series([1, 0, 2], index=["A", "B", "C"]) + self.assert_eq(pdf.round(2), psdf.round(2)) + self.assert_eq(pdf.round({"A": 1, "C": 2}), psdf.round({"A": 1, "C": 2})) + self.assert_eq(pdf.round({"A": 1, "D": 2}), psdf.round({"A": 1, "D": 2})) + self.assert_eq(pdf.round(pser), psdf.round(psser)) + msg = "decimals must be an integer, a dict-like or a Series" + with self.assertRaisesRegex(TypeError, msg): + psdf.round(1.5) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + pdf.columns = columns + psdf.columns = columns + pser = pd.Series([1, 0, 2], index=columns) + psser = ps.Series([1, 0, 2], index=columns) + self.assert_eq(pdf.round(2), psdf.round(2)) + self.assert_eq( + pdf.round({("X", "A"): 1, ("Y", "C"): 2}), psdf.round({("X", "A"): 1, ("Y", "C"): 2}) + ) + self.assert_eq(pdf.round({("X", "A"): 1, "Y": 2}), psdf.round({("X", "A"): 1, "Y": 2})) + self.assert_eq(pdf.round(pser), psdf.round(psser)) + + # non-string names + pdf = pd.DataFrame( + { + 10: [0.028208, 0.038683, 0.877076], + 20: [0.992815, 0.645646, 0.149370], + 30: [0.173891, 0.577595, 0.491027], + }, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.round({10: 1, 30: 2}), psdf.round({10: 1, 30: 2})) + + def test_diff(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.diff(), psdf.diff()) + self.assert_eq(pdf.diff().diff(-1), psdf.diff().diff(-1)) + self.assert_eq(pdf.diff().sum().astype(int), psdf.diff().sum()) + + msg = "should be an int" + with self.assertRaisesRegex(TypeError, msg): + psdf.diff(1.5) + msg = 'axis should be either 0 or "index" currently.' + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.diff(axis=1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "Col1"), ("x", "Col2"), ("y", "Col3")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(pdf.diff(), psdf.diff()) + + def test_pct_change(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0], "c": [300, 200, 400, 200]}, + index=np.random.rand(4), + ) + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.pct_change(2), pdf.pct_change(2), check_exact=False) + self.assert_eq(psdf.pct_change().sum(), pdf.pct_change().sum(), check_exact=False) + + def test_rank(self): + pdf = pd.DataFrame( + data={"col1": [1, 2, 3, 1], "col2": [3, 4, 3, 1]}, + columns=["col1", "col2"], + index=np.random.rand(4), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.rank().sort_index(), psdf.rank().sort_index()) + self.assert_eq(pdf.rank().sum(), psdf.rank().sum()) + self.assert_eq( + pdf.rank(ascending=False).sort_index(), psdf.rank(ascending=False).sort_index() + ) + self.assert_eq(pdf.rank(method="min").sort_index(), psdf.rank(method="min").sort_index()) + self.assert_eq(pdf.rank(method="max").sort_index(), psdf.rank(method="max").sort_index()) + self.assert_eq( + pdf.rank(method="first").sort_index(), psdf.rank(method="first").sort_index() + ) + self.assert_eq( + pdf.rank(method="dense").sort_index(), psdf.rank(method="dense").sort_index() + ) + + msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'" + with self.assertRaisesRegex(ValueError, msg): + psdf.rank(method="nothing") + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "col1"), ("y", "col2")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.rank().sort_index(), psdf.rank().sort_index()) + + # non-numeric columns + pdf = pd.DataFrame( + data={"col1": [1, 2, 3, 1], "col2": ["a", "b", "c", "d"]}, + index=np.random.rand(4), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.rank(numeric_only=True).sort_index(), psdf.rank(numeric_only=True).sort_index() + ) + self.assert_eq( + pdf.rank(numeric_only=False).sort_index(), psdf.rank(numeric_only=False).sort_index() + ) + self.assert_eq( + pdf.rank(numeric_only=None).sort_index(), psdf.rank(numeric_only=None).sort_index() + ) + self.assert_eq( + pdf[["col2"]].rank(numeric_only=True), + psdf[["col2"]].rank(numeric_only=True), + ) + + def test_nunique(self): + pdf = pd.DataFrame({"A": [1, 2, 3], "B": [np.nan, 3, np.nan]}, index=np.random.rand(3)) + psdf = ps.from_pandas(pdf) + + # Assert NaNs are dropped by default + self.assert_eq(psdf.nunique(), pdf.nunique()) + + # Assert including NaN values + self.assert_eq(psdf.nunique(dropna=False), pdf.nunique(dropna=False)) + + # Assert approximate counts + self.assert_eq( + ps.DataFrame({"A": range(100)}).nunique(approx=True), + pd.Series([103], index=["A"]), + ) + self.assert_eq( + ps.DataFrame({"A": range(100)}).nunique(approx=True, rsd=0.01), + pd.Series([100], index=["A"]), + ) + + # Assert unsupported axis value yet + msg = 'axis should be either 0 or "index" currently.' + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.nunique(axis=1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("Y", "B")], names=["1", "2"]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.nunique(), pdf.nunique()) + self.assert_eq(psdf.nunique(dropna=False), pdf.nunique(dropna=False)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43810): Enable DataFrameSlowTests.test_quantile for pandas 2.0.0.", + ) + def test_quantile(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf.quantile(0.5), pdf.quantile(0.5)) + self.assert_eq(psdf.quantile([0.25, 0.5, 0.75]), pdf.quantile([0.25, 0.5, 0.75])) + + self.assert_eq(psdf.loc[[]].quantile(0.5), pdf.loc[[]].quantile(0.5)) + self.assert_eq( + psdf.loc[[]].quantile([0.25, 0.5, 0.75]), pdf.loc[[]].quantile([0.25, 0.5, 0.75]) + ) + + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psdf.quantile(0.5, axis=1) + with self.assertRaisesRegex(TypeError, "accuracy must be an integer; however"): + psdf.quantile(accuracy="a") + with self.assertRaisesRegex(TypeError, "q must be a float or an array of floats;"): + psdf.quantile(q="a") + with self.assertRaisesRegex(TypeError, "q must be a float or an array of floats;"): + psdf.quantile(q=["a"]) + with self.assertRaisesRegex( + ValueError, r"percentiles should all be in the interval \[0, 1\]" + ): + psdf.quantile(q=[1.1]) + + self.assert_eq( + psdf.quantile(0.5, numeric_only=False), pdf.quantile(0.5, numeric_only=False) + ) + self.assert_eq( + psdf.quantile([0.25, 0.5, 0.75], numeric_only=False), + pdf.quantile([0.25, 0.5, 0.75], numeric_only=False), + ) + + # multi-index column + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.quantile(0.5), pdf.quantile(0.5)) + self.assert_eq(psdf.quantile([0.25, 0.5, 0.75]), pdf.quantile([0.25, 0.5, 0.75])) + + pdf = pd.DataFrame({"x": ["a", "b", "c"]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.quantile(0.5), pdf.quantile(0.5)) + self.assert_eq(psdf.quantile([0.25, 0.5, 0.75]), pdf.quantile([0.25, 0.5, 0.75])) + + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + psdf.quantile(0.5, numeric_only=False) + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + psdf.quantile([0.25, 0.5, 0.75], numeric_only=False) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43558): Enable DataFrameSlowTests.test_product for pandas 2.0.0.", + ) + def test_product(self): + pdf = pd.DataFrame( + {"A": [1, 2, 3, 4, 5], "B": [10, 20, 30, 40, 50], "C": ["a", "b", "c", "d", "e"]} + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index()) + + # Named columns + pdf.columns.name = "Koalas" + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index()) + + # MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index()) + + # Named MultiIndex columns + pdf.columns.names = ["Hello", "Koalas"] + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index()) + + # No numeric columns + pdf = pd.DataFrame({"key": ["a", "b", "c"], "val": ["x", "y", "z"]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index()) + + # No numeric named columns + pdf.columns.name = "Koalas" + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), almost=True) + + # No numeric MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), almost=True) + + # No numeric named MultiIndex columns + pdf.columns.names = ["Hello", "Koalas"] + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), almost=True) + + # All NaN columns + pdf = pd.DataFrame( + { + "A": [np.nan, np.nan, np.nan, np.nan, np.nan], + "B": [10, 20, 30, 40, 50], + "C": ["a", "b", "c", "d", "e"], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), check_exact=False) + + # All NaN named columns + pdf.columns.name = "Koalas" + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), check_exact=False) + + # All NaN MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), check_exact=False) + + # All NaN named MultiIndex columns + pdf.columns.names = ["Hello", "Koalas"] + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.prod(), psdf.prod().sort_index(), check_exact=False) + + +class FrameComputeTests(FrameComputeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_compute import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_corrwith.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_corrwith.py new file mode 100644 index 0000000..cf25f39 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_corrwith.py @@ -0,0 +1,86 @@ +# +# 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 distutils.version import LooseVersion +import unittest + + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameCorrwithMixin: + def test_corrwith(self): + df1 = ps.DataFrame( + {"A": [1, np.nan, 7, 8], "B": [False, True, True, False], "C": [10, 4, 9, 3]} + ) + df2 = df1[["A", "C"]] + df3 = df1[["B", "C"]] + self._test_corrwith(df1, df2) + self._test_corrwith(df1, df3) + self._test_corrwith((df1 + 1), df2.A) + self._test_corrwith((df1 + 1), df3.B) + self._test_corrwith((df1 + 1), (df2.C + 2)) + self._test_corrwith((df1 + 1), (df3.B + 2)) + + with self.assertRaisesRegex(TypeError, "unsupported type"): + df1.corrwith(123) + with self.assertRaisesRegex(NotImplementedError, "only works for axis=0"): + df1.corrwith(df1.A, axis=1) + with self.assertRaisesRegex(ValueError, "Invalid method"): + df1.corrwith(df1.A, method="cov") + + df_bool = ps.DataFrame({"A": [True, True, False, False], "B": [True, False, False, True]}) + self._test_corrwith(df_bool, df_bool.A) + self._test_corrwith(df_bool, df_bool.B) + + def _test_corrwith(self, psdf, psobj): + pdf = psdf._to_pandas() + pobj = psobj._to_pandas() + # There was a regression in pandas 1.5.0 + # when other is Series and method is "pearson" or "spearman", and fixed in pandas 1.5.1 + # Therefore, we only test the pandas 1.5.0 in different way. + # See https://github.com/pandas-dev/pandas/issues/48826 for the reported issue, + # and https://github.com/pandas-dev/pandas/pull/46174 for the initial PR that causes. + if LooseVersion(pd.__version__) == LooseVersion("1.5.0") and isinstance(pobj, pd.Series): + methods = ["kendall"] + else: + methods = ["pearson", "spearman", "kendall"] + for method in methods: + for drop in [True, False]: + p_corr = pdf.corrwith(pobj, drop=drop, method=method) + ps_corr = psdf.corrwith(psobj, drop=drop, method=method) + self.assert_eq(p_corr.sort_index(), ps_corr.sort_index(), almost=True) + + +class FrameCorrwithTests(FrameCorrwithMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_corrwith import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cov.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cov.py new file mode 100644 index 0000000..3bbd6ab --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cov.py @@ -0,0 +1,151 @@ +# +# 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 distutils.version import LooseVersion +import unittest +import decimal + + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameCovMixin: + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43809): Enable DataFrameSlowTests.test_cov for pandas 2.0.0.", + ) + def test_cov(self): + # SPARK-36396: Implement DataFrame.cov + + # int + pdf = pd.DataFrame([(1, 2), (0, 3), (2, 0), (1, 1)], columns=["a", "b"]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.cov(), psdf.cov(), almost=True) + self.assert_eq(pdf.cov(min_periods=4), psdf.cov(min_periods=4), almost=True) + self.assert_eq(pdf.cov(min_periods=5), psdf.cov(min_periods=5)) + + # ddof + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf.cov(ddof="ddof") + for ddof in [-1, 0, 2]: + self.assert_eq(pdf.cov(ddof=ddof), psdf.cov(ddof=ddof), almost=True) + self.assert_eq( + pdf.cov(min_periods=4, ddof=ddof), psdf.cov(min_periods=4, ddof=ddof), almost=True + ) + self.assert_eq(pdf.cov(min_periods=5, ddof=ddof), psdf.cov(min_periods=5, ddof=ddof)) + + # bool + pdf = pd.DataFrame( + { + "a": [1, np.nan, 3, 4], + "b": [True, False, False, True], + "c": [True, True, False, True], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.cov(), psdf.cov(), almost=True) + self.assert_eq(pdf.cov(min_periods=4), psdf.cov(min_periods=4), almost=True) + self.assert_eq(pdf.cov(min_periods=5), psdf.cov(min_periods=5)) + + # extension dtype + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + numeric_dtypes = ["Int8", "Int16", "Int32", "Int64", "Float32", "Float64", "float"] + boolean_dtypes = ["boolean", "bool"] + else: + numeric_dtypes = ["Int8", "Int16", "Int32", "Int64", "float"] + boolean_dtypes = ["boolean", "bool"] + + sers = [pd.Series([1, 2, 3, None], dtype=dtype) for dtype in numeric_dtypes] + sers += [pd.Series([True, False, True, None], dtype=dtype) for dtype in boolean_dtypes] + sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3), None])) + + pdf = pd.concat(sers, axis=1) + pdf.columns = [dtype for dtype in numeric_dtypes + boolean_dtypes] + ["decimal"] + psdf = ps.from_pandas(pdf) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq(pdf.cov(), psdf.cov(), almost=True) + self.assert_eq(pdf.cov(min_periods=3), psdf.cov(min_periods=3), almost=True) + self.assert_eq(pdf.cov(min_periods=4), psdf.cov(min_periods=4)) + else: + test_types = [ + "Int8", + "Int16", + "Int32", + "Int64", + "float", + "boolean", + "bool", + ] + expected = pd.DataFrame( + data=[ + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0000000, 0.0000000], + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0000000, 0.0000000], + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0000000, 0.0000000], + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0000000, 0.0000000], + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0000000, 0.0000000], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.3333333, 0.3333333], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.3333333, 0.3333333], + ], + index=test_types, + columns=test_types, + ) + self.assert_eq(expected, psdf.cov(), almost=True) + + # string column + pdf = pd.DataFrame( + [(1, 2, "a", 1), (0, 3, "b", 1), (2, 0, "c", 9), (1, 1, "d", 1)], + columns=["a", "b", "c", "d"], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.cov(), psdf.cov(), almost=True) + self.assert_eq(pdf.cov(min_periods=4), psdf.cov(min_periods=4), almost=True) + self.assert_eq(pdf.cov(min_periods=5), psdf.cov(min_periods=5)) + + # nan + np.random.seed(42) + pdf = pd.DataFrame(np.random.randn(20, 3), columns=["a", "b", "c"]) + pdf.loc[pdf.index[:5], "a"] = np.nan + pdf.loc[pdf.index[5:10], "b"] = np.nan + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.cov(min_periods=11), psdf.cov(min_periods=11), almost=True) + self.assert_eq(pdf.cov(min_periods=10), psdf.cov(min_periods=10), almost=True) + + # return empty DataFrame + pdf = pd.DataFrame([("1", "2"), ("0", "3"), ("2", "0"), ("1", "1")], columns=["a", "b"]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.cov(), psdf.cov()) + + +class FrameCovTests(FrameCovMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_cov import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cumulative.py new file mode 100644 index 0000000..8c0b147 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_cumulative.py @@ -0,0 +1,139 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameCumulativeMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def _test_cummin(self, pdf, psdf): + self.assert_eq(pdf.cummin(), psdf.cummin()) + self.assert_eq(pdf.cummin(skipna=False), psdf.cummin(skipna=False)) + self.assert_eq(pdf.cummin().sum(), psdf.cummin().sum()) + + def test_cummin(self): + pdf = pd.DataFrame( + [[2.0, 1.0], [5, None], [1.0, 0.0], [2.0, 4.0], [4.0, 9.0]], + columns=list("AB"), + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + self._test_cummin(pdf, psdf) + + def test_cummin_multiindex_columns(self): + arrays = [np.array(["A", "A", "B", "B"]), np.array(["one", "two", "one", "two"])] + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "C", "B"], columns=arrays) + pdf.at["C", ("A", "two")] = None + psdf = ps.from_pandas(pdf) + self._test_cummin(pdf, psdf) + + def _test_cummax(self, pdf, psdf): + self.assert_eq(pdf.cummax(), psdf.cummax()) + self.assert_eq(pdf.cummax(skipna=False), psdf.cummax(skipna=False)) + self.assert_eq(pdf.cummax().sum(), psdf.cummax().sum()) + + def test_cummax(self): + pdf = pd.DataFrame( + [[2.0, 1.0], [5, None], [1.0, 0.0], [2.0, 4.0], [4.0, 9.0]], + columns=list("AB"), + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + self._test_cummax(pdf, psdf) + + def test_cummax_multiindex_columns(self): + arrays = [np.array(["A", "A", "B", "B"]), np.array(["one", "two", "one", "two"])] + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "C", "B"], columns=arrays) + pdf.at["C", ("A", "two")] = None + psdf = ps.from_pandas(pdf) + self._test_cummax(pdf, psdf) + + def _test_cumsum(self, pdf, psdf): + self.assert_eq(pdf.cumsum(), psdf.cumsum()) + self.assert_eq(pdf.cumsum(skipna=False), psdf.cumsum(skipna=False)) + self.assert_eq(pdf.cumsum().sum(), psdf.cumsum().sum()) + + def test_cumsum(self): + pdf = pd.DataFrame( + [[2.0, 1.0], [5, None], [1.0, 0.0], [2.0, 4.0], [4.0, 9.0]], + columns=list("AB"), + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + self._test_cumsum(pdf, psdf) + + def test_cumsum_multiindex_columns(self): + arrays = [np.array(["A", "A", "B", "B"]), np.array(["one", "two", "one", "two"])] + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "C", "B"], columns=arrays) + pdf.at["C", ("A", "two")] = None + psdf = ps.from_pandas(pdf) + self._test_cumsum(pdf, psdf) + + def _test_cumprod(self, pdf, psdf): + self.assert_eq(pdf.cumprod(), psdf.cumprod(), almost=True) + self.assert_eq(pdf.cumprod(skipna=False), psdf.cumprod(skipna=False), almost=True) + self.assert_eq(pdf.cumprod().sum(), psdf.cumprod().sum(), almost=True) + + def test_cumprod(self): + pdf = pd.DataFrame( + [[2.0, 1.0, 1], [5, None, 2], [1.0, -1.0, -3], [2.0, 0, 4], [4.0, 9.0, 5]], + columns=list("ABC"), + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + self._test_cumprod(pdf, psdf) + + def test_cumprod_multiindex_columns(self): + arrays = [np.array(["A", "A", "B", "B"]), np.array(["one", "two", "one", "two"])] + pdf = pd.DataFrame(np.random.rand(3, 4), index=["A", "C", "B"], columns=arrays) + pdf.at["C", ("A", "two")] = None + psdf = ps.from_pandas(pdf) + self._test_cumprod(pdf, psdf) + + +class FrameCumulativeTests(FrameCumulativeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_cumulative import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_describe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_describe.py new file mode 100644 index 0000000..af98d28 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_describe.py @@ -0,0 +1,458 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameDescribeMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43556): Enable DataFrameSlowTests.test_describe for pandas 2.0.0.", + ) + def test_describe(self): + pdf, psdf = self.df_pair + + # numeric columns + self.assert_eq(psdf.describe(), pdf.describe()) + psdf.a += psdf.a + pdf.a += pdf.a + self.assert_eq(psdf.describe(), pdf.describe()) + + # string columns + psdf = ps.DataFrame({"A": ["a", "b", "b", "c"], "B": ["d", "e", "f", "f"]}) + pdf = psdf._to_pandas() + self.assert_eq(psdf.describe(), pdf.describe().astype(str)) + psdf.A += psdf.A + pdf.A += pdf.A + self.assert_eq(psdf.describe(), pdf.describe().astype(str)) + + # timestamp columns + psdf = ps.DataFrame( + { + "A": [ + pd.Timestamp("2020-10-20"), + pd.Timestamp("2021-06-02"), + pd.Timestamp("2021-06-02"), + pd.Timestamp("2022-07-11"), + ], + "B": [ + pd.Timestamp("2021-11-20"), + pd.Timestamp("2023-06-02"), + pd.Timestamp("2026-07-11"), + pd.Timestamp("2026-07-11"), + ], + } + ) + pdf = psdf._to_pandas() + # NOTE: Set `datetime_is_numeric=True` for pandas: + # FutureWarning: Treating datetime data as categorical rather than numeric in + # `.describe` is deprecated and will be removed in a future version of pandas. + # Specify `datetime_is_numeric=True` to silence this + # warning and adopt the future behavior now. + # NOTE: Compare the result except percentiles, since we use approximate percentile + # so the result is different from pandas. + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pdf.describe(datetime_is_numeric=True) + .astype(str) + .loc[["count", "mean", "min", "max"]], + ) + else: + self.assert_eq( + psdf.describe(), + ps.DataFrame( + { + "A": [ + "4", + "2021-07-16 18:00:00", + "2020-10-20 00:00:00", + "2020-10-20 00:00:00", + "2021-06-02 00:00:00", + "2021-06-02 00:00:00", + "2022-07-11 00:00:00", + ], + "B": [ + "4", + "2024-08-02 18:00:00", + "2021-11-20 00:00:00", + "2021-11-20 00:00:00", + "2023-06-02 00:00:00", + "2026-07-11 00:00:00", + "2026-07-11 00:00:00", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max"], + ), + ) + + # String & timestamp columns + psdf = ps.DataFrame( + { + "A": ["a", "b", "b", "c"], + "B": [ + pd.Timestamp("2021-11-20"), + pd.Timestamp("2023-06-02"), + pd.Timestamp("2026-07-11"), + pd.Timestamp("2026-07-11"), + ], + } + ) + pdf = psdf._to_pandas() + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pdf.describe(datetime_is_numeric=True) + .astype(str) + .loc[["count", "mean", "min", "max"]], + ) + psdf.A += psdf.A + pdf.A += pdf.A + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pdf.describe(datetime_is_numeric=True) + .astype(str) + .loc[["count", "mean", "min", "max"]], + ) + else: + expected_result = ps.DataFrame( + { + "B": [ + "4", + "2024-08-02 18:00:00", + "2021-11-20 00:00:00", + "2021-11-20 00:00:00", + "2023-06-02 00:00:00", + "2026-07-11 00:00:00", + "2026-07-11 00:00:00", + ] + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max"], + ) + self.assert_eq( + psdf.describe(), + expected_result, + ) + psdf.A += psdf.A + self.assert_eq( + psdf.describe(), + expected_result, + ) + + # Numeric & timestamp columns + psdf = ps.DataFrame( + { + "A": [1, 2, 2, 3], + "B": [ + pd.Timestamp("2021-11-20"), + pd.Timestamp("2023-06-02"), + pd.Timestamp("2026-07-11"), + pd.Timestamp("2026-07-11"), + ], + } + ) + pdf = psdf._to_pandas() + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + pandas_result = pdf.describe(datetime_is_numeric=True) + pandas_result.B = pandas_result.B.astype(str) + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pandas_result.loc[["count", "mean", "min", "max"]], + ) + psdf.A += psdf.A + pdf.A += pdf.A + pandas_result = pdf.describe(datetime_is_numeric=True) + pandas_result.B = pandas_result.B.astype(str) + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pandas_result.loc[["count", "mean", "min", "max"]], + ) + else: + self.assert_eq( + psdf.describe(), + ps.DataFrame( + { + "A": [4, 2, 1, 1, 2, 2, 3, 0.816497], + "B": [ + "4", + "2024-08-02 18:00:00", + "2021-11-20 00:00:00", + "2021-11-20 00:00:00", + "2023-06-02 00:00:00", + "2026-07-11 00:00:00", + "2026-07-11 00:00:00", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max", "std"], + ), + ) + psdf.A += psdf.A + self.assert_eq( + psdf.describe(), + ps.DataFrame( + { + "A": [4, 4, 2, 2, 4, 4, 6, 1.632993], + "B": [ + "4", + "2024-08-02 18:00:00", + "2021-11-20 00:00:00", + "2021-11-20 00:00:00", + "2023-06-02 00:00:00", + "2026-07-11 00:00:00", + "2026-07-11 00:00:00", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max", "std"], + ), + ) + + # Include None column + psdf = ps.DataFrame( + { + "a": [1, 2, 3], + "b": [pd.Timestamp(1), pd.Timestamp(1), pd.Timestamp(1)], + "c": [None, None, None], + } + ) + pdf = psdf._to_pandas() + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + pandas_result = pdf.describe(datetime_is_numeric=True) + pandas_result.b = pandas_result.b.astype(str) + self.assert_eq( + psdf.describe().loc[["count", "mean", "min", "max"]], + pandas_result.loc[["count", "mean", "min", "max"]], + ) + else: + self.assert_eq( + psdf.describe(), + ps.DataFrame( + { + "a": [3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 3.0, 1.0], + "b": [ + "3", + "1970-01-01 00:00:00.000001", + "1970-01-01 00:00:00.000001", + "1970-01-01 00:00:00.000001", + "1970-01-01 00:00:00.000001", + "1970-01-01 00:00:00.000001", + "1970-01-01 00:00:00.000001", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max", "std"], + ), + ) + + msg = r"Percentiles should all be in the interval \[0, 1\]" + with self.assertRaisesRegex(ValueError, msg): + psdf.describe(percentiles=[1.1]) + + psdf = ps.DataFrame() + msg = "Cannot describe a DataFrame without columns" + with self.assertRaisesRegex(ValueError, msg): + psdf.describe() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43556): Enable DataFrameSlowTests.test_describe for pandas 2.0.0.", + ) + def test_describe_empty(self): + # Empty DataFrame + psdf = ps.DataFrame(columns=["A", "B"]) + pdf = psdf._to_pandas() + self.assert_eq( + psdf.describe(), + pdf.describe().astype(float), + ) + + # Explicit empty DataFrame numeric only + psdf = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + pdf = psdf._to_pandas() + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf[pdf.a != pdf.a].describe(), + ) + + # Explicit empty DataFrame string only + psdf = ps.DataFrame({"a": ["a", "b", "c"], "b": ["q", "w", "e"]}) + pdf = psdf._to_pandas() + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf[pdf.a != pdf.a].describe().astype(float), + ) + + # Explicit empty DataFrame timestamp only + psdf = ps.DataFrame( + { + "a": [pd.Timestamp(1), pd.Timestamp(1), pd.Timestamp(1)], + "b": [pd.Timestamp(1), pd.Timestamp(1), pd.Timestamp(1)], + } + ) + pdf = psdf._to_pandas() + # For timestamp type, we should convert NaT to None in pandas result + # since pandas API on Spark doesn't support the NaT for object type. + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + pdf_result = pdf[pdf.a != pdf.a].describe(datetime_is_numeric=True) + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf_result.where(pdf_result.notnull(), None).astype(str), + ) + else: + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + ps.DataFrame( + { + "a": [ + "0", + "None", + "None", + "None", + "None", + "None", + "None", + ], + "b": [ + "0", + "None", + "None", + "None", + "None", + "None", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max"], + ), + ) + + # Explicit empty DataFrame numeric & timestamp + psdf = ps.DataFrame( + {"a": [1, 2, 3], "b": [pd.Timestamp(1), pd.Timestamp(1), pd.Timestamp(1)]} + ) + pdf = psdf._to_pandas() + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + pdf_result = pdf[pdf.a != pdf.a].describe(datetime_is_numeric=True) + pdf_result.b = pdf_result.b.where(pdf_result.b.notnull(), None).astype(str) + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf_result, + ) + else: + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + ps.DataFrame( + { + "a": [ + 0, + None, + None, + None, + None, + None, + None, + None, + ], + "b": [ + "0", + "None", + "None", + "None", + "None", + "None", + "None", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max", "std"], + ), + ) + + # Explicit empty DataFrame numeric & string + psdf = ps.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + pdf = psdf._to_pandas() + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf[pdf.a != pdf.a].describe(), + ) + + # Explicit empty DataFrame string & timestamp + psdf = ps.DataFrame( + {"a": ["a", "b", "c"], "b": [pd.Timestamp(1), pd.Timestamp(1), pd.Timestamp(1)]} + ) + pdf = psdf._to_pandas() + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + pdf_result = pdf[pdf.a != pdf.a].describe(datetime_is_numeric=True) + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + pdf_result.where(pdf_result.notnull(), None).astype(str), + ) + else: + self.assert_eq( + psdf[psdf.a != psdf.a].describe(), + ps.DataFrame( + { + "b": [ + "0", + "None", + "None", + "None", + "None", + "None", + "None", + ], + }, + index=["count", "mean", "min", "25%", "50%", "75%", "max"], + ), + ) + + +class FrameDescribeTests(FrameDescribeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_describe import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_eval.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_eval.py new file mode 100644 index 0000000..8101ed3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_eval.py @@ -0,0 +1,86 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameEvalMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + def test_eval(self): + pdf = pd.DataFrame({"A": range(1, 6), "B": range(10, 0, -2)}) + psdf = ps.from_pandas(pdf) + + # operation between columns (returns Series) + self.assert_eq(pdf.eval("A + B"), psdf.eval("A + B")) + self.assert_eq(pdf.eval("A + A"), psdf.eval("A + A")) + # assignment (returns DataFrame) + self.assert_eq(pdf.eval("C = A + B"), psdf.eval("C = A + B")) + self.assert_eq(pdf.eval("A = A + A"), psdf.eval("A = A + A")) + # operation between scalars (returns scalar) + self.assert_eq(pdf.eval("1 + 1"), psdf.eval("1 + 1")) + # complicated operations with assignment + self.assert_eq( + pdf.eval("B = A + B // (100 + 200) * (500 - B) - 10.5"), + psdf.eval("B = A + B // (100 + 200) * (500 - B) - 10.5"), + ) + + # inplace=True (only support for assignment) + pdf.eval("C = A + B", inplace=True) + psdf.eval("C = A + B", inplace=True) + self.assert_eq(pdf, psdf) + pser = pdf.A + psser = psdf.A + pdf.eval("A = B + C", inplace=True) + psdf.eval("A = B + C", inplace=True) + self.assert_eq(pdf, psdf) + # Skip due to pandas bug: https://github.com/pandas-dev/pandas/issues/47449 + if not (LooseVersion("1.4.0") <= LooseVersion(pd.__version__) <= LooseVersion("1.4.3")): + self.assert_eq(pser, psser) + + # doesn't support for multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b"), ("z", "c")]) + psdf.columns = columns + self.assertRaises(TypeError, lambda: psdf.eval("x.a + y.b")) + + +class FrameEvalTests(FrameEvalMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_eval import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_melt.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_melt.py new file mode 100644 index 0000000..c8969a3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_melt.py @@ -0,0 +1,202 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.pandas.utils import name_like_string + + +class FrameMeltMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_melt(self): + pdf = pd.DataFrame( + {"A": [1, 3, 5], "B": [2, 4, 6], "C": [7, 8, 9]}, index=np.random.rand(3) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.melt().sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt().sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars="A").sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt(id_vars="A").sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=["A", "B"]).sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt(id_vars=["A", "B"]).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=("A", "B")).sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt(id_vars=("A", "B")).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=["A"], value_vars=["C"]) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=["A"], value_vars=["C"]).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=["A"], value_vars=["B"], var_name="myVarname", value_name="myValname") + .sort_values(["myVarname", "myValname"]) + .reset_index(drop=True), + pdf.melt( + id_vars=["A"], value_vars=["B"], var_name="myVarname", value_name="myValname" + ).sort_values(["myVarname", "myValname"]), + ) + self.assert_eq( + psdf.melt(value_vars=("A", "B")) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(value_vars=("A", "B")).sort_values(["variable", "value"]), + ) + + self.assertRaises(KeyError, lambda: psdf.melt(id_vars="Z")) + self.assertRaises(KeyError, lambda: psdf.melt(value_vars="Z")) + + # multi-index columns + TEN = 10.0 + TWELVE = 20.0 + + columns = pd.MultiIndex.from_tuples([(TEN, "A"), (TEN, "B"), (TWELVE, "C")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.melt().sort_values(["variable_0", "variable_1", "value"]).reset_index(drop=True), + pdf.melt().sort_values(["variable_0", "variable_1", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=[(TEN, "A")]) + .sort_values(["variable_0", "variable_1", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=[(TEN, "A")]) + .sort_values(["variable_0", "variable_1", "value"]) + .rename(columns=name_like_string), + ) + self.assert_eq( + psdf.melt(id_vars=[(TEN, "A")], value_vars=[(TWELVE, "C")]) + .sort_values(["variable_0", "variable_1", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=[(TEN, "A")], value_vars=[(TWELVE, "C")]) + .sort_values(["variable_0", "variable_1", "value"]) + .rename(columns=name_like_string), + ) + self.assert_eq( + psdf.melt( + id_vars=[(TEN, "A")], + value_vars=[(TEN, "B")], + var_name=["myV1", "myV2"], + value_name="myValname", + ) + .sort_values(["myV1", "myV2", "myValname"]) + .reset_index(drop=True), + pdf.melt( + id_vars=[(TEN, "A")], + value_vars=[(TEN, "B")], + var_name=["myV1", "myV2"], + value_name="myValname", + ) + .sort_values(["myV1", "myV2", "myValname"]) + .rename(columns=name_like_string), + ) + + columns.names = ["v0", "v1"] + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.melt().sort_values(["v0", "v1", "value"]).reset_index(drop=True), + pdf.melt().sort_values(["v0", "v1", "value"]), + ) + + self.assertRaises(ValueError, lambda: psdf.melt(id_vars=(TEN, "A"))) + self.assertRaises(ValueError, lambda: psdf.melt(value_vars=(TEN, "A"))) + self.assertRaises(KeyError, lambda: psdf.melt(id_vars=[TEN])) + self.assertRaises(KeyError, lambda: psdf.melt(id_vars=[(TWELVE, "A")])) + self.assertRaises(KeyError, lambda: psdf.melt(value_vars=[TWELVE])) + self.assertRaises(KeyError, lambda: psdf.melt(value_vars=[(TWELVE, "A")])) + + # non-string names + pdf.columns = [10.0, 20.0, 30.0] + psdf.columns = [10.0, 20.0, 30.0] + + self.assert_eq( + psdf.melt().sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt().sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=10.0).sort_values(["variable", "value"]).reset_index(drop=True), + pdf.melt(id_vars=10.0).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=[10.0, 20.0]) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=[10.0, 20.0]).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=(10.0, 20.0)) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=(10.0, 20.0)).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(id_vars=[10.0], value_vars=[30.0]) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(id_vars=[10.0], value_vars=[30.0]).sort_values(["variable", "value"]), + ) + self.assert_eq( + psdf.melt(value_vars=(10.0, 20.0)) + .sort_values(["variable", "value"]) + .reset_index(drop=True), + pdf.melt(value_vars=(10.0, 20.0)).sort_values(["variable", "value"]), + ) + + +class FrameMeltTests(FrameMeltMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_melt import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_missing_data.py new file mode 100644 index 0000000..71775e2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_missing_data.py @@ -0,0 +1,520 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Missing data handling' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#missing-data-handling +class FrameMissingDataMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_backfill(self): + pdf = pd.DataFrame( + { + "A": [None, 3, None, None], + "B": [2, 4, None, 3], + "C": [None, None, None, 1], + "D": [0, 1, 5, 4], + }, + columns=["A", "B", "C", "D"], + ) + psdf = ps.from_pandas(pdf) + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(pdf.backfill(), psdf.backfill()) + + # Test `inplace=True` + pdf.backfill(inplace=True) + psdf.backfill(inplace=True) + self.assert_eq(pdf, psdf) + else: + expected = ps.DataFrame( + { + "A": [3.0, 3.0, None, None], + "B": [2.0, 4.0, 3.0, 3.0], + "C": [1.0, 1.0, 1.0, 1.0], + "D": [0, 1, 5, 4], + }, + columns=["A", "B", "C", "D"], + ) + self.assert_eq(expected, psdf.backfill()) + + # Test `inplace=True` + psdf.backfill(inplace=True) + self.assert_eq(expected, psdf) + + def _test_dropna(self, pdf, axis): + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.dropna(axis=axis), pdf.dropna(axis=axis)) + self.assert_eq(psdf.dropna(axis=axis, how="all"), pdf.dropna(axis=axis, how="all")) + self.assert_eq(psdf.dropna(axis=axis, subset=["x"]), pdf.dropna(axis=axis, subset=["x"])) + self.assert_eq(psdf.dropna(axis=axis, subset="x"), pdf.dropna(axis=axis, subset=["x"])) + self.assert_eq( + psdf.dropna(axis=axis, subset=["y", "z"]), pdf.dropna(axis=axis, subset=["y", "z"]) + ) + self.assert_eq( + psdf.dropna(axis=axis, subset=["y", "z"], how="all"), + pdf.dropna(axis=axis, subset=["y", "z"], how="all"), + ) + + self.assert_eq(psdf.dropna(axis=axis, thresh=2), pdf.dropna(axis=axis, thresh=2)) + self.assert_eq( + psdf.dropna(axis=axis, thresh=1, subset=["y", "z"]), + pdf.dropna(axis=axis, thresh=1, subset=["y", "z"]), + ) + + pdf2 = pdf.copy() + psdf2 = psdf.copy() + pser = pdf2[pdf2.columns[0]] + psser = psdf2[psdf2.columns[0]] + pdf2.dropna(inplace=True, axis=axis) + psdf2.dropna(inplace=True, axis=axis) + self.assert_eq(psdf2, pdf2) + self.assert_eq(psser, pser) + + # multi-index + columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")]) + if axis == 0: + pdf.columns = columns + else: + pdf.index = columns + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.dropna(axis=axis), pdf.dropna(axis=axis)) + self.assert_eq(psdf.dropna(axis=axis, how="all"), pdf.dropna(axis=axis, how="all")) + self.assert_eq( + psdf.dropna(axis=axis, subset=[("a", "x")]), pdf.dropna(axis=axis, subset=[("a", "x")]) + ) + self.assert_eq( + psdf.dropna(axis=axis, subset=("a", "x")), pdf.dropna(axis=axis, subset=[("a", "x")]) + ) + self.assert_eq( + psdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")]), + pdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")]), + ) + self.assert_eq( + psdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")], how="all"), + pdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")], how="all"), + ) + + self.assert_eq(psdf.dropna(axis=axis, thresh=2), pdf.dropna(axis=axis, thresh=2)) + self.assert_eq( + psdf.dropna(axis=axis, thresh=1, subset=[("a", "y"), ("b", "z")]), + pdf.dropna(axis=axis, thresh=1, subset=[("a", "y"), ("b", "z")]), + ) + + def test_dropna_axis_index(self): + pdf = pd.DataFrame( + { + "x": [np.nan, 2, 3, 4, np.nan, 6], + "y": [1, 2, np.nan, 4, np.nan, np.nan], + "z": [1, 2, 3, 4, np.nan, np.nan], + }, + index=np.random.rand(6), + ) + + self._test_dropna(pdf, axis=0) + + # empty + pdf = pd.DataFrame(index=np.random.rand(6)) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.dropna(), pdf.dropna()) + self.assert_eq(psdf.dropna(how="all"), pdf.dropna(how="all")) + self.assert_eq(psdf.dropna(thresh=0), pdf.dropna(thresh=0)) + self.assert_eq(psdf.dropna(thresh=1), pdf.dropna(thresh=1)) + + # Only NA value + pdf["a"] = [np.nan] * 6 + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.dropna(), pdf.dropna()) + self.assert_eq(psdf.dropna(how="all"), pdf.dropna(how="all")) + self.assert_eq(psdf.dropna(thresh=0), pdf.dropna(thresh=0)) + self.assert_eq(psdf.dropna(thresh=1), pdf.dropna(thresh=1)) + + with self.assertRaisesRegex(ValueError, "No axis named foo"): + psdf.dropna(axis="foo") + + self.assertRaises(KeyError, lambda: psdf.dropna(subset="1")) + with self.assertRaisesRegex(ValueError, "invalid how option: 1"): + psdf.dropna(how=1) + with self.assertRaisesRegex(TypeError, "must specify how or thresh"): + psdf.dropna(how=None) + + def test_dropna_axis_column(self): + pdf = pd.DataFrame( + { + "x": [np.nan, 2, 3, 4, np.nan, 6], + "y": [1, 2, np.nan, 4, np.nan, np.nan], + "z": [1, 2, 3, 4, np.nan, np.nan], + }, + index=[str(r) for r in np.random.rand(6)], + ).T + + self._test_dropna(pdf, axis=1) + + psdf = ps.from_pandas(pdf) + with self.assertRaisesRegex( + ValueError, "The length of each subset must be the same as the index size." + ): + psdf.dropna(subset=(["x", "y"]), axis=1) + + # empty + pdf = pd.DataFrame({"x": [], "y": [], "z": []}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.dropna(axis=1), pdf.dropna(axis=1)) + self.assert_eq(psdf.dropna(axis=1, how="all"), pdf.dropna(axis=1, how="all")) + self.assert_eq(psdf.dropna(axis=1, thresh=0), pdf.dropna(axis=1, thresh=0)) + self.assert_eq(psdf.dropna(axis=1, thresh=1), pdf.dropna(axis=1, thresh=1)) + + def test_fillna(self): + pdf = pd.DataFrame( + { + "x": [np.nan, 2, 3, 4, np.nan, 6], + "y": [1, 2, np.nan, 4, np.nan, np.nan], + "z": [1, 2, 3, 4, np.nan, np.nan], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) + self.assert_eq( + psdf.fillna({"x": -1, "y": -2, "z": -5}), pdf.fillna({"x": -1, "y": -2, "z": -5}) + ) + self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) + self.assert_eq(pdf.fillna(method="ffill", limit=2), psdf.fillna(method="ffill", limit=2)) + self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) + self.assert_eq(pdf.fillna(method="bfill", limit=2), psdf.fillna(method="bfill", limit=2)) + + pdf = pdf.set_index(["x", "y"]) + psdf = ps.from_pandas(pdf) + # check multi index + self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) + self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) + self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) + + pser = pdf.z + psser = psdf.z + pdf.fillna({"x": -1, "y": -2, "z": -5}, inplace=True) + psdf.fillna({"x": -1, "y": -2, "z": -5}, inplace=True) + self.assert_eq(psdf, pdf) + # Skip due to pandas bug: https://github.com/pandas-dev/pandas/issues/47188 + if not (LooseVersion("1.4.0") <= LooseVersion(pd.__version__) <= LooseVersion("1.4.2")): + self.assert_eq(psser, pser) + + pser = pdf.z + psser = psdf.z + pdf.fillna(0, inplace=True) + psdf.fillna(0, inplace=True) + self.assert_eq(psdf, pdf) + self.assert_eq(psser, pser) + + s_nan = pd.Series([-1, -2, -5], index=["x", "y", "z"], dtype=int) + self.assert_eq(psdf.fillna(s_nan), pdf.fillna(s_nan)) + + with self.assertRaisesRegex(NotImplementedError, "fillna currently only"): + psdf.fillna(-1, axis=1) + with self.assertRaisesRegex(NotImplementedError, "fillna currently only"): + psdf.fillna(-1, axis="columns") + with self.assertRaisesRegex(ValueError, "limit parameter for value is not support now"): + psdf.fillna(-1, limit=1) + with self.assertRaisesRegex(TypeError, "Unsupported.*DataFrame"): + psdf.fillna(pd.DataFrame({"x": [-1], "y": [-1], "z": [-1]})) + with self.assertRaisesRegex(TypeError, "Unsupported.*int64"): + psdf.fillna({"x": np.int64(-6), "y": np.int64(-4), "z": -5}) + with self.assertRaisesRegex(ValueError, "Expecting 'pad', 'ffill', 'backfill' or 'bfill'."): + psdf.fillna(method="xxx") + with self.assertRaisesRegex( + ValueError, "Must specify a fillna 'value' or 'method' parameter." + ): + psdf.fillna() + + # multi-index columns + pdf = pd.DataFrame( + { + ("x", "a"): [np.nan, 2, 3, 4, np.nan, 6], + ("x", "b"): [1, 2, np.nan, 4, np.nan, np.nan], + ("y", "c"): [1, 2, 3, 4, np.nan, np.nan], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) + self.assert_eq( + psdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), + pdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), + ) + self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) + self.assert_eq(pdf.fillna(method="ffill", limit=2), psdf.fillna(method="ffill", limit=2)) + self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) + self.assert_eq(pdf.fillna(method="bfill", limit=2), psdf.fillna(method="bfill", limit=2)) + + # See also: https://github.com/pandas-dev/pandas/issues/47649 + if LooseVersion("1.4.3") != LooseVersion(pd.__version__): + self.assert_eq(psdf.fillna({"x": -1}), pdf.fillna({"x": -1})) + self.assert_eq( + psdf.fillna({"x": -1, ("x", "b"): -2}), pdf.fillna({"x": -1, ("x", "b"): -2}) + ) + self.assert_eq( + psdf.fillna({("x", "b"): -2, "x": -1}), pdf.fillna({("x", "b"): -2, "x": -1}) + ) + + # check multi index + pdf = pdf.set_index([("x", "a"), ("x", "b")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) + self.assert_eq( + psdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), + pdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), + ) + + def test_replace(self): + pdf = pd.DataFrame( + { + "name": ["Ironman", "Captain America", "Thor", "Hulk"], + "weapon": ["Mark-45", "Shield", "Mjolnir", "Smash"], + }, + index=np.random.rand(4), + ) + psdf = ps.from_pandas(pdf) + + with self.assertRaisesRegex( + NotImplementedError, "replace currently works only for method='pad" + ): + psdf.replace(method="bfill") + with self.assertRaisesRegex( + NotImplementedError, "replace currently works only when limit=None" + ): + psdf.replace(limit=10) + with self.assertRaisesRegex( + NotImplementedError, "replace currently doesn't supports regex" + ): + psdf.replace(regex="") + + with self.assertRaisesRegex(ValueError, "Length of to_replace and value must be same"): + psdf.replace(to_replace=["Ironman"], value=["Spiderman", "Doctor Strange"]) + with self.assertRaisesRegex(TypeError, "Unsupported type function"): + psdf.replace("Ironman", lambda x: "Spiderman") + with self.assertRaisesRegex(TypeError, "Unsupported type function"): + psdf.replace(lambda x: "Ironman", "Spiderman") + + self.assert_eq(psdf.replace("Ironman", "Spiderman"), pdf.replace("Ironman", "Spiderman")) + self.assert_eq( + psdf.replace(["Ironman", "Captain America"], ["Rescue", "Hawkeye"]), + pdf.replace(["Ironman", "Captain America"], ["Rescue", "Hawkeye"]), + ) + self.assert_eq( + psdf.replace(("Ironman", "Captain America"), ("Rescue", "Hawkeye")), + pdf.replace(("Ironman", "Captain America"), ("Rescue", "Hawkeye")), + ) + + # inplace + pser = pdf.name + psser = psdf.name + pdf.replace("Ironman", "Spiderman", inplace=True) + psdf.replace("Ironman", "Spiderman", inplace=True) + self.assert_eq(psdf, pdf) + self.assert_eq(psser, pser) + + pdf = pd.DataFrame( + {"A": [0, 1, 2, 3, np.nan], "B": [5, 6, 7, 8, np.nan], "C": ["a", "b", "c", "d", None]}, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.replace([0, 1, 2, 3, 5, 6], 4), pdf.replace([0, 1, 2, 3, 5, 6], 4)) + + self.assert_eq( + psdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), + pdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), + ) + + self.assert_eq(psdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200})) + + self.assert_eq( + psdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100), + pdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100), + ) + + self.assert_eq( + psdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}), + pdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}), + ) + self.assert_eq( + psdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}), + pdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}), + ) + + self.assert_eq(psdf.replace({"C": ["a", None]}, "e"), pdf.replace({"C": ["a", None]}, "e")) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.replace([0, 1, 2, 3, 5, 6], 4), pdf.replace([0, 1, 2, 3, 5, 6], 4)) + + self.assert_eq( + psdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), + pdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), + ) + + self.assert_eq(psdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200})) + + self.assert_eq( + psdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100), + pdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100), + ) + + self.assert_eq( + psdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}), + pdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}), + ) + self.assert_eq( + psdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}), + pdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}), + ) + + self.assert_eq( + psdf.replace({("Y", "C"): ["a", None]}, "e"), + pdf.replace({("Y", "C"): ["a", None]}, "e"), + ) + + def test_ffill(self): + idx = np.random.rand(6) + pdf = pd.DataFrame( + { + "x": [np.nan, 2, 3, 4, np.nan, 6], + "y": [1, 2, np.nan, 4, np.nan, np.nan], + "z": [1, 2, 3, 4, np.nan, np.nan], + }, + index=idx, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.ffill(), pdf.ffill()) + self.assert_eq(psdf.ffill(limit=1), pdf.ffill(limit=1)) + + pser = pdf.y + psser = psdf.y + + psdf.ffill(inplace=True) + pdf.ffill(inplace=True) + + self.assert_eq(psdf, pdf) + self.assert_eq(psser, pser) + self.assert_eq(psser[idx[2]], pser[idx[2]]) + + def test_bfill(self): + idx = np.random.rand(6) + pdf = pd.DataFrame( + { + "x": [np.nan, 2, 3, 4, np.nan, 6], + "y": [1, 2, np.nan, 4, np.nan, np.nan], + "z": [1, 2, 3, 4, np.nan, np.nan], + }, + index=idx, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.bfill(), pdf.bfill()) + self.assert_eq(psdf.bfill(limit=1), pdf.bfill(limit=1)) + + pser = pdf.x + psser = psdf.x + + psdf.bfill(inplace=True) + pdf.bfill(inplace=True) + + self.assert_eq(psdf, pdf) + self.assert_eq(psser, pser) + self.assert_eq(psser[idx[0]], pser[idx[0]]) + + def test_pad(self): + pdf = pd.DataFrame( + { + "A": [None, 3, None, None], + "B": [2, 4, None, 3], + "C": [None, None, None, 1], + "D": [0, 1, 5, 4], + }, + columns=["A", "B", "C", "D"], + ) + psdf = ps.from_pandas(pdf) + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(pdf.pad(), psdf.pad()) + + # Test `inplace=True` + pdf.pad(inplace=True) + psdf.pad(inplace=True) + self.assert_eq(pdf, psdf) + else: + expected = ps.DataFrame( + { + "A": [None, 3, 3, 3], + "B": [2.0, 4.0, 4.0, 3.0], + "C": [None, None, None, 1], + "D": [0, 1, 5, 4], + }, + columns=["A", "B", "C", "D"], + ) + self.assert_eq(expected, psdf.pad()) + + # Test `inplace=True` + psdf.pad(inplace=True) + self.assert_eq(expected, psdf) + + +class FrameMissingDataTests(FrameMissingDataMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_missing_data import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_pivot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_pivot.py new file mode 100644 index 0000000..e1c55cc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/computation/test_pivot.py @@ -0,0 +1,361 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FramePivotMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_pivot_table_dtypes(self): + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, 2, 2, 4, 2, 4], + "e": [1, 2, 2, 4, 2, 4], + "c": [1, 2, 9, 4, 7, 4], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + # Skip columns comparison by reset_index + res_df = psdf.pivot_table( + index=["c"], columns="a", values=["b"], aggfunc={"b": "mean"} + ).dtypes.reset_index(drop=True) + exp_df = pdf.pivot_table( + index=["c"], columns="a", values=["b"], aggfunc={"b": "mean"} + ).dtypes.reset_index(drop=True) + self.assert_eq(res_df, exp_df) + + # Results don't have the same column's name + + # Todo: self.assert_eq(psdf.pivot_table(columns="a", values="b").dtypes, + # pdf.pivot_table(columns="a", values="b").dtypes) + + # Todo: self.assert_eq(psdf.pivot_table(index=['c'], columns="a", values="b").dtypes, + # pdf.pivot_table(index=['c'], columns="a", values="b").dtypes) + + # Todo: self.assert_eq(psdf.pivot_table(index=['e', 'c'], columns="a", values="b").dtypes, + # pdf.pivot_table(index=['e', 'c'], columns="a", values="b").dtypes) + + # Todo: self.assert_eq(psdf.pivot_table(index=['e', 'c'], + # columns="a", values="b", fill_value=999).dtypes, pdf.pivot_table(index=['e', 'c'], + # columns="a", values="b", fill_value=999).dtypes) + + def test_pivot_table(self): + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, 2, 2, 4, 2, 4], + "e": [10, 20, 20, 40, 20, 40], + "c": [1, 2, 9, 4, 7, 4], + "d": [-1, -2, -3, -4, -5, -6], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + # Checking if both DataFrames have the same results + self.assert_eq( + psdf.pivot_table(columns="a", values="b").sort_index(), + pdf.pivot_table(columns="a", values="b").sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table(index=["c"], columns="a", values="b").sort_index(), + pdf.pivot_table(index=["c"], columns="a", values="b").sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table(index=["c"], columns="a", values="b", aggfunc="sum").sort_index(), + pdf.pivot_table(index=["c"], columns="a", values="b", aggfunc="sum").sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table(index=["c"], columns="a", values=["b"], aggfunc="sum").sort_index(), + pdf.pivot_table(index=["c"], columns="a", values=["b"], aggfunc="sum").sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=["c"], columns="a", values=["b", "e"], aggfunc="sum" + ).sort_index(), + pdf.pivot_table( + index=["c"], columns="a", values=["b", "e"], aggfunc="sum" + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=["c"], columns="a", values=["b", "e", "d"], aggfunc="sum" + ).sort_index(), + pdf.pivot_table( + index=["c"], columns="a", values=["b", "e", "d"], aggfunc="sum" + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=["c"], columns="a", values=["b", "e"], aggfunc={"b": "mean", "e": "sum"} + ).sort_index(), + pdf.pivot_table( + index=["c"], columns="a", values=["b", "e"], aggfunc={"b": "mean", "e": "sum"} + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table(index=["e", "c"], columns="a", values="b").sort_index(), + pdf.pivot_table(index=["e", "c"], columns="a", values="b").sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=["e", "c"], columns="a", values="b", fill_value=999 + ).sort_index(), + pdf.pivot_table(index=["e", "c"], columns="a", values="b", fill_value=999).sort_index(), + almost=True, + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("y", "e"), ("z", "c"), ("w", "d")] + ) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.pivot_table(columns=("x", "a"), values=("x", "b")).sort_index(), + pdf.pivot_table(columns=[("x", "a")], values=[("x", "b")]).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=[("z", "c")], columns=("x", "a"), values=[("x", "b")] + ).sort_index(), + pdf.pivot_table( + index=[("z", "c")], columns=[("x", "a")], values=[("x", "b")] + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=[("z", "c")], columns=("x", "a"), values=[("x", "b"), ("y", "e")] + ).sort_index(), + pdf.pivot_table( + index=[("z", "c")], columns=[("x", "a")], values=[("x", "b"), ("y", "e")] + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=[("z", "c")], columns=("x", "a"), values=[("x", "b"), ("y", "e"), ("w", "d")] + ).sort_index(), + pdf.pivot_table( + index=[("z", "c")], + columns=[("x", "a")], + values=[("x", "b"), ("y", "e"), ("w", "d")], + ).sort_index(), + almost=True, + ) + + self.assert_eq( + psdf.pivot_table( + index=[("z", "c")], + columns=("x", "a"), + values=[("x", "b"), ("y", "e")], + aggfunc={("x", "b"): "mean", ("y", "e"): "sum"}, + ).sort_index(), + pdf.pivot_table( + index=[("z", "c")], + columns=[("x", "a")], + values=[("x", "b"), ("y", "e")], + aggfunc={("x", "b"): "mean", ("y", "e"): "sum"}, + ).sort_index(), + almost=True, + ) + + def test_pivot_table_and_index(self): + # https://github.com/databricks/koalas/issues/805 + pdf = pd.DataFrame( + { + "A": ["foo", "foo", "foo", "foo", "foo", "bar", "bar", "bar", "bar"], + "B": ["one", "one", "one", "two", "two", "one", "one", "two", "two"], + "C": [ + "small", + "large", + "large", + "small", + "small", + "large", + "small", + "small", + "large", + ], + "D": [1, 2, 2, 3, 3, 4, 5, 6, 7], + "E": [2, 4, 5, 5, 6, 6, 8, 9, 9], + }, + columns=["A", "B", "C", "D", "E"], + index=np.random.rand(9), + ) + psdf = ps.from_pandas(pdf) + + ptable = pdf.pivot_table( + values="D", index=["A", "B"], columns="C", aggfunc="sum", fill_value=0 + ).sort_index() + ktable = psdf.pivot_table( + values="D", index=["A", "B"], columns="C", aggfunc="sum", fill_value=0 + ).sort_index() + + self.assert_eq(ktable, ptable) + self.assert_eq(ktable.index, ptable.index) + self.assert_eq(repr(ktable.index), repr(ptable.index)) + + def test_pivot_errors(self): + psdf = ps.range(10) + + with self.assertRaisesRegex(ValueError, "columns should be set"): + psdf.pivot(index="id") + + with self.assertRaisesRegex(ValueError, "values should be set"): + psdf.pivot(index="id", columns="id") + + def test_pivot_table_errors(self): + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, 2, 2, 4, 2, 4], + "e": [1, 2, 2, 4, 2, 4], + "c": [1, 2, 9, 4, 7, 4], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + self.assertRaises(KeyError, lambda: psdf.pivot_table(index=["c"], columns="a", values=5)) + + msg = "index should be a None or a list of columns." + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table(index="c", columns="a", values="b") + + msg = "pivot_table doesn't support aggfunc as dict and without index." + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.pivot_table(columns="a", values=["b", "e"], aggfunc={"b": "mean", "e": "sum"}) + + msg = "columns should be one column name." + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table(columns=["a"], values=["b"], aggfunc={"b": "mean", "e": "sum"}) + + msg = "Columns in aggfunc must be the same as values." + with self.assertRaisesRegex(ValueError, msg): + psdf.pivot_table( + index=["e", "c"], columns="a", values="b", aggfunc={"b": "mean", "e": "sum"} + ) + + msg = "values can't be a list without index." + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.pivot_table(columns="a", values=["b", "e"]) + + msg = "Wrong columns A." + with self.assertRaisesRegex(ValueError, msg): + psdf.pivot_table( + index=["c"], columns="A", values=["b", "e"], aggfunc={"b": "mean", "e": "sum"} + ) + + msg = "values should be one column or list of columns." + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table(columns="a", values=(["b"], ["c"])) + + msg = "aggfunc must be a dict mapping from column name to aggregate functions" + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table(columns="a", values="b", aggfunc={"a": lambda x: sum(x)}) + + psdf = ps.DataFrame( + { + "A": ["foo", "foo", "foo", "foo", "foo", "bar", "bar", "bar", "bar"], + "B": ["one", "one", "one", "two", "two", "one", "one", "two", "two"], + "C": [ + "small", + "large", + "large", + "small", + "small", + "large", + "small", + "small", + "large", + ], + "D": [1, 2, 2, 3, 3, 4, 5, 6, 7], + "E": [2, 4, 5, 5, 6, 6, 8, 9, 9], + }, + columns=["A", "B", "C", "D", "E"], + index=np.random.rand(9), + ) + + msg = "values should be a numeric type." + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table( + index=["C"], columns="A", values=["B", "E"], aggfunc={"B": "mean", "E": "sum"} + ) + + msg = "values should be a numeric type." + with self.assertRaisesRegex(TypeError, msg): + psdf.pivot_table(index=["C"], columns="A", values="B", aggfunc={"B": "mean"}) + + +class FramePivotTests(FramePivotMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.computation.test_pivot import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_any_all.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_any_all.py new file mode 100644 index 0000000..29aa233 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_any_all.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_any_all import FrameAnyAllMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityAnyAllTests(FrameAnyAllMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_any_all import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_apply_func.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_apply_func.py new file mode 100644 index 0000000..3a72084 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_apply_func.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_apply_func import FrameApplyFunctionMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityApplyFunctionTests( + FrameApplyFunctionMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_apply_func import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_binary_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_binary_ops.py new file mode 100644 index 0000000..2b10c33 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_binary_ops.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_binary_ops import FrameBinaryOpsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityBinaryOpsTests(FrameBinaryOpsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_binary_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_combine.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_combine.py new file mode 100644 index 0000000..af23600 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_combine.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.computation.test_combine import FrameCombineMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityCombineTests(FrameCombineMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_combine import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_compute.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_compute.py new file mode 100644 index 0000000..88eeb73 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_compute.py @@ -0,0 +1,60 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_compute import FrameComputeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityComputeTests(FrameComputeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_diff(self): + super().test_diff() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_mode(self): + super().test_mode() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_pct_change(self): + super().test_pct_change() + + @unittest.skip("TODO(SPARK-43618): Fix pyspark.sq.column._unary_op to work with Spark Connect.") + def test_rank(self): + super().test_rank() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_compute import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_corrwith.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_corrwith.py new file mode 100644 index 0000000..7c6af99 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_corrwith.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_corrwith import FrameCorrwithMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityCorrwithTests(FrameCorrwithMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_corrwith import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cov.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cov.py new file mode 100644 index 0000000..22abeec --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cov.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_cov import FrameCovMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityCovTests(FrameCovMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_cov import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cumulative.py new file mode 100644 index 0000000..8015d90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_cumulative.py @@ -0,0 +1,90 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_cumulative import FrameCumulativeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityCumulativeTests( + FrameCumulativeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummax(self): + super().test_cummax() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummax_multiindex_columns(self): + super().test_cummax_multiindex_columns() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummin(self): + super().test_cummin() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummin_multiindex_columns(self): + super().test_cummin_multiindex_columns() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumprod(self): + super().test_cumprod() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumprod_multiindex_columns(self): + super().test_cumprod_multiindex_columns() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumsum(self): + super().test_cumsum() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumsum_multiindex_columns(self): + super().test_cumsum_multiindex_columns() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_cumulative import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_describe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_describe.py new file mode 100644 index 0000000..0ca8735 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_describe.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_describe import FrameDescribeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityDescribeTests(FrameDescribeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_describe import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_eval.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_eval.py new file mode 100644 index 0000000..58a347a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_eval.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_eval import FrameEvalMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityEvalTests(FrameEvalMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_eval import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_melt.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_melt.py new file mode 100644 index 0000000..cda16eb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_melt.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_melt import FrameMeltMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityMeltTests(FrameMeltMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_melt import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_missing_data.py new file mode 100644 index 0000000..d2ff09e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_missing_data.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.computation.test_missing_data import FrameMissingDataMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityMissingDataTests( + FrameMissingDataMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_missing_data import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_pivot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_pivot.py new file mode 100644 index 0000000..c8ec48e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/computation/test_parity_pivot.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.computation.test_pivot import FramePivotMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityPivotTests(FramePivotMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.computation.test_parity_pivot import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_base.py new file mode 100644 index 0000000..c277f5c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_base.py @@ -0,0 +1,36 @@ +# +# 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 unittest + +from pyspark.pandas.tests.data_type_ops.test_base import BaseTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class BaseParityTests(BaseTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_base import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py new file mode 100644 index 0000000..663c000 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark.pandas.tests.data_type_ops.test_binary_ops import BinaryOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class BinaryOpsParityTests( + BinaryOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_binary_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py new file mode 100644 index 0000000..52d5179 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_boolean_ops import BooleanOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class BooleanOpsParityTests( + BooleanOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_boolean_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py new file mode 100644 index 0000000..b680e5b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_categorical_ops import CategoricalOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class CategoricalOpsParityTests( + CategoricalOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_eq(self): + super().test_eq() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_ne(self): + super().test_ne() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_categorical_ops import * + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py new file mode 100644 index 0000000..ef58757 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.data_type_ops.test_complex_ops import ComplexOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ComplexOpsParityTests( + ComplexOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_complex_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py new file mode 100644 index 0000000..e7b1c7d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_date_ops import DateOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class DateOpsParityTests( + DateOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_date_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py new file mode 100644 index 0000000..6d081b1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_datetime_ops import DatetimeOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class DatetimeOpsParityTests( + DatetimeOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_datetime_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py new file mode 100644 index 0000000..63b53c0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark.pandas.tests.data_type_ops.test_null_ops import NullOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class NullOpsParityTests( + NullOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_null_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py new file mode 100644 index 0000000..6f5c294 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_num_arithmetic import ArithmeticTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ArithmeticParityTests( + ArithmeticTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_num_arithmetic import * # noqa + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py new file mode 100644 index 0000000..04aa24c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_num_ops import NumOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class NumOpsParityTests( + NumOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py new file mode 100644 index 0000000..4d322d8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_num_reverse import ReverseTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ReverseParityTests( + ReverseTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_num_reverse import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py new file mode 100644 index 0000000..ecbf94a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.test_string_ops import StringOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class StringOpsParityTests( + StringOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_string_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py new file mode 100644 index 0000000..058dd2b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py @@ -0,0 +1,47 @@ +# +# 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 unittest + +import pyspark.pandas as ps +from pyspark.pandas.tests.data_type_ops.test_timedelta_ops import TimedeltaOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class TimedeltaOpsParityTests( + TimedeltaOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_timedelta_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py new file mode 100644 index 0000000..70a79e4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.data_type_ops.test_udt_ops import UDTOpsTestsMixin +from pyspark.pandas.tests.connect.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class UDTOpsParityTests( + UDTOpsTestsMixin, PandasOnSparkTestUtils, OpsTestBase, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.data_type_ops.test_parity_udt_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/testing_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/testing_utils.py new file mode 100644 index 0000000..6e06f2b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/data_type_ops/testing_utils.py @@ -0,0 +1,226 @@ +# +# 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 datetime +import decimal +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.pandas.typedef import extension_dtypes + +from pyspark.pandas.typedef.typehints import ( + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) + +if extension_dtypes_available: + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + +if extension_float_dtypes_available: + from pandas import Float32Dtype, Float64Dtype + +if extension_object_dtypes_available: + from pandas import BooleanDtype, StringDtype + + +class OpsTestBase: + """The test base for arithmetic operations of different data types.""" + + @property + def numeric_pdf(self): + dtypes = [np.int32, int, np.float32, float] + sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes] + sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)])) + sers.append(pd.Series([1, 2, np.nan], dtype=float)) + # Skip decimal_nan test before v1.3.0, it not supported by pandas on spark yet. + if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + sers.append( + pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(np.nan)]) + ) + pdf = pd.concat(sers, axis=1) + if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + pdf.columns = [dtype.__name__ for dtype in dtypes] + [ + "decimal", + "float_nan", + "decimal_nan", + ] + else: + pdf.columns = [dtype.__name__ for dtype in dtypes] + ["decimal", "float_nan"] + return pdf + + @property + def numeric_psdf(self): + return ps.from_pandas(self.numeric_pdf) + + @property + def numeric_df_cols(self): + return self.numeric_pdf.columns + + @property + def integral_pdf(self): + return pd.DataFrame({"this": [1, 2, 3], "that": [2, 2, 1]}) + + @property + def integral_psdf(self): + return ps.from_pandas(self.integral_pdf) + + @property + def non_numeric_pdf(self): + psers = { + "string": pd.Series(["x", "y", "z"]), + "bool": pd.Series([True, True, False]), + "date": pd.Series( + [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)] + ), + "datetime": pd.to_datetime(pd.Series([1, 2, 3])), + "timedelta": pd.Series( + [datetime.timedelta(1), datetime.timedelta(hours=2), datetime.timedelta(weeks=3)] + ), + "categorical": pd.Series(["a", "b", "a"], dtype="category"), + } + return pd.concat(psers, axis=1) + + @property + def non_numeric_psdf(self): + return ps.from_pandas(self.non_numeric_pdf) + + @property + def non_numeric_df_cols(self): + return self.non_numeric_pdf.columns + + @property + def pdf(self): + return pd.concat([self.numeric_pdf, self.non_numeric_pdf], axis=1) + + @property + def df_cols(self): + return self.pdf.columns + + @property + def numeric_psers(self): + dtypes = [np.float32, float, int, np.int32] + sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes] + sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)])) + return sers + + @property + def numeric_pssers(self): + return [ps.from_pandas(pser) for pser in self.numeric_psers] + + @property + def numeric_pser_psser_pairs(self): + return zip(self.numeric_psers, self.numeric_pssers) + + @property + def non_numeric_psers(self): + psers = { + "string": pd.Series(["x", "y", "z"]), + "datetime": pd.to_datetime(pd.Series([1, 2, 3])), + "bool": pd.Series([True, True, False]), + "date": pd.Series( + [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)] + ), + "categorical": pd.Series(["a", "b", "a"], dtype="category"), + } + return psers + + @property + def non_numeric_pssers(self): + pssers = {} + + for k, v in self.non_numeric_psers.items(): + pssers[k] = ps.from_pandas(v) + return pssers + + @property + def non_numeric_pser_psser_pairs(self): + return zip(self.non_numeric_psers.values(), self.non_numeric_pssers.values()) + + @property + def pssers(self): + return self.numeric_pssers + list(self.non_numeric_pssers.values()) + + @property + def psers(self): + return self.numeric_psers + list(self.non_numeric_psers.values()) + + @property + def pser_psser_pairs(self): + return zip(self.psers, self.pssers) + + @property + def string_extension_dtype(self): + return ["string", StringDtype()] if extension_object_dtypes_available else [] + + @property + def object_extension_dtypes(self): + return ( + ["boolean", "string", BooleanDtype(), StringDtype()] + if extension_object_dtypes_available + else [] + ) + + @property + def fractional_extension_dtypes(self): + return ( + ["Float32", "Float64", Float32Dtype(), Float64Dtype()] + if extension_float_dtypes_available + else [] + ) + + @property + def integral_extension_dtypes(self): + return ( + [ + "Int8", + "Int16", + "Int32", + "Int64", + Int8Dtype(), + Int16Dtype(), + Int32Dtype(), + Int64Dtype(), + ] + if extension_dtypes_available + else [] + ) + + @property + def extension_dtypes(self): + return ( + self.object_extension_dtypes + + self.fractional_extension_dtypes + + self.integral_extension_dtypes + ) + + def check_extension(self, left, right): + """ + Compare `psser` and `pser` of numeric ExtensionDtypes. + + This utility is to adjust an issue for comparing numeric ExtensionDtypes in specific + pandas versions. Please refer to https://github.com/pandas-dev/pandas/issues/39410. + """ + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(left, right, check_exact=False) + self.assertTrue(isinstance(left.dtype, extension_dtypes)) + self.assertTrue(isinstance(right.dtype, extension_dtypes)) + else: + self.assert_eq(left, right) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_align.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_align.py new file mode 100644 index 0000000..5e180eb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_align.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_align import DiffFramesAlignMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityAlignTests( + DiffFramesAlignMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_align import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic_slow.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic_slow.py new file mode 100644 index 0000000..926caf5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic_slow.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_basic_slow import DiffFramesBasicSlowMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityBasicSlowTests( + DiffFramesBasicSlowMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_diff(self): + super().test_diff() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_rank(self): + super().test_rank() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_shift(self): + super().test_shift() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_basic_slow import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_cov_corrwith.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_cov_corrwith.py new file mode 100644 index 0000000..b496c0c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_cov_corrwith.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_cov_corrwith import DiffFramesCovCorrWithMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityCovCorrWithTests( + DiffFramesCovCorrWithMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_cov_corrwith import * # noqa + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_frame.py new file mode 100644 index 0000000..b71f092 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_frame.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_dot_frame import DiffFramesDotFrameMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityDotFrameTests( + DiffFramesDotFrameMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_dot_frame import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_series.py new file mode 100644 index 0000000..81c1cd9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_dot_series.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_dot_series import DiffFramesDotSeriesMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityDotSeriesTests( + DiffFramesDotSeriesMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_dot_series import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_index.py new file mode 100644 index 0000000..9cc81e0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_index.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_index import DiffFramesIndexMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParityIndexTests( + DiffFramesIndexMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_index import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_series.py new file mode 100644 index 0000000..238dabd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_series.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_series import DiffFramesSeriesMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParitySeriesTests( + DiffFramesSeriesMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_series import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_frame.py new file mode 100644 index 0000000..397d4b7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_frame.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_setitem_frame import DiffFramesSetItemFrameMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParitySetItemFrameTests( + DiffFramesSetItemFrameMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43658): Fix unexpected `SparkConnectGrpcException` from Spark Connect client." + ) + def test_frame_iloc_setitem(self): + super().test_frame_iloc_setitem() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_setitem_frame import * # noqa + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_series.py new file mode 100644 index 0000000..07aff29 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_setitem_series.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark.pandas.tests.diff_frames_ops.test_setitem_series import DiffFramesSetItemSeriesMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DiffFramesParitySetItemSeriesTests( + DiffFramesSetItemSeriesMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43658): Fix unexpected `SparkConnectGrpcException` from Spark Connect client." + ) + def test_series_iloc_setitem(self): + super().test_series_iloc_setitem() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_setitem_series import * # noqa + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_attrs.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_attrs.py new file mode 100644 index 0000000..2f13601 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_attrs.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_attrs import FrameAttrsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityAttrsTests(FrameAttrsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_attrs import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_constructor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_constructor.py new file mode 100644 index 0000000..19af61c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_constructor.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.frame.test_constructor import FrameConstructorMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityConstructorTests( + FrameConstructorMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_constructor import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_conversion.py new file mode 100644 index 0000000..a580254 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_conversion.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_conversion import FrameConversionMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityConversionTests( + FrameConversionMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_conversion import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reindexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reindexing.py new file mode 100644 index 0000000..0caa6c3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reindexing.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_reindexing import FrameReindexingMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityReindexingTests( + FrameReindexingMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_reindexing import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reshaping.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reshaping.py new file mode 100644 index 0000000..e4bac7b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_reshaping.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.frame.test_reshaping import FrameReshapingMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityReshapingTests(FrameReshapingMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_reshaping import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_spark.py new file mode 100644 index 0000000..09de0f1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_spark.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_spark import FrameSparkMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParitySparkTests(FrameSparkMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_spark import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_take.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_take.py new file mode 100644 index 0000000..4fc8b24 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_take.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_take import FrameTakeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityTakeTests(FrameTakeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_take import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_time_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_time_series.py new file mode 100644 index 0000000..ae289ed --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_time_series.py @@ -0,0 +1,48 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_time_series import FrameTimeSeriesMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityTimeSeriesTests( + FrameTimeSeriesMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_shift(self): + super().test_shift() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_time_series import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_truncate.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_truncate.py new file mode 100644 index 0000000..32adbf3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/frame/test_parity_truncate.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.frame.test_truncate import FrameTruncateMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityTruncateTests(FrameTruncateMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.frame.test_parity_truncate import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_aggregate.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_aggregate.py new file mode 100644 index 0000000..e3c9f62 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_aggregate.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_aggregate import GroupbyAggregateMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityAggregateTests( + GroupbyAggregateMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_aggregate import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_apply_func.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_apply_func.py new file mode 100644 index 0000000..8eeec76 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_apply_func.py @@ -0,0 +1,41 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_apply_func import GroupbyApplyFuncMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityApplyFuncTests( + GroupbyApplyFuncMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43628): Enable SparkContext with Spark Connect.") + def test_apply_with_side_effect(self): + super().test_apply_with_side_effect() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_apply_func import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_cumulative.py new file mode 100644 index 0000000..aaa799b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_cumulative.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. +# +import unittest + +from pyspark.pandas.tests.groupby.test_cumulative import GroupbyCumulativeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityCumulativeTests( + GroupbyCumulativeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumcount(self): + super().test_cumcount() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummax(self): + super().test_cummax() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cummin(self): + super().test_cummin() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumprod(self): + super().test_cumprod() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_cumsum(self): + super().test_cumsum() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_cumulative import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_describe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_describe.py new file mode 100644 index 0000000..826e1b0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_describe.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_describe import GroupbyDescribeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityDescribeTests( + GroupbyDescribeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_describe import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_groupby.py new file mode 100644 index 0000000..4e9f510 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_groupby.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_groupby import GroupByTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class GroupByParityTests( + GroupByTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_shift(self): + super().test_shift() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_diff(self): + super().test_diff() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_rank(self): + super().test_rank() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_groupby import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_head_tail.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_head_tail.py new file mode 100644 index 0000000..373e18e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_head_tail.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_head_tail import GroupbyHeadTailMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityHeadTailTests( + GroupbyHeadTailMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_head_tail import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_index.py new file mode 100644 index 0000000..6a0c841 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_index.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_index import GroupbyIndexMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityIndexTests(GroupbyIndexMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_index import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_missing_data.py new file mode 100644 index 0000000..1ca101e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_missing_data.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_missing_data import GroupbyMissingDataMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityMissingDataTests( + GroupbyMissingDataMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_bfill(self): + super().test_bfill() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_ffill(self): + super().test_ffill() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_fillna(self): + super().test_fillna() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_missing_data import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply.py new file mode 100644 index 0000000..895fe98 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParitySplitApplyTests( + GroupbySplitApplyMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_split_apply import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_stat.py new file mode 100644 index 0000000..75f1ed4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/groupby/test_parity_stat.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.groupby.test_stat import GroupbyStatMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class GroupbyParityStatTests(GroupbyStatMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.groupby.test_parity_stat import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_align.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_align.py new file mode 100644 index 0000000..0bf84e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_align.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_align import FrameAlignMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityAlignTests(FrameAlignMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_align import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_base.py new file mode 100644 index 0000000..3cf4dc9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_base.py @@ -0,0 +1,50 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_base import IndexesTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class IndexesParityTests( + IndexesTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_factorize(self): + super().test_factorize() + + @unittest.skip("TODO(SPARK-43704): Enable IndexesParityTests.test_to_series.") + def test_to_series(self): + super().test_to_series() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_base import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_category.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_category.py new file mode 100644 index 0000000..d99d013 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_category.py @@ -0,0 +1,73 @@ +# +# 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 unittest + +from pyspark.pandas.tests.indexes.test_category import CategoricalIndexTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class CategoricalIndexParityTests( + CategoricalIndexTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_append(self): + super().test_append() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_factorize(self): + super().test_factorize() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_intersection(self): + super().test_intersection() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_remove_categories(self): + super().test_remove_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_remove_unused_categories(self): + super().test_remove_unused_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_reorder_categories(self): + super().test_reorder_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_set_categories(self): + super().test_set_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_union(self): + super().test_union() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_category import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_datetime.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_datetime.py new file mode 100644 index 0000000..48c3e49 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_datetime.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class DatetimeIndexParityTests( + DatetimeIndexTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_datetime import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing.py new file mode 100644 index 0000000..a764893 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_indexing import FrameIndexingMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityIndexingTests(FrameIndexingMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_indexing import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reindex.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reindex.py new file mode 100644 index 0000000..7e9c535 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reindex.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_reindex import FrameReindexMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityReindexTests(FrameReindexMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_reindex import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_rename.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_rename.py new file mode 100644 index 0000000..77da128 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_rename.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_rename import FrameRenameMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityRenameTests(FrameRenameMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_rename import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reset_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reset_index.py new file mode 100644 index 0000000..6647d76 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_reset_index.py @@ -0,0 +1,48 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.indexes.test_reset_index import FrameResetIndexMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityResetIndexTests( + FrameResetIndexMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_reset_index_with_default_index_types(self): + super().test_reset_index_with_default_index_types() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_reset_index import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_timedelta.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_timedelta.py new file mode 100644 index 0000000..c980ef4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/indexes/test_parity_timedelta.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.indexes.test_timedelta import TimedeltaIndexTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class TimedeltaIndexParityTests( + TimedeltaIndexTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.indexes.test_parity_timedelta import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/test_parity_io.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/test_parity_io.py new file mode 100644 index 0000000..9aab992 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/io/test_parity_io.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.io.test_io import FrameIOMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class FrameParityIOTests(FrameIOMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.io.test_parity_io import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot.py new file mode 100644 index 0000000..24392ea --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot.py @@ -0,0 +1,45 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_frame_plot import DataFramePlotTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DataFramePlotParityTests( + DataFramePlotTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_compute_hist_multi_columns(self): + super().test_compute_hist_multi_columns() + + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_compute_hist_single_column(self): + super().test_compute_hist_single_column() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_frame_plot import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_matplotlib.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_matplotlib.py new file mode 100644 index 0000000..98da885 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_matplotlib.py @@ -0,0 +1,45 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_frame_plot_matplotlib import DataFramePlotMatplotlibTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class DataFramePlotMatplotlibParityTests( + DataFramePlotMatplotlibTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist_plot(self): + super().test_hist_plot() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_kde_plot(self): + super().test_kde_plot() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_frame_plot_matplotlib import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_plotly.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_plotly.py new file mode 100644 index 0000000..7a3efee --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_frame_plot_plotly.py @@ -0,0 +1,49 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_frame_plot_plotly import DataFramePlotPlotlyTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class DataFramePlotPlotlyParityTests( + DataFramePlotPlotlyTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist_layout_kwargs(self): + super().test_hist_layout_kwargs() + + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist_plot(self): + super().test_hist_plot() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_kde_plot(self): + super().test_kde_plot() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot.py new file mode 100644 index 0000000..abc9317 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_series_plot import SeriesPlotTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesPlotParityTests(SeriesPlotTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_series_plot import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_matplotlib.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_matplotlib.py new file mode 100644 index 0000000..975d78f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_matplotlib.py @@ -0,0 +1,53 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_series_plot_matplotlib import SeriesPlotMatplotlibTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class SeriesPlotMatplotlibParityTests( + SeriesPlotMatplotlibTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist(self): + super().test_hist() + + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist_plot(self): + super().test_hist_plot() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_kde_plot(self): + super().test_kde_plot() + + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_single_value_hist(self): + super().test_single_value_hist() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_series_plot_matplotlib import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_plotly.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_plotly.py new file mode 100644 index 0000000..5ced8cb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/plot/test_parity_series_plot_plotly.py @@ -0,0 +1,45 @@ +# +# 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 unittest + +from pyspark.pandas.tests.plot.test_series_plot_plotly import SeriesPlotPlotlyTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class SeriesPlotPlotlyParityTests( + SeriesPlotPlotlyTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43711): Fix Transformer.transform to work with Spark Connect.") + def test_hist_plot(self): + super().test_hist_plot() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_kde_plot(self): + super().test_kde_plot() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.plot.test_parity_series_plot_plotly import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_all_any.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_all_any.py new file mode 100644 index 0000000..44856fa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_all_any.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_all_any import SeriesAllAnyMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityAllAnyTests(SeriesAllAnyMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_all_any import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_arg_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_arg_ops.py new file mode 100644 index 0000000..bd17521 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_arg_ops.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_arg_ops import SeriesArgOpsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityArgOpsTests(SeriesArgOpsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_arg_ops import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_of.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_of.py new file mode 100644 index 0000000..ad4faed --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_of.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_as_of import SeriesAsOfMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityArgOpsTests(SeriesAsOfMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_as_of import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_type.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_type.py new file mode 100644 index 0000000..5fa949c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_as_type.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_as_type import SeriesAsTypeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityAsTypeTests(SeriesAsTypeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_as_type import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_compute.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_compute.py new file mode 100644 index 0000000..8876fcb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_compute.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_compute import SeriesComputeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityComputeTests(SeriesComputeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_compute import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_conversion.py new file mode 100644 index 0000000..08d2773 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_conversion.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_conversion import SeriesConversionMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityConversionTests( + SeriesConversionMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_conversion import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_cumulative.py new file mode 100644 index 0000000..c5c61e1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_cumulative.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_cumulative import SeriesCumulativeMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityCumulativeTests( + SeriesCumulativeMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_cumulative import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_index.py new file mode 100644 index 0000000..2b92cce --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_index.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_index import SeriesIndexMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityIndexTests(SeriesIndexMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_index import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_missing_data.py new file mode 100644 index 0000000..a95b312 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_missing_data.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_missing_data import SeriesMissingDataMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityMissingDataTests( + SeriesMissingDataMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_missing_data import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_series.py new file mode 100644 index 0000000..5f1d0d1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_series.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_series import SeriesTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityTests(SeriesTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_series import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_sort.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_sort.py new file mode 100644 index 0000000..bf124bc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_sort.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_sort import SeriesSortMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParitySortTests(SeriesSortMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_sort import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_stat.py new file mode 100644 index 0000000..916e120 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/series/test_parity_stat.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.pandas.tests.series.test_stat import SeriesStatMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesParityStatTests(SeriesStatMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.series.test_parity_stat import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_categorical.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_categorical.py new file mode 100644 index 0000000..210cfce --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_categorical.py @@ -0,0 +1,66 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.test_categorical import CategoricalTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class CategoricalParityTests( + CategoricalTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_astype(self): + super().test_astype() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_factorize(self): + super().test_factorize() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_remove_categories(self): + super().test_remove_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_remove_unused_categories(self): + super().test_remove_unused_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_reorder_categories(self): + super().test_reorder_categories() + + @unittest.skip("TODO(SPARK-43620): Support `Column` for SparkConnectColumn.__getitem__.") + def test_set_categories(self): + super().test_set_categories() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_categorical import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_config.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_config.py new file mode 100644 index 0000000..e394d14 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_config.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_config import ConfigTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ConfigParityTests(ConfigTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_config import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_csv.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_csv.py new file mode 100644 index 0000000..2b0c0af --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_csv.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_csv import CsvTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class CsvParityTests(CsvTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_csv import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_conversion.py new file mode 100644 index 0000000..c5a26a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_conversion.py @@ -0,0 +1,42 @@ +# +# 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 unittest + +from pyspark import pandas as ps +from pyspark.pandas.tests.test_dataframe_conversion import DataFrameConversionTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class DataFrameConversionParityTests( + DataFrameConversionTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase, TestUtils +): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_dataframe_conversion import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_spark_io.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_spark_io.py new file mode 100644 index 0000000..3b700dd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_dataframe_spark_io.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_dataframe_spark_io import DataFrameSparkIOTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class DataFrameSparkIOParityTests( + DataFrameSparkIOTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase, TestUtils +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_dataframe_spark_io import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_default_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_default_index.py new file mode 100644 index 0000000..c5410e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_default_index.py @@ -0,0 +1,49 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_default_index import DefaultIndexTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class DefaultIndexParityTests( + DefaultIndexTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_default_index_sequence(self): + super().test_default_index_sequence() + + @unittest.skip( + "TODO(SPARK-43623): Enable DefaultIndexParityTests.test_index_distributed_sequence_cleanup." + ) + def test_index_distributed_sequence_cleanup(self): + super().test_index_distributed_sequence_cleanup() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_default_index import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ewm.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ewm.py new file mode 100644 index 0000000..7487282 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ewm.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_ewm import EWMTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class EWMParityTests(EWMTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_ewm import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_expanding.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_expanding.py new file mode 100644 index 0000000..7f8b1a3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_expanding.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_expanding import ExpandingTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class ExpandingParityTests( + ExpandingTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_expanding import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_extension.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_extension.py new file mode 100644 index 0000000..8491399 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_extension.py @@ -0,0 +1,49 @@ +# +# 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 unittest + +import pandas as pd +import numpy as np +from pyspark import pandas as ps +from pyspark.pandas.tests.test_extension import ExtensionTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ExtensionParityTests(ExtensionTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_extension import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_frame_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_frame_spark.py new file mode 100644 index 0000000..986ddde --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_frame_spark.py @@ -0,0 +1,53 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_frame_spark import SparkFrameMethodsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class SparkFrameMethodsParityTests( + SparkFrameMethodsTestsMixin, TestUtils, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43628): Enable SparkContext with Spark Connect.") + def test_checkpoint(self): + super().test_checkpoint() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_coalesce(self): + super().test_coalesce() + + @unittest.skip("TODO(SPARK-43630): Implement `localCheckpoint` for Spark Connect DataFrame.") + def test_local_checkpoint(self): + super().test_local_checkpoint() + + @unittest.skip("TODO(SPARK-43629): Enable RDD with Spark Connect.") + def test_repartition(self): + super().test_repartition() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_frame_spark import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_generic_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_generic_functions.py new file mode 100644 index 0000000..1bf2650 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_generic_functions.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_generic_functions import GenericFunctionsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class GenericFunctionsParityTests( + GenericFunctionsTestsMixin, TestUtils, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_interpolate(self): + super().test_interpolate() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_generic_functions import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexing.py new file mode 100644 index 0000000..9a14978 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexing.py @@ -0,0 +1,49 @@ +# +# 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 unittest + +import pandas as pd +from pyspark import pandas as ps +from pyspark.pandas.tests.test_indexing import BasicIndexingTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class BasicIndexingParityTests( + BasicIndexingTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @property + def pdf(self): + return pd.DataFrame( + {"month": [1, 4, 7, 10], "year": [2012, 2014, 2013, 2014], "sale": [55, 40, 84, 31]} + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_indexing import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexops_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexops_spark.py new file mode 100644 index 0000000..37a2ba6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_indexops_spark.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_indexops_spark import SparkIndexOpsMethodsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SparkIndexOpsMethodsParityTests( + SparkIndexOpsMethodsTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_indexops_spark import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_internal.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_internal.py new file mode 100644 index 0000000..d586fec --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_internal.py @@ -0,0 +1,41 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_internal import InternalFrameTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class InternalFrameParityTests( + InternalFrameTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + @unittest.skip("TODO(SPARK-43654): Enable InternalFrameParityTests.test_from_pandas.") + def test_from_pandas(self): + super().test_from_pandas() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_internal import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_namespace.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_namespace.py new file mode 100644 index 0000000..db7f62f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_namespace.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_namespace import NamespaceTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class NamespaceParityTests(NamespaceTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @unittest.skip("TODO(SPARK-43655): Enable NamespaceParityTests.test_get_index_map.") + def test_get_index_map(self): + super().test_get_index_map() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_namespace import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_numpy_compat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_numpy_compat.py new file mode 100644 index 0000000..1bfeda0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_numpy_compat.py @@ -0,0 +1,60 @@ +# +# 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 unittest + +import pandas as pd +from pyspark import pandas as ps +from pyspark.pandas.tests.test_numpy_compat import NumPyCompatTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class NumPyCompatParityTests(NumPyCompatTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skip( + "TODO(SPARK-43656): Fix pyspark.sql.column._to_java_column to accept Connect Column." + ) + def test_np_spark_compat_frame(self): + super().test_np_spark_compat_frame() + + @unittest.skip( + "TODO(SPARK-43656): Fix pyspark.sql.column._to_java_column to accept Connect Column." + ) + def test_np_spark_compat_series(self): + super().test_np_spark_compat_series() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_numpy_compat import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames.py new file mode 100644 index 0000000..20d7efc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames.py @@ -0,0 +1,48 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_ops_on_diff_frames import ( + OpsOnDiffFramesDisabledTestsMixin, + OpsOnDiffFramesEnabledTestsMixin, +) +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class OpsOnDiffFramesEnabledParityTests( + OpsOnDiffFramesEnabledTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +class OpsOnDiffFramesDisabledParityTests( + OpsOnDiffFramesDisabledTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py new file mode 100644 index 0000000..685ec5c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import OpsOnDiffFramesGroupByTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class OpsOnDiffFramesGroupByParityTests( + OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_expanding.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_expanding.py new file mode 100644 index 0000000..c373268 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_expanding.py @@ -0,0 +1,44 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_ops_on_diff_frames_groupby_expanding import ( + OpsOnDiffFramesGroupByExpandingTestsMixin, +) +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class OpsOnDiffFramesGroupByExpandingParityTests( + OpsOnDiffFramesGroupByExpandingTestsMixin, + PandasOnSparkTestUtils, + TestUtils, + ReusedConnectTestCase, +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby_expanding import * + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_rolling.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_rolling.py new file mode 100644 index 0000000..dd82e44 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby_rolling.py @@ -0,0 +1,84 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_ops_on_diff_frames_groupby_rolling import ( + OpsOnDiffFramesGroupByRollingTestsMixin, +) +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class OpsOnDiffFramesGroupByRollingParityTests( + OpsOnDiffFramesGroupByRollingTestsMixin, + PandasOnSparkTestUtils, + TestUtils, + ReusedConnectTestCase, +): + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_count(self): + super().test_groupby_rolling_count() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_min(self): + super().test_groupby_rolling_min() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_max(self): + super().test_groupby_rolling_max() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_mean(self): + super().test_groupby_rolling_mean() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_sum(self): + super().test_groupby_rolling_sum() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_std(self): + super().test_groupby_rolling_std() + + @unittest.skip( + "TODO(SPARK-43611): Fix unexpected `AnalysisException` from Spark Connect client." + ) + def test_groupby_rolling_var(self): + super().test_groupby_rolling_var() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby_rolling import * + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_repr.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_repr.py new file mode 100644 index 0000000..1f558c3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_repr.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_repr import ReprTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ReprParityTests(ReprTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_repr import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_resample.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_resample.py new file mode 100644 index 0000000..caca2f9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_resample.py @@ -0,0 +1,45 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_resample import ResampleTestsMixin, ResampleWithTimezoneMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class ResampleParityTests( + ResampleTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +class ResampleWithTimezoneTests( + ResampleWithTimezoneMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_resample import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_reshape.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_reshape.py new file mode 100644 index 0000000..0773978 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_reshape.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_reshape import ReshapeTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ReshapeParityTests(ReshapeTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @unittest.skip("TODO(SPARK-43662): Enable ReshapeParityTests.test_merge_asof.") + def test_merge_asof(self): + super().test_merge_asof() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_reshape import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_rolling.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_rolling.py new file mode 100644 index 0000000..8318bed --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_rolling.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_rolling import RollingTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class RollingParityTests( + RollingTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_rolling import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_scalars.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_scalars.py new file mode 100644 index 0000000..3c93244 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_scalars.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_scalars import ScalarTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ScalarParityTests(ScalarTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_scalars import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_conversion.py new file mode 100644 index 0000000..6545b96 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_conversion.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_series_conversion import SeriesConversionTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesConversionParityTests( + SeriesConversionTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_series_conversion import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_datetime.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_datetime.py new file mode 100644 index 0000000..0842558 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_datetime.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_series_datetime import SeriesDateTimeTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesDateTimeParityTests( + SeriesDateTimeTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_series_datetime import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_string.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_string.py new file mode 100644 index 0000000..d7c0335 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_series_string.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_series_string import SeriesStringTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SeriesStringParityTests( + SeriesStringTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_series_string import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_spark_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_spark_functions.py new file mode 100644 index 0000000..00f7514 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_spark_functions.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_spark_functions import SparkFunctionsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SparkFunctionsParityTests( + SparkFunctionsTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_spark_functions import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_sql.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_sql.py new file mode 100644 index 0000000..c042de6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_sql.py @@ -0,0 +1,43 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_sql import SQLTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class SQLParityTests(SQLTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + @unittest.skip("TODO(SPARK-43664): Fix TABLE_OR_VIEW_NOT_FOUND from SQLParityTests.") + def test_sql_with_index_col(self): + super().test_sql_with_index_col() + + @unittest.skip("TODO(SPARK-43664): Fix TABLE_OR_VIEW_NOT_FOUND from SQLParityTests.") + def test_sql_with_pandas_on_spark_objects(self): + super().test_sql_with_pandas_on_spark_objects() + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_sql import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_stats.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_stats.py new file mode 100644 index 0000000..7eddc4c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_stats.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_stats import StatsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class StatsParityTests(StatsTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_stats import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_typedef.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_typedef.py new file mode 100644 index 0000000..8df36ad --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_typedef.py @@ -0,0 +1,36 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_typedef import TypeHintTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class TypeHintParityTests(TypeHintTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_typedef import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_utils.py new file mode 100644 index 0000000..67c6fad --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_utils.py @@ -0,0 +1,37 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_utils import UtilsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class UtilsParityTests(UtilsTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_utils import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_window.py new file mode 100644 index 0000000..dc54277 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/connect/test_parity_window.py @@ -0,0 +1,39 @@ +# +# 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 unittest + +from pyspark.pandas.tests.test_window import ExpandingRollingTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils, TestUtils + + +class ExpandingRollingParityTests( + ExpandingRollingTestsMixin, PandasOnSparkTestUtils, TestUtils, ReusedConnectTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.connect.test_parity_window import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_base.py new file mode 100644 index 0000000..551bbba --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_base.py @@ -0,0 +1,107 @@ +# +# 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 unittest + +from pandas.api.types import CategoricalDtype +from pandas.api.extensions import ExtensionDtype + +from pyspark.pandas.data_type_ops.base import DataTypeOps +from pyspark.pandas.data_type_ops.binary_ops import BinaryOps +from pyspark.pandas.data_type_ops.boolean_ops import BooleanOps, BooleanExtensionOps +from pyspark.pandas.data_type_ops.categorical_ops import CategoricalOps +from pyspark.pandas.data_type_ops.complex_ops import ArrayOps, MapOps, StructOps +from pyspark.pandas.data_type_ops.date_ops import DateOps +from pyspark.pandas.data_type_ops.datetime_ops import DatetimeOps, DatetimeNTZOps +from pyspark.pandas.data_type_ops.null_ops import NullOps +from pyspark.pandas.data_type_ops.num_ops import IntegralOps, FractionalOps, DecimalOps +from pyspark.pandas.data_type_ops.string_ops import StringOps +from pyspark.pandas.data_type_ops.timedelta_ops import TimedeltaOps +from pyspark.pandas.data_type_ops.udt_ops import UDTOps +from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + DataType, + DateType, + DayTimeIntervalType, + DecimalType, + FractionalType, + IntegralType, + MapType, + NullType, + StringType, + StructType, + TimestampType, + TimestampNTZType, + UserDefinedType, +) + + +class BaseTestsMixin: + def test_data_type_ops(self): + _mock_spark_type = DataType() + _mock_dtype = ExtensionDtype() + _mappings = ( + (CategoricalDtype(), _mock_spark_type, CategoricalOps), + (_mock_dtype, DecimalType(), DecimalOps), + (_mock_dtype, FractionalType(), FractionalOps), + (_mock_dtype, IntegralType(), IntegralOps), + (_mock_dtype, StringType(), StringOps), + (_mock_dtype, BooleanType(), BooleanOps), + (_mock_dtype, TimestampType(), DatetimeOps), + (_mock_dtype, TimestampNTZType(), DatetimeNTZOps), + (_mock_dtype, DateType(), DateOps), + (_mock_dtype, DayTimeIntervalType(), TimedeltaOps), + (_mock_dtype, BinaryType(), BinaryOps), + (_mock_dtype, ArrayType(StringType()), ArrayOps), + (_mock_dtype, MapType(StringType(), IntegralType()), MapOps), + (_mock_dtype, StructType(), StructOps), + (_mock_dtype, NullType(), NullOps), + (_mock_dtype, UserDefinedType(), UDTOps), + ) + for _dtype, _spark_type, _ops in _mappings: + self.assertIsInstance(DataTypeOps(_dtype, _spark_type), _ops) + + _unknow_spark_type = _mock_spark_type + self.assertRaises(TypeError, DataTypeOps, BooleanType(), _unknow_spark_type) + + def test_bool_ext_ops(self): + from pyspark.pandas.typedef.typehints import extension_object_dtypes_available + + if extension_object_dtypes_available: + from pandas import BooleanDtype + + self.assertIsInstance(DataTypeOps(BooleanDtype(), BooleanType()), BooleanExtensionOps) + else: + self.assertIsInstance(DataTypeOps(ExtensionDtype(), BooleanType()), BooleanOps) + + +class BaseTests(BaseTestsMixin, unittest.TestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.data_type_ops.test_base import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py new file mode 100644 index 0000000..732cc29 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py @@ -0,0 +1,224 @@ +# +# 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 pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class BinaryOpsTestsMixin: + @property + def pser(self): + return pd.Series([b"1", b"2", b"3"]) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def byte_pdf(self): + psers = { + "this": self.pser, + "that": pd.Series([b"2", b"3", b"4"]), + } + return pd.concat(psers, axis=1) + + @property + def byte_psdf(self): + return ps.from_pandas(self.byte_pdf) + + def test_add(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + pser, psser = byte_pdf["this"], byte_psdf["this"] + other_pser, other_psser = byte_pdf["that"], byte_psdf["that"] + + self.assert_eq(psser + b"1", pser + b"1") + self.assert_eq(psser + psser, pser + pser) + self.assert_eq(psser + psser.astype("bytes"), pser + pser.astype("bytes")) + self.assertRaises(TypeError, lambda: psser + "x") + self.assertRaises(TypeError, lambda: psser + 1) + + self.assert_eq(pser + pser, psser + psser) + self.assert_eq(pser + other_pser, psser + other_psser) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser - psser) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assert_eq(b"1" + self.psser, b"1" + self.pser) + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psser & True) + self.assertRaises(TypeError, lambda: self.psser & False) + self.assertRaises(TypeError, lambda: self.psser & self.psser) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psser) + self.assertRaises(TypeError, lambda: False & self.psser) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psser | True) + self.assertRaises(TypeError, lambda: self.psser | False) + self.assertRaises(TypeError, lambda: self.psser | self.psser) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psser) + self.assertRaises(TypeError, lambda: False | self.psser) + + def test_from_to_pandas(self): + data = [b"1", b"2", b"3"] + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + self.assert_eq(psser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=[b"2", b"3", b"1"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] == byte_pdf["that"], byte_psdf["this"] == byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] == byte_pdf["this"], byte_psdf["this"] == byte_psdf["this"]) + + def test_ne(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] != byte_pdf["that"], byte_psdf["this"] != byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] != byte_pdf["this"], byte_psdf["this"] != byte_psdf["this"]) + + def test_lt(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] < byte_pdf["that"], byte_psdf["this"] < byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] < byte_pdf["this"], byte_psdf["this"] < byte_psdf["this"]) + + def test_le(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] <= byte_pdf["that"], byte_psdf["this"] <= byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] <= byte_pdf["this"], byte_psdf["this"] <= byte_psdf["this"]) + + def test_gt(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] > byte_pdf["that"], byte_psdf["this"] > byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] > byte_pdf["this"], byte_psdf["this"] > byte_psdf["this"]) + + def test_ge(self): + byte_pdf, byte_psdf = self.byte_pdf, self.byte_psdf + self.assert_eq(byte_pdf["this"] >= byte_pdf["that"], byte_psdf["this"] >= byte_psdf["that"]) + self.assert_eq(byte_pdf["this"] >= byte_pdf["this"], byte_psdf["this"] >= byte_psdf["this"]) + + +class BinaryOpsTests(BinaryOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_binary_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py new file mode 100644 index 0000000..2226756 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py @@ -0,0 +1,825 @@ +# +# 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 datetime +import unittest +from distutils.version import LooseVersion + +import pandas as pd +import numpy as np +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas import option_context +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.pandas.typedef.typehints import ( + extension_float_dtypes_available, + extension_object_dtypes_available, +) + + +class BooleanOpsTestsMixin: + @property + def bool_pdf(self): + return pd.DataFrame({"this": [True, False, True], "that": [False, True, True]}) + + @property + def bool_psdf(self): + return ps.from_pandas(self.bool_pdf) + + def test_add(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(b_pser + 1, b_psser + 1) + self.assert_eq(b_pser + 0.1, b_psser + 0.1) + self.assert_eq(b_pser + b_pser.astype(int), b_psser + b_psser.astype(int)) + self.assert_eq(b_pser + b_pser, b_psser + b_psser) + self.assert_eq(b_pser + True, b_psser + True) + self.assert_eq(b_pser + False, b_psser + False) + + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(b_pser + pser, b_psser + psser) + for col in self.non_numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if col == "bool": + self.assert_eq(b_pser + pser, b_psser + psser) + else: + self.assertRaises(TypeError, lambda: b_psser + psser) + + def test_sub(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(b_pser - 1, b_psser - 1) + self.assert_eq(b_pser - 0.1, b_psser - 0.1) + self.assert_eq(b_pser - b_pser.astype(int), b_psser - b_psser.astype(int)) + self.assertRaises(TypeError, lambda: b_psser - b_psser) + self.assertRaises(TypeError, lambda: b_psser - True) + + for col in self.numeric_df_cols: + self.assert_eq(b_pser - pdf[col], b_psser - psdf[col]) + + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: b_psser - psdf[col]) + + def test_mul(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(b_pser * 1, b_psser * 1) + self.assert_eq(b_pser * 0.1, b_psser * 0.1) + self.assert_eq(b_pser * b_pser.astype(int), b_psser * b_psser.astype(int)) + self.assert_eq(b_pser * b_pser, b_psser * b_psser) + self.assert_eq(b_pser * True, b_psser * True) + self.assert_eq(b_pser * False, b_psser * False) + + for col in self.numeric_df_cols: + self.assert_eq(b_pser * pdf[col], b_psser * psdf[col]) + + for col in self.non_numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if col == "bool": + self.assert_eq(b_pser * pser, b_psser * psser) + else: + self.assertRaises(TypeError, lambda: b_psser * psser) + + def test_truediv(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(b_pser / 1, b_psser / 1) + self.assert_eq(b_pser / 0.1, b_psser / 0.1) + self.assert_eq(b_pser / b_pser.astype(int), b_psser / b_psser.astype(int)) + self.assertRaises(TypeError, lambda: b_psser / b_psser) + self.assertRaises(TypeError, lambda: b_psser / True) + + self.assert_eq(b_pser / pdf["float"], b_psser / psdf["float"]) + + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: b_psser / psdf[col]) + + def test_floordiv(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + + # float is always returned in pandas-on-Spark + self.assert_eq((b_pser // 1).astype("float"), b_psser // 1) + + # in pandas, 1 // 0.1 = 9.0; in pandas-on-Spark, 1 // 0.1 = 10.0 + # self.assert_eq(b_pser // 0.1, b_psser // 0.1) + + self.assert_eq(b_pser // b_pser.astype(int), b_psser // b_psser.astype(int)) + self.assertRaises(TypeError, lambda: b_psser // b_psser) + self.assertRaises(TypeError, lambda: b_psser // True) + + self.assert_eq(b_pser // pdf["float"], b_psser // psdf["float"]) + + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: b_psser // psdf[col]) + + def test_mod(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(b_pser % 1, b_psser % 1) + self.assert_eq(b_pser % 0.1, b_psser % 0.1) + self.assert_eq(b_pser % b_pser.astype(float), b_psser % b_psser.astype(float)) + self.assertRaises(TypeError, lambda: b_psser % b_psser) + self.assertRaises(TypeError, lambda: b_psser % True) + + for col in self.numeric_df_cols: + self.assert_eq(b_pser % pdf[col], b_psser % psdf[col]) + + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: b_psser % psdf[col]) + + def test_pow(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + # float is always returned in pandas-on-Spark + self.assert_eq((b_pser**1).astype("float"), b_psser**1) + self.assert_eq(b_pser**0.1, b_psser**0.1) + self.assert_eq(b_pser ** b_pser.astype(float), b_psser ** b_psser.astype(float)) + self.assertRaises(TypeError, lambda: b_psser**b_psser) + self.assertRaises(TypeError, lambda: b_psser**True) + + self.assert_eq(b_pser % pdf["float"], b_psser % psdf["float"]) + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: b_psser % psdf[col]) + + def test_radd(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(1 + b_pser, 1 + b_psser) + self.assert_eq(0.1 + b_pser, 0.1 + b_psser) + self.assert_eq(True + b_pser, True + b_psser) + self.assert_eq(False + b_pser, False + b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) + b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) + b_psser) + + def test_rsub(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(1 - b_pser, 1 - b_psser) + self.assert_eq(0.1 - b_pser, 0.1 - b_psser) + self.assertRaises(TypeError, lambda: "x" - b_psser) + self.assertRaises(TypeError, lambda: True - b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) - b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) - b_psser) + + def test_rmul(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(1 * b_pser, 1 * b_psser) + self.assert_eq(0.1 * b_pser, 0.1 * b_psser) + self.assertRaises(TypeError, lambda: "x" * b_psser) + self.assert_eq(True * b_pser, True * b_psser) + self.assert_eq(False * b_pser, False * b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) * b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) * b_psser) + + def test_rtruediv(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(1 / b_pser, 1 / b_psser) + self.assert_eq(0.1 / b_pser, 0.1 / b_psser) + self.assertRaises(TypeError, lambda: "x" / b_psser) + self.assertRaises(TypeError, lambda: True / b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) / b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) / b_psser) + + def test_rfloordiv(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + self.assert_eq(1 // b_pser, 1 // b_psser) + self.assert_eq(0.1 // b_pser, 0.1 // b_psser) + self.assertRaises(TypeError, lambda: "x" // b_psser) + self.assertRaises(TypeError, lambda: True // b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) // b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) // b_psser) + + def test_rpow(self): + pdf, psdf = self.pdf, self.psdf + + b_pser, b_psser = pdf["bool"], psdf["bool"] + # float is returned always in pandas-on-Spark + self.assert_eq((1**b_pser).astype(float), 1**b_psser) + self.assert_eq(0.1**b_pser, 0.1**b_psser) + self.assertRaises(TypeError, lambda: "x" ** b_psser) + self.assertRaises(TypeError, lambda: True**b_psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) ** b_psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) ** b_psser) + + def test_rmod(self): + psdf = self.psdf + + b_psser = psdf["bool"] + # 1 % False is 0.0 in pandas + self.assert_eq(pd.Series([0, 0, None], dtype=float, name="bool"), 1 % b_psser) + # 0.1 / True is 0.1 in pandas + self.assert_eq( + pd.Series([0.10000000000000009, 0.10000000000000009, None], dtype=float, name="bool"), + 0.1 % b_psser, + ) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) % b_psser) + self.assertRaises(TypeError, lambda: True % b_psser) + + def test_and(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + + self.assert_eq(pser & True, psser & True) + self.assert_eq(pser & False, psser & False) + self.assert_eq(pser & pser, psser & psser) + self.assert_eq(pser & other_pser, psser & other_psser) + self.assert_eq(other_pser & pser, other_psser & psser) + + def test_rand(self): + pser, psser = self.pdf["bool"], self.psdf["bool"] + self.assert_eq(True & pser, True & psser) + self.assert_eq(False & pser, False & psser) + + def test_or(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + + self.assert_eq(pser | True, psser | True) + self.assert_eq(pser | False, psser | False) + self.assert_eq(pser | pser, psser | psser) + self.assert_eq(True | pser, True | psser) + self.assert_eq(False | pser, False | psser) + + self.assert_eq(pser | other_pser, psser | other_psser) + self.assert_eq(other_pser | pser, other_psser | psser) + + def test_ror(self): + pser, psser = self.pdf["bool"], self.psdf["bool"] + self.assert_eq(True | pser, True | psser) + self.assert_eq(False | pser, False | psser) + + def test_xor(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + + self.assert_eq(pser ^ other_pser, psser ^ other_psser) + self.assert_eq(pser ^ True, psser ^ True) + self.assert_eq(pser ^ False, psser ^ False) + self.assert_eq(pser ^ 2, psser ^ 2) + self.assert_eq(pser ^ 99, psser ^ 99) + + with self.assertRaisesRegex(TypeError, "XOR can not be applied to given types."): + psser ^ "a" + + with option_context("compute.ops_on_diff_frames", True): + pser, other_pser = self.pdf["bool"], self.integral_pdf["this"] + psser, other_psser = self.psdf["bool"], self.integral_psdf["this"] + + self.assert_eq(pser ^ other_pser, psser ^ other_psser) + + def test_rxor(self): + pser, psser = self.pdf["bool"], self.psdf["bool"] + self.assert_eq(True ^ pser, True ^ psser) + self.assert_eq(False ^ pser, False ^ psser) + self.assert_eq(1 ^ pser, 1 ^ psser) + + def test_isnull(self): + self.assert_eq(self.pdf["bool"].isnull(), self.psdf["bool"].isnull()) + + def test_astype(self): + pser, psser = self.pdf["bool"], self.psdf["bool"] + self.assert_eq(pser.astype(int), psser.astype(int)) + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + self.assert_eq(pser.astype(np.int32), psser.astype(np.int32)) + self.assert_eq(pser.astype(np.int16), psser.astype(np.int16)) + self.assert_eq(pser.astype(np.int8), psser.astype(np.int8)) + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=[False, True]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + def test_neg(self): + self.assert_eq(-self.pdf["bool"], -self.psdf["bool"]) + + def test_abs(self): + self.assert_eq(abs(self.pdf["bool"]), abs(self.psdf["bool"])) + + def test_invert(self): + self.assert_eq(~self.pdf["bool"], ~self.psdf["bool"]) + + def test_eq(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser == other_pser, psser == other_psser) + self.assert_eq(pser == pser, psser == psser) + + def test_ne(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser != other_pser, psser != other_psser) + self.assert_eq(pser != pser, psser != psser) + + def test_lt(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser < other_pser, psser < other_psser) + self.assert_eq(pser < pser, psser < psser) + + def test_le(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser <= other_pser, psser <= other_psser) + self.assert_eq(pser <= pser, psser <= psser) + + def test_gt(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser > other_pser, psser > other_psser) + self.assert_eq(pser > pser, psser > psser) + + def test_ge(self): + pdf, psdf = self.bool_pdf, self.bool_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + self.assert_eq(pser >= other_pser, psser >= other_psser) + self.assert_eq(pser >= pser, psser >= psser) + + +@unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" +) +class BooleanExtensionOpsTest(OpsTestBase): + @property + def boolean_pdf(self): + return pd.DataFrame( + {"this": [True, False, None], "that": [False, None, True]}, dtype="boolean" + ) + + @property + def boolean_psdf(self): + return ps.from_pandas(self.boolean_pdf) + + @property + def boolean_numeric_pdf(self): + return pd.concat([self.boolean_pdf, self.numeric_pdf], axis=1) + + @property + def boolean_numeric_psdf(self): + return ps.from_pandas(self.boolean_numeric_pdf) + + @property + def boolean_non_numeric_pdf(self): + return pd.concat([self.boolean_pdf, self.non_numeric_pdf], axis=1) + + @property + def boolean_non_numeric_psdf(self): + return ps.from_pandas(self.boolean_non_numeric_pdf) + + def test_add(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(pser + 1, psser + 1) + if extension_float_dtypes_available: + self.check_extension(pser + 0.1, psser + 0.1) + else: + self.assert_eq(pser + 0.1, psser + 0.1) + + # In pandas, NA | True is NA, whereas NA | True is True in pandas-on-Spark + self.check_extension( + pd.Series([True, True, True], dtype="boolean", name=psser.name), psser + True + ) + self.check_extension(pser + False, psser + False) + self.check_extension(pser + pser, psser + psser) + + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf["this"] + pdf[col], psdf["this"] + psdf[col], almost=True) + + pdf, psdf = self.boolean_non_numeric_pdf, self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + if col == "bool": + self.check_extension(pdf["this"] + pdf[col], psdf["this"] + psdf[col]) + else: + self.assertRaises(TypeError, lambda: psdf["this"] + psdf[col]) + + def test_sub(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(pser - 1, psser - 1) + if extension_float_dtypes_available: + self.check_extension(pser - 0.1, psser - 0.1) + else: + self.assert_eq(pser - 0.1, psser - 0.1) + self.assertRaises(TypeError, lambda: psser - psser) + self.assertRaises(TypeError, lambda: psser - True) + + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf["this"] - pdf[col], psdf["this"] - psdf[col], almost=True) + + pdf, psdf = self.boolean_non_numeric_pdf, self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: psdf["this"] - psdf[col]) + + def test_mul(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(pser * 1, psser * 1) + if extension_float_dtypes_available: + self.check_extension(pser * 0.1, psser * 0.1) + else: + self.assert_eq(pser * 0.1, psser * 0.1) + + # In pandas, NA & False is NA, whereas NA & False is False in pandas-on-Spark + self.check_extension(pser * True, psser * True) + self.check_extension( + pd.Series([False, False, False], dtype="boolean", name=psser.name), psser * False + ) + self.check_extension(pser * pser, psser * psser) + + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf["this"] * pdf[col], psdf["this"] * psdf[col], almost=True) + + pdf, psdf = self.boolean_non_numeric_pdf, self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + if col == "bool": + self.check_extension(pdf["that"] * pdf[col], psdf["that"] * psdf[col]) + else: + self.assertRaises(TypeError, lambda: psdf["this"] * psdf[col]) + + def test_truediv(self): + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + if extension_float_dtypes_available: + self.check_extension(pser / 1, psser / 1) + self.check_extension(pser / 0.1, psser / 0.1) + else: + self.assert_eq(pser / 1, psser / 1) + self.assert_eq(pser / 0.1, psser / 0.1) + self.assertRaises(TypeError, lambda: psser / psser) + self.assertRaises(TypeError, lambda: psser / True) + + self.assert_eq( + pser / pdf["float"], + psser / psdf["float"], + almost=True, + ) + psdf = self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: psdf["this"] / psdf[col]) + + def test_floordiv(self): + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + + # float is always returned in pandas-on-Spark + if extension_float_dtypes_available: + self.check_extension((pser // 1).astype("Float64"), psser // 1) + else: + self.assert_eq((pser // 1).astype("float"), psser // 1) + + # in pandas, 1 // 0.1 = 9.0; in pandas-on-Spark, 1 // 0.1 = 10.0 + # self.assert_eq(pser // 0.1, psser // 0.1) + + self.assertRaises(TypeError, lambda: psser // psser) + self.assertRaises(TypeError, lambda: psser // True) + + self.assert_eq( + pser // pdf["float"], + psser // psdf["float"], + almost=True, + ) + psdf = self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: psdf["this"] // psdf[col]) + + def test_mod(self): + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + self.check_extension(pser % 1, psser % 1) + if extension_float_dtypes_available: + self.check_extension(pser % 0.1, psser % 0.1) + else: + self.assert_eq(pser % 0.1, psser % 0.1) + self.assertRaises(TypeError, lambda: psser % psser) + self.assertRaises(TypeError, lambda: psser % True) + + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf["this"] % pdf[col], psdf["this"] % psdf[col], almost=True) + + psdf = self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: psdf["this"] % psdf[col]) + + def test_pow(self): + pdf, psdf = self.boolean_numeric_pdf, self.boolean_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + # float is always returned in pandas-on-Spark + if extension_float_dtypes_available: + self.check_extension((pser**1).astype("Float64"), psser**1) + self.check_extension((pser**0.1).astype("Float64"), psser**0.1) + self.check_extension( + (pser ** pser.astype(float)).astype("Float64"), psser ** psser.astype(float) + ) + else: + self.assert_eq((pser**1).astype("float"), psser**1) + self.assert_eq((pser**0.1).astype("float"), psser**0.1) + self.assert_eq( + (pser ** pser.astype(float)).astype("float"), psser ** psser.astype(float) + ) + self.assertRaises(TypeError, lambda: psser**psser) + self.assertRaises(TypeError, lambda: psser**True) + + self.assert_eq( + pser ** pdf["float"], + psser ** psdf["float"], + almost=True, + ) + psdf = self.boolean_non_numeric_psdf + for col in self.non_numeric_df_cols: + self.assertRaises(TypeError, lambda: psdf["this"] % psdf[col]) + + def test_radd(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + + self.check_extension(1 + pser, 1 + psser) + if extension_float_dtypes_available: + self.check_extension(0.1 + pser, 0.1 + psser) + else: + self.assert_eq(0.1 + pser, 0.1 + psser) + self.assertRaises(TypeError, lambda: "x" + psser) + + # In pandas, NA | True is NA, whereas NA | True is True in pandas-on-Spark + self.check_extension( + ps.Series([True, True, True], dtype="boolean", name=psser.name), True + psser + ) + self.check_extension(False + pser, False + psser) + + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) + psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) + psser) + + def test_rsub(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(1 - pser, 1 - psser) + if extension_float_dtypes_available: + self.check_extension(0.1 - pser, 0.1 - psser) + else: + self.assert_eq(0.1 - pser, 0.1 - psser) + self.assertRaises(TypeError, lambda: "x" - psser) + self.assertRaises(TypeError, lambda: True - psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) - psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) - psser) + + def test_rmul(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(1 * pser, 1 * psser) + if extension_float_dtypes_available: + self.check_extension(0.1 * pser, 0.1 * psser) + else: + self.assert_eq(0.1 * pser, 0.1 * psser) + self.assertRaises(TypeError, lambda: "x" * psser) + + # In pandas, NA & False is NA, whereas NA & False is False in pandas-on-Spark + self.check_extension(True * pser, True * psser) + self.check_extension( + pd.Series([False, False, False], dtype="boolean", name=psser.name), False * psser + ) + + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) * psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) * psser) + + def test_rtruediv(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + if extension_float_dtypes_available: + self.check_extension(1 / pser, 1 / psser) + self.check_extension(0.1 / pser, 0.1 / psser) + else: + self.assert_eq(1 / pser, 1 / psser) + self.assert_eq(0.1 / pser, 0.1 / psser) + self.assertRaises(TypeError, lambda: "x" / psser) + self.assertRaises(TypeError, lambda: True / psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) / psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) / psser) + + def test_rfloordiv(self): + psser = self.boolean_psdf["this"] + self.assert_eq( + pd.Series([1.0, np.inf, np.nan], name=psser.name), (1 // psser).astype(float) + ) + self.assert_eq( + pd.Series([0.0, np.inf, np.nan], name=psser.name), (0.1 // psser).astype(float) + ) + self.assertRaises(TypeError, lambda: "x" // psser) + self.assertRaises(TypeError, lambda: True // psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) // psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) // psser) + + def test_rpow(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + if extension_float_dtypes_available: + self.check_extension(pd.Series([1, 1, 1], dtype="Float64", name=psser.name), 1**psser) + self.check_extension((0.1**pser).astype("Float64"), 0.1**psser) + else: + self.assert_eq(pd.Series([1, 1, 1], dtype="float", name=psser.name), 1**psser) + self.assert_eq((0.1**pser).astype("float"), 0.1**psser) + self.assertRaises(TypeError, lambda: "x" ** psser) + self.assertRaises(TypeError, lambda: True**psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) ** psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) ** psser) + + def test_rmod(self): + psser = self.boolean_psdf["this"] + self.check_extension( + ps.Series([0, np.nan, np.nan], dtype="Int64", name=psser.name), 1 % psser + ) + if extension_float_dtypes_available: + self.check_extension( + pd.Series([0.10000000000000009, np.nan, np.nan], dtype="Float64", name=psser.name), + 0.1 % psser, + ) + else: + self.assert_eq( + pd.Series([0.10000000000000009, np.nan, np.nan], dtype="float", name=psser.name), + 0.1 % psser, + ) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) % psser) + self.assertRaises(TypeError, lambda: True % psser) + + def test_and(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser & True, psser & True) + self.check_extension(pser & False, psser & False) + self.check_extension(pser & pser, psser & psser) + + self.check_extension(pser & other_pser, psser & other_psser) + self.check_extension(other_pser & pser, other_psser & psser) + + def test_rand(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(True & pser, True & psser) + self.check_extension(False & pser, False & psser) + + def test_or(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser | True, psser | True) + self.check_extension(pser | False, psser | False) + self.check_extension(pser | pser, psser | psser) + + self.check_extension(pser | other_pser, psser | other_psser) + self.check_extension(other_pser | pser, other_psser | psser) + + def test_ror(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(True | pser, True | psser) + self.check_extension(False | pser, False | psser) + + def test_xor(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser ^ True, psser ^ True) + self.check_extension(pser ^ False, psser ^ False) + self.check_extension(pser ^ pser, psser ^ psser) + + self.check_extension(pser ^ other_pser, psser ^ other_psser) + self.check_extension(other_pser ^ pser, other_psser ^ psser) + with self.assertRaisesRegex(TypeError, "XOR can not be applied to given types."): + psser ^ 2 + + def test_rxor(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + self.check_extension(True | pser, True | psser) + self.check_extension(False | pser, False | psser) + with self.assertRaisesRegex(TypeError, "XOR can not be applied to given types."): + 1 ^ psser + + def test_from_to_pandas(self): + data = [True, True, False, None] + pser = pd.Series(data, dtype="boolean") + psser = ps.Series(data, dtype="boolean") + self.check_extension(pser, psser._to_pandas()) + self.check_extension(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.boolean_pdf["this"].isnull(), self.boolean_psdf["this"].isnull()) + + def test_astype(self): + pser, psser = self.boolean_pdf["this"], self.boolean_psdf["this"] + + self.assert_eq(pser.astype(str).tolist(), psser.astype(str).tolist()) + + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=[False, True]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + for dtype in self.extension_dtypes: + if dtype in self.fractional_extension_dtypes: + # A pandas boolean extension series cannot be casted to fractional extension dtypes + self.assert_eq([1.0, 0.0, np.nan], psser.astype(dtype).tolist()) + elif dtype in self.string_extension_dtype: + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + # Limit pandas version due to https://github.com/pandas-dev/pandas/issues/31204 + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + else: + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.boolean_psdf["this"]) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.boolean_psdf["this"])) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.boolean_psdf["this"]) + + def test_eq(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser == other_pser, psser == other_psser) + self.check_extension(pser == pser, psser == psser) + + def test_ne(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser != other_pser, psser != other_psser) + self.check_extension(pser != pser, psser != psser) + + def test_lt(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser < other_pser, psser < other_psser) + self.check_extension(pser < pser, psser < psser) + + def test_le(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser <= other_pser, psser <= other_psser) + self.check_extension(pser <= pser, psser <= psser) + + def test_gt(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser > other_pser, psser > other_psser) + self.check_extension(pser > pser, psser > psser) + + def test_ge(self): + pdf, psdf = self.boolean_pdf, self.boolean_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.check_extension(pser >= other_pser, psser >= other_psser) + self.check_extension(pser >= pser, psser >= psser) + + +class BooleanOpsTests(BooleanOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.data_type_ops.test_boolean_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py new file mode 100644 index 0000000..e56fce4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py @@ -0,0 +1,562 @@ +# +# 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 distutils.version import LooseVersion + +import pandas as pd +import numpy as np +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class CategoricalOpsTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "this_numeric_cat": pd.Series([1, 2, 3], dtype="category"), + "that_numeric_cat": pd.Series([3, 2, 1], dtype="category"), + "this_ordered_numeric_cat": pd.Categorical( + [1, 2, 3], categories=[3, 2, 1], ordered=True + ), + "that_ordered_numeric_cat": pd.Categorical( + [2, 3, 1], categories=[3, 2, 1], ordered=True + ), + "this_string_cat": pd.Series(["x", "y", "z"], dtype="category"), + "that_string_cat": pd.Series(["z", "y", "x"], dtype="category"), + "this_ordered_string_cat": pd.Categorical( + ["x", "y", "z"], categories=["x", "z", "y"], ordered=True + ), + "that_ordered_string_cat": pd.Categorical( + ["z", "y", "x"], categories=["x", "z", "y"], ordered=True + ), + "this_given_cat_string_cat": pd.Series( + pd.Categorical(["x", "y", "z"], categories=list("zyx")) + ), + } + ) + + @property + def pser(self): + return pd.Series([1, 2, 3], dtype="category") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser - psser) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + with option_context("compute.ops_on_diff_frames", True): + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psser & True) + self.assertRaises(TypeError, lambda: self.psser & False) + self.assertRaises(TypeError, lambda: self.psser & self.psser) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psser) + self.assertRaises(TypeError, lambda: False & self.psser) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psser | True) + self.assertRaises(TypeError, lambda: self.psser | False) + self.assertRaises(TypeError, lambda: self.psser | self.psser) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psser) + self.assertRaises(TypeError, lambda: False | self.psser) + + def test_from_to_pandas(self): + data = [1, "x", "y"] + pser = pd.Series(data, dtype="category") + psser = ps.Series(data, dtype="category") + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + data = [1, 2, 3] + pser = pd.Series(data, dtype="category") + psser = ps.from_pandas(pser) + self.assert_eq(pser.astype(int), psser.astype(int)) + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + self.assert_eq(pser.astype(np.int32), psser.astype(np.int32)) + self.assert_eq(pser.astype(np.int16), psser.astype(np.int16)) + self.assert_eq(pser.astype(np.int8), psser.astype(np.int8)) + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + + cat_type = CategoricalDtype(categories=[3, 1, 2]) + # CategoricalDtype is not updated if the dtype is same from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + else: + self.assert_eq(psser.astype(cat_type), pser) + + # Empty + pser = pd.Series([], dtype="category") + psser = ps.from_pandas(pser) + self.assert_eq(pser.astype(int), psser.astype(int)) + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + self.assert_eq(pser.astype(np.int32), psser.astype(np.int32)) + self.assert_eq(pser.astype(np.int16), psser.astype(np.int16)) + self.assert_eq(pser.astype(np.int8), psser.astype(np.int8)) + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + self.assert_eq(pser.astype("category"), psser.astype("category")) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.pdf, self.psdf + + pser, psser = pdf["this_numeric_cat"], psdf["this_numeric_cat"] + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser == 1, ordered_psser == 1) + self.assert_eq(pser == pser, psser == psser) + self.assert_eq(ordered_pser == ordered_pser, ordered_psser == ordered_psser) + + pser, psser = pdf["this_string_cat"], psdf["this_string_cat"] + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(pser == "x", psser == "x") + self.assert_eq(pser == pser, psser == psser) + self.assert_eq(ordered_pser == ordered_pser, ordered_psser == ordered_psser) + + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser == 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser == "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser == ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list", + lambda: ordered_psser == [1, 2, 3], + ) + + self.assert_eq( + pdf["this_numeric_cat"] == pdf["that_numeric_cat"], + psdf["this_numeric_cat"] == psdf["that_numeric_cat"], + ) + self.assert_eq( + pdf["this_string_cat"] == pdf["that_string_cat"], + psdf["this_string_cat"] == psdf["that_string_cat"], + ) + + self.assert_eq( + pdf["this_string_cat"] == pdf["this_given_cat_string_cat"], + psdf["this_string_cat"] == psdf["this_given_cat_string_cat"], + ) + + pser1 = pd.Series(pd.Categorical(list("abca"))) + pser2 = pd.Series(pd.Categorical(list("bcaa"), categories=list("bca"))) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + with option_context("compute.ops_on_diff_frames", True): + self.assert_eq(pser1 == pser2, (psser1 == psser2).sort_index()) + + psser3 = ps.Series(pd.Categorical(list("xyzx"))) + self.assertRaisesRegex( + TypeError, + "Categoricals can only be compared if 'categories' are the same.", + lambda: psser1 == psser3, + ) + + def test_ne(self): + pdf, psdf = self.pdf, self.psdf + + pser, psser = pdf["this_numeric_cat"], psdf["this_numeric_cat"] + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser != 1, ordered_psser != 1) + self.assert_eq(pser != pser, psser != psser) + self.assert_eq(ordered_pser != ordered_pser, ordered_psser != ordered_psser) + + pser, psser = pdf["this_string_cat"], psdf["this_string_cat"] + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(pser != "x", psser != "x") + self.assert_eq(pser != pser, psser != psser) + self.assert_eq(ordered_pser != ordered_pser, ordered_psser != ordered_psser) + + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser != 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser != "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser != ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list.", + lambda: ordered_psser != [1, 2, 3], + ) + self.assert_eq( + pdf["this_numeric_cat"] != pdf["that_numeric_cat"], + psdf["this_numeric_cat"] != psdf["that_numeric_cat"], + ) + self.assert_eq( + pdf["this_string_cat"] != pdf["that_string_cat"], + psdf["this_string_cat"] != psdf["that_string_cat"], + ) + self.assert_eq( + pdf["this_string_cat"] != pdf["this_given_cat_string_cat"], + psdf["this_string_cat"] != psdf["this_given_cat_string_cat"], + ) + + pser1 = pd.Series(pd.Categorical(list("abca"))) + pser2 = pd.Series(pd.Categorical(list("bcaa"), categories=list("bca"))) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + with option_context("compute.ops_on_diff_frames", True): + self.assert_eq(pser1 != pser2, (psser1 != psser2).sort_index()) + + def test_lt(self): + pdf, psdf = self.pdf, self.psdf + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser < 1, ordered_psser < 1) + self.assert_eq(ordered_pser < ordered_pser, ordered_psser < ordered_psser) + + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(ordered_pser < "x", ordered_psser < "x") + self.assert_eq(ordered_pser < ordered_pser, ordered_psser < ordered_psser) + + self.assertRaisesRegex( + TypeError, + "Unordered Categoricals can only compare equality or not", + lambda: self.psser < ordered_psser, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser < 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser < "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser < ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list", + lambda: ordered_psser < [1, 2, 3], + ) + self.assert_eq( + pdf["this_ordered_numeric_cat"] < pdf["that_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"] < psdf["that_ordered_numeric_cat"], + ) + self.assert_eq( + pdf["this_ordered_string_cat"] < pdf["that_ordered_string_cat"], + psdf["this_ordered_string_cat"] < psdf["that_ordered_string_cat"], + ) + + def test_le(self): + pdf, psdf = self.pdf, self.psdf + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser <= 1, ordered_psser <= 1) + self.assert_eq(ordered_pser <= ordered_pser, ordered_psser <= ordered_psser) + + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(ordered_pser <= "x", ordered_psser <= "x") + self.assert_eq(ordered_pser <= ordered_pser, ordered_psser <= ordered_psser) + self.assertRaisesRegex( + TypeError, + "Unordered Categoricals can only compare equality or not", + lambda: self.psser <= ordered_psser, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser <= 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser <= "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser <= ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list", + lambda: ordered_psser <= [1, 2, 3], + ) + self.assert_eq( + pdf["this_ordered_numeric_cat"] <= pdf["that_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"] <= psdf["that_ordered_numeric_cat"], + ) + self.assert_eq( + pdf["this_ordered_string_cat"] <= pdf["that_ordered_string_cat"], + psdf["this_ordered_string_cat"] <= psdf["that_ordered_string_cat"], + ) + + def test_gt(self): + pdf, psdf = self.pdf, self.psdf + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser > 1, ordered_psser > 1) + self.assert_eq(ordered_pser > ordered_pser, ordered_psser > ordered_psser) + + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(ordered_pser > "x", ordered_psser > "x") + self.assert_eq(ordered_pser > ordered_pser, ordered_psser > ordered_psser) + self.assertRaisesRegex( + TypeError, + "Unordered Categoricals can only compare equality or not", + lambda: self.psser > ordered_psser, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser > 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser > "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser > ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list", + lambda: ordered_psser > [1, 2, 3], + ) + self.assert_eq( + pdf["this_ordered_numeric_cat"] > pdf["that_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"] > psdf["that_ordered_numeric_cat"], + ) + self.assert_eq( + pdf["this_ordered_string_cat"] > pdf["that_ordered_string_cat"], + psdf["this_ordered_string_cat"] > psdf["that_ordered_string_cat"], + ) + + def test_ge(self): + pdf, psdf = self.pdf, self.psdf + ordered_pser, ordered_psser = ( + pdf["this_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"], + ) + self.assert_eq(ordered_pser >= 1, ordered_psser >= 1) + self.assert_eq(ordered_pser >= ordered_pser, ordered_psser >= ordered_psser) + + ordered_pser, ordered_psser = ( + pdf["this_ordered_string_cat"], + psdf["this_ordered_string_cat"], + ) + self.assert_eq(ordered_pser >= "x", ordered_psser >= "x") + self.assert_eq(ordered_pser >= ordered_pser, ordered_psser >= ordered_psser) + self.assertRaisesRegex( + TypeError, + "Unordered Categoricals can only compare equality or not", + lambda: self.psser >= ordered_psser, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser >= 4, + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with a scalar, which is not a category", + lambda: ordered_psser >= "a", + ) + self.assertRaisesRegex( + TypeError, + "Cannot compare a Categorical with the given type", + lambda: ordered_psser >= ps.Series([1, 2, 3]), + ) + self.assertRaisesRegex( + TypeError, + "The operation can not be applied to list", + lambda: ordered_psser >= [1, 2, 3], + ) + self.assert_eq( + pdf["this_ordered_numeric_cat"] >= pdf["that_ordered_numeric_cat"], + psdf["this_ordered_numeric_cat"] >= psdf["that_ordered_numeric_cat"], + ) + self.assert_eq( + pdf["this_ordered_string_cat"] >= pdf["that_ordered_string_cat"], + psdf["this_ordered_string_cat"] >= psdf["that_ordered_string_cat"], + ) + + +class CategoricalOpsTests(CategoricalOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_categorical_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py new file mode 100644 index 0000000..f7c6642 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py @@ -0,0 +1,368 @@ +# +# 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 decimal +import datetime + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class ComplexOpsTestsMixin: + @property + def pser(self): + return pd.Series([[1, 2, 3]]) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def numeric_array_pdf(self): + psers = { + "int": pd.Series([[1, 2, 3]]), + "float": pd.Series([[0.1, 0.2, 0.3]]), + "decimal": pd.Series([[decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)]]), + } + return pd.concat(psers, axis=1) + + @property + def numeric_array_psdf(self): + return ps.from_pandas(self.numeric_array_pdf) + + @property + def numeric_array_df_cols(self): + return self.numeric_array_pdf.columns + + @property + def non_numeric_array_pdf(self): + psers = { + "string": pd.Series([["x", "y", "z"]]), + "date": pd.Series( + [[datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)]] + ), + "bool": pd.Series([[True, True, False]]), + } + return pd.concat(psers, axis=1) + + @property + def non_numeric_array_psdf(self): + return ps.from_pandas(self.non_numeric_array_pdf) + + @property + def non_numeric_array_df_cols(self): + return self.non_numeric_array_pdf.columns + + @property + def array_pdf(self): + return pd.concat([self.numeric_array_pdf, self.non_numeric_array_pdf], axis=1) + + @property + def array_psdf(self): + return ps.from_pandas(self.array_pdf) + + @property + def array_df_cols(self): + return self.array_pdf.columns + + @property + def complex_pdf(self): + psers = { + "this_array": self.pser, + "that_array": pd.Series([[2, 3, 4]]), + "this_struct": pd.Series([("x", 1)]), + "that_struct": pd.Series([("a", 2)]), + } + return pd.concat(psers, axis=1) + + @property + def complex_psdf(self): + pssers = { + "this_array": self.psser, + "that_array": ps.Series([[2, 3, 4]]), + "this_struct": ps.Index([("x", 1)]).to_series().reset_index(drop=True), + "that_struct": ps.Index([("a", 2)]).to_series().reset_index(drop=True), + } + return ps.concat(pssers, axis=1) + + def test_add(self): + pdf, psdf = self.array_pdf, self.array_psdf + for col in self.array_df_cols: + self.assert_eq(pdf[col] + pdf[col], psdf[col] + psdf[col]) + + # Numeric array + Numeric array + for col in self.numeric_array_df_cols: + pser1, psser1 = pdf[col], psdf[col] + for other_col in self.numeric_array_df_cols: + pser2, psser2 = pdf[other_col], psdf[other_col] + self.assert_eq((pser1 + pser2).sort_values(), (psser1 + psser2).sort_values()) + + # Non-numeric array + Non-numeric array + self.assertRaises( + TypeError, + lambda: psdf["string"] + psdf["bool"], + ) + self.assertRaises( + TypeError, + lambda: psdf["string"] + psdf["date"], + ) + self.assertRaises( + TypeError, + lambda: psdf["bool"] + psdf["date"], + ) + + for col in self.non_numeric_array_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser + pser, psser + psser) + + # Numeric array + Non-numeric array + for numeric_col in self.numeric_array_df_cols: + for non_numeric_col in self.non_numeric_array_df_cols: + self.assertRaises(TypeError, lambda: psdf[numeric_col] + psdf[non_numeric_col]) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] - psdf[other_col]) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] * psdf[other_col]) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] / psdf[other_col]) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] // psdf[other_col]) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] % psdf[other_col]) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + psdf = self.array_psdf + for col in self.array_df_cols: + for other_col in self.array_df_cols: + self.assertRaises(TypeError, lambda: psdf[col] ** psdf[other_col]) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psser & True) + self.assertRaises(TypeError, lambda: self.psser & False) + self.assertRaises(TypeError, lambda: self.psser & self.psser) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psser) + self.assertRaises(TypeError, lambda: False & self.psser) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psser | True) + self.assertRaises(TypeError, lambda: self.psser | False) + self.assertRaises(TypeError, lambda: self.psser | self.psser) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psser) + self.assertRaises(TypeError, lambda: False | self.psser) + + def test_from_to_pandas(self): + pdf, psdf = self.array_pdf, self.array_psdf + for col in self.array_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + pdf, psdf = self.array_pdf, self.array_psdf + for col in self.array_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser.isnull(), psser.isnull()) + + def test_astype(self): + self.assert_eq(self.pser.astype(str), self.psser.astype(str)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] == pdf["that_array"], psdf["this_array"] == psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] == pdf["that_struct"], psdf["this_struct"] == psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] == pdf["this_array"], psdf["this_array"] == psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] == pdf["this_struct"], psdf["this_struct"] == psdf["this_struct"] + ) + + def test_ne(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] != pdf["that_array"], psdf["this_array"] != psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] != pdf["that_struct"], psdf["this_struct"] != psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] != pdf["this_array"], psdf["this_array"] != psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] != pdf["this_struct"], psdf["this_struct"] != psdf["this_struct"] + ) + + def test_lt(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] < pdf["that_array"], psdf["this_array"] < psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] < pdf["that_struct"], psdf["this_struct"] < psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] < pdf["this_array"], psdf["this_array"] < psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] < pdf["this_struct"], psdf["this_struct"] < psdf["this_struct"] + ) + + def test_le(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] <= pdf["that_array"], psdf["this_array"] <= psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] <= pdf["that_struct"], psdf["this_struct"] <= psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] <= pdf["this_array"], psdf["this_array"] <= psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] <= pdf["this_struct"], psdf["this_struct"] <= psdf["this_struct"] + ) + + def test_gt(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] > pdf["that_array"], psdf["this_array"] > psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] > pdf["that_struct"], psdf["this_struct"] > psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] > pdf["this_array"], psdf["this_array"] > psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] > pdf["this_struct"], psdf["this_struct"] > psdf["this_struct"] + ) + + def test_ge(self): + pdf, psdf = self.complex_pdf, self.complex_pdf + self.assert_eq( + pdf["this_array"] >= pdf["that_array"], psdf["this_array"] >= psdf["that_array"] + ) + self.assert_eq( + pdf["this_struct"] >= pdf["that_struct"], psdf["this_struct"] >= psdf["that_struct"] + ) + self.assert_eq( + pdf["this_array"] >= pdf["this_array"], psdf["this_array"] >= psdf["this_array"] + ) + self.assert_eq( + pdf["this_struct"] >= pdf["this_struct"], psdf["this_struct"] >= psdf["this_struct"] + ) + + +class ComplexOpsTests(ComplexOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_complex_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py new file mode 100644 index 0000000..6b50ef0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py @@ -0,0 +1,257 @@ +# +# 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 datetime +import unittest +from distutils.version import LooseVersion + +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class DateOpsTestsMixin: + @property + def pser(self): + return pd.Series( + [datetime.date(1994, 1, 31), datetime.date(1994, 2, 1), datetime.date(1994, 2, 2)] + ) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def date_pdf(self): + psers = { + "this": self.pser, + "that": pd.Series( + [datetime.date(2000, 1, 31), datetime.date(1994, 3, 1), datetime.date(1990, 2, 2)] + ), + } + return pd.concat(psers, axis=1) + + @property + def date_psdf(self): + return ps.from_pandas(self.date_pdf) + + @property + def some_date(self): + return datetime.date(1994, 1, 1) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + self.assertRaises(TypeError, lambda: self.psser + self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43571): Enable DateOpsTests.test_sub for pandas 2.0.0.", + ) + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + self.assert_eq( + (self.pser - self.some_date).dt.days, + self.psser - self.some_date, + ) + pdf, psdf = self.pdf, self.psdf + for col in self.df_cols: + if col == "date": + self.assert_eq((pdf["date"] - pdf[col]).dt.days, psdf["date"] - psdf[col]) + else: + self.assertRaises(TypeError, lambda: psdf["date"] - psdf[col]) + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq((pdf["this"] - pdf["that"]).dt.days, psdf["this"] - psdf["that"]) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + self.assertRaises(TypeError, lambda: self.psser * self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + self.assertRaises(TypeError, lambda: self.psser / self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + self.assertRaises(TypeError, lambda: self.psser // self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + self.assertRaises(TypeError, lambda: self.psser % self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + self.assertRaises(TypeError, lambda: self.psser**self.some_date) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + self.assertRaises(TypeError, lambda: self.some_date + self.psser) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43570): Enable DateOpsTests.test_rsub for pandas 2.0.0.", + ) + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + self.assert_eq( + (self.some_date - self.pser).dt.days, + self.some_date - self.psser, + ) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 1 * self.psser) + self.assertRaises(TypeError, lambda: self.some_date * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + self.assertRaises(TypeError, lambda: self.some_date / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + self.assertRaises(TypeError, lambda: self.some_date // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + self.assertRaises(TypeError, lambda: self.some_date % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + self.assertRaises(TypeError, lambda: self.some_date**self.psser) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psser & True) + self.assertRaises(TypeError, lambda: self.psser & False) + self.assertRaises(TypeError, lambda: self.psser & self.psser) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psser) + self.assertRaises(TypeError, lambda: False & self.psser) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psser | True) + self.assertRaises(TypeError, lambda: self.psser | False) + self.assertRaises(TypeError, lambda: self.psser | self.psser) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psser) + self.assertRaises(TypeError, lambda: False | self.psser) + + def test_from_to_pandas(self): + data = [datetime.date(1994, 1, 31), datetime.date(1994, 2, 1), datetime.date(1994, 2, 2)] + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + cat_type = CategoricalDtype(categories=["a", "b", "c"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] == pdf["that"], psdf["this"] == psdf["that"]) + self.assert_eq(pdf["this"] == pdf["this"], psdf["this"] == psdf["this"]) + + def test_ne(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] != pdf["that"], psdf["this"] != psdf["that"]) + self.assert_eq(pdf["this"] != pdf["this"], psdf["this"] != psdf["this"]) + + def test_lt(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] < pdf["that"], psdf["this"] < psdf["that"]) + self.assert_eq(pdf["this"] < pdf["this"], psdf["this"] < psdf["this"]) + + def test_le(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] <= pdf["that"], psdf["this"] <= psdf["that"]) + self.assert_eq(pdf["this"] <= pdf["this"], psdf["this"] <= psdf["this"]) + + def test_gt(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] > pdf["that"], psdf["this"] > psdf["that"]) + self.assert_eq(pdf["this"] > pdf["this"], psdf["this"] > psdf["this"]) + + def test_ge(self): + pdf, psdf = self.date_pdf, self.date_psdf + self.assert_eq(pdf["this"] >= pdf["that"], psdf["this"] >= psdf["that"]) + self.assert_eq(pdf["this"] >= pdf["this"], psdf["this"] >= psdf["this"]) + + +class DateOpsTests(DateOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_date_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py new file mode 100644 index 0000000..c7bda90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py @@ -0,0 +1,260 @@ +# +# 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 datetime + +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class DatetimeOpsTestsMixin: + @property + def pser(self): + return pd.Series(pd.date_range("1994-1-31 10:30:15", periods=3, freq="D")) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def datetime_pdf(self): + psers = { + "this": self.pser, + "that": pd.Series(pd.date_range("1994-2-1 10:30:15", periods=3, freq="D")), + } + return pd.concat(psers, axis=1) + + @property + def datetime_psdf(self): + return ps.from_pandas(self.datetime_pdf) + + @property + def some_datetime(self): + return datetime.datetime(1994, 1, 31, 10, 30, 00) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + self.assertRaises(TypeError, lambda: self.psser + self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + self.assert_eq( + (self.pser - self.some_datetime).dt.total_seconds().astype("int"), + self.psser - self.some_datetime, + ) + + pdf, psdf = self.pdf, self.psdf + for col in self.df_cols: + if col == "datetime": + self.assert_eq( + (pdf["datetime"] - pdf[col]).dt.total_seconds().astype("int"), + psdf["datetime"] - psdf[col], + ) + else: + self.assertRaises(TypeError, lambda: psdf["datetime"] - psdf[col]) + + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq( + (pdf["that"] - pdf["this"]).dt.total_seconds().astype("int"), + psdf["that"] - psdf["this"], + ) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + self.assertRaises(TypeError, lambda: self.psser * self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + self.assertRaises(TypeError, lambda: self.psser / self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + self.assertRaises(TypeError, lambda: self.psser // self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + self.assertRaises(TypeError, lambda: self.psser % self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + self.assertRaises(TypeError, lambda: self.psser**self.some_datetime) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + self.assert_eq( + (self.some_datetime - self.pser).dt.total_seconds().astype("int"), + self.some_datetime - self.psser, + ) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 1 * self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + self.assertRaises(TypeError, lambda: self.some_datetime**self.psser) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psser & True) + self.assertRaises(TypeError, lambda: self.psser & False) + self.assertRaises(TypeError, lambda: self.psser & self.psser) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psser) + self.assertRaises(TypeError, lambda: False & self.psser) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psser | True) + self.assertRaises(TypeError, lambda: self.psser | False) + self.assertRaises(TypeError, lambda: self.psser | self.psser) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psser) + self.assertRaises(TypeError, lambda: False | self.psser) + + def test_from_to_pandas(self): + data = pd.date_range("1994-1-31 10:30:15", periods=3, freq="M") + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=["a", "b", "c"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + self.assertRaises(TypeError, lambda: psser.astype(bool)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] == pdf["that"], psdf["this"] == psdf["that"]) + self.assert_eq(pdf["this"] == pdf["this"], psdf["this"] == psdf["this"]) + + def test_ne(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] != pdf["that"], psdf["this"] != psdf["that"]) + self.assert_eq(pdf["this"] != pdf["this"], psdf["this"] != psdf["this"]) + + def test_lt(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] < pdf["that"], psdf["this"] < psdf["that"]) + self.assert_eq(pdf["this"] < pdf["this"], psdf["this"] < psdf["this"]) + + def test_le(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] <= pdf["that"], psdf["this"] <= psdf["that"]) + self.assert_eq(pdf["this"] <= pdf["this"], psdf["this"] <= psdf["this"]) + + def test_gt(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] > pdf["that"], psdf["this"] > psdf["that"]) + self.assert_eq(pdf["this"] > pdf["this"], psdf["this"] > psdf["this"]) + + def test_ge(self): + pdf, psdf = self.datetime_pdf, self.datetime_psdf + self.assert_eq(pdf["this"] >= pdf["that"], psdf["this"] >= psdf["that"]) + self.assert_eq(pdf["this"] >= pdf["this"], psdf["this"] >= psdf["this"]) + + +class DatetimeOpsTests(DatetimeOpsTestsMixin, OpsTestBase): + pass + + +class DatetimeNTZOpsTest(DatetimeOpsTests): + @classmethod + def setUpClass(cls): + super(DatetimeOpsTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.timestampType", "timestamp_ntz") + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_datetime_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py new file mode 100644 index 0000000..19a3e7c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py @@ -0,0 +1,178 @@ +# +# 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 pandas as pd +from pandas.api.types import CategoricalDtype + +import pyspark.pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class NullOpsTestsMixin: + @property + def pser(self): + return pd.Series([None, None, None]) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser - psser) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_from_to_pandas(self): + data = [None, None, None] + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=[1, 2, 3]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser == pser, psser == psser) + self.assert_eq(pser == [None, 1, None], psser == [None, 1, None]) + + def test_ne(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser != pser, psser != psser) + + def test_lt(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser < pser, psser < psser) + + def test_le(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser <= pser, psser <= psser) + + def test_gt(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser > pser, psser > psser) + + def test_ge(self): + pser, psser = self.pser, self.psser + self.assert_eq(pser >= pser, psser >= psser) + + +class NullOpsTests(NullOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_null_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py new file mode 100644 index 0000000..be70484 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py @@ -0,0 +1,184 @@ +# +# 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 datetime +import unittest + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class ArithmeticTestsMixin: + """Unit tests for arithmetic operations of numeric data types. + + A few test cases are disabled because pandas-on-Spark returns float64 whereas pandas + returns float32. + The underlying reason is the respective Spark operations return DoubleType always. + """ + + @property + def float_pser(self): + return pd.Series([1, 2, 3], dtype=float) + + @property + def float_psser(self): + return ps.from_pandas(self.float_pser) + + def test_add(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser + pser, psser + psser) + self.assert_eq(pser + 1, psser + 1) + # self.assert_eq(pser + 0.1, psser + 0.1) + self.assert_eq(pser + pser.astype(bool), psser + psser.astype(bool)) + self.assert_eq(pser + True, psser + True) + self.assert_eq(pser + False, psser + False) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pser + pdf[n_col], psser + psdf[n_col]) + else: + self.assertRaises(TypeError, lambda: psser + psdf[n_col]) + + def test_sub(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser - pser, psser - psser) + self.assert_eq(pser - 1, psser - 1) + # self.assert_eq(pser - 0.1, psser - 0.1) + self.assert_eq(pser - pser.astype(bool), psser - psser.astype(bool)) + self.assert_eq(pser - True, psser - True) + self.assert_eq(pser - False, psser - False) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pser - pdf[n_col], psser - psdf[n_col]) + else: + self.assertRaises(TypeError, lambda: psser - psdf[n_col]) + + def test_mul(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser * pser, psser * psser) + self.assert_eq(pser * pser.astype(bool), psser * psser.astype(bool)) + self.assert_eq(pser * True, psser * True) + self.assert_eq(pser * False, psser * False) + + if psser.dtype in [int, np.int32]: + self.assert_eq(pser * pdf["string"], psser * psdf["string"]) + else: + self.assertRaises(TypeError, lambda: psser * psdf["string"]) + + self.assert_eq(pser * pdf["bool"], psser * psdf["bool"]) + + self.assertRaises(TypeError, lambda: psser * psdf["datetime"]) + self.assertRaises(TypeError, lambda: psser * psdf["date"]) + self.assertRaises(TypeError, lambda: psser * psdf["categorical"]) + + def test_truediv(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if psser.dtype in [float, int, np.int32]: + self.assert_eq(pser / pser, psser / psser) + self.assert_eq(pser / pser.astype(bool), psser / psser.astype(bool)) + self.assert_eq(pser / True, psser / True) + self.assert_eq(pser / False, psser / False) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pdf["float"] / pdf[n_col], psdf["float"] / psdf[n_col]) + else: + self.assertRaises(TypeError, lambda: psser / psdf[n_col]) + + def test_floordiv(self): + pdf, psdf = self.pdf, self.psdf + pser, psser = pdf["float"], psdf["float"] + self.assert_eq(pser // pser, psser // psser) + self.assert_eq(pser // pser.astype(bool), psser // psser.astype(bool)) + self.assert_eq(pser // True, psser // True) + self.assert_eq(pser // False, psser // False) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pdf["float"] // pdf["bool"], psdf["float"] // psdf["bool"]) + else: + for col in self.numeric_df_cols: + psser = psdf[col] + self.assertRaises(TypeError, lambda: psser // psdf[n_col]) + + def test_mod(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser % pser, psser % psser) + self.assert_eq(pser % pser.astype(bool), psser % psser.astype(bool)) + self.assert_eq(pser % True, psser % True) + if col in ["int", "int32"]: + self.assert_eq( + pd.Series([np.nan, np.nan, np.nan], dtype=float, name=col), psser % False + ) + else: + self.assert_eq( + pd.Series([np.nan, np.nan, np.nan], dtype=pser.dtype, name=col), psser % False + ) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pdf["float"] % pdf[n_col], psdf["float"] % psdf[n_col]) + else: + self.assertRaises(TypeError, lambda: psser % psdf[n_col]) + + def test_pow(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if col in ["float", "float_w_nan"]: + self.assert_eq(pser**pser, psser**psser) + self.assert_eq(pser ** pser.astype(bool), psser ** psser.astype(bool)) + self.assert_eq(pser**True, psser**True) + self.assert_eq(pser**False, psser**False) + self.assert_eq(pser**1, psser**1) + self.assert_eq(pser**0, psser**0) + + for n_col in self.non_numeric_df_cols: + if n_col == "bool": + self.assert_eq(pdf["float"] ** pdf[n_col], psdf["float"] ** psdf[n_col]) + else: + self.assertRaises(TypeError, lambda: psser ** psdf[n_col]) + + +class ArithmeticTests(ArithmeticTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.data_type_ops.test_num_arithmetic import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py new file mode 100644 index 0000000..a799dd2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py @@ -0,0 +1,497 @@ +# +# 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 datetime +import unittest +from distutils.version import LooseVersion + +import pandas as pd +import numpy as np +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.pandas.typedef.typehints import ( + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) +from pyspark.sql.types import DecimalType, IntegralType + + +class NumOpsTestsMixin: + """Unit tests for arithmetic operations of numeric data types. + + A few test cases are disabled because pandas-on-Spark returns float64 whereas pandas + returns float32. + The underlying reason is the respective Spark operations return DoubleType always. + """ + + @property + def float_pser(self): + return pd.Series([1, 2, 3], dtype=float) + + @property + def float_psser(self): + return ps.from_pandas(self.float_pser) + + def test_and(self): + psdf = self.psdf + for col in self.numeric_df_cols: + psser = psdf[col] + self.assertRaises(TypeError, lambda: psser & True) + self.assertRaises(TypeError, lambda: psser & False) + self.assertRaises(TypeError, lambda: psser & psser) + + def test_rand(self): + psdf = self.psdf + for col in self.numeric_df_cols: + psser = psdf[col] + self.assertRaises(TypeError, lambda: True & psser) + self.assertRaises(TypeError, lambda: False & psser) + + def test_or(self): + psdf = self.psdf + for col in self.numeric_df_cols: + psser = psdf[col] + self.assertRaises(TypeError, lambda: psser | True) + self.assertRaises(TypeError, lambda: psser | False) + self.assertRaises(TypeError, lambda: psser | psser) + + def test_ror(self): + psdf = self.psdf + for col in self.numeric_df_cols: + psser = psdf[col] + self.assertRaises(TypeError, lambda: True | psser) + self.assertRaises(TypeError, lambda: False | psser) + + def test_xor(self): + pdf, psdf = self.integral_pdf, self.integral_psdf + pser, other_pser = pdf["this"], pdf["that"] + psser, other_psser = psdf["this"], psdf["that"] + + self.assert_eq(pser ^ other_pser, psser ^ other_psser) + self.assert_eq(pser ^ 2, psser ^ 2) + self.assert_eq(pser ^ 3, psser ^ 3) + self.assert_eq(pser ^ False, psser ^ False) + self.assert_eq(pser ^ True, psser ^ True) + + with self.assertRaisesRegex(TypeError, "XOR can not be applied to given types."): + psser ^ "a" + psser ^ None + + with option_context("compute.ops_on_diff_frames", True): + pser, other_pser = self.integral_pdf["this"], self.pdf["bool"] + psser, other_psser = self.integral_psdf["this"], self.psdf["bool"] + + self.assert_eq(pser ^ other_pser, psser ^ other_psser) + + def test_rxor(self): + pser, psser = self.pdf["int"], self.psdf["int"] + self.assert_eq(True ^ pser, True ^ psser) + self.assert_eq(False ^ pser, False ^ psser) + self.assert_eq(1 ^ pser, 1 ^ psser) + + def test_from_to_pandas(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf[col].isnull(), psdf[col].isnull()) + + def test_astype(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + + for int_type in [int, np.int32, np.int16, np.int8]: + if not pser.hasnans: + self.assert_eq(pser.astype(int_type), psser.astype(int_type)) + else: + self.assertRaisesRegex( + ValueError, + "Cannot convert %s with missing " + "values to integer" % psser._dtype_op.pretty_name, + lambda: psser.astype(int_type), + ) + + # TODO(SPARK-37039): the np.nan series.astype(bool) should be True + if not pser.hasnans: + self.assert_eq(pser.astype(bool), psser.astype(bool)) + + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=[2, 1, 3]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + if extension_object_dtypes_available and extension_float_dtypes_available: + pser = pd.Series(pd.Categorical([1.0, 2.0, 3.0]), dtype=pd.Float64Dtype()) + psser = ps.from_pandas(pser) + self.assert_eq(pser.astype(pd.BooleanDtype()), psser.astype(pd.BooleanDtype())) + + def test_astype_eager_check(self): + psser = self.psdf["float_nan"] + with ps.option_context("compute.eager_check", True), self.assertRaisesRegex( + ValueError, "Cannot convert" + ): + psser.astype(int) + with ps.option_context("compute.eager_check", False): + psser.astype(int) + + # Skip decimal_nan test before v1.3.0, it not supported by pandas on spark yet. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + psser = self.psdf["decimal_nan"] + with ps.option_context("compute.eager_check", True), self.assertRaisesRegex( + ValueError, "Cannot convert" + ): + psser.astype(int) + with ps.option_context("compute.eager_check", False): + psser.astype(int) + + def test_neg(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + self.assert_eq(-pdf[col], -psdf[col]) + + def test_abs(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + self.assert_eq(abs(pdf[col]), abs(psdf[col])) + + def test_invert(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if isinstance(psser.spark.data_type, IntegralType): + self.assert_eq(~pser, ~psser) + else: + self.assertRaises(TypeError, lambda: ~psser) + + def test_eq(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf[col] == pdf[col], psdf[col] == psdf[col]) + + def test_ne(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + self.assert_eq(pdf[col] != pdf[col], psdf[col] != psdf[col]) + + def test_lt(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if isinstance(psser.spark.data_type, DecimalType): + self.assertRaisesRegex(TypeError, "< can not be applied to", lambda: psser < psser) + else: + self.assert_eq(pser < pser, psser < psser) + + def test_le(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if isinstance(psser.spark.data_type, DecimalType): + self.assertRaisesRegex( + TypeError, "<= can not be applied to", lambda: psser <= psser + ) + else: + self.assert_eq(pser <= pser, psser <= psser) + + def test_gt(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if isinstance(psser.spark.data_type, DecimalType): + self.assertRaisesRegex(TypeError, "> can not be applied to", lambda: psser > psser) + else: + self.assert_eq(pser > pser, psser > psser) + + def test_ge(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + if isinstance(psser.spark.data_type, DecimalType): + self.assertRaisesRegex( + TypeError, ">= can not be applied to", lambda: psser >= psser + ) + else: + self.assert_eq(pser >= pser, psser >= psser) + + +@unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") +class IntegralExtensionOpsTest(OpsTestBase): + @property + def intergral_extension_psers(self): + return [pd.Series([1, 2, 3, None], dtype=dtype) for dtype in self.integral_extension_dtypes] + + @property + def intergral_extension_pssers(self): + return [ps.from_pandas(pser) for pser in self.intergral_extension_psers] + + @property + def intergral_extension_pser_psser_pairs(self): + return zip(self.intergral_extension_psers, self.intergral_extension_pssers) + + def test_from_to_pandas(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser, psser._to_pandas()) + self.check_extension(ps.from_pandas(pser), psser) + + def test_isnull(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.assert_eq(pser.isnull(), psser.isnull()) + + def test_astype(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + for dtype in self.extension_dtypes: + if dtype in self.string_extension_dtype: + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + # Limit pandas version due to + # https://github.com/pandas-dev/pandas/issues/31204 + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + else: + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + with ps.option_context("compute.eager_check", True): + self.assertRaisesRegex( + ValueError, + "Cannot convert integrals with missing values to bool", + lambda: psser.astype(bool), + ) + self.assertRaisesRegex( + ValueError, + "Cannot convert integrals with missing values to integer", + lambda: psser.astype(int), + ) + self.assertRaisesRegex( + ValueError, + "Cannot convert integrals with missing values to integer", + lambda: psser.astype(np.int32), + ) + with ps.option_context("compute.eager_check", False): + psser.astype(bool) + psser.astype(int) + psser.astype(np.int32) + + def test_neg(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + if LooseVersion(pd.__version__) < LooseVersion("1.1.3"): + # pandas < 1.1.0: object dtype is returned after negation + # pandas 1.1.1 and 1.1.2: + # a TypeError "bad operand type for unary -: 'IntegerArray'" is raised + # Please refer to https://github.com/pandas-dev/pandas/issues/36063. + self.check_extension(pd.Series([-1, -2, -3, None], dtype=pser.dtype), -psser) + else: + self.check_extension(-pser, -psser) + + def test_abs(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(abs(pser), abs(psser)) + + def test_invert(self): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(~pser, ~psser) + + def test_eq(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser == pser, (psser == psser).sort_index()) + + def test_ne(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser != pser, (psser != psser).sort_index()) + + def test_lt(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser < pser, (psser < psser).sort_index()) + + def test_le(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser <= pser, (psser <= psser).sort_index()) + + def test_gt(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser > pser, (psser > psser).sort_index()) + + def test_ge(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.intergral_extension_pser_psser_pairs: + self.check_extension(pser >= pser, (psser >= psser).sort_index()) + + def test_xor(self): + for psser in self.intergral_extension_pssers: + self.assertRaisesRegex( + TypeError, + "XOR can not be applied to given types.", + lambda: psser ^ 1, + ) + self.assertRaisesRegex( + TypeError, + "XOR can not be applied to given types.", + lambda: psser ^ psser, + ) + self.assertRaisesRegex( + TypeError, + "XOR can not be applied to given types.", + lambda: psser ^ False, + ) + + def test_rxor(self): + for psser in self.intergral_extension_pssers: + self.assertRaisesRegex( + TypeError, + "XOR can not be applied to given types.", + lambda: 1 ^ psser, + ) + self.assertRaisesRegex( + TypeError, + "XOR can not be applied to given types.", + lambda: False ^ psser, + ) + + +@unittest.skipIf( + not extension_float_dtypes_available, "pandas extension float dtypes are not available" +) +class FractionalExtensionOpsTest(OpsTestBase): + @property + def fractional_extension_psers(self): + return [ + pd.Series([0.1, 0.2, 0.3, None], dtype=dtype) + for dtype in self.fractional_extension_dtypes + ] + + @property + def fractional_extension_pssers(self): + return [ps.from_pandas(pser) for pser in self.fractional_extension_psers] + + @property + def fractional_extension_pser_psser_pairs(self): + return zip(self.fractional_extension_psers, self.fractional_extension_pssers) + + def test_from_to_pandas(self): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser, psser._to_pandas()) + self.check_extension(ps.from_pandas(pser), psser) + + def test_isnull(self): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.assert_eq(pser.isnull(), psser.isnull()) + + def test_astype(self): + for pser, psser in self.fractional_extension_pser_psser_pairs: + for dtype in self.extension_dtypes: + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + with ps.option_context("compute.eager_check", True): + self.assertRaisesRegex( + ValueError, + "Cannot convert fractions with missing values to bool", + lambda: psser.astype(bool), + ) + self.assertRaisesRegex( + ValueError, + "Cannot convert fractions with missing values to integer", + lambda: psser.astype(int), + ) + self.assertRaisesRegex( + ValueError, + "Cannot convert fractions with missing values to integer", + lambda: psser.astype(np.int32), + ) + with ps.option_context("compute.eager_check", False): + psser.astype(bool) + psser.astype(int) + psser.astype(np.int32) + + def test_neg(self): + # pandas raises "TypeError: bad operand type for unary -: 'FloatingArray'" + for dtype in self.fractional_extension_dtypes: + self.assert_eq( + ps.Series([-0.1, -0.2, -0.3, None], dtype=dtype), + -ps.Series([0.1, 0.2, 0.3, None], dtype=dtype), + ) + + def test_abs(self): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(abs(pser), abs(psser)) + + def test_invert(self): + for psser in self.fractional_extension_pssers: + self.assertRaises(TypeError, lambda: ~psser) + + def test_eq(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser == pser, (psser == psser).sort_index()) + + def test_ne(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser != pser, (psser != psser).sort_index()) + + def test_lt(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser < pser, (psser < psser).sort_index()) + + def test_le(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser <= pser, (psser <= psser).sort_index()) + + def test_gt(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser > pser, (psser > psser).sort_index()) + + def test_ge(self): + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.fractional_extension_pser_psser_pairs: + self.check_extension(pser >= pser, (psser >= psser).sort_index()) + + +class NumOpsTests(NumOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.data_type_ops.test_num_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py new file mode 100644 index 0000000..4e7ee17 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py @@ -0,0 +1,140 @@ +# +# 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 datetime +import unittest + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class ReverseTestsMixin: + """Unit tests for arithmetic operations of numeric data types. + + A few test cases are disabled because pandas-on-Spark returns float64 whereas pandas + returns float32. + The underlying reason is the respective Spark operations return DoubleType always. + """ + + @property + def float_pser(self): + return pd.Series([1, 2, 3], dtype=float) + + @property + def float_psser(self): + return ps.from_pandas(self.float_pser) + + def test_radd(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(1 + pser, 1 + psser) + # self.assert_eq(0.1 + pser, 0.1 + psser) + self.assertRaises(TypeError, lambda: "x" + psser) + self.assert_eq(True + pser, True + psser) + self.assert_eq(False + pser, False + psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) + psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) + psser) + + def test_rsub(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(1 - pser, 1 - psser) + # self.assert_eq(0.1 - pser, 0.1 - psser) + self.assertRaises(TypeError, lambda: "x" - psser) + self.assert_eq(True - pser, True - psser) + self.assert_eq(False - pser, False - psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) - psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) - psser) + + def test_rmul(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(1 * pser, 1 * psser) + # self.assert_eq(0.1 * pser, 0.1 * psser) + self.assertRaises(TypeError, lambda: "x" * psser) + self.assert_eq(True * pser, True * psser) + self.assert_eq(False * pser, False * psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) * psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) * psser) + + def test_rtruediv(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + # self.assert_eq(5 / pser, 5 / psser) + # self.assert_eq(0.1 / pser, 0.1 / psser) + self.assertRaises(TypeError, lambda: "x" / psser) + self.assert_eq((True / pser).astype(float), True / psser, check_exact=False) + self.assert_eq((False / pser).astype(float), False / psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) / psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) / psser) + + def test_rfloordiv(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + # self.assert_eq(5 // pser, 5 // psser) + # self.assert_eq(0.1 // pser, 0.1 // psser) + self.assertRaises(TypeError, lambda: "x" // psser) + self.assert_eq((True // pser).astype(float), True // psser) + self.assert_eq((False // pser).astype(float), False // psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) // psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) // psser) + + def test_rpow(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + # self.assert_eq(1 ** pser, 1 ** psser) + # self.assert_eq(0.1 ** pser, 0.1 ** psser) + self.assertRaises(TypeError, lambda: "x" ** psser) + self.assert_eq((True**pser).astype(float), True**psser) + self.assert_eq((False**pser).astype(float), False**psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) ** psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) ** psser) + + def test_rmod(self): + pdf, psdf = self.pdf, self.psdf + for col in self.numeric_df_cols: + pser, psser = pdf[col], psdf[col] + self.assert_eq(1 % pser, 1 % psser) + # self.assert_eq(0.1 % pser, 0.1 % psser) + self.assert_eq(True % pser, True % psser) + self.assert_eq(False % pser, False % psser) + self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) % psser) + self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) % psser) + + +class ReverseTests(ReverseTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.data_type_ops.test_num_reverse import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py new file mode 100644 index 0000000..136366d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py @@ -0,0 +1,354 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.pandas.typedef.typehints import extension_object_dtypes_available + +if extension_object_dtypes_available: + from pandas import StringDtype + + +class StringOpsTestsMixin: + @property + def bool_pdf(self): + return pd.DataFrame({"this": ["x", "y", "z"], "that": ["z", "y", "x"]}) + + @property + def bool_psdf(self): + return ps.from_pandas(self.bool_pdf) + + @property + def bool_non_numeric_pdf(self): + return pd.concat([self.bool_pdf, self.non_numeric_pdf], axis=1) + + @property + def bool_non_numeric_psdf(self): + return ps.from_pandas(self.bool_non_numeric_pdf) + + def test_add(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser + "x", psser + "x") + self.assertRaises(TypeError, lambda: psser + 1) + + self.assert_eq(pser + other_pser, psser + other_psser) + + for col in self.non_numeric_df_cols: + if col != "string": + self.assertRaises(TypeError, lambda: psser + psdf[col]) + + def test_sub(self): + psdf = self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] - "x") + self.assertRaises(TypeError, lambda: psdf["string"] - 1) + + for col in self.df_cols: + self.assertRaises(TypeError, lambda: psdf["string"] - psdf[col]) + + def test_mul(self): + pdf, psdf = self.pdf, self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] * "x") + self.assert_eq(pdf["string"] * 1, psdf["string"] * 1) + + for col in self.df_cols: + if col in ["int", "int32"]: + self.assert_eq(pdf["string"] * pdf[col], psdf["string"] * psdf[col]) + else: + self.assertRaises(TypeError, lambda: psdf["string"] * psdf[col]) + + def test_truediv(self): + psdf = self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] / "x") + self.assertRaises(TypeError, lambda: psdf["string"] / 1) + + for col in self.df_cols: + self.assertRaises(TypeError, lambda: psdf["string"] / psdf[col]) + + def test_floordiv(self): + psdf = self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] // "x") + self.assertRaises(TypeError, lambda: psdf["string"] // 1) + + for col in self.df_cols: + self.assertRaises(TypeError, lambda: psdf["string"] // psdf[col]) + + def test_mod(self): + psdf = self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] % "x") + self.assertRaises(TypeError, lambda: psdf["string"] % 1) + + for col in self.df_cols: + self.assertRaises(TypeError, lambda: psdf["string"] % psdf[col]) + + def test_pow(self): + psdf = self.psdf + self.assertRaises(TypeError, lambda: psdf["string"] ** "x") + self.assertRaises(TypeError, lambda: psdf["string"] ** 1) + + for col in self.df_cols: + self.assertRaises(TypeError, lambda: psdf["string"] ** psdf[col]) + + def test_radd(self): + self.assert_eq("x" + self.pdf["string"], "x" + self.psdf["string"]) + self.assertRaises(TypeError, lambda: 1 + self.psdf["string"]) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psdf["string"]) + self.assertRaises(TypeError, lambda: 1 - self.psdf["string"]) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psdf["string"]) + self.assert_eq(1 * self.pdf["string"], 1 * self.psdf["string"]) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psdf["string"]) + self.assertRaises(TypeError, lambda: 1 / self.psdf["string"]) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psdf["string"]) + self.assertRaises(TypeError, lambda: 1 // self.psdf["string"]) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psdf["string"]) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psdf["string"]) + self.assertRaises(TypeError, lambda: 1 ** self.psdf["string"]) + + def test_and(self): + self.assertRaises(TypeError, lambda: self.psdf["string"] & True) + self.assertRaises(TypeError, lambda: self.psdf["string"] & False) + self.assertRaises(TypeError, lambda: self.psdf["string"] & self.psdf["string"]) + + def test_rand(self): + self.assertRaises(TypeError, lambda: True & self.psdf["string"]) + self.assertRaises(TypeError, lambda: False & self.psdf["string"]) + + def test_or(self): + self.assertRaises(TypeError, lambda: self.psdf["string"] | True) + self.assertRaises(TypeError, lambda: self.psdf["string"] | False) + self.assertRaises(TypeError, lambda: self.psdf["string"] | self.psdf["string"]) + + def test_ror(self): + self.assertRaises(TypeError, lambda: True | self.psdf["string"]) + self.assertRaises(TypeError, lambda: False | self.psdf["string"]) + + def test_from_to_pandas(self): + data = ["x", "y", "z"] + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pdf["string"].isnull(), self.psdf["string"].isnull()) + + def test_astype(self): + pser = pd.Series(["1", "2", "3"]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.astype(int), psser.astype(int)) + self.assert_eq(pser.astype(float), psser.astype(float)) + self.assert_eq(pser.astype(np.float32), psser.astype(np.float32)) + self.assert_eq(pser.astype(np.int32), psser.astype(np.int32)) + self.assert_eq(pser.astype(np.int16), psser.astype(np.int16)) + self.assert_eq(pser.astype(np.int8), psser.astype(np.int8)) + self.assert_eq(pser.astype(str), psser.astype(str)) + self.assert_eq(pser.astype(bool), psser.astype(bool)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=["3", "1", "2"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psdf["string"]) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psdf["string"])) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psdf["string"]) + + def test_eq(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser == other_pser, psser == other_psser) + self.assert_eq(pser == pser, psser == psser) + + def test_ne(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser != other_pser, psser != other_psser) + self.assert_eq(pser != pser, psser != psser) + + def test_lt(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser < other_pser, psser < other_psser) + self.assert_eq(pser < pser, psser < psser) + + def test_le(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser <= other_pser, psser <= other_psser) + self.assert_eq(pser <= pser, psser <= psser) + + def test_gt(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser > other_pser, psser > other_psser) + self.assert_eq(pser > pser, psser > psser) + + def test_ge(self): + pdf, psdf = self.bool_non_numeric_pdf, self.bool_non_numeric_psdf + pser, psser = pdf["this"], psdf["this"] + other_pser, other_psser = pdf["that"], psdf["that"] + self.assert_eq(pser >= other_pser, psser >= other_psser) + self.assert_eq(pser >= pser, psser >= psser) + + +class StringOpsTests(StringOpsTestsMixin, OpsTestBase): + pass + + +@unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" +) +class StringExtensionOpsTest(StringOpsTests): + @property + def pser(self): + return pd.Series(["x", "y", "z", None], dtype="string") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def other_pser(self): + return pd.Series([None, "z", "y", "x"], dtype="string") + + @property + def other_psser(self): + return ps.from_pandas(self.other_pser) + + def test_radd(self): + self.assert_eq("x" + self.pser, ("x" + self.psser).astype("string")) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assert_eq(self.pser * 1, self.psser * 1) + + with option_context("compute.ops_on_diff_frames", True): + for pser, psser in self.pser_psser_pairs: + if psser.dtype in [np.int32, np.int64]: + self.assert_eq( + ps.Series(["x", "yy", "zzz", None]).astype("string"), + (self.psser * psser).sort_index(), + ) + else: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_from_to_pandas(self): + data = ["x", "y", "z", None] + pser = pd.Series(data, dtype="string") + psser = ps.Series(data, dtype="string") + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + + self.assert_eq(pser.astype(str).tolist(), psser.astype(str).tolist()) + + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=["x", "y"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + for dtype in self.object_extension_dtypes: + if dtype in ["string", StringDtype()]: + self.check_extension(pser.astype(dtype), psser.astype(dtype)) + + def test_eq(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser == self.other_pser, (self.psser == self.other_psser).sort_index() + ) + self.check_extension(self.pser == self.pser, (self.psser == self.psser).sort_index()) + + def test_ne(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser != self.other_pser, (self.psser != self.other_psser).sort_index() + ) + self.check_extension(self.pser != self.pser, (self.psser != self.psser).sort_index()) + + def test_lt(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser < self.other_pser, (self.psser < self.other_psser).sort_index() + ) + self.check_extension(self.pser < self.pser, (self.psser < self.psser).sort_index()) + + def test_le(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser <= self.other_pser, (self.psser <= self.other_psser).sort_index() + ) + self.check_extension(self.pser <= self.pser, (self.psser <= self.psser).sort_index()) + + def test_gt(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser > self.other_pser, (self.psser > self.other_psser).sort_index() + ) + self.check_extension(self.pser > self.pser, (self.psser > self.psser).sort_index()) + + def test_ge(self): + with option_context("compute.ops_on_diff_frames", True): + self.check_extension( + self.pser >= self.other_pser, (self.psser >= self.other_psser).sort_index() + ) + self.check_extension(self.pser >= self.pser, (self.psser >= self.psser).sort_index()) + + +if __name__ == "__main__": + + from pyspark.pandas.tests.data_type_ops.test_string_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py new file mode 100644 index 0000000..f89ec17 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py @@ -0,0 +1,219 @@ +# +# 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 datetime import timedelta + +import pandas as pd +from pandas.api.types import CategoricalDtype + +import pyspark.pandas as ps +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class TimedeltaOpsTestsMixin: + @property + def pser(self): + return pd.Series([timedelta(1), timedelta(microseconds=2), timedelta(weeks=3)]) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def timedelta_pdf(self): + psers = { + "this": self.pser, + "that": pd.Series([timedelta(0), timedelta(microseconds=1), timedelta(seconds=2)]), + } + return pd.concat(psers, axis=1) + + @property + def timedelta_psdf(self): + return ps.from_pandas(self.timedelta_pdf) + + @property + def some_timedelta(self): + return timedelta(weeks=2) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + self.assert_eq(self.pser - self.some_timedelta, self.psser - self.some_timedelta) + + pdf, psdf = self.pdf, self.psdf + for col in self.df_cols: + if col == "timedelta": + self.assert_eq(pdf["timedelta"] - pdf[col], psdf["timedelta"] - psdf[col]) + else: + self.assertRaises(TypeError, lambda: psdf["timedelta"] - psdf[col]) + + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["that"] - pdf["this"], psdf["that"] - psdf["this"]) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + self.assert_eq(self.some_timedelta - self.pser, self.some_timedelta - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_from_to_pandas(self): + data = [timedelta(1), timedelta(microseconds=2)] + pser = pd.Series(data) + psser = ps.Series(data) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + pser = self.pser + psser = self.psser + target_psser = ps.Series( + [ + "INTERVAL '1 00:00:00' DAY TO SECOND", + "INTERVAL '0 00:00:00.000002' DAY TO SECOND", + "INTERVAL '21 00:00:00' DAY TO SECOND", + ] + ) + self.assert_eq(target_psser, psser.astype(str)) + self.assert_eq(pser.astype("category"), psser.astype("category")) + cat_type = CategoricalDtype(categories=["a", "b", "c"]) + self.assert_eq(pser.astype(cat_type), psser.astype(cat_type)) + + self.assertRaises(TypeError, lambda: psser.astype(bool)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] == pdf["this"], psdf["this"] == psdf["this"]) + self.assert_eq(pdf["this"] == pdf["that"], psdf["this"] == psdf["that"]) + + def test_ne(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] != pdf["this"], psdf["this"] != psdf["this"]) + self.assert_eq(pdf["this"] != pdf["that"], psdf["this"] != psdf["that"]) + + def test_lt(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] < pdf["that"], psdf["this"] < psdf["that"]) + self.assert_eq(pdf["this"] < pdf["this"], psdf["this"] < psdf["this"]) + + def test_le(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] <= pdf["that"], psdf["this"] <= psdf["that"]) + self.assert_eq(pdf["this"] <= pdf["this"], psdf["this"] <= psdf["this"]) + + def test_gt(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] > pdf["that"], psdf["this"] > psdf["that"]) + self.assert_eq(pdf["this"] > pdf["this"], psdf["this"] > psdf["this"]) + + def test_ge(self): + pdf, psdf = self.timedelta_pdf, self.timedelta_psdf + self.assert_eq(pdf["this"] >= pdf["that"], psdf["this"] >= psdf["that"]) + self.assert_eq(pdf["this"] >= pdf["this"], psdf["this"] >= psdf["this"]) + + +class TimedeltaOpsTests(TimedeltaOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_timedelta_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py new file mode 100644 index 0000000..45f8cca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py @@ -0,0 +1,192 @@ +# +# 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 pandas as pd + +import pyspark.pandas as ps +from pyspark.ml.linalg import SparseVector +from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase + + +class UDTOpsTestsMixin: + @property + def pser(self): + sparse_values = {0: 0.1, 1: 1.1} + return pd.Series([SparseVector(len(sparse_values), sparse_values)]) + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @property + def udt_pdf(self): + sparse_values = {0: 0.2, 1: 1.0} + psers = { + "this": self.pser, + "that": pd.Series([SparseVector(len(sparse_values), sparse_values)]), + } + return pd.concat(psers, axis=1) + + @property + def udt_psdf(self): + return ps.from_pandas(self.udt_pdf) + + def test_add(self): + self.assertRaises(TypeError, lambda: self.psser + "x") + self.assertRaises(TypeError, lambda: self.psser + 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser + psser) + + def test_sub(self): + self.assertRaises(TypeError, lambda: self.psser - "x") + self.assertRaises(TypeError, lambda: self.psser - 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser - psser) + + def test_mul(self): + self.assertRaises(TypeError, lambda: self.psser * "x") + self.assertRaises(TypeError, lambda: self.psser * 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser * psser) + + def test_truediv(self): + self.assertRaises(TypeError, lambda: self.psser / "x") + self.assertRaises(TypeError, lambda: self.psser / 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser / psser) + + def test_floordiv(self): + self.assertRaises(TypeError, lambda: self.psser // "x") + self.assertRaises(TypeError, lambda: self.psser // 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser // psser) + + def test_mod(self): + self.assertRaises(TypeError, lambda: self.psser % "x") + self.assertRaises(TypeError, lambda: self.psser % 1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser % psser) + + def test_pow(self): + self.assertRaises(TypeError, lambda: self.psser ** "x") + self.assertRaises(TypeError, lambda: self.psser**1) + + for psser in self.pssers: + self.assertRaises(TypeError, lambda: self.psser**psser) + + def test_radd(self): + self.assertRaises(TypeError, lambda: "x" + self.psser) + self.assertRaises(TypeError, lambda: 1 + self.psser) + + def test_rsub(self): + self.assertRaises(TypeError, lambda: "x" - self.psser) + self.assertRaises(TypeError, lambda: 1 - self.psser) + + def test_rmul(self): + self.assertRaises(TypeError, lambda: "x" * self.psser) + self.assertRaises(TypeError, lambda: 2 * self.psser) + + def test_rtruediv(self): + self.assertRaises(TypeError, lambda: "x" / self.psser) + self.assertRaises(TypeError, lambda: 1 / self.psser) + + def test_rfloordiv(self): + self.assertRaises(TypeError, lambda: "x" // self.psser) + self.assertRaises(TypeError, lambda: 1 // self.psser) + + def test_rmod(self): + self.assertRaises(TypeError, lambda: 1 % self.psser) + + def test_rpow(self): + self.assertRaises(TypeError, lambda: "x" ** self.psser) + self.assertRaises(TypeError, lambda: 1**self.psser) + + def test_from_to_pandas(self): + sparse_values = {0: 0.1, 1: 1.1} + sparse_vector = SparseVector(len(sparse_values), sparse_values) + pser = pd.Series([sparse_vector]) + psser = ps.Series([sparse_vector]) + self.assert_eq(pser, psser._to_pandas()) + self.assert_eq(ps.from_pandas(pser), psser) + + def test_isnull(self): + self.assert_eq(self.pser.isnull(), self.psser.isnull()) + + def test_astype(self): + self.assertRaises(TypeError, lambda: self.psser.astype(str)) + + def test_neg(self): + self.assertRaises(TypeError, lambda: -self.psser) + + def test_abs(self): + self.assertRaises(TypeError, lambda: abs(self.psser)) + + def test_invert(self): + self.assertRaises(TypeError, lambda: ~self.psser) + + def test_eq(self): + pdf, psdf = self.udt_pdf, self.udt_psdf + self.assert_eq(pdf["this"] == pdf["this"], psdf["this"] == psdf["this"]) + self.assert_eq(pdf["this"] == pdf["that"], psdf["this"] == psdf["that"]) + + def test_ne(self): + pdf, psdf = self.udt_pdf, self.udt_psdf + self.assert_eq(pdf["this"] != pdf["this"], psdf["this"] != psdf["this"]) + self.assert_eq(pdf["this"] != pdf["that"], psdf["this"] != psdf["that"]) + + def test_lt(self): + self.assertRaisesRegex( + TypeError, "< can not be applied to", lambda: self.psser < self.psser + ) + + def test_le(self): + self.assertRaisesRegex( + TypeError, "<= can not be applied to", lambda: self.psser <= self.psser + ) + + def test_gt(self): + self.assertRaisesRegex( + TypeError, "> can not be applied to", lambda: self.psser > self.psser + ) + + def test_ge(self): + self.assertRaisesRegex( + TypeError, ">= can not be applied to", lambda: self.psser >= self.psser + ) + + +class UDTOpsTests(UDTOpsTestsMixin, OpsTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.data_type_ops.test_udt_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/testing_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/testing_utils.py new file mode 100644 index 0000000..222b945 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/data_type_ops/testing_utils.py @@ -0,0 +1,228 @@ +# +# 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 datetime +import decimal +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.pandas.typedef import extension_dtypes + +from pyspark.pandas.typedef.typehints import ( + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) + +from pyspark.testing.pandasutils import ComparisonTestBase + +if extension_dtypes_available: + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + +if extension_float_dtypes_available: + from pandas import Float32Dtype, Float64Dtype + +if extension_object_dtypes_available: + from pandas import BooleanDtype, StringDtype + + +class OpsTestBase(ComparisonTestBase): + """The test base for arithmetic operations of different data types.""" + + @property + def numeric_pdf(self): + dtypes = [np.int32, int, np.float32, float] + sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes] + sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)])) + sers.append(pd.Series([1, 2, np.nan], dtype=float)) + # Skip decimal_nan test before v1.3.0, it not supported by pandas on spark yet. + if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + sers.append( + pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(np.nan)]) + ) + pdf = pd.concat(sers, axis=1) + if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + pdf.columns = [dtype.__name__ for dtype in dtypes] + [ + "decimal", + "float_nan", + "decimal_nan", + ] + else: + pdf.columns = [dtype.__name__ for dtype in dtypes] + ["decimal", "float_nan"] + return pdf + + @property + def numeric_psdf(self): + return ps.from_pandas(self.numeric_pdf) + + @property + def numeric_df_cols(self): + return self.numeric_pdf.columns + + @property + def integral_pdf(self): + return pd.DataFrame({"this": [1, 2, 3], "that": [2, 2, 1]}) + + @property + def integral_psdf(self): + return ps.from_pandas(self.integral_pdf) + + @property + def non_numeric_pdf(self): + psers = { + "string": pd.Series(["x", "y", "z"]), + "bool": pd.Series([True, True, False]), + "date": pd.Series( + [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)] + ), + "datetime": pd.to_datetime(pd.Series([1, 2, 3])), + "timedelta": pd.Series( + [datetime.timedelta(1), datetime.timedelta(hours=2), datetime.timedelta(weeks=3)] + ), + "categorical": pd.Series(["a", "b", "a"], dtype="category"), + } + return pd.concat(psers, axis=1) + + @property + def non_numeric_psdf(self): + return ps.from_pandas(self.non_numeric_pdf) + + @property + def non_numeric_df_cols(self): + return self.non_numeric_pdf.columns + + @property + def pdf(self): + return pd.concat([self.numeric_pdf, self.non_numeric_pdf], axis=1) + + @property + def df_cols(self): + return self.pdf.columns + + @property + def numeric_psers(self): + dtypes = [np.float32, float, int, np.int32] + sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes] + sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)])) + return sers + + @property + def numeric_pssers(self): + return [ps.from_pandas(pser) for pser in self.numeric_psers] + + @property + def numeric_pser_psser_pairs(self): + return zip(self.numeric_psers, self.numeric_pssers) + + @property + def non_numeric_psers(self): + psers = { + "string": pd.Series(["x", "y", "z"]), + "datetime": pd.to_datetime(pd.Series([1, 2, 3])), + "bool": pd.Series([True, True, False]), + "date": pd.Series( + [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)] + ), + "categorical": pd.Series(["a", "b", "a"], dtype="category"), + } + return psers + + @property + def non_numeric_pssers(self): + pssers = {} + + for k, v in self.non_numeric_psers.items(): + pssers[k] = ps.from_pandas(v) + return pssers + + @property + def non_numeric_pser_psser_pairs(self): + return zip(self.non_numeric_psers.values(), self.non_numeric_pssers.values()) + + @property + def pssers(self): + return self.numeric_pssers + list(self.non_numeric_pssers.values()) + + @property + def psers(self): + return self.numeric_psers + list(self.non_numeric_psers.values()) + + @property + def pser_psser_pairs(self): + return zip(self.psers, self.pssers) + + @property + def string_extension_dtype(self): + return ["string", StringDtype()] if extension_object_dtypes_available else [] + + @property + def object_extension_dtypes(self): + return ( + ["boolean", "string", BooleanDtype(), StringDtype()] + if extension_object_dtypes_available + else [] + ) + + @property + def fractional_extension_dtypes(self): + return ( + ["Float32", "Float64", Float32Dtype(), Float64Dtype()] + if extension_float_dtypes_available + else [] + ) + + @property + def integral_extension_dtypes(self): + return ( + [ + "Int8", + "Int16", + "Int32", + "Int64", + Int8Dtype(), + Int16Dtype(), + Int32Dtype(), + Int64Dtype(), + ] + if extension_dtypes_available + else [] + ) + + @property + def extension_dtypes(self): + return ( + self.object_extension_dtypes + + self.fractional_extension_dtypes + + self.integral_extension_dtypes + ) + + def check_extension(self, left, right): + """ + Compare `psser` and `pser` of numeric ExtensionDtypes. + + This utility is to adjust an issue for comparing numeric ExtensionDtypes in specific + pandas versions. Please refer to https://github.com/pandas-dev/pandas/issues/39410. + """ + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(left, right, check_exact=False) + self.assertTrue(isinstance(left.dtype, extension_dtypes)) + self.assertTrue(isinstance(right.dtype, extension_dtypes)) + else: + self.assert_eq(left, right) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_align.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_align.py new file mode 100644 index 0000000..3e3f454 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_align.py @@ -0,0 +1,118 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesAlignMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_align(self): + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30]) + pdf2 = pd.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12]) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + for join in ["outer", "inner", "left", "right"]: + for axis in [None, 0]: + psdf_l, psdf_r = psdf1.align(psdf2, join=join, axis=axis) + pdf_l, pdf_r = pdf1.align(pdf2, join=join, axis=axis) + self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index()) + self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index()) + + pser1 = pd.Series([7, 8, 9], index=[10, 11, 12]) + pser2 = pd.Series(["g", "h", "i"], index=[10, 20, 30]) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + for join in ["outer", "inner", "left", "right"]: + psser_l, psser_r = psser1.align(psser2, join=join) + pser_l, pser_r = pser1.align(pser2, join=join) + self.assert_eq(psser_l.sort_index(), pser_l.sort_index()) + self.assert_eq(psser_r.sort_index(), pser_r.sort_index()) + + psdf_l, psser_r = psdf1.align(psser1, join=join, axis=0) + pdf_l, pser_r = pdf1.align(pser1, join=join, axis=0) + self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index()) + self.assert_eq(psser_r.sort_index(), pser_r.sort_index()) + + psser_l, psdf_r = psser1.align(psdf1, join=join) + pser_l, pdf_r = pser1.align(pdf1, join=join) + self.assert_eq(psser_l.sort_index(), pser_l.sort_index()) + self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index()) + + # multi-index columns + pdf3 = pd.DataFrame( + {("x", "a"): [4, 5, 6], ("y", "c"): ["d", "e", "f"]}, index=[10, 11, 12] + ) + psdf3 = ps.from_pandas(pdf3) + pser3 = pdf3[("y", "c")] + psser3 = psdf3[("y", "c")] + + for join in ["outer", "inner", "left", "right"]: + psdf_l, psdf_r = psdf1.align(psdf3, join=join, axis=0) + pdf_l, pdf_r = pdf1.align(pdf3, join=join, axis=0) + self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index()) + self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index()) + + psser_l, psser_r = psser1.align(psser3, join=join) + pser_l, pser_r = pser1.align(pser3, join=join) + self.assert_eq(psser_l.sort_index(), pser_l.sort_index()) + self.assert_eq(psser_r.sort_index(), pser_r.sort_index()) + + psdf_l, psser_r = psdf1.align(psser3, join=join, axis=0) + pdf_l, pser_r = pdf1.align(pser3, join=join, axis=0) + self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index()) + self.assert_eq(psser_r.sort_index(), pser_r.sort_index()) + + psser_l, psdf_r = psser3.align(psdf1, join=join) + pser_l, pdf_r = pser3.align(pdf1, join=join) + self.assert_eq(psser_l.sort_index(), pser_l.sort_index()) + self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index()) + + self.assertRaises(ValueError, lambda: psdf1.align(psdf3, axis=None)) + self.assertRaises(ValueError, lambda: psdf1.align(psdf3, axis=1)) + + +class DiffFramesAlignTests(DiffFramesAlignMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_align import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py new file mode 100644 index 0000000..a82c98e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py @@ -0,0 +1,198 @@ +# +# 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 pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesBasicSlowMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_update(self): + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + pser.update(pd.Series([4, 5, 6])) + psser.update(ps.Series([4, 5, 6])) + self.assert_eq(psser.sort_index(), pser.sort_index()) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + pser1 = pd.Series([None, 2, 3, 4, 5, 6, 7, 8, None]) + pser2 = pd.Series([None, 5, None, 3, 2, 1, None, 0, 0]) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + pser1.update(pser2) + psser1.update(psser2) + self.assert_eq(psser1.sort_index(), pser1) + + def test_where(self): + pdf1 = pd.DataFrame({"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame({"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.where(pdf2 > 100), psdf1.where(psdf2 > 100).sort_index()) + + pdf1 = pd.DataFrame({"A": [-1, -2, -3, -4, -5], "B": [-100, -200, -300, -400, -500]}) + pdf2 = pd.DataFrame({"A": [-10, -20, -30, -40, -50], "B": [-5, -4, -3, -2, -1]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.where(pdf2 < -250), psdf1.where(psdf2 < -250).sort_index()) + + # multi-index columns + pdf1 = pd.DataFrame({("X", "A"): [0, 1, 2, 3, 4], ("X", "B"): [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame( + {("X", "A"): [0, -1, -2, -3, -4], ("X", "B"): [-100, -200, -300, -400, -500]} + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.where(pdf2 > 100), psdf1.where(psdf2 > 100).sort_index()) + + def test_mask(self): + pdf1 = pd.DataFrame({"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame({"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.mask(pdf2 < 100), psdf1.mask(psdf2 < 100).sort_index()) + + pdf1 = pd.DataFrame({"A": [-1, -2, -3, -4, -5], "B": [-100, -200, -300, -400, -500]}) + pdf2 = pd.DataFrame({"A": [-10, -20, -30, -40, -50], "B": [-5, -4, -3, -2, -1]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.mask(pdf2 > -250), psdf1.mask(psdf2 > -250).sort_index()) + + # multi-index columns + pdf1 = pd.DataFrame({("X", "A"): [0, 1, 2, 3, 4], ("X", "B"): [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame( + {("X", "A"): [0, -1, -2, -3, -4], ("X", "B"): [-100, -200, -300, -400, -500]} + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.mask(pdf2 < 100), psdf1.mask(psdf2 < 100).sort_index()) + + def test_to_series_comparison(self): + psidx1 = ps.Index([1, 2, 3, 4, 5]) + psidx2 = ps.Index([1, 2, 3, 4, 5]) + + self.assert_eq((psidx1.to_series() == psidx2.to_series()).all(), True) + + psidx1.name = "koalas" + psidx2.name = "koalas" + + self.assert_eq((psidx1.to_series() == psidx2.to_series()).all(), True) + + def test_pow_and_rpow(self): + pser = pd.Series([1, 2, np.nan]) + psser = ps.from_pandas(pser) + pser_other = pd.Series([np.nan, 2, 3]) + psser_other = ps.from_pandas(pser_other) + + self.assert_eq(pser.pow(pser_other), psser.pow(psser_other).sort_index()) + self.assert_eq(pser**pser_other, (psser**psser_other).sort_index()) + self.assert_eq(pser.rpow(pser_other), psser.rpow(psser_other).sort_index()) + + def test_shift(self): + pdf = pd.DataFrame( + { + "Col1": [10, 20, 15, 30, 45], + "Col2": [13, 23, 18, 33, 48], + "Col3": [17, 27, 22, 37, 52], + }, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.shift().loc[pdf["Col1"] == 20].astype(int), psdf.shift().loc[psdf["Col1"] == 20] + ) + self.assert_eq( + pdf["Col2"].shift().loc[pdf["Col1"] == 20].astype(int), + psdf["Col2"].shift().loc[psdf["Col1"] == 20], + ) + + def test_diff(self): + pdf = pd.DataFrame( + { + "Col1": [10, 20, 15, 30, 45], + "Col2": [13, 23, 18, 33, 48], + "Col3": [17, 27, 22, 37, 52], + }, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.diff().loc[pdf["Col1"] == 20].astype(int), psdf.diff().loc[psdf["Col1"] == 20] + ) + self.assert_eq( + pdf["Col2"].diff().loc[pdf["Col1"] == 20].astype(int), + psdf["Col2"].diff().loc[psdf["Col1"] == 20], + ) + + def test_rank(self): + pdf = pd.DataFrame( + { + "Col1": [10, 20, 15, 30, 45], + "Col2": [13, 23, 18, 33, 48], + "Col3": [17, 27, 22, 37, 52], + }, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.rank().loc[pdf["Col1"] == 20], psdf.rank().loc[psdf["Col1"] == 20]) + self.assert_eq( + pdf["Col2"].rank().loc[pdf["Col1"] == 20], psdf["Col2"].rank().loc[psdf["Col1"] == 20] + ) + + +class DiffFramesBasicSlowTests(DiffFramesBasicSlowMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_basic_slow import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_cov_corrwith.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_cov_corrwith.py new file mode 100644 index 0000000..d3b36da --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_cov_corrwith.py @@ -0,0 +1,181 @@ +# +# 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 distutils.version import LooseVersion + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesCovCorrWithMixin: + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 10, 11], + ) + + @property + def pdf2(self): + return pd.DataFrame( + {"a": [9, 8, 7, 6, 5, 4, 3, 2, 1], "b": [0, 0, 0, 4, 5, 6, 1, 2, 3]}, + index=list(range(9)), + ) + + @property + def pdf3(self): + return pd.DataFrame( + {"b": [1, 1, 1, 1, 1, 1, 1, 1, 1], "c": [1, 1, 1, 1, 1, 1, 1, 1, 1]}, + index=list(range(9)), + ) + + @property + def pdf4(self): + return pd.DataFrame( + {"e": [2, 2, 2, 2, 2, 2, 2, 2, 2], "f": [2, 2, 2, 2, 2, 2, 2, 2, 2]}, + index=list(range(9)), + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.from_pandas(self.pdf2) + + @property + def psdf3(self): + return ps.from_pandas(self.pdf3) + + @property + def psdf4(self): + return ps.from_pandas(self.pdf4) + + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_cov(self): + pser1 = pd.Series([0.90010907, 0.13484424, 0.62036035], index=[0, 1, 2]) + pser2 = pd.Series([0.12528585, 0.26962463, 0.51111198], index=[1, 2, 3]) + self._test_cov(pser1, pser2) + + pser1 = pd.Series([0.90010907, 0.13484424, 0.62036035], index=[0, 1, 2]) + pser2 = pd.Series([0.12528585, 0.26962463, 0.51111198, 0.32076008], index=[1, 2, 3, 4]) + self._test_cov(pser1, pser2) + + pser1 = pd.Series([0.90010907, 0.13484424, 0.62036035, 0.32076008], index=[0, 1, 2, 3]) + pser2 = pd.Series([0.12528585, 0.26962463], index=[1, 2]) + self._test_cov(pser1, pser2) + + psser1 = ps.from_pandas(pser1) + with self.assertRaisesRegex(TypeError, "unsupported type: "): + psser1.cov([0.12528585, 0.26962463, 0.51111198]) + with self.assertRaisesRegex( + TypeError, "unsupported type: " + ): + psser1.cov(pser2) + + def _test_cov(self, pser1, pser2): + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + pcov = pser1.cov(pser2) + pscov = psser1.cov(psser2) + self.assert_eq(pcov, pscov, almost=True) + + pcov = pser1.cov(pser2, min_periods=2) + pscov = psser1.cov(psser2, min_periods=2) + self.assert_eq(pcov, pscov, almost=True) + + pcov = pser1.cov(pser2, min_periods=3) + pscov = psser1.cov(psser2, min_periods=3) + self.assert_eq(pcov, pscov, almost=True) + + def test_corrwith(self): + df1 = ps.DataFrame({"A": [1, np.nan, 7, 8], "X": [5, 8, np.nan, 3], "C": [10, 4, 9, 3]}) + df2 = ps.DataFrame({"A": [5, 3, 6, 4], "B": [11, 2, 4, 3], "C": [4, 3, 8, np.nan]}) + self._test_corrwith(df1, df2) + self._test_corrwith((df1 + 1), df2.B) + self._test_corrwith((df1 + 1), (df2.B + 2)) + + # There was a regression in pandas 1.5.0, and fixed in pandas 1.5.1. + # Therefore, we only test the pandas 1.5.0 in different way. + # See https://github.com/pandas-dev/pandas/issues/49141 for the reported issue, + # and https://github.com/pandas-dev/pandas/pull/46174 for the initial PR that causes. + df_bool = ps.DataFrame({"A": [True, True, False, False], "B": [True, False, False, True]}) + ser_bool = ps.Series([True, True, False, True]) + if LooseVersion(pd.__version__) == LooseVersion("1.5.0"): + expected = ps.Series([0.5773502691896257, 0.5773502691896257], index=["B", "A"]) + self.assert_eq(df_bool.corrwith(ser_bool), expected, almost=True) + else: + self._test_corrwith(df_bool, ser_bool) + + self._test_corrwith(self.psdf1, self.psdf1) + self._test_corrwith(self.psdf1, self.psdf2) + self._test_corrwith(self.psdf2, self.psdf3) + self._test_corrwith(self.psdf3, self.psdf4) + + self._test_corrwith(self.psdf1, self.psdf1.a) + # There was a regression in pandas 1.5.0, and fixed in pandas 1.5.1. + # Therefore, we only test the pandas 1.5.0 in different way. + # See https://github.com/pandas-dev/pandas/issues/49141 for the reported issue, + # and https://github.com/pandas-dev/pandas/pull/46174 for the initial PR that causes. + if LooseVersion(pd.__version__) == LooseVersion("1.5.0"): + expected = ps.Series([-0.08827348295047496, 0.4413674147523748], index=["b", "a"]) + self.assert_eq(self.psdf1.corrwith(self.psdf2.b), expected, almost=True) + else: + self._test_corrwith(self.psdf1, self.psdf2.b) + + self._test_corrwith(self.psdf2, self.psdf3.c) + self._test_corrwith(self.psdf3, self.psdf4.f) + + def _test_corrwith(self, psdf, psobj): + pdf = psdf._to_pandas() + pobj = psobj._to_pandas() + for drop in [True, False]: + p_corr = pdf.corrwith(pobj, drop=drop) + ps_corr = psdf.corrwith(psobj, drop=drop) + self.assert_eq(p_corr.sort_index(), ps_corr.sort_index(), almost=True) + + +class DiffFramesCovCorrWithTests(DiffFramesCovCorrWithMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_cov_corrwith import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py new file mode 100644 index 0000000..7eeeada --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py @@ -0,0 +1,103 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesDotFrameMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_frame_dot(self): + pdf = pd.DataFrame([[0, 1, -2, -1], [1, 1, 1, 1]]) + psdf = ps.from_pandas(pdf) + + pser = pd.Series([1, 1, 2, 1]) + psser = ps.from_pandas(pser) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # Index reorder + pser = pser.reindex([1, 0, 2, 3]) + psser = ps.from_pandas(pser) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # ser with name + pser.name = "ser" + psser = ps.from_pandas(pser) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # df with MultiIndex as column (ser with MultiIndex) + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pser = pd.Series([1, 1, 2, 1], index=pidx) + pdf = pd.DataFrame([[0, 1, -2, -1], [1, 1, 1, 1]], columns=pidx) + psdf = ps.from_pandas(pdf) + psser = ps.from_pandas(pser) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # df with Index as column (ser with Index) + pidx = pd.Index([1, 2, 3, 4], name="number") + pser = pd.Series([1, 1, 2, 1], index=pidx) + pdf = pd.DataFrame([[0, 1, -2, -1], [1, 1, 1, 1]], columns=pidx) + psdf = ps.from_pandas(pdf) + psser = ps.from_pandas(pser) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # df with Index + pdf.index = pd.Index(["x", "y"], name="char") + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + # df with MultiIndex + pdf.index = pd.MultiIndex.from_arrays([[1, 1], ["red", "blue"]], names=("number", "color")) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.dot(psser), pdf.dot(pser)) + + pdf = pd.DataFrame([[1, 2], [3, 4]]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.dot(psdf[0]), pdf.dot(pdf[0])) + self.assert_eq(psdf.dot(psdf[0] * 10), pdf.dot(pdf[0] * 10)) + self.assert_eq((psdf + 1).dot(psdf[0] * 10), (pdf + 1).dot(pdf[0] * 10)) + + +class DiffFramesDotFrameTests(DiffFramesDotFrameMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_dot_frame import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py new file mode 100644 index 0000000..46d6d8e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py @@ -0,0 +1,141 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesDotSeriesMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_series_dot(self): + pser = pd.Series([90, 91, 85], index=[2, 4, 1]) + psser = ps.from_pandas(pser) + pser_other = pd.Series([90, 91, 85], index=[2, 4, 1]) + psser_other = ps.from_pandas(pser_other) + + self.assert_eq(psser.dot(psser_other), pser.dot(pser_other)) + + psser_other = ps.Series([90, 91, 85], index=[1, 2, 4]) + pser_other = pd.Series([90, 91, 85], index=[1, 2, 4]) + + self.assert_eq(psser.dot(psser_other), pser.dot(pser_other)) + + # length of index is different + psser_other = ps.Series([90, 91, 85, 100], index=[2, 4, 1, 0]) + with self.assertRaisesRegex(ValueError, "matrices are not aligned"): + psser.dot(psser_other) + + # for MultiIndex + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + psser = ps.from_pandas(pser) + pser_other = pd.Series([-450, 20, 12, -30, -250, 15, -320, 100, 3], index=midx) + psser_other = ps.from_pandas(pser_other) + self.assert_eq(psser.dot(psser_other), pser.dot(pser_other)) + + pser = pd.Series([0, 1, 2, 3]) + psser = ps.from_pandas(pser) + + # DataFrame "other" without Index/MultiIndex as columns + pdf = pd.DataFrame([[0, 1], [-2, 3], [4, -5], [6, 7]]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + + # DataFrame "other" with Index as columns + pdf.columns = pd.Index(["x", "y"]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + pdf.columns = pd.Index(["x", "y"], name="cols_name") + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + + pdf = pdf.reindex([1, 0, 2, 3]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + + # DataFrame "other" with MultiIndex as columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + pdf.columns = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y")], names=["cols_name1", "cols_name2"] + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + + psser = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}).b + pser = psser._to_pandas() + psdf = ps.DataFrame({"c": [7, 8, 9]}) + pdf = psdf._to_pandas() + self.assert_eq(psser.dot(psdf), pser.dot(pdf)) + + # SPARK-36968: ps.Series.dot raise "matrices are not aligned" if index is not same + pser = pd.Series([90, 91, 85], index=[0, 1, 2]) + psser = ps.from_pandas(pser) + pser_other = pd.Series([90, 91, 85], index=[0, 1, 3]) + psser_other = ps.from_pandas(pser_other) + pser_other2 = pd.Series([90, 91, 85, 100], index=[0, 1, 3, 5]) + psser_other2 = ps.from_pandas(pser_other2) + + with self.assertRaisesRegex(ValueError, "matrices are not aligned"): + psser.dot(psser_other) + + with ps.option_context("compute.eager_check", False), self.assertRaisesRegex( + ValueError, "matrices are not aligned" + ): + psser.dot(psser_other2) + + with ps.option_context("compute.eager_check", True), self.assertRaisesRegex( + ValueError, "matrices are not aligned" + ): + psser.dot(psser_other) + + with ps.option_context("compute.eager_check", False): + self.assert_eq(psser.dot(psser_other), 16381) + + +class DiffFramesDotSeriesTests(DiffFramesDotSeriesMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_dot_series import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_index.py new file mode 100644 index 0000000..11d737d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_index.py @@ -0,0 +1,109 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesIndexMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_index_ops(self): + pidx1 = pd.Index([1, 2, 3, 4, 5], name="x") + pidx2 = pd.Index([6, 7, 8, 9, 10], name="x") + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + + self.assert_eq(psidx1 * 10 + psidx2, pidx1 * 10 + pidx2) + self.assert_eq(psidx1.rename(None) * 10 + psidx2, pidx1.rename(None) * 10 + pidx2) + self.assert_eq(psidx1 * 10 + psidx2.rename(None), pidx1 * 10 + pidx2.rename(None)) + + pidx3 = pd.Index([11, 12, 13]) + psidx3 = ps.from_pandas(pidx3) + + with self.assertRaisesRegex( + ValueError, "operands could not be broadcast together with shapes" + ): + psidx1 + psidx3 + + pidx1 = pd.Index([1, 2, 3, 4, 5], name="a") + pidx2 = pd.Index([6, 7, 8, 9, 10], name="a") + pidx3 = pd.Index([11, 12, 13, 14, 15], name="x") + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + psidx3 = ps.from_pandas(pidx3) + + self.assert_eq(psidx1 * 10 + psidx2, pidx1 * 10 + pidx2) + self.assert_eq(psidx1 * 10 + psidx3, pidx1 * 10 + pidx3) + + def test_multi_index_column_assignment_frame(self): + pdf = pd.DataFrame({"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0]}) + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")]) + psdf = ps.DataFrame(pdf) + + psdf["c"] = ps.Series([10, 20, 30, 20]) + pdf["c"] = pd.Series([10, 20, 30, 20]) + + psdf[("d", "x")] = ps.Series([100, 200, 300, 200], name="1") + pdf[("d", "x")] = pd.Series([100, 200, 300, 200], name="1") + + psdf[("d", "y")] = ps.Series([1000, 2000, 3000, 2000], name=("1", "2")) + pdf[("d", "y")] = pd.Series([1000, 2000, 3000, 2000], name=("1", "2")) + + psdf["e"] = ps.Series([10000, 20000, 30000, 20000], name=("1", "2", "3")) + pdf["e"] = pd.Series([10000, 20000, 30000, 20000], name=("1", "2", "3")) + + psdf[[("f", "x"), ("f", "y")]] = ps.DataFrame( + {"1": [100000, 200000, 300000, 200000], "2": [1000000, 2000000, 3000000, 2000000]} + ) + pdf[[("f", "x"), ("f", "y")]] = pd.DataFrame( + {"1": [100000, 200000, 300000, 200000], "2": [1000000, 2000000, 3000000, 2000000]} + ) + + self.assert_eq(repr(psdf.sort_index()), repr(pdf)) + + with self.assertRaisesRegex(KeyError, "Key length \\(3\\) exceeds index depth \\(2\\)"): + psdf[("1", "2", "3")] = ps.Series([100, 200, 300, 200]) + + +class DiffFramesIndexTests(DiffFramesIndexMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_index import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_series.py new file mode 100644 index 0000000..d3a8df5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_series.py @@ -0,0 +1,136 @@ +# +# 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 pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesSeriesMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_series_repeat(self): + pser1 = pd.Series(["a", "b", "c"], name="a") + pser2 = pd.Series([10, 20, 30], name="rep") + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + self.assert_eq(psser1.repeat(psser2).sort_index(), pser1.repeat(pser2).sort_index()) + + def test_series_ops(self): + pser1 = pd.Series([1, 2, 3, 4, 5, 6, 7], name="x", index=[11, 12, 13, 14, 15, 16, 17]) + pser2 = pd.Series([1, 2, 3, 4, 5, 6, 7], name="x", index=[11, 12, 13, 14, 15, 16, 17]) + pidx1 = pd.Index([10, 11, 12, 13, 14, 15, 16], name="x") + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + psidx1 = ps.from_pandas(pidx1) + + self.assert_eq( + (psser1 + 1 + 10 * psser2).sort_index(), (pser1 + 1 + 10 * pser2).sort_index() + ) + self.assert_eq( + (psser1 + 1 + 10 * psser2.rename()).sort_index(), + (pser1 + 1 + 10 * pser2.rename()).sort_index(), + ) + self.assert_eq( + (psser1.rename() + 1 + 10 * psser2).sort_index(), + (pser1.rename() + 1 + 10 * pser2).sort_index(), + ) + self.assert_eq( + (psser1.rename() + 1 + 10 * psser2.rename()).sort_index(), + (pser1.rename() + 1 + 10 * pser2.rename()).sort_index(), + ) + + self.assert_eq(psser1 + 1 + 10 * psidx1, pser1 + 1 + 10 * pidx1) + self.assert_eq(psser1.rename() + 1 + 10 * psidx1, pser1.rename() + 1 + 10 * pidx1) + self.assert_eq(psser1 + 1 + 10 * psidx1.rename(None), pser1 + 1 + 10 * pidx1.rename(None)) + self.assert_eq( + psser1.rename() + 1 + 10 * psidx1.rename(None), + pser1.rename() + 1 + 10 * pidx1.rename(None), + ) + + self.assert_eq(psidx1 + 1 + 10 * psser1, pidx1 + 1 + 10 * pser1) + self.assert_eq(psidx1 + 1 + 10 * psser1.rename(), pidx1 + 1 + 10 * pser1.rename()) + self.assert_eq(psidx1.rename(None) + 1 + 10 * psser1, pidx1.rename(None) + 1 + 10 * pser1) + self.assert_eq( + psidx1.rename(None) + 1 + 10 * psser1.rename(), + pidx1.rename(None) + 1 + 10 * pser1.rename(), + ) + + pidx2 = pd.Index([11, 12, 13]) + psidx2 = ps.from_pandas(pidx2) + + with self.assertRaisesRegex( + ValueError, "operands could not be broadcast together with shapes" + ): + psser1 + psidx2 + + with self.assertRaisesRegex( + ValueError, "operands could not be broadcast together with shapes" + ): + psidx2 + psser1 + + def test_series_eq(self): + pser = pd.Series([1, 2, 3, 4, 5, 6], name="x") + psser = ps.from_pandas(pser) + + # other = Series + pandas_other = pd.Series([np.nan, 1, 3, 4, np.nan, 6], name="x") + pandas_on_spark_other = ps.from_pandas(pandas_other) + self.assert_eq(pser.eq(pandas_other), psser.eq(pandas_on_spark_other).sort_index()) + self.assert_eq(pser == pandas_other, (psser == pandas_on_spark_other).sort_index()) + + # other = Series with different Index + pandas_other = pd.Series( + [np.nan, 1, 3, 4, np.nan, 6], index=[10, 20, 30, 40, 50, 60], name="x" + ) + pandas_on_spark_other = ps.from_pandas(pandas_other) + self.assert_eq(pser.eq(pandas_other), psser.eq(pandas_on_spark_other).sort_index()) + + # other = Index + pandas_other = pd.Index([np.nan, 1, 3, 4, np.nan, 6], name="x") + pandas_on_spark_other = ps.from_pandas(pandas_other) + self.assert_eq(pser.eq(pandas_other), psser.eq(pandas_on_spark_other).sort_index()) + self.assert_eq(pser == pandas_other, (psser == pandas_on_spark_other).sort_index()) + + +class DiffFramesSeriesTests(DiffFramesSeriesMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_series import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py new file mode 100644 index 0000000..4de3c50 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py @@ -0,0 +1,125 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesSetItemFrameMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_frame_loc_setitem(self): + pdf_orig = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf_orig = ps.DataFrame(pdf_orig) + + pdf = pdf_orig.copy() + psdf = psdf_orig.copy() + pser1 = pdf.max_speed + pser2 = pdf.shield + psser1 = psdf.max_speed + psser2 = psdf.shield + + another_psdf = ps.DataFrame(pdf_orig) + + psdf.loc[["viper", "sidewinder"], ["shield"]] = -another_psdf.max_speed + pdf.loc[["viper", "sidewinder"], ["shield"]] = -pdf.max_speed + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf = pdf_orig.copy() + psdf = psdf_orig.copy() + pser1 = pdf.max_speed + pser2 = pdf.shield + psser1 = psdf.max_speed + psser2 = psdf.shield + psdf.loc[another_psdf.max_speed < 5, ["shield"]] = -psdf.max_speed + pdf.loc[pdf.max_speed < 5, ["shield"]] = -pdf.max_speed + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf = pdf_orig.copy() + psdf = psdf_orig.copy() + pser1 = pdf.max_speed + pser2 = pdf.shield + psser1 = psdf.max_speed + psser2 = psdf.shield + psdf.loc[another_psdf.max_speed < 5, ["shield"]] = -another_psdf.max_speed + pdf.loc[pdf.max_speed < 5, ["shield"]] = -pdf.max_speed + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + def test_frame_iloc_setitem(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.DataFrame(pdf) + another_psdf = ps.DataFrame(pdf) + + psdf.iloc[[0, 1, 2], 1] = -another_psdf.max_speed + pdf.iloc[[0, 1, 2], 1] = -pdf.max_speed + self.assert_eq(psdf, pdf) + + with self.assertRaisesRegex( + ValueError, + "shape mismatch", + ): + psdf.iloc[[1, 2], [1]] = -another_psdf.max_speed + + psdf.iloc[[0, 1, 2], 1] = 10 * another_psdf.max_speed + pdf.iloc[[0, 1, 2], 1] = 10 * pdf.max_speed + self.assert_eq(psdf, pdf) + + with self.assertRaisesRegex(ValueError, "shape mismatch"): + psdf.iloc[[0], 1] = 10 * another_psdf.max_speed + + +class DiffFramesSetItemFrameTests(DiffFramesSetItemFrameMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_setitem_frame import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py new file mode 100644 index 0000000..3473118 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py @@ -0,0 +1,217 @@ +# +# 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 pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class DiffFramesSetItemSeriesMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def test_series_loc_setitem(self): + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + + pser_another = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser_another = ps.from_pandas(pser_another) + + psser.loc[psser % 2 == 1] = -psser_another + pser.loc[pser % 2 == 1] = -pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + psser.loc[psser_another % 2 == 1] = -psser + pser.loc[pser_another % 2 == 1] = -pser + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + psser.loc[psser_another % 2 == 1] = -psser + pser.loc[pser_another % 2 == 1] = -pser + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + psser.loc[psser_another % 2 == 1] = -psser_another + pser.loc[pser_another % 2 == 1] = -pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + psser.loc[["viper", "sidewinder"]] = -psser_another + pser.loc[["viper", "sidewinder"]] = -pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + psser.loc[psser_another % 2 == 1] = 10 + pser.loc[pser_another % 2 == 1] = 10 + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + def test_series_iloc_setitem(self): + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + + pser1 = pser + 1 + psser1 = psser + 1 + + pser_another = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser_another = ps.from_pandas(pser_another) + + psser.iloc[[0, 1, 2]] = -psser_another + pser.iloc[[0, 1, 2]] = -pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaisesRegex( + ValueError, + "cannot set using a list-like indexer with a different length than the value", + ): + psser.iloc[[1, 2]] = -psser_another + + psser.iloc[[0, 1, 2]] = 10 * psser_another + pser.iloc[[0, 1, 2]] = 10 * pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaisesRegex( + ValueError, + "cannot set using a list-like indexer with a different length than the value", + ): + psser.iloc[[0]] = 10 * psser_another + + psser1.iloc[[0, 1, 2]] = -psser_another + pser1.iloc[[0, 1, 2]] = -pser_another + self.assert_eq(psser1, pser1) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaisesRegex( + ValueError, + "cannot set using a list-like indexer with a different length than the value", + ): + psser1.iloc[[1, 2]] = -psser_another + + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + + piloc = pser.iloc + kiloc = psser.iloc + + kiloc[[0, 1, 2]] = -psser_another + piloc[[0, 1, 2]] = -pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaisesRegex( + ValueError, + "cannot set using a list-like indexer with a different length than the value", + ): + kiloc[[1, 2]] = -psser_another + + kiloc[[0, 1, 2]] = 10 * psser_another + piloc[[0, 1, 2]] = 10 * pser_another + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaisesRegex( + ValueError, + "cannot set using a list-like indexer with a different length than the value", + ): + kiloc[[0]] = 10 * psser_another + + +class DiffFramesSetItemSeriesTests( + DiffFramesSetItemSeriesMixin, PandasOnSparkTestCase, SQLTestUtils +): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.diff_frames_ops.test_setitem_series import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_attrs.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_attrs.py new file mode 100644 index 0000000..9b9a25d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_attrs.py @@ -0,0 +1,384 @@ +# +# 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 datetime import datetime +from distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Attributes and underlying data' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#attributes-and-underlying-data +class FrameAttrsMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_column_names(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf.columns, pdf.columns) + self.assert_eq(psdf[["b", "a"]].columns, pdf[["b", "a"]].columns) + self.assert_eq(psdf["a"].name, pdf["a"].name) + self.assert_eq((psdf["a"] + 1).name, (pdf["a"] + 1).name) + + self.assert_eq((psdf.a + psdf.b).name, (pdf.a + pdf.b).name) + self.assert_eq((psdf.a + psdf.b.rename("a")).name, (pdf.a + pdf.b.rename("a")).name) + self.assert_eq((psdf.a + psdf.b.rename()).name, (pdf.a + pdf.b.rename()).name) + self.assert_eq((psdf.a.rename() + psdf.b).name, (pdf.a.rename() + pdf.b).name) + self.assert_eq( + (psdf.a.rename() + psdf.b.rename()).name, (pdf.a.rename() + pdf.b.rename()).name + ) + + def test_rename_columns(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7) + ) + psdf = ps.from_pandas(pdf) + + psdf.columns = ["x", "y"] + pdf.columns = ["x", "y"] + self.assert_eq(psdf.columns, pd.Index(["x", "y"])) + self.assert_eq(psdf, pdf) + self.assert_eq(psdf._internal.data_spark_column_names, ["x", "y"]) + self.assert_eq(psdf.to_spark().columns, ["x", "y"]) + self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "x", "y"]) + + columns = pdf.columns + columns.name = "lvl_1" + + psdf.columns = columns + self.assert_eq(psdf.columns.names, ["lvl_1"]) + self.assert_eq(psdf, pdf) + + msg = "Length mismatch: Expected axis has 2 elements, new values have 4 elements" + with self.assertRaisesRegex(ValueError, msg): + psdf.columns = [1, 2, 3, 4] + + # Multi-index columns + pdf = pd.DataFrame( + {("A", "0"): [1, 2, 2, 3], ("B", "1"): [1, 2, 3, 4]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + + columns = pdf.columns + self.assert_eq(psdf.columns, columns) + self.assert_eq(psdf, pdf) + + pdf.columns = ["x", "y"] + psdf.columns = ["x", "y"] + self.assert_eq(psdf.columns, pd.Index(["x", "y"])) + self.assert_eq(psdf, pdf) + self.assert_eq(psdf._internal.data_spark_column_names, ["x", "y"]) + self.assert_eq(psdf.to_spark().columns, ["x", "y"]) + self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "x", "y"]) + + pdf.columns = columns + psdf.columns = columns + self.assert_eq(psdf.columns, columns) + self.assert_eq(psdf, pdf) + self.assert_eq(psdf._internal.data_spark_column_names, ["(A, 0)", "(B, 1)"]) + self.assert_eq(psdf.to_spark().columns, ["(A, 0)", "(B, 1)"]) + self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "(A, 0)", "(B, 1)"]) + + columns.names = ["lvl_1", "lvl_2"] + + psdf.columns = columns + self.assert_eq(psdf.columns.names, ["lvl_1", "lvl_2"]) + self.assert_eq(psdf, pdf) + self.assert_eq(psdf._internal.data_spark_column_names, ["(A, 0)", "(B, 1)"]) + self.assert_eq(psdf.to_spark().columns, ["(A, 0)", "(B, 1)"]) + self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "(A, 0)", "(B, 1)"]) + + def test_multi_index_dtypes(self): + # SPARK-36930: Support ps.MultiIndex.dtypes + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + pmidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + psmidx = ps.from_pandas(pmidx) + + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq(psmidx.dtypes, pmidx.dtypes) + else: + expected = pd.Series([np.dtype("int64"), np.dtype("O")], index=["number", "color"]) + self.assert_eq(psmidx.dtypes, expected) + + # multiple labels + pmidx = pd.MultiIndex.from_arrays(arrays, names=[("zero", "first"), ("one", "second")]) + psmidx = ps.from_pandas(pmidx) + + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + if LooseVersion(pd.__version__) not in (LooseVersion("1.4.1"), LooseVersion("1.4.2")): + self.assert_eq(psmidx.dtypes, pmidx.dtypes) + else: + expected = pd.Series( + [np.dtype("int64"), np.dtype("O")], + index=pd.Index([("zero", "first"), ("one", "second")]), + ) + self.assert_eq(psmidx.dtypes, expected) + + def test_multi_index_dtypes_not_unique_name(self): + # Regression test for https://github.com/pandas-dev/pandas/issues/45174 + pmidx = pd.MultiIndex.from_arrays([[1], [2]], names=[1, 1]) + psmidx = ps.from_pandas(pmidx) + + if LooseVersion(pd.__version__) < LooseVersion("1.4"): + expected = pd.Series( + [np.dtype("int64"), np.dtype("int64")], + index=[1, 1], + ) + self.assert_eq(psmidx.dtypes, expected) + else: + self.assert_eq(psmidx.dtypes, pmidx.dtypes) + + def test_dtype(self): + pdf = pd.DataFrame( + { + "a": list("abc"), + "b": list(range(1, 4)), + "c": np.arange(3, 6).astype("i1"), + "d": np.arange(4.0, 7.0, dtype="float64"), + "e": [True, False, True], + "f": pd.date_range("20130101", periods=3), + }, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + self.assertTrue((psdf.dtypes == pdf.dtypes).all()) + + # multi-index columns + columns = pd.MultiIndex.from_tuples(zip(list("xxxyyz"), list("abcdef"))) + pdf.columns = columns + psdf.columns = columns + self.assertTrue((psdf.dtypes == pdf.dtypes).all()) + + def test_axes(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.axes, psdf.axes) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.axes, psdf.axes) + + def test_inplace(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + pdf["a"] = pdf["a"] + 10 + psdf["a"] = psdf["a"] + 10 + + self.assert_eq(psdf, pdf) + # SPARK-38946: Since Spark 3.4, df.__setitem__ generate a new dataframe to follow + # pandas 1.4 behaviors + if LooseVersion(pd.__version__) >= LooseVersion("1.4.0"): + self.assert_eq(psser, pser) + else: + # Follow pandas latest behavior + with self.assertRaisesRegex(AssertionError, "Series are different"): + self.assert_eq(psser, pser) + + def test_dataframe_multiindex_columns(self): + pdf = pd.DataFrame( + { + ("x", "a", "1"): [1, 2, 3], + ("x", "b", "2"): [4, 5, 6], + ("y.z", "c.d", "3"): [7, 8, 9], + ("x", "b", "4"): [10, 11, 12], + }, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf["x"], pdf["x"]) + self.assert_eq(psdf["y.z"], pdf["y.z"]) + self.assert_eq(psdf["x"]["b"], pdf["x"]["b"]) + self.assert_eq(psdf["x"]["b"]["2"], pdf["x"]["b"]["2"]) + + self.assert_eq(psdf.x, pdf.x) + self.assert_eq(psdf.x.b, pdf.x.b) + self.assert_eq(psdf.x.b["2"], pdf.x.b["2"]) + + self.assertRaises(KeyError, lambda: psdf["z"]) + self.assertRaises(AttributeError, lambda: psdf.z) + + self.assert_eq(psdf[("x",)], pdf[("x",)]) + self.assert_eq(psdf[("x", "a")], pdf[("x", "a")]) + self.assert_eq(psdf[("x", "a", "1")], pdf[("x", "a", "1")]) + + def test_dataframe_column_level_name(self): + column = pd.Index(["A", "B", "C"], name="X") + pdf = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=column, index=np.random.rand(2)) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf.columns.names, pdf.columns.names) + self.assert_eq(psdf._to_pandas().columns.names, pdf.columns.names) + + def test_dataframe_multiindex_names_level(self): + columns = pd.MultiIndex.from_tuples( + [("X", "A", "Z"), ("X", "B", "Z"), ("Y", "C", "Z"), ("Y", "D", "Z")], + names=["lvl_1", "lvl_2", "lv_3"], + ) + pdf = pd.DataFrame( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16], [17, 18, 19, 20]], + columns=columns, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.columns.names, pdf.columns.names) + self.assert_eq(psdf._to_pandas().columns.names, pdf.columns.names) + + psdf1 = ps.from_pandas(pdf) + self.assert_eq(psdf1.columns.names, pdf.columns.names) + + self.assertRaises( + AssertionError, + lambda: ps.DataFrame(psdf1._internal.copy(column_label_names=("level",))), + ) + + self.assert_eq(psdf["X"], pdf["X"]) + self.assert_eq(psdf["X"].columns.names, pdf["X"].columns.names) + self.assert_eq(psdf["X"]._to_pandas().columns.names, pdf["X"].columns.names) + self.assert_eq(psdf["X"]["A"], pdf["X"]["A"]) + self.assert_eq(psdf["X"]["A"].columns.names, pdf["X"]["A"].columns.names) + self.assert_eq(psdf["X"]["A"]._to_pandas().columns.names, pdf["X"]["A"].columns.names) + self.assert_eq(psdf[("X", "A")], pdf[("X", "A")]) + self.assert_eq(psdf[("X", "A")].columns.names, pdf[("X", "A")].columns.names) + self.assert_eq(psdf[("X", "A")]._to_pandas().columns.names, pdf[("X", "A")].columns.names) + self.assert_eq(psdf[("X", "A", "Z")], pdf[("X", "A", "Z")]) + + def test_repr_cache_invalidation(self): + # If there is any cache, inplace operations should invalidate it. + df = ps.range(10) + df.__repr__() + df["a"] = df["id"] + self.assertEqual(df.__repr__(), df._to_pandas().__repr__()) + + def test_repr_html_cache_invalidation(self): + # If there is any cache, inplace operations should invalidate it. + df = ps.range(10) + df._repr_html_() + df["a"] = df["id"] + self.assertEqual(df._repr_html_(), df._to_pandas()._repr_html_()) + + def test_assign(self): + pdf, psdf = self.df_pair + + psdf["w"] = 1.0 + pdf["w"] = 1.0 + + self.assert_eq(psdf, pdf) + + psdf.w = 10.0 + pdf.w = 10.0 + + self.assert_eq(psdf, pdf) + + psdf[1] = 1.0 + pdf[1] = 1.0 + + self.assert_eq(psdf, pdf) + + psdf = psdf.assign(a=psdf["a"] * 2) + pdf = pdf.assign(a=pdf["a"] * 2) + + self.assert_eq(psdf, pdf) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "w"), ("y", "v")]) + pdf.columns = columns + psdf.columns = columns + + psdf[("a", "c")] = "def" + pdf[("a", "c")] = "def" + + self.assert_eq(psdf, pdf) + + psdf = psdf.assign(Z="ZZ") + pdf = pdf.assign(Z="ZZ") + + self.assert_eq(psdf, pdf) + + psdf["x"] = "ghi" + pdf["x"] = "ghi" + + self.assert_eq(psdf, pdf) + + def test_attributes(self): + psdf = self.psdf + + self.assertIn("a", dir(psdf)) + self.assertNotIn("foo", dir(psdf)) + self.assertRaises(AttributeError, lambda: psdf.foo) + + psdf = ps.DataFrame({"a b c": [1, 2, 3]}) + self.assertNotIn("a b c", dir(psdf)) + psdf = ps.DataFrame({"a": [1, 2], 5: [1, 2]}) + self.assertIn("a", dir(psdf)) + self.assertNotIn(5, dir(psdf)) + + def test_empty_timestamp(self): + pdf = pd.DataFrame( + { + "t": [ + datetime(2019, 1, 1, 0, 0, 0), + datetime(2019, 1, 2, 0, 0, 0), + datetime(2019, 1, 3, 0, 0, 0), + ] + }, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf[psdf["t"] != psdf["t"]], pdf[pdf["t"] != pdf["t"]]) + self.assert_eq(psdf[psdf["t"] != psdf["t"]].dtypes, pdf[pdf["t"] != pdf["t"]].dtypes) + + +class FrameAttrsTests(FrameAttrsMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_attrs import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_constructor.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_constructor.py new file mode 100644 index 0000000..139218b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_constructor.py @@ -0,0 +1,598 @@ +# +# 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 datetime import datetime, timedelta +from distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.typedef.typehints import ( + extension_dtypes, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) +from pyspark.pandas.utils import is_testing + +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Constructor' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#constructor +# as well as extensions. +class FrameConstructorMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_dataframe(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf["a"] + 1, pdf["a"] + 1) + + self.assert_eq(psdf.columns, pd.Index(["a", "b"])) + + self.assert_eq(psdf[psdf["b"] > 2], pdf[pdf["b"] > 2]) + self.assert_eq(-psdf[psdf["b"] > 2], -pdf[pdf["b"] > 2]) + self.assert_eq(psdf[["a", "b"]], pdf[["a", "b"]]) + self.assert_eq(psdf.a, pdf.a) + self.assert_eq(psdf.b.mean(), pdf.b.mean()) + self.assert_eq(psdf.b.var(), pdf.b.var()) + self.assert_eq(psdf.b.std(), pdf.b.std()) + + pdf, psdf = self.df_pair + self.assert_eq(psdf[["a", "b"]], pdf[["a", "b"]]) + + self.assertEqual(psdf.a.notnull().rename("x").name, "x") + + # check ps.DataFrame(ps.Series) + pser = pd.Series([1, 2, 3], name="x", index=np.random.rand(3)) + psser = ps.from_pandas(pser) + self.assert_eq(pd.DataFrame(pser), ps.DataFrame(psser)) + + # check ps.DataFrame(ps.Series) with `columns` + self.assert_eq(pd.DataFrame(pser, columns=["x"]), ps.DataFrame(psser, columns=["x"])) + self.assert_eq(pd.DataFrame(pser, columns=("x",)), ps.DataFrame(psser, columns=("x",))) + self.assert_eq( + pd.DataFrame(pser, columns={"x": None}), ps.DataFrame(psser, columns={"x": None}) + ) + + # check psdf[pd.Index] + pdf, psdf = self.df_pair + column_mask = pdf.columns.isin(["a", "b"]) + index_cols = pdf.columns[column_mask] + self.assert_eq(psdf[index_cols], pdf[index_cols]) + + if is_testing(): + err_msg = "pandas-on-Spark doesn't allow columns to be created via a new attribute name" + with self.assertRaisesRegex(AssertionError, err_msg): + psdf.X = [10, 20, 30, 40, 50, 60, 70, 80, 90] + else: + with self.assertWarns(UserWarning): + psdf.X = [10, 20, 30, 40, 50, 60, 70, 80, 90] + # If a new column is created, the following test would fail. + # It means that the pandas have changed their behavior, so we should follow. + self.assert_eq(pdf, psdf) + + def test_creation_index(self): + data = np.random.randn(5, 3) + + # test local data with pd.Index + self.assert_eq( + ps.DataFrame(data=[1, 2], index=pd.Index([1, 2])), + pd.DataFrame(data=[1, 2], index=pd.Index([1, 2])), + ) + self.assert_eq( + ps.DataFrame(data=[1, 2], index=pd.Index([2, 3])), + pd.DataFrame(data=[1, 2], index=pd.Index([2, 3])), + ) + self.assert_eq( + ps.DataFrame(data=[1, 2], index=pd.Index([3, 4])), + pd.DataFrame(data=[1, 2], index=pd.Index([3, 4])), + ) + self.assert_eq( + ps.DataFrame(data=data, index=pd.Index([1, 2, 3, 5, 6])), + pd.DataFrame(data=data, index=pd.Index([1, 2, 3, 5, 6])), + ) + + # test local data with ps.Index + self.assert_eq( + ps.DataFrame(data=[1, 2], index=ps.Index([1, 2])), + pd.DataFrame(data=[1, 2], index=pd.Index([1, 2])), + ) + self.assert_eq( + ps.DataFrame(data=[1, 2], index=ps.Index([2, 3])), + pd.DataFrame(data=[1, 2], index=pd.Index([2, 3])), + ) + self.assert_eq( + ps.DataFrame(data=[1, 2], index=ps.Index([3, 4])), + pd.DataFrame(data=[1, 2], index=pd.Index([3, 4])), + ) + self.assert_eq( + ps.DataFrame(data=data, index=ps.Index([1, 2, 3, 5, 6])), + pd.DataFrame(data=data, index=pd.Index([1, 2, 3, 5, 6])), + ) + + err_msg = "Cannot combine the series or dataframe" + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with ps.Index + ps.DataFrame(data=ps.DataFrame([1, 2]), index=ps.Index([1, 2])) + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with pd.Index + ps.DataFrame(data=ps.DataFrame([1, 2]), index=pd.Index([3, 4])) + + with ps.option_context("compute.ops_on_diff_frames", True): + # test pd.DataFrame with pd.Index + self.assert_eq( + ps.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([0, 1])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([0, 1])), + ) + self.assert_eq( + ps.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([1, 2])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([1, 2])), + ) + + # test ps.DataFrame with ps.Index + self.assert_eq( + ps.DataFrame(data=ps.DataFrame([1, 2]), index=ps.Index([0, 1])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([0, 1])), + ) + self.assert_eq( + ps.DataFrame(data=ps.DataFrame([1, 2]), index=ps.Index([1, 2])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([1, 2])), + ) + + # test ps.DataFrame with pd.Index + self.assert_eq( + ps.DataFrame(data=ps.DataFrame([1, 2]), index=pd.Index([0, 1])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([0, 1])), + ) + self.assert_eq( + ps.DataFrame(data=ps.DataFrame([1, 2]), index=pd.Index([1, 2])), + pd.DataFrame(data=pd.DataFrame([1, 2]), index=pd.Index([1, 2])), + ) + + # test with multi data columns + pdf = pd.DataFrame(data=data, columns=["A", "B", "C"]) + psdf = ps.from_pandas(pdf) + + # test with pd.DataFrame and pd.Index + self.assert_eq( + ps.DataFrame(data=pdf, index=pd.Index([2, 3, 4, 5, 6])), + pd.DataFrame(data=pdf, index=pd.Index([2, 3, 4, 5, 6])), + ) + + # test with pd.DataFrame and ps.Index + self.assert_eq( + ps.DataFrame(data=pdf, index=ps.Index([2, 3, 4, 5, 6])), + pd.DataFrame(data=pdf, index=pd.Index([2, 3, 4, 5, 6])), + ) + + with ps.option_context("compute.ops_on_diff_frames", True): + # test with ps.DataFrame and pd.Index + self.assert_eq( + ps.DataFrame(data=psdf, index=pd.Index([2, 3, 4, 5, 6])), + pd.DataFrame(data=pdf, index=pd.Index([2, 3, 4, 5, 6])), + ) + + # test with ps.DataFrame and ps.Index + self.assert_eq( + ps.DataFrame(data=psdf, index=ps.Index([2, 3, 4, 5, 6])), + pd.DataFrame(data=pdf, index=pd.Index([2, 3, 4, 5, 6])), + ) + + # test String Index + pdf = pd.DataFrame( + data={ + "s": ["Hello", "World", "Databricks"], + "x": [2002, 2003, 2004], + } + ) + pdf = pdf.set_index("s") + pdf.index.name = None + psdf = ps.from_pandas(pdf) + + # test with pd.DataFrame and pd.Index + self.assert_eq( + ps.DataFrame(data=pdf, index=pd.Index(["Hello", "Universe", "Databricks"])), + pd.DataFrame(data=pdf, index=pd.Index(["Hello", "Universe", "Databricks"])), + ) + + # test with pd.DataFrame and ps.Index + self.assert_eq( + ps.DataFrame(data=pdf, index=ps.Index(["Hello", "Universe", "Databricks"])), + pd.DataFrame(data=pdf, index=pd.Index(["Hello", "Universe", "Databricks"])), + ) + + with ps.option_context("compute.ops_on_diff_frames", True): + # test with ps.DataFrame and pd.Index + self.assert_eq( + ps.DataFrame( + data=psdf, index=pd.Index(["Hello", "Universe", "Databricks"]) + ).sort_index(), + pd.DataFrame( + data=pdf, index=pd.Index(["Hello", "Universe", "Databricks"]) + ).sort_index(), + ) + + # test with ps.DataFrame and ps.Index + self.assert_eq( + ps.DataFrame( + data=psdf, index=ps.Index(["Hello", "Universe", "Databricks"]) + ).sort_index(), + pd.DataFrame( + data=pdf, index=pd.Index(["Hello", "Universe", "Databricks"]) + ).sort_index(), + ) + + # test DatetimeIndex + pdf = pd.DataFrame( + data={ + "t": [ + datetime(2022, 9, 1, 0, 0, 0, 0), + datetime(2022, 9, 2, 0, 0, 0, 0), + datetime(2022, 9, 3, 0, 0, 0, 0), + ], + "x": [2002, 2003, 2004], + } + ) + pdf = pdf.set_index("t") + pdf.index.name = None + psdf = ps.from_pandas(pdf) + + # test with pd.DataFrame and pd.DatetimeIndex + self.assert_eq( + ps.DataFrame( + data=pdf, + index=pd.DatetimeIndex(["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"]), + ), + pd.DataFrame( + data=pdf, + index=pd.DatetimeIndex(["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"]), + ), + ) + + # test with pd.DataFrame and ps.DatetimeIndex + self.assert_eq( + ps.DataFrame( + data=pdf, + index=ps.DatetimeIndex(["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"]), + ), + pd.DataFrame( + data=pdf, + index=pd.DatetimeIndex(["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"]), + ), + ) + + with ps.option_context("compute.ops_on_diff_frames", True): + # test with ps.DataFrame and pd.DatetimeIndex + self.assert_eq( + ps.DataFrame( + data=psdf, + index=pd.DatetimeIndex( + ["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"] + ), + ), + pd.DataFrame( + data=pdf, + index=pd.DatetimeIndex( + ["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"] + ), + ), + ) + + # test with ps.DataFrame and ps.DatetimeIndex + self.assert_eq( + ps.DataFrame( + data=psdf, + index=ps.DatetimeIndex( + ["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"] + ), + ), + pd.DataFrame( + data=pdf, + index=pd.DatetimeIndex( + ["2022-08-31", "2022-09-02", "2022-09-03", "2022-09-05"] + ), + ), + ) + + # test MultiIndex + # test local data with ps.MultiIndex + self.assert_eq( + ps.DataFrame(data=[1, 2], index=ps.MultiIndex.from_tuples([(1, 3), (2, 4)])), + pd.DataFrame(data=[1, 2], index=pd.MultiIndex.from_tuples([(1, 3), (2, 4)])), + ) + + # test distributed data with ps.MultiIndex + err_msg = "Cannot combine a Distributed Dataset with a MultiIndex" + with ps.option_context("compute.ops_on_diff_frames", True): + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with ps.Index + ps.DataFrame( + data=ps.DataFrame([1, 2]), index=ps.MultiIndex.from_tuples([(1, 3), (2, 4)]) + ) + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with pd.Index + ps.DataFrame( + data=ps.DataFrame([1, 2]), index=ps.MultiIndex.from_tuples([(1, 3), (2, 4)]) + ) + + def test_creation_index_same_anchor(self): + pdf = pd.DataFrame( + { + "a": [1, 2, None, 4], + "b": [1, None, None, 4], + "c": [1, 2, None, None], + "d": [None, 2, None, 4], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.DataFrame(data=psdf, index=psdf.index), + pd.DataFrame(data=pdf, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf + 1, index=psdf.index), + pd.DataFrame(data=pdf + 1, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf[["a", "c"]] * 2, index=psdf.index), + pd.DataFrame(data=pdf[["a", "c"]] * 2, index=pdf.index), + ) + + # test String Index + pdf = pd.DataFrame( + data={"s": ["Hello", "World", "Databricks"], "x": [2002, 2003, 2004], "y": [4, 5, 6]} + ) + pdf = pdf.set_index("s") + pdf.index.name = None + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.DataFrame(data=psdf, index=psdf.index), + pd.DataFrame(data=pdf, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf + 1, index=psdf.index), + pd.DataFrame(data=pdf + 1, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf[["y"]] * 2, index=psdf.index), + pd.DataFrame(data=pdf[["y"]] * 2, index=pdf.index), + ) + + # test DatetimeIndex + pdf = pd.DataFrame( + data={ + "t": [ + datetime(2022, 9, 1, 0, 0, 0, 0), + datetime(2022, 9, 2, 0, 0, 0, 0), + datetime(2022, 9, 3, 0, 0, 0, 0), + ], + "x": [2002, 2003, 2004], + "y": [4, 5, 6], + } + ) + pdf = pdf.set_index("t") + pdf.index.name = None + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.DataFrame(data=psdf, index=psdf.index), + pd.DataFrame(data=pdf, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf + 1, index=psdf.index), + pd.DataFrame(data=pdf + 1, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf[["y"]] * 2, index=psdf.index), + pd.DataFrame(data=pdf[["y"]] * 2, index=pdf.index), + ) + + # test TimedeltaIndex + pdf = pd.DataFrame( + data={ + "t": [ + timedelta(1), + timedelta(3), + timedelta(5), + ], + "x": [2002, 2003, 2004], + "y": [4, 5, 6], + } + ) + pdf = pdf.set_index("t") + pdf.index.name = None + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.DataFrame(data=psdf, index=psdf.index), + pd.DataFrame(data=pdf, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf + 1, index=psdf.index), + pd.DataFrame(data=pdf + 1, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf[["y"]] * 2, index=psdf.index), + pd.DataFrame(data=pdf[["y"]] * 2, index=pdf.index), + ) + + # test CategoricalIndex + pdf = pd.DataFrame( + data={ + "z": [-1, -2, -3, -4], + "x": [2002, 2003, 2004, 2005], + "y": [4, 5, 6, 7], + }, + index=pd.CategoricalIndex(["a", "c", "b", "a"], categories=["a", "b", "c"]), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.DataFrame(data=psdf, index=psdf.index), + pd.DataFrame(data=pdf, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf + 1, index=psdf.index), + pd.DataFrame(data=pdf + 1, index=pdf.index), + ) + self.assert_eq( + ps.DataFrame(data=psdf[["y"]] * 2, index=psdf.index), + pd.DataFrame(data=pdf[["y"]] * 2, index=pdf.index), + ) + + # test distributed data with ps.MultiIndex + pdf = pd.DataFrame( + data={ + "z": [-1, -2, -3, -4], + "x": [2002, 2003, 2004, 2005], + "y": [4, 5, 6, 7], + }, + index=pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z"), ("a", "x")]), + ) + psdf = ps.from_pandas(pdf) + + err_msg = "Cannot combine a Distributed Dataset with a MultiIndex" + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with ps.MultiIndex + ps.DataFrame(data=psdf, index=psdf.index) + with self.assertRaisesRegex(ValueError, err_msg): + # test ps.DataFrame with pd.MultiIndex + ps.DataFrame(data=psdf, index=pdf.index) + + def _check_extension(self, psdf, pdf): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(psdf, pdf, check_exact=False) + for dtype in psdf.dtypes: + self.assertTrue(isinstance(dtype, extension_dtypes)) + else: + self.assert_eq(psdf, pdf) + + @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") + def test_extension_dtypes(self): + pdf = pd.DataFrame( + { + "a": pd.Series([1, 2, None, 4], dtype="Int8"), + "b": pd.Series([1, None, None, 4], dtype="Int16"), + "c": pd.Series([1, 2, None, None], dtype="Int32"), + "d": pd.Series([None, 2, None, 4], dtype="Int64"), + } + ) + psdf = ps.from_pandas(pdf) + + self._check_extension(psdf, pdf) + self._check_extension(psdf + psdf, pdf + pdf) + + @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") + def test_astype_extension_dtypes(self): + pdf = pd.DataFrame( + { + "a": [1, 2, None, 4], + "b": [1, None, None, 4], + "c": [1, 2, None, None], + "d": [None, 2, None, 4], + } + ) + psdf = ps.from_pandas(pdf) + + astype = {"a": "Int8", "b": "Int16", "c": "Int32", "d": "Int64"} + + self._check_extension(psdf.astype(astype), pdf.astype(astype)) + + @unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" + ) + def test_extension_object_dtypes(self): + pdf = pd.DataFrame( + { + "a": pd.Series(["a", "b", None, "c"], dtype="string"), + "b": pd.Series([True, None, False, True], dtype="boolean"), + } + ) + psdf = ps.from_pandas(pdf) + + self._check_extension(psdf, pdf) + + @unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" + ) + def test_astype_extension_object_dtypes(self): + pdf = pd.DataFrame({"a": ["a", "b", None, "c"], "b": [True, None, False, True]}) + psdf = ps.from_pandas(pdf) + + astype = {"a": "string", "b": "boolean"} + + self._check_extension(psdf.astype(astype), pdf.astype(astype)) + + @unittest.skipIf( + not extension_float_dtypes_available, "pandas extension float dtypes are not available" + ) + def test_extension_float_dtypes(self): + pdf = pd.DataFrame( + { + "a": pd.Series([1.0, 2.0, None, 4.0], dtype="Float32"), + "b": pd.Series([1.0, None, 3.0, 4.0], dtype="Float64"), + } + ) + psdf = ps.from_pandas(pdf) + + self._check_extension(psdf, pdf) + self._check_extension(psdf + 1, pdf + 1) + self._check_extension(psdf + psdf, pdf + pdf) + + @unittest.skipIf( + not extension_float_dtypes_available, "pandas extension float dtypes are not available" + ) + def test_astype_extension_float_dtypes(self): + pdf = pd.DataFrame({"a": [1.0, 2.0, None, 4.0], "b": [1.0, None, 3.0, 4.0]}) + psdf = ps.from_pandas(pdf) + + astype = {"a": "Float32", "b": "Float64"} + + self._check_extension(psdf.astype(astype), pdf.astype(astype)) + + +class FrameConstructorTests(FrameConstructorMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_constructor import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_conversion.py new file mode 100644 index 0000000..116a7d3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_conversion.py @@ -0,0 +1,73 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Conversion' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#conversion +class FrameConversionMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_astype(self): + psdf = self.psdf + + msg = "Only a column name can be used for the key in a dtype mappings argument." + with self.assertRaisesRegex(KeyError, msg): + psdf.astype({"c": float}) + + def test_isnull(self): + pdf = pd.DataFrame( + {"x": [1, 2, 3, 4, None, 6], "y": list("abdabd")}, index=np.random.rand(6) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.notnull(), pdf.notnull()) + self.assert_eq(psdf.isnull(), pdf.isnull()) + + +class FrameConversionTests(FrameConversionMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_conversion import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reindexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reindexing.py new file mode 100644 index 0000000..ea9a75b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reindexing.py @@ -0,0 +1,869 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd +from pandas.tseries.offsets import DateOffset + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Reindexing / Selection / Label manipulation' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#reindexing-selection-label-manipulation +class FrameReindexingMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_add_prefix(self): + pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [3, 4, 5, 6]}, index=np.random.rand(4)) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.add_prefix("col_"), psdf.add_prefix("col_")) + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.add_prefix("col_"), psdf.add_prefix("col_")) + + def test_add_suffix(self): + pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [3, 4, 5, 6]}, index=np.random.rand(4)) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.add_suffix("first_series"), psdf.add_suffix("first_series")) + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.add_suffix("first_series"), psdf.add_suffix("first_series")) + + def test_at_time(self): + idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") + pdf = pd.DataFrame({"A": [1, 2, 3, 4]}, index=idx) + psdf = ps.from_pandas(pdf) + psdf.at_time("0:20") + self.assert_eq( + pdf.at_time("0:20").sort_index(), + psdf.at_time("0:20").sort_index(), + ) + + # Index name is 'ts' + pdf.index.name = "ts" + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.at_time("0:20").sort_index(), + psdf.at_time("0:20").sort_index(), + ) + + # Index name is 'ts', column label is 'index' + pdf.columns = pd.Index(["index"]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.at_time("0:40").sort_index(), + psdf.at_time("0:40").sort_index(), + ) + + # Both index name and column label are 'index' + pdf.index.name = "index" + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.at_time("0:40").sort_index(), + psdf.at_time("0:40").sort_index(), + ) + + # Index name is 'index', column label is ('X', 'A') + pdf.columns = pd.MultiIndex.from_arrays([["X"], ["A"]]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.at_time("0:40").sort_index(), + psdf.at_time("0:40").sort_index(), + ) + + with self.assertRaisesRegex(NotImplementedError, "'asof' argument is not supported"): + psdf.at_time("0:15", asof=True) + + with self.assertRaisesRegex(NotImplementedError, "at_time currently only works for axis=0"): + psdf.at_time("0:15", axis=1) + + psdf = ps.DataFrame({"A": [1, 2, 3, 4]}) + with self.assertRaisesRegex(TypeError, "Index must be DatetimeIndex"): + psdf.at_time("0:15") + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43557): Enable DataFrameSlowTests.test_between_time for pandas 2.0.0.", + ) + def test_between_time(self): + idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") + pdf = pd.DataFrame({"A": [1, 2, 3, 4]}, index=idx) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.between_time("0:15", "0:45").sort_index(), + psdf.between_time("0:15", "0:45").sort_index(), + ) + + pdf.index.name = "ts" + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.between_time("0:15", "0:45").sort_index(), + psdf.between_time("0:15", "0:45").sort_index(), + ) + + # Column label is 'index' + pdf.columns = pd.Index(["index"]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.between_time("0:15", "0:45").sort_index(), + psdf.between_time("0:15", "0:45").sort_index(), + ) + + # Both index name and column label are 'index' + pdf.index.name = "index" + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.between_time("0:15", "0:45").sort_index(), + psdf.between_time("0:15", "0:45").sort_index(), + ) + + # Index name is 'index', column label is ('X', 'A') + pdf.columns = pd.MultiIndex.from_arrays([["X"], ["A"]]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.between_time("0:15", "0:45").sort_index(), + psdf.between_time("0:15", "0:45").sort_index(), + ) + + with self.assertRaisesRegex( + NotImplementedError, "between_time currently only works for axis=0" + ): + psdf.between_time("0:15", "0:45", axis=1) + + psdf = ps.DataFrame({"A": [1, 2, 3, 4]}) + with self.assertRaisesRegex(TypeError, "Index must be DatetimeIndex"): + psdf.between_time("0:15", "0:45") + + def test_drop(self): + pdf = pd.DataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]}, index=np.random.rand(2)) + psdf = ps.from_pandas(pdf) + + # Assert 'labels' or 'columns' parameter is set + expected_error_message = "Need to specify at least one of 'labels' or 'columns'" + with self.assertRaisesRegex(ValueError, expected_error_message): + psdf.drop() + + # + # Drop columns + # + + # Assert using a str for 'labels' works + self.assert_eq(psdf.drop("x", axis=1), pdf.drop("x", axis=1)) + self.assert_eq((psdf + 1).drop("x", axis=1), (pdf + 1).drop("x", axis=1)) + # Assert using a list for 'labels' works + self.assert_eq(psdf.drop(["y", "z"], axis=1), pdf.drop(["y", "z"], axis=1)) + self.assert_eq(psdf.drop(["x", "y", "z"], axis=1), pdf.drop(["x", "y", "z"], axis=1)) + # Assert using 'columns' instead of 'labels' produces the same results + self.assert_eq(psdf.drop(columns="x"), pdf.drop(columns="x")) + self.assert_eq(psdf.drop(columns=["y", "z"]), pdf.drop(columns=["y", "z"])) + self.assert_eq(psdf.drop(columns=["x", "y", "z"]), pdf.drop(columns=["x", "y", "z"])) + self.assert_eq(psdf.drop(columns=[]), pdf.drop(columns=[])) + + columns = pd.MultiIndex.from_tuples([(1, "x"), (1, "y"), (2, "z")]) + pdf.columns = columns + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.drop(columns=1), pdf.drop(columns=1)) + self.assert_eq(psdf.drop(columns=(1, "x")), pdf.drop(columns=(1, "x"))) + self.assert_eq(psdf.drop(columns=[(1, "x"), 2]), pdf.drop(columns=[(1, "x"), 2])) + self.assert_eq( + psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + ) + + self.assertRaises(KeyError, lambda: psdf.drop(columns=3)) + self.assertRaises(KeyError, lambda: psdf.drop(columns=(1, "z"))) + + pdf.index = pd.MultiIndex.from_tuples([("i", 0), ("j", 1)]) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), + ) + + # non-string names + pdf = pd.DataFrame({10: [1, 2], 20: [3, 4], 30: [5, 6]}, index=np.random.rand(2)) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.drop(10, axis=1), pdf.drop(10, axis=1)) + self.assert_eq(psdf.drop([20, 30], axis=1), pdf.drop([20, 30], axis=1)) + + # + # Drop rows + # + + pdf = pd.DataFrame({"X": [1, 2, 3], "Y": [4, 5, 6], "Z": [7, 8, 9]}, index=["A", "B", "C"]) + psdf = ps.from_pandas(pdf) + + # Given labels (and axis = 0) + self.assert_eq(psdf.drop(labels="A", axis=0), pdf.drop(labels="A", axis=0)) + self.assert_eq(psdf.drop(labels="A"), pdf.drop(labels="A")) + self.assert_eq((psdf + 1).drop(labels="A"), (pdf + 1).drop(labels="A")) + self.assert_eq(psdf.drop(labels=["A", "C"], axis=0), pdf.drop(labels=["A", "C"], axis=0)) + self.assert_eq( + psdf.drop(labels=["A", "B", "C"], axis=0), pdf.drop(labels=["A", "B", "C"], axis=0) + ) + + with ps.option_context("compute.isin_limit", 2): + self.assert_eq( + psdf.drop(labels=["A", "B", "C"], axis=0), pdf.drop(labels=["A", "B", "C"], axis=0) + ) + + # Given index + self.assert_eq(psdf.drop(index="A"), pdf.drop(index="A")) + self.assert_eq(psdf.drop(index=["A", "C"]), pdf.drop(index=["A", "C"])) + self.assert_eq(psdf.drop(index=["A", "B", "C"]), pdf.drop(index=["A", "B", "C"])) + self.assert_eq(psdf.drop(index=[]), pdf.drop(index=[])) + + with ps.option_context("compute.isin_limit", 2): + self.assert_eq(psdf.drop(index=["A", "B", "C"]), pdf.drop(index=["A", "B", "C"])) + + # Non-string names + pdf.index = [10, 20, 30] + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.drop(labels=10, axis=0), pdf.drop(labels=10, axis=0)) + self.assert_eq(psdf.drop(labels=[10, 30], axis=0), pdf.drop(labels=[10, 30], axis=0)) + self.assert_eq( + psdf.drop(labels=[10, 20, 30], axis=0), pdf.drop(labels=[10, 20, 30], axis=0) + ) + + with ps.option_context("compute.isin_limit", 2): + self.assert_eq( + psdf.drop(labels=[10, 20, 30], axis=0), pdf.drop(labels=[10, 20, 30], axis=0) + ) + + # MultiIndex + pdf.index = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + self.assertRaises(NotImplementedError, lambda: psdf.drop(labels=[("a", "x")])) + + # + # Drop rows and columns + # + pdf = pd.DataFrame({"X": [1, 2, 3], "Y": [4, 5, 6], "Z": [7, 8, 9]}, index=["A", "B", "C"]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.drop(index="A", columns="X"), pdf.drop(index="A", columns="X")) + self.assert_eq( + psdf.drop(index=["A", "C"], columns=["X", "Z"]), + pdf.drop(index=["A", "C"], columns=["X", "Z"]), + ) + self.assert_eq( + psdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), + pdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), + ) + with ps.option_context("compute.isin_limit", 2): + self.assert_eq( + psdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), + pdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), + ) + self.assert_eq( + psdf.drop(index=[], columns=["X", "Z"]), + pdf.drop(index=[], columns=["X", "Z"]), + ) + self.assert_eq( + psdf.drop(index=["A", "B", "C"], columns=[]), + pdf.drop(index=["A", "B", "C"], columns=[]), + ) + self.assert_eq( + psdf.drop(index=[], columns=[]), + pdf.drop(index=[], columns=[]), + ) + self.assertRaises( + ValueError, + lambda: psdf.drop(labels="A", axis=0, columns="X"), + ) + + def test_droplevel(self): + pdf = ( + pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + .set_index([0, 1]) + .rename_axis(["a", "b"]) + ) + pdf.columns = pd.MultiIndex.from_tuples( + [("c", "e"), ("d", "f")], names=["level_1", "level_2"] + ) + psdf = ps.from_pandas(pdf) + + self.assertRaises(ValueError, lambda: psdf.droplevel(["a", "b"])) + self.assertRaises(ValueError, lambda: psdf.droplevel([1, 1, 1, 1, 1])) + self.assertRaises(IndexError, lambda: psdf.droplevel(2)) + self.assertRaises(IndexError, lambda: psdf.droplevel(-3)) + self.assertRaises(KeyError, lambda: psdf.droplevel({"a"})) + self.assertRaises(KeyError, lambda: psdf.droplevel({"a": 1})) + + self.assertRaises(ValueError, lambda: psdf.droplevel(["level_1", "level_2"], axis=1)) + self.assertRaises(IndexError, lambda: psdf.droplevel(2, axis=1)) + self.assertRaises(IndexError, lambda: psdf.droplevel(-3, axis=1)) + self.assertRaises(KeyError, lambda: psdf.droplevel({"level_1"}, axis=1)) + self.assertRaises(KeyError, lambda: psdf.droplevel({"level_1": 1}, axis=1)) + + self.assert_eq(pdf.droplevel("a"), psdf.droplevel("a")) + self.assert_eq(pdf.droplevel(["a"]), psdf.droplevel(["a"])) + self.assert_eq(pdf.droplevel(("a",)), psdf.droplevel(("a",))) + self.assert_eq(pdf.droplevel(0), psdf.droplevel(0)) + self.assert_eq(pdf.droplevel(-1), psdf.droplevel(-1)) + + self.assert_eq(pdf.droplevel("level_1", axis=1), psdf.droplevel("level_1", axis=1)) + self.assert_eq(pdf.droplevel(["level_1"], axis=1), psdf.droplevel(["level_1"], axis=1)) + self.assert_eq(pdf.droplevel(("level_1",), axis=1), psdf.droplevel(("level_1",), axis=1)) + self.assert_eq(pdf.droplevel(0, axis=1), psdf.droplevel(0, axis=1)) + self.assert_eq(pdf.droplevel(-1, axis=1), psdf.droplevel(-1, axis=1)) + + # Tupled names + pdf.columns.names = [("level", 1), ("level", 2)] + pdf.index.names = [("a", 10), ("x", 20)] + psdf = ps.from_pandas(pdf) + + self.assertRaises(KeyError, lambda: psdf.droplevel("a")) + self.assertRaises(KeyError, lambda: psdf.droplevel(("a", 10))) + + self.assert_eq(pdf.droplevel([("a", 10)]), psdf.droplevel([("a", 10)])) + self.assert_eq( + pdf.droplevel([("level", 1)], axis=1), psdf.droplevel([("level", 1)], axis=1) + ) + + # non-string names + pdf = ( + pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + .set_index([0, 1]) + .rename_axis([10.0, 20.0]) + ) + pdf.columns = pd.MultiIndex.from_tuples([("c", "e"), ("d", "f")], names=[100.0, 200.0]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.droplevel(10.0), psdf.droplevel(10.0)) + self.assert_eq(pdf.droplevel([10.0]), psdf.droplevel([10.0])) + self.assert_eq(pdf.droplevel((10.0,)), psdf.droplevel((10.0,))) + self.assert_eq(pdf.droplevel(0), psdf.droplevel(0)) + self.assert_eq(pdf.droplevel(-1), psdf.droplevel(-1)) + self.assert_eq(pdf.droplevel(100.0, axis=1), psdf.droplevel(100.0, axis=1)) + self.assert_eq(pdf.droplevel(0, axis=1), psdf.droplevel(0, axis=1)) + + def test_drop_duplicates(self): + pdf = pd.DataFrame( + {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) + ) + psdf = ps.from_pandas(pdf) + + # inplace is False + for keep in ["first", "last", False]: + with self.subTest(keep=keep): + self.assert_eq( + pdf.drop_duplicates(keep=keep).sort_index(), + psdf.drop_duplicates(keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates("a", keep=keep).sort_index(), + psdf.drop_duplicates("a", keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates(["a", "b"], keep=keep).sort_index(), + psdf.drop_duplicates(["a", "b"], keep=keep).sort_index(), + ) + self.assert_eq( + pdf.set_index("a", append=True).drop_duplicates(keep=keep).sort_index(), + psdf.set_index("a", append=True).drop_duplicates(keep=keep).sort_index(), + ) + self.assert_eq( + pdf.set_index("a", append=True).drop_duplicates("b", keep=keep).sort_index(), + psdf.set_index("a", append=True).drop_duplicates("b", keep=keep).sort_index(), + ) + + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + + # inplace is False + for keep in ["first", "last", False]: + with self.subTest("multi-index columns", keep=keep): + self.assert_eq( + pdf.drop_duplicates(keep=keep).sort_index(), + psdf.drop_duplicates(keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates(("x", "a"), keep=keep).sort_index(), + psdf.drop_duplicates(("x", "a"), keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(), + psdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates( + [("x", "a"), ("y", "b")], keep=keep, ignore_index=True + ).sort_index(), + psdf.drop_duplicates( + [("x", "a"), ("y", "b")], keep=keep, ignore_index=True + ).sort_index(), + ) + + # inplace is True + subset_list = [None, "a", ["a", "b"]] + for subset in subset_list: + pdf = pd.DataFrame( + {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) + ) + psdf = ps.from_pandas(pdf) + pser = pdf.a + psser = psdf.a + pdf.drop_duplicates(subset=subset, inplace=True) + psdf.drop_duplicates(subset=subset, inplace=True) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser.sort_index(), pser.sort_index()) + + # multi-index columns, inplace is True + subset_list = [None, ("x", "a"), [("x", "a"), ("y", "b")]] + for subset in subset_list: + pdf = pd.DataFrame( + {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) + ) + psdf = ps.from_pandas(pdf) + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + pser = pdf[("x", "a")] + psser = psdf[("x", "a")] + pdf.drop_duplicates(subset=subset, inplace=True) + pdf.drop_duplicates(subset=subset, inplace=True, ignore_index=True) + psdf.drop_duplicates(subset=subset, inplace=True) + psdf.drop_duplicates(subset=subset, inplace=True, ignore_index=True) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser.sort_index(), pser.sort_index()) + + # non-string names + pdf = pd.DataFrame( + {10: [1, 2, 2, 2, 3], 20: ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.drop_duplicates(10, keep=keep).sort_index(), + psdf.drop_duplicates(10, keep=keep).sort_index(), + ) + self.assert_eq( + pdf.drop_duplicates([10, 20], keep=keep).sort_index(), + psdf.drop_duplicates([10, 20], keep=keep).sort_index(), + ) + + def test_duplicated(self): + pdf = pd.DataFrame( + {"a": [1, 1, 2, 3], "b": [1, 1, 1, 4], "c": [1, 1, 1, 5]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) + self.assert_eq( + pdf.duplicated(keep="last").sort_index(), + psdf.duplicated(keep="last").sort_index(), + ) + self.assert_eq( + pdf.duplicated(keep=False).sort_index(), + psdf.duplicated(keep=False).sort_index(), + ) + self.assert_eq( + pdf.duplicated(subset="b").sort_index(), + psdf.duplicated(subset="b").sort_index(), + ) + self.assert_eq( + pdf.duplicated(subset=["b"]).sort_index(), + psdf.duplicated(subset=["b"]).sort_index(), + ) + with self.assertRaisesRegex(ValueError, "'keep' only supports 'first', 'last' and False"): + psdf.duplicated(keep="false") + with self.assertRaisesRegex(KeyError, "'d'"): + psdf.duplicated(subset=["d"]) + + pdf.index.name = "x" + psdf.index.name = "x" + self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) + + # multi-index + self.assert_eq( + pdf.set_index("a", append=True).duplicated().sort_index(), + psdf.set_index("a", append=True).duplicated().sort_index(), + ) + self.assert_eq( + pdf.set_index("a", append=True).duplicated(keep=False).sort_index(), + psdf.set_index("a", append=True).duplicated(keep=False).sort_index(), + ) + self.assert_eq( + pdf.set_index("a", append=True).duplicated(subset=["b"]).sort_index(), + psdf.set_index("a", append=True).duplicated(subset=["b"]).sort_index(), + ) + + # mutli-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) + self.assert_eq( + pdf.duplicated(subset=("x", "b")).sort_index(), + psdf.duplicated(subset=("x", "b")).sort_index(), + ) + self.assert_eq( + pdf.duplicated(subset=[("x", "b")]).sort_index(), + psdf.duplicated(subset=[("x", "b")]).sort_index(), + ) + + # non-string names + pdf = pd.DataFrame( + {10: [1, 1, 2, 3], 20: [1, 1, 1, 4], 30: [1, 1, 1, 5]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) + self.assert_eq( + pdf.duplicated(subset=10).sort_index(), + psdf.duplicated(subset=10).sort_index(), + ) + + def test_filter(self): + pdf = pd.DataFrame( + { + "aa": ["aa", "bd", "bc", "ab", "ce"], + "ba": [1, 2, 3, 4, 5], + "cb": [1.0, 2.0, 3.0, 4.0, 5.0], + "db": [1.0, np.nan, 3.0, np.nan, 5.0], + } + ) + pdf = pdf.set_index("aa") + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.filter(items=["ab", "aa"], axis=0).sort_index(), + pdf.filter(items=["ab", "aa"], axis=0).sort_index(), + ) + + with option_context("compute.isin_limit", 0): + self.assert_eq( + psdf.filter(items=["ab", "aa"], axis=0).sort_index(), + pdf.filter(items=["ab", "aa"], axis=0).sort_index(), + ) + + self.assert_eq( + psdf.filter(items=["ba", "db"], axis=1).sort_index(), + pdf.filter(items=["ba", "db"], axis=1).sort_index(), + ) + + self.assert_eq(psdf.filter(like="b", axis="index"), pdf.filter(like="b", axis="index")) + self.assert_eq(psdf.filter(like="c", axis="columns"), pdf.filter(like="c", axis="columns")) + + self.assert_eq( + psdf.filter(regex="b.*", axis="index"), pdf.filter(regex="b.*", axis="index") + ) + self.assert_eq( + psdf.filter(regex="b.*", axis="columns"), pdf.filter(regex="b.*", axis="columns") + ) + + pdf = pdf.set_index("ba", append=True) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(), + pdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(), + ) + + with self.assertRaisesRegex(TypeError, "Unsupported type list"): + psdf.filter(items=[["aa", 1], ("bd", 2)], axis=0) + + with self.assertRaisesRegex(ValueError, "The item should not be empty."): + psdf.filter(items=[(), ("bd", 2)], axis=0) + + self.assert_eq(psdf.filter(like="b", axis=0), pdf.filter(like="b", axis=0)) + + self.assert_eq(psdf.filter(regex="b.*", axis=0), pdf.filter(regex="b.*", axis=0)) + + with self.assertRaisesRegex(ValueError, "items should be a list-like object"): + psdf.filter(items="b") + + with self.assertRaisesRegex(ValueError, "No axis named"): + psdf.filter(regex="b.*", axis=123) + + with self.assertRaisesRegex(TypeError, "Must pass either `items`, `like`"): + psdf.filter() + + with self.assertRaisesRegex(TypeError, "mutually exclusive"): + psdf.filter(regex="b.*", like="aaa") + + # multi-index columns + pdf = pd.DataFrame( + { + ("x", "aa"): ["aa", "ab", "bc", "bd", "ce"], + ("x", "ba"): [1, 2, 3, 4, 5], + ("y", "cb"): [1.0, 2.0, 3.0, 4.0, 5.0], + ("z", "db"): [1.0, np.nan, 3.0, np.nan, 5.0], + } + ) + pdf = pdf.set_index(("x", "aa")) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.filter(items=["ab", "aa"], axis=0).sort_index(), + pdf.filter(items=["ab", "aa"], axis=0).sort_index(), + ) + self.assert_eq( + psdf.filter(items=[("x", "ba"), ("z", "db")], axis=1).sort_index(), + pdf.filter(items=[("x", "ba"), ("z", "db")], axis=1).sort_index(), + ) + + self.assert_eq(psdf.filter(like="b", axis="index"), pdf.filter(like="b", axis="index")) + self.assert_eq(psdf.filter(like="c", axis="columns"), pdf.filter(like="c", axis="columns")) + + self.assert_eq( + psdf.filter(regex="b.*", axis="index"), pdf.filter(regex="b.*", axis="index") + ) + self.assert_eq( + psdf.filter(regex="b.*", axis="columns"), pdf.filter(regex="b.*", axis="columns") + ) + + def test_last(self): + index = pd.date_range("2018-04-09", periods=4, freq="2D") + pdf = pd.DataFrame([1, 2, 3, 4], index=index) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.last("1D"), psdf.last("1D")) + self.assert_eq(pdf.last(DateOffset(days=1)), psdf.last(DateOffset(days=1))) + with self.assertRaisesRegex(TypeError, "'last' only supports a DatetimeIndex"): + ps.DataFrame([1, 2, 3, 4]).last("1D") + + def test_first(self): + index = pd.date_range("2018-04-09", periods=4, freq="2D") + pdf = pd.DataFrame([1, 2, 3, 4], index=index) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.first("1D"), psdf.first("1D")) + self.assert_eq(pdf.first(DateOffset(days=1)), psdf.first(DateOffset(days=1))) + with self.assertRaisesRegex(TypeError, "'first' only supports a DatetimeIndex"): + ps.DataFrame([1, 2, 3, 4]).first("1D") + + def test_swaplevel(self): + # MultiIndex with two levels + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame({"x1": ["a", "b", "c", "d"], "x2": ["a", "b", "c", "d"]}, index=pidx) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.swaplevel(), psdf.swaplevel()) + self.assert_eq(pdf.swaplevel(0, 1), psdf.swaplevel(0, 1)) + self.assert_eq(pdf.swaplevel(1, 1), psdf.swaplevel(1, 1)) + self.assert_eq(pdf.swaplevel("number", "color"), psdf.swaplevel("number", "color")) + + # MultiIndex with more than two levels + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"], ["l", "m", "s", "xs"]] + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color", "size")) + pdf = pd.DataFrame({"x1": ["a", "b", "c", "d"], "x2": ["a", "b", "c", "d"]}, index=pidx) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.swaplevel(), psdf.swaplevel()) + self.assert_eq(pdf.swaplevel(0, 1), psdf.swaplevel(0, 1)) + self.assert_eq(pdf.swaplevel(0, 2), psdf.swaplevel(0, 2)) + self.assert_eq(pdf.swaplevel(1, 2), psdf.swaplevel(1, 2)) + self.assert_eq(pdf.swaplevel(1, 1), psdf.swaplevel(1, 1)) + self.assert_eq(pdf.swaplevel(-1, -2), psdf.swaplevel(-1, -2)) + self.assert_eq(pdf.swaplevel("number", "color"), psdf.swaplevel("number", "color")) + self.assert_eq(pdf.swaplevel("number", "size"), psdf.swaplevel("number", "size")) + self.assert_eq(pdf.swaplevel("color", "size"), psdf.swaplevel("color", "size")) + self.assert_eq( + pdf.swaplevel("color", "size", axis="index"), + psdf.swaplevel("color", "size", axis="index"), + ) + self.assert_eq( + pdf.swaplevel("color", "size", axis=0), psdf.swaplevel("color", "size", axis=0) + ) + + pdf = pd.DataFrame( + { + "x1": ["a", "b", "c", "d"], + "x2": ["a", "b", "c", "d"], + "x3": ["a", "b", "c", "d"], + "x4": ["a", "b", "c", "d"], + } + ) + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color", "size")) + pdf.columns = pidx + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.swaplevel(axis=1), psdf.swaplevel(axis=1)) + self.assert_eq(pdf.swaplevel(0, 1, axis=1), psdf.swaplevel(0, 1, axis=1)) + self.assert_eq(pdf.swaplevel(0, 2, axis=1), psdf.swaplevel(0, 2, axis=1)) + self.assert_eq(pdf.swaplevel(1, 2, axis=1), psdf.swaplevel(1, 2, axis=1)) + self.assert_eq(pdf.swaplevel(1, 1, axis=1), psdf.swaplevel(1, 1, axis=1)) + self.assert_eq(pdf.swaplevel(-1, -2, axis=1), psdf.swaplevel(-1, -2, axis=1)) + self.assert_eq( + pdf.swaplevel("number", "color", axis=1), psdf.swaplevel("number", "color", axis=1) + ) + self.assert_eq( + pdf.swaplevel("number", "size", axis=1), psdf.swaplevel("number", "size", axis=1) + ) + self.assert_eq( + pdf.swaplevel("color", "size", axis=1), psdf.swaplevel("color", "size", axis=1) + ) + self.assert_eq( + pdf.swaplevel("color", "size", axis="columns"), + psdf.swaplevel("color", "size", axis="columns"), + ) + + # Error conditions + self.assertRaises(AssertionError, lambda: ps.DataFrame([1, 2]).swaplevel()) + self.assertRaises(IndexError, lambda: psdf.swaplevel(0, 9, axis=1)) + self.assertRaises(KeyError, lambda: psdf.swaplevel("not_number", "color", axis=1)) + self.assertRaises(ValueError, lambda: psdf.swaplevel(axis=2)) + + def test_swapaxes(self): + pdf = pd.DataFrame( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=["x", "y", "z"], columns=["a", "b", "c"] + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.swapaxes(0, 1), pdf.swapaxes(0, 1)) + self.assert_eq(psdf.swapaxes(1, 0), pdf.swapaxes(1, 0)) + self.assert_eq(psdf.swapaxes("index", "columns"), pdf.swapaxes("index", "columns")) + self.assert_eq(psdf.swapaxes("columns", "index"), pdf.swapaxes("columns", "index")) + self.assert_eq((psdf + 1).swapaxes(0, 1), (pdf + 1).swapaxes(0, 1)) + + self.assertRaises(AssertionError, lambda: psdf.swapaxes(0, 1, copy=False)) + self.assertRaises(ValueError, lambda: psdf.swapaxes(0, -1)) + + def test_isin(self): + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, 2, 9, 4, 2, 4], + "c": ["one", "three", "six", "seven", "one", "5"], + }, + index=np.random.rand(6), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.isin([4, "six"]), pdf.isin([4, "six"])) + # Seems like pandas has a bug when passing `np.array` as parameter + self.assert_eq(psdf.isin(np.array([4, "six"])), pdf.isin([4, "six"])) + self.assert_eq( + psdf.isin({"a": [2, 8], "c": ["three", "one"]}), + pdf.isin({"a": [2, 8], "c": ["three", "one"]}), + ) + self.assert_eq( + psdf.isin({"a": np.array([2, 8]), "c": ["three", "one"]}), + pdf.isin({"a": np.array([2, 8]), "c": ["three", "one"]}), + ) + + msg = "'DataFrame' object has no attribute {'e'}" + with self.assertRaisesRegex(AttributeError, msg): + psdf.isin({"e": [5, 7], "a": [1, 6]}) + + msg = "DataFrame and Series are not supported" + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.isin(pdf) + + msg = "Values should be iterable, Series, DataFrame or dict." + with self.assertRaisesRegex(TypeError, msg): + psdf.isin(1) + + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, None, 9, 4, None, 4], + "c": [None, 5, None, 3, 2, 1], + }, + ) + psdf = ps.from_pandas(pdf) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq(psdf.isin([4, 3, 1, 1, None]), pdf.isin([4, 3, 1, 1, None])) + else: + expected = pd.DataFrame( + { + "a": [True, False, True, True, False, False], + "b": [True, False, False, True, False, True], + "c": [False, False, False, True, False, True], + } + ) + self.assert_eq(psdf.isin([4, 3, 1, 1, None]), expected) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq( + psdf.isin({"b": [4, 3, 1, 1, None]}), pdf.isin({"b": [4, 3, 1, 1, None]}) + ) + else: + expected = pd.DataFrame( + { + "a": [False, False, False, False, False, False], + "b": [True, False, False, True, False, True], + "c": [False, False, False, False, False, False], + } + ) + self.assert_eq(psdf.isin({"b": [4, 3, 1, 1, None]}), expected) + + def test_sample(self): + psdf = ps.DataFrame({"A": [0, 2, 4]}, index=["x", "y", "z"]) + + # Make sure the tests run, but we can't check the result because they are non-deterministic. + psdf.sample(frac=0.1) + psdf.sample(frac=0.2, replace=True) + psdf.sample(frac=0.2, random_state=5) + psdf["A"].sample(frac=0.2) + psdf["A"].sample(frac=0.2, replace=True) + psdf["A"].sample(frac=0.2, random_state=5) + + self.assert_eq(psdf.sample(frac=0.1, ignore_index=True).index.dtype, np.int64) + self.assert_eq(psdf.sample(frac=0.2, replace=True, ignore_index=True).index.dtype, np.int64) + self.assert_eq( + psdf.sample(frac=0.2, random_state=5, ignore_index=True).index.dtype, np.int64 + ) + self.assert_eq(psdf["A"].sample(frac=0.2, ignore_index=True).index.dtype, np.int64) + self.assert_eq( + psdf["A"].sample(frac=0.2, replace=True, ignore_index=True).index.dtype, np.int64 + ) + self.assert_eq( + psdf["A"].sample(frac=0.2, random_state=5, ignore_index=True).index.dtype, np.int64 + ) + + with self.assertRaises(ValueError): + psdf.sample() + with self.assertRaises(NotImplementedError): + psdf.sample(n=1) + + +class FrameReidexingTests(FrameReindexingMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_reindexing import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reshaping.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reshaping.py new file mode 100644 index 0000000..5ee7d2c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_reshaping.py @@ -0,0 +1,487 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context + +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Reshaping, sorting, transposing' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#reshaping-sorting-transposing +class FrameReshapingMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_sort_values(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.sort_values("b"), pdf.sort_values("b")) + self.assert_eq( + psdf.sort_values("b", ignore_index=True), pdf.sort_values("b", ignore_index=True) + ) + + for ascending in [True, False]: + for na_position in ["first", "last"]: + self.assert_eq( + psdf.sort_values("a", ascending=ascending, na_position=na_position), + pdf.sort_values("a", ascending=ascending, na_position=na_position), + ) + + self.assert_eq(psdf.sort_values(["a", "b"]), pdf.sort_values(["a", "b"])) + self.assert_eq( + psdf.sort_values(["a", "b"], ignore_index=True), + pdf.sort_values(["a", "b"], ignore_index=True), + ) + self.assert_eq( + psdf.sort_values(["a", "b"], ascending=[False, True]), + pdf.sort_values(["a", "b"], ascending=[False, True]), + ) + + self.assertRaises(ValueError, lambda: psdf.sort_values(["b", "a"], ascending=[False])) + + self.assert_eq( + psdf.sort_values(["a", "b"], na_position="first"), + pdf.sort_values(["a", "b"], na_position="first"), + ) + + self.assertRaises(ValueError, lambda: psdf.sort_values(["b", "a"], na_position="invalid")) + + pserA = pdf.a + psserA = psdf.a + self.assert_eq(psdf.sort_values("b", inplace=True), pdf.sort_values("b", inplace=True)) + self.assert_eq(psdf, pdf) + self.assert_eq(psserA, pserA) + + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7) + ) + psdf = ps.from_pandas(pdf) + pserA = pdf.a + psserA = psdf.a + self.assert_eq( + psdf.sort_values("b", inplace=True, ignore_index=True), + pdf.sort_values("b", inplace=True, ignore_index=True), + ) + self.assert_eq(psdf, pdf) + self.assert_eq(psserA, pserA) + + # multi-index indexes + + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, + index=pd.MultiIndex.from_tuples( + [ + ("bar", "one"), + ("bar", "two"), + ("baz", "one"), + ("baz", "two"), + ("foo", "one"), + ("foo", "two"), + ("qux", "one"), + ] + ), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.sort_values("b"), pdf.sort_values("b")) + self.assert_eq( + psdf.sort_values("b", ignore_index=True), pdf.sort_values("b", ignore_index=True) + ) + + # multi-index columns + pdf = pd.DataFrame( + {("X", 10): [1, 2, 3, 4, 5, None, 7], ("X", 20): [7, 6, 5, 4, 3, 2, 1]}, + index=np.random.rand(7), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.sort_values(("X", 20)), pdf.sort_values(("X", 20))) + self.assert_eq( + psdf.sort_values([("X", 20), ("X", 10)]), pdf.sort_values([("X", 20), ("X", 10)]) + ) + + self.assertRaisesRegex( + ValueError, + "For a multi-index, the label must be a tuple with elements", + lambda: psdf.sort_values(["X"]), + ) + + # non-string names + pdf = pd.DataFrame( + {10: [1, 2, 3, 4, 5, None, 7], 20: [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.sort_values(20), pdf.sort_values(20)) + self.assert_eq(psdf.sort_values([20, 10]), pdf.sort_values([20, 10])) + + def test_sort_index(self): + pdf = pd.DataFrame( + {"A": [2, 1, np.nan], "B": [np.nan, 0, np.nan]}, index=["b", "a", np.nan] + ) + psdf = ps.from_pandas(pdf) + + # Assert invalid parameters + self.assertRaises(NotImplementedError, lambda: psdf.sort_index(axis=1)) + self.assertRaises(NotImplementedError, lambda: psdf.sort_index(kind="mergesort")) + self.assertRaises(ValueError, lambda: psdf.sort_index(na_position="invalid")) + + # Assert default behavior without parameters + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + # Assert ignoring index + self.assert_eq(psdf.sort_index(ignore_index=True), pdf.sort_index(ignore_index=True)) + # Assert sorting descending + self.assert_eq(psdf.sort_index(ascending=False), pdf.sort_index(ascending=False)) + # Assert sorting NA indices first + self.assert_eq(psdf.sort_index(na_position="first"), pdf.sort_index(na_position="first")) + # Assert sorting descending and NA indices first + self.assert_eq( + psdf.sort_index(ascending=False, na_position="first"), + pdf.sort_index(ascending=False, na_position="first"), + ) + + # Assert sorting inplace + pserA = pdf.A + psserA = psdf.A + self.assertEqual(psdf.sort_index(inplace=True), pdf.sort_index(inplace=True)) + self.assert_eq(psdf, pdf) + self.assert_eq(psserA, pserA) + pserA = pdf.A + psserA = psdf.A + self.assertEqual( + psdf.sort_index(inplace=True, ascending=False, ignore_index=True), + pdf.sort_index(inplace=True, ascending=False, ignore_index=True), + ) + self.assert_eq(psdf, pdf) + self.assert_eq(psserA, pserA) + + # Assert multi-indices + pdf = pd.DataFrame( + {"A": range(4), "B": range(4)[::-1]}, index=[["b", "b", "a", "a"], [1, 0, 1, 0]] + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psdf.sort_index(level=[1, 0]), pdf.sort_index(level=[1, 0])) + self.assert_eq(psdf.reset_index().sort_index(), pdf.reset_index().sort_index()) + # Assert ignoring index + self.assert_eq(psdf.sort_index(ignore_index=True), pdf.sort_index(ignore_index=True)) + + # Assert with multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_nlargest(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1], "c": [1, 1, 2, 2, 3, 3, 3]}, + index=np.random.rand(7), + ) + psdf = ps.from_pandas(pdf) + # see also: https://github.com/pandas-dev/pandas/issues/46589 + if not (LooseVersion("1.4.0") <= LooseVersion(pd.__version__) <= LooseVersion("1.4.2")): + self.assert_eq(psdf.nlargest(5, columns="a"), pdf.nlargest(5, columns="a")) + self.assert_eq( + psdf.nlargest(5, columns=["a", "b"]), pdf.nlargest(5, columns=["a", "b"]) + ) + self.assert_eq(psdf.nlargest(5, columns=["c"]), pdf.nlargest(5, columns=["c"])) + self.assert_eq( + psdf.nlargest(5, columns=["c"], keep="first"), + pdf.nlargest(5, columns=["c"], keep="first"), + ) + self.assert_eq( + psdf.nlargest(5, columns=["c"], keep="last"), + pdf.nlargest(5, columns=["c"], keep="last"), + ) + msg = "`keep`=all is not implemented yet." + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.nlargest(5, columns=["c"], keep="all") + msg = 'keep must be either "first", "last" or "all".' + with self.assertRaisesRegex(ValueError, msg): + psdf.nlargest(5, columns=["c"], keep="xx") + + def test_nsmallest(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1], "c": [1, 1, 2, 2, 3, 3, 3]}, + index=np.random.rand(7), + ) + psdf = ps.from_pandas(pdf) + # see also: https://github.com/pandas-dev/pandas/issues/46589 + if not (LooseVersion("1.4.0") <= LooseVersion(pd.__version__) <= LooseVersion("1.4.2")): + self.assert_eq(psdf.nsmallest(n=5, columns="a"), pdf.nsmallest(5, columns="a")) + self.assert_eq( + psdf.nsmallest(n=5, columns=["a", "b"]), pdf.nsmallest(5, columns=["a", "b"]) + ) + self.assert_eq(psdf.nsmallest(n=5, columns=["c"]), pdf.nsmallest(5, columns=["c"])) + self.assert_eq( + psdf.nsmallest(n=5, columns=["c"], keep="first"), + pdf.nsmallest(5, columns=["c"], keep="first"), + ) + self.assert_eq( + psdf.nsmallest(n=5, columns=["c"], keep="last"), + pdf.nsmallest(5, columns=["c"], keep="last"), + ) + msg = "`keep`=all is not implemented yet." + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.nlargest(5, columns=["c"], keep="all") + msg = 'keep must be either "first", "last" or "all".' + with self.assertRaisesRegex(ValueError, msg): + psdf.nlargest(5, columns=["c"], keep="xx") + + def test_stack(self): + pdf_single_level_cols = pd.DataFrame( + [[0, 1], [2, 3]], index=["cat", "dog"], columns=["weight", "height"] + ) + psdf_single_level_cols = ps.from_pandas(pdf_single_level_cols) + + self.assert_eq( + psdf_single_level_cols.stack().sort_index(), pdf_single_level_cols.stack().sort_index() + ) + + multicol1 = pd.MultiIndex.from_tuples( + [("weight", "kg"), ("weight", "pounds")], names=["x", "y"] + ) + pdf_multi_level_cols1 = pd.DataFrame( + [[1, 2], [2, 4]], index=["cat", "dog"], columns=multicol1 + ) + psdf_multi_level_cols1 = ps.from_pandas(pdf_multi_level_cols1) + + self.assert_eq( + psdf_multi_level_cols1.stack().sort_index(), pdf_multi_level_cols1.stack().sort_index() + ) + + multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")]) + pdf_multi_level_cols2 = pd.DataFrame( + [[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2 + ) + psdf_multi_level_cols2 = ps.from_pandas(pdf_multi_level_cols2) + + self.assert_eq( + psdf_multi_level_cols2.stack().sort_index(), pdf_multi_level_cols2.stack().sort_index() + ) + + pdf = pd.DataFrame( + { + ("y", "c"): [True, True], + ("x", "b"): [False, False], + ("x", "c"): [True, False], + ("y", "a"): [False, True], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.stack().sort_index(), pdf.stack().sort_index()) + self.assert_eq(psdf[[]].stack().sort_index(), pdf[[]].stack().sort_index(), almost=True) + + def test_unstack(self): + pdf = pd.DataFrame( + np.random.randn(3, 3), + index=pd.MultiIndex.from_tuples([("rg1", "x"), ("rg1", "y"), ("rg2", "z")]), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.unstack().sort_index(), pdf.unstack().sort_index(), almost=True) + self.assert_eq( + psdf.unstack().unstack().sort_index(), pdf.unstack().unstack().sort_index(), almost=True + ) + + def test_explode(self): + pdf = pd.DataFrame( + {"A": [[-1.0, np.nan], [0.0, np.inf], [1.0, -np.inf]], "B": 1}, index=["a", "b", "c"] + ) + pdf.index.name = "index" + pdf.columns.name = "columns" + psdf = ps.from_pandas(pdf) + + expected_result1, result1 = pdf.explode("A"), psdf.explode("A") + expected_result2, result2 = pdf.explode("B"), psdf.explode("B") + expected_result3, result3 = pdf.explode("A", ignore_index=True), psdf.explode( + "A", ignore_index=True + ) + + self.assert_eq(result1, expected_result1, almost=True) + self.assert_eq(result2, expected_result2) + self.assert_eq(result1.index.name, expected_result1.index.name) + self.assert_eq(result1.columns.name, expected_result1.columns.name) + self.assert_eq(result3, expected_result3, almost=True) + self.assert_eq(result3.index, expected_result3.index) + + self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"])) + + # MultiIndex + midx = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("y", "c")], names=["index1", "index2"] + ) + pdf.index = midx + psdf = ps.from_pandas(pdf) + + expected_result1, result1 = pdf.explode("A"), psdf.explode("A") + expected_result2, result2 = pdf.explode("B"), psdf.explode("B") + expected_result3, result3 = pdf.explode("A", ignore_index=True), psdf.explode( + "A", ignore_index=True + ) + + self.assert_eq(result1, expected_result1, almost=True) + self.assert_eq(result2, expected_result2) + self.assert_eq(result1.index.names, expected_result1.index.names) + self.assert_eq(result1.columns.name, expected_result1.columns.name) + self.assert_eq(result3, expected_result3, almost=True) + self.assert_eq(result3.index, expected_result3.index) + + self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"])) + + # MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X")], names=["column1", "column2"]) + pdf.columns = columns + psdf.columns = columns + + expected_result1, result1 = pdf.explode(("A", "Z")), psdf.explode(("A", "Z")) + expected_result2, result2 = pdf.explode(("B", "X")), psdf.explode(("B", "X")) + expected_result3, result3 = pdf.A.explode("Z"), psdf.A.explode("Z") + + self.assert_eq(result1, expected_result1, almost=True) + self.assert_eq(result2, expected_result2) + self.assert_eq(result1.index.names, expected_result1.index.names) + self.assert_eq(result1.columns.names, expected_result1.columns.names) + self.assert_eq(result3, expected_result3, almost=True) + + self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"])) + self.assertRaises(ValueError, lambda: psdf.explode("A")) + + def test_transpose(self): + # TODO: what if with random index? + pdf1 = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}, columns=["col1", "col2"]) + psdf1 = ps.from_pandas(pdf1) + + pdf2 = pd.DataFrame( + data={"score": [9, 8], "kids": [0, 0], "age": [12, 22]}, + columns=["score", "kids", "age"], + ) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.transpose().sort_index(), psdf1.transpose().sort_index()) + self.assert_eq(pdf2.transpose().sort_index(), psdf2.transpose().sort_index()) + + with option_context("compute.max_rows", None): + self.assert_eq(pdf1.transpose().sort_index(), psdf1.transpose().sort_index()) + + self.assert_eq(pdf2.transpose().sort_index(), psdf2.transpose().sort_index()) + + pdf3 = pd.DataFrame( + { + ("cg1", "a"): [1, 2, 3], + ("cg1", "b"): [4, 5, 6], + ("cg2", "c"): [7, 8, 9], + ("cg3", "d"): [9, 9, 9], + }, + index=pd.MultiIndex.from_tuples([("rg1", "x"), ("rg1", "y"), ("rg2", "z")]), + ) + psdf3 = ps.from_pandas(pdf3) + + self.assert_eq(pdf3.transpose().sort_index(), psdf3.transpose().sort_index()) + + with option_context("compute.max_rows", None): + self.assert_eq(pdf3.transpose().sort_index(), psdf3.transpose().sort_index()) + + def test_assign_list(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + pdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90] + psdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90] + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser, pser) + + with self.assertRaisesRegex(ValueError, "Length of values does not match length of index"): + psdf["z"] = [10, 20, 30, 40, 50, 60, 70, 80] + + def test_squeeze(self): + axises = [None, 0, 1, "rows", "index", "columns"] + + # Multiple columns + pdf = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"], index=["x", "y"]) + psdf = ps.from_pandas(pdf) + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + # Multiple columns with MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X")]) + pdf.columns = columns + psdf.columns = columns + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + + # Single column with single value + pdf = pd.DataFrame([[1]], columns=["a"], index=["x"]) + psdf = ps.from_pandas(pdf) + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + # Single column with single value with MultiIndex column + columns = pd.MultiIndex.from_tuples([("A", "Z")]) + pdf.columns = columns + psdf.columns = columns + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + + # Single column with multiple values + pdf = pd.DataFrame([1, 2, 3, 4], columns=["a"]) + psdf = ps.from_pandas(pdf) + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + # Single column with multiple values with MultiIndex column + pdf.columns = columns + psdf.columns = columns + for axis in axises: + self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis)) + + +class FrameReshapingTests(FrameReshapingMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_reshaping import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_spark.py new file mode 100644 index 0000000..da16e94 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_spark.py @@ -0,0 +1,309 @@ +# +# 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 inspect +import sys +import unittest +from io import StringIO + +import numpy as np +import pandas as pd +from pyspark import StorageLevel +from pyspark.ml.linalg import SparseVector +from pyspark.sql.types import StructType + +from pyspark import pandas as ps +from pyspark.pandas.frame import CachedDataFrame +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame +from pyspark.testing.pandasutils import ComparisonTestBase, SPARK_CONF_ARROW_ENABLED +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Spark-related' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#spark-related +class FrameSparkMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_empty_dataframe(self): + pdf = pd.DataFrame({"a": pd.Series([], dtype="i1"), "b": pd.Series([], dtype="str")}) + + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + def test_all_null_dataframe(self): + pdf = pd.DataFrame( + { + "a": [None, None, None, "a"], + "b": [None, None, None, 1], + "c": [None, None, None] + list(np.arange(1, 2).astype("i1")), + "d": [None, None, None, 1.0], + "e": [None, None, None, True], + "f": [None, None, None] + list(pd.date_range("20130101", periods=1)), + }, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.iloc[:-1], pdf.iloc[:-1]) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + self.assert_eq(psdf.iloc[:-1], pdf.iloc[:-1]) + + pdf = pd.DataFrame( + { + "a": pd.Series([None, None, None], dtype="float64"), + "b": pd.Series([None, None, None], dtype="str"), + }, + ) + + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + def test_nullable_object(self): + pdf = pd.DataFrame( + { + "a": list("abc") + [np.nan, None], + "b": list(range(1, 4)) + [np.nan, None], + "c": list(np.arange(3, 6).astype("i1")) + [np.nan, None], + "d": list(np.arange(4.0, 7.0, dtype="float64")) + [np.nan, None], + "e": [True, False, True, np.nan, None], + "f": list(pd.date_range("20130101", periods=3)) + [np.nan, None], + }, + index=np.random.rand(5), + ) + + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + def test_dot_in_column_name(self): + self.assert_eq( + ps.DataFrame(ps.range(1)._internal.spark_frame.selectExpr("1L as `a.b`"))["a.b"], + ps.Series([1], name="a.b"), + ) + + def test_spark_schema(self): + psdf = ps.DataFrame( + { + "a": list("abc"), + "b": list(range(1, 4)), + "c": np.arange(3, 6).astype("i1"), + "d": np.arange(4.0, 7.0, dtype="float64"), + "e": [True, False, True], + "f": pd.date_range("20130101", periods=3), + }, + columns=["a", "b", "c", "d", "e", "f"], + ) + + actual = psdf.spark.schema() + expected = ( + StructType() + .add("a", "string", False) + .add("b", "long", False) + .add("c", "byte", False) + .add("d", "double", False) + .add("e", "boolean", False) + .add("f", "timestamp", False) + ) + self.assertEqual(actual, expected) + + actual = psdf.spark.schema("index") + expected = ( + StructType() + .add("index", "long", False) + .add("a", "string", False) + .add("b", "long", False) + .add("c", "byte", False) + .add("d", "double", False) + .add("e", "boolean", False) + .add("f", "timestamp", False) + ) + self.assertEqual(actual, expected) + + def test_print_schema(self): + psdf = ps.DataFrame( + {"a": list("abc"), "b": list(range(1, 4)), "c": np.arange(3, 6).astype("i1")}, + columns=["a", "b", "c"], + ) + + prev = sys.stdout + try: + out = StringIO() + sys.stdout = out + psdf.spark.print_schema() + actual = out.getvalue().strip() + + self.assertTrue("a: string" in actual, actual) + self.assertTrue("b: long" in actual, actual) + self.assertTrue("c: byte" in actual, actual) + + out = StringIO() + sys.stdout = out + psdf.spark.print_schema(index_col="index") + actual = out.getvalue().strip() + + self.assertTrue("index: long" in actual, actual) + self.assertTrue("a: string" in actual, actual) + self.assertTrue("b: long" in actual, actual) + self.assertTrue("c: byte" in actual, actual) + finally: + sys.stdout = prev + + def test_explain_hint(self): + psdf1 = ps.DataFrame( + {"lkey": ["foo", "bar", "baz", "foo"], "value": [1, 2, 3, 5]}, + columns=["lkey", "value"], + ) + psdf2 = ps.DataFrame( + {"rkey": ["foo", "bar", "baz", "foo"], "value": [5, 6, 7, 8]}, + columns=["rkey", "value"], + ) + merged = psdf1.merge(psdf2.spark.hint("broadcast"), left_on="lkey", right_on="rkey") + prev = sys.stdout + try: + out = StringIO() + sys.stdout = out + merged.spark.explain() + actual = out.getvalue().strip() + + self.assertTrue("Broadcast" in actual, actual) + finally: + sys.stdout = prev + + def test_cache(self): + pdf = pd.DataFrame( + [(0.2, 0.3), (0.0, 0.6), (0.6, 0.0), (0.2, 0.1)], columns=["dogs", "cats"] + ) + psdf = ps.from_pandas(pdf) + + with psdf.spark.cache() as cached_df: + self.assert_eq(isinstance(cached_df, CachedDataFrame), True) + self.assert_eq( + repr(cached_df.spark.storage_level), repr(StorageLevel(True, True, False, True)) + ) + + def test_persist(self): + pdf = pd.DataFrame( + [(0.2, 0.3), (0.0, 0.6), (0.6, 0.0), (0.2, 0.1)], columns=["dogs", "cats"] + ) + psdf = ps.from_pandas(pdf) + storage_levels = [ + StorageLevel.DISK_ONLY, + StorageLevel.MEMORY_AND_DISK, + StorageLevel.MEMORY_ONLY, + StorageLevel.OFF_HEAP, + ] + + for storage_level in storage_levels: + with psdf.spark.persist(storage_level) as cached_df: + self.assert_eq(isinstance(cached_df, CachedDataFrame), True) + self.assert_eq(repr(cached_df.spark.storage_level), repr(storage_level)) + + self.assertRaises(TypeError, lambda: psdf.spark.persist("DISK_ONLY")) + + def test_udt(self): + sparse_values = {0: 0.1, 1: 1.1} + sparse_vector = SparseVector(len(sparse_values), sparse_values) + pdf = pd.DataFrame({"a": [sparse_vector], "b": [10]}) + + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf, pdf) + + def test_missing(self): + psdf = self.psdf + + missing_functions = inspect.getmembers(MissingPandasLikeDataFrame, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*DataFrame.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*DataFrame.*{}.*is deprecated".format(name) + ): + getattr(psdf, name)() + + missing_properties = inspect.getmembers( + MissingPandasLikeDataFrame, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*DataFrame.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf, name) + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*DataFrame.*{}.*is deprecated".format(name) + ): + getattr(psdf, name) + + +class FrameSparkTests(FrameSparkMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_spark import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_take.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_take.py new file mode 100644 index 0000000..28d20e9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_take.py @@ -0,0 +1,156 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameTakeMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_take(self): + pdf = pd.DataFrame( + {"A": range(0, 50000), "B": range(100000, 0, -2), "C": range(100000, 50000, -1)} + ) + psdf = ps.from_pandas(pdf) + + # axis=0 (default) + self.assert_eq(psdf.take([1, 2]).sort_index(), pdf.take([1, 2]).sort_index()) + self.assert_eq(psdf.take([-1, -2]).sort_index(), pdf.take([-1, -2]).sort_index()) + self.assert_eq( + psdf.take(range(100, 110)).sort_index(), pdf.take(range(100, 110)).sort_index() + ) + self.assert_eq( + psdf.take(range(-110, -100)).sort_index(), pdf.take(range(-110, -100)).sort_index() + ) + self.assert_eq( + psdf.take([10, 100, 1000, 10000]).sort_index(), + pdf.take([10, 100, 1000, 10000]).sort_index(), + ) + self.assert_eq( + psdf.take([-10, -100, -1000, -10000]).sort_index(), + pdf.take([-10, -100, -1000, -10000]).sort_index(), + ) + + # axis=1 + self.assert_eq( + psdf.take([1, 2], axis=1).sort_index(), pdf.take([1, 2], axis=1).sort_index() + ) + self.assert_eq( + psdf.take([-1, -2], axis=1).sort_index(), pdf.take([-1, -2], axis=1).sort_index() + ) + self.assert_eq( + psdf.take(range(1, 3), axis=1).sort_index(), + pdf.take(range(1, 3), axis=1).sort_index(), + ) + self.assert_eq( + psdf.take(range(-1, -3), axis=1).sort_index(), + pdf.take(range(-1, -3), axis=1).sort_index(), + ) + self.assert_eq( + psdf.take([2, 1], axis=1).sort_index(), + pdf.take([2, 1], axis=1).sort_index(), + ) + self.assert_eq( + psdf.take([-1, -2], axis=1).sort_index(), + pdf.take([-1, -2], axis=1).sort_index(), + ) + + # MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "C")]) + psdf.columns = columns + pdf.columns = columns + + # MultiIndex columns with axis=0 (default) + self.assert_eq(psdf.take([1, 2]).sort_index(), pdf.take([1, 2]).sort_index()) + self.assert_eq(psdf.take([-1, -2]).sort_index(), pdf.take([-1, -2]).sort_index()) + self.assert_eq( + psdf.take(range(100, 110)).sort_index(), pdf.take(range(100, 110)).sort_index() + ) + self.assert_eq( + psdf.take(range(-110, -100)).sort_index(), pdf.take(range(-110, -100)).sort_index() + ) + self.assert_eq( + psdf.take([10, 100, 1000, 10000]).sort_index(), + pdf.take([10, 100, 1000, 10000]).sort_index(), + ) + self.assert_eq( + psdf.take([-10, -100, -1000, -10000]).sort_index(), + pdf.take([-10, -100, -1000, -10000]).sort_index(), + ) + + # axis=1 + self.assert_eq( + psdf.take([1, 2], axis=1).sort_index(), pdf.take([1, 2], axis=1).sort_index() + ) + self.assert_eq( + psdf.take([-1, -2], axis=1).sort_index(), pdf.take([-1, -2], axis=1).sort_index() + ) + self.assert_eq( + psdf.take(range(1, 3), axis=1).sort_index(), + pdf.take(range(1, 3), axis=1).sort_index(), + ) + self.assert_eq( + psdf.take(range(-1, -3), axis=1).sort_index(), + pdf.take(range(-1, -3), axis=1).sort_index(), + ) + self.assert_eq( + psdf.take([2, 1], axis=1).sort_index(), + pdf.take([2, 1], axis=1).sort_index(), + ) + self.assert_eq( + psdf.take([-1, -2], axis=1).sort_index(), + pdf.take([-1, -2], axis=1).sort_index(), + ) + + # Checking the type of indices. + self.assertRaises(TypeError, lambda: psdf.take(1)) + self.assertRaises(TypeError, lambda: psdf.take("1")) + self.assertRaises(TypeError, lambda: psdf.take({1, 2})) + self.assertRaises(TypeError, lambda: psdf.take({1: None, 2: None})) + + +class FrameTakeTests(FrameTakeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_take import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_time_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_time_series.py new file mode 100644 index 0000000..eed9086 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_time_series.py @@ -0,0 +1,149 @@ +# +# 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 datetime import datetime +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Time series-related' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#time-series-related +class FrameTimeSeriesMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_shift(self): + pdf = pd.DataFrame( + { + "Col1": [10, 20, 15, 30, 45], + "Col2": [13, 23, 18, 33, 48], + "Col3": [17, 27, 22, 37, 52], + }, + index=np.random.rand(5), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.shift(3), psdf.shift(3)) + self.assert_eq(pdf.shift().shift(-1), psdf.shift().shift(-1)) + self.assert_eq(pdf.shift().sum().astype(int), psdf.shift().sum()) + + # Need the expected result since pandas 0.23 does not support `fill_value` argument. + pdf1 = pd.DataFrame( + {"Col1": [0, 0, 0, 10, 20], "Col2": [0, 0, 0, 13, 23], "Col3": [0, 0, 0, 17, 27]}, + index=pdf.index, + ) + self.assert_eq(pdf1, psdf.shift(periods=3, fill_value=0)) + msg = "should be an int" + with self.assertRaisesRegex(TypeError, msg): + psdf.shift(1.5) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "Col1"), ("x", "Col2"), ("y", "Col3")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(pdf.shift(3), psdf.shift(3)) + self.assert_eq(pdf.shift().shift(-1), psdf.shift().shift(-1)) + self.assert_eq(pdf.shift(0), psdf.shift(0)) + + def test_first_valid_index(self): + pdf = pd.DataFrame( + {"a": [None, 2, 3, 2], "b": [None, 2.0, 3.0, 1.0], "c": [None, 200, 400, 200]}, + index=["Q", "W", "E", "R"], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.first_valid_index(), psdf.first_valid_index()) + self.assert_eq(pdf[[]].first_valid_index(), psdf[[]].first_valid_index()) + + # MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.first_valid_index(), psdf.first_valid_index()) + + # Empty DataFrame + pdf = pd.Series([]).to_frame() + psdf = ps.Series([]).to_frame() + self.assert_eq(pdf.first_valid_index(), psdf.first_valid_index()) + + pdf = pd.DataFrame( + {"a": [None, 2, 3, 2], "b": [None, 2.0, 3.0, 1.0], "c": [None, 200, 400, 200]}, + index=[ + datetime(2021, 1, 1), + datetime(2021, 2, 1), + datetime(2021, 3, 1), + datetime(2021, 4, 1), + ], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.first_valid_index(), psdf.first_valid_index()) + + def test_last_valid_index(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, None], "b": [1.0, 2.0, 3.0, None], "c": [100, 200, 400, None]}, + index=["Q", "W", "E", "R"], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.last_valid_index(), psdf.last_valid_index()) + self.assert_eq(pdf[[]].last_valid_index(), psdf[[]].last_valid_index()) + + # MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.last_valid_index(), psdf.last_valid_index()) + + # Empty DataFrame + pdf = pd.Series([]).to_frame() + psdf = ps.Series([]).to_frame() + self.assert_eq(pdf.last_valid_index(), psdf.last_valid_index()) + + def test_to_datetime(self): + pdf = pd.DataFrame( + {"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}, index=np.random.rand(2) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + + +class FrameTimeSeriesTests(FrameTimeSeriesMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_time_series import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_truncate.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_truncate.py new file mode 100644 index 0000000..ef64ba9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/frame/test_truncate.py @@ -0,0 +1,163 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameTruncateMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_truncate(self): + pdf1 = pd.DataFrame( + { + "A": ["a", "b", "c", "d", "e", "f", "g"], + "B": ["h", "i", "j", "k", "l", "m", "n"], + "C": ["o", "p", "q", "r", "s", "t", "u"], + }, + index=[-500, -20, -1, 0, 400, 550, 1000], + ) + psdf1 = ps.from_pandas(pdf1) + pdf2 = pd.DataFrame( + { + "A": ["a", "b", "c", "d", "e", "f", "g"], + "B": ["h", "i", "j", "k", "l", "m", "n"], + "C": ["o", "p", "q", "r", "s", "t", "u"], + }, + index=[1000, 550, 400, 0, -1, -20, -500], + ) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(psdf1.truncate(), pdf1.truncate()) + self.assert_eq(psdf1.truncate(before=-20), pdf1.truncate(before=-20)) + self.assert_eq(psdf1.truncate(after=400), pdf1.truncate(after=400)) + self.assert_eq(psdf1.truncate(copy=False), pdf1.truncate(copy=False)) + self.assert_eq(psdf1.truncate(-20, 400, copy=False), pdf1.truncate(-20, 400, copy=False)) + # The bug for these tests has been fixed in pandas 1.1.0. + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + self.assert_eq(psdf2.truncate(0, 550), pdf2.truncate(0, 550)) + self.assert_eq(psdf2.truncate(0, 550, copy=False), pdf2.truncate(0, 550, copy=False)) + else: + expected_psdf = ps.DataFrame( + {"A": ["b", "c", "d"], "B": ["i", "j", "k"], "C": ["p", "q", "r"]}, + index=[550, 400, 0], + ) + self.assert_eq(psdf2.truncate(0, 550), expected_psdf) + self.assert_eq(psdf2.truncate(0, 550, copy=False), expected_psdf) + + # axis = 1 + self.assert_eq(psdf1.truncate(axis=1), pdf1.truncate(axis=1)) + self.assert_eq(psdf1.truncate(before="B", axis=1), pdf1.truncate(before="B", axis=1)) + self.assert_eq(psdf1.truncate(after="A", axis=1), pdf1.truncate(after="A", axis=1)) + self.assert_eq(psdf1.truncate(copy=False, axis=1), pdf1.truncate(copy=False, axis=1)) + self.assert_eq(psdf2.truncate("B", "C", axis=1), pdf2.truncate("B", "C", axis=1)) + self.assert_eq( + psdf1.truncate("B", "C", copy=False, axis=1), + pdf1.truncate("B", "C", copy=False, axis=1), + ) + + # MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "Z")]) + pdf1.columns = columns + psdf1.columns = columns + pdf2.columns = columns + psdf2.columns = columns + + self.assert_eq(psdf1.truncate(), pdf1.truncate()) + self.assert_eq(psdf1.truncate(before=-20), pdf1.truncate(before=-20)) + self.assert_eq(psdf1.truncate(after=400), pdf1.truncate(after=400)) + self.assert_eq(psdf1.truncate(copy=False), pdf1.truncate(copy=False)) + self.assert_eq(psdf1.truncate(-20, 400, copy=False), pdf1.truncate(-20, 400, copy=False)) + # The bug for these tests has been fixed in pandas 1.1.0. + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + self.assert_eq(psdf2.truncate(0, 550), pdf2.truncate(0, 550)) + self.assert_eq(psdf2.truncate(0, 550, copy=False), pdf2.truncate(0, 550, copy=False)) + else: + expected_psdf.columns = columns + self.assert_eq(psdf2.truncate(0, 550), expected_psdf) + self.assert_eq(psdf2.truncate(0, 550, copy=False), expected_psdf) + # axis = 1 + self.assert_eq(psdf1.truncate(axis=1), pdf1.truncate(axis=1)) + self.assert_eq(psdf1.truncate(before="B", axis=1), pdf1.truncate(before="B", axis=1)) + self.assert_eq(psdf1.truncate(after="A", axis=1), pdf1.truncate(after="A", axis=1)) + self.assert_eq(psdf1.truncate(copy=False, axis=1), pdf1.truncate(copy=False, axis=1)) + self.assert_eq(psdf2.truncate("B", "C", axis=1), pdf2.truncate("B", "C", axis=1)) + self.assert_eq( + psdf1.truncate("B", "C", copy=False, axis=1), + pdf1.truncate("B", "C", copy=False, axis=1), + ) + + # Exceptions + psdf = ps.DataFrame( + { + "A": ["a", "b", "c", "d", "e", "f", "g"], + "B": ["h", "i", "j", "k", "l", "m", "n"], + "C": ["o", "p", "q", "r", "s", "t", "u"], + }, + index=[-500, 100, 400, 0, -1, 550, -20], + ) + msg = "truncate requires a sorted index" + with self.assertRaisesRegex(ValueError, msg): + psdf.truncate() + + psdf = ps.DataFrame( + { + "A": ["a", "b", "c", "d", "e", "f", "g"], + "B": ["h", "i", "j", "k", "l", "m", "n"], + "C": ["o", "p", "q", "r", "s", "t", "u"], + }, + index=[-500, -20, -1, 0, 400, 550, 1000], + ) + msg = "Truncate: -20 must be after 400" + with self.assertRaisesRegex(ValueError, msg): + psdf.truncate(400, -20) + msg = "Truncate: B must be after C" + with self.assertRaisesRegex(ValueError, msg): + psdf.truncate("C", "B", axis=1) + + +class FrameTruncateTests(FrameTruncateMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.frame.test_truncate import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_aggregate.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_aggregate.py new file mode 100644 index 0000000..6ceae82 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_aggregate.py @@ -0,0 +1,311 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyAggregateMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-44289): Enable GroupbyAggregateTests.test_aggregate for pandas 2.0.0.", + ) + def test_aggregate(self): + pdf = pd.DataFrame( + {"A": [1, 1, 2, 2], "B": [1, 2, 3, 4], "C": [0.362, 0.227, 1.267, -0.562]} + ) + psdf = ps.from_pandas(pdf) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(list(df.columns)).reset_index(drop=True) + + for kkey, pkey in [("A", "A"), (psdf.A, pdf.A)]: + with self.subTest(as_index=as_index, key=pkey): + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg("sum")), + sort(pdf.groupby(pkey, as_index=as_index).agg("sum")), + ) + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg({"B": "min", "C": "sum"})), + sort(pdf.groupby(pkey, as_index=as_index).agg({"B": "min", "C": "sum"})), + ) + self.assert_eq( + sort( + psdf.groupby(kkey, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + sort( + pdf.groupby(pkey, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + ) + + if as_index: + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg(["sum"])), + sort(pdf.groupby(pkey, as_index=as_index).agg(["sum"])), + ) + else: + # seems like a pandas' bug for as_index=False and func_or_funcs is list? + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg(["sum"])), + sort(pdf.groupby(pkey, as_index=True).agg(["sum"]).reset_index()), + ) + + for kkey, pkey in [(psdf.A + 1, pdf.A + 1), (psdf.copy().A, pdf.copy().A)]: + with self.subTest(as_index=as_index, key=pkey): + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg("sum")), + sort(pdf.groupby(pkey, as_index=as_index).agg("sum")), + ) + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg({"B": "min", "C": "sum"})), + sort(pdf.groupby(pkey, as_index=as_index).agg({"B": "min", "C": "sum"})), + ) + self.assert_eq( + sort( + psdf.groupby(kkey, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + sort( + pdf.groupby(pkey, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + ) + self.assert_eq( + sort(psdf.groupby(kkey, as_index=as_index).agg(["sum"])), + sort(pdf.groupby(pkey, as_index=as_index).agg(["sum"])), + ) + + expected_error_message = ( + r"aggs must be a dict mapping from column name to aggregate functions " + r"\(string or list of strings\)." + ) + with self.assertRaisesRegex(ValueError, expected_error_message): + psdf.groupby("A", as_index=as_index).agg(0) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([(10, "A"), (10, "B"), (20, "C")]) + pdf.columns = columns + psdf.columns = columns + + for as_index in [True, False]: + stats_psdf = psdf.groupby((10, "A"), as_index=as_index).agg( + {(10, "B"): "min", (20, "C"): "sum"} + ) + stats_pdf = pdf.groupby((10, "A"), as_index=as_index).agg( + {(10, "B"): "min", (20, "C"): "sum"} + ) + self.assert_eq( + stats_psdf.sort_values(by=[(10, "B"), (20, "C")]).reset_index(drop=True), + stats_pdf.sort_values(by=[(10, "B"), (20, "C")]).reset_index(drop=True), + ) + + stats_psdf = psdf.groupby((10, "A")).agg({(10, "B"): ["min", "max"], (20, "C"): "sum"}) + stats_pdf = pdf.groupby((10, "A")).agg({(10, "B"): ["min", "max"], (20, "C"): "sum"}) + self.assert_eq( + stats_psdf.sort_values( + by=[(10, "B", "min"), (10, "B", "max"), (20, "C", "sum")] + ).reset_index(drop=True), + stats_pdf.sort_values( + by=[(10, "B", "min"), (10, "B", "max"), (20, "C", "sum")] + ).reset_index(drop=True), + ) + + # non-string names + pdf.columns = [10, 20, 30] + psdf.columns = [10, 20, 30] + + for as_index in [True, False]: + stats_psdf = psdf.groupby(10, as_index=as_index).agg({20: "min", 30: "sum"}) + stats_pdf = pdf.groupby(10, as_index=as_index).agg({20: "min", 30: "sum"}) + self.assert_eq( + stats_psdf.sort_values(by=[20, 30]).reset_index(drop=True), + stats_pdf.sort_values(by=[20, 30]).reset_index(drop=True), + ) + + stats_psdf = psdf.groupby(10).agg({20: ["min", "max"], 30: "sum"}) + stats_pdf = pdf.groupby(10).agg({20: ["min", "max"], 30: "sum"}) + self.assert_eq( + stats_psdf.sort_values(by=[(20, "min"), (20, "max"), (30, "sum")]).reset_index( + drop=True + ), + stats_pdf.sort_values(by=[(20, "min"), (20, "max"), (30, "sum")]).reset_index( + drop=True + ), + ) + + def test_aggregate_func_str_list(self): + # this is test for cases where only string or list is assigned + pdf = pd.DataFrame( + { + "kind": ["cat", "dog", "cat", "dog"], + "height": [9.1, 6.0, 9.5, 34.0], + "weight": [7.9, 7.5, 9.9, 198.0], + } + ) + psdf = ps.from_pandas(pdf) + + agg_funcs = ["max", "min", ["min", "max"]] + for aggfunc in agg_funcs: + + # Since in Koalas groupby, the order of rows might be different + # so sort on index to ensure they have same output + sorted_agg_psdf = psdf.groupby("kind").agg(aggfunc).sort_index() + sorted_agg_pdf = pdf.groupby("kind").agg(aggfunc).sort_index() + self.assert_eq(sorted_agg_psdf, sorted_agg_pdf) + + # test on multi index column case + pdf = pd.DataFrame( + {"A": [1, 1, 2, 2], "B": [1, 2, 3, 4], "C": [0.362, 0.227, 1.267, -0.562]} + ) + psdf = ps.from_pandas(pdf) + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + pdf.columns = columns + psdf.columns = columns + + for aggfunc in agg_funcs: + sorted_agg_psdf = psdf.groupby(("X", "A")).agg(aggfunc).sort_index() + sorted_agg_pdf = pdf.groupby(("X", "A")).agg(aggfunc).sort_index() + self.assert_eq(sorted_agg_psdf, sorted_agg_pdf) + + def test_aggregate_relabel(self): + # this is to test named aggregation in groupby + pdf = pd.DataFrame({"group": ["a", "a", "b", "b"], "A": [0, 1, 2, 3], "B": [5, 6, 7, 8]}) + psdf = ps.from_pandas(pdf) + + # different agg column, same function + agg_pdf = pdf.groupby("group").agg(a_max=("A", "max"), b_max=("B", "max")).sort_index() + agg_psdf = psdf.groupby("group").agg(a_max=("A", "max"), b_max=("B", "max")).sort_index() + self.assert_eq(agg_pdf, agg_psdf) + + # same agg column, different functions + agg_pdf = pdf.groupby("group").agg(b_max=("B", "max"), b_min=("B", "min")).sort_index() + agg_psdf = psdf.groupby("group").agg(b_max=("B", "max"), b_min=("B", "min")).sort_index() + self.assert_eq(agg_pdf, agg_psdf) + + # test on NamedAgg + agg_pdf = ( + pdf.groupby("group").agg(b_max=pd.NamedAgg(column="B", aggfunc="max")).sort_index() + ) + agg_psdf = ( + psdf.groupby("group").agg(b_max=ps.NamedAgg(column="B", aggfunc="max")).sort_index() + ) + self.assert_eq(agg_psdf, agg_pdf) + + # test on NamedAgg multi columns aggregation + agg_pdf = ( + pdf.groupby("group") + .agg( + b_max=pd.NamedAgg(column="B", aggfunc="max"), + b_min=pd.NamedAgg(column="B", aggfunc="min"), + ) + .sort_index() + ) + agg_psdf = ( + psdf.groupby("group") + .agg( + b_max=ps.NamedAgg(column="B", aggfunc="max"), + b_min=ps.NamedAgg(column="B", aggfunc="min"), + ) + .sort_index() + ) + self.assert_eq(agg_psdf, agg_pdf) + + def test_aggregate_relabel_multiindex(self): + pdf = pd.DataFrame({"A": [0, 1, 2, 3], "B": [5, 6, 7, 8], "group": ["a", "a", "b", "b"]}) + pdf.columns = pd.MultiIndex.from_tuples([("y", "A"), ("y", "B"), ("x", "group")]) + psdf = ps.from_pandas(pdf) + + agg_pdf = pdf.groupby(("x", "group")).agg(a_max=(("y", "A"), "max")).sort_index() + agg_psdf = psdf.groupby(("x", "group")).agg(a_max=(("y", "A"), "max")).sort_index() + self.assert_eq(agg_pdf, agg_psdf) + + # same column, different methods + agg_pdf = ( + pdf.groupby(("x", "group")) + .agg(a_max=(("y", "A"), "max"), a_min=(("y", "A"), "min")) + .sort_index() + ) + agg_psdf = ( + psdf.groupby(("x", "group")) + .agg(a_max=(("y", "A"), "max"), a_min=(("y", "A"), "min")) + .sort_index() + ) + self.assert_eq(agg_pdf, agg_psdf) + + # different column, different methods + agg_pdf = ( + pdf.groupby(("x", "group")) + .agg(a_max=(("y", "B"), "max"), a_min=(("y", "A"), "min")) + .sort_index() + ) + agg_psdf = ( + psdf.groupby(("x", "group")) + .agg(a_max=(("y", "B"), "max"), a_min=(("y", "A"), "min")) + .sort_index() + ) + self.assert_eq(agg_pdf, agg_psdf) + + +class GroupbyAggregateTests(GroupbyAggregateMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_aggregate import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_apply_func.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_apply_func.py new file mode 100644 index 0000000..995fe04 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_apply_func.py @@ -0,0 +1,524 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyApplyFuncMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43708): Enable GroupByTests.test_apply " "for pandas 2.0.0.", + ) + def test_apply(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("b").apply(lambda x: x + x.min()).sort_index(), + pdf.groupby("b").apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby("b").apply(len).sort_index(), + pdf.groupby("b").apply(len).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")["a"] + .apply(lambda x, y, z: x + x.min() + y * z, 10, z=20) + .sort_index(), + pdf.groupby("b")["a"].apply(lambda x, y, z: x + x.min() + y * z, 10, z=20).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")[["a"]].apply(lambda x: x + x.min()).sort_index(), + pdf.groupby("b")[["a"]].apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"]) + .apply(lambda x, y, z: x + x.min() + y + z, 1, z=2) + .sort_index(), + pdf.groupby(["a", "b"]).apply(lambda x, y, z: x + x.min() + y + z, 1, z=2).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["c"].apply(lambda x: 1).sort_index(), + pdf.groupby(["b"])["c"].apply(lambda x: 1).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["c"].apply(len).sort_index(), + pdf.groupby(["b"])["c"].apply(len).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5).apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5)["a"].apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)[["a"]].apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5)[["a"]].apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)[["a"]].apply(len).sort_index(), + pdf.groupby(pdf.b // 5)[["a"]].apply(len).sort_index(), + almost=True, + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).apply(lambda x: x + x.min()).sort_index(), + pdf.a.rename().groupby(pdf.b).apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).apply(lambda x: x + x.min()).sort_index(), + pdf.a.groupby(pdf.b.rename()).apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).apply(lambda x: x + x.min()).sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).apply(lambda x: x + x.min()).sort_index(), + ) + + with self.assertRaisesRegex(TypeError, "int object is not callable"): + psdf.groupby("b").apply(1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).apply(lambda x: 1).sort_index(), + pdf.groupby(("x", "b")).apply(lambda x: 1).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).apply(lambda x: x + x.min()).sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(("x", "b")).apply(len).sort_index(), + pdf.groupby(("x", "b")).apply(len).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).apply(len).sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).apply(len).sort_index(), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43706): Enable GroupByTests.test_apply_without_shortcut " "for pandas 2.0.0.", + ) + def test_apply_without_shortcut(self): + with option_context("compute.shortcut_limit", 0): + self.test_apply() + + def test_apply_with_type_hint(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + psdf = ps.from_pandas(pdf) + + def add_max1(x) -> ps.DataFrame[int, int, int]: + return x + x.min() + + # Type hints set the default column names, and we use default index for + # pandas API on Spark. Here we ignore both diff. + actual = psdf.groupby("b").apply(add_max1).sort_index() + expected = pdf.groupby("b").apply(add_max1).sort_index() + self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy())) + self.assert_eq(sorted(actual["c2"].to_numpy()), sorted(expected["c"].to_numpy())) + + def add_max2( + x, + ) -> ps.DataFrame[slice("a", int), slice("b", int), slice("c", int)]: # noqa: F405 + return x + x.min() + + actual = psdf.groupby("b").apply(add_max2).sort_index() + expected = pdf.groupby("b").apply(add_max2).sort_index() + self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy())) + self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy())) + self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy())) + + def test_apply_negative(self): + def func(_) -> ps.Series[int]: + return pd.Series([1]) + + with self.assertRaisesRegex(TypeError, "Series as a return type hint at frame groupby"): + ps.range(10).groupby("id").apply(func) + + def test_apply_with_new_dataframe(self): + pdf = pd.DataFrame( + {"timestamp": [0.0, 0.5, 1.0, 0.0, 0.5], "car_id": ["A", "A", "A", "B", "B"]} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(), + pdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(), + ) + + self.assert_eq( + psdf.groupby("car_id") + .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]})) + .sort_index(), + pdf.groupby("car_id") + .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]})) + .sort_index(), + ) + + # dataframe with 1000+ records + pdf = pd.DataFrame( + { + "timestamp": [0.0, 0.5, 1.0, 0.0, 0.5] * 300, + "car_id": ["A", "A", "A", "B", "B"] * 300, + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(), + pdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(), + ) + + self.assert_eq( + psdf.groupby("car_id") + .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]})) + .sort_index(), + pdf.groupby("car_id") + .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]})) + .sort_index(), + ) + + def test_apply_infer_schema_without_shortcut(self): + # SPARK-39054: Ensure infer schema accuracy in GroupBy.apply + with option_context("compute.shortcut_limit", 0): + dfs = ( + {"timestamp": [0.0], "car_id": ["A"]}, + {"timestamp": [0.0, 0.0], "car_id": ["A", "A"]}, + ) + func = lambda _: pd.DataFrame({"column": [0.0]}) # noqa: E731 + for df in dfs: + pdf = pd.DataFrame(df) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("car_id").apply(func).sort_index(), + pdf.groupby("car_id").apply(func).sort_index(), + ) + + def test_apply_with_new_dataframe_without_shortcut(self): + with option_context("compute.shortcut_limit", 0): + self.test_apply_with_new_dataframe() + + def test_apply_key_handling(self): + pdf = pd.DataFrame( + {"d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("d").apply(sum).sort_index(), pdf.groupby("d").apply(sum).sort_index() + ) + + with ps.option_context("compute.shortcut_limit", 1): + self.assert_eq( + psdf.groupby("d").apply(sum).sort_index(), pdf.groupby("d").apply(sum).sort_index() + ) + + def test_apply_with_side_effect(self): + pdf = pd.DataFrame( + {"d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]} + ) + psdf = ps.from_pandas(pdf) + + acc = ps.utils.default_session().sparkContext.accumulator(0) + + def sum_with_acc_frame(x) -> ps.DataFrame[np.float64, np.float64]: + nonlocal acc + acc += 1 + return np.sum(x) + + actual = psdf.groupby("d").apply(sum_with_acc_frame) + actual.columns = ["d", "v"] + self.assert_eq( + actual._to_pandas().sort_index(), + pdf.groupby("d").apply(sum).sort_index().reset_index(drop=True), + ) + self.assert_eq(acc.value, 2) + + def sum_with_acc_series(x) -> np.float64: + nonlocal acc + acc += 1 + return np.sum(x) + + self.assert_eq( + psdf.groupby("d")["v"].apply(sum_with_acc_series)._to_pandas().sort_index(), + pdf.groupby("d")["v"].apply(sum).sort_index().reset_index(drop=True), + ) + self.assert_eq(acc.value, 4) + + def test_apply_return_series(self): + # SPARK-36907: Fix DataFrameGroupBy.apply without shortcut. + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").apply(lambda x: x.iloc[0]).sort_index(), + pdf.groupby("b").apply(lambda x: x.iloc[0]).sort_index(), + ) + self.assert_eq( + psdf.groupby("b").apply(lambda x: x["a"]).sort_index(), + pdf.groupby("b").apply(lambda x: x["a"]).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b", "c"]).apply(lambda x: x.iloc[0]).sort_index(), + pdf.groupby(["b", "c"]).apply(lambda x: x.iloc[0]).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b", "c"]).apply(lambda x: x["a"]).sort_index(), + pdf.groupby(["b", "c"]).apply(lambda x: x["a"]).sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).apply(lambda x: x.iloc[0]).sort_index(), + pdf.groupby(("x", "b")).apply(lambda x: x.iloc[0]).sort_index(), + ) + self.assert_eq( + psdf.groupby(("x", "b")).apply(lambda x: x[("x", "a")]).sort_index(), + pdf.groupby(("x", "b")).apply(lambda x: x[("x", "a")]).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "b"), ("y", "c")]).apply(lambda x: x.iloc[0]).sort_index(), + pdf.groupby([("x", "b"), ("y", "c")]).apply(lambda x: x.iloc[0]).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "b"), ("y", "c")]).apply(lambda x: x[("x", "a")]).sort_index(), + pdf.groupby([("x", "b"), ("y", "c")]).apply(lambda x: x[("x", "a")]).sort_index(), + ) + + def test_apply_return_series_without_shortcut(self): + # SPARK-36907: Fix DataFrameGroupBy.apply without shortcut. + with ps.option_context("compute.shortcut_limit", 2): + self.test_apply_return_series() + + def test_apply_explicitly_infer(self): + # SPARK-39317 + from pyspark.pandas.utils import SPARK_CONF_ARROW_ENABLED + + def plus_min(x): + return x + x.min() + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + df = ps.DataFrame({"A": ["a", "a", "b"], "B": [1, 2, 3]}, columns=["A", "B"]) + g = df.groupby("A") + g.apply(plus_min).sort_index() + + def test_transform(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("b").transform(lambda x: x + x.min()).sort_index(), + pdf.groupby("b").transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")["a"].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby("b")["a"].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")[["a"]].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby("b")[["a"]].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"]).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(["a", "b"]).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["c"].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(["b"])["c"].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5)["a"].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)[["a"]].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pdf.b // 5)[["a"]].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).transform(lambda x: x + x.min()).sort_index(), + pdf.a.rename().groupby(pdf.b).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).transform(lambda x: x + x.min()).sort_index(), + pdf.a.groupby(pdf.b.rename()).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).transform(lambda x: x + x.min()).sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).transform(lambda x: x + x.min()).sort_index(), + ) + with self.assertRaisesRegex(TypeError, "str object is not callable"): + psdf.groupby("a").transform("sum") + + def udf(col) -> int: + return col + 10 + + with self.assertRaisesRegex( + TypeError, + "Expected the return type of this function to be of Series type, " + "but found type ScalarType\\[LongType\\(\\)\\]", + ): + psdf.groupby("a").transform(udf) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(("x", "b")).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x + x.min()).sort_index(), + ) + + def test_transform_without_shortcut(self): + with option_context("compute.shortcut_limit", 0): + self.test_transform() + + def test_filter(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby("b").filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")["a"].filter(lambda x: any(x == 2)).sort_index(), + pdf.groupby("b")["a"].filter(lambda x: any(x == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby("b")[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby("b")[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"]).filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(["a", "b"]).filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf["b"] // 5).filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(pdf["b"] // 5).filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf["b"] // 5)["a"].filter(lambda x: any(x == 2)).sort_index(), + pdf.groupby(pdf["b"] // 5)["a"].filter(lambda x: any(x == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf["b"] // 5)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(pdf["b"] // 5)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).filter(lambda x: any(x == 2)).sort_index(), + pdf.a.rename().groupby(pdf.b).filter(lambda x: any(x == 2)).sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(), + pdf.a.groupby(pdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(), + ) + + with self.assertRaisesRegex(TypeError, "int object is not callable"): + psdf.groupby("b").filter(1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).filter(lambda x: any(x[("x", "a")] == 2)).sort_index(), + pdf.groupby(("x", "b")).filter(lambda x: any(x[("x", "a")] == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]) + .filter(lambda x: any(x[("x", "a")] == 2)) + .sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]) + .filter(lambda x: any(x[("x", "a")] == 2)) + .sort_index(), + ) + + +class GroupbyApplyFuncTests(GroupbyApplyFuncMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_apply_func import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_cumulative.py new file mode 100644 index 0000000..64d6812 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_cumulative.py @@ -0,0 +1,419 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import DataError +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyCumulativeMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_cumcount(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + for ascending in [True, False]: + self.assert_eq( + psdf.groupby("b").cumcount(ascending=ascending).sort_index(), + pdf.groupby("b").cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"]).cumcount(ascending=ascending).sort_index(), + pdf.groupby(["a", "b"]).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].cumcount(ascending=ascending).sort_index(), + pdf.groupby(["b"])["a"].cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].cumcount(ascending=ascending).sort_index(), + pdf.groupby(["b"])[["a", "c"]].cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).cumcount(ascending=ascending).sort_index(), + pdf.groupby(pdf.b // 5).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].cumcount(ascending=ascending).sort_index(), + pdf.groupby(pdf.b // 5)["a"].cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby("b").cumcount(ascending=ascending).sum(), + pdf.groupby("b").cumcount(ascending=ascending).sum(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).cumcount(ascending=ascending).sort_index(), + pdf.a.rename().groupby(pdf.b).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).cumcount(ascending=ascending).sort_index(), + pdf.a.groupby(pdf.b.rename()).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).cumcount(ascending=ascending).sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).cumcount(ascending=ascending).sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + for ascending in [True, False]: + self.assert_eq( + psdf.groupby(("x", "b")).cumcount(ascending=ascending).sort_index(), + pdf.groupby(("x", "b")).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).cumcount(ascending=ascending).sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).cumcount(ascending=ascending).sort_index(), + ) + + def test_cummin(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").cummin().sort_index(), pdf.groupby("b").cummin().sort_index() + ) + self.assert_eq( + psdf.groupby(["a", "b"]).cummin().sort_index(), + pdf.groupby(["a", "b"]).cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].cummin().sort_index(), + pdf.groupby(["b"])["a"].cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].cummin().sort_index(), + pdf.groupby(["b"])[["a", "c"]].cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).cummin().sort_index(), + pdf.groupby(pdf.b // 5).cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].cummin().sort_index(), + pdf.groupby(pdf.b // 5)["a"].cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby("b").cummin().sum().sort_index(), + pdf.groupby("b").cummin().sum().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).cummin().sort_index(), + pdf.a.rename().groupby(pdf.b).cummin().sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).cummin().sort_index(), + pdf.a.groupby(pdf.b.rename()).cummin().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).cummin().sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).cummin().sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).cummin().sort_index(), + pdf.groupby(("x", "b")).cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).cummin().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).cummin().sort_index(), + ) + + psdf = ps.DataFrame([["a"], ["b"], ["c"]], columns=["A"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"]).cummin()) + psdf = ps.DataFrame([[1, "a"], [2, "b"], [3, "c"]], columns=["A", "B"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"])["B"].cummin()) + + def test_cummax(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").cummax().sort_index(), pdf.groupby("b").cummax().sort_index() + ) + self.assert_eq( + psdf.groupby(["a", "b"]).cummax().sort_index(), + pdf.groupby(["a", "b"]).cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].cummax().sort_index(), + pdf.groupby(["b"])["a"].cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].cummax().sort_index(), + pdf.groupby(["b"])[["a", "c"]].cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).cummax().sort_index(), + pdf.groupby(pdf.b // 5).cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].cummax().sort_index(), + pdf.groupby(pdf.b // 5)["a"].cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby("b").cummax().sum().sort_index(), + pdf.groupby("b").cummax().sum().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).cummax().sort_index(), + pdf.a.rename().groupby(pdf.b).cummax().sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).cummax().sort_index(), + pdf.a.groupby(pdf.b.rename()).cummax().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).cummax().sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).cummax().sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).cummax().sort_index(), + pdf.groupby(("x", "b")).cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).cummax().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).cummax().sort_index(), + ) + + psdf = ps.DataFrame([["a"], ["b"], ["c"]], columns=["A"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"]).cummax()) + psdf = ps.DataFrame([[1, "a"], [2, "b"], [3, "c"]], columns=["A", "B"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"])["B"].cummax()) + + def test_cumsum(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").cumsum().sort_index(), pdf.groupby("b").cumsum().sort_index() + ) + self.assert_eq( + psdf.groupby(["a", "b"]).cumsum().sort_index(), + pdf.groupby(["a", "b"]).cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].cumsum().sort_index(), + pdf.groupby(["b"])["a"].cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].cumsum().sort_index(), + pdf.groupby(["b"])[["a", "c"]].cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).cumsum().sort_index(), + pdf.groupby(pdf.b // 5).cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].cumsum().sort_index(), + pdf.groupby(pdf.b // 5)["a"].cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby("b").cumsum().sum().sort_index(), + pdf.groupby("b").cumsum().sum().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).cumsum().sort_index(), + pdf.a.rename().groupby(pdf.b).cumsum().sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).cumsum().sort_index(), + pdf.a.groupby(pdf.b.rename()).cumsum().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).cumsum().sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).cumsum().sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).cumsum().sort_index(), + pdf.groupby(("x", "b")).cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).cumsum().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).cumsum().sort_index(), + ) + + psdf = ps.DataFrame([["a"], ["b"], ["c"]], columns=["A"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"]).cumsum()) + psdf = ps.DataFrame([[1, "a"], [2, "b"], [3, "c"]], columns=["A", "B"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"])["B"].cumsum()) + + def test_cumprod(self): + pdf = pd.DataFrame( + { + "a": [1, 2, -3, 4, -5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 0, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("b").cumprod().sort_index(), + pdf.groupby("b").cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby(["a", "b"]).cumprod().sort_index(), + pdf.groupby(["a", "b"]).cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby(["b"])["a"].cumprod().sort_index(), + pdf.groupby(["b"])["a"].cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].cumprod().sort_index(), + pdf.groupby(["b"])[["a", "c"]].cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby(psdf.b // 3).cumprod().sort_index(), + pdf.groupby(pdf.b // 3).cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby(psdf.b // 3)["a"].cumprod().sort_index(), + pdf.groupby(pdf.b // 3)["a"].cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby("b").cumprod().sum().sort_index(), + pdf.groupby("b").cumprod().sum().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).cumprod().sort_index(), + pdf.a.rename().groupby(pdf.b).cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).cumprod().sort_index(), + pdf.a.groupby(pdf.b.rename()).cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).cumprod().sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).cumprod().sort_index(), + check_exact=False, + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).cumprod().sort_index(), + pdf.groupby(("x", "b")).cumprod().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).cumprod().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).cumprod().sort_index(), + check_exact=False, + ) + + psdf = ps.DataFrame([["a"], ["b"], ["c"]], columns=["A"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"]).cumprod()) + psdf = ps.DataFrame([[1, "a"], [2, "b"], [3, "c"]], columns=["A", "B"]) + self.assertRaises(DataError, lambda: psdf.groupby(["A"])["B"].cumprod()) + + +class GroupbyCumulativeTests(GroupbyCumulativeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_cumulative import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_describe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_describe.py new file mode 100644 index 0000000..dcd0e33 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_describe.py @@ -0,0 +1,144 @@ +# +# 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 unittest + +from itertools import product +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyDescribeMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_describe(self): + # support for numeric type, not support for string type yet + datas = [] + datas.append({"a": [1, 1, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + datas.append({"a": [-1, -1, -3], "b": [-4, -5, -6], "c": [-7, -8, -9]}) + datas.append({"a": [0, 0, 0], "b": [0, 0, 0], "c": [0, 8, 0]}) + # it is okay if string type column as a group key + datas.append({"a": ["a", "a", "c"], "b": [4, 5, 6], "c": [7, 8, 9]}) + + percentiles = [0.25, 0.5, 0.75] + formatted_percentiles = ["25%", "50%", "75%"] + non_percentile_stats = ["count", "mean", "std", "min", "max"] + + for data in datas: + pdf = pd.DataFrame(data) + psdf = ps.from_pandas(pdf) + + describe_pdf = pdf.groupby("a").describe().sort_index() + describe_psdf = psdf.groupby("a").describe().sort_index() + + # since the result of percentile columns are slightly difference from pandas, + # we should check them separately: non-percentile columns & percentile columns + + # 1. Check that non-percentile columns are equal. + agg_cols = [col.name for col in psdf.groupby("a")._agg_columns] + self.assert_eq( + describe_psdf.drop(columns=list(product(agg_cols, formatted_percentiles))), + describe_pdf.drop(columns=formatted_percentiles, level=1), + check_exact=False, + ) + + # 2. Check that percentile columns are equal. + # The interpolation argument is yet to be implemented in Koalas. + quantile_pdf = pdf.groupby("a").quantile(percentiles, interpolation="nearest") + quantile_pdf = quantile_pdf.unstack(level=1).astype(float) + self.assert_eq( + describe_psdf.drop(columns=list(product(agg_cols, non_percentile_stats))), + quantile_pdf.rename(columns="{:.0%}".format, level=1), + ) + + # not support for string type yet + datas = [] + datas.append({"a": ["a", "a", "c"], "b": ["d", "e", "f"], "c": ["g", "h", "i"]}) + datas.append({"a": ["a", "a", "c"], "b": [4, 0, 1], "c": ["g", "h", "i"]}) + for data in datas: + pdf = pd.DataFrame(data) + psdf = ps.from_pandas(pdf) + + self.assertRaises( + NotImplementedError, lambda: psdf.groupby("a").describe().sort_index() + ) + + # multi-index columns + pdf = pd.DataFrame({("x", "a"): [1, 1, 3], ("x", "b"): [4, 5, 6], ("y", "c"): [7, 8, 9]}) + psdf = ps.from_pandas(pdf) + + describe_pdf = pdf.groupby(("x", "a")).describe().sort_index() + describe_psdf = psdf.groupby(("x", "a")).describe().sort_index() + + # 1. Check that non-percentile columns are equal. + agg_column_labels = [col._column_label for col in psdf.groupby(("x", "a"))._agg_columns] + self.assert_eq( + describe_psdf.drop( + columns=[ + tuple(list(label) + [s]) + for label, s in product(agg_column_labels, formatted_percentiles) + ] + ), + describe_pdf.drop(columns=formatted_percentiles, level=2), + check_exact=False, + ) + + # 2. Check that percentile columns are equal. + # The interpolation argument is yet to be implemented in Koalas. + quantile_pdf = pdf.groupby(("x", "a")).quantile(percentiles, interpolation="nearest") + quantile_pdf = quantile_pdf.unstack(level=1).astype(float) + + self.assert_eq( + describe_psdf.drop( + columns=[ + tuple(list(label) + [s]) + for label, s in product(agg_column_labels, non_percentile_stats) + ] + ), + quantile_pdf.rename(columns="{:.0%}".format, level=2), + ) + + +class GroupbyDescribeTests(GroupbyDescribeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_describe import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_groupby.py new file mode 100644 index 0000000..c353cc8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_groupby.py @@ -0,0 +1,979 @@ +# +# 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 unittest +import inspect +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError, DataError +from pyspark.pandas.missing.groupby import ( + MissingPandasLikeDataFrameGroupBy, + MissingPandasLikeSeriesGroupBy, +) +from pyspark.pandas.groupby import is_multi_agg_with_relabel, SeriesGroupBy +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class GroupByTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_groupby_simple(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 6, 4, 4, 6, 4, 3, 7], + "b": [4, 2, 7, 3, 3, 1, 1, 1, 2], + "c": [4, 2, 7, 3, None, 1, 1, 1, 2], + "d": list("abcdefght"), + "e": [True, False, True, False, True, False, True, False, True], + }, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + # TODO(SPARK-43295): Make DataFrameGroupBy.sum support for string type columns + pdf = pdf[["a", "b", "c", "e"]] + psdf = ps.from_pandas(pdf) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values("a").reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index).sum()), + sort(pdf.groupby("a", as_index=as_index).sum()), + ) + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index).b.sum()), + sort(pdf.groupby("a", as_index=as_index).b.sum()), + ) + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index)["b"].sum()), + sort(pdf.groupby("a", as_index=as_index)["b"].sum()), + ) + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index)[["b", "c"]].sum()), + sort(pdf.groupby("a", as_index=as_index)[["b", "c"]].sum()), + ) + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index)[[]].sum()), + sort(pdf.groupby("a", as_index=as_index)[[]].sum()), + ) + self.assert_eq( + sort(psdf.groupby("a", as_index=as_index)["c"].sum()), + sort(pdf.groupby("a", as_index=as_index)["c"].sum()), + ) + + self.assert_eq( + psdf.groupby("a").a.sum().sort_index(), pdf.groupby("a").a.sum().sort_index() + ) + self.assert_eq( + psdf.groupby("a")["a"].sum().sort_index(), pdf.groupby("a")["a"].sum().sort_index() + ) + self.assert_eq( + psdf.groupby("a")[["a"]].sum().sort_index(), pdf.groupby("a")[["a"]].sum().sort_index() + ) + self.assert_eq( + psdf.groupby("a")[["a", "c"]].sum().sort_index(), + pdf.groupby("a")[["a", "c"]].sum().sort_index(), + ) + + self.assert_eq( + psdf.a.groupby(psdf.b).sum().sort_index(), pdf.a.groupby(pdf.b).sum().sort_index() + ) + + for axis in [0, "index"]: + self.assert_eq( + psdf.groupby("a", axis=axis).a.sum().sort_index(), + pdf.groupby("a", axis=axis).a.sum().sort_index(), + ) + self.assert_eq( + psdf.groupby("a", axis=axis)["a"].sum().sort_index(), + pdf.groupby("a", axis=axis)["a"].sum().sort_index(), + ) + self.assert_eq( + psdf.groupby("a", axis=axis)[["a"]].sum().sort_index(), + pdf.groupby("a", axis=axis)[["a"]].sum().sort_index(), + ) + self.assert_eq( + psdf.groupby("a", axis=axis)[["a", "c"]].sum().sort_index(), + pdf.groupby("a", axis=axis)[["a", "c"]].sum().sort_index(), + ) + + self.assert_eq( + psdf.a.groupby(psdf.b, axis=axis).sum().sort_index(), + pdf.a.groupby(pdf.b, axis=axis).sum().sort_index(), + ) + + self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False).a) + self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)["a"]) + self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)[["a"]]) + self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)[["a", "c"]]) + self.assertRaises(KeyError, lambda: psdf.groupby("z", as_index=False)[["a", "c"]]) + self.assertRaises(KeyError, lambda: psdf.groupby(["z"], as_index=False)[["a", "c"]]) + + self.assertRaises(TypeError, lambda: psdf.a.groupby(psdf.b, as_index=False)) + + self.assertRaises(NotImplementedError, lambda: psdf.groupby("a", axis=1)) + self.assertRaises(NotImplementedError, lambda: psdf.groupby("a", axis="columns")) + self.assertRaises(ValueError, lambda: psdf.groupby("a", "b")) + self.assertRaises(TypeError, lambda: psdf.a.groupby(psdf.a, psdf.b)) + + # we can't use column name/names as a parameter `by` for `SeriesGroupBy`. + self.assertRaises(KeyError, lambda: psdf.a.groupby(by="a")) + self.assertRaises(KeyError, lambda: psdf.a.groupby(by=["a", "b"])) + self.assertRaises(KeyError, lambda: psdf.a.groupby(by=("a", "b"))) + self.assertRaises(KeyError, lambda: psdf.a.groupby(by=[("a", "b")])) + + # we can't use DataFrame as a parameter `by` for `DataFrameGroupBy`/`SeriesGroupBy`. + self.assertRaises(ValueError, lambda: psdf.groupby(psdf)) + self.assertRaises(ValueError, lambda: psdf.a.groupby(psdf)) + self.assertRaises(ValueError, lambda: psdf.a.groupby((psdf,))) + + with self.assertRaisesRegex(ValueError, "Grouper for 'list' not 1-dimensional"): + psdf.groupby(by=[["a", "b"]]) + + # non-string names + pdf = pd.DataFrame( + { + 10: [1, 2, 6, 4, 4, 6, 4, 3, 7], + 20: [4, 2, 7, 3, 3, 1, 1, 1, 2], + 30: [4, 2, 7, 3, None, 1, 1, 1, 2], + 40: list("abcdefght"), + }, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + psdf = ps.from_pandas(pdf) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + # TODO(SPARK-43295): Make DataFrameGroupBy.sum support for string type columns + pdf = pdf[[10, 20, 30]] + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(10).reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby(10, as_index=as_index).sum()), + sort(pdf.groupby(10, as_index=as_index).sum()), + ) + self.assert_eq( + sort(psdf.groupby(10, as_index=as_index)[20].sum()), + sort(pdf.groupby(10, as_index=as_index)[20].sum()), + ) + self.assert_eq( + sort(psdf.groupby(10, as_index=as_index)[[20, 30]].sum()), + sort(pdf.groupby(10, as_index=as_index)[[20, 30]].sum()), + ) + + def test_nsmallest(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + "c": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + "d": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + }, + index=np.random.rand(9 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby(["a"])["b"].nsmallest(1).sort_values(), + pdf.groupby(["a"])["b"].nsmallest(1).sort_values(), + ) + self.assert_eq( + psdf.groupby(["a"])["b"].nsmallest(2).sort_index(), + pdf.groupby(["a"])["b"].nsmallest(2).sort_index(), + ) + self.assert_eq( + (psdf.b * 10).groupby(psdf.a).nsmallest(2).sort_index(), + (pdf.b * 10).groupby(pdf.a).nsmallest(2).sort_index(), + ) + self.assert_eq( + psdf.b.rename().groupby(psdf.a).nsmallest(2).sort_index(), + pdf.b.rename().groupby(pdf.a).nsmallest(2).sort_index(), + ) + self.assert_eq( + psdf.b.groupby(psdf.a.rename()).nsmallest(2).sort_index(), + pdf.b.groupby(pdf.a.rename()).nsmallest(2).sort_index(), + ) + self.assert_eq( + psdf.b.rename().groupby(psdf.a.rename()).nsmallest(2).sort_index(), + pdf.b.rename().groupby(pdf.a.rename()).nsmallest(2).sort_index(), + ) + with self.assertRaisesRegex(ValueError, "nsmallest do not support multi-index now"): + psdf.set_index(["a", "b"]).groupby(["c"])["d"].nsmallest(1) + + def test_nlargest(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + "c": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + "d": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, + }, + index=np.random.rand(9 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby(["a"])["b"].nlargest(1).sort_values(), + pdf.groupby(["a"])["b"].nlargest(1).sort_values(), + ) + self.assert_eq( + psdf.groupby(["a"])["b"].nlargest(2).sort_index(), + pdf.groupby(["a"])["b"].nlargest(2).sort_index(), + ) + self.assert_eq( + (psdf.b * 10).groupby(psdf.a).nlargest(2).sort_index(), + (pdf.b * 10).groupby(pdf.a).nlargest(2).sort_index(), + ) + self.assert_eq( + psdf.b.rename().groupby(psdf.a).nlargest(2).sort_index(), + pdf.b.rename().groupby(pdf.a).nlargest(2).sort_index(), + ) + self.assert_eq( + psdf.b.groupby(psdf.a.rename()).nlargest(2).sort_index(), + pdf.b.groupby(pdf.a.rename()).nlargest(2).sort_index(), + ) + self.assert_eq( + psdf.b.rename().groupby(psdf.a.rename()).nlargest(2).sort_index(), + pdf.b.rename().groupby(pdf.a.rename()).nlargest(2).sort_index(), + ) + with self.assertRaisesRegex(ValueError, "nlargest do not support multi-index now"): + psdf.set_index(["a", "b"]).groupby(["c"])["d"].nlargest(1) + + def test_shift(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 2, 2, 3, 3] * 3, + "b": [1, 1, 2, 2, 3, 4] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("a").shift().sort_index(), pdf.groupby("a").shift().sort_index() + ) + # TODO: seems like a pandas' bug when fill_value is not None? + # self.assert_eq(psdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index(), + # pdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index()) + self.assert_eq( + psdf.groupby(["b"])["a"].shift().sort_index(), + pdf.groupby(["b"])["a"].shift().sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"])["c"].shift().sort_index(), + pdf.groupby(["a", "b"])["c"].shift().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).shift().sort_index(), + pdf.groupby(pdf.b // 5).shift().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].shift().sort_index(), + pdf.groupby(pdf.b // 5)["a"].shift().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).shift().sort_index(), + pdf.a.rename().groupby(pdf.b).shift().sort_index(), + ) + self.assert_eq( + psdf.a.groupby(psdf.b.rename()).shift().sort_index(), + pdf.a.groupby(pdf.b.rename()).shift().sort_index(), + ) + self.assert_eq( + psdf.a.rename().groupby(psdf.b.rename()).shift().sort_index(), + pdf.a.rename().groupby(pdf.b.rename()).shift().sort_index(), + ) + + self.assert_eq(psdf.groupby("a").shift().sum(), pdf.groupby("a").shift().sum().astype(int)) + self.assert_eq( + psdf.a.rename().groupby(psdf.b).shift().sum(), + pdf.a.rename().groupby(pdf.b).shift().sum(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "a")).shift().sort_index(), + pdf.groupby(("x", "a")).shift().sort_index(), + ) + # TODO: seems like a pandas' bug when fill_value is not None? + # self.assert_eq(psdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1, + # fill_value=0).sort_index(), + # pdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1, + # fill_value=0).sort_index()) + + def test_missing(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}) + + # DataFrameGroupBy functions + missing_functions = inspect.getmembers( + MissingPandasLikeDataFrameGroupBy, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a"), name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*GroupBy.*{}.*is deprecated".format(name) + ): + getattr(psdf.groupby("a"), name)() + + # SeriesGroupBy functions + missing_functions = inspect.getmembers(MissingPandasLikeSeriesGroupBy, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a), name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*GroupBy.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.groupby(psdf.a), name)() + + # DataFrameGroupBy properties + missing_properties = inspect.getmembers( + MissingPandasLikeDataFrameGroupBy, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a"), name) + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*GroupBy.*{}.*is deprecated".format(name) + ): + getattr(psdf.groupby("a"), name) + + # SeriesGroupBy properties + missing_properties = inspect.getmembers( + MissingPandasLikeSeriesGroupBy, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a), name) + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*GroupBy.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.groupby(psdf.a), name) + + @staticmethod + def test_is_multi_agg_with_relabel(): + + assert is_multi_agg_with_relabel(a="max") is False + assert is_multi_agg_with_relabel(a_min=("a", "max"), a_max=("a", "min")) is True + + def test_get_group(self): + pdf = pd.DataFrame( + [ + ("falcon", "bird", 389.0), + ("parrot", "bird", 24.0), + ("lion", "mammal", 80.5), + ("monkey", "mammal", np.nan), + ], + columns=["name", "class", "max_speed"], + index=[0, 2, 3, 1], + ) + pdf.columns.name = "Koalas" + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("class").get_group("bird"), + pdf.groupby("class").get_group("bird"), + ) + self.assert_eq( + psdf.groupby("class")["name"].get_group("mammal"), + pdf.groupby("class")["name"].get_group("mammal"), + ) + self.assert_eq( + psdf.groupby("class")[["name"]].get_group("mammal"), + pdf.groupby("class")[["name"]].get_group("mammal"), + ) + self.assert_eq( + psdf.groupby(["class", "name"]).get_group(("mammal", "lion")), + pdf.groupby(["class", "name"]).get_group(("mammal", "lion")), + ) + self.assert_eq( + psdf.groupby(["class", "name"])["max_speed"].get_group(("mammal", "lion")), + pdf.groupby(["class", "name"])["max_speed"].get_group(("mammal", "lion")), + ) + self.assert_eq( + psdf.groupby(["class", "name"])[["max_speed"]].get_group(("mammal", "lion")), + pdf.groupby(["class", "name"])[["max_speed"]].get_group(("mammal", "lion")), + ) + self.assert_eq( + (psdf.max_speed + 1).groupby(psdf["class"]).get_group("mammal"), + (pdf.max_speed + 1).groupby(pdf["class"]).get_group("mammal"), + ) + self.assert_eq( + psdf.groupby("max_speed").get_group(80.5), + pdf.groupby("max_speed").get_group(80.5), + ) + + self.assertRaises(KeyError, lambda: psdf.groupby("class").get_group("fish")) + self.assertRaises(TypeError, lambda: psdf.groupby("class").get_group(["bird", "mammal"])) + self.assertRaises(KeyError, lambda: psdf.groupby("class")["name"].get_group("fish")) + self.assertRaises( + TypeError, lambda: psdf.groupby("class")["name"].get_group(["bird", "mammal"]) + ) + self.assertRaises( + KeyError, lambda: psdf.groupby(["class", "name"]).get_group(("lion", "mammal")) + ) + self.assertRaises(ValueError, lambda: psdf.groupby(["class", "name"]).get_group(("lion",))) + self.assertRaises( + ValueError, lambda: psdf.groupby(["class", "name"]).get_group(("mammal",)) + ) + self.assertRaises(ValueError, lambda: psdf.groupby(["class", "name"]).get_group("mammal")) + + # MultiIndex columns + pdf.columns = pd.MultiIndex.from_tuples([("A", "name"), ("B", "class"), ("C", "max_speed")]) + pdf.columns.names = ["Hello", "Koalas"] + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby(("B", "class")).get_group("bird"), + pdf.groupby(("B", "class")).get_group("bird"), + ) + self.assert_eq( + psdf.groupby(("B", "class"))[[("A", "name")]].get_group("mammal"), + pdf.groupby(("B", "class"))[[("A", "name")]].get_group("mammal"), + ) + self.assert_eq( + psdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal", "lion")), + pdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal", "lion")), + ) + self.assert_eq( + psdf.groupby([("B", "class"), ("A", "name")])[[("C", "max_speed")]].get_group( + ("mammal", "lion") + ), + pdf.groupby([("B", "class"), ("A", "name")])[[("C", "max_speed")]].get_group( + ("mammal", "lion") + ), + ) + self.assert_eq( + (psdf[("C", "max_speed")] + 1).groupby(psdf[("B", "class")]).get_group("mammal"), + (pdf[("C", "max_speed")] + 1).groupby(pdf[("B", "class")]).get_group("mammal"), + ) + self.assert_eq( + psdf.groupby(("C", "max_speed")).get_group(80.5), + pdf.groupby(("C", "max_speed")).get_group(80.5), + ) + + self.assertRaises(KeyError, lambda: psdf.groupby(("B", "class")).get_group("fish")) + self.assertRaises( + TypeError, lambda: psdf.groupby(("B", "class")).get_group(["bird", "mammal"]) + ) + self.assertRaises( + KeyError, lambda: psdf.groupby(("B", "class"))[("A", "name")].get_group("fish") + ) + self.assertRaises( + KeyError, + lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("lion", "mammal")), + ) + self.assertRaises( + ValueError, + lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("lion",)), + ) + self.assertRaises( + ValueError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal",)) + ) + self.assertRaises( + ValueError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group("mammal") + ) + + def test_getitem(self): + psdf = ps.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, + }, + index=np.random.rand(10 * 3), + ) + + self.assertTrue(isinstance(psdf.groupby("a")["b"], SeriesGroupBy)) + + def test_all_any(self): + pdf = pd.DataFrame( + { + "A": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + "B": [True, True, True, False, False, False, None, True, None, False], + } + ) + psdf = ps.from_pandas(pdf) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values("A").reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index).all()), + sort(pdf.groupby("A", as_index=as_index).all()), + ) + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index).any()), + sort(pdf.groupby("A", as_index=as_index).any()), + ) + + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index).all()).B, + sort(pdf.groupby("A", as_index=as_index).all()).B, + ) + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index).any()).B, + sort(pdf.groupby("A", as_index=as_index).any()).B, + ) + + self.assert_eq( + psdf.B.groupby(psdf.A).all().sort_index(), pdf.B.groupby(pdf.A).all().sort_index() + ) + self.assert_eq( + psdf.B.groupby(psdf.A).any().sort_index(), pdf.B.groupby(pdf.A).any().sort_index() + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("Y", "B")]) + pdf.columns = columns + psdf.columns = columns + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(("X", "A")).reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby(("X", "A"), as_index=as_index).all()), + sort(pdf.groupby(("X", "A"), as_index=as_index).all()), + ) + self.assert_eq( + sort(psdf.groupby(("X", "A"), as_index=as_index).any()), + sort(pdf.groupby(("X", "A"), as_index=as_index).any()), + ) + + # Test skipna + pdf = pd.DataFrame({"A": [True, True], "B": [1, np.nan], "C": [True, None]}) + pdf.name = "x" + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("A").all(skipna=False).sort_index(), + pdf.groupby("A").all(skipna=False).sort_index(), + ) + self.assert_eq( + psdf.groupby("A").all(skipna=True).sort_index(), + pdf.groupby("A").all(skipna=True).sort_index(), + ) + + def test_raises(self): + psdf = ps.DataFrame( + {"a": [1, 2, 6, 4, 4, 6, 4, 3, 7], "b": [4, 2, 7, 3, 3, 1, 1, 1, 2]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + # test raises with incorrect key + self.assertRaises(ValueError, lambda: psdf.groupby([])) + self.assertRaises(KeyError, lambda: psdf.groupby("x")) + self.assertRaises(KeyError, lambda: psdf.groupby(["a", "x"])) + self.assertRaises(KeyError, lambda: psdf.groupby("a")["x"]) + self.assertRaises(KeyError, lambda: psdf.groupby("a")["b", "x"]) + self.assertRaises(KeyError, lambda: psdf.groupby("a")[["b", "x"]]) + + def test_nunique(self): + pdf = pd.DataFrame( + {"a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], "b": [2, 2, 2, 3, 3, 4, 4, 5, 5, 5]} + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("a").agg({"b": "nunique"}).sort_index(), + pdf.groupby("a").agg({"b": "nunique"}).sort_index(), + ) + if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): + expected = ps.DataFrame({"b": [2, 2]}, index=pd.Index([0, 1], name="a")) + self.assert_eq(psdf.groupby("a").nunique().sort_index(), expected) + self.assert_eq( + psdf.groupby("a").nunique(dropna=False).sort_index(), + expected, + ) + else: + self.assert_eq( + psdf.groupby("a").nunique().sort_index(), pdf.groupby("a").nunique().sort_index() + ) + self.assert_eq( + psdf.groupby("a").nunique(dropna=False).sort_index(), + pdf.groupby("a").nunique(dropna=False).sort_index(), + ) + self.assert_eq( + psdf.groupby("a")["b"].nunique().sort_index(), + pdf.groupby("a")["b"].nunique().sort_index(), + ) + self.assert_eq( + psdf.groupby("a")["b"].nunique(dropna=False).sort_index(), + pdf.groupby("a")["b"].nunique(dropna=False).sort_index(), + ) + + nunique_psdf = psdf.groupby("a", as_index=False).agg({"b": "nunique"}) + nunique_pdf = pdf.groupby("a", as_index=False).agg({"b": "nunique"}) + self.assert_eq( + nunique_psdf.sort_values(["a", "b"]).reset_index(drop=True), + nunique_pdf.sort_values(["a", "b"]).reset_index(drop=True), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + + if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): + expected = ps.DataFrame({("y", "b"): [2, 2]}, index=pd.Index([0, 1], name=("x", "a"))) + self.assert_eq( + psdf.groupby(("x", "a")).nunique().sort_index(), + expected, + ) + self.assert_eq( + psdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), + expected, + ) + else: + self.assert_eq( + psdf.groupby(("x", "a")).nunique().sort_index(), + pdf.groupby(("x", "a")).nunique().sort_index(), + ) + self.assert_eq( + psdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), + pdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), + ) + + def test_unique(self): + for pdf in [ + pd.DataFrame( + {"a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], "b": [2, 2, 2, 3, 3, 4, 4, 5, 5, 5]} + ), + pd.DataFrame( + { + "a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], + "b": ["w", "w", "w", "x", "x", "y", "y", "z", "z", "z"], + } + ), + ]: + with self.subTest(pdf=pdf): + psdf = ps.from_pandas(pdf) + + actual = psdf.groupby("a")["b"].unique().sort_index()._to_pandas() + expect = pdf.groupby("a")["b"].unique().sort_index() + self.assert_eq(len(actual), len(expect)) + for act, exp in zip(actual, expect): + self.assertTrue(sorted(act) == sorted(exp)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43444): Enable GroupBySlowTests.test_value_counts for pandas 2.0.0.", + ) + def test_value_counts(self): + pdf = pd.DataFrame( + {"A": [np.nan, 2, 2, 3, 3, 3], "B": [1, 1, 2, 3, 3, np.nan]}, columns=["A", "B"] + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + psdf.groupby("A")["B"].value_counts().sort_index(), + pdf.groupby("A")["B"].value_counts().sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"].value_counts(dropna=False).sort_index(), + pdf.groupby("A")["B"].value_counts(dropna=False).sort_index(), + ) + self.assert_eq( + psdf.groupby("A", dropna=False)["B"].value_counts(dropna=False).sort_index(), + pdf.groupby("A", dropna=False)["B"].value_counts(dropna=False).sort_index(), + # Returns are the same considering values and types, + # disable check_exact to pass the assert_eq + check_exact=False, + ) + self.assert_eq( + psdf.groupby("A")["B"].value_counts(sort=True, ascending=False).sort_index(), + pdf.groupby("A")["B"].value_counts(sort=True, ascending=False).sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"] + .value_counts(sort=True, ascending=False, dropna=False) + .sort_index(), + pdf.groupby("A")["B"] + .value_counts(sort=True, ascending=False, dropna=False) + .sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"] + .value_counts(sort=True, ascending=True, dropna=False) + .sort_index(), + pdf.groupby("A")["B"] + .value_counts(sort=True, ascending=True, dropna=False) + .sort_index(), + ) + self.assert_eq( + psdf.B.rename().groupby(psdf.A).value_counts().sort_index(), + pdf.B.rename().groupby(pdf.A).value_counts().sort_index(), + ) + self.assert_eq( + psdf.B.rename().groupby(psdf.A, dropna=False).value_counts().sort_index(), + pdf.B.rename().groupby(pdf.A, dropna=False).value_counts().sort_index(), + # Returns are the same considering values and types, + # disable check_exact to pass the assert_eq + check_exact=False, + ) + self.assert_eq( + psdf.B.groupby(psdf.A.rename()).value_counts().sort_index(), + pdf.B.groupby(pdf.A.rename()).value_counts().sort_index(), + ) + self.assert_eq( + psdf.B.rename().groupby(psdf.A.rename()).value_counts().sort_index(), + pdf.B.rename().groupby(pdf.A.rename()).value_counts().sort_index(), + ) + + def test_size(self): + pdf = pd.DataFrame({"A": [1, 2, 2, 3, 3, 3], "B": [1, 1, 2, 3, 3, 3]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.groupby("A").size().sort_index(), pdf.groupby("A").size().sort_index()) + self.assert_eq( + psdf.groupby("A")["B"].size().sort_index(), pdf.groupby("A")["B"].size().sort_index() + ) + self.assert_eq( + psdf.groupby("A")[["B"]].size().sort_index(), + pdf.groupby("A")[["B"]].size().sort_index(), + ) + self.assert_eq( + psdf.groupby(["A", "B"]).size().sort_index(), + pdf.groupby(["A", "B"]).size().sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("Y", "B")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("X", "A")).size().sort_index(), + pdf.groupby(("X", "A")).size().sort_index(), + ) + self.assert_eq( + psdf.groupby([("X", "A"), ("Y", "B")]).size().sort_index(), + pdf.groupby([("X", "A"), ("Y", "B")]).size().sort_index(), + ) + + def test_diff(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.groupby("b").diff().sort_index(), pdf.groupby("b").diff().sort_index()) + self.assert_eq( + psdf.groupby(["a", "b"]).diff().sort_index(), + pdf.groupby(["a", "b"]).diff().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].diff().sort_index(), + pdf.groupby(["b"])["a"].diff().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "b"]].diff().sort_index(), + pdf.groupby(["b"])[["a", "b"]].diff().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).diff().sort_index(), + pdf.groupby(pdf.b // 5).diff().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].diff().sort_index(), + pdf.groupby(pdf.b // 5)["a"].diff().sort_index(), + ) + + self.assert_eq(psdf.groupby("b").diff().sum(), pdf.groupby("b").diff().sum().astype(int)) + self.assert_eq(psdf.groupby(["b"])["a"].diff().sum(), pdf.groupby(["b"])["a"].diff().sum()) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).diff().sort_index(), + pdf.groupby(("x", "b")).diff().sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).diff().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).diff().sort_index(), + ) + + def test_rank(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + index=np.random.rand(6 * 3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.groupby("b").rank().sort_index(), pdf.groupby("b").rank().sort_index()) + self.assert_eq( + psdf.groupby(["a", "b"]).rank().sort_index(), + pdf.groupby(["a", "b"]).rank().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])["a"].rank().sort_index(), + pdf.groupby(["b"])["a"].rank().sort_index(), + ) + self.assert_eq( + psdf.groupby(["b"])[["a", "c"]].rank().sort_index(), + pdf.groupby(["b"])[["a", "c"]].rank().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5).rank().sort_index(), + pdf.groupby(pdf.b // 5).rank().sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.b // 5)["a"].rank().sort_index(), + pdf.groupby(pdf.b // 5)["a"].rank().sort_index(), + ) + + self.assert_eq(psdf.groupby("b").rank().sum(), pdf.groupby("b").rank().sum()) + self.assert_eq(psdf.groupby(["b"])["a"].rank().sum(), pdf.groupby(["b"])["a"].rank().sum()) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("x", "b")).rank().sort_index(), + pdf.groupby(("x", "b")).rank().sort_index(), + ) + self.assert_eq( + psdf.groupby([("x", "a"), ("x", "b")]).rank().sort_index(), + pdf.groupby([("x", "a"), ("x", "b")]).rank().sort_index(), + ) + + +class GroupByTests(GroupByTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_groupby import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_head_tail.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_head_tail.py new file mode 100644 index 0000000..1a22db7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_head_tail.py @@ -0,0 +1,234 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyHeadTailMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_head(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, + }, + index=np.random.rand(10 * 3), + ) + psdf = ps.from_pandas(pdf) + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby("a").head(limit).sort_index(), + psdf.groupby("a").head(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")["b"].head(limit).sort_index(), + psdf.groupby("a")["b"].head(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")[["b"]].head(limit).sort_index(), + psdf.groupby("a")[["b"]].head(limit).sort_index(), + ) + + self.assert_eq( + pdf.groupby(pdf.a // 2).head(2).sort_index(), + psdf.groupby(psdf.a // 2).head(2).sort_index(), + ) + self.assert_eq( + pdf.groupby(pdf.a // 2)["b"].head(2).sort_index(), + psdf.groupby(psdf.a // 2)["b"].head(2).sort_index(), + ) + self.assert_eq( + pdf.groupby(pdf.a // 2)[["b"]].head(2).sort_index(), + psdf.groupby(psdf.a // 2)[["b"]].head(2).sort_index(), + ) + + self.assert_eq( + pdf.b.rename().groupby(pdf.a).head(2).sort_index(), + psdf.b.rename().groupby(psdf.a).head(2).sort_index(), + ) + self.assert_eq( + pdf.b.groupby(pdf.a.rename()).head(2).sort_index(), + psdf.b.groupby(psdf.a.rename()).head(2).sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a.rename()).head(2).sort_index(), + psdf.b.rename().groupby(psdf.a.rename()).head(2).sort_index(), + ) + + # multi-index + midx = pd.MultiIndex( + [["x", "y"], ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]], + [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], + ) + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3], + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5], + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6], + }, + columns=["a", "b", "c"], + index=midx, + ) + psdf = ps.from_pandas(pdf) + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby("a").head(limit).sort_index(), + psdf.groupby("a").head(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")["b"].head(limit).sort_index(), + psdf.groupby("a")["b"].head(limit).sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby(("x", "a")).head(limit).sort_index(), + psdf.groupby(("x", "a")).head(limit).sort_index(), + ) + + def test_tail(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, + }, + index=np.random.rand(10 * 3), + ) + psdf = ps.from_pandas(pdf) + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby("a").tail(limit).sort_index(), + psdf.groupby("a").tail(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")["b"].tail(limit).sort_index(), + psdf.groupby("a")["b"].tail(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")[["b"]].tail(limit).sort_index(), + psdf.groupby("a")[["b"]].tail(limit).sort_index(), + ) + + self.assert_eq( + pdf.groupby(pdf.a // 2).tail(2).sort_index(), + psdf.groupby(psdf.a // 2).tail(2).sort_index(), + ) + self.assert_eq( + pdf.groupby(pdf.a // 2)["b"].tail(2).sort_index(), + psdf.groupby(psdf.a // 2)["b"].tail(2).sort_index(), + ) + self.assert_eq( + pdf.groupby(pdf.a // 2)[["b"]].tail(2).sort_index(), + psdf.groupby(psdf.a // 2)[["b"]].tail(2).sort_index(), + ) + + self.assert_eq( + pdf.b.rename().groupby(pdf.a).tail(2).sort_index(), + psdf.b.rename().groupby(psdf.a).tail(2).sort_index(), + ) + self.assert_eq( + pdf.b.groupby(pdf.a.rename()).tail(2).sort_index(), + psdf.b.groupby(psdf.a.rename()).tail(2).sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a.rename()).tail(2).sort_index(), + psdf.b.rename().groupby(psdf.a.rename()).tail(2).sort_index(), + ) + + # multi-index + midx = pd.MultiIndex( + [["x", "y"], ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]], + [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], + ) + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3], + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5], + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6], + }, + columns=["a", "b", "c"], + index=midx, + ) + psdf = ps.from_pandas(pdf) + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby("a").tail(limit).sort_index(), + psdf.groupby("a").tail(limit).sort_index(), + ) + self.assert_eq( + pdf.groupby("a")["b"].tail(limit).sort_index(), + psdf.groupby("a")["b"].tail(limit).sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + for limit in (2, 100000, -2, -100000, -1): + self.assert_eq( + pdf.groupby(("x", "a")).tail(limit).sort_index(), + psdf.groupby(("x", "a")).tail(limit).sort_index(), + ) + + +class GroupbyHeadTailTests(GroupbyHeadTailMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_head_tail import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_index.py new file mode 100644 index 0000000..d38c964 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_index.py @@ -0,0 +1,206 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyIndexMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43555): Enable GroupByTests.test_groupby_multiindex_columns for pandas 2.0.0.", + ) + def test_groupby_multiindex_columns(self): + pdf = pd.DataFrame( + { + (10, "a"): [1, 2, 6, 4, 4, 6, 4, 3, 7], + (10, "b"): [4, 2, 7, 3, 3, 1, 1, 1, 2], + (20, "c"): [4, 2, 7, 3, None, 1, 1, 1, 2], + (30, "d"): list("abcdefght"), + }, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby((10, "a")).sum().sort_index(), pdf.groupby((10, "a")).sum().sort_index() + ) + self.assert_eq( + psdf.groupby((10, "a"), as_index=False) + .sum() + .sort_values((10, "a")) + .reset_index(drop=True), + pdf.groupby((10, "a"), as_index=False) + .sum() + .sort_values((10, "a")) + .reset_index(drop=True), + ) + self.assert_eq( + psdf.groupby((10, "a"))[[(20, "c")]].sum().sort_index(), + pdf.groupby((10, "a"))[[(20, "c")]].sum().sort_index(), + ) + + # TODO: a pandas bug? + # expected = pdf.groupby((10, "a"))[(20, "c")].sum().sort_index() + expected = pd.Series( + [4.0, 2.0, 1.0, 4.0, 8.0, 2.0], + name=(20, "c"), + index=pd.Index([1, 2, 3, 4, 6, 7], name=(10, "a")), + ) + + self.assert_eq(psdf.groupby((10, "a"))[(20, "c")].sum().sort_index(), expected) + + if LooseVersion(pd.__version__) != LooseVersion("1.1.3") and LooseVersion( + pd.__version__ + ) != LooseVersion("1.1.4"): + self.assert_eq( + psdf[(20, "c")].groupby(psdf[(10, "a")]).sum().sort_index(), + pdf[(20, "c")].groupby(pdf[(10, "a")]).sum().sort_index(), + ) + else: + # Due to pandas bugs resolved in 1.0.4, re-introduced in 1.1.3 and resolved in 1.1.5 + self.assert_eq(psdf[(20, "c")].groupby(psdf[(10, "a")]).sum().sort_index(), expected) + + def test_idxmax(self): + pdf = pd.DataFrame( + {"a": [1, 1, 2, 2, 3] * 3, "b": [1, 2, 3, 4, 5] * 3, "c": [5, 4, 3, 2, 1] * 3} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.groupby(["a"]).idxmax().sort_index(), psdf.groupby(["a"]).idxmax().sort_index() + ) + self.assert_eq( + pdf.groupby(["a"]).idxmax(skipna=False).sort_index(), + psdf.groupby(["a"]).idxmax(skipna=False).sort_index(), + ) + self.assert_eq( + pdf.groupby(["a"])["b"].idxmax().sort_index(), + psdf.groupby(["a"])["b"].idxmax().sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a).idxmax().sort_index(), + psdf.b.rename().groupby(psdf.a).idxmax().sort_index(), + ) + self.assert_eq( + pdf.b.groupby(pdf.a.rename()).idxmax().sort_index(), + psdf.b.groupby(psdf.a.rename()).idxmax().sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a.rename()).idxmax().sort_index(), + psdf.b.rename().groupby(psdf.a.rename()).idxmax().sort_index(), + ) + + with self.assertRaisesRegex(ValueError, "idxmax only support one-level index now"): + psdf.set_index(["a", "b"]).groupby(["c"]).idxmax() + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + pdf.groupby(("x", "a")).idxmax().sort_index(), + psdf.groupby(("x", "a")).idxmax().sort_index(), + ) + self.assert_eq( + pdf.groupby(("x", "a")).idxmax(skipna=False).sort_index(), + psdf.groupby(("x", "a")).idxmax(skipna=False).sort_index(), + ) + + def test_idxmin(self): + pdf = pd.DataFrame( + {"a": [1, 1, 2, 2, 3] * 3, "b": [1, 2, 3, 4, 5] * 3, "c": [5, 4, 3, 2, 1] * 3} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.groupby(["a"]).idxmin().sort_index(), psdf.groupby(["a"]).idxmin().sort_index() + ) + self.assert_eq( + pdf.groupby(["a"]).idxmin(skipna=False).sort_index(), + psdf.groupby(["a"]).idxmin(skipna=False).sort_index(), + ) + self.assert_eq( + pdf.groupby(["a"])["b"].idxmin().sort_index(), + psdf.groupby(["a"])["b"].idxmin().sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a).idxmin().sort_index(), + psdf.b.rename().groupby(psdf.a).idxmin().sort_index(), + ) + self.assert_eq( + pdf.b.groupby(pdf.a.rename()).idxmin().sort_index(), + psdf.b.groupby(psdf.a.rename()).idxmin().sort_index(), + ) + self.assert_eq( + pdf.b.rename().groupby(pdf.a.rename()).idxmin().sort_index(), + psdf.b.rename().groupby(psdf.a.rename()).idxmin().sort_index(), + ) + + with self.assertRaisesRegex(ValueError, "idxmin only support one-level index now"): + psdf.set_index(["a", "b"]).groupby(["c"]).idxmin() + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + pdf.groupby(("x", "a")).idxmin().sort_index(), + psdf.groupby(("x", "a")).idxmin().sort_index(), + ) + self.assert_eq( + pdf.groupby(("x", "a")).idxmin(skipna=False).sort_index(), + psdf.groupby(("x", "a")).idxmin(skipna=False).sort_index(), + ) + + +class GroupbyIndexTests(GroupbyIndexMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_index import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_missing_data.py new file mode 100644 index 0000000..419a6f0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_missing_data.py @@ -0,0 +1,421 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyMissingDataMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_fillna(self): + pdf = pd.DataFrame( + { + "A": [1, 1, 2, 2] * 3, + "B": [2, 4, None, 3] * 3, + "C": [None, None, None, 1] * 3, + "D": [0, 1, 5, 4] * 3, + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("A").fillna(0).sort_index(), pdf.groupby("A").fillna(0).sort_index() + ) + self.assert_eq( + psdf.groupby("A")["C"].fillna(0).sort_index(), + pdf.groupby("A")["C"].fillna(0).sort_index(), + ) + self.assert_eq( + psdf.groupby("A")[["C"]].fillna(0).sort_index(), + pdf.groupby("A")[["C"]].fillna(0).sort_index(), + ) + self.assert_eq( + psdf.groupby("A").fillna(method="bfill").sort_index(), + pdf.groupby("A").fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["C"].fillna(method="bfill").sort_index(), + pdf.groupby("A")["C"].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby("A")[["C"]].fillna(method="bfill").sort_index(), + pdf.groupby("A")[["C"]].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby("A").fillna(method="ffill").sort_index(), + pdf.groupby("A").fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["C"].fillna(method="ffill").sort_index(), + pdf.groupby("A")["C"].fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby("A")[["C"]].fillna(method="ffill").sort_index(), + pdf.groupby("A")[["C"]].fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5).fillna(method="bfill").sort_index(), + pdf.groupby(pdf.A // 5).fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5)["C"].fillna(method="bfill").sort_index(), + pdf.groupby(pdf.A // 5)["C"].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5)[["C"]].fillna(method="bfill").sort_index(), + pdf.groupby(pdf.A // 5)[["C"]].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5).fillna(method="ffill").sort_index(), + pdf.groupby(pdf.A // 5).fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5)["C"].fillna(method="ffill").sort_index(), + pdf.groupby(pdf.A // 5)["C"].fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby(psdf.A // 5)[["C"]].fillna(method="ffill").sort_index(), + pdf.groupby(pdf.A // 5)[["C"]].fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.C.rename().groupby(psdf.A).fillna(0).sort_index(), + pdf.C.rename().groupby(pdf.A).fillna(0).sort_index(), + ) + self.assert_eq( + psdf.C.groupby(psdf.A.rename()).fillna(0).sort_index(), + pdf.C.groupby(pdf.A.rename()).fillna(0).sort_index(), + ) + self.assert_eq( + psdf.C.rename().groupby(psdf.A.rename()).fillna(0).sort_index(), + pdf.C.rename().groupby(pdf.A.rename()).fillna(0).sort_index(), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("X", "A")).fillna(0).sort_index(), + pdf.groupby(("X", "A")).fillna(0).sort_index(), + ) + self.assert_eq( + psdf.groupby(("X", "A")).fillna(method="bfill").sort_index(), + pdf.groupby(("X", "A")).fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(("X", "A")).fillna(method="ffill").sort_index(), + pdf.groupby(("X", "A")).fillna(method="ffill").sort_index(), + ) + + def test_ffill(self): + idx = np.random.rand(4 * 3) + pdf = pd.DataFrame( + { + "A": [1, 1, 2, 2] * 3, + "B": [2, 4, None, 3] * 3, + "C": [None, None, None, 1] * 3, + "D": [0, 1, 5, 4] * 3, + }, + index=idx, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("A").ffill().sort_index(), pdf.groupby("A").ffill().sort_index() + ) + self.assert_eq( + psdf.groupby("A")[["B"]].ffill().sort_index(), + pdf.groupby("A")[["B"]].ffill().sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"].ffill().sort_index(), pdf.groupby("A")["B"].ffill().sort_index() + ) + self.assert_eq( + psdf.groupby("A")["B"].ffill()[idx[6]], pdf.groupby("A")["B"].ffill()[idx[6]] + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("X", "A")).ffill().sort_index(), + pdf.groupby(("X", "A")).ffill().sort_index(), + ) + + def test_bfill(self): + idx = np.random.rand(4 * 3) + pdf = pd.DataFrame( + { + "A": [1, 1, 2, 2] * 3, + "B": [2, 4, None, 3] * 3, + "C": [None, None, None, 1] * 3, + "D": [0, 1, 5, 4] * 3, + }, + index=idx, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf.groupby("A").bfill().sort_index(), pdf.groupby("A").bfill().sort_index() + ) + self.assert_eq( + psdf.groupby("A")[["B"]].bfill().sort_index(), + pdf.groupby("A")[["B"]].bfill().sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"].bfill().sort_index(), + pdf.groupby("A")["B"].bfill().sort_index(), + ) + self.assert_eq( + psdf.groupby("A")["B"].bfill()[idx[6]], pdf.groupby("A")["B"].bfill()[idx[6]] + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")]) + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + psdf.groupby(("X", "A")).bfill().sort_index(), + pdf.groupby(("X", "A")).bfill().sort_index(), + ) + + def test_dropna(self): + pdf = pd.DataFrame( + {"A": [None, 1, None, 1, 2], "B": [1, 2, 3, None, None], "C": [4, 5, 6, 7, None]} + ) + psdf = ps.from_pandas(pdf) + + # pd.DataFrame.groupby with dropna parameter is implemented since pandas 1.1.0 + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + for dropna in [True, False]: + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values("A").reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index, dropna=dropna).std()), + sort(pdf.groupby("A", as_index=as_index, dropna=dropna).std()), + ) + + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index, dropna=dropna).B.std()), + sort(pdf.groupby("A", as_index=as_index, dropna=dropna).B.std()), + ) + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index, dropna=dropna)["B"].std()), + sort(pdf.groupby("A", as_index=as_index, dropna=dropna)["B"].std()), + ) + + self.assert_eq( + sort( + psdf.groupby("A", as_index=as_index, dropna=dropna).agg( + {"B": "min", "C": "std"} + ) + ), + sort( + pdf.groupby("A", as_index=as_index, dropna=dropna).agg( + {"B": "min", "C": "std"} + ) + ), + ) + + for dropna in [True, False]: + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(["A", "B"]).reset_index(drop=True) + + self.assert_eq( + sort( + psdf.groupby(["A", "B"], as_index=as_index, dropna=dropna).agg( + {"C": ["min", "std"]} + ) + ), + sort( + pdf.groupby(["A", "B"], as_index=as_index, dropna=dropna).agg( + {"C": ["min", "std"]} + ) + ), + almost=True, + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + pdf.columns = columns + psdf.columns = columns + + for dropna in [True, False]: + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(("X", "A")).reset_index(drop=True) + + sorted_stats_psdf = sort( + psdf.groupby(("X", "A"), as_index=as_index, dropna=dropna).agg( + {("X", "B"): "min", ("Y", "C"): "std"} + ) + ) + sorted_stats_pdf = sort( + pdf.groupby(("X", "A"), as_index=as_index, dropna=dropna).agg( + {("X", "B"): "min", ("Y", "C"): "std"} + ) + ) + self.assert_eq(sorted_stats_psdf, sorted_stats_pdf) + else: + # Testing dropna=True (pandas default behavior) + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values("A").reset_index(drop=True) + + self.assert_eq( + sort(psdf.groupby("A", as_index=as_index, dropna=True)["B"].min()), + sort(pdf.groupby("A", as_index=as_index)["B"].min()), + ) + + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(["A", "B"]).reset_index(drop=True) + + self.assert_eq( + sort( + psdf.groupby(["A", "B"], as_index=as_index, dropna=True).agg( + {"C": ["min", "std"]} + ) + ), + sort(pdf.groupby(["A", "B"], as_index=as_index).agg({"C": ["min", "std"]})), + almost=True, + ) + + # Testing dropna=False + index = pd.Index([1.0, 2.0, np.nan], name="A") + expected = pd.Series([2.0, np.nan, 1.0], index=index, name="B") + result = psdf.groupby("A", as_index=True, dropna=False)["B"].min().sort_index() + self.assert_eq(expected, result) + + expected = pd.DataFrame({"A": [1.0, 2.0, np.nan], "B": [2.0, np.nan, 1.0]}) + result = ( + psdf.groupby("A", as_index=False, dropna=False)["B"] + .min() + .sort_values("A") + .reset_index(drop=True) + ) + self.assert_eq(expected, result) + + index = pd.MultiIndex.from_tuples( + [(1.0, 2.0), (1.0, None), (2.0, None), (None, 1.0), (None, 3.0)], names=["A", "B"] + ) + expected = pd.DataFrame( + { + ("C", "min"): [5.0, 7.0, np.nan, 4.0, 6.0], + ("C", "std"): [np.nan, np.nan, np.nan, np.nan, np.nan], + }, + index=index, + ) + result = ( + psdf.groupby(["A", "B"], as_index=True, dropna=False) + .agg({"C": ["min", "std"]}) + .sort_index() + ) + self.assert_eq(expected, result) + + expected = pd.DataFrame( + { + ("A", ""): [1.0, 1.0, 2.0, np.nan, np.nan], + ("B", ""): [2.0, np.nan, np.nan, 1.0, 3.0], + ("C", "min"): [5.0, 7.0, np.nan, 4.0, 6.0], + ("C", "std"): [np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + result = ( + psdf.groupby(["A", "B"], as_index=False, dropna=False) + .agg({"C": ["min", "std"]}) + .sort_values(["A", "B"]) + .reset_index(drop=True) + ) + self.assert_eq(expected, result) + + +class GroupbyMissingDataTests(GroupbyMissingDataMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_missing_data import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_split_apply.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_split_apply.py new file mode 100644 index 0000000..070fa01 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_split_apply.py @@ -0,0 +1,187 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbySplitApplyMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43445): Enable GroupBySlowTests.test_split_apply_combine_on_series " + "for pandas 2.0.0.", + ) + def test_split_apply_combine_on_series(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 6, 4, 4, 6, 4, 3, 7], + "b": [4, 2, 7, 3, 3, 1, 1, 1, 2], + "c": [4, 2, 7, 3, None, 1, 1, 1, 2], + "d": list("abcdefght"), + }, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + psdf = ps.from_pandas(pdf) + + funcs = [ + ((True, False), ["sum", "min", "max", "count", "first", "last"]), + ((True, True), ["mean"]), + ((False, False), ["var", "std", "skew"]), + ] + funcs = [(check_exact, almost, f) for (check_exact, almost), fs in funcs for f in fs] + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(list(df.columns)).reset_index(drop=True) + + for check_exact, almost, func in funcs: + for kkey, pkey in [("b", "b"), (psdf.b, pdf.b)]: + with self.subTest(as_index=as_index, func=func, key=pkey): + if as_index is True or func != "std": + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()), + sort(getattr(pdf.groupby(pkey, as_index=as_index).a, func)()), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()), + sort(getattr(pdf.groupby(pkey, as_index=as_index), func)()), + check_exact=check_exact, + almost=almost, + ) + else: + # seems like a pandas' bug for as_index=False and func == "std"? + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()), + sort(pdf.groupby(pkey, as_index=True).a.std().reset_index()), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()), + sort(pdf.groupby(pkey, as_index=True).std().reset_index()), + check_exact=check_exact, + almost=almost, + ) + + for kkey, pkey in [(psdf.b + 1, pdf.b + 1), (psdf.copy().b, pdf.copy().b)]: + with self.subTest(as_index=as_index, func=func, key=pkey): + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()), + sort(getattr(pdf.groupby(pkey, as_index=as_index).a, func)()), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()), + sort(getattr(pdf.groupby(pkey, as_index=as_index), func)()), + check_exact=check_exact, + almost=almost, + ) + + for check_exact, almost, func in funcs: + for i in [0, 4, 7]: + with self.subTest(as_index=as_index, func=func, i=i): + self.assert_eq( + sort(getattr(psdf.groupby(psdf.b > i, as_index=as_index).a, func)()), + sort(getattr(pdf.groupby(pdf.b > i, as_index=as_index).a, func)()), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + sort(getattr(psdf.groupby(psdf.b > i, as_index=as_index), func)()), + sort(getattr(pdf.groupby(pdf.b > i, as_index=as_index), func)()), + check_exact=check_exact, + almost=almost, + ) + + for check_exact, almost, func in funcs: + for kkey, pkey in [ + (psdf.b, pdf.b), + (psdf.b + 1, pdf.b + 1), + (psdf.copy().b, pdf.copy().b), + (psdf.b.rename(), pdf.b.rename()), + ]: + with self.subTest(func=func, key=pkey): + self.assert_eq( + getattr(psdf.a.groupby(kkey), func)().sort_index(), + getattr(pdf.a.groupby(pkey), func)().sort_index(), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + getattr((psdf.a + 1).groupby(kkey), func)().sort_index(), + getattr((pdf.a + 1).groupby(pkey), func)().sort_index(), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + getattr((psdf.b + 1).groupby(kkey), func)().sort_index(), + getattr((pdf.b + 1).groupby(pkey), func)().sort_index(), + check_exact=check_exact, + almost=almost, + ) + self.assert_eq( + getattr(psdf.a.rename().groupby(kkey), func)().sort_index(), + getattr(pdf.a.rename().groupby(pkey), func)().sort_index(), + check_exact=check_exact, + almost=almost, + ) + + +class GroupbySplitApplyTests(GroupbySplitApplyMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_split_apply import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_stat.py new file mode 100644 index 0000000..bfdeeec --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/groupby/test_stat.py @@ -0,0 +1,397 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class GroupbyStatMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [3.1, 4.1, 4.1, 3.1], + "C": ["a", "b", "b", "a"], + "D": [True, False, False, True], + } + ) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + # TODO: All statistical functions should leverage this utility + def _test_stat_func(self, func, check_exact=True): + pdf, psdf = self.pdf, self.psdf + for p_groupby_obj, ps_groupby_obj in [ + # Against DataFrameGroupBy + (pdf.groupby("A"), psdf.groupby("A")), + # Against DataFrameGroupBy with an aggregation column of string type + (pdf.groupby("A")[["C"]], psdf.groupby("A")[["C"]]), + # Against SeriesGroupBy + (pdf.groupby("A")["B"], psdf.groupby("A")["B"]), + ]: + self.assert_eq( + func(p_groupby_obj).sort_index(), + func(ps_groupby_obj).sort_index(), + check_exact=check_exact, + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43554): Enable GroupByTests.test_basic_stat_funcs for pandas 2.0.0.", + ) + def test_basic_stat_funcs(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.var(), check_exact=False) + + pdf, psdf = self.pdf, self.psdf + + # Unlike pandas', the median in pandas-on-Spark is an approximated median based upon + # approximate percentile computation because computing median across a large dataset + # is extremely expensive. + expected = ps.DataFrame({"B": [3.1, 3.1], "D": [0, 0]}, index=pd.Index([1, 2], name="A")) + self.assert_eq( + psdf.groupby("A").median().sort_index(), + expected, + ) + self.assert_eq( + psdf.groupby("A").median(numeric_only=None).sort_index(), + expected, + ) + self.assert_eq( + psdf.groupby("A").median(numeric_only=False).sort_index(), + expected, + ) + self.assert_eq( + psdf.groupby("A")["B"].median().sort_index(), + expected.B, + ) + with self.assertRaises(TypeError): + psdf.groupby("A")["C"].mean() + + with self.assertRaisesRegex( + TypeError, "Unaccepted data types of aggregation columns; numeric or bool expected." + ): + psdf.groupby("A")[["C"]].std() + + with self.assertRaisesRegex( + TypeError, "Unaccepted data types of aggregation columns; numeric or bool expected." + ): + psdf.groupby("A")[["C"]].sem() + + self.assert_eq( + psdf.groupby("A").std().sort_index(), + pdf.groupby("A").std().sort_index(), + check_exact=False, + ) + self.assert_eq( + psdf.groupby("A").sem().sort_index(), + pdf.groupby("A").sem().sort_index(), + check_exact=False, + ) + + # TODO: fix bug of `sum` and re-enable the test below + # self._test_stat_func(lambda groupby_obj: groupby_obj.sum(), check_exact=False) + self.assert_eq( + psdf.groupby("A").sum().sort_index(), + pdf.groupby("A").sum().sort_index(), + check_exact=False, + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43706): Enable GroupByTests.test_mean " "for pandas 2.0.0.", + ) + def test_mean(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.mean()) + self._test_stat_func(lambda groupby_obj: groupby_obj.mean(numeric_only=None)) + self._test_stat_func(lambda groupby_obj: groupby_obj.mean(numeric_only=True)) + psdf = self.psdf + with self.assertRaises(TypeError): + psdf.groupby("A")["C"].mean() + + def test_quantile(self): + dfs = [ + pd.DataFrame( + [["a", 1], ["a", 2], ["a", 3], ["b", 1], ["b", 3], ["b", 5]], columns=["key", "val"] + ), + pd.DataFrame( + [["a", True], ["a", True], ["a", False], ["b", True], ["b", True], ["b", False]], + columns=["key", "val"], + ), + ] + for df in dfs: + psdf = ps.from_pandas(df) + # q accept float and int between 0 and 1 + for i in [0, 0.1, 0.5, 1]: + self.assert_eq( + df.groupby("key").quantile(q=i, interpolation="lower"), + psdf.groupby("key").quantile(q=i), + almost=True, + ) + self.assert_eq( + df.groupby("key")["val"].quantile(q=i, interpolation="lower"), + psdf.groupby("key")["val"].quantile(q=i), + almost=True, + ) + # raise ValueError when q not in [0, 1] + with self.assertRaises(ValueError): + psdf.groupby("key").quantile(q=1.1) + with self.assertRaises(ValueError): + psdf.groupby("key").quantile(q=-0.1) + with self.assertRaises(ValueError): + psdf.groupby("key").quantile(q=2) + with self.assertRaises(ValueError): + psdf.groupby("key").quantile(q=np.nan) + # raise TypeError when q type mismatch + with self.assertRaises(TypeError): + psdf.groupby("key").quantile(q="0.1") + # raise NotImplementedError when q is list like type + with self.assertRaises(NotImplementedError): + psdf.groupby("key").quantile(q=(0.1, 0.5)) + with self.assertRaises(NotImplementedError): + psdf.groupby("key").quantile(q=[0.1, 0.5]) + + def test_min(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.min()) + self._test_stat_func(lambda groupby_obj: groupby_obj.min(min_count=2)) + self._test_stat_func(lambda groupby_obj: groupby_obj.min(numeric_only=None)) + self._test_stat_func(lambda groupby_obj: groupby_obj.min(numeric_only=True)) + self._test_stat_func(lambda groupby_obj: groupby_obj.min(numeric_only=True, min_count=2)) + + def test_max(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.max()) + self._test_stat_func(lambda groupby_obj: groupby_obj.max(min_count=2)) + self._test_stat_func(lambda groupby_obj: groupby_obj.max(numeric_only=None)) + self._test_stat_func(lambda groupby_obj: groupby_obj.max(numeric_only=True)) + self._test_stat_func(lambda groupby_obj: groupby_obj.max(numeric_only=True, min_count=2)) + + def test_sum(self): + pdf = pd.DataFrame( + { + "A": ["a", "a", "b", "a"], + "B": [1, 2, 1, 2], + "C": [-1.5, np.nan, -3.2, 0.1], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.groupby("A").sum().sort_index(), psdf.groupby("A").sum().sort_index()) + self.assert_eq( + pdf.groupby("A").sum(min_count=2).sort_index(), + psdf.groupby("A").sum(min_count=2).sort_index(), + ) + self.assert_eq( + pdf.groupby("A").sum(min_count=3).sort_index(), + psdf.groupby("A").sum(min_count=3).sort_index(), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43553): Enable GroupByTests.test_mad for pandas 2.0.0.", + ) + def test_mad(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.mad()) + + def test_first(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.first()) + self._test_stat_func(lambda groupby_obj: groupby_obj.first(numeric_only=None)) + self._test_stat_func(lambda groupby_obj: groupby_obj.first(numeric_only=True)) + + pdf = pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [-1.5, np.nan, -3.2, 0.1], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.groupby("A").first().sort_index(), psdf.groupby("A").first().sort_index() + ) + self.assert_eq( + pdf.groupby("A").first(min_count=1).sort_index(), + psdf.groupby("A").first(min_count=1).sort_index(), + ) + self.assert_eq( + pdf.groupby("A").first(min_count=2).sort_index(), + psdf.groupby("A").first(min_count=2).sort_index(), + ) + + def test_last(self): + self._test_stat_func(lambda groupby_obj: groupby_obj.last()) + self._test_stat_func(lambda groupby_obj: groupby_obj.last(numeric_only=None)) + self._test_stat_func(lambda groupby_obj: groupby_obj.last(numeric_only=True)) + + pdf = pd.DataFrame( + { + "A": [1, 2, 1, 2], + "B": [-1.5, np.nan, -3.2, 0.1], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.groupby("A").last().sort_index(), psdf.groupby("A").last().sort_index()) + self.assert_eq( + pdf.groupby("A").last(min_count=1).sort_index(), + psdf.groupby("A").last(min_count=1).sort_index(), + ) + self.assert_eq( + pdf.groupby("A").last(min_count=2).sort_index(), + psdf.groupby("A").last(min_count=2).sort_index(), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43552): Enable GroupByTests.test_nth for pandas 2.0.0.", + ) + def test_nth(self): + for n in [0, 1, 2, 128, -1, -2, -128]: + self._test_stat_func(lambda groupby_obj: groupby_obj.nth(n)) + + with self.assertRaisesRegex(NotImplementedError, "slice or list"): + self.psdf.groupby("B").nth(slice(0, 2)) + with self.assertRaisesRegex(NotImplementedError, "slice or list"): + self.psdf.groupby("B").nth([0, 1, -1]) + with self.assertRaisesRegex(TypeError, "Invalid index"): + self.psdf.groupby("B").nth("x") + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43551): Enable GroupByTests.test_prod for pandas 2.0.0.", + ) + def test_prod(self): + pdf = pd.DataFrame( + { + "A": [1, 2, 1, 2, 1], + "B": [3.1, 4.1, 4.1, 3.1, 0.1], + "C": ["a", "b", "b", "a", "c"], + "D": [True, False, False, True, False], + "E": [-1, -2, 3, -4, -2], + "F": [-1.5, np.nan, -3.2, 0.1, 0], + "G": [np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + + for n in [0, 1, 2, 128, -1, -2, -128]: + self._test_stat_func( + lambda groupby_obj: groupby_obj.prod(min_count=n), check_exact=False + ) + self._test_stat_func( + lambda groupby_obj: groupby_obj.prod(numeric_only=None, min_count=n), + check_exact=False, + ) + self._test_stat_func( + lambda groupby_obj: groupby_obj.prod(numeric_only=True, min_count=n), + check_exact=False, + ) + self.assert_eq( + pdf.groupby("A").prod(min_count=n).sort_index(), + psdf.groupby("A").prod(min_count=n).sort_index(), + almost=True, + ) + + def test_median(self): + psdf = ps.DataFrame( + { + "a": [1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0], + "b": [2.0, 3.0, 1.0, 4.0, 6.0, 9.0, 8.0, 10.0, 7.0, 5.0], + "c": [3.0, 5.0, 2.0, 5.0, 1.0, 2.0, 6.0, 4.0, 3.0, 6.0], + }, + columns=["a", "b", "c"], + index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6], + ) + # DataFrame + expected_result = ps.DataFrame( + {"b": [2.0, 8.0, 7.0], "c": [3.0, 2.0, 4.0]}, index=pd.Index([1.0, 2.0, 3.0], name="a") + ) + self.assert_eq(expected_result, psdf.groupby("a").median().sort_index()) + # Series + expected_result = ps.Series( + [2.0, 8.0, 7.0], name="b", index=pd.Index([1.0, 2.0, 3.0], name="a") + ) + self.assert_eq(expected_result, psdf.groupby("a")["b"].median().sort_index()) + + with self.assertRaisesRegex(TypeError, "accuracy must be an integer; however"): + psdf.groupby("a").median(accuracy="a") + + def test_ddof(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, + }, + index=np.random.rand(10 * 3), + ) + psdf = ps.from_pandas(pdf) + + for ddof in [-1, 0, 1, 2, 3]: + # std + self.assert_eq( + pdf.groupby("a").std(ddof=ddof).sort_index(), + psdf.groupby("a").std(ddof=ddof).sort_index(), + check_exact=False, + ) + self.assert_eq( + pdf.groupby("a")["b"].std(ddof=ddof).sort_index(), + psdf.groupby("a")["b"].std(ddof=ddof).sort_index(), + check_exact=False, + ) + # var + self.assert_eq( + pdf.groupby("a").var(ddof=ddof).sort_index(), + psdf.groupby("a").var(ddof=ddof).sort_index(), + check_exact=False, + ) + self.assert_eq( + pdf.groupby("a")["b"].var(ddof=ddof).sort_index(), + psdf.groupby("a")["b"].var(ddof=ddof).sort_index(), + check_exact=False, + ) + # sem + self.assert_eq( + pdf.groupby("a").sem(ddof=ddof).sort_index(), + psdf.groupby("a").sem(ddof=ddof).sort_index(), + check_exact=False, + ) + self.assert_eq( + pdf.groupby("a")["b"].sem(ddof=ddof).sort_index(), + psdf.groupby("a")["b"].sem(ddof=ddof).sort_index(), + check_exact=False, + ) + + +class GroupbyStatTests(GroupbyStatMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.groupby.test_stat import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_align.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_align.py new file mode 100644 index 0000000..56fde9b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_align.py @@ -0,0 +1,100 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameAlignMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_align(self): + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30]) + psdf1 = ps.from_pandas(pdf1) + + for join in ["outer", "inner", "left", "right"]: + for axis in [None, 0, 1]: + psdf_l, psdf_r = psdf1.align(psdf1[["b"]], join=join, axis=axis) + pdf_l, pdf_r = pdf1.align(pdf1[["b"]], join=join, axis=axis) + self.assert_eq(psdf_l, pdf_l) + self.assert_eq(psdf_r, pdf_r) + + psdf_l, psdf_r = psdf1[["a"]].align(psdf1[["b", "a"]], join=join, axis=axis) + pdf_l, pdf_r = pdf1[["a"]].align(pdf1[["b", "a"]], join=join, axis=axis) + self.assert_eq(psdf_l, pdf_l) + self.assert_eq(psdf_r, pdf_r) + + psdf_l, psdf_r = psdf1[["b", "a"]].align(psdf1[["a"]], join=join, axis=axis) + pdf_l, pdf_r = pdf1[["b", "a"]].align(pdf1[["a"]], join=join, axis=axis) + self.assert_eq(psdf_l, pdf_l) + self.assert_eq(psdf_r, pdf_r) + + psdf_l, psdf_r = psdf1.align(psdf1["b"], axis=0) + pdf_l, pdf_r = pdf1.align(pdf1["b"], axis=0) + self.assert_eq(psdf_l, pdf_l) + self.assert_eq(psdf_r, pdf_r) + + psdf_l, psser_b = psdf1[["a"]].align(psdf1["b"], axis=0) + pdf_l, pser_b = pdf1[["a"]].align(pdf1["b"], axis=0) + self.assert_eq(psdf_l, pdf_l) + self.assert_eq(psser_b, pser_b) + + self.assertRaises(ValueError, lambda: psdf1.align(psdf1, join="unknown")) + self.assertRaises(ValueError, lambda: psdf1.align(psdf1["b"])) + self.assertRaises(TypeError, lambda: psdf1.align(["b"])) + self.assertRaises(NotImplementedError, lambda: psdf1.align(psdf1["b"], axis=1)) + + pdf2 = pd.DataFrame({"a": [4, 5, 6], "d": ["d", "e", "f"]}, index=[10, 11, 12]) + psdf2 = ps.from_pandas(pdf2) + + for join in ["outer", "inner", "left", "right"]: + psdf_l, psdf_r = psdf1.align(psdf2, join=join, axis=1) + pdf_l, pdf_r = pdf1.align(pdf2, join=join, axis=1) + self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index()) + self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index()) + + +class FrameAlignTests(FrameAlignMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_align import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_base.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_base.py new file mode 100644 index 0000000..6cb7c58 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_base.py @@ -0,0 +1,2743 @@ +# +# 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 inspect +import unittest +from distutils.version import LooseVersion +from datetime import datetime, timedelta + +import numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.indexes import ( + MissingPandasLikeDatetimeIndex, + MissingPandasLikeIndex, + MissingPandasLikeMultiIndex, + MissingPandasLikeTimedeltaIndex, +) +from pyspark.testing.pandasutils import ComparisonTestBase, TestUtils, SPARK_CONF_ARROW_ENABLED + + +class IndexesTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43606): Enable IndexesTests.test_index_basic for pandas 2.0.0.", + ) + def test_index_basic(self): + for pdf in [ + pd.DataFrame(np.random.randn(10, 5), index=np.random.randint(100, size=10)), + pd.DataFrame( + np.random.randn(10, 5), index=np.random.randint(100, size=10).astype(np.int32) + ), + pd.DataFrame(np.random.randn(10, 5), index=np.random.randn(10)), + pd.DataFrame(np.random.randn(10, 5), index=np.random.randn(10).astype(np.float32)), + pd.DataFrame(np.random.randn(10, 5), index=list("abcdefghij")), + pd.DataFrame( + np.random.randn(10, 5), index=pd.date_range("2011-01-01", freq="D", periods=10) + ), + pd.DataFrame(np.random.randn(10, 5), index=pd.Categorical(list("abcdefghij"))), + pd.DataFrame(np.random.randn(10, 5), columns=list("abcde")).set_index(["a", "b"]), + ]: + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.index, pdf.index) + # Int64Index is removed from pandas 2.0.0, so we should compare the dtype itself. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(psdf.index.dtype, pdf.index.dtype) + else: + self.assert_eq(type(psdf.index).__name__, type(pdf.index).__name__) + + self.assert_eq(ps.Index([])._summary(), "Index: 0 entries") + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + with self.assertRaisesRegexp(ValueError, "The truth value of a Index is ambiguous."): + bool(ps.Index([1])) + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.Index([1, 2, 3], name=[(1, 2, 3)]) + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.Index([1.0, 2.0, 3.0], name=[(1, 2, 3)]) + else: + with self.assertRaisesRegexp( + ValueError, "The truth value of a Int64Index is ambiguous." + ): + bool(ps.Index([1])) + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.Int64Index([1, 2, 3], name=[(1, 2, 3)]) + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.Float64Index([1.0, 2.0, 3.0], name=[(1, 2, 3)]) + + def test_index_from_series(self): + pser = pd.Series([1, 2, 3], name="a", index=[10, 20, 30]) + psser = ps.from_pandas(pser) + + self.assert_eq(ps.Index(psser), pd.Index(pser)) + self.assert_eq(ps.Index(psser, dtype="float"), pd.Index(pser, dtype="float")) + self.assert_eq(ps.Index(psser, name="x"), pd.Index(pser, name="x")) + + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(ps.Index(psser, dtype="int64"), pd.Index(pser, dtype="int64")) + self.assert_eq(ps.Index(psser, dtype="float64"), pd.Index(pser, dtype="float64")) + elif LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(ps.Int64Index(psser), pd.Int64Index(pser)) + self.assert_eq(ps.Float64Index(psser), pd.Float64Index(pser)) + else: + self.assert_eq(ps.Int64Index(psser), pd.Int64Index(pser).rename("a")) + self.assert_eq(ps.Float64Index(psser), pd.Float64Index(pser).rename("a")) + + pser = pd.Series([datetime(2021, 3, 1), datetime(2021, 3, 2)], name="x", index=[10, 20]) + psser = ps.from_pandas(pser) + + self.assert_eq(ps.Index(psser), pd.Index(pser)) + self.assert_eq(ps.DatetimeIndex(psser), pd.DatetimeIndex(pser)) + + def test_index_from_index(self): + pidx = pd.Index([1, 2, 3], name="a") + psidx = ps.from_pandas(pidx) + + self.assert_eq(ps.Index(psidx), pd.Index(pidx)) + self.assert_eq(ps.Index(psidx, dtype="float"), pd.Index(pidx, dtype="float")) + self.assert_eq(ps.Index(psidx, name="x"), pd.Index(pidx, name="x")) + self.assert_eq(ps.Index(psidx, copy=True), pd.Index(pidx, copy=True)) + + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(ps.Index(psidx, dtype="int64"), pd.Index(pidx, dtype="int64")) + self.assert_eq(ps.Index(psidx, dtype="float64"), pd.Index(pidx, dtype="float64")) + else: + self.assert_eq(ps.Int64Index(psidx), pd.Int64Index(pidx)) + self.assert_eq(ps.Float64Index(psidx), pd.Float64Index(pidx)) + + pidx = pd.DatetimeIndex(["2021-03-01", "2021-03-02"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(ps.Index(psidx), pd.Index(pidx)) + self.assert_eq(ps.DatetimeIndex(psidx), pd.DatetimeIndex(pidx)) + + def test_index_getattr(self): + psidx = self.psdf.index + item = "databricks" + + expected_error_message = "'.*Index' object has no attribute '{}'".format(item) + with self.assertRaisesRegex(AttributeError, expected_error_message): + psidx.__getattr__(item) + with self.assertRaisesRegex(AttributeError, expected_error_message): + ps.from_pandas(pd.date_range("2011-01-01", freq="D", periods=10)).__getattr__(item) + + def test_multi_index_getattr(self): + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame(np.random.randn(4, 5), idx) + psdf = ps.from_pandas(pdf) + psidx = psdf.index + item = "databricks" + + expected_error_message = "'MultiIndex' object has no attribute '{}'".format(item) + with self.assertRaisesRegex(AttributeError, expected_error_message): + psidx.__getattr__(item) + + def test_to_series(self): + pidx = self.pdf.index + psidx = self.psdf.index + + self.assert_eq(psidx.to_series(), pidx.to_series()) + self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a")) + + # With name + pidx.name = "Koalas" + psidx.name = "Koalas" + self.assert_eq(psidx.to_series(), pidx.to_series()) + self.assert_eq(psidx.to_series(name=("x", "a")), pidx.to_series(name=("x", "a"))) + + # With tupled name + pidx.name = ("x", "a") + psidx.name = ("x", "a") + self.assert_eq(psidx.to_series(), pidx.to_series()) + self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a")) + + self.assert_eq((psidx + 1).to_series(), (pidx + 1).to_series()) + + pidx = self.pdf.set_index("b", append=True).index + psidx = self.psdf.set_index("b", append=True).index + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + self.assert_eq(psidx.to_series(), pidx.to_series()) + self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a")) + + expected_error_message = "Series.name must be a hashable type" + with self.assertRaisesRegex(TypeError, expected_error_message): + psidx.to_series(name=["x", "a"]) + + def test_to_frame(self): + pidx = self.pdf.index + psidx = self.psdf.index + + self.assert_eq(psidx.to_frame(), pidx.to_frame()) + self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False)) + + pidx.name = "a" + psidx.name = "a" + + self.assert_eq(psidx.to_frame(), pidx.to_frame()) + self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False)) + + self.assert_eq(psidx.to_frame(name="x"), pidx.to_frame(name="x")) + self.assert_eq(psidx.to_frame(index=False, name="x"), pidx.to_frame(index=False, name="x")) + + self.assertRaises(TypeError, lambda: psidx.to_frame(name=["x"])) + + # non-string name + self.assert_eq(psidx.to_frame(name=10), pidx.to_frame(name=10)) + self.assert_eq(psidx.to_frame(name=("x", 10)), pidx.to_frame(name=("x", 10))) + + pidx = self.pdf.set_index("b", append=True).index + psidx = self.psdf.set_index("b", append=True).index + + self.assert_eq(psidx.to_frame(), pidx.to_frame()) + self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False)) + + self.assert_eq(psidx.to_frame(name=["x", "y"]), pidx.to_frame(name=["x", "y"])) + self.assert_eq(psidx.to_frame(name=("x", "y")), pidx.to_frame(name=("x", "y"))) + self.assert_eq( + psidx.to_frame(index=False, name=["x", "y"]), + pidx.to_frame(index=False, name=["x", "y"]), + ) + + self.assertRaises(TypeError, lambda: psidx.to_frame(name="x")) + self.assertRaises(ValueError, lambda: psidx.to_frame(name=["x"])) + + # non-string names + self.assert_eq(psidx.to_frame(name=[10, 20]), pidx.to_frame(name=[10, 20])) + self.assert_eq(psidx.to_frame(name=("x", 10)), pidx.to_frame(name=("x", 10))) + if LooseVersion(pd.__version__) < LooseVersion("1.5.0"): + self.assert_eq( + psidx.to_frame(name=[("x", 10), ("y", 20)]), + pidx.to_frame(name=[("x", 10), ("y", 20)]), + ) + else: + # Since pandas 1.5.0, the result is changed as below: + # (x, 10) (y, 20) + # b + # 0 4 0 4 + # 1 5 1 5 + # 3 6 3 6 + # 5 3 5 3 + # 6 2 6 2 + # 8 1 8 1 + # 9 0 9 0 + # 0 9 0 + # 0 9 0 + # + # The columns should be `Index([('x', 20), ('y', 20)], dtype='object')`, + # but pandas API on Spark doesn't support such a way for creating Index. + # So, we currently cannot follow the behavior of pandas. + expected_result = ps.DataFrame( + {("x", 10): [0, 1, 3, 5, 6, 8, 9, 9, 9], ("y", 20): [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=ps.MultiIndex.from_tuples( + [(0, 4), (1, 5), (3, 6), (5, 3), (6, 2), (8, 1), (9, 0), (9, 0), (9, 0)], + names=[None, "b"], + ), + ) + self.assert_eq(psidx.to_frame(name=[("x", 10), ("y", 20)]), expected_result) + + def test_index_names(self): + psdf = self.psdf + self.assertIsNone(psdf.index.name) + + idx = pd.Index([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], name="x") + pdf = pd.DataFrame(np.random.randn(10, 5), index=idx, columns=list("abcde")) + psdf = ps.from_pandas(pdf) + + pser = pdf.a + psser = psdf.a + + self.assertEqual(psdf.index.name, pdf.index.name) + self.assertEqual(psdf.index.names, pdf.index.names) + + pidx = pdf.index + psidx = psdf.index + pidx.name = "renamed" + psidx.name = "renamed" + self.assertEqual(psidx.name, pidx.name) + self.assertEqual(psidx.names, pidx.names) + self.assert_eq(psidx, pidx) + self.assertEqual(psdf.index.name, pdf.index.name) + self.assertEqual(psdf.index.names, pdf.index.names) + self.assertEqual(psser.index.names, pser.index.names) + + pidx.name = None + psidx.name = None + self.assertEqual(psidx.name, pidx.name) + self.assertEqual(psidx.names, pidx.names) + self.assert_eq(psidx, pidx) + self.assertEqual(psdf.index.name, pdf.index.name) + self.assertEqual(psdf.index.names, pdf.index.names) + self.assertEqual(psser.index.names, pser.index.names) + + with self.assertRaisesRegex(ValueError, "Names must be a list-like"): + psidx.names = "hi" + + expected_error_message = "Length of new names must be {}, got {}".format( + psdf._internal.index_level, len(["0", "1"]) + ) + with self.assertRaisesRegex(ValueError, expected_error_message): + psidx.names = ["0", "1"] + + expected_error_message = "Index.name must be a hashable type" + with self.assertRaisesRegex(TypeError, expected_error_message): + ps.Index([1, 2, 3], name=["0", "1"]) + with self.assertRaisesRegex(TypeError, expected_error_message): + psidx.name = ["renamed"] + with self.assertRaisesRegex(TypeError, expected_error_message): + psidx.name = ["0", "1"] + # Specifying `names` when creating Index is no longer supported from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + pass + else: + with self.assertRaisesRegex(TypeError, expected_error_message): + ps.Index([(1, 2), (3, 4)], names=["a", ["b"]]) + + def test_multi_index_names(self): + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame(np.random.randn(4, 5), idx) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.names, pdf.index.names) + + pidx = pdf.index + psidx = psdf.index + pidx.names = ["renamed_number", "renamed_color"] + psidx.names = ["renamed_number", "renamed_color"] + self.assertEqual(psidx.names, pidx.names) + + pidx.names = ["renamed_number", None] + psidx.names = ["renamed_number", None] + self.assertEqual(psidx.names, pidx.names) + self.assert_eq(psidx, pidx) + + with self.assertRaises(PandasNotImplementedError): + psidx.name + with self.assertRaises(PandasNotImplementedError): + psidx.name = "renamed" + + def test_index_rename(self): + pdf = pd.DataFrame( + np.random.randn(10, 5), index=pd.Index([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], name="x") + ) + psdf = ps.from_pandas(pdf) + + pidx = pdf.index + psidx = psdf.index + + self.assert_eq(psidx.rename("y"), pidx.rename("y")) + self.assert_eq(psdf.index.names, pdf.index.names) + + # non-string names + self.assert_eq(psidx.rename(0), pidx.rename(0)) + self.assert_eq(psidx.rename(("y", 0)), pidx.rename(("y", 0))) + + psidx.rename("z", inplace=True) + pidx.rename("z", inplace=True) + + self.assert_eq(psidx, pidx) + self.assert_eq(psdf.index.names, pdf.index.names) + + self.assert_eq(psidx.rename(None), pidx.rename(None)) + self.assert_eq(psdf.index.names, pdf.index.names) + + self.assertRaises(TypeError, lambda: psidx.rename(["x", "y"])) + + def test_multi_index_rename(self): + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame(np.random.randn(4, 5), idx) + psdf = ps.from_pandas(pdf) + + pmidx = pdf.index + psmidx = psdf.index + + self.assert_eq(psmidx.rename(["n", "c"]), pmidx.rename(["n", "c"])) + self.assert_eq(psdf.index.names, pdf.index.names) + + # non-string names + self.assert_eq(psmidx.rename([0, 1]), pmidx.rename([0, 1])) + self.assert_eq( + psmidx.rename([("x", "a"), ("y", "b")]), pmidx.rename([("x", "a"), ("y", "b")]) + ) + + psmidx.rename(["num", "col"], inplace=True) + pmidx.rename(["num", "col"], inplace=True) + + self.assert_eq(psmidx, pmidx) + self.assert_eq(psdf.index.names, pdf.index.names) + + self.assert_eq(psmidx.rename([None, None]), pmidx.rename([None, None])) + self.assert_eq(psdf.index.names, pdf.index.names) + + self.assertRaises(TypeError, lambda: psmidx.rename("number")) + self.assertRaises(TypeError, lambda: psmidx.rename(None)) + self.assertRaises(ValueError, lambda: psmidx.rename(["number"])) + + def test_multi_index_levshape(self): + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2)]) + psidx = ps.from_pandas(pidx) + self.assertEqual(pidx.levshape, psidx.levshape) + + def test_index_unique(self): + psidx = self.psdf.index + + # here the output is different than pandas in terms of order + expected = [0, 1, 3, 5, 6, 8, 9] + + self.assert_eq(expected, sorted(psidx.unique()._to_pandas())) + self.assert_eq(expected, sorted(psidx.unique(level=0)._to_pandas())) + + expected = [1, 2, 4, 6, 7, 9, 10] + self.assert_eq(expected, sorted((psidx + 1).unique()._to_pandas())) + + with self.assertRaisesRegex(IndexError, "Too many levels*"): + psidx.unique(level=1) + + with self.assertRaisesRegex(KeyError, "Requested level (hi)*"): + psidx.unique(level="hi") + + def test_multi_index_copy(self): + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pdf = pd.DataFrame(np.random.randn(4, 5), idx) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.index.copy(), pdf.index.copy()) + + def test_drop_duplicates(self): + pidx = pd.Index([4, 2, 4, 1, 4, 3]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.drop_duplicates(), pidx.drop_duplicates()) + self.assert_eq((psidx + 1).drop_duplicates(), (pidx + 1).drop_duplicates()) + + self.assert_eq(psidx.drop_duplicates(keep="first"), pidx.drop_duplicates(keep="first")) + self.assert_eq(psidx.drop_duplicates(keep="last"), pidx.drop_duplicates(keep="last")) + self.assert_eq(psidx.drop_duplicates(keep=False), pidx.drop_duplicates(keep=False)) + + arrays = [[1, 2, 3, 1, 2], ["red", "blue", "black", "red", "blue"]] + pmidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.drop_duplicates(), pmidx.drop_duplicates()) + self.assert_eq(psmidx.drop_duplicates(keep="first"), pmidx.drop_duplicates(keep="first")) + self.assert_eq(psmidx.drop_duplicates(keep="last"), pmidx.drop_duplicates(keep="last")) + self.assert_eq(psmidx.drop_duplicates(keep=False), pmidx.drop_duplicates(keep=False)) + + def test_dropna(self): + pidx = pd.Index([np.nan, 2, 4, 1, None, 3]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.dropna(), pidx.dropna()) + self.assert_eq((psidx + 1).dropna(), (pidx + 1).dropna()) + + self.assert_eq(psidx.dropna(how="any"), pidx.dropna(how="any")) + self.assert_eq(psidx.dropna(how="all"), pidx.dropna(how="all")) + + pmidx = pd.MultiIndex.from_tuples( + [(np.nan, 1.0), (2.0, 2.0), (np.nan, None), (3.0, np.nan)] + ) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.dropna(), pmidx.dropna()) + self.assert_eq(psmidx.dropna(how="any"), pmidx.dropna(how="any")) + self.assert_eq(psmidx.dropna(how="all"), pmidx.dropna(how="all")) + + invalid_how = "none" + with self.assertRaisesRegex(ValueError, "invalid how option: %s" % invalid_how): + psmidx.dropna(invalid_how) + + def test_index_symmetric_difference(self): + pidx1 = pd.Index([1, 2, 3, 4]) + pidx2 = pd.Index([2, 3, 4, 5]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + + self.assert_eq( + psidx1.symmetric_difference(psidx2).sort_values(), + pidx1.symmetric_difference(pidx2).sort_values(), + ) + self.assert_eq( + (psidx1 + 1).symmetric_difference(psidx2).sort_values(), + (pidx1 + 1).symmetric_difference(pidx2).sort_values(), + ) + # No longer supported from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + (psidx1 ^ psidx2).sort_values(), + ps.Index([1, 5], dtype="int64"), + ) + else: + self.assert_eq( + (psidx1 ^ psidx2).sort_values(), + (pidx1 ^ pidx2).sort_values(), + ) + self.assert_eq( + psidx1.symmetric_difference(psidx2, result_name="result").sort_values(), + pidx1.symmetric_difference(pidx2, result_name="result").sort_values(), + ) + + pmidx1 = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 0, 0, 0, 1, 2, 0, 1, 2]], + ) + pmidx2 = pd.MultiIndex( + [["koalas", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 0, 0, 0, 1, 2, 0, 1, 2]], + ) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + + self.assert_eq( + psmidx1.symmetric_difference(psmidx2).sort_values(), + pmidx1.symmetric_difference(pmidx2).sort_values(), + ) + + # Pandas has a bug that raise TypeError when setting `result_name` for MultiIndex. + pandas_result = pmidx1.symmetric_difference(pmidx2) + pandas_result.names = ["a", "b"] + self.assert_eq( + psmidx1.symmetric_difference(psmidx2, result_name=["a", "b"]).sort_values(), + pandas_result, + ) + + # Pandas sort the result by default, so doesn't provide the `True` for sort. + self.assert_eq( + psmidx1.symmetric_difference(psmidx2, sort=True), + pmidx1.symmetric_difference(pmidx2), + ) + + idx = ps.Index(["a", "b", "c"]) + midx = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + + with self.assertRaisesRegex(NotImplementedError, "Doesn't support*"): + idx.symmetric_difference(midx) + + def test_multi_index_symmetric_difference(self): + idx = ps.Index(["a", "b", "c"]) + midx = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + midx_ = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + + self.assert_eq( + midx.symmetric_difference(midx_), + midx._to_pandas().symmetric_difference(midx_._to_pandas()), + ) + + with self.assertRaisesRegex(NotImplementedError, "Doesn't support*"): + midx.symmetric_difference(idx) + + def test_missing(self): + psdf = ps.DataFrame( + { + "a": [1, 2, 3], + "b": [4, 5, 6], + "c": pd.date_range("2011-01-01", freq="D", periods=3), + "d": pd.Categorical(["a", "b", "c"]), + "e": [timedelta(1), timedelta(2), timedelta(3)], + } + ) + + # Index functions + missing_functions = inspect.getmembers(MissingPandasLikeIndex, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("a").index, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index("a").index, name)() + + # MultiIndex functions + missing_functions = inspect.getmembers(MissingPandasLikeMultiIndex, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index(["a", "b"]).index, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index(["a", "b"]).index, name)() + + # DatetimeIndex functions + missing_functions = inspect.getmembers(MissingPandasLikeDatetimeIndex, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("c").index, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index("c").index, name)() + + # TimedeltaIndex functions + missing_functions = inspect.getmembers(MissingPandasLikeTimedeltaIndex, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("e").index, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index("e").index, name)() + + # Index properties + missing_properties = inspect.getmembers( + MissingPandasLikeIndex, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("a").index, name) + + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index("a").index, name) + + # MultiIndex properties + missing_properties = inspect.getmembers( + MissingPandasLikeMultiIndex, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index(["a", "b"]).index, name) + + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Index.*{}.*is deprecated".format(name) + ): + getattr(psdf.set_index(["a", "b"]).index, name) + + # DatetimeIndex properties + missing_properties = inspect.getmembers( + MissingPandasLikeDatetimeIndex, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("c").index, name) + + # TimedeltaIndex properties + missing_properties = inspect.getmembers( + MissingPandasLikeDatetimeIndex, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Index.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.set_index("c").index, name) + + def test_index_has_duplicates(self): + indexes = [("a", "b", "c"), ("a", "a", "c"), (1, 3, 3), (1, 2, 3)] + names = [None, "ks", "ks", None] + has_dup = [False, True, True, False] + + for idx, name, expected in zip(indexes, names, has_dup): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=pd.Index(idx, name=name)) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.has_duplicates, expected) + + def test_multiindex_has_duplicates(self): + indexes = [ + [list("abc"), list("edf")], + [list("aac"), list("edf")], + [list("aac"), list("eef")], + [[1, 4, 4], [4, 6, 6]], + ] + has_dup = [False, False, True, True] + + for idx, expected in zip(indexes, has_dup): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=idx) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.has_duplicates, expected) + + def test_multi_index_not_supported(self): + psdf = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + + with self.assertRaisesRegex(TypeError, "cannot perform any with this index type"): + psdf.set_index(["a", "b"]).index.any() + + with self.assertRaisesRegex(TypeError, "cannot perform all with this index type"): + psdf.set_index(["a", "b"]).index.all() + + def test_index_nlevels(self): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=pd.Index(["a", "b", "c"])) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.nlevels, 1) + + def test_multiindex_nlevel(self): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=[list("abc"), list("def")]) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.nlevels, 2) + + def test_multiindex_from_arrays(self): + arrays = [["a", "a", "b", "b"], ["red", "blue", "red", "blue"]] + pidx = pd.MultiIndex.from_arrays(arrays) + psidx = ps.MultiIndex.from_arrays(arrays) + + self.assert_eq(pidx, psidx) + + def test_multiindex_swaplevel(self): + pidx = pd.MultiIndex.from_arrays([["a", "b"], [1, 2]]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.swaplevel(0, 1), psidx.swaplevel(0, 1)) + + pidx = pd.MultiIndex.from_arrays([["a", "b"], [1, 2]], names=["word", "number"]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.swaplevel(0, 1), psidx.swaplevel(0, 1)) + + pidx = pd.MultiIndex.from_arrays([["a", "b"], [1, 2]], names=["word", None]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.swaplevel(-2, -1), psidx.swaplevel(-2, -1)) + self.assert_eq(pidx.swaplevel(0, 1), psidx.swaplevel(0, 1)) + self.assert_eq(pidx.swaplevel("word", 1), psidx.swaplevel("word", 1)) + + with self.assertRaisesRegex(IndexError, "Too many levels: Index"): + psidx.swaplevel(-3, "word") + with self.assertRaisesRegex(IndexError, "Too many levels: Index"): + psidx.swaplevel(0, 2) + with self.assertRaisesRegex(IndexError, "Too many levels: Index"): + psidx.swaplevel(0, -3) + with self.assertRaisesRegex(KeyError, "Level work not found"): + psidx.swaplevel(0, "work") + + def test_multiindex_droplevel(self): + pidx = pd.MultiIndex.from_tuples( + [("a", "x", 1), ("b", "y", 2)], names=["level1", "level2", "level3"] + ) + psidx = ps.from_pandas(pidx) + with self.assertRaisesRegex(IndexError, "Too many levels: Index has only 3 levels, not 5"): + psidx.droplevel(4) + + with self.assertRaisesRegex(KeyError, "Level level4 not found"): + psidx.droplevel("level4") + + with self.assertRaisesRegex(KeyError, "Level.*level3.*level4.*not found"): + psidx.droplevel([("level3", "level4")]) + + with self.assertRaisesRegex( + ValueError, + "Cannot remove 4 levels from an index with 3 levels: at least one " + "level must be left.", + ): + psidx.droplevel([0, 0, 1, 2]) + + with self.assertRaisesRegex( + ValueError, + "Cannot remove 3 levels from an index with 3 levels: at least one " + "level must be left.", + ): + psidx.droplevel([0, 1, 2]) + + self.assert_eq(pidx.droplevel(0), psidx.droplevel(0)) + self.assert_eq(pidx.droplevel([0, 1]), psidx.droplevel([0, 1])) + self.assert_eq(pidx.droplevel((0, 1)), psidx.droplevel((0, 1))) + self.assert_eq(pidx.droplevel([0, "level2"]), psidx.droplevel([0, "level2"])) + self.assert_eq(pidx.droplevel((0, "level2")), psidx.droplevel((0, "level2"))) + + # non-string names + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2)], names=[1.0, 2.0, 3.0]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.droplevel(1.0), psidx.droplevel(1.0)) + self.assert_eq(pidx.droplevel([0, 2.0]), psidx.droplevel([0, 2.0])) + + def test_index_fillna(self): + pidx = pd.Index([1, 2, None]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.fillna(0), psidx.fillna(0), almost=True) + self.assert_eq(pidx.rename("name").fillna(0), psidx.rename("name").fillna(0), almost=True) + + with self.assertRaisesRegex(TypeError, "Unsupported type list"): + psidx.fillna([1, 2]) + + def test_index_drop(self): + pidx = pd.Index([1, 2, 3]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.drop(1), psidx.drop(1)) + self.assert_eq(pidx.drop([1, 2]), psidx.drop([1, 2])) + self.assert_eq((pidx + 1).drop([2, 3]), (psidx + 1).drop([2, 3])) + + def test_multiindex_drop(self): + pidx = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z")], names=["level1", "level2"] + ) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.drop("a"), psidx.drop("a")) + self.assert_eq(pidx.drop(["a", "b"]), psidx.drop(["a", "b"])) + self.assert_eq(pidx.drop(["x", "y"], level=1), psidx.drop(["x", "y"], level=1)) + self.assert_eq( + pidx.drop(["x", "y"], level="level2"), psidx.drop(["x", "y"], level="level2") + ) + + pidx.names = ["lv1", "lv2"] + psidx.names = ["lv1", "lv2"] + self.assert_eq(pidx.drop(["x", "y"], level="lv2"), psidx.drop(["x", "y"], level="lv2")) + + self.assertRaises(IndexError, lambda: psidx.drop(["a", "b"], level=2)) + self.assertRaises(KeyError, lambda: psidx.drop(["a", "b"], level="level")) + + psidx.names = ["lv", "lv"] + self.assertRaises(ValueError, lambda: psidx.drop(["x", "y"], level="lv")) + + def _test_sort_values(self, pidx, psidx): + self.assert_eq(pidx.sort_values(), psidx.sort_values()) + # Parameter ascending + self.assert_eq(pidx.sort_values(ascending=False), psidx.sort_values(ascending=False)) + # Parameter return_indexer + p_sorted, p_indexer = pidx.sort_values(return_indexer=True) + ps_sorted, ps_indexer = psidx.sort_values(return_indexer=True) + self.assert_eq(p_sorted, ps_sorted) + self.assert_eq(p_indexer, ps_indexer.to_list()) + self.assert_eq( + pidx.sort_values(return_indexer=False), psidx.sort_values(return_indexer=False) + ) + # Parameter return_indexer and ascending + p_sorted, p_indexer = pidx.sort_values(return_indexer=True, ascending=False) + ps_sorted, ps_indexer = psidx.sort_values(return_indexer=True, ascending=False) + self.assert_eq(p_sorted, ps_sorted) + self.assert_eq(p_indexer, ps_indexer.to_list()) + self.assert_eq( + pidx.sort_values(return_indexer=False, ascending=False), + psidx.sort_values(return_indexer=False, ascending=False), + ) + + def test_sort_values(self): + pidx = pd.Index([-10, -100, 200, 100]) + psidx = ps.from_pandas(pidx) + + self._test_sort_values(pidx, psidx) + + pidx.name = "koalas" + psidx.name = "koalas" + + self._test_sort_values(pidx, psidx) + + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psidx = ps.from_pandas(pidx) + + pidx.names = ["hello", "koalas", "goodbye"] + psidx.names = ["hello", "koalas", "goodbye"] + + self._test_sort_values(pidx, psidx) + + def test_index_drop_duplicates(self): + pidx = pd.Index([1, 1, 2]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.drop_duplicates().sort_values(), psidx.drop_duplicates().sort_values()) + + pidx = pd.MultiIndex.from_tuples([(1, 1), (1, 1), (2, 2)], names=["level1", "level2"]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.drop_duplicates().sort_values(), psidx.drop_duplicates().sort_values()) + + def test_index_sort(self): + idx = ps.Index([1, 2, 3, 4, 5]) + midx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2)]) + + with self.assertRaisesRegex( + TypeError, "cannot sort an Index object in-place, use sort_values instead" + ): + idx.sort() + with self.assertRaisesRegex( + TypeError, "cannot sort an Index object in-place, use sort_values instead" + ): + midx.sort() + + def test_multiindex_isna(self): + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + + with self.assertRaisesRegex(NotImplementedError, "isna is not defined for MultiIndex"): + psidx.isna() + + with self.assertRaisesRegex(NotImplementedError, "isna is not defined for MultiIndex"): + psidx.isnull() + + with self.assertRaisesRegex(NotImplementedError, "notna is not defined for MultiIndex"): + psidx.notna() + + with self.assertRaisesRegex(NotImplementedError, "notna is not defined for MultiIndex"): + psidx.notnull() + + def test_index_nunique(self): + pidx = pd.Index([1, 1, 2, None]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.nunique(), psidx.nunique()) + self.assert_eq(pidx.nunique(dropna=True), psidx.nunique(dropna=True)) + + def test_multiindex_nunique(self): + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + with self.assertRaisesRegex(NotImplementedError, "notna is not defined for MultiIndex"): + psidx.notnull() + + def test_multiindex_rename(self): + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psidx = ps.from_pandas(pidx) + + pidx = pidx.rename(list("ABC")) + psidx = psidx.rename(list("ABC")) + self.assert_eq(pidx, psidx) + + pidx = pidx.rename(["my", "name", "is"]) + psidx = psidx.rename(["my", "name", "is"]) + self.assert_eq(pidx, psidx) + + def test_multiindex_set_names(self): + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psidx = ps.from_pandas(pidx) + + pidx = pidx.set_names(["set", "new", "names"]) + psidx = psidx.set_names(["set", "new", "names"]) + self.assert_eq(pidx, psidx) + + pidx.set_names(["set", "new", "names"], inplace=True) + psidx.set_names(["set", "new", "names"], inplace=True) + self.assert_eq(pidx, psidx) + + pidx = pidx.set_names("first", level=0) + psidx = psidx.set_names("first", level=0) + self.assert_eq(pidx, psidx) + + pidx = pidx.set_names("second", level=1) + psidx = psidx.set_names("second", level=1) + self.assert_eq(pidx, psidx) + + pidx = pidx.set_names("third", level=2) + psidx = psidx.set_names("third", level=2) + self.assert_eq(pidx, psidx) + + pidx.set_names("first", level=0, inplace=True) + psidx.set_names("first", level=0, inplace=True) + self.assert_eq(pidx, psidx) + + pidx.set_names("second", level=1, inplace=True) + psidx.set_names("second", level=1, inplace=True) + self.assert_eq(pidx, psidx) + + pidx.set_names("third", level=2, inplace=True) + psidx.set_names("third", level=2, inplace=True) + self.assert_eq(pidx, psidx) + + def test_multiindex_from_tuples(self): + tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "blue")] + pidx = pd.MultiIndex.from_tuples(tuples) + psidx = ps.MultiIndex.from_tuples(tuples) + + self.assert_eq(pidx, psidx) + + def test_multiindex_from_product(self): + iterables = [[0, 1, 2], ["green", "purple"]] + pidx = pd.MultiIndex.from_product(iterables) + psidx = ps.MultiIndex.from_product(iterables) + + self.assert_eq(pidx, psidx) + + def test_multiindex_tuple_column_name(self): + column_labels = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")]) + pdf = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=column_labels) + pdf.set_index(("a", "x"), append=True, inplace=True) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf, psdf) + + def test_len(self): + pidx = pd.Index(range(10000)) + psidx = ps.from_pandas(pidx) + + self.assert_eq(len(pidx), len(psidx)) + + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + + self.assert_eq(len(pidx), len(psidx)) + + def test_delete(self): + pidx = pd.Index([10, 9, 8, 7, 6, 7, 8, 9, 10]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.delete(8).sort_values(), psidx.delete(8).sort_values()) + self.assert_eq(pidx.delete(-9).sort_values(), psidx.delete(-9).sort_values()) + self.assert_eq( + pidx.delete([-9, 0, 8]).sort_values(), psidx.delete([-9, 0, 8]).sort_values() + ) + + with self.assertRaisesRegex(IndexError, "index 9 is out of bounds for axis 0 with size 9"): + psidx.delete([0, 9]) + with self.assertRaisesRegex( + IndexError, "index -10 is out of bounds for axis 0 with size 9" + ): + psidx.delete([-10, 0]) + with self.assertRaisesRegex(IndexError, "index 9 is out of bounds for axis 0 with size 9"): + psidx.delete(9) + with self.assertRaisesRegex( + IndexError, "index -10 is out of bounds for axis 0 with size 9" + ): + psidx.delete(-10) + + # MultiIndex + pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + + self.assert_eq(pidx.delete(2).sort_values(), psidx.delete(2).sort_values()) + self.assert_eq(pidx.delete(-3).sort_values(), psidx.delete(-3).sort_values()) + self.assert_eq( + pidx.delete([-3, 0, 2]).sort_values(), psidx.delete([-3, 0, 2]).sort_values() + ) + + with self.assertRaisesRegex(IndexError, "index 3 is out of bounds for axis 0 with size 3"): + psidx.delete([0, 3]) + with self.assertRaisesRegex(IndexError, "index -4 is out of bounds for axis 0 with size 3"): + psidx.delete([-4, 0]) + with self.assertRaisesRegex(IndexError, "index 3 is out of bounds for axis 0 with size 3"): + psidx.delete(3) + with self.assertRaisesRegex(IndexError, "index -4 is out of bounds for axis 0 with size 3"): + psidx.delete(-4) + + def test_append(self): + # Index + pidx = pd.Index(range(10000)) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.append(pidx), psidx.append(psidx)) + + # Index with name + pidx1 = pd.Index(range(10000), name="a") + pidx2 = pd.Index(range(10000), name="b") + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + + self.assert_eq(pidx1.append(pidx2), psidx1.append(psidx2)) + + self.assert_eq(pidx2.append(pidx1), psidx2.append(psidx1)) + + # Index from DataFrame + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=["a", "b", "c"]) + pdf2 = pd.DataFrame({"a": [7, 8, 9], "d": [10, 11, None]}, index=["x", "y", "z"]) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + pidx1 = pdf1.set_index("a").index + pidx2 = pdf2.set_index("d").index + psidx1 = psdf1.set_index("a").index + psidx2 = psdf2.set_index("d").index + + self.assert_eq(pidx1.append(pidx2), psidx1.append(psidx2)) + + self.assert_eq(pidx2.append(pidx1), psidx2.append(psidx1)) + + # Index from DataFrame with MultiIndex columns + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + pdf2 = pd.DataFrame({"a": [7, 8, 9], "d": [10, 11, 12]}) + pdf1.columns = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")]) + pdf2.columns = pd.MultiIndex.from_tuples([("a", "x"), ("d", "y")]) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + pidx1 = pdf1.set_index(("a", "x")).index + pidx2 = pdf2.set_index(("d", "y")).index + psidx1 = psdf1.set_index(("a", "x")).index + psidx2 = psdf2.set_index(("d", "y")).index + + self.assert_eq(pidx1.append(pidx2), psidx1.append(psidx2)) + + self.assert_eq(pidx2.append(pidx1), psidx2.append(psidx1)) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pmidx.append(pmidx), psmidx.append(psmidx)) + + # MultiIndex with names + pmidx1 = pd.MultiIndex.from_tuples( + [("a", "x", 1), ("b", "y", 2), ("c", "z", 3)], names=["x", "y", "z"] + ) + pmidx2 = pd.MultiIndex.from_tuples( + [("a", "x", 1), ("b", "y", 2), ("c", "z", 3)], names=["p", "q", "r"] + ) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + + # TODO(SPARK-43241): MultiIndex.append not checking names for equality. + # Also refer to https://github.com/pandas-dev/pandas/pull/48288. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + pmidx1.append(pmidx2), psmidx1.append(psmidx2).rename([None, None, None]) + ) + else: + self.assert_eq(pmidx1.append(pmidx2), psmidx1.append(psmidx2)) + + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + pmidx2.append(pmidx1), psmidx2.append(psmidx1).rename([None, None, None]) + ) + else: + self.assert_eq(pmidx2.append(pmidx1), psmidx2.append(psmidx1)) + + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + pmidx1.append(pmidx2).names, + psmidx1.append(psmidx2).rename([None, None, None]).names, + ) + else: + self.assert_eq(pmidx1.append(pmidx2).names, psmidx1.append(psmidx2).names) + + # Index & MultiIndex is currently not supported + expected_error_message = r"append\(\) between Index & MultiIndex is currently not supported" + with self.assertRaisesRegex(NotImplementedError, expected_error_message): + psidx.append(psmidx) + with self.assertRaisesRegex(NotImplementedError, expected_error_message): + psmidx.append(psidx) + + # MultiIndexs with different levels is currently not supported + psmidx3 = ps.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + expected_error_message = ( + r"append\(\) between MultiIndexs with different levels is currently not supported" + ) + with self.assertRaisesRegex(NotImplementedError, expected_error_message): + psmidx.append(psmidx3) + + def test_argmin(self): + pidx = pd.Index([100, 50, 10, 20, 30, 60, 0, 50, 0, 100, 100, 100, 20, 0, 0]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.argmin(), psidx.argmin()) + + # MultiIndex + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + with self.assertRaisesRegex( + TypeError, "reduction operation 'argmin' not allowed for this dtype" + ): + psidx.argmin() + + def test_argmax(self): + pidx = pd.Index([100, 50, 10, 20, 30, 60, 0, 50, 0, 100, 100, 100, 20, 0, 0]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.argmax(), psidx.argmax()) + + # MultiIndex + psidx = ps.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) + with self.assertRaisesRegex( + TypeError, "reduction operation 'argmax' not allowed for this dtype" + ): + psidx.argmax() + + def test_min(self): + pidx = pd.Index([3, 2, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.min(), psidx.min()) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2)]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pmidx.min(), psmidx.min()) + + pidx = pd.DatetimeIndex(["2021-02-01", "2021-01-01", "2021-04-01", "2021-03-01"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.min(), psidx.min()) + + def test_max(self): + pidx = pd.Index([3, 2, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.max(), psidx.max()) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2)]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pmidx.max(), psmidx.max()) + + pidx = pd.DatetimeIndex(["2021-02-01", "2021-01-01", "2021-04-01", "2021-03-01"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.max(), psidx.max()) + + def test_monotonic(self): + # test monotonic_increasing & monotonic_decreasing for MultiIndex. + # Since the Behavior for null value was changed in pandas >= 1.0.0, + # several cases are tested differently. + datas = [] + + # increasing / decreasing ordered each index level with string + datas.append([("w", "a"), ("x", "b"), ("y", "c"), ("z", "d")]) + datas.append([("w", "d"), ("x", "c"), ("y", "b"), ("z", "a")]) + datas.append([("z", "a"), ("y", "b"), ("x", "c"), ("w", "d")]) + datas.append([("z", "d"), ("y", "c"), ("x", "b"), ("w", "a")]) + # mixed order each index level with string + datas.append([("z", "a"), ("x", "b"), ("y", "c"), ("w", "d")]) + datas.append([("z", "a"), ("y", "c"), ("x", "b"), ("w", "d")]) + + # increasing / decreasing ordered each index level with integer + datas.append([(1, 100), (2, 200), (3, 300), (4, 400), (5, 500)]) + datas.append([(1, 500), (2, 400), (3, 300), (4, 200), (5, 100)]) + datas.append([(5, 100), (4, 200), (3, 300), (2, 400), (1, 500)]) + datas.append([(5, 500), (4, 400), (3, 300), (2, 200), (1, 100)]) + # mixed order each index level with integer + datas.append([(1, 500), (3, 400), (2, 300), (4, 200), (5, 100)]) + datas.append([(1, 100), (2, 300), (3, 200), (4, 400), (5, 500)]) + + # integer / negative mixed tests + datas.append([("a", -500), ("b", -400), ("c", -300), ("d", -200), ("e", -100)]) + datas.append([("e", -500), ("d", -400), ("c", -300), ("b", -200), ("a", -100)]) + datas.append([(-5, "a"), (-4, "b"), (-3, "c"), (-2, "d"), (-1, "e")]) + datas.append([(-5, "e"), (-4, "d"), (-3, "c"), (-2, "b"), (-1, "a")]) + datas.append([(-5, "e"), (-3, "d"), (-2, "c"), (-4, "b"), (-1, "a")]) + datas.append([(-5, "e"), (-4, "c"), (-3, "b"), (-2, "d"), (-1, "a")]) + + # boolean type tests + datas.append([(True, True), (True, True)]) + datas.append([(True, True), (True, False)]) + datas.append([(True, False), (True, True)]) + datas.append([(False, True), (False, True)]) + datas.append([(False, True), (False, False)]) + datas.append([(False, False), (False, True)]) + datas.append([(True, True), (False, True)]) + datas.append([(True, True), (False, False)]) + datas.append([(True, False), (False, True)]) + datas.append([(False, True), (True, True)]) + datas.append([(False, True), (True, False)]) + datas.append([(False, False), (True, True)]) + + # duplicated index value tests + datas.append([("x", "d"), ("y", "c"), ("y", "b"), ("z", "a")]) + datas.append([("x", "d"), ("y", "b"), ("y", "c"), ("z", "a")]) + + # more depth tests + datas.append([("x", "d", "o"), ("y", "c", "p"), ("y", "c", "q"), ("z", "a", "r")]) + datas.append([("x", "d", "o"), ("y", "c", "q"), ("y", "c", "p"), ("z", "a", "r")]) + + # None type tests (None type is treated as False from pandas >= 1.1.4) + # Refer https://github.com/pandas-dev/pandas/issues/37220 + datas.append([(1, 100), (2, 200), (None, 300), (4, 400), (5, 500)]) + datas.append([(1, 100), (2, 200), (None, None), (4, 400), (5, 500)]) + datas.append([("x", "d"), ("y", "c"), ("y", None), ("z", "a")]) + datas.append([("x", "d"), ("y", "c"), ("y", "b"), (None, "a")]) + datas.append([("x", "d"), ("y", "b"), ("y", "c"), (None, "a")]) + datas.append([("x", "d", "o"), ("y", "c", "p"), ("y", "c", None), ("z", "a", "r")]) + + for data in datas: + with self.subTest(data=data): + pmidx = pd.MultiIndex.from_tuples(data) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + # datas below return different result depends on pandas version. + # Because the behavior of handling null values is changed in pandas >= 1.1.4. + # Since Koalas follows latest pandas, all of them should return `False`. + datas = [] + datas.append([(1, 100), (2, 200), (3, None), (4, 400), (5, 500)]) + datas.append([(1, None), (2, 200), (3, 300), (4, 400), (5, 500)]) + datas.append([(1, 100), (2, 200), (3, 300), (4, 400), (5, None)]) + datas.append([(False, None), (True, True)]) + datas.append([(None, False), (True, True)]) + datas.append([(False, False), (True, None)]) + datas.append([(False, False), (None, True)]) + datas.append([("x", "d"), ("y", None), ("y", None), ("z", "a")]) + datas.append([("x", "d", "o"), ("y", "c", None), ("y", "c", None), ("z", "a", "r")]) + datas.append([(1, 100), (2, 200), (3, 300), (4, 400), (None, 500)]) + datas.append([(1, 100), (2, 200), (3, 300), (4, 400), (None, None)]) + datas.append([(5, 100), (4, 200), (3, None), (2, 400), (1, 500)]) + datas.append([(5, None), (4, 200), (3, 300), (2, 400), (1, 500)]) + datas.append([(5, 100), (4, 200), (3, None), (2, 400), (1, 500)]) + datas.append([(5, 100), (4, 200), (3, 300), (2, 400), (1, None)]) + datas.append([(True, None), (True, True)]) + datas.append([(None, True), (True, True)]) + datas.append([(True, True), (None, True)]) + datas.append([(True, True), (True, None)]) + datas.append([(None, 100), (2, 200), (3, 300), (4, 400), (5, 500)]) + datas.append([(None, None), (2, 200), (3, 300), (4, 400), (5, 500)]) + datas.append([("x", "d"), ("y", None), ("y", "c"), ("z", "a")]) + datas.append([("x", "d", "o"), ("y", "c", None), ("y", "c", "q"), ("z", "a", "r")]) + + for data in datas: + with self.subTest(data=data): + pmidx = pd.MultiIndex.from_tuples(data) + psmidx = ps.from_pandas(pmidx) + if LooseVersion(pd.__version__) < LooseVersion("1.1.4"): + self.assert_eq(psmidx.is_monotonic_increasing, False) + self.assert_eq(psmidx.is_monotonic_decreasing, False) + else: + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + # The datas below are tested another way since they cannot be an arguments for + # `MultiIndex.from_tuples` in pandas >= 1.1.0. + # Refer https://github.com/databricks/koalas/pull/1688#issuecomment-667156560 for detail. + if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): + pmidx = pd.MultiIndex.from_tuples( + [(-5, None), (-4, None), (-3, None), (-2, None), (-1, None)] + ) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.is_monotonic_increasing, False) + self.assert_eq(psmidx.is_monotonic_decreasing, False) + + pmidx = pd.MultiIndex.from_tuples( + [(None, "e"), (None, "c"), (None, "b"), (None, "d"), (None, "a")] + ) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.is_monotonic_increasing, False) + self.assert_eq(psmidx.is_monotonic_decreasing, False) + + pmidx = pd.MultiIndex.from_tuples( + [(None, None), (None, None), (None, None), (None, None), (None, None)] + ) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.is_monotonic_increasing, False) + self.assert_eq(psmidx.is_monotonic_decreasing, False) + + pmidx = pd.MultiIndex.from_tuples([(None, None)]) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(psmidx.is_monotonic_increasing, False) + self.assert_eq(psmidx.is_monotonic_decreasing, False) + + else: + # For [(-5, None), (-4, None), (-3, None), (-2, None), (-1, None)] + psdf = ps.DataFrame({"a": [-5, -4, -3, -2, -1], "b": [1, 1, 1, 1, 1]}) + psdf["b"] = None + psmidx = psdf.set_index(["a", "b"]).index + pmidx = psmidx._to_pandas() + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + # For [(None, "e"), (None, "c"), (None, "b"), (None, "d"), (None, "a")] + psdf = ps.DataFrame({"a": [1, 1, 1, 1, 1], "b": ["e", "c", "b", "d", "a"]}) + psdf["a"] = None + psmidx = psdf.set_index(["a", "b"]).index + pmidx = psmidx._to_pandas() + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + # For [(None, None), (None, None), (None, None), (None, None), (None, None)] + psdf = ps.DataFrame({"a": [1, 1, 1, 1, 1], "b": [1, 1, 1, 1, 1]}) + psdf["a"] = None + psdf["b"] = None + psmidx = psdf.set_index(["a", "b"]).index + pmidx = psmidx._to_pandas() + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + # For [(None, None)] + psdf = ps.DataFrame({"a": [1], "b": [1]}) + psdf["a"] = None + psdf["b"] = None + psmidx = psdf.set_index(["a", "b"]).index + pmidx = psmidx._to_pandas() + self.assert_eq(psmidx.is_monotonic_increasing, pmidx.is_monotonic_increasing) + self.assert_eq(psmidx.is_monotonic_decreasing, pmidx.is_monotonic_decreasing) + + def test_difference(self): + # Index + pidx1 = pd.Index([1, 2, 3, 4], name="koalas") + pidx2 = pd.Index([3, 4, 5, 6], name="koalas") + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + # Series + pser = pd.Series([3, 4, 5, 6], name="koalas") + psser = ps.from_pandas(pser) + + self.assert_eq( + psidx1.difference(psidx2).sort_values(), pidx1.difference(pidx2).sort_values() + ) + self.assert_eq(psidx1.difference(psser).sort_values(), pidx1.difference(pser).sort_values()) + self.assert_eq( + psidx1.difference([3, 4, 5, 6]).sort_values(), + pidx1.difference([3, 4, 5, 6]).sort_values(), + ) + self.assert_eq( + psidx1.difference((3, 4, 5, 6)).sort_values(), + pidx1.difference((3, 4, 5, 6)).sort_values(), + ) + self.assert_eq( + psidx1.difference({3, 4, 5, 6}).sort_values(), + pidx1.difference({3, 4, 5, 6}).sort_values(), + ) + self.assert_eq( + psidx1.difference({3: 1, 4: 2, 5: 3, 6: 4}).sort_values(), + pidx1.difference({3: 1, 4: 2, 5: 3, 6: 4}).sort_values(), + ) + + # Exceptions for Index + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx1.difference("1234") + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx1.difference(1234) + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx1.difference(12.34) + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx1.difference(None) + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx1.difference(np.nan) + with self.assertRaisesRegex( + ValueError, "The 'sort' keyword only takes the values of None or True; 1 was passed." + ): + psidx1.difference(psidx2, sort=1) + + # MultiIndex + pmidx1 = pd.MultiIndex.from_tuples( + [("a", "x", 1), ("b", "y", 2), ("c", "z", 3)], names=["hello", "koalas", "world"] + ) + pmidx2 = pd.MultiIndex.from_tuples( + [("a", "x", 1), ("b", "z", 2), ("k", "z", 3)], names=["hello", "koalas", "world"] + ) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + + self.assert_eq( + psmidx1.difference(psmidx2).sort_values(), pmidx1.difference(pmidx2).sort_values() + ) + self.assert_eq( + psmidx1.difference(psidx1).sort_values(), pmidx1.difference(pidx1).sort_values() + ) + self.assert_eq( + psidx1.difference(psmidx1).sort_values(), pidx1.difference(pmidx1).sort_values() + ) + self.assert_eq(psidx1.difference(psser).sort_values(), pidx1.difference(pser).sort_values()) + self.assert_eq( + psmidx1.difference({("a", "x", 1)}).sort_values(), + pmidx1.difference({("a", "x", 1)}).sort_values(), + ) + self.assert_eq( + psmidx1.difference({("a", "x", 1): [1, 2, 3]}).sort_values(), + pmidx1.difference({("a", "x", 1): [1, 2, 3]}).sort_values(), + ) + + # Exceptions for MultiIndex + with self.assertRaisesRegex(TypeError, "other must be a MultiIndex or a list of tuples"): + psmidx1.difference(["b", "z", "2"]) + + def test_repeat(self): + pidx = pd.Index(["a", "b", "c"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.repeat(3).sort_values(), pidx.repeat(3).sort_values()) + self.assert_eq(psidx.repeat(0).sort_values(), pidx.repeat(0).sort_values()) + self.assert_eq((psidx + "x").repeat(3).sort_values(), (pidx + "x").repeat(3).sort_values()) + + self.assertRaises(ValueError, lambda: psidx.repeat(-1)) + self.assertRaises(TypeError, lambda: psidx.repeat("abc")) + + pmidx = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(psmidx.repeat(3).sort_values(), pmidx.repeat(3).sort_values()) + self.assert_eq(psmidx.repeat(0).sort_values(), pmidx.repeat(0).sort_values(), almost=True) + + self.assertRaises(ValueError, lambda: psmidx.repeat(-1)) + self.assertRaises(TypeError, lambda: psmidx.repeat("abc")) + + def test_unique(self): + pidx = pd.Index(["a", "b", "a"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.unique().sort_values(), pidx.unique().sort_values()) + self.assert_eq(psidx.unique().sort_values(), pidx.unique().sort_values()) + + pmidx = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("x", "a")]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(psmidx.unique().sort_values(), pmidx.unique().sort_values()) + self.assert_eq(psmidx.unique().sort_values(), pmidx.unique().sort_values()) + + with self.assertRaisesRegex( + IndexError, "Too many levels: Index has only 1 level, -2 is not a valid level number" + ): + psidx.unique(level=-2) + + def test_asof(self): + # Increasing values + pidx = pd.Index(["2013-12-31", "2014-01-02", "2014-01-03"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.asof("2014-01-01"), pidx.asof("2014-01-01")) + self.assert_eq(psidx.asof("2014-01-02"), pidx.asof("2014-01-02")) + self.assert_eq(repr(psidx.asof("1999-01-02")), repr(pidx.asof("1999-01-02"))) + self.assert_eq(psidx.asof("2014-01-04"), pidx.asof("2014-01-04")) + + pidx = pd.DatetimeIndex(["2013-12-31", "2014-01-02", "2014-01-03"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.asof("2014-01-01"), pidx.asof("2014-01-01")) + self.assert_eq(psidx.asof("2014-01-02"), pidx.asof("2014-01-02")) + self.assert_eq(repr(psidx.asof("1999-01-02")), repr(pidx.asof("1999-01-02"))) + + # Decreasing values + pidx = pd.Index(["2014-01-03", "2014-01-02", "2013-12-31"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.asof("2014-01-01"), pidx.asof("2014-01-01")) + self.assert_eq(psidx.asof("2014-01-02"), pidx.asof("2014-01-02")) + self.assert_eq(psidx.asof("1999-01-02"), pidx.asof("1999-01-02")) + self.assert_eq(repr(psidx.asof("2015-01-02")), repr(pidx.asof("2015-01-02"))) + + pidx = pd.DatetimeIndex(["2014-01-03", "2014-01-02", "2013-12-31"]) + psidx = ps.from_pandas(pidx) + + # TODO: a pandas bug? + # self.assert_eq(psidx.asof("2014-01-01"), pidx.asof("2014-01-01")) + # self.assert_eq(psidx.asof("2014-01-02"), pidx.asof("2014-01-02")) + # self.assert_eq(psidx.asof("1999-01-02"), pidx.asof("1999-01-02")) + # self.assert_eq(repr(psidx.asof("2015-01-02")), repr(pidx.asof("2015-01-02"))) + self.assert_eq(psidx.asof("2014-01-01"), pd.Timestamp("2014-01-02 00:00:00")) + self.assert_eq(psidx.asof("2014-01-02"), pd.Timestamp("2014-01-02 00:00:00")) + self.assert_eq(psidx.asof("1999-01-02"), pd.Timestamp("2013-12-31 00:00:00")) + self.assert_eq(repr(psidx.asof("2015-01-02")), repr(pd.NaT)) + + # Not increasing, neither decreasing (ValueError) + psidx = ps.Index(["2013-12-31", "2015-01-02", "2014-01-03"]) + self.assertRaises(ValueError, lambda: psidx.asof("2013-12-31")) + + psmidx = ps.MultiIndex.from_tuples([("a", "a"), ("a", "b"), ("a", "c")]) + self.assertRaises(NotImplementedError, lambda: psmidx.asof(("a", "b"))) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43608): Enable IndexesTests.test_union for pandas 2.0.0.", + ) + def test_union(self): + # Index + pidx1 = pd.Index([1, 2, 3, 4]) + pidx2 = pd.Index([3, 4, 5, 6]) + pidx3 = pd.Index([7.0, 8.0, 9.0, 10.0]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + psidx3 = ps.from_pandas(pidx3) + + self.assert_eq(psidx1.union(psidx2), pidx1.union(pidx2)) + self.assert_eq(psidx2.union(psidx1), pidx2.union(pidx1)) + self.assert_eq(psidx1.union(psidx3), pidx1.union(pidx3)) + # Deprecated case, but adding to track if pandas stop supporting union + # as a set operation. It should work fine until stop supporting anyway. + # No longer supported from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(psidx1 | psidx2, ps.Index([3, 4], dtype="int64")) + else: + self.assert_eq(pidx1 | pidx2, psidx1 | psidx2) + + self.assert_eq(psidx1.union([3, 4, 5, 6]), pidx1.union([3, 4, 5, 6]), almost=True) + self.assert_eq(psidx2.union([1, 2, 3, 4]), pidx2.union([1, 2, 3, 4]), almost=True) + self.assert_eq( + psidx1.union(ps.Series([3, 4, 5, 6])), pidx1.union(pd.Series([3, 4, 5, 6])), almost=True + ) + self.assert_eq( + psidx2.union(ps.Series([1, 2, 3, 4])), pidx2.union(pd.Series([1, 2, 3, 4])), almost=True + ) + + # Testing if the result is correct after sort=False. + self.assert_eq( + psidx1.union(psidx2, sort=False).sort_values(), + pidx1.union(pidx2, sort=False).sort_values(), + ) + self.assert_eq( + psidx2.union(psidx1, sort=False).sort_values(), + pidx2.union(pidx1, sort=False).sort_values(), + ) + self.assert_eq( + psidx1.union([3, 4, 5, 6], sort=False).sort_values(), + pidx1.union([3, 4, 5, 6], sort=False).sort_values(), + almost=True, + ) + self.assert_eq( + psidx2.union([1, 2, 3, 4], sort=False).sort_values(), + pidx2.union([1, 2, 3, 4], sort=False).sort_values(), + almost=True, + ) + self.assert_eq( + psidx1.union(ps.Series([3, 4, 5, 6]), sort=False).sort_values(), + pidx1.union(pd.Series([3, 4, 5, 6]), sort=False).sort_values(), + almost=True, + ) + self.assert_eq( + psidx2.union(ps.Series([1, 2, 3, 4]), sort=False).sort_values(), + pidx2.union(pd.Series([1, 2, 3, 4]), sort=False).sort_values(), + almost=True, + ) + + pidx1 = pd.Index([1, 2, 3, 4, 3, 4, 3, 4]) + pidx2 = pd.Index([3, 4, 3, 4, 5, 6]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + + self.assert_eq(psidx1.union(psidx2), pidx1.union(pidx2)) + self.assert_eq( + psidx1.union([3, 4, 3, 3, 5, 6]), pidx1.union([3, 4, 3, 4, 5, 6]), almost=True + ) + self.assert_eq( + psidx1.union(ps.Series([3, 4, 3, 3, 5, 6])), + pidx1.union(pd.Series([3, 4, 3, 4, 5, 6])), + almost=True, + ) + + # Manually create the expected result here since there is a bug in Index.union + # dropping duplicated values in pandas < 1.3. + expected = pd.Index([1, 2, 3, 3, 3, 4, 4, 4, 5, 6]) + self.assert_eq(psidx2.union(psidx1), expected) + self.assert_eq( + psidx2.union([1, 2, 3, 4, 3, 4, 3, 4]), + expected, + almost=True, + ) + self.assert_eq( + psidx2.union(ps.Series([1, 2, 3, 4, 3, 4, 3, 4])), + expected, + almost=True, + ) + + # MultiIndex + pmidx1 = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("x", "a"), ("x", "b")]) + pmidx2 = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("x", "c"), ("x", "d")]) + pmidx3 = pd.MultiIndex.from_tuples([(1, 1), (1, 2), (1, 3), (1, 4), (1, 3), (1, 4)]) + pmidx4 = pd.MultiIndex.from_tuples([(1, 3), (1, 4), (1, 5), (1, 6)]) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + psmidx3 = ps.from_pandas(pmidx3) + psmidx4 = ps.from_pandas(pmidx4) + + # Manually create the expected result here since there is a bug in MultiIndex.union + # dropping duplicated values in pandas < 1.3. + expected = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "a"), ("x", "b"), ("x", "b"), ("x", "c"), ("x", "d")] + ) + self.assert_eq(psmidx1.union(psmidx2), expected) + self.assert_eq(psmidx2.union(psmidx1), expected) + self.assert_eq( + psmidx1.union([("x", "a"), ("x", "b"), ("x", "c"), ("x", "d")]), + expected, + ) + self.assert_eq( + psmidx2.union([("x", "a"), ("x", "b"), ("x", "a"), ("x", "b")]), + expected, + ) + + expected = pd.MultiIndex.from_tuples( + [(1, 1), (1, 2), (1, 3), (1, 3), (1, 4), (1, 4), (1, 5), (1, 6)] + ) + self.assert_eq(psmidx3.union(psmidx4), expected) + self.assert_eq(psmidx4.union(psmidx3), expected) + self.assert_eq( + psmidx3.union([(1, 3), (1, 4), (1, 5), (1, 6)]), + expected, + ) + self.assert_eq( + psmidx4.union([(1, 1), (1, 2), (1, 3), (1, 4), (1, 3), (1, 4)]), + expected, + ) + + # Testing if the result is correct after sort=False. + # Manually create the expected result here since there is a bug in MultiIndex.union + # dropping duplicated values in pandas < 1.3. + expected = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "a"), ("x", "b"), ("x", "b"), ("x", "c"), ("x", "d")] + ) + self.assert_eq(psmidx1.union(psmidx2, sort=False).sort_values(), expected) + self.assert_eq(psmidx2.union(psmidx1, sort=False).sort_values(), expected) + self.assert_eq( + psmidx1.union( + [("x", "a"), ("x", "b"), ("x", "c"), ("x", "d")], sort=False + ).sort_values(), + expected, + ) + self.assert_eq( + psmidx2.union( + [("x", "a"), ("x", "b"), ("x", "a"), ("x", "b")], sort=False + ).sort_values(), + expected, + ) + + expected = pd.MultiIndex.from_tuples( + [(1, 1), (1, 2), (1, 3), (1, 3), (1, 4), (1, 4), (1, 5), (1, 6)] + ) + self.assert_eq(psmidx3.union(psmidx4, sort=False).sort_values(), expected) + self.assert_eq(psmidx4.union(psmidx3, sort=False).sort_values(), expected) + self.assert_eq( + psmidx3.union([(1, 3), (1, 4), (1, 5), (1, 6)], sort=False).sort_values(), expected + ) + self.assert_eq( + psmidx4.union( + [(1, 1), (1, 2), (1, 3), (1, 4), (1, 3), (1, 4)], sort=False + ).sort_values(), + expected, + ) + + self.assertRaises(NotImplementedError, lambda: psidx1.union(psmidx1)) + self.assertRaises(TypeError, lambda: psmidx1.union(psidx1)) + self.assertRaises(TypeError, lambda: psmidx1.union(["x", "a"])) + self.assertRaises(ValueError, lambda: psidx1.union(ps.range(2))) + + def test_take(self): + # Index + pidx = pd.Index([100, 200, 300, 400, 500], name="Koalas") + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.take([0, 2, 4]).sort_values(), pidx.take([0, 2, 4]).sort_values()) + self.assert_eq( + psidx.take(range(0, 5, 2)).sort_values(), pidx.take(range(0, 5, 2)).sort_values() + ) + self.assert_eq(psidx.take([-4, -2, 0]).sort_values(), pidx.take([-4, -2, 0]).sort_values()) + self.assert_eq( + psidx.take(range(-4, 1, 2)).sort_values(), pidx.take(range(-4, 1, 2)).sort_values() + ) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("x", "c")], names=["hello", "Koalas"] + ) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(psmidx.take([0, 2]).sort_values(), pmidx.take([0, 2]).sort_values()) + self.assert_eq( + psmidx.take(range(0, 4, 2)).sort_values(), pmidx.take(range(0, 4, 2)).sort_values() + ) + self.assert_eq(psmidx.take([-2, 0]).sort_values(), pmidx.take([-2, 0]).sort_values()) + self.assert_eq( + psmidx.take(range(-2, 1, 2)).sort_values(), pmidx.take(range(-2, 1, 2)).sort_values() + ) + + # Checking the type of indices. + self.assertRaises(TypeError, lambda: psidx.take(1)) + self.assertRaises(TypeError, lambda: psidx.take("1")) + self.assertRaises(TypeError, lambda: psidx.take({1, 2})) + self.assertRaises(TypeError, lambda: psidx.take({1: None, 2: None})) + self.assertRaises(TypeError, lambda: psmidx.take(1)) + self.assertRaises(TypeError, lambda: psmidx.take("1")) + self.assertRaises(TypeError, lambda: psmidx.take({1, 2})) + self.assertRaises(TypeError, lambda: psmidx.take({1: None, 2: None})) + + def test_index_get_level_values(self): + pidx = pd.Index([1, 2, 3], name="ks") + psidx = ps.from_pandas(pidx) + + for level in [0, "ks"]: + self.assert_eq(psidx.get_level_values(level), pidx.get_level_values(level)) + + def test_multiindex_get_level_values(self): + pmidx = pd.MultiIndex.from_tuples([("a", "d"), ("b", "e"), ("c", "f")]) + pmidx.names = ["level_1", "level_2"] + psmidx = ps.from_pandas(pmidx) + + for level in [0, 1, "level_1", "level_2"]: + self.assert_eq(psmidx.get_level_values(level), pmidx.get_level_values(level)) + + def test_index_get_level_number(self): + # name of two levels are the same, which is None + psdf = ps.DataFrame({"a": [1, 2, 3]}, index=[list("aac"), list("ddf")]) + with self.assertRaisesRegex( + ValueError, "The name None occurs multiple times, use a level number" + ): + psdf.index._get_level_number(None) + + mi = pd.MultiIndex.from_arrays((list("abc"), list("def"))) + mi.names = ["level_1", "level_2"] + psdf = ps.DataFrame({"a": [1, 2, 3]}, index=mi) + + # level is not int and not in the level name list + with self.assertRaisesRegex(KeyError, "Level lv_3 not found"): + psdf.index._get_level_number("lv_3") + + # level is int, but an invalid negative number + with self.assertRaisesRegex(IndexError, "Too many levels: Index has only"): + psdf.index._get_level_number(-3) + + # level is int, but an invalid positive number + with self.assertRaisesRegex(IndexError, "Too many levels: Index has only"): + psdf.index._get_level_number(3) + + # Correct and valid inputs in numbers + level_number = [-2, -1, 0, 1] + outputs = [0, 1, 0, 1] + + for lv, output in zip(level_number, outputs): + self.assertEqual(output, psdf.index._get_level_number(lv)) + + # Valid inputs as level names + level_names = ["level_1", "level_2"] + outputs = [0, 1] + + for lv, output in zip(level_names, outputs): + self.assertEqual(output, psdf.index._get_level_number(lv)) + + def test_holds_integer(self): + pidx = pd.Index([1, 2, 3, 4]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.holds_integer(), psidx.holds_integer()) + + pidx = pd.Index([1.1, 2.2, 3.3, 4.4]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.holds_integer(), psidx.holds_integer()) + + pidx = pd.Index(["A", "B", "C", "D"]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.holds_integer(), psidx.holds_integer()) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "a")]) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(pmidx.holds_integer(), psmidx.holds_integer()) + + pmidx = pd.MultiIndex.from_tuples([(10, 1), (10, 2), (20, 1)]) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(pmidx.holds_integer(), psmidx.holds_integer()) + + def test_abs(self): + pidx = pd.Index([-2, -1, 0, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(abs(pidx), abs(psidx)) + self.assert_eq(np.abs(pidx), np.abs(psidx)) + + psidx = ps.MultiIndex.from_tuples([(1, 2)], names=["level1", "level2"]) + with self.assertRaisesRegex(TypeError, "perform __abs__ with this index"): + abs(psidx) + + def test_hasnans(self): + # BooleanType + pidx = pd.Index([True, False, True, True]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.hasnans, psidx.hasnans) + + pidx = pd.Index([True, False, np.nan, True]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.hasnans, psidx.hasnans) + + # TimestampType + pser = pd.Series([pd.Timestamp("2020-07-30") for _ in range(3)]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + pser = pd.Series([pd.Timestamp("2020-07-30"), np.nan, pd.Timestamp("2020-07-30")]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + # empty + pidx = pd.Index([]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.hasnans, psidx.hasnans) + + # Not supported for MultiIndex + psmidx = ps.Index([("a", 1), ("b", 2)]) + self.assertRaises(NotImplementedError, lambda: psmidx.hasnans()) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43607): Enable IndexesTests.test_intersection for pandas 2.0.0.", + ) + def test_intersection(self): + pidx = pd.Index([1, 2, 3, 4], name="Koalas") + psidx = ps.from_pandas(pidx) + + # other = Index + pidx_other = pd.Index([3, 4, 5, 6], name="Koalas") + psidx_other = ps.from_pandas(pidx_other) + self.assert_eq(pidx.intersection(pidx_other), psidx.intersection(psidx_other).sort_values()) + self.assert_eq( + (pidx + 1).intersection(pidx_other), (psidx + 1).intersection(psidx_other).sort_values() + ) + # Deprecated case, but adding to track if pandas stop supporting intersection + # as a set operation. It should work fine until stop supporting anyway. + # No longer supported from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + (psidx & psidx_other).sort_values(), ps.Index([3, 1, 7, 1], dtype="int64") + ) + else: + self.assert_eq(pidx & pidx_other, (psidx & psidx_other).sort_values()) + + pidx_other_different_name = pd.Index([3, 4, 5, 6], name="Databricks") + psidx_other_different_name = ps.from_pandas(pidx_other_different_name) + self.assert_eq( + pidx.intersection(pidx_other_different_name), + psidx.intersection(psidx_other_different_name).sort_values(), + ) + self.assert_eq( + (pidx + 1).intersection(pidx_other_different_name), + (psidx + 1).intersection(psidx_other_different_name).sort_values(), + ) + + pidx_other_from_frame = pd.DataFrame({"a": [3, 4, 5, 6]}).set_index("a").index + psidx_other_from_frame = ps.from_pandas(pidx_other_from_frame) + self.assert_eq( + pidx.intersection(pidx_other_from_frame), + psidx.intersection(psidx_other_from_frame).sort_values(), + ) + self.assert_eq( + (pidx + 1).intersection(pidx_other_from_frame), + (psidx + 1).intersection(psidx_other_from_frame).sort_values(), + ) + + # other = MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psmidx = ps.from_pandas(pmidx) + if LooseVersion(pd.__version__) < LooseVersion("1.2.0"): + self.assert_eq( + psidx.intersection(psmidx).sort_values(), + psidx._psdf.head(0).index.rename(None), + almost=True, + ) + self.assert_eq( + (psidx + 1).intersection(psmidx).sort_values(), + psidx._psdf.head(0).index.rename(None), + almost=True, + ) + else: + self.assert_eq( + pidx.intersection(pmidx), psidx.intersection(psmidx).sort_values(), almost=True + ) + self.assert_eq( + (pidx + 1).intersection(pmidx), + (psidx + 1).intersection(psmidx).sort_values(), + almost=True, + ) + + # other = Series + pser = pd.Series([3, 4, 5, 6]) + psser = ps.from_pandas(pser) + if LooseVersion(pd.__version__) < LooseVersion("1.2.0"): + self.assert_eq(psidx.intersection(psser).sort_values(), ps.Index([3, 4], name="Koalas")) + self.assert_eq( + (psidx + 1).intersection(psser).sort_values(), ps.Index([3, 4, 5], name="Koalas") + ) + else: + self.assert_eq(pidx.intersection(pser), psidx.intersection(psser).sort_values()) + self.assert_eq( + (pidx + 1).intersection(pser), (psidx + 1).intersection(psser).sort_values() + ) + + pser_different_name = pd.Series([3, 4, 5, 6], name="Databricks") + psser_different_name = ps.from_pandas(pser_different_name) + if LooseVersion(pd.__version__) < LooseVersion("1.2.0"): + self.assert_eq( + psidx.intersection(psser_different_name).sort_values(), + ps.Index([3, 4], name="Koalas"), + ) + self.assert_eq( + (psidx + 1).intersection(psser_different_name).sort_values(), + ps.Index([3, 4, 5], name="Koalas"), + ) + else: + self.assert_eq( + pidx.intersection(pser_different_name), + psidx.intersection(psser_different_name).sort_values(), + ) + self.assert_eq( + (pidx + 1).intersection(pser_different_name), + (psidx + 1).intersection(psser_different_name).sort_values(), + ) + + others = ([3, 4, 5, 6], (3, 4, 5, 6), {3: None, 4: None, 5: None, 6: None}) + for other in others: + if LooseVersion(pd.__version__) < LooseVersion("1.2.0"): + self.assert_eq( + psidx.intersection(other).sort_values(), ps.Index([3, 4], name="Koalas") + ) + self.assert_eq( + (psidx + 1).intersection(other).sort_values(), + ps.Index([3, 4, 5], name="Koalas"), + ) + else: + self.assert_eq(pidx.intersection(other), psidx.intersection(other).sort_values()) + self.assert_eq( + (pidx + 1).intersection(other), (psidx + 1).intersection(other).sort_values() + ) + + # MultiIndex / other = Index + self.assert_eq( + pmidx.intersection(pidx), psmidx.intersection(psidx).sort_values(), almost=True + ) + self.assert_eq( + pmidx.intersection(pidx_other_from_frame), + psmidx.intersection(psidx_other_from_frame).sort_values(), + almost=True, + ) + + # MultiIndex / other = MultiIndex + pmidx_other = pd.MultiIndex.from_tuples([("c", "z"), ("d", "w")]) + psmidx_other = ps.from_pandas(pmidx_other) + self.assert_eq( + pmidx.intersection(pmidx_other), psmidx.intersection(psmidx_other).sort_values() + ) + + # MultiIndex / other = list + other = [("c", "z"), ("d", "w")] + self.assert_eq(pmidx.intersection(other), psmidx.intersection(other).sort_values()) + + # MultiIndex / other = tuple + other = (("c", "z"), ("d", "w")) + self.assert_eq(pmidx.intersection(other), psmidx.intersection(other).sort_values()) + + # MultiIndex / other = dict + other = {("c", "z"): None, ("d", "w"): None} + self.assert_eq(pmidx.intersection(other), psmidx.intersection(other).sort_values()) + + # MultiIndex with different names. + pmidx1 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")], names=["X", "Y"]) + pmidx2 = pd.MultiIndex.from_tuples([("c", "z"), ("d", "w")], names=["A", "B"]) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.intersection(pmidx2), psmidx1.intersection(psmidx2).sort_values()) + + with self.assertRaisesRegex(TypeError, "Input must be Index or array-like"): + psidx.intersection(4) + with self.assertRaisesRegex(TypeError, "other must be a MultiIndex or a list of tuples"): + psmidx.intersection(4) + with self.assertRaisesRegex(TypeError, "other must be a MultiIndex or a list of tuples"): + psmidx.intersection(ps.Series([3, 4, 5, 6])) + with self.assertRaisesRegex(TypeError, "other must be a MultiIndex or a list of tuples"): + psmidx.intersection([("c", "z"), ["d", "w"]]) + with self.assertRaisesRegex(ValueError, "Index data must be 1-dimensional"): + psidx.intersection(ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})) + with self.assertRaisesRegex(ValueError, "Index data must be 1-dimensional"): + psmidx.intersection(ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})) + # other = list of tuple + with self.assertRaisesRegex(ValueError, "Names should be list-like for a MultiIndex"): + psidx.intersection([(1, 2), (3, 4)]) + + def test_item(self): + pidx = pd.Index([10]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.item(), psidx.item()) + + # with timestamp + pidx = pd.Index([datetime(1990, 3, 9)]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.item(), psidx.item()) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x")]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pmidx.item(), psmidx.item()) + + # MultiIndex with timestamp + pmidx = pd.MultiIndex.from_tuples([(datetime(1990, 3, 9), datetime(2019, 8, 15))]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pidx.item(), psidx.item()) + + err_msg = "can only convert an array of size 1 to a Python scalar" + with self.assertRaisesRegex(ValueError, err_msg): + ps.Index([10, 20]).item() + with self.assertRaisesRegex(ValueError, err_msg): + ps.MultiIndex.from_tuples([("a", "x"), ("b", "y")]).item() + + def test_inferred_type(self): + # Integer + pidx = pd.Index([1, 2, 3]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.inferred_type, psidx.inferred_type) + + # Floating + pidx = pd.Index([1.0, 2.0, 3.0]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.inferred_type, psidx.inferred_type) + + # String + pidx = pd.Index(["a", "b", "c"]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.inferred_type, psidx.inferred_type) + + # Boolean + pidx = pd.Index([True, False, True, False]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.inferred_type, psidx.inferred_type) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x")]) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(pmidx.inferred_type, psmidx.inferred_type) + + def test_multi_index_from_index(self): + tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "blue")] + pmidx = pd.Index(tuples) + psmidx = ps.Index(tuples) + + self.assertTrue(isinstance(psmidx, ps.MultiIndex)) + self.assert_eq(pmidx, psmidx) + + # Specify the `names` + # Specify the `names` while Index creating is no longer supported from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + pmidx = pd.Index(tuples) + pmidx.names = ["Hello", "Koalas"] + psmidx = ps.Index(tuples) + psmidx.names = ["Hello", "Koalas"] + else: + pmidx = pd.Index(tuples, names=["Hello", "Koalas"]) + psmidx = ps.Index(tuples, names=["Hello", "Koalas"]) + + self.assertTrue(isinstance(psmidx, ps.MultiIndex)) + self.assert_eq(pmidx, psmidx) + + def test_multiindex_from_frame(self): + pdf = pd.DataFrame( + [["HI", "Temp"], ["HI", "Precip"], ["NJ", "Temp"], ["NJ", "Precip"]], columns=["a", "b"] + ) + psdf = ps.from_pandas(pdf) + pidx = pd.MultiIndex.from_frame(pdf) + psidx = ps.MultiIndex.from_frame(psdf) + + self.assert_eq(pidx, psidx) + + # Specify `names` + pidx = pd.MultiIndex.from_frame(pdf, names=["state", "observation"]) + psidx = ps.MultiIndex.from_frame(psdf, names=["state", "observation"]) + self.assert_eq(pidx, psidx) + + pidx = pd.MultiIndex.from_frame(pdf, names=("state", "observation")) + psidx = ps.MultiIndex.from_frame(psdf, names=("state", "observation")) + self.assert_eq(pidx, psidx) + + # MultiIndex columns + pidx = pd.MultiIndex.from_tuples([("a", "w"), ("b", "x")]) + pdf.columns = pidx + psdf = ps.from_pandas(pdf) + + pidx = pd.MultiIndex.from_frame(pdf) + psidx = ps.MultiIndex.from_frame(psdf) + + self.assert_eq(pidx, psidx) + + # tuples for names + pidx = pd.MultiIndex.from_frame(pdf, names=[("a", "w"), ("b", "x")]) + psidx = ps.MultiIndex.from_frame(psdf, names=[("a", "w"), ("b", "x")]) + + self.assert_eq(pidx, psidx) + + err_msg = "Input must be a DataFrame" + with self.assertRaisesRegex(TypeError, err_msg): + ps.MultiIndex.from_frame({"a": [1, 2, 3], "b": [4, 5, 6]}) + + self.assertRaises(TypeError, lambda: ps.MultiIndex.from_frame(psdf, names="ab")) + + # non-string names + self.assert_eq( + ps.MultiIndex.from_frame(psdf, names=[0, 1]), + pd.MultiIndex.from_frame(pdf, names=[0, 1]), + ) + self.assert_eq( + ps.MultiIndex.from_frame(psdf, names=[("x", 0), ("y", 1)]), + pd.MultiIndex.from_frame(pdf, names=[("x", 0), ("y", 1)]), + ) + + pdf = pd.DataFrame([["HI", "Temp"], ["HI", "Precip"], ["NJ", "Temp"], ["NJ", "Precip"]]) + psdf = ps.from_pandas(pdf) + self.assert_eq(ps.MultiIndex.from_frame(psdf), pd.MultiIndex.from_frame(pdf)) + + def test_is_type_compatible(self): + data_types = ["integer", "floating", "string", "boolean"] + # Integer + pidx = pd.Index([1, 2, 3]) + psidx = ps.from_pandas(pidx) + # is_type_compatible is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + expected_results = [True, False, False, False] + for data_type, expected_result in zip(data_types, expected_results): + self.assert_eq(psidx.is_type_compatible(data_type), expected_result) + else: + for data_type in data_types: + self.assert_eq( + pidx.is_type_compatible(data_type), psidx.is_type_compatible(data_type) + ) + + # Floating + pidx = pd.Index([1.0, 2.0, 3.0]) + psidx = ps.from_pandas(pidx) + # is_type_compatible is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + expected_results = [False, True, False, False] + for data_type, expected_result in zip(data_types, expected_results): + self.assert_eq(psidx.is_type_compatible(data_type), expected_result) + else: + for data_type in data_types: + self.assert_eq( + pidx.is_type_compatible(data_type), psidx.is_type_compatible(data_type) + ) + + # String + pidx = pd.Index(["a", "b", "c"]) + psidx = ps.from_pandas(pidx) + # is_type_compatible is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + expected_results = [False, False, True, False] + for data_type, expected_result in zip(data_types, expected_results): + self.assert_eq(psidx.is_type_compatible(data_type), expected_result) + else: + for data_type in data_types: + self.assert_eq( + pidx.is_type_compatible(data_type), psidx.is_type_compatible(data_type) + ) + + # Boolean + pidx = pd.Index([True, False, True, False]) + psidx = ps.from_pandas(pidx) + # is_type_compatible is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + expected_results = [False, False, False, True] + for data_type, expected_result in zip(data_types, expected_results): + self.assert_eq(psidx.is_type_compatible(data_type), expected_result) + else: + for data_type in data_types: + self.assert_eq( + pidx.is_type_compatible(data_type), psidx.is_type_compatible(data_type) + ) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x")]) + psmidx = ps.from_pandas(pmidx) + # is_type_compatible is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + expected_results = [False, False, False, False] + for data_type, expected_result in zip(data_types, expected_results): + self.assert_eq(psmidx.is_type_compatible(data_type), expected_result) + else: + for data_type in data_types: + self.assert_eq( + pmidx.is_type_compatible(data_type), psmidx.is_type_compatible(data_type) + ) + + def test_asi8(self): + # Integer + pidx = pd.Index([1, 2, 3]) + psidx = ps.from_pandas(pidx) + # asi8 is removed from pandas 2.0.0. + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(np.array(pidx), psidx.asi8) + self.assert_eq(np.array(pidx.astype("int")), psidx.astype("int").asi8) + self.assert_eq(np.array(pidx.astype("int16")), psidx.astype("int16").asi8) + self.assert_eq(np.array(pidx.astype("int8")), psidx.astype("int8").asi8) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + self.assert_eq(pidx.astype("int").asi8, psidx.astype("int").asi8) + self.assert_eq(pidx.astype("int16").asi8, psidx.astype("int16").asi8) + self.assert_eq(pidx.astype("int8").asi8, psidx.astype("int8").asi8) + + # Integer with missing value + pidx = pd.Index([1, 2, None, 4, 5]) + psidx = ps.from_pandas(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(None, psidx.asi8) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + + # Datetime + pidx = pd.date_range(end="1/1/2018", periods=3) + psidx = ps.from_pandas(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq( + np.array([1514592000000000000, 1514678400000000000, 1514764800000000000]), + psidx.asi8, + ) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + + # Floating + pidx = pd.Index([1.0, 2.0, 3.0]) + psidx = ps.from_pandas(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(None, psidx.asi8) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + + # String + pidx = pd.Index(["a", "b", "c"]) + psidx = ps.from_pandas(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(None, psidx.asi8) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + + # Boolean + pidx = pd.Index([True, False, True, False]) + psidx = ps.from_pandas(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(None, psidx.asi8) + else: + self.assert_eq(pidx.asi8, psidx.asi8) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([(1, 2)]) + psmidx = ps.from_pandas(pmidx) + if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): + self.assert_eq(None, psmidx.asi8) + else: + self.assert_eq(pmidx.asi8, psmidx.asi8) + + def test_index_is_unique(self): + indexes = [("a", "b", "c"), ("a", "a", "c"), (1, 3, 3), (1, 2, 3)] + names = [None, "ks", "ks", None] + is_uniq = [True, False, False, True] + + for idx, name, expected in zip(indexes, names, is_uniq): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=pd.Index(idx, name=name)) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.is_unique, expected) + + def test_multiindex_is_unique(self): + indexes = [ + [list("abc"), list("edf")], + [list("aac"), list("edf")], + [list("aac"), list("eef")], + [[1, 4, 4], [4, 6, 6]], + ] + is_uniq = [True, True, False, False] + + for idx, expected in zip(indexes, is_uniq): + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=idx) + psdf = ps.from_pandas(pdf) + + self.assertEqual(psdf.index.is_unique, expected) + + def test_view(self): + pidx = pd.Index([1, 2, 3, 4], name="Koalas") + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.view(), psidx.view()) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(pmidx.view(), psmidx.view()) + + def test_insert(self): + # Integer + pidx = pd.Index([1, 2, 3], name="Koalas") + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.insert(1, 100), psidx.insert(1, 100)) + self.assert_eq(pidx.insert(-1, 100), psidx.insert(-1, 100)) + err_msg = "index 100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(100, 100) + err_msg = "index -100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(-100, 100) + + # Floating + pidx = pd.Index([1.0, 2.0, 3.0], name="Koalas") + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.insert(1, 100.0), psidx.insert(1, 100.0)) + self.assert_eq(pidx.insert(-1, 100.0), psidx.insert(-1, 100.0)) + err_msg = "index 100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(100, 100) + err_msg = "index -100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(-100, 100) + + # String + pidx = pd.Index(["a", "b", "c"], name="Koalas") + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.insert(1, "x"), psidx.insert(1, "x")) + self.assert_eq(pidx.insert(-1, "x"), psidx.insert(-1, "x")) + err_msg = "index 100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(100, "x") + err_msg = "index -100 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(-100, "x") + + # Boolean + pidx = pd.Index([True, False, True, False], name="Koalas") + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.insert(1, True), psidx.insert(1, True)) + self.assert_eq(pidx.insert(-1, True), psidx.insert(-1, True)) + err_msg = "index 100 is out of bounds for axis 0 with size 4" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(100, True) + err_msg = "index -100 is out of bounds for axis 0 with size 4" + with self.assertRaisesRegex(IndexError, err_msg): + psidx.insert(-100, True) + + # MultiIndex + pmidx = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z")], names=["Hello", "Koalas"] + ) + psmidx = ps.from_pandas(pmidx) + self.assert_eq(pmidx.insert(2, ("h", "j")), psmidx.insert(2, ("h", "j"))) + self.assert_eq(pmidx.insert(-1, ("h", "j")), psmidx.insert(-1, ("h", "j"))) + + err_msg = "index 4 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psmidx.insert(4, ("b", "y")) + + err_msg = "index -4 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + psmidx.insert(-4, ("b", "y")) + + def test_astype(self): + pidx = pd.Index([10, 20, 15, 30, 45], name="x") + psidx = ps.Index(pidx) + + self.assert_eq(psidx.astype(int), pidx.astype(int)) + self.assert_eq(psidx.astype(np.int8), pidx.astype(np.int8)) + self.assert_eq(psidx.astype(np.int16), pidx.astype(np.int16)) + self.assert_eq(psidx.astype(np.int32), pidx.astype(np.int32)) + self.assert_eq(psidx.astype(np.int64), pidx.astype(np.int64)) + self.assert_eq(psidx.astype(np.byte), pidx.astype(np.byte)) + self.assert_eq(psidx.astype("int"), pidx.astype("int")) + self.assert_eq(psidx.astype("int8"), pidx.astype("int8")) + self.assert_eq(psidx.astype("int16"), pidx.astype("int16")) + self.assert_eq(psidx.astype("int32"), pidx.astype("int32")) + self.assert_eq(psidx.astype("int64"), pidx.astype("int64")) + self.assert_eq(psidx.astype("b"), pidx.astype("b")) + self.assert_eq(psidx.astype("byte"), pidx.astype("byte")) + self.assert_eq(psidx.astype("i"), pidx.astype("i")) + self.assert_eq(psidx.astype("long"), pidx.astype("long")) + self.assert_eq(psidx.astype("short"), pidx.astype("short")) + self.assert_eq(psidx.astype(np.float32), pidx.astype(np.float32)) + self.assert_eq(psidx.astype(np.float64), pidx.astype(np.float64)) + self.assert_eq(psidx.astype("float"), pidx.astype("float")) + self.assert_eq(psidx.astype("float32"), pidx.astype("float32")) + self.assert_eq(psidx.astype("float64"), pidx.astype("float64")) + self.assert_eq(psidx.astype("double"), pidx.astype("double")) + self.assert_eq(psidx.astype("f"), pidx.astype("f")) + self.assert_eq(psidx.astype(bool), pidx.astype(bool)) + self.assert_eq(psidx.astype("bool"), pidx.astype("bool")) + self.assert_eq(psidx.astype("?"), pidx.astype("?")) + self.assert_eq(psidx.astype(np.unicode_), pidx.astype(np.unicode_)) + self.assert_eq(psidx.astype("str"), pidx.astype("str")) + self.assert_eq(psidx.astype("U"), pidx.astype("U")) + + pidx = pd.Index([10, 20, 15, 30, 45, None], name="x") + psidx = ps.Index(pidx) + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq(psidx.astype(bool), pidx.astype(bool)) + self.assert_eq(psidx.astype(str), pidx.astype(str)) + else: + self.assert_eq( + psidx.astype(bool), ps.Index([True, True, True, True, True, True], name="x") + ) + self.assert_eq( + psidx.astype(str), + ps.Index(["10.0", "20.0", "15.0", "30.0", "45.0", "nan"], name="x"), + ) + + pidx = pd.Index(["hi", "hi ", " ", " \t", "", None], name="x") + psidx = ps.Index(pidx) + + self.assert_eq(psidx.astype(bool), pidx.astype(bool)) + self.assert_eq(psidx.astype(str), pidx.astype(str)) + + pidx = pd.Index([True, False, None], name="x") + psidx = ps.Index(pidx) + + self.assert_eq(psidx.astype(bool), pidx.astype(bool)) + + pidx = pd.Index(["2020-10-27"], name="x") + psidx = ps.Index(pidx) + + self.assert_eq(psidx.astype("datetime64[ns]"), pidx.astype("datetime64[ns]")) + + with self.assertRaisesRegex(TypeError, "not understood"): + psidx.astype("int63") + + def test_to_list(self): + # Index + pidx = pd.Index([1, 2, 3, 4, 5]) + psidx = ps.from_pandas(pidx) + # MultiIndex + tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "green")] + pmidx = pd.MultiIndex.from_tuples(tuples) + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(psidx.tolist(), pidx.tolist()) + self.assert_eq(psmidx.tolist(), pmidx.tolist()) + + def test_index_ops(self): + pidx = pd.Index([1, 2, 3, 4, 5]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx * 100 + psidx * 10 + psidx, pidx * 100 + pidx * 10 + pidx) + + pidx = pd.Index([1, 2, 3, 4, 5], name="a") + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx * 100 + psidx * 10 + psidx, pidx * 100 + pidx * 10 + pidx) + + pdf = pd.DataFrame( + index=pd.MultiIndex.from_tuples([(1, 2), (3, 4), (5, 6)], names=["a", "b"]) + ) + psdf = ps.from_pandas(pdf) + + pidx1 = pdf.index.get_level_values(0) + pidx2 = pdf.index.get_level_values(1) + psidx1 = psdf.index.get_level_values(0) + psidx2 = psdf.index.get_level_values(1) + + self.assert_eq(psidx1 * 10 + psidx2, pidx1 * 10 + pidx2) + + def test_factorize(self): + pidx = pd.Index(["a", "b", "a", "b"]) + psidx = ps.from_pandas(pidx) + pcodes, puniques = pidx.factorize(sort=True) + kcodes, kuniques = psidx.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pmidx = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) + psmidx = ps.from_pandas(pmidx) + + self.assertRaises(PandasNotImplementedError, lambda: psmidx.factorize()) + + def test_map(self): + pidx = pd.Index([1, 2, 3]) + psidx = ps.from_pandas(pidx) + + # Apply dict + self.assert_eq( + pidx.map({1: "one", 2: "two", 3: "three"}), + psidx.map({1: "one", 2: "two", 3: "three"}), + ) + self.assert_eq( + pidx.map({1: "one", 2: "two"}), + psidx.map({1: "one", 2: "two"}), + ) + self.assert_eq( + pidx.map({1: "one", 2: "two"}, na_action="ignore"), + psidx.map({1: "one", 2: "two"}, na_action="ignore"), + ) + self.assert_eq( + pidx.map({1: 10, 2: 20}), + psidx.map({1: 10, 2: 20}), + ) + self.assert_eq( + (pidx + 1).map({1: 10, 2: 20}), + (psidx + 1).map({1: 10, 2: 20}), + ) + + # Apply lambda + self.assert_eq( + pidx.map(lambda id: id + 1), + psidx.map(lambda id: id + 1), + ) + self.assert_eq( + pidx.map(lambda id: id + 1.1), + psidx.map(lambda id: id + 1.1), + ) + self.assert_eq( + pidx.map(lambda id: "{id} + 1".format(id=id)), + psidx.map(lambda id: "{id} + 1".format(id=id)), + ) + self.assert_eq( + (pidx + 1).map(lambda id: "{id} + 1".format(id=id)), + (psidx + 1).map(lambda id: "{id} + 1".format(id=id)), + ) + + # Apply series + pser = pd.Series(["one", "two", "three"], index=[1, 2, 3]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + pser = pd.Series(["one", "two", "three"]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + self.assert_eq( + pidx.map(pser, na_action="ignore"), + psidx.map(pser, na_action="ignore"), + ) + pser = pd.Series([1, 2, 3]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + self.assert_eq( + (pidx + 1).map(pser), + (psidx + 1).map(pser), + ) + + self.assertRaises( + TypeError, + lambda: psidx.map({1: 1, 2: 2.0, 3: "three"}), + ) + + def test_multiindex_equal_levels(self): + pmidx1 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + pmidx2 = pd.MultiIndex.from_tuples([("b", "y"), ("a", "x"), ("c", "z")]) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx2 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "j")]) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx2 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("a", "x")]) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx2 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y")]) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx2 = pd.MultiIndex.from_tuples([("a", "y"), ("b", "x"), ("c", "z")]) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx1 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z"), ("a", "y")]) + pmidx2 = pd.MultiIndex.from_tuples([("a", "y"), ("b", "x"), ("c", "z"), ("c", "x")]) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + pmidx1 = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + pmidx2 = pd.MultiIndex.from_tuples([("a", "x", "q"), ("b", "y", "w"), ("c", "z", "e")]) + psmidx1 = ps.from_pandas(pmidx1) + psmidx2 = ps.from_pandas(pmidx2) + self.assert_eq(pmidx1.equal_levels(pmidx2), psmidx1.equal_levels(psmidx2)) + + def test_to_numpy(self): + pidx = pd.Index([1, 2, 3, 4]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.to_numpy(copy=True), psidx.to_numpy(copy=True)) + + def test_drop_level(self): + tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "green")] + pmidx = pd.MultiIndex.from_tuples(tuples) + psmidx = ps.from_pandas(pmidx) + + with self.assertRaisesRegex( + IndexError, "Too many levels: Index has only 2 levels, -3 is not a valid level number" + ): + psmidx.droplevel(-3) + + def test_multi_index_nunique(self): + tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "green")] + pmidx = pd.MultiIndex.from_tuples(tuples) + psmidx = ps.from_pandas(pmidx) + + with self.assertRaisesRegex(NotImplementedError, "nunique is not defined for MultiIndex"): + psmidx.nunique() + + +class IndexesTests(IndexesTestsMixin, ComparisonTestBase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_base import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_category.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_category.py new file mode 100644 index 0000000..ffffae8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_category.py @@ -0,0 +1,484 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import pandas as pd +from pandas.api.types import CategoricalDtype + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class CategoricalIndexTestsMixin: + def test_categorical_index(self): + pidx = pd.CategoricalIndex([1, 2, 3]) + psidx = ps.CategoricalIndex([1, 2, 3]) + + self.assert_eq(psidx, pidx) + self.assert_eq(psidx.categories, pidx.categories) + self.assert_eq(psidx.codes, pd.Index(pidx.codes)) + self.assert_eq(psidx.ordered, pidx.ordered) + + pidx = pd.Index([1, 2, 3], dtype="category") + psidx = ps.Index([1, 2, 3], dtype="category") + + self.assert_eq(psidx, pidx) + self.assert_eq(psidx.categories, pidx.categories) + self.assert_eq(psidx.codes, pd.Index(pidx.codes)) + self.assert_eq(psidx.ordered, pidx.ordered) + + pdf = pd.DataFrame( + { + "a": pd.Categorical([1, 2, 3, 1, 2, 3]), + "b": pd.Categorical(["a", "b", "c", "a", "b", "c"], categories=["c", "b", "a"]), + }, + index=pd.Categorical([10, 20, 30, 20, 30, 10], categories=[30, 10, 20], ordered=True), + ) + psdf = ps.from_pandas(pdf) + + pidx = pdf.set_index("b").index + psidx = psdf.set_index("b").index + + self.assert_eq(psidx, pidx) + self.assert_eq(psidx.categories, pidx.categories) + self.assert_eq(psidx.codes, pd.Index(pidx.codes)) + self.assert_eq(psidx.ordered, pidx.ordered) + + pidx = pdf.set_index(["a", "b"]).index.get_level_values(0) + psidx = psdf.set_index(["a", "b"]).index.get_level_values(0) + + self.assert_eq(psidx, pidx) + self.assert_eq(psidx.categories, pidx.categories) + self.assert_eq(psidx.codes, pd.Index(pidx.codes)) + self.assert_eq(psidx.ordered, pidx.ordered) + + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.CategoricalIndex([1, 2, 3], name=[(1, 2, 3)]) + with self.assertRaisesRegexp( + TypeError, "Cannot perform 'all' with this index type: CategoricalIndex" + ): + ps.CategoricalIndex([1, 2, 3]).all() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43568): Enable CategoricalIndexTests.test_categories_setter for pandas 2.0.0.", + ) + def test_categories_setter(self): + pdf = pd.DataFrame( + { + "a": pd.Categorical([1, 2, 3, 1, 2, 3]), + "b": pd.Categorical(["a", "b", "c", "a", "b", "c"], categories=["c", "b", "a"]), + }, + index=pd.Categorical([10, 20, 30, 20, 30, 10], categories=[30, 10, 20], ordered=True), + ) + psdf = ps.from_pandas(pdf) + + pidx = pdf.index + psidx = psdf.index + + pidx.categories = ["z", "y", "x"] + psidx.categories = ["z", "y", "x"] + # Pandas deprecated all the in-place category-setting behaviors, dtypes also not be + # refreshed in categories.setter since Pandas 1.4+, we should also consider to clean up + # this test when in-place category-setting removed: + # https://github.com/pandas-dev/pandas/issues/46820 + if LooseVersion("1.4") >= LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(pidx, psidx) + self.assert_eq(pdf, psdf) + else: + pidx = pidx.set_categories(pidx.categories) + pdf.index = pidx + self.assert_eq(pidx, psidx) + self.assert_eq(pdf, psdf) + + with self.assertRaises(ValueError): + psidx.categories = [1, 2, 3, 4] + + def test_add_categories(self): + pidx = pd.CategoricalIndex([1, 2, 3], categories=[3, 2, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.add_categories(4), psidx.add_categories(4)) + self.assert_eq(pidx.add_categories([4, 5]), psidx.add_categories([4, 5])) + self.assert_eq(pidx.add_categories([]), psidx.add_categories([])) + + self.assertRaises(ValueError, lambda: psidx.add_categories(4, inplace=True)) + self.assertRaises(ValueError, lambda: psidx.add_categories(3)) + self.assertRaises(ValueError, lambda: psidx.add_categories([4, 4])) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43633): Enable CategoricalIndexTests.test_remove_categories for pandas 2.0.0.", + ) + def test_remove_categories(self): + pidx = pd.CategoricalIndex([1, 2, 3], categories=[3, 2, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.remove_categories(2), psidx.remove_categories(2)) + self.assert_eq(pidx.remove_categories([1, 3]), psidx.remove_categories([1, 3])) + self.assert_eq(pidx.remove_categories([]), psidx.remove_categories([])) + self.assert_eq(pidx.remove_categories([2, 2]), psidx.remove_categories([2, 2])) + self.assert_eq(pidx.remove_categories([1, 2, 3]), psidx.remove_categories([1, 2, 3])) + self.assert_eq(pidx.remove_categories(None), psidx.remove_categories(None)) + self.assert_eq(pidx.remove_categories([None]), psidx.remove_categories([None])) + + self.assertRaises(ValueError, lambda: psidx.remove_categories(4, inplace=True)) + self.assertRaises(ValueError, lambda: psidx.remove_categories(4)) + self.assertRaises(ValueError, lambda: psidx.remove_categories([4, None])) + + def test_remove_unused_categories(self): + pidx = pd.CategoricalIndex([1, 4, 5, 3], categories=[4, 3, 2, 1]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.remove_unused_categories(), psidx.remove_unused_categories()) + + self.assertRaises(ValueError, lambda: psidx.remove_unused_categories(inplace=True)) + + def test_reorder_categories(self): + pidx = pd.CategoricalIndex([1, 2, 3]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.reorder_categories([1, 2, 3]), psidx.reorder_categories([1, 2, 3])) + self.assert_eq( + pidx.reorder_categories([1, 2, 3], ordered=True), + psidx.reorder_categories([1, 2, 3], ordered=True), + ) + self.assert_eq(pidx.reorder_categories([3, 2, 1]), psidx.reorder_categories([3, 2, 1])) + self.assert_eq( + pidx.reorder_categories([3, 2, 1], ordered=True), + psidx.reorder_categories([3, 2, 1], ordered=True), + ) + + self.assertRaises(ValueError, lambda: psidx.reorder_categories([1, 2, 3], inplace=True)) + self.assertRaises(ValueError, lambda: psidx.reorder_categories([1, 2])) + self.assertRaises(ValueError, lambda: psidx.reorder_categories([1, 2, 4])) + self.assertRaises(ValueError, lambda: psidx.reorder_categories([1, 2, 2])) + self.assertRaises(TypeError, lambda: psidx.reorder_categories(1)) + + def test_as_ordered_unordered(self): + pidx = pd.CategoricalIndex(["x", "y", "z"], categories=["z", "y", "x"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(pidx.as_ordered(), psidx.as_ordered()) + self.assert_eq(pidx.as_unordered(), psidx.as_unordered()) + + self.assertRaises(ValueError, lambda: psidx.as_ordered(inplace=True)) + self.assertRaises(ValueError, lambda: psidx.as_unordered(inplace=True)) + + def test_astype(self): + pidx = pd.Index(["a", "b", "c"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.astype("category"), pidx.astype("category")) + self.assert_eq( + psidx.astype(CategoricalDtype(["c", "a", "b"])), + pidx.astype(CategoricalDtype(["c", "a", "b"])), + ) + + pcidx = pidx.astype(CategoricalDtype(["c", "a", "b"])) + pscidx = psidx.astype(CategoricalDtype(["c", "a", "b"])) + + self.assert_eq(pscidx.astype("category"), pcidx.astype("category")) + + # CategoricalDtype is not updated if the dtype is same from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + pscidx.astype(CategoricalDtype(["b", "c", "a"])), + pcidx.astype(CategoricalDtype(["b", "c", "a"])), + ) + else: + self.assert_eq( + pscidx.astype(CategoricalDtype(["b", "c", "a"])), + pcidx, + ) + + self.assert_eq(pscidx.astype(str), pcidx.astype(str)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43567): Enable CategoricalIndexTests.test_factorize for pandas 2.0.0.", + ) + def test_factorize(self): + pidx = pd.CategoricalIndex([1, 2, 3, None]) + psidx = ps.from_pandas(pidx) + + pcodes, puniques = pidx.factorize() + kcodes, kuniques = psidx.factorize() + + self.assert_eq(kcodes.tolist(), pcodes.tolist()) + self.assert_eq(kuniques, puniques) + + pcodes, puniques = pidx.factorize(na_sentinel=-2) + kcodes, kuniques = psidx.factorize(na_sentinel=-2) + + self.assert_eq(kcodes.tolist(), pcodes.tolist()) + self.assert_eq(kuniques, puniques) + + def test_append(self): + pidx1 = pd.CategoricalIndex(["x", "y", "z"], categories=["z", "y", "x", "w"]) + pidx2 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x", "w"]) + pidx3 = pd.Index(["y", "x", "w", "z"]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + psidx3 = ps.from_pandas(pidx3) + + self.assert_eq(psidx1.append(psidx2), pidx1.append(pidx2)) + if LooseVersion(pd.__version__) >= LooseVersion("1.5.0"): + self.assert_eq( + psidx1.append(psidx3.astype("category")), pidx1.append(pidx3.astype("category")) + ) + else: + expected_result = ps.CategoricalIndex( + ["x", "y", "z", "y", "x", "w", "z"], + categories=["z", "y", "x", "w"], + ordered=False, + dtype="category", + ) + self.assert_eq(psidx1.append(psidx3.astype("category")), expected_result) + + # TODO: append non-categorical or categorical with a different category + self.assertRaises(NotImplementedError, lambda: psidx1.append(psidx3)) + + pidx4 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x"]) + psidx4 = ps.from_pandas(pidx4) + self.assertRaises(NotImplementedError, lambda: psidx1.append(psidx4)) + + def test_union(self): + pidx1 = pd.CategoricalIndex(["x", "y", "z"], categories=["z", "y", "x", "w"]) + pidx2 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x", "w"]) + pidx3 = pd.Index(["y", "x", "w", "z"]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + psidx3 = ps.from_pandas(pidx3) + + self.assert_eq(psidx1.union(psidx2), pidx1.union(pidx2)) + self.assert_eq( + psidx1.union(psidx3.astype("category")), pidx1.union(pidx3.astype("category")) + ) + + # TODO: union non-categorical or categorical with a different category + self.assertRaises(NotImplementedError, lambda: psidx1.union(psidx3)) + + pidx4 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x"]) + psidx4 = ps.from_pandas(pidx4) + self.assertRaises(NotImplementedError, lambda: psidx1.union(psidx4)) + + def test_intersection(self): + pidx1 = pd.CategoricalIndex(["x", "y", "z"], categories=["z", "y", "x", "w"]) + pidx2 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x", "w"]) + pidx3 = pd.Index(["y", "x", "w", "z"]) + psidx1 = ps.from_pandas(pidx1) + psidx2 = ps.from_pandas(pidx2) + psidx3 = ps.from_pandas(pidx3) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq( + psidx1.intersection(psidx2).sort_values(), pidx1.intersection(pidx2).sort_values() + ) + self.assert_eq( + psidx1.intersection(psidx3.astype("category")).sort_values(), + pidx1.intersection(pidx3.astype("category")).sort_values(), + ) + else: + self.assert_eq( + psidx1.intersection(psidx2).sort_values(), + pidx1.intersection(pidx2).set_categories(pidx1.categories).sort_values(), + ) + self.assert_eq( + psidx1.intersection(psidx3.astype("category")).sort_values(), + pidx1.intersection(pidx3.astype("category")) + .set_categories(pidx1.categories) + .sort_values(), + ) + + # TODO: intersection non-categorical or categorical with a different category + self.assertRaises(NotImplementedError, lambda: psidx1.intersection(psidx3)) + + pidx4 = pd.CategoricalIndex(["y", "x", "w"], categories=["z", "y", "x"]) + psidx4 = ps.from_pandas(pidx4) + self.assertRaises(NotImplementedError, lambda: psidx1.intersection(psidx4)) + + def test_insert(self): + pidx = pd.CategoricalIndex(["x", "y", "z"], categories=["z", "y", "x", "w"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx.insert(1, "w"), pidx.insert(1, "w")) + + def test_rename_categories(self): + pidx = pd.CategoricalIndex(["a", "b", "c", "d"]) + psidx = ps.from_pandas(pidx) + self.assert_eq(pidx.rename_categories([0, 1, 3, 2]), psidx.rename_categories([0, 1, 3, 2])) + self.assert_eq( + pidx.rename_categories({"a": "A", "c": "C"}), + psidx.rename_categories({"a": "A", "c": "C"}), + ) + self.assert_eq( + pidx.rename_categories(lambda x: x.upper()), + psidx.rename_categories(lambda x: x.upper()), + ) + self.assertRaises( + TypeError, + lambda: psidx.rename_categories(None), + ) + self.assertRaises( + TypeError, + lambda: psidx.rename_categories(1), + ) + self.assertRaises( + TypeError, + lambda: psidx.rename_categories("x"), + ) + self.assertRaises( + ValueError, + lambda: psidx.rename_categories({"b": "B", "c": "C"}, inplace=True), + ) + + def test_set_categories(self): + pidx = pd.CategoricalIndex(["a", "b", "c", "d"]) + psidx = ps.from_pandas(pidx) + + self.assert_eq( + pidx.set_categories(["a", "c", "b", "o"]), + psidx.set_categories(["a", "c", "b", "o"]), + ) + self.assert_eq( + pidx.set_categories(["a", "c", "b"]), + psidx.set_categories(["a", "c", "b"]), + ) + self.assert_eq( + pidx.set_categories(["a", "c", "b", "d", "e"]), + psidx.set_categories(["a", "c", "b", "d", "e"]), + ) + + self.assert_eq( + pidx.set_categories([0, 1, 3, 2], rename=True), + psidx.set_categories([0, 1, 3, 2], rename=True), + ) + self.assert_eq( + pidx.set_categories([0, 1, 3], rename=True), + psidx.set_categories([0, 1, 3], rename=True), + ) + self.assert_eq( + pidx.set_categories([0, 1, 3, 2, 4], rename=True), + psidx.set_categories([0, 1, 3, 2, 4], rename=True), + ) + + self.assert_eq( + pidx.set_categories(["a", "c", "b", "o"], ordered=True), + psidx.set_categories(["a", "c", "b", "o"], ordered=True), + ) + self.assert_eq( + pidx.set_categories(["a", "c", "b"], ordered=True), + psidx.set_categories(["a", "c", "b"], ordered=True), + ) + self.assert_eq( + pidx.set_categories(["a", "c", "b", "d", "e"], ordered=True), + psidx.set_categories(["a", "c", "b", "d", "e"], ordered=True), + ) + + self.assertRaisesRegex( + ValueError, + "cannot use inplace with CategoricalIndex", + lambda: psidx.set_categories(["a", "c", "b", "o"], inplace=True), + ) + + def test_map(self): + pidxs = [pd.CategoricalIndex([1, 2, 3]), pd.CategoricalIndex([1, 2, 3], ordered=True)] + psidxs = [ps.from_pandas(pidx) for pidx in pidxs] + + for pidx, psidx in zip(pidxs, psidxs): + + # Apply dict + self.assert_eq( + pidx.map({1: "one", 2: "two", 3: "three"}), + psidx.map({1: "one", 2: "two", 3: "three"}), + ) + self.assert_eq( + pidx.map({1: "one", 2: "two", 3: "one"}), + psidx.map({1: "one", 2: "two", 3: "one"}), + ) + self.assert_eq( + pidx.map({1: "one", 2: "two"}), + psidx.map({1: "one", 2: "two"}), + ) + self.assert_eq( + pidx.map({1: "one", 2: "two"}), + psidx.map({1: "one", 2: "two"}), + ) + self.assert_eq( + pidx.map({1: 10, 2: 20}), + psidx.map({1: 10, 2: 20}), + ) + + # Apply lambda + self.assert_eq( + pidx.map(lambda id: id + 1), + psidx.map(lambda id: id + 1), + ) + self.assert_eq( + pidx.map(lambda id: id + 1.1), + psidx.map(lambda id: id + 1.1), + ) + self.assert_eq( + pidx.map(lambda id: "{id} + 1".format(id=id)), + psidx.map(lambda id: "{id} + 1".format(id=id)), + ) + + # Apply series + pser = pd.Series(["one", "two", "three"], index=[1, 2, 3]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + pser = pd.Series(["one", "two", "three"]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + pser = pd.Series([1, 2, 3]) + self.assert_eq( + pidx.map(pser), + psidx.map(pser), + ) + + self.assertRaises( + TypeError, + lambda: psidx.map({1: 1, 2: 2.0, 3: "three"}), + ) + + +class CategoricalIndexTests(CategoricalIndexTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.indexes.test_category import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_datetime.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_datetime.py new file mode 100644 index 0000000..4fb3561 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_datetime.py @@ -0,0 +1,276 @@ +# +# 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 datetime +import unittest + +from distutils.version import LooseVersion + +import pandas as pd + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class DatetimeIndexTestsMixin: + @property + def fixed_freqs(self): + return [ + "D", + "H", + "T", # min + "S", + "L", # ms + "U", # us + # 'N' not supported + ] + + @property + def non_fixed_freqs(self): + return ["W", "Q"] + + @property + def pidxs(self): + return [ + pd.DatetimeIndex([0]), + pd.DatetimeIndex(["2004-01-01", "2002-12-31", "2000-04-01"]), + ] + [ + pd.date_range("2000-01-01", periods=3, freq=freq) + for freq in (self.fixed_freqs + self.non_fixed_freqs) + ] + + @property + def psidxs(self): + return [ps.from_pandas(pidx) for pidx in self.pidxs] + + @property + def idx_pairs(self): + return list(zip(self.psidxs, self.pidxs)) + + def _disallow_nanoseconds(self, f): + self.assertRaises(ValueError, lambda: f(freq="ns")) + self.assertRaises(ValueError, lambda: f(freq="N")) + + def test_datetime_index(self): + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.DatetimeIndex(["2004-01-01", "2002-12-31", "2000-04-01"], name=[(1, 2)]) + with self.assertRaisesRegexp( + TypeError, "Cannot perform 'all' with this index type: DatetimeIndex" + ): + ps.DatetimeIndex(["2004-01-01", "2002-12-31", "2000-04-01"]).all() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43608): Enable DatetimeIndexTests.test_properties for pandas 2.0.0.", + ) + def test_properties(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq(psidx.year, pidx.year) + self.assert_eq(psidx.month, pidx.month) + self.assert_eq(psidx.day, pidx.day) + self.assert_eq(psidx.hour, pidx.hour) + self.assert_eq(psidx.minute, pidx.minute) + self.assert_eq(psidx.second, pidx.second) + self.assert_eq(psidx.microsecond, pidx.microsecond) + self.assert_eq(psidx.week, pidx.week) + self.assert_eq(psidx.weekofyear, pidx.weekofyear) + self.assert_eq(psidx.dayofweek, pidx.dayofweek) + self.assert_eq(psidx.weekday, pidx.weekday) + self.assert_eq(psidx.dayofyear, pidx.dayofyear) + self.assert_eq(psidx.quarter, pidx.quarter) + self.assert_eq(psidx.daysinmonth, pidx.daysinmonth) + self.assert_eq(psidx.days_in_month, pidx.days_in_month) + self.assert_eq(psidx.is_month_start, pd.Index(pidx.is_month_start)) + self.assert_eq(psidx.is_month_end, pd.Index(pidx.is_month_end)) + self.assert_eq(psidx.is_quarter_start, pd.Index(pidx.is_quarter_start)) + self.assert_eq(psidx.is_quarter_end, pd.Index(pidx.is_quarter_end)) + self.assert_eq(psidx.is_year_start, pd.Index(pidx.is_year_start)) + self.assert_eq(psidx.is_year_end, pd.Index(pidx.is_year_end)) + self.assert_eq(psidx.is_leap_year, pd.Index(pidx.is_leap_year)) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2.0"): + self.assert_eq(psidx.day_of_year, pidx.day_of_year) + self.assert_eq(psidx.day_of_week, pidx.day_of_week) + + def test_ceil(self): + for psidx, pidx in self.idx_pairs: + for freq in self.fixed_freqs: + self.assert_eq(psidx.ceil(freq), pidx.ceil(freq)) + + self._disallow_nanoseconds(self.psidxs[0].ceil) + + def test_floor(self): + for psidx, pidx in self.idx_pairs: + for freq in self.fixed_freqs: + self.assert_eq(psidx.floor(freq), pidx.floor(freq)) + + self._disallow_nanoseconds(self.psidxs[0].floor) + + def test_round(self): + for psidx, pidx in self.idx_pairs: + for freq in self.fixed_freqs: + self.assert_eq(psidx.round(freq), pidx.round(freq)) + + self._disallow_nanoseconds(self.psidxs[0].round) + + def test_day_name(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq(psidx.day_name(), pidx.day_name()) + + def test_month_name(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq(psidx.month_name(), pidx.month_name()) + + def test_normalize(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq(psidx.normalize(), pidx.normalize()) + + def test_strftime(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq( + psidx.strftime(date_format="%B %d, %Y"), pidx.strftime(date_format="%B %d, %Y") + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43644): Enable DatetimeIndexTests.test_indexer_between_time " + "for pandas 2.0.0.", + ) + def test_indexer_between_time(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq( + psidx.indexer_between_time("00:00:00", "00:01:00").sort_values(), + pd.Index(pidx.indexer_between_time("00:00:00", "00:01:00")), + ) + + self.assert_eq( + psidx.indexer_between_time( + datetime.time(0, 0, 0), datetime.time(0, 1, 0) + ).sort_values(), + pd.Index(pidx.indexer_between_time(datetime.time(0, 0, 0), datetime.time(0, 1, 0))), + ) + + self.assert_eq( + psidx.indexer_between_time("00:00:00", "00:01:00", True, False).sort_values(), + pd.Index(pidx.indexer_between_time("00:00:00", "00:01:00", True, False)), + ) + + self.assert_eq( + psidx.indexer_between_time("00:00:00", "00:01:00", False, True).sort_values(), + pd.Index(pidx.indexer_between_time("00:00:00", "00:01:00", False, True)), + ) + + self.assert_eq( + psidx.indexer_between_time("00:00:00", "00:01:00", False, False).sort_values(), + pd.Index(pidx.indexer_between_time("00:00:00", "00:01:00", False, False)), + ) + + self.assert_eq( + psidx.indexer_between_time("00:00:00", "00:01:00", True, True).sort_values(), + pd.Index(pidx.indexer_between_time("00:00:00", "00:01:00", True, True)), + ) + + def test_indexer_at_time(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq( + psidx.indexer_at_time("00:00:00").sort_values(), + pd.Index(pidx.indexer_at_time("00:00:00")), + ) + + self.assert_eq( + psidx.indexer_at_time(datetime.time(0, 1, 0)).sort_values(), + pd.Index(pidx.indexer_at_time(datetime.time(0, 1, 0))), + ) + + self.assert_eq( + psidx.indexer_at_time("00:00:01").sort_values(), + pd.Index(pidx.indexer_at_time("00:00:01")), + ) + + self.assertRaises( + NotImplementedError, + lambda: ps.DatetimeIndex([0]).indexer_at_time("00:00:00", asof=True), + ) + + def test_arithmetic_op_exceptions(self): + for psidx, pidx in self.idx_pairs: + py_datetime = pidx.to_pydatetime() + for other in [1, 0.1, psidx, psidx.to_series().reset_index(drop=True), py_datetime]: + expected_err_msg = "Addition can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx + other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other + psidx) + + expected_err_msg = "Multiplication can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx * other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other * psidx) + + expected_err_msg = "True division can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx / other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other / psidx) + + expected_err_msg = "Floor division can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx // other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other // psidx) + + expected_err_msg = "Modulo can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx % other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other % psidx) + + expected_err_msg = "Datetime subtraction can only be applied to datetime series." + + for other in [1, 0.1]: + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx - other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other - psidx) + + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psidx - other) + self.assertRaises(NotImplementedError, lambda: py_datetime - psidx) + + def test_map(self): + for psidx, pidx in self.idx_pairs: + self.assert_eq(psidx.map(lambda x: x.normalize()), pidx.map(lambda x: x.normalize())) + self.assert_eq( + psidx.map(lambda x: x.strftime("%B %d, %Y, %r")), + pidx.map(lambda x: x.strftime("%B %d, %Y, %r")), + ) + + pidx = pd.date_range(start="2020-08-08", end="2020-08-10") + psidx = ps.from_pandas(pidx) + mapper_dict = { + datetime.datetime(2020, 8, 8): datetime.datetime(2021, 8, 8), + datetime.datetime(2020, 8, 9): datetime.datetime(2021, 8, 9), + } + self.assert_eq(psidx.map(mapper_dict), pidx.map(mapper_dict)) + + mapper_pser = pd.Series([1, 2, 3], index=pidx) + self.assert_eq(psidx.map(mapper_pser), pidx.map(mapper_pser)) + + +class DatetimeIndexTests(DatetimeIndexTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.indexes.test_datetime import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_indexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_indexing.py new file mode 100644 index 0000000..64fc753 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_indexing.py @@ -0,0 +1,432 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Indexing, Iteration' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#indexing-iteration +class FrameIndexingMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_head(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf.head(2), pdf.head(2)) + self.assert_eq(psdf.head(3), pdf.head(3)) + self.assert_eq(psdf.head(0), pdf.head(0)) + self.assert_eq(psdf.head(-3), pdf.head(-3)) + self.assert_eq(psdf.head(-10), pdf.head(-10)) + with option_context("compute.ordered_head", True): + self.assert_eq(psdf.head(), pdf.head()) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43559): Enable DataFrameSlowTests.test_iteritems for pandas 2.0.0.", + ) + def test_iteritems(self): + pdf = pd.DataFrame( + {"species": ["bear", "bear", "marsupial"], "population": [1864, 22000, 80000]}, + index=["panda", "polar", "koala"], + columns=["species", "population"], + ) + psdf = ps.from_pandas(pdf) + + for (p_name, p_items), (k_name, k_items) in zip(pdf.iteritems(), psdf.iteritems()): + self.assert_eq(p_name, k_name) + self.assert_eq(p_items, k_items) + + def test_keys(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.keys(), pdf.keys()) + + def test_tail(self): + pdf = pd.DataFrame({"x": range(1000)}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.tail(), psdf.tail()) + self.assert_eq(pdf.tail(10), psdf.tail(10)) + self.assert_eq(pdf.tail(-990), psdf.tail(-990)) + self.assert_eq(pdf.tail(0), psdf.tail(0)) + self.assert_eq(pdf.tail(-1001), psdf.tail(-1001)) + self.assert_eq(pdf.tail(1001), psdf.tail(1001)) + self.assert_eq((pdf + 1).tail(), (psdf + 1).tail()) + self.assert_eq((pdf + 1).tail(10), (psdf + 1).tail(10)) + self.assert_eq((pdf + 1).tail(-990), (psdf + 1).tail(-990)) + self.assert_eq((pdf + 1).tail(0), (psdf + 1).tail(0)) + self.assert_eq((pdf + 1).tail(-1001), (psdf + 1).tail(-1001)) + self.assert_eq((pdf + 1).tail(1001), (psdf + 1).tail(1001)) + with self.assertRaisesRegex(TypeError, "bad operand type for unary -: 'str'"): + psdf.tail("10") + + def test_xs(self): + d = { + "num_legs": [4, 4, 2, 2], + "num_wings": [0, 0, 2, 2], + "class": ["mammal", "mammal", "mammal", "bird"], + "animal": ["cat", "dog", "bat", "penguin"], + "locomotion": ["walks", "walks", "flies", "walks"], + } + pdf = pd.DataFrame(data=d) + pdf = pdf.set_index(["class", "animal", "locomotion"]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.xs("mammal"), pdf.xs("mammal")) + self.assert_eq(psdf.xs(("mammal",)), pdf.xs(("mammal",))) + self.assert_eq(psdf.xs(("mammal", "dog", "walks")), pdf.xs(("mammal", "dog", "walks"))) + self.assert_eq( + ps.concat([psdf, psdf]).xs(("mammal", "dog", "walks")), + pd.concat([pdf, pdf]).xs(("mammal", "dog", "walks")), + ) + self.assert_eq(psdf.xs("cat", level=1), pdf.xs("cat", level=1)) + self.assert_eq(psdf.xs("flies", level=2), pdf.xs("flies", level=2)) + self.assert_eq(psdf.xs("mammal", level=-3), pdf.xs("mammal", level=-3)) + + msg = 'axis should be either 0 or "index" currently.' + with self.assertRaisesRegex(NotImplementedError, msg): + psdf.xs("num_wings", axis=1) + with self.assertRaises(KeyError): + psdf.xs(("mammal", "dog", "walk")) + msg = r"'Key length \(4\) exceeds index depth \(3\)'" + with self.assertRaisesRegex(KeyError, msg): + psdf.xs(("mammal", "dog", "walks", "foo")) + msg = "'key' should be a scalar value or tuple that contains scalar values" + with self.assertRaisesRegex(TypeError, msg): + psdf.xs(["mammal", "dog", "walks", "foo"]) + + self.assertRaises(IndexError, lambda: psdf.xs("foo", level=-4)) + self.assertRaises(IndexError, lambda: psdf.xs("foo", level=3)) + + self.assertRaises(KeyError, lambda: psdf.xs(("dog", "walks"), level=1)) + + # non-string names + pdf = pd.DataFrame(data=d) + pdf = pdf.set_index(["class", "animal", "num_legs", "num_wings"]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.xs(("mammal", "dog", 4)), pdf.xs(("mammal", "dog", 4))) + self.assert_eq(psdf.xs(2, level=2), pdf.xs(2, level=2)) + + self.assert_eq((psdf + "a").xs(("mammal", "dog", 4)), (pdf + "a").xs(("mammal", "dog", 4))) + self.assert_eq((psdf + "a").xs(2, level=2), (pdf + "a").xs(2, level=2)) + + def test_where(self): + pdf, psdf = self.df_pair + + # pandas requires `axis` argument when the `other` is Series. + # `axis` is not fully supported yet in pandas-on-Spark. + self.assert_eq( + psdf.where(psdf > 2, psdf.a + 10, axis=0), pdf.where(pdf > 2, pdf.a + 10, axis=0) + ) + + with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"): + psdf.where(1) + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psdf.where(psdf > 2, psdf.a + 10, axis=1) + + def test_mask(self): + psdf = ps.from_pandas(self.pdf) + + with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"): + psdf.mask(1) + + def test_query(self): + pdf = pd.DataFrame({"A": range(1, 6), "B": range(10, 0, -2), "C": range(10, 5, -1)}) + psdf = ps.from_pandas(pdf) + + exprs = ("A > B", "A < C", "C == B") + for expr in exprs: + self.assert_eq(psdf.query(expr), pdf.query(expr)) + + # test `inplace=True` + for expr in exprs: + dummy_psdf = psdf.copy() + dummy_pdf = pdf.copy() + + pser = dummy_pdf.A + psser = dummy_psdf.A + dummy_pdf.query(expr, inplace=True) + dummy_psdf.query(expr, inplace=True) + + self.assert_eq(dummy_psdf, dummy_pdf) + self.assert_eq(psser, pser) + + # invalid values for `expr` + invalid_exprs = (1, 1.0, (exprs[0],), [exprs[0]]) + for expr in invalid_exprs: + with self.assertRaisesRegex( + TypeError, + "expr must be a string to be evaluated, {} given".format(type(expr).__name__), + ): + psdf.query(expr) + + # invalid values for `inplace` + invalid_inplaces = (1, 0, "True", "False") + for inplace in invalid_inplaces: + with self.assertRaisesRegex( + TypeError, + 'For argument "inplace" expected type bool, received type {}.'.format( + type(inplace).__name__ + ), + ): + psdf.query("a < b", inplace=inplace) + + # doesn't support for MultiIndex columns + columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "C")]) + psdf.columns = columns + with self.assertRaisesRegex(TypeError, "Doesn't support for MultiIndex columns"): + psdf.query("('A', 'Z') > ('B', 'X')") + + def test_insert(self): + # + # Basic DataFrame + # + pdf = pd.DataFrame([1, 2, 3]) + psdf = ps.from_pandas(pdf) + + psdf.insert(1, "b", 10) + pdf.insert(1, "b", 10) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + psdf.insert(2, "c", 0.1) + pdf.insert(2, "c", 0.1) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + psdf.insert(3, "d", psdf.b + 1) + pdf.insert(3, "d", pdf.b + 1) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + + psser = ps.Series([4, 5, 6]) + self.assertRaises(ValueError, lambda: psdf.insert(0, "y", psser)) + self.assertRaisesRegex( + ValueError, "cannot insert b, already exists", lambda: psdf.insert(1, "b", 10) + ) + self.assertRaisesRegex( + TypeError, + '"column" should be a scalar value or tuple that contains scalar values', + lambda: psdf.insert(0, list("abc"), psser), + ) + self.assertRaisesRegex( + TypeError, + "loc must be int", + lambda: psdf.insert((1,), "b", 10), + ) + self.assertRaisesRegex( + NotImplementedError, + "Assigning column name as tuple is only supported for MultiIndex columns for now.", + lambda: psdf.insert(0, ("e",), 10), + ) + + self.assertRaises(ValueError, lambda: psdf.insert(0, "e", [7, 8, 9, 10])) + self.assertRaises(ValueError, lambda: psdf.insert(0, "f", ps.Series([7, 8]))) + self.assertRaises(AssertionError, lambda: psdf.insert(100, "y", psser)) + self.assertRaises(AssertionError, lambda: psdf.insert(1, "y", psser, allow_duplicates=True)) + + # + # DataFrame with MultiIndex as columns + # + pdf = pd.DataFrame({("x", "a", "b"): [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + + psdf.insert(1, "b", 10) + pdf.insert(1, "b", 10) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + psdf.insert(2, "c", 0.1) + pdf.insert(2, "c", 0.1) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + psdf.insert(3, "d", psdf.b + 1) + pdf.insert(3, "d", pdf.b + 1) + self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True) + + self.assertRaisesRegex( + ValueError, "cannot insert d, already exists", lambda: psdf.insert(4, "d", 11) + ) + self.assertRaisesRegex( + ValueError, + r"cannot insert \('x', 'a', 'b'\), already exists", + lambda: psdf.insert(4, ("x", "a", "b"), 11), + ) + self.assertRaisesRegex( + ValueError, + '"column" must have length equal to number of column levels.', + lambda: psdf.insert(4, ("e",), 11), + ) + + def test_itertuples(self): + pdf = pd.DataFrame({"num_legs": [4, 2], "num_wings": [0, 2]}, index=["dog", "hawk"]) + psdf = ps.from_pandas(pdf) + + for ptuple, ktuple in zip( + pdf.itertuples(index=False, name="Animal"), psdf.itertuples(index=False, name="Animal") + ): + self.assert_eq(ptuple, ktuple) + for ptuple, ktuple in zip(pdf.itertuples(name=None), psdf.itertuples(name=None)): + self.assert_eq(ptuple, ktuple) + for ptuple, ktuple in zip( + pdf.itertuples(index=False, name=None), psdf.itertuples(index=False, name=None) + ): + self.assert_eq(ptuple, ktuple) + + pdf.index = pd.MultiIndex.from_arrays( + [[1, 2], ["black", "brown"]], names=("count", "color") + ) + psdf = ps.from_pandas(pdf) + for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")): + self.assert_eq(ptuple, ktuple) + + pdf.columns = pd.MultiIndex.from_arrays( + [["CA", "WA"], ["age", "children"]], names=("origin", "info") + ) + psdf = ps.from_pandas(pdf) + for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")): + self.assert_eq(ptuple, ktuple) + + pdf = pd.DataFrame([1, 2, 3]) + psdf = ps.from_pandas(pdf) + for ptuple, ktuple in zip( + (pdf + 1).itertuples(name="num"), (psdf + 1).itertuples(name="num") + ): + self.assert_eq(ptuple, ktuple) + + # DataFrames with a large number of columns (>254) + pdf = pd.DataFrame(np.random.random((1, 255))) + psdf = ps.from_pandas(pdf) + for ptuple, ktuple in zip(pdf.itertuples(name="num"), psdf.itertuples(name="num")): + self.assert_eq(ptuple, ktuple) + + def test_iterrows(self): + pdf = pd.DataFrame( + { + ("x", "a", "1"): [1, 2, 3], + ("x", "b", "2"): [4, 5, 6], + ("y.z", "c.d", "3"): [7, 8, 9], + ("x", "b", "4"): [10, 11, 12], + }, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + + for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()): + self.assert_eq(pdf_k, psdf_k) + self.assert_eq(pdf_v, psdf_v) + + # MultiIndex + pmidx = pd.Index([(1, 2), (3, 4), (5, 6)]) + pdf.index = pmidx + psdf = ps.from_pandas(pdf) + + for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()): + self.assert_eq(pdf_k, psdf_k) + self.assert_eq(pdf_v, psdf_v) + + def test_multiindex_column_access(self): + columns = pd.MultiIndex.from_tuples( + [ + ("a", "", "", "b"), + ("c", "", "d", ""), + ("e", "", "f", ""), + ("e", "g", "", ""), + ("", "", "", "h"), + ("i", "", "", ""), + ] + ) + + pdf = pd.DataFrame( + [ + (1, "a", "x", 10, 100, 1000), + (2, "b", "y", 20, 200, 2000), + (3, "c", "z", 30, 300, 3000), + ], + columns=columns, + index=np.random.rand(3), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf["a"], pdf["a"]) + self.assert_eq(psdf["a"]["b"], pdf["a"]["b"]) + self.assert_eq(psdf["c"], pdf["c"]) + self.assert_eq(psdf["c"]["d"], pdf["c"]["d"]) + self.assert_eq(psdf["e"], pdf["e"]) + self.assert_eq(psdf["e"][""]["f"], pdf["e"][""]["f"]) + self.assert_eq(psdf["e"]["g"], pdf["e"]["g"]) + self.assert_eq(psdf[""], pdf[""]) + self.assert_eq(psdf[""]["h"], pdf[""]["h"]) + self.assert_eq(psdf["i"], pdf["i"]) + + self.assert_eq(psdf[["a", "e"]], pdf[["a", "e"]]) + self.assert_eq(psdf[["e", "a"]], pdf[["e", "a"]]) + + self.assert_eq(psdf[("a",)], pdf[("a",)]) + self.assert_eq(psdf[("e", "g")], pdf[("e", "g")]) + # self.assert_eq(psdf[("i",)], pdf[("i",)]) + self.assert_eq(psdf[("i", "")], pdf[("i", "")]) + + self.assertRaises(KeyError, lambda: psdf[("a", "b")]) + + def test_getitem_with_none_key(self): + psdf = self.psdf + + with self.assertRaisesRegex(KeyError, "none key"): + psdf[None] + + def test_iter_dataframe(self): + pdf, psdf = self.df_pair + + for value_psdf, value_pdf in zip(psdf, pdf): + self.assert_eq(value_psdf, value_pdf) + + +class FrameIndexingTests(FrameIndexingMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_indexing import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reindex.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reindex.py new file mode 100644 index 0000000..933b4a2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reindex.py @@ -0,0 +1,310 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameReindexMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43811): Enable DataFrameTests.test_reindex for pandas 2.0.0.", + ) + def test_reindex(self): + index = pd.Index(["A", "B", "C", "D", "E"]) + columns = pd.Index(["numbers"]) + pdf = pd.DataFrame([1.0, 2.0, 3.0, 4.0, None], index=index, columns=columns) + psdf = ps.from_pandas(pdf) + + columns2 = pd.Index(["numbers", "2", "3"], name="cols2") + self.assert_eq( + pdf.reindex(columns=columns2).sort_index(), + psdf.reindex(columns=columns2).sort_index(), + ) + + columns = pd.Index(["numbers"], name="cols") + pdf.columns = columns + psdf.columns = columns + + self.assert_eq( + pdf.reindex(["A", "B", "C"], columns=["numbers", "2", "3"]).sort_index(), + psdf.reindex(["A", "B", "C"], columns=["numbers", "2", "3"]).sort_index(), + ) + + self.assert_eq( + pdf.reindex(["A", "B", "C"], index=["numbers", "2", "3"]).sort_index(), + psdf.reindex(["A", "B", "C"], index=["numbers", "2", "3"]).sort_index(), + ) + + self.assert_eq( + pdf.reindex(index=["A", "B"]).sort_index(), psdf.reindex(index=["A", "B"]).sort_index() + ) + + self.assert_eq( + pdf.reindex(index=["A", "B", "2", "3"]).sort_index(), + psdf.reindex(index=["A", "B", "2", "3"]).sort_index(), + ) + + self.assert_eq( + pdf.reindex(index=["A", "E", "2", "3"], fill_value=0).sort_index(), + psdf.reindex(index=["A", "E", "2", "3"], fill_value=0).sort_index(), + ) + + self.assert_eq( + pdf.reindex(columns=["numbers"]).sort_index(), + psdf.reindex(columns=["numbers"]).sort_index(), + ) + + self.assert_eq( + pdf.reindex(columns=["numbers"], copy=True).sort_index(), + psdf.reindex(columns=["numbers"], copy=True).sort_index(), + ) + + # Using float as fill_value to avoid int64/32 clash + self.assert_eq( + pdf.reindex(columns=["numbers", "2", "3"], fill_value=0.0).sort_index(), + psdf.reindex(columns=["numbers", "2", "3"], fill_value=0.0).sort_index(), + ) + + columns2 = pd.Index(["numbers", "2", "3"]) + self.assert_eq( + pdf.reindex(columns=columns2).sort_index(), + psdf.reindex(columns=columns2).sort_index(), + ) + + columns2 = pd.Index(["numbers", "2", "3"], name="cols2") + self.assert_eq( + pdf.reindex(columns=columns2).sort_index(), + psdf.reindex(columns=columns2).sort_index(), + ) + + # Reindexing single Index on single Index + pindex2 = pd.Index(["A", "C", "D", "E", "0"], name="index2") + kindex2 = ps.from_pandas(pindex2) + + for fill_value in [None, 0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + pindex2 = pd.DataFrame({"index2": ["A", "C", "D", "E", "0"]}).set_index("index2").index + kindex2 = ps.from_pandas(pindex2) + + for fill_value in [None, 0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + # Reindexing MultiIndex on single Index + pindex = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("F", "G")], names=["name1", "name2"] + ) + kindex = ps.from_pandas(pindex) + + self.assert_eq( + pdf.reindex(index=pindex, fill_value=0.0).sort_index(), + psdf.reindex(index=kindex, fill_value=0.0).sort_index(), + ) + + # Specifying the `labels` parameter + new_index = ["V", "W", "X", "Y", "Z"] + self.assert_eq( + pdf.reindex(labels=new_index, fill_value=0.0, axis=0).sort_index(), + psdf.reindex(labels=new_index, fill_value=0.0, axis=0).sort_index(), + ) + self.assert_eq( + pdf.reindex(labels=new_index, fill_value=0.0, axis=1).sort_index(), + psdf.reindex(labels=new_index, fill_value=0.0, axis=1).sort_index(), + ) + + self.assertRaises(TypeError, lambda: psdf.reindex(columns=["numbers", "2", "3"], axis=1)) + self.assertRaises(TypeError, lambda: psdf.reindex(columns=["numbers", "2", "3"], axis=2)) + self.assertRaises(TypeError, lambda: psdf.reindex(columns="numbers")) + self.assertRaises(TypeError, lambda: psdf.reindex(index=["A", "B", "C"], axis=1)) + self.assertRaises(TypeError, lambda: psdf.reindex(index=123)) + + # Reindexing MultiIndex on MultiIndex + pdf = pd.DataFrame({"numbers": [1.0, 2.0, None]}, index=pindex) + psdf = ps.from_pandas(pdf) + pindex2 = pd.MultiIndex.from_tuples( + [("A", "G"), ("C", "D"), ("I", "J")], names=["name1", "name2"] + ) + kindex2 = ps.from_pandas(pindex2) + + for fill_value in [None, 0.0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + pindex2 = ( + pd.DataFrame({"index_level_1": ["A", "C", "I"], "index_level_2": ["G", "D", "J"]}) + .set_index(["index_level_1", "index_level_2"]) + .index + ) + kindex2 = ps.from_pandas(pindex2) + + for fill_value in [None, 0.0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + columns = pd.MultiIndex.from_tuples([("X", "numbers")], names=["cols1", "cols2"]) + pdf.columns = columns + psdf.columns = columns + + # Reindexing MultiIndex index on MultiIndex columns and MultiIndex index + for fill_value in [None, 0.0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + index = pd.Index(["A", "B", "C", "D", "E"]) + pdf = pd.DataFrame(data=[1.0, 2.0, 3.0, 4.0, None], index=index, columns=columns) + psdf = ps.from_pandas(pdf) + pindex2 = pd.Index(["A", "C", "D", "E", "0"], name="index2") + kindex2 = ps.from_pandas(pindex2) + + # Reindexing single Index on MultiIndex columns and single Index + for fill_value in [None, 0.0]: + self.assert_eq( + pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(), + psdf.reindex(index=kindex2, fill_value=fill_value).sort_index(), + ) + + for fill_value in [None, 0.0]: + self.assert_eq( + pdf.reindex( + columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=fill_value + ).sort_index(), + psdf.reindex( + columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=fill_value + ).sort_index(), + ) + + columns2 = pd.MultiIndex.from_tuples( + [("X", "numbers"), ("Y", "2"), ("Y", "3")], names=["cols3", "cols4"] + ) + self.assert_eq( + pdf.reindex(columns=columns2).sort_index(), + psdf.reindex(columns=columns2).sort_index(), + ) + + self.assertRaises(TypeError, lambda: psdf.reindex(columns=["X"])) + self.assertRaises(ValueError, lambda: psdf.reindex(columns=[("X",)])) + + def test_reindex_like(self): + data = [[1.0, 2.0], [3.0, None], [None, 4.0]] + index = pd.Index(["A", "B", "C"], name="index") + columns = pd.Index(["numbers", "values"], name="cols") + pdf = pd.DataFrame(data=data, index=index, columns=columns) + psdf = ps.from_pandas(pdf) + + # Reindexing single Index on single Index + data2 = [[5.0, None], [6.0, 7.0], [8.0, None]] + index2 = pd.Index(["A", "C", "D"], name="index2") + columns2 = pd.Index(["numbers", "F"], name="cols2") + pdf2 = pd.DataFrame(data=data2, index=index2, columns=columns2) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq( + pdf.reindex_like(pdf2).sort_index(), + psdf.reindex_like(psdf2).sort_index(), + ) + + pdf2 = pd.DataFrame({"index_level_1": ["A", "C", "I"]}) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq( + pdf.reindex_like(pdf2.set_index(["index_level_1"])).sort_index(), + psdf.reindex_like(psdf2.set_index(["index_level_1"])).sort_index(), + ) + + # Reindexing MultiIndex on single Index + index2 = pd.MultiIndex.from_tuples( + [("A", "G"), ("C", "D"), ("I", "J")], names=["name3", "name4"] + ) + pdf2 = pd.DataFrame(data=data2, index=index2) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq( + pdf.reindex_like(pdf2).sort_index(), + psdf.reindex_like(psdf2).sort_index(), + ) + + self.assertRaises(TypeError, lambda: psdf.reindex_like(index2)) + self.assertRaises(AssertionError, lambda: psdf2.reindex_like(psdf)) + + # Reindexing MultiIndex on MultiIndex + columns2 = pd.MultiIndex.from_tuples( + [("numbers", "third"), ("values", "second")], names=["cols3", "cols4"] + ) + pdf2.columns = columns2 + psdf2.columns = columns2 + + columns = pd.MultiIndex.from_tuples( + [("numbers", "first"), ("values", "second")], names=["cols1", "cols2"] + ) + index = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["name1", "name2"] + ) + pdf = pd.DataFrame(data=data, index=index, columns=columns) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pdf.reindex_like(pdf2).sort_index(), + psdf.reindex_like(psdf2).sort_index(), + ) + + +class FrameReindexTests(FrameReindexMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_reindex import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_rename.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_rename.py new file mode 100644 index 0000000..b594089 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_rename.py @@ -0,0 +1,257 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameRenameMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_rename_dataframe(self): + pdf1 = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) + psdf1 = ps.from_pandas(pdf1) + + self.assert_eq( + psdf1.rename(columns={"A": "a", "B": "b"}), pdf1.rename(columns={"A": "a", "B": "b"}) + ) + + result_psdf = psdf1.rename(index={1: 10, 2: 20}) + result_pdf = pdf1.rename(index={1: 10, 2: 20}) + self.assert_eq(result_psdf, result_pdf) + + # inplace + pser = result_pdf.A + psser = result_psdf.A + result_psdf.rename(index={10: 100, 20: 200}, inplace=True) + result_pdf.rename(index={10: 100, 20: 200}, inplace=True) + self.assert_eq(result_psdf, result_pdf) + self.assert_eq(psser, pser) + + def str_lower(s) -> str: + return str.lower(s) + + self.assert_eq( + psdf1.rename(str_lower, axis="columns"), pdf1.rename(str_lower, axis="columns") + ) + + def mul10(x) -> int: + return x * 10 + + self.assert_eq(psdf1.rename(mul10, axis="index"), pdf1.rename(mul10, axis="index")) + + self.assert_eq( + psdf1.rename(columns=str_lower, index={1: 10, 2: 20}), + pdf1.rename(columns=str_lower, index={1: 10, 2: 20}), + ) + + self.assert_eq( + psdf1.rename(columns=lambda x: str.lower(x)), + pdf1.rename(columns=lambda x: str.lower(x)), + ) + + idx = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Y", "D")]) + pdf2 = pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], columns=idx) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(psdf2.rename(columns=str_lower), pdf2.rename(columns=str_lower)) + self.assert_eq( + psdf2.rename(columns=lambda x: str.lower(x)), + pdf2.rename(columns=lambda x: str.lower(x)), + ) + + self.assert_eq( + psdf2.rename(columns=str_lower, level=0), pdf2.rename(columns=str_lower, level=0) + ) + self.assert_eq( + psdf2.rename(columns=str_lower, level=1), pdf2.rename(columns=str_lower, level=1) + ) + + pdf3 = pd.DataFrame([[1, 2], [3, 4], [5, 6], [7, 8]], index=idx, columns=list("ab")) + psdf3 = ps.from_pandas(pdf3) + + self.assert_eq(psdf3.rename(index=str_lower), pdf3.rename(index=str_lower)) + self.assert_eq( + psdf3.rename(index=str_lower, level=0), pdf3.rename(index=str_lower, level=0) + ) + self.assert_eq( + psdf3.rename(index=str_lower, level=1), pdf3.rename(index=str_lower, level=1) + ) + + pdf4 = pdf2 + 1 + psdf4 = psdf2 + 1 + self.assert_eq(psdf4.rename(columns=str_lower), pdf4.rename(columns=str_lower)) + + pdf5 = pdf3 + 1 + psdf5 = psdf3 + 1 + self.assert_eq(psdf5.rename(index=str_lower), pdf5.rename(index=str_lower)) + + msg = "Either `index` or `columns` should be provided." + with self.assertRaisesRegex(ValueError, msg): + psdf1.rename() + msg = "`mapper` or `index` or `columns` should be either dict-like or function type." + with self.assertRaisesRegex(ValueError, msg): + psdf1.rename(mapper=[str_lower], axis=1) + msg = "Mapper dict should have the same value type." + with self.assertRaisesRegex(ValueError, msg): + psdf1.rename({"A": "a", "B": 2}, axis=1) + msg = r"level should be an integer between \[0, column_labels_level\)" + with self.assertRaisesRegex(ValueError, msg): + psdf2.rename(columns=str_lower, level=2) + msg = r"level should be an integer between \[0, 2\)" + with self.assertRaisesRegex(ValueError, msg): + psdf3.rename(index=str_lower, level=2) + + def test_rename_axis(self): + index = pd.Index(["A", "B", "C"], name="index") + columns = pd.Index(["numbers", "values"], name="cols") + pdf = pd.DataFrame([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], index=index, columns=columns) + psdf = ps.from_pandas(pdf) + + for axis in [0, "index"]: + self.assert_eq( + pdf.rename_axis("index2", axis=axis).sort_index(), + psdf.rename_axis("index2", axis=axis).sort_index(), + ) + self.assert_eq( + pdf.rename_axis(["index2"], axis=axis).sort_index(), + psdf.rename_axis(["index2"], axis=axis).sort_index(), + ) + + for axis in [1, "columns"]: + self.assert_eq( + pdf.rename_axis("cols2", axis=axis).sort_index(), + psdf.rename_axis("cols2", axis=axis).sort_index(), + ) + self.assert_eq( + pdf.rename_axis(["cols2"], axis=axis).sort_index(), + psdf.rename_axis(["cols2"], axis=axis).sort_index(), + ) + + pdf2 = pdf.copy() + psdf2 = psdf.copy() + pdf2.rename_axis("index2", axis="index", inplace=True) + psdf2.rename_axis("index2", axis="index", inplace=True) + self.assert_eq(pdf2.sort_index(), psdf2.sort_index()) + + self.assertRaises(ValueError, lambda: psdf.rename_axis(["index2", "index3"], axis=0)) + self.assertRaises(ValueError, lambda: psdf.rename_axis(["cols2", "cols3"], axis=1)) + self.assertRaises(TypeError, lambda: psdf.rename_axis(mapper=["index2"], index=["index3"])) + self.assertRaises(ValueError, lambda: psdf.rename_axis(ps)) + + self.assert_eq( + pdf.rename_axis(index={"index": "index2"}, columns={"cols": "cols2"}).sort_index(), + psdf.rename_axis(index={"index": "index2"}, columns={"cols": "cols2"}).sort_index(), + ) + + self.assert_eq( + pdf.rename_axis(index={"missing": "index2"}, columns={"missing": "cols2"}).sort_index(), + psdf.rename_axis( + index={"missing": "index2"}, columns={"missing": "cols2"} + ).sort_index(), + ) + + self.assert_eq( + pdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), + psdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), + ) + + index = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] + ) + columns = pd.MultiIndex.from_tuples( + [("numbers", "first"), ("values", "second")], names=["cols1", "cols2"] + ) + pdf = pd.DataFrame([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], index=index, columns=columns) + psdf = ps.from_pandas(pdf) + + for axis in [0, "index"]: + self.assert_eq( + pdf.rename_axis(["index3", "index4"], axis=axis).sort_index(), + psdf.rename_axis(["index3", "index4"], axis=axis).sort_index(), + ) + + for axis in [1, "columns"]: + self.assert_eq( + pdf.rename_axis(["cols3", "cols4"], axis=axis).sort_index(), + psdf.rename_axis(["cols3", "cols4"], axis=axis).sort_index(), + ) + + self.assertRaises( + ValueError, lambda: psdf.rename_axis(["index3", "index4", "index5"], axis=0) + ) + self.assertRaises(ValueError, lambda: psdf.rename_axis(["cols3", "cols4", "cols5"], axis=1)) + + self.assert_eq( + pdf.rename_axis(index={"index1": "index3"}, columns={"cols1": "cols3"}).sort_index(), + psdf.rename_axis(index={"index1": "index3"}, columns={"cols1": "cols3"}).sort_index(), + ) + + self.assert_eq( + pdf.rename_axis(index={"missing": "index3"}, columns={"missing": "cols3"}).sort_index(), + psdf.rename_axis( + index={"missing": "index3"}, columns={"missing": "cols3"} + ).sort_index(), + ) + + self.assert_eq( + pdf.rename_axis( + index={"index1": "index3", "index2": "index4"}, + columns={"cols1": "cols3", "cols2": "cols4"}, + ).sort_index(), + psdf.rename_axis( + index={"index1": "index3", "index2": "index4"}, + columns={"cols1": "cols3", "cols2": "cols4"}, + ).sort_index(), + ) + + self.assert_eq( + pdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), + psdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), + ) + + +class FrameRenameTests(FrameRenameMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_rename import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reset_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reset_index.py new file mode 100644 index 0000000..17cdb64 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_reset_index.py @@ -0,0 +1,160 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class FrameResetIndexMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_reset_index(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=np.random.rand(3)) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + self.assert_eq(psdf.reset_index().index, pdf.reset_index().index) + self.assert_eq(psdf.reset_index(drop=True), pdf.reset_index(drop=True)) + + pdf.index.name = "a" + psdf.index.name = "a" + + with self.assertRaisesRegex(ValueError, "cannot insert a, already exists"): + psdf.reset_index() + + self.assert_eq(psdf.reset_index(drop=True), pdf.reset_index(drop=True)) + + # inplace + pser = pdf.a + psser = psdf.a + pdf.reset_index(drop=True, inplace=True) + psdf.reset_index(drop=True, inplace=True) + self.assert_eq(psdf, pdf) + self.assert_eq(psser, pser) + + pdf.columns = ["index", "b"] + psdf.columns = ["index", "b"] + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + + def test_reset_index_with_default_index_types(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=np.random.rand(3)) + psdf = ps.from_pandas(pdf) + + with ps.option_context("compute.default_index_type", "sequence"): + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + + with ps.option_context("compute.default_index_type", "distributed-sequence"): + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + + with ps.option_context("compute.default_index_type", "distributed"): + # the index is different. + self.assert_eq( + psdf.reset_index()._to_pandas().reset_index(drop=True), pdf.reset_index() + ) + + def test_reset_index_with_multiindex_columns(self): + index = pd.MultiIndex.from_tuples( + [("bird", "falcon"), ("bird", "parrot"), ("mammal", "lion"), ("mammal", "monkey")], + names=["class", "name"], + ) + columns = pd.MultiIndex.from_tuples([("speed", "max"), ("species", "type")]) + pdf = pd.DataFrame( + [(389.0, "fly"), (24.0, "fly"), (80.5, "run"), (np.nan, "jump")], + index=index, + columns=columns, + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + self.assert_eq(psdf.reset_index(level="class"), pdf.reset_index(level="class")) + self.assert_eq( + psdf.reset_index(level="class", col_level=1), + pdf.reset_index(level="class", col_level=1), + ) + self.assert_eq( + psdf.reset_index(level="class", col_level=1, col_fill="species"), + pdf.reset_index(level="class", col_level=1, col_fill="species"), + ) + self.assert_eq( + psdf.reset_index(level="class", col_level=1, col_fill="genus"), + pdf.reset_index(level="class", col_level=1, col_fill="genus"), + ) + + with self.assertRaisesRegex(IndexError, "Index has only 2 levels, not 3"): + psdf.reset_index(col_level=2) + + pdf.index.names = [("x", "class"), ("y", "name")] + psdf.index.names = [("x", "class"), ("y", "name")] + + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + + with self.assertRaisesRegex(ValueError, "Item must have length equal to number of levels."): + psdf.reset_index(col_level=1) + + def test_index_to_frame_reset_index(self): + def check(psdf, pdf): + self.assert_eq(psdf.reset_index(), pdf.reset_index()) + self.assert_eq(psdf.reset_index(drop=True), pdf.reset_index(drop=True)) + + pdf.reset_index(drop=True, inplace=True) + psdf.reset_index(drop=True, inplace=True) + self.assert_eq(psdf, pdf) + + pdf, psdf = self.df_pair + check(psdf.index.to_frame(), pdf.index.to_frame()) + check(psdf.index.to_frame(index=False), pdf.index.to_frame(index=False)) + + check(psdf.index.to_frame(name="a"), pdf.index.to_frame(name="a")) + check(psdf.index.to_frame(index=False, name="a"), pdf.index.to_frame(index=False, name="a")) + check(psdf.index.to_frame(name=("x", "a")), pdf.index.to_frame(name=("x", "a"))) + check( + psdf.index.to_frame(index=False, name=("x", "a")), + pdf.index.to_frame(index=False, name=("x", "a")), + ) + + +class FrameResetIndexTests(FrameResetIndexMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.indexes.test_reset_index import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_timedelta.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_timedelta.py new file mode 100644 index 0000000..a9bb93e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/indexes/test_timedelta.py @@ -0,0 +1,128 @@ +# +# 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 unittest +from datetime import timedelta +from distutils.version import LooseVersion + +import pandas as pd + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class TimedeltaIndexTestsMixin: + @property + def pidx(self): + return pd.TimedeltaIndex( + [ + timedelta(days=1), + timedelta(seconds=1), + timedelta(microseconds=1), + timedelta(milliseconds=1), + timedelta(minutes=1), + timedelta(hours=1), + timedelta(weeks=1), + ], + name="x", + ) + + @property + def neg_pidx(self): + return pd.TimedeltaIndex( + [ + timedelta(days=-1), + timedelta(seconds=-1), + timedelta(microseconds=-1), + timedelta(milliseconds=-1), + timedelta(minutes=-1), + timedelta(hours=-1), + timedelta(weeks=-1), + ], + name="x", + ) + + @property + def psidx(self): + return ps.from_pandas(self.pidx) + + @property + def neg_psidx(self): + return ps.from_pandas(self.neg_pidx) + + def test_timedelta_index(self): + # Create TimedeltaIndex from constructor + psidx = ps.TimedeltaIndex( + [ + timedelta(days=1), + timedelta(seconds=1), + timedelta(microseconds=1), + timedelta(milliseconds=1), + timedelta(minutes=1), + timedelta(hours=1), + timedelta(weeks=1), + ], + name="x", + ) + self.assert_eq(psidx, self.pidx) + # Create TimedeltaIndex from Series + self.assert_eq( + ps.TimedeltaIndex(ps.Series([timedelta(days=1)])), + pd.TimedeltaIndex(pd.Series([timedelta(days=1)])), + ) + # Create TimedeltaIndex from Index + self.assert_eq( + ps.TimedeltaIndex(ps.Index([timedelta(days=1)])), + pd.TimedeltaIndex(pd.Index([timedelta(days=1)])), + ) + + # ps.TimedeltaIndex(ps.Index([1, 2, 3])) + with self.assertRaisesRegexp(TypeError, "Index.name must be a hashable type"): + ps.TimedeltaIndex([timedelta(1), timedelta(microseconds=2)], name=[(1, 2)]) + with self.assertRaisesRegexp( + TypeError, "Cannot perform 'all' with this index type: TimedeltaIndex" + ): + psidx.all() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43705): Enable TimedeltaIndexTests.test_properties for pandas 2.0.0.", + ) + def test_properties(self): + self.assert_eq(self.psidx.days, self.pidx.days) + self.assert_eq(self.psidx.seconds, self.pidx.seconds) + self.assert_eq(self.psidx.microseconds, self.pidx.microseconds) + self.assert_eq(self.neg_psidx.days, self.neg_pidx.days) + self.assert_eq(self.neg_psidx.seconds, self.neg_pidx.seconds) + self.assert_eq(self.neg_psidx.microseconds, self.neg_pidx.microseconds) + + +class TimedeltaIndexTests(TimedeltaIndexTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.indexes.test_timedelta import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/test_io.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/test_io.py new file mode 100644 index 0000000..4eadc6a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/io/test_io.py @@ -0,0 +1,137 @@ +# +# 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 distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ( + have_tabulate, + ComparisonTestBase, + tabulate_requirement_message, +) +from pyspark.testing.sqlutils import SQLTestUtils + + +# This file contains test cases for 'Serialization / IO / Conversion' +# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#serialization-io-conversion +class FrameIOMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def df_pair(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + return pdf, psdf + + def test_to_numpy(self): + pdf = pd.DataFrame( + { + "a": [4, 2, 3, 4, 8, 6], + "b": [1, 2, 9, 4, 2, 4], + "c": ["one", "three", "six", "seven", "one", "5"], + }, + index=np.random.rand(6), + ) + + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.to_numpy(), pdf.values) + + def test_to_pandas(self): + pdf, psdf = self.df_pair + self.assert_eq(psdf._to_pandas(), pdf) + + def test_to_spark(self): + psdf = ps.from_pandas(self.pdf) + + with self.assertRaisesRegex(ValueError, "'index_col' cannot be overlapped"): + psdf.to_spark(index_col="a") + + with self.assertRaisesRegex(ValueError, "length of index columns.*1.*3"): + psdf.to_spark(index_col=["x", "y", "z"]) + + @unittest.skipIf(not have_tabulate, tabulate_requirement_message) + def test_to_markdown(self): + pdf = pd.DataFrame(data={"animal_1": ["elk", "pig"], "animal_2": ["dog", "quetzal"]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.to_markdown(), psdf.to_markdown()) + + def test_from_dict(self): + data = {"row_1": [3, 2, 1, 0], "row_2": [10, 20, 30, 40]} + pdf = pd.DataFrame.from_dict(data) + psdf = ps.DataFrame.from_dict(data) + self.assert_eq(pdf, psdf) + + pdf = pd.DataFrame.from_dict(data, dtype="int8") + psdf = ps.DataFrame.from_dict(data, dtype="int8") + self.assert_eq(pdf, psdf) + + pdf = pd.DataFrame.from_dict(data, orient="index", columns=["A", "B", "C", "D"]) + psdf = ps.DataFrame.from_dict(data, orient="index", columns=["A", "B", "C", "D"]) + self.assert_eq(pdf, psdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) < LooseVersion("1.3.0"), + "pandas support `Styler.to_latex` since 1.3.0", + ) + def test_style(self): + # Currently, the `style` function returns a pandas object `Styler` as it is, + # processing only the number of rows declared in `compute.max_rows`. + # So it's a bit vague to test, but we are doing minimal tests instead of not testing at all. + pdf = pd.DataFrame(np.random.randn(10, 4), columns=["A", "B", "C", "D"]) + psdf = ps.from_pandas(pdf) + + def style_negative(v, props=""): + return props if v < 0 else None + + def check_style(): + # If the value is negative, the text color will be displayed as red. + pdf_style = pdf.style.applymap(style_negative, props="color:red;") + psdf_style = psdf.style.applymap(style_negative, props="color:red;") + + # Test whether the same shape as pandas table is created including the color. + self.assert_eq(pdf_style.to_latex(), psdf_style.to_latex()) + + check_style() + + with ps.option_context("compute.max_rows", None): + check_style() + + +class FrameIOTests(FrameIOMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.io.test_io import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot.py new file mode 100644 index 0000000..6797a73 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot.py @@ -0,0 +1,170 @@ +# +# 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 pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option, option_context +from pyspark.pandas.plot import TopNPlotBase, SampledPlotBase, HistogramPlotBase, BoxPlotBase +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class DataFramePlotTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("plotting.max_rows", 2000) + set_option("plotting.sample_ratio", None) + + @classmethod + def tearDownClass(cls): + reset_option("plotting.max_rows") + reset_option("plotting.sample_ratio") + super().tearDownClass() + + def test_missing(self): + psdf = ps.DataFrame(np.random.rand(2500, 4), columns=["a", "b", "c", "d"]) + + unsupported_functions = ["hexbin"] + + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*DataFrame.*{}.*not implemented".format(name) + ): + getattr(psdf.plot, name)() + + def test_topn_max_rows(self): + + pdf = pd.DataFrame(np.random.rand(2500, 4), columns=["a", "b", "c", "d"]) + psdf = ps.from_pandas(pdf) + + data = TopNPlotBase().get_top_n(psdf) + self.assertEqual(len(data), 2000) + + def test_sampled_plot_with_ratio(self): + with option_context("plotting.sample_ratio", 0.5): + pdf = pd.DataFrame(np.random.rand(2500, 4), columns=["a", "b", "c", "d"]) + psdf = ps.from_pandas(pdf) + data = SampledPlotBase().get_sampled(psdf) + self.assertEqual(round(len(data) / 2500, 1), 0.5) + + def test_sampled_plot_with_max_rows(self): + # 'plotting.max_rows' is 2000 + pdf = pd.DataFrame(np.random.rand(2000, 4), columns=["a", "b", "c", "d"]) + psdf = ps.from_pandas(pdf) + data = SampledPlotBase().get_sampled(psdf) + self.assertEqual(round(len(data) / 2000, 1), 1) + + def test_compute_hist_single_column(self): + psdf = ps.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10] + ) + + expected_bins = np.linspace(1, 50, 11) + bins = HistogramPlotBase.get_bins(psdf[["a"]].to_spark(), 10) + + expected_histogram = np.array([5, 4, 1, 0, 0, 0, 0, 0, 0, 1]) + histogram = HistogramPlotBase.compute_hist(psdf[["a"]], bins)[0] + self.assert_eq(pd.Series(expected_bins), pd.Series(bins)) + self.assert_eq(pd.Series(expected_histogram, name="a"), histogram, almost=True) + + def test_compute_hist_multi_columns(self): + expected_bins = np.linspace(1, 50, 11) + psdf = ps.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], + "b": [50, 50, 30, 30, 30, 24, 10, 5, 4, 3, 1], + } + ) + + bins = HistogramPlotBase.get_bins(psdf.to_spark(), 10) + self.assert_eq(pd.Series(expected_bins), pd.Series(bins)) + + expected_histograms = [ + np.array([5, 4, 1, 0, 0, 0, 0, 0, 0, 1]), + np.array([4, 1, 0, 0, 1, 3, 0, 0, 0, 2]), + ] + histograms = HistogramPlotBase.compute_hist(psdf, bins) + expected_names = ["a", "b"] + + for histogram, expected_histogram, expected_name in zip( + histograms, expected_histograms, expected_names + ): + self.assert_eq( + pd.Series(expected_histogram, name=expected_name), histogram, almost=True + ) + + def test_compute_box_multi_columns(self): + # compare compute_multicol_stats with compute_stats + def check_box_multi_columns(psdf): + k = 1.5 + multicol_stats = BoxPlotBase.compute_multicol_stats( + psdf, ["a", "b", "c"], whis=k, precision=0.01 + ) + multicol_outliers = BoxPlotBase.multicol_outliers(psdf, multicol_stats) + multicol_whiskers = BoxPlotBase.calc_multicol_whiskers( + ["a", "b", "c"], multicol_outliers + ) + + for col in ["a", "b", "c"]: + col_stats = multicol_stats[col] + col_whiskers = multicol_whiskers[col] + + stats, fences = BoxPlotBase.compute_stats(psdf[col], col, whis=k, precision=0.01) + outliers = BoxPlotBase.outliers(psdf[col], col, *fences) + whiskers = BoxPlotBase.calc_whiskers(col, outliers) + + self.assertEqual(stats["mean"], col_stats["mean"]) + self.assertEqual(stats["med"], col_stats["med"]) + self.assertEqual(stats["q1"], col_stats["q1"]) + self.assertEqual(stats["q3"], col_stats["q3"]) + self.assertEqual(fences[0], col_stats["lfence"]) + self.assertEqual(fences[1], col_stats["ufence"]) + self.assertEqual(whiskers[0], col_whiskers["min"]) + self.assertEqual(whiskers[1], col_whiskers["max"]) + + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], + "b": [3, 2, 5, 4, 5, 6, 8, 8, 11, 60, 90], + "c": [-30, -2, 5, 4, 5, 6, -8, 8, 11, 12, 18], + }, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10], + ) + psdf = ps.from_pandas(pdf) + + check_box_multi_columns(psdf) + check_box_multi_columns(-psdf) + + +class DataFramePlotTests(DataFramePlotTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.plot.test_frame_plot import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py new file mode 100644 index 0000000..a479685 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py @@ -0,0 +1,547 @@ +# +# 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 base64 +from io import BytesIO +import unittest +from distutils.version import LooseVersion + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import ( + have_matplotlib, + matplotlib_requirement_message, + PandasOnSparkTestCase, + TestUtils, +) + +if have_matplotlib: + import matplotlib + from matplotlib import pyplot as plt + + matplotlib.use("agg") + + +@unittest.skipIf(not have_matplotlib, matplotlib_requirement_message) +class DataFramePlotMatplotlibTestsMixin: + sample_ratio_default = None + + @classmethod + def setUpClass(cls): + super().setUpClass() + pd.set_option("plotting.backend", "matplotlib") + set_option("plotting.backend", "matplotlib") + set_option("plotting.max_rows", 2000) + set_option("plotting.sample_ratio", None) + + @classmethod + def tearDownClass(cls): + pd.reset_option("plotting.backend") + reset_option("plotting.backend") + reset_option("plotting.max_rows") + reset_option("plotting.sample_ratio") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], "b": [2, 3, 4, 5, 7, 9, 10, 15, 34, 45, 49]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10], + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @staticmethod + def plot_to_base64(ax): + bytes_data = BytesIO() + ax.figure.savefig(bytes_data, format="png") + bytes_data.seek(0) + b64_data = base64.b64encode(bytes_data.read()) + plt.close(ax.figure) + return b64_data + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43641): Enable DataFramePlotMatplotlibTests.test_line_plot " + "for pandas 2.0.0.", + ) + def test_line_plot(self): + def check_line_plot(pdf, psdf): + ax1 = pdf.plot(kind="line", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="line", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.line(colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.line(colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_line_plot(pdf1, psdf1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf1.columns = columns + psdf1.columns = columns + check_line_plot(pdf1, psdf1) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43634): Enable DataFramePlotMatplotlibTests.test_area_plot for pandas 2.0.0.", + ) + def test_area_plot(self): + def check_area_plot(pdf, psdf): + ax1 = pdf.plot(kind="area", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="area", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.area(colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.area(colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + pdf = self.pdf1 + psdf = self.psdf1 + check_area_plot(pdf, psdf) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf.columns = columns + psdf.columns = columns + check_area_plot(pdf, psdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43635): Enable DataFramePlotMatplotlibTests.test_area_plot_stacked_false " + "for pandas 2.0.0.", + ) + def test_area_plot_stacked_false(self): + def check_area_plot_stacked_false(pdf, psdf): + ax1 = pdf.plot.area(stacked=False) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.area(stacked=False) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + # test if frame area plot is correct when stacked=False because default is True + + pdf = pd.DataFrame( + { + "sales": [3, 2, 3, 9, 10, 6], + "signups": [5, 5, 6, 12, 14, 13], + "visits": [20, 42, 28, 62, 81, 50], + }, + index=pd.date_range(start="2018/01/01", end="2018/07/01", freq="M"), + ) + psdf = ps.from_pandas(pdf) + check_area_plot_stacked_false(pdf, psdf) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "sales"), ("x", "signups"), ("y", "visits")]) + pdf.columns = columns + psdf.columns = columns + check_area_plot_stacked_false(pdf, psdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43636): Enable DataFramePlotMatplotlibTests.test_area_plot_y " + "for pandas 2.0.0.", + ) + def test_area_plot_y(self): + def check_area_plot_y(pdf, psdf, y): + ax1 = pdf.plot.area(y=y) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.area(y=y) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + # test if frame area plot is correct when y is specified + pdf = pd.DataFrame( + { + "sales": [3, 2, 3, 9, 10, 6], + "signups": [5, 5, 6, 12, 14, 13], + "visits": [20, 42, 28, 62, 81, 50], + }, + index=pd.date_range(start="2018/01/01", end="2018/07/01", freq="M"), + ) + psdf = ps.from_pandas(pdf) + check_area_plot_y(pdf, psdf, y="sales") + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "sales"), ("x", "signups"), ("y", "visits")]) + pdf.columns = columns + psdf.columns = columns + check_area_plot_y(pdf, psdf, y=("x", "sales")) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43639): Enable DataFramePlotMatplotlibTests.test_barh_plot_with_x_y " + "for pandas 2.0.0.", + ) + def test_barh_plot_with_x_y(self): + def check_barh_plot_with_x_y(pdf, psdf, x, y): + ax1 = pdf.plot(kind="barh", x=x, y=y, colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="barh", x=x, y=y, colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.barh(x=x, y=y, colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.barh(x=x, y=y, colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + # this is testing plot with specified x and y + pdf1 = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + psdf1 = ps.from_pandas(pdf1) + check_barh_plot_with_x_y(pdf1, psdf1, x="lab", y="val") + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "lab"), ("y", "val")]) + pdf1.columns = columns + psdf1.columns = columns + check_barh_plot_with_x_y(pdf1, psdf1, x=("x", "lab"), y=("y", "val")) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43640): Enable DataFramePlotMatplotlibTests.test_barh_plot " + "for pandas 2.0.0.", + ) + def test_barh_plot(self): + def check_barh_plot(pdf, psdf): + ax1 = pdf.plot(kind="barh", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="barh", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.barh(colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.barh(colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + # this is testing when x or y is not assigned + pdf1 = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + psdf1 = ps.from_pandas(pdf1) + check_barh_plot(pdf1, psdf1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "lab"), ("y", "val")]) + pdf1.columns = columns + psdf1.columns = columns + check_barh_plot(pdf1, psdf1) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43637): Enable DataFramePlotMatplotlibTests.test_bar_plot " "for pandas 2.0.0.", + ) + def test_bar_plot(self): + def check_bar_plot(pdf, psdf): + ax1 = pdf.plot(kind="bar", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="bar", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.bar(colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.bar(colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_bar_plot(pdf1, psdf1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "lab"), ("y", "val")]) + pdf1.columns = columns + psdf1.columns = columns + check_bar_plot(pdf1, psdf1) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43638): Enable DataFramePlotMatplotlibTests.test_bar_with_x_y " + "for pandas 2.0.0.", + ) + def test_bar_with_x_y(self): + # this is testing plot with specified x and y + pdf = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + psdf = ps.from_pandas(pdf) + + ax1 = pdf.plot(kind="bar", x="lab", y="val", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="bar", x="lab", y="val", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax3 = pdf.plot.bar(x="lab", y="val", colormap="Paired") + bin3 = self.plot_to_base64(ax3) + ax4 = psdf.plot.bar(x="lab", y="val", colormap="Paired") + bin4 = self.plot_to_base64(ax4) + self.assertEqual(bin3, bin4) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "lab"), ("y", "val")]) + pdf.columns = columns + psdf.columns = columns + + ax5 = pdf.plot(kind="bar", x=("x", "lab"), y=("y", "val"), colormap="Paired") + bin5 = self.plot_to_base64(ax5) + ax6 = psdf.plot(kind="bar", x=("x", "lab"), y=("y", "val"), colormap="Paired") + bin6 = self.plot_to_base64(ax6) + self.assertEqual(bin5, bin6) + + ax7 = pdf.plot.bar(x=("x", "lab"), y=("y", "val"), colormap="Paired") + bin7 = self.plot_to_base64(ax7) + ax8 = psdf.plot.bar(x=("x", "lab"), y=("y", "val"), colormap="Paired") + bin8 = self.plot_to_base64(ax8) + self.assertEqual(bin7, bin8) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43642): Enable DataFramePlotMatplotlibTests.test_pie_plot " "for pandas 2.0.0.", + ) + def test_pie_plot(self): + def check_pie_plot(pdf, psdf, y): + ax1 = pdf.plot.pie(y=y, figsize=(5, 5), colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.pie(y=y, figsize=(5, 5), colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf.plot(kind="pie", y=y, figsize=(5, 5), colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="pie", y=y, figsize=(5, 5), colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax11, ax12 = pdf.plot.pie(figsize=(5, 5), subplots=True, colormap="Paired") + bin11 = self.plot_to_base64(ax11) + bin12 = self.plot_to_base64(ax12) + self.assertEqual(bin11, bin12) + + ax21, ax22 = psdf.plot.pie(figsize=(5, 5), subplots=True, colormap="Paired") + bin21 = self.plot_to_base64(ax21) + bin22 = self.plot_to_base64(ax22) + self.assertEqual(bin21, bin22) + + ax11, ax12 = pdf.plot(kind="pie", figsize=(5, 5), subplots=True, colormap="Paired") + bin11 = self.plot_to_base64(ax11) + bin12 = self.plot_to_base64(ax12) + self.assertEqual(bin11, bin12) + + ax21, ax22 = psdf.plot(kind="pie", figsize=(5, 5), subplots=True, colormap="Paired") + bin21 = self.plot_to_base64(ax21) + bin22 = self.plot_to_base64(ax22) + self.assertEqual(bin21, bin22) + + pdf1 = pd.DataFrame( + {"mass": [0.330, 4.87, 5.97], "radius": [2439.7, 6051.8, 6378.1]}, + index=["Mercury", "Venus", "Earth"], + ) + psdf1 = ps.from_pandas(pdf1) + check_pie_plot(pdf1, psdf1, y="mass") + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "mass"), ("y", "radius")]) + pdf1.columns = columns + psdf1.columns = columns + check_pie_plot(pdf1, psdf1, y=("x", "mass")) + + def test_pie_plot_error_message(self): + # this is to test if error is correctly raising when y is not specified + # and subplots is not set to True + pdf = pd.DataFrame( + {"mass": [0.330, 4.87, 5.97], "radius": [2439.7, 6051.8, 6378.1]}, + index=["Mercury", "Venus", "Earth"], + ) + psdf = ps.from_pandas(pdf) + + with self.assertRaises(ValueError) as context: + psdf.plot.pie(figsize=(5, 5), colormap="Paired") + error_message = "pie requires either y column or 'subplots=True'" + self.assertTrue(error_message in str(context.exception)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43643): Enable DataFramePlotMatplotlibTests.test_scatter_plot " + "for pandas 2.0.0.", + ) + def test_scatter_plot(self): + def check_scatter_plot(pdf, psdf, x, y, c): + ax1 = pdf.plot.scatter(x=x, y=y) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.scatter(x=x, y=y) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf.plot(kind="scatter", x=x, y=y) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="scatter", x=x, y=y) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + # check when keyword c is given as name of a column + ax1 = pdf.plot.scatter(x=x, y=y, c=c, s=50) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.scatter(x=x, y=y, c=c, s=50) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + # Use pandas scatter plot example + pdf1 = pd.DataFrame(np.random.rand(50, 4), columns=["a", "b", "c", "d"]) + psdf1 = ps.from_pandas(pdf1) + check_scatter_plot(pdf1, psdf1, x="a", y="b", c="c") + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c"), ("z", "d")]) + pdf1.columns = columns + psdf1.columns = columns + check_scatter_plot(pdf1, psdf1, x=("x", "a"), y=("x", "b"), c=("y", "c")) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43720): Enable DataFramePlotMatplotlibTests.test_hist_plot for pandas 2.0.0.", + ) + def test_hist_plot(self): + def check_hist_plot(pdf, psdf): + _, ax1 = plt.subplots(1, 1) + ax1 = pdf.plot.hist() + bin1 = self.plot_to_base64(ax1) + _, ax2 = plt.subplots(1, 1) + ax2 = psdf.plot.hist() + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf.plot.hist(bins=15) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.hist(bins=15) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf.plot(kind="hist", bins=15) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot(kind="hist", bins=15) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf.plot.hist(bins=3, bottom=[2, 1, 3]) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.plot.hist(bins=3, bottom=[2, 1, 3]) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + non_numeric_pdf = self.pdf1.copy() + non_numeric_pdf.c = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"] + non_numeric_psdf = ps.from_pandas(non_numeric_pdf) + ax1 = non_numeric_pdf.plot.hist( + x=non_numeric_pdf.columns[0], y=non_numeric_pdf.columns[1], bins=3 + ) + bin1 = self.plot_to_base64(ax1) + ax2 = non_numeric_psdf.plot.hist( + x=non_numeric_pdf.columns[0], y=non_numeric_pdf.columns[1], bins=3 + ) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_hist_plot(pdf1, psdf1) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf1.columns = columns + psdf1.columns = columns + check_hist_plot(pdf1, psdf1) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43722): Enable DataFramePlotMatplotlibTests.test_kde_plot for pandas 2.0.0.", + ) + def test_kde_plot(self): + def moving_average(a, n=10): + ret = np.cumsum(a, dtype=float) + ret[n:] = ret[n:] - ret[:-n] + return ret[n - 1 :] / n + + def check_kde_plot(pdf, psdf, *args, **kwargs): + _, ax1 = plt.subplots(1, 1) + ax1 = pdf.plot.kde(*args, **kwargs) + _, ax2 = plt.subplots(1, 1) + ax2 = psdf.plot.kde(*args, **kwargs) + + try: + for i, (line1, line2) in enumerate(zip(ax1.get_lines(), ax2.get_lines())): + expected = line1.get_xydata().ravel() + actual = line2.get_xydata().ravel() + # TODO: Due to implementation difference, the output is different comparing + # to pandas'. We should identify the root cause of difference, and reduce + # the diff. + + # Note: Data is from 1 to 50. So, it smooths them by moving average and compares + # both. + self.assertTrue( + np.allclose(moving_average(actual), moving_average(expected), rtol=3.0) + ) + finally: + ax1.cla() + ax2.cla() + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_kde_plot(pdf1, psdf1, bw_method=0.3) + check_kde_plot(pdf1, psdf1, ind=[1, 2, 3], bw_method=3.0) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) + pdf1.columns = columns + pdf1.columns = columns + check_kde_plot(pdf1, psdf1, bw_method=0.3) + check_kde_plot(pdf1, psdf1, ind=[1, 2, 3], bw_method=3.0) + + +class DataFramePlotMatplotlibTests( + DataFramePlotMatplotlibTestsMixin, PandasOnSparkTestCase, TestUtils +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.plot.test_frame_plot_matplotlib import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py new file mode 100644 index 0000000..37469db --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py @@ -0,0 +1,285 @@ +# +# 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 unittest +import pprint + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import ( + have_plotly, + plotly_requirement_message, + PandasOnSparkTestCase, + TestUtils, +) +from pyspark.pandas.utils import name_like_string + +if have_plotly: + from plotly import express + import plotly.graph_objs as go + + +@unittest.skipIf(not have_plotly, plotly_requirement_message) +class DataFramePlotPlotlyTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + pd.set_option("plotting.backend", "plotly") + set_option("plotting.backend", "plotly") + set_option("plotting.max_rows", 2000) + set_option("plotting.sample_ratio", None) + + @classmethod + def tearDownClass(cls): + pd.reset_option("plotting.backend") + reset_option("plotting.backend") + reset_option("plotting.max_rows") + reset_option("plotting.sample_ratio") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], "b": [2, 3, 4, 5, 7, 9, 10, 15, 34, 45, 49]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10], + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + def test_line_plot(self): + def check_line_plot(pdf, psdf): + self.assertEqual(pdf.plot(kind="line"), psdf.plot(kind="line")) + self.assertEqual(pdf.plot.line(), psdf.plot.line()) + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_line_plot(pdf1, psdf1) + + def test_area_plot(self): + def check_area_plot(pdf, psdf): + self.assertEqual(pdf.plot(kind="area"), psdf.plot(kind="area")) + self.assertEqual(pdf.plot.area(), psdf.plot.area()) + + pdf = self.pdf1 + psdf = self.psdf1 + check_area_plot(pdf, psdf) + + def test_area_plot_y(self): + def check_area_plot_y(pdf, psdf, y): + self.assertEqual(pdf.plot.area(y=y), psdf.plot.area(y=y)) + + # test if frame area plot is correct when y is specified + pdf = pd.DataFrame( + { + "sales": [3, 2, 3, 9, 10, 6], + "signups": [5, 5, 6, 12, 14, 13], + "visits": [20, 42, 28, 62, 81, 50], + }, + index=pd.date_range(start="2018/01/01", end="2018/07/01", freq="M"), + ) + psdf = ps.from_pandas(pdf) + check_area_plot_y(pdf, psdf, y="sales") + + def test_barh_plot_with_x_y(self): + def check_barh_plot_with_x_y(pdf, psdf, x, y): + self.assertEqual(pdf.plot(kind="barh", x=x, y=y), psdf.plot(kind="barh", x=x, y=y)) + self.assertEqual(pdf.plot.barh(x=x, y=y), psdf.plot.barh(x=x, y=y)) + + # this is testing plot with specified x and y + pdf1 = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + psdf1 = ps.from_pandas(pdf1) + check_barh_plot_with_x_y(pdf1, psdf1, x="lab", y="val") + + def test_barh_plot(self): + def check_barh_plot(pdf, psdf): + self.assertEqual(pdf.plot(kind="barh"), psdf.plot(kind="barh")) + self.assertEqual(pdf.plot.barh(), psdf.plot.barh()) + + # this is testing when x or y is not assigned + pdf1 = pd.DataFrame({"lab": [20.1, 40.5, 60.6], "val": [10, 30, 20]}) + psdf1 = ps.from_pandas(pdf1) + check_barh_plot(pdf1, psdf1) + + def test_bar_plot(self): + def check_bar_plot(pdf, psdf): + self.assertEqual(pdf.plot(kind="bar"), psdf.plot(kind="bar")) + self.assertEqual(pdf.plot.bar(), psdf.plot.bar()) + + pdf1 = self.pdf1 + psdf1 = self.psdf1 + check_bar_plot(pdf1, psdf1) + + def test_bar_with_x_y(self): + # this is testing plot with specified x and y + pdf = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + psdf = ps.from_pandas(pdf) + + self.assertEqual( + pdf.plot(kind="bar", x="lab", y="val"), psdf.plot(kind="bar", x="lab", y="val") + ) + self.assertEqual(pdf.plot.bar(x="lab", y="val"), psdf.plot.bar(x="lab", y="val")) + + def test_scatter_plot(self): + def check_scatter_plot(pdf, psdf, x, y, c): + self.assertEqual(pdf.plot.scatter(x=x, y=y), psdf.plot.scatter(x=x, y=y)) + self.assertEqual( + pdf.plot(kind="scatter", x=x, y=y), psdf.plot(kind="scatter", x=x, y=y) + ) + + # check when keyword c is given as name of a column + self.assertEqual( + pdf.plot.scatter(x=x, y=y, c=c, s=50), psdf.plot.scatter(x=x, y=y, c=c, s=50) + ) + + # Use pandas scatter plot example + pdf1 = pd.DataFrame(np.random.rand(50, 4), columns=["a", "b", "c", "d"]) + psdf1 = ps.from_pandas(pdf1) + check_scatter_plot(pdf1, psdf1, x="a", y="b", c="c") + + def test_pie_plot(self): + def check_pie_plot(psdf): + pdf = psdf._to_pandas() + self.assertEqual( + psdf.plot(kind="pie", y=psdf.columns[0]), + express.pie(pdf, values="a", names=pdf.index), + ) + + self.assertEqual( + psdf.plot(kind="pie", values="a"), + express.pie(pdf, values="a"), + ) + + psdf1 = self.psdf1 + check_pie_plot(psdf1) + + # TODO: support multi-index columns + # columns = pd.MultiIndex.from_tuples([("x", "y"), ("y", "z")]) + # psdf1.columns = columns + # check_pie_plot(psdf1) + + # TODO: support multi-index + # psdf1 = ps.DataFrame( + # { + # "a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], + # "b": [2, 3, 4, 5, 7, 9, 10, 15, 34, 45, 49] + # }, + # index=pd.MultiIndex.from_tuples([("x", "y")] * 11), + # ) + # check_pie_plot(psdf1) + + def test_hist_layout_kwargs(self): + s = ps.Series([1, 3, 2]) + plt = s.plot.hist(title="Title", foo="xxx") + self.assertEqual(plt.layout.barmode, "stack") + self.assertEqual(plt.layout.title.text, "Title") + self.assertFalse(hasattr(plt.layout, "foo")) + + def test_hist_plot(self): + def check_hist_plot(psdf): + bins = np.array([1.0, 5.9, 10.8, 15.7, 20.6, 25.5, 30.4, 35.3, 40.2, 45.1, 50.0]) + data = [ + np.array([5.0, 4.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + np.array([4.0, 3.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0]), + ] + prev = bins[0] + text_bins = [] + for b in bins[1:]: + text_bins.append("[%s, %s)" % (prev, b)) + prev = b + text_bins[-1] = text_bins[-1][:-1] + "]" + bins = 0.5 * (bins[:-1] + bins[1:]) + name_a = name_like_string(psdf.columns[0]) + name_b = name_like_string(psdf.columns[1]) + bars = [ + go.Bar( + x=bins, + y=data[0], + name=name_a, + text=text_bins, + hovertemplate=("variable=" + name_a + "
    value=%{text}
    count=%{y}"), + ), + go.Bar( + x=bins, + y=data[1], + name=name_b, + text=text_bins, + hovertemplate=("variable=" + name_b + "
    value=%{text}
    count=%{y}"), + ), + ] + fig = go.Figure(data=bars, layout=go.Layout(barmode="stack")) + fig["layout"]["xaxis"]["title"] = "value" + fig["layout"]["yaxis"]["title"] = "count" + + self.assertEqual( + pprint.pformat(psdf.plot(kind="hist").to_dict()), pprint.pformat(fig.to_dict()) + ) + + psdf1 = self.psdf1 + check_hist_plot(psdf1) + + columns = pd.MultiIndex.from_tuples([("x", "y"), ("y", "z")]) + psdf1.columns = columns + check_hist_plot(psdf1) + + def test_kde_plot(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5], "b": [1, 3, 5, 7, 9], "c": [2, 4, 6, 8, 10]}) + + pdf = pd.DataFrame( + { + "Density": [ + 0.03515491, + 0.06834979, + 0.00663503, + 0.02372059, + 0.06834979, + 0.01806934, + 0.01806934, + 0.06834979, + 0.02372059, + ], + "names": ["a", "a", "a", "b", "b", "b", "c", "c", "c"], + "index": [-3.5, 5.5, 14.5, -3.5, 5.5, 14.5, -3.5, 5.5, 14.5], + } + ) + + actual = psdf.plot.kde(bw_method=5, ind=3) + + expected = express.line(pdf, x="index", y="Density", color="names") + expected["layout"]["xaxis"]["title"] = None + + self.assertEqual(pprint.pformat(actual.to_dict()), pprint.pformat(expected.to_dict())) + + +class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.plot.test_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot.py new file mode 100644 index 0000000..9daefbc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot.py @@ -0,0 +1,106 @@ +# +# 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 unittest + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.plot import PandasOnSparkPlotAccessor, BoxPlotBase +from pyspark.testing.pandasutils import have_plotly, plotly_requirement_message + + +class SeriesPlotTestsMixin: + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10] + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @unittest.skipIf(not have_plotly, plotly_requirement_message) + def test_plot_backends(self): + plot_backend = "plotly" + + with ps.option_context("plotting.backend", plot_backend): + self.assertEqual(ps.options.plotting.backend, plot_backend) + + module = PandasOnSparkPlotAccessor._get_plot_backend(plot_backend) + self.assertEqual(module.__name__, "pyspark.pandas.plot.plotly") + + def test_plot_backends_incorrect(self): + fake_plot_backend = "none_plotting_module" + + with ps.option_context("plotting.backend", fake_plot_backend): + self.assertEqual(ps.options.plotting.backend, fake_plot_backend) + + with self.assertRaises(ValueError): + PandasOnSparkPlotAccessor._get_plot_backend(fake_plot_backend) + + def test_box_summary(self): + def check_box_summary(psdf, pdf): + k = 1.5 + stats, fences = BoxPlotBase.compute_stats(psdf["a"], "a", whis=k, precision=0.01) + outliers = BoxPlotBase.outliers(psdf["a"], "a", *fences) + whiskers = BoxPlotBase.calc_whiskers("a", outliers) + fliers = BoxPlotBase.get_fliers("a", outliers, whiskers[0]) + + expected_mean = pdf["a"].mean() + expected_median = pdf["a"].median() + expected_q1 = np.percentile(pdf["a"], 25) + expected_q3 = np.percentile(pdf["a"], 75) + iqr = expected_q3 - expected_q1 + expected_fences = (expected_q1 - k * iqr, expected_q3 + k * iqr) + pdf["outlier"] = ~pdf["a"].between(fences[0], fences[1]) + expected_whiskers = ( + pdf.query("not outlier")["a"].min(), + pdf.query("not outlier")["a"].max(), + ) + expected_fliers = pdf.query("outlier")["a"].values + + self.assertEqual(expected_mean, stats["mean"]) + self.assertEqual(expected_median, stats["med"]) + self.assertEqual(expected_q1, stats["q1"] + 0.5) + self.assertEqual(expected_q3, stats["q3"] - 0.5) + self.assertEqual(expected_fences[0], fences[0] + 2.0) + self.assertEqual(expected_fences[1], fences[1] - 2.0) + self.assertEqual(expected_whiskers[0], whiskers[0]) + self.assertEqual(expected_whiskers[1], whiskers[1]) + self.assertEqual(expected_fliers, fliers) + + check_box_summary(self.psdf1, self.pdf1) + check_box_summary(-self.psdf1, -self.pdf1) + + +class SeriesPlotTests(SeriesPlotTestsMixin, unittest.TestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.plot.test_series_plot import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py new file mode 100644 index 0000000..c98c1ae --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py @@ -0,0 +1,409 @@ +# +# 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 base64 +from io import BytesIO +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import ( + have_matplotlib, + matplotlib_requirement_message, + PandasOnSparkTestCase, + TestUtils, +) + +if have_matplotlib: + import matplotlib + from matplotlib import pyplot as plt + + matplotlib.use("agg") + + +@unittest.skipIf(not have_matplotlib, matplotlib_requirement_message) +class SeriesPlotMatplotlibTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + pd.set_option("plotting.backend", "matplotlib") + set_option("plotting.backend", "matplotlib") + set_option("plotting.max_rows", 1000) + + @classmethod + def tearDownClass(cls): + pd.reset_option("plotting.backend") + reset_option("plotting.backend") + reset_option("plotting.max_rows") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10] + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.range(1002) + + @property + def pdf2(self): + return self.psdf2._to_pandas() + + @staticmethod + def plot_to_base64(ax): + bytes_data = BytesIO() + ax.figure.savefig(bytes_data, format="png") + bytes_data.seek(0) + b64_data = base64.b64encode(bytes_data.read()) + plt.close(ax.figure) + return b64_data + + def test_bar_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + ax1 = pdf["a"].plot(kind="bar", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="bar", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot(kind="bar", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="bar", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_bar_plot_limited(self): + pdf = self.pdf2 + psdf = self.psdf2 + + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["id"][:1000].plot.bar(colormap="Paired") + ax1.text( + 1, + 1, + "showing top 1000 elements only", + size=6, + ha="right", + va="bottom", + transform=ax1.transAxes, + ) + bin1 = self.plot_to_base64(ax1) + + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["id"].plot.bar(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + + self.assertEqual(bin1, bin2) + + def test_pie_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + ax1 = pdf["a"].plot.pie(colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot.pie(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot(kind="pie", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="pie", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_pie_plot_limited(self): + pdf = self.pdf2 + psdf = self.psdf2 + + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["id"][:1000].plot.pie(colormap="Paired") + ax1.text( + 1, + 1, + "showing top 1000 elements only", + size=6, + ha="right", + va="bottom", + transform=ax1.transAxes, + ) + bin1 = self.plot_to_base64(ax1) + + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["id"].plot.pie(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + + self.assertEqual(bin1, bin2) + + def test_line_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + ax1 = pdf["a"].plot(kind="line", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="line", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot.line(colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot.line(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_barh_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + ax1 = pdf["a"].plot(kind="barh", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="barh", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_barh_plot_limited(self): + pdf = self.pdf2 + psdf = self.psdf2 + + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["id"][:1000].plot.barh(colormap="Paired") + ax1.text( + 1, + 1, + "showing top 1000 elements only", + size=6, + ha="right", + va="bottom", + transform=ax1.transAxes, + ) + bin1 = self.plot_to_base64(ax1) + + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["id"].plot.barh(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + + self.assertEqual(bin1, bin2) + + def test_hist(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10] + ) + + psdf = ps.from_pandas(pdf) + + def plot_to_base64(ax): + bytes_data = BytesIO() + ax.figure.savefig(bytes_data, format="png") + bytes_data.seek(0) + b64_data = base64.b64encode(bytes_data.read()) + plt.close(ax.figure) + return b64_data + + _, ax1 = plt.subplots(1, 1) + # Using plot.hist() because pandas changes ticks props when called hist() + ax1 = pdf["a"].plot.hist() + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["a"].hist() + self.assert_eq(plot_to_base64(ax1), plot_to_base64(ax2)) + + def test_hist_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["a"].plot.hist() + bin1 = self.plot_to_base64(ax1) + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["a"].plot.hist() + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot.hist(bins=15) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot.hist(bins=15) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot(kind="hist", bins=15) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot(kind="hist", bins=15) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["a"].plot.hist(bins=3, bottom=[2, 1, 3]) + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["a"].plot.hist(bins=3, bottom=[2, 1, 3]) + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_area_plot(self): + pdf = pd.DataFrame( + { + "sales": [3, 2, 3, 9, 10, 6], + "signups": [5, 5, 6, 12, 14, 13], + "visits": [20, 42, 28, 62, 81, 50], + }, + index=pd.date_range(start="2018/01/01", end="2018/07/01", freq="M"), + ) + psdf = ps.from_pandas(pdf) + + ax1 = pdf["sales"].plot(kind="area", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["sales"].plot(kind="area", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + ax1 = pdf["sales"].plot.area(colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf["sales"].plot.area(colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + # just a sanity check for df.col type + ax1 = pdf.sales.plot(kind="area", colormap="Paired") + bin1 = self.plot_to_base64(ax1) + ax2 = psdf.sales.plot(kind="area", colormap="Paired") + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + def test_box_plot(self): + def check_box_plot(pser, psser, *args, **kwargs): + _, ax1 = plt.subplots(1, 1) + ax1 = pser.plot.box(*args, **kwargs) + _, ax2 = plt.subplots(1, 1) + ax2 = psser.plot.box(*args, **kwargs) + + diffs = [ + np.array([0, 0.5, 0, 0.5, 0, -0.5, 0, -0.5, 0, 0.5]), + np.array([0, 0.5, 0, 0]), + np.array([0, -0.5, 0, 0]), + ] + + try: + for i, (line1, line2) in enumerate(zip(ax1.get_lines(), ax2.get_lines())): + expected = line1.get_xydata().ravel() + actual = line2.get_xydata().ravel() + if i < 3: + actual += diffs[i] + self.assert_eq(pd.Series(expected), pd.Series(actual)) + finally: + ax1.cla() + ax2.cla() + + # Non-named Series + pser = pd.Series([1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], [0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10]) + psser = ps.from_pandas(pser) + + spec = [(self.pdf1.a, self.psdf1.a), (pser, psser)] + + for p, k in spec: + check_box_plot(p, k) + check_box_plot(p, k, showfliers=True) + check_box_plot(p, k, sym="") + check_box_plot(p, k, sym=".", color="r") + check_box_plot(p, k, use_index=False, labels=["Test"]) + check_box_plot(p, k, usermedians=[2.0]) + check_box_plot(p, k, conf_intervals=[(1.0, 3.0)]) + + val = (1, 3) + self.assertRaises( + ValueError, lambda: check_box_plot(self.pdf1, self.psdf1, usermedians=[2.0, 3.0]) + ) + self.assertRaises( + ValueError, lambda: check_box_plot(self.pdf1, self.psdf1, conf_intervals=[val, val]) + ) + self.assertRaises( + ValueError, lambda: check_box_plot(self.pdf1, self.psdf1, conf_intervals=[(1,)]) + ) + + def test_kde_plot(self): + def moving_average(a, n=10): + ret = np.cumsum(a, dtype=float) + ret[n:] = ret[n:] - ret[:-n] + return ret[n - 1 :] / n + + def check_kde_plot(pdf, psdf, *args, **kwargs): + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["a"].plot.kde(*args, **kwargs) + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["a"].plot.kde(*args, **kwargs) + + try: + for i, (line1, line2) in enumerate(zip(ax1.get_lines(), ax2.get_lines())): + expected = line1.get_xydata().ravel() + actual = line2.get_xydata().ravel() + # TODO: Due to implementation difference, the output is different comparing + # to pandas'. We should identify the root cause of difference, and reduce + # the diff. + + # Note: Data is from 1 to 50. So, it smooths them by moving average and compares + # both. + self.assertTrue( + np.allclose(moving_average(actual), moving_average(expected), rtol=3) + ) + finally: + ax1.cla() + ax2.cla() + + check_kde_plot(self.pdf1, self.psdf1, bw_method=0.3) + check_kde_plot(self.pdf1, self.psdf1, ind=[1, 2, 3, 4, 5], bw_method=3.0) + + def test_empty_hist(self): + pdf = self.pdf1.assign(categorical="A") + psdf = ps.from_pandas(pdf) + psser = psdf["categorical"] + + with self.assertRaisesRegex(TypeError, "Empty 'DataFrame': no numeric data to plot"): + psser.plot.hist() + + def test_single_value_hist(self): + pdf = self.pdf1.assign(single=2) + psdf = ps.from_pandas(pdf) + + _, ax1 = plt.subplots(1, 1) + ax1 = pdf["single"].plot.hist() + bin1 = self.plot_to_base64(ax1) + _, ax2 = plt.subplots(1, 1) + ax2 = psdf["single"].plot.hist() + bin2 = self.plot_to_base64(ax2) + self.assertEqual(bin1, bin2) + + +class SeriesPlotMatplotlibTests(SeriesPlotMatplotlibTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.plot.test_series_plot_matplotlib import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py new file mode 100644 index 0000000..1aa175f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py @@ -0,0 +1,247 @@ +# +# 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 unittest +import pprint + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.utils import name_like_string +from pyspark.testing.pandasutils import ( + have_plotly, + plotly_requirement_message, + PandasOnSparkTestCase, + TestUtils, +) + +if have_plotly: + from plotly import express + import plotly.graph_objs as go + + +@unittest.skipIf(not have_plotly, plotly_requirement_message) +class SeriesPlotPlotlyTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + pd.set_option("plotting.backend", "plotly") + set_option("plotting.backend", "plotly") + set_option("plotting.max_rows", 1000) + set_option("plotting.sample_ratio", None) + + @classmethod + def tearDownClass(cls): + pd.reset_option("plotting.backend") + reset_option("plotting.backend") + reset_option("plotting.max_rows") + reset_option("plotting.sample_ratio") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10] + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.range(1002) + + @property + def pdf2(self): + return self.psdf2._to_pandas() + + def test_bar_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + self.assertEqual(pdf["a"].plot(kind="bar"), psdf["a"].plot(kind="bar")) + self.assertEqual(pdf["a"].plot.bar(), psdf["a"].plot.bar()) + + def test_line_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + self.assertEqual(pdf["a"].plot(kind="line"), psdf["a"].plot(kind="line")) + self.assertEqual(pdf["a"].plot.line(), psdf["a"].plot.line()) + + def test_barh_plot(self): + pdf = self.pdf1 + psdf = self.psdf1 + + self.assertEqual(pdf["a"].plot(kind="barh"), psdf["a"].plot(kind="barh")) + + def test_area_plot(self): + pdf = pd.DataFrame( + { + "sales": [3, 2, 3, 9, 10, 6], + "signups": [5, 5, 6, 12, 14, 13], + "visits": [20, 42, 28, 62, 81, 50], + }, + index=pd.date_range(start="2018/01/01", end="2018/07/01", freq="M"), + ) + psdf = ps.from_pandas(pdf) + + self.assertEqual(pdf["sales"].plot(kind="area"), psdf["sales"].plot(kind="area")) + self.assertEqual(pdf["sales"].plot.area(), psdf["sales"].plot.area()) + + # just a sanity check for df.col type + self.assertEqual(pdf.sales.plot(kind="area"), psdf.sales.plot(kind="area")) + + def test_pie_plot(self): + psdf = self.psdf1 + pdf = psdf._to_pandas() + self.assertEqual( + psdf["a"].plot(kind="pie"), + express.pie(pdf, values=pdf.columns[0], names=pdf.index), + ) + + # TODO: support multi-index columns + # columns = pd.MultiIndex.from_tuples([("x", "y")]) + # psdf.columns = columns + # pdf.columns = columns + # self.assertEqual( + # psdf[("x", "y")].plot(kind="pie"), + # express.pie(pdf, values=pdf.iloc[:, 0].to_numpy(), names=pdf.index.to_numpy()), + # ) + + # TODO: support multi-index + # psdf = ps.DataFrame( + # { + # "a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50], + # "b": [2, 3, 4, 5, 7, 9, 10, 15, 34, 45, 49] + # }, + # index=pd.MultiIndex.from_tuples([("x", "y")] * 11), + # ) + # pdf = psdf._to_pandas() + # self.assertEqual( + # psdf["a"].plot(kind="pie"), express.pie(pdf, values=pdf.columns[0], names=pdf.index), + # ) + + def test_hist_plot(self): + def check_hist_plot(psser): + bins = np.array([1.0, 5.9, 10.8, 15.7, 20.6, 25.5, 30.4, 35.3, 40.2, 45.1, 50.0]) + data = np.array([5.0, 4.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) + prev = bins[0] + text_bins = [] + for b in bins[1:]: + text_bins.append("[%s, %s)" % (prev, b)) + prev = b + text_bins[-1] = text_bins[-1][:-1] + "]" + bins = 0.5 * (bins[:-1] + bins[1:]) + name_a = name_like_string(psser.name) + bars = [ + go.Bar( + x=bins, + y=data, + name=name_a, + text=text_bins, + hovertemplate=("variable=" + name_a + "
    value=%{text}
    count=%{y}"), + ), + ] + fig = go.Figure(data=bars, layout=go.Layout(barmode="stack")) + fig["layout"]["xaxis"]["title"] = "value" + fig["layout"]["yaxis"]["title"] = "count" + + self.assertEqual( + pprint.pformat(psser.plot(kind="hist").to_dict()), pprint.pformat(fig.to_dict()) + ) + + psdf1 = self.psdf1 + check_hist_plot(psdf1["a"]) + + columns = pd.MultiIndex.from_tuples([("x", "y")]) + psdf1.columns = columns + check_hist_plot(psdf1[("x", "y")]) + + def test_pox_plot(self): + def check_pox_plot(psser): + fig = go.Figure() + fig.add_trace( + go.Box( + name=name_like_string(psser.name), + q1=[3], + median=[6], + q3=[9], + mean=[10.0], + lowerfence=[1], + upperfence=[15], + y=[[50]], + boxpoints="suspectedoutliers", + notched=False, + ) + ) + fig["layout"]["xaxis"]["title"] = name_like_string(psser.name) + fig["layout"]["yaxis"]["title"] = "value" + + self.assertEqual( + pprint.pformat(psser.plot(kind="box").to_dict()), pprint.pformat(fig.to_dict()) + ) + + psdf1 = self.psdf1 + check_pox_plot(psdf1["a"]) + + columns = pd.MultiIndex.from_tuples([("x", "y")]) + psdf1.columns = columns + check_pox_plot(psdf1[("x", "y")]) + + def test_pox_plot_arguments(self): + with self.assertRaisesRegex(ValueError, "does not support"): + self.psdf1.a.plot.box(boxpoints="all") + with self.assertRaisesRegex(ValueError, "does not support"): + self.psdf1.a.plot.box(notched=True) + self.psdf1.a.plot.box(hovertext="abc") # other arguments should not throw an exception + + def test_kde_plot(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5]}) + pdf = pd.DataFrame( + { + "Density": [0.05709372, 0.07670272, 0.05709372], + "names": ["a", "a", "a"], + "index": [-1.0, 3.0, 7.0], + } + ) + + actual = psdf.a.plot.kde(bw_method=5, ind=3) + + expected = express.line(pdf, x="index", y="Density") + expected["layout"]["xaxis"]["title"] = None + + self.assertEqual(pprint.pformat(actual.to_dict()), pprint.pformat(expected.to_dict())) + + +class SeriesPlotPlotlyTests(SeriesPlotPlotlyTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.plot.test_series_plot_plotly import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/__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/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_all_any.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_all_any.py new file mode 100644 index 0000000..6663675 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_all_any.py @@ -0,0 +1,105 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesAllAnyMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_all(self): + for pser in [ + pd.Series([True, True], name="x"), + pd.Series([True, False], name="x"), + pd.Series([0, 1], name="x"), + pd.Series([1, 2, 3], name="x"), + pd.Series([np.nan, 0, 1], name="x"), + pd.Series([np.nan, 1, 2, 3], name="x"), + pd.Series([True, True, None], name="x"), + pd.Series([True, False, None], name="x"), + pd.Series([], name="x"), + pd.Series([np.nan], name="x"), + pd.Series([np.nan, np.nan], name="x"), + pd.Series([None], name="x"), + pd.Series([None, None], name="x"), + ]: + psser = ps.from_pandas(pser) + self.assert_eq(psser.all(), pser.all()) + self.assert_eq(psser.all(skipna=False), pser.all(skipna=False)) + self.assert_eq(psser.all(skipna=True), pser.all(skipna=True)) + + pser = pd.Series([1, 2, 3, 4], name="x") + psser = ps.from_pandas(pser) + + self.assert_eq((psser % 2 == 0).all(), (pser % 2 == 0).all()) + + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psser.all(axis=1) + + def test_any(self): + for pser in [ + pd.Series([False, False], name="x"), + pd.Series([True, False], name="x"), + pd.Series([0, 1], name="x"), + pd.Series([1, 2, 3], name="x"), + pd.Series([True, True, None], name="x"), + pd.Series([True, False, None], name="x"), + pd.Series([], name="x"), + pd.Series([np.nan], name="x"), + ]: + psser = ps.from_pandas(pser) + self.assert_eq(psser.any(), pser.any()) + + pser = pd.Series([1, 2, 3, 4], name="x") + psser = ps.from_pandas(pser) + + self.assert_eq((psser % 2 == 0).any(), (pser % 2 == 0).any()) + + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psser.any(axis=1) + + +class SeriesAllAnyTests(SeriesAllAnyMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_all_any import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_arg_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_arg_ops.py new file mode 100644 index 0000000..5b1aa24 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_arg_ops.py @@ -0,0 +1,197 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesArgOpsMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_argsort(self): + # Without null values + pser = pd.Series([0, -100, 50, 100, 20], index=["A", "B", "C", "D", "E"]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # MultiIndex + pser.index = pd.MultiIndex.from_tuples( + [("a", "v"), ("b", "w"), ("c", "x"), ("d", "y"), ("e", "z")] + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # With name + pser.name = "Koalas" + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # Series from Index + pidx = pd.Index([4.0, -6.0, 2.0, -100.0, 11.0, 20.0, 1.0, -99.0]) + psidx = ps.from_pandas(pidx) + self.assert_eq( + pidx.to_series().argsort().sort_index(), psidx.to_series().argsort().sort_index() + ) + self.assert_eq( + (-pidx.to_series()).argsort().sort_index(), (-psidx.to_series()).argsort().sort_index() + ) + + # Series from Index with name + pidx.name = "Koalas" + psidx = ps.from_pandas(pidx) + self.assert_eq( + pidx.to_series().argsort().sort_index(), psidx.to_series().argsort().sort_index() + ) + self.assert_eq( + (-pidx.to_series()).argsort().sort_index(), (-psidx.to_series()).argsort().sort_index() + ) + + # Series from DataFrame + pdf = pd.DataFrame({"A": [4.0, -6.0, 2.0, np.nan, -100.0, 11.0, 20.0, np.nan, 1.0, -99.0]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.A.argsort().sort_index(), psdf.A.argsort().sort_index()) + self.assert_eq((-pdf.A).argsort().sort_index(), (-psdf.A).argsort().sort_index()) + + # With null values + pser = pd.Series([0, -100, np.nan, 100, np.nan], index=["A", "B", "C", "D", "E"]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # MultiIndex with null values + pser.index = pd.MultiIndex.from_tuples( + [("a", "v"), ("b", "w"), ("c", "x"), ("d", "y"), ("e", "z")] + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # With name with null values + pser.name = "Koalas" + psser = ps.from_pandas(pser) + self.assert_eq(pser.argsort().sort_index(), psser.argsort().sort_index()) + self.assert_eq((-pser).argsort().sort_index(), (-psser).argsort().sort_index()) + + # Series from Index with null values + pidx = pd.Index([4.0, -6.0, 2.0, np.nan, -100.0, 11.0, 20.0, np.nan, 1.0, -99.0]) + psidx = ps.from_pandas(pidx) + self.assert_eq( + pidx.to_series().argsort().sort_index(), psidx.to_series().argsort().sort_index() + ) + self.assert_eq( + (-pidx.to_series()).argsort().sort_index(), (-psidx.to_series()).argsort().sort_index() + ) + + # Series from Index with name with null values + pidx.name = "Koalas" + psidx = ps.from_pandas(pidx) + self.assert_eq( + pidx.to_series().argsort().sort_index(), psidx.to_series().argsort().sort_index() + ) + self.assert_eq( + (-pidx.to_series()).argsort().sort_index(), (-psidx.to_series()).argsort().sort_index() + ) + + # Series from DataFrame with null values + pdf = pd.DataFrame({"A": [4.0, -6.0, 2.0, np.nan, -100.0, 11.0, 20.0, np.nan, 1.0, -99.0]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.A.argsort().sort_index(), psdf.A.argsort().sort_index()) + self.assert_eq((-pdf.A).argsort().sort_index(), (-psdf.A).argsort().sort_index()) + + def test_argmin_argmax(self): + pser = pd.Series( + { + "Corn Flakes": 100.0, + "Almond Delight": 110.0, + "Cinnamon Toast Crunch": 120.0, + "Cocoa Puff": 110.0, + "Expensive Flakes": 120.0, + "Cheap Flakes": 100.0, + }, + name="Koalas", + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argmin(), psser.argmin()) + self.assert_eq(pser.argmax(), psser.argmax()) + self.assert_eq(pser.argmin(skipna=False), psser.argmin(skipna=False)) + self.assert_eq(pser.argmax(skipna=False), psser.argmax(skipna=False)) + self.assert_eq(pser.argmax(skipna=False), psser.argmax(skipna=False)) + self.assert_eq((pser + 1).argmax(skipna=False), (psser + 1).argmax(skipna=False)) + self.assert_eq(pser.argmin(skipna=False), psser.argmin(skipna=False)) + self.assert_eq((pser + 1).argmin(skipna=False), (psser + 1).argmin(skipna=False)) + + # MultiIndex + pser.index = pd.MultiIndex.from_tuples( + [("a", "t"), ("b", "u"), ("c", "v"), ("d", "w"), ("e", "x"), ("f", "u")] + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.argmin(), psser.argmin()) + self.assert_eq(pser.argmax(), psser.argmax()) + self.assert_eq(pser.argmax(skipna=False), psser.argmax(skipna=False)) + + pser2 = pd.Series([np.NaN, 1.0, 2.0, np.NaN]) + psser2 = ps.from_pandas(pser2) + self.assert_eq(pser2.argmin(), psser2.argmin()) + self.assert_eq(pser2.argmax(), psser2.argmax()) + self.assert_eq(pser2.argmin(skipna=False), psser2.argmin(skipna=False)) + self.assert_eq(pser2.argmax(skipna=False), psser2.argmax(skipna=False)) + + # Null Series + self.assert_eq(pd.Series([np.nan]).argmin(), ps.Series([np.nan]).argmin()) + self.assert_eq(pd.Series([np.nan]).argmax(), ps.Series([np.nan]).argmax()) + self.assert_eq( + pd.Series([np.nan]).argmax(skipna=False), ps.Series([np.nan]).argmax(skipna=False) + ) + + with self.assertRaisesRegex(ValueError, "attempt to get argmin of an empty sequence"): + ps.Series([]).argmin() + with self.assertRaisesRegex(ValueError, "attempt to get argmax of an empty sequence"): + ps.Series([]).argmax() + with self.assertRaisesRegex(ValueError, "axis can only be 0 or 'index'"): + psser.argmax(axis=1) + with self.assertRaisesRegex(ValueError, "axis can only be 0 or 'index'"): + psser.argmin(axis=1) + + +class SeriesArgOpsTests(SeriesArgOpsMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_arg_ops import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_of.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_of.py new file mode 100644 index 0000000..552176a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_of.py @@ -0,0 +1,137 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesAsOfMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_asof(self): + pser = pd.Series([1, 2, np.nan, 4], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.asof(20), pser.asof(20)) + self.assert_eq(psser.asof([5, 20]).sort_index(), pser.asof([5, 20]).sort_index()) + self.assert_eq(psser.asof(100), pser.asof(100)) + self.assert_eq(repr(psser.asof(-100)), repr(pser.asof(-100))) + self.assert_eq(psser.asof([-100, 100]).sort_index(), pser.asof([-100, 100]).sort_index()) + + # where cannot be an Index, Series or a DataFrame + self.assertRaises(ValueError, lambda: psser.asof(ps.Index([-100, 100]))) + self.assertRaises(ValueError, lambda: psser.asof(ps.Series([-100, 100]))) + self.assertRaises(ValueError, lambda: psser.asof(ps.DataFrame({"A": [1, 2, 3]}))) + # asof is not supported for a MultiIndex + pser.index = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c"), ("y", "d")]) + psser = ps.from_pandas(pser) + self.assertRaises(ValueError, lambda: psser.asof(20)) + # asof requires a sorted index (More precisely, should be a monotonic increasing) + psser = ps.Series([1, 2, np.nan, 4], index=[10, 30, 20, 40], name="Koalas") + self.assertRaises(ValueError, lambda: psser.asof(20)) + psser = ps.Series([1, 2, np.nan, 4], index=[40, 30, 20, 10], name="Koalas") + self.assertRaises(ValueError, lambda: psser.asof(20)) + + pidx = pd.DatetimeIndex(["2013-12-31", "2014-01-02", "2014-01-03"]) + pser = pd.Series([1, 2, np.nan], index=pidx) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.asof("2014-01-01"), pser.asof("2014-01-01")) + self.assert_eq(psser.asof("2014-01-02"), pser.asof("2014-01-02")) + self.assert_eq(repr(psser.asof("1999-01-02")), repr(pser.asof("1999-01-02"))) + + # SPARK-37482: Skip check monotonic increasing for Series.asof with 'compute.eager_check' + pser = pd.Series([1, 2, np.nan, 4], index=[10, 30, 20, 40]) + psser = ps.from_pandas(pser) + + with ps.option_context("compute.eager_check", False): + self.assert_eq(psser.asof(20), 1.0) + + pser = pd.Series([1, 2, np.nan, 4], index=[40, 30, 20, 10]) + psser = ps.from_pandas(pser) + + with ps.option_context("compute.eager_check", False): + self.assert_eq(psser.asof(20), 4.0) + + pser = pd.Series([2, 1, np.nan, 4], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof([5, 20]), pser.asof([5, 20])) + + pser = pd.Series([4, np.nan, np.nan, 2], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof([5, 100]), pser.asof([5, 100])) + + pser = pd.Series([np.nan, 4, 1, 2], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof([5, 35]), pser.asof([5, 35])) + + pser = pd.Series([2, 1, np.nan, 4], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof([25, 25]), pser.asof([25, 25])) + + pser = pd.Series([2, 1, np.nan, 4], index=["a", "b", "c", "d"], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof(["a", "d"]), pser.asof(["a", "d"])) + + pser = pd.Series( + [2, 1, np.nan, 4], + index=[ + pd.Timestamp(2020, 1, 1), + pd.Timestamp(2020, 2, 2), + pd.Timestamp(2020, 3, 3), + pd.Timestamp(2020, 4, 4), + ], + name="Koalas", + ) + psser = ps.from_pandas(pser) + self.assert_eq( + psser.asof([pd.Timestamp(2020, 1, 1)]), + pser.asof([pd.Timestamp(2020, 1, 1)]), + ) + + pser = pd.Series([2, np.nan, 1, 4], index=[10, 20, 30, 40], name="Koalas") + psser = ps.from_pandas(pser) + self.assert_eq(psser.asof(np.nan), pser.asof(np.nan)) + self.assert_eq(psser.asof([np.nan, np.nan]), pser.asof([np.nan, np.nan])) + self.assert_eq(psser.asof([10, np.nan]), pser.asof([10, np.nan])) + + +class SeriesAsOfTests(SeriesAsOfMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_as_of import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_type.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_type.py new file mode 100644 index 0000000..d2ff9e4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_as_type.py @@ -0,0 +1,227 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.pandas.typedef.typehints import ( + extension_dtypes, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) + + +class SeriesAsTypeMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43466): Enable SeriesTests.test_astype for pandas 2.0.0.", + ) + def test_astype(self): + psers = [pd.Series([10, 20, 15, 30, 45], name="x")] + + if extension_dtypes_available: + psers.append(pd.Series([10, 20, 15, 30, 45], name="x", dtype="Int64")) + if extension_float_dtypes_available: + psers.append(pd.Series([10, 20, 15, 30, 45], name="x", dtype="Float64")) + + for pser in psers: + self._test_numeric_astype(pser) + + pser = pd.Series([10, 20, 15, 30, 45, None, np.nan], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.astype(bool), pser.astype(bool)) + self.assert_eq(psser.astype(str), pser.astype(str)) + + pser = pd.Series(["hi", "hi ", " ", " \t", "", None], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.astype(bool), pser.astype(bool)) + if LooseVersion("1.1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.1.4"): + # a pandas bug: https://github.com/databricks/koalas/pull/1818#issuecomment-703961980 + self.assert_eq(psser.astype(str).tolist(), ["hi", "hi ", " ", " \t", "", "None"]) + else: + self.assert_eq(psser.astype(str), pser.astype(str)) + self.assert_eq(psser.str.strip().astype(bool), pser.str.strip().astype(bool)) + + if extension_object_dtypes_available: + from pandas import StringDtype + + self._check_extension(psser.astype("string"), pser.astype("string")) + self._check_extension(psser.astype(StringDtype()), pser.astype(StringDtype())) + + pser = pd.Series([True, False, None], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.astype(bool), pser.astype(bool)) + self.assert_eq(psser.astype(str), pser.astype(str)) + + if extension_object_dtypes_available: + from pandas import BooleanDtype, StringDtype + + self._check_extension(psser.astype("boolean"), pser.astype("boolean")) + self._check_extension(psser.astype(BooleanDtype()), pser.astype(BooleanDtype())) + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self._check_extension(psser.astype("string"), pser.astype("string")) + self._check_extension(psser.astype(StringDtype()), pser.astype(StringDtype())) + else: + self._check_extension( + psser.astype("string"), + pd.Series(["True", "False", None], name="x", dtype="string"), + ) + self._check_extension( + psser.astype(StringDtype()), + pd.Series(["True", "False", None], name="x", dtype=StringDtype()), + ) + + pser = pd.Series(["2020-10-27 00:00:01", None], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.astype(np.datetime64), pser.astype(np.datetime64)) + self.assert_eq(psser.astype("datetime64[ns]"), pser.astype("datetime64[ns]")) + self.assert_eq(psser.astype("M"), pser.astype("M")) + self.assert_eq(psser.astype("M").astype(str), pser.astype("M").astype(str)) + # Comment out the below test cause because pandas returns `NaT` or `nan` randomly + # self.assert_eq( + # psser.astype("M").dt.date.astype(str), pser.astype("M").dt.date.astype(str) + # ) + + if extension_object_dtypes_available: + from pandas import StringDtype + + # The behavior of casting datetime to nullable string is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self._check_extension( + psser.astype("M").astype("string"), pser.astype("M").astype("string") + ) + self._check_extension( + psser.astype("M").astype(StringDtype()), pser.astype("M").astype(StringDtype()) + ) + else: + expected = ps.Series(["2020-10-27 00:00:01", None], name="x", dtype="string") + self._check_extension(psser.astype("M").astype("string"), expected) + self._check_extension(psser.astype("M").astype(StringDtype()), expected) + + with self.assertRaisesRegex(TypeError, "not understood"): + psser.astype("int63") + + def _test_numeric_astype(self, pser): + psser = ps.Series(pser) + + self.assert_eq(psser.astype(int), pser.astype(int)) + self.assert_eq(psser.astype(np.int8), pser.astype(np.int8)) + self.assert_eq(psser.astype(np.int16), pser.astype(np.int16)) + self.assert_eq(psser.astype(np.int32), pser.astype(np.int32)) + self.assert_eq(psser.astype(np.int64), pser.astype(np.int64)) + self.assert_eq(psser.astype(np.byte), pser.astype(np.byte)) + self.assert_eq(psser.astype("int"), pser.astype("int")) + self.assert_eq(psser.astype("int8"), pser.astype("int8")) + self.assert_eq(psser.astype("int16"), pser.astype("int16")) + self.assert_eq(psser.astype("int32"), pser.astype("int32")) + self.assert_eq(psser.astype("int64"), pser.astype("int64")) + self.assert_eq(psser.astype("b"), pser.astype("b")) + self.assert_eq(psser.astype("byte"), pser.astype("byte")) + self.assert_eq(psser.astype("i"), pser.astype("i")) + self.assert_eq(psser.astype("long"), pser.astype("long")) + self.assert_eq(psser.astype("short"), pser.astype("short")) + self.assert_eq(psser.astype(np.float32), pser.astype(np.float32)) + self.assert_eq(psser.astype(np.float64), pser.astype(np.float64)) + self.assert_eq(psser.astype("float"), pser.astype("float")) + self.assert_eq(psser.astype("float32"), pser.astype("float32")) + self.assert_eq(psser.astype("float64"), pser.astype("float64")) + self.assert_eq(psser.astype("double"), pser.astype("double")) + self.assert_eq(psser.astype("f"), pser.astype("f")) + self.assert_eq(psser.astype(bool), pser.astype(bool)) + self.assert_eq(psser.astype("bool"), pser.astype("bool")) + self.assert_eq(psser.astype("?"), pser.astype("?")) + self.assert_eq(psser.astype(np.unicode_), pser.astype(np.unicode_)) + self.assert_eq(psser.astype("str"), pser.astype("str")) + self.assert_eq(psser.astype("U"), pser.astype("U")) + + if extension_dtypes_available: + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + + self._check_extension(psser.astype("Int8"), pser.astype("Int8")) + self._check_extension(psser.astype("Int16"), pser.astype("Int16")) + self._check_extension(psser.astype("Int32"), pser.astype("Int32")) + self._check_extension(psser.astype("Int64"), pser.astype("Int64")) + self._check_extension(psser.astype(Int8Dtype()), pser.astype(Int8Dtype())) + self._check_extension(psser.astype(Int16Dtype()), pser.astype(Int16Dtype())) + self._check_extension(psser.astype(Int32Dtype()), pser.astype(Int32Dtype())) + self._check_extension(psser.astype(Int64Dtype()), pser.astype(Int64Dtype())) + + if extension_object_dtypes_available: + from pandas import StringDtype + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self._check_extension(psser.astype("string"), pser.astype("string")) + self._check_extension(psser.astype(StringDtype()), pser.astype(StringDtype())) + else: + self._check_extension( + psser.astype("string"), + pd.Series(["10", "20", "15", "30", "45"], name="x", dtype="string"), + ) + self._check_extension( + psser.astype(StringDtype()), + pd.Series(["10", "20", "15", "30", "45"], name="x", dtype=StringDtype()), + ) + + if extension_float_dtypes_available: + from pandas import Float32Dtype, Float64Dtype + + self._check_extension(psser.astype("Float32"), pser.astype("Float32")) + self._check_extension(psser.astype("Float64"), pser.astype("Float64")) + self._check_extension(psser.astype(Float32Dtype()), pser.astype(Float32Dtype())) + self._check_extension(psser.astype(Float64Dtype()), pser.astype(Float64Dtype())) + + def _check_extension(self, psser, pser): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(psser, pser, check_exact=False) + self.assertTrue(isinstance(psser.dtype, extension_dtypes)) + else: + self.assert_eq(psser, pser) + + +class SeriesAsTypeTests(SeriesAsTypeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_as_type import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_compute.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_compute.py new file mode 100644 index 0000000..2fbdaef --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_compute.py @@ -0,0 +1,634 @@ +# +# 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 unittest +from distutils.version import LooseVersion +from itertools import product + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesComputeMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_duplicated(self): + for pser in [ + pd.Series(["beetle", None, "beetle", None, "lama", "beetle"], name="objects"), + pd.Series([1, np.nan, 1, np.nan], name="numbers"), + pd.Series( + [ + pd.Timestamp("2022-01-01"), + pd.Timestamp("2022-02-02"), + pd.Timestamp("2022-01-01"), + pd.Timestamp("2022-02-02"), + ], + name="times", + ), + ]: + psser = ps.from_pandas(pser) + self.assert_eq(psser.duplicated().sort_index(), pser.duplicated()) + self.assert_eq( + psser.duplicated(keep="first").sort_index(), pser.duplicated(keep="first") + ) + self.assert_eq(psser.duplicated(keep="last").sort_index(), pser.duplicated(keep="last")) + self.assert_eq(psser.duplicated(keep=False).sort_index(), pser.duplicated(keep=False)) + + pser = pd.Series([1, 2, 1, 2, 3], name="numbers") + psser = ps.from_pandas(pser) + self.assert_eq((psser + 1).duplicated().sort_index(), (pser + 1).duplicated()) + + def test_drop_duplicates(self): + pdf = pd.DataFrame({"animal": ["lama", "cow", "lama", "beetle", "lama", "hippo"]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.animal + psser = psdf.animal + + self.assert_eq(psser.drop_duplicates().sort_index(), pser.drop_duplicates().sort_index()) + self.assert_eq( + psser.drop_duplicates(keep="last").sort_index(), + pser.drop_duplicates(keep="last").sort_index(), + ) + + # inplace + psser.drop_duplicates(keep=False, inplace=True) + pser.drop_duplicates(keep=False, inplace=True) + self.assert_eq(psser.sort_index(), pser.sort_index()) + self.assert_eq(psdf, pdf) + + def test_clip(self): + pdf = pd.DataFrame({"x": [0, 2, 4]}, index=np.random.rand(3)) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + # Assert list-like values are not accepted for 'lower' and 'upper' + msg = "List-like value are not supported for 'lower' and 'upper' at the moment" + with self.assertRaises(TypeError, msg=msg): + psser.clip(lower=[1]) + with self.assertRaises(TypeError, msg=msg): + psser.clip(upper=[1]) + + # Assert no lower or upper + self.assert_eq(psser.clip(), pser.clip()) + # Assert lower only + self.assert_eq(psser.clip(1), pser.clip(1)) + # Assert upper only + self.assert_eq(psser.clip(upper=3), pser.clip(upper=3)) + # Assert lower and upper + self.assert_eq(psser.clip(1, 3), pser.clip(1, 3)) + self.assert_eq((psser + 1).clip(1, 3), (pser + 1).clip(1, 3)) + + # Assert inplace is True + pser.clip(1, 3, inplace=True) + psser.clip(1, 3, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + # Assert behavior on string values + str_psser = ps.Series(["a", "b", "c"]) + self.assert_eq(str_psser.clip(1, 3), str_psser) + + def test_compare(self): + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + pser = pd.Series([1, 2]) + psser = ps.from_pandas(pser) + + res_psdf = psser.compare(psser) + self.assertTrue(res_psdf.empty) + self.assert_eq(res_psdf.columns, pd.Index(["self", "other"])) + + self.assert_eq( + pser.compare(pser + 1).sort_index(), psser.compare(psser + 1).sort_index() + ) + + pser = pd.Series([1, 2], index=["x", "y"]) + psser = ps.from_pandas(pser) + self.assert_eq( + pser.compare(pser + 1).sort_index(), psser.compare(psser + 1).sort_index() + ) + else: + psser = ps.Series([1, 2]) + res_psdf = psser.compare(psser) + self.assertTrue(res_psdf.empty) + self.assert_eq(res_psdf.columns, pd.Index(["self", "other"])) + expected = ps.DataFrame([[1, 2], [2, 3]], columns=["self", "other"]) + self.assert_eq(expected, psser.compare(psser + 1).sort_index()) + + psser = ps.Series([1, 2], index=["x", "y"]) + expected = ps.DataFrame([[1, 2], [2, 3]], index=["x", "y"], columns=["self", "other"]) + self.assert_eq(expected, psser.compare(psser + 1).sort_index()) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43465): Enable SeriesTests.test_append for pandas 2.0.0.", + ) + def test_append(self): + pser1 = pd.Series([1, 2, 3], name="0") + pser2 = pd.Series([4, 5, 6], name="0") + pser3 = pd.Series([4, 5, 6], index=[3, 4, 5], name="0") + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + psser3 = ps.from_pandas(pser3) + + self.assert_eq(psser1.append(psser2), pser1.append(pser2)) + self.assert_eq(psser1.append(psser3), pser1.append(pser3)) + self.assert_eq( + psser1.append(psser2, ignore_index=True), pser1.append(pser2, ignore_index=True) + ) + + psser1.append(psser3, verify_integrity=True) + msg = "Indices have overlapping values" + with self.assertRaises(ValueError, msg=msg): + psser1.append(psser2, verify_integrity=True) + + def test_shift(self): + pser = pd.Series([10, 20, 15, 30, 45], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.shift(2), pser.shift(2)) + self.assert_eq(psser.shift().shift(-1), pser.shift().shift(-1)) + self.assert_eq(psser.shift().sum(), pser.shift().sum()) + + self.assert_eq(psser.shift(periods=2, fill_value=0), pser.shift(periods=2, fill_value=0)) + + with self.assertRaisesRegex(TypeError, "periods should be an int; however"): + psser.shift(periods=1.5) + + self.assert_eq(psser.shift(periods=0), pser.shift(periods=0)) + + def test_diff(self): + pser = pd.Series([10, 20, 15, 30, 45], name="x") + psser = ps.Series(pser) + + self.assert_eq(psser.diff(2), pser.diff(2)) + self.assert_eq(psser.diff().diff(-1), pser.diff().diff(-1)) + self.assert_eq(psser.diff().sum(), pser.diff().sum()) + + def test_aggregate(self): + pser = pd.Series([10, 20, 15, 30, 45], name="x") + psser = ps.Series(pser) + msg = "func must be a string or list of strings" + with self.assertRaisesRegex(TypeError, msg): + psser.aggregate({"x": ["min", "max"]}) + msg = ( + "If the given function is a list, it " "should only contains function names as strings." + ) + with self.assertRaisesRegex(ValueError, msg): + psser.aggregate(["min", max]) + + def test_drop(self): + pdf = pd.DataFrame({"x": [10, 20, 15, 30, 45]}) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + self.assert_eq(psser.drop(1), pser.drop(1)) + self.assert_eq(psser.drop([1, 4]), pser.drop([1, 4])) + self.assert_eq(psser.drop(columns=1), pser.drop(columns=1)) + self.assert_eq(psser.drop(columns=[1, 4]), pser.drop(columns=[1, 4])) + + msg = "Need to specify at least one of 'labels', 'index' or 'columns'" + with self.assertRaisesRegex(ValueError, msg): + psser.drop() + self.assertRaises(KeyError, lambda: psser.drop((0, 1))) + + psser.drop([2, 3], inplace=True) + pser.drop([2, 3], inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + n_pser, n_psser = pser + 1, psser + 1 + n_psser.drop([1, 4], inplace=True) + n_pser.drop([1, 4], inplace=True) + self.assert_eq(n_psser, n_pser) + self.assert_eq(psser, pser) + + # For MultiIndex + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + + pdf = pd.DataFrame({"x": [45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3]}, index=midx) + psdf = ps.from_pandas(pdf) + psser, pser = psdf.x, pdf.x + + self.assert_eq(psser.drop("lama"), pser.drop("lama")) + self.assert_eq(psser.drop(labels="weight", level=1), pser.drop(labels="weight", level=1)) + self.assert_eq(psser.drop(("lama", "weight")), pser.drop(("lama", "weight"))) + self.assert_eq( + psser.drop([("lama", "speed"), ("falcon", "weight")]), + pser.drop([("lama", "speed"), ("falcon", "weight")]), + ) + self.assert_eq(psser.drop({"lama": "speed"}), pser.drop({"lama": "speed"})) + + msg = "'level' should be less than the number of indexes" + with self.assertRaisesRegex(ValueError, msg): + psser.drop(labels="weight", level=2) + + msg = ( + "If the given index is a list, it " + "should only contains names as all tuples or all non tuples " + "that contain index names" + ) + with self.assertRaisesRegex(ValueError, msg): + psser.drop(["lama", ["cow", "falcon"]]) + + msg = "Cannot specify both 'labels' and 'index'/'columns'" + with self.assertRaisesRegex(ValueError, msg): + psser.drop("lama", index="cow") + + with self.assertRaisesRegex(ValueError, msg): + psser.drop("lama", columns="cow") + + msg = r"'Key length \(2\) exceeds index depth \(3\)'" + with self.assertRaisesRegex(KeyError, msg): + psser.drop(("lama", "speed", "x")) + + psser.drop({"lama": "speed"}, inplace=True) + pser.drop({"lama": "speed"}, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + def test_pop(self): + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pdf = pd.DataFrame({"x": [45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3]}, index=midx) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.pop(("lama", "speed")), pser.pop(("lama", "speed"))) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + msg = r"'Key length \(3\) exceeds index depth \(2\)'" + with self.assertRaisesRegex(KeyError, msg): + psser.pop(("lama", "speed", "x")) + + msg = "'key' should be string or tuple that contains strings" + with self.assertRaisesRegex(TypeError, msg): + psser.pop(["lama", "speed"]) + + pser = pd.Series(["a", "b", "c", "a"], dtype="category") + psser = ps.from_pandas(pser) + + if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"): + self.assert_eq(psser.pop(0), pser.pop(0)) + self.assert_eq(psser, pser) + + self.assert_eq(psser.pop(3), pser.pop(3)) + self.assert_eq(psser, pser) + else: + # Before pandas 1.3.0, `pop` modifies the dtype of categorical series wrongly. + self.assert_eq(psser.pop(0), "a") + self.assert_eq( + psser, + pd.Series( + pd.Categorical(["b", "c", "a"], categories=["a", "b", "c"]), index=[1, 2, 3] + ), + ) + + self.assert_eq(psser.pop(3), "a") + self.assert_eq( + psser, + pd.Series(pd.Categorical(["b", "c"], categories=["a", "b", "c"]), index=[1, 2]), + ) + + def test_duplicates(self): + psers = { + "test on texts": pd.Series( + ["lama", "cow", "lama", "beetle", "lama", "hippo"], name="animal" + ), + "test on numbers": pd.Series([1, 1, 2, 4, 3]), + } + keeps = ["first", "last", False] + + for (msg, pser), keep in product(psers.items(), keeps): + with self.subTest(msg, keep=keep): + psser = ps.Series(pser) + + self.assert_eq( + pser.drop_duplicates(keep=keep).sort_values(), + psser.drop_duplicates(keep=keep).sort_values(), + ) + + def test_truncate(self): + pser1 = pd.Series([10, 20, 30, 40, 50, 60, 70], index=[1, 2, 3, 4, 5, 6, 7]) + psser1 = ps.Series(pser1) + pser2 = pd.Series([10, 20, 30, 40, 50, 60, 70], index=[7, 6, 5, 4, 3, 2, 1]) + psser2 = ps.Series(pser2) + + self.assert_eq(psser1.truncate(), pser1.truncate()) + self.assert_eq(psser1.truncate(before=2), pser1.truncate(before=2)) + self.assert_eq(psser1.truncate(after=5), pser1.truncate(after=5)) + self.assert_eq(psser1.truncate(copy=False), pser1.truncate(copy=False)) + self.assert_eq(psser1.truncate(2, 5, copy=False), pser1.truncate(2, 5, copy=False)) + # The bug for these tests has been fixed in pandas 1.1.0. + if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"): + self.assert_eq(psser2.truncate(4, 6), pser2.truncate(4, 6)) + self.assert_eq(psser2.truncate(4, 6, copy=False), pser2.truncate(4, 6, copy=False)) + else: + expected_psser = ps.Series([20, 30, 40], index=[6, 5, 4]) + self.assert_eq(psser2.truncate(4, 6), expected_psser) + self.assert_eq(psser2.truncate(4, 6, copy=False), expected_psser) + + psser = ps.Series([10, 20, 30, 40, 50, 60, 70], index=[1, 2, 3, 4, 3, 2, 1]) + msg = "truncate requires a sorted index" + with self.assertRaisesRegex(ValueError, msg): + psser.truncate() + + psser = ps.Series([10, 20, 30, 40, 50, 60, 70], index=[1, 2, 3, 4, 5, 6, 7]) + msg = "Truncate: 2 must be after 5" + with self.assertRaisesRegex(ValueError, msg): + psser.truncate(5, 2) + + def test_unstack(self): + pser = pd.Series( + [10, -2, 4, 7], + index=pd.MultiIndex.from_tuples( + [("one", "a", "z"), ("one", "b", "x"), ("two", "a", "c"), ("two", "b", "v")], + names=["A", "B", "C"], + ), + ) + psser = ps.from_pandas(pser) + + levels = [-3, -2, -1, 0, 1, 2] + for level in levels: + pandas_result = pser.unstack(level=level) + pandas_on_spark_result = psser.unstack(level=level).sort_index() + self.assert_eq(pandas_result, pandas_on_spark_result) + self.assert_eq(pandas_result.index.names, pandas_on_spark_result.index.names) + self.assert_eq(pandas_result.columns.names, pandas_on_spark_result.columns.names) + + # non-numeric datatypes + pser = pd.Series( + list("abcd"), index=pd.MultiIndex.from_product([["one", "two"], ["a", "b"]]) + ) + psser = ps.from_pandas(pser) + + levels = [-2, -1, 0, 1] + for level in levels: + pandas_result = pser.unstack(level=level) + pandas_on_spark_result = psser.unstack(level=level).sort_index() + self.assert_eq(pandas_result, pandas_on_spark_result) + self.assert_eq(pandas_result.index.names, pandas_on_spark_result.index.names) + self.assert_eq(pandas_result.columns.names, pandas_on_spark_result.columns.names) + + # Exceeding the range of level + self.assertRaises(IndexError, lambda: psser.unstack(level=3)) + self.assertRaises(IndexError, lambda: psser.unstack(level=-4)) + # Only support for MultiIndex + psser = ps.Series([10, -2, 4, 7]) + self.assertRaises(ValueError, lambda: psser.unstack()) + + def test_abs(self): + pser = pd.Series([-2, -1, 0, 1]) + psser = ps.from_pandas(pser) + + self.assert_eq(abs(psser), abs(pser)) + self.assert_eq(np.abs(psser), np.abs(pser)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43550): Enable SeriesTests.test_factorize for pandas 2.0.0.", + ) + def test_factorize(self): + pser = pd.Series(["a", "b", "a", "b"]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series([5, 1, 5, 1]) + psser = ps.from_pandas(pser) + pcodes, puniques = (pser + 1).factorize(sort=True) + kcodes, kuniques = (psser + 1).factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series(["a", "b", "a", "b"], name="ser", index=["w", "x", "y", "z"]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series( + ["a", "b", "a", "b"], index=pd.MultiIndex.from_arrays([[4, 3, 2, 1], [1, 2, 3, 4]]) + ) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + # + # Deals with None and np.nan + # + pser = pd.Series(["a", "b", "a", np.nan]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series([1, None, 3, 2, 1]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series(["a", None, "a"]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize(sort=True) + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pser = pd.Series([None, np.nan]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize() + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes, kcodes.to_list()) + # pandas: Float64Index([], dtype='float64') + self.assert_eq(pd.Index([]), kuniques) + + pser = pd.Series([np.nan, np.nan]) + psser = ps.from_pandas(pser) + pcodes, puniques = pser.factorize() + kcodes, kuniques = psser.factorize() + self.assert_eq(pcodes, kcodes.to_list()) + # pandas: Float64Index([], dtype='float64') + self.assert_eq(pd.Index([]), kuniques) + + # + # Deals with na_sentinel + # + # pandas >= 1.1.2 support na_sentinel=None + # + pd_below_1_1_2 = LooseVersion(pd.__version__) < LooseVersion("1.1.2") + + pser = pd.Series(["a", "b", "a", np.nan, None]) + psser = ps.from_pandas(pser) + + pcodes, puniques = pser.factorize(sort=True, na_sentinel=-2) + kcodes, kuniques = psser.factorize(na_sentinel=-2) + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + pcodes, puniques = pser.factorize(sort=True, na_sentinel=2) + kcodes, kuniques = psser.factorize(na_sentinel=2) + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + if not pd_below_1_1_2: + pcodes, puniques = pser.factorize(sort=True, na_sentinel=None) + kcodes, kuniques = psser.factorize(na_sentinel=None) + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + # puniques is Index(['a', 'b', nan], dtype='object') + self.assert_eq(ps.Index(["a", "b", None]), kuniques) + + psser = ps.Series([1, 2, np.nan, 4, 5]) # Arrow takes np.nan as null + psser.loc[3] = np.nan # Spark takes np.nan as NaN + kcodes, kuniques = psser.factorize(na_sentinel=None) + pcodes, puniques = psser._to_pandas().factorize(sort=True, na_sentinel=None) + self.assert_eq(pcodes.tolist(), kcodes.to_list()) + self.assert_eq(puniques, kuniques) + + def test_explode(self): + pser = pd.Series([[1, 2, 3], [], None, [3, 4]]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.explode(), psser.explode(), almost=True) + + # MultiIndex + pser.index = pd.MultiIndex.from_tuples([("a", "w"), ("b", "x"), ("c", "y"), ("d", "z")]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.explode(), psser.explode(), almost=True) + + # non-array type Series + pser = pd.Series([1, 2, 3, 4]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.explode(), psser.explode()) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43467): Enable SeriesTests.test_between for pandas 2.0.0.", + ) + def test_between(self): + pser = pd.Series([np.nan, 1, 2, 3, 4]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.between(1, 4), pser.between(1, 4)) + self.assert_eq(psser.between(1, 4, inclusive="both"), pser.between(1, 4, inclusive="both")) + self.assert_eq( + psser.between(1, 4, inclusive="neither"), pser.between(1, 4, inclusive="neither") + ) + self.assert_eq(psser.between(1, 4, inclusive="left"), pser.between(1, 4, inclusive="left")) + self.assert_eq( + psser.between(1, 4, inclusive="right"), pser.between(1, 4, inclusive="right") + ) + expected_err_msg = ( + "Inclusive has to be either string of 'both'," "'left', 'right', or 'neither'" + ) + with self.assertRaisesRegex(ValueError, expected_err_msg): + psser.between(1, 4, inclusive="middle") + + # Test for backward compatibility + self.assert_eq(psser.between(1, 4, inclusive=True), pser.between(1, 4, inclusive=True)) + self.assert_eq(psser.between(1, 4, inclusive=False), pser.between(1, 4, inclusive=False)) + with self.assertWarns(FutureWarning): + psser.between(1, 4, inclusive=True) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43479): Enable SeriesTests.test_between_time for pandas 2.0.0.", + ) + def test_between_time(self): + idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") + pser = pd.Series([1, 2, 3, 4], index=idx) + psser = ps.from_pandas(pser) + self.assert_eq( + pser.between_time("0:15", "0:45").sort_index(), + psser.between_time("0:15", "0:45").sort_index(), + ) + + pser.index.name = "ts" + psser = ps.from_pandas(pser) + self.assert_eq( + pser.between_time("0:15", "0:45").sort_index(), + psser.between_time("0:15", "0:45").sort_index(), + ) + + pser.index.name = "index" + psser = ps.from_pandas(pser) + self.assert_eq( + pser.between_time("0:15", "0:45").sort_index(), + psser.between_time("0:15", "0:45").sort_index(), + ) + + def test_at_time(self): + idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") + pser = pd.Series([1, 2, 3, 4], index=idx) + psser = ps.from_pandas(pser) + self.assert_eq( + pser.at_time("0:20").sort_index(), + psser.at_time("0:20").sort_index(), + ) + + pser.index.name = "ts" + psser = ps.from_pandas(pser) + self.assert_eq( + pser.at_time("0:20").sort_index(), + psser.at_time("0:20").sort_index(), + ) + + pser.index.name = "index" + psser = ps.from_pandas(pser) + self.assert_eq( + pser.at_time("0:20").sort_index(), + psser.at_time("0:20").sort_index(), + ) + + +class SeriesComputeTests(SeriesComputeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_compute import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_conversion.py new file mode 100644 index 0000000..9d98d45 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_conversion.py @@ -0,0 +1,88 @@ +# +# 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 unittest + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.testing.pandasutils import have_tabulate, tabulate_requirement_message + + +class SeriesConversionMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_to_numpy(self): + pser = pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + psser = ps.from_pandas(pser) + self.assert_eq(psser.to_numpy(), pser.values) + + def test_to_datetime(self): + pser = pd.Series(["3/11/2000", "3/12/2000", "3/13/2000"] * 100) + psser = ps.from_pandas(pser) + + self.assert_eq( + pd.to_datetime(pser, infer_datetime_format=True), + ps.to_datetime(psser, infer_datetime_format=True), + ) + + def test_to_list(self): + self.assert_eq(self.psser.tolist(), self.pser.tolist()) + + def test_to_frame(self): + pser = pd.Series(["a", "b", "c"]) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.to_frame(name="a"), psser.to_frame(name="a")) + + # for MultiIndex + midx = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + pser = pd.Series(["a", "b", "c"], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.to_frame(name="a"), psser.to_frame(name="a")) + + @unittest.skipIf(not have_tabulate, tabulate_requirement_message) + def test_to_markdown(self): + pser = pd.Series(["elk", "pig", "dog", "quetzal"], name="animal") + psser = ps.from_pandas(pser) + + self.assert_eq(pser.to_markdown(), psser.to_markdown()) + + +class SeriesConversionTests(SeriesConversionMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_conversion import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_cumulative.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_cumulative.py new file mode 100644 index 0000000..b51031b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_cumulative.py @@ -0,0 +1,139 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesCumulativeMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_cummin(self): + pser = pd.Series([1.0, None, 0.0, 4.0, 9.0]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cummin(), psser.cummin()) + self.assert_eq(pser.cummin(skipna=False), psser.cummin(skipna=False)) + self.assert_eq(pser.cummin().sum(), psser.cummin().sum()) + + # with reversed index + pser.index = [4, 3, 2, 1, 0] + psser = ps.from_pandas(pser) + self.assert_eq(pser.cummin(), psser.cummin()) + self.assert_eq(pser.cummin(skipna=False), psser.cummin(skipna=False)) + + def test_cummax(self): + pser = pd.Series([1.0, None, 0.0, 4.0, 9.0]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cummax(), psser.cummax()) + self.assert_eq(pser.cummax(skipna=False), psser.cummax(skipna=False)) + self.assert_eq(pser.cummax().sum(), psser.cummax().sum()) + + # with reversed index + pser.index = [4, 3, 2, 1, 0] + psser = ps.from_pandas(pser) + self.assert_eq(pser.cummax(), psser.cummax()) + self.assert_eq(pser.cummax(skipna=False), psser.cummax(skipna=False)) + + def test_cumsum(self): + pser = pd.Series([1.0, None, 0.0, 4.0, 9.0]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumsum(), psser.cumsum()) + self.assert_eq(pser.cumsum(skipna=False), psser.cumsum(skipna=False)) + self.assert_eq(pser.cumsum().sum(), psser.cumsum().sum()) + + # with reversed index + pser.index = [4, 3, 2, 1, 0] + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumsum(), psser.cumsum()) + self.assert_eq(pser.cumsum(skipna=False), psser.cumsum(skipna=False)) + + # bool + pser = pd.Series([True, True, False, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumsum().astype(int), psser.cumsum()) + self.assert_eq(pser.cumsum(skipna=False).astype(int), psser.cumsum(skipna=False)) + + with self.assertRaisesRegex(TypeError, r"Could not convert object \(string\) to numeric"): + ps.Series(["a", "b", "c", "d"]).cumsum() + + def test_cumprod(self): + pser = pd.Series([1.0, None, 1.0, 4.0, 9.0]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False), psser.cumprod(skipna=False)) + self.assert_eq(pser.cumprod().sum(), psser.cumprod().sum()) + + # with integer type + pser = pd.Series([1, 10, 1, 4, 9]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False), psser.cumprod(skipna=False)) + self.assert_eq(pser.cumprod().sum(), psser.cumprod().sum()) + + # with reversed index + pser.index = [4, 3, 2, 1, 0] + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False), psser.cumprod(skipna=False)) + + # including zero + pser = pd.Series([1, 2, 0, 3]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False), psser.cumprod(skipna=False)) + + # including negative values + pser = pd.Series([1, -1, -2]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False), psser.cumprod(skipna=False)) + + # bool + pser = pd.Series([True, True, False, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.cumprod(), psser.cumprod()) + self.assert_eq(pser.cumprod(skipna=False).astype(int), psser.cumprod(skipna=False)) + + with self.assertRaisesRegex(TypeError, r"Could not convert object \(string\) to numeric"): + ps.Series(["a", "b", "c", "d"]).cumprod() + + +class SeriesCumulativeTests(SeriesCumulativeMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_cumulative import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_index.py new file mode 100644 index 0000000..31b85b6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_index.py @@ -0,0 +1,475 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesIndexMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_align(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + psdf = ps.from_pandas(pdf) + + for join in ["outer", "inner", "left", "right"]: + for axis in [None, 0]: + psser_l, psser_r = psdf.a.align(psdf.b, join=join, axis=axis) + pser_l, pser_r = pdf.a.align(pdf.b, join=join, axis=axis) + self.assert_eq(psser_l, pser_l) + self.assert_eq(psser_r, pser_r) + + psser_l, psdf_r = psdf.b.align(psdf[["b", "a"]], join=join, axis=axis) + pser_l, pdf_r = pdf.b.align(pdf[["b", "a"]], join=join, axis=axis) + self.assert_eq(psser_l, pser_l) + self.assert_eq(psdf_r, pdf_r) + + self.assertRaises(ValueError, lambda: psdf.a.align(psdf.b, axis=1)) + + def test_reindex(self): + index = ["A", "B", "C", "D", "E"] + pser = pd.Series([1.0, 2.0, 3.0, 4.0, None], index=index, name="x") + psser = ps.from_pandas(pser) + + self.assert_eq(pser, psser) + + self.assert_eq( + pser.reindex(["A", "B"]).sort_index(), + psser.reindex(["A", "B"]).sort_index(), + ) + + self.assert_eq( + pser.reindex(["A", "B", "2", "3"]).sort_index(), + psser.reindex(["A", "B", "2", "3"]).sort_index(), + ) + + self.assert_eq( + pser.reindex(["A", "E", "2"], fill_value=0).sort_index(), + psser.reindex(["A", "E", "2"], fill_value=0).sort_index(), + ) + + self.assertRaises(TypeError, lambda: psser.reindex(index=123)) + + def test_reindex_like(self): + data = [1.0, 2.0, None] + index = pd.Index(["A", "B", "C"], name="index1") + pser = pd.Series(data=data, index=index, name="name1") + psser = ps.from_pandas(pser) + + # Reindexing single Index on single Index + data2 = [3.0, None, 4.0] + index2 = pd.Index(["A", "C", "D"], name="index2") + pser2 = pd.Series(data=data2, index=index2, name="name2") + psser2 = ps.from_pandas(pser2) + + self.assert_eq( + pser.reindex_like(pser2).sort_index(), + psser.reindex_like(psser2).sort_index(), + ) + + self.assert_eq( + (pser + 1).reindex_like(pser2).sort_index(), + (psser + 1).reindex_like(psser2).sort_index(), + ) + + # Reindexing MultiIndex on single Index + index2 = pd.MultiIndex.from_tuples( + [("A", "G"), ("C", "D"), ("I", "J")], names=["index3", "index4"] + ) + pser2 = pd.Series(data=data2, index=index2, name="name2") + psser2 = ps.from_pandas(pser2) + + self.assert_eq( + pser.reindex_like(pser2).sort_index(), + psser.reindex_like(psser2).sort_index(), + ) + + self.assertRaises(TypeError, lambda: psser.reindex_like(index2)) + self.assertRaises(AssertionError, lambda: psser2.reindex_like(psser)) + + # Reindexing MultiIndex on MultiIndex + index = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] + ) + pser = pd.Series(data=data, index=index, name="name1") + psser = ps.from_pandas(pser) + + self.assert_eq( + pser.reindex_like(pser2).sort_index(), + psser.reindex_like(psser2).sort_index(), + ) + + # Reindexing with DataFrame + index2 = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["name3", "name4"] + ) + pdf = pd.DataFrame(data=data, index=index2) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + pser.reindex_like(pdf).sort_index(), + psser.reindex_like(psdf).sort_index(), + ) + + def test_reset_index(self): + pdf = pd.DataFrame({"foo": [1, 2, 3, 4]}, index=pd.Index(["a", "b", "c", "d"], name="idx")) + psdf = ps.from_pandas(pdf) + + pser = pdf.foo + psser = psdf.foo + + self.assert_eq(psser.reset_index(), pser.reset_index()) + self.assert_eq(psser.reset_index(name="values"), pser.reset_index(name="values")) + self.assert_eq(psser.reset_index(drop=True), pser.reset_index(drop=True)) + + # inplace + psser.reset_index(drop=True, inplace=True) + pser.reset_index(drop=True, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + def test_reset_index_with_default_index_types(self): + pser = pd.Series([1, 2, 3], name="0", index=np.random.rand(3)) + psser = ps.from_pandas(pser) + + with ps.option_context("compute.default_index_type", "sequence"): + self.assert_eq(psser.reset_index(), pser.reset_index()) + + with ps.option_context("compute.default_index_type", "distributed-sequence"): + # the order might be changed. + self.assert_eq(psser.reset_index().sort_index(), pser.reset_index()) + + with ps.option_context("compute.default_index_type", "distributed"): + # the index is different. + self.assert_eq( + psser.reset_index()._to_pandas().reset_index(drop=True), pser.reset_index() + ) + + def test_index_to_series_reset_index(self): + def check(psser, pser): + self.assert_eq(psser.reset_index(), pser.reset_index()) + self.assert_eq(psser.reset_index(drop=True), pser.reset_index(drop=True)) + + pser.reset_index(drop=True, inplace=True) + psser.reset_index(drop=True, inplace=True) + self.assert_eq(psser, pser) + + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + psdf = ps.from_pandas(pdf) + check(psdf.index.to_series(), pdf.index.to_series()) + check(psdf.index.to_series(name="a"), pdf.index.to_series(name="a")) + check(psdf.index.to_series(name=("x", "a")), pdf.index.to_series(name=("x", "a"))) + + def test_idxmax(self): + pser = pd.Series(data=[1, 4, 5], index=["A", "B", "C"]) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmax(), pser.idxmax()) + self.assertEqual(psser.idxmax(skipna=False), pser.idxmax(skipna=False)) + + index = pd.MultiIndex.from_arrays( + [["a", "a", "b", "b"], ["c", "d", "e", "f"]], names=("first", "second") + ) + pser = pd.Series(data=[1, 2, 4, 5], index=index) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmax(), pser.idxmax()) + self.assertEqual(psser.idxmax(skipna=False), pser.idxmax(skipna=False)) + + psser = ps.Series([]) + with self.assertRaisesRegex(ValueError, "an empty sequence"): + psser.idxmax() + + pser = pd.Series([1, 100, None, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmax(), pser.idxmax()) + self.assertEqual(repr(psser.idxmax(skipna=False)), repr(pser.idxmax(skipna=False))) + + def test_idxmin(self): + pser = pd.Series(data=[1, 4, 5], index=["A", "B", "C"]) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmin(), pser.idxmin()) + self.assertEqual(psser.idxmin(skipna=False), pser.idxmin(skipna=False)) + + index = pd.MultiIndex.from_arrays( + [["a", "a", "b", "b"], ["c", "d", "e", "f"]], names=("first", "second") + ) + pser = pd.Series(data=[1, 2, 4, 5], index=index) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmin(), pser.idxmin()) + self.assertEqual(psser.idxmin(skipna=False), pser.idxmin(skipna=False)) + + psser = ps.Series([]) + with self.assertRaisesRegex(ValueError, "an empty sequence"): + psser.idxmin() + + pser = pd.Series([1, 100, None, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) + psser = ps.Series(pser) + + self.assertEqual(psser.idxmin(), pser.idxmin()) + self.assertEqual(repr(psser.idxmin(skipna=False)), repr(pser.idxmin(skipna=False))) + + def test_index(self): + # to check setting name of Index properly. + idx = pd.Index([1, 2, 3, 4, 5, 6, 7, 8, 9]) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=idx) + psser = ps.from_pandas(pser) + + psser.name = "koalas" + pser.name = "koalas" + self.assert_eq(psser.index.name, pser.index.name) + + # for check setting names of MultiIndex properly. + psser.names = ["hello", "koalas"] + pser.names = ["hello", "koalas"] + self.assert_eq(psser.index.names, pser.index.names) + + def test_axes(self): + pser = pd.Series([90, 91, 85], index=[2, 4, 1]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.axes, pser.axes) + + # for MultiIndex + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + psser = ps.from_pandas(pser) + self.assert_eq(psser.axes, pser.axes) + + def test_swaplevel(self): + # MultiIndex with two levels + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) + pser = pd.Series(["a", "b", "c", "d"], index=pidx) + psser = ps.from_pandas(pser) + self.assert_eq(pser.swaplevel(), psser.swaplevel()) + self.assert_eq(pser.swaplevel(0, 1), psser.swaplevel(0, 1)) + self.assert_eq(pser.swaplevel(1, 1), psser.swaplevel(1, 1)) + self.assert_eq(pser.swaplevel("number", "color"), psser.swaplevel("number", "color")) + + # MultiIndex with more than two levels + arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"], ["l", "m", "s", "xs"]] + pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color", "size")) + pser = pd.Series(["a", "b", "c", "d"], index=pidx) + psser = ps.from_pandas(pser) + self.assert_eq(pser.swaplevel(), psser.swaplevel()) + self.assert_eq(pser.swaplevel(0, 1), psser.swaplevel(0, 1)) + self.assert_eq(pser.swaplevel(0, 2), psser.swaplevel(0, 2)) + self.assert_eq(pser.swaplevel(1, 2), psser.swaplevel(1, 2)) + self.assert_eq(pser.swaplevel(1, 1), psser.swaplevel(1, 1)) + self.assert_eq(pser.swaplevel(-1, -2), psser.swaplevel(-1, -2)) + self.assert_eq(pser.swaplevel("number", "color"), psser.swaplevel("number", "color")) + self.assert_eq(pser.swaplevel("number", "size"), psser.swaplevel("number", "size")) + self.assert_eq(pser.swaplevel("color", "size"), psser.swaplevel("color", "size")) + + # Error conditions + self.assertRaises(AssertionError, lambda: ps.Series([1, 2]).swaplevel()) + self.assertRaises(IndexError, lambda: psser.swaplevel(0, 9)) + self.assertRaises(KeyError, lambda: psser.swaplevel("not_number", "color")) + self.assertRaises(AssertionError, lambda: psser.swaplevel(copy=False)) + + def test_swapaxes(self): + pser = pd.Series([1, 2, 3], index=["x", "y", "z"], name="ser") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.swapaxes(0, 0), pser.swapaxes(0, 0)) + self.assert_eq(psser.swapaxes("index", "index"), pser.swapaxes("index", "index")) + self.assert_eq((psser + 1).swapaxes(0, 0), (pser + 1).swapaxes(0, 0)) + + self.assertRaises(AssertionError, lambda: psser.swapaxes(0, 1, copy=False)) + self.assertRaises(ValueError, lambda: psser.swapaxes(0, 1)) + self.assertRaises(ValueError, lambda: psser.swapaxes("index", "columns")) + + def test_droplevel(self): + pser = pd.Series( + [1, 2, 3], + index=pd.MultiIndex.from_tuples( + [("x", "a", "q"), ("x", "b", "w"), ("y", "c", "e")], + names=["level_1", "level_2", "level_3"], + ), + ) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.droplevel(0), psser.droplevel(0)) + self.assert_eq(pser.droplevel("level_1"), psser.droplevel("level_1")) + self.assert_eq(pser.droplevel(-1), psser.droplevel(-1)) + self.assert_eq(pser.droplevel([0]), psser.droplevel([0])) + self.assert_eq(pser.droplevel(["level_1"]), psser.droplevel(["level_1"])) + self.assert_eq(pser.droplevel((0,)), psser.droplevel((0,))) + self.assert_eq(pser.droplevel(("level_1",)), psser.droplevel(("level_1",))) + self.assert_eq(pser.droplevel([0, 2]), psser.droplevel([0, 2])) + self.assert_eq( + pser.droplevel(["level_1", "level_3"]), psser.droplevel(["level_1", "level_3"]) + ) + self.assert_eq(pser.droplevel((1, 2)), psser.droplevel((1, 2))) + self.assert_eq( + pser.droplevel(("level_2", "level_3")), psser.droplevel(("level_2", "level_3")) + ) + + with self.assertRaisesRegex(KeyError, "Level {0, 1, 2} not found"): + psser.droplevel({0, 1, 2}) + with self.assertRaisesRegex(KeyError, "Level level_100 not found"): + psser.droplevel(["level_1", "level_100"]) + with self.assertRaisesRegex(IndexError, "Too many levels: Index has only 3 levels, not 11"): + psser.droplevel(10) + with self.assertRaisesRegex( + IndexError, + "Too many levels: Index has only 3 levels, -10 is not a valid level number", + ): + psser.droplevel(-10) + with self.assertRaisesRegex( + ValueError, + "Cannot remove 3 levels from an index with 3 levels: " + "at least one level must be left.", + ): + psser.droplevel([0, 1, 2]) + with self.assertRaisesRegex( + ValueError, + "Cannot remove 5 levels from an index with 3 levels: " + "at least one level must be left.", + ): + psser.droplevel([1, 1, 1, 1, 1]) + + # Tupled names + pser.index.names = [("a", "1"), ("b", "2"), ("c", "3")] + psser = ps.from_pandas(pser) + + self.assert_eq( + pser.droplevel([("a", "1"), ("c", "3")]), psser.droplevel([("a", "1"), ("c", "3")]) + ) + + def test_last_valid_index(self): + pser = pd.Series([250, 1.5, 320, 1, 0.3, None, None, None, None]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) + + # MultiIndex columns + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser.index = midx + psser = ps.from_pandas(pser) + self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) + + # Empty Series + pser = pd.Series([]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) + + def test_first_valid_index(self): + # Empty Series + pser = pd.Series([]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.first_valid_index(), psser.first_valid_index()) + + def test_rename_axis(self): + index = pd.Index(["A", "B", "C"], name="index") + pdf = pd.DataFrame({"x": [1.0, 2.0, 3.0]}, index=index) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + self.assert_eq( + pser.rename_axis("index2").sort_index(), + psser.rename_axis("index2").sort_index(), + ) + + self.assert_eq( + (pser + 1).rename_axis("index2").sort_index(), + (psser + 1).rename_axis("index2").sort_index(), + ) + + self.assertRaises(ValueError, lambda: psser.rename_axis(["index2", "index3"])) + self.assertRaises(TypeError, lambda: psser.rename_axis(mapper=["index2"], index=["index3"])) + + self.assert_eq( + pser.rename_axis(index={"index": "index2", "missing": "index4"}).sort_index(), + psser.rename_axis(index={"index": "index2", "missing": "index4"}).sort_index(), + ) + + self.assert_eq( + pser.rename_axis(index=str.upper).sort_index(), + psser.rename_axis(index=str.upper).sort_index(), + ) + + pser.rename_axis("index2", inplace=True) + psser.rename_axis("index2", inplace=True) + self.assert_eq(pser.sort_index(), psser.sort_index()) + # Note: in pandas, pdf.x's index is renamed, whereas pdf's index isn't due to a bug. + self.assert_eq(pdf, psdf) + + index = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] + ) + pser = pd.Series([1.0, 2.0, 3.0], index=index, name="name") + psser = ps.from_pandas(pser) + + self.assert_eq( + pser.rename_axis(["index3", "index4"]).sort_index(), + psser.rename_axis(["index3", "index4"]).sort_index(), + ) + + self.assertRaises(ValueError, lambda: psser.rename_axis(["index3", "index4", "index5"])) + + self.assert_eq( + pser.rename_axis( + index={"index1": "index3", "index2": "index4", "missing": "index5"} + ).sort_index(), + psser.rename_axis( + index={"index1": "index3", "index2": "index4", "missing": "index5"} + ).sort_index(), + ) + + self.assert_eq( + pser.rename_axis(index=str.upper).sort_index(), + psser.rename_axis(index=str.upper).sort_index(), + ) + + +class SeriesIndexTests(SeriesIndexMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_index import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_missing_data.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_missing_data.py new file mode 100644 index 0000000..5067335 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_missing_data.py @@ -0,0 +1,265 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesMissingDataMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_fillna(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6], "y": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.fillna(0), pser.fillna(0)) + self.assert_eq(psser.fillna(np.nan).fillna(0), pser.fillna(np.nan).fillna(0)) + + psser.fillna(0, inplace=True) + pser.fillna(0, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + # test considering series does not have NA/NaN values + psser.fillna(0, inplace=True) + pser.fillna(0, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + psser = psdf.x.rename("y") + pser = pdf.x.rename("y") + psser.fillna(0, inplace=True) + pser.fillna(0, inplace=True) + self.assert_eq(psser.head(), pser.head()) + self.assert_eq(psdf, pdf) + + pser = pd.Series([1, 2, 3, 4, 5, 6], name="x") + psser = ps.from_pandas(pser) + + pser.loc[3] = np.nan + psser.loc[3] = np.nan + + self.assert_eq(psser.fillna(0), pser.fillna(0)) + self.assert_eq(psser.fillna(method="ffill"), pser.fillna(method="ffill")) + self.assert_eq(psser.fillna(method="bfill"), pser.fillna(method="bfill")) + self.assert_eq(psser.fillna(method="backfill"), pser.fillna(method="backfill")) + + # inplace fillna on non-nullable column + pdf = pd.DataFrame({"a": [1, 2, None], "b": [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.b + psser = psdf.b + + self.assert_eq(psser.fillna(0), pser.fillna(0)) + self.assert_eq(psser.fillna(np.nan).fillna(0), pser.fillna(np.nan).fillna(0)) + + psser.fillna(0, inplace=True) + pser.fillna(0, inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + with self.assertRaisesRegex( + ValueError, "Must specify a fillna 'value' or 'method' parameter." + ): + psser.fillna() + with self.assertRaisesRegex(TypeError, "Unsupported type list"): + psdf.a.fillna([0]) + with self.assertRaisesRegex( + NotImplementedError, "fillna currently only works for axis=0 or axis='index'" + ): + psdf.a.fillna(0, axis=1) + with self.assertRaisesRegex( + NotImplementedError, "limit parameter for value is not support now" + ): + psdf.a.fillna(0, limit=1) + + def test_dropna(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.dropna(), pser.dropna()) + + pser.dropna(inplace=True) + psser.dropna(inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + def test_add_and_radd_fill_value(self): + pser = pd.Series([1, 2, None, 4]) + psser = ps.from_pandas(pser) + values = [-10, -0.5, 0, None, 0.5, 10] + for value in values: + self.assert_eq(pser.add(10, fill_value=value), psser.add(10, fill_value=value)) + self.assert_eq(pser.radd(10, fill_value=value), psser.radd(10, fill_value=value)) + + self.assertRaises(NotImplementedError, lambda: psser.add([1, 2, None, None], fill_value=10)) + self.assertRaises( + NotImplementedError, lambda: psser.radd([1, 2, None, None], fill_value=10) + ) + + def test_replace(self): + pser = pd.Series([10, 20, 15, 30, np.nan], name="x") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.replace(), pser.replace()) + self.assert_eq(psser.replace({}), pser.replace({})) + + self.assert_eq(psser.replace(np.nan, 45), pser.replace(np.nan, 45)) + self.assert_eq(psser.replace([10, 15], 45), pser.replace([10, 15], 45)) + self.assert_eq(psser.replace((10, 15), 45), pser.replace((10, 15), 45)) + self.assert_eq(psser.replace([10, 15], [45, 50]), pser.replace([10, 15], [45, 50])) + self.assert_eq(psser.replace((10, 15), (45, 50)), pser.replace((10, 15), (45, 50))) + + pser = pd.Series(["bat", "foo", "bait", "abc", "bar", "zoo"]) + psser = ps.from_pandas(pser) + self.assert_eq( + psser.replace(to_replace=r"^ba.$", value="new", regex=True), + pser.replace(to_replace=r"^ba.$", value="new", regex=True), + ) + self.assert_eq( + psser.replace(regex=r"^.oo$", value="new"), pser.replace(regex=r"^.oo$", value="new") + ) + self.assert_eq( + (psser + "o").replace(regex=r"^.ooo$", value="new"), + (pser + "o").replace(regex=r"^.ooo$", value="new"), + ) + + msg = "'to_replace' should be one of str, list, tuple, dict, int, float" + with self.assertRaisesRegex(TypeError, msg): + psser.replace(ps.range(5)) + msg = "Replacement lists must match in length. Expecting 3 got 2" + with self.assertRaisesRegex(ValueError, msg): + psser.replace(["bat", "foo", "bait"], ["a", "b"]) + msg = "'to_replace' must be 'None' if 'regex' is not a bool" + with self.assertRaisesRegex(ValueError, msg): + psser.replace(to_replace="foo", regex=r"^.oo$") + msg = "If 'regex' is True then 'to_replace' must be a string" + with self.assertRaisesRegex(AssertionError, msg): + psser.replace(["bat", "foo", "bait"], regex=True) + unsupported_regex = [r"^.oo$", r"^ba.$"] + msg = "'regex' of %s type is not supported" % type(unsupported_regex).__name__ + with self.assertRaisesRegex(NotImplementedError, msg): + psser.replace(regex=unsupported_regex, value="new") + + def test_bfill(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6], "y": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.bfill(), pser.bfill()) + self.assert_eq(psser.bfill()[0], pser.bfill()[0]) + + psser.bfill(inplace=True) + pser.bfill(inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psser[0], pser[0]) + self.assert_eq(psdf, pdf) + + def test_ffill(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6], "y": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.ffill(), pser.ffill()) + self.assert_eq(psser.ffill()[4], pser.ffill()[4]) + + psser.ffill(inplace=True) + pser.ffill(inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psser[4], pser[4]) + self.assert_eq(psdf, pdf) + + def test_pad(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(pser.pad(), psser.pad()) + + # Test `inplace=True` + pser.pad(inplace=True) + psser.pad(inplace=True) + self.assert_eq(pser, psser) + self.assert_eq(pdf, psdf) + else: + expected = ps.Series([np.nan, 2, 3, 4, 4, 6], name="x") + self.assert_eq(expected, psser.pad()) + + # Test `inplace=True` + psser.pad(inplace=True) + self.assert_eq(expected, psser) + + def test_backfill(self): + pdf = pd.DataFrame({"x": [np.nan, 2, 3, 4, np.nan, 6]}) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + self.assert_eq(pser.backfill(), psser.backfill()) + + # Test `inplace=True` + pser.backfill(inplace=True) + psser.backfill(inplace=True) + self.assert_eq(pser, psser) + self.assert_eq(pdf, psdf) + else: + expected = ps.Series([2.0, 2.0, 3.0, 4.0, 6.0, 6.0], name="x") + self.assert_eq(expected, psser.backfill()) + + # Test `inplace=True` + psser.backfill(inplace=True) + self.assert_eq(expected, psser) + + +class SeriesMissingDataTests(SeriesMissingDataMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_missing_data import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_series.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_series.py new file mode 100644 index 0000000..116acb2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_series.py @@ -0,0 +1,818 @@ +# +# 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 unittest +from collections import defaultdict +from distutils.version import LooseVersion +import inspect + +from datetime import datetime, timedelta + +import numpy as np +import pandas as pd +from pyspark.ml.linalg import SparseVector + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ( + have_tabulate, + PandasOnSparkTestCase, + SPARK_CONF_ARROW_ENABLED, + tabulate_requirement_message, +) +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.series import MissingPandasLikeSeries +from pyspark.pandas.typedef.typehints import ( + extension_dtypes, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) + + +class SeriesTestsMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_creation_index(self): + err_msg = ( + "The given index cannot be a pandas-on-Spark index. Try pandas index or array-like." + ) + with self.assertRaisesRegex(TypeError, err_msg): + ps.Series([1, 2], index=ps.Index([1, 2])) + + with self.assertRaisesRegex(TypeError, err_msg): + ps.Series([1, 2], index=ps.MultiIndex.from_tuples([(1, 3), (2, 4)])) + + def test_series_ops(self): + pser = self.pser + psser = self.psser + + self.assert_eq(psser + 1 + 10 * psser, pser + 1 + 10 * pser) + self.assert_eq(psser + 1 + 10 * psser.index, pser + 1 + 10 * pser.index) + self.assert_eq(psser.index + 1 + 10 * psser, pser.index + 1 + 10 * pser) + + def test_series_tuple_name(self): + pser = self.pser + pser.name = ("x", "a") + + psser = ps.from_pandas(pser) + + self.assert_eq(psser, pser) + self.assert_eq(psser.name, pser.name) + + pser.name = ("y", "z") + psser.name = ("y", "z") + + self.assert_eq(psser, pser) + self.assert_eq(psser.name, pser.name) + + def test_repr_cache_invalidation(self): + # If there is any cache, inplace operations should invalidate it. + s = ps.range(10)["id"] + s.__repr__() + s.rename("a", inplace=True) + self.assertEqual(s.__repr__(), s.rename("a").__repr__()) + + def _check_extension(self, psser, pser): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(psser, pser, check_exact=False) + self.assertTrue(isinstance(psser.dtype, extension_dtypes)) + else: + self.assert_eq(psser, pser) + + def test_empty_series(self): + pser_a = pd.Series([], dtype="i1") + pser_b = pd.Series([], dtype="str") + + self.assert_eq(ps.from_pandas(pser_a), pser_a) + + psser_b = ps.from_pandas(pser_b) + self.assert_eq(psser_b, pser_b) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + self.assert_eq(ps.from_pandas(pser_a), pser_a) + self.assert_eq(ps.from_pandas(pser_b), pser_b) + + def test_all_null_series(self): + pser_a = pd.Series([None, None, None], dtype="float64") + pser_b = pd.Series([None, None, None], dtype="str") + + self.assert_eq(ps.from_pandas(pser_a), pser_a) + + psser_b = ps.from_pandas(pser_b) + self.assert_eq(psser_b, pser_b) + + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + self.assert_eq(ps.from_pandas(pser_a), pser_a) + self.assert_eq(ps.from_pandas(pser_b), pser_b) + + def test_head(self): + psser = self.psser + pser = self.pser + + self.assert_eq(psser.head(3), pser.head(3)) + self.assert_eq(psser.head(0), pser.head(0)) + self.assert_eq(psser.head(-3), pser.head(-3)) + self.assert_eq(psser.head(-10), pser.head(-10)) + + def test_last(self): + with self.assertRaises(TypeError): + self.psser.last("1D") + + index = pd.date_range("2018-04-09", periods=4, freq="2D") + pser = pd.Series([1, 2, 3, 4], index=index) + psser = ps.from_pandas(pser) + self.assert_eq(psser.last("1D"), pser.last("1D")) + + def test_first(self): + with self.assertRaises(TypeError): + self.psser.first("1D") + + index = pd.date_range("2018-04-09", periods=4, freq="2D") + pser = pd.Series([1, 2, 3, 4], index=index) + psser = ps.from_pandas(pser) + self.assert_eq(psser.first("1D"), pser.first("1D")) + + def test_rename(self): + pser = pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + psser = ps.from_pandas(pser) + + pser.name = "renamed" + psser.name = "renamed" + self.assertEqual(psser.name, "renamed") + self.assert_eq(psser, pser) + + pser.name = None + psser.name = None + self.assertEqual(psser.name, None) + self.assert_eq(psser, pser) + + pidx = pser.index + psidx = psser.index + pidx.name = "renamed" + psidx.name = "renamed" + self.assertEqual(psidx.name, "renamed") + self.assert_eq(psidx, pidx) + + expected_error_message = "Series.name must be a hashable type" + with self.assertRaisesRegex(TypeError, expected_error_message): + psser.name = ["renamed"] + with self.assertRaisesRegex(TypeError, expected_error_message): + psser.name = ["0", "1"] + with self.assertRaisesRegex(TypeError, expected_error_message): + ps.Series([1, 2, 3], name=["0", "1"]) + + def test_rename_method(self): + # Series name + pdf = pd.DataFrame({"x": [1, 2, 3, 4, 5, 6, 7]}) + psdf = ps.from_pandas(pdf) + pser, psser = pdf.x, psdf.x + + self.assert_eq(psser.rename("y"), pser.rename("y")) + self.assertEqual(psser.name, "x") # no mutation + self.assert_eq(psser.rename(), pser.rename()) + + self.assert_eq((psser.rename("y") + 1).head(), (pser.rename("y") + 1).head()) + + psser.rename("z", inplace=True) + pser.rename("z", inplace=True) + self.assertEqual(psser.name, "z") + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + expected_error_message = "Series.name must be a hashable type" + with self.assertRaisesRegex(TypeError, expected_error_message): + psser.rename(["0", "1"]) + + # Function index + self.assert_eq(psser.rename(lambda x: x**2), pser.rename(lambda x: x**2)) + self.assert_eq((psser + 1).rename(lambda x: x**2), (pser + 1).rename(lambda x: x**2)) + + expected_error_message = "inplace True is not supported yet for a function 'index'" + with self.assertRaisesRegex(ValueError, expected_error_message): + psser.rename(lambda x: x**2, inplace=True) + + unsupported_index_inputs = (pd.Series([2, 3, 4, 5, 6, 7, 8]), {0: "zero", 1: "one"}) + for index in unsupported_index_inputs: + expected_error_message = ( + "'index' of %s type is not supported yet" % type(index).__name__ + ) + with self.assertRaisesRegex(ValueError, expected_error_message): + psser.rename(index) + + # Series index + # pser = pd.Series(['a', 'b', 'c', 'd', 'e', 'f', 'g'], name='x') + # psser = ps.from_pandas(s) + + # res = psser.rename(pser) + # self.assert_eq(res, pser.rename(pser)) + + # res = psser.rename(psser) + # self.assert_eq(res, pser.rename(pser)) + + # res = psser.rename(lambda x: x**2, inplace=True) + # self.assertis(res, psser) + # s.rename(lambda x: x**2, inplace=True) + # self.assert_eq(psser, pser) + + def test_or(self): + pdf = pd.DataFrame( + { + "left": [True, False, True, False, np.nan, np.nan, True, False, np.nan], + "right": [True, False, False, True, True, False, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["left"] | psdf["right"], pdf["left"] | pdf["right"]) + self.assert_eq(psdf["left"] | True, pdf["left"] | True) + self.assert_eq(psdf["left"] | False, pdf["left"] | False) + self.assert_eq(psdf["left"] | None, pdf["left"] | None) + self.assert_eq(True | psdf["right"], True | pdf["right"]) + self.assert_eq(False | psdf["right"], False | pdf["right"]) + self.assert_eq(None | psdf["right"], None | pdf["right"]) + + @unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" + ) + def test_or_extenstion_dtypes(self): + pdf = pd.DataFrame( + { + "left": [True, False, True, False, np.nan, np.nan, True, False, np.nan], + "right": [True, False, False, True, True, False, np.nan, np.nan, np.nan], + } + ).astype("boolean") + psdf = ps.from_pandas(pdf) + + self._check_extension(psdf["left"] | psdf["right"], pdf["left"] | pdf["right"]) + self._check_extension(psdf["left"] | True, pdf["left"] | True) + self._check_extension(psdf["left"] | False, pdf["left"] | False) + self._check_extension(psdf["left"] | pd.NA, pdf["left"] | pd.NA) + self._check_extension(True | psdf["right"], True | pdf["right"]) + self._check_extension(False | psdf["right"], False | pdf["right"]) + self._check_extension(pd.NA | psdf["right"], pd.NA | pdf["right"]) + + def test_and(self): + pdf = pd.DataFrame( + { + "left": [True, False, True, False, np.nan, np.nan, True, False, np.nan], + "right": [True, False, False, True, True, False, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf["left"] & psdf["right"], pdf["left"] & pdf["right"]) + self.assert_eq(psdf["left"] & True, pdf["left"] & True) + self.assert_eq(psdf["left"] & False, pdf["left"] & False) + self.assert_eq(psdf["left"] & None, pdf["left"] & None) + self.assert_eq(True & psdf["right"], True & pdf["right"]) + self.assert_eq(False & psdf["right"], False & pdf["right"]) + self.assert_eq(None & psdf["right"], None & pdf["right"]) + + @unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" + ) + def test_and_extenstion_dtypes(self): + pdf = pd.DataFrame( + { + "left": [True, False, True, False, np.nan, np.nan, True, False, np.nan], + "right": [True, False, False, True, True, False, np.nan, np.nan, np.nan], + } + ).astype("boolean") + psdf = ps.from_pandas(pdf) + + self._check_extension(psdf["left"] & psdf["right"], pdf["left"] & pdf["right"]) + self._check_extension(psdf["left"] & True, pdf["left"] & True) + self._check_extension(psdf["left"] & False, pdf["left"] & False) + self._check_extension(psdf["left"] & pd.NA, pdf["left"] & pd.NA) + self._check_extension(True & psdf["right"], True & pdf["right"]) + self._check_extension(False & psdf["right"], False & pdf["right"]) + self._check_extension(pd.NA & psdf["right"], pd.NA & pdf["right"]) + + def test_isin(self): + pser = pd.Series(["lama", "cow", "lama", "beetle", "lama", "hippo"], name="animal") + + psser = ps.from_pandas(pser) + + self.assert_eq(psser.isin(["cow", "lama"]), pser.isin(["cow", "lama"])) + self.assert_eq(psser.isin(np.array(["cow", "lama"])), pser.isin(np.array(["cow", "lama"]))) + self.assert_eq(psser.isin({"cow"}), pser.isin({"cow"})) + + pser = pd.Series([np.int64(1), np.int32(1), 1]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.isin([np.int64(1)]), pser.isin([np.int64(1)])) + + msg = "only list-like objects are allowed to be passed to isin()" + with self.assertRaisesRegex(TypeError, msg): + psser.isin(1) + + # when Series have NaN + pser = pd.Series(["lama", "cow", None, "lama", "beetle", "lama", "hippo", None], name="a") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.isin(["cow", "lama"]), pser.isin(["cow", "lama"])) + + pser = pd.Series([None, 5, None, 3, 2, 1, None, 0, 0], name="a") + psser = ps.from_pandas(pser) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq(psser.isin([1, 5, 0, None]), pser.isin([1, 5, 0, None])) + else: + expected = pd.Series( + [False, True, False, False, False, True, False, True, True], name="a" + ) + self.assert_eq(psser.isin([1, 5, 0, None]), expected) + + def test_notnull(self): + pser = pd.Series([1, 2, 3, 4, np.nan, 6], name="x") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.notnull(), pser.notnull()) + + pser = self.pser + psser = self.psser + + self.assert_eq(psser.notnull(), pser.notnull()) + + def test_missing(self): + psser = self.psser + + missing_functions = inspect.getmembers(MissingPandasLikeSeries, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Series.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psser, name)() + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Series.*{}.*is deprecated".format(name) + ): + getattr(psser, name)() + + missing_properties = inspect.getmembers( + MissingPandasLikeSeries, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Series.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psser, name) + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Series.*{}.*is deprecated".format(name) + ): + getattr(psser, name) + + def test_map(self): + pser = pd.Series(["cat", "dog", None, "rabbit"]) + psser = ps.from_pandas(pser) + + # dict correspondence + # Currently pandas API on Spark doesn't return NaN as pandas does. + self.assert_eq(psser.map({}), pser.map({}).replace({np.nan: None})) + + d = defaultdict(lambda: "abc") + self.assertTrue("abc" in repr(psser.map(d))) + self.assert_eq(psser.map(d), pser.map(d)) + + # series correspondence + pser_to_apply = pd.Series(["one", "two", "four"], index=["cat", "dog", "rabbit"]) + self.assert_eq(psser.map(pser_to_apply), pser.map(pser_to_apply)) + self.assert_eq( + psser.map(pser_to_apply, na_action="ignore"), + pser.map(pser_to_apply, na_action="ignore"), + ) + + # function correspondence + self.assert_eq( + psser.map(lambda x: x.upper(), na_action="ignore"), + pser.map(lambda x: x.upper(), na_action="ignore"), + ) + + def to_upper(string) -> str: + return string.upper() if string else "" + + self.assert_eq(psser.map(to_upper), pser.map(to_upper)) + + def tomorrow(date) -> datetime: + return date + timedelta(days=1) + + pser = pd.Series([datetime(2019, 10, 24)]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.map(tomorrow), pser.map(tomorrow)) + + def test_add_prefix(self): + pser = pd.Series([1, 2, 3, 4], name="0") + psser = ps.from_pandas(pser) + self.assert_eq(pser.add_prefix("item_"), psser.add_prefix("item_")) + + pser = pd.Series( + [1, 2, 3], + name="0", + index=pd.MultiIndex.from_tuples([("A", "X"), ("A", "Y"), ("B", "X")]), + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.add_prefix("item_"), psser.add_prefix("item_")) + + def test_add_suffix(self): + pser = pd.Series([1, 2, 3, 4], name="0") + psser = ps.from_pandas(pser) + self.assert_eq(pser.add_suffix("_item"), psser.add_suffix("_item")) + + pser = pd.Series( + [1, 2, 3], + name="0", + index=pd.MultiIndex.from_tuples([("A", "X"), ("A", "Y"), ("B", "X")]), + ) + psser = ps.from_pandas(pser) + self.assert_eq(pser.add_suffix("_item"), psser.add_suffix("_item")) + + def test_xs(self): + midx = pd.MultiIndex( + [["a", "b", "c"], ["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.xs(("a", "lama", "speed")), pser.xs(("a", "lama", "speed"))) + + def test_update(self): + pser = pd.Series([10, 20, 15, 30, 45], name="x") + psser = ps.Series(pser) + + msg = "'other' must be a Series" + with self.assertRaisesRegex(TypeError, msg): + psser.update(10) + + def _get_data(): + pdf = pd.DataFrame( + { + "a": [None, 2, 3, 4, 5, 6, 7, 8, None], + "b": [None, 5, None, 3, 2, 1, None, 0, 0], + "c": [1, 5, 1, 3, 2, 1, 1, 0, 0], + }, + ) + psdf = ps.from_pandas(pdf) + return pdf, psdf + + pdf, psdf = _get_data() + + psdf.a.update(psdf.a) + pdf.a.update(pdf.a) + self.assert_eq(psdf, pdf) + + pdf, psdf = _get_data() + + psdf.a.update(psdf.b) + pdf.a.update(pdf.b) + self.assert_eq(psdf, pdf) + + pdf, psdf = _get_data() + pser = pdf.a + psser = psdf.a + + pser.update(pdf.b) + psser.update(psdf.b) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + def test_where(self): + pser1 = pd.Series([0, 1, 2, 3, 4]) + psser1 = ps.from_pandas(pser1) + + self.assert_eq(pser1.where(pser1 > 3), psser1.where(psser1 > 3).sort_index()) + + def test_mask(self): + pser1 = pd.Series([0, 1, 2, 3, 4]) + psser1 = ps.from_pandas(pser1) + + self.assert_eq(pser1.mask(pser1 > 3), psser1.mask(psser1 > 3).sort_index()) + + def test_getitem(self): + pser = pd.Series([10, 20, 15, 30, 45], ["A", "A", "B", "C", "D"]) + psser = ps.Series(pser) + + self.assert_eq(psser["A"], pser["A"]) + self.assert_eq(psser["B"], pser["B"]) + self.assert_eq(psser[psser > 15], pser[pser > 15]) + + # for MultiIndex + midx = pd.MultiIndex( + [["a", "b", "c"], ["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 0, 0, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], name="0", index=midx) + psser = ps.Series(pser) + + self.assert_eq(psser["a"], pser["a"]) + self.assert_eq(psser["a", "lama"], pser["a", "lama"]) + self.assert_eq(psser[psser > 1.5], pser[pser > 1.5]) + + msg = r"'Key length \(4\) exceeds index depth \(3\)'" + with self.assertRaisesRegex(KeyError, msg): + psser[("a", "lama", "speed", "x")] + + def test_keys(self): + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.keys(), pser.keys()) + + def test_udt(self): + sparse_values = {0: 0.1, 1: 1.1} + sparse_vector = SparseVector(len(sparse_values), sparse_values) + pser = pd.Series([sparse_vector]) + psser = ps.from_pandas(pser) + self.assert_eq(psser, pser) + + def test_repeat(self): + pser = pd.Series(["a", "b", "c"], name="0", index=np.random.rand(3)) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.repeat(3).sort_index(), pser.repeat(3).sort_index()) + self.assert_eq(psser.repeat(0).sort_index(), pser.repeat(0).sort_index()) + + self.assertRaises(ValueError, lambda: psser.repeat(-1)) + self.assertRaises(TypeError, lambda: psser.repeat("abc")) + + pdf = pd.DataFrame({"a": ["a", "b", "c"], "rep": [10, 20, 30]}, index=np.random.rand(3)) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.a.repeat(psdf.rep).sort_index(), pdf.a.repeat(pdf.rep).sort_index()) + + def test_take(self): + pser = pd.Series([100, 200, 300, 400, 500], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.take([0, 2, 4]).sort_values(), pser.take([0, 2, 4]).sort_values()) + self.assert_eq( + psser.take(range(0, 5, 2)).sort_values(), pser.take(range(0, 5, 2)).sort_values() + ) + self.assert_eq(psser.take([-4, -2, 0]).sort_values(), pser.take([-4, -2, 0]).sort_values()) + self.assert_eq( + psser.take(range(-2, 1, 2)).sort_values(), pser.take(range(-2, 1, 2)).sort_values() + ) + + # Checking the type of indices. + self.assertRaises(TypeError, lambda: psser.take(1)) + self.assertRaises(TypeError, lambda: psser.take("1")) + self.assertRaises(TypeError, lambda: psser.take({1, 2})) + self.assertRaises(TypeError, lambda: psser.take({1: None, 2: None})) + + def test_squeeze(self): + # Single value + pser = pd.Series([90]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.squeeze(), pser.squeeze()) + + # Single value with MultiIndex + midx = pd.MultiIndex.from_tuples([("a", "b", "c")]) + pser = pd.Series([90], index=midx) + psser = ps.from_pandas(pser) + self.assert_eq(psser.squeeze(), pser.squeeze()) + + # Multiple values + pser = pd.Series([90, 91, 85]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.squeeze(), pser.squeeze()) + + # Multiple values with MultiIndex + midx = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + pser = pd.Series([90, 91, 85], index=midx) + psser = ps.from_pandas(pser) + self.assert_eq(psser.squeeze(), pser.squeeze()) + + def test_shape(self): + pser = pd.Series(["a", "b", "c"]) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.shape, psser.shape) + + # for MultiIndex + midx = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) + pser = pd.Series(["a", "b", "c"], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.shape, psser.shape) + + def test_item(self): + psser = ps.Series([10, 20]) + self.assertRaises(ValueError, lambda: psser.item()) + + def test_filter(self): + pser = pd.Series([0, 1, 2], index=["one", "two", "three"]) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.filter(items=["one", "three"]), psser.filter(items=["one", "three"])) + self.assert_eq(pser.filter(regex="e$"), psser.filter(regex="e$")) + self.assert_eq(pser.filter(like="hre"), psser.filter(like="hre")) + + with self.assertRaisesRegex(ValueError, "Series does not support columns axis."): + psser.filter(like="hre", axis=1) + + # for MultiIndex + midx = pd.MultiIndex.from_tuples([("one", "x"), ("two", "y"), ("three", "z")]) + pser = pd.Series([0, 1, 2], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq( + pser.filter(items=[("one", "x"), ("three", "z")]), + psser.filter(items=[("one", "x"), ("three", "z")]), + ) + + with self.assertRaisesRegex(TypeError, "Unsupported type list"): + psser.filter(items=[["one", "x"], ("three", "z")]) + + with self.assertRaisesRegex(ValueError, "The item should not be empty."): + psser.filter(items=[(), ("three", "z")]) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43480): Enable SeriesTests.test_iteritems for pandas 2.0.0.", + ) + def test_iteritems(self): + pser = pd.Series(["A", "B", "C"]) + psser = ps.from_pandas(pser) + + for (p_name, p_items), (k_name, k_items) in zip(pser.iteritems(), psser.iteritems()): + self.assert_eq(p_name, k_name) + self.assert_eq(p_items, k_items) + + def test_dot(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq((psdf["b"] * 10).dot(psdf["a"]), (pdf["b"] * 10).dot(pdf["a"])) + self.assert_eq((psdf["b"] * 10).dot(psdf), (pdf["b"] * 10).dot(pdf)) + self.assert_eq((psdf["b"] * 10).dot(psdf + 1), (pdf["b"] * 10).dot(pdf + 1)) + + psdf_other = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=["x", "y", "z"]) + with self.assertRaisesRegex(ValueError, "matrices are not aligned"): + psdf["b"].dot(psdf_other) + + def test_tail(self): + pser = pd.Series(range(1000), name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(pser.tail(), psser.tail()) + self.assert_eq(pser.tail(10), psser.tail(10)) + self.assert_eq(pser.tail(-990), psser.tail(-990)) + self.assert_eq(pser.tail(0), psser.tail(0)) + self.assert_eq(pser.tail(1001), psser.tail(1001)) + self.assert_eq(pser.tail(-1001), psser.tail(-1001)) + self.assert_eq((pser + 1).tail(), (psser + 1).tail()) + self.assert_eq((pser + 1).tail(10), (psser + 1).tail(10)) + self.assert_eq((pser + 1).tail(-990), (psser + 1).tail(-990)) + self.assert_eq((pser + 1).tail(0), (psser + 1).tail(0)) + self.assert_eq((pser + 1).tail(1001), (psser + 1).tail(1001)) + self.assert_eq((pser + 1).tail(-1001), (psser + 1).tail(-1001)) + with self.assertRaisesRegex(TypeError, "bad operand type for unary -: 'str'"): + psser.tail("10") + + def test_apply(self): + psser = self.psser + + def udf(col) -> ps.Series[int]: + return col + 10 + + with self.assertRaisesRegex( + ValueError, + r"Expected the return type of this function to be of scalar type, " + r"but found type SeriesType\[LongType\(\)\]", + ): + psser.apply(udf) + + def test_combine_first(self): + pdf = pd.DataFrame( + { + "A": {"falcon": 330.0, "eagle": 160.0}, + "B": {"falcon": 345.0, "eagle": 200.0, "duck": 30.0}, + } + ) + pser1, pser2 = pdf.A, pdf.B + psdf = ps.from_pandas(pdf) + psser1, psser2 = psdf.A, psdf.B + + self.assert_eq(psser1.combine_first(psser2), pser1.combine_first(pser2)) + + psser1.name = pser1.name = ("X", "A") + psser2.name = pser2.name = ("Y", "B") + + self.assert_eq(psser1.combine_first(psser2), pser1.combine_first(pser2)) + + def test_eq(self): + pser = pd.Series([1, 2, 3, 4, 5, 6], name="x") + psser = ps.from_pandas(pser) + + # other = Series + self.assert_eq(pser.eq(pser), psser.eq(psser)) + self.assert_eq(pser == pser, psser == psser) + + # other = dict + other = {1: None, 2: None, 3: None, 4: None, np.nan: None, 6: None} + self.assert_eq(pser.eq(other), psser.eq(other)) + self.assert_eq(pser == other, psser == other) + + # other = set + other = {1, 2, 3, 4, np.nan, 6} + self.assert_eq(pser.eq(other), psser.eq(other)) + self.assert_eq(pser == other, psser == other) + + # other = list + other = [np.nan, 1, 3, 4, np.nan, 6] + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq(pser.eq(other), psser.eq(other).sort_index()) + self.assert_eq(pser == other, (psser == other).sort_index()) + else: + self.assert_eq(pser.eq(other).rename("x"), psser.eq(other).sort_index()) + self.assert_eq((pser == other).rename("x"), (psser == other).sort_index()) + + # other = tuple + other = (np.nan, 1, 3, 4, np.nan, 6) + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq(pser.eq(other), psser.eq(other).sort_index()) + self.assert_eq(pser == other, (psser == other).sort_index()) + else: + self.assert_eq(pser.eq(other).rename("x"), psser.eq(other).sort_index()) + self.assert_eq((pser == other).rename("x"), (psser == other).sort_index()) + + # other = list with the different length + other = [np.nan, 1, 3, 4, np.nan] + with self.assertRaisesRegex(ValueError, "Lengths must be equal"): + psser.eq(other) + with self.assertRaisesRegex(ValueError, "Lengths must be equal"): + psser == other + + # other = tuple with the different length + other = (np.nan, 1, 3, 4, np.nan) + with self.assertRaisesRegex(ValueError, "Lengths must be equal"): + psser.eq(other) + with self.assertRaisesRegex(ValueError, "Lengths must be equal"): + psser == other + + def test_transform(self): + psser = self.psser + with self.assertRaisesRegex( + NotImplementedError, 'axis should be either 0 or "index" currently.' + ): + psser.transform(lambda x: x + 1, axis=1) + + +class SeriesTests(SeriesTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_series import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_sort.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_sort.py new file mode 100644 index 0000000..a3772bb --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_sort.py @@ -0,0 +1,162 @@ +# +# 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 unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesSortMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_sort_values(self): + pdf = pd.DataFrame({"x": [1, 2, 3, 4, 5, None, 7]}, index=np.random.rand(7)) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + self.assert_eq(psser.sort_values(), pser.sort_values()) + self.assert_eq(psser.sort_values(ignore_index=True), pser.sort_values(ignore_index=True)) + self.assert_eq(psser.sort_values(ascending=False), pser.sort_values(ascending=False)) + self.assert_eq( + psser.sort_values(na_position="first"), pser.sort_values(na_position="first") + ) + + self.assertRaises(ValueError, lambda: psser.sort_values(na_position="invalid")) + + # inplace + # pandas raises an exception when the Series is derived from DataFrame + psser.sort_values(inplace=True) + self.assert_eq(psser, pser.sort_values()) + self.assert_eq(psdf, pdf) + + # pandas raises an exception when the Series is derived from DataFrame + psser.sort_values(inplace=True, ascending=False, ignore_index=True) + self.assert_eq(psser, pser.sort_values(ascending=False, ignore_index=True)) + self.assert_eq(psdf, pdf) + + pser = pdf.x.copy() + psser = psdf.x.copy() + + psser.sort_values(inplace=True) + pser.sort_values(inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + def test_sort_index(self): + pdf = pd.DataFrame({"x": [2, 1, np.nan]}, index=["b", "a", np.nan]) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psser = psdf.x + + # Assert invalid parameters + self.assertRaises(NotImplementedError, lambda: psser.sort_index(axis=1)) + self.assertRaises(NotImplementedError, lambda: psser.sort_index(kind="mergesort")) + self.assertRaises(ValueError, lambda: psser.sort_index(na_position="invalid")) + + # Assert default behavior without parameters + self.assert_eq(psser.sort_index(), pser.sort_index()) + # Assert sorting descending + self.assert_eq(psser.sort_index(ascending=False), pser.sort_index(ascending=False)) + # Assert sorting NA indices first + self.assert_eq(psser.sort_index(na_position="first"), pser.sort_index(na_position="first")) + # Assert ignoring index + self.assert_eq(psser.sort_index(ignore_index=True), pser.sort_index(ignore_index=True)) + + # Assert sorting inplace + # pandas sorts pdf.x by the index and update the column only + # when the Series is derived from DataFrame. + psser.sort_index(inplace=True) + self.assert_eq(psser, pser.sort_index()) + self.assert_eq(psdf, pdf) + + # pandas sorts pdf.x by the index and update the column only + # when the Series is derived from DataFrame. + psser.sort_index(inplace=True, ascending=False, ignore_index=True) + self.assert_eq(psser, pser.sort_index(ascending=False, ignore_index=True)) + self.assert_eq(psdf, pdf) + + pser = pdf.x.copy() + psser = psdf.x.copy() + + psser.sort_index(inplace=True) + pser.sort_index(inplace=True) + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + # Assert multi-indices + pser = pd.Series(range(4), index=[["b", "b", "a", "a"], [1, 0, 1, 0]], name="0") + psser = ps.from_pandas(pser) + self.assert_eq(psser.sort_index(), pser.sort_index()) + self.assert_eq(psser.sort_index(level=[1, 0]), pser.sort_index(level=[1, 0])) + + self.assert_eq(psser.reset_index().sort_index(), pser.reset_index().sort_index()) + + def test_searchsorted(self): + pser1 = pd.Series([1, 2, 2, 3]) + + index2 = pd.date_range("2018-04-09", periods=4, freq="2D") + pser2 = pd.Series([1, 2, 3, 4], index=index2) + + index3 = pd.MultiIndex.from_tuples( + [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] + ) + pser3 = pd.Series([1.0, 2.0, 3.0], index=index3, name="name") + + pser4 = pd.Series([]) + + for pser in [pser1, pser2, pser3, pser4]: + psser = ps.from_pandas(pser) + for value in [0.5, 1, 2, 3.0, 4, 5]: + for side in ["left", "right"]: + self.assert_eq( + pser.searchsorted(value, side=side), + psser.searchsorted(value, side=side), + ) + + with self.assertRaisesRegex(ValueError, "Invalid side"): + ps.from_pandas(pser1).searchsorted(1.1, side=[1, 2]) + with self.assertRaisesRegex(ValueError, "Invalid side"): + ps.from_pandas(pser1).searchsorted(1.1, side="middle") + + +class SeriesSortTests(SeriesSortMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_sort import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_stat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_stat.py new file mode 100644 index 0000000..0d6e242 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/series/test_stat.py @@ -0,0 +1,780 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd +from decimal import Decimal + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesStatMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_nunique(self): + pser = pd.Series([1, 2, 1, np.nan]) + psser = ps.from_pandas(pser) + + # Assert NaNs are dropped by default + nunique_result = psser.nunique() + self.assertEqual(nunique_result, 2) + self.assert_eq(nunique_result, pser.nunique()) + + # Assert including NaN values + nunique_result = psser.nunique(dropna=False) + self.assertEqual(nunique_result, 3) + self.assert_eq(nunique_result, pser.nunique(dropna=False)) + + # Assert approximate counts + self.assertEqual(ps.Series(range(100)).nunique(approx=True), 103) + self.assertEqual(ps.Series(range(100)).nunique(approx=True, rsd=0.01), 100) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43464): Enable SeriesTests.test_value_counts for pandas 2.0.0.", + ) + def test_value_counts(self): + # this is also containing test for Index & MultiIndex + pser = pd.Series( + [1, 2, 1, 3, 3, np.nan, 1, 4, 2, np.nan, 3, np.nan, 3, 1, 3], + index=[1, 2, 1, 3, 3, np.nan, 1, 4, 2, np.nan, 3, np.nan, 3, 1, 3], + name="x", + ) + psser = ps.from_pandas(pser) + + exp = pser.value_counts() + res = psser.value_counts() + self.assertEqual(res.name, exp.name) + self.assert_eq(res, exp) + + self.assert_eq(psser.value_counts(normalize=True), pser.value_counts(normalize=True)) + self.assert_eq(psser.value_counts(ascending=True), pser.value_counts(ascending=True)) + self.assert_eq( + psser.value_counts(normalize=True, dropna=False), + pser.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.value_counts(ascending=True, dropna=False), + pser.value_counts(ascending=True, dropna=False), + ) + + self.assert_eq( + psser.index.value_counts(normalize=True), pser.index.value_counts(normalize=True) + ) + self.assert_eq( + psser.index.value_counts(ascending=True), pser.index.value_counts(ascending=True) + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + ) + + with self.assertRaisesRegex( + NotImplementedError, "value_counts currently does not support bins" + ): + psser.value_counts(bins=3) + + pser.name = "index" + psser.name = "index" + self.assert_eq(psser.value_counts(), pser.value_counts()) + + # Series from DataFrame + pdf = pd.DataFrame({"a": [2, 2, 3], "b": [None, 1, None]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.a.value_counts(normalize=True), pdf.a.value_counts(normalize=True)) + self.assert_eq(psdf.a.value_counts(ascending=True), pdf.a.value_counts(ascending=True)) + self.assert_eq( + psdf.a.value_counts(normalize=True, dropna=False), + pdf.a.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psdf.a.value_counts(ascending=True, dropna=False), + pdf.a.value_counts(ascending=True, dropna=False), + ) + + self.assert_eq( + psser.index.value_counts(normalize=True), pser.index.value_counts(normalize=True) + ) + self.assert_eq( + psser.index.value_counts(ascending=True), pser.index.value_counts(ascending=True) + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + ) + + # Series with NaN index + pser = pd.Series([3, 2, 3, 1, 2, 3], index=[2.0, None, 5.0, 5.0, None, 5.0]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.value_counts(normalize=True), pser.value_counts(normalize=True)) + self.assert_eq(psser.value_counts(ascending=True), pser.value_counts(ascending=True)) + self.assert_eq( + psser.value_counts(normalize=True, dropna=False), + pser.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.value_counts(ascending=True, dropna=False), + pser.value_counts(ascending=True, dropna=False), + ) + + self.assert_eq( + psser.index.value_counts(normalize=True), pser.index.value_counts(normalize=True) + ) + self.assert_eq( + psser.index.value_counts(ascending=True), pser.index.value_counts(ascending=True) + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + ) + + # Series with MultiIndex + pser.index = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("y", "c"), ("x", "a"), ("y", "c"), ("x", "a")] + ) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.value_counts(normalize=True), pser.value_counts(normalize=True)) + self.assert_eq(psser.value_counts(ascending=True), pser.value_counts(ascending=True)) + self.assert_eq( + psser.value_counts(normalize=True, dropna=False), + pser.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.value_counts(ascending=True, dropna=False), + pser.value_counts(ascending=True, dropna=False), + ) + + # FIXME: MultiIndex.value_counts returns wrong indices. + self.assert_eq( + psser.index.value_counts(normalize=True), + pser.index.value_counts(normalize=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True), + pser.index.value_counts(ascending=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + almost=True, + ) + + # Series with MultiIndex some of index has NaN + pser.index = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", None), ("y", "c"), ("x", "a"), ("y", "c"), ("x", "a")] + ) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.value_counts(normalize=True), pser.value_counts(normalize=True)) + self.assert_eq(psser.value_counts(ascending=True), pser.value_counts(ascending=True)) + self.assert_eq( + psser.value_counts(normalize=True, dropna=False), + pser.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.value_counts(ascending=True, dropna=False), + pser.value_counts(ascending=True, dropna=False), + ) + + # FIXME: MultiIndex.value_counts returns wrong indices. + self.assert_eq( + psser.index.value_counts(normalize=True), + pser.index.value_counts(normalize=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True), + pser.index.value_counts(ascending=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + almost=True, + ) + + # Series with MultiIndex some of index is NaN. + pser.index = pd.MultiIndex.from_tuples( + [("x", "a"), None, ("y", "c"), ("x", "a"), ("y", "c"), ("x", "a")] + ) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.value_counts(normalize=True), pser.value_counts(normalize=True)) + self.assert_eq(psser.value_counts(ascending=True), pser.value_counts(ascending=True)) + self.assert_eq( + psser.value_counts(normalize=True, dropna=False), + pser.value_counts(normalize=True, dropna=False), + ) + self.assert_eq( + psser.value_counts(ascending=True, dropna=False), + pser.value_counts(ascending=True, dropna=False), + ) + + # FIXME: MultiIndex.value_counts returns wrong indices. + self.assert_eq( + psser.index.value_counts(normalize=True), + pser.index.value_counts(normalize=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True), + pser.index.value_counts(ascending=True), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(normalize=True, dropna=False), + pser.index.value_counts(normalize=True, dropna=False), + almost=True, + ) + self.assert_eq( + psser.index.value_counts(ascending=True, dropna=False), + pser.index.value_counts(ascending=True, dropna=False), + almost=True, + ) + + def test_nsmallest(self): + sample_lst = [1, 2, 3, 4, np.nan, 6] + pser = pd.Series(sample_lst, name="x") + psser = ps.Series(sample_lst, name="x") + self.assert_eq(psser.nsmallest(n=3), pser.nsmallest(n=3)) + self.assert_eq(psser.nsmallest(), pser.nsmallest()) + self.assert_eq((psser + 1).nsmallest(), (pser + 1).nsmallest()) + + def test_nlargest(self): + sample_lst = [1, 2, 3, 4, np.nan, 6] + pser = pd.Series(sample_lst, name="x") + psser = ps.Series(sample_lst, name="x") + self.assert_eq(psser.nlargest(n=3), pser.nlargest(n=3)) + self.assert_eq(psser.nlargest(), pser.nlargest()) + self.assert_eq((psser + 1).nlargest(), (pser + 1).nlargest()) + + def test_is_unique(self): + # We can't use pandas' is_unique for comparison. pandas 0.23 ignores None + pser = pd.Series([1, 2, 2, None, None]) + psser = ps.from_pandas(pser) + self.assertEqual(False, psser.is_unique) + self.assertEqual(False, (psser + 1).is_unique) + + pser = pd.Series([1, None, None]) + psser = ps.from_pandas(pser) + self.assertEqual(False, psser.is_unique) + self.assertEqual(False, (psser + 1).is_unique) + + pser = pd.Series([1]) + psser = ps.from_pandas(pser) + self.assertEqual(pser.is_unique, psser.is_unique) + self.assertEqual((pser + 1).is_unique, (psser + 1).is_unique) + + pser = pd.Series([1, 1, 1]) + psser = ps.from_pandas(pser) + self.assertEqual(pser.is_unique, psser.is_unique) + self.assertEqual((pser + 1).is_unique, (psser + 1).is_unique) + + def test_median(self): + with self.assertRaisesRegex(TypeError, "accuracy must be an integer; however"): + ps.Series([24.0, 21.0, 25.0, 33.0, 26.0]).median(accuracy="a") + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43463): Enable SeriesTests.test_rank for pandas 2.0.0.", + ) + def test_rank(self): + pser = pd.Series([1, 2, 3, 1], name="x") + psser = ps.from_pandas(pser) + self.assert_eq(pser.rank(), psser.rank().sort_index()) + self.assert_eq(pser.rank().sum(), psser.rank().sum()) + self.assert_eq(pser.rank(ascending=False), psser.rank(ascending=False).sort_index()) + self.assert_eq(pser.rank(method="min"), psser.rank(method="min").sort_index()) + self.assert_eq(pser.rank(method="max"), psser.rank(method="max").sort_index()) + self.assert_eq(pser.rank(method="first"), psser.rank(method="first").sort_index()) + self.assert_eq(pser.rank(method="dense"), psser.rank(method="dense").sort_index()) + + non_numeric_pser = pd.Series(["a", "c", "b", "d"], name="x", index=[10, 11, 12, 13]) + non_numeric_psser = ps.from_pandas(non_numeric_pser) + self.assert_eq( + non_numeric_pser.rank(numeric_only=True), + non_numeric_psser.rank(numeric_only=True), + ) + self.assert_eq( + non_numeric_pser.rank(numeric_only=None), + non_numeric_psser.rank(numeric_only=None).sort_index(), + ) + self.assert_eq( + non_numeric_pser.rank(numeric_only=False), + non_numeric_psser.rank(numeric_only=False).sort_index(), + ) + self.assert_eq( + (non_numeric_pser + "x").rank(numeric_only=True), + (non_numeric_psser + "x").rank(numeric_only=True), + ) + + msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'" + with self.assertRaisesRegex(ValueError, msg): + psser.rank(method="nothing") + + msg = "method must be one of 'average', 'min', 'max', 'first', 'dense'" + with self.assertRaisesRegex(ValueError, msg): + psser.rank(method="nothing") + + midx = pd.MultiIndex.from_tuples([("a", "b"), ("a", "c"), ("b", "c"), ("c", "d")]) + pser.index = midx + psser = ps.from_pandas(pser) + msg = "rank do not support MultiIndex now" + with self.assertRaisesRegex(NotImplementedError, msg): + psser.rank(method="min") + + def test_round(self): + pser = pd.Series([0.028208, 0.038683, 0.877076], name="x") + psser = ps.from_pandas(pser) + self.assert_eq(pser.round(2), psser.round(2)) + msg = "decimals must be an integer" + with self.assertRaisesRegex(TypeError, msg): + psser.round(1.5) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43469): Enable SeriesTests.test_quantile for pandas 2.0.0.", + ) + def test_quantile(self): + pser = pd.Series([]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.quantile(0.5), pser.quantile(0.5)) + self.assert_eq(psser.quantile([0.25, 0.5, 0.75]), pser.quantile([0.25, 0.5, 0.75])) + + with self.assertRaisesRegex(TypeError, "accuracy must be an integer; however"): + ps.Series([24.0, 21.0, 25.0, 33.0, 26.0]).quantile(accuracy="a") + with self.assertRaisesRegex(TypeError, "q must be a float or an array of floats;"): + ps.Series([24.0, 21.0, 25.0, 33.0, 26.0]).quantile(q=1) + with self.assertRaisesRegex(TypeError, "q must be a float or an array of floats;"): + ps.Series([24.0, 21.0, 25.0, 33.0, 26.0]).quantile(q=["a"]) + with self.assertRaisesRegex( + ValueError, "percentiles should all be in the interval \\[0, 1\\]" + ): + ps.Series([24.0, 21.0, 25.0, 33.0, 26.0]).quantile(q=1.1) + + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + ps.Series(["a", "b", "c"]).quantile() + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + ps.Series(["a", "b", "c"]).quantile([0.25, 0.5, 0.75]) + + def test_pct_change(self): + pser = pd.Series([90, 91, 85], index=[2, 4, 1]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.pct_change(), pser.pct_change(), check_exact=False) + self.assert_eq(psser.pct_change().sum(), pser.pct_change().sum(), almost=True) + self.assert_eq(psser.pct_change(periods=2), pser.pct_change(periods=2), check_exact=False) + self.assert_eq(psser.pct_change(periods=-1), pser.pct_change(periods=-1), check_exact=False) + self.assert_eq(psser.pct_change(periods=-100000000), pser.pct_change(periods=-100000000)) + self.assert_eq(psser.pct_change(periods=100000000), pser.pct_change(periods=100000000)) + + # for MultiIndex + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.pct_change(), pser.pct_change(), check_exact=False) + self.assert_eq(psser.pct_change().sum(), pser.pct_change().sum(), almost=True) + self.assert_eq(psser.pct_change(periods=2), pser.pct_change(periods=2), check_exact=False) + self.assert_eq(psser.pct_change(periods=-1), pser.pct_change(periods=-1), check_exact=False) + self.assert_eq(psser.pct_change(periods=-100000000), pser.pct_change(periods=-100000000)) + self.assert_eq(psser.pct_change(periods=100000000), pser.pct_change(periods=100000000)) + + def test_divmod(self): + pser = pd.Series([100, None, 300, None, 500], name="Koalas") + psser = ps.from_pandas(pser) + + kdiv, kmod = psser.divmod(-100) + pdiv, pmod = pser.divmod(-100) + self.assert_eq(kdiv, pdiv) + self.assert_eq(kmod, pmod) + + kdiv, kmod = psser.divmod(100) + pdiv, pmod = pser.divmod(100) + self.assert_eq(kdiv, pdiv) + self.assert_eq(kmod, pmod) + + def test_rdivmod(self): + pser = pd.Series([100, None, 300, None, 500]) + psser = ps.from_pandas(pser) + + krdiv, krmod = psser.rdivmod(-100) + prdiv, prmod = pser.rdivmod(-100) + self.assert_eq(krdiv, prdiv) + self.assert_eq(krmod, prmod) + + krdiv, krmod = psser.rdivmod(100) + prdiv, prmod = pser.rdivmod(100) + self.assert_eq(krdiv, prdiv) + self.assert_eq(krmod, prmod) + + def test_mod(self): + pser = pd.Series([100, None, -300, None, 500, -700], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.mod(-150), pser.mod(-150)) + self.assert_eq(psser.mod(0), pser.mod(0)) + self.assert_eq(psser.mod(150), pser.mod(150)) + + pdf = pd.DataFrame({"a": [100, None, -300, None, 500, -700], "b": [150] * 6}) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.a.mod(psdf.b), pdf.a.mod(pdf.b)) + + def test_mode(self): + pser = pd.Series([0, 0, 1, 1, 1, np.nan, np.nan, np.nan]) + psser = ps.from_pandas(pser) + self.assert_eq(psser.mode(), pser.mode()) + self.assert_eq( + psser.mode(dropna=False).sort_values().reset_index(drop=True), + pser.mode(dropna=False).sort_values().reset_index(drop=True), + ) + + pser.name = "x" + psser = ps.from_pandas(pser) + if LooseVersion(pd.__version__) < LooseVersion("1.4"): + # Due to pandas bug: https://github.com/pandas-dev/pandas/issues/46737 + psser.name = None + self.assert_eq(psser.mode(), pser.mode()) + else: + self.assert_eq(psser.mode(), pser.mode()) + self.assert_eq( + psser.mode(dropna=False).sort_values().reset_index(drop=True), + pser.mode(dropna=False).sort_values().reset_index(drop=True), + ) + + def test_rmod(self): + pser = pd.Series([100, None, -300, None, 500, -700], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(psser.rmod(-150), pser.rmod(-150)) + self.assert_eq(psser.rmod(0), pser.rmod(0)) + self.assert_eq(psser.rmod(150), pser.rmod(150)) + + pdf = pd.DataFrame({"a": [100, None, -300, None, 500, -700], "b": [150] * 6}) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.a.rmod(psdf.b), pdf.a.rmod(pdf.b)) + + def test_div_zero_and_nan(self): + pser = pd.Series([100, None, -300, None, 500, -700, np.inf, -np.inf], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(pser.div(0), psser.div(0)) + self.assert_eq(pser.truediv(0), psser.truediv(0)) + self.assert_eq(pser / 0, psser / 0) + self.assert_eq(pser.div(np.nan), psser.div(np.nan)) + self.assert_eq(pser.truediv(np.nan), psser.truediv(np.nan)) + self.assert_eq(pser / np.nan, psser / np.nan) + + self.assert_eq(pser.floordiv(0), psser.floordiv(0)) + self.assert_eq(pser // 0, psser // 0) + self.assert_eq(pser.floordiv(np.nan), psser.floordiv(np.nan)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43468): Enable SeriesTests.test_mad for pandas 2.0.0.", + ) + def test_mad(self): + pser = pd.Series([1, 2, 3, 4], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(pser.mad(), psser.mad()) + + pser = pd.Series([None, -2, 5, 10, 50, np.nan, -20], name="Koalas") + psser = ps.from_pandas(pser) + + self.assert_eq(pser.mad(), psser.mad()) + + pmidx = pd.MultiIndex.from_tuples( + [("a", "1"), ("a", "2"), ("b", "1"), ("b", "2"), ("c", "1")] + ) + pser = pd.Series([1, 2, 3, 4, 5], name="Koalas") + pser.index = pmidx + psser = ps.from_pandas(pser) + + self.assert_eq(pser.mad(), psser.mad()) + + pmidx = pd.MultiIndex.from_tuples( + [("a", "1"), ("a", "2"), ("b", "1"), ("b", "2"), ("c", "1")] + ) + pser = pd.Series([None, -2, 5, 50, np.nan], name="Koalas") + pser.index = pmidx + psser = ps.from_pandas(pser) + + # Mark almost as True to avoid precision issue like: + # "21.555555555555554 != 21.555555555555557" + self.assert_eq(pser.mad(), psser.mad(), almost=True) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43481): Enable SeriesTests.test_product for pandas 2.0.0.", + ) + def test_product(self): + pser = pd.Series([10, 20, 30, 40, 50]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + # Containing NA values + pser = pd.Series([10, np.nan, 30, np.nan, 50]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod(), almost=True) + + # All-NA values + pser = pd.Series([np.nan, np.nan, np.nan]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + # Empty Series + pser = pd.Series([]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + # Boolean Series + pser = pd.Series([True, True, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + pser = pd.Series([False, False, False]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + pser = pd.Series([True, False, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(), psser.prod()) + + # With `min_count` parameter + pser = pd.Series([10, 20, 30, 40, 50]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(min_count=5), psser.prod(min_count=5)) + self.assert_eq(pser.prod(min_count=6), psser.prod(min_count=6)) + + pser = pd.Series([10, np.nan, 30, np.nan, 50]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(min_count=3), psser.prod(min_count=3), almost=True) + self.assert_eq(pser.prod(min_count=4), psser.prod(min_count=4)) + + pser = pd.Series([np.nan, np.nan, np.nan]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(min_count=1), psser.prod(min_count=1)) + + pser = pd.Series([]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.prod(min_count=1), psser.prod(min_count=1)) + + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + ps.Series(["a", "b", "c"]).prod() + with self.assertRaisesRegex( + TypeError, "Could not convert datetime64\\[ns\\] \\(timestamp.*\\) to numeric" + ): + ps.Series([pd.Timestamp("2016-01-01") for _ in range(3)]).prod() + with self.assertRaisesRegex(NotImplementedError, "Series does not support columns axis."): + psser.prod(axis=1) + + def test_hasnans(self): + # BooleanType + pser = pd.Series([True, False, True, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + pser = pd.Series([True, False, np.nan, True]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + # TimestampType + pser = pd.Series([pd.Timestamp("2020-07-30") for _ in range(3)]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + pser = pd.Series([pd.Timestamp("2020-07-30"), np.nan, pd.Timestamp("2020-07-30")]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + # DecimalType + pser = pd.Series([Decimal("0.1"), Decimal("NaN")]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + # empty + pser = pd.Series([]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.hasnans, psser.hasnans) + + def test_pow_and_rpow(self): + pser = pd.Series([1, 2, np.nan]) + psser = ps.from_pandas(pser) + + self.assert_eq(pser.pow(np.nan), psser.pow(np.nan)) + self.assert_eq(pser**np.nan, psser**np.nan) + self.assert_eq(pser.rpow(np.nan), psser.rpow(np.nan)) + self.assert_eq(1**pser, 1**psser) + + def test_autocorr(self): + pdf = pd.DataFrame({"s1": [0.90010907, 0.13484424, 0.62036035]}) + self._test_autocorr(pdf) + + pdf = pd.DataFrame({"s1": [0.90010907, np.nan, 0.13484424, 0.62036035]}) + self._test_autocorr(pdf) + + pdf = pd.DataFrame({"s1": [0.2, 0.0, 0.6, 0.2, np.nan, 0.5, 0.6]}) + self._test_autocorr(pdf) + + psser = ps.from_pandas(pdf["s1"]) + with self.assertRaisesRegex(TypeError, r"lag should be an int; however, got"): + psser.autocorr(1.0) + + def _test_autocorr(self, pdf): + psdf = ps.from_pandas(pdf) + for lag in range(-10, 10): + p_autocorr = pdf["s1"].autocorr(lag) + ps_autocorr = psdf["s1"].autocorr(lag) + self.assert_eq(p_autocorr, ps_autocorr, almost=True) + + def test_cov(self): + pdf = pd.DataFrame( + { + "s1": ["a", "b", "c"], + "s2": [0.12528585, 0.26962463, 0.51111198], + }, + index=[0, 1, 2], + ) + psdf = ps.from_pandas(pdf) + with self.assertRaisesRegex(TypeError, "unsupported dtype: object"): + psdf["s1"].cov(psdf["s2"]) + with self.assertRaisesRegex(TypeError, "unsupported dtype: object"): + psdf["s2"].cov(psdf["s1"]) + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf["s2"].cov(psdf["s2"], ddof="ddof") + + pdf = pd.DataFrame( + { + "s1": [0.90010907, 0.13484424, 0.62036035], + "s2": [0.12528585, 0.26962463, 0.51111198], + }, + index=[0, 1, 2], + ) + self._test_cov(pdf) + + pdf = pd.DataFrame( + { + "s1": [0.90010907, np.nan, 0.13484424, 0.62036035], + "s2": [0.12528585, 0.81131178, 0.26962463, 0.51111198], + }, + index=[0, 1, 2, 3], + ) + self._test_cov(pdf) + + def _test_cov(self, pdf): + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf["s1"].cov(pdf["s2"]), psdf["s1"].cov(psdf["s2"]), almost=True) + self.assert_eq( + pdf["s1"].cov(pdf["s2"], ddof=2), psdf["s1"].cov(psdf["s2"], ddof=2), almost=True + ) + + self.assert_eq( + pdf["s1"].cov(pdf["s2"], min_periods=3), + psdf["s1"].cov(psdf["s2"], min_periods=3), + almost=True, + ) + self.assert_eq( + pdf["s1"].cov(pdf["s2"], min_periods=3, ddof=-1), + psdf["s1"].cov(psdf["s2"], min_periods=3, ddof=-1), + almost=True, + ) + + self.assert_eq( + pdf["s1"].cov(pdf["s2"], min_periods=4), + psdf["s1"].cov(psdf["s2"], min_periods=4), + almost=True, + ) + self.assert_eq( + pdf["s1"].cov(pdf["s2"], min_periods=4, ddof=3), + psdf["s1"].cov(psdf["s2"], min_periods=4, ddof=3), + almost=True, + ) + + def test_series_stat_fail(self): + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).mean() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).skew() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).kurtosis() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).std() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).var() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).median() + with self.assertRaisesRegex(TypeError, "Could not convert object"): + ps.Series(["a", "b", "c"]).sem() + + +class SeriesStatTests(SeriesStatMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.series.test_stat import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_categorical.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_categorical.py new file mode 100644 index 0000000..dae882a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_categorical.py @@ -0,0 +1,741 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase, TestUtils + + +class CategoricalTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + { + "a": pd.Categorical([1, 2, 3, 1, 2, 3]), + "b": pd.Categorical( + ["b", "a", "c", "c", "b", "a"], categories=["c", "b", "d", "a"] + ), + }, + ) + + @property + def df_pair(self): + return self.pdf, self.psdf + + def test_categorical_frame(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf, pdf) + self.assert_eq(psdf.a, pdf.a) + self.assert_eq(psdf.b, pdf.b) + self.assert_eq(psdf.index, pdf.index) + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psdf.sort_values("b"), pdf.sort_values("b")) + + def test_categorical_series(self): + pser = pd.Series([1, 2, 3], dtype="category") + psser = ps.Series([1, 2, 3], dtype="category") + + self.assert_eq(psser, pser) + self.assert_eq(psser.cat.categories, pser.cat.categories) + self.assert_eq(psser.cat.codes, pser.cat.codes) + self.assert_eq(psser.cat.ordered, pser.cat.ordered) + + with self.assertRaisesRegex(ValueError, "Cannot call CategoricalAccessor on type int64"): + ps.Series([1, 2, 3]).cat + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43566): Enable CategoricalTests.test_categories_setter for pandas 2.0.0.", + ) + def test_categories_setter(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + pser.cat.categories = ["z", "y", "x"] + psser.cat.categories = ["z", "y", "x"] + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + # Bug in pandas 1.3. dtype is not updated properly with `inplace` argument. + pser = pser.astype(CategoricalDtype(categories=["x", "y", "z"])) + + self.assert_eq(pser, psser) + self.assert_eq(pdf, psdf) + + with self.assertRaises(ValueError): + psser.cat.categories = [1, 2, 3, 4] + + def test_add_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + self.assert_eq(pser.cat.add_categories(4), psser.cat.add_categories(4)) + self.assert_eq(pser.cat.add_categories([4, 5]), psser.cat.add_categories([4, 5])) + self.assert_eq(pser.cat.add_categories([]), psser.cat.add_categories([])) + + pser = pser.cat.add_categories(4) + psser = psser.cat.add_categories(4) + + self.assertRaises(ValueError, lambda: psser.cat.add_categories(4)) + self.assertRaises(ValueError, lambda: psser.cat.add_categories([5, 5])) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43605): Enable CategoricalTests.test_remove_categories for pandas 2.0.0.", + ) + def test_remove_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + self.assert_eq(pser.cat.remove_categories(2), psser.cat.remove_categories(2)) + self.assert_eq(pser.cat.remove_categories([1, 3]), psser.cat.remove_categories([1, 3])) + self.assert_eq(pser.cat.remove_categories([]), psser.cat.remove_categories([])) + self.assert_eq(pser.cat.remove_categories([2, 2]), psser.cat.remove_categories([2, 2])) + self.assert_eq( + pser.cat.remove_categories([1, 2, 3]), psser.cat.remove_categories([1, 2, 3]) + ) + self.assert_eq(pser.cat.remove_categories(None), psser.cat.remove_categories(None)) + self.assert_eq(pser.cat.remove_categories([None]), psser.cat.remove_categories([None])) + + self.assertRaises(ValueError, lambda: psser.cat.remove_categories(4)) + self.assertRaises(ValueError, lambda: psser.cat.remove_categories([4, None])) + + def test_remove_unused_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + self.assert_eq(pser.cat.remove_unused_categories(), psser.cat.remove_unused_categories()) + + pser = pser.cat.add_categories(4) + pser = pser.cat.remove_categories(2) + psser = psser.cat.add_categories(4) + psser = psser.cat.remove_categories(2) + + self.assert_eq(pser.cat.remove_unused_categories(), psser.cat.remove_unused_categories()) + + def test_reorder_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + self.assert_eq( + pser.cat.reorder_categories([1, 2, 3]), psser.cat.reorder_categories([1, 2, 3]) + ) + self.assert_eq( + pser.cat.reorder_categories([1, 2, 3], ordered=True), + psser.cat.reorder_categories([1, 2, 3], ordered=True), + ) + self.assert_eq( + pser.cat.reorder_categories([3, 2, 1]), psser.cat.reorder_categories([3, 2, 1]) + ) + self.assert_eq( + pser.cat.reorder_categories([3, 2, 1], ordered=True), + psser.cat.reorder_categories([3, 2, 1], ordered=True), + ) + + self.assertRaises(ValueError, lambda: psser.cat.reorder_categories([1, 2])) + self.assertRaises(ValueError, lambda: psser.cat.reorder_categories([1, 2, 4])) + self.assertRaises(ValueError, lambda: psser.cat.reorder_categories([1, 2, 2])) + self.assertRaises(TypeError, lambda: psser.cat.reorder_categories(1)) + self.assertRaises(TypeError, lambda: psdf.b.cat.reorder_categories("abcd")) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43565): Enable CategoricalTests.test_as_ordered_unordered for pandas 2.0.0.", + ) + def test_as_ordered_unordered(self): + pdf, psdf = self.df_pair + + pser = pdf.a + psser = psdf.a + + # as_ordered + self.assert_eq(pser.cat.as_ordered(), psser.cat.as_ordered()) + + pser.cat.as_ordered(inplace=True) + psser.cat.as_ordered(inplace=True) + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + # Bug in pandas 1.3. dtype is not updated properly with `inplace` argument. + pser = pser.astype(CategoricalDtype(categories=[1, 2, 3], ordered=True)) + + self.assert_eq(pser, psser) + self.assert_eq(pdf, psdf) + + # as_unordered + self.assert_eq(pser.cat.as_unordered(), psser.cat.as_unordered()) + + pser.cat.as_unordered(inplace=True) + psser.cat.as_unordered(inplace=True) + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + # Bug in pandas 1.3. dtype is not updated properly with `inplace` argument. + pser = pser.astype(CategoricalDtype(categories=[1, 2, 3], ordered=False)) + pdf.a = pser + + self.assert_eq(pser, psser) + self.assert_eq(pdf, psdf) + + def test_astype(self): + pser = pd.Series(["a", "b", "c"]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.astype("category"), pser.astype("category")) + self.assert_eq( + psser.astype(CategoricalDtype(["c", "a", "b"])), + pser.astype(CategoricalDtype(["c", "a", "b"])), + ) + + pcser = pser.astype(CategoricalDtype(["c", "a", "b"])) + pscser = psser.astype(CategoricalDtype(["c", "a", "b"])) + + self.assert_eq(pscser.astype("category"), pcser.astype("category")) + + # CategoricalDtype is not updated if the dtype is same from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + pscser.astype(CategoricalDtype(["b", "c", "a"])), + pcser.astype(CategoricalDtype(["b", "c", "a"])), + ) + else: + self.assert_eq( + pscser.astype(CategoricalDtype(["b", "c", "a"])), + pcser, + ) + + self.assert_eq(pscser.astype(str), pcser.astype(str)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43564): Enable CategoricalTests.test_factorize for pandas 2.0.0.", + ) + def test_factorize(self): + pser = pd.Series(["a", "b", "c", None], dtype=CategoricalDtype(["c", "a", "d", "b"])) + psser = ps.from_pandas(pser) + + pcodes, puniques = pser.factorize() + kcodes, kuniques = psser.factorize() + + self.assert_eq(kcodes.tolist(), pcodes.tolist()) + self.assert_eq(kuniques, puniques) + + pcodes, puniques = pser.factorize(na_sentinel=-2) + kcodes, kuniques = psser.factorize(na_sentinel=-2) + + self.assert_eq(kcodes.tolist(), pcodes.tolist()) + self.assert_eq(kuniques, puniques) + + def test_frame_apply(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf.apply(lambda x: x).sort_index(), pdf.apply(lambda x: x).sort_index()) + self.assert_eq( + psdf.apply(lambda x: x, axis=1).sort_index(), + pdf.apply(lambda x: x, axis=1).sort_index(), + ) + + def test_frame_apply_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_frame_apply() + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c"]) + + def categorize(ser) -> ps.Series[dtype]: + return ser.astype(dtype) + + self.assert_eq( + psdf.apply(categorize).sort_values(["a", "b"]).reset_index(drop=True), + pdf.apply(categorize).sort_values(["a", "b"]).reset_index(drop=True), + ) + + def test_frame_transform(self): + pdf, psdf = self.df_pair + + self.assert_eq(psdf.transform(lambda x: x), pdf.transform(lambda x: x)) + self.assert_eq(psdf.transform(lambda x: x.cat.codes), pdf.transform(lambda x: x.cat.codes)) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + self.assert_eq( + psdf.transform(lambda x: x.astype(dtype)).sort_index(), + pdf.transform(lambda x: x.astype(dtype)).sort_index(), + ) + + def test_frame_transform_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_frame_transform() + + pdf, psdf = self.df_pair + + def codes(pser) -> ps.Series[np.int8]: + return pser.cat.codes + + self.assert_eq(psdf.transform(codes), pdf.transform(codes)) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + def to_category(pser) -> ps.Series[dtype]: + return pser.astype(dtype) + + self.assert_eq( + psdf.transform(to_category).sort_index(), pdf.transform(to_category).sort_index() + ) + + def test_series_apply(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.a.apply(lambda x: x).sort_index(), pdf.a.apply(lambda x: x).sort_index() + ) + + def test_series_apply_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_series_apply() + + pdf, psdf = self.df_pair + ret = psdf.a.dtype + + def identity(pser) -> ret: + return pser + + self.assert_eq(psdf.a.apply(identity).sort_index(), pdf.a.apply(identity).sort_index()) + + # TODO: The return type is still category. + # def to_str(x) -> str: + # return str(x) + # + # self.assert_eq( + # psdf.a.apply(to_str).sort_index(), pdf.a.apply(to_str).sort_index() + # ) + + def test_groupby_apply(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.groupby("a").apply(lambda df: df).sort_index(), + pdf.groupby("a").apply(lambda df: df).sort_index(), + ) + self.assert_eq( + psdf.groupby("b").apply(lambda df: df[["a"]]).sort_index(), + pdf.groupby("b").apply(lambda df: df[["a"]]).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", "b"]).apply(lambda df: df).sort_index(), + pdf.groupby(["a", "b"]).apply(lambda df: df).sort_index(), + ) + self.assert_eq( + psdf.groupby("a").apply(lambda df: df.b.cat.codes).sort_index(), + pdf.groupby("a").apply(lambda df: df.b.cat.codes).sort_index(), + ) + self.assert_eq( + psdf.groupby("a")["b"].apply(lambda b: b.cat.codes).sort_index(), + pdf.groupby("a")["b"].apply(lambda b: b.cat.codes).sort_index(), + ) + + # TODO: grouping by a categorical type sometimes preserves unused categories. + # self.assert_eq( + # psdf.groupby("a").apply(len).sort_index(), pdf.groupby("a").apply(len).sort_index(), + # ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43813): Enable CategoricalTests.test_groupby_apply_without_shortcut " + "for pandas 2.0.0.", + ) + def test_groupby_apply_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_groupby_apply() + + pdf, psdf = self.df_pair + + def identity(df) -> ps.DataFrame[zip(psdf.columns, psdf.dtypes)]: + return df + + self.assert_eq( + psdf.groupby("a").apply(identity).sort_values(["a", "b"]).reset_index(drop=True), + pdf.groupby("a").apply(identity).sort_values(["a", "b"]).reset_index(drop=True), + ) + + def test_groupby_transform(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.groupby("a").transform(lambda x: x).sort_index(), + pdf.groupby("a").transform(lambda x: x).sort_index(), + ) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + self.assert_eq( + psdf.groupby("a").transform(lambda x: x.astype(dtype)).sort_index(), + pdf.groupby("a").transform(lambda x: x.astype(dtype)).sort_index(), + ) + + def test_groupby_transform_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_groupby_transform() + + pdf, psdf = self.df_pair + + def identity(x) -> ps.Series[psdf.b.dtype]: + return x + + self.assert_eq( + psdf.groupby("a").transform(identity).sort_values("b").reset_index(drop=True), + pdf.groupby("a").transform(identity).sort_values("b").reset_index(drop=True), + ) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + # The behavior for CategoricalDtype is changed from pandas 1.3 + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + ret_dtype = pdf.b.dtype + else: + ret_dtype = dtype + + def astype(x) -> ps.Series[ret_dtype]: + return x.astype(dtype) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2"): + self.assert_eq( + psdf.groupby("a").transform(astype).sort_values("b").reset_index(drop=True), + pdf.groupby("a").transform(astype).sort_values("b").reset_index(drop=True), + ) + else: + expected = pdf.groupby("a").transform(astype) + expected["b"] = dtype.categories.take(expected["b"].cat.codes).astype(dtype) + self.assert_eq( + psdf.groupby("a").transform(astype).sort_values("b").reset_index(drop=True), + expected.sort_values("b").reset_index(drop=True), + ) + + def test_frame_apply_batch(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda pdf: pdf.astype(str)).sort_index(), + pdf.astype(str).sort_index(), + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(lambda pdf: pdf.astype(dtype)).sort_index(), + pdf.astype(dtype).sort_index(), + ) + + def test_frame_apply_batch_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_frame_apply_batch() + + pdf, psdf = self.df_pair + + def to_str(pdf) -> 'ps.DataFrame["a":str, "b":str]': # noqa: F405 + return pdf.astype(str) + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(to_str).sort_values(["a", "b"]).reset_index(drop=True), + to_str(pdf).sort_values(["a", "b"]).reset_index(drop=True), + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + ret = ps.DataFrame["a":dtype, "b":dtype] + + def to_category(pdf) -> ret: + return pdf.astype(dtype) + + self.assert_eq( + psdf.pandas_on_spark.apply_batch(to_category) + .sort_values(["a", "b"]) + .reset_index(drop=True), + to_category(pdf).sort_values(["a", "b"]).reset_index(drop=True), + ) + + def test_frame_transform_batch(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.astype(str)).sort_index(), + pdf.astype(str).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.b.cat.codes).sort_index(), + pdf.b.cat.codes.sort_index(), + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.astype(dtype)).sort_index(), + pdf.astype(dtype).sort_index(), + ) + self.assert_eq( + psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.b.astype(dtype)).sort_index(), + pdf.b.astype(dtype).sort_index(), + ) + + def test_frame_transform_batch_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_frame_transform_batch() + + pdf, psdf = self.df_pair + + def to_str(pdf) -> 'ps.DataFrame["a":str, "b":str]': # noqa: F405 + return pdf.astype(str) + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(to_str).sort_index(), + to_str(pdf).sort_index(), + ) + + def to_codes(pdf) -> ps.Series[np.int8]: + return pdf.b.cat.codes + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(to_codes).sort_index(), + to_codes(pdf).sort_index(), + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + ret = ps.DataFrame["a":dtype, "b":dtype] + + def to_category(pdf) -> ret: + return pdf.astype(dtype) + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(to_category).sort_index(), + to_category(pdf).sort_index(), + ) + + def to_category(pdf) -> ps.Series[dtype]: + return pdf.b.astype(dtype) + + self.assert_eq( + psdf.pandas_on_spark.transform_batch(to_category).sort_index(), + to_category(pdf).rename().sort_index(), + ) + + def test_series_transform_batch(self): + pdf, psdf = self.df_pair + + self.assert_eq( + psdf.a.pandas_on_spark.transform_batch(lambda pser: pser.astype(str)).sort_index(), + pdf.a.astype(str).sort_index(), + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + self.assert_eq( + psdf.a.pandas_on_spark.transform_batch(lambda pser: pser.astype(dtype)).sort_index(), + pdf.a.astype(dtype).sort_index(), + ) + + def test_series_transform_batch_without_shortcut(self): + with ps.option_context("compute.shortcut_limit", 0): + self.test_series_transform_batch() + + pdf, psdf = self.df_pair + + def to_str(pser) -> ps.Series[str]: + return pser.astype(str) + + self.assert_eq( + psdf.a.pandas_on_spark.transform_batch(to_str).sort_index(), to_str(pdf.a).sort_index() + ) + + pdf = pd.DataFrame( + {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]} + ) + psdf = ps.from_pandas(pdf) + + dtype = CategoricalDtype(categories=["a", "b", "c", "d"]) + + def to_category(pser) -> ps.Series[dtype]: + return pser.astype(dtype) + + self.assert_eq( + psdf.a.pandas_on_spark.transform_batch(to_category).sort_index(), + to_category(pdf.a).sort_index(), + ) + + def test_unstack(self): + pdf = self.pdf + index = pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("x", "c"), ("y", "a"), ("y", "b"), ("y", "d")] + ) + pdf.index = index + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.a.unstack().sort_index(), pdf.a.unstack().sort_index()) + self.assert_eq(psdf.b.unstack().sort_index(), pdf.b.unstack().sort_index()) + + def test_rename_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.b + psser = psdf.b + + self.assert_eq( + pser.cat.rename_categories([0, 1, 3, 2]), psser.cat.rename_categories([0, 1, 3, 2]) + ) + self.assert_eq( + pser.cat.rename_categories({"a": "A", "c": "C"}), + psser.cat.rename_categories({"a": "A", "c": "C"}), + ) + self.assert_eq( + pser.cat.rename_categories(lambda x: x.upper()), + psser.cat.rename_categories(lambda x: x.upper()), + ) + + self.assertRaisesRegex( + ValueError, + "new categories need to have the same number of items as the old categories", + lambda: psser.cat.rename_categories([0, 1, 2]), + ) + self.assertRaises( + TypeError, + lambda: psser.cat.rename_categories(None), + ) + self.assertRaises( + TypeError, + lambda: psser.cat.rename_categories(1), + ) + self.assertRaises( + TypeError, + lambda: psser.cat.rename_categories("x"), + ) + + def test_set_categories(self): + pdf, psdf = self.df_pair + + pser = pdf.b + psser = psdf.b + + self.assert_eq( + pser.cat.set_categories(["a", "c", "b", "o"]), + psser.cat.set_categories(["a", "c", "b", "o"]), + ) + self.assert_eq( + pser.cat.set_categories(["a", "c", "b"]), + psser.cat.set_categories(["a", "c", "b"]), + ) + self.assert_eq( + pser.cat.set_categories(["a", "c", "b", "d", "e"]), + psser.cat.set_categories(["a", "c", "b", "d", "e"]), + ) + + self.assert_eq( + pser.cat.set_categories([0, 1, 3, 2], rename=True), + psser.cat.set_categories([0, 1, 3, 2], rename=True), + ) + self.assert_eq( + pser.cat.set_categories([0, 1, 3], rename=True), + psser.cat.set_categories([0, 1, 3], rename=True), + ) + self.assert_eq( + pser.cat.set_categories([0, 1, 3, 2, 4], rename=True), + psser.cat.set_categories([0, 1, 3, 2, 4], rename=True), + ) + + self.assert_eq( + pser.cat.set_categories(["a", "c", "b", "o"], ordered=True), + psser.cat.set_categories(["a", "c", "b", "o"], ordered=True), + ) + self.assert_eq( + pser.cat.set_categories(["a", "c", "b"], ordered=True), + psser.cat.set_categories(["a", "c", "b"], ordered=True), + ) + self.assert_eq( + pser.cat.set_categories(["a", "c", "b", "d", "e"], ordered=True), + psser.cat.set_categories(["a", "c", "b", "d", "e"], ordered=True), + ) + + self.assertRaisesRegex( + TypeError, + "Parameter 'new_categories' must be list-like, was", + lambda: psser.cat.set_categories(None), + ) + + +class CategoricalTests(CategoricalTestsMixin, ComparisonTestBase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_categorical import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_config.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_config.py new file mode 100644 index 0000000..f61de6e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_config.py @@ -0,0 +1,160 @@ +# +# 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 import pandas as ps +from pyspark.pandas import config +from pyspark.pandas.config import Option, DictWrapper +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class ConfigTestsMixin: + def setUp(self): + config._options_dict["test.config"] = Option(key="test.config", doc="", default="default") + + config._options_dict["test.config.list"] = Option( + key="test.config.list", doc="", default=[], types=list + ) + config._options_dict["test.config.float"] = Option( + key="test.config.float", doc="", default=1.2, types=float + ) + + config._options_dict["test.config.int"] = Option( + key="test.config.int", + doc="", + default=1, + types=int, + check_func=(lambda v: v > 0, "bigger then 0"), + ) + config._options_dict["test.config.int.none"] = Option( + key="test.config.int", doc="", default=None, types=(int, type(None)) + ) + + def tearDown(self): + ps.reset_option("test.config") + del config._options_dict["test.config"] + del config._options_dict["test.config.list"] + del config._options_dict["test.config.float"] + del config._options_dict["test.config.int"] + del config._options_dict["test.config.int.none"] + + def test_get_set_reset_option(self): + self.assertEqual(ps.get_option("test.config"), "default") + + ps.set_option("test.config", "value") + self.assertEqual(ps.get_option("test.config"), "value") + + ps.reset_option("test.config") + self.assertEqual(ps.get_option("test.config"), "default") + + def test_get_set_reset_option_different_types(self): + ps.set_option("test.config.list", [1, 2, 3, 4]) + self.assertEqual(ps.get_option("test.config.list"), [1, 2, 3, 4]) + + ps.set_option("test.config.float", 5.0) + self.assertEqual(ps.get_option("test.config.float"), 5.0) + + ps.set_option("test.config.int", 123) + self.assertEqual(ps.get_option("test.config.int"), 123) + + self.assertEqual(ps.get_option("test.config.int.none"), None) # default None + ps.set_option("test.config.int.none", 123) + self.assertEqual(ps.get_option("test.config.int.none"), 123) + ps.set_option("test.config.int.none", None) + self.assertEqual(ps.get_option("test.config.int.none"), None) + + def test_different_types(self): + with self.assertRaisesRegex(TypeError, "was "): + ps.set_option("test.config.list", 1) + + with self.assertRaisesRegex(TypeError, "however, expected types are"): + ps.set_option("test.config.float", "abc") + + with self.assertRaisesRegex(TypeError, "[]"): + ps.set_option("test.config.int", "abc") + + with self.assertRaisesRegex(TypeError, "(, )"): + ps.set_option("test.config.int.none", "abc") + + def test_check_func(self): + with self.assertRaisesRegex(ValueError, "bigger then 0"): + ps.set_option("test.config.int", -1) + + def test_unknown_option(self): + with self.assertRaisesRegex(config.OptionError, "No such option"): + ps.get_option("unknown") + + with self.assertRaisesRegex(config.OptionError, "Available options"): + ps.set_option("unknown", "value") + + with self.assertRaisesRegex(config.OptionError, "test.config"): + ps.reset_option("unknown") + + def test_namespace_access(self): + try: + self.assertEqual(ps.options.compute.max_rows, ps.get_option("compute.max_rows")) + ps.options.compute.max_rows = 0 + self.assertEqual(ps.options.compute.max_rows, 0) + self.assertTrue(isinstance(ps.options.compute, DictWrapper)) + + wrapper = ps.options.compute + self.assertEqual(wrapper.max_rows, ps.get_option("compute.max_rows")) + wrapper.max_rows = 1000 + self.assertEqual(ps.options.compute.max_rows, 1000) + + self.assertRaisesRegex(config.OptionError, "No such option", lambda: ps.options.compu) + self.assertRaisesRegex( + config.OptionError, "No such option", lambda: ps.options.compute.max + ) + self.assertRaisesRegex( + config.OptionError, "No such option", lambda: ps.options.max_rows1 + ) + + with self.assertRaisesRegex(config.OptionError, "No such option"): + ps.options.compute.max = 0 + with self.assertRaisesRegex(config.OptionError, "No such option"): + ps.options.compute = 0 + with self.assertRaisesRegex(config.OptionError, "No such option"): + ps.options.com = 0 + finally: + ps.reset_option("compute.max_rows") + + def test_dir_options(self): + self.assertTrue("compute.default_index_type" in dir(ps.options)) + self.assertTrue("plotting.sample_ratio" in dir(ps.options)) + + self.assertTrue("default_index_type" in dir(ps.options.compute)) + self.assertTrue("sample_ratio" not in dir(ps.options.compute)) + + self.assertTrue("default_index_type" not in dir(ps.options.plotting)) + self.assertTrue("sample_ratio" in dir(ps.options.plotting)) + + +class ConfigTests(ConfigTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_config import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_csv.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_csv.py new file mode 100644 index 0000000..b118f7c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_csv.py @@ -0,0 +1,453 @@ +# +# 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 os +import shutil +import tempfile +import unittest +from contextlib import contextmanager +from distutils.version import LooseVersion + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +def normalize_text(s): + return "\n".join(map(str.strip, s.strip().split("\n"))) + + +class CsvTestsMixin: + def setUp(self): + self.tmp_dir = tempfile.mkdtemp(prefix=CsvTests.__name__) + + def tearDown(self): + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + @property + def csv_text(self): + return normalize_text( + """ + name,amount + Alice,100 + Bob,-200 + Charlie,300 + Dennis,400 + Edith,-500 + Frank,600 + Alice,200 + Frank,-200 + Bob,600 + Alice,400 + Frank,200 + Alice,300 + Edith,600 + """ + ) + + @property + def csv_text_2(self): + return normalize_text( + """ + A,B + item1,1 + item2,1,2 + item3,1,2,3,4 + item4,1 + """ + ) + + @property + def csv_text_with_comments(self): + return normalize_text( + """ + # header + %s + # comment + Alice,400 + Edith,600 + # footer + """ + % self.csv_text + ) + + @property + def tab_delimited_csv_text(self): + return normalize_text( + """ + name\tamount + Alice\t100 + Bob\t-200 + Charlie\t300 + """ + ) + + @property + def q_quoted_csv_text(self): + return normalize_text( + """ + QnameQ,QamountQ + QA,liceQ,Q100Q + QB,obQ,Q-200Q + QC,harlieQ,Q300Q + """ + ) + + @property + def e_escapeted_csv_text(self): + return normalize_text( + """ + name,amount + "AE"lice",100 + "BE"ob",-200 + "CE"harlie",300 + """ + ) + + @contextmanager + def csv_file(self, csv): + with self.temp_file() as tmp: + with open(tmp, "w") as f: + f.write(csv) + yield tmp + + def test_read_csv(self): + with self.csv_file(self.csv_text) as fn: + + def check(header="infer", names=None, usecols=None, index_col=None): + expected = pd.read_csv( + fn, header=header, names=names, usecols=usecols, index_col=index_col + ) + actual = ps.read_csv( + fn, header=header, names=names, usecols=usecols, index_col=index_col + ) + self.assert_eq(expected, actual, almost=True) + + check() + check(header=0) + check(header=None) + check(names=["n", "a"]) + check(names=[("x", "n"), ("y", "a")]) + check(names=[10, 20]) + check(header=0, names=["n", "a"]) + check(usecols=[1]) + check(usecols=[1, 0]) + check(usecols=["amount"]) + check(usecols=["amount", "name"]) + check(usecols=[]) + check(usecols=[1, 1]) + check(usecols=["amount", "amount"]) + check(header=None, usecols=[1]) + check(names=["n", "a"], usecols=["a"]) + check(header=None, names=["n", "a"], usecols=["a"]) + check(index_col=["amount"]) + check(header=None, index_col=[1]) + check(names=["n", "a"], index_col=["a"]) + check(names=["n", "a"], index_col="a") + + # check with pyspark patch. + expected = pd.read_csv(fn) + actual = ps.read_csv(fn) + self.assert_eq(expected, actual, almost=True) + + self.assertRaisesRegex( + ValueError, "non-unique", lambda: ps.read_csv(fn, names=["n", "n"]) + ) + self.assertRaisesRegex( + ValueError, + "does not match the number.*3", + lambda: ps.read_csv(fn, names=["n", "a", "b"]), + ) + self.assertRaisesRegex( + ValueError, + "does not match the number.*3", + lambda: ps.read_csv(fn, header=0, names=["n", "a", "b"]), + ) + self.assertRaisesRegex( + ValueError, "Usecols do not match.*3", lambda: ps.read_csv(fn, usecols=[1, 3]) + ) + self.assertRaisesRegex( + ValueError, + "Usecols do not match.*col", + lambda: ps.read_csv(fn, usecols=["amount", "col"]), + ) + self.assertRaisesRegex( + ValueError, "Unknown header argument 1", lambda: ps.read_csv(fn, header="1") + ) + expected_error_message = ( + "'usecols' must either be list-like of all strings, " + "all unicode, all integers or a callable." + ) + self.assertRaisesRegex( + ValueError, expected_error_message, lambda: ps.read_csv(fn, usecols=[1, "amount"]) + ) + + # check with index_col + expected = pd.read_csv(fn).set_index("name") + actual = ps.read_csv(fn, index_col="name") + self.assert_eq(expected, actual, almost=True) + + def test_read_with_spark_schema(self): + with self.csv_file(self.csv_text_2) as fn: + actual = ps.read_csv(fn, names="A string, B string, C long, D long, E long") + expected = pd.read_csv(fn, names=["A", "B", "C", "D", "E"]) + self.assert_eq(expected, actual) + + def test_read_csv_with_comment(self): + with self.csv_file(self.csv_text_with_comments) as fn: + expected = pd.read_csv(fn, comment="#") + actual = ps.read_csv(fn, comment="#") + self.assert_eq(expected, actual, almost=True) + + self.assertRaisesRegex( + ValueError, + "Only length-1 comment characters supported", + lambda: ps.read_csv(fn, comment="").show(), + ) + self.assertRaisesRegex( + ValueError, + "Only length-1 comment characters supported", + lambda: ps.read_csv(fn, comment="##").show(), + ) + self.assertRaisesRegex( + ValueError, + "Only length-1 comment characters supported", + lambda: ps.read_csv(fn, comment=1), + ) + self.assertRaisesRegex( + ValueError, + "Only length-1 comment characters supported", + lambda: ps.read_csv(fn, comment=[1]), + ) + + def test_read_csv_with_limit(self): + with self.csv_file(self.csv_text_with_comments) as fn: + expected = pd.read_csv(fn, comment="#", nrows=2) + actual = ps.read_csv(fn, comment="#", nrows=2) + self.assert_eq(expected, actual, almost=True) + + def test_read_csv_with_encoding(self): + # SPARK-37181: Read csv supporting latin-1 encoding. + with self.csv_file(self.csv_text) as fn: + expected = pd.read_csv(fn, encoding="latin-1") + actual = ps.read_csv(fn, encoding="latin-1") + self.assert_eq(expected, actual, almost=True) + + def test_read_csv_with_sep(self): + with self.csv_file(self.tab_delimited_csv_text) as fn: + expected = pd.read_csv(fn, sep="\t") + actual = ps.read_csv(fn, sep="\t") + self.assert_eq(expected, actual, almost=True) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43563): Enable CsvTests.test_read_csv_with_squeeze for pandas 2.0.0.", + ) + def test_read_csv_with_squeeze(self): + with self.csv_file(self.csv_text) as fn: + expected = pd.read_csv(fn, squeeze=True, usecols=["name"]) + actual = ps.read_csv(fn, squeeze=True, usecols=["name"]) + self.assert_eq(expected, actual, almost=True) + + expected = pd.read_csv(fn, squeeze=True, usecols=["name", "amount"]) + actual = ps.read_csv(fn, squeeze=True, usecols=["name", "amount"]) + self.assert_eq(expected, actual, almost=True) + + expected = pd.read_csv(fn, squeeze=True, usecols=["name", "amount"], index_col=["name"]) + actual = ps.read_csv(fn, squeeze=True, usecols=["name", "amount"], index_col=["name"]) + self.assert_eq(expected, actual, almost=True) + + def test_read_csv_with_mangle_dupe_cols(self): + self.assertRaisesRegex( + ValueError, "mangle_dupe_cols", lambda: ps.read_csv("path", mangle_dupe_cols=False) + ) + + def test_read_csv_with_parse_dates(self): + self.assertRaisesRegex( + ValueError, "parse_dates", lambda: ps.read_csv("path", parse_dates=True) + ) + + def test_read_csv_with_dtype(self): + with self.csv_file(self.csv_text) as fn: + self.assert_eq(ps.read_csv(fn), pd.read_csv(fn), almost=True) + self.assert_eq(ps.read_csv(fn, dtype=str), pd.read_csv(fn, dtype=str)) + self.assert_eq( + ps.read_csv(fn, dtype={"amount": "int64"}), + pd.read_csv(fn, dtype={"amount": "int64"}), + ) + + def test_read_csv_with_quotechar(self): + with self.csv_file(self.q_quoted_csv_text) as fn: + self.assert_eq( + ps.read_csv(fn, quotechar="Q"), pd.read_csv(fn, quotechar="Q"), almost=True + ) + + def test_read_csv_with_escapechar(self): + with self.csv_file(self.e_escapeted_csv_text) as fn: + self.assert_eq( + ps.read_csv(fn, escapechar="E"), pd.read_csv(fn, escapechar="E"), almost=True + ) + + self.assert_eq( + ps.read_csv(fn, escapechar="ABC", escape="E"), + pd.read_csv(fn, escapechar="E"), + almost=True, + ) + + def test_to_csv(self): + pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": [4, 5, 6]}, index=[0, 1, 3]) + psdf = ps.DataFrame(pdf) + + self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) + self.assert_eq(psdf.to_csv(columns=["aa"]), pdf.to_csv(columns=["aa"], index=False)) + self.assert_eq(psdf.aa.to_csv(), pdf.aa.to_csv(index=False, header=True)) + + pdf = pd.DataFrame({"a": [1, np.nan, 3], "b": ["one", "two", None]}, index=[0, 1, 3]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.to_csv(na_rep="null"), pdf.to_csv(na_rep="null", index=False)) + self.assert_eq( + psdf.a.to_csv(na_rep="null"), pdf.a.to_csv(na_rep="null", index=False, header=True) + ) + + self.assertRaises(KeyError, lambda: psdf.to_csv(columns=["ab"])) + + pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}, index=[0, 1, 3]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) + self.assert_eq(psdf.to_csv(header=False), pdf.to_csv(header=False, index=False)) + self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) + + # non-string names + pdf = pd.DataFrame({10: [1, 2, 3], 20: [4, 5, 6]}, index=[0, 1, 3]) + psdf = ps.DataFrame(pdf) + + self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) + self.assert_eq(psdf.to_csv(columns=[10]), pdf.to_csv(columns=[10], index=False)) + + self.assertRaises(TypeError, lambda: psdf.to_csv(columns=10)) + + def _check_output(self, dir, expected): + output_paths = [path for path in os.listdir(dir) if path.startswith("part-")] + assert len(output_paths) > 0 + output_path = "%s/%s" % (dir, output_paths[0]) + with open(output_path) as f: + self.assertEqual(f.read(), expected) + + def test_to_csv_with_path(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + tmp_dir = "{}/tmp1".format(self.tmp_dir) + + psdf.to_csv(tmp_dir, num_files=1) + self._check_output(tmp_dir, pdf.to_csv(index=False)) + + tmp_dir = "{}/tmp2".format(self.tmp_dir) + + self.assertRaises(KeyError, lambda: psdf.to_csv(tmp_dir, columns=["c"], num_files=1)) + + # non-string names + pdf = pd.DataFrame({10: [1, 2, 3], 20: ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + tmp_dir = "{}/tmp3".format(self.tmp_dir) + + psdf.to_csv(tmp_dir, num_files=1) + self._check_output(tmp_dir, pdf.to_csv(index=False)) + + tmp_dir = "{}/tmp4".format(self.tmp_dir) + + psdf.to_csv(tmp_dir, columns=[10], num_files=1) + self._check_output(tmp_dir, pdf.to_csv(columns=[10], index=False)) + + tmp_dir = "{}/tmp5".format(self.tmp_dir) + + self.assertRaises(TypeError, lambda: psdf.to_csv(tmp_dir, columns=10, num_files=1)) + + def test_to_csv_with_path_and_basic_options(self): + pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + psdf.to_csv(self.tmp_dir, num_files=1, sep="|", header=False, columns=["aa"]) + expected = pdf.to_csv(index=False, sep="|", header=False, columns=["aa"]) + + self._check_output(self.tmp_dir, expected) + + def test_to_csv_with_path_and_basic_options_multiindex_columns(self): + pdf = pd.DataFrame({("x", "a"): [1, 2, 3], ("y", "b"): ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + with self.assertRaises(ValueError): + psdf.to_csv(self.tmp_dir, num_files=1, sep="|", columns=[("x", "a")]) + + psdf.to_csv(self.tmp_dir, num_files=1, sep="|", header=["a"], columns=[("x", "a")]) + pdf.columns = ["a", "b"] + expected = pdf.to_csv(index=False, sep="|", columns=["a"]) + + self._check_output(self.tmp_dir, expected) + + def test_to_csv_with_path_and_pyspark_options(self): + pdf = pd.DataFrame({"a": [1, 2, 3, None], "b": ["a", "b", "c", None]}) + psdf = ps.DataFrame(pdf) + + psdf.to_csv(self.tmp_dir, nullValue="null", num_files=1) + expected = pdf.to_csv(index=False, na_rep="null") + + self._check_output(self.tmp_dir, expected) + + def test_to_csv_with_partition_cols(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + psdf.to_csv(self.tmp_dir, partition_cols="b", num_files=1) + + partition_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("b=")] + assert len(partition_paths) > 0 + for partition_path in partition_paths: + column, value = partition_path.split("=") + expected = pdf[pdf[column] == value].drop("b", axis=1).to_csv(index=False) + + output_paths = [ + path + for path in os.listdir("%s/%s" % (self.tmp_dir, partition_path)) + if path.startswith("part-") + ] + assert len(output_paths) > 0 + output_path = "%s/%s/%s" % (self.tmp_dir, partition_path, output_paths[0]) + with open(output_path) as f: + self.assertEqual(f.read(), expected) + + +class CsvTests(CsvTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_csv import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_dataframe_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_dataframe_conversion.py new file mode 100644 index 0000000..d245b44 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_dataframe_conversion.py @@ -0,0 +1,281 @@ +# +# 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 os +import shutil +import string +import tempfile +import unittest +import sys +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import ComparisonTestBase, TestUtils +from pyspark.testing.sqlutils import SQLTestUtils + + +class DataFrameConversionTestsMixin: + """Test cases for "small data" conversion and I/O.""" + + def setUp(self): + self.tmp_dir = tempfile.mkdtemp(prefix=DataFrameConversionTests.__name__) + + def tearDown(self): + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + @property + def pdf(self): + return pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]) + + @staticmethod + def strip_all_whitespace(str): + """A helper function to remove all whitespace from a string.""" + return str.translate({ord(c): None for c in string.whitespace}) + + def test_to_html(self): + expected = self.strip_all_whitespace( + """ +
    + + + + + + + + +
    ab
    014
    125
    336
    + """ + ) + got = self.strip_all_whitespace(self.psdf.to_html()) + self.assert_eq(got, expected) + + # with max_rows set + expected = self.strip_all_whitespace( + """ + + + + + + + + +
    ab
    014
    125
    + """ + ) + got = self.strip_all_whitespace(self.psdf.to_html(max_rows=2)) + self.assert_eq(got, expected) + + @staticmethod + def get_excel_dfs(pandas_on_spark_location, pandas_location): + return { + "got": pd.read_excel(pandas_on_spark_location, index_col=0), + "expected": pd.read_excel(pandas_location, index_col=0), + } + + def test_to_excel(self): + with self.temp_dir() as dirpath: + pandas_location = dirpath + "/" + "output1.xlsx" + pandas_on_spark_location = dirpath + "/" + "output2.xlsx" + + pdf = self.pdf + psdf = self.psdf + psdf.to_excel(pandas_on_spark_location) + pdf.to_excel(pandas_location) + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + psdf.a.to_excel(pandas_on_spark_location) + pdf.a.to_excel(pandas_location) + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + pdf = pd.DataFrame({"a": [1, None, 3], "b": ["one", "two", None]}, index=[0, 1, 3]) + + psdf = ps.from_pandas(pdf) + + psdf.to_excel(pandas_on_spark_location, na_rep="null") + pdf.to_excel(pandas_location, na_rep="null") + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}, index=[0, 1, 3]) + + psdf = ps.from_pandas(pdf) + + psdf.to_excel(pandas_on_spark_location, float_format="%.1f") + pdf.to_excel(pandas_location, float_format="%.1f") + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + psdf.to_excel(pandas_on_spark_location, header=False) + pdf.to_excel(pandas_location, header=False) + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + psdf.to_excel(pandas_on_spark_location, index=False) + pdf.to_excel(pandas_location, index=False) + dataframes = self.get_excel_dfs(pandas_on_spark_location, pandas_location) + self.assert_eq(dataframes["got"], dataframes["expected"]) + + def test_to_json(self): + pdf = self.pdf + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.to_json(orient="records"), pdf.to_json(orient="records")) + + def test_to_json_negative(self): + psdf = ps.from_pandas(self.pdf) + + with self.assertRaises(NotImplementedError): + psdf.to_json(orient="table") + + with self.assertRaises(NotImplementedError): + psdf.to_json(lines=False) + + def test_read_json_negative(self): + with self.assertRaises(NotImplementedError): + ps.read_json("invalid", lines=False) + + def test_to_json_with_path(self): + pdf = pd.DataFrame({"a": [1], "b": ["a"]}) + psdf = ps.DataFrame(pdf) + + psdf.to_json(self.tmp_dir, num_files=1) + expected = pdf.to_json(orient="records") + + output_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("part-")] + assert len(output_paths) > 0 + output_path = "%s/%s" % (self.tmp_dir, output_paths[0]) + self.assertEqual("[%s]" % open(output_path).read().strip(), expected) + + def test_to_json_with_partition_cols(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + psdf = ps.DataFrame(pdf) + + psdf.to_json(self.tmp_dir, partition_cols="b", num_files=1) + + partition_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("b=")] + assert len(partition_paths) > 0 + for partition_path in partition_paths: + column, value = partition_path.split("=") + expected = pdf[pdf[column] == value].drop("b", axis=1).to_json(orient="records") + + output_paths = [ + path + for path in os.listdir("%s/%s" % (self.tmp_dir, partition_path)) + if path.startswith("part-") + ] + assert len(output_paths) > 0 + output_path = "%s/%s/%s" % (self.tmp_dir, partition_path, output_paths[0]) + self.assertEqual("[%s]" % open(output_path).read().strip(), expected) + + @unittest.skipIf( + sys.platform == "linux" or sys.platform == "linux2", + "Pyperclip could not find a copy/paste mechanism for Linux.", + ) + def test_to_clipboard(self): + pdf = self.pdf + psdf = self.psdf + + self.assert_eq(psdf.to_clipboard(), pdf.to_clipboard()) + self.assert_eq(psdf.to_clipboard(excel=False), pdf.to_clipboard(excel=False)) + self.assert_eq( + psdf.to_clipboard(sep=";", index=False), pdf.to_clipboard(sep=";", index=False) + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43561): Enable DataFrameConversionTests.test_to_latex for pandas 2.0.0.", + ) + def test_to_latex(self): + pdf = self.pdf + psdf = self.psdf + + self.assert_eq(psdf.to_latex(), pdf.to_latex()) + self.assert_eq(psdf.to_latex(col_space=2), pdf.to_latex(col_space=2)) + self.assert_eq(psdf.to_latex(header=True), pdf.to_latex(header=True)) + self.assert_eq(psdf.to_latex(index=False), pdf.to_latex(index=False)) + self.assert_eq(psdf.to_latex(na_rep="-"), pdf.to_latex(na_rep="-")) + self.assert_eq(psdf.to_latex(float_format="%.1f"), pdf.to_latex(float_format="%.1f")) + self.assert_eq(psdf.to_latex(sparsify=False), pdf.to_latex(sparsify=False)) + self.assert_eq(psdf.to_latex(index_names=False), pdf.to_latex(index_names=False)) + self.assert_eq(psdf.to_latex(bold_rows=True), pdf.to_latex(bold_rows=True)) + self.assert_eq(psdf.to_latex(decimal=","), pdf.to_latex(decimal=",")) + + def test_to_records(self): + pdf = pd.DataFrame({"A": [1, 2], "B": [0.5, 0.75]}, index=["a", "b"]) + + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.to_records(), pdf.to_records()) + self.assert_eq(psdf.to_records(index=False), pdf.to_records(index=False)) + self.assert_eq(psdf.to_records(index_dtypes=" LooseVersion("1.5.0"): + expected_psdf = ps.read_parquet(path2, pandas_metadata=True).set_index( + "index", append=True + ) + else: + path3 = "{}/file3.parquet".format(tmp) + expected3.to_parquet(path3) + expected_psdf = ps.read_parquet(path3, pandas_metadata=True) + + self.assert_eq(expected_psdf, expected3) + + def test_parquet_write(self): + with self.temp_dir() as tmp: + pdf = self.test_pdf + expected = ps.DataFrame(pdf) + + # Write out partitioned by one column + expected.to_parquet(tmp, mode="overwrite", partition_cols="i32") + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_parquet(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Write out partitioned by two columns + expected.to_parquet(tmp, mode="overwrite", partition_cols=["i32", "bhello"]) + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_parquet(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Set `compression` with string + expected.to_parquet(tmp, mode="overwrite", partition_cols="i32", compression="none") + actual = ps.read_parquet(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Test `options` parameter + expected.to_parquet( + tmp, mode="overwrite", partition_cols="i32", options={"compression": "none"} + ) + actual = ps.read_parquet(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + def test_table(self): + with self.table("test_table"): + pdf = self.test_pdf + expected = ps.DataFrame(pdf) + + # Write out partitioned by one column + expected.spark.to_table("test_table", mode="overwrite", partition_cols="i32") + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_table("test_table") + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Write out partitioned by two columns + expected.to_table("test_table", mode="overwrite", partition_cols=["i32", "bhello"]) + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_table("test_table") + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # When index columns are known + expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]] + actual_idx = ps.read_table("test_table", index_col="bhello")[["f", "i32", "i64"]] + self.assert_eq( + actual_idx.sort_values(by="f").to_spark().toPandas(), + expected_idx.sort_values(by="f").to_spark().toPandas(), + ) + + expected_idx = expected.set_index(["bhello"])[["f", "i32", "i64"]] + actual_idx = ps.read_table("test_table", index_col=["bhello"])[["f", "i32", "i64"]] + self.assert_eq( + actual_idx.sort_values(by="f").to_spark().toPandas(), + expected_idx.sort_values(by="f").to_spark().toPandas(), + ) + + expected_idx = expected.set_index(["i32", "bhello"])[["f", "i64"]] + actual_idx = ps.read_table("test_table", index_col=["i32", "bhello"])[["f", "i64"]] + self.assert_eq( + actual_idx.sort_values(by="f").to_spark().toPandas(), + expected_idx.sort_values(by="f").to_spark().toPandas(), + ) + + def test_spark_io(self): + with self.temp_dir() as tmp: + pdf = self.test_pdf + expected = ps.DataFrame(pdf) + + # Write out partitioned by one column + expected.to_spark_io(tmp, format="json", mode="overwrite", partition_cols="i32") + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_spark_io(tmp, format="json") + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Write out partitioned by two columns + expected.to_spark_io( + tmp, format="json", mode="overwrite", partition_cols=["i32", "bhello"] + ) + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_spark_io(path=tmp, format="json") + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # When index columns are known + pdf = self.test_pdf + expected = ps.DataFrame(pdf) + col_order = ["f", "i32", "i64"] + + expected_idx = expected.set_index("bhello")[col_order] + actual_idx = ps.read_spark_io(tmp, format="json", index_col="bhello")[col_order] + self.assert_eq( + actual_idx.sort_values(by="f").to_spark().toPandas(), + expected_idx.sort_values(by="f").to_spark().toPandas(), + ) + + # TODO(SPARK-40353): re-enabling the `test_read_excel`. + @unittest.skip("openpyxl") + def test_read_excel(self): + with self.temp_dir() as tmp: + + path1 = "{}/file1.xlsx".format(tmp) + self.test_pdf[["i32"]].to_excel(path1) + + self.assert_eq(ps.read_excel(open(path1, "rb")), pd.read_excel(open(path1, "rb"))) + self.assert_eq( + ps.read_excel(open(path1, "rb"), index_col=0), + pd.read_excel(open(path1, "rb"), index_col=0), + ) + self.assert_eq( + ps.read_excel(open(path1, "rb"), index_col=0, squeeze=True), + pd.read_excel(open(path1, "rb"), index_col=0, squeeze=True), + ) + + self.assert_eq(ps.read_excel(path1), pd.read_excel(path1)) + self.assert_eq(ps.read_excel(path1, index_col=0), pd.read_excel(path1, index_col=0)) + self.assert_eq( + ps.read_excel(path1, index_col=0, squeeze=True), + pd.read_excel(path1, index_col=0, squeeze=True), + ) + + self.assert_eq(ps.read_excel(tmp), pd.read_excel(path1)) + + path2 = "{}/file2.xlsx".format(tmp) + self.test_pdf[["i32"]].to_excel(path2) + self.assert_eq( + ps.read_excel(tmp, index_col=0).sort_index(), + pd.concat( + [pd.read_excel(path1, index_col=0), pd.read_excel(path2, index_col=0)] + ).sort_index(), + ) + self.assert_eq( + ps.read_excel(tmp, index_col=0, squeeze=True).sort_index(), + pd.concat( + [ + pd.read_excel(path1, index_col=0, squeeze=True), + pd.read_excel(path2, index_col=0, squeeze=True), + ] + ).sort_index(), + ) + + with self.temp_dir() as tmp: + path1 = "{}/file1.xlsx".format(tmp) + with pd.ExcelWriter(path1) as writer: + self.test_pdf.to_excel(writer, sheet_name="Sheet_name_1") + self.test_pdf[["i32"]].to_excel(writer, sheet_name="Sheet_name_2") + + sheet_names = [["Sheet_name_1", "Sheet_name_2"], None] + + pdfs1 = pd.read_excel(open(path1, "rb"), sheet_name=None, index_col=0) + pdfs1_squeezed = pd.read_excel( + open(path1, "rb"), sheet_name=None, index_col=0, squeeze=True + ) + + for sheet_name in sheet_names: + psdfs = ps.read_excel(open(path1, "rb"), sheet_name=sheet_name, index_col=0) + self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"]) + self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"]) + + psdfs = ps.read_excel( + open(path1, "rb"), sheet_name=sheet_name, index_col=0, squeeze=True + ) + self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"]) + self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"]) + + self.assert_eq( + ps.read_excel(tmp, index_col=0, sheet_name="Sheet_name_2"), + pdfs1["Sheet_name_2"], + ) + + for sheet_name in sheet_names: + psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0) + self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"]) + self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"]) + + psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True) + self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"]) + self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"]) + + path2 = "{}/file2.xlsx".format(tmp) + with pd.ExcelWriter(path2) as writer: + self.test_pdf.to_excel(writer, sheet_name="Sheet_name_1") + self.test_pdf[["i32"]].to_excel(writer, sheet_name="Sheet_name_2") + + pdfs2 = pd.read_excel(path2, sheet_name=None, index_col=0) + pdfs2_squeezed = pd.read_excel(path2, sheet_name=None, index_col=0, squeeze=True) + + self.assert_eq( + ps.read_excel(tmp, sheet_name="Sheet_name_2", index_col=0).sort_index(), + pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(), + ) + self.assert_eq( + ps.read_excel( + tmp, sheet_name="Sheet_name_2", index_col=0, squeeze=True + ).sort_index(), + pd.concat( + [pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]] + ).sort_index(), + ) + + for sheet_name in sheet_names: + psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0) + self.assert_eq( + psdfs["Sheet_name_1"].sort_index(), + pd.concat([pdfs1["Sheet_name_1"], pdfs2["Sheet_name_1"]]).sort_index(), + ) + self.assert_eq( + psdfs["Sheet_name_2"].sort_index(), + pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(), + ) + + psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True) + self.assert_eq( + psdfs["Sheet_name_1"].sort_index(), + pd.concat( + [pdfs1_squeezed["Sheet_name_1"], pdfs2_squeezed["Sheet_name_1"]] + ).sort_index(), + ) + self.assert_eq( + psdfs["Sheet_name_2"].sort_index(), + pd.concat( + [pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]] + ).sort_index(), + ) + + def test_read_orc(self): + with self.temp_dir() as tmp: + path = "{}/file1.orc".format(tmp) + data = self.test_pdf + self.spark.createDataFrame(data, "i32 int, i64 long, f double, bhello string").coalesce( + 1 + ).write.orc(path, mode="overwrite") + + expected = data.reset_index()[data.columns] + actual = ps.read_orc(path) + self.assertPandasEqual(expected, actual._to_pandas()) + + # columns + columns = ["i32", "i64"] + expected = data.reset_index()[columns] + actual = ps.read_orc(path, columns=columns) + self.assertPandasEqual(expected, actual._to_pandas()) + + # index_col + expected = data.set_index("i32") + actual = ps.read_orc(path, index_col="i32") + self.assert_eq(actual, expected) + + expected = data.set_index(["i32", "f"]) + actual = ps.read_orc(path, index_col=["i32", "f"]) + self.assert_eq(actual, expected) + + # index_col with columns + expected = data.set_index("i32")[["i64", "bhello"]] + actual = ps.read_orc(path, index_col=["i32"], columns=["i64", "bhello"]) + self.assert_eq(actual, expected) + + expected = data.set_index(["i32", "f"])[["bhello", "i64"]] + actual = ps.read_orc(path, index_col=["i32", "f"], columns=["bhello", "i64"]) + self.assert_eq(actual, expected) + + msg = "Unknown column name 'i'" + with self.assertRaises(ValueError, msg=msg): + ps.read_orc(path, columns="i32") + msg = "Unknown column name 'i34'" + with self.assertRaises(ValueError, msg=msg): + ps.read_orc(path, columns=["i34", "i64"]) + + def test_orc_write(self): + with self.temp_dir() as tmp: + pdf = self.test_pdf + expected = ps.DataFrame(pdf) + + # Write out partitioned by one column + expected.to_orc(tmp, mode="overwrite", partition_cols="i32") + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_orc(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Write out partitioned by two columns + expected.to_orc(tmp, mode="overwrite", partition_cols=["i32", "bhello"]) + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_orc(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + # Test `options` parameter + expected.to_orc( + tmp, mode="overwrite", partition_cols="i32", options={"compression": "none"} + ) + # Reset column order, as once the data is written out, Spark rearranges partition + # columns to appear first. + actual = ps.read_orc(tmp) + self.assertFalse((actual.columns == self.test_column_order).all()) + actual = actual[self.test_column_order] + self.assert_eq( + actual.sort_values(by="f").to_spark().toPandas(), + expected.sort_values(by="f").to_spark().toPandas(), + ) + + +class DataFrameSparkIOTests(DataFrameSparkIOTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_dataframe_spark_io import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_default_index.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_default_index.py new file mode 100644 index 0000000..45ceaf5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_default_index.py @@ -0,0 +1,109 @@ +# +# 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 pandas as pd + +from pyspark.sql import functions as F +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class DefaultIndexTestsMixin: + def test_default_index_sequence(self): + with ps.option_context("compute.default_index_type", "sequence"): + sdf = self.spark.range(1000) + self.assert_eq(ps.DataFrame(sdf), pd.DataFrame({"id": list(range(1000))})) + + def test_default_index_distributed_sequence(self): + with ps.option_context("compute.default_index_type", "distributed-sequence"): + sdf = self.spark.range(1000) + self.assert_eq(ps.DataFrame(sdf), pd.DataFrame({"id": list(range(1000))})) + + def test_default_index_distributed(self): + with ps.option_context("compute.default_index_type", "distributed"): + sdf = self.spark.range(1000) + pdf = ps.DataFrame(sdf)._to_pandas() + self.assertEqual(len(set(pdf.index)), len(pdf)) + + def test_index_distributed_sequence_cleanup(self): + with ps.option_context( + "compute.default_index_type", "distributed-sequence" + ), ps.option_context("compute.ops_on_diff_frames", True): + + with ps.option_context("compute.default_index_cache", "LOCAL_CHECKPOINT"): + cached_rdd_ids = [rdd_id for rdd_id in self.spark._jsc.getPersistentRDDs()] + + psdf1 = ( + self.spark.range(0, 100, 1, 10).withColumn("Key", F.col("id") % 33).pandas_api() + ) + + psdf2 = psdf1["Key"].reset_index() + psdf2["index"] = (psdf2.groupby(["Key"]).cumcount() == 0).astype(int) + psdf2["index"] = psdf2["index"].cumsum() + + psdf3 = ps.merge(psdf1, psdf2, how="inner", left_on=["Key"], right_on=["Key"]) + _ = len(psdf3) + + # newly cached rdd + self.assertTrue( + any( + rdd_id not in cached_rdd_ids + for rdd_id in self.spark._jsc.getPersistentRDDs() + ) + ) + + for storage_level in ["NONE", "DISK_ONLY_2", "MEMORY_AND_DISK_SER"]: + with ps.option_context("compute.default_index_cache", storage_level): + cached_rdd_ids = [rdd_id for rdd_id in self.spark._jsc.getPersistentRDDs()] + + psdf1 = ( + self.spark.range(0, 100, 1, 10) + .withColumn("Key", F.col("id") % 33) + .pandas_api() + ) + + psdf2 = psdf1["Key"].reset_index() + psdf2["index"] = (psdf2.groupby(["Key"]).cumcount() == 0).astype(int) + psdf2["index"] = psdf2["index"].cumsum() + + psdf3 = ps.merge(psdf1, psdf2, how="inner", left_on=["Key"], right_on=["Key"]) + _ = len(psdf3) + + # no newly cached rdd + self.assertTrue( + all( + rdd_id in cached_rdd_ids + for rdd_id in self.spark._jsc.getPersistentRDDs() + ) + ) + + +class DefaultIndexTests(DefaultIndexTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_default_index import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ewm.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ewm.py new file mode 100644 index 0000000..a8886a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ewm.py @@ -0,0 +1,434 @@ +# +# 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 numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils +from pyspark.pandas.window import ExponentialMoving + + +class EWMTestsMixin: + def test_ewm_error(self): + with self.assertRaisesRegex( + TypeError, "psdf_or_psser must be a series or dataframe; however, got:.*int" + ): + ExponentialMoving(1, 2) + + psdf = ps.range(10) + + with self.assertRaisesRegex(ValueError, "min_periods must be >= 0"): + psdf.ewm(min_periods=-1, alpha=0.5).mean() + + with self.assertRaisesRegex(ValueError, "com must be >= 0"): + psdf.ewm(com=-0.1).mean() + + with self.assertRaisesRegex(ValueError, "span must be >= 1"): + psdf.ewm(span=0.7).mean() + + with self.assertRaisesRegex(ValueError, "halflife must be > 0"): + psdf.ewm(halflife=0).mean() + + with self.assertRaisesRegex(ValueError, "alpha must be in"): + psdf.ewm(alpha=1.7).mean() + + with self.assertRaisesRegex(ValueError, "Must pass one of com, span, halflife, or alpha"): + psdf.ewm().mean() + + with self.assertRaisesRegex( + ValueError, "com, span, halflife, and alpha are mutually exclusive" + ): + psdf.ewm(com=0.5, alpha=0.7).mean() + + with self.assertRaisesRegex(ValueError, "min_periods must be >= 0"): + psdf.groupby(psdf.id).ewm(min_periods=-1, alpha=0.5).mean() + + with self.assertRaisesRegex(ValueError, "com must be >= 0"): + psdf.groupby(psdf.id).ewm(com=-0.1).mean() + + with self.assertRaisesRegex(ValueError, "span must be >= 1"): + psdf.groupby(psdf.id).ewm(span=0.7).mean() + + with self.assertRaisesRegex(ValueError, "halflife must be > 0"): + psdf.groupby(psdf.id).ewm(halflife=0).mean() + + with self.assertRaisesRegex(ValueError, "alpha must be in"): + psdf.groupby(psdf.id).ewm(alpha=1.7).mean() + + with self.assertRaisesRegex(ValueError, "Must pass one of com, span, halflife, or alpha"): + psdf.groupby(psdf.id).ewm().mean() + + with self.assertRaisesRegex( + ValueError, "com, span, halflife, and alpha are mutually exclusive" + ): + psdf.groupby(psdf.id).ewm(com=0.5, alpha=0.7).mean() + + def _test_ewm_func(self, f): + pser = pd.Series([1, 2, 3], index=np.random.rand(3), name="a") + psser = ps.from_pandas(pser) + self.assert_eq(getattr(psser.ewm(com=0.2), f)(), getattr(pser.ewm(com=0.2), f)()) + self.assert_eq( + getattr(psser.ewm(com=0.2), f)().sum(), getattr(pser.ewm(com=0.2), f)().sum() + ) + self.assert_eq(getattr(psser.ewm(span=1.7), f)(), getattr(pser.ewm(span=1.7), f)()) + self.assert_eq( + getattr(psser.ewm(span=1.7), f)().sum(), getattr(pser.ewm(span=1.7), f)().sum() + ) + self.assert_eq(getattr(psser.ewm(halflife=0.5), f)(), getattr(pser.ewm(halflife=0.5), f)()) + self.assert_eq( + getattr(psser.ewm(halflife=0.5), f)().sum(), getattr(pser.ewm(halflife=0.5), f)().sum() + ) + self.assert_eq(getattr(psser.ewm(alpha=0.7), f)(), getattr(pser.ewm(alpha=0.7), f)()) + self.assert_eq( + getattr(psser.ewm(alpha=0.7), f)().sum(), getattr(pser.ewm(alpha=0.7), f)().sum() + ) + self.assert_eq( + getattr(psser.ewm(alpha=0.7, min_periods=2), f)(), + getattr(pser.ewm(alpha=0.7, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psser.ewm(alpha=0.7, min_periods=2), f)().sum(), + getattr(pser.ewm(alpha=0.7, min_periods=2), f)().sum(), + ) + + pdf = pd.DataFrame( + {"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(getattr(psdf.ewm(com=0.2), f)(), getattr(pdf.ewm(com=0.2), f)()) + self.assert_eq(getattr(psdf.ewm(com=0.2), f)().sum(), getattr(pdf.ewm(com=0.2), f)().sum()) + self.assert_eq(getattr(psdf.ewm(span=1.7), f)(), getattr(pdf.ewm(span=1.7), f)()) + self.assert_eq( + getattr(psdf.ewm(span=1.7), f)().sum(), getattr(pdf.ewm(span=1.7), f)().sum() + ) + self.assert_eq(getattr(psdf.ewm(halflife=0.5), f)(), getattr(pdf.ewm(halflife=0.5), f)()) + self.assert_eq( + getattr(psdf.ewm(halflife=0.5), f)().sum(), getattr(pdf.ewm(halflife=0.5), f)().sum() + ) + self.assert_eq(getattr(psdf.ewm(alpha=0.7), f)(), getattr(pdf.ewm(alpha=0.7), f)()) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7), f)().sum(), getattr(pdf.ewm(alpha=0.7), f)().sum() + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, min_periods=2), f)(), + getattr(pdf.ewm(alpha=0.7, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, min_periods=2), f)().sum(), + getattr(pdf.ewm(alpha=0.7, min_periods=2), f)().sum(), + ) + + pdf = pd.DataFrame( + { + "s1": [None, 2, 3, 4], + "s2": [1, None, 3, 4], + "s3": [1, 3, 4, 5], + "s4": [1, 0, 3, 4], + "s5": [None, None, 1, None], + "s6": [None, None, None, None], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + getattr(psdf.ewm(com=0.2, ignore_na=True), f)(), + getattr(pdf.ewm(com=0.2, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(com=0.2, ignore_na=True), f)().sum(), + getattr(pdf.ewm(com=0.2, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(com=0.2, ignore_na=False), f)(), + getattr(pdf.ewm(com=0.2, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(com=0.2, ignore_na=False), f)().sum(), + getattr(pdf.ewm(com=0.2, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(span=1.7, ignore_na=True), f)(), + getattr(pdf.ewm(span=1.7, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(span=1.7, ignore_na=True), f)().sum(), + getattr(pdf.ewm(span=1.7, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(span=1.7, ignore_na=False), f)(), + getattr(pdf.ewm(span=1.7, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(span=1.7, ignore_na=False), f)().sum(), + getattr(pdf.ewm(span=1.7, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(halflife=0.5, ignore_na=True), f)(), + getattr(pdf.ewm(halflife=0.5, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(halflife=0.5, ignore_na=True), f)().sum(), + getattr(pdf.ewm(halflife=0.5, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(halflife=0.5, ignore_na=False), f)(), + getattr(pdf.ewm(halflife=0.5, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(halflife=0.5, ignore_na=False), f)().sum(), + getattr(pdf.ewm(halflife=0.5, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=True), f)(), + getattr(pdf.ewm(alpha=0.7, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=True), f)().sum(), + getattr(pdf.ewm(alpha=0.7, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=False), f)(), + getattr(pdf.ewm(alpha=0.7, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=False), f)().sum(), + getattr(pdf.ewm(alpha=0.7, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=True, min_periods=2), f)(), + getattr(pdf.ewm(alpha=0.7, ignore_na=True, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=True, min_periods=2), f)().sum(), + getattr(pdf.ewm(alpha=0.7, ignore_na=True, min_periods=2), f)().sum(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=False, min_periods=2), f)(), + getattr(pdf.ewm(alpha=0.7, ignore_na=False, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psdf.ewm(alpha=0.7, ignore_na=False, min_periods=2), f)().sum(), + getattr(pdf.ewm(alpha=0.7, ignore_na=False, min_periods=2), f)().sum(), + ) + + def test_ewm_mean(self): + self._test_ewm_func("mean") + + def _test_groupby_ewm_func(self, f): + pser = pd.Series([1, 2, 3, 2], index=np.random.rand(4), name="a") + psser = ps.from_pandas(pser) + self.assert_eq(getattr(psser.ewm(com=0.2), f)(), getattr(pser.ewm(com=0.2), f)()) + self.assert_eq( + getattr(psser.groupby(psser).ewm(com=0.2), f)().sum(), + getattr(pser.groupby(pser).ewm(com=0.2), f)().sum(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(span=1.7), f)(), + getattr(pser.groupby(pser).ewm(span=1.7), f)(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(span=1.7), f)().sum(), + getattr(pser.groupby(pser).ewm(span=1.7), f)().sum(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(halflife=0.5), f)(), + getattr(pser.groupby(pser).ewm(halflife=0.5), f)(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(halflife=0.5), f)().sum(), + getattr(pser.groupby(pser).ewm(halflife=0.5), f)().sum(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(alpha=0.7), f)(), + getattr(pser.groupby(pser).ewm(alpha=0.7), f)(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(alpha=0.7), f)().sum(), + getattr(pser.groupby(pser).ewm(alpha=0.7), f)().sum(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(alpha=0.7, min_periods=2), f)(), + getattr(pser.groupby(pser).ewm(alpha=0.7, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psser.groupby(psser).ewm(alpha=0.7, min_periods=2), f)().sum(), + getattr(pser.groupby(pser).ewm(alpha=0.7, min_periods=2), f)().sum(), + ) + + pdf = pd.DataFrame( + {"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(com=0.2), f)(), + getattr(pdf.groupby(pdf.a).ewm(com=0.2), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(com=0.2), f)().sum(), + getattr(pdf.groupby(pdf.a).ewm(com=0.2), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(span=1.7), f)(), + getattr(pdf.groupby(pdf.a).ewm(span=1.7), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(span=1.7), f)().sum(), + getattr(pdf.groupby(pdf.a).ewm(span=1.7), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(halflife=0.5), f)(), + getattr(pdf.groupby(pdf.a).ewm(halflife=0.5), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(halflife=0.5), f)().sum(), + getattr(pdf.groupby(pdf.a).ewm(halflife=0.5), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(alpha=0.7), f)(), + getattr(pdf.groupby(pdf.a).ewm(alpha=0.7), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(alpha=0.7), f)().sum(), + getattr(pdf.groupby(pdf.a).ewm(alpha=0.7), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(alpha=0.7, min_periods=2), f)(), + getattr(pdf.groupby(pdf.a).ewm(alpha=0.7, min_periods=2), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.a).ewm(alpha=0.7, min_periods=2), f)().sum(), + getattr(pdf.groupby(pdf.a).ewm(alpha=0.7, min_periods=2), f)().sum(), + ) + + pdf = pd.DataFrame( + { + "s1": [None, 2, 3, 4], + "s2": [1, None, 3, 4], + "s3": [1, 3, 4, 5], + "s4": [1, 0, 3, 4], + "s5": [None, None, 1, None], + "s6": [None, None, None, None], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(com=0.2, ignore_na=True), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(com=0.2, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(com=0.2, ignore_na=True), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(com=0.2, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(com=0.2, ignore_na=False), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(com=0.2, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(com=0.2, ignore_na=False), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(com=0.2, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(span=1.7, ignore_na=True), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(span=1.7, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(span=1.7, ignore_na=True), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(span=1.7, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(span=1.7, ignore_na=False), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(span=1.7, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(span=1.7, ignore_na=False), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(span=1.7, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(halflife=0.5, ignore_na=True), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(halflife=0.5, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(halflife=0.5, ignore_na=True), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(halflife=0.5, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(halflife=0.5, ignore_na=False), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(halflife=0.5, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(halflife=0.5, ignore_na=False), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(halflife=0.5, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=True), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=True), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=True), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=True), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=False), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=False), f)(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=False), f)().sum(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=False), f)().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=True, min_periods=2), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=True, min_periods=2), f)(), + ) + self.assert_eq( + getattr( + psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=True, min_periods=2), f + )().sum(), + getattr( + pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=True, min_periods=2), f + )().sum(), + ) + self.assert_eq( + getattr(psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=False, min_periods=2), f)(), + getattr(pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=False, min_periods=2), f)(), + ) + self.assert_eq( + getattr( + psdf.groupby(psdf.s1 + 1).ewm(alpha=0.7, ignore_na=False, min_periods=2), f + )().sum(), + getattr( + pdf.groupby(pdf.s1 + 1).ewm(alpha=0.7, ignore_na=False, min_periods=2), f + )().sum(), + ) + + def test_groupby_ewm_func(self): + self._test_groupby_ewm_func("mean") + + +class EWMTests(EWMTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_ewm import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_expanding.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_expanding.py new file mode 100644 index 0000000..10927e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_expanding.py @@ -0,0 +1,253 @@ +# +# 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 distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.pandas.window import Expanding +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class ExpandingTestsMixin: + def _test_expanding_func(self, ps_func, pd_func=None): + if not pd_func: + pd_func = ps_func + if isinstance(pd_func, str): + pd_func = self.convert_str_to_lambda(pd_func) + if isinstance(ps_func, str): + ps_func = self.convert_str_to_lambda(ps_func) + pser = pd.Series([1, 2, 3, 7, 9, 8], index=np.random.rand(6), name="a") + psser = ps.from_pandas(pser) + self.assert_eq(ps_func(psser.expanding(2)), pd_func(pser.expanding(2)), almost=True) + self.assert_eq(ps_func(psser.expanding(2)), pd_func(pser.expanding(2)), almost=True) + + # Multiindex + pser = pd.Series( + [1, 2, 3], index=pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")]) + ) + psser = ps.from_pandas(pser) + self.assert_eq(ps_func(psser.expanding(2)), pd_func(pser.expanding(2))) + + pdf = pd.DataFrame( + {"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(ps_func(psdf.expanding(2)), pd_func(pdf.expanding(2))) + self.assert_eq(ps_func(psdf.expanding(2)).sum(), pd_func(pdf.expanding(2)).sum()) + + # Multiindex column + columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(ps_func(psdf.expanding(2)), pd_func(pdf.expanding(2))) + + def test_expanding_error(self): + with self.assertRaisesRegex(ValueError, "min_periods must be >= 0"): + ps.range(10).expanding(-1) + + with self.assertRaisesRegex( + TypeError, "psdf_or_psser must be a series or dataframe; however, got:.*int" + ): + Expanding(1, 2) + + def test_expanding_repr(self): + self.assertEqual(repr(ps.range(10).expanding(5)), "Expanding [min_periods=5]") + + def test_expanding_count(self): + self._test_expanding_func("count") + + def test_expanding_min(self): + self._test_expanding_func("min") + + def test_expanding_max(self): + self._test_expanding_func("max") + + def test_expanding_mean(self): + self._test_expanding_func("mean") + + def test_expanding_quantile(self): + self._test_expanding_func(lambda x: x.quantile(0.5), lambda x: x.quantile(0.5, "lower")) + + def test_expanding_sum(self): + self._test_expanding_func("sum") + + def test_expanding_std(self): + self._test_expanding_func("std") + + def test_expanding_var(self): + self._test_expanding_func("var") + + def test_expanding_skew(self): + self._test_expanding_func("skew") + + def test_expanding_kurt(self): + self._test_expanding_func("kurt") + + def _test_groupby_expanding_func(self, ps_func, pd_func=None): + if not pd_func: + pd_func = ps_func + if isinstance(pd_func, str): + pd_func = self.convert_str_to_lambda(pd_func) + if isinstance(ps_func, str): + ps_func = self.convert_str_to_lambda(ps_func) + pser = pd.Series([1, 2, 3, 2], index=np.random.rand(4), name="a") + psser = ps.from_pandas(pser) + self.assert_eq( + ps_func(psser.groupby(psser).expanding(2)).sort_index(), + pd_func(pser.groupby(pser).expanding(2)).sort_index(), + ) + self.assert_eq( + ps_func(psser.groupby(psser).expanding(2)).sum(), + pd_func(pser.groupby(pser).expanding(2)).sum(), + ) + + # Multiindex + pser = pd.Series( + [1, 2, 3, 2], + index=pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z"), ("c", "z")]), + name="a", + ) + psser = ps.from_pandas(pser) + self.assert_eq( + ps_func(psser.groupby(psser).expanding(2)).sort_index(), + pd_func(pser.groupby(pser).expanding(2)).sort_index(), + ) + + pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}) + psdf = ps.from_pandas(pdf) + + # The behavior of GroupBy.expanding is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + ps_func(psdf.groupby(psdf.a).expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a).expanding(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a).expanding(2)).sum(), + pd_func(pdf.groupby(pdf.a).expanding(2)).sum(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a + 1).expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a + 1).expanding(2)).sort_index(), + ) + else: + self.assert_eq( + ps_func(psdf.groupby(psdf.a).expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a).expanding(2)).drop("a", axis=1).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a).expanding(2)).sum(), + pd_func(pdf.groupby(pdf.a).expanding(2)).sum().drop("a"), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a + 1).expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a + 1).expanding(2)).drop("a", axis=1).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.b.groupby(psdf.a).expanding(2)).sort_index(), + pd_func(pdf.b.groupby(pdf.a).expanding(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a)["b"].expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a)["b"].expanding(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a)[["b"]].expanding(2)).sort_index(), + pd_func(pdf.groupby(pdf.a)[["b"]].expanding(2)).sort_index(), + ) + + # Multiindex column + columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")]) + pdf.columns = columns + psdf.columns = columns + + # The behavior of GroupBy.expanding is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + ps_func(psdf.groupby(("a", "x")).expanding(2)).sort_index(), + pd_func(pdf.groupby(("a", "x")).expanding(2)).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.groupby([("a", "x"), ("a", "y")]).expanding(2)).sort_index(), + pd_func(pdf.groupby([("a", "x"), ("a", "y")]).expanding(2)).sort_index(), + ) + else: + self.assert_eq( + ps_func(psdf.groupby(("a", "x")).expanding(2)).sort_index(), + pd_func(pdf.groupby(("a", "x")).expanding(2)).drop(("a", "x"), axis=1).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.groupby([("a", "x"), ("a", "y")]).expanding(2)).sort_index(), + pd_func(pdf.groupby([("a", "x"), ("a", "y")]).expanding(2)) + .drop([("a", "x"), ("a", "y")], axis=1) + .sort_index(), + ) + + def test_groupby_expanding_count(self): + self._test_groupby_expanding_func("count") + + def test_groupby_expanding_min(self): + self._test_groupby_expanding_func("min") + + def test_groupby_expanding_max(self): + self._test_groupby_expanding_func("max") + + def test_groupby_expanding_mean(self): + self._test_groupby_expanding_func("mean") + + def test_groupby_expanding_quantile(self): + self._test_groupby_expanding_func( + lambda x: x.quantile(0.5), lambda x: x.quantile(0.5, "lower") + ) + + def test_groupby_expanding_sum(self): + self._test_groupby_expanding_func("sum") + + def test_groupby_expanding_std(self): + self._test_groupby_expanding_func("std") + + def test_groupby_expanding_var(self): + self._test_groupby_expanding_func("var") + + def test_groupby_expanding_skew(self): + self._test_groupby_expanding_func("skew") + + def test_groupby_expanding_kurt(self): + self._test_groupby_expanding_func("kurt") + + +class ExpandingTests(ExpandingTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_expanding import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_extension.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_extension.py new file mode 100644 index 0000000..fba850c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_extension.py @@ -0,0 +1,152 @@ +# +# 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 contextlib + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import assert_produces_warning, ComparisonTestBase +from pyspark.pandas.extensions import ( + register_dataframe_accessor, + register_series_accessor, + register_index_accessor, +) + + +@contextlib.contextmanager +def ensure_removed(obj, attr): + """ + Ensure attribute attached to 'obj' during testing is removed in the end + """ + try: + yield + + finally: + try: + delattr(obj, attr) + except AttributeError: + pass + + +class CustomAccessor: + def __init__(self, obj): + self.obj = obj + self.item = "item" + + @property + def prop(self): + return self.item + + def method(self): + return self.item + + def check_length(self, col=None): + if type(self.obj) == ps.DataFrame or col is not None: + return len(self.obj[col]) + else: + try: + return len(self.obj) + except Exception as e: + raise ValueError(str(e)) + + +class ExtensionTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=np.random.rand(9), + ) + + @property + def accessor(self): + return CustomAccessor(self.psdf) + + def test_setup(self): + self.assertEqual("item", self.accessor.item) + + def test_dataframe_register(self): + with ensure_removed(ps.DataFrame, "test"): + register_dataframe_accessor("test")(CustomAccessor) + assert self.psdf.test.prop == "item" + assert self.psdf.test.method() == "item" + assert len(self.psdf["a"]) == self.psdf.test.check_length("a") + + def test_series_register(self): + with ensure_removed(ps.Series, "test"): + register_series_accessor("test")(CustomAccessor) + assert self.psdf.a.test.prop == "item" + assert self.psdf.a.test.method() == "item" + assert self.psdf.a.test.check_length() == len(self.psdf["a"]) + + def test_index_register(self): + with ensure_removed(ps.Index, "test"): + register_index_accessor("test")(CustomAccessor) + assert self.psdf.index.test.prop == "item" + assert self.psdf.index.test.method() == "item" + assert self.psdf.index.test.check_length() == self.psdf.index.size + + def test_accessor_works(self): + register_series_accessor("test")(CustomAccessor) + + s = ps.Series([1, 2]) + assert s.test.obj is s + assert s.test.prop == "item" + assert s.test.method() == "item" + + def test_overwrite_warns(self): + mean = ps.Series.mean + try: + with assert_produces_warning(UserWarning, raise_on_extra_warnings=False) as w: + register_series_accessor("mean")(CustomAccessor) + s = ps.Series([1, 2]) + assert s.mean.prop == "item" + msg = str(w[0].message) + assert "mean" in msg + assert "CustomAccessor" in msg + assert "Series" in msg + finally: + ps.Series.mean = mean + + def test_raises_attr_error(self): + with ensure_removed(ps.Series, "bad"): + + class Bad: + def __init__(self, data): + raise AttributeError("whoops") + + with self.assertRaises(AttributeError): + ps.Series([1, 2], dtype=object).bad + + +class ExtensionTests(ExtensionTestsMixin, ComparisonTestBase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_extension import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_frame_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_frame_spark.py new file mode 100644 index 0000000..f1c7859 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_frame_spark.py @@ -0,0 +1,162 @@ +# +# 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 os + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils +from pyspark.testing.sqlutils import SQLTestUtils + + +class SparkFrameMethodsTestsMixin: + def test_frame_apply_negative(self): + with self.assertRaisesRegex( + ValueError, "The output of the function.* pyspark.sql.DataFrame.*int" + ): + ps.range(10).spark.apply(lambda scol: 1) + + def test_hint(self): + pdf1 = pd.DataFrame( + {"lkey": ["foo", "bar", "baz", "foo"], "value": [1, 2, 3, 5]} + ).set_index("lkey") + pdf2 = pd.DataFrame( + {"rkey": ["foo", "bar", "baz", "foo"], "value": [5, 6, 7, 8]} + ).set_index("rkey") + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + hints = ["broadcast", "merge", "shuffle_hash", "shuffle_replicate_nl"] + + for hint in hints: + self.assert_eq( + pdf1.merge(pdf2, left_index=True, right_index=True).sort_values( + ["value_x", "value_y"] + ), + psdf1.merge(psdf2.spark.hint(hint), left_index=True, right_index=True).sort_values( + ["value_x", "value_y"] + ), + almost=True, + ) + self.assert_eq( + pdf1.merge(pdf2 + 1, left_index=True, right_index=True).sort_values( + ["value_x", "value_y"] + ), + psdf1.merge( + (psdf2 + 1).spark.hint(hint), left_index=True, right_index=True + ).sort_values(["value_x", "value_y"]), + almost=True, + ) + + def test_repartition(self): + psdf = ps.DataFrame({"age": [5, 5, 2, 2], "name": ["Bob", "Bob", "Alice", "Alice"]}) + num_partitions = psdf.to_spark().rdd.getNumPartitions() + 1 + + num_partitions += 1 + new_psdf = psdf.spark.repartition(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + # Reserves Index + psdf = psdf.set_index("age") + num_partitions += 1 + new_psdf = psdf.spark.repartition(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + # Reflects internal changes + psdf = psdf.reset_index() + psdf = psdf.set_index("name") + psdf2 = psdf + 1 + num_partitions += 1 + self.assert_eq( + psdf2.sort_index(), (psdf + 1).spark.repartition(num_partitions).sort_index() + ) + + # Reserves MultiIndex + psdf = ps.DataFrame({"a": ["a", "b", "c"]}, index=[[1, 2, 3], [4, 5, 6]]) + num_partitions = psdf.to_spark().rdd.getNumPartitions() + 1 + new_psdf = psdf.spark.repartition(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + def test_coalesce(self): + num_partitions = 10 + psdf = ps.DataFrame({"age": [5, 5, 2, 2], "name": ["Bob", "Bob", "Alice", "Alice"]}) + psdf = psdf.spark.repartition(num_partitions) + + num_partitions -= 1 + new_psdf = psdf.spark.coalesce(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + # Reserves Index + psdf = psdf.set_index("age") + num_partitions -= 1 + new_psdf = psdf.spark.coalesce(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + # Reflects internal changes + psdf = psdf.reset_index() + psdf = psdf.set_index("name") + psdf2 = psdf + 1 + num_partitions -= 1 + self.assert_eq(psdf2.sort_index(), (psdf + 1).spark.coalesce(num_partitions).sort_index()) + + # Reserves MultiIndex + psdf = ps.DataFrame({"a": ["a", "b", "c"]}, index=[[1, 2, 3], [4, 5, 6]]) + num_partitions -= 1 + psdf = psdf.spark.repartition(num_partitions) + + num_partitions -= 1 + new_psdf = psdf.spark.coalesce(num_partitions) + self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions) + self.assert_eq(psdf.sort_index(), new_psdf.sort_index()) + + def test_checkpoint(self): + with self.temp_dir() as tmp: + self.spark.sparkContext.setCheckpointDir(tmp) + psdf = ps.DataFrame({"a": ["a", "b", "c"]}) + new_psdf = psdf.spark.checkpoint() + self.assertIsNotNone(os.listdir(tmp)) + self.assert_eq(psdf, new_psdf) + + def test_local_checkpoint(self): + psdf = ps.DataFrame({"a": ["a", "b", "c"]}) + new_psdf = psdf.spark.local_checkpoint() + self.assert_eq(psdf, new_psdf) + + +class SparkFrameMethodsTests( + SparkFrameMethodsTestsMixin, PandasOnSparkTestCase, SQLTestUtils, TestUtils +): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_frame_spark import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_generic_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_generic_functions.py new file mode 100644 index 0000000..f537e10 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_generic_functions.py @@ -0,0 +1,234 @@ +# +# 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 numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class GenericFunctionsTestsMixin: + def test_interpolate_error(self): + psdf = ps.range(10) + + with self.assertRaisesRegex( + NotImplementedError, "interpolate currently works only for method='linear'" + ): + psdf.interpolate(method="quadratic") + + with self.assertRaisesRegex( + NotImplementedError, "interpolate currently works only for method='linear'" + ): + psdf.id.interpolate(method="quadratic") + + with self.assertRaisesRegex(ValueError, "limit must be > 0"): + psdf.interpolate(limit=0) + + with self.assertRaisesRegex(ValueError, "limit must be > 0"): + psdf.id.interpolate(limit=0) + + with self.assertRaisesRegex(ValueError, "invalid limit_direction"): + psdf.interpolate(limit_direction="jump") + + with self.assertRaisesRegex(ValueError, "invalid limit_direction"): + psdf.id.interpolate(limit_direction="jump") + + with self.assertRaisesRegex(ValueError, "invalid limit_area"): + psdf.interpolate(limit_area="jump") + + with self.assertRaisesRegex(ValueError, "invalid limit_area"): + psdf.id.interpolate(limit_area="jump") + + def _test_interpolate(self, pobj): + psobj = ps.from_pandas(pobj) + self.assert_eq(psobj.interpolate(), pobj.interpolate()) + for limit in range(1, 5): + for limit_direction in [None, "forward", "backward", "both"]: + for limit_area in [None, "inside", "outside"]: + self.assert_eq( + psobj.interpolate( + limit=limit, limit_direction=limit_direction, limit_area=limit_area + ), + pobj.interpolate( + limit=limit, limit_direction=limit_direction, limit_area=limit_area + ), + ) + + def test_interpolate(self): + pser = pd.Series( + [ + 1, + np.nan, + 3, + ], + name="a", + ) + self._test_interpolate(pser) + + pser = pd.Series( + [ + np.nan, + np.nan, + np.nan, + ], + name="a", + ) + self._test_interpolate(pser) + + pser = pd.Series( + [ + np.nan, + np.nan, + np.nan, + 0, + 1, + np.nan, + np.nan, + np.nan, + np.nan, + 3, + np.nan, + np.nan, + np.nan, + ], + name="a", + ) + self._test_interpolate(pser) + + pdf = pd.DataFrame( + [ + (1, 0.0, np.nan), + (2, np.nan, 2.0), + (3, 2.0, 3.0), + (4, np.nan, 4.0), + (5, np.nan, 1.0), + ], + columns=list("abc"), + ) + self._test_interpolate(pdf) + + pdf = pd.DataFrame( + [ + (0.0, np.nan, -1.0, 1.0, np.nan), + (np.nan, 2.0, np.nan, np.nan, np.nan), + (2.0, 3.0, np.nan, 9.0, np.nan), + (np.nan, 4.0, -4.0, 16.0, np.nan), + (np.nan, 1.0, np.nan, 7.0, np.nan), + ], + columns=list("abcde"), + ) + self._test_interpolate(pdf) + + pdf = pd.DataFrame( + [ + (0.0, np.nan, -1.0, False, np.nan), + (np.nan, 2.0, np.nan, True, np.nan), + (2.0, 3.0, np.nan, True, np.nan), + (np.nan, 4.0, -4.0, False, np.nan), + (np.nan, 1.0, np.nan, True, np.nan), + ], + columns=list("abcde"), + ) + self._test_interpolate(pdf) + + def _test_stat_functions(self, stat_func): + pdf = pd.DataFrame({"a": [np.nan, np.nan, np.nan], "b": [1, np.nan, 2], "c": [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(stat_func(pdf.a), stat_func(psdf.a)) + self.assert_eq(stat_func(pdf.b), stat_func(psdf.b)) + self.assert_eq(stat_func(pdf), stat_func(psdf)) + + # Fix skew and kurtosis and re-enable tests below + def test_stat_functions(self): + self._test_stat_functions(lambda x: x.sum()) + self._test_stat_functions(lambda x: x.sum(skipna=False)) + self._test_stat_functions(lambda x: x.mean()) + self._test_stat_functions(lambda x: x.mean(skipna=False)) + self._test_stat_functions(lambda x: x.product()) + self._test_stat_functions(lambda x: x.product(skipna=False)) + self._test_stat_functions(lambda x: x.min()) + self._test_stat_functions(lambda x: x.min(skipna=False)) + self._test_stat_functions(lambda x: x.max()) + self._test_stat_functions(lambda x: x.max(skipna=False)) + self._test_stat_functions(lambda x: x.std()) + self._test_stat_functions(lambda x: x.std(skipna=False)) + self._test_stat_functions(lambda x: x.std(ddof=2)) + self._test_stat_functions(lambda x: x.var()) + self._test_stat_functions(lambda x: x.var(ddof=2)) + self._test_stat_functions(lambda x: x.sem()) + self._test_stat_functions(lambda x: x.sem(skipna=False)) + # self._test_stat_functions(lambda x: x.skew()) + self._test_stat_functions(lambda x: x.skew(skipna=False)) + + # Test cases below return differently from pandas (either by design or to be fixed) + pdf = pd.DataFrame({"a": [np.nan, np.nan, np.nan], "b": [1, np.nan, 2], "c": [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf.std(ddof="ddof") + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf.a.std(ddof="ddof") + + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf.var(ddof="ddof") + with self.assertRaisesRegex(TypeError, "ddof must be integer"): + psdf.a.var(ddof="ddof") + + self.assert_eq(pdf.a.median(), psdf.a.median()) + self.assert_eq(pdf.a.median(skipna=False), psdf.a.median(skipna=False)) + self.assert_eq(1.0, psdf.b.median()) + self.assert_eq(pdf.b.median(skipna=False), psdf.b.median(skipna=False)) + self.assert_eq(pdf.c.median(), psdf.c.median()) + + self.assert_eq(pdf.a.kurtosis(skipna=False), psdf.a.kurtosis(skipna=False)) + self.assert_eq(pdf.a.kurtosis(), psdf.a.kurtosis()) + self.assert_eq(pdf.b.kurtosis(skipna=False), psdf.b.kurtosis(skipna=False)) + self.assert_eq(pdf.b.kurtosis(), psdf.b.kurtosis()) + self.assert_eq(pdf.c.kurtosis(), psdf.c.kurtosis()) + + def test_prod_precision(self): + pdf = pd.DataFrame( + { + "a": [np.nan, np.nan, np.nan, np.nan], + "b": [1, np.nan, np.nan, -4], + "c": [1, -2, 3, -4], + "d": [55108, 55108, 55108, 55108], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(pdf.prod(), psdf.prod()) + self.assert_eq(pdf.prod(skipna=False), psdf.prod(skipna=False)) + self.assert_eq(pdf.prod(min_count=3), psdf.prod(min_count=3)) + self.assert_eq(pdf.prod(skipna=False, min_count=3), psdf.prod(skipna=False, min_count=3)) + + +class GenericFunctionsTests(GenericFunctionsTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_generic_functions import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexing.py new file mode 100644 index 0000000..689d4e0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexing.py @@ -0,0 +1,1339 @@ +# +# 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 datetime +from distutils.version import LooseVersion +import unittest + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import SparkPandasIndexingError +from pyspark.testing.pandasutils import ComparisonTestBase, compare_both + + +class BasicIndexingTestsMixin: + @property + def pdf(self): + return pd.DataFrame( + {"month": [1, 4, 7, 10], "year": [2012, 2014, 2013, 2014], "sale": [55, 40, 84, 31]} + ) + + @compare_both(almost=False) + def test_indexing(self, df): + df1 = df.set_index("month") + yield df1 + + yield df.set_index("month", drop=False) + yield df.set_index("month", append=True) + yield df.set_index(["year", "month"]) + yield df.set_index(["year", "month"], drop=False) + yield df.set_index(["year", "month"], append=True) + + yield df1.set_index("year", drop=False, append=True) + + df2 = df1.copy() + df2.set_index("year", append=True, inplace=True) + yield df2 + + self.assertRaisesRegex(KeyError, "unknown", lambda: df.set_index("unknown")) + self.assertRaisesRegex(KeyError, "unknown", lambda: df.set_index(["month", "unknown"])) + + for d in [df, df1, df2]: + yield d.reset_index() + yield d.reset_index(drop=True) + + yield df1.reset_index(level=0) + yield df2.reset_index(level=1) + yield df2.reset_index(level=[1, 0]) + yield df1.reset_index(level="month") + yield df2.reset_index(level="year") + yield df2.reset_index(level=["month", "year"]) + yield df2.reset_index(level="month", drop=True) + yield df2.reset_index(level=["month", "year"], drop=True) + + self.assertRaisesRegex( + IndexError, + "Too many levels: Index has only 1 level, not 3", + lambda: df1.reset_index(level=2), + ) + self.assertRaisesRegex( + IndexError, + "Too many levels: Index has only 1 level, not 4", + lambda: df1.reset_index(level=[3, 2]), + ) + self.assertRaisesRegex(KeyError, "unknown.*month", lambda: df1.reset_index(level="unknown")) + self.assertRaisesRegex( + KeyError, "Level unknown not found", lambda: df2.reset_index(level="unknown") + ) + + df3 = df2.copy() + df3.reset_index(inplace=True) + yield df3 + + yield df1.sale.reset_index() + yield df1.sale.reset_index(level=0) + yield df2.sale.reset_index(level=[1, 0]) + yield df1.sale.reset_index(drop=True) + yield df1.sale.reset_index(name="s") + yield df1.sale.reset_index(name="s", drop=True) + + s = df1.sale + self.assertRaisesRegex( + TypeError, + "Cannot reset_index inplace on a Series to create a DataFrame", + lambda: s.reset_index(inplace=True), + ) + s.reset_index(drop=True, inplace=True) + yield s + yield df1 + + # multi-index columns + df4 = df.copy() + df4.columns = pd.MultiIndex.from_tuples( + [("cal", "month"), ("cal", "year"), ("num", "sale")] + ) + df5 = df4.set_index(("cal", "month")) + yield df5 + yield df4.set_index([("cal", "month"), ("num", "sale")]) + + self.assertRaises(KeyError, lambda: df5.reset_index(level=("cal", "month"))) + + yield df5.reset_index(level=[("cal", "month")]) + + # non-string names + df6 = df.copy() + df6.columns = [10.0, 20.0, 30.0] + df7 = df6.set_index(10.0) + yield df7 + yield df6.set_index([10.0, 30.0]) + + yield df7.reset_index(level=10.0) + yield df7.reset_index(level=[10.0]) + + df8 = df.copy() + df8.columns = pd.MultiIndex.from_tuples([(10, "month"), (10, "year"), (20, "sale")]) + df9 = df8.set_index((10, "month")) + yield df9 + yield df8.set_index([(10, "month"), (20, "sale")]) + + yield df9.reset_index(level=[(10, "month")]) + + def test_from_pandas_with_explicit_index(self): + pdf = self.pdf + + df1 = ps.from_pandas(pdf.set_index("month")) + self.assertPandasEqual(df1._to_pandas(), pdf.set_index("month")) + + df2 = ps.from_pandas(pdf.set_index(["year", "month"])) + self.assertPandasEqual(df2._to_pandas(), pdf.set_index(["year", "month"])) + + def test_limitations(self): + df = self.psdf.set_index("month") + + self.assertRaisesRegex( + ValueError, + "Level should be all int or all string.", + lambda: df.reset_index([1, "month"]), + ) + + +class IndexingTest(ComparisonTestBase): + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + @property + def pdf2(self): + return pd.DataFrame( + {0: [1, 2, 3, 4, 5, 6, 7, 8, 9], 1: [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + @property + def psdf2(self): + return ps.from_pandas(self.pdf2) + + def test_at(self): + pdf = self.pdf + psdf = self.psdf + # Create the equivalent of pdf.loc[3] as a Koalas Series + # This is necessary because .loc[n] does not currently work with Koalas DataFrames (#383) + test_series = ps.Series([3, 6], index=["a", "b"], name="3") + + # Assert invalided signatures raise TypeError + with self.assertRaises(TypeError, msg="Use DataFrame.at like .at[row_index, column_name]"): + psdf.at[3] + with self.assertRaises(TypeError, msg="Use DataFrame.at like .at[row_index, column_name]"): + psdf.at["ab"] # 'ab' is of length 2 but str type instead of tuple + with self.assertRaises(TypeError, msg="Use Series.at like .at[column_name]"): + test_series.at[3, "b"] + + # Assert .at for DataFrames + self.assertEqual(psdf.at[3, "b"], 6) + self.assertEqual(psdf.at[3, "b"], pdf.at[3, "b"]) + self.assert_eq(psdf.at[9, "b"], np.array([0, 0, 0])) + self.assert_eq(psdf.at[9, "b"], pdf.at[9, "b"]) + + # Assert .at for Series + self.assertEqual(test_series.at["b"], 6) + self.assertEqual(test_series.at["b"], pdf.loc[3].at["b"]) + + # Assert multi-character indices + self.assertEqual( + ps.Series([0, 1], index=["ab", "cd"]).at["ab"], + pd.Series([0, 1], index=["ab", "cd"]).at["ab"], + ) + + # Assert invalid column or index names result in a KeyError like with pandas + with self.assertRaises(KeyError, msg="x"): + psdf.at[3, "x"] + with self.assertRaises(KeyError, msg=99): + psdf.at[99, "b"] + + with self.assertRaises(ValueError): + psdf.at[(3, 6), "b"] + with self.assertRaises(KeyError): + psdf.at[3, ("x", "b")] + + # Assert setting values fails + with self.assertRaises(TypeError): + psdf.at[3, "b"] = 10 + + # non-string column names + pdf = self.pdf2 + psdf = self.psdf2 + + # Assert .at for DataFrames + self.assertEqual(psdf.at[3, 1], 6) + self.assertEqual(psdf.at[3, 1], pdf.at[3, 1]) + self.assert_eq(psdf.at[9, 1], np.array([0, 0, 0])) + self.assert_eq(psdf.at[9, 1], pdf.at[9, 1]) + + def test_at_multiindex(self): + pdf = self.pdf.set_index("b", append=True) + psdf = self.psdf.set_index("b", append=True) + + # TODO: seems like a pandas' bug in pandas>=1.1.0 + if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): + self.assert_eq(psdf.at[(3, 6), "a"], pdf.at[(3, 6), "a"]) + self.assert_eq(psdf.at[(3,), "a"], pdf.at[(3,), "a"]) + self.assert_eq(list(psdf.at[(9, 0), "a"]), list(pdf.at[(9, 0), "a"])) + self.assert_eq(list(psdf.at[(9,), "a"]), list(pdf.at[(9,), "a"])) + else: + self.assert_eq(psdf.at[(3, 6), "a"], 3) + self.assert_eq(psdf.at[(3,), "a"], np.array([3])) + self.assert_eq(list(psdf.at[(9, 0), "a"]), [7, 8, 9]) + self.assert_eq(list(psdf.at[(9,), "a"]), [7, 8, 9]) + + with self.assertRaises(ValueError): + psdf.at[3, "a"] + + def test_at_multiindex_columns(self): + arrays = [np.array(["bar", "bar", "baz", "baz"]), np.array(["one", "two", "one", "two"])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.at["B", ("bar", "one")], pdf.at["B", ("bar", "one")]) + + with self.assertRaises(KeyError): + psdf.at["B", "bar"] + + # non-string column names + arrays = [np.array([0, 0, 1, 1]), np.array([1, 2, 1, 2])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.at["B", (0, 1)], pdf.at["B", (0, 1)]) + + def test_iat(self): + pdf = self.pdf + psdf = self.psdf + # Create the equivalent of pdf.loc[3] as a Koalas Series + # This is necessary because .loc[n] does not currently work with Koalas DataFrames (#383) + test_series = ps.Series([3, 6], index=["a", "b"], name="3") + + # Assert invalided signatures raise TypeError + with self.assertRaises( + TypeError, + msg="Use DataFrame.at like .iat[row_interget_position, column_integer_position]", + ): + psdf.iat[3] + with self.assertRaises( + ValueError, msg="iAt based indexing on multi-index can only have tuple values" + ): + psdf.iat[3, "b"] # 'ab' is of length 2 but str type instead of tuple + with self.assertRaises(TypeError, msg="Use Series.iat like .iat[row_integer_position]"): + test_series.iat[3, "b"] + + # Assert .iat for DataFrames + self.assertEqual(psdf.iat[7, 0], 8) + self.assertEqual(psdf.iat[7, 0], pdf.iat[7, 0]) + + # Assert .iat for Series + self.assertEqual(test_series.iat[1], 6) + self.assertEqual(test_series.iat[1], pdf.loc[3].iat[1]) + + # Assert invalid column or integer position result in a KeyError like with pandas + with self.assertRaises(KeyError, msg=99): + psdf.iat[0, 99] + with self.assertRaises(KeyError, msg=99): + psdf.iat[99, 0] + + with self.assertRaises(ValueError): + psdf.iat[(1, 1), 1] + with self.assertRaises(ValueError): + psdf.iat[1, (1, 1)] + + # Assert setting values fails + with self.assertRaises(TypeError): + psdf.iat[4, 1] = 10 + + def test_iat_multiindex(self): + pdf = self.pdf.set_index("b", append=True) + psdf = self.psdf.set_index("b", append=True) + + self.assert_eq(psdf.iat[7, 0], pdf.iat[7, 0]) + + with self.assertRaises(ValueError): + psdf.iat[3, "a"] + + def test_iat_multiindex_columns(self): + arrays = [np.array(["bar", "bar", "baz", "baz"]), np.array(["one", "two", "one", "two"])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.iat[1, 3], pdf.iat[1, 3]) + + with self.assertRaises(KeyError): + psdf.iat[0, 99] + with self.assertRaises(KeyError): + psdf.iat[99, 0] + + def test_loc(self): + psdf = self.psdf + pdf = self.pdf + + self.assert_eq(psdf.loc[5:5], pdf.loc[5:5]) + self.assert_eq(psdf.loc[3:8], pdf.loc[3:8]) + self.assert_eq(psdf.loc[:8], pdf.loc[:8]) + self.assert_eq(psdf.loc[3:], pdf.loc[3:]) + self.assert_eq(psdf.loc[[5]], pdf.loc[[5]]) + self.assert_eq(psdf.loc[:], pdf.loc[:]) + + # TODO?: self.assert_eq(psdf.loc[[3, 4, 1, 8]], pdf.loc[[3, 4, 1, 8]]) + # TODO?: self.assert_eq(psdf.loc[[3, 4, 1, 9]], pdf.loc[[3, 4, 1, 9]]) + # TODO?: self.assert_eq(psdf.loc[np.array([3, 4, 1, 9])], pdf.loc[np.array([3, 4, 1, 9])]) + + self.assert_eq(psdf.a.loc[5:5], pdf.a.loc[5:5]) + self.assert_eq(psdf.a.loc[3:8], pdf.a.loc[3:8]) + self.assert_eq(psdf.a.loc[:8], pdf.a.loc[:8]) + self.assert_eq(psdf.a.loc[3:], pdf.a.loc[3:]) + self.assert_eq(psdf.a.loc[[5]], pdf.a.loc[[5]]) + + # TODO?: self.assert_eq(psdf.a.loc[[3, 4, 1, 8]], pdf.a.loc[[3, 4, 1, 8]]) + # TODO?: self.assert_eq(psdf.a.loc[[3, 4, 1, 9]], pdf.a.loc[[3, 4, 1, 9]]) + # TODO?: self.assert_eq(psdf.a.loc[np.array([3, 4, 1, 9])], + # pdf.a.loc[np.array([3, 4, 1, 9])]) + + self.assert_eq(psdf.a.loc[[]], pdf.a.loc[[]]) + self.assert_eq(psdf.a.loc[np.array([])], pdf.a.loc[np.array([])]) + + self.assert_eq(psdf.loc[1000:], pdf.loc[1000:]) + self.assert_eq(psdf.loc[-2000:-1000], pdf.loc[-2000:-1000]) + + self.assert_eq(psdf.loc[5], pdf.loc[5]) + self.assert_eq(psdf.loc[9], pdf.loc[9]) + self.assert_eq(psdf.a.loc[5], pdf.a.loc[5]) + self.assert_eq(psdf.a.loc[9], pdf.a.loc[9]) + + self.assertRaises(KeyError, lambda: psdf.loc[10]) + self.assertRaises(KeyError, lambda: psdf.a.loc[10]) + + # monotonically increasing index test + pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}, index=[0, 1, 1, 2, 2, 2, 4, 5, 6]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc[:2], pdf.loc[:2]) + self.assert_eq(psdf.loc[:3], pdf.loc[:3]) + self.assert_eq(psdf.loc[3:], pdf.loc[3:]) + self.assert_eq(psdf.loc[4:], pdf.loc[4:]) + self.assert_eq(psdf.loc[3:2], pdf.loc[3:2]) + self.assert_eq(psdf.loc[-1:2], pdf.loc[-1:2]) + self.assert_eq(psdf.loc[3:10], pdf.loc[3:10]) + + # monotonically decreasing index test + pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}, index=[6, 5, 5, 4, 4, 4, 2, 1, 0]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc[:4], pdf.loc[:4]) + self.assert_eq(psdf.loc[:3], pdf.loc[:3]) + self.assert_eq(psdf.loc[3:], pdf.loc[3:]) + self.assert_eq(psdf.loc[2:], pdf.loc[2:]) + self.assert_eq(psdf.loc[2:3], pdf.loc[2:3]) + self.assert_eq(psdf.loc[2:-1], pdf.loc[2:-1]) + self.assert_eq(psdf.loc[10:3], pdf.loc[10:3]) + + # test when type of key is string and given value is not included in key + pdf = pd.DataFrame({"a": [1, 2, 3]}, index=["a", "b", "d"]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc["a":"z"], pdf.loc["a":"z"]) + + # KeyError when index is not monotonic increasing or decreasing + # and specified values don't exist in index + psdf = ps.DataFrame([[1, 2], [4, 5], [7, 8]], index=["cobra", "viper", "sidewinder"]) + + self.assertRaises(KeyError, lambda: psdf.loc["cobra":"koalas"]) + self.assertRaises(KeyError, lambda: psdf.loc["koalas":"viper"]) + + psdf = ps.DataFrame([[1, 2], [4, 5], [7, 8]], index=[10, 30, 20]) + + self.assertRaises(KeyError, lambda: psdf.loc[0:30]) + self.assertRaises(KeyError, lambda: psdf.loc[10:100]) + + def test_loc_getitem_boolean_series(self): + pdf = pd.DataFrame( + {"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}, index=[20, 10, 30, 0, 50] + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(pdf.A.loc[pdf.B > 200], psdf.A.loc[psdf.B > 200]) + self.assert_eq(pdf.B.loc[pdf.B > 200], psdf.B.loc[psdf.B > 200]) + self.assert_eq(pdf.loc[pdf.B > 200], psdf.loc[psdf.B > 200]) + + def test_loc_non_informative_index(self): + pdf = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 30, 40]) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc[20:30], pdf.loc[20:30]) + + pdf = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 20, 40]) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.loc[20:20], pdf.loc[20:20]) + + def test_loc_with_series(self): + psdf = self.psdf + pdf = self.pdf + + self.assert_eq(psdf.loc[psdf.a % 2 == 0], pdf.loc[pdf.a % 2 == 0]) + self.assert_eq(psdf.loc[psdf.a % 2 == 0, "a"], pdf.loc[pdf.a % 2 == 0, "a"]) + self.assert_eq(psdf.loc[psdf.a % 2 == 0, ["a"]], pdf.loc[pdf.a % 2 == 0, ["a"]]) + self.assert_eq(psdf.a.loc[psdf.a % 2 == 0], pdf.a.loc[pdf.a % 2 == 0]) + + self.assert_eq(psdf.loc[psdf.copy().a % 2 == 0], pdf.loc[pdf.copy().a % 2 == 0]) + self.assert_eq(psdf.loc[psdf.copy().a % 2 == 0, "a"], pdf.loc[pdf.copy().a % 2 == 0, "a"]) + self.assert_eq( + psdf.loc[psdf.copy().a % 2 == 0, ["a"]], pdf.loc[pdf.copy().a % 2 == 0, ["a"]] + ) + self.assert_eq(psdf.a.loc[psdf.copy().a % 2 == 0], pdf.a.loc[pdf.copy().a % 2 == 0]) + + def test_loc_noindex(self): + psdf = self.psdf + psdf = psdf.reset_index() + pdf = self.pdf + pdf = pdf.reset_index() + + self.assert_eq(psdf[["a"]], pdf[["a"]]) + + self.assert_eq(psdf.loc[:], pdf.loc[:]) + self.assert_eq(psdf.loc[5:5], pdf.loc[5:5]) + + def test_loc_multiindex(self): + psdf = self.psdf + psdf = psdf.set_index("b", append=True) + pdf = self.pdf + pdf = pdf.set_index("b", append=True) + + self.assert_eq(psdf.loc[:], pdf.loc[:]) + self.assert_eq(psdf.loc[5:5], pdf.loc[5:5]) + self.assert_eq(psdf.loc[5:9], pdf.loc[5:9]) + + self.assert_eq(psdf.loc[5], pdf.loc[5]) + self.assert_eq(psdf.loc[9], pdf.loc[9]) + # TODO: self.assert_eq(psdf.loc[(5, 3)], pdf.loc[(5, 3)]) + # TODO: self.assert_eq(psdf.loc[(9, 0)], pdf.loc[(9, 0)]) + self.assert_eq(psdf.a.loc[5], pdf.a.loc[5]) + self.assert_eq(psdf.a.loc[9], pdf.a.loc[9]) + self.assertTrue((psdf.a.loc[(5, 3)] == pdf.a.loc[(5, 3)]).all()) + self.assert_eq(psdf.a.loc[(9, 0)], pdf.a.loc[(9, 0)]) + + # monotonically increasing index test + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + index=pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("y", "c"), ("y", "d"), ("z", "e")] + ), + ) + psdf = ps.from_pandas(pdf) + + for rows_sel in [ + slice(None), + slice("y", None), + slice(None, "y"), + slice(("x", "b"), None), + slice(None, ("y", "c")), + slice(("x", "b"), ("y", "c")), + slice("x", ("y", "c")), + slice(("x", "b"), "y"), + ]: + with self.subTest("monotonically increasing", rows_sel=rows_sel): + self.assert_eq(psdf.loc[rows_sel], pdf.loc[rows_sel]) + self.assert_eq(psdf.a.loc[rows_sel], pdf.a.loc[rows_sel]) + + # monotonically increasing first index test + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + index=pd.MultiIndex.from_tuples( + [("x", "a"), ("x", "b"), ("y", "c"), ("y", "a"), ("z", "e")] + ), + ) + psdf = ps.from_pandas(pdf) + + for rows_sel in [ + slice(None), + slice("y", None), + slice(None, "y"), + ]: + with self.subTest("monotonically increasing first index", rows_sel=rows_sel): + self.assert_eq(psdf.loc[rows_sel], pdf.loc[rows_sel]) + self.assert_eq(psdf.a.loc[rows_sel], pdf.a.loc[rows_sel]) + + for rows_sel in [ + slice(("x", "b"), None), + slice(None, ("y", "c")), + slice(("x", "b"), ("y", "c")), + slice("x", ("y", "c")), + slice(("x", "b"), "y"), + ]: + with self.subTest("monotonically increasing first index", rows_sel=rows_sel): + self.assertRaises(KeyError, lambda: psdf.loc[rows_sel]) + self.assertRaises(KeyError, lambda: psdf.a.loc[rows_sel]) + + # not monotonically increasing index test + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5]}, + index=pd.MultiIndex.from_tuples( + [("z", "e"), ("y", "d"), ("y", "c"), ("x", "b"), ("x", "a")] + ), + ) + psdf = ps.from_pandas(pdf) + + for rows_sel in [ + slice("y", None), + slice(None, "y"), + slice(("x", "b"), None), + slice(None, ("y", "c")), + slice(("x", "b"), ("y", "c")), + slice("x", ("y", "c")), + slice(("x", "b"), "y"), + ]: + with self.subTest("monotonically decreasing", rows_sel=rows_sel): + self.assertRaises(KeyError, lambda: psdf.loc[rows_sel]) + self.assertRaises(KeyError, lambda: psdf.a.loc[rows_sel]) + + def test_loc2d_multiindex(self): + psdf = self.psdf + psdf = psdf.set_index("b", append=True) + pdf = self.pdf + pdf = pdf.set_index("b", append=True) + + self.assert_eq(psdf.loc[:, :], pdf.loc[:, :]) + self.assert_eq(psdf.loc[:, "a"], pdf.loc[:, "a"]) + self.assert_eq(psdf.loc[5:5, "a"], pdf.loc[5:5, "a"]) + + self.assert_eq(psdf.loc[:, "a":"a"], pdf.loc[:, "a":"a"]) + self.assert_eq(psdf.loc[:, "a":"c"], pdf.loc[:, "a":"c"]) + self.assert_eq(psdf.loc[:, "b":"c"], pdf.loc[:, "b":"c"]) + + def test_loc2d(self): + psdf = self.psdf + pdf = self.pdf + + # index indexer is always regarded as slice for duplicated values + self.assert_eq(psdf.loc[5:5, "a"], pdf.loc[5:5, "a"]) + self.assert_eq(psdf.loc[[5], "a"], pdf.loc[[5], "a"]) + self.assert_eq(psdf.loc[5:5, ["a"]], pdf.loc[5:5, ["a"]]) + self.assert_eq(psdf.loc[[5], ["a"]], pdf.loc[[5], ["a"]]) + self.assert_eq(psdf.loc[:, :], pdf.loc[:, :]) + + self.assert_eq(psdf.loc[3:8, "a"], pdf.loc[3:8, "a"]) + self.assert_eq(psdf.loc[:8, "a"], pdf.loc[:8, "a"]) + self.assert_eq(psdf.loc[3:, "a"], pdf.loc[3:, "a"]) + self.assert_eq(psdf.loc[[8], "a"], pdf.loc[[8], "a"]) + + self.assert_eq(psdf.loc[3:8, ["a"]], pdf.loc[3:8, ["a"]]) + self.assert_eq(psdf.loc[:8, ["a"]], pdf.loc[:8, ["a"]]) + self.assert_eq(psdf.loc[3:, ["a"]], pdf.loc[3:, ["a"]]) + # TODO?: self.assert_eq(psdf.loc[[3, 4, 3], ['a']], pdf.loc[[3, 4, 3], ['a']]) + + self.assertRaises(SparkPandasIndexingError, lambda: psdf.loc[3, 3, 3]) + self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[3, 3]) + self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[3:, 3]) + self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[psdf.a % 2 == 0, 3]) + + self.assert_eq(psdf.loc[5, "a"], pdf.loc[5, "a"]) + self.assert_eq(psdf.loc[9, "a"], pdf.loc[9, "a"]) + self.assert_eq(psdf.loc[5, ["a"]], pdf.loc[5, ["a"]]) + self.assert_eq(psdf.loc[9, ["a"]], pdf.loc[9, ["a"]]) + + self.assert_eq(psdf.loc[:, "a":"a"], pdf.loc[:, "a":"a"]) + self.assert_eq(psdf.loc[:, "a":"d"], pdf.loc[:, "a":"d"]) + self.assert_eq(psdf.loc[:, "c":"d"], pdf.loc[:, "c":"d"]) + + # bool list-like column select + bool_list = [True, False] + self.assert_eq(psdf.loc[:, bool_list], pdf.loc[:, bool_list]) + self.assert_eq(psdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)]) + + pser = pd.Series(bool_list, index=pdf.columns) + self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser]) + pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns))) + self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser]) + + self.assertRaises(IndexError, lambda: psdf.loc[:, bool_list[:-1]]) + self.assertRaises(IndexError, lambda: psdf.loc[:, np.array(bool_list + [True])]) + self.assertRaises(SparkPandasIndexingError, lambda: psdf.loc[:, pd.Series(bool_list)]) + + # non-string column names + psdf = self.psdf2 + pdf = self.pdf2 + + self.assert_eq(psdf.loc[5:5, 0], pdf.loc[5:5, 0]) + self.assert_eq(psdf.loc[5:5, [0]], pdf.loc[5:5, [0]]) + self.assert_eq(psdf.loc[3:8, 0], pdf.loc[3:8, 0]) + self.assert_eq(psdf.loc[3:8, [0]], pdf.loc[3:8, [0]]) + + self.assert_eq(psdf.loc[:, 0:0], pdf.loc[:, 0:0]) + self.assert_eq(psdf.loc[:, 0:3], pdf.loc[:, 0:3]) + self.assert_eq(psdf.loc[:, 2:3], pdf.loc[:, 2:3]) + + def test_loc2d_multiindex_columns(self): + arrays = [np.array(["bar", "bar", "baz", "baz"]), np.array(["one", "two", "one", "two"])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc["B":"B", "bar"], pdf.loc["B":"B", "bar"]) + self.assert_eq(psdf.loc["B":"B", ["bar"]], pdf.loc["B":"B", ["bar"]]) + + self.assert_eq(psdf.loc[:, "bar":"bar"], pdf.loc[:, "bar":"bar"]) + self.assert_eq(psdf.loc[:, "bar":("baz", "one")], pdf.loc[:, "bar":("baz", "one")]) + self.assert_eq( + psdf.loc[:, ("bar", "two"):("baz", "one")], pdf.loc[:, ("bar", "two"):("baz", "one")] + ) + self.assert_eq(psdf.loc[:, ("bar", "two"):"bar"], pdf.loc[:, ("bar", "two"):"bar"]) + self.assert_eq(psdf.loc[:, "a":"bax"], pdf.loc[:, "a":"bax"]) + self.assert_eq( + psdf.loc[:, ("bar", "x"):("baz", "a")], + pdf.loc[:, ("bar", "x"):("baz", "a")], + almost=True, + ) + + pdf = pd.DataFrame( + np.random.randn(3, 4), + index=["A", "B", "C"], + columns=pd.MultiIndex.from_tuples( + [("bar", "two"), ("bar", "one"), ("baz", "one"), ("baz", "two")] + ), + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc[:, "bar":"baz"], pdf.loc[:, "bar":"baz"]) + + self.assertRaises(KeyError, lambda: psdf.loc[:, "bar":("baz", "one")]) + self.assertRaises(KeyError, lambda: psdf.loc[:, ("bar", "two"):"bar"]) + + # bool list-like column select + bool_list = [True, False, True, False] + self.assert_eq(psdf.loc[:, bool_list], pdf.loc[:, bool_list]) + self.assert_eq(psdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)]) + + pser = pd.Series(bool_list, index=pdf.columns) + self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser]) + + pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns))) + self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser]) + + # non-string column names + arrays = [np.array([0, 0, 1, 1]), np.array([1, 2, 1, 2])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc["B":"B", 0], pdf.loc["B":"B", 0]) + self.assert_eq(psdf.loc["B":"B", [0]], pdf.loc["B":"B", [0]]) + self.assert_eq(psdf.loc[:, 0:0], pdf.loc[:, 0:0]) + self.assert_eq(psdf.loc[:, 0:(1, 1)], pdf.loc[:, 0:(1, 1)]) + self.assert_eq(psdf.loc[:, (0, 2):(1, 1)], pdf.loc[:, (0, 2):(1, 1)]) + self.assert_eq(psdf.loc[:, (0, 2):0], pdf.loc[:, (0, 2):0]) + self.assert_eq(psdf.loc[:, -1:2], pdf.loc[:, -1:2]) + + def test_loc2d_with_known_divisions(self): + pdf = pd.DataFrame( + np.random.randn(20, 5), index=list("abcdefghijklmnopqrst"), columns=list("ABCDE") + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc[["a"], "A"], pdf.loc[["a"], "A"]) + self.assert_eq(psdf.loc[["a"], ["A"]], pdf.loc[["a"], ["A"]]) + self.assert_eq(psdf.loc["a":"o", "A"], pdf.loc["a":"o", "A"]) + self.assert_eq(psdf.loc["a":"o", ["A"]], pdf.loc["a":"o", ["A"]]) + self.assert_eq(psdf.loc[["n"], ["A"]], pdf.loc[["n"], ["A"]]) + self.assert_eq(psdf.loc[["a", "c", "n"], ["A"]], pdf.loc[["a", "c", "n"], ["A"]]) + # TODO?: self.assert_eq(psdf.loc[['t', 'b'], ['A']], pdf.loc[['t', 'b'], ['A']]) + # TODO?: self.assert_eq(psdf.loc[['r', 'r', 'c', 'g', 'h'], ['A']], + # TODO?: pdf.loc[['r', 'r', 'c', 'g', 'h'], ['A']]) + + @unittest.skip("TODO: should handle duplicated columns properly") + def test_loc2d_duplicated_columns(self): + pdf = pd.DataFrame( + np.random.randn(20, 5), index=list("abcdefghijklmnopqrst"), columns=list("AABCD") + ) + psdf = ps.from_pandas(pdf) + + # TODO?: self.assert_eq(psdf.loc[['a'], 'A'], pdf.loc[['a'], 'A']) + # TODO?: self.assert_eq(psdf.loc[['a'], ['A']], pdf.loc[['a'], ['A']]) + self.assert_eq(psdf.loc[["j"], "B"], pdf.loc[["j"], "B"]) + self.assert_eq(psdf.loc[["j"], ["B"]], pdf.loc[["j"], ["B"]]) + + # TODO?: self.assert_eq(psdf.loc['a':'o', 'A'], pdf.loc['a':'o', 'A']) + # TODO?: self.assert_eq(psdf.loc['a':'o', ['A']], pdf.loc['a':'o', ['A']]) + self.assert_eq(psdf.loc["j":"q", "B"], pdf.loc["j":"q", "B"]) + self.assert_eq(psdf.loc["j":"q", ["B"]], pdf.loc["j":"q", ["B"]]) + + # TODO?: self.assert_eq(psdf.loc['a':'o', 'B':'D'], pdf.loc['a':'o', 'B':'D']) + # TODO?: self.assert_eq(psdf.loc['a':'o', 'B':'D'], pdf.loc['a':'o', 'B':'D']) + # TODO?: self.assert_eq(psdf.loc['j':'q', 'B':'A'], pdf.loc['j':'q', 'B':'A']) + # TODO?: self.assert_eq(psdf.loc['j':'q', 'B':'A'], pdf.loc['j':'q', 'B':'A']) + + self.assert_eq(psdf.loc[psdf.B > 0, "B"], pdf.loc[pdf.B > 0, "B"]) + # TODO?: self.assert_eq(psdf.loc[psdf.B > 0, ['A', 'C']], pdf.loc[pdf.B > 0, ['A', 'C']]) + + def test_getitem(self): + pdf = pd.DataFrame( + { + "A": [1, 2, 3, 4, 5, 6, 7, 8, 9], + "B": [9, 8, 7, 6, 5, 4, 3, 2, 1], + "C": [True, False, True] * 3, + }, + columns=list("ABC"), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf["A"], pdf["A"]) + + self.assert_eq(psdf[["A", "B"]], pdf[["A", "B"]]) + + self.assert_eq(psdf[psdf.C], pdf[pdf.C]) + + self.assertRaises(KeyError, lambda: psdf["X"]) + self.assertRaises(KeyError, lambda: psdf[["A", "X"]]) + self.assertRaises(AttributeError, lambda: psdf.X) + + # not str/unicode + # TODO?: pdf = pd.DataFrame(np.random.randn(10, 5)) + # TODO?: psdf = ps.from_pandas(pdf) + # TODO?: self.assert_eq(psdf[0], pdf[0]) + # TODO?: self.assert_eq(psdf[[1, 2]], pdf[[1, 2]]) + + # TODO?: self.assertRaises(KeyError, lambda: pdf[8]) + # TODO?: self.assertRaises(KeyError, lambda: pdf[[1, 8]]) + + # non-string column names + pdf = pd.DataFrame( + { + 10: [1, 2, 3, 4, 5, 6, 7, 8, 9], + 20: [9, 8, 7, 6, 5, 4, 3, 2, 1], + 30: [True, False, True] * 3, + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf[10], pdf[10]) + self.assert_eq(psdf[[10, 20]], pdf[[10, 20]]) + + def test_getitem_slice(self): + pdf = pd.DataFrame( + { + "A": [1, 2, 3, 4, 5, 6, 7, 8, 9], + "B": [9, 8, 7, 6, 5, 4, 3, 2, 1], + "C": [True, False, True] * 3, + }, + index=list("abcdefghi"), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf["a":"e"], pdf["a":"e"]) + self.assert_eq(psdf["a":"b"], pdf["a":"b"]) + self.assert_eq(psdf["f":], pdf["f":]) + + def test_loc_on_numpy_datetimes(self): + pdf = pd.DataFrame( + {"x": [1, 2, 3]}, index=list(map(np.datetime64, ["2014", "2015", "2016"])) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc["2014":"2015"], pdf.loc["2014":"2015"]) + + def test_loc_on_pandas_datetimes(self): + pdf = pd.DataFrame( + {"x": [1, 2, 3]}, index=list(map(pd.Timestamp, ["2014", "2015", "2016"])) + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.loc["2014":"2015"], pdf.loc["2014":"2015"]) + + @unittest.skip("TODO?: the behavior of slice for datetime") + def test_loc_datetime_no_freq(self): + datetime_index = pd.date_range("2016-01-01", "2016-01-31", freq="12h") + datetime_index.freq = None # FORGET FREQUENCY + pdf = pd.DataFrame({"num": range(len(datetime_index))}, index=datetime_index) + psdf = ps.from_pandas(pdf) + + slice_ = slice("2016-01-03", "2016-01-05") + result = psdf.loc[slice_, :] + expected = pdf.loc[slice_, :] + self.assert_eq(result, expected) + + @unittest.skip("TODO?: the behavior of slice for datetime") + def test_loc_timestamp_str(self): + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.date_range("2011-01-01", freq="H", periods=100), + ) + psdf = ps.from_pandas(pdf) + + # partial string slice + # TODO?: self.assert_eq(pdf.loc['2011-01-02'], + # TODO?: psdf.loc['2011-01-02']) + self.assert_eq(pdf.loc["2011-01-02":"2011-01-05"], psdf.loc["2011-01-02":"2011-01-05"]) + + # series + # TODO?: self.assert_eq(pdf.A.loc['2011-01-02'], + # TODO?: psdf.A.loc['2011-01-02']) + self.assert_eq(pdf.A.loc["2011-01-02":"2011-01-05"], psdf.A.loc["2011-01-02":"2011-01-05"]) + + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.date_range("2011-01-01", freq="M", periods=100), + ) + psdf = ps.from_pandas(pdf) + # TODO?: self.assert_eq(pdf.loc['2011-01'], psdf.loc['2011-01']) + # TODO?: self.assert_eq(pdf.loc['2011'], psdf.loc['2011']) + + self.assert_eq(pdf.loc["2011-01":"2012-05"], psdf.loc["2011-01":"2012-05"]) + self.assert_eq(pdf.loc["2011":"2015"], psdf.loc["2011":"2015"]) + + # series + # TODO?: self.assert_eq(pdf.B.loc['2011-01'], psdf.B.loc['2011-01']) + # TODO?: self.assert_eq(pdf.B.loc['2011'], psdf.B.loc['2011']) + + self.assert_eq(pdf.B.loc["2011-01":"2012-05"], psdf.B.loc["2011-01":"2012-05"]) + self.assert_eq(pdf.B.loc["2011":"2015"], psdf.B.loc["2011":"2015"]) + + @unittest.skip("TODO?: the behavior of slice for datetime") + def test_getitem_timestamp_str(self): + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.date_range("2011-01-01", freq="H", periods=100), + ) + psdf = ps.from_pandas(pdf) + + # partial string slice + # TODO?: self.assert_eq(pdf['2011-01-02'], + # TODO?: psdf['2011-01-02']) + self.assert_eq(pdf["2011-01-02":"2011-01-05"], psdf["2011-01-02":"2011-01-05"]) + + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.date_range("2011-01-01", freq="M", periods=100), + ) + psdf = ps.from_pandas(pdf) + + # TODO?: self.assert_eq(pdf['2011-01'], psdf['2011-01']) + # TODO?: self.assert_eq(pdf['2011'], psdf['2011']) + + self.assert_eq(pdf["2011-01":"2012-05"], psdf["2011-01":"2012-05"]) + self.assert_eq(pdf["2011":"2015"], psdf["2011":"2015"]) + + @unittest.skip("TODO?: period index can't convert to DataFrame correctly") + def test_getitem_period_str(self): + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.period_range("2011-01-01", freq="H", periods=100), + ) + psdf = ps.from_pandas(pdf) + + # partial string slice + # TODO?: self.assert_eq(pdf['2011-01-02'], + # TODO?: psdf['2011-01-02']) + self.assert_eq(pdf["2011-01-02":"2011-01-05"], psdf["2011-01-02":"2011-01-05"]) + + pdf = pd.DataFrame( + {"A": np.random.randn(100), "B": np.random.randn(100)}, + index=pd.period_range("2011-01-01", freq="M", periods=100), + ) + psdf = ps.from_pandas(pdf) + + # TODO?: self.assert_eq(pdf['2011-01'], psdf['2011-01']) + # TODO?: self.assert_eq(pdf['2011'], psdf['2011']) + + self.assert_eq(pdf["2011-01":"2012-05"], psdf["2011-01":"2012-05"]) + self.assert_eq(pdf["2011":"2015"], psdf["2011":"2015"]) + + def test_iloc(self): + pdf = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.iloc[0, 0], pdf.iloc[0, 0]) + for indexer in [0, [0], [0, 1], [1, 0], [False, True, True], slice(0, 1)]: + self.assert_eq(psdf.iloc[:, indexer], pdf.iloc[:, indexer]) + self.assert_eq(psdf.iloc[:1, indexer], pdf.iloc[:1, indexer]) + self.assert_eq(psdf.iloc[:-1, indexer], pdf.iloc[:-1, indexer]) + # self.assert_eq(psdf.iloc[psdf.index == 2, indexer], pdf.iloc[pdf.index == 2, indexer]) + + def test_iloc_multiindex_columns(self): + arrays = [np.array(["bar", "bar", "baz", "baz"]), np.array(["one", "two", "one", "two"])] + + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + + for indexer in [0, [0], [0, 1], [1, 0], [False, True, True, True], slice(0, 1)]: + self.assert_eq(psdf.iloc[:, indexer], pdf.iloc[:, indexer]) + self.assert_eq(psdf.iloc[:1, indexer], pdf.iloc[:1, indexer]) + self.assert_eq(psdf.iloc[:-1, indexer], pdf.iloc[:-1, indexer]) + # self.assert_eq(psdf.iloc[psdf.index == "B", indexer], + # pdf.iloc[pdf.index == "B", indexer]) + + def test_iloc_series(self): + pser = pd.Series([1, 2, 3]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.iloc[0], pser.iloc[0]) + self.assert_eq(psser.iloc[:], pser.iloc[:]) + self.assert_eq(psser.iloc[:1], pser.iloc[:1]) + self.assert_eq(psser.iloc[:-1], pser.iloc[:-1]) + + self.assert_eq((psser + 1).iloc[0], (pser + 1).iloc[0]) + self.assert_eq((psser + 1).iloc[:], (pser + 1).iloc[:]) + self.assert_eq((psser + 1).iloc[:1], (pser + 1).iloc[:1]) + self.assert_eq((psser + 1).iloc[:-1], (pser + 1).iloc[:-1]) + + def test_iloc_slice_rows_sel(self): + pdf = pd.DataFrame({"A": [1, 2] * 5, "B": [3, 4] * 5, "C": [5, 6] * 5}) + psdf = ps.from_pandas(pdf) + + for rows_sel in [ + slice(None), + slice(0, 1), + slice(1, 2), + slice(-3, None), + slice(None, -3), + slice(None, 0), + slice(None, None, 3), + slice(3, 8, 2), + slice(None, None, -2), + slice(8, 3, -2), + slice(8, None, -2), + slice(None, 3, -2), + ]: + with self.subTest(rows_sel=rows_sel): + self.assert_eq(psdf.iloc[rows_sel].sort_index(), pdf.iloc[rows_sel].sort_index()) + self.assert_eq( + psdf.A.iloc[rows_sel].sort_index(), pdf.A.iloc[rows_sel].sort_index() + ) + self.assert_eq( + (psdf.A + 1).iloc[rows_sel].sort_index(), + (pdf.A + 1).iloc[rows_sel].sort_index(), + ) + + def test_iloc_iterable_rows_sel(self): + pdf = pd.DataFrame({"A": [1, 2] * 5, "B": [3, 4] * 5, "C": [5, 6] * 5}) + psdf = ps.from_pandas(pdf) + + for rows_sel in [ + [], + np.array([0, 1]), + [1, 2], + np.array([-3]), + [3], + np.array([-2]), + [8, 3, -5], + ]: + with self.subTest(rows_sel=rows_sel): + self.assert_eq(psdf.iloc[rows_sel].sort_index(), pdf.iloc[rows_sel].sort_index()) + self.assert_eq( + psdf.A.iloc[rows_sel].sort_index(), pdf.A.iloc[rows_sel].sort_index() + ) + self.assert_eq( + (psdf.A + 1).iloc[rows_sel].sort_index(), + (pdf.A + 1).iloc[rows_sel].sort_index(), + ) + + with self.subTest(rows_sel=rows_sel): + self.assert_eq( + psdf.iloc[rows_sel, :].sort_index(), pdf.iloc[rows_sel, :].sort_index() + ) + + with self.subTest(rows_sel=rows_sel): + self.assert_eq( + psdf.iloc[rows_sel, :1].sort_index(), pdf.iloc[rows_sel, :1].sort_index() + ) + + def test_frame_loc_setitem(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.from_pandas(pdf) + + pser1 = pdf.max_speed + pser2 = pdf.shield + psser1 = psdf.max_speed + psser2 = psdf.shield + + pdf.loc[["viper", "sidewinder"], ["shield", "max_speed"]] = 10 + psdf.loc[["viper", "sidewinder"], ["shield", "max_speed"]] = 10 + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc[["viper", "sidewinder"], "shield"] = 50 + psdf.loc[["viper", "sidewinder"], "shield"] = 50 + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc["cobra", "max_speed"] = 30 + psdf.loc["cobra", "max_speed"] = 30 + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc[pdf.max_speed < 5, "max_speed"] = -pdf.max_speed + psdf.loc[psdf.max_speed < 5, "max_speed"] = -psdf.max_speed + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc[pdf.max_speed < 2, "max_speed"] = -pdf.max_speed + psdf.loc[psdf.max_speed < 2, "max_speed"] = -psdf.max_speed + self.assert_eq(psdf, pdf) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc[:, "min_speed"] = 0 + psdf.loc[:, "min_speed"] = 0 + self.assert_eq(psdf, pdf, almost=True) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + with self.assertRaisesRegex(ValueError, "Incompatible indexer with Series"): + psdf.loc["cobra", "max_speed"] = -psdf.max_speed + with self.assertRaisesRegex(ValueError, "shape mismatch"): + psdf.loc[:, ["shield", "max_speed"]] = -psdf.max_speed + with self.assertRaisesRegex(ValueError, "Only a dataframe with one column can be assigned"): + psdf.loc[:, "max_speed"] = psdf + + # multi-index columns + columns = pd.MultiIndex.from_tuples( + [("x", "max_speed"), ("x", "shield"), ("y", "min_speed")] + ) + pdf.columns = columns + psdf.columns = columns + + pdf.loc[:, ("y", "shield")] = -pdf[("x", "shield")] + psdf.loc[:, ("y", "shield")] = -psdf[("x", "shield")] + self.assert_eq(psdf, pdf, almost=True) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + pdf.loc[:, "z"] = 100 + psdf.loc[:, "z"] = 100 + self.assert_eq(psdf, pdf, almost=True) + self.assert_eq(psser1, pser1) + self.assert_eq(psser2, pser2) + + with self.assertRaisesRegex(KeyError, "Key length \\(3\\) exceeds index depth \\(2\\)"): + psdf.loc[:, [("x", "max_speed", "foo")]] = -psdf[("x", "shield")] + + pdf = pd.DataFrame( + [[1], [4], [7]], index=["cobra", "viper", "sidewinder"], columns=["max_speed"] + ) + psdf = ps.from_pandas(pdf) + + pdf.loc[:, "max_speed"] = pdf + psdf.loc[:, "max_speed"] = psdf + self.assert_eq(psdf, pdf) + + def test_frame_iloc_setitem(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.from_pandas(pdf) + + pdf.iloc[[1, 2], [1, 0]] = 10 + psdf.iloc[[1, 2], [1, 0]] = 10 + self.assert_eq(psdf, pdf) + + pdf.iloc[0, 1] = 50 + psdf.iloc[0, 1] = 50 + self.assert_eq(psdf, pdf) + + with self.assertRaisesRegex(ValueError, "setting an array element with a sequence."): + psdf.iloc[0, 0] = -psdf.max_speed + with self.assertRaisesRegex(ValueError, "shape mismatch"): + psdf.iloc[:, [1, 0]] = -psdf.max_speed + with self.assertRaisesRegex(ValueError, "Only a dataframe with one column can be assigned"): + psdf.iloc[:, 0] = psdf + + pdf = pd.DataFrame( + [[1], [4], [7]], index=["cobra", "viper", "sidewinder"], columns=["max_speed"] + ) + psdf = ps.from_pandas(pdf) + + pdf.iloc[:, 0] = pdf + psdf.iloc[:, 0] = psdf + self.assert_eq(psdf, pdf) + + def test_series_loc_setitem(self): + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + + pser.loc[pser % 2 == 1] = -pser + psser.loc[psser % 2 == 1] = -psser + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + for key, value in [ + (["viper", "sidewinder"], 10), + ("viper", 50), + (slice(None), 10), + (slice(None, "viper"), 20), + (slice("viper", None), 30), + ]: + with self.subTest(key=key, value=value): + pser.loc[key] = value + psser.loc[key] = value + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaises(ValueError): + psser.loc["viper"] = -psser + + # multiindex + pser = pd.Series( + [1, 2, 3], + index=pd.MultiIndex.from_tuples([("x", "cobra"), ("x", "viper"), ("y", "sidewinder")]), + ) + psser = ps.from_pandas(pser) + + pser.loc["x"] = pser * 10 + psser.loc["x"] = psser * 10 + self.assert_eq(psser, pser) + + pser.loc["y"] = pser * 10 + psser.loc["y"] = psser * 10 + self.assert_eq(psser, pser) + + def test_series_iloc_setitem(self): + pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) + psdf = ps.from_pandas(pdf) + + pser = pdf.x + psery = pdf.y + psser = psdf.x + pssery = psdf.y + + piloc = pser.iloc + kiloc = psser.iloc + + pser1 = pser + 1 + psser1 = psser + 1 + + for key, value in [ + ([1, 2], 10), + (1, 50), + (slice(None), 10), + (slice(None, 1), 20), + (slice(1, None), 30), + ]: + with self.subTest(key=key, value=value): + pser.iloc[key] = value + psser.iloc[key] = value + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + piloc[key] = -value + kiloc[key] = -value + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + pser1.iloc[key] = value + psser1.iloc[key] = value + self.assert_eq(psser1, pser1) + self.assert_eq(psdf, pdf) + self.assert_eq(pssery, psery) + + with self.assertRaises(ValueError): + psser.iloc[1] = -psser + + pser = pd.Index([1, 2, 3]).to_series() + psser = ps.Index([1, 2, 3]).to_series() + + pser1 = pser + 1 + psser1 = psser + 1 + + pser.iloc[0] = 10 + psser.iloc[0] = 10 + self.assert_eq(psser, pser) + + pser1.iloc[0] = 20 + psser1.iloc[0] = 20 + self.assert_eq(psser1, pser1) + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + psdf = ps.from_pandas(pdf) + + pser = pdf.a + psser = psdf.a + + pser.iloc[[0, 1, 2]] = -pdf.b + psser.iloc[[0, 1, 2]] = -psdf.b + self.assert_eq(psser, pser) + self.assert_eq(psdf, pdf) + + with self.assertRaisesRegex(ValueError, "setting an array element with a sequence."): + psser.iloc[1] = psdf[["b"]] + + def test_iloc_raises(self): + pdf = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]}) + psdf = ps.from_pandas(pdf) + + with self.assertRaisesRegex(SparkPandasIndexingError, "Only accepts pairs of candidates"): + psdf.iloc[[0, 1], [0, 1], [1, 2]] + + with self.assertRaisesRegex(SparkPandasIndexingError, "Too many indexers"): + psdf.A.iloc[[0, 1], [0, 1]] + + with self.assertRaisesRegex(TypeError, "cannot do slice indexing with these indexers"): + psdf.iloc[:"b", :] + + with self.assertRaisesRegex(TypeError, "cannot do slice indexing with these indexers"): + psdf.iloc[:, :"b"] + + with self.assertRaisesRegex(TypeError, "cannot perform reduce with flexible type"): + psdf.iloc[:, ["A"]] + + with self.assertRaisesRegex(ValueError, "Location based indexing can only have"): + psdf.iloc[:, "A"] + + with self.assertRaisesRegex(IndexError, "out of range"): + psdf.iloc[:, [5, 6]] + + def test_index_operator_datetime(self): + dates = pd.date_range("20130101", periods=6) + pdf = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list("ABCD")) + psdf = ps.from_pandas(pdf) + + # Positional iloc search + self.assert_eq(psdf[:4], pdf[:4], almost=True) + self.assert_eq(psdf[:3], pdf[:3], almost=True) + self.assert_eq(psdf[3:], pdf[3:], almost=True) + self.assert_eq(psdf[2:], pdf[2:], almost=True) + self.assert_eq(psdf[2:3], pdf[2:3], almost=True) + self.assert_eq(psdf[2:-1], pdf[2:-1], almost=True) + self.assert_eq(psdf[10:3], pdf[10:3], almost=True) + + # Index loc search + self.assert_eq(psdf.A[4], pdf.A[4]) + self.assert_eq(psdf.A[3], pdf.A[3]) + + # Positional iloc search + self.assert_eq(psdf.A[:4], pdf.A[:4], almost=True) + self.assert_eq(psdf.A[:3], pdf.A[:3], almost=True) + self.assert_eq(psdf.A[3:], pdf.A[3:], almost=True) + self.assert_eq(psdf.A[2:], pdf.A[2:], almost=True) + self.assert_eq(psdf.A[2:3], pdf.A[2:3], almost=True) + self.assert_eq(psdf.A[2:-1], pdf.A[2:-1], almost=True) + self.assert_eq(psdf.A[10:3], pdf.A[10:3], almost=True) + + dt1 = datetime.datetime.strptime("2013-01-02", "%Y-%m-%d") + dt2 = datetime.datetime.strptime("2013-01-04", "%Y-%m-%d") + + # Index loc search + self.assert_eq(psdf[:dt2], pdf[:dt2], almost=True) + self.assert_eq(psdf[dt1:], pdf[dt1:], almost=True) + self.assert_eq(psdf[dt1:dt2], pdf[dt1:dt2], almost=True) + self.assert_eq(psdf.A[dt2], pdf.A[dt2], almost=True) + self.assert_eq(psdf.A[:dt2], pdf.A[:dt2], almost=True) + self.assert_eq(psdf.A[dt1:], pdf.A[dt1:], almost=True) + self.assert_eq(psdf.A[dt1:dt2], pdf.A[dt1:dt2], almost=True) + + def test_index_operator_int(self): + pdf = pd.DataFrame(np.random.randn(6, 4), index=[1, 3, 5, 7, 9, 11], columns=list("ABCD")) + psdf = ps.from_pandas(pdf) + + # Positional iloc search + self.assert_eq(psdf[:4], pdf[:4]) + self.assert_eq(psdf[:3], pdf[:3]) + self.assert_eq(psdf[3:], pdf[3:]) + self.assert_eq(psdf[2:], pdf[2:]) + self.assert_eq(psdf[2:3], pdf[2:3]) + self.assert_eq(psdf[2:-1], pdf[2:-1]) + self.assert_eq(psdf[10:3], pdf[10:3]) + + # Index loc search + self.assert_eq(psdf.A[5], pdf.A[5]) + self.assert_eq(psdf.A[3], pdf.A[3]) + with self.assertRaisesRegex( + NotImplementedError, "Duplicated row selection is not currently supported" + ): + psdf.iloc[[1, 1]] + + +class BasicIndexingTests(BasicIndexingTestsMixin, ComparisonTestBase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_indexing import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexops_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexops_spark.py new file mode 100644 index 0000000..3f7691a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_indexops_spark.py @@ -0,0 +1,82 @@ +# +# 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 pandas as pd +from pyspark.errors import AnalysisException +from pyspark.sql import functions as F + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SparkIndexOpsMethodsTestsMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + def test_series_transform_negative(self): + with self.assertRaisesRegex( + ValueError, "The output of the function.* pyspark.sql.Column.*int" + ): + self.psser.spark.transform(lambda scol: 1) + + with self.assertRaisesRegex(AnalysisException, ".*UNRESOLVED_COLUMN.*`non-existent`.*"): + self.psser.spark.transform(lambda scol: F.col("non-existent")) + + def test_multiindex_transform_negative(self): + with self.assertRaisesRegex( + NotImplementedError, "MultiIndex does not support spark.transform yet" + ): + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [1, 1, 1, 1, 1, 2, 1, 2, 2]], + ) + s = ps.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + s.index.spark.transform(lambda scol: scol) + + def test_series_apply_negative(self): + with self.assertRaisesRegex( + ValueError, "The output of the function.* pyspark.sql.Column.*int" + ): + self.psser.spark.apply(lambda scol: 1) + + with self.assertRaisesRegex(AnalysisException, ".*UNRESOLVED_COLUMN.*`non-existent`.*"): + self.psser.spark.transform(lambda scol: F.col("non-existent")) + + +class SparkIndexOpsMethodsTests( + SparkIndexOpsMethodsTestsMixin, PandasOnSparkTestCase, SQLTestUtils +): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_indexops_spark import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_internal.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_internal.py new file mode 100644 index 0000000..5a936d2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_internal.py @@ -0,0 +1,124 @@ +# +# 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 pandas as pd + +from pyspark.pandas.internal import ( + InternalFrame, + SPARK_DEFAULT_INDEX_NAME, + SPARK_INDEX_NAME_FORMAT, +) +from pyspark.pandas.utils import spark_column_equals +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class InternalFrameTestsMixin: + def test_from_pandas(self): + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + + internal = InternalFrame.from_pandas(pdf) + sdf = internal.spark_frame + + self.assert_eq(internal.index_spark_column_names, [SPARK_DEFAULT_INDEX_NAME]) + self.assert_eq(internal.index_names, [None]) + self.assert_eq(internal.column_labels, [("a",), ("b",)]) + self.assert_eq(internal.data_spark_column_names, ["a", "b"]) + self.assertTrue(spark_column_equals(internal.spark_column_for(("a",)), sdf["a"])) + self.assertTrue(spark_column_equals(internal.spark_column_for(("b",)), sdf["b"])) + + self.assert_eq(internal.to_pandas_frame, pdf) + + # non-string column name + pdf1 = pd.DataFrame({0: [1, 2, 3], 1: [4, 5, 6]}) + + internal = InternalFrame.from_pandas(pdf1) + sdf = internal.spark_frame + + self.assert_eq(internal.index_spark_column_names, [SPARK_DEFAULT_INDEX_NAME]) + self.assert_eq(internal.index_names, [None]) + self.assert_eq(internal.column_labels, [(0,), (1,)]) + self.assert_eq(internal.data_spark_column_names, ["0", "1"]) + self.assertTrue(spark_column_equals(internal.spark_column_for((0,)), sdf["0"])) + self.assertTrue(spark_column_equals(internal.spark_column_for((1,)), sdf["1"])) + + self.assert_eq(internal.to_pandas_frame, pdf1) + + # categorical column + pdf2 = pd.DataFrame({0: [1, 2, 3], 1: pd.Categorical([4, 5, 6])}) + internal = InternalFrame.from_pandas(pdf2) + sdf = internal.spark_frame + + self.assert_eq(internal.index_spark_column_names, [SPARK_DEFAULT_INDEX_NAME]) + self.assert_eq(internal.index_names, [None]) + self.assert_eq(internal.column_labels, [(0,), (1,)]) + self.assert_eq(internal.data_spark_column_names, ["0", "1"]) + self.assertTrue(spark_column_equals(internal.spark_column_for((0,)), sdf["0"])) + self.assertTrue(spark_column_equals(internal.spark_column_for((1,)), sdf["1"])) + + self.assert_eq(internal.to_pandas_frame, pdf2) + + # multi-index + pdf.set_index("a", append=True, inplace=True) + + internal = InternalFrame.from_pandas(pdf) + sdf = internal.spark_frame + + self.assert_eq( + internal.index_spark_column_names, + [SPARK_INDEX_NAME_FORMAT(0), SPARK_INDEX_NAME_FORMAT(1)], + ) + self.assert_eq(internal.index_names, [None, ("a",)]) + self.assert_eq(internal.column_labels, [("b",)]) + self.assert_eq(internal.data_spark_column_names, ["b"]) + self.assertTrue(spark_column_equals(internal.spark_column_for(("b",)), sdf["b"])) + + self.assert_eq(internal.to_pandas_frame, pdf) + + # multi-index columns + pdf.columns = pd.MultiIndex.from_tuples([("x", "b")]) + + internal = InternalFrame.from_pandas(pdf) + sdf = internal.spark_frame + + self.assert_eq( + internal.index_spark_column_names, + [SPARK_INDEX_NAME_FORMAT(0), SPARK_INDEX_NAME_FORMAT(1)], + ) + self.assert_eq(internal.index_names, [None, ("a",)]) + self.assert_eq(internal.column_labels, [("x", "b")]) + self.assert_eq(internal.data_spark_column_names, ["(x, b)"]) + self.assertTrue(spark_column_equals(internal.spark_column_for(("x", "b")), sdf["(x, b)"])) + + self.assert_eq(internal.to_pandas_frame, pdf) + + +class InternalFrameTests(InternalFrameTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_internal import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_namespace.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_namespace.py new file mode 100644 index 0000000..64c58a7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_namespace.py @@ -0,0 +1,638 @@ +# +# 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 distutils.version import LooseVersion +import itertools +import inspect +import unittest + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.namespace import _get_index_map, read_delta +from pyspark.pandas.utils import spark_column_equals +from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class NamespaceTestsMixin: + def test_from_pandas(self): + pdf = pd.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf, pdf) + + pser = pdf.year + psser = ps.from_pandas(pser) + + self.assert_eq(psser, pser) + + pidx = pdf.index + psidx = ps.from_pandas(pidx) + + self.assert_eq(psidx, pidx) + + pmidx = pdf.set_index("year", append=True).index + psmidx = ps.from_pandas(pmidx) + + self.assert_eq(psmidx, pmidx) + + expected_error_message = "Unknown data type: {}".format(type(psidx).__name__) + with self.assertRaisesRegex(TypeError, expected_error_message): + ps.from_pandas(psidx) + + def test_to_datetime(self): + pdf = pd.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + self.assert_eq(pd.to_datetime(1490195805, unit="s"), ps.to_datetime(1490195805, unit="s")) + self.assert_eq( + pd.to_datetime(1490195805433502912, unit="ns"), + ps.to_datetime(1490195805433502912, unit="ns"), + ) + + self.assert_eq( + pd.to_datetime([1, 2, 3], unit="D", origin=pd.Timestamp("1960-01-01")), + ps.to_datetime([1, 2, 3], unit="D", origin=pd.Timestamp("1960-01-01")), + ) + + pdf = pd.DataFrame({"years": [2015, 2016], "month": [2, 3], "day": [4, 5]}) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame({"years": [2015, 2016], "months": [2, 3], "day": [4, 5]}) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame({"years": [2015, 2016], "months": [2, 3], "days": [4, 5]}) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + # SPARK-36946: Support time for ps.to_datetime + pdf = pd.DataFrame( + { + "year": [2015, 2016], + "month": [2, 3], + "day": [4, 5], + "hour": [2, 3], + "minute": [10, 30], + "second": [21, 25], + } + ) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame( + { + "year": [2015, 2016], + "month": [2, 3], + "day": [4, 5], + "hour": [2, 3], + "minute": [10, 30], + } + ) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5], "hour": [2, 3]}) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame( + { + "year": [2015, 2016], + "month": [2, 3], + "day": [4, 5], + "hour": [2, 3], + "minute": [10, 30], + "second": [21, 25], + "ms": [50, 69], + } + ) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame( + { + "year": [2015, 2016], + "month": [2, 3], + "day": [4, 5], + "hour": [2, 3], + "minute": [10, 30], + "second": [21, 25], + "ms": [50, 69], + "millisecond": [123, 678], + } + ) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + pdf = pd.DataFrame( + { + "Year": [2015, 2016], + "Month": [2, 3], + "Day": [4, 5], + "Hour": [2, 3], + "Minute": [10, 30], + "Second": [21, 25], + "ms": [50, 69], + "millisecond": [123, 678], + } + ) + psdf = ps.from_pandas(pdf) + dict_from_pdf = pdf.to_dict() + + self.assert_eq(pd.to_datetime(pdf), ps.to_datetime(psdf)) + self.assert_eq(pd.to_datetime(dict_from_pdf), ps.to_datetime(dict_from_pdf)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43709): Enable NamespaceTests.test_date_range for pandas 2.0.0.", + ) + def test_date_range(self): + self.assert_eq( + ps.date_range(start="1/1/2018", end="1/08/2018"), + pd.date_range(start="1/1/2018", end="1/08/2018"), + ) + self.assert_eq( + ps.date_range(start="1/1/2018", periods=8), pd.date_range(start="1/1/2018", periods=8) + ) + self.assert_eq( + ps.date_range(end="1/1/2018", periods=8), pd.date_range(end="1/1/2018", periods=8) + ) + self.assert_eq( + ps.date_range(start="2018-04-24", end="2018-04-27", periods=3), + pd.date_range(start="2018-04-24", end="2018-04-27", periods=3), + ) + + self.assert_eq( + ps.date_range(start="1/1/2018", periods=5, freq="M"), + pd.date_range(start="1/1/2018", periods=5, freq="M"), + ) + + self.assert_eq( + ps.date_range(start="1/1/2018", periods=5, freq="3M"), + pd.date_range(start="1/1/2018", periods=5, freq="3M"), + ) + + self.assert_eq( + ps.date_range(start="1/1/2018", periods=5, freq=pd.offsets.MonthEnd(3)), + pd.date_range(start="1/1/2018", periods=5, freq=pd.offsets.MonthEnd(3)), + ) + + self.assert_eq( + ps.date_range(start="2017-01-01", end="2017-01-04", closed="left"), + pd.date_range(start="2017-01-01", end="2017-01-04", closed="left"), + ) + + self.assert_eq( + ps.date_range(start="2017-01-01", end="2017-01-04", closed="right"), + pd.date_range(start="2017-01-01", end="2017-01-04", closed="right"), + ) + + self.assertRaises( + AssertionError, lambda: ps.date_range(start="1/1/2018", periods=5, tz="Asia/Tokyo") + ) + self.assertRaises( + AssertionError, lambda: ps.date_range(start="1/1/2018", periods=5, freq="ns") + ) + self.assertRaises( + AssertionError, lambda: ps.date_range(start="1/1/2018", periods=5, freq="N") + ) + + def test_to_timedelta(self): + self.assert_eq( + ps.to_timedelta("1 days 06:05:01.00003"), + pd.to_timedelta("1 days 06:05:01.00003"), + ) + self.assert_eq( + ps.to_timedelta("15.5us"), + pd.to_timedelta("15.5us"), + ) + self.assert_eq( + ps.to_timedelta(["1 days 06:05:01.00003", "15.5us", "nan"]), + pd.to_timedelta(["1 days 06:05:01.00003", "15.5us", "nan"]), + ) + self.assert_eq( + ps.to_timedelta(np.arange(5), unit="s"), + pd.to_timedelta(np.arange(5), unit="s"), + ) + self.assert_eq( + ps.to_timedelta(ps.Series([1, 2]), unit="d"), + pd.to_timedelta(pd.Series([1, 2]), unit="d"), + ) + self.assert_eq( + ps.to_timedelta(pd.Series([1, 2]), unit="d"), + pd.to_timedelta(pd.Series([1, 2]), unit="d"), + ) + + def test_timedelta_range(self): + self.assert_eq( + ps.timedelta_range(start="1 day", end="3 days"), + pd.timedelta_range(start="1 day", end="3 days"), + ) + self.assert_eq( + ps.timedelta_range(start="1 day", periods=3), + pd.timedelta_range(start="1 day", periods=3), + ) + self.assert_eq( + ps.timedelta_range(end="3 days", periods=3), + pd.timedelta_range(end="3 days", periods=3), + ) + self.assert_eq( + ps.timedelta_range(end="3 days", periods=3, closed="right"), + pd.timedelta_range(end="3 days", periods=3, closed="right"), + ) + self.assert_eq( + ps.timedelta_range(start="1 day", end="3 days", freq="6H"), + pd.timedelta_range(start="1 day", end="3 days", freq="6H"), + ) + self.assert_eq( + ps.timedelta_range(start="1 day", end="3 days", periods=4), + pd.timedelta_range(start="1 day", end="3 days", periods=4), + ) + + self.assertRaises( + AssertionError, lambda: ps.timedelta_range(start="1 day", periods=3, freq="ns") + ) + + def test_concat_multiindex_sort(self): + # SPARK-39314: Respect ps.concat sort parameter to follow pandas behavior + idx = pd.MultiIndex.from_tuples([("Y", "A"), ("Y", "B"), ("X", "C"), ("X", "D")]) + pdf = pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], columns=idx) + psdf = ps.from_pandas(pdf) + + ignore_indexes = [True, False] + joins = ["inner", "outer"] + sorts = [True] + if LooseVersion(pd.__version__) >= LooseVersion("1.4"): + sorts += [False] + objs = [ + ([psdf, psdf.reset_index()], [pdf, pdf.reset_index()]), + ([psdf.reset_index(), psdf], [pdf.reset_index(), pdf]), + ] + for ignore_index, join, sort in itertools.product(ignore_indexes, joins, sorts): + for i, (psdfs, pdfs) in enumerate(objs): + self.assert_eq( + ps.concat(psdfs, ignore_index=ignore_index, join=join, sort=sort), + pd.concat(pdfs, ignore_index=ignore_index, join=join, sort=sort), + ) + + def test_concat_index_axis(self): + pdf = pd.DataFrame({"A": [0, 2, 4], "B": [1, 3, 5], "C": [6, 7, 8]}) + # TODO: pdf.columns.names = ["ABC"] + psdf = ps.from_pandas(pdf) + + ignore_indexes = [True, False] + joins = ["inner", "outer"] + sorts = [True, False] + + objs = [ + ([psdf, psdf], [pdf, pdf]), + # no Series + ([psdf, psdf.reset_index()], [pdf, pdf.reset_index()]), + ([psdf.reset_index(), psdf], [pdf.reset_index(), pdf]), + ([psdf, psdf[["C", "A"]]], [pdf, pdf[["C", "A"]]]), + ([psdf[["C", "A"]], psdf], [pdf[["C", "A"]], pdf]), + # more than two Series + ([psdf["C"], psdf, psdf["A"]], [pdf["C"], pdf, pdf["A"]]), + # only Series + ([psdf["C"], psdf["A"]], [pdf["C"], pdf["A"]]), + ] + + # See also https://github.com/pandas-dev/pandas/issues/47127 + if LooseVersion(pd.__version__) >= LooseVersion("1.4.3"): + series_objs = [ + # more than two Series + ([psdf, psdf["C"], psdf["A"]], [pdf, pdf["C"], pdf["A"]]), + # only one Series + ([psdf, psdf["C"]], [pdf, pdf["C"]]), + ([psdf["C"], psdf], [pdf["C"], pdf]), + ] + for psdfs, pdfs in series_objs: + for ignore_index, join, sort in itertools.product(ignore_indexes, joins, sorts): + self.assert_eq( + ps.concat(psdfs, ignore_index=ignore_index, join=join, sort=sort), + pd.concat(pdfs, ignore_index=ignore_index, join=join, sort=sort), + ) + + for ignore_index, join, sort in itertools.product(ignore_indexes, joins, sorts): + for i, (psdfs, pdfs) in enumerate(objs): + with self.subTest( + ignore_index=ignore_index, join=join, sort=sort, pdfs=pdfs, pair=i + ): + self.assert_eq( + ps.concat(psdfs, ignore_index=ignore_index, join=join, sort=sort), + pd.concat(pdfs, ignore_index=ignore_index, join=join, sort=sort), + almost=(join == "outer"), + ) + + self.assertRaisesRegex(TypeError, "first argument must be", lambda: ps.concat(psdf)) + self.assertRaisesRegex(TypeError, "cannot concatenate object", lambda: ps.concat([psdf, 1])) + + psdf2 = psdf.set_index("B", append=True) + self.assertRaisesRegex( + ValueError, "Index type and names should be same", lambda: ps.concat([psdf, psdf2]) + ) + + self.assertRaisesRegex(ValueError, "No objects to concatenate", lambda: ps.concat([])) + + self.assertRaisesRegex(ValueError, "All objects passed", lambda: ps.concat([None, None])) + + pdf3 = pdf.copy() + psdf3 = psdf.copy() + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) + # TODO: colums.names = ["XYZ", "ABC"] + pdf3.columns = columns + psdf3.columns = columns + + objs = [ + ([psdf3, psdf3], [pdf3, pdf3]), + ([psdf3, psdf3.reset_index()], [pdf3, pdf3.reset_index()]), + ([psdf3, psdf3[[("Y", "C"), ("X", "A")]]], [pdf3, pdf3[[("Y", "C"), ("X", "A")]]]), + ] + + if LooseVersion(pd.__version__) >= LooseVersion("1.4"): + objs += [ + ([psdf3.reset_index(), psdf3], [pdf3.reset_index(), pdf3]), + ([psdf3[[("Y", "C"), ("X", "A")]], psdf3], [pdf3[[("Y", "C"), ("X", "A")]], pdf3]), + ] + + for ignore_index, sort in itertools.product(ignore_indexes, sorts): + for i, (psdfs, pdfs) in enumerate(objs): + with self.subTest( + ignore_index=ignore_index, join="outer", sort=sort, pdfs=pdfs, pair=i + ): + self.assert_eq( + ps.concat(psdfs, ignore_index=ignore_index, join="outer", sort=sort), + pd.concat(pdfs, ignore_index=ignore_index, join="outer", sort=sort), + ) + + # Skip tests for `join="inner" and sort=False` since pandas is flaky. + for ignore_index in ignore_indexes: + for i, (psdfs, pdfs) in enumerate(objs): + with self.subTest( + ignore_index=ignore_index, join="inner", sort=True, pdfs=pdfs, pair=i + ): + self.assert_eq( + ps.concat(psdfs, ignore_index=ignore_index, join="inner", sort=True), + pd.concat(pdfs, ignore_index=ignore_index, join="inner", sort=True), + ) + + self.assertRaisesRegex( + ValueError, + "MultiIndex columns should have the same levels", + lambda: ps.concat([psdf, psdf3]), + ) + self.assertRaisesRegex( + ValueError, + "MultiIndex columns should have the same levels", + lambda: ps.concat([psdf3[("Y", "C")], psdf3]), + ) + + pdf4 = pd.DataFrame({"A": [0, 2, 4], "B": [1, 3, 5], "C": [10, 20, 30]}) + psdf4 = ps.from_pandas(pdf4) + self.assertRaisesRegex( + ValueError, + r"Only can inner \(intersect\) or outer \(union\) join the other axis.", + lambda: ps.concat([psdf, psdf4], join=""), + ) + + self.assertRaisesRegex( + ValueError, + r"Only can inner \(intersect\) or outer \(union\) join the other axis.", + lambda: ps.concat([psdf, psdf4], join="", axis=1), + ) + + self.assertRaisesRegex( + ValueError, + r"Only can inner \(intersect\) or outer \(union\) join the other axis.", + lambda: ps.concat([psdf.A, psdf4.B], join="", axis=1), + ) + + self.assertRaisesRegex( + ValueError, + r"Labels have to be unique; however, got duplicated labels \['A'\].", + lambda: ps.concat([psdf.A, psdf4.A], join="inner", axis=1), + ) + + def test_concat_column_axis(self): + pdf1 = pd.DataFrame({"A": [0, 2, 4], "B": [1, 3, 5]}, index=[1, 2, 3]) + pdf1.columns.names = ["AB"] + pdf2 = pd.DataFrame({"C": [1, 2, 3], "D": [4, 5, 6]}, index=[1, 3, 5]) + pdf2.columns.names = ["CD"] + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + psdf3 = psdf1.copy() + psdf4 = psdf2.copy() + pdf3 = pdf1.copy() + pdf4 = pdf2.copy() + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")], names=["X", "AB"]) + pdf3.columns = columns + psdf3.columns = columns + + columns = pd.MultiIndex.from_tuples([("X", "C"), ("X", "D")], names=["Y", "CD"]) + pdf4.columns = columns + psdf4.columns = columns + + ignore_indexes = [True, False] + joins = ["inner", "outer"] + + objs = [ + ([psdf1.A, psdf1.A.rename("B")], [pdf1.A, pdf1.A.rename("B")]), + ( + [psdf3[("X", "A")], psdf3[("X", "B")]], + [pdf3[("X", "A")], pdf3[("X", "B")]], + ), + ( + [psdf3[("X", "A")], psdf3[("X", "B")].rename("ABC")], + [pdf3[("X", "A")], pdf3[("X", "B")].rename("ABC")], + ), + ( + [psdf3[("X", "A")].rename("ABC"), psdf3[("X", "B")]], + [pdf3[("X", "A")].rename("ABC"), pdf3[("X", "B")]], + ), + ] + + for ignore_index, join in itertools.product(ignore_indexes, joins): + for i, (psdfs, pdfs) in enumerate(objs): + with self.subTest(ignore_index=ignore_index, join=join, pdfs=pdfs, pair=i): + actual = ps.concat(psdfs, axis=1, ignore_index=ignore_index, join=join) + expected = pd.concat(pdfs, axis=1, ignore_index=ignore_index, join=join) + self.assert_eq( + repr(actual.sort_values(list(actual.columns)).reset_index(drop=True)), + repr(expected.sort_values(list(expected.columns)).reset_index(drop=True)), + ) + + # test dataframes equality with broadcast hint. + def test_broadcast(self): + psdf = ps.DataFrame( + {"key": ["K0", "K1", "K2", "K3"], "A": ["A0", "A1", "A2", "A3"]}, columns=["key", "A"] + ) + self.assert_eq(psdf, ps.broadcast(psdf)) + + psdf.columns = ["x", "y"] + self.assert_eq(psdf, ps.broadcast(psdf)) + + psdf.columns = [("a", "c"), ("b", "d")] + self.assert_eq(psdf, ps.broadcast(psdf)) + + psser = ps.Series([1, 2, 3]) + expected_error_message = "Invalid type : expected DataFrame got {}".format( + type(psser).__name__ + ) + with self.assertRaisesRegex(TypeError, expected_error_message): + ps.broadcast(psser) + + def test_get_index_map(self): + psdf = ps.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}) + sdf = psdf.to_spark() + self.assertEqual(_get_index_map(sdf), (None, None)) + + def check(actual, expected): + actual_scols, actual_labels = actual + expected_column_names, expected_labels = expected + self.assertEqual(len(actual_scols), len(expected_column_names)) + for actual_scol, expected_column_name in zip(actual_scols, expected_column_names): + expected_scol = sdf[expected_column_name] + self.assertTrue(spark_column_equals(actual_scol, expected_scol)) + self.assertEqual(actual_labels, expected_labels) + + check(_get_index_map(sdf, "year"), (["year"], [("year",)])) + check(_get_index_map(sdf, ["year", "month"]), (["year", "month"], [("year",), ("month",)])) + + self.assertRaises(KeyError, lambda: _get_index_map(sdf, ["year", "hour"])) + + def test_read_delta_with_wrong_input(self): + self.assertRaisesRegex( + ValueError, + "version and timestamp cannot be used together", + lambda: read_delta("fake_path", version="0", timestamp="2021-06-22"), + ) + + def test_to_numeric(self): + pser = pd.Series(["1", "2", None, "4", "hello"]) + psser = ps.from_pandas(pser) + + # "coerce" and "raise" with Series that contains un-parsable data. + self.assert_eq( + pd.to_numeric(pser, errors="coerce"), ps.to_numeric(psser, errors="coerce"), almost=True + ) + + # "raise" with Series that contains parsable data only. + pser = pd.Series(["1", "2", None, "4", "5.0"]) + psser = ps.from_pandas(pser) + + self.assert_eq( + pd.to_numeric(pser, errors="raise"), ps.to_numeric(psser, errors="raise"), almost=True + ) + + # "coerce", "ignore" and "raise" with non-Series. + data = ["1", "2", None, "4", "hello"] + self.assert_eq(pd.to_numeric(data, errors="coerce"), ps.to_numeric(data, errors="coerce")) + self.assert_eq(pd.to_numeric(data, errors="ignore"), ps.to_numeric(data, errors="ignore")) + + self.assertRaisesRegex( + ValueError, + 'Unable to parse string "hello"', + lambda: ps.to_numeric(data, errors="raise"), + ) + + # "raise" with non-Series that contains parsable data only. + data = ["1", "2", None, "4", "5.0"] + + self.assert_eq( + pd.to_numeric(data, errors="raise"), ps.to_numeric(data, errors="raise"), almost=True + ) + + # Wrong string for `errors` parameter. + self.assertRaisesRegex( + ValueError, + "invalid error value specified", + lambda: ps.to_numeric(psser, errors="errors"), + ) + # NotImplementedError + self.assertRaisesRegex( + NotImplementedError, + "'ignore' is not implemented yet, when the `arg` is Series.", + lambda: ps.to_numeric(psser, errors="ignore"), + ) + + def test_missing(self): + missing_functions = inspect.getmembers( + MissingPandasLikeGeneralFunctions, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "The method.*pd.*{}.*not implemented yet.".format(name), + ): + getattr(ps, name)() + + +class NamespaceTests(NamespaceTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_namespace import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_numpy_compat.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_numpy_compat.py new file mode 100644 index 0000000..e849932 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_numpy_compat.py @@ -0,0 +1,200 @@ +# +# 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 numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas import set_option, reset_option +from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings +from pyspark.testing.pandasutils import ComparisonTestBase +from pyspark.testing.sqlutils import SQLTestUtils + + +class NumPyCompatTestsMixin: + blacklist = [ + # Koalas does not currently support + "conj", + "conjugate", + "isnat", + "matmul", + "frexp", + # Values are close enough but tests failed. + "arccos", + "exp", + "expm1", + "log", # flaky + "log10", # flaky + "log1p", # flaky + "modf", + "floor_divide", # flaky + # Results seem inconsistent in a different version of, I (Hyukjin) suspect, PyArrow. + # From PyArrow 0.15, seems it returns the correct results via PySpark. Probably we + # can enable it later when Koalas switches to PyArrow 0.15 completely. + "left_shift", + ] + + @property + def pdf(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + def test_np_add_series(self): + psdf = self.psdf + pdf = self.pdf + + self.assert_eq(np.add(psdf.a, psdf.b), np.add(pdf.a, pdf.b)) + + psdf = self.psdf + pdf = self.pdf + self.assert_eq(np.add(psdf.a, 1), np.add(pdf.a, 1)) + + def test_np_add_index(self): + k_index = self.psdf.index + p_index = self.pdf.index + self.assert_eq(np.add(k_index, k_index), np.add(p_index, p_index)) + + def test_np_unsupported_series(self): + psdf = self.psdf + with self.assertRaisesRegex(NotImplementedError, "pandas.*not.*support.*sqrt.*"): + np.sqrt(psdf.a, psdf.b) + + def test_np_unsupported_frame(self): + psdf = self.psdf + with self.assertRaisesRegex(NotImplementedError, "on-Spark.*not.*support.*sqrt.*"): + np.sqrt(psdf, psdf) + + psdf1 = ps.DataFrame({"A": [1, 2, 3]}) + psdf2 = ps.DataFrame({("A", "B"): [4, 5, 6]}) + with self.assertRaisesRegex(ValueError, "cannot join with no overlapping index names"): + np.left_shift(psdf1, psdf2) + + def test_np_spark_compat_series(self): + # Use randomly generated dataFrame + pdf = pd.DataFrame( + np.random.randint(-100, 100, size=(np.random.randint(100), 2)), columns=["a", "b"] + ) + pdf2 = pd.DataFrame( + np.random.randint(-100, 100, size=(len(pdf), len(pdf.columns))), columns=["a", "b"] + ) + psdf = ps.from_pandas(pdf) + psdf2 = ps.from_pandas(pdf2) + + for np_name, spark_func in unary_np_spark_mappings.items(): + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # unary ufunc + self.assert_eq(np_func(pdf.a), np_func(psdf.a), almost=True) + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + + for np_name, spark_func in binary_np_spark_mappings.items(): + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # binary ufunc + self.assert_eq(np_func(pdf.a, pdf.b), np_func(psdf.a, psdf.b), almost=True) + self.assert_eq(np_func(pdf.a, 1), np_func(psdf.a, 1), almost=True) + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + + # Test only top 5 for now. 'compute.ops_on_diff_frames' option increases too much time. + try: + set_option("compute.ops_on_diff_frames", True) + for np_name, spark_func in list(binary_np_spark_mappings.items())[:5]: + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # binary ufunc + self.assert_eq( + np_func(pdf.a, pdf2.b).sort_index(), + np_func(psdf.a, psdf2.b).sort_index(), + almost=True, + ) + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + finally: + reset_option("compute.ops_on_diff_frames") + + def test_np_spark_compat_frame(self): + # Use randomly generated dataFrame + pdf = pd.DataFrame( + np.random.randint(-100, 100, size=(np.random.randint(100), 2)), columns=["a", "b"] + ) + pdf2 = pd.DataFrame( + np.random.randint(-100, 100, size=(len(pdf), len(pdf.columns))), columns=["a", "b"] + ) + psdf = ps.from_pandas(pdf) + psdf2 = ps.from_pandas(pdf2) + + for np_name, spark_func in unary_np_spark_mappings.items(): + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # unary ufunc + self.assert_eq(np_func(pdf), np_func(psdf), almost=True) + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + + for np_name, spark_func in binary_np_spark_mappings.items(): + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # binary ufunc + self.assert_eq(np_func(pdf, pdf), np_func(psdf, psdf), almost=True) + self.assert_eq(np_func(pdf, 1), np_func(psdf, 1), almost=True) + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + + # Test only top 5 for now. 'compute.ops_on_diff_frames' option increases too much time. + try: + set_option("compute.ops_on_diff_frames", True) + for np_name, spark_func in list(binary_np_spark_mappings.items())[:5]: + np_func = getattr(np, np_name) + if np_name not in self.blacklist: + try: + # binary ufunc + self.assert_eq( + np_func(pdf, pdf2).sort_index(), + np_func(psdf, psdf2).sort_index(), + almost=True, + ) + + except Exception as e: + raise AssertionError("Test in '%s' function was failed." % np_name) from e + finally: + reset_option("compute.ops_on_diff_frames") + + +class NumPyCompatTests(NumPyCompatTestsMixin, ComparisonTestBase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_numpy_compat import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames.py new file mode 100644 index 0000000..3d25788 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames.py @@ -0,0 +1,1355 @@ +# +# 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 distutils.version import LooseVersion +from itertools import product +import unittest + +import pandas as pd +import numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.frame import DataFrame +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.pandas.typedef.typehints import ( + extension_dtypes, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, +) + + +class OpsOnDiffFramesEnabledTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 10, 11], + ) + + @property + def pdf2(self): + return pd.DataFrame( + {"a": [9, 8, 7, 6, 5, 4, 3, 2, 1], "b": [0, 0, 0, 4, 5, 6, 1, 2, 3]}, + index=list(range(9)), + ) + + @property + def pdf3(self): + return pd.DataFrame( + {"b": [1, 1, 1, 1, 1, 1, 1, 1, 1], "c": [1, 1, 1, 1, 1, 1, 1, 1, 1]}, + index=list(range(9)), + ) + + @property + def pdf4(self): + return pd.DataFrame( + {"e": [2, 2, 2, 2, 2, 2, 2, 2, 2], "f": [2, 2, 2, 2, 2, 2, 2, 2, 2]}, + index=list(range(9)), + ) + + @property + def pdf5(self): + return pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7, 8, 9], + "b": [4, 5, 6, 3, 2, 1, 0, 0, 0], + "c": [4, 5, 6, 3, 2, 1, 0, 0, 0], + }, + index=[0, 1, 3, 5, 6, 8, 9, 10, 11], + ).set_index(["a", "b"]) + + @property + def pdf6(self): + return pd.DataFrame( + { + "a": [9, 8, 7, 6, 5, 4, 3, 2, 1], + "b": [0, 0, 0, 4, 5, 6, 1, 2, 3], + "c": [9, 8, 7, 6, 5, 4, 3, 2, 1], + "e": [4, 5, 6, 3, 2, 1, 0, 0, 0], + }, + index=list(range(9)), + ).set_index(["a", "b"]) + + @property + def pser1(self): + midx = pd.MultiIndex( + [["lama", "cow", "falcon", "koala"], ["speed", "weight", "length", "power"]], + [[0, 3, 1, 1, 1, 2, 2, 2], [0, 2, 0, 3, 2, 0, 1, 3]], + ) + return pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1], index=midx) + + @property + def pser2(self): + midx = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + return pd.Series([-45, 200, -1.2, 30, -250, 1.5, 320, 1, -0.3], index=midx) + + @property + def pser3(self): + midx = pd.MultiIndex( + [["koalas", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [1, 1, 2, 0, 0, 2, 2, 2, 1]], + ) + return pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.from_pandas(self.pdf2) + + @property + def psdf3(self): + return ps.from_pandas(self.pdf3) + + @property + def psdf4(self): + return ps.from_pandas(self.pdf4) + + @property + def psdf5(self): + return ps.from_pandas(self.pdf5) + + @property + def psdf6(self): + return ps.from_pandas(self.pdf6) + + @property + def psser1(self): + return ps.from_pandas(self.pser1) + + @property + def psser2(self): + return ps.from_pandas(self.pser2) + + @property + def psser3(self): + return ps.from_pandas(self.pser3) + + def test_ranges(self): + self.assert_eq( + (ps.range(10) + ps.range(10)).sort_index(), + ( + ps.DataFrame({"id": list(range(10))}) + ps.DataFrame({"id": list(range(10))}) + ).sort_index(), + ) + + def test_no_matched_index(self): + with self.assertRaisesRegex(ValueError, "Index names must be exactly matched"): + ps.DataFrame({"a": [1, 2, 3]}).set_index("a") + ps.DataFrame( + {"b": [1, 2, 3]} + ).set_index("b") + + def test_arithmetic(self): + self._test_arithmetic_frame(self.pdf1, self.pdf2, check_extension=False) + self._test_arithmetic_series(self.pser1, self.pser2, check_extension=False) + + @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") + def test_arithmetic_extension_dtypes(self): + self._test_arithmetic_frame( + self.pdf1.astype("Int64"), self.pdf2.astype("Int64"), check_extension=True + ) + self._test_arithmetic_series( + self.pser1.astype(int).astype("Int64"), + self.pser2.astype(int).astype("Int64"), + check_extension=True, + ) + + @unittest.skipIf( + not extension_float_dtypes_available, "pandas extension float dtypes are not available" + ) + def test_arithmetic_extension_float_dtypes(self): + self._test_arithmetic_frame( + self.pdf1.astype("Float64"), self.pdf2.astype("Float64"), check_extension=True + ) + self._test_arithmetic_series( + self.pser1.astype("Float64"), self.pser2.astype("Float64"), check_extension=True + ) + + def _test_arithmetic_frame(self, pdf1, pdf2, *, check_extension): + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + def assert_eq(actual, expected): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(actual, expected, check_exact=not check_extension) + if check_extension: + if isinstance(actual, DataFrame): + for dtype in actual.dtypes: + self.assertTrue(isinstance(dtype, extension_dtypes)) + else: + self.assertTrue(isinstance(actual.dtype, extension_dtypes)) + else: + self.assert_eq(actual, expected) + + # Series + assert_eq((psdf1.a - psdf2.b).sort_index(), (pdf1.a - pdf2.b).sort_index()) + + assert_eq((psdf1.a * psdf2.a).sort_index(), (pdf1.a * pdf2.a).sort_index()) + + if check_extension and not extension_float_dtypes_available: + self.assert_eq( + (psdf1["a"] / psdf2["a"]).sort_index(), (pdf1["a"] / pdf2["a"]).sort_index() + ) + else: + assert_eq((psdf1["a"] / psdf2["a"]).sort_index(), (pdf1["a"] / pdf2["a"]).sort_index()) + + # DataFrame + assert_eq((psdf1 + psdf2).sort_index(), (pdf1 + pdf2).sort_index()) + + # Multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b")]) + psdf1.columns = columns + psdf2.columns = columns + pdf1.columns = columns + pdf2.columns = columns + + # Series + assert_eq( + (psdf1[("x", "a")] - psdf2[("x", "b")]).sort_index(), + (pdf1[("x", "a")] - pdf2[("x", "b")]).sort_index(), + ) + + assert_eq( + (psdf1[("x", "a")] - psdf2["x"]["b"]).sort_index(), + (pdf1[("x", "a")] - pdf2["x"]["b"]).sort_index(), + ) + + assert_eq( + (psdf1["x"]["a"] - psdf2[("x", "b")]).sort_index(), + (pdf1["x"]["a"] - pdf2[("x", "b")]).sort_index(), + ) + + # DataFrame + assert_eq((psdf1 + psdf2).sort_index(), (pdf1 + pdf2).sort_index()) + + def _test_arithmetic_series(self, pser1, pser2, *, check_extension): + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + def assert_eq(actual, expected): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(actual, expected, check_exact=not check_extension) + if check_extension: + self.assertTrue(isinstance(actual.dtype, extension_dtypes)) + else: + self.assert_eq(actual, expected) + + # MultiIndex Series + assert_eq((psser1 + psser2).sort_index(), (pser1 + pser2).sort_index()) + + assert_eq((psser1 - psser2).sort_index(), (pser1 - pser2).sort_index()) + + assert_eq((psser1 * psser2).sort_index(), (pser1 * pser2).sort_index()) + + if check_extension and not extension_float_dtypes_available: + self.assert_eq((psser1 / psser2).sort_index(), (pser1 / pser2).sort_index()) + else: + assert_eq((psser1 / psser2).sort_index(), (pser1 / pser2).sort_index()) + + def test_arithmetic_chain(self): + self._test_arithmetic_chain_frame(self.pdf1, self.pdf2, self.pdf3, check_extension=False) + self._test_arithmetic_chain_series( + self.pser1, self.pser2, self.pser3, check_extension=False + ) + + @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") + def test_arithmetic_chain_extension_dtypes(self): + self._test_arithmetic_chain_frame( + self.pdf1.astype("Int64"), + self.pdf2.astype("Int64"), + self.pdf3.astype("Int64"), + check_extension=True, + ) + self._test_arithmetic_chain_series( + self.pser1.astype(int).astype("Int64"), + self.pser2.astype(int).astype("Int64"), + self.pser3.astype(int).astype("Int64"), + check_extension=True, + ) + + @unittest.skipIf( + not extension_float_dtypes_available, "pandas extension float dtypes are not available" + ) + def test_arithmetic_chain_extension_float_dtypes(self): + self._test_arithmetic_chain_frame( + self.pdf1.astype("Float64"), + self.pdf2.astype("Float64"), + self.pdf3.astype("Float64"), + check_extension=True, + ) + self._test_arithmetic_chain_series( + self.pser1.astype("Float64"), + self.pser2.astype("Float64"), + self.pser3.astype("Float64"), + check_extension=True, + ) + + def _test_arithmetic_chain_frame(self, pdf1, pdf2, pdf3, *, check_extension): + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + psdf3 = ps.from_pandas(pdf3) + + common_columns = set(psdf1.columns).intersection(psdf2.columns).intersection(psdf3.columns) + + def assert_eq(actual, expected): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(actual, expected, check_exact=not check_extension) + if check_extension: + if isinstance(actual, DataFrame): + for column, dtype in zip(actual.columns, actual.dtypes): + if column in common_columns: + self.assertTrue(isinstance(dtype, extension_dtypes)) + else: + self.assertFalse(isinstance(dtype, extension_dtypes)) + else: + self.assertTrue(isinstance(actual.dtype, extension_dtypes)) + else: + self.assert_eq(actual, expected) + + # Series + assert_eq( + (psdf1.a - psdf2.b - psdf3.c).sort_index(), (pdf1.a - pdf2.b - pdf3.c).sort_index() + ) + + assert_eq( + (psdf1.a * (psdf2.a * psdf3.c)).sort_index(), (pdf1.a * (pdf2.a * pdf3.c)).sort_index() + ) + + if check_extension and not extension_float_dtypes_available: + self.assert_eq( + (psdf1["a"] / psdf2["a"] / psdf3["c"]).sort_index(), + (pdf1["a"] / pdf2["a"] / pdf3["c"]).sort_index(), + ) + else: + assert_eq( + (psdf1["a"] / psdf2["a"] / psdf3["c"]).sort_index(), + (pdf1["a"] / pdf2["a"] / pdf3["c"]).sort_index(), + ) + + # DataFrame + if check_extension and LooseVersion(pd.__version__) < LooseVersion("1.1"): + self.assert_eq( + (psdf1 + psdf2 - psdf3).sort_index(), (pdf1 + pdf2 - pdf3).sort_index(), almost=True + ) + else: + assert_eq((psdf1 + psdf2 - psdf3).sort_index(), (pdf1 + pdf2 - pdf3).sort_index()) + + # Multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b")]) + psdf1.columns = columns + psdf2.columns = columns + pdf1.columns = columns + pdf2.columns = columns + columns = pd.MultiIndex.from_tuples([("x", "b"), ("y", "c")]) + psdf3.columns = columns + pdf3.columns = columns + + common_columns = set(psdf1.columns).intersection(psdf2.columns).intersection(psdf3.columns) + + # Series + assert_eq( + (psdf1[("x", "a")] - psdf2[("x", "b")] - psdf3[("y", "c")]).sort_index(), + (pdf1[("x", "a")] - pdf2[("x", "b")] - pdf3[("y", "c")]).sort_index(), + ) + + assert_eq( + (psdf1[("x", "a")] * (psdf2[("x", "b")] * psdf3[("y", "c")])).sort_index(), + (pdf1[("x", "a")] * (pdf2[("x", "b")] * pdf3[("y", "c")])).sort_index(), + ) + + # DataFrame + if check_extension and LooseVersion(pd.__version__) < LooseVersion("1.1"): + self.assert_eq( + (psdf1 + psdf2 - psdf3).sort_index(), (pdf1 + pdf2 - pdf3).sort_index(), almost=True + ) + else: + assert_eq((psdf1 + psdf2 - psdf3).sort_index(), (pdf1 + pdf2 - pdf3).sort_index()) + + def _test_arithmetic_chain_series(self, pser1, pser2, pser3, *, check_extension): + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + psser3 = ps.from_pandas(pser3) + + def assert_eq(actual, expected): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(actual, expected, check_exact=not check_extension) + if check_extension: + self.assertTrue(isinstance(actual.dtype, extension_dtypes)) + else: + self.assert_eq(actual, expected) + + # MultiIndex Series + assert_eq((psser1 + psser2 - psser3).sort_index(), (pser1 + pser2 - pser3).sort_index()) + + assert_eq((psser1 * psser2 * psser3).sort_index(), (pser1 * pser2 * pser3).sort_index()) + + if check_extension and not extension_float_dtypes_available: + self.assert_eq( + (psser1 - psser2 / psser3).sort_index(), (pser1 - pser2 / pser3).sort_index() + ) + else: + assert_eq((psser1 - psser2 / psser3).sort_index(), (pser1 - pser2 / pser3).sort_index()) + + assert_eq((psser1 + psser2 * psser3).sort_index(), (pser1 + pser2 * pser3).sort_index()) + + def test_mod(self): + pser = pd.Series([100, None, -300, None, 500, -700]) + pser_other = pd.Series([-150] * 6) + psser = ps.from_pandas(pser) + psser_other = ps.from_pandas(pser_other) + + self.assert_eq(psser.mod(psser_other).sort_index(), pser.mod(pser_other)) + self.assert_eq(psser.mod(psser_other).sort_index(), pser.mod(pser_other)) + self.assert_eq(psser.mod(psser_other).sort_index(), pser.mod(pser_other)) + + def test_rmod(self): + pser = pd.Series([100, None, -300, None, 500, -700]) + pser_other = pd.Series([-150] * 6) + psser = ps.from_pandas(pser) + psser_other = ps.from_pandas(pser_other) + + self.assert_eq(psser.rmod(psser_other).sort_index(), pser.rmod(pser_other)) + self.assert_eq(psser.rmod(psser_other).sort_index(), pser.rmod(pser_other)) + self.assert_eq(psser.rmod(psser_other).sort_index(), pser.rmod(pser_other)) + + def test_getitem_boolean_series(self): + pdf1 = pd.DataFrame( + {"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}, index=[20, 10, 30, 0, 50] + ) + pdf2 = pd.DataFrame( + {"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}, + index=[0, 30, 10, 20, 50], + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1[pdf2.A > -3].sort_index(), psdf1[psdf2.A > -3].sort_index()) + + self.assert_eq(pdf1.A[pdf2.A > -3].sort_index(), psdf1.A[psdf2.A > -3].sort_index()) + + self.assert_eq( + (pdf1.A + 1)[pdf2.A > -3].sort_index(), (psdf1.A + 1)[psdf2.A > -3].sort_index() + ) + + def test_loc_getitem_boolean_series(self): + pdf1 = pd.DataFrame( + {"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}, index=[20, 10, 30, 0, 50] + ) + pdf2 = pd.DataFrame( + {"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}, + index=[20, 10, 30, 0, 50], + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq(pdf1.loc[pdf2.A > -3].sort_index(), psdf1.loc[psdf2.A > -3].sort_index()) + + self.assert_eq(pdf1.A.loc[pdf2.A > -3].sort_index(), psdf1.A.loc[psdf2.A > -3].sort_index()) + + self.assert_eq( + (pdf1.A + 1).loc[pdf2.A > -3].sort_index(), (psdf1.A + 1).loc[psdf2.A > -3].sort_index() + ) + + pser = pd.Series([0, 1, 2, 3, 4], index=[20, 10, 30, 0, 50]) + psser = ps.from_pandas(pser) + self.assert_eq(pser.loc[pdf2.A > -3].sort_index(), psser.loc[psdf2.A > -3].sort_index()) + pser.name = psser.name = "B" + self.assert_eq(pser.loc[pdf2.A > -3].sort_index(), psser.loc[psdf2.A > -3].sort_index()) + + def test_bitwise(self): + pser1 = pd.Series([True, False, True, False, np.nan, np.nan, True, False, np.nan]) + pser2 = pd.Series([True, False, False, True, True, False, np.nan, np.nan, np.nan]) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + self.assert_eq(pser1 | pser2, (psser1 | psser2).sort_index()) + self.assert_eq(pser1 & pser2, (psser1 & psser2).sort_index()) + + pser1 = pd.Series([True, False, np.nan], index=list("ABC")) + pser2 = pd.Series([False, True, np.nan], index=list("DEF")) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + self.assert_eq(pser1 | pser2, (psser1 | psser2).sort_index()) + self.assert_eq(pser1 & pser2, (psser1 & psser2).sort_index()) + + @unittest.skipIf( + not extension_object_dtypes_available, "pandas extension object dtypes are not available" + ) + def test_bitwise_extension_dtype(self): + def assert_eq(actual, expected): + if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"): + self.assert_eq(actual, expected, check_exact=False) + self.assertTrue(isinstance(actual.dtype, extension_dtypes)) + else: + self.assert_eq(actual, expected) + + pser1 = pd.Series( + [True, False, True, False, np.nan, np.nan, True, False, np.nan], dtype="boolean" + ) + pser2 = pd.Series( + [True, False, False, True, True, False, np.nan, np.nan, np.nan], dtype="boolean" + ) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + assert_eq((psser1 | psser2).sort_index(), pser1 | pser2) + assert_eq((psser1 & psser2).sort_index(), pser1 & pser2) + + pser1 = pd.Series([True, False, np.nan], index=list("ABC"), dtype="boolean") + pser2 = pd.Series([False, True, np.nan], index=list("DEF"), dtype="boolean") + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + # a pandas bug? + # assert_eq((psser1 | psser2).sort_index(), pser1 | pser2) + # assert_eq((psser1 & psser2).sort_index(), pser1 & pser2) + assert_eq( + (psser1 | psser2).sort_index(), + pd.Series([True, None, None, None, True, None], index=list("ABCDEF"), dtype="boolean"), + ) + assert_eq( + (psser1 & psser2).sort_index(), + pd.Series( + [None, False, None, False, None, None], index=list("ABCDEF"), dtype="boolean" + ), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43453): Enable OpsOnDiffFramesEnabledTests.test_concat_column_axis " + "for pandas 2.0.0.", + ) + def test_concat_column_axis(self): + pdf1 = pd.DataFrame({"A": [0, 2, 4], "B": [1, 3, 5]}, index=[1, 2, 3]) + pdf1.columns.names = ["AB"] + pdf2 = pd.DataFrame({"C": [1, 2, 3], "D": [4, 5, 6]}, index=[1, 3, 5]) + pdf2.columns.names = ["CD"] + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + psdf3 = psdf1.copy() + psdf4 = psdf2.copy() + pdf3 = pdf1.copy() + pdf4 = pdf2.copy() + + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")], names=["X", "AB"]) + pdf3.columns = columns + psdf3.columns = columns + + columns = pd.MultiIndex.from_tuples([("X", "C"), ("X", "D")], names=["Y", "CD"]) + pdf4.columns = columns + psdf4.columns = columns + + pdf5 = pd.DataFrame({"A": [0, 2, 4], "B": [1, 3, 5]}, index=[1, 2, 3]) + pdf6 = pd.DataFrame({"C": [1, 2, 3]}, index=[1, 3, 5]) + psdf5 = ps.from_pandas(pdf5) + psdf6 = ps.from_pandas(pdf6) + + ignore_indexes = [True, False] + joins = ["inner", "outer"] + + objs = [ + ([psdf1.A, psdf2.C], [pdf1.A, pdf2.C]), + # TODO: ([psdf1, psdf2.C], [pdf1, pdf2.C]), + ([psdf1.A, psdf2], [pdf1.A, pdf2]), + ([psdf1.A, psdf2.C], [pdf1.A, pdf2.C]), + ([psdf3[("X", "A")], psdf4[("X", "C")]], [pdf3[("X", "A")], pdf4[("X", "C")]]), + ([psdf3, psdf4[("X", "C")]], [pdf3, pdf4[("X", "C")]]), + ([psdf3[("X", "A")], psdf4], [pdf3[("X", "A")], pdf4]), + ([psdf3, psdf4], [pdf3, pdf4]), + ([psdf5, psdf6], [pdf5, pdf6]), + ([psdf6, psdf5], [pdf6, pdf5]), + ] + + for ignore_index, join in product(ignore_indexes, joins): + for i, (psdfs, pdfs) in enumerate(objs): + with self.subTest(ignore_index=ignore_index, join=join, pdfs=pdfs, pair=i): + actual = ps.concat(psdfs, axis=1, ignore_index=ignore_index, join=join) + expected = pd.concat(pdfs, axis=1, ignore_index=ignore_index, join=join) + self.assert_eq( + repr(actual.sort_values(list(actual.columns)).reset_index(drop=True)), + repr(expected.sort_values(list(expected.columns)).reset_index(drop=True)), + ) + actual = ps.concat( + psdfs, axis=1, ignore_index=ignore_index, join=join, sort=True + ) + expected = pd.concat( + pdfs, axis=1, ignore_index=ignore_index, join=join, sort=True + ) + self.assert_eq( + repr(actual.reset_index(drop=True)), + repr(expected.reset_index(drop=True)), + ) + + def test_combine_first(self): + pser1 = pd.Series({"falcon": 330.0, "eagle": 160.0}) + pser2 = pd.Series({"falcon": 345.0, "eagle": 200.0, "duck": 30.0}) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + self.assert_eq( + psser1.combine_first(psser2).sort_index(), pser1.combine_first(pser2).sort_index() + ) + with self.assertRaisesRegex( + TypeError, "`combine_first` only allows `Series` for parameter `other`" + ): + psser1.combine_first(50) + + psser1.name = ("X", "A") + psser2.name = ("Y", "B") + pser1.name = ("X", "A") + pser2.name = ("Y", "B") + self.assert_eq( + psser1.combine_first(psser2).sort_index(), pser1.combine_first(pser2).sort_index() + ) + + # MultiIndex + midx1 = pd.MultiIndex( + [["lama", "cow", "falcon", "koala"], ["speed", "weight", "length", "power"]], + [[0, 3, 1, 1, 1, 2, 2, 2], [0, 2, 0, 3, 2, 0, 1, 3]], + ) + midx2 = pd.MultiIndex( + [["lama", "cow", "falcon"], ["speed", "weight", "length"]], + [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], + ) + pser1 = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1], index=midx1) + pser2 = pd.Series([-45, 200, -1.2, 30, -250, 1.5, 320, 1, -0.3], index=midx2) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + self.assert_eq( + psser1.combine_first(psser2).sort_index(), pser1.combine_first(pser2).sort_index() + ) + + # DataFrame + pdf1 = pd.DataFrame({"A": [None, 0], "B": [4, None]}) + psdf1 = ps.from_pandas(pdf1) + pdf2 = pd.DataFrame({"C": [3, 3], "B": [1, 1]}) + psdf2 = ps.from_pandas(pdf2) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2.0"): + self.assert_eq(pdf1.combine_first(pdf2), psdf1.combine_first(psdf2).sort_index()) + else: + # pandas < 1.2.0 returns unexpected dtypes, + # please refer to https://github.com/pandas-dev/pandas/issues/28481 for details + expected_pdf = pd.DataFrame({"A": [None, 0], "B": [4.0, 1.0], "C": [3, 3]}) + self.assert_eq(expected_pdf, psdf1.combine_first(psdf2).sort_index()) + + pdf1.columns = pd.MultiIndex.from_tuples([("A", "willow"), ("B", "pine")]) + psdf1 = ps.from_pandas(pdf1) + pdf2.columns = pd.MultiIndex.from_tuples([("C", "oak"), ("B", "pine")]) + psdf2 = ps.from_pandas(pdf2) + + if LooseVersion(pd.__version__) >= LooseVersion("1.2.0"): + self.assert_eq(pdf1.combine_first(pdf2), psdf1.combine_first(psdf2).sort_index()) + else: + # pandas < 1.2.0 returns unexpected dtypes, + # please refer to https://github.com/pandas-dev/pandas/issues/28481 for details + expected_pdf = pd.DataFrame({"A": [None, 0], "B": [4.0, 1.0], "C": [3, 3]}) + expected_pdf.columns = pd.MultiIndex.from_tuples( + [("A", "willow"), ("B", "pine"), ("C", "oak")] + ) + self.assert_eq(expected_pdf, psdf1.combine_first(psdf2).sort_index()) + + def test_insert(self): + # + # Basic DataFrame + # + pdf = pd.DataFrame([1, 2, 3]) + psdf = ps.from_pandas(pdf) + + pser = pd.Series([4, 5, 6]) + psser = ps.from_pandas(pser) + psdf.insert(1, "y", psser) + pdf.insert(1, "y", pser) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + # + # DataFrame with Index different from inserting Series' + # + pdf = pd.DataFrame([1, 2, 3], index=[10, 20, 30]) + psdf = ps.from_pandas(pdf) + + pser = pd.Series([4, 5, 6]) + psser = ps.from_pandas(pser) + psdf.insert(1, "y", psser) + pdf.insert(1, "y", pser) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + # + # DataFrame with Multi-index columns + # + pdf = pd.DataFrame({("x", "a"): [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + + pser = pd.Series([4, 5, 6]) + psser = ps.from_pandas(pser) + pdf = pd.DataFrame({("x", "a", "b"): [1, 2, 3]}) + psdf = ps.from_pandas(pdf) + psdf.insert(0, "a", psser) + pdf.insert(0, "a", pser) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + psdf.insert(0, ("b", "c", ""), psser) + pdf.insert(0, ("b", "c", ""), pser) + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_compare(self): + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + pser1 = pd.Series(["b", "c", np.nan, "g", np.nan]) + pser2 = pd.Series(["a", "c", np.nan, np.nan, "h"]) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + self.assert_eq( + pser1.compare(pser2).sort_index(), + psser1.compare(psser2).sort_index(), + ) + + # `keep_shape=True` + self.assert_eq( + pser1.compare(pser2, keep_shape=True).sort_index(), + psser1.compare(psser2, keep_shape=True).sort_index(), + ) + # `keep_equal=True` + self.assert_eq( + pser1.compare(pser2, keep_equal=True).sort_index(), + psser1.compare(psser2, keep_equal=True).sort_index(), + ) + # `keep_shape=True` and `keep_equal=True` + self.assert_eq( + pser1.compare(pser2, keep_shape=True, keep_equal=True).sort_index(), + psser1.compare(psser2, keep_shape=True, keep_equal=True).sort_index(), + ) + + # MultiIndex + pser1.index = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("x", "k"), ("q", "l")] + ) + pser2.index = pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("x", "k"), ("q", "l")] + ) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + self.assert_eq( + pser1.compare(pser2).sort_index(), + psser1.compare(psser2).sort_index(), + ) + + # `keep_shape=True` with MultiIndex + self.assert_eq( + pser1.compare(pser2, keep_shape=True).sort_index(), + psser1.compare(psser2, keep_shape=True).sort_index(), + ) + # `keep_equal=True` with MultiIndex + self.assert_eq( + pser1.compare(pser2, keep_equal=True).sort_index(), + psser1.compare(psser2, keep_equal=True).sort_index(), + ) + # `keep_shape=True` and `keep_equal=True` with MultiIndex + self.assert_eq( + pser1.compare(pser2, keep_shape=True, keep_equal=True).sort_index(), + psser1.compare(psser2, keep_shape=True, keep_equal=True).sort_index(), + ) + else: + psser1 = ps.Series(["b", "c", np.nan, "g", np.nan]) + psser2 = ps.Series(["a", "c", np.nan, np.nan, "h"]) + expected = ps.DataFrame( + [["b", "a"], ["g", None], [None, "h"]], index=[0, 3, 4], columns=["self", "other"] + ) + self.assert_eq(expected, psser1.compare(psser2).sort_index()) + + # `keep_shape=True` + expected = ps.DataFrame( + [["b", "a"], [None, None], [None, None], ["g", None], [None, "h"]], + index=[0, 1, 2, 3, 4], + columns=["self", "other"], + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_shape=True).sort_index(), + ) + # `keep_equal=True` + expected = ps.DataFrame( + [["b", "a"], ["g", None], [None, "h"]], index=[0, 3, 4], columns=["self", "other"] + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_equal=True).sort_index(), + ) + # `keep_shape=True` and `keep_equal=True` + expected = ps.DataFrame( + [["b", "a"], ["c", "c"], [None, None], ["g", None], [None, "h"]], + index=[0, 1, 2, 3, 4], + columns=["self", "other"], + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_shape=True, keep_equal=True).sort_index(), + ) + + # MultiIndex + psser1 = ps.Series( + ["b", "c", np.nan, "g", np.nan], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("x", "k"), ("q", "l")] + ), + ) + psser2 = ps.Series( + ["a", "c", np.nan, np.nan, "h"], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("x", "k"), ("q", "l")] + ), + ) + expected = ps.DataFrame( + [["b", "a"], [None, "h"], ["g", None]], + index=pd.MultiIndex.from_tuples([("a", "x"), ("q", "l"), ("x", "k")]), + columns=["self", "other"], + ) + self.assert_eq(expected, psser1.compare(psser2).sort_index()) + + # `keep_shape=True` + expected = ps.DataFrame( + [["b", "a"], [None, None], [None, None], [None, "h"], ["g", None]], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("q", "l"), ("x", "k")] + ), + columns=["self", "other"], + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_shape=True).sort_index(), + ) + # `keep_equal=True` + expected = ps.DataFrame( + [["b", "a"], [None, "h"], ["g", None]], + index=pd.MultiIndex.from_tuples([("a", "x"), ("q", "l"), ("x", "k")]), + columns=["self", "other"], + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_equal=True).sort_index(), + ) + # `keep_shape=True` and `keep_equal=True` + expected = ps.DataFrame( + [["b", "a"], ["c", "c"], [None, None], [None, "h"], ["g", None]], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("q", "l"), ("x", "k")] + ), + columns=["self", "other"], + ) + self.assert_eq( + expected, + psser1.compare(psser2, keep_shape=True, keep_equal=True).sort_index(), + ) + + # Different Index + with self.assertRaisesRegex( + ValueError, "Can only compare identically-labeled Series objects" + ): + psser1 = ps.Series( + [1, 2, 3, 4, 5], + index=pd.Index([1, 2, 3, 4, 5]), + ) + psser2 = ps.Series( + [2, 2, 3, 4, 1], + index=pd.Index([5, 4, 3, 2, 1]), + ) + psser1.compare(psser2) + # Different MultiIndex + with self.assertRaisesRegex( + ValueError, "Can only compare identically-labeled Series objects" + ): + psser1 = ps.Series( + [1, 2, 3, 4, 5], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "z"), ("x", "k"), ("q", "l")] + ), + ) + psser2 = ps.Series( + [2, 2, 3, 4, 1], + index=pd.MultiIndex.from_tuples( + [("a", "x"), ("b", "y"), ("c", "a"), ("x", "k"), ("q", "l")] + ), + ) + psser1.compare(psser2) + # SPARK-37495: Skip identical index checking of Series.compare when config + # 'compute.eager_check' is disabled + psser1 = ps.Series([1, 2, 3, 4, 5], index=pd.Index([1, 2, 3, 4, 5])) + psser2 = ps.Series([1, 2, 3, 4, 5, 6], index=pd.Index([1, 2, 4, 3, 6, 7])) + expected = ps.DataFrame( + {"self": [3, 4, 5, np.nan, np.nan], "other": [4, 3, np.nan, 5.0, 6.0]}, + index=[3, 4, 5, 6, 7], + ) + + with ps.option_context("compute.eager_check", False): + self.assert_eq(expected, psser1.compare(psser2)) + + def test_different_columns(self): + psdf1 = self.psdf1 + psdf4 = self.psdf4 + pdf1 = self.pdf1 + pdf4 = self.pdf4 + + self.assert_eq((psdf1 + psdf4).sort_index(), (pdf1 + pdf4).sort_index(), almost=True) + + # Multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b")]) + psdf1.columns = columns + pdf1.columns = columns + columns = pd.MultiIndex.from_tuples([("z", "e"), ("z", "f")]) + psdf4.columns = columns + pdf4.columns = columns + + self.assert_eq((psdf1 + psdf4).sort_index(), (pdf1 + pdf4).sort_index(), almost=True) + + def test_assignment_series(self): + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psser = psdf.a + pser = pdf.a + psdf["a"] = self.psdf2.a + pdf["a"] = self.pdf2.a + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser, pser) + + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psser = psdf.a + pser = pdf.a + psdf["a"] = self.psdf2.b + pdf["a"] = self.pdf2.b + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser, pser) + + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf["c"] = self.psdf2.a + pdf["c"] = self.pdf2.a + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + # Multi-index columns + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b")]) + psdf.columns = columns + pdf.columns = columns + psdf[("y", "c")] = self.psdf2.a + pdf[("y", "c")] = self.pdf2.a + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + pdf = pd.DataFrame({"a": [1, 2, 3], "Koalas": [0, 1, 2]}).set_index("Koalas", drop=False) + psdf = ps.from_pandas(pdf) + + psdf.index.name = None + psdf["NEW"] = ps.Series([100, 200, 300]) + + pdf.index.name = None + pdf["NEW"] = pd.Series([100, 200, 300]) + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_assignment_frame(self): + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psser = psdf.a + pser = pdf.a + psdf[["a", "b"]] = self.psdf1 + pdf[["a", "b"]] = self.pdf1 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser, pser) + + # 'c' does not exist in `psdf`. + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psser = psdf.a + pser = pdf.a + psdf[["b", "c"]] = self.psdf1 + pdf[["b", "c"]] = self.pdf1 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + self.assert_eq(psser, pser) + + # 'c' and 'd' do not exist in `psdf`. + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf[["c", "d"]] = self.psdf1 + pdf[["c", "d"]] = self.pdf1 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + # Multi-index columns + columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b")]) + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf.columns = columns + pdf.columns = columns + psdf[[("y", "c"), ("z", "d")]] = self.psdf1 + pdf[[("y", "c"), ("z", "d")]] = self.pdf1 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf1 = ps.from_pandas(self.pdf1) + pdf1 = self.pdf1 + psdf1.columns = columns + pdf1.columns = columns + psdf[["c", "d"]] = psdf1 + pdf[["c", "d"]] = pdf1 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_assignment_series_chain(self): + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf["a"] = self.psdf1.a + pdf["a"] = self.pdf1.a + + psdf["a"] = self.psdf2.b + pdf["a"] = self.pdf2.b + + psdf["d"] = self.psdf3.c + pdf["d"] = self.pdf3.c + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_assignment_frame_chain(self): + psdf = ps.from_pandas(self.pdf1) + pdf = self.pdf1 + psdf[["a", "b"]] = self.psdf1 + pdf[["a", "b"]] = self.pdf1 + + psdf[["e", "f"]] = self.psdf3 + pdf[["e", "f"]] = self.pdf3 + + psdf[["b", "c"]] = self.psdf2 + pdf[["b", "c"]] = self.pdf2 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_multi_index_arithmetic(self): + psdf5 = self.psdf5 + psdf6 = self.psdf6 + pdf5 = self.pdf5 + pdf6 = self.pdf6 + + # Series + self.assert_eq((psdf5.c - psdf6.e).sort_index(), (pdf5.c - pdf6.e).sort_index()) + + self.assert_eq((psdf5["c"] / psdf6["e"]).sort_index(), (pdf5["c"] / pdf6["e"]).sort_index()) + + # DataFrame + self.assert_eq((psdf5 + psdf6).sort_index(), (pdf5 + pdf6).sort_index(), almost=True) + + def test_multi_index_assignment_series(self): + psdf = ps.from_pandas(self.pdf5) + pdf = self.pdf5 + psdf["x"] = self.psdf6.e + pdf["x"] = self.pdf6.e + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + psdf = ps.from_pandas(self.pdf5) + pdf = self.pdf5 + psdf["e"] = self.psdf6.e + pdf["e"] = self.pdf6.e + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + psdf = ps.from_pandas(self.pdf5) + pdf = self.pdf5 + psdf["c"] = self.psdf6.e + pdf["c"] = self.pdf6.e + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + def test_multi_index_assignment_frame(self): + psdf = ps.from_pandas(self.pdf5) + pdf = self.pdf5 + psdf[["c"]] = self.psdf5 + pdf[["c"]] = self.pdf5 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + psdf = ps.from_pandas(self.pdf5) + pdf = self.pdf5 + psdf[["x"]] = self.psdf5 + pdf[["x"]] = self.pdf5 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + psdf = ps.from_pandas(self.pdf6) + pdf = self.pdf6 + psdf[["x", "y"]] = self.psdf6 + pdf[["x", "y"]] = self.pdf6 + + self.assert_eq(psdf.sort_index(), pdf.sort_index()) + + +class OpsOnDiffFramesDisabledTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", False) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + @property + def pdf1(self): + return pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, + index=[0, 1, 3, 5, 6, 8, 9, 9, 9], + ) + + @property + def pdf2(self): + return pd.DataFrame( + {"a": [9, 8, 7, 6, 5, 4, 3, 2, 1], "b": [0, 0, 0, 4, 5, 6, 1, 2, 3]}, + index=list(range(9)), + ) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.from_pandas(self.pdf2) + + def test_arithmetic(self): + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + self.psdf1.a - self.psdf2.b + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + self.psdf1.a - self.psdf2.a + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + self.psdf1["a"] - self.psdf2["a"] + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + self.psdf1 - self.psdf2 + + def test_assignment(self): + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf = ps.from_pandas(self.pdf1) + psdf["c"] = self.psdf1.a + + def test_frame_loc_setitem(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.DataFrame(pdf) + another_psdf = ps.DataFrame(pdf) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf.loc[["viper", "sidewinder"], ["shield"]] = another_psdf.max_speed + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf.loc[another_psdf.max_speed < 5, ["shield"]] = -psdf.max_speed + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf.loc[another_psdf.max_speed < 5, ["shield"]] = -another_psdf.max_speed + + def test_frame_iloc_setitem(self): + pdf = pd.DataFrame( + [[1, 2], [4, 5], [7, 8]], + index=["cobra", "viper", "sidewinder"], + columns=["max_speed", "shield"], + ) + psdf = ps.DataFrame(pdf) + another_psdf = ps.DataFrame(pdf) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf.iloc[[1, 2], [1]] = another_psdf.max_speed.iloc[[1, 2]] + + def test_series_loc_setitem(self): + pser = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser = ps.from_pandas(pser) + + pser_another = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser_another = ps.from_pandas(pser_another) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.loc[psser % 2 == 1] = -psser_another + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.loc[psser_another % 2 == 1] = -psser + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.loc[psser_another % 2 == 1] = -psser_another + + def test_series_iloc_setitem(self): + pser = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser = ps.from_pandas(pser) + + pser_another = pd.Series([1, 2, 3], index=["cobra", "viper", "sidewinder"]) + psser_another = ps.from_pandas(pser_another) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.iloc[[1]] = -psser_another.iloc[[1]] + + def test_where(self): + pdf1 = pd.DataFrame({"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame({"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.where(psdf2 > 100) + + pdf1 = pd.DataFrame({"A": [-1, -2, -3, -4, -5], "B": [-100, -200, -300, -400, -500]}) + pdf2 = pd.DataFrame({"A": [-10, -20, -30, -40, -50], "B": [-5, -4, -3, -2, -1]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.where(psdf2 < -250) + + def test_mask(self): + pdf1 = pd.DataFrame({"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}) + pdf2 = pd.DataFrame({"A": [0, -1, -2, -3, -4], "B": [-100, -200, -300, -400, -500]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.mask(psdf2 < 100) + + pdf1 = pd.DataFrame({"A": [-1, -2, -3, -4, -5], "B": [-100, -200, -300, -400, -500]}) + pdf2 = pd.DataFrame({"A": [-10, -20, -30, -40, -50], "B": [-5, -4, -3, -2, -1]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.mask(psdf2 > -250) + + def test_align(self): + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30]) + pdf2 = pd.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12]) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.align(psdf2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.align(psdf2, axis=0) + + def test_pow_and_rpow(self): + pser = pd.Series([1, 2, np.nan]) + psser = ps.from_pandas(pser) + pser_other = pd.Series([np.nan, 2, 3]) + psser_other = ps.from_pandas(pser_other) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.pow(psser_other) + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser**psser_other + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.rpow(psser_other) + + def test_equals(self): + psidx1 = ps.Index([1, 2, 3, 4]) + psidx2 = ps.Index([1, 2, 3, 4]) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psidx1.equals(psidx2) + + def test_combine_first(self): + pdf1 = pd.DataFrame({"A": [None, 0], "B": [4, None]}) + psdf1 = ps.from_pandas(pdf1) + + self.assertRaises(TypeError, lambda: psdf1.combine_first(ps.Series([1, 2]))) + + pser1 = pd.Series({"falcon": 330.0, "eagle": 160.0}) + pser2 = pd.Series({"falcon": 345.0, "eagle": 200.0, "duck": 30.0}) + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser1.combine_first(psser2) + + pdf1 = pd.DataFrame({"A": [None, 0], "B": [4, None]}) + psdf1 = ps.from_pandas(pdf1) + pdf2 = pd.DataFrame({"C": [3, 3], "B": [1, 1]}) + psdf2 = ps.from_pandas(pdf2) + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psdf1.combine_first(psdf2) + + def test_series_eq(self): + pser = pd.Series([1, 2, 3, 4, 5, 6], name="x") + psser = ps.from_pandas(pser) + + others = ( + ps.Series([np.nan, 1, 3, 4, np.nan, 6], name="x"), + ps.Index([np.nan, 1, 3, 4, np.nan, 6], name="x"), + ) + for other in others: + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser.eq(other) + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser == other + + +class OpsOnDiffFramesEnabledTests( + OpsOnDiffFramesEnabledTestsMixin, PandasOnSparkTestCase, SQLTestUtils +): + pass + + +class OpsOnDiffFramesDisabledTests( + OpsOnDiffFramesDisabledTestsMixin, PandasOnSparkTestCase, SQLTestUtils +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_ops_on_diff_frames import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py new file mode 100644 index 0000000..f581db4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py @@ -0,0 +1,655 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class OpsOnDiffFramesGroupByTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43460): Enable OpsOnDiffFramesGroupByTests.test_groupby_different_lengths " + "for pandas 2.0.0.", + ) + def test_groupby_different_lengths(self): + pdfs1 = [ + pd.DataFrame({"c": [4, 2, 7, 3, None, 1, 1, 1, 2], "d": list("abcdefght")}), + pd.DataFrame({"c": [4, 2, 7, None, 1, 1, 2], "d": list("abcdefg")}), + pd.DataFrame({"c": [4, 2, 7, 3, None, 1, 1, 1, 2, 2], "d": list("abcdefghti")}), + ] + pdfs2 = [ + pd.DataFrame({"a": [1, 2, 6, 4, 4, 6, 4, 3, 7], "b": [4, 2, 7, 3, 3, 1, 1, 1, 2]}), + pd.DataFrame({"a": [1, 2, 6, 4, 4, 6, 4, 7], "b": [4, 2, 7, 3, 3, 1, 1, 2]}), + pd.DataFrame({"a": [1, 2, 6, 4, 4, 6, 4, 3, 7], "b": [4, 2, 7, 3, 3, 1, 1, 1, 2]}), + ] + + for pdf1, pdf2 in zip(pdfs1, pdfs2): + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values("c").reset_index(drop=True) + + self.assert_eq( + sort(psdf1.groupby(psdf2.a, as_index=as_index).sum()), + sort(pdf1.groupby(pdf2.a, as_index=as_index).sum()), + almost=as_index, + ) + + self.assert_eq( + sort(psdf1.groupby(psdf2.a, as_index=as_index).c.sum()), + sort(pdf1.groupby(pdf2.a, as_index=as_index).c.sum()), + almost=as_index, + ) + self.assert_eq( + sort(psdf1.groupby(psdf2.a, as_index=as_index)["c"].sum()), + sort(pdf1.groupby(pdf2.a, as_index=as_index)["c"].sum()), + almost=as_index, + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43459): Enable OpsOnDiffFramesGroupByTests.test_groupby_multiindex_columns " + "for pandas 2.0.0.", + ) + def test_groupby_multiindex_columns(self): + pdf1 = pd.DataFrame( + {("y", "c"): [4, 2, 7, 3, None, 1, 1, 1, 2], ("z", "d"): list("abcdefght")} + ) + pdf2 = pd.DataFrame( + {("x", "a"): [1, 2, 6, 4, 4, 6, 4, 3, 7], ("x", "b"): [4, 2, 7, 3, 3, 1, 1, 1, 2]} + ) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq( + psdf1.groupby(psdf2[("x", "a")]).sum().sort_index(), + pdf1.groupby(pdf2[("x", "a")]).sum().sort_index(), + ) + + self.assert_eq( + psdf1.groupby(psdf2[("x", "a")], as_index=False) + .sum() + .sort_values(("y", "c")) + .reset_index(drop=True), + pdf1.groupby(pdf2[("x", "a")], as_index=False) + .sum() + .sort_values(("y", "c")) + .reset_index(drop=True), + ) + self.assert_eq( + psdf1.groupby(psdf2[("x", "a")])[[("y", "c")]].sum().sort_index(), + pdf1.groupby(pdf2[("x", "a")])[[("y", "c")]].sum().sort_index(), + ) + + def test_split_apply_combine_on_series(self): + pdf1 = pd.DataFrame({"C": [0.362, 0.227, 1.267, -0.562], "B": [1, 2, 3, 4]}) + pdf2 = pd.DataFrame({"A": [1, 1, 2, 2]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(list(df.columns)).reset_index(drop=True) + + with self.subTest(as_index=as_index): + self.assert_eq( + sort(psdf1.groupby(psdf2.A, as_index=as_index).sum()), + sort(pdf1.groupby(pdf2.A, as_index=as_index).sum()), + ) + self.assert_eq( + sort(psdf1.groupby(psdf2.A, as_index=as_index).B.sum()), + sort(pdf1.groupby(pdf2.A, as_index=as_index).B.sum()), + ) + self.assert_eq( + sort(psdf1.groupby([psdf1.C, psdf2.A], as_index=as_index).sum()), + sort(pdf1.groupby([pdf1.C, pdf2.A], as_index=as_index).sum()), + ) + self.assert_eq( + sort(psdf1.groupby([psdf1.C + 1, psdf2.A], as_index=as_index).sum()), + sort(pdf1.groupby([pdf1.C + 1, pdf2.A], as_index=as_index).sum()), + ) + + self.assert_eq( + psdf1.B.groupby(psdf2.A).sum().sort_index(), + pdf1.B.groupby(pdf2.A).sum().sort_index(), + ) + self.assert_eq( + (psdf1.B + 1).groupby(psdf2.A).sum().sort_index(), + (pdf1.B + 1).groupby(pdf2.A).sum().sort_index(), + ) + + self.assert_eq( + psdf1.B.groupby(psdf2.A.rename()).sum().sort_index(), + pdf1.B.groupby(pdf2.A.rename()).sum().sort_index(), + ) + self.assert_eq( + psdf1.B.rename().groupby(psdf2.A).sum().sort_index(), + pdf1.B.rename().groupby(pdf2.A).sum().sort_index(), + ) + self.assert_eq( + psdf1.B.rename().groupby(psdf2.A.rename()).sum().sort_index(), + pdf1.B.rename().groupby(pdf2.A.rename()).sum().sort_index(), + ) + + def test_aggregate(self): + pdf1 = pd.DataFrame({"C": [0.362, 0.227, 1.267, -0.562], "B": [1, 2, 3, 4]}) + pdf2 = pd.DataFrame({"A": [1, 1, 2, 2]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + for as_index in [True, False]: + if as_index: + + def sort(df): + return df.sort_index() + + else: + + def sort(df): + return df.sort_values(list(df.columns)).reset_index(drop=True) + + with self.subTest(as_index=as_index): + self.assert_eq( + sort(psdf1.groupby(psdf2.A, as_index=as_index).agg("sum")), + sort(pdf1.groupby(pdf2.A, as_index=as_index).agg("sum")), + ) + self.assert_eq( + sort(psdf1.groupby(psdf2.A, as_index=as_index).agg({"B": "min", "C": "sum"})), + sort(pdf1.groupby(pdf2.A, as_index=as_index).agg({"B": "min", "C": "sum"})), + ) + self.assert_eq( + sort( + psdf1.groupby(psdf2.A, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + sort( + pdf1.groupby(pdf2.A, as_index=as_index).agg( + {"B": ["min", "max"], "C": "sum"} + ) + ), + ) + self.assert_eq( + sort(psdf1.groupby([psdf1.C, psdf2.A], as_index=as_index).agg("sum")), + sort(pdf1.groupby([pdf1.C, pdf2.A], as_index=as_index).agg("sum")), + ) + self.assert_eq( + sort(psdf1.groupby([psdf1.C + 1, psdf2.A], as_index=as_index).agg("sum")), + sort(pdf1.groupby([pdf1.C + 1, pdf2.A], as_index=as_index).agg("sum")), + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("Y", "C"), ("X", "B")]) + pdf1.columns = columns + psdf1.columns = columns + + columns = pd.MultiIndex.from_tuples([("X", "A")]) + pdf2.columns = columns + psdf2.columns = columns + + for as_index in [True, False]: + stats_psdf = psdf1.groupby(psdf2[("X", "A")], as_index=as_index).agg( + {("X", "B"): "min", ("Y", "C"): "sum"} + ) + stats_pdf = pdf1.groupby(pdf2[("X", "A")], as_index=as_index).agg( + {("X", "B"): "min", ("Y", "C"): "sum"} + ) + self.assert_eq( + stats_psdf.sort_values(by=[("X", "B"), ("Y", "C")]).reset_index(drop=True), + stats_pdf.sort_values(by=[("X", "B"), ("Y", "C")]).reset_index(drop=True), + ) + + stats_psdf = psdf1.groupby(psdf2[("X", "A")]).agg( + {("X", "B"): ["min", "max"], ("Y", "C"): "sum"} + ) + stats_pdf = pdf1.groupby(pdf2[("X", "A")]).agg( + {("X", "B"): ["min", "max"], ("Y", "C"): "sum"} + ) + self.assert_eq( + stats_psdf.sort_values( + by=[("X", "B", "min"), ("X", "B", "max"), ("Y", "C", "sum")] + ).reset_index(drop=True), + stats_pdf.sort_values( + by=[("X", "B", "min"), ("X", "B", "max"), ("Y", "C", "sum")] + ).reset_index(drop=True), + ) + + def test_duplicated_labels(self): + pdf1 = pd.DataFrame({"A": [3, 2, 1]}) + pdf2 = pd.DataFrame({"A": [1, 2, 3]}) + psdf1 = ps.from_pandas(pdf1) + psdf2 = ps.from_pandas(pdf2) + + self.assert_eq( + psdf1.groupby(psdf2.A).sum().sort_index(), pdf1.groupby(pdf2.A).sum().sort_index() + ) + self.assert_eq( + psdf1.groupby(psdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True), + pdf1.groupby(pdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True), + ) + + def test_apply(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8]) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey).apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["a"].apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey)["a"].apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey)[["a"]].apply(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", kkey]).apply(lambda x: x + x.min()).sort_index(), + pdf.groupby(["a", pkey]).apply(lambda x: x + x.min()).sort_index(), + ) + + def test_transform(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8]) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey).transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["a"].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey)["a"].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(pkey)[["a"]].transform(lambda x: x + x.min()).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", kkey]).transform(lambda x: x + x.min()).sort_index(), + pdf.groupby(["a", pkey]).transform(lambda x: x + x.min()).sort_index(), + ) + + def test_filter(self): + pdf = pd.DataFrame( + {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]}, + columns=["a", "b", "c"], + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8]) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(pkey).filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["a"].filter(lambda x: any(x == 2)).sort_index(), + pdf.groupby(pkey)["a"].filter(lambda x: any(x == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(pkey)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(), + ) + self.assert_eq( + psdf.groupby(["a", kkey]).filter(lambda x: any(x.a == 2)).sort_index(), + pdf.groupby(["a", pkey]).filter(lambda x: any(x.a == 2)).sort_index(), + ) + + def test_head(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, + "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, + "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, + }, + ) + pkey = pd.Series([1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + pdf.groupby(pkey).head(2).sort_index(), psdf.groupby(kkey).head(2).sort_index() + ) + self.assert_eq( + pdf.groupby("a")["b"].head(2).sort_index(), psdf.groupby("a")["b"].head(2).sort_index() + ) + self.assert_eq( + pdf.groupby("a")[["b"]].head(2).sort_index(), + psdf.groupby("a")[["b"]].head(2).sort_index(), + ) + self.assert_eq( + pdf.groupby([pkey, "b"]).head(2).sort_index(), + psdf.groupby([kkey, "b"]).head(2).sort_index(), + ) + + def test_cumcount(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + for ascending in [True, False]: + self.assert_eq( + psdf.groupby(kkey).cumcount(ascending=ascending).sort_index(), + pdf.groupby(pkey).cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["a"].cumcount(ascending=ascending).sort_index(), + pdf.groupby(pkey)["a"].cumcount(ascending=ascending).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].cumcount(ascending=ascending).sort_index(), + pdf.groupby(pkey)[["a"]].cumcount(ascending=ascending).sort_index(), + ) + + def test_cummin(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).cummin().sort_index(), pdf.groupby(pkey).cummin().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].cummin().sort_index(), + pdf.groupby(pkey)["a"].cummin().sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].cummin().sort_index(), + pdf.groupby(pkey)[["a"]].cummin().sort_index(), + ) + + def test_cummax(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).cummax().sort_index(), pdf.groupby(pkey).cummax().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].cummax().sort_index(), + pdf.groupby(pkey)["a"].cummax().sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].cummax().sort_index(), + pdf.groupby(pkey)[["a"]].cummax().sort_index(), + ) + + def test_cumsum(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).cumsum().sort_index(), pdf.groupby(pkey).cumsum().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].cumsum().sort_index(), + pdf.groupby(pkey)["a"].cumsum().sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].cumsum().sort_index(), + pdf.groupby(pkey)[["a"]].cumsum().sort_index(), + ) + + def test_cumprod(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).cumprod().sort_index(), + pdf.groupby(pkey).cumprod().sort_index(), + almost=True, + ) + self.assert_eq( + psdf.groupby(kkey)["a"].cumprod().sort_index(), + pdf.groupby(pkey)["a"].cumprod().sort_index(), + almost=True, + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].cumprod().sort_index(), + pdf.groupby(pkey)[["a"]].cumprod().sort_index(), + almost=True, + ) + + def test_diff(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + } + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).diff().sort_index(), pdf.groupby(pkey).diff().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].diff().sort_index(), pdf.groupby(pkey)["a"].diff().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].diff().sort_index(), + pdf.groupby(pkey)[["a"]].diff().sort_index(), + ) + + self.assert_eq(psdf.groupby(kkey).diff().sum(), pdf.groupby(pkey).diff().sum().astype(int)) + self.assert_eq(psdf.groupby(kkey)["a"].diff().sum(), pdf.groupby(pkey)["a"].diff().sum()) + + def test_fillna(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6] * 3, + "b": [1, 1, 2, 3, 5, 8] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 3, 5, 8] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).rank().sort_index(), pdf.groupby(pkey).rank().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].rank().sort_index(), pdf.groupby(pkey)["a"].rank().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].rank().sort_index(), + pdf.groupby(pkey)[["a"]].rank().sort_index(), + ) + + self.assert_eq(psdf.groupby(kkey).rank().sum(), pdf.groupby(pkey).rank().sum()) + self.assert_eq(psdf.groupby(kkey)["a"].rank().sum(), pdf.groupby(pkey)["a"].rank().sum()) + + def test_shift(self): + pdf = pd.DataFrame( + { + "a": [1, 1, 2, 2, 3, 3] * 3, + "b": [1, 1, 2, 2, 3, 4] * 3, + "c": [1, 4, 9, 16, 25, 36] * 3, + }, + ) + pkey = pd.Series([1, 1, 2, 2, 3, 4] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).shift().sort_index(), pdf.groupby(pkey).shift().sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["a"].shift().sort_index(), + pdf.groupby(pkey)["a"].shift().sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["a"]].shift().sort_index(), + pdf.groupby(pkey)[["a"]].shift().sort_index(), + ) + + self.assert_eq( + psdf.groupby(kkey).shift().sum(), pdf.groupby(pkey).shift().sum().astype(int) + ) + self.assert_eq(psdf.groupby(kkey)["a"].shift().sum(), pdf.groupby(pkey)["a"].shift().sum()) + + def test_fillna(self): + pdf = pd.DataFrame( + { + "A": [1, 1, 2, 2] * 3, + "B": [2, 4, None, 3] * 3, + "C": [None, None, None, 1] * 3, + "D": [0, 1, 5, 4] * 3, + } + ) + pkey = pd.Series([1, 1, 2, 2] * 3) + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + psdf.groupby(kkey).fillna(0).sort_index(), pdf.groupby(pkey).fillna(0).sort_index() + ) + self.assert_eq( + psdf.groupby(kkey)["C"].fillna(0).sort_index(), + pdf.groupby(pkey)["C"].fillna(0).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["C"]].fillna(0).sort_index(), + pdf.groupby(pkey)[["C"]].fillna(0).sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey).fillna(method="bfill").sort_index(), + pdf.groupby(pkey).fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["C"].fillna(method="bfill").sort_index(), + pdf.groupby(pkey)["C"].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["C"]].fillna(method="bfill").sort_index(), + pdf.groupby(pkey)[["C"]].fillna(method="bfill").sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey).fillna(method="ffill").sort_index(), + pdf.groupby(pkey).fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)["C"].fillna(method="ffill").sort_index(), + pdf.groupby(pkey)["C"].fillna(method="ffill").sort_index(), + ) + self.assert_eq( + psdf.groupby(kkey)[["C"]].fillna(method="ffill").sort_index(), + pdf.groupby(pkey)[["C"]].fillna(method="ffill").sort_index(), + ) + + +class OpsOnDiffFramesGroupByTests( + OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestCase, SQLTestUtils +): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_expanding.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_expanding.py new file mode 100644 index 0000000..9987a22 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_expanding.py @@ -0,0 +1,113 @@ +# +# 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 distutils.version import LooseVersion + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class OpsOnDiffFramesGroupByExpandingTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def _test_groupby_expanding_func(self, f): + pser = pd.Series([1, 2, 3]) + pkey = pd.Series([1, 2, 3], name="a") + psser = ps.from_pandas(pser) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + getattr(psser.groupby(kkey).expanding(2), f)().sort_index(), + getattr(pser.groupby(pkey).expanding(2), f)().sort_index(), + ) + + pdf = pd.DataFrame({"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0]}) + pkey = pd.Series([1, 2, 3, 2], name="a") + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + # The behavior of GroupBy.expanding is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + getattr(psdf.groupby(kkey).expanding(2), f)().sort_index(), + getattr(pdf.groupby(pkey).expanding(2), f)().sort_index(), + ) + else: + self.assert_eq( + getattr(psdf.groupby(kkey).expanding(2), f)().sort_index(), + getattr(pdf.groupby(pkey).expanding(2), f)().drop("a", axis=1).sort_index(), + ) + + self.assert_eq( + getattr(psdf.groupby(kkey)["b"].expanding(2), f)().sort_index(), + getattr(pdf.groupby(pkey)["b"].expanding(2), f)().sort_index(), + ) + self.assert_eq( + getattr(psdf.groupby(kkey)[["b"]].expanding(2), f)().sort_index(), + getattr(pdf.groupby(pkey)[["b"]].expanding(2), f)().sort_index(), + ) + + def test_groupby_expanding_count(self): + self._test_groupby_expanding_func("count") + + def test_groupby_expanding_min(self): + self._test_groupby_expanding_func("min") + + def test_groupby_expanding_max(self): + self._test_groupby_expanding_func("max") + + def test_groupby_expanding_mean(self): + self._test_groupby_expanding_func("mean") + + def test_groupby_expanding_sum(self): + self._test_groupby_expanding_func("sum") + + def test_groupby_expanding_std(self): + self._test_groupby_expanding_func("std") + + def test_groupby_expanding_var(self): + self._test_groupby_expanding_func("var") + + +class OpsOnDiffFramesGroupByExpandingTests( + OpsOnDiffFramesGroupByExpandingTestsMixin, PandasOnSparkTestCase, TestUtils +): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_ops_on_diff_frames_groupby_expanding import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_rolling.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_rolling.py new file mode 100644 index 0000000..17e2bb8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby_rolling.py @@ -0,0 +1,118 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class OpsOnDiffFramesGroupByRollingTestsMixin: + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("compute.ops_on_diff_frames", True) + + @classmethod + def tearDownClass(cls): + reset_option("compute.ops_on_diff_frames") + super().tearDownClass() + + def _test_groupby_rolling_func(self, f): + pser = pd.Series([1, 2, 3], name="a") + pkey = pd.Series([1, 2, 3], name="a") + psser = ps.from_pandas(pser) + kkey = ps.from_pandas(pkey) + + self.assert_eq( + getattr(psser.groupby(kkey).rolling(2), f)().sort_index(), + getattr(pser.groupby(pkey).rolling(2), f)().sort_index(), + ) + + pdf = pd.DataFrame({"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0]}) + pkey = pd.Series([1, 2, 3, 2], name="a") + psdf = ps.from_pandas(pdf) + kkey = ps.from_pandas(pkey) + + # The behavior of GroupBy.rolling is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + getattr(psdf.groupby(kkey).rolling(2), f)().sort_index(), + getattr(pdf.groupby(pkey).rolling(2), f)().sort_index(), + ) + else: + self.assert_eq( + getattr(psdf.groupby(kkey).rolling(2), f)().sort_index(), + getattr(pdf.groupby(pkey).rolling(2), f)().drop("a", axis=1).sort_index(), + ) + + self.assert_eq( + getattr(psdf.groupby(kkey)["b"].rolling(2), f)().sort_index(), + getattr(pdf.groupby(pkey)["b"].rolling(2), f)().sort_index(), + ) + self.assert_eq( + getattr(psdf.groupby(kkey)[["b"]].rolling(2), f)().sort_index(), + getattr(pdf.groupby(pkey)[["b"]].rolling(2), f)().sort_index(), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43452): Enable RollingTests.test_groupby_rolling_count for pandas 2.0.0.", + ) + def test_groupby_rolling_count(self): + self._test_groupby_rolling_func("count") + + def test_groupby_rolling_min(self): + self._test_groupby_rolling_func("min") + + def test_groupby_rolling_max(self): + self._test_groupby_rolling_func("max") + + def test_groupby_rolling_mean(self): + self._test_groupby_rolling_func("mean") + + def test_groupby_rolling_sum(self): + self._test_groupby_rolling_func("sum") + + def test_groupby_rolling_std(self): + # TODO: `std` now raise error in pandas 1.0.0 + self._test_groupby_rolling_func("std") + + def test_groupby_rolling_var(self): + self._test_groupby_rolling_func("var") + + +class OpsOnDiffFramesGroupByRollingTests( + OpsOnDiffFramesGroupByRollingTestsMixin, PandasOnSparkTestCase, TestUtils +): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_ops_on_diff_frames_groupby_rolling import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_repr.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_repr.py new file mode 100644 index 0000000..fba7fa9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_repr.py @@ -0,0 +1,192 @@ +# +# 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 numpy as np + +from pyspark import pandas as ps +from pyspark.pandas.config import set_option, reset_option, option_context +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class ReprTestsMixin: + max_display_count = 23 + + @classmethod + def setUpClass(cls): + super().setUpClass() + set_option("display.max_rows", ReprTests.max_display_count) + + @classmethod + def tearDownClass(cls): + reset_option("display.max_rows") + super().tearDownClass() + + def test_repr_dataframe(self): + psdf = ps.range(ReprTests.max_display_count) + self.assertTrue("Showing only the first" not in repr(psdf)) + self.assert_eq(repr(psdf), repr(psdf._to_pandas())) + + psdf = ps.range(ReprTests.max_display_count + 1) + self.assertTrue("Showing only the first" in repr(psdf)) + self.assertTrue( + repr(psdf).startswith(repr(psdf._to_pandas().head(ReprTests.max_display_count))) + ) + + with option_context("display.max_rows", None): + psdf = ps.range(ReprTests.max_display_count + 1) + self.assert_eq(repr(psdf), repr(psdf._to_pandas())) + + def test_repr_series(self): + psser = ps.range(ReprTests.max_display_count).id + self.assertTrue("Showing only the first" not in repr(psser)) + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + psser = ps.range(ReprTests.max_display_count + 1).id + self.assertTrue("Showing only the first" in repr(psser)) + self.assertTrue( + repr(psser).startswith(repr(psser._to_pandas().head(ReprTests.max_display_count))) + ) + + with option_context("display.max_rows", None): + psser = ps.range(ReprTests.max_display_count + 1).id + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + psser = ps.range(ReprTests.max_display_count).id.rename() + self.assertTrue("Showing only the first" not in repr(psser)) + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + psser = ps.range(ReprTests.max_display_count + 1).id.rename() + self.assertTrue("Showing only the first" in repr(psser)) + self.assertTrue( + repr(psser).startswith(repr(psser._to_pandas().head(ReprTests.max_display_count))) + ) + + with option_context("display.max_rows", None): + psser = ps.range(ReprTests.max_display_count + 1).id.rename() + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + psser = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count)] + ).to_series() + self.assertTrue("Showing only the first" not in repr(psser)) + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + psser = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count + 1)] + ).to_series() + self.assertTrue("Showing only the first" in repr(psser)) + self.assertTrue( + repr(psser).startswith(repr(psser._to_pandas().head(ReprTests.max_display_count))) + ) + + with option_context("display.max_rows", None): + psser = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count + 1)] + ).to_series() + self.assert_eq(repr(psser), repr(psser._to_pandas())) + + def test_repr_indexes(self): + psidx = ps.range(ReprTests.max_display_count).index + self.assertTrue("Showing only the first" not in repr(psidx)) + self.assert_eq(repr(psidx), repr(psidx._to_pandas())) + + psidx = ps.range(ReprTests.max_display_count + 1).index + self.assertTrue("Showing only the first" in repr(psidx)) + self.assertTrue( + repr(psidx).startswith( + repr(psidx._to_pandas().to_series().head(ReprTests.max_display_count).index) + ) + ) + + with option_context("display.max_rows", None): + psidx = ps.range(ReprTests.max_display_count + 1).index + self.assert_eq(repr(psidx), repr(psidx._to_pandas())) + + psidx = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count)] + ) + self.assertTrue("Showing only the first" not in repr(psidx)) + self.assert_eq(repr(psidx), repr(psidx._to_pandas())) + + psidx = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count + 1)] + ) + self.assertTrue("Showing only the first" in repr(psidx)) + self.assertTrue( + repr(psidx).startswith( + repr(psidx._to_pandas().to_frame().head(ReprTests.max_display_count).index) + ) + ) + + with option_context("display.max_rows", None): + psidx = ps.MultiIndex.from_tuples( + [(100 * i, i) for i in range(ReprTests.max_display_count + 1)] + ) + self.assert_eq(repr(psidx), repr(psidx._to_pandas())) + + def test_html_repr(self): + psdf = ps.range(ReprTests.max_display_count) + self.assertTrue("Showing only the first" not in psdf._repr_html_()) + self.assertEqual(psdf._repr_html_(), psdf._to_pandas()._repr_html_()) + + psdf = ps.range(ReprTests.max_display_count + 1) + self.assertTrue("Showing only the first" in psdf._repr_html_()) + + with option_context("display.max_rows", None): + psdf = ps.range(ReprTests.max_display_count + 1) + self.assertEqual(psdf._repr_html_(), psdf._to_pandas()._repr_html_()) + + def test_repr_float_index(self): + psdf = ps.DataFrame( + {"a": np.random.rand(ReprTests.max_display_count)}, + index=np.random.rand(ReprTests.max_display_count), + ) + self.assertTrue("Showing only the first" not in repr(psdf)) + self.assert_eq(repr(psdf), repr(psdf._to_pandas())) + self.assertTrue("Showing only the first" not in repr(psdf.a)) + self.assert_eq(repr(psdf.a), repr(psdf.a._to_pandas())) + self.assertTrue("Showing only the first" not in repr(psdf.index)) + self.assert_eq(repr(psdf.index), repr(psdf.index._to_pandas())) + + self.assertTrue("Showing only the first" not in psdf._repr_html_()) + self.assertEqual(psdf._repr_html_(), psdf._to_pandas()._repr_html_()) + + psdf = ps.DataFrame( + {"a": np.random.rand(ReprTests.max_display_count + 1)}, + index=np.random.rand(ReprTests.max_display_count + 1), + ) + self.assertTrue("Showing only the first" in repr(psdf)) + self.assertTrue("Showing only the first" in repr(psdf.a)) + self.assertTrue("Showing only the first" in repr(psdf.index)) + self.assertTrue("Showing only the first" in psdf._repr_html_()) + + +class ReprTests(ReprTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_repr import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_resample.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_resample.py new file mode 100644 index 0000000..4061402 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_resample.py @@ -0,0 +1,346 @@ +# +# 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 unittest +import inspect +import datetime +import os + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError, DataError +from pyspark.pandas.missing.resample import ( + MissingPandasLikeDataFrameResampler, + MissingPandasLikeSeriesResampler, +) +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class ResampleTestsMixin: + @property + def pdf1(self): + np.random.seed(11) + dates = [ + pd.NaT, + datetime.datetime(2011, 12, 31), + datetime.datetime(2011, 12, 31, 0, 0, 1), + datetime.datetime(2011, 12, 31, 23, 59, 59), + datetime.datetime(2012, 1, 1), + datetime.datetime(2012, 1, 1, 0, 0, 1), + pd.NaT, + datetime.datetime(2012, 1, 1, 23, 59, 59), + datetime.datetime(2012, 1, 2), + pd.NaT, + datetime.datetime(2012, 1, 30, 23, 59, 59), + datetime.datetime(2012, 1, 31), + datetime.datetime(2012, 1, 31, 0, 0, 1), + datetime.datetime(2012, 3, 31), + datetime.datetime(2013, 5, 3), + datetime.datetime(2022, 5, 3), + ] + return pd.DataFrame( + np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=list("AB") + ) + + @property + def pdf2(self): + np.random.seed(22) + dates = [ + datetime.datetime(2022, 5, 1, 4, 5, 6), + datetime.datetime(2022, 5, 3), + datetime.datetime(2022, 5, 3, 23, 59, 59), + datetime.datetime(2022, 5, 4), + pd.NaT, + datetime.datetime(2022, 5, 4, 0, 0, 1), + datetime.datetime(2022, 5, 11), + ] + return pd.DataFrame( + np.random.rand(len(dates), 2), index=pd.DatetimeIndex(dates), columns=list("AB") + ) + + @property + def pdf3(self): + np.random.seed(22) + index = pd.date_range(start="2011-01-02", end="2022-05-01", freq="1D") + return pd.DataFrame(np.random.rand(len(index), 2), index=index, columns=list("AB")) + + @property + def pdf4(self): + np.random.seed(33) + index = pd.date_range(start="2020-12-12", end="2022-05-01", freq="1H") + return pd.DataFrame(np.random.rand(len(index), 2), index=index, columns=list("AB")) + + @property + def pdf5(self): + np.random.seed(44) + index = pd.date_range(start="2021-12-30 03:04:05", end="2022-01-02 06:07:08", freq="1T") + return pd.DataFrame(np.random.rand(len(index), 2), index=index, columns=list("AB")) + + @property + def pdf6(self): + np.random.seed(55) + index = pd.date_range(start="2022-05-02 03:04:05", end="2022-05-02 06:07:08", freq="1S") + return pd.DataFrame(np.random.rand(len(index), 2), index=index, columns=list("AB")) + + @property + def psdf1(self): + return ps.from_pandas(self.pdf1) + + @property + def psdf2(self): + return ps.from_pandas(self.pdf2) + + @property + def psdf3(self): + return ps.from_pandas(self.pdf3) + + @property + def psdf4(self): + return ps.from_pandas(self.pdf4) + + @property + def psdf5(self): + return ps.from_pandas(self.pdf5) + + @property + def psdf6(self): + return ps.from_pandas(self.pdf6) + + def test_resample_error(self): + psdf = ps.range(10) + + with self.assertRaisesRegex( + NotImplementedError, "resample currently works only for DatetimeIndex" + ): + psdf.resample("3Y").sum() + + with self.assertRaisesRegex( + NotImplementedError, "resample currently works only for DatetimeIndex" + ): + psdf.id.resample("3Y").sum() + + dates = [ + datetime.datetime(2012, 1, 2), + datetime.datetime(2012, 5, 3), + datetime.datetime(2022, 5, 3), + pd.NaT, + ] + pdf = pd.DataFrame(np.ones(len(dates)), index=pd.DatetimeIndex(dates), columns=["A"]) + psdf = ps.from_pandas(pdf) + + with self.assertRaisesRegex(ValueError, "rule code W-SUN is not supported"): + psdf.A.resample("3W").sum() + + with self.assertRaisesRegex(ValueError, "rule offset must be positive"): + psdf.A.resample("0Y").sum() + + with self.assertRaisesRegex(ValueError, "invalid closed: 'middle'"): + psdf.A.resample("3Y", closed="middle").sum() + + with self.assertRaisesRegex(ValueError, "invalid label: 'both'"): + psdf.A.resample("3Y", label="both").sum() + + with self.assertRaisesRegex( + NotImplementedError, "`on` currently works only for TimestampType" + ): + psdf.A.resample("2D", on=psdf.A).sum() + + with self.assertRaisesRegex( + NotImplementedError, "`on` currently works only for TimestampType" + ): + psdf[["A"]].resample("2D", on=psdf.A).sum() + + psdf["B"] = ["a", "b", "c", "d"] + with self.assertRaisesRegex(ValueError, "No available aggregation columns!"): + psdf.B.resample("2D").sum() + + with self.assertRaisesRegex(ValueError, "No available aggregation columns!"): + psdf[[]].resample("2D").sum() + + def test_missing(self): + pdf_r = self.psdf1.resample("3Y") + pser_r = self.psdf1.A.resample("3Y") + + # DataFrameResampler functions + missing_functions = inspect.getmembers( + MissingPandasLikeDataFrameResampler, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Resampler.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(pdf_r, name)() + + # SeriesResampler functions + missing_functions = inspect.getmembers(MissingPandasLikeSeriesResampler, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Resampler.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(pser_r, name)() + + # DataFrameResampler properties + missing_properties = inspect.getmembers( + MissingPandasLikeDataFrameResampler, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Resampler.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(pdf_r, name) + + # SeriesResampler properties + missing_properties = inspect.getmembers( + MissingPandasLikeSeriesResampler, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Resampler.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(pser_r, name) + + def _test_resample(self, pobj, psobj, rules, closed, label, func): + for rule in rules: + p_resample = pobj.resample(rule=rule, closed=closed, label=label) + ps_resample = psobj.resample(rule=rule, closed=closed, label=label) + self.assert_eq( + getattr(p_resample, func)().sort_index(), + getattr(ps_resample, func)().sort_index(), + almost=True, + ) + + def test_dataframe_resample(self): + self._test_resample(self.pdf1, self.psdf1, ["3Y", "9M", "17D"], None, None, "min") + self._test_resample(self.pdf2, self.psdf2, ["3A", "11M", "D"], None, "left", "max") + self._test_resample(self.pdf3, self.psdf3, ["20D", "1M"], None, "right", "sum") + self._test_resample(self.pdf4, self.psdf4, ["11H", "21D"], "left", None, "mean") + self._test_resample(self.pdf5, self.psdf5, ["55MIN", "2H", "D"], "left", "left", "std") + self._test_resample(self.pdf6, self.psdf6, ["29S", "10MIN", "3H"], "left", "right", "var") + + def test_series_resample(self): + self._test_resample(self.pdf1.A, self.psdf1.A, ["4Y"], "right", None, "min") + self._test_resample(self.pdf2.A, self.psdf2.A, ["13M"], "right", "left", "max") + self._test_resample(self.pdf3.A, self.psdf3.A, ["1001H"], "right", "right", "sum") + self._test_resample(self.pdf4.A, self.psdf4.A, ["6D"], None, None, "mean") + self._test_resample(self.pdf5.A, self.psdf5.A, ["47T"], "left", "left", "var") + self._test_resample(self.pdf6.A, self.psdf6.A, ["111S"], "right", "right", "std") + + def test_resample_on(self): + np.random.seed(77) + dates = [ + datetime.datetime(2022, 5, 1, 4, 5, 6), + datetime.datetime(2022, 5, 3), + datetime.datetime(2022, 5, 3, 23, 59, 59), + datetime.datetime(2022, 5, 4), + pd.NaT, + datetime.datetime(2022, 5, 4, 0, 0, 1), + datetime.datetime(2022, 5, 11), + ] + pdf = pd.DataFrame( + np.random.rand(len(dates), 3), index=pd.DatetimeIndex(dates), columns=list("ABC") + ) + pdf["X"] = pd.DatetimeIndex(dates) + psdf = ps.from_pandas(pdf) + self.assert_eq( + pdf.resample("2D", on="X").sum().sort_index(), + psdf.resample("2D", on=psdf.X).sum().sort_index(), + almost=True, + ) + + +class ResampleWithTimezoneMixin: + timezone = None + + @classmethod + def setUpClass(cls): + cls.timezone = os.environ.get("TZ", None) + os.environ["TZ"] = "America/New_York" + super(ResampleWithTimezoneMixin, cls).setUpClass() + + @classmethod + def tearDownClass(cls): + super(ResampleWithTimezoneMixin, cls).tearDownClass() + if cls.timezone is not None: + os.environ["TZ"] = cls.timezone + + @property + def pdf(self): + np.random.seed(22) + index = pd.date_range(start="2011-01-02", end="2022-05-01", freq="1D") + return pd.DataFrame(np.random.rand(len(index), 2), index=index, columns=list("AB")) + + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + def test_series_resample_with_timezone(self): + with self.sql_conf( + { + "spark.sql.session.timeZone": "Asia/Seoul", + "spark.sql.timestampType": "TIMESTAMP_NTZ", + } + ): + p_resample = self.pdf.resample(rule="1001H", closed="right", label="right") + ps_resample = self.psdf.resample(rule="1001H", closed="right", label="right") + self.assert_eq( + p_resample.sum().sort_index(), + ps_resample.sum().sort_index(), + almost=True, + ) + + +class ResampleTests(ResampleTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +class ResampleWithTimezoneTests(ResampleWithTimezoneMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_resample import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_reshape.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_reshape.py new file mode 100644 index 0000000..b4ebba7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_reshape.py @@ -0,0 +1,495 @@ +# +# 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 datetime +from decimal import Decimal +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.pandas.utils import name_like_string +from pyspark.errors import AnalysisException +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class ReshapeTestsMixin: + def test_get_dummies(self): + for pdf_or_ps in [ + pd.Series([1, 1, 1, 2, 2, 1, 3, 4]), + # pd.Series([1, 1, 1, 2, 2, 1, 3, 4], dtype='category'), + # pd.Series(pd.Categorical([1, 1, 1, 2, 2, 1, 3, 4], + # categories=[4, 3, 2, 1])), + pd.DataFrame( + { + "a": [1, 2, 3, 4, 4, 3, 2, 1], + # 'b': pd.Categorical(list('abcdabcd')), + "b": list("abcdabcd"), + } + ), + pd.DataFrame({10: [1, 2, 3, 4, 4, 3, 2, 1], 20: list("abcdabcd")}), + ]: + psdf_or_psser = ps.from_pandas(pdf_or_ps) + + self.assert_eq(ps.get_dummies(psdf_or_psser), pd.get_dummies(pdf_or_ps, dtype=np.int8)) + + psser = ps.Series([1, 1, 1, 2, 2, 1, 3, 4]) + with self.assertRaisesRegex( + NotImplementedError, "get_dummies currently does not support sparse" + ): + ps.get_dummies(psser, sparse=True) + with self.assertRaisesRegex(NotImplementedError, "get_dummies currently only accept"): + ps.get_dummies(ps.Series([b"1"])) + with self.assertRaisesRegex(NotImplementedError, "get_dummies currently only accept"): + ps.get_dummies(ps.Series([None])) + + def test_get_dummies_object(self): + pdf = pd.DataFrame( + { + "a": [1, 2, 3, 4, 4, 3, 2, 1], + # 'a': pd.Categorical([1, 2, 3, 4, 4, 3, 2, 1]), + "b": list("abcdabcd"), + # 'c': pd.Categorical(list('abcdabcd')), + "c": list("abcdabcd"), + } + ) + psdf = ps.from_pandas(pdf) + + # Explicitly exclude object columns + self.assert_eq( + ps.get_dummies(psdf, columns=["a", "c"]), + pd.get_dummies(pdf, columns=["a", "c"], dtype=np.int8), + ) + + self.assert_eq(ps.get_dummies(psdf), pd.get_dummies(pdf, dtype=np.int8)) + self.assert_eq(ps.get_dummies(psdf.b), pd.get_dummies(pdf.b, dtype=np.int8)) + self.assert_eq( + ps.get_dummies(psdf, columns=["b"]), pd.get_dummies(pdf, columns=["b"], dtype=np.int8) + ) + + self.assertRaises(KeyError, lambda: ps.get_dummies(psdf, columns=("a", "c"))) + self.assertRaises(TypeError, lambda: ps.get_dummies(psdf, columns="b")) + + # non-string names + pdf = pd.DataFrame( + {10: [1, 2, 3, 4, 4, 3, 2, 1], 20: list("abcdabcd"), 30: list("abcdabcd")} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.get_dummies(psdf, columns=[10, 30]), + pd.get_dummies(pdf, columns=[10, 30], dtype=np.int8), + ) + + self.assertRaises(TypeError, lambda: ps.get_dummies(psdf, columns=10)) + + def test_get_dummies_date_datetime(self): + pdf = pd.DataFrame( + { + "d": [ + datetime.date(2019, 1, 1), + datetime.date(2019, 1, 2), + datetime.date(2019, 1, 1), + ], + "dt": [ + datetime.datetime(2019, 1, 1, 0, 0, 0), + datetime.datetime(2019, 1, 1, 0, 0, 1), + datetime.datetime(2019, 1, 1, 0, 0, 0), + ], + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(ps.get_dummies(psdf), pd.get_dummies(pdf, dtype=np.int8)) + self.assert_eq(ps.get_dummies(psdf.d), pd.get_dummies(pdf.d, dtype=np.int8)) + self.assert_eq(ps.get_dummies(psdf.dt), pd.get_dummies(pdf.dt, dtype=np.int8)) + + def test_get_dummies_boolean(self): + pdf = pd.DataFrame({"b": [True, False, True]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(ps.get_dummies(psdf), pd.get_dummies(pdf, dtype=np.int8)) + self.assert_eq(ps.get_dummies(psdf.b), pd.get_dummies(pdf.b, dtype=np.int8)) + + def test_get_dummies_decimal(self): + pdf = pd.DataFrame({"d": [Decimal(1.0), Decimal(2.0), Decimal(1)]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(ps.get_dummies(psdf), pd.get_dummies(pdf, dtype=np.int8)) + self.assert_eq(ps.get_dummies(psdf.d), pd.get_dummies(pdf.d, dtype=np.int8), almost=True) + + def test_get_dummies_kwargs(self): + # pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4], dtype='category') + pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4]) + psser = ps.from_pandas(pser) + self.assert_eq( + ps.get_dummies(psser, prefix="X", prefix_sep="-"), + pd.get_dummies(pser, prefix="X", prefix_sep="-", dtype=np.int8), + ) + + self.assert_eq( + ps.get_dummies(psser, drop_first=True), + pd.get_dummies(pser, drop_first=True, dtype=np.int8), + ) + + # nan + # pser = pd.Series([1, 1, 1, 2, np.nan, 3, np.nan, 5], dtype='category') + pser = pd.Series([1, 1, 1, 2, np.nan, 3, np.nan, 5]) + psser = ps.from_pandas(pser) + self.assert_eq(ps.get_dummies(psser), pd.get_dummies(pser, dtype=np.int8), almost=True) + + # dummy_na + self.assert_eq( + ps.get_dummies(psser, dummy_na=True), pd.get_dummies(pser, dummy_na=True, dtype=np.int8) + ) + + def test_get_dummies_prefix(self): + pdf = pd.DataFrame({"A": ["a", "b", "a"], "B": ["b", "a", "c"], "D": [0, 0, 1]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.get_dummies(psdf, prefix=["foo", "bar"]), + pd.get_dummies(pdf, prefix=["foo", "bar"], dtype=np.int8), + ) + + self.assert_eq( + ps.get_dummies(psdf, prefix=["foo"], columns=["B"]), + pd.get_dummies(pdf, prefix=["foo"], columns=["B"], dtype=np.int8), + ) + + self.assert_eq( + ps.get_dummies(psdf, prefix={"A": "foo", "B": "bar"}), + pd.get_dummies(pdf, prefix={"A": "foo", "B": "bar"}, dtype=np.int8), + ) + + self.assert_eq( + ps.get_dummies(psdf, prefix={"B": "foo", "A": "bar"}), + pd.get_dummies(pdf, prefix={"B": "foo", "A": "bar"}, dtype=np.int8), + ) + + self.assert_eq( + ps.get_dummies(psdf, prefix={"A": "foo", "B": "bar"}, columns=["A", "B"]), + pd.get_dummies(pdf, prefix={"A": "foo", "B": "bar"}, columns=["A", "B"], dtype=np.int8), + ) + + with self.assertRaisesRegex(NotImplementedError, "string types"): + ps.get_dummies(psdf, prefix="foo") + with self.assertRaisesRegex(ValueError, "Length of 'prefix' \\(1\\) .* \\(2\\)"): + ps.get_dummies(psdf, prefix=["foo"]) + with self.assertRaisesRegex(ValueError, "Length of 'prefix' \\(2\\) .* \\(1\\)"): + ps.get_dummies(psdf, prefix=["foo", "bar"], columns=["B"]) + + pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4], name="A") + psser = ps.from_pandas(pser) + + self.assert_eq( + ps.get_dummies(psser, prefix="foo"), pd.get_dummies(pser, prefix="foo", dtype=np.int8) + ) + + # columns are ignored. + self.assert_eq( + ps.get_dummies(psser, prefix=["foo"], columns=["B"]), + pd.get_dummies(pser, prefix=["foo"], columns=["B"], dtype=np.int8), + ) + + def test_get_dummies_dtype(self): + pdf = pd.DataFrame( + { + # "A": pd.Categorical(['a', 'b', 'a'], categories=['a', 'b', 'c']), + "A": ["a", "b", "a"], + "B": [0, 0, 1], + } + ) + psdf = ps.from_pandas(pdf) + + exp = pd.get_dummies(pdf) + exp = exp.astype({"A_a": "float64", "A_b": "float64"}) + res = ps.get_dummies(psdf, dtype="float64") + self.assert_eq(res, exp) + + def test_get_dummies_multiindex_columns(self): + pdf = pd.DataFrame( + { + ("x", "a", "1"): [1, 2, 3, 4, 4, 3, 2, 1], + ("x", "b", "2"): list("abcdabcd"), + ("y", "c", "3"): list("abcdabcd"), + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.get_dummies(psdf), + pd.get_dummies(pdf, dtype=np.int8).rename(columns=name_like_string), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=[("y", "c", "3"), ("x", "a", "1")]), + pd.get_dummies(pdf, columns=[("y", "c", "3"), ("x", "a", "1")], dtype=np.int8).rename( + columns=name_like_string + ), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=["x"]), + pd.get_dummies(pdf, columns=["x"], dtype=np.int8).rename(columns=name_like_string), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=("x", "a")), + pd.get_dummies(pdf, columns=("x", "a"), dtype=np.int8).rename(columns=name_like_string), + ) + + self.assertRaises(KeyError, lambda: ps.get_dummies(psdf, columns=["z"])) + self.assertRaises(KeyError, lambda: ps.get_dummies(psdf, columns=("x", "c"))) + self.assertRaises(ValueError, lambda: ps.get_dummies(psdf, columns=[("x",), "c"])) + self.assertRaises(TypeError, lambda: ps.get_dummies(psdf, columns="x")) + + # non-string names + pdf = pd.DataFrame( + { + ("x", 1, "a"): [1, 2, 3, 4, 4, 3, 2, 1], + ("x", 2, "b"): list("abcdabcd"), + ("y", 3, "c"): list("abcdabcd"), + } + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + ps.get_dummies(psdf), + pd.get_dummies(pdf, dtype=np.int8).rename(columns=name_like_string), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=[("y", 3, "c"), ("x", 1, "a")]), + pd.get_dummies(pdf, columns=[("y", 3, "c"), ("x", 1, "a")], dtype=np.int8).rename( + columns=name_like_string + ), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=["x"]), + pd.get_dummies(pdf, columns=["x"], dtype=np.int8).rename(columns=name_like_string), + ) + self.assert_eq( + ps.get_dummies(psdf, columns=("x", 1)), + pd.get_dummies(pdf, columns=("x", 1), dtype=np.int8).rename(columns=name_like_string), + ) + + def test_merge_asof(self): + pdf_left = pd.DataFrame( + {"a": [1, 5, 10], "b": ["x", "y", "z"], "left_val": ["a", "b", "c"]}, index=[10, 20, 30] + ) + pdf_right = pd.DataFrame( + {"a": [1, 2, 3, 6, 7], "b": ["v", "w", "x", "y", "z"], "right_val": [1, 2, 3, 6, 7]}, + index=[100, 101, 102, 103, 104], + ) + psdf_left = ps.from_pandas(pdf_left) + psdf_right = ps.from_pandas(pdf_right) + + self.assert_eq( + pd.merge_asof(pdf_left, pdf_right, on="a").sort_values("a").reset_index(drop=True), + ps.merge_asof(psdf_left, psdf_right, on="a").sort_values("a").reset_index(drop=True), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, left_on="a", right_on="a") + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, left_on="a", right_on="a") + .sort_values("a") + .reset_index(drop=True) + ), + ) + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + pd.merge_asof( + pdf_left.set_index("a"), pdf_right, left_index=True, right_on="a" + ).sort_index(), + ps.merge_asof( + psdf_left.set_index("a"), psdf_right, left_index=True, right_on="a" + ).sort_index(), + ) + else: + expected = pd.DataFrame( + { + "b_x": ["x", "y", "z"], + "left_val": ["a", "b", "c"], + "a": [1, 3, 7], + "b_y": ["v", "x", "z"], + "right_val": [1, 3, 7], + }, + index=pd.Index([1, 5, 10], name="a"), + ) + self.assert_eq( + expected, + ps.merge_asof( + psdf_left.set_index("a"), psdf_right, left_index=True, right_on="a" + ).sort_index(), + ) + self.assert_eq( + pd.merge_asof( + pdf_left, pdf_right.set_index("a"), left_on="a", right_index=True + ).sort_index(), + ps.merge_asof( + psdf_left, psdf_right.set_index("a"), left_on="a", right_index=True + ).sort_index(), + ) + self.assert_eq( + pd.merge_asof( + pdf_left.set_index("a"), pdf_right.set_index("a"), left_index=True, right_index=True + ).sort_index(), + ps.merge_asof( + psdf_left.set_index("a"), + psdf_right.set_index("a"), + left_index=True, + right_index=True, + ).sort_index(), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, on="a", by="b") + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, on="a", by="b") + .sort_values("a") + .reset_index(drop=True) + ), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, on="a", tolerance=1) + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, on="a", tolerance=1) + .sort_values("a") + .reset_index(drop=True) + ), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, on="a", allow_exact_matches=False) + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, on="a", allow_exact_matches=False) + .sort_values("a") + .reset_index(drop=True) + ), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, on="a", direction="forward") + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, on="a", direction="forward") + .sort_values("a") + .reset_index(drop=True) + ), + ) + self.assert_eq( + ( + pd.merge_asof(pdf_left, pdf_right, on="a", direction="nearest") + .sort_values("a") + .reset_index(drop=True) + ), + ( + ps.merge_asof(psdf_left, psdf_right, on="a", direction="nearest") + .sort_values("a") + .reset_index(drop=True) + ), + ) + # Including Series + self.assert_eq( + pd.merge_asof(pdf_left["a"], pdf_right, on="a").sort_values("a").reset_index(drop=True), + ps.merge_asof(psdf_left["a"], psdf_right, on="a") + .sort_values("a") + .reset_index(drop=True), + ) + self.assert_eq( + pd.merge_asof(pdf_left, pdf_right["a"], on="a").sort_values("a").reset_index(drop=True), + ps.merge_asof(psdf_left, psdf_right["a"], on="a") + .sort_values("a") + .reset_index(drop=True), + ) + self.assert_eq( + pd.merge_asof(pdf_left["a"], pdf_right["a"], on="a") + .sort_values("a") + .reset_index(drop=True), + ps.merge_asof(psdf_left["a"], psdf_right["a"], on="a") + .sort_values("a") + .reset_index(drop=True), + ) + + self.assertRaises( + AnalysisException, lambda: ps.merge_asof(psdf_left, psdf_right, on="a", tolerance=-1) + ) + with self.assertRaisesRegex( + ValueError, + 'Can only pass argument "on" OR "left_on" and "right_on", not a combination of both.', + ): + ps.merge_asof(psdf_left, psdf_right, on="a", left_on="a") + psdf_multi_index = ps.DataFrame( + {"a": [1, 2, 3, 6, 7], "b": ["v", "w", "x", "y", "z"], "right_val": [1, 2, 3, 6, 7]}, + index=pd.MultiIndex.from_tuples([(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]), + ) + with self.assertRaisesRegex(ValueError, "right can only have one index"): + ps.merge_asof(psdf_left, psdf_multi_index, right_index=True) + with self.assertRaisesRegex(ValueError, "left can only have one index"): + ps.merge_asof(psdf_multi_index, psdf_right, left_index=True) + with self.assertRaisesRegex(ValueError, "Must pass right_on or right_index=True"): + ps.merge_asof(psdf_left, psdf_right, left_index=True) + with self.assertRaisesRegex(ValueError, "Must pass left_on or left_index=True"): + ps.merge_asof(psdf_left, psdf_right, right_index=True) + with self.assertRaisesRegex(ValueError, "can only asof on a key for left"): + ps.merge_asof(psdf_left, psdf_right, right_on="a", left_on=["a", "b"]) + with self.assertRaisesRegex(ValueError, "can only asof on a key for right"): + ps.merge_asof(psdf_left, psdf_right, right_on=["a", "b"], left_on="a") + with self.assertRaisesRegex( + ValueError, 'Can only pass argument "by" OR "left_by" and "right_by".' + ): + ps.merge_asof(psdf_left, psdf_right, on="a", by="b", left_by="a") + with self.assertRaisesRegex(ValueError, "missing right_by"): + ps.merge_asof(psdf_left, psdf_right, on="a", left_by="b") + with self.assertRaisesRegex(ValueError, "missing left_by"): + ps.merge_asof(psdf_left, psdf_right, on="a", right_by="b") + with self.assertRaisesRegex(ValueError, "left_by and right_by must be same length"): + ps.merge_asof(psdf_left, psdf_right, on="a", left_by="b", right_by=["a", "b"]) + psdf_right.columns = ["A", "B", "C"] + with self.assertRaisesRegex(ValueError, "No common columns to perform merge on."): + ps.merge_asof(psdf_left, psdf_right) + + +class ReshapeTests(ReshapeTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_reshape import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_rolling.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_rolling.py new file mode 100644 index 0000000..00b9de8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_rolling.py @@ -0,0 +1,263 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +import pyspark.pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils +from pyspark.pandas.window import Rolling + + +class RollingTestsMixin: + def test_rolling_error(self): + with self.assertRaisesRegex(ValueError, "window must be >= 0"): + ps.range(10).rolling(window=-1) + with self.assertRaisesRegex(ValueError, "min_periods must be >= 0"): + ps.range(10).rolling(window=1, min_periods=-1) + + with self.assertRaisesRegex( + TypeError, "psdf_or_psser must be a series or dataframe; however, got:.*int" + ): + Rolling(1, 2) + + def _test_rolling_func(self, ps_func, pd_func=None): + if not pd_func: + pd_func = ps_func + if isinstance(pd_func, str): + pd_func = self.convert_str_to_lambda(pd_func) + if isinstance(ps_func, str): + ps_func = self.convert_str_to_lambda(ps_func) + pser = pd.Series([1, 2, 3, 7, 9, 8], index=np.random.rand(6), name="a") + psser = ps.from_pandas(pser) + self.assert_eq(ps_func(psser.rolling(2)), pd_func(pser.rolling(2))) + self.assert_eq(ps_func(psser.rolling(2)).sum(), pd_func(pser.rolling(2)).sum()) + + # Multiindex + pser = pd.Series( + [1, 2, 3], + index=pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")]), + name="a", + ) + psser = ps.from_pandas(pser) + self.assert_eq(ps_func(psser.rolling(2)), pd_func(pser.rolling(2))) + + pdf = pd.DataFrame( + {"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}, index=np.random.rand(4) + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(ps_func(psdf.rolling(2)), pd_func(pdf.rolling(2))) + self.assert_eq(ps_func(psdf.rolling(2)).sum(), pd_func(pdf.rolling(2)).sum()) + + # Multiindex column + columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")]) + pdf.columns = columns + psdf.columns = columns + self.assert_eq(ps_func(psdf.rolling(2)), pd_func(pdf.rolling(2))) + + def test_rolling_min(self): + self._test_rolling_func("min") + + def test_rolling_max(self): + self._test_rolling_func("max") + + def test_rolling_mean(self): + self._test_rolling_func("mean") + + def test_rolling_quantile(self): + self._test_rolling_func(lambda x: x.quantile(0.5), lambda x: x.quantile(0.5, "lower")) + + def test_rolling_sum(self): + self._test_rolling_func("sum") + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43451): Enable RollingTests.test_rolling_count for pandas 2.0.0.", + ) + def test_rolling_count(self): + self._test_rolling_func("count") + + def test_rolling_std(self): + self._test_rolling_func("std") + + def test_rolling_var(self): + self._test_rolling_func("var") + + def test_rolling_skew(self): + self._test_rolling_func("skew") + + def test_rolling_kurt(self): + self._test_rolling_func("kurt") + + def _test_groupby_rolling_func(self, ps_func, pd_func=None): + if not pd_func: + pd_func = ps_func + if isinstance(pd_func, str): + pd_func = self.convert_str_to_lambda(pd_func) + if isinstance(ps_func, str): + ps_func = self.convert_str_to_lambda(ps_func) + pser = pd.Series([1, 2, 3, 2], index=np.random.rand(4), name="a") + psser = ps.from_pandas(pser) + self.assert_eq( + ps_func(psser.groupby(psser).rolling(2)).sort_index(), + pd_func(pser.groupby(pser).rolling(2)).sort_index(), + ) + self.assert_eq( + ps_func(psser.groupby(psser).rolling(2)).sum(), + pd_func(pser.groupby(pser).rolling(2)).sum(), + ) + + # Multiindex + pser = pd.Series( + [1, 2, 3, 2], + index=pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z"), ("c", "z")]), + name="a", + ) + psser = ps.from_pandas(pser) + self.assert_eq( + ps_func(psser.groupby(psser).rolling(2)).sort_index(), + pd_func(pser.groupby(pser).rolling(2)).sort_index(), + ) + + pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0, 2.0], "b": [4.0, 2.0, 3.0, 1.0]}) + psdf = ps.from_pandas(pdf) + + # The behavior of GroupBy.rolling is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + ps_func(psdf.groupby(psdf.a).rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a).rolling(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a).rolling(2)).sum(), + pd_func(pdf.groupby(pdf.a).rolling(2)).sum(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a + 1).rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a + 1).rolling(2)).sort_index(), + ) + else: + self.assert_eq( + ps_func(psdf.groupby(psdf.a).rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a).rolling(2)).drop("a", axis=1).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a).rolling(2)).sum(), + pd_func(pdf.groupby(pdf.a).rolling(2)).sum().drop("a"), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a + 1).rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a + 1).rolling(2)).drop("a", axis=1).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.b.groupby(psdf.a).rolling(2)).sort_index(), + pd_func(pdf.b.groupby(pdf.a).rolling(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a)["b"].rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a)["b"].rolling(2)).sort_index(), + ) + self.assert_eq( + ps_func(psdf.groupby(psdf.a)[["b"]].rolling(2)).sort_index(), + pd_func(pdf.groupby(pdf.a)[["b"]].rolling(2)).sort_index(), + ) + + # Multiindex column + columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")]) + pdf.columns = columns + psdf.columns = columns + + # The behavior of GroupBy.rolling is changed from pandas 1.3. + if LooseVersion(pd.__version__) >= LooseVersion("1.3"): + self.assert_eq( + ps_func(psdf.groupby(("a", "x")).rolling(2)).sort_index(), + pd_func(pdf.groupby(("a", "x")).rolling(2)).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.groupby([("a", "x"), ("a", "y")]).rolling(2)).sort_index(), + pd_func(pdf.groupby([("a", "x"), ("a", "y")]).rolling(2)).sort_index(), + ) + else: + self.assert_eq( + ps_func(psdf.groupby(("a", "x")).rolling(2)).sort_index(), + pd_func(pdf.groupby(("a", "x")).rolling(2)).drop(("a", "x"), axis=1).sort_index(), + ) + + self.assert_eq( + ps_func(psdf.groupby([("a", "x"), ("a", "y")]).rolling(2)).sort_index(), + pd_func(pdf.groupby([("a", "x"), ("a", "y")]).rolling(2)) + .drop([("a", "x"), ("a", "y")], axis=1) + .sort_index(), + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43452): Enable RollingTests.test_groupby_rolling_count for pandas 2.0.0.", + ) + def test_groupby_rolling_count(self): + self._test_groupby_rolling_func("count") + + def test_groupby_rolling_min(self): + self._test_groupby_rolling_func("min") + + def test_groupby_rolling_max(self): + self._test_groupby_rolling_func("max") + + def test_groupby_rolling_mean(self): + self._test_groupby_rolling_func("mean") + + def test_groupby_rolling_quantile(self): + self._test_groupby_rolling_func( + lambda x: x.quantile(0.5), lambda x: x.quantile(0.5, "lower") + ) + + def test_groupby_rolling_sum(self): + self._test_groupby_rolling_func("sum") + + def test_groupby_rolling_std(self): + # TODO: `std` now raise error in pandas 1.0.0 + self._test_groupby_rolling_func("std") + + def test_groupby_rolling_var(self): + self._test_groupby_rolling_func("var") + + def test_groupby_rolling_skew(self): + self._test_groupby_rolling_func("skew") + + def test_groupby_rolling_kurt(self): + self._test_groupby_rolling_func("kurt") + + +class RollingTests(RollingTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_rolling import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_scalars.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_scalars.py new file mode 100644 index 0000000..5dd8d4c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_scalars.py @@ -0,0 +1,59 @@ +# +# 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 inspect + +import pyspark.pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.scalars import MissingPandasLikeScalars +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class ScalarTestsMixin: + def test_missing(self): + missing_scalars = inspect.getmembers(MissingPandasLikeScalars) + + missing_scalars = [ + name + for (name, type_) in missing_scalars + if isinstance(type_, PandasNotImplementedError) + ] + + for scalar_name in missing_scalars: + with self.assertRaisesRegex( + PandasNotImplementedError, + "The scalar `ps.{0}` is not reimplemented in pyspark.pandas;" + " use `pd.{0}`.".format(scalar_name), + ): + getattr(ps, scalar_name) + + +class ScalarTests(ScalarTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_scalars import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_conversion.py new file mode 100644 index 0000000..cbdb02d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_conversion.py @@ -0,0 +1,85 @@ +# +# 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 unittest +import sys +from distutils.version import LooseVersion + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesConversionTestsMixin: + @property + def pser(self): + return pd.Series([1, 2, 3, 4, 5, 6, 7], name="x") + + @property + def psser(self): + return ps.from_pandas(self.pser) + + @unittest.skipIf( + sys.platform == "linux" or sys.platform == "linux2", + "Pyperclip could not find a copy/paste mechanism for Linux.", + ) + def test_to_clipboard(self): + pser = self.pser + psser = self.psser + + self.assert_eq(psser.to_clipboard(), pser.to_clipboard()) + self.assert_eq(psser.to_clipboard(excel=False), pser.to_clipboard(excel=False)) + self.assert_eq( + psser.to_clipboard(sep=",", index=False), pser.to_clipboard(sep=",", index=False) + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43458): Enable SeriesConversionTests.test_to_latex for pandas 2.0.0.", + ) + def test_to_latex(self): + pser = self.pser + psser = self.psser + + self.assert_eq(psser.to_latex(), pser.to_latex()) + self.assert_eq(psser.to_latex(col_space=2), pser.to_latex(col_space=2)) + self.assert_eq(psser.to_latex(header=True), pser.to_latex(header=True)) + self.assert_eq(psser.to_latex(index=False), pser.to_latex(index=False)) + self.assert_eq(psser.to_latex(na_rep="-"), pser.to_latex(na_rep="-")) + self.assert_eq(psser.to_latex(float_format="%.1f"), pser.to_latex(float_format="%.1f")) + self.assert_eq(psser.to_latex(sparsify=False), pser.to_latex(sparsify=False)) + self.assert_eq(psser.to_latex(index_names=False), pser.to_latex(index_names=False)) + self.assert_eq(psser.to_latex(bold_rows=True), pser.to_latex(bold_rows=True)) + self.assert_eq(psser.to_latex(decimal=","), pser.to_latex(decimal=",")) + + +class SeriesConversionTests(SeriesConversionTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_series_conversion import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_datetime.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_datetime.py new file mode 100644 index 0000000..918176b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_datetime.py @@ -0,0 +1,364 @@ +# +# 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 datetime +import unittest +from distutils.version import LooseVersion + +import numpy as np +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesDateTimeTestsMixin: + @property + def pdf1(self): + date1 = pd.Series(pd.date_range("2012-1-1 12:45:31", periods=3, freq="M")) + date2 = pd.Series(pd.date_range("2013-3-11 21:45:00", periods=3, freq="W")) + return pd.DataFrame(dict(start_date=date1, end_date=date2)) + + @property + def pd_start_date(self): + return self.pdf1["start_date"] + + @property + def ks_start_date(self): + return ps.from_pandas(self.pd_start_date) + + def check_func(self, func): + self.assert_eq(func(self.ks_start_date), func(self.pd_start_date)) + + @unittest.skip( + "Those fail in certain OSs presumably due to different" + "timezone behaviours inherited from C library." + ) + def test_timestamp_subtraction(self): + pdf = self.pdf1 + psdf = ps.from_pandas(pdf) + + actual = (psdf["end_date"] - psdf["start_date"] - 1)._to_pandas() + expected = (pdf["end_date"] - pdf["start_date"]) // np.timedelta64(1, "s") - 1 + self.assert_eq(actual, expected) + + actual = (psdf["end_date"] - pd.Timestamp("2012-1-1 12:45:31") - 1)._to_pandas() + expected = (pdf["end_date"] - pd.Timestamp("2012-1-1 12:45:31")) // np.timedelta64( + 1, "s" + ) - 1 + self.assert_eq(actual, expected) + + actual = (pd.Timestamp("2013-3-11 21:45:00") - psdf["start_date"] - 1)._to_pandas() + expected = (pd.Timestamp("2013-3-11 21:45:00") - pdf["start_date"]) // np.timedelta64( + 1, "s" + ) - 1 + self.assert_eq(actual, expected) + + def test_timestamp_subtraction_errors(self): + psdf = ps.DataFrame( + {"a": pd.date_range("2016-12-31", "2017-01-08", freq="D"), "b": pd.Series(range(9))} + ) + expected_error_message = "Datetime subtraction can only be applied to datetime series." + with self.assertRaisesRegex(TypeError, expected_error_message): + psdf["a"] - psdf["b"] + with self.assertRaisesRegex(TypeError, expected_error_message): + psdf["a"] - 1 + with self.assertRaisesRegex(TypeError, expected_error_message): + 1 - psdf["a"] + + def test_arithmetic_op_exceptions(self): + psser = self.ks_start_date + py_datetime = self.pd_start_date.dt.to_pydatetime() + datetime_index = ps.Index(self.pd_start_date) + + for other in [1, 0.1, psser, datetime_index, py_datetime]: + expected_err_msg = "Addition can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser + other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other + psser) + + expected_err_msg = "Multiplication can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser * other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other * psser) + + expected_err_msg = "True division can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser / other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other / psser) + + expected_err_msg = "Floor division can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser // other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other // psser) + + expected_err_msg = "Modulo can not be applied to datetimes." + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser % other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other % psser) + + expected_err_msg = "Datetime subtraction can only be applied to datetime series." + + for other in [1, 0.1]: + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser - other) + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: other - psser) + + self.assertRaisesRegex(TypeError, expected_err_msg, lambda: psser - other) + self.assertRaises(NotImplementedError, lambda: py_datetime - psser) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43462): Enable SeriesDateTimeTests.test_date_subtraction for pandas 2.0.0.", + ) + def test_date_subtraction(self): + pdf = self.pdf1 + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf["end_date"].dt.date - psdf["start_date"].dt.date, + (pdf["end_date"].dt.date - pdf["start_date"].dt.date).dt.days, + ) + + self.assert_eq( + psdf["end_date"].dt.date - datetime.date(2012, 1, 1), + (pdf["end_date"].dt.date - datetime.date(2012, 1, 1)).dt.days, + ) + + self.assert_eq( + datetime.date(2013, 3, 11) - psdf["start_date"].dt.date, + (datetime.date(2013, 3, 11) - pdf["start_date"].dt.date).dt.days, + ) + + psdf = ps.DataFrame( + {"a": pd.date_range("2016-12-31", "2017-01-08", freq="D"), "b": pd.Series(range(9))} + ) + expected_error_message = "Date subtraction can only be applied to date series." + with self.assertRaisesRegex(TypeError, expected_error_message): + psdf["a"].dt.date - psdf["b"] + with self.assertRaisesRegex(TypeError, expected_error_message): + psdf["a"].dt.date - 1 + with self.assertRaisesRegex(TypeError, expected_error_message): + 1 - psdf["a"].dt.date + + @unittest.skip( + "It fails in certain OSs presumably due to different " + "timezone behaviours inherited from C library." + ) + def test_div(self): + pdf = self.pdf1 + psdf = ps.from_pandas(pdf) + for u in "D", "s", "ms": + duration = np.timedelta64(1, u) + self.assert_eq( + (psdf["end_date"] - psdf["start_date"]) / duration, + (pdf["end_date"] - pdf["start_date"]) / duration, + ) + + @unittest.skip("It is currently failed probably for the same reason in 'test_subtraction'") + def test_date(self): + self.check_func(lambda x: x.dt.date) + + def test_time(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.dt.time) + + def test_timetz(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.dt.timetz) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43736): Enable SeriesDateTimeTests.test_year for pandas 2.0.0.", + ) + def test_year(self): + self.check_func(lambda x: x.dt.year) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43731): Enable SeriesDateTimeTests.test_month for pandas 2.0.0.", + ) + def test_month(self): + self.check_func(lambda x: x.dt.month) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43722): Enable SeriesDateTimeTests.test_day for pandas 2.0.0.", + ) + def test_day(self): + self.check_func(lambda x: x.dt.day) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43728): Enable SeriesDateTimeTests.test_hour for pandas 2.0.0.", + ) + def test_hour(self): + self.check_func(lambda x: x.dt.hour) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43730): Enable SeriesDateTimeTests.test_minute for pandas 2.0.0.", + ) + def test_minute(self): + self.check_func(lambda x: x.dt.minute) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43733): Enable SeriesDateTimeTests.test_second for pandas 2.0.0.", + ) + def test_second(self): + self.check_func(lambda x: x.dt.second) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43729): Enable SeriesDateTimeTests.test_microsecond for pandas 2.0.0.", + ) + def test_microsecond(self): + self.check_func(lambda x: x.dt.microsecond) + + def test_nanosecond(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.dt.nanosecond) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-42617): Support `isocalendar`", + ) + def test_week(self): + self.check_func(lambda x: x.dt.week) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-42617): Support `isocalendar`", + ) + def test_weekofyear(self): + self.check_func(lambda x: x.dt.weekofyear) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43723): Enable SeriesDateTimeTests.test_dayofweek for pandas 2.0.0.", + ) + def test_dayofweek(self): + self.check_func(lambda x: x.dt.dayofweek) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43735): Enable SeriesDateTimeTests.test_weekday for pandas 2.0.0.", + ) + def test_weekday(self): + self.check_func(lambda x: x.dt.weekday) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43724): Enable SeriesDateTimeTests.test_dayofyear for pandas 2.0.0.", + ) + def test_dayofyear(self): + self.check_func(lambda x: x.dt.dayofyear) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43732): Enable SeriesDateTimeTests.test_quarter for pandas 2.0.0.", + ) + def test_quarter(self): + self.check_func(lambda x: x.dt.quarter) + + def test_is_month_start(self): + self.check_func(lambda x: x.dt.is_month_start) + + def test_is_month_end(self): + self.check_func(lambda x: x.dt.is_month_end) + + def test_is_quarter_start(self): + self.check_func(lambda x: x.dt.is_quarter_start) + + def test_is_quarter_end(self): + self.check_func(lambda x: x.dt.is_quarter_end) + + def test_is_year_start(self): + self.check_func(lambda x: x.dt.is_year_start) + + def test_is_year_end(self): + self.check_func(lambda x: x.dt.is_year_end) + + def test_is_leap_year(self): + self.check_func(lambda x: x.dt.is_leap_year) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43726): Enable SeriesDateTimeTests.test_daysinmonth for pandas 2.0.0.", + ) + def test_daysinmonth(self): + self.check_func(lambda x: x.dt.daysinmonth) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43725): Enable SeriesDateTimeTests.test_days_in_month for pandas 2.0.0.", + ) + def test_days_in_month(self): + self.check_func(lambda x: x.dt.days_in_month) + + @unittest.expectedFailure + def test_tz_localize(self): + self.check_func(lambda x: x.dt.tz_localize("America/New_York")) + + @unittest.expectedFailure + def test_tz_convert(self): + self.check_func(lambda x: x.dt.tz_convert("America/New_York")) + + def test_normalize(self): + self.check_func(lambda x: x.dt.normalize()) + + def test_strftime(self): + self.check_func(lambda x: x.dt.strftime("%Y-%m-%d")) + + def test_round(self): + self.check_func(lambda x: x.dt.round(freq="min")) + self.check_func(lambda x: x.dt.round(freq="H")) + + def test_floor(self): + self.check_func(lambda x: x.dt.floor(freq="min")) + self.check_func(lambda x: x.dt.floor(freq="H")) + + def test_ceil(self): + self.check_func(lambda x: x.dt.ceil(freq="min")) + self.check_func(lambda x: x.dt.ceil(freq="H")) + + @unittest.skip("Unsupported locale setting") + def test_month_name(self): + self.check_func(lambda x: x.dt.month_name()) + self.check_func(lambda x: x.dt.month_name(locale="en_US.UTF-8")) + + @unittest.skip("Unsupported locale setting") + def test_day_name(self): + self.check_func(lambda x: x.dt.day_name()) + self.check_func(lambda x: x.dt.day_name(locale="en_US.UTF-8")) + + def test_unsupported_type(self): + self.assertRaisesRegex( + ValueError, "Cannot call DatetimeMethods on type LongType", lambda: ps.Series([0]).dt + ) + + +class SeriesDateTimeTests(SeriesDateTimeTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_series_datetime import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_string.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_string.py new file mode 100644 index 0000000..3c2bd58 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_series_string.py @@ -0,0 +1,362 @@ +# +# 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 pandas as pd +import numpy as np +import re +import unittest +from distutils.version import LooseVersion + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils + + +class SeriesStringTestsMixin: + @property + def pser(self): + return pd.Series( + [ + "apples", + "Bananas", + "carrots", + "1", + "100", + "", + "\nleading-whitespace", + "trailing-Whitespace \t", + None, + np.NaN, + ] + ) + + def check_func(self, func, almost=False): + self.check_func_on_series(func, self.pser, almost=almost) + + def check_func_on_series(self, func, pser, almost=False): + self.assert_eq(func(ps.from_pandas(pser)), func(pser), almost=almost) + + def test_string_add_str_num(self): + pdf = pd.DataFrame(dict(col1=["a"], col2=[1])) + psdf = ps.from_pandas(pdf) + with self.assertRaises(TypeError): + psdf["col1"] + psdf["col2"] + + def test_string_add_assign(self): + pdf = pd.DataFrame(dict(col1=["a", "b", "c"], col2=["1", "2", "3"])) + psdf = ps.from_pandas(pdf) + psdf["col1"] += psdf["col2"] + pdf["col1"] += pdf["col2"] + self.assert_eq(psdf["col1"], pdf["col1"]) + + def test_string_add_str_str(self): + pdf = pd.DataFrame(dict(col1=["a", "b", "c"], col2=["1", "2", "3"])) + psdf = ps.from_pandas(pdf) + + # TODO: Fix the Series names + self.assert_eq(psdf["col1"] + psdf["col2"], pdf["col1"] + pdf["col2"]) + self.assert_eq(psdf["col2"] + psdf["col1"], pdf["col2"] + pdf["col1"]) + + def test_string_add_str_lit(self): + pdf = pd.DataFrame(dict(col1=["a", "b", "c"])) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf["col1"] + "_lit", pdf["col1"] + "_lit") + self.assert_eq("_lit" + psdf["col1"], "_lit" + pdf["col1"]) + + def test_string_capitalize(self): + self.check_func(lambda x: x.str.capitalize()) + + def test_string_title(self): + self.check_func(lambda x: x.str.title()) + + def test_string_lower(self): + self.check_func(lambda x: x.str.lower()) + + def test_string_upper(self): + self.check_func(lambda x: x.str.upper()) + + def test_string_swapcase(self): + self.check_func(lambda x: x.str.swapcase()) + + def test_string_startswith(self): + pattern = "car" + self.check_func(lambda x: x.str.startswith(pattern)) + self.check_func(lambda x: x.str.startswith(pattern, na=False)) + + def test_string_endswith(self): + pattern = "s" + self.check_func(lambda x: x.str.endswith(pattern)) + self.check_func(lambda x: x.str.endswith(pattern, na=False)) + + def test_string_strip(self): + self.check_func(lambda x: x.str.strip()) + self.check_func(lambda x: x.str.strip("es\t")) + self.check_func(lambda x: x.str.strip("1")) + + def test_string_lstrip(self): + self.check_func(lambda x: x.str.lstrip()) + self.check_func(lambda x: x.str.lstrip("\n1le")) + self.check_func(lambda x: x.str.lstrip("s")) + + def test_string_rstrip(self): + self.check_func(lambda x: x.str.rstrip()) + self.check_func(lambda x: x.str.rstrip("\t ec")) + self.check_func(lambda x: x.str.rstrip("0")) + + def test_string_get(self): + self.check_func(lambda x: x.str.get(6)) + self.check_func(lambda x: x.str.get(-1)) + + def test_string_isalnum(self): + self.check_func(lambda x: x.str.isalnum()) + + def test_string_isalpha(self): + self.check_func(lambda x: x.str.isalpha()) + + def test_string_isdigit(self): + self.check_func(lambda x: x.str.isdigit()) + + def test_string_isspace(self): + self.check_func(lambda x: x.str.isspace()) + + def test_string_islower(self): + self.check_func(lambda x: x.str.islower()) + + def test_string_isupper(self): + self.check_func(lambda x: x.str.isupper()) + + def test_string_istitle(self): + self.check_func(lambda x: x.str.istitle()) + + def test_string_isnumeric(self): + self.check_func(lambda x: x.str.isnumeric()) + + def test_string_isdecimal(self): + self.check_func(lambda x: x.str.isdecimal()) + + def test_string_cat(self): + psser = ps.from_pandas(self.pser) + with self.assertRaises(NotImplementedError): + psser.str.cat() + + def test_string_center(self): + self.check_func(lambda x: x.str.center(0)) + self.check_func(lambda x: x.str.center(10)) + self.check_func(lambda x: x.str.center(10, "x")) + + def test_string_contains(self): + self.check_func(lambda x: x.str.contains("le", regex=False)) + self.check_func(lambda x: x.str.contains("White", case=True, regex=False)) + self.check_func(lambda x: x.str.contains("apples|carrots", regex=True)) + self.check_func(lambda x: x.str.contains("BANANAS", flags=re.IGNORECASE, na=False)) + + def test_string_count(self): + self.check_func(lambda x: x.str.count("wh|Wh")) + self.check_func(lambda x: x.str.count("WH", flags=re.IGNORECASE)) + + def test_string_decode(self): + psser = ps.from_pandas(self.pser) + with self.assertRaises(NotImplementedError): + psser.str.decode("utf-8") + + def test_string_encode(self): + psser = ps.from_pandas(self.pser) + with self.assertRaises(NotImplementedError): + psser.str.encode("utf-8") + + def test_string_extract(self): + psser = ps.from_pandas(self.pser) + with self.assertRaises(NotImplementedError): + psser.str.extract("pat") + + def test_string_extractall(self): + psser = ps.from_pandas(self.pser) + with self.assertRaises(NotImplementedError): + psser.str.extractall("pat") + + def test_string_find(self): + self.check_func(lambda x: x.str.find("a")) + self.check_func(lambda x: x.str.find("a", start=3)) + self.check_func(lambda x: x.str.find("a", start=0, end=1)) + + def test_string_findall(self): + self.check_func_on_series(lambda x: x.str.findall("es|as").apply(str), self.pser[:-1]) + self.check_func_on_series( + lambda x: x.str.findall("wh.*", flags=re.IGNORECASE).apply(str), self.pser[:-1] + ) + + def test_string_index(self): + pser = pd.Series(["tea", "eat"]) + self.check_func_on_series(lambda x: x.str.index("ea"), pser) + with self.assertRaises(Exception): + self.check_func_on_series(lambda x: x.str.index("ea", start=0, end=2), pser) + with self.assertRaises(Exception): + self.check_func(lambda x: x.str.index("not-found")) + + def test_string_join(self): + pser = pd.Series([["a", "b", "c"], ["xx", "yy", "zz"]]) + self.check_func_on_series(lambda x: x.str.join("-"), pser) + self.check_func(lambda x: x.str.join("-")) + + def test_string_len(self): + self.check_func(lambda x: x.str.len()) + pser = pd.Series([["a", "b", "c"], ["xx"], []]) + self.check_func_on_series(lambda x: x.str.len(), pser) + + def test_string_ljust(self): + self.check_func(lambda x: x.str.ljust(0)) + self.check_func(lambda x: x.str.ljust(10)) + self.check_func(lambda x: x.str.ljust(30, "x")) + + def test_string_match(self): + self.check_func(lambda x: x.str.match("in")) + self.check_func(lambda x: x.str.match("apples|carrots", na=False)) + self.check_func(lambda x: x.str.match("White", case=True)) + self.check_func(lambda x: x.str.match("BANANAS", flags=re.IGNORECASE)) + + def test_string_normalize(self): + self.check_func(lambda x: x.str.normalize("NFC")) + self.check_func(lambda x: x.str.normalize("NFKD")) + + def test_string_pad(self): + self.check_func(lambda x: x.str.pad(10)) + self.check_func(lambda x: x.str.pad(10, side="both")) + self.check_func(lambda x: x.str.pad(10, side="right", fillchar="-")) + + def test_string_partition(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.str.partition()) + + def test_string_repeat(self): + self.check_func(lambda x: x.str.repeat(repeats=3)) + with self.assertRaises(TypeError): + self.check_func(lambda x: x.str.repeat(repeats=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43476): Enable SeriesStringTests.test_string_replace for pandas 2.0.0.", + ) + def test_string_replace(self): + self.check_func(lambda x: x.str.replace("a.", "xx", regex=True)) + self.check_func(lambda x: x.str.replace("a.", "xx", regex=False)) + self.check_func(lambda x: x.str.replace("ing", "0", flags=re.IGNORECASE)) + + # reverse every lowercase word + def repl(m): + return m.group(0)[::-1] + + self.check_func(lambda x: x.str.replace(r"[a-z]+", repl)) + # compiled regex with flags + regex_pat = re.compile(r"WHITESPACE", flags=re.IGNORECASE) + self.check_func(lambda x: x.str.replace(regex_pat, "---")) + + def test_string_rfind(self): + self.check_func(lambda x: x.str.rfind("a")) + self.check_func(lambda x: x.str.rfind("a", start=3)) + self.check_func(lambda x: x.str.rfind("a", start=0, end=1)) + + def test_string_rindex(self): + pser = pd.Series(["teatea", "eateat"]) + self.check_func_on_series(lambda x: x.str.rindex("ea"), pser) + with self.assertRaises(Exception): + self.check_func_on_series(lambda x: x.str.rindex("ea", start=0, end=2), pser) + with self.assertRaises(Exception): + self.check_func(lambda x: x.str.rindex("not-found")) + + def test_string_rjust(self): + self.check_func(lambda x: x.str.rjust(0)) + self.check_func(lambda x: x.str.rjust(10)) + self.check_func(lambda x: x.str.rjust(30, "x")) + + def test_string_rpartition(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.str.rpartition()) + + def test_string_slice(self): + self.check_func(lambda x: x.str.slice(start=1)) + self.check_func(lambda x: x.str.slice(stop=3)) + self.check_func(lambda x: x.str.slice(step=2)) + self.check_func(lambda x: x.str.slice(start=0, stop=5, step=3)) + + def test_string_slice_replace(self): + self.check_func(lambda x: x.str.slice_replace(1, repl="X")) + self.check_func(lambda x: x.str.slice_replace(stop=2, repl="X")) + self.check_func(lambda x: x.str.slice_replace(start=1, stop=3, repl="X")) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43478): Enable SeriesStringTests.test_string_split for pandas 2.0.0.", + ) + def test_string_split(self): + self.check_func_on_series(lambda x: repr(x.str.split()), self.pser[:-1]) + self.check_func_on_series(lambda x: repr(x.str.split(r"p*")), self.pser[:-1]) + pser = pd.Series(["This is a sentence.", "This-is-a-long-word."]) + self.check_func_on_series(lambda x: repr(x.str.split(n=2)), pser) + self.check_func_on_series(lambda x: repr(x.str.split(pat="-", n=2)), pser) + self.check_func_on_series(lambda x: x.str.split(n=2, expand=True), pser, almost=True) + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.str.split(expand=True)) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43477): Enable SeriesStringTests.test_string_rsplit for pandas 2.0.0.", + ) + def test_string_rsplit(self): + self.check_func_on_series(lambda x: repr(x.str.rsplit()), self.pser[:-1]) + self.check_func_on_series(lambda x: repr(x.str.rsplit(r"p*")), self.pser[:-1]) + pser = pd.Series(["This is a sentence.", "This-is-a-long-word."]) + self.check_func_on_series(lambda x: repr(x.str.rsplit(n=2)), pser) + self.check_func_on_series(lambda x: repr(x.str.rsplit(pat="-", n=2)), pser) + self.check_func_on_series(lambda x: x.str.rsplit(n=2, expand=True), pser, almost=True) + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.str.rsplit(expand=True)) + + def test_string_translate(self): + m = str.maketrans({"a": "X", "e": "Y", "i": None}) + self.check_func(lambda x: x.str.translate(m)) + + def test_string_wrap(self): + self.check_func(lambda x: x.str.wrap(5)) + self.check_func(lambda x: x.str.wrap(5, expand_tabs=False)) + self.check_func(lambda x: x.str.wrap(5, replace_whitespace=False)) + self.check_func(lambda x: x.str.wrap(5, drop_whitespace=False)) + self.check_func(lambda x: x.str.wrap(5, break_long_words=False)) + self.check_func(lambda x: x.str.wrap(5, break_on_hyphens=False)) + + def test_string_zfill(self): + self.check_func(lambda x: x.str.zfill(10)) + + def test_string_get_dummies(self): + with self.assertRaises(NotImplementedError): + self.check_func(lambda x: x.str.get_dummies()) + + +class SeriesStringTests(SeriesStringTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_series_string import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_spark_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_spark_functions.py new file mode 100644 index 0000000..3e2281c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_spark_functions.py @@ -0,0 +1,46 @@ +# +# 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 numpy as np + +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.utils import spark_column_equals +from pyspark.sql import functions as F +from pyspark.testing.pandasutils import PandasOnSparkTestCase + + +class SparkFunctionsTestsMixin: + def test_repeat(self): + # TODO: Placeholder + pass + + +class SparkFunctionsTests(SparkFunctionsTestsMixin, PandasOnSparkTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_spark_functions import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_sql.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_sql.py new file mode 100644 index 0000000..7800d57 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_sql.py @@ -0,0 +1,123 @@ +# +# 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 import pandas as ps +from pyspark.errors import ParseException +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.testing.utils import assertDataFrameEqual + + +class SQLTestsMixin: + def test_error_variable_not_exist(self): + with self.assertRaisesRegex(KeyError, "variable_foo"): + ps.sql("select * from {variable_foo}") + + def test_error_bad_sql(self): + with self.assertRaises(ParseException): + ps.sql("this is not valid sql") + + def test_series_not_referred(self): + psdf = ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) + with self.assertRaisesRegex(ValueError, "The series in {ser}"): + ps.sql("SELECT {ser} FROM range(10)", ser=psdf.A) + + def test_sql_with_index_col(self): + import pandas as pd + + # Index + psdf = ps.DataFrame( + {"A": [1, 2, 3], "B": [4, 5, 6]}, index=pd.Index(["a", "b", "c"], name="index") + ) + psdf_reset_index = psdf.reset_index() + actual = ps.sql( + "select * from {psdf_reset_index} where A > 1", + index_col="index", + psdf_reset_index=psdf_reset_index, + ) + expected = psdf.iloc[[1, 2]] + assertDataFrameEqual(actual, expected) + + # MultiIndex + psdf = ps.DataFrame( + {"A": [1, 2, 3], "B": [4, 5, 6]}, + index=pd.MultiIndex.from_tuples( + [("a", "b"), ("c", "d"), ("e", "f")], names=["index1", "index2"] + ), + ) + psdf_reset_index = psdf.reset_index() + actual = ps.sql( + "select * from {psdf_reset_index} where A > 1", + index_col=["index1", "index2"], + psdf_reset_index=psdf_reset_index, + ) + expected = psdf.iloc[[1, 2]] + assertDataFrameEqual(actual, expected) + + def test_sql_with_pandas_objects(self): + import pandas as pd + + pdf = pd.DataFrame({"a": [1, 2, 3, 4]}) + assertDataFrameEqual( + ps.sql("SELECT {col} + 1 as a FROM {tbl}", col=pdf.a, tbl=pdf), pdf + 1 + ) + + def test_sql_with_python_objects(self): + assertDataFrameEqual( + ps.sql("SELECT {col} as a FROM range(1)", col="lit"), ps.DataFrame({"a": ["lit"]}) + ) + assertDataFrameEqual( + ps.sql("SELECT id FROM range(10) WHERE id IN {pred}", col="lit", pred=(1, 2, 3)), + ps.DataFrame({"id": [1, 2, 3]}), + ) + assertDataFrameEqual( + ps.sql("SELECT {col} as a FROM range(1)", col="a'''c''d"), + ps.DataFrame({"a": ["a'''c''d"]}), + ) + assertDataFrameEqual( + ps.sql("SELECT id FROM range(10) WHERE id IN {pred}", col="a'''c''d", pred=(1, 2, 3)), + ps.DataFrame({"id": [1, 2, 3]}), + ) + + def test_sql_with_pandas_on_spark_objects(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4]}) + + assertDataFrameEqual(ps.sql("SELECT {col} FROM {tbl}", col=psdf.a, tbl=psdf), psdf) + assertDataFrameEqual(ps.sql("SELECT {tbl.a} FROM {tbl}", tbl=psdf), psdf) + + psdf = ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) + assertDataFrameEqual( + ps.sql("SELECT {col}, {col2} FROM {tbl}", col=psdf.A, col2=psdf.B, tbl=psdf), psdf + ) + assertDataFrameEqual(ps.sql("SELECT {tbl.A}, {tbl.B} FROM {tbl}", tbl=psdf), psdf) + + +class SQLTests(SQLTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_sql import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_stats.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_stats.py new file mode 100644 index 0000000..ec56fa7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_stats.py @@ -0,0 +1,581 @@ +# +# 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 unittest +from distutils.version import LooseVersion +import numpy as np +import pandas as pd + +try: + from pandas._testing import makeMissingDataframe +except ImportError: + from pandas.util.testing import makeMissingDataframe + +from pyspark import pandas as ps +from pyspark.pandas.config import option_context +from pyspark.testing.pandasutils import PandasOnSparkTestCase, SPARK_CONF_ARROW_ENABLED +from pyspark.testing.sqlutils import SQLTestUtils + + +class StatsTestsMixin: + def _test_stat_functions(self, pdf_or_pser, psdf_or_psser): + functions = ["max", "min", "mean", "sum", "count"] + for funcname in functions: + self.assert_eq(getattr(psdf_or_psser, funcname)(), getattr(pdf_or_pser, funcname)()) + + functions = ["std", "var", "product", "sem"] + for funcname in functions: + self.assert_eq( + getattr(psdf_or_psser, funcname)(), + getattr(pdf_or_pser, funcname)(), + check_exact=False, + ) + + functions = ["std", "var", "sem"] + for funcname in functions: + self.assert_eq( + getattr(psdf_or_psser, funcname)(ddof=0), + getattr(pdf_or_pser, funcname)(ddof=0), + check_exact=False, + ) + + # NOTE: To test skew, kurt, and median, just make sure they run. + # The numbers are different in spark and pandas. + functions = ["skew", "kurt", "median"] + for funcname in functions: + getattr(psdf_or_psser, funcname)() + + def test_stat_functions(self): + pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [1, 2, 3, 4], "C": [1, np.nan, 3, np.nan]}) + psdf = ps.from_pandas(pdf) + self._test_stat_functions(pdf.A, psdf.A) + self._test_stat_functions(pdf, psdf) + + # empty + self._test_stat_functions(pdf.A.loc[[]], psdf.A.loc[[]]) + self._test_stat_functions(pdf.loc[[]], psdf.loc[[]]) + + def test_stat_functions_multiindex_column(self): + arrays = [np.array(["A", "A", "B", "B"]), np.array(["one", "two", "one", "two"])] + pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays) + psdf = ps.from_pandas(pdf) + self._test_stat_functions(pdf.A, psdf.A) + self._test_stat_functions(pdf, psdf) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43499): Enable SeriesTests.test_stat_functions_with_no_numeric_columns " + "for pandas 2.0.0.", + ) + def test_stat_functions_with_no_numeric_columns(self): + pdf = pd.DataFrame( + { + "A": ["a", None, "c", "d", None, "f", "g"], + "B": ["A", "B", "C", None, "E", "F", None], + } + ) + psdf = ps.from_pandas(pdf) + + self._test_stat_functions(pdf, psdf) + + def test_sum(self): + pdf = pd.DataFrame({"a": [1, 2, 3, np.nan], "b": [0.1, np.nan, 0.3, np.nan]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.sum(), pdf.sum()) + self.assert_eq(psdf.sum(axis=1), pdf.sum(axis=1)) + self.assert_eq(psdf.sum(min_count=3), pdf.sum(min_count=3)) + self.assert_eq(psdf.sum(axis=1, min_count=1), pdf.sum(axis=1, min_count=1)) + self.assert_eq(psdf.loc[[]].sum(), pdf.loc[[]].sum()) + self.assert_eq(psdf.loc[[]].sum(min_count=1), pdf.loc[[]].sum(min_count=1)) + + self.assert_eq(psdf["a"].sum(), pdf["a"].sum()) + self.assert_eq(psdf["a"].sum(min_count=3), pdf["a"].sum(min_count=3)) + self.assert_eq(psdf["b"].sum(min_count=3), pdf["b"].sum(min_count=3)) + self.assert_eq(psdf["a"].loc[[]].sum(), pdf["a"].loc[[]].sum()) + self.assert_eq(psdf["a"].loc[[]].sum(min_count=1), pdf["a"].loc[[]].sum(min_count=1)) + + def test_product(self): + pdf = pd.DataFrame( + {"a": [1, -2, -3, np.nan], "b": [0.1, np.nan, -0.3, np.nan], "c": [10, 20, 0, -10]} + ) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.product(), pdf.product(), check_exact=False) + self.assert_eq(psdf.product(axis=1), pdf.product(axis=1)) + self.assert_eq(psdf.product(min_count=3), pdf.product(min_count=3), check_exact=False) + self.assert_eq(psdf.product(axis=1, min_count=1), pdf.product(axis=1, min_count=1)) + self.assert_eq(psdf.loc[[]].product(), pdf.loc[[]].product()) + self.assert_eq(psdf.loc[[]].product(min_count=1), pdf.loc[[]].product(min_count=1)) + + self.assert_eq(psdf["a"].product(), pdf["a"].product(), check_exact=False) + self.assert_eq( + psdf["a"].product(min_count=3), pdf["a"].product(min_count=3), check_exact=False + ) + self.assert_eq(psdf["b"].product(min_count=3), pdf["b"].product(min_count=3)) + self.assert_eq(psdf["c"].product(min_count=3), pdf["c"].product(min_count=3)) + self.assert_eq(psdf["a"].loc[[]].product(), pdf["a"].loc[[]].product()) + self.assert_eq( + psdf["a"].loc[[]].product(min_count=1), pdf["a"].loc[[]].product(min_count=1) + ) + + def test_abs(self): + pdf = pd.DataFrame( + { + "A": [1, -2, np.nan, -4, 5], + "B": [1.0, -2, np.nan, -4, 5], + "C": [-6.0, -7, -8, np.nan, 10], + "D": ["a", "b", "c", "d", np.nan], + "E": [True, np.nan, False, True, True], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.A.abs(), pdf.A.abs()) + self.assert_eq(psdf.B.abs(), pdf.B.abs()) + self.assert_eq(psdf.E.abs(), pdf.E.abs()) + # pandas' bug? + # self.assert_eq(psdf[["B", "C", "E"]].abs(), pdf[["B", "C", "E"]].abs()) + self.assert_eq(psdf[["B", "C"]].abs(), pdf[["B", "C"]].abs()) + self.assert_eq(psdf[["E"]].abs(), pdf[["E"]].abs()) + + with self.assertRaisesRegex( + TypeError, "bad operand type for abs\\(\\): object \\(string\\)" + ): + psdf.abs() + with self.assertRaisesRegex( + TypeError, "bad operand type for abs\\(\\): object \\(string\\)" + ): + psdf.D.abs() + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43498): Enable SeriesTests.test_axis_on_dataframe for pandas 2.0.0.", + ) + def test_axis_on_dataframe(self): + # The number of each count is intentionally big + # because when data is small, it executes a shortcut. + # Less than 'compute.shortcut_limit' will execute a shortcut + # by using collected pandas dataframe directly. + # now we set the 'compute.shortcut_limit' as 1000 explicitly + with option_context("compute.shortcut_limit", 1000): + pdf = pd.DataFrame( + { + "A": [1, -2, 3, -4, 5] * 300, + "B": [1.0, -2, 3, -4, 5] * 300, + "C": [-6.0, -7, -8, -9, 10] * 300, + "D": [True, False, True, False, False] * 300, + }, + index=range(10, 15001, 10), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.count(axis=1), pdf.count(axis=1)) + self.assert_eq(psdf.var(axis=1), pdf.var(axis=1)) + self.assert_eq(psdf.var(axis=1, ddof=0), pdf.var(axis=1, ddof=0)) + self.assert_eq(psdf.std(axis=1), pdf.std(axis=1)) + self.assert_eq(psdf.std(axis=1, ddof=0), pdf.std(axis=1, ddof=0)) + self.assert_eq(psdf.max(axis=1), pdf.max(axis=1)) + self.assert_eq(psdf.min(axis=1), pdf.min(axis=1)) + self.assert_eq(psdf.sum(axis=1), pdf.sum(axis=1)) + self.assert_eq(psdf.product(axis=1), pdf.product(axis=1)) + self.assert_eq(psdf.kurtosis(axis=0), pdf.kurtosis(axis=0), almost=True) + self.assert_eq(psdf.kurtosis(axis=1), pdf.kurtosis(axis=1)) + self.assert_eq(psdf.skew(axis=0), pdf.skew(axis=0), almost=True) + self.assert_eq(psdf.skew(axis=1), pdf.skew(axis=1)) + self.assert_eq(psdf.mean(axis=1), pdf.mean(axis=1)) + self.assert_eq(psdf.sem(axis=1), pdf.sem(axis=1)) + self.assert_eq(psdf.sem(axis=1, ddof=0), pdf.sem(axis=1, ddof=0)) + + self.assert_eq( + psdf.count(axis=1, numeric_only=True), pdf.count(axis=1, numeric_only=True) + ) + self.assert_eq(psdf.var(axis=1, numeric_only=True), pdf.var(axis=1, numeric_only=True)) + self.assert_eq( + psdf.var(axis=1, ddof=0, numeric_only=True), + pdf.var(axis=1, ddof=0, numeric_only=True), + ) + self.assert_eq(psdf.std(axis=1, numeric_only=True), pdf.std(axis=1, numeric_only=True)) + self.assert_eq( + psdf.std(axis=1, ddof=0, numeric_only=True), + pdf.std(axis=1, ddof=0, numeric_only=True), + ) + self.assert_eq( + psdf.max(axis=1, numeric_only=True), + pdf.max(axis=1, numeric_only=True).astype(float), + ) + self.assert_eq( + psdf.min(axis=1, numeric_only=True), + pdf.min(axis=1, numeric_only=True).astype(float), + ) + self.assert_eq( + psdf.sum(axis=1, numeric_only=True), + pdf.sum(axis=1, numeric_only=True).astype(float), + ) + self.assert_eq( + psdf.product(axis=1, numeric_only=True), + pdf.product(axis=1, numeric_only=True).astype(float), + ) + self.assert_eq( + psdf.kurtosis(axis=0, numeric_only=True), + pdf.kurtosis(axis=0, numeric_only=True), + almost=True, + ) + self.assert_eq( + psdf.kurtosis(axis=1, numeric_only=True), pdf.kurtosis(axis=1, numeric_only=True) + ) + self.assert_eq( + psdf.skew(axis=0, numeric_only=True), + pdf.skew(axis=0, numeric_only=True), + almost=True, + ) + self.assert_eq( + psdf.skew(axis=1, numeric_only=True), pdf.skew(axis=1, numeric_only=True) + ) + self.assert_eq( + psdf.mean(axis=1, numeric_only=True), pdf.mean(axis=1, numeric_only=True) + ) + self.assert_eq(psdf.sem(axis=1, numeric_only=True), pdf.sem(axis=1, numeric_only=True)) + self.assert_eq( + psdf.sem(axis=1, ddof=0, numeric_only=True), + pdf.sem(axis=1, ddof=0, numeric_only=True), + ) + + def test_skew_kurt_numerical_stability(self): + pdf = pd.DataFrame( + { + "A": [1, 1, 1, 1, 1], + "B": [1.0, np.nan, 4, 2, 5], + "C": [-6.0, -7, np.nan, np.nan, 10], + "D": [1.2, np.nan, np.nan, 9.8, np.nan], + "E": [1, np.nan, np.nan, np.nan, np.nan], + "F": [np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.skew(), pdf.skew(), almost=True) + self.assert_eq(psdf.kurt(), pdf.kurt(), almost=True) + + def test_dataframe_corr(self): + pdf = makeMissingDataframe(0.3, 42) + psdf = ps.from_pandas(pdf) + + with self.assertRaisesRegex(ValueError, "Invalid method"): + psdf.corr("std") + with self.assertRaisesRegex(TypeError, "Invalid min_periods type"): + psdf.corr(min_periods="3") + + for method in ["pearson", "spearman", "kendall"]: + self.assert_eq(psdf.corr(method=method), pdf.corr(method=method), check_exact=False) + self.assert_eq( + psdf.corr(method=method, min_periods=1), + pdf.corr(method=method, min_periods=1), + check_exact=False, + ) + self.assert_eq( + psdf.corr(method=method, min_periods=3), + pdf.corr(method=method, min_periods=3), + check_exact=False, + ) + self.assert_eq( + (psdf + 1).corr(method=method, min_periods=2), + (pdf + 1).corr(method=method, min_periods=2), + check_exact=False, + ) + + # multi-index columns + columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")]) + pdf.columns = columns + psdf.columns = columns + + for method in ["pearson", "spearman", "kendall"]: + self.assert_eq(psdf.corr(method=method), pdf.corr(method=method), check_exact=False) + self.assert_eq( + psdf.corr(method=method, min_periods=1), + pdf.corr(method=method, min_periods=1), + check_exact=False, + ) + self.assert_eq( + psdf.corr(method=method, min_periods=3), + pdf.corr(method=method, min_periods=3), + check_exact=False, + ) + self.assert_eq( + (psdf + 1).corr(method=method, min_periods=2), + (pdf + 1).corr(method=method, min_periods=2), + check_exact=False, + ) + + # test with identical values + pdf = pd.DataFrame( + { + "a": [0, 1, 1, 1, 0], + "b": [2, 2, -1, 1, np.nan], + "c": [3, 3, 3, 3, 3], + "d": [np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + + for method in ["pearson", "spearman", "kendall"]: + self.assert_eq(psdf.corr(method=method), pdf.corr(method=method), check_exact=False) + self.assert_eq( + psdf.corr(method=method, min_periods=1), + pdf.corr(method=method, min_periods=1), + check_exact=False, + ) + self.assert_eq( + psdf.corr(method=method, min_periods=3), + pdf.corr(method=method, min_periods=3), + check_exact=False, + ) + + def test_series_corr(self): + pdf = makeMissingDataframe(0.3, 42) + pser1 = pdf.A + pser2 = pdf.B + psdf = ps.from_pandas(pdf) + psser1 = psdf.A + psser2 = psdf.B + + with self.assertRaisesRegex(ValueError, "Invalid method"): + psser1.corr(psser2, method="std") + with self.assertRaisesRegex(TypeError, "Invalid min_periods type"): + psser1.corr(psser2, min_periods="3") + + for method in ["pearson", "spearman", "kendall"]: + self.assert_eq( + psser1.corr(psser2, method=method), + pser1.corr(pser2, method=method), + almost=True, + ) + self.assert_eq( + psser1.corr(psser2, method=method, min_periods=1), + pser1.corr(pser2, method=method, min_periods=1), + almost=True, + ) + self.assert_eq( + psser1.corr(psser2, method=method, min_periods=3), + pser1.corr(pser2, method=method, min_periods=3), + almost=True, + ) + self.assert_eq( + (psser1 + 1).corr(psser2 - 2, method=method, min_periods=2), + (pser1 + 1).corr(pser2 - 2, method=method, min_periods=2), + almost=True, + ) + + # different anchors + psser1 = ps.from_pandas(pser1) + psser2 = ps.from_pandas(pser2) + + with self.assertRaisesRegex(ValueError, "Cannot combine the series or dataframe"): + psser1.corr(psser2) + + for method in ["pearson", "spearman", "kendall"]: + with ps.option_context("compute.ops_on_diff_frames", True): + self.assert_eq( + psser1.corr(psser2, method=method), + pser1.corr(pser2, method=method), + almost=True, + ) + self.assert_eq( + psser1.corr(psser2, method=method, min_periods=1), + pser1.corr(pser2, method=method, min_periods=1), + almost=True, + ) + self.assert_eq( + psser1.corr(psser2, method=method, min_periods=3), + pser1.corr(pser2, method=method, min_periods=3), + almost=True, + ) + self.assert_eq( + (psser1 + 1).corr(psser2 - 2, method=method, min_periods=2), + (pser1 + 1).corr(pser2 - 2, method=method, min_periods=2), + almost=True, + ) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43497): Enable SeriesTests.test_cov_corr_meta for pandas 2.0.0.", + ) + def test_cov_corr_meta(self): + # Disable arrow execution since corr() is using UDT internally which is not supported. + with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}): + pdf = pd.DataFrame( + { + "a": np.array([1, 2, 3], dtype="i1"), + "b": np.array([1, 2, 3], dtype="i2"), + "c": np.array([1, 2, 3], dtype="i4"), + "d": np.array([1, 2, 3]), + "e": np.array([1.0, 2.0, 3.0], dtype="f4"), + "f": np.array([1.0, 2.0, 3.0]), + "g": np.array([True, False, True]), + "h": np.array(list("abc")), + }, + index=pd.Index([1, 2, 3], name="myindex"), + ) + psdf = ps.from_pandas(pdf) + self.assert_eq(psdf.corr(), pdf.corr(), check_exact=False) + + def test_stats_on_boolean_dataframe(self): + pdf = pd.DataFrame({"A": [True, False, True], "B": [False, False, True]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.min(), pdf.min()) + self.assert_eq(psdf.max(), pdf.max()) + self.assert_eq(psdf.count(), pdf.count()) + + self.assert_eq(psdf.sum(), pdf.sum()) + self.assert_eq(psdf.product(), pdf.product()) + self.assert_eq(psdf.mean(), pdf.mean()) + + self.assert_eq(psdf.var(), pdf.var(), check_exact=False) + self.assert_eq(psdf.var(ddof=0), pdf.var(ddof=0), check_exact=False) + self.assert_eq(psdf.std(), pdf.std(), check_exact=False) + self.assert_eq(psdf.std(ddof=0), pdf.std(ddof=0), check_exact=False) + self.assert_eq(psdf.sem(), pdf.sem(), check_exact=False) + self.assert_eq(psdf.sem(ddof=0), pdf.sem(ddof=0), check_exact=False) + + def test_stats_on_boolean_series(self): + pser = pd.Series([True, False, True]) + psser = ps.from_pandas(pser) + + self.assert_eq(psser.min(), pser.min()) + self.assert_eq(psser.max(), pser.max()) + self.assert_eq(psser.count(), pser.count()) + + self.assert_eq(psser.sum(), pser.sum()) + self.assert_eq(psser.product(), pser.product()) + self.assert_eq(psser.mean(), pser.mean()) + + self.assert_eq(psser.var(), pser.var(), almost=True) + self.assert_eq(psser.var(ddof=0), pser.var(ddof=0), almost=True) + self.assert_eq(psser.var(ddof=2), pser.var(ddof=2), almost=True) + self.assert_eq(psser.std(), pser.std(), almost=True) + self.assert_eq(psser.std(ddof=0), pser.std(ddof=0), almost=True) + self.assert_eq(psser.std(ddof=2), pser.std(ddof=2), almost=True) + self.assert_eq(psser.sem(), pser.sem(), almost=True) + self.assert_eq(psser.sem(ddof=0), pser.sem(ddof=0), almost=True) + self.assert_eq(psser.sem(ddof=2), pser.sem(ddof=2), almost=True) + + def test_stats_on_non_numeric_columns_should_be_discarded_if_numeric_only_is_true(self): + pdf = pd.DataFrame({"i": [0, 1, 2], "b": [False, False, True], "s": ["x", "y", "z"]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq( + psdf[["i", "s"]].max(numeric_only=True), pdf[["i", "s"]].max(numeric_only=True) + ) + self.assert_eq( + psdf[["b", "s"]].max(numeric_only=True), pdf[["b", "s"]].max(numeric_only=True) + ) + self.assert_eq( + psdf[["i", "s"]].min(numeric_only=True), pdf[["i", "s"]].min(numeric_only=True) + ) + self.assert_eq( + psdf[["b", "s"]].min(numeric_only=True), pdf[["b", "s"]].min(numeric_only=True) + ) + self.assert_eq(psdf.count(numeric_only=True), pdf.count(numeric_only=True)) + + self.assert_eq(psdf.sum(numeric_only=True), pdf.sum(numeric_only=True)) + self.assert_eq(psdf.product(numeric_only=True), pdf.product(numeric_only=True)) + + self.assert_eq(psdf.mean(numeric_only=True), pdf.mean(numeric_only=True)) + + self.assert_eq(psdf.var(numeric_only=True), pdf.var(numeric_only=True), check_exact=False) + self.assert_eq( + psdf.var(ddof=0, numeric_only=True), + pdf.var(ddof=0, numeric_only=True), + check_exact=False, + ) + self.assert_eq( + psdf.var(ddof=2, numeric_only=True), + pdf.var(ddof=2, numeric_only=True), + check_exact=False, + ) + self.assert_eq(psdf.std(numeric_only=True), pdf.std(numeric_only=True), check_exact=False) + self.assert_eq( + psdf.std(ddof=0, numeric_only=True), + pdf.std(ddof=0, numeric_only=True), + check_exact=False, + ) + self.assert_eq( + psdf.std(ddof=2, numeric_only=True), + pdf.std(ddof=2, numeric_only=True), + check_exact=False, + ) + self.assert_eq(psdf.sem(numeric_only=True), pdf.sem(numeric_only=True), check_exact=False) + self.assert_eq( + psdf.sem(ddof=0, numeric_only=True), + pdf.sem(ddof=0, numeric_only=True), + check_exact=False, + ) + self.assert_eq( + psdf.sem(ddof=2, numeric_only=True), + pdf.sem(ddof=2, numeric_only=True), + check_exact=False, + ) + + self.assert_eq(len(psdf.median(numeric_only=True)), len(pdf.median(numeric_only=True))) + self.assert_eq(len(psdf.kurtosis(numeric_only=True)), len(pdf.kurtosis(numeric_only=True))) + self.assert_eq(len(psdf.skew(numeric_only=True)), len(pdf.skew(numeric_only=True))) + + # Boolean was excluded because of a behavior change in NumPy + # https://github.com/numpy/numpy/pull/16273#discussion_r641264085 which pandas inherits + # but this behavior is inconsistent in pandas context. + # Boolean column in quantile tests are excluded for now. + # TODO(SPARK-35555): track and match the behavior of quantile to pandas' + pdf = pd.DataFrame({"i": [0, 1, 2], "s": ["x", "y", "z"]}) + psdf = ps.from_pandas(pdf) + self.assert_eq( + len(psdf.quantile(q=0.5, numeric_only=True)), + len(pdf.quantile(q=0.5, numeric_only=True)), + ) + self.assert_eq( + len(psdf.quantile(q=[0.25, 0.5, 0.75], numeric_only=True)), + len(pdf.quantile(q=[0.25, 0.5, 0.75], numeric_only=True)), + ) + + def test_numeric_only_unsupported(self): + pdf = pd.DataFrame({"i": [0, 1, 2], "b": [False, False, True], "s": ["x", "y", "z"]}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(psdf.sum(numeric_only=True), pdf.sum(numeric_only=True)) + self.assert_eq( + psdf[["i", "b"]].sum(numeric_only=False), pdf[["i", "b"]].sum(numeric_only=False) + ) + + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + psdf.sum(numeric_only=False) + + with self.assertRaisesRegex(TypeError, "Could not convert object \\(string\\) to numeric"): + psdf.s.sum() + + +class StatsTests(StatsTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_stats import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_typedef.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_typedef.py new file mode 100644 index 0000000..97e400d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_typedef.py @@ -0,0 +1,447 @@ +# +# 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 sys +import unittest +import datetime +import decimal +from distutils.version import LooseVersion +from typing import List + +import pandas +import pandas as pd +from pandas.api.types import CategoricalDtype +import numpy as np +from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + FloatType, + IntegerType, + LongType, + StringType, + StructField, + StructType, + ByteType, + ShortType, + DateType, + DecimalType, + DoubleType, + TimestampType, +) + +from pyspark.pandas.typedef import ( + as_spark_type, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, + infer_return_type, + pandas_on_spark_type, +) +from pyspark import pandas as ps + + +class TypeHintTestsMixin: + def test_infer_schema_with_no_return(self): + def try_infer_return_type(): + def f(): + pass + + infer_return_type(f) + + self.assertRaisesRegex( + ValueError, "A return value is required for the input function", try_infer_return_type + ) + + def try_infer_return_type(): + def f() -> None: + pass + + infer_return_type(f) + + self.assertRaisesRegex( + TypeError, "Type was not understood", try_infer_return_type + ) + + def test_infer_schema_from_pandas_instances(self): + def func() -> pd.Series[int]: + pass + + inferred = infer_return_type(func) + self.assertEqual(inferred.dtype, np.int64) + self.assertEqual(inferred.spark_type, LongType()) + + def func() -> pd.Series[float]: + pass + + inferred = infer_return_type(func) + self.assertEqual(inferred.dtype, np.float64) + self.assertEqual(inferred.spark_type, DoubleType()) + + def func() -> "pd.DataFrame[np.float_, str]": + pass + + expected = StructType([StructField("c0", DoubleType()), StructField("c1", StringType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64, np.unicode_]) + self.assertEqual(inferred.spark_type, expected) + + def func() -> "pandas.DataFrame[float]": + pass + + expected = StructType([StructField("c0", DoubleType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64]) + self.assertEqual(inferred.spark_type, expected) + + def func() -> "pd.Series[int]": + pass + + inferred = infer_return_type(func) + self.assertEqual(inferred.dtype, np.int64) + self.assertEqual(inferred.spark_type, LongType()) + + def func() -> pd.DataFrame[np.float64, str]: + pass + + expected = StructType([StructField("c0", DoubleType()), StructField("c1", StringType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64, np.unicode_]) + self.assertEqual(inferred.spark_type, expected) + + def func() -> pd.DataFrame[np.float_]: + pass + + expected = StructType([StructField("c0", DoubleType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64]) + self.assertEqual(inferred.spark_type, expected) + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) + + def func() -> pd.DataFrame[pdf.dtypes]: + pass + + expected = StructType([StructField("c0", LongType()), StructField("c1", LongType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.int64, np.int64]) + self.assertEqual(inferred.spark_type, expected) + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": pd.Categorical(["a", "b", "c"])}) + + def func() -> pd.Series[pdf.b.dtype]: + pass + + inferred = infer_return_type(func) + self.assertEqual(inferred.dtype, CategoricalDtype(categories=["a", "b", "c"])) + self.assertEqual(inferred.spark_type, LongType()) + + def func() -> pd.DataFrame[pdf.dtypes]: + pass + + expected = StructType([StructField("c0", LongType()), StructField("c1", LongType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.int64, CategoricalDtype(categories=["a", "b", "c"])]) + self.assertEqual(inferred.spark_type, expected) + + def test_if_pandas_implements_class_getitem(self): + # the current type hint implementation of pandas DataFrame assumes pandas doesn't + # implement '__class_getitem__'. This test case is to make sure pandas + # doesn't implement them. + assert not ps._frame_has_class_getitem + assert not ps._series_has_class_getitem + + def test_infer_schema_with_names_pandas_instances(self): + def func() -> 'pd.DataFrame["a" : np.float_, "b":str]': # noqa: F405 + pass + + expected = StructType([StructField("a", DoubleType()), StructField("b", StringType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64, np.unicode_]) + self.assertEqual(inferred.spark_type, expected) + + def func() -> "pd.DataFrame['a': float, 'b': int]": # noqa: F405 + pass + + expected = StructType([StructField("a", DoubleType()), StructField("b", LongType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.float64, np.int64]) + self.assertEqual(inferred.spark_type, expected) + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) + + def func() -> pd.DataFrame[zip(pdf.columns, pdf.dtypes)]: + pass + + expected = StructType([StructField("a", LongType()), StructField("b", LongType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.int64, np.int64]) + self.assertEqual(inferred.spark_type, expected) + + pdf = pd.DataFrame({("x", "a"): [1, 2, 3], ("y", "b"): [3, 4, 5]}) + + def func() -> pd.DataFrame[zip(pdf.columns, pdf.dtypes)]: + pass + + expected = StructType( + [StructField("(x, a)", LongType()), StructField("(y, b)", LongType())] + ) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.int64, np.int64]) + self.assertEqual(inferred.spark_type, expected) + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": pd.Categorical(["a", "b", "c"])}) + + def func() -> pd.DataFrame[zip(pdf.columns, pdf.dtypes)]: + pass + + expected = StructType([StructField("a", LongType()), StructField("b", LongType())]) + inferred = infer_return_type(func) + self.assertEqual(inferred.dtypes, [np.int64, CategoricalDtype(categories=["a", "b", "c"])]) + self.assertEqual(inferred.spark_type, expected) + + def test_infer_schema_with_names_pandas_instances_negative(self): + def try_infer_return_type(): + def f() -> 'pd.DataFrame["a" : np.float_ : 1, "b":str:2]': # noqa: F405 + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type) + + class A: + pass + + def try_infer_return_type(): + def f() -> pd.DataFrame[A]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "not understood", try_infer_return_type) + + def try_infer_return_type(): + def f() -> 'pd.DataFrame["a" : float : 1, "b":str:2]': # noqa: F405 + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type) + + # object type + pdf = pd.DataFrame({"a": ["a", 2, None]}) + + def try_infer_return_type(): + def f() -> pd.DataFrame[pdf.dtypes]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type) + + def try_infer_return_type(): + def f() -> pd.Series[pdf.a.dtype]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type) + + def test_infer_schema_with_names_negative(self): + def try_infer_return_type(): + def f() -> 'ps.DataFrame["a" : float : 1, "b":str:2]': # noqa: F405 + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type) + + class A: + pass + + def try_infer_return_type(): + def f() -> ps.DataFrame[A]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "not understood", try_infer_return_type) + + def try_infer_return_type(): + def f() -> 'ps.DataFrame["a" : np.float_ : 1, "b":str:2]': # noqa: F405 + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type) + + # object type + pdf = pd.DataFrame({"a": ["a", 2, None]}) + + def try_infer_return_type(): + def f() -> ps.DataFrame[pdf.dtypes]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type) + + def try_infer_return_type(): + def f() -> ps.Series[pdf.a.dtype]: + pass + + infer_return_type(f) + + self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type) + + def test_as_spark_type_pandas_on_spark_dtype(self): + type_mapper = { + # binary + np.character: (np.character, BinaryType()), + np.bytes_: (np.bytes_, BinaryType()), + np.string_: (np.bytes_, BinaryType()), + bytes: (np.bytes_, BinaryType()), + # integer + np.int8: (np.int8, ByteType()), + np.byte: (np.int8, ByteType()), + np.int16: (np.int16, ShortType()), + np.int32: (np.int32, IntegerType()), + np.int64: (np.int64, LongType()), + int: (np.int64, LongType()), + # floating + np.float32: (np.float32, FloatType()), + np.float64: (np.float64, DoubleType()), + float: (np.float64, DoubleType()), + # string + np.unicode_: (np.unicode_, StringType()), + str: (np.unicode_, StringType()), + # bool + bool: (np.bool_, BooleanType()), + # datetime + np.datetime64: (np.datetime64, TimestampType()), + datetime.datetime: (np.dtype("datetime64[ns]"), TimestampType()), + # DateType + datetime.date: (np.dtype("object"), DateType()), + # DecimalType + decimal.Decimal: (np.dtype("object"), DecimalType(38, 18)), + # ArrayType + np.ndarray: (np.dtype("object"), ArrayType(StringType())), + # CategoricalDtype + CategoricalDtype(categories=["a", "b", "c"]): ( + CategoricalDtype(categories=["a", "b", "c"]), + LongType(), + ), + } + + for numpy_or_python_type, (dtype, spark_type) in type_mapper.items(): + self.assertEqual(as_spark_type(numpy_or_python_type), spark_type) + self.assertEqual(pandas_on_spark_type(numpy_or_python_type), (dtype, spark_type)) + + if isinstance(numpy_or_python_type, CategoricalDtype): + # Nested CategoricalDtype is not yet supported. + continue + + self.assertEqual(as_spark_type(List[numpy_or_python_type]), ArrayType(spark_type)) + self.assertEqual( + pandas_on_spark_type(List[numpy_or_python_type]), + (np.dtype("object"), ArrayType(spark_type)), + ) + + # For NumPy typing, NumPy version should be 1.21+ and Python version should be 3.8+ + if sys.version_info >= (3, 8) and LooseVersion(np.__version__) >= LooseVersion("1.21"): + import numpy.typing as ntp + + self.assertEqual( + as_spark_type(ntp.NDArray[numpy_or_python_type]), ArrayType(spark_type) + ) + self.assertEqual( + pandas_on_spark_type(ntp.NDArray[numpy_or_python_type]), + (np.dtype("object"), ArrayType(spark_type)), + ) + + with self.assertRaisesRegex(TypeError, "Type uint64 was not understood."): + as_spark_type(np.dtype("uint64")) + + with self.assertRaisesRegex(TypeError, "Type object was not understood."): + as_spark_type(np.dtype("object")) + + with self.assertRaisesRegex(TypeError, "Type uint64 was not understood."): + pandas_on_spark_type(np.dtype("uint64")) + + with self.assertRaisesRegex(TypeError, "Type object was not understood."): + pandas_on_spark_type(np.dtype("object")) + + @unittest.skipIf(not extension_dtypes_available, "The pandas extension types are not available") + def test_as_spark_type_extension_dtypes(self): + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + + type_mapper = { + Int8Dtype(): ByteType(), + Int16Dtype(): ShortType(), + Int32Dtype(): IntegerType(), + Int64Dtype(): LongType(), + } + + for extension_dtype, spark_type in type_mapper.items(): + self.assertEqual(as_spark_type(extension_dtype), spark_type) + self.assertEqual(pandas_on_spark_type(extension_dtype), (extension_dtype, spark_type)) + + @unittest.skipIf( + not extension_object_dtypes_available, "The pandas extension object types are not available" + ) + def test_as_spark_type_extension_object_dtypes(self): + from pandas import BooleanDtype, StringDtype + + type_mapper = { + BooleanDtype(): BooleanType(), + StringDtype(): StringType(), + } + + for extension_dtype, spark_type in type_mapper.items(): + self.assertEqual(as_spark_type(extension_dtype), spark_type) + self.assertEqual(pandas_on_spark_type(extension_dtype), (extension_dtype, spark_type)) + + @unittest.skipIf( + not extension_float_dtypes_available, "The pandas extension float types are not available" + ) + def test_as_spark_type_extension_float_dtypes(self): + from pandas import Float32Dtype, Float64Dtype + + type_mapper = { + Float32Dtype(): FloatType(), + Float64Dtype(): DoubleType(), + } + + for extension_dtype, spark_type in type_mapper.items(): + self.assertEqual(as_spark_type(extension_dtype), spark_type) + self.assertEqual(pandas_on_spark_type(extension_dtype), (extension_dtype, spark_type)) + + +class TypeHintTests(TypeHintTestsMixin, unittest.TestCase): + pass + + +if __name__ == "__main__": + from pyspark.pandas.tests.test_typedef import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_utils.py new file mode 100644 index 0000000..60961dc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_utils.py @@ -0,0 +1,301 @@ +# +# 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 pandas as pd + +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.utils import ( + lazy_property, + validate_arguments_and_invoke_function, + validate_bool_kwarg, + validate_index_loc, + validate_mode, +) +from pyspark.testing.pandasutils import ( + PandasOnSparkTestCase, + assertPandasOnSparkEqual, + _assert_pandas_equal, + _assert_pandas_almost_equal, +) +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.errors import PySparkAssertionError + +some_global_variable = 0 + + +class UtilsTestsMixin: + + # a dummy to_html version with an extra parameter that pandas does not support + # used in test_validate_arguments_and_invoke_function + def to_html(self, max_rows=None, unsupported_param=None): + args = locals() + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]) + validate_arguments_and_invoke_function(pdf, self.to_html, pd.DataFrame.to_html, args) + + def to_clipboard(self, sep=",", **kwargs): + args = locals() + + pdf = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]) + validate_arguments_and_invoke_function( + pdf, self.to_clipboard, pd.DataFrame.to_clipboard, args + ) + + # Support for **kwargs + self.to_clipboard(sep=",", index=False) + + def test_validate_arguments_and_invoke_function(self): + # This should pass and run fine + self.to_html() + self.to_html(unsupported_param=None) + self.to_html(max_rows=5) + + # This should fail because we are explicitly setting an unsupported param + # to a non-default value + with self.assertRaises(TypeError): + self.to_html(unsupported_param=1) + + def test_lazy_property(self): + obj = TestClassForLazyProp() + # If lazy prop is not working, the second test would fail (because it'd be 2) + self.assert_eq(obj.lazy_prop, 1) + self.assert_eq(obj.lazy_prop, 1) + + def test_validate_bool_kwarg(self): + # This should pass and run fine + pandas_on_spark = True + self.assert_eq(validate_bool_kwarg(pandas_on_spark, "pandas_on_spark"), True) + pandas_on_spark = False + self.assert_eq(validate_bool_kwarg(pandas_on_spark, "pandas_on_spark"), False) + pandas_on_spark = None + self.assert_eq(validate_bool_kwarg(pandas_on_spark, "pandas_on_spark"), None) + + # This should fail because we are explicitly setting a non-boolean value + pandas_on_spark = "true" + with self.assertRaisesRegex( + TypeError, 'For argument "pandas_on_spark" expected type bool, received type str.' + ): + validate_bool_kwarg(pandas_on_spark, "pandas_on_spark") + + def test_validate_mode(self): + self.assert_eq(validate_mode("a"), "append") + self.assert_eq(validate_mode("w"), "overwrite") + self.assert_eq(validate_mode("a+"), "append") + self.assert_eq(validate_mode("w+"), "overwrite") + + with self.assertRaises(ValueError): + validate_mode("r") + + def test_validate_index_loc(self): + psidx = Index([1, 2, 3]) + validate_index_loc(psidx, -1) + validate_index_loc(psidx, -3) + err_msg = "index 4 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + validate_index_loc(psidx, 4) + err_msg = "index -4 is out of bounds for axis 0 with size 3" + with self.assertRaisesRegex(IndexError, err_msg): + validate_index_loc(psidx, -4) + + def test_assert_df_assert_pandas_on_spark_equal(self): + import pyspark.pandas as ps + + psdf1 = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + psdf2 = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + + assertPandasOnSparkEqual(psdf1, psdf2, checkRowOrder=False) + assertPandasOnSparkEqual(psdf1, psdf2, checkRowOrder=True) + + def test_assert_pandas_on_spark_equal_ignore_order(self): + import pyspark.pandas as ps + + psdf1 = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + psdf2 = ps.DataFrame({"a": [2, 1, 3], "b": [5, 4, 6], "c": [8, 7, 9]}) + + assertPandasOnSparkEqual(psdf1, psdf2, checkRowOrder=False) + + def test_assert_series_assert_pandas_on_spark_equal(self): + import pyspark.pandas as ps + + s1 = ps.Series([212.32, 100.0001]) + s2 = ps.Series([212.32, 100.0001]) + + assertPandasOnSparkEqual(s1, s2, checkExact=False) + + def test_assert_index_assert_pandas_on_spark_equal(self): + import pyspark.pandas as ps + + s1 = ps.Index([212.300001, 100.000]) + s2 = ps.Index([212.3, 100.0001]) + + assertPandasOnSparkEqual(s1, s2, almost=True) + + def test_assert_error_assert_pandas_on_spark_equal(self): + import pyspark.pandas as ps + + list1 = [10, 20, 30] + list2 = [10, 20, 30] + + with self.assertRaises(PySparkAssertionError) as pe: + assertPandasOnSparkEqual(list1, list2) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{ps.DataFrame.__name__}, " + f"{ps.Series.__name__}, " + f"{ps.Index.__name__}", + "arg_name": "actual", + "actual_type": type(list1), + }, + ) + + def test_assert_None_assert_pandas_on_spark_equal(self): + psdf1 = None + psdf2 = None + + assertPandasOnSparkEqual(psdf1, psdf2) + + def test_assert_empty_assert_pandas_on_spark_equal(self): + import pyspark.pandas as ps + + psdf1 = ps.DataFrame() + psdf2 = ps.DataFrame() + + assertPandasOnSparkEqual(psdf1, psdf2) + + def test_dataframe_error_assert_pandas_equal(self): + pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]) + pdf2 = pd.DataFrame({"a": [1, 3, 3], "b": [4, 5, 6]}, index=[0, 1, 3]) + + with self.assertRaises(PySparkAssertionError) as pe: + _assert_pandas_equal(pdf1, pdf2, True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": pdf1.to_string(), + "left_dtype": str(pdf1.dtypes), + "right": pdf2.to_string(), + "right_dtype": str(pdf2.dtypes), + }, + ) + + def test_series_error_assert_pandas_equal(self): + series1 = pd.Series([1, 2, 3]) + series2 = pd.Series([4, 5, 6]) + + with self.assertRaises(PySparkAssertionError) as pe: + _assert_pandas_equal(series1, series2, True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_SERIES", + message_parameters={ + "left": series1.to_string(), + "left_dtype": str(series1.dtype), + "right": series2.to_string(), + "right_dtype": str(series2.dtype), + }, + ) + + def test_index_error_assert_pandas_equal(self): + index1 = pd.Index([1, 2, 3]) + index2 = pd.Index([4, 5, 6]) + + with self.assertRaises(PySparkAssertionError) as pe: + _assert_pandas_equal(index1, index2, True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_INDEX", + message_parameters={ + "left": index1, + "left_dtype": str(index1.dtype), + "right": index2, + "right_dtype": str(index2.dtype), + }, + ) + + def test_multiindex_error_assert_pandas_almost_equal(self): + pdf1 = pd.DataFrame({"a": [1, 2], "b": [4, 10]}, index=[0, 1]) + pdf2 = pd.DataFrame({"a": [1, 5, 3], "b": [1, 5, 6]}, index=[0, 1, 3]) + multiindex1 = pd.MultiIndex.from_frame(pdf1) + multiindex2 = pd.MultiIndex.from_frame(pdf2) + + with self.assertRaises(PySparkAssertionError) as pe: + _assert_pandas_almost_equal(multiindex1, multiindex2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_MULTIINDEX", + message_parameters={ + "left": multiindex1, + "left_dtype": str(multiindex1.dtype), + "right": multiindex2, + "right_dtype": str(multiindex1.dtype), + }, + ) + + def test_dataframe_error_assert_pandas_on_spark_almost_equal(self): + import pyspark.pandas as ps + + psdf1 = ps.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + psdf2 = ps.DataFrame({"a": [1, 2], "b": [4, 5], "c": [7, 8]}) + + with self.assertRaises(PySparkAssertionError) as pe: + assertPandasOnSparkEqual(psdf1, psdf2, almost=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": psdf1.to_string(), + "left_dtype": str(psdf1.dtypes), + "right": psdf2.to_string(), + "right_dtype": str(psdf2.dtypes), + }, + ) + + +class TestClassForLazyProp: + def __init__(self): + self.some_variable = 0 + + @lazy_property + def lazy_prop(self): + self.some_variable += 1 + return self.some_variable + + +class UtilsTests(UtilsTestsMixin, PandasOnSparkTestCase, SQLTestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_utils import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_window.py new file mode 100644 index 0000000..33f06a1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/tests/test_window.py @@ -0,0 +1,465 @@ +# +# 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 inspect + +from pyspark import pandas as ps +from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.window import ( + MissingPandasLikeExpanding, + MissingPandasLikeRolling, + MissingPandasLikeExpandingGroupby, + MissingPandasLikeRollingGroupby, + MissingPandasLikeExponentialMoving, + MissingPandasLikeExponentialMovingGroupby, +) +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils + + +class ExpandingRollingTestsMixin: + def test_missing(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}) + + # Expanding functions + missing_functions = inspect.getmembers(MissingPandasLikeExpanding, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.expanding(1), name)() # Frame + + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.expanding(1), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.expanding(1), name)() # Frame + + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.expanding(1), name)() # Series + + # Rolling functions + missing_functions = inspect.getmembers(MissingPandasLikeRolling, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.rolling(1), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.rolling(1), name)() # Series + + # ExponentialMoving functions + missing_functions = inspect.getmembers( + MissingPandasLikeExponentialMoving, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + + # Expanding properties + missing_properties = inspect.getmembers( + MissingPandasLikeExpanding, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.expanding(1), name) # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.expanding(1), name) # Series + + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.expanding(1), name) # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.expanding(1), name) # Series + + # Rolling properties + missing_properties = inspect.getmembers( + MissingPandasLikeRolling, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.rolling(1), name)() # Series + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.rolling(1), name)() # Series + + # ExponentialMoving properties + missing_properties = inspect.getmembers( + MissingPandasLikeExponentialMoving, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + + def test_missing_groupby(self): + psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}) + + # Expanding functions + missing_functions = inspect.getmembers( + MissingPandasLikeExpandingGroupby, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").expanding(1), name)() # Frame + + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).expanding(1), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.groupby("a").expanding(1), name)() # Frame + + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.groupby(psdf.a).expanding(1), name)() # Series + + # Rolling functions + missing_functions = inspect.getmembers(MissingPandasLikeRollingGroupby, inspect.isfunction) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).rolling(1), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "method.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.rolling(1), name)() # Series + + # ExponentialMoving functions + missing_functions = inspect.getmembers( + MissingPandasLikeExponentialMovingGroupby, inspect.isfunction + ) + unsupported_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" + ] + for name in unsupported_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).ewm(com=0.5), name)() # Series + + deprecated_functions = [ + name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" + ] + for name in deprecated_functions: + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "method.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + + # Expanding properties + missing_properties = inspect.getmembers( + MissingPandasLikeExpandingGroupby, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").expanding(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Expanding.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).expanding(1), name)() # Series + + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.expanding(1), name) # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Expanding.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.expanding(1), name) # Series + + # Rolling properties + missing_properties = inspect.getmembers( + MissingPandasLikeRollingGroupby, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*Rolling.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).rolling(1), name)() # Series + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.rolling(1), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, "property.*Rolling.*{}.*is deprecated".format(name) + ): + getattr(psdf.a.rolling(1), name)() # Series + + # ExponentialMoving properties + missing_properties = inspect.getmembers( + MissingPandasLikeExponentialMovingGroupby, lambda o: isinstance(o, property) + ) + unsupported_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "unsupported_property" + ] + for name in unsupported_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.groupby("a").ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*not implemented( yet\\.|\\. .+)".format(name), + ): + getattr(psdf.a.groupby(psdf.a).ewm(com=0.5), name)() # Series + deprecated_properties = [ + name + for (name, type_) in missing_properties + if type_.fget.__name__ == "deprecated_property" + ] + for name in deprecated_properties: + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.ewm(com=0.5), name)() # Frame + with self.assertRaisesRegex( + PandasNotImplementedError, + "property.*ExponentialMoving.*{}.*is deprecated".format(name), + ): + getattr(psdf.a.ewm(com=0.5), name)() # Series + + +class ExpandingRollingTests(ExpandingRollingTestsMixin, PandasOnSparkTestCase, TestUtils): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.pandas.tests.test_window import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/__init__.py new file mode 100644 index 0000000..4949067 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/__init__.py @@ -0,0 +1,18 @@ +# +# 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.pandas.typedef.typehints import * # noqa: F401,F403,F405 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/typehints.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/typehints.py new file mode 100644 index 0000000..6e41395 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/typedef/typehints.py @@ -0,0 +1,862 @@ +# +# 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. +# + +""" +Utilities to deal with types. This is mostly focused on python3. +""" +import datetime +import decimal +import sys +import typing +from collections.abc import Iterable +from distutils.version import LooseVersion +from inspect import isclass +from typing import Any, Callable, Generic, List, Tuple, Union, Type, get_type_hints + +import numpy as np +import pandas as pd +from pandas.api.types import CategoricalDtype, pandas_dtype # type: ignore[attr-defined] +from pandas.api.extensions import ExtensionDtype + +extension_dtypes: Tuple[type, ...] +try: + from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype + + extension_dtypes_available = True + extension_dtypes = (Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype) + + try: + from pandas import BooleanDtype, StringDtype + + extension_object_dtypes_available = True + extension_dtypes += (BooleanDtype, StringDtype) + except ImportError: + extension_object_dtypes_available = False + + try: + from pandas import Float32Dtype, Float64Dtype + + extension_float_dtypes_available = True + extension_dtypes += (Float32Dtype, Float64Dtype) + except ImportError: + extension_float_dtypes_available = False + +except ImportError: + extension_dtypes_available = False + extension_object_dtypes_available = False + extension_float_dtypes_available = False + extension_dtypes = () + +import pyarrow as pa +import pyspark.sql.types as types +from pyspark.sql.pandas.types import to_arrow_type, from_arrow_type + +# For running doctests and reference resolution in PyCharm. +from pyspark import pandas as ps # noqa: F401 +from pyspark.pandas._typing import Dtype, T + +if typing.TYPE_CHECKING: + from pyspark.pandas.internal import InternalField + + +# A column of data, with the data type. +class SeriesType(Generic[T]): + def __init__(self, dtype: Dtype, spark_type: types.DataType): + self.dtype = dtype + self.spark_type = spark_type + + def __repr__(self) -> str: + return "SeriesType[{}]".format(self.spark_type) + + +class DataFrameType: + def __init__( + self, + index_fields: List["InternalField"], + data_fields: List["InternalField"], + ): + self.index_fields = index_fields + self.data_fields = data_fields + self.fields = index_fields + data_fields + + @property + def dtypes(self) -> List[Dtype]: + return [field.dtype for field in self.fields] + + @property + def spark_type(self) -> types.StructType: + return types.StructType([field.struct_field for field in self.fields]) + + def __repr__(self) -> str: + return "DataFrameType[{}]".format(self.spark_type) + + +# The type is a scalar type that is furthermore understood by Spark. +class ScalarType: + def __init__(self, dtype: Dtype, spark_type: types.DataType): + self.dtype = dtype + self.spark_type = spark_type + + def __repr__(self) -> str: + return "ScalarType[{}]".format(self.spark_type) + + +# The type is left unspecified or we do not know about this type. +class UnknownType: + def __init__(self, tpe: Any): + self.tpe = tpe + + def __repr__(self) -> str: + return "UnknownType[{}]".format(self.tpe) + + +class IndexNameTypeHolder: + name = None + tpe = None + short_name = "IndexNameType" + + +class NameTypeHolder: + name = None + tpe = None + short_name = "NameType" + + +def as_spark_type( + tpe: Union[str, type, Dtype], *, raise_error: bool = True, prefer_timestamp_ntz: bool = False +) -> types.DataType: + """ + Given a Python type, returns the equivalent spark type. + Accepts: + - the built-in types in Python + - the built-in types in numpy + - list of pairs of (field_name, type) + - dictionaries of field_name -> type + - Python3's typing system + """ + # For NumPy typing, NumPy version should be 1.21+ and Python version should be 3.8+ + if sys.version_info >= (3, 8) and LooseVersion(np.__version__) >= LooseVersion("1.21"): + if ( + hasattr(tpe, "__origin__") + and tpe.__origin__ is np.ndarray # type: ignore[union-attr] + and hasattr(tpe, "__args__") + and len(tpe.__args__) > 1 # type: ignore[union-attr] + ): + # numpy.typing.NDArray + return types.ArrayType( + as_spark_type( + tpe.__args__[1].__args__[0], raise_error=raise_error # type: ignore[union-attr] + ) + ) + + if isinstance(tpe, np.dtype) and tpe == np.dtype("object"): + pass + # ArrayType + elif tpe in (np.ndarray,): + return types.ArrayType(types.StringType()) + elif hasattr(tpe, "__origin__") and issubclass( + tpe.__origin__, list # type: ignore[union-attr] + ): + element_type = as_spark_type( + tpe.__args__[0], raise_error=raise_error # type: ignore[union-attr] + ) + if element_type is None: + return None + return types.ArrayType(element_type) + # BinaryType + elif tpe in (bytes, np.character, np.bytes_, np.string_): + return types.BinaryType() + # BooleanType + elif tpe in (bool, np.bool_, "bool", "?"): + return types.BooleanType() + # DateType + elif tpe in (datetime.date,): + return types.DateType() + # NumericType + elif tpe in (np.int8, np.byte, "int8", "byte", "b"): + return types.ByteType() + elif tpe in (decimal.Decimal,): + # TODO: considering the precision & scale for decimal type. + return types.DecimalType(38, 18) + elif tpe in (float, np.float_, np.float64, "float", "float64", "double"): + return types.DoubleType() + elif tpe in (np.float32, "float32", "f"): + return types.FloatType() + elif tpe in (np.int32, "int32", "i"): + return types.IntegerType() + elif tpe in (int, np.int64, "int", "int64", "long"): + return types.LongType() + elif tpe in (np.int16, "int16", "short"): + return types.ShortType() + # StringType + elif tpe in (str, np.unicode_, "str", "U"): + return types.StringType() + # TimestampType or TimestampNTZType if timezone is not specified. + elif tpe in (datetime.datetime, np.datetime64, "datetime64[ns]", "M", pd.Timestamp): + return types.TimestampNTZType() if prefer_timestamp_ntz else types.TimestampType() + + # DayTimeIntervalType + elif tpe in (datetime.timedelta, np.timedelta64, "timedelta64[ns]"): + return types.DayTimeIntervalType() + + # categorical types + elif isinstance(tpe, CategoricalDtype) or (isinstance(tpe, str) and tpe == "category"): + return types.LongType() + + # extension types + elif extension_dtypes_available: + # IntegralType + if isinstance(tpe, Int8Dtype) or (isinstance(tpe, str) and tpe == "Int8"): + return types.ByteType() + elif isinstance(tpe, Int16Dtype) or (isinstance(tpe, str) and tpe == "Int16"): + return types.ShortType() + elif isinstance(tpe, Int32Dtype) or (isinstance(tpe, str) and tpe == "Int32"): + return types.IntegerType() + elif isinstance(tpe, Int64Dtype) or (isinstance(tpe, str) and tpe == "Int64"): + return types.LongType() + + if extension_object_dtypes_available: + # BooleanType + if isinstance(tpe, BooleanDtype) or (isinstance(tpe, str) and tpe == "boolean"): + return types.BooleanType() + # StringType + elif isinstance(tpe, StringDtype) or (isinstance(tpe, str) and tpe == "string"): + return types.StringType() + + if extension_float_dtypes_available: + # FractionalType + if isinstance(tpe, Float32Dtype) or (isinstance(tpe, str) and tpe == "Float32"): + return types.FloatType() + elif isinstance(tpe, Float64Dtype) or (isinstance(tpe, str) and tpe == "Float64"): + return types.DoubleType() + + if raise_error: + raise TypeError("Type %s was not understood." % tpe) + else: + return None + + +def spark_type_to_pandas_dtype( + spark_type: types.DataType, *, use_extension_dtypes: bool = False +) -> Dtype: + """Return the given Spark DataType to pandas dtype.""" + + if use_extension_dtypes and extension_dtypes_available: + # IntegralType + if isinstance(spark_type, types.ByteType): + return Int8Dtype() + elif isinstance(spark_type, types.ShortType): + return Int16Dtype() + elif isinstance(spark_type, types.IntegerType): + return Int32Dtype() + elif isinstance(spark_type, types.LongType): + return Int64Dtype() + + if extension_object_dtypes_available: + # BooleanType + if isinstance(spark_type, types.BooleanType): + return BooleanDtype() + # StringType + elif isinstance(spark_type, types.StringType): + return StringDtype() + + # FractionalType + if extension_float_dtypes_available: + if isinstance(spark_type, types.FloatType): + return Float32Dtype() + elif isinstance(spark_type, types.DoubleType): + return Float64Dtype() + + if isinstance( + spark_type, + ( + types.DateType, + types.NullType, + types.ArrayType, + types.MapType, + types.StructType, + types.UserDefinedType, + ), + ): + return np.dtype("object") + elif isinstance(spark_type, types.TimestampType): + return np.dtype("datetime64[ns]") + else: + return np.dtype(to_arrow_type(spark_type).to_pandas_dtype()) + + +def pandas_on_spark_type(tpe: Union[str, type, Dtype]) -> Tuple[Dtype, types.DataType]: + """ + Convert input into a pandas only dtype object or a numpy dtype object, + and its corresponding Spark DataType. + + Parameters + ---------- + tpe : object to be converted + + Returns + ------- + tuple of np.dtype or a pandas dtype, and Spark DataType + + Raises + ------ + TypeError if not a dtype + + Examples + -------- + >>> pandas_on_spark_type(int) + (dtype('int64'), LongType()) + >>> pandas_on_spark_type(str) + (dtype('>> pandas_on_spark_type(datetime.date) + (dtype('O'), DateType()) + >>> pandas_on_spark_type(datetime.datetime) + (dtype('>> pandas_on_spark_type(datetime.timedelta) + (dtype('>> pandas_on_spark_type(List[bool]) + (dtype('O'), ArrayType(BooleanType(), True)) + """ + try: + dtype = pandas_dtype(tpe) + spark_type = as_spark_type(dtype) + except TypeError: + spark_type = as_spark_type(tpe) + dtype = spark_type_to_pandas_dtype(spark_type) + return dtype, spark_type + + +def infer_pd_series_spark_type( + pser: pd.Series, dtype: Dtype, prefer_timestamp_ntz: bool = False +) -> types.DataType: + """Infer Spark DataType from pandas Series dtype. + + :param pser: :class:`pandas.Series` to be inferred + :param dtype: the Series' dtype + :param prefer_timestamp_ntz: if true, infers datetime without timezone as + TimestampNTZType type. If false, infers it as TimestampType. + :return: the inferred Spark data type + """ + if dtype == np.dtype("object"): + if len(pser) == 0 or pser.isnull().all(): + return types.NullType() + elif hasattr(pser.iloc[0], "__UDT__"): + return pser.iloc[0].__UDT__ + else: + return from_arrow_type(pa.Array.from_pandas(pser).type, prefer_timestamp_ntz) + elif isinstance(dtype, CategoricalDtype): + if isinstance(pser.dtype, CategoricalDtype): + return as_spark_type(pser.cat.codes.dtype, prefer_timestamp_ntz=prefer_timestamp_ntz) + else: + # `pser` must already be converted to codes. + return as_spark_type(pser.dtype, prefer_timestamp_ntz=prefer_timestamp_ntz) + else: + return as_spark_type(dtype, prefer_timestamp_ntz=prefer_timestamp_ntz) + + +def infer_return_type(f: Callable) -> Union[SeriesType, DataFrameType, ScalarType, UnknownType]: + """ + Infer the return type from the return type annotation of the given function. + + The returned type class indicates both dtypes (a pandas only dtype object + or a numpy dtype object) and its corresponding Spark DataType. + + >>> def func() -> int: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtype + dtype('int64') + >>> inferred.spark_type + LongType() + + >>> def func() -> ps.Series[int]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtype + dtype('int64') + >>> inferred.spark_type + LongType() + + >>> def func() -> ps.DataFrame[float, str]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64'), dtype('>> inferred.spark_type + StructType([StructField('c0', DoubleType(), True), StructField('c1', StringType(), True)]) + + >>> def func() -> ps.DataFrame[float]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64')] + >>> inferred.spark_type + StructType([StructField('c0', DoubleType(), True)]) + + >>> def func() -> 'int': + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtype + dtype('int64') + >>> inferred.spark_type + LongType() + + >>> def func() -> 'ps.Series[int]': + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtype + dtype('int64') + >>> inferred.spark_type + LongType() + + >>> def func() -> 'ps.DataFrame[float, str]': + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64'), dtype('>> inferred.spark_type + StructType([StructField('c0', DoubleType(), True), StructField('c1', StringType(), True)]) + + >>> def func() -> 'ps.DataFrame[float]': + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64')] + >>> inferred.spark_type + StructType([StructField('c0', DoubleType(), True)]) + + >>> def func() -> ps.DataFrame['a': float, 'b': int]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64'), dtype('int64')] + >>> inferred.spark_type + StructType([StructField('a', DoubleType(), True), StructField('b', LongType(), True)]) + + >>> def func() -> "ps.DataFrame['a': float, 'b': int]": + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('float64'), dtype('int64')] + >>> inferred.spark_type + StructType([StructField('a', DoubleType(), True), StructField('b', LongType(), True)]) + + >>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) + >>> def func() -> ps.DataFrame[pdf.dtypes]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64')] + >>> inferred.spark_type + StructType([StructField('c0', LongType(), True), StructField('c1', LongType(), True)]) + + >>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]}) + >>> def func() -> ps.DataFrame[zip(pdf.columns, pdf.dtypes)]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64')] + >>> inferred.spark_type + StructType([StructField('a', LongType(), True), StructField('b', LongType(), True)]) + + >>> pdf = pd.DataFrame({("x", "a"): [1, 2, 3], ("y", "b"): [3, 4, 5]}) + >>> def func() -> ps.DataFrame[zip(pdf.columns, pdf.dtypes)]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64')] + >>> inferred.spark_type + StructType([StructField('(x, a)', LongType(), True), StructField('(y, b)', LongType(), True)]) + + >>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": pd.Categorical([3, 4, 5])}) + >>> def func() -> ps.DataFrame[pdf.dtypes]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), CategoricalDtype(categories=[3, 4, 5], ordered=False)] + >>> inferred.spark_type + StructType([StructField('c0', LongType(), True), StructField('c1', LongType(), True)]) + + >>> def func() -> ps.DataFrame[zip(pdf.columns, pdf.dtypes)]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), CategoricalDtype(categories=[3, 4, 5], ordered=False)] + >>> inferred.spark_type + StructType([StructField('a', LongType(), True), StructField('b', LongType(), True)]) + + >>> def func() -> ps.Series[pdf.b.dtype]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtype + CategoricalDtype(categories=[3, 4, 5], ordered=False) + >>> inferred.spark_type + LongType() + + >>> def func() -> ps.DataFrame[int, [int, int]]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64'), dtype('int64')] + >>> inferred.spark_type.simpleString() + 'struct<__index_level_0__:bigint,c0:bigint,c1:bigint>' + >>> inferred.index_fields + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', LongType(), True))] + + >>> def func() -> ps.DataFrame[pdf.index.dtype, pdf.dtypes]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64'), CategoricalDtype(categories=[3, 4, 5], ordered=False)] + >>> inferred.spark_type.simpleString() + 'struct<__index_level_0__:bigint,c0:bigint,c1:bigint>' + >>> inferred.index_fields + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', LongType(), True))] + + >>> def func() -> ps.DataFrame[ + ... ("index", CategoricalDtype(categories=[3, 4, 5], ordered=False)), + ... [("id", int), ("A", int)]]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [CategoricalDtype(categories=[3, 4, 5], ordered=False), dtype('int64'), dtype('int64')] + >>> inferred.spark_type.simpleString() + 'struct' + >>> inferred.index_fields + [InternalField(dtype=category, struct_field=StructField('index', LongType(), True))] + + >>> def func() -> ps.DataFrame[ + ... (pdf.index.name, pdf.index.dtype), zip(pdf.columns, pdf.dtypes)]: + ... pass + >>> inferred = infer_return_type(func) + >>> inferred.dtypes + [dtype('int64'), dtype('int64'), CategoricalDtype(categories=[3, 4, 5], ordered=False)] + >>> inferred.spark_type.simpleString() + 'struct<__index_level_0__:bigint,a:bigint,b:bigint>' + >>> inferred.index_fields + [InternalField(dtype=int64, struct_field=StructField('__index_level_0__', LongType(), True))] + """ + # We should re-import to make sure the class 'SeriesType' is not treated as a class + # within this module locally. See Series.__class_getitem__ which imports this class + # canonically. + from pyspark.pandas.internal import InternalField, SPARK_INDEX_NAME_FORMAT + from pyspark.pandas.typedef import SeriesType, NameTypeHolder, IndexNameTypeHolder + from pyspark.pandas.utils import name_like_string + + tpe = get_type_hints(f).get("return", None) + + if tpe is None: + raise ValueError("A return value is required for the input function") + + if hasattr(tpe, "__origin__") and issubclass(tpe.__origin__, SeriesType): + tpe = tpe.__args__[0] + if issubclass(tpe, NameTypeHolder): + tpe = tpe.tpe + dtype, spark_type = pandas_on_spark_type(tpe) + return SeriesType(dtype, spark_type) + + # Note that, DataFrame type hints will create a Tuple. + # Tuple has _name but other types have __name__ + name = getattr(tpe, "_name", getattr(tpe, "__name__", None)) + # Check if the name is Tuple. + if name == "Tuple": + tuple_type = tpe + parameters = getattr(tuple_type, "__args__") + + index_parameters = [ + p for p in parameters if isclass(p) and issubclass(p, IndexNameTypeHolder) + ] + data_parameters = [p for p in parameters if p not in index_parameters] + assert len(data_parameters) > 0, "Type hints for data must not be empty." + + index_fields = [] + if len(index_parameters) >= 1: + for level, index_parameter in enumerate(index_parameters): + index_name = index_parameter.name + index_dtype, index_spark_type = pandas_on_spark_type(index_parameter.tpe) + index_fields.append( + InternalField( + dtype=index_dtype, + struct_field=types.StructField( + name=index_name + if index_name is not None + else SPARK_INDEX_NAME_FORMAT(level), + dataType=index_spark_type, + ), + ) + ) + else: + # No type hint for index. + assert len(index_parameters) == 0 + + data_dtypes, data_spark_types = zip( + *( + pandas_on_spark_type(p.tpe) + if isclass(p) and issubclass(p, NameTypeHolder) + else pandas_on_spark_type(p) + for p in data_parameters + ) + ) + data_names = [ + p.name if isclass(p) and issubclass(p, NameTypeHolder) else None + for p in data_parameters + ] + data_fields = [] + for i, (data_name, data_dtype, data_spark_type) in enumerate( + zip(data_names, data_dtypes, data_spark_types) + ): + data_fields.append( + InternalField( + dtype=data_dtype, + struct_field=types.StructField( + name=name_like_string(data_name) if data_name is not None else ("c%s" % i), + dataType=data_spark_type, + ), + ) + ) + + return DataFrameType(index_fields=index_fields, data_fields=data_fields) + + tpes = pandas_on_spark_type(tpe) + if tpes is None: + return UnknownType(tpe) + else: + return ScalarType(*tpes) + + +# TODO: once pandas exposes a typing module like numpy.typing, we should deprecate +# this logic and migrate to it by implementing the typing module in pandas API on Spark. + + +def create_type_for_series_type(param: Any) -> Type[SeriesType]: + """ + Supported syntax: + + >>> str(ps.Series[float]).endswith("SeriesType[float]") + True + """ + from pyspark.pandas.typedef import NameTypeHolder + + new_class: Type[NameTypeHolder] + if isinstance(param, ExtensionDtype): + new_class = type(NameTypeHolder.short_name, (NameTypeHolder,), {}) + new_class.tpe = param # type: ignore[assignment] + else: + new_class = param.type if isinstance(param, np.dtype) else param + + return SeriesType[new_class] # type: ignore[valid-type] + + +# TODO: Remove this variadic-generic hack by tuple once ww drop Python up to 3.9. +# See also PEP 646. One problem is that pandas doesn't inherits Generic[T] +# so we might have to leave this hack only for monkey-patching pandas DataFrame. +def create_tuple_for_frame_type(params: Any) -> object: + """ + This is a workaround to support variadic generic in DataFrame. + + See https://github.com/python/typing/issues/193 + we always wraps the given type hints by a tuple to mimic the variadic generic. + + Supported syntax: + + >>> import pandas as pd + >>> pdf = pd.DataFrame({'a': range(1)}) + + Typing data columns only: + + >>> ps.DataFrame[float, float] # doctest: +ELLIPSIS + typing.Tuple[...NameType, ...NameType] + >>> ps.DataFrame[pdf.dtypes] # doctest: +ELLIPSIS + typing.Tuple[...NameType] + >>> ps.DataFrame["id": int, "A": int] # doctest: +ELLIPSIS + typing.Tuple[...NameType, ...NameType] + >>> ps.DataFrame[zip(pdf.columns, pdf.dtypes)] # doctest: +ELLIPSIS + typing.Tuple[...NameType] + + Typing data columns with an index: + + >>> ps.DataFrame[int, [int, int]] # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...NameType, ...NameType] + >>> ps.DataFrame[pdf.index.dtype, pdf.dtypes] # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...NameType] + >>> ps.DataFrame[("index", int), [("id", int), ("A", int)]] # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...NameType, ...NameType] + >>> ps.DataFrame[(pdf.index.name, pdf.index.dtype), zip(pdf.columns, pdf.dtypes)] + ... # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...NameType] + + Typing data columns with an Multi-index: + >>> arrays = [[1, 1, 2], ['red', 'blue', 'red']] + >>> idx = pd.MultiIndex.from_arrays(arrays, names=('number', 'color')) + >>> pdf = pd.DataFrame({'a': range(3)}, index=idx) + >>> ps.DataFrame[[int, int], [int, int]] # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...IndexNameType, ...NameType, ...NameType] + >>> ps.DataFrame[pdf.index.dtypes, pdf.dtypes] # doctest: +ELLIPSIS, +SKIP + typing.Tuple[...IndexNameType, ...NameType] + >>> ps.DataFrame[[("index-1", int), ("index-2", int)], [("id", int), ("A", int)]] + ... # doctest: +ELLIPSIS + typing.Tuple[...IndexNameType, ...IndexNameType, ...NameType, ...NameType] + >>> ps.DataFrame[zip(pdf.index.names, pdf.index.dtypes), zip(pdf.columns, pdf.dtypes)] + ... # doctest: +ELLIPSIS, +SKIP + typing.Tuple[...IndexNameType, ...NameType] + """ + return Tuple[_to_type_holders(params)] + + +def _to_type_holders(params: Any) -> Tuple: + from pyspark.pandas.typedef import NameTypeHolder, IndexNameTypeHolder + + is_with_index = ( + isinstance(params, tuple) + and len(params) == 2 + and isinstance(params[1], (zip, list, pd.Series)) + ) + + if is_with_index: + # With index + # DataFrame[index_type, [type, ...]] + # DataFrame[dtype instance, dtypes instance] + # DataFrame[[index_type, ...], [type, ...]] + # DataFrame[dtypes instance, dtypes instance] + # DataFrame[(index_name, index_type), [(name, type), ...]] + # DataFrame[(index_name, index_type), zip(names, types)] + # DataFrame[[(index_name, index_type), ...], [(name, type), ...]] + # DataFrame[zip(index_names, index_types), zip(names, types)] + def is_list_of_pairs(p: Any) -> bool: + return ( + isinstance(p, list) + and len(p) >= 1 + and all(isinstance(param, tuple) and (len(param) == 2) for param in p) + ) + + index_params = params[0] + if isinstance(index_params, tuple) and len(index_params) == 2: + # DataFrame[("index", int), ...] + index_params = [index_params] + + if is_list_of_pairs(index_params): + # DataFrame[[("index", int), ("index-2", int)], ...] + index_params = tuple(slice(name, tpe) for name, tpe in index_params) + + index_types = _new_type_holders(index_params, IndexNameTypeHolder) + + data_types = params[1] + if is_list_of_pairs(data_types): + # DataFrame[..., [("id", int), ("A", int)]] + data_types = tuple(slice(*data_type) for data_type in data_types) + + data_types = _new_type_holders(data_types, NameTypeHolder) + + return index_types + data_types + else: + # Without index + # DataFrame[type, type, ...] + # DataFrame[name: type, name: type, ...] + # DataFrame[dtypes instance] + # DataFrame[zip(names, types)] + return _new_type_holders(params, NameTypeHolder) + + +def _new_type_holders( + params: Any, holder_clazz: Type[Union[NameTypeHolder, IndexNameTypeHolder]] +) -> Tuple: + if isinstance(params, zip): + # DataFrame[zip(names, types)] + params = tuple(slice(name, tpe) for name, tpe in params) # type: ignore[misc, has-type] + + if isinstance(params, Iterable): + # DataFrame[type, type, ...] + # DataFrame[name: type, name: type, ...] + # DataFrame[dtypes instance] + params = tuple(params) + else: + # DataFrame[type, type] + # DataFrame[name: type] + params = (params,) + + is_named_params = all( + isinstance(param, slice) and param.step is None and param.stop is not None + for param in params + ) + is_unnamed_params = all( + not isinstance(param, slice) and not isinstance(param, Iterable) for param in params + ) + + if is_named_params: + # DataFrame["id": int, "A": int] + new_params = [] + for param in params: + new_param: Type[Union[NameTypeHolder, IndexNameTypeHolder]] = type( + holder_clazz.short_name, (holder_clazz,), {} + ) + new_param.name = param.start + if isinstance(param.stop, ExtensionDtype): + new_param.tpe = param.stop # type: ignore[assignment] + else: + # When the given argument is a numpy's dtype instance. + new_param.tpe = param.stop.type if isinstance(param.stop, np.dtype) else param.stop + new_params.append(new_param) + return tuple(new_params) + elif is_unnamed_params: + # DataFrame[float, float] + new_types = [] + for param in params: + new_type: Type[Union[NameTypeHolder, IndexNameTypeHolder]] = type( + holder_clazz.short_name, (holder_clazz,), {} + ) + if isinstance(param, ExtensionDtype): + new_type.tpe = param # type: ignore[assignment] + else: + new_type.tpe = param.type if isinstance(param, np.dtype) else param + new_types.append(new_type) + return tuple(new_types) + else: + raise TypeError( + """Type hints should be specified as one of: + - DataFrame[type, type, ...] + - DataFrame[name: type, name: type, ...] + - DataFrame[dtypes instance] + - DataFrame[zip(names, types)] + - DataFrame[index_type, [type, ...]] + - DataFrame[(index_name, index_type), [(name, type), ...]] + - DataFrame[dtype instance, dtypes instance] + - DataFrame[(index_name, index_type), zip(names, types)] + - DataFrame[[index_type, ...], [type, ...]] + - DataFrame[[(index_name, index_type), ...], [(name, type), ...]] + - DataFrame[dtypes instance, dtypes instance] + - DataFrame[zip(index_names, index_types), zip(names, types)]\n""" + + "However, got %s." % str(params) + ) + + +def _test() -> None: + import doctest + import sys + import pyspark.pandas.typedef.typehints + + globs = pyspark.pandas.typedef.typehints.__dict__.copy() + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.typedef.typehints, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/__init__.py new file mode 100644 index 0000000..e14a905 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/__init__.py @@ -0,0 +1,143 @@ +# +# 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 types import ModuleType +from typing import Union + +import pandas as pd + +from pyspark.pandas import config, namespace, sql_formatter +from pyspark.pandas.accessors import PandasOnSparkFrameMethods +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.datetimes import DatetimeMethods +from pyspark.pandas.groupby import DataFrameGroupBy, SeriesGroupBy +from pyspark.pandas.indexes.base import Index +from pyspark.pandas.indexes.category import CategoricalIndex +from pyspark.pandas.indexes.datetimes import DatetimeIndex +from pyspark.pandas.indexes.multi import MultiIndex +from pyspark.pandas.indexes.numeric import Float64Index, Int64Index +from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame +from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions +from pyspark.pandas.missing.groupby import ( + MissingPandasLikeDataFrameGroupBy, + MissingPandasLikeSeriesGroupBy, +) +from pyspark.pandas.missing.indexes import ( + MissingPandasLikeDatetimeIndex, + MissingPandasLikeIndex, + MissingPandasLikeMultiIndex, +) +from pyspark.pandas.missing.series import MissingPandasLikeSeries +from pyspark.pandas.missing.window import ( + MissingPandasLikeExpanding, + MissingPandasLikeRolling, + MissingPandasLikeExpandingGroupby, + MissingPandasLikeRollingGroupby, + MissingPandasLikeExponentialMoving, + MissingPandasLikeExponentialMovingGroupby, +) +from pyspark.pandas.series import Series +from pyspark.pandas.spark.accessors import ( + CachedSparkFrameMethods, + SparkFrameMethods, + SparkIndexOpsMethods, +) +from pyspark.pandas.strings import StringMethods +from pyspark.pandas.window import ( + Expanding, + ExpandingGroupby, + Rolling, + RollingGroupby, + ExponentialMoving, + ExponentialMovingGroupby, +) +from pyspark.instrumentation_utils import _attach + + +def attach(logger_module: Union[str, ModuleType]) -> None: + """ + Attach the usage logger. + + Parameters + ---------- + logger_module : the module or module name contains the usage logger. + The module needs to provide `get_logger` function as an entry point of the plug-in + returning the usage logger. + + See Also + -------- + usage_logger : the reference implementation of the usage logger. + """ + + modules = [config, namespace] + classes = [ + DataFrame, + Series, + Index, + MultiIndex, + Int64Index, + Float64Index, + CategoricalIndex, + DatetimeIndex, + DataFrameGroupBy, + SeriesGroupBy, + DatetimeMethods, + StringMethods, + Expanding, + ExpandingGroupby, + Rolling, + RollingGroupby, + ExponentialMoving, + ExponentialMovingGroupby, + CachedSparkFrameMethods, + SparkFrameMethods, + SparkIndexOpsMethods, + PandasOnSparkFrameMethods, + ] + + try: + from pyspark.pandas import mlflow + + modules.append(mlflow) + classes.append(mlflow.PythonModelWrapper) + except ImportError: + pass + + sql_formatter._CAPTURE_SCOPES = 4 + modules.append(sql_formatter) + + missings = [ + (pd, MissingPandasLikeGeneralFunctions), + (pd.DataFrame, MissingPandasLikeDataFrame), + (pd.Series, MissingPandasLikeSeries), + (pd.Index, MissingPandasLikeIndex), + (pd.MultiIndex, MissingPandasLikeMultiIndex), + (pd.DatetimeIndex, MissingPandasLikeDatetimeIndex), + (pd.core.groupby.DataFrameGroupBy, MissingPandasLikeDataFrameGroupBy), + (pd.core.groupby.SeriesGroupBy, MissingPandasLikeSeriesGroupBy), + (pd.core.window.Expanding, MissingPandasLikeExpanding), + (pd.core.window.Rolling, MissingPandasLikeRolling), + (pd.core.window.ExpandingGroupby, MissingPandasLikeExpandingGroupby), + (pd.core.window.RollingGroupby, MissingPandasLikeRollingGroupby), + (pd.core.window.ExponentialMovingWindow, MissingPandasLikeExponentialMoving), + ( + pd.core.window.ExponentialMovingWindowGroupby, # type: ignore[attr-defined] + MissingPandasLikeExponentialMovingGroupby, + ), + ] + + _attach(logger_module, modules, classes, missings) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/usage_logger.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/usage_logger.py new file mode 100644 index 0000000..a17c52a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/usage_logging/usage_logger.py @@ -0,0 +1,132 @@ +# +# 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. +# + +""" +The reference implementation of usage logger using the Python standard logging library. +""" + +from inspect import Signature +import logging +from typing import Any, Optional + + +def get_logger() -> Any: + """An entry point of the plug-in and return the usage logger.""" + return PandasOnSparkUsageLogger() + + +def _format_signature(signature): + return ( + "({})".format(", ".join([p.name for p in signature.parameters.values()])) + if signature is not None + else "" + ) + + +class PandasOnSparkUsageLogger: + """ + The reference implementation of usage logger. + + The usage logger needs to provide the following methods: + + - log_success(self, class_name, name, duration, signature=None) + - log_failure(self, class_name, name, ex, duration, signature=None) + - log_missing(self, class_name, name, is_deprecated=False, signature=None) + """ + + def __init__(self): + self.logger = logging.getLogger("pyspark.pandas.usage_logger") + + def log_success( + self, class_name: str, name: str, duration: float, signature: Optional[Signature] = None + ) -> None: + """ + Log the function or property call is successfully finished. + + :param class_name: the target class name + :param name: the target function or property name + :param duration: the duration to finish the function or property call + :param signature: the signature if the target is a function, else None + """ + if self.logger.isEnabledFor(logging.INFO): + msg = ( + "A {function} `{class_name}.{name}{signature}` was successfully finished " + "after {duration:.3f} ms." + ).format( + class_name=class_name, + name=name, + signature=_format_signature(signature), + duration=duration * 1000, + function="function" if signature is not None else "property", + ) + self.logger.info(msg) + + def log_failure( + self, + class_name: str, + name: str, + ex: Exception, + duration: float, + signature: Optional[Signature] = None, + ) -> None: + """ + Log the function or property call failed. + + :param class_name: the target class name + :param name: the target function or property name + :param ex: the exception causing the failure + :param duration: the duration until the function or property call fails + :param signature: the signature if the target is a function, else None + """ + if self.logger.isEnabledFor(logging.WARNING): + msg = ( + "A {function} `{class_name}.{name}{signature}` was failed " + "after {duration:.3f} ms: {msg}" + ).format( + class_name=class_name, + name=name, + signature=_format_signature(signature), + msg=str(ex), + duration=duration * 1000, + function="function" if signature is not None else "property", + ) + self.logger.warning(msg) + + def log_missing( + self, + class_name: str, + name: str, + is_deprecated: bool = False, + signature: Optional[Signature] = None, + ) -> None: + """ + Log the missing or deprecated function or property is called. + + :param class_name: the target class name + :param name: the target function or property name + :param is_deprecated: True if the function or property is marked as deprecated + :param signature: the original function signature if the target is a function, else None + """ + if self.logger.isEnabledFor(logging.INFO): + msg = "A {deprecated} {function} `{class_name}.{name}{signature}` was called.".format( + class_name=class_name, + name=name, + signature=_format_signature(signature), + function="function" if signature is not None else "property", + deprecated="deprecated" if is_deprecated else "missing", + ) + self.logger.info(msg) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/utils.py new file mode 100644 index 0000000..ebeb1d6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/utils.py @@ -0,0 +1,1063 @@ +# +# 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. +# +""" +Commonly used utils in pandas-on-Spark. +""" + +import functools +from contextlib import contextmanager +import os +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Optional, + Tuple, + Union, + TYPE_CHECKING, + cast, + no_type_check, + overload, +) +import warnings + +from pyspark.sql import functions as F, Column, DataFrame as PySparkDataFrame, SparkSession +from pyspark.sql.types import DoubleType +from pyspark.sql.utils import is_remote, get_dataframe_class +from pyspark.errors import PySparkTypeError +import pandas as pd +from pandas.api.types import is_list_like # type: ignore[attr-defined] + +# For running doctests and reference resolution in PyCharm. +from pyspark import pandas as ps # noqa: F401 +from pyspark.pandas._typing import ( + Axis, + Label, + Name, + DataFrameOrSeries, +) +from pyspark.pandas.typedef.typehints import as_spark_type + + +if TYPE_CHECKING: + from pyspark.pandas.indexes.base import Index + from pyspark.pandas.base import IndexOpsMixin + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.internal import InternalFrame + from pyspark.pandas.series import Series + + +ERROR_MESSAGE_CANNOT_COMBINE = ( + "Cannot combine the series or dataframe because it comes from a different dataframe. " + "In order to allow this operation, enable 'compute.ops_on_diff_frames' option." +) + + +SPARK_CONF_ARROW_ENABLED = "spark.sql.execution.arrow.pyspark.enabled" + + +class PandasAPIOnSparkAdviceWarning(Warning): + pass + + +def same_anchor( + this: Union["DataFrame", "IndexOpsMixin", "InternalFrame"], + that: Union["DataFrame", "IndexOpsMixin", "InternalFrame"], +) -> bool: + """ + Check if the anchors of the given DataFrame or Series are the same or not. + """ + from pyspark.pandas.base import IndexOpsMixin + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.internal import InternalFrame + + if isinstance(this, InternalFrame): + this_internal = this + else: + assert isinstance(this, (DataFrame, IndexOpsMixin)), type(this) + this_internal = this._internal + + if isinstance(that, InternalFrame): + that_internal = that + else: + assert isinstance(that, (DataFrame, IndexOpsMixin)), type(that) + that_internal = that._internal + + return ( + this_internal.spark_frame is that_internal.spark_frame + and this_internal.index_level == that_internal.index_level + and all( + spark_column_equals(this_scol, that_scol) + for this_scol, that_scol in zip( + this_internal.index_spark_columns, that_internal.index_spark_columns + ) + ) + ) + + +def combine_frames( + this: "DataFrame", + *args: DataFrameOrSeries, + how: str = "full", + preserve_order_column: bool = False, +) -> "DataFrame": + """ + This method combines `this` DataFrame with a different `that` DataFrame or + Series from a different DataFrame. + + It returns a DataFrame that has prefix `this_` and `that_` to distinct + the columns names from both DataFrames + + It internally performs a join operation which can be expensive in general. + So, if `compute.ops_on_diff_frames` option is False, + this method throws an exception. + """ + from pyspark.pandas.config import get_option + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.internal import ( + InternalField, + InternalFrame, + HIDDEN_COLUMNS, + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, + ) + from pyspark.pandas.series import Series + + if all(isinstance(arg, Series) for arg in args): + assert all( + same_anchor(arg, args[0]) for arg in args + ), "Currently only one different DataFrame (from given Series) is supported" + assert not same_anchor(this, args[0]), "We don't need to combine. All series is in this." + that = args[0]._psdf[list(args)] + elif len(args) == 1 and isinstance(args[0], DataFrame): + assert isinstance(args[0], DataFrame) + assert not same_anchor( + this, args[0] + ), "We don't need to combine. `this` and `that` are same." + that = args[0] + else: + raise AssertionError("args should be single DataFrame or " "single/multiple Series") + + if get_option("compute.ops_on_diff_frames"): + + def resolve(internal: InternalFrame, side: str) -> InternalFrame: + def rename(col: str) -> str: + return "__{}_{}".format(side, col) + + internal = internal.resolved_copy + sdf = internal.spark_frame + sdf = internal.spark_frame.select( + *[ + scol_for(sdf, col).alias(rename(col)) + for col in sdf.columns + if col not in HIDDEN_COLUMNS + ], + *HIDDEN_COLUMNS, + ) + return internal.copy( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.index_spark_column_names + ], + index_fields=[ + field.copy(name=rename(field.name)) for field in internal.index_fields + ], + data_spark_columns=[ + scol_for(sdf, rename(col)) for col in internal.data_spark_column_names + ], + data_fields=[field.copy(name=rename(field.name)) for field in internal.data_fields], + ) + + this_internal = resolve(this._internal, "this") + that_internal = resolve(that._internal, "that") + + this_index_map = list( + zip( + this_internal.index_spark_column_names, + this_internal.index_names, + this_internal.index_fields, + ) + ) + that_index_map = list( + zip( + that_internal.index_spark_column_names, + that_internal.index_names, + that_internal.index_fields, + ) + ) + assert len(this_index_map) == len(that_index_map) + + join_scols = [] + merged_index_scols = [] + + # Note that the order of each element in index_map is guaranteed according to the index + # level. + this_and_that_index_map = list(zip(this_index_map, that_index_map)) + + this_sdf = this_internal.spark_frame.alias("this") + that_sdf = that_internal.spark_frame.alias("that") + + # If the same named index is found, that's used. + index_column_names = [] + index_use_extension_dtypes = [] + for ( + i, + ((this_column, this_name, this_field), (that_column, that_name, that_field)), + ) in enumerate(this_and_that_index_map): + if this_name == that_name: + # We should merge the Spark columns into one + # to mimic pandas' behavior. + this_scol = scol_for(this_sdf, this_column) + that_scol = scol_for(that_sdf, that_column) + join_scol = this_scol == that_scol + join_scols.append(join_scol) + + column_name = SPARK_INDEX_NAME_FORMAT(i) + index_column_names.append(column_name) + index_use_extension_dtypes.append( + any(field.is_extension_dtype for field in [this_field, that_field]) + ) + merged_index_scols.append( + F.when(this_scol.isNotNull(), this_scol).otherwise(that_scol).alias(column_name) + ) + else: + raise ValueError("Index names must be exactly matched currently.") + + assert len(join_scols) > 0, "cannot join with no overlapping index names" + + joined_df = this_sdf.join(that_sdf, on=join_scols, how=how) + + if preserve_order_column: + order_column = [scol_for(this_sdf, NATURAL_ORDER_COLUMN_NAME)] + else: + order_column = [] + + joined_df = joined_df.select( + *merged_index_scols, + *( + scol_for(this_sdf, this_internal.spark_column_name_for(label)) + for label in this_internal.column_labels + ), + *( + scol_for(that_sdf, that_internal.spark_column_name_for(label)) + for label in that_internal.column_labels + ), + *order_column, + ) + + index_spark_columns = [scol_for(joined_df, col) for col in index_column_names] + + index_columns = set(index_column_names) + new_data_columns = [ + col + for col in joined_df.columns + if col not in index_columns and col != NATURAL_ORDER_COLUMN_NAME + ] + + schema = joined_df.select(*index_spark_columns, *new_data_columns).schema + + index_fields = [ + InternalField.from_struct_field(struct_field, use_extension_dtypes=use_extension_dtypes) + for struct_field, use_extension_dtypes in zip( + schema.fields[: len(index_spark_columns)], index_use_extension_dtypes + ) + ] + data_fields = [ + InternalField.from_struct_field( + struct_field, use_extension_dtypes=field.is_extension_dtype + ) + for struct_field, field in zip( + schema.fields[len(index_spark_columns) :], + this_internal.data_fields + that_internal.data_fields, + ) + ] + + level = max(this_internal.column_labels_level, that_internal.column_labels_level) + + def fill_label(label: Optional[Label]) -> List: + if label is None: + return ([""] * (level - 1)) + [None] + else: + return ([""] * (level - len(label))) + list(label) + + column_labels = [ + tuple(["this"] + fill_label(label)) for label in this_internal.column_labels + ] + [tuple(["that"] + fill_label(label)) for label in that_internal.column_labels] + column_label_names = ( + cast(List[Optional[Label]], [None]) * (1 + level - this_internal.column_labels_level) + ) + this_internal.column_label_names + return DataFrame( + InternalFrame( + spark_frame=joined_df, + index_spark_columns=index_spark_columns, + index_names=this_internal.index_names, + index_fields=index_fields, + column_labels=column_labels, + data_spark_columns=[scol_for(joined_df, col) for col in new_data_columns], + data_fields=data_fields, + column_label_names=column_label_names, + ) + ) + else: + raise ValueError(ERROR_MESSAGE_CANNOT_COMBINE) + + +def align_diff_frames( + resolve_func: Callable[ + ["DataFrame", List[Label], List[Label]], Iterator[Tuple["Series", Label]] + ], + this: "DataFrame", + that: "DataFrame", + fillna: bool = True, + how: str = "full", + preserve_order_column: bool = False, +) -> "DataFrame": + """ + This method aligns two different DataFrames with a given `func`. Columns are resolved and + handled within the given `func`. + To use this, `compute.ops_on_diff_frames` should be True, for now. + + :param resolve_func: Takes aligned (joined) DataFrame, the column of the current DataFrame, and + the column of another DataFrame. It returns an iterable that produces Series. + + >>> from pyspark.pandas.config import set_option, reset_option + >>> + >>> set_option("compute.ops_on_diff_frames", True) + >>> + >>> psdf1 = ps.DataFrame({'a': [9, 8, 7, 6, 5, 4, 3, 2, 1]}) + >>> psdf2 = ps.DataFrame({'a': [9, 8, 7, 6, 5, 4, 3, 2, 1]}) + >>> + >>> def func(psdf, this_column_labels, that_column_labels): + ... psdf # conceptually this is A + B. + ... + ... # Within this function, Series from A or B can be performed against `psdf`. + ... this_label = this_column_labels[0] # this is ('a',) from psdf1. + ... that_label = that_column_labels[0] # this is ('a',) from psdf2. + ... new_series = (psdf[this_label] - psdf[that_label]).rename(str(this_label)) + ... + ... # This new series will be placed in new DataFrame. + ... yield (new_series, this_label) + >>> + >>> + >>> align_diff_frames(func, psdf1, psdf2).sort_index() + a + 0 0 + 1 0 + 2 0 + 3 0 + 4 0 + 5 0 + 6 0 + 7 0 + 8 0 + >>> reset_option("compute.ops_on_diff_frames") + + :param this: a DataFrame to align + :param that: another DataFrame to align + :param fillna: If True, it fills missing values in non-common columns in both `this` and `that`. + Otherwise, it returns as are. + :param how: join way. In addition, it affects how `resolve_func` resolves the column conflict. + - full: `resolve_func` should resolve only common columns from 'this' and 'that' DataFrames. + For instance, if 'this' has columns A, B, C and that has B, C, D, `this_columns` and + 'that_columns' in this function are B, C and B, C. + - left: `resolve_func` should resolve columns including `that` column. + For instance, if 'this' has columns A, B, C and that has B, C, D, `this_columns` is + B, C but `that_columns` are B, C, D. + - inner: Same as 'full' mode; however, internally performs inner join instead. + :return: Aligned DataFrame + """ + from pyspark.pandas.frame import DataFrame + + assert how == "full" or how == "left" or how == "inner" + + this_column_labels = this._internal.column_labels + that_column_labels = that._internal.column_labels + common_column_labels = set(this_column_labels).intersection(that_column_labels) + + # 1. Perform the join given two dataframes. + combined = combine_frames(this, that, how=how, preserve_order_column=preserve_order_column) + + # 2. Apply the given function to transform the columns in a batch and keep the new columns. + combined_column_labels = combined._internal.column_labels + + that_columns_to_apply: List[Label] = [] + this_columns_to_apply: List[Label] = [] + additional_that_columns: List[Label] = [] + if is_remote(): + from pyspark.sql.connect.column import Column as ConnectColumn + + Column = ConnectColumn + columns_to_keep: List[Union[Series, Column]] = [] # type: ignore[valid-type] + column_labels_to_keep: List[Label] = [] + + for combined_label in combined_column_labels: + for common_label in common_column_labels: + if combined_label == tuple(["this", *common_label]): + this_columns_to_apply.append(combined_label) + break + elif combined_label == tuple(["that", *common_label]): + that_columns_to_apply.append(combined_label) + break + else: + if how == "left" and combined_label in [ + tuple(["that", *label]) for label in that_column_labels + ]: + # In this case, we will drop `that_columns` in `columns_to_keep` but passes + # it later to `func`. `func` should resolve it. + # Note that adding this into a separate list (`additional_that_columns`) + # is intentional so that `this_columns` and `that_columns` can be paired. + additional_that_columns.append(combined_label) + elif fillna: + columns_to_keep.append(F.lit(None).cast(DoubleType()).alias(str(combined_label))) + column_labels_to_keep.append(combined_label) + else: + columns_to_keep.append(combined._psser_for(combined_label)) + column_labels_to_keep.append(combined_label) + + that_columns_to_apply += additional_that_columns + + # Should extract columns to apply and do it in a batch in case + # it adds new columns for example. + columns_applied: List[Union[Series, Column]] # type: ignore[valid-type] + column_labels_applied: List[Label] + if len(this_columns_to_apply) > 0 or len(that_columns_to_apply) > 0: + psser_set, column_labels_set = zip( + *resolve_func(combined, this_columns_to_apply, that_columns_to_apply) + ) + columns_applied = list(psser_set) + column_labels_applied = list(column_labels_set) + else: + columns_applied = [] + column_labels_applied = [] + + applied: DataFrame = DataFrame( + combined._internal.with_new_columns( + columns_applied + columns_to_keep, + column_labels=column_labels_applied + column_labels_to_keep, + ) + ) + + # 3. Restore the names back and deduplicate columns. + this_labels: Dict[Label, Label] = {} + # Add columns in an order of its original frame. + for this_label in this_column_labels: + for new_label in applied._internal.column_labels: + if new_label[1:] not in this_labels and this_label == new_label[1:]: + this_labels[new_label[1:]] = new_label + + # After that, we will add the rest columns. + other_labels: Dict[Label, Label] = {} + for new_label in applied._internal.column_labels: + if new_label[1:] not in this_labels: + other_labels[new_label[1:]] = new_label + + psdf = applied[list(this_labels.values()) + list(other_labels.values())] + psdf.columns = psdf.columns.droplevel() + return psdf + + +def is_testing() -> bool: + """Indicates whether Spark is currently running tests.""" + return "SPARK_TESTING" in os.environ + + +def default_session() -> SparkSession: + spark = SparkSession.getActiveSession() + if spark is None: + spark = SparkSession.builder.appName("pandas-on-Spark").getOrCreate() + + # Turn ANSI off when testing the pandas API on Spark since + # the behavior of pandas API on Spark follows pandas, not SQL. + if is_testing(): + spark.conf.set("spark.sql.ansi.enabled", False) + if spark.conf.get("spark.sql.ansi.enabled") == "true": + log_advice( + "The config 'spark.sql.ansi.enabled' is set to True. " + "This can cause unexpected behavior " + "from pandas API on Spark since pandas API on Spark follows " + "the behavior of pandas, not SQL." + ) + + return spark + + +@contextmanager +def sql_conf(pairs: Dict[str, Any], *, spark: Optional[SparkSession] = None) -> Iterator[None]: + """ + A convenient context manager to set `value` to the Spark SQL configuration `key` and + then restores it back when it exits. + """ + assert isinstance(pairs, dict), "pairs should be a dictionary." + + if spark is None: + spark = default_session() + + keys = pairs.keys() + new_values = pairs.values() + old_values = [spark.conf.get(key, None) for key in keys] + for key, new_value in zip(keys, new_values): + spark.conf.set(key, new_value) + try: + yield + finally: + for key, old_value in zip(keys, old_values): + if old_value is None: + spark.conf.unset(key) + else: + spark.conf.set(key, old_value) + + +def validate_arguments_and_invoke_function( + pobj: Union[pd.DataFrame, pd.Series], + pandas_on_spark_func: Callable, + pandas_func: Callable, + input_args: Dict, +) -> Any: + """ + Invokes a pandas function. + + This is created because different versions of pandas support different parameters, and as a + result when we code against the latest version, our users might get a confusing + "got an unexpected keyword argument" error if they are using an older version of pandas. + + This function validates all the arguments, removes the ones that are not supported if they + are simply the default value (i.e. most likely the user didn't explicitly specify it). It + throws a TypeError if the user explicitly specifies an argument that is not supported by the + pandas version available. + + For example usage, look at DataFrame.to_html(). + + :param pobj: the pandas DataFrame or Series to operate on + :param pandas_on_spark_func: pandas-on-Spark function, used to get default parameter values + :param pandas_func: pandas function, used to check whether pandas supports all the arguments + :param input_args: arguments to pass to the pandas function, often created by using locals(). + Make sure locals() call is at the top of the function so it captures only + input parameters, rather than local variables. + :return: whatever pandas_func returns + """ + import inspect + + # Makes a copy since whatever passed in is likely created by locals(), and we can't delete + # 'self' key from that. + args = input_args.copy() + del args["self"] + + if "kwargs" in args: + # explode kwargs + kwargs = args["kwargs"] + del args["kwargs"] + args = {**args, **kwargs} + + pandas_on_spark_params = inspect.signature(pandas_on_spark_func).parameters + pandas_params = inspect.signature(pandas_func).parameters + + for param in pandas_on_spark_params.values(): + if param.name not in pandas_params: + if args[param.name] == param.default: + del args[param.name] + else: + raise TypeError( + ( + "The pandas version [%s] available does not support parameter '%s' " + + "for function '%s'." + ) + % (pd.__version__, param.name, pandas_func.__name__) + ) + + args["self"] = pobj + return pandas_func(**args) + + +@no_type_check +def lazy_property(fn: Callable[[Any], Any]) -> property: + """ + Decorator that makes a property lazy-evaluated. + + Copied from https://stevenloria.com/lazy-properties/ + """ + attr_name = "_lazy_" + fn.__name__ + + @property + @functools.wraps(fn) + def wrapped_lazy_property(self): + if not hasattr(self, attr_name): + setattr(self, attr_name, fn(self)) + return getattr(self, attr_name) + + def deleter(self): + if hasattr(self, attr_name): + delattr(self, attr_name) + + return wrapped_lazy_property.deleter(deleter) + + +def scol_for(sdf: PySparkDataFrame, column_name: str) -> Column: + """Return Spark Column for the given column name.""" + return sdf["`{}`".format(column_name)] + + +def column_labels_level(column_labels: List[Label]) -> int: + """Return the level of the column index.""" + if len(column_labels) == 0: + return 1 + else: + levels = set(1 if label is None else len(label) for label in column_labels) + assert len(levels) == 1, levels + return list(levels)[0] + + +def name_like_string(name: Optional[Name]) -> str: + """ + Return the name-like strings from str or tuple of str + + Examples + -------- + >>> name = 'abc' + >>> name_like_string(name) + 'abc' + + >>> name = ('abc',) + >>> name_like_string(name) + 'abc' + + >>> name = ('a', 'b', 'c') + >>> name_like_string(name) + '(a, b, c)' + """ + label: Label + if name is None: + label = ("__none__",) + elif is_list_like(name): + label = tuple([str(n) for n in name]) + else: + label = (str(name),) + return ("(%s)" % ", ".join(label)) if len(label) > 1 else label[0] + + +def is_name_like_tuple(value: Any, allow_none: bool = True, check_type: bool = False) -> bool: + """ + Check the given tuple is to be able to be used as a name. + + Examples + -------- + >>> is_name_like_tuple(('abc',)) + True + >>> is_name_like_tuple((1,)) + True + >>> is_name_like_tuple(('abc', 1, None)) + True + >>> is_name_like_tuple(('abc', 1, None), check_type=True) + True + >>> is_name_like_tuple((1.0j,)) + True + >>> is_name_like_tuple(tuple()) + False + >>> is_name_like_tuple((list('abc'),)) + False + >>> is_name_like_tuple(('abc', 1, None), allow_none=False) + False + >>> is_name_like_tuple((1.0j,), check_type=True) + False + """ + if value is None: + return allow_none + elif not isinstance(value, tuple): + return False + elif len(value) == 0: + return False + elif not allow_none and any(v is None for v in value): + return False + elif any(is_list_like(v) or isinstance(v, slice) for v in value): + return False + elif check_type: + return all( + v is None or as_spark_type(type(v), raise_error=False) is not None for v in value + ) + else: + return True + + +def is_name_like_value( + value: Any, allow_none: bool = True, allow_tuple: bool = True, check_type: bool = False +) -> bool: + """ + Check the given value is like a name. + + Examples + -------- + >>> is_name_like_value('abc') + True + >>> is_name_like_value(1) + True + >>> is_name_like_value(None) + True + >>> is_name_like_value(('abc',)) + True + >>> is_name_like_value(1.0j) + True + >>> is_name_like_value(list('abc')) + False + >>> is_name_like_value(None, allow_none=False) + False + >>> is_name_like_value(('abc',), allow_tuple=False) + False + >>> is_name_like_value(1.0j, check_type=True) + False + """ + if value is None: + return allow_none + elif isinstance(value, tuple): + return allow_tuple and is_name_like_tuple( + value, allow_none=allow_none, check_type=check_type + ) + elif is_list_like(value) or isinstance(value, slice): + return False + elif check_type: + return as_spark_type(type(value), raise_error=False) is not None + else: + return True + + +def validate_axis(axis: Optional[Axis] = 0, none_axis: int = 0) -> int: + """Check the given axis is valid.""" + # convert to numeric axis + axis = cast(Dict[Optional[Axis], int], {None: none_axis, "index": 0, "columns": 1}).get( + axis, axis + ) + if axis in (none_axis, 0, 1): + return cast(int, axis) + else: + raise ValueError("No axis named {0}".format(axis)) + + +def validate_bool_kwarg(value: Any, arg_name: str) -> Optional[bool]: + """Ensures that argument passed in arg_name is of type bool.""" + if not (isinstance(value, bool) or value is None): + raise TypeError( + 'For argument "{}" expected type bool, received ' + "type {}.".format(arg_name, type(value).__name__) + ) + return value + + +def validate_how(how: str) -> str: + """Check the given how for join is valid.""" + if how == "full": + warnings.warn( + "Warning: While pandas-on-Spark will accept 'full', you should use 'outer' " + + "instead to be compatible with the pandas merge API", + UserWarning, + ) + if how == "outer": + # 'outer' in pandas equals 'full' in Spark + how = "full" + if how not in ("inner", "left", "right", "full"): + raise ValueError( + "The 'how' parameter has to be amongst the following values: ", + "['inner', 'left', 'right', 'outer']", + ) + return how + + +def validate_mode(mode: str) -> str: + """Check the given mode for writing is valid.""" + if mode in ("w", "w+"): + # 'w' in pandas equals 'overwrite' in Spark + # '+' is meaningless for writing methods, but pandas just pass it as 'w'. + mode = "overwrite" + if mode in ("a", "a+"): + # 'a' in pandas equals 'append' in Spark + # '+' is meaningless for writing methods, but pandas just pass it as 'a'. + mode = "append" + if mode not in ( + "w", + "a", + "w+", + "a+", + "overwrite", + "append", + "ignore", + "error", + "errorifexists", + ): + raise ValueError( + "The 'mode' parameter has to be amongst the following values: ", + "['w', 'a', 'w+', 'a+', 'overwrite', 'append', 'ignore', 'error', 'errorifexists']", + ) + return mode + + +@overload +def verify_temp_column_name(df: PySparkDataFrame, column_name_or_label: str) -> str: + ... + + +@overload +def verify_temp_column_name(df: "DataFrame", column_name_or_label: Name) -> Label: + ... + + +def verify_temp_column_name( + df: Union["DataFrame", PySparkDataFrame], + column_name_or_label: Union[str, Name], +) -> Union[str, Label]: + """ + Verify that the given column name does not exist in the given pandas-on-Spark or + Spark DataFrame. + + The temporary column names should start and end with `__`. In addition, `column_name_or_label` + expects a single string, or column labels when `df` is a pandas-on-Spark DataFrame. + + >>> psdf = ps.DataFrame({("x", "a"): ['a', 'b', 'c']}) + >>> psdf["__dummy__"] = 0 + >>> psdf[("", "__dummy__")] = 1 + >>> psdf # doctest: +NORMALIZE_WHITESPACE + x __dummy__ + a __dummy__ + 0 a 0 1 + 1 b 0 1 + 2 c 0 1 + + >>> verify_temp_column_name(psdf, '__tmp__') + ('__tmp__', '') + >>> verify_temp_column_name(psdf, ('', '__tmp__')) + ('', '__tmp__') + >>> verify_temp_column_name(psdf, '__dummy__') + Traceback (most recent call last): + ... + AssertionError: ... `(__dummy__, )` ... + >>> verify_temp_column_name(psdf, ('', '__dummy__')) + Traceback (most recent call last): + ... + AssertionError: ... `(, __dummy__)` ... + >>> verify_temp_column_name(psdf, 'dummy') + Traceback (most recent call last): + ... + AssertionError: ... should be empty or start and end with `__`: ('dummy', '') + >>> verify_temp_column_name(psdf, ('', 'dummy')) + Traceback (most recent call last): + ... + AssertionError: ... should be empty or start and end with `__`: ('', 'dummy') + + >>> internal = psdf._internal.resolved_copy + >>> sdf = internal.spark_frame + >>> sdf.select(internal.data_spark_columns).show() # doctest: +NORMALIZE_WHITESPACE + +------+---------+-------------+ + |(x, a)|__dummy__|(, __dummy__)| + +------+---------+-------------+ + | a| 0| 1| + | b| 0| 1| + | c| 0| 1| + +------+---------+-------------+ + + >>> verify_temp_column_name(sdf, '__tmp__') + '__tmp__' + >>> verify_temp_column_name(sdf, '__dummy__') + Traceback (most recent call last): + ... + AssertionError: ... `__dummy__` ... '(x, a)', '__dummy__', '(, __dummy__)', ... + >>> verify_temp_column_name(sdf, ('', '__dummy__')) + Traceback (most recent call last): + ... + AssertionError: + >>> verify_temp_column_name(sdf, 'dummy') + Traceback (most recent call last): + ... + AssertionError: ... should start and end with `__`: dummy + """ + from pyspark.pandas.frame import DataFrame + + if isinstance(df, DataFrame): + if isinstance(column_name_or_label, str): + column_name = column_name_or_label + + level = df._internal.column_labels_level + column_name_or_label = tuple([column_name_or_label] + ([""] * (level - 1))) + else: + column_name = name_like_string(column_name_or_label) + + assert any(len(label) > 0 for label in column_name_or_label) and all( + label == "" or (label.startswith("__") and label.endswith("__")) + for label in column_name_or_label + ), "The temporary column name should be empty or start and end with `__`: {}".format( + column_name_or_label + ) + assert all( + column_name_or_label != label for label in df._internal.column_labels + ), "The given column name `{}` already exists in the pandas-on-Spark DataFrame: {}".format( + name_like_string(column_name_or_label), df.columns + ) + df = df._internal.resolved_copy.spark_frame + else: + assert isinstance(column_name_or_label, str), type(column_name_or_label) + assert column_name_or_label.startswith("__") and column_name_or_label.endswith( + "__" + ), "The temporary column name should start and end with `__`: {}".format( + column_name_or_label + ) + column_name = column_name_or_label + + SparkDataFrame = get_dataframe_class() + assert isinstance(df, SparkDataFrame), type(df) + assert ( + column_name not in df.columns + ), "The given column name `{}` already exists in the Spark DataFrame: {}".format( + column_name, df.columns + ) + + return column_name_or_label + + +def spark_column_equals(left: Column, right: Column) -> bool: + """ + Check both `left` and `right` have the same expressions. + + >>> spark_column_equals(sf.lit(0), sf.lit(0)) + True + >>> spark_column_equals(sf.lit(0) + 1, sf.lit(0) + 1) + True + >>> spark_column_equals(sf.lit(0) + 1, sf.lit(0) + 2) + False + >>> sdf1 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark() + >>> spark_column_equals(sdf1["x"] + 1, sdf1["x"] + 1) + True + >>> sdf2 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark() + >>> spark_column_equals(sdf1["x"] + 1, sdf2["x"] + 1) + False + """ + if is_remote(): + from pyspark.sql.connect.column import Column as ConnectColumn + + if not isinstance(left, ConnectColumn): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "left", "arg_type": type(left).__name__}, + ) + if not isinstance(right, ConnectColumn): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "right", "arg_type": type(right).__name__}, + ) + return repr(left) == repr(right) + else: + return left._jc.equals(right._jc) + + +def compare_null_first( + left: Column, + right: Column, + comp: Callable[ + [Column, Column], + Column, + ], +) -> Column: + return (left.isNotNull() & right.isNotNull() & comp(left, right)) | ( + left.isNull() & right.isNotNull() + ) + + +def compare_null_last( + left: Column, + right: Column, + comp: Callable[ + [Column, Column], + Column, + ], +) -> Column: + return (left.isNotNull() & right.isNotNull() & comp(left, right)) | ( + left.isNotNull() & right.isNull() + ) + + +def compare_disallow_null( + left: Column, + right: Column, + comp: Callable[ + [Column, Column], + Column, + ], +) -> Column: + return left.isNotNull() & right.isNotNull() & comp(left, right) + + +def compare_allow_null( + left: Column, + right: Column, + comp: Callable[ + [Column, Column], + Column, + ], +) -> Column: + return left.isNull() | right.isNull() | comp(left, right) + + +def log_advice(message: str) -> None: + """ + Display advisory logs for functions to be aware of when using pandas API on Spark + for the existing pandas/PySpark users who may not be familiar with distributed environments + or the behavior of pandas. + """ + warnings.warn(message, PandasAPIOnSparkAdviceWarning) + + +def validate_index_loc(index: "Index", loc: int) -> None: + """ + Raises IndexError if index is out of bounds + """ + length = len(index) + if loc < 0: + loc = loc + length + if loc < 0: + raise IndexError( + "index {} is out of bounds for axis 0 with size {}".format((loc - length), length) + ) + else: + if loc > length: + raise IndexError( + "index {} is out of bounds for axis 0 with size {}".format(loc, length) + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.utils + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.utils.__dict__.copy() + globs["ps"] = pyspark.pandas + globs["sf"] = F + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.utils tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.utils, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/window.py new file mode 100644 index 0000000..cc793f0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/pandas/window.py @@ -0,0 +1,2702 @@ +# +# 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 functools import partial +from typing import Any, Callable, Generic, List, Optional + +import numpy as np + +from pyspark.sql import Window +from pyspark.sql import functions as F +from pyspark.pandas.missing.window import ( + MissingPandasLikeRolling, + MissingPandasLikeRollingGroupby, + MissingPandasLikeExpanding, + MissingPandasLikeExpandingGroupby, + MissingPandasLikeExponentialMoving, + MissingPandasLikeExponentialMovingGroupby, +) + +# For running doctests and reference resolution in PyCharm. +from pyspark import pandas as ps # noqa: F401 +from pyspark.pandas._typing import FrameLike +from pyspark.pandas.groupby import GroupBy, DataFrameGroupBy +from pyspark.pandas.internal import NATURAL_ORDER_COLUMN_NAME, SPARK_INDEX_NAME_FORMAT +from pyspark.pandas.spark import functions as SF +from pyspark.pandas.utils import scol_for +from pyspark.sql.column import Column +from pyspark.sql.types import ( + DoubleType, +) +from pyspark.sql.window import WindowSpec + + +class RollingAndExpanding(Generic[FrameLike], metaclass=ABCMeta): + def __init__(self, window: WindowSpec, min_periods: int): + self._window = window + # This unbounded Window is later used to handle 'min_periods' for now. + self._unbounded_window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + self._min_periods = min_periods + + @abstractmethod + def _apply_as_series_or_frame(self, func: Callable[[Column], Column]) -> FrameLike: + """ + Wraps a function that handles Spark column in order + to support it in both pandas-on-Spark Series and DataFrame. + Note that the given `func` name should be same as the API's method name. + """ + pass + + @abstractmethod + def count(self) -> FrameLike: + pass + + def sum(self) -> FrameLike: + def sum(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.sum(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(sum) + + def min(self) -> FrameLike: + def min(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.min(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(min) + + def max(self) -> FrameLike: + def max(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.max(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(max) + + def mean(self) -> FrameLike: + def mean(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.mean(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(mean) + + def quantile(self, q: float, accuracy: int = 10000) -> FrameLike: + def quantile(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.percentile_approx(scol.cast(DoubleType()), q, accuracy).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(quantile) + + def std(self) -> FrameLike: + def std(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.stddev(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(std) + + def var(self) -> FrameLike: + def var(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.variance(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(var) + + def skew(self) -> FrameLike: + def skew(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + SF.skew(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(skew) + + def kurt(self) -> FrameLike: + def kurt(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + SF.kurt(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(kurt) + + +class RollingLike(RollingAndExpanding[FrameLike]): + def __init__( + self, + window: int, + min_periods: Optional[int] = None, + ): + if window < 0: + raise ValueError("window must be >= 0") + if (min_periods is not None) and (min_periods < 0): + raise ValueError("min_periods must be >= 0") + if min_periods is None: + # TODO: 'min_periods' is not equivalent in pandas because it does not count NA as + # a value. + min_periods = window + + window_spec = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.currentRow - (window - 1), Window.currentRow + ) + + super().__init__(window_spec, min_periods) + + def count(self) -> FrameLike: + def count(scol: Column) -> Column: + return F.count(scol).over(self._window) + + return self._apply_as_series_or_frame(count).astype("float64") # type: ignore[attr-defined] + + +class Rolling(RollingLike[FrameLike]): + def __init__( + self, + psdf_or_psser: FrameLike, + window: int, + min_periods: Optional[int] = None, + ): + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series + + super().__init__(window, min_periods) + + self._psdf_or_psser = psdf_or_psser + + if not isinstance(psdf_or_psser, (DataFrame, Series)): + raise TypeError( + "psdf_or_psser must be a series or dataframe; however, got: %s" + % type(psdf_or_psser) + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeRolling, item): + property_or_func = getattr(MissingPandasLikeRolling, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + def _apply_as_series_or_frame(self, func: Callable[[Column], Column]) -> FrameLike: + return self._psdf_or_psser._apply_series_op( + lambda psser: psser._with_new_scol(func(psser.spark.column)), # TODO: dtype? + should_resolve=True, + ) + + def count(self) -> FrameLike: + """ + The rolling count of any non-NaN observations inside the window. + + .. 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 + Return type is the same as the original object with `np.float64` dtype. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.count : Count of the full Series. + pyspark.pandas.DataFrame.count : Count of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 3, float("nan"), 10]) + >>> s.rolling(1).count() + 0 1.0 + 1 1.0 + 2 0.0 + 3 1.0 + dtype: float64 + + >>> s.rolling(3).count() + 0 1.0 + 1 2.0 + 2 2.0 + 3 2.0 + dtype: float64 + + >>> s.to_frame().rolling(1).count() + 0 + 0 1.0 + 1 1.0 + 2 0.0 + 3 1.0 + + >>> s.to_frame().rolling(3).count() + 0 + 0 1.0 + 1 2.0 + 2 2.0 + 3 2.0 + """ + return super().count() + + def sum(self) -> FrameLike: + """ + Calculate rolling summation of given DataFrame or Series. + + .. 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 + Same type as the input, with the same index, containing the + rolling summation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.sum : Reducing sum for Series. + pyspark.pandas.DataFrame.sum : Reducing sum for DataFrame. + + 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).sum() + 0 NaN + 1 7.0 + 2 8.0 + 3 7.0 + 4 8.0 + dtype: float64 + + >>> s.rolling(3).sum() + 0 NaN + 1 NaN + 2 12.0 + 3 10.0 + 4 13.0 + dtype: float64 + + For DataFrame, each rolling summation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 4 16 + 1 3 9 + 2 5 25 + 3 2 4 + 4 6 36 + + >>> df.rolling(2).sum() + A B + 0 NaN NaN + 1 7.0 25.0 + 2 8.0 34.0 + 3 7.0 29.0 + 4 8.0 40.0 + + >>> df.rolling(3).sum() + A B + 0 NaN NaN + 1 NaN NaN + 2 12.0 50.0 + 3 10.0 38.0 + 4 13.0 65.0 + """ + return super().sum() + + def min(self) -> FrameLike: + """ + Calculate the rolling minimum. + + .. 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 a Series. + pyspark.pandas.DataFrame.rolling : Calling object with a DataFrame. + pyspark.pandas.Series.min : Similar method for Series. + pyspark.pandas.DataFrame.min : Similar method for DataFrame. + + 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).min() + 0 NaN + 1 3.0 + 2 3.0 + 3 2.0 + 4 2.0 + dtype: float64 + + >>> s.rolling(3).min() + 0 NaN + 1 NaN + 2 3.0 + 3 2.0 + 4 2.0 + dtype: float64 + + For DataFrame, each rolling minimum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 4 16 + 1 3 9 + 2 5 25 + 3 2 4 + 4 6 36 + + >>> df.rolling(2).min() + A B + 0 NaN NaN + 1 3.0 9.0 + 2 3.0 9.0 + 3 2.0 4.0 + 4 2.0 4.0 + + >>> df.rolling(3).min() + A B + 0 NaN NaN + 1 NaN NaN + 2 3.0 9.0 + 3 2.0 4.0 + 4 2.0 4.0 + """ + return super().min() + + def max(self) -> FrameLike: + """ + Calculate the rolling maximum. + + .. 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 + Return type is determined by the caller. + + See Also + -------- + pyspark.pandas.Series.rolling : Series rolling. + pyspark.pandas.DataFrame.rolling : DataFrame rolling. + pyspark.pandas.Series.max : Similar method for Series. + pyspark.pandas.DataFrame.max : Similar method for DataFrame. + + 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).max() + 0 NaN + 1 4.0 + 2 5.0 + 3 5.0 + 4 6.0 + dtype: float64 + + >>> s.rolling(3).max() + 0 NaN + 1 NaN + 2 5.0 + 3 5.0 + 4 6.0 + dtype: float64 + + For DataFrame, each rolling maximum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 4 16 + 1 3 9 + 2 5 25 + 3 2 4 + 4 6 36 + + >>> df.rolling(2).max() + A B + 0 NaN NaN + 1 4.0 16.0 + 2 5.0 25.0 + 3 5.0 25.0 + 4 6.0 36.0 + + >>> df.rolling(3).max() + A B + 0 NaN NaN + 1 NaN NaN + 2 5.0 25.0 + 3 5.0 25.0 + 4 6.0 36.0 + """ + return super().max() + + def mean(self) -> FrameLike: + """ + Calculate the rolling mean of the values. + + .. 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.mean : Equivalent method for Series. + pyspark.pandas.DataFrame.mean : Equivalent method for DataFrame. + + 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).mean() + 0 NaN + 1 3.5 + 2 4.0 + 3 3.5 + 4 4.0 + dtype: float64 + + >>> s.rolling(3).mean() + 0 NaN + 1 NaN + 2 4.000000 + 3 3.333333 + 4 4.333333 + dtype: float64 + + For DataFrame, each rolling mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 4 16 + 1 3 9 + 2 5 25 + 3 2 4 + 4 6 36 + + >>> df.rolling(2).mean() + A B + 0 NaN NaN + 1 3.5 12.5 + 2 4.0 17.0 + 3 3.5 14.5 + 4 4.0 20.0 + + >>> df.rolling(3).mean() + A B + 0 NaN NaN + 1 NaN NaN + 2 4.000000 16.666667 + 3 3.333333 12.666667 + 4 4.333333 21.666667 + """ + return super().mean() + + def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: + """ + Calculate the rolling quantile of the values. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + quantile : float + Value between 0 and 1 providing the quantile to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + This is a panda-on-Spark specific parameter. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the rolling + calculation. + + Notes + ----- + `quantile` in pandas-on-Spark are using distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas, also `interpolation` + parameter is not supported yet. + + 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. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling rolling with Series data. + pyspark.pandas.DataFrame.rolling : Calling rolling with DataFrames. + pyspark.pandas.Series.quantile : Aggregating quantile for Series. + pyspark.pandas.DataFrame.quantile : Aggregating quantile for DataFrame. + + 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).quantile(0.5) + 0 NaN + 1 3.0 + 2 3.0 + 3 2.0 + 4 2.0 + dtype: float64 + + >>> s.rolling(3).quantile(0.5) + 0 NaN + 1 NaN + 2 4.0 + 3 3.0 + 4 5.0 + dtype: float64 + + For DataFrame, each rolling quantile is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 4 16 + 1 3 9 + 2 5 25 + 3 2 4 + 4 6 36 + + >>> df.rolling(2).quantile(0.5) + A B + 0 NaN NaN + 1 3.0 9.0 + 2 3.0 9.0 + 3 2.0 4.0 + 4 2.0 4.0 + + >>> df.rolling(3).quantile(0.5) + A B + 0 NaN NaN + 1 NaN NaN + 2 4.0 16.0 + 3 3.0 9.0 + 4 5.0 25.0 + """ + return super().quantile(quantile, accuracy) + + def std(self) -> FrameLike: + """ + Calculate rolling standard deviation. + + .. 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 + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.rolling(3).std() + 0 NaN + 1 NaN + 2 0.577350 + 3 1.000000 + 4 1.000000 + 5 1.154701 + 6 0.000000 + dtype: float64 + + For DataFrame, each rolling standard deviation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.rolling(2).std() + 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().std() + + def var(self) -> FrameLike: + """ + Calculate unbiased rolling variance. + + .. 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 + Returns the same object type as the caller of the rolling calculation. + + See Also + -------- + Series.rolling : Calling object with Series data. + DataFrame.rolling : Calling object with DataFrames. + Series.var : Equivalent method for Series. + DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.rolling(3).var() + 0 NaN + 1 NaN + 2 0.333333 + 3 1.000000 + 4 1.000000 + 5 1.333333 + 6 0.000000 + dtype: float64 + + For DataFrame, each unbiased rolling variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.rolling(2).var() + A B + 0 NaN NaN + 1 0.0 0.0 + 2 0.5 60.5 + 3 0.5 84.5 + 4 2.0 288.0 + 5 0.0 0.0 + 6 0.0 0.0 + """ + return super().var() + + def skew(self) -> FrameLike: + """ + Calculate unbiased rolling skew. + + .. 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 + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 1, 5, 9]) + >>> s.rolling(3).skew() + 0 NaN + 1 NaN + 2 1.732051 + 3 0.000000 + 4 0.000000 + 5 -0.935220 + 6 -1.732051 + 7 0.000000 + dtype: float64 + + For DataFrame, each rolling standard deviation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.rolling(5).skew() + A B + 0 NaN NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 1.257788 1.369456 + 5 -1.492685 -0.526039 + 6 -1.492685 -0.526039 + 7 -0.551618 0.686072 + """ + return super().skew() + + def kurt(self) -> FrameLike: + """ + Calculate unbiased rolling kurtosis. + + .. 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 + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 1, 5, 9]) + >>> s.rolling(4).kurt() + 0 NaN + 1 NaN + 2 NaN + 3 -1.289256 + 4 -1.289256 + 5 2.234867 + 6 2.227147 + 7 1.500000 + dtype: float64 + + For DataFrame, each unbiased rolling variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.rolling(5).kurt() + A B + 0 NaN NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 0.312500 0.906336 + 5 2.818047 1.016942 + 6 2.818047 1.016942 + 7 0.867769 0.389750 + """ + return super().kurt() + + +class RollingGroupby(RollingLike[FrameLike]): + def __init__( + self, + groupby: GroupBy[FrameLike], + window: int, + min_periods: Optional[int] = None, + ): + super().__init__(window, min_periods) + + self._groupby = groupby + self._window = self._window.partitionBy(*[ser.spark.column for ser in groupby._groupkeys]) + self._unbounded_window = self._unbounded_window.partitionBy( + *[ser.spark.column for ser in groupby._groupkeys] + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeRollingGroupby, item): + property_or_func = getattr(MissingPandasLikeRollingGroupby, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + def _apply_as_series_or_frame(self, func: Callable[[Column], Column]) -> FrameLike: + """ + Wraps a function that handles Spark column in order + to support it in both pandas-on-Spark Series and DataFrame. + Note that the given `func` name should be same as the API's method name. + """ + from pyspark.pandas import DataFrame + + groupby = self._groupby + psdf = groupby._psdf + + # Here we need to include grouped key as an index, and shift previous index. + # [index_column0, index_column1] -> [grouped key, index_column0, index_column1] + new_index_scols: List[Column] = [] + new_index_spark_column_names = [] + new_index_names = [] + new_index_fields = [] + for groupkey in groupby._groupkeys: + index_column_name = SPARK_INDEX_NAME_FORMAT(len(new_index_scols)) + new_index_scols.append(groupkey.spark.column.alias(index_column_name)) + new_index_spark_column_names.append(index_column_name) + new_index_names.append(groupkey._column_label) + new_index_fields.append(groupkey._internal.data_fields[0].copy(name=index_column_name)) + + for new_index_scol, index_name, index_field in zip( + psdf._internal.index_spark_columns, + psdf._internal.index_names, + psdf._internal.index_fields, + ): + index_column_name = SPARK_INDEX_NAME_FORMAT(len(new_index_scols)) + new_index_scols.append(new_index_scol.alias(index_column_name)) + new_index_spark_column_names.append(index_column_name) + new_index_names.append(index_name) + new_index_fields.append(index_field.copy(name=index_column_name)) + + if groupby._agg_columns_selected: + agg_columns = groupby._agg_columns + else: + # pandas doesn't keep the groupkey as a column from 1.3 for DataFrameGroupBy + column_labels_to_exclude = groupby._column_labels_to_exclude.copy() + if isinstance(groupby, DataFrameGroupBy): + for groupkey in groupby._groupkeys: # type: ignore[attr-defined] + column_labels_to_exclude.add(groupkey._internal.column_labels[0]) + agg_columns = [ + psdf._psser_for(label) + for label in psdf._internal.column_labels + if label not in column_labels_to_exclude + ] + + applied = [] + for agg_column in agg_columns: + applied.append(agg_column._with_new_scol(func(agg_column.spark.column))) # TODO: dtype? + + # Seems like pandas filters out when grouped key is NA. + cond = groupby._groupkeys[0].spark.column.isNotNull() + for c in groupby._groupkeys[1:]: + cond = cond | c.spark.column.isNotNull() + + sdf = psdf._internal.spark_frame.filter(cond).select( + new_index_scols + [c.spark.column for c in applied] + ) + + internal = psdf._internal.copy( + spark_frame=sdf, + index_spark_columns=[scol_for(sdf, col) for col in new_index_spark_column_names], + index_names=new_index_names, + index_fields=new_index_fields, + column_labels=[c._column_label for c in applied], + data_spark_columns=[ + scol_for(sdf, c._internal.data_spark_column_names[0]) for c in applied + ], + data_fields=[c._internal.data_fields[0] for c in applied], + ) + + return groupby._handle_output(DataFrame(internal)) + + def count(self) -> FrameLike: + """ + The rolling count of any non-NaN observations inside the window. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the expanding + calculation. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling object with Series data. + pyspark.pandas.DataFrame.rolling : Calling object with DataFrames. + pyspark.pandas.Series.count : Count of the full Series. + pyspark.pandas.DataFrame.count : Count of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).count().sort_index() + 2 0 1.0 + 1 2.0 + 3 2 1.0 + 3 2.0 + 4 3.0 + 4 5 1.0 + 6 2.0 + 7 3.0 + 8 3.0 + 5 9 1.0 + 10 2.0 + dtype: float64 + + For DataFrame, each rolling count is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).count().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 2 0 1.0 + 1 2.0 + 3 2 1.0 + 3 2.0 + 4 2.0 + 4 5 1.0 + 6 2.0 + 7 2.0 + 8 2.0 + 5 9 1.0 + 10 2.0 + """ + return super().count() + + def sum(self) -> FrameLike: + """ + The rolling summation of any non-NaN observations inside the window. + + 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.sum : Sum of the full Series. + pyspark.pandas.DataFrame.sum : Sum of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).sum().sort_index() + 2 0 NaN + 1 NaN + 3 2 NaN + 3 NaN + 4 9.0 + 4 5 NaN + 6 NaN + 7 12.0 + 8 12.0 + 5 9 NaN + 10 NaN + dtype: float64 + + For DataFrame, each rolling summation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).sum().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 2 0 NaN + 1 8.0 + 3 2 NaN + 3 18.0 + 4 18.0 + 4 5 NaN + 6 32.0 + 7 32.0 + 8 32.0 + 5 9 NaN + 10 50.0 + """ + return super().sum() + + def min(self) -> FrameLike: + """ + The rolling minimum of any non-NaN observations inside the window. + + 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.min : Min of the full Series. + pyspark.pandas.DataFrame.min : Min of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).min().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 minimum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).min().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().min() + + def max(self) -> FrameLike: + """ + The rolling maximum of any non-NaN observations inside the window. + + 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.max : Max of the full Series. + pyspark.pandas.DataFrame.max : Max of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).max().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 maximum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).max().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().max() + + def mean(self) -> FrameLike: + """ + The rolling mean of any non-NaN observations inside the window. + + 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.mean : Mean of the full Series. + pyspark.pandas.DataFrame.mean : Mean of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).mean().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 mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).mean().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().mean() + + def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: + """ + Calculate rolling quantile. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + quantile : float + Value between 0 and 1 providing the quantile to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + This is a panda-on-Spark specific parameter. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the rolling + calculation. + + Notes + ----- + `quantile` in pandas-on-Spark are using distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas, also `interpolation` + parameter is not supported yet. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling rolling with Series data. + pyspark.pandas.DataFrame.rolling : Calling rolling with DataFrames. + pyspark.pandas.Series.quantile : Aggregating quantile for Series. + pyspark.pandas.DataFrame.quantile : Aggregating quantile for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).quantile(0.5).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 quantile is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).quantile(0.5).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().quantile(quantile, accuracy) + + def std(self) -> FrameLike: + """ + Calculate rolling standard deviation. + + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + """ + return super().std() + + def var(self) -> FrameLike: + """ + Calculate unbiased rolling variance. + + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + """ + return super().var() + + def skew(self) -> FrameLike: + """ + Calculate unbiased rolling skew. + + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + """ + return super().skew() + + def kurt(self) -> FrameLike: + """ + Calculate unbiased rolling kurtosis. + + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + """ + return super().kurt() + + +class ExpandingLike(RollingAndExpanding[FrameLike]): + def __init__(self, min_periods: int = 1): + if min_periods < 0: + raise ValueError("min_periods must be >= 0") + + window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + + super().__init__(window, min_periods) + + def count(self) -> FrameLike: + def count(scol: Column) -> Column: + return F.when( + F.row_number().over(self._unbounded_window) >= self._min_periods, + F.count(scol).over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(count).astype("float64") # type: ignore[attr-defined] + + +class Expanding(ExpandingLike[FrameLike]): + def __init__(self, psdf_or_psser: FrameLike, min_periods: int = 1): + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series + + super().__init__(min_periods) + + if not isinstance(psdf_or_psser, (DataFrame, Series)): + raise TypeError( + "psdf_or_psser must be a series or dataframe; however, got: %s" + % type(psdf_or_psser) + ) + self._psdf_or_psser = psdf_or_psser + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeExpanding, item): + property_or_func = getattr(MissingPandasLikeExpanding, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + # TODO: when add 'axis' parameter, should add to here too. + def __repr__(self) -> str: + return "Expanding [min_periods={}]".format(self._min_periods) + + _apply_as_series_or_frame = Rolling._apply_as_series_or_frame + + def count(self) -> FrameLike: + """ + The expanding count of any non-NaN observations inside the window. + + .. 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.count : Count of the full Series. + pyspark.pandas.DataFrame.count : Count of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 3, float("nan"), 10]) + >>> s.expanding().count() + 0 1.0 + 1 2.0 + 2 2.0 + 3 3.0 + dtype: float64 + + >>> s.to_frame().expanding().count() + 0 + 0 1.0 + 1 2.0 + 2 2.0 + 3 3.0 + """ + return super().count() + + def sum(self) -> FrameLike: + """ + Calculate expanding summation of given DataFrame or Series. + + .. 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 + Same type as the input, with the same index, containing the + expanding summation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.sum : Reducing sum for Series. + pyspark.pandas.DataFrame.sum : Reducing sum for DataFrame. + + Examples + -------- + >>> s = ps.Series([1, 2, 3, 4, 5]) + >>> s + 0 1 + 1 2 + 2 3 + 3 4 + 4 5 + dtype: int64 + + >>> s.expanding(3).sum() + 0 NaN + 1 NaN + 2 6.0 + 3 10.0 + 4 15.0 + dtype: float64 + + For DataFrame, each expanding summation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df + A B + 0 1 1 + 1 2 4 + 2 3 9 + 3 4 16 + 4 5 25 + + >>> df.expanding(3).sum() + A B + 0 NaN NaN + 1 NaN NaN + 2 6.0 14.0 + 3 10.0 30.0 + 4 15.0 55.0 + """ + return super().sum() + + def min(self) -> FrameLike: + """ + Calculate the expanding minimum. + + .. 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 a Series. + pyspark.pandas.DataFrame.expanding : Calling object with a DataFrame. + pyspark.pandas.Series.min : Similar method for Series. + pyspark.pandas.DataFrame.min : Similar method for DataFrame. + + Examples + -------- + Performing a expanding minimum with a window size of 3. + + >>> s = ps.Series([4, 3, 5, 2, 6]) + >>> s.expanding(3).min() + 0 NaN + 1 NaN + 2 3.0 + 3 2.0 + 4 2.0 + dtype: float64 + """ + return super().min() + + def max(self) -> FrameLike: + """ + Calculate the expanding maximum. + + .. 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 + Return type is determined by the caller. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.max : Similar method for Series. + pyspark.pandas.DataFrame.max : Similar method for DataFrame. + + Examples + -------- + Performing a expanding minimum with a window size of 3. + + >>> s = ps.Series([4, 3, 5, 2, 6]) + >>> s.expanding(3).max() + 0 NaN + 1 NaN + 2 5.0 + 3 5.0 + 4 6.0 + dtype: float64 + """ + return super().max() + + def mean(self) -> FrameLike: + """ + Calculate the expanding mean of the values. + + .. 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.mean : Equivalent method for Series. + pyspark.pandas.DataFrame.mean : Equivalent method for DataFrame. + + Examples + -------- + The below examples will show expanding mean calculations with window sizes of + two and three, respectively. + + >>> s = ps.Series([1, 2, 3, 4]) + >>> s.expanding(2).mean() + 0 NaN + 1 1.5 + 2 2.0 + 3 2.5 + dtype: float64 + + >>> s.expanding(3).mean() + 0 NaN + 1 NaN + 2 2.0 + 3 2.5 + dtype: float64 + """ + return super().mean() + + def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: + """ + Calculate the expanding quantile of the values. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the expanding + calculation. + + Parameters + ---------- + quantile : float + Value between 0 and 1 providing the quantile to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + This is a panda-on-Spark specific parameter. + + Notes + ----- + `quantile` in pandas-on-Spark are using distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas (the result is + similar to the interpolation set to `lower`), also `interpolation` parameter is + not supported yet. + + 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. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling expanding with Series data. + pyspark.pandas.DataFrame.expanding : Calling expanding with DataFrames. + pyspark.pandas.Series.quantile : Aggregating quantile for Series. + pyspark.pandas.DataFrame.quantile : Aggregating quantile for DataFrame. + + Examples + -------- + The below examples will show expanding quantile calculations with window sizes of + two and three, respectively. + + >>> s = ps.Series([1, 2, 3, 4]) + >>> s.expanding(2).quantile(0.5) + 0 NaN + 1 1.0 + 2 2.0 + 3 2.0 + dtype: float64 + + >>> s.expanding(3).quantile(0.5) + 0 NaN + 1 NaN + 2 2.0 + 3 2.0 + dtype: float64 + """ + return super().quantile(quantile, accuracy) + + def std(self) -> FrameLike: + """ + Calculate expanding standard deviation. + + .. 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 + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.expanding(3).std() + 0 NaN + 1 NaN + 2 0.577350 + 3 0.957427 + 4 0.894427 + 5 0.836660 + 6 0.786796 + dtype: float64 + + For DataFrame, each expanding standard deviation variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.expanding(2).std() + A B + 0 NaN NaN + 1 0.000000 0.000000 + 2 0.577350 6.350853 + 3 0.957427 11.412712 + 4 0.894427 10.630146 + 5 0.836660 9.928075 + 6 0.786796 9.327379 + """ + return super().std() + + def var(self) -> FrameLike: + """ + Calculate unbiased expanding variance. + + .. 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 + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.expanding(3).var() + 0 NaN + 1 NaN + 2 0.333333 + 3 0.916667 + 4 0.800000 + 5 0.700000 + 6 0.619048 + dtype: float64 + + For DataFrame, each unbiased expanding variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.expanding(2).var() + A B + 0 NaN NaN + 1 0.000000 0.000000 + 2 0.333333 40.333333 + 3 0.916667 130.250000 + 4 0.800000 113.000000 + 5 0.700000 98.566667 + 6 0.619048 87.000000 + """ + return super().var() + + def skew(self) -> FrameLike: + """ + Calculate unbiased expanding skew. + + .. 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 + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 1, 5, 9]) + >>> s.expanding(3).skew() + 0 NaN + 1 NaN + 2 1.732051 + 3 0.854563 + 4 1.257788 + 5 -1.571593 + 6 -1.657542 + 7 -0.521760 + dtype: float64 + + For DataFrame, each expanding standard deviation variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.expanding(5).skew() + A B + 0 NaN NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 1.257788 1.369456 + 5 -1.571593 -0.423309 + 6 -1.657542 -0.355737 + 7 -0.521760 1.116874 + """ + return super().skew() + + def kurt(self) -> FrameLike: + """ + Calculate unbiased expanding kurtosis. + + .. 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 + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 1, 5, 9]) + >>> s.expanding(4).kurt() + 0 NaN + 1 NaN + 2 NaN + 3 -1.289256 + 4 0.312500 + 5 3.419520 + 6 4.028185 + 7 2.230373 + dtype: float64 + + For DataFrame, each unbiased expanding variance is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.expanding(5).kurt() + A B + 0 NaN NaN + 1 NaN NaN + 2 NaN NaN + 3 NaN NaN + 4 0.312500 0.906336 + 5 3.419520 1.486581 + 6 4.028185 1.936169 + 7 2.230373 2.273792 + """ + return super().kurt() + + +class ExpandingGroupby(ExpandingLike[FrameLike]): + def __init__(self, groupby: GroupBy[FrameLike], min_periods: int = 1): + super().__init__(min_periods) + + self._groupby = groupby + self._window = self._window.partitionBy(*[ser.spark.column for ser in groupby._groupkeys]) + self._unbounded_window = self._window.partitionBy( + *[ser.spark.column for ser in groupby._groupkeys] + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeExpandingGroupby, item): + property_or_func = getattr(MissingPandasLikeExpandingGroupby, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + _apply_as_series_or_frame = RollingGroupby._apply_as_series_or_frame + + def count(self) -> FrameLike: + """ + The expanding count of any non-NaN observations inside the window. + + 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.count : Count of the full Series. + pyspark.pandas.DataFrame.count : Count of the full DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).count().sort_index() + 2 0 NaN + 1 NaN + 3 2 NaN + 3 NaN + 4 3.0 + 4 5 NaN + 6 NaN + 7 3.0 + 8 4.0 + 5 9 NaN + 10 NaN + dtype: float64 + + For DataFrame, each expanding count is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).count().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 2 0 NaN + 1 2.0 + 3 2 NaN + 3 2.0 + 4 3.0 + 4 5 NaN + 6 2.0 + 7 3.0 + 8 4.0 + 5 9 NaN + 10 2.0 + """ + return super().count() + + def sum(self) -> FrameLike: + """ + Calculate expanding summation of given DataFrame or Series. + + Returns + ------- + Series or DataFrame + Same type as the input, with the same index, containing the + expanding summation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.sum : Reducing sum for Series. + pyspark.pandas.DataFrame.sum : Reducing sum for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).sum().sort_index() + 2 0 NaN + 1 NaN + 3 2 NaN + 3 NaN + 4 9.0 + 4 5 NaN + 6 NaN + 7 12.0 + 8 16.0 + 5 9 NaN + 10 NaN + dtype: float64 + + For DataFrame, each expanding summation is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).sum().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 2 0 NaN + 1 8.0 + 3 2 NaN + 3 18.0 + 4 27.0 + 4 5 NaN + 6 32.0 + 7 48.0 + 8 64.0 + 5 9 NaN + 10 50.0 + """ + return super().sum() + + def min(self) -> FrameLike: + """ + Calculate the expanding minimum. + + 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 a Series. + pyspark.pandas.DataFrame.expanding : Calling object with a DataFrame. + pyspark.pandas.Series.min : Similar method for Series. + pyspark.pandas.DataFrame.min : Similar method for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).min().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 expanding minimum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).min().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().min() + + def max(self) -> FrameLike: + """ + Calculate the expanding maximum. + + Returns + ------- + Series or DataFrame + Return type is determined by the caller. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.max : Similar method for Series. + pyspark.pandas.DataFrame.max : Similar method for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).max().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 expanding maximum is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).max().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().max() + + def mean(self) -> FrameLike: + """ + Calculate the expanding mean of the values. + + 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.mean : Equivalent method for Series. + pyspark.pandas.DataFrame.mean : Equivalent method for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).mean().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 expanding mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).mean().sort_index() # doctest: +NORMALIZE_WHITESPACE + 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().mean() + + def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: + """ + Calculate the expanding quantile of the values. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + quantile : float + Value between 0 and 1 providing the quantile to compute. + accuracy : int, optional + Default accuracy of approximation. Larger value means better accuracy. + The relative error can be deduced by 1.0 / accuracy. + This is a panda-on-Spark specific parameter. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the expanding + calculation. + + Notes + ----- + `quantile` in pandas-on-Spark are using distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas, also `interpolation` + parameter is not supported yet. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling expanding with Series data. + pyspark.pandas.DataFrame.expanding : Calling expanding with DataFrames. + pyspark.pandas.Series.quantile : Aggregating quantile for Series. + pyspark.pandas.DataFrame.quantile : Aggregating quantile for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).quantile(0.5).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 expanding quantile is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).expanding(2).quantile(0.5).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().quantile(quantile, accuracy) + + def std(self) -> FrameLike: + """ + Calculate expanding standard deviation. + + + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + """ + return super().std() + + def var(self) -> FrameLike: + """ + Calculate unbiased expanding variance. + + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + """ + return super().var() + + def skew(self) -> FrameLike: + """ + Calculate expanding standard skew. + + + 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.std : Equivalent method for Series. + pyspark.pandas.DataFrame.std : Equivalent method for DataFrame. + numpy.std : Equivalent method for Numpy array. + """ + return super().skew() + + def kurt(self) -> FrameLike: + """ + Calculate unbiased expanding kurtosis. + + 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.var : Equivalent method for Series. + pyspark.pandas.DataFrame.var : Equivalent method for DataFrame. + numpy.var : Equivalent method for Numpy array. + """ + return super().kurt() + + +class ExponentialMovingLike(Generic[FrameLike], metaclass=ABCMeta): + def __init__( + self, + window: WindowSpec, + com: Optional[float] = None, + span: Optional[float] = None, + halflife: Optional[float] = None, + alpha: Optional[float] = None, + min_periods: Optional[int] = None, + ignore_na: bool = False, + ): + if (min_periods is not None) and (min_periods < 0): + raise ValueError("min_periods must be >= 0") + if min_periods is None: + min_periods = 0 + self._min_periods = min_periods + self._ignore_na = ignore_na + + self._window = window + # This unbounded Window is later used to handle 'min_periods' for now. + self._unbounded_window = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + + if (com is not None) and (not com >= 0): + raise ValueError("com must be >= 0") + self._com = com + + if (span is not None) and (not span >= 1): + raise ValueError("span must be >= 1") + self._span = span + + if (halflife is not None) and (not halflife > 0): + raise ValueError("halflife must be > 0") + self._halflife = halflife + + if (alpha is not None) and (not 0 < alpha <= 1): + raise ValueError("alpha must be in (0, 1]") + self._alpha = alpha + + def _compute_unified_alpha(self) -> float: + unified_alpha = np.nan + opt_count = 0 + + if self._com is not None: + unified_alpha = 1.0 / (1 + self._com) + opt_count += 1 + if self._span is not None: + unified_alpha = 2.0 / (1 + self._span) + opt_count += 1 + if self._halflife is not None: + unified_alpha = 1.0 - np.exp(-np.log(2) / self._halflife) + opt_count += 1 + if self._alpha is not None: + unified_alpha = self._alpha + opt_count += 1 + + if opt_count == 0: + raise ValueError("Must pass one of com, span, halflife, or alpha") + if opt_count != 1: + raise ValueError("com, span, halflife, and alpha are mutually exclusive") + + return unified_alpha + + @abstractmethod + def _apply_as_series_or_frame(self, func: Callable[[Column], Column]) -> FrameLike: + """ + Wraps a function that handles Spark column in order + to support it in both pandas-on-Spark Series and DataFrame. + Note that the given `func` name should be same as the API's method name. + """ + pass + + def mean(self) -> FrameLike: + unified_alpha = self._compute_unified_alpha() + + def mean(scol: Column) -> Column: + col_ewm = SF.ewm(scol, unified_alpha, self._ignore_na) + return F.when( + F.count(F.when(~scol.isNull(), 1).otherwise(None)).over(self._unbounded_window) + >= self._min_periods, + col_ewm.over(self._window), + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(mean) + + +class ExponentialMoving(ExponentialMovingLike[FrameLike]): + def __init__( + self, + psdf_or_psser: FrameLike, + com: Optional[float] = None, + span: Optional[float] = None, + halflife: Optional[float] = None, + alpha: Optional[float] = None, + min_periods: Optional[int] = None, + ignore_na: bool = False, + ): + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.series import Series + + if not isinstance(psdf_or_psser, (DataFrame, Series)): + raise TypeError( + "psdf_or_psser must be a series or dataframe; however, got: %s" + % type(psdf_or_psser) + ) + self._psdf_or_psser = psdf_or_psser + + window_spec = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + + super().__init__(window_spec, com, span, halflife, alpha, min_periods, ignore_na) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeExponentialMoving, item): + property_or_func = getattr(MissingPandasLikeExponentialMoving, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + _apply_as_series_or_frame = Rolling._apply_as_series_or_frame + + def mean(self) -> FrameLike: + """ + Calculate an online exponentially weighted mean. + + Notes + ----- + There are behavior differences between pandas-on-Spark and pandas. + + * 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 exponentially + calculation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.mean : Equivalent method for Series. + pyspark.pandas.DataFrame.mean : Equivalent method for DataFrame. + + Examples + -------- + The below examples will show computing exponentially weighted moving average. + + >>> df = ps.DataFrame({'s1': [.2, .0, .6, .2, .4, .5, .6], 's2': [2, 1, 3, 1, 0, 0, 0]}) + >>> df.ewm(com=0.1).mean() + s1 s2 + 0 0.200000 2.000000 + 1 0.016667 1.083333 + 2 0.547368 2.827068 + 3 0.231557 1.165984 + 4 0.384688 0.105992 + 5 0.489517 0.009636 + 6 0.589956 0.000876 + + >>> df.s2.ewm(halflife=1.5, min_periods=3).mean() + 0 NaN + 1 NaN + 2 2.182572 + 3 1.663174 + 4 0.979949 + 5 0.593155 + 6 0.364668 + Name: s2, dtype: float64 + """ + return super().mean() + + # TODO: when add 'adjust' parameter, should add to here too. + def __repr__(self) -> str: + return ( + "ExponentialMoving [com={}, span={}, halflife={}, alpha={}, " + "min_periods={}, ignore_na={}]".format( + self._com, + self._span, + self._halflife, + self._alpha, + self._min_periods, + self._ignore_na, + ) + ) + + +class ExponentialMovingGroupby(ExponentialMovingLike[FrameLike]): + def __init__( + self, + groupby: GroupBy[FrameLike], + com: Optional[float] = None, + span: Optional[float] = None, + halflife: Optional[float] = None, + alpha: Optional[float] = None, + min_periods: Optional[int] = None, + ignore_na: bool = False, + ): + window_spec = Window.orderBy(NATURAL_ORDER_COLUMN_NAME).rowsBetween( + Window.unboundedPreceding, Window.currentRow + ) + super().__init__(window_spec, com, span, halflife, alpha, min_periods, ignore_na) + + self._groupby = groupby + self._window = self._window.partitionBy(*[ser.spark.column for ser in groupby._groupkeys]) + self._unbounded_window = self._unbounded_window.partitionBy( + *[ser.spark.column for ser in groupby._groupkeys] + ) + + def __getattr__(self, item: str) -> Any: + if hasattr(MissingPandasLikeExponentialMovingGroupby, item): + property_or_func = getattr(MissingPandasLikeExponentialMovingGroupby, item) + if isinstance(property_or_func, property): + return property_or_func.fget(self) + else: + return partial(property_or_func, self) + raise AttributeError(item) + + _apply_as_series_or_frame = RollingGroupby._apply_as_series_or_frame + + def mean(self) -> FrameLike: + """ + Calculate an online exponentially weighted mean. + + Notes + ----- + There are behavior differences between pandas-on-Spark and pandas. + + * 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 exponentially + calculation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.mean : Equivalent method for Series. + pyspark.pandas.DataFrame.mean : Equivalent method for DataFrame. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).ewm(alpha=0.5).mean().sort_index() + 2 0 2.0 + 1 2.0 + 3 2 3.0 + 3 3.0 + 4 3.0 + 4 5 4.0 + 6 4.0 + 7 4.0 + 8 4.0 + 5 9 5.0 + 10 5.0 + dtype: float64 + + For DataFrame, each ewm mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).ewm(alpha=0.5).mean().sort_index() # doctest: +NORMALIZE_WHITESPACE + B + A + 2 0 4.0 + 1 4.0 + 3 2 9.0 + 3 9.0 + 4 9.0 + 4 5 16.0 + 6 16.0 + 7 16.0 + 8 16.0 + 5 9 25.0 + 10 25.0 + """ + return super().mean() + + # TODO: when add 'adjust' parameter, should add to here too. + def __repr__(self) -> str: + return ( + "ExponentialMovingGroupby [com={}, span={}, halflife={}, alpha={}, " + "min_periods={}, ignore_na={}]".format( + self._com, + self._span, + self._halflife, + self._alpha, + self._min_periods, + self._ignore_na, + ) + ) + + +def _test() -> None: + import os + import doctest + import sys + from pyspark.sql import SparkSession + import pyspark.pandas.window + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.pandas.window.__dict__.copy() + globs["ps"] = pyspark.pandas + spark = ( + SparkSession.builder.master("local[4]").appName("pyspark.pandas.window tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.pandas.window, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/profiler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/profiler.py new file mode 100644 index 0000000..d7990d2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/profiler.py @@ -0,0 +1,489 @@ +# +# 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 ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + TYPE_CHECKING, + Union, + cast, +) + +import cProfile +import inspect +import pstats +import linecache +import os +import atexit +import sys +import warnings + +try: + from memory_profiler import choose_backend, CodeMap, LineProfiler # type: ignore[import] + + has_memory_profiler = True +except Exception: + has_memory_profiler = False + +from pyspark.accumulators import AccumulatorParam +from pyspark.errors import PySparkRuntimeError + +if TYPE_CHECKING: + from pyspark.context import SparkContext + +MemoryTuple = Tuple[float, float, int] +LineProfile = Tuple[int, Optional[MemoryTuple]] +CodeMapDict = Dict[str, List[LineProfile]] + + +class ProfilerCollector: + """ + This class keeps track of different profilers on a per + stage/UDF basis. Also this is used to create new profilers for + the different stages/UDFs. + """ + + def __init__( + self, + profiler_cls: Type["Profiler"], + udf_profiler_cls: Type["Profiler"], + memory_profiler_cls: Type["Profiler"], + dump_path: Optional[str] = None, + ): + self.profiler_cls: Type[Profiler] = profiler_cls + self.udf_profiler_cls: Type[Profiler] = udf_profiler_cls + self.memory_profiler_cls: Type[Profiler] = memory_profiler_cls + self.profile_dump_path: Optional[str] = dump_path + self.profilers: List[List[Any]] = [] + + def new_profiler(self, ctx: "SparkContext") -> "Profiler": + """Create a new profiler using class `profiler_cls`""" + return self.profiler_cls(ctx) + + def new_udf_profiler(self, ctx: "SparkContext") -> "Profiler": + """Create a new profiler using class `udf_profiler_cls`""" + return self.udf_profiler_cls(ctx) + + def new_memory_profiler(self, ctx: "SparkContext") -> "Profiler": + """Create a new profiler using class `memory_profiler_cls`""" + return self.memory_profiler_cls(ctx) + + def add_profiler(self, id: int, profiler: "Profiler") -> None: + """Add a profiler for RDD/UDF `id`""" + if not self.profilers: + if self.profile_dump_path: + atexit.register(self.dump_profiles, self.profile_dump_path) + else: + atexit.register(self.show_profiles) + + self.profilers.append([id, profiler, False]) + + def dump_profiles(self, path: str) -> None: + """Dump the profile stats into directory `path`""" + for id, profiler, _ in self.profilers: + profiler.dump(id, path) + self.profilers = [] + + def show_profiles(self) -> None: + """Print the profile stats to stdout""" + for i, (id, profiler, showed) in enumerate(self.profilers): + if not showed and profiler: + profiler.show(id) + # mark it as showed + self.profilers[i][2] = True + + +class Profiler: + """ + PySpark supports custom profilers, this is to allow for different profilers to + be used as well as outputting to different formats than what is provided in the + BasicProfiler. + + A custom profiler has to define or inherit the following methods: + profile - will produce a system profile of some sort. + stats - return the collected stats. + dump - dumps the profiles to a path + add - adds a profile to the existing accumulated profile + + The profiler class is chosen when creating a SparkContext + + Examples + -------- + >>> from pyspark import SparkConf, SparkContext + >>> from pyspark import BasicProfiler + >>> class MyCustomProfiler(BasicProfiler): + ... def show(self, id): + ... print("My custom profiles for RDD:%s" % id) + ... + >>> conf = SparkConf().set("spark.python.profile", "true") + >>> sc = SparkContext('local', 'test', conf=conf, profiler_cls=MyCustomProfiler) + >>> sc.parallelize(range(1000)).map(lambda x: 2 * x).take(10) + [0, 2, 4, 6, 8, 10, 12, 14, 16, 18] + >>> sc.parallelize(range(1000)).count() + 1000 + >>> sc.show_profiles() + My custom profiles for RDD:1 + My custom profiles for RDD:3 + >>> sc.stop() + + Notes + ----- + This API is a developer API. + """ + + def __init__(self, ctx: "SparkContext") -> None: + pass + + def profile(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + """Do profiling on the function `func`""" + raise NotImplementedError + + def stats(self) -> Union[pstats.Stats, Dict]: + """Return the collected profiling stats""" + raise NotImplementedError + + def show(self, id: int) -> None: + """Print the profile stats to stdout""" + raise NotImplementedError + + def dump(self, id: int, path: str) -> None: + """Dump the profile into path""" + raise NotImplementedError + + +if has_memory_profiler: + + class CodeMapForUDF(CodeMap): + def add( + self, + code: Any, + toplevel_code: Optional[Any] = None, + *, + sub_lines: Optional[List] = None, + start_line: Optional[int] = None, + ) -> None: + if code in self: + return + + if toplevel_code is None: + toplevel_code = code + filename = code.co_filename + if sub_lines is None or start_line is None: + (sub_lines, start_line) = inspect.getsourcelines(code) + linenos = range(start_line, start_line + len(sub_lines)) + self._toplevel.append((filename, code, linenos)) + self[code] = {} + else: + self[code] = self[toplevel_code] + for subcode in filter(inspect.iscode, code.co_consts): + self.add(subcode, toplevel_code=toplevel_code) + + class UDFLineProfiler(LineProfiler): + def __init__(self, **kw: Any) -> None: + include_children = kw.get("include_children", False) + backend = kw.get("backend", "psutil") + self.code_map = CodeMapForUDF(include_children=include_children, backend=backend) + self.enable_count = 0 + self.max_mem = kw.get("max_mem", None) + self.prevlines: List = [] + self.backend = choose_backend(kw.get("backend", None)) + self.prev_lineno = None + + def __call__( + self, + func: Optional[Callable[..., Any]] = None, + precision: int = 1, + *, + sub_lines: Optional[List] = None, + start_line: Optional[int] = None, + ) -> Callable[..., Any]: + if func is not None: + self.add_function(func, sub_lines=sub_lines, start_line=start_line) + f = self.wrap_function(func) + f.__module__ = func.__module__ + f.__name__ = func.__name__ + f.__doc__ = func.__doc__ + f.__dict__.update(getattr(func, "__dict__", {})) + return f + else: + + def inner_partial(f: Callable[..., Any]) -> Any: + return self.__call__(f, precision=precision) + + return inner_partial + + def add_function( + self, + func: Callable[..., Any], + *, + sub_lines: Optional[List] = None, + start_line: Optional[int] = None, + ) -> None: + """Record line profiling information for the given Python function.""" + try: + # func_code does not exist in Python3 + code = func.__code__ + except AttributeError: + warnings.warn("Could not extract a code object for the object %r" % func) + else: + self.code_map.add(code, sub_lines=sub_lines, start_line=start_line) + + +class PStatsParam(AccumulatorParam[Optional[pstats.Stats]]): + """PStatsParam is used to merge pstats.Stats""" + + @staticmethod + def zero(value: Optional[pstats.Stats]) -> None: + return None + + @staticmethod + def addInPlace( + value1: Optional[pstats.Stats], value2: Optional[pstats.Stats] + ) -> Optional[pstats.Stats]: + if value1 is None: + return value2 + value1.add(value2) + return value1 + + +class MemUsageParam(AccumulatorParam[Optional[CodeMapDict]]): + """MemUsageParam is used to merge memory usage code map""" + + @staticmethod + def zero(value: Optional[CodeMapDict]) -> None: + return None + + @staticmethod + def addInPlace( + value1: Optional[CodeMapDict], value2: Optional[CodeMapDict] + ) -> Optional[CodeMapDict]: + # An example value looks as below + # {'': [(3, (144.2578125, 144.2578125, 1)), + # (4, (0.0, 144.2578125, 1))]} + if value1 is None or len(value1) == 0: + return value2 + if value2 is None or len(value2) == 0: + return value1 + + # value1, value2 should have same keys - file name + for filename in value1: + l1 = cast(List[LineProfile], value1.get(filename)) + l2 = cast(List[LineProfile], value2.get(filename)) + c1 = dict((k, v) for k, v in l1) + c2 = dict((k, v) for k, v in l2) + udf_code_map: Dict[int, Optional[MemoryTuple]] = {} + for lineno in c1: + if c1[lineno] and c2[lineno]: + # c1, c2 should have same keys - line number + udf_code_map[lineno] = ( + cast(MemoryTuple, c1[lineno])[0] + + cast(MemoryTuple, c2[lineno])[0], # increment + cast(MemoryTuple, c1[lineno])[1] + + cast(MemoryTuple, c2[lineno])[1], # mem_usage + cast(MemoryTuple, c1[lineno])[2] + + cast(MemoryTuple, c2[lineno])[2], # occurrences + ) + elif c1[lineno]: + udf_code_map[lineno] = cast(MemoryTuple, c1[lineno]) + elif c2[lineno]: + udf_code_map[lineno] = cast(MemoryTuple, c2[lineno]) + else: + udf_code_map[lineno] = None + value1[filename] = [(k, v) for k, v in udf_code_map.items()] + return value1 + + +class BasicProfiler(Profiler): + """ + BasicProfiler is the default profiler, which is implemented based on + cProfile and Accumulator + """ + + def __init__(self, ctx: "SparkContext") -> None: + super().__init__(ctx) + # Creates a new accumulator for combining the profiles of different + # partitions of a stage + self._accumulator = ctx.accumulator(None, PStatsParam) # type: ignore[arg-type] + + def profile(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + """Runs and profiles the method to_profile passed in. A profile object is returned.""" + pr = cProfile.Profile() + ret = pr.runcall(func, *args, **kwargs) + st = pstats.Stats(pr) + st.stream = None # type: ignore[attr-defined] # make it picklable + st.strip_dirs() + + # Adds a new profile to the existing accumulated value + self._accumulator.add(st) # type: ignore[arg-type] + + return ret + + def stats(self) -> pstats.Stats: + return cast(pstats.Stats, self._accumulator.value) + + def show(self, id: int) -> None: + """Print the profile stats to stdout, id is the RDD id""" + stats = self.stats() + if stats: + print("=" * 60) + print("Profile of RDD" % id) + print("=" * 60) + stats.sort_stats("time", "cumulative").print_stats() + + def dump(self, id: int, path: str) -> None: + """Dump the profile into path, id is the RDD id""" + if not os.path.exists(path): + os.makedirs(path) + stats = self.stats() + if stats: + p = os.path.join(path, "rdd_%d.pstats" % id) + stats.dump_stats(p) + + +class UDFBasicProfiler(BasicProfiler): + """ + UDFBasicProfiler is the profiler for Python/Pandas UDFs. + """ + + def show(self, id: int) -> None: + """Print the profile stats to stdout, id is the PythonUDF id""" + stats = self.stats() + if stats: + print("=" * 60) + print("Profile of UDF" % id) + print("=" * 60) + stats.sort_stats("time", "cumulative").print_stats() + + def dump(self, id: int, path: str) -> None: + """Dump the profile into path, id is the PythonUDF id""" + if not os.path.exists(path): + os.makedirs(path) + stats = self.stats() + if stats: + p = os.path.join(path, "udf_%d.pstats" % id) + stats.dump_stats(p) + + +class MemoryProfiler(Profiler): + """ + MemoryProfiler, which is implemented based on memory profiler and Accumulator + """ + + def __init__(self, ctx: "SparkContext") -> None: + super().__init__(ctx) + # Creates a new accumulator for combining the profiles + self._accumulator = ctx.accumulator(None, MemUsageParam) # type: ignore[arg-type] + + def profile( # type: ignore + self, + sub_lines: Optional[List], + start_line: Optional[int], + func: Callable[..., Any], + *args: Any, + **kwargs: Any, + ) -> Any: + """Runs and profiles the method func passed in. A profile object is returned.""" + if has_memory_profiler: + profiler = UDFLineProfiler() + wrapped = profiler(func, sub_lines=sub_lines, start_line=start_line) + ret = wrapped(*args, **kwargs) + codemap_dict = { + filename: list(line_iterator) + for filename, line_iterator in profiler.code_map.items() + } + # Adds a new profile to the existing accumulated value + self._accumulator.add(codemap_dict) # type: ignore[arg-type] + return ret + else: + raise PySparkRuntimeError( + error_class="MISSING_LIBRARY_FOR_PROFILER", + message_parameters={}, + ) + + def stats(self) -> CodeMapDict: + """Return the collected memory profiles""" + return cast(CodeMapDict, self._accumulator.value) + + def _show_results( + self, code_map: CodeMapDict, stream: Optional[Any] = None, precision: int = 1 + ) -> None: + if stream is None: + stream = sys.stdout + template = "{0:>6} {1:>12} {2:>12} {3:>10} {4:<}" + + for (filename, lines) in code_map.items(): + header = template.format( + "Line #", "Mem usage", "Increment", "Occurrences", "Line Contents" + ) + + stream.write("Filename: " + filename + "\n\n") + stream.write(header + "\n") + stream.write("=" * len(header) + "\n") + + all_lines = linecache.getlines(filename) + + float_format = "{0}.{1}f".format(precision + 4, precision) + template_mem = "{0:" + float_format + "} MiB" + for (lineno, mem) in lines: + total_mem: Union[float, str] + inc: Union[float, str] + occurrences: Union[float, str] + if mem: + inc = mem[0] + total_mem = mem[1] + total_mem = template_mem.format(total_mem) + occurrences = mem[2] + inc = template_mem.format(inc) + else: + total_mem = "" + inc = "" + occurrences = "" + tmp = template.format(lineno, total_mem, inc, occurrences, all_lines[lineno - 1]) + stream.write(tmp) + stream.write("\n\n") + + def show(self, id: int) -> None: + """Print the profile stats to stdout, id is the PythonUDF id""" + code_map = self.stats() + if code_map: + print("=" * 60) + print("Profile of UDF" % id) + print("=" * 60) + self._show_results(code_map) + + def dump(self, id: int, path: str) -> None: + """Dump the memory profile into path, id is the PythonUDF id""" + if not os.path.exists(path): + os.makedirs(path) + stats = self.stats() # dict + if stats: + p = os.path.join(path, "udf_%d_memory.txt" % id) + with open(p, "w+") as f: + self._show_results(stats, stream=f) + + +if __name__ == "__main__": + import doctest + + (failure_count, test_count) = doctest.testmod() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/py.typed b/spark-3.5.0-bin-hadoop3/python/pyspark/py.typed new file mode 100644 index 0000000..b648ac9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/py.typed @@ -0,0 +1 @@ +partial diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/python/pyspark/shell.py b/spark-3.5.0-bin-hadoop3/python/pyspark/python/pyspark/shell.py new file mode 100644 index 0000000..7e2093c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/python/pyspark/shell.py @@ -0,0 +1,123 @@ +# +# 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. +# + +""" +An interactive shell. + +This file is designed to be launched as a PYTHONSTARTUP script. +""" + +import atexit +import builtins +import os +import platform +import warnings +import sys + +import pyspark +from pyspark.context import SparkContext +from pyspark.sql import SparkSession +from pyspark.sql.context import SQLContext +from pyspark.sql.utils import is_remote +from urllib.parse import urlparse + +if getattr(builtins, "__IPYTHON__", False): + # (Only) during PYTHONSTARTUP execution, IPython temporarily adds the parent + # directory of the script into the Python path, which results in searching + # packages under `pyspark` directory. + # For example, `import pandas` attempts to import `pyspark.pandas`, see also SPARK-42266. + if "__file__" in globals(): + parent_dir = os.path.abspath(os.path.dirname(__file__)) + if parent_dir in sys.path: + sys.path.remove(parent_dir) + + +if is_remote(): + try: + # Creates pyspark.sql.connect.SparkSession. + spark = SparkSession.builder.getOrCreate() + except Exception: + import sys + import traceback + + warnings.warn("Failed to initialize Spark session.") + traceback.print_exc(file=sys.stderr) + sys.exit(1) + version = pyspark.__version__ + sc = None +else: + if os.environ.get("SPARK_EXECUTOR_URI"): + SparkContext.setSystemProperty("spark.executor.uri", os.environ["SPARK_EXECUTOR_URI"]) + + SparkContext._ensure_initialized() + + try: + spark = SparkSession._create_shell_session() + except Exception: + import sys + import traceback + + warnings.warn("Failed to initialize Spark session.") + traceback.print_exc(file=sys.stderr) + sys.exit(1) + + sc = spark.sparkContext + atexit.register((lambda sc: lambda: sc.stop())(sc)) + + # for compatibility + sqlContext = SQLContext._get_or_create(sc) + sqlCtx = sqlContext + version = sc.version + +sql = spark.sql + +print( + r"""Welcome to + ____ __ + / __/__ ___ _____/ /__ + _\ \/ _ \/ _ `/ __/ '_/ + /__ / .__/\_,_/_/ /_/\_\ version %s + /_/ +""" + % version +) +print( + "Using Python version %s (%s, %s)" + % (platform.python_version(), platform.python_build()[0], platform.python_build()[1]) +) +if is_remote(): + url = os.environ.get("SPARK_REMOTE", None) + assert url is not None + if url.startswith("local"): + url = "sc://localhost" # only for display in the console. + print("Client connected to the Spark Connect server at %s" % urlparse(url).netloc) +else: + print("Spark context Web UI available at %s" % (sc.uiWebUrl)) # type: ignore[union-attr] + print( + "Spark context available as 'sc' (master = %s, app id = %s)." + % (sc.master, sc.applicationId) # type: ignore[union-attr] + ) + +print("SparkSession available as 'spark'.") + +# The ./bin/pyspark script stores the old PYTHONSTARTUP value in OLD_PYTHONSTARTUP, +# which allows us to execute the user's PYTHONSTARTUP file: +_pythonstartup = os.environ.get("OLD_PYTHONSTARTUP") +if _pythonstartup and os.path.isfile(_pythonstartup): + with open(_pythonstartup) as f: + code = compile(f.read(), _pythonstartup, "exec") + exec(code) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/rdd.py b/spark-3.5.0-bin-hadoop3/python/pyspark/rdd.py new file mode 100644 index 0000000..8ea9a31 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/rdd.py @@ -0,0 +1,5518 @@ +# +# 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 copy +import sys +import os +import operator +import shlex +import warnings +import heapq +import bisect +import random +from subprocess import Popen, PIPE +from threading import Thread +from collections import defaultdict +from itertools import chain +from functools import reduce +from math import sqrt, log, isinf, isnan, pow, ceil +from typing import ( + Any, + Callable, + Dict, + Generic, + Hashable, + Iterable, + Iterator, + IO, + List, + NoReturn, + Optional, + Sequence, + Tuple, + Union, + TypeVar, + cast, + overload, + TYPE_CHECKING, +) + +from pyspark.java_gateway import local_connect_and_auth +from pyspark.serializers import ( + AutoBatchedSerializer, + BatchedSerializer, + NoOpSerializer, + CartesianDeserializer, + CloudPickleSerializer, + PairDeserializer, + CPickleSerializer, + Serializer, + pack_long, + read_int, + write_int, +) +from pyspark.join import ( + python_join, + python_left_outer_join, + python_right_outer_join, + python_full_outer_join, + python_cogroup, +) +from pyspark.statcounter import StatCounter +from pyspark.rddsampler import RDDSampler, RDDRangeSampler, RDDStratifiedSampler +from pyspark.storagelevel import StorageLevel +from pyspark.resource.requests import ExecutorResourceRequests, TaskResourceRequests +from pyspark.resource.profile import ResourceProfile +from pyspark.resultiterable import ResultIterable +from pyspark.shuffle import ( + Aggregator, + ExternalMerger, + get_used_memory, + ExternalSorter, + ExternalGroupBy, +) +from pyspark.traceback_utils import SCCallSiteSync +from pyspark.util import fail_on_stopiteration, _parse_memory +from pyspark.errors import PySparkRuntimeError + + +if TYPE_CHECKING: + import socket + import io + + from pyspark._typing import NonUDFType + from pyspark._typing import S, NumberOrArray + from pyspark.context import SparkContext + from pyspark.sql.pandas._typing import ( + PandasScalarUDFType, + PandasGroupedMapUDFType, + PandasGroupedAggUDFType, + PandasWindowAggUDFType, + PandasScalarIterUDFType, + PandasMapIterUDFType, + PandasCogroupedMapUDFType, + ArrowMapIterUDFType, + PandasGroupedMapUDFWithStateType, + ) + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.types import AtomicType, StructType + from pyspark.sql._typing import ( + AtomicValue, + RowLike, + SQLArrowBatchedUDFType, + SQLArrowTableUDFType, + SQLBatchedUDFType, + SQLTableUDFType, + ) + + from py4j.java_gateway import JavaObject + from py4j.java_collections import JavaArray + +T = TypeVar("T") +T_co = TypeVar("T_co", covariant=True) +U = TypeVar("U") +K = TypeVar("K", bound=Hashable) +V = TypeVar("V") +V1 = TypeVar("V1") +V2 = TypeVar("V2") +V3 = TypeVar("V3") + + +__all__ = ["RDD"] + + +class PythonEvalType: + """ + Evaluation type of python rdd. + + These values are internal to PySpark. + + These values should match values in org.apache.spark.api.python.PythonEvalType. + """ + + NON_UDF: "NonUDFType" = 0 + + SQL_BATCHED_UDF: "SQLBatchedUDFType" = 100 + SQL_ARROW_BATCHED_UDF: "SQLArrowBatchedUDFType" = 101 + + SQL_SCALAR_PANDAS_UDF: "PandasScalarUDFType" = 200 + SQL_GROUPED_MAP_PANDAS_UDF: "PandasGroupedMapUDFType" = 201 + SQL_GROUPED_AGG_PANDAS_UDF: "PandasGroupedAggUDFType" = 202 + SQL_WINDOW_AGG_PANDAS_UDF: "PandasWindowAggUDFType" = 203 + SQL_SCALAR_PANDAS_ITER_UDF: "PandasScalarIterUDFType" = 204 + SQL_MAP_PANDAS_ITER_UDF: "PandasMapIterUDFType" = 205 + SQL_COGROUPED_MAP_PANDAS_UDF: "PandasCogroupedMapUDFType" = 206 + SQL_MAP_ARROW_ITER_UDF: "ArrowMapIterUDFType" = 207 + SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: "PandasGroupedMapUDFWithStateType" = 208 + + SQL_TABLE_UDF: "SQLTableUDFType" = 300 + SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301 + + +def portable_hash(x: Hashable) -> int: + """ + This function returns consistent hash code for builtin types, especially + for None and tuple with None. + + The algorithm is similar to that one used by CPython 2.7 + + Examples + -------- + >>> portable_hash(None) + 0 + >>> portable_hash((None, 1)) & 0xffffffff + 219750521 + """ + + if "PYTHONHASHSEED" not in os.environ: + raise PySparkRuntimeError( + error_class="PYTHON_HASH_SEED_NOT_SET", + message_parameters={}, + ) + + if x is None: + return 0 + if isinstance(x, tuple): + h = 0x345678 + for i in x: + h ^= portable_hash(i) + h *= 1000003 + h &= sys.maxsize + h ^= len(x) + if h == -1: + h = -2 + return int(h) + return hash(x) + + +class BoundedFloat(float): + """ + Bounded value is generated by approximate job, with confidence and low + bound and high bound. + + Examples + -------- + >>> BoundedFloat(100.0, 0.95, 95.0, 105.0) + 100.0 + """ + + confidence: float + low: float + high: float + + def __new__(cls, mean: float, confidence: float, low: float, high: float) -> "BoundedFloat": + obj = float.__new__(cls, mean) + obj.confidence = confidence + obj.low = low + obj.high = high + return obj + + +def _create_local_socket(sock_info: "JavaArray") -> "io.BufferedRWPair": + """ + Create a local socket that can be used to load deserialized data from the JVM + + Parameters + ---------- + sock_info : tuple + Tuple containing port number and authentication secret for a local socket. + + Returns + ------- + sockfile file descriptor of the local socket + """ + sockfile: "io.BufferedRWPair" + sock: "socket.socket" + port: int = sock_info[0] + auth_secret: str = sock_info[1] + sockfile, sock = local_connect_and_auth(port, auth_secret) + # The RDD materialization time is unpredictable, if we set a timeout for socket reading + # operation, it will very possibly fail. See SPARK-18281. + sock.settimeout(None) + return sockfile + + +def _load_from_socket(sock_info: "JavaArray", serializer: Serializer) -> Iterator[Any]: + """ + Connect to a local socket described by sock_info and use the given serializer to yield data + + Parameters + ---------- + sock_info : tuple + Tuple containing port number and authentication secret for a local socket. + serializer : class:`Serializer` + The PySpark serializer to use + + Returns + ------- + result of meth:`Serializer.load_stream`, + usually a generator that yields deserialized data + """ + sockfile = _create_local_socket(sock_info) + # The socket will be automatically closed when garbage-collected. + return serializer.load_stream(sockfile) + + +def _local_iterator_from_socket(sock_info: "JavaArray", serializer: Serializer) -> Iterator[Any]: + class PyLocalIterable: + """Create a synchronous local iterable over a socket""" + + def __init__(self, _sock_info: "JavaArray", _serializer: Serializer): + port: int + auth_secret: str + jsocket_auth_server: "JavaObject" + port, auth_secret, self.jsocket_auth_server = _sock_info + self._sockfile = _create_local_socket((port, auth_secret)) + self._serializer = _serializer + self._read_iter: Iterator[Any] = iter([]) # Initialize as empty iterator + self._read_status = 1 + + def __iter__(self) -> Iterator[Any]: + while self._read_status == 1: + # Request next partition data from Java + write_int(1, self._sockfile) + self._sockfile.flush() + + # If response is 1 then there is a partition to read, if 0 then fully consumed + self._read_status = read_int(self._sockfile) + if self._read_status == 1: + + # Load the partition data as a stream and read each item + self._read_iter = self._serializer.load_stream(self._sockfile) + for item in self._read_iter: + yield item + + # An error occurred, join serving thread and raise any exceptions from the JVM + elif self._read_status == -1: + self.jsocket_auth_server.getResult() + + def __del__(self) -> None: + # If local iterator is not fully consumed, + if self._read_status == 1: + try: + # Finish consuming partition data stream + for _ in self._read_iter: + pass + # Tell Java to stop sending data and close connection + write_int(0, self._sockfile) + self._sockfile.flush() + except Exception: + # Ignore any errors, socket is automatically closed when garbage-collected + pass + + return iter(PyLocalIterable(sock_info, serializer)) + + +class Partitioner: + def __init__(self, numPartitions: int, partitionFunc: Callable[[Any], int]): + self.numPartitions = numPartitions + self.partitionFunc = partitionFunc + + def __eq__(self, other: Any) -> bool: + return ( + isinstance(other, Partitioner) + and self.numPartitions == other.numPartitions + and self.partitionFunc == other.partitionFunc + ) + + def __call__(self, k: Any) -> int: + return self.partitionFunc(k) % self.numPartitions + + +class RDD(Generic[T_co]): + + """ + A Resilient Distributed Dataset (RDD), the basic abstraction in Spark. + Represents an immutable, partitioned collection of elements that can be + operated on in parallel. + """ + + def __init__( + self, + jrdd: "JavaObject", + ctx: "SparkContext", + jrdd_deserializer: Serializer = AutoBatchedSerializer(CPickleSerializer()), + ): + self._jrdd = jrdd + self.is_cached = False + self.is_checkpointed = False + self.has_resource_profile = False + self.ctx = ctx + self._jrdd_deserializer = jrdd_deserializer + self._id = jrdd.id() + self.partitioner: Optional[Partitioner] = None + + def _pickled(self: "RDD[T]") -> "RDD[T]": + return self._reserialize(AutoBatchedSerializer(CPickleSerializer())) + + def id(self) -> int: + """ + A unique ID for this RDD (within its SparkContext). + + .. versionadded:: 0.7.0 + + Returns + ------- + int + The unique ID for this :class:`RDD` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.id() # doctest: +SKIP + 3 + """ + return self._id + + def __repr__(self) -> str: + return self._jrdd.toString() + + def __getnewargs__(self) -> NoReturn: + # This method is called when attempting to pickle an RDD, which is always an error: + raise PySparkRuntimeError( + error_class="RDD_TRANSFORM_ONLY_VALID_ON_DRIVER", + message_parameters={}, + ) + + @property + def context(self) -> "SparkContext": + """ + The :class:`SparkContext` that this RDD was created on. + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`SparkContext` + The :class:`SparkContext` that this RDD was created on + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.context + + >>> rdd.context is sc + True + """ + return self.ctx + + def cache(self: "RDD[T]") -> "RDD[T]": + """ + Persist this RDD with the default storage level (`MEMORY_ONLY`). + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`RDD` + The same :class:`RDD` with storage level set to `MEMORY_ONLY` + + See Also + -------- + :meth:`RDD.persist` + :meth:`RDD.unpersist` + :meth:`RDD.getStorageLevel` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd2 = rdd.cache() + >>> rdd2 is rdd + True + >>> str(rdd.getStorageLevel()) + 'Memory Serialized 1x Replicated' + >>> _ = rdd.unpersist() + """ + self.is_cached = True + self.persist(StorageLevel.MEMORY_ONLY) + return self + + def persist(self: "RDD[T]", storageLevel: StorageLevel = StorageLevel.MEMORY_ONLY) -> "RDD[T]": + """ + Set this RDD's storage level to persist its values across operations + after the first time it is computed. This can only be used to assign + a new storage level if the RDD does not have a storage level set yet. + If no storage level is specified defaults to (`MEMORY_ONLY`). + + .. versionadded:: 0.9.1 + + Parameters + ---------- + storageLevel : :class:`StorageLevel`, default `MEMORY_ONLY` + the target storage level + + Returns + ------- + :class:`RDD` + The same :class:`RDD` with storage level set to `storageLevel`. + + See Also + -------- + :meth:`RDD.cache` + :meth:`RDD.unpersist` + :meth:`RDD.getStorageLevel` + + Examples + -------- + >>> rdd = sc.parallelize(["b", "a", "c"]) + >>> rdd.persist().is_cached + True + >>> str(rdd.getStorageLevel()) + 'Memory Serialized 1x Replicated' + >>> _ = rdd.unpersist() + >>> rdd.is_cached + False + + >>> from pyspark import StorageLevel + >>> rdd2 = sc.range(5) + >>> _ = rdd2.persist(StorageLevel.MEMORY_AND_DISK) + >>> rdd2.is_cached + True + >>> str(rdd2.getStorageLevel()) + 'Disk Memory Serialized 1x Replicated' + + Can not override existing storage level + + >>> _ = rdd2.persist(StorageLevel.MEMORY_ONLY_2) + Traceback (most recent call last): + ... + py4j.protocol.Py4JJavaError: ... + + Assign another storage level after `unpersist` + + >>> _ = rdd2.unpersist() + >>> rdd2.is_cached + False + >>> _ = rdd2.persist(StorageLevel.MEMORY_ONLY_2) + >>> str(rdd2.getStorageLevel()) + 'Memory Serialized 2x Replicated' + >>> rdd2.is_cached + True + >>> _ = rdd2.unpersist() + """ + self.is_cached = True + javaStorageLevel = self.ctx._getJavaStorageLevel(storageLevel) + self._jrdd.persist(javaStorageLevel) + return self + + def unpersist(self: "RDD[T]", blocking: bool = False) -> "RDD[T]": + """ + Mark the RDD as non-persistent, and remove all blocks for it from + memory and disk. + + .. versionadded:: 0.9.1 + + Parameters + ---------- + blocking : bool, optional, default False + whether to block until all blocks are deleted + + .. versionadded:: 3.0.0 + + Returns + ------- + :class:`RDD` + The same :class:`RDD` + + See Also + -------- + :meth:`RDD.cache` + :meth:`RDD.persist` + :meth:`RDD.getStorageLevel` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.is_cached + False + >>> _ = rdd.unpersist() + >>> rdd.is_cached + False + >>> _ = rdd.cache() + >>> rdd.is_cached + True + >>> _ = rdd.unpersist() + >>> rdd.is_cached + False + >>> _ = rdd.unpersist() + """ + self.is_cached = False + self._jrdd.unpersist(blocking) + return self + + def checkpoint(self) -> None: + """ + Mark this RDD for checkpointing. It will be saved to a file inside the + checkpoint directory set with :meth:`SparkContext.setCheckpointDir` and + all references to its parent RDDs will be removed. This function must + be called before any job has been executed on this RDD. It is strongly + recommended that this RDD is persisted in memory, otherwise saving it + on a file will require recomputation. + + .. versionadded:: 0.7.0 + + See Also + -------- + :meth:`RDD.isCheckpointed` + :meth:`RDD.getCheckpointFile` + :meth:`RDD.localCheckpoint` + :meth:`SparkContext.setCheckpointDir` + :meth:`SparkContext.getCheckpointDir` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.is_checkpointed + False + >>> rdd.getCheckpointFile() == None + True + + >>> rdd.checkpoint() + >>> rdd.is_checkpointed + True + >>> rdd.getCheckpointFile() == None + True + + >>> rdd.count() + 5 + >>> rdd.is_checkpointed + True + >>> rdd.getCheckpointFile() == None + False + """ + self.is_checkpointed = True + self._jrdd.rdd().checkpoint() + + def isCheckpointed(self) -> bool: + """ + Return whether this RDD is checkpointed and materialized, either reliably or locally. + + .. versionadded:: 0.7.0 + + Returns + ------- + bool + whether this :class:`RDD` is checkpointed and materialized, either reliably or locally + + See Also + -------- + :meth:`RDD.checkpoint` + :meth:`RDD.getCheckpointFile` + :meth:`SparkContext.setCheckpointDir` + :meth:`SparkContext.getCheckpointDir` + """ + return self._jrdd.rdd().isCheckpointed() + + def localCheckpoint(self) -> None: + """ + Mark this RDD for local checkpointing using Spark's existing caching layer. + + This method is for users who wish to truncate RDD lineages while skipping the expensive + step of replicating the materialized data in a reliable distributed file system. This is + useful for RDDs with long lineages that need to be truncated periodically (e.g. GraphX). + + Local checkpointing sacrifices fault-tolerance for performance. In particular, checkpointed + data is written to ephemeral local storage in the executors instead of to a reliable, + fault-tolerant storage. The effect is that if an executor fails during the computation, + the checkpointed data may no longer be accessible, causing an irrecoverable job failure. + + This is NOT safe to use with dynamic allocation, which removes executors along + with their cached blocks. If you must use both features, you are advised to set + `spark.dynamicAllocation.cachedExecutorIdleTimeout` to a high value. + + The checkpoint directory set through :meth:`SparkContext.setCheckpointDir` is not used. + + .. versionadded:: 2.2.0 + + See Also + -------- + :meth:`RDD.checkpoint` + :meth:`RDD.isLocallyCheckpointed` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.isLocallyCheckpointed() + False + + >>> rdd.localCheckpoint() + >>> rdd.isLocallyCheckpointed() + True + """ + self._jrdd.rdd().localCheckpoint() + + def isLocallyCheckpointed(self) -> bool: + """ + Return whether this RDD is marked for local checkpointing. + + Exposed for testing. + + .. versionadded:: 2.2.0 + + Returns + ------- + bool + whether this :class:`RDD` is marked for local checkpointing + + See Also + -------- + :meth:`RDD.localCheckpoint` + """ + return self._jrdd.rdd().isLocallyCheckpointed() + + def getCheckpointFile(self) -> Optional[str]: + """ + Gets the name of the file to which this RDD was checkpointed + + Not defined if RDD is checkpointed locally. + + .. versionadded:: 0.7.0 + + Returns + ------- + str + the name of the file to which this :class:`RDD` was checkpointed + + See Also + -------- + :meth:`RDD.checkpoint` + :meth:`SparkContext.setCheckpointDir` + :meth:`SparkContext.getCheckpointDir` + """ + checkpointFile = self._jrdd.rdd().getCheckpointFile() + + return checkpointFile.get() if checkpointFile.isDefined() else None + + def cleanShuffleDependencies(self, blocking: bool = False) -> None: + """ + Removes an RDD's shuffles and it's non-persisted ancestors. + + When running without a shuffle service, cleaning up shuffle files enables downscaling. + If you use the RDD after this call, you should checkpoint and materialize it first. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + blocking : bool, optional, default False + whether to block on shuffle cleanup tasks + + Notes + ----- + This API is a developer API. + """ + self._jrdd.rdd().cleanShuffleDependencies(blocking) + + def map(self: "RDD[T]", f: Callable[[T], U], preservesPartitioning: bool = False) -> "RDD[U]": + """ + Return a new RDD by applying a function to each element of this RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to run on each element of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to all elements + + See Also + -------- + :meth:`RDD.flatMap` + :meth:`RDD.mapPartitions` + :meth:`RDD.mapPartitionsWithIndex` + :meth:`RDD.mapPartitionsWithSplit` + + Examples + -------- + >>> rdd = sc.parallelize(["b", "a", "c"]) + >>> sorted(rdd.map(lambda x: (x, 1)).collect()) + [('a', 1), ('b', 1), ('c', 1)] + """ + + def func(_: int, iterator: Iterable[T]) -> Iterable[U]: + return map(fail_on_stopiteration(f), iterator) + + return self.mapPartitionsWithIndex(func, preservesPartitioning) + + def flatMap( + self: "RDD[T]", f: Callable[[T], Iterable[U]], preservesPartitioning: bool = False + ) -> "RDD[U]": + """ + Return a new RDD by first applying a function to all elements of this + RDD, and then flattening the results. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to turn a T into a sequence of U + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to all elements + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.mapPartitions` + :meth:`RDD.mapPartitionsWithIndex` + :meth:`RDD.mapPartitionsWithSplit` + + Examples + -------- + >>> rdd = sc.parallelize([2, 3, 4]) + >>> sorted(rdd.flatMap(lambda x: range(1, x)).collect()) + [1, 1, 1, 2, 2, 3] + >>> sorted(rdd.flatMap(lambda x: [(x, x), (x, x)]).collect()) + [(2, 2), (2, 2), (3, 3), (3, 3), (4, 4), (4, 4)] + """ + + def func(_: int, iterator: Iterable[T]) -> Iterable[U]: + return chain.from_iterable(map(fail_on_stopiteration(f), iterator)) + + return self.mapPartitionsWithIndex(func, preservesPartitioning) + + def mapPartitions( + self: "RDD[T]", f: Callable[[Iterable[T]], Iterable[U]], preservesPartitioning: bool = False + ) -> "RDD[U]": + """ + Return a new RDD by applying a function to each partition of this RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to run on each partition of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each partition + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.flatMap` + :meth:`RDD.mapPartitionsWithIndex` + :meth:`RDD.mapPartitionsWithSplit` + :meth:`RDDBarrier.mapPartitions` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 2) + >>> def f(iterator): yield sum(iterator) + ... + >>> rdd.mapPartitions(f).collect() + [3, 7] + """ + + def func(_: int, iterator: Iterable[T]) -> Iterable[U]: + return f(iterator) + + return self.mapPartitionsWithIndex(func, preservesPartitioning) + + def mapPartitionsWithIndex( + self: "RDD[T]", + f: Callable[[int, Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + ) -> "RDD[U]": + """ + Return a new RDD by applying a function to each partition of this RDD, + while tracking the index of the original partition. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to run on each partition of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each partition + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.flatMap` + :meth:`RDD.mapPartitions` + :meth:`RDD.mapPartitionsWithSplit` + :meth:`RDDBarrier.mapPartitionsWithIndex` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 4) + >>> def f(splitIndex, iterator): yield splitIndex + ... + >>> rdd.mapPartitionsWithIndex(f).sum() + 6 + """ + return PipelinedRDD(self, f, preservesPartitioning) + + def mapPartitionsWithSplit( + self: "RDD[T]", + f: Callable[[int, Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + ) -> "RDD[U]": + """ + Return a new RDD by applying a function to each partition of this RDD, + while tracking the index of the original partition. + + .. versionadded:: 0.7.0 + + .. deprecated:: 0.9.0 + use meth:`RDD.mapPartitionsWithIndex` instead. + + Parameters + ---------- + f : function + a function to run on each partition of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each partition + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.flatMap` + :meth:`RDD.mapPartitions` + :meth:`RDD.mapPartitionsWithIndex` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 4) + >>> def f(splitIndex, iterator): yield splitIndex + ... + >>> rdd.mapPartitionsWithSplit(f).sum() + 6 + """ + warnings.warn( + "mapPartitionsWithSplit is deprecated; use mapPartitionsWithIndex instead", + FutureWarning, + stacklevel=2, + ) + return self.mapPartitionsWithIndex(f, preservesPartitioning) + + def getNumPartitions(self) -> int: + """ + Returns the number of partitions in RDD + + .. versionadded:: 1.1.0 + + Returns + ------- + int + number of partitions + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 2) + >>> rdd.getNumPartitions() + 2 + """ + return self._jrdd.partitions().size() + + def filter(self: "RDD[T]", f: Callable[[T], bool]) -> "RDD[T]": + """ + Return a new RDD containing only the elements that satisfy a predicate. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to run on each element of the RDD + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each element + + See Also + -------- + :meth:`RDD.map` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4, 5]) + >>> rdd.filter(lambda x: x % 2 == 0).collect() + [2, 4] + """ + + def func(iterator: Iterable[T]) -> Iterable[T]: + return filter(fail_on_stopiteration(f), iterator) + + return self.mapPartitions(func, True) + + def distinct(self: "RDD[T]", numPartitions: Optional[int] = None) -> "RDD[T]": + """ + Return a new RDD containing the distinct elements in this RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a new :class:`RDD` containing the distinct elements + + See Also + -------- + :meth:`RDD.countApproxDistinct` + + Examples + -------- + >>> sorted(sc.parallelize([1, 1, 2, 3]).distinct().collect()) + [1, 2, 3] + """ + return ( + self.map(lambda x: (x, None)) + .reduceByKey(lambda x, _: x, numPartitions) + .map(lambda x: x[0]) + ) + + def sample( + self: "RDD[T]", withReplacement: bool, fraction: float, seed: Optional[int] = None + ) -> "RDD[T]": + """ + Return a sampled subset of this RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + withReplacement : bool + can elements be sampled multiple times (replaced when sampled out) + fraction : float + expected size of the sample as a fraction of this RDD's size + without replacement: probability that each element is chosen; fraction must be [0, 1] + with replacement: expected number of times each element is chosen; fraction must be >= 0 + seed : int, optional + seed for the random number generator + + Returns + ------- + :class:`RDD` + a new :class:`RDD` containing a sampled subset of elements + + See Also + -------- + :meth:`RDD.takeSample` + :meth:`RDD.sampleByKey` + :meth:`pyspark.sql.DataFrame.sample` + + Notes + ----- + This is not guaranteed to provide exactly the fraction specified of the total + count of the given :class:`DataFrame`. + + Examples + -------- + >>> rdd = sc.parallelize(range(100), 4) + >>> 6 <= rdd.sample(False, 0.1, 81).count() <= 14 + True + """ + if not fraction >= 0: + raise ValueError("Fraction must be nonnegative.") + return self.mapPartitionsWithIndex(RDDSampler(withReplacement, fraction, seed).func, True) + + def randomSplit( + self: "RDD[T]", weights: Sequence[Union[int, float]], seed: Optional[int] = None + ) -> "List[RDD[T]]": + """ + Randomly splits this RDD with the provided weights. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + weights : list + weights for splits, will be normalized if they don't sum to 1 + seed : int, optional + random seed + + Returns + ------- + list + split :class:`RDD`\\s in a list + + See Also + -------- + :meth:`pyspark.sql.DataFrame.randomSplit` + + Examples + -------- + >>> rdd = sc.parallelize(range(500), 1) + >>> rdd1, rdd2 = rdd.randomSplit([2, 3], 17) + >>> len(rdd1.collect() + rdd2.collect()) + 500 + >>> 150 < rdd1.count() < 250 + True + >>> 250 < rdd2.count() < 350 + True + """ + if not all(w >= 0 for w in weights): + raise ValueError("Weights must be nonnegative") + s = float(sum(weights)) + if not s > 0: + raise ValueError("Sum of weights must be positive") + cweights = [0.0] + for w in weights: + cweights.append(cweights[-1] + w / s) + if seed is None: + seed = random.randint(0, 2**32 - 1) + return [ + self.mapPartitionsWithIndex(RDDRangeSampler(lb, ub, seed).func, True) + for lb, ub in zip(cweights, cweights[1:]) + ] + + # this is ported from scala/spark/RDD.scala + def takeSample( + self: "RDD[T]", withReplacement: bool, num: int, seed: Optional[int] = None + ) -> List[T]: + """ + Return a fixed-size sampled subset of this RDD. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + withReplacement : list + whether sampling is done with replacement + num : int + size of the returned sample + seed : int, optional + random seed + + Returns + ------- + list + a fixed-size sampled subset of this :class:`RDD` in an array + + See Also + -------- + :meth:`RDD.sample` + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> import sys + >>> rdd = sc.parallelize(range(0, 10)) + >>> len(rdd.takeSample(True, 20, 1)) + 20 + >>> len(rdd.takeSample(False, 5, 2)) + 5 + >>> len(rdd.takeSample(False, 15, 3)) + 10 + >>> sc.range(0, 10).takeSample(False, sys.maxsize) + Traceback (most recent call last): + ... + ValueError: Sample size cannot be greater than ... + """ + numStDev = 10.0 + maxSampleSize = sys.maxsize - int(numStDev * sqrt(sys.maxsize)) + if num < 0: + raise ValueError("Sample size cannot be negative.") + elif num > maxSampleSize: + raise ValueError("Sample size cannot be greater than %d." % maxSampleSize) + + if num == 0 or self.getNumPartitions() == 0: + return [] + + initialCount = self.count() + if initialCount == 0: + return [] + + rand = random.Random(seed) + + if (not withReplacement) and num >= initialCount: + # shuffle current RDD and return + samples = self.collect() + rand.shuffle(samples) + return samples + + fraction = RDD._computeFractionForSampleSize(num, initialCount, withReplacement) + samples = self.sample(withReplacement, fraction, seed).collect() + + # If the first sample didn't turn out large enough, keep trying to take samples; + # this shouldn't happen often because we use a big multiplier for their initial size. + # See: scala/spark/RDD.scala + while len(samples) < num: + # TODO: add log warning for when more than one iteration was run + seed = rand.randint(0, sys.maxsize) + samples = self.sample(withReplacement, fraction, seed).collect() + + rand.shuffle(samples) + + return samples[0:num] + + @staticmethod + def _computeFractionForSampleSize( + sampleSizeLowerBound: int, total: int, withReplacement: bool + ) -> float: + """ + Returns a sampling rate that guarantees a sample of + size >= sampleSizeLowerBound 99.99% of the time. + + How the sampling rate is determined: + Let p = num / total, where num is the sample size and total is the + total number of data points in the RDD. We're trying to compute + q > p such that + - when sampling with replacement, we're drawing each data point + with prob_i ~ Pois(q), where we want to guarantee + Pr[s < num] < 0.0001 for s = sum(prob_i for i from 0 to + total), i.e. the failure rate of not having a sufficiently large + sample < 0.0001. Setting q = p + 5 * sqrt(p/total) is sufficient + to guarantee 0.9999 success rate for num > 12, but we need a + slightly larger q (9 empirically determined). + - when sampling without replacement, we're drawing each data point + with prob_i ~ Binomial(total, fraction) and our choice of q + guarantees 1-delta, or 0.9999 success rate, where success rate is + defined the same as in sampling with replacement. + """ + fraction = float(sampleSizeLowerBound) / total + if withReplacement: + numStDev = 5 + if sampleSizeLowerBound < 12: + numStDev = 9 + return fraction + numStDev * sqrt(fraction / total) + else: + delta = 0.00005 + gamma = -log(delta) / total + return min(1, fraction + gamma + sqrt(gamma * gamma + 2 * gamma * fraction)) + + def union(self: "RDD[T]", other: "RDD[U]") -> "RDD[Union[T, U]]": + """ + Return the union of this RDD and another one. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + + Returns + ------- + :class:`RDD` + the union of this :class:`RDD` and another one + + See Also + -------- + :meth:`SparkContext.union` + :meth:`pyspark.sql.DataFrame.union` + + Examples + -------- + >>> rdd = sc.parallelize([1, 1, 2, 3]) + >>> rdd.union(rdd).collect() + [1, 1, 2, 3, 1, 1, 2, 3] + """ + if self._jrdd_deserializer == other._jrdd_deserializer: + rdd: "RDD[Union[T, U]]" = RDD( + self._jrdd.union(other._jrdd), self.ctx, self._jrdd_deserializer + ) + else: + # These RDDs contain data in different serialized formats, so we + # must normalize them to the default serializer. + self_copy = self._reserialize() + other_copy = other._reserialize() + rdd = RDD(self_copy._jrdd.union(other_copy._jrdd), self.ctx, self.ctx.serializer) + if ( + self.partitioner == other.partitioner + and self.getNumPartitions() == rdd.getNumPartitions() + ): + rdd.partitioner = self.partitioner + return rdd + + def intersection(self: "RDD[T]", other: "RDD[T]") -> "RDD[T]": + """ + Return the intersection of this RDD and another one. The output will + not contain any duplicate elements, even if the input RDDs did. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + + Returns + ------- + :class:`RDD` + the intersection of this :class:`RDD` and another one + + See Also + -------- + :meth:`pyspark.sql.DataFrame.intersect` + + Notes + ----- + This method performs a shuffle internally. + + Examples + -------- + >>> rdd1 = sc.parallelize([1, 10, 2, 3, 4, 5]) + >>> rdd2 = sc.parallelize([1, 6, 2, 3, 7, 8]) + >>> rdd1.intersection(rdd2).collect() + [1, 2, 3] + """ + return ( + self.map(lambda v: (v, None)) + .cogroup(other.map(lambda v: (v, None))) + .filter(lambda k_vs: all(k_vs[1])) + .keys() + ) + + def _reserialize(self: "RDD[T]", serializer: Optional[Serializer] = None) -> "RDD[T]": + serializer = serializer or self.ctx.serializer + if self._jrdd_deserializer != serializer: + self = self.map(lambda x: x, preservesPartitioning=True) + self._jrdd_deserializer = serializer + return self + + def __add__(self: "RDD[T]", other: "RDD[U]") -> "RDD[Union[T, U]]": + """ + Return the union of this RDD and another one. + + Examples + -------- + >>> rdd = sc.parallelize([1, 1, 2, 3]) + >>> (rdd + rdd).collect() + [1, 1, 2, 3, 1, 1, 2, 3] + """ + if not isinstance(other, RDD): + raise TypeError + return self.union(other) + + @overload + def repartitionAndSortWithinPartitions( + self: "RDD[Tuple[S, V]]", + numPartitions: Optional[int] = ..., + partitionFunc: Callable[["S"], int] = ..., + ascending: bool = ..., + ) -> "RDD[Tuple[S, V]]": + ... + + @overload + def repartitionAndSortWithinPartitions( + self: "RDD[Tuple[K, V]]", + numPartitions: Optional[int], + partitionFunc: Callable[[K], int], + ascending: bool, + keyfunc: Callable[[K], "S"], + ) -> "RDD[Tuple[K, V]]": + ... + + @overload + def repartitionAndSortWithinPartitions( + self: "RDD[Tuple[K, V]]", + numPartitions: Optional[int] = ..., + partitionFunc: Callable[[K], int] = ..., + ascending: bool = ..., + *, + keyfunc: Callable[[K], "S"], + ) -> "RDD[Tuple[K, V]]": + ... + + def repartitionAndSortWithinPartitions( + self: "RDD[Tuple[Any, Any]]", + numPartitions: Optional[int] = None, + partitionFunc: Callable[[Any], int] = portable_hash, + ascending: bool = True, + keyfunc: Callable[[Any], Any] = lambda x: x, + ) -> "RDD[Tuple[Any, Any]]": + """ + Repartition the RDD according to the given partitioner and, within each resulting partition, + sort records by their keys. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + a function to compute the partition index + ascending : bool, optional, default True + sort the keys in ascending or descending order + keyfunc : function, optional, default identity mapping + a function to compute the key + + Returns + ------- + :class:`RDD` + a new :class:`RDD` + + See Also + -------- + :meth:`RDD.repartition` + :meth:`RDD.partitionBy` + :meth:`RDD.sortBy` + :meth:`RDD.sortByKey` + + Examples + -------- + >>> rdd = sc.parallelize([(0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)]) + >>> rdd2 = rdd.repartitionAndSortWithinPartitions(2, lambda x: x % 2, True) + >>> rdd2.glom().collect() + [[(0, 5), (0, 8), (2, 6)], [(1, 3), (3, 8), (3, 8)]] + """ + if numPartitions is None: + numPartitions = self._defaultReducePartitions() + + memory = self._memory_limit() + serializer = self._jrdd_deserializer + + def sortPartition(iterator: Iterable[Tuple[K, V]]) -> Iterable[Tuple[K, V]]: + sort = ExternalSorter(memory * 0.9, serializer).sorted + return iter(sort(iterator, key=lambda k_v: keyfunc(k_v[0]), reverse=(not ascending))) + + return self.partitionBy(numPartitions, partitionFunc).mapPartitions(sortPartition, True) + + @overload + def sortByKey( + self: "RDD[Tuple[S, V]]", + ascending: bool = ..., + numPartitions: Optional[int] = ..., + ) -> "RDD[Tuple[K, V]]": + ... + + @overload + def sortByKey( + self: "RDD[Tuple[K, V]]", + ascending: bool, + numPartitions: int, + keyfunc: Callable[[K], "S"], + ) -> "RDD[Tuple[K, V]]": + ... + + @overload + def sortByKey( + self: "RDD[Tuple[K, V]]", + ascending: bool = ..., + numPartitions: Optional[int] = ..., + *, + keyfunc: Callable[[K], "S"], + ) -> "RDD[Tuple[K, V]]": + ... + + def sortByKey( + self: "RDD[Tuple[K, V]]", + ascending: Optional[bool] = True, + numPartitions: Optional[int] = None, + keyfunc: Callable[[Any], Any] = lambda x: x, + ) -> "RDD[Tuple[K, V]]": + """ + Sorts this RDD, which is assumed to consist of (key, value) pairs. + + .. versionadded:: 0.9.1 + + Parameters + ---------- + ascending : bool, optional, default True + sort the keys in ascending or descending order + numPartitions : int, optional + the number of partitions in new :class:`RDD` + keyfunc : function, optional, default identity mapping + a function to compute the key + + Returns + ------- + :class:`RDD` + a new :class:`RDD` + + See Also + -------- + :meth:`RDD.sortBy` + :meth:`pyspark.sql.DataFrame.sort` + + Examples + -------- + >>> tmp = [('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)] + >>> sc.parallelize(tmp).sortByKey().first() + ('1', 3) + >>> sc.parallelize(tmp).sortByKey(True, 1).collect() + [('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)] + >>> sc.parallelize(tmp).sortByKey(True, 2).collect() + [('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)] + >>> tmp2 = [('Mary', 1), ('had', 2), ('a', 3), ('little', 4), ('lamb', 5)] + >>> tmp2.extend([('whose', 6), ('fleece', 7), ('was', 8), ('white', 9)]) + >>> sc.parallelize(tmp2).sortByKey(True, 3, keyfunc=lambda k: k.lower()).collect() + [('a', 3), ('fleece', 7), ('had', 2), ('lamb', 5),...('white', 9), ('whose', 6)] + """ + if numPartitions is None: + numPartitions = self._defaultReducePartitions() + + memory = self._memory_limit() + serializer = self._jrdd_deserializer + + def sortPartition(iterator: Iterable[Tuple[K, V]]) -> Iterable[Tuple[K, V]]: + sort = ExternalSorter(memory * 0.9, serializer).sorted + return iter(sort(iterator, key=lambda kv: keyfunc(kv[0]), reverse=(not ascending))) + + if numPartitions == 1: + if self.getNumPartitions() > 1: + self = self.coalesce(1) + return self.mapPartitions(sortPartition, True) + + # first compute the boundary of each part via sampling: we want to partition + # the key-space into bins such that the bins have roughly the same + # number of (key, value) pairs falling into them + rddSize = self.count() + if not rddSize: + return self # empty RDD + maxSampleSize = numPartitions * 20.0 # constant from Spark's RangePartitioner + fraction = min(maxSampleSize / max(rddSize, 1), 1.0) + samples = self.sample(False, fraction, 1).map(lambda kv: kv[0]).collect() + samples = sorted(samples, key=keyfunc) + + # we have numPartitions many parts but one of the them has + # an implicit boundary + bounds = [ + samples[int(len(samples) * (i + 1) / numPartitions)] + for i in range(0, numPartitions - 1) + ] + + def rangePartitioner(k: K) -> int: + p = bisect.bisect_left(bounds, keyfunc(k)) + if ascending: + return p + else: + return numPartitions - 1 - p # type: ignore[operator] + + return self.partitionBy(numPartitions, rangePartitioner).mapPartitions(sortPartition, True) + + def sortBy( + self: "RDD[T]", + keyfunc: Callable[[T], "S"], + ascending: bool = True, + numPartitions: Optional[int] = None, + ) -> "RDD[T]": + """ + Sorts this RDD by the given keyfunc + + .. versionadded:: 1.1.0 + + Parameters + ---------- + keyfunc : function + a function to compute the key + ascending : bool, optional, default True + sort the keys in ascending or descending order + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a new :class:`RDD` + + See Also + -------- + :meth:`RDD.sortByKey` + :meth:`pyspark.sql.DataFrame.sort` + + Examples + -------- + >>> tmp = [('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)] + >>> sc.parallelize(tmp).sortBy(lambda x: x[0]).collect() + [('1', 3), ('2', 5), ('a', 1), ('b', 2), ('d', 4)] + >>> sc.parallelize(tmp).sortBy(lambda x: x[1]).collect() + [('a', 1), ('b', 2), ('1', 3), ('d', 4), ('2', 5)] + """ + return ( + self.keyBy(keyfunc) # type: ignore[type-var] + .sortByKey(ascending, numPartitions) + .values() + ) + + def glom(self: "RDD[T]") -> "RDD[List[T]]": + """ + Return an RDD created by coalescing all elements within each partition + into a list. + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`RDD` + a new :class:`RDD` coalescing all elements within each partition into a list + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 2) + >>> sorted(rdd.glom().collect()) + [[1, 2], [3, 4]] + """ + + def func(iterator: Iterable[T]) -> Iterable[List[T]]: + yield list(iterator) + + return self.mapPartitions(func) + + def cartesian(self: "RDD[T]", other: "RDD[U]") -> "RDD[Tuple[T, U]]": + """ + Return the Cartesian product of this RDD and another one, that is, the + RDD of all pairs of elements ``(a, b)`` where ``a`` is in `self` and + ``b`` is in `other`. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + + Returns + ------- + :class:`RDD` + the Cartesian product of this :class:`RDD` and another one + + See Also + -------- + :meth:`pyspark.sql.DataFrame.crossJoin` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2]) + >>> sorted(rdd.cartesian(rdd).collect()) + [(1, 1), (1, 2), (2, 1), (2, 2)] + """ + # Due to batching, we can't use the Java cartesian method. + deserializer = CartesianDeserializer(self._jrdd_deserializer, other._jrdd_deserializer) + return RDD(self._jrdd.cartesian(other._jrdd), self.ctx, deserializer) + + def groupBy( + self: "RDD[T]", + f: Callable[[T], K], + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, Iterable[T]]]": + """ + Return an RDD of grouped items. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to compute the key + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + a function to compute the partition index + + Returns + ------- + :class:`RDD` + a new :class:`RDD` of grouped items + + See Also + -------- + :meth:`RDD.groupByKey` + :meth:`pyspark.sql.DataFrame.groupBy` + + Examples + -------- + >>> rdd = sc.parallelize([1, 1, 2, 3, 5, 8]) + >>> result = rdd.groupBy(lambda x: x % 2).collect() + >>> sorted([(x, sorted(y)) for (x, y) in result]) + [(0, [2, 8]), (1, [1, 1, 3, 5])] + """ + return self.map(lambda x: (f(x), x)).groupByKey(numPartitions, partitionFunc) + + def pipe( + self, command: str, env: Optional[Dict[str, str]] = None, checkCode: bool = False + ) -> "RDD[str]": + """ + Return an RDD created by piping elements to a forked external process. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + command : str + command to run. + env : dict, optional + environment variables to set. + checkCode : bool, optional + whether to check the return value of the shell command. + + Returns + ------- + :class:`RDD` + a new :class:`RDD` of strings + + Examples + -------- + >>> sc.parallelize(['1', '2', '', '3']).pipe('cat').collect() + ['1', '2', '', '3'] + """ + if env is None: + env = dict() + + def func(iterator: Iterable[T]) -> Iterable[str]: + pipe = Popen(shlex.split(command), env=env, stdin=PIPE, stdout=PIPE) + + def pipe_objs(out: IO[bytes]) -> None: + for obj in iterator: + s = str(obj).rstrip("\n") + "\n" + out.write(s.encode("utf-8")) + out.close() + + Thread(target=pipe_objs, args=[pipe.stdin]).start() + + def check_return_code() -> Iterable[int]: + pipe.wait() + if checkCode and pipe.returncode: + raise PySparkRuntimeError( + error_class="PIPE_FUNCTION_EXITED", + message_parameters={ + "func_name": command, + "error_code": str(pipe.returncode), + }, + ) + else: + for i in range(0): + yield i + + return ( + cast(bytes, x).rstrip(b"\n").decode("utf-8") + for x in chain( + iter(cast(IO[bytes], pipe.stdout).readline, b""), check_return_code() + ) + ) + + return self.mapPartitions(func) + + def foreach(self: "RDD[T]", f: Callable[[T], None]) -> None: + """ + Applies a function to all elements of this RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function applied to each element + + See Also + -------- + :meth:`RDD.foreachPartition` + :meth:`pyspark.sql.DataFrame.foreach` + :meth:`pyspark.sql.DataFrame.foreachPartition` + + Examples + -------- + >>> def f(x): print(x) + ... + >>> sc.parallelize([1, 2, 3, 4, 5]).foreach(f) + """ + f = fail_on_stopiteration(f) + + def processPartition(iterator: Iterable[T]) -> Iterable[Any]: + for x in iterator: + f(x) + return iter([]) + + self.mapPartitions(processPartition).count() # Force evaluation + + def foreachPartition(self: "RDD[T]", f: Callable[[Iterable[T]], None]) -> None: + """ + Applies a function to each partition of this RDD. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + f : function + a function applied to each partition + + See Also + -------- + :meth:`RDD.foreach` + :meth:`pyspark.sql.DataFrame.foreach` + :meth:`pyspark.sql.DataFrame.foreachPartition` + + Examples + -------- + >>> def f(iterator): + ... for x in iterator: + ... print(x) + ... + >>> sc.parallelize([1, 2, 3, 4, 5]).foreachPartition(f) + """ + + def func(it: Iterable[T]) -> Iterable[Any]: + r = f(it) + try: + return iter(r) # type: ignore[call-overload] + except TypeError: + return iter([]) + + self.mapPartitions(func).count() # Force evaluation + + def collect(self: "RDD[T]") -> List[T]: + """ + Return a list that contains all the elements in this RDD. + + .. versionadded:: 0.7.0 + + Returns + ------- + list + a list containing all the elements + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + See Also + -------- + :meth:`RDD.toLocalIterator` + :meth:`pyspark.sql.DataFrame.collect` + + Examples + -------- + >>> sc.range(5).collect() + [0, 1, 2, 3, 4] + >>> sc.parallelize(["x", "y", "z"]).collect() + ['x', 'y', 'z'] + """ + with SCCallSiteSync(self.context): + assert self.ctx._jvm is not None + sock_info = self.ctx._jvm.PythonRDD.collectAndServe(self._jrdd.rdd()) + return list(_load_from_socket(sock_info, self._jrdd_deserializer)) + + def collectWithJobGroup( + self: "RDD[T]", groupId: str, description: str, interruptOnCancel: bool = False + ) -> "List[T]": + """ + When collect rdd, use this method to specify job group. + + .. versionadded:: 3.0.0 + + .. deprecated:: 3.1.0 + Use :class:`pyspark.InheritableThread` with the pinned thread mode enabled. + + Parameters + ---------- + groupId : str + The group ID to assign. + description : str + The description to set for the job group. + interruptOnCancel : bool, optional, default False + whether to interrupt jobs on job cancellation. + + Returns + ------- + list + a list containing all the elements + + See Also + -------- + :meth:`RDD.collect` + :meth:`SparkContext.setJobGroup` + """ + warnings.warn( + "Deprecated in 3.1, Use pyspark.InheritableThread with " + "the pinned thread mode enabled.", + FutureWarning, + ) + + with SCCallSiteSync(self.context): + assert self.ctx._jvm is not None + sock_info = self.ctx._jvm.PythonRDD.collectAndServeWithJobGroup( + self._jrdd.rdd(), groupId, description, interruptOnCancel + ) + return list(_load_from_socket(sock_info, self._jrdd_deserializer)) + + def reduce(self: "RDD[T]", f: Callable[[T, T], T]) -> T: + """ + Reduces the elements of this RDD using the specified commutative and + associative binary operator. Currently reduces partitions locally. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + the reduce function + + Returns + ------- + T + the aggregated result + + See Also + -------- + :meth:`RDD.treeReduce` + :meth:`RDD.aggregate` + :meth:`RDD.treeAggregate` + + Examples + -------- + >>> from operator import add + >>> sc.parallelize([1, 2, 3, 4, 5]).reduce(add) + 15 + >>> sc.parallelize((2 for _ in range(10))).map(lambda x: 1).cache().reduce(add) + 10 + >>> sc.parallelize([]).reduce(add) + Traceback (most recent call last): + ... + ValueError: Can not reduce() empty RDD + """ + f = fail_on_stopiteration(f) + + def func(iterator: Iterable[T]) -> Iterable[T]: + iterator = iter(iterator) + try: + initial = next(iterator) + except StopIteration: + return + yield reduce(f, iterator, initial) + + vals = self.mapPartitions(func).collect() + if vals: + return reduce(f, vals) + raise ValueError("Can not reduce() empty RDD") + + def treeReduce(self: "RDD[T]", f: Callable[[T, T], T], depth: int = 2) -> T: + """ + Reduces the elements of this RDD in a multi-level tree pattern. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + f : function + the reduce function + depth : int, optional, default 2 + suggested depth of the tree (default: 2) + + Returns + ------- + T + the aggregated result + + See Also + -------- + :meth:`RDD.reduce` + :meth:`RDD.aggregate` + :meth:`RDD.treeAggregate` + + Examples + -------- + >>> add = lambda x, y: x + y + >>> rdd = sc.parallelize([-5, -4, -3, -2, -1, 1, 2, 3, 4], 10) + >>> rdd.treeReduce(add) + -5 + >>> rdd.treeReduce(add, 1) + -5 + >>> rdd.treeReduce(add, 2) + -5 + >>> rdd.treeReduce(add, 5) + -5 + >>> rdd.treeReduce(add, 10) + -5 + """ + if depth < 1: + raise ValueError("Depth cannot be smaller than 1 but got %d." % depth) + + # Use the second entry to indicate whether this is a dummy value. + zeroValue: Tuple[T, bool] = ( # type: ignore[assignment] + None, + True, + ) + + def op(x: Tuple[T, bool], y: Tuple[T, bool]) -> Tuple[T, bool]: + if x[1]: + return y + elif y[1]: + return x + else: + return f(x[0], y[0]), False + + reduced = self.map(lambda x: (x, False)).treeAggregate(zeroValue, op, op, depth) + if reduced[1]: + raise ValueError("Cannot reduce empty RDD.") + return reduced[0] + + def fold(self: "RDD[T]", zeroValue: T, op: Callable[[T, T], T]) -> T: + """ + Aggregate the elements of each partition, and then the results for all + the partitions, using a given associative function and a neutral "zero value." + + The function ``op(t1, t2)`` is allowed to modify ``t1`` and return it + as its result value to avoid object allocation; however, it should not + modify ``t2``. + + This behaves somewhat differently from fold operations implemented + for non-distributed collections in functional languages like Scala. + This fold operation may be applied to partitions individually, and then + fold those results into the final result, rather than apply the fold + to each element sequentially in some defined ordering. For functions + that are not commutative, the result may differ from that of a fold + applied to a non-distributed collection. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + zeroValue : T + the initial value for the accumulated result of each partition + op : function + a function used to both accumulate results within a partition and combine + results from different partitions + + Returns + ------- + T + the aggregated result + + See Also + -------- + :meth:`RDD.reduce` + :meth:`RDD.aggregate` + + Examples + -------- + >>> from operator import add + >>> sc.parallelize([1, 2, 3, 4, 5]).fold(0, add) + 15 + """ + op = fail_on_stopiteration(op) + + def func(iterator: Iterable[T]) -> Iterable[T]: + acc = zeroValue + for obj in iterator: + acc = op(acc, obj) + yield acc + + # collecting result of mapPartitions here ensures that the copy of + # zeroValue provided to each partition is unique from the one provided + # to the final reduce call + vals = self.mapPartitions(func).collect() + return reduce(op, vals, zeroValue) + + def aggregate( + self: "RDD[T]", zeroValue: U, seqOp: Callable[[U, T], U], combOp: Callable[[U, U], U] + ) -> U: + """ + Aggregate the elements of each partition, and then the results for all + the partitions, using a given combine functions and a neutral "zero + value." + + The functions ``op(t1, t2)`` is allowed to modify ``t1`` and return it + as its result value to avoid object allocation; however, it should not + modify ``t2``. + + The first function (seqOp) can return a different result type, U, than + the type of this RDD. Thus, we need one operation for merging a T into + an U and one operation for merging two U + + .. versionadded:: 1.1.0 + + Parameters + ---------- + zeroValue : U + the initial value for the accumulated result of each partition + seqOp : function + a function used to accumulate results within a partition + combOp : function + an associative function used to combine results from different partitions + + Returns + ------- + U + the aggregated result + + See Also + -------- + :meth:`RDD.reduce` + :meth:`RDD.fold` + + Examples + -------- + >>> seqOp = (lambda x, y: (x[0] + y, x[1] + 1)) + >>> combOp = (lambda x, y: (x[0] + y[0], x[1] + y[1])) + >>> sc.parallelize([1, 2, 3, 4]).aggregate((0, 0), seqOp, combOp) + (10, 4) + >>> sc.parallelize([]).aggregate((0, 0), seqOp, combOp) + (0, 0) + """ + seqOp = fail_on_stopiteration(seqOp) + combOp = fail_on_stopiteration(combOp) + + def func(iterator: Iterable[T]) -> Iterable[U]: + acc = zeroValue + for obj in iterator: + acc = seqOp(acc, obj) + yield acc + + # collecting result of mapPartitions here ensures that the copy of + # zeroValue provided to each partition is unique from the one provided + # to the final reduce call + vals = self.mapPartitions(func).collect() + return reduce(combOp, vals, zeroValue) + + def treeAggregate( + self: "RDD[T]", + zeroValue: U, + seqOp: Callable[[U, T], U], + combOp: Callable[[U, U], U], + depth: int = 2, + ) -> U: + """ + Aggregates the elements of this RDD in a multi-level tree + pattern. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + zeroValue : U + the initial value for the accumulated result of each partition + seqOp : function + a function used to accumulate results within a partition + combOp : function + an associative function used to combine results from different partitions + depth : int, optional, default 2 + suggested depth of the tree + + Returns + ------- + U + the aggregated result + + See Also + -------- + :meth:`RDD.aggregate` + :meth:`RDD.treeReduce` + + Examples + -------- + >>> add = lambda x, y: x + y + >>> rdd = sc.parallelize([-5, -4, -3, -2, -1, 1, 2, 3, 4], 10) + >>> rdd.treeAggregate(0, add, add) + -5 + >>> rdd.treeAggregate(0, add, add, 1) + -5 + >>> rdd.treeAggregate(0, add, add, 2) + -5 + >>> rdd.treeAggregate(0, add, add, 5) + -5 + >>> rdd.treeAggregate(0, add, add, 10) + -5 + """ + if depth < 1: + raise ValueError("Depth cannot be smaller than 1 but got %d." % depth) + + if self.getNumPartitions() == 0: + return zeroValue + + def aggregatePartition(iterator: Iterable[T]) -> Iterable[U]: + acc = zeroValue + for obj in iterator: + acc = seqOp(acc, obj) + yield acc + + partiallyAggregated = self.mapPartitions(aggregatePartition) + numPartitions = partiallyAggregated.getNumPartitions() + scale = max(int(ceil(pow(numPartitions, 1.0 / depth))), 2) + # If creating an extra level doesn't help reduce the wall-clock time, we stop the tree + # aggregation. + while numPartitions > scale + numPartitions / scale: + numPartitions /= scale # type: ignore[assignment] + curNumPartitions = int(numPartitions) + + def mapPartition(i: int, iterator: Iterable[U]) -> Iterable[Tuple[int, U]]: + for obj in iterator: + yield (i % curNumPartitions, obj) + + partiallyAggregated = ( + partiallyAggregated.mapPartitionsWithIndex(mapPartition) + .reduceByKey(combOp, curNumPartitions) + .values() + ) + + return partiallyAggregated.reduce(combOp) + + @overload + def max(self: "RDD[S]") -> "S": + ... + + @overload + def max(self: "RDD[T]", key: Callable[[T], "S"]) -> T: + ... + + def max(self: "RDD[T]", key: Optional[Callable[[T], "S"]] = None) -> T: + """ + Find the maximum item in this RDD. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + key : function, optional + A function used to generate key for comparing + + Returns + ------- + T + the maximum item + + See Also + -------- + :meth:`RDD.min` + + Examples + -------- + >>> rdd = sc.parallelize([1.0, 5.0, 43.0, 10.0]) + >>> rdd.max() + 43.0 + >>> rdd.max(key=str) + 5.0 + """ + if key is None: + return self.reduce(max) # type: ignore[arg-type] + return self.reduce(lambda a, b: max(a, b, key=key)) # type: ignore[arg-type] + + @overload + def min(self: "RDD[S]") -> "S": + ... + + @overload + def min(self: "RDD[T]", key: Callable[[T], "S"]) -> T: + ... + + def min(self: "RDD[T]", key: Optional[Callable[[T], "S"]] = None) -> T: + """ + Find the minimum item in this RDD. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + key : function, optional + A function used to generate key for comparing + + Returns + ------- + T + the minimum item + + See Also + -------- + :meth:`RDD.max` + + Examples + -------- + >>> rdd = sc.parallelize([2.0, 5.0, 43.0, 10.0]) + >>> rdd.min() + 2.0 + >>> rdd.min(key=str) + 10.0 + """ + if key is None: + return self.reduce(min) # type: ignore[arg-type] + return self.reduce(lambda a, b: min(a, b, key=key)) # type: ignore[arg-type] + + def sum(self: "RDD[NumberOrArray]") -> "NumberOrArray": + """ + Add up the elements in this RDD. + + .. versionadded:: 0.7.0 + + Returns + ------- + float, int, or complex + the sum of all elements + + See Also + -------- + :meth:`RDD.mean` + :meth:`RDD.sumApprox` + + Examples + -------- + >>> sc.parallelize([1.0, 2.0, 3.0]).sum() + 6.0 + """ + return self.mapPartitions(lambda x: [sum(x)]).fold( # type: ignore[return-value] + 0, operator.add + ) + + def count(self) -> int: + """ + Return the number of elements in this RDD. + + .. versionadded:: 0.7.0 + + Returns + ------- + int + the number of elements + + See Also + -------- + :meth:`RDD.countApprox` + :meth:`pyspark.sql.DataFrame.count` + + Examples + -------- + >>> sc.parallelize([2, 3, 4]).count() + 3 + """ + return self.mapPartitions(lambda i: [sum(1 for _ in i)]).sum() + + def stats(self: "RDD[NumberOrArray]") -> StatCounter: + """ + Return a :class:`StatCounter` object that captures the mean, variance + and count of the RDD's elements in one operation. + + .. versionadded:: 0.9.1 + + Returns + ------- + :class:`StatCounter` + a :class:`StatCounter` capturing the mean, variance and count of all elements + + See Also + -------- + :meth:`RDD.stdev` + :meth:`RDD.sampleStdev` + :meth:`RDD.variance` + :meth:`RDD.sampleVariance` + :meth:`RDD.histogram` + :meth:`pyspark.sql.DataFrame.stat` + """ + + def redFunc(left_counter: StatCounter, right_counter: StatCounter) -> StatCounter: + return left_counter.mergeStats(right_counter) + + return self.mapPartitions(lambda i: [StatCounter(i)]).reduce( # type: ignore[arg-type] + redFunc + ) + + def histogram( + self: "RDD[S]", buckets: Union[int, List["S"], Tuple["S", ...]] + ) -> Tuple[Sequence["S"], List[int]]: + """ + Compute a histogram using the provided buckets. The buckets + are all open to the right except for the last which is closed. + e.g. [1,10,20,50] means the buckets are [1,10) [10,20) [20,50], + which means 1<=x<10, 10<=x<20, 20<=x<=50. And on the input of 1 + and 50 we would have a histogram of 1,0,1. + + If your histogram is evenly spaced (e.g. [0, 10, 20, 30]), + this can be switched from an O(log n) insertion to O(1) per + element (where n is the number of buckets). + + Buckets must be sorted, not contain any duplicates, and have + at least two elements. + + If `buckets` is a number, it will generate buckets which are + evenly spaced between the minimum and maximum of the RDD. For + example, if the min value is 0 and the max is 100, given `buckets` + as 2, the resulting buckets will be [0,50) [50,100]. `buckets` must + be at least 1. An exception is raised if the RDD contains infinity. + If the elements in the RDD do not vary (max == min), a single bucket + will be used. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + buckets : int, or list, or tuple + if `buckets` is a number, it computes a histogram of the data using + `buckets` number of buckets evenly, otherwise, `buckets` is the provided + buckets to bin the data. + + Returns + ------- + tuple + a tuple of buckets and histogram + + See Also + -------- + :meth:`RDD.stats` + + Examples + -------- + >>> rdd = sc.parallelize(range(51)) + >>> rdd.histogram(2) + ([0, 25, 50], [25, 26]) + >>> rdd.histogram([0, 5, 25, 50]) + ([0, 5, 25, 50], [5, 20, 26]) + >>> rdd.histogram([0, 15, 30, 45, 60]) # evenly spaced buckets + ([0, 15, 30, 45, 60], [15, 15, 15, 6]) + >>> rdd = sc.parallelize(["ab", "ac", "b", "bd", "ef"]) + >>> rdd.histogram(("a", "b", "c")) + (('a', 'b', 'c'), [2, 2]) + """ + if isinstance(buckets, int): + if buckets < 1: + raise ValueError("number of buckets must be >= 1") + + # filter out non-comparable elements + def comparable(x: Any) -> bool: + if x is None: + return False + if type(x) is float and isnan(x): + return False + return True + + filtered = self.filter(comparable) + + # faster than stats() + def minmax(a: Tuple["S", "S"], b: Tuple["S", "S"]) -> Tuple["S", "S"]: + return min(a[0], b[0]), max(a[1], b[1]) + + try: + minv, maxv = filtered.map(lambda x: (x, x)).reduce(minmax) + except TypeError as e: + if " empty " in str(e): + raise ValueError("can not generate buckets from empty RDD") + raise + + if minv == maxv or buckets == 1: + return [minv, maxv], [filtered.count()] + + try: + inc = (maxv - minv) / buckets # type: ignore[operator] + except TypeError: + raise TypeError("Can not generate buckets with non-number in RDD") + + if isinf(inc): + raise ValueError("Can not generate buckets with infinite value") + + # keep them as integer if possible + inc = int(inc) + if inc * buckets != maxv - minv: # type: ignore[operator] + inc = (maxv - minv) * 1.0 / buckets # type: ignore[operator] + + buckets = [i * inc + minv for i in range(buckets)] + buckets.append(maxv) # fix accumulated error + even = True + + elif isinstance(buckets, (list, tuple)): + if len(buckets) < 2: + raise ValueError("buckets should have more than one value") + + if any(i is None or isinstance(i, float) and isnan(i) for i in buckets): + raise ValueError("can not have None or NaN in buckets") + + if sorted(buckets) != list(buckets): + raise ValueError("buckets should be sorted") + + if len(set(buckets)) != len(buckets): + raise ValueError("buckets should not contain duplicated values") + + minv = buckets[0] + maxv = buckets[-1] + even = False + inc = None + try: + steps = [ + buckets[i + 1] - buckets[i] # type: ignore[operator] + for i in range(len(buckets) - 1) + ] + except TypeError: + pass # objects in buckets do not support '-' + else: + if max(steps) - min(steps) < 1e-10: # handle precision errors + even = True + inc = (maxv - minv) / (len(buckets) - 1) # type: ignore[operator] + + else: + raise TypeError("buckets should be a list or tuple or number(int or long)") + + def histogram(iterator: Iterable["S"]) -> Iterable[List[int]]: + counters = [0] * len(buckets) # type: ignore[arg-type] + for i in iterator: + if i is None or (isinstance(i, float) and isnan(i)) or i > maxv or i < minv: + continue + t = ( + int((i - minv) / inc) # type: ignore[operator] + if even + else bisect.bisect_right(buckets, i) - 1 # type: ignore[arg-type] + ) + counters[t] += 1 + # add last two together + last = counters.pop() + counters[-1] += last + return [counters] + + def mergeCounters(a: List[int], b: List[int]) -> List[int]: + return [i + j for i, j in zip(a, b)] + + return buckets, self.mapPartitions(histogram).reduce(mergeCounters) + + def mean(self: "RDD[NumberOrArray]") -> float: + """ + Compute the mean of this RDD's elements. + + .. versionadded:: 0.9.1 + + Returns + ------- + float + the mean of all elements + + See Also + -------- + :meth:`RDD.stats` + :meth:`RDD.sum` + :meth:`RDD.meanApprox` + + Examples + -------- + >>> sc.parallelize([1, 2, 3]).mean() + 2.0 + """ + return self.stats().mean() + + def variance(self: "RDD[NumberOrArray]") -> float: + """ + Compute the variance of this RDD's elements. + + .. versionadded:: 0.9.1 + + Returns + ------- + float + the variance of all elements + + See Also + -------- + :meth:`RDD.stats` + :meth:`RDD.sampleVariance` + :meth:`RDD.stdev` + :meth:`RDD.sampleStdev` + + Examples + -------- + >>> sc.parallelize([1, 2, 3]).variance() + 0.666... + """ + return self.stats().variance() + + def stdev(self: "RDD[NumberOrArray]") -> float: + """ + Compute the standard deviation of this RDD's elements. + + .. versionadded:: 0.9.1 + + Returns + ------- + float + the standard deviation of all elements + + See Also + -------- + :meth:`RDD.stats` + :meth:`RDD.sampleStdev` + :meth:`RDD.variance` + :meth:`RDD.sampleVariance` + + Examples + -------- + >>> sc.parallelize([1, 2, 3]).stdev() + 0.816... + """ + return self.stats().stdev() + + def sampleStdev(self: "RDD[NumberOrArray]") -> float: + """ + Compute the sample standard deviation of this RDD's elements (which + corrects for bias in estimating the standard deviation by dividing by + N-1 instead of N). + + .. versionadded:: 0.9.1 + + Returns + ------- + float + the sample standard deviation of all elements + + See Also + -------- + :meth:`RDD.stats` + :meth:`RDD.stdev` + :meth:`RDD.variance` + :meth:`RDD.sampleVariance` + + Examples + -------- + >>> sc.parallelize([1, 2, 3]).sampleStdev() + 1.0 + """ + return self.stats().sampleStdev() + + def sampleVariance(self: "RDD[NumberOrArray]") -> float: + """ + Compute the sample variance of this RDD's elements (which corrects + for bias in estimating the variance by dividing by N-1 instead of N). + + .. versionadded:: 0.9.1 + + Returns + ------- + float + the sample variance of all elements + + See Also + -------- + :meth:`RDD.stats` + :meth:`RDD.variance` + :meth:`RDD.stdev` + :meth:`RDD.sampleStdev` + + Examples + -------- + >>> sc.parallelize([1, 2, 3]).sampleVariance() + 1.0 + """ + return self.stats().sampleVariance() + + def countByValue(self: "RDD[K]") -> Dict[K, int]: + """ + Return the count of each unique value in this RDD as a dictionary of + (value, count) pairs. + + .. versionadded:: 0.7.0 + + Returns + ------- + dict + a dictionary of (value, count) pairs + + See Also + -------- + :meth:`RDD.collectAsMap` + :meth:`RDD.countByKey` + + Examples + -------- + >>> sorted(sc.parallelize([1, 2, 1, 2, 2], 2).countByValue().items()) + [(1, 2), (2, 3)] + """ + + def countPartition(iterator: Iterable[K]) -> Iterable[Dict[K, int]]: + counts: Dict[K, int] = defaultdict(int) + for obj in iterator: + counts[obj] += 1 + yield counts + + def mergeMaps(m1: Dict[K, int], m2: Dict[K, int]) -> Dict[K, int]: + for k, v in m2.items(): + m1[k] += v + return m1 + + return self.mapPartitions(countPartition).reduce(mergeMaps) + + @overload + def top(self: "RDD[S]", num: int) -> List["S"]: + ... + + @overload + def top(self: "RDD[T]", num: int, key: Callable[[T], "S"]) -> List[T]: + ... + + def top(self: "RDD[T]", num: int, key: Optional[Callable[[T], "S"]] = None) -> List[T]: + """ + Get the top N elements from an RDD. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + num : int + top N + key : function, optional + a function used to generate key for comparing + + Returns + ------- + list + the top N elements + + See Also + -------- + :meth:`RDD.takeOrdered` + :meth:`RDD.max` + :meth:`RDD.min` + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + It returns the list sorted in descending order. + + Examples + -------- + >>> sc.parallelize([10, 4, 2, 12, 3]).top(1) + [12] + >>> sc.parallelize([2, 3, 4, 5, 6], 2).top(2) + [6, 5] + >>> sc.parallelize([10, 4, 2, 12, 3]).top(3, key=str) + [4, 3, 2] + """ + + def topIterator(iterator: Iterable[T]) -> Iterable[List[T]]: + yield heapq.nlargest(num, iterator, key=key) + + def merge(a: List[T], b: List[T]) -> List[T]: + return heapq.nlargest(num, a + b, key=key) + + return self.mapPartitions(topIterator).reduce(merge) + + @overload + def takeOrdered(self: "RDD[S]", num: int) -> List["S"]: + ... + + @overload + def takeOrdered(self: "RDD[T]", num: int, key: Callable[[T], "S"]) -> List[T]: + ... + + def takeOrdered(self: "RDD[T]", num: int, key: Optional[Callable[[T], "S"]] = None) -> List[T]: + """ + Get the N elements from an RDD ordered in ascending order or as + specified by the optional key function. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + num : int + top N + key : function, optional + a function used to generate key for comparing + + Returns + ------- + list + the top N elements + + See Also + -------- + :meth:`RDD.top` + :meth:`RDD.max` + :meth:`RDD.min` + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> sc.parallelize([10, 1, 2, 9, 3, 4, 5, 6, 7]).takeOrdered(6) + [1, 2, 3, 4, 5, 6] + >>> sc.parallelize([10, 1, 2, 9, 3, 4, 5, 6, 7], 2).takeOrdered(6, key=lambda x: -x) + [10, 9, 7, 6, 5, 4] + >>> sc.emptyRDD().takeOrdered(3) + [] + """ + if num < 0: + raise ValueError("top N cannot be negative.") + + if num == 0 or self.getNumPartitions() == 0: + return [] + else: + + def merge(a: List[T], b: List[T]) -> List[T]: + return heapq.nsmallest(num, a + b, key) + + return self.mapPartitions(lambda it: [heapq.nsmallest(num, it, key)]).reduce(merge) + + def take(self: "RDD[T]", num: int) -> List[T]: + """ + Take the first num elements of the RDD. + + It works by first scanning one partition, and use the results from + that partition to estimate the number of additional partitions needed + to satisfy the limit. + + Translated from the Scala implementation in RDD#take(). + + .. versionadded:: 0.7.0 + + Parameters + ---------- + num : int + first number of elements + + Returns + ------- + list + the first `num` elements + + See Also + -------- + :meth:`RDD.first` + :meth:`pyspark.sql.DataFrame.take` + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> sc.parallelize([2, 3, 4, 5, 6]).cache().take(2) + [2, 3] + >>> sc.parallelize([2, 3, 4, 5, 6]).take(10) + [2, 3, 4, 5, 6] + >>> sc.parallelize(range(100), 100).filter(lambda x: x > 90).take(3) + [91, 92, 93] + """ + items: List[T] = [] + totalParts = self.getNumPartitions() + partsScanned = 0 + + while len(items) < num and partsScanned < totalParts: + # The number of partitions to try in this iteration. + # It is ok for this number to be greater than totalParts because + # we actually cap it at totalParts in runJob. + numPartsToTry = 1 + if partsScanned > 0: + # If we didn't find any rows after the previous iteration, + # quadruple and retry. Otherwise, interpolate the number of + # partitions we need to try, but overestimate it by 50%. + # We also cap the estimation in the end. + if len(items) == 0: + numPartsToTry = partsScanned * 4 + else: + # the first parameter of max is >=1 whenever partsScanned >= 2 + numPartsToTry = int(1.5 * num * partsScanned / len(items)) - partsScanned + numPartsToTry = min(max(numPartsToTry, 1), partsScanned * 4) + + left = num - len(items) + + def takeUpToNumLeft(iterator: Iterable[T]) -> Iterable[T]: + iterator = iter(iterator) + taken = 0 + while taken < left: + try: + yield next(iterator) + except StopIteration: + return + taken += 1 + + p = range(partsScanned, min(partsScanned + numPartsToTry, totalParts)) + res = self.context.runJob(self, takeUpToNumLeft, p) + + items += res + partsScanned += numPartsToTry + + return items[:num] + + def first(self: "RDD[T]") -> T: + """ + Return the first element in this RDD. + + .. versionadded:: 0.7.0 + + Returns + ------- + T + the first element + + See Also + -------- + :meth:`RDD.take` + :meth:`pyspark.sql.DataFrame.first` + :meth:`pyspark.sql.DataFrame.head` + + Examples + -------- + >>> sc.parallelize([2, 3, 4]).first() + 2 + >>> sc.parallelize([]).first() + Traceback (most recent call last): + ... + ValueError: RDD is empty + """ + rs = self.take(1) + if rs: + return rs[0] + raise ValueError("RDD is empty") + + def isEmpty(self) -> bool: + """ + Returns true if and only if the RDD contains no elements at all. + + .. versionadded:: 1.3.0 + + Returns + ------- + bool + whether the :class:`RDD` is empty + + See Also + -------- + :meth:`RDD.first` + :meth:`pyspark.sql.DataFrame.isEmpty` + + Notes + ----- + An RDD may be empty even when it has at least 1 partition. + + Examples + -------- + >>> sc.parallelize([]).isEmpty() + True + >>> sc.parallelize([1]).isEmpty() + False + """ + return self.getNumPartitions() == 0 or len(self.take(1)) == 0 + + def saveAsNewAPIHadoopDataset( + self: "RDD[Tuple[K, V]]", + conf: Dict[str, str], + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + ) -> None: + """ + Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to any Hadoop file + system, using the new Hadoop OutputFormat API (mapreduce package). Keys/values are + converted for output using either user specified converters or, by default, + "org.apache.spark.api.python.JavaToWritableConverter". + + .. versionadded:: 1.1.0 + + Parameters + ---------- + conf : dict + Hadoop job configuration + keyConverter : str, optional + fully qualified classname of key converter (None by default) + valueConverter : str, optional + fully qualified classname of value converter (None by default) + + See Also + -------- + :meth:`SparkContext.newAPIHadoopRDD` + :meth:`RDD.saveAsHadoopDataset` + :meth:`RDD.saveAsHadoopFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsSequenceFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "new_hadoop_file") + ... + ... # Create the conf for writing + ... write_conf = { + ... "mapreduce.job.outputformat.class": (output_format_class), + ... "mapreduce.job.output.key.class": key_class, + ... "mapreduce.job.output.value.class": value_class, + ... "mapreduce.output.fileoutputformat.outputdir": path, + ... } + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsNewAPIHadoopDataset(conf=write_conf) + ... + ... # Create the conf for reading + ... read_conf = {"mapreduce.input.fileinputformat.inputdir": path} + ... + ... # Load this Hadoop file as an RDD + ... loaded = sc.newAPIHadoopRDD(input_format_class, + ... key_class, value_class, conf=read_conf) + ... sorted(loaded.collect()) + [(1, ''), (1, 'a'), (3, 'x')] + """ + jconf = self.ctx._dictToJavaMap(conf) + pickledRDD = self._pickled() + assert self.ctx._jvm is not None + + self.ctx._jvm.PythonRDD.saveAsHadoopDataset( + pickledRDD._jrdd, True, jconf, keyConverter, valueConverter, True + ) + + def saveAsNewAPIHadoopFile( + self: "RDD[Tuple[K, V]]", + path: str, + outputFormatClass: str, + keyClass: Optional[str] = None, + valueClass: Optional[str] = None, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + ) -> None: + """ + Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to any Hadoop file + system, using the new Hadoop OutputFormat API (mapreduce package). Key and value types + will be inferred if not specified. Keys and values are converted for output using either + user specified converters or "org.apache.spark.api.python.JavaToWritableConverter". The + `conf` is applied on top of the base Hadoop conf associated with the SparkContext + of this RDD to create a merged Hadoop MapReduce job configuration for saving the data. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + path : str + path to Hadoop file + outputFormatClass : str + fully qualified classname of Hadoop OutputFormat + (e.g. "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat") + keyClass : str, optional + fully qualified classname of key Writable class + (e.g. "org.apache.hadoop.io.IntWritable", None by default) + valueClass : str, optional + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.Text", None by default) + keyConverter : str, optional + fully qualified classname of key converter (None by default) + valueConverter : str, optional + fully qualified classname of value converter (None by default) + conf : dict, optional + Hadoop job configuration (None by default) + + See Also + -------- + :meth:`SparkContext.newAPIHadoopFile` + :meth:`RDD.saveAsHadoopDataset` + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsHadoopFile` + :meth:`RDD.saveAsSequenceFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the class of output format + + >>> output_format_class = "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "hadoop_file") + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, {3.0: "bb"}), (2, {1.0: "aa"}), (3, {2.0: "dd"})]) + ... rdd.saveAsNewAPIHadoopFile(path, output_format_class) + ... + ... # Load this Hadoop file as an RDD + ... sorted(sc.sequenceFile(path).collect()) + [(1, {3.0: 'bb'}), (2, {1.0: 'aa'}), (3, {2.0: 'dd'})] + """ + jconf = self.ctx._dictToJavaMap(conf) + pickledRDD = self._pickled() + assert self.ctx._jvm is not None + + self.ctx._jvm.PythonRDD.saveAsNewAPIHadoopFile( + pickledRDD._jrdd, + True, + path, + outputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + ) + + def saveAsHadoopDataset( + self: "RDD[Tuple[K, V]]", + conf: Dict[str, str], + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + ) -> None: + """ + Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to any Hadoop file + system, using the old Hadoop OutputFormat API (mapred package). Keys/values are + converted for output using either user specified converters or, by default, + "org.apache.spark.api.python.JavaToWritableConverter". + + .. versionadded:: 1.1.0 + + Parameters + ---------- + conf : dict + Hadoop job configuration + keyConverter : str, optional + fully qualified classname of key converter (None by default) + valueConverter : str, optional + fully qualified classname of value converter (None by default) + + See Also + -------- + :meth:`SparkContext.hadoopRDD` + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsHadoopFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsSequenceFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapred.TextOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapred.TextInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "old_hadoop_file") + ... + ... # Create the conf for writing + ... write_conf = { + ... "mapred.output.format.class": output_format_class, + ... "mapreduce.job.output.key.class": key_class, + ... "mapreduce.job.output.value.class": value_class, + ... "mapreduce.output.fileoutputformat.outputdir": path, + ... } + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsHadoopDataset(conf=write_conf) + ... + ... # Create the conf for reading + ... read_conf = {"mapreduce.input.fileinputformat.inputdir": path} + ... + ... # Load this Hadoop file as an RDD + ... loaded = sc.hadoopRDD(input_format_class, key_class, value_class, conf=read_conf) + ... sorted(loaded.collect()) + [(0, '1\\t'), (0, '1\\ta'), (0, '3\\tx')] + """ + jconf = self.ctx._dictToJavaMap(conf) + pickledRDD = self._pickled() + assert self.ctx._jvm is not None + + self.ctx._jvm.PythonRDD.saveAsHadoopDataset( + pickledRDD._jrdd, True, jconf, keyConverter, valueConverter, False + ) + + def saveAsHadoopFile( + self: "RDD[Tuple[K, V]]", + path: str, + outputFormatClass: str, + keyClass: Optional[str] = None, + valueClass: Optional[str] = None, + keyConverter: Optional[str] = None, + valueConverter: Optional[str] = None, + conf: Optional[Dict[str, str]] = None, + compressionCodecClass: Optional[str] = None, + ) -> None: + """ + Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to any Hadoop file + system, using the old Hadoop OutputFormat API (mapred package). Key and value types + will be inferred if not specified. Keys and values are converted for output using either + user specified converters or "org.apache.spark.api.python.JavaToWritableConverter". The + `conf` is applied on top of the base Hadoop conf associated with the SparkContext + of this RDD to create a merged Hadoop MapReduce job configuration for saving the data. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + path : str + path to Hadoop file + outputFormatClass : str + fully qualified classname of Hadoop OutputFormat + (e.g. "org.apache.hadoop.mapred.SequenceFileOutputFormat") + keyClass : str, optional + fully qualified classname of key Writable class + (e.g. "org.apache.hadoop.io.IntWritable", None by default) + valueClass : str, optional + fully qualified classname of value Writable class + (e.g. "org.apache.hadoop.io.Text", None by default) + keyConverter : str, optional + fully qualified classname of key converter (None by default) + valueConverter : str, optional + fully qualified classname of value converter (None by default) + conf : dict, optional + (None by default) + compressionCodecClass : str + fully qualified classname of the compression codec class + i.e. "org.apache.hadoop.io.compress.GzipCodec" (None by default) + + See Also + -------- + :meth:`SparkContext.hadoopFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsHadoopDataset` + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsSequenceFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> output_format_class = "org.apache.hadoop.mapred.TextOutputFormat" + >>> input_format_class = "org.apache.hadoop.mapred.TextInputFormat" + >>> key_class = "org.apache.hadoop.io.IntWritable" + >>> value_class = "org.apache.hadoop.io.Text" + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "old_hadoop_file") + ... + ... # Write a temporary Hadoop file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsHadoopFile(path, output_format_class, key_class, value_class) + ... + ... # Load this Hadoop file as an RDD + ... loaded = sc.hadoopFile(path, input_format_class, key_class, value_class) + ... sorted(loaded.collect()) + [(0, '1\\t'), (0, '1\\ta'), (0, '3\\tx')] + """ + jconf = self.ctx._dictToJavaMap(conf) + pickledRDD = self._pickled() + assert self.ctx._jvm is not None + + self.ctx._jvm.PythonRDD.saveAsHadoopFile( + pickledRDD._jrdd, + True, + path, + outputFormatClass, + keyClass, + valueClass, + keyConverter, + valueConverter, + jconf, + compressionCodecClass, + ) + + def saveAsSequenceFile( + self: "RDD[Tuple[K, V]]", path: str, compressionCodecClass: Optional[str] = None + ) -> None: + """ + Output a Python RDD of key-value pairs (of form ``RDD[(K, V)]``) to any Hadoop file + system, using the "org.apache.hadoop.io.Writable" types that we convert from the + RDD's key and value types. The mechanism is as follows: + + 1. Pickle is used to convert pickled Python RDD into RDD of Java objects. + 2. Keys and values of this Java RDD are converted to Writables and written out. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + path : str + path to sequence file + compressionCodecClass : str, optional + fully qualified classname of the compression codec class + i.e. "org.apache.hadoop.io.compress.GzipCodec" (None by default) + + See Also + -------- + :meth:`SparkContext.sequenceFile` + :meth:`RDD.saveAsHadoopFile` + :meth:`RDD.saveAsNewAPIHadoopFile` + :meth:`RDD.saveAsHadoopDataset` + :meth:`RDD.saveAsNewAPIHadoopDataset` + :meth:`RDD.saveAsSequenceFile` + + Examples + -------- + >>> import os + >>> import tempfile + + Set the related classes + + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "sequence_file") + ... + ... # Write a temporary sequence file + ... rdd = sc.parallelize([(1, ""), (1, "a"), (3, "x")]) + ... rdd.saveAsSequenceFile(path) + ... + ... # Load this sequence file as an RDD + ... loaded = sc.sequenceFile(path) + ... sorted(loaded.collect()) + [(1, ''), (1, 'a'), (3, 'x')] + """ + pickledRDD = self._pickled() + assert self.ctx._jvm is not None + + self.ctx._jvm.PythonRDD.saveAsSequenceFile( + pickledRDD._jrdd, True, path, compressionCodecClass + ) + + def saveAsPickleFile(self, path: str, batchSize: int = 10) -> None: + """ + Save this RDD as a SequenceFile of serialized objects. The serializer + used is :class:`pyspark.serializers.CPickleSerializer`, default batch size + is 10. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + path : str + path to pickled file + batchSize : int, optional, default 10 + the number of Python objects represented as a single Java object. + + See Also + -------- + :meth:`SparkContext.pickleFile` + + Examples + -------- + >>> import os + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... path = os.path.join(d, "pickle_file") + ... + ... # Write a temporary pickled file + ... sc.parallelize(range(10)).saveAsPickleFile(path, 3) + ... + ... # Load picked file as an RDD + ... sorted(sc.pickleFile(path, 3).collect()) + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + """ + ser: Serializer + if batchSize == 0: + ser = AutoBatchedSerializer(CPickleSerializer()) + else: + ser = BatchedSerializer(CPickleSerializer(), batchSize) + self._reserialize(ser)._jrdd.saveAsObjectFile(path) + + def saveAsTextFile(self, path: str, compressionCodecClass: Optional[str] = None) -> None: + """ + Save this RDD as a text file, using string representations of elements. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + path : str + path to text file + compressionCodecClass : str, optional + fully qualified classname of the compression codec class + i.e. "org.apache.hadoop.io.compress.GzipCodec" (None by default) + + See Also + -------- + :meth:`SparkContext.textFile` + :meth:`SparkContext.wholeTextFiles` + + Examples + -------- + >>> import os + >>> import tempfile + >>> from fileinput import input + >>> from glob import glob + >>> with tempfile.TemporaryDirectory() as d1: + ... path1 = os.path.join(d1, "text_file1") + ... + ... # Write a temporary text file + ... sc.parallelize(range(10)).saveAsTextFile(path1) + ... + ... # Load text file as an RDD + ... ''.join(sorted(input(glob(path1 + "/part-0000*")))) + '0\\n1\\n2\\n3\\n4\\n5\\n6\\n7\\n8\\n9\\n' + + Empty lines are tolerated when saving to text files. + + >>> with tempfile.TemporaryDirectory() as d2: + ... path2 = os.path.join(d2, "text2_file2") + ... + ... # Write another temporary text file + ... sc.parallelize(['', 'foo', '', 'bar', '']).saveAsTextFile(path2) + ... + ... # Load text file as an RDD + ... ''.join(sorted(input(glob(path2 + "/part-0000*")))) + '\\n\\n\\nbar\\nfoo\\n' + + Using compressionCodecClass + + >>> from fileinput import input, hook_compressed + >>> with tempfile.TemporaryDirectory() as d3: + ... path3 = os.path.join(d3, "text3") + ... codec = "org.apache.hadoop.io.compress.GzipCodec" + ... + ... # Write another temporary text file with specified codec + ... sc.parallelize(['foo', 'bar']).saveAsTextFile(path3, codec) + ... + ... # Load text file as an RDD + ... result = sorted(input(glob(path3 + "/part*.gz"), openhook=hook_compressed)) + ... ''.join([r.decode('utf-8') if isinstance(r, bytes) else r for r in result]) + 'bar\\nfoo\\n' + """ + + def func(split: int, iterator: Iterable[Any]) -> Iterable[bytes]: + for x in iterator: + if isinstance(x, bytes): + yield x + elif isinstance(x, str): + yield x.encode("utf-8") + else: + yield str(x).encode("utf-8") + + keyed = self.mapPartitionsWithIndex(func) + keyed._bypass_serializer = True # type: ignore[attr-defined] + + assert self.ctx._jvm is not None + + if compressionCodecClass: + compressionCodec = self.ctx._jvm.java.lang.Class.forName(compressionCodecClass) + keyed._jrdd.map(self.ctx._jvm.BytesToString()).saveAsTextFile(path, compressionCodec) + else: + keyed._jrdd.map(self.ctx._jvm.BytesToString()).saveAsTextFile(path) + + # Pair functions + + def collectAsMap(self: "RDD[Tuple[K, V]]") -> Dict[K, V]: + """ + Return the key-value pairs in this RDD to the master as a dictionary. + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`dict` + a dictionary of (key, value) pairs + + See Also + -------- + :meth:`RDD.countByValue` + + Notes + ----- + This method should only be used if the resulting data is expected + to be small, as all the data is loaded into the driver's memory. + + Examples + -------- + >>> m = sc.parallelize([(1, 2), (3, 4)]).collectAsMap() + >>> m[1] + 2 + >>> m[3] + 4 + """ + return dict(self.collect()) + + def keys(self: "RDD[Tuple[K, V]]") -> "RDD[K]": + """ + Return an RDD with the keys of each tuple. + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`RDD` + a :class:`RDD` only containing the keys + + See Also + -------- + :meth:`RDD.values` + + Examples + -------- + >>> rdd = sc.parallelize([(1, 2), (3, 4)]).keys() + >>> rdd.collect() + [1, 3] + """ + return self.map(lambda x: x[0]) + + def values(self: "RDD[Tuple[K, V]]") -> "RDD[V]": + """ + Return an RDD with the values of each tuple. + + .. versionadded:: 0.7.0 + + Returns + ------- + :class:`RDD` + a :class:`RDD` only containing the values + + See Also + -------- + :meth:`RDD.keys` + + Examples + -------- + >>> rdd = sc.parallelize([(1, 2), (3, 4)]).values() + >>> rdd.collect() + [2, 4] + """ + return self.map(lambda x: x[1]) + + def reduceByKey( + self: "RDD[Tuple[K, V]]", + func: Callable[[V, V], V], + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, V]]": + """ + Merge the values for each key using an associative and commutative reduce function. + + This will also perform the merging locally on each mapper before + sending results to a reducer, similarly to a "combiner" in MapReduce. + + Output will be partitioned with `numPartitions` partitions, or + the default parallelism level if `numPartitions` is not specified. + Default partitioner is hash-partition. + + .. versionadded:: 1.6.0 + + Parameters + ---------- + func : function + the reduce function + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the aggregated result for each key + + See Also + -------- + :meth:`RDD.reduceByKeyLocally` + :meth:`RDD.combineByKey` + :meth:`RDD.aggregateByKey` + :meth:`RDD.foldByKey` + :meth:`RDD.groupByKey` + + Examples + -------- + >>> from operator import add + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 1)]) + >>> sorted(rdd.reduceByKey(add).collect()) + [('a', 2), ('b', 1)] + """ + return self.combineByKey(lambda x: x, func, func, numPartitions, partitionFunc) + + def reduceByKeyLocally(self: "RDD[Tuple[K, V]]", func: Callable[[V, V], V]) -> Dict[K, V]: + """ + Merge the values for each key using an associative and commutative reduce function, but + return the results immediately to the master as a dictionary. + + This will also perform the merging locally on each mapper before + sending results to a reducer, similarly to a "combiner" in MapReduce. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + func : function + the reduce function + + Returns + ------- + dict + a dict containing the keys and the aggregated result for each key + + See Also + -------- + :meth:`RDD.reduceByKey` + :meth:`RDD.aggregateByKey` + + Examples + -------- + >>> from operator import add + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 1)]) + >>> sorted(rdd.reduceByKeyLocally(add).items()) + [('a', 2), ('b', 1)] + """ + func = fail_on_stopiteration(func) + + def reducePartition(iterator: Iterable[Tuple[K, V]]) -> Iterable[Dict[K, V]]: + m: Dict[K, V] = {} + for k, v in iterator: + m[k] = func(m[k], v) if k in m else v + yield m + + def mergeMaps(m1: Dict[K, V], m2: Dict[K, V]) -> Dict[K, V]: + for k, v in m2.items(): + m1[k] = func(m1[k], v) if k in m1 else v + return m1 + + return self.mapPartitions(reducePartition).reduce(mergeMaps) + + def countByKey(self: "RDD[Tuple[K, V]]") -> Dict[K, int]: + """ + Count the number of elements for each key, and return the result to the + master as a dictionary. + + .. versionadded:: 0.7.0 + + Returns + ------- + dict + a dictionary of (key, count) pairs + + See Also + -------- + :meth:`RDD.collectAsMap` + :meth:`RDD.countByValue` + + Examples + -------- + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 1)]) + >>> sorted(rdd.countByKey().items()) + [('a', 2), ('b', 1)] + """ + return self.map(lambda x: x[0]).countByValue() + + def join( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, Tuple[V, U]]]": + """ + Return an RDD containing all pairs of elements with matching keys in + `self` and `other`. + + Each pair of elements will be returned as a (k, (v1, v2)) tuple, where + (k, v1) is in `self` and (k, v2) is in `other`. + + Performs a hash join across the cluster. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing all pairs of elements with matching keys + + See Also + -------- + :meth:`RDD.leftOuterJoin` + :meth:`RDD.rightOuterJoin` + :meth:`RDD.fullOuterJoin` + :meth:`RDD.cogroup` + :meth:`RDD.groupWith` + :meth:`pyspark.sql.DataFrame.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd2 = sc.parallelize([("a", 2), ("a", 3)]) + >>> sorted(rdd1.join(rdd2).collect()) + [('a', (1, 2)), ('a', (1, 3))] + """ + return python_join(self, other, numPartitions) + + def leftOuterJoin( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, Tuple[V, Optional[U]]]]": + """ + Perform a left outer join of `self` and `other`. + + For each element (k, v) in `self`, the resulting RDD will either + contain all pairs (k, (v, w)) for w in `other`, or the pair + (k, (v, None)) if no elements in `other` have key k. + + Hash-partitions the resulting RDD into the given number of partitions. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing all pairs of elements with matching keys + + See Also + -------- + :meth:`RDD.join` + :meth:`RDD.rightOuterJoin` + :meth:`RDD.fullOuterJoin` + :meth:`pyspark.sql.DataFrame.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd2 = sc.parallelize([("a", 2)]) + >>> sorted(rdd1.leftOuterJoin(rdd2).collect()) + [('a', (1, 2)), ('b', (4, None))] + """ + return python_left_outer_join(self, other, numPartitions) + + def rightOuterJoin( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, Tuple[Optional[V], U]]]": + """ + Perform a right outer join of `self` and `other`. + + For each element (k, w) in `other`, the resulting RDD will either + contain all pairs (k, (v, w)) for v in this, or the pair (k, (None, w)) + if no elements in `self` have key k. + + Hash-partitions the resulting RDD into the given number of partitions. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing all pairs of elements with matching keys + + See Also + -------- + :meth:`RDD.join` + :meth:`RDD.leftOuterJoin` + :meth:`RDD.fullOuterJoin` + :meth:`pyspark.sql.DataFrame.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd2 = sc.parallelize([("a", 2)]) + >>> sorted(rdd2.rightOuterJoin(rdd1).collect()) + [('a', (2, 1)), ('b', (None, 4))] + """ + return python_right_outer_join(self, other, numPartitions) + + def fullOuterJoin( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, Tuple[Optional[V], Optional[U]]]]": + """ + Perform a right outer join of `self` and `other`. + + For each element (k, v) in `self`, the resulting RDD will either + contain all pairs (k, (v, w)) for w in `other`, or the pair + (k, (v, None)) if no elements in `other` have key k. + + Similarly, for each element (k, w) in `other`, the resulting RDD will + either contain all pairs (k, (v, w)) for v in `self`, or the pair + (k, (None, w)) if no elements in `self` have key k. + + Hash-partitions the resulting RDD into the given number of partitions. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing all pairs of elements with matching keys + + See Also + -------- + :meth:`RDD.join` + :meth:`RDD.leftOuterJoin` + :meth:`RDD.fullOuterJoin` + :meth:`pyspark.sql.DataFrame.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd2 = sc.parallelize([("a", 2), ("c", 8)]) + >>> sorted(rdd1.fullOuterJoin(rdd2).collect()) + [('a', (1, 2)), ('b', (4, None)), ('c', (None, 8))] + """ + return python_full_outer_join(self, other, numPartitions) + + # TODO: add option to control map-side combining + # portable_hash is used as default, because builtin hash of None is different + # cross machines. + def partitionBy( + self: "RDD[Tuple[K, V]]", + numPartitions: Optional[int], + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, V]]": + """ + Return a copy of the RDD partitioned using the specified partitioner. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` partitioned using the specified partitioner + + See Also + -------- + :meth:`RDD.repartition` + :meth:`RDD.repartitionAndSortWithinPartitions` + + Examples + -------- + >>> pairs = sc.parallelize([1, 2, 3, 4, 2, 4, 1]).map(lambda x: (x, x)) + >>> sets = pairs.partitionBy(2).glom().collect() + >>> len(set(sets[0]).intersection(set(sets[1]))) + 0 + """ + if numPartitions is None: + numPartitions = self._defaultReducePartitions() + partitioner = Partitioner(numPartitions, partitionFunc) + if self.partitioner == partitioner: + return self + + # Transferring O(n) objects to Java is too expensive. + # Instead, we'll form the hash buckets in Python, + # transferring O(numPartitions) objects to Java. + # Each object is a (splitNumber, [objects]) pair. + # In order to avoid too huge objects, the objects are + # grouped into chunks. + outputSerializer = self.ctx._unbatched_serializer + + limit = self._memory_limit() / 2 + + def add_shuffle_key(split: int, iterator: Iterable[Tuple[K, V]]) -> Iterable[bytes]: + + buckets = defaultdict(list) + c, batch = 0, min(10 * numPartitions, 1000) # type: ignore[operator] + + for k, v in iterator: + buckets[partitionFunc(k) % numPartitions].append((k, v)) # type: ignore[operator] + c += 1 + + # check used memory and avg size of chunk of objects + if c % 1000 == 0 and get_used_memory() > limit or c > batch: + n, size = len(buckets), 0 + for split in list(buckets.keys()): + yield pack_long(split) + d = outputSerializer.dumps(buckets[split]) + del buckets[split] + yield d + size += len(d) + + avg = int(size / n) >> 20 + # let 1M < avg < 10M + if avg < 1: + batch = min(sys.maxsize, batch * 1.5) # type: ignore[assignment] + elif avg > 10: + batch = max(int(batch / 1.5), 1) + c = 0 + + for split, items in buckets.items(): + yield pack_long(split) + yield outputSerializer.dumps(items) + + keyed = self.mapPartitionsWithIndex(add_shuffle_key, preservesPartitioning=True) + keyed._bypass_serializer = True # type: ignore[attr-defined] + assert self.ctx._jvm is not None + + with SCCallSiteSync(self.context): + pairRDD = self.ctx._jvm.PairwiseRDD(keyed._jrdd.rdd()).asJavaPairRDD() + jpartitioner = self.ctx._jvm.PythonPartitioner(numPartitions, id(partitionFunc)) + jrdd = self.ctx._jvm.PythonRDD.valueOfPair(pairRDD.partitionBy(jpartitioner)) + rdd: "RDD[Tuple[K, V]]" = RDD(jrdd, self.ctx, BatchedSerializer(outputSerializer)) + rdd.partitioner = partitioner + return rdd + + # TODO: add control over map-side aggregation + def combineByKey( + self: "RDD[Tuple[K, V]]", + createCombiner: Callable[[V], U], + mergeValue: Callable[[U, V], U], + mergeCombiners: Callable[[U, U], U], + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, U]]": + """ + Generic function to combine the elements for each key using a custom + set of aggregation functions. + + Turns an RDD[(K, V)] into a result of type RDD[(K, C)], for a "combined + type" C. + + To avoid memory allocation, both mergeValue and mergeCombiners are allowed to + modify and return their first argument instead of creating a new C. + + In addition, users can control the partitioning of the output RDD. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + createCombiner : function + a function to turns a V into a C + mergeValue : function + a function to merge a V into a C + mergeCombiners : function + a function to combine two C's into a single one + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the aggregated result for each key + + See Also + -------- + :meth:`RDD.reduceByKey` + :meth:`RDD.aggregateByKey` + :meth:`RDD.foldByKey` + :meth:`RDD.groupByKey` + + Notes + ----- + V and C can be different -- for example, one might group an RDD of type + (Int, Int) into an RDD of type (Int, List[Int]). + + Examples + -------- + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 2)]) + >>> def to_list(a): + ... return [a] + ... + >>> def append(a, b): + ... a.append(b) + ... return a + ... + >>> def extend(a, b): + ... a.extend(b) + ... return a + ... + >>> sorted(rdd.combineByKey(to_list, append, extend).collect()) + [('a', [1, 2]), ('b', [1])] + """ + if numPartitions is None: + numPartitions = self._defaultReducePartitions() + + serializer = self.ctx.serializer + memory = self._memory_limit() + agg = Aggregator(createCombiner, mergeValue, mergeCombiners) + + def combineLocally(iterator: Iterable[Tuple[K, V]]) -> Iterable[Tuple[K, U]]: + merger = ExternalMerger(agg, memory * 0.9, serializer) + merger.mergeValues(iterator) + return merger.items() + + locally_combined = self.mapPartitions(combineLocally, preservesPartitioning=True) + shuffled = locally_combined.partitionBy(numPartitions, partitionFunc) + + def _mergeCombiners(iterator: Iterable[Tuple[K, U]]) -> Iterable[Tuple[K, U]]: + merger = ExternalMerger(agg, memory, serializer) + merger.mergeCombiners(iterator) + return merger.items() + + return shuffled.mapPartitions(_mergeCombiners, preservesPartitioning=True) + + def aggregateByKey( + self: "RDD[Tuple[K, V]]", + zeroValue: U, + seqFunc: Callable[[U, V], U], + combFunc: Callable[[U, U], U], + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, U]]": + """ + Aggregate the values of each key, using given combine functions and a neutral + "zero value". This function can return a different result type, U, than the type + of the values in this RDD, V. Thus, we need one operation for merging a V into + a U and one operation for merging two U's, The former operation is used for merging + values within a partition, and the latter is used for merging values between + partitions. To avoid memory allocation, both of these functions are + allowed to modify and return their first argument instead of creating a new U. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + zeroValue : U + the initial value for the accumulated result of each partition + seqFunc : function + a function to merge a V into a U + combFunc : function + a function to combine two U's into a single one + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the aggregated result for each key + + See Also + -------- + :meth:`RDD.reduceByKey` + :meth:`RDD.combineByKey` + :meth:`RDD.foldByKey` + :meth:`RDD.groupByKey` + + Examples + -------- + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 2)]) + >>> seqFunc = (lambda x, y: (x[0] + y, x[1] + 1)) + >>> combFunc = (lambda x, y: (x[0] + y[0], x[1] + y[1])) + >>> sorted(rdd.aggregateByKey((0, 0), seqFunc, combFunc).collect()) + [('a', (3, 2)), ('b', (1, 1))] + """ + + def createZero() -> U: + return copy.deepcopy(zeroValue) + + return self.combineByKey( + lambda v: seqFunc(createZero(), v), seqFunc, combFunc, numPartitions, partitionFunc + ) + + def foldByKey( + self: "RDD[Tuple[K, V]]", + zeroValue: V, + func: Callable[[V, V], V], + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, V]]": + """ + Merge the values for each key using an associative function "func" + and a neutral "zeroValue" which may be added to the result an + arbitrary number of times, and must not change the result + (e.g., 0 for addition, or 1 for multiplication.). + + .. versionadded:: 1.1.0 + + Parameters + ---------- + zeroValue : V + the initial value for the accumulated result of each partition + func : function + a function to combine two V's into a single one + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the aggregated result for each key + + See Also + -------- + :meth:`RDD.reduceByKey` + :meth:`RDD.combineByKey` + :meth:`RDD.aggregateByKey` + :meth:`RDD.groupByKey` + + Examples + -------- + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 1)]) + >>> from operator import add + >>> sorted(rdd.foldByKey(0, add).collect()) + [('a', 2), ('b', 1)] + """ + + def createZero() -> V: + return copy.deepcopy(zeroValue) + + return self.combineByKey( + lambda v: func(createZero(), v), func, func, numPartitions, partitionFunc + ) + + def _memory_limit(self) -> int: + return _parse_memory(self.ctx._conf.get("spark.python.worker.memory", "512m")) + + # TODO: support variant with custom partitioner + def groupByKey( + self: "RDD[Tuple[K, V]]", + numPartitions: Optional[int] = None, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "RDD[Tuple[K, Iterable[V]]]": + """ + Group the values for each key in the RDD into a single sequence. + Hash-partitions the resulting RDD with numPartitions partitions. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + partitionFunc : function, optional, default `portable_hash` + function to compute the partition index + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the grouped result for each key + + See Also + -------- + :meth:`RDD.reduceByKey` + :meth:`RDD.combineByKey` + :meth:`RDD.aggregateByKey` + :meth:`RDD.foldByKey` + + Notes + ----- + If you are grouping in order to perform an aggregation (such as a + sum or average) over each key, using reduceByKey or aggregateByKey will + provide much better performance. + + Examples + -------- + >>> rdd = sc.parallelize([("a", 1), ("b", 1), ("a", 1)]) + >>> sorted(rdd.groupByKey().mapValues(len).collect()) + [('a', 2), ('b', 1)] + >>> sorted(rdd.groupByKey().mapValues(list).collect()) + [('a', [1, 1]), ('b', [1])] + """ + + def createCombiner(x: V) -> List[V]: + return [x] + + def mergeValue(xs: List[V], x: V) -> List[V]: + xs.append(x) + return xs + + def mergeCombiners(a: List[V], b: List[V]) -> List[V]: + a.extend(b) + return a + + memory = self._memory_limit() + serializer = self._jrdd_deserializer + agg = Aggregator(createCombiner, mergeValue, mergeCombiners) + + def combine(iterator: Iterable[Tuple[K, V]]) -> Iterable[Tuple[K, List[V]]]: + merger = ExternalMerger(agg, memory * 0.9, serializer) + merger.mergeValues(iterator) + return merger.items() + + locally_combined = self.mapPartitions(combine, preservesPartitioning=True) + shuffled = locally_combined.partitionBy(numPartitions, partitionFunc) + + def groupByKey(it: Iterable[Tuple[K, List[V]]]) -> Iterable[Tuple[K, List[V]]]: + merger = ExternalGroupBy(agg, memory, serializer) + merger.mergeCombiners(it) + return merger.items() + + return shuffled.mapPartitions(groupByKey, True).mapValues(ResultIterable) + + def flatMapValues( + self: "RDD[Tuple[K, V]]", f: Callable[[V], Iterable[U]] + ) -> "RDD[Tuple[K, U]]": + """ + Pass each value in the key-value pair RDD through a flatMap function + without changing the keys; this also retains the original RDD's + partitioning. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to turn a V into a sequence of U + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the flat-mapped value + + See Also + -------- + :meth:`RDD.flatMap` + :meth:`RDD.mapValues` + + Examples + -------- + >>> rdd = sc.parallelize([("a", ["x", "y", "z"]), ("b", ["p", "r"])]) + >>> def f(x): return x + ... + >>> rdd.flatMapValues(f).collect() + [('a', 'x'), ('a', 'y'), ('a', 'z'), ('b', 'p'), ('b', 'r')] + """ + + def flat_map_fn(kv: Tuple[K, V]) -> Iterable[Tuple[K, U]]: + return ((kv[0], x) for x in f(kv[1])) + + return self.flatMap(flat_map_fn, preservesPartitioning=True) + + def mapValues(self: "RDD[Tuple[K, V]]", f: Callable[[V], U]) -> "RDD[Tuple[K, U]]": + """ + Pass each value in the key-value pair RDD through a map function + without changing the keys; this also retains the original RDD's + partitioning. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + f : function + a function to turn a V into a U + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and the mapped value + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.flatMapValues` + + Examples + -------- + >>> rdd = sc.parallelize([("a", ["apple", "banana", "lemon"]), ("b", ["grapes"])]) + >>> def f(x): return len(x) + ... + >>> rdd.mapValues(f).collect() + [('a', 3), ('b', 1)] + """ + + def map_values_fn(kv: Tuple[K, V]) -> Tuple[K, U]: + return kv[0], f(kv[1]) + + return self.map(map_values_fn, preservesPartitioning=True) + + @overload + def groupWith( + self: "RDD[Tuple[K, V]]", other: "RDD[Tuple[K, V1]]" + ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1]]]]": + ... + + @overload + def groupWith( + self: "RDD[Tuple[K, V]]", other: "RDD[Tuple[K, V1]]", __o1: "RDD[Tuple[K, V2]]" + ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[V1], ResultIterable[V2]]]]": + ... + + @overload + def groupWith( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, V1]]", + _o1: "RDD[Tuple[K, V2]]", + _o2: "RDD[Tuple[K, V3]]", + ) -> """RDD[ + Tuple[ + K, + Tuple[ + ResultIterable[V], + ResultIterable[V1], + ResultIterable[V2], + ResultIterable[V3], + ], + ] + ]""": + ... + + def groupWith( # type: ignore[misc] + self: "RDD[Tuple[Any, Any]]", other: "RDD[Tuple[Any, Any]]", *others: "RDD[Tuple[Any, Any]]" + ) -> "RDD[Tuple[Any, Tuple[ResultIterable[Any], ...]]]": + """ + Alias for cogroup but with support for multiple RDDs. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + others : :class:`RDD` + other :class:`RDD`\\s + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and cogrouped values + + See Also + -------- + :meth:`RDD.cogroup` + :meth:`RDD.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 5), ("b", 6)]) + >>> rdd2 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd3 = sc.parallelize([("a", 2)]) + >>> rdd4 = sc.parallelize([("b", 42)]) + >>> [(x, tuple(map(list, y))) for x, y in + ... sorted(list(rdd1.groupWith(rdd2, rdd3, rdd4).collect()))] + [('a', ([5], [1], [2], [])), ('b', ([6], [4], [], [42]))] + + """ + return python_cogroup((self, other) + others, numPartitions=None) + + # TODO: add variant with custom partitioner + def cogroup( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]": + """ + For each key k in `self` or `other`, return a resulting RDD that + contains a tuple with the list of values for that key in `self` as + well as `other`. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the keys and cogrouped values + + See Also + -------- + :meth:`RDD.groupWith` + :meth:`RDD.join` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4)]) + >>> rdd2 = sc.parallelize([("a", 2)]) + >>> [(x, tuple(map(list, y))) for x, y in sorted(list(rdd1.cogroup(rdd2).collect()))] + [('a', ([1], [2])), ('b', ([4], []))] + """ + return python_cogroup((self, other), numPartitions) + + def sampleByKey( + self: "RDD[Tuple[K, V]]", + withReplacement: bool, + fractions: Dict[K, Union[float, int]], + seed: Optional[int] = None, + ) -> "RDD[Tuple[K, V]]": + """ + Return a subset of this RDD sampled by key (via stratified sampling). + Create a sample of this RDD using variable sampling rates for + different keys as specified by fractions, a key to sampling rate map. + + .. versionadded:: 0.7.0 + + Parameters + ---------- + withReplacement : bool + whether to sample with or without replacement + fractions : dict + map of specific keys to sampling rates + seed : int, optional + seed for the random number generator + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the stratified sampling result + + See Also + -------- + :meth:`RDD.sample` + + Examples + -------- + >>> fractions = {"a": 0.2, "b": 0.1} + >>> rdd = sc.parallelize(fractions.keys()).cartesian(sc.parallelize(range(0, 1000))) + >>> sample = dict(rdd.sampleByKey(False, fractions, 2).groupByKey().collect()) + >>> 100 < len(sample["a"]) < 300 and 50 < len(sample["b"]) < 150 + True + >>> max(sample["a"]) <= 999 and min(sample["a"]) >= 0 + True + >>> max(sample["b"]) <= 999 and min(sample["b"]) >= 0 + True + """ + for fraction in fractions.values(): + assert fraction >= 0.0, "Negative fraction value: %s" % fraction + return self.mapPartitionsWithIndex( + RDDStratifiedSampler(withReplacement, fractions, seed).func, True + ) + + def subtractByKey( + self: "RDD[Tuple[K, V]]", + other: "RDD[Tuple[K, Any]]", + numPartitions: Optional[int] = None, + ) -> "RDD[Tuple[K, V]]": + """ + Return each (key, value) pair in `self` that has no pair with matching + key in `other`. + + .. versionadded:: 0.9.1 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` with the pairs from this whose keys are not in `other` + + See Also + -------- + :meth:`RDD.subtract` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4), ("b", 5), ("a", 2)]) + >>> rdd2 = sc.parallelize([("a", 3), ("c", None)]) + >>> sorted(rdd1.subtractByKey(rdd2).collect()) + [('b', 4), ('b', 5)] + """ + + def filter_func(pair: Tuple[K, Tuple[V, Any]]) -> bool: + key, (val1, val2) = pair + return val1 and not val2 # type: ignore[return-value] + + return ( + self.cogroup(other, numPartitions) + .filter(filter_func) # type: ignore[arg-type] + .flatMapValues(lambda x: x[0]) + ) + + def subtract(self: "RDD[T]", other: "RDD[T]", numPartitions: Optional[int] = None) -> "RDD[T]": + """ + Return each value in `self` that is not contained in `other`. + + .. versionadded:: 0.9.1 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` with the elements from this that are not in `other` + + See Also + -------- + :meth:`RDD.subtractByKey` + + Examples + -------- + >>> rdd1 = sc.parallelize([("a", 1), ("b", 4), ("b", 5), ("a", 3)]) + >>> rdd2 = sc.parallelize([("a", 3), ("c", None)]) + >>> sorted(rdd1.subtract(rdd2).collect()) + [('a', 1), ('b', 4), ('b', 5)] + """ + # note: here 'True' is just a placeholder + rdd = other.map(lambda x: (x, True)) + return self.map(lambda x: (x, True)).subtractByKey(rdd, numPartitions).keys() + + def keyBy(self: "RDD[T]", f: Callable[[T], K]) -> "RDD[Tuple[K, T]]": + """ + Creates tuples of the elements in this RDD by applying `f`. + + .. versionadded:: 0.9.1 + + Parameters + ---------- + f : function + a function to compute the key + + Returns + ------- + :class:`RDD` + a :class:`RDD` with the elements from this that are not in `other` + + See Also + -------- + :meth:`RDD.map` + :meth:`RDD.keys` + :meth:`RDD.values` + + Examples + -------- + >>> rdd1 = sc.parallelize(range(0,3)).keyBy(lambda x: x*x) + >>> rdd2 = sc.parallelize(zip(range(0,5), range(0,5))) + >>> [(x, list(map(list, y))) for x, y in sorted(rdd1.cogroup(rdd2).collect())] + [(0, [[0], [0]]), (1, [[1], [1]]), (2, [[], [2]]), (3, [[], [3]]), (4, [[2], [4]])] + """ + return self.map(lambda x: (f(x), x)) + + def repartition(self: "RDD[T]", numPartitions: int) -> "RDD[T]": + """ + Return a new RDD that has exactly numPartitions partitions. + + Can increase or decrease the level of parallelism in this RDD. + Internally, this uses a shuffle to redistribute data. + If you are decreasing the number of partitions in this RDD, consider + using `coalesce`, which can avoid performing a shuffle. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` with exactly numPartitions partitions + + See Also + -------- + :meth:`RDD.coalesce` + :meth:`RDD.partitionBy` + :meth:`RDD.repartitionAndSortWithinPartitions` + + Examples + -------- + >>> rdd = sc.parallelize([1,2,3,4,5,6,7], 4) + >>> sorted(rdd.glom().collect()) + [[1], [2, 3], [4, 5], [6, 7]] + >>> len(rdd.repartition(2).glom().collect()) + 2 + >>> len(rdd.repartition(10).glom().collect()) + 10 + """ + return self.coalesce(numPartitions, shuffle=True) + + def coalesce(self: "RDD[T]", numPartitions: int, shuffle: bool = False) -> "RDD[T]": + """ + Return a new RDD that is reduced into `numPartitions` partitions. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + numPartitions : int, optional + the number of partitions in new :class:`RDD` + shuffle : bool, optional, default False + whether to add a shuffle step + + Returns + ------- + :class:`RDD` + a :class:`RDD` that is reduced into `numPartitions` partitions + + See Also + -------- + :meth:`RDD.repartition` + + Examples + -------- + >>> sc.parallelize([1, 2, 3, 4, 5], 3).glom().collect() + [[1], [2, 3], [4, 5]] + >>> sc.parallelize([1, 2, 3, 4, 5], 3).coalesce(1).glom().collect() + [[1, 2, 3, 4, 5]] + """ + if not numPartitions > 0: + raise ValueError("Number of partitions must be positive.") + if shuffle: + # Decrease the batch size in order to distribute evenly the elements across output + # partitions. Otherwise, repartition will possibly produce highly skewed partitions. + batchSize = min(10, self.ctx._batchSize or 1024) + ser = BatchedSerializer(CPickleSerializer(), batchSize) + selfCopy = self._reserialize(ser) + jrdd_deserializer = selfCopy._jrdd_deserializer + jrdd = selfCopy._jrdd.coalesce(numPartitions, shuffle) + else: + jrdd_deserializer = self._jrdd_deserializer + jrdd = self._jrdd.coalesce(numPartitions, shuffle) + return RDD(jrdd, self.ctx, jrdd_deserializer) + + def zip(self: "RDD[T]", other: "RDD[U]") -> "RDD[Tuple[T, U]]": + """ + Zips this RDD with another one, returning key-value pairs with the + first element in each RDD second element in each RDD, etc. Assumes + that the two RDDs have the same number of partitions and the same + number of elements in each partition (e.g. one was made through + a map on the other). + + .. versionadded:: 1.0.0 + + Parameters + ---------- + other : :class:`RDD` + another :class:`RDD` + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the zipped key-value pairs + + See Also + -------- + :meth:`RDD.zipWithIndex` + :meth:`RDD.zipWithUniqueId` + + Examples + -------- + >>> rdd1 = sc.parallelize(range(0,5)) + >>> rdd2 = sc.parallelize(range(1000, 1005)) + >>> rdd1.zip(rdd2).collect() + [(0, 1000), (1, 1001), (2, 1002), (3, 1003), (4, 1004)] + """ + + def get_batch_size(ser: Serializer) -> int: + if isinstance(ser, BatchedSerializer): + return ser.batchSize + return 1 # not batched + + def batch_as(rdd: "RDD[V]", batchSize: int) -> "RDD[V]": + return rdd._reserialize(BatchedSerializer(CPickleSerializer(), batchSize)) + + my_batch = get_batch_size(self._jrdd_deserializer) + other_batch = get_batch_size(other._jrdd_deserializer) + if my_batch != other_batch or not my_batch: + # use the smallest batchSize for both of them + batchSize = min(my_batch, other_batch) + if batchSize <= 0: + # auto batched or unlimited + batchSize = 100 + other = batch_as(other, batchSize) + self = batch_as(self, batchSize) + + if self.getNumPartitions() != other.getNumPartitions(): + raise ValueError("Can only zip with RDD which has the same number of partitions") + + # There will be an Exception in JVM if there are different number + # of items in each partitions. + pairRDD = self._jrdd.zip(other._jrdd) + deserializer = PairDeserializer(self._jrdd_deserializer, other._jrdd_deserializer) + return RDD(pairRDD, self.ctx, deserializer) + + def zipWithIndex(self: "RDD[T]") -> "RDD[Tuple[T, int]]": + """ + Zips this RDD with its element indices. + + The ordering is first based on the partition index and then the + ordering of items within each partition. So the first item in + the first partition gets index 0, and the last item in the last + partition receives the largest index. + + This method needs to trigger a spark job when this RDD contains + more than one partitions. + + .. versionadded:: 1.2.0 + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the zipped key-index pairs + + See Also + -------- + :meth:`RDD.zip` + :meth:`RDD.zipWithUniqueId` + + Examples + -------- + >>> sc.parallelize(["a", "b", "c", "d"], 3).zipWithIndex().collect() + [('a', 0), ('b', 1), ('c', 2), ('d', 3)] + """ + starts = [0] + if self.getNumPartitions() > 1: + nums = self.mapPartitions(lambda it: [sum(1 for i in it)]).collect() + for i in range(len(nums) - 1): + starts.append(starts[-1] + nums[i]) + + def func(k: int, it: Iterable[T]) -> Iterable[Tuple[T, int]]: + for i, v in enumerate(it, starts[k]): + yield v, i + + return self.mapPartitionsWithIndex(func) + + def zipWithUniqueId(self: "RDD[T]") -> "RDD[Tuple[T, int]]": + """ + Zips this RDD with generated unique Long ids. + + Items in the kth partition will get ids k, n+k, 2*n+k, ..., where + n is the number of partitions. So there may exist gaps, but this + method won't trigger a spark job, which is different from + :meth:`zipWithIndex`. + + .. versionadded:: 1.2.0 + + Returns + ------- + :class:`RDD` + a :class:`RDD` containing the zipped key-UniqueId pairs + + See Also + -------- + :meth:`RDD.zip` + :meth:`RDD.zipWithIndex` + + Examples + -------- + >>> sc.parallelize(["a", "b", "c", "d", "e"], 3).zipWithUniqueId().collect() + [('a', 0), ('b', 1), ('c', 4), ('d', 2), ('e', 5)] + """ + n = self.getNumPartitions() + + def func(k: int, it: Iterable[T]) -> Iterable[Tuple[T, int]]: + for i, v in enumerate(it): + yield v, i * n + k + + return self.mapPartitionsWithIndex(func) + + def name(self) -> Optional[str]: + """ + Return the name of this RDD. + + .. versionadded:: 1.0.0 + + Returns + ------- + str + :class:`RDD` name + + See Also + -------- + :meth:`RDD.setName` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.name() == None + True + """ + n = self._jrdd.name() + return n if n else None + + def setName(self: "RDD[T]", name: str) -> "RDD[T]": + """ + Assign a name to this RDD. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + name : str + new name + + Returns + ------- + :class:`RDD` + the same :class:`RDD` with name updated + + See Also + -------- + :meth:`RDD.name` + + Examples + -------- + >>> rdd = sc.parallelize([1, 2]) + >>> rdd.setName('I am an RDD').name() + 'I am an RDD' + """ + self._jrdd.setName(name) + return self + + def toDebugString(self) -> Optional[bytes]: + """ + A description of this RDD and its recursive dependencies for debugging. + + .. versionadded:: 1.0.0 + + Returns + ------- + bytes + debugging information of this :class:`RDD` + + Examples + -------- + >>> rdd = sc.range(5) + >>> rdd.toDebugString() + b'...PythonRDD...ParallelCollectionRDD...' + """ + debug_string = self._jrdd.toDebugString() + + return debug_string.encode("utf-8") if debug_string else None + + def getStorageLevel(self) -> StorageLevel: + """ + Get the RDD's current storage level. + + .. versionadded:: 1.0.0 + + Returns + ------- + :class:`StorageLevel` + current :class:`StorageLevel` + + See Also + -------- + :meth:`RDD.name` + + Examples + -------- + >>> rdd = sc.parallelize([1,2]) + >>> rdd.getStorageLevel() + StorageLevel(False, False, False, False, 1) + >>> print(rdd.getStorageLevel()) + Serialized 1x Replicated + """ + java_storage_level = self._jrdd.getStorageLevel() + storage_level = StorageLevel( + java_storage_level.useDisk(), + java_storage_level.useMemory(), + java_storage_level.useOffHeap(), + java_storage_level.deserialized(), + java_storage_level.replication(), + ) + return storage_level + + def _defaultReducePartitions(self) -> int: + """ + Returns the default number of partitions to use during reduce tasks (e.g., groupBy). + If spark.default.parallelism is set, then we'll use the value from SparkContext + defaultParallelism, otherwise we'll use the number of partitions in this RDD. + + This mirrors the behavior of the Scala Partitioner#defaultPartitioner, intended to reduce + the likelihood of OOMs. Once PySpark adopts Partitioner-based APIs, this behavior will + be inherent. + """ + if self.ctx._conf.contains("spark.default.parallelism"): + return self.ctx.defaultParallelism + else: + return self.getNumPartitions() + + def lookup(self: "RDD[Tuple[K, V]]", key: K) -> List[V]: + """ + Return the list of values in the RDD for key `key`. This operation + is done efficiently if the RDD has a known partitioner by only + searching the partition that the key maps to. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + key : K + the key to look up + + Returns + ------- + list + the list of values in the :class:`RDD` for key `key` + + Examples + -------- + >>> l = range(1000) + >>> rdd = sc.parallelize(zip(l, l), 10) + >>> rdd.lookup(42) # slow + [42] + >>> sorted = rdd.sortByKey() + >>> sorted.lookup(42) # fast + [42] + >>> sorted.lookup(1024) + [] + >>> rdd2 = sc.parallelize([(('a', 'b'), 'c')]).groupByKey() + >>> list(rdd2.lookup(('a', 'b'))[0]) + ['c'] + """ + values = self.filter(lambda kv: kv[0] == key).values() + + if self.partitioner is not None: + return self.ctx.runJob(values, lambda x: x, [self.partitioner(key)]) + + return values.collect() + + def _to_java_object_rdd(self) -> "JavaObject": + """Return a JavaRDD of Object by unpickling + + It will convert each Python object into Java object by Pickle, whenever the + RDD is serialized in batch or not. + """ + rdd = self._pickled() + assert self.ctx._jvm is not None + + return self.ctx._jvm.SerDeUtil.pythonToJava(rdd._jrdd, True) + + def countApprox(self, timeout: int, confidence: float = 0.95) -> int: + """ + Approximate version of count() that returns a potentially incomplete + result within a timeout, even if not all tasks have finished. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + timeout : int + maximum time to wait for the job, in milliseconds + confidence : float + the desired statistical confidence in the result + + Returns + ------- + int + a potentially incomplete result, with error bounds + + See Also + -------- + :meth:`RDD.count` + + Examples + -------- + >>> rdd = sc.parallelize(range(1000), 10) + >>> rdd.countApprox(1000, 1.0) + 1000 + """ + drdd = self.mapPartitions(lambda it: [float(sum(1 for i in it))]) + return int(drdd.sumApprox(timeout, confidence)) + + def sumApprox( + self: "RDD[Union[float, int]]", timeout: int, confidence: float = 0.95 + ) -> BoundedFloat: + """ + Approximate operation to return the sum within a timeout + or meet the confidence. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + timeout : int + maximum time to wait for the job, in milliseconds + confidence : float + the desired statistical confidence in the result + + Returns + ------- + :class:`BoundedFloat` + a potentially incomplete result, with error bounds + + See Also + -------- + :meth:`RDD.sum` + + Examples + -------- + >>> rdd = sc.parallelize(range(1000), 10) + >>> r = sum(range(1000)) + >>> abs(rdd.sumApprox(1000) - r) / r < 0.05 + True + """ + jrdd = self.mapPartitions(lambda it: [float(sum(it))])._to_java_object_rdd() + assert self.ctx._jvm is not None + jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd()) + r = jdrdd.sumApprox(timeout, confidence).getFinalValue() + return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high()) + + def meanApprox( + self: "RDD[Union[float, int]]", timeout: int, confidence: float = 0.95 + ) -> BoundedFloat: + """ + Approximate operation to return the mean within a timeout + or meet the confidence. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + timeout : int + maximum time to wait for the job, in milliseconds + confidence : float + the desired statistical confidence in the result + + Returns + ------- + :class:`BoundedFloat` + a potentially incomplete result, with error bounds + + See Also + -------- + :meth:`RDD.mean` + + Examples + -------- + >>> rdd = sc.parallelize(range(1000), 10) + >>> r = sum(range(1000)) / 1000.0 + >>> abs(rdd.meanApprox(1000) - r) / r < 0.05 + True + """ + jrdd = self.map(float)._to_java_object_rdd() + assert self.ctx._jvm is not None + jdrdd = self.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd()) + r = jdrdd.meanApprox(timeout, confidence).getFinalValue() + return BoundedFloat(r.mean(), r.confidence(), r.low(), r.high()) + + def countApproxDistinct(self: "RDD[T]", relativeSD: float = 0.05) -> int: + """ + Return approximate number of distinct elements in the RDD. + + .. versionadded:: 1.2.0 + + Parameters + ---------- + relativeSD : float, optional + Relative accuracy. Smaller values create + counters that require more space. + It must be greater than 0.000017. + + Returns + ------- + int + approximate number of distinct elements + + See Also + -------- + :meth:`RDD.distinct` + + Notes + ----- + The algorithm used is based on streamlib's implementation of + `"HyperLogLog in Practice: Algorithmic Engineering of a State + of The Art Cardinality Estimation Algorithm", available here + `_. + + Examples + -------- + >>> n = sc.parallelize(range(1000)).map(str).countApproxDistinct() + >>> 900 < n < 1100 + True + >>> n = sc.parallelize([i % 20 for i in range(1000)]).countApproxDistinct() + >>> 16 < n < 24 + True + """ + if relativeSD < 0.000017: + raise ValueError("relativeSD should be greater than 0.000017") + # the hash space in Java is 2^32 + hashRDD = self.map(lambda x: portable_hash(x) & 0xFFFFFFFF) + return hashRDD._to_java_object_rdd().countApproxDistinct(relativeSD) + + def toLocalIterator(self: "RDD[T]", prefetchPartitions: bool = False) -> Iterator[T]: + """ + Return an iterator that contains all of the elements in this RDD. + The iterator will consume as much memory as the largest partition in this RDD. + With prefetch it may consume up to the memory of the 2 largest partitions. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + prefetchPartitions : bool, optional + If Spark should pre-fetch the next partition + before it is needed. + + Returns + ------- + :class:`collections.abc.Iterator` + an iterator that contains all of the elements in this :class:`RDD` + + See Also + -------- + :meth:`RDD.collect` + :meth:`pyspark.sql.DataFrame.toLocalIterator` + + Examples + -------- + >>> rdd = sc.parallelize(range(10)) + >>> [x for x in rdd.toLocalIterator()] + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + """ + assert self.ctx._jvm is not None + + with SCCallSiteSync(self.context): + sock_info = self.ctx._jvm.PythonRDD.toLocalIteratorAndServe( + self._jrdd.rdd(), prefetchPartitions + ) + return _local_iterator_from_socket(sock_info, self._jrdd_deserializer) + + def barrier(self: "RDD[T]") -> "RDDBarrier[T]": + """ + Marks the current stage as a barrier stage, where Spark must launch all tasks together. + In case of a task failure, instead of only restarting the failed task, Spark will abort the + entire stage and relaunch all tasks for this stage. + The barrier execution mode feature is experimental and it only handles limited scenarios. + Please read the linked SPIP and design docs to understand the limitations and future plans. + + .. versionadded:: 2.4.0 + + Returns + ------- + :class:`RDDBarrier` + instance that provides actions within a barrier stage. + + See Also + -------- + :class:`pyspark.BarrierTaskContext` + + Notes + ----- + For additional information see + + - `SPIP: Barrier Execution Mode `_ + - `Design Doc `_ + + This API is experimental + """ + return RDDBarrier(self) + + def _is_barrier(self) -> bool: + """ + Whether this RDD is in a barrier stage. + """ + return self._jrdd.rdd().isBarrier() + + def withResources(self: "RDD[T]", profile: ResourceProfile) -> "RDD[T]": + """ + Specify a :class:`pyspark.resource.ResourceProfile` to use when calculating this RDD. + This is only supported on certain cluster managers and currently requires dynamic + allocation to be enabled. It will result in new executors with the resources specified + being acquired to calculate the RDD. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + profile : :class:`pyspark.resource.ResourceProfile` + a resource profile + + Returns + ------- + :class:`RDD` + the same :class:`RDD` with user specified profile + + See Also + -------- + :meth:`RDD.getResourceProfile` + + Notes + ----- + This API is experimental + """ + self.has_resource_profile = True + if profile._java_resource_profile is not None: + jrp = profile._java_resource_profile + else: + assert self.ctx._jvm is not None + + builder = self.ctx._jvm.org.apache.spark.resource.ResourceProfileBuilder() + ereqs = ExecutorResourceRequests(self.ctx._jvm, profile._executor_resource_requests) + treqs = TaskResourceRequests(self.ctx._jvm, profile._task_resource_requests) + builder.require(ereqs._java_executor_resource_requests) + builder.require(treqs._java_task_resource_requests) + jrp = builder.build() + + self._jrdd.withResources(jrp) + return self + + def getResourceProfile(self) -> Optional[ResourceProfile]: + """ + Get the :class:`pyspark.resource.ResourceProfile` specified with this RDD or None + if it wasn't specified. + + .. versionadded:: 3.1.0 + + Returns + ------- + class:`pyspark.resource.ResourceProfile` + The user specified profile or None if none were specified + + See Also + -------- + :meth:`RDD.withResources` + + Notes + ----- + This API is experimental + """ + rp = self._jrdd.getResourceProfile() + if rp is not None: + return ResourceProfile(_java_resource_profile=rp) + else: + return None + + @overload + def toDF( + self: "RDD[RowLike]", + schema: Optional[Union[List[str], Tuple[str, ...]]] = None, + sampleRatio: Optional[float] = None, + ) -> "DataFrame": + ... + + @overload + def toDF( + self: "RDD[RowLike]", schema: Optional[Union["StructType", str]] = None + ) -> "DataFrame": + ... + + @overload + def toDF( + self: "RDD[AtomicValue]", + schema: Union["AtomicType", str], + ) -> "DataFrame": + ... + + def toDF( + self: "RDD[Any]", schema: Optional[Any] = None, sampleRatio: Optional[float] = None + ) -> "DataFrame": + raise PySparkRuntimeError( + error_class="CALL_BEFORE_INITIALIZE", + message_parameters={ + "func_name": "RDD.toDF", + "object": "SparkSession", + }, + ) + + +def _prepare_for_python_RDD(sc: "SparkContext", command: Any) -> Tuple[bytes, Any, Any, Any]: + # the serialized command will be compressed by broadcast + ser = CloudPickleSerializer() + pickled_command = ser.dumps(command) + assert sc._jvm is not None + if len(pickled_command) > sc._jvm.PythonUtils.getBroadcastThreshold(sc._jsc): # Default 1M + # The broadcast will have same life cycle as created PythonRDD + broadcast = sc.broadcast(pickled_command) + pickled_command = ser.dumps(broadcast) + broadcast_vars = [x._jbroadcast for x in sc._pickled_broadcast_vars] + sc._pickled_broadcast_vars.clear() + return pickled_command, broadcast_vars, sc.environment, sc._python_includes + + +def _wrap_function( + sc: "SparkContext", func: Callable, deserializer: Any, serializer: Any, profiler: Any = None +) -> "JavaObject": + assert deserializer, "deserializer should not be empty" + assert serializer, "serializer should not be empty" + command = (func, profiler, deserializer, serializer) + pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command) + assert sc._jvm is not None + return sc._jvm.SimplePythonFunction( + bytearray(pickled_command), + env, + includes, + sc.pythonExec, + sc.pythonVer, + broadcast_vars, + sc._javaAccumulator, + ) + + +class RDDBarrier(Generic[T]): + + """ + Wraps an RDD in a barrier stage, which forces Spark to launch tasks of this stage together. + :class:`RDDBarrier` instances are created by :meth:`RDD.barrier`. + + .. versionadded:: 2.4.0 + + Notes + ----- + This API is experimental + """ + + def __init__(self, rdd: RDD[T]): + self.rdd = rdd + + def mapPartitions( + self, f: Callable[[Iterable[T]], Iterable[U]], preservesPartitioning: bool = False + ) -> RDD[U]: + """ + Returns a new RDD by applying a function to each partition of the wrapped RDD, + where tasks are launched together in a barrier stage. + The interface is the same as :meth:`RDD.mapPartitions`. + Please see the API doc there. + + .. versionadded:: 2.4.0 + + Parameters + ---------- + f : function + a function to run on each partition of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each partition + + See Also + -------- + :meth:`RDD.mapPartitions` + + Notes + ----- + This API is experimental + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 2) + >>> def f(iterator): yield sum(iterator) + ... + >>> barrier = rdd.barrier() + >>> barrier + + >>> barrier.mapPartitions(f).collect() + [3, 7] + """ + + def func(s: int, iterator: Iterable[T]) -> Iterable[U]: + return f(iterator) + + return PipelinedRDD(self.rdd, func, preservesPartitioning, isFromBarrier=True) + + def mapPartitionsWithIndex( + self, + f: Callable[[int, Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + ) -> RDD[U]: + """ + Returns a new RDD by applying a function to each partition of the wrapped RDD, while + tracking the index of the original partition. And all tasks are launched together + in a barrier stage. + The interface is the same as :meth:`RDD.mapPartitionsWithIndex`. + Please see the API doc there. + + .. versionadded:: 3.0.0 + + Parameters + ---------- + f : function + a function to run on each partition of the RDD + preservesPartitioning : bool, optional, default False + indicates whether the input function preserves the partitioner, + which should be False unless this is a pair RDD and the input + + Returns + ------- + :class:`RDD` + a new :class:`RDD` by applying a function to each partition + + See Also + -------- + :meth:`RDD.mapPartitionsWithIndex` + + Notes + ----- + This API is experimental + + Examples + -------- + >>> rdd = sc.parallelize([1, 2, 3, 4], 4) + >>> def f(splitIndex, iterator): yield splitIndex + ... + >>> barrier = rdd.barrier() + >>> barrier + + >>> barrier.mapPartitionsWithIndex(f).sum() + 6 + """ + return PipelinedRDD(self.rdd, f, preservesPartitioning, isFromBarrier=True) + + +class PipelinedRDD(RDD[U], Generic[T, U]): + + """ + Examples + -------- + Pipelined maps: + + >>> rdd = sc.parallelize([1, 2, 3, 4]) + >>> rdd.map(lambda x: 2 * x).cache().map(lambda x: 2 * x).collect() + [4, 8, 12, 16] + >>> rdd.map(lambda x: 2 * x).map(lambda x: 2 * x).collect() + [4, 8, 12, 16] + + Pipelined reduces: + + >>> from operator import add + >>> rdd.map(lambda x: 2 * x).reduce(add) + 20 + >>> rdd.flatMap(lambda x: [x, x]).reduce(add) + 20 + """ + + def __init__( + self, + prev: RDD[T], + func: Callable[[int, Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + isFromBarrier: bool = False, + ): + if not isinstance(prev, PipelinedRDD) or not prev._is_pipelinable(): + # This transformation is the first in its stage: + self.func = func + self.preservesPartitioning = preservesPartitioning + self._prev_jrdd = prev._jrdd + self._prev_jrdd_deserializer = prev._jrdd_deserializer + else: + prev_func: Callable[[int, Iterable[V]], Iterable[T]] = prev.func + + def pipeline_func(split: int, iterator: Iterable[V]) -> Iterable[U]: + return func(split, prev_func(split, iterator)) + + self.func = pipeline_func + self.preservesPartitioning = prev.preservesPartitioning and preservesPartitioning + self._prev_jrdd = prev._prev_jrdd # maintain the pipeline + self._prev_jrdd_deserializer = prev._prev_jrdd_deserializer + self.is_cached = False + self.has_resource_profile = False + self.is_checkpointed = False + self.ctx = prev.ctx + self.prev = prev + self._jrdd_val: Optional["JavaObject"] = None + self._id = None + self._jrdd_deserializer = self.ctx.serializer + self._bypass_serializer = False + self.partitioner = prev.partitioner if self.preservesPartitioning else None + self.is_barrier = isFromBarrier or prev._is_barrier() + + def getNumPartitions(self) -> int: + return self._prev_jrdd.partitions().size() + + @property + def _jrdd(self) -> "JavaObject": + if self._jrdd_val: + return self._jrdd_val + if self._bypass_serializer: + self._jrdd_deserializer = NoOpSerializer() + + if ( + self.ctx.profiler_collector + and self.ctx._conf.get("spark.python.profile", "false") == "true" + ): + profiler = self.ctx.profiler_collector.new_profiler(self.ctx) + else: + profiler = None + + wrapped_func = _wrap_function( + self.ctx, self.func, self._prev_jrdd_deserializer, self._jrdd_deserializer, profiler + ) + + assert self.ctx._jvm is not None + python_rdd = self.ctx._jvm.PythonRDD( + self._prev_jrdd.rdd(), wrapped_func, self.preservesPartitioning, self.is_barrier + ) + self._jrdd_val = python_rdd.asJavaRDD() + + if profiler: + assert self._jrdd_val is not None + self._id = self._jrdd_val.id() + self.ctx.profiler_collector.add_profiler(self._id, profiler) + return self._jrdd_val + + def id(self) -> int: + if self._id is None: + self._id = self._jrdd.id() + return self._id + + def _is_pipelinable(self) -> bool: + return not (self.is_cached or self.is_checkpointed or self.has_resource_profile) + + def _is_barrier(self) -> bool: + return self.is_barrier + + +def _test() -> None: + import doctest + import tempfile + from pyspark.context import SparkContext + + tmp_dir = tempfile.TemporaryDirectory() + globs = globals().copy() + # The small batch size here ensures that we see multiple batches, + # even in these small test examples: + globs["sc"] = SparkContext("local[4]", "PythonTest") + globs["sc"].setCheckpointDir(tmp_dir.name) + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["sc"].stop() + tmp_dir.cleanup() + if failure_count: + tmp_dir.cleanup() + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/rddsampler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/rddsampler.py new file mode 100644 index 0000000..bf82818 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/rddsampler.py @@ -0,0 +1,115 @@ +# +# 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 sys +import random +import math + + +class RDDSamplerBase: + def __init__(self, withReplacement, seed=None): + self._seed = seed if seed is not None else random.randint(0, sys.maxsize) + self._withReplacement = withReplacement + self._random = None + + def initRandomGenerator(self, split): + self._random = random.Random(self._seed ^ split) + + # mixing because the initial seeds are close to each other + for _ in range(10): + self._random.randint(0, 1) + + def getUniformSample(self): + return self._random.random() + + def getPoissonSample(self, mean): + # Using Knuth's algorithm described in + # http://en.wikipedia.org/wiki/Poisson_distribution + if mean < 20.0: + # one exp and k+1 random calls + lda = math.exp(-mean) + p = self._random.random() + k = 0 + while p > lda: + k += 1 + p *= self._random.random() + else: + # switch to the log domain, k+1 expovariate (random + log) calls + p = self._random.expovariate(mean) + k = 0 + while p < 1.0: + k += 1 + p += self._random.expovariate(mean) + return k + + def func(self, split, iterator): + raise NotImplementedError + + +class RDDSampler(RDDSamplerBase): + def __init__(self, withReplacement, fraction, seed=None): + RDDSamplerBase.__init__(self, withReplacement, seed) + self._fraction = fraction + + def func(self, split, iterator): + self.initRandomGenerator(split) + if self._withReplacement: + for obj in iterator: + # For large datasets, the expected number of occurrences of each element in + # a sample with replacement is Poisson(frac). We use that to get a count for + # each element. + count = self.getPoissonSample(self._fraction) + for _ in range(0, count): + yield obj + else: + for obj in iterator: + if self.getUniformSample() < self._fraction: + yield obj + + +class RDDRangeSampler(RDDSamplerBase): + def __init__(self, lowerBound, upperBound, seed=None): + RDDSamplerBase.__init__(self, False, seed) + self._lowerBound = lowerBound + self._upperBound = upperBound + + def func(self, split, iterator): + self.initRandomGenerator(split) + for obj in iterator: + if self._lowerBound <= self.getUniformSample() < self._upperBound: + yield obj + + +class RDDStratifiedSampler(RDDSamplerBase): + def __init__(self, withReplacement, fractions, seed=None): + RDDSamplerBase.__init__(self, withReplacement, seed) + self._fractions = fractions + + def func(self, split, iterator): + self.initRandomGenerator(split) + if self._withReplacement: + for key, val in iterator: + # For large datasets, the expected number of occurrences of each element in + # a sample with replacement is Poisson(frac). We use that to get a count for + # each element. + count = self.getPoissonSample(self._fractions[key]) + for _ in range(0, count): + yield key, val + else: + for key, val in iterator: + if self.getUniformSample() < self._fractions[key]: + yield key, val diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resource/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/__init__.py new file mode 100644 index 0000000..b4c4634 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/__init__.py @@ -0,0 +1,38 @@ +# +# 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. +# + +""" +APIs to let users manipulate resource requirements. +""" +from pyspark.resource.information import ResourceInformation +from pyspark.resource.requests import ( + TaskResourceRequest, + TaskResourceRequests, + ExecutorResourceRequest, + ExecutorResourceRequests, +) +from pyspark.resource.profile import ResourceProfile, ResourceProfileBuilder + +__all__ = [ + "TaskResourceRequest", + "TaskResourceRequests", + "ExecutorResourceRequest", + "ExecutorResourceRequests", + "ResourceProfile", + "ResourceInformation", + "ResourceProfileBuilder", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resource/information.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/information.py new file mode 100644 index 0000000..92cfc5a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/information.py @@ -0,0 +1,69 @@ +# +# 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 List + + +class ResourceInformation: + + """ + Class to hold information about a type of Resource. A resource could be a GPU, FPGA, etc. + The array of addresses are resource specific and its up to the user to interpret the address. + + One example is GPUs, where the addresses would be the indices of the GPUs + + .. versionadded:: 3.0.0 + + Parameters + ---------- + name : str + the name of the resource + addresses : list + a list of strings describing the addresses of the resource + + Notes + ----- + This API is evolving. + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + """ + + def __init__(self, name: str, addresses: List[str]): + self._name = name + self._addresses = addresses + + @property + def name(self) -> str: + """ + Returns + ------- + str + the name of the resource + """ + return self._name + + @property + def addresses(self) -> List[str]: + """ + Returns + ------- + list + a list of strings describing the addresses of the resource + """ + return self._addresses diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resource/profile.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/profile.py new file mode 100644 index 0000000..0b2de44 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/profile.py @@ -0,0 +1,317 @@ +# +# 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 overload, Dict, Union, Optional +from py4j.java_gateway import JavaObject + +from pyspark.resource.requests import ( + TaskResourceRequest, + TaskResourceRequests, + ExecutorResourceRequests, + ExecutorResourceRequest, +) + + +class ResourceProfile: + + """ + Resource profile to associate with an RDD. A :class:`pyspark.resource.ResourceProfile` + allows the user to specify executor and task requirements for an RDD that will get + applied during a stage. This allows the user to change the resource requirements between + stages. This is meant to be immutable so user cannot change it after building. + + .. versionadded:: 3.1.0 + + Notes + ----- + This API is evolving. + + Examples + -------- + Create Executor resource requests. + + >>> executor_requests = ( + ... ExecutorResourceRequests() + ... .cores(2) + ... .memory("6g") + ... .memoryOverhead("1g") + ... .pysparkMemory("2g") + ... .offheapMemory("3g") + ... .resource("gpu", 2, "testGpus", "nvidia.com") + ... ) + + Create task resource requasts. + + >>> task_requests = TaskResourceRequests().cpus(2).resource("gpu", 2) + + Create a resource profile. + + >>> builder = ResourceProfileBuilder() + >>> resource_profile = builder.require(executor_requests).require(task_requests).build + + Create an RDD with the resource profile. + + >>> rdd = sc.parallelize(range(10)).withResources(resource_profile) + >>> rdd.getResourceProfile() + + >>> rdd.getResourceProfile().taskResources + {'cpus': <...TaskResourceRequest...>, 'gpu': <...TaskResourceRequest...>} + >>> rdd.getResourceProfile().executorResources + {'gpu': <...ExecutorResourceRequest...>, + 'cores': <...ExecutorResourceRequest...>, + 'offHeap': <...ExecutorResourceRequest...>, + 'memoryOverhead': <...ExecutorResourceRequest...>, + 'pyspark.memory': <...ExecutorResourceRequest...>, + 'memory': <...ExecutorResourceRequest...>} + """ + + @overload + def __init__(self, _java_resource_profile: JavaObject): + ... + + @overload + def __init__( + self, + _java_resource_profile: None = ..., + _exec_req: Optional[Dict[str, ExecutorResourceRequest]] = ..., + _task_req: Optional[Dict[str, TaskResourceRequest]] = ..., + ): + ... + + def __init__( + self, + _java_resource_profile: Optional[JavaObject] = None, + _exec_req: Optional[Dict[str, ExecutorResourceRequest]] = None, + _task_req: Optional[Dict[str, TaskResourceRequest]] = None, + ): + if _java_resource_profile is not None: + self._java_resource_profile = _java_resource_profile + else: + self._java_resource_profile = None + self._executor_resource_requests = _exec_req or {} + self._task_resource_requests = _task_req or {} + + @property + def id(self) -> int: + """ + Returns + ------- + int + A unique id of this :class:`ResourceProfile` + """ + + if self._java_resource_profile is not None: + return self._java_resource_profile.id() + else: + raise RuntimeError( + "SparkContext must be created to get the id, get the id " + "after adding the ResourceProfile to an RDD" + ) + + @property + def taskResources(self) -> Dict[str, TaskResourceRequest]: + """ + Returns + ------- + dict + a dictionary of resources to :class:`TaskResourceRequest` + """ + + if self._java_resource_profile is not None: + taskRes = self._java_resource_profile.taskResourcesJMap() + result = {} + for k, v in taskRes.items(): + result[k] = TaskResourceRequest(v.resourceName(), v.amount()) + return result + else: + return self._task_resource_requests + + @property + def executorResources(self) -> Dict[str, ExecutorResourceRequest]: + """ + Returns + ------- + dict + a dictionary of resources to :class:`ExecutorResourceRequest` + """ + if self._java_resource_profile is not None: + execRes = self._java_resource_profile.executorResourcesJMap() + result = {} + for k, v in execRes.items(): + result[k] = ExecutorResourceRequest( + v.resourceName(), v.amount(), v.discoveryScript(), v.vendor() + ) + return result + else: + return self._executor_resource_requests + + +class ResourceProfileBuilder: + + """ + Resource profile Builder to build a resource profile to associate with an RDD. + A ResourceProfile allows the user to specify executor and task requirements for + an RDD that will get applied during a stage. This allows the user to change the + resource requirements between stages. + + .. versionadded:: 3.1.0 + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + + Notes + ----- + This API is evolving. + """ + + def __init__(self) -> None: + from pyspark.context import SparkContext + + # TODO: ignore[attr-defined] will be removed, once SparkContext is inlined + _jvm = SparkContext._jvm + if _jvm is not None: + self._jvm = _jvm + self._java_resource_profile_builder = ( + _jvm.org.apache.spark.resource.ResourceProfileBuilder() + ) + else: + self._jvm = None + self._java_resource_profile_builder = None + self._executor_resource_requests: Dict[str, ExecutorResourceRequest] = {} + self._task_resource_requests: Dict[str, TaskResourceRequest] = {} + + def require( + self, resourceRequest: Union[ExecutorResourceRequests, TaskResourceRequests] + ) -> "ResourceProfileBuilder": + """ + Add executor resource requests + + Parameters + ---------- + resourceRequest : :class:`ExecutorResourceRequests` or :class:`TaskResourceRequests` + The detailed executor resource requests, see :class:`ExecutorResourceRequests` + + Returns + ------- + dict + a dictionary of resources to :class:`ExecutorResourceRequest` + """ + + if isinstance(resourceRequest, TaskResourceRequests): + if self._java_resource_profile_builder is not None: + if resourceRequest._java_task_resource_requests is not None: + self._java_resource_profile_builder.require( + resourceRequest._java_task_resource_requests + ) + else: + taskReqs = TaskResourceRequests(self._jvm, resourceRequest.requests) + self._java_resource_profile_builder.require( + taskReqs._java_task_resource_requests + ) + else: + self._task_resource_requests.update(resourceRequest.requests) + else: + if self._java_resource_profile_builder is not None: + r = resourceRequest._java_executor_resource_requests + if r is not None: + self._java_resource_profile_builder.require(r) + else: + execReqs = ExecutorResourceRequests(self._jvm, resourceRequest.requests) + self._java_resource_profile_builder.require( + execReqs._java_executor_resource_requests + ) + else: + self._executor_resource_requests.update(resourceRequest.requests) + return self + + def clearExecutorResourceRequests(self) -> None: + if self._java_resource_profile_builder is not None: + self._java_resource_profile_builder.clearExecutorResourceRequests() + else: + self._executor_resource_requests = {} + + def clearTaskResourceRequests(self) -> None: + if self._java_resource_profile_builder is not None: + self._java_resource_profile_builder.clearTaskResourceRequests() + else: + self._task_resource_requests = {} + + @property + def taskResources(self) -> Dict[str, TaskResourceRequest]: + """ + Returns + ------- + dict + a dictionary of resources to :class:`TaskResourceRequest` + """ + if self._java_resource_profile_builder is not None: + taskRes = self._java_resource_profile_builder.taskResourcesJMap() + result = {} + for k, v in taskRes.items(): + result[k] = TaskResourceRequest(v.resourceName(), v.amount()) + return result + else: + return self._task_resource_requests + + @property + def executorResources(self) -> Dict[str, ExecutorResourceRequest]: + """ + Returns + ------- + dict + a dictionary of resources to :class:`ExecutorResourceRequest` + """ + if self._java_resource_profile_builder is not None: + result = {} + execRes = self._java_resource_profile_builder.executorResourcesJMap() + for k, v in execRes.items(): + result[k] = ExecutorResourceRequest( + v.resourceName(), v.amount(), v.discoveryScript(), v.vendor() + ) + return result + else: + return self._executor_resource_requests + + @property + def build(self) -> ResourceProfile: + if self._java_resource_profile_builder is not None: + jresourceProfile = self._java_resource_profile_builder.build() + return ResourceProfile(_java_resource_profile=jresourceProfile) + else: + return ResourceProfile( + _exec_req=self._executor_resource_requests, _task_req=self._task_resource_requests + ) + + +def _test() -> None: + import doctest + import sys + from pyspark import SparkContext + + globs = globals().copy() + globs["sc"] = SparkContext("local[4]", "profile tests") + (failure_count, test_count) = doctest.testmod( + globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resource/requests.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/requests.py new file mode 100644 index 0000000..043124e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/requests.py @@ -0,0 +1,539 @@ +# +# 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 overload, Optional, Dict + +from py4j.java_gateway import JavaObject, JVMView + +from pyspark.util import _parse_memory + + +class ExecutorResourceRequest: + """ + An Executor resource request. This is used in conjunction with the ResourceProfile to + programmatically specify the resources needed for an RDD that will be applied at the + stage level. + + This is used to specify what the resource requirements are for an Executor and how + Spark can find out specific details about those resources. Not all the parameters are + required for every resource type. Resources like GPUs are supported and have same limitations + as using the global spark configs spark.executor.resource.gpu.*. The amount, discoveryScript, + and vendor parameters for resources are all the same parameters a user would specify through the + configs: spark.executor.resource.{resourceName}.{amount, discoveryScript, vendor}. + + For instance, a user wants to allocate an Executor with GPU resources on YARN. The user has + to specify the resource name (gpu), the amount or number of GPUs per Executor, + the discovery script would be specified so that when the Executor starts up it can + discovery what GPU addresses are available for it to use because YARN doesn't tell + Spark that, then vendor would not be used because its specific for Kubernetes. + + See the configuration and cluster specific docs for more details. + + Use :class:`pyspark.ExecutorResourceRequests` class as a convenience API. + + .. versionadded:: 3.1.0 + + Parameters + ---------- + resourceName : str + Name of the resource + amount : str + Amount requesting + discoveryScript : str, optional + Optional script used to discover the resources. This is required on some + cluster managers that don't tell Spark the addresses of the resources + allocated. The script runs on Executors startup to discover the addresses + of the resources available. + vendor : str, optional + Vendor, required for some cluster managers + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + resourceName: str, + amount: int, + discoveryScript: str = "", + vendor: str = "", + ): + self._name = resourceName + self._amount = amount + self._discovery_script = discoveryScript + self._vendor = vendor + + @property + def resourceName(self) -> str: + """ + Returns + ------- + str + Name of the resource + """ + return self._name + + @property + def amount(self) -> int: + """ + Returns + ------- + str + Amount requesting + """ + return self._amount + + @property + def discoveryScript(self) -> str: + """ + Returns + ------- + str + Amount requesting + """ + return self._discovery_script + + @property + def vendor(self) -> str: + """ + Returns + ------- + str + Vendor, required for some cluster managers + """ + return self._vendor + + +class ExecutorResourceRequests: + + """ + A set of Executor resource requests. This is used in conjunction with the + :class:`pyspark.resource.ResourceProfileBuilder` to programmatically specify the + resources needed for an RDD that will be applied at the stage level. + + .. versionadded:: 3.1.0 + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + + Notes + ----- + This API is evolving. + """ + + _CORES = "cores" + _MEMORY = "memory" + _OVERHEAD_MEM = "memoryOverhead" + _PYSPARK_MEM = "pyspark.memory" + _OFFHEAP_MEM = "offHeap" + + @overload + def __init__(self, _jvm: JVMView): + ... + + @overload + def __init__( + self, + _jvm: None = ..., + _requests: Optional[Dict[str, ExecutorResourceRequest]] = ..., + ): + ... + + def __init__( + self, + _jvm: Optional[JVMView] = None, + _requests: Optional[Dict[str, ExecutorResourceRequest]] = None, + ): + from pyspark import SparkContext + + _jvm = _jvm or SparkContext._jvm + if _jvm is not None: + self._java_executor_resource_requests = ( + _jvm.org.apache.spark.resource.ExecutorResourceRequests() + ) + if _requests is not None: + for k, v in _requests.items(): + if k == self._MEMORY: + self._java_executor_resource_requests.memory(str(v.amount)) + elif k == self._OVERHEAD_MEM: + self._java_executor_resource_requests.memoryOverhead(str(v.amount)) + elif k == self._PYSPARK_MEM: + self._java_executor_resource_requests.pysparkMemory(str(v.amount)) + elif k == self._CORES: + self._java_executor_resource_requests.cores(v.amount) + else: + self._java_executor_resource_requests.resource( + v.resourceName, v.amount, v.discoveryScript, v.vendor + ) + else: + self._java_executor_resource_requests = None + self._executor_resources: Dict[str, ExecutorResourceRequest] = {} + + def memory(self, amount: str) -> "ExecutorResourceRequests": + """ + Specify heap memory. The value specified will be converted to MiB. + This is a convenient API to add :class:`ExecutorResourceRequest` for "memory" resource. + + Parameters + ---------- + amount : str + Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g). + Default unit is MiB if not specified. + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.memory(amount) + else: + self._executor_resources[self._MEMORY] = ExecutorResourceRequest( + self._MEMORY, _parse_memory(amount) + ) + return self + + def memoryOverhead(self, amount: str) -> "ExecutorResourceRequests": + """ + Specify overhead memory. The value specified will be converted to MiB. + This is a convenient API to add :class:`ExecutorResourceRequest` for "memoryOverhead" + resource. + + Parameters + ---------- + amount : str + Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g). + Default unit is MiB if not specified. + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.memoryOverhead(amount) + else: + self._executor_resources[self._OVERHEAD_MEM] = ExecutorResourceRequest( + self._OVERHEAD_MEM, _parse_memory(amount) + ) + return self + + def pysparkMemory(self, amount: str) -> "ExecutorResourceRequests": + """ + Specify pyspark memory. The value specified will be converted to MiB. + This is a convenient API to add :class:`ExecutorResourceRequest` for "pyspark.memory" + resource. + + Parameters + ---------- + amount : str + Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g). + Default unit is MiB if not specified. + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.pysparkMemory(amount) + else: + self._executor_resources[self._PYSPARK_MEM] = ExecutorResourceRequest( + self._PYSPARK_MEM, _parse_memory(amount) + ) + return self + + def offheapMemory(self, amount: str) -> "ExecutorResourceRequests": + """ + Specify off heap memory. The value specified will be converted to MiB. + This value only take effect when MEMORY_OFFHEAP_ENABLED is true. + This is a convenient API to add :class:`ExecutorResourceRequest` for "offHeap" + resource. + + Parameters + ---------- + amount : str + Amount of memory. In the same format as JVM memory strings (e.g. 512m, 2g). + Default unit is MiB if not specified. + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.offHeapMemory(amount) + else: + self._executor_resources[self._OFFHEAP_MEM] = ExecutorResourceRequest( + self._OFFHEAP_MEM, _parse_memory(amount) + ) + return self + + def cores(self, amount: int) -> "ExecutorResourceRequests": + """ + Specify number of cores per Executor. + This is a convenient API to add :class:`ExecutorResourceRequest` for "cores" resource. + + Parameters + ---------- + amount : int + Number of cores to allocate per Executor. + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.cores(amount) + else: + self._executor_resources[self._CORES] = ExecutorResourceRequest(self._CORES, amount) + return self + + def resource( + self, + resourceName: str, + amount: int, + discoveryScript: str = "", + vendor: str = "", + ) -> "ExecutorResourceRequests": + """ + Amount of a particular custom resource(GPU, FPGA, etc) to use. The resource names supported + correspond to the regular Spark configs with the prefix removed. For instance, resources + like GPUs are gpu (spark configs `spark.executor.resource.gpu.*`). If you pass in a resource + that the cluster manager doesn't support the result is undefined, it may error or may just + be ignored. + This is a convenient API to add :class:`ExecutorResourceRequest` for custom resources. + + Parameters + ---------- + resourceName : str + Name of the resource. + amount : str + amount of that resource per executor to use. + discoveryScript : str, optional + Optional script used to discover the resources. This is required on + some cluster managers that don't tell Spark the addresses of + the resources allocated. The script runs on Executors startup to + of the resources available. + vendor : str + Optional vendor, required for some cluster managers + + Returns + ------- + :class:`ExecutorResourceRequests` + """ + if self._java_executor_resource_requests is not None: + self._java_executor_resource_requests.resource( + resourceName, amount, discoveryScript, vendor + ) + else: + self._executor_resources[resourceName] = ExecutorResourceRequest( + resourceName, amount, discoveryScript, vendor + ) + return self + + @property + def requests(self) -> Dict[str, ExecutorResourceRequest]: + """ + Returns + ------- + dict + Returns all the resource requests for the executor. + """ + if self._java_executor_resource_requests is not None: + result = {} + execRes = self._java_executor_resource_requests.requestsJMap() + for k, v in execRes.items(): + result[k] = ExecutorResourceRequest( + v.resourceName(), v.amount(), v.discoveryScript(), v.vendor() + ) + return result + else: + return self._executor_resources + + +class TaskResourceRequest: + """ + A task resource request. This is used in conjunction with the + :class:`pyspark.resource.ResourceProfile` to programmatically specify the resources + needed for an RDD that will be applied at the stage level. The amount is specified + as a float to allow for saying you want more than 1 task per resource. Valid values + are less than or equal to 0.5 or whole numbers. + Use :class:`pyspark.resource.TaskResourceRequests` class as a convenience API. + + Parameters + ---------- + resourceName : str + Name of the resource + amount : float + Amount requesting as a float to support fractional resource requests. + Valid values are less than or equal to 0.5 or whole numbers. This essentially + lets you configure X number of tasks to run on a single resource, + ie amount equals 0.5 translates into 2 tasks per resource address. + + .. versionadded:: 3.1.0 + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + + Notes + ----- + This API is evolving. + """ + + def __init__(self, resourceName: str, amount: float): + self._name = resourceName + self._amount = float(amount) + + @property + def resourceName(self) -> str: + """ + Returns + ------- + str + Name of the resource. + """ + return self._name + + @property + def amount(self) -> float: + """ + Returns + ------- + str + Amount requesting as a float to support fractional resource requests. + """ + return self._amount + + +class TaskResourceRequests: + + """ + A set of task resource requests. This is used in conjunction with the + :class:`pyspark.resource.ResourceProfileBuilder` to programmatically specify the resources + needed for an RDD that will be applied at the stage level. + + .. versionadded:: 3.1.0 + + See Also + -------- + :class:`pyspark.resource.ResourceProfile` + + Notes + ----- + This API is evolving. + """ + + _CPUS = "cpus" + + @overload + def __init__(self, _jvm: JVMView): + ... + + @overload + def __init__( + self, + _jvm: None = ..., + _requests: Optional[Dict[str, TaskResourceRequest]] = ..., + ): + ... + + def __init__( + self, + _jvm: Optional[JVMView] = None, + _requests: Optional[Dict[str, TaskResourceRequest]] = None, + ): + from pyspark import SparkContext + + _jvm = _jvm or SparkContext._jvm + if _jvm is not None: + self._java_task_resource_requests: Optional[ + JavaObject + ] = _jvm.org.apache.spark.resource.TaskResourceRequests() + if _requests is not None: + for k, v in _requests.items(): + if k == self._CPUS: + self._java_task_resource_requests.cpus(int(v.amount)) + else: + self._java_task_resource_requests.resource(v.resourceName, v.amount) + else: + self._java_task_resource_requests = None + self._task_resources: Dict[str, TaskResourceRequest] = {} + + def cpus(self, amount: int) -> "TaskResourceRequests": + """ + Specify number of cpus per Task. + This is a convenient API to add :class:`TaskResourceRequest` for cpus. + + Parameters + ---------- + amount : int + Number of cpus to allocate per Task. + + Returns + ------- + :class:`TaskResourceRequests` + """ + if self._java_task_resource_requests is not None: + self._java_task_resource_requests.cpus(amount) + else: + self._task_resources[self._CPUS] = TaskResourceRequest(self._CPUS, amount) + return self + + def resource(self, resourceName: str, amount: float) -> "TaskResourceRequests": + """ + Amount of a particular custom resource(GPU, FPGA, etc) to use. + This is a convenient API to add :class:`TaskResourceRequest` for custom resources. + + Parameters + ---------- + resourceName : str + Name of the resource. + amount : float + Amount requesting as a float to support fractional resource requests. + Valid values are less than or equal to 0.5 or whole numbers. This essentially + lets you configure X number of tasks to run on a single resource, + ie amount equals 0.5 translates into 2 tasks per resource address. + + Returns + ------- + :class:`TaskResourceRequests` + """ + if self._java_task_resource_requests is not None: + self._java_task_resource_requests.resource(resourceName, float(amount)) + else: + self._task_resources[resourceName] = TaskResourceRequest(resourceName, amount) + return self + + @property + def requests(self) -> Dict[str, TaskResourceRequest]: + """ + Returns + ------- + dict + Returns all the resource requests for the task. + """ + if self._java_task_resource_requests is not None: + result = {} + taskRes = self._java_task_resource_requests.requestsJMap() + for k, v in taskRes.items(): + result[k] = TaskResourceRequest(v.resourceName(), v.amount()) + return result + else: + return self._task_resources diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/test_resources.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/test_resources.py new file mode 100644 index 0000000..81a4ea4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resource/tests/test_resources.py @@ -0,0 +1,83 @@ +# +# 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 unittest + +from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests + + +class ResourceProfileTests(unittest.TestCase): + def test_profile_before_sc(self): + rpb = ResourceProfileBuilder() + ereqs = ExecutorResourceRequests().cores(2).memory("6g").memoryOverhead("1g") + ereqs.pysparkMemory("2g").offheapMemory("3g").resource("gpu", 2, "testGpus", "nvidia.com") + treqs = TaskResourceRequests().cpus(2).resource("gpu", 2) + + def assert_request_contents(exec_reqs, task_reqs): + self.assertEqual(len(exec_reqs), 6) + self.assertEqual(exec_reqs["cores"].amount, 2) + self.assertEqual(exec_reqs["memory"].amount, 6144) + self.assertEqual(exec_reqs["memoryOverhead"].amount, 1024) + self.assertEqual(exec_reqs["pyspark.memory"].amount, 2048) + self.assertEqual(exec_reqs["offHeap"].amount, 3072) + self.assertEqual(exec_reqs["gpu"].amount, 2) + self.assertEqual(exec_reqs["gpu"].discoveryScript, "testGpus") + self.assertEqual(exec_reqs["gpu"].resourceName, "gpu") + self.assertEqual(exec_reqs["gpu"].vendor, "nvidia.com") + self.assertEqual(len(task_reqs), 2) + self.assertEqual(task_reqs["cpus"].amount, 2.0) + self.assertEqual(task_reqs["gpu"].amount, 2.0) + + assert_request_contents(ereqs.requests, treqs.requests) + rp = rpb.require(ereqs).require(treqs).build + assert_request_contents(rp.executorResources, rp.taskResources) + from pyspark import SparkContext, SparkConf + + sc = SparkContext(conf=SparkConf()) + rdd = sc.parallelize(range(10)).withResources(rp) + return_rp = rdd.getResourceProfile() + assert_request_contents(return_rp.executorResources, return_rp.taskResources) + # intermix objects created before SparkContext init and after + rpb2 = ResourceProfileBuilder() + # use reqs created before SparkContext with Builder after + rpb2.require(ereqs) + rpb2.require(treqs) + rp2 = rpb2.build + self.assertTrue(rp2.id > 0) + rdd2 = sc.parallelize(range(10)).withResources(rp2) + return_rp2 = rdd2.getResourceProfile() + assert_request_contents(return_rp2.executorResources, return_rp2.taskResources) + ereqs2 = ExecutorResourceRequests().cores(2).memory("6g").memoryOverhead("1g") + ereqs.pysparkMemory("2g").resource("gpu", 2, "testGpus", "nvidia.com") + treqs2 = TaskResourceRequests().cpus(2).resource("gpu", 2) + # use reqs created after SparkContext with Builder before + rpb.require(ereqs2) + rpb.require(treqs2) + rp3 = rpb.build + assert_request_contents(rp3.executorResources, rp3.taskResources) + sc.stop() + + +if __name__ == "__main__": + from pyspark.resource.tests.test_resources import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/resultiterable.py b/spark-3.5.0-bin-hadoop3/python/pyspark/resultiterable.py new file mode 100644 index 0000000..7f35cb6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/resultiterable.py @@ -0,0 +1,45 @@ +# +# 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 TypeVar, TYPE_CHECKING, Iterator, Iterable + +if TYPE_CHECKING: + from pyspark._typing import SizedIterable + + +__all__ = ["ResultIterable"] + +T = TypeVar("T") + + +class ResultIterable(Iterable[T]): + + """ + A special result iterable. This is used because the standard + iterator can not be pickled + """ + + def __init__(self, data: "SizedIterable[T]"): + self.data: "SizedIterable[T]" = data + self.index: int = 0 + self.maxindex: int = len(data) + + def __iter__(self) -> Iterator[T]: + return iter(self.data) + + def __len__(self) -> int: + return len(self.data) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/serializers.py b/spark-3.5.0-bin-hadoop3/python/pyspark/serializers.py new file mode 100644 index 0000000..80b6114 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/serializers.py @@ -0,0 +1,681 @@ +# +# 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. +# + +""" +PySpark supports custom serializers for transferring data; this can improve +performance. + +By default, PySpark uses :class:`CloudPickleSerializer` to serialize objects using Python's +`cPickle` serializer, which can serialize nearly any Python object. +Other serializers, like :class:`MarshalSerializer`, support fewer datatypes but can be +faster. + + +Examples +-------- +The serializer is chosen when creating :class:`SparkContext`: + +>>> from pyspark.context import SparkContext +>>> from pyspark.serializers import MarshalSerializer +>>> sc = SparkContext('local', 'test', serializer=MarshalSerializer()) +>>> sc.parallelize(list(range(1000))).map(lambda x: 2 * x).take(10) +[0, 2, 4, 6, 8, 10, 12, 14, 16, 18] +>>> sc.stop() + +PySpark serializes objects in batches; by default, the batch size is chosen based +on the size of objects and is also configurable by SparkContext's `batchSize` +parameter: + +>>> sc = SparkContext('local', 'test', batchSize=2) +>>> rdd = sc.parallelize(range(16), 4).map(lambda x: x) + +Behind the scenes, this creates a JavaRDD with four partitions, each of +which contains two batches of two objects: + +>>> rdd.glom().collect() +[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]] +>>> int(rdd._jrdd.count()) +8 +>>> sc.stop() +""" + +import sys +import os +from itertools import chain, product +import marshal +import struct +import types +import collections +import zlib +import itertools +import pickle + +pickle_protocol = pickle.HIGHEST_PROTOCOL + +from pyspark import cloudpickle +from pyspark.util import print_exec + + +__all__ = [ + "PickleSerializer", + "CPickleSerializer", + "CloudPickleSerializer", + "MarshalSerializer", + "UTF8Deserializer", +] + + +class SpecialLengths: + END_OF_DATA_SECTION = -1 + PYTHON_EXCEPTION_THROWN = -2 + TIMING_DATA = -3 + END_OF_STREAM = -4 + NULL = -5 + START_ARROW_STREAM = -6 + + +class Serializer: + def dump_stream(self, iterator, stream): + """ + Serialize an iterator of objects to the output stream. + """ + raise NotImplementedError + + def load_stream(self, stream): + """ + Return an iterator of deserialized objects from the input stream. + """ + raise NotImplementedError + + def dumps(self, obj): + """ + Serialize an object into a byte array. + When batching is used, this will be called with an array of objects. + """ + raise NotImplementedError + + def _load_stream_without_unbatching(self, stream): + """ + Return an iterator of deserialized batches (iterable) of objects from the input stream. + If the serializer does not operate on batches the default implementation returns an + iterator of single element lists. + """ + return map(lambda x: [x], self.load_stream(stream)) + + # Note: our notion of "equality" is that output generated by + # equal serializers can be deserialized using the same serializer. + + # This default implementation handles the simple cases; + # subclasses should override __eq__ as appropriate. + + def __eq__(self, other): + return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ + + def __ne__(self, other): + return not self.__eq__(other) + + def __repr__(self): + return "%s()" % self.__class__.__name__ + + def __hash__(self): + return hash(str(self)) + + +class FramedSerializer(Serializer): + + """ + Serializer that writes objects as a stream of (length, data) pairs, + where `length` is a 32-bit integer and data is `length` bytes. + """ + + def dump_stream(self, iterator, stream): + for obj in iterator: + self._write_with_length(obj, stream) + + def load_stream(self, stream): + while True: + try: + yield self._read_with_length(stream) + except EOFError: + return + + def _write_with_length(self, obj, stream): + serialized = self.dumps(obj) + if serialized is None: + raise ValueError("serialized value should not be None") + if len(serialized) > (1 << 31): + raise ValueError("can not serialize object larger than 2G") + write_int(len(serialized), stream) + stream.write(serialized) + + def _read_with_length(self, stream): + length = read_int(stream) + if length == SpecialLengths.END_OF_DATA_SECTION: + raise EOFError + elif length == SpecialLengths.NULL: + return None + obj = stream.read(length) + if len(obj) < length: + raise EOFError + return self.loads(obj) + + def dumps(self, obj): + """ + Serialize an object into a byte array. + When batching is used, this will be called with an array of objects. + """ + raise NotImplementedError + + def loads(self, obj): + """ + Deserialize an object from a byte array. + """ + raise NotImplementedError + + +class BatchedSerializer(Serializer): + + """ + Serializes a stream of objects in batches by calling its wrapped + Serializer with streams of objects. + """ + + UNLIMITED_BATCH_SIZE = -1 + UNKNOWN_BATCH_SIZE = 0 + + def __init__(self, serializer, batchSize=UNLIMITED_BATCH_SIZE): + self.serializer = serializer + self.batchSize = batchSize + + def _batched(self, iterator): + if self.batchSize == self.UNLIMITED_BATCH_SIZE: + yield list(iterator) + elif hasattr(iterator, "__len__") and hasattr(iterator, "__getslice__"): + n = len(iterator) + for i in range(0, n, self.batchSize): + yield iterator[i : i + self.batchSize] + else: + items = [] + count = 0 + for item in iterator: + items.append(item) + count += 1 + if count == self.batchSize: + yield items + items = [] + count = 0 + if items: + yield items + + def dump_stream(self, iterator, stream): + self.serializer.dump_stream(self._batched(iterator), stream) + + def load_stream(self, stream): + return chain.from_iterable(self._load_stream_without_unbatching(stream)) + + def _load_stream_without_unbatching(self, stream): + return self.serializer.load_stream(stream) + + def __repr__(self): + return "BatchedSerializer(%s, %d)" % (str(self.serializer), self.batchSize) + + +class FlattenedValuesSerializer(BatchedSerializer): + + """ + Serializes a stream of list of pairs, split the list of values + which contain more than a certain number of objects to make them + have similar sizes. + """ + + def __init__(self, serializer, batchSize=10): + BatchedSerializer.__init__(self, serializer, batchSize) + + def _batched(self, iterator): + n = self.batchSize + for key, values in iterator: + for i in range(0, len(values), n): + yield key, values[i : i + n] + + def load_stream(self, stream): + return self.serializer.load_stream(stream) + + def __repr__(self): + return "FlattenedValuesSerializer(%s, %d)" % (self.serializer, self.batchSize) + + +class AutoBatchedSerializer(BatchedSerializer): + """ + Choose the size of batch automatically based on the size of object + """ + + def __init__(self, serializer, bestSize=1 << 16): + BatchedSerializer.__init__(self, serializer, self.UNKNOWN_BATCH_SIZE) + self.bestSize = bestSize + + def dump_stream(self, iterator, stream): + batch, best = 1, self.bestSize + iterator = iter(iterator) + while True: + vs = list(itertools.islice(iterator, batch)) + if not vs: + break + + bytes = self.serializer.dumps(vs) + write_int(len(bytes), stream) + stream.write(bytes) + + size = len(bytes) + if size < best: + batch *= 2 + elif size > best * 10 and batch > 1: + batch //= 2 + + def __repr__(self): + return "AutoBatchedSerializer(%s)" % self.serializer + + +class CartesianDeserializer(Serializer): + + """ + Deserializes the JavaRDD cartesian() of two PythonRDDs. + Due to pyspark batching we cannot simply use the result of the Java RDD cartesian, + we additionally need to do the cartesian within each pair of batches. + """ + + def __init__(self, key_ser, val_ser): + self.key_ser = key_ser + self.val_ser = val_ser + + def _load_stream_without_unbatching(self, stream): + key_batch_stream = self.key_ser._load_stream_without_unbatching(stream) + val_batch_stream = self.val_ser._load_stream_without_unbatching(stream) + for (key_batch, val_batch) in zip(key_batch_stream, val_batch_stream): + # for correctness with repeated cartesian/zip this must be returned as one batch + yield product(key_batch, val_batch) + + def load_stream(self, stream): + return chain.from_iterable(self._load_stream_without_unbatching(stream)) + + def __repr__(self): + return "CartesianDeserializer(%s, %s)" % (str(self.key_ser), str(self.val_ser)) + + +class PairDeserializer(Serializer): + + """ + Deserializes the JavaRDD zip() of two PythonRDDs. + Due to pyspark batching we cannot simply use the result of the Java RDD zip, + we additionally need to do the zip within each pair of batches. + """ + + def __init__(self, key_ser, val_ser): + self.key_ser = key_ser + self.val_ser = val_ser + + def _load_stream_without_unbatching(self, stream): + key_batch_stream = self.key_ser._load_stream_without_unbatching(stream) + val_batch_stream = self.val_ser._load_stream_without_unbatching(stream) + for (key_batch, val_batch) in zip(key_batch_stream, val_batch_stream): + # For double-zipped RDDs, the batches can be iterators from other PairDeserializer, + # instead of lists. We need to convert them to lists if needed. + key_batch = key_batch if hasattr(key_batch, "__len__") else list(key_batch) + val_batch = val_batch if hasattr(val_batch, "__len__") else list(val_batch) + if len(key_batch) != len(val_batch): + raise ValueError( + "Can not deserialize PairRDD with different number of items" + " in batches: (%d, %d)" % (len(key_batch), len(val_batch)) + ) + # for correctness with repeated cartesian/zip this must be returned as one batch + yield zip(key_batch, val_batch) + + def load_stream(self, stream): + return chain.from_iterable(self._load_stream_without_unbatching(stream)) + + def __repr__(self): + return "PairDeserializer(%s, %s)" % (str(self.key_ser), str(self.val_ser)) + + +class NoOpSerializer(FramedSerializer): + def loads(self, obj): + return obj + + def dumps(self, obj): + return obj + + +if os.environ.get("PYSPARK_ENABLE_NAMEDTUPLE_PATCH") == "1": + # Hack namedtuple, make it picklable. + # For Python 3.8+, we use CPickle-based cloudpickle. + # SPARK-41189: There are still behaviour differences between regular pickle + # and Cloudpickle e.g., bug fixes from the upstream. It's safer to have + # a switch to turn on and off for the time being. + + __cls = {} # type: ignore[var-annotated] + + def _restore(name, fields, value): + """Restore an object of namedtuple""" + k = (name, fields) + cls = __cls.get(k) + if cls is None: + cls = collections.namedtuple(name, fields) + __cls[k] = cls + return cls(*value) + + def _hack_namedtuple(cls): + """Make class generated by namedtuple picklable""" + name = cls.__name__ + fields = cls._fields + + def __reduce__(self): + return (_restore, (name, fields, tuple(self))) + + cls.__reduce__ = __reduce__ + cls._is_namedtuple_ = True + return cls + + def _hijack_namedtuple(): + """Hack namedtuple() to make it picklable""" + # hijack only one time + if hasattr(collections.namedtuple, "__hijack"): + return + + global _old_namedtuple # or it will put in closure + global _old_namedtuple_kwdefaults # or it will put in closure too + + def _copy_func(f): + return types.FunctionType( + f.__code__, f.__globals__, f.__name__, f.__defaults__, f.__closure__ + ) + + _old_namedtuple = _copy_func(collections.namedtuple) + _old_namedtuple_kwdefaults = collections.namedtuple.__kwdefaults__ + + def namedtuple(*args, **kwargs): + for k, v in _old_namedtuple_kwdefaults.items(): + kwargs[k] = kwargs.get(k, v) + cls = _old_namedtuple(*args, **kwargs) + return _hack_namedtuple(cls) + + # replace namedtuple with the new one + collections.namedtuple.__globals__[ + "_old_namedtuple_kwdefaults" + ] = _old_namedtuple_kwdefaults + collections.namedtuple.__globals__["_old_namedtuple"] = _old_namedtuple + collections.namedtuple.__globals__["_hack_namedtuple"] = _hack_namedtuple + collections.namedtuple.__code__ = namedtuple.__code__ + collections.namedtuple.__hijack = 1 + + # hack the cls already generated by namedtuple. + # Those created in other modules can be pickled as normal, + # so only hack those in __main__ module + for n, o in sys.modules["__main__"].__dict__.items(): + if ( + type(o) is type + and o.__base__ is tuple + and hasattr(o, "_fields") + and "__reduce__" not in o.__dict__ + ): + _hack_namedtuple(o) # hack inplace + + _hijack_namedtuple() + + +class PickleSerializer(FramedSerializer): + + """ + Serializes objects using Python's pickle serializer: + + http://docs.python.org/2/library/pickle.html + + This serializer supports nearly any Python object, but may + not be as fast as more specialized serializers. + """ + + def dumps(self, obj): + return pickle.dumps(obj, pickle_protocol) + + def loads(self, obj, encoding="bytes"): + return pickle.loads(obj, encoding=encoding) + + +class CloudPickleSerializer(FramedSerializer): + def dumps(self, obj): + try: + return cloudpickle.dumps(obj, pickle_protocol) + except pickle.PickleError: + raise + except Exception as e: + emsg = str(e) + if "'i' format requires" in emsg: + msg = "Object too large to serialize: %s" % emsg + else: + msg = "Could not serialize object: %s: %s" % (e.__class__.__name__, emsg) + print_exec(sys.stderr) + raise pickle.PicklingError(msg) + + def loads(self, obj, encoding="bytes"): + return cloudpickle.loads(obj, encoding=encoding) + + +if os.environ.get("PYSPARK_ENABLE_NAMEDTUPLE_PATCH") == "1": + CPickleSerializer = PickleSerializer +else: + CPickleSerializer = CloudPickleSerializer # type: ignore[misc, assignment] + + +class MarshalSerializer(FramedSerializer): + + """ + Serializes objects using Python's Marshal serializer: + + http://docs.python.org/2/library/marshal.html + + This serializer is faster than CloudPickleSerializer but supports fewer datatypes. + """ + + def dumps(self, obj): + return marshal.dumps(obj) + + def loads(self, obj): + return marshal.loads(obj) + + +class AutoSerializer(FramedSerializer): + + """ + Choose marshal or pickle as serialization protocol automatically + """ + + def __init__(self): + FramedSerializer.__init__(self) + self._type = None + + def dumps(self, obj): + if self._type is not None: + return b"P" + pickle.dumps(obj, -1) + try: + return b"M" + marshal.dumps(obj) + except Exception: + self._type = b"P" + return b"P" + pickle.dumps(obj, -1) + + def loads(self, obj): + _type = obj[0] + if _type == b"M": + return marshal.loads(obj[1:]) + elif _type == b"P": + return pickle.loads(obj[1:]) + else: + raise ValueError("invalid serialization type: %s" % _type) + + +class CompressedSerializer(FramedSerializer): + """ + Compress the serialized data + """ + + def __init__(self, serializer): + FramedSerializer.__init__(self) + assert isinstance(serializer, FramedSerializer), "serializer must be a FramedSerializer" + self.serializer = serializer + + def dumps(self, obj): + return zlib.compress(self.serializer.dumps(obj), 1) + + def loads(self, obj): + return self.serializer.loads(zlib.decompress(obj)) + + def __repr__(self): + return "CompressedSerializer(%s)" % self.serializer + + +class UTF8Deserializer(Serializer): + + """ + Deserializes streams written by String.getBytes. + """ + + def __init__(self, use_unicode=True): + self.use_unicode = use_unicode + + def loads(self, stream): + length = read_int(stream) + if length == SpecialLengths.END_OF_DATA_SECTION: + raise EOFError + elif length == SpecialLengths.NULL: + return None + s = stream.read(length) + return s.decode("utf-8") if self.use_unicode else s + + def load_stream(self, stream): + try: + while True: + yield self.loads(stream) + except struct.error: + return + except EOFError: + return + + def __repr__(self): + return "UTF8Deserializer(%s)" % self.use_unicode + + +def read_long(stream): + length = stream.read(8) + if not length: + raise EOFError + return struct.unpack("!q", length)[0] + + +def write_long(value, stream): + stream.write(struct.pack("!q", value)) + + +def pack_long(value): + return struct.pack("!q", value) + + +def read_int(stream): + length = stream.read(4) + if not length: + raise EOFError + return struct.unpack("!i", length)[0] + + +def write_int(value, stream): + stream.write(struct.pack("!i", value)) + + +def read_bool(stream): + length = stream.read(1) + if not length: + raise EOFError + return struct.unpack("!?", length)[0] + + +def write_with_length(obj, stream): + write_int(len(obj), stream) + stream.write(obj) + + +class ChunkedStream: + + """ + This is a file-like object takes a stream of data, of unknown length, and breaks it into fixed + length frames. The intended use case is serializing large data and sending it immediately over + a socket -- we do not want to buffer the entire data before sending it, but the receiving end + needs to know whether or not there is more data coming. + + It works by buffering the incoming data in some fixed-size chunks. If the buffer is full, it + first sends the buffer size, then the data. This repeats as long as there is more data to send. + When this is closed, it sends the length of whatever data is in the buffer, then that data, and + finally a "length" of -1 to indicate the stream has completed. + """ + + def __init__(self, wrapped, buffer_size): + self.buffer_size = buffer_size + self.buffer = bytearray(buffer_size) + self.current_pos = 0 + self.wrapped = wrapped + + def write(self, bytes): + byte_pos = 0 + byte_remaining = len(bytes) + while byte_remaining > 0: + new_pos = byte_remaining + self.current_pos + if new_pos < self.buffer_size: + # just put it in our buffer + self.buffer[self.current_pos : new_pos] = bytes[byte_pos:] + self.current_pos = new_pos + byte_remaining = 0 + else: + # fill the buffer, send the length then the contents, and start filling again + space_left = self.buffer_size - self.current_pos + new_byte_pos = byte_pos + space_left + self.buffer[self.current_pos : self.buffer_size] = bytes[byte_pos:new_byte_pos] + write_int(self.buffer_size, self.wrapped) + self.wrapped.write(self.buffer) + byte_remaining -= space_left + byte_pos = new_byte_pos + self.current_pos = 0 + + def close(self): + # if there is anything left in the buffer, write it out first + if self.current_pos > 0: + write_int(self.current_pos, self.wrapped) + self.wrapped.write(self.buffer[: self.current_pos]) + # -1 length indicates to the receiving end that we're done. + write_int(-1, self.wrapped) + self.wrapped.close() + + @property + def closed(self): + """ + Return True if the `wrapped` object has been closed. + NOTE: this property is required by pyarrow to be used as a file-like object in + pyarrow.RecordBatchStreamWriter from ArrowStreamSerializer + """ + return self.wrapped.closed + + +if __name__ == "__main__": + import doctest + + (failure_count, test_count) = doctest.testmod() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/shell.py b/spark-3.5.0-bin-hadoop3/python/pyspark/shell.py new file mode 100644 index 0000000..7e2093c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/shell.py @@ -0,0 +1,123 @@ +# +# 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. +# + +""" +An interactive shell. + +This file is designed to be launched as a PYTHONSTARTUP script. +""" + +import atexit +import builtins +import os +import platform +import warnings +import sys + +import pyspark +from pyspark.context import SparkContext +from pyspark.sql import SparkSession +from pyspark.sql.context import SQLContext +from pyspark.sql.utils import is_remote +from urllib.parse import urlparse + +if getattr(builtins, "__IPYTHON__", False): + # (Only) during PYTHONSTARTUP execution, IPython temporarily adds the parent + # directory of the script into the Python path, which results in searching + # packages under `pyspark` directory. + # For example, `import pandas` attempts to import `pyspark.pandas`, see also SPARK-42266. + if "__file__" in globals(): + parent_dir = os.path.abspath(os.path.dirname(__file__)) + if parent_dir in sys.path: + sys.path.remove(parent_dir) + + +if is_remote(): + try: + # Creates pyspark.sql.connect.SparkSession. + spark = SparkSession.builder.getOrCreate() + except Exception: + import sys + import traceback + + warnings.warn("Failed to initialize Spark session.") + traceback.print_exc(file=sys.stderr) + sys.exit(1) + version = pyspark.__version__ + sc = None +else: + if os.environ.get("SPARK_EXECUTOR_URI"): + SparkContext.setSystemProperty("spark.executor.uri", os.environ["SPARK_EXECUTOR_URI"]) + + SparkContext._ensure_initialized() + + try: + spark = SparkSession._create_shell_session() + except Exception: + import sys + import traceback + + warnings.warn("Failed to initialize Spark session.") + traceback.print_exc(file=sys.stderr) + sys.exit(1) + + sc = spark.sparkContext + atexit.register((lambda sc: lambda: sc.stop())(sc)) + + # for compatibility + sqlContext = SQLContext._get_or_create(sc) + sqlCtx = sqlContext + version = sc.version + +sql = spark.sql + +print( + r"""Welcome to + ____ __ + / __/__ ___ _____/ /__ + _\ \/ _ \/ _ `/ __/ '_/ + /__ / .__/\_,_/_/ /_/\_\ version %s + /_/ +""" + % version +) +print( + "Using Python version %s (%s, %s)" + % (platform.python_version(), platform.python_build()[0], platform.python_build()[1]) +) +if is_remote(): + url = os.environ.get("SPARK_REMOTE", None) + assert url is not None + if url.startswith("local"): + url = "sc://localhost" # only for display in the console. + print("Client connected to the Spark Connect server at %s" % urlparse(url).netloc) +else: + print("Spark context Web UI available at %s" % (sc.uiWebUrl)) # type: ignore[union-attr] + print( + "Spark context available as 'sc' (master = %s, app id = %s)." + % (sc.master, sc.applicationId) # type: ignore[union-attr] + ) + +print("SparkSession available as 'spark'.") + +# The ./bin/pyspark script stores the old PYTHONSTARTUP value in OLD_PYTHONSTARTUP, +# which allows us to execute the user's PYTHONSTARTUP file: +_pythonstartup = os.environ.get("OLD_PYTHONSTARTUP") +if _pythonstartup and os.path.isfile(_pythonstartup): + with open(_pythonstartup) as f: + code = compile(f.read(), _pythonstartup, "exec") + exec(code) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/shuffle.py b/spark-3.5.0-bin-hadoop3/python/pyspark/shuffle.py new file mode 100644 index 0000000..da03110 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/shuffle.py @@ -0,0 +1,854 @@ +# +# 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 os +import platform +import shutil +import warnings +import gc +import itertools +import operator +import random +import sys + +import heapq +from pyspark.serializers import ( + BatchedSerializer, + CPickleSerializer, + FlattenedValuesSerializer, + CompressedSerializer, + AutoBatchedSerializer, +) +from pyspark.util import fail_on_stopiteration + + +try: + import psutil + + process = None + + def get_used_memory(): + """Return the used memory in MiB""" + global process + if process is None or process._pid != os.getpid(): + process = psutil.Process(os.getpid()) + if hasattr(process, "memory_info"): + info = process.memory_info() + else: + info = process.get_memory_info() + return info.rss >> 20 + +except ImportError: + + def get_used_memory(): + """Return the used memory in MiB""" + if platform.system() == "Linux": + for line in open("/proc/self/status"): + if line.startswith("VmRSS:"): + return int(line.split()[1]) >> 10 + + else: + warnings.warn("Please install psutil to have better " "support with spilling") + if platform.system() == "Darwin": + import resource + + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return rss >> 20 + # TODO: support windows + + return 0 + + +def _get_local_dirs(sub): + """Get all the directories""" + path = os.environ.get("SPARK_LOCAL_DIRS", "/tmp") + dirs = path.split(",") + if len(dirs) > 1: + if sys.version_info < (3, 11): + # different order in different processes and instances + rnd = random.Random(os.getpid() + id(dirs)) + random.shuffle(dirs, rnd.random) + else: + random.shuffle(dirs) + return [os.path.join(d, "python", str(os.getpid()), sub) for d in dirs] + + +# global stats +MemoryBytesSpilled = 0 +DiskBytesSpilled = 0 + + +class Aggregator: + + """ + Aggregator has tree functions to merge values into combiner. + + createCombiner: (value) -> combiner + mergeValue: (combine, value) -> combiner + mergeCombiners: (combiner, combiner) -> combiner + """ + + def __init__(self, createCombiner, mergeValue, mergeCombiners): + self.createCombiner = fail_on_stopiteration(createCombiner) + self.mergeValue = fail_on_stopiteration(mergeValue) + self.mergeCombiners = fail_on_stopiteration(mergeCombiners) + + +class SimpleAggregator(Aggregator): + + """ + SimpleAggregator is useful for the cases that combiners have + same type with values + """ + + def __init__(self, combiner): + Aggregator.__init__(self, lambda x: x, combiner, combiner) + + +class Merger: + + """ + Merge shuffled data together by aggregator + """ + + def __init__(self, aggregator): + self.agg = aggregator + + def mergeValues(self, iterator): + """Combine the items by creator and combiner""" + raise NotImplementedError + + def mergeCombiners(self, iterator): + """Merge the combined items by mergeCombiner""" + raise NotImplementedError + + def items(self): + """Return the merged items ad iterator""" + raise NotImplementedError + + +def _compressed_serializer(self, serializer=None): + # always use CPickleSerializer to simplify implementation + ser = CPickleSerializer() + return AutoBatchedSerializer(CompressedSerializer(ser)) + + +class ExternalMerger(Merger): + + """ + External merger will dump the aggregated data into disks when + memory usage goes above the limit, then merge them together. + + This class works as follows: + + - It repeatedly combine the items and save them in one dict in + memory. + + - When the used memory goes above memory limit, it will split + the combined data into partitions by hash code, dump them + into disk, one file per partition. + + - Then it goes through the rest of the iterator, combine items + into different dict by hash. Until the used memory goes over + memory limit, it dump all the dicts into disks, one file per + dict. Repeat this again until combine all the items. + + - Before return any items, it will load each partition and + combine them separately. Yield them before loading next + partition. + + - During loading a partition, if the memory goes over limit, + it will partition the loaded data and dump them into disks + and load them partition by partition again. + + `data` and `pdata` are used to hold the merged items in memory. + At first, all the data are merged into `data`. Once the used + memory goes over limit, the items in `data` are dumped into + disks, `data` will be cleared, all rest of items will be merged + into `pdata` and then dumped into disks. Before returning, all + the items in `pdata` will be dumped into disks. + + Finally, if any items were spilled into disks, each partition + will be merged into `data` and be yielded, then cleared. + + Examples + -------- + >>> agg = SimpleAggregator(lambda x, y: x + y) + >>> merger = ExternalMerger(agg, 10) + >>> N = 10000 + >>> merger.mergeValues(zip(range(N), range(N))) + >>> assert merger.spills > 0 + >>> sum(v for k,v in merger.items()) + 49995000 + + >>> merger = ExternalMerger(agg, 10) + >>> merger.mergeCombiners(zip(range(N), range(N))) + >>> assert merger.spills > 0 + >>> sum(v for k,v in merger.items()) + 49995000 + """ + + # the max total partitions created recursively + MAX_TOTAL_PARTITIONS = 4096 + + def __init__( + self, + aggregator, + memory_limit=512, + serializer=None, + localdirs=None, + scale=1, + partitions=59, + batch=1000, + ): + Merger.__init__(self, aggregator) + self.memory_limit = memory_limit + self.serializer = _compressed_serializer(serializer) + self.localdirs = localdirs or _get_local_dirs(str(id(self))) + # number of partitions when spill data into disks + self.partitions = partitions + # check the memory after # of items merged + self.batch = batch + # scale is used to scale down the hash of key for recursive hash map + self.scale = scale + # un-partitioned merged data + self.data = {} + # partitioned merged data, list of dicts + self.pdata = [] + # number of chunks dumped into disks + self.spills = 0 + # randomize the hash of key, id(o) is the address of o (aligned by 8) + self._seed = id(self) + 7 + + def _get_spill_dir(self, n): + """Choose one directory for spill by number n""" + return os.path.join(self.localdirs[n % len(self.localdirs)], str(n)) + + def _next_limit(self): + """ + Return the next memory limit. If the memory is not released + after spilling, it will dump the data only when the used memory + starts to increase. + """ + return max(self.memory_limit, get_used_memory() * 1.05) + + def mergeValues(self, iterator): + """Combine the items by creator and combiner""" + # speedup attribute lookup + creator, comb = self.agg.createCombiner, self.agg.mergeValue + c, data, pdata, hfun, batch = 0, self.data, self.pdata, self._partition, self.batch + limit = self.memory_limit + + for k, v in iterator: + d = pdata[hfun(k)] if pdata else data + d[k] = comb(d[k], v) if k in d else creator(v) + + c += 1 + if c >= batch: + if get_used_memory() >= limit: + self._spill() + limit = self._next_limit() + batch /= 2 + c = 0 + else: + batch *= 1.5 + + if get_used_memory() >= limit: + self._spill() + + def _partition(self, key): + """Return the partition for key""" + return hash((key, self._seed)) % self.partitions + + def _object_size(self, obj): + """How much of memory for this obj, assume that all the objects + consume similar bytes of memory + """ + return 1 + + def mergeCombiners(self, iterator, limit=None): + """Merge (K,V) pair by mergeCombiner""" + if limit is None: + limit = self.memory_limit + # speedup attribute lookup + comb, hfun, objsize = self.agg.mergeCombiners, self._partition, self._object_size + c, data, pdata, batch = 0, self.data, self.pdata, self.batch + for k, v in iterator: + d = pdata[hfun(k)] if pdata else data + d[k] = comb(d[k], v) if k in d else v + if not limit: + continue + + c += objsize(v) + if c > batch: + if get_used_memory() > limit: + self._spill() + limit = self._next_limit() + batch /= 2 + c = 0 + else: + batch *= 1.5 + + if limit and get_used_memory() >= limit: + self._spill() + + def _spill(self): + """ + dump already partitioned data into disks. + + It will dump the data in batch for better performance. + """ + global MemoryBytesSpilled, DiskBytesSpilled + path = self._get_spill_dir(self.spills) + if not os.path.exists(path): + os.makedirs(path) + + used_memory = get_used_memory() + if not self.pdata: + # The data has not been partitioned, it will iterator the + # dataset once, write them into different files, has no + # additional memory. It only called when the memory goes + # above limit at the first time. + + # open all the files for writing + streams = [open(os.path.join(path, str(i)), "wb") for i in range(self.partitions)] + + for k, v in self.data.items(): + h = self._partition(k) + # put one item in batch, make it compatible with load_stream + # it will increase the memory if dump them in batch + self.serializer.dump_stream([(k, v)], streams[h]) + + for s in streams: + DiskBytesSpilled += s.tell() + s.close() + + self.data.clear() + self.pdata.extend([{} for i in range(self.partitions)]) + + else: + for i in range(self.partitions): + p = os.path.join(path, str(i)) + with open(p, "wb") as f: + # dump items in batch + self.serializer.dump_stream(iter(self.pdata[i].items()), f) + self.pdata[i].clear() + DiskBytesSpilled += os.path.getsize(p) + + self.spills += 1 + gc.collect() # release the memory as much as possible + MemoryBytesSpilled += max(used_memory - get_used_memory(), 0) << 20 + + def items(self): + """Return all merged items as iterator""" + if not self.pdata and not self.spills: + return iter(self.data.items()) + return self._external_items() + + def _external_items(self): + """Return all partitioned items as iterator""" + assert not self.data + if any(self.pdata): + self._spill() + # disable partitioning and spilling when merge combiners from disk + self.pdata = [] + + try: + for i in range(self.partitions): + for v in self._merged_items(i): + yield v + self.data.clear() + + # remove the merged partition + for j in range(self.spills): + path = self._get_spill_dir(j) + os.remove(os.path.join(path, str(i))) + finally: + self._cleanup() + + def _merged_items(self, index): + self.data = {} + limit = self._next_limit() + for j in range(self.spills): + path = self._get_spill_dir(j) + p = os.path.join(path, str(index)) + # do not check memory during merging + with open(p, "rb") as f: + self.mergeCombiners(self.serializer.load_stream(f), 0) + + # limit the total partitions + if ( + self.scale * self.partitions < self.MAX_TOTAL_PARTITIONS + and j < self.spills - 1 + and get_used_memory() > limit + ): + self.data.clear() # will read from disk again + gc.collect() # release the memory as much as possible + return self._recursive_merged_items(index) + + return self.data.items() + + def _recursive_merged_items(self, index): + """ + merge the partitioned items and return the as iterator + + If one partition can not be fit in memory, then them will be + partitioned and merged recursively. + """ + subdirs = [os.path.join(d, "parts", str(index)) for d in self.localdirs] + m = ExternalMerger( + self.agg, + self.memory_limit, + self.serializer, + subdirs, + self.scale * self.partitions, + self.partitions, + self.batch, + ) + m.pdata = [{} for _ in range(self.partitions)] + limit = self._next_limit() + + for j in range(self.spills): + path = self._get_spill_dir(j) + p = os.path.join(path, str(index)) + with open(p, "rb") as f: + m.mergeCombiners(self.serializer.load_stream(f), 0) + + if get_used_memory() > limit: + m._spill() + limit = self._next_limit() + + return m._external_items() + + def _cleanup(self): + """Clean up all the files in disks""" + for d in self.localdirs: + shutil.rmtree(d, True) + + +class ExternalSorter: + """ + ExternalSorter will divide the elements into chunks, sort them in + memory and dump them into disks, finally merge them back. + + The spilling will only happen when the used memory goes above + the limit. + + Examples + -------- + >>> sorter = ExternalSorter(1) # 1M + >>> import random + >>> l = list(range(1024)) + >>> random.shuffle(l) + >>> sorted(l) == list(sorter.sorted(l)) + True + >>> sorted(l) == list(sorter.sorted(l, key=lambda x: -x, reverse=True)) + True + """ + + def __init__(self, memory_limit, serializer=None): + self.memory_limit = memory_limit + self.local_dirs = _get_local_dirs("sort") + self.serializer = _compressed_serializer(serializer) + + def _get_path(self, n): + """Choose one directory for spill by number n""" + d = self.local_dirs[n % len(self.local_dirs)] + if not os.path.exists(d): + os.makedirs(d) + return os.path.join(d, str(n)) + + def _next_limit(self): + """ + Return the next memory limit. If the memory is not released + after spilling, it will dump the data only when the used memory + starts to increase. + """ + return max(self.memory_limit, get_used_memory() * 1.05) + + def sorted(self, iterator, key=None, reverse=False): + """ + Sort the elements in iterator, do external sort when the memory + goes above the limit. + """ + global MemoryBytesSpilled, DiskBytesSpilled + batch, limit = 100, self._next_limit() + chunks, current_chunk = [], [] + iterator = iter(iterator) + while True: + # pick elements in batch + chunk = list(itertools.islice(iterator, batch)) + current_chunk.extend(chunk) + if len(chunk) < batch: + break + + used_memory = get_used_memory() + if used_memory > limit: + # sort them inplace will save memory + current_chunk.sort(key=key, reverse=reverse) + path = self._get_path(len(chunks)) + with open(path, "wb") as f: + self.serializer.dump_stream(current_chunk, f) + + def load(f): + for v in self.serializer.load_stream(f): + yield v + # close the file explicit once we consume all the items + # to avoid ResourceWarning in Python3 + f.close() + + chunks.append(load(open(path, "rb"))) + current_chunk = [] + MemoryBytesSpilled += max(used_memory - get_used_memory(), 0) << 20 + DiskBytesSpilled += os.path.getsize(path) + os.unlink(path) # data will be deleted after close + + elif not chunks: + batch = min(int(batch * 1.5), 10000) + + current_chunk.sort(key=key, reverse=reverse) + if not chunks: + return current_chunk + + if current_chunk: + chunks.append(iter(current_chunk)) + + return heapq.merge(*chunks, key=key, reverse=reverse) + + +class ExternalList: + """ + ExternalList can have many items which cannot be hold in memory in + the same time. + + Examples + -------- + >>> l = ExternalList(list(range(100))) + >>> len(l) + 100 + >>> l.append(10) + >>> len(l) + 101 + >>> for i in range(20240): + ... l.append(i) + >>> len(l) + 20341 + >>> import pickle + >>> l2 = pickle.loads(pickle.dumps(l)) + >>> len(l2) + 20341 + >>> list(l2)[100] + 10 + """ + + LIMIT = 10240 + + def __init__(self, values): + self.values = values + self.count = len(values) + self._file = None + self._ser = None + + def __getstate__(self): + if self._file is not None: + self._file.flush() + with os.fdopen(os.dup(self._file.fileno()), "rb") as f: + f.seek(0) + serialized = f.read() + else: + serialized = b"" + return self.values, self.count, serialized + + def __setstate__(self, item): + self.values, self.count, serialized = item + if serialized: + self._open_file() + self._file.write(serialized) + else: + self._file = None + self._ser = None + + def __iter__(self): + if self._file is not None: + self._file.flush() + # read all items from disks first + with os.fdopen(os.dup(self._file.fileno()), "rb") as f: + f.seek(0) + for v in self._ser.load_stream(f): + yield v + + for v in self.values: + yield v + + def __len__(self): + return self.count + + def append(self, value): + self.values.append(value) + self.count += 1 + # dump them into disk if the key is huge + if len(self.values) >= self.LIMIT: + self._spill() + + def _open_file(self): + dirs = _get_local_dirs("objects") + d = dirs[id(self) % len(dirs)] + if not os.path.exists(d): + os.makedirs(d) + p = os.path.join(d, str(id(self))) + self._file = open(p, "w+b", 65536) + self._ser = BatchedSerializer(CompressedSerializer(CPickleSerializer()), 1024) + os.unlink(p) + + def __del__(self): + if self._file: + self._file.close() + self._file = None + + def _spill(self): + """dump the values into disk""" + global MemoryBytesSpilled, DiskBytesSpilled + if self._file is None: + self._open_file() + + used_memory = get_used_memory() + pos = self._file.tell() + self._ser.dump_stream(self.values, self._file) + self.values = [] + gc.collect() + DiskBytesSpilled += self._file.tell() - pos + MemoryBytesSpilled += max(used_memory - get_used_memory(), 0) << 20 + + +class ExternalListOfList(ExternalList): + """ + An external list for list. + + Examples + -------- + >>> l = ExternalListOfList([[i, i] for i in range(100)]) + >>> len(l) + 200 + >>> l.append(range(10)) + >>> len(l) + 210 + >>> len(list(l)) + 210 + """ + + def __init__(self, values): + ExternalList.__init__(self, values) + self.count = sum(len(i) for i in values) + + def append(self, value): + ExternalList.append(self, value) + # already counted 1 in ExternalList.append + self.count += len(value) - 1 + + def __iter__(self): + for values in ExternalList.__iter__(self): + for v in values: + yield v + + +class GroupByKey: + """ + Group a sorted iterator as [(k1, it1), (k2, it2), ...] + + Examples + -------- + >>> k = [i // 3 for i in range(6)] + >>> v = [[i] for i in range(6)] + >>> g = GroupByKey(zip(k, v)) + >>> [(k, list(it)) for k, it in g] + [(0, [0, 1, 2]), (1, [3, 4, 5])] + """ + + def __init__(self, iterator): + self.iterator = iterator + + def __iter__(self): + key, values = None, None + for k, v in self.iterator: + if values is not None and k == key: + values.append(v) + else: + if values is not None: + yield (key, values) + key = k + values = ExternalListOfList([v]) + if values is not None: + yield (key, values) + + +class ExternalGroupBy(ExternalMerger): + + """ + Group by the items by key. If any partition of them can not been + hold in memory, it will do sort based group by. + + This class works as follows: + + - It repeatedly group the items by key and save them in one dict in + memory. + + - When the used memory goes above memory limit, it will split + the combined data into partitions by hash code, dump them + into disk, one file per partition. If the number of keys + in one partitions is smaller than 1000, it will sort them + by key before dumping into disk. + + - Then it goes through the rest of the iterator, group items + by key into different dict by hash. Until the used memory goes over + memory limit, it dump all the dicts into disks, one file per + dict. Repeat this again until combine all the items. It + also will try to sort the items by key in each partition + before dumping into disks. + + - It will yield the grouped items partitions by partitions. + If the data in one partitions can be hold in memory, then it + will load and combine them in memory and yield. + + - If the dataset in one partition cannot be hold in memory, + it will sort them first. If all the files are already sorted, + it merge them by heap.merge(), so it will do external sort + for all the files. + + - After sorting, `GroupByKey` class will put all the continuous + items with the same key as a group, yield the values as + an iterator. + """ + + SORT_KEY_LIMIT = 1000 + + def flattened_serializer(self): + assert isinstance(self.serializer, BatchedSerializer) + ser = self.serializer + return FlattenedValuesSerializer(ser, 20) + + def _object_size(self, obj): + return len(obj) + + def _spill(self): + """ + dump already partitioned data into disks. + """ + global MemoryBytesSpilled, DiskBytesSpilled + path = self._get_spill_dir(self.spills) + if not os.path.exists(path): + os.makedirs(path) + + used_memory = get_used_memory() + if not self.pdata: + # The data has not been partitioned, it will iterator the + # data once, write them into different files, has no + # additional memory. It only called when the memory goes + # above limit at the first time. + + # open all the files for writing + streams = [open(os.path.join(path, str(i)), "wb") for i in range(self.partitions)] + + # If the number of keys is small, then the overhead of sort is small + # sort them before dumping into disks + self._sorted = len(self.data) < self.SORT_KEY_LIMIT + if self._sorted: + self.serializer = self.flattened_serializer() + for k in sorted(self.data.keys()): + h = self._partition(k) + self.serializer.dump_stream([(k, self.data[k])], streams[h]) + else: + for k, v in self.data.items(): + h = self._partition(k) + self.serializer.dump_stream([(k, v)], streams[h]) + + for s in streams: + DiskBytesSpilled += s.tell() + s.close() + + self.data.clear() + # self.pdata is cached in `mergeValues` and `mergeCombiners` + self.pdata.extend([{} for i in range(self.partitions)]) + + else: + for i in range(self.partitions): + p = os.path.join(path, str(i)) + with open(p, "wb") as f: + # dump items in batch + if self._sorted: + # sort by key only (stable) + sorted_items = sorted(self.pdata[i].items(), key=operator.itemgetter(0)) + self.serializer.dump_stream(sorted_items, f) + else: + self.serializer.dump_stream(self.pdata[i].items(), f) + self.pdata[i].clear() + DiskBytesSpilled += os.path.getsize(p) + + self.spills += 1 + gc.collect() # release the memory as much as possible + MemoryBytesSpilled += max(used_memory - get_used_memory(), 0) << 20 + + def _merged_items(self, index): + size = sum( + os.path.getsize(os.path.join(self._get_spill_dir(j), str(index))) + for j in range(self.spills) + ) + # if the memory can not hold all the partition, + # then use sort based merge. Because of compression, + # the data on disks will be much smaller than needed memory + if size >= self.memory_limit << 17: # * 1M / 8 + return self._merge_sorted_items(index) + + self.data = {} + for j in range(self.spills): + path = self._get_spill_dir(j) + p = os.path.join(path, str(index)) + # do not check memory during merging + with open(p, "rb") as f: + self.mergeCombiners(self.serializer.load_stream(f), 0) + return self.data.items() + + def _merge_sorted_items(self, index): + """load a partition from disk, then sort and group by key""" + + def load_partition(j): + path = self._get_spill_dir(j) + p = os.path.join(path, str(index)) + with open(p, "rb", 65536) as f: + for v in self.serializer.load_stream(f): + yield v + + disk_items = [load_partition(j) for j in range(self.spills)] + + if self._sorted: + # all the partitions are already sorted + sorted_items = heapq.merge(*disk_items, key=operator.itemgetter(0)) + + else: + # Flatten the combined values, so it will not consume huge + # memory during merging sort. + ser = self.flattened_serializer() + sorter = ExternalSorter(self.memory_limit, ser) + sorted_items = sorter.sorted(itertools.chain(*disk_items), key=operator.itemgetter(0)) + return ((k, vs) for k, vs in GroupByKey(sorted_items)) + + +if __name__ == "__main__": + import doctest + + (failure_count, test_count) = doctest.testmod() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/__init__.py new file mode 100644 index 0000000..d0d6948 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/__init__.py @@ -0,0 +1,75 @@ +# +# 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. +# + +""" +Important classes of Spark SQL and DataFrames: + + - :class:`pyspark.sql.SparkSession` + Main entry point for :class:`DataFrame` and SQL functionality. + - :class:`pyspark.sql.DataFrame` + A distributed collection of data grouped into named columns. + - :class:`pyspark.sql.Column` + A column expression in a :class:`DataFrame`. + - :class:`pyspark.sql.Row` + A row of data in a :class:`DataFrame`. + - :class:`pyspark.sql.GroupedData` + Aggregation methods, returned by :func:`DataFrame.groupBy`. + - :class:`pyspark.sql.DataFrameNaFunctions` + Methods for handling missing data (null values). + - :class:`pyspark.sql.DataFrameStatFunctions` + Methods for statistics functionality. + - :class:`pyspark.sql.functions` + List of built-in functions available for :class:`DataFrame`. + - :class:`pyspark.sql.types` + List of data types available. + - :class:`pyspark.sql.Window` + For working with window functions. +""" +from pyspark.sql.types import Row +from pyspark.sql.context import SQLContext, HiveContext, UDFRegistration, UDTFRegistration +from pyspark.sql.session import SparkSession +from pyspark.sql.column import Column +from pyspark.sql.catalog import Catalog +from pyspark.sql.dataframe import DataFrame, DataFrameNaFunctions, DataFrameStatFunctions +from pyspark.sql.group import GroupedData +from pyspark.sql.observation import Observation +from pyspark.sql.readwriter import DataFrameReader, DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.window import Window, WindowSpec +from pyspark.sql.pandas.group_ops import PandasCogroupedOps + + +__all__ = [ + "SparkSession", + "SQLContext", + "HiveContext", + "UDFRegistration", + "UDTFRegistration", + "DataFrame", + "GroupedData", + "Column", + "Catalog", + "Observation", + "Row", + "DataFrameNaFunctions", + "DataFrameStatFunctions", + "Window", + "WindowSpec", + "DataFrameReader", + "DataFrameWriter", + "DataFrameWriterV2", + "PandasCogroupedOps", +] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/_typing.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/_typing.pyi new file mode 100644 index 0000000..3d095f5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/_typing.pyi @@ -0,0 +1,80 @@ +# +# 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 ( + Any, + Callable, + List, + Optional, + Tuple, + TypeVar, + Union, +) +from typing_extensions import Literal, Protocol + +import datetime +import decimal + +from pyspark._typing import PrimitiveType +import pyspark.sql.types +from pyspark.sql.column import Column + +ColumnOrName = Union[Column, str] +ColumnOrName_ = TypeVar("ColumnOrName_", bound=ColumnOrName) +DecimalLiteral = decimal.Decimal +DateTimeLiteral = Union[datetime.datetime, datetime.date] +LiteralType = PrimitiveType +AtomicDataTypeOrString = Union[pyspark.sql.types.AtomicType, str] +DataTypeOrString = Union[pyspark.sql.types.DataType, str] +OptionalPrimitiveType = Optional[PrimitiveType] + +AtomicValue = TypeVar( + "AtomicValue", + datetime.datetime, + datetime.date, + decimal.Decimal, + bool, + str, + int, + float, +) + +RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...], pyspark.sql.types.Row) + +SQLBatchedUDFType = Literal[100] +SQLArrowBatchedUDFType = Literal[101] +SQLTableUDFType = Literal[300] +SQLArrowTableUDFType = Literal[301] + +class SupportsOpen(Protocol): + def open(self, partition_id: int, epoch_id: int) -> bool: ... + +class SupportsProcess(Protocol): + def process(self, row: pyspark.sql.types.Row) -> None: ... + +class SupportsClose(Protocol): + def close(self, error: Exception) -> None: ... + +class UserDefinedFunctionLike(Protocol): + func: Callable[..., Any] + evalType: int + deterministic: bool + @property + def returnType(self) -> pyspark.sql.types.DataType: ... + def __call__(self, *args: ColumnOrName) -> Column: ... + def asNondeterministic(self) -> UserDefinedFunctionLike: ... diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/__init__.py new file mode 100644 index 0000000..ac530a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +__all__ = ["functions"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/functions.py new file mode 100644 index 0000000..e49953e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/avro/functions.py @@ -0,0 +1,188 @@ +# +# 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. +# + +""" +A collections of builtin avro functions +""" + + +from typing import Dict, Optional, TYPE_CHECKING, cast + +from py4j.java_gateway import JVMView + +from pyspark.sql.column import Column, _to_java_column +from pyspark.sql.utils import get_active_spark_context, try_remote_avro_functions +from pyspark.util import _print_missing_jar + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName + + +@try_remote_avro_functions +def from_avro( + data: "ColumnOrName", jsonFormatSchema: str, options: Optional[Dict[str, str]] = None +) -> Column: + """ + Converts a binary column of Avro format into its corresponding catalyst value. + The specified schema must match the read data, otherwise the behavior is undefined: + it may fail or return arbitrary result. + To deserialize the data with a compatible and evolved schema, the expected Avro schema can be + set via the option avroSchema. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + data : :class:`~pyspark.sql.Column` or str + the binary column. + jsonFormatSchema : str + the avro schema in JSON string format. + options : dict, optional + options to control how the Avro record is parsed. + + Notes + ----- + Avro is built-in but external data source module since Spark 2.4. Please deploy the + application as per the deployment section of "Apache Avro Data Source Guide". + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.avro.functions import from_avro, to_avro + >>> data = [(1, Row(age=2, name='Alice'))] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> avroDf = df.select(to_avro(df.value).alias("avro")) + >>> avroDf.collect() + [Row(avro=bytearray(b'\\x00\\x00\\x04\\x00\\nAlice'))] + + >>> jsonFormatSchema = '''{"type":"record","name":"topLevelRecord","fields": + ... [{"name":"avro","type":[{"type":"record","name":"value","namespace":"topLevelRecord", + ... "fields":[{"name":"age","type":["long","null"]}, + ... {"name":"name","type":["string","null"]}]},"null"]}]}''' + >>> avroDf.select(from_avro(avroDf.avro, jsonFormatSchema).alias("value")).collect() + [Row(value=Row(avro=Row(age=2, name='Alice')))] + """ + + sc = get_active_spark_context() + try: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.avro.functions.from_avro( + _to_java_column(data), jsonFormatSchema, options or {} + ) + except TypeError as e: + if str(e) == "'JavaPackage' object is not callable": + _print_missing_jar("Avro", "avro", "avro", sc.version) + raise + return Column(jc) + + +@try_remote_avro_functions +def to_avro(data: "ColumnOrName", jsonFormatSchema: str = "") -> Column: + """ + Converts a column into binary of avro format. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + data : :class:`~pyspark.sql.Column` or str + the data column. + jsonFormatSchema : str, optional + user-specified output avro schema in JSON string format. + + Notes + ----- + Avro is built-in but external data source module since Spark 2.4. Please deploy the + application as per the deployment section of "Apache Avro Data Source Guide". + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.avro.functions import to_avro + >>> data = ['SPADES'] + >>> df = spark.createDataFrame(data, "string") + >>> df.select(to_avro(df.value).alias("suite")).collect() + [Row(suite=bytearray(b'\\x00\\x0cSPADES'))] + + >>> jsonFormatSchema = '''["null", {"type": "enum", "name": "value", + ... "symbols": ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"]}]''' + >>> df.select(to_avro(df.value, jsonFormatSchema).alias("suite")).collect() + [Row(suite=bytearray(b'\\x02\\x00'))] + """ + + sc = get_active_spark_context() + try: + if jsonFormatSchema == "": + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.avro.functions.to_avro( + _to_java_column(data) + ) + else: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.avro.functions.to_avro( + _to_java_column(data), jsonFormatSchema + ) + except TypeError as e: + if str(e) == "'JavaPackage' object is not callable": + _print_missing_jar("Avro", "avro", "avro", sc.version) + raise + return Column(jc) + + +def _test() -> None: + import os + import sys + from pyspark.testing.utils import search_jar + + avro_jar = search_jar("connector/avro", "spark-avro", "spark-avro") + if avro_jar is None: + print( + "Skipping all Avro Python tests as the optional Avro project was " + "not compiled into a JAR. To run these tests, " + "you need to build Spark with 'build/sbt -Pavro package' or " + "'build/mvn -Pavro package' before running this test." + ) + sys.exit(0) + else: + existing_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + jars_args = "--jars %s" % avro_jar + os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) + + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.avro.functions + + globs = pyspark.sql.avro.functions.__dict__.copy() + spark = ( + SparkSession.builder.master("local[4]").appName("sql.avro.functions tests").getOrCreate() + ) + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.avro.functions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/catalog.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/catalog.py new file mode 100644 index 0000000..2c6ed28 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/catalog.py @@ -0,0 +1,1271 @@ +# +# 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 sys +import warnings +from typing import Any, Callable, NamedTuple, List, Optional, TYPE_CHECKING + +from pyspark.storagelevel import StorageLevel +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.session import SparkSession +from pyspark.sql.types import StructType + +if TYPE_CHECKING: + from pyspark.sql._typing import UserDefinedFunctionLike + from pyspark.sql.types import DataType + + +class CatalogMetadata(NamedTuple): + name: str + description: Optional[str] + + +class Database(NamedTuple): + name: str + catalog: Optional[str] + description: Optional[str] + locationUri: str + + +class Table(NamedTuple): + name: str + catalog: Optional[str] + namespace: Optional[List[str]] + description: Optional[str] + tableType: str + isTemporary: bool + + @property + def database(self) -> Optional[str]: + if self.namespace is not None and len(self.namespace) == 1: + return self.namespace[0] + else: + return None + + +class Column(NamedTuple): + name: str + description: Optional[str] + dataType: str + nullable: bool + isPartition: bool + isBucket: bool + + +class Function(NamedTuple): + name: str + catalog: Optional[str] + namespace: Optional[List[str]] + description: Optional[str] + className: str + isTemporary: bool + + +class Catalog: + """User-facing catalog API, accessible through `SparkSession.catalog`. + + This is a thin wrapper around its Scala implementation org.apache.spark.sql.catalog.Catalog. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, sparkSession: SparkSession) -> None: + """Create a new Catalog that wraps the underlying JVM object.""" + self._sparkSession = sparkSession + self._jsparkSession = sparkSession._jsparkSession + self._sc = sparkSession._sc + self._jcatalog = sparkSession._jsparkSession.catalog() + + def currentCatalog(self) -> str: + """Returns the current default catalog in this session. + + .. versionadded:: 3.4.0 + + Examples + -------- + >>> spark.catalog.currentCatalog() + 'spark_catalog' + """ + return self._jcatalog.currentCatalog() + + def setCurrentCatalog(self, catalogName: str) -> None: + """Sets the current default catalog in this session. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + catalogName : str + name of the catalog to set + + Examples + -------- + >>> spark.catalog.setCurrentCatalog("spark_catalog") + """ + return self._jcatalog.setCurrentCatalog(catalogName) + + def listCatalogs(self, pattern: Optional[str] = None) -> List[CatalogMetadata]: + """Returns a list of catalogs in this session. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + pattern : str + The pattern that the catalog name needs to match. + + .. versionchanged: 3.5.0 + Added ``pattern`` argument. + + Returns + ------- + list + A list of :class:`CatalogMetadata`. + + Examples + -------- + >>> spark.catalog.listCatalogs() + [CatalogMetadata(name='spark_catalog', description=None)] + + >>> spark.catalog.listCatalogs("spark*") + [CatalogMetadata(name='spark_catalog', description=None)] + + >>> spark.catalog.listCatalogs("hive*") + [] + """ + if pattern is None: + iter = self._jcatalog.listCatalogs().toLocalIterator() + else: + iter = self._jcatalog.listCatalogs(pattern).toLocalIterator() + catalogs = [] + while iter.hasNext(): + jcatalog = iter.next() + catalogs.append( + CatalogMetadata(name=jcatalog.name(), description=jcatalog.description()) + ) + return catalogs + + def currentDatabase(self) -> str: + """ + Returns the current default database in this session. + + .. versionadded:: 2.0.0 + + Returns + ------- + str + The current default database name. + + Examples + -------- + >>> spark.catalog.currentDatabase() + 'default' + """ + return self._jcatalog.currentDatabase() + + def setCurrentDatabase(self, dbName: str) -> None: + """ + Sets the current default database in this session. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> spark.catalog.setCurrentDatabase("default") + """ + return self._jcatalog.setCurrentDatabase(dbName) + + def listDatabases(self, pattern: Optional[str] = None) -> List[Database]: + """ + Returns a list of databases available across all sessions. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + pattern : str + The pattern that the database name needs to match. + + .. versionchanged: 3.5.0 + Adds ``pattern`` argument. + + Returns + ------- + list + A list of :class:`Database`. + + Examples + -------- + >>> spark.catalog.listDatabases() + [Database(name='default', catalog='spark_catalog', description='default database', ... + + >>> spark.catalog.listDatabases("def*") + [Database(name='default', catalog='spark_catalog', description='default database', ... + + >>> spark.catalog.listDatabases("def2*") + [] + """ + if pattern is None: + iter = self._jcatalog.listDatabases().toLocalIterator() + else: + iter = self._jcatalog.listDatabases(pattern).toLocalIterator() + databases = [] + while iter.hasNext(): + jdb = iter.next() + databases.append( + Database( + name=jdb.name(), + catalog=jdb.catalog(), + description=jdb.description(), + locationUri=jdb.locationUri(), + ) + ) + return databases + + def getDatabase(self, dbName: str) -> Database: + """Get the database with the specified name. + This throws an :class:`AnalysisException` when the database cannot be found. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + dbName : str + name of the database to get. + + Returns + ------- + :class:`Database` + The database found by the name. + + Examples + -------- + >>> spark.catalog.getDatabase("default") + Database(name='default', catalog='spark_catalog', description='default database', ... + + Using the fully qualified name with the catalog name. + + >>> spark.catalog.getDatabase("spark_catalog.default") + Database(name='default', catalog='spark_catalog', description='default database', ... + """ + jdb = self._jcatalog.getDatabase(dbName) + return Database( + name=jdb.name(), + catalog=jdb.catalog(), + description=jdb.description(), + locationUri=jdb.locationUri(), + ) + + def databaseExists(self, dbName: str) -> bool: + """Check if the database with the specified name exists. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + dbName : str + name of the database to check existence + + .. versionchanged:: 3.4.0 + Allow ``dbName`` to be qualified with catalog name. + + Returns + ------- + bool + Indicating whether the database exists + + Examples + -------- + Check if 'test_new_database' database exists + + >>> spark.catalog.databaseExists("test_new_database") + False + >>> _ = spark.sql("CREATE DATABASE test_new_database") + >>> spark.catalog.databaseExists("test_new_database") + True + + Using the fully qualified name with the catalog name. + + >>> spark.catalog.databaseExists("spark_catalog.test_new_database") + True + >>> _ = spark.sql("DROP DATABASE test_new_database") + """ + return self._jcatalog.databaseExists(dbName) + + def listTables( + self, dbName: Optional[str] = None, pattern: Optional[str] = None + ) -> List[Table]: + """Returns a list of tables/views in the specified database. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + dbName : str + name of the database to list the tables. + + .. versionchanged:: 3.4.0 + Allow ``dbName`` to be qualified with catalog name. + + pattern : str + The pattern that the database name needs to match. + + .. versionchanged: 3.5.0 + Adds ``pattern`` argument. + + Returns + ------- + list + A list of :class:`Table`. + + Notes + ----- + If no database is specified, the current database and catalog + are used. This API includes all temporary views. + + Examples + -------- + >>> spark.range(1).createTempView("test_view") + >>> spark.catalog.listTables() + [Table(name='test_view', catalog=None, namespace=[], description=None, ... + + >>> spark.catalog.listTables(pattern="test*") + [Table(name='test_view', catalog=None, namespace=[], description=None, ... + + >>> spark.catalog.listTables(pattern="table*") + [] + + >>> _ = spark.catalog.dropTempView("test_view") + >>> spark.catalog.listTables() + [] + """ + if dbName is None: + dbName = self.currentDatabase() + + if pattern is None: + iter = self._jcatalog.listTables(dbName).toLocalIterator() + else: + iter = self._jcatalog.listTables(dbName, pattern).toLocalIterator() + tables = [] + while iter.hasNext(): + jtable = iter.next() + + jnamespace = jtable.namespace() + if jnamespace is not None: + namespace = [jnamespace[i] for i in range(0, len(jnamespace))] + else: + namespace = None + + tables.append( + Table( + name=jtable.name(), + catalog=jtable.catalog(), + namespace=namespace, + description=jtable.description(), + tableType=jtable.tableType(), + isTemporary=jtable.isTemporary(), + ) + ) + return tables + + def getTable(self, tableName: str) -> Table: + """Get the table or view with the specified name. This table can be a temporary view or a + table/view. This throws an :class:`AnalysisException` when no Table can be found. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + tableName : str + name of the table to get. + + .. versionchanged:: 3.4.0 + Allow `tableName` to be qualified with catalog name. + + Returns + ------- + :class:`Table` + The table found by the name. + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.getTable("tbl1") + Table(name='tbl1', catalog='spark_catalog', namespace=['default'], ... + + Using the fully qualified name with the catalog name. + + >>> spark.catalog.getTable("default.tbl1") + Table(name='tbl1', catalog='spark_catalog', namespace=['default'], ... + >>> spark.catalog.getTable("spark_catalog.default.tbl1") + Table(name='tbl1', catalog='spark_catalog', namespace=['default'], ... + >>> _ = spark.sql("DROP TABLE tbl1") + + Throw an analysis exception when the table does not exist. + + >>> spark.catalog.getTable("tbl1") + Traceback (most recent call last): + ... + AnalysisException: ... + """ + jtable = self._jcatalog.getTable(tableName) + jnamespace = jtable.namespace() + if jnamespace is not None: + namespace = [jnamespace[i] for i in range(0, len(jnamespace))] + else: + namespace = None + return Table( + name=jtable.name(), + catalog=jtable.catalog(), + namespace=namespace, + description=jtable.description(), + tableType=jtable.tableType(), + isTemporary=jtable.isTemporary(), + ) + + def listFunctions( + self, dbName: Optional[str] = None, pattern: Optional[str] = None + ) -> List[Function]: + """ + Returns a list of functions registered in the specified database. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + dbName : str + name of the database to list the functions. + ``dbName`` can be qualified with catalog name. + pattern : str + The pattern that the function name needs to match. + + .. versionchanged: 3.5.0 + Adds ``pattern`` argument. + + Returns + ------- + list + A list of :class:`Function`. + + Notes + ----- + If no database is specified, the current database and catalog + are used. This API includes all temporary functions. + + Examples + -------- + >>> spark.catalog.listFunctions() + [Function(name=... + + >>> spark.catalog.listFunctions(pattern="to_*") + [Function(name=... + + >>> spark.catalog.listFunctions(pattern="*not_existing_func*") + [] + """ + if dbName is None: + dbName = self.currentDatabase() + iter = self._jcatalog.listFunctions(dbName).toLocalIterator() + if pattern is None: + iter = self._jcatalog.listFunctions(dbName).toLocalIterator() + else: + iter = self._jcatalog.listFunctions(dbName, pattern).toLocalIterator() + functions = [] + while iter.hasNext(): + jfunction = iter.next() + jnamespace = jfunction.namespace() + if jnamespace is not None: + namespace = [jnamespace[i] for i in range(0, len(jnamespace))] + else: + namespace = None + + functions.append( + Function( + name=jfunction.name(), + catalog=jfunction.catalog(), + namespace=namespace, + description=jfunction.description(), + className=jfunction.className(), + isTemporary=jfunction.isTemporary(), + ) + ) + return functions + + def functionExists(self, functionName: str, dbName: Optional[str] = None) -> bool: + """Check if the function with the specified name exists. + This can either be a temporary function or a function. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + functionName : str + name of the function to check existence + + .. versionchanged:: 3.4.0 + Allow ``functionName`` to be qualified with catalog name + + dbName : str, optional + name of the database to check function existence in. + + Returns + ------- + bool + Indicating whether the function exists + + Notes + ----- + If no database is specified, the current database and catalog + are used. This API includes all temporary functions. + + Examples + -------- + >>> spark.catalog.functionExists("count") + True + + Using the fully qualified name for function name. + + >>> spark.catalog.functionExists("default.unexisting_function") + False + >>> spark.catalog.functionExists("spark_catalog.default.unexisting_function") + False + """ + if dbName is None: + return self._jcatalog.functionExists(functionName) + else: + warnings.warn( + "`dbName` has been deprecated since Spark 3.4 and might be removed in " + "a future version. Use functionExists(`dbName.tableName`) instead.", + FutureWarning, + ) + return self._jcatalog.functionExists(dbName, functionName) + + def getFunction(self, functionName: str) -> Function: + """Get the function with the specified name. This function can be a temporary function or a + function. This throws an :class:`AnalysisException` when the function cannot be found. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + functionName : str + name of the function to check existence. + + Returns + ------- + :class:`Function` + The function found by the name. + + Examples + -------- + >>> _ = spark.sql( + ... "CREATE FUNCTION my_func1 AS 'test.org.apache.spark.sql.MyDoubleAvg'") + >>> spark.catalog.getFunction("my_func1") + Function(name='my_func1', catalog='spark_catalog', namespace=['default'], ... + + Using the fully qualified name for function name. + + >>> spark.catalog.getFunction("default.my_func1") + Function(name='my_func1', catalog='spark_catalog', namespace=['default'], ... + >>> spark.catalog.getFunction("spark_catalog.default.my_func1") + Function(name='my_func1', catalog='spark_catalog', namespace=['default'], ... + + Throw an analysis exception when the function does not exists. + + >>> spark.catalog.getFunction("my_func2") + Traceback (most recent call last): + ... + AnalysisException: ... + """ + jfunction = self._jcatalog.getFunction(functionName) + jnamespace = jfunction.namespace() + if jnamespace is not None: + namespace = [jnamespace[i] for i in range(0, len(jnamespace))] + else: + namespace = None + return Function( + name=jfunction.name(), + catalog=jfunction.catalog(), + namespace=namespace, + description=jfunction.description(), + className=jfunction.className(), + isTemporary=jfunction.isTemporary(), + ) + + def listColumns(self, tableName: str, dbName: Optional[str] = None) -> List[Column]: + """Returns a list of columns for the given table/view in the specified database. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + tableName : str + name of the table to list columns. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name when ``dbName`` is None. + + dbName : str, optional + name of the database to find the table to list columns. + + Returns + ------- + list + A list of :class:`Column`. + + Notes + ----- + The order of arguments here is different from that of its JVM counterpart + because Python does not support method overloading. + + If no database is specified, the current database and catalog + are used. This API includes all temporary views. + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tblA (name STRING, age INT) USING parquet") + >>> spark.catalog.listColumns("tblA") + [Column(name='name', description=None, dataType='string', nullable=True, ... + >>> _ = spark.sql("DROP TABLE tblA") + """ + if dbName is None: + iter = self._jcatalog.listColumns(tableName).toLocalIterator() + else: + warnings.warn( + "`dbName` has been deprecated since Spark 3.4 and might be removed in " + "a future version. Use listColumns(`dbName.tableName`) instead.", + FutureWarning, + ) + iter = self._jcatalog.listColumns(dbName, tableName).toLocalIterator() + + columns = [] + while iter.hasNext(): + jcolumn = iter.next() + columns.append( + Column( + name=jcolumn.name(), + description=jcolumn.description(), + dataType=jcolumn.dataType(), + nullable=jcolumn.nullable(), + isPartition=jcolumn.isPartition(), + isBucket=jcolumn.isBucket(), + ) + ) + return columns + + def tableExists(self, tableName: str, dbName: Optional[str] = None) -> bool: + """Check if the table or view with the specified name exists. + This can either be a temporary view or a table/view. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + tableName : str + name of the table to check existence. + If no database is specified, first try to treat ``tableName`` as a + multi-layer-namespace identifier, then try ``tableName`` as a normal table + name in the current database if necessary. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name when ``dbName`` is None. + + dbName : str, optional + name of the database to check table existence in. + + Returns + ------- + bool + Indicating whether the table/view exists + + Examples + -------- + This function can check if a table is defined or not: + + >>> spark.catalog.tableExists("unexisting_table") + False + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.tableExists("tbl1") + True + + Using the fully qualified names for tables. + + >>> spark.catalog.tableExists("default.tbl1") + True + >>> spark.catalog.tableExists("spark_catalog.default.tbl1") + True + >>> spark.catalog.tableExists("tbl1", "default") + True + >>> _ = spark.sql("DROP TABLE tbl1") + + Check if views exist: + + >>> spark.catalog.tableExists("view1") + False + >>> _ = spark.sql("CREATE VIEW view1 AS SELECT 1") + >>> spark.catalog.tableExists("view1") + True + + Using the fully qualified names for views. + + >>> spark.catalog.tableExists("default.view1") + True + >>> spark.catalog.tableExists("spark_catalog.default.view1") + True + >>> spark.catalog.tableExists("view1", "default") + True + >>> _ = spark.sql("DROP VIEW view1") + + Check if temporary views exist: + + >>> _ = spark.sql("CREATE TEMPORARY VIEW view1 AS SELECT 1") + >>> spark.catalog.tableExists("view1") + True + >>> df = spark.sql("DROP VIEW view1") + >>> spark.catalog.tableExists("view1") + False + """ + if dbName is None: + return self._jcatalog.tableExists(tableName) + else: + warnings.warn( + "`dbName` has been deprecated since Spark 3.4 and might be removed in " + "a future version. Use tableExists(`dbName.tableName`) instead.", + FutureWarning, + ) + return self._jcatalog.tableExists(dbName, tableName) + + def createExternalTable( + self, + tableName: str, + path: Optional[str] = None, + source: Optional[str] = None, + schema: Optional[StructType] = None, + **options: str, + ) -> DataFrame: + """Creates a table based on the dataset in a data source. + + It returns the DataFrame associated with the external table. + + The data source is specified by the ``source`` and a set of ``options``. + If ``source`` is not specified, the default data source configured by + ``spark.sql.sources.default`` will be used. + + Optionally, a schema can be provided as the schema of the returned :class:`DataFrame` and + created external table. + + .. versionadded:: 2.0.0 + + Returns + ------- + :class:`DataFrame` + """ + warnings.warn( + "createExternalTable is deprecated since Spark 2.2, please use createTable instead.", + FutureWarning, + ) + return self.createTable(tableName, path, source, schema, **options) + + def createTable( + self, + tableName: str, + path: Optional[str] = None, + source: Optional[str] = None, + schema: Optional[StructType] = None, + description: Optional[str] = None, + **options: str, + ) -> DataFrame: + """Creates a table based on the dataset in a data source. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + tableName : str + name of the table to create. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name. + + path : str, optional + the path in which the data for this table exists. + When ``path`` is specified, an external table is + created from the data at the given path. Otherwise a managed table is created. + source : str, optional + the source of this table such as 'parquet, 'orc', etc. + If ``source`` is not specified, the default data source configured by + ``spark.sql.sources.default`` will be used. + schema : class:`StructType`, optional + the schema for this table. + description : str, optional + the description of this table. + + .. versionchanged:: 3.1.0 + Added the ``description`` parameter. + + **options : dict, optional + extra options to specify in the table. + + Returns + ------- + :class:`DataFrame` + The DataFrame associated with the table. + + Examples + -------- + Creating a managed table. + + >>> _ = spark.catalog.createTable("tbl1", schema=spark.range(1).schema, source='parquet') + >>> _ = spark.sql("DROP TABLE tbl1") + + Creating an external table + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... _ = spark.catalog.createTable( + ... "tbl2", schema=spark.range(1).schema, path=d, source='parquet') + >>> _ = spark.sql("DROP TABLE tbl2") + """ + if path is not None: + options["path"] = path + if source is None: + c = self._sparkSession._jconf + source = c.defaultDataSourceName() + if description is None: + description = "" + if schema is None: + df = self._jcatalog.createTable(tableName, source, description, options) + else: + if not isinstance(schema, StructType): + raise TypeError("schema should be StructType") + scala_datatype = self._jsparkSession.parseDataType(schema.json()) + df = self._jcatalog.createTable(tableName, source, scala_datatype, description, options) + return DataFrame(df, self._sparkSession) + + def dropTempView(self, viewName: str) -> bool: + """Drops the local temporary view with the given view name in the catalog. + If the view has been cached before, then it will also be uncached. + Returns true if this view is dropped successfully, false otherwise. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + viewName : str + name of the temporary view to drop. + + Returns + ------- + bool + If the temporary view was successfully dropped or not. + + .. versionadded:: 2.1.0 + The return type of this method was ``None`` in Spark 2.0, but changed to ``bool`` + in Spark 2.1. + + Examples + -------- + >>> spark.createDataFrame([(1, 1)]).createTempView("my_table") + + Dropping the temporary view. + + >>> spark.catalog.dropTempView("my_table") + True + + Throw an exception if the temporary view does not exists. + + >>> spark.table("my_table") + Traceback (most recent call last): + ... + AnalysisException: ... + """ + return self._jcatalog.dropTempView(viewName) + + def dropGlobalTempView(self, viewName: str) -> bool: + """Drops the global temporary view with the given view name in the catalog. + + .. versionadded:: 2.1.0 + + Parameters + ---------- + viewName : str + name of the global view to drop. + + Returns + ------- + bool + If the global view was successfully dropped or not. + + Notes + ----- + If the view has been cached before, then it will also be uncached. + + Examples + -------- + >>> spark.createDataFrame([(1, 1)]).createGlobalTempView("my_table") + + Dropping the global view. + + >>> spark.catalog.dropGlobalTempView("my_table") + True + + Throw an exception if the global view does not exists. + + >>> spark.table("global_temp.my_table") + Traceback (most recent call last): + ... + AnalysisException: ... + """ + return self._jcatalog.dropGlobalTempView(viewName) + + def registerFunction( + self, name: str, f: Callable[..., Any], returnType: Optional["DataType"] = None + ) -> "UserDefinedFunctionLike": + """An alias for :func:`spark.udf.register`. + See :meth:`pyspark.sql.UDFRegistration.register`. + + .. versionadded:: 2.0.0 + + .. deprecated:: 2.3.0 + Use :func:`spark.udf.register` instead. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + warnings.warn("Deprecated in 2.3.0. Use spark.udf.register instead.", FutureWarning) + return self._sparkSession.udf.register(name, f, returnType) + + def isCached(self, tableName: str) -> bool: + """ + Returns true if the table is currently cached in-memory. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + tableName : str + name of the table to get. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name. + + Returns + ------- + bool + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.cacheTable("tbl1") + >>> spark.catalog.isCached("tbl1") + True + + Throw an analysis exception when the table does not exist. + + >>> spark.catalog.isCached("not_existing_table") + Traceback (most recent call last): + ... + AnalysisException: ... + + Using the fully qualified name for the table. + + >>> spark.catalog.isCached("spark_catalog.default.tbl1") + True + >>> spark.catalog.uncacheTable("tbl1") + >>> _ = spark.sql("DROP TABLE tbl1") + """ + return self._jcatalog.isCached(tableName) + + def cacheTable(self, tableName: str, storageLevel: Optional[StorageLevel] = None) -> None: + """Caches the specified table in-memory or with given storage level. + Default MEMORY_AND_DISK. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + tableName : str + name of the table to get. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name. + + storageLevel : :class:`StorageLevel` + storage level to set for persistence. + + .. versionchanged:: 3.5.0 + Allow to specify storage level. + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.cacheTable("tbl1") + + or + + >>> spark.catalog.cacheTable("tbl1", StorageLevel.OFF_HEAP) + + Throw an analysis exception when the table does not exist. + + >>> spark.catalog.cacheTable("not_existing_table") + Traceback (most recent call last): + ... + AnalysisException: ... + + Using the fully qualified name for the table. + + >>> spark.catalog.cacheTable("spark_catalog.default.tbl1") + >>> spark.catalog.uncacheTable("tbl1") + >>> _ = spark.sql("DROP TABLE tbl1") + """ + if storageLevel: + javaStorageLevel = self._sc._getJavaStorageLevel(storageLevel) + self._jcatalog.cacheTable(tableName, javaStorageLevel) + else: + self._jcatalog.cacheTable(tableName) + + def uncacheTable(self, tableName: str) -> None: + """Removes the specified table from the in-memory cache. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + tableName : str + name of the table to get. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name. + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.cacheTable("tbl1") + >>> spark.catalog.uncacheTable("tbl1") + >>> spark.catalog.isCached("tbl1") + False + + Throw an analysis exception when the table does not exist. + + >>> spark.catalog.uncacheTable("not_existing_table") + Traceback (most recent call last): + ... + AnalysisException: ... + + Using the fully qualified name for the table. + + >>> spark.catalog.uncacheTable("spark_catalog.default.tbl1") + >>> spark.catalog.isCached("tbl1") + False + >>> _ = spark.sql("DROP TABLE tbl1") + """ + self._jcatalog.uncacheTable(tableName) + + def clearCache(self) -> None: + """Removes all cached tables from the in-memory cache. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1") + >>> _ = spark.sql("CREATE TABLE tbl1 (name STRING, age INT) USING parquet") + >>> spark.catalog.clearCache() + >>> spark.catalog.isCached("tbl1") + False + >>> _ = spark.sql("DROP TABLE tbl1") + """ + self._jcatalog.clearCache() + + def refreshTable(self, tableName: str) -> None: + """Invalidates and refreshes all the cached data and metadata of the given table. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + tableName : str + name of the table to get. + + .. versionchanged:: 3.4.0 + Allow ``tableName`` to be qualified with catalog name. + + Examples + -------- + The example below caches a table, and then removes the data. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... _ = spark.sql("DROP TABLE IF EXISTS tbl1") + ... _ = spark.sql( + ... "CREATE TABLE tbl1 (col STRING) USING TEXT LOCATION '{}'".format(d)) + ... _ = spark.sql("INSERT INTO tbl1 SELECT 'abc'") + ... spark.catalog.cacheTable("tbl1") + ... spark.table("tbl1").show() + +---+ + |col| + +---+ + |abc| + +---+ + + Because the table is cached, it computes from the cached data as below. + + >>> spark.table("tbl1").count() + 1 + + After refreshing the table, it shows 0 because the data does not exist anymore. + + >>> spark.catalog.refreshTable("tbl1") + >>> spark.table("tbl1").count() + 0 + + Using the fully qualified name for the table. + + >>> spark.catalog.refreshTable("spark_catalog.default.tbl1") + >>> _ = spark.sql("DROP TABLE tbl1") + """ + self._jcatalog.refreshTable(tableName) + + def recoverPartitions(self, tableName: str) -> None: + """Recovers all the partitions of the given table and updates the catalog. + + .. versionadded:: 2.1.1 + + Parameters + ---------- + tableName : str + name of the table to get. + + Notes + ----- + Only works with a partitioned table, and not a view. + + Examples + -------- + The example below creates a partitioned table against the existing directory of + the partitioned table. After that, it recovers the partitions. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... _ = spark.sql("DROP TABLE IF EXISTS tbl1") + ... spark.range(1).selectExpr( + ... "id as key", "id as value").write.partitionBy("key").mode("overwrite").save(d) + ... _ = spark.sql( + ... "CREATE TABLE tbl1 (key LONG, value LONG)" + ... "USING parquet OPTIONS (path '{}') PARTITIONED BY (key)".format(d)) + ... spark.table("tbl1").show() + ... spark.catalog.recoverPartitions("tbl1") + ... spark.table("tbl1").show() + +-----+---+ + |value|key| + +-----+---+ + +-----+---+ + +-----+---+ + |value|key| + +-----+---+ + | 0| 0| + +-----+---+ + >>> _ = spark.sql("DROP TABLE tbl1") + """ + self._jcatalog.recoverPartitions(tableName) + + def refreshByPath(self, path: str) -> None: + """Invalidates and refreshes all the cached data (and the associated metadata) for any + DataFrame that contains the given data source path. + + .. versionadded:: 2.2.0 + + Parameters + ---------- + path : str + the path to refresh the cache. + + Examples + -------- + The example below caches a table, and then removes the data. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... _ = spark.sql("DROP TABLE IF EXISTS tbl1") + ... _ = spark.sql( + ... "CREATE TABLE tbl1 (col STRING) USING TEXT LOCATION '{}'".format(d)) + ... _ = spark.sql("INSERT INTO tbl1 SELECT 'abc'") + ... spark.catalog.cacheTable("tbl1") + ... spark.table("tbl1").show() + +---+ + |col| + +---+ + |abc| + +---+ + + Because the table is cached, it computes from the cached data as below. + + >>> spark.table("tbl1").count() + 1 + + After refreshing the table by path, it shows 0 because the data does not exist anymore. + + >>> spark.catalog.refreshByPath(d) + >>> spark.table("tbl1").count() + 0 + + >>> _ = spark.sql("DROP TABLE tbl1") + """ + self._jcatalog.refreshByPath(path) + + def _reset(self) -> None: + """(Internal use only) Drop all existing databases (except "default"), tables, + partitions and functions, and set the current database to "default". + + This is mainly used for tests. + """ + self._jsparkSession.sessionState().catalog().reset() + + +def _test() -> None: + import os + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.catalog + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.catalog.__dict__.copy() + globs["spark"] = ( + SparkSession.builder.master("local[4]").appName("sql.catalog tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.sql.catalog, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + globs["spark"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/column.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/column.py new file mode 100644 index 0000000..a559cf2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/column.py @@ -0,0 +1,1431 @@ +# +# 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 sys +import json +import warnings +from typing import ( + cast, + overload, + Any, + Callable, + Iterable, + List, + Optional, + Tuple, + TYPE_CHECKING, + Union, +) + +from py4j.java_gateway import JavaObject, JVMView + +from pyspark import copy_func +from pyspark.context import SparkContext +from pyspark.errors import PySparkAttributeError, PySparkTypeError, PySparkValueError +from pyspark.sql.types import DataType +from pyspark.sql.utils import get_active_spark_context + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName, LiteralType, DecimalLiteral, DateTimeLiteral + from pyspark.sql.window import WindowSpec + +__all__ = ["Column"] + + +def _create_column_from_literal(literal: Union["LiteralType", "DecimalLiteral"]) -> "Column": + sc = get_active_spark_context() + return cast(JVMView, sc._jvm).functions.lit(literal) + + +def _create_column_from_name(name: str) -> "Column": + sc = get_active_spark_context() + return cast(JVMView, sc._jvm).functions.col(name) + + +def _to_java_column(col: "ColumnOrName") -> JavaObject: + if isinstance(col, Column): + jcol = col._jc + elif isinstance(col, str): + jcol = _create_column_from_name(col) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + return jcol + + +def _to_java_expr(col: "ColumnOrName") -> JavaObject: + return _to_java_column(col).expr() + + +def _to_seq( + sc: SparkContext, + cols: Iterable["ColumnOrName"], + converter: Optional[Callable[["ColumnOrName"], JavaObject]] = None, +) -> JavaObject: + """ + Convert a list of Columns (or names) into a JVM Seq of Column. + + An optional `converter` could be used to convert items in `cols` + into JVM Column objects. + """ + if converter: + cols = [converter(c) for c in cols] + assert sc._jvm is not None + return sc._jvm.PythonUtils.toSeq(cols) + + +def _to_list( + sc: SparkContext, + cols: List["ColumnOrName"], + converter: Optional[Callable[["ColumnOrName"], JavaObject]] = None, +) -> JavaObject: + """ + Convert a list of Columns (or names) into a JVM (Scala) List of Columns. + + An optional `converter` could be used to convert items in `cols` + into JVM Column objects. + """ + if converter: + cols = [converter(c) for c in cols] + assert sc._jvm is not None + return sc._jvm.PythonUtils.toList(cols) + + +def _unary_op( + name: str, + doc: str = "unary operator", +) -> Callable[["Column"], "Column"]: + """Create a method for given unary operator""" + + def _(self: "Column") -> "Column": + jc = getattr(self._jc, name)() + return Column(jc) + + _.__doc__ = doc + return _ + + +def _func_op(name: str, doc: str = "") -> Callable[["Column"], "Column"]: + def _(self: "Column") -> "Column": + sc = get_active_spark_context() + jc = getattr(cast(JVMView, sc._jvm).functions, name)(self._jc) + return Column(jc) + + _.__doc__ = doc + return _ + + +def _bin_func_op( + name: str, + reverse: bool = False, + doc: str = "binary function", +) -> Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"]: + def _(self: "Column", other: Union["Column", "LiteralType", "DecimalLiteral"]) -> "Column": + sc = get_active_spark_context() + fn = getattr(cast(JVMView, sc._jvm).functions, name) + jc = other._jc if isinstance(other, Column) else _create_column_from_literal(other) + njc = fn(self._jc, jc) if not reverse else fn(jc, self._jc) + return Column(njc) + + _.__doc__ = doc + return _ + + +def _bin_op( + name: str, + doc: str = "binary operator", +) -> Callable[ + ["Column", Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"]], "Column" +]: + """Create a method for given binary operator""" + + def _( + self: "Column", + other: Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"], + ) -> "Column": + jc = other._jc if isinstance(other, Column) else other + njc = getattr(self._jc, name)(jc) + return Column(njc) + + _.__doc__ = doc + return _ + + +def _reverse_op( + name: str, + doc: str = "binary operator", +) -> Callable[["Column", Union["LiteralType", "DecimalLiteral"]], "Column"]: + """Create a method for binary operator (this object is on right side)""" + + def _(self: "Column", other: Union["LiteralType", "DecimalLiteral"]) -> "Column": + jother = _create_column_from_literal(other) + jc = getattr(jother, name)(self._jc) + return Column(jc) + + _.__doc__ = doc + return _ + + +class Column: + + """ + A column in a DataFrame. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + Column instances can be created by + + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + + Select a column out of a DataFrame + >>> df.name + Column<'name'> + >>> df["name"] + Column<'name'> + + Create from an expression + + >>> df.age + 1 + Column<...> + >>> 1 / df.age + Column<...> + """ + + def __init__(self, jc: JavaObject) -> None: + self._jc = jc + + # arithmetic operators + __neg__ = _func_op("negate") + __add__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("plus"), + ) + __sub__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("minus"), + ) + __mul__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("multiply"), + ) + __div__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("divide"), + ) + __truediv__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("divide"), + ) + __mod__ = cast( + Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral"]], "Column"], + _bin_op("mod"), + ) + __radd__ = cast( + Callable[["Column", Union["LiteralType", "DecimalLiteral"]], "Column"], _bin_op("plus") + ) + __rsub__ = _reverse_op("minus") + __rmul__ = cast( + Callable[["Column", Union["LiteralType", "DecimalLiteral"]], "Column"], _bin_op("multiply") + ) + __rdiv__ = _reverse_op("divide") + __rtruediv__ = _reverse_op("divide") + __rmod__ = _reverse_op("mod") + + __pow__ = _bin_func_op("pow") + __rpow__ = cast( + Callable[["Column", Union["LiteralType", "DecimalLiteral"]], "Column"], + _bin_func_op("pow", reverse=True), + ) + + # logistic operators + def __eq__( # type: ignore[override] + self, + other: Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"], + ) -> "Column": + """binary function""" + return _bin_op("equalTo")(self, other) + + def __ne__( # type: ignore[override] + self, + other: Any, + ) -> "Column": + """binary function""" + return _bin_op("notEqual")(self, other) + + __lt__ = _bin_op("lt") + __le__ = _bin_op("leq") + __ge__ = _bin_op("geq") + __gt__ = _bin_op("gt") + + _eqNullSafe_doc = """ + Equality test that is safe for null values. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other + a value or :class:`Column` + + Examples + -------- + >>> from pyspark.sql import Row + >>> df1 = spark.createDataFrame([ + ... Row(id=1, value='foo'), + ... Row(id=2, value=None) + ... ]) + >>> df1.select( + ... df1['value'] == 'foo', + ... df1['value'].eqNullSafe('foo'), + ... df1['value'].eqNullSafe(None) + ... ).show() + +-------------+---------------+----------------+ + |(value = foo)|(value <=> foo)|(value <=> NULL)| + +-------------+---------------+----------------+ + | true| true| false| + | NULL| false| true| + +-------------+---------------+----------------+ + >>> df2 = spark.createDataFrame([ + ... Row(value = 'bar'), + ... Row(value = None) + ... ]) + >>> df1.join(df2, df1["value"] == df2["value"]).count() + 0 + >>> df1.join(df2, df1["value"].eqNullSafe(df2["value"])).count() + 1 + >>> df2 = spark.createDataFrame([ + ... Row(id=1, value=float('NaN')), + ... Row(id=2, value=42.0), + ... Row(id=3, value=None) + ... ]) + >>> df2.select( + ... df2['value'].eqNullSafe(None), + ... df2['value'].eqNullSafe(float('NaN')), + ... df2['value'].eqNullSafe(42.0) + ... ).show() + +----------------+---------------+----------------+ + |(value <=> NULL)|(value <=> NaN)|(value <=> 42.0)| + +----------------+---------------+----------------+ + | false| true| false| + | false| false| true| + | true| false| false| + +----------------+---------------+----------------+ + + Notes + ----- + Unlike Pandas, PySpark doesn't consider NaN values to be NULL. See the + `NaN Semantics `_ + for details. + """ + eqNullSafe = _bin_op("eqNullSafe", _eqNullSafe_doc) + + # `and`, `or`, `not` cannot be overloaded in Python, + # so use bitwise operators as boolean operators + __and__ = _bin_op("and") + __or__ = _bin_op("or") + __invert__ = _func_op("not") + __rand__ = _bin_op("and") + __ror__ = _bin_op("or") + + # container operators + def __contains__(self, item: Any) -> None: + raise PySparkValueError( + error_class="CANNOT_APPLY_IN_FOR_COLUMN", + message_parameters={}, + ) + + # bitwise operators + _bitwiseOR_doc = """ + Compute bitwise OR of this expression with another expression. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other + a value or :class:`Column` to calculate bitwise or(|) with + this :class:`Column`. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(a=170, b=75)]) + >>> df.select(df.a.bitwiseOR(df.b)).collect() + [Row((a | b)=235)] + """ + _bitwiseAND_doc = """ + Compute bitwise AND of this expression with another expression. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other + a value or :class:`Column` to calculate bitwise and(&) with + this :class:`Column`. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(a=170, b=75)]) + >>> df.select(df.a.bitwiseAND(df.b)).collect() + [Row((a & b)=10)] + """ + _bitwiseXOR_doc = """ + Compute bitwise XOR of this expression with another expression. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other + a value or :class:`Column` to calculate bitwise xor(^) with + this :class:`Column`. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(a=170, b=75)]) + >>> df.select(df.a.bitwiseXOR(df.b)).collect() + [Row((a ^ b)=225)] + """ + + bitwiseOR = _bin_op("bitwiseOR", _bitwiseOR_doc) + bitwiseAND = _bin_op("bitwiseAND", _bitwiseAND_doc) + bitwiseXOR = _bin_op("bitwiseXOR", _bitwiseXOR_doc) + + def getItem(self, key: Any) -> "Column": + """ + An expression that gets an item at position ``ordinal`` out of a list, + or gets an item by key out of a dict. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + key + a literal value, or a :class:`Column` expression. + The result will only be true at a location if the item matches in the column. + + .. deprecated:: 3.0.0 + :class:`Column` as a parameter is deprecated. + + Returns + ------- + :class:`Column` + Column representing the item(s) got at position out of a list or by key out of a dict. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 2], {"key": "value"})], ["l", "d"]) + >>> df.select(df.l.getItem(0), df.d.getItem("key")).show() + +----+------+ + |l[0]|d[key]| + +----+------+ + | 1| value| + +----+------+ + """ + if isinstance(key, Column): + warnings.warn( + "A column as 'key' in getItem is deprecated as of Spark 3.0, and will not " + "be supported in the future release. Use `column[key]` or `column.key` syntax " + "instead.", + FutureWarning, + ) + return self[key] + + def getField(self, name: Any) -> "Column": + """ + An expression that gets a field by name in a :class:`StructType`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name + a literal value, or a :class:`Column` expression. + The result will only be true at a location if the field matches in the Column. + + .. deprecated:: 3.0.0 + :class:`Column` as a parameter is deprecated. + Returns + ------- + :class:`Column` + Column representing whether each element of Column got by name. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(r=Row(a=1, b="b"))]) + >>> df.select(df.r.getField("b")).show() + +---+ + |r.b| + +---+ + | b| + +---+ + >>> df.select(df.r.a).show() + +---+ + |r.a| + +---+ + | 1| + +---+ + """ + if isinstance(name, Column): + warnings.warn( + "A column as 'name' in getField is deprecated as of Spark 3.0, and will not " + "be supported in the future release. Use `column[name]` or `column.name` syntax " + "instead.", + FutureWarning, + ) + return self[name] + + def withField(self, fieldName: str, col: "Column") -> "Column": + """ + An expression that adds/replaces a field in :class:`StructType` by name. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + fieldName : str + a literal value. + The result will only be true at a location if any field matches in the Column. + col : :class:`Column` + A :class:`Column` expression for the column with `fieldName`. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column + which field was added/replaced by fieldName. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import lit + >>> df = spark.createDataFrame([Row(a=Row(b=1, c=2))]) + >>> df.withColumn('a', df['a'].withField('b', lit(3))).select('a.b').show() + +---+ + | b| + +---+ + | 3| + +---+ + >>> df.withColumn('a', df['a'].withField('d', lit(4))).select('a.d').show() + +---+ + | d| + +---+ + | 4| + +---+ + """ + if not isinstance(fieldName, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "fieldName", "arg_type": type(fieldName).__name__}, + ) + + if not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + + return Column(self._jc.withField(fieldName, col._jc)) + + def dropFields(self, *fieldNames: str) -> "Column": + """ + An expression that drops fields in :class:`StructType` by name. + This is a no-op if the schema doesn't contain field name(s). + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + fieldNames : str + Desired field names (collects all positional arguments passed) + The result will drop at a location if any field matches in the Column. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column with field dropped by fieldName. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import col, lit + >>> df = spark.createDataFrame([ + ... Row(a=Row(b=1, c=2, d=3, e=Row(f=4, g=5, h=6)))]) + >>> df.withColumn('a', df['a'].dropFields('b')).show() + +-----------------+ + | a| + +-----------------+ + |{2, 3, {4, 5, 6}}| + +-----------------+ + + >>> df.withColumn('a', df['a'].dropFields('b', 'c')).show() + +--------------+ + | a| + +--------------+ + |{3, {4, 5, 6}}| + +--------------+ + + This method supports dropping multiple nested fields directly e.g. + + >>> df.withColumn("a", col("a").dropFields("e.g", "e.h")).show() + +--------------+ + | a| + +--------------+ + |{1, 2, 3, {4}}| + +--------------+ + + However, if you are going to add/replace multiple nested fields, + it is preferred to extract out the nested struct before + adding/replacing multiple fields e.g. + + >>> df.select(col("a").withField( + ... "e", col("a.e").dropFields("g", "h")).alias("a") + ... ).show() + +--------------+ + | a| + +--------------+ + |{1, 2, 3, {4}}| + +--------------+ + + """ + sc = get_active_spark_context() + jc = self._jc.dropFields(_to_seq(sc, fieldNames)) + return Column(jc) + + def __getattr__(self, item: Any) -> "Column": + """ + An expression that gets an item at position ``ordinal`` out of a list, + or gets an item by key out of a dict. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + item + a literal value. + + Returns + ------- + :class:`Column` + Column representing the item got by key out of a dict. + + Examples + -------- + >>> df = spark.createDataFrame([('abcedfg', {"key": "value"})], ["l", "d"]) + >>> df.select(df.d.key).show() + +------+ + |d[key]| + +------+ + | value| + +------+ + """ + if item.startswith("__"): + raise PySparkAttributeError( + error_class="CANNOT_ACCESS_TO_DUNDER", + message_parameters={}, + ) + return self[item] + + def __getitem__(self, k: Any) -> "Column": + """ + An expression that gets an item at position ``ordinal`` out of a list, + or gets an item by key out of a dict. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + k + a literal value, or a slice object without step. + + Returns + ------- + :class:`Column` + Column representing the item got by key out of a dict, or substrings sliced by + the given slice object. + + Examples + -------- + >>> df = spark.createDataFrame([('abcedfg', {"key": "value"})], ["l", "d"]) + >>> df.select(df.l[slice(1, 3)], df.d['key']).show() + +------------------+------+ + |substring(l, 1, 3)|d[key]| + +------------------+------+ + | abc| value| + +------------------+------+ + """ + if isinstance(k, slice): + if k.step is not None: + raise PySparkValueError( + error_class="SLICE_WITH_STEP", + message_parameters={}, + ) + return self.substr(k.start, k.stop) + else: + return _bin_op("apply")(self, k) + + def __iter__(self) -> None: + raise PySparkTypeError( + error_class="NOT_ITERABLE", message_parameters={"objectName": "Column"} + ) + + # string methods + _contains_doc = """ + Contains the other element. Returns a boolean :class:`Column` based on a string match. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other + string in line. A value as a literal or a :class:`Column`. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.contains('o')).collect() + [Row(age=5, name='Bob')] + """ + _startswith_doc = """ + String starts with. Returns a boolean :class:`Column` based on a string match. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`Column` or str + string at start of line (do not use a regex `^`) + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.startswith('Al')).collect() + [Row(age=2, name='Alice')] + >>> df.filter(df.name.startswith('^Al')).collect() + [] + """ + _endswith_doc = """ + String ends with. Returns a boolean :class:`Column` based on a string match. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`Column` or str + string at end of line (do not use a regex `$`) + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.endswith('ice')).collect() + [Row(age=2, name='Alice')] + >>> df.filter(df.name.endswith('ice$')).collect() + [] + """ + + contains = _bin_op("contains", _contains_doc) + startswith = _bin_op("startsWith", _startswith_doc) + endswith = _bin_op("endsWith", _endswith_doc) + + def like(self: "Column", other: str) -> "Column": + """ + SQL like expression. Returns a boolean :class:`Column` based on a SQL LIKE match. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : str + a SQL LIKE pattern + + See Also + -------- + pyspark.sql.Column.rlike + + Returns + ------- + :class:`Column` + Column of booleans showing whether each element + in the Column is matched by SQL LIKE pattern. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.like('Al%')).collect() + [Row(age=2, name='Alice')] + """ + njc = getattr(self._jc, "like")(other) + return Column(njc) + + def rlike(self: "Column", other: str) -> "Column": + """ + SQL RLIKE expression (LIKE with Regex). Returns a boolean :class:`Column` based on a regex + match. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : str + an extended regex expression + + Returns + ------- + :class:`Column` + Column of booleans showing whether each element + in the Column is matched by extended regex expression. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.rlike('ice$')).collect() + [Row(age=2, name='Alice')] + """ + njc = getattr(self._jc, "rlike")(other) + return Column(njc) + + def ilike(self: "Column", other: str) -> "Column": + """ + SQL ILIKE expression (case insensitive LIKE). Returns a boolean :class:`Column` + based on a case insensitive match. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : str + a SQL LIKE pattern + + See Also + -------- + pyspark.sql.Column.rlike + + Returns + ------- + :class:`Column` + Column of booleans showing whether each element + in the Column is matched by SQL LIKE pattern. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.filter(df.name.ilike('%Ice')).collect() + [Row(age=2, name='Alice')] + """ + njc = getattr(self._jc, "ilike")(other) + return Column(njc) + + @overload + def substr(self, startPos: int, length: int) -> "Column": + ... + + @overload + def substr(self, startPos: "Column", length: "Column") -> "Column": + ... + + def substr(self, startPos: Union[int, "Column"], length: Union[int, "Column"]) -> "Column": + """ + Return a :class:`Column` which is a substring of the column. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + startPos : :class:`Column` or int + start position + length : :class:`Column` or int + length of the substring + + Returns + ------- + :class:`Column` + Column representing whether each element of Column is substr of origin Column. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.name.substr(1, 3).alias("col")).collect() + [Row(col='Ali'), Row(col='Bob')] + """ + if type(startPos) != type(length): + raise PySparkTypeError( + error_class="NOT_SAME_TYPE", + message_parameters={ + "arg_name1": "startPos", + "arg_name2": "length", + "arg_type1": type(startPos).__name__, + "arg_type2": type(length).__name__, + }, + ) + if isinstance(startPos, int): + jc = self._jc.substr(startPos, length) + elif isinstance(startPos, Column): + jc = self._jc.substr(startPos._jc, cast("Column", length)._jc) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT", + message_parameters={"arg_name": "startPos", "arg_type": type(startPos).__name__}, + ) + return Column(jc) + + def isin(self, *cols: Any) -> "Column": + """ + A boolean expression that is evaluated to true if the value of this + expression is contained by the evaluated values of the arguments. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols + The result will only be true at a location if any value matches in the Column. + + Returns + ------- + :class:`Column` + Column of booleans showing whether each element in the Column is contained in cols. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df[df.name.isin("Bob", "Mike")].collect() + [Row(age=5, name='Bob')] + >>> df[df.age.isin([1, 2, 3])].collect() + [Row(age=2, name='Alice')] + """ + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cast(Tuple, cols[0]) + cols = cast( + Tuple, + [c._jc if isinstance(c, Column) else _create_column_from_literal(c) for c in cols], + ) + sc = get_active_spark_context() + jc = getattr(self._jc, "isin")(_to_seq(sc, cols)) + return Column(jc) + + # order + _asc_doc = """ + Returns a sort expression based on the ascending order of the column. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.asc()).collect() + [Row(name='Alice'), Row(name='Tom')] + """ + _asc_nulls_first_doc = """ + Returns a sort expression based on ascending order of the column, and null values + return before non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), (None, 60), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.asc_nulls_first()).collect() + [Row(name=None), Row(name='Alice'), Row(name='Tom')] + + """ + _asc_nulls_last_doc = """ + Returns a sort expression based on ascending order of the column, and null values + appear after non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), (None, 60), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.asc_nulls_last()).collect() + [Row(name='Alice'), Row(name='Tom'), Row(name=None)] + + """ + _desc_doc = """ + Returns a sort expression based on the descending order of the column. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.desc()).collect() + [Row(name='Tom'), Row(name='Alice')] + """ + _desc_nulls_first_doc = """ + Returns a sort expression based on the descending order of the column, and null values + appear before non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), (None, 60), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.desc_nulls_first()).collect() + [Row(name=None), Row(name='Tom'), Row(name='Alice')] + + """ + _desc_nulls_last_doc = """ + Returns a sort expression based on the descending order of the column, and null values + appear after non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([('Tom', 80), (None, 60), ('Alice', None)], ["name", "height"]) + >>> df.select(df.name).orderBy(df.name.desc_nulls_last()).collect() + [Row(name='Tom'), Row(name='Alice'), Row(name=None)] + """ + + asc = _unary_op("asc", _asc_doc) + asc_nulls_first = _unary_op("asc_nulls_first", _asc_nulls_first_doc) + asc_nulls_last = _unary_op("asc_nulls_last", _asc_nulls_last_doc) + desc = _unary_op("desc", _desc_doc) + desc_nulls_first = _unary_op("desc_nulls_first", _desc_nulls_first_doc) + desc_nulls_last = _unary_op("desc_nulls_last", _desc_nulls_last_doc) + + _isNull_doc = """ + True if the current expression is null. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(name='Tom', height=80), Row(name='Alice', height=None)]) + >>> df.filter(df.height.isNull()).collect() + [Row(name='Alice', height=None)] + """ + _isNotNull_doc = """ + True if the current expression is NOT null. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(name='Tom', height=80), Row(name='Alice', height=None)]) + >>> df.filter(df.height.isNotNull()).collect() + [Row(name='Tom', height=80)] + """ + + isNull = _unary_op("isNull", _isNull_doc) + isNotNull = _unary_op("isNotNull", _isNotNull_doc) + + def alias(self, *alias: str, **kwargs: Any) -> "Column": + """ + Returns this column aliased with a new name or names (in the case of expressions that + return more than one column, such as explode). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + alias : str + desired column names (collects all positional arguments passed) + + Other Parameters + ---------------- + metadata: dict + a dict of information to be stored in ``metadata`` attribute of the + corresponding :class:`StructField ` (optional, keyword + only argument) + + .. versionchanged:: 2.2.0 + Added optional ``metadata`` argument. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column is aliased with new name or names. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.age.alias("age2")).collect() + [Row(age2=2), Row(age2=5)] + >>> df.select(df.age.alias("age3", metadata={'max': 99})).schema['age3'].metadata['max'] + 99 + """ + + metadata = kwargs.pop("metadata", None) + assert not kwargs, "Unexpected kwargs where passed: %s" % kwargs + + sc = get_active_spark_context() + if len(alias) == 1: + if metadata: + assert sc._jvm is not None + jmeta = sc._jvm.org.apache.spark.sql.types.Metadata.fromJson(json.dumps(metadata)) + return Column(getattr(self._jc, "as")(alias[0], jmeta)) + else: + return Column(getattr(self._jc, "as")(alias[0])) + else: + if metadata: + raise PySparkValueError( + error_class="ONLY_ALLOWED_FOR_SINGLE_COLUMN", + message_parameters={"arg_name": "metadata"}, + ) + return Column(getattr(self._jc, "as")(_to_seq(sc, list(alias)))) + + name = copy_func(alias, sinceversion=2.0, doc=":func:`name` is an alias for :func:`alias`.") + + def cast(self, dataType: Union[DataType, str]) -> "Column": + """ + Casts the column into type ``dataType``. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + dataType : :class:`DataType` or str + a DataType or Python string literal with a DDL-formatted string + to use when parsing the column to the same type. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column is cast into new type. + + Examples + -------- + >>> from pyspark.sql.types import StringType + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.age.cast("string").alias('ages')).collect() + [Row(ages='2'), Row(ages='5')] + >>> df.select(df.age.cast(StringType()).alias('ages')).collect() + [Row(ages='2'), Row(ages='5')] + """ + if isinstance(dataType, str): + jc = self._jc.cast(dataType) + elif isinstance(dataType, DataType): + from pyspark.sql import SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + jdt = spark._jsparkSession.parseDataType(dataType.json()) + jc = self._jc.cast(jdt) + else: + raise PySparkTypeError( + error_class="NOT_DATATYPE_OR_STR", + message_parameters={"arg_name": "dataType", "arg_type": type(dataType).__name__}, + ) + return Column(jc) + + astype = copy_func(cast, sinceversion=1.4, doc=":func:`astype` is an alias for :func:`cast`.") + + def between( + self, + lowerBound: Union["Column", "LiteralType", "DateTimeLiteral", "DecimalLiteral"], + upperBound: Union["Column", "LiteralType", "DateTimeLiteral", "DecimalLiteral"], + ) -> "Column": + """ + True if the current column is between the lower bound and upper bound, inclusive. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + lowerBound : :class:`Column`, int, float, string, bool, datetime, date or Decimal + a boolean expression that boundary start, inclusive. + upperBound : :class:`Column`, int, float, string, bool, datetime, date or Decimal + a boolean expression that boundary end, inclusive. + + Returns + ------- + :class:`Column` + Column of booleans showing whether each element of Column + is between left and right (inclusive). + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.name, df.age.between(2, 4)).show() + +-----+---------------------------+ + | name|((age >= 2) AND (age <= 4))| + +-----+---------------------------+ + |Alice| true| + | Bob| false| + +-----+---------------------------+ + """ + return (self >= lowerBound) & (self <= upperBound) + + def when(self, condition: "Column", value: Any) -> "Column": + """ + Evaluates a list of conditions and returns one of multiple possible result expressions. + If :func:`Column.otherwise` is not invoked, None is returned for unmatched conditions. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + condition : :class:`Column` + a boolean :class:`Column` expression. + value + a literal value, or a :class:`Column` expression. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column is in conditions. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.name, sf.when(df.age > 4, 1).when(df.age < 3, -1).otherwise(0)).show() + +-----+------------------------------------------------------------+ + | name|CASE WHEN (age > 4) THEN 1 WHEN (age < 3) THEN -1 ELSE 0 END| + +-----+------------------------------------------------------------+ + |Alice| -1| + | Bob| 1| + +-----+------------------------------------------------------------+ + + See Also + -------- + pyspark.sql.functions.when + """ + if not isinstance(condition, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + v = value._jc if isinstance(value, Column) else value + jc = self._jc.when(condition._jc, v) + return Column(jc) + + def otherwise(self, value: Any) -> "Column": + """ + Evaluates a list of conditions and returns one of multiple possible result expressions. + If :func:`Column.otherwise` is not invoked, None is returned for unmatched conditions. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + value + a literal value, or a :class:`Column` expression. + + Returns + ------- + :class:`Column` + Column representing whether each element of Column is unmatched conditions. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.select(df.name, sf.when(df.age > 3, 1).otherwise(0)).show() + +-----+-------------------------------------+ + | name|CASE WHEN (age > 3) THEN 1 ELSE 0 END| + +-----+-------------------------------------+ + |Alice| 0| + | Bob| 1| + +-----+-------------------------------------+ + + See Also + -------- + pyspark.sql.functions.when + """ + v = value._jc if isinstance(value, Column) else value + jc = self._jc.otherwise(v) + return Column(jc) + + def over(self, window: "WindowSpec") -> "Column": + """ + Define a windowing column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + window : :class:`WindowSpec` + + Returns + ------- + :class:`Column` + + Examples + -------- + >>> from pyspark.sql import Window + >>> window = ( + ... Window.partitionBy("name") + ... .orderBy("age") + ... .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ... ) + >>> from pyspark.sql.functions import rank, min, desc + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df.withColumn( + ... "rank", rank().over(window) + ... ).withColumn( + ... "min", min('age').over(window) + ... ).sort(desc("age")).show() + +---+-----+----+---+ + |age| name|rank|min| + +---+-----+----+---+ + | 5| Bob| 1| 5| + | 2|Alice| 1| 2| + +---+-----+----+---+ + """ + from pyspark.sql.window import WindowSpec + + if not isinstance(window, WindowSpec): + raise PySparkTypeError( + error_class="NOT_WINDOWSPEC", + message_parameters={"arg_name": "window", "arg_type": type(window).__name__}, + ) + jc = self._jc.over(window._jspec) + return Column(jc) + + def __nonzero__(self) -> None: + raise PySparkValueError( + error_class="CANNOT_CONVERT_COLUMN_INTO_BOOL", + message_parameters={}, + ) + + __bool__ = __nonzero__ + + def __repr__(self) -> str: + return "Column<'%s'>" % self._jc.toString() + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.column + + globs = pyspark.sql.column.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.column tests").getOrCreate() + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.column, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/conf.py new file mode 100644 index 0000000..599bb35 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/conf.py @@ -0,0 +1,99 @@ +# +# 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 sys +from typing import Any, Optional, Union + +from py4j.java_gateway import JavaObject + +from pyspark import since, _NoValue +from pyspark._globals import _NoValueType + + +class RuntimeConfig: + """User-facing configuration API, accessible through `SparkSession.conf`. + + Options set here are automatically propagated to the Hadoop configuration during I/O. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, jconf: JavaObject) -> None: + """Create a new RuntimeConfig that wraps the underlying JVM object.""" + self._jconf = jconf + + @since(2.0) + def set(self, key: str, value: Union[str, int, bool]) -> None: + """Sets the given Spark runtime configuration property.""" + self._jconf.set(key, value) + + @since(2.0) + def get( + self, key: str, default: Union[Optional[str], _NoValueType] = _NoValue + ) -> Optional[str]: + """Returns the value of Spark runtime configuration property for the given key, + assuming it is set. + """ + self._checkType(key, "key") + if default is _NoValue: + return self._jconf.get(key) + else: + if default is not None: + self._checkType(default, "default") + return self._jconf.get(key, default) + + @since(2.0) + def unset(self, key: str) -> None: + """Resets the configuration property for the given key.""" + self._jconf.unset(key) + + def _checkType(self, obj: Any, identifier: str) -> None: + """Assert that an object is of type str.""" + if not isinstance(obj, str): + raise TypeError( + "expected %s '%s' to be a string (was '%s')" % (identifier, obj, type(obj).__name__) + ) + + @since(2.4) + def isModifiable(self, key: str) -> bool: + """Indicates whether the configuration property with the given key + is modifiable in the current session. + """ + return self._jconf.isModifiable(key) + + +def _test() -> None: + import os + import doctest + from pyspark.sql.session import SparkSession + import pyspark.sql.conf + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.conf.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.conf tests").getOrCreate() + globs["sc"] = spark.sparkContext + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod(pyspark.sql.conf, globs=globs) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/__init__.py new file mode 100644 index 0000000..93236e8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +"""Spark Connect client""" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/_typing.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/_typing.py new file mode 100644 index 0000000..4c76e37 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/_typing.py @@ -0,0 +1,90 @@ +# +# 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 sys + +if sys.version_info >= (3, 8): + from typing import Protocol +else: + from typing_extensions import Protocol + +from types import FunctionType +from typing import Any, Callable, Iterable, Union, Optional, NewType +import datetime +import decimal + +import pyarrow +from pandas.core.frame import DataFrame as PandasDataFrame + +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.types import DataType +from pyspark.sql.streaming.state import GroupState + + +ColumnOrName = Union[Column, str] + +PrimitiveType = Union[bool, float, int, str] + +OptionalPrimitiveType = Optional[PrimitiveType] + +LiteralType = PrimitiveType + +DecimalLiteral = decimal.Decimal + +DateTimeLiteral = Union[datetime.datetime, datetime.date] + +DataTypeOrString = Union[DataType, str] + +DataFrameLike = PandasDataFrame + +PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]] + +ArrowMapIterFunction = Callable[[Iterable[pyarrow.RecordBatch]], Iterable[pyarrow.RecordBatch]] + +PandasGroupedMapFunction = Union[ + Callable[[DataFrameLike], DataFrameLike], + Callable[[Any, DataFrameLike], DataFrameLike], +] + +GroupedMapPandasUserDefinedFunction = NewType("GroupedMapPandasUserDefinedFunction", FunctionType) + +PandasCogroupedMapFunction = Callable[[DataFrameLike, DataFrameLike], DataFrameLike] + +PandasGroupedMapFunctionWithState = Callable[ + [Any, Iterable[DataFrameLike], GroupState], Iterable[DataFrameLike] +] + + +class UserDefinedFunctionLike(Protocol): + func: Callable[..., Any] + evalType: int + deterministic: bool + + @property + def returnType(self) -> DataType: + ... + + def __call__(self, *args: ColumnOrName) -> Column: + ... + + def asNondeterministic(self) -> "UserDefinedFunctionLike": + ... + + +class UserDefinedFunctionCallable(Protocol): + def __call__(self, *_: ColumnOrName) -> Column: + ... diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/__init__.py new file mode 100644 index 0000000..6d29d44 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +"""Spark Connect Python Client - Avro Functions""" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/functions.py new file mode 100644 index 0000000..bf019ef --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/avro/functions.py @@ -0,0 +1,107 @@ +# +# 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. +# + +""" +A collections of builtin avro functions +""" + +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import Dict, Optional, TYPE_CHECKING + +from pyspark.sql.avro import functions as PyAvroFunctions + +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.functions import _invoke_function, _to_col, _options_to_col, lit + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ColumnOrName + + +def from_avro( + data: "ColumnOrName", jsonFormatSchema: str, options: Optional[Dict[str, str]] = None +) -> Column: + if options is None: + return _invoke_function("from_avro", _to_col(data), lit(jsonFormatSchema)) + else: + return _invoke_function( + "from_avro", _to_col(data), lit(jsonFormatSchema), _options_to_col(options) + ) + + +from_avro.__doc__ = PyAvroFunctions.from_avro.__doc__ + + +def to_avro(data: "ColumnOrName", jsonFormatSchema: str = "") -> Column: + if jsonFormatSchema == "": + return _invoke_function("to_avro", _to_col(data)) + else: + return _invoke_function("to_avro", _to_col(data), lit(jsonFormatSchema)) + + +to_avro.__doc__ = PyAvroFunctions.to_avro.__doc__ + + +def _test() -> None: + import os + import sys + from pyspark.testing.utils import search_jar + + avro_jar = search_jar("connector/avro", "spark-avro", "spark-avro") + if avro_jar is None: + print( + "Skipping all Avro Python tests as the optional Avro project was " + "not compiled into a JAR. To run these tests, " + "you need to build Spark with 'build/sbt -Pavro package' or " + "'build/mvn -Pavro package' before running this test." + ) + sys.exit(0) + else: + existing_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + jars_args = "--jars %s" % avro_jar + os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) + + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.avro.functions + + globs = pyspark.sql.connect.avro.functions.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.avro.functions tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.avro.functions, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/catalog.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/catalog.py new file mode 100644 index 0000000..2a54a0d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/catalog.py @@ -0,0 +1,353 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import Any, Callable, List, Optional, TYPE_CHECKING + +import warnings +import pandas as pd + +from pyspark.storagelevel import StorageLevel +from pyspark.sql.types import StructType +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.catalog import ( + Catalog as PySparkCatalog, + CatalogMetadata, + Database, + Table, + Function, + Column, +) +from pyspark.sql.connect import plan + +if TYPE_CHECKING: + from pyspark.sql.connect.session import SparkSession + from pyspark.sql.connect._typing import DataTypeOrString, UserDefinedFunctionLike + + +class Catalog: + def __init__(self, sparkSession: "SparkSession") -> None: + self._sparkSession = sparkSession + + def _execute_and_fetch(self, catalog: plan.LogicalPlan) -> pd.DataFrame: + pdf = DataFrame.withPlan(catalog, session=self._sparkSession).toPandas() + assert pdf is not None + return pdf + + def currentCatalog(self) -> str: + pdf = self._execute_and_fetch(plan.CurrentCatalog()) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + currentCatalog.__doc__ = PySparkCatalog.currentCatalog.__doc__ + + def setCurrentCatalog(self, catalogName: str) -> None: + self._execute_and_fetch(plan.SetCurrentCatalog(catalog_name=catalogName)) + + setCurrentCatalog.__doc__ = PySparkCatalog.setCurrentCatalog.__doc__ + + def listCatalogs(self, pattern: Optional[str] = None) -> List[CatalogMetadata]: + pdf = self._execute_and_fetch(plan.ListCatalogs(pattern=pattern)) + return [ + CatalogMetadata(name=row.iloc[0], description=row.iloc[1]) for _, row in pdf.iterrows() + ] + + listCatalogs.__doc__ = PySparkCatalog.listCatalogs.__doc__ + + def currentDatabase(self) -> str: + pdf = self._execute_and_fetch(plan.CurrentDatabase()) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + currentDatabase.__doc__ = PySparkCatalog.currentDatabase.__doc__ + + def setCurrentDatabase(self, dbName: str) -> None: + self._execute_and_fetch(plan.SetCurrentDatabase(db_name=dbName)) + + setCurrentDatabase.__doc__ = PySparkCatalog.setCurrentDatabase.__doc__ + + def listDatabases(self, pattern: Optional[str] = None) -> List[Database]: + pdf = self._execute_and_fetch(plan.ListDatabases(pattern=pattern)) + return [ + Database( + name=row.iloc[0], + catalog=row.iloc[1], + description=row.iloc[2], + locationUri=row.iloc[3], + ) + for _, row in pdf.iterrows() + ] + + listDatabases.__doc__ = PySparkCatalog.listDatabases.__doc__ + + def getDatabase(self, dbName: str) -> Database: + pdf = self._execute_and_fetch(plan.GetDatabase(db_name=dbName)) + assert pdf is not None + row = pdf.iloc[0] + return Database( + name=row[0], + catalog=row[1], + description=row[2], + locationUri=row[3], + ) + + getDatabase.__doc__ = PySparkCatalog.getDatabase.__doc__ + + def databaseExists(self, dbName: str) -> bool: + pdf = self._execute_and_fetch(plan.DatabaseExists(db_name=dbName)) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + databaseExists.__doc__ = PySparkCatalog.databaseExists.__doc__ + + def listTables( + self, dbName: Optional[str] = None, pattern: Optional[str] = None + ) -> List[Table]: + pdf = self._execute_and_fetch(plan.ListTables(db_name=dbName, pattern=pattern)) + return [ + Table( + name=row.iloc[0], + catalog=row.iloc[1], + # If None, returns None. + namespace=None if row.iloc[2] is None else list(row.iloc[2]), + description=row.iloc[3], + tableType=row.iloc[4], + isTemporary=row.iloc[5], + ) + for _, row in pdf.iterrows() + ] + + listTables.__doc__ = PySparkCatalog.listTables.__doc__ + + def getTable(self, tableName: str) -> Table: + pdf = self._execute_and_fetch(plan.GetTable(table_name=tableName)) + assert pdf is not None + row = pdf.iloc[0] + return Table( + name=row.iloc[0], + catalog=row.iloc[1], + # If None, returns None. + namespace=None if row.iloc[2] is None else list(row.iloc[2]), + description=row.iloc[3], + tableType=row.iloc[4], + isTemporary=row.iloc[5], + ) + + getTable.__doc__ = PySparkCatalog.getTable.__doc__ + + def listFunctions( + self, dbName: Optional[str] = None, pattern: Optional[str] = None + ) -> List[Function]: + pdf = self._execute_and_fetch(plan.ListFunctions(db_name=dbName, pattern=pattern)) + return [ + Function( + name=row.iloc[0], + catalog=row.iloc[1], + # If None, returns None. + namespace=None if row.iloc[2] is None else list(row.iloc[2]), + description=row.iloc[3], + className=row.iloc[4], + isTemporary=row.iloc[5], + ) + for _, row in pdf.iterrows() + ] + + listFunctions.__doc__ = PySparkCatalog.listFunctions.__doc__ + + def functionExists(self, functionName: str, dbName: Optional[str] = None) -> bool: + pdf = self._execute_and_fetch( + plan.FunctionExists(function_name=functionName, db_name=dbName) + ) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + functionExists.__doc__ = PySparkCatalog.functionExists.__doc__ + + def getFunction(self, functionName: str) -> Function: + pdf = self._execute_and_fetch(plan.GetFunction(function_name=functionName)) + assert pdf is not None + row = pdf.iloc[0] + return Function( + name=row.iloc[0], + catalog=row.iloc[1], + # If None, returns None. + namespace=None if row.iloc[2] is None else list(row.iloc[2]), + description=row.iloc[3], + className=row.iloc[4], + isTemporary=row.iloc[5], + ) + + getFunction.__doc__ = PySparkCatalog.getFunction.__doc__ + + def listColumns(self, tableName: str, dbName: Optional[str] = None) -> List[Column]: + pdf = self._execute_and_fetch(plan.ListColumns(table_name=tableName, db_name=dbName)) + return [ + Column( + name=row.iloc[0], + description=row.iloc[1], + dataType=row.iloc[2], + nullable=row.iloc[3], + isPartition=row.iloc[4], + isBucket=row.iloc[5], + ) + for _, row in pdf.iterrows() + ] + + listColumns.__doc__ = PySparkCatalog.listColumns.__doc__ + + def tableExists(self, tableName: str, dbName: Optional[str] = None) -> bool: + pdf = self._execute_and_fetch(plan.TableExists(table_name=tableName, db_name=dbName)) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + tableExists.__doc__ = PySparkCatalog.tableExists.__doc__ + + def createExternalTable( + self, + tableName: str, + path: Optional[str] = None, + source: Optional[str] = None, + schema: Optional[StructType] = None, + **options: str, + ) -> DataFrame: + catalog = plan.CreateExternalTable( + table_name=tableName, + path=path, # type: ignore[arg-type] + source=source, + schema=schema, + options=options, + ) + df = DataFrame.withPlan(catalog, session=self._sparkSession) + df.toPandas() # Eager execution. + return df + + createExternalTable.__doc__ = PySparkCatalog.createExternalTable.__doc__ + + def createTable( + self, + tableName: str, + path: Optional[str] = None, + source: Optional[str] = None, + schema: Optional[StructType] = None, + description: Optional[str] = None, + **options: str, + ) -> DataFrame: + catalog = plan.CreateTable( + table_name=tableName, + path=path, # type: ignore[arg-type] + source=source, + schema=schema, + description=description, + options=options, + ) + df = DataFrame.withPlan(catalog, session=self._sparkSession) + df.toPandas() # Eager execution. + return df + + createTable.__doc__ = PySparkCatalog.createTable.__doc__ + + def dropTempView(self, viewName: str) -> bool: + pdf = self._execute_and_fetch(plan.DropTempView(view_name=viewName)) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + dropTempView.__doc__ = PySparkCatalog.dropTempView.__doc__ + + def dropGlobalTempView(self, viewName: str) -> bool: + pdf = self._execute_and_fetch(plan.DropGlobalTempView(view_name=viewName)) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + dropGlobalTempView.__doc__ = PySparkCatalog.dropGlobalTempView.__doc__ + + def isCached(self, tableName: str) -> bool: + pdf = self._execute_and_fetch(plan.IsCached(table_name=tableName)) + assert pdf is not None + return pdf.iloc[0].iloc[0] + + isCached.__doc__ = PySparkCatalog.isCached.__doc__ + + def cacheTable(self, tableName: str, storageLevel: Optional[StorageLevel] = None) -> None: + self._execute_and_fetch(plan.CacheTable(table_name=tableName, storage_level=storageLevel)) + + cacheTable.__doc__ = PySparkCatalog.cacheTable.__doc__ + + def uncacheTable(self, tableName: str) -> None: + self._execute_and_fetch(plan.UncacheTable(table_name=tableName)) + + uncacheTable.__doc__ = PySparkCatalog.uncacheTable.__doc__ + + def clearCache(self) -> None: + self._execute_and_fetch(plan.ClearCache()) + + clearCache.__doc__ = PySparkCatalog.clearCache.__doc__ + + def refreshTable(self, tableName: str) -> None: + self._execute_and_fetch(plan.RefreshTable(table_name=tableName)) + + refreshTable.__doc__ = PySparkCatalog.refreshTable.__doc__ + + def recoverPartitions(self, tableName: str) -> None: + self._execute_and_fetch(plan.RecoverPartitions(table_name=tableName)) + + recoverPartitions.__doc__ = PySparkCatalog.recoverPartitions.__doc__ + + def refreshByPath(self, path: str) -> None: + self._execute_and_fetch(plan.RefreshByPath(path=path)) + + refreshByPath.__doc__ = PySparkCatalog.refreshByPath.__doc__ + + def registerFunction( + self, name: str, f: Callable[..., Any], returnType: Optional["DataTypeOrString"] = None + ) -> "UserDefinedFunctionLike": + warnings.warn("Deprecated in 2.3.0. Use spark.udf.register instead.", FutureWarning) + return self._sparkSession.udf.register(name, f, returnType) + + registerFunction.__doc__ = PySparkCatalog.registerFunction.__doc__ + + +Catalog.__doc__ = PySparkCatalog.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.catalog + + globs = pyspark.sql.connect.catalog.__dict__.copy() + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.catalog tests").remote("local[4]").getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.catalog, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/__init__.py new file mode 100644 index 0000000..469d1c5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/__init__.py @@ -0,0 +1,22 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from pyspark.sql.connect.client.core import * # noqa: F401,F403 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/artifact.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/artifact.py new file mode 100644 index 0000000..c858768 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/artifact.py @@ -0,0 +1,412 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import hashlib +import importlib +import io +import sys +import os +import zlib +from itertools import chain +from typing import List, Iterable, BinaryIO, Iterator, Optional, Tuple +import abc +from pathlib import Path +from urllib.parse import urlparse +from urllib.request import url2pathname +from functools import cached_property + +import grpc + +import pyspark.sql.connect.proto as proto +import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib + + +JAR_PREFIX: str = "jars" +PYFILE_PREFIX: str = "pyfiles" +ARCHIVE_PREFIX: str = "archives" +FILE_PREFIX: str = "files" +FORWARD_TO_FS_PREFIX: str = "forward_to_fs" +CACHE_PREFIX: str = "cache" + + +class LocalData(metaclass=abc.ABCMeta): + """ + Payload stored on this machine. + """ + + @cached_property + @abc.abstractmethod + def stream(self) -> BinaryIO: + pass + + @cached_property + @abc.abstractmethod + def size(self) -> int: + pass + + +class LocalFile(LocalData): + """ + Payload stored in a local file. + """ + + def __init__(self, path: str): + self.path = path + self._size: int + self._stream: int + + @cached_property + def size(self) -> int: + return os.path.getsize(self.path) + + @cached_property + def stream(self) -> BinaryIO: + return open(self.path, "rb") + + +class InMemory(LocalData): + """ + Payload stored in memory. + """ + + def __init__(self, blob: bytes): + self.blob = blob + self._size: int + self._stream: int + + @cached_property + def size(self) -> int: + return len(self.blob) + + @cached_property + def stream(self) -> BinaryIO: + return io.BytesIO(self.blob) + + +class Artifact: + def __init__(self, path: str, storage: LocalData): + assert not Path(path).is_absolute(), f"Bad path: {path}" + self.path = path + self.storage = storage + + @cached_property + def size(self) -> int: + if isinstance(self.storage, LocalData): + return self.storage.size + else: + raise RuntimeError(f"Unsupported storage {type(self.storage)}") + + +def new_jar_artifact(file_name: str, storage: LocalData) -> Artifact: + return _new_artifact(JAR_PREFIX, ".jar", file_name, storage) + + +def new_pyfile_artifact(file_name: str, storage: LocalData) -> Artifact: + assert any(file_name.endswith(s) for s in (".py", ".zip", ".egg", ".jar")) + return _new_artifact(PYFILE_PREFIX, "", file_name, storage) + + +def new_archive_artifact(file_name: str, storage: LocalData) -> Artifact: + return _new_artifact(ARCHIVE_PREFIX, "", file_name, storage) + + +def new_file_artifact(file_name: str, storage: LocalData) -> Artifact: + return _new_artifact(FILE_PREFIX, "", file_name, storage) + + +def new_cache_artifact(id: str, storage: LocalData) -> Artifact: + return _new_artifact(CACHE_PREFIX, "", id, storage) + + +def _new_artifact( + prefix: str, required_suffix: str, file_name: str, storage: LocalData +) -> Artifact: + assert not Path(file_name).is_absolute() + assert file_name.endswith(required_suffix) + return Artifact(os.path.join(prefix, file_name), storage) + + +class ArtifactManager: + """ + The Artifact Manager is responsible for handling and transferring artifacts from the local + client to the server (local/remote). + + Parameters + ---------- + user_id : str, optional + User ID. + session_id: str + An unique identifier of the session which the artifact manager belongs to. + channel: grpc.Channel + GRPC Channel instance. + """ + + # Using the midpoint recommendation of 32KiB for chunk size as specified in + # https://github.com/grpc/grpc.github.io/issues/371. + CHUNK_SIZE: int = 32 * 1024 + + def __init__( + self, + user_id: Optional[str], + session_id: str, + channel: grpc.Channel, + metadata: Iterable[Tuple[str, str]], + ): + self._user_context = proto.UserContext() + if user_id is not None: + self._user_context.user_id = user_id + self._stub = grpc_lib.SparkConnectServiceStub(channel) + self._session_id = session_id + self._metadata = metadata + + def _parse_artifacts( + self, path_or_uri: str, pyfile: bool, archive: bool, file: bool + ) -> List[Artifact]: + # Currently only local files with .jar extension is supported. + parsed = urlparse(path_or_uri) + # Check if it is a file from the scheme + if parsed.scheme == "": + # Similar with Utils.resolveURI. + fragment = parsed.fragment + parsed = urlparse(Path(url2pathname(parsed.path)).absolute().as_uri()) + parsed = parsed._replace(fragment=fragment) + + if parsed.scheme == "file": + local_path = url2pathname(parsed.path) + name = Path(local_path).name + if pyfile and name.endswith(".py"): + artifact = new_pyfile_artifact(name, LocalFile(local_path)) + importlib.invalidate_caches() + elif pyfile and ( + name.endswith(".zip") or name.endswith(".egg") or name.endswith(".jar") + ): + sys.path.insert(1, local_path) + artifact = new_pyfile_artifact(name, LocalFile(local_path)) + importlib.invalidate_caches() + elif archive and ( + name.endswith(".zip") + or name.endswith(".jar") + or name.endswith(".tar.gz") + or name.endswith(".tgz") + or name.endswith(".tar") + ): + assert any(name.endswith(s) for s in (".zip", ".jar", ".tar.gz", ".tgz", ".tar")) + + if parsed.fragment != "": + # Minimal fix for the workaround of fragment handling in URI. + # This has a limitation - hash(#) in the file name would not work. + if "#" in local_path: + raise ValueError("'#' in the path is not supported for adding an archive.") + name = f"{name}#{parsed.fragment}" + + artifact = new_archive_artifact(name, LocalFile(local_path)) + elif file: + artifact = new_file_artifact(name, LocalFile(local_path)) + elif name.endswith(".jar"): + artifact = new_jar_artifact(name, LocalFile(local_path)) + else: + raise RuntimeError(f"Unsupported file format: {local_path}") + return [artifact] + raise RuntimeError(f"Unsupported scheme: {parsed.scheme}") + + def _parse_forward_to_fs_artifacts(self, local_path: str, dest_path: str) -> List[Artifact]: + abs_path: Path = Path(local_path).absolute() + # TODO: Support directory path. + assert abs_path.is_file(), "local path must be a file path." + storage = LocalFile(str(abs_path)) + + assert Path(dest_path).is_absolute(), "destination FS path must be an absolute path." + + # The `dest_path` is an absolute path, to add the FORWARD_TO_FS_PREFIX, + # we cannot use `os.path.join` + artifact_path = FORWARD_TO_FS_PREFIX + dest_path + return [Artifact(artifact_path, storage)] + + def _create_requests( + self, *path: str, pyfile: bool, archive: bool, file: bool + ) -> Iterator[proto.AddArtifactsRequest]: + """Separated for the testing purpose.""" + return self._add_artifacts( + chain( + *(self._parse_artifacts(p, pyfile=pyfile, archive=archive, file=file) for p in path) + ) + ) + + def _retrieve_responses( + self, requests: Iterator[proto.AddArtifactsRequest] + ) -> proto.AddArtifactsResponse: + """Separated for the testing purpose.""" + return self._stub.AddArtifacts(requests, metadata=self._metadata) + + def _request_add_artifacts(self, requests: Iterator[proto.AddArtifactsRequest]) -> None: + response: proto.AddArtifactsResponse = self._retrieve_responses(requests) + summaries: List[proto.AddArtifactsResponse.ArtifactSummary] = [] + + for summary in response.artifacts: + summaries.append(summary) + # TODO(SPARK-42658): Handle responses containing CRC failures. + + def add_artifacts(self, *path: str, pyfile: bool, archive: bool, file: bool) -> None: + """ + Add a single artifact to the session. + Currently only local files with .jar extension is supported. + """ + requests: Iterator[proto.AddArtifactsRequest] = self._create_requests( + *path, pyfile=pyfile, archive=archive, file=file + ) + self._request_add_artifacts(requests) + + def _add_forward_to_fs_artifacts(self, local_path: str, dest_path: str) -> None: + requests: Iterator[proto.AddArtifactsRequest] = self._add_artifacts( + self._parse_forward_to_fs_artifacts(local_path, dest_path) + ) + self._request_add_artifacts(requests) + + def _add_artifacts(self, artifacts: Iterable[Artifact]) -> Iterator[proto.AddArtifactsRequest]: + """ + Add a number of artifacts to the session. + """ + + current_batch: List[Artifact] = [] + current_batch_size = 0 + + def add_to_batch(dep: Artifact, size: int) -> None: + nonlocal current_batch + nonlocal current_batch_size + + current_batch.append(dep) + current_batch_size += size + + def write_batch() -> Iterator[proto.AddArtifactsRequest]: + nonlocal current_batch + nonlocal current_batch_size + + yield from self._add_batched_artifacts(current_batch) + current_batch = [] + current_batch_size = 0 + + for artifact in artifacts: + data = artifact.storage + size = data.size + if size > ArtifactManager.CHUNK_SIZE: + # Payload can either be a batch OR a single chunked artifact. + # Write batch if non-empty before chunking current artifact. + if len(current_batch) > 0: + yield from write_batch() + yield from self._add_chunked_artifact(artifact) + else: + if current_batch_size + size > ArtifactManager.CHUNK_SIZE: + yield from write_batch() + add_to_batch(artifact, size) + + if len(current_batch) > 0: + yield from write_batch() + + def _add_batched_artifacts( + self, artifacts: Iterable[Artifact] + ) -> Iterator[proto.AddArtifactsRequest]: + """ + Add a batch of artifacts to the stream. All the artifacts in this call are packaged into a + single :class:`proto.AddArtifactsRequest`. + """ + artifact_chunks = [] + + for artifact in artifacts: + binary = artifact.storage.stream.read() + crc32 = zlib.crc32(binary) + data = proto.AddArtifactsRequest.ArtifactChunk(data=binary, crc=crc32) + artifact_chunks.append( + proto.AddArtifactsRequest.SingleChunkArtifact(name=artifact.path, data=data) + ) + + # Write the request once + yield proto.AddArtifactsRequest( + session_id=self._session_id, + user_context=self._user_context, + batch=proto.AddArtifactsRequest.Batch(artifacts=artifact_chunks), + ) + + def _add_chunked_artifact(self, artifact: Artifact) -> Iterator[proto.AddArtifactsRequest]: + """ + Add a artifact in chunks to the stream. The artifact's data is spread out over multiple + :class:`proto.AddArtifactsRequest requests`. + """ + initial_batch = True + # Integer division that rounds up to the nearest whole number. + get_num_chunks = int( + (artifact.size + (ArtifactManager.CHUNK_SIZE - 1)) / ArtifactManager.CHUNK_SIZE + ) + + # Consume stream in chunks until there is no data left to read. + for chunk in iter(lambda: artifact.storage.stream.read(ArtifactManager.CHUNK_SIZE), b""): + if initial_batch: + # First RPC contains the `BeginChunkedArtifact` payload (`begin_chunk`). + yield proto.AddArtifactsRequest( + session_id=self._session_id, + user_context=self._user_context, + begin_chunk=proto.AddArtifactsRequest.BeginChunkedArtifact( + name=artifact.path, + total_bytes=artifact.size, + num_chunks=get_num_chunks, + initial_chunk=proto.AddArtifactsRequest.ArtifactChunk( + data=chunk, crc=zlib.crc32(chunk) + ), + ), + ) + initial_batch = False + else: + # Subsequent RPCs contains the `ArtifactChunk` payload (`chunk`). + yield proto.AddArtifactsRequest( + session_id=self._session_id, + user_context=self._user_context, + chunk=proto.AddArtifactsRequest.ArtifactChunk( + data=chunk, crc=zlib.crc32(chunk) + ), + ) + + def is_cached_artifact(self, hash: str) -> bool: + """ + Ask the server either any artifact with `hash` has been cached at the server side or not. + """ + artifactName = CACHE_PREFIX + "/" + hash + request = proto.ArtifactStatusesRequest( + user_context=self._user_context, session_id=self._session_id, names=[artifactName] + ) + resp: proto.ArtifactStatusesResponse = self._stub.ArtifactStatus( + request, metadata=self._metadata + ) + status = resp.statuses.get(artifactName) + return status.exists if status is not None else False + + def cache_artifact(self, blob: bytes) -> str: + """ + Cache the give blob at the session. + """ + hash = hashlib.sha256(blob).hexdigest() + if not self.is_cached_artifact(hash): + requests = self._add_artifacts([new_cache_artifact(hash, InMemory(blob))]) + response: proto.AddArtifactsResponse = self._retrieve_responses(requests) + summaries: List[proto.AddArtifactsResponse.ArtifactSummary] = [] + + for summary in response.artifacts: + summaries.append(summary) + # TODO(SPARK-42658): Handle responses containing CRC failures. + + return hash diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/core.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/core.py new file mode 100644 index 0000000..7b3299d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/core.py @@ -0,0 +1,1688 @@ +# +# 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. +# +__all__ = [ + "ChannelBuilder", + "SparkConnectClient", + "getLogLevel", +] + +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import threading +import logging +import os +import platform +import random +import time +import urllib.parse +import uuid +import sys +from types import TracebackType +from typing import ( + Iterable, + Iterator, + Optional, + Any, + Union, + List, + Tuple, + Dict, + Set, + NoReturn, + cast, + Callable, + Generator, + Type, + TYPE_CHECKING, + Sequence, +) + +import pandas as pd +import pyarrow as pa + +import google.protobuf.message +from grpc_status import rpc_status +import grpc +from google.protobuf import text_format +from google.rpc import error_details_pb2 + +from pyspark.version import __version__ +from pyspark.resource.information import ResourceInformation +from pyspark.sql.connect.client.artifact import ArtifactManager +from pyspark.sql.connect.client.reattach import ( + ExecutePlanResponseReattachableIterator, + RetryException, +) +from pyspark.sql.connect.conversion import storage_level_to_proto, proto_to_storage_level +import pyspark.sql.connect.proto as pb2 +import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib +import pyspark.sql.connect.types as types +from pyspark.errors.exceptions.connect import ( + convert_exception, + SparkConnectException, + SparkConnectGrpcException, +) +from pyspark.sql.connect.expressions import ( + PythonUDF, + CommonInlineUserDefinedFunction, + JavaUDF, +) +from pyspark.sql.connect.plan import ( + CommonInlineUserDefinedTableFunction, + PythonUDTF, +) +from pyspark.sql.connect.utils import get_python_ver +from pyspark.sql.pandas.types import _create_converter_to_pandas, from_arrow_schema +from pyspark.sql.types import DataType, StructType, TimestampType, _has_type +from pyspark.rdd import PythonEvalType +from pyspark.storagelevel import StorageLevel +from pyspark.errors import PySparkValueError + + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import DataTypeOrString + + +def _configure_logging() -> logging.Logger: + """Configure logging for the Spark Connect clients.""" + logger = logging.getLogger(__name__) + handler = logging.StreamHandler() + handler.setFormatter( + logging.Formatter(fmt="%(asctime)s %(process)d %(levelname)s %(funcName)s %(message)s") + ) + logger.addHandler(handler) + + # Check the environment variables for log levels: + if "SPARK_CONNECT_LOG_LEVEL" in os.environ: + logger.setLevel(os.environ["SPARK_CONNECT_LOG_LEVEL"].upper()) + else: + logger.disabled = True + return logger + + +# Instantiate the logger based on the environment configuration. +logger = _configure_logging() + + +def getLogLevel() -> Optional[int]: + """ + This returns this log level as integer, or none (if no logging is enabled). + + Spark Connect logging can be configured with environment variable 'SPARK_CONNECT_LOG_LEVEL' + + .. versionadded:: 3.5.0 + """ + + if not logger.disabled: + return logger.level + return None + + +class ChannelBuilder: + """ + This is a helper class that is used to create a GRPC channel based on the given + connection string per the documentation of Spark Connect. + + .. versionadded:: 3.4.0 + + Examples + -------- + >>> cb = ChannelBuilder("sc://localhost") + ... cb.endpoint + "localhost:15002" + + >>> cb = ChannelBuilder("sc://localhost/;use_ssl=true;token=aaa") + ... cb.secure + True + """ + + PARAM_USE_SSL = "use_ssl" + PARAM_TOKEN = "token" + PARAM_USER_ID = "user_id" + PARAM_USER_AGENT = "user_agent" + PARAM_SESSION_ID = "session_id" + MAX_MESSAGE_LENGTH = 128 * 1024 * 1024 + + @staticmethod + def default_port() -> int: + if "SPARK_TESTING" in os.environ: + from pyspark.sql.session import SparkSession as PySparkSession + + # In the case when Spark Connect uses the local mode, it starts the regular Spark + # session that starts Spark Connect server that sets `SparkSession._instantiatedSession` + # via SparkSession.__init__. + # + # We are getting the actual server port from the Spark session via Py4J to address + # the case when the server port is set to 0 (in which allocates an ephemeral port). + # + # This is only used in the test/development mode. + session = PySparkSession._instantiatedSession + + # 'spark.local.connect' is set when we use the local mode in Spark Connect. + if session is not None and session.conf.get("spark.local.connect", "0") == "1": + + jvm = PySparkSession._instantiatedSession._jvm # type: ignore[union-attr] + return getattr( + getattr( + jvm.org.apache.spark.sql.connect.service, # type: ignore[union-attr] + "SparkConnectService$", + ), + "MODULE$", + ).localPort() + return 15002 + + def __init__(self, url: str, channelOptions: Optional[List[Tuple[str, Any]]] = None) -> None: + """ + Constructs a new channel builder. This is used to create the proper GRPC channel from + the connection string. + + Parameters + ---------- + url : str + Spark Connect connection string + channelOptions: list of tuple, optional + Additional options that can be passed to the GRPC channel construction. + """ + # Explicitly check the scheme of the URL. + if url[:5] != "sc://": + raise PySparkValueError( + error_class="INVALID_CONNECT_URL", + message_parameters={ + "detail": "The URL must start with 'sc://'. Please update the URL to " + "follow the correct format, e.g., 'sc://hostname:port'.", + }, + ) + # Rewrite the URL to use http as the scheme so that we can leverage + # Python's built-in parser. + tmp_url = "http" + url[2:] + self.url = urllib.parse.urlparse(tmp_url) + self.params: Dict[str, str] = {} + if len(self.url.path) > 0 and self.url.path != "/": + raise PySparkValueError( + error_class="INVALID_CONNECT_URL", + message_parameters={ + "detail": f"The path component '{self.url.path}' must be empty. Please update " + f"the URL to follow the correct format, e.g., 'sc://hostname:port'.", + }, + ) + self._extract_attributes() + + GRPC_DEFAULT_OPTIONS = [ + ("grpc.max_send_message_length", ChannelBuilder.MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", ChannelBuilder.MAX_MESSAGE_LENGTH), + ] + + if channelOptions is None: + self._channel_options = GRPC_DEFAULT_OPTIONS + else: + self._channel_options = GRPC_DEFAULT_OPTIONS + channelOptions + + def _extract_attributes(self) -> None: + if len(self.url.params) > 0: + parts = self.url.params.split(";") + for p in parts: + kv = p.split("=") + if len(kv) != 2: + raise PySparkValueError( + error_class="INVALID_CONNECT_URL", + message_parameters={ + "detail": f"Parameter '{p}' should be provided as a " + f"key-value pair separated by an equal sign (=). Please update " + f"the parameter to follow the correct format, e.g., 'key=value'.", + }, + ) + self.params[kv[0]] = urllib.parse.unquote(kv[1]) + + netloc = self.url.netloc.split(":") + if len(netloc) == 1: + self.host = netloc[0] + self.port = ChannelBuilder.default_port() + elif len(netloc) == 2: + self.host = netloc[0] + self.port = int(netloc[1]) + else: + raise PySparkValueError( + error_class="INVALID_CONNECT_URL", + message_parameters={ + "detail": f"Target destination '{self.url.netloc}' should match the " + f"':' pattern. Please update the destination to follow " + f"the correct format, e.g., 'hostname:port'.", + }, + ) + + def metadata(self) -> Iterable[Tuple[str, str]]: + """ + Builds the GRPC specific metadata list to be injected into the request. All + parameters will be converted to metadata except ones that are explicitly used + by the channel. + + Returns + ------- + A list of tuples (key, value) + """ + return [ + (k, self.params[k]) + for k in self.params + if k + not in [ + ChannelBuilder.PARAM_TOKEN, + ChannelBuilder.PARAM_USE_SSL, + ChannelBuilder.PARAM_USER_ID, + ChannelBuilder.PARAM_USER_AGENT, + ChannelBuilder.PARAM_SESSION_ID, + ] + ] + + @property + def secure(self) -> bool: + if self._token is not None: + return True + + value = self.params.get(ChannelBuilder.PARAM_USE_SSL, "") + return value.lower() == "true" + + @property + def endpoint(self) -> str: + return f"{self.host}:{self.port}" + + @property + def _token(self) -> Optional[str]: + return self.params.get(ChannelBuilder.PARAM_TOKEN, None) + + @property + def userId(self) -> Optional[str]: + """ + Returns + ------- + The user_id extracted from the parameters of the connection string or `None` if not + specified. + """ + return self.params.get(ChannelBuilder.PARAM_USER_ID, None) + + @property + def userAgent(self) -> str: + """ + Returns + ------- + user_agent : str + The user_agent parameter specified in the connection string, + or "_SPARK_CONNECT_PYTHON" when not specified. + The returned value will be percent encoded. + """ + user_agent = self.params.get( + ChannelBuilder.PARAM_USER_AGENT, + os.getenv("SPARK_CONNECT_USER_AGENT", "_SPARK_CONNECT_PYTHON"), + ) + ua_len = len(urllib.parse.quote(user_agent)) + if ua_len > 2048: + raise SparkConnectException( + f"'user_agent' parameter should not exceed 2048 characters, found {len} characters." + ) + return " ".join( + [ + user_agent, + f"spark/{__version__}", + f"os/{platform.uname().system.lower()}", + f"python/{platform.python_version()}", + ] + ) + + def get(self, key: str) -> Any: + """ + Parameters + ---------- + key : str + Parameter key name. + + Returns + ------- + The parameter value if present, raises exception otherwise. + """ + return self.params[key] + + @property + def session_id(self) -> Optional[str]: + """ + Returns + ------- + The session_id extracted from the parameters of the connection string or `None` if not + specified. + """ + session_id = self.params.get(ChannelBuilder.PARAM_SESSION_ID, None) + if session_id is not None: + try: + uuid.UUID(session_id, version=4) + except ValueError as ve: + raise ValueError("Parameter value 'session_id' must be a valid UUID format.", ve) + return session_id + + def toChannel(self) -> grpc.Channel: + """ + Applies the parameters of the connection string and creates a new + GRPC channel according to the configuration. Passes optional channel options to + construct the channel. + + Returns + ------- + GRPC Channel instance. + """ + destination = f"{self.host}:{self.port}" + + # Setting a token implicitly sets the `use_ssl` to True. + if not self.secure and self._token is not None: + use_secure = True + elif self.secure: + use_secure = True + else: + use_secure = False + + if not use_secure: + return grpc.insecure_channel(destination, options=self._channel_options) + else: + # Default SSL Credentials. + opt_token = self.params.get(ChannelBuilder.PARAM_TOKEN, None) + # When a token is present, pass the token to the channel. + if opt_token is not None: + ssl_creds = grpc.ssl_channel_credentials() + composite_creds = grpc.composite_channel_credentials( + ssl_creds, grpc.access_token_call_credentials(opt_token) + ) + return grpc.secure_channel( + destination, credentials=composite_creds, options=self._channel_options + ) + else: + return grpc.secure_channel( + destination, + credentials=grpc.ssl_channel_credentials(), + options=self._channel_options, + ) + + +class MetricValue: + def __init__(self, name: str, value: Union[int, float], type: str): + self._name = name + self._type = type + self._value = value + + def __repr__(self) -> str: + return f"<{self._name}={self._value} ({self._type})>" + + @property + def name(self) -> str: + return self._name + + @property + def value(self) -> Union[int, float]: + return self._value + + @property + def metric_type(self) -> str: + return self._type + + +class PlanMetrics: + def __init__(self, name: str, id: int, parent: int, metrics: List[MetricValue]): + self._name = name + self._id = id + self._parent_id = parent + self._metrics = metrics + + def __repr__(self) -> str: + return f"Plan({self._name})={self._metrics}" + + @property + def name(self) -> str: + return self._name + + @property + def plan_id(self) -> int: + return self._id + + @property + def parent_plan_id(self) -> int: + return self._parent_id + + @property + def metrics(self) -> List[MetricValue]: + return self._metrics + + +class PlanObservedMetrics: + def __init__(self, name: str, metrics: List[pb2.Expression.Literal]): + self._name = name + self._metrics = metrics + + def __repr__(self) -> str: + return f"Plan observed({self._name}={self._metrics})" + + @property + def name(self) -> str: + return self._name + + @property + def metrics(self) -> List[pb2.Expression.Literal]: + return self._metrics + + +class AnalyzeResult: + def __init__( + self, + schema: Optional[DataType], + explain_string: Optional[str], + tree_string: Optional[str], + is_local: Optional[bool], + is_streaming: Optional[bool], + input_files: Optional[List[str]], + spark_version: Optional[str], + parsed: Optional[DataType], + is_same_semantics: Optional[bool], + semantic_hash: Optional[int], + storage_level: Optional[StorageLevel], + ): + self.schema = schema + self.explain_string = explain_string + self.tree_string = tree_string + self.is_local = is_local + self.is_streaming = is_streaming + self.input_files = input_files + self.spark_version = spark_version + self.parsed = parsed + self.is_same_semantics = is_same_semantics + self.semantic_hash = semantic_hash + self.storage_level = storage_level + + @classmethod + def fromProto(cls, pb: Any) -> "AnalyzeResult": + schema: Optional[DataType] = None + explain_string: Optional[str] = None + tree_string: Optional[str] = None + is_local: Optional[bool] = None + is_streaming: Optional[bool] = None + input_files: Optional[List[str]] = None + spark_version: Optional[str] = None + parsed: Optional[DataType] = None + is_same_semantics: Optional[bool] = None + semantic_hash: Optional[int] = None + storage_level: Optional[StorageLevel] = None + + if pb.HasField("schema"): + schema = types.proto_schema_to_pyspark_data_type(pb.schema.schema) + elif pb.HasField("explain"): + explain_string = pb.explain.explain_string + elif pb.HasField("tree_string"): + tree_string = pb.tree_string.tree_string + elif pb.HasField("is_local"): + is_local = pb.is_local.is_local + elif pb.HasField("is_streaming"): + is_streaming = pb.is_streaming.is_streaming + elif pb.HasField("input_files"): + input_files = pb.input_files.files + elif pb.HasField("spark_version"): + spark_version = pb.spark_version.version + elif pb.HasField("ddl_parse"): + parsed = types.proto_schema_to_pyspark_data_type(pb.ddl_parse.parsed) + elif pb.HasField("same_semantics"): + is_same_semantics = pb.same_semantics.result + elif pb.HasField("semantic_hash"): + semantic_hash = pb.semantic_hash.result + elif pb.HasField("persist"): + pass + elif pb.HasField("unpersist"): + pass + elif pb.HasField("get_storage_level"): + storage_level = proto_to_storage_level(pb.get_storage_level.storage_level) + else: + raise SparkConnectException("No analyze result found!") + + return AnalyzeResult( + schema, + explain_string, + tree_string, + is_local, + is_streaming, + input_files, + spark_version, + parsed, + is_same_semantics, + semantic_hash, + storage_level, + ) + + +class ConfigResult: + def __init__(self, pairs: List[Tuple[str, Optional[str]]], warnings: List[str]): + self.pairs = pairs + self.warnings = warnings + + @classmethod + def fromProto(cls, pb: pb2.ConfigResponse) -> "ConfigResult": + return ConfigResult( + pairs=[(pair.key, pair.value if pair.HasField("value") else None) for pair in pb.pairs], + warnings=list(pb.warnings), + ) + + +class SparkConnectClient(object): + """ + Conceptually the remote spark session that communicates with the server + """ + + @classmethod + def retry_exception(cls, e: Exception) -> bool: + """ + Helper function that is used to identify if an exception thrown by the server + can be retried or not. + + Parameters + ---------- + e : Exception + The GRPC error as received from the server. Typed as Exception, because other exception + thrown during client processing can be passed here as well. + + Returns + ------- + True if the exception can be retried, False otherwise. + + """ + if not isinstance(e, grpc.RpcError): + return False + + if e.code() in [grpc.StatusCode.INTERNAL]: + msg = str(e) + + # This error happens if another RPC preempts this RPC. + if "INVALID_CURSOR.DISCONNECTED" in msg: + return True + + if e.code() == grpc.StatusCode.UNAVAILABLE: + return True + + return False + + def __init__( + self, + connection: Union[str, ChannelBuilder], + user_id: Optional[str] = None, + channel_options: Optional[List[Tuple[str, Any]]] = None, + retry_policy: Optional[Dict[str, Any]] = None, + use_reattachable_execute: bool = True, + ): + """ + Creates a new SparkSession for the Spark Connect interface. + + Parameters + ---------- + connection : str or :class:`ChannelBuilder` + Connection string that is used to extract the connection parameters and configure + the GRPC connection. Or instance of ChannelBuilder that creates GRPC connection. + Defaults to `sc://localhost`. + user_id : str, optional + Optional unique user ID that is used to differentiate multiple users and + isolate their Spark Sessions. If the `user_id` is not set, will default to + the $USER environment. Defining the user ID as part of the connection string + takes precedence. + channel_options: list of tuple, optional + Additional options that can be passed to the GRPC channel construction. + retry_policy: dict of str and any, optional + Additional configuration for retrying. There are four configurations as below + * ``max_retries`` + Maximum number of tries default 15 + * ``backoff_multiplier`` + Backoff multiplier for the policy. Default: 4(ms) + * ``initial_backoff`` + Backoff to wait before the first retry. Default: 50(ms) + * ``max_backoff`` + Maximum backoff controls the maximum amount of time to wait before retrying + a failed request. Default: 60000(ms). + use_reattachable_execute: bool + Enable reattachable execution. + """ + self.thread_local = threading.local() + + # Parse the connection string. + self._builder = ( + connection + if isinstance(connection, ChannelBuilder) + else ChannelBuilder(connection, channel_options) + ) + self._user_id = None + self._retry_policy = { + # Please synchronize changes here with Scala side + # GrpcRetryHandler.scala + # + # Note: the number of retries is selected so that the maximum tolerated wait + # is guaranteed to be at least 10 minutes + "max_retries": 15, + "backoff_multiplier": 4.0, + "initial_backoff": 50, + "max_backoff": 60000, + "jitter": 500, + "min_jitter_threshold": 2000, + } + if retry_policy: + self._retry_policy.update(retry_policy) + + if self._builder.session_id is None: + # Generate a unique session ID for this client. This UUID must be unique to allow + # concurrent Spark sessions of the same user. If the channel is closed, creating + # a new client will create a new session ID. + self._session_id = str(uuid.uuid4()) + else: + # Use the pre-defined session ID. + self._session_id = str(self._builder.session_id) + + if self._builder.userId is not None: + self._user_id = self._builder.userId + elif user_id is not None: + self._user_id = user_id + else: + self._user_id = os.getenv("USER", None) + + self._channel = self._builder.toChannel() + self._closed = False + self._stub = grpc_lib.SparkConnectServiceStub(self._channel) + self._artifact_manager = ArtifactManager( + self._user_id, self._session_id, self._channel, self._builder.metadata() + ) + self._use_reattachable_execute = use_reattachable_execute + # Configure logging for the SparkConnect client. + + def _retrying(self) -> "Retrying": + return Retrying( + can_retry=SparkConnectClient.retry_exception, **self._retry_policy # type: ignore + ) + + def disable_reattachable_execute(self) -> "SparkConnectClient": + self._use_reattachable_execute = False + return self + + def enable_reattachable_execute(self) -> "SparkConnectClient": + self._use_reattachable_execute = True + return self + + def register_udf( + self, + function: Any, + return_type: "DataTypeOrString", + name: Optional[str] = None, + eval_type: int = PythonEvalType.SQL_BATCHED_UDF, + deterministic: bool = True, + ) -> str: + """ + Create a temporary UDF in the session catalog on the other side. We generate a + temporary name for it. + """ + + if name is None: + name = f"fun_{uuid.uuid4().hex}" + + # construct a PythonUDF + py_udf = PythonUDF( + output_type=return_type, + eval_type=eval_type, + func=function, + python_ver="%d.%d" % sys.version_info[:2], + ) + + # construct a CommonInlineUserDefinedFunction + fun = CommonInlineUserDefinedFunction( + function_name=name, + arguments=[], + function=py_udf, + deterministic=deterministic, + ).to_plan_udf(self) + + # construct the request + req = self._execute_plan_request_with_metadata() + req.plan.command.register_function.CopyFrom(fun) + + self._execute(req) + return name + + def register_udtf( + self, + function: Any, + return_type: "DataTypeOrString", + name: str, + eval_type: int = PythonEvalType.SQL_TABLE_UDF, + deterministic: bool = True, + ) -> str: + """ + Register a user-defined table function (UDTF) in the session catalog + as a temporary function. The return type, if specified, must be a + struct type and it's validated when building the proto message + for the PythonUDTF. + """ + udtf = PythonUDTF( + func=function, + return_type=return_type, + eval_type=eval_type, + python_ver=get_python_ver(), + ) + + func = CommonInlineUserDefinedTableFunction( + function_name=name, + function=udtf, + deterministic=deterministic, + arguments=[], + ).udtf_plan(self) + + req = self._execute_plan_request_with_metadata() + req.plan.command.register_table_function.CopyFrom(func) + + self._execute(req) + return name + + def register_java( + self, + name: str, + javaClassName: str, + return_type: Optional["DataTypeOrString"] = None, + aggregate: bool = False, + ) -> None: + # construct a JavaUDF + if return_type is None: + java_udf = JavaUDF(class_name=javaClassName, aggregate=aggregate) + else: + java_udf = JavaUDF(class_name=javaClassName, output_type=return_type) + fun = CommonInlineUserDefinedFunction( + function_name=name, + function=java_udf, + ).to_plan_judf(self) + # construct the request + req = self._execute_plan_request_with_metadata() + req.plan.command.register_function.CopyFrom(fun) + + self._execute(req) + + def _build_metrics(self, metrics: "pb2.ExecutePlanResponse.Metrics") -> Iterator[PlanMetrics]: + return ( + PlanMetrics( + x.name, + x.plan_id, + x.parent, + [MetricValue(k, v.value, v.metric_type) for k, v in x.execution_metrics.items()], + ) + for x in metrics.metrics + ) + + def _resources(self) -> Dict[str, ResourceInformation]: + logger.info("Fetching the resources") + cmd = pb2.Command() + cmd.get_resources_command.SetInParent() + (_, properties) = self.execute_command(cmd) + resources = properties["get_resources_command_result"] + return resources + + def _build_observed_metrics( + self, metrics: Sequence["pb2.ExecutePlanResponse.ObservedMetrics"] + ) -> Iterator[PlanObservedMetrics]: + return (PlanObservedMetrics(x.name, [v for v in x.values]) for x in metrics) + + def to_table_as_iterator(self, plan: pb2.Plan) -> Iterator[Union[StructType, "pa.Table"]]: + """ + Return given plan as a PyArrow Table iterator. + """ + logger.info(f"Executing plan {self._proto_to_string(plan)}") + req = self._execute_plan_request_with_metadata() + req.plan.CopyFrom(plan) + for response in self._execute_and_fetch_as_iterator(req): + if isinstance(response, StructType): + yield response + elif isinstance(response, pa.RecordBatch): + yield pa.Table.from_batches([response]) + + def to_table(self, plan: pb2.Plan) -> Tuple["pa.Table", Optional[StructType]]: + """ + Return given plan as a PyArrow Table. + """ + logger.info(f"Executing plan {self._proto_to_string(plan)}") + req = self._execute_plan_request_with_metadata() + req.plan.CopyFrom(plan) + table, schema, _, _, _ = self._execute_and_fetch(req) + assert table is not None + return table, schema + + def to_pandas(self, plan: pb2.Plan) -> "pd.DataFrame": + """ + Return given plan as a pandas DataFrame. + """ + logger.info(f"Executing plan {self._proto_to_string(plan)}") + req = self._execute_plan_request_with_metadata() + req.plan.CopyFrom(plan) + (self_destruct_conf,) = self.get_config_with_defaults( + ("spark.sql.execution.arrow.pyspark.selfDestruct.enabled", "false"), + ) + self_destruct = cast(str, self_destruct_conf).lower() == "true" + table, schema, metrics, observed_metrics, _ = self._execute_and_fetch( + req, self_destruct=self_destruct + ) + assert table is not None + + schema = schema or from_arrow_schema(table.schema, prefer_timestamp_ntz=True) + assert schema is not None and isinstance(schema, StructType) + + # Rename columns to avoid duplicated column names. + renamed_table = table.rename_columns([f"col_{i}" for i in range(table.num_columns)]) + if self_destruct: + # Configure PyArrow to use as little memory as possible: + # self_destruct - free columns as they are converted + # split_blocks - create a separate Pandas block for each column + # use_threads - convert one column at a time + pandas_options = { + "self_destruct": True, + "split_blocks": True, + "use_threads": False, + } + pdf = renamed_table.to_pandas(**pandas_options) + else: + pdf = renamed_table.to_pandas() + pdf.columns = schema.names + + if len(pdf.columns) > 0: + timezone: Optional[str] = None + if any(_has_type(f.dataType, TimestampType) for f in schema.fields): + (timezone,) = self.get_configs("spark.sql.session.timeZone") + + struct_in_pandas: Optional[str] = None + error_on_duplicated_field_names: bool = False + if any(_has_type(f.dataType, StructType) for f in schema.fields): + (struct_in_pandas,) = self.get_config_with_defaults( + ("spark.sql.execution.pandas.structHandlingMode", "legacy"), + ) + + if struct_in_pandas == "legacy": + error_on_duplicated_field_names = True + struct_in_pandas = "dict" + + pdf = pd.concat( + [ + _create_converter_to_pandas( + field.dataType, + field.nullable, + timezone=timezone, + struct_in_pandas=struct_in_pandas, + error_on_duplicated_field_names=error_on_duplicated_field_names, + )(pser) + for (_, pser), field, pa_field in zip(pdf.items(), schema.fields, table.schema) + ], + axis="columns", + ) + + if len(metrics) > 0: + pdf.attrs["metrics"] = metrics + if len(observed_metrics) > 0: + pdf.attrs["observed_metrics"] = observed_metrics + return pdf + + def _proto_to_string(self, p: google.protobuf.message.Message) -> str: + """ + Helper method to generate a one line string representation of the plan. + + Parameters + ---------- + p : google.protobuf.message.Message + Generic Message type + + Returns + ------- + Single line string of the serialized proto message. + """ + return text_format.MessageToString(p, as_one_line=True) + + def schema(self, plan: pb2.Plan) -> StructType: + """ + Return schema for given plan. + """ + logger.info(f"Schema for plan: {self._proto_to_string(plan)}") + schema = self._analyze(method="schema", plan=plan).schema + assert schema is not None + # Server side should populate the struct field which is the schema. + assert isinstance(schema, StructType) + return schema + + def explain_string(self, plan: pb2.Plan, explain_mode: str = "extended") -> str: + """ + Return explain string for given plan. + """ + logger.info(f"Explain (mode={explain_mode}) for plan {self._proto_to_string(plan)}") + result = self._analyze( + method="explain", plan=plan, explain_mode=explain_mode + ).explain_string + assert result is not None + return result + + def execute_command( + self, command: pb2.Command + ) -> Tuple[Optional[pd.DataFrame], Dict[str, Any]]: + """ + Execute given command. + """ + logger.info(f"Execute command for command {self._proto_to_string(command)}") + req = self._execute_plan_request_with_metadata() + if self._user_id: + req.user_context.user_id = self._user_id + req.plan.command.CopyFrom(command) + data, _, _, _, properties = self._execute_and_fetch(req) + if data is not None: + return (data.to_pandas(), properties) + else: + return (None, properties) + + def same_semantics(self, plan: pb2.Plan, other: pb2.Plan) -> bool: + """ + return if two plans have the same semantics. + """ + result = self._analyze(method="same_semantics", plan=plan, other=other).is_same_semantics + assert result is not None + return result + + def semantic_hash(self, plan: pb2.Plan) -> int: + """ + returns a `hashCode` of the logical query plan. + """ + result = self._analyze(method="semantic_hash", plan=plan).semantic_hash + assert result is not None + return result + + def close(self) -> None: + """ + Close the channel. + """ + self._channel.close() + self._closed = True + + @property + def is_closed(self) -> bool: + """ + Returns if the channel was closed previously using close() method + """ + return self._closed + + @property + def host(self) -> str: + """ + The hostname where this client intends to connect. + """ + return self._builder.host + + @property + def token(self) -> Optional[str]: + """ + The authentication bearer token during connection. + If authentication is not using a bearer token, None will be returned. + """ + return self._builder._token + + def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest: + req = pb2.ExecutePlanRequest( + session_id=self._session_id, + client_type=self._builder.userAgent, + tags=list(self.get_tags()), + ) + if self._user_id: + req.user_context.user_id = self._user_id + return req + + def _analyze_plan_request_with_metadata(self) -> pb2.AnalyzePlanRequest: + req = pb2.AnalyzePlanRequest() + req.session_id = self._session_id + req.client_type = self._builder.userAgent + if self._user_id: + req.user_context.user_id = self._user_id + return req + + def _analyze(self, method: str, **kwargs: Any) -> AnalyzeResult: + """ + Call the analyze RPC of Spark Connect. + + Returns + ------- + The result of the analyze call. + """ + req = self._analyze_plan_request_with_metadata() + if method == "schema": + req.schema.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + elif method == "explain": + req.explain.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + explain_mode = kwargs.get("explain_mode") + if explain_mode not in ["simple", "extended", "codegen", "cost", "formatted"]: + raise PySparkValueError( + error_class="UNKNOWN_EXPLAIN_MODE", + message_parameters={ + "explain_mode": str(explain_mode), + }, + ) + if explain_mode == "simple": + req.explain.explain_mode = ( + pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_SIMPLE + ) + elif explain_mode == "extended": + req.explain.explain_mode = ( + pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_EXTENDED + ) + elif explain_mode == "cost": + req.explain.explain_mode = ( + pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_COST + ) + elif explain_mode == "codegen": + req.explain.explain_mode = ( + pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_CODEGEN + ) + else: # formatted + req.explain.explain_mode = ( + pb2.AnalyzePlanRequest.Explain.ExplainMode.EXPLAIN_MODE_FORMATTED + ) + elif method == "tree_string": + req.tree_string.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + level = kwargs.get("level") + if level and isinstance(level, int): + req.tree_string.level = level + elif method == "is_local": + req.is_local.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + elif method == "is_streaming": + req.is_streaming.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + elif method == "input_files": + req.input_files.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + elif method == "spark_version": + req.spark_version.SetInParent() + elif method == "ddl_parse": + req.ddl_parse.ddl_string = cast(str, kwargs.get("ddl_string")) + elif method == "same_semantics": + req.same_semantics.target_plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + req.same_semantics.other_plan.CopyFrom(cast(pb2.Plan, kwargs.get("other"))) + elif method == "semantic_hash": + req.semantic_hash.plan.CopyFrom(cast(pb2.Plan, kwargs.get("plan"))) + elif method == "persist": + req.persist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation"))) + if kwargs.get("storage_level", None) is not None: + storage_level = cast(StorageLevel, kwargs.get("storage_level")) + req.persist.storage_level.CopyFrom(storage_level_to_proto(storage_level)) + elif method == "unpersist": + req.unpersist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation"))) + if kwargs.get("blocking", None) is not None: + req.unpersist.blocking = cast(bool, kwargs.get("blocking")) + elif method == "get_storage_level": + req.get_storage_level.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation"))) + else: + raise PySparkValueError( + error_class="UNSUPPORTED_OPERATION", + message_parameters={ + "operation": method, + }, + ) + + try: + for attempt in self._retrying(): + with attempt: + resp = self._stub.AnalyzePlan(req, metadata=self._builder.metadata()) + if resp.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request:" + f"{resp.session_id} != {self._session_id}" + ) + return AnalyzeResult.fromProto(resp) + raise SparkConnectException("Invalid state during retry exception handling.") + except Exception as error: + self._handle_error(error) + + def _execute(self, req: pb2.ExecutePlanRequest) -> None: + """ + Execute the passed request `req` and drop all results. + + Parameters + ---------- + req : pb2.ExecutePlanRequest + Proto representation of the plan. + + """ + logger.info("Execute") + + def handle_response(b: pb2.ExecutePlanResponse) -> None: + if b.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request: " + f"{b.session_id} != {self._session_id}" + ) + + try: + if self._use_reattachable_execute: + # Don't use retryHandler - own retry handling is inside. + generator = ExecutePlanResponseReattachableIterator( + req, self._stub, self._retry_policy, self._builder.metadata() + ) + for b in generator: + handle_response(b) + else: + for attempt in self._retrying(): + with attempt: + for b in self._stub.ExecutePlan(req, metadata=self._builder.metadata()): + handle_response(b) + except Exception as error: + self._handle_error(error) + + def _execute_and_fetch_as_iterator( + self, req: pb2.ExecutePlanRequest + ) -> Iterator[ + Union[ + "pa.RecordBatch", + StructType, + PlanMetrics, + PlanObservedMetrics, + Dict[str, Any], + ] + ]: + logger.info("ExecuteAndFetchAsIterator") + + def handle_response( + b: pb2.ExecutePlanResponse, + ) -> Iterator[ + Union[ + "pa.RecordBatch", + StructType, + PlanMetrics, + PlanObservedMetrics, + Dict[str, Any], + ] + ]: + if b.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request: " + f"{b.session_id} != {self._session_id}" + ) + if b.HasField("metrics"): + logger.debug("Received metric batch.") + yield from self._build_metrics(b.metrics) + if b.observed_metrics: + logger.debug("Received observed metric batch.") + yield from self._build_observed_metrics(b.observed_metrics) + if b.HasField("schema"): + logger.debug("Received the schema.") + dt = types.proto_schema_to_pyspark_data_type(b.schema) + assert isinstance(dt, StructType) + yield dt + if b.HasField("sql_command_result"): + logger.debug("Received the SQL command result.") + yield {"sql_command_result": b.sql_command_result.relation} + if b.HasField("write_stream_operation_start_result"): + field = "write_stream_operation_start_result" + yield {field: b.write_stream_operation_start_result} + if b.HasField("streaming_query_command_result"): + yield {"streaming_query_command_result": b.streaming_query_command_result} + if b.HasField("streaming_query_manager_command_result"): + cmd_result = b.streaming_query_manager_command_result + yield {"streaming_query_manager_command_result": cmd_result} + if b.HasField("get_resources_command_result"): + resources = {} + for key, resource in b.get_resources_command_result.resources.items(): + name = resource.name + addresses = [address for address in resource.addresses] + resources[key] = ResourceInformation(name, addresses) + yield {"get_resources_command_result": resources} + if b.HasField("arrow_batch"): + logger.debug( + f"Received arrow batch rows={b.arrow_batch.row_count} " + f"size={len(b.arrow_batch.data)}" + ) + + with pa.ipc.open_stream(b.arrow_batch.data) as reader: + for batch in reader: + assert isinstance(batch, pa.RecordBatch) + yield batch + + try: + if self._use_reattachable_execute: + # Don't use retryHandler - own retry handling is inside. + generator = ExecutePlanResponseReattachableIterator( + req, self._stub, self._retry_policy, self._builder.metadata() + ) + for b in generator: + yield from handle_response(b) + else: + for attempt in self._retrying(): + with attempt: + for b in self._stub.ExecutePlan(req, metadata=self._builder.metadata()): + yield from handle_response(b) + except Exception as error: + self._handle_error(error) + + def _execute_and_fetch( + self, req: pb2.ExecutePlanRequest, self_destruct: bool = False + ) -> Tuple[ + Optional["pa.Table"], + Optional[StructType], + List[PlanMetrics], + List[PlanObservedMetrics], + Dict[str, Any], + ]: + logger.info("ExecuteAndFetch") + + observed_metrics: List[PlanObservedMetrics] = [] + metrics: List[PlanMetrics] = [] + batches: List[pa.RecordBatch] = [] + schema: Optional[StructType] = None + properties: Dict[str, Any] = {} + + for response in self._execute_and_fetch_as_iterator(req): + if isinstance(response, StructType): + schema = response + elif isinstance(response, pa.RecordBatch): + batches.append(response) + elif isinstance(response, PlanMetrics): + metrics.append(response) + elif isinstance(response, PlanObservedMetrics): + observed_metrics.append(response) + elif isinstance(response, dict): + properties.update(**response) + else: + raise PySparkValueError( + error_class="UNKNOWN_RESPONSE", + message_parameters={ + "response": response, + }, + ) + + if len(batches) > 0: + if self_destruct: + results = [] + for batch in batches: + # self_destruct frees memory column-wise, but Arrow record batches are + # oriented row-wise, so copies each column into its own allocation + batch = pa.RecordBatch.from_arrays( + [ + # This call actually reallocates the array + pa.concat_arrays([array]) + for array in batch + ], + schema=batch.schema, + ) + results.append(batch) + table = pa.Table.from_batches(batches=results) + # Ensure only the table has a reference to the batches, so that + # self_destruct (if enabled) is effective + del results + del batches + else: + table = pa.Table.from_batches(batches=batches) + return table, schema, metrics, observed_metrics, properties + else: + return None, schema, metrics, observed_metrics, properties + + def _config_request_with_metadata(self) -> pb2.ConfigRequest: + req = pb2.ConfigRequest() + req.session_id = self._session_id + req.client_type = self._builder.userAgent + if self._user_id: + req.user_context.user_id = self._user_id + return req + + def get_configs(self, *keys: str) -> Tuple[Optional[str], ...]: + op = pb2.ConfigRequest.Operation(get=pb2.ConfigRequest.Get(keys=keys)) + configs = dict(self.config(op).pairs) + return tuple(configs.get(key) for key in keys) + + def get_config_with_defaults( + self, *pairs: Tuple[str, Optional[str]] + ) -> Tuple[Optional[str], ...]: + op = pb2.ConfigRequest.Operation( + get_with_default=pb2.ConfigRequest.GetWithDefault( + pairs=[pb2.KeyValue(key=key, value=default) for key, default in pairs] + ) + ) + configs = dict(self.config(op).pairs) + return tuple(configs.get(key) for key, _ in pairs) + + def config(self, operation: pb2.ConfigRequest.Operation) -> ConfigResult: + """ + Call the config RPC of Spark Connect. + + Parameters + ---------- + operation : str + Operation kind + + Returns + ------- + The result of the config call. + """ + req = self._config_request_with_metadata() + req.operation.CopyFrom(operation) + try: + for attempt in self._retrying(): + with attempt: + resp = self._stub.Config(req, metadata=self._builder.metadata()) + if resp.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request:" + f"{resp.session_id} != {self._session_id}" + ) + return ConfigResult.fromProto(resp) + raise SparkConnectException("Invalid state during retry exception handling.") + except Exception as error: + self._handle_error(error) + + def _interrupt_request( + self, interrupt_type: str, id_or_tag: Optional[str] = None + ) -> pb2.InterruptRequest: + req = pb2.InterruptRequest() + req.session_id = self._session_id + req.client_type = self._builder.userAgent + if interrupt_type == "all": + req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_ALL + elif interrupt_type == "tag": + assert id_or_tag is not None + req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_TAG + req.operation_tag = id_or_tag + elif interrupt_type == "operation": + assert id_or_tag is not None + req.interrupt_type = pb2.InterruptRequest.InterruptType.INTERRUPT_TYPE_OPERATION_ID + req.operation_id = id_or_tag + else: + raise PySparkValueError( + error_class="UNKNOWN_INTERRUPT_TYPE", + message_parameters={ + "interrupt_type": str(interrupt_type), + }, + ) + if self._user_id: + req.user_context.user_id = self._user_id + return req + + def interrupt_all(self) -> Optional[List[str]]: + req = self._interrupt_request("all") + try: + for attempt in self._retrying(): + with attempt: + resp = self._stub.Interrupt(req, metadata=self._builder.metadata()) + if resp.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request:" + f"{resp.session_id} != {self._session_id}" + ) + return list(resp.interrupted_ids) + raise SparkConnectException("Invalid state during retry exception handling.") + except Exception as error: + self._handle_error(error) + + def interrupt_tag(self, tag: str) -> Optional[List[str]]: + req = self._interrupt_request("tag", tag) + try: + for attempt in self._retrying(): + with attempt: + resp = self._stub.Interrupt(req, metadata=self._builder.metadata()) + if resp.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request:" + f"{resp.session_id} != {self._session_id}" + ) + return list(resp.interrupted_ids) + raise SparkConnectException("Invalid state during retry exception handling.") + except Exception as error: + self._handle_error(error) + + def interrupt_operation(self, op_id: str) -> Optional[List[str]]: + req = self._interrupt_request("operation", op_id) + try: + for attempt in self._retrying(): + with attempt: + resp = self._stub.Interrupt(req, metadata=self._builder.metadata()) + if resp.session_id != self._session_id: + raise SparkConnectException( + "Received incorrect session identifier for request:" + f"{resp.session_id} != {self._session_id}" + ) + return list(resp.interrupted_ids) + raise SparkConnectException("Invalid state during retry exception handling.") + except Exception as error: + self._handle_error(error) + + def add_tag(self, tag: str) -> None: + self._throw_if_invalid_tag(tag) + if not hasattr(self.thread_local, "tags"): + self.thread_local.tags = set() + self.thread_local.tags.add(tag) + + def remove_tag(self, tag: str) -> None: + self._throw_if_invalid_tag(tag) + if not hasattr(self.thread_local, "tags"): + self.thread_local.tags = set() + self.thread_local.tags.remove(tag) + + def get_tags(self) -> Set[str]: + if not hasattr(self.thread_local, "tags"): + self.thread_local.tags = set() + return self.thread_local.tags + + def clear_tags(self) -> None: + self.thread_local.tags = set() + + def _throw_if_invalid_tag(self, tag: str) -> None: + """ + Validate if a tag for ExecutePlanRequest.tags is valid. Throw ``ValueError`` if + not. + """ + spark_job_tags_sep = "," + if tag is None: + raise ValueError("Spark Connect tag cannot be null.") + if spark_job_tags_sep in tag: + raise ValueError(f"Spark Connect tag cannot contain '{spark_job_tags_sep}'.") + if len(tag) == 0: + raise ValueError("Spark Connect tag cannot be an empty string.") + + def _handle_error(self, error: Exception) -> NoReturn: + """ + Handle errors that occur during RPC calls. + + Parameters + ---------- + error : Exception + An exception thrown during RPC calls. + + Returns + ------- + Throws the appropriate internal Python exception. + """ + if isinstance(error, grpc.RpcError): + self._handle_rpc_error(error) + elif isinstance(error, ValueError): + if "Cannot invoke RPC" in str(error) and "closed" in str(error): + raise SparkConnectException( + error_class="NO_ACTIVE_SESSION", message_parameters=dict() + ) from None + raise error + + def _handle_rpc_error(self, rpc_error: grpc.RpcError) -> NoReturn: + """ + Error handling helper for dealing with GRPC Errors. On the server side, certain + exceptions are enriched with additional RPC Status information. These are + unpacked in this function and put into the exception. + + To avoid overloading the user with GRPC errors, this message explicitly + swallows the error context from the call. This GRPC Error is logged however, + and can be enabled. + + Parameters + ---------- + rpc_error : grpc.RpcError + RPC Error containing the details of the exception. + + Returns + ------- + Throws the appropriate internal Python exception. + """ + logger.exception("GRPC Error received") + # We have to cast the value here because, a RpcError is a Call as well. + # https://grpc.github.io/grpc/python/grpc.html#grpc.UnaryUnaryMultiCallable.__call__ + status = rpc_status.from_call(cast(grpc.Call, rpc_error)) + if status: + for d in status.details: + if d.Is(error_details_pb2.ErrorInfo.DESCRIPTOR): + info = error_details_pb2.ErrorInfo() + d.Unpack(info) + raise convert_exception(info, status.message) from None + + raise SparkConnectGrpcException(status.message) from None + else: + raise SparkConnectGrpcException(str(rpc_error)) from None + + def add_artifacts(self, *path: str, pyfile: bool, archive: bool, file: bool) -> None: + self._artifact_manager.add_artifacts(*path, pyfile=pyfile, archive=archive, file=file) + + def copy_from_local_to_fs(self, local_path: str, dest_path: str) -> None: + self._artifact_manager._add_forward_to_fs_artifacts(local_path, dest_path) + + def cache_artifact(self, blob: bytes) -> str: + return self._artifact_manager.cache_artifact(blob) + + +class RetryState: + """ + Simple state helper that captures the state between retries of the exceptions. It + keeps track of the last exception thrown and how many in total. When the task + finishes successfully done() returns True. + """ + + def __init__(self) -> None: + self._exception: Optional[BaseException] = None + self._done = False + self._count = 0 + + def set_exception(self, exc: BaseException) -> None: + self._exception = exc + self._count += 1 + + def throw(self) -> None: + if self._exception is None: + raise RuntimeError("No exception is set") + raise self._exception + + def set_done(self) -> None: + self._done = True + + def count(self) -> int: + return self._count + + def done(self) -> bool: + return self._done + + +class AttemptManager: + """ + Simple ContextManager that is used to capture the exception thrown inside the context. + """ + + def __init__(self, check: Callable[..., bool], retry_state: RetryState) -> None: + self._retry_state = retry_state + self._can_retry = check + + def __enter__(self) -> None: + pass + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> Optional[bool]: + if isinstance(exc_val, BaseException): + # Swallow the exception. + if self._can_retry(exc_val) or isinstance(exc_val, RetryException): + self._retry_state.set_exception(exc_val) + return True + # Bubble up the exception. + return False + else: + self._retry_state.set_done() + return None + + def is_first_try(self) -> bool: + return self._retry_state._count == 0 + + +class Retrying: + """ + This helper class is used as a generator together with a context manager to + allow retrying exceptions in particular code blocks. The Retrying can be configured + with a lambda function that is can be filtered what kind of exceptions should be + retried. + + In addition, there are several parameters that are used to configure the exponential + backoff behavior. + + An example to use this class looks like this: + + .. code-block:: python + + for attempt in Retrying(can_retry=lambda x: isinstance(x, TransientError)): + with attempt: + # do the work. + + """ + + def __init__( + self, + max_retries: int, + initial_backoff: int, + max_backoff: int, + backoff_multiplier: float, + jitter: int, + min_jitter_threshold: int, + can_retry: Callable[..., bool] = lambda x: True, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + self._can_retry = can_retry + self._max_retries = max_retries + self._initial_backoff = initial_backoff + self._max_backoff = max_backoff + self._backoff_multiplier = backoff_multiplier + self._jitter = jitter + self._min_jitter_threshold = min_jitter_threshold + self._sleep = sleep + + def __iter__(self) -> Generator[AttemptManager, None, None]: + """ + Generator function to wrap the exception producing code block. + + Returns + ------- + A generator that yields the current attempt. + """ + retry_state = RetryState() + next_backoff: float = self._initial_backoff + + if self._max_retries < 0: + raise ValueError("Can't have negative number of retries") + + while not retry_state.done() and retry_state.count() <= self._max_retries: + # Do backoff + if retry_state.count() > 0: + # Randomize backoff for this iteration + backoff = next_backoff + next_backoff = min(self._max_backoff, next_backoff * self._backoff_multiplier) + + if backoff >= self._min_jitter_threshold: + backoff += random.uniform(0, self._jitter) + + logger.debug(f"Retrying call after {backoff} ms sleep") + self._sleep(backoff / 1000.0) + yield AttemptManager(self._can_retry, retry_state) + + if not retry_state.done(): + # Exceeded number of retries, throw last exception we had + retry_state.throw() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/reattach.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/reattach.py new file mode 100644 index 0000000..7e1e722 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/client/reattach.py @@ -0,0 +1,314 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import warnings +import uuid +from collections.abc import Generator +from typing import Optional, Dict, Any, Iterator, Iterable, Tuple, Callable, cast +from multiprocessing.pool import ThreadPool +import os + +import grpc +from grpc_status import rpc_status + +import pyspark.sql.connect.proto as pb2 +import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib + + +class ExecutePlanResponseReattachableIterator(Generator): + """ + Retryable iterator of ExecutePlanResponses to an ExecutePlan call. + + It can handle situations when: + - the ExecutePlanResponse stream was broken by retryable network error (governed by + retryPolicy) + - the ExecutePlanResponse was gracefully ended by the server without a ResultComplete + message; this tells the client that there is more, and it should reattach to continue. + + Initial iterator is the result of an ExecutePlan on the request, but it can be reattached with + ReattachExecute request. ReattachExecute request is provided the responseId of last returned + ExecutePlanResponse on the iterator to return a new iterator from server that continues after + that. If the initial ExecutePlan did not even reach the server, and hence reattach fails with + INVALID_HANDLE.OPERATION_NOT_FOUND, we attempt to retry ExecutePlan. + + In reattachable execute the server does buffer some responses in case the client needs to + backtrack. To let server release this buffer sooner, this iterator asynchronously sends + ReleaseExecute RPCs that instruct the server to release responses that it already processed. + """ + + _release_thread_pool = ThreadPool(os.cpu_count() if os.cpu_count() else 8) + + def __init__( + self, + request: pb2.ExecutePlanRequest, + stub: grpc_lib.SparkConnectServiceStub, + retry_policy: Dict[str, Any], + metadata: Iterable[Tuple[str, str]], + ): + self._request = request + self._retry_policy = retry_policy + if request.operation_id: + self._operation_id = request.operation_id + else: + # Add operation id, if not present. + # with operationId set by the client, the client can use it to try to reattach on error + # even before getting the first response. If the operation in fact didn't even reach the + # server, that will end with INVALID_HANDLE.OPERATION_NOT_FOUND error. + self._operation_id = str(uuid.uuid4()) + + self._stub = stub + request.request_options.append( + pb2.ExecutePlanRequest.RequestOption( + reattach_options=pb2.ReattachOptions(reattachable=True) + ) + ) + request.operation_id = self._operation_id + self._initial_request = request + + # ResponseId of the last response returned by next() + self._last_returned_response_id: Optional[str] = None + + # True after ResponseComplete message was seen in the stream. + # Server will always send this message at the end of the stream, if the underlying iterator + # finishes without producing one, another iterator needs to be reattached. + self._result_complete = False + + # Initial iterator comes from ExecutePlan request. + # Note: This is not retried, because no error would ever be thrown here, and GRPC will only + # throw error on first self._has_next(). + self._metadata = metadata + self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter( + self._stub.ExecutePlan(self._initial_request, metadata=metadata) + ) + + # Current item from this iterator. + self._current: Optional[pb2.ExecutePlanResponse] = None + + def send(self, value: Any) -> pb2.ExecutePlanResponse: + # will trigger reattach in case the stream completed without result_complete + if not self._has_next(): + raise StopIteration() + + ret = self._current + assert ret is not None + + self._last_returned_response_id = ret.response_id + if ret.HasField("result_complete"): + self._result_complete = True + self._release_all() + else: + self._release_until(self._last_returned_response_id) + self._current = None + return ret + + def _has_next(self) -> bool: + from pyspark.sql.connect.client.core import SparkConnectClient + from pyspark.sql.connect.client.core import Retrying + + if self._result_complete: + # After response complete response + return False + else: + try: + for attempt in Retrying( + can_retry=SparkConnectClient.retry_exception, **self._retry_policy + ): + with attempt: + if self._current is None: + try: + self._current = self._call_iter( + lambda: next(self._iterator) # type: ignore[arg-type] + ) + except StopIteration: + pass + + has_next = self._current is not None + + # Graceful reattach: + # If iterator ended, but there was no ResponseComplete, it means that + # there is more, and we need to reattach. While ResponseComplete didn't + # arrive, we keep reattaching. + if not self._result_complete and not has_next: + while not has_next: + # unset iterator for new ReattachExecute to be called in _call_iter + self._iterator = None + # shouldn't change + assert not self._result_complete + try: + self._current = self._call_iter( + lambda: next(self._iterator) # type: ignore[arg-type] + ) + except StopIteration: + pass + has_next = self._current is not None + return has_next + except Exception as e: + self._release_all() + raise e + return False + + def _release_until(self, until_response_id: str) -> None: + """ + Inform the server to release the buffered execution results until and including given + result. + + This will send an asynchronous RPC which will not block this iterator, the iterator can + continue to be consumed. + """ + if self._result_complete: + return + + from pyspark.sql.connect.client.core import SparkConnectClient + from pyspark.sql.connect.client.core import Retrying + + request = self._create_release_execute_request(until_response_id) + + def target() -> None: + try: + for attempt in Retrying( + can_retry=SparkConnectClient.retry_exception, **self._retry_policy + ): + with attempt: + self._stub.ReleaseExecute(request, metadata=self._metadata) + except Exception as e: + warnings.warn(f"ReleaseExecute failed with exception: {e}.") + + ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target) + + def _release_all(self) -> None: + """ + Inform the server to release the execution, either because all results were consumed, + or the execution finished with error and the error was received. + + This will send an asynchronous RPC which will not block this. The client continues + executing, and if the release fails, server is equipped to deal with abandoned executions. + """ + if self._result_complete: + return + + from pyspark.sql.connect.client.core import SparkConnectClient + from pyspark.sql.connect.client.core import Retrying + + request = self._create_release_execute_request(None) + + def target() -> None: + try: + for attempt in Retrying( + can_retry=SparkConnectClient.retry_exception, **self._retry_policy + ): + with attempt: + self._stub.ReleaseExecute(request, metadata=self._metadata) + except Exception as e: + warnings.warn(f"ReleaseExecute failed with exception: {e}.") + + ExecutePlanResponseReattachableIterator._release_thread_pool.apply_async(target) + self._result_complete = True + + def _call_iter(self, iter_fun: Callable) -> Any: + """ + Call next() on the iterator. If this fails with this operationId not existing + on the server, this means that the initial ExecutePlan request didn't even reach the + server. In that case, attempt to start again with ExecutePlan. + + Called inside retry block, so retryable failure will get handled upstream. + """ + if self._iterator is None: + # we get a new iterator with ReattachExecute if it was unset. + self._iterator = iter( + self._stub.ReattachExecute( + self._create_reattach_execute_request(), metadata=self._metadata + ) + ) + + try: + return iter_fun() + except grpc.RpcError as e: + status = rpc_status.from_call(cast(grpc.Call, e)) + if status is not None and "INVALID_HANDLE.OPERATION_NOT_FOUND" in status.message: + if self._last_returned_response_id is not None: + raise RuntimeError( + "OPERATION_NOT_FOUND on the server but " + "responses were already received from it.", + e, + ) + # Try a new ExecutePlan, and throw upstream for retry. + self._iterator = iter( + self._stub.ExecutePlan(self._initial_request, metadata=self._metadata) + ) + raise RetryException() + else: + # Remove the iterator, so that a new one will be created after retry. + self._iterator = None + raise e + except Exception as e: + # Remove the iterator, so that a new one will be created after retry. + self._iterator = None + raise e + + def _create_reattach_execute_request(self) -> pb2.ReattachExecuteRequest: + reattach = pb2.ReattachExecuteRequest( + session_id=self._initial_request.session_id, + user_context=self._initial_request.user_context, + operation_id=self._initial_request.operation_id, + ) + + if self._initial_request.client_type: + reattach.client_type = self._initial_request.client_type + + if self._last_returned_response_id: + reattach.last_response_id = self._last_returned_response_id + + return reattach + + def _create_release_execute_request( + self, until_response_id: Optional[str] + ) -> pb2.ReleaseExecuteRequest: + release = pb2.ReleaseExecuteRequest( + session_id=self._initial_request.session_id, + user_context=self._initial_request.user_context, + operation_id=self._initial_request.operation_id, + ) + + if self._initial_request.client_type: + release.client_type = self._initial_request.client_type + + if not until_response_id: + release.release_all.CopyFrom(pb2.ReleaseExecuteRequest.ReleaseAll()) + else: + release.release_until.response_id = until_response_id + + return release + + def throw(self, type: Any = None, value: Any = None, traceback: Any = None) -> Any: + super().throw(type, value, traceback) + + def close(self) -> None: + self._release_all() + return super().close() + + def __del__(self) -> None: + return self.close() + + +class RetryException(Exception): + """ + An exception that can be thrown upstream when inside retry and which will be retryable + regardless of policy. + """ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/column.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/column.py new file mode 100644 index 0000000..0529293 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/column.py @@ -0,0 +1,511 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import datetime +import decimal +import warnings + +from typing import ( + TYPE_CHECKING, + Callable, + Any, + Union, + overload, + Optional, +) + +from pyspark.errors import PySparkTypeError, PySparkAttributeError, PySparkValueError +from pyspark.sql.types import DataType +from pyspark.sql.column import Column as PySparkColumn + +import pyspark.sql.connect.proto as proto +from pyspark.sql.connect.expressions import ( + Expression, + UnresolvedFunction, + UnresolvedExtractValue, + LiteralExpression, + CaseWhen, + SortOrder, + CastExpression, + WindowExpression, + WithField, + DropField, +) + + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ( + LiteralType, + DateTimeLiteral, + DecimalLiteral, + ) + from pyspark.sql.connect.client import SparkConnectClient + from pyspark.sql.connect.window import WindowSpec + + +def _func_op(name: str, doc: Optional[str] = "") -> Callable[["Column"], "Column"]: + def wrapped(self: "Column") -> "Column": + return Column(UnresolvedFunction(name, [self._expr])) + + wrapped.__doc__ = doc + return wrapped + + +def _bin_op( + name: str, doc: Optional[str] = "binary function", reverse: bool = False +) -> Callable[["Column", Any], "Column"]: + def wrapped(self: "Column", other: Any) -> "Column": + if other is None or isinstance( + other, + ( + bool, + float, + int, + str, + datetime.datetime, + datetime.date, + decimal.Decimal, + datetime.timedelta, + ), + ): + other_expr = LiteralExpression._from_value(other) + else: + other_expr = other._expr + + if not reverse: + return Column(UnresolvedFunction(name, [self._expr, other_expr])) + else: + return Column(UnresolvedFunction(name, [other_expr, self._expr])) + + wrapped.__doc__ = doc + return wrapped + + +def _unary_op(name: str, doc: Optional[str] = "unary function") -> Callable[["Column"], "Column"]: + def wrapped(self: "Column") -> "Column": + return Column(UnresolvedFunction(name, [self._expr])) + + wrapped.__doc__ = doc + return wrapped + + +class Column: + def __init__(self, expr: "Expression") -> None: + if not isinstance(expr, Expression): + raise PySparkTypeError( + error_class="NOT_EXPRESSION", + message_parameters={"arg_name": "expr", "arg_type": type(expr).__name__}, + ) + self._expr = expr + + __gt__ = _bin_op(">") + __lt__ = _bin_op("<") + __add__ = _bin_op("+") + __sub__ = _bin_op("-") + __mul__ = _bin_op("*") + __div__ = _bin_op("/") + __truediv__ = _bin_op("/") + __mod__ = _bin_op("%") + __radd__ = _bin_op("+", reverse=True) + __rsub__ = _bin_op("-", reverse=True) + __rmul__ = _bin_op("*", reverse=True) + __rdiv__ = _bin_op("/", reverse=True) + __rtruediv__ = _bin_op("/", reverse=True) + __rmod__ = _bin_op("%", reverse=True) + __pow__ = _bin_op("power") + __rpow__ = _bin_op("power", reverse=True) + __ge__ = _bin_op(">=") + __le__ = _bin_op("<=") + + eqNullSafe = _bin_op("<=>", PySparkColumn.eqNullSafe.__doc__) + + __neg__ = _func_op("negative") + + # `and`, `or`, `not` cannot be overloaded in Python, + # so use bitwise operators as boolean operators + __and__ = _bin_op("and") + __or__ = _bin_op("or") + __invert__ = _func_op("not") + __rand__ = _bin_op("and") + __ror__ = _bin_op("or") + + # container operators + def __contains__(self, item: Any) -> None: + raise PySparkValueError( + error_class="CANNOT_APPLY_IN_FOR_COLUMN", + message_parameters={}, + ) + + # bitwise operators + bitwiseOR = _bin_op("|", PySparkColumn.bitwiseOR.__doc__) + bitwiseAND = _bin_op("&", PySparkColumn.bitwiseAND.__doc__) + bitwiseXOR = _bin_op("^", PySparkColumn.bitwiseXOR.__doc__) + + isNull = _unary_op("isnull", PySparkColumn.isNull.__doc__) + isNotNull = _unary_op("isnotnull", PySparkColumn.isNotNull.__doc__) + + def __ne__( # type: ignore[override] + self, + other: Any, + ) -> "Column": + """binary function""" + return _func_op("not")(_bin_op("==")(self, other)) + + # string methods + contains = _bin_op("contains", PySparkColumn.contains.__doc__) + startswith = _bin_op("startswith", PySparkColumn.startswith.__doc__) + endswith = _bin_op("endswith", PySparkColumn.endswith.__doc__) + + def when(self, condition: "Column", value: Any) -> "Column": + if not isinstance(condition, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + + if not isinstance(self._expr, CaseWhen): + raise PySparkTypeError( + error_class="INVALID_WHEN_USAGE", + message_parameters={}, + ) + + if self._expr._else_value is not None: + raise PySparkTypeError( + error_class="INVALID_WHEN_USAGE", + message_parameters={}, + ) + + if isinstance(value, Column): + _value = value._expr + else: + _value = LiteralExpression._from_value(value) + + _branches = self._expr._branches + [(condition._expr, _value)] + + return Column(CaseWhen(branches=_branches, else_value=None)) + + when.__doc__ = PySparkColumn.when.__doc__ + + def otherwise(self, value: Any) -> "Column": + if not isinstance(self._expr, CaseWhen): + raise PySparkTypeError( + "otherwise() can only be applied on a Column previously generated by when()" + ) + + if self._expr._else_value is not None: + raise PySparkTypeError( + "otherwise() can only be applied once on a Column previously generated by when()" + ) + + if isinstance(value, Column): + _value = value._expr + else: + _value = LiteralExpression._from_value(value) + + return Column(CaseWhen(branches=self._expr._branches, else_value=_value)) + + otherwise.__doc__ = PySparkColumn.otherwise.__doc__ + + like = _bin_op("like", PySparkColumn.like.__doc__) + rlike = _bin_op("rlike", PySparkColumn.rlike.__doc__) + ilike = _bin_op("ilike", PySparkColumn.ilike.__doc__) + + @overload + def substr(self, startPos: int, length: int) -> "Column": + ... + + @overload + def substr(self, startPos: "Column", length: "Column") -> "Column": + ... + + def substr(self, startPos: Union[int, "Column"], length: Union[int, "Column"]) -> "Column": + if type(startPos) != type(length): + raise PySparkTypeError( + error_class="NOT_SAME_TYPE", + message_parameters={ + "arg_name1": "startPos", + "arg_name2": "length", + "arg_type1": type(startPos).__name__, + "arg_type2": type(length).__name__, + }, + ) + + if isinstance(length, Column): + length_expr = length._expr + start_expr = startPos._expr # type: ignore[union-attr] + elif isinstance(length, int): + length_expr = LiteralExpression._from_value(length) + start_expr = LiteralExpression._from_value(startPos) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT", + message_parameters={"arg_name": "length", "arg_type": type(length).__name__}, + ) + return Column(UnresolvedFunction("substring", [self._expr, start_expr, length_expr])) + + substr.__doc__ = PySparkColumn.substr.__doc__ + + def __eq__(self, other: Any) -> "Column": # type: ignore[override] + """Returns a binary expression with the current column as the left + side and the other expression as the right side. + """ + if other is None or isinstance( + other, (bool, float, int, str, datetime.datetime, datetime.date, decimal.Decimal) + ): + other_expr = LiteralExpression._from_value(other) + else: + other_expr = other._expr + + return Column(UnresolvedFunction("==", [self._expr, other_expr])) + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + return self._expr.to_plan(session) + + def alias(self, *alias: str, **kwargs: Any) -> "Column": + return Column(self._expr.alias(*alias, **kwargs)) + + alias.__doc__ = PySparkColumn.alias.__doc__ + + name = alias + + name.__doc__ = PySparkColumn.name.__doc__ + + def asc(self) -> "Column": + return self.asc_nulls_first() + + asc.__doc__ = PySparkColumn.asc.__doc__ + + def asc_nulls_first(self) -> "Column": + return Column(SortOrder(self._expr, ascending=True, nullsFirst=True)) + + asc_nulls_first.__doc__ = PySparkColumn.asc_nulls_first.__doc__ + + def asc_nulls_last(self) -> "Column": + return Column(SortOrder(self._expr, ascending=True, nullsFirst=False)) + + asc_nulls_last.__doc__ = PySparkColumn.asc_nulls_last.__doc__ + + def desc(self) -> "Column": + return self.desc_nulls_last() + + desc.__doc__ = PySparkColumn.desc.__doc__ + + def desc_nulls_first(self) -> "Column": + return Column(SortOrder(self._expr, ascending=False, nullsFirst=True)) + + desc_nulls_first.__doc__ = PySparkColumn.desc_nulls_first.__doc__ + + def desc_nulls_last(self) -> "Column": + return Column(SortOrder(self._expr, ascending=False, nullsFirst=False)) + + desc_nulls_last.__doc__ = PySparkColumn.desc_nulls_last.__doc__ + + def cast(self, dataType: Union[DataType, str]) -> "Column": + if isinstance(dataType, (DataType, str)): + return Column(CastExpression(expr=self._expr, data_type=dataType)) + else: + raise PySparkTypeError( + error_class="NOT_DATATYPE_OR_STR", + message_parameters={"arg_name": "dataType", "arg_type": type(dataType).__name__}, + ) + + cast.__doc__ = PySparkColumn.cast.__doc__ + + astype = cast + + def __repr__(self) -> str: + return "Column<'%s'>" % self._expr.__repr__() + + def over(self, window: "WindowSpec") -> "Column": + from pyspark.sql.connect.window import WindowSpec + + if not isinstance(window, WindowSpec): + raise PySparkTypeError( + error_class="NOT_WINDOWSPEC", + message_parameters={"arg_name": "window", "arg_type": type(window).__name__}, + ) + + return Column(WindowExpression(windowFunction=self._expr, windowSpec=window)) + + over.__doc__ = PySparkColumn.over.__doc__ + + def isin(self, *cols: Any) -> "Column": + if len(cols) == 1 and isinstance(cols[0], (list, set)): + _cols = list(cols[0]) + else: + _cols = list(cols) + + _exprs = [self._expr] + for c in _cols: + if isinstance(c, Column): + _exprs.append(c._expr) + else: + _exprs.append(LiteralExpression._from_value(c)) + + return Column(UnresolvedFunction("in", _exprs)) + + isin.__doc__ = PySparkColumn.isin.__doc__ + + def between( + self, + lowerBound: Union["Column", "LiteralType", "DateTimeLiteral", "DecimalLiteral"], + upperBound: Union["Column", "LiteralType", "DateTimeLiteral", "DecimalLiteral"], + ) -> "Column": + return (self >= lowerBound) & (self <= upperBound) + + between.__doc__ = PySparkColumn.between.__doc__ + + def getItem(self, key: Any) -> "Column": + if isinstance(key, Column): + warnings.warn( + "A column as 'key' in getItem is deprecated as of Spark 3.0, and will not " + "be supported in the future release. Use `column[key]` or `column.key` syntax " + "instead.", + FutureWarning, + ) + return self[key] + + getItem.__doc__ = PySparkColumn.getItem.__doc__ + + def getField(self, name: Any) -> "Column": + if isinstance(name, Column): + warnings.warn( + "A column as 'name' in getField is deprecated as of Spark 3.0, and will not " + "be supported in the future release. Use `column[name]` or `column.name` syntax " + "instead.", + FutureWarning, + ) + return self[name] + + getField.__doc__ = PySparkColumn.getField.__doc__ + + def withField(self, fieldName: str, col: "Column") -> "Column": + if not isinstance(fieldName, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "fieldName", "arg_type": type(fieldName).__name__}, + ) + + if not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + + return Column(WithField(self._expr, fieldName, col._expr)) + + withField.__doc__ = PySparkColumn.withField.__doc__ + + def dropFields(self, *fieldNames: str) -> "Column": + dropField: Optional[DropField] = None + for fieldName in fieldNames: + if not isinstance(fieldName, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "fieldName", + "arg_type": type(fieldName).__name__, + }, + ) + + if dropField is None: + dropField = DropField(self._expr, fieldName) + else: + dropField = DropField(dropField, fieldName) + + if dropField is None: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={ + "item": "dropFields", + }, + ) + + return Column(dropField) + + dropFields.__doc__ = PySparkColumn.dropFields.__doc__ + + def __getattr__(self, item: Any) -> "Column": + if item == "_jc": + raise PySparkAttributeError( + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", message_parameters={"attr_name": "_jc"} + ) + if item.startswith("__"): + raise PySparkAttributeError( + error_class="ATTRIBUTE_NOT_SUPPORTED", message_parameters={"attr_name": item} + ) + return self[item] + + def __getitem__(self, k: Any) -> "Column": + if isinstance(k, slice): + if k.step is not None: + raise PySparkValueError( + error_class="SLICE_WITH_STEP", + message_parameters={}, + ) + return self.substr(k.start, k.stop) + else: + return Column(UnresolvedExtractValue(self._expr, LiteralExpression._from_value(k))) + + def __iter__(self) -> None: + raise PySparkTypeError( + error_class="NOT_ITERABLE", + message_parameters={"objectName": "Column"}, + ) + + def __nonzero__(self) -> None: + raise PySparkValueError( + error_class="CANNOT_CONVERT_COLUMN_INTO_BOOL", + message_parameters={}, + ) + + __bool__ = __nonzero__ + + +Column.__doc__ = PySparkColumn.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.column + + globs = pyspark.sql.connect.column.__dict__.copy() + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.column tests").remote("local[4]").getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.column, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conf.py new file mode 100644 index 0000000..d323de7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conf.py @@ -0,0 +1,125 @@ +# +# 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 Any, Optional, Union, cast +import warnings + +from pyspark import _NoValue +from pyspark._globals import _NoValueType +from pyspark.sql.conf import RuntimeConfig as PySparkRuntimeConfig +from pyspark.sql.connect import proto +from pyspark.sql.connect.client import SparkConnectClient + + +class RuntimeConf: + def __init__(self, client: SparkConnectClient) -> None: + """Create a new RuntimeConfig.""" + self._client = client + + __init__.__doc__ = PySparkRuntimeConfig.__init__.__doc__ + + def set(self, key: str, value: Union[str, int, bool]) -> None: + if isinstance(value, bool): + value = "true" if value else "false" + elif isinstance(value, int): + value = str(value) + op_set = proto.ConfigRequest.Set(pairs=[proto.KeyValue(key=key, value=value)]) + operation = proto.ConfigRequest.Operation(set=op_set) + result = self._client.config(operation) + for warn in result.warnings: + warnings.warn(warn) + + set.__doc__ = PySparkRuntimeConfig.set.__doc__ + + def get( + self, key: str, default: Union[Optional[str], _NoValueType] = _NoValue + ) -> Optional[str]: + self._checkType(key, "key") + if default is _NoValue: + op_get = proto.ConfigRequest.Get(keys=[key]) + operation = proto.ConfigRequest.Operation(get=op_get) + else: + if default is not None: + self._checkType(default, "default") + op_get_with_default = proto.ConfigRequest.GetWithDefault( + pairs=[proto.KeyValue(key=key, value=cast(Optional[str], default))] + ) + operation = proto.ConfigRequest.Operation(get_with_default=op_get_with_default) + result = self._client.config(operation) + return result.pairs[0][1] + + get.__doc__ = PySparkRuntimeConfig.get.__doc__ + + def unset(self, key: str) -> None: + op_unset = proto.ConfigRequest.Unset(keys=[key]) + operation = proto.ConfigRequest.Operation(unset=op_unset) + result = self._client.config(operation) + for warn in result.warnings: + warnings.warn(warn) + + unset.__doc__ = PySparkRuntimeConfig.unset.__doc__ + + def isModifiable(self, key: str) -> bool: + op_is_modifiable = proto.ConfigRequest.IsModifiable(keys=[key]) + operation = proto.ConfigRequest.Operation(is_modifiable=op_is_modifiable) + result = self._client.config(operation) + if result.pairs[0][1] == "true": + return True + elif result.pairs[0][1] == "false": + return False + else: + raise ValueError(f"Unknown boolean value: {result.pairs[0][1]}") + + isModifiable.__doc__ = PySparkRuntimeConfig.isModifiable.__doc__ + + def _checkType(self, obj: Any, identifier: str) -> None: + """Assert that an object is of type str.""" + if not isinstance(obj, str): + raise TypeError( + "expected %s '%s' to be a string (was '%s')" % (identifier, obj, type(obj).__name__) + ) + + +RuntimeConf.__doc__ = PySparkRuntimeConfig.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.conf + + globs = pyspark.sql.connect.conf.__dict__.copy() + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.conf tests").remote("local[4]").getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.conf, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conversion.py new file mode 100644 index 0000000..1afeb3d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/conversion.py @@ -0,0 +1,490 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import array +import datetime +import decimal + +import pyarrow as pa + +from pyspark.sql.types import ( + _create_row, + Row, + DataType, + TimestampType, + TimestampNTZType, + MapType, + StructField, + StructType, + ArrayType, + BinaryType, + NullType, + DecimalType, + StringType, + UserDefinedType, +) + +from pyspark.storagelevel import StorageLevel +import pyspark.sql.connect.proto as pb2 +from pyspark.sql.pandas.types import to_arrow_schema, _dedup_names, _deduplicate_field_names + +from typing import ( + Any, + Callable, + Sequence, + List, +) + + +class LocalDataToArrowConversion: + """ + Conversion from local data (except pandas DataFrame and numpy ndarray) to Arrow. + Currently, only :class:`SparkSession` in Spark Connect can use this class. + """ + + @staticmethod + def _need_converter(dataType: DataType) -> bool: + if isinstance(dataType, NullType): + return True + elif isinstance(dataType, StructType): + # Struct maybe rows, should convert to dict. + return True + elif isinstance(dataType, ArrayType): + return LocalDataToArrowConversion._need_converter(dataType.elementType) + elif isinstance(dataType, MapType): + # Different from PySpark, here always needs conversion, + # since an Arrow Map requires a list of tuples. + return True + elif isinstance(dataType, BinaryType): + return True + elif isinstance(dataType, (TimestampType, TimestampNTZType)): + # Always truncate + return True + elif isinstance(dataType, DecimalType): + # Convert Decimal('NaN') to None + return True + elif isinstance(dataType, StringType): + # Coercion to StringType is allowed + return True + elif isinstance(dataType, UserDefinedType): + return True + else: + return False + + @staticmethod + def _create_converter(dataType: DataType) -> Callable: + assert dataType is not None and isinstance(dataType, DataType) + + if not LocalDataToArrowConversion._need_converter(dataType): + return lambda value: value + + if isinstance(dataType, NullType): + return lambda value: None + + elif isinstance(dataType, StructType): + + field_names = dataType.fieldNames() + dedup_field_names = _dedup_names(dataType.names) + + field_convs = [ + LocalDataToArrowConversion._create_converter(field.dataType) + for field in dataType.fields + ] + + def convert_struct(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, (tuple, dict)) or hasattr( + value, "__dict__" + ), f"{type(value)} {value}" + + _dict = {} + if ( + not isinstance(value, Row) + and not isinstance(value, tuple) # inherited namedtuple + and hasattr(value, "__dict__") + ): + value = value.__dict__ + if isinstance(value, dict): + for i, field in enumerate(field_names): + _dict[dedup_field_names[i]] = field_convs[i](value.get(field)) + else: + if len(value) != len(field_names): + raise ValueError( + f"Length mismatch: Expected axis has {len(field_names)} elements, " + f"new values have {len(value)} elements" + ) + for i in range(len(field_names)): + _dict[dedup_field_names[i]] = field_convs[i](value[i]) + + return _dict + + return convert_struct + + elif isinstance(dataType, ArrayType): + + element_conv = LocalDataToArrowConversion._create_converter(dataType.elementType) + + def convert_array(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, (list, array.array)) + return [element_conv(v) for v in value] + + return convert_array + + elif isinstance(dataType, MapType): + + key_conv = LocalDataToArrowConversion._create_converter(dataType.keyType) + value_conv = LocalDataToArrowConversion._create_converter(dataType.valueType) + + def convert_map(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, dict) + + _tuples = [] + for k, v in value.items(): + _tuples.append((key_conv(k), value_conv(v))) + + return _tuples + + return convert_map + + elif isinstance(dataType, BinaryType): + + def convert_binary(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, (bytes, bytearray)) + return bytes(value) + + return convert_binary + + elif isinstance(dataType, TimestampType): + + def convert_timestamp(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, datetime.datetime) + return value.astimezone(datetime.timezone.utc) + + return convert_timestamp + + elif isinstance(dataType, TimestampNTZType): + + def convert_timestamp_ntz(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, datetime.datetime) and value.tzinfo is None + return value + + return convert_timestamp_ntz + + elif isinstance(dataType, DecimalType): + + def convert_decimal(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, decimal.Decimal) + return None if value.is_nan() else value + + return convert_decimal + + elif isinstance(dataType, StringType): + + def convert_string(value: Any) -> Any: + if value is None: + return None + else: + # only atomic types are supported + assert isinstance( + value, + ( + bool, + int, + float, + str, + bytes, + bytearray, + decimal.Decimal, + datetime.date, + datetime.datetime, + datetime.timedelta, + ), + ) + if isinstance(value, bool): + # To match the PySpark which convert bool to string in + # the JVM side (python.EvaluatePython.makeFromJava) + return str(value).lower() + else: + return str(value) + + return convert_string + + elif isinstance(dataType, UserDefinedType): + udt: UserDefinedType = dataType + + conv = LocalDataToArrowConversion._create_converter(udt.sqlType()) + + def convert_udt(value: Any) -> Any: + if value is None: + return None + else: + return conv(udt.serialize(value)) + + return convert_udt + + else: + + return lambda value: value + + @staticmethod + def convert(data: Sequence[Any], schema: StructType) -> "pa.Table": + assert isinstance(data, list) and len(data) > 0 + + assert schema is not None and isinstance(schema, StructType) + + column_names = schema.fieldNames() + + column_convs = [ + LocalDataToArrowConversion._create_converter(field.dataType) for field in schema.fields + ] + + pylist: List[List] = [[] for _ in range(len(column_names))] + + for item in data: + if ( + not isinstance(item, Row) + and not isinstance(item, tuple) # inherited namedtuple + and hasattr(item, "__dict__") + ): + item = item.__dict__ + if isinstance(item, dict): + for i, col in enumerate(column_names): + pylist[i].append(column_convs[i](item.get(col))) + else: + if len(item) != len(column_names): + raise ValueError( + f"Length mismatch: Expected axis has {len(column_names)} elements, " + f"new values have {len(item)} elements" + ) + for i in range(len(column_names)): + pylist[i].append(column_convs[i](item[i])) + + pa_schema = to_arrow_schema( + StructType( + [ + StructField( + field.name, _deduplicate_field_names(field.dataType), field.nullable + ) + for field in schema.fields + ] + ) + ) + + return pa.Table.from_arrays(pylist, schema=pa_schema) + + +class ArrowTableToRowsConversion: + """ + Conversion from Arrow Table to Rows. + Currently, only :class:`DataFrame` in Spark Connect can use this class. + """ + + @staticmethod + def _need_converter(dataType: DataType) -> bool: + if isinstance(dataType, NullType): + return True + elif isinstance(dataType, StructType): + return True + elif isinstance(dataType, ArrayType): + return ArrowTableToRowsConversion._need_converter(dataType.elementType) + elif isinstance(dataType, MapType): + # Different from PySpark, here always needs conversion, + # since the input from Arrow is a list of tuples. + return True + elif isinstance(dataType, BinaryType): + return True + elif isinstance(dataType, (TimestampType, TimestampNTZType)): + # Always remove the time zone info for now + return True + elif isinstance(dataType, UserDefinedType): + return True + else: + return False + + @staticmethod + def _create_converter(dataType: DataType) -> Callable: + assert dataType is not None and isinstance(dataType, DataType) + + if not ArrowTableToRowsConversion._need_converter(dataType): + return lambda value: value + + if isinstance(dataType, NullType): + return lambda value: None + + elif isinstance(dataType, StructType): + + field_names = dataType.names + dedup_field_names = _dedup_names(field_names) + + field_convs = [ + ArrowTableToRowsConversion._create_converter(f.dataType) for f in dataType.fields + ] + + def convert_struct(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, dict) + + _values = [ + field_convs[i](value.get(name, None)) + for i, name in enumerate(dedup_field_names) + ] + return _create_row(field_names, _values) + + return convert_struct + + elif isinstance(dataType, ArrayType): + + element_conv = ArrowTableToRowsConversion._create_converter(dataType.elementType) + + def convert_array(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, list) + return [element_conv(v) for v in value] + + return convert_array + + elif isinstance(dataType, MapType): + + key_conv = ArrowTableToRowsConversion._create_converter(dataType.keyType) + value_conv = ArrowTableToRowsConversion._create_converter(dataType.valueType) + + def convert_map(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, list) + assert all(isinstance(t, tuple) and len(t) == 2 for t in value) + return dict((key_conv(t[0]), value_conv(t[1])) for t in value) + + return convert_map + + elif isinstance(dataType, BinaryType): + + def convert_binary(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, bytes) + return bytearray(value) + + return convert_binary + + elif isinstance(dataType, TimestampType): + + def convert_timestample(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, datetime.datetime) + return value.astimezone().replace(tzinfo=None) + + return convert_timestample + + elif isinstance(dataType, TimestampNTZType): + + def convert_timestample_ntz(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, datetime.datetime) + return value + + return convert_timestample_ntz + + elif isinstance(dataType, UserDefinedType): + udt: UserDefinedType = dataType + + conv = ArrowTableToRowsConversion._create_converter(udt.sqlType()) + + def convert_udt(value: Any) -> Any: + if value is None: + return None + else: + return udt.deserialize(conv(value)) + + return convert_udt + + else: + + return lambda value: value + + @staticmethod + def convert(table: "pa.Table", schema: StructType) -> List[Row]: + assert isinstance(table, pa.Table) + + assert schema is not None and isinstance(schema, StructType) + + field_converters = [ + ArrowTableToRowsConversion._create_converter(f.dataType) for f in schema.fields + ] + + columnar_data = [column.to_pylist() for column in table.columns] + + rows: List[Row] = [] + for i in range(0, table.num_rows): + values = [field_converters[j](columnar_data[j][i]) for j in range(table.num_columns)] + rows.append(_create_row(fields=schema.fieldNames(), values=values)) + return rows + + +def storage_level_to_proto(storage_level: StorageLevel) -> pb2.StorageLevel: + assert storage_level is not None and isinstance(storage_level, StorageLevel) + return pb2.StorageLevel( + use_disk=storage_level.useDisk, + use_memory=storage_level.useMemory, + use_off_heap=storage_level.useOffHeap, + deserialized=storage_level.deserialized, + replication=storage_level.replication, + ) + + +def proto_to_storage_level(storage_level: pb2.StorageLevel) -> StorageLevel: + assert storage_level is not None and isinstance(storage_level, pb2.StorageLevel) + return StorageLevel( + useDisk=storage_level.use_disk, + useMemory=storage_level.use_memory, + useOffHeap=storage_level.use_off_heap, + deserialized=storage_level.deserialized, + replication=storage_level.replication, + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/dataframe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/dataframe.py new file mode 100644 index 0000000..7b32653 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/dataframe.py @@ -0,0 +1,2172 @@ +# +# 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.errors.exceptions.base import SessionNotSameException +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import ( + Any, + Dict, + Iterator, + List, + Optional, + Tuple, + Union, + Sequence, + TYPE_CHECKING, + overload, + Callable, + cast, + Type, +) + +import sys +import random +import pandas +import pyarrow as pa +import json +import warnings +from collections.abc import Iterable + +from pyspark import _NoValue +from pyspark._globals import _NoValueType +from pyspark.sql.observation import Observation +from pyspark.sql.types import Row, StructType +from pyspark.sql.dataframe import ( + DataFrame as PySparkDataFrame, + DataFrameNaFunctions as PySparkDataFrameNaFunctions, + DataFrameStatFunctions as PySparkDataFrameStatFunctions, +) + +from pyspark.errors import ( + PySparkTypeError, + PySparkAttributeError, + PySparkValueError, + PySparkNotImplementedError, +) +from pyspark.errors.exceptions.connect import SparkConnectException +from pyspark.rdd import PythonEvalType +from pyspark.storagelevel import StorageLevel +import pyspark.sql.connect.plan as plan +from pyspark.sql.connect.group import GroupedData +from pyspark.sql.connect.readwriter import DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.connect.streaming.readwriter import DataStreamWriter +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.expressions import UnresolvedRegex +from pyspark.sql.connect.functions import ( + _to_col_with_plan_id, + _to_col, + _invoke_function, + col, + lit, + expr as sql_expression, +) +from pyspark.sql.pandas.types import from_arrow_schema + + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ( + ColumnOrName, + LiteralType, + PrimitiveType, + OptionalPrimitiveType, + PandasMapIterFunction, + ArrowMapIterFunction, + ) + from pyspark.sql.connect.session import SparkSession + from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame + + +class DataFrame: + def __init__( + self, + session: "SparkSession", + schema: Optional[StructType] = None, + ): + """Creates a new data frame""" + self._schema = schema + self._plan: Optional[plan.LogicalPlan] = None + self._session: "SparkSession" = session + # Check whether _repr_html is supported or not, we use it to avoid calling RPC twice + # by __repr__ and _repr_html_ while eager evaluation opens. + self._support_repr_html = False + + def __repr__(self) -> str: + if not self._support_repr_html: + ( + repl_eager_eval_enabled, + repl_eager_eval_max_num_rows, + repl_eager_eval_truncate, + ) = self._session._client.get_configs( + "spark.sql.repl.eagerEval.enabled", + "spark.sql.repl.eagerEval.maxNumRows", + "spark.sql.repl.eagerEval.truncate", + ) + if repl_eager_eval_enabled == "true": + return self._show_string( + n=int(cast(str, repl_eager_eval_max_num_rows)), + truncate=int(cast(str, repl_eager_eval_truncate)), + vertical=False, + ) + return "DataFrame[%s]" % (", ".join("%s: %s" % c for c in self.dtypes)) + + def _repr_html_(self) -> Optional[str]: + if not self._support_repr_html: + self._support_repr_html = True + ( + repl_eager_eval_enabled, + repl_eager_eval_max_num_rows, + repl_eager_eval_truncate, + ) = self._session._client.get_configs( + "spark.sql.repl.eagerEval.enabled", + "spark.sql.repl.eagerEval.maxNumRows", + "spark.sql.repl.eagerEval.truncate", + ) + if repl_eager_eval_enabled == "true": + pdf = DataFrame.withPlan( + plan.HtmlString( + child=self._plan, + num_rows=int(cast(str, repl_eager_eval_max_num_rows)), + truncate=int(cast(str, repl_eager_eval_truncate)), + ), + session=self._session, + ).toPandas() + assert pdf is not None + return pdf["html_string"][0] + else: + return None + + _repr_html_.__doc__ = PySparkDataFrame._repr_html_.__doc__ + + @property + def write(self) -> "DataFrameWriter": + assert self._plan is not None + return DataFrameWriter(self._plan, self._session) + + write.__doc__ = PySparkDataFrame.write.__doc__ + + def isEmpty(self) -> bool: + return len(self.take(1)) == 0 + + isEmpty.__doc__ = PySparkDataFrame.isEmpty.__doc__ + + def select(self, *cols: "ColumnOrName") -> "DataFrame": + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] + + return DataFrame.withPlan(plan.Project(self._plan, *cols), session=self._session) + + select.__doc__ = PySparkDataFrame.select.__doc__ + + def selectExpr(self, *expr: Union[str, List[str]]) -> "DataFrame": + sql_expr = [] + if len(expr) == 1 and isinstance(expr[0], list): + expr = expr[0] # type: ignore[assignment] + for element in expr: + if isinstance(element, str): + sql_expr.append(sql_expression(element)) + else: + sql_expr.extend([sql_expression(e) for e in element]) + + return DataFrame.withPlan(plan.Project(self._plan, *sql_expr), session=self._session) + + selectExpr.__doc__ = PySparkDataFrame.selectExpr.__doc__ + + def agg(self, *exprs: Union[Column, Dict[str, str]]) -> "DataFrame": + if not exprs: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "exprs"}, + ) + + if len(exprs) == 1 and isinstance(exprs[0], dict): + measures = [_invoke_function(f, col(e)) for e, f in exprs[0].items()] + return self.groupBy().agg(*measures) + else: + # other expressions + assert all(isinstance(c, Column) for c in exprs), "all exprs should be Expression" + exprs = cast(Tuple[Column, ...], exprs) + return self.groupBy().agg(*exprs) + + agg.__doc__ = PySparkDataFrame.agg.__doc__ + + def alias(self, alias: str) -> "DataFrame": + return DataFrame.withPlan(plan.SubqueryAlias(self._plan, alias), session=self._session) + + alias.__doc__ = PySparkDataFrame.alias.__doc__ + + def colRegex(self, colName: str) -> Column: + if not isinstance(colName, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "colName", "arg_type": type(colName).__name__}, + ) + if self._plan is not None: + return Column(UnresolvedRegex(colName, self._plan._plan_id)) + else: + return Column(UnresolvedRegex(colName)) + + colRegex.__doc__ = PySparkDataFrame.colRegex.__doc__ + + @property + def dtypes(self) -> List[Tuple[str, str]]: + return [(str(f.name), f.dataType.simpleString()) for f in self.schema.fields] + + dtypes.__doc__ = PySparkDataFrame.dtypes.__doc__ + + @property + def columns(self) -> List[str]: + if self._plan is None: + return [] + + return self.schema.names + + columns.__doc__ = PySparkDataFrame.columns.__doc__ + + @property + def sparkSession(self) -> "SparkSession": + return self._session + + sparkSession.__doc__ = PySparkDataFrame.sparkSession.__doc__ + + def count(self) -> int: + pdd = self.agg(_invoke_function("count", lit(1))).toPandas() + return pdd.iloc[0, 0] + + count.__doc__ = PySparkDataFrame.count.__doc__ + + def crossJoin(self, other: "DataFrame") -> "DataFrame": + if self._plan is None: + raise Exception("Cannot cartesian join when self._plan is empty.") + if other._plan is None: + raise Exception("Cannot cartesian join when other._plan is empty.") + self.checkSameSparkSession(other) + return DataFrame.withPlan( + plan.Join(left=self._plan, right=other._plan, on=None, how="cross"), + session=self._session, + ) + + crossJoin.__doc__ = PySparkDataFrame.crossJoin.__doc__ + + def checkSameSparkSession(self, other: "DataFrame") -> None: + if self._session.session_id != other._session.session_id: + raise SessionNotSameException( + error_class="SESSION_NOT_SAME", + message_parameters={}, + ) + + def coalesce(self, numPartitions: int) -> "DataFrame": + if not numPartitions > 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "numPartitions", "arg_value": str(numPartitions)}, + ) + return DataFrame.withPlan( + plan.Repartition(self._plan, num_partitions=numPartitions, shuffle=False), + self._session, + ) + + coalesce.__doc__ = PySparkDataFrame.coalesce.__doc__ + + @overload + def repartition(self, numPartitions: int, *cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def repartition(self, *cols: "ColumnOrName") -> "DataFrame": + ... + + def repartition( # type: ignore[misc] + self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName" + ) -> "DataFrame": + if isinstance(numPartitions, int): + if not numPartitions > 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={ + "arg_name": "numPartitions", + "arg_value": str(numPartitions), + }, + ) + if len(cols) == 0: + return DataFrame.withPlan( + plan.Repartition(self._plan, num_partitions=numPartitions, shuffle=True), + self._session, + ) + else: + return DataFrame.withPlan( + plan.RepartitionByExpression(self._plan, numPartitions, list(cols)), + self.sparkSession, + ) + elif isinstance(numPartitions, (str, Column)): + cols = (numPartitions,) + cols + return DataFrame.withPlan( + plan.RepartitionByExpression(self._plan, None, list(cols)), + self.sparkSession, + ) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={ + "arg_name": "numPartitions", + "arg_type": type(numPartitions).__name__, + }, + ) + + repartition.__doc__ = PySparkDataFrame.repartition.__doc__ + + @overload + def repartitionByRange(self, numPartitions: int, *cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def repartitionByRange(self, *cols: "ColumnOrName") -> "DataFrame": + ... + + def repartitionByRange( # type: ignore[misc] + self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName" + ) -> "DataFrame": + def _convert_col(col: "ColumnOrName") -> "ColumnOrName": + from pyspark.sql.connect.expressions import SortOrder, ColumnReference + + if isinstance(col, Column): + if isinstance(col._expr, SortOrder): + return col + else: + return Column(SortOrder(col._expr)) + else: + return Column(SortOrder(ColumnReference(col))) + + if isinstance(numPartitions, int): + if not numPartitions > 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={ + "arg_name": "numPartitions", + "arg_value": str(numPartitions), + }, + ) + if len(cols) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "cols"}, + ) + else: + sort = [] + sort.extend([_convert_col(c) for c in cols]) + return DataFrame.withPlan( + plan.RepartitionByExpression(self._plan, numPartitions, sort), + self.sparkSession, + ) + elif isinstance(numPartitions, (str, Column)): + cols = (numPartitions,) + cols + sort = [] + sort.extend([_convert_col(c) for c in cols]) + return DataFrame.withPlan( + plan.RepartitionByExpression(self._plan, None, sort), + self.sparkSession, + ) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={ + "arg_name": "numPartitions", + "arg_type": type(numPartitions).__name__, + }, + ) + + repartitionByRange.__doc__ = PySparkDataFrame.repartitionByRange.__doc__ + + def dropDuplicates(self, subset: Optional[List[str]] = None) -> "DataFrame": + if subset is not None and not isinstance(subset, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if subset is None: + return DataFrame.withPlan( + plan.Deduplicate(child=self._plan, all_columns_as_keys=True), session=self._session + ) + else: + return DataFrame.withPlan( + plan.Deduplicate(child=self._plan, column_names=subset), session=self._session + ) + + dropDuplicates.__doc__ = PySparkDataFrame.dropDuplicates.__doc__ + + drop_duplicates = dropDuplicates + + def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> "DataFrame": + if subset is not None and not isinstance(subset, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if subset is None: + return DataFrame.withPlan( + plan.Deduplicate(child=self._plan, all_columns_as_keys=True, within_watermark=True), + session=self._session, + ) + else: + return DataFrame.withPlan( + plan.Deduplicate(child=self._plan, column_names=subset, within_watermark=True), + session=self._session, + ) + + dropDuplicatesWithinWatermark.__doc__ = PySparkDataFrame.dropDuplicatesWithinWatermark.__doc__ + + drop_duplicates_within_watermark = dropDuplicatesWithinWatermark + + def distinct(self) -> "DataFrame": + return DataFrame.withPlan( + plan.Deduplicate(child=self._plan, all_columns_as_keys=True), session=self._session + ) + + distinct.__doc__ = PySparkDataFrame.distinct.__doc__ + + def drop(self, *cols: "ColumnOrName") -> "DataFrame": + _cols = list(cols) + if any(not isinstance(c, (str, Column)) for c in _cols): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "cols", "arg_type": type(cols).__name__}, + ) + + return DataFrame.withPlan( + plan.Drop( + child=self._plan, + columns=_cols, + ), + session=self._session, + ) + + drop.__doc__ = PySparkDataFrame.drop.__doc__ + + def filter(self, condition: Union[Column, str]) -> "DataFrame": + if isinstance(condition, str): + expr = sql_expression(condition) + else: + expr = condition + return DataFrame.withPlan(plan.Filter(child=self._plan, filter=expr), session=self._session) + + filter.__doc__ = PySparkDataFrame.filter.__doc__ + + def first(self) -> Optional[Row]: + return self.head() + + first.__doc__ = PySparkDataFrame.first.__doc__ + + def groupBy(self, *cols: "ColumnOrName") -> GroupedData: + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] + + _cols: List[Column] = [] + for c in cols: + if isinstance(c, Column): + _cols.append(c) + elif isinstance(c, str): + _cols.append(self[c]) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "groupBy", "arg_type": type(c).__name__}, + ) + + return GroupedData(df=self, group_type="groupby", grouping_cols=_cols) + + groupBy.__doc__ = PySparkDataFrame.groupBy.__doc__ + + groupby = groupBy + + def rollup(self, *cols: "ColumnOrName") -> "GroupedData": + _cols: List[Column] = [] + for c in cols: + if isinstance(c, Column): + _cols.append(c) + elif isinstance(c, str): + _cols.append(self[c]) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "rollup", "arg_type": type(c).__name__}, + ) + + return GroupedData(df=self, group_type="rollup", grouping_cols=_cols) + + rollup.__doc__ = PySparkDataFrame.rollup.__doc__ + + def cube(self, *cols: "ColumnOrName") -> "GroupedData": + _cols: List[Column] = [] + for c in cols: + if isinstance(c, Column): + _cols.append(c) + elif isinstance(c, str): + _cols.append(self[c]) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "cube", "arg_type": type(c).__name__}, + ) + + return GroupedData(df=self, group_type="cube", grouping_cols=_cols) + + cube.__doc__ = PySparkDataFrame.cube.__doc__ + + @overload + def head(self) -> Optional[Row]: + ... + + @overload + def head(self, n: int) -> List[Row]: + ... + + def head(self, n: Optional[int] = None) -> Union[Optional[Row], List[Row]]: + if n is None: + rs = self.head(1) + return rs[0] if rs else None + return self.take(n) + + head.__doc__ = PySparkDataFrame.head.__doc__ + + def take(self, num: int) -> List[Row]: + return self.limit(num).collect() + + take.__doc__ = PySparkDataFrame.take.__doc__ + + # TODO: extend `on` to also be type List[Column]. + def join( + self, + other: "DataFrame", + on: Optional[Union[str, List[str], Column, List[Column]]] = None, + how: Optional[str] = None, + ) -> "DataFrame": + if self._plan is None: + raise Exception("Cannot join when self._plan is empty.") + if other._plan is None: + raise Exception("Cannot join when other._plan is empty.") + if how is not None and isinstance(how, str): + how = how.lower().replace("_", "") + self.checkSameSparkSession(other) + return DataFrame.withPlan( + plan.Join(left=self._plan, right=other._plan, on=on, how=how), + session=self._session, + ) + + join.__doc__ = PySparkDataFrame.join.__doc__ + + def limit(self, n: int) -> "DataFrame": + return DataFrame.withPlan(plan.Limit(child=self._plan, limit=n), session=self._session) + + limit.__doc__ = PySparkDataFrame.limit.__doc__ + + def tail(self, num: int) -> List[Row]: + return DataFrame.withPlan( + plan.Tail(child=self._plan, limit=num), session=self._session + ).collect() + + tail.__doc__ = PySparkDataFrame.tail.__doc__ + + def _sort_cols( + self, cols: Sequence[Union[str, Column, List[Union[str, Column]]]], kwargs: Dict[str, Any] + ) -> List[Column]: + """Return a JVM Seq of Columns that describes the sort order""" + if cols is None: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "cols"}, + ) + + _cols: List[Column] = [] + if len(cols) == 1 and isinstance(cols[0], list): + _cols = [_to_col(c) for c in cols[0]] + else: + _cols = [_to_col(cast("ColumnOrName", c)) for c in cols] + + ascending = kwargs.get("ascending", True) + if isinstance(ascending, (bool, int)): + if not ascending: + _cols = [c.desc() for c in _cols] + elif isinstance(ascending, list): + _cols = [c if asc else c.desc() for asc, c in zip(ascending, _cols)] + else: + raise PySparkTypeError( + error_class="NOT_BOOL_OR_LIST", + message_parameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, + ) + + return _cols + + def sort( + self, *cols: Union[str, Column, List[Union[str, Column]]], **kwargs: Any + ) -> "DataFrame": + return DataFrame.withPlan( + plan.Sort( + self._plan, + columns=self._sort_cols(cols, kwargs), + is_global=True, + ), + session=self._session, + ) + + sort.__doc__ = PySparkDataFrame.sort.__doc__ + + orderBy = sort + + def sortWithinPartitions( + self, *cols: Union[str, Column, List[Union[str, Column]]], **kwargs: Any + ) -> "DataFrame": + return DataFrame.withPlan( + plan.Sort( + self._plan, + columns=self._sort_cols(cols, kwargs), + is_global=False, + ), + session=self._session, + ) + + sortWithinPartitions.__doc__ = PySparkDataFrame.sortWithinPartitions.__doc__ + + def sample( + self, + withReplacement: Optional[Union[float, bool]] = None, + fraction: Optional[Union[int, float]] = None, + seed: Optional[int] = None, + ) -> "DataFrame": + # For the cases below: + # sample(True, 0.5 [, seed]) + # sample(True, fraction=0.5 [, seed]) + # sample(withReplacement=False, fraction=0.5 [, seed]) + is_withReplacement_set = type(withReplacement) == bool and isinstance(fraction, float) + + # For the case below: + # sample(faction=0.5 [, seed]) + is_withReplacement_omitted_kwargs = withReplacement is None and isinstance(fraction, float) + + # For the case below: + # sample(0.5 [, seed]) + is_withReplacement_omitted_args = isinstance(withReplacement, float) + + if not ( + is_withReplacement_set + or is_withReplacement_omitted_kwargs + or is_withReplacement_omitted_args + ): + argtypes = [type(arg).__name__ for arg in [withReplacement, fraction, seed]] + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT", + message_parameters={ + "arg_name": "withReplacement (optional), " + + "fraction (required) and seed (optional)", + "arg_type": ", ".join(argtypes), + }, + ) + + if is_withReplacement_omitted_args: + if fraction is not None: + seed = cast(int, fraction) + fraction = withReplacement + withReplacement = None + + if withReplacement is None: + withReplacement = False + + seed = int(seed) if seed is not None else None + + return DataFrame.withPlan( + plan.Sample( + child=self._plan, + lower_bound=0.0, + upper_bound=fraction, # type: ignore[arg-type] + with_replacement=withReplacement, # type: ignore[arg-type] + seed=seed, + ), + session=self._session, + ) + + sample.__doc__ = PySparkDataFrame.sample.__doc__ + + def withColumnRenamed(self, existing: str, new: str) -> "DataFrame": + return self.withColumnsRenamed({existing: new}) + + withColumnRenamed.__doc__ = PySparkDataFrame.withColumnRenamed.__doc__ + + def withColumnsRenamed(self, colsMap: Dict[str, str]) -> "DataFrame": + if not isinstance(colsMap, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__}, + ) + + return DataFrame.withPlan(plan.WithColumnsRenamed(self._plan, colsMap), self._session) + + withColumnsRenamed.__doc__ = PySparkDataFrame.withColumnsRenamed.__doc__ + + def _show_string( + self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False + ) -> str: + if not isinstance(n, int) or isinstance(n, bool): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "n", "arg_type": type(n).__name__}, + ) + if not isinstance(vertical, bool): + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={"arg_name": "vertical", "arg_type": type(vertical).__name__}, + ) + + _truncate: int = -1 + if isinstance(truncate, bool) and truncate: + _truncate = 20 + else: + try: + _truncate = int(truncate) + except ValueError: + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={ + "arg_name": "truncate", + "arg_type": type(truncate).__name__, + }, + ) + + pdf = DataFrame.withPlan( + plan.ShowString(child=self._plan, num_rows=n, truncate=_truncate, vertical=vertical), + session=self._session, + ).toPandas() + assert pdf is not None + return pdf["show_string"][0] + + def withColumns(self, colsMap: Dict[str, Column]) -> "DataFrame": + if not isinstance(colsMap, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__}, + ) + + names: List[str] = [] + columns: List[Column] = [] + for columnName, column in colsMap.items(): + names.append(columnName) + columns.append(column) + + return DataFrame.withPlan( + plan.WithColumns( + self._plan, + columnNames=names, + columns=columns, + ), + session=self._session, + ) + + withColumns.__doc__ = PySparkDataFrame.withColumns.__doc__ + + def withColumn(self, colName: str, col: Column) -> "DataFrame": + if not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + return DataFrame.withPlan( + plan.WithColumns( + self._plan, + columnNames=[colName], + columns=[col], + ), + session=self._session, + ) + + withColumn.__doc__ = PySparkDataFrame.withColumn.__doc__ + + def withMetadata(self, columnName: str, metadata: Dict[str, Any]) -> "DataFrame": + if not isinstance(metadata, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "metadata", "arg_type": type(metadata).__name__}, + ) + + return DataFrame.withPlan( + plan.WithColumns( + self._plan, + columnNames=[columnName], + columns=[self[columnName]], + metadata=[json.dumps(metadata)], + ), + session=self._session, + ) + + withMetadata.__doc__ = PySparkDataFrame.withMetadata.__doc__ + + def unpivot( + self, + ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]], + values: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]], + variableColumnName: str, + valueColumnName: str, + ) -> "DataFrame": + assert ids is not None, "ids must not be None" + + def to_jcols( + cols: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]] + ) -> List["ColumnOrName"]: + if cols is None: + lst = [] + elif isinstance(cols, tuple): + lst = list(cols) + elif isinstance(cols, list): + lst = cols + else: + lst = [cols] + return lst + + return DataFrame.withPlan( + plan.Unpivot( + self._plan, + to_jcols(ids), + to_jcols(values) if values is not None else None, + variableColumnName, + valueColumnName, + ), + self._session, + ) + + unpivot.__doc__ = PySparkDataFrame.unpivot.__doc__ + + melt = unpivot + + def withWatermark(self, eventTime: str, delayThreshold: str) -> "DataFrame": + # TODO: reuse error handling code in sql.DataFrame.withWatermark() + if not eventTime or type(eventTime) is not str: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "eventTime", "arg_type": type(eventTime).__name__}, + ) + if not delayThreshold or type(delayThreshold) is not str: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "delayThreshold", + "arg_type": type(delayThreshold).__name__, + }, + ) + + return DataFrame.withPlan( + plan.WithWatermark( + self._plan, + event_time=eventTime, + delay_threshold=delayThreshold, + ), + session=self._session, + ) + + withWatermark.__doc__ = PySparkDataFrame.withWatermark.__doc__ + + def hint( + self, name: str, *parameters: Union["PrimitiveType", List["PrimitiveType"]] + ) -> "DataFrame": + if len(parameters) == 1 and isinstance(parameters[0], list): + parameters = parameters[0] # type: ignore[assignment] + + if not isinstance(name, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "name", "arg_type": type(name).__name__}, + ) + + allowed_types = (str, list, float, int) + for p in parameters: + if not isinstance(p, allowed_types): + raise PySparkTypeError( + error_class="INVALID_ITEM_FOR_CONTAINER", + message_parameters={ + "arg_name": "parameters", + "allowed_types": ", ".join([t.__name__ for t in allowed_types]), + "item_type": type(p).__name__, + }, + ) + + return DataFrame.withPlan( + plan.Hint(self._plan, name, list(parameters)), + session=self._session, + ) + + hint.__doc__ = PySparkDataFrame.hint.__doc__ + + def randomSplit( + self, + weights: List[float], + seed: Optional[int] = None, + ) -> List["DataFrame"]: + for w in weights: + if w < 0.0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "weights", "arg_value": str(w)}, + ) + seed = seed if seed is not None else random.randint(0, sys.maxsize) + total = sum(weights) + if total <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "sum(weights)", "arg_value": str(total)}, + ) + proportions = list(map(lambda x: x / total, weights)) + normalizedCumWeights = [0.0] + for v in proportions: + normalizedCumWeights.append(normalizedCumWeights[-1] + v) + j = 1 + length = len(normalizedCumWeights) + splits = [] + while j < length: + lowerBound = normalizedCumWeights[j - 1] + upperBound = normalizedCumWeights[j] + samplePlan = DataFrame.withPlan( + plan.Sample( + child=self._plan, + lower_bound=lowerBound, + upper_bound=upperBound, + with_replacement=False, + seed=int(seed), + deterministic_order=True, + ), + session=self._session, + ) + splits.append(samplePlan) + j += 1 + + return splits + + randomSplit.__doc__ = PySparkDataFrame.randomSplit.__doc__ + + def observe( + self, + observation: Union["Observation", str], + *exprs: Column, + ) -> "DataFrame": + if len(exprs) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "exprs"}, + ) + if not all(isinstance(c, Column) for c in exprs): + raise PySparkTypeError( + error_class="NOT_LIST_OF_COLUMN", + message_parameters={"arg_name": "exprs"}, + ) + + if isinstance(observation, Observation): + return DataFrame.withPlan( + plan.CollectMetrics(self._plan, str(observation._name), list(exprs)), + self._session, + ) + elif isinstance(observation, str): + return DataFrame.withPlan( + plan.CollectMetrics(self._plan, observation, list(exprs)), + self._session, + ) + else: + raise PySparkTypeError( + error_class="NOT_OBSERVATION_OR_STR", + message_parameters={ + "arg_name": "observation", + "arg_type": type(observation).__name__, + }, + ) + + observe.__doc__ = PySparkDataFrame.observe.__doc__ + + def show(self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False) -> None: + print(self._show_string(n, truncate, vertical)) + + show.__doc__ = PySparkDataFrame.show.__doc__ + + def union(self, other: "DataFrame") -> "DataFrame": + return self.unionAll(other) + + union.__doc__ = PySparkDataFrame.union.__doc__ + + def unionAll(self, other: "DataFrame") -> "DataFrame": + if other._plan is None: + raise PySparkValueError( + error_class="MISSING_VALID_PLAN", + message_parameters={"operator": "Union"}, + ) + self.checkSameSparkSession(other) + return DataFrame.withPlan( + plan.SetOperation(self._plan, other._plan, "union", is_all=True), session=self._session + ) + + unionAll.__doc__ = PySparkDataFrame.unionAll.__doc__ + + def unionByName(self, other: "DataFrame", allowMissingColumns: bool = False) -> "DataFrame": + if other._plan is None: + raise PySparkValueError( + error_class="MISSING_VALID_PLAN", + message_parameters={"operator": "UnionByName"}, + ) + self.checkSameSparkSession(other) + return DataFrame.withPlan( + plan.SetOperation( + self._plan, + other._plan, + "union", + by_name=True, + allow_missing_columns=allowMissingColumns, + ), + session=self._session, + ) + + unionByName.__doc__ = PySparkDataFrame.unionByName.__doc__ + + def subtract(self, other: "DataFrame") -> "DataFrame": + return DataFrame.withPlan( + plan.SetOperation(self._plan, other._plan, "except", is_all=False), + session=self._session, + ) + + subtract.__doc__ = PySparkDataFrame.subtract.__doc__ + + def exceptAll(self, other: "DataFrame") -> "DataFrame": + return DataFrame.withPlan( + plan.SetOperation(self._plan, other._plan, "except", is_all=True), session=self._session + ) + + exceptAll.__doc__ = PySparkDataFrame.exceptAll.__doc__ + + def intersect(self, other: "DataFrame") -> "DataFrame": + return DataFrame.withPlan( + plan.SetOperation(self._plan, other._plan, "intersect", is_all=False), + session=self._session, + ) + + intersect.__doc__ = PySparkDataFrame.intersect.__doc__ + + def intersectAll(self, other: "DataFrame") -> "DataFrame": + return DataFrame.withPlan( + plan.SetOperation(self._plan, other._plan, "intersect", is_all=True), + session=self._session, + ) + + intersectAll.__doc__ = PySparkDataFrame.intersectAll.__doc__ + + def where(self, condition: Union[Column, str]) -> "DataFrame": + if not isinstance(condition, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + return self.filter(condition) + + where.__doc__ = PySparkDataFrame.where.__doc__ + + @property + def na(self) -> "DataFrameNaFunctions": + return DataFrameNaFunctions(self) + + na.__doc__ = PySparkDataFrame.na.__doc__ + + def fillna( + self, + value: Union["LiteralType", Dict[str, "LiteralType"]], + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> "DataFrame": + if not isinstance(value, (float, int, str, bool, dict)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_STR", + message_parameters={"arg_name": "value", "arg_type": type(value).__name__}, + ) + if isinstance(value, dict): + if len(value) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "value"}, + ) + for c, v in value.items(): + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "key type of dict", + "arg_type": type(c).__name__, + }, + ) + if not isinstance(v, (bool, int, float, str)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT_OR_STR", + message_parameters={ + "arg_name": "value type of dict", + "arg_type": type(v).__name__, + }, + ) + + _cols: List[str] = [] + if subset is not None: + if isinstance(subset, str): + _cols = [subset] + elif isinstance(subset, (tuple, list)): + for c in subset: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "cols", "arg_type": type(c).__name__}, + ) + _cols = list(subset) + else: + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if isinstance(value, dict): + _cols = list(value.keys()) + _values = [value[c] for c in _cols] + else: + _values = [value] + + return DataFrame.withPlan( + plan.NAFill(child=self._plan, cols=_cols, values=_values), + session=self._session, + ) + + fillna.__doc__ = PySparkDataFrame.fillna.__doc__ + + def dropna( + self, + how: str = "any", + thresh: Optional[int] = None, + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> "DataFrame": + min_non_nulls: Optional[int] = None + + if how is not None: + if not isinstance(how, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "how", "arg_type": type(how).__name__}, + ) + if how == "all": + min_non_nulls = 1 + elif how == "any": + min_non_nulls = None + else: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"arg_name": "how", "arg_value": str(how)}, + ) + + if thresh is not None: + if not isinstance(thresh, int): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "thresh", "arg_type": type(thresh).__name__}, + ) + + # 'thresh' overwrites 'how' + min_non_nulls = thresh + + _cols: List[str] = [] + if subset is not None: + if isinstance(subset, str): + _cols = [subset] + elif isinstance(subset, (tuple, list)): + for c in subset: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "cols", "arg_type": type(c).__name__}, + ) + _cols = list(subset) + else: + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + return DataFrame.withPlan( + plan.NADrop(child=self._plan, cols=_cols, min_non_nulls=min_non_nulls), + session=self._session, + ) + + dropna.__doc__ = PySparkDataFrame.dropna.__doc__ + + def replace( + self, + to_replace: Union[ + "LiteralType", List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"] + ], + value: Optional[ + Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType] + ] = _NoValue, + subset: Optional[List[str]] = None, + ) -> "DataFrame": + if value is _NoValue: + if isinstance(to_replace, dict): + value = None + else: + raise PySparkTypeError( + error_class="ARGUMENT_REQUIRED", + message_parameters={"arg_name": "value", "condition": "`to_replace` is dict"}, + ) + + # Helper functions + def all_of(types: Union[Type, Tuple[Type, ...]]) -> Callable[[Iterable], bool]: + """Given a type or tuple of types and a sequence of xs + check if each x is instance of type(s) + + >>> all_of(bool)([True, False]) + True + >>> all_of(str)(["a", 1]) + False + """ + + def all_of_(xs: Iterable) -> bool: + return all(isinstance(x, types) for x in xs) + + return all_of_ + + all_of_bool = all_of(bool) + all_of_str = all_of(str) + all_of_numeric = all_of((float, int)) + + # Validate input types + valid_types = (bool, float, int, str, list, tuple) + if not isinstance(to_replace, valid_types + (dict,)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "to_replace", + "arg_type": type(to_replace).__name__, + }, + ) + + if ( + not isinstance(value, valid_types) + and value is not None + and not isinstance(to_replace, dict) + ): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT_OR_LIST_OR_NONE_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "value", "arg_type": type(value).__name__}, + ) + + if isinstance(to_replace, (list, tuple)) and isinstance(value, (list, tuple)): + if len(to_replace) != len(value): + raise PySparkValueError( + error_class="LENGTH_SHOULD_BE_THE_SAME", + message_parameters={ + "arg1": "to_replace", + "arg2": "value", + "arg1_length": str(len(to_replace)), + "arg2_length": str(len(value)), + }, + ) + + if not (subset is None or isinstance(subset, (list, tuple, str))): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + # Reshape input arguments if necessary + if isinstance(to_replace, (float, int, str)): + to_replace = [to_replace] + + if isinstance(to_replace, dict): + rep_dict = to_replace + if value is not None: + warnings.warn("to_replace is a dict and value is not None. value will be ignored.") + else: + if isinstance(value, (float, int, str)) or value is None: + value = [value for _ in range(len(to_replace))] + rep_dict = dict(zip(to_replace, cast("Iterable[Optional[Union[float, str]]]", value))) + + if isinstance(subset, str): + subset = [subset] + + # Verify we were not passed in mixed type generics. + if not any( + all_of_type(rep_dict.keys()) + and all_of_type(x for x in rep_dict.values() if x is not None) + for all_of_type in [all_of_bool, all_of_str, all_of_numeric] + ): + raise PySparkValueError( + error_class="MIXED_TYPE_REPLACEMENT", + message_parameters={}, + ) + + return DataFrame.withPlan( + plan.NAReplace(child=self._plan, cols=subset, replacements=rep_dict), + session=self._session, + ) + + replace.__doc__ = PySparkDataFrame.replace.__doc__ + + @property + def stat(self) -> "DataFrameStatFunctions": + return DataFrameStatFunctions(self) + + stat.__doc__ = PySparkDataFrame.stat.__doc__ + + def summary(self, *statistics: str) -> "DataFrame": + _statistics: List[str] = list(statistics) + for s in _statistics: + if not isinstance(s, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={"arg_name": "statistics", "arg_type": type(s).__name__}, + ) + return DataFrame.withPlan( + plan.StatSummary(child=self._plan, statistics=_statistics), + session=self._session, + ) + + summary.__doc__ = PySparkDataFrame.summary.__doc__ + + def describe(self, *cols: Union[str, List[str]]) -> "DataFrame": + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] # type: ignore[assignment] + + _cols = [] + for column in cols: + if isinstance(column, str): + _cols.append(column) + else: + _cols.extend([s for s in column]) + return DataFrame.withPlan( + plan.StatDescribe(child=self._plan, cols=_cols), + session=self._session, + ) + + describe.__doc__ = PySparkDataFrame.describe.__doc__ + + def cov(self, col1: str, col2: str) -> float: + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + pdf = DataFrame.withPlan( + plan.StatCov(child=self._plan, col1=col1, col2=col2), + session=self._session, + ).toPandas() + + assert pdf is not None + return pdf["cov"][0] + + cov.__doc__ = PySparkDataFrame.cov.__doc__ + + def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float: + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + if not method: + method = "pearson" + if not method == "pearson": + raise PySparkValueError( + error_class="VALUE_NOT_PEARSON", + message_parameters={"arg_name": "method", "arg_value": method}, + ) + pdf = DataFrame.withPlan( + plan.StatCorr(child=self._plan, col1=col1, col2=col2, method=method), + session=self._session, + ).toPandas() + + assert pdf is not None + return pdf["corr"][0] + + corr.__doc__ = PySparkDataFrame.corr.__doc__ + + def approxQuantile( + self, + col: Union[str, List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> Union[List[float], List[List[float]]]: + if not isinstance(col, (str, list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + + isStr = isinstance(col, str) + + if isinstance(col, tuple): + col = list(col) + elif isStr: + col = [cast(str, col)] + + for c in col: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={"arg_name": "columns", "arg_type": type(c).__name__}, + ) + + if not isinstance(probabilities, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={ + "arg_name": "probabilities", + "arg_type": type(probabilities).__name__, + }, + ) + if isinstance(probabilities, tuple): + probabilities = list(probabilities) + for p in probabilities: + if not isinstance(p, (float, int)) or p < 0 or p > 1: + raise PySparkTypeError( + error_class="NOT_LIST_OF_FLOAT_OR_INT", + message_parameters={ + "arg_name": "probabilities", + "arg_type": type(p).__name__, + }, + ) + + if not isinstance(relativeError, (float, int)): + raise PySparkTypeError( + error_class="NOT_FLOAT_OR_INT", + message_parameters={ + "arg_name": "relativeError", + "arg_type": type(relativeError).__name__, + }, + ) + if relativeError < 0: + raise PySparkValueError( + error_class="NEGATIVE_VALUE", + message_parameters={ + "arg_name": "relativeError", + "arg_value": str(relativeError), + }, + ) + relativeError = float(relativeError) + pdf = DataFrame.withPlan( + plan.StatApproxQuantile( + child=self._plan, + cols=list(col), + probabilities=probabilities, + relativeError=relativeError, + ), + session=self._session, + ).toPandas() + + assert pdf is not None + jaq = pdf["approx_quantile"][0] + jaq_list = [list(j) for j in jaq] + return jaq_list[0] if isStr else jaq_list + + approxQuantile.__doc__ = PySparkDataFrame.approxQuantile.__doc__ + + def crosstab(self, col1: str, col2: str) -> "DataFrame": + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + return DataFrame.withPlan( + plan.StatCrosstab(child=self._plan, col1=col1, col2=col2), + session=self._session, + ) + + crosstab.__doc__ = PySparkDataFrame.crosstab.__doc__ + + def freqItems( + self, cols: Union[List[str], Tuple[str]], support: Optional[float] = None + ) -> "DataFrame": + if isinstance(cols, tuple): + cols = list(cols) + if not isinstance(cols, list): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "cols", "arg_type": type(cols).__name__}, + ) + if not support: + support = 0.01 + return DataFrame.withPlan( + plan.StatFreqItems(child=self._plan, cols=cols, support=support), + session=self._session, + ) + + freqItems.__doc__ = PySparkDataFrame.freqItems.__doc__ + + def sampleBy( + self, col: "ColumnOrName", fractions: Dict[Any, float], seed: Optional[int] = None + ) -> "DataFrame": + from pyspark.sql.connect.expressions import ColumnReference + + if isinstance(col, str): + col = Column(ColumnReference(col)) + elif not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + if not isinstance(fractions, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "fractions", "arg_type": type(fractions).__name__}, + ) + for k, v in fractions.items(): + if not isinstance(k, (float, int, str)): + raise PySparkTypeError( + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "fractions", + "arg_type": type(fractions).__name__, + "allowed_types": "float, int, str", + "return_type": type(k).__name__, + }, + ) + fractions[k] = float(v) + seed = seed if seed is not None else random.randint(0, sys.maxsize) + return DataFrame.withPlan( + plan.StatSampleBy(child=self._plan, col=col, fractions=fractions, seed=seed), + session=self._session, + ) + + sampleBy.__doc__ = PySparkDataFrame.sampleBy.__doc__ + + def _get_alias(self) -> Optional[str]: + p = self._plan + while p is not None: + if isinstance(p, plan.Project) and p.alias: + return p.alias + p = p._child + return None + + def __getattr__(self, name: str) -> "Column": + if name in ["_jseq", "_jdf", "_jmap", "_jcols"]: + raise PySparkAttributeError( + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", message_parameters={"attr_name": name} + ) + elif name in [ + "rdd", + "toJSON", + "foreach", + "foreachPartition", + "checkpoint", + "localCheckpoint", + ]: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": f"{name}()"}, + ) + + if name not in self.columns: + raise AttributeError( + "'%s' object has no attribute '%s'" % (self.__class__.__name__, name) + ) + + return self[name] + + __getattr__.__doc__ = PySparkDataFrame.__getattr__.__doc__ + + @overload + def __getitem__(self, item: Union[int, str]) -> Column: + ... + + @overload + def __getitem__(self, item: Union[Column, List, Tuple]) -> "DataFrame": + ... + + def __getitem__(self, item: Union[int, str, Column, List, Tuple]) -> Union[Column, "DataFrame"]: + if isinstance(item, str): + # Check for alias + alias = self._get_alias() + if self._plan is None: + raise SparkConnectException("Cannot analyze on empty plan.") + return _to_col_with_plan_id( + col=alias if alias is not None else item, + plan_id=self._plan._plan_id, + ) + elif isinstance(item, Column): + return self.filter(item) + elif isinstance(item, (list, tuple)): + return self.select(*item) + elif isinstance(item, int): + return col(self.columns[item]) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "item", "arg_type": type(item).__name__}, + ) + + def __dir__(self) -> List[str]: + attrs = set(super().__dir__()) + attrs.update(self.columns) + return sorted(attrs) + + __dir__.__doc__ = PySparkDataFrame.__dir__.__doc__ + + def _print_plan(self) -> str: + if self._plan: + return self._plan.print() + return "" + + def collect(self) -> List[Row]: + if self._plan is None: + raise Exception("Cannot collect on empty plan.") + if self._session is None: + raise Exception("Cannot collect on empty session.") + query = self._plan.to_proto(self._session.client) + table, schema = self._session.client.to_table(query) + + schema = schema or from_arrow_schema(table.schema, prefer_timestamp_ntz=True) + + assert schema is not None and isinstance(schema, StructType) + + from pyspark.sql.connect.conversion import ArrowTableToRowsConversion + + return ArrowTableToRowsConversion.convert(table, schema) + + collect.__doc__ = PySparkDataFrame.collect.__doc__ + + def toPandas(self) -> "pandas.DataFrame": + if self._plan is None: + raise Exception("Cannot collect on empty plan.") + if self._session is None: + raise Exception("Cannot collect on empty session.") + query = self._plan.to_proto(self._session.client) + return self._session.client.to_pandas(query) + + toPandas.__doc__ = PySparkDataFrame.toPandas.__doc__ + + @property + def schema(self) -> StructType: + if self._plan is not None: + query = self._plan.to_proto(self._session.client) + if self._session is None: + raise Exception("Cannot analyze without SparkSession.") + return self._session.client.schema(query) + else: + raise Exception("Empty plan.") + + schema.__doc__ = PySparkDataFrame.schema.__doc__ + + def isLocal(self) -> bool: + if self._plan is None: + raise Exception("Cannot analyze on empty plan.") + query = self._plan.to_proto(self._session.client) + result = self._session.client._analyze(method="is_local", plan=query).is_local + assert result is not None + return result + + isLocal.__doc__ = PySparkDataFrame.isLocal.__doc__ + + @property + def isStreaming(self) -> bool: + if self._plan is None: + raise Exception("Cannot analyze on empty plan.") + query = self._plan.to_proto(self._session.client) + result = self._session.client._analyze(method="is_streaming", plan=query).is_streaming + assert result is not None + return result + + isStreaming.__doc__ = PySparkDataFrame.isStreaming.__doc__ + + def _tree_string(self, level: Optional[int] = None) -> str: + if self._plan is None: + raise Exception("Cannot analyze on empty plan.") + query = self._plan.to_proto(self._session.client) + result = self._session.client._analyze( + method="tree_string", plan=query, level=level + ).tree_string + assert result is not None + return result + + def printSchema(self, level: Optional[int] = None) -> None: + print(self._tree_string(level)) + + printSchema.__doc__ = PySparkDataFrame.printSchema.__doc__ + + def inputFiles(self) -> List[str]: + if self._plan is None: + raise Exception("Cannot analyze on empty plan.") + query = self._plan.to_proto(self._session.client) + result = self._session.client._analyze(method="input_files", plan=query).input_files + assert result is not None + return result + + inputFiles.__doc__ = PySparkDataFrame.inputFiles.__doc__ + + def to(self, schema: StructType) -> "DataFrame": + assert schema is not None + return DataFrame.withPlan( + plan.ToSchema(child=self._plan, schema=schema), + session=self._session, + ) + + to.__doc__ = PySparkDataFrame.to.__doc__ + + def toDF(self, *cols: str) -> "DataFrame": + for col_ in cols: + if not isinstance(col_, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={"arg_name": "cols", "arg_type": type(col_).__name__}, + ) + return DataFrame.withPlan(plan.ToDF(self._plan, list(cols)), self._session) + + toDF.__doc__ = PySparkDataFrame.toDF.__doc__ + + def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any) -> "DataFrame": + result = func(self, *args, **kwargs) + assert isinstance( + result, DataFrame + ), "Func returned an instance of type [%s], " "should have been DataFrame." % type(result) + return result + + transform.__doc__ = PySparkDataFrame.transform.__doc__ + + def _explain_string( + self, extended: Optional[Union[bool, str]] = None, mode: Optional[str] = None + ) -> str: + if extended is not None and mode is not None: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={"arg_list": "extended and mode"}, + ) + + # For the no argument case: df.explain() + is_no_argument = extended is None and mode is None + + # For the cases below: + # explain(True) + # explain(extended=False) + is_extended_case = isinstance(extended, bool) and mode is None + + # For the case when extended is mode: + # df.explain("formatted") + is_extended_as_mode = isinstance(extended, str) and mode is None + + # For the mode specified: + # df.explain(mode="formatted") + is_mode_case = extended is None and isinstance(mode, str) + + if not (is_no_argument or is_extended_case or is_extended_as_mode or is_mode_case): + argtypes = [str(type(arg)) for arg in [extended, mode] if arg is not None] + raise PySparkTypeError( + error_class="NOT_BOOL_OR_STR", + message_parameters={ + "arg_name": "extended (optional) and mode (optional)", + "arg_type": ", ".join(argtypes), + }, + ) + + # Sets an explain mode depending on a given argument + if is_no_argument: + explain_mode = "simple" + elif is_extended_case: + explain_mode = "extended" if extended else "simple" + elif is_mode_case: + explain_mode = cast(str, mode) + elif is_extended_as_mode: + explain_mode = cast(str, extended) + + if self._plan is not None: + query = self._plan.to_proto(self._session.client) + if self._session is None: + raise Exception("Cannot analyze without SparkSession.") + return self._session.client.explain_string(query, explain_mode) + else: + return "" + + def explain( + self, extended: Optional[Union[bool, str]] = None, mode: Optional[str] = None + ) -> None: + print(self._explain_string(extended=extended, mode=mode)) + + explain.__doc__ = PySparkDataFrame.explain.__doc__ + + def createTempView(self, name: str) -> None: + command = plan.CreateView( + child=self._plan, name=name, is_global=False, replace=False + ).command(session=self._session.client) + self._session.client.execute_command(command) + + createTempView.__doc__ = PySparkDataFrame.createTempView.__doc__ + + def createOrReplaceTempView(self, name: str) -> None: + command = plan.CreateView( + child=self._plan, name=name, is_global=False, replace=True + ).command(session=self._session.client) + self._session.client.execute_command(command) + + createOrReplaceTempView.__doc__ = PySparkDataFrame.createOrReplaceTempView.__doc__ + + def createGlobalTempView(self, name: str) -> None: + command = plan.CreateView( + child=self._plan, name=name, is_global=True, replace=False + ).command(session=self._session.client) + self._session.client.execute_command(command) + + createGlobalTempView.__doc__ = PySparkDataFrame.createGlobalTempView.__doc__ + + def createOrReplaceGlobalTempView(self, name: str) -> None: + command = plan.CreateView( + child=self._plan, name=name, is_global=True, replace=True + ).command(session=self._session.client) + self._session.client.execute_command(command) + + createOrReplaceGlobalTempView.__doc__ = PySparkDataFrame.createOrReplaceGlobalTempView.__doc__ + + def cache(self) -> "DataFrame": + if self._plan is None: + raise Exception("Cannot cache on empty plan.") + return self.persist() + + cache.__doc__ = PySparkDataFrame.cache.__doc__ + + def persist( + self, + storageLevel: StorageLevel = (StorageLevel.MEMORY_AND_DISK_DESER), + ) -> "DataFrame": + if self._plan is None: + raise Exception("Cannot persist on empty plan.") + relation = self._plan.plan(self._session.client) + self._session.client._analyze( + method="persist", relation=relation, storage_level=storageLevel + ) + return self + + persist.__doc__ = PySparkDataFrame.persist.__doc__ + + @property + def storageLevel(self) -> StorageLevel: + if self._plan is None: + raise Exception("Cannot persist on empty plan.") + relation = self._plan.plan(self._session.client) + storage_level = self._session.client._analyze( + method="get_storage_level", relation=relation + ).storage_level + assert storage_level is not None + return storage_level + + storageLevel.__doc__ = PySparkDataFrame.storageLevel.__doc__ + + def unpersist(self, blocking: bool = False) -> "DataFrame": + if self._plan is None: + raise Exception("Cannot unpersist on empty plan.") + relation = self._plan.plan(self._session.client) + self._session.client._analyze(method="unpersist", relation=relation, blocking=blocking) + return self + + unpersist.__doc__ = PySparkDataFrame.unpivot.__doc__ + + @property + def is_cached(self) -> bool: + return self.storageLevel != StorageLevel.NONE + + def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]: + from pyspark.sql.connect.conversion import ArrowTableToRowsConversion + + if self._plan is None: + raise Exception("Cannot collect on empty plan.") + if self._session is None: + raise Exception("Cannot collect on empty session.") + query = self._plan.to_proto(self._session.client) + + schema: Optional[StructType] = None + for schema_or_table in self._session.client.to_table_as_iterator(query): + if isinstance(schema_or_table, StructType): + assert schema is None + schema = schema_or_table + else: + assert isinstance(schema_or_table, pa.Table) + table = schema_or_table + if schema is None: + schema = from_arrow_schema(table.schema, prefer_timestamp_ntz=True) + yield from ArrowTableToRowsConversion.convert(table, schema) + + toLocalIterator.__doc__ = PySparkDataFrame.toLocalIterator.__doc__ + + def to_pandas_on_spark( + self, index_col: Optional[Union[str, List[str]]] = None + ) -> "PandasOnSparkDataFrame": + warnings.warn( + "DataFrame.to_pandas_on_spark is deprecated. Use DataFrame.pandas_api instead.", + FutureWarning, + ) + return self.pandas_api(index_col) + + def pandas_api( + self, index_col: Optional[Union[str, List[str]]] = None + ) -> "PandasOnSparkDataFrame": + from pyspark.pandas.namespace import _get_index_map + from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame + from pyspark.pandas.internal import InternalFrame + + index_spark_columns, index_names = _get_index_map(self, index_col) # type: ignore[arg-type] + internal = InternalFrame( + spark_frame=self, # type: ignore[arg-type] + index_spark_columns=index_spark_columns, + index_names=index_names, # type: ignore[arg-type] + ) + return PandasOnSparkDataFrame(internal) + + pandas_api.__doc__ = PySparkDataFrame.pandas_api.__doc__ + + def registerTempTable(self, name: str) -> None: + warnings.warn("Deprecated in 2.0, use createOrReplaceTempView instead.", FutureWarning) + self.createOrReplaceTempView(name) + + registerTempTable.__doc__ = PySparkDataFrame.registerTempTable.__doc__ + + def _map_partitions( + self, + func: "PandasMapIterFunction", + schema: Union[StructType, str], + evalType: int, + barrier: bool, + ) -> "DataFrame": + from pyspark.sql.connect.udf import UserDefinedFunction + + if self._plan is None: + raise Exception("Cannot mapInPandas when self._plan is empty.") + + udf_obj = UserDefinedFunction( + func, + returnType=schema, + evalType=evalType, + ) + + return DataFrame.withPlan( + plan.MapPartitions( + child=self._plan, function=udf_obj, cols=self.columns, is_barrier=barrier + ), + session=self._session, + ) + + def mapInPandas( + self, + func: "PandasMapIterFunction", + schema: Union[StructType, str], + barrier: bool = False, + ) -> "DataFrame": + return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, barrier) + + mapInPandas.__doc__ = PySparkDataFrame.mapInPandas.__doc__ + + def mapInArrow( + self, + func: "ArrowMapIterFunction", + schema: Union[StructType, str], + barrier: bool = False, + ) -> "DataFrame": + return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, barrier) + + mapInArrow.__doc__ = PySparkDataFrame.mapInArrow.__doc__ + + @property + def writeStream(self) -> DataStreamWriter: + assert self._plan is not None + return DataStreamWriter(plan=self._plan, session=self._session) + + writeStream.__doc__ = PySparkDataFrame.writeStream.__doc__ + + def sameSemantics(self, other: "DataFrame") -> bool: + assert self._plan is not None + assert other._plan is not None + return self._session.client.same_semantics( + plan=self._plan.to_proto(self._session.client), + other=other._plan.to_proto(other._session.client), + ) + + sameSemantics.__doc__ = PySparkDataFrame.sameSemantics.__doc__ + + def semanticHash(self) -> int: + assert self._plan is not None + return self._session.client.semantic_hash( + plan=self._plan.to_proto(self._session.client), + ) + + semanticHash.__doc__ = PySparkDataFrame.semanticHash.__doc__ + + def writeTo(self, table: str) -> "DataFrameWriterV2": + assert self._plan is not None + return DataFrameWriterV2(self._plan, self._session, table) + + writeTo.__doc__ = PySparkDataFrame.writeTo.__doc__ + + # SparkConnect specific API + def offset(self, n: int) -> "DataFrame": + return DataFrame.withPlan(plan.Offset(child=self._plan, offset=n), session=self._session) + + offset.__doc__ = PySparkDataFrame.offset.__doc__ + + @classmethod + def withPlan(cls, plan: plan.LogicalPlan, session: "SparkSession") -> "DataFrame": + """ + Main initialization method used to construct a new data frame with a child plan. + This is for internal purpose. + """ + new_frame = DataFrame(session=session) + new_frame._plan = plan + return new_frame + + +class DataFrameNaFunctions: + def __init__(self, df: DataFrame): + self.df = df + + def fill( + self, + value: Union["LiteralType", Dict[str, "LiteralType"]], + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> DataFrame: + return self.df.fillna(value=value, subset=subset) + + fill.__doc__ = DataFrame.fillna.__doc__ + + def drop( + self, + how: str = "any", + thresh: Optional[int] = None, + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> DataFrame: + return self.df.dropna(how=how, thresh=thresh, subset=subset) + + drop.__doc__ = DataFrame.dropna.__doc__ + + def replace( + self, + to_replace: Union[List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]], + value: Optional[ + Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType] + ] = _NoValue, + subset: Optional[List[str]] = None, + ) -> DataFrame: + return self.df.replace(to_replace, value, subset) + + replace.__doc__ = DataFrame.replace.__doc__ + + +DataFrameNaFunctions.__doc__ = PySparkDataFrameNaFunctions.__doc__ + + +class DataFrameStatFunctions: + def __init__(self, df: DataFrame): + self.df = df + + def cov(self, col1: str, col2: str) -> float: + return self.df.cov(col1, col2) + + cov.__doc__ = DataFrame.cov.__doc__ + + def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float: + return self.df.corr(col1, col2, method) + + corr.__doc__ = DataFrame.corr.__doc__ + + def approxQuantile( + self, + col: Union[str, List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> Union[List[float], List[List[float]]]: + return self.df.approxQuantile(col, probabilities, relativeError) + + approxQuantile.__doc__ = DataFrame.approxQuantile.__doc__ + + def crosstab(self, col1: str, col2: str) -> DataFrame: + return self.df.crosstab(col1, col2) + + crosstab.__doc__ = DataFrame.crosstab.__doc__ + + def freqItems( + self, cols: Union[List[str], Tuple[str]], support: Optional[float] = None + ) -> DataFrame: + return self.df.freqItems(cols, support) + + freqItems.__doc__ = DataFrame.freqItems.__doc__ + + def sampleBy( + self, col: str, fractions: Dict[Any, float], seed: Optional[int] = None + ) -> DataFrame: + return self.df.sampleBy(col, fractions, seed) + + sampleBy.__doc__ = DataFrame.sampleBy.__doc__ + + +DataFrameStatFunctions.__doc__ = PySparkDataFrameStatFunctions.__doc__ + + +def _test() -> None: + import os + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.dataframe + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.connect.dataframe.__dict__.copy() + # Spark Connect does not support RDD but the tests depend on them. + del pyspark.sql.connect.dataframe.DataFrame.coalesce.__doc__ + del pyspark.sql.connect.dataframe.DataFrame.repartition.__doc__ + del pyspark.sql.connect.dataframe.DataFrame.repartitionByRange.__doc__ + + # TODO(SPARK-41625): Support Structured Streaming + del pyspark.sql.connect.dataframe.DataFrame.isStreaming.__doc__ + + # TODO(SPARK-41888): Support StreamingQueryListener for DataFrame.observe + del pyspark.sql.connect.dataframe.DataFrame.observe.__doc__ + + # TODO(SPARK-43435): should reenable this test + del pyspark.sql.connect.dataframe.DataFrame.writeStream.__doc__ + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.dataframe tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.dataframe, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/expressions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/expressions.py new file mode 100644 index 0000000..d0a9b1d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/expressions.py @@ -0,0 +1,1056 @@ +# +# 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.sql.connect.utils import check_dependencies +from pyspark.sql.utils import is_timestamp_ntz_preferred + +check_dependencies(__name__) + +from typing import ( + cast, + TYPE_CHECKING, + Any, + Callable, + Union, + Sequence, + Tuple, + Optional, +) + +import json +import decimal +import datetime +import warnings +from threading import Lock + +import numpy as np + +from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.types import ( + _from_numpy_type, + DateType, + ArrayType, + NullType, + BooleanType, + BinaryType, + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + DecimalType, + StringType, + DataType, + TimestampType, + TimestampNTZType, + DayTimeIntervalType, +) + +import pyspark.sql.connect.proto as proto +from pyspark.sql.connect.types import ( + JVM_BYTE_MIN, + JVM_BYTE_MAX, + JVM_SHORT_MIN, + JVM_SHORT_MAX, + JVM_INT_MIN, + JVM_INT_MAX, + JVM_LONG_MIN, + JVM_LONG_MAX, + UnparsedDataType, + pyspark_types_to_proto_types, + proto_schema_to_pyspark_data_type, +) +from pyspark.errors import PySparkTypeError, PySparkValueError + +if TYPE_CHECKING: + from pyspark.sql.connect.client import SparkConnectClient + from pyspark.sql.connect.window import WindowSpec + + +class Expression: + """ + Expression base class. + """ + + def __init__(self) -> None: + pass + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + ... + + def __repr__(self) -> str: + ... + + def alias(self, *alias: str, **kwargs: Any) -> "ColumnAlias": + metadata = kwargs.pop("metadata", None) + assert not kwargs, "Unexpected kwargs where passed: %s" % kwargs + return ColumnAlias(self, list(alias), metadata) + + def name(self) -> str: + ... + + +class CaseWhen(Expression): + def __init__( + self, branches: Sequence[Tuple[Expression, Expression]], else_value: Optional[Expression] + ): + super().__init__() + + assert isinstance(branches, list) + for branch in branches: + assert ( + isinstance(branch, tuple) + and len(branch) == 2 + and all(isinstance(expr, Expression) for expr in branch) + ) + self._branches = branches + + if else_value is not None: + assert isinstance(else_value, Expression) + + self._else_value = else_value + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + args = [] + for condition, value in self._branches: + args.append(condition) + args.append(value) + + if self._else_value is not None: + args.append(self._else_value) + + unresolved_function = UnresolvedFunction(name="when", args=args) + + return unresolved_function.to_plan(session) + + def __repr__(self) -> str: + _cases = "".join([f" WHEN {c} THEN {v}" for c, v in self._branches]) + _else = f" ELSE {self._else_value}" if self._else_value is not None else "" + return "CASE" + _cases + _else + " END" + + +class ColumnAlias(Expression): + def __init__(self, parent: Expression, alias: Sequence[str], metadata: Any): + super().__init__() + + self._alias = alias + self._metadata = metadata + self._parent = parent + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + if len(self._alias) == 1: + exp = proto.Expression() + exp.alias.name.append(self._alias[0]) + exp.alias.expr.CopyFrom(self._parent.to_plan(session)) + + if self._metadata: + exp.alias.metadata = json.dumps(self._metadata) + return exp + else: + if self._metadata: + raise PySparkValueError( + error_class="CANNOT_PROVIDE_METADATA", + message_parameters={}, + ) + exp = proto.Expression() + exp.alias.name.extend(self._alias) + exp.alias.expr.CopyFrom(self._parent.to_plan(session)) + return exp + + def __repr__(self) -> str: + return f"{self._parent} AS {','.join(self._alias)}" + + +class LiteralExpression(Expression): + """A literal expression. + + The Python types are converted best effort into the relevant proto types. On the Spark Connect + server side, the proto types are converted to the Catalyst equivalents.""" + + def __init__(self, value: Any, dataType: DataType) -> None: + super().__init__() + + assert isinstance( + dataType, + ( + NullType, + BinaryType, + BooleanType, + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + DecimalType, + StringType, + DateType, + TimestampType, + TimestampNTZType, + DayTimeIntervalType, + ArrayType, + ), + ) + + if isinstance(dataType, NullType): + assert value is None + + if value is not None: + if isinstance(dataType, BinaryType): + assert isinstance(value, (bytes, bytearray)) + elif isinstance(dataType, BooleanType): + assert isinstance(value, (bool, np.bool_)) + value = bool(value) + elif isinstance(dataType, ByteType): + assert isinstance(value, (int, np.int8)) + assert JVM_BYTE_MIN <= int(value) <= JVM_BYTE_MAX + value = int(value) + elif isinstance(dataType, ShortType): + assert isinstance(value, (int, np.int8, np.int16)) + assert JVM_SHORT_MIN <= int(value) <= JVM_SHORT_MAX + value = int(value) + elif isinstance(dataType, IntegerType): + assert isinstance(value, (int, np.int8, np.int16, np.int32)) + assert JVM_INT_MIN <= int(value) <= JVM_INT_MAX + value = int(value) + elif isinstance(dataType, LongType): + assert isinstance(value, (int, np.int8, np.int16, np.int32, np.int64)) + assert JVM_LONG_MIN <= int(value) <= JVM_LONG_MAX + value = int(value) + elif isinstance(dataType, FloatType): + assert isinstance(value, (float, np.float32)) + value = float(value) + elif isinstance(dataType, DoubleType): + assert isinstance(value, (float, np.float32, np.float64)) + value = float(value) + elif isinstance(dataType, DecimalType): + assert isinstance(value, decimal.Decimal) + elif isinstance(dataType, StringType): + assert isinstance(value, str) + elif isinstance(dataType, DateType): + assert isinstance(value, (datetime.date, datetime.datetime)) + if isinstance(value, datetime.date): + value = DateType().toInternal(value) + else: + value = DateType().toInternal(value.date()) + elif isinstance(dataType, TimestampType): + assert isinstance(value, datetime.datetime) + value = TimestampType().toInternal(value) + elif isinstance(dataType, TimestampNTZType): + assert isinstance(value, datetime.datetime) + value = TimestampNTZType().toInternal(value) + elif isinstance(dataType, DayTimeIntervalType): + assert isinstance(value, datetime.timedelta) + value = DayTimeIntervalType().toInternal(value) + assert value is not None + elif isinstance(dataType, ArrayType): + assert isinstance(value, list) + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": str(dataType)}, + ) + + self._value = value + self._dataType = dataType + + @classmethod + def _infer_type(cls, value: Any) -> DataType: + if value is None: + return NullType() + elif isinstance(value, (bytes, bytearray)): + return BinaryType() + elif isinstance(value, bool): + return BooleanType() + elif isinstance(value, int): + if JVM_INT_MIN <= value <= JVM_INT_MAX: + return IntegerType() + elif JVM_LONG_MIN <= value <= JVM_LONG_MAX: + return LongType() + else: + raise PySparkValueError( + error_class="VALUE_NOT_BETWEEN", + message_parameters={ + "arg_name": "value", + "min": str(JVM_LONG_MIN), + "max": str(JVM_SHORT_MAX), + }, + ) + elif isinstance(value, float): + return DoubleType() + elif isinstance(value, str): + return StringType() + elif isinstance(value, decimal.Decimal): + return DecimalType() + elif isinstance(value, datetime.datetime) and is_timestamp_ntz_preferred(): + return TimestampNTZType() + elif isinstance(value, datetime.datetime): + return TimestampType() + elif isinstance(value, datetime.date): + return DateType() + elif isinstance(value, datetime.timedelta): + return DayTimeIntervalType() + elif isinstance(value, np.generic): + dt = _from_numpy_type(value.dtype) + if dt is not None: + return dt + elif isinstance(value, np.bool_): + return BooleanType() + elif isinstance(value, list): + # follow the 'infer_array_from_first_element' strategy in 'sql.types._infer_type' + # right now, it's dedicated for pyspark.ml params like array<...>, array> + if len(value) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "value"}, + ) + first = value[0] + if first is None: + raise PySparkTypeError( + error_class="CANNOT_INFER_ARRAY_TYPE", + message_parameters={}, + ) + return ArrayType(LiteralExpression._infer_type(first), True) + + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(value).__name__}, + ) + + @classmethod + def _from_value(cls, value: Any) -> "LiteralExpression": + return LiteralExpression(value=value, dataType=LiteralExpression._infer_type(value)) + + @classmethod + def _to_value( + cls, literal: "proto.Expression.Literal", dataType: Optional[DataType] = None + ) -> Any: + if literal.HasField("null"): + return None + elif literal.HasField("binary"): + assert dataType is None or isinstance(dataType, BinaryType) + return literal.binary + elif literal.HasField("boolean"): + assert dataType is None or isinstance(dataType, BooleanType) + return literal.boolean + elif literal.HasField("byte"): + assert dataType is None or isinstance(dataType, ByteType) + return literal.byte + elif literal.HasField("short"): + assert dataType is None or isinstance(dataType, ShortType) + return literal.short + elif literal.HasField("integer"): + assert dataType is None or isinstance(dataType, IntegerType) + return literal.integer + elif literal.HasField("long"): + assert dataType is None or isinstance(dataType, LongType) + return literal.long + elif literal.HasField("float"): + assert dataType is None or isinstance(dataType, FloatType) + return literal.float + elif literal.HasField("double"): + assert dataType is None or isinstance(dataType, DoubleType) + return literal.double + elif literal.HasField("decimal"): + assert dataType is None or isinstance(dataType, DecimalType) + return decimal.Decimal(literal.decimal.value) + elif literal.HasField("string"): + assert dataType is None or isinstance(dataType, StringType) + return literal.string + elif literal.HasField("date"): + assert dataType is None or isinstance(dataType, DataType) + return DateType().fromInternal(literal.date) + elif literal.HasField("timestamp"): + assert dataType is None or isinstance(dataType, TimestampType) + return TimestampType().fromInternal(literal.timestamp) + elif literal.HasField("timestamp_ntz"): + assert dataType is None or isinstance(dataType, TimestampNTZType) + return TimestampNTZType().fromInternal(literal.timestamp_ntz) + elif literal.HasField("day_time_interval"): + assert dataType is None or isinstance(dataType, DayTimeIntervalType) + return DayTimeIntervalType().fromInternal(literal.day_time_interval) + elif literal.HasField("array"): + elementType = proto_schema_to_pyspark_data_type(literal.array.element_type) + if dataType is not None: + assert isinstance(dataType, ArrayType) + assert elementType == dataType.elementType + return [LiteralExpression._to_value(v, elementType) for v in literal.array.elements] + + raise PySparkTypeError( + error_class="UNSUPPORTED_LITERAL", + message_parameters={"literal": str(literal)}, + ) + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + """Converts the literal expression to the literal in proto.""" + + expr = proto.Expression() + + if self._value is None: + expr.literal.null.CopyFrom(pyspark_types_to_proto_types(self._dataType)) + elif isinstance(self._dataType, BinaryType): + expr.literal.binary = bytes(self._value) + elif isinstance(self._dataType, BooleanType): + expr.literal.boolean = bool(self._value) + elif isinstance(self._dataType, ByteType): + expr.literal.byte = int(self._value) + elif isinstance(self._dataType, ShortType): + expr.literal.short = int(self._value) + elif isinstance(self._dataType, IntegerType): + expr.literal.integer = int(self._value) + elif isinstance(self._dataType, LongType): + expr.literal.long = int(self._value) + elif isinstance(self._dataType, FloatType): + expr.literal.float = float(self._value) + elif isinstance(self._dataType, DoubleType): + expr.literal.double = float(self._value) + elif isinstance(self._dataType, DecimalType): + expr.literal.decimal.value = str(self._value) + expr.literal.decimal.precision = self._dataType.precision + expr.literal.decimal.scale = self._dataType.scale + elif isinstance(self._dataType, StringType): + expr.literal.string = str(self._value) + elif isinstance(self._dataType, DateType): + expr.literal.date = int(self._value) + elif isinstance(self._dataType, TimestampType): + expr.literal.timestamp = int(self._value) + elif isinstance(self._dataType, TimestampNTZType): + expr.literal.timestamp_ntz = int(self._value) + elif isinstance(self._dataType, DayTimeIntervalType): + expr.literal.day_time_interval = int(self._value) + elif isinstance(self._dataType, ArrayType): + element_type = self._dataType.elementType + expr.literal.array.element_type.CopyFrom(pyspark_types_to_proto_types(element_type)) + for v in self._value: + expr.literal.array.elements.append( + LiteralExpression(v, element_type).to_plan(session).literal + ) + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": str(self._dataType)}, + ) + + return expr + + def __repr__(self) -> str: + return f"{self._value}" + + +class ColumnReference(Expression): + """Represents a column reference. There is no guarantee that this column + actually exists. In the context of this project, we refer by its name and + treat it as an unresolved attribute. Attributes that have the same fully + qualified name are identical""" + + def __init__(self, unparsed_identifier: str, plan_id: Optional[int] = None) -> None: + super().__init__() + assert isinstance(unparsed_identifier, str) + self._unparsed_identifier = unparsed_identifier + + assert plan_id is None or isinstance(plan_id, int) + self._plan_id = plan_id + + def name(self) -> str: + """Returns the qualified name of the column reference.""" + return self._unparsed_identifier + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + """Returns the Proto representation of the expression.""" + expr = proto.Expression() + expr.unresolved_attribute.unparsed_identifier = self._unparsed_identifier + if self._plan_id is not None: + expr.unresolved_attribute.plan_id = self._plan_id + return expr + + def __repr__(self) -> str: + return f"{self._unparsed_identifier}" + + def __eq__(self, other: Any) -> bool: + return ( + other is not None + and isinstance(other, ColumnReference) + and other._unparsed_identifier == self._unparsed_identifier + ) + + +class UnresolvedStar(Expression): + def __init__(self, unparsed_target: Optional[str]): + super().__init__() + + if unparsed_target is not None: + assert isinstance(unparsed_target, str) and unparsed_target.endswith(".*") + + self._unparsed_target = unparsed_target + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + expr = proto.Expression() + expr.unresolved_star.SetInParent() + if self._unparsed_target is not None: + expr.unresolved_star.unparsed_target = self._unparsed_target + return expr + + def __repr__(self) -> str: + if self._unparsed_target is not None: + return f"unresolvedstar({self._unparsed_target})" + else: + return "unresolvedstar()" + + def __eq__(self, other: Any) -> bool: + return ( + other is not None + and isinstance(other, UnresolvedStar) + and other._unparsed_target == self._unparsed_target + ) + + +class SQLExpression(Expression): + """Returns Expression which contains a string which is a SQL expression + and server side will parse it by Catalyst + """ + + def __init__(self, expr: str) -> None: + super().__init__() + self._expr: str = expr + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + """Returns the Proto representation of the SQL expression.""" + expr = proto.Expression() + expr.expression_string.expression = self._expr + return expr + + def __eq__(self, other: Any) -> bool: + return other is not None and isinstance(other, SQLExpression) and other._expr == self._expr + + +class SortOrder(Expression): + def __init__(self, child: Expression, ascending: bool = True, nullsFirst: bool = True) -> None: + super().__init__() + self._child = child + self._ascending = ascending + self._nullsFirst = nullsFirst + + def __repr__(self) -> str: + return ( + str(self._child) + + (" ASC" if self._ascending else " DESC") + + (" NULLS FIRST" if self._nullsFirst else " NULLS LAST") + ) + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + sort = proto.Expression() + sort.sort_order.child.CopyFrom(self._child.to_plan(session)) + + if self._ascending: + sort.sort_order.direction = ( + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_ASCENDING + ) + else: + sort.sort_order.direction = ( + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_DESCENDING + ) + + if self._nullsFirst: + sort.sort_order.null_ordering = proto.Expression.SortOrder.NullOrdering.SORT_NULLS_FIRST + else: + sort.sort_order.null_ordering = proto.Expression.SortOrder.NullOrdering.SORT_NULLS_LAST + + return sort + + +class UnresolvedFunction(Expression): + def __init__( + self, + name: str, + args: Sequence["Expression"], + is_distinct: bool = False, + ) -> None: + super().__init__() + + assert isinstance(name, str) + self._name = name + + assert isinstance(args, list) and all(isinstance(arg, Expression) for arg in args) + self._args = args + + assert isinstance(is_distinct, bool) + self._is_distinct = is_distinct + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + fun = proto.Expression() + fun.unresolved_function.function_name = self._name + if len(self._args) > 0: + fun.unresolved_function.arguments.extend([arg.to_plan(session) for arg in self._args]) + fun.unresolved_function.is_distinct = self._is_distinct + return fun + + def __repr__(self) -> str: + # Default print handling: + if self._is_distinct: + return f"{self._name}(distinct {', '.join([str(arg) for arg in self._args])})" + else: + return f"{self._name}({', '.join([str(arg) for arg in self._args])})" + + +class PythonUDF: + """Represents a Python user-defined function.""" + + def __init__( + self, + output_type: Union[DataType, str], + eval_type: int, + func: Callable[..., Any], + python_ver: str, + ) -> None: + self._output_type: DataType = ( + UnparsedDataType(output_type) if isinstance(output_type, str) else output_type + ) + self._eval_type = eval_type + self._func = func + self._python_ver = python_ver + + def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF: + if isinstance(self._output_type, UnparsedDataType): + parsed = session._analyze( + method="ddl_parse", ddl_string=self._output_type.data_type_string + ).parsed + assert isinstance(parsed, DataType) + output_type = parsed + else: + output_type = self._output_type + expr = proto.PythonUDF() + expr.output_type.CopyFrom(pyspark_types_to_proto_types(output_type)) + expr.eval_type = self._eval_type + expr.command = CloudPickleSerializer().dumps((self._func, output_type)) + expr.python_ver = self._python_ver + return expr + + def __repr__(self) -> str: + return f"{self._output_type}, {self._eval_type}, {self._func}, f{self._python_ver}" + + +class JavaUDF: + """Represents a Java (aggregate) user-defined function.""" + + def __init__( + self, + class_name: str, + output_type: Optional[Union[DataType, str]] = None, + aggregate: bool = False, + ) -> None: + self._class_name = class_name + self._output_type: Optional[DataType] = ( + UnparsedDataType(output_type) if isinstance(output_type, str) else output_type + ) + self._aggregate = aggregate + + def to_plan(self, session: "SparkConnectClient") -> proto.JavaUDF: + expr = proto.JavaUDF() + expr.class_name = self._class_name + if self._output_type is not None: + expr.output_type.CopyFrom(pyspark_types_to_proto_types(self._output_type)) + expr.aggregate = self._aggregate + return expr + + def __repr__(self) -> str: + return f"{self._class_name}, {self._output_type}" + + +class CommonInlineUserDefinedFunction(Expression): + """Represents a user-defined function with an inlined defined function body of any programming + languages.""" + + def __init__( + self, + function_name: str, + function: Union[PythonUDF, JavaUDF], + deterministic: bool = False, + arguments: Sequence[Expression] = [], + ): + super().__init__() + self._function_name = function_name + self._deterministic = deterministic + self._arguments = arguments + self._function = function + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + expr = proto.Expression() + expr.common_inline_user_defined_function.function_name = self._function_name + expr.common_inline_user_defined_function.deterministic = self._deterministic + if len(self._arguments) > 0: + expr.common_inline_user_defined_function.arguments.extend( + [arg.to_plan(session) for arg in self._arguments] + ) + expr.common_inline_user_defined_function.python_udf.CopyFrom( + cast(proto.PythonUDF, self._function.to_plan(session)) + ) + return expr + + def to_plan_udf(self, session: "SparkConnectClient") -> "proto.CommonInlineUserDefinedFunction": + """Compared to `to_plan`, it returns a CommonInlineUserDefinedFunction instead of an + Expression.""" + expr = proto.CommonInlineUserDefinedFunction() + expr.function_name = self._function_name + expr.deterministic = self._deterministic + if len(self._arguments) > 0: + expr.arguments.extend([arg.to_plan(session) for arg in self._arguments]) + expr.python_udf.CopyFrom(cast(proto.PythonUDF, self._function.to_plan(session))) + return expr + + def to_plan_judf( + self, session: "SparkConnectClient" + ) -> "proto.CommonInlineUserDefinedFunction": + expr = proto.CommonInlineUserDefinedFunction() + expr.function_name = self._function_name + expr.java_udf.CopyFrom(cast(proto.JavaUDF, self._function.to_plan(session))) + return expr + + def __repr__(self) -> str: + return f"{self._function_name}({', '.join([str(arg) for arg in self._arguments])})" + + +class WithField(Expression): + def __init__( + self, + structExpr: Expression, + fieldName: str, + valueExpr: Expression, + ) -> None: + super().__init__() + + assert isinstance(structExpr, Expression) + self._structExpr = structExpr + + assert isinstance(fieldName, str) + self._fieldName = fieldName + + assert isinstance(valueExpr, Expression) + self._valueExpr = valueExpr + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.update_fields.struct_expression.CopyFrom(self._structExpr.to_plan(session)) + expr.update_fields.field_name = self._fieldName + expr.update_fields.value_expression.CopyFrom(self._valueExpr.to_plan(session)) + return expr + + def __repr__(self) -> str: + return f"WithField({self._structExpr}, {self._fieldName}, {self._valueExpr})" + + +class DropField(Expression): + def __init__( + self, + structExpr: Expression, + fieldName: str, + ) -> None: + super().__init__() + + assert isinstance(structExpr, Expression) + self._structExpr = structExpr + + assert isinstance(fieldName, str) + self._fieldName = fieldName + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.update_fields.struct_expression.CopyFrom(self._structExpr.to_plan(session)) + expr.update_fields.field_name = self._fieldName + return expr + + def __repr__(self) -> str: + return f"DropField({self._structExpr}, {self._fieldName})" + + +class UnresolvedExtractValue(Expression): + def __init__( + self, + child: Expression, + extraction: Expression, + ) -> None: + super().__init__() + + assert isinstance(child, Expression) + self._child = child + + assert isinstance(extraction, Expression) + self._extraction = extraction + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.unresolved_extract_value.child.CopyFrom(self._child.to_plan(session)) + expr.unresolved_extract_value.extraction.CopyFrom(self._extraction.to_plan(session)) + return expr + + def __repr__(self) -> str: + return f"UnresolvedExtractValue({str(self._child)}, {str(self._extraction)})" + + +class UnresolvedRegex(Expression): + def __init__(self, col_name: str, plan_id: Optional[int] = None) -> None: + super().__init__() + + assert isinstance(col_name, str) + self.col_name = col_name + + assert plan_id is None or isinstance(plan_id, int) + self._plan_id = plan_id + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.unresolved_regex.col_name = self.col_name + if self._plan_id is not None: + expr.unresolved_regex.plan_id = self._plan_id + return expr + + def __repr__(self) -> str: + return f"UnresolvedRegex({self.col_name})" + + +class CastExpression(Expression): + def __init__( + self, + expr: Expression, + data_type: Union[DataType, str], + ) -> None: + super().__init__() + self._expr = expr + self._data_type = data_type + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + fun = proto.Expression() + fun.cast.expr.CopyFrom(self._expr.to_plan(session)) + if isinstance(self._data_type, str): + fun.cast.type_str = self._data_type + else: + fun.cast.type.CopyFrom(pyspark_types_to_proto_types(self._data_type)) + return fun + + def __repr__(self) -> str: + return f"({self._expr} ({self._data_type}))" + + +class UnresolvedNamedLambdaVariable(Expression): + + _lock: Lock = Lock() + _nextVarNameId: int = 0 + + def __init__( + self, + name_parts: Sequence[str], + ) -> None: + super().__init__() + + assert ( + isinstance(name_parts, list) + and len(name_parts) > 0 + and all(isinstance(p, str) for p in name_parts) + ) + + self._name_parts = name_parts + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.unresolved_named_lambda_variable.name_parts.extend(self._name_parts) + return expr + + def __repr__(self) -> str: + return f"(UnresolvedNamedLambdaVariable({', '.join(self._name_parts)})" + + @staticmethod + def fresh_var_name(name: str) -> str: + assert isinstance(name, str) and str != "" + + _id: Optional[int] = None + + with UnresolvedNamedLambdaVariable._lock: + _id = UnresolvedNamedLambdaVariable._nextVarNameId + UnresolvedNamedLambdaVariable._nextVarNameId += 1 + + assert _id is not None + + return f"{name}_{_id}" + + +class LambdaFunction(Expression): + def __init__( + self, + function: Expression, + arguments: Sequence[UnresolvedNamedLambdaVariable], + ) -> None: + super().__init__() + + assert isinstance(function, Expression) + + assert ( + isinstance(arguments, list) + and len(arguments) > 0 + and all(isinstance(arg, UnresolvedNamedLambdaVariable) for arg in arguments) + ) + + self._function = function + self._arguments = arguments + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + expr.lambda_function.function.CopyFrom(self._function.to_plan(session)) + expr.lambda_function.arguments.extend( + [arg.to_plan(session).unresolved_named_lambda_variable for arg in self._arguments] + ) + return expr + + def __repr__(self) -> str: + return f"(LambdaFunction({str(self._function)}, {', '.join(self._arguments)})" + + +class WindowExpression(Expression): + def __init__( + self, + windowFunction: Expression, + windowSpec: "WindowSpec", + ) -> None: + super().__init__() + + from pyspark.sql.connect.window import WindowSpec + + assert windowFunction is not None and isinstance(windowFunction, Expression) + + assert windowSpec is not None and isinstance(windowSpec, WindowSpec) + + self._windowFunction = windowFunction + + self._windowSpec = windowSpec + + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + expr = proto.Expression() + + expr.window.window_function.CopyFrom(self._windowFunction.to_plan(session)) + + if len(self._windowSpec._partitionSpec) > 0: + expr.window.partition_spec.extend( + [p.to_plan(session) for p in self._windowSpec._partitionSpec] + ) + else: + warnings.warn( + "WARN WindowExpression: No Partition Defined for Window operation! " + "Moving all data to a single partition, this can cause serious " + "performance degradation." + ) + + if len(self._windowSpec._orderSpec) > 0: + expr.window.order_spec.extend( + [s.to_plan(session).sort_order for s in self._windowSpec._orderSpec] + ) + + if self._windowSpec._frame is not None: + if self._windowSpec._frame._isRowFrame: + expr.window.frame_spec.frame_type = ( + proto.Expression.Window.WindowFrame.FrameType.FRAME_TYPE_ROW + ) + + start = self._windowSpec._frame._start + if start == 0: + expr.window.frame_spec.lower.current_row = True + elif start == JVM_LONG_MIN: + expr.window.frame_spec.lower.unbounded = True + elif JVM_INT_MIN <= start <= JVM_INT_MAX: + expr.window.frame_spec.lower.value.literal.integer = start + else: + raise PySparkValueError( + error_class="VALUE_NOT_BETWEEN", + message_parameters={ + "arg_name": "start", + "min": str(JVM_INT_MIN), + "max": str(JVM_INT_MAX), + }, + ) + + end = self._windowSpec._frame._end + if end == 0: + expr.window.frame_spec.upper.current_row = True + elif end == JVM_LONG_MAX: + expr.window.frame_spec.upper.unbounded = True + elif JVM_INT_MIN <= end <= JVM_INT_MAX: + expr.window.frame_spec.upper.value.literal.integer = end + else: + raise PySparkValueError( + error_class="VALUE_NOT_BETWEEN", + message_parameters={ + "arg_name": "end", + "min": str(JVM_INT_MIN), + "max": str(JVM_INT_MAX), + }, + ) + + else: + expr.window.frame_spec.frame_type = ( + proto.Expression.Window.WindowFrame.FrameType.FRAME_TYPE_RANGE + ) + + start = self._windowSpec._frame._start + if start == 0: + expr.window.frame_spec.lower.current_row = True + elif start == JVM_LONG_MIN: + expr.window.frame_spec.lower.unbounded = True + else: + expr.window.frame_spec.lower.value.literal.long = start + + end = self._windowSpec._frame._end + if end == 0: + expr.window.frame_spec.upper.current_row = True + elif end == JVM_LONG_MAX: + expr.window.frame_spec.upper.unbounded = True + else: + expr.window.frame_spec.upper.value.literal.long = end + + return expr + + def __repr__(self) -> str: + return f"WindowExpression({str(self._windowFunction)}, ({str(self._windowSpec)}))" + + +class DistributedSequenceID(Expression): + def to_plan(self, session: "SparkConnectClient") -> proto.Expression: + unresolved_function = UnresolvedFunction(name="distributed_sequence_id", args=[]) + return unresolved_function.to_plan(session) + + def __repr__(self) -> str: + return "DistributedSequenceID()" + + +class CallFunction(Expression): + def __init__(self, name: str, args: Sequence["Expression"]): + super().__init__() + + assert isinstance(name, str) + self._name = name + + assert isinstance(args, list) and all(isinstance(arg, Expression) for arg in args) + self._args = args + + def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": + expr = proto.Expression() + expr.call_function.function_name = self._name + if len(self._args) > 0: + expr.call_function.arguments.extend([arg.to_plan(session) for arg in self._args]) + return expr + + def __repr__(self) -> str: + if len(self._args) > 0: + return f"CallFunction('{self._name}', {', '.join([str(arg) for arg in self._args])})" + else: + return f"CallFunction('{self._name}')" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/functions.py new file mode 100644 index 0000000..e2583f8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/functions.py @@ -0,0 +1,3936 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import decimal +import inspect +import warnings +import functools +from typing import ( + Any, + Dict, + TYPE_CHECKING, + Union, + List, + overload, + Optional, + Tuple, + Type, + Callable, + ValuesView, + cast, +) + +import numpy as np + +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.expressions import ( + CaseWhen, + Expression, + LiteralExpression, + ColumnReference, + UnresolvedFunction, + UnresolvedStar, + SQLExpression, + LambdaFunction, + UnresolvedNamedLambdaVariable, + CallFunction, +) +from pyspark.sql.connect.udf import _create_py_udf +from pyspark.sql.connect.udtf import _create_py_udtf +from pyspark.sql import functions as pysparkfuncs +from pyspark.sql.types import _from_numpy_type, DataType, StructType, ArrayType, StringType + +# The implementation of pandas_udf is embedded in pyspark.sql.function.pandas_udf +# for code reuse. +from pyspark.sql.functions import pandas_udf # noqa: F401 + + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ( + ColumnOrName, + DataTypeOrString, + UserDefinedFunctionLike, + ) + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udtf import UserDefinedTableFunction + + +def _to_col_with_plan_id(col: str, plan_id: Optional[int]) -> Column: + if col == "*": + return Column(UnresolvedStar(unparsed_target=None)) + elif col.endswith(".*"): + return Column(UnresolvedStar(unparsed_target=col)) + else: + return Column(ColumnReference(unparsed_identifier=col, plan_id=plan_id)) + + +def _to_col(col: "ColumnOrName") -> Column: + assert isinstance(col, (Column, str)) + return col if isinstance(col, Column) else column(col) + + +def _invoke_function(name: str, *args: Union[Column, Expression]) -> Column: + """ + Simple wrapper function that converts the arguments into the appropriate types. + Parameters + ---------- + name Name of the function to be called. + args The list of arguments. + + Returns + ------- + :class:`Column` + """ + expressions: List[Expression] = [] + for arg in args: + assert isinstance(arg, (Column, Expression)) + if isinstance(arg, Column): + expressions.append(arg._expr) + else: + expressions.append(arg) + return Column(UnresolvedFunction(name, expressions)) + + +def _invoke_function_over_columns(name: str, *cols: "ColumnOrName") -> Column: + """ + Invokes n-ary function identified by name + and wraps the result with :class:`~pyspark.sql.Column`. + """ + _cols = [_to_col(c) for c in cols] + return _invoke_function(name, *_cols) + + +def _invoke_binary_math_function(name: str, col1: Any, col2: Any) -> Column: + """ + Invokes binary math function identified by name + and wraps the result with :class:`~pyspark.sql.Column`. + """ + + # For legacy reasons, the arguments here can be implicitly converted into column + _cols = [_to_col(c) if isinstance(c, (str, Column)) else lit(c) for c in (col1, col2)] + return _invoke_function(name, *_cols) + + +def _get_lambda_parameters(f: Callable) -> ValuesView[inspect.Parameter]: + signature = inspect.signature(f) + parameters = signature.parameters.values() + + # We should exclude functions that use, variable args and keyword argument + # names, as well as keyword only args. + supported_parameter_types = { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY, + } + + # Validate that the function arity is between 1 and 3. + if not (1 <= len(parameters) <= 3): + raise PySparkValueError( + error_class="WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": f.__name__, "num_args": str(len(parameters))}, + ) + + # Verify that all arguments can be used as positional arguments. + if not all(p.kind in supported_parameter_types for p in parameters): + raise PySparkValueError( + error_class="UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": f.__name__}, + ) + + return parameters + + +def _create_lambda(f: Callable) -> LambdaFunction: + """ + Create `o.a.s.sql.expressions.LambdaFunction` corresponding + to transformation described by f + + :param f: A Python of one of the following forms: + - (Column) -> Column: ... + - (Column, Column) -> Column: ... + - (Column, Column, Column) -> Column: ... + """ + parameters = _get_lambda_parameters(f) + + arg_names = ["x", "y", "z"][: len(parameters)] + arg_exprs = [ + UnresolvedNamedLambdaVariable([UnresolvedNamedLambdaVariable.fresh_var_name(arg_name)]) + for arg_name in arg_names + ] + arg_cols = [Column(arg_expr) for arg_expr in arg_exprs] + + result = f(*arg_cols) + + if not isinstance(result, Column): + raise PySparkValueError( + error_class="HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN", + message_parameters={"func_name": f.__name__, "return_type": type(result).__name__}, + ) + + return LambdaFunction(result._expr, arg_exprs) + + +def _invoke_higher_order_function( + name: str, + cols: List["ColumnOrName"], + funs: List[Callable], +) -> Column: + """ + Invokes expression identified by name, + (relative to ```org.apache.spark.sql.catalyst.expressions``) + and wraps the result with Column (first Scala one, then Python). + + :param name: Name of the expression + :param cols: a list of columns + :param funs: a list of (*Column) -> Column functions. + + :return: a Column + """ + _cols = [_to_col(c) for c in cols] + _funs = [_create_lambda(f) for f in funs] + + return _invoke_function(name, *_cols, *_funs) + + +def _options_to_col(options: Dict[str, Any]) -> Column: + _options: List[Column] = [] + for k, v in options.items(): + _options.append(lit(str(k))) + _options.append(lit(str(v))) + return create_map(*_options) + + +# Normal Functions + + +def col(col: str) -> Column: + return _to_col_with_plan_id(col=col, plan_id=None) + + +col.__doc__ = pysparkfuncs.col.__doc__ + + +column = col + + +def lit(col: Any) -> Column: + if isinstance(col, Column): + return col + elif isinstance(col, list): + if any(isinstance(c, Column) for c in col): + raise PySparkValueError( + error_class="COLUMN_IN_LIST", message_parameters={"func_name": "lit"} + ) + return array(*[lit(c) for c in col]) + elif isinstance(col, np.ndarray) and col.ndim == 1: + if _from_numpy_type(col.dtype) is None: + raise PySparkTypeError( + error_class="UNSUPPORTED_NUMPY_ARRAY_SCALAR", + message_parameters={"dtype": col.dtype.name}, + ) + + # NumpyArrayConverter for Py4J can not support ndarray with int8 values. + # Actually this is not a problem for Connect, but here still convert it + # to int16 for compatibility. + if col.dtype == np.int8: + col = col.astype(np.int16) + + return array(*[lit(c) for c in col]) + else: + return Column(LiteralExpression._from_value(col)) + + +lit.__doc__ = pysparkfuncs.lit.__doc__ + + +def bitwiseNOT(col: "ColumnOrName") -> Column: + warnings.warn("Deprecated in 3.4, use bitwise_not instead.", FutureWarning) + return bitwise_not(col) + + +bitwiseNOT.__doc__ = pysparkfuncs.bitwiseNOT.__doc__ + + +def bitwise_not(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("~", col) + + +bitwise_not.__doc__ = pysparkfuncs.bitwise_not.__doc__ + + +def bit_count(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_count", col) + + +bit_count.__doc__ = pysparkfuncs.bit_count.__doc__ + + +def bit_get(col: "ColumnOrName", pos: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_get", col, pos) + + +bit_get.__doc__ = pysparkfuncs.bit_get.__doc__ + + +def getbit(col: "ColumnOrName", pos: "ColumnOrName") -> Column: + return _invoke_function_over_columns("getbit", col, pos) + + +getbit.__doc__ = pysparkfuncs.getbit.__doc__ + + +def broadcast(df: "DataFrame") -> "DataFrame": + from pyspark.sql.connect.dataframe import DataFrame + + if not isinstance(df, DataFrame): + raise PySparkTypeError( + error_class="NOT_DATAFRAME", + message_parameters={"arg_name": "df", "arg_type": type(df).__name__}, + ) + return df.hint("broadcast") + + +broadcast.__doc__ = pysparkfuncs.broadcast.__doc__ + + +def coalesce(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("coalesce", *cols) + + +coalesce.__doc__ = pysparkfuncs.coalesce.__doc__ + + +def expr(str: str) -> Column: + return Column(SQLExpression(str)) + + +expr.__doc__ = pysparkfuncs.expr.__doc__ + + +def greatest(*cols: "ColumnOrName") -> Column: + if len(cols) < 2: + raise PySparkValueError( + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "greatest", "num_cols": "2"}, + ) + return _invoke_function_over_columns("greatest", *cols) + + +greatest.__doc__ = pysparkfuncs.greatest.__doc__ + + +def input_file_name() -> Column: + return _invoke_function("input_file_name") + + +input_file_name.__doc__ = pysparkfuncs.input_file_name.__doc__ + + +def least(*cols: "ColumnOrName") -> Column: + if len(cols) < 2: + raise PySparkValueError( + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "least", "num_cols": "2"}, + ) + return _invoke_function_over_columns("least", *cols) + + +least.__doc__ = pysparkfuncs.least.__doc__ + + +def isnan(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("isnan", col) + + +isnan.__doc__ = pysparkfuncs.isnan.__doc__ + + +def isnull(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("isnull", col) + + +isnull.__doc__ = pysparkfuncs.isnull.__doc__ + + +def monotonically_increasing_id() -> Column: + return _invoke_function("monotonically_increasing_id") + + +monotonically_increasing_id.__doc__ = pysparkfuncs.monotonically_increasing_id.__doc__ + + +def nanvl(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("nanvl", col1, col2) + + +nanvl.__doc__ = pysparkfuncs.nanvl.__doc__ + + +def rand(seed: Optional[int] = None) -> Column: + if seed is not None: + return _invoke_function("rand", lit(seed)) + else: + return _invoke_function("rand") + + +rand.__doc__ = pysparkfuncs.rand.__doc__ + + +def randn(seed: Optional[int] = None) -> Column: + if seed is not None: + return _invoke_function("randn", lit(seed)) + else: + return _invoke_function("randn") + + +randn.__doc__ = pysparkfuncs.randn.__doc__ + + +def spark_partition_id() -> Column: + return _invoke_function("spark_partition_id") + + +spark_partition_id.__doc__ = pysparkfuncs.spark_partition_id.__doc__ + + +def when(condition: Column, value: Any) -> Column: + # Explicitly not using ColumnOrName type here to make reading condition less opaque + if not isinstance(condition, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + + value_col = value if isinstance(value, Column) else lit(value) + + return Column(CaseWhen(branches=[(condition._expr, value_col._expr)], else_value=None)) + + +when.__doc__ = pysparkfuncs.when.__doc__ + + +# Sort Functions + + +def asc(col: "ColumnOrName") -> Column: + return _to_col(col).asc() + + +asc.__doc__ = pysparkfuncs.asc.__doc__ + + +def asc_nulls_first(col: "ColumnOrName") -> Column: + return _to_col(col).asc_nulls_first() + + +asc_nulls_first.__doc__ = pysparkfuncs.asc_nulls_first.__doc__ + + +def asc_nulls_last(col: "ColumnOrName") -> Column: + return _to_col(col).asc_nulls_last() + + +asc_nulls_last.__doc__ = pysparkfuncs.asc_nulls_last.__doc__ + + +def desc(col: "ColumnOrName") -> Column: + return _to_col(col).desc() + + +desc.__doc__ = pysparkfuncs.desc.__doc__ + + +def desc_nulls_first(col: "ColumnOrName") -> Column: + return _to_col(col).desc_nulls_first() + + +desc_nulls_first.__doc__ = pysparkfuncs.desc_nulls_first.__doc__ + + +def desc_nulls_last(col: "ColumnOrName") -> Column: + return _to_col(col).desc_nulls_last() + + +desc_nulls_last.__doc__ = pysparkfuncs.desc_nulls_last.__doc__ + + +# Math Functions + + +def abs(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("abs", col) + + +abs.__doc__ = pysparkfuncs.abs.__doc__ + + +def acos(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("acos", col) + + +acos.__doc__ = pysparkfuncs.acos.__doc__ + + +def acosh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("acosh", col) + + +acosh.__doc__ = pysparkfuncs.acosh.__doc__ + + +def asin(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("asin", col) + + +asin.__doc__ = pysparkfuncs.asin.__doc__ + + +def asinh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("asinh", col) + + +asinh.__doc__ = pysparkfuncs.asinh.__doc__ + + +def atan(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("atan", col) + + +atan.__doc__ = pysparkfuncs.atan.__doc__ + + +def atan2(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + return _invoke_binary_math_function("atan2", col1, col2) + + +atan2.__doc__ = pysparkfuncs.atan2.__doc__ + + +def atanh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("atanh", col) + + +atanh.__doc__ = pysparkfuncs.atanh.__doc__ + + +def bin(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bin", col) + + +bin.__doc__ = pysparkfuncs.bin.__doc__ + + +def bround(col: "ColumnOrName", scale: int = 0) -> Column: + return _invoke_function("bround", _to_col(col), lit(scale)) + + +bround.__doc__ = pysparkfuncs.bround.__doc__ + + +def cbrt(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("cbrt", col) + + +cbrt.__doc__ = pysparkfuncs.cbrt.__doc__ + + +def ceil(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ceil", col) + + +ceil.__doc__ = pysparkfuncs.ceil.__doc__ + + +def ceiling(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ceiling", col) + + +ceiling.__doc__ = pysparkfuncs.ceiling.__doc__ + + +def conv(col: "ColumnOrName", fromBase: int, toBase: int) -> Column: + return _invoke_function("conv", _to_col(col), lit(fromBase), lit(toBase)) + + +conv.__doc__ = pysparkfuncs.conv.__doc__ + + +def cos(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("cos", col) + + +cos.__doc__ = pysparkfuncs.cos.__doc__ + + +def cosh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("cosh", col) + + +cosh.__doc__ = pysparkfuncs.cosh.__doc__ + + +def cot(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("cot", col) + + +cot.__doc__ = pysparkfuncs.cot.__doc__ + + +def csc(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("csc", col) + + +csc.__doc__ = pysparkfuncs.csc.__doc__ + + +def degrees(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("degrees", col) + + +degrees.__doc__ = pysparkfuncs.degrees.__doc__ + + +def e() -> Column: + return _invoke_function("e") + + +e.__doc__ = pysparkfuncs.e.__doc__ + + +def exp(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("exp", col) + + +exp.__doc__ = pysparkfuncs.exp.__doc__ + + +def expm1(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("expm1", col) + + +expm1.__doc__ = pysparkfuncs.expm1.__doc__ + + +def factorial(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("factorial", col) + + +factorial.__doc__ = pysparkfuncs.factorial.__doc__ + + +def floor(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("floor", col) + + +floor.__doc__ = pysparkfuncs.floor.__doc__ + + +def hex(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("hex", col) + + +hex.__doc__ = pysparkfuncs.hex.__doc__ + + +def hypot(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + return _invoke_binary_math_function("hypot", col1, col2) + + +hypot.__doc__ = pysparkfuncs.hypot.__doc__ + + +def log(arg1: Union["ColumnOrName", float], arg2: Optional["ColumnOrName"] = None) -> Column: + if arg2 is None: + # in this case, arg1 should be "ColumnOrName" + return _invoke_function("ln", _to_col(cast("ColumnOrName", arg1))) + else: + # in this case, arg1 should be a float + return _invoke_function("log", lit(cast(float, arg1)), _to_col(arg2)) + + +log.__doc__ = pysparkfuncs.log.__doc__ + + +def log10(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("log10", col) + + +log10.__doc__ = pysparkfuncs.log10.__doc__ + + +def log1p(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("log1p", col) + + +log1p.__doc__ = pysparkfuncs.log1p.__doc__ + + +def ln(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ln", col) + + +ln.__doc__ = pysparkfuncs.ln.__doc__ + + +def log2(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("log2", col) + + +log2.__doc__ = pysparkfuncs.log2.__doc__ + + +def negative(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("negative", col) + + +negative.__doc__ = pysparkfuncs.negative.__doc__ + + +negate = negative + + +def pi() -> Column: + return _invoke_function("pi") + + +pi.__doc__ = pysparkfuncs.pi.__doc__ + + +def positive(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("positive", col) + + +positive.__doc__ = pysparkfuncs.positive.__doc__ + + +def pmod(dividend: Union["ColumnOrName", float], divisor: Union["ColumnOrName", float]) -> Column: + return _invoke_binary_math_function("pmod", dividend, divisor) + + +pmod.__doc__ = pysparkfuncs.pmod.__doc__ + + +def width_bucket( + v: "ColumnOrName", + min: "ColumnOrName", + max: "ColumnOrName", + numBucket: Union["ColumnOrName", int], +) -> Column: + numBucket = lit(numBucket) if isinstance(numBucket, int) else numBucket + return _invoke_function_over_columns("width_bucket", v, min, max, numBucket) + + +width_bucket.__doc__ = pysparkfuncs.width_bucket.__doc__ + + +def pow(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + return _invoke_binary_math_function("power", col1, col2) + + +pow.__doc__ = pysparkfuncs.pow.__doc__ + + +def radians(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("radians", col) + + +radians.__doc__ = pysparkfuncs.radians.__doc__ + + +def rint(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("rint", col) + + +rint.__doc__ = pysparkfuncs.rint.__doc__ + + +def round(col: "ColumnOrName", scale: int = 0) -> Column: + return _invoke_function("round", _to_col(col), lit(scale)) + + +round.__doc__ = pysparkfuncs.round.__doc__ + + +def sec(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sec", col) + + +sec.__doc__ = pysparkfuncs.sec.__doc__ + + +def shiftLeft(col: "ColumnOrName", numBits: int) -> Column: + warnings.warn("Deprecated in 3.4, use shiftleft instead.", FutureWarning) + return shiftleft(col, numBits) + + +shiftLeft.__doc__ = pysparkfuncs.shiftLeft.__doc__ + + +def shiftleft(col: "ColumnOrName", numBits: int) -> Column: + return _invoke_function("shiftleft", _to_col(col), lit(numBits)) + + +shiftleft.__doc__ = pysparkfuncs.shiftleft.__doc__ + + +def shiftRight(col: "ColumnOrName", numBits: int) -> Column: + warnings.warn("Deprecated in 3.4, use shiftright instead.", FutureWarning) + return shiftright(col, numBits) + + +shiftRight.__doc__ = pysparkfuncs.shiftRight.__doc__ + + +def shiftright(col: "ColumnOrName", numBits: int) -> Column: + return _invoke_function("shiftright", _to_col(col), lit(numBits)) + + +shiftright.__doc__ = pysparkfuncs.shiftright.__doc__ + + +def shiftRightUnsigned(col: "ColumnOrName", numBits: int) -> Column: + warnings.warn("Deprecated in 3.4, use shiftrightunsigned instead.", FutureWarning) + return shiftrightunsigned(col, numBits) + + +shiftRightUnsigned.__doc__ = pysparkfuncs.shiftRightUnsigned.__doc__ + + +def shiftrightunsigned(col: "ColumnOrName", numBits: int) -> Column: + return _invoke_function("shiftrightunsigned", _to_col(col), lit(numBits)) + + +shiftrightunsigned.__doc__ = pysparkfuncs.shiftrightunsigned.__doc__ + + +def signum(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("signum", col) + + +signum.__doc__ = pysparkfuncs.signum.__doc__ + + +def sign(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sign", col) + + +sign.__doc__ = pysparkfuncs.sign.__doc__ + + +def sin(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sin", col) + + +sin.__doc__ = pysparkfuncs.sin.__doc__ + + +def sinh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sinh", col) + + +sinh.__doc__ = pysparkfuncs.sinh.__doc__ + + +def sqrt(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sqrt", col) + + +sqrt.__doc__ = pysparkfuncs.sqrt.__doc__ + + +def try_add(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_add", left, right) + + +try_add.__doc__ = pysparkfuncs.try_add.__doc__ + + +def try_avg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_avg", col) + + +try_avg.__doc__ = pysparkfuncs.try_avg.__doc__ + + +def try_divide(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_divide", left, right) + + +try_divide.__doc__ = pysparkfuncs.try_divide.__doc__ + + +def try_multiply(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_multiply", left, right) + + +try_multiply.__doc__ = pysparkfuncs.try_multiply.__doc__ + + +def try_subtract(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_subtract", left, right) + + +try_subtract.__doc__ = pysparkfuncs.try_subtract.__doc__ + + +def try_sum(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_sum", col) + + +try_sum.__doc__ = pysparkfuncs.try_sum.__doc__ + + +def tan(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("tan", col) + + +tan.__doc__ = pysparkfuncs.tan.__doc__ + + +def tanh(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("tanh", col) + + +tanh.__doc__ = pysparkfuncs.tanh.__doc__ + + +def toDegrees(col: "ColumnOrName") -> Column: + warnings.warn("Deprecated in 3.4, use degrees instead.", FutureWarning) + return degrees(col) + + +toDegrees.__doc__ = pysparkfuncs.toDegrees.__doc__ + + +def toRadians(col: "ColumnOrName") -> Column: + warnings.warn("Deprecated in 3.4, use radians instead.", FutureWarning) + return radians(col) + + +toRadians.__doc__ = pysparkfuncs.toRadians.__doc__ + + +def unhex(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unhex", col) + + +unhex.__doc__ = pysparkfuncs.unhex.__doc__ + + +def approxCountDistinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column: + warnings.warn("Deprecated in 3.4, use approx_count_distinct instead.", FutureWarning) + return approx_count_distinct(col, rsd) + + +approxCountDistinct.__doc__ = pysparkfuncs.approxCountDistinct.__doc__ + + +def approx_count_distinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column: + if rsd is None: + return _invoke_function("approx_count_distinct", _to_col(col)) + else: + return _invoke_function("approx_count_distinct", _to_col(col), lit(rsd)) + + +approx_count_distinct.__doc__ = pysparkfuncs.approx_count_distinct.__doc__ + + +def avg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("avg", col) + + +avg.__doc__ = pysparkfuncs.avg.__doc__ + + +def collect_list(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("collect_list", col) + + +collect_list.__doc__ = pysparkfuncs.collect_list.__doc__ + + +def array_agg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_agg", col) + + +array_agg.__doc__ = pysparkfuncs.array_agg.__doc__ + + +def collect_set(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("collect_set", col) + + +collect_set.__doc__ = pysparkfuncs.collect_set.__doc__ + + +def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("corr", col1, col2) + + +corr.__doc__ = pysparkfuncs.corr.__doc__ + + +def count(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("count", col) + + +count.__doc__ = pysparkfuncs.count.__doc__ + + +def countDistinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column: + return count_distinct(col, *cols) + + +countDistinct.__doc__ = pysparkfuncs.countDistinct.__doc__ + + +def count_distinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column: + _exprs = [_to_col(c)._expr for c in [col] + list(cols)] + return Column(UnresolvedFunction("count", _exprs, is_distinct=True)) + + +count_distinct.__doc__ = pysparkfuncs.count_distinct.__doc__ + + +def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("covar_pop", col1, col2) + + +covar_pop.__doc__ = pysparkfuncs.covar_pop.__doc__ + + +def covar_samp(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("covar_samp", col1, col2) + + +covar_samp.__doc__ = pysparkfuncs.covar_samp.__doc__ + + +def first(col: "ColumnOrName", ignorenulls: bool = False) -> Column: + return _invoke_function("first", _to_col(col), lit(ignorenulls)) + + +first.__doc__ = pysparkfuncs.first.__doc__ + + +def grouping(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("grouping", col) + + +grouping.__doc__ = pysparkfuncs.grouping.__doc__ + + +def grouping_id(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("grouping_id", *cols) + + +grouping_id.__doc__ = pysparkfuncs.grouping_id.__doc__ + + +def count_min_sketch( + col: "ColumnOrName", + eps: "ColumnOrName", + confidence: "ColumnOrName", + seed: "ColumnOrName", +) -> Column: + return _invoke_function_over_columns("count_min_sketch", col, eps, confidence, seed) + + +count_min_sketch.__doc__ = pysparkfuncs.count_min_sketch.__doc__ + + +def kurtosis(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("kurtosis", col) + + +kurtosis.__doc__ = pysparkfuncs.kurtosis.__doc__ + + +def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column: + return _invoke_function("last", _to_col(col), lit(ignorenulls)) + + +last.__doc__ = pysparkfuncs.last.__doc__ + + +def max(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("max", col) + + +max.__doc__ = pysparkfuncs.max.__doc__ + + +def max_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column: + return _invoke_function_over_columns("max_by", col, ord) + + +max_by.__doc__ = pysparkfuncs.max_by.__doc__ + + +def mean(col: "ColumnOrName") -> Column: + return avg(col) + + +mean.__doc__ = pysparkfuncs.mean.__doc__ + + +def median(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("median", col) + + +median.__doc__ = pysparkfuncs.median.__doc__ + + +def min(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("min", col) + + +min.__doc__ = pysparkfuncs.min.__doc__ + + +def min_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column: + return _invoke_function_over_columns("min_by", col, ord) + + +min_by.__doc__ = pysparkfuncs.min_by.__doc__ + + +def mode(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("mode", col) + + +mode.__doc__ = pysparkfuncs.mode.__doc__ + + +def percentile( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + frequency: Union[Column, int] = 1, +) -> Column: + if isinstance(percentage, Column): + _percentage = percentage + elif isinstance(percentage, (list, tuple)): + # Convert tuple to list + _percentage = lit(list(percentage)) + else: + # Probably scalar + _percentage = lit(percentage) + + if isinstance(frequency, int): + _frequency = lit(frequency) + elif isinstance(frequency, Column): + _frequency = frequency + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT", + message_parameters={ + "arg_name": "frequency", + "arg_type": type(frequency).__name__, + }, + ) + + return _invoke_function("percentile", _to_col(col), _percentage, _frequency) + + +percentile.__doc__ = pysparkfuncs.percentile.__doc__ + + +def percentile_approx( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + accuracy: Union[Column, float] = 10000, +) -> Column: + if isinstance(percentage, Column): + percentage_col = percentage + elif isinstance(percentage, (list, tuple)): + # Convert tuple to list + percentage_col = lit(list(percentage)) + else: + # Probably scalar + percentage_col = lit(percentage) + + return _invoke_function("percentile_approx", _to_col(col), percentage_col, lit(accuracy)) + + +percentile_approx.__doc__ = pysparkfuncs.percentile_approx.__doc__ + + +def approx_percentile( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + accuracy: Union[Column, float] = 10000, +) -> Column: + if isinstance(percentage, Column): + percentage_col = percentage + elif isinstance(percentage, (list, tuple)): + # Convert tuple to list + percentage_col = lit(list(percentage)) + else: + # Probably scalar + percentage_col = lit(percentage) + + return _invoke_function("approx_percentile", _to_col(col), percentage_col, lit(accuracy)) + + +approx_percentile.__doc__ = pysparkfuncs.approx_percentile.__doc__ + + +def product(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("product", col) + + +product.__doc__ = pysparkfuncs.product.__doc__ + + +def skewness(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("skewness", col) + + +skewness.__doc__ = pysparkfuncs.skewness.__doc__ + + +def stddev(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("stddev", col) + + +stddev.__doc__ = pysparkfuncs.stddev.__doc__ + + +def std(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("std", col) + + +std.__doc__ = pysparkfuncs.std.__doc__ + + +def stddev_samp(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("stddev_samp", col) + + +stddev_samp.__doc__ = pysparkfuncs.stddev_samp.__doc__ + + +def stddev_pop(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("stddev_pop", col) + + +stddev_pop.__doc__ = pysparkfuncs.stddev_pop.__doc__ + + +def sum(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sum", col) + + +sum.__doc__ = pysparkfuncs.sum.__doc__ + + +def sumDistinct(col: "ColumnOrName") -> Column: + warnings.warn("Deprecated in 3.4, use sum_distinct instead.", FutureWarning) + return sum_distinct(col) + + +sumDistinct.__doc__ = pysparkfuncs.sumDistinct.__doc__ + + +def sum_distinct(col: "ColumnOrName") -> Column: + return Column(UnresolvedFunction("sum", [_to_col(col)._expr], is_distinct=True)) + + +sum_distinct.__doc__ = pysparkfuncs.sum_distinct.__doc__ + + +def var_pop(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("var_pop", col) + + +var_pop.__doc__ = pysparkfuncs.var_pop.__doc__ + + +def regr_avgx(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_avgx", y, x) + + +regr_avgx.__doc__ = pysparkfuncs.regr_avgx.__doc__ + + +def regr_avgy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_avgy", y, x) + + +regr_avgy.__doc__ = pysparkfuncs.regr_avgy.__doc__ + + +def regr_count(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_count", y, x) + + +regr_count.__doc__ = pysparkfuncs.regr_count.__doc__ + + +def regr_intercept(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_intercept", y, x) + + +regr_intercept.__doc__ = pysparkfuncs.regr_intercept.__doc__ + + +def regr_r2(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_r2", y, x) + + +regr_r2.__doc__ = pysparkfuncs.regr_r2.__doc__ + + +def regr_slope(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_slope", y, x) + + +regr_slope.__doc__ = pysparkfuncs.regr_slope.__doc__ + + +def regr_sxx(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_sxx", y, x) + + +regr_sxx.__doc__ = pysparkfuncs.regr_sxx.__doc__ + + +def regr_sxy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_sxy", y, x) + + +regr_sxy.__doc__ = pysparkfuncs.regr_sxy.__doc__ + + +def regr_syy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regr_syy", y, x) + + +regr_syy.__doc__ = pysparkfuncs.regr_syy.__doc__ + + +def var_samp(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("var_samp", col) + + +var_samp.__doc__ = pysparkfuncs.var_samp.__doc__ + + +def variance(col: "ColumnOrName") -> Column: + return var_samp(col) + + +variance.__doc__ = pysparkfuncs.variance.__doc__ + + +def every(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("every", col) + + +every.__doc__ = pysparkfuncs.every.__doc__ + + +def bool_and(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bool_and", col) + + +bool_and.__doc__ = pysparkfuncs.bool_and.__doc__ + + +def some(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("some", col) + + +some.__doc__ = pysparkfuncs.some.__doc__ + + +def bool_or(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bool_or", col) + + +bool_or.__doc__ = pysparkfuncs.bool_or.__doc__ + + +def bit_and(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_and", col) + + +bit_and.__doc__ = pysparkfuncs.bit_and.__doc__ + + +def bit_or(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_or", col) + + +bit_or.__doc__ = pysparkfuncs.bit_or.__doc__ + + +def bit_xor(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_xor", col) + + +bit_xor.__doc__ = pysparkfuncs.bit_xor.__doc__ + + +# Window Functions + + +def cume_dist() -> Column: + return _invoke_function("cume_dist") + + +cume_dist.__doc__ = pysparkfuncs.cume_dist.__doc__ + + +def dense_rank() -> Column: + return _invoke_function("dense_rank") + + +dense_rank.__doc__ = pysparkfuncs.dense_rank.__doc__ + + +def lag(col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None) -> Column: + if default is None: + return _invoke_function("lag", _to_col(col), lit(offset)) + else: + return _invoke_function("lag", _to_col(col), lit(offset), lit(default)) + + +lag.__doc__ = pysparkfuncs.lag.__doc__ + + +def lead(col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None) -> Column: + if default is None: + return _invoke_function("lead", _to_col(col), lit(offset)) + else: + return _invoke_function("lead", _to_col(col), lit(offset), lit(default)) + + +lead.__doc__ = pysparkfuncs.lead.__doc__ + + +def nth_value(col: "ColumnOrName", offset: int, ignoreNulls: Optional[bool] = None) -> Column: + if ignoreNulls is None: + return _invoke_function("nth_value", _to_col(col), lit(offset)) + else: + return _invoke_function("nth_value", _to_col(col), lit(offset), lit(ignoreNulls)) + + +nth_value.__doc__ = pysparkfuncs.nth_value.__doc__ + + +def any_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + if ignoreNulls is None: + return _invoke_function_over_columns("any_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("any_value", col, ignoreNulls) + + +any_value.__doc__ = pysparkfuncs.any_value.__doc__ + + +def first_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + if ignoreNulls is None: + return _invoke_function_over_columns("first_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("first_value", col, ignoreNulls) + + +first_value.__doc__ = pysparkfuncs.first_value.__doc__ + + +def last_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + if ignoreNulls is None: + return _invoke_function_over_columns("last_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("last_value", col, ignoreNulls) + + +last_value.__doc__ = pysparkfuncs.last_value.__doc__ + + +def count_if(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("count_if", col) + + +count_if.__doc__ = pysparkfuncs.count_if.__doc__ + + +def histogram_numeric(col: "ColumnOrName", nBins: "ColumnOrName") -> Column: + return _invoke_function_over_columns("histogram_numeric", col, nBins) + + +histogram_numeric.__doc__ = pysparkfuncs.histogram_numeric.__doc__ + + +def ntile(n: int) -> Column: + return _invoke_function("ntile", lit(n)) + + +ntile.__doc__ = pysparkfuncs.ntile.__doc__ + + +def percent_rank() -> Column: + return _invoke_function("percent_rank") + + +percent_rank.__doc__ = pysparkfuncs.percent_rank.__doc__ + + +def rank() -> Column: + return _invoke_function("rank") + + +rank.__doc__ = pysparkfuncs.rank.__doc__ + + +def row_number() -> Column: + return _invoke_function("row_number") + + +row_number.__doc__ = pysparkfuncs.row_number.__doc__ + + +def aggregate( + col: "ColumnOrName", + initialValue: "ColumnOrName", + merge: Callable[[Column, Column], Column], + finish: Optional[Callable[[Column], Column]] = None, +) -> Column: + if finish is not None: + return _invoke_higher_order_function("aggregate", [col, initialValue], [merge, finish]) + + else: + return _invoke_higher_order_function("aggregate", [col, initialValue], [merge]) + + +aggregate.__doc__ = pysparkfuncs.aggregate.__doc__ + + +def reduce( + col: "ColumnOrName", + initialValue: "ColumnOrName", + merge: Callable[[Column, Column], Column], + finish: Optional[Callable[[Column], Column]] = None, +) -> Column: + if finish is not None: + return _invoke_higher_order_function("reduce", [col, initialValue], [merge, finish]) + + else: + return _invoke_higher_order_function("reduce", [col, initialValue], [merge]) + + +reduce.__doc__ = pysparkfuncs.reduce.__doc__ + + +def array(*cols: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]) -> Column: + if len(cols) == 1 and isinstance(cols[0], (list, set, tuple)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_columns("array", *cols) # type: ignore[arg-type] + + +array.__doc__ = pysparkfuncs.array.__doc__ + + +def array_append(col: "ColumnOrName", value: Any) -> Column: + return _invoke_function("array_append", _to_col(col), lit(value)) + + +array_append.__doc__ = pysparkfuncs.array_append.__doc__ + + +def array_contains(col: "ColumnOrName", value: Any) -> Column: + return _invoke_function("array_contains", _to_col(col), lit(value)) + + +array_contains.__doc__ = pysparkfuncs.array_contains.__doc__ + + +def array_distinct(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_distinct", col) + + +array_distinct.__doc__ = pysparkfuncs.array_distinct.__doc__ + + +def array_except(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_except", col1, col2) + + +array_except.__doc__ = pysparkfuncs.array_except.__doc__ + + +def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Any) -> Column: + _pos = lit(pos) if isinstance(pos, int) else _to_col(pos) + return _invoke_function("array_insert", _to_col(arr), _pos, lit(value)) + + +array_insert.__doc__ = pysparkfuncs.array_insert.__doc__ + + +def array_intersect(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_intersect", col1, col2) + + +array_intersect.__doc__ = pysparkfuncs.array_intersect.__doc__ + + +def array_compact(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_compact", col) + + +array_compact.__doc__ = pysparkfuncs.array_compact.__doc__ + + +def array_join( + col: "ColumnOrName", delimiter: str, null_replacement: Optional[str] = None +) -> Column: + if null_replacement is None: + return _invoke_function("array_join", _to_col(col), lit(delimiter)) + else: + return _invoke_function("array_join", _to_col(col), lit(delimiter), lit(null_replacement)) + + +array_join.__doc__ = pysparkfuncs.array_join.__doc__ + + +def array_max(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_max", col) + + +array_max.__doc__ = pysparkfuncs.array_max.__doc__ + + +def array_min(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_min", col) + + +array_min.__doc__ = pysparkfuncs.array_min.__doc__ + + +def array_size(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_size", col) + + +array_size.__doc__ = pysparkfuncs.array_size.__doc__ + + +def cardinality(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("cardinality", col) + + +cardinality.__doc__ = pysparkfuncs.cardinality.__doc__ + + +def array_position(col: "ColumnOrName", value: Any) -> Column: + return _invoke_function("array_position", _to_col(col), lit(value)) + + +array_position.__doc__ = pysparkfuncs.array_position.__doc__ + + +def array_prepend(col: "ColumnOrName", value: Any) -> Column: + return _invoke_function("array_prepend", _to_col(col), lit(value)) + + +array_prepend.__doc__ = pysparkfuncs.array_prepend.__doc__ + + +def array_remove(col: "ColumnOrName", element: Any) -> Column: + return _invoke_function("array_remove", _to_col(col), lit(element)) + + +array_remove.__doc__ = pysparkfuncs.array_remove.__doc__ + + +def array_repeat(col: "ColumnOrName", count: Union["ColumnOrName", int]) -> Column: + _count = lit(count) if isinstance(count, int) else _to_col(count) + return _invoke_function("array_repeat", _to_col(col), _count) + + +array_repeat.__doc__ = pysparkfuncs.array_repeat.__doc__ + + +def array_sort( + col: "ColumnOrName", comparator: Optional[Callable[[Column, Column], Column]] = None +) -> Column: + if comparator is None: + return _invoke_function_over_columns("array_sort", col) + else: + return _invoke_higher_order_function("array_sort", [col], [comparator]) + + +array_sort.__doc__ = pysparkfuncs.array_sort.__doc__ + + +def array_union(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("array_union", col1, col2) + + +array_union.__doc__ = pysparkfuncs.array_union.__doc__ + + +def arrays_overlap(a1: "ColumnOrName", a2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("arrays_overlap", a1, a2) + + +arrays_overlap.__doc__ = pysparkfuncs.arrays_overlap.__doc__ + + +def arrays_zip(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("arrays_zip", *cols) + + +arrays_zip.__doc__ = pysparkfuncs.arrays_zip.__doc__ + + +def concat(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("concat", *cols) + + +concat.__doc__ = pysparkfuncs.concat.__doc__ + + +def create_map( + *cols: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]] +) -> Column: + if len(cols) == 1 and isinstance(cols[0], (list, set, tuple)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_columns("map", *cols) # type: ignore[arg-type] + + +create_map.__doc__ = pysparkfuncs.create_map.__doc__ + + +def element_at(col: "ColumnOrName", extraction: Any) -> Column: + return _invoke_function("element_at", _to_col(col), lit(extraction)) + + +element_at.__doc__ = pysparkfuncs.element_at.__doc__ + + +def try_element_at(col: "ColumnOrName", extraction: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_element_at", col, extraction) + + +try_element_at.__doc__ = pysparkfuncs.try_element_at.__doc__ + + +def exists(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + return _invoke_higher_order_function("exists", [col], [f]) + + +exists.__doc__ = pysparkfuncs.exists.__doc__ + + +def explode(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("explode", col) + + +explode.__doc__ = pysparkfuncs.explode.__doc__ + + +def explode_outer(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("explode_outer", col) + + +explode_outer.__doc__ = pysparkfuncs.explode_outer.__doc__ + + +def filter( + col: "ColumnOrName", + f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]], +) -> Column: + return _invoke_higher_order_function("filter", [col], [f]) + + +filter.__doc__ = pysparkfuncs.filter.__doc__ + + +def flatten(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("flatten", col) + + +flatten.__doc__ = pysparkfuncs.flatten.__doc__ + + +def forall(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + return _invoke_higher_order_function("forall", [col], [f]) + + +forall.__doc__ = pysparkfuncs.forall.__doc__ + + +# TODO: support options +def from_csv( + col: "ColumnOrName", + schema: Union[Column, str], + options: Optional[Dict[str, str]] = None, +) -> Column: + if isinstance(schema, Column): + _schema = schema + elif isinstance(schema, str): + _schema = lit(schema) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + + if options is None: + return _invoke_function("from_csv", _to_col(col), _schema) + else: + return _invoke_function("from_csv", _to_col(col), _schema, _options_to_col(options)) + + +from_csv.__doc__ = pysparkfuncs.from_csv.__doc__ + + +def from_json( + col: "ColumnOrName", + schema: Union[ArrayType, StructType, Column, str], + options: Optional[Dict[str, str]] = None, +) -> Column: + if isinstance(schema, Column): + _schema = schema + elif isinstance(schema, DataType): + _schema = lit(schema.json()) + elif isinstance(schema, str): + _schema = lit(schema) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_DATATYPE_OR_STR", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + + if options is None: + return _invoke_function("from_json", _to_col(col), _schema) + else: + return _invoke_function("from_json", _to_col(col), _schema, _options_to_col(options)) + + +from_json.__doc__ = pysparkfuncs.from_json.__doc__ + + +def get(col: "ColumnOrName", index: Union["ColumnOrName", int]) -> Column: + index = lit(index) if isinstance(index, int) else index + + return _invoke_function_over_columns("get", col, index) + + +get.__doc__ = pysparkfuncs.get.__doc__ + + +def get_json_object(col: "ColumnOrName", path: str) -> Column: + return _invoke_function("get_json_object", _to_col(col), lit(path)) + + +get_json_object.__doc__ = pysparkfuncs.get_json_object.__doc__ + + +def json_array_length(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("json_array_length", col) + + +json_array_length.__doc__ = pysparkfuncs.json_array_length.__doc__ + + +def json_object_keys(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("json_object_keys", col) + + +json_object_keys.__doc__ = pysparkfuncs.json_object_keys.__doc__ + + +def inline(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("inline", col) + + +inline.__doc__ = pysparkfuncs.inline.__doc__ + + +def inline_outer(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("inline_outer", col) + + +inline_outer.__doc__ = pysparkfuncs.inline_outer.__doc__ + + +def json_tuple(col: "ColumnOrName", *fields: str) -> Column: + return _invoke_function("json_tuple", _to_col(col), *[lit(field) for field in fields]) + + +json_tuple.__doc__ = pysparkfuncs.json_tuple.__doc__ + + +def map_concat( + *cols: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]] +) -> Column: + if len(cols) == 1 and isinstance(cols[0], (list, set, tuple)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_columns("map_concat", *cols) # type: ignore[arg-type] + + +map_concat.__doc__ = pysparkfuncs.map_concat.__doc__ + + +def map_contains_key(col: "ColumnOrName", value: Any) -> Column: + return array_contains(map_keys(col), lit(value)) + + +map_contains_key.__doc__ = pysparkfuncs.map_contains_key.__doc__ + + +def map_entries(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("map_entries", col) + + +map_entries.__doc__ = pysparkfuncs.map_entries.__doc__ + + +def map_filter(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + return _invoke_higher_order_function("map_filter", [col], [f]) + + +map_filter.__doc__ = pysparkfuncs.map_filter.__doc__ + + +def map_from_arrays(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("map_from_arrays", col1, col2) + + +map_from_arrays.__doc__ = pysparkfuncs.map_from_arrays.__doc__ + + +def map_from_entries(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("map_from_entries", col) + + +map_from_entries.__doc__ = pysparkfuncs.map_from_entries.__doc__ + + +def map_keys(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("map_keys", col) + + +map_keys.__doc__ = pysparkfuncs.map_keys.__doc__ + + +def map_values(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("map_values", col) + + +map_values.__doc__ = pysparkfuncs.map_values.__doc__ + + +def map_zip_with( + col1: "ColumnOrName", + col2: "ColumnOrName", + f: Callable[[Column, Column, Column], Column], +) -> Column: + return _invoke_higher_order_function("map_zip_with", [col1, col2], [f]) + + +map_zip_with.__doc__ = pysparkfuncs.map_zip_with.__doc__ + + +def str_to_map( + text: "ColumnOrName", + pairDelim: Optional["ColumnOrName"] = None, + keyValueDelim: Optional["ColumnOrName"] = None, +) -> Column: + _pairDelim = lit(",") if pairDelim is None else _to_col(pairDelim) + _keyValueDelim = lit(":") if keyValueDelim is None else _to_col(keyValueDelim) + + return _invoke_function("str_to_map", _to_col(text), _pairDelim, _keyValueDelim) + + +str_to_map.__doc__ = pysparkfuncs.str_to_map.__doc__ + + +def posexplode(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("posexplode", col) + + +posexplode.__doc__ = pysparkfuncs.posexplode.__doc__ + + +def posexplode_outer(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("posexplode_outer", col) + + +posexplode_outer.__doc__ = pysparkfuncs.posexplode_outer.__doc__ + + +def reverse(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("reverse", col) + + +reverse.__doc__ = pysparkfuncs.reverse.__doc__ + + +def sequence( + start: "ColumnOrName", stop: "ColumnOrName", step: Optional["ColumnOrName"] = None +) -> Column: + if step is None: + return _invoke_function_over_columns("sequence", start, stop) + else: + return _invoke_function_over_columns("sequence", start, stop, step) + + +sequence.__doc__ = pysparkfuncs.sequence.__doc__ + + +def schema_of_csv(csv: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + if isinstance(csv, Column): + _csv = csv + elif isinstance(csv, str): + _csv = lit(csv) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "csv", "arg_type": type(csv).__name__}, + ) + + if options is None: + return _invoke_function("schema_of_csv", _csv) + else: + return _invoke_function("schema_of_csv", _csv, _options_to_col(options)) + + +schema_of_csv.__doc__ = pysparkfuncs.schema_of_csv.__doc__ + + +def schema_of_json(json: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + if isinstance(json, Column): + _json = json + elif isinstance(json, str): + _json = lit(json) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "json", "arg_type": type(json).__name__}, + ) + + if options is None: + return _invoke_function("schema_of_json", _json) + else: + return _invoke_function("schema_of_json", _json, _options_to_col(options)) + + +schema_of_json.__doc__ = pysparkfuncs.schema_of_json.__doc__ + + +def shuffle(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("shuffle", col) + + +shuffle.__doc__ = pysparkfuncs.shuffle.__doc__ + + +def size(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("size", col) + + +size.__doc__ = pysparkfuncs.size.__doc__ + + +def slice( + col: "ColumnOrName", start: Union["ColumnOrName", int], length: Union["ColumnOrName", int] +) -> Column: + if isinstance(start, (Column, str)): + _start = start + elif isinstance(start, int): + _start = lit(start) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "start", "arg_type": type(start).__name__}, + ) + + if isinstance(length, (Column, str)): + _length = length + elif isinstance(length, int): + _length = lit(length) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "length", "arg_type": type(length).__name__}, + ) + + return _invoke_function_over_columns("slice", col, _start, _length) + + +slice.__doc__ = pysparkfuncs.slice.__doc__ + + +def sort_array(col: "ColumnOrName", asc: bool = True) -> Column: + return _invoke_function("sort_array", _to_col(col), lit(asc)) + + +sort_array.__doc__ = pysparkfuncs.sort_array.__doc__ + + +def struct( + *cols: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]] +) -> Column: + if len(cols) == 1 and isinstance(cols[0], (list, set, tuple)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_columns("struct", *cols) # type: ignore[arg-type] + + +struct.__doc__ = pysparkfuncs.struct.__doc__ + + +def named_struct(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("named_struct", *cols) + + +named_struct.__doc__ = pysparkfuncs.named_struct.__doc__ + + +def to_csv(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + if options is None: + return _invoke_function("to_csv", _to_col(col)) + else: + return _invoke_function("to_csv", _to_col(col), _options_to_col(options)) + + +to_csv.__doc__ = pysparkfuncs.to_csv.__doc__ + + +def to_json(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + if options is None: + return _invoke_function("to_json", _to_col(col)) + else: + return _invoke_function("to_json", _to_col(col), _options_to_col(options)) + + +to_json.__doc__ = pysparkfuncs.to_json.__doc__ + + +def transform( + col: "ColumnOrName", + f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]], +) -> Column: + return _invoke_higher_order_function("transform", [col], [f]) + + +transform.__doc__ = pysparkfuncs.transform.__doc__ + + +def transform_keys(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + return _invoke_higher_order_function("transform_keys", [col], [f]) + + +transform_keys.__doc__ = pysparkfuncs.transform_keys.__doc__ + + +def transform_values(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + return _invoke_higher_order_function("transform_values", [col], [f]) + + +transform_values.__doc__ = pysparkfuncs.transform_values.__doc__ + + +def zip_with( + left: "ColumnOrName", + right: "ColumnOrName", + f: Callable[[Column, Column], Column], +) -> Column: + return _invoke_higher_order_function("zip_with", [left, right], [f]) + + +zip_with.__doc__ = pysparkfuncs.zip_with.__doc__ + + +# String/Binary functions + + +def upper(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("upper", col) + + +upper.__doc__ = pysparkfuncs.upper.__doc__ + + +def lower(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("lower", col) + + +lower.__doc__ = pysparkfuncs.lower.__doc__ + + +def ascii(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ascii", col) + + +ascii.__doc__ = pysparkfuncs.ascii.__doc__ + + +def base64(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("base64", col) + + +base64.__doc__ = pysparkfuncs.base64.__doc__ + + +def unbase64(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unbase64", col) + + +unbase64.__doc__ = pysparkfuncs.unbase64.__doc__ + + +def ltrim(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ltrim", col) + + +ltrim.__doc__ = pysparkfuncs.ltrim.__doc__ + + +def rtrim(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("rtrim", col) + + +rtrim.__doc__ = pysparkfuncs.rtrim.__doc__ + + +def trim(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("trim", col) + + +trim.__doc__ = pysparkfuncs.trim.__doc__ + + +def concat_ws(sep: str, *cols: "ColumnOrName") -> Column: + return _invoke_function("concat_ws", lit(sep), *[_to_col(c) for c in cols]) + + +concat_ws.__doc__ = pysparkfuncs.concat_ws.__doc__ + + +def decode(col: "ColumnOrName", charset: str) -> Column: + return _invoke_function("decode", _to_col(col), lit(charset)) + + +decode.__doc__ = pysparkfuncs.decode.__doc__ + + +def encode(col: "ColumnOrName", charset: str) -> Column: + return _invoke_function("encode", _to_col(col), lit(charset)) + + +encode.__doc__ = pysparkfuncs.encode.__doc__ + + +def format_number(col: "ColumnOrName", d: int) -> Column: + return _invoke_function("format_number", _to_col(col), lit(d)) + + +format_number.__doc__ = pysparkfuncs.format_number.__doc__ + + +def format_string(format: str, *cols: "ColumnOrName") -> Column: + return _invoke_function("format_string", lit(format), *[_to_col(c) for c in cols]) + + +format_string.__doc__ = pysparkfuncs.format_string.__doc__ + + +def instr(str: "ColumnOrName", substr: str) -> Column: + return _invoke_function("instr", _to_col(str), lit(substr)) + + +instr.__doc__ = pysparkfuncs.instr.__doc__ + + +def overlay( + src: "ColumnOrName", + replace: "ColumnOrName", + pos: Union["ColumnOrName", int], + len: Union["ColumnOrName", int] = -1, +) -> Column: + if not isinstance(pos, (int, str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "pos", "arg_type": type(pos).__name__}, + ) + if len is not None and not isinstance(len, (int, str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "len", "arg_type": type(len).__name__}, + ) + + if isinstance(pos, int): + pos = lit(pos) + if isinstance(len, int): + len = lit(len) + + return _invoke_function_over_columns("overlay", src, replace, pos, len) + + +overlay.__doc__ = pysparkfuncs.overlay.__doc__ + + +def sentences( + string: "ColumnOrName", + language: Optional["ColumnOrName"] = None, + country: Optional["ColumnOrName"] = None, +) -> Column: + _language = lit("") if language is None else _to_col(language) + _country = lit("") if country is None else _to_col(country) + + return _invoke_function("sentences", _to_col(string), _language, _country) + + +sentences.__doc__ = pysparkfuncs.sentences.__doc__ + + +def substring(str: "ColumnOrName", pos: int, len: int) -> Column: + return _invoke_function("substring", _to_col(str), lit(pos), lit(len)) + + +substring.__doc__ = pysparkfuncs.substring.__doc__ + + +def substring_index(str: "ColumnOrName", delim: str, count: int) -> Column: + return _invoke_function("substring_index", _to_col(str), lit(delim), lit(count)) + + +substring_index.__doc__ = pysparkfuncs.substring_index.__doc__ + + +def levenshtein( + left: "ColumnOrName", right: "ColumnOrName", threshold: Optional[int] = None +) -> Column: + if threshold is None: + return _invoke_function_over_columns("levenshtein", left, right) + else: + return _invoke_function("levenshtein", _to_col(left), _to_col(right), lit(threshold)) + + +levenshtein.__doc__ = pysparkfuncs.levenshtein.__doc__ + + +def locate(substr: str, str: "ColumnOrName", pos: int = 1) -> Column: + return _invoke_function("locate", lit(substr), _to_col(str), lit(pos)) + + +locate.__doc__ = pysparkfuncs.locate.__doc__ + + +def lpad(col: "ColumnOrName", len: int, pad: str) -> Column: + return _invoke_function("lpad", _to_col(col), lit(len), lit(pad)) + + +lpad.__doc__ = pysparkfuncs.lpad.__doc__ + + +def rpad(col: "ColumnOrName", len: int, pad: str) -> Column: + return _invoke_function("rpad", _to_col(col), lit(len), lit(pad)) + + +rpad.__doc__ = pysparkfuncs.rpad.__doc__ + + +def repeat(col: "ColumnOrName", n: int) -> Column: + return _invoke_function("repeat", _to_col(col), lit(n)) + + +repeat.__doc__ = pysparkfuncs.repeat.__doc__ + + +def split(str: "ColumnOrName", pattern: str, limit: int = -1) -> Column: + return _invoke_function("split", _to_col(str), lit(pattern), lit(limit)) + + +split.__doc__ = pysparkfuncs.split.__doc__ + + +def rlike(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + return _invoke_function_over_columns("rlike", str, regexp) + + +rlike.__doc__ = pysparkfuncs.rlike.__doc__ + + +def regexp(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regexp", str, regexp) + + +regexp.__doc__ = pysparkfuncs.regexp.__doc__ + + +def regexp_like(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regexp_like", str, regexp) + + +regexp_like.__doc__ = pysparkfuncs.regexp_like.__doc__ + + +def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regexp_count", str, regexp) + + +regexp_count.__doc__ = pysparkfuncs.regexp_count.__doc__ + + +def regexp_extract(str: "ColumnOrName", pattern: str, idx: int) -> Column: + return _invoke_function("regexp_extract", _to_col(str), lit(pattern), lit(idx)) + + +regexp_extract.__doc__ = pysparkfuncs.regexp_extract.__doc__ + + +def regexp_extract_all( + str: "ColumnOrName", regexp: "ColumnOrName", idx: Optional[Union[int, Column]] = None +) -> Column: + if idx is None: + return _invoke_function_over_columns("regexp_extract_all", str, regexp) + else: + if isinstance(idx, int): + idx = lit(idx) + return _invoke_function_over_columns("regexp_extract_all", str, regexp, idx) + + +regexp_extract_all.__doc__ = pysparkfuncs.regexp_extract_all.__doc__ + + +def regexp_replace( + string: "ColumnOrName", pattern: Union[str, Column], replacement: Union[str, Column] +) -> Column: + if isinstance(pattern, str): + pattern = lit(pattern) + + if isinstance(replacement, str): + replacement = lit(replacement) + + return _invoke_function("regexp_replace", _to_col(string), pattern, replacement) + + +regexp_replace.__doc__ = pysparkfuncs.regexp_replace.__doc__ + + +def regexp_substr(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + return _invoke_function_over_columns("regexp_substr", str, regexp) + + +regexp_substr.__doc__ = pysparkfuncs.regexp_substr.__doc__ + + +def regexp_instr( + str: "ColumnOrName", regexp: "ColumnOrName", idx: Optional[Union[int, Column]] = None +) -> Column: + if idx is None: + return _invoke_function_over_columns("regexp_instr", str, regexp) + else: + if isinstance(idx, int): + idx = lit(idx) + return _invoke_function_over_columns("regexp_instr", str, regexp, idx) + + +regexp_instr.__doc__ = pysparkfuncs.regexp_instr.__doc__ + + +def initcap(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("initcap", col) + + +initcap.__doc__ = pysparkfuncs.initcap.__doc__ + + +def soundex(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("soundex", col) + + +soundex.__doc__ = pysparkfuncs.soundex.__doc__ + + +def length(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("length", col) + + +length.__doc__ = pysparkfuncs.length.__doc__ + + +def octet_length(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("octet_length", col) + + +octet_length.__doc__ = pysparkfuncs.octet_length.__doc__ + + +def bit_length(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bit_length", col) + + +bit_length.__doc__ = pysparkfuncs.bit_length.__doc__ + + +def translate(srcCol: "ColumnOrName", matching: str, replace: str) -> Column: + return _invoke_function("translate", _to_col(srcCol), lit(matching), lit(replace)) + + +translate.__doc__ = pysparkfuncs.translate.__doc__ + + +def to_binary(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + if format is not None: + return _invoke_function_over_columns("to_binary", col, format) + else: + return _invoke_function_over_columns("to_binary", col) + + +to_binary.__doc__ = pysparkfuncs.to_binary.__doc__ + + +def to_char(col: "ColumnOrName", format: "ColumnOrName") -> Column: + return _invoke_function_over_columns("to_char", col, format) + + +to_char.__doc__ = pysparkfuncs.to_char.__doc__ + + +def to_varchar(col: "ColumnOrName", format: "ColumnOrName") -> Column: + return _invoke_function_over_columns("to_varchar", col, format) + + +to_varchar.__doc__ = pysparkfuncs.to_varchar.__doc__ + + +def to_number(col: "ColumnOrName", format: "ColumnOrName") -> Column: + return _invoke_function_over_columns("to_number", col, format) + + +to_number.__doc__ = pysparkfuncs.to_number.__doc__ + + +def replace( + src: "ColumnOrName", search: "ColumnOrName", replace: Optional["ColumnOrName"] = None +) -> Column: + if replace is not None: + return _invoke_function_over_columns("replace", src, search, replace) + else: + return _invoke_function_over_columns("replace", src, search) + + +replace.__doc__ = pysparkfuncs.replace.__doc__ + + +def split_part(src: "ColumnOrName", delimiter: "ColumnOrName", partNum: "ColumnOrName") -> Column: + return _invoke_function_over_columns("split_part", src, delimiter, partNum) + + +split_part.__doc__ = pysparkfuncs.split_part.__doc__ + + +def substr( + str: "ColumnOrName", pos: "ColumnOrName", len: Optional["ColumnOrName"] = None +) -> Column: + if len is not None: + return _invoke_function_over_columns("substr", str, pos, len) + else: + return _invoke_function_over_columns("substr", str, pos) + + +substr.__doc__ = pysparkfuncs.substr.__doc__ + + +def parse_url( + url: "ColumnOrName", partToExtract: "ColumnOrName", key: Optional["ColumnOrName"] = None +) -> Column: + if key is not None: + return _invoke_function_over_columns("parse_url", url, partToExtract, key) + else: + return _invoke_function_over_columns("parse_url", url, partToExtract) + + +parse_url.__doc__ = pysparkfuncs.parse_url.__doc__ + + +def printf(format: "ColumnOrName", *cols: "ColumnOrName") -> Column: + return _invoke_function("printf", _to_col(format), *[_to_col(c) for c in cols]) + + +printf.__doc__ = pysparkfuncs.printf.__doc__ + + +def url_decode(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("url_decode", str) + + +url_decode.__doc__ = pysparkfuncs.url_decode.__doc__ + + +def url_encode(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("url_encode", str) + + +url_encode.__doc__ = pysparkfuncs.url_encode.__doc__ + + +def position( + substr: "ColumnOrName", str: "ColumnOrName", start: Optional["ColumnOrName"] = None +) -> Column: + if start is not None: + return _invoke_function_over_columns("position", substr, str, start) + else: + return _invoke_function_over_columns("position", substr, str) + + +position.__doc__ = pysparkfuncs.position.__doc__ + + +def endswith(str: "ColumnOrName", suffix: "ColumnOrName") -> Column: + return _invoke_function_over_columns("endswith", str, suffix) + + +endswith.__doc__ = pysparkfuncs.endswith.__doc__ + + +def startswith(str: "ColumnOrName", prefix: "ColumnOrName") -> Column: + return _invoke_function_over_columns("startswith", str, prefix) + + +startswith.__doc__ = pysparkfuncs.startswith.__doc__ + + +def char(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("char", col) + + +char.__doc__ = pysparkfuncs.char.__doc__ + + +def try_to_binary(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + if format is not None: + return _invoke_function_over_columns("try_to_binary", col, format) + else: + return _invoke_function_over_columns("try_to_binary", col) + + +try_to_binary.__doc__ = pysparkfuncs.try_to_binary.__doc__ + + +def try_to_number(col: "ColumnOrName", format: "ColumnOrName") -> Column: + return _invoke_function_over_columns("try_to_number", col, format) + + +try_to_number.__doc__ = pysparkfuncs.try_to_number.__doc__ + + +def btrim(str: "ColumnOrName", trim: Optional["ColumnOrName"] = None) -> Column: + if trim is not None: + return _invoke_function_over_columns("btrim", str, trim) + else: + return _invoke_function_over_columns("btrim", str) + + +btrim.__doc__ = pysparkfuncs.btrim.__doc__ + + +def char_length(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("char_length", str) + + +char_length.__doc__ = pysparkfuncs.char_length.__doc__ + + +def character_length(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("character_length", str) + + +character_length.__doc__ = pysparkfuncs.character_length.__doc__ + + +def contains(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("contains", left, right) + + +contains.__doc__ = pysparkfuncs.contains.__doc__ + + +def elt(*inputs: "ColumnOrName") -> Column: + return _invoke_function("elt", *[_to_col(input) for input in inputs]) + + +elt.__doc__ = pysparkfuncs.elt.__doc__ + + +def find_in_set(str: "ColumnOrName", str_array: "ColumnOrName") -> Column: + return _invoke_function_over_columns("find_in_set", str, str_array) + + +find_in_set.__doc__ = pysparkfuncs.find_in_set.__doc__ + + +def like( + str: "ColumnOrName", pattern: "ColumnOrName", escapeChar: Optional["Column"] = None +) -> Column: + if escapeChar is not None: + return _invoke_function_over_columns("like", str, pattern, escapeChar) + else: + return _invoke_function_over_columns("like", str, pattern) + + +like.__doc__ = pysparkfuncs.like.__doc__ + + +def ilike( + str: "ColumnOrName", pattern: "ColumnOrName", escapeChar: Optional["Column"] = None +) -> Column: + if escapeChar is not None: + return _invoke_function_over_columns("ilike", str, pattern, escapeChar) + else: + return _invoke_function_over_columns("ilike", str, pattern) + + +ilike.__doc__ = pysparkfuncs.ilike.__doc__ + + +def lcase(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("lcase", str) + + +lcase.__doc__ = pysparkfuncs.lcase.__doc__ + + +def ucase(str: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ucase", str) + + +ucase.__doc__ = pysparkfuncs.ucase.__doc__ + + +def left(str: "ColumnOrName", len: "ColumnOrName") -> Column: + return _invoke_function_over_columns("left", str, len) + + +left.__doc__ = pysparkfuncs.left.__doc__ + + +def right(str: "ColumnOrName", len: "ColumnOrName") -> Column: + return _invoke_function_over_columns("right", str, len) + + +right.__doc__ = pysparkfuncs.right.__doc__ + + +def mask( + col: "ColumnOrName", + upperChar: Optional["ColumnOrName"] = None, + lowerChar: Optional["ColumnOrName"] = None, + digitChar: Optional["ColumnOrName"] = None, + otherChar: Optional["ColumnOrName"] = None, +) -> Column: + _upperChar = lit("X") if upperChar is None else upperChar + _lowerChar = lit("x") if lowerChar is None else lowerChar + _digitChar = lit("n") if digitChar is None else digitChar + _otherChar = lit(None) if otherChar is None else otherChar + + return _invoke_function_over_columns( + "mask", col, _upperChar, _lowerChar, _digitChar, _otherChar + ) + + +mask.__doc__ = pysparkfuncs.mask.__doc__ + + +# Date/Timestamp functions +# TODO(SPARK-41455): Resolve dtypes inconsistencies for: +# to_timestamp, from_utc_timestamp, to_utc_timestamp, +# timestamp_seconds, current_timestamp, date_trunc + + +def curdate() -> Column: + return _invoke_function("curdate") + + +curdate.__doc__ = pysparkfuncs.curdate.__doc__ + + +def current_date() -> Column: + return _invoke_function("current_date") + + +current_date.__doc__ = pysparkfuncs.current_date.__doc__ + + +def current_timestamp() -> Column: + return _invoke_function("current_timestamp") + + +current_timestamp.__doc__ = pysparkfuncs.current_timestamp.__doc__ + + +def now() -> Column: + return _invoke_function("current_timestamp") + + +now.__doc__ = pysparkfuncs.now.__doc__ + + +def current_timezone() -> Column: + return _invoke_function("current_timezone") + + +current_timezone.__doc__ = pysparkfuncs.current_timezone.__doc__ + + +def localtimestamp() -> Column: + return _invoke_function("localtimestamp") + + +localtimestamp.__doc__ = pysparkfuncs.localtimestamp.__doc__ + + +def date_format(date: "ColumnOrName", format: str) -> Column: + return _invoke_function("date_format", _to_col(date), lit(format)) + + +date_format.__doc__ = pysparkfuncs.date_format.__doc__ + + +def year(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("year", col) + + +year.__doc__ = pysparkfuncs.year.__doc__ + + +def quarter(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("quarter", col) + + +quarter.__doc__ = pysparkfuncs.quarter.__doc__ + + +def month(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("month", col) + + +month.__doc__ = pysparkfuncs.month.__doc__ + + +def dayofweek(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("dayofweek", col) + + +dayofweek.__doc__ = pysparkfuncs.dayofweek.__doc__ + + +def dayofmonth(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("dayofmonth", col) + + +dayofmonth.__doc__ = pysparkfuncs.dayofmonth.__doc__ + + +def day(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("day", col) + + +day.__doc__ = pysparkfuncs.day.__doc__ + + +def dayofyear(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("dayofyear", col) + + +dayofyear.__doc__ = pysparkfuncs.dayofyear.__doc__ + + +def hour(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("hour", col) + + +hour.__doc__ = pysparkfuncs.hour.__doc__ + + +def minute(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("minute", col) + + +minute.__doc__ = pysparkfuncs.minute.__doc__ + + +def second(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("second", col) + + +second.__doc__ = pysparkfuncs.second.__doc__ + + +def weekofyear(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("weekofyear", col) + + +weekofyear.__doc__ = pysparkfuncs.weekofyear.__doc__ + + +def weekday(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("weekday", col) + + +weekday.__doc__ = pysparkfuncs.weekday.__doc__ + + +def extract(field: "ColumnOrName", source: "ColumnOrName") -> Column: + return _invoke_function_over_columns("extract", field, source) + + +extract.__doc__ = pysparkfuncs.extract.__doc__ + + +def date_part(field: "ColumnOrName", source: "ColumnOrName") -> Column: + return _invoke_function_over_columns("date_part", field, source) + + +extract.__doc__ = pysparkfuncs.extract.__doc__ + + +def datepart(field: "ColumnOrName", source: "ColumnOrName") -> Column: + return _invoke_function_over_columns("datepart", field, source) + + +extract.__doc__ = pysparkfuncs.extract.__doc__ + + +def make_date(year: "ColumnOrName", month: "ColumnOrName", day: "ColumnOrName") -> Column: + return _invoke_function_over_columns("make_date", year, month, day) + + +make_date.__doc__ = pysparkfuncs.make_date.__doc__ + + +def date_add(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("date_add", start, days) + + +date_add.__doc__ = pysparkfuncs.date_add.__doc__ + + +def dateadd(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("dateadd", start, days) + + +dateadd.__doc__ = pysparkfuncs.dateadd.__doc__ + + +def date_sub(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("date_sub", start, days) + + +date_sub.__doc__ = pysparkfuncs.date_sub.__doc__ + + +def datediff(end: "ColumnOrName", start: "ColumnOrName") -> Column: + return _invoke_function_over_columns("datediff", end, start) + + +datediff.__doc__ = pysparkfuncs.datediff.__doc__ + + +def date_diff(end: "ColumnOrName", start: "ColumnOrName") -> Column: + return _invoke_function_over_columns("date_diff", end, start) + + +date_diff.__doc__ = pysparkfuncs.date_diff.__doc__ + + +def date_from_unix_date(days: "ColumnOrName") -> Column: + return _invoke_function_over_columns("date_from_unix_date", days) + + +date_from_unix_date.__doc__ = pysparkfuncs.date_from_unix_date.__doc__ + + +def add_months(start: "ColumnOrName", months: Union["ColumnOrName", int]) -> Column: + months = lit(months) if isinstance(months, int) else months + return _invoke_function_over_columns("add_months", start, months) + + +add_months.__doc__ = pysparkfuncs.add_months.__doc__ + + +def months_between(date1: "ColumnOrName", date2: "ColumnOrName", roundOff: bool = True) -> Column: + return _invoke_function("months_between", _to_col(date1), _to_col(date2), lit(roundOff)) + + +months_between.__doc__ = pysparkfuncs.months_between.__doc__ + + +def to_date(col: "ColumnOrName", format: Optional[str] = None) -> Column: + if format is None: + return _invoke_function_over_columns("to_date", col) + else: + return _invoke_function("to_date", _to_col(col), lit(format)) + + +to_date.__doc__ = pysparkfuncs.to_date.__doc__ + + +def unix_date(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unix_date", col) + + +unix_date.__doc__ = pysparkfuncs.unix_date.__doc__ + + +def unix_micros(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unix_micros", col) + + +unix_micros.__doc__ = pysparkfuncs.unix_micros.__doc__ + + +def unix_millis(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unix_millis", col) + + +unix_millis.__doc__ = pysparkfuncs.unix_millis.__doc__ + + +def unix_seconds(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("unix_seconds", col) + + +unix_seconds.__doc__ = pysparkfuncs.unix_seconds.__doc__ + + +@overload +def to_timestamp(col: "ColumnOrName") -> Column: + ... + + +@overload +def to_timestamp(col: "ColumnOrName", format: str) -> Column: + ... + + +def to_timestamp(col: "ColumnOrName", format: Optional[str] = None) -> Column: + if format is None: + return _invoke_function_over_columns("to_timestamp", col) + else: + return _invoke_function("to_timestamp", _to_col(col), lit(format)) + + +to_timestamp.__doc__ = pysparkfuncs.to_timestamp.__doc__ + + +def try_to_timestamp(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + if format is not None: + return _invoke_function_over_columns("try_to_timestamp", col, format) + else: + return _invoke_function_over_columns("try_to_timestamp", col) + + +try_to_timestamp.__doc__ = pysparkfuncs.try_to_timestamp.__doc__ + + +def xpath(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath", xml, path) + + +xpath.__doc__ = pysparkfuncs.xpath.__doc__ + + +def xpath_boolean(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_boolean", xml, path) + + +xpath_boolean.__doc__ = pysparkfuncs.xpath_boolean.__doc__ + + +def xpath_double(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_double", xml, path) + + +xpath_double.__doc__ = pysparkfuncs.xpath_double.__doc__ + + +def xpath_number(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_number", xml, path) + + +xpath_number.__doc__ = pysparkfuncs.xpath_number.__doc__ + + +def xpath_float(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_float", xml, path) + + +xpath_float.__doc__ = pysparkfuncs.xpath_float.__doc__ + + +def xpath_int(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_int", xml, path) + + +xpath_int.__doc__ = pysparkfuncs.xpath_int.__doc__ + + +def xpath_long(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_long", xml, path) + + +xpath_long.__doc__ = pysparkfuncs.xpath_long.__doc__ + + +def xpath_short(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_short", xml, path) + + +xpath_short.__doc__ = pysparkfuncs.xpath_short.__doc__ + + +def xpath_string(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xpath_string", xml, path) + + +xpath_string.__doc__ = pysparkfuncs.xpath_string.__doc__ + + +def trunc(date: "ColumnOrName", format: str) -> Column: + return _invoke_function("trunc", _to_col(date), lit(format)) + + +trunc.__doc__ = pysparkfuncs.trunc.__doc__ + + +def date_trunc(format: str, timestamp: "ColumnOrName") -> Column: + return _invoke_function("date_trunc", lit(format), _to_col(timestamp)) + + +date_trunc.__doc__ = pysparkfuncs.date_trunc.__doc__ + + +def next_day(date: "ColumnOrName", dayOfWeek: str) -> Column: + return _invoke_function("next_day", _to_col(date), lit(dayOfWeek)) + + +next_day.__doc__ = pysparkfuncs.next_day.__doc__ + + +def last_day(date: "ColumnOrName") -> Column: + return _invoke_function_over_columns("last_day", date) + + +last_day.__doc__ = pysparkfuncs.last_day.__doc__ + + +def from_unixtime(timestamp: "ColumnOrName", format: str = "yyyy-MM-dd HH:mm:ss") -> Column: + return _invoke_function("from_unixtime", _to_col(timestamp), lit(format)) + + +from_unixtime.__doc__ = pysparkfuncs.from_unixtime.__doc__ + + +@overload +def unix_timestamp(timestamp: "ColumnOrName", format: str = ...) -> Column: + ... + + +@overload +def unix_timestamp() -> Column: + ... + + +def unix_timestamp( + timestamp: Optional["ColumnOrName"] = None, format: str = "yyyy-MM-dd HH:mm:ss" +) -> Column: + if timestamp is None: + return _invoke_function("unix_timestamp") + return _invoke_function("unix_timestamp", _to_col(timestamp), lit(format)) + + +unix_timestamp.__doc__ = pysparkfuncs.unix_timestamp.__doc__ + + +def from_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") -> Column: + if isinstance(tz, str): + tz = lit(tz) + return _invoke_function_over_columns("from_utc_timestamp", timestamp, tz) + + +from_utc_timestamp.__doc__ = pysparkfuncs.from_utc_timestamp.__doc__ + + +def to_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") -> Column: + if isinstance(tz, str): + tz = lit(tz) + return _invoke_function_over_columns("to_utc_timestamp", timestamp, tz) + + +to_utc_timestamp.__doc__ = pysparkfuncs.to_utc_timestamp.__doc__ + + +def timestamp_seconds(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("timestamp_seconds", col) + + +timestamp_seconds.__doc__ = pysparkfuncs.timestamp_seconds.__doc__ + + +def timestamp_millis(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("timestamp_millis", col) + + +timestamp_millis.__doc__ = pysparkfuncs.timestamp_millis.__doc__ + + +def timestamp_micros(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("timestamp_micros", col) + + +timestamp_micros.__doc__ = pysparkfuncs.timestamp_micros.__doc__ + + +def window( + timeColumn: "ColumnOrName", + windowDuration: str, + slideDuration: Optional[str] = None, + startTime: Optional[str] = None, +) -> Column: + if windowDuration is None or not isinstance(windowDuration, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "windowDuration", + "arg_type": type(windowDuration).__name__, + }, + ) + if slideDuration is not None and not isinstance(slideDuration, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "slideDuration", + "arg_type": type(slideDuration).__name__, + }, + ) + if startTime is not None and not isinstance(startTime, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "startTime", "arg_type": type(startTime).__name__}, + ) + + time_col = _to_col(timeColumn) + + if slideDuration is not None and startTime is not None: + return _invoke_function( + "window", time_col, lit(windowDuration), lit(slideDuration), lit(startTime) + ) + elif slideDuration is not None: + return _invoke_function("window", time_col, lit(windowDuration), lit(slideDuration)) + elif startTime is not None: + return _invoke_function( + "window", time_col, lit(windowDuration), lit(windowDuration), lit(startTime) + ) + else: + return _invoke_function("window", time_col, lit(windowDuration)) + + +window.__doc__ = pysparkfuncs.window.__doc__ + + +def window_time( + windowColumn: "ColumnOrName", +) -> Column: + return _invoke_function("window_time", _to_col(windowColumn)) + + +window_time.__doc__ = pysparkfuncs.window_time.__doc__ + + +def session_window(timeColumn: "ColumnOrName", gapDuration: Union[Column, str]) -> Column: + if gapDuration is None or not isinstance(gapDuration, (Column, str)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "gapDuration", "arg_type": type(gapDuration).__name__}, + ) + + time_col = _to_col(timeColumn) + + if isinstance(gapDuration, Column): + return _invoke_function("session_window", time_col, gapDuration) + else: + return _invoke_function("session_window", time_col, lit(gapDuration)) + + +session_window.__doc__ = pysparkfuncs.session_window.__doc__ + + +def to_unix_timestamp( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + if format is not None: + return _invoke_function_over_columns("to_unix_timestamp", timestamp, format) + else: + return _invoke_function_over_columns("to_unix_timestamp", timestamp) + + +to_unix_timestamp.__doc__ = pysparkfuncs.to_unix_timestamp.__doc__ + + +def to_timestamp_ltz( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + if format is not None: + return _invoke_function_over_columns("to_timestamp_ltz", timestamp, format) + else: + return _invoke_function_over_columns("to_timestamp_ltz", timestamp) + + +to_timestamp_ltz.__doc__ = pysparkfuncs.to_timestamp_ltz.__doc__ + + +def to_timestamp_ntz( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + if format is not None: + return _invoke_function_over_columns("to_timestamp_ntz", timestamp, format) + else: + return _invoke_function_over_columns("to_timestamp_ntz", timestamp) + + +to_timestamp_ntz.__doc__ = pysparkfuncs.to_timestamp_ntz.__doc__ + + +# Partition Transformation Functions + + +def bucket(numBuckets: Union[Column, int], col: "ColumnOrName") -> Column: + if isinstance(numBuckets, int): + _numBuckets = lit(numBuckets) + elif isinstance(numBuckets, Column): + _numBuckets = numBuckets + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT", + message_parameters={ + "arg_name": "numBuckets", + "arg_type": type(numBuckets).__name__, + }, + ) + + return _invoke_function("bucket", _numBuckets, _to_col(col)) + + +bucket.__doc__ = pysparkfuncs.bucket.__doc__ + + +def years(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("years", col) + + +years.__doc__ = pysparkfuncs.years.__doc__ + + +def months(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("months", col) + + +months.__doc__ = pysparkfuncs.months.__doc__ + + +def days(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("days", col) + + +days.__doc__ = pysparkfuncs.days.__doc__ + + +def hours(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("hours", col) + + +hours.__doc__ = pysparkfuncs.hours.__doc__ + + +def convert_timezone( + sourceTz: Optional[Column], targetTz: Column, sourceTs: "ColumnOrName" +) -> Column: + if sourceTz is None: + return _invoke_function_over_columns("convert_timezone", targetTz, sourceTs) + else: + return _invoke_function_over_columns("convert_timezone", sourceTz, targetTz, sourceTs) + + +convert_timezone.__doc__ = pysparkfuncs.convert_timezone.__doc__ + + +def make_dt_interval( + days: Optional["ColumnOrName"] = None, + hours: Optional["ColumnOrName"] = None, + mins: Optional["ColumnOrName"] = None, + secs: Optional["ColumnOrName"] = None, +) -> Column: + _days = lit(0) if days is None else _to_col(days) + _hours = lit(0) if hours is None else _to_col(hours) + _mins = lit(0) if mins is None else _to_col(mins) + _secs = lit(decimal.Decimal(0)) if secs is None else _to_col(secs) + + return _invoke_function_over_columns("make_dt_interval", _days, _hours, _mins, _secs) + + +make_dt_interval.__doc__ = pysparkfuncs.make_dt_interval.__doc__ + + +def make_interval( + years: Optional["ColumnOrName"] = None, + months: Optional["ColumnOrName"] = None, + weeks: Optional["ColumnOrName"] = None, + days: Optional["ColumnOrName"] = None, + hours: Optional["ColumnOrName"] = None, + mins: Optional["ColumnOrName"] = None, + secs: Optional["ColumnOrName"] = None, +) -> Column: + _years = lit(0) if years is None else _to_col(years) + _months = lit(0) if months is None else _to_col(months) + _weeks = lit(0) if weeks is None else _to_col(weeks) + _days = lit(0) if days is None else _to_col(days) + _hours = lit(0) if hours is None else _to_col(hours) + _mins = lit(0) if mins is None else _to_col(mins) + _secs = lit(decimal.Decimal(0)) if secs is None else _to_col(secs) + + return _invoke_function_over_columns( + "make_interval", _years, _months, _weeks, _days, _hours, _mins, _secs + ) + + +make_interval.__doc__ = pysparkfuncs.make_interval.__doc__ + + +def make_timestamp( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", + timezone: Optional["ColumnOrName"] = None, +) -> Column: + if timezone is not None: + return _invoke_function_over_columns( + "make_timestamp", years, months, days, hours, mins, secs, timezone + ) + else: + return _invoke_function_over_columns( + "make_timestamp", years, months, days, hours, mins, secs + ) + + +make_timestamp.__doc__ = pysparkfuncs.make_timestamp.__doc__ + + +def make_timestamp_ltz( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", + timezone: Optional["ColumnOrName"] = None, +) -> Column: + if timezone is not None: + return _invoke_function_over_columns( + "make_timestamp_ltz", years, months, days, hours, mins, secs, timezone + ) + else: + return _invoke_function_over_columns( + "make_timestamp_ltz", years, months, days, hours, mins, secs + ) + + +make_timestamp_ltz.__doc__ = pysparkfuncs.make_timestamp_ltz.__doc__ + + +def make_timestamp_ntz( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", +) -> Column: + return _invoke_function_over_columns( + "make_timestamp_ntz", years, months, days, hours, mins, secs + ) + + +make_timestamp_ntz.__doc__ = pysparkfuncs.make_timestamp_ntz.__doc__ + + +def make_ym_interval( + years: Optional["ColumnOrName"] = None, + months: Optional["ColumnOrName"] = None, +) -> Column: + _years = lit(0) if years is None else _to_col(years) + _months = lit(0) if months is None else _to_col(months) + return _invoke_function_over_columns("make_ym_interval", _years, _months) + + +make_ym_interval.__doc__ = pysparkfuncs.make_ym_interval.__doc__ + +# Misc Functions + + +def current_catalog() -> Column: + return _invoke_function("current_catalog") + + +current_catalog.__doc__ = pysparkfuncs.current_catalog.__doc__ + + +def current_database() -> Column: + return _invoke_function("current_database") + + +current_database.__doc__ = pysparkfuncs.current_database.__doc__ + + +def current_schema() -> Column: + return _invoke_function("current_schema") + + +current_schema.__doc__ = pysparkfuncs.current_schema.__doc__ + + +def current_user() -> Column: + return _invoke_function("current_user") + + +current_user.__doc__ = pysparkfuncs.current_user.__doc__ + + +def user() -> Column: + return _invoke_function("user") + + +user.__doc__ = pysparkfuncs.user.__doc__ + + +def assert_true(col: "ColumnOrName", errMsg: Optional[Union[Column, str]] = None) -> Column: + if errMsg is None: + return _invoke_function_over_columns("assert_true", col) + if not isinstance(errMsg, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": type(errMsg).__name__}, + ) + _err_msg = lit(errMsg) if isinstance(errMsg, str) else _to_col(errMsg) + return _invoke_function("assert_true", _to_col(col), _err_msg) + + +assert_true.__doc__ = pysparkfuncs.assert_true.__doc__ + + +def raise_error(errMsg: Union[Column, str]) -> Column: + if not isinstance(errMsg, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": type(errMsg).__name__}, + ) + _err_msg = lit(errMsg) if isinstance(errMsg, str) else _to_col(errMsg) + return _invoke_function("raise_error", _err_msg) + + +raise_error.__doc__ = pysparkfuncs.raise_error.__doc__ + + +def crc32(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("crc32", col) + + +crc32.__doc__ = pysparkfuncs.crc32.__doc__ + + +def hash(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("hash", *cols) + + +hash.__doc__ = pysparkfuncs.hash.__doc__ + + +def xxhash64(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xxhash64", *cols) + + +xxhash64.__doc__ = pysparkfuncs.xxhash64.__doc__ + + +def md5(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("md5", col) + + +md5.__doc__ = pysparkfuncs.md5.__doc__ + + +def sha1(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sha1", col) + + +sha1.__doc__ = pysparkfuncs.sha1.__doc__ + + +def sha2(col: "ColumnOrName", numBits: int) -> Column: + return _invoke_function("sha2", _to_col(col), lit(numBits)) + + +sha2.__doc__ = pysparkfuncs.sha2.__doc__ + + +def hll_sketch_agg(col: "ColumnOrName", lgConfigK: Optional[Union[int, Column]] = None) -> Column: + if lgConfigK is None: + return _invoke_function_over_columns("hll_sketch_agg", col) + else: + _lgConfigK = lit(lgConfigK) if isinstance(lgConfigK, int) else lgConfigK + return _invoke_function_over_columns("hll_sketch_agg", col, _lgConfigK) + + +hll_sketch_agg.__doc__ = pysparkfuncs.hll_sketch_agg.__doc__ + + +def hll_union_agg(col: "ColumnOrName", allowDifferentLgConfigK: Optional[bool] = None) -> Column: + if allowDifferentLgConfigK is None: + return _invoke_function_over_columns("hll_union_agg", col) + else: + _allowDifferentLgConfigK = ( + lit(allowDifferentLgConfigK) + if isinstance(allowDifferentLgConfigK, bool) + else allowDifferentLgConfigK + ) + return _invoke_function_over_columns("hll_union_agg", col, _allowDifferentLgConfigK) + + +hll_union_agg.__doc__ = pysparkfuncs.hll_union_agg.__doc__ + + +def hll_sketch_estimate(col: "ColumnOrName") -> Column: + return _invoke_function("hll_sketch_estimate", _to_col(col)) + + +hll_sketch_estimate.__doc__ = pysparkfuncs.hll_sketch_estimate.__doc__ + + +def hll_union( + col1: "ColumnOrName", col2: "ColumnOrName", allowDifferentLgConfigK: Optional[bool] = None +) -> Column: + if allowDifferentLgConfigK is not None: + return _invoke_function( + "hll_union", _to_col(col1), _to_col(col2), lit(allowDifferentLgConfigK) + ) + else: + return _invoke_function("hll_union", _to_col(col1), _to_col(col2)) + + +hll_union.__doc__ = pysparkfuncs.hll_union.__doc__ + + +# Predicates Function + + +def ifnull(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("ifnull", col1, col2) + + +ifnull.__doc__ = pysparkfuncs.ifnull.__doc__ + + +def isnotnull(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("isnotnull", col) + + +isnotnull.__doc__ = pysparkfuncs.isnotnull.__doc__ + + +def equal_null(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("equal_null", col1, col2) + + +equal_null.__doc__ = pysparkfuncs.equal_null.__doc__ + + +def nullif(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("nullif", col1, col2) + + +nullif.__doc__ = pysparkfuncs.nullif.__doc__ + + +def nvl(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + return _invoke_function_over_columns("nvl", col1, col2) + + +nvl.__doc__ = pysparkfuncs.nvl.__doc__ + + +def nvl2(col1: "ColumnOrName", col2: "ColumnOrName", col3: "ColumnOrName") -> Column: + return _invoke_function_over_columns("nvl2", col1, col2, col3) + + +nvl2.__doc__ = pysparkfuncs.nvl2.__doc__ + + +def aes_encrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + iv: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + _mode = lit("GCM") if mode is None else _to_col(mode) + _padding = lit("DEFAULT") if padding is None else _to_col(padding) + _iv = lit("") if iv is None else _to_col(iv) + _aad = lit("") if aad is None else _to_col(aad) + + return _invoke_function_over_columns("aes_encrypt", input, key, _mode, _padding, _iv, _aad) + + +aes_encrypt.__doc__ = pysparkfuncs.aes_encrypt.__doc__ + + +def aes_decrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + _mode = lit("GCM") if mode is None else _to_col(mode) + _padding = lit("DEFAULT") if padding is None else _to_col(padding) + _aad = lit("") if aad is None else _to_col(aad) + + return _invoke_function_over_columns("aes_decrypt", input, key, _mode, _padding, _aad) + + +aes_decrypt.__doc__ = pysparkfuncs.aes_decrypt.__doc__ + + +def try_aes_decrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + _mode = lit("GCM") if mode is None else _to_col(mode) + _padding = lit("DEFAULT") if padding is None else _to_col(padding) + _aad = lit("") if aad is None else _to_col(aad) + + return _invoke_function_over_columns("try_aes_decrypt", input, key, _mode, _padding, _aad) + + +try_aes_decrypt.__doc__ = pysparkfuncs.try_aes_decrypt.__doc__ + + +def sha(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("sha", col) + + +sha.__doc__ = pysparkfuncs.sha.__doc__ + + +def input_file_block_length() -> Column: + return _invoke_function_over_columns("input_file_block_length") + + +input_file_block_length.__doc__ = pysparkfuncs.input_file_block_length.__doc__ + + +def input_file_block_start() -> Column: + return _invoke_function_over_columns("input_file_block_start") + + +input_file_block_start.__doc__ = pysparkfuncs.input_file_block_start.__doc__ + + +def reflect(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("reflect", *cols) + + +reflect.__doc__ = pysparkfuncs.reflect.__doc__ + + +def java_method(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("java_method", *cols) + + +java_method.__doc__ = pysparkfuncs.java_method.__doc__ + + +def version() -> Column: + return _invoke_function_over_columns("version") + + +version.__doc__ = pysparkfuncs.version.__doc__ + + +def typeof(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("typeof", col) + + +typeof.__doc__ = pysparkfuncs.typeof.__doc__ + + +def stack(*cols: "ColumnOrName") -> Column: + return _invoke_function_over_columns("stack", *cols) + + +stack.__doc__ = pysparkfuncs.stack.__doc__ + + +def bitmap_bit_position(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_bit_position", col) + + +bitmap_bit_position.__doc__ = pysparkfuncs.bitmap_bit_position.__doc__ + + +def bitmap_bucket_number(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_bucket_number", col) + + +bitmap_bucket_number.__doc__ = pysparkfuncs.bitmap_bucket_number.__doc__ + + +def bitmap_construct_agg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_construct_agg", col) + + +bitmap_construct_agg.__doc__ = pysparkfuncs.bitmap_construct_agg.__doc__ + + +def bitmap_count(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_count", col) + + +bitmap_count.__doc__ = pysparkfuncs.bitmap_count.__doc__ + + +def bitmap_or_agg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_or_agg", col) + + +bitmap_or_agg.__doc__ = pysparkfuncs.bitmap_or_agg.__doc__ + + +# Call Functions + + +def call_udf(udfName: str, *cols: "ColumnOrName") -> Column: + return _invoke_function(udfName, *[_to_col(c) for c in cols]) + + +call_udf.__doc__ = pysparkfuncs.call_udf.__doc__ + + +def unwrap_udt(col: "ColumnOrName") -> Column: + return _invoke_function("unwrap_udt", _to_col(col)) + + +unwrap_udt.__doc__ = pysparkfuncs.unwrap_udt.__doc__ + + +def udf( + f: Optional[Union[Callable[..., Any], "DataTypeOrString"]] = None, + returnType: "DataTypeOrString" = StringType(), + useArrow: Optional[bool] = None, +) -> Union["UserDefinedFunctionLike", Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]]: + if f is None or isinstance(f, (str, DataType)): + # If DataType has been passed as a positional argument + # for decorator use it as a returnType + return_type = f or returnType + return functools.partial( + _create_py_udf, + returnType=return_type, + useArrow=useArrow, + ) + else: + return _create_py_udf(f=f, returnType=returnType, useArrow=useArrow) + + +udf.__doc__ = pysparkfuncs.udf.__doc__ + + +def udtf( + cls: Optional[Type] = None, + *, + returnType: Union[StructType, str], + useArrow: Optional[bool] = None, +) -> Union["UserDefinedTableFunction", Callable[[Type], "UserDefinedTableFunction"]]: + if cls is None: + return functools.partial(_create_py_udtf, returnType=returnType, useArrow=useArrow) + else: + return _create_py_udtf(cls=cls, returnType=returnType, useArrow=useArrow) + + +udtf.__doc__ = pysparkfuncs.udtf.__doc__ + + +def call_function(funcName: str, *cols: "ColumnOrName") -> Column: + expressions = [_to_col(c)._expr for c in cols] + return Column(CallFunction(funcName, expressions)) + + +call_function.__doc__ = pysparkfuncs.call_function.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.functions + + globs = pyspark.sql.connect.functions.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.functions tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.functions, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/group.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/group.py new file mode 100644 index 0000000..a393d2c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/group.py @@ -0,0 +1,415 @@ +# +# 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 warnings + +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import ( + Any, + Dict, + List, + Sequence, + Union, + TYPE_CHECKING, + Optional, + overload, + cast, +) + +from pyspark.rdd import PythonEvalType +from pyspark.sql.group import GroupedData as PySparkGroupedData +from pyspark.sql.pandas.group_ops import PandasCogroupedOps as PySparkPandasCogroupedOps +from pyspark.sql.types import NumericType +from pyspark.sql.types import StructType + +import pyspark.sql.connect.plan as plan +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.functions import _invoke_function, col, lit +from pyspark.errors import PySparkNotImplementedError, PySparkTypeError + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ( + LiteralType, + PandasGroupedMapFunction, + GroupedMapPandasUserDefinedFunction, + PandasCogroupedMapFunction, + PandasGroupedMapFunctionWithState, + ) + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.types import StructType + + +class GroupedData: + def __init__( + self, + df: "DataFrame", + group_type: str, + grouping_cols: Sequence["Column"], + pivot_col: Optional["Column"] = None, + pivot_values: Optional[Sequence["LiteralType"]] = None, + ) -> None: + from pyspark.sql.connect.dataframe import DataFrame + + assert isinstance(df, DataFrame) + self._df = df + + assert isinstance(group_type, str) and group_type in ["groupby", "rollup", "cube", "pivot"] + self._group_type = group_type + + assert isinstance(grouping_cols, list) and all(isinstance(g, Column) for g in grouping_cols) + self._grouping_cols: List[Column] = grouping_cols + + self._pivot_col: Optional["Column"] = None + self._pivot_values: Optional[List[Any]] = None + if group_type == "pivot": + assert pivot_col is not None and isinstance(pivot_col, Column) + assert pivot_values is None or isinstance(pivot_values, list) + self._pivot_col = pivot_col + self._pivot_values = pivot_values + + def __repr__(self) -> str: + # the expressions are not resolved here, + # so the string representation can be different from vanilla PySpark. + grouping_str = ", ".join(str(e._expr) for e in self._grouping_cols) + grouping_str = f"grouping expressions: [{grouping_str}]" + + value_str = ", ".join("%s: %s" % c for c in self._df.dtypes) + + if self._group_type == "groupby": + type_str = "GroupBy" + elif self._group_type == "rollup": + type_str = "RollUp" + elif self._group_type == "cube": + type_str = "Cube" + else: + type_str = "Pivot" + + return f"GroupedData[{grouping_str}, value: [{value_str}], type: {type_str}]" + + @overload + def agg(self, *exprs: Column) -> "DataFrame": + ... + + @overload + def agg(self, __exprs: Dict[str, str]) -> "DataFrame": + ... + + def agg(self, *exprs: Union[Column, Dict[str, str]]) -> "DataFrame": + from pyspark.sql.connect.dataframe import DataFrame + + assert exprs, "exprs should not be empty" + if len(exprs) == 1 and isinstance(exprs[0], dict): + # Convert the dict into key value pairs + aggregate_cols = [_invoke_function(exprs[0][k], col(k)) for k in exprs[0]] + else: + # Columns + assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column" + aggregate_cols = cast(List[Column], list(exprs)) + + return DataFrame.withPlan( + plan.Aggregate( + child=self._df._plan, + group_type=self._group_type, + grouping_cols=self._grouping_cols, + aggregate_cols=aggregate_cols, + pivot_col=self._pivot_col, + pivot_values=self._pivot_values, + ), + session=self._df._session, + ) + + agg.__doc__ = PySparkGroupedData.agg.__doc__ + + def _numeric_agg(self, function: str, cols: Sequence[str]) -> "DataFrame": + from pyspark.sql.connect.dataframe import DataFrame + + assert isinstance(function, str) and function in ["min", "max", "avg", "sum"] + + assert isinstance(cols, list) and all(isinstance(c, str) for c in cols) + + schema = self._df.schema + + numerical_cols: List[str] = [ + field.name for field in schema.fields if isinstance(field.dataType, NumericType) + ] + + agg_cols: List[str] = [] + + if len(cols) > 0: + invalid_cols = [c for c in cols if c not in numerical_cols] + if len(invalid_cols) > 0: + raise PySparkTypeError( + error_class="NOT_NUMERIC_COLUMNS", + message_parameters={"invalid_columns": str(invalid_cols)}, + ) + agg_cols = cols + else: + # if no column is provided, then all numerical columns are selected + agg_cols = numerical_cols + + return DataFrame.withPlan( + plan.Aggregate( + child=self._df._plan, + group_type=self._group_type, + grouping_cols=self._grouping_cols, + aggregate_cols=[_invoke_function(function, col(c)) for c in agg_cols], + pivot_col=self._pivot_col, + pivot_values=self._pivot_values, + ), + session=self._df._session, + ) + + def min(self, *cols: str) -> "DataFrame": + return self._numeric_agg("min", list(cols)) + + min.__doc__ = PySparkGroupedData.min.__doc__ + + def max(self, *cols: str) -> "DataFrame": + return self._numeric_agg("max", list(cols)) + + max.__doc__ = PySparkGroupedData.max.__doc__ + + def sum(self, *cols: str) -> "DataFrame": + return self._numeric_agg("sum", list(cols)) + + sum.__doc__ = PySparkGroupedData.sum.__doc__ + + def avg(self, *cols: str) -> "DataFrame": + return self._numeric_agg("avg", list(cols)) + + avg.__doc__ = PySparkGroupedData.avg.__doc__ + + mean = avg + + def count(self) -> "DataFrame": + return self.agg(_invoke_function("count", lit(1)).alias("count")) + + count.__doc__ = PySparkGroupedData.count.__doc__ + + def pivot(self, pivot_col: str, values: Optional[List["LiteralType"]] = None) -> "GroupedData": + if self._group_type != "groupby": + if self._group_type == "pivot": + raise PySparkNotImplementedError( + error_class="UNSUPPORTED_OPERATION", + message_parameters={"operation": "Repeated PIVOT operation"}, + ) + else: + raise PySparkNotImplementedError( + error_class="UNSUPPORTED_OPERATION", + message_parameters={"operation": f"PIVOT after {self._group_type.upper()}"}, + ) + + if not isinstance(pivot_col, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "pivot_col", "arg_type": type(pivot_col).__name__}, + ) + + if values is not None: + if not isinstance(values, list): + raise PySparkTypeError( + error_class="NOT_LIST", + message_parameters={"arg_name": "values", "arg_type": type(values).__name__}, + ) + for v in values: + if not isinstance(v, (bool, float, int, str)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT_OR_STR", + message_parameters={"arg_name": "value", "arg_type": type(v).__name__}, + ) + + return GroupedData( + df=self._df, + group_type="pivot", + grouping_cols=self._grouping_cols, + pivot_col=self._df[pivot_col], + pivot_values=values, + ) + + pivot.__doc__ = PySparkGroupedData.pivot.__doc__ + + def apply(self, udf: "GroupedMapPandasUserDefinedFunction") -> "DataFrame": + # Columns are special because hasattr always return True + if ( + isinstance(udf, Column) + or not hasattr(udf, "func") + or ( + udf.evalType # type: ignore[attr-defined] + != PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF + ) + ): + raise PySparkTypeError( + error_class="INVALID_UDF_EVAL_TYPE", + message_parameters={"eval_type": "SQL_GROUPED_MAP_PANDAS_UDF"}, + ) + + warnings.warn( + "It is preferred to use 'applyInPandas' over this " + "API. This API will be deprecated in the future releases. See SPARK-28264 for " + "more details.", + UserWarning, + ) + + return self.applyInPandas(udf.func, schema=udf.returnType) # type: ignore[attr-defined] + + apply.__doc__ = PySparkGroupedData.apply.__doc__ + + def applyInPandas( + self, func: "PandasGroupedMapFunction", schema: Union["StructType", str] + ) -> "DataFrame": + from pyspark.sql.connect.udf import UserDefinedFunction + from pyspark.sql.connect.dataframe import DataFrame + + udf_obj = UserDefinedFunction( + func, + returnType=schema, + evalType=PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, + ) + + return DataFrame.withPlan( + plan.GroupMap( + child=self._df._plan, + grouping_cols=self._grouping_cols, + function=udf_obj, + cols=self._df.columns, + ), + session=self._df._session, + ) + + applyInPandas.__doc__ = PySparkGroupedData.applyInPandas.__doc__ + + def applyInPandasWithState( + self, + func: "PandasGroupedMapFunctionWithState", + outputStructType: Union[StructType, str], + stateStructType: Union[StructType, str], + outputMode: str, + timeoutConf: str, + ) -> "DataFrame": + from pyspark.sql.connect.udf import UserDefinedFunction + from pyspark.sql.connect.dataframe import DataFrame + + udf_obj = UserDefinedFunction( + func, + returnType=outputStructType, + evalType=PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + ) + + output_schema: str = ( + outputStructType.json() + if isinstance(outputStructType, StructType) + else outputStructType + ) + + state_schema: str = ( + stateStructType.json() if isinstance(stateStructType, StructType) else stateStructType + ) + + return DataFrame.withPlan( + plan.ApplyInPandasWithState( + child=self._df._plan, + grouping_cols=self._grouping_cols, + function=udf_obj, + output_schema=output_schema, + state_schema=state_schema, + output_mode=outputMode, + timeout_conf=timeoutConf, + cols=self._df.columns, + ), + session=self._df._session, + ) + + applyInPandasWithState.__doc__ = PySparkGroupedData.applyInPandasWithState.__doc__ + + def cogroup(self, other: "GroupedData") -> "PandasCogroupedOps": + return PandasCogroupedOps(self, other) + + cogroup.__doc__ = PySparkGroupedData.cogroup.__doc__ + + +GroupedData.__doc__ = PySparkGroupedData.__doc__ + + +class PandasCogroupedOps: + def __init__(self, gd1: "GroupedData", gd2: "GroupedData"): + self._gd1 = gd1 + self._gd2 = gd2 + + def applyInPandas( + self, func: "PandasCogroupedMapFunction", schema: Union["StructType", str] + ) -> "DataFrame": + from pyspark.sql.connect.udf import UserDefinedFunction + from pyspark.sql.connect.dataframe import DataFrame + + udf_obj = UserDefinedFunction( + func, + returnType=schema, + evalType=PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, + ) + + all_cols = self._extract_cols(self._gd1) + self._extract_cols(self._gd2) + return DataFrame.withPlan( + plan.CoGroupMap( + input=self._gd1._df._plan, + input_grouping_cols=self._gd1._grouping_cols, + other=self._gd2._df._plan, + other_grouping_cols=self._gd2._grouping_cols, + function=udf_obj, + cols=all_cols, + ), + session=self._gd1._df._session, + ) + + applyInPandas.__doc__ = PySparkPandasCogroupedOps.applyInPandas.__doc__ + + @staticmethod + def _extract_cols(gd: "GroupedData") -> List[Column]: + df = gd._df + return [df[col] for col in df.columns] + + +PandasCogroupedOps.__doc__ = PySparkPandasCogroupedOps.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.group + + globs = pyspark.sql.connect.group.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.group tests").remote("local[4]").getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.group, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/plan.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/plan.py new file mode 100644 index 0000000..84fd013 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/plan.py @@ -0,0 +1,2289 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import Any, List, Optional, Type, Sequence, Union, cast, TYPE_CHECKING, Mapping, Dict +import functools +import json +import pickle +from threading import Lock +from inspect import signature, isclass + +import pyarrow as pa + +from pyspark.serializers import CloudPickleSerializer +from pyspark.storagelevel import StorageLevel +from pyspark.sql.types import DataType + +import pyspark.sql.connect.proto as proto +from pyspark.sql.connect.conversion import storage_level_to_proto +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.expressions import ( + Expression, + SortOrder, + ColumnReference, + LiteralExpression, +) +from pyspark.sql.connect.types import pyspark_types_to_proto_types, UnparsedDataType +from pyspark.errors import ( + PySparkTypeError, + PySparkNotImplementedError, + PySparkRuntimeError, + IllegalArgumentException, +) + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ColumnOrName + from pyspark.sql.connect.client import SparkConnectClient + from pyspark.sql.connect.udf import UserDefinedFunction + + +class LogicalPlan: + + _lock: Lock = Lock() + _nextPlanId: int = 0 + + INDENT = 2 + + def __init__(self, child: Optional["LogicalPlan"]) -> None: + self._child = child + + plan_id: Optional[int] = None + with LogicalPlan._lock: + plan_id = LogicalPlan._nextPlanId + LogicalPlan._nextPlanId += 1 + + assert plan_id is not None + self._plan_id = plan_id + + def _create_proto_relation(self) -> proto.Relation: + plan = proto.Relation() + plan.common.plan_id = self._plan_id + return plan + + def unresolved_attr(self, colName: str) -> proto.Expression: + """Creates an unresolved attribute from a column name.""" + exp = proto.Expression() + exp.unresolved_attribute.unparsed_identifier = colName + return exp + + def to_attr_or_expression( + self, col: "ColumnOrName", session: "SparkConnectClient" + ) -> proto.Expression: + """Returns either an instance of an unresolved attribute or the serialized + expression value of the column.""" + if type(col) is str: + return self.unresolved_attr(col) + else: + return cast(Column, col).to_plan(session) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + ... + + def command(self, session: "SparkConnectClient") -> proto.Command: + ... + + def _verify(self, session: "SparkConnectClient") -> bool: + """This method is used to verify that the current logical plan + can be serialized to Proto and back and afterwards is identical.""" + plan = proto.Plan() + plan.root.CopyFrom(self.plan(session)) + + serialized_plan = plan.SerializeToString() + test_plan = proto.Plan() + test_plan.ParseFromString(serialized_plan) + + return test_plan == plan + + def to_proto(self, session: "SparkConnectClient", debug: bool = False) -> proto.Plan: + """ + Generates connect proto plan based on this LogicalPlan. + + Parameters + ---------- + session : :class:`SparkConnectClient`, optional. + a session that connects remote spark cluster. + debug: bool + if enabled, the proto plan will be printed. + """ + plan = proto.Plan() + plan.root.CopyFrom(self.plan(session)) + + if debug: + print(plan) + + return plan + + def _parameters_to_print(self, parameters: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Extracts the parameters that are able to be printed. It looks up the signature + in the constructor of this :class:`LogicalPlan`, and retrieves the variables + from this instance by the same name (or the name with prefix `_`) defined + in the constructor. + + Parameters + ---------- + parameters : map + Parameter mapping from ``inspect.signature(...).parameters`` + + Returns + ------- + dict + A dictionary consisting of a string name and variable found in this + :class:`LogicalPlan`. + + Notes + ----- + :class:`LogicalPlan` itself is filtered out and considered as a non-printable + parameter. + + Examples + -------- + The example below returns a dictionary from `self._start`, `self._end`, + `self._num_partitions`. + + >>> rg = Range(0, 10, 1) + >>> rg._parameters_to_print(signature(rg.__class__.__init__).parameters) + {'start': 0, 'end': 10, 'step': 1, 'num_partitions': None} + + If the child is defined, it is not considered as a printable instance + + >>> project = Project(rg, "value") + >>> project._parameters_to_print(signature(project.__class__.__init__).parameters) + {'columns': ['value']} + """ + params = {} + for name, tpe in parameters.items(): + # LogicalPlan is not to print, e.g., LogicalPlan + is_logical_plan = isclass(tpe.annotation) and isinstance(tpe.annotation, LogicalPlan) + # Look up the string argument defined as a forward reference e.g., "LogicalPlan" + is_forwardref_logical_plan = getattr(tpe.annotation, "__forward_arg__", "").endswith( + "LogicalPlan" + ) + # Wrapped LogicalPlan, e.g., Optional[LogicalPlan] + is_nested_logical_plan = any( + isclass(a) and issubclass(a, LogicalPlan) + for a in getattr(tpe.annotation, "__args__", ()) + ) + # Wrapped forward reference of LogicalPlan, e.g., Optional["LogicalPlan"]. + is_nested_forwardref_logical_plan = any( + getattr(a, "__forward_arg__", "").endswith("LogicalPlan") + for a in getattr(tpe.annotation, "__args__", ()) + ) + if ( + not is_logical_plan + and not is_forwardref_logical_plan + and not is_nested_logical_plan + and not is_nested_forwardref_logical_plan + ): + # Searches self.name or self._name + try: + params[name] = getattr(self, name) + except AttributeError: + try: + params[name] = getattr(self, "_" + name) + except AttributeError: + pass # Simpy ignore + return params + + def print(self, indent: int = 0) -> str: + """ + Print the simple string representation of the current :class:`LogicalPlan`. + + Parameters + ---------- + indent : int + The number of leading spaces for the output string. + + Returns + ------- + str + Simple string representation of this :class:`LogicalPlan`. + """ + params = self._parameters_to_print(signature(self.__class__.__init__).parameters) + pretty_params = [f"{name}='{param}'" for name, param in params.items()] + if len(pretty_params) == 0: + pretty_str = "" + else: + pretty_str = " " + ", ".join(pretty_params) + return f"{' ' * indent}<{self.__class__.__name__}{pretty_str}>\n{self._child_print(indent)}" + + def _repr_html_(self) -> str: + """Returns a :class:`LogicalPlan` with HTML code. This is generally called in third-party + systems such as Jupyter. + + Returns + ------- + str + HTML representation of this :class:`LogicalPlan`. + """ + params = self._parameters_to_print(signature(self.__class__.__init__).parameters) + pretty_params = [ + f"\n {name}: " f"{param}
    " for name, param in params.items() + ] + if len(pretty_params) == 0: + pretty_str = "" + else: + pretty_str = "".join(pretty_params) + return f""" +
      +
    • + {self.__class__.__name__}
      {pretty_str} + {self._child_repr()} +
    • +
    + """ + + def _child_print(self, indent: int) -> str: + return self._child.print(indent + LogicalPlan.INDENT) if self._child else "" + + def _child_repr(self) -> str: + return self._child._repr_html_() if self._child is not None else "" + + +class DataSource(LogicalPlan): + """A datasource with a format and optional a schema from which Spark reads data""" + + def __init__( + self, + format: Optional[str] = None, + schema: Optional[str] = None, + options: Optional[Mapping[str, str]] = None, + paths: Optional[List[str]] = None, + predicates: Optional[List[str]] = None, + is_streaming: Optional[bool] = None, + ) -> None: + super().__init__(None) + + assert format is None or isinstance(format, str) + assert schema is None or isinstance(schema, str) + + if options is not None: + for k, v in options.items(): + assert isinstance(k, str) + assert isinstance(v, str) + + if paths is not None: + assert isinstance(paths, list) + assert all(isinstance(path, str) for path in paths) + + if predicates is not None: + assert isinstance(predicates, list) + assert all(isinstance(predicate, str) for predicate in predicates) + + self._format = format + self._schema = schema + self._options = options + self._paths = paths + self._predicates = predicates + self._is_streaming = is_streaming + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + if self._format is not None: + plan.read.data_source.format = self._format + if self._schema is not None: + plan.read.data_source.schema = self._schema + if self._options is not None and len(self._options) > 0: + for k, v in self._options.items(): + plan.read.data_source.options[k] = v + if self._paths is not None and len(self._paths) > 0: + plan.read.data_source.paths.extend(self._paths) + if self._predicates is not None and len(self._predicates) > 0: + plan.read.data_source.predicates.extend(self._predicates) + if self._is_streaming is not None: + plan.read.is_streaming = self._is_streaming + return plan + + +class Read(LogicalPlan): + def __init__( + self, + table_name: str, + options: Optional[Dict[str, str]] = None, + is_streaming: Optional[bool] = None, + ) -> None: + super().__init__(None) + self.table_name = table_name + self.options = options or {} + self._is_streaming = is_streaming + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.read.named_table.unparsed_identifier = self.table_name + if self._is_streaming is not None: + plan.read.is_streaming = self._is_streaming + for k, v in self.options.items(): + plan.read.named_table.options[k] = v + return plan + + def print(self, indent: int = 0) -> str: + return f"{' ' * indent}\n" + + +class LocalRelation(LogicalPlan): + """Creates a LocalRelation plan object based on a PyArrow Table.""" + + def __init__( + self, + table: Optional["pa.Table"], + schema: Optional[str] = None, + ) -> None: + super().__init__(None) + + if table is None: + assert schema is not None + else: + assert isinstance(table, pa.Table) + + assert schema is None or isinstance(schema, str) + + self._table = table + + self._schema = schema + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + if self._table is not None: + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, self._table.schema) as writer: + for b in self._table.to_batches(): + writer.write_batch(b) + plan.local_relation.data = sink.getvalue().to_pybytes() + + if self._schema is not None: + plan.local_relation.schema = self._schema + return plan + + def serialize(self, session: "SparkConnectClient") -> bytes: + p = self.plan(session) + return bytes(p.local_relation.SerializeToString()) + + def print(self, indent: int = 0) -> str: + return f"{' ' * indent}\n" + + def _repr_html_(self) -> str: + return """ +
      +
    • LocalRelation
    • +
    + """ + + +class CachedLocalRelation(LogicalPlan): + """Creates a CachedLocalRelation plan object based on a hash of a LocalRelation.""" + + def __init__(self, hash: str) -> None: + super().__init__(None) + + self._hash = hash + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + clr = plan.cached_local_relation + + if session._user_id: + clr.userId = session._user_id + clr.sessionId = session._session_id + clr.hash = self._hash + + return plan + + def print(self, indent: int = 0) -> str: + return f"{' ' * indent}\n" + + def _repr_html_(self) -> str: + return """ +
      +
    • CachedLocalRelation
    • +
    + """ + + +class ShowString(LogicalPlan): + def __init__( + self, child: Optional["LogicalPlan"], num_rows: int, truncate: int, vertical: bool + ) -> None: + super().__init__(child) + self.num_rows = num_rows + self.truncate = truncate + self.vertical = vertical + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.show_string.input.CopyFrom(self._child.plan(session)) + plan.show_string.num_rows = self.num_rows + plan.show_string.truncate = self.truncate + plan.show_string.vertical = self.vertical + return plan + + +class HtmlString(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], num_rows: int, truncate: int) -> None: + super().__init__(child) + self.num_rows = num_rows + self.truncate = truncate + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.html_string.input.CopyFrom(self._child.plan(session)) + plan.html_string.num_rows = self.num_rows + plan.html_string.truncate = self.truncate + return plan + + +class Project(LogicalPlan): + """Logical plan object for a projection. + + All input arguments are directly serialized into the corresponding protocol buffer + objects. This class only provides very limited error handling and input validation. + + To be compatible with PySpark, we validate that the input arguments are all + expressions to be able to serialize them to the server. + + """ + + def __init__(self, child: Optional["LogicalPlan"], *columns: "ColumnOrName") -> None: + super().__init__(child) + self._columns = list(columns) + self.alias: Optional[str] = None + self._verify_expressions() + + def _verify_expressions(self) -> None: + """Ensures that all input arguments are instances of Expression or String.""" + for c in self._columns: + if not isinstance(c, (Column, str)): + raise PySparkTypeError( + error_class="NOT_LIST_OF_COLUMN_OR_STR", + message_parameters={"arg_name": "columns"}, + ) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + from pyspark.sql.connect.functions import col + + assert self._child is not None + plan = self._create_proto_relation() + plan.project.input.CopyFrom(self._child.plan(session)) + + proj_exprs = [] + for c in self._columns: + if isinstance(c, Column): + proj_exprs.append(c.to_plan(session)) + else: + proj_exprs.append(col(c).to_plan(session)) + + plan.project.expressions.extend(proj_exprs) + return plan + + +class WithColumns(LogicalPlan): + """Logical plan object for a withColumns operation.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + columnNames: Sequence[str], + columns: Sequence[Column], + metadata: Optional[Sequence[str]] = None, + ) -> None: + super().__init__(child) + + assert isinstance(columnNames, list) + assert len(columnNames) > 0 + assert all(isinstance(c, str) for c in columnNames) + + assert isinstance(columns, list) + assert len(columns) == len(columnNames) + assert all(isinstance(c, Column) for c in columns) + + if metadata is not None: + assert isinstance(metadata, list) + assert len(metadata) == len(columnNames) + for m in metadata: + assert isinstance(m, str) + # validate json string + assert m == "" or json.loads(m) is not None + + self._columnNames = columnNames + self._columns = columns + self._metadata = metadata + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.with_columns.input.CopyFrom(self._child.plan(session)) + + for i in range(0, len(self._columnNames)): + alias = proto.Expression.Alias() + alias.expr.CopyFrom(self._columns[i].to_plan(session)) + alias.name.append(self._columnNames[i]) + if self._metadata is not None: + alias.metadata = self._metadata[i] + plan.with_columns.aliases.append(alias) + + return plan + + +class WithWatermark(LogicalPlan): + """Logical plan object for a WithWatermark operation.""" + + def __init__(self, child: Optional["LogicalPlan"], event_time: str, delay_threshold: str): + super().__init__(child) + self._event_time = event_time + self._delay_threshold = delay_threshold + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.with_watermark.input.CopyFrom(self._child.plan(session)) + plan.with_watermark.event_time = self._event_time + plan.with_watermark.delay_threshold = self._delay_threshold + return plan + + +class CachedRemoteRelation(LogicalPlan): + """Logical plan object for a DataFrame reference which represents a DataFrame that's been + cached on the server with a given id.""" + + def __init__(self, relationId: str): + super().__init__(None) + self._relationId = relationId + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.cached_remote_relation.relation_id = self._relationId + return plan + + +class Hint(LogicalPlan): + """Logical plan object for a Hint operation.""" + + def __init__(self, child: Optional["LogicalPlan"], name: str, parameters: List[Any]) -> None: + super().__init__(child) + + assert isinstance(name, str) + + self._name = name + + for param in parameters: + assert isinstance(param, (list, str, float, int)) + if isinstance(param, list): + assert all(isinstance(p, (str, float, int)) for p in param) + + self._parameters = parameters + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + from pyspark.sql.connect.functions import array, lit + + assert self._child is not None + plan = self._create_proto_relation() + plan.hint.input.CopyFrom(self._child.plan(session)) + plan.hint.name = self._name + for param in self._parameters: + if isinstance(param, list): + plan.hint.parameters.append(array(*[lit(p) for p in param]).to_plan(session)) + else: + plan.hint.parameters.append(lit(param).to_plan(session)) + return plan + + +class Filter(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], filter: Column) -> None: + super().__init__(child) + self.filter = filter + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.filter.input.CopyFrom(self._child.plan(session)) + plan.filter.condition.CopyFrom(self.filter.to_plan(session)) + return plan + + +class Limit(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], limit: int) -> None: + super().__init__(child) + self.limit = limit + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.limit.input.CopyFrom(self._child.plan(session)) + plan.limit.limit = self.limit + return plan + + +class Tail(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], limit: int) -> None: + super().__init__(child) + self.limit = limit + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.tail.input.CopyFrom(self._child.plan(session)) + plan.tail.limit = self.limit + return plan + + +class Offset(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], offset: int = 0) -> None: + super().__init__(child) + self.offset = offset + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.offset.input.CopyFrom(self._child.plan(session)) + plan.offset.offset = self.offset + return plan + + +class Deduplicate(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + all_columns_as_keys: bool = False, + column_names: Optional[List[str]] = None, + within_watermark: bool = False, + ) -> None: + super().__init__(child) + self.all_columns_as_keys = all_columns_as_keys + self.column_names = column_names + self.within_watermark = within_watermark + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.deduplicate.input.CopyFrom(self._child.plan(session)) + plan.deduplicate.all_columns_as_keys = self.all_columns_as_keys + plan.deduplicate.within_watermark = self.within_watermark + if self.column_names is not None: + plan.deduplicate.column_names.extend(self.column_names) + return plan + + +class Sort(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + columns: List[Column], + is_global: bool, + ) -> None: + super().__init__(child) + + assert all(isinstance(c, Column) for c in columns) + assert isinstance(is_global, bool) + + self.columns = columns + self.is_global = is_global + + def _convert_col( + self, col: Column, session: "SparkConnectClient" + ) -> proto.Expression.SortOrder: + if isinstance(col._expr, SortOrder): + return col._expr.to_plan(session).sort_order + else: + return SortOrder(col._expr).to_plan(session).sort_order + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.sort.input.CopyFrom(self._child.plan(session)) + plan.sort.order.extend([self._convert_col(c, session) for c in self.columns]) + plan.sort.is_global = self.is_global + return plan + + +class Drop(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + columns: List[Union[Column, str]], + ) -> None: + super().__init__(child) + if len(columns) > 0: + assert all(isinstance(c, (Column, str)) for c in columns) + self._columns = columns + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.drop.input.CopyFrom(self._child.plan(session)) + for c in self._columns: + if isinstance(c, Column): + plan.drop.columns.append(c.to_plan(session)) + else: + plan.drop.column_names.append(c) + return plan + + +class Sample(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + lower_bound: float, + upper_bound: float, + with_replacement: bool, + seed: Optional[int], + deterministic_order: bool = False, + ) -> None: + super().__init__(child) + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.with_replacement = with_replacement + self.seed = seed + self.deterministic_order = deterministic_order + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.sample.input.CopyFrom(self._child.plan(session)) + plan.sample.lower_bound = self.lower_bound + plan.sample.upper_bound = self.upper_bound + plan.sample.with_replacement = self.with_replacement + if self.seed is not None: + plan.sample.seed = self.seed + plan.sample.deterministic_order = self.deterministic_order + return plan + + +class Aggregate(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + group_type: str, + grouping_cols: Sequence[Column], + aggregate_cols: Sequence[Column], + pivot_col: Optional[Column], + pivot_values: Optional[Sequence[Any]], + ) -> None: + super().__init__(child) + + assert isinstance(group_type, str) and group_type in ["groupby", "rollup", "cube", "pivot"] + self._group_type = group_type + + assert isinstance(grouping_cols, list) and all(isinstance(c, Column) for c in grouping_cols) + self._grouping_cols = grouping_cols + + assert isinstance(aggregate_cols, list) and all( + isinstance(c, Column) for c in aggregate_cols + ) + self._aggregate_cols = aggregate_cols + + if group_type == "pivot": + assert pivot_col is not None and isinstance(pivot_col, Column) + assert pivot_values is None or isinstance(pivot_values, list) + else: + assert pivot_col is None + assert pivot_values is None + + self._pivot_col = pivot_col + self._pivot_values = pivot_values + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + from pyspark.sql.connect.functions import lit + + assert self._child is not None + plan = self._create_proto_relation() + plan.aggregate.input.CopyFrom(self._child.plan(session)) + plan.aggregate.grouping_expressions.extend( + [c.to_plan(session) for c in self._grouping_cols] + ) + plan.aggregate.aggregate_expressions.extend( + [c.to_plan(session) for c in self._aggregate_cols] + ) + + if self._group_type == "groupby": + plan.aggregate.group_type = proto.Aggregate.GroupType.GROUP_TYPE_GROUPBY + elif self._group_type == "rollup": + plan.aggregate.group_type = proto.Aggregate.GroupType.GROUP_TYPE_ROLLUP + elif self._group_type == "cube": + plan.aggregate.group_type = proto.Aggregate.GroupType.GROUP_TYPE_CUBE + elif self._group_type == "pivot": + plan.aggregate.group_type = proto.Aggregate.GroupType.GROUP_TYPE_PIVOT + assert self._pivot_col is not None + plan.aggregate.pivot.col.CopyFrom(self._pivot_col.to_plan(session)) + if self._pivot_values is not None and len(self._pivot_values) > 0: + plan.aggregate.pivot.values.extend( + [lit(v).to_plan(session).literal for v in self._pivot_values] + ) + + return plan + + +class Join(LogicalPlan): + def __init__( + self, + left: Optional["LogicalPlan"], + right: "LogicalPlan", + on: Optional[Union[str, List[str], Column, List[Column]]], + how: Optional[str], + ) -> None: + super().__init__(left) + self.left = cast(LogicalPlan, left) + self.right = right + self.on = on + if how is None: + join_type = proto.Join.JoinType.JOIN_TYPE_INNER + elif how == "inner": + join_type = proto.Join.JoinType.JOIN_TYPE_INNER + elif how in ["outer", "full", "fullouter"]: + join_type = proto.Join.JoinType.JOIN_TYPE_FULL_OUTER + elif how in ["leftouter", "left"]: + join_type = proto.Join.JoinType.JOIN_TYPE_LEFT_OUTER + elif how in ["rightouter", "right"]: + join_type = proto.Join.JoinType.JOIN_TYPE_RIGHT_OUTER + elif how in ["leftsemi", "semi"]: + join_type = proto.Join.JoinType.JOIN_TYPE_LEFT_SEMI + elif how in ["leftanti", "anti"]: + join_type = proto.Join.JoinType.JOIN_TYPE_LEFT_ANTI + elif how == "cross": + join_type = proto.Join.JoinType.JOIN_TYPE_CROSS + else: + raise IllegalArgumentException( + error_class="UNSUPPORTED_JOIN_TYPE", + message_parameters={"join_type": how}, + ) + self.how = join_type + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.join.left.CopyFrom(self.left.plan(session)) + plan.join.right.CopyFrom(self.right.plan(session)) + if self.on is not None: + if not isinstance(self.on, list): + if isinstance(self.on, str): + plan.join.using_columns.append(self.on) + else: + plan.join.join_condition.CopyFrom(self.to_attr_or_expression(self.on, session)) + elif len(self.on) > 0: + if isinstance(self.on[0], str): + plan.join.using_columns.extend(cast(str, self.on)) + else: + merge_column = functools.reduce(lambda c1, c2: c1 & c2, self.on) + plan.join.join_condition.CopyFrom(cast(Column, merge_column).to_plan(session)) + plan.join.join_type = self.how + return plan + + def print(self, indent: int = 0) -> str: + i = " " * indent + o = " " * (indent + LogicalPlan.INDENT) + n = indent + LogicalPlan.INDENT * 2 + return ( + f"{i}\n{o}" + f"left=\n{self.left.print(n)}\n{o}right=\n{self.right.print(n)}" + ) + + def _repr_html_(self) -> str: + return f""" +
      +
    • + Join
      + Left: {self.left._repr_html_()} + Right: {self.right._repr_html_()} +
    • +
    + """ + + +class SetOperation(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + other: Optional["LogicalPlan"], + set_op: str, + is_all: bool = True, + by_name: bool = False, + allow_missing_columns: bool = False, + ) -> None: + super().__init__(child) + self.other = other + self.by_name = by_name + self.is_all = is_all + self.set_op = set_op + self.allow_missing_columns = allow_missing_columns + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + if self._child is not None: + plan.set_op.left_input.CopyFrom(self._child.plan(session)) + if self.other is not None: + plan.set_op.right_input.CopyFrom(self.other.plan(session)) + if self.set_op == "union": + plan.set_op.set_op_type = proto.SetOperation.SET_OP_TYPE_UNION + elif self.set_op == "intersect": + plan.set_op.set_op_type = proto.SetOperation.SET_OP_TYPE_INTERSECT + elif self.set_op == "except": + plan.set_op.set_op_type = proto.SetOperation.SET_OP_TYPE_EXCEPT + else: + raise PySparkNotImplementedError( + error_class="UNSUPPORTED_OPERATION", + message_parameters={"feature": self.set_op}, + ) + + plan.set_op.is_all = self.is_all + plan.set_op.by_name = self.by_name + plan.set_op.allow_missing_columns = self.allow_missing_columns + return plan + + def print(self, indent: int = 0) -> str: + assert self._child is not None + assert self.other is not None + + i = " " * indent + o = " " * (indent + LogicalPlan.INDENT) + n = indent + LogicalPlan.INDENT * 2 + return ( + f"{i}SetOperation\n{o}child1=\n{self._child.print(n)}" + f"\n{o}child2=\n{self.other.print(n)}" + ) + + def _repr_html_(self) -> str: + assert self._child is not None + assert self.other is not None + + return f""" +
      +
    • + SetOperation
      + Left: {self._child._repr_html_()} + Right: {self.other._repr_html_()} +
    • +
    + """ + + +class Repartition(LogicalPlan): + """Repartition Relation into a different number of partitions.""" + + def __init__(self, child: Optional["LogicalPlan"], num_partitions: int, shuffle: bool) -> None: + super().__init__(child) + self._num_partitions = num_partitions + self._shuffle = shuffle + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + if self._child is not None: + plan.repartition.input.CopyFrom(self._child.plan(session)) + plan.repartition.shuffle = self._shuffle + plan.repartition.num_partitions = self._num_partitions + return plan + + +class RepartitionByExpression(LogicalPlan): + """Repartition Relation into a different number of partitions using Expression""" + + def __init__( + self, + child: Optional["LogicalPlan"], + num_partitions: Optional[int], + columns: List["ColumnOrName"], + ) -> None: + super().__init__(child) + self.num_partitions = num_partitions + self.columns = columns + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + + part_exprs = [] + for c in self.columns: + if isinstance(c, Column): + part_exprs.append(c.to_plan(session)) + elif c == "*": + exp = proto.Expression() + exp.unresolved_star.SetInParent() + part_exprs.append(exp) + else: + part_exprs.append(self.unresolved_attr(c)) + plan.repartition_by_expression.partition_exprs.extend(part_exprs) + + if self._child is not None: + plan.repartition_by_expression.input.CopyFrom(self._child.plan(session)) + if self.num_partitions is not None: + plan.repartition_by_expression.num_partitions = self.num_partitions + return plan + + +class SubqueryAlias(LogicalPlan): + """Alias for a relation.""" + + def __init__(self, child: Optional["LogicalPlan"], alias: str) -> None: + super().__init__(child) + self._alias = alias + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + if self._child is not None: + plan.subquery_alias.input.CopyFrom(self._child.plan(session)) + plan.subquery_alias.alias = self._alias + return plan + + +class SQL(LogicalPlan): + def __init__(self, query: str, args: Optional[Union[Dict[str, Any], List]] = None) -> None: + super().__init__(None) + + if args is not None: + if isinstance(args, Dict): + for k, v in args.items(): + assert isinstance(k, str) + else: + assert isinstance(args, List) + + self._query = query + self._args = args + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.sql.query = self._query + + if self._args is not None and len(self._args) > 0: + if isinstance(self._args, Dict): + for k, v in self._args.items(): + plan.sql.args[k].CopyFrom( + LiteralExpression._from_value(v).to_plan(session).literal + ) + else: + for v in self._args: + plan.sql.pos_args.append( + LiteralExpression._from_value(v).to_plan(session).literal + ) + + return plan + + def command(self, session: "SparkConnectClient") -> proto.Command: + cmd = proto.Command() + cmd.sql_command.sql = self._query + if self._args is not None and len(self._args) > 0: + if isinstance(self._args, Dict): + for k, v in self._args.items(): + cmd.sql_command.args[k].CopyFrom( + LiteralExpression._from_value(v).to_plan(session).literal + ) + else: + for v in self._args: + cmd.sql_command.pos_args.append( + LiteralExpression._from_value(v).to_plan(session).literal + ) + + return cmd + + +class Range(LogicalPlan): + def __init__( + self, + start: int, + end: int, + step: int, + num_partitions: Optional[int] = None, + ) -> None: + super().__init__(None) + self._start = start + self._end = end + self._step = step + self._num_partitions = num_partitions + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.range.start = self._start + plan.range.end = self._end + plan.range.step = self._step + if self._num_partitions is not None: + plan.range.num_partitions = self._num_partitions + return plan + + +class ToSchema(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], schema: DataType) -> None: + super().__init__(child) + self._schema = schema + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.to_schema.input.CopyFrom(self._child.plan(session)) + plan.to_schema.schema.CopyFrom(pyspark_types_to_proto_types(self._schema)) + return plan + + +class WithColumnsRenamed(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], colsMap: Mapping[str, str]) -> None: + super().__init__(child) + self._colsMap = colsMap + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.with_columns_renamed.input.CopyFrom(self._child.plan(session)) + for k, v in self._colsMap.items(): + plan.with_columns_renamed.rename_columns_map[k] = v + return plan + + +class Unpivot(LogicalPlan): + """Logical plan object for a unpivot operation.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + ids: List["ColumnOrName"], + values: Optional[List["ColumnOrName"]], + variable_column_name: str, + value_column_name: str, + ) -> None: + super().__init__(child) + self.ids = ids + self.values = values + self.variable_column_name = variable_column_name + self.value_column_name = value_column_name + + def col_to_expr(self, col: "ColumnOrName", session: "SparkConnectClient") -> proto.Expression: + if isinstance(col, Column): + return col.to_plan(session) + else: + return self.unresolved_attr(col) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.unpivot.input.CopyFrom(self._child.plan(session)) + plan.unpivot.ids.extend([self.col_to_expr(x, session) for x in self.ids]) + if self.values is not None: + plan.unpivot.values.values.extend([self.col_to_expr(x, session) for x in self.values]) + plan.unpivot.variable_column_name = self.variable_column_name + plan.unpivot.value_column_name = self.value_column_name + return plan + + +class CollectMetrics(LogicalPlan): + """Logical plan object for a CollectMetrics operation.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + name: str, + exprs: List["ColumnOrName"], + ) -> None: + super().__init__(child) + self._name = name + self._exprs = exprs + + def col_to_expr(self, col: "ColumnOrName", session: "SparkConnectClient") -> proto.Expression: + if isinstance(col, Column): + return col.to_plan(session) + else: + return self.unresolved_attr(col) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + + plan = proto.Relation() + plan.collect_metrics.input.CopyFrom(self._child.plan(session)) + plan.collect_metrics.name = self._name + plan.collect_metrics.metrics.extend([self.col_to_expr(x, session) for x in self._exprs]) + return plan + + +class NAFill(LogicalPlan): + def __init__( + self, child: Optional["LogicalPlan"], cols: Optional[List[str]], values: List[Any] + ) -> None: + super().__init__(child) + + assert ( + isinstance(values, list) + and len(values) > 0 + and all(isinstance(v, (bool, int, float, str)) for v in values) + ) + + if cols is not None and len(cols) > 0: + assert isinstance(cols, list) and all(isinstance(c, str) for c in cols) + if len(values) > 1: + assert len(cols) == len(values) + + self.cols = cols + self.values = values + + def _convert_value(self, v: Any) -> proto.Expression.Literal: + value = proto.Expression.Literal() + if isinstance(v, bool): + value.boolean = v + elif isinstance(v, int): + value.long = v + elif isinstance(v, float): + value.double = v + else: + value.string = v + return value + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.fill_na.input.CopyFrom(self._child.plan(session)) + if self.cols is not None and len(self.cols) > 0: + plan.fill_na.cols.extend(self.cols) + plan.fill_na.values.extend([self._convert_value(v) for v in self.values]) + return plan + + +class NADrop(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + cols: Optional[List[str]], + min_non_nulls: Optional[int], + ) -> None: + super().__init__(child) + + self.cols = cols + self.min_non_nulls = min_non_nulls + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.drop_na.input.CopyFrom(self._child.plan(session)) + if self.cols is not None and len(self.cols) > 0: + plan.drop_na.cols.extend(self.cols) + if self.min_non_nulls is not None: + plan.drop_na.min_non_nulls = self.min_non_nulls + return plan + + +class NAReplace(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + cols: Optional[List[str]], + replacements: Dict[Any, Any], + ) -> None: + super().__init__(child) + + for old_value, new_value in replacements.items(): + if old_value is not None: + assert isinstance(old_value, (bool, int, float, str)) + if new_value is not None: + assert isinstance(new_value, (bool, int, float, str)) + + self.cols = cols + self.replacements = replacements + + def _convert_int_to_float(self, v: Any) -> Any: + # a bool is also an int + if v is not None and not isinstance(v, bool) and isinstance(v, int): + return float(v) + else: + return v + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.replace.input.CopyFrom(self._child.plan(session)) + if self.cols is not None and len(self.cols) > 0: + plan.replace.cols.extend(self.cols) + if len(self.replacements) > 0: + for old_value, new_value in self.replacements.items(): + replacement = proto.NAReplace.Replacement() + replacement.old_value.CopyFrom( + LiteralExpression._from_value(self._convert_int_to_float(old_value)) + .to_plan(session) + .literal + ) + replacement.new_value.CopyFrom( + LiteralExpression._from_value(self._convert_int_to_float(new_value)) + .to_plan(session) + .literal + ) + plan.replace.replacements.append(replacement) + return plan + + +class StatSummary(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], statistics: List[str]) -> None: + super().__init__(child) + self.statistics = statistics + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.summary.input.CopyFrom(self._child.plan(session)) + plan.summary.statistics.extend(self.statistics) + return plan + + +class StatDescribe(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], cols: List[str]) -> None: + super().__init__(child) + self.cols = cols + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.describe.input.CopyFrom(self._child.plan(session)) + plan.describe.cols.extend(self.cols) + return plan + + +class StatCov(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], col1: str, col2: str) -> None: + super().__init__(child) + self._col1 = col1 + self._col2 = col2 + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.cov.input.CopyFrom(self._child.plan(session)) + plan.cov.col1 = self._col1 + plan.cov.col2 = self._col2 + return plan + + +class StatApproxQuantile(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + cols: List[str], + probabilities: List[float], + relativeError: float, + ) -> None: + super().__init__(child) + self._cols = cols + self._probabilities = probabilities + self._relativeError = relativeError + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.approx_quantile.input.CopyFrom(self._child.plan(session)) + plan.approx_quantile.cols.extend(self._cols) + plan.approx_quantile.probabilities.extend(self._probabilities) + plan.approx_quantile.relative_error = self._relativeError + return plan + + +class StatCrosstab(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], col1: str, col2: str) -> None: + super().__init__(child) + self.col1 = col1 + self.col2 = col2 + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.crosstab.input.CopyFrom(self._child.plan(session)) + plan.crosstab.col1 = self.col1 + plan.crosstab.col2 = self.col2 + return plan + + +class StatFreqItems(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + cols: List[str], + support: float, + ) -> None: + super().__init__(child) + self._cols = cols + self._support = support + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.freq_items.input.CopyFrom(self._child.plan(session)) + plan.freq_items.cols.extend(self._cols) + plan.freq_items.support = self._support + return plan + + +class StatSampleBy(LogicalPlan): + def __init__( + self, + child: Optional["LogicalPlan"], + col: "ColumnOrName", + fractions: Dict[Any, float], + seed: Optional[int], + ) -> None: + super().__init__(child) + + assert col is not None and isinstance(col, (Column, str)) + + assert fractions is not None and isinstance(fractions, dict) + for k, v in fractions.items(): + assert v is not None and isinstance(v, float) + + assert seed is None or isinstance(seed, int) + + if isinstance(col, Column): + self._col = col + else: + self._col = Column(ColumnReference(col)) + + self._fractions = fractions + + self._seed = seed + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.sample_by.input.CopyFrom(self._child.plan(session)) + plan.sample_by.col.CopyFrom(self._col._expr.to_plan(session)) + if len(self._fractions) > 0: + for k, v in self._fractions.items(): + fraction = proto.StatSampleBy.Fraction() + fraction.stratum.CopyFrom(LiteralExpression._from_value(k).to_plan(session).literal) + fraction.fraction = float(v) + plan.sample_by.fractions.append(fraction) + if self._seed is not None: + plan.sample_by.seed = self._seed + return plan + + +class StatCorr(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], col1: str, col2: str, method: str) -> None: + super().__init__(child) + self._col1 = col1 + self._col2 = col2 + self._method = method + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.corr.input.CopyFrom(self._child.plan(session)) + plan.corr.col1 = self._col1 + plan.corr.col2 = self._col2 + plan.corr.method = self._method + return plan + + +class ToDF(LogicalPlan): + def __init__(self, child: Optional["LogicalPlan"], cols: Sequence[str]) -> None: + super().__init__(child) + self._cols = cols + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.to_df.input.CopyFrom(self._child.plan(session)) + plan.to_df.column_names.extend(self._cols) + return plan + + +class CreateView(LogicalPlan): + def __init__( + self, child: Optional["LogicalPlan"], name: str, is_global: bool, replace: bool + ) -> None: + super().__init__(child) + self._name = name + self._is_global = is_global + self._replace = replace + + def command(self, session: "SparkConnectClient") -> proto.Command: + assert self._child is not None + plan = proto.Command() + + plan.create_dataframe_view.replace = self._replace + plan.create_dataframe_view.is_global = self._is_global + plan.create_dataframe_view.name = self._name + plan.create_dataframe_view.input.CopyFrom(self._child.plan(session)) + return plan + + +class WriteOperation(LogicalPlan): + def __init__(self, child: "LogicalPlan") -> None: + super(WriteOperation, self).__init__(child) + self.source: Optional[str] = None + self.path: Optional[str] = None + self.table_name: Optional[str] = None + self.table_save_method: Optional[str] = None + self.mode: Optional[str] = None + self.sort_cols: List[str] = [] + self.partitioning_cols: List[str] = [] + self.options: Dict[str, Optional[str]] = {} + self.num_buckets: int = -1 + self.bucket_cols: List[str] = [] + + def command(self, session: "SparkConnectClient") -> proto.Command: + assert self._child is not None + plan = proto.Command() + + plan.write_operation.input.CopyFrom(self._child.plan(session)) + if self.source is not None: + plan.write_operation.source = self.source + plan.write_operation.sort_column_names.extend(self.sort_cols) + plan.write_operation.partitioning_columns.extend(self.partitioning_cols) + + if self.num_buckets > 0: + plan.write_operation.bucket_by.bucket_column_names.extend(self.bucket_cols) + plan.write_operation.bucket_by.num_buckets = self.num_buckets + + for k in self.options: + if self.options[k] is None: + plan.write_operation.options.pop(k, None) + else: + plan.write_operation.options[k] = cast(str, self.options[k]) + + if self.table_name is not None: + plan.write_operation.table.table_name = self.table_name + if self.table_save_method is not None: + tsm = self.table_save_method.lower() + if tsm == "save_as_table": + plan.write_operation.table.save_method = ( + proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_SAVE_AS_TABLE # noqa: E501 + ) + elif tsm == "insert_into": + plan.write_operation.table.save_method = ( + proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_INSERT_INTO + ) + else: + raise ValueError( + f"Unknown TestSaveMethod value for DataFrame: {self.table_save_method}" + ) + elif self.path is not None: + plan.write_operation.path = self.path + + if self.mode is not None: + wm = self.mode.lower() + if wm == "append": + plan.write_operation.mode = proto.WriteOperation.SaveMode.SAVE_MODE_APPEND + elif wm == "overwrite": + plan.write_operation.mode = proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE + elif wm == "error": + plan.write_operation.mode = proto.WriteOperation.SaveMode.SAVE_MODE_ERROR_IF_EXISTS + elif wm == "ignore": + plan.write_operation.mode = proto.WriteOperation.SaveMode.SAVE_MODE_IGNORE + else: + raise ValueError(f"Unknown SaveMode value for DataFrame: {self.mode}") + return plan + + def print(self, indent: int = 0) -> str: + i = " " * indent + return ( + f"{i}" + f"" + ) + + def _repr_html_(self) -> str: + return ( + f"
    • WriteOperation
      source='{self.source}'
      " + f"path: '{self.path}
      " + f"table_name: '{self.table_name}'
      " + f"table_save_method: '{self.table_save_method}'
      " + f"mode: '{self.mode}'
      " + f"sort_cols: '{self.sort_cols}'
      " + f"partitioning_cols: '{self.partitioning_cols}'
      " + f"num_buckets: '{self.num_buckets}'
      " + f"bucket_cols: '{self.bucket_cols}'
      " + f"options: '{self.options}'
      " + f"
    " + ) + + +class WriteOperationV2(LogicalPlan): + def __init__(self, child: "LogicalPlan", table_name: str) -> None: + super(WriteOperationV2, self).__init__(child) + self.table_name: Optional[str] = table_name + self.provider: Optional[str] = None + self.partitioning_columns: List["ColumnOrName"] = [] + self.options: dict[str, Optional[str]] = {} + self.table_properties: dict[str, Optional[str]] = {} + self.mode: Optional[str] = None + self.overwrite_condition: Optional["ColumnOrName"] = None + + def col_to_expr(self, col: "ColumnOrName", session: "SparkConnectClient") -> proto.Expression: + if isinstance(col, Column): + return col.to_plan(session) + else: + return self.unresolved_attr(col) + + def command(self, session: "SparkConnectClient") -> proto.Command: + assert self._child is not None + plan = proto.Command() + plan.write_operation_v2.input.CopyFrom(self._child.plan(session)) + if self.table_name is not None: + plan.write_operation_v2.table_name = self.table_name + if self.provider is not None: + plan.write_operation_v2.provider = self.provider + + plan.write_operation_v2.partitioning_columns.extend( + [self.col_to_expr(x, session) for x in self.partitioning_columns] + ) + + for k in self.options: + if self.options[k] is None: + plan.write_operation_v2.options.pop(k, None) + else: + plan.write_operation_v2.options[k] = cast(str, self.options[k]) + + for k in self.table_properties: + if self.table_properties[k] is None: + plan.write_operation_v2.table_properties.pop(k, None) + else: + plan.write_operation_v2.table_properties[k] = cast(str, self.table_properties[k]) + + if self.mode is not None: + wm = self.mode.lower() + if wm == "create": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_CREATE + elif wm == "overwrite": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_OVERWRITE + elif wm == "overwrite_partition": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_OVERWRITE_PARTITIONS + elif wm == "append": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_APPEND + elif wm == "replace": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_REPLACE + if self.overwrite_condition is not None: + plan.write_operation_v2.overwrite_condition.CopyFrom( + self.col_to_expr(self.overwrite_condition, session) + ) + elif wm == "create_or_replace": + plan.write_operation_v2.mode = proto.WriteOperationV2.Mode.MODE_CREATE_OR_REPLACE + else: + raise ValueError(f"Unknown Mode value for DataFrame: {self.mode}") + return plan + + +class WriteStreamOperation(LogicalPlan): + def __init__(self, child: "LogicalPlan") -> None: + super(WriteStreamOperation, self).__init__(child) + self.write_op = proto.WriteStreamOperationStart() + + def command(self, session: "SparkConnectClient") -> proto.Command: + assert self._child is not None + self.write_op.input.CopyFrom(self._child.plan(session)) + cmd = proto.Command() + cmd.write_stream_operation_start.CopyFrom(self.write_op) + return cmd + + +# Catalog API (internal-only) + + +class CurrentDatabase(LogicalPlan): + def __init__(self) -> None: + super().__init__(None) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + return proto.Relation(catalog=proto.Catalog(current_database=proto.CurrentDatabase())) + + +class SetCurrentDatabase(LogicalPlan): + def __init__(self, db_name: str) -> None: + super().__init__(None) + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation() + plan.catalog.set_current_database.db_name = self._db_name + return plan + + +class ListDatabases(LogicalPlan): + def __init__(self, pattern: Optional[str] = None) -> None: + super().__init__(None) + self._pattern = pattern + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(list_databases=proto.ListDatabases())) + if self._pattern is not None: + plan.catalog.list_databases.pattern = self._pattern + return plan + + +class ListTables(LogicalPlan): + def __init__(self, db_name: Optional[str] = None, pattern: Optional[str] = None) -> None: + super().__init__(None) + self._db_name = db_name + self._pattern = pattern + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(list_tables=proto.ListTables())) + if self._db_name is not None: + plan.catalog.list_tables.db_name = self._db_name + if self._pattern is not None: + plan.catalog.list_tables.pattern = self._pattern + return plan + + +class ListFunctions(LogicalPlan): + def __init__(self, db_name: Optional[str] = None, pattern: Optional[str] = None) -> None: + super().__init__(None) + self._db_name = db_name + self._pattern = pattern + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(list_functions=proto.ListFunctions())) + if self._db_name is not None: + plan.catalog.list_functions.db_name = self._db_name + if self._pattern is not None: + plan.catalog.list_functions.pattern = self._pattern + return plan + + +class ListColumns(LogicalPlan): + def __init__(self, table_name: str, db_name: Optional[str] = None) -> None: + super().__init__(None) + self._table_name = table_name + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(list_columns=proto.ListColumns())) + plan.catalog.list_columns.table_name = self._table_name + if self._db_name is not None: + plan.catalog.list_columns.db_name = self._db_name + return plan + + +class GetDatabase(LogicalPlan): + def __init__(self, db_name: str) -> None: + super().__init__(None) + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(get_database=proto.GetDatabase())) + plan.catalog.get_database.db_name = self._db_name + return plan + + +class GetTable(LogicalPlan): + def __init__(self, table_name: str, db_name: Optional[str] = None) -> None: + super().__init__(None) + self._table_name = table_name + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(get_table=proto.GetTable())) + plan.catalog.get_table.table_name = self._table_name + if self._db_name is not None: + plan.catalog.get_table.db_name = self._db_name + return plan + + +class GetFunction(LogicalPlan): + def __init__(self, function_name: str, db_name: Optional[str] = None) -> None: + super().__init__(None) + self._function_name = function_name + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(get_function=proto.GetFunction())) + plan.catalog.get_function.function_name = self._function_name + if self._db_name is not None: + plan.catalog.get_function.db_name = self._db_name + return plan + + +class DatabaseExists(LogicalPlan): + def __init__(self, db_name: str) -> None: + super().__init__(None) + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(database_exists=proto.DatabaseExists())) + plan.catalog.database_exists.db_name = self._db_name + return plan + + +class TableExists(LogicalPlan): + def __init__(self, table_name: str, db_name: Optional[str] = None) -> None: + super().__init__(None) + self._table_name = table_name + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(table_exists=proto.TableExists())) + plan.catalog.table_exists.table_name = self._table_name + if self._db_name is not None: + plan.catalog.table_exists.db_name = self._db_name + return plan + + +class FunctionExists(LogicalPlan): + def __init__(self, function_name: str, db_name: Optional[str] = None) -> None: + super().__init__(None) + self._function_name = function_name + self._db_name = db_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(function_exists=proto.FunctionExists())) + plan.catalog.function_exists.function_name = self._function_name + if self._db_name is not None: + plan.catalog.function_exists.db_name = self._db_name + return plan + + +class CreateExternalTable(LogicalPlan): + def __init__( + self, + table_name: str, + path: str, + source: Optional[str] = None, + schema: Optional[DataType] = None, + options: Mapping[str, str] = {}, + ) -> None: + super().__init__(None) + self._table_name = table_name + self._path = path + self._source = source + self._schema = schema + self._options = options + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation( + catalog=proto.Catalog(create_external_table=proto.CreateExternalTable()) + ) + plan.catalog.create_external_table.table_name = self._table_name + if self._path is not None: + plan.catalog.create_external_table.path = self._path + if self._source is not None: + plan.catalog.create_external_table.source = self._source + if self._schema is not None: + plan.catalog.create_external_table.schema.CopyFrom( + pyspark_types_to_proto_types(self._schema) + ) + for k in self._options.keys(): + v = self._options.get(k) + if v is not None: + plan.catalog.create_external_table.options[k] = v + return plan + + +class CreateTable(LogicalPlan): + def __init__( + self, + table_name: str, + path: str, + source: Optional[str] = None, + description: Optional[str] = None, + schema: Optional[DataType] = None, + options: Mapping[str, str] = {}, + ) -> None: + super().__init__(None) + self._table_name = table_name + self._path = path + self._source = source + self._description = description + self._schema = schema + self._options = options + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(create_table=proto.CreateTable())) + plan.catalog.create_table.table_name = self._table_name + if self._path is not None: + plan.catalog.create_table.path = self._path + if self._source is not None: + plan.catalog.create_table.source = self._source + if self._description is not None: + plan.catalog.create_table.description = self._description + if self._schema is not None: + plan.catalog.create_table.schema.CopyFrom(pyspark_types_to_proto_types(self._schema)) + for k in self._options.keys(): + v = self._options.get(k) + if v is not None: + plan.catalog.create_table.options[k] = v + return plan + + +class DropTempView(LogicalPlan): + def __init__(self, view_name: str) -> None: + super().__init__(None) + self._view_name = view_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(drop_temp_view=proto.DropTempView())) + plan.catalog.drop_temp_view.view_name = self._view_name + return plan + + +class DropGlobalTempView(LogicalPlan): + def __init__(self, view_name: str) -> None: + super().__init__(None) + self._view_name = view_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation( + catalog=proto.Catalog(drop_global_temp_view=proto.DropGlobalTempView()) + ) + plan.catalog.drop_global_temp_view.view_name = self._view_name + return plan + + +class RecoverPartitions(LogicalPlan): + def __init__(self, table_name: str) -> None: + super().__init__(None) + self._table_name = table_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation( + catalog=proto.Catalog( + recover_partitions=proto.RecoverPartitions(table_name=self._table_name) + ) + ) + return plan + + +class IsCached(LogicalPlan): + def __init__(self, table_name: str) -> None: + super().__init__(None) + self._table_name = table_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation( + catalog=proto.Catalog(is_cached=proto.IsCached(table_name=self._table_name)) + ) + return plan + + +class CacheTable(LogicalPlan): + def __init__(self, table_name: str, storage_level: Optional[StorageLevel] = None) -> None: + super().__init__(None) + self._table_name = table_name + self._storage_level = storage_level + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + _cache_table = proto.CacheTable(table_name=self._table_name) + if self._storage_level: + _cache_table.storage_level.CopyFrom(storage_level_to_proto(self._storage_level)) + plan = proto.Relation(catalog=proto.Catalog(cache_table=_cache_table)) + return plan + + +class UncacheTable(LogicalPlan): + def __init__(self, table_name: str) -> None: + super().__init__(None) + self._table_name = table_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(uncache_table=proto.UncacheTable())) + plan.catalog.uncache_table.table_name = self._table_name + return plan + + +class ClearCache(LogicalPlan): + def __init__(self) -> None: + super().__init__(None) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + return proto.Relation(catalog=proto.Catalog(clear_cache=proto.ClearCache())) + + +class RefreshTable(LogicalPlan): + def __init__(self, table_name: str) -> None: + super().__init__(None) + self._table_name = table_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(refresh_table=proto.RefreshTable())) + plan.catalog.refresh_table.table_name = self._table_name + return plan + + +class RefreshByPath(LogicalPlan): + def __init__(self, path: str) -> None: + super().__init__(None) + self._path = path + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(refresh_by_path=proto.RefreshByPath())) + plan.catalog.refresh_by_path.path = self._path + return plan + + +class CurrentCatalog(LogicalPlan): + def __init__(self) -> None: + super().__init__(None) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + return proto.Relation(catalog=proto.Catalog(current_catalog=proto.CurrentCatalog())) + + +class SetCurrentCatalog(LogicalPlan): + def __init__(self, catalog_name: str) -> None: + super().__init__(None) + self._catalog_name = catalog_name + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(set_current_catalog=proto.SetCurrentCatalog())) + plan.catalog.set_current_catalog.catalog_name = self._catalog_name + return plan + + +class ListCatalogs(LogicalPlan): + def __init__(self, pattern: Optional[str] = None) -> None: + super().__init__(None) + self._pattern = pattern + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = proto.Relation(catalog=proto.Catalog(list_catalogs=proto.ListCatalogs())) + if self._pattern is not None: + plan.catalog.list_catalogs.pattern = self._pattern + return plan + + +class MapPartitions(LogicalPlan): + """Logical plan object for a mapPartitions-equivalent API: mapInPandas, mapInArrow.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + function: "UserDefinedFunction", + cols: List[str], + is_barrier: bool, + ) -> None: + super().__init__(child) + + self._func = function._build_common_inline_user_defined_function(*cols) + self._is_barrier = is_barrier + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.map_partitions.input.CopyFrom(self._child.plan(session)) + plan.map_partitions.func.CopyFrom(self._func.to_plan_udf(session)) + plan.map_partitions.is_barrier = self._is_barrier + return plan + + +class GroupMap(LogicalPlan): + """Logical plan object for a Group Map API: apply, applyInPandas.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + grouping_cols: Sequence[Column], + function: "UserDefinedFunction", + cols: List[str], + ): + assert isinstance(grouping_cols, list) and all(isinstance(c, Column) for c in grouping_cols) + + super().__init__(child) + self._grouping_cols = grouping_cols + self._func = function._build_common_inline_user_defined_function(*cols) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.group_map.input.CopyFrom(self._child.plan(session)) + plan.group_map.grouping_expressions.extend( + [c.to_plan(session) for c in self._grouping_cols] + ) + plan.group_map.func.CopyFrom(self._func.to_plan_udf(session)) + return plan + + +class CoGroupMap(LogicalPlan): + """Logical plan object for a CoGroup Map API: applyInPandas.""" + + def __init__( + self, + input: Optional["LogicalPlan"], + input_grouping_cols: Sequence[Column], + other: Optional["LogicalPlan"], + other_grouping_cols: Sequence[Column], + function: "UserDefinedFunction", + cols: List[Column], + ): + assert isinstance(input_grouping_cols, list) and all( + isinstance(c, Column) for c in input_grouping_cols + ) + assert isinstance(other_grouping_cols, list) and all( + isinstance(c, Column) for c in other_grouping_cols + ) + + super().__init__(input) + self._input_grouping_cols = input_grouping_cols + self._other_grouping_cols = other_grouping_cols + self._other = cast(LogicalPlan, other) + self._func = function._build_common_inline_user_defined_function(*cols) + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.co_group_map.input.CopyFrom(self._child.plan(session)) + plan.co_group_map.input_grouping_expressions.extend( + [c.to_plan(session) for c in self._input_grouping_cols] + ) + plan.co_group_map.other.CopyFrom(self._other.plan(session)) + plan.co_group_map.other_grouping_expressions.extend( + [c.to_plan(session) for c in self._other_grouping_cols] + ) + plan.co_group_map.func.CopyFrom(self._func.to_plan_udf(session)) + return plan + + +class ApplyInPandasWithState(LogicalPlan): + """Logical plan object for a applyInPandasWithState.""" + + def __init__( + self, + child: Optional["LogicalPlan"], + grouping_cols: Sequence[Column], + function: "UserDefinedFunction", + output_schema: str, + state_schema: str, + output_mode: str, + timeout_conf: str, + cols: List[str], + ): + assert isinstance(grouping_cols, list) and all(isinstance(c, Column) for c in grouping_cols) + + super().__init__(child) + self._grouping_cols = grouping_cols + self._func = function._build_common_inline_user_defined_function(*cols) + self._output_schema = output_schema + self._state_schema = state_schema + self._output_mode = output_mode + self._timeout_conf = timeout_conf + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + assert self._child is not None + plan = self._create_proto_relation() + plan.apply_in_pandas_with_state.input.CopyFrom(self._child.plan(session)) + plan.apply_in_pandas_with_state.grouping_expressions.extend( + [c.to_plan(session) for c in self._grouping_cols] + ) + plan.apply_in_pandas_with_state.func.CopyFrom(self._func.to_plan_udf(session)) + plan.apply_in_pandas_with_state.output_schema = self._output_schema + plan.apply_in_pandas_with_state.state_schema = self._state_schema + plan.apply_in_pandas_with_state.output_mode = self._output_mode + plan.apply_in_pandas_with_state.timeout_conf = self._timeout_conf + return plan + + +class PythonUDTF: + """Represents a Python user-defined table function.""" + + def __init__( + self, + func: Type, + return_type: Union[DataType, str], + eval_type: int, + python_ver: str, + ) -> None: + self._func = func + self._name = func.__name__ + self._return_type: DataType = ( + UnparsedDataType(return_type) if isinstance(return_type, str) else return_type + ) + self._eval_type = eval_type + self._python_ver = python_ver + + def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDTF: + udtf = proto.PythonUDTF() + # Currently the return type cannot be None. + # TODO(SPARK-44380): support `analyze` in Python UDTFs + assert self._return_type is not None + udtf.return_type.CopyFrom(pyspark_types_to_proto_types(self._return_type)) + udtf.eval_type = self._eval_type + try: + udtf.command = CloudPickleSerializer().dumps(self._func) + except pickle.PicklingError: + raise PySparkRuntimeError( + error_class="UDTF_SERIALIZATION_ERROR", + message_parameters={ + "name": self._name, + "message": "Please check the stack trace and " + "make sure the function is serializable.", + }, + ) + udtf.python_ver = self._python_ver + return udtf + + def __repr__(self) -> str: + return ( + f"PythonUDTF({self._name}, {self._return_type}, " + f"{self._eval_type}, {self._python_ver})" + ) + + +class CommonInlineUserDefinedTableFunction(LogicalPlan): + """ + Logical plan object for a user-defined table function with + an inlined defined function body. + """ + + def __init__( + self, + function_name: str, + function: PythonUDTF, + deterministic: bool, + arguments: Sequence[Expression], + ) -> None: + super().__init__(None) + self._function_name = function_name + self._deterministic = deterministic + self._arguments = arguments + self._function = function + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + plan = self._create_proto_relation() + plan.common_inline_user_defined_table_function.function_name = self._function_name + plan.common_inline_user_defined_table_function.deterministic = self._deterministic + if len(self._arguments) > 0: + plan.common_inline_user_defined_table_function.arguments.extend( + [arg.to_plan(session) for arg in self._arguments] + ) + plan.common_inline_user_defined_table_function.python_udtf.CopyFrom( + self._function.to_plan(session) + ) + return plan + + def udtf_plan( + self, session: "SparkConnectClient" + ) -> "proto.CommonInlineUserDefinedTableFunction": + """ + Compared to `plan`, it returns a `proto.CommonInlineUserDefinedTableFunction` + instead of a `proto.Relation`. + """ + plan = proto.CommonInlineUserDefinedTableFunction() + plan.function_name = self._function_name + plan.deterministic = self._deterministic + if len(self._arguments) > 0: + plan.arguments.extend([arg.to_plan(session) for arg in self._arguments]) + plan.python_udtf.CopyFrom(cast(proto.PythonUDF, self._function.to_plan(session))) + return plan + + def __repr__(self) -> str: + return f"{self._function_name}({', '.join([str(arg) for arg in self._arguments])})" + + +class CachedRelation(LogicalPlan): + def __init__(self, plan: proto.Relation) -> None: + super(CachedRelation, self).__init__(None) + self._plan = plan + # Update the plan ID based on the incremented counter. + self._plan.common.plan_id = self._plan_id + + def plan(self, session: "SparkConnectClient") -> proto.Relation: + return self._plan diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/__init__.py new file mode 100644 index 0000000..3e8d074 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/__init__.py @@ -0,0 +1,25 @@ +# +# 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.sql.connect.proto.base_pb2_grpc import * +from pyspark.sql.connect.proto.base_pb2 import * +from pyspark.sql.connect.proto.types_pb2 import * +from pyspark.sql.connect.proto.commands_pb2 import * +from pyspark.sql.connect.proto.expressions_pb2 import * +from pyspark.sql.connect.proto.relations_pb2 import * +from pyspark.sql.connect.proto.catalog_pb2 import * +from pyspark.sql.connect.proto.common_pb2 import * diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.py new file mode 100644 index 0000000..fa1868b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.py @@ -0,0 +1,203 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/base.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2 +from pyspark.sql.connect.proto import commands_pb2 as spark_dot_connect_dot_commands__pb2 +from pyspark.sql.connect.proto import common_pb2 as spark_dot_connect_dot_common__pb2 +from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2 +from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2 +from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto"t\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.CommandH\x00R\x07\x63ommandB\t\n\x07op_type"z\n\x0bUserContext\x12\x17\n\x07user_id\x18\x01 \x01(\tR\x06userId\x12\x1b\n\tuser_name\x18\x02 \x01(\tR\x08userName\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions"\xf5\x12\n\x12\x41nalyzePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12\x42\n\x06schema\x18\x04 \x01(\x0b\x32(.spark.connect.AnalyzePlanRequest.SchemaH\x00R\x06schema\x12\x45\n\x07\x65xplain\x18\x05 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.ExplainH\x00R\x07\x65xplain\x12O\n\x0btree_string\x18\x06 \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.TreeStringH\x00R\ntreeString\x12\x46\n\x08is_local\x18\x07 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.IsLocalH\x00R\x07isLocal\x12R\n\x0cis_streaming\x18\x08 \x01(\x0b\x32-.spark.connect.AnalyzePlanRequest.IsStreamingH\x00R\x0bisStreaming\x12O\n\x0binput_files\x18\t \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.InputFilesH\x00R\ninputFiles\x12U\n\rspark_version\x18\n \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SparkVersionH\x00R\x0csparkVersion\x12I\n\tddl_parse\x18\x0b \x01(\x0b\x32*.spark.connect.AnalyzePlanRequest.DDLParseH\x00R\x08\x64\x64lParse\x12X\n\x0esame_semantics\x18\x0c \x01(\x0b\x32/.spark.connect.AnalyzePlanRequest.SameSemanticsH\x00R\rsameSemantics\x12U\n\rsemantic_hash\x18\r \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SemanticHashH\x00R\x0csemanticHash\x12\x45\n\x07persist\x18\x0e \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.PersistH\x00R\x07persist\x12K\n\tunpersist\x18\x0f \x01(\x0b\x32+.spark.connect.AnalyzePlanRequest.UnpersistH\x00R\tunpersist\x12_\n\x11get_storage_level\x18\x10 \x01(\x0b\x32\x31.spark.connect.AnalyzePlanRequest.GetStorageLevelH\x00R\x0fgetStorageLevel\x1a\x31\n\x06Schema\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\xbb\x02\n\x07\x45xplain\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12X\n\x0c\x65xplain_mode\x18\x02 \x01(\x0e\x32\x35.spark.connect.AnalyzePlanRequest.Explain.ExplainModeR\x0b\x65xplainMode"\xac\x01\n\x0b\x45xplainMode\x12\x1c\n\x18\x45XPLAIN_MODE_UNSPECIFIED\x10\x00\x12\x17\n\x13\x45XPLAIN_MODE_SIMPLE\x10\x01\x12\x19\n\x15\x45XPLAIN_MODE_EXTENDED\x10\x02\x12\x18\n\x14\x45XPLAIN_MODE_CODEGEN\x10\x03\x12\x15\n\x11\x45XPLAIN_MODE_COST\x10\x04\x12\x1a\n\x16\x45XPLAIN_MODE_FORMATTED\x10\x05\x1aZ\n\nTreeString\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12\x19\n\x05level\x18\x02 \x01(\x05H\x00R\x05level\x88\x01\x01\x42\x08\n\x06_level\x1a\x32\n\x07IsLocal\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x36\n\x0bIsStreaming\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x35\n\nInputFiles\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x0e\n\x0cSparkVersion\x1a)\n\x08\x44\x44LParse\x12\x1d\n\nddl_string\x18\x01 \x01(\tR\tddlString\x1ay\n\rSameSemantics\x12\x34\n\x0btarget_plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\ntargetPlan\x12\x32\n\nother_plan\x18\x02 \x01(\x0b\x32\x13.spark.connect.PlanR\totherPlan\x1a\x37\n\x0cSemanticHash\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x97\x01\n\x07Persist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x45\n\rstorage_level\x18\x02 \x01(\x0b\x32\x1b.spark.connect.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level\x1an\n\tUnpersist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x1f\n\x08\x62locking\x18\x02 \x01(\x08H\x00R\x08\x62locking\x88\x01\x01\x42\x0b\n\t_blocking\x1a\x46\n\x0fGetStorageLevel\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relationB\t\n\x07\x61nalyzeB\x0e\n\x0c_client_type"\x99\r\n\x13\x41nalyzePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x43\n\x06schema\x18\x02 \x01(\x0b\x32).spark.connect.AnalyzePlanResponse.SchemaH\x00R\x06schema\x12\x46\n\x07\x65xplain\x18\x03 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.ExplainH\x00R\x07\x65xplain\x12P\n\x0btree_string\x18\x04 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.TreeStringH\x00R\ntreeString\x12G\n\x08is_local\x18\x05 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.IsLocalH\x00R\x07isLocal\x12S\n\x0cis_streaming\x18\x06 \x01(\x0b\x32..spark.connect.AnalyzePlanResponse.IsStreamingH\x00R\x0bisStreaming\x12P\n\x0binput_files\x18\x07 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.InputFilesH\x00R\ninputFiles\x12V\n\rspark_version\x18\x08 \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SparkVersionH\x00R\x0csparkVersion\x12J\n\tddl_parse\x18\t \x01(\x0b\x32+.spark.connect.AnalyzePlanResponse.DDLParseH\x00R\x08\x64\x64lParse\x12Y\n\x0esame_semantics\x18\n \x01(\x0b\x32\x30.spark.connect.AnalyzePlanResponse.SameSemanticsH\x00R\rsameSemantics\x12V\n\rsemantic_hash\x18\x0b \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SemanticHashH\x00R\x0csemanticHash\x12\x46\n\x07persist\x18\x0c \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.PersistH\x00R\x07persist\x12L\n\tunpersist\x18\r \x01(\x0b\x32,.spark.connect.AnalyzePlanResponse.UnpersistH\x00R\tunpersist\x12`\n\x11get_storage_level\x18\x0e \x01(\x0b\x32\x32.spark.connect.AnalyzePlanResponse.GetStorageLevelH\x00R\x0fgetStorageLevel\x1a\x39\n\x06Schema\x12/\n\x06schema\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1a\x30\n\x07\x45xplain\x12%\n\x0e\x65xplain_string\x18\x01 \x01(\tR\rexplainString\x1a-\n\nTreeString\x12\x1f\n\x0btree_string\x18\x01 \x01(\tR\ntreeString\x1a$\n\x07IsLocal\x12\x19\n\x08is_local\x18\x01 \x01(\x08R\x07isLocal\x1a\x30\n\x0bIsStreaming\x12!\n\x0cis_streaming\x18\x01 \x01(\x08R\x0bisStreaming\x1a"\n\nInputFiles\x12\x14\n\x05\x66iles\x18\x01 \x03(\tR\x05\x66iles\x1a(\n\x0cSparkVersion\x12\x18\n\x07version\x18\x01 \x01(\tR\x07version\x1a;\n\x08\x44\x44LParse\x12/\n\x06parsed\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06parsed\x1a\'\n\rSameSemantics\x12\x16\n\x06result\x18\x01 \x01(\x08R\x06result\x1a&\n\x0cSemanticHash\x12\x16\n\x06result\x18\x01 \x01(\x05R\x06result\x1a\t\n\x07Persist\x1a\x0b\n\tUnpersist\x1aS\n\x0fGetStorageLevel\x12@\n\rstorage_level\x18\x01 \x01(\x0b\x32\x1b.spark.connect.StorageLevelR\x0cstorageLevelB\x08\n\x06result"\xa0\x04\n\x12\x45xecutePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12&\n\x0coperation_id\x18\x06 \x01(\tH\x00R\x0boperationId\x88\x01\x01\x12\'\n\x04plan\x18\x03 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x12X\n\x0frequest_options\x18\x05 \x03(\x0b\x32/.spark.connect.ExecutePlanRequest.RequestOptionR\x0erequestOptions\x12\x12\n\x04tags\x18\x07 \x03(\tR\x04tags\x1a\xa5\x01\n\rRequestOption\x12K\n\x10reattach_options\x18\x01 \x01(\x0b\x32\x1e.spark.connect.ReattachOptionsH\x00R\x0freattachOptions\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textensionB\x10\n\x0erequest_optionB\x0f\n\r_operation_idB\x0e\n\x0c_client_type"\x99\x0f\n\x13\x45xecutePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12!\n\x0coperation_id\x18\x0c \x01(\tR\x0boperationId\x12\x1f\n\x0bresponse_id\x18\r \x01(\tR\nresponseId\x12P\n\x0b\x61rrow_batch\x18\x02 \x01(\x0b\x32-.spark.connect.ExecutePlanResponse.ArrowBatchH\x00R\narrowBatch\x12\x63\n\x12sql_command_result\x18\x05 \x01(\x0b\x32\x33.spark.connect.ExecutePlanResponse.SqlCommandResultH\x00R\x10sqlCommandResult\x12~\n#write_stream_operation_start_result\x18\x08 \x01(\x0b\x32..spark.connect.WriteStreamOperationStartResultH\x00R\x1fwriteStreamOperationStartResult\x12q\n\x1estreaming_query_command_result\x18\t \x01(\x0b\x32*.spark.connect.StreamingQueryCommandResultH\x00R\x1bstreamingQueryCommandResult\x12k\n\x1cget_resources_command_result\x18\n \x01(\x0b\x32(.spark.connect.GetResourcesCommandResultH\x00R\x19getResourcesCommandResult\x12\x87\x01\n&streaming_query_manager_command_result\x18\x0b \x01(\x0b\x32\x31.spark.connect.StreamingQueryManagerCommandResultH\x00R"streamingQueryManagerCommandResult\x12\\\n\x0fresult_complete\x18\x0e \x01(\x0b\x32\x31.spark.connect.ExecutePlanResponse.ResultCompleteH\x00R\x0eresultComplete\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x44\n\x07metrics\x18\x04 \x01(\x0b\x32*.spark.connect.ExecutePlanResponse.MetricsR\x07metrics\x12]\n\x10observed_metrics\x18\x06 \x03(\x0b\x32\x32.spark.connect.ExecutePlanResponse.ObservedMetricsR\x0fobservedMetrics\x12/\n\x06schema\x18\x07 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1aG\n\x10SqlCommandResult\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x1a=\n\nArrowBatch\x12\x1b\n\trow_count\x18\x01 \x01(\x03R\x08rowCount\x12\x12\n\x04\x64\x61ta\x18\x02 \x01(\x0cR\x04\x64\x61ta\x1a\x85\x04\n\x07Metrics\x12Q\n\x07metrics\x18\x01 \x03(\x0b\x32\x37.spark.connect.ExecutePlanResponse.Metrics.MetricObjectR\x07metrics\x1a\xcc\x02\n\x0cMetricObject\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x17\n\x07plan_id\x18\x02 \x01(\x03R\x06planId\x12\x16\n\x06parent\x18\x03 \x01(\x03R\x06parent\x12z\n\x11\x65xecution_metrics\x18\x04 \x03(\x0b\x32M.spark.connect.ExecutePlanResponse.Metrics.MetricObject.ExecutionMetricsEntryR\x10\x65xecutionMetrics\x1a{\n\x15\x45xecutionMetricsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ExecutePlanResponse.Metrics.MetricValueR\x05value:\x02\x38\x01\x1aX\n\x0bMetricValue\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x14\n\x05value\x18\x02 \x01(\x03R\x05value\x12\x1f\n\x0bmetric_type\x18\x03 \x01(\tR\nmetricType\x1a`\n\x0fObservedMetrics\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1a\x10\n\x0eResultCompleteB\x0f\n\rresponse_type"A\n\x08KeyValue\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x19\n\x05value\x18\x02 \x01(\tH\x00R\x05value\x88\x01\x01\x42\x08\n\x06_value"\x84\x08\n\rConfigRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12\x44\n\toperation\x18\x03 \x01(\x0b\x32&.spark.connect.ConfigRequest.OperationR\toperation\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x00R\nclientType\x88\x01\x01\x1a\xf2\x03\n\tOperation\x12\x34\n\x03set\x18\x01 \x01(\x0b\x32 .spark.connect.ConfigRequest.SetH\x00R\x03set\x12\x34\n\x03get\x18\x02 \x01(\x0b\x32 .spark.connect.ConfigRequest.GetH\x00R\x03get\x12W\n\x10get_with_default\x18\x03 \x01(\x0b\x32+.spark.connect.ConfigRequest.GetWithDefaultH\x00R\x0egetWithDefault\x12G\n\nget_option\x18\x04 \x01(\x0b\x32&.spark.connect.ConfigRequest.GetOptionH\x00R\tgetOption\x12>\n\x07get_all\x18\x05 \x01(\x0b\x32#.spark.connect.ConfigRequest.GetAllH\x00R\x06getAll\x12:\n\x05unset\x18\x06 \x01(\x0b\x32".spark.connect.ConfigRequest.UnsetH\x00R\x05unset\x12P\n\ris_modifiable\x18\x07 \x01(\x0b\x32).spark.connect.ConfigRequest.IsModifiableH\x00R\x0cisModifiableB\t\n\x07op_type\x1a\x34\n\x03Set\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x1a\x19\n\x03Get\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a?\n\x0eGetWithDefault\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x1a\x1f\n\tGetOption\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a\x30\n\x06GetAll\x12\x1b\n\x06prefix\x18\x01 \x01(\tH\x00R\x06prefix\x88\x01\x01\x42\t\n\x07_prefix\x1a\x1b\n\x05Unset\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a"\n\x0cIsModifiable\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keysB\x0e\n\x0c_client_type"z\n\x0e\x43onfigResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12-\n\x05pairs\x18\x02 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x12\x1a\n\x08warnings\x18\x03 \x03(\tR\x08warnings"\xe7\x06\n\x13\x41\x64\x64\x41rtifactsRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x06 \x01(\tH\x01R\nclientType\x88\x01\x01\x12@\n\x05\x62\x61tch\x18\x03 \x01(\x0b\x32(.spark.connect.AddArtifactsRequest.BatchH\x00R\x05\x62\x61tch\x12Z\n\x0b\x62\x65gin_chunk\x18\x04 \x01(\x0b\x32\x37.spark.connect.AddArtifactsRequest.BeginChunkedArtifactH\x00R\nbeginChunk\x12H\n\x05\x63hunk\x18\x05 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkH\x00R\x05\x63hunk\x1a\x35\n\rArtifactChunk\x12\x12\n\x04\x64\x61ta\x18\x01 \x01(\x0cR\x04\x64\x61ta\x12\x10\n\x03\x63rc\x18\x02 \x01(\x03R\x03\x63rc\x1ao\n\x13SingleChunkArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x44\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x04\x64\x61ta\x1a]\n\x05\x42\x61tch\x12T\n\tartifacts\x18\x01 \x03(\x0b\x32\x36.spark.connect.AddArtifactsRequest.SingleChunkArtifactR\tartifacts\x1a\xc1\x01\n\x14\x42\x65ginChunkedArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x1f\n\x0btotal_bytes\x18\x02 \x01(\x03R\ntotalBytes\x12\x1d\n\nnum_chunks\x18\x03 \x01(\x03R\tnumChunks\x12U\n\rinitial_chunk\x18\x04 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x0cinitialChunkB\t\n\x07payloadB\x0e\n\x0c_client_type"\xbc\x01\n\x14\x41\x64\x64\x41rtifactsResponse\x12Q\n\tartifacts\x18\x01 \x03(\x0b\x32\x33.spark.connect.AddArtifactsResponse.ArtifactSummaryR\tartifacts\x1aQ\n\x0f\x41rtifactSummary\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12*\n\x11is_crc_successful\x18\x02 \x01(\x08R\x0fisCrcSuccessful"\xc3\x01\n\x17\x41rtifactStatusesRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x00R\nclientType\x88\x01\x01\x12\x14\n\x05names\x18\x04 \x03(\tR\x05namesB\x0e\n\x0c_client_type"\x8c\x02\n\x18\x41rtifactStatusesResponse\x12Q\n\x08statuses\x18\x01 \x03(\x0b\x32\x35.spark.connect.ArtifactStatusesResponse.StatusesEntryR\x08statuses\x1a(\n\x0e\x41rtifactStatus\x12\x16\n\x06\x65xists\x18\x01 \x01(\x08R\x06\x65xists\x1as\n\rStatusesEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ArtifactStatusesResponse.ArtifactStatusR\x05value:\x02\x38\x01"\xd8\x03\n\x10InterruptRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12T\n\x0einterrupt_type\x18\x04 \x01(\x0e\x32-.spark.connect.InterruptRequest.InterruptTypeR\rinterruptType\x12%\n\roperation_tag\x18\x05 \x01(\tH\x00R\x0coperationTag\x12#\n\x0coperation_id\x18\x06 \x01(\tH\x00R\x0boperationId"\x80\x01\n\rInterruptType\x12\x1e\n\x1aINTERRUPT_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12INTERRUPT_TYPE_ALL\x10\x01\x12\x16\n\x12INTERRUPT_TYPE_TAG\x10\x02\x12\x1f\n\x1bINTERRUPT_TYPE_OPERATION_ID\x10\x03\x42\x0b\n\tinterruptB\x0e\n\x0c_client_type"[\n\x11InterruptResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\'\n\x0finterrupted_ids\x18\x02 \x03(\tR\x0einterruptedIds"5\n\x0fReattachOptions\x12"\n\x0creattachable\x18\x01 \x01(\x08R\x0creattachable"\x93\x02\n\x16ReattachExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x00R\nclientType\x88\x01\x01\x12-\n\x10last_response_id\x18\x05 \x01(\tH\x01R\x0elastResponseId\x88\x01\x01\x42\x0e\n\x0c_client_typeB\x13\n\x11_last_response_id"\xc6\x03\n\x15ReleaseExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x12R\n\x0brelease_all\x18\x05 \x01(\x0b\x32/.spark.connect.ReleaseExecuteRequest.ReleaseAllH\x00R\nreleaseAll\x12X\n\rrelease_until\x18\x06 \x01(\x0b\x32\x31.spark.connect.ReleaseExecuteRequest.ReleaseUntilH\x00R\x0creleaseUntil\x1a\x0c\n\nReleaseAll\x1a/\n\x0cReleaseUntil\x12\x1f\n\x0bresponse_id\x18\x01 \x01(\tR\nresponseIdB\t\n\x07releaseB\x0e\n\x0c_client_type"p\n\x16ReleaseExecuteResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12&\n\x0coperation_id\x18\x02 \x01(\tH\x00R\x0boperationId\x88\x01\x01\x42\x0f\n\r_operation_id2\xe7\x05\n\x13SparkConnectService\x12X\n\x0b\x45xecutePlan\x12!.spark.connect.ExecutePlanRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12V\n\x0b\x41nalyzePlan\x12!.spark.connect.AnalyzePlanRequest\x1a".spark.connect.AnalyzePlanResponse"\x00\x12G\n\x06\x43onfig\x12\x1c.spark.connect.ConfigRequest\x1a\x1d.spark.connect.ConfigResponse"\x00\x12[\n\x0c\x41\x64\x64\x41rtifacts\x12".spark.connect.AddArtifactsRequest\x1a#.spark.connect.AddArtifactsResponse"\x00(\x01\x12\x63\n\x0e\x41rtifactStatus\x12&.spark.connect.ArtifactStatusesRequest\x1a\'.spark.connect.ArtifactStatusesResponse"\x00\x12P\n\tInterrupt\x12\x1f.spark.connect.InterruptRequest\x1a .spark.connect.InterruptResponse"\x00\x12`\n\x0fReattachExecute\x12%.spark.connect.ReattachExecuteRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12_\n\x0eReleaseExecute\x12$.spark.connect.ReleaseExecuteRequest\x1a%.spark.connect.ReleaseExecuteResponse"\x00\x42\x36\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.base_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY._options = None + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY._serialized_options = b"8\001" + _ARTIFACTSTATUSESRESPONSE_STATUSESENTRY._options = None + _ARTIFACTSTATUSESRESPONSE_STATUSESENTRY._serialized_options = b"8\001" + _PLAN._serialized_start = 219 + _PLAN._serialized_end = 335 + _USERCONTEXT._serialized_start = 337 + _USERCONTEXT._serialized_end = 459 + _ANALYZEPLANREQUEST._serialized_start = 462 + _ANALYZEPLANREQUEST._serialized_end = 2883 + _ANALYZEPLANREQUEST_SCHEMA._serialized_start = 1657 + _ANALYZEPLANREQUEST_SCHEMA._serialized_end = 1706 + _ANALYZEPLANREQUEST_EXPLAIN._serialized_start = 1709 + _ANALYZEPLANREQUEST_EXPLAIN._serialized_end = 2024 + _ANALYZEPLANREQUEST_EXPLAIN_EXPLAINMODE._serialized_start = 1852 + _ANALYZEPLANREQUEST_EXPLAIN_EXPLAINMODE._serialized_end = 2024 + _ANALYZEPLANREQUEST_TREESTRING._serialized_start = 2026 + _ANALYZEPLANREQUEST_TREESTRING._serialized_end = 2116 + _ANALYZEPLANREQUEST_ISLOCAL._serialized_start = 2118 + _ANALYZEPLANREQUEST_ISLOCAL._serialized_end = 2168 + _ANALYZEPLANREQUEST_ISSTREAMING._serialized_start = 2170 + _ANALYZEPLANREQUEST_ISSTREAMING._serialized_end = 2224 + _ANALYZEPLANREQUEST_INPUTFILES._serialized_start = 2226 + _ANALYZEPLANREQUEST_INPUTFILES._serialized_end = 2279 + _ANALYZEPLANREQUEST_SPARKVERSION._serialized_start = 2281 + _ANALYZEPLANREQUEST_SPARKVERSION._serialized_end = 2295 + _ANALYZEPLANREQUEST_DDLPARSE._serialized_start = 2297 + _ANALYZEPLANREQUEST_DDLPARSE._serialized_end = 2338 + _ANALYZEPLANREQUEST_SAMESEMANTICS._serialized_start = 2340 + _ANALYZEPLANREQUEST_SAMESEMANTICS._serialized_end = 2461 + _ANALYZEPLANREQUEST_SEMANTICHASH._serialized_start = 2463 + _ANALYZEPLANREQUEST_SEMANTICHASH._serialized_end = 2518 + _ANALYZEPLANREQUEST_PERSIST._serialized_start = 2521 + _ANALYZEPLANREQUEST_PERSIST._serialized_end = 2672 + _ANALYZEPLANREQUEST_UNPERSIST._serialized_start = 2674 + _ANALYZEPLANREQUEST_UNPERSIST._serialized_end = 2784 + _ANALYZEPLANREQUEST_GETSTORAGELEVEL._serialized_start = 2786 + _ANALYZEPLANREQUEST_GETSTORAGELEVEL._serialized_end = 2856 + _ANALYZEPLANRESPONSE._serialized_start = 2886 + _ANALYZEPLANRESPONSE._serialized_end = 4575 + _ANALYZEPLANRESPONSE_SCHEMA._serialized_start = 3994 + _ANALYZEPLANRESPONSE_SCHEMA._serialized_end = 4051 + _ANALYZEPLANRESPONSE_EXPLAIN._serialized_start = 4053 + _ANALYZEPLANRESPONSE_EXPLAIN._serialized_end = 4101 + _ANALYZEPLANRESPONSE_TREESTRING._serialized_start = 4103 + _ANALYZEPLANRESPONSE_TREESTRING._serialized_end = 4148 + _ANALYZEPLANRESPONSE_ISLOCAL._serialized_start = 4150 + _ANALYZEPLANRESPONSE_ISLOCAL._serialized_end = 4186 + _ANALYZEPLANRESPONSE_ISSTREAMING._serialized_start = 4188 + _ANALYZEPLANRESPONSE_ISSTREAMING._serialized_end = 4236 + _ANALYZEPLANRESPONSE_INPUTFILES._serialized_start = 4238 + _ANALYZEPLANRESPONSE_INPUTFILES._serialized_end = 4272 + _ANALYZEPLANRESPONSE_SPARKVERSION._serialized_start = 4274 + _ANALYZEPLANRESPONSE_SPARKVERSION._serialized_end = 4314 + _ANALYZEPLANRESPONSE_DDLPARSE._serialized_start = 4316 + _ANALYZEPLANRESPONSE_DDLPARSE._serialized_end = 4375 + _ANALYZEPLANRESPONSE_SAMESEMANTICS._serialized_start = 4377 + _ANALYZEPLANRESPONSE_SAMESEMANTICS._serialized_end = 4416 + _ANALYZEPLANRESPONSE_SEMANTICHASH._serialized_start = 4418 + _ANALYZEPLANRESPONSE_SEMANTICHASH._serialized_end = 4456 + _ANALYZEPLANRESPONSE_PERSIST._serialized_start = 2521 + _ANALYZEPLANRESPONSE_PERSIST._serialized_end = 2530 + _ANALYZEPLANRESPONSE_UNPERSIST._serialized_start = 2674 + _ANALYZEPLANRESPONSE_UNPERSIST._serialized_end = 2685 + _ANALYZEPLANRESPONSE_GETSTORAGELEVEL._serialized_start = 4482 + _ANALYZEPLANRESPONSE_GETSTORAGELEVEL._serialized_end = 4565 + _EXECUTEPLANREQUEST._serialized_start = 4578 + _EXECUTEPLANREQUEST._serialized_end = 5122 + _EXECUTEPLANREQUEST_REQUESTOPTION._serialized_start = 4924 + _EXECUTEPLANREQUEST_REQUESTOPTION._serialized_end = 5089 + _EXECUTEPLANRESPONSE._serialized_start = 5125 + _EXECUTEPLANRESPONSE._serialized_end = 7070 + _EXECUTEPLANRESPONSE_SQLCOMMANDRESULT._serialized_start = 6283 + _EXECUTEPLANRESPONSE_SQLCOMMANDRESULT._serialized_end = 6354 + _EXECUTEPLANRESPONSE_ARROWBATCH._serialized_start = 6356 + _EXECUTEPLANRESPONSE_ARROWBATCH._serialized_end = 6417 + _EXECUTEPLANRESPONSE_METRICS._serialized_start = 6420 + _EXECUTEPLANRESPONSE_METRICS._serialized_end = 6937 + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT._serialized_start = 6515 + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT._serialized_end = 6847 + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY._serialized_start = 6724 + _EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY._serialized_end = 6847 + _EXECUTEPLANRESPONSE_METRICS_METRICVALUE._serialized_start = 6849 + _EXECUTEPLANRESPONSE_METRICS_METRICVALUE._serialized_end = 6937 + _EXECUTEPLANRESPONSE_OBSERVEDMETRICS._serialized_start = 6939 + _EXECUTEPLANRESPONSE_OBSERVEDMETRICS._serialized_end = 7035 + _EXECUTEPLANRESPONSE_RESULTCOMPLETE._serialized_start = 7037 + _EXECUTEPLANRESPONSE_RESULTCOMPLETE._serialized_end = 7053 + _KEYVALUE._serialized_start = 7072 + _KEYVALUE._serialized_end = 7137 + _CONFIGREQUEST._serialized_start = 7140 + _CONFIGREQUEST._serialized_end = 8168 + _CONFIGREQUEST_OPERATION._serialized_start = 7360 + _CONFIGREQUEST_OPERATION._serialized_end = 7858 + _CONFIGREQUEST_SET._serialized_start = 7860 + _CONFIGREQUEST_SET._serialized_end = 7912 + _CONFIGREQUEST_GET._serialized_start = 7914 + _CONFIGREQUEST_GET._serialized_end = 7939 + _CONFIGREQUEST_GETWITHDEFAULT._serialized_start = 7941 + _CONFIGREQUEST_GETWITHDEFAULT._serialized_end = 8004 + _CONFIGREQUEST_GETOPTION._serialized_start = 8006 + _CONFIGREQUEST_GETOPTION._serialized_end = 8037 + _CONFIGREQUEST_GETALL._serialized_start = 8039 + _CONFIGREQUEST_GETALL._serialized_end = 8087 + _CONFIGREQUEST_UNSET._serialized_start = 8089 + _CONFIGREQUEST_UNSET._serialized_end = 8116 + _CONFIGREQUEST_ISMODIFIABLE._serialized_start = 8118 + _CONFIGREQUEST_ISMODIFIABLE._serialized_end = 8152 + _CONFIGRESPONSE._serialized_start = 8170 + _CONFIGRESPONSE._serialized_end = 8292 + _ADDARTIFACTSREQUEST._serialized_start = 8295 + _ADDARTIFACTSREQUEST._serialized_end = 9166 + _ADDARTIFACTSREQUEST_ARTIFACTCHUNK._serialized_start = 8682 + _ADDARTIFACTSREQUEST_ARTIFACTCHUNK._serialized_end = 8735 + _ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT._serialized_start = 8737 + _ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT._serialized_end = 8848 + _ADDARTIFACTSREQUEST_BATCH._serialized_start = 8850 + _ADDARTIFACTSREQUEST_BATCH._serialized_end = 8943 + _ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT._serialized_start = 8946 + _ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT._serialized_end = 9139 + _ADDARTIFACTSRESPONSE._serialized_start = 9169 + _ADDARTIFACTSRESPONSE._serialized_end = 9357 + _ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY._serialized_start = 9276 + _ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY._serialized_end = 9357 + _ARTIFACTSTATUSESREQUEST._serialized_start = 9360 + _ARTIFACTSTATUSESREQUEST._serialized_end = 9555 + _ARTIFACTSTATUSESRESPONSE._serialized_start = 9558 + _ARTIFACTSTATUSESRESPONSE._serialized_end = 9826 + _ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS._serialized_start = 9669 + _ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS._serialized_end = 9709 + _ARTIFACTSTATUSESRESPONSE_STATUSESENTRY._serialized_start = 9711 + _ARTIFACTSTATUSESRESPONSE_STATUSESENTRY._serialized_end = 9826 + _INTERRUPTREQUEST._serialized_start = 9829 + _INTERRUPTREQUEST._serialized_end = 10301 + _INTERRUPTREQUEST_INTERRUPTTYPE._serialized_start = 10144 + _INTERRUPTREQUEST_INTERRUPTTYPE._serialized_end = 10272 + _INTERRUPTRESPONSE._serialized_start = 10303 + _INTERRUPTRESPONSE._serialized_end = 10394 + _REATTACHOPTIONS._serialized_start = 10396 + _REATTACHOPTIONS._serialized_end = 10449 + _REATTACHEXECUTEREQUEST._serialized_start = 10452 + _REATTACHEXECUTEREQUEST._serialized_end = 10727 + _RELEASEEXECUTEREQUEST._serialized_start = 10730 + _RELEASEEXECUTEREQUEST._serialized_end = 11184 + _RELEASEEXECUTEREQUEST_RELEASEALL._serialized_start = 11096 + _RELEASEEXECUTEREQUEST_RELEASEALL._serialized_end = 11108 + _RELEASEEXECUTEREQUEST_RELEASEUNTIL._serialized_start = 11110 + _RELEASEEXECUTEREQUEST_RELEASEUNTIL._serialized_end = 11157 + _RELEASEEXECUTERESPONSE._serialized_start = 11186 + _RELEASEEXECUTERESPONSE._serialized_end = 11298 + _SPARKCONNECTSERVICE._serialized_start = 11301 + _SPARKCONNECTSERVICE._serialized_end = 12044 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.pyi new file mode 100644 index 0000000..8fd5fa7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2.pyi @@ -0,0 +1,2718 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.any_pb2 +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.internal.enum_type_wrapper +import google.protobuf.message +import pyspark.sql.connect.proto.commands_pb2 +import pyspark.sql.connect.proto.common_pb2 +import pyspark.sql.connect.proto.expressions_pb2 +import pyspark.sql.connect.proto.relations_pb2 +import pyspark.sql.connect.proto.types_pb2 +import sys +import typing + +if sys.version_info >= (3, 10): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class Plan(google.protobuf.message.Message): + """A [[Plan]] is the structure that carries the runtime information for the execution from the + client to the server. A [[Plan]] can either be of the type [[Relation]] which is a reference + to the underlying logical plan or it can be of the [[Command]] type that is used to execute + commands on the server. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ROOT_FIELD_NUMBER: builtins.int + COMMAND_FIELD_NUMBER: builtins.int + @property + def root(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: ... + @property + def command(self) -> pyspark.sql.connect.proto.commands_pb2.Command: ... + def __init__( + self, + *, + root: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + command: pyspark.sql.connect.proto.commands_pb2.Command | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "command", b"command", "op_type", b"op_type", "root", b"root" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "command", b"command", "op_type", b"op_type", "root", b"root" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["op_type", b"op_type"] + ) -> typing_extensions.Literal["root", "command"] | None: ... + +global___Plan = Plan + +class UserContext(google.protobuf.message.Message): + """User Context is used to refer to one particular user session that is executing + queries in the backend. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USER_ID_FIELD_NUMBER: builtins.int + USER_NAME_FIELD_NUMBER: builtins.int + EXTENSIONS_FIELD_NUMBER: builtins.int + user_id: builtins.str + user_name: builtins.str + @property + def extensions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + google.protobuf.any_pb2.Any + ]: + """To extend the existing user context message that is used to identify incoming requests, + Spark Connect leverages the Any protobuf type that can be used to inject arbitrary other + messages into this message. Extensions are stored as a `repeated` type to be able to + handle multiple active extensions. + """ + def __init__( + self, + *, + user_id: builtins.str = ..., + user_name: builtins.str = ..., + extensions: collections.abc.Iterable[google.protobuf.any_pb2.Any] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "extensions", b"extensions", "user_id", b"user_id", "user_name", b"user_name" + ], + ) -> None: ... + +global___UserContext = UserContext + +class AnalyzePlanRequest(google.protobuf.message.Message): + """Request to perform plan analyze, optionally to explain the plan.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Schema(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["plan", b"plan"]) -> None: ... + + class Explain(google.protobuf.message.Message): + """Explains the input plan based on a configurable mode.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _ExplainMode: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _ExplainModeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + AnalyzePlanRequest.Explain._ExplainMode.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + EXPLAIN_MODE_UNSPECIFIED: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 0 + EXPLAIN_MODE_SIMPLE: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 1 + """Generates only physical plan.""" + EXPLAIN_MODE_EXTENDED: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 2 + """Generates parsed logical plan, analyzed logical plan, optimized logical plan and physical plan. + Parsed Logical plan is a unresolved plan that extracted from the query. Analyzed logical plans + transforms which translates unresolvedAttribute and unresolvedRelation into fully typed objects. + The optimized logical plan transforms through a set of optimization rules, resulting in the + physical plan. + """ + EXPLAIN_MODE_CODEGEN: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 3 + """Generates code for the statement, if any and a physical plan.""" + EXPLAIN_MODE_COST: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 4 + """If plan node statistics are available, generates a logical plan and also the statistics.""" + EXPLAIN_MODE_FORMATTED: AnalyzePlanRequest.Explain._ExplainMode.ValueType # 5 + """Generates a physical plan outline and also node details.""" + + class ExplainMode(_ExplainMode, metaclass=_ExplainModeEnumTypeWrapper): + """Plan explanation mode.""" + + EXPLAIN_MODE_UNSPECIFIED: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 0 + EXPLAIN_MODE_SIMPLE: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 1 + """Generates only physical plan.""" + EXPLAIN_MODE_EXTENDED: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 2 + """Generates parsed logical plan, analyzed logical plan, optimized logical plan and physical plan. + Parsed Logical plan is a unresolved plan that extracted from the query. Analyzed logical plans + transforms which translates unresolvedAttribute and unresolvedRelation into fully typed objects. + The optimized logical plan transforms through a set of optimization rules, resulting in the + physical plan. + """ + EXPLAIN_MODE_CODEGEN: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 3 + """Generates code for the statement, if any and a physical plan.""" + EXPLAIN_MODE_COST: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 4 + """If plan node statistics are available, generates a logical plan and also the statistics.""" + EXPLAIN_MODE_FORMATTED: AnalyzePlanRequest.Explain.ExplainMode.ValueType # 5 + """Generates a physical plan outline and also node details.""" + + PLAN_FIELD_NUMBER: builtins.int + EXPLAIN_MODE_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + explain_mode: global___AnalyzePlanRequest.Explain.ExplainMode.ValueType + """(Required) For analyzePlan rpc calls, configure the mode to explain plan in strings.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + explain_mode: global___AnalyzePlanRequest.Explain.ExplainMode.ValueType = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["explain_mode", b"explain_mode", "plan", b"plan"], + ) -> None: ... + + class TreeString(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + LEVEL_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + level: builtins.int + """(Optional) Max level of the schema.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + level: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_level", b"_level", "level", b"level", "plan", b"plan" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_level", b"_level", "level", b"level", "plan", b"plan" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_level", b"_level"] + ) -> typing_extensions.Literal["level"] | None: ... + + class IsLocal(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["plan", b"plan"]) -> None: ... + + class IsStreaming(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["plan", b"plan"]) -> None: ... + + class InputFiles(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be analyzed.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["plan", b"plan"]) -> None: ... + + class SparkVersion(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + + class DDLParse(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DDL_STRING_FIELD_NUMBER: builtins.int + ddl_string: builtins.str + """(Required) The DDL formatted string to be parsed.""" + def __init__( + self, + *, + ddl_string: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["ddl_string", b"ddl_string"] + ) -> None: ... + + class SameSemantics(google.protobuf.message.Message): + """Returns `true` when the logical query plans are equal and therefore return same results.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TARGET_PLAN_FIELD_NUMBER: builtins.int + OTHER_PLAN_FIELD_NUMBER: builtins.int + @property + def target_plan(self) -> global___Plan: + """(Required) The plan to be compared.""" + @property + def other_plan(self) -> global___Plan: + """(Required) The other plan to be compared.""" + def __init__( + self, + *, + target_plan: global___Plan | None = ..., + other_plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "other_plan", b"other_plan", "target_plan", b"target_plan" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "other_plan", b"other_plan", "target_plan", b"target_plan" + ], + ) -> None: ... + + class SemanticHash(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PLAN_FIELD_NUMBER: builtins.int + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to get a hashCode.""" + def __init__( + self, + *, + plan: global___Plan | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["plan", b"plan"] + ) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["plan", b"plan"]) -> None: ... + + class Persist(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RELATION_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + @property + def relation(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The logical plan to persist.""" + @property + def storage_level(self) -> pyspark.sql.connect.proto.common_pb2.StorageLevel: + """(Optional) The storage level.""" + def __init__( + self, + *, + relation: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + storage_level: pyspark.sql.connect.proto.common_pb2.StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "relation", + b"relation", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "relation", + b"relation", + "storage_level", + b"storage_level", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + + class Unpersist(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RELATION_FIELD_NUMBER: builtins.int + BLOCKING_FIELD_NUMBER: builtins.int + @property + def relation(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The logical plan to unpersist.""" + blocking: builtins.bool + """(Optional) Whether to block until all blocks are deleted.""" + def __init__( + self, + *, + relation: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + blocking: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_blocking", b"_blocking", "blocking", b"blocking", "relation", b"relation" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_blocking", b"_blocking", "blocking", b"blocking", "relation", b"relation" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_blocking", b"_blocking"] + ) -> typing_extensions.Literal["blocking"] | None: ... + + class GetStorageLevel(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RELATION_FIELD_NUMBER: builtins.int + @property + def relation(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The logical plan to get the storage level.""" + def __init__( + self, + *, + relation: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["relation", b"relation"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["relation", b"relation"] + ) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + EXPLAIN_FIELD_NUMBER: builtins.int + TREE_STRING_FIELD_NUMBER: builtins.int + IS_LOCAL_FIELD_NUMBER: builtins.int + IS_STREAMING_FIELD_NUMBER: builtins.int + INPUT_FILES_FIELD_NUMBER: builtins.int + SPARK_VERSION_FIELD_NUMBER: builtins.int + DDL_PARSE_FIELD_NUMBER: builtins.int + SAME_SEMANTICS_FIELD_NUMBER: builtins.int + SEMANTIC_HASH_FIELD_NUMBER: builtins.int + PERSIST_FIELD_NUMBER: builtins.int + UNPERSIST_FIELD_NUMBER: builtins.int + GET_STORAGE_LEVEL_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + @property + def schema(self) -> global___AnalyzePlanRequest.Schema: ... + @property + def explain(self) -> global___AnalyzePlanRequest.Explain: ... + @property + def tree_string(self) -> global___AnalyzePlanRequest.TreeString: ... + @property + def is_local(self) -> global___AnalyzePlanRequest.IsLocal: ... + @property + def is_streaming(self) -> global___AnalyzePlanRequest.IsStreaming: ... + @property + def input_files(self) -> global___AnalyzePlanRequest.InputFiles: ... + @property + def spark_version(self) -> global___AnalyzePlanRequest.SparkVersion: ... + @property + def ddl_parse(self) -> global___AnalyzePlanRequest.DDLParse: ... + @property + def same_semantics(self) -> global___AnalyzePlanRequest.SameSemantics: ... + @property + def semantic_hash(self) -> global___AnalyzePlanRequest.SemanticHash: ... + @property + def persist(self) -> global___AnalyzePlanRequest.Persist: ... + @property + def unpersist(self) -> global___AnalyzePlanRequest.Unpersist: ... + @property + def get_storage_level(self) -> global___AnalyzePlanRequest.GetStorageLevel: ... + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + client_type: builtins.str | None = ..., + schema: global___AnalyzePlanRequest.Schema | None = ..., + explain: global___AnalyzePlanRequest.Explain | None = ..., + tree_string: global___AnalyzePlanRequest.TreeString | None = ..., + is_local: global___AnalyzePlanRequest.IsLocal | None = ..., + is_streaming: global___AnalyzePlanRequest.IsStreaming | None = ..., + input_files: global___AnalyzePlanRequest.InputFiles | None = ..., + spark_version: global___AnalyzePlanRequest.SparkVersion | None = ..., + ddl_parse: global___AnalyzePlanRequest.DDLParse | None = ..., + same_semantics: global___AnalyzePlanRequest.SameSemantics | None = ..., + semantic_hash: global___AnalyzePlanRequest.SemanticHash | None = ..., + persist: global___AnalyzePlanRequest.Persist | None = ..., + unpersist: global___AnalyzePlanRequest.Unpersist | None = ..., + get_storage_level: global___AnalyzePlanRequest.GetStorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "analyze", + b"analyze", + "client_type", + b"client_type", + "ddl_parse", + b"ddl_parse", + "explain", + b"explain", + "get_storage_level", + b"get_storage_level", + "input_files", + b"input_files", + "is_local", + b"is_local", + "is_streaming", + b"is_streaming", + "persist", + b"persist", + "same_semantics", + b"same_semantics", + "schema", + b"schema", + "semantic_hash", + b"semantic_hash", + "spark_version", + b"spark_version", + "tree_string", + b"tree_string", + "unpersist", + b"unpersist", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "analyze", + b"analyze", + "client_type", + b"client_type", + "ddl_parse", + b"ddl_parse", + "explain", + b"explain", + "get_storage_level", + b"get_storage_level", + "input_files", + b"input_files", + "is_local", + b"is_local", + "is_streaming", + b"is_streaming", + "persist", + b"persist", + "same_semantics", + b"same_semantics", + "schema", + b"schema", + "semantic_hash", + b"semantic_hash", + "session_id", + b"session_id", + "spark_version", + b"spark_version", + "tree_string", + b"tree_string", + "unpersist", + b"unpersist", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["analyze", b"analyze"] + ) -> typing_extensions.Literal[ + "schema", + "explain", + "tree_string", + "is_local", + "is_streaming", + "input_files", + "spark_version", + "ddl_parse", + "same_semantics", + "semantic_hash", + "persist", + "unpersist", + "get_storage_level", + ] | None: ... + +global___AnalyzePlanRequest = AnalyzePlanRequest + +class AnalyzePlanResponse(google.protobuf.message.Message): + """Response to performing analysis of the query. Contains relevant metadata to be able to + reason about the performance. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Schema(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SCHEMA_FIELD_NUMBER: builtins.int + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + def __init__( + self, + *, + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["schema", b"schema"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["schema", b"schema"] + ) -> None: ... + + class Explain(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXPLAIN_STRING_FIELD_NUMBER: builtins.int + explain_string: builtins.str + def __init__( + self, + *, + explain_string: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["explain_string", b"explain_string"] + ) -> None: ... + + class TreeString(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TREE_STRING_FIELD_NUMBER: builtins.int + tree_string: builtins.str + def __init__( + self, + *, + tree_string: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["tree_string", b"tree_string"] + ) -> None: ... + + class IsLocal(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + IS_LOCAL_FIELD_NUMBER: builtins.int + is_local: builtins.bool + def __init__( + self, + *, + is_local: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["is_local", b"is_local"] + ) -> None: ... + + class IsStreaming(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + IS_STREAMING_FIELD_NUMBER: builtins.int + is_streaming: builtins.bool + def __init__( + self, + *, + is_streaming: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["is_streaming", b"is_streaming"] + ) -> None: ... + + class InputFiles(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FILES_FIELD_NUMBER: builtins.int + @property + def files( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """A best-effort snapshot of the files that compose this Dataset""" + def __init__( + self, + *, + files: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["files", b"files"]) -> None: ... + + class SparkVersion(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VERSION_FIELD_NUMBER: builtins.int + version: builtins.str + def __init__( + self, + *, + version: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["version", b"version"] + ) -> None: ... + + class DDLParse(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PARSED_FIELD_NUMBER: builtins.int + @property + def parsed(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + def __init__( + self, + *, + parsed: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["parsed", b"parsed"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["parsed", b"parsed"] + ) -> None: ... + + class SameSemantics(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESULT_FIELD_NUMBER: builtins.int + result: builtins.bool + def __init__( + self, + *, + result: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["result", b"result"] + ) -> None: ... + + class SemanticHash(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESULT_FIELD_NUMBER: builtins.int + result: builtins.int + def __init__( + self, + *, + result: builtins.int = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["result", b"result"] + ) -> None: ... + + class Persist(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + + class Unpersist(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + + class GetStorageLevel(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + @property + def storage_level(self) -> pyspark.sql.connect.proto.common_pb2.StorageLevel: + """(Required) The StorageLevel as a result of get_storage_level request.""" + def __init__( + self, + *, + storage_level: pyspark.sql.connect.proto.common_pb2.StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["storage_level", b"storage_level"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["storage_level", b"storage_level"] + ) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + EXPLAIN_FIELD_NUMBER: builtins.int + TREE_STRING_FIELD_NUMBER: builtins.int + IS_LOCAL_FIELD_NUMBER: builtins.int + IS_STREAMING_FIELD_NUMBER: builtins.int + INPUT_FILES_FIELD_NUMBER: builtins.int + SPARK_VERSION_FIELD_NUMBER: builtins.int + DDL_PARSE_FIELD_NUMBER: builtins.int + SAME_SEMANTICS_FIELD_NUMBER: builtins.int + SEMANTIC_HASH_FIELD_NUMBER: builtins.int + PERSIST_FIELD_NUMBER: builtins.int + UNPERSIST_FIELD_NUMBER: builtins.int + GET_STORAGE_LEVEL_FIELD_NUMBER: builtins.int + session_id: builtins.str + @property + def schema(self) -> global___AnalyzePlanResponse.Schema: ... + @property + def explain(self) -> global___AnalyzePlanResponse.Explain: ... + @property + def tree_string(self) -> global___AnalyzePlanResponse.TreeString: ... + @property + def is_local(self) -> global___AnalyzePlanResponse.IsLocal: ... + @property + def is_streaming(self) -> global___AnalyzePlanResponse.IsStreaming: ... + @property + def input_files(self) -> global___AnalyzePlanResponse.InputFiles: ... + @property + def spark_version(self) -> global___AnalyzePlanResponse.SparkVersion: ... + @property + def ddl_parse(self) -> global___AnalyzePlanResponse.DDLParse: ... + @property + def same_semantics(self) -> global___AnalyzePlanResponse.SameSemantics: ... + @property + def semantic_hash(self) -> global___AnalyzePlanResponse.SemanticHash: ... + @property + def persist(self) -> global___AnalyzePlanResponse.Persist: ... + @property + def unpersist(self) -> global___AnalyzePlanResponse.Unpersist: ... + @property + def get_storage_level(self) -> global___AnalyzePlanResponse.GetStorageLevel: ... + def __init__( + self, + *, + session_id: builtins.str = ..., + schema: global___AnalyzePlanResponse.Schema | None = ..., + explain: global___AnalyzePlanResponse.Explain | None = ..., + tree_string: global___AnalyzePlanResponse.TreeString | None = ..., + is_local: global___AnalyzePlanResponse.IsLocal | None = ..., + is_streaming: global___AnalyzePlanResponse.IsStreaming | None = ..., + input_files: global___AnalyzePlanResponse.InputFiles | None = ..., + spark_version: global___AnalyzePlanResponse.SparkVersion | None = ..., + ddl_parse: global___AnalyzePlanResponse.DDLParse | None = ..., + same_semantics: global___AnalyzePlanResponse.SameSemantics | None = ..., + semantic_hash: global___AnalyzePlanResponse.SemanticHash | None = ..., + persist: global___AnalyzePlanResponse.Persist | None = ..., + unpersist: global___AnalyzePlanResponse.Unpersist | None = ..., + get_storage_level: global___AnalyzePlanResponse.GetStorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "ddl_parse", + b"ddl_parse", + "explain", + b"explain", + "get_storage_level", + b"get_storage_level", + "input_files", + b"input_files", + "is_local", + b"is_local", + "is_streaming", + b"is_streaming", + "persist", + b"persist", + "result", + b"result", + "same_semantics", + b"same_semantics", + "schema", + b"schema", + "semantic_hash", + b"semantic_hash", + "spark_version", + b"spark_version", + "tree_string", + b"tree_string", + "unpersist", + b"unpersist", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "ddl_parse", + b"ddl_parse", + "explain", + b"explain", + "get_storage_level", + b"get_storage_level", + "input_files", + b"input_files", + "is_local", + b"is_local", + "is_streaming", + b"is_streaming", + "persist", + b"persist", + "result", + b"result", + "same_semantics", + b"same_semantics", + "schema", + b"schema", + "semantic_hash", + b"semantic_hash", + "session_id", + b"session_id", + "spark_version", + b"spark_version", + "tree_string", + b"tree_string", + "unpersist", + b"unpersist", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["result", b"result"] + ) -> typing_extensions.Literal[ + "schema", + "explain", + "tree_string", + "is_local", + "is_streaming", + "input_files", + "spark_version", + "ddl_parse", + "same_semantics", + "semantic_hash", + "persist", + "unpersist", + "get_storage_level", + ] | None: ... + +global___AnalyzePlanResponse = AnalyzePlanResponse + +class ExecutePlanRequest(google.protobuf.message.Message): + """A request to be executed by the service.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class RequestOption(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + REATTACH_OPTIONS_FIELD_NUMBER: builtins.int + EXTENSION_FIELD_NUMBER: builtins.int + @property + def reattach_options(self) -> global___ReattachOptions: ... + @property + def extension(self) -> google.protobuf.any_pb2.Any: + """Extension type for request options""" + def __init__( + self, + *, + reattach_options: global___ReattachOptions | None = ..., + extension: google.protobuf.any_pb2.Any | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "extension", + b"extension", + "reattach_options", + b"reattach_options", + "request_option", + b"request_option", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "extension", + b"extension", + "reattach_options", + b"reattach_options", + "request_option", + b"request_option", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["request_option", b"request_option"] + ) -> typing_extensions.Literal["reattach_options", "extension"] | None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + PLAN_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + REQUEST_OPTIONS_FIELD_NUMBER: builtins.int + TAGS_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context + + user_context.user_id and session+id both identify a unique remote spark session on the + server side. + """ + operation_id: builtins.str + """(Optional) + Provide an id for this request. If not provided, it will be generated by the server. + It is returned in every ExecutePlanResponse.operation_id of the ExecutePlan response stream. + The id must be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def plan(self) -> global___Plan: + """(Required) The logical plan to be executed / analyzed.""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + @property + def request_options( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ExecutePlanRequest.RequestOption + ]: + """Repeated element for options that can be passed to the request. This element is currently + unused but allows to pass in an extension value used for arbitrary options. + """ + @property + def tags( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Tags to tag the given execution with. + Tags cannot contain ',' character and cannot be empty strings. + Used by Interrupt with interrupt.tag. + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + operation_id: builtins.str | None = ..., + plan: global___Plan | None = ..., + client_type: builtins.str | None = ..., + request_options: collections.abc.Iterable[global___ExecutePlanRequest.RequestOption] + | None = ..., + tags: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "_operation_id", + b"_operation_id", + "client_type", + b"client_type", + "operation_id", + b"operation_id", + "plan", + b"plan", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "_operation_id", + b"_operation_id", + "client_type", + b"client_type", + "operation_id", + b"operation_id", + "plan", + b"plan", + "request_options", + b"request_options", + "session_id", + b"session_id", + "tags", + b"tags", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_operation_id", b"_operation_id"] + ) -> typing_extensions.Literal["operation_id"] | None: ... + +global___ExecutePlanRequest = ExecutePlanRequest + +class ExecutePlanResponse(google.protobuf.message.Message): + """The response of a query, can be one or more for each request. Responses belonging to the + same input query, carry the same `session_id`. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class SqlCommandResult(google.protobuf.message.Message): + """A SQL command returns an opaque Relation that can be directly used as input for the next + call. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RELATION_FIELD_NUMBER: builtins.int + @property + def relation(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: ... + def __init__( + self, + *, + relation: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["relation", b"relation"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["relation", b"relation"] + ) -> None: ... + + class ArrowBatch(google.protobuf.message.Message): + """Batch results of metrics.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ROW_COUNT_FIELD_NUMBER: builtins.int + DATA_FIELD_NUMBER: builtins.int + row_count: builtins.int + data: builtins.bytes + def __init__( + self, + *, + row_count: builtins.int = ..., + data: builtins.bytes = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["data", b"data", "row_count", b"row_count"] + ) -> None: ... + + class Metrics(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class MetricObject(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ExecutionMetricsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + @property + def value(self) -> global___ExecutePlanResponse.Metrics.MetricValue: ... + def __init__( + self, + *, + key: builtins.str = ..., + value: global___ExecutePlanResponse.Metrics.MetricValue | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + NAME_FIELD_NUMBER: builtins.int + PLAN_ID_FIELD_NUMBER: builtins.int + PARENT_FIELD_NUMBER: builtins.int + EXECUTION_METRICS_FIELD_NUMBER: builtins.int + name: builtins.str + plan_id: builtins.int + parent: builtins.int + @property + def execution_metrics( + self, + ) -> google.protobuf.internal.containers.MessageMap[ + builtins.str, global___ExecutePlanResponse.Metrics.MetricValue + ]: ... + def __init__( + self, + *, + name: builtins.str = ..., + plan_id: builtins.int = ..., + parent: builtins.int = ..., + execution_metrics: collections.abc.Mapping[ + builtins.str, global___ExecutePlanResponse.Metrics.MetricValue + ] + | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "execution_metrics", + b"execution_metrics", + "name", + b"name", + "parent", + b"parent", + "plan_id", + b"plan_id", + ], + ) -> None: ... + + class MetricValue(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + METRIC_TYPE_FIELD_NUMBER: builtins.int + name: builtins.str + value: builtins.int + metric_type: builtins.str + def __init__( + self, + *, + name: builtins.str = ..., + value: builtins.int = ..., + metric_type: builtins.str = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "metric_type", b"metric_type", "name", b"name", "value", b"value" + ], + ) -> None: ... + + METRICS_FIELD_NUMBER: builtins.int + @property + def metrics( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ExecutePlanResponse.Metrics.MetricObject + ]: ... + def __init__( + self, + *, + metrics: collections.abc.Iterable[global___ExecutePlanResponse.Metrics.MetricObject] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["metrics", b"metrics"] + ) -> None: ... + + class ObservedMetrics(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + VALUES_FIELD_NUMBER: builtins.int + name: builtins.str + @property + def values( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: ... + def __init__( + self, + *, + name: builtins.str = ..., + values: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["name", b"name", "values", b"values"] + ) -> None: ... + + class ResultComplete(google.protobuf.message.Message): + """If present, in a reattachable execution this means that after server sends onComplete, + the execution is complete. If the server sends onComplete without sending a ResultComplete, + it means that there is more, and the client should use ReattachExecute RPC to continue. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + RESPONSE_ID_FIELD_NUMBER: builtins.int + ARROW_BATCH_FIELD_NUMBER: builtins.int + SQL_COMMAND_RESULT_FIELD_NUMBER: builtins.int + WRITE_STREAM_OPERATION_START_RESULT_FIELD_NUMBER: builtins.int + STREAMING_QUERY_COMMAND_RESULT_FIELD_NUMBER: builtins.int + GET_RESOURCES_COMMAND_RESULT_FIELD_NUMBER: builtins.int + STREAMING_QUERY_MANAGER_COMMAND_RESULT_FIELD_NUMBER: builtins.int + RESULT_COMPLETE_FIELD_NUMBER: builtins.int + EXTENSION_FIELD_NUMBER: builtins.int + METRICS_FIELD_NUMBER: builtins.int + OBSERVED_METRICS_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + session_id: builtins.str + operation_id: builtins.str + """Identifies the ExecutePlan execution. + If set by the client in ExecutePlanRequest.operationId, that value is returned. + Otherwise generated by the server. + It is an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + response_id: builtins.str + """Identified the response in the stream. + The id is an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def arrow_batch(self) -> global___ExecutePlanResponse.ArrowBatch: ... + @property + def sql_command_result(self) -> global___ExecutePlanResponse.SqlCommandResult: + """Special case for executing SQL commands.""" + @property + def write_stream_operation_start_result( + self, + ) -> pyspark.sql.connect.proto.commands_pb2.WriteStreamOperationStartResult: + """Response for a streaming query.""" + @property + def streaming_query_command_result( + self, + ) -> pyspark.sql.connect.proto.commands_pb2.StreamingQueryCommandResult: + """Response for commands on a streaming query.""" + @property + def get_resources_command_result( + self, + ) -> pyspark.sql.connect.proto.commands_pb2.GetResourcesCommandResult: + """Response for 'SparkContext.resources'.""" + @property + def streaming_query_manager_command_result( + self, + ) -> pyspark.sql.connect.proto.commands_pb2.StreamingQueryManagerCommandResult: + """Response for commands on the streaming query manager.""" + @property + def result_complete(self) -> global___ExecutePlanResponse.ResultComplete: + """Response type informing if the stream is complete in reattachable execution.""" + @property + def extension(self) -> google.protobuf.any_pb2.Any: + """Support arbitrary result objects.""" + @property + def metrics(self) -> global___ExecutePlanResponse.Metrics: + """Metrics for the query execution. Typically, this field is only present in the last + batch of results and then represent the overall state of the query execution. + """ + @property + def observed_metrics( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ExecutePlanResponse.ObservedMetrics + ]: + """The metrics observed during the execution of the query plan.""" + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional) The Spark schema. This field is available when `collect` is called.""" + def __init__( + self, + *, + session_id: builtins.str = ..., + operation_id: builtins.str = ..., + response_id: builtins.str = ..., + arrow_batch: global___ExecutePlanResponse.ArrowBatch | None = ..., + sql_command_result: global___ExecutePlanResponse.SqlCommandResult | None = ..., + write_stream_operation_start_result: pyspark.sql.connect.proto.commands_pb2.WriteStreamOperationStartResult + | None = ..., + streaming_query_command_result: pyspark.sql.connect.proto.commands_pb2.StreamingQueryCommandResult + | None = ..., + get_resources_command_result: pyspark.sql.connect.proto.commands_pb2.GetResourcesCommandResult + | None = ..., + streaming_query_manager_command_result: pyspark.sql.connect.proto.commands_pb2.StreamingQueryManagerCommandResult + | None = ..., + result_complete: global___ExecutePlanResponse.ResultComplete | None = ..., + extension: google.protobuf.any_pb2.Any | None = ..., + metrics: global___ExecutePlanResponse.Metrics | None = ..., + observed_metrics: collections.abc.Iterable[global___ExecutePlanResponse.ObservedMetrics] + | None = ..., + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "arrow_batch", + b"arrow_batch", + "extension", + b"extension", + "get_resources_command_result", + b"get_resources_command_result", + "metrics", + b"metrics", + "response_type", + b"response_type", + "result_complete", + b"result_complete", + "schema", + b"schema", + "sql_command_result", + b"sql_command_result", + "streaming_query_command_result", + b"streaming_query_command_result", + "streaming_query_manager_command_result", + b"streaming_query_manager_command_result", + "write_stream_operation_start_result", + b"write_stream_operation_start_result", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arrow_batch", + b"arrow_batch", + "extension", + b"extension", + "get_resources_command_result", + b"get_resources_command_result", + "metrics", + b"metrics", + "observed_metrics", + b"observed_metrics", + "operation_id", + b"operation_id", + "response_id", + b"response_id", + "response_type", + b"response_type", + "result_complete", + b"result_complete", + "schema", + b"schema", + "session_id", + b"session_id", + "sql_command_result", + b"sql_command_result", + "streaming_query_command_result", + b"streaming_query_command_result", + "streaming_query_manager_command_result", + b"streaming_query_manager_command_result", + "write_stream_operation_start_result", + b"write_stream_operation_start_result", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["response_type", b"response_type"] + ) -> typing_extensions.Literal[ + "arrow_batch", + "sql_command_result", + "write_stream_operation_start_result", + "streaming_query_command_result", + "get_resources_command_result", + "streaming_query_manager_command_result", + "result_complete", + "extension", + ] | None: ... + +global___ExecutePlanResponse = ExecutePlanResponse + +class KeyValue(google.protobuf.message.Message): + """The key-value pair for the config request and response.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + """(Required) The key.""" + value: builtins.str + """(Optional) The value.""" + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_value", b"_value", "value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_value", b"_value", "key", b"key", "value", b"value" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_value", b"_value"] + ) -> typing_extensions.Literal["value"] | None: ... + +global___KeyValue = KeyValue + +class ConfigRequest(google.protobuf.message.Message): + """Request to update or fetch the configurations.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Operation(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SET_FIELD_NUMBER: builtins.int + GET_FIELD_NUMBER: builtins.int + GET_WITH_DEFAULT_FIELD_NUMBER: builtins.int + GET_OPTION_FIELD_NUMBER: builtins.int + GET_ALL_FIELD_NUMBER: builtins.int + UNSET_FIELD_NUMBER: builtins.int + IS_MODIFIABLE_FIELD_NUMBER: builtins.int + @property + def set(self) -> global___ConfigRequest.Set: ... + @property + def get(self) -> global___ConfigRequest.Get: ... + @property + def get_with_default(self) -> global___ConfigRequest.GetWithDefault: ... + @property + def get_option(self) -> global___ConfigRequest.GetOption: ... + @property + def get_all(self) -> global___ConfigRequest.GetAll: ... + @property + def unset(self) -> global___ConfigRequest.Unset: ... + @property + def is_modifiable(self) -> global___ConfigRequest.IsModifiable: ... + def __init__( + self, + *, + set: global___ConfigRequest.Set | None = ..., + get: global___ConfigRequest.Get | None = ..., + get_with_default: global___ConfigRequest.GetWithDefault | None = ..., + get_option: global___ConfigRequest.GetOption | None = ..., + get_all: global___ConfigRequest.GetAll | None = ..., + unset: global___ConfigRequest.Unset | None = ..., + is_modifiable: global___ConfigRequest.IsModifiable | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "get", + b"get", + "get_all", + b"get_all", + "get_option", + b"get_option", + "get_with_default", + b"get_with_default", + "is_modifiable", + b"is_modifiable", + "op_type", + b"op_type", + "set", + b"set", + "unset", + b"unset", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "get", + b"get", + "get_all", + b"get_all", + "get_option", + b"get_option", + "get_with_default", + b"get_with_default", + "is_modifiable", + b"is_modifiable", + "op_type", + b"op_type", + "set", + b"set", + "unset", + b"unset", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["op_type", b"op_type"] + ) -> typing_extensions.Literal[ + "set", "get", "get_with_default", "get_option", "get_all", "unset", "is_modifiable" + ] | None: ... + + class Set(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PAIRS_FIELD_NUMBER: builtins.int + @property + def pairs( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___KeyValue]: + """(Required) The config key-value pairs to set.""" + def __init__( + self, + *, + pairs: collections.abc.Iterable[global___KeyValue] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["pairs", b"pairs"]) -> None: ... + + class Get(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEYS_FIELD_NUMBER: builtins.int + @property + def keys( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The config keys to get.""" + def __init__( + self, + *, + keys: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["keys", b"keys"]) -> None: ... + + class GetWithDefault(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PAIRS_FIELD_NUMBER: builtins.int + @property + def pairs( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___KeyValue]: + """(Required) The config key-value paris to get. The value will be used as the default value.""" + def __init__( + self, + *, + pairs: collections.abc.Iterable[global___KeyValue] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["pairs", b"pairs"]) -> None: ... + + class GetOption(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEYS_FIELD_NUMBER: builtins.int + @property + def keys( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The config keys to get optionally.""" + def __init__( + self, + *, + keys: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["keys", b"keys"]) -> None: ... + + class GetAll(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PREFIX_FIELD_NUMBER: builtins.int + prefix: builtins.str + """(Optional) The prefix of the config key to get.""" + def __init__( + self, + *, + prefix: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_prefix", b"_prefix", "prefix", b"prefix"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["_prefix", b"_prefix", "prefix", b"prefix"] + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_prefix", b"_prefix"] + ) -> typing_extensions.Literal["prefix"] | None: ... + + class Unset(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEYS_FIELD_NUMBER: builtins.int + @property + def keys( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The config keys to unset.""" + def __init__( + self, + *, + keys: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["keys", b"keys"]) -> None: ... + + class IsModifiable(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEYS_FIELD_NUMBER: builtins.int + @property + def keys( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The config keys to check the config is modifiable.""" + def __init__( + self, + *, + keys: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["keys", b"keys"]) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + OPERATION_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context""" + @property + def operation(self) -> global___ConfigRequest.Operation: + """(Required) The operation for the config.""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + operation: global___ConfigRequest.Operation | None = ..., + client_type: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "operation", + b"operation", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "operation", + b"operation", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + +global___ConfigRequest = ConfigRequest + +class ConfigResponse(google.protobuf.message.Message): + """Response to the config request.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SESSION_ID_FIELD_NUMBER: builtins.int + PAIRS_FIELD_NUMBER: builtins.int + WARNINGS_FIELD_NUMBER: builtins.int + session_id: builtins.str + @property + def pairs( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___KeyValue]: + """(Optional) The result key-value pairs. + + Available when the operation is 'Get', 'GetWithDefault', 'GetOption', 'GetAll'. + Also available for the operation 'IsModifiable' with boolean string "true" and "false". + """ + @property + def warnings( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) + + Warning messages for deprecated or unsupported configurations. + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + pairs: collections.abc.Iterable[global___KeyValue] | None = ..., + warnings: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "pairs", b"pairs", "session_id", b"session_id", "warnings", b"warnings" + ], + ) -> None: ... + +global___ConfigResponse = ConfigResponse + +class AddArtifactsRequest(google.protobuf.message.Message): + """Request to transfer client-local artifacts.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ArtifactChunk(google.protobuf.message.Message): + """A chunk of an Artifact.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DATA_FIELD_NUMBER: builtins.int + CRC_FIELD_NUMBER: builtins.int + data: builtins.bytes + """Data chunk.""" + crc: builtins.int + """CRC to allow server to verify integrity of the chunk.""" + def __init__( + self, + *, + data: builtins.bytes = ..., + crc: builtins.int = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["crc", b"crc", "data", b"data"] + ) -> None: ... + + class SingleChunkArtifact(google.protobuf.message.Message): + """An artifact that is contained in a single `ArtifactChunk`. + Generally, this message represents tiny artifacts such as REPL-generated class files. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + DATA_FIELD_NUMBER: builtins.int + name: builtins.str + """The name of the artifact is expected in the form of a "Relative Path" that is made up of a + sequence of directories and the final file element. + Examples of "Relative Path"s: "jars/test.jar", "classes/xyz.class", "abc.xyz", "a/b/X.jar". + The server is expected to maintain the hierarchy of files as defined by their name. (i.e + The relative path of the file on the server's filesystem will be the same as the name of + the provided artifact) + """ + @property + def data(self) -> global___AddArtifactsRequest.ArtifactChunk: + """A single data chunk.""" + def __init__( + self, + *, + name: builtins.str = ..., + data: global___AddArtifactsRequest.ArtifactChunk | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["data", b"data"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["data", b"data", "name", b"name"] + ) -> None: ... + + class Batch(google.protobuf.message.Message): + """A number of `SingleChunkArtifact` batched into a single RPC.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ARTIFACTS_FIELD_NUMBER: builtins.int + @property + def artifacts( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___AddArtifactsRequest.SingleChunkArtifact + ]: ... + def __init__( + self, + *, + artifacts: collections.abc.Iterable[global___AddArtifactsRequest.SingleChunkArtifact] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["artifacts", b"artifacts"] + ) -> None: ... + + class BeginChunkedArtifact(google.protobuf.message.Message): + """Signals the beginning/start of a chunked artifact. + A large artifact is transferred through a payload of `BeginChunkedArtifact` followed by a + sequence of `ArtifactChunk`s. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + TOTAL_BYTES_FIELD_NUMBER: builtins.int + NUM_CHUNKS_FIELD_NUMBER: builtins.int + INITIAL_CHUNK_FIELD_NUMBER: builtins.int + name: builtins.str + """Name of the artifact undergoing chunking. Follows the same conventions as the `name` in + the `Artifact` message. + """ + total_bytes: builtins.int + """Total size of the artifact in bytes.""" + num_chunks: builtins.int + """Number of chunks the artifact is split into. + This includes the `initial_chunk`. + """ + @property + def initial_chunk(self) -> global___AddArtifactsRequest.ArtifactChunk: + """The first/initial chunk.""" + def __init__( + self, + *, + name: builtins.str = ..., + total_bytes: builtins.int = ..., + num_chunks: builtins.int = ..., + initial_chunk: global___AddArtifactsRequest.ArtifactChunk | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["initial_chunk", b"initial_chunk"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "initial_chunk", + b"initial_chunk", + "name", + b"name", + "num_chunks", + b"num_chunks", + "total_bytes", + b"total_bytes", + ], + ) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + BATCH_FIELD_NUMBER: builtins.int + BEGIN_CHUNK_FIELD_NUMBER: builtins.int + CHUNK_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """User context""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + @property + def batch(self) -> global___AddArtifactsRequest.Batch: ... + @property + def begin_chunk(self) -> global___AddArtifactsRequest.BeginChunkedArtifact: + """The metadata and the initial chunk of a large artifact chunked into multiple requests. + The server side is notified about the total size of the large artifact as well as the + number of chunks to expect. + """ + @property + def chunk(self) -> global___AddArtifactsRequest.ArtifactChunk: + """A chunk of an artifact excluding metadata. This can be any chunk of a large artifact + excluding the first chunk (which is included in `BeginChunkedArtifact`). + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + client_type: builtins.str | None = ..., + batch: global___AddArtifactsRequest.Batch | None = ..., + begin_chunk: global___AddArtifactsRequest.BeginChunkedArtifact | None = ..., + chunk: global___AddArtifactsRequest.ArtifactChunk | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "batch", + b"batch", + "begin_chunk", + b"begin_chunk", + "chunk", + b"chunk", + "client_type", + b"client_type", + "payload", + b"payload", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "batch", + b"batch", + "begin_chunk", + b"begin_chunk", + "chunk", + b"chunk", + "client_type", + b"client_type", + "payload", + b"payload", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["payload", b"payload"] + ) -> typing_extensions.Literal["batch", "begin_chunk", "chunk"] | None: ... + +global___AddArtifactsRequest = AddArtifactsRequest + +class AddArtifactsResponse(google.protobuf.message.Message): + """Response to adding an artifact. Contains relevant metadata to verify successful transfer of + artifact(s). + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ArtifactSummary(google.protobuf.message.Message): + """Metadata of an artifact.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + IS_CRC_SUCCESSFUL_FIELD_NUMBER: builtins.int + name: builtins.str + is_crc_successful: builtins.bool + """Whether the CRC (Cyclic Redundancy Check) is successful on server verification. + The server discards any artifact that fails the CRC. + If false, the client may choose to resend the artifact specified by `name`. + """ + def __init__( + self, + *, + name: builtins.str = ..., + is_crc_successful: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "is_crc_successful", b"is_crc_successful", "name", b"name" + ], + ) -> None: ... + + ARTIFACTS_FIELD_NUMBER: builtins.int + @property + def artifacts( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___AddArtifactsResponse.ArtifactSummary + ]: + """The list of artifact(s) seen by the server.""" + def __init__( + self, + *, + artifacts: collections.abc.Iterable[global___AddArtifactsResponse.ArtifactSummary] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["artifacts", b"artifacts"] + ) -> None: ... + +global___AddArtifactsResponse = AddArtifactsResponse + +class ArtifactStatusesRequest(google.protobuf.message.Message): + """Request to get current statuses of artifacts at the server side.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + NAMES_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """User context""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + @property + def names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """The name of the artifact is expected in the form of a "Relative Path" that is made up of a + sequence of directories and the final file element. + Examples of "Relative Path"s: "jars/test.jar", "classes/xyz.class", "abc.xyz", "a/b/X.jar". + The server is expected to maintain the hierarchy of files as defined by their name. (i.e + The relative path of the file on the server's filesystem will be the same as the name of + the provided artifact) + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + client_type: builtins.str | None = ..., + names: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "names", + b"names", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + +global___ArtifactStatusesRequest = ArtifactStatusesRequest + +class ArtifactStatusesResponse(google.protobuf.message.Message): + """Response to checking artifact statuses.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ArtifactStatus(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXISTS_FIELD_NUMBER: builtins.int + exists: builtins.bool + """Exists or not particular artifact at the server.""" + def __init__( + self, + *, + exists: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["exists", b"exists"] + ) -> None: ... + + class StatusesEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + @property + def value(self) -> global___ArtifactStatusesResponse.ArtifactStatus: ... + def __init__( + self, + *, + key: builtins.str = ..., + value: global___ArtifactStatusesResponse.ArtifactStatus | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + STATUSES_FIELD_NUMBER: builtins.int + @property + def statuses( + self, + ) -> google.protobuf.internal.containers.MessageMap[ + builtins.str, global___ArtifactStatusesResponse.ArtifactStatus + ]: + """A map of artifact names to their statuses.""" + def __init__( + self, + *, + statuses: collections.abc.Mapping[ + builtins.str, global___ArtifactStatusesResponse.ArtifactStatus + ] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["statuses", b"statuses"] + ) -> None: ... + +global___ArtifactStatusesResponse = ArtifactStatusesResponse + +class InterruptRequest(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _InterruptType: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _InterruptTypeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + InterruptRequest._InterruptType.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + INTERRUPT_TYPE_UNSPECIFIED: InterruptRequest._InterruptType.ValueType # 0 + INTERRUPT_TYPE_ALL: InterruptRequest._InterruptType.ValueType # 1 + """Interrupt all running executions within the session with the provided session_id.""" + INTERRUPT_TYPE_TAG: InterruptRequest._InterruptType.ValueType # 2 + """Interrupt all running executions within the session with the provided operation_tag.""" + INTERRUPT_TYPE_OPERATION_ID: InterruptRequest._InterruptType.ValueType # 3 + """Interrupt the running execution within the session with the provided operation_id.""" + + class InterruptType(_InterruptType, metaclass=_InterruptTypeEnumTypeWrapper): ... + INTERRUPT_TYPE_UNSPECIFIED: InterruptRequest.InterruptType.ValueType # 0 + INTERRUPT_TYPE_ALL: InterruptRequest.InterruptType.ValueType # 1 + """Interrupt all running executions within the session with the provided session_id.""" + INTERRUPT_TYPE_TAG: InterruptRequest.InterruptType.ValueType # 2 + """Interrupt all running executions within the session with the provided operation_tag.""" + INTERRUPT_TYPE_OPERATION_ID: InterruptRequest.InterruptType.ValueType # 3 + """Interrupt the running execution within the session with the provided operation_id.""" + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + INTERRUPT_TYPE_FIELD_NUMBER: builtins.int + OPERATION_TAG_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id specifies a spark session for a user id (which is specified + by user_context.user_id). The session_id is set by the client to be able to + collate streaming responses from different queries within the dedicated session. + The id should be an UUID string of the format `00112233-4455-6677-8899-aabbccddeeff` + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context""" + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + interrupt_type: global___InterruptRequest.InterruptType.ValueType + """(Required) The type of interrupt to execute.""" + operation_tag: builtins.str + """if interrupt_tag == INTERRUPT_TYPE_TAG, interrupt operation with this tag.""" + operation_id: builtins.str + """if interrupt_tag == INTERRUPT_TYPE_OPERATION_ID, interrupt operation with this operation_id.""" + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + client_type: builtins.str | None = ..., + interrupt_type: global___InterruptRequest.InterruptType.ValueType = ..., + operation_tag: builtins.str = ..., + operation_id: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "interrupt", + b"interrupt", + "operation_id", + b"operation_id", + "operation_tag", + b"operation_tag", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "interrupt", + b"interrupt", + "interrupt_type", + b"interrupt_type", + "operation_id", + b"operation_id", + "operation_tag", + b"operation_tag", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["interrupt", b"interrupt"] + ) -> typing_extensions.Literal["operation_tag", "operation_id"] | None: ... + +global___InterruptRequest = InterruptRequest + +class InterruptResponse(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SESSION_ID_FIELD_NUMBER: builtins.int + INTERRUPTED_IDS_FIELD_NUMBER: builtins.int + session_id: builtins.str + """Session id in which the interrupt was running.""" + @property + def interrupted_ids( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Operation ids of the executions which were interrupted.""" + def __init__( + self, + *, + session_id: builtins.str = ..., + interrupted_ids: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "interrupted_ids", b"interrupted_ids", "session_id", b"session_id" + ], + ) -> None: ... + +global___InterruptResponse = InterruptResponse + +class ReattachOptions(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + REATTACHABLE_FIELD_NUMBER: builtins.int + reattachable: builtins.bool + """If true, the request can be reattached to using ReattachExecute. + ReattachExecute can be used either if the stream broke with a GRPC network error, + or if the server closed the stream without sending a response with StreamStatus.complete=true. + The server will keep a buffer of responses in case a response is lost, and + ReattachExecute needs to back-track. + + If false, the execution response stream will will not be reattachable, and all responses are + immediately released by the server after being sent. + """ + def __init__( + self, + *, + reattachable: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["reattachable", b"reattachable"] + ) -> None: ... + +global___ReattachOptions = ReattachOptions + +class ReattachExecuteRequest(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + LAST_RESPONSE_ID_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id of the request to reattach to. + This must be an id of existing session. + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context + + user_context.user_id and session+id both identify a unique remote spark session on the + server side. + """ + operation_id: builtins.str + """(Required) + Provide an id of the request to reattach to. + This must be an id of existing operation. + """ + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + last_response_id: builtins.str + """(Optional) + Last already processed response id from the response stream. + After reattach, server will resume the response stream after that response. + If not specified, server will restart the stream from the start. + + Note: server controls the amount of responses that it buffers and it may drop responses, + that are far behind the latest returned response, so this can't be used to arbitrarily + scroll back the cursor. If the response is no longer available, this will result in an error. + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + operation_id: builtins.str = ..., + client_type: builtins.str | None = ..., + last_response_id: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "_last_response_id", + b"_last_response_id", + "client_type", + b"client_type", + "last_response_id", + b"last_response_id", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "_last_response_id", + b"_last_response_id", + "client_type", + b"client_type", + "last_response_id", + b"last_response_id", + "operation_id", + b"operation_id", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_last_response_id", b"_last_response_id"] + ) -> typing_extensions.Literal["last_response_id"] | None: ... + +global___ReattachExecuteRequest = ReattachExecuteRequest + +class ReleaseExecuteRequest(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ReleaseAll(google.protobuf.message.Message): + """Release and close operation completely. + This will also interrupt the query if it is running execution, and wait for it to be torn down. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + + class ReleaseUntil(google.protobuf.message.Message): + """Release all responses from the operation response stream up to and including + the response with the given by response_id. + While server determines by itself how much of a buffer of responses to keep, client providing + explicit release calls will help reduce resource consumption. + Noop if response_id not found in cached responses. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESPONSE_ID_FIELD_NUMBER: builtins.int + response_id: builtins.str + def __init__( + self, + *, + response_id: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["response_id", b"response_id"] + ) -> None: ... + + SESSION_ID_FIELD_NUMBER: builtins.int + USER_CONTEXT_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + CLIENT_TYPE_FIELD_NUMBER: builtins.int + RELEASE_ALL_FIELD_NUMBER: builtins.int + RELEASE_UNTIL_FIELD_NUMBER: builtins.int + session_id: builtins.str + """(Required) + + The session_id of the request to reattach to. + This must be an id of existing session. + """ + @property + def user_context(self) -> global___UserContext: + """(Required) User context + + user_context.user_id and session+id both identify a unique remote spark session on the + server side. + """ + operation_id: builtins.str + """(Required) + Provide an id of the request to reattach to. + This must be an id of existing operation. + """ + client_type: builtins.str + """Provides optional information about the client sending the request. This field + can be used for language or version specific information and is only intended for + logging purposes and will not be interpreted by the server. + """ + @property + def release_all(self) -> global___ReleaseExecuteRequest.ReleaseAll: ... + @property + def release_until(self) -> global___ReleaseExecuteRequest.ReleaseUntil: ... + def __init__( + self, + *, + session_id: builtins.str = ..., + user_context: global___UserContext | None = ..., + operation_id: builtins.str = ..., + client_type: builtins.str | None = ..., + release_all: global___ReleaseExecuteRequest.ReleaseAll | None = ..., + release_until: global___ReleaseExecuteRequest.ReleaseUntil | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "release", + b"release", + "release_all", + b"release_all", + "release_until", + b"release_until", + "user_context", + b"user_context", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_client_type", + b"_client_type", + "client_type", + b"client_type", + "operation_id", + b"operation_id", + "release", + b"release", + "release_all", + b"release_all", + "release_until", + b"release_until", + "session_id", + b"session_id", + "user_context", + b"user_context", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_client_type", b"_client_type"] + ) -> typing_extensions.Literal["client_type"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["release", b"release"] + ) -> typing_extensions.Literal["release_all", "release_until"] | None: ... + +global___ReleaseExecuteRequest = ReleaseExecuteRequest + +class ReleaseExecuteResponse(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SESSION_ID_FIELD_NUMBER: builtins.int + OPERATION_ID_FIELD_NUMBER: builtins.int + session_id: builtins.str + """Session id in which the release was running.""" + operation_id: builtins.str + """Operation id of the operation on which the release executed. + If the operation couldn't be found (because e.g. it was concurrently released), will be unset. + Otherwise, it will be equal to the operation_id from request. + """ + def __init__( + self, + *, + session_id: builtins.str = ..., + operation_id: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_operation_id", b"_operation_id", "operation_id", b"operation_id" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_operation_id", + b"_operation_id", + "operation_id", + b"operation_id", + "session_id", + b"session_id", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_operation_id", b"_operation_id"] + ) -> typing_extensions.Literal["operation_id"] | None: ... + +global___ReleaseExecuteResponse = ReleaseExecuteResponse diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2_grpc.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2_grpc.py new file mode 100644 index 0000000..e6bfda8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/base_pb2_grpc.py @@ -0,0 +1,423 @@ +# +# 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. +# +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +from pyspark.sql.connect.proto import base_pb2 as spark_dot_connect_dot_base__pb2 + + +class SparkConnectServiceStub(object): + """Main interface for the SparkConnect service.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.ExecutePlan = channel.unary_stream( + "/spark.connect.SparkConnectService/ExecutePlan", + request_serializer=spark_dot_connect_dot_base__pb2.ExecutePlanRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.ExecutePlanResponse.FromString, + ) + self.AnalyzePlan = channel.unary_unary( + "/spark.connect.SparkConnectService/AnalyzePlan", + request_serializer=spark_dot_connect_dot_base__pb2.AnalyzePlanRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.AnalyzePlanResponse.FromString, + ) + self.Config = channel.unary_unary( + "/spark.connect.SparkConnectService/Config", + request_serializer=spark_dot_connect_dot_base__pb2.ConfigRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.ConfigResponse.FromString, + ) + self.AddArtifacts = channel.stream_unary( + "/spark.connect.SparkConnectService/AddArtifacts", + request_serializer=spark_dot_connect_dot_base__pb2.AddArtifactsRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.AddArtifactsResponse.FromString, + ) + self.ArtifactStatus = channel.unary_unary( + "/spark.connect.SparkConnectService/ArtifactStatus", + request_serializer=spark_dot_connect_dot_base__pb2.ArtifactStatusesRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.ArtifactStatusesResponse.FromString, + ) + self.Interrupt = channel.unary_unary( + "/spark.connect.SparkConnectService/Interrupt", + request_serializer=spark_dot_connect_dot_base__pb2.InterruptRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.InterruptResponse.FromString, + ) + self.ReattachExecute = channel.unary_stream( + "/spark.connect.SparkConnectService/ReattachExecute", + request_serializer=spark_dot_connect_dot_base__pb2.ReattachExecuteRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.ExecutePlanResponse.FromString, + ) + self.ReleaseExecute = channel.unary_unary( + "/spark.connect.SparkConnectService/ReleaseExecute", + request_serializer=spark_dot_connect_dot_base__pb2.ReleaseExecuteRequest.SerializeToString, + response_deserializer=spark_dot_connect_dot_base__pb2.ReleaseExecuteResponse.FromString, + ) + + +class SparkConnectServiceServicer(object): + """Main interface for the SparkConnect service.""" + + def ExecutePlan(self, request, context): + """Executes a request that contains the query and returns a stream of [[Response]]. + + It is guaranteed that there is at least one ARROW batch returned even if the result set is empty. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def AnalyzePlan(self, request, context): + """Analyzes a query and returns a [[AnalyzeResponse]] containing metadata about the query.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def Config(self, request, context): + """Update or fetch the configurations and returns a [[ConfigResponse]] containing the result.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def AddArtifacts(self, request_iterator, context): + """Add artifacts to the session and returns a [[AddArtifactsResponse]] containing metadata about + the added artifacts. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def ArtifactStatus(self, request, context): + """Check statuses of artifacts in the session and returns them in a [[ArtifactStatusesResponse]]""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def Interrupt(self, request, context): + """Interrupts running executions""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def ReattachExecute(self, request, context): + """Reattach to an existing reattachable execution. + The ExecutePlan must have been started with ReattachOptions.reattachable=true. + If the ExecutePlanResponse stream ends without a ResultComplete message, there is more to + continue. If there is a ResultComplete, the client should use ReleaseExecute with + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + def ReleaseExecute(self, request, context): + """Release an reattachable execution, or parts thereof. + The ExecutePlan must have been started with ReattachOptions.reattachable=true. + Non reattachable executions are released automatically and immediately after the ExecutePlan + RPC and ReleaseExecute may not be used. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details("Method not implemented!") + raise NotImplementedError("Method not implemented!") + + +def add_SparkConnectServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + "ExecutePlan": grpc.unary_stream_rpc_method_handler( + servicer.ExecutePlan, + request_deserializer=spark_dot_connect_dot_base__pb2.ExecutePlanRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.ExecutePlanResponse.SerializeToString, + ), + "AnalyzePlan": grpc.unary_unary_rpc_method_handler( + servicer.AnalyzePlan, + request_deserializer=spark_dot_connect_dot_base__pb2.AnalyzePlanRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.AnalyzePlanResponse.SerializeToString, + ), + "Config": grpc.unary_unary_rpc_method_handler( + servicer.Config, + request_deserializer=spark_dot_connect_dot_base__pb2.ConfigRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.ConfigResponse.SerializeToString, + ), + "AddArtifacts": grpc.stream_unary_rpc_method_handler( + servicer.AddArtifacts, + request_deserializer=spark_dot_connect_dot_base__pb2.AddArtifactsRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.AddArtifactsResponse.SerializeToString, + ), + "ArtifactStatus": grpc.unary_unary_rpc_method_handler( + servicer.ArtifactStatus, + request_deserializer=spark_dot_connect_dot_base__pb2.ArtifactStatusesRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.ArtifactStatusesResponse.SerializeToString, + ), + "Interrupt": grpc.unary_unary_rpc_method_handler( + servicer.Interrupt, + request_deserializer=spark_dot_connect_dot_base__pb2.InterruptRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.InterruptResponse.SerializeToString, + ), + "ReattachExecute": grpc.unary_stream_rpc_method_handler( + servicer.ReattachExecute, + request_deserializer=spark_dot_connect_dot_base__pb2.ReattachExecuteRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.ExecutePlanResponse.SerializeToString, + ), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler( + servicer.ReleaseExecute, + request_deserializer=spark_dot_connect_dot_base__pb2.ReleaseExecuteRequest.FromString, + response_serializer=spark_dot_connect_dot_base__pb2.ReleaseExecuteResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", rpc_method_handlers + ) + server.add_generic_rpc_handlers((generic_handler,)) + + +# This class is part of an EXPERIMENTAL API. +class SparkConnectService(object): + """Main interface for the SparkConnect service.""" + + @staticmethod + def ExecutePlan( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_stream( + request, + target, + "/spark.connect.SparkConnectService/ExecutePlan", + spark_dot_connect_dot_base__pb2.ExecutePlanRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.ExecutePlanResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def AnalyzePlan( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/spark.connect.SparkConnectService/AnalyzePlan", + spark_dot_connect_dot_base__pb2.AnalyzePlanRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.AnalyzePlanResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def Config( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/spark.connect.SparkConnectService/Config", + spark_dot_connect_dot_base__pb2.ConfigRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.ConfigResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def AddArtifacts( + request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.stream_unary( + request_iterator, + target, + "/spark.connect.SparkConnectService/AddArtifacts", + spark_dot_connect_dot_base__pb2.AddArtifactsRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.AddArtifactsResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def ArtifactStatus( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/spark.connect.SparkConnectService/ArtifactStatus", + spark_dot_connect_dot_base__pb2.ArtifactStatusesRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.ArtifactStatusesResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def Interrupt( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/spark.connect.SparkConnectService/Interrupt", + spark_dot_connect_dot_base__pb2.InterruptRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.InterruptResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def ReattachExecute( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_stream( + request, + target, + "/spark.connect.SparkConnectService/ReattachExecute", + spark_dot_connect_dot_base__pb2.ReattachExecuteRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.ExecutePlanResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) + + @staticmethod + def ReleaseExecute( + request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None, + ): + return grpc.experimental.unary_unary( + request, + target, + "/spark.connect.SparkConnectService/ReleaseExecute", + spark_dot_connect_dot_base__pb2.ReleaseExecuteRequest.SerializeToString, + spark_dot_connect_dot_base__pb2.ReleaseExecuteResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.py new file mode 100644 index 0000000..709f0f0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.py @@ -0,0 +1,109 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/catalog.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from pyspark.sql.connect.proto import common_pb2 as spark_dot_connect_dot_common__pb2 +from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1bspark/connect/catalog.proto\x12\rspark.connect\x1a\x1aspark/connect/common.proto\x1a\x19spark/connect/types.proto"\xc6\x0e\n\x07\x43\x61talog\x12K\n\x10\x63urrent_database\x18\x01 \x01(\x0b\x32\x1e.spark.connect.CurrentDatabaseH\x00R\x0f\x63urrentDatabase\x12U\n\x14set_current_database\x18\x02 \x01(\x0b\x32!.spark.connect.SetCurrentDatabaseH\x00R\x12setCurrentDatabase\x12\x45\n\x0elist_databases\x18\x03 \x01(\x0b\x32\x1c.spark.connect.ListDatabasesH\x00R\rlistDatabases\x12<\n\x0blist_tables\x18\x04 \x01(\x0b\x32\x19.spark.connect.ListTablesH\x00R\nlistTables\x12\x45\n\x0elist_functions\x18\x05 \x01(\x0b\x32\x1c.spark.connect.ListFunctionsH\x00R\rlistFunctions\x12?\n\x0clist_columns\x18\x06 \x01(\x0b\x32\x1a.spark.connect.ListColumnsH\x00R\x0blistColumns\x12?\n\x0cget_database\x18\x07 \x01(\x0b\x32\x1a.spark.connect.GetDatabaseH\x00R\x0bgetDatabase\x12\x36\n\tget_table\x18\x08 \x01(\x0b\x32\x17.spark.connect.GetTableH\x00R\x08getTable\x12?\n\x0cget_function\x18\t \x01(\x0b\x32\x1a.spark.connect.GetFunctionH\x00R\x0bgetFunction\x12H\n\x0f\x64\x61tabase_exists\x18\n \x01(\x0b\x32\x1d.spark.connect.DatabaseExistsH\x00R\x0e\x64\x61tabaseExists\x12?\n\x0ctable_exists\x18\x0b \x01(\x0b\x32\x1a.spark.connect.TableExistsH\x00R\x0btableExists\x12H\n\x0f\x66unction_exists\x18\x0c \x01(\x0b\x32\x1d.spark.connect.FunctionExistsH\x00R\x0e\x66unctionExists\x12X\n\x15\x63reate_external_table\x18\r \x01(\x0b\x32".spark.connect.CreateExternalTableH\x00R\x13\x63reateExternalTable\x12?\n\x0c\x63reate_table\x18\x0e \x01(\x0b\x32\x1a.spark.connect.CreateTableH\x00R\x0b\x63reateTable\x12\x43\n\x0e\x64rop_temp_view\x18\x0f \x01(\x0b\x32\x1b.spark.connect.DropTempViewH\x00R\x0c\x64ropTempView\x12V\n\x15\x64rop_global_temp_view\x18\x10 \x01(\x0b\x32!.spark.connect.DropGlobalTempViewH\x00R\x12\x64ropGlobalTempView\x12Q\n\x12recover_partitions\x18\x11 \x01(\x0b\x32 .spark.connect.RecoverPartitionsH\x00R\x11recoverPartitions\x12\x36\n\tis_cached\x18\x12 \x01(\x0b\x32\x17.spark.connect.IsCachedH\x00R\x08isCached\x12<\n\x0b\x63\x61\x63he_table\x18\x13 \x01(\x0b\x32\x19.spark.connect.CacheTableH\x00R\ncacheTable\x12\x42\n\runcache_table\x18\x14 \x01(\x0b\x32\x1b.spark.connect.UncacheTableH\x00R\x0cuncacheTable\x12<\n\x0b\x63lear_cache\x18\x15 \x01(\x0b\x32\x19.spark.connect.ClearCacheH\x00R\nclearCache\x12\x42\n\rrefresh_table\x18\x16 \x01(\x0b\x32\x1b.spark.connect.RefreshTableH\x00R\x0crefreshTable\x12\x46\n\x0frefresh_by_path\x18\x17 \x01(\x0b\x32\x1c.spark.connect.RefreshByPathH\x00R\rrefreshByPath\x12H\n\x0f\x63urrent_catalog\x18\x18 \x01(\x0b\x32\x1d.spark.connect.CurrentCatalogH\x00R\x0e\x63urrentCatalog\x12R\n\x13set_current_catalog\x18\x19 \x01(\x0b\x32 .spark.connect.SetCurrentCatalogH\x00R\x11setCurrentCatalog\x12\x42\n\rlist_catalogs\x18\x1a \x01(\x0b\x32\x1b.spark.connect.ListCatalogsH\x00R\x0clistCatalogsB\n\n\x08\x63\x61t_type"\x11\n\x0f\x43urrentDatabase"-\n\x12SetCurrentDatabase\x12\x17\n\x07\x64\x62_name\x18\x01 \x01(\tR\x06\x64\x62Name":\n\rListDatabases\x12\x1d\n\x07pattern\x18\x01 \x01(\tH\x00R\x07pattern\x88\x01\x01\x42\n\n\x08_pattern"a\n\nListTables\x12\x1c\n\x07\x64\x62_name\x18\x01 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x12\x1d\n\x07pattern\x18\x02 \x01(\tH\x01R\x07pattern\x88\x01\x01\x42\n\n\x08_db_nameB\n\n\x08_pattern"d\n\rListFunctions\x12\x1c\n\x07\x64\x62_name\x18\x01 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x12\x1d\n\x07pattern\x18\x02 \x01(\tH\x01R\x07pattern\x88\x01\x01\x42\n\n\x08_db_nameB\n\n\x08_pattern"V\n\x0bListColumns\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x1c\n\x07\x64\x62_name\x18\x02 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x42\n\n\x08_db_name"&\n\x0bGetDatabase\x12\x17\n\x07\x64\x62_name\x18\x01 \x01(\tR\x06\x64\x62Name"S\n\x08GetTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x1c\n\x07\x64\x62_name\x18\x02 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x42\n\n\x08_db_name"\\\n\x0bGetFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x1c\n\x07\x64\x62_name\x18\x02 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x42\n\n\x08_db_name")\n\x0e\x44\x61tabaseExists\x12\x17\n\x07\x64\x62_name\x18\x01 \x01(\tR\x06\x64\x62Name"V\n\x0bTableExists\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x1c\n\x07\x64\x62_name\x18\x02 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x42\n\n\x08_db_name"_\n\x0e\x46unctionExists\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x1c\n\x07\x64\x62_name\x18\x02 \x01(\tH\x00R\x06\x64\x62Name\x88\x01\x01\x42\n\n\x08_db_name"\xc6\x02\n\x13\x43reateExternalTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x17\n\x04path\x18\x02 \x01(\tH\x00R\x04path\x88\x01\x01\x12\x1b\n\x06source\x18\x03 \x01(\tH\x01R\x06source\x88\x01\x01\x12\x34\n\x06schema\x18\x04 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x02R\x06schema\x88\x01\x01\x12I\n\x07options\x18\x05 \x03(\x0b\x32/.spark.connect.CreateExternalTable.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\x07\n\x05_pathB\t\n\x07_sourceB\t\n\x07_schema"\xed\x02\n\x0b\x43reateTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x17\n\x04path\x18\x02 \x01(\tH\x00R\x04path\x88\x01\x01\x12\x1b\n\x06source\x18\x03 \x01(\tH\x01R\x06source\x88\x01\x01\x12%\n\x0b\x64\x65scription\x18\x04 \x01(\tH\x02R\x0b\x64\x65scription\x88\x01\x01\x12\x34\n\x06schema\x18\x05 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x03R\x06schema\x88\x01\x01\x12\x41\n\x07options\x18\x06 \x03(\x0b\x32\'.spark.connect.CreateTable.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\x07\n\x05_pathB\t\n\x07_sourceB\x0e\n\x0c_descriptionB\t\n\x07_schema"+\n\x0c\x44ropTempView\x12\x1b\n\tview_name\x18\x01 \x01(\tR\x08viewName"1\n\x12\x44ropGlobalTempView\x12\x1b\n\tview_name\x18\x01 \x01(\tR\x08viewName"2\n\x11RecoverPartitions\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName")\n\x08IsCached\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName"\x84\x01\n\nCacheTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12\x45\n\rstorage_level\x18\x02 \x01(\x0b\x32\x1b.spark.connect.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"-\n\x0cUncacheTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName"\x0c\n\nClearCache"-\n\x0cRefreshTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName"#\n\rRefreshByPath\x12\x12\n\x04path\x18\x01 \x01(\tR\x04path"\x10\n\x0e\x43urrentCatalog"6\n\x11SetCurrentCatalog\x12!\n\x0c\x63\x61talog_name\x18\x01 \x01(\tR\x0b\x63\x61talogName"9\n\x0cListCatalogs\x12\x1d\n\x07pattern\x18\x01 \x01(\tH\x00R\x07pattern\x88\x01\x01\x42\n\n\x08_patternB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.catalog_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _CREATEEXTERNALTABLE_OPTIONSENTRY._options = None + _CREATEEXTERNALTABLE_OPTIONSENTRY._serialized_options = b"8\001" + _CREATETABLE_OPTIONSENTRY._options = None + _CREATETABLE_OPTIONSENTRY._serialized_options = b"8\001" + _CATALOG._serialized_start = 102 + _CATALOG._serialized_end = 1964 + _CURRENTDATABASE._serialized_start = 1966 + _CURRENTDATABASE._serialized_end = 1983 + _SETCURRENTDATABASE._serialized_start = 1985 + _SETCURRENTDATABASE._serialized_end = 2030 + _LISTDATABASES._serialized_start = 2032 + _LISTDATABASES._serialized_end = 2090 + _LISTTABLES._serialized_start = 2092 + _LISTTABLES._serialized_end = 2189 + _LISTFUNCTIONS._serialized_start = 2191 + _LISTFUNCTIONS._serialized_end = 2291 + _LISTCOLUMNS._serialized_start = 2293 + _LISTCOLUMNS._serialized_end = 2379 + _GETDATABASE._serialized_start = 2381 + _GETDATABASE._serialized_end = 2419 + _GETTABLE._serialized_start = 2421 + _GETTABLE._serialized_end = 2504 + _GETFUNCTION._serialized_start = 2506 + _GETFUNCTION._serialized_end = 2598 + _DATABASEEXISTS._serialized_start = 2600 + _DATABASEEXISTS._serialized_end = 2641 + _TABLEEXISTS._serialized_start = 2643 + _TABLEEXISTS._serialized_end = 2729 + _FUNCTIONEXISTS._serialized_start = 2731 + _FUNCTIONEXISTS._serialized_end = 2826 + _CREATEEXTERNALTABLE._serialized_start = 2829 + _CREATEEXTERNALTABLE._serialized_end = 3155 + _CREATEEXTERNALTABLE_OPTIONSENTRY._serialized_start = 3066 + _CREATEEXTERNALTABLE_OPTIONSENTRY._serialized_end = 3124 + _CREATETABLE._serialized_start = 3158 + _CREATETABLE._serialized_end = 3523 + _CREATETABLE_OPTIONSENTRY._serialized_start = 3066 + _CREATETABLE_OPTIONSENTRY._serialized_end = 3124 + _DROPTEMPVIEW._serialized_start = 3525 + _DROPTEMPVIEW._serialized_end = 3568 + _DROPGLOBALTEMPVIEW._serialized_start = 3570 + _DROPGLOBALTEMPVIEW._serialized_end = 3619 + _RECOVERPARTITIONS._serialized_start = 3621 + _RECOVERPARTITIONS._serialized_end = 3671 + _ISCACHED._serialized_start = 3673 + _ISCACHED._serialized_end = 3714 + _CACHETABLE._serialized_start = 3717 + _CACHETABLE._serialized_end = 3849 + _UNCACHETABLE._serialized_start = 3851 + _UNCACHETABLE._serialized_end = 3896 + _CLEARCACHE._serialized_start = 3898 + _CLEARCACHE._serialized_end = 3910 + _REFRESHTABLE._serialized_start = 3912 + _REFRESHTABLE._serialized_end = 3957 + _REFRESHBYPATH._serialized_start = 3959 + _REFRESHBYPATH._serialized_end = 3994 + _CURRENTCATALOG._serialized_start = 3996 + _CURRENTCATALOG._serialized_end = 4012 + _SETCURRENTCATALOG._serialized_start = 4014 + _SETCURRENTCATALOG._serialized_end = 4068 + _LISTCATALOGS._serialized_start = 4070 + _LISTCATALOGS._serialized_end = 4127 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.pyi new file mode 100644 index 0000000..3d14961 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/catalog_pb2.pyi @@ -0,0 +1,1130 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.message +import pyspark.sql.connect.proto.common_pb2 +import pyspark.sql.connect.proto.types_pb2 +import sys +import typing + +if sys.version_info >= (3, 8): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class Catalog(google.protobuf.message.Message): + """Catalog messages are marked as unstable.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CURRENT_DATABASE_FIELD_NUMBER: builtins.int + SET_CURRENT_DATABASE_FIELD_NUMBER: builtins.int + LIST_DATABASES_FIELD_NUMBER: builtins.int + LIST_TABLES_FIELD_NUMBER: builtins.int + LIST_FUNCTIONS_FIELD_NUMBER: builtins.int + LIST_COLUMNS_FIELD_NUMBER: builtins.int + GET_DATABASE_FIELD_NUMBER: builtins.int + GET_TABLE_FIELD_NUMBER: builtins.int + GET_FUNCTION_FIELD_NUMBER: builtins.int + DATABASE_EXISTS_FIELD_NUMBER: builtins.int + TABLE_EXISTS_FIELD_NUMBER: builtins.int + FUNCTION_EXISTS_FIELD_NUMBER: builtins.int + CREATE_EXTERNAL_TABLE_FIELD_NUMBER: builtins.int + CREATE_TABLE_FIELD_NUMBER: builtins.int + DROP_TEMP_VIEW_FIELD_NUMBER: builtins.int + DROP_GLOBAL_TEMP_VIEW_FIELD_NUMBER: builtins.int + RECOVER_PARTITIONS_FIELD_NUMBER: builtins.int + IS_CACHED_FIELD_NUMBER: builtins.int + CACHE_TABLE_FIELD_NUMBER: builtins.int + UNCACHE_TABLE_FIELD_NUMBER: builtins.int + CLEAR_CACHE_FIELD_NUMBER: builtins.int + REFRESH_TABLE_FIELD_NUMBER: builtins.int + REFRESH_BY_PATH_FIELD_NUMBER: builtins.int + CURRENT_CATALOG_FIELD_NUMBER: builtins.int + SET_CURRENT_CATALOG_FIELD_NUMBER: builtins.int + LIST_CATALOGS_FIELD_NUMBER: builtins.int + @property + def current_database(self) -> global___CurrentDatabase: ... + @property + def set_current_database(self) -> global___SetCurrentDatabase: ... + @property + def list_databases(self) -> global___ListDatabases: ... + @property + def list_tables(self) -> global___ListTables: ... + @property + def list_functions(self) -> global___ListFunctions: ... + @property + def list_columns(self) -> global___ListColumns: ... + @property + def get_database(self) -> global___GetDatabase: ... + @property + def get_table(self) -> global___GetTable: ... + @property + def get_function(self) -> global___GetFunction: ... + @property + def database_exists(self) -> global___DatabaseExists: ... + @property + def table_exists(self) -> global___TableExists: ... + @property + def function_exists(self) -> global___FunctionExists: ... + @property + def create_external_table(self) -> global___CreateExternalTable: ... + @property + def create_table(self) -> global___CreateTable: ... + @property + def drop_temp_view(self) -> global___DropTempView: ... + @property + def drop_global_temp_view(self) -> global___DropGlobalTempView: ... + @property + def recover_partitions(self) -> global___RecoverPartitions: ... + @property + def is_cached(self) -> global___IsCached: ... + @property + def cache_table(self) -> global___CacheTable: ... + @property + def uncache_table(self) -> global___UncacheTable: ... + @property + def clear_cache(self) -> global___ClearCache: ... + @property + def refresh_table(self) -> global___RefreshTable: ... + @property + def refresh_by_path(self) -> global___RefreshByPath: ... + @property + def current_catalog(self) -> global___CurrentCatalog: ... + @property + def set_current_catalog(self) -> global___SetCurrentCatalog: ... + @property + def list_catalogs(self) -> global___ListCatalogs: ... + def __init__( + self, + *, + current_database: global___CurrentDatabase | None = ..., + set_current_database: global___SetCurrentDatabase | None = ..., + list_databases: global___ListDatabases | None = ..., + list_tables: global___ListTables | None = ..., + list_functions: global___ListFunctions | None = ..., + list_columns: global___ListColumns | None = ..., + get_database: global___GetDatabase | None = ..., + get_table: global___GetTable | None = ..., + get_function: global___GetFunction | None = ..., + database_exists: global___DatabaseExists | None = ..., + table_exists: global___TableExists | None = ..., + function_exists: global___FunctionExists | None = ..., + create_external_table: global___CreateExternalTable | None = ..., + create_table: global___CreateTable | None = ..., + drop_temp_view: global___DropTempView | None = ..., + drop_global_temp_view: global___DropGlobalTempView | None = ..., + recover_partitions: global___RecoverPartitions | None = ..., + is_cached: global___IsCached | None = ..., + cache_table: global___CacheTable | None = ..., + uncache_table: global___UncacheTable | None = ..., + clear_cache: global___ClearCache | None = ..., + refresh_table: global___RefreshTable | None = ..., + refresh_by_path: global___RefreshByPath | None = ..., + current_catalog: global___CurrentCatalog | None = ..., + set_current_catalog: global___SetCurrentCatalog | None = ..., + list_catalogs: global___ListCatalogs | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "cache_table", + b"cache_table", + "cat_type", + b"cat_type", + "clear_cache", + b"clear_cache", + "create_external_table", + b"create_external_table", + "create_table", + b"create_table", + "current_catalog", + b"current_catalog", + "current_database", + b"current_database", + "database_exists", + b"database_exists", + "drop_global_temp_view", + b"drop_global_temp_view", + "drop_temp_view", + b"drop_temp_view", + "function_exists", + b"function_exists", + "get_database", + b"get_database", + "get_function", + b"get_function", + "get_table", + b"get_table", + "is_cached", + b"is_cached", + "list_catalogs", + b"list_catalogs", + "list_columns", + b"list_columns", + "list_databases", + b"list_databases", + "list_functions", + b"list_functions", + "list_tables", + b"list_tables", + "recover_partitions", + b"recover_partitions", + "refresh_by_path", + b"refresh_by_path", + "refresh_table", + b"refresh_table", + "set_current_catalog", + b"set_current_catalog", + "set_current_database", + b"set_current_database", + "table_exists", + b"table_exists", + "uncache_table", + b"uncache_table", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "cache_table", + b"cache_table", + "cat_type", + b"cat_type", + "clear_cache", + b"clear_cache", + "create_external_table", + b"create_external_table", + "create_table", + b"create_table", + "current_catalog", + b"current_catalog", + "current_database", + b"current_database", + "database_exists", + b"database_exists", + "drop_global_temp_view", + b"drop_global_temp_view", + "drop_temp_view", + b"drop_temp_view", + "function_exists", + b"function_exists", + "get_database", + b"get_database", + "get_function", + b"get_function", + "get_table", + b"get_table", + "is_cached", + b"is_cached", + "list_catalogs", + b"list_catalogs", + "list_columns", + b"list_columns", + "list_databases", + b"list_databases", + "list_functions", + b"list_functions", + "list_tables", + b"list_tables", + "recover_partitions", + b"recover_partitions", + "refresh_by_path", + b"refresh_by_path", + "refresh_table", + b"refresh_table", + "set_current_catalog", + b"set_current_catalog", + "set_current_database", + b"set_current_database", + "table_exists", + b"table_exists", + "uncache_table", + b"uncache_table", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["cat_type", b"cat_type"] + ) -> typing_extensions.Literal[ + "current_database", + "set_current_database", + "list_databases", + "list_tables", + "list_functions", + "list_columns", + "get_database", + "get_table", + "get_function", + "database_exists", + "table_exists", + "function_exists", + "create_external_table", + "create_table", + "drop_temp_view", + "drop_global_temp_view", + "recover_partitions", + "is_cached", + "cache_table", + "uncache_table", + "clear_cache", + "refresh_table", + "refresh_by_path", + "current_catalog", + "set_current_catalog", + "list_catalogs", + ] | None: ... + +global___Catalog = Catalog + +class CurrentDatabase(google.protobuf.message.Message): + """See `spark.catalog.currentDatabase`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___CurrentDatabase = CurrentDatabase + +class SetCurrentDatabase(google.protobuf.message.Message): + """See `spark.catalog.setCurrentDatabase`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DB_NAME_FIELD_NUMBER: builtins.int + db_name: builtins.str + """(Required)""" + def __init__( + self, + *, + db_name: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["db_name", b"db_name"]) -> None: ... + +global___SetCurrentDatabase = SetCurrentDatabase + +class ListDatabases(google.protobuf.message.Message): + """See `spark.catalog.listDatabases`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PATTERN_FIELD_NUMBER: builtins.int + pattern: builtins.str + """(Optional) The pattern that the database name needs to match""" + def __init__( + self, + *, + pattern: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_pattern", b"_pattern", "pattern", b"pattern"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["_pattern", b"_pattern", "pattern", b"pattern"] + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_pattern", b"_pattern"] + ) -> typing_extensions.Literal["pattern"] | None: ... + +global___ListDatabases = ListDatabases + +class ListTables(google.protobuf.message.Message): + """See `spark.catalog.listTables`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DB_NAME_FIELD_NUMBER: builtins.int + PATTERN_FIELD_NUMBER: builtins.int + db_name: builtins.str + """(Optional)""" + pattern: builtins.str + """(Optional) The pattern that the table name needs to match""" + def __init__( + self, + *, + db_name: builtins.str | None = ..., + pattern: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_db_name", + b"_db_name", + "_pattern", + b"_pattern", + "db_name", + b"db_name", + "pattern", + b"pattern", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", + b"_db_name", + "_pattern", + b"_pattern", + "db_name", + b"db_name", + "pattern", + b"pattern", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_pattern", b"_pattern"] + ) -> typing_extensions.Literal["pattern"] | None: ... + +global___ListTables = ListTables + +class ListFunctions(google.protobuf.message.Message): + """See `spark.catalog.listFunctions`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DB_NAME_FIELD_NUMBER: builtins.int + PATTERN_FIELD_NUMBER: builtins.int + db_name: builtins.str + """(Optional)""" + pattern: builtins.str + """(Optional) The pattern that the function name needs to match""" + def __init__( + self, + *, + db_name: builtins.str | None = ..., + pattern: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_db_name", + b"_db_name", + "_pattern", + b"_pattern", + "db_name", + b"db_name", + "pattern", + b"pattern", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", + b"_db_name", + "_pattern", + b"_pattern", + "db_name", + b"db_name", + "pattern", + b"pattern", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_pattern", b"_pattern"] + ) -> typing_extensions.Literal["pattern"] | None: ... + +global___ListFunctions = ListFunctions + +class ListColumns(google.protobuf.message.Message): + """See `spark.catalog.listColumns`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + DB_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + db_name: builtins.str + """(Optional)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + db_name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_db_name", b"_db_name", "db_name", b"db_name"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", b"_db_name", "db_name", b"db_name", "table_name", b"table_name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + +global___ListColumns = ListColumns + +class GetDatabase(google.protobuf.message.Message): + """See `spark.catalog.getDatabase`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DB_NAME_FIELD_NUMBER: builtins.int + db_name: builtins.str + """(Required)""" + def __init__( + self, + *, + db_name: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["db_name", b"db_name"]) -> None: ... + +global___GetDatabase = GetDatabase + +class GetTable(google.protobuf.message.Message): + """See `spark.catalog.getTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + DB_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + db_name: builtins.str + """(Optional)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + db_name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_db_name", b"_db_name", "db_name", b"db_name"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", b"_db_name", "db_name", b"db_name", "table_name", b"table_name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + +global___GetTable = GetTable + +class GetFunction(google.protobuf.message.Message): + """See `spark.catalog.getFunction`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + DB_NAME_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required)""" + db_name: builtins.str + """(Optional)""" + def __init__( + self, + *, + function_name: builtins.str = ..., + db_name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_db_name", b"_db_name", "db_name", b"db_name"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", b"_db_name", "db_name", b"db_name", "function_name", b"function_name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + +global___GetFunction = GetFunction + +class DatabaseExists(google.protobuf.message.Message): + """See `spark.catalog.databaseExists`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DB_NAME_FIELD_NUMBER: builtins.int + db_name: builtins.str + """(Required)""" + def __init__( + self, + *, + db_name: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["db_name", b"db_name"]) -> None: ... + +global___DatabaseExists = DatabaseExists + +class TableExists(google.protobuf.message.Message): + """See `spark.catalog.tableExists`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + DB_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + db_name: builtins.str + """(Optional)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + db_name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_db_name", b"_db_name", "db_name", b"db_name"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", b"_db_name", "db_name", b"db_name", "table_name", b"table_name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + +global___TableExists = TableExists + +class FunctionExists(google.protobuf.message.Message): + """See `spark.catalog.functionExists`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + DB_NAME_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required)""" + db_name: builtins.str + """(Optional)""" + def __init__( + self, + *, + function_name: builtins.str = ..., + db_name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_db_name", b"_db_name", "db_name", b"db_name"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_db_name", b"_db_name", "db_name", b"db_name", "function_name", b"function_name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_db_name", b"_db_name"] + ) -> typing_extensions.Literal["db_name"] | None: ... + +global___FunctionExists = FunctionExists + +class CreateExternalTable(google.protobuf.message.Message): + """See `spark.catalog.createExternalTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + TABLE_NAME_FIELD_NUMBER: builtins.int + PATH_FIELD_NUMBER: builtins.int + SOURCE_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + path: builtins.str + """(Optional)""" + source: builtins.str + """(Optional)""" + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional)""" + @property + def options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Options could be empty for valid data source format. + The map key is case insensitive. + """ + def __init__( + self, + *, + table_name: builtins.str = ..., + path: builtins.str | None = ..., + source: builtins.str | None = ..., + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_path", + b"_path", + "_schema", + b"_schema", + "_source", + b"_source", + "path", + b"path", + "schema", + b"schema", + "source", + b"source", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_path", + b"_path", + "_schema", + b"_schema", + "_source", + b"_source", + "options", + b"options", + "path", + b"path", + "schema", + b"schema", + "source", + b"source", + "table_name", + b"table_name", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_path", b"_path"] + ) -> typing_extensions.Literal["path"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_schema", b"_schema"] + ) -> typing_extensions.Literal["schema"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_source", b"_source"] + ) -> typing_extensions.Literal["source"] | None: ... + +global___CreateExternalTable = CreateExternalTable + +class CreateTable(google.protobuf.message.Message): + """See `spark.catalog.createTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + TABLE_NAME_FIELD_NUMBER: builtins.int + PATH_FIELD_NUMBER: builtins.int + SOURCE_FIELD_NUMBER: builtins.int + DESCRIPTION_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + path: builtins.str + """(Optional)""" + source: builtins.str + """(Optional)""" + description: builtins.str + """(Optional)""" + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional)""" + @property + def options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Options could be empty for valid data source format. + The map key is case insensitive. + """ + def __init__( + self, + *, + table_name: builtins.str = ..., + path: builtins.str | None = ..., + source: builtins.str | None = ..., + description: builtins.str | None = ..., + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_description", + b"_description", + "_path", + b"_path", + "_schema", + b"_schema", + "_source", + b"_source", + "description", + b"description", + "path", + b"path", + "schema", + b"schema", + "source", + b"source", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_description", + b"_description", + "_path", + b"_path", + "_schema", + b"_schema", + "_source", + b"_source", + "description", + b"description", + "options", + b"options", + "path", + b"path", + "schema", + b"schema", + "source", + b"source", + "table_name", + b"table_name", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_description", b"_description"] + ) -> typing_extensions.Literal["description"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_path", b"_path"] + ) -> typing_extensions.Literal["path"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_schema", b"_schema"] + ) -> typing_extensions.Literal["schema"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_source", b"_source"] + ) -> typing_extensions.Literal["source"] | None: ... + +global___CreateTable = CreateTable + +class DropTempView(google.protobuf.message.Message): + """See `spark.catalog.dropTempView`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VIEW_NAME_FIELD_NUMBER: builtins.int + view_name: builtins.str + """(Required)""" + def __init__( + self, + *, + view_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["view_name", b"view_name"] + ) -> None: ... + +global___DropTempView = DropTempView + +class DropGlobalTempView(google.protobuf.message.Message): + """See `spark.catalog.dropGlobalTempView`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VIEW_NAME_FIELD_NUMBER: builtins.int + view_name: builtins.str + """(Required)""" + def __init__( + self, + *, + view_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["view_name", b"view_name"] + ) -> None: ... + +global___DropGlobalTempView = DropGlobalTempView + +class RecoverPartitions(google.protobuf.message.Message): + """See `spark.catalog.recoverPartitions`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["table_name", b"table_name"] + ) -> None: ... + +global___RecoverPartitions = RecoverPartitions + +class IsCached(google.protobuf.message.Message): + """See `spark.catalog.isCached`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["table_name", b"table_name"] + ) -> None: ... + +global___IsCached = IsCached + +class CacheTable(google.protobuf.message.Message): + """See `spark.catalog.cacheTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + @property + def storage_level(self) -> pyspark.sql.connect.proto.common_pb2.StorageLevel: + """(Optional)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + storage_level: pyspark.sql.connect.proto.common_pb2.StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "storage_level", + b"storage_level", + "table_name", + b"table_name", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___CacheTable = CacheTable + +class UncacheTable(google.protobuf.message.Message): + """See `spark.catalog.uncacheTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["table_name", b"table_name"] + ) -> None: ... + +global___UncacheTable = UncacheTable + +class ClearCache(google.protobuf.message.Message): + """See `spark.catalog.clearCache`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___ClearCache = ClearCache + +class RefreshTable(google.protobuf.message.Message): + """See `spark.catalog.refreshTable`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TABLE_NAME_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required)""" + def __init__( + self, + *, + table_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["table_name", b"table_name"] + ) -> None: ... + +global___RefreshTable = RefreshTable + +class RefreshByPath(google.protobuf.message.Message): + """See `spark.catalog.refreshByPath`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PATH_FIELD_NUMBER: builtins.int + path: builtins.str + """(Required)""" + def __init__( + self, + *, + path: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["path", b"path"]) -> None: ... + +global___RefreshByPath = RefreshByPath + +class CurrentCatalog(google.protobuf.message.Message): + """See `spark.catalog.currentCatalog`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___CurrentCatalog = CurrentCatalog + +class SetCurrentCatalog(google.protobuf.message.Message): + """See `spark.catalog.setCurrentCatalog`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CATALOG_NAME_FIELD_NUMBER: builtins.int + catalog_name: builtins.str + """(Required)""" + def __init__( + self, + *, + catalog_name: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["catalog_name", b"catalog_name"] + ) -> None: ... + +global___SetCurrentCatalog = SetCurrentCatalog + +class ListCatalogs(google.protobuf.message.Message): + """See `spark.catalog.listCatalogs`""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PATTERN_FIELD_NUMBER: builtins.int + pattern: builtins.str + """(Optional) The pattern that the catalog name needs to match""" + def __init__( + self, + *, + pattern: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_pattern", b"_pattern", "pattern", b"pattern"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["_pattern", b"_pattern", "pattern", b"pattern"] + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_pattern", b"_pattern"] + ) -> typing_extensions.Literal["pattern"] | None: ... + +global___ListCatalogs = ListCatalogs diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.py new file mode 100644 index 0000000..90911e3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.py @@ -0,0 +1,141 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/commands.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2 +from pyspark.sql.connect.proto import common_pb2 as spark_dot_connect_dot_common__pb2 +from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2 +from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1cspark/connect/commands.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto"\xf5\x07\n\x07\x43ommand\x12]\n\x11register_function\x18\x01 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionH\x00R\x10registerFunction\x12H\n\x0fwrite_operation\x18\x02 \x01(\x0b\x32\x1d.spark.connect.WriteOperationH\x00R\x0ewriteOperation\x12_\n\x15\x63reate_dataframe_view\x18\x03 \x01(\x0b\x32).spark.connect.CreateDataFrameViewCommandH\x00R\x13\x63reateDataframeView\x12O\n\x12write_operation_v2\x18\x04 \x01(\x0b\x32\x1f.spark.connect.WriteOperationV2H\x00R\x10writeOperationV2\x12<\n\x0bsql_command\x18\x05 \x01(\x0b\x32\x19.spark.connect.SqlCommandH\x00R\nsqlCommand\x12k\n\x1cwrite_stream_operation_start\x18\x06 \x01(\x0b\x32(.spark.connect.WriteStreamOperationStartH\x00R\x19writeStreamOperationStart\x12^\n\x17streaming_query_command\x18\x07 \x01(\x0b\x32$.spark.connect.StreamingQueryCommandH\x00R\x15streamingQueryCommand\x12X\n\x15get_resources_command\x18\x08 \x01(\x0b\x32".spark.connect.GetResourcesCommandH\x00R\x13getResourcesCommand\x12t\n\x1fstreaming_query_manager_command\x18\t \x01(\x0b\x32+.spark.connect.StreamingQueryManagerCommandH\x00R\x1cstreamingQueryManagerCommand\x12m\n\x17register_table_function\x18\n \x01(\x0b\x32\x33.spark.connect.CommonInlineUserDefinedTableFunctionH\x00R\x15registerTableFunction\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textensionB\x0e\n\x0c\x63ommand_type"\xf1\x01\n\nSqlCommand\x12\x10\n\x03sql\x18\x01 \x01(\tR\x03sql\x12\x37\n\x04\x61rgs\x18\x02 \x03(\x0b\x32#.spark.connect.SqlCommand.ArgsEntryR\x04\x61rgs\x12<\n\x08pos_args\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x07posArgs\x1aZ\n\tArgsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x05value:\x02\x38\x01"\x96\x01\n\x1a\x43reateDataFrameViewCommand\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x1b\n\tis_global\x18\x03 \x01(\x08R\x08isGlobal\x12\x18\n\x07replace\x18\x04 \x01(\x08R\x07replace"\x9b\x08\n\x0eWriteOperation\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1b\n\x06source\x18\x02 \x01(\tH\x01R\x06source\x88\x01\x01\x12\x14\n\x04path\x18\x03 \x01(\tH\x00R\x04path\x12?\n\x05table\x18\x04 \x01(\x0b\x32\'.spark.connect.WriteOperation.SaveTableH\x00R\x05table\x12:\n\x04mode\x18\x05 \x01(\x0e\x32&.spark.connect.WriteOperation.SaveModeR\x04mode\x12*\n\x11sort_column_names\x18\x06 \x03(\tR\x0fsortColumnNames\x12\x31\n\x14partitioning_columns\x18\x07 \x03(\tR\x13partitioningColumns\x12\x43\n\tbucket_by\x18\x08 \x01(\x0b\x32&.spark.connect.WriteOperation.BucketByR\x08\x62ucketBy\x12\x44\n\x07options\x18\t \x03(\x0b\x32*.spark.connect.WriteOperation.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\x82\x02\n\tSaveTable\x12\x1d\n\ntable_name\x18\x01 \x01(\tR\ttableName\x12X\n\x0bsave_method\x18\x02 \x01(\x0e\x32\x37.spark.connect.WriteOperation.SaveTable.TableSaveMethodR\nsaveMethod"|\n\x0fTableSaveMethod\x12!\n\x1dTABLE_SAVE_METHOD_UNSPECIFIED\x10\x00\x12#\n\x1fTABLE_SAVE_METHOD_SAVE_AS_TABLE\x10\x01\x12!\n\x1dTABLE_SAVE_METHOD_INSERT_INTO\x10\x02\x1a[\n\x08\x42ucketBy\x12.\n\x13\x62ucket_column_names\x18\x01 \x03(\tR\x11\x62ucketColumnNames\x12\x1f\n\x0bnum_buckets\x18\x02 \x01(\x05R\nnumBuckets"\x89\x01\n\x08SaveMode\x12\x19\n\x15SAVE_MODE_UNSPECIFIED\x10\x00\x12\x14\n\x10SAVE_MODE_APPEND\x10\x01\x12\x17\n\x13SAVE_MODE_OVERWRITE\x10\x02\x12\x1d\n\x19SAVE_MODE_ERROR_IF_EXISTS\x10\x03\x12\x14\n\x10SAVE_MODE_IGNORE\x10\x04\x42\x0b\n\tsave_typeB\t\n\x07_source"\xad\x06\n\x10WriteOperationV2\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1d\n\ntable_name\x18\x02 \x01(\tR\ttableName\x12\x1f\n\x08provider\x18\x03 \x01(\tH\x00R\x08provider\x88\x01\x01\x12L\n\x14partitioning_columns\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13partitioningColumns\x12\x46\n\x07options\x18\x05 \x03(\x0b\x32,.spark.connect.WriteOperationV2.OptionsEntryR\x07options\x12_\n\x10table_properties\x18\x06 \x03(\x0b\x32\x34.spark.connect.WriteOperationV2.TablePropertiesEntryR\x0ftableProperties\x12\x38\n\x04mode\x18\x07 \x01(\x0e\x32$.spark.connect.WriteOperationV2.ModeR\x04mode\x12J\n\x13overwrite_condition\x18\x08 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x12overwriteCondition\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\x42\n\x14TablePropertiesEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"\x9f\x01\n\x04Mode\x12\x14\n\x10MODE_UNSPECIFIED\x10\x00\x12\x0f\n\x0bMODE_CREATE\x10\x01\x12\x12\n\x0eMODE_OVERWRITE\x10\x02\x12\x1d\n\x19MODE_OVERWRITE_PARTITIONS\x10\x03\x12\x0f\n\x0bMODE_APPEND\x10\x04\x12\x10\n\x0cMODE_REPLACE\x10\x05\x12\x1a\n\x16MODE_CREATE_OR_REPLACE\x10\x06\x42\x0b\n\t_provider"\xa0\x06\n\x19WriteStreamOperationStart\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x16\n\x06\x66ormat\x18\x02 \x01(\tR\x06\x66ormat\x12O\n\x07options\x18\x03 \x03(\x0b\x32\x35.spark.connect.WriteStreamOperationStart.OptionsEntryR\x07options\x12:\n\x19partitioning_column_names\x18\x04 \x03(\tR\x17partitioningColumnNames\x12:\n\x18processing_time_interval\x18\x05 \x01(\tH\x00R\x16processingTimeInterval\x12%\n\ravailable_now\x18\x06 \x01(\x08H\x00R\x0c\x61vailableNow\x12\x14\n\x04once\x18\x07 \x01(\x08H\x00R\x04once\x12\x46\n\x1e\x63ontinuous_checkpoint_interval\x18\x08 \x01(\tH\x00R\x1c\x63ontinuousCheckpointInterval\x12\x1f\n\x0boutput_mode\x18\t \x01(\tR\noutputMode\x12\x1d\n\nquery_name\x18\n \x01(\tR\tqueryName\x12\x14\n\x04path\x18\x0b \x01(\tH\x01R\x04path\x12\x1f\n\ntable_name\x18\x0c \x01(\tH\x01R\ttableName\x12N\n\x0e\x66oreach_writer\x18\r \x01(\x0b\x32\'.spark.connect.StreamingForeachFunctionR\rforeachWriter\x12L\n\rforeach_batch\x18\x0e \x01(\x0b\x32\'.spark.connect.StreamingForeachFunctionR\x0c\x66oreachBatch\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\t\n\x07triggerB\x12\n\x10sink_destination"\xb3\x01\n\x18StreamingForeachFunction\x12\x43\n\x0fpython_function\x18\x01 \x01(\x0b\x32\x18.spark.connect.PythonUDFH\x00R\x0epythonFunction\x12\x46\n\x0escala_function\x18\x02 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFH\x00R\rscalaFunctionB\n\n\x08\x66unction"y\n\x1fWriteStreamOperationStartResult\x12\x42\n\x08query_id\x18\x01 \x01(\x0b\x32\'.spark.connect.StreamingQueryInstanceIdR\x07queryId\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name"A\n\x18StreamingQueryInstanceId\x12\x0e\n\x02id\x18\x01 \x01(\tR\x02id\x12\x15\n\x06run_id\x18\x02 \x01(\tR\x05runId"\xf8\x04\n\x15StreamingQueryCommand\x12\x42\n\x08query_id\x18\x01 \x01(\x0b\x32\'.spark.connect.StreamingQueryInstanceIdR\x07queryId\x12\x18\n\x06status\x18\x02 \x01(\x08H\x00R\x06status\x12%\n\rlast_progress\x18\x03 \x01(\x08H\x00R\x0clastProgress\x12)\n\x0frecent_progress\x18\x04 \x01(\x08H\x00R\x0erecentProgress\x12\x14\n\x04stop\x18\x05 \x01(\x08H\x00R\x04stop\x12\x34\n\x15process_all_available\x18\x06 \x01(\x08H\x00R\x13processAllAvailable\x12O\n\x07\x65xplain\x18\x07 \x01(\x0b\x32\x33.spark.connect.StreamingQueryCommand.ExplainCommandH\x00R\x07\x65xplain\x12\x1e\n\texception\x18\x08 \x01(\x08H\x00R\texception\x12k\n\x11\x61wait_termination\x18\t \x01(\x0b\x32<.spark.connect.StreamingQueryCommand.AwaitTerminationCommandH\x00R\x10\x61waitTermination\x1a,\n\x0e\x45xplainCommand\x12\x1a\n\x08\x65xtended\x18\x01 \x01(\x08R\x08\x65xtended\x1aL\n\x17\x41waitTerminationCommand\x12"\n\ntimeout_ms\x18\x02 \x01(\x03H\x00R\ttimeoutMs\x88\x01\x01\x42\r\n\x0b_timeout_msB\t\n\x07\x63ommand"\xf5\x08\n\x1bStreamingQueryCommandResult\x12\x42\n\x08query_id\x18\x01 \x01(\x0b\x32\'.spark.connect.StreamingQueryInstanceIdR\x07queryId\x12Q\n\x06status\x18\x02 \x01(\x0b\x32\x37.spark.connect.StreamingQueryCommandResult.StatusResultH\x00R\x06status\x12j\n\x0frecent_progress\x18\x03 \x01(\x0b\x32?.spark.connect.StreamingQueryCommandResult.RecentProgressResultH\x00R\x0erecentProgress\x12T\n\x07\x65xplain\x18\x04 \x01(\x0b\x32\x38.spark.connect.StreamingQueryCommandResult.ExplainResultH\x00R\x07\x65xplain\x12Z\n\texception\x18\x05 \x01(\x0b\x32:.spark.connect.StreamingQueryCommandResult.ExceptionResultH\x00R\texception\x12p\n\x11\x61wait_termination\x18\x06 \x01(\x0b\x32\x41.spark.connect.StreamingQueryCommandResult.AwaitTerminationResultH\x00R\x10\x61waitTermination\x1a\xaa\x01\n\x0cStatusResult\x12%\n\x0estatus_message\x18\x01 \x01(\tR\rstatusMessage\x12*\n\x11is_data_available\x18\x02 \x01(\x08R\x0fisDataAvailable\x12*\n\x11is_trigger_active\x18\x03 \x01(\x08R\x0fisTriggerActive\x12\x1b\n\tis_active\x18\x04 \x01(\x08R\x08isActive\x1aH\n\x14RecentProgressResult\x12\x30\n\x14recent_progress_json\x18\x05 \x03(\tR\x12recentProgressJson\x1a\'\n\rExplainResult\x12\x16\n\x06result\x18\x01 \x01(\tR\x06result\x1a\xc5\x01\n\x0f\x45xceptionResult\x12\x30\n\x11\x65xception_message\x18\x01 \x01(\tH\x00R\x10\x65xceptionMessage\x88\x01\x01\x12$\n\x0b\x65rror_class\x18\x02 \x01(\tH\x01R\nerrorClass\x88\x01\x01\x12$\n\x0bstack_trace\x18\x03 \x01(\tH\x02R\nstackTrace\x88\x01\x01\x42\x14\n\x12_exception_messageB\x0e\n\x0c_error_classB\x0e\n\x0c_stack_trace\x1a\x38\n\x16\x41waitTerminationResult\x12\x1e\n\nterminated\x18\x01 \x01(\x08R\nterminatedB\r\n\x0bresult_type"\xbd\x06\n\x1cStreamingQueryManagerCommand\x12\x18\n\x06\x61\x63tive\x18\x01 \x01(\x08H\x00R\x06\x61\x63tive\x12\x1d\n\tget_query\x18\x02 \x01(\tH\x00R\x08getQuery\x12|\n\x15\x61wait_any_termination\x18\x03 \x01(\x0b\x32\x46.spark.connect.StreamingQueryManagerCommand.AwaitAnyTerminationCommandH\x00R\x13\x61waitAnyTermination\x12+\n\x10reset_terminated\x18\x04 \x01(\x08H\x00R\x0fresetTerminated\x12n\n\x0c\x61\x64\x64_listener\x18\x05 \x01(\x0b\x32I.spark.connect.StreamingQueryManagerCommand.StreamingQueryListenerCommandH\x00R\x0b\x61\x64\x64Listener\x12t\n\x0fremove_listener\x18\x06 \x01(\x0b\x32I.spark.connect.StreamingQueryManagerCommand.StreamingQueryListenerCommandH\x00R\x0eremoveListener\x12\'\n\x0elist_listeners\x18\x07 \x01(\x08H\x00R\rlistListeners\x1aO\n\x1a\x41waitAnyTerminationCommand\x12"\n\ntimeout_ms\x18\x01 \x01(\x03H\x00R\ttimeoutMs\x88\x01\x01\x42\r\n\x0b_timeout_ms\x1a\xcd\x01\n\x1dStreamingQueryListenerCommand\x12)\n\x10listener_payload\x18\x01 \x01(\x0cR\x0flistenerPayload\x12U\n\x17python_listener_payload\x18\x02 \x01(\x0b\x32\x18.spark.connect.PythonUDFH\x00R\x15pythonListenerPayload\x88\x01\x01\x12\x0e\n\x02id\x18\x03 \x01(\tR\x02idB\x1a\n\x18_python_listener_payloadB\t\n\x07\x63ommand"\xb4\x08\n"StreamingQueryManagerCommandResult\x12X\n\x06\x61\x63tive\x18\x01 \x01(\x0b\x32>.spark.connect.StreamingQueryManagerCommandResult.ActiveResultH\x00R\x06\x61\x63tive\x12`\n\x05query\x18\x02 \x01(\x0b\x32H.spark.connect.StreamingQueryManagerCommandResult.StreamingQueryInstanceH\x00R\x05query\x12\x81\x01\n\x15\x61wait_any_termination\x18\x03 \x01(\x0b\x32K.spark.connect.StreamingQueryManagerCommandResult.AwaitAnyTerminationResultH\x00R\x13\x61waitAnyTermination\x12+\n\x10reset_terminated\x18\x04 \x01(\x08H\x00R\x0fresetTerminated\x12#\n\x0c\x61\x64\x64_listener\x18\x05 \x01(\x08H\x00R\x0b\x61\x64\x64Listener\x12)\n\x0fremove_listener\x18\x06 \x01(\x08H\x00R\x0eremoveListener\x12{\n\x0elist_listeners\x18\x07 \x01(\x0b\x32R.spark.connect.StreamingQueryManagerCommandResult.ListStreamingQueryListenerResultH\x00R\rlistListeners\x1a\x7f\n\x0c\x41\x63tiveResult\x12o\n\x0e\x61\x63tive_queries\x18\x01 \x03(\x0b\x32H.spark.connect.StreamingQueryManagerCommandResult.StreamingQueryInstanceR\ractiveQueries\x1as\n\x16StreamingQueryInstance\x12\x37\n\x02id\x18\x01 \x01(\x0b\x32\'.spark.connect.StreamingQueryInstanceIdR\x02id\x12\x17\n\x04name\x18\x02 \x01(\tH\x00R\x04name\x88\x01\x01\x42\x07\n\x05_name\x1a;\n\x19\x41waitAnyTerminationResult\x12\x1e\n\nterminated\x18\x01 \x01(\x08R\nterminated\x1aK\n\x1eStreamingQueryListenerInstance\x12)\n\x10listener_payload\x18\x01 \x01(\x0cR\x0flistenerPayload\x1a\x45\n ListStreamingQueryListenerResult\x12!\n\x0clistener_ids\x18\x01 \x03(\tR\x0blistenerIdsB\r\n\x0bresult_type"\x15\n\x13GetResourcesCommand"\xd4\x01\n\x19GetResourcesCommandResult\x12U\n\tresources\x18\x01 \x03(\x0b\x32\x37.spark.connect.GetResourcesCommandResult.ResourcesEntryR\tresources\x1a`\n\x0eResourcesEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x38\n\x05value\x18\x02 \x01(\x0b\x32".spark.connect.ResourceInformationR\x05value:\x02\x38\x01\x42\x36\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.commands_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _SQLCOMMAND_ARGSENTRY._options = None + _SQLCOMMAND_ARGSENTRY._serialized_options = b"8\001" + _WRITEOPERATION_OPTIONSENTRY._options = None + _WRITEOPERATION_OPTIONSENTRY._serialized_options = b"8\001" + _WRITEOPERATIONV2_OPTIONSENTRY._options = None + _WRITEOPERATIONV2_OPTIONSENTRY._serialized_options = b"8\001" + _WRITEOPERATIONV2_TABLEPROPERTIESENTRY._options = None + _WRITEOPERATIONV2_TABLEPROPERTIESENTRY._serialized_options = b"8\001" + _WRITESTREAMOPERATIONSTART_OPTIONSENTRY._options = None + _WRITESTREAMOPERATIONSTART_OPTIONSENTRY._serialized_options = b"8\001" + _GETRESOURCESCOMMANDRESULT_RESOURCESENTRY._options = None + _GETRESOURCESCOMMANDRESULT_RESOURCESENTRY._serialized_options = b"8\001" + _COMMAND._serialized_start = 167 + _COMMAND._serialized_end = 1180 + _SQLCOMMAND._serialized_start = 1183 + _SQLCOMMAND._serialized_end = 1424 + _SQLCOMMAND_ARGSENTRY._serialized_start = 1334 + _SQLCOMMAND_ARGSENTRY._serialized_end = 1424 + _CREATEDATAFRAMEVIEWCOMMAND._serialized_start = 1427 + _CREATEDATAFRAMEVIEWCOMMAND._serialized_end = 1577 + _WRITEOPERATION._serialized_start = 1580 + _WRITEOPERATION._serialized_end = 2631 + _WRITEOPERATION_OPTIONSENTRY._serialized_start = 2055 + _WRITEOPERATION_OPTIONSENTRY._serialized_end = 2113 + _WRITEOPERATION_SAVETABLE._serialized_start = 2116 + _WRITEOPERATION_SAVETABLE._serialized_end = 2374 + _WRITEOPERATION_SAVETABLE_TABLESAVEMETHOD._serialized_start = 2250 + _WRITEOPERATION_SAVETABLE_TABLESAVEMETHOD._serialized_end = 2374 + _WRITEOPERATION_BUCKETBY._serialized_start = 2376 + _WRITEOPERATION_BUCKETBY._serialized_end = 2467 + _WRITEOPERATION_SAVEMODE._serialized_start = 2470 + _WRITEOPERATION_SAVEMODE._serialized_end = 2607 + _WRITEOPERATIONV2._serialized_start = 2634 + _WRITEOPERATIONV2._serialized_end = 3447 + _WRITEOPERATIONV2_OPTIONSENTRY._serialized_start = 2055 + _WRITEOPERATIONV2_OPTIONSENTRY._serialized_end = 2113 + _WRITEOPERATIONV2_TABLEPROPERTIESENTRY._serialized_start = 3206 + _WRITEOPERATIONV2_TABLEPROPERTIESENTRY._serialized_end = 3272 + _WRITEOPERATIONV2_MODE._serialized_start = 3275 + _WRITEOPERATIONV2_MODE._serialized_end = 3434 + _WRITESTREAMOPERATIONSTART._serialized_start = 3450 + _WRITESTREAMOPERATIONSTART._serialized_end = 4250 + _WRITESTREAMOPERATIONSTART_OPTIONSENTRY._serialized_start = 2055 + _WRITESTREAMOPERATIONSTART_OPTIONSENTRY._serialized_end = 2113 + _STREAMINGFOREACHFUNCTION._serialized_start = 4253 + _STREAMINGFOREACHFUNCTION._serialized_end = 4432 + _WRITESTREAMOPERATIONSTARTRESULT._serialized_start = 4434 + _WRITESTREAMOPERATIONSTARTRESULT._serialized_end = 4555 + _STREAMINGQUERYINSTANCEID._serialized_start = 4557 + _STREAMINGQUERYINSTANCEID._serialized_end = 4622 + _STREAMINGQUERYCOMMAND._serialized_start = 4625 + _STREAMINGQUERYCOMMAND._serialized_end = 5257 + _STREAMINGQUERYCOMMAND_EXPLAINCOMMAND._serialized_start = 5124 + _STREAMINGQUERYCOMMAND_EXPLAINCOMMAND._serialized_end = 5168 + _STREAMINGQUERYCOMMAND_AWAITTERMINATIONCOMMAND._serialized_start = 5170 + _STREAMINGQUERYCOMMAND_AWAITTERMINATIONCOMMAND._serialized_end = 5246 + _STREAMINGQUERYCOMMANDRESULT._serialized_start = 5260 + _STREAMINGQUERYCOMMANDRESULT._serialized_end = 6401 + _STREAMINGQUERYCOMMANDRESULT_STATUSRESULT._serialized_start = 5843 + _STREAMINGQUERYCOMMANDRESULT_STATUSRESULT._serialized_end = 6013 + _STREAMINGQUERYCOMMANDRESULT_RECENTPROGRESSRESULT._serialized_start = 6015 + _STREAMINGQUERYCOMMANDRESULT_RECENTPROGRESSRESULT._serialized_end = 6087 + _STREAMINGQUERYCOMMANDRESULT_EXPLAINRESULT._serialized_start = 6089 + _STREAMINGQUERYCOMMANDRESULT_EXPLAINRESULT._serialized_end = 6128 + _STREAMINGQUERYCOMMANDRESULT_EXCEPTIONRESULT._serialized_start = 6131 + _STREAMINGQUERYCOMMANDRESULT_EXCEPTIONRESULT._serialized_end = 6328 + _STREAMINGQUERYCOMMANDRESULT_AWAITTERMINATIONRESULT._serialized_start = 6330 + _STREAMINGQUERYCOMMANDRESULT_AWAITTERMINATIONRESULT._serialized_end = 6386 + _STREAMINGQUERYMANAGERCOMMAND._serialized_start = 6404 + _STREAMINGQUERYMANAGERCOMMAND._serialized_end = 7233 + _STREAMINGQUERYMANAGERCOMMAND_AWAITANYTERMINATIONCOMMAND._serialized_start = 6935 + _STREAMINGQUERYMANAGERCOMMAND_AWAITANYTERMINATIONCOMMAND._serialized_end = 7014 + _STREAMINGQUERYMANAGERCOMMAND_STREAMINGQUERYLISTENERCOMMAND._serialized_start = 7017 + _STREAMINGQUERYMANAGERCOMMAND_STREAMINGQUERYLISTENERCOMMAND._serialized_end = 7222 + _STREAMINGQUERYMANAGERCOMMANDRESULT._serialized_start = 7236 + _STREAMINGQUERYMANAGERCOMMANDRESULT._serialized_end = 8312 + _STREAMINGQUERYMANAGERCOMMANDRESULT_ACTIVERESULT._serialized_start = 7844 + _STREAMINGQUERYMANAGERCOMMANDRESULT_ACTIVERESULT._serialized_end = 7971 + _STREAMINGQUERYMANAGERCOMMANDRESULT_STREAMINGQUERYINSTANCE._serialized_start = 7973 + _STREAMINGQUERYMANAGERCOMMANDRESULT_STREAMINGQUERYINSTANCE._serialized_end = 8088 + _STREAMINGQUERYMANAGERCOMMANDRESULT_AWAITANYTERMINATIONRESULT._serialized_start = 8090 + _STREAMINGQUERYMANAGERCOMMANDRESULT_AWAITANYTERMINATIONRESULT._serialized_end = 8149 + _STREAMINGQUERYMANAGERCOMMANDRESULT_STREAMINGQUERYLISTENERINSTANCE._serialized_start = 8151 + _STREAMINGQUERYMANAGERCOMMANDRESULT_STREAMINGQUERYLISTENERINSTANCE._serialized_end = 8226 + _STREAMINGQUERYMANAGERCOMMANDRESULT_LISTSTREAMINGQUERYLISTENERRESULT._serialized_start = 8228 + _STREAMINGQUERYMANAGERCOMMANDRESULT_LISTSTREAMINGQUERYLISTENERRESULT._serialized_end = 8297 + _GETRESOURCESCOMMAND._serialized_start = 8314 + _GETRESOURCESCOMMAND._serialized_end = 8335 + _GETRESOURCESCOMMANDRESULT._serialized_start = 8338 + _GETRESOURCESCOMMANDRESULT._serialized_end = 8550 + _GETRESOURCESCOMMANDRESULT_RESOURCESENTRY._serialized_start = 8454 + _GETRESOURCESCOMMANDRESULT_RESOURCESENTRY._serialized_end = 8550 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.pyi new file mode 100644 index 0000000..f3dca7a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/commands_pb2.pyi @@ -0,0 +1,1766 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.any_pb2 +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.internal.enum_type_wrapper +import google.protobuf.message +import pyspark.sql.connect.proto.common_pb2 +import pyspark.sql.connect.proto.expressions_pb2 +import pyspark.sql.connect.proto.relations_pb2 +import sys +import typing + +if sys.version_info >= (3, 10): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class Command(google.protobuf.message.Message): + """A [[Command]] is an operation that is executed by the server that does not directly consume or + produce a relational result. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + REGISTER_FUNCTION_FIELD_NUMBER: builtins.int + WRITE_OPERATION_FIELD_NUMBER: builtins.int + CREATE_DATAFRAME_VIEW_FIELD_NUMBER: builtins.int + WRITE_OPERATION_V2_FIELD_NUMBER: builtins.int + SQL_COMMAND_FIELD_NUMBER: builtins.int + WRITE_STREAM_OPERATION_START_FIELD_NUMBER: builtins.int + STREAMING_QUERY_COMMAND_FIELD_NUMBER: builtins.int + GET_RESOURCES_COMMAND_FIELD_NUMBER: builtins.int + STREAMING_QUERY_MANAGER_COMMAND_FIELD_NUMBER: builtins.int + REGISTER_TABLE_FUNCTION_FIELD_NUMBER: builtins.int + EXTENSION_FIELD_NUMBER: builtins.int + @property + def register_function( + self, + ) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: ... + @property + def write_operation(self) -> global___WriteOperation: ... + @property + def create_dataframe_view(self) -> global___CreateDataFrameViewCommand: ... + @property + def write_operation_v2(self) -> global___WriteOperationV2: ... + @property + def sql_command(self) -> global___SqlCommand: ... + @property + def write_stream_operation_start(self) -> global___WriteStreamOperationStart: ... + @property + def streaming_query_command(self) -> global___StreamingQueryCommand: ... + @property + def get_resources_command(self) -> global___GetResourcesCommand: ... + @property + def streaming_query_manager_command(self) -> global___StreamingQueryManagerCommand: ... + @property + def register_table_function( + self, + ) -> pyspark.sql.connect.proto.relations_pb2.CommonInlineUserDefinedTableFunction: ... + @property + def extension(self) -> google.protobuf.any_pb2.Any: + """This field is used to mark extensions to the protocol. When plugins generate arbitrary + Commands they can add them here. During the planning the correct resolution is done. + """ + def __init__( + self, + *, + register_function: pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction + | None = ..., + write_operation: global___WriteOperation | None = ..., + create_dataframe_view: global___CreateDataFrameViewCommand | None = ..., + write_operation_v2: global___WriteOperationV2 | None = ..., + sql_command: global___SqlCommand | None = ..., + write_stream_operation_start: global___WriteStreamOperationStart | None = ..., + streaming_query_command: global___StreamingQueryCommand | None = ..., + get_resources_command: global___GetResourcesCommand | None = ..., + streaming_query_manager_command: global___StreamingQueryManagerCommand | None = ..., + register_table_function: pyspark.sql.connect.proto.relations_pb2.CommonInlineUserDefinedTableFunction + | None = ..., + extension: google.protobuf.any_pb2.Any | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "command_type", + b"command_type", + "create_dataframe_view", + b"create_dataframe_view", + "extension", + b"extension", + "get_resources_command", + b"get_resources_command", + "register_function", + b"register_function", + "register_table_function", + b"register_table_function", + "sql_command", + b"sql_command", + "streaming_query_command", + b"streaming_query_command", + "streaming_query_manager_command", + b"streaming_query_manager_command", + "write_operation", + b"write_operation", + "write_operation_v2", + b"write_operation_v2", + "write_stream_operation_start", + b"write_stream_operation_start", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "command_type", + b"command_type", + "create_dataframe_view", + b"create_dataframe_view", + "extension", + b"extension", + "get_resources_command", + b"get_resources_command", + "register_function", + b"register_function", + "register_table_function", + b"register_table_function", + "sql_command", + b"sql_command", + "streaming_query_command", + b"streaming_query_command", + "streaming_query_manager_command", + b"streaming_query_manager_command", + "write_operation", + b"write_operation", + "write_operation_v2", + b"write_operation_v2", + "write_stream_operation_start", + b"write_stream_operation_start", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["command_type", b"command_type"] + ) -> typing_extensions.Literal[ + "register_function", + "write_operation", + "create_dataframe_view", + "write_operation_v2", + "sql_command", + "write_stream_operation_start", + "streaming_query_command", + "get_resources_command", + "streaming_query_manager_command", + "register_table_function", + "extension", + ] | None: ... + +global___Command = Command + +class SqlCommand(google.protobuf.message.Message): + """A SQL Command is used to trigger the eager evaluation of SQL commands in Spark. + + When the SQL provide as part of the message is a command it will be immediately evaluated + and the result will be collected and returned as part of a LocalRelation. If the result is + not a command, the operation will simply return a SQL Relation. This allows the client to be + almost oblivious to the server-side behavior. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ArgsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + @property + def value(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression.Literal: ... + def __init__( + self, + *, + key: builtins.str = ..., + value: pyspark.sql.connect.proto.expressions_pb2.Expression.Literal | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + SQL_FIELD_NUMBER: builtins.int + ARGS_FIELD_NUMBER: builtins.int + POS_ARGS_FIELD_NUMBER: builtins.int + sql: builtins.str + """(Required) SQL Query.""" + @property + def args( + self, + ) -> google.protobuf.internal.containers.MessageMap[ + builtins.str, pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Optional) A map of parameter names to literal expressions.""" + @property + def pos_args( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Optional) A sequence of literal expressions for positional parameters in the SQL query text.""" + def __init__( + self, + *, + sql: builtins.str = ..., + args: collections.abc.Mapping[ + builtins.str, pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + pos_args: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "args", b"args", "pos_args", b"pos_args", "sql", b"sql" + ], + ) -> None: ... + +global___SqlCommand = SqlCommand + +class CreateDataFrameViewCommand(google.protobuf.message.Message): + """A command that can create DataFrame global temp view or local temp view.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + IS_GLOBAL_FIELD_NUMBER: builtins.int + REPLACE_FIELD_NUMBER: builtins.int + @property + def input(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The relation that this view will be built on.""" + name: builtins.str + """(Required) View name.""" + is_global: builtins.bool + """(Required) Whether this is global temp view or local temp view.""" + replace: builtins.bool + """(Required) + + If true, and if the view already exists, updates it; if false, and if the view + already exists, throws exception. + """ + def __init__( + self, + *, + input: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + name: builtins.str = ..., + is_global: builtins.bool = ..., + replace: builtins.bool = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", b"input", "is_global", b"is_global", "name", b"name", "replace", b"replace" + ], + ) -> None: ... + +global___CreateDataFrameViewCommand = CreateDataFrameViewCommand + +class WriteOperation(google.protobuf.message.Message): + """As writes are not directly handled during analysis and planning, they are modeled as commands.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _SaveMode: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _SaveModeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + WriteOperation._SaveMode.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + SAVE_MODE_UNSPECIFIED: WriteOperation._SaveMode.ValueType # 0 + SAVE_MODE_APPEND: WriteOperation._SaveMode.ValueType # 1 + SAVE_MODE_OVERWRITE: WriteOperation._SaveMode.ValueType # 2 + SAVE_MODE_ERROR_IF_EXISTS: WriteOperation._SaveMode.ValueType # 3 + SAVE_MODE_IGNORE: WriteOperation._SaveMode.ValueType # 4 + + class SaveMode(_SaveMode, metaclass=_SaveModeEnumTypeWrapper): ... + SAVE_MODE_UNSPECIFIED: WriteOperation.SaveMode.ValueType # 0 + SAVE_MODE_APPEND: WriteOperation.SaveMode.ValueType # 1 + SAVE_MODE_OVERWRITE: WriteOperation.SaveMode.ValueType # 2 + SAVE_MODE_ERROR_IF_EXISTS: WriteOperation.SaveMode.ValueType # 3 + SAVE_MODE_IGNORE: WriteOperation.SaveMode.ValueType # 4 + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + class SaveTable(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _TableSaveMethod: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _TableSaveMethodEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + WriteOperation.SaveTable._TableSaveMethod.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + TABLE_SAVE_METHOD_UNSPECIFIED: WriteOperation.SaveTable._TableSaveMethod.ValueType # 0 + TABLE_SAVE_METHOD_SAVE_AS_TABLE: WriteOperation.SaveTable._TableSaveMethod.ValueType # 1 + TABLE_SAVE_METHOD_INSERT_INTO: WriteOperation.SaveTable._TableSaveMethod.ValueType # 2 + + class TableSaveMethod(_TableSaveMethod, metaclass=_TableSaveMethodEnumTypeWrapper): ... + TABLE_SAVE_METHOD_UNSPECIFIED: WriteOperation.SaveTable.TableSaveMethod.ValueType # 0 + TABLE_SAVE_METHOD_SAVE_AS_TABLE: WriteOperation.SaveTable.TableSaveMethod.ValueType # 1 + TABLE_SAVE_METHOD_INSERT_INTO: WriteOperation.SaveTable.TableSaveMethod.ValueType # 2 + + TABLE_NAME_FIELD_NUMBER: builtins.int + SAVE_METHOD_FIELD_NUMBER: builtins.int + table_name: builtins.str + """(Required) The table name.""" + save_method: global___WriteOperation.SaveTable.TableSaveMethod.ValueType + """(Required) The method to be called to write to the table.""" + def __init__( + self, + *, + table_name: builtins.str = ..., + save_method: global___WriteOperation.SaveTable.TableSaveMethod.ValueType = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "save_method", b"save_method", "table_name", b"table_name" + ], + ) -> None: ... + + class BucketBy(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + BUCKET_COLUMN_NAMES_FIELD_NUMBER: builtins.int + NUM_BUCKETS_FIELD_NUMBER: builtins.int + @property + def bucket_column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ... + num_buckets: builtins.int + def __init__( + self, + *, + bucket_column_names: collections.abc.Iterable[builtins.str] | None = ..., + num_buckets: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "bucket_column_names", b"bucket_column_names", "num_buckets", b"num_buckets" + ], + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + SOURCE_FIELD_NUMBER: builtins.int + PATH_FIELD_NUMBER: builtins.int + TABLE_FIELD_NUMBER: builtins.int + MODE_FIELD_NUMBER: builtins.int + SORT_COLUMN_NAMES_FIELD_NUMBER: builtins.int + PARTITIONING_COLUMNS_FIELD_NUMBER: builtins.int + BUCKET_BY_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + @property + def input(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The output of the `input` relation will be persisted according to the options.""" + source: builtins.str + """(Optional) Format value according to the Spark documentation. Examples are: text, parquet, delta.""" + path: builtins.str + @property + def table(self) -> global___WriteOperation.SaveTable: ... + mode: global___WriteOperation.SaveMode.ValueType + """(Required) the save mode.""" + @property + def sort_column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) List of columns to sort the output by.""" + @property + def partitioning_columns( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) List of columns for partitioning.""" + @property + def bucket_by(self) -> global___WriteOperation.BucketBy: + """(Optional) Bucketing specification. Bucketing must set the number of buckets and the columns + to bucket by. + """ + @property + def options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """(Optional) A list of configuration options.""" + def __init__( + self, + *, + input: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + source: builtins.str | None = ..., + path: builtins.str = ..., + table: global___WriteOperation.SaveTable | None = ..., + mode: global___WriteOperation.SaveMode.ValueType = ..., + sort_column_names: collections.abc.Iterable[builtins.str] | None = ..., + partitioning_columns: collections.abc.Iterable[builtins.str] | None = ..., + bucket_by: global___WriteOperation.BucketBy | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_source", + b"_source", + "bucket_by", + b"bucket_by", + "input", + b"input", + "path", + b"path", + "save_type", + b"save_type", + "source", + b"source", + "table", + b"table", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_source", + b"_source", + "bucket_by", + b"bucket_by", + "input", + b"input", + "mode", + b"mode", + "options", + b"options", + "partitioning_columns", + b"partitioning_columns", + "path", + b"path", + "save_type", + b"save_type", + "sort_column_names", + b"sort_column_names", + "source", + b"source", + "table", + b"table", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_source", b"_source"] + ) -> typing_extensions.Literal["source"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["save_type", b"save_type"] + ) -> typing_extensions.Literal["path", "table"] | None: ... + +global___WriteOperation = WriteOperation + +class WriteOperationV2(google.protobuf.message.Message): + """As writes are not directly handled during analysis and planning, they are modeled as commands.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _Mode: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _ModeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + WriteOperationV2._Mode.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + MODE_UNSPECIFIED: WriteOperationV2._Mode.ValueType # 0 + MODE_CREATE: WriteOperationV2._Mode.ValueType # 1 + MODE_OVERWRITE: WriteOperationV2._Mode.ValueType # 2 + MODE_OVERWRITE_PARTITIONS: WriteOperationV2._Mode.ValueType # 3 + MODE_APPEND: WriteOperationV2._Mode.ValueType # 4 + MODE_REPLACE: WriteOperationV2._Mode.ValueType # 5 + MODE_CREATE_OR_REPLACE: WriteOperationV2._Mode.ValueType # 6 + + class Mode(_Mode, metaclass=_ModeEnumTypeWrapper): ... + MODE_UNSPECIFIED: WriteOperationV2.Mode.ValueType # 0 + MODE_CREATE: WriteOperationV2.Mode.ValueType # 1 + MODE_OVERWRITE: WriteOperationV2.Mode.ValueType # 2 + MODE_OVERWRITE_PARTITIONS: WriteOperationV2.Mode.ValueType # 3 + MODE_APPEND: WriteOperationV2.Mode.ValueType # 4 + MODE_REPLACE: WriteOperationV2.Mode.ValueType # 5 + MODE_CREATE_OR_REPLACE: WriteOperationV2.Mode.ValueType # 6 + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + class TablePropertiesEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + TABLE_NAME_FIELD_NUMBER: builtins.int + PROVIDER_FIELD_NUMBER: builtins.int + PARTITIONING_COLUMNS_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + TABLE_PROPERTIES_FIELD_NUMBER: builtins.int + MODE_FIELD_NUMBER: builtins.int + OVERWRITE_CONDITION_FIELD_NUMBER: builtins.int + @property + def input(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The output of the `input` relation will be persisted according to the options.""" + table_name: builtins.str + """(Required) The destination of the write operation must be either a path or a table.""" + provider: builtins.str + """(Optional) A provider for the underlying output data source. Spark's default catalog supports + "parquet", "json", etc. + """ + @property + def partitioning_columns( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) List of columns for partitioning for output table created by `create`, + `createOrReplace`, or `replace` + """ + @property + def options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """(Optional) A list of configuration options.""" + @property + def table_properties( + self, + ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """(Optional) A list of table properties.""" + mode: global___WriteOperationV2.Mode.ValueType + """(Required) Write mode.""" + @property + def overwrite_condition(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: + """(Optional) A condition for overwrite saving mode""" + def __init__( + self, + *, + input: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + table_name: builtins.str = ..., + provider: builtins.str | None = ..., + partitioning_columns: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + table_properties: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + mode: global___WriteOperationV2.Mode.ValueType = ..., + overwrite_condition: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_provider", + b"_provider", + "input", + b"input", + "overwrite_condition", + b"overwrite_condition", + "provider", + b"provider", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_provider", + b"_provider", + "input", + b"input", + "mode", + b"mode", + "options", + b"options", + "overwrite_condition", + b"overwrite_condition", + "partitioning_columns", + b"partitioning_columns", + "provider", + b"provider", + "table_name", + b"table_name", + "table_properties", + b"table_properties", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_provider", b"_provider"] + ) -> typing_extensions.Literal["provider"] | None: ... + +global___WriteOperationV2 = WriteOperationV2 + +class WriteStreamOperationStart(google.protobuf.message.Message): + """Starts write stream operation as streaming query. Query ID and Run ID of the streaming + query are returned. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + FORMAT_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + PARTITIONING_COLUMN_NAMES_FIELD_NUMBER: builtins.int + PROCESSING_TIME_INTERVAL_FIELD_NUMBER: builtins.int + AVAILABLE_NOW_FIELD_NUMBER: builtins.int + ONCE_FIELD_NUMBER: builtins.int + CONTINUOUS_CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + OUTPUT_MODE_FIELD_NUMBER: builtins.int + QUERY_NAME_FIELD_NUMBER: builtins.int + PATH_FIELD_NUMBER: builtins.int + TABLE_NAME_FIELD_NUMBER: builtins.int + FOREACH_WRITER_FIELD_NUMBER: builtins.int + FOREACH_BATCH_FIELD_NUMBER: builtins.int + @property + def input(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: + """(Required) The output of the `input` streaming relation will be written.""" + format: builtins.str + """The following fields directly map to API for DataStreamWriter(). + Consult API documentation unless explicitly documented here. + """ + @property + def options( + self, + ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: ... + @property + def partitioning_column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ... + processing_time_interval: builtins.str + available_now: builtins.bool + once: builtins.bool + continuous_checkpoint_interval: builtins.str + output_mode: builtins.str + query_name: builtins.str + path: builtins.str + table_name: builtins.str + @property + def foreach_writer(self) -> global___StreamingForeachFunction: ... + @property + def foreach_batch(self) -> global___StreamingForeachFunction: ... + def __init__( + self, + *, + input: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + format: builtins.str = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + partitioning_column_names: collections.abc.Iterable[builtins.str] | None = ..., + processing_time_interval: builtins.str = ..., + available_now: builtins.bool = ..., + once: builtins.bool = ..., + continuous_checkpoint_interval: builtins.str = ..., + output_mode: builtins.str = ..., + query_name: builtins.str = ..., + path: builtins.str = ..., + table_name: builtins.str = ..., + foreach_writer: global___StreamingForeachFunction | None = ..., + foreach_batch: global___StreamingForeachFunction | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "available_now", + b"available_now", + "continuous_checkpoint_interval", + b"continuous_checkpoint_interval", + "foreach_batch", + b"foreach_batch", + "foreach_writer", + b"foreach_writer", + "input", + b"input", + "once", + b"once", + "path", + b"path", + "processing_time_interval", + b"processing_time_interval", + "sink_destination", + b"sink_destination", + "table_name", + b"table_name", + "trigger", + b"trigger", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "available_now", + b"available_now", + "continuous_checkpoint_interval", + b"continuous_checkpoint_interval", + "foreach_batch", + b"foreach_batch", + "foreach_writer", + b"foreach_writer", + "format", + b"format", + "input", + b"input", + "once", + b"once", + "options", + b"options", + "output_mode", + b"output_mode", + "partitioning_column_names", + b"partitioning_column_names", + "path", + b"path", + "processing_time_interval", + b"processing_time_interval", + "query_name", + b"query_name", + "sink_destination", + b"sink_destination", + "table_name", + b"table_name", + "trigger", + b"trigger", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["sink_destination", b"sink_destination"] + ) -> typing_extensions.Literal["path", "table_name"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["trigger", b"trigger"] + ) -> typing_extensions.Literal[ + "processing_time_interval", "available_now", "once", "continuous_checkpoint_interval" + ] | None: ... + +global___WriteStreamOperationStart = WriteStreamOperationStart + +class StreamingForeachFunction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PYTHON_FUNCTION_FIELD_NUMBER: builtins.int + SCALA_FUNCTION_FIELD_NUMBER: builtins.int + @property + def python_function(self) -> pyspark.sql.connect.proto.expressions_pb2.PythonUDF: ... + @property + def scala_function(self) -> pyspark.sql.connect.proto.expressions_pb2.ScalarScalaUDF: ... + def __init__( + self, + *, + python_function: pyspark.sql.connect.proto.expressions_pb2.PythonUDF | None = ..., + scala_function: pyspark.sql.connect.proto.expressions_pb2.ScalarScalaUDF | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "function", + b"function", + "python_function", + b"python_function", + "scala_function", + b"scala_function", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "function", + b"function", + "python_function", + b"python_function", + "scala_function", + b"scala_function", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["function", b"function"] + ) -> typing_extensions.Literal["python_function", "scala_function"] | None: ... + +global___StreamingForeachFunction = StreamingForeachFunction + +class WriteStreamOperationStartResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + QUERY_ID_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + @property + def query_id(self) -> global___StreamingQueryInstanceId: + """(Required) Query instance. See `StreamingQueryInstanceId`.""" + name: builtins.str + """An optional query name.""" + def __init__( + self, + *, + query_id: global___StreamingQueryInstanceId | None = ..., + name: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["query_id", b"query_id"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["name", b"name", "query_id", b"query_id"] + ) -> None: ... + +global___WriteStreamOperationStartResult = WriteStreamOperationStartResult + +class StreamingQueryInstanceId(google.protobuf.message.Message): + """A tuple that uniquely identifies an instance of streaming query run. It consists of `id` that + persists across the streaming runs and `run_id` that changes between each run of the + streaming query that resumes from the checkpoint. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ID_FIELD_NUMBER: builtins.int + RUN_ID_FIELD_NUMBER: builtins.int + id: builtins.str + """(Required) The unique id of this query that persists across restarts from checkpoint data. + That is, this id is generated when a query is started for the first time, and + will be the same every time it is restarted from checkpoint data. + """ + run_id: builtins.str + """(Required) The unique id of this run of the query. That is, every start/restart of a query + will generate a unique run_id. Therefore, every time a query is restarted from + checkpoint, it will have the same `id` but different `run_id`s. + """ + def __init__( + self, + *, + id: builtins.str = ..., + run_id: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["id", b"id", "run_id", b"run_id"] + ) -> None: ... + +global___StreamingQueryInstanceId = StreamingQueryInstanceId + +class StreamingQueryCommand(google.protobuf.message.Message): + """Commands for a streaming query.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ExplainCommand(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXTENDED_FIELD_NUMBER: builtins.int + extended: builtins.bool + """TODO: Consider reusing Explain from AnalyzePlanRequest message. + We can not do this right now since it base.proto imports this file. + """ + def __init__( + self, + *, + extended: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["extended", b"extended"] + ) -> None: ... + + class AwaitTerminationCommand(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TIMEOUT_MS_FIELD_NUMBER: builtins.int + timeout_ms: builtins.int + def __init__( + self, + *, + timeout_ms: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_timeout_ms", b"_timeout_ms", "timeout_ms", b"timeout_ms" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_timeout_ms", b"_timeout_ms", "timeout_ms", b"timeout_ms" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_timeout_ms", b"_timeout_ms"] + ) -> typing_extensions.Literal["timeout_ms"] | None: ... + + QUERY_ID_FIELD_NUMBER: builtins.int + STATUS_FIELD_NUMBER: builtins.int + LAST_PROGRESS_FIELD_NUMBER: builtins.int + RECENT_PROGRESS_FIELD_NUMBER: builtins.int + STOP_FIELD_NUMBER: builtins.int + PROCESS_ALL_AVAILABLE_FIELD_NUMBER: builtins.int + EXPLAIN_FIELD_NUMBER: builtins.int + EXCEPTION_FIELD_NUMBER: builtins.int + AWAIT_TERMINATION_FIELD_NUMBER: builtins.int + @property + def query_id(self) -> global___StreamingQueryInstanceId: + """(Required) Query instance. See `StreamingQueryInstanceId`.""" + status: builtins.bool + """status() API.""" + last_progress: builtins.bool + """lastProgress() API.""" + recent_progress: builtins.bool + """recentProgress() API.""" + stop: builtins.bool + """stop() API. Stops the query.""" + process_all_available: builtins.bool + """processAllAvailable() API. Waits till all the available data is processed""" + @property + def explain(self) -> global___StreamingQueryCommand.ExplainCommand: + """explain() API. Returns logical and physical plans.""" + exception: builtins.bool + """exception() API. Returns the exception in the query if any.""" + @property + def await_termination(self) -> global___StreamingQueryCommand.AwaitTerminationCommand: + """awaitTermination() API. Waits for the termination of the query.""" + def __init__( + self, + *, + query_id: global___StreamingQueryInstanceId | None = ..., + status: builtins.bool = ..., + last_progress: builtins.bool = ..., + recent_progress: builtins.bool = ..., + stop: builtins.bool = ..., + process_all_available: builtins.bool = ..., + explain: global___StreamingQueryCommand.ExplainCommand | None = ..., + exception: builtins.bool = ..., + await_termination: global___StreamingQueryCommand.AwaitTerminationCommand | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "await_termination", + b"await_termination", + "command", + b"command", + "exception", + b"exception", + "explain", + b"explain", + "last_progress", + b"last_progress", + "process_all_available", + b"process_all_available", + "query_id", + b"query_id", + "recent_progress", + b"recent_progress", + "status", + b"status", + "stop", + b"stop", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "await_termination", + b"await_termination", + "command", + b"command", + "exception", + b"exception", + "explain", + b"explain", + "last_progress", + b"last_progress", + "process_all_available", + b"process_all_available", + "query_id", + b"query_id", + "recent_progress", + b"recent_progress", + "status", + b"status", + "stop", + b"stop", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["command", b"command"] + ) -> typing_extensions.Literal[ + "status", + "last_progress", + "recent_progress", + "stop", + "process_all_available", + "explain", + "exception", + "await_termination", + ] | None: ... + +global___StreamingQueryCommand = StreamingQueryCommand + +class StreamingQueryCommandResult(google.protobuf.message.Message): + """Response for commands on a streaming query.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class StatusResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STATUS_MESSAGE_FIELD_NUMBER: builtins.int + IS_DATA_AVAILABLE_FIELD_NUMBER: builtins.int + IS_TRIGGER_ACTIVE_FIELD_NUMBER: builtins.int + IS_ACTIVE_FIELD_NUMBER: builtins.int + status_message: builtins.str + """See documentation for these Scala 'StreamingQueryStatus' struct""" + is_data_available: builtins.bool + is_trigger_active: builtins.bool + is_active: builtins.bool + def __init__( + self, + *, + status_message: builtins.str = ..., + is_data_available: builtins.bool = ..., + is_trigger_active: builtins.bool = ..., + is_active: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "is_active", + b"is_active", + "is_data_available", + b"is_data_available", + "is_trigger_active", + b"is_trigger_active", + "status_message", + b"status_message", + ], + ) -> None: ... + + class RecentProgressResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RECENT_PROGRESS_JSON_FIELD_NUMBER: builtins.int + @property + def recent_progress_json( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Progress reports as an array of json strings.""" + def __init__( + self, + *, + recent_progress_json: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal["recent_progress_json", b"recent_progress_json"], + ) -> None: ... + + class ExplainResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESULT_FIELD_NUMBER: builtins.int + result: builtins.str + """Logical and physical plans as string""" + def __init__( + self, + *, + result: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["result", b"result"] + ) -> None: ... + + class ExceptionResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXCEPTION_MESSAGE_FIELD_NUMBER: builtins.int + ERROR_CLASS_FIELD_NUMBER: builtins.int + STACK_TRACE_FIELD_NUMBER: builtins.int + exception_message: builtins.str + """(Optional) Exception message as string, maps to the return value of original + StreamingQueryException's toString method + """ + error_class: builtins.str + """(Optional) Exception error class as string""" + stack_trace: builtins.str + """(Optional) Exception stack trace as string""" + def __init__( + self, + *, + exception_message: builtins.str | None = ..., + error_class: builtins.str | None = ..., + stack_trace: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_error_class", + b"_error_class", + "_exception_message", + b"_exception_message", + "_stack_trace", + b"_stack_trace", + "error_class", + b"error_class", + "exception_message", + b"exception_message", + "stack_trace", + b"stack_trace", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_error_class", + b"_error_class", + "_exception_message", + b"_exception_message", + "_stack_trace", + b"_stack_trace", + "error_class", + b"error_class", + "exception_message", + b"exception_message", + "stack_trace", + b"stack_trace", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_error_class", b"_error_class"] + ) -> typing_extensions.Literal["error_class"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_exception_message", b"_exception_message"], + ) -> typing_extensions.Literal["exception_message"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_stack_trace", b"_stack_trace"] + ) -> typing_extensions.Literal["stack_trace"] | None: ... + + class AwaitTerminationResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TERMINATED_FIELD_NUMBER: builtins.int + terminated: builtins.bool + def __init__( + self, + *, + terminated: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["terminated", b"terminated"] + ) -> None: ... + + QUERY_ID_FIELD_NUMBER: builtins.int + STATUS_FIELD_NUMBER: builtins.int + RECENT_PROGRESS_FIELD_NUMBER: builtins.int + EXPLAIN_FIELD_NUMBER: builtins.int + EXCEPTION_FIELD_NUMBER: builtins.int + AWAIT_TERMINATION_FIELD_NUMBER: builtins.int + @property + def query_id(self) -> global___StreamingQueryInstanceId: + """(Required) Query instance id. See `StreamingQueryInstanceId`.""" + @property + def status(self) -> global___StreamingQueryCommandResult.StatusResult: ... + @property + def recent_progress(self) -> global___StreamingQueryCommandResult.RecentProgressResult: ... + @property + def explain(self) -> global___StreamingQueryCommandResult.ExplainResult: ... + @property + def exception(self) -> global___StreamingQueryCommandResult.ExceptionResult: ... + @property + def await_termination(self) -> global___StreamingQueryCommandResult.AwaitTerminationResult: ... + def __init__( + self, + *, + query_id: global___StreamingQueryInstanceId | None = ..., + status: global___StreamingQueryCommandResult.StatusResult | None = ..., + recent_progress: global___StreamingQueryCommandResult.RecentProgressResult | None = ..., + explain: global___StreamingQueryCommandResult.ExplainResult | None = ..., + exception: global___StreamingQueryCommandResult.ExceptionResult | None = ..., + await_termination: global___StreamingQueryCommandResult.AwaitTerminationResult | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "await_termination", + b"await_termination", + "exception", + b"exception", + "explain", + b"explain", + "query_id", + b"query_id", + "recent_progress", + b"recent_progress", + "result_type", + b"result_type", + "status", + b"status", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "await_termination", + b"await_termination", + "exception", + b"exception", + "explain", + b"explain", + "query_id", + b"query_id", + "recent_progress", + b"recent_progress", + "result_type", + b"result_type", + "status", + b"status", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["result_type", b"result_type"] + ) -> typing_extensions.Literal[ + "status", "recent_progress", "explain", "exception", "await_termination" + ] | None: ... + +global___StreamingQueryCommandResult = StreamingQueryCommandResult + +class StreamingQueryManagerCommand(google.protobuf.message.Message): + """Commands for the streaming query manager.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class AwaitAnyTerminationCommand(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TIMEOUT_MS_FIELD_NUMBER: builtins.int + timeout_ms: builtins.int + """(Optional) The waiting time in milliseconds to wait for any query to terminate.""" + def __init__( + self, + *, + timeout_ms: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_timeout_ms", b"_timeout_ms", "timeout_ms", b"timeout_ms" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_timeout_ms", b"_timeout_ms", "timeout_ms", b"timeout_ms" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_timeout_ms", b"_timeout_ms"] + ) -> typing_extensions.Literal["timeout_ms"] | None: ... + + class StreamingQueryListenerCommand(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LISTENER_PAYLOAD_FIELD_NUMBER: builtins.int + PYTHON_LISTENER_PAYLOAD_FIELD_NUMBER: builtins.int + ID_FIELD_NUMBER: builtins.int + listener_payload: builtins.bytes + @property + def python_listener_payload( + self, + ) -> pyspark.sql.connect.proto.expressions_pb2.PythonUDF: ... + id: builtins.str + def __init__( + self, + *, + listener_payload: builtins.bytes = ..., + python_listener_payload: pyspark.sql.connect.proto.expressions_pb2.PythonUDF + | None = ..., + id: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_python_listener_payload", + b"_python_listener_payload", + "python_listener_payload", + b"python_listener_payload", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_python_listener_payload", + b"_python_listener_payload", + "id", + b"id", + "listener_payload", + b"listener_payload", + "python_listener_payload", + b"python_listener_payload", + ], + ) -> None: ... + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_python_listener_payload", b"_python_listener_payload" + ], + ) -> typing_extensions.Literal["python_listener_payload"] | None: ... + + ACTIVE_FIELD_NUMBER: builtins.int + GET_QUERY_FIELD_NUMBER: builtins.int + AWAIT_ANY_TERMINATION_FIELD_NUMBER: builtins.int + RESET_TERMINATED_FIELD_NUMBER: builtins.int + ADD_LISTENER_FIELD_NUMBER: builtins.int + REMOVE_LISTENER_FIELD_NUMBER: builtins.int + LIST_LISTENERS_FIELD_NUMBER: builtins.int + active: builtins.bool + """active() API, returns a list of active queries.""" + get_query: builtins.str + """get() API, returns the StreamingQuery identified by id.""" + @property + def await_any_termination( + self, + ) -> global___StreamingQueryManagerCommand.AwaitAnyTerminationCommand: + """awaitAnyTermination() API, wait until any query terminates or timeout.""" + reset_terminated: builtins.bool + """resetTerminated() API.""" + @property + def add_listener(self) -> global___StreamingQueryManagerCommand.StreamingQueryListenerCommand: + """addListener API.""" + @property + def remove_listener( + self, + ) -> global___StreamingQueryManagerCommand.StreamingQueryListenerCommand: + """removeListener API.""" + list_listeners: builtins.bool + """listListeners() API, returns a list of streaming query listeners.""" + def __init__( + self, + *, + active: builtins.bool = ..., + get_query: builtins.str = ..., + await_any_termination: global___StreamingQueryManagerCommand.AwaitAnyTerminationCommand + | None = ..., + reset_terminated: builtins.bool = ..., + add_listener: global___StreamingQueryManagerCommand.StreamingQueryListenerCommand + | None = ..., + remove_listener: global___StreamingQueryManagerCommand.StreamingQueryListenerCommand + | None = ..., + list_listeners: builtins.bool = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "active", + b"active", + "add_listener", + b"add_listener", + "await_any_termination", + b"await_any_termination", + "command", + b"command", + "get_query", + b"get_query", + "list_listeners", + b"list_listeners", + "remove_listener", + b"remove_listener", + "reset_terminated", + b"reset_terminated", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "active", + b"active", + "add_listener", + b"add_listener", + "await_any_termination", + b"await_any_termination", + "command", + b"command", + "get_query", + b"get_query", + "list_listeners", + b"list_listeners", + "remove_listener", + b"remove_listener", + "reset_terminated", + b"reset_terminated", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["command", b"command"] + ) -> typing_extensions.Literal[ + "active", + "get_query", + "await_any_termination", + "reset_terminated", + "add_listener", + "remove_listener", + "list_listeners", + ] | None: ... + +global___StreamingQueryManagerCommand = StreamingQueryManagerCommand + +class StreamingQueryManagerCommandResult(google.protobuf.message.Message): + """Response for commands on the streaming query manager.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ActiveResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ACTIVE_QUERIES_FIELD_NUMBER: builtins.int + @property + def active_queries( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___StreamingQueryManagerCommandResult.StreamingQueryInstance + ]: ... + def __init__( + self, + *, + active_queries: collections.abc.Iterable[ + global___StreamingQueryManagerCommandResult.StreamingQueryInstance + ] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["active_queries", b"active_queries"] + ) -> None: ... + + class StreamingQueryInstance(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ID_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + @property + def id(self) -> global___StreamingQueryInstanceId: + """(Required) The id and runId of this query.""" + name: builtins.str + """(Optional) The name of this query.""" + def __init__( + self, + *, + id: global___StreamingQueryInstanceId | None = ..., + name: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal["_name", b"_name", "id", b"id", "name", b"name"], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["_name", b"_name", "id", b"id", "name", b"name"], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_name", b"_name"] + ) -> typing_extensions.Literal["name"] | None: ... + + class AwaitAnyTerminationResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TERMINATED_FIELD_NUMBER: builtins.int + terminated: builtins.bool + def __init__( + self, + *, + terminated: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["terminated", b"terminated"] + ) -> None: ... + + class StreamingQueryListenerInstance(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LISTENER_PAYLOAD_FIELD_NUMBER: builtins.int + listener_payload: builtins.bytes + def __init__( + self, + *, + listener_payload: builtins.bytes = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["listener_payload", b"listener_payload"] + ) -> None: ... + + class ListStreamingQueryListenerResult(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LISTENER_IDS_FIELD_NUMBER: builtins.int + @property + def listener_ids( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) Reference IDs of listener instances.""" + def __init__( + self, + *, + listener_ids: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["listener_ids", b"listener_ids"] + ) -> None: ... + + ACTIVE_FIELD_NUMBER: builtins.int + QUERY_FIELD_NUMBER: builtins.int + AWAIT_ANY_TERMINATION_FIELD_NUMBER: builtins.int + RESET_TERMINATED_FIELD_NUMBER: builtins.int + ADD_LISTENER_FIELD_NUMBER: builtins.int + REMOVE_LISTENER_FIELD_NUMBER: builtins.int + LIST_LISTENERS_FIELD_NUMBER: builtins.int + @property + def active(self) -> global___StreamingQueryManagerCommandResult.ActiveResult: ... + @property + def query(self) -> global___StreamingQueryManagerCommandResult.StreamingQueryInstance: ... + @property + def await_any_termination( + self, + ) -> global___StreamingQueryManagerCommandResult.AwaitAnyTerminationResult: ... + reset_terminated: builtins.bool + add_listener: builtins.bool + remove_listener: builtins.bool + @property + def list_listeners( + self, + ) -> global___StreamingQueryManagerCommandResult.ListStreamingQueryListenerResult: ... + def __init__( + self, + *, + active: global___StreamingQueryManagerCommandResult.ActiveResult | None = ..., + query: global___StreamingQueryManagerCommandResult.StreamingQueryInstance | None = ..., + await_any_termination: global___StreamingQueryManagerCommandResult.AwaitAnyTerminationResult + | None = ..., + reset_terminated: builtins.bool = ..., + add_listener: builtins.bool = ..., + remove_listener: builtins.bool = ..., + list_listeners: global___StreamingQueryManagerCommandResult.ListStreamingQueryListenerResult + | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "active", + b"active", + "add_listener", + b"add_listener", + "await_any_termination", + b"await_any_termination", + "list_listeners", + b"list_listeners", + "query", + b"query", + "remove_listener", + b"remove_listener", + "reset_terminated", + b"reset_terminated", + "result_type", + b"result_type", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "active", + b"active", + "add_listener", + b"add_listener", + "await_any_termination", + b"await_any_termination", + "list_listeners", + b"list_listeners", + "query", + b"query", + "remove_listener", + b"remove_listener", + "reset_terminated", + b"reset_terminated", + "result_type", + b"result_type", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["result_type", b"result_type"] + ) -> typing_extensions.Literal[ + "active", + "query", + "await_any_termination", + "reset_terminated", + "add_listener", + "remove_listener", + "list_listeners", + ] | None: ... + +global___StreamingQueryManagerCommandResult = StreamingQueryManagerCommandResult + +class GetResourcesCommand(google.protobuf.message.Message): + """Command to get the output of 'SparkContext.resources'""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___GetResourcesCommand = GetResourcesCommand + +class GetResourcesCommandResult(google.protobuf.message.Message): + """Response for command 'GetResourcesCommand'.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ResourcesEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + @property + def value(self) -> pyspark.sql.connect.proto.common_pb2.ResourceInformation: ... + def __init__( + self, + *, + key: builtins.str = ..., + value: pyspark.sql.connect.proto.common_pb2.ResourceInformation | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + RESOURCES_FIELD_NUMBER: builtins.int + @property + def resources( + self, + ) -> google.protobuf.internal.containers.MessageMap[ + builtins.str, pyspark.sql.connect.proto.common_pb2.ResourceInformation + ]: ... + def __init__( + self, + *, + resources: collections.abc.Mapping[ + builtins.str, pyspark.sql.connect.proto.common_pb2.ResourceInformation + ] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["resources", b"resources"] + ) -> None: ... + +global___GetResourcesCommandResult = GetResourcesCommandResult diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.py new file mode 100644 index 0000000..9564f63 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.py @@ -0,0 +1,47 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/common.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1aspark/connect/common.proto\x12\rspark.connect"\xb0\x01\n\x0cStorageLevel\x12\x19\n\x08use_disk\x18\x01 \x01(\x08R\x07useDisk\x12\x1d\n\nuse_memory\x18\x02 \x01(\x08R\tuseMemory\x12 \n\x0cuse_off_heap\x18\x03 \x01(\x08R\nuseOffHeap\x12"\n\x0c\x64\x65serialized\x18\x04 \x01(\x08R\x0c\x64\x65serialized\x12 \n\x0breplication\x18\x05 \x01(\x05R\x0breplication"G\n\x13ResourceInformation\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x1c\n\taddresses\x18\x02 \x03(\tR\taddressesB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.common_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _STORAGELEVEL._serialized_start = 46 + _STORAGELEVEL._serialized_end = 222 + _RESOURCEINFORMATION._serialized_start = 224 + _RESOURCEINFORMATION._serialized_end = 295 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.pyi new file mode 100644 index 0000000..bb7bded --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/common_pb2.pyi @@ -0,0 +1,123 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.message +import sys + +if sys.version_info >= (3, 8): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class StorageLevel(google.protobuf.message.Message): + """StorageLevel for persisting Datasets/Tables.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USE_DISK_FIELD_NUMBER: builtins.int + USE_MEMORY_FIELD_NUMBER: builtins.int + USE_OFF_HEAP_FIELD_NUMBER: builtins.int + DESERIALIZED_FIELD_NUMBER: builtins.int + REPLICATION_FIELD_NUMBER: builtins.int + use_disk: builtins.bool + """(Required) Whether the cache should use disk or not.""" + use_memory: builtins.bool + """(Required) Whether the cache should use memory or not.""" + use_off_heap: builtins.bool + """(Required) Whether the cache should use off-heap or not.""" + deserialized: builtins.bool + """(Required) Whether the cached data is deserialized or not.""" + replication: builtins.int + """(Required) The number of replicas.""" + def __init__( + self, + *, + use_disk: builtins.bool = ..., + use_memory: builtins.bool = ..., + use_off_heap: builtins.bool = ..., + deserialized: builtins.bool = ..., + replication: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "deserialized", + b"deserialized", + "replication", + b"replication", + "use_disk", + b"use_disk", + "use_memory", + b"use_memory", + "use_off_heap", + b"use_off_heap", + ], + ) -> None: ... + +global___StorageLevel = StorageLevel + +class ResourceInformation(google.protobuf.message.Message): + """ResourceInformation to hold information about a type of Resource. + The corresponding class is 'org.apache.spark.resource.ResourceInformation' + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + ADDRESSES_FIELD_NUMBER: builtins.int + name: builtins.str + """(Required) The name of the resource""" + @property + def addresses( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) An array of strings describing the addresses of the resource.""" + def __init__( + self, + *, + name: builtins.str = ..., + addresses: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["addresses", b"addresses", "name", b"name"] + ) -> None: ... + +global___ResourceInformation = ResourceInformation diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.py new file mode 100644 index 0000000..285edd4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.py @@ -0,0 +1,53 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/example_plugins.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2 +from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n#spark/connect/example_plugins.proto\x12\rspark.connect\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto"i\n\x15\x45xamplePluginRelation\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63ustom_field\x18\x02 \x01(\tR\x0b\x63ustomField"m\n\x17\x45xamplePluginExpression\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12!\n\x0c\x63ustom_field\x18\x02 \x01(\tR\x0b\x63ustomField"9\n\x14\x45xamplePluginCommand\x12!\n\x0c\x63ustom_field\x18\x01 \x01(\tR\x0b\x63ustomFieldB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.example_plugins_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _EXAMPLEPLUGINRELATION._serialized_start = 118 + _EXAMPLEPLUGINRELATION._serialized_end = 223 + _EXAMPLEPLUGINEXPRESSION._serialized_start = 225 + _EXAMPLEPLUGINEXPRESSION._serialized_end = 334 + _EXAMPLEPLUGINCOMMAND._serialized_start = 336 + _EXAMPLEPLUGINCOMMAND._serialized_end = 393 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.pyi new file mode 100644 index 0000000..1be966f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/example_plugins_pb2.pyi @@ -0,0 +1,112 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import google.protobuf.descriptor +import google.protobuf.message +import pyspark.sql.connect.proto.expressions_pb2 +import pyspark.sql.connect.proto.relations_pb2 +import sys + +if sys.version_info >= (3, 8): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class ExamplePluginRelation(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + CUSTOM_FIELD_FIELD_NUMBER: builtins.int + @property + def input(self) -> pyspark.sql.connect.proto.relations_pb2.Relation: ... + custom_field: builtins.str + def __init__( + self, + *, + input: pyspark.sql.connect.proto.relations_pb2.Relation | None = ..., + custom_field: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["custom_field", b"custom_field", "input", b"input"], + ) -> None: ... + +global___ExamplePluginRelation = ExamplePluginRelation + +class ExamplePluginExpression(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CHILD_FIELD_NUMBER: builtins.int + CUSTOM_FIELD_FIELD_NUMBER: builtins.int + @property + def child(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: ... + custom_field: builtins.str + def __init__( + self, + *, + child: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + custom_field: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["child", b"child"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["child", b"child", "custom_field", b"custom_field"], + ) -> None: ... + +global___ExamplePluginExpression = ExamplePluginExpression + +class ExamplePluginCommand(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CUSTOM_FIELD_FIELD_NUMBER: builtins.int + custom_field: builtins.str + def __init__( + self, + *, + custom_field: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["custom_field", b"custom_field"] + ) -> None: ... + +global___ExamplePluginCommand = ExamplePluginCommand diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.py new file mode 100644 index 0000000..51d1a5d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.py @@ -0,0 +1,107 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/expressions.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2 +from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1fspark/connect/expressions.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x19spark/connect/types.proto"\xd9+\n\nExpression\x12=\n\x07literal\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x07literal\x12\x62\n\x14unresolved_attribute\x18\x02 \x01(\x0b\x32-.spark.connect.Expression.UnresolvedAttributeH\x00R\x13unresolvedAttribute\x12_\n\x13unresolved_function\x18\x03 \x01(\x0b\x32,.spark.connect.Expression.UnresolvedFunctionH\x00R\x12unresolvedFunction\x12Y\n\x11\x65xpression_string\x18\x04 \x01(\x0b\x32*.spark.connect.Expression.ExpressionStringH\x00R\x10\x65xpressionString\x12S\n\x0funresolved_star\x18\x05 \x01(\x0b\x32(.spark.connect.Expression.UnresolvedStarH\x00R\x0eunresolvedStar\x12\x37\n\x05\x61lias\x18\x06 \x01(\x0b\x32\x1f.spark.connect.Expression.AliasH\x00R\x05\x61lias\x12\x34\n\x04\x63\x61st\x18\x07 \x01(\x0b\x32\x1e.spark.connect.Expression.CastH\x00R\x04\x63\x61st\x12V\n\x10unresolved_regex\x18\x08 \x01(\x0b\x32).spark.connect.Expression.UnresolvedRegexH\x00R\x0funresolvedRegex\x12\x44\n\nsort_order\x18\t \x01(\x0b\x32#.spark.connect.Expression.SortOrderH\x00R\tsortOrder\x12S\n\x0flambda_function\x18\n \x01(\x0b\x32(.spark.connect.Expression.LambdaFunctionH\x00R\x0elambdaFunction\x12:\n\x06window\x18\x0b \x01(\x0b\x32 .spark.connect.Expression.WindowH\x00R\x06window\x12l\n\x18unresolved_extract_value\x18\x0c \x01(\x0b\x32\x30.spark.connect.Expression.UnresolvedExtractValueH\x00R\x16unresolvedExtractValue\x12M\n\rupdate_fields\x18\r \x01(\x0b\x32&.spark.connect.Expression.UpdateFieldsH\x00R\x0cupdateFields\x12\x82\x01\n unresolved_named_lambda_variable\x18\x0e \x01(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableH\x00R\x1dunresolvedNamedLambdaVariable\x12~\n#common_inline_user_defined_function\x18\x0f \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionH\x00R\x1f\x63ommonInlineUserDefinedFunction\x12\x42\n\rcall_function\x18\x10 \x01(\x0b\x32\x1b.spark.connect.CallFunctionH\x00R\x0c\x63\x61llFunction\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x1a\x8f\x06\n\x06Window\x12\x42\n\x0fwindow_function\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0ewindowFunction\x12@\n\x0epartition_spec\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\rpartitionSpec\x12\x42\n\norder_spec\x18\x03 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\torderSpec\x12K\n\nframe_spec\x18\x04 \x01(\x0b\x32,.spark.connect.Expression.Window.WindowFrameR\tframeSpec\x1a\xed\x03\n\x0bWindowFrame\x12U\n\nframe_type\x18\x01 \x01(\x0e\x32\x36.spark.connect.Expression.Window.WindowFrame.FrameTypeR\tframeType\x12P\n\x05lower\x18\x02 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05lower\x12P\n\x05upper\x18\x03 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05upper\x1a\x91\x01\n\rFrameBoundary\x12!\n\x0b\x63urrent_row\x18\x01 \x01(\x08H\x00R\ncurrentRow\x12\x1e\n\tunbounded\x18\x02 \x01(\x08H\x00R\tunbounded\x12\x31\n\x05value\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionH\x00R\x05valueB\n\n\x08\x62oundary"O\n\tFrameType\x12\x18\n\x14\x46RAME_TYPE_UNDEFINED\x10\x00\x12\x12\n\x0e\x46RAME_TYPE_ROW\x10\x01\x12\x14\n\x10\x46RAME_TYPE_RANGE\x10\x02\x1a\xa9\x03\n\tSortOrder\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12O\n\tdirection\x18\x02 \x01(\x0e\x32\x31.spark.connect.Expression.SortOrder.SortDirectionR\tdirection\x12U\n\rnull_ordering\x18\x03 \x01(\x0e\x32\x30.spark.connect.Expression.SortOrder.NullOrderingR\x0cnullOrdering"l\n\rSortDirection\x12\x1e\n\x1aSORT_DIRECTION_UNSPECIFIED\x10\x00\x12\x1c\n\x18SORT_DIRECTION_ASCENDING\x10\x01\x12\x1d\n\x19SORT_DIRECTION_DESCENDING\x10\x02"U\n\x0cNullOrdering\x12\x1a\n\x16SORT_NULLS_UNSPECIFIED\x10\x00\x12\x14\n\x10SORT_NULLS_FIRST\x10\x01\x12\x13\n\x0fSORT_NULLS_LAST\x10\x02\x1a\x91\x01\n\x04\x43\x61st\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12-\n\x04type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04type\x12\x1b\n\x08type_str\x18\x03 \x01(\tH\x00R\x07typeStrB\x0e\n\x0c\x63\x61st_to_type\x1a\x9b\x0c\n\x07Literal\x12-\n\x04null\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04null\x12\x18\n\x06\x62inary\x18\x02 \x01(\x0cH\x00R\x06\x62inary\x12\x1a\n\x07\x62oolean\x18\x03 \x01(\x08H\x00R\x07\x62oolean\x12\x14\n\x04\x62yte\x18\x04 \x01(\x05H\x00R\x04\x62yte\x12\x16\n\x05short\x18\x05 \x01(\x05H\x00R\x05short\x12\x1a\n\x07integer\x18\x06 \x01(\x05H\x00R\x07integer\x12\x14\n\x04long\x18\x07 \x01(\x03H\x00R\x04long\x12\x16\n\x05\x66loat\x18\n \x01(\x02H\x00R\x05\x66loat\x12\x18\n\x06\x64ouble\x18\x0b \x01(\x01H\x00R\x06\x64ouble\x12\x45\n\x07\x64\x65\x63imal\x18\x0c \x01(\x0b\x32).spark.connect.Expression.Literal.DecimalH\x00R\x07\x64\x65\x63imal\x12\x18\n\x06string\x18\r \x01(\tH\x00R\x06string\x12\x14\n\x04\x64\x61te\x18\x10 \x01(\x05H\x00R\x04\x64\x61te\x12\x1e\n\ttimestamp\x18\x11 \x01(\x03H\x00R\ttimestamp\x12%\n\rtimestamp_ntz\x18\x12 \x01(\x03H\x00R\x0ctimestampNtz\x12\x61\n\x11\x63\x61lendar_interval\x18\x13 \x01(\x0b\x32\x32.spark.connect.Expression.Literal.CalendarIntervalH\x00R\x10\x63\x61lendarInterval\x12\x30\n\x13year_month_interval\x18\x14 \x01(\x05H\x00R\x11yearMonthInterval\x12,\n\x11\x64\x61y_time_interval\x18\x15 \x01(\x03H\x00R\x0f\x64\x61yTimeInterval\x12?\n\x05\x61rray\x18\x16 \x01(\x0b\x32\'.spark.connect.Expression.Literal.ArrayH\x00R\x05\x61rray\x12\x39\n\x03map\x18\x17 \x01(\x0b\x32%.spark.connect.Expression.Literal.MapH\x00R\x03map\x12\x42\n\x06struct\x18\x18 \x01(\x0b\x32(.spark.connect.Expression.Literal.StructH\x00R\x06struct\x1au\n\x07\x44\x65\x63imal\x12\x14\n\x05value\x18\x01 \x01(\tR\x05value\x12!\n\tprecision\x18\x02 \x01(\x05H\x00R\tprecision\x88\x01\x01\x12\x19\n\x05scale\x18\x03 \x01(\x05H\x01R\x05scale\x88\x01\x01\x42\x0c\n\n_precisionB\x08\n\x06_scale\x1a\x62\n\x10\x43\x61lendarInterval\x12\x16\n\x06months\x18\x01 \x01(\x05R\x06months\x12\x12\n\x04\x64\x61ys\x18\x02 \x01(\x05R\x04\x64\x61ys\x12"\n\x0cmicroseconds\x18\x03 \x01(\x03R\x0cmicroseconds\x1a\x82\x01\n\x05\x41rray\x12:\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x0b\x65lementType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lements\x1a\xe3\x01\n\x03Map\x12\x32\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x07keyType\x12\x36\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\tvalueType\x12\x35\n\x04keys\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x04keys\x12\x39\n\x06values\x18\x04 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1a\x81\x01\n\x06Struct\x12\x38\n\x0bstruct_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\nstructType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lementsB\x0e\n\x0cliteral_type\x1ap\n\x13UnresolvedAttribute\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x42\n\n\x08_plan_id\x1a\xcc\x01\n\x12UnresolvedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x1f\n\x0bis_distinct\x18\x03 \x01(\x08R\nisDistinct\x12\x37\n\x18is_user_defined_function\x18\x04 \x01(\x08R\x15isUserDefinedFunction\x1a\x32\n\x10\x45xpressionString\x12\x1e\n\nexpression\x18\x01 \x01(\tR\nexpression\x1aR\n\x0eUnresolvedStar\x12,\n\x0funparsed_target\x18\x01 \x01(\tH\x00R\x0eunparsedTarget\x88\x01\x01\x42\x12\n\x10_unparsed_target\x1aV\n\x0fUnresolvedRegex\x12\x19\n\x08\x63ol_name\x18\x01 \x01(\tR\x07\x63olName\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x42\n\n\x08_plan_id\x1a\x84\x01\n\x16UnresolvedExtractValue\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12\x39\n\nextraction\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\nextraction\x1a\xbb\x01\n\x0cUpdateFields\x12\x46\n\x11struct_expression\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x10structExpression\x12\x1d\n\nfield_name\x18\x02 \x01(\tR\tfieldName\x12\x44\n\x10value_expression\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0fvalueExpression\x1ax\n\x05\x41lias\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12\x12\n\x04name\x18\x02 \x03(\tR\x04name\x12\x1f\n\x08metadata\x18\x03 \x01(\tH\x00R\x08metadata\x88\x01\x01\x42\x0b\n\t_metadata\x1a\x9e\x01\n\x0eLambdaFunction\x12\x35\n\x08\x66unction\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08\x66unction\x12U\n\targuments\x18\x02 \x03(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableR\targuments\x1a>\n\x1dUnresolvedNamedLambdaVariable\x12\x1d\n\nname_parts\x18\x01 \x03(\tR\tnamePartsB\x0b\n\texpr_type"\xec\x02\n\x1f\x43ommonInlineUserDefinedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x39\n\npython_udf\x18\x04 \x01(\x0b\x32\x18.spark.connect.PythonUDFH\x00R\tpythonUdf\x12I\n\x10scalar_scala_udf\x18\x05 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFH\x00R\x0escalarScalaUdf\x12\x33\n\x08java_udf\x18\x06 \x01(\x0b\x32\x16.spark.connect.JavaUDFH\x00R\x07javaUdfB\n\n\x08\x66unction"\x9b\x01\n\tPythonUDF\x12\x38\n\x0boutput_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVer"\xb8\x01\n\x0eScalarScalaUDF\x12\x18\n\x07payload\x18\x01 \x01(\x0cR\x07payload\x12\x37\n\ninputTypes\x18\x02 \x03(\x0b\x32\x17.spark.connect.DataTypeR\ninputTypes\x12\x37\n\noutputType\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1a\n\x08nullable\x18\x04 \x01(\x08R\x08nullable"\x95\x01\n\x07JavaUDF\x12\x1d\n\nclass_name\x18\x01 \x01(\tR\tclassName\x12=\n\x0boutput_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\noutputType\x88\x01\x01\x12\x1c\n\taggregate\x18\x03 \x01(\x08R\taggregateB\x0e\n\x0c_output_type"l\n\x0c\x43\x61llFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targumentsB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.expressions_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _EXPRESSION._serialized_start = 105 + _EXPRESSION._serialized_end = 5698 + _EXPRESSION_WINDOW._serialized_start = 1543 + _EXPRESSION_WINDOW._serialized_end = 2326 + _EXPRESSION_WINDOW_WINDOWFRAME._serialized_start = 1833 + _EXPRESSION_WINDOW_WINDOWFRAME._serialized_end = 2326 + _EXPRESSION_WINDOW_WINDOWFRAME_FRAMEBOUNDARY._serialized_start = 2100 + _EXPRESSION_WINDOW_WINDOWFRAME_FRAMEBOUNDARY._serialized_end = 2245 + _EXPRESSION_WINDOW_WINDOWFRAME_FRAMETYPE._serialized_start = 2247 + _EXPRESSION_WINDOW_WINDOWFRAME_FRAMETYPE._serialized_end = 2326 + _EXPRESSION_SORTORDER._serialized_start = 2329 + _EXPRESSION_SORTORDER._serialized_end = 2754 + _EXPRESSION_SORTORDER_SORTDIRECTION._serialized_start = 2559 + _EXPRESSION_SORTORDER_SORTDIRECTION._serialized_end = 2667 + _EXPRESSION_SORTORDER_NULLORDERING._serialized_start = 2669 + _EXPRESSION_SORTORDER_NULLORDERING._serialized_end = 2754 + _EXPRESSION_CAST._serialized_start = 2757 + _EXPRESSION_CAST._serialized_end = 2902 + _EXPRESSION_LITERAL._serialized_start = 2905 + _EXPRESSION_LITERAL._serialized_end = 4468 + _EXPRESSION_LITERAL_DECIMAL._serialized_start = 3740 + _EXPRESSION_LITERAL_DECIMAL._serialized_end = 3857 + _EXPRESSION_LITERAL_CALENDARINTERVAL._serialized_start = 3859 + _EXPRESSION_LITERAL_CALENDARINTERVAL._serialized_end = 3957 + _EXPRESSION_LITERAL_ARRAY._serialized_start = 3960 + _EXPRESSION_LITERAL_ARRAY._serialized_end = 4090 + _EXPRESSION_LITERAL_MAP._serialized_start = 4093 + _EXPRESSION_LITERAL_MAP._serialized_end = 4320 + _EXPRESSION_LITERAL_STRUCT._serialized_start = 4323 + _EXPRESSION_LITERAL_STRUCT._serialized_end = 4452 + _EXPRESSION_UNRESOLVEDATTRIBUTE._serialized_start = 4470 + _EXPRESSION_UNRESOLVEDATTRIBUTE._serialized_end = 4582 + _EXPRESSION_UNRESOLVEDFUNCTION._serialized_start = 4585 + _EXPRESSION_UNRESOLVEDFUNCTION._serialized_end = 4789 + _EXPRESSION_EXPRESSIONSTRING._serialized_start = 4791 + _EXPRESSION_EXPRESSIONSTRING._serialized_end = 4841 + _EXPRESSION_UNRESOLVEDSTAR._serialized_start = 4843 + _EXPRESSION_UNRESOLVEDSTAR._serialized_end = 4925 + _EXPRESSION_UNRESOLVEDREGEX._serialized_start = 4927 + _EXPRESSION_UNRESOLVEDREGEX._serialized_end = 5013 + _EXPRESSION_UNRESOLVEDEXTRACTVALUE._serialized_start = 5016 + _EXPRESSION_UNRESOLVEDEXTRACTVALUE._serialized_end = 5148 + _EXPRESSION_UPDATEFIELDS._serialized_start = 5151 + _EXPRESSION_UPDATEFIELDS._serialized_end = 5338 + _EXPRESSION_ALIAS._serialized_start = 5340 + _EXPRESSION_ALIAS._serialized_end = 5460 + _EXPRESSION_LAMBDAFUNCTION._serialized_start = 5463 + _EXPRESSION_LAMBDAFUNCTION._serialized_end = 5621 + _EXPRESSION_UNRESOLVEDNAMEDLAMBDAVARIABLE._serialized_start = 5623 + _EXPRESSION_UNRESOLVEDNAMEDLAMBDAVARIABLE._serialized_end = 5685 + _COMMONINLINEUSERDEFINEDFUNCTION._serialized_start = 5701 + _COMMONINLINEUSERDEFINEDFUNCTION._serialized_end = 6065 + _PYTHONUDF._serialized_start = 6068 + _PYTHONUDF._serialized_end = 6223 + _SCALARSCALAUDF._serialized_start = 6226 + _SCALARSCALAUDF._serialized_end = 6410 + _JAVAUDF._serialized_start = 6413 + _JAVAUDF._serialized_end = 6562 + _CALLFUNCTION._serialized_start = 6564 + _CALLFUNCTION._serialized_end = 6672 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.pyi new file mode 100644 index 0000000..b9b16ce --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/expressions_pb2.pyi @@ -0,0 +1,1507 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.any_pb2 +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.internal.enum_type_wrapper +import google.protobuf.message +import pyspark.sql.connect.proto.types_pb2 +import sys +import typing + +if sys.version_info >= (3, 10): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class Expression(google.protobuf.message.Message): + """Expression used to refer to fields, functions and similar. This can be used everywhere + expressions in SQL appear. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Window(google.protobuf.message.Message): + """Expression for the OVER clause or WINDOW clause.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class WindowFrame(google.protobuf.message.Message): + """The window frame""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _FrameType: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _FrameTypeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + Expression.Window.WindowFrame._FrameType.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + FRAME_TYPE_UNDEFINED: Expression.Window.WindowFrame._FrameType.ValueType # 0 + FRAME_TYPE_ROW: Expression.Window.WindowFrame._FrameType.ValueType # 1 + """RowFrame treats rows in a partition individually.""" + FRAME_TYPE_RANGE: Expression.Window.WindowFrame._FrameType.ValueType # 2 + """RangeFrame treats rows in a partition as groups of peers. + All rows having the same 'ORDER BY' ordering are considered as peers. + """ + + class FrameType(_FrameType, metaclass=_FrameTypeEnumTypeWrapper): ... + FRAME_TYPE_UNDEFINED: Expression.Window.WindowFrame.FrameType.ValueType # 0 + FRAME_TYPE_ROW: Expression.Window.WindowFrame.FrameType.ValueType # 1 + """RowFrame treats rows in a partition individually.""" + FRAME_TYPE_RANGE: Expression.Window.WindowFrame.FrameType.ValueType # 2 + """RangeFrame treats rows in a partition as groups of peers. + All rows having the same 'ORDER BY' ordering are considered as peers. + """ + + class FrameBoundary(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CURRENT_ROW_FIELD_NUMBER: builtins.int + UNBOUNDED_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + current_row: builtins.bool + """CURRENT ROW boundary""" + unbounded: builtins.bool + """UNBOUNDED boundary. + For lower bound, it will be converted to 'UnboundedPreceding'. + for upper bound, it will be converted to 'UnboundedFollowing'. + """ + @property + def value(self) -> global___Expression: + """This is an expression for future proofing. We are expecting literals on the server side.""" + def __init__( + self, + *, + current_row: builtins.bool = ..., + unbounded: builtins.bool = ..., + value: global___Expression | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "boundary", + b"boundary", + "current_row", + b"current_row", + "unbounded", + b"unbounded", + "value", + b"value", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "boundary", + b"boundary", + "current_row", + b"current_row", + "unbounded", + b"unbounded", + "value", + b"value", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["boundary", b"boundary"] + ) -> typing_extensions.Literal["current_row", "unbounded", "value"] | None: ... + + FRAME_TYPE_FIELD_NUMBER: builtins.int + LOWER_FIELD_NUMBER: builtins.int + UPPER_FIELD_NUMBER: builtins.int + frame_type: global___Expression.Window.WindowFrame.FrameType.ValueType + """(Required) The type of the frame.""" + @property + def lower(self) -> global___Expression.Window.WindowFrame.FrameBoundary: + """(Required) The lower bound of the frame.""" + @property + def upper(self) -> global___Expression.Window.WindowFrame.FrameBoundary: + """(Required) The upper bound of the frame.""" + def __init__( + self, + *, + frame_type: global___Expression.Window.WindowFrame.FrameType.ValueType = ..., + lower: global___Expression.Window.WindowFrame.FrameBoundary | None = ..., + upper: global___Expression.Window.WindowFrame.FrameBoundary | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["lower", b"lower", "upper", b"upper"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "frame_type", b"frame_type", "lower", b"lower", "upper", b"upper" + ], + ) -> None: ... + + WINDOW_FUNCTION_FIELD_NUMBER: builtins.int + PARTITION_SPEC_FIELD_NUMBER: builtins.int + ORDER_SPEC_FIELD_NUMBER: builtins.int + FRAME_SPEC_FIELD_NUMBER: builtins.int + @property + def window_function(self) -> global___Expression: + """(Required) The window function.""" + @property + def partition_spec( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression + ]: + """(Optional) The way that input rows are partitioned.""" + @property + def order_spec( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.SortOrder + ]: + """(Optional) Ordering of rows in a partition.""" + @property + def frame_spec(self) -> global___Expression.Window.WindowFrame: + """(Optional) Window frame in a partition. + + If not set, it will be treated as 'UnspecifiedFrame'. + """ + def __init__( + self, + *, + window_function: global___Expression | None = ..., + partition_spec: collections.abc.Iterable[global___Expression] | None = ..., + order_spec: collections.abc.Iterable[global___Expression.SortOrder] | None = ..., + frame_spec: global___Expression.Window.WindowFrame | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "frame_spec", b"frame_spec", "window_function", b"window_function" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "frame_spec", + b"frame_spec", + "order_spec", + b"order_spec", + "partition_spec", + b"partition_spec", + "window_function", + b"window_function", + ], + ) -> None: ... + + class SortOrder(google.protobuf.message.Message): + """SortOrder is used to specify the data ordering, it is normally used in Sort and Window. + It is an unevaluable expression and cannot be evaluated, so can not be used in Projection. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _SortDirection: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _SortDirectionEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + Expression.SortOrder._SortDirection.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + SORT_DIRECTION_UNSPECIFIED: Expression.SortOrder._SortDirection.ValueType # 0 + SORT_DIRECTION_ASCENDING: Expression.SortOrder._SortDirection.ValueType # 1 + SORT_DIRECTION_DESCENDING: Expression.SortOrder._SortDirection.ValueType # 2 + + class SortDirection(_SortDirection, metaclass=_SortDirectionEnumTypeWrapper): ... + SORT_DIRECTION_UNSPECIFIED: Expression.SortOrder.SortDirection.ValueType # 0 + SORT_DIRECTION_ASCENDING: Expression.SortOrder.SortDirection.ValueType # 1 + SORT_DIRECTION_DESCENDING: Expression.SortOrder.SortDirection.ValueType # 2 + + class _NullOrdering: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _NullOrderingEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + Expression.SortOrder._NullOrdering.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + SORT_NULLS_UNSPECIFIED: Expression.SortOrder._NullOrdering.ValueType # 0 + SORT_NULLS_FIRST: Expression.SortOrder._NullOrdering.ValueType # 1 + SORT_NULLS_LAST: Expression.SortOrder._NullOrdering.ValueType # 2 + + class NullOrdering(_NullOrdering, metaclass=_NullOrderingEnumTypeWrapper): ... + SORT_NULLS_UNSPECIFIED: Expression.SortOrder.NullOrdering.ValueType # 0 + SORT_NULLS_FIRST: Expression.SortOrder.NullOrdering.ValueType # 1 + SORT_NULLS_LAST: Expression.SortOrder.NullOrdering.ValueType # 2 + + CHILD_FIELD_NUMBER: builtins.int + DIRECTION_FIELD_NUMBER: builtins.int + NULL_ORDERING_FIELD_NUMBER: builtins.int + @property + def child(self) -> global___Expression: + """(Required) The expression to be sorted.""" + direction: global___Expression.SortOrder.SortDirection.ValueType + """(Required) The sort direction, should be ASCENDING or DESCENDING.""" + null_ordering: global___Expression.SortOrder.NullOrdering.ValueType + """(Required) How to deal with NULLs, should be NULLS_FIRST or NULLS_LAST.""" + def __init__( + self, + *, + child: global___Expression | None = ..., + direction: global___Expression.SortOrder.SortDirection.ValueType = ..., + null_ordering: global___Expression.SortOrder.NullOrdering.ValueType = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["child", b"child"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "child", b"child", "direction", b"direction", "null_ordering", b"null_ordering" + ], + ) -> None: ... + + class Cast(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXPR_FIELD_NUMBER: builtins.int + TYPE_FIELD_NUMBER: builtins.int + TYPE_STR_FIELD_NUMBER: builtins.int + @property + def expr(self) -> global___Expression: + """(Required) the expression to be casted.""" + @property + def type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + type_str: builtins.str + """If this is set, Server will use Catalyst parser to parse this string to DataType.""" + def __init__( + self, + *, + expr: global___Expression | None = ..., + type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + type_str: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "cast_to_type", + b"cast_to_type", + "expr", + b"expr", + "type", + b"type", + "type_str", + b"type_str", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "cast_to_type", + b"cast_to_type", + "expr", + b"expr", + "type", + b"type", + "type_str", + b"type_str", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["cast_to_type", b"cast_to_type"] + ) -> typing_extensions.Literal["type", "type_str"] | None: ... + + class Literal(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Decimal(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VALUE_FIELD_NUMBER: builtins.int + PRECISION_FIELD_NUMBER: builtins.int + SCALE_FIELD_NUMBER: builtins.int + value: builtins.str + """the string representation.""" + precision: builtins.int + """The maximum number of digits allowed in the value. + the maximum precision is 38. + """ + scale: builtins.int + """declared scale of decimal literal""" + def __init__( + self, + *, + value: builtins.str = ..., + precision: builtins.int | None = ..., + scale: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_precision", + b"_precision", + "_scale", + b"_scale", + "precision", + b"precision", + "scale", + b"scale", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_precision", + b"_precision", + "_scale", + b"_scale", + "precision", + b"precision", + "scale", + b"scale", + "value", + b"value", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_precision", b"_precision"] + ) -> typing_extensions.Literal["precision"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_scale", b"_scale"] + ) -> typing_extensions.Literal["scale"] | None: ... + + class CalendarInterval(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + MONTHS_FIELD_NUMBER: builtins.int + DAYS_FIELD_NUMBER: builtins.int + MICROSECONDS_FIELD_NUMBER: builtins.int + months: builtins.int + days: builtins.int + microseconds: builtins.int + def __init__( + self, + *, + months: builtins.int = ..., + days: builtins.int = ..., + microseconds: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "days", b"days", "microseconds", b"microseconds", "months", b"months" + ], + ) -> None: ... + + class Array(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ELEMENT_TYPE_FIELD_NUMBER: builtins.int + ELEMENTS_FIELD_NUMBER: builtins.int + @property + def element_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + @property + def elements( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.Literal + ]: ... + def __init__( + self, + *, + element_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + elements: collections.abc.Iterable[global___Expression.Literal] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["element_type", b"element_type"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "element_type", b"element_type", "elements", b"elements" + ], + ) -> None: ... + + class Map(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_TYPE_FIELD_NUMBER: builtins.int + VALUE_TYPE_FIELD_NUMBER: builtins.int + KEYS_FIELD_NUMBER: builtins.int + VALUES_FIELD_NUMBER: builtins.int + @property + def key_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + @property + def value_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + @property + def keys( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.Literal + ]: ... + @property + def values( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.Literal + ]: ... + def __init__( + self, + *, + key_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + value_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + keys: collections.abc.Iterable[global___Expression.Literal] | None = ..., + values: collections.abc.Iterable[global___Expression.Literal] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "key_type", b"key_type", "value_type", b"value_type" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "key_type", + b"key_type", + "keys", + b"keys", + "value_type", + b"value_type", + "values", + b"values", + ], + ) -> None: ... + + class Struct(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STRUCT_TYPE_FIELD_NUMBER: builtins.int + ELEMENTS_FIELD_NUMBER: builtins.int + @property + def struct_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + @property + def elements( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.Literal + ]: ... + def __init__( + self, + *, + struct_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + elements: collections.abc.Iterable[global___Expression.Literal] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["struct_type", b"struct_type"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "elements", b"elements", "struct_type", b"struct_type" + ], + ) -> None: ... + + NULL_FIELD_NUMBER: builtins.int + BINARY_FIELD_NUMBER: builtins.int + BOOLEAN_FIELD_NUMBER: builtins.int + BYTE_FIELD_NUMBER: builtins.int + SHORT_FIELD_NUMBER: builtins.int + INTEGER_FIELD_NUMBER: builtins.int + LONG_FIELD_NUMBER: builtins.int + FLOAT_FIELD_NUMBER: builtins.int + DOUBLE_FIELD_NUMBER: builtins.int + DECIMAL_FIELD_NUMBER: builtins.int + STRING_FIELD_NUMBER: builtins.int + DATE_FIELD_NUMBER: builtins.int + TIMESTAMP_FIELD_NUMBER: builtins.int + TIMESTAMP_NTZ_FIELD_NUMBER: builtins.int + CALENDAR_INTERVAL_FIELD_NUMBER: builtins.int + YEAR_MONTH_INTERVAL_FIELD_NUMBER: builtins.int + DAY_TIME_INTERVAL_FIELD_NUMBER: builtins.int + ARRAY_FIELD_NUMBER: builtins.int + MAP_FIELD_NUMBER: builtins.int + STRUCT_FIELD_NUMBER: builtins.int + @property + def null(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ... + binary: builtins.bytes + boolean: builtins.bool + byte: builtins.int + short: builtins.int + integer: builtins.int + long: builtins.int + float: builtins.float + double: builtins.float + @property + def decimal(self) -> global___Expression.Literal.Decimal: ... + string: builtins.str + date: builtins.int + """Date in units of days since the UNIX epoch.""" + timestamp: builtins.int + """Timestamp in units of microseconds since the UNIX epoch.""" + timestamp_ntz: builtins.int + """Timestamp in units of microseconds since the UNIX epoch (without timezone information).""" + @property + def calendar_interval(self) -> global___Expression.Literal.CalendarInterval: ... + year_month_interval: builtins.int + day_time_interval: builtins.int + @property + def array(self) -> global___Expression.Literal.Array: ... + @property + def map(self) -> global___Expression.Literal.Map: ... + @property + def struct(self) -> global___Expression.Literal.Struct: ... + def __init__( + self, + *, + null: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + binary: builtins.bytes = ..., + boolean: builtins.bool = ..., + byte: builtins.int = ..., + short: builtins.int = ..., + integer: builtins.int = ..., + long: builtins.int = ..., + float: builtins.float = ..., + double: builtins.float = ..., + decimal: global___Expression.Literal.Decimal | None = ..., + string: builtins.str = ..., + date: builtins.int = ..., + timestamp: builtins.int = ..., + timestamp_ntz: builtins.int = ..., + calendar_interval: global___Expression.Literal.CalendarInterval | None = ..., + year_month_interval: builtins.int = ..., + day_time_interval: builtins.int = ..., + array: global___Expression.Literal.Array | None = ..., + map: global___Expression.Literal.Map | None = ..., + struct: global___Expression.Literal.Struct | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "array", + b"array", + "binary", + b"binary", + "boolean", + b"boolean", + "byte", + b"byte", + "calendar_interval", + b"calendar_interval", + "date", + b"date", + "day_time_interval", + b"day_time_interval", + "decimal", + b"decimal", + "double", + b"double", + "float", + b"float", + "integer", + b"integer", + "literal_type", + b"literal_type", + "long", + b"long", + "map", + b"map", + "null", + b"null", + "short", + b"short", + "string", + b"string", + "struct", + b"struct", + "timestamp", + b"timestamp", + "timestamp_ntz", + b"timestamp_ntz", + "year_month_interval", + b"year_month_interval", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "array", + b"array", + "binary", + b"binary", + "boolean", + b"boolean", + "byte", + b"byte", + "calendar_interval", + b"calendar_interval", + "date", + b"date", + "day_time_interval", + b"day_time_interval", + "decimal", + b"decimal", + "double", + b"double", + "float", + b"float", + "integer", + b"integer", + "literal_type", + b"literal_type", + "long", + b"long", + "map", + b"map", + "null", + b"null", + "short", + b"short", + "string", + b"string", + "struct", + b"struct", + "timestamp", + b"timestamp", + "timestamp_ntz", + b"timestamp_ntz", + "year_month_interval", + b"year_month_interval", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["literal_type", b"literal_type"] + ) -> typing_extensions.Literal[ + "null", + "binary", + "boolean", + "byte", + "short", + "integer", + "long", + "float", + "double", + "decimal", + "string", + "date", + "timestamp", + "timestamp_ntz", + "calendar_interval", + "year_month_interval", + "day_time_interval", + "array", + "map", + "struct", + ] | None: ... + + class UnresolvedAttribute(google.protobuf.message.Message): + """An unresolved attribute that is not explicitly bound to a specific column, but the column + is resolved during analysis by name. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + UNPARSED_IDENTIFIER_FIELD_NUMBER: builtins.int + PLAN_ID_FIELD_NUMBER: builtins.int + unparsed_identifier: builtins.str + """(Required) An identifier that will be parsed by Catalyst parser. This should follow the + Spark SQL identifier syntax. + """ + plan_id: builtins.int + """(Optional) The id of corresponding connect plan.""" + def __init__( + self, + *, + unparsed_identifier: builtins.str = ..., + plan_id: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal["_plan_id", b"_plan_id", "plan_id", b"plan_id"], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_plan_id", + b"_plan_id", + "plan_id", + b"plan_id", + "unparsed_identifier", + b"unparsed_identifier", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_plan_id", b"_plan_id"] + ) -> typing_extensions.Literal["plan_id"] | None: ... + + class UnresolvedFunction(google.protobuf.message.Message): + """An unresolved function is not explicitly bound to one explicit function, but the function + is resolved during analysis following Sparks name resolution rules. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + ARGUMENTS_FIELD_NUMBER: builtins.int + IS_DISTINCT_FIELD_NUMBER: builtins.int + IS_USER_DEFINED_FUNCTION_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required) name (or unparsed name for user defined function) for the unresolved function.""" + @property + def arguments( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression + ]: + """(Optional) Function arguments. Empty arguments are allowed.""" + is_distinct: builtins.bool + """(Required) Indicate if this function should be applied on distinct values.""" + is_user_defined_function: builtins.bool + """(Required) Indicate if this is a user defined function. + + When it is not a user defined function, Connect will use the function name directly. + When it is a user defined function, Connect will parse the function name first. + """ + def __init__( + self, + *, + function_name: builtins.str = ..., + arguments: collections.abc.Iterable[global___Expression] | None = ..., + is_distinct: builtins.bool = ..., + is_user_defined_function: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arguments", + b"arguments", + "function_name", + b"function_name", + "is_distinct", + b"is_distinct", + "is_user_defined_function", + b"is_user_defined_function", + ], + ) -> None: ... + + class ExpressionString(google.protobuf.message.Message): + """Expression as string.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXPRESSION_FIELD_NUMBER: builtins.int + expression: builtins.str + """(Required) A SQL expression that will be parsed by Catalyst parser.""" + def __init__( + self, + *, + expression: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["expression", b"expression"] + ) -> None: ... + + class UnresolvedStar(google.protobuf.message.Message): + """UnresolvedStar is used to expand all the fields of a relation or struct.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + UNPARSED_TARGET_FIELD_NUMBER: builtins.int + unparsed_target: builtins.str + """(Optional) The target of the expansion. + + If set, it should end with '.*' and will be parsed by 'parseAttributeName' + in the server side. + """ + def __init__( + self, + *, + unparsed_target: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_unparsed_target", b"_unparsed_target", "unparsed_target", b"unparsed_target" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_unparsed_target", b"_unparsed_target", "unparsed_target", b"unparsed_target" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_unparsed_target", b"_unparsed_target"] + ) -> typing_extensions.Literal["unparsed_target"] | None: ... + + class UnresolvedRegex(google.protobuf.message.Message): + """Represents all of the input attributes to a given relational operator, for example in + "SELECT `(id)?+.+` FROM ...". + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + COL_NAME_FIELD_NUMBER: builtins.int + PLAN_ID_FIELD_NUMBER: builtins.int + col_name: builtins.str + """(Required) The column name used to extract column with regex.""" + plan_id: builtins.int + """(Optional) The id of corresponding connect plan.""" + def __init__( + self, + *, + col_name: builtins.str = ..., + plan_id: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal["_plan_id", b"_plan_id", "plan_id", b"plan_id"], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_plan_id", b"_plan_id", "col_name", b"col_name", "plan_id", b"plan_id" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_plan_id", b"_plan_id"] + ) -> typing_extensions.Literal["plan_id"] | None: ... + + class UnresolvedExtractValue(google.protobuf.message.Message): + """Extracts a value or values from an Expression""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CHILD_FIELD_NUMBER: builtins.int + EXTRACTION_FIELD_NUMBER: builtins.int + @property + def child(self) -> global___Expression: + """(Required) The expression to extract value from, can be + Map, Array, Struct or array of Structs. + """ + @property + def extraction(self) -> global___Expression: + """(Required) The expression to describe the extraction, can be + key of Map, index of Array, field name of Struct. + """ + def __init__( + self, + *, + child: global___Expression | None = ..., + extraction: global___Expression | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal["child", b"child", "extraction", b"extraction"], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["child", b"child", "extraction", b"extraction"], + ) -> None: ... + + class UpdateFields(google.protobuf.message.Message): + """Add, replace or drop a field of `StructType` expression by name.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STRUCT_EXPRESSION_FIELD_NUMBER: builtins.int + FIELD_NAME_FIELD_NUMBER: builtins.int + VALUE_EXPRESSION_FIELD_NUMBER: builtins.int + @property + def struct_expression(self) -> global___Expression: + """(Required) The struct expression.""" + field_name: builtins.str + """(Required) The field name.""" + @property + def value_expression(self) -> global___Expression: + """(Optional) The expression to add or replace. + + When not set, it means this field will be dropped. + """ + def __init__( + self, + *, + struct_expression: global___Expression | None = ..., + field_name: builtins.str = ..., + value_expression: global___Expression | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "struct_expression", b"struct_expression", "value_expression", b"value_expression" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "field_name", + b"field_name", + "struct_expression", + b"struct_expression", + "value_expression", + b"value_expression", + ], + ) -> None: ... + + class Alias(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + EXPR_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + METADATA_FIELD_NUMBER: builtins.int + @property + def expr(self) -> global___Expression: + """(Required) The expression that alias will be added on.""" + @property + def name( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) a list of name parts for the alias. + + Scalar columns only has one name that presents. + """ + metadata: builtins.str + """(Optional) Alias metadata expressed as a JSON map.""" + def __init__( + self, + *, + expr: global___Expression | None = ..., + name: collections.abc.Iterable[builtins.str] | None = ..., + metadata: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_metadata", b"_metadata", "expr", b"expr", "metadata", b"metadata" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_metadata", b"_metadata", "expr", b"expr", "metadata", b"metadata", "name", b"name" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_metadata", b"_metadata"] + ) -> typing_extensions.Literal["metadata"] | None: ... + + class LambdaFunction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_FIELD_NUMBER: builtins.int + ARGUMENTS_FIELD_NUMBER: builtins.int + @property + def function(self) -> global___Expression: + """(Required) The lambda function. + + The function body should use 'UnresolvedAttribute' as arguments, the sever side will + replace 'UnresolvedAttribute' with 'UnresolvedNamedLambdaVariable'. + """ + @property + def arguments( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___Expression.UnresolvedNamedLambdaVariable + ]: + """(Required) Function variables. Must contains 1 ~ 3 variables.""" + def __init__( + self, + *, + function: global___Expression | None = ..., + arguments: collections.abc.Iterable[global___Expression.UnresolvedNamedLambdaVariable] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["function", b"function"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arguments", b"arguments", "function", b"function" + ], + ) -> None: ... + + class UnresolvedNamedLambdaVariable(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_PARTS_FIELD_NUMBER: builtins.int + @property + def name_parts( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) a list of name parts for the variable. Must not be empty.""" + def __init__( + self, + *, + name_parts: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["name_parts", b"name_parts"] + ) -> None: ... + + LITERAL_FIELD_NUMBER: builtins.int + UNRESOLVED_ATTRIBUTE_FIELD_NUMBER: builtins.int + UNRESOLVED_FUNCTION_FIELD_NUMBER: builtins.int + EXPRESSION_STRING_FIELD_NUMBER: builtins.int + UNRESOLVED_STAR_FIELD_NUMBER: builtins.int + ALIAS_FIELD_NUMBER: builtins.int + CAST_FIELD_NUMBER: builtins.int + UNRESOLVED_REGEX_FIELD_NUMBER: builtins.int + SORT_ORDER_FIELD_NUMBER: builtins.int + LAMBDA_FUNCTION_FIELD_NUMBER: builtins.int + WINDOW_FIELD_NUMBER: builtins.int + UNRESOLVED_EXTRACT_VALUE_FIELD_NUMBER: builtins.int + UPDATE_FIELDS_FIELD_NUMBER: builtins.int + UNRESOLVED_NAMED_LAMBDA_VARIABLE_FIELD_NUMBER: builtins.int + COMMON_INLINE_USER_DEFINED_FUNCTION_FIELD_NUMBER: builtins.int + CALL_FUNCTION_FIELD_NUMBER: builtins.int + EXTENSION_FIELD_NUMBER: builtins.int + @property + def literal(self) -> global___Expression.Literal: ... + @property + def unresolved_attribute(self) -> global___Expression.UnresolvedAttribute: ... + @property + def unresolved_function(self) -> global___Expression.UnresolvedFunction: ... + @property + def expression_string(self) -> global___Expression.ExpressionString: ... + @property + def unresolved_star(self) -> global___Expression.UnresolvedStar: ... + @property + def alias(self) -> global___Expression.Alias: ... + @property + def cast(self) -> global___Expression.Cast: ... + @property + def unresolved_regex(self) -> global___Expression.UnresolvedRegex: ... + @property + def sort_order(self) -> global___Expression.SortOrder: ... + @property + def lambda_function(self) -> global___Expression.LambdaFunction: ... + @property + def window(self) -> global___Expression.Window: ... + @property + def unresolved_extract_value(self) -> global___Expression.UnresolvedExtractValue: ... + @property + def update_fields(self) -> global___Expression.UpdateFields: ... + @property + def unresolved_named_lambda_variable( + self, + ) -> global___Expression.UnresolvedNamedLambdaVariable: ... + @property + def common_inline_user_defined_function(self) -> global___CommonInlineUserDefinedFunction: ... + @property + def call_function(self) -> global___CallFunction: ... + @property + def extension(self) -> google.protobuf.any_pb2.Any: + """This field is used to mark extensions to the protocol. When plugins generate arbitrary + relations they can add them here. During the planning the correct resolution is done. + """ + def __init__( + self, + *, + literal: global___Expression.Literal | None = ..., + unresolved_attribute: global___Expression.UnresolvedAttribute | None = ..., + unresolved_function: global___Expression.UnresolvedFunction | None = ..., + expression_string: global___Expression.ExpressionString | None = ..., + unresolved_star: global___Expression.UnresolvedStar | None = ..., + alias: global___Expression.Alias | None = ..., + cast: global___Expression.Cast | None = ..., + unresolved_regex: global___Expression.UnresolvedRegex | None = ..., + sort_order: global___Expression.SortOrder | None = ..., + lambda_function: global___Expression.LambdaFunction | None = ..., + window: global___Expression.Window | None = ..., + unresolved_extract_value: global___Expression.UnresolvedExtractValue | None = ..., + update_fields: global___Expression.UpdateFields | None = ..., + unresolved_named_lambda_variable: global___Expression.UnresolvedNamedLambdaVariable + | None = ..., + common_inline_user_defined_function: global___CommonInlineUserDefinedFunction | None = ..., + call_function: global___CallFunction | None = ..., + extension: google.protobuf.any_pb2.Any | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "alias", + b"alias", + "call_function", + b"call_function", + "cast", + b"cast", + "common_inline_user_defined_function", + b"common_inline_user_defined_function", + "expr_type", + b"expr_type", + "expression_string", + b"expression_string", + "extension", + b"extension", + "lambda_function", + b"lambda_function", + "literal", + b"literal", + "sort_order", + b"sort_order", + "unresolved_attribute", + b"unresolved_attribute", + "unresolved_extract_value", + b"unresolved_extract_value", + "unresolved_function", + b"unresolved_function", + "unresolved_named_lambda_variable", + b"unresolved_named_lambda_variable", + "unresolved_regex", + b"unresolved_regex", + "unresolved_star", + b"unresolved_star", + "update_fields", + b"update_fields", + "window", + b"window", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "alias", + b"alias", + "call_function", + b"call_function", + "cast", + b"cast", + "common_inline_user_defined_function", + b"common_inline_user_defined_function", + "expr_type", + b"expr_type", + "expression_string", + b"expression_string", + "extension", + b"extension", + "lambda_function", + b"lambda_function", + "literal", + b"literal", + "sort_order", + b"sort_order", + "unresolved_attribute", + b"unresolved_attribute", + "unresolved_extract_value", + b"unresolved_extract_value", + "unresolved_function", + b"unresolved_function", + "unresolved_named_lambda_variable", + b"unresolved_named_lambda_variable", + "unresolved_regex", + b"unresolved_regex", + "unresolved_star", + b"unresolved_star", + "update_fields", + b"update_fields", + "window", + b"window", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["expr_type", b"expr_type"] + ) -> typing_extensions.Literal[ + "literal", + "unresolved_attribute", + "unresolved_function", + "expression_string", + "unresolved_star", + "alias", + "cast", + "unresolved_regex", + "sort_order", + "lambda_function", + "window", + "unresolved_extract_value", + "update_fields", + "unresolved_named_lambda_variable", + "common_inline_user_defined_function", + "call_function", + "extension", + ] | None: ... + +global___Expression = Expression + +class CommonInlineUserDefinedFunction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + DETERMINISTIC_FIELD_NUMBER: builtins.int + ARGUMENTS_FIELD_NUMBER: builtins.int + PYTHON_UDF_FIELD_NUMBER: builtins.int + SCALAR_SCALA_UDF_FIELD_NUMBER: builtins.int + JAVA_UDF_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required) Name of the user-defined function.""" + deterministic: builtins.bool + """(Optional) Indicate if the user-defined function is deterministic.""" + @property + def arguments( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Expression]: + """(Optional) Function arguments. Empty arguments are allowed.""" + @property + def python_udf(self) -> global___PythonUDF: ... + @property + def scalar_scala_udf(self) -> global___ScalarScalaUDF: ... + @property + def java_udf(self) -> global___JavaUDF: ... + def __init__( + self, + *, + function_name: builtins.str = ..., + deterministic: builtins.bool = ..., + arguments: collections.abc.Iterable[global___Expression] | None = ..., + python_udf: global___PythonUDF | None = ..., + scalar_scala_udf: global___ScalarScalaUDF | None = ..., + java_udf: global___JavaUDF | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "function", + b"function", + "java_udf", + b"java_udf", + "python_udf", + b"python_udf", + "scalar_scala_udf", + b"scalar_scala_udf", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arguments", + b"arguments", + "deterministic", + b"deterministic", + "function", + b"function", + "function_name", + b"function_name", + "java_udf", + b"java_udf", + "python_udf", + b"python_udf", + "scalar_scala_udf", + b"scalar_scala_udf", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["function", b"function"] + ) -> typing_extensions.Literal["python_udf", "scalar_scala_udf", "java_udf"] | None: ... + +global___CommonInlineUserDefinedFunction = CommonInlineUserDefinedFunction + +class PythonUDF(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + OUTPUT_TYPE_FIELD_NUMBER: builtins.int + EVAL_TYPE_FIELD_NUMBER: builtins.int + COMMAND_FIELD_NUMBER: builtins.int + PYTHON_VER_FIELD_NUMBER: builtins.int + @property + def output_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Required) Output type of the Python UDF""" + eval_type: builtins.int + """(Required) EvalType of the Python UDF""" + command: builtins.bytes + """(Required) The encoded commands of the Python UDF""" + python_ver: builtins.str + """(Required) Python version being used in the client.""" + def __init__( + self, + *, + output_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + eval_type: builtins.int = ..., + command: builtins.bytes = ..., + python_ver: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["output_type", b"output_type"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "command", + b"command", + "eval_type", + b"eval_type", + "output_type", + b"output_type", + "python_ver", + b"python_ver", + ], + ) -> None: ... + +global___PythonUDF = PythonUDF + +class ScalarScalaUDF(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PAYLOAD_FIELD_NUMBER: builtins.int + INPUTTYPES_FIELD_NUMBER: builtins.int + OUTPUTTYPE_FIELD_NUMBER: builtins.int + NULLABLE_FIELD_NUMBER: builtins.int + payload: builtins.bytes + """(Required) Serialized JVM object containing UDF definition, input encoders and output encoder""" + @property + def inputTypes( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.types_pb2.DataType + ]: + """(Optional) Input type(s) of the UDF""" + @property + def outputType(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Required) Output type of the UDF""" + nullable: builtins.bool + """(Required) True if the UDF can return null value""" + def __init__( + self, + *, + payload: builtins.bytes = ..., + inputTypes: collections.abc.Iterable[pyspark.sql.connect.proto.types_pb2.DataType] + | None = ..., + outputType: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + nullable: builtins.bool = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["outputType", b"outputType"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "inputTypes", + b"inputTypes", + "nullable", + b"nullable", + "outputType", + b"outputType", + "payload", + b"payload", + ], + ) -> None: ... + +global___ScalarScalaUDF = ScalarScalaUDF + +class JavaUDF(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CLASS_NAME_FIELD_NUMBER: builtins.int + OUTPUT_TYPE_FIELD_NUMBER: builtins.int + AGGREGATE_FIELD_NUMBER: builtins.int + class_name: builtins.str + """(Required) Fully qualified name of Java class""" + @property + def output_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional) Output type of the Java UDF""" + aggregate: builtins.bool + """(Required) Indicate if the Java user-defined function is an aggregate function""" + def __init__( + self, + *, + class_name: builtins.str = ..., + output_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + aggregate: builtins.bool = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_output_type", b"_output_type", "output_type", b"output_type" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_output_type", + b"_output_type", + "aggregate", + b"aggregate", + "class_name", + b"class_name", + "output_type", + b"output_type", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_output_type", b"_output_type"] + ) -> typing_extensions.Literal["output_type"] | None: ... + +global___JavaUDF = JavaUDF + +class CallFunction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + ARGUMENTS_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required) Unparsed name of the SQL function.""" + @property + def arguments( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Expression]: + """(Optional) Function arguments. Empty arguments are allowed.""" + def __init__( + self, + *, + function_name: builtins.str = ..., + arguments: collections.abc.Iterable[global___Expression] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arguments", b"arguments", "function_name", b"function_name" + ], + ) -> None: ... + +global___CallFunction = CallFunction diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.py new file mode 100644 index 0000000..3a0a7ff --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.py @@ -0,0 +1,195 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/relations.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2 +from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2 +from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__pb2 +from pyspark.sql.connect.proto import catalog_pb2 as spark_dot_connect_dot_catalog__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1dspark/connect/relations.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1fspark/connect/expressions.proto\x1a\x19spark/connect/types.proto\x1a\x1bspark/connect/catalog.proto"\xe1\x18\n\x08Relation\x12\x35\n\x06\x63ommon\x18\x01 \x01(\x0b\x32\x1d.spark.connect.RelationCommonR\x06\x63ommon\x12)\n\x04read\x18\x02 \x01(\x0b\x32\x13.spark.connect.ReadH\x00R\x04read\x12\x32\n\x07project\x18\x03 \x01(\x0b\x32\x16.spark.connect.ProjectH\x00R\x07project\x12/\n\x06\x66ilter\x18\x04 \x01(\x0b\x32\x15.spark.connect.FilterH\x00R\x06\x66ilter\x12)\n\x04join\x18\x05 \x01(\x0b\x32\x13.spark.connect.JoinH\x00R\x04join\x12\x34\n\x06set_op\x18\x06 \x01(\x0b\x32\x1b.spark.connect.SetOperationH\x00R\x05setOp\x12)\n\x04sort\x18\x07 \x01(\x0b\x32\x13.spark.connect.SortH\x00R\x04sort\x12,\n\x05limit\x18\x08 \x01(\x0b\x32\x14.spark.connect.LimitH\x00R\x05limit\x12\x38\n\taggregate\x18\t \x01(\x0b\x32\x18.spark.connect.AggregateH\x00R\taggregate\x12&\n\x03sql\x18\n \x01(\x0b\x32\x12.spark.connect.SQLH\x00R\x03sql\x12\x45\n\x0elocal_relation\x18\x0b \x01(\x0b\x32\x1c.spark.connect.LocalRelationH\x00R\rlocalRelation\x12/\n\x06sample\x18\x0c \x01(\x0b\x32\x15.spark.connect.SampleH\x00R\x06sample\x12/\n\x06offset\x18\r \x01(\x0b\x32\x15.spark.connect.OffsetH\x00R\x06offset\x12>\n\x0b\x64\x65\x64uplicate\x18\x0e \x01(\x0b\x32\x1a.spark.connect.DeduplicateH\x00R\x0b\x64\x65\x64uplicate\x12,\n\x05range\x18\x0f \x01(\x0b\x32\x14.spark.connect.RangeH\x00R\x05range\x12\x45\n\x0esubquery_alias\x18\x10 \x01(\x0b\x32\x1c.spark.connect.SubqueryAliasH\x00R\rsubqueryAlias\x12>\n\x0brepartition\x18\x11 \x01(\x0b\x32\x1a.spark.connect.RepartitionH\x00R\x0brepartition\x12*\n\x05to_df\x18\x12 \x01(\x0b\x32\x13.spark.connect.ToDFH\x00R\x04toDf\x12U\n\x14with_columns_renamed\x18\x13 \x01(\x0b\x32!.spark.connect.WithColumnsRenamedH\x00R\x12withColumnsRenamed\x12<\n\x0bshow_string\x18\x14 \x01(\x0b\x32\x19.spark.connect.ShowStringH\x00R\nshowString\x12)\n\x04\x64rop\x18\x15 \x01(\x0b\x32\x13.spark.connect.DropH\x00R\x04\x64rop\x12)\n\x04tail\x18\x16 \x01(\x0b\x32\x13.spark.connect.TailH\x00R\x04tail\x12?\n\x0cwith_columns\x18\x17 \x01(\x0b\x32\x1a.spark.connect.WithColumnsH\x00R\x0bwithColumns\x12)\n\x04hint\x18\x18 \x01(\x0b\x32\x13.spark.connect.HintH\x00R\x04hint\x12\x32\n\x07unpivot\x18\x19 \x01(\x0b\x32\x16.spark.connect.UnpivotH\x00R\x07unpivot\x12\x36\n\tto_schema\x18\x1a \x01(\x0b\x32\x17.spark.connect.ToSchemaH\x00R\x08toSchema\x12\x64\n\x19repartition_by_expression\x18\x1b \x01(\x0b\x32&.spark.connect.RepartitionByExpressionH\x00R\x17repartitionByExpression\x12\x45\n\x0emap_partitions\x18\x1c \x01(\x0b\x32\x1c.spark.connect.MapPartitionsH\x00R\rmapPartitions\x12H\n\x0f\x63ollect_metrics\x18\x1d \x01(\x0b\x32\x1d.spark.connect.CollectMetricsH\x00R\x0e\x63ollectMetrics\x12,\n\x05parse\x18\x1e \x01(\x0b\x32\x14.spark.connect.ParseH\x00R\x05parse\x12\x36\n\tgroup_map\x18\x1f \x01(\x0b\x32\x17.spark.connect.GroupMapH\x00R\x08groupMap\x12=\n\x0c\x63o_group_map\x18 \x01(\x0b\x32\x19.spark.connect.CoGroupMapH\x00R\ncoGroupMap\x12\x45\n\x0ewith_watermark\x18! \x01(\x0b\x32\x1c.spark.connect.WithWatermarkH\x00R\rwithWatermark\x12\x63\n\x1a\x61pply_in_pandas_with_state\x18" \x01(\x0b\x32%.spark.connect.ApplyInPandasWithStateH\x00R\x16\x61pplyInPandasWithState\x12<\n\x0bhtml_string\x18# \x01(\x0b\x32\x19.spark.connect.HtmlStringH\x00R\nhtmlString\x12X\n\x15\x63\x61\x63hed_local_relation\x18$ \x01(\x0b\x32".spark.connect.CachedLocalRelationH\x00R\x13\x63\x61\x63hedLocalRelation\x12[\n\x16\x63\x61\x63hed_remote_relation\x18% \x01(\x0b\x32#.spark.connect.CachedRemoteRelationH\x00R\x14\x63\x61\x63hedRemoteRelation\x12\x8e\x01\n)common_inline_user_defined_table_function\x18& \x01(\x0b\x32\x33.spark.connect.CommonInlineUserDefinedTableFunctionH\x00R$commonInlineUserDefinedTableFunction\x12\x30\n\x07\x66ill_na\x18Z \x01(\x0b\x32\x15.spark.connect.NAFillH\x00R\x06\x66illNa\x12\x30\n\x07\x64rop_na\x18[ \x01(\x0b\x32\x15.spark.connect.NADropH\x00R\x06\x64ropNa\x12\x34\n\x07replace\x18\\ \x01(\x0b\x32\x18.spark.connect.NAReplaceH\x00R\x07replace\x12\x36\n\x07summary\x18\x64 \x01(\x0b\x32\x1a.spark.connect.StatSummaryH\x00R\x07summary\x12\x39\n\x08\x63rosstab\x18\x65 \x01(\x0b\x32\x1b.spark.connect.StatCrosstabH\x00R\x08\x63rosstab\x12\x39\n\x08\x64\x65scribe\x18\x66 \x01(\x0b\x32\x1b.spark.connect.StatDescribeH\x00R\x08\x64\x65scribe\x12*\n\x03\x63ov\x18g \x01(\x0b\x32\x16.spark.connect.StatCovH\x00R\x03\x63ov\x12-\n\x04\x63orr\x18h \x01(\x0b\x32\x17.spark.connect.StatCorrH\x00R\x04\x63orr\x12L\n\x0f\x61pprox_quantile\x18i \x01(\x0b\x32!.spark.connect.StatApproxQuantileH\x00R\x0e\x61pproxQuantile\x12=\n\nfreq_items\x18j \x01(\x0b\x32\x1c.spark.connect.StatFreqItemsH\x00R\tfreqItems\x12:\n\tsample_by\x18k \x01(\x0b\x32\x1b.spark.connect.StatSampleByH\x00R\x08sampleBy\x12\x33\n\x07\x63\x61talog\x18\xc8\x01 \x01(\x0b\x32\x16.spark.connect.CatalogH\x00R\x07\x63\x61talog\x12\x35\n\textension\x18\xe6\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x33\n\x07unknown\x18\xe7\x07 \x01(\x0b\x32\x16.spark.connect.UnknownH\x00R\x07unknownB\n\n\x08rel_type"\t\n\x07Unknown"[\n\x0eRelationCommon\x12\x1f\n\x0bsource_info\x18\x01 \x01(\tR\nsourceInfo\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x42\n\n\x08_plan_id"\xe7\x01\n\x03SQL\x12\x14\n\x05query\x18\x01 \x01(\tR\x05query\x12\x30\n\x04\x61rgs\x18\x02 \x03(\x0b\x32\x1c.spark.connect.SQL.ArgsEntryR\x04\x61rgs\x12<\n\x08pos_args\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x07posArgs\x1aZ\n\tArgsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x05value:\x02\x38\x01"\x97\x05\n\x04Read\x12\x41\n\x0bnamed_table\x18\x01 \x01(\x0b\x32\x1e.spark.connect.Read.NamedTableH\x00R\nnamedTable\x12\x41\n\x0b\x64\x61ta_source\x18\x02 \x01(\x0b\x32\x1e.spark.connect.Read.DataSourceH\x00R\ndataSource\x12!\n\x0cis_streaming\x18\x03 \x01(\x08R\x0bisStreaming\x1a\xc0\x01\n\nNamedTable\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x45\n\x07options\x18\x02 \x03(\x0b\x32+.spark.connect.Read.NamedTable.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\x95\x02\n\nDataSource\x12\x1b\n\x06\x66ormat\x18\x01 \x01(\tH\x00R\x06\x66ormat\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x12\x45\n\x07options\x18\x03 \x03(\x0b\x32+.spark.connect.Read.DataSource.OptionsEntryR\x07options\x12\x14\n\x05paths\x18\x04 \x03(\tR\x05paths\x12\x1e\n\npredicates\x18\x05 \x03(\tR\npredicates\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\t\n\x07_formatB\t\n\x07_schemaB\x0b\n\tread_type"u\n\x07Project\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12;\n\x0b\x65xpressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0b\x65xpressions"p\n\x06\x46ilter\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x37\n\tcondition\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\tcondition"\x95\x05\n\x04Join\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12@\n\x0ejoin_condition\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\rjoinCondition\x12\x39\n\tjoin_type\x18\x04 \x01(\x0e\x32\x1c.spark.connect.Join.JoinTypeR\x08joinType\x12#\n\rusing_columns\x18\x05 \x03(\tR\x0cusingColumns\x12K\n\x0ejoin_data_type\x18\x06 \x01(\x0b\x32 .spark.connect.Join.JoinDataTypeH\x00R\x0cjoinDataType\x88\x01\x01\x1a\\\n\x0cJoinDataType\x12$\n\x0eis_left_struct\x18\x01 \x01(\x08R\x0cisLeftStruct\x12&\n\x0fis_right_struct\x18\x02 \x01(\x08R\risRightStruct"\xd0\x01\n\x08JoinType\x12\x19\n\x15JOIN_TYPE_UNSPECIFIED\x10\x00\x12\x13\n\x0fJOIN_TYPE_INNER\x10\x01\x12\x18\n\x14JOIN_TYPE_FULL_OUTER\x10\x02\x12\x18\n\x14JOIN_TYPE_LEFT_OUTER\x10\x03\x12\x19\n\x15JOIN_TYPE_RIGHT_OUTER\x10\x04\x12\x17\n\x13JOIN_TYPE_LEFT_ANTI\x10\x05\x12\x17\n\x13JOIN_TYPE_LEFT_SEMI\x10\x06\x12\x13\n\x0fJOIN_TYPE_CROSS\x10\x07\x42\x11\n\x0f_join_data_type"\xdf\x03\n\x0cSetOperation\x12\x36\n\nleft_input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\tleftInput\x12\x38\n\x0bright_input\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\nrightInput\x12\x45\n\x0bset_op_type\x18\x03 \x01(\x0e\x32%.spark.connect.SetOperation.SetOpTypeR\tsetOpType\x12\x1a\n\x06is_all\x18\x04 \x01(\x08H\x00R\x05isAll\x88\x01\x01\x12\x1c\n\x07\x62y_name\x18\x05 \x01(\x08H\x01R\x06\x62yName\x88\x01\x01\x12\x37\n\x15\x61llow_missing_columns\x18\x06 \x01(\x08H\x02R\x13\x61llowMissingColumns\x88\x01\x01"r\n\tSetOpType\x12\x1b\n\x17SET_OP_TYPE_UNSPECIFIED\x10\x00\x12\x19\n\x15SET_OP_TYPE_INTERSECT\x10\x01\x12\x15\n\x11SET_OP_TYPE_UNION\x10\x02\x12\x16\n\x12SET_OP_TYPE_EXCEPT\x10\x03\x42\t\n\x07_is_allB\n\n\x08_by_nameB\x18\n\x16_allow_missing_columns"L\n\x05Limit\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"O\n\x06Offset\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x16\n\x06offset\x18\x02 \x01(\x05R\x06offset"K\n\x04Tail\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"\xc6\x04\n\tAggregate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x41\n\ngroup_type\x18\x02 \x01(\x0e\x32".spark.connect.Aggregate.GroupTypeR\tgroupType\x12L\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12N\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x14\x61ggregateExpressions\x12\x34\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.spark.connect.Aggregate.PivotR\x05pivot\x1ao\n\x05Pivot\x12+\n\x03\x63ol\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values"\x81\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04"\xa0\x01\n\x04Sort\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x05order\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\x05order\x12 \n\tis_global\x18\x03 \x01(\x08H\x00R\x08isGlobal\x88\x01\x01\x42\x0c\n\n_is_global"\x8d\x01\n\x04\x44rop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x33\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07\x63olumns\x12!\n\x0c\x63olumn_names\x18\x03 \x03(\tR\x0b\x63olumnNames"\xf0\x01\n\x0b\x44\x65\x64uplicate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames\x12\x32\n\x13\x61ll_columns_as_keys\x18\x03 \x01(\x08H\x00R\x10\x61llColumnsAsKeys\x88\x01\x01\x12.\n\x10within_watermark\x18\x04 \x01(\x08H\x01R\x0fwithinWatermark\x88\x01\x01\x42\x16\n\x14_all_columns_as_keysB\x13\n\x11_within_watermark"Y\n\rLocalRelation\x12\x17\n\x04\x64\x61ta\x18\x01 \x01(\x0cH\x00R\x04\x64\x61ta\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x42\x07\n\x05_dataB\t\n\x07_schema"_\n\x13\x43\x61\x63hedLocalRelation\x12\x16\n\x06userId\x18\x01 \x01(\tR\x06userId\x12\x1c\n\tsessionId\x18\x02 \x01(\tR\tsessionId\x12\x12\n\x04hash\x18\x03 \x01(\tR\x04hash"7\n\x14\x43\x61\x63hedRemoteRelation\x12\x1f\n\x0brelation_id\x18\x01 \x01(\tR\nrelationId"\x91\x02\n\x06Sample\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1f\n\x0blower_bound\x18\x02 \x01(\x01R\nlowerBound\x12\x1f\n\x0bupper_bound\x18\x03 \x01(\x01R\nupperBound\x12.\n\x10with_replacement\x18\x04 \x01(\x08H\x00R\x0fwithReplacement\x88\x01\x01\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x01R\x04seed\x88\x01\x01\x12/\n\x13\x64\x65terministic_order\x18\x06 \x01(\x08R\x12\x64\x65terministicOrderB\x13\n\x11_with_replacementB\x07\n\x05_seed"\x91\x01\n\x05Range\x12\x19\n\x05start\x18\x01 \x01(\x03H\x00R\x05start\x88\x01\x01\x12\x10\n\x03\x65nd\x18\x02 \x01(\x03R\x03\x65nd\x12\x12\n\x04step\x18\x03 \x01(\x03R\x04step\x12*\n\x0enum_partitions\x18\x04 \x01(\x05H\x01R\rnumPartitions\x88\x01\x01\x42\x08\n\x06_startB\x11\n\x0f_num_partitions"r\n\rSubqueryAlias\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05\x61lias\x18\x02 \x01(\tR\x05\x61lias\x12\x1c\n\tqualifier\x18\x03 \x03(\tR\tqualifier"\x8e\x01\n\x0bRepartition\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12%\n\x0enum_partitions\x18\x02 \x01(\x05R\rnumPartitions\x12\x1d\n\x07shuffle\x18\x03 \x01(\x08H\x00R\x07shuffle\x88\x01\x01\x42\n\n\x08_shuffle"\x8e\x01\n\nShowString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate\x12\x1a\n\x08vertical\x18\x04 \x01(\x08R\x08vertical"r\n\nHtmlString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate"\\\n\x0bStatSummary\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1e\n\nstatistics\x18\x02 \x03(\tR\nstatistics"Q\n\x0cStatDescribe\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols"e\n\x0cStatCrosstab\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"`\n\x07StatCov\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"\x89\x01\n\x08StatCorr\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2\x12\x1b\n\x06method\x18\x04 \x01(\tH\x00R\x06method\x88\x01\x01\x42\t\n\x07_method"\xa4\x01\n\x12StatApproxQuantile\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12$\n\rprobabilities\x18\x03 \x03(\x01R\rprobabilities\x12%\n\x0erelative_error\x18\x04 \x01(\x01R\rrelativeError"}\n\rStatFreqItems\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x1d\n\x07support\x18\x03 \x01(\x01H\x00R\x07support\x88\x01\x01\x42\n\n\x08_support"\xb5\x02\n\x0cStatSampleBy\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03\x63ol\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x42\n\tfractions\x18\x03 \x03(\x0b\x32$.spark.connect.StatSampleBy.FractionR\tfractions\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x00R\x04seed\x88\x01\x01\x1a\x63\n\x08\x46raction\x12;\n\x07stratum\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x07stratum\x12\x1a\n\x08\x66raction\x18\x02 \x01(\x01R\x08\x66ractionB\x07\n\x05_seed"\x86\x01\n\x06NAFill\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x39\n\x06values\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values"\x86\x01\n\x06NADrop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\'\n\rmin_non_nulls\x18\x03 \x01(\x05H\x00R\x0bminNonNulls\x88\x01\x01\x42\x10\n\x0e_min_non_nulls"\xa8\x02\n\tNAReplace\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12H\n\x0creplacements\x18\x03 \x03(\x0b\x32$.spark.connect.NAReplace.ReplacementR\x0creplacements\x1a\x8d\x01\n\x0bReplacement\x12>\n\told_value\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08oldValue\x12>\n\tnew_value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08newValue"X\n\x04ToDF\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames"\xef\x01\n\x12WithColumnsRenamed\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x65\n\x12rename_columns_map\x18\x02 \x03(\x0b\x32\x37.spark.connect.WithColumnsRenamed.RenameColumnsMapEntryR\x10renameColumnsMap\x1a\x43\n\x15RenameColumnsMapEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"w\n\x0bWithColumns\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x07\x61liases\x18\x02 \x03(\x0b\x32\x1f.spark.connect.Expression.AliasR\x07\x61liases"\x86\x01\n\rWithWatermark\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1d\n\nevent_time\x18\x02 \x01(\tR\teventTime\x12\'\n\x0f\x64\x65lay_threshold\x18\x03 \x01(\tR\x0e\x64\x65layThreshold"\x84\x01\n\x04Hint\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x39\n\nparameters\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\nparameters"\xc7\x02\n\x07Unpivot\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03ids\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x03ids\x12:\n\x06values\x18\x03 \x01(\x0b\x32\x1d.spark.connect.Unpivot.ValuesH\x00R\x06values\x88\x01\x01\x12\x30\n\x14variable_column_name\x18\x04 \x01(\tR\x12variableColumnName\x12*\n\x11value_column_name\x18\x05 \x01(\tR\x0fvalueColumnName\x1a;\n\x06Values\x12\x31\n\x06values\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x06valuesB\t\n\x07_values"j\n\x08ToSchema\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12/\n\x06schema\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema"\xcb\x01\n\x17RepartitionByExpression\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x0fpartition_exprs\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0epartitionExprs\x12*\n\x0enum_partitions\x18\x03 \x01(\x05H\x00R\rnumPartitions\x88\x01\x01\x42\x11\n\x0f_num_partitions"\xb5\x01\n\rMapPartitions\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x04\x66unc\x18\x02 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12"\n\nis_barrier\x18\x03 \x01(\x08H\x00R\tisBarrier\x88\x01\x01\x42\r\n\x0b_is_barrier"\xfb\x04\n\x08GroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12J\n\x13sorting_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x12sortingExpressions\x12<\n\rinitial_input\x18\x05 \x01(\x0b\x32\x17.spark.connect.RelationR\x0cinitialInput\x12[\n\x1cinitial_grouping_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x1ainitialGroupingExpressions\x12;\n\x18is_map_groups_with_state\x18\x07 \x01(\x08H\x00R\x14isMapGroupsWithState\x88\x01\x01\x12$\n\x0boutput_mode\x18\x08 \x01(\tH\x01R\noutputMode\x88\x01\x01\x12&\n\x0ctimeout_conf\x18\t \x01(\tH\x02R\x0btimeoutConf\x88\x01\x01\x42\x1b\n\x19_is_map_groups_with_stateB\x0e\n\x0c_output_modeB\x0f\n\r_timeout_conf"\x8e\x04\n\nCoGroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12W\n\x1ainput_grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18inputGroupingExpressions\x12-\n\x05other\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationR\x05other\x12W\n\x1aother_grouping_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18otherGroupingExpressions\x12\x42\n\x04\x66unc\x18\x05 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12U\n\x19input_sorting_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17inputSortingExpressions\x12U\n\x19other_sorting_expressions\x18\x07 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17otherSortingExpressions"\xe5\x02\n\x16\x41pplyInPandasWithState\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12#\n\routput_schema\x18\x04 \x01(\tR\x0coutputSchema\x12!\n\x0cstate_schema\x18\x05 \x01(\tR\x0bstateSchema\x12\x1f\n\x0boutput_mode\x18\x06 \x01(\tR\noutputMode\x12!\n\x0ctimeout_conf\x18\x07 \x01(\tR\x0btimeoutConf"\xf4\x01\n$CommonInlineUserDefinedTableFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12<\n\x0bpython_udtf\x18\x04 \x01(\x0b\x32\x19.spark.connect.PythonUDTFH\x00R\npythonUdtfB\n\n\x08\x66unction"\xb1\x01\n\nPythonUDTF\x12=\n\x0breturn_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\nreturnType\x88\x01\x01\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVerB\x0e\n\x0c_return_type"\x88\x01\n\x0e\x43ollectMetrics\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x33\n\x07metrics\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07metrics"\x84\x03\n\x05Parse\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x38\n\x06\x66ormat\x18\x02 \x01(\x0e\x32 .spark.connect.Parse.ParseFormatR\x06\x66ormat\x12\x34\n\x06schema\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x06schema\x88\x01\x01\x12;\n\x07options\x18\x04 \x03(\x0b\x32!.spark.connect.Parse.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"X\n\x0bParseFormat\x12\x1c\n\x18PARSE_FORMAT_UNSPECIFIED\x10\x00\x12\x14\n\x10PARSE_FORMAT_CSV\x10\x01\x12\x15\n\x11PARSE_FORMAT_JSON\x10\x02\x42\t\n\x07_schemaB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.relations_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _SQL_ARGSENTRY._options = None + _SQL_ARGSENTRY._serialized_options = b"8\001" + _READ_NAMEDTABLE_OPTIONSENTRY._options = None + _READ_NAMEDTABLE_OPTIONSENTRY._serialized_options = b"8\001" + _READ_DATASOURCE_OPTIONSENTRY._options = None + _READ_DATASOURCE_OPTIONSENTRY._serialized_options = b"8\001" + _WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY._options = None + _WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY._serialized_options = b"8\001" + _PARSE_OPTIONSENTRY._options = None + _PARSE_OPTIONSENTRY._serialized_options = b"8\001" + _RELATION._serialized_start = 165 + _RELATION._serialized_end = 3334 + _UNKNOWN._serialized_start = 3336 + _UNKNOWN._serialized_end = 3345 + _RELATIONCOMMON._serialized_start = 3347 + _RELATIONCOMMON._serialized_end = 3438 + _SQL._serialized_start = 3441 + _SQL._serialized_end = 3672 + _SQL_ARGSENTRY._serialized_start = 3582 + _SQL_ARGSENTRY._serialized_end = 3672 + _READ._serialized_start = 3675 + _READ._serialized_end = 4338 + _READ_NAMEDTABLE._serialized_start = 3853 + _READ_NAMEDTABLE._serialized_end = 4045 + _READ_NAMEDTABLE_OPTIONSENTRY._serialized_start = 3987 + _READ_NAMEDTABLE_OPTIONSENTRY._serialized_end = 4045 + _READ_DATASOURCE._serialized_start = 4048 + _READ_DATASOURCE._serialized_end = 4325 + _READ_DATASOURCE_OPTIONSENTRY._serialized_start = 3987 + _READ_DATASOURCE_OPTIONSENTRY._serialized_end = 4045 + _PROJECT._serialized_start = 4340 + _PROJECT._serialized_end = 4457 + _FILTER._serialized_start = 4459 + _FILTER._serialized_end = 4571 + _JOIN._serialized_start = 4574 + _JOIN._serialized_end = 5235 + _JOIN_JOINDATATYPE._serialized_start = 4913 + _JOIN_JOINDATATYPE._serialized_end = 5005 + _JOIN_JOINTYPE._serialized_start = 5008 + _JOIN_JOINTYPE._serialized_end = 5216 + _SETOPERATION._serialized_start = 5238 + _SETOPERATION._serialized_end = 5717 + _SETOPERATION_SETOPTYPE._serialized_start = 5554 + _SETOPERATION_SETOPTYPE._serialized_end = 5668 + _LIMIT._serialized_start = 5719 + _LIMIT._serialized_end = 5795 + _OFFSET._serialized_start = 5797 + _OFFSET._serialized_end = 5876 + _TAIL._serialized_start = 5878 + _TAIL._serialized_end = 5953 + _AGGREGATE._serialized_start = 5956 + _AGGREGATE._serialized_end = 6538 + _AGGREGATE_PIVOT._serialized_start = 6295 + _AGGREGATE_PIVOT._serialized_end = 6406 + _AGGREGATE_GROUPTYPE._serialized_start = 6409 + _AGGREGATE_GROUPTYPE._serialized_end = 6538 + _SORT._serialized_start = 6541 + _SORT._serialized_end = 6701 + _DROP._serialized_start = 6704 + _DROP._serialized_end = 6845 + _DEDUPLICATE._serialized_start = 6848 + _DEDUPLICATE._serialized_end = 7088 + _LOCALRELATION._serialized_start = 7090 + _LOCALRELATION._serialized_end = 7179 + _CACHEDLOCALRELATION._serialized_start = 7181 + _CACHEDLOCALRELATION._serialized_end = 7276 + _CACHEDREMOTERELATION._serialized_start = 7278 + _CACHEDREMOTERELATION._serialized_end = 7333 + _SAMPLE._serialized_start = 7336 + _SAMPLE._serialized_end = 7609 + _RANGE._serialized_start = 7612 + _RANGE._serialized_end = 7757 + _SUBQUERYALIAS._serialized_start = 7759 + _SUBQUERYALIAS._serialized_end = 7873 + _REPARTITION._serialized_start = 7876 + _REPARTITION._serialized_end = 8018 + _SHOWSTRING._serialized_start = 8021 + _SHOWSTRING._serialized_end = 8163 + _HTMLSTRING._serialized_start = 8165 + _HTMLSTRING._serialized_end = 8279 + _STATSUMMARY._serialized_start = 8281 + _STATSUMMARY._serialized_end = 8373 + _STATDESCRIBE._serialized_start = 8375 + _STATDESCRIBE._serialized_end = 8456 + _STATCROSSTAB._serialized_start = 8458 + _STATCROSSTAB._serialized_end = 8559 + _STATCOV._serialized_start = 8561 + _STATCOV._serialized_end = 8657 + _STATCORR._serialized_start = 8660 + _STATCORR._serialized_end = 8797 + _STATAPPROXQUANTILE._serialized_start = 8800 + _STATAPPROXQUANTILE._serialized_end = 8964 + _STATFREQITEMS._serialized_start = 8966 + _STATFREQITEMS._serialized_end = 9091 + _STATSAMPLEBY._serialized_start = 9094 + _STATSAMPLEBY._serialized_end = 9403 + _STATSAMPLEBY_FRACTION._serialized_start = 9295 + _STATSAMPLEBY_FRACTION._serialized_end = 9394 + _NAFILL._serialized_start = 9406 + _NAFILL._serialized_end = 9540 + _NADROP._serialized_start = 9543 + _NADROP._serialized_end = 9677 + _NAREPLACE._serialized_start = 9680 + _NAREPLACE._serialized_end = 9976 + _NAREPLACE_REPLACEMENT._serialized_start = 9835 + _NAREPLACE_REPLACEMENT._serialized_end = 9976 + _TODF._serialized_start = 9978 + _TODF._serialized_end = 10066 + _WITHCOLUMNSRENAMED._serialized_start = 10069 + _WITHCOLUMNSRENAMED._serialized_end = 10308 + _WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY._serialized_start = 10241 + _WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY._serialized_end = 10308 + _WITHCOLUMNS._serialized_start = 10310 + _WITHCOLUMNS._serialized_end = 10429 + _WITHWATERMARK._serialized_start = 10432 + _WITHWATERMARK._serialized_end = 10566 + _HINT._serialized_start = 10569 + _HINT._serialized_end = 10701 + _UNPIVOT._serialized_start = 10704 + _UNPIVOT._serialized_end = 11031 + _UNPIVOT_VALUES._serialized_start = 10961 + _UNPIVOT_VALUES._serialized_end = 11020 + _TOSCHEMA._serialized_start = 11033 + _TOSCHEMA._serialized_end = 11139 + _REPARTITIONBYEXPRESSION._serialized_start = 11142 + _REPARTITIONBYEXPRESSION._serialized_end = 11345 + _MAPPARTITIONS._serialized_start = 11348 + _MAPPARTITIONS._serialized_end = 11529 + _GROUPMAP._serialized_start = 11532 + _GROUPMAP._serialized_end = 12167 + _COGROUPMAP._serialized_start = 12170 + _COGROUPMAP._serialized_end = 12696 + _APPLYINPANDASWITHSTATE._serialized_start = 12699 + _APPLYINPANDASWITHSTATE._serialized_end = 13056 + _COMMONINLINEUSERDEFINEDTABLEFUNCTION._serialized_start = 13059 + _COMMONINLINEUSERDEFINEDTABLEFUNCTION._serialized_end = 13303 + _PYTHONUDTF._serialized_start = 13306 + _PYTHONUDTF._serialized_end = 13483 + _COLLECTMETRICS._serialized_start = 13486 + _COLLECTMETRICS._serialized_end = 13622 + _PARSE._serialized_start = 13625 + _PARSE._serialized_end = 14013 + _PARSE_OPTIONSENTRY._serialized_start = 3987 + _PARSE_OPTIONSENTRY._serialized_end = 4045 + _PARSE_PARSEFORMAT._serialized_start = 13914 + _PARSE_PARSEFORMAT._serialized_end = 14002 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.pyi new file mode 100644 index 0000000..9cadd4a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/relations_pb2.pyi @@ -0,0 +1,3626 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.any_pb2 +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.internal.enum_type_wrapper +import google.protobuf.message +import pyspark.sql.connect.proto.catalog_pb2 +import pyspark.sql.connect.proto.expressions_pb2 +import pyspark.sql.connect.proto.types_pb2 +import sys +import typing + +if sys.version_info >= (3, 10): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class Relation(google.protobuf.message.Message): + """The main [[Relation]] type. Fundamentally, a relation is a typed container + that has exactly one explicit relation type set. + + When adding new relation types, they have to be registered here. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + COMMON_FIELD_NUMBER: builtins.int + READ_FIELD_NUMBER: builtins.int + PROJECT_FIELD_NUMBER: builtins.int + FILTER_FIELD_NUMBER: builtins.int + JOIN_FIELD_NUMBER: builtins.int + SET_OP_FIELD_NUMBER: builtins.int + SORT_FIELD_NUMBER: builtins.int + LIMIT_FIELD_NUMBER: builtins.int + AGGREGATE_FIELD_NUMBER: builtins.int + SQL_FIELD_NUMBER: builtins.int + LOCAL_RELATION_FIELD_NUMBER: builtins.int + SAMPLE_FIELD_NUMBER: builtins.int + OFFSET_FIELD_NUMBER: builtins.int + DEDUPLICATE_FIELD_NUMBER: builtins.int + RANGE_FIELD_NUMBER: builtins.int + SUBQUERY_ALIAS_FIELD_NUMBER: builtins.int + REPARTITION_FIELD_NUMBER: builtins.int + TO_DF_FIELD_NUMBER: builtins.int + WITH_COLUMNS_RENAMED_FIELD_NUMBER: builtins.int + SHOW_STRING_FIELD_NUMBER: builtins.int + DROP_FIELD_NUMBER: builtins.int + TAIL_FIELD_NUMBER: builtins.int + WITH_COLUMNS_FIELD_NUMBER: builtins.int + HINT_FIELD_NUMBER: builtins.int + UNPIVOT_FIELD_NUMBER: builtins.int + TO_SCHEMA_FIELD_NUMBER: builtins.int + REPARTITION_BY_EXPRESSION_FIELD_NUMBER: builtins.int + MAP_PARTITIONS_FIELD_NUMBER: builtins.int + COLLECT_METRICS_FIELD_NUMBER: builtins.int + PARSE_FIELD_NUMBER: builtins.int + GROUP_MAP_FIELD_NUMBER: builtins.int + CO_GROUP_MAP_FIELD_NUMBER: builtins.int + WITH_WATERMARK_FIELD_NUMBER: builtins.int + APPLY_IN_PANDAS_WITH_STATE_FIELD_NUMBER: builtins.int + HTML_STRING_FIELD_NUMBER: builtins.int + CACHED_LOCAL_RELATION_FIELD_NUMBER: builtins.int + CACHED_REMOTE_RELATION_FIELD_NUMBER: builtins.int + COMMON_INLINE_USER_DEFINED_TABLE_FUNCTION_FIELD_NUMBER: builtins.int + FILL_NA_FIELD_NUMBER: builtins.int + DROP_NA_FIELD_NUMBER: builtins.int + REPLACE_FIELD_NUMBER: builtins.int + SUMMARY_FIELD_NUMBER: builtins.int + CROSSTAB_FIELD_NUMBER: builtins.int + DESCRIBE_FIELD_NUMBER: builtins.int + COV_FIELD_NUMBER: builtins.int + CORR_FIELD_NUMBER: builtins.int + APPROX_QUANTILE_FIELD_NUMBER: builtins.int + FREQ_ITEMS_FIELD_NUMBER: builtins.int + SAMPLE_BY_FIELD_NUMBER: builtins.int + CATALOG_FIELD_NUMBER: builtins.int + EXTENSION_FIELD_NUMBER: builtins.int + UNKNOWN_FIELD_NUMBER: builtins.int + @property + def common(self) -> global___RelationCommon: ... + @property + def read(self) -> global___Read: ... + @property + def project(self) -> global___Project: ... + @property + def filter(self) -> global___Filter: ... + @property + def join(self) -> global___Join: ... + @property + def set_op(self) -> global___SetOperation: ... + @property + def sort(self) -> global___Sort: ... + @property + def limit(self) -> global___Limit: ... + @property + def aggregate(self) -> global___Aggregate: ... + @property + def sql(self) -> global___SQL: ... + @property + def local_relation(self) -> global___LocalRelation: ... + @property + def sample(self) -> global___Sample: ... + @property + def offset(self) -> global___Offset: ... + @property + def deduplicate(self) -> global___Deduplicate: ... + @property + def range(self) -> global___Range: ... + @property + def subquery_alias(self) -> global___SubqueryAlias: ... + @property + def repartition(self) -> global___Repartition: ... + @property + def to_df(self) -> global___ToDF: ... + @property + def with_columns_renamed(self) -> global___WithColumnsRenamed: ... + @property + def show_string(self) -> global___ShowString: ... + @property + def drop(self) -> global___Drop: ... + @property + def tail(self) -> global___Tail: ... + @property + def with_columns(self) -> global___WithColumns: ... + @property + def hint(self) -> global___Hint: ... + @property + def unpivot(self) -> global___Unpivot: ... + @property + def to_schema(self) -> global___ToSchema: ... + @property + def repartition_by_expression(self) -> global___RepartitionByExpression: ... + @property + def map_partitions(self) -> global___MapPartitions: ... + @property + def collect_metrics(self) -> global___CollectMetrics: ... + @property + def parse(self) -> global___Parse: ... + @property + def group_map(self) -> global___GroupMap: ... + @property + def co_group_map(self) -> global___CoGroupMap: ... + @property + def with_watermark(self) -> global___WithWatermark: ... + @property + def apply_in_pandas_with_state(self) -> global___ApplyInPandasWithState: ... + @property + def html_string(self) -> global___HtmlString: ... + @property + def cached_local_relation(self) -> global___CachedLocalRelation: ... + @property + def cached_remote_relation(self) -> global___CachedRemoteRelation: ... + @property + def common_inline_user_defined_table_function( + self, + ) -> global___CommonInlineUserDefinedTableFunction: ... + @property + def fill_na(self) -> global___NAFill: + """NA functions""" + @property + def drop_na(self) -> global___NADrop: ... + @property + def replace(self) -> global___NAReplace: ... + @property + def summary(self) -> global___StatSummary: + """stat functions""" + @property + def crosstab(self) -> global___StatCrosstab: ... + @property + def describe(self) -> global___StatDescribe: ... + @property + def cov(self) -> global___StatCov: ... + @property + def corr(self) -> global___StatCorr: ... + @property + def approx_quantile(self) -> global___StatApproxQuantile: ... + @property + def freq_items(self) -> global___StatFreqItems: ... + @property + def sample_by(self) -> global___StatSampleBy: ... + @property + def catalog(self) -> pyspark.sql.connect.proto.catalog_pb2.Catalog: + """Catalog API (experimental / unstable)""" + @property + def extension(self) -> google.protobuf.any_pb2.Any: + """This field is used to mark extensions to the protocol. When plugins generate arbitrary + relations they can add them here. During the planning the correct resolution is done. + """ + @property + def unknown(self) -> global___Unknown: ... + def __init__( + self, + *, + common: global___RelationCommon | None = ..., + read: global___Read | None = ..., + project: global___Project | None = ..., + filter: global___Filter | None = ..., + join: global___Join | None = ..., + set_op: global___SetOperation | None = ..., + sort: global___Sort | None = ..., + limit: global___Limit | None = ..., + aggregate: global___Aggregate | None = ..., + sql: global___SQL | None = ..., + local_relation: global___LocalRelation | None = ..., + sample: global___Sample | None = ..., + offset: global___Offset | None = ..., + deduplicate: global___Deduplicate | None = ..., + range: global___Range | None = ..., + subquery_alias: global___SubqueryAlias | None = ..., + repartition: global___Repartition | None = ..., + to_df: global___ToDF | None = ..., + with_columns_renamed: global___WithColumnsRenamed | None = ..., + show_string: global___ShowString | None = ..., + drop: global___Drop | None = ..., + tail: global___Tail | None = ..., + with_columns: global___WithColumns | None = ..., + hint: global___Hint | None = ..., + unpivot: global___Unpivot | None = ..., + to_schema: global___ToSchema | None = ..., + repartition_by_expression: global___RepartitionByExpression | None = ..., + map_partitions: global___MapPartitions | None = ..., + collect_metrics: global___CollectMetrics | None = ..., + parse: global___Parse | None = ..., + group_map: global___GroupMap | None = ..., + co_group_map: global___CoGroupMap | None = ..., + with_watermark: global___WithWatermark | None = ..., + apply_in_pandas_with_state: global___ApplyInPandasWithState | None = ..., + html_string: global___HtmlString | None = ..., + cached_local_relation: global___CachedLocalRelation | None = ..., + cached_remote_relation: global___CachedRemoteRelation | None = ..., + common_inline_user_defined_table_function: global___CommonInlineUserDefinedTableFunction + | None = ..., + fill_na: global___NAFill | None = ..., + drop_na: global___NADrop | None = ..., + replace: global___NAReplace | None = ..., + summary: global___StatSummary | None = ..., + crosstab: global___StatCrosstab | None = ..., + describe: global___StatDescribe | None = ..., + cov: global___StatCov | None = ..., + corr: global___StatCorr | None = ..., + approx_quantile: global___StatApproxQuantile | None = ..., + freq_items: global___StatFreqItems | None = ..., + sample_by: global___StatSampleBy | None = ..., + catalog: pyspark.sql.connect.proto.catalog_pb2.Catalog | None = ..., + extension: google.protobuf.any_pb2.Any | None = ..., + unknown: global___Unknown | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "aggregate", + b"aggregate", + "apply_in_pandas_with_state", + b"apply_in_pandas_with_state", + "approx_quantile", + b"approx_quantile", + "cached_local_relation", + b"cached_local_relation", + "cached_remote_relation", + b"cached_remote_relation", + "catalog", + b"catalog", + "co_group_map", + b"co_group_map", + "collect_metrics", + b"collect_metrics", + "common", + b"common", + "common_inline_user_defined_table_function", + b"common_inline_user_defined_table_function", + "corr", + b"corr", + "cov", + b"cov", + "crosstab", + b"crosstab", + "deduplicate", + b"deduplicate", + "describe", + b"describe", + "drop", + b"drop", + "drop_na", + b"drop_na", + "extension", + b"extension", + "fill_na", + b"fill_na", + "filter", + b"filter", + "freq_items", + b"freq_items", + "group_map", + b"group_map", + "hint", + b"hint", + "html_string", + b"html_string", + "join", + b"join", + "limit", + b"limit", + "local_relation", + b"local_relation", + "map_partitions", + b"map_partitions", + "offset", + b"offset", + "parse", + b"parse", + "project", + b"project", + "range", + b"range", + "read", + b"read", + "rel_type", + b"rel_type", + "repartition", + b"repartition", + "repartition_by_expression", + b"repartition_by_expression", + "replace", + b"replace", + "sample", + b"sample", + "sample_by", + b"sample_by", + "set_op", + b"set_op", + "show_string", + b"show_string", + "sort", + b"sort", + "sql", + b"sql", + "subquery_alias", + b"subquery_alias", + "summary", + b"summary", + "tail", + b"tail", + "to_df", + b"to_df", + "to_schema", + b"to_schema", + "unknown", + b"unknown", + "unpivot", + b"unpivot", + "with_columns", + b"with_columns", + "with_columns_renamed", + b"with_columns_renamed", + "with_watermark", + b"with_watermark", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "aggregate", + b"aggregate", + "apply_in_pandas_with_state", + b"apply_in_pandas_with_state", + "approx_quantile", + b"approx_quantile", + "cached_local_relation", + b"cached_local_relation", + "cached_remote_relation", + b"cached_remote_relation", + "catalog", + b"catalog", + "co_group_map", + b"co_group_map", + "collect_metrics", + b"collect_metrics", + "common", + b"common", + "common_inline_user_defined_table_function", + b"common_inline_user_defined_table_function", + "corr", + b"corr", + "cov", + b"cov", + "crosstab", + b"crosstab", + "deduplicate", + b"deduplicate", + "describe", + b"describe", + "drop", + b"drop", + "drop_na", + b"drop_na", + "extension", + b"extension", + "fill_na", + b"fill_na", + "filter", + b"filter", + "freq_items", + b"freq_items", + "group_map", + b"group_map", + "hint", + b"hint", + "html_string", + b"html_string", + "join", + b"join", + "limit", + b"limit", + "local_relation", + b"local_relation", + "map_partitions", + b"map_partitions", + "offset", + b"offset", + "parse", + b"parse", + "project", + b"project", + "range", + b"range", + "read", + b"read", + "rel_type", + b"rel_type", + "repartition", + b"repartition", + "repartition_by_expression", + b"repartition_by_expression", + "replace", + b"replace", + "sample", + b"sample", + "sample_by", + b"sample_by", + "set_op", + b"set_op", + "show_string", + b"show_string", + "sort", + b"sort", + "sql", + b"sql", + "subquery_alias", + b"subquery_alias", + "summary", + b"summary", + "tail", + b"tail", + "to_df", + b"to_df", + "to_schema", + b"to_schema", + "unknown", + b"unknown", + "unpivot", + b"unpivot", + "with_columns", + b"with_columns", + "with_columns_renamed", + b"with_columns_renamed", + "with_watermark", + b"with_watermark", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["rel_type", b"rel_type"] + ) -> typing_extensions.Literal[ + "read", + "project", + "filter", + "join", + "set_op", + "sort", + "limit", + "aggregate", + "sql", + "local_relation", + "sample", + "offset", + "deduplicate", + "range", + "subquery_alias", + "repartition", + "to_df", + "with_columns_renamed", + "show_string", + "drop", + "tail", + "with_columns", + "hint", + "unpivot", + "to_schema", + "repartition_by_expression", + "map_partitions", + "collect_metrics", + "parse", + "group_map", + "co_group_map", + "with_watermark", + "apply_in_pandas_with_state", + "html_string", + "cached_local_relation", + "cached_remote_relation", + "common_inline_user_defined_table_function", + "fill_na", + "drop_na", + "replace", + "summary", + "crosstab", + "describe", + "cov", + "corr", + "approx_quantile", + "freq_items", + "sample_by", + "catalog", + "extension", + "unknown", + ] | None: ... + +global___Relation = Relation + +class Unknown(google.protobuf.message.Message): + """Used for testing purposes only.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___Unknown = Unknown + +class RelationCommon(google.protobuf.message.Message): + """Common metadata of all relations.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SOURCE_INFO_FIELD_NUMBER: builtins.int + PLAN_ID_FIELD_NUMBER: builtins.int + source_info: builtins.str + """(Required) Shared relation metadata.""" + plan_id: builtins.int + """(Optional) A per-client globally unique id for a given connect plan.""" + def __init__( + self, + *, + source_info: builtins.str = ..., + plan_id: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["_plan_id", b"_plan_id", "plan_id", b"plan_id"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_plan_id", b"_plan_id", "plan_id", b"plan_id", "source_info", b"source_info" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_plan_id", b"_plan_id"] + ) -> typing_extensions.Literal["plan_id"] | None: ... + +global___RelationCommon = RelationCommon + +class SQL(google.protobuf.message.Message): + """Relation that uses a SQL query to generate the output.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ArgsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + @property + def value(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression.Literal: ... + def __init__( + self, + *, + key: builtins.str = ..., + value: pyspark.sql.connect.proto.expressions_pb2.Expression.Literal | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["value", b"value"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + QUERY_FIELD_NUMBER: builtins.int + ARGS_FIELD_NUMBER: builtins.int + POS_ARGS_FIELD_NUMBER: builtins.int + query: builtins.str + """(Required) The SQL query.""" + @property + def args( + self, + ) -> google.protobuf.internal.containers.MessageMap[ + builtins.str, pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Optional) A map of parameter names to literal expressions.""" + @property + def pos_args( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Optional) A sequence of literal expressions for positional parameters in the SQL query text.""" + def __init__( + self, + *, + query: builtins.str = ..., + args: collections.abc.Mapping[ + builtins.str, pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + pos_args: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "args", b"args", "pos_args", b"pos_args", "query", b"query" + ], + ) -> None: ... + +global___SQL = SQL + +class Read(google.protobuf.message.Message): + """Relation that reads from a file / table or other data source. Does not have additional + inputs. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class NamedTable(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + UNPARSED_IDENTIFIER_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + unparsed_identifier: builtins.str + """(Required) Unparsed identifier for the table.""" + @property + def options( + self, + ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Options for the named table. The map key is case insensitive.""" + def __init__( + self, + *, + unparsed_identifier: builtins.str = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "options", b"options", "unparsed_identifier", b"unparsed_identifier" + ], + ) -> None: ... + + class DataSource(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + FORMAT_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + PATHS_FIELD_NUMBER: builtins.int + PREDICATES_FIELD_NUMBER: builtins.int + format: builtins.str + """(Optional) Supported formats include: parquet, orc, text, json, parquet, csv, avro. + + If not set, the value from SQL conf 'spark.sql.sources.default' will be used. + """ + schema: builtins.str + """(Optional) If not set, Spark will infer the schema. + + This schema string should be either DDL-formatted or JSON-formatted. + """ + @property + def options( + self, + ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Options for the data source. The context of this map varies based on the + data source format. This options could be empty for valid data source format. + The map key is case insensitive. + """ + @property + def paths( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) A list of path for file-system backed data sources.""" + @property + def predicates( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Condition in the where clause for each partition. + + This is only supported by the JDBC data source. + """ + def __init__( + self, + *, + format: builtins.str | None = ..., + schema: builtins.str | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + paths: collections.abc.Iterable[builtins.str] | None = ..., + predicates: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_format", + b"_format", + "_schema", + b"_schema", + "format", + b"format", + "schema", + b"schema", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_format", + b"_format", + "_schema", + b"_schema", + "format", + b"format", + "options", + b"options", + "paths", + b"paths", + "predicates", + b"predicates", + "schema", + b"schema", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_format", b"_format"] + ) -> typing_extensions.Literal["format"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_schema", b"_schema"] + ) -> typing_extensions.Literal["schema"] | None: ... + + NAMED_TABLE_FIELD_NUMBER: builtins.int + DATA_SOURCE_FIELD_NUMBER: builtins.int + IS_STREAMING_FIELD_NUMBER: builtins.int + @property + def named_table(self) -> global___Read.NamedTable: ... + @property + def data_source(self) -> global___Read.DataSource: ... + is_streaming: builtins.bool + """(Optional) Indicates if this is a streaming read.""" + def __init__( + self, + *, + named_table: global___Read.NamedTable | None = ..., + data_source: global___Read.DataSource | None = ..., + is_streaming: builtins.bool = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "data_source", b"data_source", "named_table", b"named_table", "read_type", b"read_type" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "data_source", + b"data_source", + "is_streaming", + b"is_streaming", + "named_table", + b"named_table", + "read_type", + b"read_type", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["read_type", b"read_type"] + ) -> typing_extensions.Literal["named_table", "data_source"] | None: ... + +global___Read = Read + +class Project(google.protobuf.message.Message): + """Projection of a bag of expressions for a given input relation. + + The input relation must be specified. + The projected expression can be an arbitrary expression. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + EXPRESSIONS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Optional) Input relation is optional for Project. + + For example, `SELECT ABS(-1)` is valid plan without an input plan. + """ + @property + def expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) A Project requires at least one expression.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + expressions: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["expressions", b"expressions", "input", b"input"], + ) -> None: ... + +global___Project = Project + +class Filter(google.protobuf.message.Message): + """Relation that applies a boolean expression `condition` on each row of `input` to produce + the output result. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + CONDITION_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a Filter.""" + @property + def condition(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: + """(Required) A Filter must have a condition expression.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + condition: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["condition", b"condition", "input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["condition", b"condition", "input", b"input"] + ) -> None: ... + +global___Filter = Filter + +class Join(google.protobuf.message.Message): + """Relation of type [[Join]]. + + `left` and `right` must be present. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _JoinType: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _JoinTypeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Join._JoinType.ValueType], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + JOIN_TYPE_UNSPECIFIED: Join._JoinType.ValueType # 0 + JOIN_TYPE_INNER: Join._JoinType.ValueType # 1 + JOIN_TYPE_FULL_OUTER: Join._JoinType.ValueType # 2 + JOIN_TYPE_LEFT_OUTER: Join._JoinType.ValueType # 3 + JOIN_TYPE_RIGHT_OUTER: Join._JoinType.ValueType # 4 + JOIN_TYPE_LEFT_ANTI: Join._JoinType.ValueType # 5 + JOIN_TYPE_LEFT_SEMI: Join._JoinType.ValueType # 6 + JOIN_TYPE_CROSS: Join._JoinType.ValueType # 7 + + class JoinType(_JoinType, metaclass=_JoinTypeEnumTypeWrapper): ... + JOIN_TYPE_UNSPECIFIED: Join.JoinType.ValueType # 0 + JOIN_TYPE_INNER: Join.JoinType.ValueType # 1 + JOIN_TYPE_FULL_OUTER: Join.JoinType.ValueType # 2 + JOIN_TYPE_LEFT_OUTER: Join.JoinType.ValueType # 3 + JOIN_TYPE_RIGHT_OUTER: Join.JoinType.ValueType # 4 + JOIN_TYPE_LEFT_ANTI: Join.JoinType.ValueType # 5 + JOIN_TYPE_LEFT_SEMI: Join.JoinType.ValueType # 6 + JOIN_TYPE_CROSS: Join.JoinType.ValueType # 7 + + class JoinDataType(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + IS_LEFT_STRUCT_FIELD_NUMBER: builtins.int + IS_RIGHT_STRUCT_FIELD_NUMBER: builtins.int + is_left_struct: builtins.bool + """If the left data type is a struct.""" + is_right_struct: builtins.bool + """If the right data type is a struct.""" + def __init__( + self, + *, + is_left_struct: builtins.bool = ..., + is_right_struct: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "is_left_struct", b"is_left_struct", "is_right_struct", b"is_right_struct" + ], + ) -> None: ... + + LEFT_FIELD_NUMBER: builtins.int + RIGHT_FIELD_NUMBER: builtins.int + JOIN_CONDITION_FIELD_NUMBER: builtins.int + JOIN_TYPE_FIELD_NUMBER: builtins.int + USING_COLUMNS_FIELD_NUMBER: builtins.int + JOIN_DATA_TYPE_FIELD_NUMBER: builtins.int + @property + def left(self) -> global___Relation: + """(Required) Left input relation for a Join.""" + @property + def right(self) -> global___Relation: + """(Required) Right input relation for a Join.""" + @property + def join_condition(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: + """(Optional) The join condition. Could be unset when `using_columns` is utilized. + + This field does not co-exist with using_columns. + """ + join_type: global___Join.JoinType.ValueType + """(Required) The join type.""" + @property + def using_columns( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Optional. using_columns provides a list of columns that should present on both sides of + the join inputs that this Join will join on. For example A JOIN B USING col_name is + equivalent to A JOIN B on A.col_name = B.col_name. + + This field does not co-exist with join_condition. + """ + @property + def join_data_type(self) -> global___Join.JoinDataType: + """(Optional) Only used by joinWith. Set the left and right join data types.""" + def __init__( + self, + *, + left: global___Relation | None = ..., + right: global___Relation | None = ..., + join_condition: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + join_type: global___Join.JoinType.ValueType = ..., + using_columns: collections.abc.Iterable[builtins.str] | None = ..., + join_data_type: global___Join.JoinDataType | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_join_data_type", + b"_join_data_type", + "join_condition", + b"join_condition", + "join_data_type", + b"join_data_type", + "left", + b"left", + "right", + b"right", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_join_data_type", + b"_join_data_type", + "join_condition", + b"join_condition", + "join_data_type", + b"join_data_type", + "join_type", + b"join_type", + "left", + b"left", + "right", + b"right", + "using_columns", + b"using_columns", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_join_data_type", b"_join_data_type"] + ) -> typing_extensions.Literal["join_data_type"] | None: ... + +global___Join = Join + +class SetOperation(google.protobuf.message.Message): + """Relation of type [[SetOperation]]""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _SetOpType: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _SetOpTypeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ + SetOperation._SetOpType.ValueType + ], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + SET_OP_TYPE_UNSPECIFIED: SetOperation._SetOpType.ValueType # 0 + SET_OP_TYPE_INTERSECT: SetOperation._SetOpType.ValueType # 1 + SET_OP_TYPE_UNION: SetOperation._SetOpType.ValueType # 2 + SET_OP_TYPE_EXCEPT: SetOperation._SetOpType.ValueType # 3 + + class SetOpType(_SetOpType, metaclass=_SetOpTypeEnumTypeWrapper): ... + SET_OP_TYPE_UNSPECIFIED: SetOperation.SetOpType.ValueType # 0 + SET_OP_TYPE_INTERSECT: SetOperation.SetOpType.ValueType # 1 + SET_OP_TYPE_UNION: SetOperation.SetOpType.ValueType # 2 + SET_OP_TYPE_EXCEPT: SetOperation.SetOpType.ValueType # 3 + + LEFT_INPUT_FIELD_NUMBER: builtins.int + RIGHT_INPUT_FIELD_NUMBER: builtins.int + SET_OP_TYPE_FIELD_NUMBER: builtins.int + IS_ALL_FIELD_NUMBER: builtins.int + BY_NAME_FIELD_NUMBER: builtins.int + ALLOW_MISSING_COLUMNS_FIELD_NUMBER: builtins.int + @property + def left_input(self) -> global___Relation: + """(Required) Left input relation for a Set operation.""" + @property + def right_input(self) -> global___Relation: + """(Required) Right input relation for a Set operation.""" + set_op_type: global___SetOperation.SetOpType.ValueType + """(Required) The Set operation type.""" + is_all: builtins.bool + """(Optional) If to remove duplicate rows. + + True to preserve all results. + False to remove duplicate rows. + """ + by_name: builtins.bool + """(Optional) If to perform the Set operation based on name resolution. + + Only UNION supports this option. + """ + allow_missing_columns: builtins.bool + """(Optional) If to perform the Set operation and allow missing columns. + + Only UNION supports this option. + """ + def __init__( + self, + *, + left_input: global___Relation | None = ..., + right_input: global___Relation | None = ..., + set_op_type: global___SetOperation.SetOpType.ValueType = ..., + is_all: builtins.bool | None = ..., + by_name: builtins.bool | None = ..., + allow_missing_columns: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_allow_missing_columns", + b"_allow_missing_columns", + "_by_name", + b"_by_name", + "_is_all", + b"_is_all", + "allow_missing_columns", + b"allow_missing_columns", + "by_name", + b"by_name", + "is_all", + b"is_all", + "left_input", + b"left_input", + "right_input", + b"right_input", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_allow_missing_columns", + b"_allow_missing_columns", + "_by_name", + b"_by_name", + "_is_all", + b"_is_all", + "allow_missing_columns", + b"allow_missing_columns", + "by_name", + b"by_name", + "is_all", + b"is_all", + "left_input", + b"left_input", + "right_input", + b"right_input", + "set_op_type", + b"set_op_type", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_allow_missing_columns", b"_allow_missing_columns"], + ) -> typing_extensions.Literal["allow_missing_columns"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_by_name", b"_by_name"] + ) -> typing_extensions.Literal["by_name"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_is_all", b"_is_all"] + ) -> typing_extensions.Literal["is_all"] | None: ... + +global___SetOperation = SetOperation + +class Limit(google.protobuf.message.Message): + """Relation of type [[Limit]] that is used to `limit` rows from the input relation.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + LIMIT_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a Limit.""" + limit: builtins.int + """(Required) the limit.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + limit: builtins.int = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["input", b"input", "limit", b"limit"] + ) -> None: ... + +global___Limit = Limit + +class Offset(google.protobuf.message.Message): + """Relation of type [[Offset]] that is used to read rows staring from the `offset` on + the input relation. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + OFFSET_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for an Offset.""" + offset: builtins.int + """(Required) the limit.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + offset: builtins.int = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["input", b"input", "offset", b"offset"] + ) -> None: ... + +global___Offset = Offset + +class Tail(google.protobuf.message.Message): + """Relation of type [[Tail]] that is used to fetch `limit` rows from the last of the input relation.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + LIMIT_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for an Tail.""" + limit: builtins.int + """(Required) the limit.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + limit: builtins.int = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["input", b"input", "limit", b"limit"] + ) -> None: ... + +global___Tail = Tail + +class Aggregate(google.protobuf.message.Message): + """Relation of type [[Aggregate]].""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _GroupType: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _GroupTypeEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Aggregate._GroupType.ValueType], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + GROUP_TYPE_UNSPECIFIED: Aggregate._GroupType.ValueType # 0 + GROUP_TYPE_GROUPBY: Aggregate._GroupType.ValueType # 1 + GROUP_TYPE_ROLLUP: Aggregate._GroupType.ValueType # 2 + GROUP_TYPE_CUBE: Aggregate._GroupType.ValueType # 3 + GROUP_TYPE_PIVOT: Aggregate._GroupType.ValueType # 4 + + class GroupType(_GroupType, metaclass=_GroupTypeEnumTypeWrapper): ... + GROUP_TYPE_UNSPECIFIED: Aggregate.GroupType.ValueType # 0 + GROUP_TYPE_GROUPBY: Aggregate.GroupType.ValueType # 1 + GROUP_TYPE_ROLLUP: Aggregate.GroupType.ValueType # 2 + GROUP_TYPE_CUBE: Aggregate.GroupType.ValueType # 3 + GROUP_TYPE_PIVOT: Aggregate.GroupType.ValueType # 4 + + class Pivot(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + COL_FIELD_NUMBER: builtins.int + VALUES_FIELD_NUMBER: builtins.int + @property + def col(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: + """(Required) The column to pivot""" + @property + def values( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Optional) List of values that will be translated to columns in the output DataFrame. + + Note that if it is empty, the server side will immediately trigger a job to collect + the distinct values of the column. + """ + def __init__( + self, + *, + col: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + values: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["col", b"col"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["col", b"col", "values", b"values"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + GROUP_TYPE_FIELD_NUMBER: builtins.int + GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + AGGREGATE_EXPRESSIONS_FIELD_NUMBER: builtins.int + PIVOT_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a RelationalGroupedDataset.""" + group_type: global___Aggregate.GroupType.ValueType + """(Required) How the RelationalGroupedDataset was built.""" + @property + def grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) Expressions for grouping keys""" + @property + def aggregate_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) List of values that will be translated to columns in the output DataFrame.""" + @property + def pivot(self) -> global___Aggregate.Pivot: + """(Optional) Pivots a column of the current `DataFrame` and performs the specified aggregation.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + group_type: global___Aggregate.GroupType.ValueType = ..., + grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + aggregate_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + pivot: global___Aggregate.Pivot | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input", "pivot", b"pivot"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "aggregate_expressions", + b"aggregate_expressions", + "group_type", + b"group_type", + "grouping_expressions", + b"grouping_expressions", + "input", + b"input", + "pivot", + b"pivot", + ], + ) -> None: ... + +global___Aggregate = Aggregate + +class Sort(google.protobuf.message.Message): + """Relation of type [[Sort]].""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + ORDER_FIELD_NUMBER: builtins.int + IS_GLOBAL_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a Sort.""" + @property + def order( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.SortOrder + ]: + """(Required) The ordering expressions""" + is_global: builtins.bool + """(Optional) if this is a global sort.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + order: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.SortOrder + ] + | None = ..., + is_global: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_is_global", b"_is_global", "input", b"input", "is_global", b"is_global" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_is_global", + b"_is_global", + "input", + b"input", + "is_global", + b"is_global", + "order", + b"order", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_is_global", b"_is_global"] + ) -> typing_extensions.Literal["is_global"] | None: ... + +global___Sort = Sort + +class Drop(google.protobuf.message.Message): + """Drop specified columns.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLUMNS_FIELD_NUMBER: builtins.int + COLUMN_NAMES_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def columns( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) columns to drop.""" + @property + def column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) names of columns to drop.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + columns: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + column_names: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "column_names", b"column_names", "columns", b"columns", "input", b"input" + ], + ) -> None: ... + +global___Drop = Drop + +class Deduplicate(google.protobuf.message.Message): + """Relation of type [[Deduplicate]] which have duplicate rows removed, could consider either only + the subset of columns or all the columns. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLUMN_NAMES_FIELD_NUMBER: builtins.int + ALL_COLUMNS_AS_KEYS_FIELD_NUMBER: builtins.int + WITHIN_WATERMARK_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a Deduplicate.""" + @property + def column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Deduplicate based on a list of column names. + + This field does not co-use with `all_columns_as_keys`. + """ + all_columns_as_keys: builtins.bool + """(Optional) Deduplicate based on all the columns of the input relation. + + This field does not co-use with `column_names`. + """ + within_watermark: builtins.bool + """(Optional) Deduplicate within the time range of watermark.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + column_names: collections.abc.Iterable[builtins.str] | None = ..., + all_columns_as_keys: builtins.bool | None = ..., + within_watermark: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_all_columns_as_keys", + b"_all_columns_as_keys", + "_within_watermark", + b"_within_watermark", + "all_columns_as_keys", + b"all_columns_as_keys", + "input", + b"input", + "within_watermark", + b"within_watermark", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_all_columns_as_keys", + b"_all_columns_as_keys", + "_within_watermark", + b"_within_watermark", + "all_columns_as_keys", + b"all_columns_as_keys", + "column_names", + b"column_names", + "input", + b"input", + "within_watermark", + b"within_watermark", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_all_columns_as_keys", b"_all_columns_as_keys"], + ) -> typing_extensions.Literal["all_columns_as_keys"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_within_watermark", b"_within_watermark"] + ) -> typing_extensions.Literal["within_watermark"] | None: ... + +global___Deduplicate = Deduplicate + +class LocalRelation(google.protobuf.message.Message): + """A relation that does not need to be qualified by name.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DATA_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + data: builtins.bytes + """(Optional) Local collection data serialized into Arrow IPC streaming format which contains + the schema of the data. + """ + schema: builtins.str + """(Optional) The schema of local data. + It should be either a DDL-formatted type string or a JSON string. + + The server side will update the column names and data types according to this schema. + If the 'data' is not provided, then this schema will be required. + """ + def __init__( + self, + *, + data: builtins.bytes | None = ..., + schema: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_data", b"_data", "_schema", b"_schema", "data", b"data", "schema", b"schema" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_data", b"_data", "_schema", b"_schema", "data", b"data", "schema", b"schema" + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_data", b"_data"] + ) -> typing_extensions.Literal["data"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_schema", b"_schema"] + ) -> typing_extensions.Literal["schema"] | None: ... + +global___LocalRelation = LocalRelation + +class CachedLocalRelation(google.protobuf.message.Message): + """A local relation that has been cached already.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USERID_FIELD_NUMBER: builtins.int + SESSIONID_FIELD_NUMBER: builtins.int + HASH_FIELD_NUMBER: builtins.int + userId: builtins.str + """(Required) An identifier of the user which created the local relation""" + sessionId: builtins.str + """(Required) An identifier of the Spark SQL session in which the user created the local relation.""" + hash: builtins.str + """(Required) A sha-256 hash of the serialized local relation in proto, see LocalRelation.""" + def __init__( + self, + *, + userId: builtins.str = ..., + sessionId: builtins.str = ..., + hash: builtins.str = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "hash", b"hash", "sessionId", b"sessionId", "userId", b"userId" + ], + ) -> None: ... + +global___CachedLocalRelation = CachedLocalRelation + +class CachedRemoteRelation(google.protobuf.message.Message): + """Represents a remote relation that has been cached on server.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RELATION_ID_FIELD_NUMBER: builtins.int + relation_id: builtins.str + """(Required) ID of the remote related (assigned by the service).""" + def __init__( + self, + *, + relation_id: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["relation_id", b"relation_id"] + ) -> None: ... + +global___CachedRemoteRelation = CachedRemoteRelation + +class Sample(google.protobuf.message.Message): + """Relation of type [[Sample]] that samples a fraction of the dataset.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + LOWER_BOUND_FIELD_NUMBER: builtins.int + UPPER_BOUND_FIELD_NUMBER: builtins.int + WITH_REPLACEMENT_FIELD_NUMBER: builtins.int + SEED_FIELD_NUMBER: builtins.int + DETERMINISTIC_ORDER_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a Sample.""" + lower_bound: builtins.float + """(Required) lower bound.""" + upper_bound: builtins.float + """(Required) upper bound.""" + with_replacement: builtins.bool + """(Optional) Whether to sample with replacement.""" + seed: builtins.int + """(Optional) The random seed.""" + deterministic_order: builtins.bool + """(Required) Explicitly sort the underlying plan to make the ordering deterministic or cache it. + This flag is true when invoking `dataframe.randomSplit` to randomly splits DataFrame with the + provided weights. Otherwise, it is false. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + lower_bound: builtins.float = ..., + upper_bound: builtins.float = ..., + with_replacement: builtins.bool | None = ..., + seed: builtins.int | None = ..., + deterministic_order: builtins.bool = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_seed", + b"_seed", + "_with_replacement", + b"_with_replacement", + "input", + b"input", + "seed", + b"seed", + "with_replacement", + b"with_replacement", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_seed", + b"_seed", + "_with_replacement", + b"_with_replacement", + "deterministic_order", + b"deterministic_order", + "input", + b"input", + "lower_bound", + b"lower_bound", + "seed", + b"seed", + "upper_bound", + b"upper_bound", + "with_replacement", + b"with_replacement", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_seed", b"_seed"] + ) -> typing_extensions.Literal["seed"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_with_replacement", b"_with_replacement"] + ) -> typing_extensions.Literal["with_replacement"] | None: ... + +global___Sample = Sample + +class Range(google.protobuf.message.Message): + """Relation of type [[Range]] that generates a sequence of integers.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + START_FIELD_NUMBER: builtins.int + END_FIELD_NUMBER: builtins.int + STEP_FIELD_NUMBER: builtins.int + NUM_PARTITIONS_FIELD_NUMBER: builtins.int + start: builtins.int + """(Optional) Default value = 0""" + end: builtins.int + """(Required)""" + step: builtins.int + """(Required)""" + num_partitions: builtins.int + """Optional. Default value is assigned by 1) SQL conf "spark.sql.leafNodeDefaultParallelism" if + it is set, or 2) spark default parallelism. + """ + def __init__( + self, + *, + start: builtins.int | None = ..., + end: builtins.int = ..., + step: builtins.int = ..., + num_partitions: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_num_partitions", + b"_num_partitions", + "_start", + b"_start", + "num_partitions", + b"num_partitions", + "start", + b"start", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_num_partitions", + b"_num_partitions", + "_start", + b"_start", + "end", + b"end", + "num_partitions", + b"num_partitions", + "start", + b"start", + "step", + b"step", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_num_partitions", b"_num_partitions"] + ) -> typing_extensions.Literal["num_partitions"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_start", b"_start"] + ) -> typing_extensions.Literal["start"] | None: ... + +global___Range = Range + +class SubqueryAlias(google.protobuf.message.Message): + """Relation alias.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + ALIAS_FIELD_NUMBER: builtins.int + QUALIFIER_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation of SubqueryAlias.""" + alias: builtins.str + """(Required) The alias.""" + @property + def qualifier( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Qualifier of the alias.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + alias: builtins.str = ..., + qualifier: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "alias", b"alias", "input", b"input", "qualifier", b"qualifier" + ], + ) -> None: ... + +global___SubqueryAlias = SubqueryAlias + +class Repartition(google.protobuf.message.Message): + """Relation repartition.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NUM_PARTITIONS_FIELD_NUMBER: builtins.int + SHUFFLE_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation of Repartition.""" + num_partitions: builtins.int + """(Required) Must be positive.""" + shuffle: builtins.bool + """(Optional) Default value is false.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + num_partitions: builtins.int = ..., + shuffle: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_shuffle", b"_shuffle", "input", b"input", "shuffle", b"shuffle" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_shuffle", + b"_shuffle", + "input", + b"input", + "num_partitions", + b"num_partitions", + "shuffle", + b"shuffle", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_shuffle", b"_shuffle"] + ) -> typing_extensions.Literal["shuffle"] | None: ... + +global___Repartition = Repartition + +class ShowString(google.protobuf.message.Message): + """Compose the string representing rows for output. + It will invoke 'Dataset.showString' to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NUM_ROWS_FIELD_NUMBER: builtins.int + TRUNCATE_FIELD_NUMBER: builtins.int + VERTICAL_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + num_rows: builtins.int + """(Required) Number of rows to show.""" + truncate: builtins.int + """(Required) If set to more than 0, truncates strings to + `truncate` characters and all cells will be aligned right. + """ + vertical: builtins.bool + """(Required) If set to true, prints output rows vertically (one line per column value).""" + def __init__( + self, + *, + input: global___Relation | None = ..., + num_rows: builtins.int = ..., + truncate: builtins.int = ..., + vertical: builtins.bool = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", + b"input", + "num_rows", + b"num_rows", + "truncate", + b"truncate", + "vertical", + b"vertical", + ], + ) -> None: ... + +global___ShowString = ShowString + +class HtmlString(google.protobuf.message.Message): + """Compose the string representing rows for output. + It will invoke 'Dataset.htmlString' to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NUM_ROWS_FIELD_NUMBER: builtins.int + TRUNCATE_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + num_rows: builtins.int + """(Required) Number of rows to show.""" + truncate: builtins.int + """(Required) If set to more than 0, truncates strings to + `truncate` characters and all cells will be aligned right. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + num_rows: builtins.int = ..., + truncate: builtins.int = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", b"input", "num_rows", b"num_rows", "truncate", b"truncate" + ], + ) -> None: ... + +global___HtmlString = HtmlString + +class StatSummary(google.protobuf.message.Message): + """Computes specified statistics for numeric and string columns. + It will invoke 'Dataset.summary' (same as 'StatFunctions.summary') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + STATISTICS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def statistics( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Statistics from to be computed. + + Available statistics are: + count + mean + stddev + min + max + arbitrary approximate percentiles specified as a percentage (e.g. 75%) + count_distinct + approx_count_distinct + + If no statistics are given, this function computes 'count', 'mean', 'stddev', 'min', + 'approximate quartiles' (percentiles at 25%, 50%, and 75%), and 'max'. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + statistics: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["input", b"input", "statistics", b"statistics"] + ) -> None: ... + +global___StatSummary = StatSummary + +class StatDescribe(google.protobuf.message.Message): + """Computes basic statistics for numeric and string columns, including count, mean, stddev, min, + and max. If no columns are given, this function computes statistics for all numerical or + string columns. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Columns to compute statistics on.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["cols", b"cols", "input", b"input"] + ) -> None: ... + +global___StatDescribe = StatDescribe + +class StatCrosstab(google.protobuf.message.Message): + """Computes a pair-wise frequency table of the given columns. Also known as a contingency table. + It will invoke 'Dataset.stat.crosstab' (same as 'StatFunctions.crossTabulate') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COL1_FIELD_NUMBER: builtins.int + COL2_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + col1: builtins.str + """(Required) The name of the first column. + + Distinct items will make the first item of each row. + """ + col2: builtins.str + """(Required) The name of the second column. + + Distinct items will make the column names of the DataFrame. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + col1: builtins.str = ..., + col2: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["col1", b"col1", "col2", b"col2", "input", b"input"], + ) -> None: ... + +global___StatCrosstab = StatCrosstab + +class StatCov(google.protobuf.message.Message): + """Calculate the sample covariance of two numerical columns of a DataFrame. + It will invoke 'Dataset.stat.cov' (same as 'StatFunctions.calculateCov') to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COL1_FIELD_NUMBER: builtins.int + COL2_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + col1: builtins.str + """(Required) The name of the first column.""" + col2: builtins.str + """(Required) The name of the second column.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + col1: builtins.str = ..., + col2: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["col1", b"col1", "col2", b"col2", "input", b"input"], + ) -> None: ... + +global___StatCov = StatCov + +class StatCorr(google.protobuf.message.Message): + """Calculates the correlation of two columns of a DataFrame. Currently only supports the Pearson + Correlation Coefficient. It will invoke 'Dataset.stat.corr' (same as + 'StatFunctions.pearsonCorrelation') to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COL1_FIELD_NUMBER: builtins.int + COL2_FIELD_NUMBER: builtins.int + METHOD_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + col1: builtins.str + """(Required) The name of the first column.""" + col2: builtins.str + """(Required) The name of the second column.""" + method: builtins.str + """(Optional) Default value is 'pearson'. + + Currently only supports the Pearson Correlation Coefficient. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + col1: builtins.str = ..., + col2: builtins.str = ..., + method: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_method", b"_method", "input", b"input", "method", b"method" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_method", + b"_method", + "col1", + b"col1", + "col2", + b"col2", + "input", + b"input", + "method", + b"method", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_method", b"_method"] + ) -> typing_extensions.Literal["method"] | None: ... + +global___StatCorr = StatCorr + +class StatApproxQuantile(google.protobuf.message.Message): + """Calculates the approximate quantiles of numerical columns of a DataFrame. + It will invoke 'Dataset.stat.approxQuantile' (same as 'StatFunctions.approxQuantile') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + PROBABILITIES_FIELD_NUMBER: builtins.int + RELATIVE_ERROR_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The names of the numerical columns.""" + @property + def probabilities( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]: + """(Required) A list of quantile probabilities. + + Each number must belong to [0, 1]. + For example 0 is the minimum, 0.5 is the median, 1 is the maximum. + """ + relative_error: builtins.float + """(Required) The relative target precision to achieve (greater than or equal to 0). + + If set to zero, the exact quantiles are computed, which could be very expensive. + Note that values greater than 1 are accepted but give the same result as 1. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + probabilities: collections.abc.Iterable[builtins.float] | None = ..., + relative_error: builtins.float = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "cols", + b"cols", + "input", + b"input", + "probabilities", + b"probabilities", + "relative_error", + b"relative_error", + ], + ) -> None: ... + +global___StatApproxQuantile = StatApproxQuantile + +class StatFreqItems(google.protobuf.message.Message): + """Finding frequent items for columns, possibly with false positives. + It will invoke 'Dataset.stat.freqItems' (same as 'StatFunctions.freqItems') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + SUPPORT_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) The names of the columns to search frequent items in.""" + support: builtins.float + """(Optional) The minimum frequency for an item to be considered `frequent`. + Should be greater than 1e-4. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + support: builtins.float | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_support", b"_support", "input", b"input", "support", b"support" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_support", b"_support", "cols", b"cols", "input", b"input", "support", b"support" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_support", b"_support"] + ) -> typing_extensions.Literal["support"] | None: ... + +global___StatFreqItems = StatFreqItems + +class StatSampleBy(google.protobuf.message.Message): + """Returns a stratified sample without replacement based on the fraction + given on each stratum. + It will invoke 'Dataset.stat.freqItems' (same as 'StatFunctions.freqItems') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Fraction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STRATUM_FIELD_NUMBER: builtins.int + FRACTION_FIELD_NUMBER: builtins.int + @property + def stratum(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression.Literal: + """(Required) The stratum.""" + fraction: builtins.float + """(Required) The fraction value. Must be in [0, 1].""" + def __init__( + self, + *, + stratum: pyspark.sql.connect.proto.expressions_pb2.Expression.Literal | None = ..., + fraction: builtins.float = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["stratum", b"stratum"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["fraction", b"fraction", "stratum", b"stratum"], + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + COL_FIELD_NUMBER: builtins.int + FRACTIONS_FIELD_NUMBER: builtins.int + SEED_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def col(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression: + """(Required) The column that defines strata.""" + @property + def fractions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___StatSampleBy.Fraction + ]: + """(Required) Sampling fraction for each stratum. + + If a stratum is not specified, we treat its fraction as zero. + """ + seed: builtins.int + """(Optional) The random seed.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + col: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ..., + fractions: collections.abc.Iterable[global___StatSampleBy.Fraction] | None = ..., + seed: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_seed", b"_seed", "col", b"col", "input", b"input", "seed", b"seed" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_seed", + b"_seed", + "col", + b"col", + "fractions", + b"fractions", + "input", + b"input", + "seed", + b"seed", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_seed", b"_seed"] + ) -> typing_extensions.Literal["seed"] | None: ... + +global___StatSampleBy = StatSampleBy + +class NAFill(google.protobuf.message.Message): + """Replaces null values. + It will invoke 'Dataset.na.fill' (same as 'DataFrameNaFunctions.fill') to compute the results. + Following 3 parameter combinations are supported: + 1, 'values' only contains 1 item, 'cols' is empty: + replaces null values in all type-compatible columns. + 2, 'values' only contains 1 item, 'cols' is not empty: + replaces null values in specified columns. + 3, 'values' contains more than 1 items, then 'cols' is required to have the same length: + replaces each specified column with corresponding value. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + VALUES_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Optional list of column names to consider.""" + @property + def values( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ]: + """(Required) Values to replace null values with. + + Should contain at least 1 item. + Only 4 data types are supported now: bool, long, double, string + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + values: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Literal + ] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "cols", b"cols", "input", b"input", "values", b"values" + ], + ) -> None: ... + +global___NAFill = NAFill + +class NADrop(google.protobuf.message.Message): + """Drop rows containing null values. + It will invoke 'Dataset.na.drop' (same as 'DataFrameNaFunctions.drop') to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + MIN_NON_NULLS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) Optional list of column names to consider. + + When it is empty, all the columns in the input relation will be considered. + """ + min_non_nulls: builtins.int + """(Optional) The minimum number of non-null and non-NaN values required to keep. + + When not set, it is equivalent to the number of considered columns, which means + a row will be kept only if all columns are non-null. + + 'how' options ('all', 'any') can be easily converted to this field: + - 'all' -> set 'min_non_nulls' 1; + - 'any' -> keep 'min_non_nulls' unset; + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + min_non_nulls: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_min_non_nulls", + b"_min_non_nulls", + "input", + b"input", + "min_non_nulls", + b"min_non_nulls", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_min_non_nulls", + b"_min_non_nulls", + "cols", + b"cols", + "input", + b"input", + "min_non_nulls", + b"min_non_nulls", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_min_non_nulls", b"_min_non_nulls"] + ) -> typing_extensions.Literal["min_non_nulls"] | None: ... + +global___NADrop = NADrop + +class NAReplace(google.protobuf.message.Message): + """Replaces old values with the corresponding values. + It will invoke 'Dataset.na.replace' (same as 'DataFrameNaFunctions.replace') + to compute the results. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Replacement(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + OLD_VALUE_FIELD_NUMBER: builtins.int + NEW_VALUE_FIELD_NUMBER: builtins.int + @property + def old_value(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression.Literal: + """(Required) The old value. + + Only 4 data types are supported now: null, bool, double, string. + """ + @property + def new_value(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression.Literal: + """(Required) The new value. + + Should be of the same data type with the old value. + """ + def __init__( + self, + *, + old_value: pyspark.sql.connect.proto.expressions_pb2.Expression.Literal | None = ..., + new_value: pyspark.sql.connect.proto.expressions_pb2.Expression.Literal | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "new_value", b"new_value", "old_value", b"old_value" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "new_value", b"new_value", "old_value", b"old_value" + ], + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + COLS_FIELD_NUMBER: builtins.int + REPLACEMENTS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def cols( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Optional) List of column names to consider. + + When it is empty, all the type-compatible columns in the input relation will be considered. + """ + @property + def replacements( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___NAReplace.Replacement + ]: + """(Optional) The value replacement mapping.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + cols: collections.abc.Iterable[builtins.str] | None = ..., + replacements: collections.abc.Iterable[global___NAReplace.Replacement] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "cols", b"cols", "input", b"input", "replacements", b"replacements" + ], + ) -> None: ... + +global___NAReplace = NAReplace + +class ToDF(google.protobuf.message.Message): + """Rename columns on the input relation by the same length of names.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + COLUMN_NAMES_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation of RenameColumnsBySameLengthNames.""" + @property + def column_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """(Required) + + The number of columns of the input relation must be equal to the length + of this field. If this is not true, an exception will be returned. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + column_names: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal["column_names", b"column_names", "input", b"input"], + ) -> None: ... + +global___ToDF = ToDF + +class WithColumnsRenamed(google.protobuf.message.Message): + """Rename columns on the input relation by a map with name to name mapping.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class RenameColumnsMapEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + RENAME_COLUMNS_MAP_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def rename_columns_map( + self, + ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """(Required) + + Renaming column names of input relation from A to B where A is the map key + and B is the map value. This is a no-op if schema doesn't contain any A. It + does not require that all input relation column names to present as keys. + duplicated B are not allowed. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + rename_columns_map: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", b"input", "rename_columns_map", b"rename_columns_map" + ], + ) -> None: ... + +global___WithColumnsRenamed = WithColumnsRenamed + +class WithColumns(google.protobuf.message.Message): + """Adding columns or replacing the existing columns that have the same names.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + ALIASES_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def aliases( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Alias + ]: + """(Required) + + Given a column name, apply the corresponding expression on the column. If column + name exists in the input relation, then replace the column. If the column name + does not exist in the input relation, then adds it as a new column. + + Only one name part is expected from each Expression.Alias. + + An exception is thrown when duplicated names are present in the mapping. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + aliases: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression.Alias + ] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["aliases", b"aliases", "input", b"input"] + ) -> None: ... + +global___WithColumns = WithColumns + +class WithWatermark(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + EVENT_TIME_FIELD_NUMBER: builtins.int + DELAY_THRESHOLD_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation""" + event_time: builtins.str + """(Required) Name of the column containing event time.""" + delay_threshold: builtins.str + """(Required)""" + def __init__( + self, + *, + input: global___Relation | None = ..., + event_time: builtins.str = ..., + delay_threshold: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "delay_threshold", b"delay_threshold", "event_time", b"event_time", "input", b"input" + ], + ) -> None: ... + +global___WithWatermark = WithWatermark + +class Hint(google.protobuf.message.Message): + """Specify a hint over a relation. Hint should have a name and optional parameters.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + PARAMETERS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + name: builtins.str + """(Required) Hint name. + + Supported Join hints include BROADCAST, MERGE, SHUFFLE_HASH, SHUFFLE_REPLICATE_NL. + + Supported partitioning hints include COALESCE, REPARTITION, REPARTITION_BY_RANGE. + """ + @property + def parameters( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Hint parameters.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + name: builtins.str = ..., + parameters: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", b"input", "name", b"name", "parameters", b"parameters" + ], + ) -> None: ... + +global___Hint = Hint + +class Unpivot(google.protobuf.message.Message): + """Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns set.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Values(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VALUES_FIELD_NUMBER: builtins.int + @property + def values( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: ... + def __init__( + self, + *, + values: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["values", b"values"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + IDS_FIELD_NUMBER: builtins.int + VALUES_FIELD_NUMBER: builtins.int + VARIABLE_COLUMN_NAME_FIELD_NUMBER: builtins.int + VALUE_COLUMN_NAME_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def ids( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) Id columns.""" + @property + def values(self) -> global___Unpivot.Values: + """(Optional) Value columns to unpivot.""" + variable_column_name: builtins.str + """(Required) Name of the variable column.""" + value_column_name: builtins.str + """(Required) Name of the value column.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + ids: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + values: global___Unpivot.Values | None = ..., + variable_column_name: builtins.str = ..., + value_column_name: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_values", b"_values", "input", b"input", "values", b"values" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_values", + b"_values", + "ids", + b"ids", + "input", + b"input", + "value_column_name", + b"value_column_name", + "values", + b"values", + "variable_column_name", + b"variable_column_name", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_values", b"_values"] + ) -> typing_extensions.Literal["values"] | None: ... + +global___Unpivot = Unpivot + +class ToSchema(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Required) The user provided schema. + + The Sever side will update the dataframe with this schema. + """ + def __init__( + self, + *, + input: global___Relation | None = ..., + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input", "schema", b"schema"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["input", b"input", "schema", b"schema"] + ) -> None: ... + +global___ToSchema = ToSchema + +class RepartitionByExpression(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + PARTITION_EXPRS_FIELD_NUMBER: builtins.int + NUM_PARTITIONS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + @property + def partition_exprs( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) The partitioning expressions.""" + num_partitions: builtins.int + """(Optional) number of partitions, must be positive.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + partition_exprs: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + num_partitions: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_num_partitions", + b"_num_partitions", + "input", + b"input", + "num_partitions", + b"num_partitions", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_num_partitions", + b"_num_partitions", + "input", + b"input", + "num_partitions", + b"num_partitions", + "partition_exprs", + b"partition_exprs", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_num_partitions", b"_num_partitions"] + ) -> typing_extensions.Literal["num_partitions"] | None: ... + +global___RepartitionByExpression = RepartitionByExpression + +class MapPartitions(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + FUNC_FIELD_NUMBER: builtins.int + IS_BARRIER_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for a mapPartitions-equivalent API: mapInPandas, mapInArrow.""" + @property + def func(self) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: + """(Required) Input user-defined function.""" + is_barrier: builtins.bool + """(Optional) Whether to use barrier mode execution or not.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + func: pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction + | None = ..., + is_barrier: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_is_barrier", + b"_is_barrier", + "func", + b"func", + "input", + b"input", + "is_barrier", + b"is_barrier", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_is_barrier", + b"_is_barrier", + "func", + b"func", + "input", + b"input", + "is_barrier", + b"is_barrier", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_is_barrier", b"_is_barrier"] + ) -> typing_extensions.Literal["is_barrier"] | None: ... + +global___MapPartitions = MapPartitions + +class GroupMap(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + FUNC_FIELD_NUMBER: builtins.int + SORTING_EXPRESSIONS_FIELD_NUMBER: builtins.int + INITIAL_INPUT_FIELD_NUMBER: builtins.int + INITIAL_GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + IS_MAP_GROUPS_WITH_STATE_FIELD_NUMBER: builtins.int + OUTPUT_MODE_FIELD_NUMBER: builtins.int + TIMEOUT_CONF_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for Group Map API: apply, applyInPandas.""" + @property + def grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) Expressions for grouping keys.""" + @property + def func(self) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: + """(Required) Input user-defined function.""" + @property + def sorting_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Expressions for sorting. Only used by Scala Sorted Group Map API.""" + @property + def initial_input(self) -> global___Relation: + """Below fields are only used by (Flat)MapGroupsWithState + (Optional) Input relation for initial State. + """ + @property + def initial_grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Expressions for grouping keys of the initial state input relation.""" + is_map_groups_with_state: builtins.bool + """(Optional) True if MapGroupsWithState, false if FlatMapGroupsWithState.""" + output_mode: builtins.str + """(Optional) The output mode of the function.""" + timeout_conf: builtins.str + """(Optional) Timeout configuration for groups that do not receive data for a while.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + func: pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction + | None = ..., + sorting_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + initial_input: global___Relation | None = ..., + initial_grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + is_map_groups_with_state: builtins.bool | None = ..., + output_mode: builtins.str | None = ..., + timeout_conf: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_is_map_groups_with_state", + b"_is_map_groups_with_state", + "_output_mode", + b"_output_mode", + "_timeout_conf", + b"_timeout_conf", + "func", + b"func", + "initial_input", + b"initial_input", + "input", + b"input", + "is_map_groups_with_state", + b"is_map_groups_with_state", + "output_mode", + b"output_mode", + "timeout_conf", + b"timeout_conf", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_is_map_groups_with_state", + b"_is_map_groups_with_state", + "_output_mode", + b"_output_mode", + "_timeout_conf", + b"_timeout_conf", + "func", + b"func", + "grouping_expressions", + b"grouping_expressions", + "initial_grouping_expressions", + b"initial_grouping_expressions", + "initial_input", + b"initial_input", + "input", + b"input", + "is_map_groups_with_state", + b"is_map_groups_with_state", + "output_mode", + b"output_mode", + "sorting_expressions", + b"sorting_expressions", + "timeout_conf", + b"timeout_conf", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_is_map_groups_with_state", b"_is_map_groups_with_state" + ], + ) -> typing_extensions.Literal["is_map_groups_with_state"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_output_mode", b"_output_mode"] + ) -> typing_extensions.Literal["output_mode"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_timeout_conf", b"_timeout_conf"] + ) -> typing_extensions.Literal["timeout_conf"] | None: ... + +global___GroupMap = GroupMap + +class CoGroupMap(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + INPUT_GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + OTHER_FIELD_NUMBER: builtins.int + OTHER_GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + FUNC_FIELD_NUMBER: builtins.int + INPUT_SORTING_EXPRESSIONS_FIELD_NUMBER: builtins.int + OTHER_SORTING_EXPRESSIONS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) One input relation for CoGroup Map API - applyInPandas.""" + @property + def input_grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """Expressions for grouping keys of the first input relation.""" + @property + def other(self) -> global___Relation: + """(Required) The other input relation.""" + @property + def other_grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """Expressions for grouping keys of the other input relation.""" + @property + def func(self) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: + """(Required) Input user-defined function.""" + @property + def input_sorting_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Expressions for sorting. Only used by Scala Sorted CoGroup Map API.""" + @property + def other_sorting_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Expressions for sorting. Only used by Scala Sorted CoGroup Map API.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + input_grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + other: global___Relation | None = ..., + other_grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + func: pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction + | None = ..., + input_sorting_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + other_sorting_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "func", b"func", "input", b"input", "other", b"other" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "func", + b"func", + "input", + b"input", + "input_grouping_expressions", + b"input_grouping_expressions", + "input_sorting_expressions", + b"input_sorting_expressions", + "other", + b"other", + "other_grouping_expressions", + b"other_grouping_expressions", + "other_sorting_expressions", + b"other_sorting_expressions", + ], + ) -> None: ... + +global___CoGroupMap = CoGroupMap + +class ApplyInPandasWithState(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + GROUPING_EXPRESSIONS_FIELD_NUMBER: builtins.int + FUNC_FIELD_NUMBER: builtins.int + OUTPUT_SCHEMA_FIELD_NUMBER: builtins.int + STATE_SCHEMA_FIELD_NUMBER: builtins.int + OUTPUT_MODE_FIELD_NUMBER: builtins.int + TIMEOUT_CONF_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation for applyInPandasWithState.""" + @property + def grouping_expressions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) Expressions for grouping keys.""" + @property + def func(self) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: + """(Required) Input user-defined function.""" + output_schema: builtins.str + """(Required) Schema for the output DataFrame.""" + state_schema: builtins.str + """(Required) Schema for the state.""" + output_mode: builtins.str + """(Required) The output mode of the function.""" + timeout_conf: builtins.str + """(Required) Timeout configuration for groups that do not receive data for a while.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + grouping_expressions: collections.abc.Iterable[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ] + | None = ..., + func: pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction + | None = ..., + output_schema: builtins.str = ..., + state_schema: builtins.str = ..., + output_mode: builtins.str = ..., + timeout_conf: builtins.str = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["func", b"func", "input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "func", + b"func", + "grouping_expressions", + b"grouping_expressions", + "input", + b"input", + "output_mode", + b"output_mode", + "output_schema", + b"output_schema", + "state_schema", + b"state_schema", + "timeout_conf", + b"timeout_conf", + ], + ) -> None: ... + +global___ApplyInPandasWithState = ApplyInPandasWithState + +class CommonInlineUserDefinedTableFunction(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FUNCTION_NAME_FIELD_NUMBER: builtins.int + DETERMINISTIC_FIELD_NUMBER: builtins.int + ARGUMENTS_FIELD_NUMBER: builtins.int + PYTHON_UDTF_FIELD_NUMBER: builtins.int + function_name: builtins.str + """(Required) Name of the user-defined table function.""" + deterministic: builtins.bool + """(Optional) Whether the user-defined table function is deterministic.""" + @property + def arguments( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Optional) Function input arguments. Empty arguments are allowed.""" + @property + def python_udtf(self) -> global___PythonUDTF: ... + def __init__( + self, + *, + function_name: builtins.str = ..., + deterministic: builtins.bool = ..., + arguments: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + python_udtf: global___PythonUDTF | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "function", b"function", "python_udtf", b"python_udtf" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "arguments", + b"arguments", + "deterministic", + b"deterministic", + "function", + b"function", + "function_name", + b"function_name", + "python_udtf", + b"python_udtf", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["function", b"function"] + ) -> typing_extensions.Literal["python_udtf"] | None: ... + +global___CommonInlineUserDefinedTableFunction = CommonInlineUserDefinedTableFunction + +class PythonUDTF(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RETURN_TYPE_FIELD_NUMBER: builtins.int + EVAL_TYPE_FIELD_NUMBER: builtins.int + COMMAND_FIELD_NUMBER: builtins.int + PYTHON_VER_FIELD_NUMBER: builtins.int + @property + def return_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional) Return type of the Python UDTF.""" + eval_type: builtins.int + """(Required) EvalType of the Python UDTF.""" + command: builtins.bytes + """(Required) The encoded commands of the Python UDTF.""" + python_ver: builtins.str + """(Required) Python version being used in the client.""" + def __init__( + self, + *, + return_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + eval_type: builtins.int = ..., + command: builtins.bytes = ..., + python_ver: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_return_type", b"_return_type", "return_type", b"return_type" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_return_type", + b"_return_type", + "command", + b"command", + "eval_type", + b"eval_type", + "python_ver", + b"python_ver", + "return_type", + b"return_type", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_return_type", b"_return_type"] + ) -> typing_extensions.Literal["return_type"] | None: ... + +global___PythonUDTF = PythonUDTF + +class CollectMetrics(google.protobuf.message.Message): + """Collect arbitrary (named) metrics from a dataset.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + INPUT_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + METRICS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) The input relation.""" + name: builtins.str + """(Required) Name of the metrics.""" + @property + def metrics( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + pyspark.sql.connect.proto.expressions_pb2.Expression + ]: + """(Required) The metric sequence.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + name: builtins.str = ..., + metrics: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression] + | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["input", b"input"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "input", b"input", "metrics", b"metrics", "name", b"name" + ], + ) -> None: ... + +global___CollectMetrics = CollectMetrics + +class Parse(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class _ParseFormat: + ValueType = typing.NewType("ValueType", builtins.int) + V: typing_extensions.TypeAlias = ValueType + + class _ParseFormatEnumTypeWrapper( + google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Parse._ParseFormat.ValueType], + builtins.type, + ): # noqa: F821 + DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor + PARSE_FORMAT_UNSPECIFIED: Parse._ParseFormat.ValueType # 0 + PARSE_FORMAT_CSV: Parse._ParseFormat.ValueType # 1 + PARSE_FORMAT_JSON: Parse._ParseFormat.ValueType # 2 + + class ParseFormat(_ParseFormat, metaclass=_ParseFormatEnumTypeWrapper): ... + PARSE_FORMAT_UNSPECIFIED: Parse.ParseFormat.ValueType # 0 + PARSE_FORMAT_CSV: Parse.ParseFormat.ValueType # 1 + PARSE_FORMAT_JSON: Parse.ParseFormat.ValueType # 2 + + class OptionsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"] + ) -> None: ... + + INPUT_FIELD_NUMBER: builtins.int + FORMAT_FIELD_NUMBER: builtins.int + SCHEMA_FIELD_NUMBER: builtins.int + OPTIONS_FIELD_NUMBER: builtins.int + @property + def input(self) -> global___Relation: + """(Required) Input relation to Parse. The input is expected to have single text column.""" + format: global___Parse.ParseFormat.ValueType + """(Required) The expected format of the text.""" + @property + def schema(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional) DataType representing the schema. If not set, Spark will infer the schema.""" + @property + def options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Options for the csv/json parser. The map key is case insensitive.""" + def __init__( + self, + *, + input: global___Relation | None = ..., + format: global___Parse.ParseFormat.ValueType = ..., + schema: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., + options: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_schema", b"_schema", "input", b"input", "schema", b"schema" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_schema", + b"_schema", + "format", + b"format", + "input", + b"input", + "options", + b"options", + "schema", + b"schema", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_schema", b"_schema"] + ) -> typing_extensions.Literal["schema"] | None: ... + +global___Parse = Parse diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.py new file mode 100644 index 0000000..3619703 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.py @@ -0,0 +1,95 @@ +# +# 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. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: spark/connect/types.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b"\n\x19spark/connect/types.proto\x12\rspark.connect\"\xc7 \n\x08\x44\x61taType\x12\x32\n\x04null\x18\x01 \x01(\x0b\x32\x1c.spark.connect.DataType.NULLH\x00R\x04null\x12\x38\n\x06\x62inary\x18\x02 \x01(\x0b\x32\x1e.spark.connect.DataType.BinaryH\x00R\x06\x62inary\x12;\n\x07\x62oolean\x18\x03 \x01(\x0b\x32\x1f.spark.connect.DataType.BooleanH\x00R\x07\x62oolean\x12\x32\n\x04\x62yte\x18\x04 \x01(\x0b\x32\x1c.spark.connect.DataType.ByteH\x00R\x04\x62yte\x12\x35\n\x05short\x18\x05 \x01(\x0b\x32\x1d.spark.connect.DataType.ShortH\x00R\x05short\x12;\n\x07integer\x18\x06 \x01(\x0b\x32\x1f.spark.connect.DataType.IntegerH\x00R\x07integer\x12\x32\n\x04long\x18\x07 \x01(\x0b\x32\x1c.spark.connect.DataType.LongH\x00R\x04long\x12\x35\n\x05\x66loat\x18\x08 \x01(\x0b\x32\x1d.spark.connect.DataType.FloatH\x00R\x05\x66loat\x12\x38\n\x06\x64ouble\x18\t \x01(\x0b\x32\x1e.spark.connect.DataType.DoubleH\x00R\x06\x64ouble\x12;\n\x07\x64\x65\x63imal\x18\n \x01(\x0b\x32\x1f.spark.connect.DataType.DecimalH\x00R\x07\x64\x65\x63imal\x12\x38\n\x06string\x18\x0b \x01(\x0b\x32\x1e.spark.connect.DataType.StringH\x00R\x06string\x12\x32\n\x04\x63har\x18\x0c \x01(\x0b\x32\x1c.spark.connect.DataType.CharH\x00R\x04\x63har\x12<\n\x08var_char\x18\r \x01(\x0b\x32\x1f.spark.connect.DataType.VarCharH\x00R\x07varChar\x12\x32\n\x04\x64\x61te\x18\x0e \x01(\x0b\x32\x1c.spark.connect.DataType.DateH\x00R\x04\x64\x61te\x12\x41\n\ttimestamp\x18\x0f \x01(\x0b\x32!.spark.connect.DataType.TimestampH\x00R\ttimestamp\x12K\n\rtimestamp_ntz\x18\x10 \x01(\x0b\x32$.spark.connect.DataType.TimestampNTZH\x00R\x0ctimestampNtz\x12W\n\x11\x63\x61lendar_interval\x18\x11 \x01(\x0b\x32(.spark.connect.DataType.CalendarIntervalH\x00R\x10\x63\x61lendarInterval\x12[\n\x13year_month_interval\x18\x12 \x01(\x0b\x32).spark.connect.DataType.YearMonthIntervalH\x00R\x11yearMonthInterval\x12U\n\x11\x64\x61y_time_interval\x18\x13 \x01(\x0b\x32'.spark.connect.DataType.DayTimeIntervalH\x00R\x0f\x64\x61yTimeInterval\x12\x35\n\x05\x61rray\x18\x14 \x01(\x0b\x32\x1d.spark.connect.DataType.ArrayH\x00R\x05\x61rray\x12\x38\n\x06struct\x18\x15 \x01(\x0b\x32\x1e.spark.connect.DataType.StructH\x00R\x06struct\x12/\n\x03map\x18\x16 \x01(\x0b\x32\x1b.spark.connect.DataType.MapH\x00R\x03map\x12/\n\x03udt\x18\x17 \x01(\x0b\x32\x1b.spark.connect.DataType.UDTH\x00R\x03udt\x12>\n\x08unparsed\x18\x18 \x01(\x0b\x32 .spark.connect.DataType.UnparsedH\x00R\x08unparsed\x1a\x43\n\x07\x42oolean\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a@\n\x04\x42yte\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x41\n\x05Short\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x43\n\x07Integer\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a@\n\x04Long\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x41\n\x05\x46loat\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x42\n\x06\x44ouble\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x42\n\x06String\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x42\n\x06\x42inary\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a@\n\x04NULL\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\x45\n\tTimestamp\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a@\n\x04\x44\x61te\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1aH\n\x0cTimestampNTZ\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1aL\n\x10\x43\x61lendarInterval\x12\x38\n\x18type_variation_reference\x18\x01 \x01(\rR\x16typeVariationReference\x1a\xb3\x01\n\x11YearMonthInterval\x12$\n\x0bstart_field\x18\x01 \x01(\x05H\x00R\nstartField\x88\x01\x01\x12 \n\tend_field\x18\x02 \x01(\x05H\x01R\x08\x65ndField\x88\x01\x01\x12\x38\n\x18type_variation_reference\x18\x03 \x01(\rR\x16typeVariationReferenceB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1a\xb1\x01\n\x0f\x44\x61yTimeInterval\x12$\n\x0bstart_field\x18\x01 \x01(\x05H\x00R\nstartField\x88\x01\x01\x12 \n\tend_field\x18\x02 \x01(\x05H\x01R\x08\x65ndField\x88\x01\x01\x12\x38\n\x18type_variation_reference\x18\x03 \x01(\rR\x16typeVariationReferenceB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1aX\n\x04\x43har\x12\x16\n\x06length\x18\x01 \x01(\x05R\x06length\x12\x38\n\x18type_variation_reference\x18\x02 \x01(\rR\x16typeVariationReference\x1a[\n\x07VarChar\x12\x16\n\x06length\x18\x01 \x01(\x05R\x06length\x12\x38\n\x18type_variation_reference\x18\x02 \x01(\rR\x16typeVariationReference\x1a\x99\x01\n\x07\x44\x65\x63imal\x12\x19\n\x05scale\x18\x01 \x01(\x05H\x00R\x05scale\x88\x01\x01\x12!\n\tprecision\x18\x02 \x01(\x05H\x01R\tprecision\x88\x01\x01\x12\x38\n\x18type_variation_reference\x18\x03 \x01(\rR\x16typeVariationReferenceB\x08\n\x06_scaleB\x0c\n\n_precision\x1a\xa1\x01\n\x0bStructField\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x34\n\tdata_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x08\x64\x61taType\x12\x1a\n\x08nullable\x18\x03 \x01(\x08R\x08nullable\x12\x1f\n\x08metadata\x18\x04 \x01(\tH\x00R\x08metadata\x88\x01\x01\x42\x0b\n\t_metadata\x1a\x7f\n\x06Struct\x12;\n\x06\x66ields\x18\x01 \x03(\x0b\x32#.spark.connect.DataType.StructFieldR\x06\x66ields\x12\x38\n\x18type_variation_reference\x18\x02 \x01(\rR\x16typeVariationReference\x1a\xa2\x01\n\x05\x41rray\x12:\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x0b\x65lementType\x12#\n\rcontains_null\x18\x02 \x01(\x08R\x0c\x63ontainsNull\x12\x38\n\x18type_variation_reference\x18\x03 \x01(\rR\x16typeVariationReference\x1a\xdb\x01\n\x03Map\x12\x32\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x07keyType\x12\x36\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\tvalueType\x12.\n\x13value_contains_null\x18\x03 \x01(\x08R\x11valueContainsNull\x12\x38\n\x18type_variation_reference\x18\x04 \x01(\rR\x16typeVariationReference\x1a\x8f\x02\n\x03UDT\x12\x12\n\x04type\x18\x01 \x01(\tR\x04type\x12 \n\tjvm_class\x18\x02 \x01(\tH\x00R\x08jvmClass\x88\x01\x01\x12&\n\x0cpython_class\x18\x03 \x01(\tH\x01R\x0bpythonClass\x88\x01\x01\x12;\n\x17serialized_python_class\x18\x04 \x01(\tH\x02R\x15serializedPythonClass\x88\x01\x01\x12\x32\n\x08sql_type\x18\x05 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x07sqlTypeB\x0c\n\n_jvm_classB\x0f\n\r_python_classB\x1a\n\x18_serialized_python_class\x1a\x34\n\x08Unparsed\x12(\n\x10\x64\x61ta_type_string\x18\x01 \x01(\tR\x0e\x64\x61taTypeStringB\x06\n\x04kindB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3" +) + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "spark.connect.types_pb2", globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = ( + b"\n\036org.apache.spark.connect.protoP\001Z\022internal/generated" + ) + _DATATYPE._serialized_start = 45 + _DATATYPE._serialized_end = 4212 + _DATATYPE_BOOLEAN._serialized_start = 1534 + _DATATYPE_BOOLEAN._serialized_end = 1601 + _DATATYPE_BYTE._serialized_start = 1603 + _DATATYPE_BYTE._serialized_end = 1667 + _DATATYPE_SHORT._serialized_start = 1669 + _DATATYPE_SHORT._serialized_end = 1734 + _DATATYPE_INTEGER._serialized_start = 1736 + _DATATYPE_INTEGER._serialized_end = 1803 + _DATATYPE_LONG._serialized_start = 1805 + _DATATYPE_LONG._serialized_end = 1869 + _DATATYPE_FLOAT._serialized_start = 1871 + _DATATYPE_FLOAT._serialized_end = 1936 + _DATATYPE_DOUBLE._serialized_start = 1938 + _DATATYPE_DOUBLE._serialized_end = 2004 + _DATATYPE_STRING._serialized_start = 2006 + _DATATYPE_STRING._serialized_end = 2072 + _DATATYPE_BINARY._serialized_start = 2074 + _DATATYPE_BINARY._serialized_end = 2140 + _DATATYPE_NULL._serialized_start = 2142 + _DATATYPE_NULL._serialized_end = 2206 + _DATATYPE_TIMESTAMP._serialized_start = 2208 + _DATATYPE_TIMESTAMP._serialized_end = 2277 + _DATATYPE_DATE._serialized_start = 2279 + _DATATYPE_DATE._serialized_end = 2343 + _DATATYPE_TIMESTAMPNTZ._serialized_start = 2345 + _DATATYPE_TIMESTAMPNTZ._serialized_end = 2417 + _DATATYPE_CALENDARINTERVAL._serialized_start = 2419 + _DATATYPE_CALENDARINTERVAL._serialized_end = 2495 + _DATATYPE_YEARMONTHINTERVAL._serialized_start = 2498 + _DATATYPE_YEARMONTHINTERVAL._serialized_end = 2677 + _DATATYPE_DAYTIMEINTERVAL._serialized_start = 2680 + _DATATYPE_DAYTIMEINTERVAL._serialized_end = 2857 + _DATATYPE_CHAR._serialized_start = 2859 + _DATATYPE_CHAR._serialized_end = 2947 + _DATATYPE_VARCHAR._serialized_start = 2949 + _DATATYPE_VARCHAR._serialized_end = 3040 + _DATATYPE_DECIMAL._serialized_start = 3043 + _DATATYPE_DECIMAL._serialized_end = 3196 + _DATATYPE_STRUCTFIELD._serialized_start = 3199 + _DATATYPE_STRUCTFIELD._serialized_end = 3360 + _DATATYPE_STRUCT._serialized_start = 3362 + _DATATYPE_STRUCT._serialized_end = 3489 + _DATATYPE_ARRAY._serialized_start = 3492 + _DATATYPE_ARRAY._serialized_end = 3654 + _DATATYPE_MAP._serialized_start = 3657 + _DATATYPE_MAP._serialized_end = 3876 + _DATATYPE_UDT._serialized_start = 3879 + _DATATYPE_UDT._serialized_end = 4150 + _DATATYPE_UNPARSED._serialized_start = 4152 + _DATATYPE_UNPARSED._serialized_end = 4204 +# @@protoc_insertion_point(module_scope) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.pyi new file mode 100644 index 0000000..956701b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/proto/types_pb2.pyi @@ -0,0 +1,980 @@ +# +# 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. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +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 builtins +import collections.abc +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.message +import sys +import typing + +if sys.version_info >= (3, 8): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class DataType(google.protobuf.message.Message): + """This message describes the logical [[DataType]] of something. It does not carry the value + itself but only describes it. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class Boolean(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Byte(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Short(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Integer(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Long(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Float(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Double(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class String(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Binary(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class NULL(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Timestamp(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Date(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class TimestampNTZ(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class CalendarInterval(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class YearMonthInterval(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + START_FIELD_FIELD_NUMBER: builtins.int + END_FIELD_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + start_field: builtins.int + end_field: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + start_field: builtins.int | None = ..., + end_field: builtins.int | None = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_end_field", + b"_end_field", + "_start_field", + b"_start_field", + "end_field", + b"end_field", + "start_field", + b"start_field", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_end_field", + b"_end_field", + "_start_field", + b"_start_field", + "end_field", + b"end_field", + "start_field", + b"start_field", + "type_variation_reference", + b"type_variation_reference", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_end_field", b"_end_field"] + ) -> typing_extensions.Literal["end_field"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_start_field", b"_start_field"] + ) -> typing_extensions.Literal["start_field"] | None: ... + + class DayTimeInterval(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + START_FIELD_FIELD_NUMBER: builtins.int + END_FIELD_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + start_field: builtins.int + end_field: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + start_field: builtins.int | None = ..., + end_field: builtins.int | None = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_end_field", + b"_end_field", + "_start_field", + b"_start_field", + "end_field", + b"end_field", + "start_field", + b"start_field", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_end_field", + b"_end_field", + "_start_field", + b"_start_field", + "end_field", + b"end_field", + "start_field", + b"start_field", + "type_variation_reference", + b"type_variation_reference", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_end_field", b"_end_field"] + ) -> typing_extensions.Literal["end_field"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_start_field", b"_start_field"] + ) -> typing_extensions.Literal["start_field"] | None: ... + + class Char(google.protobuf.message.Message): + """Start compound types.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LENGTH_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + length: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + length: builtins.int = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "length", b"length", "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class VarChar(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LENGTH_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + length: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + length: builtins.int = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "length", b"length", "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Decimal(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + SCALE_FIELD_NUMBER: builtins.int + PRECISION_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + scale: builtins.int + precision: builtins.int + type_variation_reference: builtins.int + def __init__( + self, + *, + scale: builtins.int | None = ..., + precision: builtins.int | None = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_precision", + b"_precision", + "_scale", + b"_scale", + "precision", + b"precision", + "scale", + b"scale", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_precision", + b"_precision", + "_scale", + b"_scale", + "precision", + b"precision", + "scale", + b"scale", + "type_variation_reference", + b"type_variation_reference", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_precision", b"_precision"] + ) -> typing_extensions.Literal["precision"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_scale", b"_scale"] + ) -> typing_extensions.Literal["scale"] | None: ... + + class StructField(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + NAME_FIELD_NUMBER: builtins.int + DATA_TYPE_FIELD_NUMBER: builtins.int + NULLABLE_FIELD_NUMBER: builtins.int + METADATA_FIELD_NUMBER: builtins.int + name: builtins.str + @property + def data_type(self) -> global___DataType: ... + nullable: builtins.bool + metadata: builtins.str + def __init__( + self, + *, + name: builtins.str = ..., + data_type: global___DataType | None = ..., + nullable: builtins.bool = ..., + metadata: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_metadata", b"_metadata", "data_type", b"data_type", "metadata", b"metadata" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_metadata", + b"_metadata", + "data_type", + b"data_type", + "metadata", + b"metadata", + "name", + b"name", + "nullable", + b"nullable", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_metadata", b"_metadata"] + ) -> typing_extensions.Literal["metadata"] | None: ... + + class Struct(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FIELDS_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + @property + def fields( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___DataType.StructField + ]: ... + type_variation_reference: builtins.int + def __init__( + self, + *, + fields: collections.abc.Iterable[global___DataType.StructField] | None = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "fields", b"fields", "type_variation_reference", b"type_variation_reference" + ], + ) -> None: ... + + class Array(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ELEMENT_TYPE_FIELD_NUMBER: builtins.int + CONTAINS_NULL_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + @property + def element_type(self) -> global___DataType: ... + contains_null: builtins.bool + type_variation_reference: builtins.int + def __init__( + self, + *, + element_type: global___DataType | None = ..., + contains_null: builtins.bool = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["element_type", b"element_type"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "contains_null", + b"contains_null", + "element_type", + b"element_type", + "type_variation_reference", + b"type_variation_reference", + ], + ) -> None: ... + + class Map(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_TYPE_FIELD_NUMBER: builtins.int + VALUE_TYPE_FIELD_NUMBER: builtins.int + VALUE_CONTAINS_NULL_FIELD_NUMBER: builtins.int + TYPE_VARIATION_REFERENCE_FIELD_NUMBER: builtins.int + @property + def key_type(self) -> global___DataType: ... + @property + def value_type(self) -> global___DataType: ... + value_contains_null: builtins.bool + type_variation_reference: builtins.int + def __init__( + self, + *, + key_type: global___DataType | None = ..., + value_type: global___DataType | None = ..., + value_contains_null: builtins.bool = ..., + type_variation_reference: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "key_type", b"key_type", "value_type", b"value_type" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "key_type", + b"key_type", + "type_variation_reference", + b"type_variation_reference", + "value_contains_null", + b"value_contains_null", + "value_type", + b"value_type", + ], + ) -> None: ... + + class UDT(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + TYPE_FIELD_NUMBER: builtins.int + JVM_CLASS_FIELD_NUMBER: builtins.int + PYTHON_CLASS_FIELD_NUMBER: builtins.int + SERIALIZED_PYTHON_CLASS_FIELD_NUMBER: builtins.int + SQL_TYPE_FIELD_NUMBER: builtins.int + type: builtins.str + jvm_class: builtins.str + python_class: builtins.str + serialized_python_class: builtins.str + @property + def sql_type(self) -> global___DataType: ... + def __init__( + self, + *, + type: builtins.str = ..., + jvm_class: builtins.str | None = ..., + python_class: builtins.str | None = ..., + serialized_python_class: builtins.str | None = ..., + sql_type: global___DataType | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_jvm_class", + b"_jvm_class", + "_python_class", + b"_python_class", + "_serialized_python_class", + b"_serialized_python_class", + "jvm_class", + b"jvm_class", + "python_class", + b"python_class", + "serialized_python_class", + b"serialized_python_class", + "sql_type", + b"sql_type", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_jvm_class", + b"_jvm_class", + "_python_class", + b"_python_class", + "_serialized_python_class", + b"_serialized_python_class", + "jvm_class", + b"jvm_class", + "python_class", + b"python_class", + "serialized_python_class", + b"serialized_python_class", + "sql_type", + b"sql_type", + "type", + b"type", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_jvm_class", b"_jvm_class"] + ) -> typing_extensions.Literal["jvm_class"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_python_class", b"_python_class"] + ) -> typing_extensions.Literal["python_class"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_serialized_python_class", b"_serialized_python_class" + ], + ) -> typing_extensions.Literal["serialized_python_class"] | None: ... + + class Unparsed(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DATA_TYPE_STRING_FIELD_NUMBER: builtins.int + data_type_string: builtins.str + """(Required) The unparsed data type string""" + def __init__( + self, + *, + data_type_string: builtins.str = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["data_type_string", b"data_type_string"] + ) -> None: ... + + NULL_FIELD_NUMBER: builtins.int + BINARY_FIELD_NUMBER: builtins.int + BOOLEAN_FIELD_NUMBER: builtins.int + BYTE_FIELD_NUMBER: builtins.int + SHORT_FIELD_NUMBER: builtins.int + INTEGER_FIELD_NUMBER: builtins.int + LONG_FIELD_NUMBER: builtins.int + FLOAT_FIELD_NUMBER: builtins.int + DOUBLE_FIELD_NUMBER: builtins.int + DECIMAL_FIELD_NUMBER: builtins.int + STRING_FIELD_NUMBER: builtins.int + CHAR_FIELD_NUMBER: builtins.int + VAR_CHAR_FIELD_NUMBER: builtins.int + DATE_FIELD_NUMBER: builtins.int + TIMESTAMP_FIELD_NUMBER: builtins.int + TIMESTAMP_NTZ_FIELD_NUMBER: builtins.int + CALENDAR_INTERVAL_FIELD_NUMBER: builtins.int + YEAR_MONTH_INTERVAL_FIELD_NUMBER: builtins.int + DAY_TIME_INTERVAL_FIELD_NUMBER: builtins.int + ARRAY_FIELD_NUMBER: builtins.int + STRUCT_FIELD_NUMBER: builtins.int + MAP_FIELD_NUMBER: builtins.int + UDT_FIELD_NUMBER: builtins.int + UNPARSED_FIELD_NUMBER: builtins.int + @property + def null(self) -> global___DataType.NULL: ... + @property + def binary(self) -> global___DataType.Binary: ... + @property + def boolean(self) -> global___DataType.Boolean: ... + @property + def byte(self) -> global___DataType.Byte: + """Numeric types""" + @property + def short(self) -> global___DataType.Short: ... + @property + def integer(self) -> global___DataType.Integer: ... + @property + def long(self) -> global___DataType.Long: ... + @property + def float(self) -> global___DataType.Float: ... + @property + def double(self) -> global___DataType.Double: ... + @property + def decimal(self) -> global___DataType.Decimal: ... + @property + def string(self) -> global___DataType.String: + """String types""" + @property + def char(self) -> global___DataType.Char: ... + @property + def var_char(self) -> global___DataType.VarChar: ... + @property + def date(self) -> global___DataType.Date: + """Datatime types""" + @property + def timestamp(self) -> global___DataType.Timestamp: ... + @property + def timestamp_ntz(self) -> global___DataType.TimestampNTZ: ... + @property + def calendar_interval(self) -> global___DataType.CalendarInterval: + """Interval types""" + @property + def year_month_interval(self) -> global___DataType.YearMonthInterval: ... + @property + def day_time_interval(self) -> global___DataType.DayTimeInterval: ... + @property + def array(self) -> global___DataType.Array: + """Complex types""" + @property + def struct(self) -> global___DataType.Struct: ... + @property + def map(self) -> global___DataType.Map: ... + @property + def udt(self) -> global___DataType.UDT: + """UserDefinedType""" + @property + def unparsed(self) -> global___DataType.Unparsed: + """UnparsedDataType""" + def __init__( + self, + *, + null: global___DataType.NULL | None = ..., + binary: global___DataType.Binary | None = ..., + boolean: global___DataType.Boolean | None = ..., + byte: global___DataType.Byte | None = ..., + short: global___DataType.Short | None = ..., + integer: global___DataType.Integer | None = ..., + long: global___DataType.Long | None = ..., + float: global___DataType.Float | None = ..., + double: global___DataType.Double | None = ..., + decimal: global___DataType.Decimal | None = ..., + string: global___DataType.String | None = ..., + char: global___DataType.Char | None = ..., + var_char: global___DataType.VarChar | None = ..., + date: global___DataType.Date | None = ..., + timestamp: global___DataType.Timestamp | None = ..., + timestamp_ntz: global___DataType.TimestampNTZ | None = ..., + calendar_interval: global___DataType.CalendarInterval | None = ..., + year_month_interval: global___DataType.YearMonthInterval | None = ..., + day_time_interval: global___DataType.DayTimeInterval | None = ..., + array: global___DataType.Array | None = ..., + struct: global___DataType.Struct | None = ..., + map: global___DataType.Map | None = ..., + udt: global___DataType.UDT | None = ..., + unparsed: global___DataType.Unparsed | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "array", + b"array", + "binary", + b"binary", + "boolean", + b"boolean", + "byte", + b"byte", + "calendar_interval", + b"calendar_interval", + "char", + b"char", + "date", + b"date", + "day_time_interval", + b"day_time_interval", + "decimal", + b"decimal", + "double", + b"double", + "float", + b"float", + "integer", + b"integer", + "kind", + b"kind", + "long", + b"long", + "map", + b"map", + "null", + b"null", + "short", + b"short", + "string", + b"string", + "struct", + b"struct", + "timestamp", + b"timestamp", + "timestamp_ntz", + b"timestamp_ntz", + "udt", + b"udt", + "unparsed", + b"unparsed", + "var_char", + b"var_char", + "year_month_interval", + b"year_month_interval", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "array", + b"array", + "binary", + b"binary", + "boolean", + b"boolean", + "byte", + b"byte", + "calendar_interval", + b"calendar_interval", + "char", + b"char", + "date", + b"date", + "day_time_interval", + b"day_time_interval", + "decimal", + b"decimal", + "double", + b"double", + "float", + b"float", + "integer", + b"integer", + "kind", + b"kind", + "long", + b"long", + "map", + b"map", + "null", + b"null", + "short", + b"short", + "string", + b"string", + "struct", + b"struct", + "timestamp", + b"timestamp", + "timestamp_ntz", + b"timestamp_ntz", + "udt", + b"udt", + "unparsed", + b"unparsed", + "var_char", + b"var_char", + "year_month_interval", + b"year_month_interval", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["kind", b"kind"] + ) -> typing_extensions.Literal[ + "null", + "binary", + "boolean", + "byte", + "short", + "integer", + "long", + "float", + "double", + "decimal", + "string", + "char", + "var_char", + "date", + "timestamp", + "timestamp_ntz", + "calendar_interval", + "year_month_interval", + "day_time_interval", + "array", + "struct", + "map", + "udt", + "unparsed", + ] | None: ... + +global___DataType = DataType diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/__init__.py new file mode 100644 index 0000000..dc81e9f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +"""Spark Connect Python Client - Protobuf Functions""" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/functions.py new file mode 100644 index 0000000..56119f4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/protobuf/functions.py @@ -0,0 +1,166 @@ +# +# 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. +# + +""" +A collections of builtin protobuf functions +""" + +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import Dict, Optional, TYPE_CHECKING + +from pyspark.sql.protobuf import functions as PyProtobufFunctions + +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.functions import _invoke_function, _to_col, _options_to_col, lit + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ColumnOrName + + +def from_protobuf( + data: "ColumnOrName", + messageName: str, + descFilePath: Optional[str] = None, + options: Optional[Dict[str, str]] = None, + binaryDescriptorSet: Optional[bytes] = None, +) -> Column: + binary_proto = None + if binaryDescriptorSet is not None: + binary_proto = binaryDescriptorSet + elif descFilePath is not None: + binary_proto = _read_descriptor_set_file(descFilePath) + + # TODO: simplify the code when _invoke_function() supports None as input. + if binary_proto is not None: + if options is None: + return _invoke_function( + "from_protobuf", _to_col(data), lit(messageName), lit(binary_proto) + ) + else: + return _invoke_function( + "from_protobuf", + _to_col(data), + lit(messageName), + lit(binary_proto), + _options_to_col(options), + ) + else: + if options is None: + return _invoke_function("from_protobuf", _to_col(data), lit(messageName)) + else: + return _invoke_function( + "from_protobuf", _to_col(data), lit(messageName), _options_to_col(options) + ) + + +from_protobuf.__doc__ = PyProtobufFunctions.from_protobuf.__doc__ + + +def to_protobuf( + data: "ColumnOrName", + messageName: str, + descFilePath: Optional[str] = None, + options: Optional[Dict[str, str]] = None, + binaryDescriptorSet: Optional[bytes] = None, +) -> Column: + binary_proto = None + if binaryDescriptorSet is not None: + binary_proto = binaryDescriptorSet + elif descFilePath is not None: + binary_proto = _read_descriptor_set_file(descFilePath) + + # TODO: simplify the code when _invoke_function() supports None as input. + if binary_proto is not None: + if options is None: + return _invoke_function( + "to_protobuf", _to_col(data), lit(messageName), lit(binary_proto) + ) + else: + return _invoke_function( + "to_protobuf", + _to_col(data), + lit(messageName), + lit(binary_proto), + _options_to_col(options), + ) + else: + if options is None: + return _invoke_function("to_protobuf", _to_col(data), lit(messageName)) + else: + return _invoke_function( + "to_protobuf", _to_col(data), lit(messageName), _options_to_col(options) + ) + + +to_protobuf.__doc__ = PyProtobufFunctions.to_protobuf.__doc__ + + +def _read_descriptor_set_file(filePath: str) -> bytes: + with open(filePath, "rb") as f: + return f.read() + + +def _test() -> None: + import os + import sys + from pyspark.testing.utils import search_jar + + protobuf_jar = search_jar("connector/protobuf", "spark-protobuf-assembly-", "spark-protobuf") + if protobuf_jar is None: + print( + "Skipping all Protobuf Python tests as the optional Protobuf project was " + "not compiled into a JAR. To run these tests, " + "you need to build Spark with 'build/sbt package' or " + "'build/mvn package' before running this test." + ) + sys.exit(0) + else: + existing_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + jars_args = "--jars %s" % protobuf_jar + os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) + + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.protobuf.functions + + globs = pyspark.sql.connect.protobuf.functions.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.protobuf.functions tests") + .remote("local[2]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.protobuf.functions, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/readwriter.py new file mode 100644 index 0000000..cfcbced --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/readwriter.py @@ -0,0 +1,860 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +from typing import Dict +from typing import Optional, Union, List, overload, Tuple, cast +from typing import TYPE_CHECKING + +from pyspark.sql.connect.plan import Read, DataSource, LogicalPlan, WriteOperation, WriteOperationV2 +from pyspark.sql.types import StructType +from pyspark.sql.utils import to_str +from pyspark.sql.readwriter import ( + DataFrameWriter as PySparkDataFrameWriter, + DataFrameReader as PySparkDataFrameReader, + DataFrameWriterV2 as PySparkDataFrameWriterV2, +) +from pyspark.errors import PySparkAttributeError, PySparkTypeError, PySparkValueError + +if TYPE_CHECKING: + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType + from pyspark.sql.connect.session import SparkSession + +__all__ = ["DataFrameReader", "DataFrameWriter"] + +PathOrPaths = Union[str, List[str]] +TupleOrListOfString = Union[List[str], Tuple[str, ...]] + + +class OptionUtils: + def _set_opts( + self, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """ + Set named options (filter out those the value is None) + """ + if schema is not None: + self.schema(schema) # type: ignore[attr-defined] + for k, v in options.items(): + if v is not None: + self.option(k, v) # type: ignore[attr-defined] + + +class DataFrameReader(OptionUtils): + # TODO(SPARK-40539) Achieve parity with PySpark. + + def __init__(self, client: "SparkSession"): + self._client = client + self._format: Optional[str] = None + self._schema = "" + self._options: Dict[str, str] = {} + + def format(self, source: str) -> "DataFrameReader": + self._format = source + return self + + format.__doc__ = PySparkDataFrameReader.format.__doc__ + + def schema(self, schema: Union[StructType, str]) -> "DataFrameReader": + if isinstance(schema, StructType): + self._schema = schema.json() + elif isinstance(schema, str): + self._schema = schema + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_STRUCT", + message_parameters={ + "arg_name": "schema", + "arg_type": type(schema).__name__, + }, + ) + return self + + schema.__doc__ = PySparkDataFrameReader.schema.__doc__ + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameReader": + self._options[key] = str(value) + return self + + option.__doc__ = PySparkDataFrameReader.option.__doc__ + + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameReader": + for k in options: + self.option(k, to_str(options[k])) + return self + + options.__doc__ = PySparkDataFrameReader.options.__doc__ + + def load( + self, + path: Optional[PathOrPaths] = None, + format: Optional[str] = None, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> "DataFrame": + if format is not None: + self.format(format) + if schema is not None: + self.schema(schema) + self.options(**options) + + paths = path + if isinstance(path, str): + paths = [path] + + plan = DataSource( + format=self._format, + schema=self._schema, + options=self._options, + paths=paths, # type: ignore[arg-type] + ) + return self._df(plan) + + load.__doc__ = PySparkDataFrameReader.load.__doc__ + + def _df(self, plan: LogicalPlan) -> "DataFrame": + from pyspark.sql.connect.dataframe import DataFrame + + return DataFrame.withPlan(plan, self._client) + + def table(self, tableName: str) -> "DataFrame": + return self._df(Read(tableName, self._options)) + + table.__doc__ = PySparkDataFrameReader.table.__doc__ + + def json( + self, + path: PathOrPaths, + schema: Optional[Union[StructType, str]] = None, + primitivesAsString: Optional[Union[bool, str]] = None, + prefersDecimal: Optional[Union[bool, str]] = None, + allowComments: Optional[Union[bool, str]] = None, + allowUnquotedFieldNames: Optional[Union[bool, str]] = None, + allowSingleQuotes: Optional[Union[bool, str]] = None, + allowNumericLeadingZero: Optional[Union[bool, str]] = None, + allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + allowUnquotedControlChars: Optional[Union[bool, str]] = None, + lineSep: Optional[str] = None, + samplingRatio: Optional[Union[float, str]] = None, + dropFieldIfAllNull: Optional[Union[bool, str]] = None, + encoding: Optional[str] = None, + locale: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + allowNonNumericNumbers: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + primitivesAsString=primitivesAsString, + prefersDecimal=prefersDecimal, + allowComments=allowComments, + allowUnquotedFieldNames=allowUnquotedFieldNames, + allowSingleQuotes=allowSingleQuotes, + allowNumericLeadingZero=allowNumericLeadingZero, + allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + multiLine=multiLine, + allowUnquotedControlChars=allowUnquotedControlChars, + lineSep=lineSep, + samplingRatio=samplingRatio, + dropFieldIfAllNull=dropFieldIfAllNull, + encoding=encoding, + locale=locale, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + allowNonNumericNumbers=allowNonNumericNumbers, + ) + if isinstance(path, str): + path = [path] + return self.load(path=path, format="json", schema=schema) + + json.__doc__ = PySparkDataFrameReader.json.__doc__ + + def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame": + mergeSchema = options.get("mergeSchema", None) + pathGlobFilter = options.get("pathGlobFilter", None) + modifiedBefore = options.get("modifiedBefore", None) + modifiedAfter = options.get("modifiedAfter", None) + recursiveFileLookup = options.get("recursiveFileLookup", None) + datetimeRebaseMode = options.get("datetimeRebaseMode", None) + int96RebaseMode = options.get("int96RebaseMode", None) + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + datetimeRebaseMode=datetimeRebaseMode, + int96RebaseMode=int96RebaseMode, + ) + + return self.load(path=list(paths), format="parquet") + + parquet.__doc__ = PySparkDataFrameReader.parquet.__doc__ + + def text( + self, + paths: PathOrPaths, + wholetext: Optional[bool] = None, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + wholetext=wholetext, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + ) + + if isinstance(paths, str): + paths = [paths] + return self.load(path=paths, format="text") + + text.__doc__ = PySparkDataFrameReader.text.__doc__ + + def csv( + self, + path: PathOrPaths, + schema: Optional[Union[StructType, str]] = None, + sep: Optional[str] = None, + encoding: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + comment: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + inferSchema: Optional[Union[bool, str]] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + nanValue: Optional[str] = None, + positiveInf: Optional[str] = None, + negativeInf: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + maxColumns: Optional[Union[int, str]] = None, + maxCharsPerColumn: Optional[Union[int, str]] = None, + maxMalformedLogPerPartition: Optional[Union[int, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[str] = None, + samplingRatio: Optional[Union[float, str]] = None, + enforceSchema: Optional[Union[bool, str]] = None, + emptyValue: Optional[str] = None, + locale: Optional[str] = None, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + unescapedQuoteHandling: Optional[str] = None, + ) -> "DataFrame": + self._set_opts( + sep=sep, + encoding=encoding, + quote=quote, + escape=escape, + comment=comment, + header=header, + inferSchema=inferSchema, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + nullValue=nullValue, + nanValue=nanValue, + positiveInf=positiveInf, + negativeInf=negativeInf, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + maxColumns=maxColumns, + maxCharsPerColumn=maxCharsPerColumn, + maxMalformedLogPerPartition=maxMalformedLogPerPartition, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + multiLine=multiLine, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + samplingRatio=samplingRatio, + enforceSchema=enforceSchema, + emptyValue=emptyValue, + locale=locale, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + unescapedQuoteHandling=unescapedQuoteHandling, + ) + if isinstance(path, str): + path = [path] + return self.load(path=path, format="csv", schema=schema) + + csv.__doc__ = PySparkDataFrameReader.csv.__doc__ + + def orc( + self, + path: PathOrPaths, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + path = [path] + return self.load(path=path, format="orc") + + orc.__doc__ = PySparkDataFrameReader.orc.__doc__ + + @overload + def jdbc( + self, url: str, table: str, *, properties: Optional[Dict[str, str]] = None + ) -> "DataFrame": + ... + + @overload + def jdbc( + self, + url: str, + table: str, + column: str, + lowerBound: Union[int, str], + upperBound: Union[int, str], + numPartitions: int, + *, + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + ... + + @overload + def jdbc( + self, + url: str, + table: str, + *, + predicates: List[str], + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + ... + + def jdbc( + self, + url: str, + table: str, + column: Optional[str] = None, + lowerBound: Optional[Union[int, str]] = None, + upperBound: Optional[Union[int, str]] = None, + numPartitions: Optional[int] = None, + predicates: Optional[List[str]] = None, + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + if properties is None: + properties = dict() + + self.format("jdbc") + + if column is not None: + assert lowerBound is not None, "lowerBound can not be None when ``column`` is specified" + assert upperBound is not None, "upperBound can not be None when ``column`` is specified" + assert ( + numPartitions is not None + ), "numPartitions can not be None when ``column`` is specified" + self.options( + partitionColumn=column, + lowerBound=lowerBound, + upperBound=upperBound, + numPartitions=numPartitions, + ) + self.options(**properties) + self.options(url=url, dbtable=table) + return self.load() + else: + self.options(**properties) + self.options(url=url, dbtable=table) + if predicates is not None: + plan = DataSource( + format=self._format, + schema=self._schema, + options=self._options, + predicates=predicates, + ) + return self._df(plan) + else: + return self.load() + + jdbc.__doc__ = PySparkDataFrameReader.jdbc.__doc__ + + @property + def _jreader(self) -> None: + raise PySparkAttributeError( + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", message_parameters={"attr_name": "_jreader"} + ) + + +DataFrameReader.__doc__ = PySparkDataFrameReader.__doc__ + + +class DataFrameWriter(OptionUtils): + def __init__(self, plan: "LogicalPlan", session: "SparkSession"): + self._df: "LogicalPlan" = plan + self._spark: "SparkSession" = session + self._write: "WriteOperation" = WriteOperation(self._df) + + def mode(self, saveMode: Optional[str]) -> "DataFrameWriter": + # At the JVM side, the default value of mode is already set to "error". + # So, if the given saveMode is None, we will not call JVM-side's mode method. + if saveMode is not None: + self._write.mode = saveMode + return self + + mode.__doc__ = PySparkDataFrameWriter.mode.__doc__ + + def format(self, source: str) -> "DataFrameWriter": + self._write.source = source + return self + + format.__doc__ = PySparkDataFrameWriter.format.__doc__ + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriter": + self._write.options[key] = to_str(value) + return self + + option.__doc__ = PySparkDataFrameWriter.option.__doc__ + + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter": + for k in options: + self._write.options[k] = to_str(options[k]) + return self + + options.__doc__ = PySparkDataFrameWriter.options.__doc__ + + @overload + def partitionBy(self, *cols: str) -> "DataFrameWriter": + ... + + @overload + def partitionBy(self, *cols: List[str]) -> "DataFrameWriter": + ... + + def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": + if len(cols) == 1 and isinstance(cols[0], (list, tuple)): + cols = cols[0] # type: ignore[assignment] + + self._write.partitioning_cols = cast(List[str], cols) + return self + + partitionBy.__doc__ = PySparkDataFrameWriter.partitionBy.__doc__ + + @overload + def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter": + ... + + @overload + def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter": + ... + + def bucketBy( + self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str] + ) -> "DataFrameWriter": + if not isinstance(numBuckets, int): + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={ + "arg_list": f"`col` of type {type(col).__name__} and `cols`", + }, + ) + + if isinstance(col, (list, tuple)): + if cols: + raise PySparkValueError( + error_class="NOT_INT", + message_parameters={ + "arg_list": "numBuckets", + }, + ) + + col, cols = col[0], col[1:] # type: ignore[assignment] + + for c in cols: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "cols", + "arg_type": type(c).__name__, + }, + ) + if not isinstance(col, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "col", + "arg_type": type(col).__name__, + }, + ) + + self._write.num_buckets = numBuckets + self._write.bucket_cols = cast(List[str], [col, *cols]) + return self + + bucketBy.__doc__ = PySparkDataFrameWriter.bucketBy.__doc__ + + @overload + def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": + ... + + @overload + def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter": + ... + + def sortBy( + self, col: Union[str, TupleOrListOfString], *cols: Optional[str] + ) -> "DataFrameWriter": + if isinstance(col, (list, tuple)): + if cols: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={ + "arg_list": f"`col` of type {type(col).__name__} and `cols`", + }, + ) + + col, cols = col[0], col[1:] # type: ignore[assignment] + + for c in cols: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "cols", + "arg_type": type(c).__name__, + }, + ) + if not isinstance(col, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "col", + "arg_type": type(col).__name__, + }, + ) + + self._write.sort_cols = cast(List[str], [col, *cols]) + return self + + sortBy.__doc__ = PySparkDataFrameWriter.sortBy.__doc__ + + def save( + self, + path: Optional[str] = None, + format: Optional[str] = None, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + self.mode(mode).options(**options) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + self._write.path = path + self._spark.client.execute_command(self._write.command(self._spark.client)) + + save.__doc__ = PySparkDataFrameWriter.save.__doc__ + + def insertInto(self, tableName: str, overwrite: Optional[bool] = None) -> None: + if overwrite is not None: + self.mode("overwrite" if overwrite else "append") + self._write.table_name = tableName + self._write.table_save_method = "insert_into" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + insertInto.__doc__ = PySparkDataFrameWriter.insertInto.__doc__ + + def saveAsTable( + self, + name: str, + format: Optional[str] = None, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + self.mode(mode).options(**options) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + self._write.table_name = name + self._write.table_save_method = "save_as_table" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + saveAsTable.__doc__ = PySparkDataFrameWriter.saveAsTable.__doc__ + + def json( + self, + path: str, + mode: Optional[str] = None, + compression: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + lineSep: Optional[str] = None, + encoding: Optional[str] = None, + ignoreNullFields: Optional[Union[bool, str]] = None, + ) -> None: + self.mode(mode) + self._set_opts( + compression=compression, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + lineSep=lineSep, + encoding=encoding, + ignoreNullFields=ignoreNullFields, + ) + self.format("json").save(path) + + json.__doc__ = PySparkDataFrameWriter.json.__doc__ + + def parquet( + self, + path: str, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + compression: Optional[str] = None, + ) -> None: + self.mode(mode) + if partitionBy is not None: + self.partitionBy(partitionBy) + self.option("compression", compression) + self.format("parquet").save(path) + + parquet.__doc__ = PySparkDataFrameWriter.parquet.__doc__ + + def text( + self, path: str, compression: Optional[str] = None, lineSep: Optional[str] = None + ) -> None: + self._set_opts(compression=compression, lineSep=lineSep) + self.format("text").save(path) + + text.__doc__ = PySparkDataFrameWriter.text.__doc__ + + def csv( + self, + path: str, + mode: Optional[str] = None, + compression: Optional[str] = None, + sep: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + escapeQuotes: Optional[Union[bool, str]] = None, + quoteAll: Optional[Union[bool, str]] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[str] = None, + encoding: Optional[str] = None, + emptyValue: Optional[str] = None, + lineSep: Optional[str] = None, + ) -> None: + self.mode(mode) + self._set_opts( + compression=compression, + sep=sep, + quote=quote, + escape=escape, + header=header, + nullValue=nullValue, + escapeQuotes=escapeQuotes, + quoteAll=quoteAll, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + encoding=encoding, + emptyValue=emptyValue, + lineSep=lineSep, + ) + self.format("csv").save(path) + + csv.__doc__ = PySparkDataFrameWriter.csv.__doc__ + + def orc( + self, + path: str, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + compression: Optional[str] = None, + ) -> None: + self.mode(mode) + if partitionBy is not None: + self.partitionBy(partitionBy) + self._set_opts(compression=compression) + self.format("orc").save(path) + + orc.__doc__ = PySparkDataFrameWriter.orc.__doc__ + + def jdbc( + self, + url: str, + table: str, + mode: Optional[str] = None, + properties: Optional[Dict[str, str]] = None, + ) -> None: + if properties is None: + properties = dict() + + self.format("jdbc").mode(mode).options(**properties).options(url=url, dbtable=table).save() + + jdbc.__doc__ = PySparkDataFrameWriter.jdbc.__doc__ + + +class DataFrameWriterV2(OptionUtils): + def __init__(self, plan: "LogicalPlan", session: "SparkSession", table: str): + self._df: "LogicalPlan" = plan + self._spark: "SparkSession" = session + self._table_name: str = table + self._write: "WriteOperationV2" = WriteOperationV2(self._df, self._table_name) + + def using(self, provider: str) -> "DataFrameWriterV2": + self._write.provider = provider + return self + + using.__doc__ = PySparkDataFrameWriterV2.using.__doc__ + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriterV2": + self._write.options[key] = to_str(value) + return self + + option.__doc__ = PySparkDataFrameWriterV2.option.__doc__ + + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriterV2": + for k in options: + self._write.options[k] = to_str(options[k]) + return self + + options.__doc__ = PySparkDataFrameWriterV2.options.__doc__ + + def tableProperty(self, property: str, value: str) -> "DataFrameWriterV2": + self._write.table_properties[property] = value + return self + + tableProperty.__doc__ = PySparkDataFrameWriterV2.tableProperty.__doc__ + + def partitionedBy(self, col: "ColumnOrName", *cols: "ColumnOrName") -> "DataFrameWriterV2": + self._write.partitioning_columns = [col] + self._write.partitioning_columns.extend(cols) + return self + + partitionedBy.__doc__ = PySparkDataFrameWriterV2.partitionedBy.__doc__ + + def create(self) -> None: + self._write.mode = "create" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + create.__doc__ = PySparkDataFrameWriterV2.create.__doc__ + + def replace(self) -> None: + self._write.mode = "replace" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + replace.__doc__ = PySparkDataFrameWriterV2.replace.__doc__ + + def createOrReplace(self) -> None: + self._write.mode = "create_or_replace" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + createOrReplace.__doc__ = PySparkDataFrameWriterV2.createOrReplace.__doc__ + + def append(self) -> None: + self._write.mode = "append" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + append.__doc__ = PySparkDataFrameWriterV2.append.__doc__ + + def overwrite(self, condition: "ColumnOrName") -> None: + self._write.mode = "overwrite" + self._write.overwrite_condition = condition + self._spark.client.execute_command(self._write.command(self._spark.client)) + + overwrite.__doc__ = PySparkDataFrameWriterV2.overwrite.__doc__ + + def overwritePartitions(self) -> None: + self._write.mode = "overwrite_partitions" + self._spark.client.execute_command(self._write.command(self._spark.client)) + + overwritePartitions.__doc__ = PySparkDataFrameWriterV2.overwritePartitions.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.readwriter + + globs = pyspark.sql.connect.readwriter.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.readwriter tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.readwriter, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/session.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/session.py new file mode 100644 index 0000000..1307c8b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/session.py @@ -0,0 +1,949 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import threading +import os +import warnings +from collections.abc import Sized +from distutils.version import LooseVersion +from functools import reduce +from threading import RLock +from typing import ( + Optional, + Any, + Union, + Dict, + List, + Tuple, + Set, + cast, + overload, + Iterable, + TYPE_CHECKING, + ClassVar, +) + +import numpy as np +import pandas as pd +import pyarrow as pa +from pandas.api.types import ( # type: ignore[attr-defined] + is_datetime64_dtype, + is_datetime64tz_dtype, + is_timedelta64_dtype, +) +import urllib + +from pyspark import SparkContext, SparkConf, __version__ +from pyspark.sql.connect.client import SparkConnectClient, ChannelBuilder +from pyspark.sql.connect.conf import RuntimeConf +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.connect.plan import ( + SQL, + Range, + LocalRelation, + LogicalPlan, + CachedLocalRelation, + CachedRelation, + CachedRemoteRelation, +) +from pyspark.sql.connect.readwriter import DataFrameReader +from pyspark.sql.connect.streaming import DataStreamReader, StreamingQueryManager +from pyspark.sql.pandas.serializers import ArrowStreamPandasSerializer +from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type, _deduplicate_field_names +from pyspark.sql.session import classproperty, SparkSession as PySparkSession +from pyspark.sql.types import ( + _infer_schema, + _has_nulltype, + _merge_type, + Row, + DataType, + DayTimeIntervalType, + StructType, + AtomicType, + TimestampType, +) +from pyspark.sql.utils import to_str +from pyspark.errors import ( + PySparkAttributeError, + PySparkNotImplementedError, + PySparkRuntimeError, + PySparkValueError, + PySparkTypeError, +) + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import OptionalPrimitiveType + from pyspark.sql.connect.catalog import Catalog + from pyspark.sql.connect.udf import UDFRegistration + from pyspark.sql.connect.udtf import UDTFRegistration + + +class SparkSession: + # The active SparkSession for the current thread + _active_session: ClassVar[threading.local] = threading.local() + # Reference to the root SparkSession + _default_session: ClassVar[Optional["SparkSession"]] = None + _lock: ClassVar[RLock] = RLock() + + class Builder: + """Builder for :class:`SparkSession`.""" + + _lock = RLock() + + def __init__(self) -> None: + self._options: Dict[str, Any] = {} + self._channel_builder: Optional[ChannelBuilder] = None + + @overload + def config(self, key: str, value: Any) -> "SparkSession.Builder": + ... + + @overload + def config(self, *, map: Dict[str, "OptionalPrimitiveType"]) -> "SparkSession.Builder": + ... + + def config( + self, + key: Optional[str] = None, + value: Optional[Any] = None, + *, + map: Optional[Dict[str, "OptionalPrimitiveType"]] = None, + ) -> "SparkSession.Builder": + with self._lock: + if map is not None: + for k, v in map.items(): + self._options[k] = to_str(v) + else: + self._options[cast(str, key)] = to_str(value) + return self + + def master(self, master: str) -> "SparkSession.Builder": + return self + + def appName(self, name: str) -> "SparkSession.Builder": + return self.config("spark.app.name", name) + + def remote(self, location: str = "sc://localhost") -> "SparkSession.Builder": + return self.config("spark.remote", location) + + def channelBuilder(self, channelBuilder: ChannelBuilder) -> "SparkSession.Builder": + """Uses custom :class:`ChannelBuilder` implementation, when there is a need + to customize the behavior for creation of GRPC connections. + + .. versionadded:: 3.5.0 + + An example to use this class looks like this: + + .. code-block:: python + + from pyspark.sql.connect import SparkSession, ChannelBuilder + + class CustomChannelBuilder(ChannelBuilder): + ... + + custom_channel_builder = CustomChannelBuilder(...) + spark = SparkSession.builder().channelBuilder(custom_channel_builder).getOrCreate() + + Returns + ------- + :class:`SparkSession.Builder` + """ + with self._lock: + # self._channel_builder is a separate field, because it may hold the state + # and cannot be serialized with to_str() + self._channel_builder = channelBuilder + return self + + def enableHiveSupport(self) -> "SparkSession.Builder": + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", message_parameters={"feature": "enableHiveSupport"} + ) + + def _apply_options(self, session: "SparkSession") -> None: + with self._lock: + for k, v in self._options.items(): + try: + session.conf.set(k, v) + except Exception as e: + warnings.warn(str(e)) + + def create(self) -> "SparkSession": + has_channel_builder = self._channel_builder is not None + has_spark_remote = "spark.remote" in self._options + + if has_channel_builder and has_spark_remote: + raise ValueError( + "Only one of connection string or channelBuilder " + "can be used to create a new SparkSession." + ) + + if not has_channel_builder and not has_spark_remote: + raise ValueError( + "Needs either connection string or channelBuilder to create a new SparkSession." + ) + + if has_channel_builder: + assert self._channel_builder is not None + session = SparkSession(connection=self._channel_builder) + else: + spark_remote = to_str(self._options.get("spark.remote")) + assert spark_remote is not None + session = SparkSession(connection=spark_remote) + + SparkSession._set_default_and_active_session(session) + self._apply_options(session) + return session + + def getOrCreate(self) -> "SparkSession": + with SparkSession._lock: + session = SparkSession.getActiveSession() + if session is None: + session = SparkSession._default_session + if session is None: + session = self.create() + self._apply_options(session) + return session + + _client: SparkConnectClient + + @classproperty + def builder(cls) -> Builder: + return cls.Builder() + + builder.__doc__ = PySparkSession.builder.__doc__ + + def __init__(self, connection: Union[str, ChannelBuilder], userId: Optional[str] = None): + """ + Creates a new SparkSession for the Spark Connect interface. + + Parameters + ---------- + connection: Union[str,ChannelBuilder] + Connection string that is used to extract the connection parameters and configure + the GRPC connection. Or instance of ChannelBuilder that creates GRPC connection. + Defaults to `sc://localhost`. + userId : str, optional + Optional unique user ID that is used to differentiate multiple users and + isolate their Spark Sessions. If the `user_id` is not set, will default to + the $USER environment. Defining the user ID as part of the connection string + takes precedence. + """ + self._client = SparkConnectClient(connection=connection, user_id=userId) + self._session_id = self._client._session_id + + @classmethod + def _set_default_and_active_session(cls, session: "SparkSession") -> None: + """ + Set the (global) default :class:`SparkSession`, and (thread-local) + active :class:`SparkSession` when they are not set yet. + """ + with cls._lock: + if cls._default_session is None: + cls._default_session = session + if getattr(cls._active_session, "session", None) is None: + cls._active_session.session = session + + @classmethod + def getActiveSession(cls) -> Optional["SparkSession"]: + return getattr(cls._active_session, "session", None) + + getActiveSession.__doc__ = PySparkSession.getActiveSession.__doc__ + + @classmethod + def active(cls) -> "SparkSession": + session = cls.getActiveSession() + if session is None: + session = cls._default_session + if session is None: + raise PySparkRuntimeError( + error_class="NO_ACTIVE_OR_DEFAULT_SESSION", + message_parameters={}, + ) + return session + + active.__doc__ = PySparkSession.active.__doc__ + + def table(self, tableName: str) -> DataFrame: + return self.read.table(tableName) + + table.__doc__ = PySparkSession.table.__doc__ + + @property + def read(self) -> "DataFrameReader": + return DataFrameReader(self) + + read.__doc__ = PySparkSession.read.__doc__ + + @property + def readStream(self) -> "DataStreamReader": + return DataStreamReader(self) + + readStream.__doc__ = PySparkSession.readStream.__doc__ + + def _inferSchemaFromList( + self, data: Iterable[Any], names: Optional[List[str]] = None + ) -> StructType: + """ + Infer schema from list of Row, dict, or tuple. + """ + if not data: + raise PySparkValueError( + error_class="CANNOT_INFER_EMPTY_SCHEMA", + message_parameters={}, + ) + + ( + infer_dict_as_struct, + infer_array_from_first_element, + prefer_timestamp_ntz, + ) = self._client.get_configs( + "spark.sql.pyspark.inferNestedDictAsStruct.enabled", + "spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled", + "spark.sql.timestampType", + ) + return reduce( + _merge_type, + ( + _infer_schema( + row, + names, + infer_dict_as_struct=(infer_dict_as_struct == "true"), + infer_array_from_first_element=(infer_array_from_first_element == "true"), + prefer_timestamp_ntz=(prefer_timestamp_ntz == "TIMESTAMP_NTZ"), + ) + for row in data + ), + ) + + def createDataFrame( + self, + data: Union["pd.DataFrame", "np.ndarray", Iterable[Any]], + schema: Optional[Union[AtomicType, StructType, str, List[str], Tuple[str, ...]]] = None, + ) -> "DataFrame": + assert data is not None + if isinstance(data, DataFrame): + raise PySparkTypeError( + error_class="INVALID_TYPE", + message_parameters={"arg_name": "data", "data_type": "DataFrame"}, + ) + + _schema: Optional[Union[AtomicType, StructType]] = None + _cols: Optional[List[str]] = None + _num_cols: Optional[int] = None + + if isinstance(schema, str): + schema = self.client._analyze( # type: ignore[assignment] + method="ddl_parse", ddl_string=schema + ).parsed + + if isinstance(schema, (AtomicType, StructType)): + _schema = schema + if isinstance(schema, StructType): + _num_cols = len(schema.fields) + else: + _num_cols = 1 + + elif isinstance(schema, (list, tuple)): + # Must re-encode any unicode strings to be consistent with StructField names + _cols = [x.encode("utf-8") if not isinstance(x, str) else x for x in schema] + _num_cols = len(_cols) + + if isinstance(data, np.ndarray) and data.ndim not in [1, 2]: + raise PySparkValueError( + error_class="INVALID_NDARRAY_DIMENSION", + message_parameters={"dimensions": "1 or 2"}, + ) + elif isinstance(data, Sized) and len(data) == 0: + if _schema is not None: + return DataFrame.withPlan(LocalRelation(table=None, schema=_schema.json()), self) + else: + raise PySparkValueError( + error_class="CANNOT_INFER_EMPTY_SCHEMA", + message_parameters={}, + ) + + _table: Optional[pa.Table] = None + + if isinstance(data, pd.DataFrame): + # Logic was borrowed from `_create_from_pandas_with_arrow` in + # `pyspark.sql.pandas.conversion.py`. Should ideally deduplicate the logics. + + # If no schema supplied by user then get the names of columns only + if schema is None: + _cols = [str(x) if not isinstance(x, str) else x for x in data.columns] + elif isinstance(schema, (list, tuple)) and cast(int, _num_cols) < len(data.columns): + assert isinstance(_cols, list) + _cols.extend([f"_{i + 1}" for i in range(cast(int, _num_cols), len(data.columns))]) + _num_cols = len(_cols) + + # Determine arrow types to coerce data when creating batches + arrow_schema: Optional[pa.Schema] = None + spark_types: List[Optional[DataType]] + arrow_types: List[Optional[pa.DataType]] + if isinstance(schema, StructType): + deduped_schema = cast(StructType, _deduplicate_field_names(schema)) + spark_types = [field.dataType for field in deduped_schema.fields] + arrow_schema = to_arrow_schema(deduped_schema) + arrow_types = [field.type for field in arrow_schema] + _cols = [str(x) if not isinstance(x, str) else x for x in schema.fieldNames()] + elif isinstance(schema, DataType): + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW", + message_parameters={"data_type": str(schema)}, + ) + else: + # Any timestamps must be coerced to be compatible with Spark + spark_types = [ + TimestampType() + if is_datetime64_dtype(t) or is_datetime64tz_dtype(t) + else DayTimeIntervalType() + if is_timedelta64_dtype(t) + else None + for t in data.dtypes + ] + arrow_types = [to_arrow_type(dt) if dt is not None else None for dt in spark_types] + + timezone, safecheck = self._client.get_configs( + "spark.sql.session.timeZone", "spark.sql.execution.pandas.convertToArrowArraySafely" + ) + + ser = ArrowStreamPandasSerializer(cast(str, timezone), safecheck == "true") + + _table = pa.Table.from_batches( + [ + ser._create_batch( + [ + (c, at, st) + for (_, c), at, st in zip(data.items(), arrow_types, spark_types) + ] + ) + ] + ) + + if isinstance(schema, StructType): + assert arrow_schema is not None + _table = _table.rename_columns( + cast(StructType, _deduplicate_field_names(schema)).names + ).cast(arrow_schema) + + elif isinstance(data, np.ndarray): + if _cols is None: + if data.ndim == 1 or data.shape[1] == 1: + _cols = ["value"] + else: + _cols = ["_%s" % i for i in range(1, data.shape[1] + 1)] + + if data.ndim == 1: + if 1 != len(_cols): + raise PySparkValueError( + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={ + "expected_length": str(len(_cols)), + "actual_length": "1", + }, + ) + + _table = pa.Table.from_arrays([pa.array(data)], _cols) + else: + if data.shape[1] != len(_cols): + raise PySparkValueError( + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={ + "expected_length": str(len(_cols)), + "actual_length": str(data.shape[1]), + }, + ) + + _table = pa.Table.from_arrays( + [pa.array(data[::, i]) for i in range(0, data.shape[1])], _cols + ) + + # The _table should already have the proper column names. + _cols = None + + else: + _data = list(data) + + if isinstance(_data[0], dict): + # Sort the data to respect inferred schema. + # For dictionaries, we sort the schema in alphabetical order. + _data = [dict(sorted(d.items())) if d is not None else None for d in _data] + + elif not isinstance(_data[0], (Row, tuple, list, dict)) and not hasattr( + _data[0], "__dict__" + ): + # input data can be [1, 2, 3] + # we need to convert it to [[1], [2], [3]] to be able to infer schema. + _data = [[d] for d in _data] + + if _schema is not None: + if not isinstance(_schema, StructType): + _schema = StructType().add("value", _schema) + else: + _schema = self._inferSchemaFromList(_data, _cols) + + if _cols is not None and cast(int, _num_cols) < len(_cols): + _num_cols = len(_cols) + + if _has_nulltype(_schema): + # For cases like createDataFrame([("Alice", None, 80.1)], schema) + # we can not infer the schema from the data itself. + raise ValueError( + "Some of types cannot be determined after inferring, " + "a StructType Schema is required in this case" + ) + + from pyspark.sql.connect.conversion import LocalDataToArrowConversion + + # Spark Connect will try its best to build the Arrow table with the + # inferred schema in the client side, and then rename the columns and + # cast the datatypes in the server side. + _table = LocalDataToArrowConversion.convert(_data, _schema) + + # TODO: Beside the validation on number of columns, we should also check + # whether the Arrow Schema is compatible with the user provided Schema. + if _num_cols is not None and _num_cols != _table.shape[1]: + raise PySparkValueError( + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={ + "expected_length": str(_num_cols), + "actual_length": str(_table.shape[1]), + }, + ) + + if _schema is not None: + local_relation = LocalRelation(_table, schema=_schema.json()) + else: + local_relation = LocalRelation(_table) + + cache_threshold = self._client.get_configs("spark.sql.session.localRelationCacheThreshold") + plan: LogicalPlan = local_relation + if cache_threshold[0] is not None and int(cache_threshold[0]) <= _table.nbytes: + plan = CachedLocalRelation(self._cache_local_relation(local_relation)) + + df = DataFrame.withPlan(plan, self) + if _cols is not None and len(_cols) > 0: + df = df.toDF(*_cols) + return df + + createDataFrame.__doc__ = PySparkSession.createDataFrame.__doc__ + + def sql(self, sqlQuery: str, args: Optional[Union[Dict[str, Any], List]] = None) -> "DataFrame": + cmd = SQL(sqlQuery, args) + data, properties = self.client.execute_command(cmd.command(self._client)) + if "sql_command_result" in properties: + return DataFrame.withPlan(CachedRelation(properties["sql_command_result"]), self) + else: + return DataFrame.withPlan(SQL(sqlQuery, args), self) + + sql.__doc__ = PySparkSession.sql.__doc__ + + def range( + self, + start: int, + end: Optional[int] = None, + step: int = 1, + numPartitions: Optional[int] = None, + ) -> DataFrame: + if end is None: + actual_end = start + start = 0 + else: + actual_end = end + + if numPartitions is not None: + numPartitions = int(numPartitions) + + return DataFrame.withPlan( + Range( + start=int(start), end=int(actual_end), step=int(step), num_partitions=numPartitions + ), + self, + ) + + range.__doc__ = PySparkSession.range.__doc__ + + @property + def catalog(self) -> "Catalog": + from pyspark.sql.connect.catalog import Catalog + + if not hasattr(self, "_catalog"): + self._catalog = Catalog(self) + return self._catalog + + catalog.__doc__ = PySparkSession.catalog.__doc__ + + def __del__(self) -> None: + try: + # Try its best to close. + self.client.close() + except Exception: + pass + + def interruptAll(self) -> List[str]: + op_ids = self.client.interrupt_all() + assert op_ids is not None + return op_ids + + interruptAll.__doc__ = PySparkSession.interruptAll.__doc__ + + def interruptTag(self, tag: str) -> List[str]: + op_ids = self.client.interrupt_tag(tag) + assert op_ids is not None + return op_ids + + interruptTag.__doc__ = PySparkSession.interruptTag.__doc__ + + def interruptOperation(self, op_id: str) -> List[str]: + op_ids = self.client.interrupt_operation(op_id) + assert op_ids is not None + return op_ids + + interruptOperation.__doc__ = PySparkSession.interruptOperation.__doc__ + + def addTag(self, tag: str) -> None: + self.client.add_tag(tag) + + addTag.__doc__ = PySparkSession.addTag.__doc__ + + def removeTag(self, tag: str) -> None: + self.client.remove_tag(tag) + + removeTag.__doc__ = PySparkSession.removeTag.__doc__ + + def getTags(self) -> Set[str]: + return self.client.get_tags() + + getTags.__doc__ = PySparkSession.getTags.__doc__ + + def clearTags(self) -> None: + return self.client.clear_tags() + + clearTags.__doc__ = PySparkSession.clearTags.__doc__ + + def stop(self) -> None: + # Stopping the session will only close the connection to the current session (and + # the life cycle of the session is maintained by the server), + # whereas the regular PySpark session immediately terminates the Spark Context + # itself, meaning that stopping all Spark sessions. + # It is controversial to follow the existing the regular Spark session's behavior + # specifically in Spark Connect the Spark Connect server is designed for + # multi-tenancy - the remote client side cannot just stop the server and stop + # other remote clients being used from other users. + with SparkSession._lock: + self.client.close() + if self is SparkSession._default_session: + SparkSession._default_session = None + if self is getattr(SparkSession._active_session, "session", None): + SparkSession._active_session.session = None + + if "SPARK_LOCAL_REMOTE" in os.environ: + # When local mode is in use, follow the regular Spark session's + # behavior by terminating the Spark Connect server, + # meaning that you can stop local mode, and restart the Spark Connect + # client with a different remote address. + if PySparkSession._activeSession is not None: + PySparkSession._activeSession.stop() + del os.environ["SPARK_LOCAL_REMOTE"] + del os.environ["SPARK_CONNECT_MODE_ENABLED"] + if "SPARK_REMOTE" in os.environ: + del os.environ["SPARK_REMOTE"] + + stop.__doc__ = PySparkSession.stop.__doc__ + + @property + def is_stopped(self) -> bool: + """ + Returns if this session was stopped + """ + return self.client.is_closed + + @property + def conf(self) -> RuntimeConf: + return RuntimeConf(self.client) + + conf.__doc__ = PySparkSession.conf.__doc__ + + @property + def streams(self) -> "StreamingQueryManager": + return StreamingQueryManager(self) + + streams.__doc__ = PySparkSession.streams.__doc__ + + def __getattr__(self, name: str) -> Any: + if name in ["_jsc", "_jconf", "_jvm", "_jsparkSession"]: + raise PySparkAttributeError( + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", message_parameters={"attr_name": name} + ) + elif name in ["newSession", "sparkContext"]: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", message_parameters={"feature": f"{name}()"} + ) + return object.__getattribute__(self, name) + + @property + def udf(self) -> "UDFRegistration": + from pyspark.sql.connect.udf import UDFRegistration + + return UDFRegistration(self) + + udf.__doc__ = PySparkSession.udf.__doc__ + + @property + def udtf(self) -> "UDTFRegistration": + from pyspark.sql.connect.udtf import UDTFRegistration + + return UDTFRegistration(self) + + udtf.__doc__ = PySparkSession.udtf.__doc__ + + @property + def version(self) -> str: + result = self._client._analyze(method="spark_version").spark_version + assert result is not None + return result + + version.__doc__ = PySparkSession.version.__doc__ + + @property + def client(self) -> "SparkConnectClient": + return self._client + + client.__doc__ = PySparkSession.client.__doc__ + + def addArtifacts( + self, *path: str, pyfile: bool = False, archive: bool = False, file: bool = False + ) -> None: + if sum([file, pyfile, archive]) > 1: + raise ValueError("'pyfile', 'archive' and/or 'file' cannot be True together.") + self._client.add_artifacts(*path, pyfile=pyfile, archive=archive, file=file) + + addArtifacts.__doc__ = PySparkSession.addArtifacts.__doc__ + + addArtifact = addArtifacts + + def _cache_local_relation(self, local_relation: LocalRelation) -> str: + """ + Cache the local relation at the server side if it has not been cached yet. + """ + serialized = local_relation.serialize(self._client) + return self._client.cache_artifact(serialized) + + def copyFromLocalToFs(self, local_path: str, dest_path: str) -> None: + if urllib.parse.urlparse(dest_path).scheme: + raise ValueError( + "`spark_session.copyFromLocalToFs` API only allows `dest_path` to be a path " + "without scheme, and spark driver uses the default scheme to " + "determine the destination file system." + ) + self._client.copy_from_local_to_fs(local_path, dest_path) + + copyFromLocalToFs.__doc__ = PySparkSession.copyFromLocalToFs.__doc__ + + def _create_remote_dataframe(self, remote_id: str) -> "DataFrame": + """ + In internal API to reference a runtime DataFrame on the server side. + This is used in ForeachBatch() runner, where the remote DataFrame refers to the + output of a micro batch. + """ + return DataFrame.withPlan(CachedRemoteRelation(remote_id), self) + + @staticmethod + def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: + """ + Starts the Spark Connect server given the master (thread-unsafe). + + At the high level, there are two cases. The first case is development case, e.g., + you locally build Apache Spark, and run ``SparkSession.builder.remote("local")``: + + 1. This method automatically finds the jars for Spark Connect (because the jars for + Spark Connect are not bundled in the regular Apache Spark release). + + 2. Temporarily remove all states for Spark Connect, for example, ``SPARK_REMOTE`` + environment variable. + + 3. Starts a JVM (without Spark Context) first, and adds the Spark Connect server jars + into the current class loader. Otherwise, Spark Context with ``spark.plugins`` + cannot be initialized because the JVM is already running without the jars in + the classpath before executing this Python process for driver side (in case of + PySpark application submission). + + 4. Starts a regular Spark session that automatically starts a Spark Connect server + via ``spark.plugins`` feature. + + The second case is when you use Apache Spark release: + + 1. Users must specify either the jars or package, e.g., ``--packages + org.apache.spark:spark-connect_2.12:3.4.0``. The jars or packages would be specified + in SparkSubmit automatically. This method does not do anything related to this. + + 2. Temporarily remove all states for Spark Connect, for example, ``SPARK_REMOTE`` + environment variable. It does not do anything for PySpark application submission as + well because jars or packages were already specified before executing this Python + process for driver side. + + 3. Starts a regular Spark session that automatically starts a Spark Connect server + with JVM via ``spark.plugins`` feature. + """ + session = PySparkSession._instantiatedSession + if session is None or session._sc._jsc is None: + + # Configurations to be overwritten + overwrite_conf = opts + overwrite_conf["spark.master"] = master + overwrite_conf["spark.local.connect"] = "1" + + # Configurations to be set if unset. + default_conf = {"spark.plugins": "org.apache.spark.sql.connect.SparkConnectPlugin"} + + if "SPARK_TESTING" in os.environ: + # For testing, we use 0 to use an ephemeral port to allow parallel testing. + # See also SPARK-42272. + overwrite_conf["spark.connect.grpc.binding.port"] = "0" + + def create_conf(**kwargs: Any) -> SparkConf: + conf = SparkConf(**kwargs) + for k, v in overwrite_conf.items(): + conf.set(k, v) + for k, v in default_conf.items(): + if not conf.contains(k): + conf.set(k, v) + return conf + + # Check if we're using unreleased version that is in development. + # Also checks SPARK_TESTING for RC versions. + is_dev_mode = ( + "dev" in LooseVersion(__version__).version or "SPARK_TESTING" in os.environ + ) + + origin_remote = os.environ.get("SPARK_REMOTE", None) + try: + if origin_remote is not None: + # So SparkSubmit thinks no remote is set in order to + # start the regular PySpark session. + del os.environ["SPARK_REMOTE"] + + SparkContext._ensure_initialized(conf=create_conf(loadDefaults=False)) + + if is_dev_mode: + # Try and catch for a possibility in production because pyspark.testing + # does not exist in the canonical release. + try: + from pyspark.testing.utils import search_jar + + # Note that, in production, spark.jars.packages configuration should be + # set by users. Here we're automatically searching the jars locally built. + connect_jar = search_jar( + "connector/connect/server", "spark-connect-assembly-", "spark-connect" + ) + if connect_jar is None: + warnings.warn( + "Attempted to automatically find the Spark Connect jars because " + "'SPARK_TESTING' environment variable is set, or the current " + f"PySpark version is dev version ({__version__}). However, the jar" + " was not found. Manually locate the jars and specify them, e.g., " + "'spark.jars' configuration." + ) + else: + pyutils = SparkContext._jvm.PythonSQLUtils # type: ignore[union-attr] + pyutils.addJarToCurrentClassLoader(connect_jar) + + # Required for local-cluster testing as their executors need the jars + # to load the Spark plugin for Spark Connect. + if master.startswith("local-cluster"): + if "spark.jars" in overwrite_conf: + overwrite_conf[ + "spark.jars" + ] = f"{overwrite_conf['spark.jars']},{connect_jar}" + else: + overwrite_conf["spark.jars"] = connect_jar + + except ImportError: + pass + + # The regular PySpark session is registered as an active session + # so would not be garbage-collected. + PySparkSession( + SparkContext.getOrCreate(create_conf(loadDefaults=True, _jvm=SparkContext._jvm)) + ) + + # Lastly only keep runtime configurations because other configurations are + # disallowed to set in the regular Spark Connect session. + utl = SparkContext._jvm.PythonSQLUtils # type: ignore[union-attr] + runtime_conf_keys = [c._1() for c in utl.listRuntimeSQLConfigs()] + new_opts = {k: opts[k] for k in opts if k in runtime_conf_keys} + opts.clear() + opts.update(new_opts) + + finally: + if origin_remote is not None: + os.environ["SPARK_REMOTE"] = origin_remote + else: + raise PySparkRuntimeError( + error_class="SESSION_OR_CONTEXT_EXISTS", + message_parameters={}, + ) + + @property + def session_id(self) -> str: + return self._session_id + + +SparkSession.__doc__ = PySparkSession.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.session + + globs = pyspark.sql.connect.session.__dict__.copy() + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.session tests").remote("local[4]").getOrCreate() + ) + + # Uses PySpark session to test builder. + globs["SparkSession"] = PySparkSession + # Spark Connect does not support to set master together. + pyspark.sql.connect.session.SparkSession.__doc__ = None + del pyspark.sql.connect.session.SparkSession.Builder.master.__doc__ + # RDD API is not supported in Spark Connect. + del pyspark.sql.connect.session.SparkSession.createDataFrame.__doc__ + + # TODO(SPARK-41811): Implement SparkSession.sql's string formatter + del pyspark.sql.connect.session.SparkSession.sql.__doc__ + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.session, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/__init__.py new file mode 100644 index 0000000..cc50ff1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/__init__.py @@ -0,0 +1,22 @@ +# +# 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.sql.connect.streaming.query import StreamingQuery # noqa: F401 +from pyspark.sql.connect.streaming.readwriter import DataStreamReader # noqa: F401 +from pyspark.sql.connect.streaming.readwriter import DataStreamWriter # noqa: F401 +from pyspark.sql.connect.streaming.query import StreamingQueryManager # noqa: F401 +from pyspark.errors import StreamingQueryException # noqa: F401 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/query.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/query.py new file mode 100644 index 0000000..021d27e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/query.py @@ -0,0 +1,295 @@ +# +# 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 json +import sys +from typing import TYPE_CHECKING, Any, cast, Dict, List, Optional + +from pyspark.errors import StreamingQueryException, PySparkValueError +import pyspark.sql.connect.proto as pb2 +from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.connect import proto +from pyspark.sql.connect.utils import get_python_ver +from pyspark.sql.streaming import StreamingQueryListener +from pyspark.sql.streaming.query import ( + StreamingQuery as PySparkStreamingQuery, + StreamingQueryManager as PySparkStreamingQueryManager, +) +from pyspark.errors.exceptions.connect import ( + StreamingQueryException as CapturedStreamingQueryException, +) + +__all__ = ["StreamingQuery", "StreamingQueryManager"] + +if TYPE_CHECKING: + from pyspark.sql.connect.session import SparkSession + + +class StreamingQuery: + def __init__( + self, session: "SparkSession", queryId: str, runId: str, name: Optional[str] = None + ) -> None: + self._session = session + self._query_id = queryId + self._run_id = runId + self._name = name + + @property + def id(self) -> str: + return self._query_id + + id.__doc__ = PySparkStreamingQuery.id.__doc__ + + @property + def runId(self) -> str: + return self._run_id + + runId.__doc__ = PySparkStreamingQuery.runId.__doc__ + + @property + def name(self) -> Optional[str]: + return self._name + + name.__doc__ = PySparkStreamingQuery.name.__doc__ + + @property + def isActive(self) -> bool: + return self._fetch_status().is_active + + isActive.__doc__ = PySparkStreamingQuery.isActive.__doc__ + + def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]: + cmd = pb2.StreamingQueryCommand() + if timeout is not None: + if not isinstance(timeout, (int, float)) or timeout <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "timeout", "arg_value": type(timeout).__name__}, + ) + cmd.await_termination.timeout_ms = int(timeout * 1000) + terminated = self._execute_streaming_query_cmd(cmd).await_termination.terminated + return terminated + else: + await_termination_cmd = pb2.StreamingQueryCommand.AwaitTerminationCommand() + cmd.await_termination.CopyFrom(await_termination_cmd) + self._execute_streaming_query_cmd(cmd) + return None + + awaitTermination.__doc__ = PySparkStreamingQuery.awaitTermination.__doc__ + + @property + def status(self) -> Dict[str, Any]: + proto = self._fetch_status() + return { + "message": proto.status_message, + "isDataAvailable": proto.is_data_available, + "isTriggerActive": proto.is_trigger_active, + } + + status.__doc__ = PySparkStreamingQuery.status.__doc__ + + @property + def recentProgress(self) -> List[Dict[str, Any]]: + cmd = pb2.StreamingQueryCommand() + cmd.recent_progress = True + progress = self._execute_streaming_query_cmd(cmd).recent_progress.recent_progress_json + return [json.loads(p) for p in progress] + + recentProgress.__doc__ = PySparkStreamingQuery.recentProgress.__doc__ + + @property + def lastProgress(self) -> Optional[Dict[str, Any]]: + cmd = pb2.StreamingQueryCommand() + cmd.last_progress = True + progress = self._execute_streaming_query_cmd(cmd).recent_progress.recent_progress_json + if len(progress) > 0: + return json.loads(progress[-1]) + else: + return None + + lastProgress.__doc__ = PySparkStreamingQuery.lastProgress.__doc__ + + def processAllAvailable(self) -> None: + cmd = pb2.StreamingQueryCommand() + cmd.process_all_available = True + self._execute_streaming_query_cmd(cmd) + + processAllAvailable.__doc__ = PySparkStreamingQuery.processAllAvailable.__doc__ + + def stop(self) -> None: + cmd = pb2.StreamingQueryCommand() + cmd.stop = True + self._execute_streaming_query_cmd(cmd) + + stop.__doc__ = PySparkStreamingQuery.stop.__doc__ + + def explain(self, extended: bool = False) -> None: + cmd = pb2.StreamingQueryCommand() + cmd.explain.extended = extended + result = self._execute_streaming_query_cmd(cmd).explain.result + print(result) + + explain.__doc__ = PySparkStreamingQuery.explain.__doc__ + + def exception(self) -> Optional[StreamingQueryException]: + cmd = pb2.StreamingQueryCommand() + cmd.exception = True + exception = self._execute_streaming_query_cmd(cmd).exception + if not exception.HasField("exception_message"): + return None + else: + # Drop the Java StreamingQueryException type info + # exception_message maps to the return value of original + # StreamingQueryException's toString method + msg = exception.exception_message.split(": ", 1)[1] + if exception.HasField("stack_trace"): + msg += f"\n\nJVM stacktrace:\n{exception.stack_trace}" + return CapturedStreamingQueryException(msg, reason=exception.error_class) + + exception.__doc__ = PySparkStreamingQuery.exception.__doc__ + + def _fetch_status(self) -> pb2.StreamingQueryCommandResult.StatusResult: + cmd = pb2.StreamingQueryCommand() + cmd.status = True + return self._execute_streaming_query_cmd(cmd).status + + def _execute_streaming_query_cmd( + self, cmd: pb2.StreamingQueryCommand + ) -> pb2.StreamingQueryCommandResult: + cmd.query_id.id = self._query_id + cmd.query_id.run_id = self._run_id + exec_cmd = pb2.Command() + exec_cmd.streaming_query_command.CopyFrom(cmd) + (_, properties) = self._session.client.execute_command(exec_cmd) + return cast(pb2.StreamingQueryCommandResult, properties["streaming_query_command_result"]) + + +class StreamingQueryManager: + def __init__(self, session: "SparkSession") -> None: + self._session = session + + @property + def active(self) -> List[StreamingQuery]: + cmd = pb2.StreamingQueryManagerCommand() + cmd.active = True + queries = self._execute_streaming_query_manager_cmd(cmd).active.active_queries + return [StreamingQuery(self._session, q.id.id, q.id.run_id, q.name) for q in queries] + + active.__doc__ = PySparkStreamingQueryManager.active.__doc__ + + def get(self, id: str) -> Optional[StreamingQuery]: + cmd = pb2.StreamingQueryManagerCommand() + cmd.get_query = id + response = self._execute_streaming_query_manager_cmd(cmd) + if response.HasField("query"): + query = response.query + return StreamingQuery(self._session, query.id.id, query.id.run_id, query.name) + else: + return None + + get.__doc__ = PySparkStreamingQueryManager.get.__doc__ + + def awaitAnyTermination(self, timeout: Optional[int] = None) -> Optional[bool]: + cmd = pb2.StreamingQueryManagerCommand() + if timeout is not None: + if not isinstance(timeout, (int, float)) or timeout <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "timeout", "arg_value": type(timeout).__name__}, + ) + cmd.await_any_termination.timeout_ms = int(timeout * 1000) + terminated = self._execute_streaming_query_manager_cmd( + cmd + ).await_any_termination.terminated + return terminated + else: + await_any_termination_cmd = ( + pb2.StreamingQueryManagerCommand.AwaitAnyTerminationCommand() + ) + cmd.await_any_termination.CopyFrom(await_any_termination_cmd) + self._execute_streaming_query_manager_cmd(cmd) + return None + + awaitAnyTermination.__doc__ = PySparkStreamingQueryManager.awaitAnyTermination.__doc__ + + def resetTerminated(self) -> None: + cmd = pb2.StreamingQueryManagerCommand() + cmd.reset_terminated = True + self._execute_streaming_query_manager_cmd(cmd) + + resetTerminated.__doc__ = PySparkStreamingQueryManager.resetTerminated.__doc__ + + def addListener(self, listener: StreamingQueryListener) -> None: + listener._init_listener_id() + cmd = pb2.StreamingQueryManagerCommand() + expr = proto.PythonUDF() + expr.command = CloudPickleSerializer().dumps(listener) + expr.python_ver = get_python_ver() + cmd.add_listener.python_listener_payload.CopyFrom(expr) + cmd.add_listener.id = listener._id + self._execute_streaming_query_manager_cmd(cmd) + + addListener.__doc__ = PySparkStreamingQueryManager.addListener.__doc__ + + def removeListener(self, listener: StreamingQueryListener) -> None: + cmd = pb2.StreamingQueryManagerCommand() + cmd.remove_listener.id = listener._id + self._execute_streaming_query_manager_cmd(cmd) + + removeListener.__doc__ = PySparkStreamingQueryManager.removeListener.__doc__ + + def _execute_streaming_query_manager_cmd( + self, cmd: pb2.StreamingQueryManagerCommand + ) -> pb2.StreamingQueryManagerCommandResult: + exec_cmd = pb2.Command() + exec_cmd.streaming_query_manager_command.CopyFrom(cmd) + (_, properties) = self._session.client.execute_command(exec_cmd) + return cast( + pb2.StreamingQueryManagerCommandResult, + properties["streaming_query_manager_command_result"], + ) + + +def _test() -> None: + import doctest + import os + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.streaming.query + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.connect.streaming.query.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.streaming.query tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.streaming.query, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/readwriter.py new file mode 100644 index 0000000..89097fc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/readwriter.py @@ -0,0 +1,617 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import sys +from typing import cast, overload, Callable, Dict, List, Optional, TYPE_CHECKING, Union + +from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.connect.plan import DataSource, LogicalPlan, Read, WriteStreamOperation +import pyspark.sql.connect.proto as pb2 +from pyspark.sql.connect.readwriter import OptionUtils, to_str +from pyspark.sql.connect.streaming.query import StreamingQuery +from pyspark.sql.streaming.readwriter import ( + DataStreamReader as PySparkDataStreamReader, + DataStreamWriter as PySparkDataStreamWriter, +) +from pyspark.sql.connect.utils import get_python_ver +from pyspark.sql.types import Row, StructType +from pyspark.errors import PySparkTypeError, PySparkValueError + +if TYPE_CHECKING: + from pyspark.sql.connect.session import SparkSession + from pyspark.sql.connect._typing import OptionalPrimitiveType + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql._typing import SupportsProcess + +__all__ = ["DataStreamReader", "DataStreamWriter"] + + +class DataStreamReader(OptionUtils): + def __init__(self, client: "SparkSession") -> None: + self._format: Optional[str] = None + self._schema = "" + self._client = client + self._options: Dict[str, str] = {} + + def _df(self, plan: LogicalPlan) -> "DataFrame": + from pyspark.sql.connect.dataframe import DataFrame + + return DataFrame.withPlan(plan, self._client) + + def format(self, source: str) -> "DataStreamReader": + self._format = source + return self + + format.__doc__ = PySparkDataStreamReader.format.__doc__ + + def schema(self, schema: Union[StructType, str]) -> "DataStreamReader": + if isinstance(schema, StructType): + self._schema = schema.json() + elif isinstance(schema, str): + self._schema = schema + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_STRUCT", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + return self + + schema.__doc__ = PySparkDataStreamReader.schema.__doc__ + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataStreamReader": + self._options[key] = str(value) + return self + + option.__doc__ = PySparkDataStreamReader.option.__doc__ + + def options(self, **options: "OptionalPrimitiveType") -> "DataStreamReader": + for k in options: + self.option(k, to_str(options[k])) + return self + + options.__doc__ = PySparkDataStreamReader.options.__doc__ + + def load( + self, + path: Optional[str] = None, + format: Optional[str] = None, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> "DataFrame": + if format is not None: + self.format(format) + if schema is not None: + self.schema(schema) + self.options(**options) + if path is not None and (type(path) != str or len(path.strip()) == 0): + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "path", "arg_value": str(path)}, + ) + + plan = DataSource( + format=self._format, + schema=self._schema, + options=self._options, + paths=[path] if path else None, + is_streaming=True, + ) + + return self._df(plan) + + load.__doc__ = PySparkDataStreamReader.load.__doc__ + + def json( + self, + path: str, + schema: Optional[Union[StructType, str]] = None, + primitivesAsString: Optional[Union[bool, str]] = None, + prefersDecimal: Optional[Union[bool, str]] = None, + allowComments: Optional[Union[bool, str]] = None, + allowUnquotedFieldNames: Optional[Union[bool, str]] = None, + allowSingleQuotes: Optional[Union[bool, str]] = None, + allowNumericLeadingZero: Optional[Union[bool, str]] = None, + allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + allowUnquotedControlChars: Optional[Union[bool, str]] = None, + lineSep: Optional[str] = None, + locale: Optional[str] = None, + dropFieldIfAllNull: Optional[Union[bool, str]] = None, + encoding: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + allowNonNumericNumbers: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + schema=schema, + primitivesAsString=primitivesAsString, + prefersDecimal=prefersDecimal, + allowComments=allowComments, + allowUnquotedFieldNames=allowUnquotedFieldNames, + allowSingleQuotes=allowSingleQuotes, + allowNumericLeadingZero=allowNumericLeadingZero, + allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + multiLine=multiLine, + allowUnquotedControlChars=allowUnquotedControlChars, + lineSep=lineSep, + locale=locale, + dropFieldIfAllNull=dropFieldIfAllNull, + encoding=encoding, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + allowNonNumericNumbers=allowNonNumericNumbers, + ) + if isinstance(path, str): + return self.load(path=path, format="json") + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + json.__doc__ = PySparkDataStreamReader.json.__doc__ + + def orc( + self, + path: str, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + return self.load(path=path, format="orc") + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + orc.__doc__ = PySparkDataStreamReader.orc.__doc__ + + def parquet( + self, + path: str, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + datetimeRebaseMode: Optional[Union[bool, str]] = None, + int96RebaseMode: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + datetimeRebaseMode=datetimeRebaseMode, + int96RebaseMode=int96RebaseMode, + ) + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + datetimeRebaseMode=datetimeRebaseMode, + int96RebaseMode=int96RebaseMode, + ) + if isinstance(path, str): + return self.load(path=path, format="parquet") + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + parquet.__doc__ = PySparkDataStreamReader.parquet.__doc__ + + def text( + self, + path: str, + wholetext: bool = False, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + self._set_opts( + wholetext=wholetext, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + return self.load(path=path, format="text") + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + text.__doc__ = PySparkDataStreamReader.text.__doc__ + + def csv( + self, + path: str, + schema: Optional[Union[StructType, str]] = None, + sep: Optional[str] = None, + encoding: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + comment: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + inferSchema: Optional[Union[bool, str]] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + nanValue: Optional[str] = None, + positiveInf: Optional[str] = None, + negativeInf: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + maxColumns: Optional[Union[int, str]] = None, + maxCharsPerColumn: Optional[Union[int, str]] = None, + maxMalformedLogPerPartition: Optional[Union[int, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[Union[bool, str]] = None, + enforceSchema: Optional[Union[bool, str]] = None, + emptyValue: Optional[str] = None, + locale: Optional[str] = None, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + unescapedQuoteHandling: Optional[str] = None, + ) -> "DataFrame": + self._set_opts( + schema=schema, + sep=sep, + encoding=encoding, + quote=quote, + escape=escape, + comment=comment, + header=header, + inferSchema=inferSchema, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + nullValue=nullValue, + nanValue=nanValue, + positiveInf=positiveInf, + negativeInf=negativeInf, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + maxColumns=maxColumns, + maxCharsPerColumn=maxCharsPerColumn, + maxMalformedLogPerPartition=maxMalformedLogPerPartition, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + multiLine=multiLine, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + enforceSchema=enforceSchema, + emptyValue=emptyValue, + locale=locale, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + unescapedQuoteHandling=unescapedQuoteHandling, + ) + if isinstance(path, str): + return self.load(path=path, format="csv") + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + csv.__doc__ = PySparkDataStreamReader.csv.__doc__ + + def table(self, tableName: str) -> "DataFrame": + return self._df(Read(tableName, self._options, is_streaming=True)) + + table.__doc__ = PySparkDataStreamReader.table.__doc__ + + +DataStreamReader.__doc__ = PySparkDataStreamReader.__doc__ + + +class DataStreamWriter: + def __init__(self, plan: "LogicalPlan", session: "SparkSession") -> None: + self._session = session + self._write_stream = WriteStreamOperation(plan) + self._write_proto = self._write_stream.write_op + + def outputMode(self, outputMode: str) -> "DataStreamWriter": + self._write_proto.output_mode = outputMode + return self + + outputMode.__doc__ = PySparkDataStreamWriter.outputMode.__doc__ + + def format(self, source: str) -> "DataStreamWriter": + self._write_proto.format = source + return self + + format.__doc__ = PySparkDataStreamWriter.format.__doc__ + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataStreamWriter": + self._write_proto.options[key] = cast(str, to_str(value)) + return self + + option.__doc__ = PySparkDataStreamWriter.option.__doc__ + + def options(self, **options: "OptionalPrimitiveType") -> "DataStreamWriter": + for k in options: + self.option(k, options[k]) + return self + + options.__doc__ = PySparkDataStreamWriter.options.__doc__ + + @overload + def partitionBy(self, *cols: str) -> "DataStreamWriter": + ... + + @overload + def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": + ... + + def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + if len(cols) == 1 and isinstance(cols[0], (list, tuple)): + cols = cols[0] + # Clear any existing columns (if any). + while len(self._write_proto.partitioning_column_names) > 0: + self._write_proto.partitioning_column_names.pop() + self._write_proto.partitioning_column_names.extend(cast(List[str], cols)) + return self + + partitionBy.__doc__ = PySparkDataStreamWriter.partitionBy.__doc__ + + def queryName(self, queryName: str) -> "DataStreamWriter": + self._write_proto.query_name = queryName + return self + + queryName.__doc__ = PySparkDataStreamWriter.queryName.__doc__ + + @overload + def trigger(self, *, processingTime: str) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, once: bool) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, continuous: str) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, availableNow: bool) -> "DataStreamWriter": + ... + + def trigger( + self, + *, + processingTime: Optional[str] = None, + once: Optional[bool] = None, + continuous: Optional[str] = None, + availableNow: Optional[bool] = None, + ) -> "DataStreamWriter": + params = [processingTime, once, continuous, availableNow] + + if params.count(None) == 4: + raise PySparkValueError( + error_class="ONLY_ALLOW_SINGLE_TRIGGER", + message_parameters={}, + ) + elif params.count(None) < 3: + raise PySparkValueError( + error_class="ONLY_ALLOW_SINGLE_TRIGGER", + message_parameters={}, + ) + + if processingTime is not None: + if type(processingTime) != str or len(processingTime.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={ + "arg_name": "processingTime", + "arg_value": str(processingTime), + }, + ) + self._write_proto.processing_time_interval = processingTime.strip() + + elif once is not None: + if once is not True: + raise PySparkValueError( + error_class="VALUE_NOT_TRUE", + message_parameters={"arg_name": "once", "arg_value": str(once)}, + ) + self._write_proto.once = True + + elif continuous is not None: + if type(continuous) != str or len(continuous.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "continuous", "arg_value": str(continuous)}, + ) + self._write_proto.continuous_checkpoint_interval = continuous.strip() + + else: + if availableNow is not True: + raise PySparkValueError( + error_class="VALUE_NOT_TRUE", + message_parameters={"arg_name": "availableNow", "arg_value": str(availableNow)}, + ) + self._write_proto.available_now = True + + return self + + trigger.__doc__ = PySparkDataStreamWriter.trigger.__doc__ + + @overload + def foreach(self, f: Callable[[Row], None]) -> "DataStreamWriter": + ... + + @overload + def foreach(self, f: "SupportsProcess") -> "DataStreamWriter": + ... + + def foreach(self, f: Union[Callable[[Row], None], "SupportsProcess"]) -> "DataStreamWriter": + from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer + + func = PySparkDataStreamWriter._construct_foreach_function(f) + serializer = AutoBatchedSerializer(CPickleSerializer()) + command = (func, None, serializer, serializer) + # Python ForeachWriter isn't really a PythonUDF. But we reuse it for simplicity. + self._write_proto.foreach_writer.python_function.command = CloudPickleSerializer().dumps( + command + ) + self._write_proto.foreach_writer.python_function.python_ver = "%d.%d" % sys.version_info[:2] + return self + + foreach.__doc__ = PySparkDataStreamWriter.foreach.__doc__ + + def foreachBatch(self, func: Callable[["DataFrame", int], None]) -> "DataStreamWriter": + self._write_proto.foreach_batch.python_function.command = CloudPickleSerializer().dumps( + func + ) + self._write_proto.foreach_batch.python_function.python_ver = get_python_ver() + return self + + foreachBatch.__doc__ = PySparkDataStreamWriter.foreachBatch.__doc__ + + def _start_internal( + self, + path: Optional[str] = None, + tableName: Optional[str] = None, + format: Optional[str] = None, + outputMode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + queryName: Optional[str] = None, + **options: "OptionalPrimitiveType", + ) -> StreamingQuery: + self.options(**options) + if outputMode is not None: + self.outputMode(outputMode) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + if queryName is not None: + self.queryName(queryName) + if path: + self._write_proto.path = path + if tableName: + self._write_proto.table_name = tableName + + cmd = self._write_stream.command(self._session.client) + (_, properties) = self._session.client.execute_command(cmd) + + start_result = cast( + pb2.WriteStreamOperationStartResult, properties["write_stream_operation_start_result"] + ) + return StreamingQuery( + session=self._session, + queryId=start_result.query_id.id, + runId=start_result.query_id.run_id, + name=start_result.name, + ) + + def start( + self, + path: Optional[str] = None, + format: Optional[str] = None, + outputMode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + queryName: Optional[str] = None, + **options: "OptionalPrimitiveType", + ) -> StreamingQuery: + return self._start_internal( + path=path, + tableName=None, + format=format, + outputMode=outputMode, + partitionBy=partitionBy, + queryName=queryName, + **options, + ) + + start.__doc__ = PySparkDataStreamWriter.start.__doc__ + + def toTable( + self, + tableName: str, + format: Optional[str] = None, + outputMode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + queryName: Optional[str] = None, + **options: "OptionalPrimitiveType", + ) -> StreamingQuery: + return self._start_internal( + path=None, + tableName=tableName, + format=format, + outputMode=outputMode, + partitionBy=partitionBy, + queryName=queryName, + **options, + ) + + toTable.__doc__ = PySparkDataStreamWriter.toTable.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.streaming.readwriter + + globs = pyspark.sql.connect.readwriter.__dict__.copy() + + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.streaming.readwriter tests") + .remote("local[4]") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.streaming.readwriter, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/__init__.py new file mode 100644 index 0000000..a5c9801 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +"""Spark Connect Streaming Server-side Worker""" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py new file mode 100644 index 0000000..cf61463 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py @@ -0,0 +1,87 @@ +# +# 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. +# + +""" +A worker for streaming foreachBatch in Spark Connect. +Usually this is ran on the driver side of the Spark Connect Server. +""" +import os + +from pyspark.java_gateway import local_connect_and_auth +from pyspark.serializers import ( + write_int, + read_long, + UTF8Deserializer, + CPickleSerializer, +) +from pyspark import worker +from pyspark.sql import SparkSession +from typing import IO +from pyspark.worker_util import check_python_version + +pickle_ser = CPickleSerializer() +utf8_deserializer = UTF8Deserializer() + + +def main(infile: IO, outfile: IO) -> None: + check_python_version(infile) + + connect_url = os.environ["SPARK_CONNECT_LOCAL_URL"] + session_id = utf8_deserializer.loads(infile) + + print( + "Streaming foreachBatch worker is starting with " + f"url {connect_url} and sessionId {session_id}." + ) + + spark_connect_session = SparkSession.builder.remote(connect_url).getOrCreate() + spark_connect_session._client._session_id = session_id # type: ignore[attr-defined] + + # TODO(SPARK-44460): Pass credentials. + # TODO(SPARK-44461): Enable Process Isolation + + func = worker.read_command(pickle_ser, infile) + write_int(0, outfile) # Indicate successful initialization + + outfile.flush() + + log_name = "Streaming ForeachBatch worker" + + def process(df_id, batch_id): # type: ignore[no-untyped-def] + print(f"{log_name} Started batch {batch_id} with DF id {df_id}") + batch_df = spark_connect_session._create_remote_dataframe(df_id) + func(batch_df, batch_id) + print(f"{log_name} Completed batch {batch_id} with DF id {df_id}") + + while True: + df_ref_id = utf8_deserializer.loads(infile) + batch_id = read_long(infile) + process(df_ref_id, int(batch_id)) # TODO(SPARK-44463): Propagate error to the user. + write_int(0, outfile) + outfile.flush() + + +if __name__ == "__main__": + # Read information about how to connect back to the JVM from the environment. + java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"]) + auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"] + (sock_file, sock) = local_connect_and_auth(java_port, auth_secret) + # There could be a long time between each micro batch. + sock.settimeout(None) + write_int(os.getpid(), sock_file) + sock_file.flush() + main(sock_file, sock_file) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/listener_worker.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/listener_worker.py new file mode 100644 index 0000000..e1f4678 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/streaming/worker/listener_worker.py @@ -0,0 +1,100 @@ +# +# 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. +# + +""" +A worker for streaming query listener in Spark Connect. +Usually this is ran on the driver side of the Spark Connect Server. +""" +import os +import json + +from pyspark.java_gateway import local_connect_and_auth +from pyspark.serializers import ( + read_int, + write_int, + UTF8Deserializer, + CPickleSerializer, +) +from pyspark import worker +from pyspark.sql import SparkSession +from typing import IO + +from pyspark.sql.streaming.listener import ( + QueryStartedEvent, + QueryProgressEvent, + QueryTerminatedEvent, + QueryIdleEvent, +) +from pyspark.worker_util import check_python_version + +pickle_ser = CPickleSerializer() +utf8_deserializer = UTF8Deserializer() + + +def main(infile: IO, outfile: IO) -> None: + check_python_version(infile) + + connect_url = os.environ["SPARK_CONNECT_LOCAL_URL"] + session_id = utf8_deserializer.loads(infile) + + print( + "Streaming query listener worker is starting with " + f"url {connect_url} and sessionId {session_id}." + ) + + spark_connect_session = SparkSession.builder.remote(connect_url).getOrCreate() + spark_connect_session._client._session_id = session_id # type: ignore[attr-defined] + + # TODO(SPARK-44460): Pass credentials. + # TODO(SPARK-44461): Enable Process Isolation + + listener = worker.read_command(pickle_ser, infile) + write_int(0, outfile) # Indicate successful initialization + + outfile.flush() + + listener._set_spark_session(spark_connect_session) + assert listener.spark == spark_connect_session + + def process(listener_event_str, listener_event_type): # type: ignore[no-untyped-def] + listener_event = json.loads(listener_event_str) + if listener_event_type == 0: + listener.onQueryStarted(QueryStartedEvent.fromJson(listener_event)) + elif listener_event_type == 1: + listener.onQueryProgress(QueryProgressEvent.fromJson(listener_event)) + elif listener_event_type == 2: + listener.onQueryIdle(QueryIdleEvent.fromJson(listener_event)) + elif listener_event_type == 3: + listener.onQueryTerminated(QueryTerminatedEvent.fromJson(listener_event)) + + while True: + event = utf8_deserializer.loads(infile) + event_type = read_int(infile) + process(event, int(event_type)) # TODO(SPARK-44463): Propagate error to the user. + outfile.flush() + + +if __name__ == "__main__": + # Read information about how to connect back to the JVM from the environment. + java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"]) + auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"] + (sock_file, sock) = local_connect_and_auth(java_port, auth_secret) + # There could be a long time between each listener event. + sock.settimeout(None) + write_int(os.getpid(), sock_file) + sock_file.flush() + main(sock_file, sock_file) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/types.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/types.py new file mode 100644 index 0000000..0db2833 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/types.py @@ -0,0 +1,301 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import json + +from typing import Any, Dict, Optional + +from pyspark.sql.types import ( + DataType, + ByteType, + ShortType, + IntegerType, + FloatType, + DateType, + TimestampType, + TimestampNTZType, + DayTimeIntervalType, + YearMonthIntervalType, + MapType, + StringType, + CharType, + VarcharType, + StructType, + StructField, + ArrayType, + DoubleType, + LongType, + DecimalType, + BinaryType, + BooleanType, + NullType, + UserDefinedType, +) +from pyspark.errors import PySparkAssertionError + +import pyspark.sql.connect.proto as pb2 + + +JVM_BYTE_MIN: int = -(1 << 7) +JVM_BYTE_MAX: int = (1 << 7) - 1 +JVM_SHORT_MIN: int = -(1 << 15) +JVM_SHORT_MAX: int = (1 << 15) - 1 +JVM_INT_MIN: int = -(1 << 31) +JVM_INT_MAX: int = (1 << 31) - 1 +JVM_LONG_MIN: int = -(1 << 63) +JVM_LONG_MAX: int = (1 << 63) - 1 + + +class UnparsedDataType(DataType): + """ + Unparsed data type. + + The data type string will be parsed later. + + Parameters + ---------- + data_type_string : str + The data type string format equals :class:`DataType.simpleString`, + except that the top level struct type can omit the ``struct<>``. + This also supports a schema in a DDL-formatted string and case-insensitive strings. + + Examples + -------- + >>> from pyspark.sql.connect.types import UnparsedDataType + + >>> UnparsedDataType("int ") + UnparsedDataType('int ') + >>> UnparsedDataType("INT ") + UnparsedDataType('INT ') + >>> UnparsedDataType("a: byte, b: decimal( 16 , 8 ) ") + UnparsedDataType('a: byte, b: decimal( 16 , 8 ) ') + >>> UnparsedDataType("a DOUBLE, b STRING") + UnparsedDataType('a DOUBLE, b STRING') + >>> UnparsedDataType("a DOUBLE, b CHAR( 50 )") + UnparsedDataType('a DOUBLE, b CHAR( 50 )') + >>> UnparsedDataType("a DOUBLE, b VARCHAR( 50 )") + UnparsedDataType('a DOUBLE, b VARCHAR( 50 )') + >>> UnparsedDataType("a: array< short>") + UnparsedDataType('a: array< short>') + >>> UnparsedDataType(" map ") + UnparsedDataType(' map ') + """ + + def __init__(self, data_type_string: str): + self.data_type_string = data_type_string + + def simpleString(self) -> str: + return "unparsed(%s)" % repr(self.data_type_string) + + def __repr__(self) -> str: + return "UnparsedDataType(%s)" % repr(self.data_type_string) + + def jsonValue(self) -> Dict[str, Any]: + raise PySparkAssertionError( + error_class="INVALID_CALL_ON_UNRESOLVED_OBJECT", + message_parameters={"func_name": "jsonValue"}, + ) + + def needConversion(self) -> bool: + raise PySparkAssertionError( + error_class="INVALID_CALL_ON_UNRESOLVED_OBJECT", + message_parameters={"func_name": "needConversion"}, + ) + + def toInternal(self, obj: Any) -> Any: + raise PySparkAssertionError( + error_class="INVALID_CALL_ON_UNRESOLVED_OBJECT", + message_parameters={"func_name": "toInternal"}, + ) + + def fromInternal(self, obj: Any) -> Any: + raise PySparkAssertionError( + error_class="INVALID_CALL_ON_UNRESOLVED_OBJECT", + message_parameters={"func_name": "fromInternal"}, + ) + + +def pyspark_types_to_proto_types(data_type: DataType) -> pb2.DataType: + ret = pb2.DataType() + if isinstance(data_type, NullType): + ret.null.CopyFrom(pb2.DataType.NULL()) + elif isinstance(data_type, StringType): + ret.string.CopyFrom(pb2.DataType.String()) + elif isinstance(data_type, BooleanType): + ret.boolean.CopyFrom(pb2.DataType.Boolean()) + elif isinstance(data_type, BinaryType): + ret.binary.CopyFrom(pb2.DataType.Binary()) + elif isinstance(data_type, ByteType): + ret.byte.CopyFrom(pb2.DataType.Byte()) + elif isinstance(data_type, ShortType): + ret.short.CopyFrom(pb2.DataType.Short()) + elif isinstance(data_type, IntegerType): + ret.integer.CopyFrom(pb2.DataType.Integer()) + elif isinstance(data_type, LongType): + ret.long.CopyFrom(pb2.DataType.Long()) + elif isinstance(data_type, FloatType): + ret.float.CopyFrom(pb2.DataType.Float()) + elif isinstance(data_type, DoubleType): + ret.double.CopyFrom(pb2.DataType.Double()) + elif isinstance(data_type, DecimalType): + ret.decimal.scale = data_type.scale + ret.decimal.precision = data_type.precision + elif isinstance(data_type, DateType): + ret.date.CopyFrom(pb2.DataType.Date()) + elif isinstance(data_type, TimestampType): + ret.timestamp.CopyFrom(pb2.DataType.Timestamp()) + elif isinstance(data_type, TimestampNTZType): + ret.timestamp_ntz.CopyFrom(pb2.DataType.TimestampNTZ()) + elif isinstance(data_type, DayTimeIntervalType): + ret.day_time_interval.start_field = data_type.startField + ret.day_time_interval.end_field = data_type.endField + elif isinstance(data_type, YearMonthIntervalType): + ret.year_month_interval.start_field = data_type.startField + ret.year_month_interval.end_field = data_type.endField + elif isinstance(data_type, StructType): + struct = pb2.DataType.Struct() + for field in data_type.fields: + struct_field = pb2.DataType.StructField() + struct_field.name = field.name + struct_field.data_type.CopyFrom(pyspark_types_to_proto_types(field.dataType)) + struct_field.nullable = field.nullable + if field.metadata is not None and len(field.metadata) > 0: + struct_field.metadata = json.dumps(field.metadata) + struct.fields.append(struct_field) + ret.struct.CopyFrom(struct) + elif isinstance(data_type, MapType): + ret.map.key_type.CopyFrom(pyspark_types_to_proto_types(data_type.keyType)) + ret.map.value_type.CopyFrom(pyspark_types_to_proto_types(data_type.valueType)) + ret.map.value_contains_null = data_type.valueContainsNull + elif isinstance(data_type, ArrayType): + ret.array.element_type.CopyFrom(pyspark_types_to_proto_types(data_type.elementType)) + ret.array.contains_null = data_type.containsNull + elif isinstance(data_type, UserDefinedType): + json_value = data_type.jsonValue() + ret.udt.type = "udt" + if "class" in json_value: + # Scala/Java UDT + ret.udt.jvm_class = json_value["class"] + else: + # Python UDT + ret.udt.serialized_python_class = json_value["serializedClass"] + ret.udt.python_class = json_value["pyClass"] + ret.udt.sql_type.CopyFrom(pyspark_types_to_proto_types(data_type.sqlType())) + elif isinstance(data_type, UnparsedDataType): + data_type_string = data_type.data_type_string + ret.unparsed.data_type_string = data_type_string + else: + raise Exception(f"Unsupported data type {data_type}") + return ret + + +def proto_schema_to_pyspark_data_type(schema: pb2.DataType) -> DataType: + if schema.HasField("null"): + return NullType() + elif schema.HasField("boolean"): + return BooleanType() + elif schema.HasField("binary"): + return BinaryType() + elif schema.HasField("byte"): + return ByteType() + elif schema.HasField("short"): + return ShortType() + elif schema.HasField("integer"): + return IntegerType() + elif schema.HasField("long"): + return LongType() + elif schema.HasField("float"): + return FloatType() + elif schema.HasField("double"): + return DoubleType() + elif schema.HasField("decimal"): + p = schema.decimal.precision if schema.decimal.HasField("precision") else 10 + s = schema.decimal.scale if schema.decimal.HasField("scale") else 0 + return DecimalType(precision=p, scale=s) + elif schema.HasField("string"): + return StringType() + elif schema.HasField("char"): + return CharType(schema.char.length) + elif schema.HasField("var_char"): + return VarcharType(schema.var_char.length) + elif schema.HasField("date"): + return DateType() + elif schema.HasField("timestamp"): + return TimestampType() + elif schema.HasField("timestamp_ntz"): + return TimestampNTZType() + elif schema.HasField("day_time_interval"): + start: Optional[int] = ( + schema.day_time_interval.start_field + if schema.day_time_interval.HasField("start_field") + else None + ) + end: Optional[int] = ( + schema.day_time_interval.end_field + if schema.day_time_interval.HasField("end_field") + else None + ) + return DayTimeIntervalType(startField=start, endField=end) + elif schema.HasField("year_month_interval"): + start: Optional[int] = ( # type: ignore[no-redef] + schema.year_month_interval.start_field + if schema.year_month_interval.HasField("start_field") + else None + ) + end: Optional[int] = ( # type: ignore[no-redef] + schema.year_month_interval.end_field + if schema.year_month_interval.HasField("end_field") + else None + ) + return YearMonthIntervalType(startField=start, endField=end) + elif schema.HasField("array"): + return ArrayType( + proto_schema_to_pyspark_data_type(schema.array.element_type), + schema.array.contains_null, + ) + elif schema.HasField("struct"): + fields = [] + for f in schema.struct.fields: + if f.HasField("metadata"): + metadata = json.loads(f.metadata) + else: + metadata = None + fields.append( + StructField( + f.name, proto_schema_to_pyspark_data_type(f.data_type), f.nullable, metadata + ) + ) + return StructType(fields) + elif schema.HasField("map"): + return MapType( + proto_schema_to_pyspark_data_type(schema.map.key_type), + proto_schema_to_pyspark_data_type(schema.map.value_type), + schema.map.value_contains_null, + ) + elif schema.HasField("udt"): + assert schema.udt.type == "udt" + json_value = {} + if schema.udt.HasField("python_class"): + json_value["pyClass"] = schema.udt.python_class + if schema.udt.HasField("serialized_python_class"): + json_value["serializedClass"] = schema.udt.serialized_python_class + return UserDefinedType.fromJson(json_value) + else: + raise Exception(f"Unsupported data type {schema}") diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udf.py new file mode 100644 index 0000000..2636777 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udf.py @@ -0,0 +1,296 @@ +# +# 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. +# +""" +User-defined function related classes and functions +""" +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import sys +import functools +import warnings +from inspect import getfullargspec +from typing import cast, Callable, Any, TYPE_CHECKING, Optional, Union + +from pyspark.rdd import PythonEvalType +from pyspark.sql.connect.expressions import ( + ColumnReference, + PythonUDF, + CommonInlineUserDefinedFunction, +) +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.types import UnparsedDataType +from pyspark.sql.types import DataType, StringType +from pyspark.sql.udf import UDFRegistration as PySparkUDFRegistration +from pyspark.errors import PySparkTypeError, PySparkRuntimeError + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ( + ColumnOrName, + DataTypeOrString, + UserDefinedFunctionLike, + ) + from pyspark.sql.connect.session import SparkSession + from pyspark.sql.types import StringType + + +def _create_py_udf( + f: Callable[..., Any], + returnType: "DataTypeOrString", + useArrow: Optional[bool] = None, +) -> "UserDefinedFunctionLike": + if useArrow is None: + is_arrow_enabled = False + try: + from pyspark.sql.connect.session import SparkSession + + session = SparkSession.active() + is_arrow_enabled = ( + str(session.conf.get("spark.sql.execution.pythonUDF.arrow.enabled")).lower() + == "true" + ) + except PySparkRuntimeError as e: + if e.error_class == "NO_ACTIVE_OR_DEFAULT_SESSION": + pass # Just uses the default if no session found. + else: + raise e + else: + is_arrow_enabled = useArrow + + eval_type: int = PythonEvalType.SQL_BATCHED_UDF + + if is_arrow_enabled: + try: + is_func_with_args = len(getfullargspec(f).args) > 0 + except TypeError: + is_func_with_args = False + if is_func_with_args: + eval_type = PythonEvalType.SQL_ARROW_BATCHED_UDF + else: + warnings.warn( + "Arrow optimization for Python UDFs cannot be enabled.", + UserWarning, + ) + + return _create_udf(f, returnType, eval_type) + + +def _create_udf( + f: Callable[..., Any], + returnType: "DataTypeOrString", + evalType: int, + name: Optional[str] = None, + deterministic: bool = True, +) -> "UserDefinedFunctionLike": + # Set the name of the UserDefinedFunction object to be the name of function f + udf_obj = UserDefinedFunction( + f, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + ) + return udf_obj._wrapped() + + +class UserDefinedFunction: + """ + User defined function in Python + + Notes + ----- + The constructor of this class is not supposed to be directly called. + Use :meth:`pyspark.sql.functions.udf` or :meth:`pyspark.sql.functions.pandas_udf` + to create this instance. + """ + + def __init__( + self, + func: Callable[..., Any], + returnType: "DataTypeOrString" = StringType(), + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_BATCHED_UDF, + deterministic: bool = True, + ): + if not callable(func): + raise PySparkTypeError( + error_class="NOT_CALLABLE", + message_parameters={"arg_name": "func", "arg_type": type(func).__name__}, + ) + + if not isinstance(returnType, (DataType, str)): + raise PySparkTypeError( + error_class="NOT_DATATYPE_OR_STR", + message_parameters={ + "arg_name": "returnType", + "arg_type": type(returnType).__name__, + }, + ) + + if not isinstance(evalType, int): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "evalType", "arg_type": type(evalType).__name__}, + ) + + self.func = func + self.returnType: DataType = ( + UnparsedDataType(returnType) if isinstance(returnType, str) else returnType + ) + self._name = name or ( + func.__name__ if hasattr(func, "__name__") else func.__class__.__name__ + ) + self.evalType = evalType + self.deterministic = deterministic + + def _build_common_inline_user_defined_function( + self, *cols: "ColumnOrName" + ) -> CommonInlineUserDefinedFunction: + arg_cols = [ + col if isinstance(col, Column) else Column(ColumnReference(col)) for col in cols + ] + arg_exprs = [col._expr for col in arg_cols] + + py_udf = PythonUDF( + output_type=self.returnType, + eval_type=self.evalType, + func=self.func, + python_ver="%d.%d" % sys.version_info[:2], + ) + return CommonInlineUserDefinedFunction( + function_name=self._name, + function=py_udf, + deterministic=self.deterministic, + arguments=arg_exprs, + ) + + def __call__(self, *cols: "ColumnOrName") -> Column: + return Column(self._build_common_inline_user_defined_function(*cols)) + + # This function is for improving the online help system in the interactive interpreter. + # For example, the built-in help / pydoc.help. It wraps the UDF with the docstring and + # argument annotation. (See: SPARK-19161) + def _wrapped(self) -> "UserDefinedFunctionLike": + """ + Wrap this udf with a function and attach docstring from func + """ + + # It is possible for a callable instance without __name__ attribute or/and + # __module__ attribute to be wrapped here. For example, functools.partial. In this case, + # we should avoid wrapping the attributes from the wrapped function to the wrapper + # function. So, we take out these attribute names from the default names to set and + # then manually assign it after being wrapped. + assignments = tuple( + a for a in functools.WRAPPER_ASSIGNMENTS if a != "__name__" and a != "__module__" + ) + + @functools.wraps(self.func, assigned=assignments) + def wrapper(*args: "ColumnOrName") -> Column: + return self(*args) + + wrapper.__name__ = self._name + wrapper.__module__ = ( + self.func.__module__ + if hasattr(self.func, "__module__") + else self.func.__class__.__module__ + ) + + wrapper.func = self.func # type: ignore[attr-defined] + wrapper.returnType = self.returnType # type: ignore[attr-defined] + wrapper.evalType = self.evalType # type: ignore[attr-defined] + wrapper.deterministic = self.deterministic # type: ignore[attr-defined] + wrapper.asNondeterministic = functools.wraps( # type: ignore[attr-defined] + self.asNondeterministic + )(lambda: self.asNondeterministic()._wrapped()) + wrapper._unwrapped = self # type: ignore[attr-defined] + return wrapper # type: ignore[return-value] + + def asNondeterministic(self) -> "UserDefinedFunction": + """ + Updates UserDefinedFunction to nondeterministic. + + .. versionadded:: 3.4.0 + """ + self.deterministic = False + return self + + +class UDFRegistration: + """ + Wrapper for user-defined function registration. + """ + + def __init__(self, sparkSession: "SparkSession"): + self.sparkSession = sparkSession + + def register( + self, + name: str, + f: Union[Callable[..., Any], "UserDefinedFunctionLike"], + returnType: Optional["DataTypeOrString"] = None, + ) -> "UserDefinedFunctionLike": + # This is to check whether the input function is from a user-defined function or + # Python function. + if hasattr(f, "asNondeterministic"): + if returnType is not None: + raise PySparkTypeError( + error_class="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF", + message_parameters={"arg_name": "f", "return_type": str(returnType)}, + ) + f = cast("UserDefinedFunctionLike", f) + if f.evalType not in [ + PythonEvalType.SQL_BATCHED_UDF, + PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + ]: + raise PySparkTypeError( + error_class="INVALID_UDF_EVAL_TYPE", + message_parameters={ + "eval_type": "SQL_BATCHED_UDF, SQL_ARROW_BATCHED_UDF, " + "SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_PANDAS_ITER_UDF or " + "SQL_GROUPED_AGG_PANDAS_UDF" + }, + ) + self.sparkSession._client.register_udf( + f.func, f.returnType, name, f.evalType, f.deterministic + ) + return f + else: + if returnType is None: + returnType = StringType() + py_udf = _create_udf( + f, returnType=returnType, evalType=PythonEvalType.SQL_BATCHED_UDF, name=name + ) + + self.sparkSession._client.register_udf(py_udf.func, returnType, name) + return py_udf + + register.__doc__ = PySparkUDFRegistration.register.__doc__ + + def registerJavaFunction( + self, + name: str, + javaClassName: str, + returnType: Optional["DataTypeOrString"] = None, + ) -> None: + self.sparkSession._client.register_java(name, javaClassName, returnType) + + registerJavaFunction.__doc__ = PySparkUDFRegistration.registerJavaFunction.__doc__ + + def registerJavaUDAF(self, name: str, javaClassName: str) -> None: + self.sparkSession._client.register_java(name, javaClassName, aggregate=True) + + registerJavaUDAF.__doc__ = PySparkUDFRegistration.registerJavaUDAF.__doc__ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udtf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udtf.py new file mode 100644 index 0000000..edf3ad4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/udtf.py @@ -0,0 +1,200 @@ +# +# 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. +# +""" +User-defined table function related classes and functions +""" +from pyspark.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import warnings +from typing import Type, TYPE_CHECKING, Optional, Union + +from pyspark.rdd import PythonEvalType +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.expressions import ColumnReference +from pyspark.sql.connect.plan import ( + CommonInlineUserDefinedTableFunction, + PythonUDTF, +) +from pyspark.sql.connect.types import UnparsedDataType +from pyspark.sql.connect.utils import get_python_ver +from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration +from pyspark.sql.udtf import _validate_udtf_handler +from pyspark.sql.types import DataType, StructType +from pyspark.errors import PySparkRuntimeError, PySparkTypeError + + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ColumnOrName + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.session import SparkSession + + +def _create_udtf( + cls: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_TABLE_UDF, + deterministic: bool = False, +) -> "UserDefinedTableFunction": + udtf_obj = UserDefinedTableFunction( + cls, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + ) + return udtf_obj + + +def _create_py_udtf( + cls: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + deterministic: bool = False, + useArrow: Optional[bool] = None, +) -> "UserDefinedTableFunction": + if useArrow is not None: + arrow_enabled = useArrow + else: + from pyspark.sql.connect.session import SparkSession + + arrow_enabled = False + try: + session = SparkSession.active() + arrow_enabled = ( + str(session.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled")).lower() + == "true" + ) + except PySparkRuntimeError as e: + if e.error_class == "NO_ACTIVE_OR_DEFAULT_SESSION": + pass # Just uses the default if no session found. + else: + raise e + + eval_type: int = PythonEvalType.SQL_TABLE_UDF + + if arrow_enabled: + from pyspark.sql.pandas.utils import ( + require_minimum_pandas_version, + require_minimum_pyarrow_version, + ) + + try: + require_minimum_pandas_version() + require_minimum_pyarrow_version() + eval_type = PythonEvalType.SQL_ARROW_TABLE_UDF + except ImportError as e: + warnings.warn( + f"Arrow optimization for Python UDTFs cannot be enabled: {str(e)}. " + f"Falling back to using regular Python UDTFs.", + UserWarning, + ) + + return _create_udtf(cls, returnType, name, eval_type, deterministic) + + +class UserDefinedTableFunction: + """ + User defined function in Python + + Notes + ----- + The constructor of this class is not supposed to be directly called. + Use :meth:`pyspark.sql.functions.udtf` to create this instance. + """ + + def __init__( + self, + func: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_TABLE_UDF, + deterministic: bool = False, + ) -> None: + _validate_udtf_handler(func) + + self.func = func + self.returnType: DataType = ( + UnparsedDataType(returnType) if isinstance(returnType, str) else returnType + ) + self._name = name or func.__name__ + self.evalType = evalType + self.deterministic = deterministic + + def _build_common_inline_user_defined_table_function( + self, *cols: "ColumnOrName" + ) -> CommonInlineUserDefinedTableFunction: + arg_cols = [ + col if isinstance(col, Column) else Column(ColumnReference(col)) for col in cols + ] + arg_exprs = [col._expr for col in arg_cols] + + udtf = PythonUDTF( + func=self.func, + return_type=self.returnType, + eval_type=self.evalType, + python_ver=get_python_ver(), + ) + return CommonInlineUserDefinedTableFunction( + function_name=self._name, + function=udtf, + deterministic=self.deterministic, + arguments=arg_exprs, + ) + + def __call__(self, *cols: "ColumnOrName") -> "DataFrame": + from pyspark.sql.connect.session import SparkSession + from pyspark.sql.connect.dataframe import DataFrame + + session = SparkSession.active() + + plan = self._build_common_inline_user_defined_table_function(*cols) + return DataFrame.withPlan(plan, session) + + def asDeterministic(self) -> "UserDefinedTableFunction": + self.deterministic = True + return self + + +class UDTFRegistration: + """ + Wrapper for user-defined table function registration. + + .. versionadded:: 3.5.0 + """ + + def __init__(self, sparkSession: "SparkSession"): + self.sparkSession = sparkSession + + def register( + self, + name: str, + f: "UserDefinedTableFunction", + ) -> "UserDefinedTableFunction": + if f.evalType not in [PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF]: + raise PySparkTypeError( + error_class="INVALID_UDTF_EVAL_TYPE", + message_parameters={ + "name": name, + "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF", + }, + ) + + self.sparkSession._client.register_udtf( + f.func, f.returnType, name, f.evalType, f.deterministic + ) + return f + + register.__doc__ = PySparkUDTFRegistration.register.__doc__ diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/utils.py new file mode 100644 index 0000000..8872ba5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/utils.py @@ -0,0 +1,58 @@ +# +# 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 sys + +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version + + +def check_dependencies(mod_name: str) -> None: + if mod_name == "__main__": + from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + + if not should_test_connect: + print( + f"Skipping {mod_name} doctests: {connect_requirement_message}", + file=sys.stderr, + ) + sys.exit(0) + else: + require_minimum_pandas_version() + require_minimum_pyarrow_version() + require_minimum_grpc_version() + + +def require_minimum_grpc_version() -> None: + """Raise ImportError if minimum version of grpc is not installed""" + minimum_grpc_version = "1.48.1" + + from distutils.version import LooseVersion + + try: + import grpc + except ImportError as error: + raise ImportError( + "grpcio >= %s must be installed; however, " "it was not found." % minimum_grpc_version + ) from error + if LooseVersion(grpc.__version__) < LooseVersion(minimum_grpc_version): + raise ImportError( + "grpcio >= %s must be installed; however, " + "your version was %s." % (minimum_grpc_version, grpc.__version__) + ) + + +def get_python_ver() -> str: + return "%d.%d" % sys.version_info[:2] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/window.py new file mode 100644 index 0000000..ad082c6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/connect/window.py @@ -0,0 +1,263 @@ +# +# 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.sql.connect.utils import check_dependencies + +check_dependencies(__name__) + +import sys +from typing import TYPE_CHECKING, Union, Sequence, List, Optional + +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.expressions import ( + ColumnReference, + Expression, + SortOrder, +) +from pyspark.sql.connect.types import ( + JVM_LONG_MIN, + JVM_LONG_MAX, +) +from pyspark.sql.window import Window as PySparkWindow, WindowSpec as PySparkWindowSpec +from pyspark.errors import PySparkTypeError + +if TYPE_CHECKING: + from pyspark.sql.connect._typing import ColumnOrName + +__all__ = ["Window", "WindowSpec"] + + +class WindowFrame: + def __init__(self, isRowFrame: bool, start: int, end: int) -> None: + super().__init__() + + assert isinstance(isRowFrame, bool) + + assert isinstance(start, int) + + assert isinstance(end, int) + + self._isRowFrame = isRowFrame + + self._start = start + + self._end = end + + def __repr__(self) -> str: + if self._isRowFrame: + return f"WindowFrame(ROW_FRAME, {self._start}, {self._end})" + else: + return f"WindowFrame(RANGE_FRAME, {self._start}, {self._end})" + + +class WindowSpec: + def __init__( + self, + partitionSpec: Sequence[Expression], + orderSpec: Sequence[SortOrder], + frame: Optional[WindowFrame], + ) -> None: + + assert isinstance(partitionSpec, list) and all( + isinstance(p, Expression) for p in partitionSpec + ) + + assert isinstance(orderSpec, list) and all(isinstance(s, SortOrder) for s in orderSpec) + + assert frame is None or isinstance(frame, WindowFrame) + + self._partitionSpec = partitionSpec + + self._orderSpec = orderSpec + + self._frame = frame + + def partitionBy(self, *cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + _cols: List[ColumnOrName] = [] + for col in cols: + if isinstance(col, (str, Column)): + _cols.append(col) + elif isinstance(col, list): + for c in col: + if isinstance(c, (str, Column)): + _cols.append(c) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_LIST_OR_STR", + message_parameters={"arg_name": "cols", "arg_type": type(c).__name__}, + ) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_LIST_OR_STR", + message_parameters={"arg_name": "cols", "arg_type": type(col).__name__}, + ) + + newPartitionSpec: List[Expression] = [] + for c in _cols: + if isinstance(c, Column): + newPartitionSpec.append(c._expr) + else: + newPartitionSpec.append(ColumnReference(c)) + + return WindowSpec( + partitionSpec=newPartitionSpec, + orderSpec=self._orderSpec, + frame=self._frame, + ) + + def orderBy(self, *cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + _cols: List[ColumnOrName] = [] + for col in cols: + if isinstance(col, (str, Column)): + _cols.append(col) + elif isinstance(col, list): + for c in col: + if isinstance(c, (str, Column)): + _cols.append(c) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_LIST_OR_STR", + message_parameters={"arg_name": "cols", "arg_type": type(c).__name__}, + ) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_LIST_OR_STR", + message_parameters={"arg_name": "cols", "arg_type": type(col).__name__}, + ) + + newOrderSpec: List[SortOrder] = [] + for c in _cols: + if isinstance(c, Column): + if isinstance(c._expr, SortOrder): + newOrderSpec.append(c._expr) + else: + newOrderSpec.append(SortOrder(c._expr)) + else: + newOrderSpec.append(SortOrder(ColumnReference(c))) + + return WindowSpec( + partitionSpec=self._partitionSpec, + orderSpec=newOrderSpec, + frame=self._frame, + ) + + def rowsBetween(self, start: int, end: int) -> "WindowSpec": + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + + return WindowSpec( + partitionSpec=self._partitionSpec, + orderSpec=self._orderSpec, + frame=WindowFrame(isRowFrame=True, start=start, end=end), + ) + + def rangeBetween(self, start: int, end: int) -> "WindowSpec": + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + + return WindowSpec( + partitionSpec=self._partitionSpec, + orderSpec=self._orderSpec, + frame=WindowFrame(isRowFrame=False, start=start, end=end), + ) + + def __repr__(self) -> str: + strs: List[str] = [] + if len(self._partitionSpec) > 0: + str_p = ", ".join([str(p) for p in self._partitionSpec]) + strs.append(f"PartitionBy({str_p})") + if len(self._orderSpec) > 0: + str_s = ", ".join([str(s) for s in self._orderSpec]) + strs.append(f"OrderBy({str_s})") + if self._frame is not None: + strs.append(str(self._frame)) + return "WindowSpec(" + ", ".join(strs) + ")" + + +WindowSpec.rangeBetween.__doc__ = PySparkWindowSpec.rangeBetween.__doc__ +WindowSpec.rowsBetween.__doc__ = PySparkWindowSpec.rowsBetween.__doc__ +WindowSpec.orderBy.__doc__ = PySparkWindowSpec.orderBy.__doc__ +WindowSpec.partitionBy.__doc__ = PySparkWindowSpec.partitionBy.__doc__ +WindowSpec.__doc__ = PySparkWindowSpec.__doc__ + + +class Window: + _PRECEDING_THRESHOLD = max(-sys.maxsize, JVM_LONG_MIN) + _FOLLOWING_THRESHOLD = min(sys.maxsize, JVM_LONG_MAX) + + unboundedPreceding: int = JVM_LONG_MIN + + unboundedFollowing: int = JVM_LONG_MAX + + currentRow: int = 0 + + _spec = WindowSpec(partitionSpec=[], orderSpec=[], frame=None) + + @staticmethod + def partitionBy(*cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + return Window._spec.partitionBy(*cols) + + @staticmethod + def orderBy(*cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + return Window._spec.orderBy(*cols) + + @staticmethod + def rowsBetween(start: int, end: int) -> "WindowSpec": + return Window._spec.rowsBetween(start, end) + + @staticmethod + def rangeBetween(start: int, end: int) -> "WindowSpec": + return Window._spec.rangeBetween(start, end) + + +Window.orderBy.__doc__ = PySparkWindow.orderBy.__doc__ +Window.rowsBetween.__doc__ = PySparkWindow.rowsBetween.__doc__ +Window.rangeBetween.__doc__ = PySparkWindow.rangeBetween.__doc__ +Window.partitionBy.__doc__ = PySparkWindow.partitionBy.__doc__ +Window.__doc__ = PySparkWindow.__doc__ + + +def _test() -> None: + import sys + import doctest + from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.window + + globs = pyspark.sql.connect.window.__dict__.copy() + globs["spark"] = ( + PySparkSession.builder.appName("sql.connect.window tests").remote("local[4]").getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.connect.window, + globs=globs, + optionflags=doctest.ELLIPSIS + | doctest.NORMALIZE_WHITESPACE + | doctest.IGNORE_EXCEPTION_DETAIL, + ) + + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/context.py new file mode 100644 index 0000000..817c3b9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/context.py @@ -0,0 +1,818 @@ +# +# 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 sys +import warnings +from typing import ( + Optional, + Union, + Callable, + Any, + Iterable, + List, + Tuple, + overload, + Type, + ClassVar, + TYPE_CHECKING, + cast, +) + +from py4j.java_gateway import JavaObject + +from pyspark import since, _NoValue +from pyspark._globals import _NoValueType +from pyspark.sql.session import _monkey_patch_RDD, SparkSession +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.readwriter import DataFrameReader +from pyspark.sql.streaming import DataStreamReader +from pyspark.sql.udf import UDFRegistration # noqa: F401 +from pyspark.sql.udtf import UDTFRegistration +from pyspark.errors.exceptions.captured import install_exception_handler +from pyspark.context import SparkContext +from pyspark.rdd import RDD +from pyspark.sql.types import AtomicType, DataType, StructType +from pyspark.sql.streaming import StreamingQueryManager + +if TYPE_CHECKING: + from pyspark.sql._typing import ( + AtomicValue, + RowLike, + UserDefinedFunctionLike, + ) + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike + +__all__ = ["SQLContext", "HiveContext"] + + +# TODO: ignore[attr-defined] will be removed, once SparkContext is inlined +class SQLContext: + """The entry point for working with structured data (rows and columns) in Spark, in Spark 1.x. + + As of Spark 2.0, this is replaced by :class:`SparkSession`. However, we are keeping the class + here for backward compatibility. + + A SQLContext can be used to create :class:`DataFrame`, register :class:`DataFrame` as + tables, execute SQL over tables, cache tables, and read parquet files. + + .. deprecated:: 3.0.0 + Use :func:`SparkSession.builder.getOrCreate()` instead. + + Parameters + ---------- + sparkContext : :class:`SparkContext` + The :class:`SparkContext` backing this SQLContext. + sparkSession : :class:`SparkSession` + The :class:`SparkSession` around which this SQLContext wraps. + jsqlContext : optional + An optional JVM Scala SQLContext. If set, we do not instantiate a new + SQLContext in the JVM, instead we make all calls to this object. + This is only for internal. + + Examples + -------- + >>> from datetime import datetime + >>> from pyspark.sql import Row + >>> sqlContext = SQLContext(sc) + >>> allTypes = sc.parallelize([Row(i=1, s="string", d=1.0, l=1, + ... b=True, list=[1, 2, 3], dict={"s": 0}, row=Row(a=1), + ... time=datetime(2014, 8, 1, 14, 1, 5))]) + >>> df = allTypes.toDF() + >>> df.createOrReplaceTempView("allTypes") + >>> sqlContext.sql('select i+1, d+1, not b, list[1], dict["s"], time, row.a ' + ... 'from allTypes where b and i > 0').collect() + [Row((i + 1)=2, (d + 1)=2.0, (NOT b)=False, list[1]=2, \ + dict[s]=0, time=datetime.datetime(2014, 8, 1, 14, 1, 5), a=1)] + >>> df.rdd.map(lambda x: (x.i, x.s, x.d, x.l, x.b, x.time, x.row.a, x.list)).collect() + [(1, 'string', 1.0, 1, True, datetime.datetime(2014, 8, 1, 14, 1, 5), 1, [1, 2, 3])] + """ + + _instantiatedContext: ClassVar[Optional["SQLContext"]] = None + + def __init__( + self, + sparkContext: SparkContext, + sparkSession: Optional[SparkSession] = None, + jsqlContext: Optional[JavaObject] = None, + ): + if sparkSession is None: + warnings.warn( + "Deprecated in 3.0.0. Use SparkSession.builder.getOrCreate() instead.", + FutureWarning, + ) + + self._sc = sparkContext + self._jsc = self._sc._jsc + self._jvm = self._sc._jvm + if sparkSession is None: + sparkSession = SparkSession._getActiveSessionOrCreate() + if jsqlContext is None: + jsqlContext = sparkSession._jsparkSession.sqlContext() + self.sparkSession = sparkSession + self._jsqlContext = jsqlContext + _monkey_patch_RDD(self.sparkSession) + install_exception_handler() + if ( + SQLContext._instantiatedContext is None + or SQLContext._instantiatedContext._sc._jsc is None + ): + SQLContext._instantiatedContext = self + + @property + def _ssql_ctx(self) -> JavaObject: + """Accessor for the JVM Spark SQL context. + + Subclasses can override this property to provide their own + JVM Contexts. + """ + return self._jsqlContext + + @classmethod + def getOrCreate(cls: Type["SQLContext"], sc: SparkContext) -> "SQLContext": + """ + Get the existing SQLContext or create a new one with given SparkContext. + + .. versionadded:: 1.6.0 + + .. deprecated:: 3.0.0 + Use :func:`SparkSession.builder.getOrCreate()` instead. + + Parameters + ---------- + sc : :class:`SparkContext` + """ + warnings.warn( + "Deprecated in 3.0.0. Use SparkSession.builder.getOrCreate() instead.", + FutureWarning, + ) + return cls._get_or_create(sc) + + @classmethod + def _get_or_create( + cls: Type["SQLContext"], sc: SparkContext, **static_conf: Any + ) -> "SQLContext": + + if ( + cls._instantiatedContext is None + or SQLContext._instantiatedContext._sc._jsc is None # type: ignore[union-attr] + ): + assert sc._jvm is not None + # There can be only one running Spark context. That will automatically + # be used in the Spark session internally. + session = SparkSession._getActiveSessionOrCreate(**static_conf) + cls(sc, session, session._jsparkSession.sqlContext()) + return cast(SQLContext, cls._instantiatedContext) + + def newSession(self) -> "SQLContext": + """ + Returns a new SQLContext as new session, that has separate SQLConf, + registered temporary views and UDFs, but shared SparkContext and + table cache. + + .. versionadded:: 1.6.0 + """ + return self.__class__(self._sc, self.sparkSession.newSession()) + + def setConf(self, key: str, value: Union[bool, int, str]) -> None: + """Sets the given Spark SQL configuration property. + + .. versionadded:: 1.3.0 + """ + self.sparkSession.conf.set(key, value) + + def getConf( + self, key: str, defaultValue: Union[Optional[str], _NoValueType] = _NoValue + ) -> Optional[str]: + """Returns the value of Spark SQL configuration property for the given key. + + If the key is not set and defaultValue is set, return + defaultValue. If the key is not set and defaultValue is not set, return + the system default value. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> sqlContext.getConf("spark.sql.shuffle.partitions") + '200' + >>> sqlContext.getConf("spark.sql.shuffle.partitions", "10") + '10' + >>> sqlContext.setConf("spark.sql.shuffle.partitions", "50") + >>> sqlContext.getConf("spark.sql.shuffle.partitions", "10") + '50' + """ + return self.sparkSession.conf.get(key, defaultValue) + + @property + def udf(self) -> UDFRegistration: + """Returns a :class:`UDFRegistration` for UDF registration. + + .. versionadded:: 1.3.1 + + Returns + ------- + :class:`UDFRegistration` + """ + return self.sparkSession.udf + + @property + def udtf(self) -> UDTFRegistration: + """Returns a :class:`UDTFRegistration` for UDTF registration. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`UDTFRegistration` + """ + return self.sparkSession.udtf + + def range( + self, + start: int, + end: Optional[int] = None, + step: int = 1, + numPartitions: Optional[int] = None, + ) -> DataFrame: + """ + Create a :class:`DataFrame` with single :class:`pyspark.sql.types.LongType` column named + ``id``, containing elements in a range from ``start`` to ``end`` (exclusive) with + step value ``step``. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + start : int + the start value + end : int, optional + the end value (exclusive) + step : int, optional + the incremental step (default: 1) + numPartitions : int, optional + the number of partitions of the DataFrame + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> sqlContext.range(1, 7, 2).collect() + [Row(id=1), Row(id=3), Row(id=5)] + + If only one argument is specified, it will be used as the end value. + + >>> sqlContext.range(3).collect() + [Row(id=0), Row(id=1), Row(id=2)] + """ + return self.sparkSession.range(start, end, step, numPartitions) + + def registerFunction( + self, name: str, f: Callable[..., Any], returnType: Optional[DataType] = None + ) -> "UserDefinedFunctionLike": + """An alias for :func:`spark.udf.register`. + See :meth:`pyspark.sql.UDFRegistration.register`. + + .. versionadded:: 1.2.0 + + .. deprecated:: 2.3.0 + Use :func:`spark.udf.register` instead. + """ + warnings.warn("Deprecated in 2.3.0. Use spark.udf.register instead.", FutureWarning) + return self.sparkSession.udf.register(name, f, returnType) + + def registerJavaFunction( + self, name: str, javaClassName: str, returnType: Optional[DataType] = None + ) -> None: + """An alias for :func:`spark.udf.registerJavaFunction`. + See :meth:`pyspark.sql.UDFRegistration.registerJavaFunction`. + + .. versionadded:: 2.1.0 + + .. deprecated:: 2.3.0 + Use :func:`spark.udf.registerJavaFunction` instead. + """ + warnings.warn( + "Deprecated in 2.3.0. Use spark.udf.registerJavaFunction instead.", + FutureWarning, + ) + return self.sparkSession.udf.registerJavaFunction(name, javaClassName, returnType) + + # TODO(andrew): delete this once we refactor things to take in SparkSession + def _inferSchema(self, rdd: RDD, samplingRatio: Optional[float] = None) -> StructType: + """ + Infer schema from an RDD of Row or tuple. + + Parameters + ---------- + rdd : :class:`RDD` + an RDD of Row or tuple + samplingRatio : float, optional + sampling ratio, or no sampling (default) + + Returns + ------- + :class:`pyspark.sql.types.StructType` + """ + return self.sparkSession._inferSchema(rdd, samplingRatio) + + @overload + def createDataFrame( + self, + data: Union["RDD[RowLike]", Iterable["RowLike"]], + schema: Union[List[str], Tuple[str, ...]] = ..., + samplingRatio: Optional[float] = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: Union["RDD[RowLike]", Iterable["RowLike"]], + schema: Union[StructType, str], + *, + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: Union[ + "RDD[AtomicValue]", + Iterable["AtomicValue"], + ], + schema: Union[AtomicType, str], + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, data: "PandasDataFrameLike", samplingRatio: Optional[float] = ... + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: "PandasDataFrameLike", + schema: Union[StructType, str], + verifySchema: bool = ..., + ) -> DataFrame: + ... + + def createDataFrame( # type: ignore[misc] + self, + data: Union[RDD[Any], Iterable[Any], "PandasDataFrameLike"], + schema: Optional[Union[AtomicType, StructType, str]] = None, + samplingRatio: Optional[float] = None, + verifySchema: bool = True, + ) -> DataFrame: + """ + Creates a :class:`DataFrame` from an :class:`RDD`, a list or a :class:`pandas.DataFrame`. + + When ``schema`` is a list of column names, the type of each column + will be inferred from ``data``. + + When ``schema`` is ``None``, it will try to infer the schema (column names and types) + from ``data``, which should be an RDD of :class:`Row`, + or :class:`namedtuple`, or :class:`dict`. + + When ``schema`` is :class:`pyspark.sql.types.DataType` or a datatype string it must match + the real data, or an exception will be thrown at runtime. If the given schema is not + :class:`pyspark.sql.types.StructType`, it will be wrapped into a + :class:`pyspark.sql.types.StructType` as its only field, and the field name will be "value", + each record will also be wrapped into a tuple, which can be converted to row later. + + If schema inference is needed, ``samplingRatio`` is used to determine the ratio of + rows used for schema inference. The first row will be used if ``samplingRatio`` is ``None``. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 2.0.0 + The ``schema`` parameter can be a :class:`pyspark.sql.types.DataType` or a + datatype string after 2.0. + If it's not a :class:`pyspark.sql.types.StructType`, it will be wrapped into a + :class:`pyspark.sql.types.StructType` and each record will also be wrapped into a tuple. + + .. versionchanged:: 2.1.0 + Added verifySchema. + + Parameters + ---------- + data : :class:`RDD` or iterable + an RDD of any kind of SQL data representation (:class:`Row`, + :class:`tuple`, ``int``, ``boolean``, etc.), or :class:`list`, or + :class:`pandas.DataFrame`. + schema : :class:`pyspark.sql.types.DataType`, str or list, optional + a :class:`pyspark.sql.types.DataType` or a datatype string or a list of + column names, default is None. The data type string format equals to + :class:`pyspark.sql.types.DataType.simpleString`, except that top level struct type can + omit the ``struct<>``. + samplingRatio : float, optional + the sample ratio of rows used for inferring + verifySchema : bool, optional + verify data types of every row against schema. Enabled by default. + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> l = [('Alice', 1)] + >>> sqlContext.createDataFrame(l).collect() + [Row(_1='Alice', _2=1)] + >>> sqlContext.createDataFrame(l, ['name', 'age']).collect() + [Row(name='Alice', age=1)] + + >>> d = [{'name': 'Alice', 'age': 1}] + >>> sqlContext.createDataFrame(d).collect() + [Row(age=1, name='Alice')] + + >>> rdd = sc.parallelize(l) + >>> sqlContext.createDataFrame(rdd).collect() + [Row(_1='Alice', _2=1)] + >>> df = sqlContext.createDataFrame(rdd, ['name', 'age']) + >>> df.collect() + [Row(name='Alice', age=1)] + + >>> from pyspark.sql import Row + >>> Person = Row('name', 'age') + >>> person = rdd.map(lambda r: Person(*r)) + >>> df2 = sqlContext.createDataFrame(person) + >>> df2.collect() + [Row(name='Alice', age=1)] + + >>> from pyspark.sql.types import * + >>> schema = StructType([ + ... StructField("name", StringType(), True), + ... StructField("age", IntegerType(), True)]) + >>> df3 = sqlContext.createDataFrame(rdd, schema) + >>> df3.collect() + [Row(name='Alice', age=1)] + + >>> sqlContext.createDataFrame(df.toPandas()).collect() # doctest: +SKIP + [Row(name='Alice', age=1)] + >>> sqlContext.createDataFrame(pandas.DataFrame([[1, 2]])).collect() # doctest: +SKIP + [Row(0=1, 1=2)] + + >>> sqlContext.createDataFrame(rdd, "a: string, b: int").collect() + [Row(a='Alice', b=1)] + >>> rdd = rdd.map(lambda row: row[1]) + >>> sqlContext.createDataFrame(rdd, "int").collect() + [Row(value=1)] + >>> sqlContext.createDataFrame(rdd, "boolean").collect() # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + Py4JJavaError: ... + """ + return self.sparkSession.createDataFrame( # type: ignore[call-overload] + data, schema, samplingRatio, verifySchema + ) + + def registerDataFrameAsTable(self, df: DataFrame, tableName: str) -> None: + """Registers the given :class:`DataFrame` as a temporary table in the catalog. + + Temporary tables exist only during the lifetime of this instance of :class:`SQLContext`. + + .. versionadded:: 1.3.0 + + Examples + -------- + >>> sqlContext.registerDataFrameAsTable(df, "table1") + """ + df.createOrReplaceTempView(tableName) + + def dropTempTable(self, tableName: str) -> None: + """Remove the temporary table from catalog. + + .. versionadded:: 1.6.0 + + Examples + -------- + >>> sqlContext.registerDataFrameAsTable(df, "table1") + >>> sqlContext.dropTempTable("table1") + """ + self.sparkSession.catalog.dropTempView(tableName) + + def createExternalTable( + self, + tableName: str, + path: Optional[str] = None, + source: Optional[str] = None, + schema: Optional[StructType] = None, + **options: str, + ) -> DataFrame: + """Creates an external table based on the dataset in a data source. + + It returns the DataFrame associated with the external table. + + The data source is specified by the ``source`` and a set of ``options``. + If ``source`` is not specified, the default data source configured by + ``spark.sql.sources.default`` will be used. + + Optionally, a schema can be provided as the schema of the returned :class:`DataFrame` and + created external table. + + .. versionadded:: 1.3.0 + + Returns + ------- + :class:`DataFrame` + """ + return self.sparkSession.catalog.createExternalTable( + tableName, path, source, schema, **options + ) + + def sql(self, sqlQuery: str) -> DataFrame: + """Returns a :class:`DataFrame` representing the result of the given query. + + .. versionadded:: 1.0.0 + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> sqlContext.registerDataFrameAsTable(df, "table1") + >>> df2 = sqlContext.sql("SELECT field1 AS f1, field2 as f2 from table1") + >>> df2.collect() + [Row(f1=1, f2='row1'), Row(f1=2, f2='row2'), Row(f1=3, f2='row3')] + """ + return self.sparkSession.sql(sqlQuery) + + def table(self, tableName: str) -> DataFrame: + """Returns the specified table or view as a :class:`DataFrame`. + + .. versionadded:: 1.0.0 + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> sqlContext.registerDataFrameAsTable(df, "table1") + >>> df2 = sqlContext.table("table1") + >>> sorted(df.collect()) == sorted(df2.collect()) + True + """ + return self.sparkSession.table(tableName) + + def tables(self, dbName: Optional[str] = None) -> DataFrame: + """Returns a :class:`DataFrame` containing names of tables in the given database. + + If ``dbName`` is not specified, the current database will be used. + + The returned DataFrame has two columns: ``tableName`` and ``isTemporary`` + (a column with :class:`BooleanType` indicating if a table is a temporary one or not). + + .. versionadded:: 1.3.0 + + Parameters + ---------- + dbName: str, optional + name of the database to use. + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> sqlContext.registerDataFrameAsTable(df, "table1") + >>> df2 = sqlContext.tables() + >>> df2.filter("tableName = 'table1'").first() + Row(namespace='', tableName='table1', isTemporary=True) + """ + if dbName is None: + return DataFrame(self._ssql_ctx.tables(), self.sparkSession) + else: + return DataFrame(self._ssql_ctx.tables(dbName), self.sparkSession) + + def tableNames(self, dbName: Optional[str] = None) -> List[str]: + """Returns a list of names of tables in the database ``dbName``. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + dbName: str + name of the database to use. Default to the current database. + + Returns + ------- + list + list of table names, in string + + >>> sqlContext.registerDataFrameAsTable(df, "table1") + >>> "table1" in sqlContext.tableNames() + True + >>> "table1" in sqlContext.tableNames("default") + True + """ + if dbName is None: + return [name for name in self._ssql_ctx.tableNames()] + else: + return [name for name in self._ssql_ctx.tableNames(dbName)] + + @since(1.0) + def cacheTable(self, tableName: str) -> None: + """Caches the specified table in-memory.""" + self._ssql_ctx.cacheTable(tableName) + + @since(1.0) + def uncacheTable(self, tableName: str) -> None: + """Removes the specified table from the in-memory cache.""" + self._ssql_ctx.uncacheTable(tableName) + + @since(1.3) + def clearCache(self) -> None: + """Removes all cached tables from the in-memory cache.""" + self._ssql_ctx.clearCache() + + @property + def read(self) -> DataFrameReader: + """ + Returns a :class:`DataFrameReader` that can be used to read data + in as a :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + Returns + ------- + :class:`DataFrameReader` + """ + return DataFrameReader(self.sparkSession) + + @property + def readStream(self) -> DataStreamReader: + """ + Returns a :class:`DataStreamReader` that can be used to read data streams + as a streaming :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + Notes + ----- + This API is evolving. + + Returns + ------- + :class:`DataStreamReader` + + >>> text_sdf = sqlContext.readStream.text(tempfile.mkdtemp()) + >>> text_sdf.isStreaming + True + """ + return DataStreamReader(self.sparkSession) + + @property + def streams(self) -> StreamingQueryManager: + """Returns a :class:`StreamingQueryManager` that allows managing all the + :class:`StreamingQuery` StreamingQueries active on `this` context. + + .. versionadded:: 2.0.0 + + Notes + ----- + This API is evolving. + """ + from pyspark.sql.streaming import StreamingQueryManager + + return StreamingQueryManager(self._ssql_ctx.streams()) + + +# TODO: ignore[attr-defined] will be removed, once SparkContext is inlined +class HiveContext(SQLContext): + """A variant of Spark SQL that integrates with data stored in Hive. + + Configuration for Hive is read from ``hive-site.xml`` on the classpath. + It supports running both SQL and HiveQL commands. + + .. deprecated:: 2.0.0 + Use SparkSession.builder.enableHiveSupport().getOrCreate(). + + Parameters + ---------- + sparkContext : :class:`SparkContext` + The SparkContext to wrap. + jhiveContext : optional + An optional JVM Scala HiveContext. If set, we do not instantiate a new + :class:`HiveContext` in the JVM, instead we make all calls to this object. + This is only for internal use. + + """ + + _static_conf = {"spark.sql.catalogImplementation": "hive"} + + def __init__( + self, + sparkContext: SparkContext, + sparkSession: Optional[SparkSession] = None, + jhiveContext: Optional[JavaObject] = None, + ): + warnings.warn( + "HiveContext is deprecated in Spark 2.0.0. Please use " + + "SparkSession.builder.enableHiveSupport().getOrCreate() instead.", + FutureWarning, + ) + static_conf = {} + if jhiveContext is None: + static_conf = HiveContext._static_conf + # There can be only one running Spark context. That will automatically + # be used in the Spark session internally. + if sparkSession is not None: + sparkSession = SparkSession._getActiveSessionOrCreate(**static_conf) + SQLContext.__init__(self, sparkContext, sparkSession, jhiveContext) + + @classmethod + def _get_or_create( + cls: Type["SQLContext"], sc: SparkContext, **static_conf: Any + ) -> "SQLContext": + return SQLContext._get_or_create(sc, **HiveContext._static_conf) + + @classmethod + def _createForTesting(cls, sparkContext: SparkContext) -> "HiveContext": + """(Internal use only) Create a new HiveContext for testing. + + All test code that touches HiveContext *must* go through this method. Otherwise, + you may end up launching multiple derby instances and encounter with incredibly + confusing error messages. + """ + jsc = sparkContext._jsc.sc() + assert sparkContext._jvm is not None + jtestHive = sparkContext._jvm.org.apache.spark.sql.hive.test.TestHiveContext(jsc, False) + return cls(sparkContext, jtestHive) + + def refreshTable(self, tableName: str) -> None: + """Invalidate and refresh all the cached metadata of the given + table. For performance reasons, Spark SQL or the external data source + library it uses might cache certain metadata about a table, such as the + location of blocks. When those change outside of Spark SQL, users should + call this function to invalidate the cache. + """ + self._ssql_ctx.refreshTable(tableName) + + +def _test() -> None: + import os + import doctest + import tempfile + from pyspark.context import SparkContext + from pyspark.sql import Row, SQLContext + import pyspark.sql.context + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.context.__dict__.copy() + sc = SparkContext("local[4]", "PythonTest") + globs["tempfile"] = tempfile + globs["os"] = os + globs["sc"] = sc + globs["sqlContext"] = SQLContext(sc) + globs["rdd"] = rdd = sc.parallelize( + [ + Row(field1=1, field2="row1"), + Row(field1=2, field2="row2"), + Row(field1=3, field2="row3"), + ] + ) + globs["df"] = rdd.toDF() + jsonStrings = [ + '{"field1": 1, "field2": "row1", "field3":{"field4":11}}', + '{"field1" : 2, "field3":{"field4":22, "field5": [10, 11]},"field6":[{"field7": "row2"}]}', + '{"field1" : null, "field2": "row3", "field3":{"field4":33, "field5": []}}', + ] + globs["jsonStrings"] = jsonStrings + globs["json"] = sc.parallelize(jsonStrings) + (failure_count, test_count) = doctest.testmod( + pyspark.sql.context, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + globs["sc"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/dataframe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/dataframe.py new file mode 100644 index 0000000..30ed73d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/dataframe.py @@ -0,0 +1,5967 @@ +# +# 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 json +import os +import sys +import random +import warnings +from collections.abc import Iterable +from functools import reduce +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, + overload, + TYPE_CHECKING, +) + +from py4j.java_gateway import JavaObject, JVMView + +from pyspark import copy_func, _NoValue +from pyspark._globals import _NoValueType +from pyspark.context import SparkContext +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.rdd import ( + RDD, + _load_from_socket, + _local_iterator_from_socket, +) +from pyspark.serializers import BatchedSerializer, CPickleSerializer, UTF8Deserializer +from pyspark.storagelevel import StorageLevel +from pyspark.traceback_utils import SCCallSiteSync +from pyspark.sql.column import Column, _to_seq, _to_list, _to_java_column +from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.streaming import DataStreamWriter +from pyspark.sql.types import ( + StructType, + Row, + _parse_datatype_json_string, +) +from pyspark.sql.utils import get_active_spark_context +from pyspark.sql.pandas.conversion import PandasConversionMixin +from pyspark.sql.pandas.map_ops import PandasMapOpsMixin + +if TYPE_CHECKING: + from pyspark._typing import PrimitiveType + from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame + from pyspark.sql._typing import ColumnOrName, LiteralType, OptionalPrimitiveType + from pyspark.sql.context import SQLContext + from pyspark.sql.session import SparkSession + from pyspark.sql.group import GroupedData + from pyspark.sql.observation import Observation + + +__all__ = ["DataFrame", "DataFrameNaFunctions", "DataFrameStatFunctions"] + + +class DataFrame(PandasMapOpsMixin, PandasConversionMixin): + """A distributed collection of data grouped into named columns. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + A :class:`DataFrame` is equivalent to a relational table in Spark SQL, + and can be created using various functions in :class:`SparkSession`: + + >>> people = spark.createDataFrame([ + ... {"deptId": 1, "age": 40, "name": "Hyukjin Kwon", "gender": "M", "salary": 50}, + ... {"deptId": 1, "age": 50, "name": "Takuya Ueshin", "gender": "M", "salary": 100}, + ... {"deptId": 2, "age": 60, "name": "Xinrong Meng", "gender": "F", "salary": 150}, + ... {"deptId": 3, "age": 20, "name": "Haejoon Lee", "gender": "M", "salary": 200} + ... ]) + + Once created, it can be manipulated using the various domain-specific-language + (DSL) functions defined in: :class:`DataFrame`, :class:`Column`. + + To select a column from the :class:`DataFrame`, use the apply method: + + >>> age_col = people.age + + A more concrete example: + + >>> # To create DataFrame using SparkSession + ... department = spark.createDataFrame([ + ... {"id": 1, "name": "PySpark"}, + ... {"id": 2, "name": "ML"}, + ... {"id": 3, "name": "Spark SQL"} + ... ]) + + >>> people.filter(people.age > 30).join( + ... department, people.deptId == department.id).groupBy( + ... department.name, "gender").agg({"salary": "avg", "age": "max"}).show() + +-------+------+-----------+--------+ + | name|gender|avg(salary)|max(age)| + +-------+------+-----------+--------+ + | ML| F| 150.0| 60| + |PySpark| M| 75.0| 50| + +-------+------+-----------+--------+ + + Notes + ----- + A DataFrame should only be created as described above. It should not be directly + created via using the constructor. + """ + + def __init__( + self, + jdf: JavaObject, + sql_ctx: Union["SQLContext", "SparkSession"], + ): + from pyspark.sql.context import SQLContext + + self._sql_ctx: Optional["SQLContext"] = None + + if isinstance(sql_ctx, SQLContext): + assert not os.environ.get("SPARK_TESTING") # Sanity check for our internal usage. + assert isinstance(sql_ctx, SQLContext) + # We should remove this if-else branch in the future release, and rename + # sql_ctx to session in the constructor. This is an internal code path but + # was kept with a warning because it's used intensively by third-party libraries. + warnings.warn("DataFrame constructor is internal. Do not directly use it.") + self._sql_ctx = sql_ctx + session = sql_ctx.sparkSession + else: + session = sql_ctx + self._session: "SparkSession" = session + + self._sc: SparkContext = sql_ctx._sc + self._jdf: JavaObject = jdf + self.is_cached = False + # initialized lazily + self._schema: Optional[StructType] = None + self._lazy_rdd: Optional[RDD[Row]] = None + # Check whether _repr_html is supported or not, we use it to avoid calling _jdf twice + # by __repr__ and _repr_html_ while eager evaluation opens. + self._support_repr_html = False + + @property + def sql_ctx(self) -> "SQLContext": + from pyspark.sql.context import SQLContext + + warnings.warn( + "DataFrame.sql_ctx is an internal property, and will be removed " + "in future releases. Use DataFrame.sparkSession instead." + ) + if self._sql_ctx is None: + self._sql_ctx = SQLContext._get_or_create(self._sc) + return self._sql_ctx + + @property + def sparkSession(self) -> "SparkSession": + """Returns Spark session that created this :class:`DataFrame`. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`SparkSession` + + Examples + -------- + >>> df = spark.range(1) + >>> type(df.sparkSession) + + """ + return self._session + + @property + def rdd(self) -> "RDD[Row]": + """Returns the content as an :class:`pyspark.RDD` of :class:`Row`. + + .. versionadded:: 1.3.0 + + Returns + ------- + :class:`RDD` + + Examples + -------- + >>> df = spark.range(1) + >>> type(df.rdd) + + """ + if self._lazy_rdd is None: + jrdd = self._jdf.javaToPython() + self._lazy_rdd = RDD( + jrdd, self.sparkSession._sc, BatchedSerializer(CPickleSerializer()) + ) + return self._lazy_rdd + + @property + def na(self) -> "DataFrameNaFunctions": + """Returns a :class:`DataFrameNaFunctions` for handling missing values. + + .. versionadded:: 1.3.1 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrameNaFunctions` + + Examples + -------- + >>> df = spark.sql("SELECT 1 AS c1, int(NULL) AS c2") + >>> type(df.na) + + + Replace the missing values as 2. + + >>> df.na.fill(2).show() + +---+---+ + | c1| c2| + +---+---+ + | 1| 2| + +---+---+ + """ + return DataFrameNaFunctions(self) + + @property + def stat(self) -> "DataFrameStatFunctions": + """Returns a :class:`DataFrameStatFunctions` for statistic functions. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrameStatFunctions` + + Examples + -------- + >>> import pyspark.sql.functions as f + >>> df = spark.range(3).withColumn("c", f.expr("id + 1")) + >>> type(df.stat) + + >>> df.stat.corr("id", "c") + 1.0 + """ + return DataFrameStatFunctions(self) + + def toJSON(self, use_unicode: bool = True) -> RDD[str]: + """Converts a :class:`DataFrame` into a :class:`RDD` of string. + + Each row is turned into a JSON document as one element in the returned RDD. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + use_unicode : bool, optional, default True + Whether to convert to unicode or not. + + Returns + ------- + :class:`RDD` + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.toJSON().first() + '{"age":2,"name":"Alice"}' + """ + rdd = self._jdf.toJSON() + return RDD(rdd.toJavaRDD(), self._sc, UTF8Deserializer(use_unicode)) + + def registerTempTable(self, name: str) -> None: + """Registers this :class:`DataFrame` as a temporary table using the given name. + + The lifetime of this temporary table is tied to the :class:`SparkSession` + that was used to create this :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 2.0.0 + Use :meth:`DataFrame.createOrReplaceTempView` instead. + + Parameters + ---------- + name : str + Name of the temporary table to register. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.registerTempTable("people") + >>> df2 = spark.sql("SELECT * FROM people") + >>> sorted(df.collect()) == sorted(df2.collect()) + True + >>> spark.catalog.dropTempView("people") + True + + """ + warnings.warn("Deprecated in 2.0, use createOrReplaceTempView instead.", FutureWarning) + self._jdf.createOrReplaceTempView(name) + + def createTempView(self, name: str) -> None: + """Creates a local temporary view with this :class:`DataFrame`. + + The lifetime of this temporary table is tied to the :class:`SparkSession` + that was used to create this :class:`DataFrame`. + throws :class:`TempTableAlreadyExistsException`, if the view name already exists in the + catalog. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + Name of the view. + + Examples + -------- + Create a local temporary view. + + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.createTempView("people") + >>> df2 = spark.sql("SELECT * FROM people") + >>> sorted(df.collect()) == sorted(df2.collect()) + True + + Throw an exception if the table already exists. + + >>> df.createTempView("people") # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + AnalysisException: "Temporary table 'people' already exists;" + >>> spark.catalog.dropTempView("people") + True + + """ + self._jdf.createTempView(name) + + def createOrReplaceTempView(self, name: str) -> None: + """Creates or replaces a local temporary view with this :class:`DataFrame`. + + The lifetime of this temporary table is tied to the :class:`SparkSession` + that was used to create this :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + Name of the view. + + Examples + -------- + Create a local temporary view named 'people'. + + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.createOrReplaceTempView("people") + + Replace the local temporary view. + + >>> df2 = df.filter(df.age > 3) + >>> df2.createOrReplaceTempView("people") + >>> df3 = spark.sql("SELECT * FROM people") + >>> sorted(df3.collect()) == sorted(df2.collect()) + True + >>> spark.catalog.dropTempView("people") + True + + """ + self._jdf.createOrReplaceTempView(name) + + def createGlobalTempView(self, name: str) -> None: + """Creates a global temporary view with this :class:`DataFrame`. + + The lifetime of this temporary view is tied to this Spark application. + throws :class:`TempTableAlreadyExistsException`, if the view name already exists in the + catalog. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + Name of the view. + + Examples + -------- + Create a global temporary view. + + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.createGlobalTempView("people") + >>> df2 = spark.sql("SELECT * FROM global_temp.people") + >>> sorted(df.collect()) == sorted(df2.collect()) + True + + Throws an exception if the global temporary view already exists. + + >>> df.createGlobalTempView("people") # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + AnalysisException: "Temporary table 'people' already exists;" + >>> spark.catalog.dropGlobalTempView("people") + True + + """ + self._jdf.createGlobalTempView(name) + + def createOrReplaceGlobalTempView(self, name: str) -> None: + """Creates or replaces a global temporary view using the given name. + + The lifetime of this temporary view is tied to this Spark application. + + .. versionadded:: 2.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + Name of the view. + + Examples + -------- + Create a global temporary view. + + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.createOrReplaceGlobalTempView("people") + + Replace the global temporary view. + + >>> df2 = df.filter(df.age > 3) + >>> df2.createOrReplaceGlobalTempView("people") + >>> df3 = spark.sql("SELECT * FROM global_temp.people") + >>> sorted(df3.collect()) == sorted(df2.collect()) + True + >>> spark.catalog.dropGlobalTempView("people") + True + + """ + self._jdf.createOrReplaceGlobalTempView(name) + + @property + def write(self) -> DataFrameWriter: + """ + Interface for saving the content of the non-streaming :class:`DataFrame` out into external + storage. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrameWriter` + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> type(df.write) + + + Write the DataFrame as a table. + + >>> _ = spark.sql("DROP TABLE IF EXISTS tab2") + >>> df.write.saveAsTable("tab2") + >>> _ = spark.sql("DROP TABLE tab2") + """ + return DataFrameWriter(self) + + @property + def writeStream(self) -> DataStreamWriter: + """ + Interface for saving the content of the streaming :class:`DataFrame` out into external + storage. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Returns + ------- + :class:`DataStreamWriter` + + Examples + -------- + >>> import tempfile + >>> df = spark.readStream.format("rate").load() + >>> type(df.writeStream) + + + >>> with tempfile.TemporaryDirectory() as d: + ... # Create a table with Rate source. + ... df.writeStream.toTable( + ... "my_table", checkpointLocation=d) + <...streaming.query.StreamingQuery object at 0x...> + """ + return DataStreamWriter(self) + + @property + def schema(self) -> StructType: + """Returns the schema of this :class:`DataFrame` as a :class:`pyspark.sql.types.StructType`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`StructType` + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Retrieve the schema of the current DataFrame. + + >>> df.schema + StructType([StructField('age', LongType(), True), + StructField('name', StringType(), True)]) + """ + if self._schema is None: + try: + self._schema = cast( + StructType, _parse_datatype_json_string(self._jdf.schema().json()) + ) + except Exception as e: + raise PySparkValueError( + error_class="CANNOT_PARSE_DATATYPE", + message_parameters={"error": str(e)}, + ) + return self._schema + + def printSchema(self, level: Optional[int] = None) -> None: + """Prints out the schema in the tree format. + Optionally allows to specify how many levels to print if schema is nested. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + level : int, optional, default None + How many levels to print for nested schemas. + + .. versionchanged:: 3.5.0 + Added Level parameter. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.printSchema() + root + |-- age: long (nullable = true) + |-- name: string (nullable = true) + + >>> df = spark.createDataFrame([(1, (2,2))], ["a", "b"]) + >>> df.printSchema(1) + root + |-- a: long (nullable = true) + |-- b: struct (nullable = true) + + >>> df.printSchema(2) + root + |-- a: long (nullable = true) + |-- b: struct (nullable = true) + | |-- _1: long (nullable = true) + | |-- _2: long (nullable = true) + """ + if level: + print(self._jdf.schema().treeString(level)) + else: + print(self._jdf.schema().treeString()) + + def explain( + self, extended: Optional[Union[bool, str]] = None, mode: Optional[str] = None + ) -> None: + """Prints the (logical and physical) plans to the console for debugging purposes. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + extended : bool, optional + default ``False``. If ``False``, prints only the physical plan. + When this is a string without specifying the ``mode``, it works as the mode is + specified. + mode : str, optional + specifies the expected output format of plans. + + * ``simple``: Print only a physical plan. + * ``extended``: Print both logical and physical plans. + * ``codegen``: Print a physical plan and generated codes if they are available. + * ``cost``: Print a logical plan and statistics if they are available. + * ``formatted``: Split explain output into two sections: a physical plan outline \ + and node details. + + .. versionchanged:: 3.0.0 + Added optional argument `mode` to specify the expected output format of plans. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Print out the physical plan only (default). + + >>> df.explain() # doctest: +SKIP + == Physical Plan == + *(1) Scan ExistingRDD[age...,name...] + + Print out all of the parsed, analyzed, optimized and physical plans. + + >>> df.explain(True) + == Parsed Logical Plan == + ... + == Analyzed Logical Plan == + ... + == Optimized Logical Plan == + ... + == Physical Plan == + ... + + Print out the plans with two sections: a physical plan outline and node details + + >>> df.explain(mode="formatted") # doctest: +SKIP + == Physical Plan == + * Scan ExistingRDD (...) + (1) Scan ExistingRDD [codegen id : ...] + Output [2]: [age..., name...] + ... + + Print a logical plan and statistics if they are available. + + >>> df.explain("cost") + == Optimized Logical Plan == + ...Statistics... + ... + """ + + if extended is not None and mode is not None: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={"arg_list": "extended and mode"}, + ) + + # For the no argument case: df.explain() + is_no_argument = extended is None and mode is None + + # For the cases below: + # explain(True) + # explain(extended=False) + is_extended_case = isinstance(extended, bool) and mode is None + + # For the case when extended is mode: + # df.explain("formatted") + is_extended_as_mode = isinstance(extended, str) and mode is None + + # For the mode specified: + # df.explain(mode="formatted") + is_mode_case = extended is None and isinstance(mode, str) + + if not (is_no_argument or is_extended_case or is_extended_as_mode or is_mode_case): + if (extended is not None) and (not isinstance(extended, (bool, str))): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_STR", + message_parameters={ + "arg_name": "extended", + "arg_type": type(extended).__name__, + }, + ) + if (mode is not None) and (not isinstance(mode, str)): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "mode", "arg_type": type(mode).__name__}, + ) + + # Sets an explain mode depending on a given argument + if is_no_argument: + explain_mode = "simple" + elif is_extended_case: + explain_mode = "extended" if extended else "simple" + elif is_mode_case: + explain_mode = cast(str, mode) + elif is_extended_as_mode: + explain_mode = cast(str, extended) + assert self._sc._jvm is not None + print(self._sc._jvm.PythonSQLUtils.explainString(self._jdf.queryExecution(), explain_mode)) + + def exceptAll(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing rows in this :class:`DataFrame` but + not in another :class:`DataFrame` while preserving duplicates. + + This is equivalent to `EXCEPT ALL` in SQL. + As standard in SQL, this function resolves columns by position (not by name). + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + The other :class:`DataFrame` to compare to. + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> df1 = spark.createDataFrame( + ... [("a", 1), ("a", 1), ("a", 1), ("a", 2), ("b", 3), ("c", 4)], ["C1", "C2"]) + >>> df2 = spark.createDataFrame([("a", 1), ("b", 3)], ["C1", "C2"]) + >>> df1.exceptAll(df2).show() + +---+---+ + | C1| C2| + +---+---+ + | a| 1| + | a| 1| + | a| 2| + | c| 4| + +---+---+ + + """ + return DataFrame(self._jdf.exceptAll(other._jdf), self.sparkSession) + + def isLocal(self) -> bool: + """Returns ``True`` if the :func:`collect` and :func:`take` methods can be run locally + (without any Spark executors). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + bool + + Examples + -------- + >>> df = spark.sql("SHOW TABLES") + >>> df.isLocal() + True + """ + return self._jdf.isLocal() + + @property + def isStreaming(self) -> bool: + """Returns ``True`` if this :class:`DataFrame` contains one or more sources that + continuously return data as it arrives. A :class:`DataFrame` that reads data from a + streaming source must be executed as a :class:`StreamingQuery` using the :func:`start` + method in :class:`DataStreamWriter`. Methods that return a single answer, (e.g., + :func:`count` or :func:`collect`) will throw an :class:`AnalysisException` when there + is a streaming source present. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Returns + ------- + bool + Whether it's streaming DataFrame or not. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.isStreaming + True + """ + return self._jdf.isStreaming() + + def isEmpty(self) -> bool: + """ + Checks if the :class:`DataFrame` is empty and returns a boolean value. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + bool + Returns ``True`` if the DataFrame is empty, ``False`` otherwise. + + See Also + -------- + DataFrame.count : Counts the number of rows in DataFrame. + + Notes + ----- + - Unlike `count()`, this method does not trigger any computation. + - An empty DataFrame has no rows. It may have columns, but no data. + + Examples + -------- + Example 1: Checking if an empty DataFrame is empty + + >>> df_empty = spark.createDataFrame([], 'a STRING') + >>> df_empty.isEmpty() + True + + Example 2: Checking if a non-empty DataFrame is empty + + >>> df_non_empty = spark.createDataFrame(["a"], 'STRING') + >>> df_non_empty.isEmpty() + False + + Example 3: Checking if a DataFrame with null values is empty + + >>> df_nulls = spark.createDataFrame([(None, None)], 'a STRING, b INT') + >>> df_nulls.isEmpty() + False + + Example 4: Checking if a DataFrame with no rows but with columns is empty + + >>> df_no_rows = spark.createDataFrame([], 'id INT, value STRING') + >>> df_no_rows.isEmpty() + True + """ + return self._jdf.isEmpty() + + def show(self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False) -> None: + """Prints the first ``n`` rows to the console. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + n : int, optional + Number of rows to show. + truncate : bool or int, optional + If set to ``True``, truncate strings longer than 20 chars by default. + If set to a number greater than one, truncates long strings to length ``truncate`` + and align cells right. + vertical : bool, optional + If set to ``True``, print output rows vertically (one line + per column value). + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Show only top 2 rows. + + >>> df.show(2) + +---+-----+ + |age| name| + +---+-----+ + | 14| Tom| + | 23|Alice| + +---+-----+ + only showing top 2 rows + + Show :class:`DataFrame` where the maximum number of characters is 3. + + >>> df.show(truncate=3) + +---+----+ + |age|name| + +---+----+ + | 14| Tom| + | 23| Ali| + | 16| Bob| + +---+----+ + + Show :class:`DataFrame` vertically. + + >>> df.show(vertical=True) + -RECORD 0----- + age | 14 + name | Tom + -RECORD 1----- + age | 23 + name | Alice + -RECORD 2----- + age | 16 + name | Bob + """ + + if not isinstance(n, int) or isinstance(n, bool): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "n", "arg_type": type(n).__name__}, + ) + + if not isinstance(vertical, bool): + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={"arg_name": "vertical", "arg_type": type(vertical).__name__}, + ) + + if isinstance(truncate, bool) and truncate: + print(self._jdf.showString(n, 20, vertical)) + else: + try: + int_truncate = int(truncate) + except ValueError: + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={ + "arg_name": "truncate", + "arg_type": type(truncate).__name__, + }, + ) + + print(self._jdf.showString(n, int_truncate, vertical)) + + def __repr__(self) -> str: + if not self._support_repr_html and self.sparkSession._jconf.isReplEagerEvalEnabled(): + vertical = False + return self._jdf.showString( + self.sparkSession._jconf.replEagerEvalMaxNumRows(), + self.sparkSession._jconf.replEagerEvalTruncate(), + vertical, + ) + else: + return "DataFrame[%s]" % (", ".join("%s: %s" % c for c in self.dtypes)) + + def _repr_html_(self) -> Optional[str]: + """Returns a :class:`DataFrame` with html code when you enabled eager evaluation + by 'spark.sql.repl.eagerEval.enabled', this only called by REPL you are + using support eager evaluation with HTML. + """ + if not self._support_repr_html: + self._support_repr_html = True + if self.sparkSession._jconf.isReplEagerEvalEnabled(): + return self._jdf.htmlString( + self.sparkSession._jconf.replEagerEvalMaxNumRows(), + self.sparkSession._jconf.replEagerEvalTruncate(), + ) + else: + return None + + def checkpoint(self, eager: bool = True) -> "DataFrame": + """Returns a checkpointed version of this :class:`DataFrame`. Checkpointing can be used to + truncate the logical plan of this :class:`DataFrame`, which is especially useful in + iterative algorithms where the plan may grow exponentially. It will be saved to files + inside the checkpoint directory set with :meth:`SparkContext.setCheckpointDir`. + + .. versionadded:: 2.1.0 + + Parameters + ---------- + eager : bool, optional, default True + Whether to checkpoint this :class:`DataFrame` immediately. + + Returns + ------- + :class:`DataFrame` + Checkpointed DataFrame. + + Notes + ----- + This API is experimental. + + Examples + -------- + >>> import tempfile + >>> df = spark.createDataFrame([ + ... (14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> with tempfile.TemporaryDirectory() as d: + ... spark.sparkContext.setCheckpointDir("/tmp/bb") + ... df.checkpoint(False) + DataFrame[age: bigint, name: string] + """ + jdf = self._jdf.checkpoint(eager) + return DataFrame(jdf, self.sparkSession) + + def localCheckpoint(self, eager: bool = True) -> "DataFrame": + """Returns a locally checkpointed version of this :class:`DataFrame`. Checkpointing can be + used to truncate the logical plan of this :class:`DataFrame`, which is especially useful in + iterative algorithms where the plan may grow exponentially. Local checkpoints are + stored in the executors using the caching subsystem and therefore they are not reliable. + + .. versionadded:: 2.3.0 + + Parameters + ---------- + eager : bool, optional, default True + Whether to checkpoint this :class:`DataFrame` immediately. + + Returns + ------- + :class:`DataFrame` + Checkpointed DataFrame. + + Notes + ----- + This API is experimental. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.localCheckpoint(False) + DataFrame[age: bigint, name: string] + """ + jdf = self._jdf.localCheckpoint(eager) + return DataFrame(jdf, self.sparkSession) + + def withWatermark(self, eventTime: str, delayThreshold: str) -> "DataFrame": + """Defines an event time watermark for this :class:`DataFrame`. A watermark tracks a point + in time before which we assume no more late data is going to arrive. + + Spark will use this watermark for several purposes: + - To know when a given time window aggregation can be finalized and thus can be emitted + when using output modes that do not allow updates. + + - To minimize the amount of state that we need to keep for on-going aggregations. + + The current watermark is computed by looking at the `MAX(eventTime)` seen across + all of the partitions in the query minus a user specified `delayThreshold`. Due to the cost + of coordinating this value across partitions, the actual watermark used is only guaranteed + to be at least `delayThreshold` behind the actual event time. In some cases we may still + process records that arrive more than `delayThreshold` late. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + eventTime : str + the name of the column that contains the event time of the row. + delayThreshold : str + the minimum delay to wait to data to arrive late, relative to the + latest record that has been processed in the form of an interval + (e.g. "1 minute" or "5 hours"). + + Returns + ------- + :class:`DataFrame` + Watermarked DataFrame + + Notes + ----- + This is a feature only for Structured Streaming. + + This API is evolving. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import timestamp_seconds + >>> df = spark.readStream.format("rate").load().selectExpr( + ... "value % 5 AS value", "timestamp") + >>> df.select("value", df.timestamp.alias("time")).withWatermark("time", '10 minutes') + DataFrame[value: bigint, time: timestamp] + + Group the data by window and value (0 - 4), and compute the count of each group. + + >>> import time + >>> from pyspark.sql.functions import window + >>> query = (df + ... .withWatermark("timestamp", "10 minutes") + ... .groupBy( + ... window(df.timestamp, "10 minutes", "5 minutes"), + ... df.value) + ... ).count().writeStream.outputMode("complete").format("console").start() + >>> time.sleep(3) + >>> query.stop() + """ + if not eventTime or type(eventTime) is not str: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "eventTime", "arg_type": type(eventTime).__name__}, + ) + if not delayThreshold or type(delayThreshold) is not str: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={ + "arg_name": "delayThreshold", + "arg_type": type(delayThreshold).__name__, + }, + ) + jdf = self._jdf.withWatermark(eventTime, delayThreshold) + return DataFrame(jdf, self.sparkSession) + + def hint( + self, name: str, *parameters: Union["PrimitiveType", List["PrimitiveType"]] + ) -> "DataFrame": + """Specifies some hint on the current :class:`DataFrame`. + + .. versionadded:: 2.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + A name of the hint. + parameters : str, list, float or int + Optional parameters. + + Returns + ------- + :class:`DataFrame` + Hinted DataFrame + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df2 = spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")]) + >>> df.join(df2, "name").explain() # doctest: +SKIP + == Physical Plan == + ... + ... +- SortMergeJoin ... + ... + + Explicitly trigger the broadcast hashjoin by providing the hint in ``df2``. + + >>> df.join(df2.hint("broadcast"), "name").explain() + == Physical Plan == + ... + ... +- BroadcastHashJoin ... + ... + """ + if len(parameters) == 1 and isinstance(parameters[0], list): + parameters = parameters[0] # type: ignore[assignment] + + if not isinstance(name, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "name", "arg_type": type(name).__name__}, + ) + + allowed_types = (str, list, float, int) + for p in parameters: + if not isinstance(p, allowed_types): + raise PySparkTypeError( + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "parameters", + "arg_type": type(parameters).__name__, + "allowed_types": ", ".join(map(lambda x: x.__name__, allowed_types)), + "return_type": type(p).__name__, + }, + ) + + jdf = self._jdf.hint(name, self._jseq(parameters)) + return DataFrame(jdf, self.sparkSession) + + def count(self) -> int: + """Returns the number of rows in this :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + int + Number of rows. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Return the number of rows in the :class:`DataFrame`. + + >>> df.count() + 3 + """ + return int(self._jdf.count()) + + def collect(self) -> List[Row]: + """Returns all the records as a list of :class:`Row`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + list + List of rows. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.collect() + [Row(age=14, name='Tom'), Row(age=23, name='Alice'), Row(age=16, name='Bob')] + """ + with SCCallSiteSync(self._sc): + sock_info = self._jdf.collectToPython() + return list(_load_from_socket(sock_info, BatchedSerializer(CPickleSerializer()))) + + def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]: + """ + Returns an iterator that contains all of the rows in this :class:`DataFrame`. + The iterator will consume as much memory as the largest partition in this + :class:`DataFrame`. With prefetch it may consume up to the memory of the 2 largest + partitions. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + prefetchPartitions : bool, optional + If Spark should pre-fetch the next partition before it is needed. + + .. versionchanged:: 3.4.0 + This argument does not take effect for Spark Connect. + + Returns + ------- + Iterator + Iterator of rows. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> list(df.toLocalIterator()) + [Row(age=14, name='Tom'), Row(age=23, name='Alice'), Row(age=16, name='Bob')] + """ + with SCCallSiteSync(self._sc): + sock_info = self._jdf.toPythonIterator(prefetchPartitions) + return _local_iterator_from_socket(sock_info, BatchedSerializer(CPickleSerializer())) + + def limit(self, num: int) -> "DataFrame": + """Limits the result count to the number specified. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + num : int + Number of records to return. Will return this number of records + or all records if the DataFrame contains less than this number of records. + + Returns + ------- + :class:`DataFrame` + Subset of the records + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.limit(1).show() + +---+----+ + |age|name| + +---+----+ + | 14| Tom| + +---+----+ + >>> df.limit(0).show() + +---+----+ + |age|name| + +---+----+ + +---+----+ + """ + jdf = self._jdf.limit(num) + return DataFrame(jdf, self.sparkSession) + + def offset(self, num: int) -> "DataFrame": + """Returns a new :class: `DataFrame` by skipping the first `n` rows. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports vanilla PySpark. + + Parameters + ---------- + num : int + Number of records to skip. + + Returns + ------- + :class:`DataFrame` + Subset of the records + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.offset(1).show() + +---+-----+ + |age| name| + +---+-----+ + | 23|Alice| + | 16| Bob| + +---+-----+ + >>> df.offset(10).show() + +---+----+ + |age|name| + +---+----+ + +---+----+ + """ + jdf = self._jdf.offset(num) + return DataFrame(jdf, self.sparkSession) + + def take(self, num: int) -> List[Row]: + """Returns the first ``num`` rows as a :class:`list` of :class:`Row`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + num : int + Number of records to return. Will return this number of records + or all records if the DataFrame contains less than this number of records.. + + Returns + ------- + list + List of rows + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Return the first 2 rows of the :class:`DataFrame`. + + >>> df.take(2) + [Row(age=14, name='Tom'), Row(age=23, name='Alice')] + """ + return self.limit(num).collect() + + def tail(self, num: int) -> List[Row]: + """ + Returns the last ``num`` rows as a :class:`list` of :class:`Row`. + + Running tail requires moving data into the application's driver process, and doing so with + a very large ``num`` can crash the driver process with OutOfMemoryError. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + num : int + Number of records to return. Will return this number of records + or all records if the DataFrame contains less than this number of records. + + Returns + ------- + list + List of rows + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + >>> df.tail(2) + [Row(age=23, name='Alice'), Row(age=16, name='Bob')] + """ + with SCCallSiteSync(self._sc): + sock_info = self._jdf.tailToPython(num) + return list(_load_from_socket(sock_info, BatchedSerializer(CPickleSerializer()))) + + def foreach(self, f: Callable[[Row], None]) -> None: + """Applies the ``f`` function to all :class:`Row` of this :class:`DataFrame`. + + This is a shorthand for ``df.rdd.foreach()``. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + f : function + A function that accepts one parameter which will + receive each row to process. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> def func(person): + ... print(person.name) + ... + >>> df.foreach(func) + """ + self.rdd.foreach(f) + + def foreachPartition(self, f: Callable[[Iterator[Row]], None]) -> None: + """Applies the ``f`` function to each partition of this :class:`DataFrame`. + + This a shorthand for ``df.rdd.foreachPartition()``. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + f : function + A function that accepts one parameter which will receive + each partition to process. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> def func(itr): + ... for person in itr: + ... print(person.name) + ... + >>> df.foreachPartition(func) + """ + self.rdd.foreachPartition(f) # type: ignore[arg-type] + + def cache(self) -> "DataFrame": + """Persists the :class:`DataFrame` with the default storage level (`MEMORY_AND_DISK`). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The default storage level has changed to `MEMORY_AND_DISK` to match Scala in 2.0. + + Returns + ------- + :class:`DataFrame` + Cached DataFrame. + + Examples + -------- + >>> df = spark.range(1) + >>> df.cache() + DataFrame[id: bigint] + + >>> df.explain() + == Physical Plan == + AdaptiveSparkPlan isFinalPlan=false + +- InMemoryTableScan ... + """ + self.is_cached = True + self._jdf.cache() + return self + + def persist( + self, + storageLevel: StorageLevel = (StorageLevel.MEMORY_AND_DISK_DESER), + ) -> "DataFrame": + """Sets the storage level to persist the contents of the :class:`DataFrame` across + operations after the first time it is computed. This can only be used to assign + a new storage level if the :class:`DataFrame` does not have a storage level set yet. + If no storage level is specified defaults to (`MEMORY_AND_DISK_DESER`) + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The default storage level has changed to `MEMORY_AND_DISK_DESER` to match Scala in 3.0. + + Parameters + ---------- + storageLevel : :class:`StorageLevel` + Storage level to set for persistence. Default is MEMORY_AND_DISK_DESER. + + Returns + ------- + :class:`DataFrame` + Persisted DataFrame. + + Examples + -------- + >>> df = spark.range(1) + >>> df.persist() + DataFrame[id: bigint] + + >>> df.explain() + == Physical Plan == + AdaptiveSparkPlan isFinalPlan=false + +- InMemoryTableScan ... + + Persists the data in the disk by specifying the storage level. + + >>> from pyspark.storagelevel import StorageLevel + >>> df.persist(StorageLevel.DISK_ONLY) + DataFrame[id: bigint] + """ + self.is_cached = True + javaStorageLevel = self._sc._getJavaStorageLevel(storageLevel) + self._jdf.persist(javaStorageLevel) + return self + + @property + def storageLevel(self) -> StorageLevel: + """Get the :class:`DataFrame`'s current storage level. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`StorageLevel` + Currently defined storage level. + + Examples + -------- + >>> df1 = spark.range(10) + >>> df1.storageLevel + StorageLevel(False, False, False, False, 1) + >>> df1.cache().storageLevel + StorageLevel(True, True, False, True, 1) + + >>> df2 = spark.range(5) + >>> df2.persist(StorageLevel.DISK_ONLY_2).storageLevel + StorageLevel(True, False, False, False, 2) + """ + java_storage_level = self._jdf.storageLevel() + storage_level = StorageLevel( + java_storage_level.useDisk(), + java_storage_level.useMemory(), + java_storage_level.useOffHeap(), + java_storage_level.deserialized(), + java_storage_level.replication(), + ) + return storage_level + + def unpersist(self, blocking: bool = False) -> "DataFrame": + """Marks the :class:`DataFrame` as non-persistent, and remove all blocks for it from + memory and disk. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + `blocking` default has changed to ``False`` to match Scala in 2.0. + + Parameters + ---------- + blocking : bool + Whether to block until all blocks are deleted. + + Returns + ------- + :class:`DataFrame` + Unpersisted DataFrame. + + Examples + -------- + >>> df = spark.range(1) + >>> df.persist() + DataFrame[id: bigint] + >>> df.unpersist() + DataFrame[id: bigint] + >>> df = spark.range(1) + >>> df.unpersist(True) + DataFrame[id: bigint] + """ + self.is_cached = False + self._jdf.unpersist(blocking) + return self + + def coalesce(self, numPartitions: int) -> "DataFrame": + """ + Returns a new :class:`DataFrame` that has exactly `numPartitions` partitions. + + Similar to coalesce defined on an :class:`RDD`, this operation results in a + narrow dependency, e.g. if you go from 1000 partitions to 100 partitions, + there will not be a shuffle, instead each of the 100 new partitions will + claim 10 of the current partitions. If a larger number of partitions is requested, + it will stay at the current number of partitions. + + However, if you're doing a drastic coalesce, e.g. to numPartitions = 1, + this may result in your computation taking place on fewer nodes than + you like (e.g. one node in the case of numPartitions = 1). To avoid this, + you can call repartition(). This will add a shuffle step, but means the + current upstream partitions will be executed in parallel (per whatever + the current partitioning is). + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + numPartitions : int + specify the target number of partitions + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> df = spark.range(10) + >>> df.coalesce(1).rdd.getNumPartitions() + 1 + """ + return DataFrame(self._jdf.coalesce(numPartitions), self.sparkSession) + + @overload + def repartition(self, numPartitions: int, *cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def repartition(self, *cols: "ColumnOrName") -> "DataFrame": + ... + + def repartition( # type: ignore[misc] + self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName" + ) -> "DataFrame": + """ + Returns a new :class:`DataFrame` partitioned by the given partitioning expressions. The + resulting :class:`DataFrame` is hash partitioned. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + numPartitions : int + can be an int to specify the target number of partitions or a Column. + If it is a Column, it will be used as the first partitioning column. If not specified, + the default number of partitions is used. + cols : str or :class:`Column` + partitioning columns. + + .. versionchanged:: 1.6.0 + Added optional arguments to specify the partitioning columns. Also made numPartitions + optional if partitioning columns are specified. + + Returns + ------- + :class:`DataFrame` + Repartitioned DataFrame. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Repartition the data into 10 partitions. + + >>> df.repartition(10).rdd.getNumPartitions() + 10 + + Repartition the data into 7 partitions by 'age' column. + + >>> df.repartition(7, "age").rdd.getNumPartitions() + 7 + + Repartition the data into 7 partitions by 'age' and 'name columns. + + >>> df.repartition(3, "name", "age").rdd.getNumPartitions() + 3 + """ + if isinstance(numPartitions, int): + if len(cols) == 0: + return DataFrame(self._jdf.repartition(numPartitions), self.sparkSession) + else: + return DataFrame( + self._jdf.repartition(numPartitions, self._jcols(*cols)), + self.sparkSession, + ) + elif isinstance(numPartitions, (str, Column)): + cols = (numPartitions,) + cols + return DataFrame(self._jdf.repartition(self._jcols(*cols)), self.sparkSession) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={ + "arg_name": "numPartitions", + "arg_type": type(numPartitions).__name__, + }, + ) + + @overload + def repartitionByRange(self, numPartitions: int, *cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def repartitionByRange(self, *cols: "ColumnOrName") -> "DataFrame": + ... + + def repartitionByRange( # type: ignore[misc] + self, numPartitions: Union[int, "ColumnOrName"], *cols: "ColumnOrName" + ) -> "DataFrame": + """ + Returns a new :class:`DataFrame` partitioned by the given partitioning expressions. The + resulting :class:`DataFrame` is range partitioned. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + numPartitions : int + can be an int to specify the target number of partitions or a Column. + If it is a Column, it will be used as the first partitioning column. If not specified, + the default number of partitions is used. + cols : str or :class:`Column` + partitioning columns. + + Returns + ------- + :class:`DataFrame` + Repartitioned DataFrame. + + Notes + ----- + At least one partition-by expression must be specified. + When no explicit sort order is specified, "ascending nulls first" is assumed. + + Due to performance reasons this method uses sampling to estimate the ranges. + Hence, the output may not be consistent, since sampling can return different values. + The sample size can be controlled by the config + `spark.sql.execution.rangeExchange.sampleSizePerPartition`. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + Repartition the data into 2 partitions by range in 'age' column. + For example, the first partition can have ``(14, "Tom")``, and the second + partition would have ``(16, "Bob")`` and ``(23, "Alice")``. + + >>> df.repartitionByRange(2, "age").rdd.getNumPartitions() + 2 + """ + if isinstance(numPartitions, int): + if len(cols) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "partition-by expression"}, + ) + else: + return DataFrame( + self._jdf.repartitionByRange(numPartitions, self._jcols(*cols)), + self.sparkSession, + ) + elif isinstance(numPartitions, (str, Column)): + cols = (numPartitions,) + cols + return DataFrame(self._jdf.repartitionByRange(self._jcols(*cols)), self.sparkSession) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={ + "arg_name": "numPartitions", + "arg_type": type(numPartitions).__name__, + }, + ) + + def distinct(self) -> "DataFrame": + """Returns a new :class:`DataFrame` containing the distinct rows in this :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrame` + DataFrame with distinct records. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (23, "Alice")], ["age", "name"]) + + Return the number of distinct rows in the :class:`DataFrame` + + >>> df.distinct().count() + 2 + """ + return DataFrame(self._jdf.distinct(), self.sparkSession) + + @overload + def sample(self, fraction: float, seed: Optional[int] = ...) -> "DataFrame": + ... + + @overload + def sample( + self, + withReplacement: Optional[bool], + fraction: float, + seed: Optional[int] = ..., + ) -> "DataFrame": + ... + + def sample( # type: ignore[misc] + self, + withReplacement: Optional[Union[float, bool]] = None, + fraction: Optional[Union[int, float]] = None, + seed: Optional[int] = None, + ) -> "DataFrame": + """Returns a sampled subset of this :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + withReplacement : bool, optional + Sample with replacement or not (default ``False``). + fraction : float, optional + Fraction of rows to generate, range [0.0, 1.0]. + seed : int, optional + Seed for sampling (default a random seed). + + Returns + ------- + :class:`DataFrame` + Sampled rows from given DataFrame. + + Notes + ----- + This is not guaranteed to provide exactly the fraction specified of the total + count of the given :class:`DataFrame`. + + `fraction` is required and, `withReplacement` and `seed` are optional. + + Examples + -------- + >>> df = spark.range(10) + >>> df.sample(0.5, 3).count() # doctest: +SKIP + 7 + >>> df.sample(fraction=0.5, seed=3).count() # doctest: +SKIP + 7 + >>> df.sample(withReplacement=True, fraction=0.5, seed=3).count() # doctest: +SKIP + 1 + >>> df.sample(1.0).count() + 10 + >>> df.sample(fraction=1.0).count() + 10 + >>> df.sample(False, fraction=1.0).count() + 10 + """ + + # For the cases below: + # sample(True, 0.5 [, seed]) + # sample(True, fraction=0.5 [, seed]) + # sample(withReplacement=False, fraction=0.5 [, seed]) + is_withReplacement_set = type(withReplacement) == bool and isinstance(fraction, float) + + # For the case below: + # sample(faction=0.5 [, seed]) + is_withReplacement_omitted_kwargs = withReplacement is None and isinstance(fraction, float) + + # For the case below: + # sample(0.5 [, seed]) + is_withReplacement_omitted_args = isinstance(withReplacement, float) + + if not ( + is_withReplacement_set + or is_withReplacement_omitted_kwargs + or is_withReplacement_omitted_args + ): + argtypes = [type(arg).__name__ for arg in [withReplacement, fraction, seed]] + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT", + message_parameters={ + "arg_name": "withReplacement (optional), " + + "fraction (required) and seed (optional)", + "arg_type": ", ".join(argtypes), + }, + ) + + if is_withReplacement_omitted_args: + if fraction is not None: + seed = cast(int, fraction) + fraction = withReplacement + withReplacement = None + + seed = int(seed) if seed is not None else None + args = [arg for arg in [withReplacement, fraction, seed] if arg is not None] + jdf = self._jdf.sample(*args) + return DataFrame(jdf, self.sparkSession) + + def sampleBy( + self, col: "ColumnOrName", fractions: Dict[Any, float], seed: Optional[int] = None + ) -> "DataFrame": + """ + Returns a stratified sample without replacement based on the + fraction given on each stratum. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`Column` or str + column that defines strata + + .. versionchanged:: 3.0.0 + Added sampling by a column of :class:`Column` + fractions : dict + sampling fraction for each stratum. If a stratum is not + specified, we treat its fraction as zero. + seed : int, optional + random seed + + Returns + ------- + a new :class:`DataFrame` that represents the stratified sample + + Examples + -------- + >>> from pyspark.sql.functions import col + >>> dataset = spark.range(0, 100).select((col("id") % 3).alias("key")) + >>> sampled = dataset.sampleBy("key", fractions={0: 0.1, 1: 0.2}, seed=0) + >>> sampled.groupBy("key").count().orderBy("key").show() + +---+-----+ + |key|count| + +---+-----+ + | 0| 3| + | 1| 6| + +---+-----+ + >>> dataset.sampleBy(col("key"), fractions={2: 1.0}, seed=0).count() + 33 + """ + if isinstance(col, str): + col = Column(col) + elif not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + if not isinstance(fractions, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "fractions", "arg_type": type(fractions).__name__}, + ) + for k, v in fractions.items(): + if not isinstance(k, (float, int, str)): + raise PySparkTypeError( + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "fractions", + "arg_type": type(fractions).__name__, + "allowed_types": "float, int, str", + "return_type": type(k).__name__, + }, + ) + fractions[k] = float(v) + col = col._jc + seed = seed if seed is not None else random.randint(0, sys.maxsize) + return DataFrame( + self._jdf.stat().sampleBy(col, self._jmap(fractions), seed), self.sparkSession + ) + + def randomSplit(self, weights: List[float], seed: Optional[int] = None) -> List["DataFrame"]: + """Randomly splits this :class:`DataFrame` with the provided weights. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + weights : list + list of doubles as weights with which to split the :class:`DataFrame`. + Weights will be normalized if they don't sum up to 1.0. + seed : int, optional + The seed for sampling. + + Returns + ------- + list + List of DataFrames. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([ + ... Row(age=10, height=80, name="Alice"), + ... Row(age=5, height=None, name="Bob"), + ... Row(age=None, height=None, name="Tom"), + ... Row(age=None, height=None, name=None), + ... ]) + + >>> splits = df.randomSplit([1.0, 2.0], 24) + >>> splits[0].count() + 2 + >>> splits[1].count() + 2 + """ + for w in weights: + if w < 0.0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "weights", "arg_value": str(w)}, + ) + seed = seed if seed is not None else random.randint(0, sys.maxsize) + df_array = self._jdf.randomSplit( + _to_list(self.sparkSession._sc, cast(List["ColumnOrName"], weights)), int(seed) + ) + return [DataFrame(df, self.sparkSession) for df in df_array] + + @property + def dtypes(self) -> List[Tuple[str, str]]: + """Returns all column names and their data types as a list. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + list + List of columns as tuple pairs. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.dtypes + [('age', 'bigint'), ('name', 'string')] + """ + return [(str(f.name), f.dataType.simpleString()) for f in self.schema.fields] + + @property + def columns(self) -> List[str]: + """ + Retrieves the names of all columns in the :class:`DataFrame` as a list. + + The order of the column names in the list reflects their order in the DataFrame. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + list + List of column names in the DataFrame. + + Examples + -------- + Example 1: Retrieve column names of a DataFrame + + >>> df = spark.createDataFrame( + ... [(14, "Tom", "CA"), (23, "Alice", "NY"), (16, "Bob", "TX")], + ... ["age", "name", "state"] + ... ) + >>> df.columns + ['age', 'name', 'state'] + + Example 2: Using column names to project specific columns + + >>> selected_cols = [col for col in df.columns if col != "age"] + >>> df.select(selected_cols).show() + +-----+-----+ + | name|state| + +-----+-----+ + | Tom| CA| + |Alice| NY| + | Bob| TX| + +-----+-----+ + + Example 3: Checking if a specific column exists in a DataFrame + + >>> "state" in df.columns + True + >>> "salary" in df.columns + False + + Example 4: Iterating over columns to apply a transformation + + >>> import pyspark.sql.functions as f + >>> for col_name in df.columns: + ... df = df.withColumn(col_name, f.upper(f.col(col_name))) + >>> df.show() + +---+-----+-----+ + |age| name|state| + +---+-----+-----+ + | 14| TOM| CA| + | 23|ALICE| NY| + | 16| BOB| TX| + +---+-----+-----+ + + Example 5: Renaming columns and checking the updated column names + + >>> df = df.withColumnRenamed("name", "first_name") + >>> df.columns + ['age', 'first_name', 'state'] + + Example 6: Using the `columns` property to ensure two DataFrames have the + same columns before a union + + >>> df2 = spark.createDataFrame( + ... [(30, "Eve", "FL"), (40, "Sam", "WA")], ["age", "name", "location"]) + >>> df.columns == df2.columns + False + """ + return [f.name for f in self.schema.fields] + + def colRegex(self, colName: str) -> Column: + """ + Selects column based on the column name specified as a regex and returns it + as :class:`Column`. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + colName : str + string, column name specified as a regex. + + Returns + ------- + :class:`Column` + + Examples + -------- + >>> df = spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["Col1", "Col2"]) + >>> df.select(df.colRegex("`(Col1)?+.+`")).show() + +----+ + |Col2| + +----+ + | 1| + | 2| + | 3| + +----+ + """ + if not isinstance(colName, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "colName", "arg_type": type(colName).__name__}, + ) + jc = self._jdf.colRegex(colName) + return Column(jc) + + def to(self, schema: StructType) -> "DataFrame": + """ + Returns a new :class:`DataFrame` where each row is reconciled to match the specified + schema. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + schema : :class:`StructType` + Specified schema. + + Returns + ------- + :class:`DataFrame` + Reconciled DataFrame. + + Notes + ----- + * Reorder columns and/or inner fields by name to match the specified schema. + + * Project away columns and/or inner fields that are not needed by the specified schema. + Missing columns and/or inner fields (present in the specified schema but not input + DataFrame) lead to failures. + + * Cast the columns and/or inner fields to match the data types in the specified schema, + if the types are compatible, e.g., numeric to numeric (error if overflows), but + not string to int. + + * Carry over the metadata from the specified schema, while the columns and/or inner fields + still keep their own metadata if not overwritten by the specified schema. + + * Fail if the nullability is not compatible. For example, the column and/or inner field + is nullable but the specified schema requires them to be not nullable. + + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql.types import StructField, StringType + >>> df = spark.createDataFrame([("a", 1)], ["i", "j"]) + >>> df.schema + StructType([StructField('i', StringType(), True), StructField('j', LongType(), True)]) + + >>> schema = StructType([StructField("j", StringType()), StructField("i", StringType())]) + >>> df2 = df.to(schema) + >>> df2.schema + StructType([StructField('j', StringType(), True), StructField('i', StringType(), True)]) + >>> df2.show() + +---+---+ + | j| i| + +---+---+ + | 1| a| + +---+---+ + """ + assert schema is not None + jschema = self._jdf.sparkSession().parseDataType(schema.json()) + return DataFrame(self._jdf.to(jschema), self.sparkSession) + + def alias(self, alias: str) -> "DataFrame": + """Returns a new :class:`DataFrame` with an alias set. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + alias : str + an alias name to be set for the :class:`DataFrame`. + + Returns + ------- + :class:`DataFrame` + Aliased DataFrame. + + Examples + -------- + >>> from pyspark.sql.functions import col, desc + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df_as1 = df.alias("df_as1") + >>> df_as2 = df.alias("df_as2") + >>> joined_df = df_as1.join(df_as2, col("df_as1.name") == col("df_as2.name"), 'inner') + >>> joined_df.select( + ... "df_as1.name", "df_as2.name", "df_as2.age").sort(desc("df_as1.name")).show() + +-----+-----+---+ + | name| name|age| + +-----+-----+---+ + | Tom| Tom| 14| + | Bob| Bob| 16| + |Alice|Alice| 23| + +-----+-----+---+ + """ + assert isinstance(alias, str), "alias should be a string" + return DataFrame(getattr(self._jdf, "as")(alias), self.sparkSession) + + def crossJoin(self, other: "DataFrame") -> "DataFrame": + """Returns the cartesian product with another :class:`DataFrame`. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Right side of the cartesian product. + + Returns + ------- + :class:`DataFrame` + Joined DataFrame. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df2 = spark.createDataFrame( + ... [Row(height=80, name="Tom"), Row(height=85, name="Bob")]) + >>> df.crossJoin(df2.select("height")).select("age", "name", "height").show() + +---+-----+------+ + |age| name|height| + +---+-----+------+ + | 14| Tom| 80| + | 14| Tom| 85| + | 23|Alice| 80| + | 23|Alice| 85| + | 16| Bob| 80| + | 16| Bob| 85| + +---+-----+------+ + """ + + jdf = self._jdf.crossJoin(other._jdf) + return DataFrame(jdf, self.sparkSession) + + def join( + self, + other: "DataFrame", + on: Optional[Union[str, List[str], Column, List[Column]]] = None, + how: Optional[str] = None, + ) -> "DataFrame": + """Joins with another :class:`DataFrame`, using the given join expression. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Right side of the join + on : str, list or :class:`Column`, optional + a string for the join column name, a list of column names, + a join expression (Column), or a list of Columns. + If `on` is a string or a list of strings indicating the name of the join column(s), + the column(s) must exist on both sides, and this performs an equi-join. + how : str, optional + default ``inner``. Must be one of: ``inner``, ``cross``, ``outer``, + ``full``, ``fullouter``, ``full_outer``, ``left``, ``leftouter``, ``left_outer``, + ``right``, ``rightouter``, ``right_outer``, ``semi``, ``leftsemi``, ``left_semi``, + ``anti``, ``leftanti`` and ``left_anti``. + + Returns + ------- + :class:`DataFrame` + Joined DataFrame. + + Examples + -------- + The following performs a full outer join between ``df1`` and ``df2``. + + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import desc + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")]).toDF("age", "name") + >>> df2 = spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")]) + >>> df3 = spark.createDataFrame([Row(age=2, name="Alice"), Row(age=5, name="Bob")]) + >>> df4 = spark.createDataFrame([ + ... Row(age=10, height=80, name="Alice"), + ... Row(age=5, height=None, name="Bob"), + ... Row(age=None, height=None, name="Tom"), + ... Row(age=None, height=None, name=None), + ... ]) + + Inner join on columns (default) + + >>> df.join(df2, 'name').select(df.name, df2.height).show() + +----+------+ + |name|height| + +----+------+ + | Bob| 85| + +----+------+ + >>> df.join(df4, ['name', 'age']).select(df.name, df.age).show() + +----+---+ + |name|age| + +----+---+ + | Bob| 5| + +----+---+ + + Outer join for both DataFrames on the 'name' column. + + >>> df.join(df2, df.name == df2.name, 'outer').select( + ... df.name, df2.height).sort(desc("name")).show() + +-----+------+ + | name|height| + +-----+------+ + | Bob| 85| + |Alice| NULL| + | NULL| 80| + +-----+------+ + >>> df.join(df2, 'name', 'outer').select('name', 'height').sort(desc("name")).show() + +-----+------+ + | name|height| + +-----+------+ + | Tom| 80| + | Bob| 85| + |Alice| NULL| + +-----+------+ + + Outer join for both DataFrams with multiple columns. + + >>> df.join( + ... df3, + ... [df.name == df3.name, df.age == df3.age], + ... 'outer' + ... ).select(df.name, df3.age).show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 2| + | Bob| 5| + +-----+---+ + """ + + if on is not None and not isinstance(on, list): + on = [on] # type: ignore[assignment] + + if on is not None: + if isinstance(on[0], str): + on = self._jseq(cast(List[str], on)) + else: + assert isinstance(on[0], Column), "on should be Column or list of Column" + on = reduce(lambda x, y: x.__and__(y), cast(List[Column], on)) + on = on._jc + + if on is None and how is None: + jdf = self._jdf.join(other._jdf) + else: + if how is None: + how = "inner" + if on is None: + on = self._jseq([]) + assert isinstance(how, str), "how should be a string" + jdf = self._jdf.join(other._jdf, on, how) + return DataFrame(jdf, self.sparkSession) + + # TODO(SPARK-22947): Fix the DataFrame API. + def _joinAsOf( + self, + other: "DataFrame", + leftAsOfColumn: Union[str, Column], + rightAsOfColumn: Union[str, Column], + on: Optional[Union[str, List[str], Column, List[Column]]] = None, + how: Optional[str] = None, + *, + tolerance: Optional[Column] = None, + allowExactMatches: bool = True, + direction: str = "backward", + ) -> "DataFrame": + """ + Perform an as-of join. + + This is similar to a left-join except that we match on the nearest + key rather than equal keys. + + Parameters + ---------- + other : :class:`DataFrame` + Right side of the join + leftAsOfColumn : str or :class:`Column` + a string for the as-of join column name, or a Column + rightAsOfColumn : str or :class:`Column` + a string for the as-of join column name, or a Column + on : str, list or :class:`Column`, optional + a string for the join column name, a list of column names, + a join expression (Column), or a list of Columns. + If `on` is a string or a list of strings indicating the name of the join column(s), + the column(s) must exist on both sides, and this performs an equi-join. + how : str, optional + default ``inner``. Must be one of: ``inner`` and ``left``. + tolerance : :class:`Column`, optional + an asof tolerance within this range; must be compatible + with the merge index. + allowExactMatches : bool, optional + default ``True``. + direction : str, optional + default ``backward``. Must be one of: ``backward``, ``forward``, and ``nearest``. + + Examples + -------- + The following performs an as-of join between ``left`` and ``right``. + + >>> left = spark.createDataFrame([(1, "a"), (5, "b"), (10, "c")], ["a", "left_val"]) + >>> right = spark.createDataFrame([(1, 1), (2, 2), (3, 3), (6, 6), (7, 7)], + ... ["a", "right_val"]) + >>> left._joinAsOf( + ... right, leftAsOfColumn="a", rightAsOfColumn="a" + ... ).select(left.a, 'left_val', 'right_val').sort("a").collect() + [Row(a=1, left_val='a', right_val=1), + Row(a=5, left_val='b', right_val=3), + Row(a=10, left_val='c', right_val=7)] + + >>> from pyspark.sql import functions as sf + >>> left._joinAsOf( + ... right, leftAsOfColumn="a", rightAsOfColumn="a", tolerance=sf.lit(1) + ... ).select(left.a, 'left_val', 'right_val').sort("a").collect() + [Row(a=1, left_val='a', right_val=1)] + + >>> left._joinAsOf( + ... right, leftAsOfColumn="a", rightAsOfColumn="a", how="left", tolerance=sf.lit(1) + ... ).select(left.a, 'left_val', 'right_val').sort("a").collect() + [Row(a=1, left_val='a', right_val=1), + Row(a=5, left_val='b', right_val=None), + Row(a=10, left_val='c', right_val=None)] + + >>> left._joinAsOf( + ... right, leftAsOfColumn="a", rightAsOfColumn="a", allowExactMatches=False + ... ).select(left.a, 'left_val', 'right_val').sort("a").collect() + [Row(a=5, left_val='b', right_val=3), + Row(a=10, left_val='c', right_val=7)] + + >>> left._joinAsOf( + ... right, leftAsOfColumn="a", rightAsOfColumn="a", direction="forward" + ... ).select(left.a, 'left_val', 'right_val').sort("a").collect() + [Row(a=1, left_val='a', right_val=1), + Row(a=5, left_val='b', right_val=6)] + """ + if isinstance(leftAsOfColumn, str): + leftAsOfColumn = self[leftAsOfColumn] + left_as_of_jcol = leftAsOfColumn._jc + if isinstance(rightAsOfColumn, str): + rightAsOfColumn = other[rightAsOfColumn] + right_as_of_jcol = rightAsOfColumn._jc + + if on is not None and not isinstance(on, list): + on = [on] # type: ignore[assignment] + + if on is not None: + if isinstance(on[0], str): + on = self._jseq(cast(List[str], on)) + else: + assert isinstance(on[0], Column), "on should be Column or list of Column" + on = reduce(lambda x, y: x.__and__(y), cast(List[Column], on)) + on = on._jc + + if how is None: + how = "inner" + assert isinstance(how, str), "how should be a string" + + if tolerance is not None: + assert isinstance(tolerance, Column), "tolerance should be Column" + tolerance = tolerance._jc + + jdf = self._jdf.joinAsOf( + other._jdf, + left_as_of_jcol, + right_as_of_jcol, + on, + how, + tolerance, + allowExactMatches, + direction, + ) + return DataFrame(jdf, self.sparkSession) + + def sortWithinPartitions( + self, *cols: Union[str, Column, List[Union[str, Column]]], **kwargs: Any + ) -> "DataFrame": + """Returns a new :class:`DataFrame` with each partition sorted by the specified column(s). + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str, list or :class:`Column`, optional + list of :class:`Column` or column names to sort by. + + Other Parameters + ---------------- + ascending : bool or list, optional, default True + boolean or list of boolean. + Sort ascending vs. descending. Specify list for multiple sort orders. + If a list is specified, the length of the list must equal the length of the `cols`. + + Returns + ------- + :class:`DataFrame` + DataFrame sorted by partitions. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.sortWithinPartitions("age", ascending=False) + DataFrame[age: bigint, name: string] + """ + jdf = self._jdf.sortWithinPartitions(self._sort_cols(cols, kwargs)) + return DataFrame(jdf, self.sparkSession) + + def sort( + self, *cols: Union[str, Column, List[Union[str, Column]]], **kwargs: Any + ) -> "DataFrame": + """Returns a new :class:`DataFrame` sorted by the specified column(s). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str, list, or :class:`Column`, optional + list of :class:`Column` or column names to sort by. + + Other Parameters + ---------------- + ascending : bool or list, optional, default True + boolean or list of boolean. + Sort ascending vs. descending. Specify list for multiple sort orders. + If a list is specified, the length of the list must equal the length of the `cols`. + + Returns + ------- + :class:`DataFrame` + Sorted DataFrame. + + Examples + -------- + >>> from pyspark.sql.functions import desc, asc + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + + Sort the DataFrame in ascending order. + + >>> df.sort(asc("age")).show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + | 5| Bob| + +---+-----+ + + Sort the DataFrame in descending order. + + >>> df.sort(df.age.desc()).show() + +---+-----+ + |age| name| + +---+-----+ + | 5| Bob| + | 2|Alice| + +---+-----+ + >>> df.orderBy(df.age.desc()).show() + +---+-----+ + |age| name| + +---+-----+ + | 5| Bob| + | 2|Alice| + +---+-----+ + >>> df.sort("age", ascending=False).show() + +---+-----+ + |age| name| + +---+-----+ + | 5| Bob| + | 2|Alice| + +---+-----+ + + Specify multiple columns + + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]) + >>> df.orderBy(desc("age"), "name").show() + +---+-----+ + |age| name| + +---+-----+ + | 5| Bob| + | 2|Alice| + | 2| Bob| + +---+-----+ + + Specify multiple columns for sorting order at `ascending`. + + >>> df.orderBy(["age", "name"], ascending=[False, False]).show() + +---+-----+ + |age| name| + +---+-----+ + | 5| Bob| + | 2| Bob| + | 2|Alice| + +---+-----+ + """ + jdf = self._jdf.sort(self._sort_cols(cols, kwargs)) + return DataFrame(jdf, self.sparkSession) + + orderBy = sort + + def _jseq( + self, + cols: Sequence, + converter: Optional[Callable[..., Union["PrimitiveType", JavaObject]]] = None, + ) -> JavaObject: + """Return a JVM Seq of Columns from a list of Column or names""" + return _to_seq(self.sparkSession._sc, cols, converter) + + def _jmap(self, jm: Dict) -> JavaObject: + """Return a JVM Scala Map from a dict""" + return _to_scala_map(self.sparkSession._sc, jm) + + def _jcols(self, *cols: "ColumnOrName") -> JavaObject: + """Return a JVM Seq of Columns from a list of Column or column names + + If `cols` has only one list in it, cols[0] will be used as the list. + """ + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] + return self._jseq(cols, _to_java_column) + + def _sort_cols( + self, cols: Sequence[Union[str, Column, List[Union[str, Column]]]], kwargs: Dict[str, Any] + ) -> JavaObject: + """Return a JVM Seq of Columns that describes the sort order""" + if not cols: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "column"}, + ) + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] + jcols = [_to_java_column(cast("ColumnOrName", c)) for c in cols] + ascending = kwargs.get("ascending", True) + if isinstance(ascending, (bool, int)): + if not ascending: + jcols = [jc.desc() for jc in jcols] + elif isinstance(ascending, list): + jcols = [jc if asc else jc.desc() for asc, jc in zip(ascending, jcols)] + else: + raise PySparkTypeError( + error_class="NOT_BOOL_OR_LIST", + message_parameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, + ) + return self._jseq(jcols) + + def describe(self, *cols: Union[str, List[str]]) -> "DataFrame": + """Computes basic statistics for numeric and string columns. + + .. versionadded:: 1.3.1 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + This includes count, mean, stddev, min, and max. If no columns are + given, this function computes statistics for all numerical or string columns. + + Notes + ----- + This function is meant for exploratory data analysis, as we make no + guarantee about the backward compatibility of the schema of the resulting + :class:`DataFrame`. + + Use summary for expanded statistics and control over which statistics to compute. + + Parameters + ---------- + cols : str, list, optional + Column name or list of column names to describe by (default All columns). + + Returns + ------- + :class:`DataFrame` + A new DataFrame that describes (provides statistics) given DataFrame. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [("Bob", 13, 40.3, 150.5), ("Alice", 12, 37.8, 142.3), ("Tom", 11, 44.1, 142.2)], + ... ["name", "age", "weight", "height"], + ... ) + >>> df.describe(['age']).show() + +-------+----+ + |summary| age| + +-------+----+ + | count| 3| + | mean|12.0| + | stddev| 1.0| + | min| 11| + | max| 13| + +-------+----+ + + >>> df.describe(['age', 'weight', 'height']).show() + +-------+----+------------------+-----------------+ + |summary| age| weight| height| + +-------+----+------------------+-----------------+ + | count| 3| 3| 3| + | mean|12.0| 40.73333333333333| 145.0| + | stddev| 1.0|3.1722757341273704|4.763402145525822| + | min| 11| 37.8| 142.2| + | max| 13| 44.1| 150.5| + +-------+----+------------------+-----------------+ + + See Also + -------- + DataFrame.summary + """ + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] # type: ignore[assignment] + jdf = self._jdf.describe(self._jseq(cols)) + return DataFrame(jdf, self.sparkSession) + + def summary(self, *statistics: str) -> "DataFrame": + """Computes specified statistics for numeric and string columns. Available statistics are: + - count + - mean + - stddev + - min + - max + - arbitrary approximate percentiles specified as a percentage (e.g., 75%) + + If no statistics are given, this function computes count, mean, stddev, min, + approximate quartiles (percentiles at 25%, 50%, and 75%), and max. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + statistics : str, optional + Column names to calculate statistics by (default All columns). + + Returns + ------- + :class:`DataFrame` + A new DataFrame that provides statistics for the given DataFrame. + + Notes + ----- + This function is meant for exploratory data analysis, as we make no + guarantee about the backward compatibility of the schema of the resulting + :class:`DataFrame`. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [("Bob", 13, 40.3, 150.5), ("Alice", 12, 37.8, 142.3), ("Tom", 11, 44.1, 142.2)], + ... ["name", "age", "weight", "height"], + ... ) + >>> df.select("age", "weight", "height").summary().show() + +-------+----+------------------+-----------------+ + |summary| age| weight| height| + +-------+----+------------------+-----------------+ + | count| 3| 3| 3| + | mean|12.0| 40.73333333333333| 145.0| + | stddev| 1.0|3.1722757341273704|4.763402145525822| + | min| 11| 37.8| 142.2| + | 25%| 11| 37.8| 142.2| + | 50%| 12| 40.3| 142.3| + | 75%| 13| 44.1| 150.5| + | max| 13| 44.1| 150.5| + +-------+----+------------------+-----------------+ + + >>> df.select("age", "weight", "height").summary("count", "min", "25%", "75%", "max").show() + +-------+---+------+------+ + |summary|age|weight|height| + +-------+---+------+------+ + | count| 3| 3| 3| + | min| 11| 37.8| 142.2| + | 25%| 11| 37.8| 142.2| + | 75%| 13| 44.1| 150.5| + | max| 13| 44.1| 150.5| + +-------+---+------+------+ + + See Also + -------- + DataFrame.display + """ + if len(statistics) == 1 and isinstance(statistics[0], list): + statistics = statistics[0] + jdf = self._jdf.summary(self._jseq(statistics)) + return DataFrame(jdf, self.sparkSession) + + @overload + def head(self) -> Optional[Row]: + ... + + @overload + def head(self, n: int) -> List[Row]: + ... + + def head(self, n: Optional[int] = None) -> Union[Optional[Row], List[Row]]: + """Returns the first ``n`` rows. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + This method should only be used if the resulting array is expected + to be small, as all the data is loaded into the driver's memory. + + Parameters + ---------- + n : int, optional + default 1. Number of rows to return. + + Returns + ------- + If n is greater than 1, return a list of :class:`Row`. + If n is 1, return a single Row. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.head() + Row(age=2, name='Alice') + >>> df.head(1) + [Row(age=2, name='Alice')] + """ + if n is None: + rs = self.head(1) + return rs[0] if rs else None + return self.take(n) + + def first(self) -> Optional[Row]: + """Returns the first row as a :class:`Row`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`Row` + First row if :class:`DataFrame` is not empty, otherwise ``None``. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.first() + Row(age=2, name='Alice') + """ + return self.head() + + @overload + def __getitem__(self, item: Union[int, str]) -> Column: + ... + + @overload + def __getitem__(self, item: Union[Column, List, Tuple]) -> "DataFrame": + ... + + def __getitem__(self, item: Union[int, str, Column, List, Tuple]) -> Union[Column, "DataFrame"]: + """Returns the column as a :class:`Column`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + item : int, str, :class:`Column`, list or tuple + column index, column name, column, or a list or tuple of columns + + Returns + ------- + :class:`Column` or :class:`DataFrame` + a specified column, or a filtered or projected dataframe. + + * If the input `item` is an int or str, the output is a :class:`Column`. + + * If the input `item` is a :class:`Column`, the output is a :class:`DataFrame` + filtered by this given :class:`Column`. + + * If the input `item` is a list or tuple, the output is a :class:`DataFrame` + projected by this given list or tuple. + + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + + Retrieve a column instance. + + >>> df.select(df['age']).show() + +---+ + |age| + +---+ + | 2| + | 5| + +---+ + + >>> df.select(df[1]).show() + +-----+ + | name| + +-----+ + |Alice| + | Bob| + +-----+ + + Select multiple string columns as index. + + >>> df[["name", "age"]].show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 2| + | Bob| 5| + +-----+---+ + >>> df[df.age > 3].show() + +---+----+ + |age|name| + +---+----+ + | 5| Bob| + +---+----+ + >>> df[df[0] > 3].show() + +---+----+ + |age|name| + +---+----+ + | 5| Bob| + +---+----+ + """ + if isinstance(item, str): + jc = self._jdf.apply(item) + return Column(jc) + elif isinstance(item, Column): + return self.filter(item) + elif isinstance(item, (list, tuple)): + return self.select(*item) + elif isinstance(item, int): + jc = self._jdf.apply(self.columns[item]) + return Column(jc) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_FLOAT_OR_INT_OR_LIST_OR_STR", + message_parameters={"arg_name": "item", "arg_type": type(item).__name__}, + ) + + def __getattr__(self, name: str) -> Column: + """Returns the :class:`Column` denoted by ``name``. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + Column name to return as :class:`Column`. + + Returns + ------- + :class:`Column` + Requested column. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + + Retrieve a column instance. + + >>> df.select(df.age).show() + +---+ + |age| + +---+ + | 2| + | 5| + +---+ + """ + if name not in self.columns: + raise AttributeError( + "'%s' object has no attribute '%s'" % (self.__class__.__name__, name) + ) + jc = self._jdf.apply(name) + return Column(jc) + + def __dir__(self) -> List[str]: + """ + Examples + -------- + >>> from pyspark.sql.functions import lit + + Create a dataframe with a column named 'id'. + + >>> df = spark.range(3) + >>> [attr for attr in dir(df) if attr[0] == 'i'][:7] # Includes column id + ['id', 'inputFiles', 'intersect', 'intersectAll', 'isEmpty', 'isLocal', 'isStreaming'] + + Add a column named 'i_like_pancakes'. + + >>> df = df.withColumn('i_like_pancakes', lit(1)) + >>> [attr for attr in dir(df) if attr[0] == 'i'][:7] # Includes columns i_like_pancakes, id + ['i_like_pancakes', 'id', 'inputFiles', 'intersect', 'intersectAll', 'isEmpty', 'isLocal'] + + Try to add an existed column 'inputFiles'. + + >>> df = df.withColumn('inputFiles', lit(2)) + >>> [attr for attr in dir(df) if attr[0] == 'i'][:7] # Doesn't duplicate inputFiles + ['i_like_pancakes', 'id', 'inputFiles', 'intersect', 'intersectAll', 'isEmpty', 'isLocal'] + + Try to add a column named 'id2'. + + >>> df = df.withColumn('id2', lit(3)) + >>> [attr for attr in dir(df) if attr[0] == 'i'][:7] # result includes id2 and sorted + ['i_like_pancakes', 'id', 'id2', 'inputFiles', 'intersect', 'intersectAll', 'isEmpty'] + + Don't include columns that are not valid python identifiers. + + >>> df = df.withColumn('1', lit(4)) + >>> df = df.withColumn('name 1', lit(5)) + >>> [attr for attr in dir(df) if attr[0] == 'i'][:7] # Doesn't include 1 or name 1 + ['i_like_pancakes', 'id', 'id2', 'inputFiles', 'intersect', 'intersectAll', 'isEmpty'] + """ + attrs = set(super().__dir__()) + attrs.update(filter(lambda s: s.isidentifier(), self.columns)) + return sorted(attrs) + + @overload + def select(self, *cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def select(self, __cols: Union[List[Column], List[str]]) -> "DataFrame": + ... + + def select(self, *cols: "ColumnOrName") -> "DataFrame": # type: ignore[misc] + """Projects a set of expressions and returns a new :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str, :class:`Column`, or list + column names (string) or expressions (:class:`Column`). + If one of the column names is '*', that column is expanded to include all columns + in the current :class:`DataFrame`. + + Returns + ------- + :class:`DataFrame` + A DataFrame with subset (or all) of columns. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + + Select all columns in the DataFrame. + + >>> df.select('*').show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + | 5| Bob| + +---+-----+ + + Select a column with other expressions in the DataFrame. + + >>> df.select(df.name, (df.age + 10).alias('age')).show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 12| + | Bob| 15| + +-----+---+ + """ + jdf = self._jdf.select(self._jcols(*cols)) + return DataFrame(jdf, self.sparkSession) + + @overload + def selectExpr(self, *expr: str) -> "DataFrame": + ... + + @overload + def selectExpr(self, *expr: List[str]) -> "DataFrame": + ... + + def selectExpr(self, *expr: Union[str, List[str]]) -> "DataFrame": + """Projects a set of SQL expressions and returns a new :class:`DataFrame`. + + This is a variant of :func:`select` that accepts SQL expressions. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrame` + A DataFrame with new/old columns transformed by expressions. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.selectExpr("age * 2", "abs(age)").show() + +---------+--------+ + |(age * 2)|abs(age)| + +---------+--------+ + | 4| 2| + | 10| 5| + +---------+--------+ + """ + if len(expr) == 1 and isinstance(expr[0], list): + expr = expr[0] # type: ignore[assignment] + jdf = self._jdf.selectExpr(self._jseq(expr)) + return DataFrame(jdf, self.sparkSession) + + def filter(self, condition: "ColumnOrName") -> "DataFrame": + """Filters rows using the given condition. + + :func:`where` is an alias for :func:`filter`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + condition : :class:`Column` or str + a :class:`Column` of :class:`types.BooleanType` + or a string of SQL expressions. + + Returns + ------- + :class:`DataFrame` + Filtered DataFrame. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (5, "Bob")], schema=["age", "name"]) + + Filter by :class:`Column` instances. + + >>> df.filter(df.age > 3).show() + +---+----+ + |age|name| + +---+----+ + | 5| Bob| + +---+----+ + >>> df.where(df.age == 2).show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + +---+-----+ + + Filter by SQL expression in a string. + + >>> df.filter("age > 3").show() + +---+----+ + |age|name| + +---+----+ + | 5| Bob| + +---+----+ + >>> df.where("age = 2").show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + +---+-----+ + """ + if isinstance(condition, str): + jdf = self._jdf.filter(condition) + elif isinstance(condition, Column): + jdf = self._jdf.filter(condition._jc) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + return DataFrame(jdf, self.sparkSession) + + @overload + def groupBy(self, *cols: "ColumnOrName") -> "GroupedData": + ... + + @overload + def groupBy(self, __cols: Union[List[Column], List[str]]) -> "GroupedData": + ... + + def groupBy(self, *cols: "ColumnOrName") -> "GroupedData": # type: ignore[misc] + """Groups the :class:`DataFrame` using the specified columns, + so we can run aggregation on them. See :class:`GroupedData` + for all the available aggregate functions. + + :func:`groupby` is an alias for :func:`groupBy`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : list, str or :class:`Column` + columns to group by. + Each element should be a column name (string) or an expression (:class:`Column`) + or list of them. + + Returns + ------- + :class:`GroupedData` + Grouped data by given columns. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice"), (2, "Bob"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]) + + Empty grouping columns triggers a global aggregation. + + >>> df.groupBy().avg().show() + +--------+ + |avg(age)| + +--------+ + | 2.75| + +--------+ + + Group-by 'name', and specify a dictionary to calculate the summation of 'age'. + + >>> df.groupBy("name").agg({"age": "sum"}).sort("name").show() + +-----+--------+ + | name|sum(age)| + +-----+--------+ + |Alice| 2| + | Bob| 9| + +-----+--------+ + + Group-by 'name', and calculate maximum values. + + >>> df.groupBy(df.name).max().sort("name").show() + +-----+--------+ + | name|max(age)| + +-----+--------+ + |Alice| 2| + | Bob| 5| + +-----+--------+ + + Group-by 'name' and 'age', and calculate the number of rows in each group. + + >>> df.groupBy(["name", df.age]).count().sort("name", "age").show() + +-----+---+-----+ + | name|age|count| + +-----+---+-----+ + |Alice| 2| 1| + | Bob| 2| 2| + | Bob| 5| 1| + +-----+---+-----+ + """ + jgd = self._jdf.groupBy(self._jcols(*cols)) + from pyspark.sql.group import GroupedData + + return GroupedData(jgd, self) + + @overload + def rollup(self, *cols: "ColumnOrName") -> "GroupedData": + ... + + @overload + def rollup(self, __cols: Union[List[Column], List[str]]) -> "GroupedData": + ... + + def rollup(self, *cols: "ColumnOrName") -> "GroupedData": # type: ignore[misc] + """ + Create a multi-dimensional rollup for the current :class:`DataFrame` using + the specified columns, so we can run aggregation on them. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : list, str or :class:`Column` + Columns to roll-up by. + Each element should be a column name (string) or an expression (:class:`Column`) + or list of them. + + Returns + ------- + :class:`GroupedData` + Rolled-up data by given columns. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.rollup("name", df.age).count().orderBy("name", "age").show() + +-----+----+-----+ + | name| age|count| + +-----+----+-----+ + | NULL|NULL| 2| + |Alice|NULL| 1| + |Alice| 2| 1| + | Bob|NULL| 1| + | Bob| 5| 1| + +-----+----+-----+ + """ + jgd = self._jdf.rollup(self._jcols(*cols)) + from pyspark.sql.group import GroupedData + + return GroupedData(jgd, self) + + @overload + def cube(self, *cols: "ColumnOrName") -> "GroupedData": + ... + + @overload + def cube(self, __cols: Union[List[Column], List[str]]) -> "GroupedData": + ... + + def cube(self, *cols: "ColumnOrName") -> "GroupedData": # type: ignore[misc] + """ + Create a multi-dimensional cube for the current :class:`DataFrame` using + the specified columns, so we can run aggregations on them. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : list, str or :class:`Column` + columns to create cube by. + Each element should be a column name (string) or an expression (:class:`Column`) + or list of them. + + Returns + ------- + :class:`GroupedData` + Cube of the data by given columns. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.cube("name", df.age).count().orderBy("name", "age").show() + +-----+----+-----+ + | name| age|count| + +-----+----+-----+ + | NULL|NULL| 2| + | NULL| 2| 1| + | NULL| 5| 1| + |Alice|NULL| 1| + |Alice| 2| 1| + | Bob|NULL| 1| + | Bob| 5| 1| + +-----+----+-----+ + """ + jgd = self._jdf.cube(self._jcols(*cols)) + from pyspark.sql.group import GroupedData + + return GroupedData(jgd, self) + + def unpivot( + self, + ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]], + values: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]], + variableColumnName: str, + valueColumnName: str, + ) -> "DataFrame": + """ + Unpivot a DataFrame from wide format to long format, optionally leaving + identifier columns set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, + except for the aggregation, which cannot be reversed. + + This function is useful to massage a DataFrame into a format where some + columns are identifier columns ("ids"), while all other columns ("values") + are "unpivoted" to the rows, leaving just two non-id columns, named as given + by `variableColumnName` and `valueColumnName`. + + When no "id" columns are given, the unpivoted DataFrame consists of only the + "variable" and "value" columns. + + The `values` columns must not be empty so at least one value must be given to be unpivoted. + When `values` is `None`, all non-id columns will be unpivoted. + + All "value" columns must share a least common data type. Unless they are the same data type, + all "value" columns are cast to the nearest common data type. For instance, types + `IntegerType` and `LongType` are cast to `LongType`, while `IntegerType` and `StringType` + do not have a common data type and `unpivot` fails. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + ids : str, Column, tuple, list + Column(s) to use as identifiers. Can be a single column or column name, + or a list or tuple for multiple columns. + values : str, Column, tuple, list, optional + Column(s) to unpivot. Can be a single column or column name, or a list or tuple + for multiple columns. If specified, must not be empty. If not specified, uses all + columns that are not set as `ids`. + variableColumnName : str + Name of the variable column. + valueColumnName : str + Name of the value column. + + Returns + ------- + :class:`DataFrame` + Unpivoted DataFrame. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(1, 11, 1.1), (2, 12, 1.2)], + ... ["id", "int", "double"], + ... ) + >>> df.show() + +---+---+------+ + | id|int|double| + +---+---+------+ + | 1| 11| 1.1| + | 2| 12| 1.2| + +---+---+------+ + + >>> df.unpivot("id", ["int", "double"], "var", "val").show() + +---+------+----+ + | id| var| val| + +---+------+----+ + | 1| int|11.0| + | 1|double| 1.1| + | 2| int|12.0| + | 2|double| 1.2| + +---+------+----+ + + See Also + -------- + DataFrame.melt + """ + assert ids is not None, "ids must not be None" + + def to_jcols( + cols: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]] + ) -> JavaObject: + if isinstance(cols, list): + return self._jcols(*cols) + if isinstance(cols, tuple): + return self._jcols(*list(cols)) + return self._jcols(cols) + + jids = to_jcols(ids) + if values is None: + jdf = self._jdf.unpivotWithSeq(jids, variableColumnName, valueColumnName) + else: + jvals = to_jcols(values) + jdf = self._jdf.unpivotWithSeq(jids, jvals, variableColumnName, valueColumnName) + + return DataFrame(jdf, self.sparkSession) + + def melt( + self, + ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]], + values: Optional[Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]]], + variableColumnName: str, + valueColumnName: str, + ) -> "DataFrame": + """ + Unpivot a DataFrame from wide format to long format, optionally leaving + identifier columns set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, + except for the aggregation, which cannot be reversed. + + :func:`melt` is an alias for :func:`unpivot`. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + ids : str, Column, tuple, list, optional + Column(s) to use as identifiers. Can be a single column or column name, + or a list or tuple for multiple columns. + values : str, Column, tuple, list, optional + Column(s) to unpivot. Can be a single column or column name, or a list or tuple + for multiple columns. If not specified or empty, use all columns that + are not set as `ids`. + variableColumnName : str + Name of the variable column. + valueColumnName : str + Name of the value column. + + Returns + ------- + :class:`DataFrame` + Unpivoted DataFrame. + + See Also + -------- + DataFrame.unpivot + + Notes + ----- + Supports Spark Connect. + """ + return self.unpivot(ids, values, variableColumnName, valueColumnName) + + def agg(self, *exprs: Union[Column, Dict[str, str]]) -> "DataFrame": + """Aggregate on the entire :class:`DataFrame` without groups + (shorthand for ``df.groupBy().agg()``). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + exprs : :class:`Column` or dict of key and value strings + Columns or expressions to aggregate DataFrame by. + + Returns + ------- + :class:`DataFrame` + Aggregated DataFrame. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.agg({"age": "max"}).show() + +--------+ + |max(age)| + +--------+ + | 5| + +--------+ + >>> df.agg(sf.min(df.age)).show() + +--------+ + |min(age)| + +--------+ + | 2| + +--------+ + """ + return self.groupBy().agg(*exprs) # type: ignore[arg-type] + + def observe( + self, + observation: Union["Observation", str], + *exprs: Column, + ) -> "DataFrame": + """Define (named) metrics to observe on the DataFrame. This method returns an 'observed' + DataFrame that returns the same result as the input, with the following guarantees: + + * It will compute the defined aggregates (metrics) on all the data that is flowing through + the Dataset at that point. + + * It will report the value of the defined aggregate columns as soon as we reach a completion + point. A completion point is either the end of a query (batch mode) or the end of a + streaming epoch. The value of the aggregates only reflects the data processed since + the previous completion point. + + The metrics columns must either contain a literal (e.g. lit(42)), or should contain one or + more aggregate functions (e.g. sum(a) or sum(a + b) + avg(c) - lit(1)). Expressions that + contain references to the input Dataset's columns must always be wrapped in an aggregate + function. + + A user can observe these metrics by adding + Python's :class:`~pyspark.sql.streaming.StreamingQueryListener`, + Scala/Java's ``org.apache.spark.sql.streaming.StreamingQueryListener`` or Scala/Java's + ``org.apache.spark.sql.util.QueryExecutionListener`` to the spark session. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + observation : :class:`Observation` or str + `str` to specify the name, or an :class:`Observation` instance to obtain the metric. + + .. versionchanged:: 3.4.0 + Added support for `str` in this parameter. + exprs : :class:`Column` + column expressions (:class:`Column`). + + Returns + ------- + :class:`DataFrame` + the observed :class:`DataFrame`. + + Notes + ----- + When ``observation`` is :class:`Observation`, this method only supports batch queries. + When ``observation`` is a string, this method works for both batch and streaming queries. + Continuous execution is currently not supported yet. + + Examples + -------- + When ``observation`` is :class:`Observation`, only batch queries work as below. + + >>> from pyspark.sql.functions import col, count, lit, max + >>> from pyspark.sql import Observation + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> observation = Observation("my metrics") + >>> observed_df = df.observe(observation, count(lit(1)).alias("count"), max(col("age"))) + >>> observed_df.count() + 2 + >>> observation.get + {'count': 2, 'max(age)': 5} + + When ``observation`` is a string, streaming queries also work as below. + + >>> from pyspark.sql.streaming import StreamingQueryListener + >>> class MyErrorListener(StreamingQueryListener): + ... def onQueryStarted(self, event): + ... pass + ... + ... def onQueryProgress(self, event): + ... row = event.progress.observedMetrics.get("my_event") + ... # Trigger if the number of errors exceeds 5 percent + ... num_rows = row.rc + ... num_error_rows = row.erc + ... ratio = num_error_rows / num_rows + ... if ratio > 0.05: + ... # Trigger alert + ... pass + ... + ... def onQueryIdle(self, event): + ... pass + ... + ... def onQueryTerminated(self, event): + ... pass + ... + >>> spark.streams.addListener(MyErrorListener()) + >>> # Observe row count (rc) and error row count (erc) in the streaming Dataset + ... observed_ds = df.observe( + ... "my_event", + ... count(lit(1)).alias("rc"), + ... count(col("error")).alias("erc")) # doctest: +SKIP + >>> observed_ds.writeStream.format("console").start() # doctest: +SKIP + """ + from pyspark.sql import Observation + + if len(exprs) == 0: + raise PySparkValueError( + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "exprs"}, + ) + if not all(isinstance(c, Column) for c in exprs): + raise PySparkTypeError( + error_class="NOT_LIST_OF_COLUMN", + message_parameters={"arg_name": "exprs"}, + ) + + if isinstance(observation, Observation): + return observation._on(self, *exprs) + elif isinstance(observation, str): + return DataFrame( + self._jdf.observe( + observation, exprs[0]._jc, _to_seq(self._sc, [c._jc for c in exprs[1:]]) + ), + self.sparkSession, + ) + else: + raise PySparkTypeError( + error_class="NOT_LIST_OF_COLUMN", + message_parameters={ + "arg_name": "observation", + "arg_type": type(observation).__name__, + }, + ) + + def union(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing the union of rows in this and another + :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be unioned. + + Returns + ------- + :class:`DataFrame` + A new :class:`DataFrame` containing the combined rows with corresponding columns. + + See Also + -------- + DataFrame.unionAll + + Notes + ----- + This method performs a SQL-style set union of the rows from both `DataFrame` objects, + with no automatic deduplication of elements. + + Use the `distinct()` method to perform deduplication of rows. + + The method resolves columns by position (not by name), following the standard behavior + in SQL. + + Examples + -------- + Example 1: Combining two DataFrames with the same schema + + >>> df1 = spark.createDataFrame([(1, 'A'), (2, 'B')], ['id', 'value']) + >>> df2 = spark.createDataFrame([(3, 'C'), (4, 'D')], ['id', 'value']) + >>> df3 = df1.union(df2) + >>> df3.show() + +---+-----+ + | id|value| + +---+-----+ + | 1| A| + | 2| B| + | 3| C| + | 4| D| + +---+-----+ + + Example 2: Combining two DataFrames with different schemas + + >>> from pyspark.sql.functions import lit + >>> df1 = spark.createDataFrame([("Alice", 1), ("Bob", 2)], ["name", "id"]) + >>> df2 = spark.createDataFrame([(3, "Charlie"), (4, "Dave")], ["id", "name"]) + >>> df1 = df1.withColumn("age", lit(30)) + >>> df2 = df2.withColumn("age", lit(40)) + >>> df3 = df1.union(df2) + >>> df3.show() + +-----+-------+---+ + | name| id|age| + +-----+-------+---+ + |Alice| 1| 30| + | Bob| 2| 30| + | 3|Charlie| 40| + | 4| Dave| 40| + +-----+-------+---+ + + Example 3: Combining two DataFrames with mismatched columns + + >>> df1 = spark.createDataFrame([(1, 2)], ["A", "B"]) + >>> df2 = spark.createDataFrame([(3, 4)], ["C", "D"]) + >>> df3 = df1.union(df2) + >>> df3.show() + +---+---+ + | A| B| + +---+---+ + | 1| 2| + | 3| 4| + +---+---+ + + Example 4: Combining duplicate rows from two different DataFrames + + >>> df1 = spark.createDataFrame([(1, 'A'), (2, 'B'), (3, 'C')], ['id', 'value']) + >>> df2 = spark.createDataFrame([(3, 'C'), (4, 'D')], ['id', 'value']) + >>> df3 = df1.union(df2).distinct().sort("id") + >>> df3.show() + +---+-----+ + | id|value| + +---+-----+ + | 1| A| + | 2| B| + | 3| C| + | 4| D| + +---+-----+ + """ + return DataFrame(self._jdf.union(other._jdf), self.sparkSession) + + def unionAll(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing the union of rows in this and another + :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be combined + + Returns + ------- + :class:`DataFrame` + A new :class:`DataFrame` containing combined rows from both dataframes. + + Notes + ----- + This method combines all rows from both `DataFrame` objects with no automatic + deduplication of elements. + + Use the `distinct()` method to perform deduplication of rows. + + :func:`unionAll` is an alias to :func:`union` + + See Also + -------- + DataFrame.union + """ + return self.union(other) + + def unionByName(self, other: "DataFrame", allowMissingColumns: bool = False) -> "DataFrame": + """Returns a new :class:`DataFrame` containing union of rows in this and another + :class:`DataFrame`. + + This method performs a union operation on both input DataFrames, resolving columns by + name (rather than position). When `allowMissingColumns` is True, missing columns will + be filled with null. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be combined. + allowMissingColumns : bool, optional, default False + Specify whether to allow missing columns. + + .. versionadded:: 3.1.0 + + Returns + ------- + :class:`DataFrame` + A new :class:`DataFrame` containing the combined rows with corresponding + columns of the two given DataFrames. + + Examples + -------- + Example 1: Union of two DataFrames with same columns in different order. + + >>> df1 = spark.createDataFrame([[1, 2, 3]], ["col0", "col1", "col2"]) + >>> df2 = spark.createDataFrame([[4, 5, 6]], ["col1", "col2", "col0"]) + >>> df1.unionByName(df2).show() + +----+----+----+ + |col0|col1|col2| + +----+----+----+ + | 1| 2| 3| + | 6| 4| 5| + +----+----+----+ + + Example 2: Union with missing columns and setting `allowMissingColumns=True`. + + >>> df1 = spark.createDataFrame([[1, 2, 3]], ["col0", "col1", "col2"]) + >>> df2 = spark.createDataFrame([[4, 5, 6]], ["col1", "col2", "col3"]) + >>> df1.unionByName(df2, allowMissingColumns=True).show() + +----+----+----+----+ + |col0|col1|col2|col3| + +----+----+----+----+ + | 1| 2| 3|NULL| + |NULL| 4| 5| 6| + +----+----+----+----+ + + Example 3: Union of two DataFrames with few common columns. + + >>> df1 = spark.createDataFrame([[1, 2, 3]], ["col0", "col1", "col2"]) + >>> df2 = spark.createDataFrame([[4, 5, 6, 7]], ["col1", "col2", "col3", "col4"]) + >>> df1.unionByName(df2, allowMissingColumns=True).show() + +----+----+----+----+----+ + |col0|col1|col2|col3|col4| + +----+----+----+----+----+ + | 1| 2| 3|NULL|NULL| + |NULL| 4| 5| 6| 7| + +----+----+----+----+----+ + + Example 4: Union of two DataFrames with completely different columns. + + >>> df1 = spark.createDataFrame([[0, 1, 2]], ["col0", "col1", "col2"]) + >>> df2 = spark.createDataFrame([[3, 4, 5]], ["col3", "col4", "col5"]) + >>> df1.unionByName(df2, allowMissingColumns=True).show() + +----+----+----+----+----+----+ + |col0|col1|col2|col3|col4|col5| + +----+----+----+----+----+----+ + | 0| 1| 2|NULL|NULL|NULL| + |NULL|NULL|NULL| 3| 4| 5| + +----+----+----+----+----+----+ + """ + return DataFrame(self._jdf.unionByName(other._jdf, allowMissingColumns), self.sparkSession) + + def intersect(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing rows only in + both this :class:`DataFrame` and another :class:`DataFrame`. + Note that any duplicates are removed. To preserve duplicates + use :func:`intersectAll`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be combined. + + Returns + ------- + :class:`DataFrame` + Combined DataFrame. + + Notes + ----- + This is equivalent to `INTERSECT` in SQL. + + Examples + -------- + >>> df1 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3), ("c", 4)], ["C1", "C2"]) + >>> df2 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3)], ["C1", "C2"]) + >>> df1.intersect(df2).sort(df1.C1.desc()).show() + +---+---+ + | C1| C2| + +---+---+ + | b| 3| + | a| 1| + +---+---+ + """ + return DataFrame(self._jdf.intersect(other._jdf), self.sparkSession) + + def intersectAll(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing rows in both this :class:`DataFrame` + and another :class:`DataFrame` while preserving duplicates. + + This is equivalent to `INTERSECT ALL` in SQL. As standard in SQL, this function + resolves columns by position (not by name). + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be combined. + + Returns + ------- + :class:`DataFrame` + Combined DataFrame. + + Examples + -------- + >>> df1 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3), ("c", 4)], ["C1", "C2"]) + >>> df2 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3)], ["C1", "C2"]) + >>> df1.intersectAll(df2).sort("C1", "C2").show() + +---+---+ + | C1| C2| + +---+---+ + | a| 1| + | a| 1| + | b| 3| + +---+---+ + """ + return DataFrame(self._jdf.intersectAll(other._jdf), self.sparkSession) + + def subtract(self, other: "DataFrame") -> "DataFrame": + """Return a new :class:`DataFrame` containing rows in this :class:`DataFrame` + but not in another :class:`DataFrame`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + other : :class:`DataFrame` + Another :class:`DataFrame` that needs to be subtracted. + + Returns + ------- + :class:`DataFrame` + Subtracted DataFrame. + + Notes + ----- + This is equivalent to `EXCEPT DISTINCT` in SQL. + + Examples + -------- + >>> df1 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3), ("c", 4)], ["C1", "C2"]) + >>> df2 = spark.createDataFrame([("a", 1), ("a", 1), ("b", 3)], ["C1", "C2"]) + >>> df1.subtract(df2).show() + +---+---+ + | C1| C2| + +---+---+ + | c| 4| + +---+---+ + """ + return DataFrame(getattr(self._jdf, "except")(other._jdf), self.sparkSession) + + def dropDuplicates(self, subset: Optional[List[str]] = None) -> "DataFrame": + """Return a new :class:`DataFrame` with duplicate rows removed, + optionally only considering certain columns. + + For a static batch :class:`DataFrame`, it just drops duplicate rows. For a streaming + :class:`DataFrame`, it will keep all data across triggers as intermediate state to drop + duplicates rows. You can use :func:`withWatermark` to limit how late the duplicate data can + be and the system will accordingly limit the state. In addition, data older than + watermark will be dropped to avoid any possibility of duplicates. + + :func:`drop_duplicates` is an alias for :func:`dropDuplicates`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + subset : List of column names, optional + List of columns to use for duplicate comparison (default All columns). + + Returns + ------- + :class:`DataFrame` + DataFrame without duplicates. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([ + ... Row(name='Alice', age=5, height=80), + ... Row(name='Alice', age=5, height=80), + ... Row(name='Alice', age=10, height=80) + ... ]) + + Deduplicate the same rows. + + >>> df.dropDuplicates().show() + +-----+---+------+ + | name|age|height| + +-----+---+------+ + |Alice| 5| 80| + |Alice| 10| 80| + +-----+---+------+ + + Deduplicate values on 'name' and 'height' columns. + + >>> df.dropDuplicates(['name', 'height']).show() + +-----+---+------+ + | name|age|height| + +-----+---+------+ + |Alice| 5| 80| + +-----+---+------+ + """ + if subset is not None and (not isinstance(subset, Iterable) or isinstance(subset, str)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if subset is None: + jdf = self._jdf.dropDuplicates() + else: + jdf = self._jdf.dropDuplicates(self._jseq(subset)) + return DataFrame(jdf, self.sparkSession) + + def dropDuplicatesWithinWatermark(self, subset: Optional[List[str]] = None) -> "DataFrame": + """Return a new :class:`DataFrame` with duplicate rows removed, + optionally only considering certain columns, within watermark. + + This only works with streaming :class:`DataFrame`, and watermark for the input + :class:`DataFrame` must be set via :func:`withWatermark`. + + For a streaming :class:`DataFrame`, this will keep all data across triggers as intermediate + state to drop duplicated rows. The state will be kept to guarantee the semantic, "Events + are deduplicated as long as the time distance of earliest and latest events are smaller + than the delay threshold of watermark." Users are encouraged to set the delay threshold of + watermark longer than max timestamp differences among duplicated events. + + Note: too late data older than watermark will be dropped. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + subset : List of column names, optional + List of columns to use for duplicate comparison (default All columns). + + Returns + ------- + :class:`DataFrame` + DataFrame without duplicates. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import timestamp_seconds + >>> df = spark.readStream.format("rate").load().selectExpr( + ... "value % 5 AS value", "timestamp") + >>> df.select("value", df.timestamp.alias("time")).withWatermark("time", '10 minutes') + DataFrame[value: bigint, time: timestamp] + + Deduplicate the same rows. + + >>> df.dropDuplicatesWithinWatermark() # doctest: +SKIP + + Deduplicate values on 'value' columns. + + >>> df.dropDuplicatesWithinWatermark(['value']) # doctest: +SKIP + """ + if subset is not None and (not isinstance(subset, Iterable) or isinstance(subset, str)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if subset is None: + jdf = self._jdf.dropDuplicatesWithinWatermark() + else: + jdf = self._jdf.dropDuplicatesWithinWatermark(self._jseq(subset)) + return DataFrame(jdf, self.sparkSession) + + def dropna( + self, + how: str = "any", + thresh: Optional[int] = None, + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> "DataFrame": + """Returns a new :class:`DataFrame` omitting rows with null values. + :func:`DataFrame.dropna` and :func:`DataFrameNaFunctions.drop` are aliases of each other. + + .. versionadded:: 1.3.1 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + how : str, optional + 'any' or 'all'. + If 'any', drop a row if it contains any nulls. + If 'all', drop a row only if all its values are null. + thresh: int, optional + default None + If specified, drop rows that have less than `thresh` non-null values. + This overwrites the `how` parameter. + subset : str, tuple or list, optional + optional list of column names to consider. + + Returns + ------- + :class:`DataFrame` + DataFrame with null only rows excluded. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([ + ... Row(age=10, height=80, name="Alice"), + ... Row(age=5, height=None, name="Bob"), + ... Row(age=None, height=None, name="Tom"), + ... Row(age=None, height=None, name=None), + ... ]) + >>> df.na.drop().show() + +---+------+-----+ + |age|height| name| + +---+------+-----+ + | 10| 80|Alice| + +---+------+-----+ + """ + if how is not None and how not in ["any", "all"]: + raise PySparkValueError( + error_class="VALUE_NOT_ANY_OR_ALL", + message_parameters={"arg_name": "how", "arg_type": how}, + ) + + if subset is None: + subset = self.columns + elif isinstance(subset, str): + subset = [subset] + elif not isinstance(subset, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + if thresh is None: + thresh = len(subset) if how == "any" else 1 + + return DataFrame(self._jdf.na().drop(thresh, self._jseq(subset)), self.sparkSession) + + @overload + def fillna( + self, + value: "LiteralType", + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = ..., + ) -> "DataFrame": + ... + + @overload + def fillna(self, value: Dict[str, "LiteralType"]) -> "DataFrame": + ... + + def fillna( + self, + value: Union["LiteralType", Dict[str, "LiteralType"]], + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> "DataFrame": + """Replace null values, alias for ``na.fill()``. + :func:`DataFrame.fillna` and :func:`DataFrameNaFunctions.fill` are aliases of each other. + + .. versionadded:: 1.3.1 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + value : int, float, string, bool or dict + Value to replace null values with. + If the value is a dict, then `subset` is ignored and `value` must be a mapping + from column name (string) to replacement value. The replacement value must be + an int, float, boolean, or string. + subset : str, tuple or list, optional + optional list of column names to consider. + Columns specified in subset that do not have matching data types are ignored. + For example, if `value` is a string, and subset contains a non-string column, + then the non-string column is simply ignored. + + Returns + ------- + :class:`DataFrame` + DataFrame with replaced null values. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (10, 80.5, "Alice", None), + ... (5, None, "Bob", None), + ... (None, None, "Tom", None), + ... (None, None, None, True)], + ... schema=["age", "height", "name", "bool"]) + + Fill all null values with 50 for numeric columns. + + >>> df.na.fill(50).show() + +---+------+-----+----+ + |age|height| name|bool| + +---+------+-----+----+ + | 10| 80.5|Alice|NULL| + | 5| 50.0| Bob|NULL| + | 50| 50.0| Tom|NULL| + | 50| 50.0| NULL|true| + +---+------+-----+----+ + + Fill all null values with ``False`` for boolean columns. + + >>> df.na.fill(False).show() + +----+------+-----+-----+ + | age|height| name| bool| + +----+------+-----+-----+ + | 10| 80.5|Alice|false| + | 5| NULL| Bob|false| + |NULL| NULL| Tom|false| + |NULL| NULL| NULL| true| + +----+------+-----+-----+ + + Fill all null values with to 50 and "unknown" for 'age' and 'name' column respectively. + + >>> df.na.fill({'age': 50, 'name': 'unknown'}).show() + +---+------+-------+----+ + |age|height| name|bool| + +---+------+-------+----+ + | 10| 80.5| Alice|NULL| + | 5| NULL| Bob|NULL| + | 50| NULL| Tom|NULL| + | 50| NULL|unknown|true| + +---+------+-------+----+ + """ + if not isinstance(value, (float, int, str, bool, dict)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_STR", + message_parameters={"arg_name": "value", "arg_type": type(value).__name__}, + ) + + # Note that bool validates isinstance(int), but we don't want to + # convert bools to floats + + if not isinstance(value, bool) and isinstance(value, int): + value = float(value) + + if isinstance(value, dict): + return DataFrame(self._jdf.na().fill(value), self.sparkSession) + elif subset is None: + return DataFrame(self._jdf.na().fill(value), self.sparkSession) + else: + if isinstance(subset, str): + subset = [subset] + elif not isinstance(subset, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + return DataFrame(self._jdf.na().fill(value, self._jseq(subset)), self.sparkSession) + + @overload + def replace( + self, + to_replace: "LiteralType", + value: "OptionalPrimitiveType", + subset: Optional[List[str]] = ..., + ) -> "DataFrame": + ... + + @overload + def replace( + self, + to_replace: List["LiteralType"], + value: List["OptionalPrimitiveType"], + subset: Optional[List[str]] = ..., + ) -> "DataFrame": + ... + + @overload + def replace( + self, + to_replace: Dict["LiteralType", "OptionalPrimitiveType"], + subset: Optional[List[str]] = ..., + ) -> "DataFrame": + ... + + @overload + def replace( + self, + to_replace: List["LiteralType"], + value: "OptionalPrimitiveType", + subset: Optional[List[str]] = ..., + ) -> "DataFrame": + ... + + def replace( # type: ignore[misc] + self, + to_replace: Union[ + "LiteralType", List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"] + ], + value: Optional[ + Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType] + ] = _NoValue, + subset: Optional[List[str]] = None, + ) -> "DataFrame": + """Returns a new :class:`DataFrame` replacing a value with another value. + :func:`DataFrame.replace` and :func:`DataFrameNaFunctions.replace` are + aliases of each other. + Values to_replace and value must have the same type and can only be numerics, booleans, + or strings. Value can have None. When replacing, the new value will be cast + to the type of the existing column. + For numeric replacements all values to be replaced should have unique + floating point representation. In case of conflicts (for example with `{42: -1, 42.0: 1}`) + and arbitrary replacement will be used. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + to_replace : bool, int, float, string, list or dict + Value to be replaced. + If the value is a dict, then `value` is ignored or can be omitted, and `to_replace` + must be a mapping between a value and a replacement. + value : bool, int, float, string or None, optional + The replacement value must be a bool, int, float, string or None. If `value` is a + list, `value` should be of the same length and type as `to_replace`. + If `value` is a scalar and `to_replace` is a sequence, then `value` is + used as a replacement for each item in `to_replace`. + subset : list, optional + optional list of column names to consider. + Columns specified in subset that do not have matching data types are ignored. + For example, if `value` is a string, and subset contains a non-string column, + then the non-string column is simply ignored. + + Returns + ------- + :class:`DataFrame` + DataFrame with replaced values. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (10, 80, "Alice"), + ... (5, None, "Bob"), + ... (None, 10, "Tom"), + ... (None, None, None)], + ... schema=["age", "height", "name"]) + + Replace 10 to 20 in all columns. + + >>> df.na.replace(10, 20).show() + +----+------+-----+ + | age|height| name| + +----+------+-----+ + | 20| 80|Alice| + | 5| NULL| Bob| + |NULL| 20| Tom| + |NULL| NULL| NULL| + +----+------+-----+ + + Replace 'Alice' to null in all columns. + + >>> df.na.replace('Alice', None).show() + +----+------+----+ + | age|height|name| + +----+------+----+ + | 10| 80|NULL| + | 5| NULL| Bob| + |NULL| 10| Tom| + |NULL| NULL|NULL| + +----+------+----+ + + Replace 'Alice' to 'A', and 'Bob' to 'B' in the 'name' column. + + >>> df.na.replace(['Alice', 'Bob'], ['A', 'B'], 'name').show() + +----+------+----+ + | age|height|name| + +----+------+----+ + | 10| 80| A| + | 5| NULL| B| + |NULL| 10| Tom| + |NULL| NULL|NULL| + +----+------+----+ + """ + if value is _NoValue: + if isinstance(to_replace, dict): + value = None + else: + raise PySparkTypeError( + error_class="ARGUMENT_REQUIRED", + message_parameters={"arg_name": "value", "condition": "`to_replace` is dict"}, + ) + + # Helper functions + def all_of(types: Union[Type, Tuple[Type, ...]]) -> Callable[[Iterable], bool]: + """Given a type or tuple of types and a sequence of xs + check if each x is instance of type(s) + + >>> all_of(bool)([True, False]) + True + >>> all_of(str)(["a", 1]) + False + """ + + def all_of_(xs: Iterable) -> bool: + return all(isinstance(x, types) for x in xs) + + return all_of_ + + all_of_bool = all_of(bool) + all_of_str = all_of(str) + all_of_numeric = all_of((float, int)) + + # Validate input types + valid_types = (bool, float, int, str, list, tuple) + if not isinstance(to_replace, valid_types + (dict,)): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "to_replace", + "arg_type": type(to_replace).__name__, + }, + ) + + if ( + not isinstance(value, valid_types) + and value is not None + and not isinstance(to_replace, dict) + ): + raise PySparkTypeError( + error_class="NOT_BOOL_OR_FLOAT_OR_INT_OR_LIST_OR_NONE_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "value", + "arg_type": type(value).__name__, + }, + ) + + if isinstance(to_replace, (list, tuple)) and isinstance(value, (list, tuple)): + if len(to_replace) != len(value): + raise PySparkValueError( + error_class="LENGTH_SHOULD_BE_THE_SAME", + message_parameters={ + "arg1": "to_replace", + "arg2": "value", + "arg1_length": str(len(to_replace)), + "arg2_length": str(len(value)), + }, + ) + + if not (subset is None or isinstance(subset, (list, tuple, str))): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": type(subset).__name__}, + ) + + # Reshape input arguments if necessary + if isinstance(to_replace, (float, int, str)): + to_replace = [to_replace] + + if isinstance(to_replace, dict): + rep_dict = to_replace + if value is not None: + warnings.warn("to_replace is a dict and value is not None. value will be ignored.") + else: + if isinstance(value, (float, int, str)) or value is None: + value = [value for _ in range(len(to_replace))] + rep_dict = dict(zip(to_replace, cast("Iterable[Optional[Union[float, str]]]", value))) + + if isinstance(subset, str): + subset = [subset] + + # Verify we were not passed in mixed type generics. + if not any( + all_of_type(rep_dict.keys()) + and all_of_type(x for x in rep_dict.values() if x is not None) + for all_of_type in [all_of_bool, all_of_str, all_of_numeric] + ): + raise PySparkValueError( + error_class="MIXED_TYPE_REPLACEMENT", + message_parameters={}, + ) + + if subset is None: + return DataFrame(self._jdf.na().replace("*", rep_dict), self.sparkSession) + else: + return DataFrame( + self._jdf.na().replace(self._jseq(subset), self._jmap(rep_dict)), + self.sparkSession, + ) + + @overload + def approxQuantile( + self, + col: str, + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> List[float]: + ... + + @overload + def approxQuantile( + self, + col: Union[List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> List[List[float]]: + ... + + def approxQuantile( + self, + col: Union[str, List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> Union[List[float], List[List[float]]]: + """ + Calculates the approximate quantiles of numerical columns of a + :class:`DataFrame`. + + The result of this algorithm has the following deterministic bound: + If the :class:`DataFrame` has N elements and if we request the quantile at + probability `p` up to error `err`, then the algorithm will return + a sample `x` from the :class:`DataFrame` so that the *exact* rank of `x` is + close to (p * N). More precisely, + + floor((p - err) * N) <= rank(x) <= ceil((p + err) * N). + + This method implements a variation of the Greenwald-Khanna + algorithm (with some speed optimizations). The algorithm was first + present in [[https://doi.org/10.1145/375663.375670 + Space-efficient Online Computation of Quantile Summaries]] + by Greenwald and Khanna. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col: str, tuple or list + Can be a single column name, or a list of names for multiple columns. + + .. versionchanged:: 2.2.0 + Added support for multiple columns. + probabilities : list or tuple + a list of quantile probabilities + Each number must belong to [0, 1]. + For example 0 is the minimum, 0.5 is the median, 1 is the maximum. + relativeError : float + The relative target precision to achieve + (>= 0). If set to zero, the exact quantiles are computed, which + could be very expensive. Note that values greater than 1 are + accepted but gives the same result as 1. + + Returns + ------- + list + the approximate quantiles at the given probabilities. + + * If the input `col` is a string, the output is a list of floats. + + * If the input `col` is a list or tuple of strings, the output is also a + list, but each element in it is a list of floats, i.e., the output + is a list of list of floats. + + Notes + ----- + Null values will be ignored in numerical columns before calculation. + For columns only containing null values, an empty list is returned. + """ + + if not isinstance(col, (str, list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + + isStr = isinstance(col, str) + + if isinstance(col, tuple): + col = list(col) + elif isStr: + col = [cast(str, col)] + + for c in col: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "col", + "arg_type": type(col).__name__, + "allowed_types": "str", + "return_type": type(c).__name__, + }, + ) + col = _to_list(self._sc, cast(List["ColumnOrName"], col)) + + if not isinstance(probabilities, (list, tuple)): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={ + "arg_name": "probabilities", + "arg_type": type(probabilities).__name__, + }, + ) + if isinstance(probabilities, tuple): + probabilities = list(probabilities) + for p in probabilities: + if not isinstance(p, (float, int)) or p < 0 or p > 1: + raise PySparkTypeError( + error_class="NOT_LIST_OF_FLOAT_OR_INT", + message_parameters={ + "arg_name": "probabilities", + "arg_type": type(p).__name__, + }, + ) + probabilities = _to_list(self._sc, cast(List["ColumnOrName"], probabilities)) + + if not isinstance(relativeError, (float, int)): + raise PySparkTypeError( + error_class="NOT_FLOAT_OR_INT", + message_parameters={ + "arg_name": "relativeError", + "arg_type": type(relativeError).__name__, + }, + ) + if relativeError < 0: + raise PySparkValueError( + error_class="NEGATIVE_VALUE", + message_parameters={ + "arg_name": "relativeError", + "arg_value": str(relativeError), + }, + ) + relativeError = float(relativeError) + + jaq = self._jdf.stat().approxQuantile(col, probabilities, relativeError) + jaq_list = [list(j) for j in jaq] + return jaq_list[0] if isStr else jaq_list + + def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float: + """ + Calculates the correlation of two columns of a :class:`DataFrame` as a double value. + Currently only supports the Pearson Correlation Coefficient. + :func:`DataFrame.corr` and :func:`DataFrameStatFunctions.corr` are aliases of each other. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str + The name of the first column + col2 : str + The name of the second column + method : str, optional + The correlation method. Currently only supports "pearson" + + Returns + ------- + float + Pearson Correlation Coefficient of two columns. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 12), (10, 1), (19, 8)], ["c1", "c2"]) + >>> df.corr("c1", "c2") + -0.3592106040535498 + >>> df = spark.createDataFrame([(11, 12), (10, 11), (9, 10)], ["small", "bigger"]) + >>> df.corr("small", "bigger") + 1.0 + + """ + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + if not method: + method = "pearson" + if not method == "pearson": + raise PySparkValueError( + error_class="VALUE_NOT_PEARSON", + message_parameters={"arg_name": "method", "arg_value": method}, + ) + return self._jdf.stat().corr(col1, col2, method) + + def cov(self, col1: str, col2: str) -> float: + """ + Calculate the sample covariance for the given columns, specified by their names, as a + double value. :func:`DataFrame.cov` and :func:`DataFrameStatFunctions.cov` are aliases. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str + The name of the first column + col2 : str + The name of the second column + + Returns + ------- + float + Covariance of two columns. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 12), (10, 1), (19, 8)], ["c1", "c2"]) + >>> df.cov("c1", "c2") + -18.0 + >>> df = spark.createDataFrame([(11, 12), (10, 11), (9, 10)], ["small", "bigger"]) + >>> df.cov("small", "bigger") + 1.0 + + """ + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + return self._jdf.stat().cov(col1, col2) + + def crosstab(self, col1: str, col2: str) -> "DataFrame": + """ + Computes a pair-wise frequency table of the given columns. Also known as a contingency + table. + The first column of each row will be the distinct values of `col1` and the column names + will be the distinct values of `col2`. The name of the first column will be `$col1_$col2`. + Pairs that have no occurrences will have zero as their counts. + :func:`DataFrame.crosstab` and :func:`DataFrameStatFunctions.crosstab` are aliases. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str + The name of the first column. Distinct items will make the first item of + each row. + col2 : str + The name of the second column. Distinct items will make the column names + of the :class:`DataFrame`. + + Returns + ------- + :class:`DataFrame` + Frequency matrix of two columns. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 11), (1, 11), (3, 10), (4, 8), (4, 8)], ["c1", "c2"]) + >>> df.crosstab("c1", "c2").sort("c1_c2").show() + +-----+---+---+---+ + |c1_c2| 10| 11| 8| + +-----+---+---+---+ + | 1| 0| 2| 0| + | 3| 1| 0| 0| + | 4| 0| 0| 2| + +-----+---+---+---+ + + """ + if not isinstance(col1, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": type(col1).__name__}, + ) + if not isinstance(col2, str): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": type(col2).__name__}, + ) + return DataFrame(self._jdf.stat().crosstab(col1, col2), self.sparkSession) + + def freqItems( + self, cols: Union[List[str], Tuple[str]], support: Optional[float] = None + ) -> "DataFrame": + """ + Finding frequent items for columns, possibly with false positives. Using the + frequent element count algorithm described in + "https://doi.org/10.1145/762471.762473, proposed by Karp, Schenker, and Papadimitriou". + :func:`DataFrame.freqItems` and :func:`DataFrameStatFunctions.freqItems` are aliases. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : list or tuple + Names of the columns to calculate frequent items for as a list or tuple of + strings. + support : float, optional + The frequency with which to consider an item 'frequent'. Default is 1%. + The support must be greater than 1e-4. + + Returns + ------- + :class:`DataFrame` + DataFrame with frequent items. + + Notes + ----- + This function is meant for exploratory data analysis, as we make no + guarantee about the backward compatibility of the schema of the resulting + :class:`DataFrame`. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 11), (1, 11), (3, 10), (4, 8), (4, 8)], ["c1", "c2"]) + >>> df.freqItems(["c1", "c2"]).show() # doctest: +SKIP + +------------+------------+ + |c1_freqItems|c2_freqItems| + +------------+------------+ + | [4, 1, 3]| [8, 11, 10]| + +------------+------------+ + """ + if isinstance(cols, tuple): + cols = list(cols) + if not isinstance(cols, list): + raise PySparkTypeError( + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "cols", "arg_type": type(cols).__name__}, + ) + if not support: + support = 0.01 + return DataFrame( + self._jdf.stat().freqItems(_to_seq(self._sc, cols), support), self.sparkSession + ) + + def _ipython_key_completions_(self) -> List[str]: + """Returns the names of columns in this :class:`DataFrame`. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], ["age", "name"]) + >>> df._ipython_key_completions_() + ['age', 'name'] + + Would return illegal identifiers. + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], ["age 1", "name?1"]) + >>> df._ipython_key_completions_() + ['age 1', 'name?1'] + """ + return self.columns + + def withColumns(self, *colsMap: Dict[str, Column]) -> "DataFrame": + """ + Returns a new :class:`DataFrame` by adding multiple columns or replacing the + existing columns that have the same names. + + The colsMap is a map of column name and column, the column must only refer to attributes + supplied by this Dataset. It is an error to add columns that refer to some other Dataset. + + .. versionadded:: 3.3.0 + Added support for multiple columns adding + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + colsMap : dict + a dict of column name and :class:`Column`. Currently, only a single map is supported. + + Returns + ------- + :class:`DataFrame` + DataFrame with new or replaced columns. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.withColumns({'age2': df.age + 2, 'age3': df.age + 3}).show() + +---+-----+----+----+ + |age| name|age2|age3| + +---+-----+----+----+ + | 2|Alice| 4| 5| + | 5| Bob| 7| 8| + +---+-----+----+----+ + """ + # Below code is to help enable kwargs in future. + assert len(colsMap) == 1 + colsMap = colsMap[0] # type: ignore[assignment] + + if not isinstance(colsMap, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__}, + ) + + col_names = list(colsMap.keys()) + cols = list(colsMap.values()) + + return DataFrame( + self._jdf.withColumns(_to_seq(self._sc, col_names), self._jcols(*cols)), + self.sparkSession, + ) + + def withColumn(self, colName: str, col: Column) -> "DataFrame": + """ + Returns a new :class:`DataFrame` by adding a column or replacing the + existing column that has the same name. + + The column expression must be an expression over this :class:`DataFrame`; attempting to add + a column from some other :class:`DataFrame` will raise an error. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + colName : str + string, name of the new column. + col : :class:`Column` + a :class:`Column` expression for the new column. + + Returns + ------- + :class:`DataFrame` + DataFrame with new or replaced column. + + Notes + ----- + This method introduces a projection internally. Therefore, calling it multiple + times, for instance, via loops in order to add multiple columns can generate big + plans which can cause performance issues and even `StackOverflowException`. + To avoid this, use :func:`select` with multiple columns at once. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.withColumn('age2', df.age + 2).show() + +---+-----+----+ + |age| name|age2| + +---+-----+----+ + | 2|Alice| 4| + | 5| Bob| 7| + +---+-----+----+ + """ + if not isinstance(col, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, + ) + return DataFrame(self._jdf.withColumn(colName, col._jc), self.sparkSession) + + def withColumnRenamed(self, existing: str, new: str) -> "DataFrame": + """Returns a new :class:`DataFrame` by renaming an existing column. + This is a no-op if the schema doesn't contain the given column name. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + existing : str + string, name of the existing column to rename. + new : str + string, new name of the column. + + Returns + ------- + :class:`DataFrame` + DataFrame with renamed column. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df.withColumnRenamed('age', 'age2').show() + +----+-----+ + |age2| name| + +----+-----+ + | 2|Alice| + | 5| Bob| + +----+-----+ + """ + return DataFrame(self._jdf.withColumnRenamed(existing, new), self.sparkSession) + + def withColumnsRenamed(self, colsMap: Dict[str, str]) -> "DataFrame": + """ + Returns a new :class:`DataFrame` by renaming multiple columns. + This is a no-op if the schema doesn't contain the given column names. + + .. versionadded:: 3.4.0 + Added support for multiple columns renaming + + Parameters + ---------- + colsMap : dict + a dict of existing column names and corresponding desired column names. + Currently, only a single map is supported. + + Returns + ------- + :class:`DataFrame` + DataFrame with renamed columns. + + See Also + -------- + :meth:`withColumnRenamed` + + Notes + ----- + Support Spark Connect + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df = df.withColumns({'age2': df.age + 2, 'age3': df.age + 3}) + >>> df.withColumnsRenamed({'age2': 'age4', 'age3': 'age5'}).show() + +---+-----+----+----+ + |age| name|age4|age5| + +---+-----+----+----+ + | 2|Alice| 4| 5| + | 5| Bob| 7| 8| + +---+-----+----+----+ + """ + if not isinstance(colsMap, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__}, + ) + + return DataFrame(self._jdf.withColumnsRenamed(colsMap), self.sparkSession) + + def withMetadata(self, columnName: str, metadata: Dict[str, Any]) -> "DataFrame": + """Returns a new :class:`DataFrame` by updating an existing column with metadata. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + columnName : str + string, name of the existing column to update the metadata. + metadata : dict + dict, new metadata to be assigned to df.schema[columnName].metadata + + Returns + ------- + :class:`DataFrame` + DataFrame with updated metadata column. + + Examples + -------- + >>> df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + >>> df_meta = df.withMetadata('age', {'foo': 'bar'}) + >>> df_meta.schema['age'].metadata + {'foo': 'bar'} + """ + if not isinstance(metadata, dict): + raise PySparkTypeError( + error_class="NOT_DICT", + message_parameters={"arg_name": "metadata", "arg_type": type(metadata).__name__}, + ) + sc = get_active_spark_context() + jmeta = cast(JVMView, sc._jvm).org.apache.spark.sql.types.Metadata.fromJson( + json.dumps(metadata) + ) + return DataFrame(self._jdf.withMetadata(columnName, jmeta), self.sparkSession) + + @overload + def drop(self, cols: "ColumnOrName") -> "DataFrame": + ... + + @overload + def drop(self, *cols: str) -> "DataFrame": + ... + + def drop(self, *cols: "ColumnOrName") -> "DataFrame": # type: ignore[misc] + """Returns a new :class:`DataFrame` without specified columns. + This is a no-op if the schema doesn't contain the given column name(s). + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols: str or :class:`Column` + a name of the column, or the :class:`Column` to drop + + Returns + ------- + :class:`DataFrame` + DataFrame without given columns. + + Notes + ----- + When an input is a column name, it is treated literally without further interpretation. + Otherwise, will try to match the equivalent expression. + So that dropping column by its name `drop(colName)` has different semantic with directly + dropping the column `drop(col(colName))`. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import col, lit + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df2 = spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")]) + + >>> df.drop('age').show() + +-----+ + | name| + +-----+ + | Tom| + |Alice| + | Bob| + +-----+ + >>> df.drop(df.age).show() + +-----+ + | name| + +-----+ + | Tom| + |Alice| + | Bob| + +-----+ + + Drop the column that joined both DataFrames on. + + >>> df.join(df2, df.name == df2.name, 'inner').drop('name').sort('age').show() + +---+------+ + |age|height| + +---+------+ + | 14| 80| + | 16| 85| + +---+------+ + + >>> df3 = df.join(df2) + >>> df3.show() + +---+-----+------+----+ + |age| name|height|name| + +---+-----+------+----+ + | 14| Tom| 80| Tom| + | 14| Tom| 85| Bob| + | 23|Alice| 80| Tom| + | 23|Alice| 85| Bob| + | 16| Bob| 80| Tom| + | 16| Bob| 85| Bob| + +---+-----+------+----+ + + Drop two column by the same name. + + >>> df3.drop("name").show() + +---+------+ + |age|height| + +---+------+ + | 14| 80| + | 14| 85| + | 23| 80| + | 23| 85| + | 16| 80| + | 16| 85| + +---+------+ + + Can not drop col('name') due to ambiguous reference. + + >>> df3.drop(col("name")).show() + Traceback (most recent call last): + ... + pyspark.errors.exceptions.captured.AnalysisException: [AMBIGUOUS_REFERENCE] Reference... + + >>> df4 = df.withColumn("a.b.c", lit(1)) + >>> df4.show() + +---+-----+-----+ + |age| name|a.b.c| + +---+-----+-----+ + | 14| Tom| 1| + | 23|Alice| 1| + | 16| Bob| 1| + +---+-----+-----+ + + >>> df4.drop("a.b.c").show() + +---+-----+ + |age| name| + +---+-----+ + | 14| Tom| + | 23|Alice| + | 16| Bob| + +---+-----+ + + Can not find a column matching the expression "a.b.c". + + >>> df4.drop(col("a.b.c")).show() + +---+-----+-----+ + |age| name|a.b.c| + +---+-----+-----+ + | 14| Tom| 1| + | 23|Alice| 1| + | 16| Bob| 1| + +---+-----+-----+ + """ + column_names: List[str] = [] + java_columns: List[JavaObject] = [] + + for c in cols: + if isinstance(c, str): + column_names.append(c) + elif isinstance(c, Column): + java_columns.append(c._jc) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": type(c).__name__}, + ) + + jdf = self._jdf + if len(java_columns) > 0: + first_column, *remaining_columns = java_columns + jdf = jdf.drop(first_column, self._jseq(remaining_columns)) + if len(column_names) > 0: + jdf = jdf.drop(self._jseq(column_names)) + + return DataFrame(jdf, self.sparkSession) + + def toDF(self, *cols: str) -> "DataFrame": + """Returns a new :class:`DataFrame` that with new specified column names + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + *cols : tuple + a tuple of string new column name. The length of the + list needs to be the same as the number of columns in the initial + :class:`DataFrame` + + Returns + ------- + :class:`DataFrame` + DataFrame with new column names. + + Examples + -------- + >>> df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), + ... (16, "Bob")], ["age", "name"]) + >>> df.toDF('f1', 'f2').show() + +---+-----+ + | f1| f2| + +---+-----+ + | 14| Tom| + | 23|Alice| + | 16| Bob| + +---+-----+ + """ + for col in cols: + if not isinstance(col, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={"arg_name": "cols", "arg_type": type(col).__name__}, + ) + jdf = self._jdf.toDF(self._jseq(cols)) + return DataFrame(jdf, self.sparkSession) + + def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any) -> "DataFrame": + """Returns a new :class:`DataFrame`. Concise syntax for chaining custom transformations. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + func : function + a function that takes and returns a :class:`DataFrame`. + *args + Positional arguments to pass to func. + + .. versionadded:: 3.3.0 + **kwargs + Keyword arguments to pass to func. + + .. versionadded:: 3.3.0 + + Returns + ------- + :class:`DataFrame` + Transformed DataFrame. + + Examples + -------- + >>> from pyspark.sql.functions import col + >>> df = spark.createDataFrame([(1, 1.0), (2, 2.0)], ["int", "float"]) + >>> def cast_all_to_int(input_df): + ... return input_df.select([col(col_name).cast("int") for col_name in input_df.columns]) + ... + >>> def sort_columns_asc(input_df): + ... return input_df.select(*sorted(input_df.columns)) + ... + >>> df.transform(cast_all_to_int).transform(sort_columns_asc).show() + +-----+---+ + |float|int| + +-----+---+ + | 1| 1| + | 2| 2| + +-----+---+ + + >>> def add_n(input_df, n): + ... return input_df.select([(col(col_name) + n).alias(col_name) + ... for col_name in input_df.columns]) + >>> df.transform(add_n, 1).transform(add_n, n=10).show() + +---+-----+ + |int|float| + +---+-----+ + | 12| 12.0| + | 13| 13.0| + +---+-----+ + """ + result = func(self, *args, **kwargs) + assert isinstance( + result, DataFrame + ), "Func returned an instance of type [%s], " "should have been DataFrame." % type(result) + return result + + def sameSemantics(self, other: "DataFrame") -> bool: + """ + Returns `True` when the logical query plans inside both :class:`DataFrame`\\s are equal and + therefore return the same results. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + The equality comparison here is simplified by tolerating the cosmetic differences + such as attribute names. + + This API can compare both :class:`DataFrame`\\s very fast but can still return + `False` on the :class:`DataFrame` that return the same results, for instance, from + different plans. Such false negative semantic can be useful when caching as an example. + + This API is a developer API. + + Parameters + ---------- + other : :class:`DataFrame` + The other DataFrame to compare against. + + Returns + ------- + bool + Whether these two DataFrames are similar. + + Examples + -------- + >>> df1 = spark.range(10) + >>> df2 = spark.range(10) + >>> df1.withColumn("col1", df1.id * 2).sameSemantics(df2.withColumn("col1", df2.id * 2)) + True + >>> df1.withColumn("col1", df1.id * 2).sameSemantics(df2.withColumn("col1", df2.id + 2)) + False + >>> df1.withColumn("col1", df1.id * 2).sameSemantics(df2.withColumn("col0", df2.id * 2)) + True + """ + if not isinstance(other, DataFrame): + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "other", "arg_type": type(other).__name__}, + ) + return self._jdf.sameSemantics(other._jdf) + + def semanticHash(self) -> int: + """ + Returns a hash code of the logical query plan against this :class:`DataFrame`. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + Unlike the standard hash code, the hash is calculated against the query plan + simplified by tolerating the cosmetic differences such as attribute names. + + This API is a developer API. + + Returns + ------- + int + Hash value. + + Examples + -------- + >>> spark.range(10).selectExpr("id as col0").semanticHash() # doctest: +SKIP + 1855039936 + >>> spark.range(10).selectExpr("id as col1").semanticHash() # doctest: +SKIP + 1855039936 + """ + return self._jdf.semanticHash() + + def inputFiles(self) -> List[str]: + """ + Returns a best-effort snapshot of the files that compose this :class:`DataFrame`. + This method simply asks each constituent BaseRelation for its respective files and + takes the union of all results. Depending on the source relations, this may not find + all input files. Duplicates are removed. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + list + List of file paths. + + Examples + -------- + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a single-row DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).repartition(1).write.json(d, mode="overwrite") + ... + ... # Read the JSON file as a DataFrame. + ... df = spark.read.format("json").load(d) + ... + ... # Returns the number of input files. + ... len(df.inputFiles()) + 1 + """ + return list(self._jdf.inputFiles()) + + where = copy_func(filter, sinceversion=1.3, doc=":func:`where` is an alias for :func:`filter`.") + + # Two aliases below were added for pandas compatibility many years ago. + # There are too many differences compared to pandas and we cannot just + # make it "compatible" by adding aliases. Therefore, we stop adding such + # aliases as of Spark 3.0. Two methods below remain just + # for legacy users currently. + groupby = copy_func( + groupBy, sinceversion=1.4, doc=":func:`groupby` is an alias for :func:`groupBy`." + ) + + drop_duplicates = copy_func( + dropDuplicates, + sinceversion=1.4, + doc=":func:`drop_duplicates` is an alias for :func:`dropDuplicates`.", + ) + + def writeTo(self, table: str) -> DataFrameWriterV2: + """ + Create a write configuration builder for v2 sources. + + This builder is used to configure and execute write operations. + + For example, to append or create or replace existing tables. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + table : str + Target table name to write to. + + Returns + ------- + :class:`DataFrameWriterV2` + DataFrameWriterV2 to use further to specify how to save the data + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.writeTo("catalog.db.table").append() # doctest: +SKIP + >>> df.writeTo( # doctest: +SKIP + ... "catalog.db.table" + ... ).partitionedBy("col").createOrReplace() + """ + return DataFrameWriterV2(self, table) + + # Keep to_pandas_on_spark for backward compatibility for now. + def to_pandas_on_spark( + self, index_col: Optional[Union[str, List[str]]] = None + ) -> "PandasOnSparkDataFrame": + warnings.warn( + "DataFrame.to_pandas_on_spark is deprecated. Use DataFrame.pandas_api instead.", + FutureWarning, + ) + return self.pandas_api(index_col) + + def pandas_api( + self, index_col: Optional[Union[str, List[str]]] = None + ) -> "PandasOnSparkDataFrame": + """ + Converts the existing DataFrame into a pandas-on-Spark DataFrame. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + If a pandas-on-Spark DataFrame is converted to a Spark DataFrame and then back + to pandas-on-Spark, it will lose the index information and the original index + will be turned into a normal column. + + This is only available if Pandas is installed and available. + + Parameters + ---------- + index_col: str or list of str, optional, default: None + Index column of table in Spark. + + Returns + ------- + :class:`PandasOnSparkDataFrame` + + See Also + -------- + pyspark.pandas.frame.DataFrame.to_spark + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + >>> df.pandas_api() # doctest: +SKIP + age name + 0 14 Tom + 1 23 Alice + 2 16 Bob + + We can specify the index columns. + + >>> df.pandas_api(index_col="age") # doctest: +SKIP + name + age + 14 Tom + 23 Alice + 16 Bob + """ + from pyspark.pandas.namespace import _get_index_map + from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame + from pyspark.pandas.internal import InternalFrame + + index_spark_columns, index_names = _get_index_map(self, index_col) + internal = InternalFrame( + spark_frame=self, + index_spark_columns=index_spark_columns, + index_names=index_names, # type: ignore[arg-type] + ) + return PandasOnSparkDataFrame(internal) + + # Keep to_koalas for backward compatibility for now. + def to_koalas( + self, index_col: Optional[Union[str, List[str]]] = None + ) -> "PandasOnSparkDataFrame": + return self.pandas_api(index_col) + + +def _to_scala_map(sc: SparkContext, jm: Dict) -> JavaObject: + """ + Convert a dict into a JVM Map. + """ + assert sc._jvm is not None + return sc._jvm.PythonUtils.toScalaMap(jm) + + +class DataFrameNaFunctions: + """Functionality for working with missing data in :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, df: DataFrame): + self.df = df + + def drop( + self, + how: str = "any", + thresh: Optional[int] = None, + subset: Optional[Union[str, Tuple[str, ...], List[str]]] = None, + ) -> DataFrame: + return self.df.dropna(how=how, thresh=thresh, subset=subset) + + drop.__doc__ = DataFrame.dropna.__doc__ + + @overload + def fill(self, value: "LiteralType", subset: Optional[List[str]] = ...) -> DataFrame: + ... + + @overload + def fill(self, value: Dict[str, "LiteralType"]) -> DataFrame: + ... + + def fill( + self, + value: Union["LiteralType", Dict[str, "LiteralType"]], + subset: Optional[List[str]] = None, + ) -> DataFrame: + return self.df.fillna(value=value, subset=subset) # type: ignore[arg-type] + + fill.__doc__ = DataFrame.fillna.__doc__ + + @overload + def replace( + self, + to_replace: List["LiteralType"], + value: List["OptionalPrimitiveType"], + subset: Optional[List[str]] = ..., + ) -> DataFrame: + ... + + @overload + def replace( + self, + to_replace: Dict["LiteralType", "OptionalPrimitiveType"], + subset: Optional[List[str]] = ..., + ) -> DataFrame: + ... + + @overload + def replace( + self, + to_replace: List["LiteralType"], + value: "OptionalPrimitiveType", + subset: Optional[List[str]] = ..., + ) -> DataFrame: + ... + + def replace( # type: ignore[misc] + self, + to_replace: Union[List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]], + value: Optional[ + Union["OptionalPrimitiveType", List["OptionalPrimitiveType"], _NoValueType] + ] = _NoValue, + subset: Optional[List[str]] = None, + ) -> DataFrame: + return self.df.replace(to_replace, value, subset) # type: ignore[arg-type] + + replace.__doc__ = DataFrame.replace.__doc__ + + +class DataFrameStatFunctions: + """Functionality for statistic functions with :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, df: DataFrame): + self.df = df + + @overload + def approxQuantile( + self, + col: str, + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> List[float]: + ... + + @overload + def approxQuantile( + self, + col: Union[List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> List[List[float]]: + ... + + def approxQuantile( + self, + col: Union[str, List[str], Tuple[str]], + probabilities: Union[List[float], Tuple[float]], + relativeError: float, + ) -> Union[List[float], List[List[float]]]: + return self.df.approxQuantile(col, probabilities, relativeError) + + approxQuantile.__doc__ = DataFrame.approxQuantile.__doc__ + + def corr(self, col1: str, col2: str, method: Optional[str] = None) -> float: + return self.df.corr(col1, col2, method) + + corr.__doc__ = DataFrame.corr.__doc__ + + def cov(self, col1: str, col2: str) -> float: + return self.df.cov(col1, col2) + + cov.__doc__ = DataFrame.cov.__doc__ + + def crosstab(self, col1: str, col2: str) -> DataFrame: + return self.df.crosstab(col1, col2) + + crosstab.__doc__ = DataFrame.crosstab.__doc__ + + def freqItems(self, cols: List[str], support: Optional[float] = None) -> DataFrame: + return self.df.freqItems(cols, support) + + freqItems.__doc__ = DataFrame.freqItems.__doc__ + + def sampleBy( + self, col: str, fractions: Dict[Any, float], seed: Optional[int] = None + ) -> DataFrame: + return self.df.sampleBy(col, fractions, seed) + + sampleBy.__doc__ = DataFrame.sampleBy.__doc__ + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.dataframe + + globs = pyspark.sql.dataframe.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.dataframe tests").getOrCreate() + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.dataframe, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/functions.py new file mode 100644 index 0000000..06cb306 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/functions.py @@ -0,0 +1,15889 @@ +# +# 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. +# + +""" +A collections of builtin functions +""" +import inspect +import decimal +import sys +import functools +import warnings +from typing import ( + Any, + cast, + Callable, + Dict, + List, + Iterable, + overload, + Optional, + Tuple, + Type, + TYPE_CHECKING, + Union, + ValuesView, +) + +from py4j.java_gateway import JVMView + +from pyspark import SparkContext +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql.column import Column, _to_java_column, _to_seq, _create_column_from_literal +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.types import ArrayType, DataType, StringType, StructType, _from_numpy_type + +# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409 +from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401 +from pyspark.sql.udtf import UserDefinedTableFunction, _create_py_udtf + +# Keep pandas_udf and PandasUDFType import for backwards compatible import; moved in SPARK-28264 +from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType # noqa: F401 +from pyspark.sql.utils import ( + to_str, + has_numpy, + try_remote_functions, + get_active_spark_context, +) + +if TYPE_CHECKING: + from pyspark.sql._typing import ( + ColumnOrName, + ColumnOrName_, + DataTypeOrString, + UserDefinedFunctionLike, + ) + +if has_numpy: + import numpy as np + +# Note to developers: all of PySpark functions here take string as column names whenever possible. +# Namely, if columns are referred as arguments, they can always be both Column or string, +# even though there might be few exceptions for legacy or inevitable reasons. +# If you are fixing other language APIs together, also please note that Scala side is not the case +# since it requires making every single overridden definition. + + +def _get_jvm_function(name: str, sc: SparkContext) -> Callable: + """ + Retrieves JVM function identified by name from + Java gateway associated with sc. + """ + assert sc._jvm is not None + return getattr(sc._jvm.functions, name) + + +def _invoke_function(name: str, *args: Any) -> Column: + """ + Invokes JVM function identified by name with args + and wraps the result with :class:`~pyspark.sql.Column`. + """ + assert SparkContext._active_spark_context is not None + jf = _get_jvm_function(name, SparkContext._active_spark_context) + return Column(jf(*args)) + + +def _invoke_function_over_columns(name: str, *cols: "ColumnOrName") -> Column: + """ + Invokes n-ary JVM function identified by name + and wraps the result with :class:`~pyspark.sql.Column`. + """ + return _invoke_function(name, *(_to_java_column(col) for col in cols)) + + +def _invoke_function_over_seq_of_columns(name: str, cols: "Iterable[ColumnOrName]") -> Column: + """ + Invokes unary JVM function identified by name with + and wraps the result with :class:`~pyspark.sql.Column`. + """ + sc = get_active_spark_context() + return _invoke_function(name, _to_seq(sc, cols, _to_java_column)) + + +def _invoke_binary_math_function(name: str, col1: Any, col2: Any) -> Column: + """ + Invokes binary JVM math function identified by name + and wraps the result with :class:`~pyspark.sql.Column`. + """ + + # For legacy reasons, the arguments here can be implicitly converted into column + cols = [ + _to_java_column(c) if isinstance(c, (str, Column)) else _create_column_from_literal(c) + for c in (col1, col2) + ] + return _invoke_function(name, *cols) + + +def _options_to_str(options: Optional[Dict[str, Any]] = None) -> Dict[str, Optional[str]]: + if options: + return {key: to_str(value) for (key, value) in options.items()} + return {} + + +@try_remote_functions +def lit(col: Any) -> Column: + """ + Creates a :class:`~pyspark.sql.Column` of literal value. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column`, str, int, float, bool or list, NumPy literals or ndarray. + the value to make it as a PySpark literal. If a column is passed, + it returns the column as is. + + .. versionchanged:: 3.4.0 + Since 3.4.0, it supports the list type. + + Returns + ------- + :class:`~pyspark.sql.Column` + the literal instance. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(lit(5).alias('height'), df.id).show() + +------+---+ + |height| id| + +------+---+ + | 5| 0| + +------+---+ + + Create a literal from a list. + + >>> spark.range(1).select(lit([1, 2, 3])).show() + +--------------+ + |array(1, 2, 3)| + +--------------+ + | [1, 2, 3]| + +--------------+ + """ + if isinstance(col, Column): + return col + elif isinstance(col, list): + if any(isinstance(c, Column) for c in col): + raise PySparkValueError( + error_class="COLUMN_IN_LIST", message_parameters={"func_name": "lit"} + ) + return array(*[lit(item) for item in col]) + else: + if has_numpy and isinstance(col, np.generic): + dt = _from_numpy_type(col.dtype) + if dt is not None: + return _invoke_function("lit", col).astype(dt).alias(str(col)) + return _invoke_function("lit", col) + + +@try_remote_functions +def col(col: str) -> Column: + """ + Returns a :class:`~pyspark.sql.Column` based on the given column name. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : str + the name for the column + + Returns + ------- + :class:`~pyspark.sql.Column` + the corresponding column instance. + + Examples + -------- + >>> col('x') + Column<'x'> + >>> column('x') + Column<'x'> + """ + return _invoke_function("col", col) + + +column = col + + +@try_remote_functions +def asc(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the ascending order of the given column name. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the ascending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + Sort by the column 'id' in the descending order. + + >>> df = spark.range(5) + >>> df = df.sort(desc("id")) + >>> df.show() + +---+ + | id| + +---+ + | 4| + | 3| + | 2| + | 1| + | 0| + +---+ + + Sort by the column 'id' in the ascending order. + + >>> df.orderBy(asc("id")).show() + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + +---+ + """ + return col.asc() if isinstance(col, Column) else _invoke_function("asc", col) + + +@try_remote_functions +def desc(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the descending order of the given column name. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the descending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + Sort by the column 'id' in the descending order. + + >>> spark.range(5).orderBy(desc("id")).show() + +---+ + | id| + +---+ + | 4| + | 3| + | 2| + | 1| + | 0| + +---+ + """ + return col.desc() if isinstance(col, Column) else _invoke_function("desc", col) + + +@try_remote_functions +def sqrt(col: "ColumnOrName") -> Column: + """ + Computes the square root of the specified float value. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(sqrt(lit(4))).show() + +-------+ + |SQRT(4)| + +-------+ + | 2.0| + +-------+ + """ + return _invoke_function_over_columns("sqrt", col) + + +@try_remote_functions +def try_add(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns the sum of `left`and `right` and the result is null on overflow. + The acceptable input types are the same with the `+` operator. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + right : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(1982, 15), (1990, 2)], ["birth", "age"]) + >>> df.select(try_add(df.birth, df.age).alias('r')).collect() + [Row(r=1997), Row(r=1992)] + + >>> from pyspark.sql.types import StructType, StructField, IntegerType, StringType + >>> schema = StructType([ + ... StructField("i", IntegerType(), True), + ... StructField("d", StringType(), True), + ... ]) + >>> df = spark.createDataFrame([(1, '2015-09-30')], schema) + >>> df = df.select(df.i, to_date(df.d).alias('d')) + >>> df.select(try_add(df.d, df.i).alias('r')).collect() + [Row(r=datetime.date(2015, 10, 1))] + + >>> df.select(try_add(df.d, make_interval(df.i)).alias('r')).collect() + [Row(r=datetime.date(2016, 9, 30))] + + >>> df.select( + ... try_add(df.d, make_interval(lit(0), lit(0), lit(0), df.i)).alias('r') + ... ).collect() + [Row(r=datetime.date(2015, 10, 1))] + + >>> df.select( + ... try_add(make_interval(df.i), make_interval(df.i)).alias('r') + ... ).show(truncate=False) + +-------+ + |r | + +-------+ + |2 years| + +-------+ + """ + return _invoke_function_over_columns("try_add", left, right) + + +@try_remote_functions +def try_avg(col: "ColumnOrName") -> Column: + """ + Returns the mean calculated from values of a group and the result is null on overflow. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [(1982, 15), (1990, 2)], ["birth", "age"] + ... ).select(sf.try_avg("age")).show() + +------------+ + |try_avg(age)| + +------------+ + | 8.5| + +------------+ + """ + return _invoke_function_over_columns("try_avg", col) + + +@try_remote_functions +def try_divide(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns `dividend`/`divisor`. It always performs floating point division. Its result is + always null if `divisor` is 0. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + dividend + right : :class:`~pyspark.sql.Column` or str + divisor + + Examples + -------- + >>> df = spark.createDataFrame([(6000, 15), (1990, 2)], ["a", "b"]) + >>> df.select(try_divide(df.a, df.b).alias('r')).collect() + [Row(r=400.0), Row(r=995.0)] + + >>> df = spark.createDataFrame([(1, 2)], ["year", "month"]) + >>> df.select( + ... try_divide(make_interval(df.year), df.month).alias('r') + ... ).show(truncate=False) + +--------+ + |r | + +--------+ + |6 months| + +--------+ + + >>> df.select( + ... try_divide(make_interval(df.year, df.month), lit(2)).alias('r') + ... ).show(truncate=False) + +--------+ + |r | + +--------+ + |7 months| + +--------+ + + >>> df.select( + ... try_divide(make_interval(df.year, df.month), lit(0)).alias('r') + ... ).show(truncate=False) + +----+ + |r | + +----+ + |NULL| + +----+ + """ + return _invoke_function_over_columns("try_divide", left, right) + + +@try_remote_functions +def try_multiply(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns `left`*`right` and the result is null on overflow. The acceptable input types are the + same with the `*` operator. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + multiplicand + right : :class:`~pyspark.sql.Column` or str + multiplier + + Examples + -------- + >>> df = spark.createDataFrame([(6000, 15), (1990, 2)], ["a", "b"]) + >>> df.select(try_multiply(df.a, df.b).alias('r')).collect() + [Row(r=90000), Row(r=3980)] + + >>> df = spark.createDataFrame([(2, 3),], ["a", "b"]) + >>> df.select(try_multiply(make_interval(df.a), df.b).alias('r')).show(truncate=False) + +-------+ + |r | + +-------+ + |6 years| + +-------+ + """ + return _invoke_function_over_columns("try_multiply", left, right) + + +@try_remote_functions +def try_subtract(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns `left`-`right` and the result is null on overflow. The acceptable input types are the + same with the `-` operator. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + right : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(6000, 15), (1990, 2)], ["a", "b"]) + >>> df.select(try_subtract(df.a, df.b).alias('r')).collect() + [Row(r=5985), Row(r=1988)] + + >>> from pyspark.sql.types import StructType, StructField, IntegerType, StringType + >>> schema = StructType([ + ... StructField("i", IntegerType(), True), + ... StructField("d", StringType(), True), + ... ]) + >>> df = spark.createDataFrame([(1, '2015-09-30')], schema) + >>> df = df.select(df.i, to_date(df.d).alias('d')) + >>> df.select(try_subtract(df.d, df.i).alias('r')).collect() + [Row(r=datetime.date(2015, 9, 29))] + + >>> df.select(try_subtract(df.d, make_interval(df.i)).alias('r')).collect() + [Row(r=datetime.date(2014, 9, 30))] + + >>> df.select( + ... try_subtract(df.d, make_interval(lit(0), lit(0), lit(0), df.i)).alias('r') + ... ).collect() + [Row(r=datetime.date(2015, 9, 29))] + + >>> df.select( + ... try_subtract(make_interval(df.i), make_interval(df.i)).alias('r') + ... ).show(truncate=False) + +---------+ + |r | + +---------+ + |0 seconds| + +---------+ + """ + return _invoke_function_over_columns("try_subtract", left, right) + + +@try_remote_functions +def try_sum(col: "ColumnOrName") -> Column: + """ + Returns the sum calculated from values of a group and the result is null on overflow. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(10).select(sf.try_sum("id")).show() + +-----------+ + |try_sum(id)| + +-----------+ + | 45| + +-----------+ + """ + return _invoke_function_over_columns("try_sum", col) + + +@try_remote_functions +def abs(col: "ColumnOrName") -> Column: + """ + Computes the absolute value. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(abs(lit(-1))).show() + +-------+ + |abs(-1)| + +-------+ + | 1| + +-------+ + """ + return _invoke_function_over_columns("abs", col) + + +@try_remote_functions +def mode(col: "ColumnOrName") -> Column: + """ + Returns the most frequent value in a group. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the most frequent value in a group. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... ("Java", 2012, 20000), ("dotNET", 2012, 5000), + ... ("Java", 2012, 20000), ("dotNET", 2012, 5000), + ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)], + ... schema=("course", "year", "earnings")) + >>> df.groupby("course").agg(mode("year")).show() + +------+----------+ + |course|mode(year)| + +------+----------+ + | Java| 2012| + |dotNET| 2012| + +------+----------+ + """ + return _invoke_function_over_columns("mode", col) + + +@try_remote_functions +def max(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the maximum value of the expression in a group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + column for computed results. + + Examples + -------- + >>> df = spark.range(10) + >>> df.select(max(col("id"))).show() + +-------+ + |max(id)| + +-------+ + | 9| + +-------+ + """ + return _invoke_function_over_columns("max", col) + + +@try_remote_functions +def min(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the minimum value of the expression in a group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + column for computed results. + + Examples + -------- + >>> df = spark.range(10) + >>> df.select(min(df.id)).show() + +-------+ + |min(id)| + +-------+ + | 0| + +-------+ + """ + return _invoke_function_over_columns("min", col) + + +@try_remote_functions +def max_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column: + """ + Returns the value associated with the maximum value of ord. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + ord : :class:`~pyspark.sql.Column` or str + column to be maximized + + Returns + ------- + :class:`~pyspark.sql.Column` + value associated with the maximum value of ord. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... ("Java", 2012, 20000), ("dotNET", 2012, 5000), + ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)], + ... schema=("course", "year", "earnings")) + >>> df.groupby("course").agg(max_by("year", "earnings")).show() + +------+----------------------+ + |course|max_by(year, earnings)| + +------+----------------------+ + | Java| 2013| + |dotNET| 2013| + +------+----------------------+ + """ + return _invoke_function_over_columns("max_by", col, ord) + + +@try_remote_functions +def min_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column: + """ + Returns the value associated with the minimum value of ord. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + ord : :class:`~pyspark.sql.Column` or str + column to be minimized + + Returns + ------- + :class:`~pyspark.sql.Column` + value associated with the minimum value of ord. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... ("Java", 2012, 20000), ("dotNET", 2012, 5000), + ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)], + ... schema=("course", "year", "earnings")) + >>> df.groupby("course").agg(min_by("year", "earnings")).show() + +------+----------------------+ + |course|min_by(year, earnings)| + +------+----------------------+ + | Java| 2012| + |dotNET| 2012| + +------+----------------------+ + """ + return _invoke_function_over_columns("min_by", col, ord) + + +@try_remote_functions +def count(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the number of items in a group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + column for computed results. + + Examples + -------- + Count by all columns (start), and by a column that does not count ``None``. + + >>> df = spark.createDataFrame([(None,), ("a",), ("b",), ("c",)], schema=["alphabets"]) + >>> df.select(count(expr("*")), count(df.alphabets)).show() + +--------+----------------+ + |count(1)|count(alphabets)| + +--------+----------------+ + | 4| 3| + +--------+----------------+ + """ + return _invoke_function_over_columns("count", col) + + +@try_remote_functions +def sum(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the sum of all values in the expression. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(10) + >>> df.select(sum(df["id"])).show() + +-------+ + |sum(id)| + +-------+ + | 45| + +-------+ + """ + return _invoke_function_over_columns("sum", col) + + +@try_remote_functions +def avg(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the average of the values in a group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(10) + >>> df.select(avg(col("id"))).show() + +-------+ + |avg(id)| + +-------+ + | 4.5| + +-------+ + """ + return _invoke_function_over_columns("avg", col) + + +@try_remote_functions +def mean(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the average of the values in a group. + An alias of :func:`avg`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(10) + >>> df.select(mean(df.id)).show() + +-------+ + |avg(id)| + +-------+ + | 4.5| + +-------+ + """ + return _invoke_function_over_columns("mean", col) + + +@try_remote_functions +def median(col: "ColumnOrName") -> Column: + """ + Returns the median of the values in a group. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the median of the values in a group. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... ("Java", 2012, 20000), ("dotNET", 2012, 5000), + ... ("Java", 2012, 22000), ("dotNET", 2012, 10000), + ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)], + ... schema=("course", "year", "earnings")) + >>> df.groupby("course").agg(median("earnings")).show() + +------+----------------+ + |course|median(earnings)| + +------+----------------+ + | Java| 22000.0| + |dotNET| 10000.0| + +------+----------------+ + """ + return _invoke_function_over_columns("median", col) + + +@try_remote_functions +def sumDistinct(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the sum of distinct values in the expression. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 3.2.0 + Use :func:`sum_distinct` instead. + """ + warnings.warn("Deprecated in 3.2, use sum_distinct instead.", FutureWarning) + return sum_distinct(col) + + +@try_remote_functions +def sum_distinct(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the sum of distinct values in the expression. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([(None,), (1,), (1,), (2,)], schema=["numbers"]) + >>> df.select(sum_distinct(col("numbers"))).show() + +---------------------+ + |sum(DISTINCT numbers)| + +---------------------+ + | 3| + +---------------------+ + """ + return _invoke_function_over_columns("sum_distinct", col) + + +@try_remote_functions +def product(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the product of the values in a group. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : str, :class:`Column` + column containing values to be multiplied together + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1, 10).toDF('x').withColumn('mod3', col('x') % 3) + >>> prods = df.groupBy('mod3').agg(product('x').alias('product')) + >>> prods.orderBy('mod3').show() + +----+-------+ + |mod3|product| + +----+-------+ + | 0| 162.0| + | 1| 28.0| + | 2| 80.0| + +----+-------+ + """ + return _invoke_function_over_columns("product", col) + + +@try_remote_functions +def acos(col: "ColumnOrName") -> Column: + """ + Computes inverse cosine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + inverse cosine of `col`, as if computed by `java.lang.Math.acos()` + + Examples + -------- + >>> df = spark.range(1, 3) + >>> df.select(acos(df.id)).show() + +--------+ + |ACOS(id)| + +--------+ + | 0.0| + | NaN| + +--------+ + """ + return _invoke_function_over_columns("acos", col) + + +@try_remote_functions +def acosh(col: "ColumnOrName") -> Column: + """ + Computes inverse hyperbolic cosine of the input column. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(2) + >>> df.select(acosh(col("id"))).show() + +---------+ + |ACOSH(id)| + +---------+ + | NaN| + | 0.0| + +---------+ + """ + return _invoke_function_over_columns("acosh", col) + + +@try_remote_functions +def asin(col: "ColumnOrName") -> Column: + """ + Computes inverse sine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + inverse sine of `col`, as if computed by `java.lang.Math.asin()` + + Examples + -------- + >>> df = spark.createDataFrame([(0,), (2,)]) + >>> df.select(asin(df.schema.fieldNames()[0])).show() + +--------+ + |ASIN(_1)| + +--------+ + | 0.0| + | NaN| + +--------+ + """ + return _invoke_function_over_columns("asin", col) + + +@try_remote_functions +def asinh(col: "ColumnOrName") -> Column: + """ + Computes inverse hyperbolic sine of the input column. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(asinh(col("id"))).show() + +---------+ + |ASINH(id)| + +---------+ + | 0.0| + +---------+ + """ + return _invoke_function_over_columns("asinh", col) + + +@try_remote_functions +def atan(col: "ColumnOrName") -> Column: + """ + Compute inverse tangent of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + inverse tangent of `col`, as if computed by `java.lang.Math.atan()` + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(atan(df.id)).show() + +--------+ + |ATAN(id)| + +--------+ + | 0.0| + +--------+ + """ + return _invoke_function_over_columns("atan", col) + + +@try_remote_functions +def atanh(col: "ColumnOrName") -> Column: + """ + Computes inverse hyperbolic tangent of the input column. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([(0,), (2,)], schema=["numbers"]) + >>> df.select(atanh(df["numbers"])).show() + +--------------+ + |ATANH(numbers)| + +--------------+ + | 0.0| + | NaN| + +--------------+ + """ + return _invoke_function_over_columns("atanh", col) + + +@try_remote_functions +def cbrt(col: "ColumnOrName") -> Column: + """ + Computes the cube-root of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(cbrt(lit(27))).show() + +--------+ + |CBRT(27)| + +--------+ + | 3.0| + +--------+ + """ + return _invoke_function_over_columns("cbrt", col) + + +@try_remote_functions +def ceil(col: "ColumnOrName") -> Column: + """ + Computes the ceiling of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(ceil(lit(-0.1))).show() + +----------+ + |CEIL(-0.1)| + +----------+ + | 0| + +----------+ + """ + return _invoke_function_over_columns("ceil", col) + + +@try_remote_functions +def ceiling(col: "ColumnOrName") -> Column: + """ + Computes the ceiling of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.ceil(sf.lit(-0.1))).show() + +----------+ + |CEIL(-0.1)| + +----------+ + | 0| + +----------+ + """ + return _invoke_function_over_columns("ceiling", col) + + +@try_remote_functions +def cos(col: "ColumnOrName") -> Column: + """ + Computes cosine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in radians + + Returns + ------- + :class:`~pyspark.sql.Column` + cosine of the angle, as if computed by `java.lang.Math.cos()`. + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(cos(lit(math.pi))).first() + Row(COS(3.14159...)=-1.0) + """ + return _invoke_function_over_columns("cos", col) + + +@try_remote_functions +def cosh(col: "ColumnOrName") -> Column: + """ + Computes hyperbolic cosine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + hyperbolic angle + + Returns + ------- + :class:`~pyspark.sql.Column` + hyperbolic cosine of the angle, as if computed by `java.lang.Math.cosh()` + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(cosh(lit(1))).first() + Row(COSH(1)=1.54308...) + """ + return _invoke_function_over_columns("cosh", col) + + +@try_remote_functions +def cot(col: "ColumnOrName") -> Column: + """ + Computes cotangent of the input column. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in radians. + + Returns + ------- + :class:`~pyspark.sql.Column` + cotangent of the angle. + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(cot(lit(math.radians(45)))).first() + Row(COT(0.78539...)=1.00000...) + """ + return _invoke_function_over_columns("cot", col) + + +@try_remote_functions +def csc(col: "ColumnOrName") -> Column: + """ + Computes cosecant of the input column. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in radians. + + Returns + ------- + :class:`~pyspark.sql.Column` + cosecant of the angle. + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(csc(lit(math.radians(90)))).first() + Row(CSC(1.57079...)=1.0) + """ + return _invoke_function_over_columns("csc", col) + + +@try_remote_functions +def e() -> Column: + """Returns Euler's number. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.range(1).select(e()).show() + +-----------------+ + | E()| + +-----------------+ + |2.718281828459045| + +-----------------+ + """ + return _invoke_function("e") + + +@try_remote_functions +def exp(col: "ColumnOrName") -> Column: + """ + Computes the exponential of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate exponential for. + + Returns + ------- + :class:`~pyspark.sql.Column` + exponential of the given value. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(exp(lit(0))).show() + +------+ + |EXP(0)| + +------+ + | 1.0| + +------+ + """ + return _invoke_function_over_columns("exp", col) + + +@try_remote_functions +def expm1(col: "ColumnOrName") -> Column: + """ + Computes the exponential of the given value minus one. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate exponential for. + + Returns + ------- + :class:`~pyspark.sql.Column` + exponential less one. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(expm1(lit(1))).first() + Row(EXPM1(1)=1.71828...) + """ + return _invoke_function_over_columns("expm1", col) + + +@try_remote_functions +def floor(col: "ColumnOrName") -> Column: + """ + Computes the floor of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to find floor for. + + Returns + ------- + :class:`~pyspark.sql.Column` + nearest integer that is less than or equal to given value. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(floor(lit(2.5))).show() + +----------+ + |FLOOR(2.5)| + +----------+ + | 2| + +----------+ + """ + return _invoke_function_over_columns("floor", col) + + +@try_remote_functions +def log(col: "ColumnOrName") -> Column: + """ + Computes the natural logarithm of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate natural logarithm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + natural logarithm of the given value. + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(log(lit(math.e))).first() + Row(ln(2.71828...)=1.0) + """ + return _invoke_function_over_columns("log", col) + + +@try_remote_functions +def log10(col: "ColumnOrName") -> Column: + """ + Computes the logarithm of the given value in Base 10. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate logarithm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + logarithm of the given value in Base 10. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(log10(lit(100))).show() + +----------+ + |LOG10(100)| + +----------+ + | 2.0| + +----------+ + """ + return _invoke_function_over_columns("log10", col) + + +@try_remote_functions +def log1p(col: "ColumnOrName") -> Column: + """ + Computes the natural logarithm of the "given value plus one". + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate natural logarithm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + natural logarithm of the "given value plus one". + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(log1p(lit(math.e))).first() + Row(LOG1P(2.71828...)=1.31326...) + + Same as: + + >>> df.select(log(lit(math.e+1))).first() + Row(ln(3.71828...)=1.31326...) + """ + return _invoke_function_over_columns("log1p", col) + + +@try_remote_functions +def negative(col: "ColumnOrName") -> Column: + """ + Returns the negative value. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to calculate negative value for. + + Returns + ------- + :class:`~pyspark.sql.Column` + negative value. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(3).select(sf.negative("id")).show() + +------------+ + |negative(id)| + +------------+ + | 0| + | -1| + | -2| + +------------+ + """ + return _invoke_function_over_columns("negative", col) + + +negate = negative + + +@try_remote_functions +def pi() -> Column: + """Returns Pi. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.range(1).select(pi()).show() + +-----------------+ + | PI()| + +-----------------+ + |3.141592653589793| + +-----------------+ + """ + return _invoke_function("pi") + + +@try_remote_functions +def positive(col: "ColumnOrName") -> Column: + """ + Returns the value. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input value column. + + Returns + ------- + :class:`~pyspark.sql.Column` + value. + + Examples + -------- + >>> df = spark.createDataFrame([(-1,), (0,), (1,)], ['v']) + >>> df.select(positive("v").alias("p")).show() + +---+ + | p| + +---+ + | -1| + | 0| + | 1| + +---+ + """ + return _invoke_function_over_columns("positive", col) + + +@try_remote_functions +def rint(col: "ColumnOrName") -> Column: + """ + Returns the double value that is closest in value to the argument and + is equal to a mathematical integer. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(rint(lit(10.6))).show() + +----------+ + |rint(10.6)| + +----------+ + | 11.0| + +----------+ + + >>> df.select(rint(lit(10.3))).show() + +----------+ + |rint(10.3)| + +----------+ + | 10.0| + +----------+ + """ + return _invoke_function_over_columns("rint", col) + + +@try_remote_functions +def sec(col: "ColumnOrName") -> Column: + """ + Computes secant of the input column. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Angle in radians + + Returns + ------- + :class:`~pyspark.sql.Column` + Secant of the angle. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(sec(lit(1.5))).first() + Row(SEC(1.5)=14.13683...) + """ + return _invoke_function_over_columns("sec", col) + + +@try_remote_functions +def signum(col: "ColumnOrName") -> Column: + """ + Computes the signum of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select( + ... sf.signum(sf.lit(-5)), + ... sf.signum(sf.lit(6)) + ... ).show() + +----------+---------+ + |SIGNUM(-5)|SIGNUM(6)| + +----------+---------+ + | -1.0| 1.0| + +----------+---------+ + """ + return _invoke_function_over_columns("signum", col) + + +@try_remote_functions +def sign(col: "ColumnOrName") -> Column: + """ + Computes the signum of the given value. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select( + ... sf.sign(sf.lit(-5)), + ... sf.sign(sf.lit(6)) + ... ).show() + +--------+-------+ + |sign(-5)|sign(6)| + +--------+-------+ + | -1.0| 1.0| + +--------+-------+ + """ + return _invoke_function_over_columns("sign", col) + + +@try_remote_functions +def sin(col: "ColumnOrName") -> Column: + """ + Computes sine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + sine of the angle, as if computed by `java.lang.Math.sin()` + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(sin(lit(math.radians(90)))).first() + Row(SIN(1.57079...)=1.0) + """ + return _invoke_function_over_columns("sin", col) + + +@try_remote_functions +def sinh(col: "ColumnOrName") -> Column: + """ + Computes hyperbolic sine of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + hyperbolic angle. + + Returns + ------- + :class:`~pyspark.sql.Column` + hyperbolic sine of the given value, + as if computed by `java.lang.Math.sinh()` + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(sinh(lit(1.1))).first() + Row(SINH(1.1)=1.33564...) + """ + return _invoke_function_over_columns("sinh", col) + + +@try_remote_functions +def tan(col: "ColumnOrName") -> Column: + """ + Computes tangent of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in radians + + Returns + ------- + :class:`~pyspark.sql.Column` + tangent of the given value, as if computed by `java.lang.Math.tan()` + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(tan(lit(math.radians(45)))).first() + Row(TAN(0.78539...)=0.99999...) + """ + return _invoke_function_over_columns("tan", col) + + +@try_remote_functions +def tanh(col: "ColumnOrName") -> Column: + """ + Computes hyperbolic tangent of the input column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + hyperbolic angle + + Returns + ------- + :class:`~pyspark.sql.Column` + hyperbolic tangent of the given value + as if computed by `java.lang.Math.tanh()` + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(tanh(lit(math.radians(90)))).first() + Row(TANH(1.57079...)=0.91715...) + """ + return _invoke_function_over_columns("tanh", col) + + +@try_remote_functions +def toDegrees(col: "ColumnOrName") -> Column: + """ + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 2.1.0 + Use :func:`degrees` instead. + """ + warnings.warn("Deprecated in 2.1, use degrees instead.", FutureWarning) + return degrees(col) + + +@try_remote_functions +def toRadians(col: "ColumnOrName") -> Column: + """ + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 2.1.0 + Use :func:`radians` instead. + """ + warnings.warn("Deprecated in 2.1, use radians instead.", FutureWarning) + return radians(col) + + +@try_remote_functions +def bitwiseNOT(col: "ColumnOrName") -> Column: + """ + Computes bitwise not. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 3.2.0 + Use :func:`bitwise_not` instead. + """ + warnings.warn("Deprecated in 3.2, use bitwise_not instead.", FutureWarning) + return bitwise_not(col) + + +@try_remote_functions +def bitwise_not(col: "ColumnOrName") -> Column: + """ + Computes bitwise not. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(bitwise_not(lit(0))).show() + +---+ + | ~0| + +---+ + | -1| + +---+ + >>> df.select(bitwise_not(lit(1))).show() + +---+ + | ~1| + +---+ + | -2| + +---+ + """ + return _invoke_function_over_columns("bitwise_not", col) + + +@try_remote_functions +def bit_count(col: "ColumnOrName") -> Column: + """ + Returns the number of bits that are set in the argument expr as an unsigned 64-bit integer, + or NULL if the argument is NULL. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the number of bits that are set in the argument expr as an unsigned 64-bit integer, + or NULL if the argument is NULL. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(bit_count("c")).show() + +------------+ + |bit_count(c)| + +------------+ + | 1| + | 1| + | 1| + +------------+ + """ + return _invoke_function_over_columns("bit_count", col) + + +@try_remote_functions +def bit_get(col: "ColumnOrName", pos: "ColumnOrName") -> Column: + """ + Returns the value of the bit (0 or 1) at the specified position. + The positions are numbered from right to left, starting at zero. + The position argument cannot be negative. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + pos : :class:`~pyspark.sql.Column` or str + The positions are numbered from right to left, starting at zero. + + Returns + ------- + :class:`~pyspark.sql.Column` + the value of the bit (0 or 1) at the specified position. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(bit_get("c", lit(1))).show() + +-------------+ + |bit_get(c, 1)| + +-------------+ + | 0| + | 0| + | 1| + +-------------+ + """ + return _invoke_function_over_columns("bit_get", col, pos) + + +@try_remote_functions +def getbit(col: "ColumnOrName", pos: "ColumnOrName") -> Column: + """ + Returns the value of the bit (0 or 1) at the specified position. + The positions are numbered from right to left, starting at zero. + The position argument cannot be negative. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + pos : :class:`~pyspark.sql.Column` or str + The positions are numbered from right to left, starting at zero. + + Returns + ------- + :class:`~pyspark.sql.Column` + the value of the bit (0 or 1) at the specified position. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[1], [1], [2]], ["c"] + ... ).select(sf.getbit("c", sf.lit(1))).show() + +------------+ + |getbit(c, 1)| + +------------+ + | 0| + | 0| + | 1| + +------------+ + """ + return _invoke_function_over_columns("getbit", col, pos) + + +@try_remote_functions +def asc_nulls_first(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the ascending order of the given + column name, and null values return before non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the ascending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + >>> df1 = spark.createDataFrame([(1, "Bob"), + ... (0, None), + ... (2, "Alice")], ["age", "name"]) + >>> df1.sort(asc_nulls_first(df1.name)).show() + +---+-----+ + |age| name| + +---+-----+ + | 0| NULL| + | 2|Alice| + | 1| Bob| + +---+-----+ + + """ + return ( + col.asc_nulls_first() + if isinstance(col, Column) + else _invoke_function("asc_nulls_first", col) + ) + + +@try_remote_functions +def asc_nulls_last(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the ascending order of the given + column name, and null values appear after non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the ascending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + >>> df1 = spark.createDataFrame([(0, None), + ... (1, "Bob"), + ... (2, "Alice")], ["age", "name"]) + >>> df1.sort(asc_nulls_last(df1.name)).show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + | 1| Bob| + | 0| NULL| + +---+-----+ + + """ + return ( + col.asc_nulls_last() if isinstance(col, Column) else _invoke_function("asc_nulls_last", col) + ) + + +@try_remote_functions +def desc_nulls_first(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the descending order of the given + column name, and null values appear before non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the descending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + >>> df1 = spark.createDataFrame([(0, None), + ... (1, "Bob"), + ... (2, "Alice")], ["age", "name"]) + >>> df1.sort(desc_nulls_first(df1.name)).show() + +---+-----+ + |age| name| + +---+-----+ + | 0| NULL| + | 1| Bob| + | 2|Alice| + +---+-----+ + + """ + return ( + col.desc_nulls_first() + if isinstance(col, Column) + else _invoke_function("desc_nulls_first", col) + ) + + +@try_remote_functions +def desc_nulls_last(col: "ColumnOrName") -> Column: + """ + Returns a sort expression based on the descending order of the given + column name, and null values appear after non-null values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to sort by in the descending order. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column specifying the order. + + Examples + -------- + >>> df1 = spark.createDataFrame([(0, None), + ... (1, "Bob"), + ... (2, "Alice")], ["age", "name"]) + >>> df1.sort(desc_nulls_last(df1.name)).show() + +---+-----+ + |age| name| + +---+-----+ + | 1| Bob| + | 2|Alice| + | 0| NULL| + +---+-----+ + + """ + return ( + col.desc_nulls_last() + if isinstance(col, Column) + else _invoke_function("desc_nulls_last", col) + ) + + +@try_remote_functions +def stddev(col: "ColumnOrName") -> Column: + """ + Aggregate function: alias for stddev_samp. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + standard deviation of given column. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(6).select(sf.stddev("id")).show() + +------------------+ + | stddev(id)| + +------------------+ + |1.8708286933869...| + +------------------+ + """ + return _invoke_function_over_columns("stddev", col) + + +@try_remote_functions +def std(col: "ColumnOrName") -> Column: + """ + Aggregate function: alias for stddev_samp. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + standard deviation of given column. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(6).select(sf.std("id")).show() + +------------------+ + | std(id)| + +------------------+ + |1.8708286933869...| + +------------------+ + """ + return _invoke_function_over_columns("std", col) + + +@try_remote_functions +def stddev_samp(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the unbiased sample standard deviation of + the expression in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + standard deviation of given column. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(6).select(sf.stddev_samp("id")).show() + +------------------+ + | stddev_samp(id)| + +------------------+ + |1.8708286933869...| + +------------------+ + """ + return _invoke_function_over_columns("stddev_samp", col) + + +@try_remote_functions +def stddev_pop(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns population standard deviation of + the expression in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + standard deviation of given column. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(6).select(sf.stddev_pop("id")).show() + +-----------------+ + | stddev_pop(id)| + +-----------------+ + |1.707825127659...| + +-----------------+ + """ + return _invoke_function_over_columns("stddev_pop", col) + + +@try_remote_functions +def variance(col: "ColumnOrName") -> Column: + """ + Aggregate function: alias for var_samp + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + variance of given column. + + Examples + -------- + >>> df = spark.range(6) + >>> df.select(variance(df.id)).show() + +------------+ + |var_samp(id)| + +------------+ + | 3.5| + +------------+ + """ + return _invoke_function_over_columns("variance", col) + + +@try_remote_functions +def var_samp(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the unbiased sample variance of + the values in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + variance of given column. + + Examples + -------- + >>> df = spark.range(6) + >>> df.select(var_samp(df.id)).show() + +------------+ + |var_samp(id)| + +------------+ + | 3.5| + +------------+ + """ + return _invoke_function_over_columns("var_samp", col) + + +@try_remote_functions +def var_pop(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the population variance of the values in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + variance of given column. + + Examples + -------- + >>> df = spark.range(6) + >>> df.select(var_pop(df.id)).first() + Row(var_pop(id)=2.91666...) + """ + return _invoke_function_over_columns("var_pop", col) + + +@try_remote_functions +def regr_avgx(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the average of the independent variable for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the average of the independent variable for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_avgx("y", "x")).first() + Row(regr_avgx(y, x)=0.999) + """ + return _invoke_function_over_columns("regr_avgx", y, x) + + +@try_remote_functions +def regr_avgy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the average of the dependent variable for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the average of the dependent variable for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_avgy("y", "x")).first() + Row(regr_avgy(y, x)=9.980732994136464) + """ + return _invoke_function_over_columns("regr_avgy", y, x) + + +@try_remote_functions +def regr_count(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the number of non-null number pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the number of non-null number pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_count("y", "x")).first() + Row(regr_count(y, x)=1000) + """ + return _invoke_function_over_columns("regr_count", y, x) + + +@try_remote_functions +def regr_intercept(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the intercept of the univariate linear regression line + for non-null pairs in a group, where `y` is the dependent variable and + `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the intercept of the univariate linear regression line for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_intercept("y", "x")).first() + Row(regr_intercept(y, x)=-0.04961745990969568) + """ + return _invoke_function_over_columns("regr_intercept", y, x) + + +@try_remote_functions +def regr_r2(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the coefficient of determination for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the coefficient of determination for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_r2("y", "x")).first() + Row(regr_r2(y, x)=0.9851908293645436) + """ + return _invoke_function_over_columns("regr_r2", y, x) + + +@try_remote_functions +def regr_slope(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns the slope of the linear regression line for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + the slope of the linear regression line for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_slope("y", "x")).first() + Row(regr_slope(y, x)=10.040390844891048) + """ + return _invoke_function_over_columns("regr_slope", y, x) + + +@try_remote_functions +def regr_sxx(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns REGR_COUNT(y, x) * VAR_POP(x) for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + REGR_COUNT(y, x) * VAR_POP(x) for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_sxx("y", "x")).first() + Row(regr_sxx(y, x)=666.9989999999996) + """ + return _invoke_function_over_columns("regr_sxx", y, x) + + +@try_remote_functions +def regr_sxy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns REGR_COUNT(y, x) * COVAR_POP(y, x) for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + REGR_COUNT(y, x) * COVAR_POP(y, x) for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_sxy("y", "x")).first() + Row(regr_sxy(y, x)=6696.93065315148) + """ + return _invoke_function_over_columns("regr_sxy", y, x) + + +@try_remote_functions +def regr_syy(y: "ColumnOrName", x: "ColumnOrName") -> Column: + """ + Aggregate function: returns REGR_COUNT(y, x) * VAR_POP(y) for non-null pairs + in a group, where `y` is the dependent variable and `x` is the independent variable. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + y : :class:`~pyspark.sql.Column` or str + the dependent variable. + x : :class:`~pyspark.sql.Column` or str + the independent variable. + + Returns + ------- + :class:`~pyspark.sql.Column` + REGR_COUNT(y, x) * VAR_POP(y) for non-null pairs in a group. + + Examples + -------- + >>> x = (col("id") % 3).alias("x") + >>> y = (randn(42) + x * 10).alias("y") + >>> df = spark.range(0, 1000, 1, 1).select(x, y) + >>> df.select(regr_syy("y", "x")).first() + Row(regr_syy(y, x)=68250.53503811295) + """ + return _invoke_function_over_columns("regr_syy", y, x) + + +@try_remote_functions +def every(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns true if all values of `col` are true. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to check if all values are true. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if all values of `col` are true, false otherwise. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[True], [True], [True]], ["flag"] + ... ).select(sf.every("flag")).show() + +-----------+ + |every(flag)| + +-----------+ + | true| + +-----------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[True], [False], [True]], ["flag"] + ... ).select(sf.every("flag")).show() + +-----------+ + |every(flag)| + +-----------+ + | false| + +-----------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[False], [False], [False]], ["flag"] + ... ).select(sf.every("flag")).show() + +-----------+ + |every(flag)| + +-----------+ + | false| + +-----------+ + """ + return _invoke_function_over_columns("every", col) + + +@try_remote_functions +def bool_and(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns true if all values of `col` are true. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to check if all values are true. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if all values of `col` are true, false otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([[True], [True], [True]], ["flag"]) + >>> df.select(bool_and("flag")).show() + +--------------+ + |bool_and(flag)| + +--------------+ + | true| + +--------------+ + >>> df = spark.createDataFrame([[True], [False], [True]], ["flag"]) + >>> df.select(bool_and("flag")).show() + +--------------+ + |bool_and(flag)| + +--------------+ + | false| + +--------------+ + >>> df = spark.createDataFrame([[False], [False], [False]], ["flag"]) + >>> df.select(bool_and("flag")).show() + +--------------+ + |bool_and(flag)| + +--------------+ + | false| + +--------------+ + """ + return _invoke_function_over_columns("bool_and", col) + + +@try_remote_functions +def some(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns true if at least one value of `col` is true. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to check if at least one value is true. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if at least one value of `col` is true, false otherwise. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[True], [True], [True]], ["flag"] + ... ).select(sf.some("flag")).show() + +----------+ + |some(flag)| + +----------+ + | true| + +----------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[True], [False], [True]], ["flag"] + ... ).select(sf.some("flag")).show() + +----------+ + |some(flag)| + +----------+ + | true| + +----------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [[False], [False], [False]], ["flag"] + ... ).select(sf.some("flag")).show() + +----------+ + |some(flag)| + +----------+ + | false| + +----------+ + """ + return _invoke_function_over_columns("some", col) + + +@try_remote_functions +def bool_or(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns true if at least one value of `col` is true. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to check if at least one value is true. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if at least one value of `col` is true, false otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([[True], [True], [True]], ["flag"]) + >>> df.select(bool_or("flag")).show() + +-------------+ + |bool_or(flag)| + +-------------+ + | true| + +-------------+ + >>> df = spark.createDataFrame([[True], [False], [True]], ["flag"]) + >>> df.select(bool_or("flag")).show() + +-------------+ + |bool_or(flag)| + +-------------+ + | true| + +-------------+ + >>> df = spark.createDataFrame([[False], [False], [False]], ["flag"]) + >>> df.select(bool_or("flag")).show() + +-------------+ + |bool_or(flag)| + +-------------+ + | false| + +-------------+ + """ + return _invoke_function_over_columns("bool_or", col) + + +@try_remote_functions +def bit_and(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the bitwise AND of all non-null input values, or null if none. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the bitwise AND of all non-null input values, or null if none. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(bit_and("c")).first() + Row(bit_and(c)=0) + """ + return _invoke_function_over_columns("bit_and", col) + + +@try_remote_functions +def bit_or(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the bitwise OR of all non-null input values, or null if none. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the bitwise OR of all non-null input values, or null if none. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(bit_or("c")).first() + Row(bit_or(c)=3) + """ + return _invoke_function_over_columns("bit_or", col) + + +@try_remote_functions +def bit_xor(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the bitwise XOR of all non-null input values, or null if none. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the bitwise XOR of all non-null input values, or null if none. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(bit_xor("c")).first() + Row(bit_xor(c)=2) + """ + return _invoke_function_over_columns("bit_xor", col) + + +@try_remote_functions +def skewness(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the skewness of the values in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + skewness of given column. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(skewness(df.c)).first() + Row(skewness(c)=0.70710...) + """ + return _invoke_function_over_columns("skewness", col) + + +@try_remote_functions +def kurtosis(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns the kurtosis of the values in a group. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + kurtosis of given column. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.select(kurtosis(df.c)).show() + +-----------+ + |kurtosis(c)| + +-----------+ + | -1.5| + +-----------+ + """ + return _invoke_function_over_columns("kurtosis", col) + + +@try_remote_functions +def collect_list(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns a list of objects with duplicates. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic because the order of collected results depends + on the order of the rows which may be non-deterministic after a shuffle. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + list of objects with duplicates. + + Examples + -------- + >>> df2 = spark.createDataFrame([(2,), (5,), (5,)], ('age',)) + >>> df2.agg(collect_list('age')).collect() + [Row(collect_list(age)=[2, 5, 5])] + """ + return _invoke_function_over_columns("collect_list", col) + + +@try_remote_functions +def array_agg(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns a list of objects with duplicates. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + list of objects with duplicates. + + Examples + -------- + >>> df = spark.createDataFrame([[1],[1],[2]], ["c"]) + >>> df.agg(array_agg('c').alias('r')).collect() + [Row(r=[1, 1, 2])] + """ + return _invoke_function_over_columns("array_agg", col) + + +@try_remote_functions +def collect_set(col: "ColumnOrName") -> Column: + """ + Aggregate function: returns a set of objects with duplicate elements eliminated. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic because the order of collected results depends + on the order of the rows which may be non-deterministic after a shuffle. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + list of objects with no duplicates. + + Examples + -------- + >>> df2 = spark.createDataFrame([(2,), (5,), (5,)], ('age',)) + >>> df2.agg(array_sort(collect_set('age')).alias('c')).collect() + [Row(c=[2, 5])] + """ + return _invoke_function_over_columns("collect_set", col) + + +@try_remote_functions +def degrees(col: "ColumnOrName") -> Column: + """ + Converts an angle measured in radians to an approximately equivalent angle + measured in degrees. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in radians + + Returns + ------- + :class:`~pyspark.sql.Column` + angle in degrees, as if computed by `java.lang.Math.toDegrees()` + + Examples + -------- + >>> import math + >>> df = spark.range(1) + >>> df.select(degrees(lit(math.pi))).first() + Row(DEGREES(3.14159...)=180.0) + """ + return _invoke_function_over_columns("degrees", col) + + +@try_remote_functions +def radians(col: "ColumnOrName") -> Column: + """ + Converts an angle measured in degrees to an approximately equivalent angle + measured in radians. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + angle in degrees + + Returns + ------- + :class:`~pyspark.sql.Column` + angle in radians, as if computed by `java.lang.Math.toRadians()` + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(radians(lit(180))).first() + Row(RADIANS(180)=3.14159...) + """ + return _invoke_function_over_columns("radians", col) + + +@try_remote_functions +def atan2(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + """ + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str, :class:`~pyspark.sql.Column` or float + coordinate on y-axis + col2 : str, :class:`~pyspark.sql.Column` or float + coordinate on x-axis + + Returns + ------- + :class:`~pyspark.sql.Column` + the `theta` component of the point + (`r`, `theta`) + in polar coordinates that corresponds to the point + (`x`, `y`) in Cartesian coordinates, + as if computed by `java.lang.Math.atan2()` + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(atan2(lit(1), lit(2))).first() + Row(ATAN2(1, 2)=0.46364...) + """ + return _invoke_binary_math_function("atan2", col1, col2) + + +@try_remote_functions +def hypot(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + """ + Computes ``sqrt(a^2 + b^2)`` without intermediate overflow or underflow. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str, :class:`~pyspark.sql.Column` or float + a leg. + col2 : str, :class:`~pyspark.sql.Column` or float + b leg. + + Returns + ------- + :class:`~pyspark.sql.Column` + length of the hypotenuse. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(hypot(lit(1), lit(2))).first() + Row(HYPOT(1, 2)=2.23606...) + """ + return _invoke_binary_math_function("hypot", col1, col2) + + +@try_remote_functions +def pow(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) -> Column: + """ + Returns the value of the first argument raised to the power of the second argument. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : str, :class:`~pyspark.sql.Column` or float + the base number. + col2 : str, :class:`~pyspark.sql.Column` or float + the exponent number. + + Returns + ------- + :class:`~pyspark.sql.Column` + the base rased to the power the argument. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(pow(lit(3), lit(2))).first() + Row(POWER(3, 2)=9.0) + """ + return _invoke_binary_math_function("pow", col1, col2) + + +power = pow + + +@try_remote_functions +def pmod(dividend: Union["ColumnOrName", float], divisor: Union["ColumnOrName", float]) -> Column: + """ + Returns the positive value of dividend mod divisor. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + dividend : str, :class:`~pyspark.sql.Column` or float + the column that contains dividend, or the specified dividend value + divisor : str, :class:`~pyspark.sql.Column` or float + the column that contains divisor, or the specified divisor value + + Returns + ------- + :class:`~pyspark.sql.Column` + positive value of dividend mod divisor. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql.functions import pmod + >>> df = spark.createDataFrame([ + ... (1.0, float('nan')), (float('nan'), 2.0), (10.0, 3.0), + ... (float('nan'), float('nan')), (-3.0, 4.0), (-10.0, 3.0), + ... (-5.0, -6.0), (7.0, -8.0), (1.0, 2.0)], + ... ("a", "b")) + >>> df.select(pmod("a", "b")).show() + +----------+ + |pmod(a, b)| + +----------+ + | NaN| + | NaN| + | 1.0| + | NaN| + | 1.0| + | 2.0| + | -5.0| + | 7.0| + | 1.0| + +----------+ + """ + return _invoke_binary_math_function("pmod", dividend, divisor) + + +@try_remote_functions +def width_bucket( + v: "ColumnOrName", + min: "ColumnOrName", + max: "ColumnOrName", + numBucket: Union["ColumnOrName", int], +) -> Column: + """ + Returns the bucket number into which the value of this expression would fall + after being evaluated. Note that input arguments must follow conditions listed below; + otherwise, the method will return null. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + v : str or :class:`~pyspark.sql.Column` + value to compute a bucket number in the histogram + min : str or :class:`~pyspark.sql.Column` + minimum value of the histogram + max : str or :class:`~pyspark.sql.Column` + maximum value of the histogram + numBucket : str, :class:`~pyspark.sql.Column` or int + the number of buckets + + Returns + ------- + :class:`~pyspark.sql.Column` + the bucket number into which the value would fall after being evaluated + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (5.3, 0.2, 10.6, 5), + ... (-2.1, 1.3, 3.4, 3), + ... (8.1, 0.0, 5.7, 4), + ... (-0.9, 5.2, 0.5, 2)], + ... ['v', 'min', 'max', 'n']) + >>> df.select(width_bucket('v', 'min', 'max', 'n')).show() + +----------------------------+ + |width_bucket(v, min, max, n)| + +----------------------------+ + | 3| + | 0| + | 5| + | 3| + +----------------------------+ + """ + numBucket = lit(numBucket) if isinstance(numBucket, int) else numBucket + return _invoke_function_over_columns("width_bucket", v, min, max, numBucket) + + +@try_remote_functions +def row_number() -> Column: + """ + Window function: returns a sequential number starting at 1 within a window partition. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for calculating row numbers. + + Examples + -------- + >>> from pyspark.sql import Window + >>> df = spark.range(3) + >>> w = Window.orderBy(df.id.desc()) + >>> df.withColumn("desc_order", row_number().over(w)).show() + +---+----------+ + | id|desc_order| + +---+----------+ + | 2| 1| + | 1| 2| + | 0| 3| + +---+----------+ + """ + return _invoke_function("row_number") + + +@try_remote_functions +def dense_rank() -> Column: + """ + Window function: returns the rank of rows within a window partition, without any gaps. + + The difference between rank and dense_rank is that dense_rank leaves no gaps in ranking + sequence when there are ties. That is, if you were ranking a competition using dense_rank + and had three people tie for second place, you would say that all three were in second + place and that the next person came in third. Rank would give me sequential numbers, making + the person that came in third place (after the ties) would register as coming in fifth. + + This is equivalent to the DENSE_RANK function in SQL. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for calculating ranks. + + Examples + -------- + >>> from pyspark.sql import Window, types + >>> df = spark.createDataFrame([1, 1, 2, 3, 3, 4], types.IntegerType()) + >>> w = Window.orderBy("value") + >>> df.withColumn("drank", dense_rank().over(w)).show() + +-----+-----+ + |value|drank| + +-----+-----+ + | 1| 1| + | 1| 1| + | 2| 2| + | 3| 3| + | 3| 3| + | 4| 4| + +-----+-----+ + """ + return _invoke_function("dense_rank") + + +@try_remote_functions +def rank() -> Column: + """ + Window function: returns the rank of rows within a window partition. + + The difference between rank and dense_rank is that dense_rank leaves no gaps in ranking + sequence when there are ties. That is, if you were ranking a competition using dense_rank + and had three people tie for second place, you would say that all three were in second + place and that the next person came in third. Rank would give me sequential numbers, making + the person that came in third place (after the ties) would register as coming in fifth. + + This is equivalent to the RANK function in SQL. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for calculating ranks. + + Examples + -------- + >>> from pyspark.sql import Window, types + >>> df = spark.createDataFrame([1, 1, 2, 3, 3, 4], types.IntegerType()) + >>> w = Window.orderBy("value") + >>> df.withColumn("drank", rank().over(w)).show() + +-----+-----+ + |value|drank| + +-----+-----+ + | 1| 1| + | 1| 1| + | 2| 3| + | 3| 4| + | 3| 4| + | 4| 6| + +-----+-----+ + """ + return _invoke_function("rank") + + +@try_remote_functions +def cume_dist() -> Column: + """ + Window function: returns the cumulative distribution of values within a window partition, + i.e. the fraction of rows that are below the current row. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for calculating cumulative distribution. + + Examples + -------- + >>> from pyspark.sql import Window, types + >>> df = spark.createDataFrame([1, 2, 3, 3, 4], types.IntegerType()) + >>> w = Window.orderBy("value") + >>> df.withColumn("cd", cume_dist().over(w)).show() + +-----+---+ + |value| cd| + +-----+---+ + | 1|0.2| + | 2|0.4| + | 3|0.8| + | 3|0.8| + | 4|1.0| + +-----+---+ + """ + return _invoke_function("cume_dist") + + +@try_remote_functions +def percent_rank() -> Column: + """ + Window function: returns the relative rank (i.e. percentile) of rows within a window partition. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for calculating relative rank. + + Examples + -------- + >>> from pyspark.sql import Window, types + >>> df = spark.createDataFrame([1, 1, 2, 3, 3, 4], types.IntegerType()) + >>> w = Window.orderBy("value") + >>> df.withColumn("pr", percent_rank().over(w)).show() + +-----+---+ + |value| pr| + +-----+---+ + | 1|0.0| + | 1|0.0| + | 2|0.4| + | 3|0.6| + | 3|0.6| + | 4|1.0| + +-----+---+ + """ + return _invoke_function("percent_rank") + + +@try_remote_functions +def approxCountDistinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column: + """ + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 2.1.0 + Use :func:`approx_count_distinct` instead. + """ + warnings.warn("Deprecated in 2.1, use approx_count_distinct instead.", FutureWarning) + return approx_count_distinct(col, rsd) + + +@try_remote_functions +def approx_count_distinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column: + """Aggregate function: returns a new :class:`~pyspark.sql.Column` for approximate distinct count + of column `col`. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + rsd : float, optional + maximum relative standard deviation allowed (default = 0.05). + For rsd < 0.01, it is more efficient to use :func:`count_distinct` + + Returns + ------- + :class:`~pyspark.sql.Column` + the column of computed results. + + Examples + -------- + >>> df = spark.createDataFrame([1,2,2,3], "INT") + >>> df.agg(approx_count_distinct("value").alias('distinct_values')).show() + +---------------+ + |distinct_values| + +---------------+ + | 3| + +---------------+ + """ + if rsd is None: + return _invoke_function_over_columns("approx_count_distinct", col) + else: + return _invoke_function("approx_count_distinct", _to_java_column(col), rsd) + + +@try_remote_functions +def broadcast(df: DataFrame) -> DataFrame: + """ + Marks a DataFrame as small enough for use in broadcast joins. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.DataFrame` + DataFrame marked as ready for broadcast join. + + Examples + -------- + >>> from pyspark.sql import types + >>> df = spark.createDataFrame([1, 2, 3, 3, 4], types.IntegerType()) + >>> df_small = spark.range(3) + >>> df_b = broadcast(df_small) + >>> df.join(df_b, df.value == df_small.id).show() + +-----+---+ + |value| id| + +-----+---+ + | 1| 1| + | 2| 2| + +-----+---+ + """ + + sc = get_active_spark_context() + return DataFrame(cast(JVMView, sc._jvm).functions.broadcast(df._jdf), df.sparkSession) + + +@try_remote_functions +def coalesce(*cols: "ColumnOrName") -> Column: + """Returns the first column that is not null. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + list of columns to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + value of the first column that is not null. + + Examples + -------- + >>> cDf = spark.createDataFrame([(None, None), (1, None), (None, 2)], ("a", "b")) + >>> cDf.show() + +----+----+ + | a| b| + +----+----+ + |NULL|NULL| + | 1|NULL| + |NULL| 2| + +----+----+ + + >>> cDf.select(coalesce(cDf["a"], cDf["b"])).show() + +--------------+ + |coalesce(a, b)| + +--------------+ + | NULL| + | 1| + | 2| + +--------------+ + + >>> cDf.select('*', coalesce(cDf["a"], lit(0.0))).show() + +----+----+----------------+ + | a| b|coalesce(a, 0.0)| + +----+----+----------------+ + |NULL|NULL| 0.0| + | 1|NULL| 1.0| + |NULL| 2| 0.0| + +----+----+----------------+ + """ + return _invoke_function_over_seq_of_columns("coalesce", cols) + + +@try_remote_functions +def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """Returns a new :class:`~pyspark.sql.Column` for the Pearson Correlation Coefficient for + ``col1`` and ``col2``. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + first column to calculate correlation. + col1 : :class:`~pyspark.sql.Column` or str + second column to calculate correlation. + + Returns + ------- + :class:`~pyspark.sql.Column` + Pearson Correlation Coefficient of these two column values. + + Examples + -------- + >>> a = range(20) + >>> b = [2 * x for x in range(20)] + >>> df = spark.createDataFrame(zip(a, b), ["a", "b"]) + >>> df.agg(corr("a", "b").alias('c')).collect() + [Row(c=1.0)] + """ + return _invoke_function_over_columns("corr", col1, col2) + + +@try_remote_functions +def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """Returns a new :class:`~pyspark.sql.Column` for the population covariance of ``col1`` and + ``col2``. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + first column to calculate covariance. + col1 : :class:`~pyspark.sql.Column` or str + second column to calculate covariance. + + Returns + ------- + :class:`~pyspark.sql.Column` + covariance of these two column values. + + Examples + -------- + >>> a = [1] * 10 + >>> b = [1] * 10 + >>> df = spark.createDataFrame(zip(a, b), ["a", "b"]) + >>> df.agg(covar_pop("a", "b").alias('c')).collect() + [Row(c=0.0)] + """ + return _invoke_function_over_columns("covar_pop", col1, col2) + + +@try_remote_functions +def covar_samp(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """Returns a new :class:`~pyspark.sql.Column` for the sample covariance of ``col1`` and + ``col2``. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + first column to calculate covariance. + col1 : :class:`~pyspark.sql.Column` or str + second column to calculate covariance. + + Returns + ------- + :class:`~pyspark.sql.Column` + sample covariance of these two column values. + + Examples + -------- + >>> a = [1] * 10 + >>> b = [1] * 10 + >>> df = spark.createDataFrame(zip(a, b), ["a", "b"]) + >>> df.agg(covar_samp("a", "b").alias('c')).collect() + [Row(c=0.0)] + """ + return _invoke_function_over_columns("covar_samp", col1, col2) + + +@try_remote_functions +def countDistinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column: + """Returns a new :class:`~pyspark.sql.Column` for distinct count of ``col`` or ``cols``. + + An alias of :func:`count_distinct`, and it is encouraged to use :func:`count_distinct` + directly. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + return count_distinct(col, *cols) + + +@try_remote_functions +def count_distinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column: + """Returns a new :class:`Column` for distinct count of ``col`` or ``cols``. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + first column to compute on. + cols : :class:`~pyspark.sql.Column` or str + other columns to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + distinct values of these two column values. + + Examples + -------- + >>> from pyspark.sql import types + >>> df1 = spark.createDataFrame([1, 1, 3], types.IntegerType()) + >>> df2 = spark.createDataFrame([1, 2], types.IntegerType()) + >>> df1.join(df2).show() + +-----+-----+ + |value|value| + +-----+-----+ + | 1| 1| + | 1| 2| + | 1| 1| + | 1| 2| + | 3| 1| + | 3| 2| + +-----+-----+ + >>> df1.join(df2).select(count_distinct(df1.value, df2.value)).show() + +----------------------------+ + |count(DISTINCT value, value)| + +----------------------------+ + | 4| + +----------------------------+ + """ + sc = get_active_spark_context() + return _invoke_function( + "count_distinct", _to_java_column(col), _to_seq(sc, cols, _to_java_column) + ) + + +@try_remote_functions +def first(col: "ColumnOrName", ignorenulls: bool = False) -> Column: + """Aggregate function: returns the first value in a group. + + The function by default returns the first values it sees. It will return the first non-null + value it sees when ignoreNulls is set to true. If all values are null, then null is returned. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic because its results depends on the order of the + rows which may be non-deterministic after a shuffle. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to fetch first value for. + ignorenulls : :class:`~pyspark.sql.Column` or str + if first value is null then look for first non-null value. + + Returns + ------- + :class:`~pyspark.sql.Column` + first value of the group. + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age")) + >>> df = df.orderBy(df.age) + >>> df.groupby("name").agg(first("age")).orderBy("name").show() + +-----+----------+ + | name|first(age)| + +-----+----------+ + |Alice| NULL| + | Bob| 5| + +-----+----------+ + + Now, to ignore any nulls we needs to set ``ignorenulls`` to `True` + + >>> df.groupby("name").agg(first("age", ignorenulls=True)).orderBy("name").show() + +-----+----------+ + | name|first(age)| + +-----+----------+ + |Alice| 2| + | Bob| 5| + +-----+----------+ + """ + return _invoke_function("first", _to_java_column(col), ignorenulls) + + +@try_remote_functions +def grouping(col: "ColumnOrName") -> Column: + """ + Aggregate function: indicates whether a specified column in a GROUP BY list is aggregated + or not, returns 1 for aggregated or 0 for not aggregated in the result set. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to check if it's aggregated. + + Returns + ------- + :class:`~pyspark.sql.Column` + returns 1 for aggregated or 0 for not aggregated in the result set. + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.cube("name").agg(grouping("name"), sum("age")).orderBy("name").show() + +-----+--------------+--------+ + | name|grouping(name)|sum(age)| + +-----+--------------+--------+ + | NULL| 1| 7| + |Alice| 0| 2| + | Bob| 0| 5| + +-----+--------------+--------+ + """ + return _invoke_function_over_columns("grouping", col) + + +@try_remote_functions +def grouping_id(*cols: "ColumnOrName") -> Column: + """ + Aggregate function: returns the level of grouping, equals to + + (grouping(c1) << (n-1)) + (grouping(c2) << (n-2)) + ... + grouping(cn) + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The list of columns should match with grouping columns exactly, or empty (means all + the grouping columns). + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + columns to check for. + + Returns + ------- + :class:`~pyspark.sql.Column` + returns level of the grouping it relates to. + + Examples + -------- + >>> df = spark.createDataFrame([(1, "a", "a"), + ... (3, "a", "a"), + ... (4, "b", "c")], ["c1", "c2", "c3"]) + >>> df.cube("c2", "c3").agg(grouping_id(), sum("c1")).orderBy("c2", "c3").show() + +----+----+-------------+-------+ + | c2| c3|grouping_id()|sum(c1)| + +----+----+-------------+-------+ + |NULL|NULL| 3| 8| + |NULL| a| 2| 4| + |NULL| c| 2| 4| + | a|NULL| 1| 4| + | a| a| 0| 4| + | b|NULL| 1| 4| + | b| c| 0| 4| + +----+----+-------------+-------+ + """ + return _invoke_function_over_seq_of_columns("grouping_id", cols) + + +@try_remote_functions +def count_min_sketch( + col: "ColumnOrName", + eps: "ColumnOrName", + confidence: "ColumnOrName", + seed: "ColumnOrName", +) -> Column: + """ + Returns a count-min sketch of a column with the given esp, confidence and seed. + The result is an array of bytes, which can be deserialized to a `CountMinSketch` before usage. + Count-min sketch is a probabilistic data structure used for cardinality estimation + using sub-linear space. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + eps : :class:`~pyspark.sql.Column` or str + relative error, must be positive + confidence : :class:`~pyspark.sql.Column` or str + confidence, must be positive and less than 1.0 + seed : :class:`~pyspark.sql.Column` or str + random seed + + Returns + ------- + :class:`~pyspark.sql.Column` + count-min sketch of the column + + Examples + -------- + >>> df = spark.createDataFrame([[1], [2], [1]], ['data']) + >>> df = df.agg(count_min_sketch(df.data, lit(0.5), lit(0.5), lit(1)).alias('sketch')) + >>> df.select(hex(df.sketch).alias('r')).collect() + [Row(r='0000000100000000000000030000000100000004000000005D8D6AB90000000000000000000000000000000200000000000000010000000000000000')] + """ + return _invoke_function_over_columns("count_min_sketch", col, eps, confidence, seed) + + +@try_remote_functions +def input_file_name() -> Column: + """ + Creates a string column for the file name of the current Spark task. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + file names. + + Examples + -------- + >>> import os + >>> path = os.path.abspath(__file__) + >>> df = spark.read.text(path) + >>> df.select(input_file_name()).first() + Row(input_file_name()='file:///...') + """ + return _invoke_function("input_file_name") + + +@try_remote_functions +def isnan(col: "ColumnOrName") -> Column: + """An expression that returns true if the column is NaN. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + True if value is NaN and False otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([(1.0, float('nan')), (float('nan'), 2.0)], ("a", "b")) + >>> df.select("a", "b", isnan("a").alias("r1"), isnan(df.b).alias("r2")).show() + +---+---+-----+-----+ + | a| b| r1| r2| + +---+---+-----+-----+ + |1.0|NaN|false| true| + |NaN|2.0| true|false| + +---+---+-----+-----+ + """ + return _invoke_function_over_columns("isnan", col) + + +@try_remote_functions +def isnull(col: "ColumnOrName") -> Column: + """An expression that returns true if the column is null. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + True if value is null and False otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([(1, None), (None, 2)], ("a", "b")) + >>> df.select("a", "b", isnull("a").alias("r1"), isnull(df.b).alias("r2")).show() + +----+----+-----+-----+ + | a| b| r1| r2| + +----+----+-----+-----+ + | 1|NULL|false| true| + |NULL| 2| true|false| + +----+----+-----+-----+ + """ + return _invoke_function_over_columns("isnull", col) + + +@try_remote_functions +def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column: + """Aggregate function: returns the last value in a group. + + The function by default returns the last values it sees. It will return the last non-null + value it sees when ignoreNulls is set to true. If all values are null, then null is returned. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic because its results depends on the order of the + rows which may be non-deterministic after a shuffle. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column to fetch last value for. + ignorenulls : :class:`~pyspark.sql.Column` or str + if last value is null then look for non-null value. + + Returns + ------- + :class:`~pyspark.sql.Column` + last value of the group. + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age")) + >>> df = df.orderBy(df.age.desc()) + >>> df.groupby("name").agg(last("age")).orderBy("name").show() + +-----+---------+ + | name|last(age)| + +-----+---------+ + |Alice| NULL| + | Bob| 5| + +-----+---------+ + + Now, to ignore any nulls we needs to set ``ignorenulls`` to `True` + + >>> df.groupby("name").agg(last("age", ignorenulls=True)).orderBy("name").show() + +-----+---------+ + | name|last(age)| + +-----+---------+ + |Alice| 2| + | Bob| 5| + +-----+---------+ + """ + return _invoke_function("last", _to_java_column(col), ignorenulls) + + +@try_remote_functions +def monotonically_increasing_id() -> Column: + """A column that generates monotonically increasing 64-bit integers. + + The generated ID is guaranteed to be monotonically increasing and unique, but not consecutive. + The current implementation puts the partition ID in the upper 31 bits, and the record number + within each partition in the lower 33 bits. The assumption is that the data frame has + less than 1 billion partitions, and each partition has less than 8 billion records. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic because its result depends on partition IDs. + + As an example, consider a :class:`DataFrame` with two partitions, each with 3 records. + This expression would return the following IDs: + 0, 1, 2, 8589934592 (1L << 33), 8589934593, 8589934594. + + Returns + ------- + :class:`~pyspark.sql.Column` + last value of the group. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> spark.range(0, 10, 1, 2).select(sf.monotonically_increasing_id()).show() + +-----------------------------+ + |monotonically_increasing_id()| + +-----------------------------+ + | 0| + | 1| + | 2| + | 3| + | 4| + | 8589934592| + | 8589934593| + | 8589934594| + | 8589934595| + | 8589934596| + +-----------------------------+ + """ + return _invoke_function("monotonically_increasing_id") + + +@try_remote_functions +def nanvl(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """Returns col1 if it is not NaN, or col2 if col1 is NaN. + + Both inputs should be floating point columns (:class:`DoubleType` or :class:`FloatType`). + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + first column to check. + col2 : :class:`~pyspark.sql.Column` or str + second column to return if first is NaN. + + Returns + ------- + :class:`~pyspark.sql.Column` + value from first column or second if first is NaN . + + Examples + -------- + >>> df = spark.createDataFrame([(1.0, float('nan')), (float('nan'), 2.0)], ("a", "b")) + >>> df.select(nanvl("a", "b").alias("r1"), nanvl(df.a, df.b).alias("r2")).collect() + [Row(r1=1.0, r2=1.0), Row(r1=2.0, r2=2.0)] + """ + return _invoke_function_over_columns("nanvl", col1, col2) + + +@try_remote_functions +def percentile( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + frequency: Union[Column, int] = 1, +) -> Column: + """Returns the exact percentile(s) of numeric column `expr` at the given percentage(s) + with value range in [0.0, 1.0]. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str input column. + percentage : :class:`~pyspark.sql.Column`, float, list of floats or tuple of floats + percentage in decimal (must be between 0.0 and 1.0). + frequency : :class:`~pyspark.sql.Column` or int is a positive numeric literal which + controls frequency. + + Returns + ------- + :class:`~pyspark.sql.Column` + the exact `percentile` of the numeric column. + + Examples + -------- + >>> key = (col("id") % 3).alias("key") + >>> value = (randn(42) + key * 10).alias("value") + >>> df = spark.range(0, 1000, 1, 1).select(key, value) + >>> df.select( + ... percentile("value", [0.25, 0.5, 0.75], lit(1)).alias("quantiles") + ... ).show() + +--------------------+ + | quantiles| + +--------------------+ + |[0.74419914941216...| + +--------------------+ + + >>> df.groupBy("key").agg( + ... percentile("value", 0.5, lit(1)).alias("median") + ... ).show() + +---+--------------------+ + |key| median| + +---+--------------------+ + | 0|-0.03449962216667901| + | 1| 9.990389751837329| + | 2| 19.967859769284075| + +---+--------------------+ + """ + sc = get_active_spark_context() + + if isinstance(percentage, (list, tuple)): + # A local list + percentage = _invoke_function( + "array", _to_seq(sc, [_create_column_from_literal(x) for x in percentage]) + )._jc + elif isinstance(percentage, Column): + # Already a Column + percentage = _to_java_column(percentage) + else: + # Probably scalar + percentage = _create_column_from_literal(percentage) + + frequency = ( + _to_java_column(frequency) + if isinstance(frequency, Column) + else _create_column_from_literal(frequency) + ) + + return _invoke_function("percentile", _to_java_column(col), percentage, frequency) + + +@try_remote_functions +def percentile_approx( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + accuracy: Union[Column, float] = 10000, +) -> Column: + """Returns the approximate `percentile` of the numeric column `col` which is the smallest value + in the ordered `col` values (sorted from least to greatest) such that no more than `percentage` + of `col` values is less than the value or equal to that value. + + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column. + percentage : :class:`~pyspark.sql.Column`, float, list of floats or tuple of floats + percentage in decimal (must be between 0.0 and 1.0). + When percentage is an array, each value of the percentage array must be between 0.0 and 1.0. + In this case, returns the approximate percentile array of column col + at the given percentage array. + accuracy : :class:`~pyspark.sql.Column` or float + is a positive numeric literal which controls approximation accuracy + at the cost of memory. Higher value of accuracy yields better accuracy, + 1.0/accuracy is the relative error of the approximation. (default: 10000). + + Returns + ------- + :class:`~pyspark.sql.Column` + approximate `percentile` of the numeric column. + + Examples + -------- + >>> key = (col("id") % 3).alias("key") + >>> value = (randn(42) + key * 10).alias("value") + >>> df = spark.range(0, 1000, 1, 1).select(key, value) + >>> df.select( + ... percentile_approx("value", [0.25, 0.5, 0.75], 1000000).alias("quantiles") + ... ).printSchema() + root + |-- quantiles: array (nullable = true) + | |-- element: double (containsNull = false) + + >>> df.groupBy("key").agg( + ... percentile_approx("value", 0.5, lit(1000000)).alias("median") + ... ).printSchema() + root + |-- key: long (nullable = true) + |-- median: double (nullable = true) + """ + sc = get_active_spark_context() + + if isinstance(percentage, (list, tuple)): + # A local list + percentage = _invoke_function( + "array", _to_seq(sc, [_create_column_from_literal(x) for x in percentage]) + )._jc + elif isinstance(percentage, Column): + # Already a Column + percentage = _to_java_column(percentage) + else: + # Probably scalar + percentage = _create_column_from_literal(percentage) + + accuracy = ( + _to_java_column(accuracy) + if isinstance(accuracy, Column) + else _create_column_from_literal(accuracy) + ) + + return _invoke_function("percentile_approx", _to_java_column(col), percentage, accuracy) + + +@try_remote_functions +def approx_percentile( + col: "ColumnOrName", + percentage: Union[Column, float, List[float], Tuple[float]], + accuracy: Union[Column, float] = 10000, +) -> Column: + """Returns the approximate `percentile` of the numeric column `col` which is the smallest value + in the ordered `col` values (sorted from least to greatest) such that no more than `percentage` + of `col` values is less than the value or equal to that value. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column. + percentage : :class:`~pyspark.sql.Column`, float, list of floats or tuple of floats + percentage in decimal (must be between 0.0 and 1.0). + When percentage is an array, each value of the percentage array must be between 0.0 and 1.0. + In this case, returns the approximate percentile array of column col + at the given percentage array. + accuracy : :class:`~pyspark.sql.Column` or float + is a positive numeric literal which controls approximation accuracy + at the cost of memory. Higher value of accuracy yields better accuracy, + 1.0/accuracy is the relative error of the approximation. (default: 10000). + + Returns + ------- + :class:`~pyspark.sql.Column` + approximate `percentile` of the numeric column. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> key = (sf.col("id") % 3).alias("key") + >>> value = (sf.randn(42) + key * 10).alias("value") + >>> df = spark.range(0, 1000, 1, 1).select(key, value) + >>> df.select( + ... sf.approx_percentile("value", [0.25, 0.5, 0.75], 1000000) + ... ).printSchema() + root + |-- approx_percentile(value, array(0.25, 0.5, 0.75), 1000000): array (nullable = true) + | |-- element: double (containsNull = false) + + >>> df.groupBy("key").agg( + ... sf.approx_percentile("value", 0.5, sf.lit(1000000)) + ... ).printSchema() + root + |-- key: long (nullable = true) + |-- approx_percentile(value, 0.5, 1000000): double (nullable = true) + """ + sc = get_active_spark_context() + + if isinstance(percentage, (list, tuple)): + # A local list + percentage = _invoke_function( + "array", _to_seq(sc, [_create_column_from_literal(x) for x in percentage]) + )._jc + elif isinstance(percentage, Column): + # Already a Column + percentage = _to_java_column(percentage) + else: + # Probably scalar + percentage = _create_column_from_literal(percentage) + + accuracy = ( + _to_java_column(accuracy) + if isinstance(accuracy, Column) + else _create_column_from_literal(accuracy) + ) + + return _invoke_function("approx_percentile", _to_java_column(col), percentage, accuracy) + + +@try_remote_functions +def rand(seed: Optional[int] = None) -> Column: + """Generates a random column with independent and identically distributed (i.i.d.) samples + uniformly distributed in [0.0, 1.0). + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic in general case. + + Parameters + ---------- + seed : int (default: None) + seed value for random generator. + + Returns + ------- + :class:`~pyspark.sql.Column` + random values. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> spark.range(0, 2, 1, 1).withColumn('rand', sf.rand(seed=42) * 3).show() + +---+------------------+ + | id| rand| + +---+------------------+ + | 0|1.8575681106759028| + | 1|1.5288056527339444| + +---+------------------+ + """ + if seed is not None: + return _invoke_function("rand", seed) + else: + return _invoke_function("rand") + + +@try_remote_functions +def randn(seed: Optional[int] = None) -> Column: + """Generates a column with independent and identically distributed (i.i.d.) samples from + the standard normal distribution. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic in general case. + + Parameters + ---------- + seed : int (default: None) + seed value for random generator. + + Returns + ------- + :class:`~pyspark.sql.Column` + random values. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> spark.range(0, 2, 1, 1).withColumn('randn', sf.randn(seed=42)).show() + +---+------------------+ + | id| randn| + +---+------------------+ + | 0| 2.384479054241165| + | 1|0.1920934041293524| + +---+------------------+ + """ + if seed is not None: + return _invoke_function("randn", seed) + else: + return _invoke_function("randn") + + +@try_remote_functions +def round(col: "ColumnOrName", scale: int = 0) -> Column: + """ + Round the given value to `scale` decimal places using HALF_UP rounding mode if `scale` >= 0 + or at integral part when `scale` < 0. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column to round. + scale : int optional default 0 + scale value. + + Returns + ------- + :class:`~pyspark.sql.Column` + rounded values. + + Examples + -------- + >>> spark.createDataFrame([(2.5,)], ['a']).select(round('a', 0).alias('r')).collect() + [Row(r=3.0)] + """ + return _invoke_function("round", _to_java_column(col), scale) + + +@try_remote_functions +def bround(col: "ColumnOrName", scale: int = 0) -> Column: + """ + Round the given value to `scale` decimal places using HALF_EVEN rounding mode if `scale` >= 0 + or at integral part when `scale` < 0. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column to round. + scale : int optional default 0 + scale value. + + Returns + ------- + :class:`~pyspark.sql.Column` + rounded values. + + Examples + -------- + >>> spark.createDataFrame([(2.5,)], ['a']).select(bround('a', 0).alias('r')).collect() + [Row(r=2.0)] + """ + return _invoke_function("bround", _to_java_column(col), scale) + + +@try_remote_functions +def shiftLeft(col: "ColumnOrName", numBits: int) -> Column: + """Shift the given value numBits left. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 3.2.0 + Use :func:`shiftleft` instead. + """ + warnings.warn("Deprecated in 3.2, use shiftleft instead.", FutureWarning) + return shiftleft(col, numBits) + + +@try_remote_functions +def shiftleft(col: "ColumnOrName", numBits: int) -> Column: + """Shift the given value numBits left. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to shift. + numBits : int + number of bits to shift. + + Returns + ------- + :class:`~pyspark.sql.Column` + shifted value. + + Examples + -------- + >>> spark.createDataFrame([(21,)], ['a']).select(shiftleft('a', 1).alias('r')).collect() + [Row(r=42)] + """ + return _invoke_function("shiftleft", _to_java_column(col), numBits) + + +@try_remote_functions +def shiftRight(col: "ColumnOrName", numBits: int) -> Column: + """(Signed) shift the given value numBits right. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 3.2.0 + Use :func:`shiftright` instead. + """ + warnings.warn("Deprecated in 3.2, use shiftright instead.", FutureWarning) + return shiftright(col, numBits) + + +@try_remote_functions +def shiftright(col: "ColumnOrName", numBits: int) -> Column: + """(Signed) shift the given value numBits right. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to shift. + numBits : int + number of bits to shift. + + Returns + ------- + :class:`~pyspark.sql.Column` + shifted values. + + Examples + -------- + >>> spark.createDataFrame([(42,)], ['a']).select(shiftright('a', 1).alias('r')).collect() + [Row(r=21)] + """ + return _invoke_function("shiftright", _to_java_column(col), numBits) + + +@try_remote_functions +def shiftRightUnsigned(col: "ColumnOrName", numBits: int) -> Column: + """Unsigned shift the given value numBits right. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. deprecated:: 3.2.0 + Use :func:`shiftrightunsigned` instead. + """ + warnings.warn("Deprecated in 3.2, use shiftrightunsigned instead.", FutureWarning) + return shiftrightunsigned(col, numBits) + + +@try_remote_functions +def shiftrightunsigned(col: "ColumnOrName", numBits: int) -> Column: + """Unsigned shift the given value numBits right. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to shift. + numBits : int + number of bits to shift. + + Returns + ------- + :class:`~pyspark.sql.Column` + shifted value. + + Examples + -------- + >>> df = spark.createDataFrame([(-42,)], ['a']) + >>> df.select(shiftrightunsigned('a', 1).alias('r')).collect() + [Row(r=9223372036854775787)] + """ + return _invoke_function("shiftrightunsigned", _to_java_column(col), numBits) + + +@try_remote_functions +def spark_partition_id() -> Column: + """A column for partition ID. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + This is non deterministic because it depends on data partitioning and task scheduling. + + Returns + ------- + :class:`~pyspark.sql.Column` + partition id the record belongs to. + + Examples + -------- + >>> df = spark.range(2) + >>> df.repartition(1).select(spark_partition_id().alias("pid")).collect() + [Row(pid=0), Row(pid=0)] + """ + return _invoke_function("spark_partition_id") + + +@try_remote_functions +def expr(str: str) -> Column: + """Parses the expression string into the column that it represents + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + str : str + expression defined in string. + + Returns + ------- + :class:`~pyspark.sql.Column` + column representing the expression. + + Examples + -------- + >>> df = spark.createDataFrame([["Alice"], ["Bob"]], ["name"]) + >>> df.select("name", expr("length(name)")).show() + +-----+------------+ + | name|length(name)| + +-----+------------+ + |Alice| 5| + | Bob| 3| + +-----+------------+ + """ + return _invoke_function("expr", str) + + +@overload +def struct(*cols: "ColumnOrName") -> Column: + ... + + +@overload +def struct(__cols: Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]) -> Column: + ... + + +@try_remote_functions +def struct( + *cols: Union["ColumnOrName", Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]] +) -> Column: + """Creates a new struct column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : list, set, str or :class:`~pyspark.sql.Column` + column names or :class:`~pyspark.sql.Column`\\s to contain in the output struct. + + Returns + ------- + :class:`~pyspark.sql.Column` + a struct type column of given columns. + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.select(struct('age', 'name').alias("struct")).collect() + [Row(struct=Row(age=2, name='Alice')), Row(struct=Row(age=5, name='Bob'))] + >>> df.select(struct([df.age, df.name]).alias("struct")).collect() + [Row(struct=Row(age=2, name='Alice')), Row(struct=Row(age=5, name='Bob'))] + """ + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_seq_of_columns("struct", cols) # type: ignore[arg-type] + + +@try_remote_functions +def named_struct(*cols: "ColumnOrName") -> Column: + """ + Creates a struct with the given field names and values. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + list of columns to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + + Examples + -------- + >>> df = spark.createDataFrame([(1, 2, 3)], ['a', 'b', 'c']) + >>> df.select(named_struct(lit('x'), df.a, lit('y'), df.b).alias('r')).collect() + [Row(r=Row(x=1, y=2))] + """ + return _invoke_function_over_seq_of_columns("named_struct", cols) + + +@try_remote_functions +def greatest(*cols: "ColumnOrName") -> Column: + """ + Returns the greatest value of the list of column names, skipping null values. + This function takes at least 2 parameters. It will return null if all parameters are null. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + columns to check for gratest value. + + Returns + ------- + :class:`~pyspark.sql.Column` + gratest value. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 4, 3)], ['a', 'b', 'c']) + >>> df.select(greatest(df.a, df.b, df.c).alias("greatest")).collect() + [Row(greatest=4)] + """ + if len(cols) < 2: + raise PySparkValueError( + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "greatest", "num_cols": "2"}, + ) + return _invoke_function_over_seq_of_columns("greatest", cols) + + +@try_remote_functions +def least(*cols: "ColumnOrName") -> Column: + """ + Returns the least value of the list of column names, skipping null values. + This function takes at least 2 parameters. It will return null if all parameters are null. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + column names or columns to be compared + + Returns + ------- + :class:`~pyspark.sql.Column` + least value. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 4, 3)], ['a', 'b', 'c']) + >>> df.select(least(df.a, df.b, df.c).alias("least")).collect() + [Row(least=1)] + """ + if len(cols) < 2: + raise PySparkValueError( + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "least", "num_cols": "2"}, + ) + return _invoke_function_over_seq_of_columns("least", cols) + + +@try_remote_functions +def when(condition: Column, value: Any) -> Column: + """Evaluates a list of conditions and returns one of multiple possible result expressions. + If :func:`pyspark.sql.Column.otherwise` is not invoked, None is returned for unmatched + conditions. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + condition : :class:`~pyspark.sql.Column` + a boolean :class:`~pyspark.sql.Column` expression. + value : + a literal value, or a :class:`~pyspark.sql.Column` expression. + + Returns + ------- + :class:`~pyspark.sql.Column` + column representing when expression. + + Examples + -------- + >>> df = spark.range(3) + >>> df.select(when(df['id'] == 2, 3).otherwise(4).alias("age")).show() + +---+ + |age| + +---+ + | 4| + | 4| + | 3| + +---+ + + >>> df.select(when(df.id == 2, df.id + 1).alias("age")).show() + +----+ + | age| + +----+ + |NULL| + |NULL| + | 3| + +----+ + """ + # Explicitly not using ColumnOrName type here to make reading condition less opaque + if not isinstance(condition, Column): + raise PySparkTypeError( + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, + ) + v = value._jc if isinstance(value, Column) else value + + return _invoke_function("when", condition._jc, v) + + +@overload # type: ignore[no-redef] +def log(arg1: "ColumnOrName") -> Column: + ... + + +@overload +def log(arg1: float, arg2: "ColumnOrName") -> Column: + ... + + +@try_remote_functions +def log(arg1: Union["ColumnOrName", float], arg2: Optional["ColumnOrName"] = None) -> Column: + """Returns the first argument-based logarithm of the second argument. + + If there is only one argument, then this takes the natural logarithm of the argument. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + arg1 : :class:`~pyspark.sql.Column`, str or float + base number or actual number (in this case base is `e`) + arg2 : :class:`~pyspark.sql.Column`, str or float + number to calculate logariphm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + logariphm of given value. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.sql("SELECT * FROM VALUES (1), (2), (4) AS t(value)") + >>> df.select(sf.log(2.0, df.value).alias('log2_value')).show() + +----------+ + |log2_value| + +----------+ + | 0.0| + | 1.0| + | 2.0| + +----------+ + + And Natural logarithm + + >>> df.select(sf.log(df.value).alias('ln_value')).show() + +------------------+ + | ln_value| + +------------------+ + | 0.0| + |0.6931471805599453| + |1.3862943611198906| + +------------------+ + """ + if arg2 is None: + return _invoke_function_over_columns("log", cast("ColumnOrName", arg1)) + else: + return _invoke_function("log", arg1, _to_java_column(arg2)) + + +@try_remote_functions +def ln(col: "ColumnOrName") -> Column: + """Returns the natural logarithm of the argument. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column to calculate logariphm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + natural logarithm of given value. + + Examples + -------- + >>> df = spark.createDataFrame([(4,)], ['a']) + >>> df.select(ln('a')).show() + +------------------+ + | ln(a)| + +------------------+ + |1.3862943611198906| + +------------------+ + """ + return _invoke_function_over_columns("ln", col) + + +@try_remote_functions +def log2(col: "ColumnOrName") -> Column: + """Returns the base-2 logarithm of the argument. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column to calculate logariphm for. + + Returns + ------- + :class:`~pyspark.sql.Column` + logariphm of given value. + + Examples + -------- + >>> df = spark.createDataFrame([(4,)], ['a']) + >>> df.select(log2('a').alias('log2')).show() + +----+ + |log2| + +----+ + | 2.0| + +----+ + """ + return _invoke_function_over_columns("log2", col) + + +@try_remote_functions +def conv(col: "ColumnOrName", fromBase: int, toBase: int) -> Column: + """ + Convert a number in a string column from one base to another. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column to convert base for. + fromBase: int + from base number. + toBase: int + to base number. + + Returns + ------- + :class:`~pyspark.sql.Column` + logariphm of given value. + + Examples + -------- + >>> df = spark.createDataFrame([("010101",)], ['n']) + >>> df.select(conv(df.n, 2, 16).alias('hex')).collect() + [Row(hex='15')] + """ + return _invoke_function("conv", _to_java_column(col), fromBase, toBase) + + +@try_remote_functions +def factorial(col: "ColumnOrName") -> Column: + """ + Computes the factorial of the given value. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column to calculate factorial for. + + Returns + ------- + :class:`~pyspark.sql.Column` + factorial of given value. + + Examples + -------- + >>> df = spark.createDataFrame([(5,)], ['n']) + >>> df.select(factorial(df.n).alias('f')).collect() + [Row(f=120)] + """ + return _invoke_function_over_columns("factorial", col) + + +# --------------- Window functions ------------------------ + + +@try_remote_functions +def lag(col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None) -> Column: + """ + Window function: returns the value that is `offset` rows before the current row, and + `default` if there is less than `offset` rows before the current row. For example, + an `offset` of one will return the previous row at any given point in the window partition. + + This is equivalent to the LAG function in SQL. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + offset : int, optional default 1 + number of row to extend + default : optional + default value + + Returns + ------- + :class:`~pyspark.sql.Column` + value before current row based on `offset`. + + Examples + -------- + >>> from pyspark.sql import Window + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.show() + +---+---+ + | c1| c2| + +---+---+ + | a| 1| + | a| 2| + | a| 3| + | b| 8| + | b| 2| + +---+---+ + >>> w = Window.partitionBy("c1").orderBy("c2") + >>> df.withColumn("previos_value", lag("c2").over(w)).show() + +---+---+-------------+ + | c1| c2|previos_value| + +---+---+-------------+ + | a| 1| NULL| + | a| 2| 1| + | a| 3| 2| + | b| 2| NULL| + | b| 8| 2| + +---+---+-------------+ + >>> df.withColumn("previos_value", lag("c2", 1, 0).over(w)).show() + +---+---+-------------+ + | c1| c2|previos_value| + +---+---+-------------+ + | a| 1| 0| + | a| 2| 1| + | a| 3| 2| + | b| 2| 0| + | b| 8| 2| + +---+---+-------------+ + >>> df.withColumn("previos_value", lag("c2", 2, -1).over(w)).show() + +---+---+-------------+ + | c1| c2|previos_value| + +---+---+-------------+ + | a| 1| -1| + | a| 2| -1| + | a| 3| 1| + | b| 2| -1| + | b| 8| -1| + +---+---+-------------+ + """ + return _invoke_function("lag", _to_java_column(col), offset, default) + + +@try_remote_functions +def lead(col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None) -> Column: + """ + Window function: returns the value that is `offset` rows after the current row, and + `default` if there is less than `offset` rows after the current row. For example, + an `offset` of one will return the next row at any given point in the window partition. + + This is equivalent to the LEAD function in SQL. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + offset : int, optional default 1 + number of row to extend + default : optional + default value + + Returns + ------- + :class:`~pyspark.sql.Column` + value after current row based on `offset`. + + Examples + -------- + >>> from pyspark.sql import Window + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.show() + +---+---+ + | c1| c2| + +---+---+ + | a| 1| + | a| 2| + | a| 3| + | b| 8| + | b| 2| + +---+---+ + >>> w = Window.partitionBy("c1").orderBy("c2") + >>> df.withColumn("next_value", lead("c2").over(w)).show() + +---+---+----------+ + | c1| c2|next_value| + +---+---+----------+ + | a| 1| 2| + | a| 2| 3| + | a| 3| NULL| + | b| 2| 8| + | b| 8| NULL| + +---+---+----------+ + >>> df.withColumn("next_value", lead("c2", 1, 0).over(w)).show() + +---+---+----------+ + | c1| c2|next_value| + +---+---+----------+ + | a| 1| 2| + | a| 2| 3| + | a| 3| 0| + | b| 2| 8| + | b| 8| 0| + +---+---+----------+ + >>> df.withColumn("next_value", lead("c2", 2, -1).over(w)).show() + +---+---+----------+ + | c1| c2|next_value| + +---+---+----------+ + | a| 1| 3| + | a| 2| -1| + | a| 3| -1| + | b| 2| -1| + | b| 8| -1| + +---+---+----------+ + """ + return _invoke_function("lead", _to_java_column(col), offset, default) + + +@try_remote_functions +def nth_value(col: "ColumnOrName", offset: int, ignoreNulls: Optional[bool] = False) -> Column: + """ + Window function: returns the value that is the `offset`\\th row of the window frame + (counting from 1), and `null` if the size of window frame is less than `offset` rows. + + It will return the `offset`\\th non-null value it sees when `ignoreNulls` is set to + true. If all values are null, then null is returned. + + This is equivalent to the nth_value function in SQL. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + offset : int + number of row to use as the value + ignoreNulls : bool, optional + indicates the Nth value should skip null in the + determination of which row to use + + Returns + ------- + :class:`~pyspark.sql.Column` + value of nth row. + + Examples + -------- + >>> from pyspark.sql import Window + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.show() + +---+---+ + | c1| c2| + +---+---+ + | a| 1| + | a| 2| + | a| 3| + | b| 8| + | b| 2| + +---+---+ + >>> w = Window.partitionBy("c1").orderBy("c2") + >>> df.withColumn("nth_value", nth_value("c2", 1).over(w)).show() + +---+---+---------+ + | c1| c2|nth_value| + +---+---+---------+ + | a| 1| 1| + | a| 2| 1| + | a| 3| 1| + | b| 2| 2| + | b| 8| 2| + +---+---+---------+ + >>> df.withColumn("nth_value", nth_value("c2", 2).over(w)).show() + +---+---+---------+ + | c1| c2|nth_value| + +---+---+---------+ + | a| 1| NULL| + | a| 2| 2| + | a| 3| 2| + | b| 2| NULL| + | b| 8| 8| + +---+---+---------+ + """ + return _invoke_function("nth_value", _to_java_column(col), offset, ignoreNulls) + + +@try_remote_functions +def any_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + """Returns some value of `col` for a group of rows. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + ignorenulls : :class:`~pyspark.sql.Column` or bool + if first value is null then look for first non-null value. + + Returns + ------- + :class:`~pyspark.sql.Column` + some value of `col` for a group of rows. + + Examples + -------- + >>> df = spark.createDataFrame([(None, 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.select(any_value('c1'), any_value('c2')).collect() + [Row(any_value(c1)=None, any_value(c2)=1)] + >>> df.select(any_value('c1', True), any_value('c2', True)).collect() + [Row(any_value(c1)='a', any_value(c2)=1)] + """ + if ignoreNulls is None: + return _invoke_function_over_columns("any_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("any_value", col, ignoreNulls) + + +@try_remote_functions +def first_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + """Returns the first value of `col` for a group of rows. It will return the first non-null + value it sees when `ignoreNulls` is set to true. If all values are null, then null is returned. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + ignorenulls : :class:`~pyspark.sql.Column` or bool + if first value is null then look for first non-null value. + + Returns + ------- + :class:`~pyspark.sql.Column` + some value of `col` for a group of rows. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [(None, 1), ("a", 2), ("a", 3), ("b", 8), ("b", 2)], ["a", "b"] + ... ).select(sf.first_value('a'), sf.first_value('b')).show() + +--------------+--------------+ + |first_value(a)|first_value(b)| + +--------------+--------------+ + | NULL| 1| + +--------------+--------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [(None, 1), ("a", 2), ("a", 3), ("b", 8), ("b", 2)], ["a", "b"] + ... ).select(sf.first_value('a', True), sf.first_value('b', True)).show() + +--------------+--------------+ + |first_value(a)|first_value(b)| + +--------------+--------------+ + | a| 1| + +--------------+--------------+ + """ + if ignoreNulls is None: + return _invoke_function_over_columns("first_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("first_value", col, ignoreNulls) + + +@try_remote_functions +def last_value(col: "ColumnOrName", ignoreNulls: Optional[Union[bool, Column]] = None) -> Column: + """Returns the last value of `col` for a group of rows. It will return the last non-null + value it sees when `ignoreNulls` is set to true. If all values are null, then null is returned. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + ignorenulls : :class:`~pyspark.sql.Column` or bool + if first value is null then look for first non-null value. + + Returns + ------- + :class:`~pyspark.sql.Column` + some value of `col` for a group of rows. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("a", 1), ("a", 2), ("a", 3), ("b", 8), (None, 2)], ["a", "b"] + ... ).select(sf.last_value('a'), sf.last_value('b')).show() + +-------------+-------------+ + |last_value(a)|last_value(b)| + +-------------+-------------+ + | NULL| 2| + +-------------+-------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("a", 1), ("a", 2), ("a", 3), ("b", 8), (None, 2)], ["a", "b"] + ... ).select(sf.last_value('a', True), sf.last_value('b', True)).show() + +-------------+-------------+ + |last_value(a)|last_value(b)| + +-------------+-------------+ + | b| 2| + +-------------+-------------+ + """ + if ignoreNulls is None: + return _invoke_function_over_columns("last_value", col) + else: + ignoreNulls = lit(ignoreNulls) if isinstance(ignoreNulls, bool) else ignoreNulls + return _invoke_function_over_columns("last_value", col, ignoreNulls) + + +@try_remote_functions +def count_if(col: "ColumnOrName") -> Column: + """Returns the number of `TRUE` values for the `col`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the number of `TRUE` values for the `col`. + + Examples + -------- + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.select(count_if(col('c2') % 2 == 0)).show() + +------------------------+ + |count_if(((c2 % 2) = 0))| + +------------------------+ + | 3| + +------------------------+ + """ + return _invoke_function_over_columns("count_if", col) + + +@try_remote_functions +def histogram_numeric(col: "ColumnOrName", nBins: "ColumnOrName") -> Column: + """Computes a histogram on numeric 'col' using nb bins. + The return value is an array of (x,y) pairs representing the centers of the + histogram's bins. As the value of 'nb' is increased, the histogram approximation + gets finer-grained, but may yield artifacts around outliers. In practice, 20-40 + histogram bins appear to work well, with more bins being required for skewed or + smaller datasets. Note that this function creates a histogram with non-uniform + bin widths. It offers no guarantees in terms of the mean-squared-error of the + histogram, but in practice is comparable to the histograms produced by the R/S-Plus + statistical computing packages. Note: the output type of the 'x' field in the return value is + propagated from the input value consumed in the aggregate function. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + nBins : :class:`~pyspark.sql.Column` or str + number of Histogram columns. + + Returns + ------- + :class:`~pyspark.sql.Column` + a histogram on numeric 'col' using nb bins. + + Examples + -------- + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.select(histogram_numeric('c2', lit(5))).show() + +------------------------+ + |histogram_numeric(c2, 5)| + +------------------------+ + | [{1, 1.0}, {2, 1....| + +------------------------+ + """ + return _invoke_function_over_columns("histogram_numeric", col, nBins) + + +@try_remote_functions +def ntile(n: int) -> Column: + """ + Window function: returns the ntile group id (from 1 to `n` inclusive) + in an ordered window partition. For example, if `n` is 4, the first + quarter of the rows will get value 1, the second quarter will get 2, + the third quarter will get 3, and the last quarter will get 4. + + This is equivalent to the NTILE function in SQL. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + n : int + an integer + + Returns + ------- + :class:`~pyspark.sql.Column` + portioned group id. + + Examples + -------- + >>> from pyspark.sql import Window + >>> df = spark.createDataFrame([("a", 1), + ... ("a", 2), + ... ("a", 3), + ... ("b", 8), + ... ("b", 2)], ["c1", "c2"]) + >>> df.show() + +---+---+ + | c1| c2| + +---+---+ + | a| 1| + | a| 2| + | a| 3| + | b| 8| + | b| 2| + +---+---+ + >>> w = Window.partitionBy("c1").orderBy("c2") + >>> df.withColumn("ntile", ntile(2).over(w)).show() + +---+---+-----+ + | c1| c2|ntile| + +---+---+-----+ + | a| 1| 1| + | a| 2| 1| + | a| 3| 2| + | b| 2| 1| + | b| 8| 2| + +---+---+-----+ + """ + return _invoke_function("ntile", int(n)) + + +# ---------------------- Date/Timestamp functions ------------------------------ + + +@try_remote_functions +def curdate() -> Column: + """ + Returns the current date at the start of query evaluation as a :class:`DateType` column. + All calls of current_date within the same query return the same value. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`~pyspark.sql.Column` + current date. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.curdate()).show() # doctest: +SKIP + +--------------+ + |current_date()| + +--------------+ + | 2022-08-26| + +--------------+ + """ + return _invoke_function("curdate") + + +@try_remote_functions +def current_date() -> Column: + """ + Returns the current date at the start of query evaluation as a :class:`DateType` column. + All calls of current_date within the same query return the same value. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + current date. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(current_date()).show() # doctest: +SKIP + +--------------+ + |current_date()| + +--------------+ + | 2022-08-26| + +--------------+ + """ + return _invoke_function("current_date") + + +@try_remote_functions +def current_timezone() -> Column: + """ + Returns the current session local timezone. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`~pyspark.sql.Column` + current session local timezone. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> spark.range(1).select(current_timezone()).show() + +-------------------+ + | current_timezone()| + +-------------------+ + |America/Los_Angeles| + +-------------------+ + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function("current_timezone") + + +@try_remote_functions +def current_timestamp() -> Column: + """ + Returns the current timestamp at the start of query evaluation as a :class:`TimestampType` + column. All calls of current_timestamp within the same query return the same value. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + current date and time. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(current_timestamp()).show(truncate=False) # doctest: +SKIP + +-----------------------+ + |current_timestamp() | + +-----------------------+ + |2022-08-26 21:23:22.716| + +-----------------------+ + """ + return _invoke_function("current_timestamp") + + +@try_remote_functions +def now() -> Column: + """ + Returns the current timestamp at the start of query evaluation. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`~pyspark.sql.Column` + current timestamp at the start of query evaluation. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(now()).show(truncate=False) # doctest: +SKIP + +-----------------------+ + |now() | + +-----------------------+ + |2022-08-26 21:23:22.716| + +-----------------------+ + """ + return _invoke_function("current_timestamp") + + +@try_remote_functions +def localtimestamp() -> Column: + """ + Returns the current timestamp without time zone at the start of query evaluation + as a timestamp without time zone column. All calls of localtimestamp within the + same query return the same value. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + current local date and time. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(localtimestamp()).show(truncate=False) # doctest: +SKIP + +-----------------------+ + |localtimestamp() | + +-----------------------+ + |2022-08-26 21:28:34.639| + +-----------------------+ + """ + return _invoke_function("localtimestamp") + + +@try_remote_functions +def date_format(date: "ColumnOrName", format: str) -> Column: + """ + Converts a date/timestamp/string to a value of string in the format specified by the date + format given by the second argument. + + A pattern could be for instance `dd.MM.yyyy` and could return a string like '18.03.1993'. All + pattern letters of `datetime pattern`_. can be used. + + .. _datetime pattern: https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + Whenever possible, use specialized functions like `year`. + + Parameters + ---------- + date : :class:`~pyspark.sql.Column` or str + input column of values to format. + format: str + format to use to represent datetime values. + + Returns + ------- + :class:`~pyspark.sql.Column` + string value representing formatted datetime. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(date_format('dt', 'MM/dd/yyy').alias('date')).collect() + [Row(date='04/08/2015')] + """ + return _invoke_function("date_format", _to_java_column(date), format) + + +@try_remote_functions +def year(col: "ColumnOrName") -> Column: + """ + Extract the year of a given date/timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + year part of the date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(year('dt').alias('year')).collect() + [Row(year=2015)] + """ + return _invoke_function_over_columns("year", col) + + +@try_remote_functions +def quarter(col: "ColumnOrName") -> Column: + """ + Extract the quarter of a given date/timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + quarter of the date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(quarter('dt').alias('quarter')).collect() + [Row(quarter=2)] + """ + return _invoke_function_over_columns("quarter", col) + + +@try_remote_functions +def month(col: "ColumnOrName") -> Column: + """ + Extract the month of a given date/timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + month part of the date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(month('dt').alias('month')).collect() + [Row(month=4)] + """ + return _invoke_function_over_columns("month", col) + + +@try_remote_functions +def dayofweek(col: "ColumnOrName") -> Column: + """ + Extract the day of the week of a given date/timestamp as integer. + Ranges from 1 for a Sunday through to 7 for a Saturday + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + day of the week for given date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(dayofweek('dt').alias('day')).collect() + [Row(day=4)] + """ + return _invoke_function_over_columns("dayofweek", col) + + +@try_remote_functions +def dayofmonth(col: "ColumnOrName") -> Column: + """ + Extract the day of the month of a given date/timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + day of the month for given date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(dayofmonth('dt').alias('day')).collect() + [Row(day=8)] + """ + return _invoke_function_over_columns("dayofmonth", col) + + +@try_remote_functions +def day(col: "ColumnOrName") -> Column: + """ + Extract the day of the month of a given date/timestamp as integer. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + day of the month for given date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(day('dt').alias('day')).collect() + [Row(day=8)] + """ + return _invoke_function_over_columns("day", col) + + +@try_remote_functions +def dayofyear(col: "ColumnOrName") -> Column: + """ + Extract the day of the year of a given date/timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + day of the year for given date/timestamp as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(dayofyear('dt').alias('day')).collect() + [Row(day=98)] + """ + return _invoke_function_over_columns("dayofyear", col) + + +@try_remote_functions +def hour(col: "ColumnOrName") -> Column: + """ + Extract the hours of a given timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + hour part of the timestamp as integer. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select(hour('ts').alias('hour')).collect() + [Row(hour=13)] + """ + return _invoke_function_over_columns("hour", col) + + +@try_remote_functions +def minute(col: "ColumnOrName") -> Column: + """ + Extract the minutes of a given timestamp as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + minutes part of the timestamp as integer. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select(minute('ts').alias('minute')).collect() + [Row(minute=8)] + """ + return _invoke_function_over_columns("minute", col) + + +@try_remote_functions +def second(col: "ColumnOrName") -> Column: + """ + Extract the seconds of a given date as integer. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + `seconds` part of the timestamp as integer. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select(second('ts').alias('second')).collect() + [Row(second=15)] + """ + return _invoke_function_over_columns("second", col) + + +@try_remote_functions +def weekofyear(col: "ColumnOrName") -> Column: + """ + Extract the week number of a given date as integer. + A week is considered to start on a Monday and week 1 is the first week with more than 3 days, + as defined by ISO 8601 + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + `week` of the year for given date as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(weekofyear(df.dt).alias('week')).collect() + [Row(week=15)] + """ + return _invoke_function_over_columns("weekofyear", col) + + +@try_remote_functions +def weekday(col: "ColumnOrName") -> Column: + """ + Returns the day of the week for date/timestamp (0 = Monday, 1 = Tuesday, ..., 6 = Sunday). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date/timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the day of the week for date/timestamp (0 = Monday, 1 = Tuesday, ..., 6 = Sunday). + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(weekday('dt').alias('day')).show() + +---+ + |day| + +---+ + | 2| + +---+ + """ + return _invoke_function_over_columns("weekday", col) + + +@try_remote_functions +def extract(field: "ColumnOrName", source: "ColumnOrName") -> Column: + """ + Extracts a part of the date/timestamp or interval source. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + field : :class:`~pyspark.sql.Column` or str + selects which part of the source should be extracted. + source : :class:`~pyspark.sql.Column` or str + a date/timestamp or interval column from where `field` should be extracted. + + Returns + ------- + :class:`~pyspark.sql.Column` + a part of the date/timestamp or interval source. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select( + ... extract(lit('YEAR'), 'ts').alias('year'), + ... extract(lit('month'), 'ts').alias('month'), + ... extract(lit('WEEK'), 'ts').alias('week'), + ... extract(lit('D'), 'ts').alias('day'), + ... extract(lit('M'), 'ts').alias('minute'), + ... extract(lit('S'), 'ts').alias('second') + ... ).collect() + [Row(year=2015, month=4, week=15, day=8, minute=8, second=Decimal('15.000000'))] + """ + return _invoke_function_over_columns("extract", field, source) + + +@try_remote_functions +def date_part(field: "ColumnOrName", source: "ColumnOrName") -> Column: + """ + Extracts a part of the date/timestamp or interval source. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + field : :class:`~pyspark.sql.Column` or str + selects which part of the source should be extracted, and supported string values + are as same as the fields of the equivalent function `extract`. + source : :class:`~pyspark.sql.Column` or str + a date/timestamp or interval column from where `field` should be extracted. + + Returns + ------- + :class:`~pyspark.sql.Column` + a part of the date/timestamp or interval source. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select( + ... date_part(lit('YEAR'), 'ts').alias('year'), + ... date_part(lit('month'), 'ts').alias('month'), + ... date_part(lit('WEEK'), 'ts').alias('week'), + ... date_part(lit('D'), 'ts').alias('day'), + ... date_part(lit('M'), 'ts').alias('minute'), + ... date_part(lit('S'), 'ts').alias('second') + ... ).collect() + [Row(year=2015, month=4, week=15, day=8, minute=8, second=Decimal('15.000000'))] + """ + return _invoke_function_over_columns("date_part", field, source) + + +@try_remote_functions +def datepart(field: "ColumnOrName", source: "ColumnOrName") -> Column: + """ + Extracts a part of the date/timestamp or interval source. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + field : :class:`~pyspark.sql.Column` or str + selects which part of the source should be extracted, and supported string values + are as same as the fields of the equivalent function `extract`. + source : :class:`~pyspark.sql.Column` or str + a date/timestamp or interval column from where `field` should be extracted. + + Returns + ------- + :class:`~pyspark.sql.Column` + a part of the date/timestamp or interval source. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame([(datetime.datetime(2015, 4, 8, 13, 8, 15),)], ['ts']) + >>> df.select( + ... datepart(lit('YEAR'), 'ts').alias('year'), + ... datepart(lit('month'), 'ts').alias('month'), + ... datepart(lit('WEEK'), 'ts').alias('week'), + ... datepart(lit('D'), 'ts').alias('day'), + ... datepart(lit('M'), 'ts').alias('minute'), + ... datepart(lit('S'), 'ts').alias('second') + ... ).collect() + [Row(year=2015, month=4, week=15, day=8, minute=8, second=Decimal('15.000000'))] + """ + return _invoke_function_over_columns("datepart", field, source) + + +@try_remote_functions +def make_date(year: "ColumnOrName", month: "ColumnOrName", day: "ColumnOrName") -> Column: + """ + Returns a column with a date built from the year, month and day columns. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + year : :class:`~pyspark.sql.Column` or str + The year to build the date + month : :class:`~pyspark.sql.Column` or str + The month to build the date + day : :class:`~pyspark.sql.Column` or str + The day to build the date + + Returns + ------- + :class:`~pyspark.sql.Column` + a date built from given parts. + + Examples + -------- + >>> df = spark.createDataFrame([(2020, 6, 26)], ['Y', 'M', 'D']) + >>> df.select(make_date(df.Y, df.M, df.D).alias("datefield")).collect() + [Row(datefield=datetime.date(2020, 6, 26))] + """ + return _invoke_function_over_columns("make_date", year, month, day) + + +@try_remote_functions +def date_add(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + """ + Returns the date that is `days` days after `start`. If `days` is a negative value + then these amount of days will be deducted from `start`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + start : :class:`~pyspark.sql.Column` or str + date column to work on. + days : :class:`~pyspark.sql.Column` or str or int + how many days after the given date to calculate. + Accepts negative value as well to calculate backwards in time. + + Returns + ------- + :class:`~pyspark.sql.Column` + a date after/before given number of days. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08', 2,)], ['dt', 'add']) + >>> df.select(date_add(df.dt, 1).alias('next_date')).collect() + [Row(next_date=datetime.date(2015, 4, 9))] + >>> df.select(date_add(df.dt, df.add.cast('integer')).alias('next_date')).collect() + [Row(next_date=datetime.date(2015, 4, 10))] + >>> df.select(date_add('dt', -1).alias('prev_date')).collect() + [Row(prev_date=datetime.date(2015, 4, 7))] + """ + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("date_add", start, days) + + +@try_remote_functions +def dateadd(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + """ + Returns the date that is `days` days after `start`. If `days` is a negative value + then these amount of days will be deducted from `start`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + start : :class:`~pyspark.sql.Column` or str + date column to work on. + days : :class:`~pyspark.sql.Column` or str or int + how many days after the given date to calculate. + Accepts negative value as well to calculate backwards in time. + + Returns + ------- + :class:`~pyspark.sql.Column` + a date after/before given number of days. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [('2015-04-08', 2,)], ['dt', 'add'] + ... ).select(sf.dateadd("dt", 1)).show() + +---------------+ + |date_add(dt, 1)| + +---------------+ + | 2015-04-09| + +---------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [('2015-04-08', 2,)], ['dt', 'add'] + ... ).select(sf.dateadd("dt", sf.lit(2))).show() + +---------------+ + |date_add(dt, 2)| + +---------------+ + | 2015-04-10| + +---------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [('2015-04-08', 2,)], ['dt', 'add'] + ... ).select(sf.dateadd("dt", -1)).show() + +----------------+ + |date_add(dt, -1)| + +----------------+ + | 2015-04-07| + +----------------+ + """ + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("dateadd", start, days) + + +@try_remote_functions +def date_sub(start: "ColumnOrName", days: Union["ColumnOrName", int]) -> Column: + """ + Returns the date that is `days` days before `start`. If `days` is a negative value + then these amount of days will be added to `start`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + start : :class:`~pyspark.sql.Column` or str + date column to work on. + days : :class:`~pyspark.sql.Column` or str or int + how many days before the given date to calculate. + Accepts negative value as well to calculate forward in time. + + Returns + ------- + :class:`~pyspark.sql.Column` + a date before/after given number of days. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08', 2,)], ['dt', 'sub']) + >>> df.select(date_sub(df.dt, 1).alias('prev_date')).collect() + [Row(prev_date=datetime.date(2015, 4, 7))] + >>> df.select(date_sub(df.dt, df.sub.cast('integer')).alias('prev_date')).collect() + [Row(prev_date=datetime.date(2015, 4, 6))] + >>> df.select(date_sub('dt', -1).alias('next_date')).collect() + [Row(next_date=datetime.date(2015, 4, 9))] + """ + days = lit(days) if isinstance(days, int) else days + return _invoke_function_over_columns("date_sub", start, days) + + +@try_remote_functions +def datediff(end: "ColumnOrName", start: "ColumnOrName") -> Column: + """ + Returns the number of days from `start` to `end`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + end : :class:`~pyspark.sql.Column` or str + to date column to work on. + start : :class:`~pyspark.sql.Column` or str + from date column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + difference in days between two dates. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08','2015-05-10')], ['d1', 'd2']) + >>> df.select(datediff(df.d2, df.d1).alias('diff')).collect() + [Row(diff=32)] + """ + return _invoke_function_over_columns("datediff", end, start) + + +@try_remote_functions +def date_diff(end: "ColumnOrName", start: "ColumnOrName") -> Column: + """ + Returns the number of days from `start` to `end`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + end : :class:`~pyspark.sql.Column` or str + to date column to work on. + start : :class:`~pyspark.sql.Column` or str + from date column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + difference in days between two dates. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08','2015-05-10')], ['d1', 'd2']) + >>> df.select(date_diff(df.d2, df.d1).alias('diff')).collect() + [Row(diff=32)] + """ + return _invoke_function_over_columns("date_diff", end, start) + + +@try_remote_functions +def date_from_unix_date(days: "ColumnOrName") -> Column: + """ + Create date from the number of `days` since 1970-01-01. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + days : :class:`~pyspark.sql.Column` or str + the target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the date from the number of days since 1970-01-01. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(date_from_unix_date(lit(1))).show() + +----------------------+ + |date_from_unix_date(1)| + +----------------------+ + | 1970-01-02| + +----------------------+ + """ + return _invoke_function_over_columns("date_from_unix_date", days) + + +@try_remote_functions +def add_months(start: "ColumnOrName", months: Union["ColumnOrName", int]) -> Column: + """ + Returns the date that is `months` months after `start`. If `months` is a negative value + then these amount of months will be deducted from the `start`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + start : :class:`~pyspark.sql.Column` or str + date column to work on. + months : :class:`~pyspark.sql.Column` or str or int + how many months after the given date to calculate. + Accepts negative value as well to calculate backwards. + + Returns + ------- + :class:`~pyspark.sql.Column` + a date after/before given number of months. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08', 2)], ['dt', 'add']) + >>> df.select(add_months(df.dt, 1).alias('next_month')).collect() + [Row(next_month=datetime.date(2015, 5, 8))] + >>> df.select(add_months(df.dt, df.add.cast('integer')).alias('next_month')).collect() + [Row(next_month=datetime.date(2015, 6, 8))] + >>> df.select(add_months('dt', -2).alias('prev_month')).collect() + [Row(prev_month=datetime.date(2015, 2, 8))] + """ + months = lit(months) if isinstance(months, int) else months + return _invoke_function_over_columns("add_months", start, months) + + +@try_remote_functions +def months_between(date1: "ColumnOrName", date2: "ColumnOrName", roundOff: bool = True) -> Column: + """ + Returns number of months between dates date1 and date2. + If date1 is later than date2, then the result is positive. + A whole number is returned if both inputs have the same day of month or both are the last day + of their respective months. Otherwise, the difference is calculated assuming 31 days per month. + The result is rounded off to 8 digits unless `roundOff` is set to `False`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + date1 : :class:`~pyspark.sql.Column` or str + first date column. + date2 : :class:`~pyspark.sql.Column` or str + second date column. + roundOff : bool, optional + whether to round (to 8 digits) the final value or not (default: True). + + Returns + ------- + :class:`~pyspark.sql.Column` + number of months between two dates. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00', '1996-10-30')], ['date1', 'date2']) + >>> df.select(months_between(df.date1, df.date2).alias('months')).collect() + [Row(months=3.94959677)] + >>> df.select(months_between(df.date1, df.date2, False).alias('months')).collect() + [Row(months=3.9495967741935485)] + """ + return _invoke_function( + "months_between", _to_java_column(date1), _to_java_column(date2), roundOff + ) + + +@try_remote_functions +def to_date(col: "ColumnOrName", format: Optional[str] = None) -> Column: + """Converts a :class:`~pyspark.sql.Column` into :class:`pyspark.sql.types.DateType` + using the optionally specified format. Specify formats according to `datetime pattern`_. + By default, it follows casting rules to :class:`pyspark.sql.types.DateType` if the format + is omitted. Equivalent to ``col.cast("date")``. + + .. _datetime pattern: https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html + + .. versionadded:: 2.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to convert. + format: str, optional + format to use to convert date values. + + Returns + ------- + :class:`~pyspark.sql.Column` + date value as :class:`pyspark.sql.types.DateType` type. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) + >>> df.select(to_date(df.t).alias('date')).collect() + [Row(date=datetime.date(1997, 2, 28))] + + >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) + >>> df.select(to_date(df.t, 'yyyy-MM-dd HH:mm:ss').alias('date')).collect() + [Row(date=datetime.date(1997, 2, 28))] + """ + if format is None: + return _invoke_function_over_columns("to_date", col) + else: + return _invoke_function("to_date", _to_java_column(col), format) + + +@try_remote_functions +def unix_date(col: "ColumnOrName") -> Column: + """Returns the number of days since 1970-01-01. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([('1970-01-02',)], ['t']) + >>> df.select(unix_date(to_date(df.t)).alias('n')).collect() + [Row(n=1)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("unix_date", col) + + +@try_remote_functions +def unix_micros(col: "ColumnOrName") -> Column: + """Returns the number of microseconds since 1970-01-01 00:00:00 UTC. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([('2015-07-22 10:00:00',)], ['t']) + >>> df.select(unix_micros(to_timestamp(df.t)).alias('n')).collect() + [Row(n=1437584400000000)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("unix_micros", col) + + +@try_remote_functions +def unix_millis(col: "ColumnOrName") -> Column: + """Returns the number of milliseconds since 1970-01-01 00:00:00 UTC. + Truncates higher levels of precision. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([('2015-07-22 10:00:00',)], ['t']) + >>> df.select(unix_millis(to_timestamp(df.t)).alias('n')).collect() + [Row(n=1437584400000)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("unix_millis", col) + + +@try_remote_functions +def unix_seconds(col: "ColumnOrName") -> Column: + """Returns the number of seconds since 1970-01-01 00:00:00 UTC. + Truncates higher levels of precision. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([('2015-07-22 10:00:00',)], ['t']) + >>> df.select(unix_seconds(to_timestamp(df.t)).alias('n')).collect() + [Row(n=1437584400)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("unix_seconds", col) + + +@overload +def to_timestamp(col: "ColumnOrName") -> Column: + ... + + +@overload +def to_timestamp(col: "ColumnOrName", format: str) -> Column: + ... + + +@try_remote_functions +def to_timestamp(col: "ColumnOrName", format: Optional[str] = None) -> Column: + """Converts a :class:`~pyspark.sql.Column` into :class:`pyspark.sql.types.TimestampType` + using the optionally specified format. Specify formats according to `datetime pattern`_. + By default, it follows casting rules to :class:`pyspark.sql.types.TimestampType` if the format + is omitted. Equivalent to ``col.cast("timestamp")``. + + .. _datetime pattern: https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html + + .. versionadded:: 2.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column values to convert. + format: str, optional + format to use to convert timestamp values. + + Returns + ------- + :class:`~pyspark.sql.Column` + timestamp value as :class:`pyspark.sql.types.TimestampType` type. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) + >>> df.select(to_timestamp(df.t).alias('dt')).collect() + [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] + + >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) + >>> df.select(to_timestamp(df.t, 'yyyy-MM-dd HH:mm:ss').alias('dt')).collect() + [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] + """ + if format is None: + return _invoke_function_over_columns("to_timestamp", col) + else: + return _invoke_function("to_timestamp", _to_java_column(col), format) + + +@try_remote_functions +def try_to_timestamp(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + """ + Parses the `col` with the `format` to a timestamp. The function always + returns null on an invalid input with/without ANSI SQL mode enabled. The result data type is + consistent with the value of configuration `spark.sql.timestampType`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column values to convert. + format: str, optional + format to use to convert timestamp values. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) + >>> df.select(try_to_timestamp(df.t).alias('dt')).collect() + [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] + + >>> df.select(try_to_timestamp(df.t, lit('yyyy-MM-dd HH:mm:ss')).alias('dt')).collect() + [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] + """ + if format is not None: + return _invoke_function_over_columns("try_to_timestamp", col, format) + else: + return _invoke_function_over_columns("try_to_timestamp", col) + + +@try_remote_functions +def xpath(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a string array of values within the nodes of xml that match the XPath expression. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame( + ... [('b1b2b3c1c2',)], ['x']) + >>> df.select(xpath(df.x, lit('a/b/text()')).alias('r')).collect() + [Row(r=['b1', 'b2', 'b3'])] + """ + return _invoke_function_over_columns("xpath", xml, path) + + +@try_remote_functions +def xpath_boolean(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns true if the XPath expression evaluates to true, or if a matching node is found. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('1',)], ['x']) + >>> df.select(xpath_boolean(df.x, lit('a/b')).alias('r')).collect() + [Row(r=True)] + """ + return _invoke_function_over_columns("xpath_boolean", xml, path) + + +@try_remote_functions +def xpath_double(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a double value, the value zero if no match is found, + or NaN if a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('12',)], ['x']) + >>> df.select(xpath_double(df.x, lit('sum(a/b)')).alias('r')).collect() + [Row(r=3.0)] + """ + return _invoke_function_over_columns("xpath_double", xml, path) + + +@try_remote_functions +def xpath_number(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a double value, the value zero if no match is found, + or NaN if a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [('12',)], ['x'] + ... ).select(sf.xpath_number('x', sf.lit('sum(a/b)'))).show() + +-------------------------+ + |xpath_number(x, sum(a/b))| + +-------------------------+ + | 3.0| + +-------------------------+ + """ + return _invoke_function_over_columns("xpath_number", xml, path) + + +@try_remote_functions +def xpath_float(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a float value, the value zero if no match is found, + or NaN if a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('12',)], ['x']) + >>> df.select(xpath_float(df.x, lit('sum(a/b)')).alias('r')).collect() + [Row(r=3.0)] + """ + return _invoke_function_over_columns("xpath_float", xml, path) + + +@try_remote_functions +def xpath_int(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns an integer value, or the value zero if no match is found, + or a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('12',)], ['x']) + >>> df.select(xpath_int(df.x, lit('sum(a/b)')).alias('r')).collect() + [Row(r=3)] + """ + return _invoke_function_over_columns("xpath_int", xml, path) + + +@try_remote_functions +def xpath_long(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a long integer value, or the value zero if no match is found, + or a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('12',)], ['x']) + >>> df.select(xpath_long(df.x, lit('sum(a/b)')).alias('r')).collect() + [Row(r=3)] + """ + return _invoke_function_over_columns("xpath_long", xml, path) + + +@try_remote_functions +def xpath_short(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns a short integer value, or the value zero if no match is found, + or a match is found but the value is non-numeric. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('12',)], ['x']) + >>> df.select(xpath_short(df.x, lit('sum(a/b)')).alias('r')).collect() + [Row(r=3)] + """ + return _invoke_function_over_columns("xpath_short", xml, path) + + +@try_remote_functions +def xpath_string(xml: "ColumnOrName", path: "ColumnOrName") -> Column: + """ + Returns the text contents of the first xml node that matches the XPath expression. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.createDataFrame([('bcc',)], ['x']) + >>> df.select(xpath_string(df.x, lit('a/c')).alias('r')).collect() + [Row(r='cc')] + """ + return _invoke_function_over_columns("xpath_string", xml, path) + + +@try_remote_functions +def trunc(date: "ColumnOrName", format: str) -> Column: + """ + Returns date truncated to the unit specified by the format. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + date : :class:`~pyspark.sql.Column` or str + input column of values to truncate. + format : str + 'year', 'yyyy', 'yy' to truncate by year, + or 'month', 'mon', 'mm' to truncate by month + Other options are: 'week', 'quarter' + + Returns + ------- + :class:`~pyspark.sql.Column` + truncated date. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28',)], ['d']) + >>> df.select(trunc(df.d, 'year').alias('year')).collect() + [Row(year=datetime.date(1997, 1, 1))] + >>> df.select(trunc(df.d, 'mon').alias('month')).collect() + [Row(month=datetime.date(1997, 2, 1))] + """ + return _invoke_function("trunc", _to_java_column(date), format) + + +@try_remote_functions +def date_trunc(format: str, timestamp: "ColumnOrName") -> Column: + """ + Returns timestamp truncated to the unit specified by the format. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + format : str + 'year', 'yyyy', 'yy' to truncate by year, + 'month', 'mon', 'mm' to truncate by month, + 'day', 'dd' to truncate by day, + Other options are: + 'microsecond', 'millisecond', 'second', 'minute', 'hour', 'week', 'quarter' + timestamp : :class:`~pyspark.sql.Column` or str + input column of values to truncate. + + Returns + ------- + :class:`~pyspark.sql.Column` + truncated timestamp. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 05:02:11',)], ['t']) + >>> df.select(date_trunc('year', df.t).alias('year')).collect() + [Row(year=datetime.datetime(1997, 1, 1, 0, 0))] + >>> df.select(date_trunc('mon', df.t).alias('month')).collect() + [Row(month=datetime.datetime(1997, 2, 1, 0, 0))] + """ + return _invoke_function("date_trunc", format, _to_java_column(timestamp)) + + +@try_remote_functions +def next_day(date: "ColumnOrName", dayOfWeek: str) -> Column: + """ + Returns the first date which is later than the value of the date column + based on second `week day` argument. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + date : :class:`~pyspark.sql.Column` or str + target column to compute on. + dayOfWeek : str + day of the week, case-insensitive, accepts: + "Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun" + + Returns + ------- + :class:`~pyspark.sql.Column` + the column of computed results. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-07-27',)], ['d']) + >>> df.select(next_day(df.d, 'Sun').alias('date')).collect() + [Row(date=datetime.date(2015, 8, 2))] + """ + return _invoke_function("next_day", _to_java_column(date), dayOfWeek) + + +@try_remote_functions +def last_day(date: "ColumnOrName") -> Column: + """ + Returns the last day of the month which the given date belongs to. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + date : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + last day of the month. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-10',)], ['d']) + >>> df.select(last_day(df.d).alias('date')).collect() + [Row(date=datetime.date(1997, 2, 28))] + """ + return _invoke_function("last_day", _to_java_column(date)) + + +@try_remote_functions +def from_unixtime(timestamp: "ColumnOrName", format: str = "yyyy-MM-dd HH:mm:ss") -> Column: + """ + Converts the number of seconds from unix epoch (1970-01-01 00:00:00 UTC) to a string + representing the timestamp of that moment in the current system time zone in the given + format. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + column of unix time values. + format : str, optional + format to use to convert to (default: yyyy-MM-dd HH:mm:ss) + + Returns + ------- + :class:`~pyspark.sql.Column` + formatted timestamp as string. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> time_df = spark.createDataFrame([(1428476400,)], ['unix_time']) + >>> time_df.select(from_unixtime('unix_time').alias('ts')).collect() + [Row(ts='2015-04-08 00:00:00')] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function("from_unixtime", _to_java_column(timestamp), format) + + +@overload +def unix_timestamp(timestamp: "ColumnOrName", format: str = ...) -> Column: + ... + + +@overload +def unix_timestamp() -> Column: + ... + + +@try_remote_functions +def unix_timestamp( + timestamp: Optional["ColumnOrName"] = None, format: str = "yyyy-MM-dd HH:mm:ss" +) -> Column: + """ + Convert time string with given pattern ('yyyy-MM-dd HH:mm:ss', by default) + to Unix time stamp (in seconds), using the default timezone and the default + locale, returns null if failed. + + if `timestamp` is None, then it returns current timestamp. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str, optional + timestamps of string values. + format : str, optional + alternative format to use for converting (default: yyyy-MM-dd HH:mm:ss). + + Returns + ------- + :class:`~pyspark.sql.Column` + unix time as long integer. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> time_df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> time_df.select(unix_timestamp('dt', 'yyyy-MM-dd').alias('unix_time')).collect() + [Row(unix_time=1428476400)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + if timestamp is None: + return _invoke_function("unix_timestamp") + return _invoke_function("unix_timestamp", _to_java_column(timestamp), format) + + +@try_remote_functions +def from_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") -> Column: + """ + This is a common function for databases supporting TIMESTAMP WITHOUT TIMEZONE. This function + takes a timestamp which is timezone-agnostic, and interprets it as a timestamp in UTC, and + renders that timestamp as a timestamp in the given time zone. + + However, timestamp in Spark represents number of microseconds from the Unix epoch, which is not + timezone-agnostic. So in Spark this function just shift the timestamp value from UTC timezone to + the given timezone. + + This function may return confusing result if the input is a string with timezone, e.g. + '2018-03-13T06:18:23+00:00'. The reason is that, Spark firstly cast the string to timestamp + according to the timezone in the string, and finally display the result by converting the + timestamp to string according to the session local timezone. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + the column that contains timestamps + tz : :class:`~pyspark.sql.Column` or str + A string detailing the time zone ID that the input should be adjusted to. It should + be in the format of either region-based zone IDs or zone offsets. Region IDs must + have the form 'area/city', such as 'America/Los_Angeles'. Zone offsets must be in + the format '(+|-)HH:mm', for example '-08:00' or '+01:00'. Also 'UTC' and 'Z' are + supported as aliases of '+00:00'. Other short names are not recommended to use + because they can be ambiguous. + + .. versionchanged:: 2.4 + `tz` can take a :class:`~pyspark.sql.Column` containing timezone ID strings. + + Returns + ------- + :class:`~pyspark.sql.Column` + timestamp value represented in given timezone. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00', 'JST')], ['ts', 'tz']) + >>> df.select(from_utc_timestamp(df.ts, "PST").alias('local_time')).collect() + [Row(local_time=datetime.datetime(1997, 2, 28, 2, 30))] + >>> df.select(from_utc_timestamp(df.ts, df.tz).alias('local_time')).collect() + [Row(local_time=datetime.datetime(1997, 2, 28, 19, 30))] + """ + if isinstance(tz, Column): + tz = _to_java_column(tz) + return _invoke_function("from_utc_timestamp", _to_java_column(timestamp), tz) + + +@try_remote_functions +def to_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") -> Column: + """ + This is a common function for databases supporting TIMESTAMP WITHOUT TIMEZONE. This function + takes a timestamp which is timezone-agnostic, and interprets it as a timestamp in the given + timezone, and renders that timestamp as a timestamp in UTC. + + However, timestamp in Spark represents number of microseconds from the Unix epoch, which is not + timezone-agnostic. So in Spark this function just shift the timestamp value from the given + timezone to UTC timezone. + + This function may return confusing result if the input is a string with timezone, e.g. + '2018-03-13T06:18:23+00:00'. The reason is that, Spark firstly cast the string to timestamp + according to the timezone in the string, and finally display the result by converting the + timestamp to string according to the session local timezone. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + the column that contains timestamps + tz : :class:`~pyspark.sql.Column` or str + A string detailing the time zone ID that the input should be adjusted to. It should + be in the format of either region-based zone IDs or zone offsets. Region IDs must + have the form 'area/city', such as 'America/Los_Angeles'. Zone offsets must be in + the format '(+|-)HH:mm', for example '-08:00' or '+01:00'. Also 'UTC' and 'Z' are + supported as aliases of '+00:00'. Other short names are not recommended to use + because they can be ambiguous. + + .. versionchanged:: 2.4.0 + `tz` can take a :class:`~pyspark.sql.Column` containing timezone ID strings. + + Returns + ------- + :class:`~pyspark.sql.Column` + timestamp value represented in UTC timezone. + + Examples + -------- + >>> df = spark.createDataFrame([('1997-02-28 10:30:00', 'JST')], ['ts', 'tz']) + >>> df.select(to_utc_timestamp(df.ts, "PST").alias('utc_time')).collect() + [Row(utc_time=datetime.datetime(1997, 2, 28, 18, 30))] + >>> df.select(to_utc_timestamp(df.ts, df.tz).alias('utc_time')).collect() + [Row(utc_time=datetime.datetime(1997, 2, 28, 1, 30))] + """ + if isinstance(tz, Column): + tz = _to_java_column(tz) + return _invoke_function("to_utc_timestamp", _to_java_column(timestamp), tz) + + +@try_remote_functions +def timestamp_seconds(col: "ColumnOrName") -> Column: + """ + Converts the number of seconds from the Unix epoch (1970-01-01T00:00:00Z) + to a timestamp. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + unix time values. + + Returns + ------- + :class:`~pyspark.sql.Column` + converted timestamp value. + + Examples + -------- + >>> from pyspark.sql.functions import timestamp_seconds + >>> spark.conf.set("spark.sql.session.timeZone", "UTC") + >>> time_df = spark.createDataFrame([(1230219000,)], ['unix_time']) + >>> time_df.select(timestamp_seconds(time_df.unix_time).alias('ts')).show() + +-------------------+ + | ts| + +-------------------+ + |2008-12-25 15:30:00| + +-------------------+ + >>> time_df.select(timestamp_seconds('unix_time').alias('ts')).printSchema() + root + |-- ts: timestamp (nullable = true) + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + + return _invoke_function_over_columns("timestamp_seconds", col) + + +@try_remote_functions +def timestamp_millis(col: "ColumnOrName") -> Column: + """ + Creates timestamp from the number of milliseconds since UTC epoch. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + unix time values. + + Returns + ------- + :class:`~pyspark.sql.Column` + converted timestamp value. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "UTC") + >>> time_df = spark.createDataFrame([(1230219000,)], ['unix_time']) + >>> time_df.select(timestamp_millis(time_df.unix_time).alias('ts')).show() + +-------------------+ + | ts| + +-------------------+ + |1970-01-15 05:43:39| + +-------------------+ + >>> time_df.select(timestamp_millis('unix_time').alias('ts')).printSchema() + root + |-- ts: timestamp (nullable = true) + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("timestamp_millis", col) + + +@try_remote_functions +def timestamp_micros(col: "ColumnOrName") -> Column: + """ + Creates timestamp from the number of microseconds since UTC epoch. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + unix time values. + + Returns + ------- + :class:`~pyspark.sql.Column` + converted timestamp value. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "UTC") + >>> time_df = spark.createDataFrame([(1230219000,)], ['unix_time']) + >>> time_df.select(timestamp_micros(time_df.unix_time).alias('ts')).show() + +--------------------+ + | ts| + +--------------------+ + |1970-01-01 00:20:...| + +--------------------+ + >>> time_df.select(timestamp_micros('unix_time').alias('ts')).printSchema() + root + |-- ts: timestamp (nullable = true) + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns("timestamp_micros", col) + + +@try_remote_functions +def window( + timeColumn: "ColumnOrName", + windowDuration: str, + slideDuration: Optional[str] = None, + startTime: Optional[str] = None, +) -> Column: + """Bucketize rows into one or more time windows given a timestamp specifying column. Window + starts are inclusive but the window ends are exclusive, e.g. 12:05 will be in the window + [12:05,12:10) but not in [12:00,12:05). Windows can support microsecond precision. Windows in + the order of months are not supported. + + The time column must be of :class:`pyspark.sql.types.TimestampType`. + + Durations are provided as strings, e.g. '1 second', '1 day 12 hours', '2 minutes'. Valid + interval strings are 'week', 'day', 'hour', 'minute', 'second', 'millisecond', 'microsecond'. + If the ``slideDuration`` is not provided, the windows will be tumbling windows. + + The startTime is the offset with respect to 1970-01-01 00:00:00 UTC with which to start + window intervals. For example, in order to have hourly tumbling windows that start 15 minutes + past the hour, e.g. 12:15-13:15, 13:15-14:15... provide `startTime` as `15 minutes`. + + The output column will be a struct called 'window' by default with the nested columns 'start' + and 'end', where 'start' and 'end' will be of :class:`pyspark.sql.types.TimestampType`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timeColumn : :class:`~pyspark.sql.Column` + The column or the expression to use as the timestamp for windowing by time. + The time column must be of TimestampType or TimestampNTZType. + windowDuration : str + A string specifying the width of the window, e.g. `10 minutes`, + `1 second`. Check `org.apache.spark.unsafe.types.CalendarInterval` for + valid duration identifiers. Note that the duration is a fixed length of + time, and does not vary over time according to a calendar. For example, + `1 day` always means 86,400,000 milliseconds, not a calendar day. + slideDuration : str, optional + A new window will be generated every `slideDuration`. Must be less than + or equal to the `windowDuration`. Check + `org.apache.spark.unsafe.types.CalendarInterval` for valid duration + identifiers. This duration is likewise absolute, and does not vary + according to a calendar. + startTime : str, optional + The offset with respect to 1970-01-01 00:00:00 UTC with which to start + window intervals. For example, in order to have hourly tumbling windows that + start 15 minutes past the hour, e.g. 12:15-13:15, 13:15-14:15... provide + `startTime` as `15 minutes`. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame( + ... [(datetime.datetime(2016, 3, 11, 9, 0, 7), 1)], + ... ).toDF("date", "val") + >>> w = df.groupBy(window("date", "5 seconds")).agg(sum("val").alias("sum")) + >>> w.select(w.window.start.cast("string").alias("start"), + ... w.window.end.cast("string").alias("end"), "sum").collect() + [Row(start='2016-03-11 09:00:05', end='2016-03-11 09:00:10', sum=1)] + """ + + def check_string_field(field, fieldName): # type: ignore[no-untyped-def] + if not field or type(field) is not str: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": fieldName, "arg_type": type(field).__name__}, + ) + + time_col = _to_java_column(timeColumn) + check_string_field(windowDuration, "windowDuration") + if slideDuration and startTime: + check_string_field(slideDuration, "slideDuration") + check_string_field(startTime, "startTime") + return _invoke_function("window", time_col, windowDuration, slideDuration, startTime) + elif slideDuration: + check_string_field(slideDuration, "slideDuration") + return _invoke_function("window", time_col, windowDuration, slideDuration) + elif startTime: + check_string_field(startTime, "startTime") + return _invoke_function("window", time_col, windowDuration, windowDuration, startTime) + else: + return _invoke_function("window", time_col, windowDuration) + + +@try_remote_functions +def window_time( + windowColumn: "ColumnOrName", +) -> Column: + """Computes the event time from a window column. The column window values are produced + by window aggregating operators and are of type `STRUCT` + where start is inclusive and end is exclusive. The event time of records produced by window + aggregating operators can be computed as ``window_time(window)`` and are + ``window.end - lit(1).alias("microsecond")`` (as microsecond is the minimal supported event + time precision). The window column must be one produced by a window aggregating operator. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + windowColumn : :class:`~pyspark.sql.Column` + The window column of a window aggregate records. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> import datetime + >>> df = spark.createDataFrame( + ... [(datetime.datetime(2016, 3, 11, 9, 0, 7), 1)], + ... ).toDF("date", "val") + + Group the data into 5 second time windows and aggregate as sum. + + >>> w = df.groupBy(window("date", "5 seconds")).agg(sum("val").alias("sum")) + + Extract the window event time using the window_time function. + + >>> w.select( + ... w.window.end.cast("string").alias("end"), + ... window_time(w.window).cast("string").alias("window_time"), + ... "sum" + ... ).collect() + [Row(end='2016-03-11 09:00:10', window_time='2016-03-11 09:00:09.999999', sum=1)] + """ + window_col = _to_java_column(windowColumn) + return _invoke_function("window_time", window_col) + + +@try_remote_functions +def session_window(timeColumn: "ColumnOrName", gapDuration: Union[Column, str]) -> Column: + """ + Generates session window given a timestamp specifying column. + Session window is one of dynamic windows, which means the length of window is varying + according to the given inputs. The length of session window is defined as "the timestamp + of latest input of the session + gap duration", so when the new inputs are bound to the + current session window, the end time of session window can be expanded according to the new + inputs. + Windows can support microsecond precision. Windows in the order of months are not supported. + For a streaming query, you may use the function `current_timestamp` to generate windows on + processing time. + gapDuration is provided as strings, e.g. '1 second', '1 day 12 hours', '2 minutes'. Valid + interval strings are 'week', 'day', 'hour', 'minute', 'second', 'millisecond', 'microsecond'. + It could also be a Column which can be evaluated to gap duration dynamically based on the + input row. + The output column will be a struct called 'session_window' by default with the nested columns + 'start' and 'end', where 'start' and 'end' will be of :class:`pyspark.sql.types.TimestampType`. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + timeColumn : :class:`~pyspark.sql.Column` or str + The column name or column to use as the timestamp for windowing by time. + The time column must be of TimestampType or TimestampNTZType. + gapDuration : :class:`~pyspark.sql.Column` or str + A Python string literal or column specifying the timeout of the session. It could be + static value, e.g. `10 minutes`, `1 second`, or an expression/UDF that specifies gap + duration dynamically based on the input row. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([("2016-03-11 09:00:07", 1)]).toDF("date", "val") + >>> w = df.groupBy(session_window("date", "5 seconds")).agg(sum("val").alias("sum")) + >>> w.select(w.session_window.start.cast("string").alias("start"), + ... w.session_window.end.cast("string").alias("end"), "sum").collect() + [Row(start='2016-03-11 09:00:07', end='2016-03-11 09:00:12', sum=1)] + >>> w = df.groupBy(session_window("date", lit("5 seconds"))).agg(sum("val").alias("sum")) + >>> w.select(w.session_window.start.cast("string").alias("start"), + ... w.session_window.end.cast("string").alias("end"), "sum").collect() + [Row(start='2016-03-11 09:00:07', end='2016-03-11 09:00:12', sum=1)] + """ + + def check_field(field: Union[Column, str], fieldName: str) -> None: + if field is None or not isinstance(field, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": fieldName, "arg_type": type(field).__name__}, + ) + + time_col = _to_java_column(timeColumn) + check_field(gapDuration, "gapDuration") + gap_duration = gapDuration if isinstance(gapDuration, str) else _to_java_column(gapDuration) + return _invoke_function("session_window", time_col, gap_duration) + + +@try_remote_functions +def to_unix_timestamp( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + """ + Returns the UNIX timestamp of the given time. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert UNIX timestamp values. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([("2016-04-08",)], ["e"]) + >>> df.select(to_unix_timestamp(df.e, lit("yyyy-MM-dd")).alias('r')).collect() + [Row(r=1460098800)] + >>> spark.conf.unset("spark.sql.session.timeZone") + + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([("2016-04-08",)], ["e"]) + >>> df.select(to_unix_timestamp(df.e).alias('r')).collect() + [Row(r=None)] + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + if format is not None: + return _invoke_function_over_columns("to_unix_timestamp", timestamp, format) + else: + return _invoke_function_over_columns("to_unix_timestamp", timestamp) + + +@try_remote_functions +def to_timestamp_ltz( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + """ + Parses the `timestamp` with the `format` to a timestamp without time zone. + Returns null with invalid input. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert type `TimestampType` timestamp values. + + Examples + -------- + >>> df = spark.createDataFrame([("2016-12-31",)], ["e"]) + >>> df.select(to_timestamp_ltz(df.e, lit("yyyy-MM-dd")).alias('r')).collect() + ... # doctest: +SKIP + [Row(r=datetime.datetime(2016, 12, 31, 0, 0))] + + >>> df = spark.createDataFrame([("2016-12-31",)], ["e"]) + >>> df.select(to_timestamp_ltz(df.e).alias('r')).collect() + ... # doctest: +SKIP + [Row(r=datetime.datetime(2016, 12, 31, 0, 0))] + """ + if format is not None: + return _invoke_function_over_columns("to_timestamp_ltz", timestamp, format) + else: + return _invoke_function_over_columns("to_timestamp_ltz", timestamp) + + +@try_remote_functions +def to_timestamp_ntz( + timestamp: "ColumnOrName", + format: Optional["ColumnOrName"] = None, +) -> Column: + """ + Parses the `timestamp` with the `format` to a timestamp without time zone. + Returns null with invalid input. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + timestamp : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert type `TimestampNTZType` timestamp values. + + Examples + -------- + >>> df = spark.createDataFrame([("2016-04-08",)], ["e"]) + >>> df.select(to_timestamp_ntz(df.e, lit("yyyy-MM-dd")).alias('r')).collect() + ... # doctest: +SKIP + [Row(r=datetime.datetime(2016, 4, 8, 0, 0))] + + >>> df = spark.createDataFrame([("2016-04-08",)], ["e"]) + >>> df.select(to_timestamp_ntz(df.e).alias('r')).collect() + ... # doctest: +SKIP + [Row(r=datetime.datetime(2016, 4, 8, 0, 0))] + """ + if format is not None: + return _invoke_function_over_columns("to_timestamp_ntz", timestamp, format) + else: + return _invoke_function_over_columns("to_timestamp_ntz", timestamp) + + +# ---------------------------- misc functions ---------------------------------- + + +@try_remote_functions +def current_catalog() -> Column: + """Returns the current catalog. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.range(1).select(current_catalog()).show() + +-----------------+ + |current_catalog()| + +-----------------+ + | spark_catalog| + +-----------------+ + """ + return _invoke_function("current_catalog") + + +@try_remote_functions +def current_database() -> Column: + """Returns the current database. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.range(1).select(current_database()).show() + +------------------+ + |current_database()| + +------------------+ + | default| + +------------------+ + """ + return _invoke_function("current_database") + + +@try_remote_functions +def current_schema() -> Column: + """Returns the current database. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.current_schema()).show() + +------------------+ + |current_database()| + +------------------+ + | default| + +------------------+ + """ + return _invoke_function("current_schema") + + +@try_remote_functions +def current_user() -> Column: + """Returns the current database. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> spark.range(1).select(current_user()).show() # doctest: +SKIP + +--------------+ + |current_user()| + +--------------+ + | ruifeng.zheng| + +--------------+ + """ + return _invoke_function("current_user") + + +@try_remote_functions +def user() -> Column: + """Returns the current database. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.user()).show() # doctest: +SKIP + +--------------+ + |current_user()| + +--------------+ + | ruifeng.zheng| + +--------------+ + """ + return _invoke_function("user") + + +@try_remote_functions +def crc32(col: "ColumnOrName") -> Column: + """ + Calculates the cyclic redundancy check value (CRC32) of a binary column and + returns the value as a bigint. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> spark.createDataFrame([('ABC',)], ['a']).select(crc32('a').alias('crc32')).collect() + [Row(crc32=2743272264)] + """ + return _invoke_function_over_columns("crc32", col) + + +@try_remote_functions +def md5(col: "ColumnOrName") -> Column: + """Calculates the MD5 digest and returns the value as a 32 character hex string. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> spark.createDataFrame([('ABC',)], ['a']).select(md5('a').alias('hash')).collect() + [Row(hash='902fbdd2b1df0c4f70b4a5d23525e932')] + """ + return _invoke_function_over_columns("md5", col) + + +@try_remote_functions +def sha1(col: "ColumnOrName") -> Column: + """Returns the hex string result of SHA-1. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> spark.createDataFrame([('ABC',)], ['a']).select(sha1('a').alias('hash')).collect() + [Row(hash='3c01bdbb26f358bab27f267924aa2c9a03fcfdb8')] + """ + return _invoke_function_over_columns("sha1", col) + + +@try_remote_functions +def sha2(col: "ColumnOrName", numBits: int) -> Column: + """Returns the hex string result of SHA-2 family of hash functions (SHA-224, SHA-256, SHA-384, + and SHA-512). The numBits indicates the desired bit length of the result, which must have a + value of 224, 256, 384, 512, or 0 (which is equivalent to 256). + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + numBits : int + the desired bit length of the result, which must have a + value of 224, 256, 384, 512, or 0 (which is equivalent to 256). + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([["Alice"], ["Bob"]], ["name"]) + >>> df.withColumn("sha2", sha2(df.name, 256)).show(truncate=False) + +-----+----------------------------------------------------------------+ + |name |sha2 | + +-----+----------------------------------------------------------------+ + |Alice|3bc51062973c458d5a6f2d8d64a023246354ad7e064b1e4e009ec8a0699a3043| + |Bob |cd9fb1e148ccd8442e5aa74904cc73bf6fb54d1d54d333bd596aa9bb4bb4e961| + +-----+----------------------------------------------------------------+ + """ + return _invoke_function("sha2", _to_java_column(col), numBits) + + +@try_remote_functions +def hash(*cols: "ColumnOrName") -> Column: + """Calculates the hash code of given columns, and returns the result as an int column. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + one or more columns to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + hash value as int column. + + Examples + -------- + >>> df = spark.createDataFrame([('ABC', 'DEF')], ['c1', 'c2']) + + Hash for one column + + >>> df.select(hash('c1').alias('hash')).show() + +----------+ + | hash| + +----------+ + |-757602832| + +----------+ + + Two or more columns + + >>> df.select(hash('c1', 'c2').alias('hash')).show() + +---------+ + | hash| + +---------+ + |599895104| + +---------+ + """ + return _invoke_function_over_seq_of_columns("hash", cols) + + +@try_remote_functions +def xxhash64(*cols: "ColumnOrName") -> Column: + """Calculates the hash code of given columns using the 64-bit variant of the xxHash algorithm, + and returns the result as a long column. The hash computation uses an initial seed of 42. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + one or more columns to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + hash value as long column. + + Examples + -------- + >>> df = spark.createDataFrame([('ABC', 'DEF')], ['c1', 'c2']) + + Hash for one column + + >>> df.select(xxhash64('c1').alias('hash')).show() + +-------------------+ + | hash| + +-------------------+ + |4105715581806190027| + +-------------------+ + + Two or more columns + + >>> df.select(xxhash64('c1', 'c2').alias('hash')).show() + +-------------------+ + | hash| + +-------------------+ + |3233247871021311208| + +-------------------+ + """ + return _invoke_function_over_seq_of_columns("xxhash64", cols) + + +@try_remote_functions +def assert_true(col: "ColumnOrName", errMsg: Optional[Union[Column, str]] = None) -> Column: + """ + Returns `null` if the input column is `true`; throws an exception + with the provided error message otherwise. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column name or column that represents the input column to test + errMsg : :class:`~pyspark.sql.Column` or str, optional + A Python string literal or column containing the error message + + Returns + ------- + :class:`~pyspark.sql.Column` + `null` if the input column is `true` otherwise throws an error with specified message. + + Examples + -------- + >>> df = spark.createDataFrame([(0,1)], ['a', 'b']) + >>> df.select(assert_true(df.a < df.b).alias('r')).collect() + [Row(r=None)] + >>> df.select(assert_true(df.a < df.b, df.a).alias('r')).collect() + [Row(r=None)] + >>> df.select(assert_true(df.a < df.b, 'error').alias('r')).collect() + [Row(r=None)] + >>> df.select(assert_true(df.a > df.b, 'My error msg').alias('r')).collect() # doctest: +SKIP + ... + java.lang.RuntimeException: My error msg + ... + """ + if errMsg is None: + return _invoke_function_over_columns("assert_true", col) + if not isinstance(errMsg, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": type(errMsg).__name__}, + ) + + errMsg = ( + _create_column_from_literal(errMsg) if isinstance(errMsg, str) else _to_java_column(errMsg) + ) + return _invoke_function("assert_true", _to_java_column(col), errMsg) + + +@try_remote_functions +def raise_error(errMsg: Union[Column, str]) -> Column: + """ + Throws an exception with the provided error message. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + errMsg : :class:`~pyspark.sql.Column` or str + A Python string literal or column containing the error message + + Returns + ------- + :class:`~pyspark.sql.Column` + throws an error with specified message. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(raise_error("My error message")).show() # doctest: +SKIP + ... + java.lang.RuntimeException: My error message + ... + """ + if not isinstance(errMsg, (str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": type(errMsg).__name__}, + ) + + errMsg = ( + _create_column_from_literal(errMsg) if isinstance(errMsg, str) else _to_java_column(errMsg) + ) + return _invoke_function("raise_error", errMsg) + + +# ---------------------- String/Binary functions ------------------------------ + + +@try_remote_functions +def upper(col: "ColumnOrName") -> Column: + """ + Converts a string expression to upper case. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + upper case values. + + Examples + -------- + >>> df = spark.createDataFrame(["Spark", "PySpark", "Pandas API"], "STRING") + >>> df.select(upper("value")).show() + +------------+ + |upper(value)| + +------------+ + | SPARK| + | PYSPARK| + | PANDAS API| + +------------+ + """ + return _invoke_function_over_columns("upper", col) + + +@try_remote_functions +def lower(col: "ColumnOrName") -> Column: + """ + Converts a string expression to lower case. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + lower case values. + + Examples + -------- + >>> df = spark.createDataFrame(["Spark", "PySpark", "Pandas API"], "STRING") + >>> df.select(lower("value")).show() + +------------+ + |lower(value)| + +------------+ + | spark| + | pyspark| + | pandas api| + +------------+ + """ + return _invoke_function_over_columns("lower", col) + + +@try_remote_functions +def ascii(col: "ColumnOrName") -> Column: + """ + Computes the numeric value of the first character of the string column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + numeric value. + + Examples + -------- + >>> df = spark.createDataFrame(["Spark", "PySpark", "Pandas API"], "STRING") + >>> df.select(ascii("value")).show() + +------------+ + |ascii(value)| + +------------+ + | 83| + | 80| + | 80| + +------------+ + """ + return _invoke_function_over_columns("ascii", col) + + +@try_remote_functions +def base64(col: "ColumnOrName") -> Column: + """ + Computes the BASE64 encoding of a binary column and returns it as a string column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + BASE64 encoding of string value. + + Examples + -------- + >>> df = spark.createDataFrame(["Spark", "PySpark", "Pandas API"], "STRING") + >>> df.select(base64("value")).show() + +----------------+ + | base64(value)| + +----------------+ + | U3Bhcms=| + | UHlTcGFyaw==| + |UGFuZGFzIEFQSQ==| + +----------------+ + """ + return _invoke_function_over_columns("base64", col) + + +@try_remote_functions +def unbase64(col: "ColumnOrName") -> Column: + """ + Decodes a BASE64 encoded string column and returns it as a binary column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + encoded string value. + + Examples + -------- + >>> df = spark.createDataFrame(["U3Bhcms=", + ... "UHlTcGFyaw==", + ... "UGFuZGFzIEFQSQ=="], "STRING") + >>> df.select(unbase64("value")).show() + +--------------------+ + | unbase64(value)| + +--------------------+ + | [53 70 61 72 6B]| + |[50 79 53 70 61 7...| + |[50 61 6E 64 61 7...| + +--------------------+ + """ + return _invoke_function_over_columns("unbase64", col) + + +@try_remote_functions +def ltrim(col: "ColumnOrName") -> Column: + """ + Trim the spaces from left end for the specified string value. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + left trimmed values. + + Examples + -------- + >>> df = spark.createDataFrame([" Spark", "Spark ", " Spark"], "STRING") + >>> df.select(ltrim("value").alias("r")).withColumn("length", length("r")).show() + +-------+------+ + | r|length| + +-------+------+ + | Spark| 5| + |Spark | 7| + | Spark| 5| + +-------+------+ + """ + return _invoke_function_over_columns("ltrim", col) + + +@try_remote_functions +def rtrim(col: "ColumnOrName") -> Column: + """ + Trim the spaces from right end for the specified string value. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + right trimmed values. + + Examples + -------- + >>> df = spark.createDataFrame([" Spark", "Spark ", " Spark"], "STRING") + >>> df.select(rtrim("value").alias("r")).withColumn("length", length("r")).show() + +--------+------+ + | r|length| + +--------+------+ + | Spark| 8| + | Spark| 5| + | Spark| 6| + +--------+------+ + """ + return _invoke_function_over_columns("rtrim", col) + + +@try_remote_functions +def trim(col: "ColumnOrName") -> Column: + """ + Trim the spaces from both ends for the specified string column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + trimmed values from both sides. + + Examples + -------- + >>> df = spark.createDataFrame([" Spark", "Spark ", " Spark"], "STRING") + >>> df.select(trim("value").alias("r")).withColumn("length", length("r")).show() + +-----+------+ + | r|length| + +-----+------+ + |Spark| 5| + |Spark| 5| + |Spark| 5| + +-----+------+ + """ + return _invoke_function_over_columns("trim", col) + + +@try_remote_functions +def concat_ws(sep: str, *cols: "ColumnOrName") -> Column: + """ + Concatenates multiple input string columns together into a single string column, + using the given separator. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + sep : str + words separator. + cols : :class:`~pyspark.sql.Column` or str + list of columns to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + string of concatenated words. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd','123')], ['s', 'd']) + >>> df.select(concat_ws('-', df.s, df.d).alias('s')).collect() + [Row(s='abcd-123')] + """ + sc = get_active_spark_context() + return _invoke_function("concat_ws", sep, _to_seq(sc, cols, _to_java_column)) + + +@try_remote_functions +def decode(col: "ColumnOrName", charset: str) -> Column: + """ + Computes the first argument into a string from a binary using the provided character set + (one of 'US-ASCII', 'ISO-8859-1', 'UTF-8', 'UTF-16BE', 'UTF-16LE', 'UTF-16'). + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + charset : str + charset to use to decode to. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['a']) + >>> df.select(decode("a", "UTF-8")).show() + +----------------+ + |decode(a, UTF-8)| + +----------------+ + | abcd| + +----------------+ + """ + return _invoke_function("decode", _to_java_column(col), charset) + + +@try_remote_functions +def encode(col: "ColumnOrName", charset: str) -> Column: + """ + Computes the first argument into a binary from a string using the provided character set + (one of 'US-ASCII', 'ISO-8859-1', 'UTF-8', 'UTF-16BE', 'UTF-16LE', 'UTF-16'). + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + charset : str + charset to use to encode. + + Returns + ------- + :class:`~pyspark.sql.Column` + the column for computed results. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['c']) + >>> df.select(encode("c", "UTF-8")).show() + +----------------+ + |encode(c, UTF-8)| + +----------------+ + | [61 62 63 64]| + +----------------+ + """ + return _invoke_function("encode", _to_java_column(col), charset) + + +@try_remote_functions +def format_number(col: "ColumnOrName", d: int) -> Column: + """ + Formats the number X to a format like '#,--#,--#.--', rounded to d decimal places + with HALF_EVEN round mode, and returns the result as a string. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + the column name of the numeric value to be formatted + d : int + the N decimal places + + Returns + ------- + :class:`~pyspark.sql.Column` + the column of formatted results. + + >>> spark.createDataFrame([(5,)], ['a']).select(format_number('a', 4).alias('v')).collect() + [Row(v='5.0000')] + """ + return _invoke_function("format_number", _to_java_column(col), d) + + +@try_remote_functions +def format_string(format: str, *cols: "ColumnOrName") -> Column: + """ + Formats the arguments in printf-style and returns the result as a string column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + format : str + string that can contain embedded format tags and used as result column's value + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s to be used in formatting + + Returns + ------- + :class:`~pyspark.sql.Column` + the column of formatted results. + + Examples + -------- + >>> df = spark.createDataFrame([(5, "hello")], ['a', 'b']) + >>> df.select(format_string('%d %s', df.a, df.b).alias('v')).collect() + [Row(v='5 hello')] + """ + sc = get_active_spark_context() + return _invoke_function("format_string", format, _to_seq(sc, cols, _to_java_column)) + + +@try_remote_functions +def instr(str: "ColumnOrName", substr: str) -> Column: + """ + Locate the position of the first occurrence of substr column in the given string. + Returns null if either of the arguments are null. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The position is not zero based, but 1 based index. Returns 0 if substr + could not be found in str. + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + substr : str + substring to look for. + + Returns + ------- + :class:`~pyspark.sql.Column` + location of the first occurrence of the substring as integer. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['s',]) + >>> df.select(instr(df.s, 'b').alias('s')).collect() + [Row(s=2)] + """ + return _invoke_function("instr", _to_java_column(str), substr) + + +@try_remote_functions +def overlay( + src: "ColumnOrName", + replace: "ColumnOrName", + pos: Union["ColumnOrName", int], + len: Union["ColumnOrName", int] = -1, +) -> Column: + """ + Overlay the specified portion of `src` with `replace`, + starting from byte position `pos` of `src` and proceeding for `len` bytes. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + src : :class:`~pyspark.sql.Column` or str + column name or column containing the string that will be replaced + replace : :class:`~pyspark.sql.Column` or str + column name or column containing the substitution string + pos : :class:`~pyspark.sql.Column` or str or int + column name, column, or int containing the starting position in src + len : :class:`~pyspark.sql.Column` or str or int, optional + column name, column, or int containing the number of bytes to replace in src + string by 'replace' defaults to -1, which represents the length of the 'replace' string + + Returns + ------- + :class:`~pyspark.sql.Column` + string with replaced values. + + Examples + -------- + >>> df = spark.createDataFrame([("SPARK_SQL", "CORE")], ("x", "y")) + >>> df.select(overlay("x", "y", 7).alias("overlayed")).collect() + [Row(overlayed='SPARK_CORE')] + >>> df.select(overlay("x", "y", 7, 0).alias("overlayed")).collect() + [Row(overlayed='SPARK_CORESQL')] + >>> df.select(overlay("x", "y", 7, 2).alias("overlayed")).collect() + [Row(overlayed='SPARK_COREL')] + """ + if not isinstance(pos, (int, str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "pos", "arg_type": type(pos).__name__}, + ) + if len is not None and not isinstance(len, (int, str, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "len", "arg_type": type(len).__name__}, + ) + + pos = _create_column_from_literal(pos) if isinstance(pos, int) else _to_java_column(pos) + len = _create_column_from_literal(len) if isinstance(len, int) else _to_java_column(len) + + return _invoke_function("overlay", _to_java_column(src), _to_java_column(replace), pos, len) + + +@try_remote_functions +def sentences( + string: "ColumnOrName", + language: Optional["ColumnOrName"] = None, + country: Optional["ColumnOrName"] = None, +) -> Column: + """ + Splits a string into arrays of sentences, where each sentence is an array of words. + The 'language' and 'country' arguments are optional, and if omitted, the default locale is used. + + .. versionadded:: 3.2.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + string : :class:`~pyspark.sql.Column` or str + a string to be split + language : :class:`~pyspark.sql.Column` or str, optional + a language of the locale + country : :class:`~pyspark.sql.Column` or str, optional + a country of the locale + + Returns + ------- + :class:`~pyspark.sql.Column` + arrays of split sentences. + + Examples + -------- + >>> df = spark.createDataFrame([["This is an example sentence."]], ["string"]) + >>> df.select(sentences(df.string, lit("en"), lit("US"))).show(truncate=False) + +-----------------------------------+ + |sentences(string, en, US) | + +-----------------------------------+ + |[[This, is, an, example, sentence]]| + +-----------------------------------+ + >>> df = spark.createDataFrame([["Hello world. How are you?"]], ["s"]) + >>> df.select(sentences("s")).show(truncate=False) + +---------------------------------+ + |sentences(s, , ) | + +---------------------------------+ + |[[Hello, world], [How, are, you]]| + +---------------------------------+ + """ + if language is None: + language = lit("") + if country is None: + country = lit("") + + return _invoke_function_over_columns("sentences", string, language, country) + + +@try_remote_functions +def substring(str: "ColumnOrName", pos: int, len: int) -> Column: + """ + Substring starts at `pos` and is of length `len` when str is String type or + returns the slice of byte array that starts at `pos` in byte and is of length `len` + when str is Binary type. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The position is not zero based, but 1 based index. + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + pos : int + starting position in str. + len : int + length of chars. + + Returns + ------- + :class:`~pyspark.sql.Column` + substring of given value. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['s',]) + >>> df.select(substring(df.s, 1, 2).alias('s')).collect() + [Row(s='ab')] + """ + return _invoke_function("substring", _to_java_column(str), pos, len) + + +@try_remote_functions +def substring_index(str: "ColumnOrName", delim: str, count: int) -> Column: + """ + Returns the substring from string str before count occurrences of the delimiter delim. + If count is positive, everything the left of the final delimiter (counting from left) is + returned. If count is negative, every to the right of the final delimiter (counting from the + right) is returned. substring_index performs a case-sensitive match when searching for delim. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + delim : str + delimiter of values. + count : int + number of occurrences. + + Returns + ------- + :class:`~pyspark.sql.Column` + substring of given value. + + Examples + -------- + >>> df = spark.createDataFrame([('a.b.c.d',)], ['s']) + >>> df.select(substring_index(df.s, '.', 2).alias('s')).collect() + [Row(s='a.b')] + >>> df.select(substring_index(df.s, '.', -3).alias('s')).collect() + [Row(s='b.c.d')] + """ + return _invoke_function("substring_index", _to_java_column(str), delim, count) + + +@try_remote_functions +def levenshtein( + left: "ColumnOrName", right: "ColumnOrName", threshold: Optional[int] = None +) -> Column: + """Computes the Levenshtein distance of the two given strings. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + first column value. + right : :class:`~pyspark.sql.Column` or str + second column value. + threshold : int, optional + if set when the levenshtein distance of the two given strings + less than or equal to a given threshold then return result distance, or -1 + + .. versionchanged: 3.5.0 + Added ``threshold`` argument. + + Returns + ------- + :class:`~pyspark.sql.Column` + Levenshtein distance as integer value. + + Examples + -------- + >>> df0 = spark.createDataFrame([('kitten', 'sitting',)], ['l', 'r']) + >>> df0.select(levenshtein('l', 'r').alias('d')).collect() + [Row(d=3)] + >>> df0.select(levenshtein('l', 'r', 2).alias('d')).collect() + [Row(d=-1)] + """ + if threshold is None: + return _invoke_function_over_columns("levenshtein", left, right) + else: + return _invoke_function( + "levenshtein", _to_java_column(left), _to_java_column(right), threshold + ) + + +@try_remote_functions +def locate(substr: str, str: "ColumnOrName", pos: int = 1) -> Column: + """ + Locate the position of the first occurrence of substr in a string column, after position pos. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + substr : str + a string + str : :class:`~pyspark.sql.Column` or str + a Column of :class:`pyspark.sql.types.StringType` + pos : int, optional + start position (zero based) + + Returns + ------- + :class:`~pyspark.sql.Column` + position of the substring. + + Notes + ----- + The position is not zero based, but 1 based index. Returns 0 if substr + could not be found in str. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['s',]) + >>> df.select(locate('b', df.s, 1).alias('s')).collect() + [Row(s=2)] + """ + return _invoke_function("locate", substr, _to_java_column(str), pos) + + +@try_remote_functions +def lpad(col: "ColumnOrName", len: int, pad: str) -> Column: + """ + Left-pad the string column to width `len` with `pad`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + len : int + length of the final string. + pad : str + chars to prepend. + + Returns + ------- + :class:`~pyspark.sql.Column` + left padded result. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['s',]) + >>> df.select(lpad(df.s, 6, '#').alias('s')).collect() + [Row(s='##abcd')] + """ + return _invoke_function("lpad", _to_java_column(col), len, pad) + + +@try_remote_functions +def rpad(col: "ColumnOrName", len: int, pad: str) -> Column: + """ + Right-pad the string column to width `len` with `pad`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + len : int + length of the final string. + pad : str + chars to append. + + Returns + ------- + :class:`~pyspark.sql.Column` + right padded result. + + Examples + -------- + >>> df = spark.createDataFrame([('abcd',)], ['s',]) + >>> df.select(rpad(df.s, 6, '#').alias('s')).collect() + [Row(s='abcd##')] + """ + return _invoke_function("rpad", _to_java_column(col), len, pad) + + +@try_remote_functions +def repeat(col: "ColumnOrName", n: int) -> Column: + """ + Repeats a string column n times, and returns it as a new string column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + n : int + number of times to repeat value. + + Returns + ------- + :class:`~pyspark.sql.Column` + string with repeated values. + + Examples + -------- + >>> df = spark.createDataFrame([('ab',)], ['s',]) + >>> df.select(repeat(df.s, 3).alias('s')).collect() + [Row(s='ababab')] + """ + return _invoke_function("repeat", _to_java_column(col), n) + + +@try_remote_functions +def split(str: "ColumnOrName", pattern: str, limit: int = -1) -> Column: + """ + Splits str around matches of the given pattern. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + a string expression to split + pattern : str + a string representing a regular expression. The regex string should be + a Java regular expression. + limit : int, optional + an integer which controls the number of times `pattern` is applied. + + * ``limit > 0``: The resulting array's length will not be more than `limit`, and the + resulting array's last entry will contain all input beyond the last + matched pattern. + * ``limit <= 0``: `pattern` will be applied as many times as possible, and the resulting + array can be of any size. + + .. versionchanged:: 3.0 + `split` now takes an optional `limit` field. If not provided, default limit value is -1. + + Returns + ------- + :class:`~pyspark.sql.Column` + array of separated strings. + + Examples + -------- + >>> df = spark.createDataFrame([('oneAtwoBthreeC',)], ['s',]) + >>> df.select(split(df.s, '[ABC]', 2).alias('s')).collect() + [Row(s=['one', 'twoBthreeC'])] + >>> df.select(split(df.s, '[ABC]', -1).alias('s')).collect() + [Row(s=['one', 'two', 'three', ''])] + """ + return _invoke_function("split", _to_java_column(str), pattern, limit) + + +@try_remote_functions +def rlike(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + r"""Returns true if `str` matches the Java regex `regexp`, or false otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if `str` matches a Java regex, or false otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([("1a 2b 14m", r"(\d+)")], ["str", "regexp"]) + >>> df.select(rlike('str', lit(r'(\d+)')).alias('d')).collect() + [Row(d=True)] + >>> df.select(rlike('str', lit(r'\d{2}b')).alias('d')).collect() + [Row(d=False)] + >>> df.select(rlike("str", col("regexp")).alias('d')).collect() + [Row(d=True)] + """ + return _invoke_function_over_columns("rlike", str, regexp) + + +@try_remote_functions +def regexp(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + r"""Returns true if `str` matches the Java regex `regexp`, or false otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if `str` matches a Java regex, or false otherwise. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp('str', sf.lit(r'(\d+)'))).show() + +------------------+ + |REGEXP(str, (\d+))| + +------------------+ + | true| + +------------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp('str', sf.lit(r'\d{2}b'))).show() + +-------------------+ + |REGEXP(str, \d{2}b)| + +-------------------+ + | false| + +-------------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp('str', sf.col("regexp"))).show() + +-------------------+ + |REGEXP(str, regexp)| + +-------------------+ + | true| + +-------------------+ + """ + return _invoke_function_over_columns("regexp", str, regexp) + + +@try_remote_functions +def regexp_like(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + r"""Returns true if `str` matches the Java regex `regexp`, or false otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + + Returns + ------- + :class:`~pyspark.sql.Column` + true if `str` matches a Java regex, or false otherwise. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp_like('str', sf.lit(r'(\d+)'))).show() + +-----------------------+ + |REGEXP_LIKE(str, (\d+))| + +-----------------------+ + | true| + +-----------------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp_like('str', sf.lit(r'\d{2}b'))).show() + +------------------------+ + |REGEXP_LIKE(str, \d{2}b)| + +------------------------+ + | false| + +------------------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("1a 2b 14m", r"(\d+)")], ["str", "regexp"] + ... ).select(sf.regexp_like('str', sf.col("regexp"))).show() + +------------------------+ + |REGEXP_LIKE(str, regexp)| + +------------------------+ + | true| + +------------------------+ + """ + return _invoke_function_over_columns("regexp_like", str, regexp) + + +@try_remote_functions +def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + r"""Returns a count of the number of times that the Java regex pattern `regexp` is matched + in the string `str`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + + Returns + ------- + :class:`~pyspark.sql.Column` + the number of times that a Java regex pattern is matched in the string. + + Examples + -------- + >>> df = spark.createDataFrame([("1a 2b 14m", r"\d+")], ["str", "regexp"]) + >>> df.select(regexp_count('str', lit(r'\d+')).alias('d')).collect() + [Row(d=3)] + >>> df.select(regexp_count('str', lit(r'mmm')).alias('d')).collect() + [Row(d=0)] + >>> df.select(regexp_count("str", col("regexp")).alias('d')).collect() + [Row(d=3)] + """ + return _invoke_function_over_columns("regexp_count", str, regexp) + + +@try_remote_functions +def regexp_extract(str: "ColumnOrName", pattern: str, idx: int) -> Column: + r"""Extract a specific group matched by the Java regex `regexp`, from the specified string column. + If the regex did not match, or the specified group did not match, an empty string is returned. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + pattern : str + regex pattern to apply. + idx : int + matched group id. + + Returns + ------- + :class:`~pyspark.sql.Column` + matched value specified by `idx` group id. + + Examples + -------- + >>> df = spark.createDataFrame([('100-200',)], ['str']) + >>> df.select(regexp_extract('str', r'(\d+)-(\d+)', 1).alias('d')).collect() + [Row(d='100')] + >>> df = spark.createDataFrame([('foo',)], ['str']) + >>> df.select(regexp_extract('str', r'(\d+)', 1).alias('d')).collect() + [Row(d='')] + >>> df = spark.createDataFrame([('aaaac',)], ['str']) + >>> df.select(regexp_extract('str', '(a+)(b)?(c)', 2).alias('d')).collect() + [Row(d='')] + """ + return _invoke_function("regexp_extract", _to_java_column(str), pattern, idx) + + +@try_remote_functions +def regexp_extract_all( + str: "ColumnOrName", regexp: "ColumnOrName", idx: Optional[Union[int, Column]] = None +) -> Column: + r"""Extract all strings in the `str` that match the Java regex `regexp` + and corresponding to the regex group index. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + idx : int + matched group id. + + Returns + ------- + :class:`~pyspark.sql.Column` + all strings in the `str` that match a Java regex and corresponding to the regex group index. + + Examples + -------- + >>> df = spark.createDataFrame([("100-200, 300-400", r"(\d+)-(\d+)")], ["str", "regexp"]) + >>> df.select(regexp_extract_all('str', lit(r'(\d+)-(\d+)')).alias('d')).collect() + [Row(d=['100', '300'])] + >>> df.select(regexp_extract_all('str', lit(r'(\d+)-(\d+)'), 1).alias('d')).collect() + [Row(d=['100', '300'])] + >>> df.select(regexp_extract_all('str', lit(r'(\d+)-(\d+)'), 2).alias('d')).collect() + [Row(d=['200', '400'])] + >>> df.select(regexp_extract_all('str', col("regexp")).alias('d')).collect() + [Row(d=['100', '300'])] + """ + if idx is None: + return _invoke_function_over_columns("regexp_extract_all", str, regexp) + else: + idx = lit(idx) if isinstance(idx, int) else idx + return _invoke_function_over_columns("regexp_extract_all", str, regexp, idx) + + +@try_remote_functions +def regexp_replace( + string: "ColumnOrName", pattern: Union[str, Column], replacement: Union[str, Column] +) -> Column: + r"""Replace all substrings of the specified string value that match regexp with replacement. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + string : :class:`~pyspark.sql.Column` or str + column name or column containing the string value + pattern : :class:`~pyspark.sql.Column` or str + column object or str containing the regexp pattern + replacement : :class:`~pyspark.sql.Column` or str + column object or str containing the replacement + + Returns + ------- + :class:`~pyspark.sql.Column` + string with all substrings replaced. + + Examples + -------- + >>> df = spark.createDataFrame([("100-200", r"(\d+)", "--")], ["str", "pattern", "replacement"]) + >>> df.select(regexp_replace('str', r'(\d+)', '--').alias('d')).collect() + [Row(d='-----')] + >>> df.select(regexp_replace("str", col("pattern"), col("replacement")).alias('d')).collect() + [Row(d='-----')] + """ + if isinstance(pattern, str): + pattern_col = _create_column_from_literal(pattern) + else: + pattern_col = _to_java_column(pattern) + if isinstance(replacement, str): + replacement_col = _create_column_from_literal(replacement) + else: + replacement_col = _to_java_column(replacement) + return _invoke_function("regexp_replace", _to_java_column(string), pattern_col, replacement_col) + + +@try_remote_functions +def regexp_substr(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: + r"""Returns the substring that matches the Java regex `regexp` within the string `str`. + If the regular expression is not found, the result is null. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + + Returns + ------- + :class:`~pyspark.sql.Column` + the substring that matches a Java regex within the string `str`. + + Examples + -------- + >>> df = spark.createDataFrame([("1a 2b 14m", r"\d+")], ["str", "regexp"]) + >>> df.select(regexp_substr('str', lit(r'\d+')).alias('d')).collect() + [Row(d='1')] + >>> df.select(regexp_substr('str', lit(r'mmm')).alias('d')).collect() + [Row(d=None)] + >>> df.select(regexp_substr("str", col("regexp")).alias('d')).collect() + [Row(d='1')] + """ + return _invoke_function_over_columns("regexp_substr", str, regexp) + + +@try_remote_functions +def regexp_instr( + str: "ColumnOrName", regexp: "ColumnOrName", idx: Optional[Union[int, Column]] = None +) -> Column: + r"""Extract all strings in the `str` that match the Java regex `regexp` + and corresponding to the regex group index. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + target column to work on. + regexp : :class:`~pyspark.sql.Column` or str + regex pattern to apply. + idx : int + matched group id. + + Returns + ------- + :class:`~pyspark.sql.Column` + all strings in the `str` that match a Java regex and corresponding to the regex group index. + + Examples + -------- + >>> df = spark.createDataFrame([("1a 2b 14m", r"\d+(a|b|m)")], ["str", "regexp"]) + >>> df.select(regexp_instr('str', lit(r'\d+(a|b|m)')).alias('d')).collect() + [Row(d=1)] + >>> df.select(regexp_instr('str', lit(r'\d+(a|b|m)'), 1).alias('d')).collect() + [Row(d=1)] + >>> df.select(regexp_instr('str', lit(r'\d+(a|b|m)'), 2).alias('d')).collect() + [Row(d=1)] + >>> df.select(regexp_instr('str', col("regexp")).alias('d')).collect() + [Row(d=1)] + """ + if idx is None: + return _invoke_function_over_columns("regexp_instr", str, regexp) + else: + idx = lit(idx) if isinstance(idx, int) else idx + return _invoke_function_over_columns("regexp_instr", str, regexp, idx) + + +@try_remote_functions +def initcap(col: "ColumnOrName") -> Column: + """Translate the first letter of each word to upper case in the sentence. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + string with all first letters are uppercase in each word. + + Examples + -------- + >>> spark.createDataFrame([('ab cd',)], ['a']).select(initcap("a").alias('v')).collect() + [Row(v='Ab Cd')] + """ + return _invoke_function_over_columns("initcap", col) + + +@try_remote_functions +def soundex(col: "ColumnOrName") -> Column: + """ + Returns the SoundEx encoding for a string + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + SoundEx encoded string. + + Examples + -------- + >>> df = spark.createDataFrame([("Peters",),("Uhrbach",)], ['name']) + >>> df.select(soundex(df.name).alias("soundex")).collect() + [Row(soundex='P362'), Row(soundex='U612')] + """ + return _invoke_function_over_columns("soundex", col) + + +@try_remote_functions +def bin(col: "ColumnOrName") -> Column: + """Returns the string representation of the binary value of the given column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + binary representation of given value as string. + + Examples + -------- + >>> df = spark.createDataFrame([2,5], "INT") + >>> df.select(bin(df.value).alias('c')).collect() + [Row(c='10'), Row(c='101')] + """ + return _invoke_function_over_columns("bin", col) + + +@try_remote_functions +def hex(col: "ColumnOrName") -> Column: + """Computes hex value of the given column, which could be :class:`pyspark.sql.types.StringType`, + :class:`pyspark.sql.types.BinaryType`, :class:`pyspark.sql.types.IntegerType` or + :class:`pyspark.sql.types.LongType`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + hexadecimal representation of given value as string. + + Examples + -------- + >>> spark.createDataFrame([('ABC', 3)], ['a', 'b']).select(hex('a'), hex('b')).collect() + [Row(hex(a)='414243', hex(b)='3')] + """ + return _invoke_function_over_columns("hex", col) + + +@try_remote_functions +def unhex(col: "ColumnOrName") -> Column: + """Inverse of hex. Interprets each pair of characters as a hexadecimal number + and converts to the byte representation of number. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + string representation of given hexadecimal value. + + Examples + -------- + >>> spark.createDataFrame([('414243',)], ['a']).select(unhex('a')).collect() + [Row(unhex(a)=bytearray(b'ABC'))] + """ + return _invoke_function_over_columns("unhex", col) + + +@try_remote_functions +def length(col: "ColumnOrName") -> Column: + """Computes the character length of string data or number of bytes of binary data. + The length of character data includes the trailing spaces. The length of binary data + includes binary zeros. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + length of the value. + + Examples + -------- + >>> spark.createDataFrame([('ABC ',)], ['a']).select(length('a').alias('length')).collect() + [Row(length=4)] + """ + return _invoke_function_over_columns("length", col) + + +@try_remote_functions +def octet_length(col: "ColumnOrName") -> Column: + """ + Calculates the byte length for the specified string column. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Source column or strings + + Returns + ------- + :class:`~pyspark.sql.Column` + Byte length of the col + + Examples + -------- + >>> from pyspark.sql.functions import octet_length + >>> spark.createDataFrame([('cat',), ( '\U0001F408',)], ['cat']) \\ + ... .select(octet_length('cat')).collect() + [Row(octet_length(cat)=3), Row(octet_length(cat)=4)] + """ + return _invoke_function_over_columns("octet_length", col) + + +@try_remote_functions +def bit_length(col: "ColumnOrName") -> Column: + """ + Calculates the bit length for the specified string column. + + .. versionadded:: 3.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Source column or strings + + Returns + ------- + :class:`~pyspark.sql.Column` + Bit length of the col + + Examples + -------- + >>> from pyspark.sql.functions import bit_length + >>> spark.createDataFrame([('cat',), ( '\U0001F408',)], ['cat']) \\ + ... .select(bit_length('cat')).collect() + [Row(bit_length(cat)=24), Row(bit_length(cat)=32)] + """ + return _invoke_function_over_columns("bit_length", col) + + +@try_remote_functions +def translate(srcCol: "ColumnOrName", matching: str, replace: str) -> Column: + """A function translate any character in the `srcCol` by a character in `matching`. + The characters in `replace` is corresponding to the characters in `matching`. + Translation will happen whenever any character in the string is matching with the character + in the `matching`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + srcCol : :class:`~pyspark.sql.Column` or str + Source column or strings + matching : str + matching characters. + replace : str + characters for replacement. If this is shorter than `matching` string then + those chars that don't have replacement will be dropped. + + Returns + ------- + :class:`~pyspark.sql.Column` + replaced value. + + Examples + -------- + >>> spark.createDataFrame([('translate',)], ['a']).select(translate('a', "rnlt", "123") \\ + ... .alias('r')).collect() + [Row(r='1a2s3ae')] + """ + return _invoke_function("translate", _to_java_column(srcCol), matching, replace) + + +@try_remote_functions +def to_binary(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + """ + Converts the input `col` to a binary value based on the supplied `format`. + The `format` can be a case-insensitive string literal of "hex", "utf-8", "utf8", + or "base64". By default, the binary format for conversion is "hex" if + `format` is omitted. The function returns NULL if at least one of the + input parameters is NULL. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert binary values. + + Examples + -------- + >>> df = spark.createDataFrame([("abc",)], ["e"]) + >>> df.select(to_binary(df.e, lit("utf-8")).alias('r')).collect() + [Row(r=bytearray(b'abc'))] + + >>> df = spark.createDataFrame([("414243",)], ["e"]) + >>> df.select(to_binary(df.e).alias('r')).collect() + [Row(r=bytearray(b'ABC'))] + """ + if format is not None: + return _invoke_function_over_columns("to_binary", col, format) + else: + return _invoke_function_over_columns("to_binary", col) + + +@try_remote_functions +def to_char(col: "ColumnOrName", format: "ColumnOrName") -> Column: + """ + Convert `col` to a string based on the `format`. + Throws an exception if the conversion fails. The format can consist of the following + characters, case insensitive: + '0' or '9': Specifies an expected digit between 0 and 9. A sequence of 0 or 9 in the + format string matches a sequence of digits in the input value, generating a result + string of the same length as the corresponding sequence in the format string. + The result string is left-padded with zeros if the 0/9 sequence comprises more digits + than the matching part of the decimal value, starts with 0, and is before the decimal + point. Otherwise, it is padded with spaces. + '.' or 'D': Specifies the position of the decimal point (optional, only allowed once). + ',' or 'G': Specifies the position of the grouping (thousands) separator (,). + There must be a 0 or 9 to the left and right of each grouping separator. + '$': Specifies the location of the $ currency sign. This character may only be specified once. + 'S' or 'MI': Specifies the position of a '-' or '+' sign (optional, only allowed once at + the beginning or end of the format string). Note that 'S' prints '+' for positive + values but 'MI' prints a space. + 'PR': Only allowed at the end of the format string; specifies that the result string + will be wrapped by angle brackets if the input value is negative. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert char values. + + Examples + -------- + >>> df = spark.createDataFrame([(78.12,)], ["e"]) + >>> df.select(to_char(df.e, lit("$99.99")).alias('r')).collect() + [Row(r='$78.12')] + """ + return _invoke_function_over_columns("to_char", col, format) + + +@try_remote_functions +def to_varchar(col: "ColumnOrName", format: "ColumnOrName") -> Column: + """ + Convert `col` to a string based on the `format`. + Throws an exception if the conversion fails. The format can consist of the following + characters, case insensitive: + '0' or '9': Specifies an expected digit between 0 and 9. A sequence of 0 or 9 in the + format string matches a sequence of digits in the input value, generating a result + string of the same length as the corresponding sequence in the format string. + The result string is left-padded with zeros if the 0/9 sequence comprises more digits + than the matching part of the decimal value, starts with 0, and is before the decimal + point. Otherwise, it is padded with spaces. + '.' or 'D': Specifies the position of the decimal point (optional, only allowed once). + ',' or 'G': Specifies the position of the grouping (thousands) separator (,). + There must be a 0 or 9 to the left and right of each grouping separator. + '$': Specifies the location of the $ currency sign. This character may only be specified once. + 'S' or 'MI': Specifies the position of a '-' or '+' sign (optional, only allowed once at + the beginning or end of the format string). Note that 'S' prints '+' for positive + values but 'MI' prints a space. + 'PR': Only allowed at the end of the format string; specifies that the result string + will be wrapped by angle brackets if the input value is negative. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert char values. + + Examples + -------- + >>> df = spark.createDataFrame([(78.12,)], ["e"]) + >>> df.select(to_varchar(df.e, lit("$99.99")).alias('r')).collect() + [Row(r='$78.12')] + """ + return _invoke_function_over_columns("to_varchar", col, format) + + +@try_remote_functions +def to_number(col: "ColumnOrName", format: "ColumnOrName") -> Column: + """ + Convert string 'col' to a number based on the string format 'format'. + Throws an exception if the conversion fails. The format can consist of the following + characters, case insensitive: + '0' or '9': Specifies an expected digit between 0 and 9. A sequence of 0 or 9 in the + format string matches a sequence of digits in the input string. If the 0/9 + sequence starts with 0 and is before the decimal point, it can only match a digit + sequence of the same size. Otherwise, if the sequence starts with 9 or is after + the decimal point, it can match a digit sequence that has the same or smaller size. + '.' or 'D': Specifies the position of the decimal point (optional, only allowed once). + ',' or 'G': Specifies the position of the grouping (thousands) separator (,). + There must be a 0 or 9 to the left and right of each grouping separator. + 'col' must match the grouping separator relevant for the size of the number. + '$': Specifies the location of the $ currency sign. This character may only be + specified once. + 'S' or 'MI': Specifies the position of a '-' or '+' sign (optional, only allowed + once at the beginning or end of the format string). Note that 'S' allows '-' + but 'MI' does not. + 'PR': Only allowed at the end of the format string; specifies that 'col' indicates a + negative number with wrapping angled brackets. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert number values. + + Examples + -------- + >>> df = spark.createDataFrame([("$78.12",)], ["e"]) + >>> df.select(to_number(df.e, lit("$99.99")).alias('r')).collect() + [Row(r=Decimal('78.12'))] + """ + return _invoke_function_over_columns("to_number", col, format) + + +@try_remote_functions +def replace( + src: "ColumnOrName", search: "ColumnOrName", replace: Optional["ColumnOrName"] = None +) -> Column: + """ + Replaces all occurrences of `search` with `replace`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + src : :class:`~pyspark.sql.Column` or str + A column of string to be replaced. + search : :class:`~pyspark.sql.Column` or str + A column of string, If `search` is not found in `str`, `str` is returned unchanged. + replace : :class:`~pyspark.sql.Column` or str, optional + A column of string, If `replace` is not specified or is an empty string, + nothing replaces the string that is removed from `str`. + + Examples + -------- + >>> df = spark.createDataFrame([("ABCabc", "abc", "DEF",)], ["a", "b", "c"]) + >>> df.select(replace(df.a, df.b, df.c).alias('r')).collect() + [Row(r='ABCDEF')] + + >>> df.select(replace(df.a, df.b).alias('r')).collect() + [Row(r='ABC')] + """ + if replace is not None: + return _invoke_function_over_columns("replace", src, search, replace) + else: + return _invoke_function_over_columns("replace", src, search) + + +@try_remote_functions +def split_part(src: "ColumnOrName", delimiter: "ColumnOrName", partNum: "ColumnOrName") -> Column: + """ + Splits `str` by delimiter and return requested part of the split (1-based). + If any input is null, returns null. if `partNum` is out of range of split parts, + returns empty string. If `partNum` is 0, throws an error. If `partNum` is negative, + the parts are counted backward from the end of the string. + If the `delimiter` is an empty string, the `str` is not split. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + src : :class:`~pyspark.sql.Column` or str + A column of string to be splited. + delimiter : :class:`~pyspark.sql.Column` or str + A column of string, the delimiter used for split. + partNum : :class:`~pyspark.sql.Column` or str + A column of string, requested part of the split (1-based). + + Examples + -------- + >>> df = spark.createDataFrame([("11.12.13", ".", 3,)], ["a", "b", "c"]) + >>> df.select(split_part(df.a, df.b, df.c).alias('r')).collect() + [Row(r='13')] + """ + return _invoke_function_over_columns("split_part", src, delimiter, partNum) + + +@try_remote_functions +def substr( + str: "ColumnOrName", pos: "ColumnOrName", len: Optional["ColumnOrName"] = None +) -> Column: + """ + Returns the substring of `str` that starts at `pos` and is of length `len`, + or the slice of byte array that starts at `pos` and is of length `len`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + src : :class:`~pyspark.sql.Column` or str + A column of string. + pos : :class:`~pyspark.sql.Column` or str + A column of string, the substring of `str` that starts at `pos`. + len : :class:`~pyspark.sql.Column` or str, optional + A column of string, the substring of `str` is of length `len`. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("Spark SQL", 5, 1,)], ["a", "b", "c"] + ... ).select(sf.substr("a", "b", "c")).show() + +---------------+ + |substr(a, b, c)| + +---------------+ + | k| + +---------------+ + + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("Spark SQL", 5, 1,)], ["a", "b", "c"] + ... ).select(sf.substr("a", "b")).show() + +------------------------+ + |substr(a, b, 2147483647)| + +------------------------+ + | k SQL| + +------------------------+ + """ + if len is not None: + return _invoke_function_over_columns("substr", str, pos, len) + else: + return _invoke_function_over_columns("substr", str, pos) + + +@try_remote_functions +def parse_url( + url: "ColumnOrName", partToExtract: "ColumnOrName", key: Optional["ColumnOrName"] = None +) -> Column: + """ + Extracts a part from a URL. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + url : :class:`~pyspark.sql.Column` or str + A column of string. + partToExtract : :class:`~pyspark.sql.Column` or str + A column of string, the path. + key : :class:`~pyspark.sql.Column` or str, optional + A column of string, the key. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [("http://spark.apache.org/path?query=1", "QUERY", "query",)], + ... ["a", "b", "c"] + ... ) + >>> df.select(parse_url(df.a, df.b, df.c).alias('r')).collect() + [Row(r='1')] + + >>> df.select(parse_url(df.a, df.b).alias('r')).collect() + [Row(r='query=1')] + """ + if key is not None: + return _invoke_function_over_columns("parse_url", url, partToExtract, key) + else: + return _invoke_function_over_columns("parse_url", url, partToExtract) + + +@try_remote_functions +def printf(format: "ColumnOrName", *cols: "ColumnOrName") -> Column: + """ + Formats the arguments in printf-style and returns the result as a string column. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + format : :class:`~pyspark.sql.Column` or str + string that can contain embedded format tags and used as result column's value + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s to be used in formatting + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("aa%d%s", 123, "cc",)], ["a", "b", "c"] + ... ).select(sf.printf("a", "b", "c")).show() + +---------------+ + |printf(a, b, c)| + +---------------+ + | aa123cc| + +---------------+ + """ + sc = get_active_spark_context() + return _invoke_function("printf", _to_java_column(format), _to_seq(sc, cols, _to_java_column)) + + +@try_remote_functions +def url_decode(str: "ColumnOrName") -> Column: + """ + Decodes a `str` in 'application/x-www-form-urlencoded' format + using a specific encoding scheme. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A column of string to decode. + + Examples + -------- + >>> df = spark.createDataFrame([("https%3A%2F%2Fspark.apache.org",)], ["a"]) + >>> df.select(url_decode(df.a).alias('r')).collect() + [Row(r='https://spark.apache.org')] + """ + return _invoke_function_over_columns("url_decode", str) + + +@try_remote_functions +def url_encode(str: "ColumnOrName") -> Column: + """ + Translates a string into 'application/x-www-form-urlencoded' format + using a specific encoding scheme. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A column of string to encode. + + Examples + -------- + >>> df = spark.createDataFrame([("https://spark.apache.org",)], ["a"]) + >>> df.select(url_encode(df.a).alias('r')).collect() + [Row(r='https%3A%2F%2Fspark.apache.org')] + """ + return _invoke_function_over_columns("url_encode", str) + + +@try_remote_functions +def position( + substr: "ColumnOrName", str: "ColumnOrName", start: Optional["ColumnOrName"] = None +) -> Column: + """ + Returns the position of the first occurrence of `substr` in `str` after position `start`. + The given `start` and return value are 1-based. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + substr : :class:`~pyspark.sql.Column` or str + A column of string, substring. + str : :class:`~pyspark.sql.Column` or str + A column of string. + start : :class:`~pyspark.sql.Column` or str, optional + A column of string, start position. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [("bar", "foobarbar", 5,)], ["a", "b", "c"] + ... ).select(sf.position("a", "b", "c")).show() + +-----------------+ + |position(a, b, c)| + +-----------------+ + | 7| + +-----------------+ + + >>> spark.createDataFrame( + ... [("bar", "foobarbar", 5,)], ["a", "b", "c"] + ... ).select(sf.position("a", "b")).show() + +-----------------+ + |position(a, b, 1)| + +-----------------+ + | 4| + +-----------------+ + """ + if start is not None: + return _invoke_function_over_columns("position", substr, str, start) + else: + return _invoke_function_over_columns("position", substr, str) + + +@try_remote_functions +def endswith(str: "ColumnOrName", suffix: "ColumnOrName") -> Column: + """ + Returns a boolean. The value is True if str ends with suffix. + Returns NULL if either input expression is NULL. Otherwise, returns False. + Both str or suffix must be of STRING or BINARY type. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A column of string. + suffix : :class:`~pyspark.sql.Column` or str + A column of string, the suffix. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark SQL", "Spark",)], ["a", "b"]) + >>> df.select(endswith(df.a, df.b).alias('r')).collect() + [Row(r=False)] + + >>> df = spark.createDataFrame([("414243", "4243",)], ["e", "f"]) + >>> df = df.select(to_binary("e").alias("e"), to_binary("f").alias("f")) + >>> df.printSchema() + root + |-- e: binary (nullable = true) + |-- f: binary (nullable = true) + >>> df.select(endswith("e", "f"), endswith("f", "e")).show() + +--------------+--------------+ + |endswith(e, f)|endswith(f, e)| + +--------------+--------------+ + | true| false| + +--------------+--------------+ + """ + return _invoke_function_over_columns("endswith", str, suffix) + + +@try_remote_functions +def startswith(str: "ColumnOrName", prefix: "ColumnOrName") -> Column: + """ + Returns a boolean. The value is True if str starts with prefix. + Returns NULL if either input expression is NULL. Otherwise, returns False. + Both str or prefix must be of STRING or BINARY type. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A column of string. + prefix : :class:`~pyspark.sql.Column` or str + A column of string, the prefix. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark SQL", "Spark",)], ["a", "b"]) + >>> df.select(startswith(df.a, df.b).alias('r')).collect() + [Row(r=True)] + + >>> df = spark.createDataFrame([("414243", "4142",)], ["e", "f"]) + >>> df = df.select(to_binary("e").alias("e"), to_binary("f").alias("f")) + >>> df.printSchema() + root + |-- e: binary (nullable = true) + |-- f: binary (nullable = true) + >>> df.select(startswith("e", "f"), startswith("f", "e")).show() + +----------------+----------------+ + |startswith(e, f)|startswith(f, e)| + +----------------+----------------+ + | true| false| + +----------------+----------------+ + """ + return _invoke_function_over_columns("startswith", str, prefix) + + +@try_remote_functions +def char(col: "ColumnOrName") -> Column: + """ + Returns the ASCII character having the binary equivalent to `col`. If col is larger than 256 the + result is equivalent to char(col % 256) + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.char(sf.lit(65))).show() + +--------+ + |char(65)| + +--------+ + | A| + +--------+ + """ + return _invoke_function_over_columns("char", col) + + +@try_remote_functions +def btrim(str: "ColumnOrName", trim: Optional["ColumnOrName"] = None) -> Column: + """ + Remove the leading and trailing `trim` characters from `str`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + trim : :class:`~pyspark.sql.Column` or str + The trim string characters to trim, the default value is a single space + + Examples + -------- + >>> df = spark.createDataFrame([("SSparkSQLS", "SL", )], ['a', 'b']) + >>> df.select(btrim(df.a, df.b).alias('r')).collect() + [Row(r='parkSQ')] + + >>> df = spark.createDataFrame([(" SparkSQL ",)], ['a']) + >>> df.select(btrim(df.a).alias('r')).collect() + [Row(r='SparkSQL')] + """ + if trim is not None: + return _invoke_function_over_columns("btrim", str, trim) + else: + return _invoke_function_over_columns("btrim", str) + + +@try_remote_functions +def char_length(str: "ColumnOrName") -> Column: + """ + Returns the character length of string data or number of bytes of binary data. + The length of string data includes the trailing spaces. + The length of binary data includes binary zeros. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.char_length(sf.lit("SparkSQL"))).show() + +---------------------+ + |char_length(SparkSQL)| + +---------------------+ + | 8| + +---------------------+ + """ + return _invoke_function_over_columns("char_length", str) + + +@try_remote_functions +def character_length(str: "ColumnOrName") -> Column: + """ + Returns the character length of string data or number of bytes of binary data. + The length of string data includes the trailing spaces. + The length of binary data includes binary zeros. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.character_length(sf.lit("SparkSQL"))).show() + +--------------------------+ + |character_length(SparkSQL)| + +--------------------------+ + | 8| + +--------------------------+ + """ + return _invoke_function_over_columns("character_length", str) + + +@try_remote_functions +def try_to_binary(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: + """ + This is a special version of `to_binary` that performs the same operation, but returns a NULL + value instead of raising an error if the conversion cannot be performed. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert binary values. + + Examples + -------- + >>> df = spark.createDataFrame([("abc",)], ["e"]) + >>> df.select(try_to_binary(df.e, lit("utf-8")).alias('r')).collect() + [Row(r=bytearray(b'abc'))] + + >>> df = spark.createDataFrame([("414243",)], ["e"]) + >>> df.select(try_to_binary(df.e).alias('r')).collect() + [Row(r=bytearray(b'ABC'))] + """ + if format is not None: + return _invoke_function_over_columns("try_to_binary", col, format) + else: + return _invoke_function_over_columns("try_to_binary", col) + + +@try_remote_functions +def try_to_number(col: "ColumnOrName", format: "ColumnOrName") -> Column: + """ + Convert string 'col' to a number based on the string format `format`. Returns NULL if the + string 'col' does not match the expected format. The format follows the same semantics as the + to_number function. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + Input column or strings. + format : :class:`~pyspark.sql.Column` or str, optional + format to use to convert number values. + + Examples + -------- + >>> df = spark.createDataFrame([("$78.12",)], ["e"]) + >>> df.select(try_to_number(df.e, lit("$99.99")).alias('r')).collect() + [Row(r=Decimal('78.12'))] + """ + return _invoke_function_over_columns("try_to_number", col, format) + + +@try_remote_functions +def contains(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns a boolean. The value is True if right is found inside left. + Returns NULL if either input expression is NULL. Otherwise, returns False. + Both left or right must be of STRING or BINARY type. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + The input column or strings to check, may be NULL. + right : :class:`~pyspark.sql.Column` or str + The input column or strings to find, may be NULL. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark SQL", "Spark")], ['a', 'b']) + >>> df.select(contains(df.a, df.b).alias('r')).collect() + [Row(r=True)] + + >>> df = spark.createDataFrame([("414243", "4243",)], ["c", "d"]) + >>> df = df.select(to_binary("c").alias("c"), to_binary("d").alias("d")) + >>> df.printSchema() + root + |-- c: binary (nullable = true) + |-- d: binary (nullable = true) + >>> df.select(contains("c", "d"), contains("d", "c")).show() + +--------------+--------------+ + |contains(c, d)|contains(d, c)| + +--------------+--------------+ + | true| false| + +--------------+--------------+ + """ + return _invoke_function_over_columns("contains", left, right) + + +@try_remote_functions +def elt(*inputs: "ColumnOrName") -> Column: + """ + Returns the `n`-th input, e.g., returns `input2` when `n` is 2. + The function returns NULL if the index exceeds the length of the array + and `spark.sql.ansi.enabled` is set to false. If `spark.sql.ansi.enabled` is set to true, + it throws ArrayIndexOutOfBoundsException for invalid indices. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + inputs : :class:`~pyspark.sql.Column` or str + Input columns or strings. + + Examples + -------- + >>> df = spark.createDataFrame([(1, "scala", "java")], ['a', 'b', 'c']) + >>> df.select(elt(df.a, df.b, df.c).alias('r')).collect() + [Row(r='scala')] + """ + sc = get_active_spark_context() + return _invoke_function("elt", _to_seq(sc, inputs, _to_java_column)) + + +@try_remote_functions +def find_in_set(str: "ColumnOrName", str_array: "ColumnOrName") -> Column: + """ + Returns the index (1-based) of the given string (`str`) in the comma-delimited + list (`strArray`). Returns 0, if the string was not found or if the given string (`str`) + contains a comma. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + The given string to be found. + str_array : :class:`~pyspark.sql.Column` or str + The comma-delimited list. + + Examples + -------- + >>> df = spark.createDataFrame([("ab", "abc,b,ab,c,def")], ['a', 'b']) + >>> df.select(find_in_set(df.a, df.b).alias('r')).collect() + [Row(r=3)] + """ + return _invoke_function_over_columns("find_in_set", str, str_array) + + +@try_remote_functions +def like( + str: "ColumnOrName", pattern: "ColumnOrName", escapeChar: Optional["Column"] = None +) -> Column: + """ + Returns true if str matches `pattern` with `escape`, + null if any arguments are null, false otherwise. + The default escape character is the '\'. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A string. + pattern : :class:`~pyspark.sql.Column` or str + A string. The pattern is a string which is matched literally, with + exception to the following special symbols: + _ matches any one character in the input (similar to . in posix regular expressions) + % matches zero or more characters in the input (similar to .* in posix regular + expressions) + Since Spark 2.0, string literals are unescaped in our SQL parser. For example, in order + to match "\abc", the pattern should be "\\abc". + When SQL config 'spark.sql.parser.escapedStringLiterals' is enabled, it falls back + to Spark 1.6 behavior regarding string literal parsing. For example, if the config is + enabled, the pattern to match "\abc" should be "\abc". + escape : :class:`~pyspark.sql.Column` + An character added since Spark 3.0. The default escape character is the '\'. + If an escape character precedes a special symbol or another escape character, the + following character is matched literally. It is invalid to escape any other character. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark", "_park")], ['a', 'b']) + >>> df.select(like(df.a, df.b).alias('r')).collect() + [Row(r=True)] + + >>> df = spark.createDataFrame( + ... [("%SystemDrive%/Users/John", "/%SystemDrive/%//Users%")], + ... ['a', 'b'] + ... ) + >>> df.select(like(df.a, df.b, lit('/')).alias('r')).collect() + [Row(r=True)] + """ + if escapeChar is not None: + return _invoke_function_over_columns("like", str, pattern, escapeChar) + else: + return _invoke_function_over_columns("like", str, pattern) + + +@try_remote_functions +def ilike( + str: "ColumnOrName", pattern: "ColumnOrName", escapeChar: Optional["Column"] = None +) -> Column: + """ + Returns true if str matches `pattern` with `escape` case-insensitively, + null if any arguments are null, false otherwise. + The default escape character is the '\'. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + A string. + pattern : :class:`~pyspark.sql.Column` or str + A string. The pattern is a string which is matched literally, with + exception to the following special symbols: + _ matches any one character in the input (similar to . in posix regular expressions) + % matches zero or more characters in the input (similar to .* in posix regular + expressions) + Since Spark 2.0, string literals are unescaped in our SQL parser. For example, in order + to match "\abc", the pattern should be "\\abc". + When SQL config 'spark.sql.parser.escapedStringLiterals' is enabled, it falls back + to Spark 1.6 behavior regarding string literal parsing. For example, if the config is + enabled, the pattern to match "\abc" should be "\abc". + escape : :class:`~pyspark.sql.Column` + An character added since Spark 3.0. The default escape character is the '\'. + If an escape character precedes a special symbol or another escape character, the + following character is matched literally. It is invalid to escape any other character. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark", "_park")], ['a', 'b']) + >>> df.select(ilike(df.a, df.b).alias('r')).collect() + [Row(r=True)] + + >>> df = spark.createDataFrame( + ... [("%SystemDrive%/Users/John", "/%SystemDrive/%//Users%")], + ... ['a', 'b'] + ... ) + >>> df.select(ilike(df.a, df.b, lit('/')).alias('r')).collect() + [Row(r=True)] + """ + if escapeChar is not None: + return _invoke_function_over_columns("ilike", str, pattern, escapeChar) + else: + return _invoke_function_over_columns("ilike", str, pattern) + + +@try_remote_functions +def lcase(str: "ColumnOrName") -> Column: + """ + Returns `str` with all characters changed to lowercase. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.lcase(sf.lit("Spark"))).show() + +------------+ + |lcase(Spark)| + +------------+ + | spark| + +------------+ + """ + return _invoke_function_over_columns("lcase", str) + + +@try_remote_functions +def ucase(str: "ColumnOrName") -> Column: + """ + Returns `str` with all characters changed to uppercase. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.ucase(sf.lit("Spark"))).show() + +------------+ + |ucase(Spark)| + +------------+ + | SPARK| + +------------+ + """ + return _invoke_function_over_columns("ucase", str) + + +@try_remote_functions +def left(str: "ColumnOrName", len: "ColumnOrName") -> Column: + """ + Returns the leftmost `len`(`len` can be string type) characters from the string `str`, + if `len` is less or equal than 0 the result is an empty string. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + len : :class:`~pyspark.sql.Column` or str + Input column or strings, the leftmost `len`. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark SQL", 3,)], ['a', 'b']) + >>> df.select(left(df.a, df.b).alias('r')).collect() + [Row(r='Spa')] + """ + return _invoke_function_over_columns("left", str, len) + + +@try_remote_functions +def right(str: "ColumnOrName", len: "ColumnOrName") -> Column: + """ + Returns the rightmost `len`(`len` can be string type) characters from the string `str`, + if `len` is less or equal than 0 the result is an empty string. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or str + Input column or strings. + len : :class:`~pyspark.sql.Column` or str + Input column or strings, the rightmost `len`. + + Examples + -------- + >>> df = spark.createDataFrame([("Spark SQL", 3,)], ['a', 'b']) + >>> df.select(right(df.a, df.b).alias('r')).collect() + [Row(r='SQL')] + """ + return _invoke_function_over_columns("right", str, len) + + +@try_remote_functions +def mask( + col: "ColumnOrName", + upperChar: Optional["ColumnOrName"] = None, + lowerChar: Optional["ColumnOrName"] = None, + digitChar: Optional["ColumnOrName"] = None, + otherChar: Optional["ColumnOrName"] = None, +) -> Column: + """ + Masks the given string value. This can be useful for creating copies of tables with sensitive + information removed. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col: :class:`~pyspark.sql.Column` or str + target column to compute on. + upperChar: :class:`~pyspark.sql.Column` or str + character to replace upper-case characters with. Specify NULL to retain original character. + lowerChar: :class:`~pyspark.sql.Column` or str + character to replace lower-case characters with. Specify NULL to retain original character. + digitChar: :class:`~pyspark.sql.Column` or str + character to replace digit characters with. Specify NULL to retain original character. + otherChar: :class:`~pyspark.sql.Column` or str + character to replace all other characters with. Specify NULL to retain original character. + + Returns + ------- + :class:`~pyspark.sql.Column` + + Examples + -------- + >>> df = spark.createDataFrame([("AbCD123-@$#",), ("abcd-EFGH-8765-4321",)], ['data']) + >>> df.select(mask(df.data).alias('r')).collect() + [Row(r='XxXXnnn-@$#'), Row(r='xxxx-XXXX-nnnn-nnnn')] + >>> df.select(mask(df.data, lit('Y')).alias('r')).collect() + [Row(r='YxYYnnn-@$#'), Row(r='xxxx-YYYY-nnnn-nnnn')] + >>> df.select(mask(df.data, lit('Y'), lit('y')).alias('r')).collect() + [Row(r='YyYYnnn-@$#'), Row(r='yyyy-YYYY-nnnn-nnnn')] + >>> df.select(mask(df.data, lit('Y'), lit('y'), lit('d')).alias('r')).collect() + [Row(r='YyYYddd-@$#'), Row(r='yyyy-YYYY-dddd-dddd')] + >>> df.select(mask(df.data, lit('Y'), lit('y'), lit('d'), lit('*')).alias('r')).collect() + [Row(r='YyYYddd****'), Row(r='yyyy*YYYY*dddd*dddd')] + """ + + _upperChar = lit("X") if upperChar is None else upperChar + _lowerChar = lit("x") if lowerChar is None else lowerChar + _digitChar = lit("n") if digitChar is None else digitChar + _otherChar = lit(None) if otherChar is None else otherChar + return _invoke_function_over_columns( + "mask", col, _upperChar, _lowerChar, _digitChar, _otherChar + ) + + +# ---------------------- Collection functions ------------------------------ + + +@overload +def create_map(*cols: "ColumnOrName") -> Column: + ... + + +@overload +def create_map(__cols: Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]) -> Column: + ... + + +@try_remote_functions +def create_map( + *cols: Union["ColumnOrName", Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]] +) -> Column: + """Creates a new map column. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s that are + grouped as key-value pairs, e.g. (key1, value1, key2, value2, ...). + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.select(create_map('name', 'age').alias("map")).collect() + [Row(map={'Alice': 2}), Row(map={'Bob': 5})] + >>> df.select(create_map([df.name, df.age]).alias("map")).collect() + [Row(map={'Alice': 2}), Row(map={'Bob': 5})] + """ + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_seq_of_columns("map", cols) # type: ignore[arg-type] + + +@try_remote_functions +def map_from_arrays(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """Creates a new map from two arrays. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + name of column containing a set of keys. All elements should not be null + col2 : :class:`~pyspark.sql.Column` or str + name of column containing a set of values + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of map type. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 5], ['a', 'b'])], ['k', 'v']) + >>> df = df.select(map_from_arrays(df.k, df.v).alias("col")) + >>> df.show() + +----------------+ + | col| + +----------------+ + |{2 -> a, 5 -> b}| + +----------------+ + >>> df.printSchema() + root + |-- col: map (nullable = true) + | |-- key: long + | |-- value: string (valueContainsNull = true) + """ + return _invoke_function_over_columns("map_from_arrays", col1, col2) + + +@overload +def array(*cols: "ColumnOrName") -> Column: + ... + + +@overload +def array(__cols: Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]) -> Column: + ... + + +@try_remote_functions +def array( + *cols: Union["ColumnOrName", Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]] +) -> Column: + """Creates a new array column. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s that have + the same data type. + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of array type. + + Examples + -------- + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.select(array('age', 'age').alias("arr")).collect() + [Row(arr=[2, 2]), Row(arr=[5, 5])] + >>> df.select(array([df.age, df.age]).alias("arr")).collect() + [Row(arr=[2, 2]), Row(arr=[5, 5])] + >>> df.select(array('age', 'age').alias("col")).printSchema() + root + |-- col: array (nullable = false) + | |-- element: long (containsNull = true) + """ + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_seq_of_columns("array", cols) # type: ignore[arg-type] + + +@try_remote_functions +def array_contains(col: "ColumnOrName", value: Any) -> Column: + """ + Collection function: returns null if the array is null, true if the array contains the + given value, and false otherwise. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array + value : + value or column to check for in array + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of Boolean type. + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],), ([],)], ['data']) + >>> df.select(array_contains(df.data, "a")).collect() + [Row(array_contains(data, a)=True), Row(array_contains(data, a)=False)] + >>> df.select(array_contains(df.data, lit("a"))).collect() + [Row(array_contains(data, a)=True), Row(array_contains(data, a)=False)] + """ + value = value._jc if isinstance(value, Column) else value + return _invoke_function("array_contains", _to_java_column(col), value) + + +@try_remote_functions +def arrays_overlap(a1: "ColumnOrName", a2: "ColumnOrName") -> Column: + """ + Collection function: returns true if the arrays contain any common non-null element; if not, + returns null if both the arrays are non-empty and any of them contains a null element; returns + false otherwise. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of Boolean type. + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b"], ["b", "c"]), (["a"], ["b", "c"])], ['x', 'y']) + >>> df.select(arrays_overlap(df.x, df.y).alias("overlap")).collect() + [Row(overlap=True), Row(overlap=False)] + """ + return _invoke_function_over_columns("arrays_overlap", a1, a2) + + +@try_remote_functions +def slice( + x: "ColumnOrName", start: Union["ColumnOrName", int], length: Union["ColumnOrName", int] +) -> Column: + """ + Collection function: returns an array containing all the elements in `x` from index `start` + (array indices start at 1, or from the end if `start` is negative) with the specified `length`. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + x : :class:`~pyspark.sql.Column` or str + column name or column containing the array to be sliced + start : :class:`~pyspark.sql.Column` or str or int + column name, column, or int containing the starting index + length : :class:`~pyspark.sql.Column` or str or int + column name, column, or int containing the length of the slice + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of array type. Subset of array. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 2, 3],), ([4, 5],)], ['x']) + >>> df.select(slice(df.x, 2, 2).alias("sliced")).collect() + [Row(sliced=[2, 3]), Row(sliced=[5])] + """ + start = lit(start) if isinstance(start, int) else start + length = lit(length) if isinstance(length, int) else length + + return _invoke_function_over_columns("slice", x, start, length) + + +@try_remote_functions +def array_join( + col: "ColumnOrName", delimiter: str, null_replacement: Optional[str] = None +) -> Column: + """ + Concatenates the elements of `column` using the `delimiter`. Null values are replaced with + `null_replacement` if set, otherwise they are ignored. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + delimiter : str + delimiter used to concatenate elements + null_replacement : str, optional + if set then null values will be replaced by this value + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of string type. Concatenated values. + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],), (["a", None],)], ['data']) + >>> df.select(array_join(df.data, ",").alias("joined")).collect() + [Row(joined='a,b,c'), Row(joined='a')] + >>> df.select(array_join(df.data, ",", "NULL").alias("joined")).collect() + [Row(joined='a,b,c'), Row(joined='a,NULL')] + """ + get_active_spark_context() + if null_replacement is None: + return _invoke_function("array_join", _to_java_column(col), delimiter) + else: + return _invoke_function("array_join", _to_java_column(col), delimiter, null_replacement) + + +@try_remote_functions +def concat(*cols: "ColumnOrName") -> Column: + """ + Concatenates multiple input columns together into a single column. + The function works with strings, numeric, binary and compatible array columns. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + target column or columns to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + concatenated values. Type of the `Column` depends on input columns' type. + + See Also + -------- + :meth:`pyspark.sql.functions.array_join` : to concatenate string columns with delimiter + + Examples + -------- + >>> df = spark.createDataFrame([('abcd','123')], ['s', 'd']) + >>> df = df.select(concat(df.s, df.d).alias('s')) + >>> df.collect() + [Row(s='abcd123')] + >>> df + DataFrame[s: string] + + >>> df = spark.createDataFrame([([1, 2], [3, 4], [5]), ([1, 2], None, [3])], ['a', 'b', 'c']) + >>> df = df.select(concat(df.a, df.b, df.c).alias("arr")) + >>> df.collect() + [Row(arr=[1, 2, 3, 4, 5]), Row(arr=None)] + >>> df + DataFrame[arr: array] + """ + return _invoke_function_over_seq_of_columns("concat", cols) + + +@try_remote_functions +def array_position(col: "ColumnOrName", value: Any) -> Column: + """ + Collection function: Locates the position of the first occurrence of the given value + in the given array. Returns null if either of the arguments are null. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The position is not zero based, but 1 based index. Returns 0 if the given + value could not be found in the array. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + value : Any + value to look for. + + Returns + ------- + :class:`~pyspark.sql.Column` + position of the value in the given array if found and 0 otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([(["c", "b", "a"],), ([],)], ['data']) + >>> df.select(array_position(df.data, "a")).collect() + [Row(array_position(data, a)=3), Row(array_position(data, a)=0)] + """ + return _invoke_function("array_position", _to_java_column(col), value) + + +@try_remote_functions +def element_at(col: "ColumnOrName", extraction: Any) -> Column: + """ + Collection function: Returns element of array at given index in `extraction` if col is array. + Returns value for the given key in `extraction` if col is map. If position is negative + then location of the element will start from end, if number is outside the + array boundaries then None will be returned. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array or map + extraction : + index to check for in array or key to check for in map + + Returns + ------- + :class:`~pyspark.sql.Column` + value at given position. + + Notes + ----- + The position is not zero based, but 1 based index. + + See Also + -------- + :meth:`get` + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],)], ['data']) + >>> df.select(element_at(df.data, 1)).collect() + [Row(element_at(data, 1)='a')] + >>> df.select(element_at(df.data, -1)).collect() + [Row(element_at(data, -1)='c')] + + >>> df = spark.createDataFrame([({"a": 1.0, "b": 2.0},)], ['data']) + >>> df.select(element_at(df.data, lit("a"))).collect() + [Row(element_at(data, a)=1.0)] + """ + return _invoke_function_over_columns("element_at", col, lit(extraction)) + + +@try_remote_functions +def try_element_at(col: "ColumnOrName", extraction: "ColumnOrName") -> Column: + """ + (array, index) - Returns element of array at given (1-based) index. If Index is 0, Spark will + throw an error. If index < 0, accesses elements from the last to the first. The function + always returns NULL if the index exceeds the length of the array. + + (map, key) - Returns value for given key. The function always returns NULL if the key is not + contained in the map. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array or map + extraction : + index to check for in array or key to check for in map + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"],)], ['data']) + >>> df.select(try_element_at(df.data, lit(1)).alias('r')).collect() + [Row(r='a')] + >>> df.select(try_element_at(df.data, lit(-1)).alias('r')).collect() + [Row(r='c')] + + >>> df = spark.createDataFrame([({"a": 1.0, "b": 2.0},)], ['data']) + >>> df.select(try_element_at(df.data, lit("a")).alias('r')).collect() + [Row(r=1.0)] + """ + return _invoke_function_over_columns("try_element_at", col, extraction) + + +@try_remote_functions +def get(col: "ColumnOrName", index: Union["ColumnOrName", int]) -> Column: + """ + Collection function: Returns element of array at given (0-based) index. + If the index points outside of the array boundaries, then this function + returns NULL. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array + index : :class:`~pyspark.sql.Column` or str or int + index to check for in array + + Returns + ------- + :class:`~pyspark.sql.Column` + value at given position. + + Notes + ----- + The position is not 1 based, but 0 based index. + Supports Spark Connect. + + See Also + -------- + :meth:`element_at` + + Examples + -------- + >>> df = spark.createDataFrame([(["a", "b", "c"], 1)], ['data', 'index']) + >>> df.select(get(df.data, 1)).show() + +------------+ + |get(data, 1)| + +------------+ + | b| + +------------+ + + >>> df.select(get(df.data, -1)).show() + +-------------+ + |get(data, -1)| + +-------------+ + | NULL| + +-------------+ + + >>> df.select(get(df.data, 3)).show() + +------------+ + |get(data, 3)| + +------------+ + | NULL| + +------------+ + + >>> df.select(get(df.data, "index")).show() + +----------------+ + |get(data, index)| + +----------------+ + | b| + +----------------+ + + >>> df.select(get(df.data, col("index") - 1)).show() + +----------------------+ + |get(data, (index - 1))| + +----------------------+ + | a| + +----------------------+ + """ + index = lit(index) if isinstance(index, int) else index + + return _invoke_function_over_columns("get", col, index) + + +@try_remote_functions +def array_prepend(col: "ColumnOrName", value: Any) -> Column: + """ + Collection function: Returns an array containing element as + well as all elements from array. The new element is positioned + at the beginning of the array. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array + value : + a literal value, or a :class:`~pyspark.sql.Column` expression. + + Returns + ------- + :class:`~pyspark.sql.Column` + an array excluding given value. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 3, 4],), ([],)], ['data']) + >>> df.select(array_prepend(df.data, 1)).collect() + [Row(array_prepend(data, 1)=[1, 2, 3, 4]), Row(array_prepend(data, 1)=[1])] + """ + return _invoke_function_over_columns("array_prepend", col, lit(value)) + + +@try_remote_functions +def array_remove(col: "ColumnOrName", element: Any) -> Column: + """ + Collection function: Remove all elements that equal to element from the given array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array + element : + element to be removed from the array + + Returns + ------- + :class:`~pyspark.sql.Column` + an array excluding given value. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 2, 3, 1, 1],), ([],)], ['data']) + >>> df.select(array_remove(df.data, 1)).collect() + [Row(array_remove(data, 1)=[2, 3]), Row(array_remove(data, 1)=[])] + """ + return _invoke_function("array_remove", _to_java_column(col), element) + + +@try_remote_functions +def array_distinct(col: "ColumnOrName") -> Column: + """ + Collection function: removes duplicate values from the array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of unique values. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 2, 3, 2],), ([4, 5, 5, 4],)], ['data']) + >>> df.select(array_distinct(df.data)).collect() + [Row(array_distinct(data)=[1, 2, 3]), Row(array_distinct(data)=[4, 5])] + """ + return _invoke_function_over_columns("array_distinct", col) + + +@try_remote_functions +def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Any) -> Column: + """ + Collection function: adds an item into a given array at a specified array index. + Array indices start at 1, or start from the end if index is negative. + Index above array size appends the array, or prepends the array if index is negative, + with 'null' elements. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + arr : :class:`~pyspark.sql.Column` or str + name of column containing an array + pos : :class:`~pyspark.sql.Column` or str or int + name of Numeric type column indicating position of insertion + (starting at index 1, negative position is a start from the back of the array) + value : + a literal value, or a :class:`~pyspark.sql.Column` expression. + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of values, including the new specified value + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(['a', 'b', 'c'], 2, 'd'), (['c', 'b', 'a'], -2, 'd')], + ... ['data', 'pos', 'val'] + ... ) + >>> df.select(array_insert(df.data, df.pos.cast('integer'), df.val).alias('data')).collect() + [Row(data=['a', 'd', 'b', 'c']), Row(data=['c', 'b', 'd', 'a'])] + >>> df.select(array_insert(df.data, 5, 'hello').alias('data')).collect() + [Row(data=['a', 'b', 'c', None, 'hello']), Row(data=['c', 'b', 'a', None, 'hello'])] + """ + pos = lit(pos) if isinstance(pos, int) else pos + + return _invoke_function_over_columns("array_insert", arr, pos, lit(value)) + + +@try_remote_functions +def array_intersect(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Collection function: returns an array of the elements in the intersection of col1 and col2, + without duplicates. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + name of column containing array + col2 : :class:`~pyspark.sql.Column` or str + name of column containing array + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of values in the intersection of two arrays. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(c1=["b", "a", "c"], c2=["c", "d", "a", "f"])]) + >>> df.select(array_intersect(df.c1, df.c2)).collect() + [Row(array_intersect(c1, c2)=['a', 'c'])] + """ + return _invoke_function_over_columns("array_intersect", col1, col2) + + +@try_remote_functions +def array_union(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Collection function: returns an array of the elements in the union of col1 and col2, + without duplicates. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + name of column containing array + col2 : :class:`~pyspark.sql.Column` or str + name of column containing array + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of values in union of two arrays. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(c1=["b", "a", "c"], c2=["c", "d", "a", "f"])]) + >>> df.select(array_union(df.c1, df.c2)).collect() + [Row(array_union(c1, c2)=['b', 'a', 'c', 'd', 'f'])] + """ + return _invoke_function_over_columns("array_union", col1, col2) + + +@try_remote_functions +def array_except(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Collection function: returns an array of the elements in col1 but not in col2, + without duplicates. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + name of column containing array + col2 : :class:`~pyspark.sql.Column` or str + name of column containing array + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of values from first array that are not in the second. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(c1=["b", "a", "c"], c2=["c", "d", "a", "f"])]) + >>> df.select(array_except(df.c1, df.c2)).collect() + [Row(array_except(c1, c2)=['b'])] + """ + return _invoke_function_over_columns("array_except", col1, col2) + + +@try_remote_functions +def array_compact(col: "ColumnOrName") -> Column: + """ + Collection function: removes null values from the array. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + an array by excluding the null values. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame([([1, None, 2, 3],), ([4, 5, None, 4],)], ['data']) + >>> df.select(array_compact(df.data)).collect() + [Row(array_compact(data)=[1, 2, 3]), Row(array_compact(data)=[4, 5, 4])] + """ + return _invoke_function_over_columns("array_compact", col) + + +@try_remote_functions +def array_append(col: "ColumnOrName", value: Any) -> Column: + """ + Collection function: returns an array of the elements in col1 along + with the added element in col2 at the last of the array. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing array + value : + a literal value, or a :class:`~pyspark.sql.Column` expression. + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of values from first array along with the element. + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(c1=["b", "a", "c"], c2="c")]) + >>> df.select(array_append(df.c1, df.c2)).collect() + [Row(array_append(c1, c2)=['b', 'a', 'c', 'c'])] + >>> df.select(array_append(df.c1, 'x')).collect() + [Row(array_append(c1, x)=['b', 'a', 'c', 'x'])] + """ + return _invoke_function_over_columns("array_append", col, lit(value)) + + +@try_remote_functions +def explode(col: "ColumnOrName") -> Column: + """ + Returns a new row for each element in the given array or map. + Uses the default column name `col` for elements in the array and + `key` and `value` for elements in the map unless specified otherwise. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + one row per array item or map key value. + + See Also + -------- + :meth:`pyspark.functions.posexplode` + :meth:`pyspark.functions.explode_outer` + :meth:`pyspark.functions.posexplode_outer` + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(a=1, intlist=[1,2,3], mapfield={"a": "b"})]) + >>> df.select(explode(df.intlist).alias("anInt")).collect() + [Row(anInt=1), Row(anInt=2), Row(anInt=3)] + + >>> df.select(explode(df.mapfield).alias("key", "value")).show() + +---+-----+ + |key|value| + +---+-----+ + | a| b| + +---+-----+ + """ + return _invoke_function_over_columns("explode", col) + + +@try_remote_functions +def posexplode(col: "ColumnOrName") -> Column: + """ + Returns a new row for each element with position in the given array or map. + Uses the default column name `pos` for position, and `col` for elements in the + array and `key` and `value` for elements in the map unless specified otherwise. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + one row per array item or map key value including positions as a separate column. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(a=1, intlist=[1,2,3], mapfield={"a": "b"})]) + >>> df.select(posexplode(df.intlist)).collect() + [Row(pos=0, col=1), Row(pos=1, col=2), Row(pos=2, col=3)] + + >>> df.select(posexplode(df.mapfield)).show() + +---+---+-----+ + |pos|key|value| + +---+---+-----+ + | 0| a| b| + +---+---+-----+ + """ + return _invoke_function_over_columns("posexplode", col) + + +@try_remote_functions +def inline(col: "ColumnOrName") -> Column: + """ + Explodes an array of structs into a table. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to explode. + + Returns + ------- + :class:`~pyspark.sql.Column` + generator expression with the inline exploded result. + + See Also + -------- + :meth:`explode` + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([Row(structlist=[Row(a=1, b=2), Row(a=3, b=4)])]) + >>> df.select(inline(df.structlist)).show() + +---+---+ + | a| b| + +---+---+ + | 1| 2| + | 3| 4| + +---+---+ + """ + return _invoke_function_over_columns("inline", col) + + +@try_remote_functions +def explode_outer(col: "ColumnOrName") -> Column: + """ + Returns a new row for each element in the given array or map. + Unlike explode, if the array/map is null or empty then null is produced. + Uses the default column name `col` for elements in the array and + `key` and `value` for elements in the map unless specified otherwise. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + one row per array item or map key value. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(1, ["foo", "bar"], {"x": 1.0}), (2, [], {}), (3, None, None)], + ... ("id", "an_array", "a_map") + ... ) + >>> df.select("id", "an_array", explode_outer("a_map")).show() + +---+----------+----+-----+ + | id| an_array| key|value| + +---+----------+----+-----+ + | 1|[foo, bar]| x| 1.0| + | 2| []|NULL| NULL| + | 3| NULL|NULL| NULL| + +---+----------+----+-----+ + + >>> df.select("id", "a_map", explode_outer("an_array")).show() + +---+----------+----+ + | id| a_map| col| + +---+----------+----+ + | 1|{x -> 1.0}| foo| + | 1|{x -> 1.0}| bar| + | 2| {}|NULL| + | 3| NULL|NULL| + +---+----------+----+ + """ + return _invoke_function_over_columns("explode_outer", col) + + +@try_remote_functions +def posexplode_outer(col: "ColumnOrName") -> Column: + """ + Returns a new row for each element with position in the given array or map. + Unlike posexplode, if the array/map is null or empty then the row (null, null) is produced. + Uses the default column name `pos` for position, and `col` for elements in the + array and `key` and `value` for elements in the map unless specified otherwise. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + one row per array item or map key value including positions as a separate column. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(1, ["foo", "bar"], {"x": 1.0}), (2, [], {}), (3, None, None)], + ... ("id", "an_array", "a_map") + ... ) + >>> df.select("id", "an_array", posexplode_outer("a_map")).show() + +---+----------+----+----+-----+ + | id| an_array| pos| key|value| + +---+----------+----+----+-----+ + | 1|[foo, bar]| 0| x| 1.0| + | 2| []|NULL|NULL| NULL| + | 3| NULL|NULL|NULL| NULL| + +---+----------+----+----+-----+ + >>> df.select("id", "a_map", posexplode_outer("an_array")).show() + +---+----------+----+----+ + | id| a_map| pos| col| + +---+----------+----+----+ + | 1|{x -> 1.0}| 0| foo| + | 1|{x -> 1.0}| 1| bar| + | 2| {}|NULL|NULL| + | 3| NULL|NULL|NULL| + +---+----------+----+----+ + """ + return _invoke_function_over_columns("posexplode_outer", col) + + +@try_remote_functions +def inline_outer(col: "ColumnOrName") -> Column: + """ + Explodes an array of structs into a table. + Unlike inline, if the array is null or empty then null is produced for each nested column. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + input column of values to explode. + + Returns + ------- + :class:`~pyspark.sql.Column` + generator expression with the inline exploded result. + + See Also + -------- + :meth:`explode_outer` + :meth:`inline` + + Notes + ----- + Supports Spark Connect. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df = spark.createDataFrame([ + ... Row(id=1, structlist=[Row(a=1, b=2), Row(a=3, b=4)]), + ... Row(id=2, structlist=[]) + ... ]) + >>> df.select('id', inline_outer(df.structlist)).show() + +---+----+----+ + | id| a| b| + +---+----+----+ + | 1| 1| 2| + | 1| 3| 4| + | 2|NULL|NULL| + +---+----+----+ + """ + return _invoke_function_over_columns("inline_outer", col) + + +@try_remote_functions +def get_json_object(col: "ColumnOrName", path: str) -> Column: + """ + Extracts json object from a json string based on json `path` specified, and returns json string + of the extracted json object. It will return null if the input json string is invalid. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + string column in json format + path : str + path to the json object to extract + + Returns + ------- + :class:`~pyspark.sql.Column` + string representation of given JSON object value. + + Examples + -------- + >>> data = [("1", '''{"f1": "value1", "f2": "value2"}'''), ("2", '''{"f1": "value12"}''')] + >>> df = spark.createDataFrame(data, ("key", "jstring")) + >>> df.select(df.key, get_json_object(df.jstring, '$.f1').alias("c0"), \\ + ... get_json_object(df.jstring, '$.f2').alias("c1") ).collect() + [Row(key='1', c0='value1', c1='value2'), Row(key='2', c0='value12', c1=None)] + """ + return _invoke_function("get_json_object", _to_java_column(col), path) + + +@try_remote_functions +def json_tuple(col: "ColumnOrName", *fields: str) -> Column: + """Creates a new row for a json column according to the given field names. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + string column in json format + fields : str + a field or fields to extract + + Returns + ------- + :class:`~pyspark.sql.Column` + a new row for each given field value from json object + + Examples + -------- + >>> data = [("1", '''{"f1": "value1", "f2": "value2"}'''), ("2", '''{"f1": "value12"}''')] + >>> df = spark.createDataFrame(data, ("key", "jstring")) + >>> df.select(df.key, json_tuple(df.jstring, 'f1', 'f2')).collect() + [Row(key='1', c0='value1', c1='value2'), Row(key='2', c0='value12', c1=None)] + """ + sc = get_active_spark_context() + return _invoke_function("json_tuple", _to_java_column(col), _to_seq(sc, fields)) + + +@try_remote_functions +def from_json( + col: "ColumnOrName", + schema: Union[ArrayType, StructType, Column, str], + options: Optional[Dict[str, str]] = None, +) -> Column: + """ + Parses a column containing a JSON string into a :class:`MapType` with :class:`StringType` + as keys type, :class:`StructType` or :class:`ArrayType` with + the specified schema. Returns `null`, in the case of an unparseable string. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column or column name in JSON format + schema : :class:`DataType` or str + a StructType, ArrayType of StructType or Python string literal with a DDL-formatted string + to use when parsing the json column + options : dict, optional + options to control parsing. accepts the same options as the json datasource. + See `Data Source Option `_ + for the version you use. + + .. # noqa + + Returns + ------- + :class:`~pyspark.sql.Column` + a new column of complex type from given JSON object. + + Examples + -------- + >>> from pyspark.sql.types import * + >>> data = [(1, '''{"a": 1}''')] + >>> schema = StructType([StructField("a", IntegerType())]) + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(from_json(df.value, schema).alias("json")).collect() + [Row(json=Row(a=1))] + >>> df.select(from_json(df.value, "a INT").alias("json")).collect() + [Row(json=Row(a=1))] + >>> df.select(from_json(df.value, "MAP").alias("json")).collect() + [Row(json={'a': 1})] + >>> data = [(1, '''[{"a": 1}]''')] + >>> schema = ArrayType(StructType([StructField("a", IntegerType())])) + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(from_json(df.value, schema).alias("json")).collect() + [Row(json=[Row(a=1)])] + >>> schema = schema_of_json(lit('''{"a": 0}''')) + >>> df.select(from_json(df.value, schema).alias("json")).collect() + [Row(json=Row(a=None))] + >>> data = [(1, '''[1, 2, 3]''')] + >>> schema = ArrayType(IntegerType()) + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(from_json(df.value, schema).alias("json")).collect() + [Row(json=[1, 2, 3])] + """ + + if isinstance(schema, DataType): + schema = schema.json() + elif isinstance(schema, Column): + schema = _to_java_column(schema) + return _invoke_function("from_json", _to_java_column(col), schema, _options_to_str(options)) + + +@try_remote_functions +def to_json(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + """ + Converts a column containing a :class:`StructType`, :class:`ArrayType` or a :class:`MapType` + into a JSON string. Throws an exception, in the case of an unsupported type. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing a struct, an array or a map. + options : dict, optional + options to control converting. accepts the same options as the JSON datasource. + See `Data Source Option `_ + for the version you use. + Additionally the function supports the `pretty` option which enables + pretty JSON generation. + + .. # noqa + + Returns + ------- + :class:`~pyspark.sql.Column` + JSON object as string column. + + Examples + -------- + >>> from pyspark.sql import Row + >>> from pyspark.sql.types import * + >>> data = [(1, Row(age=2, name='Alice'))] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_json(df.value).alias("json")).collect() + [Row(json='{"age":2,"name":"Alice"}')] + >>> data = [(1, [Row(age=2, name='Alice'), Row(age=3, name='Bob')])] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_json(df.value).alias("json")).collect() + [Row(json='[{"age":2,"name":"Alice"},{"age":3,"name":"Bob"}]')] + >>> data = [(1, {"name": "Alice"})] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_json(df.value).alias("json")).collect() + [Row(json='{"name":"Alice"}')] + >>> data = [(1, [{"name": "Alice"}, {"name": "Bob"}])] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_json(df.value).alias("json")).collect() + [Row(json='[{"name":"Alice"},{"name":"Bob"}]')] + >>> data = [(1, ["Alice", "Bob"])] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_json(df.value).alias("json")).collect() + [Row(json='["Alice","Bob"]')] + """ + + return _invoke_function("to_json", _to_java_column(col), _options_to_str(options)) + + +@try_remote_functions +def schema_of_json(json: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + """ + Parses a JSON string and infers its schema in DDL format. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + json : :class:`~pyspark.sql.Column` or str + a JSON string or a foldable string column containing a JSON string. + options : dict, optional + options to control parsing. accepts the same options as the JSON datasource. + See `Data Source Option `_ + for the version you use. + + .. # noqa + + .. versionchanged:: 3.0.0 + It accepts `options` parameter to control schema inferring. + + Returns + ------- + :class:`~pyspark.sql.Column` + a string representation of a :class:`StructType` parsed from given JSON. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(schema_of_json(lit('{"a": 0}')).alias("json")).collect() + [Row(json='STRUCT')] + >>> schema = schema_of_json('{a: 1}', {'allowUnquotedFieldNames':'true'}) + >>> df.select(schema.alias("json")).collect() + [Row(json='STRUCT')] + """ + if isinstance(json, str): + col = _create_column_from_literal(json) + elif isinstance(json, Column): + col = _to_java_column(json) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "json", "arg_type": type(json).__name__}, + ) + + return _invoke_function("schema_of_json", col, _options_to_str(options)) + + +@try_remote_functions +def json_array_length(col: "ColumnOrName") -> Column: + """ + Returns the number of elements in the outermost JSON array. `NULL` is returned in case of + any other valid JSON string, `NULL` or an invalid JSON. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col: :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + length of json array. + + Examples + -------- + >>> df = spark.createDataFrame([(None,), ('[1, 2, 3]',), ('[]',)], ['data']) + >>> df.select(json_array_length(df.data).alias('r')).collect() + [Row(r=None), Row(r=3), Row(r=0)] + """ + return _invoke_function_over_columns("json_array_length", col) + + +@try_remote_functions +def json_object_keys(col: "ColumnOrName") -> Column: + """ + Returns all the keys of the outermost JSON object as an array. If a valid JSON object is + given, all the keys of the outermost object will be returned as an array. If it is any + other valid JSON string, an invalid JSON string or an empty string, the function returns null. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col: :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + all the keys of the outermost JSON object. + + Examples + -------- + >>> df = spark.createDataFrame([(None,), ('{}',), ('{"key1":1, "key2":2}',)], ['data']) + >>> df.select(json_object_keys(df.data).alias('r')).collect() + [Row(r=None), Row(r=[]), Row(r=['key1', 'key2'])] + """ + return _invoke_function_over_columns("json_object_keys", col) + + +@try_remote_functions +def schema_of_csv(csv: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + """ + Parses a CSV string and infers its schema in DDL format. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + csv : :class:`~pyspark.sql.Column` or str + a CSV string or a foldable string column containing a CSV string. + options : dict, optional + options to control parsing. accepts the same options as the CSV datasource. + See `Data Source Option `_ + for the version you use. + + .. # noqa + + Returns + ------- + :class:`~pyspark.sql.Column` + a string representation of a :class:`StructType` parsed from given CSV. + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(schema_of_csv(lit('1|a'), {'sep':'|'}).alias("csv")).collect() + [Row(csv='STRUCT<_c0: INT, _c1: STRING>')] + >>> df.select(schema_of_csv('1|a', {'sep':'|'}).alias("csv")).collect() + [Row(csv='STRUCT<_c0: INT, _c1: STRING>')] + """ + if isinstance(csv, str): + col = _create_column_from_literal(csv) + elif isinstance(csv, Column): + col = _to_java_column(csv) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "csv", "arg_type": type(csv).__name__}, + ) + + return _invoke_function("schema_of_csv", col, _options_to_str(options)) + + +@try_remote_functions +def to_csv(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) -> Column: + """ + Converts a column containing a :class:`StructType` into a CSV string. + Throws an exception, in the case of an unsupported type. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column containing a struct. + options: dict, optional + options to control converting. accepts the same options as the CSV datasource. + See `Data Source Option `_ + for the version you use. + + .. # noqa + + Returns + ------- + :class:`~pyspark.sql.Column` + a CSV string converted from given :class:`StructType`. + + Examples + -------- + >>> from pyspark.sql import Row + >>> data = [(1, Row(age=2, name='Alice'))] + >>> df = spark.createDataFrame(data, ("key", "value")) + >>> df.select(to_csv(df.value).alias("csv")).collect() + [Row(csv='2,Alice')] + """ + + return _invoke_function("to_csv", _to_java_column(col), _options_to_str(options)) + + +@try_remote_functions +def size(col: "ColumnOrName") -> Column: + """ + Collection function: returns the length of the array or map stored in the column. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + length of the array/map. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 2, 3],),([1],),([],)], ['data']) + >>> df.select(size(df.data)).collect() + [Row(size(data)=3), Row(size(data)=1), Row(size(data)=0)] + """ + return _invoke_function_over_columns("size", col) + + +@try_remote_functions +def array_min(col: "ColumnOrName") -> Column: + """ + Collection function: returns the minimum value of the array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + minimum value of array. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 1, 3],), ([None, 10, -1],)], ['data']) + >>> df.select(array_min(df.data).alias('min')).collect() + [Row(min=1), Row(min=-1)] + """ + return _invoke_function_over_columns("array_min", col) + + +@try_remote_functions +def array_max(col: "ColumnOrName") -> Column: + """ + Collection function: returns the maximum value of the array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + maximum value of an array. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 1, 3],), ([None, 10, -1],)], ['data']) + >>> df.select(array_max(df.data).alias('max')).collect() + [Row(max=3), Row(max=10)] + """ + return _invoke_function_over_columns("array_max", col) + + +@try_remote_functions +def array_size(col: "ColumnOrName") -> Column: + """ + Returns the total number of elements in the array. The function returns null for null input. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + total number of elements in the array. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 1, 3],), (None,)], ['data']) + >>> df.select(array_size(df.data).alias('r')).collect() + [Row(r=3), Row(r=None)] + """ + return _invoke_function_over_columns("array_size", col) + + +@try_remote_functions +def cardinality(col: "ColumnOrName") -> Column: + """ + Collection function: returns the length of the array or map stored in the column. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target column to compute on. + + Returns + ------- + :class:`~pyspark.sql.Column` + length of the array/map. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.createDataFrame( + ... [([1, 2, 3],),([1],),([],)], ['data'] + ... ).select(sf.cardinality("data")).show() + +-----------------+ + |cardinality(data)| + +-----------------+ + | 3| + | 1| + | 0| + +-----------------+ + """ + return _invoke_function_over_columns("cardinality", col) + + +@try_remote_functions +def sort_array(col: "ColumnOrName", asc: bool = True) -> Column: + """ + Collection function: sorts the input array in ascending or descending order according + to the natural ordering of the array elements. Null elements will be placed at the beginning + of the returned array in ascending order or at the end of the returned array in descending + order. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + asc : bool, optional + whether to sort in ascending or descending order. If `asc` is True (default) + then ascending and if False then descending. + + Returns + ------- + :class:`~pyspark.sql.Column` + sorted array. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 1, None, 3],),([1],),([],)], ['data']) + >>> df.select(sort_array(df.data).alias('r')).collect() + [Row(r=[None, 1, 2, 3]), Row(r=[1]), Row(r=[])] + >>> df.select(sort_array(df.data, asc=False).alias('r')).collect() + [Row(r=[3, 2, 1, None]), Row(r=[1]), Row(r=[])] + """ + return _invoke_function("sort_array", _to_java_column(col), asc) + + +@try_remote_functions +def array_sort( + col: "ColumnOrName", comparator: Optional[Callable[[Column, Column], Column]] = None +) -> Column: + """ + Collection function: sorts the input array in ascending order. The elements of the input array + must be orderable. Null elements will be placed at the end of the returned array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Can take a `comparator` function. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + comparator : callable, optional + A binary ``(Column, Column) -> Column: ...``. + The comparator will take two + arguments representing two elements of the array. It returns a negative integer, 0, or a + positive integer as the first element is less than, equal to, or greater than the second + element. If the comparator function returns null, the function will fail and raise an error. + + Returns + ------- + :class:`~pyspark.sql.Column` + sorted array. + + Examples + -------- + >>> df = spark.createDataFrame([([2, 1, None, 3],),([1],),([],)], ['data']) + >>> df.select(array_sort(df.data).alias('r')).collect() + [Row(r=[1, 2, 3, None]), Row(r=[1]), Row(r=[])] + >>> df = spark.createDataFrame([(["foo", "foobar", None, "bar"],),(["foo"],),([],)], ['data']) + >>> df.select(array_sort( + ... "data", + ... lambda x, y: when(x.isNull() | y.isNull(), lit(0)).otherwise(length(y) - length(x)) + ... ).alias("r")).collect() + [Row(r=['foobar', 'foo', None, 'bar']), Row(r=['foo']), Row(r=[])] + """ + if comparator is None: + return _invoke_function_over_columns("array_sort", col) + else: + return _invoke_higher_order_function("ArraySort", [col], [comparator]) + + +@try_remote_functions +def shuffle(col: "ColumnOrName") -> Column: + """ + Collection function: Generates a random permutation of the given array. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + The function is non-deterministic. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of elements in random order. + + Examples + -------- + >>> df = spark.createDataFrame([([1, 20, 3, 5],), ([1, 20, None, 3],)], ['data']) + >>> df.select(shuffle(df.data).alias('s')).collect() # doctest: +SKIP + [Row(s=[3, 1, 5, 20]), Row(s=[20, None, 3, 1])] + """ + return _invoke_function_over_columns("shuffle", col) + + +@try_remote_functions +def reverse(col: "ColumnOrName") -> Column: + """ + Collection function: returns a reversed string or an array with reverse order of elements. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + array of elements in reverse order. + + Examples + -------- + >>> df = spark.createDataFrame([('Spark SQL',)], ['data']) + >>> df.select(reverse(df.data).alias('s')).collect() + [Row(s='LQS krapS')] + >>> df = spark.createDataFrame([([2, 1, 3],) ,([1],) ,([],)], ['data']) + >>> df.select(reverse(df.data).alias('r')).collect() + [Row(r=[3, 1, 2]), Row(r=[1]), Row(r=[])] + """ + return _invoke_function_over_columns("reverse", col) + + +@try_remote_functions +def flatten(col: "ColumnOrName") -> Column: + """ + Collection function: creates a single array from an array of arrays. + If a structure of nested arrays is deeper than two levels, + only one level of nesting is removed. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + flattened array. + + Examples + -------- + >>> df = spark.createDataFrame([([[1, 2, 3], [4, 5], [6]],), ([None, [4, 5]],)], ['data']) + >>> df.show(truncate=False) + +------------------------+ + |data | + +------------------------+ + |[[1, 2, 3], [4, 5], [6]]| + |[NULL, [4, 5]] | + +------------------------+ + >>> df.select(flatten(df.data).alias('r')).show() + +------------------+ + | r| + +------------------+ + |[1, 2, 3, 4, 5, 6]| + | NULL| + +------------------+ + """ + return _invoke_function_over_columns("flatten", col) + + +@try_remote_functions +def map_contains_key(col: "ColumnOrName", value: Any) -> Column: + """ + Returns true if the map contains the key. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + value : + a literal value + + Returns + ------- + :class:`~pyspark.sql.Column` + True if key is in the map and False otherwise. + + Examples + -------- + >>> from pyspark.sql.functions import map_contains_key + >>> df = spark.sql("SELECT map(1, 'a', 2, 'b') as data") + >>> df.select(map_contains_key("data", 1)).show() + +---------------------------------+ + |array_contains(map_keys(data), 1)| + +---------------------------------+ + | true| + +---------------------------------+ + >>> df.select(map_contains_key("data", -1)).show() + +----------------------------------+ + |array_contains(map_keys(data), -1)| + +----------------------------------+ + | false| + +----------------------------------+ + """ + return _invoke_function("map_contains_key", _to_java_column(col), value) + + +@try_remote_functions +def map_keys(col: "ColumnOrName") -> Column: + """ + Collection function: Returns an unordered array containing the keys of the map. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + keys of the map as an array. + + Examples + -------- + >>> from pyspark.sql.functions import map_keys + >>> df = spark.sql("SELECT map(1, 'a', 2, 'b') as data") + >>> df.select(map_keys("data").alias("keys")).show() + +------+ + | keys| + +------+ + |[1, 2]| + +------+ + """ + return _invoke_function_over_columns("map_keys", col) + + +@try_remote_functions +def map_values(col: "ColumnOrName") -> Column: + """ + Collection function: Returns an unordered array containing the values of the map. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + values of the map as an array. + + Examples + -------- + >>> from pyspark.sql.functions import map_values + >>> df = spark.sql("SELECT map(1, 'a', 2, 'b') as data") + >>> df.select(map_values("data").alias("values")).show() + +------+ + |values| + +------+ + |[a, b]| + +------+ + """ + return _invoke_function_over_columns("map_values", col) + + +@try_remote_functions +def map_entries(col: "ColumnOrName") -> Column: + """ + Collection function: Returns an unordered array of all entries in the given map. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of key value pairs as a struct type + + Examples + -------- + >>> from pyspark.sql.functions import map_entries + >>> df = spark.sql("SELECT map(1, 'a', 2, 'b') as data") + >>> df = df.select(map_entries("data").alias("entries")) + >>> df.show() + +----------------+ + | entries| + +----------------+ + |[{1, a}, {2, b}]| + +----------------+ + >>> df.printSchema() + root + |-- entries: array (nullable = false) + | |-- element: struct (containsNull = false) + | | |-- key: integer (nullable = false) + | | |-- value: string (nullable = false) + """ + return _invoke_function_over_columns("map_entries", col) + + +@try_remote_functions +def map_from_entries(col: "ColumnOrName") -> Column: + """ + Collection function: Converts an array of entries (key value struct types) to a map + of values. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + + Returns + ------- + :class:`~pyspark.sql.Column` + a map created from the given array of entries. + + Examples + -------- + >>> from pyspark.sql.functions import map_from_entries + >>> df = spark.sql("SELECT array(struct(1, 'a'), struct(2, 'b')) as data") + >>> df.select(map_from_entries("data").alias("map")).show() + +----------------+ + | map| + +----------------+ + |{1 -> a, 2 -> b}| + +----------------+ + """ + return _invoke_function_over_columns("map_from_entries", col) + + +@try_remote_functions +def array_repeat(col: "ColumnOrName", count: Union["ColumnOrName", int]) -> Column: + """ + Collection function: creates an array containing a column repeated count times. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + column name or column that contains the element to be repeated + count : :class:`~pyspark.sql.Column` or str or int + column name, column, or int containing the number of times to repeat the first argument + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of repeated elements. + + Examples + -------- + >>> df = spark.createDataFrame([('ab',)], ['data']) + >>> df.select(array_repeat(df.data, 3).alias('r')).collect() + [Row(r=['ab', 'ab', 'ab'])] + """ + count = lit(count) if isinstance(count, int) else count + + return _invoke_function_over_columns("array_repeat", col, count) + + +@try_remote_functions +def arrays_zip(*cols: "ColumnOrName") -> Column: + """ + Collection function: Returns a merged array of structs in which the N-th struct contains all + N-th values of input arrays. If one of the arrays is shorter than others then + resulting struct type value will be a `null` for missing elements. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + columns of arrays to be merged. + + Returns + ------- + :class:`~pyspark.sql.Column` + merged array of entries. + + Examples + -------- + >>> from pyspark.sql.functions import arrays_zip + >>> df = spark.createDataFrame([([1, 2, 3], [2, 4, 6], [3, 6])], ['vals1', 'vals2', 'vals3']) + >>> df = df.select(arrays_zip(df.vals1, df.vals2, df.vals3).alias('zipped')) + >>> df.show(truncate=False) + +------------------------------------+ + |zipped | + +------------------------------------+ + |[{1, 2, 3}, {2, 4, 6}, {3, 6, NULL}]| + +------------------------------------+ + >>> df.printSchema() + root + |-- zipped: array (nullable = true) + | |-- element: struct (containsNull = false) + | | |-- vals1: long (nullable = true) + | | |-- vals2: long (nullable = true) + | | |-- vals3: long (nullable = true) + """ + return _invoke_function_over_seq_of_columns("arrays_zip", cols) + + +@overload +def map_concat(*cols: "ColumnOrName") -> Column: + ... + + +@overload +def map_concat(__cols: Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]) -> Column: + ... + + +@try_remote_functions +def map_concat( + *cols: Union["ColumnOrName", Union[List["ColumnOrName_"], Tuple["ColumnOrName_", ...]]] +) -> Column: + """Returns the union of all the given maps. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s + + Returns + ------- + :class:`~pyspark.sql.Column` + a map of merged entries from other maps. + + Examples + -------- + >>> from pyspark.sql.functions import map_concat + >>> df = spark.sql("SELECT map(1, 'a', 2, 'b') as map1, map(3, 'c') as map2") + >>> df.select(map_concat("map1", "map2").alias("map3")).show(truncate=False) + +------------------------+ + |map3 | + +------------------------+ + |{1 -> a, 2 -> b, 3 -> c}| + +------------------------+ + """ + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cols[0] # type: ignore[assignment] + return _invoke_function_over_seq_of_columns("map_concat", cols) # type: ignore[arg-type] + + +@try_remote_functions +def sequence( + start: "ColumnOrName", stop: "ColumnOrName", step: Optional["ColumnOrName"] = None +) -> Column: + """ + Generate a sequence of integers from `start` to `stop`, incrementing by `step`. + If `step` is not set, incrementing by 1 if `start` is less than or equal to `stop`, + otherwise -1. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + start : :class:`~pyspark.sql.Column` or str + starting value (inclusive) + stop : :class:`~pyspark.sql.Column` or str + last values (inclusive) + step : :class:`~pyspark.sql.Column` or str, optional + value to add to current to get next element (default is 1) + + Returns + ------- + :class:`~pyspark.sql.Column` + an array of sequence values + + Examples + -------- + >>> df1 = spark.createDataFrame([(-2, 2)], ('C1', 'C2')) + >>> df1.select(sequence('C1', 'C2').alias('r')).collect() + [Row(r=[-2, -1, 0, 1, 2])] + >>> df2 = spark.createDataFrame([(4, -4, -2)], ('C1', 'C2', 'C3')) + >>> df2.select(sequence('C1', 'C2', 'C3').alias('r')).collect() + [Row(r=[4, 2, 0, -2, -4])] + """ + if step is None: + return _invoke_function_over_columns("sequence", start, stop) + else: + return _invoke_function_over_columns("sequence", start, stop, step) + + +@try_remote_functions +def from_csv( + col: "ColumnOrName", + schema: Union[Column, str], + options: Optional[Dict[str, str]] = None, +) -> Column: + """ + Parses a column containing a CSV string to a row with the specified schema. + Returns `null`, in the case of an unparseable string. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + a column or column name in CSV format + schema :class:`~pyspark.sql.Column` or str + a column, or Python string literal with schema in DDL format, to use when parsing the CSV column. + options : dict, optional + options to control parsing. accepts the same options as the CSV datasource. + See `Data Source Option `_ + for the version you use. + + .. # noqa + + Returns + ------- + :class:`~pyspark.sql.Column` + a column of parsed CSV values + + Examples + -------- + >>> data = [("1,2,3",)] + >>> df = spark.createDataFrame(data, ("value",)) + >>> df.select(from_csv(df.value, "a INT, b INT, c INT").alias("csv")).collect() + [Row(csv=Row(a=1, b=2, c=3))] + >>> value = data[0][0] + >>> df.select(from_csv(df.value, schema_of_csv(value)).alias("csv")).collect() + [Row(csv=Row(_c0=1, _c1=2, _c2=3))] + >>> data = [(" abc",)] + >>> df = spark.createDataFrame(data, ("value",)) + >>> options = {'ignoreLeadingWhiteSpace': True} + >>> df.select(from_csv(df.value, "s string", options).alias("csv")).collect() + [Row(csv=Row(s='abc'))] + """ + + get_active_spark_context() + if isinstance(schema, str): + schema = _create_column_from_literal(schema) + elif isinstance(schema, Column): + schema = _to_java_column(schema) + else: + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + + return _invoke_function("from_csv", _to_java_column(col), schema, _options_to_str(options)) + + +def _unresolved_named_lambda_variable(*name_parts: Any) -> Column: + """ + Create `o.a.s.sql.expressions.UnresolvedNamedLambdaVariable`, + convert it to o.s.sql.Column and wrap in Python `Column` + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name_parts : str + """ + sc = get_active_spark_context() + name_parts_seq = _to_seq(sc, name_parts) + expressions = cast(JVMView, sc._jvm).org.apache.spark.sql.catalyst.expressions + return Column( + cast(JVMView, sc._jvm).Column(expressions.UnresolvedNamedLambdaVariable(name_parts_seq)) + ) + + +def _get_lambda_parameters(f: Callable) -> ValuesView[inspect.Parameter]: + signature = inspect.signature(f) + parameters = signature.parameters.values() + + # We should exclude functions that use + # variable args and keyword argnames + # as well as keyword only args + supported_parameter_types = { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY, + } + + # Validate that + # function arity is between 1 and 3 + if not (1 <= len(parameters) <= 3): + raise PySparkValueError( + error_class="WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": f.__name__, "num_args": str(len(parameters))}, + ) + + # and all arguments can be used as positional + if not all(p.kind in supported_parameter_types for p in parameters): + raise PySparkValueError( + error_class="UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": f.__name__}, + ) + + return parameters + + +def _create_lambda(f: Callable) -> Callable: + """ + Create `o.a.s.sql.expressions.LambdaFunction` corresponding + to transformation described by f + + :param f: A Python of one of the following forms: + - (Column) -> Column: ... + - (Column, Column) -> Column: ... + - (Column, Column, Column) -> Column: ... + """ + parameters = _get_lambda_parameters(f) + + sc = get_active_spark_context() + expressions = cast(JVMView, sc._jvm).org.apache.spark.sql.catalyst.expressions + + argnames = ["x", "y", "z"] + args = [ + _unresolved_named_lambda_variable( + expressions.UnresolvedNamedLambdaVariable.freshVarName(arg) + ) + for arg in argnames[: len(parameters)] + ] + + result = f(*args) + + if not isinstance(result, Column): + raise PySparkValueError( + error_class="HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN", + message_parameters={"func_name": f.__name__, "return_type": type(result).__name__}, + ) + + jexpr = result._jc.expr() + jargs = _to_seq(sc, [arg._jc.expr() for arg in args]) + + return expressions.LambdaFunction(jexpr, jargs, False) + + +def _invoke_higher_order_function( + name: str, + cols: List["ColumnOrName"], + funs: List[Callable], +) -> Column: + """ + Invokes expression identified by name, + (relative to ```org.apache.spark.sql.catalyst.expressions``) + and wraps the result with Column (first Scala one, then Python). + + :param name: Name of the expression + :param cols: a list of columns + :param funs: a list of (*Column) -> Column functions. + + :return: a Column + """ + sc = get_active_spark_context() + expressions = cast(JVMView, sc._jvm).org.apache.spark.sql.catalyst.expressions + expr = getattr(expressions, name) + + jcols = [_to_java_column(col).expr() for col in cols] + jfuns = [_create_lambda(f) for f in funs] + + return Column(cast(JVMView, sc._jvm).Column(expr(*jcols + jfuns))) + + +@overload +def transform(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + ... + + +@overload +def transform(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + ... + + +@try_remote_functions +def transform( + col: "ColumnOrName", + f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]], +) -> Column: + """ + Returns an array of elements after applying a transformation to each element in the input array. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + a function that is applied to each element of the input array. + Can take one of the following forms: + + - Unary ``(x: Column) -> Column: ...`` + - Binary ``(x: Column, i: Column) -> Column...``, where the second argument is + a 0-based index of the element. + + and can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + a new array of transformed elements. + + Examples + -------- + >>> df = spark.createDataFrame([(1, [1, 2, 3, 4])], ("key", "values")) + >>> df.select(transform("values", lambda x: x * 2).alias("doubled")).show() + +------------+ + | doubled| + +------------+ + |[2, 4, 6, 8]| + +------------+ + + >>> def alternate(x, i): + ... return when(i % 2 == 0, x).otherwise(-x) + ... + >>> df.select(transform("values", alternate).alias("alternated")).show() + +--------------+ + | alternated| + +--------------+ + |[1, -2, 3, -4]| + +--------------+ + """ + return _invoke_higher_order_function("ArrayTransform", [col], [f]) + + +@try_remote_functions +def exists(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + """ + Returns whether a predicate holds for one or more elements in the array. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + ``(x: Column) -> Column: ...`` returning the Boolean expression. + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + True if "any" element of an array evaluates to True when passed as an argument to + given function and False otherwise. + + Examples + -------- + >>> df = spark.createDataFrame([(1, [1, 2, 3, 4]), (2, [3, -1, 0])],("key", "values")) + >>> df.select(exists("values", lambda x: x < 0).alias("any_negative")).show() + +------------+ + |any_negative| + +------------+ + | false| + | true| + +------------+ + """ + return _invoke_higher_order_function("ArrayExists", [col], [f]) + + +@try_remote_functions +def forall(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + """ + Returns whether a predicate holds for every element in the array. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + ``(x: Column) -> Column: ...`` returning the Boolean expression. + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + True if "all" elements of an array evaluates to True when passed as an argument to + given function and False otherwise. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(1, ["bar"]), (2, ["foo", "bar"]), (3, ["foobar", "foo"])], + ... ("key", "values") + ... ) + >>> df.select(forall("values", lambda x: x.rlike("foo")).alias("all_foo")).show() + +-------+ + |all_foo| + +-------+ + | false| + | false| + | true| + +-------+ + """ + return _invoke_higher_order_function("ArrayForAll", [col], [f]) + + +@overload +def filter(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column: + ... + + +@overload +def filter(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + ... + + +@try_remote_functions +def filter( + col: "ColumnOrName", + f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]], +) -> Column: + """ + Returns an array of elements for which a predicate holds in a given array. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + A function that returns the Boolean expression. + Can take one of the following forms: + + - Unary ``(x: Column) -> Column: ...`` + - Binary ``(x: Column, i: Column) -> Column...``, where the second argument is + a 0-based index of the element. + + and can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + filtered array of elements where given function evaluated to True + when passed as an argument. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(1, ["2018-09-20", "2019-02-03", "2019-07-01", "2020-06-01"])], + ... ("key", "values") + ... ) + >>> def after_second_quarter(x): + ... return month(to_date(x)) > 6 + ... + >>> df.select( + ... filter("values", after_second_quarter).alias("after_second_quarter") + ... ).show(truncate=False) + +------------------------+ + |after_second_quarter | + +------------------------+ + |[2018-09-20, 2019-07-01]| + +------------------------+ + """ + return _invoke_higher_order_function("ArrayFilter", [col], [f]) + + +@try_remote_functions +def aggregate( + col: "ColumnOrName", + initialValue: "ColumnOrName", + merge: Callable[[Column, Column], Column], + finish: Optional[Callable[[Column], Column]] = None, +) -> Column: + """ + Applies a binary operator to an initial state and all elements in the array, + and reduces this to a single state. The final state is converted into the final result + by applying a finish function. + + Both functions can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + initialValue : :class:`~pyspark.sql.Column` or str + initial value. Name of column or expression + merge : function + a binary function ``(acc: Column, x: Column) -> Column...`` returning expression + of the same type as ``zero`` + finish : function + an optional unary function ``(x: Column) -> Column: ...`` + used to convert accumulated value. + + Returns + ------- + :class:`~pyspark.sql.Column` + final value after aggregate function is applied. + + Examples + -------- + >>> df = spark.createDataFrame([(1, [20.0, 4.0, 2.0, 6.0, 10.0])], ("id", "values")) + >>> df.select(aggregate("values", lit(0.0), lambda acc, x: acc + x).alias("sum")).show() + +----+ + | sum| + +----+ + |42.0| + +----+ + + >>> def merge(acc, x): + ... count = acc.count + 1 + ... sum = acc.sum + x + ... return struct(count.alias("count"), sum.alias("sum")) + ... + >>> df.select( + ... aggregate( + ... "values", + ... struct(lit(0).alias("count"), lit(0.0).alias("sum")), + ... merge, + ... lambda acc: acc.sum / acc.count, + ... ).alias("mean") + ... ).show() + +----+ + |mean| + +----+ + | 8.4| + +----+ + """ + if finish is not None: + return _invoke_higher_order_function("ArrayAggregate", [col, initialValue], [merge, finish]) + + else: + return _invoke_higher_order_function("ArrayAggregate", [col, initialValue], [merge]) + + +@try_remote_functions +def reduce( + col: "ColumnOrName", + initialValue: "ColumnOrName", + merge: Callable[[Column, Column], Column], + finish: Optional[Callable[[Column], Column]] = None, +) -> Column: + """ + Applies a binary operator to an initial state and all elements in the array, + and reduces this to a single state. The final state is converted into the final result + by applying a finish function. + + Both functions can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + initialValue : :class:`~pyspark.sql.Column` or str + initial value. Name of column or expression + merge : function + a binary function ``(acc: Column, x: Column) -> Column...`` returning expression + of the same type as ``zero`` + finish : function + an optional unary function ``(x: Column) -> Column: ...`` + used to convert accumulated value. + + Returns + ------- + :class:`~pyspark.sql.Column` + final value after aggregate function is applied. + + Examples + -------- + >>> df = spark.createDataFrame([(1, [20.0, 4.0, 2.0, 6.0, 10.0])], ("id", "values")) + >>> df.select(reduce("values", lit(0.0), lambda acc, x: acc + x).alias("sum")).show() + +----+ + | sum| + +----+ + |42.0| + +----+ + + >>> def merge(acc, x): + ... count = acc.count + 1 + ... sum = acc.sum + x + ... return struct(count.alias("count"), sum.alias("sum")) + ... + >>> df.select( + ... reduce( + ... "values", + ... struct(lit(0).alias("count"), lit(0.0).alias("sum")), + ... merge, + ... lambda acc: acc.sum / acc.count, + ... ).alias("mean") + ... ).show() + +----+ + |mean| + +----+ + | 8.4| + +----+ + """ + if finish is not None: + return _invoke_higher_order_function("ArrayAggregate", [col, initialValue], [merge, finish]) + + else: + return _invoke_higher_order_function("ArrayAggregate", [col, initialValue], [merge]) + + +@try_remote_functions +def zip_with( + left: "ColumnOrName", + right: "ColumnOrName", + f: Callable[[Column, Column], Column], +) -> Column: + """ + Merge two given arrays, element-wise, into a single array using a function. + If one array is shorter, nulls are appended at the end to match the length of the longer + array, before applying the function. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or str + name of the first column or expression + right : :class:`~pyspark.sql.Column` or str + name of the second column or expression + f : function + a binary function ``(x1: Column, x2: Column) -> Column...`` + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + array of calculated values derived by applying given function to each pair of arguments. + + Examples + -------- + >>> df = spark.createDataFrame([(1, [1, 3, 5, 8], [0, 2, 4, 6])], ("id", "xs", "ys")) + >>> df.select(zip_with("xs", "ys", lambda x, y: x ** y).alias("powers")).show(truncate=False) + +---------------------------+ + |powers | + +---------------------------+ + |[1.0, 9.0, 625.0, 262144.0]| + +---------------------------+ + + >>> df = spark.createDataFrame([(1, ["foo", "bar"], [1, 2, 3])], ("id", "xs", "ys")) + >>> df.select(zip_with("xs", "ys", lambda x, y: concat_ws("_", x, y)).alias("xs_ys")).show() + +-----------------+ + | xs_ys| + +-----------------+ + |[foo_1, bar_2, 3]| + +-----------------+ + """ + return _invoke_higher_order_function("ZipWith", [left, right], [f]) + + +@try_remote_functions +def transform_keys(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + """ + Applies a function to every key-value pair in a map and returns + a map with the results of those applications as the new keys for the pairs. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + a binary function ``(k: Column, v: Column) -> Column...`` + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + a new map of enties where new keys were calculated by applying given function to + each key value argument. + + Examples + -------- + >>> df = spark.createDataFrame([(1, {"foo": -2.0, "bar": 2.0})], ("id", "data")) + >>> row = df.select(transform_keys( + ... "data", lambda k, _: upper(k)).alias("data_upper") + ... ).head() + >>> sorted(row["data_upper"].items()) + [('BAR', 2.0), ('FOO', -2.0)] + """ + return _invoke_higher_order_function("TransformKeys", [col], [f]) + + +@try_remote_functions +def transform_values(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + """ + Applies a function to every key-value pair in a map and returns + a map with the results of those applications as the new values for the pairs. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + a binary function ``(k: Column, v: Column) -> Column...`` + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + a new map of enties where new values were calculated by applying given function to + each key value argument. + + Examples + -------- + >>> df = spark.createDataFrame([(1, {"IT": 10.0, "SALES": 2.0, "OPS": 24.0})], ("id", "data")) + >>> row = df.select(transform_values( + ... "data", lambda k, v: when(k.isin("IT", "OPS"), v + 10.0).otherwise(v) + ... ).alias("new_data")).head() + >>> sorted(row["new_data"].items()) + [('IT', 20.0), ('OPS', 34.0), ('SALES', 2.0)] + """ + return _invoke_higher_order_function("TransformValues", [col], [f]) + + +@try_remote_functions +def map_filter(col: "ColumnOrName", f: Callable[[Column, Column], Column]) -> Column: + """ + Returns a map whose key-value pairs satisfy a predicate. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + name of column or expression + f : function + a binary function ``(k: Column, v: Column) -> Column...`` + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + filtered map. + + Examples + -------- + >>> df = spark.createDataFrame([(1, {"foo": 42.0, "bar": 1.0, "baz": 32.0})], ("id", "data")) + >>> row = df.select(map_filter( + ... "data", lambda _, v: v > 30.0).alias("data_filtered") + ... ).head() + >>> sorted(row["data_filtered"].items()) + [('baz', 32.0), ('foo', 42.0)] + """ + return _invoke_higher_order_function("MapFilter", [col], [f]) + + +@try_remote_functions +def map_zip_with( + col1: "ColumnOrName", + col2: "ColumnOrName", + f: Callable[[Column, Column, Column], Column], +) -> Column: + """ + Merge two given maps, key-wise into a single map using a function. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + name of the first column or expression + col2 : :class:`~pyspark.sql.Column` or str + name of the second column or expression + f : function + a ternary function ``(k: Column, v1: Column, v2: Column) -> Column...`` + Can use methods of :class:`~pyspark.sql.Column`, functions defined in + :py:mod:`pyspark.sql.functions` and Scala ``UserDefinedFunctions``. + Python ``UserDefinedFunctions`` are not supported + (`SPARK-27052 `__). + + Returns + ------- + :class:`~pyspark.sql.Column` + zipped map where entries are calculated by applying given function to each + pair of arguments. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (1, {"IT": 24.0, "SALES": 12.00}, {"IT": 2.0, "SALES": 1.4})], + ... ("id", "base", "ratio") + ... ) + >>> row = df.select(map_zip_with( + ... "base", "ratio", lambda k, v1, v2: round(v1 * v2, 2)).alias("updated_data") + ... ).head() + >>> sorted(row["updated_data"].items()) + [('IT', 48.0), ('SALES', 16.8)] + """ + return _invoke_higher_order_function("MapZipWith", [col1, col2], [f]) + + +@try_remote_functions +def str_to_map( + text: "ColumnOrName", + pairDelim: Optional["ColumnOrName"] = None, + keyValueDelim: Optional["ColumnOrName"] = None, +) -> Column: + """ + Creates a map after splitting the text into key/value pairs using delimiters. + Both `pairDelim` and `keyValueDelim` are treated as regular expressions. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + text : :class:`~pyspark.sql.Column` or str + Input column or strings. + pairDelim : :class:`~pyspark.sql.Column` or str, optional + delimiter to use to split pair. + keyValueDelim : :class:`~pyspark.sql.Column` or str, optional + delimiter to use to split key/value. + + Examples + -------- + >>> df = spark.createDataFrame([("a:1,b:2,c:3",)], ["e"]) + >>> df.select(str_to_map(df.e, lit(","), lit(":")).alias('r')).collect() + [Row(r={'a': '1', 'b': '2', 'c': '3'})] + + >>> df = spark.createDataFrame([("a:1,b:2,c:3",)], ["e"]) + >>> df.select(str_to_map(df.e, lit(",")).alias('r')).collect() + [Row(r={'a': '1', 'b': '2', 'c': '3'})] + + >>> df = spark.createDataFrame([("a:1,b:2,c:3",)], ["e"]) + >>> df.select(str_to_map(df.e).alias('r')).collect() + [Row(r={'a': '1', 'b': '2', 'c': '3'})] + """ + if pairDelim is None: + pairDelim = lit(",") + if keyValueDelim is None: + keyValueDelim = lit(":") + return _invoke_function_over_columns("str_to_map", text, pairDelim, keyValueDelim) + + +# ---------------------- Partition transform functions -------------------------------- + + +@try_remote_functions +def years(col: "ColumnOrName") -> Column: + """ + Partition transform function: A transform for timestamps and dates + to partition data into years. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date or timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + data partitioned by years. + + Examples + -------- + >>> df.writeTo("catalog.db.table").partitionedBy( # doctest: +SKIP + ... years("ts") + ... ).createOrReplace() + + Notes + ----- + This function can be used only in combination with + :py:meth:`~pyspark.sql.readwriter.DataFrameWriterV2.partitionedBy` + method of the `DataFrameWriterV2`. + + """ + return _invoke_function_over_columns("years", col) + + +@try_remote_functions +def months(col: "ColumnOrName") -> Column: + """ + Partition transform function: A transform for timestamps and dates + to partition data into months. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date or timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + data partitioned by months. + + Examples + -------- + >>> df.writeTo("catalog.db.table").partitionedBy( + ... months("ts") + ... ).createOrReplace() # doctest: +SKIP + + Notes + ----- + This function can be used only in combination with + :py:meth:`~pyspark.sql.readwriter.DataFrameWriterV2.partitionedBy` + method of the `DataFrameWriterV2`. + + """ + return _invoke_function_over_columns("months", col) + + +@try_remote_functions +def days(col: "ColumnOrName") -> Column: + """ + Partition transform function: A transform for timestamps and dates + to partition data into days. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date or timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + data partitioned by days. + + Examples + -------- + >>> df.writeTo("catalog.db.table").partitionedBy( # doctest: +SKIP + ... days("ts") + ... ).createOrReplace() + + Notes + ----- + This function can be used only in combination with + :py:meth:`~pyspark.sql.readwriter.DataFrameWriterV2.partitionedBy` + method of the `DataFrameWriterV2`. + + """ + return _invoke_function_over_columns("days", col) + + +@try_remote_functions +def hours(col: "ColumnOrName") -> Column: + """ + Partition transform function: A transform for timestamps + to partition data into hours. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date or timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + data partitioned by hours. + + Examples + -------- + >>> df.writeTo("catalog.db.table").partitionedBy( # doctest: +SKIP + ... hours("ts") + ... ).createOrReplace() + + Notes + ----- + This function can be used only in combination with + :py:meth:`~pyspark.sql.readwriter.DataFrameWriterV2.partitionedBy` + method of the `DataFrameWriterV2`. + + """ + return _invoke_function_over_columns("hours", col) + + +@try_remote_functions +def convert_timezone( + sourceTz: Optional[Column], targetTz: Column, sourceTs: "ColumnOrName" +) -> Column: + """ + Converts the timestamp without time zone `sourceTs` + from the `sourceTz` time zone to `targetTz`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + sourceTz : :class:`~pyspark.sql.Column` + the time zone for the input timestamp. If it is missed, + the current session time zone is used as the source time zone. + targetTz : :class:`~pyspark.sql.Column` + the time zone to which the input timestamp should be converted. + sourceTs : :class:`~pyspark.sql.Column` + a timestamp without time zone. + + Returns + ------- + :class:`~pyspark.sql.Column` + timestamp for converted time zone. + + Examples + -------- + >>> df = spark.createDataFrame([('2015-04-08',)], ['dt']) + >>> df.select(convert_timezone( # doctest: +SKIP + ... None, lit('Asia/Hong_Kong'), 'dt').alias('ts') + ... ).show() + +-------------------+ + | ts| + +-------------------+ + |2015-04-08 00:00:00| + +-------------------+ + >>> df.select(convert_timezone( + ... lit('America/Los_Angeles'), lit('Asia/Hong_Kong'), 'dt').alias('ts') + ... ).show() + +-------------------+ + | ts| + +-------------------+ + |2015-04-08 15:00:00| + +-------------------+ + """ + if sourceTz is None: + return _invoke_function_over_columns("convert_timezone", targetTz, sourceTs) + else: + return _invoke_function_over_columns("convert_timezone", sourceTz, targetTz, sourceTs) + + +@try_remote_functions +def make_dt_interval( + days: Optional["ColumnOrName"] = None, + hours: Optional["ColumnOrName"] = None, + mins: Optional["ColumnOrName"] = None, + secs: Optional["ColumnOrName"] = None, +) -> Column: + """ + Make DayTimeIntervalType duration from days, hours, mins and secs. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + days : :class:`~pyspark.sql.Column` or str + the number of days, positive or negative + hours : :class:`~pyspark.sql.Column` or str + the number of hours, positive or negative + mins : :class:`~pyspark.sql.Column` or str + the number of minutes, positive or negative + secs : :class:`~pyspark.sql.Column` or str + the number of seconds with the fractional part in microsecond precision. + + Examples + -------- + >>> df = spark.createDataFrame([[1, 12, 30, 01.001001]], + ... ["day", "hour", "min", "sec"]) + >>> df.select(make_dt_interval( + ... df.day, df.hour, df.min, df.sec).alias('r') + ... ).show(truncate=False) + +------------------------------------------+ + |r | + +------------------------------------------+ + |INTERVAL '1 12:30:01.001001' DAY TO SECOND| + +------------------------------------------+ + + >>> df.select(make_dt_interval( + ... df.day, df.hour, df.min).alias('r') + ... ).show(truncate=False) + +-----------------------------------+ + |r | + +-----------------------------------+ + |INTERVAL '1 12:30:00' DAY TO SECOND| + +-----------------------------------+ + + >>> df.select(make_dt_interval( + ... df.day, df.hour).alias('r') + ... ).show(truncate=False) + +-----------------------------------+ + |r | + +-----------------------------------+ + |INTERVAL '1 12:00:00' DAY TO SECOND| + +-----------------------------------+ + + >>> df.select(make_dt_interval(df.day).alias('r')).show(truncate=False) + +-----------------------------------+ + |r | + +-----------------------------------+ + |INTERVAL '1 00:00:00' DAY TO SECOND| + +-----------------------------------+ + + >>> df.select(make_dt_interval().alias('r')).show(truncate=False) + +-----------------------------------+ + |r | + +-----------------------------------+ + |INTERVAL '0 00:00:00' DAY TO SECOND| + +-----------------------------------+ + """ + _days = lit(0) if days is None else days + _hours = lit(0) if hours is None else hours + _mins = lit(0) if mins is None else mins + _secs = lit(decimal.Decimal(0)) if secs is None else secs + return _invoke_function_over_columns("make_dt_interval", _days, _hours, _mins, _secs) + + +@try_remote_functions +def make_interval( + years: Optional["ColumnOrName"] = None, + months: Optional["ColumnOrName"] = None, + weeks: Optional["ColumnOrName"] = None, + days: Optional["ColumnOrName"] = None, + hours: Optional["ColumnOrName"] = None, + mins: Optional["ColumnOrName"] = None, + secs: Optional["ColumnOrName"] = None, +) -> Column: + """ + Make interval from years, months, weeks, days, hours, mins and secs. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + years : :class:`~pyspark.sql.Column` or str + the number of years, positive or negative + months : :class:`~pyspark.sql.Column` or str + the number of months, positive or negative + weeks : :class:`~pyspark.sql.Column` or str + the number of weeks, positive or negative + days : :class:`~pyspark.sql.Column` or str + the number of days, positive or negative + hours : :class:`~pyspark.sql.Column` or str + the number of hours, positive or negative + mins : :class:`~pyspark.sql.Column` or str + the number of minutes, positive or negative + secs : :class:`~pyspark.sql.Column` or str + the number of seconds with the fractional part in microsecond precision. + + Examples + -------- + >>> df = spark.createDataFrame([[100, 11, 1, 1, 12, 30, 01.001001]], + ... ["year", "month", "week", "day", "hour", "min", "sec"]) + >>> df.select(make_interval( + ... df.year, df.month, df.week, df.day, df.hour, df.min, df.sec).alias('r') + ... ).show(truncate=False) + +---------------------------------------------------------------+ + |r | + +---------------------------------------------------------------+ + |100 years 11 months 8 days 12 hours 30 minutes 1.001001 seconds| + +---------------------------------------------------------------+ + + >>> df.select(make_interval( + ... df.year, df.month, df.week, df.day, df.hour, df.min).alias('r') + ... ).show(truncate=False) + +----------------------------------------------+ + |r | + +----------------------------------------------+ + |100 years 11 months 8 days 12 hours 30 minutes| + +----------------------------------------------+ + + >>> df.select(make_interval( + ... df.year, df.month, df.week, df.day, df.hour).alias('r') + ... ).show(truncate=False) + +-----------------------------------+ + |r | + +-----------------------------------+ + |100 years 11 months 8 days 12 hours| + +-----------------------------------+ + + >>> df.select(make_interval( + ... df.year, df.month, df.week, df.day).alias('r') + ... ).show(truncate=False) + +--------------------------+ + |r | + +--------------------------+ + |100 years 11 months 8 days| + +--------------------------+ + + >>> df.select(make_interval( + ... df.year, df.month, df.week).alias('r') + ... ).show(truncate=False) + +--------------------------+ + |r | + +--------------------------+ + |100 years 11 months 7 days| + +--------------------------+ + + >>> df.select(make_interval(df.year, df.month).alias('r')).show(truncate=False) + +-------------------+ + |r | + +-------------------+ + |100 years 11 months| + +-------------------+ + + >>> df.select(make_interval(df.year).alias('r')).show(truncate=False) + +---------+ + |r | + +---------+ + |100 years| + +---------+ + """ + _years = lit(0) if years is None else years + _months = lit(0) if months is None else months + _weeks = lit(0) if weeks is None else weeks + _days = lit(0) if days is None else days + _hours = lit(0) if hours is None else hours + _mins = lit(0) if mins is None else mins + _secs = lit(decimal.Decimal(0)) if secs is None else secs + return _invoke_function_over_columns( + "make_interval", _years, _months, _weeks, _days, _hours, _mins, _secs + ) + + +@try_remote_functions +def make_timestamp( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", + timezone: Optional["ColumnOrName"] = None, +) -> Column: + """ + Create timestamp from years, months, days, hours, mins, secs and timezone fields. + The result data type is consistent with the value of configuration `spark.sql.timestampType`. + If the configuration `spark.sql.ansi.enabled` is false, the function returns NULL + on invalid inputs. Otherwise, it will throw an error instead. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + years : :class:`~pyspark.sql.Column` or str + the year to represent, from 1 to 9999 + months : :class:`~pyspark.sql.Column` or str + the month-of-year to represent, from 1 (January) to 12 (December) + days : :class:`~pyspark.sql.Column` or str + the day-of-month to represent, from 1 to 31 + hours : :class:`~pyspark.sql.Column` or str + the hour-of-day to represent, from 0 to 23 + mins : :class:`~pyspark.sql.Column` or str + the minute-of-hour to represent, from 0 to 59 + secs : :class:`~pyspark.sql.Column` or str + the second-of-minute and its micro-fraction to represent, from 0 to 60. + The value can be either an integer like 13 , or a fraction like 13.123. + If the sec argument equals to 60, the seconds field is set + to 0 and 1 minute is added to the final timestamp. + timezone : :class:`~pyspark.sql.Column` or str + the time zone identifier. For example, CET, UTC and etc. + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([[2014, 12, 28, 6, 30, 45.887, 'CET']], + ... ["year", "month", "day", "hour", "min", "sec", "timezone"]) + >>> df.select(make_timestamp( + ... df.year, df.month, df.day, df.hour, df.min, df.sec, df.timezone).alias('r') + ... ).show(truncate=False) + +-----------------------+ + |r | + +-----------------------+ + |2014-12-27 21:30:45.887| + +-----------------------+ + + >>> df.select(make_timestamp( + ... df.year, df.month, df.day, df.hour, df.min, df.sec).alias('r') + ... ).show(truncate=False) + +-----------------------+ + |r | + +-----------------------+ + |2014-12-28 06:30:45.887| + +-----------------------+ + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + if timezone is not None: + return _invoke_function_over_columns( + "make_timestamp", years, months, days, hours, mins, secs, timezone + ) + else: + return _invoke_function_over_columns( + "make_timestamp", years, months, days, hours, mins, secs + ) + + +@try_remote_functions +def make_timestamp_ltz( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", + timezone: Optional["ColumnOrName"] = None, +) -> Column: + """ + Create the current timestamp with local time zone from years, months, days, hours, mins, + secs and timezone fields. If the configuration `spark.sql.ansi.enabled` is false, + the function returns NULL on invalid inputs. Otherwise, it will throw an error instead. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + years : :class:`~pyspark.sql.Column` or str + the year to represent, from 1 to 9999 + months : :class:`~pyspark.sql.Column` or str + the month-of-year to represent, from 1 (January) to 12 (December) + days : :class:`~pyspark.sql.Column` or str + the day-of-month to represent, from 1 to 31 + hours : :class:`~pyspark.sql.Column` or str + the hour-of-day to represent, from 0 to 23 + mins : :class:`~pyspark.sql.Column` or str + the minute-of-hour to represent, from 0 to 59 + secs : :class:`~pyspark.sql.Column` or str + the second-of-minute and its micro-fraction to represent, from 0 to 60. + The value can be either an integer like 13 , or a fraction like 13.123. + If the sec argument equals to 60, the seconds field is set + to 0 and 1 minute is added to the final timestamp. + timezone : :class:`~pyspark.sql.Column` or str + the time zone identifier. For example, CET, UTC and etc. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([[2014, 12, 28, 6, 30, 45.887, 'CET']], + ... ["year", "month", "day", "hour", "min", "sec", "timezone"]) + >>> df.select(sf.make_timestamp_ltz( + ... df.year, df.month, df.day, df.hour, df.min, df.sec, df.timezone) + ... ).show(truncate=False) + +--------------------------------------------------------------+ + |make_timestamp_ltz(year, month, day, hour, min, sec, timezone)| + +--------------------------------------------------------------+ + |2014-12-27 21:30:45.887 | + +--------------------------------------------------------------+ + + >>> df.select(sf.make_timestamp_ltz( + ... df.year, df.month, df.day, df.hour, df.min, df.sec) + ... ).show(truncate=False) + +----------------------------------------------------+ + |make_timestamp_ltz(year, month, day, hour, min, sec)| + +----------------------------------------------------+ + |2014-12-28 06:30:45.887 | + +----------------------------------------------------+ + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + if timezone is not None: + return _invoke_function_over_columns( + "make_timestamp_ltz", years, months, days, hours, mins, secs, timezone + ) + else: + return _invoke_function_over_columns( + "make_timestamp_ltz", years, months, days, hours, mins, secs + ) + + +@try_remote_functions +def make_timestamp_ntz( + years: "ColumnOrName", + months: "ColumnOrName", + days: "ColumnOrName", + hours: "ColumnOrName", + mins: "ColumnOrName", + secs: "ColumnOrName", +) -> Column: + """ + Create local date-time from years, months, days, hours, mins, secs fields. + If the configuration `spark.sql.ansi.enabled` is false, the function returns NULL + on invalid inputs. Otherwise, it will throw an error instead. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + years : :class:`~pyspark.sql.Column` or str + the year to represent, from 1 to 9999 + months : :class:`~pyspark.sql.Column` or str + the month-of-year to represent, from 1 (January) to 12 (December) + days : :class:`~pyspark.sql.Column` or str + the day-of-month to represent, from 1 to 31 + hours : :class:`~pyspark.sql.Column` or str + the hour-of-day to represent, from 0 to 23 + mins : :class:`~pyspark.sql.Column` or str + the minute-of-hour to represent, from 0 to 59 + secs : :class:`~pyspark.sql.Column` or str + the second-of-minute and its micro-fraction to represent, from 0 to 60. + The value can be either an integer like 13 , or a fraction like 13.123. + If the sec argument equals to 60, the seconds field is set + to 0 and 1 minute is added to the final timestamp. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([[2014, 12, 28, 6, 30, 45.887]], + ... ["year", "month", "day", "hour", "min", "sec"]) + >>> df.select(sf.make_timestamp_ntz( + ... df.year, df.month, df.day, df.hour, df.min, df.sec) + ... ).show(truncate=False) + +----------------------------------------------------+ + |make_timestamp_ntz(year, month, day, hour, min, sec)| + +----------------------------------------------------+ + |2014-12-28 06:30:45.887 | + +----------------------------------------------------+ + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + return _invoke_function_over_columns( + "make_timestamp_ntz", years, months, days, hours, mins, secs + ) + + +@try_remote_functions +def make_ym_interval( + years: Optional["ColumnOrName"] = None, + months: Optional["ColumnOrName"] = None, +) -> Column: + """ + Make year-month interval from years, months. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + years : :class:`~pyspark.sql.Column` or str + the number of years, positive or negative + months : :class:`~pyspark.sql.Column` or str + the number of months, positive or negative + + Examples + -------- + >>> spark.conf.set("spark.sql.session.timeZone", "America/Los_Angeles") + >>> df = spark.createDataFrame([[2014, 12]], ["year", "month"]) + >>> df.select(make_ym_interval(df.year, df.month).alias('r')).show(truncate=False) + +-------------------------------+ + |r | + +-------------------------------+ + |INTERVAL '2015-0' YEAR TO MONTH| + +-------------------------------+ + >>> spark.conf.unset("spark.sql.session.timeZone") + """ + _years = lit(0) if years is None else years + _months = lit(0) if months is None else months + return _invoke_function_over_columns("make_ym_interval", _years, _months) + + +@try_remote_functions +def bucket(numBuckets: Union[Column, int], col: "ColumnOrName") -> Column: + """ + Partition transform function: A transform for any type that partitions + by a hash of the input column. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> df.writeTo("catalog.db.table").partitionedBy( # doctest: +SKIP + ... bucket(42, "ts") + ... ).createOrReplace() + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + target date or timestamp column to work on. + + Returns + ------- + :class:`~pyspark.sql.Column` + data partitioned by given columns. + + Notes + ----- + This function can be used only in combination with + :py:meth:`~pyspark.sql.readwriter.DataFrameWriterV2.partitionedBy` + method of the `DataFrameWriterV2`. + + """ + if not isinstance(numBuckets, (int, Column)): + raise PySparkTypeError( + error_class="NOT_COLUMN_OR_INT", + message_parameters={"arg_name": "numBuckets", "arg_type": type(numBuckets).__name__}, + ) + + get_active_spark_context() + numBuckets = ( + _create_column_from_literal(numBuckets) + if isinstance(numBuckets, int) + else _to_java_column(numBuckets) + ) + return _invoke_function("bucket", numBuckets, _to_java_column(col)) + + +@try_remote_functions +def call_udf(udfName: str, *cols: "ColumnOrName") -> Column: + """ + Call an user-defined function. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + udfName : str + name of the user defined function (UDF) + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s to be used in the UDF + + Returns + ------- + :class:`~pyspark.sql.Column` + result of executed udf. + + Examples + -------- + >>> from pyspark.sql.functions import call_udf, col + >>> from pyspark.sql.types import IntegerType, StringType + >>> df = spark.createDataFrame([(1, "a"),(2, "b"), (3, "c")],["id", "name"]) + >>> _ = spark.udf.register("intX2", lambda i: i * 2, IntegerType()) + >>> df.select(call_udf("intX2", "id")).show() + +---------+ + |intX2(id)| + +---------+ + | 2| + | 4| + | 6| + +---------+ + >>> _ = spark.udf.register("strX2", lambda s: s * 2, StringType()) + >>> df.select(call_udf("strX2", col("name"))).show() + +-----------+ + |strX2(name)| + +-----------+ + | aa| + | bb| + | cc| + +-----------+ + """ + sc = get_active_spark_context() + return _invoke_function("call_udf", udfName, _to_seq(sc, cols, _to_java_column)) + + +@try_remote_functions +def call_function(funcName: str, *cols: "ColumnOrName") -> Column: + """ + Call a SQL function. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + funcName : str + function name that follows the SQL identifier syntax (can be quoted, can be qualified) + cols : :class:`~pyspark.sql.Column` or str + column names or :class:`~pyspark.sql.Column`\\s to be used in the function + + Returns + ------- + :class:`~pyspark.sql.Column` + result of executed function. + + Examples + -------- + >>> from pyspark.sql.functions import call_udf, col + >>> from pyspark.sql.types import IntegerType, StringType + >>> df = spark.createDataFrame([(1, "a"),(2, "b"), (3, "c")],["id", "name"]) + >>> _ = spark.udf.register("intX2", lambda i: i * 2, IntegerType()) + >>> df.select(call_function("intX2", "id")).show() + +---------+ + |intX2(id)| + +---------+ + | 2| + | 4| + | 6| + +---------+ + >>> _ = spark.udf.register("strX2", lambda s: s * 2, StringType()) + >>> df.select(call_function("strX2", col("name"))).show() + +-----------+ + |strX2(name)| + +-----------+ + | aa| + | bb| + | cc| + +-----------+ + >>> df.select(call_function("avg", col("id"))).show() + +-------+ + |avg(id)| + +-------+ + | 2.0| + +-------+ + >>> _ = spark.sql("CREATE FUNCTION custom_avg AS 'test.org.apache.spark.sql.MyDoubleAvg'") + ... # doctest: +SKIP + >>> df.select(call_function("custom_avg", col("id"))).show() + ... # doctest: +SKIP + +------------------------------------+ + |spark_catalog.default.custom_avg(id)| + +------------------------------------+ + | 102.0| + +------------------------------------+ + >>> df.select(call_function("spark_catalog.default.custom_avg", col("id"))).show() + ... # doctest: +SKIP + +------------------------------------+ + |spark_catalog.default.custom_avg(id)| + +------------------------------------+ + | 102.0| + +------------------------------------+ + """ + sc = get_active_spark_context() + return _invoke_function("call_function", funcName, _to_seq(sc, cols, _to_java_column)) + + +@try_remote_functions +def unwrap_udt(col: "ColumnOrName") -> Column: + """ + Unwrap UDT data type column into its underlying type. + + .. versionadded:: 3.4.0 + + Notes + ----- + Supports Spark Connect. + """ + return _invoke_function("unwrap_udt", _to_java_column(col)) + + +@try_remote_functions +def hll_sketch_agg(col: "ColumnOrName", lgConfigK: Optional[Union[int, Column]] = None) -> Column: + """ + Aggregate function: returns the updatable binary representation of the Datasketches + HllSketch configured with lgConfigK arg. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str or int + lgConfigK : int, optional + The log-base-2 of K, where K is the number of buckets or slots for the HllSketch + + Returns + ------- + :class:`~pyspark.sql.Column` + The binary representation of the HllSketch. + + Examples + -------- + >>> df = spark.createDataFrame([1,2,2,3], "INT") + >>> df1 = df.agg(hll_sketch_estimate(hll_sketch_agg("value")).alias("distinct_cnt")) + >>> df1.show() + +------------+ + |distinct_cnt| + +------------+ + | 3| + +------------+ + >>> df2 = df.agg(hll_sketch_estimate( + ... hll_sketch_agg("value", lit(12)) + ... ).alias("distinct_cnt")) + >>> df2.show() + +------------+ + |distinct_cnt| + +------------+ + | 3| + +------------+ + >>> df3 = df.agg(hll_sketch_estimate( + ... hll_sketch_agg(col("value"), lit(12))).alias("distinct_cnt")) + >>> df3.show() + +------------+ + |distinct_cnt| + +------------+ + | 3| + +------------+ + """ + if lgConfigK is None: + return _invoke_function_over_columns("hll_sketch_agg", col) + else: + _lgConfigK = lit(lgConfigK) if isinstance(lgConfigK, int) else lgConfigK + return _invoke_function_over_columns("hll_sketch_agg", col, _lgConfigK) + + +@try_remote_functions +def hll_union_agg( + col: "ColumnOrName", allowDifferentLgConfigK: Optional[Union[bool, Column]] = None +) -> Column: + """ + Aggregate function: returns the updatable binary representation of the Datasketches + HllSketch, generated by merging previously created Datasketches HllSketch instances + via a Datasketches Union instance. Throws an exception if sketches have different + lgConfigK values and allowDifferentLgConfigK is unset or set to false. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str or bool + allowDifferentLgConfigK : bool, optional + Allow sketches with different lgConfigK values to be merged (defaults to false). + + Returns + ------- + :class:`~pyspark.sql.Column` + The binary representation of the merged HllSketch. + + Examples + -------- + >>> df1 = spark.createDataFrame([1,2,2,3], "INT") + >>> df1 = df1.agg(hll_sketch_agg("value").alias("sketch")) + >>> df2 = spark.createDataFrame([4,5,5,6], "INT") + >>> df2 = df2.agg(hll_sketch_agg("value").alias("sketch")) + >>> df3 = df1.union(df2).agg(hll_sketch_estimate( + ... hll_union_agg("sketch") + ... ).alias("distinct_cnt")) + >>> df3.drop("sketch").show() + +------------+ + |distinct_cnt| + +------------+ + | 6| + +------------+ + >>> df4 = df1.union(df2).agg(hll_sketch_estimate( + ... hll_union_agg("sketch", lit(False)) + ... ).alias("distinct_cnt")) + >>> df4.drop("sketch").show() + +------------+ + |distinct_cnt| + +------------+ + | 6| + +------------+ + >>> df5 = df1.union(df2).agg(hll_sketch_estimate( + ... hll_union_agg(col("sketch"), lit(False)) + ... ).alias("distinct_cnt")) + >>> df5.drop("sketch").show() + +------------+ + |distinct_cnt| + +------------+ + | 6| + +------------+ + """ + if allowDifferentLgConfigK is None: + return _invoke_function_over_columns("hll_union_agg", col) + else: + _allowDifferentLgConfigK = ( + lit(allowDifferentLgConfigK) + if isinstance(allowDifferentLgConfigK, bool) + else allowDifferentLgConfigK + ) + return _invoke_function_over_columns("hll_union_agg", col, _allowDifferentLgConfigK) + + +@try_remote_functions +def hll_sketch_estimate(col: "ColumnOrName") -> Column: + """ + Returns the estimated number of unique values given the binary representation + of a Datasketches HllSketch. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Returns + ------- + :class:`~pyspark.sql.Column` + The estimated number of unique values for the HllSketch. + + Examples + -------- + >>> df = spark.createDataFrame([1,2,2,3], "INT") + >>> df = df.agg(hll_sketch_estimate(hll_sketch_agg("value")).alias("distinct_cnt")) + >>> df.show() + +------------+ + |distinct_cnt| + +------------+ + | 3| + +------------+ + """ + return _invoke_function("hll_sketch_estimate", _to_java_column(col)) + + +@try_remote_functions +def hll_union( + col1: "ColumnOrName", col2: "ColumnOrName", allowDifferentLgConfigK: Optional[bool] = None +) -> Column: + """ + Merges two binary representations of Datasketches HllSketch objects, using a + Datasketches Union object. Throws an exception if sketches have different + lgConfigK values and allowDifferentLgConfigK is unset or set to false. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + allowDifferentLgConfigK : bool, optional + Allow sketches with different lgConfigK values to be merged (defaults to false). + + Returns + ------- + :class:`~pyspark.sql.Column` + The binary representation of the merged HllSketch. + + Examples + -------- + >>> df = spark.createDataFrame([(1,4),(2,5),(2,5),(3,6)], "struct") + >>> df = df.agg(hll_sketch_agg("v1").alias("sketch1"), hll_sketch_agg("v2").alias("sketch2")) + >>> df = df.withColumn("distinct_cnt", hll_sketch_estimate(hll_union("sketch1", "sketch2"))) + >>> df.drop("sketch1", "sketch2").show() + +------------+ + |distinct_cnt| + +------------+ + | 6| + +------------+ + """ + if allowDifferentLgConfigK is not None: + return _invoke_function( + "hll_union", _to_java_column(col1), _to_java_column(col2), allowDifferentLgConfigK + ) + else: + return _invoke_function("hll_union", _to_java_column(col1), _to_java_column(col2)) + + +# ---------------------- Predicates functions ------------------------------ + + +@try_remote_functions +def ifnull(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Returns `col2` if `col1` is null, or `col1` otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([(None,), (1,)], ["e"]) + >>> df.select(sf.ifnull(df.e, sf.lit(8))).show() + +------------+ + |ifnull(e, 8)| + +------------+ + | 8| + | 1| + +------------+ + """ + return _invoke_function_over_columns("ifnull", col1, col2) + + +@try_remote_functions +def isnotnull(col: "ColumnOrName") -> Column: + """ + Returns true if `col` is not null, or false otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(None,), (1,)], ["e"]) + >>> df.select(isnotnull(df.e).alias('r')).collect() + [Row(r=False), Row(r=True)] + """ + return _invoke_function_over_columns("isnotnull", col) + + +@try_remote_functions +def equal_null(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Returns same result as the EQUAL(=) operator for non-null operands, + but returns true if both are null, false if one of the them is null. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(None, None,), (1, 9,)], ["a", "b"]) + >>> df.select(equal_null(df.a, df.b).alias('r')).collect() + [Row(r=True), Row(r=False)] + """ + return _invoke_function_over_columns("equal_null", col1, col2) + + +@try_remote_functions +def nullif(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Returns null if `col1` equals to `col2`, or `col1` otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(None, None,), (1, 9,)], ["a", "b"]) + >>> df.select(nullif(df.a, df.b).alias('r')).collect() + [Row(r=None), Row(r=1)] + """ + return _invoke_function_over_columns("nullif", col1, col2) + + +@try_remote_functions +def nvl(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: + """ + Returns `col2` if `col1` is null, or `col1` otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(None, 8,), (1, 9,)], ["a", "b"]) + >>> df.select(nvl(df.a, df.b).alias('r')).collect() + [Row(r=8), Row(r=1)] + """ + return _invoke_function_over_columns("nvl", col1, col2) + + +@try_remote_functions +def nvl2(col1: "ColumnOrName", col2: "ColumnOrName", col3: "ColumnOrName") -> Column: + """ + Returns `col2` if `col1` is not null, or `col3` otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col1 : :class:`~pyspark.sql.Column` or str + col2 : :class:`~pyspark.sql.Column` or str + col3 : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(None, 8, 6,), (1, 9, 9,)], ["a", "b", "c"]) + >>> df.select(nvl2(df.a, df.b, df.c).alias('r')).collect() + [Row(r=6), Row(r=9)] + """ + return _invoke_function_over_columns("nvl2", col1, col2, col3) + + +@try_remote_functions +def aes_encrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + iv: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + """ + Returns an encrypted value of `input` using AES in given `mode` with the specified `padding`. + Key lengths of 16, 24 and 32 bits are supported. Supported combinations of (`mode`, + `padding`) are ('ECB', 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS'). Optional initialization + vectors (IVs) are only supported for CBC and GCM modes. These must be 16 bytes for CBC and 12 + bytes for GCM. If not provided, a random vector will be generated and prepended to the + output. Optional additional authenticated data (AAD) is only supported for GCM. If provided + for encryption, the identical AAD value must be provided for decryption. The default mode is + GCM. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + input : :class:`~pyspark.sql.Column` or str + The binary value to encrypt. + key : :class:`~pyspark.sql.Column` or str + The passphrase to use to encrypt the data. + mode : :class:`~pyspark.sql.Column` or str, optional + Specifies which block cipher mode should be used to encrypt messages. Valid modes: ECB, + GCM, CBC. + padding : :class:`~pyspark.sql.Column` or str, optional + Specifies how to pad messages whose length is not a multiple of the block size. Valid + values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB, NONE for GCM and PKCS + for CBC. + iv : :class:`~pyspark.sql.Column` or str, optional + Optional initialization vector. Only supported for CBC and GCM modes. Valid values: None or + "". 16-byte array for CBC mode. 12-byte array for GCM mode. + aad : :class:`~pyspark.sql.Column` or str, optional + Optional additional authenticated data. Only supported for GCM mode. This can be any + free-form input and must be provided for both encryption and decryption. + + Examples + -------- + >>> df = spark.createDataFrame([( + ... "Spark", "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT", + ... "000000000000000000000000", "This is an AAD mixed into the input",)], + ... ["input", "key", "mode", "padding", "iv", "aad"] + ... ) + >>> df.select(base64(aes_encrypt( + ... df.input, df.key, df.mode, df.padding, to_binary(df.iv, lit("hex")), df.aad) + ... ).alias('r')).collect() + [Row(r='AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4')] + + >>> df.select(base64(aes_encrypt( + ... df.input, df.key, df.mode, df.padding, to_binary(df.iv, lit("hex"))) + ... ).alias('r')).collect() + [Row(r='AAAAAAAAAAAAAAAAQiYi+sRNYDAOTjdSEcYBFsAWPL1f')] + + >>> df = spark.createDataFrame([( + ... "Spark SQL", "1234567890abcdef", "ECB", "PKCS",)], + ... ["input", "key", "mode", "padding"] + ... ) + >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode, df.padding), + ... df.key, df.mode, df.padding).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark SQL'))] + + >>> df = spark.createDataFrame([( + ... "Spark SQL", "0000111122223333", "ECB",)], + ... ["input", "key", "mode"] + ... ) + >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode), + ... df.key, df.mode).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark SQL'))] + + >>> df = spark.createDataFrame([( + ... "Spark SQL", "abcdefghijklmnop",)], + ... ["input", "key"] + ... ) + >>> df.select(aes_decrypt( + ... unbase64(base64(aes_encrypt(df.input, df.key))), df.key + ... ).cast("STRING").alias('r')).collect() + [Row(r='Spark SQL')] + """ + _mode = lit("GCM") if mode is None else mode + _padding = lit("DEFAULT") if padding is None else padding + _iv = lit("") if iv is None else iv + _aad = lit("") if aad is None else aad + return _invoke_function_over_columns("aes_encrypt", input, key, _mode, _padding, _iv, _aad) + + +@try_remote_functions +def aes_decrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + """ + Returns a decrypted value of `input` using AES in `mode` with `padding`. Key lengths of 16, + 24 and 32 bits are supported. Supported combinations of (`mode`, `padding`) are ('ECB', + 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS'). Optional additional authenticated data (AAD) is + only supported for GCM. If provided for encryption, the identical AAD value must be provided + for decryption. The default mode is GCM. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + input : :class:`~pyspark.sql.Column` or str + The binary value to decrypt. + key : :class:`~pyspark.sql.Column` or str + The passphrase to use to decrypt the data. + mode : :class:`~pyspark.sql.Column` or str, optional + Specifies which block cipher mode should be used to decrypt messages. Valid modes: ECB, + GCM, CBC. + padding : :class:`~pyspark.sql.Column` or str, optional + Specifies how to pad messages whose length is not a multiple of the block size. Valid + values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB, NONE for GCM and PKCS + for CBC. + aad : :class:`~pyspark.sql.Column` or str, optional + Optional additional authenticated data. Only supported for GCM mode. This can be any + free-form input and must be provided for both encryption and decryption. + + Examples + -------- + >>> df = spark.createDataFrame([( + ... "AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4", + ... "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT", + ... "This is an AAD mixed into the input",)], + ... ["input", "key", "mode", "padding", "aad"] + ... ) + >>> df.select(aes_decrypt( + ... unbase64(df.input), df.key, df.mode, df.padding, df.aad).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df = spark.createDataFrame([( + ... "AAAAAAAAAAAAAAAAAAAAAPSd4mWyMZ5mhvjiAPQJnfg=", + ... "abcdefghijklmnop12345678ABCDEFGH", "CBC", "DEFAULT",)], + ... ["input", "key", "mode", "padding"] + ... ) + >>> df.select(aes_decrypt( + ... unbase64(df.input), df.key, df.mode, df.padding).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df.select(aes_decrypt(unbase64(df.input), df.key, df.mode).alias('r')).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df = spark.createDataFrame([( + ... "83F16B2AA704794132802D248E6BFD4E380078182D1544813898AC97E709B28A94", + ... "0000111122223333",)], + ... ["input", "key"] + ... ) + >>> df.select(aes_decrypt(unhex(df.input), df.key).alias('r')).collect() + [Row(r=bytearray(b'Spark'))] + """ + _mode = lit("GCM") if mode is None else mode + _padding = lit("DEFAULT") if padding is None else padding + _aad = lit("") if aad is None else aad + return _invoke_function_over_columns("aes_decrypt", input, key, _mode, _padding, _aad) + + +@try_remote_functions +def try_aes_decrypt( + input: "ColumnOrName", + key: "ColumnOrName", + mode: Optional["ColumnOrName"] = None, + padding: Optional["ColumnOrName"] = None, + aad: Optional["ColumnOrName"] = None, +) -> Column: + """ + This is a special version of `aes_decrypt` that performs the same operation, + but returns a NULL value instead of raising an error if the decryption cannot be performed. + Returns a decrypted value of `input` using AES in `mode` with `padding`. Key lengths of 16, + 24 and 32 bits are supported. Supported combinations of (`mode`, `padding`) are ('ECB', + 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS'). Optional additional authenticated data (AAD) is + only supported for GCM. If provided for encryption, the identical AAD value must be provided + for decryption. The default mode is GCM. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + input : :class:`~pyspark.sql.Column` or str + The binary value to decrypt. + key : :class:`~pyspark.sql.Column` or str + The passphrase to use to decrypt the data. + mode : :class:`~pyspark.sql.Column` or str, optional + Specifies which block cipher mode should be used to decrypt messages. Valid modes: ECB, + GCM, CBC. + padding : :class:`~pyspark.sql.Column` or str, optional + Specifies how to pad messages whose length is not a multiple of the block size. Valid + values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB, NONE for GCM and PKCS + for CBC. + aad : :class:`~pyspark.sql.Column` or str, optional + Optional additional authenticated data. Only supported for GCM mode. This can be any + free-form input and must be provided for both encryption and decryption. + + Examples + -------- + >>> df = spark.createDataFrame([( + ... "AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4", + ... "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT", + ... "This is an AAD mixed into the input",)], + ... ["input", "key", "mode", "padding", "aad"] + ... ) + >>> df.select(try_aes_decrypt( + ... unbase64(df.input), df.key, df.mode, df.padding, df.aad).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df = spark.createDataFrame([( + ... "AAAAAAAAAAAAAAAAAAAAAPSd4mWyMZ5mhvjiAPQJnfg=", + ... "abcdefghijklmnop12345678ABCDEFGH", "CBC", "DEFAULT",)], + ... ["input", "key", "mode", "padding"] + ... ) + >>> df.select(try_aes_decrypt( + ... unbase64(df.input), df.key, df.mode, df.padding).alias('r') + ... ).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df.select(try_aes_decrypt(unbase64(df.input), df.key, df.mode).alias('r')).collect() + [Row(r=bytearray(b'Spark'))] + + >>> df = spark.createDataFrame([( + ... "83F16B2AA704794132802D248E6BFD4E380078182D1544813898AC97E709B28A94", + ... "0000111122223333",)], + ... ["input", "key"] + ... ) + >>> df.select(try_aes_decrypt(unhex(df.input), df.key).alias('r')).collect() + [Row(r=bytearray(b'Spark'))] + """ + _mode = lit("GCM") if mode is None else mode + _padding = lit("DEFAULT") if padding is None else padding + _aad = lit("") if aad is None else aad + return _invoke_function_over_columns("try_aes_decrypt", input, key, _mode, _padding, _aad) + + +@try_remote_functions +def sha(col: "ColumnOrName") -> Column: + """ + Returns a sha1 hash value as a hex string of the `col`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.sha(sf.lit("Spark"))).show() + +--------------------+ + | sha(Spark)| + +--------------------+ + |85f5955f4b27a9a4c...| + +--------------------+ + """ + return _invoke_function_over_columns("sha", col) + + +@try_remote_functions +def input_file_block_length() -> Column: + """ + Returns the length of the block being read, or -1 if not available. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.read.text("python/test_support/sql/ages_newlines.csv", lineSep=",") + >>> df.select(input_file_block_length().alias('r')).first() + Row(r=87) + """ + return _invoke_function_over_columns("input_file_block_length") + + +@try_remote_functions +def input_file_block_start() -> Column: + """ + Returns the start offset of the block being read, or -1 if not available. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.read.text("python/test_support/sql/ages_newlines.csv", lineSep=",") + >>> df.select(input_file_block_start().alias('r')).first() + Row(r=0) + """ + return _invoke_function_over_columns("input_file_block_start") + + +@try_remote_functions +def reflect(*cols: "ColumnOrName") -> Column: + """ + Calls a method with reflection. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + the first element should be a literal string for the class name, + and the second element should be a literal string for the method name, + and the remaining are input arguments to the Java method. + + Examples + -------- + >>> df = spark.createDataFrame([("a5cf6c42-0c85-418f-af6c-3e4e5b1328f2",)], ["a"]) + >>> df.select( + ... reflect(lit("java.util.UUID"), lit("fromString"), df.a).alias('r') + ... ).collect() + [Row(r='a5cf6c42-0c85-418f-af6c-3e4e5b1328f2')] + """ + return _invoke_function_over_seq_of_columns("reflect", cols) + + +@try_remote_functions +def java_method(*cols: "ColumnOrName") -> Column: + """ + Calls a method with reflection. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + the first element should be a literal string for the class name, + and the second element should be a literal string for the method name, + and the remaining are input arguments to the Java method. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select( + ... sf.java_method( + ... sf.lit("java.util.UUID"), + ... sf.lit("fromString"), + ... sf.lit("a5cf6c42-0c85-418f-af6c-3e4e5b1328f2") + ... ) + ... ).show(truncate=False) + +-----------------------------------------------------------------------------+ + |java_method(java.util.UUID, fromString, a5cf6c42-0c85-418f-af6c-3e4e5b1328f2)| + +-----------------------------------------------------------------------------+ + |a5cf6c42-0c85-418f-af6c-3e4e5b1328f2 | + +-----------------------------------------------------------------------------+ + """ + return _invoke_function_over_seq_of_columns("java_method", cols) + + +@try_remote_functions +def version() -> Column: + """ + Returns the Spark version. The string contains 2 fields, the first being a release version + and the second being a git revision. + + .. versionadded:: 3.5.0 + + Examples + -------- + >>> df = spark.range(1) + >>> df.select(version()).show(truncate=False) # doctest: +SKIP + +----------------------------------------------+ + |version() | + +----------------------------------------------+ + |3.5.0 cafbea5b13623276517a9d716f75745eff91f616| + +----------------------------------------------+ + """ + return _invoke_function_over_columns("version") + + +@try_remote_functions +def typeof(col: "ColumnOrName") -> Column: + """ + Return DDL-formatted type string for the data type of the input. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + + Examples + -------- + >>> df = spark.createDataFrame([(1,)], ["a"]) + >>> df.select(typeof(df.a).alias('r')).collect() + [Row(r='bigint')] + """ + return _invoke_function_over_columns("typeof", col) + + +@try_remote_functions +def stack(*cols: "ColumnOrName") -> Column: + """ + Separates `col1`, ..., `colk` into `n` rows. Uses column names col0, col1, etc. by default + unless specified otherwise. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + cols : :class:`~pyspark.sql.Column` or str + the first element should be a literal int for the number of rows to be separated, + and the remaining are input elements to be separated. + + Examples + -------- + >>> df = spark.createDataFrame([(1, 2, 3)], ["a", "b", "c"]) + >>> df.select(stack(lit(2), df.a, df.b, df.c)).show(truncate=False) + +----+----+ + |col0|col1| + +----+----+ + |1 |2 | + |3 |NULL| + +----+----+ + """ + return _invoke_function_over_seq_of_columns("stack", cols) + + +@try_remote_functions +def bitmap_bit_position(col: "ColumnOrName") -> Column: + """ + Returns the bit position for the given input column. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + The input column. + + Examples + -------- + >>> df = spark.createDataFrame([(123,)], ["a"]) + >>> df.select(bitmap_bit_position(df.a).alias("r")).collect() + [Row(r=122)] + """ + return _invoke_function_over_columns("bitmap_bit_position", col) + + +@try_remote_functions +def bitmap_bucket_number(col: "ColumnOrName") -> Column: + """ + Returns the bucket number for the given input column. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + The input column. + + Examples + -------- + >>> df = spark.createDataFrame([(123,)], ["a"]) + >>> df.select(bitmap_bucket_number(df.a).alias("r")).collect() + [Row(r=1)] + """ + return _invoke_function_over_columns("bitmap_bucket_number", col) + + +@try_remote_functions +def bitmap_construct_agg(col: "ColumnOrName") -> Column: + """ + Returns a bitmap with the positions of the bits set from all the values from the input column. + The input column will most likely be bitmap_bit_position(). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + The input column will most likely be bitmap_bit_position(). + + Examples + -------- + >>> df = spark.createDataFrame([(1,),(2,),(3,)], ["a"]) + >>> df.select(substring(hex( + ... bitmap_construct_agg(bitmap_bit_position(df.a)) + ... ), 0, 6).alias("r")).collect() + [Row(r='070000')] + """ + return _invoke_function_over_columns("bitmap_construct_agg", col) + + +@try_remote_functions +def bitmap_count(col: "ColumnOrName") -> Column: + """ + Returns the number of set bits in the input bitmap. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + The input bitmap. + + Examples + -------- + >>> df = spark.createDataFrame([("FFFF",)], ["a"]) + >>> df.select(bitmap_count(to_binary(df.a, lit("hex"))).alias('r')).collect() + [Row(r=16)] + """ + return _invoke_function_over_columns("bitmap_count", col) + + +@try_remote_functions +def bitmap_or_agg(col: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise OR of all of the bitmaps from the input column. + The input column should be bitmaps created from bitmap_construct_agg(). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or str + The input column should be bitmaps created from bitmap_construct_agg(). + + Examples + -------- + >>> df = spark.createDataFrame([("10",),("20",),("40",)], ["a"]) + >>> df.select(substring(hex( + ... bitmap_or_agg(to_binary(df.a, lit("hex"))) + ... ), 0, 6).alias("r")).collect() + [Row(r='700000')] + """ + return _invoke_function_over_columns("bitmap_or_agg", col) + + +# ---------------------------- User Defined Function ---------------------------------- + + +@overload +def udf( + f: Callable[..., Any], + returnType: "DataTypeOrString" = StringType(), + *, + useArrow: Optional[bool] = None, +) -> "UserDefinedFunctionLike": + ... + + +@overload +def udf( + f: Optional["DataTypeOrString"] = None, + *, + useArrow: Optional[bool] = None, +) -> Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]: + ... + + +@overload +def udf( + *, + returnType: "DataTypeOrString" = StringType(), + useArrow: Optional[bool] = None, +) -> Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]: + ... + + +@try_remote_functions +def udf( + f: Optional[Union[Callable[..., Any], "DataTypeOrString"]] = None, + returnType: "DataTypeOrString" = StringType(), + *, + useArrow: Optional[bool] = None, +) -> Union["UserDefinedFunctionLike", Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]]: + """Creates a user defined function (UDF). + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + f : function + python function if used as a standalone function + returnType : :class:`pyspark.sql.types.DataType` or str + the return type of the user-defined function. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + useArrow : bool or None + whether to use Arrow to optimize the (de)serialization. When it is None, the + Spark config "spark.sql.execution.pythonUDF.arrow.enabled" takes effect. + + Examples + -------- + >>> from pyspark.sql.types import IntegerType + >>> slen = udf(lambda s: len(s), IntegerType()) + >>> @udf + ... def to_upper(s): + ... if s is not None: + ... return s.upper() + ... + >>> @udf(returnType=IntegerType()) + ... def add_one(x): + ... if x is not None: + ... return x + 1 + ... + >>> df = spark.createDataFrame([(1, "John Doe", 21)], ("id", "name", "age")) + >>> df.select(slen("name").alias("slen(name)"), to_upper("name"), add_one("age")).show() + +----------+--------------+------------+ + |slen(name)|to_upper(name)|add_one(age)| + +----------+--------------+------------+ + | 8| JOHN DOE| 22| + +----------+--------------+------------+ + + Notes + ----- + The user-defined functions are considered deterministic by default. Due to + optimization, duplicate invocations may be eliminated or the function may even be invoked + more times than it is present in the query. If your function is not deterministic, call + `asNondeterministic` on the user defined function. E.g.: + + >>> from pyspark.sql.types import IntegerType + >>> import random + >>> random_udf = udf(lambda: int(random.random() * 100), IntegerType()).asNondeterministic() + + The user-defined functions do not support conditional expressions or short circuiting + in boolean expressions and it ends up with being executed all internally. If the functions + can fail on special rows, the workaround is to incorporate the condition into the functions. + + The user-defined functions do not take keyword arguments on the calling side. + """ + + # The following table shows most of Python data and SQL type conversions in normal UDFs that + # are not yet visible to the user. Some of behaviors are buggy and might be changed in the near + # future. The table might have to be eventually documented externally. + # Please see SPARK-28131's PR to see the codes in order to generate the table below. + # + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+------------------+----------------------+ # noqa + # |SQL Type \ Python Value(Type)|None(NoneType)|True(bool)|1(int)| a(str)| 1970-01-01(date)|1970-01-01 00:00:00(datetime)|1.0(float)|array('i', [1])(array)|[1](list)| (1,)(tuple)|bytearray(b'ABC')(bytearray)| 1(Decimal)|{'a': 1}(dict)|Row(kwargs=1)(Row)|Row(namedtuple=1)(Row)| # noqa + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+------------------+----------------------+ # noqa + # | boolean| None| True| None| None| None| None| None| None| None| None| None| None| None| X| X| # noqa + # | tinyint| None| None| 1| None| None| None| None| None| None| None| None| None| None| X| X| # noqa + # | smallint| None| None| 1| None| None| None| None| None| None| None| None| None| None| X| X| # noqa + # | int| None| None| 1| None| None| None| None| None| None| None| None| None| None| X| X| # noqa + # | bigint| None| None| 1| None| None| None| None| None| None| None| None| None| None| X| X| # noqa + # | string| None| 'true'| '1'| 'a'|'java.util.Gregor...| 'java.util.Gregor...| '1.0'| '[I@66cbb73a'| '[1]'|'[Ljava.lang.Obje...| '[B@5a51eb1a'| '1'| '{a=1}'| X| X| # noqa + # | date| None| X| X| X|datetime.date(197...| datetime.date(197...| X| X| X| X| X| X| X| X| X| # noqa + # | timestamp| None| X| X| X| X| datetime.datetime...| X| X| X| X| X| X| X| X| X| # noqa + # | float| None| None| None| None| None| None| 1.0| None| None| None| None| None| None| X| X| # noqa + # | double| None| None| None| None| None| None| 1.0| None| None| None| None| None| None| X| X| # noqa + # | array| None| None| None| None| None| None| None| [1]| [1]| [1]| [65, 66, 67]| None| None| X| X| # noqa + # | binary| None| None| None|bytearray(b'a')| None| None| None| None| None| None| bytearray(b'ABC')| None| None| X| X| # noqa + # | decimal(10,0)| None| None| None| None| None| None| None| None| None| None| None|Decimal('1')| None| X| X| # noqa + # | map| None| None| None| None| None| None| None| None| None| None| None| None| {'a': 1}| X| X| # noqa + # | struct<_1:int>| None| X| X| X| X| X| X| X|Row(_1=1)| Row(_1=1)| X| X| Row(_1=None)| Row(_1=1)| Row(_1=1)| # noqa + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+------------------+----------------------+ # noqa + # + # Note: DDL formatted string is used for 'SQL Type' for simplicity. This string can be + # used in `returnType`. + # Note: The values inside of the table are generated by `repr`. + # Note: 'X' means it throws an exception during the conversion. + + # decorator @udf, @udf(), @udf(dataType()) + if f is None or isinstance(f, (str, DataType)): + # If DataType has been passed as a positional argument + # for decorator use it as a returnType + return_type = f or returnType + return functools.partial( + _create_py_udf, + returnType=return_type, + useArrow=useArrow, + ) + else: + return _create_py_udf(f=f, returnType=returnType, useArrow=useArrow) + + +@try_remote_functions +def udtf( + cls: Optional[Type] = None, + *, + returnType: Union[StructType, str], + useArrow: Optional[bool] = None, +) -> Union["UserDefinedTableFunction", Callable[[Type], "UserDefinedTableFunction"]]: + """Creates a user defined table function (UDTF). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + cls : class + the Python user-defined table function handler class. + returnType : :class:`pyspark.sql.types.StructType` or str + the return type of the user-defined table function. The value can be either a + :class:`pyspark.sql.types.StructType` object or a DDL-formatted struct type string. + useArrow : bool or None, optional + whether to use Arrow to optimize the (de)serializations. When it's set to None, the + Spark config "spark.sql.execution.pythonUDTF.arrow.enabled" is used. + + Examples + -------- + Implement the UDTF class and create a UDTF: + + >>> class TestUDTF: + ... def eval(self, *args: Any): + ... yield "hello", "world" + ... + >>> from pyspark.sql.functions import udtf + >>> test_udtf = udtf(TestUDTF, returnType="c1: string, c2: string") + >>> test_udtf().show() + +-----+-----+ + | c1| c2| + +-----+-----+ + |hello|world| + +-----+-----+ + + UDTF can also be created using the decorator syntax: + + >>> @udtf(returnType="c1: int, c2: int") + ... class PlusOne: + ... def eval(self, x: int): + ... yield x, x + 1 + ... + >>> from pyspark.sql.functions import lit + >>> PlusOne(lit(1)).show() + +---+---+ + | c1| c2| + +---+---+ + | 1| 2| + +---+---+ + + Arrow optimization can be explicitly enabled when creating UDTFs: + + >>> @udtf(returnType="c1: int, c2: int", useArrow=True) + ... class ArrowPlusOne: + ... def eval(self, x: int): + ... yield x, x + 1 + ... + >>> ArrowPlusOne(lit(1)).show() + +---+---+ + | c1| c2| + +---+---+ + | 1| 2| + +---+---+ + + Notes + ----- + User-defined table functions (UDTFs) are considered non-deterministic by default. + Use `asDeterministic()` to mark a function as deterministic. E.g.: + + >>> class PlusOne: + ... def eval(self, a: int): + ... yield a + 1, + >>> plus_one = udtf(PlusOne, returnType="r: int").asDeterministic() + + Use "yield" to produce one row for the UDTF result relation as many times + as needed. In the context of a lateral join, each such result row will be + associated with the most recent input row consumed from the "eval" method. + + User-defined table functions are considered opaque to the optimizer by default. + As a result, operations like filters from WHERE clauses or limits from + LIMIT/OFFSET clauses that appear after the UDTF call will execute on the + UDTF's result relation. By the same token, any relations forwarded as input + to UDTFs will plan as full table scans in the absence of any explicit such + filtering or other logic explicitly written in a table subquery surrounding the + provided input relation. + + User-defined table functions do not accept keyword arguments on the calling side. + """ + if cls is None: + return functools.partial(_create_py_udtf, returnType=returnType, useArrow=useArrow) + else: + return _create_py_udtf(cls=cls, returnType=returnType, useArrow=useArrow) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.functions + + globs = pyspark.sql.functions.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.functions tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.functions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/group.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/group.py new file mode 100644 index 0000000..730f473 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/group.py @@ -0,0 +1,547 @@ +# +# 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 sys + +from typing import Callable, List, Optional, TYPE_CHECKING, overload, Dict, Union, cast, Tuple + +from py4j.java_gateway import JavaObject + +from pyspark.sql.column import Column, _to_seq +from pyspark.sql.session import SparkSession +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.pandas.group_ops import PandasGroupedOpsMixin + +if TYPE_CHECKING: + from pyspark.sql._typing import LiteralType + +__all__ = ["GroupedData"] + + +def dfapi(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]: + def _api(self: "GroupedData") -> DataFrame: + name = f.__name__ + jdf = getattr(self._jgd, name)() + return DataFrame(jdf, self.session) + + _api.__name__ = f.__name__ + _api.__doc__ = f.__doc__ + return _api + + +def df_varargs_api(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]: + def _api(self: "GroupedData", *cols: str) -> DataFrame: + name = f.__name__ + jdf = getattr(self._jgd, name)(_to_seq(self.session._sc, cols)) + return DataFrame(jdf, self.session) + + _api.__name__ = f.__name__ + _api.__doc__ = f.__doc__ + return _api + + +class GroupedData(PandasGroupedOpsMixin): + """ + A set of methods for aggregations on a :class:`DataFrame`, + created by :func:`DataFrame.groupBy`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, jgd: JavaObject, df: DataFrame): + self._jgd = jgd + self._df = df + self.session: SparkSession = df.sparkSession + + def __repr__(self) -> str: + index = 26 # index to truncate string from the JVM side + jvm_string = self._jgd.toString() + if jvm_string is not None and len(jvm_string) > index and jvm_string[index] == "[": + return f"GroupedData{jvm_string[index:]}" + else: + return super().__repr__() + + @overload + def agg(self, *exprs: Column) -> DataFrame: + ... + + @overload + def agg(self, __exprs: Dict[str, str]) -> DataFrame: + ... + + def agg(self, *exprs: Union[Column, Dict[str, str]]) -> DataFrame: + """Compute aggregates and returns the result as a :class:`DataFrame`. + + The available aggregate functions can be: + + 1. built-in aggregation functions, such as `avg`, `max`, `min`, `sum`, `count` + + 2. group aggregate pandas UDFs, created with :func:`pyspark.sql.functions.pandas_udf` + + .. note:: There is no partial aggregation with group aggregate UDFs, i.e., + a full shuffle is required. Also, all the data of a group will be loaded into + memory, so the user should be aware of the potential OOM risk if data is skewed + and certain groups are too large to fit in memory. + + .. seealso:: :func:`pyspark.sql.functions.pandas_udf` + + If ``exprs`` is a single :class:`dict` mapping from string to string, then the key + is the column to perform aggregation on, and the value is the aggregate function. + + Alternatively, ``exprs`` can also be a list of aggregate :class:`Column` expressions. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + exprs : dict + a dict mapping from column name (string) to aggregate functions (string), + or a list of :class:`Column`. + + Notes + ----- + Built-in aggregation functions and group aggregate pandas UDFs cannot be mixed + in a single call to this function. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> from pyspark.sql.functions import pandas_udf, PandasUDFType + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"]) + >>> df.show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + | 3|Alice| + | 5| Bob| + | 10| Bob| + +---+-----+ + + Group-by name, and count each group. + + >>> df.groupBy(df.name) + GroupedData[grouping...: [name...], value: [age: bigint, name: string], type: GroupBy] + + >>> df.groupBy(df.name).agg({"*": "count"}).sort("name").show() + +-----+--------+ + | name|count(1)| + +-----+--------+ + |Alice| 2| + | Bob| 2| + +-----+--------+ + + Group-by name, and calculate the minimum age. + + >>> df.groupBy(df.name).agg(sf.min(df.age)).sort("name").show() + +-----+--------+ + | name|min(age)| + +-----+--------+ + |Alice| 2| + | Bob| 5| + +-----+--------+ + + Same as above but uses pandas UDF. + + >>> @pandas_udf('int', PandasUDFType.GROUPED_AGG) # doctest: +SKIP + ... def min_udf(v): + ... return v.min() + ... + >>> df.groupBy(df.name).agg(min_udf(df.age)).sort("name").show() # doctest: +SKIP + +-----+------------+ + | name|min_udf(age)| + +-----+------------+ + |Alice| 2| + | Bob| 5| + +-----+------------+ + """ + assert exprs, "exprs should not be empty" + if len(exprs) == 1 and isinstance(exprs[0], dict): + jdf = self._jgd.agg(exprs[0]) + else: + # Columns + assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column" + exprs = cast(Tuple[Column, ...], exprs) + jdf = self._jgd.agg(exprs[0]._jc, _to_seq(self.session._sc, [c._jc for c in exprs[1:]])) + return DataFrame(jdf, self.session) + + @dfapi + def count(self) -> DataFrame: + """Counts the number of records for each group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame( + ... [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"]) + >>> df.show() + +---+-----+ + |age| name| + +---+-----+ + | 2|Alice| + | 3|Alice| + | 5| Bob| + | 10| Bob| + +---+-----+ + + Group-by name, and count each group. + + >>> df.groupBy(df.name).count().sort("name").show() + +-----+-----+ + | name|count| + +-----+-----+ + |Alice| 2| + | Bob| 2| + +-----+-----+ + """ + + @df_varargs_api + def mean(self, *cols: str) -> DataFrame: + """Computes average values for each numeric columns for each group. + + :func:`mean` is an alias for :func:`avg`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str + column names. Non-numeric columns are ignored. + """ + + @df_varargs_api + def avg(self, *cols: str) -> DataFrame: + """Computes average values for each numeric columns for each group. + + :func:`mean` is an alias for :func:`avg`. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str + column names. Non-numeric columns are ignored. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice", 80), (3, "Alice", 100), + ... (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"]) + >>> df.show() + +---+-----+------+ + |age| name|height| + +---+-----+------+ + | 2|Alice| 80| + | 3|Alice| 100| + | 5| Bob| 120| + | 10| Bob| 140| + +---+-----+------+ + + Group-by name, and calculate the mean of the age in each group. + + >>> df.groupBy("name").avg('age').sort("name").show() + +-----+--------+ + | name|avg(age)| + +-----+--------+ + |Alice| 2.5| + | Bob| 7.5| + +-----+--------+ + + Calculate the mean of the age and height in all data. + + >>> df.groupBy().avg('age', 'height').show() + +--------+-----------+ + |avg(age)|avg(height)| + +--------+-----------+ + | 5.0| 110.0| + +--------+-----------+ + """ + + @df_varargs_api + def max(self, *cols: str) -> DataFrame: + """Computes the max value for each numeric columns for each group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice", 80), (3, "Alice", 100), + ... (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"]) + >>> df.show() + +---+-----+------+ + |age| name|height| + +---+-----+------+ + | 2|Alice| 80| + | 3|Alice| 100| + | 5| Bob| 120| + | 10| Bob| 140| + +---+-----+------+ + + Group-by name, and calculate the max of the age in each group. + + >>> df.groupBy("name").max("age").sort("name").show() + +-----+--------+ + | name|max(age)| + +-----+--------+ + |Alice| 3| + | Bob| 10| + +-----+--------+ + + Calculate the max of the age and height in all data. + + >>> df.groupBy().max("age", "height").show() + +--------+-----------+ + |max(age)|max(height)| + +--------+-----------+ + | 10| 140| + +--------+-----------+ + """ + + @df_varargs_api + def min(self, *cols: str) -> DataFrame: + """Computes the min value for each numeric column for each group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str + column names. Non-numeric columns are ignored. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice", 80), (3, "Alice", 100), + ... (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"]) + >>> df.show() + +---+-----+------+ + |age| name|height| + +---+-----+------+ + | 2|Alice| 80| + | 3|Alice| 100| + | 5| Bob| 120| + | 10| Bob| 140| + +---+-----+------+ + + Group-by name, and calculate the min of the age in each group. + + >>> df.groupBy("name").min("age").sort("name").show() + +-----+--------+ + | name|min(age)| + +-----+--------+ + |Alice| 2| + | Bob| 5| + +-----+--------+ + + Calculate the min of the age and height in all data. + + >>> df.groupBy().min("age", "height").show() + +--------+-----------+ + |min(age)|min(height)| + +--------+-----------+ + | 2| 80| + +--------+-----------+ + """ + + @df_varargs_api + def sum(self, *cols: str) -> DataFrame: + """Computes the sum for each numeric columns for each group. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str + column names. Non-numeric columns are ignored. + + Examples + -------- + >>> df = spark.createDataFrame([ + ... (2, "Alice", 80), (3, "Alice", 100), + ... (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"]) + >>> df.show() + +---+-----+------+ + |age| name|height| + +---+-----+------+ + | 2|Alice| 80| + | 3|Alice| 100| + | 5| Bob| 120| + | 10| Bob| 140| + +---+-----+------+ + + Group-by name, and calculate the sum of the age in each group. + + >>> df.groupBy("name").sum("age").sort("name").show() + +-----+--------+ + | name|sum(age)| + +-----+--------+ + |Alice| 5| + | Bob| 15| + +-----+--------+ + + Calculate the sum of the age and height in all data. + + >>> df.groupBy().sum("age", "height").show() + +--------+-----------+ + |sum(age)|sum(height)| + +--------+-----------+ + | 20| 440| + +--------+-----------+ + """ + + # TODO(SPARK-41746): SparkSession.createDataFrame does not support nested datatypes + def pivot(self, pivot_col: str, values: Optional[List["LiteralType"]] = None) -> "GroupedData": + """ + Pivots a column of the current :class:`DataFrame` and perform the specified aggregation. + There are two versions of the pivot function: one that requires the caller + to specify the list of distinct values to pivot on, and one that does not. + The latter is more concise but less efficient, + because Spark needs to first compute the list of distinct values internally. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + pivot_col : str + Name of the column to pivot. + values : list, optional + List of values that will be translated to columns in the output DataFrame. + + Examples + -------- + >>> from pyspark.sql import Row + >>> df1 = spark.createDataFrame([ + ... Row(course="dotNET", year=2012, earnings=10000), + ... Row(course="Java", year=2012, earnings=20000), + ... Row(course="dotNET", year=2012, earnings=5000), + ... Row(course="dotNET", year=2013, earnings=48000), + ... Row(course="Java", year=2013, earnings=30000), + ... ]) + >>> df1.show() + +------+----+--------+ + |course|year|earnings| + +------+----+--------+ + |dotNET|2012| 10000| + | Java|2012| 20000| + |dotNET|2012| 5000| + |dotNET|2013| 48000| + | Java|2013| 30000| + +------+----+--------+ + >>> df2 = spark.createDataFrame([ + ... Row(training="expert", sales=Row(course="dotNET", year=2012, earnings=10000)), + ... Row(training="junior", sales=Row(course="Java", year=2012, earnings=20000)), + ... Row(training="expert", sales=Row(course="dotNET", year=2012, earnings=5000)), + ... Row(training="junior", sales=Row(course="dotNET", year=2013, earnings=48000)), + ... Row(training="expert", sales=Row(course="Java", year=2013, earnings=30000)), + ... ]) # doctest: +SKIP + >>> df2.show() # doctest: +SKIP + +--------+--------------------+ + |training| sales| + +--------+--------------------+ + | expert|{dotNET, 2012, 10...| + | junior| {Java, 2012, 20000}| + | expert|{dotNET, 2012, 5000}| + | junior|{dotNET, 2013, 48...| + | expert| {Java, 2013, 30000}| + +--------+--------------------+ + + Compute the sum of earnings for each year by course with each course as a separate column + + >>> df1.groupBy("year").pivot("course", ["dotNET", "Java"]).sum("earnings").show() + +----+------+-----+ + |year|dotNET| Java| + +----+------+-----+ + |2012| 15000|20000| + |2013| 48000|30000| + +----+------+-----+ + + Or without specifying column values (less efficient) + + >>> df1.groupBy("year").pivot("course").sum("earnings").show() + +----+-----+------+ + |year| Java|dotNET| + +----+-----+------+ + |2012|20000| 15000| + |2013|30000| 48000| + +----+-----+------+ + >>> df2.groupBy("sales.year").pivot("sales.course").sum("sales.earnings").show() + ... # doctest: +SKIP + +----+-----+------+ + |year| Java|dotNET| + +----+-----+------+ + |2012|20000| 15000| + |2013|30000| 48000| + +----+-----+------+ + """ + if values is None: + jgd = self._jgd.pivot(pivot_col) + else: + jgd = self._jgd.pivot(pivot_col, values) + return GroupedData(jgd, self._df) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.group + + globs = pyspark.sql.group.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.group tests").getOrCreate() + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.group, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/observation.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/observation.py new file mode 100644 index 0000000..686b036 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/observation.py @@ -0,0 +1,152 @@ +# +# 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 Any, Dict, Optional + +from py4j.java_gateway import JavaObject, JVMView + +from pyspark.sql import column +from pyspark.sql.column import Column +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.utils import try_remote_observation + +__all__ = ["Observation"] + + +class Observation: + """Class to observe (named) metrics on a :class:`DataFrame`. + + Metrics are aggregation expressions, which are applied to the DataFrame while it is being + processed by an action. + + The metrics have the following guarantees: + + - It will compute the defined aggregates (metrics) on all the data that is flowing through + the Dataset during the action. + - It will report the value of the defined aggregate columns as soon as we reach the end of + the action. + + The metrics columns must either contain a literal (e.g. lit(42)), or should contain one or + more aggregate functions (e.g. sum(a) or sum(a + b) + avg(c) - lit(1)). Expressions that + contain references to the input Dataset's columns must always be wrapped in an aggregate + function. + + An Observation instance collects the metrics while the first action is executed. Subsequent + actions do not modify the metrics returned by `Observation.get`. Retrieval of the metric via + `Observation.get` blocks until the first action has finished and metrics become available. + + .. versionadded:: 3.3.0 + + Notes + ----- + This class does not support streaming datasets. + + Examples + -------- + >>> from pyspark.sql.functions import col, count, lit, max + >>> from pyspark.sql import Observation + >>> df = spark.createDataFrame([["Alice", 2], ["Bob", 5]], ["name", "age"]) + >>> observation = Observation("my metrics") + >>> observed_df = df.observe(observation, count(lit(1)).alias("count"), max(col("age"))) + >>> observed_df.count() + 2 + >>> observation.get + {'count': 2, 'max(age)': 5} + """ + + def __init__(self, name: Optional[str] = None) -> None: + """Constructs a named or unnamed Observation instance. + + Parameters + ---------- + name : str, optional + default is a random UUID string. This is the name of the Observation and the metric. + """ + if name is not None: + if not isinstance(name, str): + raise TypeError("name should be a string") + if name == "": + raise ValueError("name should not be empty") + self._name = name + self._jvm: Optional[JVMView] = None + self._jo: Optional[JavaObject] = None + + @try_remote_observation + def _on(self, df: DataFrame, *exprs: Column) -> DataFrame: + """Attaches this observation to the given :class:`DataFrame` to observe aggregations. + + Parameters + ---------- + df : :class:`DataFrame` + the :class:`DataFrame` to be observed + exprs : list of :class:`Column` + column expressions (:class:`Column`). + + Returns + ------- + :class:`DataFrame` + the observed :class:`DataFrame`. + """ + assert self._jo is None, "an Observation can be used with a DataFrame only once" + + self._jvm = df._sc._jvm + assert self._jvm is not None + cls = self._jvm.org.apache.spark.sql.Observation + self._jo = cls(self._name) if self._name is not None else cls() + observed_df = self._jo.on( + df._jdf, exprs[0]._jc, column._to_seq(df._sc, [c._jc for c in exprs[1:]]) + ) + return DataFrame(observed_df, df.sparkSession) + + # Note that decorated property only works with Python 3.9+ which Spark Connect requires. + @property + @try_remote_observation + def get(self) -> Dict[str, Any]: + """Get the observed metrics. + + Waits until the observed dataset finishes its first action. Only the result of the + first action is available. Subsequent actions do not modify the result. + + Returns + ------- + dict + the observed metrics + """ + assert self._jo is not None, "call DataFrame.observe" + jmap = self._jo.getAsJava() + # return a pure Python dict, not jmap which is a py4j JavaMap + return {k: v for k, v in jmap.items()} + + +def _test() -> None: + import doctest + import sys + from pyspark.context import SparkContext + from pyspark.sql import SparkSession + import pyspark.sql.observation + + globs = pyspark.sql.observation.__dict__.copy() + sc = SparkContext("local[4]", "PythonTest") + globs["spark"] = SparkSession(sc) + + (failure_count, test_count) = doctest.testmod(pyspark.sql.observation, globs=globs) + sc.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/__init__.py new file mode 100644 index 0000000..32a88e9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/__init__.py @@ -0,0 +1,21 @@ +# +# 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. +# +""" +This package includes the internal APIs for PySpark about interoperability +between pandas, PySpark and PyArrow. This package should not be directly +imported and used. +""" diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/__init__.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/__init__.pyi new file mode 100644 index 0000000..6927972 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/__init__.pyi @@ -0,0 +1,344 @@ +# +# 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 ( + Any, + Callable, + Iterable, + NewType, + Tuple, + Type, + TypeVar, + Union, +) +from typing_extensions import Protocol, Literal +from types import FunctionType + +from pyspark.sql._typing import LiteralType +from pyspark.sql.streaming.state import GroupState +from pandas.core.frame import DataFrame as PandasDataFrame +from pandas.core.series import Series as PandasSeries +from numpy import ndarray as NDArray + +import pyarrow + +ArrayLike = NDArray +DataFrameLike = PandasDataFrame +SeriesLike = PandasSeries +DataFrameOrSeriesLike = Union[DataFrameLike, SeriesLike] +DataFrameOrSeriesLike_ = TypeVar("DataFrameOrSeriesLike_", bound=DataFrameOrSeriesLike) + +# UDF annotations +PandasScalarUDFType = Literal[200] +PandasGroupedMapUDFType = Literal[201] +PandasGroupedAggUDFType = Literal[202] +PandasWindowAggUDFType = Literal[203] +PandasScalarIterUDFType = Literal[204] +PandasMapIterUDFType = Literal[205] +PandasCogroupedMapUDFType = Literal[206] +ArrowMapIterUDFType = Literal[207] +PandasGroupedMapUDFWithStateType = Literal[208] + +class PandasVariadicScalarToScalarFunction(Protocol): + def __call__(self, *_: DataFrameOrSeriesLike_) -> DataFrameOrSeriesLike_: ... + +PandasScalarToScalarFunction = Union[ + PandasVariadicScalarToScalarFunction, + Callable[[DataFrameOrSeriesLike_], DataFrameOrSeriesLike_], + Callable[[DataFrameOrSeriesLike_, DataFrameOrSeriesLike_], DataFrameOrSeriesLike_], + Callable[ + [DataFrameOrSeriesLike_, DataFrameOrSeriesLike_, DataFrameOrSeriesLike_], + DataFrameOrSeriesLike_, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + SeriesLike, + ], +] + +class PandasVariadicScalarToStructFunction(Protocol): + def __call__(self, *_: DataFrameOrSeriesLike_) -> DataFrameLike: ... + +PandasScalarToStructFunction = Union[ + PandasVariadicScalarToStructFunction, + Callable[[DataFrameOrSeriesLike_], DataFrameLike], + Callable[[DataFrameOrSeriesLike_, DataFrameOrSeriesLike_], DataFrameLike], + Callable[ + [DataFrameOrSeriesLike_, DataFrameOrSeriesLike_, DataFrameOrSeriesLike_], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], + Callable[ + [ + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + DataFrameOrSeriesLike_, + ], + DataFrameLike, + ], +] + +PandasScalarIterFunction = Union[ + Callable[[Iterable[DataFrameOrSeriesLike_]], Iterable[SeriesLike]], + Callable[[Tuple[DataFrameOrSeriesLike_, ...]], Iterable[SeriesLike]], +] + +PandasGroupedMapFunction = Union[ + Callable[[DataFrameLike], DataFrameLike], + Callable[[Any, DataFrameLike], DataFrameLike], +] + +PandasGroupedMapFunctionWithState = Callable[ + [Any, Iterable[DataFrameLike], GroupState], Iterable[DataFrameLike] +] + +class PandasVariadicGroupedAggFunction(Protocol): + def __call__(self, *_: SeriesLike) -> LiteralType: ... + +PandasGroupedAggFunction = Union[ + Callable[[SeriesLike], LiteralType], + Callable[[SeriesLike, SeriesLike], LiteralType], + Callable[[SeriesLike, SeriesLike, SeriesLike], LiteralType], + Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], + Callable[[SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], LiteralType], + Callable[ + [SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike, SeriesLike], + LiteralType, + ], + Callable[ + [ + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + ], + LiteralType, + ], + Callable[ + [ + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + ], + LiteralType, + ], + Callable[ + [ + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + ], + LiteralType, + ], + Callable[ + [ + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + SeriesLike, + ], + LiteralType, + ], + PandasVariadicGroupedAggFunction, +] + +PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]] + +ArrowMapIterFunction = Callable[[Iterable[pyarrow.RecordBatch]], Iterable[pyarrow.RecordBatch]] + +PandasCogroupedMapFunction = Union[ + Callable[[DataFrameLike, DataFrameLike], DataFrameLike], + Callable[[Any, DataFrameLike, DataFrameLike], DataFrameLike], +] + +GroupedMapPandasUserDefinedFunction = NewType("GroupedMapPandasUserDefinedFunction", FunctionType) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/__init__.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/__init__.pyi new file mode 100644 index 0000000..217e5db --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/__init__.pyi @@ -0,0 +1,17 @@ +# +# 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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/frame.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/frame.pyi new file mode 100644 index 0000000..c1d6274 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/frame.pyi @@ -0,0 +1,20 @@ +# +# 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. + +# This alias is kept for legacy reasons +from pyspark.sql.pandas._typing import DataFrameLike as DataFrameLike diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/series.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/series.pyi new file mode 100644 index 0000000..333235a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/_typing/protocols/series.pyi @@ -0,0 +1,20 @@ +# +# 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. + +# This alias is kept for legacy reasons +from pyspark.sql.pandas._typing import SeriesLike as SeriesLike diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/conversion.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/conversion.py new file mode 100644 index 0000000..8664c4d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/conversion.py @@ -0,0 +1,670 @@ +# +# 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 sys +from typing import ( + Any, + Callable, + List, + Optional, + Union, + no_type_check, + overload, + TYPE_CHECKING, +) +from warnings import warn + +from pyspark.errors.exceptions.captured import unwrap_spark_exception +from pyspark.rdd import _load_from_socket +from pyspark.sql.pandas.serializers import ArrowCollectSerializer +from pyspark.sql.pandas.types import _dedup_names +from pyspark.sql.types import ArrayType, MapType, TimestampType, StructType, DataType, _create_row +from pyspark.sql.utils import is_timestamp_ntz_preferred +from pyspark.traceback_utils import SCCallSiteSync +from pyspark.errors import PySparkTypeError + +if TYPE_CHECKING: + import numpy as np + import pyarrow as pa + from py4j.java_gateway import JavaObject + + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike + from pyspark.sql import DataFrame + + +class PandasConversionMixin: + """ + Mix-in for the conversion from Spark to pandas. Currently, only :class:`DataFrame` + can use this class. + """ + + def toPandas(self) -> "PandasDataFrameLike": + """ + Returns the contents of this :class:`DataFrame` as Pandas ``pandas.DataFrame``. + + This is only available if Pandas is installed and available. + + .. versionadded:: 1.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + This method should only be used if the resulting Pandas ``pandas.DataFrame`` is + expected to be small, as all the data is loaded into the driver's memory. + + Usage with ``spark.sql.execution.arrow.pyspark.enabled=True`` is experimental. + + Examples + -------- + >>> df.toPandas() # doctest: +SKIP + age name + 0 2 Alice + 1 5 Bob + """ + from pyspark.sql.dataframe import DataFrame + + assert isinstance(self, DataFrame) + + from pyspark.sql.pandas.types import _create_converter_to_pandas + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + + import pandas as pd + + jconf = self.sparkSession._jconf + + if jconf.arrowPySparkEnabled(): + use_arrow = True + try: + from pyspark.sql.pandas.types import to_arrow_schema + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + + require_minimum_pyarrow_version() + to_arrow_schema(self.schema) + except Exception as e: + + if jconf.arrowPySparkFallbackEnabled(): + msg = ( + "toPandas attempted Arrow optimization because " + "'spark.sql.execution.arrow.pyspark.enabled' is set to true; however, " + "failed by the reason below:\n %s\n" + "Attempting non-optimization as " + "'spark.sql.execution.arrow.pyspark.fallback.enabled' is set to " + "true." % str(e) + ) + warn(msg) + use_arrow = False + else: + msg = ( + "toPandas attempted Arrow optimization because " + "'spark.sql.execution.arrow.pyspark.enabled' is set to true, but has " + "reached the error below and will not continue because automatic fallback " + "with 'spark.sql.execution.arrow.pyspark.fallback.enabled' has been set to " + "false.\n %s" % str(e) + ) + warn(msg) + raise + + # Try to use Arrow optimization when the schema is supported and the required version + # of PyArrow is found, if 'spark.sql.execution.arrow.pyspark.enabled' is enabled. + if use_arrow: + try: + import pyarrow + + self_destruct = jconf.arrowPySparkSelfDestructEnabled() + batches = self._collect_as_arrow(split_batches=self_destruct) + if len(batches) > 0: + table = pyarrow.Table.from_batches(batches) + # Ensure only the table has a reference to the batches, so that + # self_destruct (if enabled) is effective + del batches + # Pandas DataFrame created from PyArrow uses datetime64[ns] for date type + # values, but we should use datetime.date to match the behavior with when + # Arrow optimization is disabled. + pandas_options = {"date_as_object": True} + if self_destruct: + # Configure PyArrow to use as little memory as possible: + # self_destruct - free columns as they are converted + # split_blocks - create a separate Pandas block for each column + # use_threads - convert one column at a time + pandas_options.update( + { + "self_destruct": True, + "split_blocks": True, + "use_threads": False, + } + ) + # Rename columns to avoid duplicated column names. + pdf = table.rename_columns( + [f"col_{i}" for i in range(table.num_columns)] + ).to_pandas(**pandas_options) + + # Rename back to the original column names. + pdf.columns = self.columns + else: + pdf = pd.DataFrame(columns=self.columns) + + if len(pdf.columns) > 0: + timezone = jconf.sessionLocalTimeZone() + struct_in_pandas = jconf.pandasStructHandlingMode() + + error_on_duplicated_field_names = False + if struct_in_pandas == "legacy": + error_on_duplicated_field_names = True + struct_in_pandas = "dict" + + return pd.concat( + [ + _create_converter_to_pandas( + field.dataType, + field.nullable, + timezone=timezone, + struct_in_pandas=struct_in_pandas, + error_on_duplicated_field_names=error_on_duplicated_field_names, + )(pser) + for (_, pser), field in zip(pdf.items(), self.schema.fields) + ], + axis="columns", + ) + else: + return pdf + except Exception as e: + # We might have to allow fallback here as well but multiple Spark jobs can + # be executed. So, simply fail in this case for now. + msg = ( + "toPandas attempted Arrow optimization because " + "'spark.sql.execution.arrow.pyspark.enabled' is set to true, but has " + "reached the error below and can not continue. Note that " + "'spark.sql.execution.arrow.pyspark.fallback.enabled' does not have an " + "effect on failures in the middle of " + "computation.\n %s" % str(e) + ) + warn(msg) + raise + + # Below is toPandas without Arrow optimization. + rows = self.collect() + if len(rows) > 0: + pdf = pd.DataFrame.from_records( + rows, index=range(len(rows)), columns=self.columns # type: ignore[arg-type] + ) + else: + pdf = pd.DataFrame(columns=self.columns) + + if len(pdf.columns) > 0: + timezone = jconf.sessionLocalTimeZone() + struct_in_pandas = jconf.pandasStructHandlingMode() + + return pd.concat( + [ + _create_converter_to_pandas( + field.dataType, + field.nullable, + timezone=timezone, + struct_in_pandas=( + "row" if struct_in_pandas == "legacy" else struct_in_pandas + ), + error_on_duplicated_field_names=False, + timestamp_utc_localized=False, + )(pser) + for (_, pser), field in zip(pdf.items(), self.schema.fields) + ], + axis="columns", + ) + else: + return pdf + + def _collect_as_arrow(self, split_batches: bool = False) -> List["pa.RecordBatch"]: + """ + Returns all records as a list of ArrowRecordBatches, pyarrow must be installed + and available on driver and worker Python environments. + This is an experimental feature. + + :param split_batches: split batches such that each column is in its own allocation, so + that the selfDestruct optimization is effective; default False. + + .. note:: Experimental. + """ + from pyspark.sql.dataframe import DataFrame + + assert isinstance(self, DataFrame) + + with SCCallSiteSync(self._sc): + ( + port, + auth_secret, + jsocket_auth_server, + ) = self._jdf.collectAsArrowToPython() + + # Collect list of un-ordered batches where last element is a list of correct order indices + try: + batch_stream = _load_from_socket((port, auth_secret), ArrowCollectSerializer()) + if split_batches: + # When spark.sql.execution.arrow.pyspark.selfDestruct.enabled, ensure + # each column in each record batch is contained in its own allocation. + # Otherwise, selfDestruct does nothing; it frees each column as its + # converted, but each column will actually be a list of slices of record + # batches, and so no memory is actually freed until all columns are + # 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) + else: + results = 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] + + # Re-order the batch list using the correct order + return [batches[i] for i in batch_order] + + +class SparkConversionMixin: + """ + Min-in for the conversion from pandas to Spark. Currently, only :class:`SparkSession` + can use this class. + """ + + _jsparkSession: "JavaObject" + + @overload + def createDataFrame( + self, data: "PandasDataFrameLike", samplingRatio: Optional[float] = ... + ) -> "DataFrame": + ... + + @overload + def createDataFrame( + self, + data: "PandasDataFrameLike", + schema: Union[StructType, str], + verifySchema: bool = ..., + ) -> "DataFrame": + ... + + def createDataFrame( # type: ignore[misc] + self, + data: "PandasDataFrameLike", + schema: Optional[Union[StructType, List[str]]] = None, + samplingRatio: Optional[float] = None, + verifySchema: bool = True, + ) -> "DataFrame": + from pyspark.sql import SparkSession + + assert isinstance(self, SparkSession) + + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + + timezone = self._jconf.sessionLocalTimeZone() + + # If no schema supplied by user then get the names of columns only + if schema is None: + schema = [str(x) if not isinstance(x, str) else x for x in data.columns] + + if self._jconf.arrowPySparkEnabled() and len(data) > 0: + try: + return self._create_from_pandas_with_arrow(data, schema, timezone) + except Exception as e: + if self._jconf.arrowPySparkFallbackEnabled(): + msg = ( + "createDataFrame attempted Arrow optimization because " + "'spark.sql.execution.arrow.pyspark.enabled' is set to true; however, " + "failed by the reason below:\n %s\n" + "Attempting non-optimization as " + "'spark.sql.execution.arrow.pyspark.fallback.enabled' is set to " + "true." % str(e) + ) + warn(msg) + else: + msg = ( + "createDataFrame attempted Arrow optimization because " + "'spark.sql.execution.arrow.pyspark.enabled' is set to true, but has " + "reached the error below and will not continue because automatic " + "fallback with 'spark.sql.execution.arrow.pyspark.fallback.enabled' " + "has been set to false.\n %s" % str(e) + ) + warn(msg) + raise + converted_data = self._convert_from_pandas(data, schema, timezone) + return self._create_dataframe(converted_data, schema, samplingRatio, verifySchema) + + def _convert_from_pandas( + self, pdf: "PandasDataFrameLike", schema: Union[StructType, str, List[str]], timezone: str + ) -> List: + """ + Convert a pandas.DataFrame to list of records that can be used to make a DataFrame + + Returns + ------- + list + list of records + """ + import pandas as pd + from pyspark.sql import SparkSession + + assert isinstance(self, SparkSession) + + if timezone is not None: + from pyspark.sql.pandas.types import ( + _check_series_convert_timestamps_tz_local, + _get_local_timezone, + ) + from pandas.core.dtypes.common import is_datetime64tz_dtype, is_timedelta64_dtype + + copied = False + if isinstance(schema, StructType): + + def _create_converter(data_type: DataType) -> Callable[[pd.Series], pd.Series]: + if isinstance(data_type, TimestampType): + + def correct_timestamp(pser: pd.Series) -> pd.Series: + return _check_series_convert_timestamps_tz_local(pser, timezone) + + return correct_timestamp + + def _converter(dt: DataType) -> Optional[Callable[[Any], Any]]: + if isinstance(dt, ArrayType): + element_conv = _converter(dt.elementType) or (lambda x: x) + + def convert_array(value: Any) -> Any: + if value is None: + return None + else: + return [element_conv(v) for v in value] + + return convert_array + + elif isinstance(dt, MapType): + key_conv = _converter(dt.keyType) or (lambda x: x) + value_conv = _converter(dt.valueType) or (lambda x: x) + + def convert_map(value: Any) -> Any: + if value is None: + return None + else: + return {key_conv(k): value_conv(v) for k, v in value.items()} + + return convert_map + + elif isinstance(dt, StructType): + field_names = dt.names + dedup_field_names = _dedup_names(field_names) + field_convs = [ + _converter(f.dataType) or (lambda x: x) for f in dt.fields + ] + + def convert_struct(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + _values = [ + field_convs[i](value.get(name, None)) + for i, name in enumerate(dedup_field_names) + ] + return _create_row(field_names, _values) + else: + _values = [ + field_convs[i](value[i]) for i, name in enumerate(value) + ] + return _create_row(field_names, _values) + + return convert_struct + + elif isinstance(dt, TimestampType): + + def convert_timestamp(value: Any) -> Any: + if value is None: + return None + else: + return ( + pd.Timestamp(value) + .tz_localize(timezone, ambiguous=False) # type: ignore + .tz_convert(_get_local_timezone()) + .tz_localize(None) + .to_pydatetime() + ) + + return convert_timestamp + + else: + return None + + conv = _converter(data_type) + if conv is not None: + return lambda pser: pser.apply(conv) # type: ignore[return-value] + else: + return lambda pser: pser + + if len(pdf.columns) > 0: + pdf = pd.concat( + [ + _create_converter(field.dataType)(pser) + for (_, pser), field in zip(pdf.items(), schema.fields) + ], + axis="columns", + ) + copied = True + else: + should_localize = not is_timestamp_ntz_preferred() + for column, series in pdf.items(): + s = series + if should_localize and is_datetime64tz_dtype(s.dtype) and s.dt.tz is not None: + s = _check_series_convert_timestamps_tz_local(series, timezone) + if s is not series: + if not copied: + # Copy once if the series is modified to prevent the original + # Pandas DataFrame from being updated + pdf = pdf.copy() + copied = True + pdf[column] = s + + for column, series in pdf.items(): + if is_timedelta64_dtype(series): + if not copied: + pdf = pdf.copy() + copied = True + # Explicitly set the timedelta as object so the output of numpy records can + # hold the timedelta instances as are. Otherwise, it converts to the internal + # numeric values. + ser = pdf[column] + pdf[column] = pd.Series( + ser.dt.to_pytimedelta(), index=ser.index, dtype="object", name=ser.name + ) + + # Convert pandas.DataFrame to list of numpy records + np_records = pdf.set_axis( + [f"col_{i}" for i in range(len(pdf.columns))], axis="columns" # type: ignore[arg-type] + ).to_records(index=False) + + # Check if any columns need to be fixed for Spark to infer properly + if len(np_records) > 0: + record_dtype = self._get_numpy_record_dtype(np_records[0]) + if record_dtype is not None: + return [r.astype(record_dtype).tolist() for r in np_records] + + # Convert list of numpy records to python lists + return [r.tolist() for r in np_records] + + def _get_numpy_record_dtype(self, rec: "np.recarray") -> Optional["np.dtype"]: + """ + Used when converting a pandas.DataFrame to Spark using to_records(), this will correct + the dtypes of fields in a record so they can be properly loaded into Spark. + + Parameters + ---------- + rec : numpy.record + a numpy record to check field dtypes + + Returns + ------- + numpy.dtype + corrected dtype for a numpy.record or None if no correction needed + """ + import numpy as np + + cur_dtypes = rec.dtype + col_names = cur_dtypes.names + record_type_list = [] + has_rec_fix = False + for i in range(len(cur_dtypes)): + curr_type = cur_dtypes[i] + # If type is a datetime64 timestamp, convert to microseconds + # NOTE: if dtype is datetime[ns] then np.record.tolist() will output values as longs, + # conversion from [us] or lower will lead to py datetime objects, see SPARK-22417 + if curr_type == np.dtype("datetime64[ns]"): + curr_type = "datetime64[us]" + has_rec_fix = True + record_type_list.append((str(col_names[i]), curr_type)) + return np.dtype(record_type_list) if has_rec_fix else None + + def _create_from_pandas_with_arrow( + self, pdf: "PandasDataFrameLike", schema: Union[StructType, List[str]], timezone: str + ) -> "DataFrame": + """ + Create a DataFrame from a given pandas.DataFrame by slicing it into partitions, converting + to Arrow data, then sending to the JVM to parallelize. If a schema is passed in, the + data types will be used to coerce the data in Pandas to Arrow conversion. + """ + from pyspark.sql import SparkSession + from pyspark.sql.dataframe import DataFrame + + assert isinstance(self, SparkSession) + + from pyspark.sql.pandas.serializers import ArrowStreamPandasSerializer + from pyspark.sql.types import TimestampType + from pyspark.sql.pandas.types import ( + from_arrow_type, + to_arrow_type, + _deduplicate_field_names, + ) + from pyspark.sql.pandas.utils import ( + require_minimum_pandas_version, + require_minimum_pyarrow_version, + ) + + require_minimum_pandas_version() + require_minimum_pyarrow_version() + + from pandas.api.types import ( # type: ignore[attr-defined] + is_datetime64_dtype, + is_datetime64tz_dtype, + ) + import pyarrow as pa + + # Create the Spark schema from list of names passed in with Arrow types + if isinstance(schema, (list, tuple)): + arrow_schema = pa.Schema.from_pandas(pdf, preserve_index=False) + struct = StructType() + prefer_timestamp_ntz = is_timestamp_ntz_preferred() + for name, field in zip(schema, arrow_schema): + struct.add( + name, from_arrow_type(field.type, prefer_timestamp_ntz), nullable=field.nullable + ) + schema = struct + + # Determine arrow types to coerce data when creating batches + if isinstance(schema, StructType): + spark_types = [_deduplicate_field_names(f.dataType) for f in schema.fields] + elif isinstance(schema, DataType): + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW", + message_parameters={"data_type": str(schema)}, + ) + else: + # Any timestamps must be coerced to be compatible with Spark + spark_types = [ + TimestampType() if is_datetime64_dtype(t) or is_datetime64tz_dtype(t) else None + for t in pdf.dtypes + ] + + # Slice the DataFrame to be batched + step = self._jconf.arrowMaxRecordsPerBatch() + pdf_slices = (pdf.iloc[start : start + step] for start in range(0, len(pdf), step)) + + # Create list of Arrow (columns, arrow_type, spark_type) for serializer dump_stream + arrow_data = [ + [ + (c, to_arrow_type(t) if t is not None else None, t) + for (_, c), t in zip(pdf_slice.items(), spark_types) + ] + for pdf_slice in pdf_slices + ] + + jsparkSession = self._jsparkSession + + safecheck = self._jconf.arrowSafeTypeConversion() + ser = ArrowStreamPandasSerializer(timezone, safecheck) + + @no_type_check + def reader_func(temp_filename): + return self._jvm.PythonSQLUtils.readArrowStreamFromFile(temp_filename) + + @no_type_check + def create_iter_server(): + return self._jvm.ArrowIteratorServer() + + # Create Spark DataFrame from Arrow stream file, using one batch per partition + jiter = self._sc._serialize_to_jvm(arrow_data, ser, reader_func, create_iter_server) + assert self._jvm is not None + jdf = self._jvm.PythonSQLUtils.toDataFrame(jiter, schema.json(), jsparkSession) + df = DataFrame(jdf, self) + df._schema = schema + return df + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.pandas.conversion + + globs = pyspark.sql.pandas.conversion.__dict__.copy() + spark = ( + SparkSession.builder.master("local[4]").appName("sql.pandas.conversion tests").getOrCreate() + ) + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.pandas.conversion, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.py new file mode 100644 index 0000000..b7d381f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.py @@ -0,0 +1,480 @@ +# +# 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 functools +import warnings +from inspect import getfullargspec, signature +from typing import get_type_hints + +from pyspark.rdd import PythonEvalType +from pyspark.sql.pandas.typehints import infer_eval_type +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version +from pyspark.sql.types import DataType +from pyspark.sql.udf import _create_udf +from pyspark.sql.utils import is_remote +from pyspark.errors import PySparkTypeError, PySparkValueError + + +class PandasUDFType: + """Pandas UDF Types. See :meth:`pyspark.sql.functions.pandas_udf`.""" + + SCALAR = PythonEvalType.SQL_SCALAR_PANDAS_UDF + + SCALAR_ITER = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + + GROUPED_MAP = PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF + + GROUPED_AGG = PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF + + +def pandas_udf(f=None, returnType=None, functionType=None): + """ + Creates a pandas user defined function (a.k.a. vectorized user defined function). + + Pandas UDFs are user defined functions that are executed by Spark using Arrow to transfer + data and Pandas to work with the data, which allows vectorized operations. A Pandas UDF + is defined using the `pandas_udf` as a decorator or to wrap the function, and no + additional configuration is required. A Pandas UDF behaves as a regular PySpark function + API in general. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + f : function, optional + user-defined function. A python function if used as a standalone function + returnType : :class:`pyspark.sql.types.DataType` or str, optional + the return type of the user-defined function. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + functionType : int, optional + an enum value in :class:`pyspark.sql.functions.PandasUDFType`. + Default: SCALAR. This parameter exists for compatibility. + Using Python type hints is encouraged. + + Examples + -------- + In order to use this API, customarily the below are imported: + + >>> import pandas as pd + >>> from pyspark.sql.functions import pandas_udf + + From Spark 3.0 with Python 3.6+, `Python type hints `_ + detect the function types as below: + + >>> @pandas_udf(IntegerType()) + ... def slen(s: pd.Series) -> pd.Series: + ... return s.str.len() + + Prior to Spark 3.0, the pandas UDF used `functionType` to decide the execution type as below: + + >>> from pyspark.sql.functions import PandasUDFType + >>> from pyspark.sql.types import IntegerType + >>> @pandas_udf(IntegerType(), PandasUDFType.SCALAR) + ... def slen(s): + ... return s.str.len() + + It is preferred to specify type hints for the pandas UDF instead of specifying pandas UDF + type via `functionType` which will be deprecated in the future releases. + + Note that the type hint should use `pandas.Series` in all cases but there is one variant + that `pandas.DataFrame` should be used for its input or output type hint instead when the input + or output column is of :class:`pyspark.sql.types.StructType`. The following example shows + a Pandas UDF which takes long column, string column and struct column, and outputs a struct + column. It requires the function to specify the type hints of `pandas.Series` and + `pandas.DataFrame` as below: + + >>> @pandas_udf("col1 string, col2 long") + >>> def func(s1: pd.Series, s2: pd.Series, s3: pd.DataFrame) -> pd.DataFrame: + ... s3['col2'] = s1 + s2.str.len() + ... return s3 + ... + >>> # Create a Spark DataFrame that has three columns including a struct column. + ... df = spark.createDataFrame( + ... [[1, "a string", ("a nested string",)]], + ... "long_col long, string_col string, struct_col struct") + >>> df.printSchema() + root + |-- long_column: long (nullable = true) + |-- string_column: string (nullable = true) + |-- struct_column: struct (nullable = true) + | |-- col1: string (nullable = true) + >>> df.select(func("long_col", "string_col", "struct_col")).printSchema() + |-- func(long_col, string_col, struct_col): struct (nullable = true) + | |-- col1: string (nullable = true) + | |-- col2: long (nullable = true) + + In the following sections, it describes the combinations of the supported type hints. For + simplicity, `pandas.DataFrame` variant is omitted. + + * Series to Series + `pandas.Series`, ... -> `pandas.Series` + + The function takes one or more `pandas.Series` and outputs one `pandas.Series`. + The output of the function should always be of the same length as the input. + + >>> @pandas_udf("string") + ... def to_upper(s: pd.Series) -> pd.Series: + ... return s.str.upper() + ... + >>> df = spark.createDataFrame([("John Doe",)], ("name",)) + >>> df.select(to_upper("name")).show() + +--------------+ + |to_upper(name)| + +--------------+ + | JOHN DOE| + +--------------+ + + >>> @pandas_udf("first string, last string") + ... def split_expand(s: pd.Series) -> pd.DataFrame: + ... return s.str.split(expand=True) + ... + >>> df = spark.createDataFrame([("John Doe",)], ("name",)) + >>> df.select(split_expand("name")).show() + +------------------+ + |split_expand(name)| + +------------------+ + | [John, Doe]| + +------------------+ + + .. note:: The length of the input is not that of the whole input column, but is the + length of an internal batch used for each call to the function. + + * Iterator of Series to Iterator of Series + `Iterator[pandas.Series]` -> `Iterator[pandas.Series]` + + The function takes an iterator of `pandas.Series` and outputs an iterator of + `pandas.Series`. In this case, the created pandas UDF instance requires one input + column when this is called as a PySpark column. The length of the entire output from + the function should be the same length of the entire input; therefore, it can + prefetch the data from the input iterator as long as the lengths are the same. + + It is also useful when the UDF execution + requires initializing some states although internally it works identically as + Series to Series case. The pseudocode below illustrates the example. + + .. highlight:: python + .. code-block:: python + + @pandas_udf("long") + def calculate(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: + # Do some expensive initialization with a state + state = very_expensive_initialization() + for x in iterator: + # Use that state for whole iterator. + yield calculate_with_state(x, state) + + df.select(calculate("value")).show() + + >>> from typing import Iterator + >>> @pandas_udf("long") + ... def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: + ... for s in iterator: + ... yield s + 1 + ... + >>> df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"])) + >>> df.select(plus_one(df.v)).show() + +-----------+ + |plus_one(v)| + +-----------+ + | 2| + | 3| + | 4| + +-----------+ + + .. note:: The length of each series is the length of a batch internally used. + + * Iterator of Multiple Series to Iterator of Series + `Iterator[Tuple[pandas.Series, ...]]` -> `Iterator[pandas.Series]` + + The function takes an iterator of a tuple of multiple `pandas.Series` and outputs an + iterator of `pandas.Series`. In this case, the created pandas UDF instance requires + input columns as many as the series when this is called as a PySpark column. + Otherwise, it has the same characteristics and restrictions as Iterator of Series + to Iterator of Series case. + + >>> from typing import Iterator, Tuple + >>> from pyspark.sql.functions import struct, col + >>> @pandas_udf("long") + ... def multiply(iterator: Iterator[Tuple[pd.Series, pd.DataFrame]]) -> Iterator[pd.Series]: + ... for s1, df in iterator: + ... yield s1 * df.v + ... + >>> df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"])) + >>> df.withColumn('output', multiply(col("v"), struct(col("v")))).show() + +---+------+ + | v|output| + +---+------+ + | 1| 1| + | 2| 4| + | 3| 9| + +---+------+ + + .. note:: The length of each series is the length of a batch internally used. + + * Series to Scalar + `pandas.Series`, ... -> `Any` + + The function takes `pandas.Series` and returns a scalar value. The `returnType` + should be a primitive data type, and the returned scalar can be either a python primitive + type, e.g., int or float or a numpy data type, e.g., numpy.int64 or numpy.float64. + `Any` should ideally be a specific scalar type accordingly. + + >>> @pandas_udf("double") + ... def mean_udf(v: pd.Series) -> float: + ... return v.mean() + ... + >>> df = spark.createDataFrame( + ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v")) + >>> df.groupby("id").agg(mean_udf(df['v'])).show() + +---+-----------+ + | id|mean_udf(v)| + +---+-----------+ + | 1| 1.5| + | 2| 6.0| + +---+-----------+ + + This UDF can also be used as window functions as below: + + >>> from pyspark.sql import Window + >>> @pandas_udf("double") + ... def mean_udf(v: pd.Series) -> float: + ... return v.mean() + ... + >>> df = spark.createDataFrame( + ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v")) + >>> w = Window.partitionBy('id').orderBy('v').rowsBetween(-1, 0) + >>> df.withColumn('mean_v', mean_udf("v").over(w)).show() + +---+----+------+ + | id| v|mean_v| + +---+----+------+ + | 1| 1.0| 1.0| + | 1| 2.0| 1.5| + | 2| 3.0| 3.0| + | 2| 5.0| 4.0| + | 2|10.0| 7.5| + +---+----+------+ + + .. note:: For performance reasons, the input series to window functions are not copied. + Therefore, mutating the input series is not allowed and will cause incorrect results. + For the same reason, users should also not rely on the index of the input series. + + Notes + ----- + The user-defined functions do not support conditional expressions or short circuiting + in boolean expressions and it ends up with being executed all internally. If the functions + can fail on special rows, the workaround is to incorporate the condition into the functions. + + The user-defined functions do not take keyword arguments on the calling side. + + The data type of returned `pandas.Series` from the user-defined functions should be + matched with defined `returnType` (see :meth:`types.to_arrow_type` and + :meth:`types.from_arrow_type`). When there is mismatch between them, Spark might do + conversion on returned data. The conversion is not guaranteed to be correct and results + should be checked for accuracy by users. + + Currently, + :class:`pyspark.sql.types.ArrayType` of :class:`pyspark.sql.types.TimestampType` and + nested :class:`pyspark.sql.types.StructType` + are currently not supported as output types. + + See Also + -------- + pyspark.sql.GroupedData.agg + pyspark.sql.DataFrame.mapInPandas + pyspark.sql.GroupedData.applyInPandas + pyspark.sql.PandasCogroupedOps.applyInPandas + pyspark.sql.UDFRegistration.register + """ + + # The following table shows most of Pandas data and SQL type conversions in Pandas UDFs that + # are not yet visible to the user. Some of behaviors are buggy and might be changed in the near + # future. The table might have to be eventually documented externally. + # Please see SPARK-28132's PR to see the codes in order to generate the table below. + # + # +-----------------------------+----------------------+------------------+------------------+------------------+--------------------+--------------------+------------------+------------------+------------------+------------------+--------------+--------------+--------------+-----------------------------------+-----------------------------------------------------+-----------------+--------------------+-----------------------------+--------------+-----------------+------------------+---------------+--------------------------------+ # noqa + # |SQL Type \ Pandas Value(Type)|None(object(NoneType))| True(bool)| 1(int8)| 1(int16)| 1(int32)| 1(int64)| 1(uint8)| 1(uint16)| 1(uint32)| 1(uint64)| 1.0(float16)| 1.0(float32)| 1.0(float64)|1970-01-01 00:00:00(datetime64[ns])|1970-01-01 00:00:00-05:00(datetime64[ns, US/Eastern])|a(object(string))| 1(object(Decimal))|[1 2 3](object(array[int32]))| 1.0(float128)|(1+0j)(complex64)|(1+0j)(complex128)| A(category)|1 days 00:00:00(timedelta64[ns])| # noqa + # +-----------------------------+----------------------+------------------+------------------+------------------+--------------------+--------------------+------------------+------------------+------------------+------------------+--------------+--------------+--------------+-----------------------------------+-----------------------------------------------------+-----------------+--------------------+-----------------------------+--------------+-----------------+------------------+---------------+--------------------------------+ # noqa + # | boolean| None| True| True| True| True| True| True| True| True| True| True| True| True| X| X| X| X| X| X| X| X| X| X| # noqa + # | tinyint| None| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| X| X| X| 1| X| X| X| X| X| X| # noqa + # | smallint| None| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| X| X| X| 1| X| X| X| X| X| X| # noqa + # | int| None| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| X| X| X| 1| X| X| X| X| X| X| # noqa + # | bigint| None| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 1| 0| 18000000000000| X| 1| X| X| X| X| X| 86400000000000| # noqa + # | float| None| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| X| X| X| X| X| X| X| X| X| X| # noqa + # | double| None| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| 1.0| X| X| X| X| X| X| X| X| X| X| # noqa + # | date| None| X| X| X|datetime.date(197...| X| X| X| X| X| X| X| X| datetime.date(197...| datetime.date(197...| X|datetime.date(197...| X| X| X| X| X| X| # noqa + # | timestamp| None| X| X| X| X|datetime.datetime...| X| X| X| X| X| X| X| datetime.datetime...| datetime.datetime...| X|datetime.datetime...| X| X| X| X| X| X| # noqa + # | string| None| X| X| X| X| X| X| X| X| X| X| X| X| X| X| 'a'| X| X| X| X| X| 'A'| X| # noqa + # | decimal(10,0)| None| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| Decimal('1')| X| X| X| X| X| X| # noqa + # | array| None| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| [1, 2, 3]| X| X| X| X| X| # noqa + # | map| None| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| # noqa + # | struct<_1:int>| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| X| # noqa + # | binary| None|bytearray(b'\x01')|bytearray(b'\x01')|bytearray(b'\x01')| bytearray(b'\x01')| bytearray(b'\x01')|bytearray(b'\x01')|bytearray(b'\x01')|bytearray(b'\x01')|bytearray(b'\x01')|bytearray(b'')|bytearray(b'')|bytearray(b'')| bytearray(b'')| bytearray(b'')| bytearray(b'a')| X| X|bytearray(b'')| bytearray(b'')| bytearray(b'')|bytearray(b'A')| bytearray(b'')| # noqa + # +-----------------------------+----------------------+------------------+------------------+------------------+--------------------+--------------------+------------------+------------------+------------------+------------------+--------------+--------------+--------------+-----------------------------------+-----------------------------------------------------+-----------------+--------------------+-----------------------------+--------------+-----------------+------------------+---------------+--------------------------------+ # noqa + # + # Note: DDL formatted string is used for 'SQL Type' for simplicity. This string can be + # used in `returnType`. + # Note: The values inside of the table are generated by `repr`. + # Note: Python 3.9.5, Pandas 1.4.0 and PyArrow 6.0.1 are used. + # Note: Timezone is KST. + # Note: 'X' means it throws an exception during the conversion. + require_minimum_pandas_version() + require_minimum_pyarrow_version() + + # decorator @pandas_udf(returnType, functionType) + is_decorator = f is None or isinstance(f, (str, DataType)) + + if is_decorator: + # If DataType has been passed as a positional argument + # for decorator use it as a returnType + return_type = f or returnType + + if functionType is not None: + # @pandas_udf(dataType, functionType=functionType) + # @pandas_udf(returnType=dataType, functionType=functionType) + eval_type = functionType + elif returnType is not None and isinstance(returnType, int): + # @pandas_udf(dataType, functionType) + eval_type = returnType + else: + # @pandas_udf(dataType) or @pandas_udf(returnType=dataType) + eval_type = None + else: + return_type = returnType + + if functionType is not None: + eval_type = functionType + else: + eval_type = None + + if return_type is None: + raise PySparkTypeError( + error_class="CANNOT_BE_NONE", + message_parameters={"arg_name": "returnType"}, + ) + + if eval_type not in [ + PythonEvalType.SQL_SCALAR_PANDAS_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_ARROW_ITER_UDF, + PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + None, + ]: # None means it should infer the type from type hints. + + raise PySparkTypeError( + error_class="INVALID_PANDAS_UDF_TYPE", + message_parameters={ + "arg_name": "functionType", + "arg_type": str(eval_type), + }, + ) + + if is_decorator: + return functools.partial(_create_pandas_udf, returnType=return_type, evalType=eval_type) + else: + return _create_pandas_udf(f=f, returnType=return_type, evalType=eval_type) + + +def _create_pandas_udf(f, returnType, evalType): + argspec = getfullargspec(f) + + # pandas UDF by type hints. + if evalType in [ + PythonEvalType.SQL_SCALAR_PANDAS_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + ]: + warnings.warn( + "In Python 3.6+ and Spark 3.0+, it is preferred to specify type hints for " + "pandas UDF instead of specifying pandas UDF type which will be deprecated " + "in the future releases. See SPARK-28264 for more details.", + UserWarning, + ) + elif evalType in [ + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_ARROW_ITER_UDF, + PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + PythonEvalType.SQL_ARROW_BATCHED_UDF, + ]: + # In case of 'SQL_GROUPED_MAP_PANDAS_UDF', deprecation warning is being triggered + # at `apply` instead. + # In case of 'SQL_MAP_PANDAS_ITER_UDF', 'SQL_MAP_ARROW_ITER_UDF' and + # 'SQL_COGROUPED_MAP_PANDAS_UDF', the evaluation type will always be set. + # In case of 'SQL_ARROW_BATCHED_UDF', no deprecation warning is required since it is not + # exposed to users. + pass + elif len(argspec.annotations) > 0: + try: + type_hints = get_type_hints(f) + except NameError: + type_hints = {} + evalType = infer_eval_type(signature(f), type_hints) + assert evalType is not None + + if evalType is None: + # Set default is scalar UDF. + evalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF + + if ( + ( + evalType == PythonEvalType.SQL_SCALAR_PANDAS_UDF + or evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + ) + and len(argspec.args) == 0 + and argspec.varargs is None + ): + raise PySparkValueError( + error_class="INVALID_PANDAS_UDF", + message_parameters={ + "detail": "0-arg pandas_udfs are not supported. " + "Instead, create a 1-arg pandas_udf and ignore the arg in your function.", + }, + ) + + if evalType == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF and len(argspec.args) not in (1, 2): + raise PySparkValueError( + error_class="INVALID_PANDAS_UDF", + message_parameters={ + "detail": "pandas_udf with function type GROUPED_MAP or the function in " + "groupby.applyInPandas must take either one argument (data) or " + "two arguments (key, data).", + }, + ) + + if evalType == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF and len(argspec.args) not in (2, 3): + raise PySparkValueError( + error_class="INVALID_PANDAS_UDF", + message_parameters={ + "detail": "the function in cogroup.applyInPandas must take either two arguments " + "(left, right) or three arguments (key, left, right).", + }, + ) + + if is_remote(): + from pyspark.sql.connect.udf import _create_udf as _create_connect_udf + + return _create_connect_udf(f, returnType, evalType) + else: + return _create_udf(f, returnType, evalType) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.pyi b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.pyi new file mode 100644 index 0000000..1af6f86 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/functions.pyi @@ -0,0 +1,132 @@ +# +# 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 overload +from typing import Union, Callable + +from pyspark.sql._typing import ( + AtomicDataTypeOrString, + UserDefinedFunctionLike, +) +from pyspark.sql.pandas._typing import ( + GroupedMapPandasUserDefinedFunction, + PandasGroupedAggFunction, + PandasGroupedAggUDFType, + PandasGroupedMapFunction, + PandasGroupedMapUDFType, + PandasScalarIterFunction, + PandasScalarIterUDFType, + PandasScalarToScalarFunction, + PandasScalarToStructFunction, + PandasScalarUDFType, +) + +from pyspark import since as since # noqa: F401 +from pyspark.rdd import PythonEvalType as PythonEvalType # noqa: F401 +from pyspark.sql.types import ArrayType, StructType + +class PandasUDFType: + SCALAR: PandasScalarUDFType + SCALAR_ITER: PandasScalarIterUDFType + GROUPED_MAP: PandasGroupedMapUDFType + GROUPED_AGG: PandasGroupedAggUDFType + +@overload +def pandas_udf( + f: PandasScalarToScalarFunction, + returnType: Union[AtomicDataTypeOrString, ArrayType], + functionType: PandasScalarUDFType, +) -> UserDefinedFunctionLike: ... +@overload +def pandas_udf(f: Union[AtomicDataTypeOrString, ArrayType], returnType: PandasScalarUDFType) -> Callable[[PandasScalarToScalarFunction], UserDefinedFunctionLike]: ... # type: ignore[misc] +@overload +def pandas_udf(f: Union[AtomicDataTypeOrString, ArrayType], *, functionType: PandasScalarUDFType) -> Callable[[PandasScalarToScalarFunction], UserDefinedFunctionLike]: ... # type: ignore[misc] +@overload +def pandas_udf(*, returnType: Union[AtomicDataTypeOrString, ArrayType], functionType: PandasScalarUDFType) -> Callable[[PandasScalarToScalarFunction], UserDefinedFunctionLike]: ... # type: ignore[misc] +@overload +def pandas_udf( + f: PandasScalarToStructFunction, + returnType: Union[StructType, str], + functionType: PandasScalarUDFType, +) -> UserDefinedFunctionLike: ... +@overload +def pandas_udf( + f: Union[StructType, str], returnType: PandasScalarUDFType +) -> Callable[[PandasScalarToStructFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + f: Union[StructType, str], *, functionType: PandasScalarUDFType +) -> Callable[[PandasScalarToStructFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + *, returnType: Union[StructType, str], functionType: PandasScalarUDFType +) -> Callable[[PandasScalarToStructFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + f: PandasScalarIterFunction, + returnType: Union[AtomicDataTypeOrString, ArrayType], + functionType: PandasScalarIterUDFType, +) -> UserDefinedFunctionLike: ... +@overload +def pandas_udf( + f: Union[AtomicDataTypeOrString, ArrayType], returnType: PandasScalarIterUDFType +) -> Callable[[PandasScalarIterFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + *, returnType: Union[AtomicDataTypeOrString, ArrayType], functionType: PandasScalarIterUDFType +) -> Callable[[PandasScalarIterFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + f: Union[AtomicDataTypeOrString, ArrayType], *, functionType: PandasScalarIterUDFType +) -> Callable[[PandasScalarIterFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + f: PandasGroupedMapFunction, + returnType: Union[StructType, str], + functionType: PandasGroupedMapUDFType, +) -> GroupedMapPandasUserDefinedFunction: ... +@overload +def pandas_udf( + f: Union[StructType, str], returnType: PandasGroupedMapUDFType +) -> Callable[[PandasGroupedMapFunction], GroupedMapPandasUserDefinedFunction]: ... +@overload +def pandas_udf( + *, returnType: Union[StructType, str], functionType: PandasGroupedMapUDFType +) -> Callable[[PandasGroupedMapFunction], GroupedMapPandasUserDefinedFunction]: ... +@overload +def pandas_udf( + f: Union[StructType, str], *, functionType: PandasGroupedMapUDFType +) -> Callable[[PandasGroupedMapFunction], GroupedMapPandasUserDefinedFunction]: ... +@overload +def pandas_udf( + f: PandasGroupedAggFunction, + returnType: Union[AtomicDataTypeOrString, ArrayType], + functionType: PandasGroupedAggUDFType, +) -> UserDefinedFunctionLike: ... +@overload +def pandas_udf( + f: Union[AtomicDataTypeOrString, ArrayType], returnType: PandasGroupedAggUDFType +) -> Callable[[PandasGroupedAggFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + *, returnType: Union[AtomicDataTypeOrString, ArrayType], functionType: PandasGroupedAggUDFType +) -> Callable[[PandasGroupedAggFunction], UserDefinedFunctionLike]: ... +@overload +def pandas_udf( + f: Union[AtomicDataTypeOrString, ArrayType], *, functionType: PandasGroupedAggUDFType +) -> Callable[[PandasGroupedAggFunction], UserDefinedFunctionLike]: ... diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/group_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/group_ops.py new file mode 100644 index 0000000..5640348 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/group_ops.py @@ -0,0 +1,523 @@ +# +# 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 sys +from typing import List, Union, TYPE_CHECKING, cast +import warnings + +from pyspark.rdd import PythonEvalType +from pyspark.sql.column import Column +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.streaming.state import GroupStateTimeout +from pyspark.sql.types import StructType, _parse_datatype_string + +if TYPE_CHECKING: + from pyspark.sql.pandas._typing import ( + GroupedMapPandasUserDefinedFunction, + PandasGroupedMapFunction, + PandasGroupedMapFunctionWithState, + PandasCogroupedMapFunction, + ) + from pyspark.sql.group import GroupedData + + +class PandasGroupedOpsMixin: + """ + Min-in for pandas grouped operations. Currently, only :class:`GroupedData` + can use this class. + """ + + def apply(self, udf: "GroupedMapPandasUserDefinedFunction") -> DataFrame: + """ + It is an alias of :meth:`pyspark.sql.GroupedData.applyInPandas`; however, it takes a + :meth:`pyspark.sql.functions.pandas_udf` whereas + :meth:`pyspark.sql.GroupedData.applyInPandas` takes a Python native function. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Support Spark Connect. + + Parameters + ---------- + udf : :func:`pyspark.sql.functions.pandas_udf` + a grouped map user-defined function returned by + :func:`pyspark.sql.functions.pandas_udf`. + + Notes + ----- + It is preferred to use :meth:`pyspark.sql.GroupedData.applyInPandas` over this + API. This API will be deprecated in the future releases. + + Examples + -------- + >>> from pyspark.sql.functions import pandas_udf, PandasUDFType + >>> df = spark.createDataFrame( + ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], + ... ("id", "v")) + >>> @pandas_udf("id long, v double", PandasUDFType.GROUPED_MAP) # doctest: +SKIP + ... def normalize(pdf): + ... v = pdf.v + ... return pdf.assign(v=(v - v.mean()) / v.std()) + ... + >>> df.groupby("id").apply(normalize).show() # doctest: +SKIP + +---+-------------------+ + | id| v| + +---+-------------------+ + | 1|-0.7071067811865475| + | 1| 0.7071067811865475| + | 2|-0.8320502943378437| + | 2|-0.2773500981126146| + | 2| 1.1094003924504583| + +---+-------------------+ + + See Also + -------- + pyspark.sql.functions.pandas_udf + """ + # Columns are special because hasattr always return True + if ( + isinstance(udf, Column) + or not hasattr(udf, "func") + or ( + udf.evalType # type: ignore[attr-defined] + != PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF + ) + ): + raise ValueError( + "Invalid udf: the udf argument must be a pandas_udf of type " "GROUPED_MAP." + ) + + warnings.warn( + "It is preferred to use 'applyInPandas' over this " + "API. This API will be deprecated in the future releases. See SPARK-28264 for " + "more details.", + UserWarning, + ) + + return self.applyInPandas(udf.func, schema=udf.returnType) # type: ignore[attr-defined] + + def applyInPandas( + self, func: "PandasGroupedMapFunction", schema: Union[StructType, str] + ) -> DataFrame: + """ + Maps each group of the current :class:`DataFrame` using a pandas udf and returns the result + as a `DataFrame`. + + The function should take a `pandas.DataFrame` and return another + `pandas.DataFrame`. Alternatively, the user can pass a function that takes + a tuple of the grouping key(s) and a `pandas.DataFrame`. + For each group, all columns are passed together as a `pandas.DataFrame` + to the user-function and the returned `pandas.DataFrame` are combined as a + :class:`DataFrame`. + + The `schema` should be a :class:`StructType` describing the schema of the returned + `pandas.DataFrame`. The column labels of the returned `pandas.DataFrame` must either match + the field names in the defined schema if specified as strings, or match the + field data types by position if not strings, e.g. integer indices. + The length of the returned `pandas.DataFrame` can be arbitrary. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Support Spark Connect. + + Parameters + ---------- + func : function + a Python native function that takes a `pandas.DataFrame` and outputs a + `pandas.DataFrame`, or that takes one tuple (grouping keys) and a + `pandas.DataFrame` and outputs a `pandas.DataFrame`. + schema : :class:`pyspark.sql.types.DataType` or str + the return type of the `func` in PySpark. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + + Examples + -------- + >>> import pandas as pd # doctest: +SKIP + >>> from pyspark.sql.functions import pandas_udf, ceil + >>> df = spark.createDataFrame( + ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], + ... ("id", "v")) # doctest: +SKIP + >>> def normalize(pdf): + ... v = pdf.v + ... return pdf.assign(v=(v - v.mean()) / v.std()) + ... + >>> df.groupby("id").applyInPandas( + ... normalize, schema="id long, v double").show() # doctest: +SKIP + +---+-------------------+ + | id| v| + +---+-------------------+ + | 1|-0.7071067811865475| + | 1| 0.7071067811865475| + | 2|-0.8320502943378437| + | 2|-0.2773500981126146| + | 2| 1.1094003924504583| + +---+-------------------+ + + Alternatively, the user can pass a function that takes two arguments. + In this case, the grouping key(s) will be passed as the first argument and the data will + be passed as the second argument. The grouping key(s) will be passed as a tuple of numpy + data types, e.g., `numpy.int32` and `numpy.float64`. The data will still be passed in + as a `pandas.DataFrame` containing all columns from the original Spark DataFrame. + This is useful when the user does not want to hardcode grouping key(s) in the function. + + >>> df = spark.createDataFrame( + ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], + ... ("id", "v")) # doctest: +SKIP + >>> def mean_func(key, pdf): + ... # key is a tuple of one numpy.int64, which is the value + ... # of 'id' for the current group + ... return pd.DataFrame([key + (pdf.v.mean(),)]) + ... + >>> df.groupby('id').applyInPandas( + ... mean_func, schema="id long, v double").show() # doctest: +SKIP + +---+---+ + | id| v| + +---+---+ + | 1|1.5| + | 2|6.0| + +---+---+ + + >>> def sum_func(key, pdf): + ... # key is a tuple of two numpy.int64s, which is the values + ... # of 'id' and 'ceil(df.v / 2)' for the current group + ... return pd.DataFrame([key + (pdf.v.sum(),)]) + ... + >>> df.groupby(df.id, ceil(df.v / 2)).applyInPandas( + ... sum_func, schema="id long, `ceil(v / 2)` long, v double").show() # doctest: +SKIP + +---+-----------+----+ + | id|ceil(v / 2)| v| + +---+-----------+----+ + | 2| 5|10.0| + | 1| 1| 3.0| + | 2| 3| 5.0| + | 2| 2| 3.0| + +---+-----------+----+ + + Notes + ----- + This function requires a full shuffle. All the data of a group will be loaded + into memory, so the user should be aware of the potential OOM risk if data is skewed + and certain groups are too large to fit in memory. + + This API is experimental. + + See Also + -------- + pyspark.sql.functions.pandas_udf + """ + from pyspark.sql import GroupedData + from pyspark.sql.functions import pandas_udf, PandasUDFType + + assert isinstance(self, GroupedData) + + udf = pandas_udf(func, returnType=schema, functionType=PandasUDFType.GROUPED_MAP) + df = self._df + udf_column = udf(*[df[col] for col in df.columns]) + jdf = self._jgd.flatMapGroupsInPandas(udf_column._jc.expr()) + return DataFrame(jdf, self.session) + + def applyInPandasWithState( + self, + func: "PandasGroupedMapFunctionWithState", + outputStructType: Union[StructType, str], + stateStructType: Union[StructType, str], + outputMode: str, + timeoutConf: str, + ) -> DataFrame: + """ + Applies the given function to each group of data, while maintaining a user-defined + per-group state. The result Dataset will represent the flattened record returned by the + function. + + For a streaming :class:`DataFrame`, the function will be invoked first for all input groups + and then for all timed out states where the input data is set to be empty. Updates to each + group's state will be saved across invocations. + + The function should take parameters (key, Iterator[`pandas.DataFrame`], state) and + return another Iterator[`pandas.DataFrame`]. The grouping key(s) will be passed as a tuple + of numpy data types, e.g., `numpy.int32` and `numpy.float64`. The state will be passed as + :class:`pyspark.sql.streaming.state.GroupState`. + + For each group, all columns are passed together as `pandas.DataFrame` to the user-function, + and the returned `pandas.DataFrame` across all invocations are combined as a + :class:`DataFrame`. Note that the user function should not make a guess of the number of + elements in the iterator. To process all data, the user function needs to iterate all + elements and process them. On the other hand, the user function is not strictly required to + iterate through all elements in the iterator if it intends to read a part of data. + + The `outputStructType` should be a :class:`StructType` describing the schema of all + elements in the returned value, `pandas.DataFrame`. The column labels of all elements in + returned `pandas.DataFrame` must either match the field names in the defined schema if + specified as strings, or match the field data types by position if not strings, + e.g. integer indices. + + The `stateStructType` should be :class:`StructType` describing the schema of the + user-defined state. The value of the state will be presented as a tuple, as well as the + update should be performed with the tuple. The corresponding Python types for + :class:DataType are supported. Please refer to the page + https://spark.apache.org/docs/latest/sql-ref-datatypes.html (Python tab). + + The size of each `pandas.DataFrame` in both the input and output can be arbitrary. The + number of `pandas.DataFrame` in both the input and output can also be arbitrary. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + func : function + a Python native function to be called on every group. It should take parameters + (key, Iterator[`pandas.DataFrame`], state) and return Iterator[`pandas.DataFrame`]. + Note that the type of the key is tuple and the type of the state is + :class:`pyspark.sql.streaming.state.GroupState`. + outputStructType : :class:`pyspark.sql.types.DataType` or str + the type of the output records. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + stateStructType : :class:`pyspark.sql.types.DataType` or str + the type of the user-defined state. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + outputMode : str + the output mode of the function. + timeoutConf : str + timeout configuration for groups that do not receive data for a while. valid values + are defined in :class:`pyspark.sql.streaming.state.GroupStateTimeout`. + + Examples + -------- + >>> import pandas as pd # doctest: +SKIP + >>> from pyspark.sql.streaming.state import GroupStateTimeout + >>> def count_fn(key, pdf_iter, state): + ... assert isinstance(state, GroupStateImpl) + ... total_len = 0 + ... for pdf in pdf_iter: + ... total_len += len(pdf) + ... state.update((total_len,)) + ... yield pd.DataFrame({"id": [key[0]], "countAsString": [str(total_len)]}) + ... + >>> df.groupby("id").applyInPandasWithState( + ... count_fn, outputStructType="id long, countAsString string", + ... stateStructType="len long", outputMode="Update", + ... timeoutConf=GroupStateTimeout.NoTimeout) # doctest: +SKIP + + Notes + ----- + This function requires a full shuffle. + + This API is experimental. + """ + + from pyspark.sql import GroupedData + from pyspark.sql.functions import pandas_udf + + assert isinstance(self, GroupedData) + assert timeoutConf in [ + GroupStateTimeout.NoTimeout, + GroupStateTimeout.ProcessingTimeTimeout, + GroupStateTimeout.EventTimeTimeout, + ] + + if isinstance(outputStructType, str): + outputStructType = cast(StructType, _parse_datatype_string(outputStructType)) + if isinstance(stateStructType, str): + stateStructType = cast(StructType, _parse_datatype_string(stateStructType)) + + udf = pandas_udf( + func, # type: ignore[call-overload] + returnType=outputStructType, + functionType=PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + ) + df = self._df + udf_column = udf(*[df[col] for col in df.columns]) + jdf = self._jgd.applyInPandasWithState( + udf_column._jc.expr(), + self.session._jsparkSession.parseDataType(outputStructType.json()), + self.session._jsparkSession.parseDataType(stateStructType.json()), + outputMode, + timeoutConf, + ) + return DataFrame(jdf, self.session) + + def cogroup(self, other: "GroupedData") -> "PandasCogroupedOps": + """ + Cogroups this group with another group so that we can run cogrouped operations. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Support Spark Connect. + + See :class:`PandasCogroupedOps` for the operations that can be run. + """ + from pyspark.sql import GroupedData + + assert isinstance(self, GroupedData) + + return PandasCogroupedOps(self, other) + + +class PandasCogroupedOps: + """ + A logical grouping of two :class:`GroupedData`, + created by :func:`GroupedData.cogroup`. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Support Spark Connect. + + Notes + ----- + This API is experimental. + """ + + def __init__(self, gd1: "GroupedData", gd2: "GroupedData"): + self._gd1 = gd1 + self._gd2 = gd2 + + def applyInPandas( + self, func: "PandasCogroupedMapFunction", schema: Union[StructType, str] + ) -> DataFrame: + """ + Applies a function to each cogroup using pandas and returns the result + as a `DataFrame`. + + The function should take two `pandas.DataFrame`\\s and return another + `pandas.DataFrame`. Alternatively, the user can pass a function that takes + a tuple of the grouping key(s) and the two `pandas.DataFrame`\\s. + For each side of the cogroup, all columns are passed together as a + `pandas.DataFrame` to the user-function and the returned `pandas.DataFrame` are combined as + a :class:`DataFrame`. + + The `schema` should be a :class:`StructType` describing the schema of the returned + `pandas.DataFrame`. The column labels of the returned `pandas.DataFrame` must either match + the field names in the defined schema if specified as strings, or match the + field data types by position if not strings, e.g. integer indices. + The length of the returned `pandas.DataFrame` can be arbitrary. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Support Spark Connect. + + Parameters + ---------- + func : function + a Python native function that takes two `pandas.DataFrame`\\s, and + outputs a `pandas.DataFrame`, or that takes one tuple (grouping keys) and two + ``pandas.DataFrame``\\s, and outputs a ``pandas.DataFrame``. + schema : :class:`pyspark.sql.types.DataType` or str + the return type of the `func` in PySpark. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + + Examples + -------- + >>> from pyspark.sql.functions import pandas_udf + >>> df1 = spark.createDataFrame( + ... [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)], + ... ("time", "id", "v1")) + >>> df2 = spark.createDataFrame( + ... [(20000101, 1, "x"), (20000101, 2, "y")], + ... ("time", "id", "v2")) + >>> def asof_join(l, r): + ... return pd.merge_asof(l, r, on="time", by="id") + ... + >>> df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas( + ... asof_join, schema="time int, id int, v1 double, v2 string" + ... ).show() # doctest: +SKIP + +--------+---+---+---+ + | time| id| v1| v2| + +--------+---+---+---+ + |20000101| 1|1.0| x| + |20000102| 1|3.0| x| + |20000101| 2|2.0| y| + |20000102| 2|4.0| y| + +--------+---+---+---+ + + Alternatively, the user can define a function that takes three arguments. In this case, + the grouping key(s) will be passed as the first argument and the data will be passed as the + second and third arguments. The grouping key(s) will be passed as a tuple of numpy data + types, e.g., `numpy.int32` and `numpy.float64`. The data will still be passed in as two + `pandas.DataFrame` containing all columns from the original Spark DataFrames. + + >>> def asof_join(k, l, r): + ... if k == (1,): + ... return pd.merge_asof(l, r, on="time", by="id") + ... else: + ... return pd.DataFrame(columns=['time', 'id', 'v1', 'v2']) + ... + >>> df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas( + ... asof_join, "time int, id int, v1 double, v2 string").show() # doctest: +SKIP + +--------+---+---+---+ + | time| id| v1| v2| + +--------+---+---+---+ + |20000101| 1|1.0| x| + |20000102| 1|3.0| x| + +--------+---+---+---+ + + Notes + ----- + This function requires a full shuffle. All the data of a cogroup will be loaded + into memory, so the user should be aware of the potential OOM risk if data is skewed + and certain groups are too large to fit in memory. + + This API is experimental. + + See Also + -------- + pyspark.sql.functions.pandas_udf + """ + from pyspark.sql.pandas.functions import pandas_udf + + # The usage of the pandas_udf is internal so type checking is disabled. + udf = pandas_udf( + func, returnType=schema, functionType=PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF + ) # type: ignore[call-overload] + + all_cols = self._extract_cols(self._gd1) + self._extract_cols(self._gd2) + udf_column = udf(*all_cols) + jdf = self._gd1._jgd.flatMapCoGroupsInPandas(self._gd2._jgd, udf_column._jc.expr()) + return DataFrame(jdf, self._gd1.session) + + @staticmethod + def _extract_cols(gd: "GroupedData") -> List[Column]: + df = gd._df + return [df[col] for col in df.columns] + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.pandas.group_ops + + globs = pyspark.sql.pandas.group_ops.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.pandas.group tests").getOrCreate() + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.pandas.group_ops, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/map_ops.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/map_ops.py new file mode 100644 index 0000000..bc26fde --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/map_ops.py @@ -0,0 +1,218 @@ +# +# 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 sys +from typing import Union, TYPE_CHECKING + +from pyspark.rdd import PythonEvalType +from pyspark.sql.types import StructType + +if TYPE_CHECKING: + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.pandas._typing import PandasMapIterFunction, ArrowMapIterFunction + + +class PandasMapOpsMixin: + """ + Min-in for pandas map operations. Currently, only :class:`DataFrame` + can use this class. + """ + + def mapInPandas( + self, func: "PandasMapIterFunction", schema: Union[StructType, str], barrier: bool = False + ) -> "DataFrame": + """ + Maps an iterator of batches in the current :class:`DataFrame` using a Python native + function that takes and outputs a pandas DataFrame, and returns the result as a + :class:`DataFrame`. + + The function should take an iterator of `pandas.DataFrame`\\s and return + another iterator of `pandas.DataFrame`\\s. All columns are passed + together as an iterator of `pandas.DataFrame`\\s to the function and the + returned iterator of `pandas.DataFrame`\\s are combined as a :class:`DataFrame`. + Each `pandas.DataFrame` size can be controlled by + `spark.sql.execution.arrow.maxRecordsPerBatch`. The size of the function's input and + output can be different. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + func : function + a Python native function that takes an iterator of `pandas.DataFrame`\\s, and + outputs an iterator of `pandas.DataFrame`\\s. + schema : :class:`pyspark.sql.types.DataType` or str + the return type of the `func` in PySpark. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + barrier : bool, optional, default True + Use barrier mode execution. + + .. versionchanged: 3.5.0 + Added ``barrier`` argument. + + Examples + -------- + >>> from pyspark.sql.functions import pandas_udf + >>> df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age")) + >>> def filter_func(iterator): + ... for pdf in iterator: + ... yield pdf[pdf.id == 1] + ... + >>> df.mapInPandas(filter_func, df.schema).show() # doctest: +SKIP + +---+---+ + | id|age| + +---+---+ + | 1| 21| + +---+---+ + + Set ``barrier`` to ``True`` to force the ``mapInPandas`` stage running in the + barrier mode, it ensures all Python workers in the stage will be + launched concurrently. + + >>> df.mapInPandas(filter_func, df.schema, barrier=True).show() # doctest: +SKIP + +---+---+ + | id|age| + +---+---+ + | 1| 21| + +---+---+ + + Notes + ----- + This API is experimental + + See Also + -------- + pyspark.sql.functions.pandas_udf + """ + from pyspark.sql import DataFrame + from pyspark.sql.pandas.functions import pandas_udf + + assert isinstance(self, DataFrame) + + # The usage of the pandas_udf is internal so type checking is disabled. + udf = pandas_udf( + func, returnType=schema, functionType=PythonEvalType.SQL_MAP_PANDAS_ITER_UDF + ) # type: ignore[call-overload] + udf_column = udf(*[self[col] for col in self.columns]) + jdf = self._jdf.mapInPandas(udf_column._jc.expr(), barrier) + return DataFrame(jdf, self.sparkSession) + + def mapInArrow( + self, func: "ArrowMapIterFunction", schema: Union[StructType, str], barrier: bool = False + ) -> "DataFrame": + """ + Maps an iterator of batches in the current :class:`DataFrame` using a Python native + function that takes and outputs a PyArrow's `RecordBatch`, and returns the result as a + :class:`DataFrame`. + + The function should take an iterator of `pyarrow.RecordBatch`\\s and return + another iterator of `pyarrow.RecordBatch`\\s. All columns are passed + together as an iterator of `pyarrow.RecordBatch`\\s to the function and the + returned iterator of `pyarrow.RecordBatch`\\s are combined as a :class:`DataFrame`. + Each `pyarrow.RecordBatch` size can be controlled by + `spark.sql.execution.arrow.maxRecordsPerBatch`. The size of the function's input and + output can be different. + + .. versionadded:: 3.3.0 + + Parameters + ---------- + func : function + a Python native function that takes an iterator of `pyarrow.RecordBatch`\\s, and + outputs an iterator of `pyarrow.RecordBatch`\\s. + schema : :class:`pyspark.sql.types.DataType` or str + the return type of the `func` in PySpark. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + barrier : bool, optional, default True + Use barrier mode execution. + + .. versionchanged: 3.5.0 + Added ``barrier`` argument. + + Examples + -------- + >>> import pyarrow # doctest: +SKIP + >>> df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age")) + >>> def filter_func(iterator): + ... for batch in iterator: + ... pdf = batch.to_pandas() + ... yield pyarrow.RecordBatch.from_pandas(pdf[pdf.id == 1]) + >>> df.mapInArrow(filter_func, df.schema).show() # doctest: +SKIP + +---+---+ + | id|age| + +---+---+ + | 1| 21| + +---+---+ + + Set ``barrier`` to ``True`` to force the ``mapInArrow`` stage running in the + barrier mode, it ensures all Python workers in the stage will be + launched concurrently. + + >>> df.mapInArrow(filter_func, df.schema, barrier=True).show() # doctest: +SKIP + +---+---+ + | id|age| + +---+---+ + | 1| 21| + +---+---+ + + Notes + ----- + This API is unstable, and for developers. + + See Also + -------- + pyspark.sql.functions.pandas_udf + pyspark.sql.DataFrame.mapInPandas + """ + from pyspark.sql import DataFrame + from pyspark.sql.pandas.functions import pandas_udf + + assert isinstance(self, DataFrame) + + # The usage of the pandas_udf is internal so type checking is disabled. + udf = pandas_udf( + func, returnType=schema, functionType=PythonEvalType.SQL_MAP_ARROW_ITER_UDF + ) # type: ignore[call-overload] + udf_column = udf(*[self[col] for col in self.columns]) + jdf = self._jdf.pythonMapInArrow(udf_column._jc.expr(), barrier) + return DataFrame(jdf, self.sparkSession) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.pandas.map_ops + + globs = pyspark.sql.pandas.map_ops.__dict__.copy() + spark = ( + SparkSession.builder.master("local[4]").appName("sql.pandas.map_ops tests").getOrCreate() + ) + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.pandas.map_ops, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/serializers.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/serializers.py new file mode 100644 index 0000000..2cc3db1 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/serializers.py @@ -0,0 +1,1019 @@ +# +# 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. +# + +""" +Serializers for PyArrow and pandas conversions. See `pyspark.serializers` for more details. +""" + +from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError +from pyspark.serializers import Serializer, read_int, write_int, UTF8Deserializer, CPickleSerializer +from pyspark.sql.pandas.types import ( + from_arrow_type, + to_arrow_type, + _create_converter_from_pandas, + _create_converter_to_pandas, +) +from pyspark.sql.types import ( + DataType, + StringType, + StructType, + BinaryType, + StructField, + LongType, + IntegerType, +) + + +class SpecialLengths: + END_OF_DATA_SECTION = -1 + PYTHON_EXCEPTION_THROWN = -2 + TIMING_DATA = -3 + END_OF_STREAM = -4 + NULL = -5 + 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 RuntimeError( + "An error occurred while calling " + "ArrowCollectSerializer.load_stream: {}".format(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 stream. + """ + + def dump_stream(self, iterator, stream): + import pyarrow as pa + + writer = None + try: + for batch in iterator: + if writer is None: + writer = pa.RecordBatchStreamWriter(stream, batch.schema) + writer.write_batch(batch) + finally: + if writer is not None: + writer.close() + + def load_stream(self, stream): + import pyarrow as pa + + reader = pa.ipc.open_stream(stream) + for batch in reader: + yield batch + + def __repr__(self): + return "ArrowStreamSerializer" + + +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. + """ + import pyarrow as pa + + batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) + for batch in batches: + struct = batch.column(0) + yield [pa.RecordBatch.from_arrays(struct.flatten(), schema=pa.schema(struct.type))] + + 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. + """ + import pyarrow as pa + + def wrap_and_init_stream(): + should_write_start_length = True + for batch, _ in iterator: + assert isinstance(batch, pa.RecordBatch) + + # Wrap the root struct + struct = pa.StructArray.from_arrays( + batch.columns, fields=pa.struct(list(batch.schema)) + ) + batch = pa.RecordBatch.from_arrays([struct], ["_0"]) + + # Write the first record batch with initialization. + if should_write_start_length: + write_int(SpecialLengths.START_ARROW_STREAM, stream) + should_write_start_length = False + yield batch + + return super(ArrowStreamUDFSerializer, self).dump_stream(wrap_and_init_stream(), 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 + assign_cols_by_name : bool + If True, then Pandas DataFrames will get columns by name + """ + + def __init__(self, timezone, safecheck): + super(ArrowStreamPandasSerializer, self).__init__() + self._timezone = timezone + self._safecheck = safecheck + + def arrow_to_pandas(self, arrow_column, struct_in_pandas="dict", ndarray_as_list=False): + # If the given column is a date type column, creates a series of datetime.date directly + # instead of creating datetime64[ns] as intermediate data to avoid overflow caused by + # datetime64[ns] type handling. + # Cast dates to objects instead of datetime64[ns] dtype to avoid overflow. + s = arrow_column.to_pandas(date_as_object=True) + + # TODO(SPARK-43579): cache the converter for reuse + converter = _create_converter_to_pandas( + data_type=from_arrow_type(arrow_column.type, prefer_timestamp_ntz=True), + nullable=True, + timezone=self._timezone, + struct_in_pandas=struct_in_pandas, + error_on_duplicated_field_names=True, + ndarray_as_list=ndarray_as_list, + ) + return converter(s) + + def _create_array(self, series, arrow_type, spark_type=None, arrow_cast=False): + """ + Create an Arrow Array from the given pandas.Series and optional type. + + Parameters + ---------- + series : pandas.Series + A single series + arrow_type : pyarrow.DataType, optional + If None, pyarrow's inferred type will be used + spark_type : DataType, optional + If None, spark type converted from arrow_type will be used + arrow_cast: bool, optional + Whether to apply Arrow casting when the user-specified return type mismatches the + actual return values. + + Returns + ------- + pyarrow.Array + """ + import pyarrow as pa + from pandas.api.types import is_categorical_dtype + + if is_categorical_dtype(series.dtype): + series = series.astype(series.dtypes.categories.dtype) + + if arrow_type is not None: + dt = spark_type or from_arrow_type(arrow_type, prefer_timestamp_ntz=True) + # TODO(SPARK-43579): cache the converter for reuse + conv = _create_converter_from_pandas( + dt, timezone=self._timezone, error_on_duplicated_field_names=False + ) + series = conv(series) + + if hasattr(series.array, "__arrow_array__"): + mask = None + else: + mask = series.isnull() + try: + try: + return pa.Array.from_pandas( + series, mask=mask, type=arrow_type, safe=self._safecheck + ) + except pa.lib.ArrowInvalid: + if arrow_cast: + return pa.Array.from_pandas(series, mask=mask).cast( + target_type=arrow_type, safe=self._safecheck + ) + else: + raise + except TypeError as e: + error_msg = ( + "Exception thrown when converting pandas.Series (%s) " + "with name '%s' to Arrow Array (%s)." + ) + raise PySparkTypeError(error_msg % (series.dtype, series.name, arrow_type)) from e + except ValueError as e: + error_msg = ( + "Exception thrown when converting pandas.Series (%s) " + "with name '%s' to Arrow Array (%s)." + ) + if self._safecheck: + error_msg = error_msg + ( + " It can be caused by overflows or other " + "unsafe conversions warned by Arrow. Arrow safe type check " + "can be disabled by using SQL config " + "`spark.sql.execution.pandas.convertToArrowArraySafely`." + ) + raise PySparkValueError(error_msg % (series.dtype, series.name, arrow_type)) from e + + def _create_batch(self, series): + """ + Create an Arrow record batch from the given pandas.Series or list of Series, + with optional type. + + Parameters + ---------- + series : pandas.Series or list + A single series, list of series, or list of (series, arrow_type) + + Returns + ------- + pyarrow.RecordBatch + Arrow RecordBatch + """ + import pyarrow as pa + + # Make input conform to + # [(series1, arrow_type1, spark_type1), (series2, arrow_type2, spark_type2), ...] + if ( + not isinstance(series, (list, tuple)) + or (len(series) == 2 and isinstance(series[1], pa.DataType)) + or ( + len(series) == 3 + and isinstance(series[1], pa.DataType) + and isinstance(series[2], DataType) + ) + ): + series = [series] + series = ((s, None) if not isinstance(s, (list, tuple)) else s for s in series) + series = ((s[0], s[1], None) if len(s) == 2 else s for s in series) + + arrs = [ + self._create_array(s, arrow_type, spark_type) for s, arrow_type, spark_type in series + ] + return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))]) + + def dump_stream(self, iterator, stream): + """ + Make ArrowRecordBatches from Pandas Series and serialize. Input is a single series or + a list of series accompanied by an optional pyarrow type to coerce the data to. + """ + batches = (self._create_batch(series) for series in iterator) + super(ArrowStreamPandasSerializer, self).dump_stream(batches, stream) + + def load_stream(self, stream): + """ + Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series. + """ + batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) + import pyarrow as pa + + for batch in batches: + yield [self.arrow_to_pandas(c) for c in pa.Table.from_batches([batch]).itercolumns()] + + 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=False, + struct_in_pandas="dict", + ndarray_as_list=False, + arrow_cast=False, + ): + super(ArrowStreamPandasUDFSerializer, self).__init__(timezone, safecheck) + self._assign_cols_by_name = assign_cols_by_name + self._df_for_struct = df_for_struct + self._struct_in_pandas = struct_in_pandas + self._ndarray_as_list = ndarray_as_list + self._arrow_cast = arrow_cast + + def arrow_to_pandas(self, arrow_column): + import pyarrow.types as types + + if self._df_for_struct and types.is_struct(arrow_column.type): + import pandas as pd + + series = [ + super(ArrowStreamPandasUDFSerializer, self) + .arrow_to_pandas(column, self._struct_in_pandas, self._ndarray_as_list) + .rename(field.name) + for column, field in zip(arrow_column.flatten(), arrow_column.type) + ] + s = pd.concat(series, axis=1) + else: + s = super(ArrowStreamPandasUDFSerializer, self).arrow_to_pandas( + arrow_column, self._struct_in_pandas, self._ndarray_as_list + ) + return s + + def _create_struct_array(self, df, arrow_struct_type): + """ + Create an Arrow StructArray from the given pandas.DataFrame and arrow struct type. + + Parameters + ---------- + df : pandas.DataFrame + A pandas DataFrame + arrow_struct_type : pyarrow.DataType + pyarrow struct type + + Returns + ------- + pyarrow.Array + """ + import pyarrow as pa + + if len(df.columns) == 0: + return pa.array([{}] * len(df), arrow_struct_type) + # Assign result columns by schema name if user labeled with strings + if self._assign_cols_by_name and any(isinstance(name, str) for name in df.columns): + struct_arrs = [ + self._create_array(df[field.name], field.type, arrow_cast=self._arrow_cast) + for field in arrow_struct_type + ] + # Assign result columns by position + else: + struct_arrs = [ + # the selected series has name '1', so we rename it to field.name + # as the name is used by _create_array to provide a meaningful error message + self._create_array( + df[df.columns[i]].rename(field.name), + field.type, + arrow_cast=self._arrow_cast, + ) + for i, field in enumerate(arrow_struct_type) + ] + + struct_names = [field.name for field in arrow_struct_type] + return pa.StructArray.from_arrays(struct_arrs, struct_names) + + def _create_batch(self, series): + """ + Create an Arrow record batch from the given pandas.Series pandas.DataFrame + or list of Series or DataFrame, with optional type. + + Parameters + ---------- + series : pandas.Series or pandas.DataFrame or list + A single series or dataframe, list of series or dataframe, + or list of (series or dataframe, arrow_type) + + Returns + ------- + pyarrow.RecordBatch + Arrow RecordBatch + """ + import pandas as pd + import pyarrow as pa + + # Make input conform to [(series1, type1), (series2, type2), ...] + if not isinstance(series, (list, tuple)) or ( + len(series) == 2 and isinstance(series[1], pa.DataType) + ): + series = [series] + series = ((s, None) if not isinstance(s, (list, tuple)) else s for s in series) + + arrs = [] + for s, t in series: + if self._struct_in_pandas == "dict" and t is not None and pa.types.is_struct(t): + # A pandas UDF should return pd.DataFrame when the return type is a struct type. + # If it returns a pd.Series, it should throw an error. + if not isinstance(s, pd.DataFrame): + raise PySparkValueError( + "A field of type StructType expects a pandas.DataFrame, " + "but got: %s" % str(type(s)) + ) + arrs.append(self._create_struct_array(s, t)) + else: + arrs.append(self._create_array(s, t, arrow_cast=self._arrow_cast)) + + return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))]) + + 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. + """ + + def init_stream_yield_batches(): + should_write_start_length = True + for series in iterator: + batch = self._create_batch(series) + if should_write_start_length: + write_int(SpecialLengths.START_ARROW_STREAM, stream) + should_write_start_length = False + yield batch + + return ArrowStreamSerializer.dump_stream(self, init_stream_yield_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): + super(ArrowStreamPandasUDTFSerializer, self).__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, + # Enables explicit casting for mismatched return types of Arrow Python UDTFs. + arrow_cast=True, + ) + self._converter_map = dict() + + def _create_batch(self, series): + """ + Create an Arrow record batch from the given pandas.Series pandas.DataFrame + or list of Series or DataFrame, with optional type. + + Parameters + ---------- + series : pandas.Series or pandas.DataFrame or list + A single series or dataframe, list of series or dataframe, + or list of (series or dataframe, arrow_type) + + Returns + ------- + pyarrow.RecordBatch + Arrow RecordBatch + """ + import pandas as pd + import pyarrow as pa + + # Make input conform to [(series1, type1), (series2, type2), ...] + if not isinstance(series, (list, tuple)) or ( + len(series) == 2 and isinstance(series[1], pa.DataType) + ): + series = [series] + series = ((s, None) if not isinstance(s, (list, tuple)) else s for s in series) + + arrs = [] + for s, t in series: + if not isinstance(s, pd.DataFrame): + raise PySparkValueError( + "Output of an arrow-optimized Python UDTFs expects " + f"a pandas.DataFrame but got: {type(s)}" + ) + + arrs.append(self._create_struct_array(s, t)) + + return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))]) + + def _get_or_create_converter_from_pandas(self, dt): + if dt not in self._converter_map: + conv = _create_converter_from_pandas( + dt, + timezone=self._timezone, + error_on_duplicated_field_names=False, + ignore_unexpected_complex_type_values=True, + ) + self._converter_map[dt] = conv + return self._converter_map[dt] + + def _create_array(self, series, arrow_type, spark_type=None, arrow_cast=False): + """ + Override the `_create_array` method in the superclass to create an Arrow Array + from a given pandas.Series and an arrow type. The difference here is that we always + use arrow cast when creating the arrow array. Also, the error messages are specific + to arrow-optimized Python UDTFs. + + Parameters + ---------- + series : pandas.Series + A single series + arrow_type : pyarrow.DataType, optional + If None, pyarrow's inferred type will be used + spark_type : DataType, optional + If None, spark type converted from arrow_type will be used + arrow_cast: bool, optional + Whether to apply Arrow casting when the user-specified return type mismatches the + actual return values. + + Returns + ------- + pyarrow.Array + """ + import pyarrow as pa + from pandas.api.types import is_categorical_dtype + + if is_categorical_dtype(series.dtype): + series = series.astype(series.dtypes.categories.dtype) + + if arrow_type is not None: + dt = spark_type or from_arrow_type(arrow_type, prefer_timestamp_ntz=True) + conv = self._get_or_create_converter_from_pandas(dt) + series = conv(series) + + if hasattr(series.array, "__arrow_array__"): + mask = None + else: + mask = series.isnull() + + try: + try: + return pa.Array.from_pandas( + series, mask=mask, type=arrow_type, safe=self._safecheck + ) + except pa.lib.ArrowException: + if arrow_cast: + return pa.Array.from_pandas(series, mask=mask).cast( + target_type=arrow_type, safe=self._safecheck + ) + else: + raise + except pa.lib.ArrowException: + # Display the most user-friendly error messages instead of showing + # arrow's error message. This also works better with Spark Connect + # where the exception messages are by default truncated. + raise PySparkRuntimeError( + error_class="UDTF_ARROW_TYPE_CAST_ERROR", + message_parameters={ + "col_name": series.name, + "col_type": str(series.dtype), + "arrow_type": arrow_type, + }, + ) from None + + def __repr__(self): + return "ArrowStreamPandasUDTFSerializer" + + +class CogroupUDFSerializer(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 + + dataframes_in_group = None + + while dataframes_in_group is None or dataframes_in_group > 0: + dataframes_in_group = read_int(stream) + + if dataframes_in_group == 2: + batch1 = [batch for batch in ArrowStreamSerializer.load_stream(self, stream)] + batch2 = [batch for batch in ArrowStreamSerializer.load_stream(self, stream)] + yield ( + [self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch1).itercolumns()], + [self.arrow_to_pandas(c) for c in pa.Table.from_batches(batch2).itercolumns()], + ) + + elif dataframes_in_group != 0: + raise ValueError( + "Invalid number of pandas.DataFrames in group {0}".format(dataframes_in_group) + ) + + +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, + state_object_schema, + arrow_max_records_per_batch, + ): + super(ApplyInPandasWithStateSerializer, self).__init__( + timezone, safecheck, assign_cols_by_name + ) + 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) + + 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) + self.arrow_max_records_per_batch = arrow_max_records_per_batch + + 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_arrow = pa.Table.from_batches([state_batch]).itercolumns() + state_pandas = [self.arrow_to_pandas(c) for c in state_arrow][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_arrow = pa.Table.from_batches([data_batch_for_group]).itercolumns() + + data_pandas = [self.arrow_to_pandas(c) for c in data_arrow] + + # 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 + import pyarrow as pa + + 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(pa.schema(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) + + return self._create_batch( + [ + (count_pdf, self.result_count_pdf_arrow_type), + (merged_pdf, pdf_schema), + (merged_state_pdf, self.result_state_pdf_arrow_type), + ] + ) + + 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 + + def init_stream_yield_batches(batches): + """ + This function helps to ensure the requirement for Pandas UDFs - Pandas UDFs require a + START_ARROW_STREAM before the Arrow stream is sent. + + START_ARROW_STREAM 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. + """ + should_write_start_length = True + + for batch in batches: + if should_write_start_length: + write_int(SpecialLengths.START_ARROW_STREAM, stream) + should_write_start_length = False + + yield batch + + batches_to_write = init_stream_yield_batches(serialize_batches()) + + return ArrowStreamSerializer.dump_stream(self, batches_to_write, stream) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/typehints.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/typehints.py new file mode 100644 index 0000000..f0c13e6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/typehints.py @@ -0,0 +1,172 @@ +# +# 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 inspect import Signature +from typing import Any, Callable, Dict, Optional, Union, TYPE_CHECKING + +from pyspark.sql.pandas.utils import require_minimum_pandas_version +from pyspark.errors import PySparkNotImplementedError + +if TYPE_CHECKING: + from pyspark.sql.pandas._typing import ( + PandasScalarUDFType, + PandasScalarIterUDFType, + PandasGroupedAggUDFType, + ) + + +def infer_eval_type( + sig: Signature, type_hints: Dict[str, Any] +) -> Union["PandasScalarUDFType", "PandasScalarIterUDFType", "PandasGroupedAggUDFType"]: + """ + Infers the evaluation type in :class:`pyspark.rdd.PythonEvalType` from + :class:`inspect.Signature` instance and type hints. + """ + from pyspark.sql.pandas.functions import PandasUDFType + + require_minimum_pandas_version() + + import pandas as pd + + annotations = {} + for param in sig.parameters.values(): + if param.annotation is not param.empty: + annotations[param.name] = type_hints.get(param.name, param.annotation) + + # Check if all arguments have type hints + parameters_sig = [ + annotations[parameter] for parameter in sig.parameters if parameter in annotations + ] + if len(parameters_sig) != len(sig.parameters): + raise ValueError("Type hints for all parameters should be specified; however, got %s" % sig) + + # Check if the return has a type hint + return_annotation = type_hints.get("return", sig.return_annotation) + if sig.empty is return_annotation: + raise ValueError("Type hint for the return type should be specified; however, got %s" % sig) + + # Series, Frame or Union[DataFrame, Series], ... -> Series or Frame + is_series_or_frame = all( + a == pd.Series + or a == pd.DataFrame # Series + or check_union_annotation( # DataFrame # Union[DataFrame, Series] + a, parameter_check_func=lambda na: na == pd.Series or na == pd.DataFrame + ) + for a in parameters_sig + ) and (return_annotation == pd.Series or return_annotation == pd.DataFrame) + + # Iterator[Tuple[Series, Frame or Union[DataFrame, Series], ...] -> Iterator[Series or Frame] + is_iterator_tuple_series_or_frame = ( + len(parameters_sig) == 1 + and check_iterator_annotation( # Iterator + parameters_sig[0], + parameter_check_func=lambda a: check_tuple_annotation( # Tuple + a, + parameter_check_func=lambda ta: ( + ta == Ellipsis + or ta == pd.Series # ... + or ta == pd.DataFrame # Series + or check_union_annotation( # DataFrame # Union[DataFrame, Series] + ta, parameter_check_func=lambda na: (na == pd.Series or na == pd.DataFrame) + ) + ), + ), + ) + and check_iterator_annotation( + return_annotation, parameter_check_func=lambda a: a == pd.DataFrame or a == pd.Series + ) + ) + + # Iterator[Series, Frame or Union[DataFrame, Series]] -> Iterator[Series or Frame] + is_iterator_series_or_frame = ( + len(parameters_sig) == 1 + and check_iterator_annotation( + parameters_sig[0], + parameter_check_func=lambda a: ( + a == pd.Series + or a == pd.DataFrame # Series + or check_union_annotation( # DataFrame # Union[DataFrame, Series] + a, parameter_check_func=lambda ua: ua == pd.Series or ua == pd.DataFrame + ) + ), + ) + and check_iterator_annotation( + return_annotation, parameter_check_func=lambda a: a == pd.DataFrame or a == pd.Series + ) + ) + + # Series, Frame or Union[DataFrame, Series], ... -> Any + is_series_or_frame_agg = all( + a == pd.Series + or a == pd.DataFrame # Series + or check_union_annotation( # DataFrame # Union[DataFrame, Series] + a, parameter_check_func=lambda ua: ua == pd.Series or ua == pd.DataFrame + ) + for a in parameters_sig + ) and ( + # It's tricky to include only types which pd.Series constructor can take. + # Simply exclude common types used here for now (which becomes object + # types Spark can't recognize). + return_annotation != pd.Series + and return_annotation != pd.DataFrame + and not check_iterator_annotation(return_annotation) + and not check_tuple_annotation(return_annotation) + ) + + if is_series_or_frame: + return PandasUDFType.SCALAR + elif is_iterator_tuple_series_or_frame or is_iterator_series_or_frame: + return PandasUDFType.SCALAR_ITER + elif is_series_or_frame_agg: + return PandasUDFType.GROUPED_AGG + else: + raise PySparkNotImplementedError( + error_class="UNSUPPORTED_SIGNATURE", + message_parameters={"signature": str(sig)}, + ) + + +def check_tuple_annotation( + annotation: Any, parameter_check_func: Optional[Callable[[Any], bool]] = None +) -> bool: + # Tuple has _name but other types have __name__ + # Check if the name is Tuple first. After that, check the generic types. + name = getattr(annotation, "_name", getattr(annotation, "__name__", None)) + return name in ("Tuple", "tuple") and ( + parameter_check_func is None or all(map(parameter_check_func, annotation.__args__)) + ) + + +def check_iterator_annotation( + annotation: Any, parameter_check_func: Optional[Callable[[Any], bool]] = None +) -> bool: + name = getattr(annotation, "_name", getattr(annotation, "__name__", None)) + return name == "Iterator" and ( + parameter_check_func is None or all(map(parameter_check_func, annotation.__args__)) + ) + + +def check_union_annotation( + annotation: Any, parameter_check_func: Optional[Callable[[Any], bool]] = None +) -> bool: + import typing + + # Note that we cannot rely on '__origin__' in other type hints as it has changed from version + # to version. + origin = getattr(annotation, "__origin__", None) + return origin == typing.Union and ( + parameter_check_func is None or all(map(parameter_check_func, annotation.__args__)) + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/types.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/types.py new file mode 100644 index 0000000..b02a003 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/types.py @@ -0,0 +1,972 @@ +# +# 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. +# + +""" +Type-specific codes between pandas and PyArrow. Also contains some utils to correct +pandas instances during the type conversion. +""" +import datetime +import itertools +from typing import Any, Callable, Iterable, List, Optional, Union, TYPE_CHECKING + +from pyspark.sql.types import ( + cast, + BooleanType, + ByteType, + ShortType, + IntegerType, + IntegralType, + LongType, + FloatType, + DoubleType, + DecimalType, + StringType, + BinaryType, + DateType, + TimestampType, + TimestampNTZType, + DayTimeIntervalType, + ArrayType, + MapType, + StructType, + StructField, + NullType, + DataType, + UserDefinedType, + Row, + _create_row, +) +from pyspark.errors import PySparkTypeError, UnsupportedOperationException + +if TYPE_CHECKING: + import pandas as pd + import pyarrow as pa + + from pyspark.sql.pandas._typing import SeriesLike as PandasSeriesLike + + +def to_arrow_type(dt: DataType) -> "pa.DataType": + """Convert Spark data type to pyarrow type""" + from distutils.version import LooseVersion + import pyarrow as pa + + if type(dt) == BooleanType: + arrow_type = pa.bool_() + elif type(dt) == ByteType: + arrow_type = pa.int8() + elif type(dt) == ShortType: + arrow_type = pa.int16() + elif type(dt) == IntegerType: + arrow_type = pa.int32() + elif type(dt) == LongType: + arrow_type = pa.int64() + elif type(dt) == FloatType: + arrow_type = pa.float32() + elif type(dt) == DoubleType: + arrow_type = pa.float64() + elif type(dt) == DecimalType: + arrow_type = pa.decimal128(dt.precision, dt.scale) + elif type(dt) == StringType: + arrow_type = pa.string() + elif type(dt) == BinaryType: + arrow_type = pa.binary() + elif type(dt) == DateType: + arrow_type = pa.date32() + elif type(dt) == TimestampType: + # Timestamps should be in UTC, JVM Arrow timestamps require a timezone to be read + arrow_type = pa.timestamp("us", tz="UTC") + elif type(dt) == TimestampNTZType: + arrow_type = pa.timestamp("us", tz=None) + elif type(dt) == DayTimeIntervalType: + arrow_type = pa.duration("us") + elif type(dt) == ArrayType: + if type(dt.elementType) == StructType and LooseVersion(pa.__version__) < LooseVersion( + "2.0.0" + ): + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION", + message_parameters={"data_type": "Array of StructType"}, + ) + field = pa.field("element", to_arrow_type(dt.elementType), nullable=dt.containsNull) + arrow_type = pa.list_(field) + elif type(dt) == MapType: + if LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION", + message_parameters={"data_type": "MapType"}, + ) + key_field = pa.field("key", to_arrow_type(dt.keyType), nullable=False) + value_field = pa.field("value", to_arrow_type(dt.valueType), nullable=dt.valueContainsNull) + arrow_type = pa.map_(key_field, value_field) + elif type(dt) == StructType: + fields = [ + pa.field(field.name, to_arrow_type(field.dataType), nullable=field.nullable) + for field in dt + ] + arrow_type = pa.struct(fields) + elif type(dt) == NullType: + arrow_type = pa.null() + elif isinstance(dt, UserDefinedType): + arrow_type = to_arrow_type(dt.sqlType()) + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", + message_parameters={"data_type": str(dt)}, + ) + return arrow_type + + +def to_arrow_schema(schema: StructType) -> "pa.Schema": + """Convert a schema from Spark to Arrow""" + import pyarrow as pa + + fields = [ + pa.field(field.name, to_arrow_type(field.dataType), nullable=field.nullable) + for field in schema + ] + return pa.schema(fields) + + +def from_arrow_type(at: "pa.DataType", prefer_timestamp_ntz: bool = False) -> DataType: + """Convert pyarrow type to Spark data type.""" + from distutils.version import LooseVersion + import pyarrow as pa + import pyarrow.types as types + + spark_type: DataType + if types.is_boolean(at): + spark_type = BooleanType() + elif types.is_int8(at): + spark_type = ByteType() + elif types.is_int16(at): + spark_type = ShortType() + elif types.is_int32(at): + spark_type = IntegerType() + elif types.is_int64(at): + spark_type = LongType() + elif types.is_float32(at): + spark_type = FloatType() + elif types.is_float64(at): + spark_type = DoubleType() + elif types.is_decimal(at): + spark_type = DecimalType(precision=at.precision, scale=at.scale) + elif types.is_string(at): + spark_type = StringType() + elif types.is_large_string(at): + spark_type = StringType() + elif types.is_binary(at): + spark_type = BinaryType() + elif types.is_large_binary(at): + spark_type = BinaryType() + elif types.is_date32(at): + spark_type = DateType() + elif types.is_timestamp(at) and prefer_timestamp_ntz and at.tz is None: + spark_type = TimestampNTZType() + elif types.is_timestamp(at): + spark_type = TimestampType() + elif types.is_duration(at): + spark_type = DayTimeIntervalType() + elif types.is_list(at): + spark_type = ArrayType(from_arrow_type(at.value_type, prefer_timestamp_ntz)) + elif types.is_map(at): + if LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION", + message_parameters={"data_type": "MapType"}, + ) + spark_type = MapType( + from_arrow_type(at.key_type, prefer_timestamp_ntz), + from_arrow_type(at.item_type, prefer_timestamp_ntz), + ) + elif types.is_struct(at): + return StructType( + [ + StructField( + field.name, + from_arrow_type(field.type, prefer_timestamp_ntz), + nullable=field.nullable, + ) + for field in at + ] + ) + elif types.is_dictionary(at): + spark_type = from_arrow_type(at.value_type, prefer_timestamp_ntz) + elif types.is_null(at): + spark_type = NullType() + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", + message_parameters={"data_type": str(at)}, + ) + return spark_type + + +def from_arrow_schema(arrow_schema: "pa.Schema", prefer_timestamp_ntz: bool = False) -> StructType: + """Convert schema from Arrow to Spark.""" + return StructType( + [ + StructField( + field.name, + from_arrow_type(field.type, prefer_timestamp_ntz), + nullable=field.nullable, + ) + for field in arrow_schema + ] + ) + + +def _get_local_timezone() -> str: + """Get local timezone using pytz with environment variable, or dateutil. + + If there is a 'TZ' environment variable, pass it to pandas to use pytz and use it as timezone + string, otherwise use the special word 'dateutil/:' which means that pandas uses dateutil and + it reads system configuration to know the system local timezone. + + See also: + - https://github.com/pandas-dev/pandas/blob/0.19.x/pandas/tslib.pyx#L1753 + - https://github.com/dateutil/dateutil/blob/2.6.1/dateutil/tz/tz.py#L1338 + """ + import os + + return os.environ.get("TZ", "dateutil/:") + + +def _check_series_localize_timestamps(s: "PandasSeriesLike", timezone: str) -> "PandasSeriesLike": + """ + Convert timezone aware timestamps to timezone-naive in the specified timezone or local timezone. + + If the input series is not a timestamp series, then the same series is returned. If the input + series is a timestamp series, then a converted series is returned. + + Parameters + ---------- + s : pandas.Series + timezone : str + the timezone to convert. if None then use local timezone + + Returns + ------- + pandas.Series + `pandas.Series` that have been converted to tz-naive + """ + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + + from pandas.api.types import is_datetime64tz_dtype # type: ignore[attr-defined] + + tz = timezone or _get_local_timezone() + # TODO: handle nested timestamps, such as ArrayType(TimestampType())? + if is_datetime64tz_dtype(s.dtype): + return s.dt.tz_convert(tz).dt.tz_localize(None) + else: + return s + + +def _check_series_convert_timestamps_internal( + s: "PandasSeriesLike", timezone: str +) -> "PandasSeriesLike": + """ + Convert a tz-naive timestamp in the specified timezone or local timezone to UTC normalized for + Spark internal storage + + Parameters + ---------- + s : pandas.Series + timezone : str + the timezone to convert. if None then use local timezone + + Returns + ------- + pandas.Series + `pandas.Series` where if it is a timestamp, has been UTC normalized without a time zone + """ + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + + from pandas.api.types import ( # type: ignore[attr-defined] + is_datetime64_dtype, + is_datetime64tz_dtype, + ) + + # TODO: handle nested timestamps, such as ArrayType(TimestampType())? + if is_datetime64_dtype(s.dtype): + # When tz_localize a tz-naive timestamp, the result is ambiguous if the tz-naive + # timestamp is during the hour when the clock is adjusted backward during due to + # daylight saving time (dst). + # E.g., for America/New_York, the clock is adjusted backward on 2015-11-01 2:00 to + # 2015-11-01 1:00 from dst-time to standard time, and therefore, when tz_localize + # a tz-naive timestamp 2015-11-01 1:30 with America/New_York timezone, it can be either + # dst time (2015-01-01 1:30-0400) or standard time (2015-11-01 1:30-0500). + # + # Here we explicit choose to use standard time. This matches the default behavior of + # pytz. + # + # Here are some code to help understand this behavior: + # >>> import datetime + # >>> import pandas as pd + # >>> import pytz + # >>> + # >>> t = datetime.datetime(2015, 11, 1, 1, 30) + # >>> ts = pd.Series([t]) + # >>> tz = pytz.timezone('America/New_York') + # >>> + # >>> ts.dt.tz_localize(tz, ambiguous=True) + # 0 2015-11-01 01:30:00-04:00 + # dtype: datetime64[ns, America/New_York] + # >>> + # >>> ts.dt.tz_localize(tz, ambiguous=False) + # 0 2015-11-01 01:30:00-05:00 + # dtype: datetime64[ns, America/New_York] + # >>> + # >>> str(tz.localize(t)) + # '2015-11-01 01:30:00-05:00' + tz = timezone or _get_local_timezone() + return s.dt.tz_localize(tz, ambiguous=False).dt.tz_convert("UTC") + elif is_datetime64tz_dtype(s.dtype): + return s.dt.tz_convert("UTC") + else: + return s + + +def _check_series_convert_timestamps_localize( + s: "PandasSeriesLike", from_timezone: Optional[str], to_timezone: Optional[str] +) -> "PandasSeriesLike": + """ + Convert timestamp to timezone-naive in the specified timezone or local timezone + + Parameters + ---------- + s : pandas.Series + from_timezone : str + the timezone to convert from. if None then use local timezone + to_timezone : str + the timezone to convert to. if None then use local timezone + + Returns + ------- + pandas.Series + `pandas.Series` where if it is a timestamp, has been converted to tz-naive + """ + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + + import pandas as pd + from pandas.api.types import ( # type: ignore[attr-defined] + is_datetime64tz_dtype, + is_datetime64_dtype, + ) + + from_tz = from_timezone or _get_local_timezone() + to_tz = to_timezone or _get_local_timezone() + # TODO: handle nested timestamps, such as ArrayType(TimestampType())? + if is_datetime64tz_dtype(s.dtype): + return s.dt.tz_convert(to_tz).dt.tz_localize(None) + elif is_datetime64_dtype(s.dtype) and from_tz != to_tz: + # `s.dt.tz_localize('tzlocal()')` doesn't work properly when including NaT. + return cast( + "PandasSeriesLike", + s.apply( + lambda ts: ts.tz_localize(from_tz, ambiguous=False) + .tz_convert(to_tz) + .tz_localize(None) + if ts is not pd.NaT + else pd.NaT + ), + ) + else: + return s + + +def _check_series_convert_timestamps_local_tz( + s: "PandasSeriesLike", timezone: str +) -> "PandasSeriesLike": + """ + Convert timestamp to timezone-naive in the specified timezone or local timezone + + Parameters + ---------- + s : pandas.Series + timezone : str + the timezone to convert to. if None then use local timezone + + Returns + ------- + pandas.Series + `pandas.Series` where if it is a timestamp, has been converted to tz-naive + """ + return _check_series_convert_timestamps_localize(s, None, timezone) + + +def _check_series_convert_timestamps_tz_local( + s: "PandasSeriesLike", timezone: str +) -> "PandasSeriesLike": + """ + Convert timestamp to timezone-naive in the specified timezone or local timezone + + Parameters + ---------- + s : pandas.Series + timezone : str + the timezone to convert from. if None then use local timezone + + Returns + ------- + pandas.Series + `pandas.Series` where if it is a timestamp, has been converted to tz-naive + """ + return _check_series_convert_timestamps_localize(s, timezone, None) + + +def _convert_map_items_to_dict(s: "PandasSeriesLike") -> "PandasSeriesLike": + """ + Convert a series with items as list of (key, value), as made from an Arrow column of map type, + to dict for compatibility with non-arrow MapType columns. + :param s: pandas.Series of lists of (key, value) pairs + :return: pandas.Series of dictionaries + """ + return cast("PandasSeriesLike", s.apply(lambda m: None if m is None else {k: v for k, v in m})) + + +def _convert_dict_to_map_items(s: "PandasSeriesLike") -> "PandasSeriesLike": + """ + Convert a series of dictionaries to list of (key, value) pairs to match expected data + for Arrow column of map type. + :param s: pandas.Series of dictionaries + :return: pandas.Series of lists of (key, value) pairs + """ + return cast("PandasSeriesLike", s.apply(lambda d: list(d.items()) if d is not None else None)) + + +def _to_corrected_pandas_type(dt: DataType) -> Optional[Any]: + """ + When converting Spark SQL records to Pandas `pandas.DataFrame`, the inferred data type + may be wrong. This method gets the corrected data type for Pandas if that type may be + inferred incorrectly. + """ + import numpy as np + + if type(dt) == ByteType: + return np.int8 + elif type(dt) == ShortType: + return np.int16 + elif type(dt) == IntegerType: + return np.int32 + elif type(dt) == LongType: + return np.int64 + elif type(dt) == FloatType: + return np.float32 + elif type(dt) == DoubleType: + return np.float64 + elif type(dt) == BooleanType: + return bool + elif type(dt) == TimestampType: + return np.dtype("datetime64[ns]") + elif type(dt) == TimestampNTZType: + return np.dtype("datetime64[ns]") + elif type(dt) == DayTimeIntervalType: + return np.dtype("timedelta64[ns]") + else: + return None + + +def _create_converter_to_pandas( + data_type: DataType, + nullable: bool = True, + *, + timezone: Optional[str] = None, + struct_in_pandas: Optional[str] = None, + error_on_duplicated_field_names: bool = True, + timestamp_utc_localized: bool = True, + ndarray_as_list: bool = False, +) -> Callable[["pd.Series"], "pd.Series"]: + """ + Create a converter of pandas Series that is created from Spark's Python objects, + or `pyarrow.Table.to_pandas` method. + + Parameters + ---------- + data_type : :class:`DataType` + The data type corresponding to the pandas Series to be converted. + nullable : bool, optional + Whether the column is nullable or not. (default ``True``) + timezone : str, optional + The timezone to convert from. If there is a timestamp type, it's required. + struct_in_pandas : str, optional + How to handle struct type. If there is a struct type, it's required. + When ``row``, :class:`Row` object will be used. + When ``dict``, :class:`dict` will be used. If there are duplicated field names, + The fields will be suffixed, like `a_0`, `a_1`. + Must be one of: ``row``, ``dict``. + error_on_duplicated_field_names : bool, optional + Whether raise an exception when there are duplicated field names. + (default ``True``) + timestamp_utc_localized : bool, optional + Whether the timestamp values are localized to UTC or not. + The timestamp values from Arrow are localized to UTC, + whereas the ones from `df.collect()` are localized to the local timezone. + ndarray_as_list : bool, optional + Whether `np.ndarray` is converted to a list or not (default ``False``). + + Returns + ------- + The converter of `pandas.Series` + """ + import numpy as np + import pandas as pd + from pandas.core.dtypes.common import is_datetime64tz_dtype + + pandas_type = _to_corrected_pandas_type(data_type) + + if pandas_type is not None: + # SPARK-21766: if an integer field is nullable and has null values, it can be + # inferred by pandas as a float column. If we convert the column with NaN back + # to integer type e.g., np.int16, we will hit an exception. So we use the + # pandas-inferred float type, rather than the corrected type from the schema + # in this case. + if isinstance(data_type, IntegralType) and nullable: + + def correct_dtype(pser: pd.Series) -> pd.Series: + if pser.isnull().any(): + return pser.astype(np.float64, copy=False) + else: + return pser.astype(pandas_type, copy=False) + + elif isinstance(data_type, BooleanType) and nullable: + + def correct_dtype(pser: pd.Series) -> pd.Series: + if pser.isnull().any(): + return pser.astype(object, copy=False) + else: + return pser.astype(pandas_type, copy=False) + + elif isinstance(data_type, TimestampType): + assert timezone is not None + + def correct_dtype(pser: pd.Series) -> pd.Series: + if not is_datetime64tz_dtype(pser.dtype): + pser = pser.astype(pandas_type, copy=False) + return _check_series_convert_timestamps_local_tz(pser, timezone=cast(str, timezone)) + + else: + + def correct_dtype(pser: pd.Series) -> pd.Series: + return pser.astype(pandas_type, copy=False) + + return correct_dtype + + def _converter( + dt: DataType, _struct_in_pandas: Optional[str], _ndarray_as_list: bool + ) -> Optional[Callable[[Any], Any]]: + + if isinstance(dt, ArrayType): + _element_conv = _converter(dt.elementType, _struct_in_pandas, _ndarray_as_list) + + if _ndarray_as_list: + if _element_conv is None: + _element_conv = lambda x: x # noqa: E731 + + def convert_array_ndarray_as_list(value: Any) -> Any: + if value is None: + return None + else: + # In Arrow Python UDF, ArrayType is converted to `np.ndarray` + # whereas a list is expected. + return [_element_conv(v) for v in value] # type: ignore[misc] + + return convert_array_ndarray_as_list + else: + if _element_conv is None: + return None + + def convert_array_ndarray_as_ndarray(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, np.ndarray): + # `pyarrow.Table.to_pandas` uses `np.ndarray`. + return np.array([_element_conv(v) for v in value]) # type: ignore[misc] + else: + assert isinstance(value, list) + # otherwise, `list` should be used. + return [_element_conv(v) for v in value] # type: ignore[misc] + + return convert_array_ndarray_as_ndarray + + elif isinstance(dt, MapType): + _key_conv = _converter(dt.keyType, _struct_in_pandas, _ndarray_as_list) or (lambda x: x) + _value_conv = _converter(dt.valueType, _struct_in_pandas, _ndarray_as_list) or ( + lambda x: x + ) + + def convert_map(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, list): + # `pyarrow.Table.to_pandas` uses `list` of key-value tuple. + return {_key_conv(k): _value_conv(v) for k, v in value} + else: + assert isinstance(value, dict) + # otherwise, `dict` should be used. + return {_key_conv(k): _value_conv(v) for k, v in value.items()} + + return convert_map + + elif isinstance(dt, StructType): + assert _struct_in_pandas is not None + + field_names = dt.names + + if error_on_duplicated_field_names and len(set(field_names)) != len(field_names): + raise UnsupportedOperationException( + error_class="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT", + message_parameters={"field_names": str(field_names)}, + ) + + dedup_field_names = _dedup_names(field_names) + + field_convs = [ + _converter(f.dataType, _struct_in_pandas, _ndarray_as_list) or (lambda x: x) + for f in dt.fields + ] + + if _struct_in_pandas == "row": + + def convert_struct_as_row(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + # `pyarrow.Table.to_pandas` uses `dict`. + _values = [ + field_convs[i](value.get(name, None)) + for i, name in enumerate(dedup_field_names) + ] + return _create_row(field_names, _values) + else: + assert isinstance(value, Row) + # otherwise, `Row` should be used. + _values = [field_convs[i](value[i]) for i, name in enumerate(value)] + return _create_row(field_names, _values) + + return convert_struct_as_row + + elif _struct_in_pandas == "dict": + + def convert_struct_as_dict(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + # `pyarrow.Table.to_pandas` uses `dict`. + return { + name: field_convs[i](value.get(name, None)) + for i, name in enumerate(dedup_field_names) + } + else: + assert isinstance(value, Row) + # otherwise, `Row` should be used. + return { + dedup_field_names[i]: field_convs[i](v) for i, v in enumerate(value) + } + + return convert_struct_as_dict + + else: + raise ValueError(f"Unknown value for `struct_in_pandas`: {_struct_in_pandas}") + + elif isinstance(dt, TimestampType): + assert timezone is not None + + local_tz: Union[datetime.tzinfo, str] = ( + datetime.timezone.utc if timestamp_utc_localized else _get_local_timezone() + ) + + def convert_timestamp(value: Any) -> Any: + if value is None: + return None + else: + if isinstance(value, datetime.datetime) and value.tzinfo is not None: + ts = pd.Timestamp(value) + else: + ts = pd.Timestamp(value).tz_localize(local_tz) + return ts.tz_convert(timezone).tz_localize(None) + + return convert_timestamp + + elif isinstance(dt, TimestampNTZType): + + def convert_timestamp_ntz(value: Any) -> Any: + if value is None: + return None + else: + return pd.Timestamp(value) + + return convert_timestamp_ntz + + elif isinstance(dt, UserDefinedType): + udt: UserDefinedType = dt + + conv = _converter(udt.sqlType(), _struct_in_pandas="row", _ndarray_as_list=True) or ( + lambda x: x + ) + + def convert_udt(value: Any) -> Any: + if value is None: + return None + elif hasattr(value, "__UDT__"): + assert isinstance(value.__UDT__, type(udt)) + return value + else: + return udt.deserialize(conv(value)) + + return convert_udt + + else: + return None + + conv = _converter(data_type, struct_in_pandas, ndarray_as_list) + if conv is not None: + return lambda pser: pser.apply(conv) # type: ignore[return-value] + else: + return lambda pser: pser + + +def _create_converter_from_pandas( + data_type: DataType, + *, + timezone: Optional[str], + error_on_duplicated_field_names: bool = True, + ignore_unexpected_complex_type_values: bool = False, +) -> Callable[["pd.Series"], "pd.Series"]: + """ + Create a converter of pandas Series to create Spark DataFrame with Arrow optimization. + + Parameters + ---------- + data_type : :class:`DataType` + The data type corresponding to the pandas Series to be converted. + timezone : str, optional + The timezone to convert from. If there is a timestamp type, it's required. + error_on_duplicated_field_names : bool, optional + Whether raise an exception when there are duplicated field names. + (default ``True``) + ignore_unexpected_complex_type_values : bool, optional + Whether ignore the case where unexpected values are given for complex types. + If ``False``, each complex type expects: + + * array type: :class:`Iterable` + * map type: :class:`dict` + * struct type: :class:`dict` or :class:`tuple` + + and raise an AssertionError when the given value is not the expected type. + If ``True``, just ignore and return the give value. + (default ``False``) + + Returns + ------- + The converter of `pandas.Series` + """ + import pandas as pd + + if isinstance(data_type, TimestampType): + assert timezone is not None + + def correct_timestamp(pser: pd.Series) -> pd.Series: + return _check_series_convert_timestamps_internal(pser, cast(str, timezone)) + + return correct_timestamp + + def _converter(dt: DataType) -> Optional[Callable[[Any], Any]]: + + if isinstance(dt, ArrayType): + _element_conv = _converter(dt.elementType) or (lambda x: x) + + if ignore_unexpected_complex_type_values: + + def convert_array(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, Iterable): + return [_element_conv(v) for v in value] + else: + return value + + else: + + def convert_array(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, Iterable) + return [_element_conv(v) for v in value] + + return convert_array + + elif isinstance(dt, MapType): + _key_conv = _converter(dt.keyType) or (lambda x: x) + _value_conv = _converter(dt.valueType) or (lambda x: x) + + if ignore_unexpected_complex_type_values: + + def convert_map(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + return [(_key_conv(k), _value_conv(v)) for k, v in value.items()] + else: + return value + + else: + + def convert_map(value: Any) -> Any: + if value is None: + return None + else: + assert isinstance(value, dict) + return [(_key_conv(k), _value_conv(v)) for k, v in value.items()] + + return convert_map + + elif isinstance(dt, StructType): + + field_names = dt.names + + if error_on_duplicated_field_names and len(set(field_names)) != len(field_names): + raise UnsupportedOperationException( + error_class="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT", + message_parameters={"field_names": str(field_names)}, + ) + + dedup_field_names = _dedup_names(field_names) + + field_convs = [_converter(f.dataType) or (lambda x: x) for f in dt.fields] + + if ignore_unexpected_complex_type_values: + + def convert_struct(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + return { + dedup_field_names[i]: field_convs[i](value.get(key, None)) + for i, key in enumerate(field_names) + } + elif isinstance(value, tuple): + return { + dedup_field_names[i]: field_convs[i](v) for i, v in enumerate(value) + } + else: + return value + + else: + + def convert_struct(value: Any) -> Any: + if value is None: + return None + elif isinstance(value, dict): + return { + dedup_field_names[i]: field_convs[i](value.get(key, None)) + for i, key in enumerate(field_names) + } + else: + assert isinstance(value, tuple) + return { + dedup_field_names[i]: field_convs[i](v) for i, v in enumerate(value) + } + + return convert_struct + + elif isinstance(dt, TimestampType): + assert timezone is not None + + def convert_timestamp(value: Any) -> Any: + if value is None: + return None + else: + if isinstance(value, datetime.datetime) and value.tzinfo is not None: + ts = pd.Timstamp(value) + else: + ts = pd.Timestamp(value).tz_localize(timezone) + return ts.to_pydatetime() + + return convert_timestamp + + elif isinstance(dt, UserDefinedType): + udt: UserDefinedType = dt + + conv = _converter(udt.sqlType()) or (lambda x: x) + + def convert_udt(value: Any) -> Any: + if value is None: + return None + else: + return conv(udt.serialize(value)) + + return convert_udt + + return None + + conv = _converter(data_type) + if conv is not None: + return lambda pser: pser.apply(conv) # type: ignore[return-value] + else: + return lambda pser: pser + + +def _dedup_names(names: List[str]) -> List[str]: + if len(set(names)) == len(names): + return names + else: + + def _gen_dedup(_name: str) -> Callable[[], str]: + _i = itertools.count() + return lambda: f"{_name}_{next(_i)}" + + def _gen_identity(_name: str) -> Callable[[], str]: + return lambda: _name + + gen_new_name = { + name: _gen_dedup(name) if len(list(group)) > 1 else _gen_identity(name) + for name, group in itertools.groupby(sorted(names)) + } + return [gen_new_name[name]() for name in names] + + +def _deduplicate_field_names(dt: DataType) -> DataType: + if isinstance(dt, StructType): + dedup_field_names = _dedup_names(dt.names) + + return StructType( + [ + StructField( + dedup_field_names[i], + _deduplicate_field_names(field.dataType), + nullable=field.nullable, + ) + for i, field in enumerate(dt.fields) + ] + ) + elif isinstance(dt, ArrayType): + return ArrayType(_deduplicate_field_names(dt.elementType), containsNull=dt.containsNull) + elif isinstance(dt, MapType): + return MapType( + _deduplicate_field_names(dt.keyType), + _deduplicate_field_names(dt.valueType), + valueContainsNull=dt.valueContainsNull, + ) + else: + return dt diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/utils.py new file mode 100644 index 0000000..67a571f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/pandas/utils.py @@ -0,0 +1,86 @@ +# +# 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. +# + + +def require_minimum_pandas_version() -> None: + """Raise ImportError if minimum version of Pandas is not installed""" + # TODO(HyukjinKwon): Relocate and deduplicate the version specification. + minimum_pandas_version = "1.0.5" + + from distutils.version import LooseVersion + + try: + import pandas + + have_pandas = True + except ImportError as error: + have_pandas = False + raised_error = error + if not have_pandas: + raise ImportError( + "Pandas >= %s must be installed; however, " "it was not found." % minimum_pandas_version + ) from raised_error + if LooseVersion(pandas.__version__) < LooseVersion(minimum_pandas_version): + raise ImportError( + "Pandas >= %s must be installed; however, " + "your version was %s." % (minimum_pandas_version, pandas.__version__) + ) + + +def require_minimum_pyarrow_version() -> None: + """Raise ImportError if minimum version of pyarrow is not installed""" + # TODO(HyukjinKwon): Relocate and deduplicate the version specification. + minimum_pyarrow_version = "4.0.0" + + from distutils.version import LooseVersion + import os + + try: + import pyarrow + + have_arrow = True + except ImportError as error: + have_arrow = False + raised_error = error + if not have_arrow: + raise ImportError( + "PyArrow >= %s must be installed; however, " + "it was not found." % minimum_pyarrow_version + ) from raised_error + if LooseVersion(pyarrow.__version__) < LooseVersion(minimum_pyarrow_version): + raise ImportError( + "PyArrow >= %s must be installed; however, " + "your version was %s." % (minimum_pyarrow_version, pyarrow.__version__) + ) + if os.environ.get("ARROW_PRE_0_15_IPC_FORMAT", "0") == "1": + raise RuntimeError( + "Arrow legacy IPC format is not supported in PySpark, " + "please unset ARROW_PRE_0_15_IPC_FORMAT" + ) + + +def pyarrow_version_less_than_minimum(minimum_pyarrow_version: str) -> bool: + """Return False if the installed pyarrow version is less than minimum_pyarrow_version + or if pyarrow is not installed.""" + from distutils.version import LooseVersion + + try: + import pyarrow + except ImportError: + return False + + return LooseVersion(pyarrow.__version__) < LooseVersion(minimum_pyarrow_version) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/__init__.py new file mode 100644 index 0000000..ac530a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/__init__.py @@ -0,0 +1,18 @@ +# +# 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. +# + +__all__ = ["functions"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/functions.py new file mode 100644 index 0000000..acb1a17 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/protobuf/functions.py @@ -0,0 +1,334 @@ +# +# 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. +# + +""" +A collections of builtin protobuf functions +""" + + +from typing import Dict, Optional, TYPE_CHECKING, cast + +from py4j.java_gateway import JVMView + +from pyspark.sql.column import Column, _to_java_column +from pyspark.sql.utils import get_active_spark_context, try_remote_protobuf_functions +from pyspark.util import _print_missing_jar + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName + + +@try_remote_protobuf_functions +def from_protobuf( + data: "ColumnOrName", + messageName: str, + descFilePath: Optional[str] = None, + options: Optional[Dict[str, str]] = None, + binaryDescriptorSet: Optional[bytes] = None, +) -> Column: + """ + Converts a binary column of Protobuf format into its corresponding catalyst value. + The Protobuf definition is provided in one of these ways: + + - Protobuf descriptor file: E.g. a descriptor file created with + `protoc --include_imports --descriptor_set_out=abc.desc abc.proto` + - Protobuf descriptor as binary: Rather than file path as in previous option, + we can provide the binary content of the file. This allows flexibility in how the + descriptor set is created and fetched. + - Jar containing Protobuf Java class: The jar containing Java class should be shaded. + Specifically, `com.google.protobuf.*` should be shaded to + `org.sparkproject.spark_protobuf.protobuf.*`. + https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from + Protobuf files. The jar file can be added with spark-submit option --jars. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports `binaryDescriptorSet` arg to pass binary descriptor directly. + Supports Spark Connect. + + Parameters + ---------- + data : :class:`~pyspark.sql.Column` or str + the binary column. + messageName: str, optional + the protobuf message name to look for in descriptor file, or + The Protobuf class name when descFilePath parameter is not set. + E.g. `com.example.protos.ExampleEvent`. + descFilePath : str, optional + The Protobuf descriptor file. + options : dict, optional + options to control how the protobuf record is parsed. + binaryDescriptorSet: bytes, optional + The Protobuf `FileDescriptorSet` serialized as binary. + + Notes + ----- + Protobuf functionality is provided as an pluggable external module. + + Examples + -------- + >>> import tempfile + >>> data = [("1", (2, "Alice", 109200))] + >>> ddl_schema = "key STRING, value STRUCT" + >>> df = spark.createDataFrame(data, ddl_schema) + >>> desc_hex = str('0ACE010A41636F6E6E6563746F722F70726F746F6275662F7372632F746573742F726' + ... '5736F75726365732F70726F746F6275662F7079737061726B5F746573742E70726F746F121D6F72672E61' + ... '70616368652E737061726B2E73716C2E70726F746F627566224B0A0D53696D706C654D657373616765121' + ... '00A03616765180120012805520361676512120A046E616D6518022001280952046E616D6512140A057363' + ... '6F7265180320012803520573636F72654215421353696D706C654D65737361676550726F746F736206707' + ... '26F746F33') + >>> # Writing a protobuf description into a file, generated by using + >>> # connector/protobuf/src/test/resources/protobuf/pyspark_test.proto file + >>> with tempfile.TemporaryDirectory() as tmp_dir: + ... desc_file_path = "%s/pyspark_test.desc" % tmp_dir + ... with open(desc_file_path, "wb") as f: + ... _ = f.write(bytearray.fromhex(desc_hex)) + ... f.flush() + ... message_name = 'SimpleMessage' + ... proto_df = df.select( + ... to_protobuf(df.value, message_name, desc_file_path).alias("value")) + ... proto_df.show(truncate=False) + ... proto_df_1 = proto_df.select( # With file name for descriptor + ... from_protobuf(proto_df.value, message_name, desc_file_path).alias("value")) + ... proto_df_1.show(truncate=False) + ... proto_df_2 = proto_df.select( # With binary for descriptor + ... from_protobuf(proto_df.value, message_name, + ... binaryDescriptorSet = bytearray.fromhex(desc_hex)) + ... .alias("value")) + ... proto_df_2.show(truncate=False) + +----------------------------------------+ + |value | + +----------------------------------------+ + |[08 02 12 05 41 6C 69 63 65 18 90 D5 06]| + +----------------------------------------+ + +------------------+ + |value | + +------------------+ + |{2, Alice, 109200}| + +------------------+ + +------------------+ + |value | + +------------------+ + |{2, Alice, 109200}| + +------------------+ + >>> data = [([(1668035962, 2020)])] + >>> ddl_schema = "value struct" + >>> df = spark.createDataFrame(data, ddl_schema) + >>> message_class_name = "org.sparkproject.spark_protobuf.protobuf.Timestamp" + >>> to_proto_df = df.select(to_protobuf(df.value, message_class_name).alias("value")) + >>> from_proto_df = to_proto_df.select( + ... from_protobuf(to_proto_df.value, message_class_name).alias("value")) + >>> from_proto_df.show(truncate=False) + +------------------+ + |value | + +------------------+ + |{1668035962, 2020}| + +------------------+ + """ + + sc = get_active_spark_context() + try: + binary_proto = None + if binaryDescriptorSet is not None: + binary_proto = binaryDescriptorSet + elif descFilePath is not None: + binary_proto = _read_descriptor_set_file(descFilePath) + if binary_proto is not None: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.protobuf.functions.from_protobuf( + _to_java_column(data), messageName, binary_proto, options or {} + ) + else: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.protobuf.functions.from_protobuf( + _to_java_column(data), messageName, options or {} + ) + except TypeError as e: + if str(e) == "'JavaPackage' object is not callable": + _print_missing_jar("Protobuf", "protobuf", "protobuf", sc.version) + raise + return Column(jc) + + +@try_remote_protobuf_functions +def to_protobuf( + data: "ColumnOrName", + messageName: str, + descFilePath: Optional[str] = None, + options: Optional[Dict[str, str]] = None, + binaryDescriptorSet: Optional[bytes] = None, +) -> Column: + """ + Converts a column into binary of protobuf format. The Protobuf definition is provided in one + of these ways: + + - Protobuf descriptor file: E.g. a descriptor file created with + `protoc --include_imports --descriptor_set_out=abc.desc abc.proto` + - Protobuf descriptor as binary: Rather than file path as in previous option, + we can provide the binary content of the file. This allows flexibility in how the + descriptor set is created and fetched. + - Jar containing Protobuf Java class: The jar containing Java class should be shaded. + Specifically, `com.google.protobuf.*` should be shaded to + `org.sparkproject.spark_protobuf.protobuf.*`. + https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from + Protobuf files. The jar file can be added with spark-submit option --jars. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports `binaryDescriptorSet` arg to pass binary descriptor directly. + Supports Spark Connect. + + Parameters + ---------- + data : :class:`~pyspark.sql.Column` or str + the data column. + messageName: str, optional + the protobuf message name to look for in descriptor file, or + The Protobuf class name when descFilePath parameter is not set. + E.g. `com.example.protos.ExampleEvent`. + descFilePath : str, optional + the Protobuf descriptor file. + options : dict, optional + binaryDescriptorSet: bytes, optional + The Protobuf `FileDescriptorSet` serialized as binary. + + Notes + ----- + Protobuf functionality is provided as a pluggable external module + + Examples + -------- + >>> import tempfile + >>> data = [([(2, "Alice", 13093020)])] + >>> ddl_schema = "value struct" + >>> df = spark.createDataFrame(data, ddl_schema) + >>> desc_hex = str('0ACE010A41636F6E6E6563746F722F70726F746F6275662F7372632F746573742F726' + ... '5736F75726365732F70726F746F6275662F7079737061726B5F746573742E70726F746F121D6F72672E61' + ... '70616368652E737061726B2E73716C2E70726F746F627566224B0A0D53696D706C654D657373616765121' + ... '00A03616765180120012805520361676512120A046E616D6518022001280952046E616D6512140A057363' + ... '6F7265180320012803520573636F72654215421353696D706C654D65737361676550726F746F736206707' + ... '26F746F33') + >>> # Writing a protobuf description into a file, generated by using + >>> # connector/protobuf/src/test/resources/protobuf/pyspark_test.proto file + >>> with tempfile.TemporaryDirectory() as tmp_dir: + ... desc_file_path = "%s/pyspark_test.desc" % tmp_dir + ... with open(desc_file_path, "wb") as f: + ... _ = f.write(bytearray.fromhex(desc_hex)) + ... f.flush() + ... message_name = 'SimpleMessage' + ... proto_df = df.select( # With file name for descriptor + ... to_protobuf(df.value, message_name, desc_file_path).alias("suite")) + ... proto_df.show(truncate=False) + ... proto_df_2 = df.select( # With binary for descriptor + ... to_protobuf(df.value, message_name, + ... binaryDescriptorSet=bytearray.fromhex(desc_hex)) + ... .alias("suite")) + ... proto_df_2.show(truncate=False) + +-------------------------------------------+ + |suite | + +-------------------------------------------+ + |[08 02 12 05 41 6C 69 63 65 18 9C 91 9F 06]| + +-------------------------------------------+ + +-------------------------------------------+ + |suite | + +-------------------------------------------+ + |[08 02 12 05 41 6C 69 63 65 18 9C 91 9F 06]| + +-------------------------------------------+ + >>> data = [([(1668035962, 2020)])] + >>> ddl_schema = "value struct" + >>> df = spark.createDataFrame(data, ddl_schema) + >>> message_class_name = "org.sparkproject.spark_protobuf.protobuf.Timestamp" + >>> proto_df = df.select(to_protobuf(df.value, message_class_name).alias("suite")) + >>> proto_df.show(truncate=False) + +----------------------------+ + |suite | + +----------------------------+ + |[08 FA EA B0 9B 06 10 E4 0F]| + +----------------------------+ + """ + + sc = get_active_spark_context() + try: + binary_proto = None + if binaryDescriptorSet is not None: + binary_proto = binaryDescriptorSet + elif descFilePath is not None: + binary_proto = _read_descriptor_set_file(descFilePath) + if binary_proto is not None: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.protobuf.functions.to_protobuf( + _to_java_column(data), messageName, binary_proto, options or {} + ) + else: + jc = cast(JVMView, sc._jvm).org.apache.spark.sql.protobuf.functions.to_protobuf( + _to_java_column(data), messageName, options or {} + ) + + except TypeError as e: + if str(e) == "'JavaPackage' object is not callable": + _print_missing_jar("Protobuf", "protobuf", "protobuf", sc.version) + raise + return Column(jc) + + +def _read_descriptor_set_file(filePath: str) -> bytes: + # TODO(SPARK-43847): Throw structured errors like "PROTOBUF_DESCRIPTOR_FILE_NOT_FOUND" etc. + with open(filePath, "rb") as f: + return f.read() + + +def _test() -> None: + import os + import sys + from pyspark.testing.utils import search_jar + + protobuf_jar = search_jar("connector/protobuf", "spark-protobuf-assembly-", "spark-protobuf") + if protobuf_jar is None: + print( + "Skipping all Protobuf Python tests as the optional Protobuf project was " + "not compiled into a JAR. To run these tests, " + "you need to build Spark with 'build/sbt package' or " + "'build/mvn package' before running this test." + ) + sys.exit(0) + else: + existing_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + jars_args = "--jars %s" % protobuf_jar + os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) + + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.protobuf.functions + + globs = pyspark.sql.protobuf.functions.__dict__.copy() + spark = ( + SparkSession.builder.master("local[2]") + .appName("sql.protobuf.functions tests") + .getOrCreate() + ) + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.protobuf.functions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/readwriter.py new file mode 100644 index 0000000..cfac8fd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/readwriter.py @@ -0,0 +1,2159 @@ +# +# 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 sys +from typing import cast, overload, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING, Union + +from py4j.java_gateway import JavaClass, JavaObject + +from pyspark import RDD, since +from pyspark.sql.column import _to_seq, _to_java_column, Column +from pyspark.sql.types import StructType +from pyspark.sql import utils +from pyspark.sql.utils import to_str +from pyspark.errors import PySparkTypeError, PySparkValueError + +if TYPE_CHECKING: + from pyspark.sql._typing import OptionalPrimitiveType, ColumnOrName + from pyspark.sql.session import SparkSession + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.streaming import StreamingQuery + +__all__ = ["DataFrameReader", "DataFrameWriter", "DataFrameWriterV2"] + +PathOrPaths = Union[str, List[str]] +TupleOrListOfString = Union[List[str], Tuple[str, ...]] + + +class OptionUtils: + def _set_opts( + self, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """ + Set named options (filter out those the value is None) + """ + if schema is not None: + self.schema(schema) # type: ignore[attr-defined] + for k, v in options.items(): + if v is not None: + self.option(k, v) # type: ignore[attr-defined] + + +class DataFrameReader(OptionUtils): + """ + Interface used to load a :class:`DataFrame` from external storage systems + (e.g. file systems, key-value stores, etc). Use :attr:`SparkSession.read` + to access this. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, spark: "SparkSession"): + self._jreader = spark._jsparkSession.read() + self._spark = spark + + def _df(self, jdf: JavaObject) -> "DataFrame": + from pyspark.sql.dataframe import DataFrame + + return DataFrame(jdf, self._spark) + + def format(self, source: str) -> "DataFrameReader": + """Specifies the input data source format. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + source : str + string, name of the data source, e.g. 'json', 'parquet'. + + Examples + -------- + >>> spark.read.format('json') + <...readwriter.DataFrameReader object ...> + + Write a DataFrame into a JSON file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Read the JSON file as a DataFrame. + ... spark.read.format('json').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self._jreader = self._jreader.format(source) + return self + + def schema(self, schema: Union[StructType, str]) -> "DataFrameReader": + """Specifies the input schema. + + Some data sources (e.g. JSON) can infer the input schema automatically from data. + By specifying the schema here, the underlying data source can skip the schema + inference step, and thus speed up data loading. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + schema : :class:`pyspark.sql.types.StructType` or str + a :class:`pyspark.sql.types.StructType` object or a DDL-formatted string + (For example ``col0 INT, col1 DOUBLE``). + + Examples + -------- + >>> spark.read.schema("col0 INT, col1 DOUBLE") + <...readwriter.DataFrameReader object ...> + + Specify the schema with reading a CSV file. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... spark.read.schema("col0 INT, col1 DOUBLE").format("csv").load(d).printSchema() + root + |-- col0: integer (nullable = true) + |-- col1: double (nullable = true) + """ + from pyspark.sql import SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + if isinstance(schema, StructType): + jschema = spark._jsparkSession.parseDataType(schema.json()) + self._jreader = self._jreader.schema(jschema) + elif isinstance(schema, str): + self._jreader = self._jreader.schema(schema) + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_STRUCT", + message_parameters={ + "arg_name": "schema", + "arg_type": type(schema).__name__, + }, + ) + return self + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameReader": + """ + Adds an input option for the underlying data source. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + key : str + The key for the option to set. + value + The value for the option to set. + + Examples + -------- + >>> spark.read.option("key", "value") + <...readwriter.DataFrameReader object ...> + + Specify the option 'nullValue' with reading a CSV file. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file + ... df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]) + ... df.write.mode("overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon'. + ... spark.read.schema(df.schema).option( + ... "nullValue", "Hyukjin Kwon").format('csv').load(d).show() + +---+----+ + |age|name| + +---+----+ + |100|NULL| + +---+----+ + """ + self._jreader = self._jreader.option(key, to_str(value)) + return self + + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameReader": + """ + Adds input options for the underlying data source. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + **options : dict + The dictionary of string keys and prmitive-type values. + + Examples + -------- + >>> spark.read.option("key", "value") + <...readwriter.DataFrameReader object ...> + + Specify the option 'nullValue' and 'header' with reading a CSV file. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file with a header. + ... df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]) + ... df.write.option("header", True).mode("overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon', + ... # and 'header' option set to `True`. + ... spark.read.options( + ... nullValue="Hyukjin Kwon", + ... header=True + ... ).format('csv').load(d).show() + +---+----+ + |age|name| + +---+----+ + |100|NULL| + +---+----+ + """ + for k in options: + self._jreader = self._jreader.option(k, to_str(options[k])) + return self + + def load( + self, + path: Optional[PathOrPaths] = None, + format: Optional[str] = None, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> "DataFrame": + """Loads data from a data source and returns it as a :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str or list, optional + optional string or a list of string for file-system backed data sources. + format : str, optional + optional string for format of the data source. Default to 'parquet'. + schema : :class:`pyspark.sql.types.StructType` or str, optional + optional :class:`pyspark.sql.types.StructType` for the input schema + or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + **options : dict + all other string options + + Examples + -------- + Load a CSV file with format, schema and options specified. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file with a header + ... df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]) + ... df.write.option("header", True).mode("overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon', + ... # and 'header' option set to `True`. + ... df = spark.read.load( + ... d, schema=df.schema, format="csv", nullValue="Hyukjin Kwon", header=True) + ... df.printSchema() + ... df.show() + root + |-- age: long (nullable = true) + |-- name: string (nullable = true) + +---+----+ + |age|name| + +---+----+ + |100|NULL| + +---+----+ + """ + if format is not None: + self.format(format) + if schema is not None: + self.schema(schema) + self.options(**options) + if isinstance(path, str): + return self._df(self._jreader.load(path)) + elif path is not None: + if type(path) != list: + path = [path] # type: ignore[list-item] + assert self._spark._sc._jvm is not None + return self._df(self._jreader.load(self._spark._sc._jvm.PythonUtils.toSeq(path))) + else: + return self._df(self._jreader.load()) + + def json( + self, + path: Union[str, List[str], RDD[str]], + schema: Optional[Union[StructType, str]] = None, + primitivesAsString: Optional[Union[bool, str]] = None, + prefersDecimal: Optional[Union[bool, str]] = None, + allowComments: Optional[Union[bool, str]] = None, + allowUnquotedFieldNames: Optional[Union[bool, str]] = None, + allowSingleQuotes: Optional[Union[bool, str]] = None, + allowNumericLeadingZero: Optional[Union[bool, str]] = None, + allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + allowUnquotedControlChars: Optional[Union[bool, str]] = None, + lineSep: Optional[str] = None, + samplingRatio: Optional[Union[float, str]] = None, + dropFieldIfAllNull: Optional[Union[bool, str]] = None, + encoding: Optional[str] = None, + locale: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + allowNonNumericNumbers: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """ + Loads JSON files and returns the results as a :class:`DataFrame`. + + `JSON Lines `_ (newline-delimited JSON) is supported by default. + For JSON (one record per file), set the ``multiLine`` parameter to ``true``. + + If the ``schema`` parameter is not specified, this function goes + through the input once to determine the input schema. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str, list or :class:`RDD` + string represents path to the JSON dataset, or a list of paths, + or RDD of Strings storing JSON objects. + schema : :class:`pyspark.sql.types.StructType` or str, optional + an optional :class:`pyspark.sql.types.StructType` for the input schema or + a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a JSON file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Read the JSON file as a DataFrame. + ... spark.read.json(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self._set_opts( + schema=schema, + primitivesAsString=primitivesAsString, + prefersDecimal=prefersDecimal, + allowComments=allowComments, + allowUnquotedFieldNames=allowUnquotedFieldNames, + allowSingleQuotes=allowSingleQuotes, + allowNumericLeadingZero=allowNumericLeadingZero, + allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + multiLine=multiLine, + allowUnquotedControlChars=allowUnquotedControlChars, + lineSep=lineSep, + samplingRatio=samplingRatio, + dropFieldIfAllNull=dropFieldIfAllNull, + encoding=encoding, + locale=locale, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + allowNonNumericNumbers=allowNonNumericNumbers, + ) + if isinstance(path, str): + path = [path] + if type(path) == list: + assert self._spark._sc._jvm is not None + return self._df(self._jreader.json(self._spark._sc._jvm.PythonUtils.toSeq(path))) + elif isinstance(path, RDD): + + def func(iterator: Iterable) -> Iterable: + for x in iterator: + if not isinstance(x, str): + x = str(x) + if isinstance(x, str): + x = x.encode("utf-8") + yield x + + keyed = path.mapPartitions(func) + keyed._bypass_serializer = True # type: ignore[attr-defined] + assert self._spark._jvm is not None + jrdd = keyed._jrdd.map(self._spark._jvm.BytesToString()) + return self._df(self._jreader.json(jrdd)) + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_LIST_OF_RDD", + message_parameters={ + "arg_name": "path", + "arg_type": type(path).__name__, + }, + ) + + def table(self, tableName: str) -> "DataFrame": + """Returns the specified table as a :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + tableName : str + string, name of the table. + + Examples + -------- + >>> df = spark.range(10) + >>> df.createOrReplaceTempView('tblA') + >>> spark.read.table('tblA').show() + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + | 5| + | 6| + | 7| + | 8| + | 9| + +---+ + >>> _ = spark.sql("DROP TABLE tblA") + """ + return self._df(self._jreader.table(tableName)) + + def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame": + """ + Loads Parquet files, returning the result as a :class:`DataFrame`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + paths : str + + Other Parameters + ---------------- + **options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a Parquet file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a Parquet file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("parquet").save(d) + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.parquet(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + mergeSchema = options.get("mergeSchema", None) + pathGlobFilter = options.get("pathGlobFilter", None) + modifiedBefore = options.get("modifiedBefore", None) + modifiedAfter = options.get("modifiedAfter", None) + recursiveFileLookup = options.get("recursiveFileLookup", None) + datetimeRebaseMode = options.get("datetimeRebaseMode", None) + int96RebaseMode = options.get("int96RebaseMode", None) + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + datetimeRebaseMode=datetimeRebaseMode, + int96RebaseMode=int96RebaseMode, + ) + + return self._df(self._jreader.parquet(_to_seq(self._spark._sc, paths))) + + def text( + self, + paths: PathOrPaths, + wholetext: bool = False, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """ + Loads text files and returns a :class:`DataFrame` whose schema starts with a + string column named "value", and followed by partitioned columns if there + are any. + The text files must be encoded as UTF-8. + + By default, each line in the text file is a new row in the resulting DataFrame. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + paths : str or list + string, or list of strings, for input path(s). + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a text file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a text file + ... df = spark.createDataFrame([("a",), ("b",), ("c",)], schema=["alphabets"]) + ... df.write.mode("overwrite").format("text").save(d) + ... + ... # Read the text file as a DataFrame. + ... spark.read.schema(df.schema).text(d).sort("alphabets").show() + +---------+ + |alphabets| + +---------+ + | a| + | b| + | c| + +---------+ + """ + self._set_opts( + wholetext=wholetext, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + ) + + if isinstance(paths, str): + paths = [paths] + assert self._spark._sc._jvm is not None + return self._df(self._jreader.text(self._spark._sc._jvm.PythonUtils.toSeq(paths))) + + def csv( + self, + path: PathOrPaths, + schema: Optional[Union[StructType, str]] = None, + sep: Optional[str] = None, + encoding: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + comment: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + inferSchema: Optional[Union[bool, str]] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + nanValue: Optional[str] = None, + positiveInf: Optional[str] = None, + negativeInf: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + maxColumns: Optional[Union[int, str]] = None, + maxCharsPerColumn: Optional[Union[int, str]] = None, + maxMalformedLogPerPartition: Optional[Union[int, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[str] = None, + samplingRatio: Optional[Union[float, str]] = None, + enforceSchema: Optional[Union[bool, str]] = None, + emptyValue: Optional[str] = None, + locale: Optional[str] = None, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + unescapedQuoteHandling: Optional[str] = None, + ) -> "DataFrame": + r"""Loads a CSV file and returns the result as a :class:`DataFrame`. + + This function will go through the input once to determine the input schema if + ``inferSchema`` is enabled. To avoid going through the entire data once, disable + ``inferSchema`` option or specify the schema explicitly using ``schema``. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str or list + string, or list of strings, for input path(s), + or RDD of Strings storing CSV rows. + schema : :class:`pyspark.sql.types.StructType` or str, optional + an optional :class:`pyspark.sql.types.StructType` for the input schema + or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a CSV file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file + ... df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]) + ... df.write.mode("overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon'. + ... spark.read.csv(d, schema=df.schema, nullValue="Hyukjin Kwon").show() + +---+----+ + |age|name| + +---+----+ + |100|NULL| + +---+----+ + """ + self._set_opts( + schema=schema, + sep=sep, + encoding=encoding, + quote=quote, + escape=escape, + comment=comment, + header=header, + inferSchema=inferSchema, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + nullValue=nullValue, + nanValue=nanValue, + positiveInf=positiveInf, + negativeInf=negativeInf, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + maxColumns=maxColumns, + maxCharsPerColumn=maxCharsPerColumn, + maxMalformedLogPerPartition=maxMalformedLogPerPartition, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + multiLine=multiLine, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + samplingRatio=samplingRatio, + enforceSchema=enforceSchema, + emptyValue=emptyValue, + locale=locale, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + unescapedQuoteHandling=unescapedQuoteHandling, + ) + if isinstance(path, str): + path = [path] + if type(path) == list: + assert self._spark._sc._jvm is not None + return self._df(self._jreader.csv(self._spark._sc._jvm.PythonUtils.toSeq(path))) + elif isinstance(path, RDD): + + def func(iterator): + for x in iterator: + if not isinstance(x, str): + x = str(x) + if isinstance(x, str): + x = x.encode("utf-8") + yield x + + keyed = path.mapPartitions(func) + keyed._bypass_serializer = True + jrdd = keyed._jrdd.map(self._spark._jvm.BytesToString()) + # see SPARK-22112 + # There aren't any jvm api for creating a dataframe from rdd storing csv. + # We can do it through creating a jvm dataset firstly and using the jvm api + # for creating a dataframe from dataset storing csv. + jdataset = self._spark._jsparkSession.createDataset( + jrdd.rdd(), self._spark._jvm.Encoders.STRING() + ) + return self._df(self._jreader.csv(jdataset)) + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_LIST_OF_RDD", + message_parameters={ + "arg_name": "path", + "arg_type": type(path).__name__, + }, + ) + + def orc( + self, + path: PathOrPaths, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + modifiedBefore: Optional[Union[bool, str]] = None, + modifiedAfter: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """Loads ORC files, returning the result as a :class:`DataFrame`. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str or list + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a ORC file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a ORC file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("orc").save(d) + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.orc(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + modifiedBefore=modifiedBefore, + modifiedAfter=modifiedAfter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + path = [path] + return self._df(self._jreader.orc(_to_seq(self._spark._sc, path))) + + @overload + def jdbc( + self, url: str, table: str, *, properties: Optional[Dict[str, str]] = None + ) -> "DataFrame": + ... + + @overload + def jdbc( + self, + url: str, + table: str, + column: str, + lowerBound: Union[int, str], + upperBound: Union[int, str], + numPartitions: int, + *, + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + ... + + @overload + def jdbc( + self, + url: str, + table: str, + *, + predicates: List[str], + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + ... + + def jdbc( + self, + url: str, + table: str, + column: Optional[str] = None, + lowerBound: Optional[Union[int, str]] = None, + upperBound: Optional[Union[int, str]] = None, + numPartitions: Optional[int] = None, + predicates: Optional[List[str]] = None, + properties: Optional[Dict[str, str]] = None, + ) -> "DataFrame": + """ + Construct a :class:`DataFrame` representing the database table named ``table`` + accessible via JDBC URL ``url`` and connection ``properties``. + + Partitions of the table will be retrieved in parallel if either ``column`` or + ``predicates`` is specified. ``lowerBound``, ``upperBound`` and ``numPartitions`` + is needed when ``column`` is specified. + + If both ``column`` and ``predicates`` are specified, ``column`` will be used. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + table : str + the name of the table + column : str, optional + alias of ``partitionColumn`` option. Refer to ``partitionColumn`` in + `Data Source Option `_ + for the version you use. + predicates : list, optional + a list of expressions suitable for inclusion in WHERE clauses; + each one defines one partition of the :class:`DataFrame` + properties : dict, optional + a dictionary of JDBC database connection arguments. Normally at + least properties "user" and "password" with their corresponding values. + For example { 'user' : 'SYSTEM', 'password' : 'mypassword' } + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Notes + ----- + Don't create too many partitions in parallel on a large cluster; + otherwise Spark might crash your external database systems. + + Returns + ------- + :class:`DataFrame` + """ + if properties is None: + properties = dict() + assert self._spark._sc._gateway is not None + jprop = JavaClass( + "java.util.Properties", + self._spark._sc._gateway._gateway_client, + )() + for k in properties: + jprop.setProperty(k, properties[k]) + if column is not None: + assert lowerBound is not None, "lowerBound can not be None when ``column`` is specified" + assert upperBound is not None, "upperBound can not be None when ``column`` is specified" + assert ( + numPartitions is not None + ), "numPartitions can not be None when ``column`` is specified" + return self._df( + self._jreader.jdbc( + url, table, column, int(lowerBound), int(upperBound), int(numPartitions), jprop + ) + ) + if predicates is not None: + gateway = self._spark._sc._gateway + assert gateway is not None + jpredicates = utils.toJArray(gateway, gateway.jvm.java.lang.String, predicates) + return self._df(self._jreader.jdbc(url, table, jpredicates, jprop)) + return self._df(self._jreader.jdbc(url, table, jprop)) + + +class DataFrameWriter(OptionUtils): + """ + Interface used to write a :class:`DataFrame` to external storage systems + (e.g. file systems, key-value stores, etc). Use :attr:`DataFrame.write` + to access this. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, df: "DataFrame"): + self._df = df + self._spark = df.sparkSession + self._jwrite = df._jdf.write() + + def _sq(self, jsq: JavaObject) -> "StreamingQuery": + from pyspark.sql.streaming import StreamingQuery + + return StreamingQuery(jsq) + + def mode(self, saveMode: Optional[str]) -> "DataFrameWriter": + """Specifies the behavior when data or table already exists. + + Options include: + + * `append`: Append contents of this :class:`DataFrame` to existing data. + * `overwrite`: Overwrite existing data. + * `error` or `errorifexists`: Throw an exception if data already exists. + * `ignore`: Silently ignore this operation if data already exists. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + Raise an error when writing to an existing path. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... spark.createDataFrame( + ... [{"age": 80, "name": "Xinrong Meng"}] + ... ).write.mode("error").format("parquet").save(d) # doctest: +SKIP + Traceback (most recent call last): + ... + ...AnalysisException: ... + + Write a Parquet file back with various options, and read it back. + + >>> with tempfile.TemporaryDirectory() as d: + ... # Overwrite the path with a new Parquet file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("parquet").save(d) + ... + ... # Append another DataFrame into the Parquet file + ... spark.createDataFrame( + ... [{"age": 120, "name": "Takuya Ueshin"}] + ... ).write.mode("append").format("parquet").save(d) + ... + ... # Append another DataFrame into the Parquet file + ... spark.createDataFrame( + ... [{"age": 140, "name": "Haejoon Lee"}] + ... ).write.mode("ignore").format("parquet").save(d) + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.parquet(d).show() + +---+-------------+ + |age| name| + +---+-------------+ + |120|Takuya Ueshin| + |100| Hyukjin Kwon| + +---+-------------+ + """ + # At the JVM side, the default value of mode is already set to "error". + # So, if the given saveMode is None, we will not call JVM-side's mode method. + if saveMode is not None: + self._jwrite = self._jwrite.mode(saveMode) + return self + + def format(self, source: str) -> "DataFrameWriter": + """Specifies the underlying output data source. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + source : str + string, name of the data source, e.g. 'json', 'parquet'. + + Examples + -------- + >>> spark.range(1).write.format('parquet') + <...readwriter.DataFrameWriter object ...> + + Write a DataFrame into a Parquet file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a Parquet file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("parquet").save(d) + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.format('parquet').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self._jwrite = self._jwrite.format(source) + return self + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriter": + """ + Adds an output option for the underlying data source. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + key : str + The key for the option to set. + value + The value for the option to set. + + Examples + -------- + >>> spark.range(1).write.option("key", "value") + <...readwriter.DataFrameWriter object ...> + + Specify the option 'nullValue' with writing a CSV file. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file with 'nullValue' option set to 'Hyukjin Kwon'. + ... df = spark.createDataFrame([(100, None)], "age INT, name STRING") + ... df.write.option("nullValue", "Hyukjin Kwon").mode("overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame. + ... spark.read.schema(df.schema).format('csv').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + + self._jwrite = self._jwrite.option(key, to_str(value)) + return self + + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter": + """ + Adds output options for the underlying data source. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + **options : dict + The dictionary of string keys and primitive-type values. + + Examples + -------- + >>> spark.range(1).write.option("key", "value") + <...readwriter.DataFrameWriter object ...> + + Specify the option 'nullValue' and 'header' with writing a CSV file. + + >>> from pyspark.sql.types import StructType,StructField, StringType, IntegerType + >>> schema = StructType([ + ... StructField("age",IntegerType(),True), + ... StructField("name",StringType(),True), + ... ]) + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file with 'nullValue' option set to 'Hyukjin Kwon', + ... # and 'header' option set to `True`. + ... df = spark.createDataFrame([(100, None)], schema=schema) + ... df.write.options(nullValue="Hyukjin Kwon", header=True).mode( + ... "overwrite").format("csv").save(d) + ... + ... # Read the CSV file as a DataFrame. + ... spark.read.option("header", True).format('csv').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + for k in options: + self._jwrite = self._jwrite.option(k, to_str(options[k])) + return self + + @overload + def partitionBy(self, *cols: str) -> "DataFrameWriter": + ... + + @overload + def partitionBy(self, *cols: List[str]) -> "DataFrameWriter": + ... + + def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": + """Partitions the output by the given columns on the file system. + + If specified, the output is laid out on the file system similar + to Hive's partitioning scheme. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str or list + name of columns + + Examples + -------- + Write a DataFrame into a Parquet file in a partitioned manner, and read it back. + + >>> import tempfile + >>> import os + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a Parquet file in a partitioned manner. + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}, {"age": 120, "name": "Ruifeng Zheng"}] + ... ).write.partitionBy("name").mode("overwrite").format("parquet").save(d) + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.parquet(d).sort("age").show() + ... + ... # Read one partition as a DataFrame. + ... spark.read.parquet(f"{d}{os.path.sep}name=Hyukjin Kwon").show() + +---+-------------+ + |age| name| + +---+-------------+ + |100| Hyukjin Kwon| + |120|Ruifeng Zheng| + +---+-------------+ + +---+ + |age| + +---+ + |100| + +---+ + """ + if len(cols) == 1 and isinstance(cols[0], (list, tuple)): + cols = cols[0] # type: ignore[assignment] + self._jwrite = self._jwrite.partitionBy( + _to_seq(self._spark._sc, cast(Iterable["ColumnOrName"], cols)) + ) + return self + + @overload + def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter": + ... + + @overload + def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter": + ... + + def bucketBy( + self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str] + ) -> "DataFrameWriter": + """Buckets the output by the given columns. If specified, + the output is laid out on the file system similar to Hive's bucketing scheme, + but with a different bucket hash function and is not compatible with Hive's bucketing. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + numBuckets : int + the number of buckets to save + col : str, list or tuple + a name of a column, or a list of names. + cols : str + additional names (optional). If `col` is a list it should be empty. + + Notes + ----- + Applicable for file-based data sources in combination with + :py:meth:`DataFrameWriter.saveAsTable`. + + Examples + -------- + Write a DataFrame into a Parquet file in a buckted manner, and read it back. + + >>> from pyspark.sql.functions import input_file_name + >>> # Write a DataFrame into a Parquet file in a bucketed manner. + ... _ = spark.sql("DROP TABLE IF EXISTS bucketed_table") + >>> spark.createDataFrame([ + ... (100, "Hyukjin Kwon"), (120, "Hyukjin Kwon"), (140, "Haejoon Lee")], + ... schema=["age", "name"] + ... ).write.bucketBy(2, "name").mode("overwrite").saveAsTable("bucketed_table") + >>> # Read the Parquet file as a DataFrame. + ... spark.read.table("bucketed_table").sort("age").show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + |120|Hyukjin Kwon| + |140| Haejoon Lee| + +---+------------+ + >>> _ = spark.sql("DROP TABLE bucketed_table") + """ + if not isinstance(numBuckets, int): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={ + "arg_name": "numBuckets", + "arg_type": type(numBuckets).__name__, + }, + ) + + if isinstance(col, (list, tuple)): + if cols: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={ + "arg_list": f"`col` of type {type(col).__name__} and `cols`", + }, + ) + + col, cols = col[0], col[1:] # type: ignore[assignment] + + for c in cols: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "cols", + "arg_type": type(c).__name__, + }, + ) + if not isinstance(col, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "col", + "arg_type": type(col).__name__, + }, + ) + + self._jwrite = self._jwrite.bucketBy( + numBuckets, col, _to_seq(self._spark._sc, cast(Iterable["ColumnOrName"], cols)) + ) + return self + + @overload + def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": + ... + + @overload + def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter": + ... + + def sortBy( + self, col: Union[str, TupleOrListOfString], *cols: Optional[str] + ) -> "DataFrameWriter": + """Sorts the output in each bucket by the given columns on the file system. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + col : str, tuple or list + a name of a column, or a list of names. + cols : str + additional names (optional). If `col` is a list it should be empty. + + Examples + -------- + Write a DataFrame into a Parquet file in a sorted-buckted manner, and read it back. + + >>> from pyspark.sql.functions import input_file_name + >>> # Write a DataFrame into a Parquet file in a sorted-bucketed manner. + ... _ = spark.sql("DROP TABLE IF EXISTS sorted_bucketed_table") + >>> spark.createDataFrame([ + ... (100, "Hyukjin Kwon"), (120, "Hyukjin Kwon"), (140, "Haejoon Lee")], + ... schema=["age", "name"] + ... ).write.bucketBy(1, "name").sortBy("age").mode( + ... "overwrite").saveAsTable("sorted_bucketed_table") + >>> # Read the Parquet file as a DataFrame. + ... spark.read.table("sorted_bucketed_table").sort("age").show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + |120|Hyukjin Kwon| + |140| Haejoon Lee| + +---+------------+ + >>> _ = spark.sql("DROP TABLE sorted_bucketed_table") + """ + if isinstance(col, (list, tuple)): + if cols: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={ + "arg_list": f"`col` of type {type(col).__name__} and `cols`", + }, + ) + + col, cols = col[0], col[1:] # type: ignore[assignment] + + for c in cols: + if not isinstance(c, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "cols", + "arg_type": type(c).__name__, + }, + ) + if not isinstance(col, str): + raise PySparkTypeError( + error_class="NOT_LIST_OF_STR", + message_parameters={ + "arg_name": "col", + "arg_type": type(col).__name__, + }, + ) + + self._jwrite = self._jwrite.sortBy( + col, _to_seq(self._spark._sc, cast(Iterable["ColumnOrName"], cols)) + ) + return self + + def save( + self, + path: Optional[str] = None, + format: Optional[str] = None, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """Saves the contents of the :class:`DataFrame` to a data source. + + The data source is specified by the ``format`` and a set of ``options``. + If ``format`` is not specified, the default data source configured by + ``spark.sql.sources.default`` will be used. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str, optional + the path in a Hadoop supported file system + format : str, optional + the format used to save + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + partitionBy : list, optional + names of partitioning columns + **options : dict + all other string options + + Examples + -------- + Write a DataFrame into a JSON file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Read the JSON file as a DataFrame. + ... spark.read.format('json').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self.mode(mode).options(**options) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + if path is None: + self._jwrite.save() + else: + self._jwrite.save(path) + + def insertInto(self, tableName: str, overwrite: Optional[bool] = None) -> None: + """Inserts the content of the :class:`DataFrame` to the specified table. + + It requires that the schema of the :class:`DataFrame` is the same as the + schema of the table. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + overwrite : bool, optional + If true, overwrites existing data. Disabled by default + + Notes + ----- + Unlike :meth:`DataFrameWriter.saveAsTable`, :meth:`DataFrameWriter.insertInto` ignores + the column names and just uses position-based resolution. + + Examples + -------- + >>> _ = spark.sql("DROP TABLE IF EXISTS tblA") + >>> df = spark.createDataFrame([ + ... (100, "Hyukjin Kwon"), (120, "Hyukjin Kwon"), (140, "Haejoon Lee")], + ... schema=["age", "name"] + ... ) + >>> df.write.saveAsTable("tblA") + + Insert the data into 'tblA' table but with different column names. + + >>> df.selectExpr("age AS col1", "name AS col2").write.insertInto("tblA") + >>> spark.read.table("tblA").sort("age").show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + |100|Hyukjin Kwon| + |120|Hyukjin Kwon| + |120|Hyukjin Kwon| + |140| Haejoon Lee| + |140| Haejoon Lee| + +---+------------+ + >>> _ = spark.sql("DROP TABLE tblA") + """ + if overwrite is not None: + self.mode("overwrite" if overwrite else "append") + self._jwrite.insertInto(tableName) + + def saveAsTable( + self, + name: str, + format: Optional[str] = None, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + **options: "OptionalPrimitiveType", + ) -> None: + """Saves the content of the :class:`DataFrame` as the specified table. + + In the case the table already exists, behavior of this function depends on the + save mode, specified by the `mode` function (default to throwing an exception). + When `mode` is `Overwrite`, the schema of the :class:`DataFrame` does not need to be + the same as that of the existing table. + + * `append`: Append contents of this :class:`DataFrame` to existing data. + * `overwrite`: Overwrite existing data. + * `error` or `errorifexists`: Throw an exception if data already exists. + * `ignore`: Silently ignore this operation if data already exists. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + When `mode` is `Append`, if there is an existing table, we will use the format and + options of the existing table. The column order in the schema of the :class:`DataFrame` + doesn't need to be the same as that of the existing table. Unlike + :meth:`DataFrameWriter.insertInto`, :meth:`DataFrameWriter.saveAsTable` will use the + column names to find the correct column positions. + + Parameters + ---------- + name : str + the table name + format : str, optional + the format used to save + mode : str, optional + one of `append`, `overwrite`, `error`, `errorifexists`, `ignore` \ + (default: error) + partitionBy : str or list + names of partitioning columns + **options : dict + all other string options + + Examples + -------- + Creates a table from a DataFrame, and read it back. + + >>> _ = spark.sql("DROP TABLE IF EXISTS tblA") + >>> spark.createDataFrame([ + ... (100, "Hyukjin Kwon"), (120, "Hyukjin Kwon"), (140, "Haejoon Lee")], + ... schema=["age", "name"] + ... ).write.saveAsTable("tblA") + >>> spark.read.table("tblA").sort("age").show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + |120|Hyukjin Kwon| + |140| Haejoon Lee| + +---+------------+ + >>> _ = spark.sql("DROP TABLE tblA") + """ + self.mode(mode).options(**options) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + self._jwrite.saveAsTable(name) + + def json( + self, + path: str, + mode: Optional[str] = None, + compression: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + lineSep: Optional[str] = None, + encoding: Optional[str] = None, + ignoreNullFields: Optional[Union[bool, str]] = None, + ) -> None: + """Saves the content of the :class:`DataFrame` in JSON format + (`JSON Lines text format or newline-delimited JSON `_) at the + specified path. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a JSON file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.json(d, mode="overwrite") + ... + ... # Read the JSON file as a DataFrame. + ... spark.read.format("json").load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self.mode(mode) + self._set_opts( + compression=compression, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + lineSep=lineSep, + encoding=encoding, + ignoreNullFields=ignoreNullFields, + ) + self._jwrite.json(path) + + def parquet( + self, + path: str, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + compression: Optional[str] = None, + ) -> None: + """Saves the content of the :class:`DataFrame` in Parquet format at the specified path. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + partitionBy : str or list, optional + names of partitioning columns + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a Parquet file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a Parquet file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.parquet(d, mode="overwrite") + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.format("parquet").load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self.mode(mode) + if partitionBy is not None: + self.partitionBy(partitionBy) + self._set_opts(compression=compression) + self._jwrite.parquet(path) + + def text( + self, path: str, compression: Optional[str] = None, lineSep: Optional[str] = None + ) -> None: + """Saves the content of the DataFrame in a text file at the specified path. + The text files will be encoded as UTF-8. + + .. versionadded:: 1.6.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Notes + ----- + The DataFrame must have only one column that is of string type. + Each row becomes a new line in the output file. + + Examples + -------- + Write a DataFrame into a text file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a text file + ... df = spark.createDataFrame([("a",), ("b",), ("c",)], schema=["alphabets"]) + ... df.write.mode("overwrite").text(d) + ... + ... # Read the text file as a DataFrame. + ... spark.read.schema(df.schema).format("text").load(d).sort("alphabets").show() + +---------+ + |alphabets| + +---------+ + | a| + | b| + | c| + +---------+ + """ + self._set_opts(compression=compression, lineSep=lineSep) + self._jwrite.text(path) + + def csv( + self, + path: str, + mode: Optional[str] = None, + compression: Optional[str] = None, + sep: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + escapeQuotes: Optional[Union[bool, str]] = None, + quoteAll: Optional[Union[bool, str]] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[str] = None, + encoding: Optional[str] = None, + emptyValue: Optional[str] = None, + lineSep: Optional[str] = None, + ) -> None: + r"""Saves the content of the :class:`DataFrame` in CSV format at the specified path. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a CSV file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a CSV file + ... df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]) + ... df.write.csv(d, mode="overwrite") + ... + ... # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon'. + ... spark.read.schema(df.schema).format("csv").option( + ... "nullValue", "Hyukjin Kwon").load(d).show() + +---+----+ + |age|name| + +---+----+ + |100|NULL| + +---+----+ + """ + self.mode(mode) + self._set_opts( + compression=compression, + sep=sep, + quote=quote, + escape=escape, + header=header, + nullValue=nullValue, + escapeQuotes=escapeQuotes, + quoteAll=quoteAll, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + encoding=encoding, + emptyValue=emptyValue, + lineSep=lineSep, + ) + self._jwrite.csv(path) + + def orc( + self, + path: str, + mode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + compression: Optional[str] = None, + ) -> None: + """Saves the content of the :class:`DataFrame` in ORC format at the specified path. + + .. versionadded:: 1.5.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + partitionBy : str or list, optional + names of partitioning columns + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Examples + -------- + Write a DataFrame into a ORC file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a ORC file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.orc(d, mode="overwrite") + ... + ... # Read the Parquet file as a DataFrame. + ... spark.read.format("orc").load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + self.mode(mode) + if partitionBy is not None: + self.partitionBy(partitionBy) + self._set_opts(compression=compression) + self._jwrite.orc(path) + + def jdbc( + self, + url: str, + table: str, + mode: Optional[str] = None, + properties: Optional[Dict[str, str]] = None, + ) -> None: + """Saves the content of the :class:`DataFrame` to an external database table via JDBC. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + table : str + Name of the table in the external database. + mode : str, optional + specifies the behavior of the save operation when data already exists. + + * ``append``: Append contents of this :class:`DataFrame` to existing data. + * ``overwrite``: Overwrite existing data. + * ``ignore``: Silently ignore this operation if data already exists. + * ``error`` or ``errorifexists`` (default case): Throw an exception if data already \ + exists. + properties : dict + a dictionary of JDBC database connection arguments. Normally at + least properties "user" and "password" with their corresponding values. + For example { 'user' : 'SYSTEM', 'password' : 'mypassword' } + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + for the version you use. + + .. # noqa + + Notes + ----- + Don't create too many partitions in parallel on a large cluster; + otherwise Spark might crash your external database systems. + """ + if properties is None: + properties = dict() + + assert self._spark._sc._gateway is not None + jprop = JavaClass( + "java.util.Properties", + self._spark._sc._gateway._gateway_client, + )() + for k in properties: + jprop.setProperty(k, properties[k]) + self.mode(mode)._jwrite.jdbc(url, table, jprop) + + +class DataFrameWriterV2: + """ + Interface used to write a class:`pyspark.sql.dataframe.DataFrame` + to external storage using the v2 API. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, df: "DataFrame", table: str): + self._df = df + self._spark = df.sparkSession + self._jwriter = df._jdf.writeTo(table) + + @since(3.1) + def using(self, provider: str) -> "DataFrameWriterV2": + """ + Specifies a provider for the underlying output data source. + Spark's default catalog supports "parquet", "json", etc. + """ + self._jwriter.using(provider) + return self + + @since(3.1) + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriterV2": + """ + Add a write option. + """ + self._jwriter.option(key, to_str(value)) + return self + + @since(3.1) + def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriterV2": + """ + Add write options. + """ + options = {k: to_str(v) for k, v in options.items()} + self._jwriter.options(options) + return self + + @since(3.1) + def tableProperty(self, property: str, value: str) -> "DataFrameWriterV2": + """ + Add table property. + """ + self._jwriter.tableProperty(property, value) + return self + + @since(3.1) + def partitionedBy(self, col: Column, *cols: Column) -> "DataFrameWriterV2": + """ + Partition the output table created by `create`, `createOrReplace`, or `replace` using + the given columns or transforms. + + When specified, the table data will be stored by these values for efficient reads. + + For example, when a table is partitioned by day, it may be stored + in a directory layout like: + + * `table/day=2019-06-01/` + * `table/day=2019-06-02/` + + Partitioning is one of the most widely used techniques to optimize physical data layout. + It provides a coarse-grained index for skipping unnecessary data reads when queries have + predicates on the partitioned columns. In order for partitioning to work well, the number + of distinct values in each column should typically be less than tens of thousands. + + `col` and `cols` support only the following functions: + + * :py:func:`pyspark.sql.functions.years` + * :py:func:`pyspark.sql.functions.months` + * :py:func:`pyspark.sql.functions.days` + * :py:func:`pyspark.sql.functions.hours` + * :py:func:`pyspark.sql.functions.bucket` + + """ + col = _to_java_column(col) + cols = _to_seq(self._spark._sc, [_to_java_column(c) for c in cols]) + self._jwriter.partitionedBy(col, cols) + return self + + @since(3.1) + def create(self) -> None: + """ + Create a new table from the contents of the data frame. + + The new table's schema, partition layout, properties, and other configuration will be + based on the configuration set on this writer. + """ + self._jwriter.create() + + @since(3.1) + def replace(self) -> None: + """ + Replace an existing table with the contents of the data frame. + + The existing table's schema, partition layout, properties, and other configuration will be + replaced with the contents of the data frame and the configuration set on this writer. + """ + self._jwriter.replace() + + @since(3.1) + def createOrReplace(self) -> None: + """ + Create a new table or replace an existing table with the contents of the data frame. + + The output table's schema, partition layout, properties, + and other configuration will be based on the contents of the data frame + and the configuration set on this writer. + If the table exists, its configuration and data will be replaced. + """ + self._jwriter.createOrReplace() + + @since(3.1) + def append(self) -> None: + """ + Append the contents of the data frame to the output table. + """ + self._jwriter.append() + + @since(3.1) + def overwrite(self, condition: Column) -> None: + """ + Overwrite rows matching the given filter condition with the contents of the data frame in + the output table. + """ + condition = _to_java_column(condition) + self._jwriter.overwrite(condition) + + @since(3.1) + def overwritePartitions(self) -> None: + """ + Overwrite all partition for which the data frame contains at least one row with the contents + of the data frame in the output table. + + This operation is equivalent to Hive's `INSERT OVERWRITE ... PARTITION`, which replaces + partitions dynamically depending on the contents of the data frame. + """ + self._jwriter.overwritePartitions() + + +def _test() -> None: + import doctest + import os + import py4j + from pyspark.context import SparkContext + from pyspark.sql import SparkSession + import pyspark.sql.readwriter + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.readwriter.__dict__.copy() + sc = SparkContext("local[4]", "PythonTest") + try: + spark = SparkSession._getActiveSessionOrCreate() + except py4j.protocol.Py4JError: + spark = SparkSession(sc) + + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.readwriter, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/session.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/session.py new file mode 100644 index 0000000..d3f3359 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/session.py @@ -0,0 +1,2088 @@ +# +# 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 os +import sys +import warnings +from collections.abc import Sized +from functools import reduce +from threading import RLock +from types import TracebackType +from typing import ( + Any, + ClassVar, + Dict, + Iterable, + List, + Optional, + Tuple, + Type, + Union, + Set, + cast, + no_type_check, + overload, + TYPE_CHECKING, +) + +from py4j.java_gateway import JavaObject + +from pyspark import SparkConf, SparkContext +from pyspark.rdd import RDD +from pyspark.sql.column import _to_java_column +from pyspark.sql.conf import RuntimeConfig +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.functions import lit +from pyspark.sql.pandas.conversion import SparkConversionMixin +from pyspark.sql.readwriter import DataFrameReader +from pyspark.sql.sql_formatter import SQLStringFormatter +from pyspark.sql.streaming import DataStreamReader +from pyspark.sql.types import ( + AtomicType, + DataType, + StructField, + StructType, + _make_type_verifier, + _infer_schema, + _has_nulltype, + _merge_type, + _create_converter, + _parse_datatype_string, + _from_numpy_type, +) +from pyspark.errors.exceptions.captured import install_exception_handler +from pyspark.sql.utils import is_timestamp_ntz_preferred, to_str, try_remote_session_classmethod +from pyspark.errors import PySparkValueError, PySparkTypeError, PySparkRuntimeError + +if TYPE_CHECKING: + from pyspark.sql._typing import AtomicValue, RowLike, OptionalPrimitiveType + from pyspark.sql.catalog import Catalog + from pyspark.sql.pandas._typing import ArrayLike, DataFrameLike as PandasDataFrameLike + from pyspark.sql.streaming import StreamingQueryManager + from pyspark.sql.udf import UDFRegistration + from pyspark.sql.udtf import UDTFRegistration + + # Running MyPy type checks will always require pandas and + # other dependencies so importing here is fine. + from pyspark.sql.connect.client import SparkConnectClient + + +__all__ = ["SparkSession"] + + +def _monkey_patch_RDD(sparkSession: "SparkSession") -> None: + @no_type_check + def toDF(self, schema=None, sampleRatio=None): + """ + Converts current :class:`RDD` into a :class:`DataFrame` + + This is a shorthand for ``spark.createDataFrame(rdd, schema, sampleRatio)`` + + Parameters + ---------- + schema : :class:`pyspark.sql.types.DataType`, str or list, optional + a :class:`pyspark.sql.types.DataType` or a datatype string or a list of + column names, default is None. The data type string format equals to + :class:`pyspark.sql.types.DataType.simpleString`, except that top level struct type can + omit the ``struct<>`` and atomic types use ``typeName()`` as their format, e.g. use + ``byte`` instead of ``tinyint`` for :class:`pyspark.sql.types.ByteType`. + We can also use ``int`` as a short name for :class:`pyspark.sql.types.IntegerType`. + sampleRatio : float, optional + the sample ratio of rows used for inferring + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> rdd = spark.range(1).rdd.map(lambda x: tuple(x)) + >>> rdd.collect() + [(0,)] + >>> rdd.toDF().show() + +---+ + | _1| + +---+ + | 0| + +---+ + """ + return sparkSession.createDataFrame(self, schema, sampleRatio) + + RDD.toDF = toDF # type: ignore[assignment] + + +# TODO(SPARK-38912): This method can be dropped once support for Python 3.8 is dropped +# In Python 3.9, the @property decorator has been made compatible with the +# @classmethod decorator (https://docs.python.org/3.9/library/functions.html#classmethod) +# +# @classmethod + @property is also affected by a bug in Python's docstring which was backported +# to Python 3.9.6 (https://github.com/python/cpython/pull/28838) +class classproperty(property): + """Same as Python's @property decorator, but for class attributes. + + Examples + -------- + >>> class Builder: + ... def build(self): + ... return MyClass() + ... + >>> class MyClass: + ... @classproperty + ... def builder(cls): + ... print("instantiating new builder") + ... return Builder() + ... + >>> c1 = MyClass.builder + instantiating new builder + >>> c2 = MyClass.builder + instantiating new builder + >>> c1 == c2 + False + >>> isinstance(c1.build(), MyClass) + True + """ + + def __get__(self, instance: Any, owner: Any = None) -> "SparkSession.Builder": + # The "type: ignore" below silences the following error from mypy: + # error: Argument 1 to "classmethod" has incompatible + # type "Optional[Callable[[Any], Any]]"; + # expected "Callable[..., Any]" [arg-type] + return classmethod(self.fget).__get__(None, owner)() # type: ignore + + +class SparkSession(SparkConversionMixin): + """The entry point to programming Spark with the Dataset and DataFrame API. + + A SparkSession can be used to create :class:`DataFrame`, register :class:`DataFrame` as + tables, execute SQL over tables, cache tables, and read parquet files. + To create a :class:`SparkSession`, use the following builder pattern: + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + .. autoattribute:: builder + :annotation: + + Examples + -------- + Create a Spark session. + + >>> spark = ( + ... SparkSession.builder + ... .master("local") + ... .appName("Word Count") + ... .config("spark.some.config.option", "some-value") + ... .getOrCreate() + ... ) + + Create a Spark session with Spark Connect. + + >>> spark = ( + ... SparkSession.builder + ... .remote("sc://localhost") + ... .appName("Word Count") + ... .config("spark.some.config.option", "some-value") + ... .getOrCreate() + ... ) # doctest: +SKIP + """ + + class Builder: + """Builder for :class:`SparkSession`.""" + + _lock = RLock() + + def __init__(self) -> None: + self._options: Dict[str, Any] = {} + + @overload + def config(self, *, conf: SparkConf) -> "SparkSession.Builder": + ... + + @overload + def config(self, key: str, value: Any) -> "SparkSession.Builder": + ... + + @overload + def config(self, *, map: Dict[str, "OptionalPrimitiveType"]) -> "SparkSession.Builder": + ... + + def config( + self, + key: Optional[str] = None, + value: Optional[Any] = None, + conf: Optional[SparkConf] = None, + *, + map: Optional[Dict[str, "OptionalPrimitiveType"]] = None, + ) -> "SparkSession.Builder": + """Sets a config option. Options set using this method are automatically propagated to + both :class:`SparkConf` and :class:`SparkSession`'s own configuration. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + key : str, optional + a key name string for configuration property + value : str, optional + a value for configuration property + conf : :class:`SparkConf`, optional + an instance of :class:`SparkConf` + map: dictionary, optional + a dictionary of configurations to set + + .. versionadded:: 3.4.0 + + Returns + ------- + :class:`SparkSession.Builder` + + Examples + -------- + For an existing class:`SparkConf`, use `conf` parameter. + + >>> from pyspark.conf import SparkConf + >>> SparkSession.builder.config(conf=SparkConf()) + >> SparkSession.builder.config("spark.some.config.option", "some-value") + >> SparkSession.builder.config( + ... map={"spark.some.config.number": 123, "spark.some.config.float": 0.123}) + None: + """ + Helper function that validates the combination of startup URLs and raises an exception + if incompatible options are selected. + """ + if "spark.master" in self._options and ( + "spark.remote" in self._options or "SPARK_REMOTE" in os.environ + ): + raise RuntimeError( + "Spark master cannot be configured with Spark Connect server; " + "however, found URL for Spark Connect [%s]" + % self._options.get("spark.remote", os.environ.get("SPARK_REMOTE")) + ) + if "spark.remote" in self._options and ( + "spark.master" in self._options or "MASTER" in os.environ + ): + raise RuntimeError( + "Spark Connect server cannot be configured with Spark master; " + "however, found URL for Spark master [%s]" + % self._options.get("spark.master", os.environ.get("MASTER")) + ) + + if "spark.remote" in self._options: + remote = cast(str, self._options.get("spark.remote")) + if ("SPARK_REMOTE" in os.environ and os.environ["SPARK_REMOTE"] != remote) and ( + "SPARK_LOCAL_REMOTE" in os.environ and not remote.startswith("local") + ): + raise RuntimeError( + "Only one Spark Connect client URL can be set; however, got a " + "different URL [%s] from the existing [%s]" + % (os.environ["SPARK_REMOTE"], remote) + ) + + def master(self, master: str) -> "SparkSession.Builder": + """Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]" + to run locally with 4 cores, or "spark://master:7077" to run on a Spark standalone + cluster. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + master : str + a url for spark master + + Returns + ------- + :class:`SparkSession.Builder` + + Examples + -------- + >>> SparkSession.builder.master("local") + "SparkSession.Builder": + """Sets the Spark remote URL to connect to, such as "sc://host:port" to run + it via Spark Connect server. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + url : str + URL to Spark Connect server + + Returns + ------- + :class:`SparkSession.Builder` + + Examples + -------- + >>> SparkSession.builder.remote("sc://localhost") # doctest: +SKIP + "SparkSession.Builder": + """Sets a name for the application, which will be shown in the Spark web UI. + + If no application name is set, a randomly generated name will be used. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + an application name + + Returns + ------- + :class:`SparkSession.Builder` + + Examples + -------- + >>> SparkSession.builder.appName("My app") + "SparkSession.Builder": + """Enables Hive support, including connectivity to a persistent Hive metastore, support + for Hive SerDes, and Hive user-defined functions. + + .. versionadded:: 2.0.0 + + Returns + ------- + :class:`SparkSession.Builder` + + Examples + -------- + >>> SparkSession.builder.enableHiveSupport() + "SparkSession": + """Gets an existing :class:`SparkSession` or, if there is no existing one, creates a + new one based on the options set in this builder. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`SparkSession` + + Examples + -------- + This method first checks whether there is a valid global default SparkSession, and if + yes, return that one. If no valid global default SparkSession exists, the method + creates a new SparkSession and assigns the newly created SparkSession as the global + default. + + >>> s1 = SparkSession.builder.config("k1", "v1").getOrCreate() + >>> s1.conf.get("k1") == "v1" + True + + The configuration of the SparkSession can be changed afterwards + + >>> s1.conf.set("k1", "v1_new") + >>> s1.conf.get("k1") == "v1_new" + True + + In case an existing SparkSession is returned, the config options specified + in this builder will be applied to the existing SparkSession. + + >>> s2 = SparkSession.builder.config("k2", "v2").getOrCreate() + >>> s1.conf.get("k1") == s2.conf.get("k1") == "v1_new" + True + >>> s1.conf.get("k2") == s2.conf.get("k2") == "v2" + True + """ + from pyspark.context import SparkContext + from pyspark.conf import SparkConf + + opts = dict(self._options) + + with self._lock: + if ( + "SPARK_CONNECT_MODE_ENABLED" in os.environ + or "SPARK_REMOTE" in os.environ + or "spark.remote" in opts + ): + with SparkContext._lock: + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + + if ( + SparkContext._active_spark_context is None + and SparkSession._instantiatedSession is None + ): + url = opts.get("spark.remote", os.environ.get("SPARK_REMOTE")) + + if url.startswith("local"): + os.environ["SPARK_LOCAL_REMOTE"] = "1" + RemoteSparkSession._start_connect_server(url, opts) + url = "sc://localhost" + + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" + opts["spark.remote"] = url + return RemoteSparkSession.builder.config(map=opts).getOrCreate() + elif "SPARK_LOCAL_REMOTE" in os.environ: + url = "sc://localhost" + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" + opts["spark.remote"] = url + return RemoteSparkSession.builder.config(map=opts).getOrCreate() + else: + raise RuntimeError( + "Cannot start a remote Spark session because there " + "is a regular Spark session already running." + ) + + session = SparkSession._instantiatedSession + if session is None or session._sc._jsc is None: + sparkConf = SparkConf() + for key, value in self._options.items(): + sparkConf.set(key, value) + # This SparkContext may be an existing one. + sc = SparkContext.getOrCreate(sparkConf) + # Do not update `SparkConf` for existing `SparkContext`, as it's shared + # by all sessions. + session = SparkSession(sc, options=self._options) + else: + getattr( + getattr(session._jvm, "SparkSession$"), "MODULE$" + ).applyModifiableSettings(session._jsparkSession, self._options) + return session + + # Spark Connect-specific API + def create(self) -> "SparkSession": + """Creates a new SparkSession. Can only be used in the context of Spark Connect + and will throw an exception otherwise. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`SparkSession` + + Notes + ----- + This method will update the default and/or active session if they are not set. + """ + opts = dict(self._options) + if "SPARK_REMOTE" in os.environ or "spark.remote" in opts: + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + + # Validate that no incompatible configuration options are selected. + self._validate_startup_urls() + + url = opts.get("spark.remote", os.environ.get("SPARK_REMOTE")) + if url.startswith("local"): + raise RuntimeError( + "Creating new SparkSessions with `local` " + "connection string is not supported." + ) + + # Mark this Spark Session as Spark Connect. This prevents that local PySpark is + # used in conjunction with Spark Connect mode. + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" + opts["spark.remote"] = url + return RemoteSparkSession.builder.config(map=opts).create() + else: + raise RuntimeError( + "SparkSession.builder.create() can only be used with Spark Connect; " + "however, spark.remote was not found." + ) + + # TODO(SPARK-38912): Replace @classproperty with @classmethod + @property once support for + # Python 3.8 is dropped. + # + # In Python 3.9, the @property decorator has been made compatible with the + # @classmethod decorator (https://docs.python.org/3.9/library/functions.html#classmethod) + # + # @classmethod + @property is also affected by a bug in Python's docstring which was backported + # to Python 3.9.6 (https://github.com/python/cpython/pull/28838) + @classproperty + def builder(cls) -> Builder: + """Creates a :class:`Builder` for constructing a :class:`SparkSession`. + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + return cls.Builder() + + _instantiatedSession: ClassVar[Optional["SparkSession"]] = None + _activeSession: ClassVar[Optional["SparkSession"]] = None + + def __init__( + self, + sparkContext: SparkContext, + jsparkSession: Optional[JavaObject] = None, + options: Dict[str, Any] = {}, + ): + self._sc = sparkContext + self._jsc = self._sc._jsc + self._jvm = self._sc._jvm + + assert self._jvm is not None + + if jsparkSession is None: + if ( + self._jvm.SparkSession.getDefaultSession().isDefined() + and not self._jvm.SparkSession.getDefaultSession().get().sparkContext().isStopped() + ): + jsparkSession = self._jvm.SparkSession.getDefaultSession().get() + getattr(getattr(self._jvm, "SparkSession$"), "MODULE$").applyModifiableSettings( + jsparkSession, options + ) + else: + jsparkSession = self._jvm.SparkSession(self._jsc.sc(), options) + else: + getattr(getattr(self._jvm, "SparkSession$"), "MODULE$").applyModifiableSettings( + jsparkSession, options + ) + self._jsparkSession = jsparkSession + _monkey_patch_RDD(self) + install_exception_handler() + # If we had an instantiated SparkSession attached with a SparkContext + # which is stopped now, we need to renew the instantiated SparkSession. + # Otherwise, we will use invalid SparkSession when we call Builder.getOrCreate. + if ( + SparkSession._instantiatedSession is None + or SparkSession._instantiatedSession._sc._jsc is None + ): + SparkSession._instantiatedSession = self + SparkSession._activeSession = self + assert self._jvm is not None + self._jvm.SparkSession.setDefaultSession(self._jsparkSession) + self._jvm.SparkSession.setActiveSession(self._jsparkSession) + + def _repr_html_(self) -> str: + return """ +
    +

    SparkSession - {catalogImplementation}

    + {sc_HTML} +
    + """.format( + catalogImplementation=self.conf.get("spark.sql.catalogImplementation"), + sc_HTML=self.sparkContext._repr_html_(), + ) + + @property + def _jconf(self) -> "JavaObject": + """Accessor for the JVM SQL-specific configurations""" + return self._jsparkSession.sessionState().conf() + + def newSession(self) -> "SparkSession": + """ + Returns a new :class:`SparkSession` as new session, that has separate SQLConf, + registered temporary views and UDFs, but shared :class:`SparkContext` and + table cache. + + .. versionadded:: 2.0.0 + + Returns + ------- + :class:`SparkSession` + Spark session if an active session exists for the current thread + + Examples + -------- + >>> spark.newSession() + <...SparkSession object ...> + """ + return self.__class__(self._sc, self._jsparkSession.newSession()) + + @classmethod + @try_remote_session_classmethod + def getActiveSession(cls) -> Optional["SparkSession"]: + """ + Returns the active :class:`SparkSession` for the current thread, returned by the builder + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + :class:`SparkSession` + Spark session if an active session exists for the current thread + + Examples + -------- + >>> s = SparkSession.getActiveSession() + >>> df = s.createDataFrame([('Alice', 1)], ['name', 'age']) + >>> df.select("age").show() + +---+ + |age| + +---+ + | 1| + +---+ + """ + from pyspark import SparkContext + + sc = SparkContext._active_spark_context + if sc is None: + return None + else: + assert sc._jvm is not None + if sc._jvm.SparkSession.getActiveSession().isDefined(): + SparkSession(sc, sc._jvm.SparkSession.getActiveSession().get()) + return SparkSession._activeSession + else: + return None + + @classmethod + @try_remote_session_classmethod + def active(cls) -> "SparkSession": + """ + Returns the active or default :class:`SparkSession` for the current thread, returned by + the builder. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`SparkSession` + Spark session if an active or default session exists for the current thread. + """ + session = cls.getActiveSession() + if session is None: + session = cls._instantiatedSession + if session is None: + raise PySparkRuntimeError( + error_class="NO_ACTIVE_OR_DEFAULT_SESSION", + message_parameters={}, + ) + return session + + @property + def sparkContext(self) -> SparkContext: + """ + Returns the underlying :class:`SparkContext`. + + .. versionadded:: 2.0.0 + + Returns + ------- + :class:`SparkContext` + + Examples + -------- + >>> spark.sparkContext + + + Create an RDD from the Spark context + + >>> rdd = spark.sparkContext.parallelize([1, 2, 3]) + >>> rdd.collect() + [1, 2, 3] + """ + return self._sc + + @property + def version(self) -> str: + """ + The version of Spark on which this application is running. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + str + the version of Spark in string. + + Examples + -------- + >>> _ = spark.version + """ + return self._jsparkSession.version() + + @property + def conf(self) -> RuntimeConfig: + """Runtime configuration interface for Spark. + + This is the interface through which the user can get and set all Spark and Hadoop + configurations that are relevant to Spark SQL. When getting the value of a config, + this defaults to the value set in the underlying :class:`SparkContext`, if any. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`pyspark.sql.conf.RuntimeConfig` + + Examples + -------- + >>> spark.conf + + + Set a runtime configuration for the session + + >>> spark.conf.set("key", "value") + >>> spark.conf.get("key") + 'value' + """ + if not hasattr(self, "_conf"): + self._conf = RuntimeConfig(self._jsparkSession.conf()) + return self._conf + + @property + def catalog(self) -> "Catalog": + """Interface through which the user may create, drop, alter or query underlying + databases, tables, functions, etc. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`Catalog` + + Examples + -------- + >>> spark.catalog + <...Catalog object ...> + + Create a temp view, show the list, and drop it. + + >>> spark.range(1).createTempView("test_view") + >>> spark.catalog.listTables() + [Table(name='test_view', catalog=None, namespace=[], description=None, ... + >>> _ = spark.catalog.dropTempView("test_view") + """ + from pyspark.sql.catalog import Catalog + + if not hasattr(self, "_catalog"): + self._catalog = Catalog(self) + return self._catalog + + @property + def udf(self) -> "UDFRegistration": + """Returns a :class:`UDFRegistration` for UDF registration. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`UDFRegistration` + + Examples + -------- + Register a Python UDF, and use it in SQL. + + >>> strlen = spark.udf.register("strlen", lambda x: len(x)) + >>> spark.sql("SELECT strlen('test')").show() + +------------+ + |strlen(test)| + +------------+ + | 4| + +------------+ + """ + from pyspark.sql.udf import UDFRegistration + + return UDFRegistration(self) + + @property + def udtf(self) -> "UDTFRegistration": + """Returns a :class:`UDTFRegistration` for UDTF registration. + + .. versionadded:: 3.5.0 + + Returns + ------- + :class:`UDTFRegistration` + + Notes + ----- + Supports Spark Connect. + """ + from pyspark.sql.udtf import UDTFRegistration + + return UDTFRegistration(self) + + def range( + self, + start: int, + end: Optional[int] = None, + step: int = 1, + numPartitions: Optional[int] = None, + ) -> DataFrame: + """ + Create a :class:`DataFrame` with single :class:`pyspark.sql.types.LongType` column named + ``id``, containing elements in a range from ``start`` to ``end`` (exclusive) with + step value ``step``. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + start : int + the start value + end : int, optional + the end value (exclusive) + step : int, optional + the incremental step (default: 1) + numPartitions : int, optional + the number of partitions of the DataFrame + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> spark.range(1, 7, 2).show() + +---+ + | id| + +---+ + | 1| + | 3| + | 5| + +---+ + + If only one argument is specified, it will be used as the end value. + + >>> spark.range(3).show() + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + +---+ + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + + if end is None: + jdf = self._jsparkSession.range(0, int(start), int(step), int(numPartitions)) + else: + jdf = self._jsparkSession.range(int(start), int(end), int(step), int(numPartitions)) + + return DataFrame(jdf, self) + + def _inferSchemaFromList( + self, data: Iterable[Any], names: Optional[List[str]] = None + ) -> StructType: + """ + Infer schema from list of Row, dict, or tuple. + + Parameters + ---------- + data : iterable + list of Row, dict, or tuple + names : list, optional + list of column names + + Returns + ------- + :class:`pyspark.sql.types.StructType` + """ + if not data: + raise PySparkValueError( + error_class="CANNOT_INFER_EMPTY_SCHEMA", + message_parameters={}, + ) + infer_dict_as_struct = self._jconf.inferDictAsStruct() + infer_array_from_first_element = self._jconf.legacyInferArrayTypeFromFirstElement() + prefer_timestamp_ntz = is_timestamp_ntz_preferred() + schema = reduce( + _merge_type, + ( + _infer_schema( + row, + names, + infer_dict_as_struct=infer_dict_as_struct, + infer_array_from_first_element=infer_array_from_first_element, + prefer_timestamp_ntz=prefer_timestamp_ntz, + ) + for row in data + ), + ) + if _has_nulltype(schema): + raise PySparkValueError( + error_class="CANNOT_DETERMINE_TYPE", + message_parameters={}, + ) + return schema + + def _inferSchema( + self, + rdd: RDD[Any], + samplingRatio: Optional[float] = None, + names: Optional[List[str]] = None, + ) -> StructType: + """ + Infer schema from an RDD of Row, dict, or tuple. + + Parameters + ---------- + rdd : :class:`RDD` + an RDD of Row, dict, or tuple + samplingRatio : float, optional + sampling ratio, or no sampling (default) + names : list, optional + + Returns + ------- + :class:`pyspark.sql.types.StructType` + """ + first = rdd.first() + if isinstance(first, Sized) and len(first) == 0: + raise PySparkValueError( + error_class="CANNOT_INFER_EMPTY_SCHEMA", + message_parameters={}, + ) + + infer_dict_as_struct = self._jconf.inferDictAsStruct() + infer_array_from_first_element = self._jconf.legacyInferArrayTypeFromFirstElement() + prefer_timestamp_ntz = is_timestamp_ntz_preferred() + if samplingRatio is None: + schema = _infer_schema( + first, + names=names, + infer_dict_as_struct=infer_dict_as_struct, + prefer_timestamp_ntz=prefer_timestamp_ntz, + ) + if _has_nulltype(schema): + for row in rdd.take(100)[1:]: + schema = _merge_type( + schema, + _infer_schema( + row, + names=names, + infer_dict_as_struct=infer_dict_as_struct, + infer_array_from_first_element=infer_array_from_first_element, + prefer_timestamp_ntz=prefer_timestamp_ntz, + ), + ) + if not _has_nulltype(schema): + break + else: + raise PySparkValueError( + error_class="CANNOT_DETERMINE_TYPE", + message_parameters={}, + ) + else: + if samplingRatio < 0.99: + rdd = rdd.sample(False, float(samplingRatio)) + schema = rdd.map( + lambda row: _infer_schema( + row, + names, + infer_dict_as_struct=infer_dict_as_struct, + infer_array_from_first_element=infer_array_from_first_element, + prefer_timestamp_ntz=prefer_timestamp_ntz, + ) + ).reduce(_merge_type) + return schema + + def _createFromRDD( + self, + rdd: RDD[Any], + schema: Optional[Union[DataType, List[str]]], + samplingRatio: Optional[float], + ) -> Tuple[RDD[Tuple], StructType]: + """ + Create an RDD for DataFrame from an existing RDD, returns the RDD and schema. + """ + if schema is None or isinstance(schema, (list, tuple)): + struct = self._inferSchema(rdd, samplingRatio, names=schema) + converter = _create_converter(struct) + tupled_rdd = rdd.map(converter) + if isinstance(schema, (list, tuple)): + for i, name in enumerate(schema): + struct.fields[i].name = name + struct.names[i] = name + + elif isinstance(schema, StructType): + struct = schema + tupled_rdd = rdd + + else: + raise PySparkTypeError( + error_class="NOT_LIST_OR_NONE_OR_STRUCT", + message_parameters={ + "arg_name": "schema", + "arg_type": type(schema).__name__, + }, + ) + + # convert python objects to sql data + internal_rdd = tupled_rdd.map(struct.toInternal) + return internal_rdd, struct + + def _createFromLocal( + self, data: Iterable[Any], schema: Optional[Union[DataType, List[str]]] + ) -> Tuple[RDD[Tuple], StructType]: + """ + Create an RDD for DataFrame from a list or pandas.DataFrame, returns + the RDD and schema. + """ + # make sure data could consumed multiple times + if not isinstance(data, list): + data = list(data) + + if schema is None or isinstance(schema, (list, tuple)): + struct = self._inferSchemaFromList(data, names=schema) + converter = _create_converter(struct) + tupled_data: Iterable[Tuple] = map(converter, data) + if isinstance(schema, (list, tuple)): + for i, name in enumerate(schema): + struct.fields[i].name = name + struct.names[i] = name + + elif isinstance(schema, StructType): + struct = schema + tupled_data = data + + else: + raise PySparkTypeError( + error_class="NOT_LIST_OR_NONE_OR_STRUCT", + message_parameters={ + "arg_name": "schema", + "arg_type": type(schema).__name__, + }, + ) + + # convert python objects to sql data + internal_data = [struct.toInternal(row) for row in tupled_data] + return self._sc.parallelize(internal_data), struct + + @staticmethod + def _create_shell_session() -> "SparkSession": + """ + Initialize a :class:`SparkSession` for a pyspark shell session. This is called from + shell.py to make error handling simpler without needing to declare local variables in + that script, which would expose those to users. + """ + import py4j + from pyspark.conf import SparkConf + from pyspark.context import SparkContext + + try: + # Try to access HiveConf, it will raise exception if Hive is not added + conf = SparkConf() + assert SparkContext._jvm is not None + if conf.get("spark.sql.catalogImplementation", "hive").lower() == "hive": + SparkContext._jvm.org.apache.hadoop.hive.conf.HiveConf() + return SparkSession.builder.enableHiveSupport().getOrCreate() + else: + return SparkSession._getActiveSessionOrCreate() + except (py4j.protocol.Py4JError, TypeError): + if conf.get("spark.sql.catalogImplementation", "").lower() == "hive": + warnings.warn( + "Fall back to non-hive support because failing to access HiveConf, " + "please make sure you build spark with hive" + ) + + return SparkSession._getActiveSessionOrCreate() + + @staticmethod + def _getActiveSessionOrCreate(**static_conf: Any) -> "SparkSession": + """ + Returns the active :class:`SparkSession` for the current thread, returned by the builder, + or if there is no existing one, creates a new one based on the options set in the builder. + + NOTE that 'static_conf' might not be set if there's an active or default Spark session + running. + """ + spark = SparkSession.getActiveSession() + if spark is None: + builder = SparkSession.builder + for k, v in static_conf.items(): + builder = builder.config(k, v) + spark = builder.getOrCreate() + return spark + + @overload + def createDataFrame( + self, + data: Iterable["RowLike"], + schema: Union[List[str], Tuple[str, ...]] = ..., + samplingRatio: Optional[float] = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: "RDD[RowLike]", + schema: Union[List[str], Tuple[str, ...]] = ..., + samplingRatio: Optional[float] = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: Iterable["RowLike"], + schema: Union[StructType, str], + *, + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: "RDD[RowLike]", + schema: Union[StructType, str], + *, + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: "RDD[AtomicValue]", + schema: Union[AtomicType, str], + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: Iterable["AtomicValue"], + schema: Union[AtomicType, str], + verifySchema: bool = ..., + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, data: "PandasDataFrameLike", samplingRatio: Optional[float] = ... + ) -> DataFrame: + ... + + @overload + def createDataFrame( + self, + data: "PandasDataFrameLike", + schema: Union[StructType, str], + verifySchema: bool = ..., + ) -> DataFrame: + ... + + def createDataFrame( # type: ignore[misc] + self, + data: Union[RDD[Any], Iterable[Any], "PandasDataFrameLike", "ArrayLike"], + schema: Optional[Union[AtomicType, StructType, str]] = None, + samplingRatio: Optional[float] = None, + verifySchema: bool = True, + ) -> DataFrame: + """ + Creates a :class:`DataFrame` from an :class:`RDD`, a list, a :class:`pandas.DataFrame` + or a :class:`numpy.ndarray`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + data : :class:`RDD` or iterable + an RDD of any kind of SQL data representation (:class:`Row`, + :class:`tuple`, ``int``, ``boolean``, etc.), or :class:`list`, + :class:`pandas.DataFrame` or :class:`numpy.ndarray`. + schema : :class:`pyspark.sql.types.DataType`, str or list, optional + a :class:`pyspark.sql.types.DataType` or a datatype string or a list of + column names, default is None. The data type string format equals to + :class:`pyspark.sql.types.DataType.simpleString`, except that top level struct type can + omit the ``struct<>``. + + When ``schema`` is a list of column names, the type of each column + will be inferred from ``data``. + + When ``schema`` is ``None``, it will try to infer the schema (column names and types) + from ``data``, which should be an RDD of either :class:`Row`, + :class:`namedtuple`, or :class:`dict`. + + When ``schema`` is :class:`pyspark.sql.types.DataType` or a datatype string, it must + match the real data, or an exception will be thrown at runtime. If the given schema is + not :class:`pyspark.sql.types.StructType`, it will be wrapped into a + :class:`pyspark.sql.types.StructType` as its only field, and the field name will be + "value". Each record will also be wrapped into a tuple, which can be converted to row + later. + samplingRatio : float, optional + the sample ratio of rows used for inferring. The first few rows will be used + if ``samplingRatio`` is ``None``. + verifySchema : bool, optional + verify data types of every row against schema. Enabled by default. + + .. versionadded:: 2.1.0 + + Returns + ------- + :class:`DataFrame` + + Notes + ----- + Usage with `spark.sql.execution.arrow.pyspark.enabled=True` is experimental. + + Examples + -------- + Create a DataFrame from a list of tuples. + + >>> spark.createDataFrame([('Alice', 1)]).show() + +-----+---+ + | _1| _2| + +-----+---+ + |Alice| 1| + +-----+---+ + + Create a DataFrame from a list of dictionaries. + + >>> d = [{'name': 'Alice', 'age': 1}] + >>> spark.createDataFrame(d).show() + +---+-----+ + |age| name| + +---+-----+ + | 1|Alice| + +---+-----+ + + Create a DataFrame with column names specified. + + >>> spark.createDataFrame([('Alice', 1)], ['name', 'age']).show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 1| + +-----+---+ + + Create a DataFrame with the explicit schema specified. + + >>> from pyspark.sql.types import * + >>> schema = StructType([ + ... StructField("name", StringType(), True), + ... StructField("age", IntegerType(), True)]) + >>> spark.createDataFrame([('Alice', 1)], schema).show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 1| + +-----+---+ + + Create a DataFrame with the schema in DDL formatted string. + + >>> spark.createDataFrame([('Alice', 1)], "name: string, age: int").show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 1| + +-----+---+ + + Create an empty DataFrame. + When initializing an empty DataFrame in PySpark, it's mandatory to specify its schema, + as the DataFrame lacks data from which the schema can be inferred. + + >>> spark.createDataFrame([], "name: string, age: int").show() + +----+---+ + |name|age| + +----+---+ + +----+---+ + + Create a DataFrame from Row objects. + + >>> from pyspark.sql import Row + >>> Person = Row('name', 'age') + >>> df = spark.createDataFrame([Person("Alice", 1)]) + >>> df.show() + +-----+---+ + | name|age| + +-----+---+ + |Alice| 1| + +-----+---+ + + Create a DataFrame from a pandas DataFrame. + + >>> spark.createDataFrame(df.toPandas()).show() # doctest: +SKIP + +-----+---+ + | name|age| + +-----+---+ + |Alice| 1| + +-----+---+ + >>> spark.createDataFrame(pandas.DataFrame([[1, 2]])).collect() # doctest: +SKIP + +---+---+ + | 0| 1| + +---+---+ + | 1| 2| + +---+---+ + """ + SparkSession._activeSession = self + assert self._jvm is not None + self._jvm.SparkSession.setActiveSession(self._jsparkSession) + if isinstance(data, DataFrame): + raise PySparkTypeError( + error_class="SHOULD_NOT_DATAFRAME", + message_parameters={"arg_name": "data"}, + ) + + if isinstance(schema, str): + schema = cast(Union[AtomicType, StructType, str], _parse_datatype_string(schema)) + elif isinstance(schema, (list, tuple)): + # Must re-encode any unicode strings to be consistent with StructField names + schema = [x.encode("utf-8") if not isinstance(x, str) else x for x in schema] + + try: + import pandas as pd + + has_pandas = True + except Exception: + has_pandas = False + + try: + import numpy as np + + has_numpy = True + except Exception: + has_numpy = False + + if has_numpy and isinstance(data, np.ndarray): + # `data` of numpy.ndarray type will be converted to a pandas DataFrame, + # so pandas is required. + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + if data.ndim not in [1, 2]: + raise PySparkValueError( + error_class="INVALID_NDARRAY_DIMENSION", + message_parameters={"dimensions": "1 or 2"}, + ) + + if data.ndim == 1 or data.shape[1] == 1: + column_names = ["value"] + else: + column_names = ["_%s" % i for i in range(1, data.shape[1] + 1)] + + if schema is None and not self._jconf.arrowPySparkEnabled(): + # Construct `schema` from `np.dtype` of the input NumPy array + # TODO: Apply the logic below when self._jconf.arrowPySparkEnabled() is True + spark_type = _from_numpy_type(data.dtype) + if spark_type is not None: + schema = StructType( + [StructField(name, spark_type, nullable=True) for name in column_names] + ) + + data = pd.DataFrame(data, columns=column_names) + + if has_pandas and isinstance(data, pd.DataFrame): + # Create a DataFrame from pandas DataFrame. + return super(SparkSession, self).createDataFrame( # type: ignore[call-overload] + data, schema, samplingRatio, verifySchema + ) + return self._create_dataframe( + data, schema, samplingRatio, verifySchema # type: ignore[arg-type] + ) + + def _create_dataframe( + self, + data: Union[RDD[Any], Iterable[Any]], + schema: Optional[Union[DataType, List[str]]], + samplingRatio: Optional[float], + verifySchema: bool, + ) -> DataFrame: + if isinstance(schema, StructType): + verify_func = _make_type_verifier(schema) if verifySchema else lambda _: True + + @no_type_check + def prepare(obj): + verify_func(obj) + return obj + + elif isinstance(schema, DataType): + dataType = schema + schema = StructType().add("value", schema) + + verify_func = ( + _make_type_verifier(dataType, name="field value") + if verifySchema + else lambda _: True + ) + + @no_type_check + def prepare(obj): + verify_func(obj) + return (obj,) + + else: + + def prepare(obj: Any) -> Any: + return obj + + if isinstance(data, RDD): + rdd, struct = self._createFromRDD(data.map(prepare), schema, samplingRatio) + else: + rdd, struct = self._createFromLocal(map(prepare, data), schema) + assert self._jvm is not None + jrdd = self._jvm.SerDeUtil.toJavaArray(rdd._to_java_object_rdd()) + jdf = self._jsparkSession.applySchemaToPythonRDD(jrdd.rdd(), struct.json()) + df = DataFrame(jdf, self) + df._schema = struct + return df + + def sql( + self, sqlQuery: str, args: Optional[Union[Dict[str, Any], List]] = None, **kwargs: Any + ) -> DataFrame: + """Returns a :class:`DataFrame` representing the result of the given query. + When ``kwargs`` is specified, this method formats the given string by using the Python + standard formatter. The method binds named parameters to SQL literals or + positional parameters from `args`. It doesn't support named and positional parameters + in the same SQL query. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect and parameterized SQL. + + .. versionchanged:: 3.5.0 + Added positional parameters. + + Parameters + ---------- + sqlQuery : str + SQL query string. + args : dict or list + A dictionary of parameter names to Python objects or a list of Python objects + that can be converted to SQL literal expressions. See + + Supported Data Types for supported value types in Python. + For example, dictionary keys: "rank", "name", "birthdate"; + dictionary or list values: 1, "Steven", datetime.date(2023, 4, 2). + A value can be also a `Column` of literal expression, in that case it is taken as is. + + .. versionadded:: 3.4.0 + + kwargs : dict + Other variables that the user wants to set that can be referenced in the query + + .. versionchanged:: 3.3.0 + Added optional argument ``kwargs`` to specify the mapping of variables in the query. + This feature is experimental and unstable. + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + Executing a SQL query. + + >>> spark.sql("SELECT * FROM range(10) where id > 7").show() + +---+ + | id| + +---+ + | 8| + | 9| + +---+ + + Executing a SQL query with variables as Python formatter standard. + + >>> spark.sql( + ... "SELECT * FROM range(10) WHERE id > {bound1} AND id < {bound2}", bound1=7, bound2=9 + ... ).show() + +---+ + | id| + +---+ + | 8| + +---+ + + >>> mydf = spark.range(10) + >>> spark.sql( + ... "SELECT {col} FROM {mydf} WHERE id IN {x}", + ... col=mydf.id, mydf=mydf, x=tuple(range(4))).show() + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + +---+ + + >>> spark.sql(''' + ... SELECT m1.a, m2.b + ... FROM {table1} m1 INNER JOIN {table2} m2 + ... ON m1.key = m2.key + ... ORDER BY m1.a, m2.b''', + ... table1=spark.createDataFrame([(1, "a"), (2, "b")], ["a", "key"]), + ... table2=spark.createDataFrame([(3, "a"), (4, "b"), (5, "b")], ["b", "key"])).show() + +---+---+ + | a| b| + +---+---+ + | 1| 3| + | 2| 4| + | 2| 5| + +---+---+ + + Also, it is possible to query using class:`Column` from :class:`DataFrame`. + + >>> mydf = spark.createDataFrame([(1, 4), (2, 4), (3, 6)], ["A", "B"]) + >>> spark.sql("SELECT {df.A}, {df[B]} FROM {df}", df=mydf).show() + +---+---+ + | A| B| + +---+---+ + | 1| 4| + | 2| 4| + | 3| 6| + +---+---+ + + And substitude named parameters with the `:` prefix by SQL literals. + + >>> spark.sql("SELECT * FROM {df} WHERE {df[B]} > :minB", {"minB" : 5}, df=mydf).show() + +---+---+ + | A| B| + +---+---+ + | 3| 6| + +---+---+ + + Or positional parameters marked by `?` in the SQL query by SQL literals. + + >>> spark.sql( + ... "SELECT * FROM {df} WHERE {df[B]} > ? and ? < {df[A]}", + ... args=[5, 2], df=mydf).show() + +---+---+ + | A| B| + +---+---+ + | 3| 6| + +---+---+ + """ + + formatter = SQLStringFormatter(self) + if len(kwargs) > 0: + sqlQuery = formatter.format(sqlQuery, **kwargs) + try: + if isinstance(args, Dict): + litArgs = {k: _to_java_column(lit(v)) for k, v in (args or {}).items()} + else: + assert self._jvm is not None + litArgs = self._jvm.PythonUtils.toArray( + [_to_java_column(lit(v)) for v in (args or [])] + ) + return DataFrame(self._jsparkSession.sql(sqlQuery, litArgs), self) + finally: + if len(kwargs) > 0: + formatter.clear() + + def table(self, tableName: str) -> DataFrame: + """Returns the specified table as a :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + tableName : str + the table name to retrieve. + + Returns + ------- + :class:`DataFrame` + + Examples + -------- + >>> spark.range(5).createOrReplaceTempView("table1") + >>> spark.table("table1").sort("id").show() + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + +---+ + """ + return DataFrame(self._jsparkSession.table(tableName), self) + + @property + def read(self) -> DataFrameReader: + """ + Returns a :class:`DataFrameReader` that can be used to read data + in as a :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Returns + ------- + :class:`DataFrameReader` + + Examples + -------- + >>> spark.read + <...DataFrameReader object ...> + + Write a DataFrame into a JSON file and read it back. + + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a DataFrame into a JSON file + ... spark.createDataFrame( + ... [{"age": 100, "name": "Hyukjin Kwon"}] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Read the JSON file as a DataFrame. + ... spark.read.format('json').load(d).show() + +---+------------+ + |age| name| + +---+------------+ + |100|Hyukjin Kwon| + +---+------------+ + """ + return DataFrameReader(self) + + @property + def readStream(self) -> DataStreamReader: + """ + Returns a :class:`DataStreamReader` that can be used to read data streams + as a streaming :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Returns + ------- + :class:`DataStreamReader` + + Examples + -------- + >>> spark.readStream + + + The example below uses Rate source that generates rows continuously. + After that, we operate a modulo by 3, and then write the stream out to the console. + The streaming query stops in 3 seconds. + + >>> import time + >>> df = spark.readStream.format("rate").load() + >>> df = df.selectExpr("value % 3 as v") + >>> q = df.writeStream.format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + return DataStreamReader(self) + + @property + def streams(self) -> "StreamingQueryManager": + """Returns a :class:`StreamingQueryManager` that allows managing all the + :class:`StreamingQuery` instances active on `this` context. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Returns + ------- + :class:`StreamingQueryManager` + + Examples + -------- + >>> spark.streams + + + Get the list of active streaming queries + + >>> sq = spark.readStream.format( + ... "rate").load().writeStream.format('memory').queryName('this_query').start() + >>> sqm = spark.streams + >>> [q.name for q in sqm.active] + ['this_query'] + >>> sq.stop() + """ + from pyspark.sql.streaming import StreamingQueryManager + + return StreamingQueryManager(self._jsparkSession.streams()) + + def stop(self) -> None: + """ + Stop the underlying :class:`SparkContext`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Examples + -------- + >>> spark.stop() # doctest: +SKIP + """ + from pyspark.sql.context import SQLContext + + self._sc.stop() + # We should clean the default session up. See SPARK-23228. + assert self._jvm is not None + self._jvm.SparkSession.clearDefaultSession() + self._jvm.SparkSession.clearActiveSession() + SparkSession._instantiatedSession = None + SparkSession._activeSession = None + SQLContext._instantiatedContext = None + + def __enter__(self) -> "SparkSession": + """ + Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> with SparkSession.builder.master("local").getOrCreate() as session: + ... session.range(5).show() # doctest: +SKIP + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + +---+ + """ + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + """ + Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. + + Specifically stop the SparkSession on exit of the with block. + + .. versionadded:: 2.0.0 + + Examples + -------- + >>> with SparkSession.builder.master("local").getOrCreate() as session: + ... session.range(5).show() # doctest: +SKIP + +---+ + | id| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + +---+ + """ + self.stop() + + # SparkConnect-specific API + @property + def client(self) -> "SparkConnectClient": + """ + Gives access to the Spark Connect client. In normal cases this is not necessary to be used + and only relevant for testing. + + .. versionadded:: 3.4.0 + + Returns + ------- + :class:`SparkConnectClient` + + Notes + ----- + This API is unstable, and a developer API. It returns non-API instance + :class:`SparkConnectClient`. + This is an API dedicated to Spark Connect client only. With regular Spark Session, it throws + an exception. + """ + raise RuntimeError( + "SparkSession.client is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def addArtifacts( + self, *path: str, pyfile: bool = False, archive: bool = False, file: bool = False + ) -> None: + """ + Add artifact(s) to the client session. Currently only local files are supported. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + *path : tuple of str + Artifact's URIs to add. + pyfile : bool + Whether to add them as Python dependencies such as .py, .egg, .zip or .jar files. + The pyfiles are directly inserted into the path when executing Python functions + in executors. + archive : bool + Whether to add them as archives such as .zip, .jar, .tar.gz, .tgz, or .tar files. + The archives are unpacked on the executor side automatically. + file : bool + Add a file to be downloaded with this Spark job on every node. + The ``path`` passed can only be a local file for now. + + Notes + ----- + This is an API dedicated to Spark Connect client only. With regular Spark Session, it throws + an exception. + """ + raise RuntimeError( + "SparkSession.addArtifact(s) is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + addArtifact = addArtifacts + + def copyFromLocalToFs(self, local_path: str, dest_path: str) -> None: + """ + Copy file from local to cloud storage file system. + If the file already exits in destination path, old file is overwritten. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + local_path: str + Path to a local file. Directories are not supported. + The path can be either an absolute path or a relative path. + dest_path: str + The cloud storage path to the destination the file will + be copied to. + The path must be an an absolute path. + + Notes + ----- + This API is a developer API. + Also, this is an API dedicated to Spark Connect client only. With regular + Spark Session, it throws an exception. + """ + raise RuntimeError( + "SparkSession.copyFromLocalToFs is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def interruptAll(self) -> List[str]: + """ + Interrupt all operations of this session currently running on the connected server. + + .. versionadded:: 3.5.0 + + Returns + ------- + list of str + List of operationIds of interrupted operations. + + Notes + ----- + There is still a possibility of operation finishing just as it is interrupted. + """ + raise RuntimeError( + "SparkSession.interruptAll is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def interruptTag(self, tag: str) -> List[str]: + """ + Interrupt all operations of this session with the given operation tag. + + .. versionadded:: 3.5.0 + + Returns + ------- + list of str + List of operationIds of interrupted operations. + + Notes + ----- + There is still a possibility of operation finishing just as it is interrupted. + """ + raise RuntimeError( + "SparkSession.interruptTag is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def interruptOperation(self, op_id: str) -> List[str]: + """ + Interrupt an operation of this session with the given operationId. + + .. versionadded:: 3.5.0 + + Returns + ------- + list of str + List of operationIds of interrupted operations. + + Notes + ----- + There is still a possibility of operation finishing just as it is interrupted. + """ + raise RuntimeError( + "SparkSession.interruptOperation is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def addTag(self, tag: str) -> None: + """ + Add a tag to be assigned to all the operations started by this thread in this session. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + tag : list of str + The tag to be added. Cannot contain ',' (comma) character or be an empty string. + """ + raise RuntimeError( + "SparkSession.addTag is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def removeTag(self, tag: str) -> None: + """ + Remove a tag previously added to be assigned to all the operations started by this thread in + this session. Noop if such a tag was not added earlier. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + tag : list of str + The tag to be removed. Cannot contain ',' (comma) character or be an empty string. + """ + raise RuntimeError( + "SparkSession.removeTag is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def getTags(self) -> Set[str]: + """ + Get the tags that are currently set to be assigned to all the operations started by this + thread. + + .. versionadded:: 3.5.0 + + Returns + ------- + set of str + Set of tags of interrupted operations. + """ + raise RuntimeError( + "SparkSession.getTags is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + def clearTags(self) -> None: + """ + Clear the current thread's operation tags. + + .. versionadded:: 3.5.0 + """ + raise RuntimeError( + "SparkSession.clearTags is only supported with Spark Connect; " + "however, the current Spark session does not use Spark Connect." + ) + + +def _test() -> None: + import os + import doctest + import pyspark.sql.session + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.session.__dict__.copy() + globs["spark"] = ( + SparkSession.builder.master("local[4]").appName("sql.session tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod( + pyspark.sql.session, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + globs["spark"].stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/sql_formatter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/sql_formatter.py new file mode 100644 index 0000000..5e79b9f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/sql_formatter.py @@ -0,0 +1,84 @@ +# +# 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 string +import typing +from typing import Any, Optional, List, Tuple, Sequence, Mapping +import uuid + +from py4j.java_gateway import is_instance_of + +if typing.TYPE_CHECKING: + from pyspark.sql import SparkSession, DataFrame +from pyspark.sql.functions import lit + + +class SQLStringFormatter(string.Formatter): + """ + A standard ``string.Formatter`` in Python that can understand PySpark instances + with basic Python objects. This object has to be clear after the use for single SQL + query; cannot be reused across multiple SQL queries without cleaning. + """ + + def __init__(self, session: "SparkSession") -> None: + self._session: "SparkSession" = session + self._temp_views: List[Tuple[DataFrame, str]] = [] + + def get_field(self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, Any]) -> Any: + obj, first = super(SQLStringFormatter, self).get_field(field_name, args, kwargs) + return self._convert_value(obj, field_name), first + + def _convert_value(self, val: Any, field_name: str) -> Optional[str]: + """ + Converts the given value into a SQL string. + """ + from pyspark import SparkContext + from pyspark.sql import Column, DataFrame + + if isinstance(val, Column): + assert SparkContext._gateway is not None + + gw = SparkContext._gateway + jexpr = val._jc.expr() + if is_instance_of( + gw, jexpr, "org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute" + ) or is_instance_of( + gw, jexpr, "org.apache.spark.sql.catalyst.expressions.AttributeReference" + ): + return jexpr.sql() + else: + raise ValueError( + "%s in %s should be a plain column reference such as `df.col` " + "or `col('column')`" % (val, field_name) + ) + elif isinstance(val, DataFrame): + for df, n in self._temp_views: + if df is val: + return n + df_name = "_pyspark_%s" % str(uuid.uuid4()).replace("-", "") + self._temp_views.append((val, df_name)) + val.createOrReplaceTempView(df_name) + return df_name + elif isinstance(val, str): + return lit(val)._jc.expr().sql() # for escaped characters. + else: + return val + + def clear(self) -> None: + for _, n in self._temp_views: + self._session.catalog.dropTempView(n) + self._temp_views = [] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/__init__.py new file mode 100644 index 0000000..6b5723d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/__init__.py @@ -0,0 +1,21 @@ +# +# 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.sql.streaming.query import StreamingQuery, StreamingQueryManager # noqa: F401 +from pyspark.sql.streaming.readwriter import DataStreamReader, DataStreamWriter # noqa: F401 +from pyspark.sql.streaming.listener import StreamingQueryListener # noqa: F401 +from pyspark.errors import StreamingQueryException # noqa: F401 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/listener.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/listener.py new file mode 100644 index 0000000..16f4039 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/listener.py @@ -0,0 +1,1048 @@ +# +# 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 uuid +import json +from typing import Any, Dict, List, Optional +from abc import ABC, abstractmethod + +from py4j.java_gateway import JavaObject + +from pyspark.sql import Row +from pyspark import cloudpickle + +__all__ = ["StreamingQueryListener"] + + +class StreamingQueryListener(ABC): + """ + Interface for listening to events related to :class:`~pyspark.sql.streaming.StreamingQuery`. + + .. versionadded:: 3.4.0 + + Notes + ----- + The methods are not thread-safe as they may be called from different threads. + The events received are identical with Scala API. Refer to its documentation. + + This API is evolving. + + Examples + -------- + >>> class MyListener(StreamingQueryListener): + ... def onQueryStarted(self, event: QueryStartedEvent) -> None: + ... # Do something with event. + ... pass + ... + ... def onQueryProgress(self, event: QueryProgressEvent) -> None: + ... # Do something with event. + ... pass + ... + ... def onQueryIdle(self, event: QueryIdleEvent) -> None: + ... # Do something with event. + ... pass + ... + ... def onQueryTerminated(self, event: QueryTerminatedEvent) -> None: + ... # Do something with event. + ... pass + ... + >>> spark.streams.addListener(MyListener()) + """ + + def _set_spark_session( + self, spark: "SparkSession" # type: ignore[name-defined] # noqa: F821 + ) -> None: + self._sparkSession = spark + + @property + def spark(self) -> Optional["SparkSession"]: # type: ignore[name-defined] # noqa: F821 + if hasattr(self, "_sparkSession"): + return self._sparkSession + else: + return None + + def _init_listener_id(self) -> None: + self._id = str(uuid.uuid4()) + + @abstractmethod + def onQueryStarted(self, event: "QueryStartedEvent") -> None: + """ + Called when a query is started. + + Notes + ----- + This is called synchronously with :py:meth:`~pyspark.sql.streaming.DataStreamWriter.start`, + that is, `onQueryStart` will be called on all listeners before `DataStreamWriter.start()` + returns the corresponding :class:`~pyspark.sql.streaming.StreamingQuery`. + Please don't block this method as it will block your query. + """ + pass + + @abstractmethod + def onQueryProgress(self, event: "QueryProgressEvent") -> None: + """ + Called when there is some status update (ingestion rate updated, etc.) + + Notes + ----- + This method is asynchronous. The status in :class:`~pyspark.sql.streaming.StreamingQuery` + will always be latest no matter when this method is called. Therefore, the status of + :class:`~pyspark.sql.streaming.StreamingQuery`. + may be changed before/when you process the event. E.g., you may find + :class:`~pyspark.sql.streaming.StreamingQuery` is terminated when you are + processing `QueryProgressEvent`. + """ + pass + + @abstractmethod + def onQueryIdle(self, event: "QueryIdleEvent") -> None: + """ + Called when the query is idle and waiting for new data to process. + """ + pass + + @abstractmethod + def onQueryTerminated(self, event: "QueryTerminatedEvent") -> None: + """ + Called when a query is stopped, with or without error. + """ + pass + + @property + def _jlistener(self) -> JavaObject: + from pyspark import SparkContext + + if hasattr(self, "_jlistenerobj"): + return self._jlistenerobj + + self._jlistenerobj: JavaObject = ( + SparkContext._jvm.PythonStreamingQueryListenerWrapper( # type: ignore[union-attr] + JStreamingQueryListener(self) + ) + ) + return self._jlistenerobj + + +class JStreamingQueryListener: + """ + Python class that implements Java interface by Py4J. + """ + + def __init__(self, pylistener: StreamingQueryListener) -> None: + self.pylistener = pylistener + + def onQueryStarted(self, jevent: JavaObject) -> None: + self.pylistener.onQueryStarted(QueryStartedEvent.fromJObject(jevent)) + + def onQueryProgress(self, jevent: JavaObject) -> None: + self.pylistener.onQueryProgress(QueryProgressEvent.fromJObject(jevent)) + + def onQueryIdle(self, jevent: JavaObject) -> None: + self.pylistener.onQueryIdle(QueryIdleEvent.fromJObject(jevent)) + + def onQueryTerminated(self, jevent: JavaObject) -> None: + self.pylistener.onQueryTerminated(QueryTerminatedEvent.fromJObject(jevent)) + + class Java: + implements = ["org.apache.spark.sql.streaming.PythonStreamingQueryListener"] + + +class QueryStartedEvent: + """ + Event representing the start of a query. + + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, id: uuid.UUID, runId: uuid.UUID, name: Optional[str], timestamp: str + ) -> None: + self._id: uuid.UUID = id + self._runId: uuid.UUID = runId + self._name: Optional[str] = name + self._timestamp: str = timestamp + + @classmethod + def fromJObject(cls, jevent: JavaObject) -> "QueryStartedEvent": + return cls( + id=uuid.UUID(jevent.id().toString()), + runId=uuid.UUID(jevent.runId().toString()), + name=jevent.name(), + timestamp=jevent.timestamp(), + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "QueryStartedEvent": + return cls( + id=uuid.UUID(j["id"]), + runId=uuid.UUID(j["runId"]), + name=j["name"], + timestamp=j["timestamp"], + ) + + @property + def id(self) -> uuid.UUID: + """ + A unique query id that persists across restarts. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.id`. + """ + return self._id + + @property + def runId(self) -> uuid.UUID: + """ + A query id that is unique for every start/restart. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.runId`. + """ + return self._runId + + @property + def name(self) -> Optional[str]: + """ + User-specified name of the query, `None` if not specified. + """ + return self._name + + @property + def timestamp(self) -> str: + """ + The timestamp to start a query. + """ + return self._timestamp + + +class QueryProgressEvent: + """ + Event representing any progress updates in a query. + + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__(self, progress: "StreamingQueryProgress") -> None: + self._progress: StreamingQueryProgress = progress + + @classmethod + def fromJObject(cls, jevent: JavaObject) -> "QueryProgressEvent": + return cls(progress=StreamingQueryProgress.fromJObject(jevent.progress())) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "QueryProgressEvent": + return cls(progress=StreamingQueryProgress.fromJson(j["progress"])) + + @property + def progress(self) -> "StreamingQueryProgress": + """ + The query progress updates. + """ + return self._progress + + +class QueryIdleEvent: + """ + Event representing that query is idle and waiting for new data to process. + + .. versionadded:: 3.5.0 + + Notes + ----- + This API is evolving. + """ + + def __init__(self, id: uuid.UUID, runId: uuid.UUID, timestamp: str) -> None: + self._id: uuid.UUID = id + self._runId: uuid.UUID = runId + self._timestamp: str = timestamp + + @classmethod + def fromJObject(cls, jevent: JavaObject) -> "QueryIdleEvent": + return cls( + id=uuid.UUID(jevent.id().toString()), + runId=uuid.UUID(jevent.runId().toString()), + timestamp=jevent.timestamp(), + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "QueryIdleEvent": + return cls(id=uuid.UUID(j["id"]), runId=uuid.UUID(j["runId"]), timestamp=j["timestamp"]) + + @property + def id(self) -> uuid.UUID: + """ + A unique query id that persists across restarts. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.id`. + """ + return self._id + + @property + def runId(self) -> uuid.UUID: + """ + A query id that is unique for every start/restart. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.runId`. + """ + return self._runId + + @property + def timestamp(self) -> str: + """ + The timestamp when the latest no-batch trigger happened. + """ + return self._timestamp + + +class QueryTerminatedEvent: + """ + Event representing that termination of a query. + + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + id: uuid.UUID, + runId: uuid.UUID, + exception: Optional[str], + errorClassOnException: Optional[str], + ) -> None: + self._id: uuid.UUID = id + self._runId: uuid.UUID = runId + self._exception: Optional[str] = exception + self._errorClassOnException: Optional[str] = errorClassOnException + + @classmethod + def fromJObject(cls, jevent: JavaObject) -> "QueryTerminatedEvent": + jexception = jevent.exception() + jerrorclass = jevent.errorClassOnException() + return cls( + id=uuid.UUID(jevent.id().toString()), + runId=uuid.UUID(jevent.runId().toString()), + exception=jexception.get() if jexception.isDefined() else None, + errorClassOnException=jerrorclass.get() if jerrorclass.isDefined() else None, + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "QueryTerminatedEvent": + return cls( + id=uuid.UUID(j["id"]), + runId=uuid.UUID(j["runId"]), + exception=j["exception"], + errorClassOnException=j["errorClassOnException"], + ) + + @property + def id(self) -> uuid.UUID: + """ + A unique query id that persists across restarts. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.id`. + """ + return self._id + + @property + def runId(self) -> uuid.UUID: + """ + A query id that is unique for every start/restart. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.runId`. + """ + return self._runId + + @property + def exception(self) -> Optional[str]: + """ + The exception message of the query if the query was terminated + with an exception. Otherwise, it will be `None`. + """ + return self._exception + + @property + def errorClassOnException(self) -> Optional[str]: + """ + The error class from the exception if the query was terminated + with an exception which is a part of error class framework. + If the query was terminated without an exception, or the + exception is not a part of error class framework, it will be + `None`. + + .. versionadded:: 3.5.0 + """ + return self._errorClassOnException + + +class StreamingQueryProgress: + """ + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + id: uuid.UUID, + runId: uuid.UUID, + name: Optional[str], + timestamp: str, + batchId: int, + batchDuration: int, + durationMs: Dict[str, int], + eventTime: Dict[str, str], + stateOperators: List["StateOperatorProgress"], + sources: List["SourceProgress"], + sink: "SinkProgress", + numInputRows: int, + inputRowsPerSecond: float, + processedRowsPerSecond: float, + observedMetrics: Dict[str, Row], + jprogress: Optional[JavaObject] = None, + jdict: Optional[Dict[str, Any]] = None, + ): + self._jprogress: Optional[JavaObject] = jprogress + self._jdict: Optional[Dict[str, Any]] = jdict + self._id: uuid.UUID = id + self._runId: uuid.UUID = runId + self._name: Optional[str] = name + self._timestamp: str = timestamp + self._batchId: int = batchId + self._batchDuration: int = batchDuration + self._durationMs: Dict[str, int] = durationMs + self._eventTime: Dict[str, str] = eventTime + self._stateOperators: List[StateOperatorProgress] = stateOperators + self._sources: List[SourceProgress] = sources + self._sink: SinkProgress = sink + self._numInputRows: int = numInputRows + self._inputRowsPerSecond: float = inputRowsPerSecond + self._processedRowsPerSecond: float = processedRowsPerSecond + self._observedMetrics: Dict[str, Row] = observedMetrics + + @classmethod + def fromJObject(cls, jprogress: JavaObject) -> "StreamingQueryProgress": + from pyspark import SparkContext + + return cls( + jprogress=jprogress, + id=uuid.UUID(jprogress.id().toString()), + runId=uuid.UUID(jprogress.runId().toString()), + name=jprogress.name(), + timestamp=jprogress.timestamp(), + batchId=jprogress.batchId(), + batchDuration=jprogress.batchDuration(), + durationMs=dict(jprogress.durationMs()), + eventTime=dict(jprogress.eventTime()), + stateOperators=[ + StateOperatorProgress.fromJObject(js) for js in jprogress.stateOperators() + ], + sources=[SourceProgress.fromJObject(js) for js in jprogress.sources()], + sink=SinkProgress.fromJObject(jprogress.sink()), + numInputRows=jprogress.numInputRows(), + inputRowsPerSecond=jprogress.inputRowsPerSecond(), + processedRowsPerSecond=jprogress.processedRowsPerSecond(), + observedMetrics={ + k: cloudpickle.loads( + SparkContext._jvm.PythonSQLUtils.toPyRow(jr) # type: ignore[union-attr] + ) + for k, jr in dict(jprogress.observedMetrics()).items() + }, + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "StreamingQueryProgress": + return cls( + jdict=j, + id=uuid.UUID(j["id"]), + runId=uuid.UUID(j["runId"]), + name=j["name"], + timestamp=j["timestamp"], + batchId=j["batchId"], + batchDuration=j.get("batchDuration", None), + durationMs=dict(j["durationMs"]) if "durationMs" in j else {}, + eventTime=dict(j["eventTime"]) if "eventTime" in j else {}, + stateOperators=[StateOperatorProgress.fromJson(s) for s in j["stateOperators"]], + sources=[SourceProgress.fromJson(s) for s in j["sources"]], + sink=SinkProgress.fromJson(j["sink"]), + numInputRows=j["numInputRows"], + inputRowsPerSecond=j["inputRowsPerSecond"], + processedRowsPerSecond=j["processedRowsPerSecond"], + observedMetrics={ + k: Row(*row_dict.keys())(*row_dict.values()) # Assume no nested rows + for k, row_dict in j["observedMetrics"].items() + } + if "observedMetrics" in j + else {}, + ) + + @property + def id(self) -> uuid.UUID: + """ + A unique query id that persists across restarts. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.id`. + """ + return self._id + + @property + def runId(self) -> uuid.UUID: + """ + A query id that is unique for every start/restart. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.runId`. + """ + return self._runId + + @property + def name(self) -> Optional[str]: + """ + User-specified name of the query, `None` if not specified. + """ + return self._name + + @property + def timestamp(self) -> str: + """ + The timestamp to start a query. + """ + return self._timestamp + + @property + def batchId(self) -> int: + """ + A unique id for the current batch of data being processed. Note that in the + case of retries after a failure a given batchId my be executed more than once. + Similarly, when there is no data to be processed, the batchId will not be + incremented. + """ + return self._batchId + + @property + def batchDuration(self) -> int: + """ + The process duration of each batch. + """ + return self._batchDuration + + @property + def durationMs(self) -> Dict[str, int]: + """ + The amount of time taken to perform various operations in milliseconds. + """ + return self._durationMs + + @property + def eventTime(self) -> Dict[str, str]: + """ + Statistics of event time seen in this batch. It may contain the following keys: + + .. code-block:: python + + { + "max": "2016-12-05T20:54:20.827Z", # maximum event time seen in this trigger + "min": "2016-12-05T20:54:20.827Z", # minimum event time seen in this trigger + "avg": "2016-12-05T20:54:20.827Z", # average event time seen in this trigger + "watermark": "2016-12-05T20:54:20.827Z" # watermark used in this trigger + } + + All timestamps are in ISO8601 format, i.e. UTC timestamps. + """ + return self._eventTime + + @property + def stateOperators(self) -> List["StateOperatorProgress"]: + """ + Information about operators in the query that store state. + """ + return self._stateOperators + + @property + def sources(self) -> List["SourceProgress"]: + """ + detailed statistics on data being read from each of the streaming sources. + """ + return self._sources + + @property + def sink(self) -> "SinkProgress": + """ + A unique query id that persists across restarts. See + py:meth:`~pyspark.sql.streaming.StreamingQuery.id`. + """ + return self._sink + + @property + def observedMetrics(self) -> Dict[str, Row]: + return self._observedMetrics + + @property + def numInputRows(self) -> int: + """ + The aggregate (across all sources) number of records processed in a trigger. + """ + return self._numInputRows + + @property + def inputRowsPerSecond(self) -> float: + """ + The aggregate (across all sources) rate of data arriving. + """ + return self._inputRowsPerSecond + + @property + def processedRowsPerSecond(self) -> float: + """ + The aggregate (across all sources) rate at which Spark is processing data. + """ + return self._processedRowsPerSecond + + @property + def json(self) -> str: + """ + The compact JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.json() + else: + return json.dumps(self._jdict) + + @property + def prettyJson(self) -> str: + """ + The pretty (i.e. indented) JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.prettyJson() + else: + return json.dumps(self._jdict, indent=4) + + def __str__(self) -> str: + return self.prettyJson + + +class StateOperatorProgress: + """ + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + operatorName: str, + numRowsTotal: int, + numRowsUpdated: int, + numRowsRemoved: int, + allUpdatesTimeMs: int, + allRemovalsTimeMs: int, + commitTimeMs: int, + memoryUsedBytes: int, + numRowsDroppedByWatermark: int, + numShufflePartitions: int, + numStateStoreInstances: int, + customMetrics: Dict[str, int], + jprogress: Optional[JavaObject] = None, + jdict: Optional[Dict[str, Any]] = None, + ): + self._jprogress: Optional[JavaObject] = jprogress + self._jdict: Optional[Dict[str, Any]] = jdict + self._operatorName: str = operatorName + self._numRowsTotal: int = numRowsTotal + self._numRowsUpdated: int = numRowsUpdated + self._numRowsRemoved: int = numRowsRemoved + self._allUpdatesTimeMs: int = allUpdatesTimeMs + self._allRemovalsTimeMs: int = allRemovalsTimeMs + self._commitTimeMs: int = commitTimeMs + self._memoryUsedBytes: int = memoryUsedBytes + self._numRowsDroppedByWatermark: int = numRowsDroppedByWatermark + self._numShufflePartitions: int = numShufflePartitions + self._numStateStoreInstances: int = numStateStoreInstances + self._customMetrics: Dict[str, int] = customMetrics + + @classmethod + def fromJObject(cls, jprogress: JavaObject) -> "StateOperatorProgress": + return cls( + jprogress=jprogress, + operatorName=jprogress.operatorName(), + numRowsTotal=jprogress.numRowsTotal(), + numRowsUpdated=jprogress.numRowsUpdated(), + allUpdatesTimeMs=jprogress.allUpdatesTimeMs(), + numRowsRemoved=jprogress.numRowsRemoved(), + allRemovalsTimeMs=jprogress.allRemovalsTimeMs(), + commitTimeMs=jprogress.commitTimeMs(), + memoryUsedBytes=jprogress.memoryUsedBytes(), + numRowsDroppedByWatermark=jprogress.numRowsDroppedByWatermark(), + numShufflePartitions=jprogress.numShufflePartitions(), + numStateStoreInstances=jprogress.numStateStoreInstances(), + customMetrics=dict(jprogress.customMetrics()), + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "StateOperatorProgress": + return cls( + jdict=j, + operatorName=j["operatorName"], + numRowsTotal=j["numRowsTotal"], + numRowsUpdated=j["numRowsUpdated"], + numRowsRemoved=j["numRowsRemoved"], + allUpdatesTimeMs=j["allUpdatesTimeMs"], + allRemovalsTimeMs=j["allRemovalsTimeMs"], + commitTimeMs=j["commitTimeMs"], + memoryUsedBytes=j["memoryUsedBytes"], + numRowsDroppedByWatermark=j["numRowsDroppedByWatermark"], + numShufflePartitions=j["numShufflePartitions"], + numStateStoreInstances=j["numStateStoreInstances"], + customMetrics=dict(j["customMetrics"]) if "customMetrics" in j else {}, + ) + + @property + def operatorName(self) -> str: + return self._operatorName + + @property + def numRowsTotal(self) -> int: + return self._numRowsTotal + + @property + def numRowsUpdated(self) -> int: + return self._numRowsUpdated + + @property + def allUpdatesTimeMs(self) -> int: + return self._allUpdatesTimeMs + + @property + def numRowsRemoved(self) -> int: + return self._numRowsRemoved + + @property + def allRemovalsTimeMs(self) -> int: + return self._allRemovalsTimeMs + + @property + def commitTimeMs(self) -> int: + return self._commitTimeMs + + @property + def memoryUsedBytes(self) -> int: + return self._memoryUsedBytes + + @property + def numRowsDroppedByWatermark(self) -> int: + return self._numRowsDroppedByWatermark + + @property + def numShufflePartitions(self) -> int: + return self._numShufflePartitions + + @property + def numStateStoreInstances(self) -> int: + return self._numStateStoreInstances + + @property + def customMetrics(self) -> Dict[str, int]: + return self._customMetrics + + @property + def json(self) -> str: + """ + The compact JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.json() + else: + return json.dumps(self._jdict) + + @property + def prettyJson(self) -> str: + """ + The pretty (i.e. indented) JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.prettyJson() + else: + return json.dumps(self._jdict, indent=4) + + def __str__(self) -> str: + return self.prettyJson + + +class SourceProgress: + """ + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + description: str, + startOffset: str, + endOffset: str, + latestOffset: str, + numInputRows: int, + inputRowsPerSecond: float, + processedRowsPerSecond: float, + metrics: Dict[str, str], + jprogress: Optional[JavaObject] = None, + jdict: Optional[Dict[str, Any]] = None, + ) -> None: + self._jprogress: Optional[JavaObject] = jprogress + self._jdict: Optional[Dict[str, Any]] = jdict + self._description: str = description + self._startOffset: str = startOffset + self._endOffset: str = endOffset + self._latestOffset: str = latestOffset + self._numInputRows: int = numInputRows + self._inputRowsPerSecond: float = inputRowsPerSecond + self._processedRowsPerSecond: float = processedRowsPerSecond + self._metrics: Dict[str, str] = metrics + + @classmethod + def fromJObject(cls, jprogress: JavaObject) -> "SourceProgress": + return cls( + jprogress=jprogress, + description=jprogress.description(), + startOffset=str(jprogress.startOffset()), + endOffset=str(jprogress.endOffset()), + latestOffset=str(jprogress.latestOffset()), + numInputRows=jprogress.numInputRows(), + inputRowsPerSecond=jprogress.inputRowsPerSecond(), + processedRowsPerSecond=jprogress.processedRowsPerSecond(), + metrics=dict(jprogress.metrics()), + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "SourceProgress": + return cls( + jdict=j, + description=j["description"], + startOffset=str(j["startOffset"]), + endOffset=str(j["endOffset"]), + latestOffset=str(j["latestOffset"]), + numInputRows=j["numInputRows"], + inputRowsPerSecond=j["inputRowsPerSecond"], + processedRowsPerSecond=j["processedRowsPerSecond"], + metrics=dict(j["metrics"]) if "metrics" in j else {}, + ) + + @property + def description(self) -> str: + """ + Description of the source. + """ + return self._description + + @property + def startOffset(self) -> str: + """ + The starting offset for data being read. + """ + return self._startOffset + + @property + def endOffset(self) -> str: + """ + The ending offset for data being read. + """ + return self._endOffset + + @property + def latestOffset(self) -> str: + """ + The latest offset from this source. + """ + return self._latestOffset + + @property + def numInputRows(self) -> int: + """ + The number of records read from this source. + """ + return self._numInputRows + + @property + def inputRowsPerSecond(self) -> float: + """ + The rate at which data is arriving from this source. + """ + return self._inputRowsPerSecond + + @property + def processedRowsPerSecond(self) -> float: + """ + The rate at which data from this source is being processed by Spark. + """ + return self._processedRowsPerSecond + + @property + def metrics(self) -> Dict[str, str]: + return self._metrics + + @property + def json(self) -> str: + """ + The compact JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.json() + else: + return json.dumps(self._jdict) + + @property + def prettyJson(self) -> str: + """ + The pretty (i.e. indented) JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.prettyJson() + else: + return json.dumps(self._jdict, indent=4) + + def __str__(self) -> str: + return self.prettyJson + + +class SinkProgress: + """ + .. versionadded:: 3.4.0 + + Notes + ----- + This API is evolving. + """ + + def __init__( + self, + description: str, + numOutputRows: int, + metrics: Dict[str, str], + jprogress: Optional[JavaObject] = None, + jdict: Optional[Dict[str, Any]] = None, + ) -> None: + self._jprogress: Optional[JavaObject] = jprogress + self._jdict: Optional[Dict[str, Any]] = jdict + self._description: str = description + self._numOutputRows: int = numOutputRows + self._metrics: Dict[str, str] = metrics + + @classmethod + def fromJObject(cls, jprogress: JavaObject) -> "SinkProgress": + return cls( + jprogress=jprogress, + description=jprogress.description(), + numOutputRows=jprogress.numOutputRows(), + metrics=dict(jprogress.metrics()), + ) + + @classmethod + def fromJson(cls, j: Dict[str, Any]) -> "SinkProgress": + return cls( + jdict=j, + description=j["description"], + numOutputRows=j["numOutputRows"], + metrics=dict(j["metrics"]) if "metrics" in j else {}, + ) + + @property + def description(self) -> str: + """ + Description of the source. + """ + return self._description + + @property + def numOutputRows(self) -> int: + """ + Number of rows written to the sink or -1 for Continuous Mode (temporarily) + or Sink V1 (until decommissioned). + """ + return self._numOutputRows + + @property + def metrics(self) -> Dict[str, str]: + return self._metrics + + @property + def json(self) -> str: + """ + The compact JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.json() + else: + return json.dumps(self._jdict) + + @property + def prettyJson(self) -> str: + """ + The pretty (i.e. indented) JSON representation of this progress. + """ + assert self._jdict is not None or self._jprogress is not None + if self._jprogress: + return self._jprogress.prettyJson() + else: + return json.dumps(self._jdict, indent=4) + + def __str__(self) -> str: + return self.prettyJson + + +def _test() -> None: + import sys + import doctest + import os + from pyspark.sql import SparkSession + import pyspark.sql.streaming.listener + from py4j.protocol import Py4JError + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.streaming.listener.__dict__.copy() + try: + spark = SparkSession._getActiveSessionOrCreate() + except Py4JError: # noqa: F821 + spark = SparkSession(sc) # type: ignore[name-defined] # noqa: F821 + + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.streaming.listener, + globs=globs, + ) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/query.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/query.py new file mode 100644 index 0000000..db104e3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/query.py @@ -0,0 +1,746 @@ +# +# 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 json +from typing import Any, Dict, List, Optional + +from py4j.java_gateway import JavaObject, java_import + +from pyspark.errors import StreamingQueryException, PySparkValueError +from pyspark.errors.exceptions.captured import ( + StreamingQueryException as CapturedStreamingQueryException, +) +from pyspark.sql.streaming.listener import StreamingQueryListener + +__all__ = ["StreamingQuery", "StreamingQueryManager"] + + +class StreamingQuery: + """ + A handle to a query that is executing continuously in the background as new data arrives. + All these methods are thread-safe. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + """ + + def __init__(self, jsq: JavaObject) -> None: + self._jsq = jsq + + @property + def id(self) -> str: + """ + Returns the unique id of this query that persists across restarts from checkpoint data. + That is, this id is generated when a query is started for the first time, and + will be the same every time it is restarted from checkpoint data. + There can only be one query with the same id active in a Spark cluster. + Also see, `runId`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + str + The unique id of query that persists across restarts from checkpoint data. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get the unique id of this query that persists across restarts from checkpoint data + + >>> sq.id + '...' + + >>> sq.stop() + """ + return self._jsq.id().toString() + + @property + def runId(self) -> str: + """ + Returns the unique id of this query that does not persist across restarts. That is, every + query that is started (or restarted from checkpoint) will have a different runId. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + str + The unique id of query that does not persist across restarts. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get the unique id of this query that does not persist across restarts + + >>> sq.runId + '...' + + >>> sq.stop() + """ + return self._jsq.runId().toString() + + @property + def name(self) -> str: + """ + Returns the user-specified name of the query, or null if not specified. + This name can be specified in the `org.apache.spark.sql.streaming.DataStreamWriter` + as `dataframe.writeStream.queryName("query").start()`. + This name, if set, must be unique across all active queries. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + str + The user-specified name of the query, or null if not specified. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get the user-specified name of the query, or null if not specified. + + >>> sq.name + 'this_query' + + >>> sq.stop() + """ + return self._jsq.name() + + @property + def isActive(self) -> bool: + """ + Whether this streaming query is currently active or not. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + bool + The result whether specified streaming query is currently active or not. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + >>> sq.isActive + True + + >>> sq.stop() + """ + return self._jsq.isActive() + + def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]: + """ + Waits for the termination of `this` query, either by :func:`query.stop()` or by an + exception. If the query has terminated with an exception, then the exception will be thrown. + If `timeout` is set, it returns whether the query has terminated or not within the + `timeout` seconds. + + If the query has terminated, then all subsequent calls to this method will either return + immediately (if the query was terminated by :func:`stop()`), or throw the exception + immediately (if the query has terminated with exception). + + throws :class:`StreamingQueryException`, if `this` query has terminated with an exception + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + timeout : int, optional + default ``None``. The waiting time for specified streaming query to terminate. + + Returns + ------- + bool, optional + The result whether specified streaming query has terminated or not within the `timeout` + seconds if `timeout` is set. The :class:`StreamingQueryException` will be thrown if the + query has terminated with an exception. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('query_awaitTermination').start() + + Return whether the query has terminated or not within 5 seconds + + >>> sq.awaitTermination(5) + False + + >>> sq.stop() + """ + if timeout is not None: + if not isinstance(timeout, (int, float)) or timeout <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "timeout", "arg_value": type(timeout).__name__}, + ) + return self._jsq.awaitTermination(int(timeout * 1000)) + else: + return self._jsq.awaitTermination() + + @property + def status(self) -> Dict[str, Any]: + """ + Returns the current status of the query. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + dict + The current status of the specified query. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get the current status of the query + + >>> sq.status + {'message': '...', 'isDataAvailable': ..., 'isTriggerActive': ...} + + >>> sq.stop() + """ + return json.loads(self._jsq.status().json()) + + @property + def recentProgress(self) -> List[Dict[str, Any]]: + """ + Returns an array of the most recent [[StreamingQueryProgress]] updates for this query. + The number of progress updates retained for each stream is configured by Spark session + configuration `spark.sql.streaming.numRecentProgressUpdates`. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + list + List of dict which is the most recent :class:`StreamingQueryProgress` updates + for this query. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get an array of the most recent query progress updates for this query + + >>> sq.recentProgress + [...] + + >>> sq.stop() + """ + return [json.loads(p.json()) for p in self._jsq.recentProgress()] + + @property + def lastProgress(self) -> Optional[Dict[str, Any]]: + """ + Returns the most recent :class:`StreamingQueryProgress` update of this streaming query or + None if there were no progress updates + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + dict, optional + The most recent :class:`StreamingQueryProgress` update of this streaming query or + None if there were no progress updates. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Get the most recent query progress updates for this query + + >>> sq.lastProgress + >>> sq.stop() + """ + lastProgress = self._jsq.lastProgress() + if lastProgress: + return json.loads(lastProgress.json()) + else: + return None + + def processAllAvailable(self) -> None: + """ + Blocks until all available data in the source has been processed and committed to the + sink. This method is intended for testing. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + In the case of continually arriving data, this method may block forever. + Additionally, this method is only guaranteed to block until data that has been + synchronously appended data to a stream source prior to invocation. + (i.e. `getOffset` must immediately reflect the addition). + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Blocks query until all available data in the source + has been processed and committed to the sink + + >>> sq.processAllAvailable + + + >>> sq.stop() + """ + return self._jsq.processAllAvailable() + + def stop(self) -> None: + """ + Stop this streaming query. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + >>> sq.isActive + True + + Stop streaming query + + >>> sq.stop() + + >>> sq.isActive + False + """ + self._jsq.stop() + + def explain(self, extended: bool = False) -> None: + """ + Prints the (logical and physical) plans to the console for debugging purpose. + + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + extended : bool, optional + default ``False``. If ``False``, prints only the physical plan. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sdf.printSchema() + root + |-- timestamp: timestamp (nullable = true) + |-- value: long (nullable = true) + + >>> sq = sdf.writeStream.format('memory').queryName('query_explain').start() + >>> sq.processAllAvailable() # Wait a bit to generate the runtime plans. + + Explain the runtime plans + + >>> sq.explain() + == Physical Plan == + ... + >>> sq.explain(True) + == Parsed Logical Plan == + ... + == Analyzed Logical Plan == + ... + == Optimized Logical Plan == + ... + == Physical Plan == + ... + >>> sq.stop() + """ + # Cannot call `_jsq.explain(...)` because it will print in the JVM process. + # We should print it in the Python process. + print(self._jsq.explainInternal(extended)) + + def exception(self) -> Optional[StreamingQueryException]: + """ + .. versionadded:: 2.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + :class:`StreamingQueryException` + the StreamingQueryException if the query was terminated by an exception, or None. + """ + if self._jsq.exception().isDefined(): + je = self._jsq.exception().get() + msg = je.toString().split(": ", 1)[1] # Drop the Java StreamingQueryException type info + stackTrace = "\n\t at ".join(map(lambda x: x.toString(), je.getStackTrace())) + return CapturedStreamingQueryException(msg, stackTrace, je.getCause()) + else: + return None + + +class StreamingQueryManager: + """A class to manage all the :class:`StreamingQuery` StreamingQueries active. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + """ + + def __init__(self, jsqm: JavaObject) -> None: + self._jsqm = jsqm + + @property + def active(self) -> List[StreamingQuery]: + """ + Returns a list of active queries associated with this SparkSession + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Returns + ------- + list + The active queries associated with this :class:`SparkSession`. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sdf.printSchema() + root + |-- timestamp: timestamp (nullable = true) + |-- value: long (nullable = true) + + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + >>> sqm = spark.streams + + Get the list of active streaming queries + + >>> [q.name for q in sqm.active] + ['this_query'] + >>> sq.stop() + """ + return [StreamingQuery(jsq) for jsq in self._jsqm.active()] + + def get(self, id: str) -> Optional[StreamingQuery]: + """ + Returns an active query from this :class:`SparkSession`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + id : str + The unique id of specified query. + + Returns + ------- + :class:`StreamingQuery` + An active query with `id` from this SparkSession. + + Notes + ----- + Exception will be thrown if an active query with this id does not exist. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sdf.printSchema() + root + |-- timestamp: timestamp (nullable = true) + |-- value: long (nullable = true) + + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + >>> sq.name + 'this_query' + + Get an active query by id + + >>> sq = spark.streams.get(sq.id) + >>> sq.isActive + True + >>> sq.stop() + """ + query = self._jsqm.get(id) + if query is not None: + return StreamingQuery(query) + else: + return None + + def awaitAnyTermination(self, timeout: Optional[int] = None) -> Optional[bool]: + """ + Wait until any of the queries on the associated SparkSession has terminated since the + creation of the context, or since :func:`resetTerminated()` was called. If any query was + terminated with an exception, then the exception will be thrown. + If `timeout` is set, it returns whether the query has terminated or not within the + `timeout` seconds. + + If a query has terminated, then subsequent calls to :func:`awaitAnyTermination()` will + either return immediately (if the query was terminated by :func:`query.stop()`), + or throw the exception immediately (if the query was terminated with exception). Use + :func:`resetTerminated()` to clear past terminations and wait for new terminations. + + In the case where multiple queries have terminated since :func:`resetTermination()` + was called, if any query has terminated with exception, then :func:`awaitAnyTermination()` + will throw any of the exception. For correctly documenting exceptions across multiple + queries, users need to stop all of them after any of them terminates with exception, and + then check the `query.exception()` for each query. + + throws :class:`StreamingQueryException`, if `this` query has terminated with an exception + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + timeout : int, optional + default ``None``. The waiting time for any streaming query to terminate. + + Returns + ------- + bool, optional + The result whether any streaming query has terminated or not within the `timeout` + seconds if `timeout` is set. The :class:`StreamingQueryException` will be thrown if any + query has terminated with an exception. + + Examples + -------- + >>> sdf = spark.readStream.format("rate").load() + >>> sq = sdf.writeStream.format('memory').queryName('this_query').start() + + Return whether any of the query on the associated SparkSession + has terminated or not within 5 seconds + + >>> spark.streams.awaitAnyTermination(5) + True + >>> sq.stop() + """ + if timeout is not None: + if not isinstance(timeout, (int, float)) or timeout < 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "timeout", "arg_value": type(timeout).__name__}, + ) + return self._jsqm.awaitAnyTermination(int(timeout * 1000)) + else: + return self._jsqm.awaitAnyTermination() + + def resetTerminated(self) -> None: + """ + Forget about past terminated queries so that :func:`awaitAnyTermination()` can be used + again to wait for new terminations. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Examples + -------- + >>> spark.streams.resetTerminated() + """ + self._jsqm.resetTerminated() + + def addListener(self, listener: StreamingQueryListener) -> None: + """ + Register a :class:`StreamingQueryListener` to receive up-calls for life cycle events of + :class:`~pyspark.sql.streaming.StreamingQuery`. + + .. versionadded:: 3.4.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + listener : :class:`StreamingQueryListener` + A :class:`StreamingQueryListener` to receive up-calls for life cycle events of + :class:`~pyspark.sql.streaming.StreamingQuery`. + + Notes + ----- + This function behaves differently in Spark Connect mode. + In Connect, the provided functions doesn't have access to variables defined outside of it. + Also in Connect, you need to use `self.spark` to access spark session. + Using `spark` would throw an exception. + In short, if you want to use spark session inside the listener, + please use `self.spark` in Connect mode, and use `spark` otherwise. + + Examples + -------- + >>> from pyspark.sql.streaming import StreamingQueryListener + >>> class TestListener(StreamingQueryListener): + ... def onQueryStarted(self, event): + ... pass + ... + ... def onQueryProgress(self, event): + ... pass + ... + ... def onQueryIdle(self, event): + ... pass + ... + ... def onQueryTerminated(self, event): + ... pass + ... + >>> test_listener = TestListener() + + Register streaming query listener + + >>> spark.streams.addListener(test_listener) + + Deregister streaming query listener + + >>> spark.streams.removeListener(test_listener) + """ + from pyspark import SparkContext + from pyspark.java_gateway import ensure_callback_server_started + + gw = SparkContext._gateway + assert gw is not None + java_import(gw.jvm, "org.apache.spark.sql.streaming.*") + ensure_callback_server_started(gw) + + self._jsqm.addListener(listener._jlistener) + + def removeListener(self, listener: StreamingQueryListener) -> None: + """ + Deregister a :class:`StreamingQueryListener`. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + listener : :class:`StreamingQueryListener` + A :class:`StreamingQueryListener` to receive up-calls for life cycle events of + :class:`~pyspark.sql.streaming.StreamingQuery`. + + Examples + -------- + >>> from pyspark.sql.streaming import StreamingQueryListener + >>> class TestListener(StreamingQueryListener): + ... def onQueryStarted(self, event): + ... pass + ... + ... def onQueryProgress(self, event): + ... pass + ... + ... def onQueryIdle(self, event): + ... pass + ... + ... def onQueryTerminated(self, event): + ... pass + ... + >>> test_listener = TestListener() + + Register streaming query listener + + >>> spark.streams.addListener(test_listener) + + Deregister streaming query listener + + >>> spark.streams.removeListener(test_listener) + """ + self._jsqm.removeListener(listener._jlistener) + + +def _test() -> None: + import doctest + import os + import sys + from pyspark.sql import SparkSession + import pyspark.sql.streaming.query + from py4j.protocol import Py4JError + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.streaming.query.__dict__.copy() + try: + spark = SparkSession._getActiveSessionOrCreate() + except Py4JError: # noqa: F821 + spark = SparkSession(sc) # type: ignore[name-defined] # noqa: F821 + + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.streaming.query, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/readwriter.py new file mode 100644 index 0000000..2026651 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/readwriter.py @@ -0,0 +1,1652 @@ +# +# 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 sys +from collections.abc import Iterator +from typing import cast, overload, Any, Callable, List, Optional, TYPE_CHECKING, Union + +from py4j.java_gateway import java_import, JavaObject + +from pyspark.sql.column import _to_seq +from pyspark.sql.readwriter import OptionUtils, to_str +from pyspark.sql.streaming.query import StreamingQuery +from pyspark.sql.types import Row, StructType +from pyspark.sql.utils import ForeachBatchFunction +from pyspark.errors import ( + PySparkTypeError, + PySparkValueError, + PySparkAttributeError, + PySparkRuntimeError, +) + +if TYPE_CHECKING: + from pyspark.sql.session import SparkSession + from pyspark.sql._typing import SupportsProcess, OptionalPrimitiveType + from pyspark.sql.dataframe import DataFrame + +__all__ = ["DataStreamReader", "DataStreamWriter"] + + +class DataStreamReader(OptionUtils): + """ + Interface used to load a streaming :class:`DataFrame ` from external + storage systems (e.g. file systems, key-value stores, etc). + Use :attr:`SparkSession.readStream ` to access this. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> spark.readStream + <...streaming.readwriter.DataStreamReader object ...> + + The example below uses Rate source that generates rows continuously. + After that, we operate a modulo by 3, and then writes the stream out to the console. + The streaming query stops in 3 seconds. + + >>> import time + >>> df = spark.readStream.format("rate").load() + >>> df = df.selectExpr("value % 3 as v") + >>> q = df.writeStream.format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + + def __init__(self, spark: "SparkSession") -> None: + self._jreader = spark._jsparkSession.readStream() + self._spark = spark + + def _df(self, jdf: JavaObject) -> "DataFrame": + from pyspark.sql.dataframe import DataFrame + + return DataFrame(jdf, self._spark) + + def format(self, source: str) -> "DataStreamReader": + """Specifies the input data source format. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + source : str + name of the data source, e.g. 'json', 'parquet'. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> spark.readStream.format("text") + <...streaming.readwriter.DataStreamReader object ...> + + This API allows to configure other sources to read. The example below writes a small text + file, and reads it back via Text source. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary text file to read it. + ... spark.createDataFrame( + ... [("hello",), ("this",)]).write.mode("overwrite").format("text").save(d) + ... + ... # Start a streaming query to read the text file. + ... q = spark.readStream.format("text").load(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._jreader = self._jreader.format(source) + return self + + def schema(self, schema: Union[StructType, str]) -> "DataStreamReader": + """Specifies the input schema. + + Some data sources (e.g. JSON) can infer the input schema automatically from data. + By specifying the schema here, the underlying data source can skip the schema + inference step, and thus speed up data loading. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + schema : :class:`pyspark.sql.types.StructType` or str + a :class:`pyspark.sql.types.StructType` object or a DDL-formatted string + (For example ``col0 INT, col1 DOUBLE``). + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> from pyspark.sql.types import StructField, StructType, StringType + >>> spark.readStream.schema(StructType([StructField("data", StringType(), True)])) + <...streaming.readwriter.DataStreamReader object ...> + >>> spark.readStream.schema("col0 INT, col1 DOUBLE") + <...streaming.readwriter.DataStreamReader object ...> + + The example below specifies a different schema to CSV file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Start a streaming query to read the CSV file. + ... spark.readStream.schema("col0 INT, col1 STRING").format("csv").load(d).printSchema() + root + |-- col0: integer (nullable = true) + |-- col1: string (nullable = true) + """ + from pyspark.sql import SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + if isinstance(schema, StructType): + jschema = spark._jsparkSession.parseDataType(schema.json()) + self._jreader = self._jreader.schema(jschema) + elif isinstance(schema, str): + self._jreader = self._jreader.schema(schema) + else: + raise PySparkTypeError( + error_class="NOT_STR_OR_STRUCT", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + return self + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataStreamReader": + """Adds an input option for the underlying data source. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> spark.readStream.option("x", 1) + <...streaming.readwriter.DataStreamReader object ...> + + The example below specifies 'rowsPerSecond' option to Rate source in order to generate + 10 rows every second. + + >>> import time + >>> q = spark.readStream.format( + ... "rate").option("rowsPerSecond", 10).load().writeStream.format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + self._jreader = self._jreader.option(key, to_str(value)) + return self + + def options(self, **options: "OptionalPrimitiveType") -> "DataStreamReader": + """Adds input options for the underlying data source. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> spark.readStream.options(x="1", y=2) + <...streaming.readwriter.DataStreamReader object ...> + + The example below specifies 'rowsPerSecond' and 'numPartitions' options to + Rate source in order to generate 10 rows with 10 partitions every second. + + >>> import time + >>> q = spark.readStream.format("rate").options( + ... rowsPerSecond=10, numPartitions=10 + ... ).load().writeStream.format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + for k in options: + self._jreader = self._jreader.option(k, to_str(options[k])) + return self + + def load( + self, + path: Optional[str] = None, + format: Optional[str] = None, + schema: Optional[Union[StructType, str]] = None, + **options: "OptionalPrimitiveType", + ) -> "DataFrame": + """Loads a data stream from a data source and returns it as a + :class:`DataFrame `. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + path : str, optional + optional string for file-system backed data sources. + format : str, optional + optional string for format of the data source. Default to 'parquet'. + schema : :class:`pyspark.sql.types.StructType` or str, optional + optional :class:`pyspark.sql.types.StructType` for the input schema + or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + **options : dict + all other string options + + Notes + ----- + This API is evolving. + + Examples + -------- + Load a data stream from a temporary JSON file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary JSON file to read it. + ... spark.createDataFrame( + ... [(100, "Hyukjin Kwon"),], ["age", "name"] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Start a streaming query to read the JSON file. + ... q = spark.readStream.schema( + ... "age INT, name STRING" + ... ).format("json").load(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + if format is not None: + self.format(format) + if schema is not None: + self.schema(schema) + self.options(**options) + if path is not None: + if type(path) != str or len(path.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "path", "arg_value": str(path)}, + ) + return self._df(self._jreader.load(path)) + else: + return self._df(self._jreader.load()) + + def json( + self, + path: str, + schema: Optional[Union[StructType, str]] = None, + primitivesAsString: Optional[Union[bool, str]] = None, + prefersDecimal: Optional[Union[bool, str]] = None, + allowComments: Optional[Union[bool, str]] = None, + allowUnquotedFieldNames: Optional[Union[bool, str]] = None, + allowSingleQuotes: Optional[Union[bool, str]] = None, + allowNumericLeadingZero: Optional[Union[bool, str]] = None, + allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + allowUnquotedControlChars: Optional[Union[bool, str]] = None, + lineSep: Optional[str] = None, + locale: Optional[str] = None, + dropFieldIfAllNull: Optional[Union[bool, str]] = None, + encoding: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + allowNonNumericNumbers: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """ + Loads a JSON file stream and returns the results as a :class:`DataFrame`. + + `JSON Lines `_ (newline-delimited JSON) is supported by default. + For JSON (one record per file), set the ``multiLine`` parameter to ``true``. + + If the ``schema`` parameter is not specified, this function goes + through the input once to determine the input schema. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + string represents path to the JSON dataset, + or RDD of Strings storing JSON objects. + schema : :class:`pyspark.sql.types.StructType` or str, optional + an optional :class:`pyspark.sql.types.StructType` for the input schema + or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + in the version you use. + + .. # noqa + + Notes + ----- + This API is evolving. + + Examples + -------- + Load a data stream from a temporary JSON file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary JSON file to read it. + ... spark.createDataFrame( + ... [(100, "Hyukjin Kwon"),], ["age", "name"] + ... ).write.mode("overwrite").format("json").save(d) + ... + ... # Start a streaming query to read the JSON file. + ... q = spark.readStream.schema( + ... "age INT, name STRING" + ... ).json(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._set_opts( + schema=schema, + primitivesAsString=primitivesAsString, + prefersDecimal=prefersDecimal, + allowComments=allowComments, + allowUnquotedFieldNames=allowUnquotedFieldNames, + allowSingleQuotes=allowSingleQuotes, + allowNumericLeadingZero=allowNumericLeadingZero, + allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + multiLine=multiLine, + allowUnquotedControlChars=allowUnquotedControlChars, + lineSep=lineSep, + locale=locale, + dropFieldIfAllNull=dropFieldIfAllNull, + encoding=encoding, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + allowNonNumericNumbers=allowNonNumericNumbers, + ) + if isinstance(path, str): + return self._df(self._jreader.json(path)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + def orc( + self, + path: str, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """Loads a ORC file stream, returning the result as a :class:`DataFrame`. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + in the version you use. + + .. # noqa + + Examples + -------- + Load a data stream from a temporary ORC file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary ORC file to read it. + ... spark.range(10).write.mode("overwrite").format("orc").save(d) + ... + ... # Start a streaming query to read the ORC file. + ... q = spark.readStream.schema("id LONG").orc(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + return self._df(self._jreader.orc(path)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + def parquet( + self, + path: str, + mergeSchema: Optional[bool] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + datetimeRebaseMode: Optional[Union[bool, str]] = None, + int96RebaseMode: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """ + Loads a Parquet file stream, returning the result as a :class:`DataFrame`. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + path : str + the path in any Hadoop supported file system + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_. + in the version you use. + + .. # noqa + + Examples + -------- + Load a data stream from a temporary Parquet file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary Parquet file to read it. + ... spark.range(10).write.mode("overwrite").format("parquet").save(d) + ... + ... # Start a streaming query to read the Parquet file. + ... q = spark.readStream.schema( + ... "id LONG").parquet(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._set_opts( + mergeSchema=mergeSchema, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + datetimeRebaseMode=datetimeRebaseMode, + int96RebaseMode=int96RebaseMode, + ) + if isinstance(path, str): + return self._df(self._jreader.parquet(path)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + def text( + self, + path: str, + wholetext: bool = False, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + ) -> "DataFrame": + """ + Loads a text file stream and returns a :class:`DataFrame` whose schema starts with a + string column named "value", and followed by partitioned columns if there + are any. + The text files must be encoded as UTF-8. + + By default, each line in the text file is a new row in the resulting DataFrame. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + path : str or list + string, or list of strings, for input path(s). + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + in the version you use. + + .. # noqa + + Notes + ----- + This API is evolving. + + Examples + -------- + Load a data stream from a temporary text file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary text file to read it. + ... spark.createDataFrame( + ... [("hello",), ("this",)]).write.mode("overwrite").format("text").save(d) + ... + ... # Start a streaming query to read the text file. + ... q = spark.readStream.text(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._set_opts( + wholetext=wholetext, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + ) + if isinstance(path, str): + return self._df(self._jreader.text(path)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + def csv( + self, + path: str, + schema: Optional[Union[StructType, str]] = None, + sep: Optional[str] = None, + encoding: Optional[str] = None, + quote: Optional[str] = None, + escape: Optional[str] = None, + comment: Optional[str] = None, + header: Optional[Union[bool, str]] = None, + inferSchema: Optional[Union[bool, str]] = None, + ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, + ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, + nullValue: Optional[str] = None, + nanValue: Optional[str] = None, + positiveInf: Optional[str] = None, + negativeInf: Optional[str] = None, + dateFormat: Optional[str] = None, + timestampFormat: Optional[str] = None, + maxColumns: Optional[Union[int, str]] = None, + maxCharsPerColumn: Optional[Union[int, str]] = None, + maxMalformedLogPerPartition: Optional[Union[int, str]] = None, + mode: Optional[str] = None, + columnNameOfCorruptRecord: Optional[str] = None, + multiLine: Optional[Union[bool, str]] = None, + charToEscapeQuoteEscaping: Optional[Union[bool, str]] = None, + enforceSchema: Optional[Union[bool, str]] = None, + emptyValue: Optional[str] = None, + locale: Optional[str] = None, + lineSep: Optional[str] = None, + pathGlobFilter: Optional[Union[bool, str]] = None, + recursiveFileLookup: Optional[Union[bool, str]] = None, + unescapedQuoteHandling: Optional[str] = None, + ) -> "DataFrame": + r"""Loads a CSV file stream and returns the result as a :class:`DataFrame`. + + This function will go through the input once to determine the input schema if + ``inferSchema`` is enabled. To avoid going through the entire data once, disable + ``inferSchema`` option or specify the schema explicitly using ``schema``. + + Parameters + ---------- + path : str or list + string, or list of strings, for input path(s). + schema : :class:`pyspark.sql.types.StructType` or str, optional + an optional :class:`pyspark.sql.types.StructType` for the input schema + or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``). + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Other Parameters + ---------------- + Extra options + For the extra options, refer to + `Data Source Option `_ + in the version you use. + + .. # noqa + + Notes + ----- + This API is evolving. + + Examples + -------- + Load a data stream from a temporary CSV file. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d: + ... # Write a temporary text file to read it. + ... spark.createDataFrame([(1, "2"),]).write.mode("overwrite").format("csv").save(d) + ... + ... # Start a streaming query to read the CSV file. + ... q = spark.readStream.schema( + ... "col0 INT, col1 STRING" + ... ).format("csv").load(d).writeStream.format("console").start() + ... time.sleep(3) + ... q.stop() + """ + self._set_opts( + schema=schema, + sep=sep, + encoding=encoding, + quote=quote, + escape=escape, + comment=comment, + header=header, + inferSchema=inferSchema, + ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, + ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, + nullValue=nullValue, + nanValue=nanValue, + positiveInf=positiveInf, + negativeInf=negativeInf, + dateFormat=dateFormat, + timestampFormat=timestampFormat, + maxColumns=maxColumns, + maxCharsPerColumn=maxCharsPerColumn, + maxMalformedLogPerPartition=maxMalformedLogPerPartition, + mode=mode, + columnNameOfCorruptRecord=columnNameOfCorruptRecord, + multiLine=multiLine, + charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, + enforceSchema=enforceSchema, + emptyValue=emptyValue, + locale=locale, + lineSep=lineSep, + pathGlobFilter=pathGlobFilter, + recursiveFileLookup=recursiveFileLookup, + unescapedQuoteHandling=unescapedQuoteHandling, + ) + if isinstance(path, str): + return self._df(self._jreader.csv(path)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "path", "arg_type": type(path).__name__}, + ) + + def table(self, tableName: str) -> "DataFrame": + """Define a Streaming DataFrame on a Table. The DataSource corresponding to the table should + support streaming mode. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + tableName : str + string, for the name of the table. + + Returns + ------- + :class:`DataFrame` + + Notes + ----- + This API is evolving. + + Examples + -------- + Load a data stream from a table. + + >>> import tempfile + >>> import time + >>> _ = spark.sql("DROP TABLE IF EXISTS my_table") + >>> with tempfile.TemporaryDirectory() as d: + ... # Create a table with Rate source. + ... q1 = spark.readStream.format("rate").load().writeStream.toTable( + ... "my_table", checkpointLocation=d) + ... + ... # Read the table back and print out in the console. + ... q2 = spark.readStream.table("my_table").writeStream.format("console").start() + ... time.sleep(3) + ... q1.stop() + ... q2.stop() + ... _ = spark.sql("DROP TABLE my_table") + """ + if isinstance(tableName, str): + return self._df(self._jreader.table(tableName)) + else: + raise PySparkTypeError( + error_class="NOT_STR", + message_parameters={"arg_name": "tableName", "arg_type": type(tableName).__name__}, + ) + + +class DataStreamWriter: + """ + Interface used to write a streaming :class:`DataFrame ` to external + storage systems (e.g. file systems, key-value stores, etc). + Use :attr:`DataFrame.writeStream ` + to access this. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + The example below uses Rate source that generates rows continuously. + After that, we operate a modulo by 3, and then writes the stream out to the console. + The streaming query stops in 3 seconds. + + >>> import time + >>> df = spark.readStream.format("rate").load() + >>> df = df.selectExpr("value % 3 as v") + >>> q = df.writeStream.format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + + def __init__(self, df: "DataFrame") -> None: + self._df = df + self._spark = df.sparkSession + self._jwrite = df._jdf.writeStream() + + def _sq(self, jsq: JavaObject) -> StreamingQuery: + return StreamingQuery(jsq) + + def outputMode(self, outputMode: str) -> "DataStreamWriter": + """Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Options include: + + * `append`: Only the new rows in the streaming DataFrame/Dataset will be written to + the sink + * `complete`: All the rows in the streaming DataFrame/Dataset will be written to the sink + every time these are some updates + * `update`: only the rows that were updated in the streaming DataFrame/Dataset will be + written to the sink every time there are some updates. If the query doesn't contain + aggregations, it will be equivalent to `append` mode. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.writeStream.outputMode('append') + <...streaming.readwriter.DataStreamWriter object ...> + + The example below uses Complete mode that the entire aggregated counts are printed out. + + >>> import time + >>> df = spark.readStream.format("rate").option("rowsPerSecond", 10).load() + >>> df = df.groupby().count() + >>> q = df.writeStream.outputMode("complete").format("console").start() + >>> time.sleep(3) + >>> q.stop() + """ + if not outputMode or type(outputMode) != str or len(outputMode.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "outputMode", "arg_value": str(outputMode)}, + ) + self._jwrite = self._jwrite.outputMode(outputMode) + return self + + def format(self, source: str) -> "DataStreamWriter": + """Specifies the underlying output data source. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + source : str + string, name of the data source, which for now can be 'parquet'. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.writeStream.format("text") + <...streaming.readwriter.DataStreamWriter object ...> + + This API allows to configure the source to write. The example below writes a CSV + file from Rate source in a streaming manner. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d, tempfile.TemporaryDirectory() as cp: + ... df = spark.readStream.format("rate").load() + ... q = df.writeStream.format("csv").option("checkpointLocation", cp).start(d) + ... time.sleep(5) + ... q.stop() + ... spark.read.schema("timestamp TIMESTAMP, value STRING").csv(d).show() + +...---------+-----+ + |...timestamp|value| + +...---------+-----+ + ... + """ + self._jwrite = self._jwrite.format(source) + return self + + def option(self, key: str, value: "OptionalPrimitiveType") -> "DataStreamWriter": + """Adds an output option for the underlying data source. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.writeStream.option("x", 1) + <...streaming.readwriter.DataStreamWriter object ...> + + The example below specifies 'numRows' option to Console source in order to print + 3 rows for every batch. + + >>> import time + >>> q = spark.readStream.format( + ... "rate").option("rowsPerSecond", 10).load().writeStream.format( + ... "console").option("numRows", 3).start() + >>> time.sleep(3) + >>> q.stop() + """ + self._jwrite = self._jwrite.option(key, to_str(value)) + return self + + def options(self, **options: "OptionalPrimitiveType") -> "DataStreamWriter": + """Adds output options for the underlying data source. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.writeStream.option("x", 1) + <...streaming.readwriter.DataStreamWriter object ...> + + The example below specifies 'numRows' and 'truncate' options to Console source in order + to print 3 rows for every batch without truncating the results. + + >>> import time + >>> q = spark.readStream.format( + ... "rate").option("rowsPerSecond", 10).load().writeStream.format( + ... "console").options(numRows=3, truncate=False).start() + >>> time.sleep(3) + >>> q.stop() + """ + for k in options: + self._jwrite = self._jwrite.option(k, to_str(options[k])) + return self + + @overload + def partitionBy(self, *cols: str) -> "DataStreamWriter": + ... + + @overload + def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": + ... + + def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + """Partitions the output by the given columns on the file system. + + If specified, the output is laid out on the file system similar + to Hive's partitioning scheme. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + cols : str or list + name of columns + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + >>> df.writeStream.partitionBy("value") + <...streaming.readwriter.DataStreamWriter object ...> + + Partition-by timestamp column from Rate source. + + >>> import tempfile + >>> import time + >>> with tempfile.TemporaryDirectory() as d, tempfile.TemporaryDirectory() as cp: + ... df = spark.readStream.format("rate").option("rowsPerSecond", 10).load() + ... q = df.writeStream.partitionBy( + ... "timestamp").format("parquet").option("checkpointLocation", cp).start(d) + ... time.sleep(5) + ... q.stop() + ... spark.read.schema(df.schema).parquet(d).show() + +...---------+-----+ + |...timestamp|value| + +...---------+-----+ + ... + """ + if len(cols) == 1 and isinstance(cols[0], (list, tuple)): + cols = cols[0] + self._jwrite = self._jwrite.partitionBy(_to_seq(self._spark._sc, cols)) + return self + + def queryName(self, queryName: str) -> "DataStreamWriter": + """Specifies the name of the :class:`StreamingQuery` that can be started with + :func:`start`. This name must be unique among all the currently active queries + in the associated SparkSession. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + queryName : str + unique name for the query + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> import time + >>> df = spark.readStream.format("rate").load() + >>> q = df.writeStream.queryName("streaming_query").format("console").start() + >>> q.stop() + >>> q.name + 'streaming_query' + """ + if not queryName or type(queryName) != str or len(queryName.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "queryName", "arg_value": str(queryName)}, + ) + self._jwrite = self._jwrite.queryName(queryName) + return self + + @overload + def trigger(self, *, processingTime: str) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, once: bool) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, continuous: str) -> "DataStreamWriter": + ... + + @overload + def trigger(self, *, availableNow: bool) -> "DataStreamWriter": + ... + + def trigger( + self, + *, + processingTime: Optional[str] = None, + once: Optional[bool] = None, + continuous: Optional[str] = None, + availableNow: Optional[bool] = None, + ) -> "DataStreamWriter": + """Set the trigger for the stream query. If this is not set it will run the query as fast + as possible, which is equivalent to setting the trigger to ``processingTime='0 seconds'``. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + processingTime : str, optional + a processing time interval as a string, e.g. '5 seconds', '1 minute'. + Set a trigger that runs a microbatch query periodically based on the + processing time. Only one trigger can be set. + once : bool, optional + if set to True, set a trigger that processes only one batch of data in a + streaming query then terminates the query. Only one trigger can be set. + continuous : str, optional + a time interval as a string, e.g. '5 seconds', '1 minute'. + Set a trigger that runs a continuous query with a given checkpoint + interval. Only one trigger can be set. + availableNow : bool, optional + if set to True, set a trigger that processes all available data in multiple + batches then terminates the query. Only one trigger can be set. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + + Trigger the query for execution every 5 seconds + + >>> df.writeStream.trigger(processingTime='5 seconds') + <...streaming.readwriter.DataStreamWriter object ...> + + Trigger the query for execution every 5 seconds + + >>> df.writeStream.trigger(continuous='5 seconds') + <...streaming.readwriter.DataStreamWriter object ...> + + Trigger the query for reading all available data with multiple batches + + >>> df.writeStream.trigger(availableNow=True) + <...streaming.readwriter.DataStreamWriter object ...> + """ + params = [processingTime, once, continuous, availableNow] + + if params.count(None) == 4: + raise PySparkValueError( + error_class="ONLY_ALLOW_SINGLE_TRIGGER", + message_parameters={}, + ) + elif params.count(None) < 3: + raise PySparkValueError( + error_class="ONLY_ALLOW_SINGLE_TRIGGER", + message_parameters={}, + ) + + jTrigger = None + assert self._spark._sc._jvm is not None + if processingTime is not None: + if type(processingTime) != str or len(processingTime.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={ + "arg_name": "processingTime", + "arg_value": str(processingTime), + }, + ) + interval = processingTime.strip() + jTrigger = self._spark._sc._jvm.org.apache.spark.sql.streaming.Trigger.ProcessingTime( + interval + ) + + elif once is not None: + if once is not True: + raise PySparkValueError( + error_class="VALUE_NOT_TRUE", + message_parameters={"arg_name": "once", "arg_value": str(once)}, + ) + + jTrigger = self._spark._sc._jvm.org.apache.spark.sql.streaming.Trigger.Once() + + elif continuous is not None: + if type(continuous) != str or len(continuous.strip()) == 0: + raise PySparkValueError( + error_class="VALUE_NOT_NON_EMPTY_STR", + message_parameters={"arg_name": "continuous", "arg_value": str(continuous)}, + ) + interval = continuous.strip() + jTrigger = self._spark._sc._jvm.org.apache.spark.sql.streaming.Trigger.Continuous( + interval + ) + else: + if availableNow is not True: + raise PySparkValueError( + error_class="VALUE_NOT_TRUE", + message_parameters={"arg_name": "availableNow", "arg_value": str(availableNow)}, + ) + jTrigger = self._spark._sc._jvm.org.apache.spark.sql.streaming.Trigger.AvailableNow() + + self._jwrite = self._jwrite.trigger(jTrigger) + return self + + @staticmethod + def _construct_foreach_function( + f: Union[Callable[[Row], None], "SupportsProcess"] + ) -> Callable[[Any, Iterator], Iterator]: + from pyspark.taskcontext import TaskContext + + if callable(f): + # The provided object is a callable function that is supposed to be called on each row. + # Construct a function that takes an iterator and calls the provided function on each + # row. + def func_without_process(_: Any, iterator: Iterator) -> Iterator: + for x in iterator: + f(x) # type: ignore[operator] + return iter([]) + + return func_without_process + + else: + # The provided object is not a callable function. Then it is expected to have a + # 'process(row)' method, and optional 'open(partition_id, epoch_id)' and + # 'close(error)' methods. + + if not hasattr(f, "process"): + raise PySparkAttributeError( + error_class="ATTRIBUTE_NOT_CALLABLE", + message_parameters={"attr_name": "process", "obj_name": "f"}, + ) + + if not callable(getattr(f, "process")): + raise PySparkAttributeError( + error_class="ATTRIBUTE_NOT_CALLABLE", + message_parameters={"attr_name": "process", "obj_name": "f"}, + ) + + def doesMethodExist(method_name: str) -> bool: + exists = hasattr(f, method_name) + if exists and not callable(getattr(f, method_name)): + raise PySparkAttributeError( + error_class="ATTRIBUTE_NOT_CALLABLE", + message_parameters={"attr_name": method_name, "obj_name": "f"}, + ) + return exists + + open_exists = doesMethodExist("open") + close_exists = doesMethodExist("close") + + def func_with_open_process_close(partition_id: Any, iterator: Iterator) -> Iterator: + epoch_id = cast(TaskContext, TaskContext.get()).getLocalProperty( + "streaming.sql.batchId" + ) + if epoch_id: + int_epoch_id = int(epoch_id) + else: + raise PySparkRuntimeError( + error_class="CANNOT_GET_BATCH_ID", + message_parameters={"obj_name": "TaskContext"}, + ) + + # Check if the data should be processed + should_process = True + if open_exists: + should_process = f.open(partition_id, int_epoch_id) # type: ignore[union-attr] + + error = None + + try: + if should_process: + for x in iterator: + cast("SupportsProcess", f).process(x) + except Exception as ex: + error = ex + finally: + if close_exists: + f.close(error) # type: ignore[union-attr] + if error: + raise error + + return iter([]) + + return func_with_open_process_close + + @overload + def foreach(self, f: Callable[[Row], None]) -> "DataStreamWriter": + ... + + @overload + def foreach(self, f: "SupportsProcess") -> "DataStreamWriter": + ... + + def foreach(self, f: Union[Callable[[Row], None], "SupportsProcess"]) -> "DataStreamWriter": + """ + Sets the output of the streaming query to be processed using the provided writer ``f``. + This is often used to write the output of a streaming query to arbitrary storage systems. + The processing logic can be specified in two ways. + + #. A **function** that takes a row as input. + This is a simple way to express your processing logic. Note that this does + not allow you to deduplicate generated data when failures cause reprocessing of + some input data. That would require you to specify the processing logic in the next + way. + + #. An **object** with a ``process`` method and optional ``open`` and ``close`` methods. + The object can have the following methods. + + * ``open(partition_id, epoch_id)``: *Optional* method that initializes the processing + (for example, open a connection, start a transaction, etc). Additionally, you can + use the `partition_id` and `epoch_id` to deduplicate regenerated data + (discussed later). + + * ``process(row)``: *Non-optional* method that processes each :class:`Row`. + + * ``close(error)``: *Optional* method that finalizes and cleans up (for example, + close connection, commit transaction, etc.) after all rows have been processed. + + The object will be used by Spark in the following way. + + * A single copy of this object is responsible of all the data generated by a + single task in a query. In other words, one instance is responsible for + processing one partition of the data generated in a distributed manner. + + * This object must be serializable because each task will get a fresh + serialized-deserialized copy of the provided object. Hence, it is strongly + recommended that any initialization for writing data (e.g. opening a + connection or starting a transaction) is done after the `open(...)` + method has been called, which signifies that the task is ready to generate data. + + * The lifecycle of the methods are as follows. + + For each partition with ``partition_id``: + + ... For each batch/epoch of streaming data with ``epoch_id``: + + ....... Method ``open(partitionId, epochId)`` is called. + + ....... If ``open(...)`` returns true, for each row in the partition and + batch/epoch, method ``process(row)`` is called. + + ....... Method ``close(errorOrNull)`` is called with error (if any) seen while + processing rows. + + Important points to note: + + * The `partitionId` and `epochId` can be used to deduplicate generated data when + failures cause reprocessing of some input data. This depends on the execution + mode of the query. If the streaming query is being executed in the micro-batch + mode, then every partition represented by a unique tuple (partition_id, epoch_id) + is guaranteed to have the same data. Hence, (partition_id, epoch_id) can be used + to deduplicate and/or transactionally commit data and achieve exactly-once + guarantees. However, if the streaming query is being executed in the continuous + mode, then this guarantee does not hold and therefore should not be used for + deduplication. + + * The ``close()`` method (if exists) will be called if `open()` method exists and + returns successfully (irrespective of the return value), except if the Python + crashes in the middle. + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> import time + >>> df = spark.readStream.format("rate").load() + + Print every row using a function + + >>> def print_row(row): + ... print(row) + ... + >>> q = df.writeStream.foreach(print_row).start() + >>> time.sleep(3) + >>> q.stop() + + Print every row using a object with process() method + + >>> class RowPrinter: + ... def open(self, partition_id, epoch_id): + ... print("Opened %d, %d" % (partition_id, epoch_id)) + ... return True + ... + ... def process(self, row): + ... print(row) + ... + ... def close(self, error): + ... print("Closed with error: %s" % str(error)) + ... + >>> q = df.writeStream.foreach(print_row).start() + >>> time.sleep(3) + >>> q.stop() + """ + + from pyspark.rdd import _wrap_function + from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer + + func = self._construct_foreach_function(f) + serializer = AutoBatchedSerializer(CPickleSerializer()) + wrapped_func = _wrap_function(self._spark._sc, func, serializer, serializer) + assert self._spark._sc._jvm is not None + jForeachWriter = ( + self._spark._sc._jvm.org.apache.spark.sql.execution.python.PythonForeachWriter( + wrapped_func, self._df._jdf.schema() + ) + ) + self._jwrite.foreach(jForeachWriter) + return self + + def foreachBatch(self, func: Callable[["DataFrame", int], None]) -> "DataStreamWriter": + """ + Sets the output of the streaming query to be processed using the provided + function. This is supported only the in the micro-batch execution modes (that is, when the + trigger is not continuous). In every micro-batch, the provided function will be called in + every micro-batch with (i) the output rows as a DataFrame and (ii) the batch identifier. + The batchId can be used deduplicate and transactionally write the output + (that is, the provided Dataset) to external systems. The output DataFrame is guaranteed + to exactly same for the same batchId (assuming all operations are deterministic in the + query). + + .. versionadded:: 2.4.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Notes + ----- + This API is evolving. + This function behaves differently in Spark Connect mode. See examples. + In Connect, the provided function doesn't have access to variables defined outside of it. + + Examples + -------- + >>> import time + >>> df = spark.readStream.format("rate").load() + >>> my_value = -1 + >>> def func(batch_df, batch_id): + ... global my_value + ... my_value = 100 + ... batch_df.collect() + ... + >>> q = df.writeStream.foreachBatch(func).start() + >>> time.sleep(3) + >>> q.stop() + >>> # if in Spark Connect, my_value = -1, else my_value = 100 + """ + + from pyspark.java_gateway import ensure_callback_server_started + + gw = self._spark._sc._gateway + assert gw is not None + java_import(gw.jvm, "org.apache.spark.sql.execution.streaming.sources.*") + + wrapped_func = ForeachBatchFunction(self._spark, func) + gw.jvm.PythonForeachBatchHelper.callForeachBatch(self._jwrite, wrapped_func) + ensure_callback_server_started(gw) + return self + + def start( + self, + path: Optional[str] = None, + format: Optional[str] = None, + outputMode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + queryName: Optional[str] = None, + **options: "OptionalPrimitiveType", + ) -> StreamingQuery: + """Streams the contents of the :class:`DataFrame` to a data source. + + The data source is specified by the ``format`` and a set of ``options``. + If ``format`` is not specified, the default data source configured by + ``spark.sql.sources.default`` will be used. + + .. versionadded:: 2.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + path : str, optional + the path in a Hadoop supported file system + format : str, optional + the format used to save + outputMode : str, optional + specifies how data of a streaming DataFrame/Dataset is written to a + streaming sink. + + * `append`: Only the new rows in the streaming DataFrame/Dataset will be written to the + sink + * `complete`: All the rows in the streaming DataFrame/Dataset will be written to the + sink every time these are some updates + * `update`: only the rows that were updated in the streaming DataFrame/Dataset will be + written to the sink every time there are some updates. If the query doesn't contain + aggregations, it will be equivalent to `append` mode. + partitionBy : str or list, optional + names of partitioning columns + queryName : str, optional + unique name for the query + **options : dict + All other string options. You may want to provide a `checkpointLocation` + for most streams, however it is not required for a `memory` stream. + + Notes + ----- + This API is evolving. + + Examples + -------- + >>> df = spark.readStream.format("rate").load() + + Basic example. + + >>> q = df.writeStream.format('memory').queryName('this_query').start() + >>> q.isActive + True + >>> q.name + 'this_query' + >>> q.stop() + >>> q.isActive + False + + Example with using other parameters with a trigger. + + >>> q = df.writeStream.trigger(processingTime='5 seconds').start( + ... queryName='that_query', outputMode="append", format='memory') + >>> q.name + 'that_query' + >>> q.isActive + True + >>> q.stop() + """ + self.options(**options) + if outputMode is not None: + self.outputMode(outputMode) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + if queryName is not None: + self.queryName(queryName) + if path is None: + return self._sq(self._jwrite.start()) + else: + return self._sq(self._jwrite.start(path)) + + def toTable( + self, + tableName: str, + format: Optional[str] = None, + outputMode: Optional[str] = None, + partitionBy: Optional[Union[str, List[str]]] = None, + queryName: Optional[str] = None, + **options: "OptionalPrimitiveType", + ) -> StreamingQuery: + """ + Starts the execution of the streaming query, which will continually output results to the + given table as new data arrives. + + The returned :class:`StreamingQuery` object can be used to interact with the stream. + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + tableName : str + string, for the name of the table. + format : str, optional + the format used to save. + outputMode : str, optional + specifies how data of a streaming DataFrame/Dataset is written to a + streaming sink. + + * `append`: Only the new rows in the streaming DataFrame/Dataset will be written to the + sink + * `complete`: All the rows in the streaming DataFrame/Dataset will be written to the + sink every time these are some updates + * `update`: only the rows that were updated in the streaming DataFrame/Dataset will be + written to the sink every time there are some updates. If the query doesn't contain + aggregations, it will be equivalent to `append` mode. + partitionBy : str or list, optional + names of partitioning columns + queryName : str, optional + unique name for the query + **options : dict + All other string options. You may want to provide a `checkpointLocation`. + + Notes + ----- + This API is evolving. + + For v1 table, partitioning columns provided by `partitionBy` will be respected no matter + the table exists or not. A new table will be created if the table not exists. + + For v2 table, `partitionBy` will be ignored if the table already exists. `partitionBy` will + be respected only if the v2 table does not exist. Besides, the v2 table created by this API + lacks some functionalities (e.g., customized properties, options, and serde info). If you + need them, please create the v2 table manually before the execution to avoid creating a + table with incomplete information. + + Examples + -------- + Save a data stream to a table. + + >>> import tempfile + >>> import time + >>> _ = spark.sql("DROP TABLE IF EXISTS my_table2") + >>> with tempfile.TemporaryDirectory() as d: + ... # Create a table with Rate source. + ... q = spark.readStream.format("rate").option( + ... "rowsPerSecond", 10).load().writeStream.toTable( + ... "my_table2", + ... queryName='that_query', + ... outputMode="append", + ... format='parquet', + ... checkpointLocation=d) + ... time.sleep(3) + ... q.stop() + ... spark.read.table("my_table2").show() + ... _ = spark.sql("DROP TABLE my_table2") + +...---------+-----+ + |...timestamp|value| + +...---------+-----+ + ... + """ + self.options(**options) + if outputMode is not None: + self.outputMode(outputMode) + if partitionBy is not None: + self.partitionBy(partitionBy) + if format is not None: + self.format(format) + if queryName is not None: + self.queryName(queryName) + return self._sq(self._jwrite.toTable(tableName)) + + +def _test() -> None: + import doctest + import os + from pyspark.sql import SparkSession + import pyspark.sql.streaming.readwriter + + os.chdir(os.environ["SPARK_HOME"]) + + globs = pyspark.sql.streaming.readwriter.__dict__.copy() + globs["spark"] = ( + SparkSession.builder.master("local[4]") + .appName("sql.streaming.readwriter tests") + .getOrCreate() + ) + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.streaming.readwriter, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, + ) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/state.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/state.py new file mode 100644 index 0000000..8bf01b3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/streaming/state.py @@ -0,0 +1,288 @@ +# +# 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 datetime +import json +from typing import Tuple, Optional + +from pyspark.sql.types import DateType, Row, StructType +from pyspark.sql.utils import has_numpy +from pyspark.errors import PySparkTypeError, PySparkValueError + +__all__ = ["GroupState", "GroupStateTimeout"] + + +class GroupStateTimeout: + """ + Represents the type of timeouts possible for the Dataset operations applyInPandasWithState. + """ + + NoTimeout: str = "NoTimeout" + ProcessingTimeTimeout: str = "ProcessingTimeTimeout" + EventTimeTimeout: str = "EventTimeTimeout" + + +class GroupState: + """ + Wrapper class for interacting with per-group state data in `applyInPandasWithState`. + """ + + NO_TIMESTAMP: int = -1 + + def __init__( + self, + # JVM Constructor + optionalValue: Row, + batchProcessingTimeMs: int, + eventTimeWatermarkMs: int, + timeoutConf: str, + hasTimedOut: bool, + watermarkPresent: bool, + # JVM internal state. + defined: bool, + updated: bool, + removed: bool, + timeoutTimestamp: int, + # Python internal state. + keyAsUnsafe: bytes, + valueSchema: StructType, + ) -> None: + self._keyAsUnsafe = keyAsUnsafe + self._value = optionalValue + self._batch_processing_time_ms = batchProcessingTimeMs + self._event_time_watermark_ms = eventTimeWatermarkMs + + assert timeoutConf in [ + GroupStateTimeout.NoTimeout, + GroupStateTimeout.ProcessingTimeTimeout, + GroupStateTimeout.EventTimeTimeout, + ] + self._timeout_conf = timeoutConf + + self._has_timed_out = hasTimedOut + self._watermark_present = watermarkPresent + + self._defined = defined + self._updated = updated + self._removed = removed + self._timeout_timestamp = timeoutTimestamp + # Python internal state. + self._old_timeout_timestamp = timeoutTimestamp + + self._value_schema = valueSchema + + @property + def exists(self) -> bool: + """ + Whether state exists or not. + """ + return self._defined + + @property + def get(self) -> Tuple: + """ + Get the state value if it exists, or throw ValueError. + """ + if self.exists: + return tuple(self._value) + else: + raise PySparkValueError( + error_class="STATE_NOT_EXISTS", + message_parameters={}, + ) + + @property + def getOption(self) -> Optional[Tuple]: + """ + Get the state value if it exists, or return None. + """ + if self.exists: + return tuple(self._value) + else: + return None + + @property + def hasTimedOut(self) -> bool: + """ + Whether the function has been called because the key has timed out. + This can return true only when timeouts are enabled. + """ + return self._has_timed_out + + # NOTE: this function is only available to PySpark implementation due to underlying + # implementation, do not port to Scala implementation! + @property + def oldTimeoutTimestamp(self) -> int: + return self._old_timeout_timestamp + + def update(self, newValue: Tuple) -> None: + """ + Update the value of the state. The value of the state cannot be null. + """ + if newValue is None: + raise PySparkTypeError( + error_class="CANNOT_BE_NONE", + message_parameters={"arg_name": "newValue"}, + ) + + converted = [] + if has_numpy: + import numpy as np + + # In order to convert NumPy types to Python primitive types. + for v in newValue: + if isinstance(v, np.generic): + converted.append(v.tolist()) + # Address a couple of pandas dtypes too. + elif hasattr(v, "to_pytimedelta"): + converted.append(v.to_pytimedelta()) + elif hasattr(v, "to_pydatetime"): + converted.append(v.to_pydatetime()) + else: + converted.append(v) + else: + converted = list(newValue) + + self._value = Row(*converted) + self._defined = True + self._updated = True + self._removed = False + + def remove(self) -> None: + """ + Remove this state. + """ + self._defined = False + self._updated = False + self._removed = True + + def setTimeoutDuration(self, durationMs: int) -> None: + """ + Set the timeout duration in ms for this key. + Processing time timeout must be enabled. + """ + if isinstance(durationMs, str): + # TODO(SPARK-40437): Support string representation of durationMs. + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={ + "arg_name": "durationMs", + "arg_type": type(durationMs).__name__, + }, + ) + + if self._timeout_conf != GroupStateTimeout.ProcessingTimeTimeout: + raise RuntimeError( + "Cannot set timeout duration without enabling processing time timeout in " + "applyInPandasWithState" + ) + + if durationMs <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={ + "arg_name": "durationMs", + "arg_type": type(durationMs).__name__, + }, + ) + self._timeout_timestamp = durationMs + self._batch_processing_time_ms + + # TODO(SPARK-40438): Implement additionalDuration parameter. + def setTimeoutTimestamp(self, timestampMs: int) -> None: + """ + Set the timeout timestamp for this key as milliseconds in epoch time. + This timestamp cannot be older than the current watermark. + Event time timeout must be enabled. + """ + if self._timeout_conf != GroupStateTimeout.EventTimeTimeout: + raise RuntimeError( + "Cannot set timeout duration without enabling processing time timeout in " + "applyInPandasWithState" + ) + + if isinstance(timestampMs, datetime.datetime): + timestampMs = DateType().toInternal(timestampMs) + + if timestampMs <= 0: + raise PySparkValueError( + error_class="VALUE_NOT_POSITIVE", + message_parameters={ + "arg_name": "timestampMs", + "arg_type": type(timestampMs).__name__, + }, + ) + + if ( + self._event_time_watermark_ms != GroupState.NO_TIMESTAMP + and timestampMs < self._event_time_watermark_ms + ): + raise PySparkValueError( + error_class="INVALID_TIMEOUT_TIMESTAMP", + message_parameters={ + "timestamp": str(timestampMs), + "watermark": str(self._event_time_watermark_ms), + }, + ) + + self._timeout_timestamp = timestampMs + + def getCurrentWatermarkMs(self) -> int: + """ + Get the current event time watermark as milliseconds in epoch time. + In a streaming query, this can be called only when watermark is set. + """ + if not self._watermark_present: + raise RuntimeError( + "Cannot get event time watermark timestamp without setting watermark before " + "applyInPandasWithState" + ) + return self._event_time_watermark_ms + + def getCurrentProcessingTimeMs(self) -> int: + """ + Get the current processing time as milliseconds in epoch time. + In a streaming query, this will return a constant value throughout the duration of a + trigger, even if the trigger is re-executed. + """ + return self._batch_processing_time_ms + + def __str__(self) -> str: + if self.exists: + return "GroupState(%s)" % (self.get,) + else: + return "GroupState()" + + def json(self) -> str: + """ + Convert the internal values of instance into JSON. This is used to send out the update + from Python worker to executor. + """ + return json.dumps( + { + # Constructor + "optionalValue": None, # Note that optionalValue will be manually serialized. + "batchProcessingTimeMs": self._batch_processing_time_ms, + "eventTimeWatermarkMs": self._event_time_watermark_ms, + "timeoutConf": self._timeout_conf, + "hasTimedOut": self._has_timed_out, + "watermarkPresent": self._watermark_present, + # JVM internal state. + "defined": self._defined, + "updated": self._updated, + "removed": self._removed, + "timeoutTimestamp": self._timeout_timestamp, + } + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_artifact.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_artifact.py new file mode 100644 index 0000000..d45230e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_artifact.py @@ -0,0 +1,418 @@ +# +# 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 hashlib +import shutil +import tempfile +import unittest +import os + +from pyspark.sql import SparkSession +from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect +from pyspark.testing.utils import SPARK_HOME +from pyspark import SparkFiles +from pyspark.sql.functions import udf + +if should_test_connect: + from pyspark.sql.connect.client.artifact import ArtifactManager + from pyspark.sql.connect.client import ChannelBuilder + + +class ArtifactTestsMixin: + def check_add_pyfile(self, spark_session): + with tempfile.TemporaryDirectory() as d: + pyfile_path = os.path.join(d, "my_pyfile.py") + with open(pyfile_path, "w") as f: + f.write("my_func = lambda: 10") + + @udf("int") + def func(x): + import my_pyfile + + return my_pyfile.my_func() + + spark_session.addArtifacts(pyfile_path, pyfile=True) + self.assertEqual(spark_session.range(1).select(func("id")).first()[0], 10) + + def test_add_pyfile(self): + self.check_add_pyfile(self.spark) + + # Test multi sessions. Should be able to add the same + # file from different session. + self.check_add_pyfile( + SparkSession.builder.remote(f"sc://localhost:{ChannelBuilder.default_port()}").create() + ) + + def check_add_zipped_package(self, spark_session): + with tempfile.TemporaryDirectory() as d: + package_path = os.path.join(d, "my_zipfile") + os.mkdir(package_path) + pyfile_path = os.path.join(package_path, "__init__.py") + with open(pyfile_path, "w") as f: + _ = f.write("my_func = lambda: 5") + shutil.make_archive(package_path, "zip", d, "my_zipfile") + + @udf("long") + def func(x): + import my_zipfile + + return my_zipfile.my_func() + + spark_session.addArtifacts(f"{package_path}.zip", pyfile=True) + self.assertEqual(spark_session.range(1).select(func("id")).first()[0], 5) + + def test_add_zipped_package(self): + self.check_add_zipped_package(self.spark) + + # Test multi sessions. Should be able to add the same + # file from different session. + self.check_add_zipped_package( + SparkSession.builder.remote(f"sc://localhost:{ChannelBuilder.default_port()}").create() + ) + + def check_add_archive(self, spark_session): + with tempfile.TemporaryDirectory() as d: + archive_path = os.path.join(d, "my_archive") + os.mkdir(archive_path) + pyfile_path = os.path.join(archive_path, "my_file.txt") + with open(pyfile_path, "w") as f: + _ = f.write("hello world!") + shutil.make_archive(archive_path, "zip", d, "my_archive") + + # Should addArtifact first to make sure state is set, + # and 'root' can be found properly. + spark_session.addArtifacts(f"{archive_path}.zip#my_files", archive=True) + + root = self.root() + + @udf("string") + def func(x): + with open( + os.path.join(root, "my_files", "my_archive", "my_file.txt"), + "r", + ) as my_file: + return my_file.read().strip() + + self.assertEqual(spark_session.range(1).select(func("id")).first()[0], "hello world!") + + def test_add_archive(self): + self.check_add_archive(self.spark) + + # Test multi sessions. Should be able to add the same + # file from different session. + self.check_add_archive( + SparkSession.builder.remote(f"sc://localhost:{ChannelBuilder.default_port()}").create() + ) + + def check_add_file(self, spark_session): + with tempfile.TemporaryDirectory() as d: + file_path = os.path.join(d, "my_file.txt") + with open(file_path, "w") as f: + f.write("Hello world!!") + + # Should addArtifact first to make sure state is set, + # and 'root' can be found properly. + spark_session.addArtifacts(file_path, file=True) + + root = self.root() + + @udf("string") + def func(x): + with open(os.path.join(root, "my_file.txt"), "r") as my_file: + return my_file.read().strip() + + self.assertEqual(spark_session.range(1).select(func("id")).first()[0], "Hello world!!") + + def test_add_file(self): + self.check_add_file(self.spark) + + # Test multi sessions. Should be able to add the same + # file from different session. + self.check_add_file( + SparkSession.builder.remote(f"sc://localhost:{ChannelBuilder.default_port()}").create() + ) + + +class ArtifactTests(ReusedConnectTestCase, ArtifactTestsMixin): + @classmethod + def root(cls): + # In local mode, the file location is the same as Driver + # The executors are running in a thread. + jvm = SparkSession._instantiatedSession._jvm + current_uuid = ( + getattr( + getattr( + jvm.org.apache.spark, # type: ignore[union-attr] + "JobArtifactSet$", + ), + "MODULE$", + ) + .lastSeenState() + .get() + .uuid() + ) + return os.path.join(SparkFiles.getRootDirectory(), current_uuid) + + @classmethod + def setUpClass(cls): + super(ArtifactTests, cls).setUpClass() + cls.artifact_manager: ArtifactManager = cls.spark._client._artifact_manager + cls.base_resource_dir = os.path.join(SPARK_HOME, "data") + cls.artifact_file_path = os.path.join( + cls.base_resource_dir, + "artifact-tests", + ) + cls.artifact_crc_path = os.path.join( + cls.artifact_file_path, + "crc", + ) + + @classmethod + def conf(cls): + conf = super().conf() + conf.set("spark.connect.copyFromLocalToFs.allowDestLocal", "true") + return conf + + def test_basic_requests(self): + file_name = "smallJar" + small_jar_path = os.path.join(self.artifact_file_path, f"{file_name}.jar") + response = self.artifact_manager._retrieve_responses( + self.artifact_manager._create_requests( + small_jar_path, pyfile=False, archive=False, file=False + ) + ) + self.assertTrue(response.artifacts[0].name.endswith(f"{file_name}.jar")) + + def test_single_chunk_artifact(self): + file_name = "smallJar" + small_jar_path = os.path.join(self.artifact_file_path, f"{file_name}.jar") + small_jar_crc_path = os.path.join(self.artifact_crc_path, f"{file_name}.txt") + + requests = list( + self.artifact_manager._create_requests( + small_jar_path, pyfile=False, archive=False, file=False + ) + ) + self.assertEqual(len(requests), 1) + + request = requests[0] + self.assertIsNotNone(request.batch) + + batch = request.batch + self.assertEqual(len(batch.artifacts), 1) + + single_artifact = batch.artifacts[0] + self.assertTrue(single_artifact.name.endswith(".jar")) + + self.assertEqual(os.path.join("jars", f"{file_name}.jar"), single_artifact.name) + with open(small_jar_crc_path) as f1, open(small_jar_path, "rb") as f2: + self.assertEqual(single_artifact.data.crc, int(f1.readline())) + self.assertEqual(single_artifact.data.data, f2.read()) + + def test_chunked_artifacts(self): + file_name = "junitLargeJar" + large_jar_path = os.path.join(self.artifact_file_path, f"{file_name}.jar") + large_jar_crc_path = os.path.join(self.artifact_crc_path, f"{file_name}.txt") + + requests = list( + self.artifact_manager._create_requests( + large_jar_path, pyfile=False, archive=False, file=False + ) + ) + # Expected chunks = roundUp( file_size / chunk_size) = 12 + # File size of `junitLargeJar.jar` is 384581 bytes. + large_jar_size = os.path.getsize(large_jar_path) + expected_chunks = int( + (large_jar_size + (ArtifactManager.CHUNK_SIZE - 1)) / ArtifactManager.CHUNK_SIZE + ) + self.assertEqual(len(requests), expected_chunks) + request = requests[0] + self.assertIsNotNone(request.begin_chunk) + begin_chunk = request.begin_chunk + self.assertEqual(begin_chunk.name, os.path.join("jars", f"{file_name}.jar")) + self.assertEqual(begin_chunk.total_bytes, large_jar_size) + self.assertEqual(begin_chunk.num_chunks, expected_chunks) + other_requests = requests[1:] + data_chunks = [begin_chunk.initial_chunk] + [req.chunk for req in other_requests] + + with open(large_jar_crc_path) as f1, open(large_jar_path, "rb") as f2: + cscs = [chunk.crc for chunk in data_chunks] + expected_cscs = [int(line.rstrip()) for line in f1] + self.assertEqual(cscs, expected_cscs) + + binaries = [chunk.data for chunk in data_chunks] + expected_binaries = list(iter(lambda: f2.read(ArtifactManager.CHUNK_SIZE), b"")) + self.assertEqual(binaries, expected_binaries) + + def test_batched_artifacts(self): + file_name = "smallJar" + small_jar_path = os.path.join(self.artifact_file_path, f"{file_name}.jar") + small_jar_crc_path = os.path.join(self.artifact_crc_path, f"{file_name}.txt") + + requests = list( + self.artifact_manager._create_requests( + small_jar_path, small_jar_path, pyfile=False, archive=False, file=False + ) + ) + # Single request containing 2 artifacts. + self.assertEqual(len(requests), 1) + + request = requests[0] + self.assertIsNotNone(request.batch) + + batch = request.batch + self.assertEqual(len(batch.artifacts), 2) + + artifact1 = batch.artifacts[0] + self.assertTrue(artifact1.name.endswith(".jar")) + artifact2 = batch.artifacts[1] + self.assertTrue(artifact2.name.endswith(".jar")) + + self.assertEqual(os.path.join("jars", f"{file_name}.jar"), artifact1.name) + with open(small_jar_crc_path) as f1, open(small_jar_path, "rb") as f2: + crc = int(f1.readline()) + data = f2.read() + self.assertEqual(artifact1.data.crc, crc) + self.assertEqual(artifact1.data.data, data) + self.assertEqual(artifact2.data.crc, crc) + self.assertEqual(artifact2.data.data, data) + + def test_single_chunked_and_chunked_artifact(self): + file_name1 = "smallJar" + file_name2 = "junitLargeJar" + small_jar_path = os.path.join(self.artifact_file_path, f"{file_name1}.jar") + small_jar_crc_path = os.path.join(self.artifact_crc_path, f"{file_name1}.txt") + large_jar_path = os.path.join(self.artifact_file_path, f"{file_name2}.jar") + large_jar_crc_path = os.path.join(self.artifact_crc_path, f"{file_name2}.txt") + large_jar_size = os.path.getsize(large_jar_path) + + requests = list( + self.artifact_manager._create_requests( + small_jar_path, + large_jar_path, + small_jar_path, + small_jar_path, + pyfile=False, + archive=False, + file=False, + ) + ) + # There are a total of 14 requests. + # The 1st request contains a single artifact - smallJar.jar (There are no + # other artifacts batched with it since the next one is large multi-chunk artifact) + # Requests 2-13 (1-indexed) belong to the transfer of junitLargeJar.jar. This includes + # the first "beginning chunk" and the subsequent data chunks. + # The last request (14) contains two smallJar.jar batched + # together. + self.assertEqual(len(requests), 1 + 12 + 1) + + first_req_batch = requests[0].batch.artifacts + self.assertEqual(len(first_req_batch), 1) + self.assertEqual(first_req_batch[0].name, os.path.join("jars", f"{file_name1}.jar")) + with open(small_jar_crc_path) as f1, open(small_jar_path, "rb") as f2: + self.assertEqual(first_req_batch[0].data.crc, int(f1.readline())) + self.assertEqual(first_req_batch[0].data.data, f2.read()) + + second_req_batch = requests[1] + self.assertIsNotNone(second_req_batch.begin_chunk) + begin_chunk = second_req_batch.begin_chunk + self.assertEqual(begin_chunk.name, os.path.join("jars", f"{file_name2}.jar")) + self.assertEqual(begin_chunk.total_bytes, large_jar_size) + self.assertEqual(begin_chunk.num_chunks, 12) + other_requests = requests[2:-1] + data_chunks = [begin_chunk.initial_chunk] + [req.chunk for req in other_requests] + + with open(large_jar_crc_path) as f1, open(large_jar_path, "rb") as f2: + cscs = [chunk.crc for chunk in data_chunks] + expected_cscs = [int(line.rstrip()) for line in f1] + self.assertEqual(cscs, expected_cscs) + + binaries = [chunk.data for chunk in data_chunks] + expected_binaries = list(iter(lambda: f2.read(ArtifactManager.CHUNK_SIZE), b"")) + self.assertEqual(binaries, expected_binaries) + + last_request = requests[-1] + self.assertIsNotNone(last_request.batch) + + batch = last_request.batch + self.assertEqual(len(batch.artifacts), 2) + + artifact1 = batch.artifacts[0] + self.assertTrue(artifact1.name.endswith(".jar")) + artifact2 = batch.artifacts[1] + self.assertTrue(artifact2.name.endswith(".jar")) + + self.assertEqual(os.path.join("jars", f"{file_name1}.jar"), artifact1.name) + with open(small_jar_crc_path) as f1, open(small_jar_path, "rb") as f2: + crc = int(f1.readline()) + data = f2.read() + self.assertEqual(artifact1.data.crc, crc) + self.assertEqual(artifact1.data.data, data) + self.assertEqual(artifact2.data.crc, crc) + self.assertEqual(artifact2.data.data, data) + + def test_copy_from_local_to_fs(self): + with tempfile.TemporaryDirectory() as d: + with tempfile.TemporaryDirectory() as d2: + file_path = os.path.join(d, "file1") + dest_path = os.path.join(d2, "file1_dest") + file_content = "test_copy_from_local_to_FS" + + with open(file_path, "w") as f: + f.write(file_content) + + self.spark.copyFromLocalToFs(file_path, dest_path) + + with open(dest_path, "r") as f: + self.assertEqual(f.read(), file_content) + + def test_cache_artifact(self): + s = "Hello, World!" + blob = bytearray(s, "utf-8") + expected_hash = hashlib.sha256(blob).hexdigest() + self.assertEqual(self.artifact_manager.is_cached_artifact(expected_hash), False) + actualHash = self.artifact_manager.cache_artifact(blob) + self.assertEqual(actualHash, expected_hash) + self.assertEqual(self.artifact_manager.is_cached_artifact(expected_hash), True) + + +class LocalClusterArtifactTests(ReusedConnectTestCase, ArtifactTestsMixin): + @classmethod + def conf(cls): + return ( + super().conf().set("spark.driver.memory", "512M").set("spark.executor.memory", "512M") + ) + + @classmethod + def root(cls): + # In local cluster, we can mimic the production usage. + return "." + + @classmethod + def master(cls): + return "local-cluster[2,2,512]" + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.client.test_artifact import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_client.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_client.py new file mode 100644 index 0000000..98f6876 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/client/test_client.py @@ -0,0 +1,167 @@ +# +# 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 unittest +import uuid +from typing import Optional + +from pyspark.sql.connect.client import SparkConnectClient, ChannelBuilder +import pyspark.sql.connect.proto as proto +from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + +from pyspark.sql.connect.client.core import Retrying +from pyspark.sql.connect.client.reattach import RetryException + +if should_test_connect: + import pandas as pd + import pyarrow as pa + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class SparkConnectClientTestCase(unittest.TestCase): + def test_user_agent_passthrough(self): + client = SparkConnectClient("sc://foo/;user_agent=bar", use_reattachable_execute=False) + mock = MockService(client._session_id) + client._stub = mock + + command = proto.Command() + client.execute_command(command) + + self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected") + self.assertRegex(mock.req.client_type, r"^bar spark/[^ ]+ os/[^ ]+ python/[^ ]+$") + + def test_user_agent_default(self): + client = SparkConnectClient("sc://foo/", use_reattachable_execute=False) + mock = MockService(client._session_id) + client._stub = mock + + command = proto.Command() + client.execute_command(command) + + self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected") + self.assertRegex( + mock.req.client_type, r"^_SPARK_CONNECT_PYTHON spark/[^ ]+ os/[^ ]+ python/[^ ]+$" + ) + + def test_properties(self): + client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False) + self.assertEqual(client.token, "bar") + self.assertEqual(client.host, "foo") + + client = SparkConnectClient("sc://foo/", use_reattachable_execute=False) + self.assertIsNone(client.token) + + def test_channel_builder(self): + class CustomChannelBuilder(ChannelBuilder): + @property + def userId(self) -> Optional[str]: + return "abc" + + client = SparkConnectClient( + CustomChannelBuilder("sc://foo/"), use_reattachable_execute=False + ) + + self.assertEqual(client._user_id, "abc") + + def test_interrupt_all(self): + client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False) + mock = MockService(client._session_id) + client._stub = mock + + client.interrupt_all() + self.assertIsNotNone(mock.req, "Interrupt API was not called when expected") + + def test_is_closed(self): + client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False) + + self.assertFalse(client.is_closed) + client.close() + self.assertTrue(client.is_closed) + + def test_retry(self): + client = SparkConnectClient("sc://foo/;token=bar") + + total_sleep = 0 + + def sleep(t): + nonlocal total_sleep + total_sleep += t + + try: + for attempt in Retrying( + can_retry=SparkConnectClient.retry_exception, sleep=sleep, **client._retry_policy + ): + with attempt: + raise RetryException() + except RetryException: + pass + + # tolerated at least 10 mins of fails + self.assertGreaterEqual(total_sleep, 600) + + def test_channel_builder_with_session(self): + dummy = str(uuid.uuid4()) + chan = ChannelBuilder(f"sc://foo/;session_id={dummy}") + client = SparkConnectClient(chan) + self.assertEqual(client._session_id, chan.session_id) + + +class MockService: + # Simplest mock of the SparkConnectService. + # If this needs more complex logic, it needs to be replaced with Python mocking. + + req: Optional[proto.ExecutePlanRequest] + + def __init__(self, session_id: str): + self._session_id = session_id + self.req = None + + def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata): + self.req = req + resp = proto.ExecutePlanResponse() + resp.session_id = self._session_id + + pdf = pd.DataFrame(data={"col1": [1, 2]}) + schema = pa.Schema.from_pandas(pdf) + table = pa.Table.from_pandas(pdf) + sink = pa.BufferOutputStream() + + writer = pa.ipc.new_stream(sink, schema=schema) + writer.write(table) + writer.close() + + buf = sink.getvalue() + resp.arrow_batch.data = buf.to_pybytes() + return [resp] + + def Interrupt(self, req: proto.InterruptRequest, metadata): + self.req = req + resp = proto.InterruptResponse() + resp.session_id = self._session_id + return resp + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.client.test_client import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach.py new file mode 100644 index 0000000..8bc6d8b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.streaming.test_streaming_foreach import StreamingTestsForeachMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class StreamingForeachParityTests(StreamingTestsForeachMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.streaming.test_parity_foreach import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py new file mode 100644 index 0000000..74371b6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py @@ -0,0 +1,44 @@ +# +# 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 unittest + +from pyspark.sql.tests.streaming.test_streaming_foreach_batch import StreamingTestsForeachBatchMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class StreamingForeachBatchParityTests(StreamingTestsForeachBatchMixin, ReusedConnectTestCase): + @unittest.skip("SPARK-44463: Error handling needs improvement in connect foreachBatch") + def test_streaming_foreach_batch_propagates_python_errors(self): + super().test_streaming_foreach_batch_propagates_python_errors() + + @unittest.skip("This seems specific to py4j and pinned threads. The intention is unclear") + def test_streaming_foreach_batch_graceful_stop(self): + super().test_streaming_foreach_batch_graceful_stop() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.streaming.test_parity_foreach_batch import * # noqa: F401,E501 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py new file mode 100644 index 0000000..5069a76 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py @@ -0,0 +1,108 @@ +# +# 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 unittest +import time + +import pyspark.cloudpickle +from pyspark.sql.tests.streaming.test_streaming_listener import StreamingListenerTestsMixin +from pyspark.sql.streaming.listener import StreamingQueryListener +from pyspark.sql.functions import count, lit +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class TestListener(StreamingQueryListener): + def onQueryStarted(self, event): + e = pyspark.cloudpickle.dumps(event) + df = self.spark.createDataFrame(data=[(e,)]) + df.write.mode("append").saveAsTable("listener_start_events") + + def onQueryProgress(self, event): + e = pyspark.cloudpickle.dumps(event) + df = self.spark.createDataFrame(data=[(e,)]) + df.write.mode("append").saveAsTable("listener_progress_events") + + def onQueryIdle(self, event): + pass + + def onQueryTerminated(self, event): + e = pyspark.cloudpickle.dumps(event) + df = self.spark.createDataFrame(data=[(e,)]) + df.write.mode("append").saveAsTable("listener_terminated_events") + + +class StreamingListenerParityTests(StreamingListenerTestsMixin, ReusedConnectTestCase): + def test_listener_events(self): + test_listener = TestListener() + + try: + self.spark.streams.addListener(test_listener) + + # This ensures the read socket on the server won't crash (i.e. because of timeout) + # when there hasn't been a new event for a long time + time.sleep(30) + + df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load() + df_observe = df.observe("my_event", count(lit(1)).alias("rc")) + df_stateful = df_observe.groupBy().count() # make query stateful + q = ( + df_stateful.writeStream.format("noop") + .queryName("test") + .outputMode("complete") + .start() + ) + + self.assertTrue(q.isActive) + time.sleep(10) + self.assertTrue(q.lastProgress["batchId"] > 0) # ensure at least one batch is ran + q.stop() + self.assertFalse(q.isActive) + + start_event = pyspark.cloudpickle.loads( + self.spark.read.table("listener_start_events").collect()[0][0] + ) + + progress_event = pyspark.cloudpickle.loads( + self.spark.read.table("listener_progress_events").collect()[0][0] + ) + + terminated_event = pyspark.cloudpickle.loads( + self.spark.read.table("listener_terminated_events").collect()[0][0] + ) + + self.check_start_event(start_event) + self.check_progress_event(progress_event) + self.check_terminated_event(terminated_event) + + finally: + self.spark.streams.removeListener(test_listener) + + # Remove again to verify this won't throw any error + self.spark.streams.removeListener(test_listener) + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.streaming.test_parity_listener import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_streaming.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_streaming.py new file mode 100644 index 0000000..6fe2b89 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/streaming/test_parity_streaming.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.streaming.test_streaming import StreamingTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class StreamingParityTests(StreamingTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.streaming.test_parity_streaming import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_basic.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_basic.py new file mode 100644 index 0000000..2904eb4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_basic.py @@ -0,0 +1,3600 @@ +# +# 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 array +import datetime +import os +import unittest +import random +import shutil +import string +import tempfile +import uuid +from collections import defaultdict + +from pyspark.errors import ( + PySparkAttributeError, + PySparkTypeError, + PySparkException, + PySparkValueError, +) +from pyspark.errors.exceptions.base import SessionNotSameException +from pyspark.sql import SparkSession as PySparkSession, Row +from pyspark.sql.types import ( + StructType, + StructField, + LongType, + StringType, + IntegerType, + MapType, + ArrayType, + Row, +) + +from pyspark.testing.sqlutils import ( + MyObject, + SQLTestUtils, + PythonOnlyUDT, + ExamplePoint, + PythonOnlyPoint, +) +from pyspark.testing.connectutils import ( + should_test_connect, + ReusedConnectTestCase, + connect_requirement_message, +) +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.errors.exceptions.connect import ( + AnalysisException, + ParseException, + SparkConnectException, +) + +if should_test_connect: + import grpc + import pandas as pd + import numpy as np + from pyspark.sql.connect.proto import Expression as ProtoExpression + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + from pyspark.sql.connect.client import ChannelBuilder + from pyspark.sql.connect.column import Column + from pyspark.sql.connect.readwriter import DataFrameWriterV2 + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.connect.dataframe import DataFrame as CDataFrame + from pyspark.sql import functions as SF + from pyspark.sql.connect import functions as CF + from pyspark.sql.connect.client.core import Retrying, SparkConnectClient + + +class SparkConnectSQLTestCase(ReusedConnectTestCase, SQLTestUtils, PandasOnSparkTestUtils): + """Parent test fixture class for all Spark Connect related + test cases.""" + + @classmethod + def setUpClass(cls): + super(SparkConnectSQLTestCase, cls).setUpClass() + # Disable the shared namespace so pyspark.sql.functions, etc point the regular + # PySpark libraries. + os.environ["PYSPARK_NO_NAMESPACE_SHARE"] = "1" + + cls.connect = cls.spark # Switch Spark Connect session and regular PySpark session. + cls.spark = PySparkSession._instantiatedSession + assert cls.spark is not None + + cls.testData = [Row(key=i, value=str(i)) for i in range(100)] + cls.testDataStr = [Row(key=str(i)) for i in range(100)] + cls.df = cls.spark.sparkContext.parallelize(cls.testData).toDF() + cls.df_text = cls.spark.sparkContext.parallelize(cls.testDataStr).toDF() + + cls.tbl_name = "test_connect_basic_table_1" + cls.tbl_name2 = "test_connect_basic_table_2" + cls.tbl_name3 = "test_connect_basic_table_3" + cls.tbl_name4 = "test_connect_basic_table_4" + cls.tbl_name_empty = "test_connect_basic_table_empty" + + # Cleanup test data + cls.spark_connect_clean_up_test_data() + # Load test data + cls.spark_connect_load_test_data() + + @classmethod + def tearDownClass(cls): + try: + cls.spark_connect_clean_up_test_data() + # Stopping Spark Connect closes the session in JVM at the server. + cls.spark = cls.connect + del os.environ["PYSPARK_NO_NAMESPACE_SHARE"] + finally: + super(SparkConnectSQLTestCase, cls).tearDownClass() + + @classmethod + def spark_connect_load_test_data(cls): + df = cls.spark.createDataFrame([(x, f"{x}") for x in range(100)], ["id", "name"]) + # Since we might create multiple Spark sessions, we need to create global temporary view + # that is specifically maintained in the "global_temp" schema. + df.write.saveAsTable(cls.tbl_name) + df2 = cls.spark.createDataFrame( + [(x, f"{x}", 2 * x) for x in range(100)], ["col1", "col2", "col3"] + ) + df2.write.saveAsTable(cls.tbl_name2) + df3 = cls.spark.createDataFrame([(x, f"{x}") for x in range(100)], ["id", "test\n_column"]) + df3.write.saveAsTable(cls.tbl_name3) + df4 = cls.spark.createDataFrame( + [(x, {"a": x}, [x, x * 2]) for x in range(100)], ["id", "map_column", "array_column"] + ) + df4.write.saveAsTable(cls.tbl_name4) + empty_table_schema = StructType( + [ + StructField("firstname", StringType(), True), + StructField("middlename", StringType(), True), + StructField("lastname", StringType(), True), + ] + ) + emptyRDD = cls.spark.sparkContext.emptyRDD() + empty_df = cls.spark.createDataFrame(emptyRDD, empty_table_schema) + empty_df.write.saveAsTable(cls.tbl_name_empty) + + @classmethod + def spark_connect_clean_up_test_data(cls): + cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name)) + cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name2)) + cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name3)) + cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name4)) + cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name_empty)) + + +class SparkConnectBasicTests(SparkConnectSQLTestCase): + def test_df_getattr_behavior(self): + cdf = self.connect.range(10) + sdf = self.spark.range(10) + + sdf._simple_extension = 10 + cdf._simple_extension = 10 + + self.assertEqual(sdf._simple_extension, cdf._simple_extension) + self.assertEqual(type(sdf._simple_extension), type(cdf._simple_extension)) + + self.assertTrue(hasattr(cdf, "_simple_extension")) + self.assertFalse(hasattr(cdf, "_simple_extension_does_not_exsit")) + + def test_df_get_item(self): + # SPARK-41779: test __getitem__ + + query = """ + SELECT * FROM VALUES + (true, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # | true| 1|NULL| + # |false|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # filter + self.assert_eq( + cdf[cdf.a].toPandas(), + sdf[sdf.a].toPandas(), + ) + self.assert_eq( + cdf[cdf.b.isin(2, 3)].toPandas(), + sdf[sdf.b.isin(2, 3)].toPandas(), + ) + self.assert_eq( + cdf[cdf.c > 1.5].toPandas(), + sdf[sdf.c > 1.5].toPandas(), + ) + + # select + self.assert_eq( + cdf[[cdf.a, "b", cdf.c]].toPandas(), + sdf[[sdf.a, "b", sdf.c]].toPandas(), + ) + self.assert_eq( + cdf[(cdf.a, "b", cdf.c)].toPandas(), + sdf[(sdf.a, "b", sdf.c)].toPandas(), + ) + + # select by index + self.assertTrue(isinstance(cdf[0], Column)) + self.assertTrue(isinstance(cdf[1], Column)) + self.assertTrue(isinstance(cdf[2], Column)) + + self.assert_eq( + cdf[[cdf[0], cdf[1], cdf[2]]].toPandas(), + sdf[[sdf[0], sdf[1], sdf[2]]].toPandas(), + ) + + # check error + with self.assertRaises(PySparkTypeError) as pe: + cdf[1.5] + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "item", + "arg_type": "float", + }, + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf[None] + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "item", + "arg_type": "NoneType", + }, + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf[cdf] + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "item", + "arg_type": "DataFrame", + }, + ) + + def test_error_handling(self): + # SPARK-41533 Proper error handling for Spark Connect + df = self.connect.range(10).select("id2") + with self.assertRaises(AnalysisException): + df.collect() + + def test_simple_read(self): + df = self.connect.read.table(self.tbl_name) + data = df.limit(10).toPandas() + # Check that the limit is applied + self.assertEqual(len(data.index), 10) + + def test_json(self): + with tempfile.TemporaryDirectory() as d: + # Write a DataFrame into a JSON file + self.spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]).write.mode( + "overwrite" + ).format("json").save(d) + # Read the JSON file as a DataFrame. + self.assert_eq(self.connect.read.json(d).toPandas(), self.spark.read.json(d).toPandas()) + + for schema in [ + "age INT, name STRING", + StructType( + [ + StructField("age", IntegerType()), + StructField("name", StringType()), + ] + ), + ]: + self.assert_eq( + self.connect.read.json(path=d, schema=schema).toPandas(), + self.spark.read.json(path=d, schema=schema).toPandas(), + ) + + self.assert_eq( + self.connect.read.json(path=d, primitivesAsString=True).toPandas(), + self.spark.read.json(path=d, primitivesAsString=True).toPandas(), + ) + + def test_parquet(self): + # SPARK-41445: Implement DataFrameReader.parquet + with tempfile.TemporaryDirectory() as d: + # Write a DataFrame into a JSON file + self.spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]).write.mode( + "overwrite" + ).format("parquet").save(d) + # Read the Parquet file as a DataFrame. + self.assert_eq( + self.connect.read.parquet(d).toPandas(), self.spark.read.parquet(d).toPandas() + ) + + def test_text(self): + # SPARK-41849: Implement DataFrameReader.text + with tempfile.TemporaryDirectory() as d: + # Write a DataFrame into a text file + self.spark.createDataFrame( + [{"name": "Sandeep Singh"}, {"name": "Hyukjin Kwon"}] + ).write.mode("overwrite").format("text").save(d) + # Read the text file as a DataFrame. + self.assert_eq(self.connect.read.text(d).toPandas(), self.spark.read.text(d).toPandas()) + + def test_csv(self): + # SPARK-42011: Implement DataFrameReader.csv + with tempfile.TemporaryDirectory() as d: + # Write a DataFrame into a text file + self.spark.createDataFrame( + [{"name": "Sandeep Singh"}, {"name": "Hyukjin Kwon"}] + ).write.mode("overwrite").format("csv").save(d) + # Read the text file as a DataFrame. + self.assert_eq(self.connect.read.csv(d).toPandas(), self.spark.read.csv(d).toPandas()) + + def test_multi_paths(self): + # SPARK-42041: DataFrameReader should support list of paths + + with tempfile.TemporaryDirectory() as d: + text_files = [] + for i in range(0, 3): + text_file = f"{d}/text-{i}.text" + shutil.copyfile("python/test_support/sql/text-test.txt", text_file) + text_files.append(text_file) + + self.assertEqual( + self.connect.read.text(text_files).collect(), + self.spark.read.text(text_files).collect(), + ) + + with tempfile.TemporaryDirectory() as d: + json_files = [] + for i in range(0, 5): + json_file = f"{d}/json-{i}.json" + shutil.copyfile("python/test_support/sql/people.json", json_file) + json_files.append(json_file) + + self.assertEqual( + self.connect.read.json(json_files).collect(), + self.spark.read.json(json_files).collect(), + ) + + def test_orc(self): + # SPARK-42012: Implement DataFrameReader.orc + with tempfile.TemporaryDirectory() as d: + # Write a DataFrame into a text file + self.spark.createDataFrame( + [{"name": "Sandeep Singh"}, {"name": "Hyukjin Kwon"}] + ).write.mode("overwrite").format("orc").save(d) + # Read the text file as a DataFrame. + self.assert_eq(self.connect.read.orc(d).toPandas(), self.spark.read.orc(d).toPandas()) + + def test_join_condition_column_list_columns(self): + left_connect_df = self.connect.read.table(self.tbl_name) + right_connect_df = self.connect.read.table(self.tbl_name2) + left_spark_df = self.spark.read.table(self.tbl_name) + right_spark_df = self.spark.read.table(self.tbl_name2) + joined_plan = left_connect_df.join( + other=right_connect_df, on=left_connect_df.id == right_connect_df.col1, how="inner" + ) + joined_plan2 = left_spark_df.join( + other=right_spark_df, on=left_spark_df.id == right_spark_df.col1, how="inner" + ) + self.assert_eq(joined_plan.toPandas(), joined_plan2.toPandas()) + + joined_plan3 = left_connect_df.join( + other=right_connect_df, + on=[ + left_connect_df.id == right_connect_df.col1, + left_connect_df.name == right_connect_df.col2, + ], + how="inner", + ) + joined_plan4 = left_spark_df.join( + other=right_spark_df, + on=[left_spark_df.id == right_spark_df.col1, left_spark_df.name == right_spark_df.col2], + how="inner", + ) + self.assert_eq(joined_plan3.toPandas(), joined_plan4.toPandas()) + + def test_join_ambiguous_cols(self): + # SPARK-41812: test join with ambiguous columns + data1 = [Row(id=1, value="foo"), Row(id=2, value=None)] + cdf1 = self.connect.createDataFrame(data1) + sdf1 = self.spark.createDataFrame(data1) + + data2 = [Row(value="bar"), Row(value=None), Row(value="foo")] + cdf2 = self.connect.createDataFrame(data2) + sdf2 = self.spark.createDataFrame(data2) + + cdf3 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"]) + sdf3 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"]) + + self.assertEqual(cdf3.schema, sdf3.schema) + self.assertEqual(cdf3.collect(), sdf3.collect()) + + cdf4 = cdf1.join(cdf2, cdf1["value"].eqNullSafe(cdf2["value"])) + sdf4 = sdf1.join(sdf2, sdf1["value"].eqNullSafe(sdf2["value"])) + + self.assertEqual(cdf4.schema, sdf4.schema) + self.assertEqual(cdf4.collect(), sdf4.collect()) + + cdf5 = cdf1.join( + cdf2, (cdf1["value"] == cdf2["value"]) & (cdf1["value"].eqNullSafe(cdf2["value"])) + ) + sdf5 = sdf1.join( + sdf2, (sdf1["value"] == sdf2["value"]) & (sdf1["value"].eqNullSafe(sdf2["value"])) + ) + + self.assertEqual(cdf5.schema, sdf5.schema) + self.assertEqual(cdf5.collect(), sdf5.collect()) + + cdf6 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"]).select(cdf1.value) + sdf6 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"]).select(sdf1.value) + + self.assertEqual(cdf6.schema, sdf6.schema) + self.assertEqual(cdf6.collect(), sdf6.collect()) + + cdf7 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"]).select(cdf2.value) + sdf7 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"]).select(sdf2.value) + + self.assertEqual(cdf7.schema, sdf7.schema) + self.assertEqual(cdf7.collect(), sdf7.collect()) + + def test_invalid_column(self): + # SPARK-41812: fail df1.select(df2.col) + data1 = [Row(a=1, b=2, c=3)] + cdf1 = self.connect.createDataFrame(data1) + + data2 = [Row(a=2, b=0)] + cdf2 = self.connect.createDataFrame(data2) + + with self.assertRaises(AnalysisException): + cdf1.select(cdf2.a).schema + + with self.assertRaises(AnalysisException): + cdf2.withColumn("x", cdf1.a + 1).schema + + with self.assertRaisesRegex(AnalysisException, "attribute.*missing"): + cdf3 = cdf1.select(cdf1.a) + cdf3.select(cdf1.b).schema + + def test_collect(self): + cdf = self.connect.read.table(self.tbl_name) + sdf = self.spark.read.table(self.tbl_name) + + data = cdf.limit(10).collect() + self.assertEqual(len(data), 10) + # Check Row has schema column names. + self.assertTrue("name" in data[0]) + self.assertTrue("id" in data[0]) + + cdf = cdf.select( + CF.log("id"), CF.log("id"), CF.struct("id", "name"), CF.struct("id", "name") + ).limit(10) + sdf = sdf.select( + SF.log("id"), SF.log("id"), SF.struct("id", "name"), SF.struct("id", "name") + ).limit(10) + + self.assertEqual( + cdf.collect(), + sdf.collect(), + ) + + def test_collect_timestamp(self): + query = """ + SELECT * FROM VALUES + (TIMESTAMP('2022-12-25 10:30:00'), 1), + (TIMESTAMP('2022-12-25 10:31:00'), 2), + (TIMESTAMP('2022-12-25 10:32:00'), 1), + (TIMESTAMP('2022-12-25 10:33:00'), 2), + (TIMESTAMP('2022-12-26 09:30:00'), 1), + (TIMESTAMP('2022-12-26 09:35:00'), 3) + AS tab(date, val) + """ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assertEqual(cdf.schema, sdf.schema) + + self.assertEqual(cdf.collect(), sdf.collect()) + + self.assertEqual( + cdf.select(CF.date_trunc("year", cdf.date).alias("year")).collect(), + sdf.select(SF.date_trunc("year", sdf.date).alias("year")).collect(), + ) + + def test_with_columns_renamed(self): + # SPARK-41312: test DataFrame.withColumnsRenamed() + self.assertEqual( + self.connect.read.table(self.tbl_name).withColumnRenamed("id", "id_new").schema, + self.spark.read.table(self.tbl_name).withColumnRenamed("id", "id_new").schema, + ) + self.assertEqual( + self.connect.read.table(self.tbl_name) + .withColumnsRenamed({"id": "id_new", "name": "name_new"}) + .schema, + self.spark.read.table(self.tbl_name) + .withColumnsRenamed({"id": "id_new", "name": "name_new"}) + .schema, + ) + + def test_with_local_data(self): + """SPARK-41114: Test creating a dataframe using local data""" + pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) + df = self.connect.createDataFrame(pdf) + rows = df.filter(df.a == CF.lit(3)).collect() + self.assertTrue(len(rows) == 1) + self.assertEqual(rows[0][0], 3) + self.assertEqual(rows[0][1], "c") + + # Check correct behavior for empty DataFrame + pdf = pd.DataFrame({"a": []}) + with self.assertRaises(ValueError): + self.connect.createDataFrame(pdf) + + def test_with_local_ndarray(self): + """SPARK-41446: Test creating a dataframe using local list""" + data = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) + + sdf = self.spark.createDataFrame(data) + cdf = self.connect.createDataFrame(data) + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + for schema in [ + StructType( + [ + StructField("col1", IntegerType(), True), + StructField("col2", IntegerType(), True), + StructField("col3", IntegerType(), True), + StructField("col4", IntegerType(), True), + ] + ), + "struct", + "col1 int, col2 int, col3 int, col4 int", + "col1 int, col2 long, col3 string, col4 long", + "col1 int, col2 string, col3 short, col4 long", + ["a", "b", "c", "d"], + ("x1", "x2", "x3", "x4"), + ]: + with self.subTest(schema=schema): + sdf = self.spark.createDataFrame(data, schema=schema) + cdf = self.connect.createDataFrame(data, schema=schema) + + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + with self.assertRaises(PySparkValueError) as pe: + self.connect.createDataFrame(data, ["a", "b", "c", "d", "e"]) + + self.check_error( + exception=pe.exception, + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={"expected_length": "5", "actual_length": "4"}, + ) + + with self.assertRaises(ParseException): + self.connect.createDataFrame(data, "col1 magic_type, col2 int, col3 int, col4 int") + + with self.assertRaises(PySparkValueError) as pe: + self.connect.createDataFrame(data, "col1 int, col2 int, col3 int") + + self.check_error( + exception=pe.exception, + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={"expected_length": "3", "actual_length": "4"}, + ) + + # test 1 dim ndarray + data = np.array([1.0, 2.0, np.nan, 3.0, 4.0, float("NaN"), 5.0]) + self.assertEqual(data.ndim, 1) + + sdf = self.spark.createDataFrame(data) + cdf = self.connect.createDataFrame(data) + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + def test_with_local_list(self): + """SPARK-41446: Test creating a dataframe using local list""" + data = [[1, 2, 3, 4]] + + sdf = self.spark.createDataFrame(data) + cdf = self.connect.createDataFrame(data) + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + for schema in [ + "struct", + "col1 int, col2 int, col3 int, col4 int", + "col1 int, col2 long, col3 string, col4 long", + "col1 int, col2 string, col3 short, col4 long", + ["a", "b", "c", "d"], + ("x1", "x2", "x3", "x4"), + ]: + sdf = self.spark.createDataFrame(data, schema=schema) + cdf = self.connect.createDataFrame(data, schema=schema) + + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + with self.assertRaises(PySparkValueError) as pe: + self.connect.createDataFrame(data, ["a", "b", "c", "d", "e"]) + + self.check_error( + exception=pe.exception, + error_class="AXIS_LENGTH_MISMATCH", + message_parameters={"expected_length": "5", "actual_length": "4"}, + ) + + with self.assertRaises(ParseException): + self.connect.createDataFrame(data, "col1 magic_type, col2 int, col3 int, col4 int") + + with self.assertRaisesRegex( + ValueError, + "Length mismatch: Expected axis has 3 elements, new values have 4 elements", + ): + self.connect.createDataFrame(data, "col1 int, col2 int, col3 int") + + def test_with_local_rows(self): + # SPARK-41789, SPARK-41810: Test creating a dataframe with list of rows and dictionaries + rows = [ + Row(course="dotNET", year=2012, earnings=10000), + Row(course="Java", year=2012, earnings=20000), + Row(course="dotNET", year=2012, earnings=5000), + Row(course="dotNET", year=2013, earnings=48000), + Row(course="Java", year=2013, earnings=30000), + Row(course="Scala", year=2022, earnings=None), + ] + dicts = [row.asDict() for row in rows] + + for data in [rows, dicts]: + sdf = self.spark.createDataFrame(data) + cdf = self.connect.createDataFrame(data) + + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + # test with rename + sdf = self.spark.createDataFrame(data, schema=["a", "b", "c"]) + cdf = self.connect.createDataFrame(data, schema=["a", "b", "c"]) + + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + def test_streaming_local_relation(self): + threshold_conf = "spark.sql.session.localRelationCacheThreshold" + old_threshold = self.connect.conf.get(threshold_conf) + threshold = 1024 * 1024 + self.connect.conf.set(threshold_conf, threshold) + try: + suffix = "abcdef" + letters = string.ascii_lowercase + str = "".join(random.choice(letters) for i in range(threshold)) + suffix + data = [[0, str], [1, str]] + for i in range(0, 2): + cdf = self.connect.createDataFrame(data, ["a", "b"]) + self.assert_eq(cdf.count(), len(data)) + self.assert_eq(cdf.filter(f"endsWith(b, '{suffix}')").isEmpty(), False) + finally: + self.connect.conf.set(threshold_conf, old_threshold) + + def test_with_atom_type(self): + for data in [[(1), (2), (3)], [1, 2, 3]]: + for schema in ["long", "int", "short"]: + sdf = self.spark.createDataFrame(data, schema=schema) + cdf = self.connect.createDataFrame(data, schema=schema) + + self.assertEqual(sdf.schema, cdf.schema) + self.assert_eq(sdf.toPandas(), cdf.toPandas()) + + def test_with_none_and_nan(self): + # SPARK-41855: make createDataFrame support None and NaN + # SPARK-41814: test with eqNullSafe + data1 = [Row(id=1, value=float("NaN")), Row(id=2, value=42.0), Row(id=3, value=None)] + data2 = [Row(id=1, value=np.nan), Row(id=2, value=42.0), Row(id=3, value=None)] + data3 = [ + {"id": 1, "value": float("NaN")}, + {"id": 2, "value": 42.0}, + {"id": 3, "value": None}, + ] + data4 = [{"id": 1, "value": np.nan}, {"id": 2, "value": 42.0}, {"id": 3, "value": None}] + data5 = [(1, float("NaN")), (2, 42.0), (3, None)] + data6 = [(1, np.nan), (2, 42.0), (3, None)] + data7 = np.array([[1, float("NaN")], [2, 42.0], [3, None]]) + data8 = np.array([[1, np.nan], [2, 42.0], [3, None]]) + + # +---+-----+ + # | id|value| + # +---+-----+ + # | 1| NaN| + # | 2| 42.0| + # | 3| NULL| + # +---+-----+ + + for data in [data1, data2, data3, data4, data5, data6, data7, data8]: + if isinstance(data[0], (Row, dict)): + # data1, data2, data3, data4 + cdf = self.connect.createDataFrame(data) + sdf = self.spark.createDataFrame(data) + else: + # data5, data6, data7, data8 + cdf = self.connect.createDataFrame(data, schema=["id", "value"]) + sdf = self.spark.createDataFrame(data, schema=["id", "value"]) + + self.assert_eq(cdf.toPandas(), sdf.toPandas()) + + self.assert_eq( + cdf.select( + cdf["value"].eqNullSafe(None), + cdf["value"].eqNullSafe(float("NaN")), + cdf["value"].eqNullSafe(42.0), + ).toPandas(), + sdf.select( + sdf["value"].eqNullSafe(None), + sdf["value"].eqNullSafe(float("NaN")), + sdf["value"].eqNullSafe(42.0), + ).toPandas(), + ) + + # SPARK-41851: test with nanvl + data = [(1.0, float("nan")), (float("nan"), 2.0)] + + cdf = self.connect.createDataFrame(data, ("a", "b")) + sdf = self.spark.createDataFrame(data, ("a", "b")) + + self.assert_eq(cdf.toPandas(), sdf.toPandas()) + + self.assert_eq( + cdf.select( + CF.nanvl("a", "b").alias("r1"), CF.nanvl(cdf.a, cdf.b).alias("r2") + ).toPandas(), + sdf.select( + SF.nanvl("a", "b").alias("r1"), SF.nanvl(sdf.a, sdf.b).alias("r2") + ).toPandas(), + ) + + # SPARK-41852: test with pmod + data = [ + (1.0, float("nan")), + (float("nan"), 2.0), + (10.0, 3.0), + (float("nan"), float("nan")), + (-3.0, 4.0), + (-10.0, 3.0), + (-5.0, -6.0), + (7.0, -8.0), + (1.0, 2.0), + ] + + cdf = self.connect.createDataFrame(data, ("a", "b")) + sdf = self.spark.createDataFrame(data, ("a", "b")) + + self.assert_eq(cdf.toPandas(), sdf.toPandas()) + + self.assert_eq( + cdf.select(CF.pmod("a", "b")).toPandas(), + sdf.select(SF.pmod("a", "b")).toPandas(), + ) + + def test_cast_with_ddl(self): + data = [Row(date=datetime.date(2021, 12, 27), add=2)] + + cdf = self.connect.createDataFrame(data, "date date, add integer") + sdf = self.spark.createDataFrame(data, "date date, add integer") + + self.assertEqual(cdf.schema, sdf.schema) + + def test_create_empty_df(self): + for schema in [ + "STRING", + "x STRING", + "x STRING, y INTEGER", + StringType(), + StructType( + [ + StructField("x", StringType(), True), + StructField("y", IntegerType(), True), + ] + ), + ]: + cdf = self.connect.createDataFrame(data=[], schema=schema) + sdf = self.spark.createDataFrame(data=[], schema=schema) + + self.assert_eq(cdf.toPandas(), sdf.toPandas()) + + # check error + with self.assertRaises(PySparkValueError) as pe: + self.connect.createDataFrame(data=[]) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_INFER_EMPTY_SCHEMA", + message_parameters={}, + ) + + def test_create_dataframe_from_arrays(self): + # SPARK-42021: createDataFrame support array.array + data1 = [Row(a=1, b=array.array("i", [1, 2, 3]), c=array.array("d", [4, 5, 6]))] + data2 = [(array.array("d", [1, 2, 3]), 2, "3")] + data3 = [{"a": 1, "b": array.array("i", [1, 2, 3])}] + + for data in [data1, data2, data3]: + cdf = self.connect.createDataFrame(data) + sdf = self.spark.createDataFrame(data) + + # TODO: the nullability is different, need to fix + # self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_timestampe_create_from_rows(self): + data = [(datetime.datetime(2016, 3, 11, 9, 0, 7), 1)] + + cdf = self.connect.createDataFrame(data, ["date", "val"]) + sdf = self.spark.createDataFrame(data, ["date", "val"]) + + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_create_dataframe_with_coercion(self): + data1 = [[1.33, 1], ["2.1", 1]] + data2 = [[True, 1], ["false", 1]] + + for data in [data1, data2]: + cdf = self.connect.createDataFrame(data, ["a", "b"]) + sdf = self.spark.createDataFrame(data, ["a", "b"]) + + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_nested_type_create_from_rows(self): + data1 = [Row(a=1, b=Row(c=2, d=Row(e=3, f=Row(g=4, h=Row(i=5)))))] + # root + # |-- a: long (nullable = true) + # |-- b: struct (nullable = true) + # | |-- c: long (nullable = true) + # | |-- d: struct (nullable = true) + # | | |-- e: long (nullable = true) + # | | |-- f: struct (nullable = true) + # | | | |-- g: long (nullable = true) + # | | | |-- h: struct (nullable = true) + # | | | | |-- i: long (nullable = true) + + data2 = [ + ( + 1, + "a", + Row( + a=1, + b=[1, 2, 3], + c={"a": "b"}, + d=Row(x=1, y="y", z=Row(o=1, p=2, q=Row(g=1.5))), + ), + ) + ] + # root + # |-- _1: long (nullable = true) + # |-- _2: string (nullable = true) + # |-- _3: struct (nullable = true) + # | |-- a: long (nullable = true) + # | |-- b: array (nullable = true) + # | | |-- element: long (containsNull = true) + # | |-- c: map (nullable = true) + # | | |-- key: string + # | | |-- value: string (valueContainsNull = true) + # | |-- d: struct (nullable = true) + # | | |-- x: long (nullable = true) + # | | |-- y: string (nullable = true) + # | | |-- z: struct (nullable = true) + # | | | |-- o: long (nullable = true) + # | | | |-- p: long (nullable = true) + # | | | |-- q: struct (nullable = true) + # | | | | |-- g: double (nullable = true) + + data3 = [ + Row( + a=1, + b=[1, 2, 3], + c={"a": "b"}, + d=Row(x=1, y="y", z=Row(1, 2, 3)), + e=list("hello connect"), + ) + ] + # root + # |-- a: long (nullable = true) + # |-- b: array (nullable = true) + # | |-- element: long (containsNull = true) + # |-- c: map (nullable = true) + # | |-- key: string + # | |-- value: string (valueContainsNull = true) + # |-- d: struct (nullable = true) + # | |-- x: long (nullable = true) + # | |-- y: string (nullable = true) + # | |-- z: struct (nullable = true) + # | | |-- _1: long (nullable = true) + # | | |-- _2: long (nullable = true) + # | | |-- _3: long (nullable = true) + # |-- e: array (nullable = true) + # | |-- element: string (containsNull = true) + + data4 = [ + { + "a": 1, + "b": Row(x=1, y=Row(z=2)), + "c": {"x": -1, "y": 2}, + "d": [1, 2, 3, 4, 5], + } + ] + # root + # |-- a: long (nullable = true) + # |-- b: struct (nullable = true) + # | |-- x: long (nullable = true) + # | |-- y: struct (nullable = true) + # | | |-- z: long (nullable = true) + # |-- c: map (nullable = true) + # | |-- key: string + # | |-- value: long (valueContainsNull = true) + # |-- d: array (nullable = true) + # | |-- element: long (containsNull = true) + + data5 = [ + { + "a": [Row(x=1, y="2"), Row(x=-1, y="-2")], + "b": [[1, 2, 3], [4, 5], [6]], + "c": {3: {4: {5: 6}}, 7: {8: {9: 0}}}, + } + ] + # root + # |-- a: array (nullable = true) + # | |-- element: struct (containsNull = true) + # | | |-- x: long (nullable = true) + # | | |-- y: string (nullable = true) + # |-- b: array (nullable = true) + # | |-- element: array (containsNull = true) + # | | |-- element: long (containsNull = true) + # |-- c: map (nullable = true) + # | |-- key: long + # | |-- value: map (valueContainsNull = true) + # | | |-- key: long + # | | |-- value: map (valueContainsNull = true) + # | | | |-- key: long + # | | | |-- value: long (valueContainsNull = true) + + for data in [data1, data2, data3, data4, data5]: + with self.subTest(data=data): + cdf = self.connect.createDataFrame(data) + sdf = self.spark.createDataFrame(data) + + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_create_df_from_objects(self): + data = [MyObject(1, "1"), MyObject(2, "2")] + + # +---+-----+ + # |key|value| + # +---+-----+ + # | 1| 1| + # | 2| 2| + # +---+-----+ + + cdf = self.connect.createDataFrame(data) + sdf = self.spark.createDataFrame(data) + + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_simple_explain_string(self): + df = self.connect.read.table(self.tbl_name).limit(10) + result = df._explain_string() + self.assertGreater(len(result), 0) + + def test_schema(self): + schema = self.connect.read.table(self.tbl_name).schema + self.assertEqual( + StructType( + [StructField("id", LongType(), True), StructField("name", StringType(), True)] + ), + schema, + ) + + # test FloatType, DoubleType, DecimalType, StringType, BooleanType, NullType + query = """ + SELECT * FROM VALUES + (float(1.0), double(1.0), 1.0, "1", true, NULL), + (float(2.0), double(2.0), 2.0, "2", false, NULL), + (float(3.0), double(3.0), NULL, "3", false, NULL) + AS tab(a, b, c, d, e, f) + """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + # test TimestampType, DateType + query = """ + SELECT * FROM VALUES + (TIMESTAMP('2019-04-12 15:50:00'), DATE('2022-02-22')), + (TIMESTAMP('2019-04-12 15:50:00'), NULL), + (NULL, DATE('2022-02-22')) + AS tab(a, b) + """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + # test DayTimeIntervalType + query = """ SELECT INTERVAL '100 10:30' DAY TO MINUTE AS interval """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + # test MapType + query = """ + SELECT * FROM VALUES + (MAP('a', 'ab'), MAP('a', 'ab'), MAP(1, 2, 3, 4)), + (MAP('x', 'yz'), MAP('x', NULL), NULL), + (MAP('c', 'de'), NULL, MAP(-1, NULL, -3, -4)) + AS tab(a, b, c) + """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + # test ArrayType + query = """ + SELECT * FROM VALUES + (ARRAY('a', 'ab'), ARRAY(1, 2, 3), ARRAY(1, NULL, 3)), + (ARRAY('x', NULL), NULL, ARRAY(1, 3)), + (NULL, ARRAY(-1, -2, -3), Array()) + AS tab(a, b, c) + """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + # test StructType + query = """ + SELECT STRUCT(a, b, c, d), STRUCT(e, f, g), STRUCT(STRUCT(a, b), STRUCT(h)) FROM VALUES + (float(1.0), double(1.0), 1.0, "1", true, NULL, ARRAY(1, NULL, 3), MAP(1, 2, 3, 4)), + (float(2.0), double(2.0), 2.0, "2", false, NULL, ARRAY(1, 3), MAP(1, NULL, 3, 4)), + (float(3.0), double(3.0), NULL, "3", false, NULL, ARRAY(NULL), NULL) + AS tab(a, b, c, d, e, f, g, h) + """ + self.assertEqual( + self.spark.sql(query).schema, + self.connect.sql(query).schema, + ) + + def test_to(self): + # SPARK-41464: test DataFrame.to() + + cdf = self.connect.read.table(self.tbl_name) + df = self.spark.read.table(self.tbl_name) + + def assert_eq_schema(cdf: CDataFrame, df: DataFrame, schema: StructType): + cdf_to = cdf.to(schema) + df_to = df.to(schema) + self.assertEqual(cdf_to.schema, df_to.schema) + self.assert_eq(cdf_to.toPandas(), df_to.toPandas()) + + # The schema has not changed + schema = StructType( + [ + StructField("id", IntegerType(), True), + StructField("name", StringType(), True), + ] + ) + + assert_eq_schema(cdf, df, schema) + + # Change schema with struct + schema2 = StructType([StructField("struct", schema, False)]) + + cdf_to = cdf.select(CF.struct("id", "name").alias("struct")).to(schema2) + df_to = df.select(SF.struct("id", "name").alias("struct")).to(schema2) + + self.assertEqual(cdf_to.schema, df_to.schema) + + # Change the column name + schema = StructType( + [ + StructField("col1", IntegerType(), True), + StructField("col2", StringType(), True), + ] + ) + + assert_eq_schema(cdf, df, schema) + + # Change the column data type + schema = StructType( + [ + StructField("id", StringType(), True), + StructField("name", StringType(), True), + ] + ) + + assert_eq_schema(cdf, df, schema) + + # Reduce the column quantity and change data type + schema = StructType( + [ + StructField("id", LongType(), True), + ] + ) + + assert_eq_schema(cdf, df, schema) + + # incompatible field nullability + schema = StructType([StructField("id", LongType(), False)]) + self.assertRaisesRegex( + AnalysisException, + "NULLABLE_COLUMN_OR_FIELD", + lambda: cdf.to(schema).toPandas(), + ) + + # field cannot upcast + schema = StructType([StructField("name", LongType())]) + self.assertRaisesRegex( + AnalysisException, + "INVALID_COLUMN_OR_FIELD_DATA_TYPE", + lambda: cdf.to(schema).toPandas(), + ) + + schema = StructType( + [ + StructField("id", IntegerType(), True), + StructField("name", IntegerType(), True), + ] + ) + self.assertRaisesRegex( + AnalysisException, + "INVALID_COLUMN_OR_FIELD_DATA_TYPE", + lambda: cdf.to(schema).toPandas(), + ) + + # Test map type and array type + schema = StructType( + [ + StructField("id", StringType(), True), + StructField("my_map", MapType(StringType(), IntegerType(), False), True), + StructField("my_array", ArrayType(IntegerType(), False), True), + ] + ) + cdf = self.connect.read.table(self.tbl_name4) + df = self.spark.read.table(self.tbl_name4) + + assert_eq_schema(cdf, df, schema) + + def test_toDF(self): + # SPARK-41310: test DataFrame.toDF() + self.assertEqual( + self.connect.read.table(self.tbl_name).toDF("col1", "col2").schema, + self.spark.read.table(self.tbl_name).toDF("col1", "col2").schema, + ) + + def test_print_schema(self): + # SPARK-41216: Test print schema + tree_str = self.connect.sql("SELECT 1 AS X, 2 AS Y")._tree_string() + # root + # |-- X: integer (nullable = false) + # |-- Y: integer (nullable = false) + expected = "root\n |-- X: integer (nullable = false)\n |-- Y: integer (nullable = false)\n" + self.assertEqual(tree_str, expected) + + def test_is_local(self): + # SPARK-41216: Test is local + self.assertTrue(self.connect.sql("SHOW DATABASES").isLocal()) + self.assertFalse(self.connect.read.table(self.tbl_name).isLocal()) + + def test_is_streaming(self): + # SPARK-41216: Test is streaming + self.assertFalse(self.connect.read.table(self.tbl_name).isStreaming) + self.assertFalse(self.connect.sql("SELECT 1 AS X LIMIT 0").isStreaming) + + def test_input_files(self): + # SPARK-41216: Test input files + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + try: + self.df_text.write.text(tmpPath) + + input_files_list1 = ( + self.spark.read.format("text").schema("id STRING").load(path=tmpPath).inputFiles() + ) + input_files_list2 = ( + self.connect.read.format("text").schema("id STRING").load(path=tmpPath).inputFiles() + ) + + self.assertTrue(len(input_files_list1) > 0) + self.assertEqual(len(input_files_list1), len(input_files_list2)) + for file_path in input_files_list2: + self.assertTrue(file_path in input_files_list1) + finally: + shutil.rmtree(tmpPath) + + def test_limit_offset(self): + df = self.connect.read.table(self.tbl_name) + pd = df.limit(10).offset(1).toPandas() + self.assertEqual(9, len(pd.index)) + pd2 = df.offset(98).limit(10).toPandas() + self.assertEqual(2, len(pd2.index)) + + def test_tail(self): + df = self.connect.read.table(self.tbl_name) + df2 = self.spark.read.table(self.tbl_name) + self.assertEqual(df.tail(10), df2.tail(10)) + + def test_sql(self): + pdf = self.connect.sql("SELECT 1").toPandas() + self.assertEqual(1, len(pdf.index)) + + def test_sql_with_named_args(self): + df = self.connect.sql("SELECT * FROM range(10) WHERE id > :minId", args={"minId": 7}) + df2 = self.spark.sql("SELECT * FROM range(10) WHERE id > :minId", args={"minId": 7}) + self.assert_eq(df.toPandas(), df2.toPandas()) + + def test_sql_with_pos_args(self): + df = self.connect.sql("SELECT * FROM range(10) WHERE id > ?", args=[7]) + df2 = self.spark.sql("SELECT * FROM range(10) WHERE id > ?", args=[7]) + self.assert_eq(df.toPandas(), df2.toPandas()) + + def test_head(self): + # SPARK-41002: test `head` API in Python Client + df = self.connect.read.table(self.tbl_name) + self.assertIsNotNone(len(df.head())) + self.assertIsNotNone(len(df.head(1))) + self.assertIsNotNone(len(df.head(5))) + df2 = self.connect.read.table(self.tbl_name_empty) + self.assertIsNone(df2.head()) + + def test_deduplicate(self): + # SPARK-41326: test distinct and dropDuplicates. + df = self.connect.read.table(self.tbl_name) + df2 = self.spark.read.table(self.tbl_name) + self.assert_eq(df.distinct().toPandas(), df2.distinct().toPandas()) + self.assert_eq(df.dropDuplicates().toPandas(), df2.dropDuplicates().toPandas()) + self.assert_eq( + df.dropDuplicates(["name"]).toPandas(), df2.dropDuplicates(["name"]).toPandas() + ) + + def test_deduplicate_within_watermark_in_batch(self): + df = self.connect.read.table(self.tbl_name) + with self.assertRaisesRegex( + AnalysisException, + "dropDuplicatesWithinWatermark is not supported with batch DataFrames/DataSets", + ): + df.dropDuplicatesWithinWatermark().toPandas() + + def test_first(self): + # SPARK-41002: test `first` API in Python Client + df = self.connect.read.table(self.tbl_name) + self.assertIsNotNone(len(df.first())) + df2 = self.connect.read.table(self.tbl_name_empty) + self.assertIsNone(df2.first()) + + def test_take(self) -> None: + # SPARK-41002: test `take` API in Python Client + df = self.connect.read.table(self.tbl_name) + self.assertEqual(5, len(df.take(5))) + df2 = self.connect.read.table(self.tbl_name_empty) + self.assertEqual(0, len(df2.take(5))) + + def test_drop(self): + # SPARK-41169: test drop + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (false, NULL, 2), (NULL, 3, 3) + AS tab(a, b, c) + """ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + self.assert_eq( + cdf.drop("a").toPandas(), + sdf.drop("a").toPandas(), + ) + self.assert_eq( + cdf.drop("a", "b").toPandas(), + sdf.drop("a", "b").toPandas(), + ) + self.assert_eq( + cdf.drop("a", "x").toPandas(), + sdf.drop("a", "x").toPandas(), + ) + self.assert_eq( + cdf.drop(cdf.a, "x").toPandas(), + sdf.drop(sdf.a, "x").toPandas(), + ) + + def test_subquery_alias(self) -> None: + # SPARK-40938: test subquery alias. + plan_text = ( + self.connect.read.table(self.tbl_name) + .alias("special_alias") + ._explain_string(extended=True) + ) + self.assertTrue("special_alias" in plan_text) + + def test_sort(self): + # SPARK-41332: test sort + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # |false|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + self.assert_eq( + cdf.sort("a").toPandas(), + sdf.sort("a").toPandas(), + ) + self.assert_eq( + cdf.sort("c").toPandas(), + sdf.sort("c").toPandas(), + ) + self.assert_eq( + cdf.sort("b").toPandas(), + sdf.sort("b").toPandas(), + ) + self.assert_eq( + cdf.sort(cdf.c, "b").toPandas(), + sdf.sort(sdf.c, "b").toPandas(), + ) + self.assert_eq( + cdf.sort(cdf.c.desc(), "b").toPandas(), + sdf.sort(sdf.c.desc(), "b").toPandas(), + ) + self.assert_eq( + cdf.sort(cdf.c.desc(), cdf.a.asc()).toPandas(), + sdf.sort(sdf.c.desc(), sdf.a.asc()).toPandas(), + ) + + def test_range(self): + self.assert_eq( + self.connect.range(start=0, end=10).toPandas(), + self.spark.range(start=0, end=10).toPandas(), + ) + self.assert_eq( + self.connect.range(start=0, end=10, step=3).toPandas(), + self.spark.range(start=0, end=10, step=3).toPandas(), + ) + self.assert_eq( + self.connect.range(start=0, end=10, step=3, numPartitions=2).toPandas(), + self.spark.range(start=0, end=10, step=3, numPartitions=2).toPandas(), + ) + # SPARK-41301 + self.assert_eq( + self.connect.range(10).toPandas(), self.connect.range(start=0, end=10).toPandas() + ) + + def test_create_global_temp_view(self): + # SPARK-41127: test global temp view creation. + with self.tempView("view_1"): + self.connect.sql("SELECT 1 AS X LIMIT 0").createGlobalTempView("view_1") + self.connect.sql("SELECT 2 AS X LIMIT 1").createOrReplaceGlobalTempView("view_1") + self.assertTrue(self.spark.catalog.tableExists("global_temp.view_1")) + + # Test when creating a view which is already exists but + self.assertTrue(self.spark.catalog.tableExists("global_temp.view_1")) + with self.assertRaises(AnalysisException): + self.connect.sql("SELECT 1 AS X LIMIT 0").createGlobalTempView("view_1") + + def test_create_session_local_temp_view(self): + # SPARK-41372: test session local temp view creation. + with self.tempView("view_local_temp"): + self.connect.sql("SELECT 1 AS X").createTempView("view_local_temp") + self.assertEqual(self.connect.sql("SELECT * FROM view_local_temp").count(), 1) + self.connect.sql("SELECT 1 AS X LIMIT 0").createOrReplaceTempView("view_local_temp") + self.assertEqual(self.connect.sql("SELECT * FROM view_local_temp").count(), 0) + + # Test when creating a view which is already exists but + with self.assertRaises(AnalysisException): + self.connect.sql("SELECT 1 AS X LIMIT 0").createTempView("view_local_temp") + + def test_to_pandas(self): + # SPARK-41005: Test to pandas + query = """ + SELECT * FROM VALUES + (false, 1, NULL), + (false, NULL, float(2.0)), + (NULL, 3, float(3.0)) + AS tab(a, b, c) + """ + + self.assert_eq( + self.connect.sql(query).toPandas(), + self.spark.sql(query).toPandas(), + ) + + query = """ + SELECT * FROM VALUES + (1, 1, NULL), + (2, NULL, float(2.0)), + (3, 3, float(3.0)) + AS tab(a, b, c) + """ + + self.assert_eq( + self.connect.sql(query).toPandas(), + self.spark.sql(query).toPandas(), + ) + + query = """ + SELECT * FROM VALUES + (double(1.0), 1, "1"), + (NULL, NULL, NULL), + (double(2.0), 3, "3") + AS tab(a, b, c) + """ + + self.assert_eq( + self.connect.sql(query).toPandas(), + self.spark.sql(query).toPandas(), + ) + + query = """ + SELECT * FROM VALUES + (float(1.0), double(1.0), 1, "1"), + (float(2.0), double(2.0), 2, "2"), + (float(3.0), double(3.0), 3, "3") + AS tab(a, b, c, d) + """ + + self.assert_eq( + self.connect.sql(query).toPandas(), + self.spark.sql(query).toPandas(), + ) + + def test_create_dataframe_from_pandas_with_ns_timestamp(self): + """Truncate the timestamps for nanoseconds.""" + from datetime import datetime, timezone, timedelta + from pandas import Timestamp + import pandas as pd + + pdf = pd.DataFrame( + { + "naive": [datetime(2019, 1, 1, 0)], + "aware": [ + Timestamp( + year=2019, month=1, day=1, nanosecond=500, tz=timezone(timedelta(hours=-8)) + ) + ], + } + ) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + self.assertEqual( + self.connect.createDataFrame(pdf).collect(), + self.spark.createDataFrame(pdf).collect(), + ) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): + self.assertEqual( + self.connect.createDataFrame(pdf).collect(), + self.spark.createDataFrame(pdf).collect(), + ) + + def test_select_expr(self): + # SPARK-41201: test selectExpr API. + self.assert_eq( + self.connect.read.table(self.tbl_name).selectExpr("id * 2").toPandas(), + self.spark.read.table(self.tbl_name).selectExpr("id * 2").toPandas(), + ) + self.assert_eq( + self.connect.read.table(self.tbl_name) + .selectExpr(["id * 2", "cast(name as long) as name"]) + .toPandas(), + self.spark.read.table(self.tbl_name) + .selectExpr(["id * 2", "cast(name as long) as name"]) + .toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name) + .selectExpr("id * 2", "cast(name as long) as name") + .toPandas(), + self.spark.read.table(self.tbl_name) + .selectExpr("id * 2", "cast(name as long) as name") + .toPandas(), + ) + + def test_select_star(self): + data = [Row(a=1, b=Row(c=2, d=Row(e=3)))] + + # +---+--------+ + # | a| b| + # +---+--------+ + # | 1|{2, {3}}| + # +---+--------+ + + cdf = self.connect.createDataFrame(data=data) + sdf = self.spark.createDataFrame(data=data) + + self.assertEqual( + cdf.select("*").collect(), + sdf.select("*").collect(), + ) + self.assertEqual( + cdf.select("a", "*").collect(), + sdf.select("a", "*").collect(), + ) + self.assertEqual( + cdf.select("a", "b").collect(), + sdf.select("a", "b").collect(), + ) + self.assertEqual( + cdf.select("a", "b.*").collect(), + sdf.select("a", "b.*").collect(), + ) + + def test_fill_na(self): + # SPARK-41128: Test fill na + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # |false|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + self.assert_eq( + self.connect.sql(query).fillna(True).toPandas(), + self.spark.sql(query).fillna(True).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).fillna(2).toPandas(), + self.spark.sql(query).fillna(2).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).fillna(2, ["a", "b"]).toPandas(), + self.spark.sql(query).fillna(2, ["a", "b"]).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.fill({"a": True, "b": 2}).toPandas(), + self.spark.sql(query).na.fill({"a": True, "b": 2}).toPandas(), + ) + + def test_drop_na(self): + # SPARK-41148: Test drop na + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # |false|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + self.assert_eq( + self.connect.sql(query).dropna().toPandas(), + self.spark.sql(query).dropna().toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.drop(how="all", thresh=1).toPandas(), + self.spark.sql(query).na.drop(how="all", thresh=1).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).dropna(thresh=1, subset=("a", "b")).toPandas(), + self.spark.sql(query).dropna(thresh=1, subset=("a", "b")).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.drop(how="any", thresh=2, subset="a").toPandas(), + self.spark.sql(query).na.drop(how="any", thresh=2, subset="a").toPandas(), + ) + + def test_replace(self): + # SPARK-41315: Test replace + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # |false|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + self.assert_eq( + self.connect.sql(query).replace(2, 3).toPandas(), + self.spark.sql(query).replace(2, 3).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.replace(False, True).toPandas(), + self.spark.sql(query).na.replace(False, True).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).replace({1: 2, 3: -1}, subset=("a", "b")).toPandas(), + self.spark.sql(query).replace({1: 2, 3: -1}, subset=("a", "b")).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.replace((1, 2), (3, 1)).toPandas(), + self.spark.sql(query).na.replace((1, 2), (3, 1)).toPandas(), + ) + self.assert_eq( + self.connect.sql(query).na.replace((1, 2), (3, 1), subset=("c", "b")).toPandas(), + self.spark.sql(query).na.replace((1, 2), (3, 1), subset=("c", "b")).toPandas(), + ) + + with self.assertRaises(ValueError) as context: + self.connect.sql(query).replace({None: 1}, subset="a").toPandas() + self.assertTrue("Mixed type replacements are not supported" in str(context.exception)) + + with self.assertRaises(AnalysisException) as context: + self.connect.sql(query).replace({1: 2, 3: -1}, subset=("a", "x")).toPandas() + self.assertIn( + """Cannot resolve column name "x" among (a, b, c)""", str(context.exception) + ) + + def test_unpivot(self): + self.assert_eq( + self.connect.read.table(self.tbl_name) + .filter("id > 3") + .unpivot(["id"], ["name"], "variable", "value") + .toPandas(), + self.spark.read.table(self.tbl_name) + .filter("id > 3") + .unpivot(["id"], ["name"], "variable", "value") + .toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name) + .filter("id > 3") + .unpivot("id", None, "variable", "value") + .toPandas(), + self.spark.read.table(self.tbl_name) + .filter("id > 3") + .unpivot("id", None, "variable", "value") + .toPandas(), + ) + + def test_union_by_name(self): + # SPARK-41832: Test unionByName + data1 = [(1, 2, 3)] + data2 = [(6, 2, 5)] + df1_connect = self.connect.createDataFrame(data1, ["a", "b", "c"]) + df2_connect = self.connect.createDataFrame(data2, ["a", "b", "c"]) + union_df_connect = df1_connect.unionByName(df2_connect) + + df1_spark = self.spark.createDataFrame(data1, ["a", "b", "c"]) + df2_spark = self.spark.createDataFrame(data2, ["a", "b", "c"]) + union_df_spark = df1_spark.unionByName(df2_spark) + + self.assert_eq(union_df_connect.toPandas(), union_df_spark.toPandas()) + + df2_connect = self.connect.createDataFrame(data2, ["a", "B", "C"]) + union_df_connect = df1_connect.unionByName(df2_connect, allowMissingColumns=True) + + df2_spark = self.spark.createDataFrame(data2, ["a", "B", "C"]) + union_df_spark = df1_spark.unionByName(df2_spark, allowMissingColumns=True) + + self.assert_eq(union_df_connect.toPandas(), union_df_spark.toPandas()) + + def test_random_split(self): + # SPARK-41440: test randomSplit(weights, seed). + relations = ( + self.connect.read.table(self.tbl_name).filter("id > 3").randomSplit([1.0, 2.0, 3.0], 2) + ) + datasets = ( + self.spark.read.table(self.tbl_name).filter("id > 3").randomSplit([1.0, 2.0, 3.0], 2) + ) + + self.assertTrue(len(relations) == len(datasets)) + i = 0 + while i < len(relations): + self.assert_eq(relations[i].toPandas(), datasets[i].toPandas()) + i += 1 + + def test_observe(self): + # SPARK-41527: test DataFrame.observe() + observation_name = "my_metric" + + self.assert_eq( + self.connect.read.table(self.tbl_name) + .filter("id > 3") + .observe(observation_name, CF.min("id"), CF.max("id"), CF.sum("id")) + .toPandas(), + self.spark.read.table(self.tbl_name) + .filter("id > 3") + .observe(observation_name, SF.min("id"), SF.max("id"), SF.sum("id")) + .toPandas(), + ) + + from pyspark.sql.observation import Observation + + observation = Observation(observation_name) + + cdf = ( + self.connect.read.table(self.tbl_name) + .filter("id > 3") + .observe(observation, CF.min("id"), CF.max("id"), CF.sum("id")) + .toPandas() + ) + df = ( + self.spark.read.table(self.tbl_name) + .filter("id > 3") + .observe(observation, SF.min("id"), SF.max("id"), SF.sum("id")) + .toPandas() + ) + + self.assert_eq(cdf, df) + + observed_metrics = cdf.attrs["observed_metrics"] + self.assert_eq(len(observed_metrics), 1) + self.assert_eq(observed_metrics[0].name, observation_name) + self.assert_eq(len(observed_metrics[0].metrics), 3) + for metric in observed_metrics[0].metrics: + self.assertIsInstance(metric, ProtoExpression.Literal) + values = list(map(lambda metric: metric.long, observed_metrics[0].metrics)) + self.assert_eq(values, [4, 99, 4944]) + + with self.assertRaises(PySparkValueError) as pe: + self.connect.read.table(self.tbl_name).observe(observation_name) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "exprs"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name).observe(observation_name, CF.lit(1), "id") + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OF_COLUMN", + message_parameters={"arg_name": "exprs"}, + ) + + def test_with_columns(self): + # SPARK-41256: test withColumn(s). + self.assert_eq( + self.connect.read.table(self.tbl_name).withColumn("id", CF.lit(False)).toPandas(), + self.spark.read.table(self.tbl_name).withColumn("id", SF.lit(False)).toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name) + .withColumns({"id": CF.lit(False), "col_not_exist": CF.lit(False)}) + .toPandas(), + self.spark.read.table(self.tbl_name) + .withColumns( + { + "id": SF.lit(False), + "col_not_exist": SF.lit(False), + } + ) + .toPandas(), + ) + + def test_hint(self): + # SPARK-41349: Test hint + self.assert_eq( + self.connect.read.table(self.tbl_name).hint("COALESCE", 3000).toPandas(), + self.spark.read.table(self.tbl_name).hint("COALESCE", 3000).toPandas(), + ) + + # Hint with unsupported name will be ignored + self.assert_eq( + self.connect.read.table(self.tbl_name).hint("illegal").toPandas(), + self.spark.read.table(self.tbl_name).hint("illegal").toPandas(), + ) + + # Hint with all supported parameter values + such_a_nice_list = ["itworks1", "itworks2", "itworks3"] + self.assert_eq( + self.connect.read.table(self.tbl_name).hint("my awesome hint", 1.2345, 2).toPandas(), + self.spark.read.table(self.tbl_name).hint("my awesome hint", 1.2345, 2).toPandas(), + ) + + # Hint with unsupported parameter values + with self.assertRaises(AnalysisException): + self.connect.read.table(self.tbl_name).hint("REPARTITION", "id+1").toPandas() + + # Hint with unsupported parameter types + with self.assertRaises(TypeError): + self.connect.read.table(self.tbl_name).hint("REPARTITION", range(5)).toPandas() + + # Hint with unsupported parameter types + with self.assertRaises(TypeError): + self.connect.read.table(self.tbl_name).hint( + "my awesome hint", 1.2345, 2, such_a_nice_list, range(6) + ).toPandas() + + # Hint with wrong combination + with self.assertRaises(AnalysisException): + self.connect.read.table(self.tbl_name).hint("REPARTITION", "id", 3).toPandas() + + def test_join_hint(self): + cdf1 = self.connect.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + cdf2 = self.connect.createDataFrame( + [Row(height=80, name="Tom"), Row(height=85, name="Bob")] + ) + + self.assertTrue( + "BroadcastHashJoin" in cdf1.join(cdf2.hint("BROADCAST"), "name")._explain_string() + ) + self.assertTrue("SortMergeJoin" in cdf1.join(cdf2.hint("MERGE"), "name")._explain_string()) + self.assertTrue( + "ShuffledHashJoin" in cdf1.join(cdf2.hint("SHUFFLE_HASH"), "name")._explain_string() + ) + + def test_different_spark_session_join_or_union(self): + df = self.connect.range(10).limit(3) + + spark2 = RemoteSparkSession(connection="sc://localhost") + df2 = spark2.range(10).limit(3) + + with self.assertRaises(SessionNotSameException) as e1: + df.union(df2).collect() + self.check_error( + exception=e1.exception, + error_class="SESSION_NOT_SAME", + message_parameters={}, + ) + + with self.assertRaises(SessionNotSameException) as e2: + df.unionByName(df2).collect() + self.check_error( + exception=e2.exception, + error_class="SESSION_NOT_SAME", + message_parameters={}, + ) + + with self.assertRaises(SessionNotSameException) as e3: + df.join(df2).collect() + self.check_error( + exception=e3.exception, + error_class="SESSION_NOT_SAME", + message_parameters={}, + ) + + def test_extended_hint_types(self): + cdf = self.connect.range(100).toDF("id") + + cdf.hint( + "my awesome hint", + 1.2345, + "what", + ["itworks1", "itworks2", "itworks3"], + ).show() + + with self.assertRaises(PySparkTypeError) as pe: + cdf.hint( + "my awesome hint", + 1.2345, + "what", + {"itworks1": "itworks2"}, + ).show() + + self.check_error( + exception=pe.exception, + error_class="INVALID_ITEM_FOR_CONTAINER", + message_parameters={ + "arg_name": "parameters", + "allowed_types": "str, list, float, int", + "item_type": "dict", + }, + ) + + def test_empty_dataset(self): + # SPARK-41005: Test arrow based collection with empty dataset. + self.assertTrue( + self.connect.sql("SELECT 1 AS X LIMIT 0") + .toPandas() + .equals(self.spark.sql("SELECT 1 AS X LIMIT 0").toPandas()) + ) + pdf = self.connect.sql("SELECT 1 AS X LIMIT 0").toPandas() + self.assertEqual(0, len(pdf)) # empty dataset + self.assertEqual(1, len(pdf.columns)) # one column + self.assertEqual("X", pdf.columns[0]) + + def test_is_empty(self): + # SPARK-41212: Test is empty + self.assertFalse(self.connect.sql("SELECT 1 AS X").isEmpty()) + self.assertTrue(self.connect.sql("SELECT 1 AS X LIMIT 0").isEmpty()) + + def test_session(self): + self.assertEqual(self.connect, self.connect.sql("SELECT 1").sparkSession) + + def test_show(self): + # SPARK-41111: Test the show method + show_str = self.connect.sql("SELECT 1 AS X, 2 AS Y")._show_string() + # +---+---+ + # | X| Y| + # +---+---+ + # | 1| 2| + # +---+---+ + expected = "+---+---+\n| X| Y|\n+---+---+\n| 1| 2|\n+---+---+\n" + self.assertEqual(show_str, expected) + + def test_describe(self): + # SPARK-41403: Test the describe method + self.assert_eq( + self.connect.read.table(self.tbl_name).describe("id").toPandas(), + self.spark.read.table(self.tbl_name).describe("id").toPandas(), + ) + self.assert_eq( + self.connect.read.table(self.tbl_name).describe("id", "name").toPandas(), + self.spark.read.table(self.tbl_name).describe("id", "name").toPandas(), + ) + self.assert_eq( + self.connect.read.table(self.tbl_name).describe(["id", "name"]).toPandas(), + self.spark.read.table(self.tbl_name).describe(["id", "name"]).toPandas(), + ) + + def test_stat_cov(self): + # SPARK-41067: Test the stat.cov method + self.assertEqual( + self.connect.read.table(self.tbl_name2).stat.cov("col1", "col3"), + self.spark.read.table(self.tbl_name2).stat.cov("col1", "col3"), + ) + + def test_stat_corr(self): + # SPARK-41068: Test the stat.corr method + self.assertEqual( + self.connect.read.table(self.tbl_name2).stat.corr("col1", "col3"), + self.spark.read.table(self.tbl_name2).stat.corr("col1", "col3"), + ) + + self.assertEqual( + self.connect.read.table(self.tbl_name2).stat.corr("col1", "col3", "pearson"), + self.spark.read.table(self.tbl_name2).stat.corr("col1", "col3", "pearson"), + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.corr(1, "col3", "pearson") + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={ + "arg_name": "col1", + "arg_type": "int", + }, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name).stat.corr("col1", 1, "pearson") + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={ + "arg_name": "col2", + "arg_type": "int", + }, + ) + with self.assertRaises(ValueError) as context: + self.connect.read.table(self.tbl_name2).stat.corr("col1", "col3", "spearman"), + self.assertTrue( + "Currently only the calculation of the Pearson Correlation " + + "coefficient is supported." + in str(context.exception) + ) + + def test_stat_approx_quantile(self): + # SPARK-41069: Test the stat.approxQuantile method + result = self.connect.read.table(self.tbl_name2).stat.approxQuantile( + ["col1", "col3"], [0.1, 0.5, 0.9], 0.1 + ) + self.assertEqual(len(result), 2) + self.assertEqual(len(result[0]), 3) + self.assertEqual(len(result[1]), 3) + + result = self.connect.read.table(self.tbl_name2).stat.approxQuantile( + ["col1"], [0.1, 0.5, 0.9], 0.1 + ) + self.assertEqual(len(result), 1) + self.assertEqual(len(result[0]), 3) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.approxQuantile(1, [0.1, 0.5, 0.9], 0.1) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={ + "arg_name": "col", + "arg_type": "int", + }, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.approxQuantile(["col1", "col3"], 0.1, 0.1) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_TUPLE", + message_parameters={ + "arg_name": "probabilities", + "arg_type": "float", + }, + ) + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.approxQuantile( + ["col1", "col3"], [-0.1], 0.1 + ) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OF_FLOAT_OR_INT", + message_parameters={"arg_name": "probabilities", "arg_type": "float"}, + ) + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.approxQuantile( + ["col1", "col3"], [0.1, 0.5, 0.9], "str" + ) + + self.check_error( + exception=pe.exception, + error_class="NOT_FLOAT_OR_INT", + message_parameters={ + "arg_name": "relativeError", + "arg_type": "str", + }, + ) + with self.assertRaises(PySparkValueError) as pe: + self.connect.read.table(self.tbl_name2).stat.approxQuantile( + ["col1", "col3"], [0.1, 0.5, 0.9], -0.1 + ) + + self.check_error( + exception=pe.exception, + error_class="NEGATIVE_VALUE", + message_parameters={ + "arg_name": "relativeError", + "arg_value": "-0.1", + }, + ) + + def test_stat_freq_items(self): + # SPARK-41065: Test the stat.freqItems method + self.assert_eq( + self.connect.read.table(self.tbl_name2).stat.freqItems(["col1", "col3"]).toPandas(), + self.spark.read.table(self.tbl_name2).stat.freqItems(["col1", "col3"]).toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name2) + .stat.freqItems(["col1", "col3"], 0.4) + .toPandas(), + self.spark.read.table(self.tbl_name2).stat.freqItems(["col1", "col3"], 0.4).toPandas(), + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.connect.read.table(self.tbl_name2).stat.freqItems("col1") + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_TUPLE", + message_parameters={ + "arg_name": "cols", + "arg_type": "str", + }, + ) + + def test_stat_sample_by(self): + # SPARK-41069: Test stat.sample_by + + cdf = self.connect.range(0, 100).select((CF.col("id") % 3).alias("key")) + sdf = self.spark.range(0, 100).select((SF.col("id") % 3).alias("key")) + + self.assert_eq( + cdf.sampleBy(cdf.key, fractions={0: 0.1, 1: 0.2}, seed=0) + .groupBy("key") + .agg(CF.count(CF.lit(1))) + .orderBy("key") + .toPandas(), + sdf.sampleBy(sdf.key, fractions={0: 0.1, 1: 0.2}, seed=0) + .groupBy("key") + .agg(SF.count(SF.lit(1))) + .orderBy("key") + .toPandas(), + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.stat.sampleBy(cdf.key, fractions={0: 0.1, None: 0.2}, seed=0) + + self.check_error( + exception=pe.exception, + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "fractions", + "arg_type": "dict", + "allowed_types": "float, int, str", + "return_type": "NoneType", + }, + ) + + with self.assertRaises(SparkConnectException): + cdf.sampleBy(cdf.key, fractions={0: 0.1, 1: 1.2}, seed=0).show() + + def test_repr(self): + # SPARK-41213: Test the __repr__ method + query = """SELECT * FROM VALUES (1L, NULL), (3L, "Z") AS tab(a, b)""" + self.assertEqual( + self.connect.sql(query).__repr__(), + self.spark.sql(query).__repr__(), + ) + + def test_explain_string(self): + # SPARK-41122: test explain API. + plan_str = self.connect.sql("SELECT 1")._explain_string(extended=True) + self.assertTrue("Parsed Logical Plan" in plan_str) + self.assertTrue("Analyzed Logical Plan" in plan_str) + self.assertTrue("Optimized Logical Plan" in plan_str) + self.assertTrue("Physical Plan" in plan_str) + + with self.assertRaises(PySparkValueError) as pe: + self.connect.sql("SELECT 1")._explain_string(mode="unknown") + self.check_error( + exception=pe.exception, + error_class="UNKNOWN_EXPLAIN_MODE", + message_parameters={"explain_mode": "unknown"}, + ) + + def test_simple_datasource_read(self) -> None: + writeDf = self.df_text + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + writeDf.write.text(tmpPath) + + for schema in [ + "id STRING", + StructType([StructField("id", StringType())]), + ]: + readDf = self.connect.read.format("text").schema(schema).load(path=tmpPath) + expectResult = writeDf.collect() + pandasResult = readDf.toPandas() + if pandasResult is None: + self.assertTrue(False, "Empty pandas dataframe") + else: + actualResult = pandasResult.values.tolist() + self.assertEqual(len(expectResult), len(actualResult)) + + def test_simple_read_without_schema(self) -> None: + """SPARK-41300: Schema not set when reading CSV.""" + writeDf = self.df_text + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + writeDf.write.csv(tmpPath, header=True) + + readDf = self.connect.read.format("csv").option("header", True).load(path=tmpPath) + expectResult = set(writeDf.collect()) + pandasResult = set(readDf.collect()) + self.assertEqual(expectResult, pandasResult) + + def test_count(self) -> None: + # SPARK-41308: test count() API. + self.assertEqual( + self.connect.read.table(self.tbl_name).count(), + self.spark.read.table(self.tbl_name).count(), + ) + + def test_simple_transform(self) -> None: + """SPARK-41203: Support DF.transform""" + + def transform_df(input_df: CDataFrame) -> CDataFrame: + return input_df.select((CF.col("id") + CF.lit(10)).alias("id")) + + df = self.connect.range(1, 100) + result_left = df.transform(transform_df).collect() + result_right = self.connect.range(11, 110).collect() + self.assertEqual(result_right, result_left) + + # Check assertion. + with self.assertRaises(AssertionError): + df.transform(lambda x: 2) # type: ignore + + def test_alias(self) -> None: + """Testing supported and unsupported alias""" + col0 = ( + self.connect.range(1, 10) + .select(CF.col("id").alias("name", metadata={"max": 99})) + .schema.names[0] + ) + self.assertEqual("name", col0) + + with self.assertRaises(SparkConnectException) as exc: + self.connect.range(1, 10).select(CF.col("id").alias("this", "is", "not")).collect() + self.assertIn("(this, is, not)", str(exc.exception)) + + def test_column_regexp(self) -> None: + # SPARK-41438: test dataframe.colRegex() + ndf = self.connect.read.table(self.tbl_name3) + df = self.spark.read.table(self.tbl_name3) + + self.assert_eq( + ndf.select(ndf.colRegex("`tes.*\n.*mn`")).toPandas(), + df.select(df.colRegex("`tes.*\n.*mn`")).toPandas(), + ) + + def test_repartition(self) -> None: + # SPARK-41354: test dataframe.repartition(numPartitions) + self.assert_eq( + self.connect.read.table(self.tbl_name).repartition(10).toPandas(), + self.spark.read.table(self.tbl_name).repartition(10).toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name).coalesce(10).toPandas(), + self.spark.read.table(self.tbl_name).coalesce(10).toPandas(), + ) + + def test_repartition_by_expression(self) -> None: + # SPARK-41354: test dataframe.repartition(expressions) + self.assert_eq( + self.connect.read.table(self.tbl_name).repartition(10, "id").toPandas(), + self.spark.read.table(self.tbl_name).repartition(10, "id").toPandas(), + ) + + self.assert_eq( + self.connect.read.table(self.tbl_name).repartition("id").toPandas(), + self.spark.read.table(self.tbl_name).repartition("id").toPandas(), + ) + + # repartition with unsupported parameter values + with self.assertRaises(AnalysisException): + self.connect.read.table(self.tbl_name).repartition("id+1").toPandas() + + def test_repartition_by_range(self) -> None: + # SPARK-41354: test dataframe.repartitionByRange(expressions) + cdf = self.connect.read.table(self.tbl_name) + sdf = self.spark.read.table(self.tbl_name) + + self.assert_eq( + cdf.repartitionByRange(10, "id").toPandas(), + sdf.repartitionByRange(10, "id").toPandas(), + ) + + self.assert_eq( + cdf.repartitionByRange("id").toPandas(), + sdf.repartitionByRange("id").toPandas(), + ) + + self.assert_eq( + cdf.repartitionByRange(cdf.id.desc()).toPandas(), + sdf.repartitionByRange(sdf.id.desc()).toPandas(), + ) + + # repartitionByRange with unsupported parameter values + with self.assertRaises(AnalysisException): + self.connect.read.table(self.tbl_name).repartitionByRange("id+1").toPandas() + + def test_agg_with_two_agg_exprs(self) -> None: + # SPARK-41230: test dataframe.agg() + self.assert_eq( + self.connect.read.table(self.tbl_name).agg({"name": "min", "id": "max"}).toPandas(), + self.spark.read.table(self.tbl_name).agg({"name": "min", "id": "max"}).toPandas(), + ) + + def test_subtract(self): + # SPARK-41453: test dataframe.subtract() + ndf1 = self.connect.read.table(self.tbl_name) + ndf2 = ndf1.filter("id > 3") + df1 = self.spark.read.table(self.tbl_name) + df2 = df1.filter("id > 3") + + self.assert_eq( + ndf1.subtract(ndf2).toPandas(), + df1.subtract(df2).toPandas(), + ) + + def test_write_operations(self): + with tempfile.TemporaryDirectory() as d: + df = self.connect.range(50) + df.write.mode("overwrite").format("csv").save(d) + + ndf = self.connect.read.schema("id int").load(d, format="csv") + self.assertEqual(50, len(ndf.collect())) + cd = ndf.collect() + self.assertEqual(set(df.collect()), set(cd)) + + with tempfile.TemporaryDirectory() as d: + df = self.connect.range(50) + df.write.mode("overwrite").csv(d, lineSep="|") + + ndf = self.connect.read.schema("id int").load(d, format="csv", lineSep="|") + self.assertEqual(set(df.collect()), set(ndf.collect())) + + df = self.connect.range(50) + df.write.format("parquet").saveAsTable("parquet_test") + + ndf = self.connect.read.table("parquet_test") + self.assertEqual(set(df.collect()), set(ndf.collect())) + + def test_writeTo_operations(self): + # SPARK-42002: Implement DataFrameWriterV2 + import datetime + from pyspark.sql.connect.functions import col, years, months, days, hours, bucket + + df = self.connect.createDataFrame( + [(1, datetime.datetime(2000, 1, 1), "foo")], ("id", "ts", "value") + ) + writer = df.writeTo("table1") + self.assertIsInstance(writer.option("property", "value"), DataFrameWriterV2) + self.assertIsInstance(writer.options(property="value"), DataFrameWriterV2) + self.assertIsInstance(writer.using("source"), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(col("id")), DataFrameWriterV2) + self.assertIsInstance(writer.tableProperty("foo", "bar"), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(years("ts")), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(months("ts")), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(days("ts")), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(hours("ts")), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(bucket(11, "id")), DataFrameWriterV2) + self.assertIsInstance(writer.partitionedBy(bucket(3, "id"), hours("ts")), DataFrameWriterV2) + + def test_agg_with_avg(self): + # SPARK-41325: groupby.avg() + df = ( + self.connect.range(10) + .groupBy((CF.col("id") % CF.lit(2)).alias("moded")) + .avg("id") + .sort("moded") + ) + res = df.collect() + self.assertEqual(2, len(res)) + self.assertEqual(4.0, res[0][1]) + self.assertEqual(5.0, res[1][1]) + + # Additional GroupBy tests with 3 rows + + df_a = self.connect.range(10).groupBy((CF.col("id") % CF.lit(3)).alias("moded")) + df_b = self.spark.range(10).groupBy((SF.col("id") % SF.lit(3)).alias("moded")) + self.assertEqual( + set(df_b.agg(SF.sum("id")).collect()), set(df_a.agg(CF.sum("id")).collect()) + ) + + # Dict agg + measures = {"id": "sum"} + self.assertEqual( + set(df_a.agg(measures).select("sum(id)").collect()), + set(df_b.agg(measures).select("sum(id)").collect()), + ) + + def test_column_cannot_be_constructed_from_string(self): + with self.assertRaises(TypeError): + Column("col") + + def test_crossjoin(self): + # SPARK-41227: Test CrossJoin + connect_df = self.connect.read.table(self.tbl_name) + spark_df = self.spark.read.table(self.tbl_name) + self.assert_eq( + set( + connect_df.select("id") + .join(other=connect_df.select("name"), how="cross") + .toPandas() + ), + set(spark_df.select("id").join(other=spark_df.select("name"), how="cross").toPandas()), + ) + self.assert_eq( + set(connect_df.select("id").crossJoin(other=connect_df.select("name")).toPandas()), + set(spark_df.select("id").crossJoin(other=spark_df.select("name")).toPandas()), + ) + + def test_grouped_data(self): + query = """ + SELECT * FROM VALUES + ('James', 'Sales', 3000, 2020), + ('Michael', 'Sales', 4600, 2020), + ('Robert', 'Sales', 4100, 2020), + ('Maria', 'Finance', 3000, 2020), + ('James', 'Sales', 3000, 2019), + ('Scott', 'Finance', 3300, 2020), + ('Jen', 'Finance', 3900, 2020), + ('Jeff', 'Marketing', 3000, 2020), + ('Kumar', 'Marketing', 2000, 2020), + ('Saif', 'Sales', 4100, 2020) + AS T(name, department, salary, year) + """ + + # +-------+----------+------+----+ + # | name|department|salary|year| + # +-------+----------+------+----+ + # | James| Sales| 3000|2020| + # |Michael| Sales| 4600|2020| + # | Robert| Sales| 4100|2020| + # | Maria| Finance| 3000|2020| + # | James| Sales| 3000|2019| + # | Scott| Finance| 3300|2020| + # | Jen| Finance| 3900|2020| + # | Jeff| Marketing| 3000|2020| + # | Kumar| Marketing| 2000|2020| + # | Saif| Sales| 4100|2020| + # +-------+----------+------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test groupby + self.assert_eq( + cdf.groupBy("name").agg(CF.sum(cdf.salary)).toPandas(), + sdf.groupBy("name").agg(SF.sum(sdf.salary)).toPandas(), + ) + self.assert_eq( + cdf.groupBy("name", cdf.department).agg(CF.max("year"), CF.min(cdf.salary)).toPandas(), + sdf.groupBy("name", sdf.department).agg(SF.max("year"), SF.min(sdf.salary)).toPandas(), + ) + + # test rollup + self.assert_eq( + cdf.rollup("name").agg(CF.sum(cdf.salary)).toPandas(), + sdf.rollup("name").agg(SF.sum(sdf.salary)).toPandas(), + ) + self.assert_eq( + cdf.rollup("name", cdf.department).agg(CF.max("year"), CF.min(cdf.salary)).toPandas(), + sdf.rollup("name", sdf.department).agg(SF.max("year"), SF.min(sdf.salary)).toPandas(), + ) + + # test cube + self.assert_eq( + cdf.cube("name").agg(CF.sum(cdf.salary)).toPandas(), + sdf.cube("name").agg(SF.sum(sdf.salary)).toPandas(), + ) + self.assert_eq( + cdf.cube("name", cdf.department).agg(CF.max("year"), CF.min(cdf.salary)).toPandas(), + sdf.cube("name", sdf.department).agg(SF.max("year"), SF.min(sdf.salary)).toPandas(), + ) + + # test pivot + # pivot with values + self.assert_eq( + cdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .agg(CF.sum(cdf.salary)) + .toPandas(), + sdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .agg(SF.sum(sdf.salary)) + .toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name) + .pivot("department", ["Sales", "Finance", "Marketing"]) + .agg(CF.sum(cdf.salary)) + .toPandas(), + sdf.groupBy(sdf.name) + .pivot("department", ["Sales", "Finance", "Marketing"]) + .agg(SF.sum(sdf.salary)) + .toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .agg(CF.sum(cdf.salary)) + .toPandas(), + sdf.groupBy(sdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .agg(SF.sum(sdf.salary)) + .toPandas(), + ) + + # pivot without values + self.assert_eq( + cdf.groupBy("name").pivot("department").agg(CF.sum(cdf.salary)).toPandas(), + sdf.groupBy("name").pivot("department").agg(SF.sum(sdf.salary)).toPandas(), + ) + + self.assert_eq( + cdf.groupBy("name").pivot("year").agg(CF.sum(cdf.salary)).toPandas(), + sdf.groupBy("name").pivot("year").agg(SF.sum(sdf.salary)).toPandas(), + ) + + # check error + with self.assertRaisesRegex( + Exception, + "PIVOT after ROLLUP is not supported", + ): + cdf.rollup("name").pivot("department").agg(CF.sum(cdf.salary)) + + with self.assertRaisesRegex( + Exception, + "PIVOT after CUBE is not supported", + ): + cdf.cube("name").pivot("department").agg(CF.sum(cdf.salary)) + + with self.assertRaisesRegex( + Exception, + "Repeated PIVOT operation is not supported", + ): + cdf.groupBy("name").pivot("year").pivot("year").agg(CF.sum(cdf.salary)) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.groupBy("name").pivot("department", ["Sales", b"Marketing"]).agg(CF.sum(cdf.salary)) + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL_OR_FLOAT_OR_INT_OR_STR", + message_parameters={ + "arg_name": "value", + "arg_type": "bytes", + }, + ) + + def test_numeric_aggregation(self): + # SPARK-41737: test numeric aggregation + query = """ + SELECT * FROM VALUES + ('James', 'Sales', 3000, 2020), + ('Michael', 'Sales', 4600, 2020), + ('Robert', 'Sales', 4100, 2020), + ('Maria', 'Finance', 3000, 2020), + ('James', 'Sales', 3000, 2019), + ('Scott', 'Finance', 3300, 2020), + ('Jen', 'Finance', 3900, 2020), + ('Jeff', 'Marketing', 3000, 2020), + ('Kumar', 'Marketing', 2000, 2020), + ('Saif', 'Sales', 4100, 2020) + AS T(name, department, salary, year) + """ + + # +-------+----------+------+----+ + # | name|department|salary|year| + # +-------+----------+------+----+ + # | James| Sales| 3000|2020| + # |Michael| Sales| 4600|2020| + # | Robert| Sales| 4100|2020| + # | Maria| Finance| 3000|2020| + # | James| Sales| 3000|2019| + # | Scott| Finance| 3300|2020| + # | Jen| Finance| 3900|2020| + # | Jeff| Marketing| 3000|2020| + # | Kumar| Marketing| 2000|2020| + # | Saif| Sales| 4100|2020| + # +-------+----------+------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test groupby + self.assert_eq( + cdf.groupBy("name").min().toPandas(), + sdf.groupBy("name").min().toPandas(), + ) + self.assert_eq( + cdf.groupBy("name").min("salary").toPandas(), + sdf.groupBy("name").min("salary").toPandas(), + ) + self.assert_eq( + cdf.groupBy("name").max("salary").toPandas(), + sdf.groupBy("name").max("salary").toPandas(), + ) + self.assert_eq( + cdf.groupBy("name", cdf.department).avg("salary", "year").toPandas(), + sdf.groupBy("name", sdf.department).avg("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.groupBy("name", cdf.department).mean("salary", "year").toPandas(), + sdf.groupBy("name", sdf.department).mean("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.groupBy("name", cdf.department).sum("salary", "year").toPandas(), + sdf.groupBy("name", sdf.department).sum("salary", "year").toPandas(), + ) + + # test rollup + self.assert_eq( + cdf.rollup("name").max().toPandas(), + sdf.rollup("name").max().toPandas(), + ) + self.assert_eq( + cdf.rollup("name").min("salary").toPandas(), + sdf.rollup("name").min("salary").toPandas(), + ) + self.assert_eq( + cdf.rollup("name").max("salary").toPandas(), + sdf.rollup("name").max("salary").toPandas(), + ) + self.assert_eq( + cdf.rollup("name", cdf.department).avg("salary", "year").toPandas(), + sdf.rollup("name", sdf.department).avg("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.rollup("name", cdf.department).mean("salary", "year").toPandas(), + sdf.rollup("name", sdf.department).mean("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.rollup("name", cdf.department).sum("salary", "year").toPandas(), + sdf.rollup("name", sdf.department).sum("salary", "year").toPandas(), + ) + + # test cube + self.assert_eq( + cdf.cube("name").avg().toPandas(), + sdf.cube("name").avg().toPandas(), + ) + self.assert_eq( + cdf.cube("name").mean().toPandas(), + sdf.cube("name").mean().toPandas(), + ) + self.assert_eq( + cdf.cube("name").min("salary").toPandas(), + sdf.cube("name").min("salary").toPandas(), + ) + self.assert_eq( + cdf.cube("name").max("salary").toPandas(), + sdf.cube("name").max("salary").toPandas(), + ) + self.assert_eq( + cdf.cube("name", cdf.department).avg("salary", "year").toPandas(), + sdf.cube("name", sdf.department).avg("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.cube("name", cdf.department).sum("salary", "year").toPandas(), + sdf.cube("name", sdf.department).sum("salary", "year").toPandas(), + ) + + # test pivot + # pivot with values + self.assert_eq( + cdf.groupBy("name").pivot("department", ["Sales", "Marketing"]).sum().toPandas(), + sdf.groupBy("name").pivot("department", ["Sales", "Marketing"]).sum().toPandas(), + ) + self.assert_eq( + cdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .min("salary") + .toPandas(), + sdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .min("salary") + .toPandas(), + ) + self.assert_eq( + cdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .max("salary") + .toPandas(), + sdf.groupBy("name") + .pivot("department", ["Sales", "Marketing"]) + .max("salary") + .toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .avg("salary", "year") + .toPandas(), + sdf.groupBy(sdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .avg("salary", "year") + .toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .sum("salary", "year") + .toPandas(), + sdf.groupBy(sdf.name) + .pivot("department", ["Sales", "Finance", "Unknown"]) + .sum("salary", "year") + .toPandas(), + ) + + # pivot without values + self.assert_eq( + cdf.groupBy("name").pivot("department").min().toPandas(), + sdf.groupBy("name").pivot("department").min().toPandas(), + ) + self.assert_eq( + cdf.groupBy("name").pivot("department").min("salary").toPandas(), + sdf.groupBy("name").pivot("department").min("salary").toPandas(), + ) + self.assert_eq( + cdf.groupBy("name").pivot("department").max("salary").toPandas(), + sdf.groupBy("name").pivot("department").max("salary").toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name).pivot("department").avg("salary", "year").toPandas(), + sdf.groupBy(sdf.name).pivot("department").avg("salary", "year").toPandas(), + ) + self.assert_eq( + cdf.groupBy(cdf.name).pivot("department").sum("salary", "year").toPandas(), + sdf.groupBy(sdf.name).pivot("department").sum("salary", "year").toPandas(), + ) + + # check error + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.groupBy("name").min("department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.groupBy("name").max("salary", "department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.rollup("name").avg("department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.rollup("name").sum("salary", "department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.cube("name").min("department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.cube("name").max("salary", "department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.groupBy("name").pivot("department").avg("department").show() + + with self.assertRaisesRegex( + TypeError, + "Numeric aggregation function can only be applied on numeric columns", + ): + cdf.groupBy("name").pivot("department").sum("salary", "department").show() + + def test_with_metadata(self): + cdf = self.connect.createDataFrame(data=[(2, "Alice"), (5, "Bob")], schema=["age", "name"]) + self.assertEqual(cdf.schema["age"].metadata, {}) + self.assertEqual(cdf.schema["name"].metadata, {}) + + cdf1 = cdf.withMetadata(columnName="age", metadata={"max_age": 5}) + self.assertEqual(cdf1.schema["age"].metadata, {"max_age": 5}) + + cdf2 = cdf.withMetadata(columnName="name", metadata={"names": ["Alice", "Bob"]}) + self.assertEqual(cdf2.schema["name"].metadata, {"names": ["Alice", "Bob"]}) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.withMetadata(columnName="name", metadata=["magic"]) + + self.check_error( + exception=pe.exception, + error_class="NOT_DICT", + message_parameters={ + "arg_name": "metadata", + "arg_type": "list", + }, + ) + + def test_collect_nested_type(self): + query = """ + SELECT * FROM VALUES + (1, 4, 0, 8, true, true, ARRAY(1, NULL, 3), MAP(1, 2, 3, 4)), + (2, 5, -1, NULL, false, NULL, ARRAY(1, 3), MAP(1, NULL, 3, 4)), + (3, 6, NULL, 0, false, NULL, ARRAY(NULL), NULL) + AS tab(a, b, c, d, e, f, g, h) + """ + + # +---+---+----+----+-----+----+------------+-------------------+ + # | a| b| c| d| e| f| g| h| + # +---+---+----+----+-----+----+------------+-------------------+ + # | 1| 4| 0| 8| true|true|[1, null, 3]| {1 -> 2, 3 -> 4}| + # | 2| 5| -1|NULL|false|NULL| [1, 3]|{1 -> null, 3 -> 4}| + # | 3| 6|NULL| 0|false|NULL| [null]| NULL| + # +---+---+----+----+-----+----+------------+-------------------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test collect array + # +--------------+-------------+------------+ + # |array(a, b, c)| array(e, f)| g| + # +--------------+-------------+------------+ + # | [1, 4, 0]| [true, true]|[1, null, 3]| + # | [2, 5, -1]|[false, null]| [1, 3]| + # | [3, 6, null]|[false, null]| [null]| + # +--------------+-------------+------------+ + self.assertEqual( + cdf.select(CF.array("a", "b", "c"), CF.array("e", "f"), CF.col("g")).collect(), + sdf.select(SF.array("a", "b", "c"), SF.array("e", "f"), SF.col("g")).collect(), + ) + + # test collect nested array + # +-----------------------------------+-------------------------+ + # |array(array(a), array(b), array(c))|array(array(e), array(f))| + # +-----------------------------------+-------------------------+ + # | [[1], [4], [0]]| [[true], [true]]| + # | [[2], [5], [-1]]| [[false], [null]]| + # | [[3], [6], [null]]| [[false], [null]]| + # +-----------------------------------+-------------------------+ + self.assertEqual( + cdf.select( + CF.array(CF.array("a"), CF.array("b"), CF.array("c")), + CF.array(CF.array("e"), CF.array("f")), + ).collect(), + sdf.select( + SF.array(SF.array("a"), SF.array("b"), SF.array("c")), + SF.array(SF.array("e"), SF.array("f")), + ).collect(), + ) + + # test collect array of struct, map + # +----------------+---------------------+ + # |array(struct(a))| array(h)| + # +----------------+---------------------+ + # | [{1}]| [{1 -> 2, 3 -> 4}]| + # | [{2}]|[{1 -> null, 3 -> 4}]| + # | [{3}]| [null]| + # +----------------+---------------------+ + self.assertEqual( + cdf.select(CF.array(CF.struct("a")), CF.array("h")).collect(), + sdf.select(SF.array(SF.struct("a")), SF.array("h")).collect(), + ) + + # test collect map + # +-------------------+-------------------+ + # | h| map(a, b, b, c)| + # +-------------------+-------------------+ + # | {1 -> 2, 3 -> 4}| {1 -> 4, 4 -> 0}| + # |{1 -> null, 3 -> 4}| {2 -> 5, 5 -> -1}| + # | NULL|{3 -> 6, 6 -> null}| + # +-------------------+-------------------+ + self.assertEqual( + cdf.select(CF.col("h"), CF.create_map("a", "b", "b", "c")).collect(), + sdf.select(SF.col("h"), SF.create_map("a", "b", "b", "c")).collect(), + ) + + # test collect map of struct, array + # +-------------------+------------------------+ + # | map(a, g)| map(a, struct(b, g))| + # +-------------------+------------------------+ + # |{1 -> [1, null, 3]}|{1 -> {4, [1, null, 3]}}| + # | {2 -> [1, 3]}| {2 -> {5, [1, 3]}}| + # | {3 -> [null]}| {3 -> {6, [null]}}| + # +-------------------+------------------------+ + self.assertEqual( + cdf.select(CF.create_map("a", "g"), CF.create_map("a", CF.struct("b", "g"))).collect(), + sdf.select(SF.create_map("a", "g"), SF.create_map("a", SF.struct("b", "g"))).collect(), + ) + + # test collect struct + # +------------------+--------------------------+ + # |struct(a, b, c, d)| struct(e, f, g)| + # +------------------+--------------------------+ + # | {1, 4, 0, 8}|{true, true, [1, null, 3]}| + # | {2, 5, -1, null}| {false, null, [1, 3]}| + # | {3, 6, null, 0}| {false, null, [null]}| + # +------------------+--------------------------+ + self.assertEqual( + cdf.select(CF.struct("a", "b", "c", "d"), CF.struct("e", "f", "g")).collect(), + sdf.select(SF.struct("a", "b", "c", "d"), SF.struct("e", "f", "g")).collect(), + ) + + # test collect nested struct + # +------------------------------------------+--------------------------+----------------------------+ # noqa + # |struct(a, struct(a, struct(c, struct(d))))|struct(a, b, struct(c, d))| struct(e, f, struct(g))| # noqa + # +------------------------------------------+--------------------------+----------------------------+ # noqa + # | {1, {1, {0, {8}}}}| {1, 4, {0, 8}}|{true, true, {[1, null, 3]}}| # noqa + # | {2, {2, {-1, {null}}}}| {2, 5, {-1, null}}| {false, null, {[1, 3]}}| # noqa + # | {3, {3, {null, {0}}}}| {3, 6, {null, 0}}| {false, null, {[null]}}| # noqa + # +------------------------------------------+--------------------------+----------------------------+ # noqa + self.assertEqual( + cdf.select( + CF.struct("a", CF.struct("a", CF.struct("c", CF.struct("d")))), + CF.struct("a", "b", CF.struct("c", "d")), + CF.struct("e", "f", CF.struct("g")), + ).collect(), + sdf.select( + SF.struct("a", SF.struct("a", SF.struct("c", SF.struct("d")))), + SF.struct("a", "b", SF.struct("c", "d")), + SF.struct("e", "f", SF.struct("g")), + ).collect(), + ) + + # test collect struct containing array, map + # +--------------------------------------------+ + # | struct(a, struct(a, struct(g, struct(h))))| + # +--------------------------------------------+ + # |{1, {1, {[1, null, 3], {{1 -> 2, 3 -> 4}}}}}| + # | {2, {2, {[1, 3], {{1 -> null, 3 -> 4}}}}}| + # | {3, {3, {[null], {null}}}}| + # +--------------------------------------------+ + self.assertEqual( + cdf.select( + CF.struct("a", CF.struct("a", CF.struct("g", CF.struct("h")))), + ).collect(), + sdf.select( + SF.struct("a", SF.struct("a", SF.struct("g", SF.struct("h")))), + ).collect(), + ) + + def test_simple_udt(self): + from pyspark.ml.linalg import MatrixUDT, VectorUDT + + for schema in [ + StructType().add("key", LongType()).add("val", PythonOnlyUDT()), + StructType().add("key", LongType()).add("val", ArrayType(PythonOnlyUDT())), + StructType().add("key", LongType()).add("val", MapType(LongType(), PythonOnlyUDT())), + StructType().add("key", LongType()).add("val", PythonOnlyUDT()), + StructType().add("key", LongType()).add("vec", VectorUDT()), + StructType().add("key", LongType()).add("mat", MatrixUDT()), + ]: + cdf = self.connect.createDataFrame(data=[], schema=schema) + sdf = self.spark.createDataFrame(data=[], schema=schema) + + self.assertEqual(cdf.schema, sdf.schema) + + def test_simple_udt_from_read(self): + from pyspark.ml.linalg import Matrices, Vectors + + with tempfile.TemporaryDirectory() as d: + path1 = f"{d}/df1.parquet" + self.spark.createDataFrame( + [(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)], + schema=StructType().add("key", LongType()).add("val", PythonOnlyUDT()), + ).write.parquet(path1) + + path2 = f"{d}/df2.parquet" + self.spark.createDataFrame( + [(i % 3, [PythonOnlyPoint(float(i), float(i))]) for i in range(10)], + schema=StructType().add("key", LongType()).add("val", ArrayType(PythonOnlyUDT())), + ).write.parquet(path2) + + path3 = f"{d}/df3.parquet" + self.spark.createDataFrame( + [(i % 3, {i % 3: PythonOnlyPoint(float(i + 1), float(i + 1))}) for i in range(10)], + schema=StructType() + .add("key", LongType()) + .add("val", MapType(LongType(), PythonOnlyUDT())), + ).write.parquet(path3) + + path4 = f"{d}/df4.parquet" + self.spark.createDataFrame( + [(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)], + schema=StructType().add("key", LongType()).add("val", PythonOnlyUDT()), + ).write.parquet(path4) + + path5 = f"{d}/df5.parquet" + self.spark.createDataFrame( + [Row(label=1.0, point=ExamplePoint(1.0, 2.0))] + ).write.parquet(path5) + + path6 = f"{d}/df6.parquet" + self.spark.createDataFrame( + [(Vectors.dense(1.0, 2.0, 3.0),), (Vectors.sparse(3, {1: 1.0, 2: 5.5}),)], + ["vec"], + ).write.parquet(path6) + + path7 = f"{d}/df7.parquet" + self.spark.createDataFrame( + [ + (Matrices.dense(3, 2, [0, 1, 4, 5, 9, 10]),), + (Matrices.sparse(1, 1, [0, 1], [0], [2.0]),), + ], + ["mat"], + ).write.parquet(path7) + + for path in [path1, path2, path3, path4, path5, path6, path7]: + self.assertEqual( + self.connect.read.parquet(path).schema, + self.spark.read.parquet(path).schema, + ) + + def test_version(self): + self.assertEqual( + self.connect.version, + self.spark.version, + ) + + def test_same_semantics(self): + plan = self.connect.sql("SELECT 1") + other = self.connect.sql("SELECT 1") + self.assertTrue(plan.sameSemantics(other)) + + def test_semantic_hash(self): + plan = self.connect.sql("SELECT 1") + other = self.connect.sql("SELECT 1") + self.assertEqual( + plan.semanticHash(), + other.semanticHash(), + ) + + def test_unsupported_functions(self): + # SPARK-41225: Disable unsupported functions. + df = self.connect.read.table(self.tbl_name) + for f in ( + "rdd", + "foreach", + "foreachPartition", + "checkpoint", + "localCheckpoint", + ): + with self.assertRaises(NotImplementedError): + getattr(df, f)() + + def test_unsupported_session_functions(self): + # SPARK-41934: Disable unsupported functions. + + with self.assertRaises(NotImplementedError): + RemoteSparkSession.builder.enableHiveSupport() + + for f in ( + "newSession", + "sparkContext", + ): + with self.assertRaises(NotImplementedError): + getattr(self.connect, f)() + + def test_sql_with_command(self): + # SPARK-42705: spark.sql should return values from the command. + self.assertEqual( + self.connect.sql("show functions").collect(), self.spark.sql("show functions").collect() + ) + + def test_schema_has_nullable(self): + schema_false = StructType().add("id", IntegerType(), False) + cdf1 = self.connect.createDataFrame([[1]], schema=schema_false) + sdf1 = self.spark.createDataFrame([[1]], schema=schema_false) + self.assertEqual(cdf1.schema, sdf1.schema) + self.assertEqual(cdf1.collect(), sdf1.collect()) + + schema_true = StructType().add("id", IntegerType(), True) + cdf2 = self.connect.createDataFrame([[1]], schema=schema_true) + sdf2 = self.spark.createDataFrame([[1]], schema=schema_true) + self.assertEqual(cdf2.schema, sdf2.schema) + self.assertEqual(cdf2.collect(), sdf2.collect()) + + pdf1 = cdf1.toPandas() + cdf3 = self.connect.createDataFrame(pdf1, cdf1.schema) + sdf3 = self.spark.createDataFrame(pdf1, sdf1.schema) + self.assertEqual(cdf3.schema, sdf3.schema) + self.assertEqual(cdf3.collect(), sdf3.collect()) + + pdf2 = cdf2.toPandas() + cdf4 = self.connect.createDataFrame(pdf2, cdf2.schema) + sdf4 = self.spark.createDataFrame(pdf2, sdf2.schema) + self.assertEqual(cdf4.schema, sdf4.schema) + self.assertEqual(cdf4.collect(), sdf4.collect()) + + def test_array_has_nullable(self): + for schemas, data in [ + ( + [StructType().add("arr", ArrayType(IntegerType(), False), True)], + [Row([1, 2]), Row([3]), Row(None)], + ), + ( + [ + StructType().add("arr", ArrayType(IntegerType(), True), True), + "arr array", + ], + [Row([1, None]), Row([3]), Row(None)], + ), + ( + [StructType().add("arr", ArrayType(IntegerType(), False), False)], + [Row([1, 2]), Row([3])], + ), + ( + [ + StructType().add("arr", ArrayType(IntegerType(), True), False), + "arr array not null", + ], + [Row([1, None]), Row([3])], + ), + ]: + for schema in schemas: + with self.subTest(schema=schema): + cdf = self.connect.createDataFrame(data, schema=schema) + sdf = self.spark.createDataFrame(data, schema=schema) + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_map_has_nullable(self): + for schemas, data in [ + ( + [StructType().add("map", MapType(StringType(), IntegerType(), False), True)], + [Row({"a": 1, "b": 2}), Row({"a": 3}), Row(None)], + ), + ( + [ + StructType().add("map", MapType(StringType(), IntegerType(), True), True), + "map map", + ], + [Row({"a": 1, "b": None}), Row({"a": 3}), Row(None)], + ), + ( + [StructType().add("map", MapType(StringType(), IntegerType(), False), False)], + [Row({"a": 1, "b": 2}), Row({"a": 3})], + ), + ( + [ + StructType().add("map", MapType(StringType(), IntegerType(), True), False), + "map map not null", + ], + [Row({"a": 1, "b": None}), Row({"a": 3})], + ), + ]: + for schema in schemas: + with self.subTest(schema=schema): + cdf = self.connect.createDataFrame(data, schema=schema) + sdf = self.spark.createDataFrame(data, schema=schema) + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_struct_has_nullable(self): + for schemas, data in [ + ( + [ + StructType().add("struct", StructType().add("i", IntegerType(), False), True), + "struct struct", + ], + [Row(Row(1)), Row(Row(2)), Row(None)], + ), + ( + [ + StructType().add("struct", StructType().add("i", IntegerType(), True), True), + "struct struct", + ], + [Row(Row(1)), Row(Row(2)), Row(Row(None)), Row(None)], + ), + ( + [ + StructType().add("struct", StructType().add("i", IntegerType(), False), False), + "struct struct not null", + ], + [Row(Row(1)), Row(Row(2))], + ), + ( + [ + StructType().add("struct", StructType().add("i", IntegerType(), True), False), + "struct struct not null", + ], + [Row(Row(1)), Row(Row(2)), Row(Row(None))], + ), + ]: + for schema in schemas: + with self.subTest(schema=schema): + cdf = self.connect.createDataFrame(data, schema=schema) + sdf = self.spark.createDataFrame(data, schema=schema) + self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_large_client_data(self): + # SPARK-42816 support more than 4MB message size. + # ~200bytes + cols = ["abcdefghijklmnoprstuvwxyz" for x in range(10)] + # 100k rows => 20MB + row_count = 100 * 1000 + rows = [cols] * row_count + self.assertEqual(row_count, self.connect.createDataFrame(data=rows).count()) + + def test_unsupported_jvm_attribute(self): + # Unsupported jvm attributes for Spark session. + unsupported_attrs = ["_jsc", "_jconf", "_jvm", "_jsparkSession"] + spark_session = self.connect + for attr in unsupported_attrs: + with self.assertRaises(PySparkAttributeError) as pe: + getattr(spark_session, attr) + + self.check_error( + exception=pe.exception, + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", + message_parameters={"attr_name": attr}, + ) + + # Unsupported jvm attributes for DataFrame. + unsupported_attrs = ["_jseq", "_jdf", "_jmap", "_jcols"] + cdf = self.connect.range(10) + for attr in unsupported_attrs: + with self.assertRaises(PySparkAttributeError) as pe: + getattr(cdf, attr) + + self.check_error( + exception=pe.exception, + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", + message_parameters={"attr_name": attr}, + ) + + # Unsupported jvm attributes for Column. + with self.assertRaises(PySparkAttributeError) as pe: + getattr(cdf.id, "_jc") + + self.check_error( + exception=pe.exception, + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", + message_parameters={"attr_name": "_jc"}, + ) + + # Unsupported jvm attributes for DataFrameReader. + with self.assertRaises(PySparkAttributeError) as pe: + getattr(spark_session.read, "_jreader") + + self.check_error( + exception=pe.exception, + error_class="JVM_ATTRIBUTE_NOT_SUPPORTED", + message_parameters={"attr_name": "_jreader"}, + ) + + def test_df_caache(self): + df = self.connect.range(10) + df.cache() + self.assert_eq(10, df.count()) + self.assertTrue(df.is_cached) + + +class SparkConnectSessionTests(ReusedConnectTestCase): + def setUp(self) -> None: + self.spark = ( + PySparkSession.builder.config(conf=self.conf()) + .appName(self.__class__.__name__) + .remote("local[4]") + .getOrCreate() + ) + + def tearDown(self): + self.spark.stop() + + def _check_no_active_session_error(self, e: PySparkException): + self.check_error(exception=e, error_class="NO_ACTIVE_SESSION", message_parameters=dict()) + + def test_stop_session(self): + df = self.spark.sql("select 1 as a, 2 as b") + catalog = self.spark.catalog + self.spark.stop() + + # _execute_and_fetch + with self.assertRaises(SparkConnectException) as e: + self.spark.sql("select 1") + self._check_no_active_session_error(e.exception) + + with self.assertRaises(SparkConnectException) as e: + catalog.tableExists("table") + self._check_no_active_session_error(e.exception) + + # _execute + with self.assertRaises(SparkConnectException) as e: + self.spark.udf.register("test_func", lambda x: x + 1) + self._check_no_active_session_error(e.exception) + + # _analyze + with self.assertRaises(SparkConnectException) as e: + df._explain_string(extended=True) + self._check_no_active_session_error(e.exception) + + # Config + with self.assertRaises(SparkConnectException) as e: + self.spark.conf.get("some.conf") + self._check_no_active_session_error(e.exception) + + def test_error_stack_trace(self): + with self.sql_conf({"spark.sql.pyspark.jvmStacktrace.enabled": True}): + with self.assertRaises(AnalysisException) as e: + self.spark.sql("select x").collect() + self.assertTrue("JVM stacktrace" in e.exception.message) + self.assertTrue( + "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in e.exception.message + ) + + with self.sql_conf({"spark.sql.pyspark.jvmStacktrace.enabled": False}): + with self.assertRaises(AnalysisException) as e: + self.spark.sql("select x").collect() + self.assertFalse("JVM stacktrace" in e.exception.message) + self.assertFalse( + "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in e.exception.message + ) + + # Create a new session with a different stack trace size. + self.spark.stop() + spark = ( + PySparkSession.builder.config(conf=self.conf()) + .config("spark.connect.jvmStacktrace.maxSize", 128) + .remote("local[4]") + .getOrCreate() + ) + spark.conf.set("spark.sql.pyspark.jvmStacktrace.enabled", "true") + with self.assertRaises(AnalysisException) as e: + spark.sql("select x").collect() + self.assertTrue("JVM stacktrace" in e.exception.message) + self.assertFalse( + "at org.apache.spark.sql.catalyst.analysis.CheckAnalysis" in e.exception.message + ) + spark.stop() + + def test_can_create_multiple_sessions_to_different_remotes(self): + self.spark.stop() + self.assertIsNotNone(self.spark._client) + # Creates a new remote session. + other = PySparkSession.builder.remote("sc://other.remote:114/").create() + self.assertNotEquals(self.spark, other) + + # Gets currently active session. + same = PySparkSession.builder.remote("sc://other.remote.host:114/").getOrCreate() + self.assertEquals(other, same) + same.stop() + + # Make sure the environment is clean. + self.spark.stop() + with self.assertRaises(RuntimeError) as e: + PySparkSession.builder.create() + self.assertIn("Create a new SparkSession is only supported with SparkConnect.", str(e)) + + +class SparkConnectSessionWithOptionsTest(unittest.TestCase): + def setUp(self) -> None: + self.spark = ( + PySparkSession.builder.config("string", "foo") + .config("integer", 1) + .config("boolean", False) + .appName(self.__class__.__name__) + .remote("local[4]") + .getOrCreate() + ) + + def tearDown(self): + self.spark.stop() + + def test_config(self): + # Config + self.assertEqual(self.spark.conf.get("string"), "foo") + self.assertEqual(self.spark.conf.get("boolean"), "false") + self.assertEqual(self.spark.conf.get("integer"), "1") + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class ClientTests(unittest.TestCase): + def test_retry_error_handling(self): + # Helper class for wrapping the test. + class TestError(grpc.RpcError, Exception): + def __init__(self, code: grpc.StatusCode): + self._code = code + + def code(self): + return self._code + + def stub(retries, w, code): + w["attempts"] += 1 + if w["attempts"] < retries: + w["raised"] += 1 + raise TestError(code) + + # Check that max_retries 1 is only one retry so two attempts. + call_wrap = defaultdict(int) + for attempt in Retrying( + can_retry=lambda x: True, + max_retries=1, + backoff_multiplier=1, + initial_backoff=1, + max_backoff=10, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(2, call_wrap, grpc.StatusCode.INTERNAL) + + self.assertEqual(2, call_wrap["attempts"]) + self.assertEqual(1, call_wrap["raised"]) + + # Check that if we have less than 4 retries all is ok. + call_wrap = defaultdict(int) + for attempt in Retrying( + can_retry=lambda x: True, + max_retries=4, + backoff_multiplier=1, + initial_backoff=1, + max_backoff=10, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(2, call_wrap, grpc.StatusCode.INTERNAL) + + self.assertTrue(call_wrap["attempts"] < 4) + self.assertEqual(call_wrap["raised"], 1) + + # Exceed the retries. + call_wrap = defaultdict(int) + with self.assertRaises(TestError): + for attempt in Retrying( + can_retry=lambda x: True, + max_retries=2, + max_backoff=50, + backoff_multiplier=1, + initial_backoff=50, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(5, call_wrap, grpc.StatusCode.INTERNAL) + + self.assertTrue(call_wrap["attempts"] < 5) + self.assertEqual(call_wrap["raised"], 3) + + # Check that only specific exceptions are retried. + # Check that if we have less than 4 retries all is ok. + call_wrap = defaultdict(int) + for attempt in Retrying( + can_retry=lambda x: x.code() == grpc.StatusCode.UNAVAILABLE, + max_retries=4, + backoff_multiplier=1, + initial_backoff=1, + max_backoff=10, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(2, call_wrap, grpc.StatusCode.UNAVAILABLE) + + self.assertTrue(call_wrap["attempts"] < 4) + self.assertEqual(call_wrap["raised"], 1) + + # Exceed the retries. + call_wrap = defaultdict(int) + with self.assertRaises(TestError): + for attempt in Retrying( + can_retry=lambda x: x.code() == grpc.StatusCode.UNAVAILABLE, + max_retries=2, + max_backoff=50, + backoff_multiplier=1, + initial_backoff=50, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(5, call_wrap, grpc.StatusCode.UNAVAILABLE) + + self.assertTrue(call_wrap["attempts"] < 4) + self.assertEqual(call_wrap["raised"], 3) + + # Test that another error is always thrown. + call_wrap = defaultdict(int) + with self.assertRaises(TestError): + for attempt in Retrying( + can_retry=lambda x: x.code() == grpc.StatusCode.UNAVAILABLE, + max_retries=4, + backoff_multiplier=1, + initial_backoff=1, + max_backoff=10, + jitter=0, + min_jitter_threshold=0, + ): + with attempt: + stub(5, call_wrap, grpc.StatusCode.INTERNAL) + + self.assertEqual(call_wrap["attempts"], 1) + self.assertEqual(call_wrap["raised"], 1) + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class ChannelBuilderTests(unittest.TestCase): + def test_invalid_connection_strings(self): + invalid = [ + "scc://host:12", + "http://host", + "sc:/host:1234/path", + "sc://host/path", + "sc://host/;parm1;param2", + ] + for i in invalid: + self.assertRaises(PySparkValueError, ChannelBuilder, i) + + def test_sensible_defaults(self): + chan = ChannelBuilder("sc://host") + self.assertFalse(chan.secure, "Default URL is not secure") + + chan = ChannelBuilder("sc://host/;token=abcs") + self.assertTrue(chan.secure, "specifying a token must set the channel to secure") + self.assertRegex( + chan.userAgent, r"^_SPARK_CONNECT_PYTHON spark/[^ ]+ os/[^ ]+ python/[^ ]+$" + ) + chan = ChannelBuilder("sc://host/;use_ssl=abcs") + self.assertFalse(chan.secure, "Garbage in, false out") + + def test_user_agent(self): + chan = ChannelBuilder("sc://host/;user_agent=Agent123%20%2F3.4") + self.assertIn("Agent123 /3.4", chan.userAgent) + + def test_user_agent_len(self): + user_agent = "x" * 2049 + chan = ChannelBuilder(f"sc://host/;user_agent={user_agent}") + with self.assertRaises(SparkConnectException) as err: + chan.userAgent + self.assertRegex(err.exception.message, "'user_agent' parameter should not exceed") + + user_agent = "%C3%A4" * 341 # "%C3%A4" -> "ä"; (341 * 6 = 2046) < 2048 + expected = "ä" * 341 + chan = ChannelBuilder(f"sc://host/;user_agent={user_agent}") + self.assertIn(expected, chan.userAgent) + + def test_valid_channel_creation(self): + chan = ChannelBuilder("sc://host").toChannel() + self.assertIsInstance(chan, grpc.Channel) + + # Sets up a channel without tokens because ssl is not used. + chan = ChannelBuilder("sc://host/;use_ssl=true;token=abc").toChannel() + self.assertIsInstance(chan, grpc.Channel) + + chan = ChannelBuilder("sc://host/;use_ssl=true").toChannel() + self.assertIsInstance(chan, grpc.Channel) + + def test_channel_properties(self): + chan = ChannelBuilder("sc://host/;use_ssl=true;token=abc;user_agent=foo;param1=120%2021") + self.assertEqual("host:15002", chan.endpoint) + self.assertIn("foo", chan.userAgent.split(" ")) + self.assertEqual(True, chan.secure) + self.assertEqual("120 21", chan.get("param1")) + + def test_metadata(self): + chan = ChannelBuilder("sc://host/;use_ssl=true;token=abc;param1=120%2021;x-my-header=abcd") + md = chan.metadata() + self.assertEqual([("param1", "120 21"), ("x-my-header", "abcd")], md) + + def test_metadata(self): + id = str(uuid.uuid4()) + chan = ChannelBuilder(f"sc://host/;session_id={id}") + self.assertEqual(id, chan.session_id) + + chan = ChannelBuilder(f"sc://host/;session_id={id};user_agent=acbd;token=abcd;use_ssl=true") + md = chan.metadata() + for kv in md: + self.assertNotIn( + kv[0], + [ + ChannelBuilder.PARAM_SESSION_ID, + ChannelBuilder.PARAM_TOKEN, + ChannelBuilder.PARAM_USER_ID, + ChannelBuilder.PARAM_USER_AGENT, + ChannelBuilder.PARAM_USE_SSL, + ], + "Metadata must not contain fixed params", + ) + + with self.assertRaises(ValueError) as ve: + chan = ChannelBuilder("sc://host/;session_id=abcd") + SparkConnectClient(chan) + self.assertIn( + "Parameter value 'session_id' must be a valid UUID format.", str(ve.exception) + ) + + chan = ChannelBuilder("sc://host/") + self.assertIsNone(chan.session_id) + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_connect_basic import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_column.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_column.py new file mode 100644 index 0000000..d838260 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_column.py @@ -0,0 +1,1042 @@ +# +# 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 decimal +import datetime + +from pyspark.sql.types import ( + Row, + StructField, + StructType, + MapType, + NullType, + DateType, + TimestampType, + TimestampNTZType, + ByteType, + BinaryType, + ShortType, + IntegerType, + FloatType, + DayTimeIntervalType, + StringType, + DoubleType, + LongType, + DecimalType, + BooleanType, +) +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.errors.exceptions.connect import SparkConnectException +from pyspark.testing.connectutils import should_test_connect +from pyspark.sql.tests.connect.test_connect_basic import SparkConnectSQLTestCase + + +if should_test_connect: + import pandas as pd + from pyspark.sql import functions as SF + from pyspark.sql.connect import functions as CF + from pyspark.sql.connect.column import Column + from pyspark.sql.connect.expressions import DistributedSequenceID, LiteralExpression + from pyspark.sql.connect.types import ( + JVM_BYTE_MIN, + JVM_BYTE_MAX, + JVM_SHORT_MIN, + JVM_SHORT_MAX, + JVM_INT_MIN, + JVM_INT_MAX, + JVM_LONG_MIN, + JVM_LONG_MAX, + ) + + +class SparkConnectColumnTests(SparkConnectSQLTestCase): + def compare_by_show(self, df1, df2, n: int = 20, truncate: int = 20): + from pyspark.sql.dataframe import DataFrame as SDF + from pyspark.sql.connect.dataframe import DataFrame as CDF + + assert isinstance(df1, (SDF, CDF)) + if isinstance(df1, SDF): + str1 = df1._jdf.showString(n, truncate, False) + else: + str1 = df1._show_string(n, truncate, False) + + assert isinstance(df2, (SDF, CDF)) + if isinstance(df2, SDF): + str2 = df2._jdf.showString(n, truncate, False) + else: + str2 = df2._show_string(n, truncate, False) + + self.assertEqual(str1, str2) + + def test_column_operator(self): + # SPARK-41351: Column needs to support != + df = self.connect.range(10) + self.assertEqual(9, len(df.filter(df.id != CF.lit(1)).collect())) + + def test_columns(self): + # SPARK-41036: test `columns` API for python client. + df = self.connect.read.table(self.tbl_name) + df2 = self.spark.read.table(self.tbl_name) + self.assertEqual(["id", "name"], df.columns) + + self.assert_eq( + df.filter(df.name.rlike("20")).toPandas(), df2.filter(df2.name.rlike("20")).toPandas() + ) + self.assert_eq( + df.filter(df.name.like("20")).toPandas(), df2.filter(df2.name.like("20")).toPandas() + ) + self.assert_eq( + df.filter(df.name.ilike("20")).toPandas(), df2.filter(df2.name.ilike("20")).toPandas() + ) + self.assert_eq( + df.filter(df.name.contains("20")).toPandas(), + df2.filter(df2.name.contains("20")).toPandas(), + ) + self.assert_eq( + df.filter(df.name.startswith("2")).toPandas(), + df2.filter(df2.name.startswith("2")).toPandas(), + ) + self.assert_eq( + df.filter(df.name.endswith("0")).toPandas(), + df2.filter(df2.name.endswith("0")).toPandas(), + ) + self.assert_eq( + df.select(df.name.substr(0, 1).alias("col")).toPandas(), + df2.select(df2.name.substr(0, 1).alias("col")).toPandas(), + ) + self.assert_eq( + df.select(df.name.substr(0, 1).name("col")).toPandas(), + df2.select(df2.name.substr(0, 1).name("col")).toPandas(), + ) + df3 = self.connect.sql("SELECT cast(null as int) as name") + df4 = self.spark.sql("SELECT cast(null as int) as name") + self.assert_eq( + df3.filter(df3.name.isNull()).toPandas(), + df4.filter(df4.name.isNull()).toPandas(), + ) + self.assert_eq( + df3.filter(df3.name.isNotNull()).toPandas(), + df4.filter(df4.name.isNotNull()).toPandas(), + ) + + # check error + with self.assertRaises(PySparkTypeError) as pe: + df.name.substr(df.id, 10) + + self.check_error( + exception=pe.exception, + error_class="NOT_SAME_TYPE", + message_parameters={ + "arg_name1": "startPos", + "arg_name2": "length", + "arg_type1": "Column", + "arg_type2": "int", + }, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.name.substr(10.5, 10.5) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT", + message_parameters={ + "arg_name": "length", + "arg_type": "float", + }, + ) + + def test_column_with_null(self): + # SPARK-41751: test isNull, isNotNull, eqNullSafe + + query = """ + SELECT * FROM VALUES + (1, 1, NULL), (2, NULL, NULL), (3, 3, 1) + AS tab(a, b, c) + """ + + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 1| 1|NULL| + # | 2|NULL|NULL| + # | 3| 3| 1| + # +---+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test isNull + self.assert_eq( + cdf.select(cdf.a.isNull(), cdf["b"].isNull(), CF.col("c").isNull()).toPandas(), + sdf.select(sdf.a.isNull(), sdf["b"].isNull(), SF.col("c").isNull()).toPandas(), + ) + + # test isNotNull + self.assert_eq( + cdf.select(cdf.a.isNotNull(), cdf["b"].isNotNull(), CF.col("c").isNotNull()).toPandas(), + sdf.select(sdf.a.isNotNull(), sdf["b"].isNotNull(), SF.col("c").isNotNull()).toPandas(), + ) + + # test eqNullSafe + self.assert_eq( + cdf.select(cdf.a.eqNullSafe(cdf.b), cdf["b"].eqNullSafe(CF.col("c"))).toPandas(), + sdf.select(sdf.a.eqNullSafe(sdf.b), sdf["b"].eqNullSafe(SF.col("c"))).toPandas(), + ) + + def test_invalid_ops(self): + query = """ + SELECT * FROM VALUES + (1, 1, 0, NULL), (2, NULL, 1, 2.0), (3, 3, 4, 3.5) + AS tab(a, b, c, d) + """ + cdf = self.connect.sql(query) + + with self.assertRaisesRegex( + ValueError, + "Cannot apply 'in' operator against a column", + ): + 1 in cdf.a + + with self.assertRaisesRegex( + ValueError, + "Cannot convert column into bool", + ): + cdf.a > 2 and cdf.b < 1 + + with self.assertRaisesRegex( + ValueError, + "Cannot convert column into bool", + ): + cdf.a > 2 or cdf.b < 1 + + with self.assertRaisesRegex( + ValueError, + "Cannot convert column into bool", + ): + not (cdf.a > 2) + + with self.assertRaisesRegex( + TypeError, + "Column is not iterable", + ): + for x in cdf.a: + pass + + def test_datetime(self): + query = """ + SELECT * FROM VALUES + (TIMESTAMP('2022-12-22 15:50:00'), DATE('2022-12-25'), 1.1), + (TIMESTAMP('2022-12-22 18:50:00'), NULL, 2.2), + (TIMESTAMP('2022-12-23 15:50:00'), DATE('2022-12-24'), 3.3), + (NULL, DATE('2022-12-22'), NULL) + AS tab(a, b, c) + """ + # +-------------------+----------+----+ + # | a| b| c| + # +-------------------+----------+----+ + # |2022-12-22 15:50:00|2022-12-25| 1.1| + # |2022-12-22 18:50:00| NULL| 2.2| + # |2022-12-23 15:50:00|2022-12-24| 3.3| + # | NULL|2022-12-22|NULL| + # +-------------------+----------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # datetime.date + self.assert_eq( + cdf.select(cdf.a < datetime.date(2022, 12, 23)).toPandas(), + sdf.select(sdf.a < datetime.date(2022, 12, 23)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a != datetime.date(2022, 12, 23)).toPandas(), + sdf.select(sdf.a != datetime.date(2022, 12, 23)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a == datetime.date(2022, 12, 22)).toPandas(), + sdf.select(sdf.a == datetime.date(2022, 12, 22)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b < datetime.date(2022, 12, 23)).toPandas(), + sdf.select(sdf.b < datetime.date(2022, 12, 23)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b >= datetime.date(2022, 12, 23)).toPandas(), + sdf.select(sdf.b >= datetime.date(2022, 12, 23)).toPandas(), + ) + + # datetime.datetime + self.assert_eq( + cdf.select(cdf.a < datetime.datetime(2022, 12, 22, 17, 0, 0)).toPandas(), + sdf.select(sdf.a < datetime.datetime(2022, 12, 22, 17, 0, 0)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a > datetime.datetime(2022, 12, 22, 17, 0, 0)).toPandas(), + sdf.select(sdf.a > datetime.datetime(2022, 12, 22, 17, 0, 0)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b >= datetime.datetime(2022, 12, 23, 17, 0, 0)).toPandas(), + sdf.select(sdf.b >= datetime.datetime(2022, 12, 23, 17, 0, 0)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b < datetime.datetime(2022, 12, 23, 17, 0, 0)).toPandas(), + sdf.select(sdf.b < datetime.datetime(2022, 12, 23, 17, 0, 0)).toPandas(), + ) + + def test_decimal(self): + # SPARK-41701: test decimal + query = """ + SELECT * FROM VALUES + (1, 1, 0, NULL), (2, NULL, 1, 2.0), (3, 3, 4, 3.5) + AS tab(a, b, c, d) + """ + # +---+----+---+----+ + # | a| b| c| d| + # +---+----+---+----+ + # | 1| 1| 0|NULL| + # | 2|NULL| 1| 2.0| + # | 3| 3| 4| 3.5| + # +---+----+---+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(cdf.a < decimal.Decimal(3)).toPandas(), + sdf.select(sdf.a < decimal.Decimal(3)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a != decimal.Decimal(2)).toPandas(), + sdf.select(sdf.a != decimal.Decimal(2)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a == decimal.Decimal(2)).toPandas(), + sdf.select(sdf.a == decimal.Decimal(2)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b < decimal.Decimal(2.5)).toPandas(), + sdf.select(sdf.b < decimal.Decimal(2.5)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.d >= decimal.Decimal(3.0)).toPandas(), + sdf.select(sdf.d >= decimal.Decimal(3.0)).toPandas(), + ) + + def test_none(self): + # SPARK-41783: test none + + query = """ + SELECT * FROM VALUES + (1, 1, NULL), (2, NULL, 1), (NULL, 3, 4) + AS tab(a, b, c) + """ + + # +----+----+----+ + # | a| b| c| + # +----+----+----+ + # | 1| 1|NULL| + # | 2|NULL| 1| + # |NULL| 3| 4| + # +----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(cdf.b > None, CF.col("c") >= None).toPandas(), + sdf.select(sdf.b > None, SF.col("c") >= None).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b < None, CF.col("c") <= None).toPandas(), + sdf.select(sdf.b < None, SF.col("c") <= None).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b.eqNullSafe(None), CF.col("c").eqNullSafe(None)).toPandas(), + sdf.select(sdf.b.eqNullSafe(None), SF.col("c").eqNullSafe(None)).toPandas(), + ) + + def test_simple_binary_expressions(self): + """Test complex expression""" + cdf = self.connect.read.table(self.tbl_name) + pdf = ( + cdf.select(cdf.id).where(cdf.id % CF.lit(30) == CF.lit(0)).sort(cdf.id.asc()).toPandas() + ) + self.assertEqual(len(pdf.index), 4) + + res = pd.DataFrame(data={"id": [0, 30, 60, 90]}) + self.assert_(pdf.equals(res), f"{pdf.to_string()} != {res.to_string()}") + + def test_literal_with_acceptable_type(self): + for value, dataType in [ + (b"binary\0\0asas", BinaryType()), + (True, BooleanType()), + (False, BooleanType()), + (0, ByteType()), + (JVM_BYTE_MIN, ByteType()), + (JVM_BYTE_MAX, ByteType()), + (0, ShortType()), + (JVM_SHORT_MIN, ShortType()), + (JVM_SHORT_MAX, ShortType()), + (0, IntegerType()), + (JVM_INT_MIN, IntegerType()), + (JVM_INT_MAX, IntegerType()), + (0, LongType()), + (JVM_LONG_MIN, LongType()), + (JVM_LONG_MAX, LongType()), + (0.0, FloatType()), + (1.234567, FloatType()), + (float("nan"), FloatType()), + (float("inf"), FloatType()), + (float("-inf"), FloatType()), + (0.0, DoubleType()), + (1.234567, DoubleType()), + (float("nan"), DoubleType()), + (float("inf"), DoubleType()), + (float("-inf"), DoubleType()), + (decimal.Decimal(0.0), DecimalType()), + (decimal.Decimal(1.234567), DecimalType()), + ("sss", StringType()), + (datetime.date(2022, 12, 13), DateType()), + (datetime.datetime.now(), DateType()), + (datetime.datetime.now(), TimestampType()), + (datetime.datetime.now(), TimestampNTZType()), + (datetime.timedelta(1, 2, 3), DayTimeIntervalType()), + ]: + lit = LiteralExpression(value=value, dataType=dataType) + self.assertEqual(dataType, lit._dataType) + + def test_literal_with_unsupported_type(self): + for value, dataType in [ + (b"binary\0\0asas", BooleanType()), + (True, StringType()), + (False, DoubleType()), + (JVM_BYTE_MIN - 1, ByteType()), + (JVM_BYTE_MAX + 1, ByteType()), + (JVM_SHORT_MIN - 1, ShortType()), + (JVM_SHORT_MAX + 1, ShortType()), + (JVM_INT_MIN - 1, IntegerType()), + (JVM_INT_MAX + 1, IntegerType()), + (JVM_LONG_MIN - 1, LongType()), + (JVM_LONG_MAX + 1, LongType()), + (0.1, DecimalType()), + (datetime.date(2022, 12, 13), TimestampType()), + (datetime.timedelta(1, 2, 3), DateType()), + ({1: 2}, MapType(IntegerType(), IntegerType())), + ( + {"a": "xyz", "b": 1}, + StructType([StructField("a", StringType()), StructField("b", IntegerType())]), + ), + ]: + with self.assertRaises(AssertionError): + LiteralExpression(value=value, dataType=dataType) + + def test_literal_null(self): + for dataType in [ + NullType(), + BinaryType(), + BooleanType(), + ByteType(), + ShortType(), + IntegerType(), + LongType(), + FloatType(), + DoubleType(), + DecimalType(), + DateType(), + TimestampType(), + TimestampNTZType(), + DayTimeIntervalType(), + ]: + lit_null = LiteralExpression(value=None, dataType=dataType) + self.assertTrue(lit_null._value is None) + self.assertEqual(dataType, lit_null._dataType) + + cdf = self.connect.range(0, 1).select(Column(lit_null)) + self.assertEqual(dataType, cdf.schema.fields[0].dataType) + + for value, dataType in [ + ("123", NullType()), + (123, NullType()), + (None, MapType(IntegerType(), IntegerType())), + (None, StructType([StructField("a", StringType())])), + ]: + with self.assertRaises(AssertionError): + LiteralExpression(value=value, dataType=dataType) + + def test_literal_integers(self): + cdf = self.connect.range(0, 1) + sdf = self.spark.range(0, 1) + + cdf1 = cdf.select( + CF.lit(0), + CF.lit(1), + CF.lit(-1), + CF.lit(JVM_INT_MAX), + CF.lit(JVM_INT_MIN), + CF.lit(JVM_INT_MAX + 1), + CF.lit(JVM_INT_MIN - 1), + CF.lit(JVM_LONG_MAX), + CF.lit(JVM_LONG_MIN), + CF.lit(JVM_LONG_MAX - 1), + CF.lit(JVM_LONG_MIN + 1), + ) + + sdf1 = sdf.select( + SF.lit(0), + SF.lit(1), + SF.lit(-1), + SF.lit(JVM_INT_MAX), + SF.lit(JVM_INT_MIN), + SF.lit(JVM_INT_MAX + 1), + SF.lit(JVM_INT_MIN - 1), + SF.lit(JVM_LONG_MAX), + SF.lit(JVM_LONG_MIN), + SF.lit(JVM_LONG_MAX - 1), + SF.lit(JVM_LONG_MIN + 1), + ) + + self.assertEqual(cdf1.schema, sdf1.schema) + self.assert_eq(cdf1.toPandas(), sdf1.toPandas()) + + # negative test for incorrect type + with self.assertRaises(PySparkValueError) as pe: + cdf.select(CF.lit(JVM_LONG_MAX + 1)).show() + + self.check_error( + exception=pe.exception, + error_class="VALUE_NOT_BETWEEN", + message_parameters={"arg_name": "value", "min": "-9223372036854775808", "max": "32767"}, + ) + + with self.assertRaises(PySparkValueError) as pe: + cdf.select(CF.lit(JVM_LONG_MIN - 1)).show() + + self.check_error( + exception=pe.exception, + error_class="VALUE_NOT_BETWEEN", + message_parameters={"arg_name": "value", "min": "-9223372036854775808", "max": "32767"}, + ) + + def test_cast(self): + # SPARK-41412: test basic Column.cast + df = self.connect.read.table(self.tbl_name) + df2 = self.spark.read.table(self.tbl_name) + + self.assert_eq( + df.select(df.id.cast("string")).toPandas(), df2.select(df2.id.cast("string")).toPandas() + ) + self.assert_eq( + df.select(df.id.astype("string")).toPandas(), + df2.select(df2.id.astype("string")).toPandas(), + ) + + for x in [ + StringType(), + ShortType(), + IntegerType(), + LongType(), + FloatType(), + DoubleType(), + ByteType(), + DecimalType(10, 2), + BooleanType(), + DayTimeIntervalType(), + ]: + self.assert_eq( + df.select(df.id.cast(x)).toPandas(), df2.select(df2.id.cast(x)).toPandas() + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.id.cast(10) + + self.check_error( + exception=pe.exception, + error_class="NOT_DATATYPE_OR_STR", + message_parameters={"arg_name": "dataType", "arg_type": "int"}, + ) + + def test_isin(self): + # SPARK-41526: test Column.isin + query = """ + SELECT * FROM VALUES + (1, 1, 0, NULL), (2, NULL, 1, 2.0), (3, 3, 4, 3.5) + AS tab(a, b, c, d) + """ + # +---+----+---+----+ + # | a| b| c| d| + # +---+----+---+----+ + # | 1| 1| 0|NULL| + # | 2|NULL| 1| 2.0| + # | 3| 3| 4| 3.5| + # +---+----+---+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test literals + self.assert_eq( + cdf.select(cdf.b.isin(1, 2, 3)).toPandas(), + sdf.select(sdf.b.isin(1, 2, 3)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b.isin([1, 2, 3])).toPandas(), + sdf.select(sdf.b.isin([1, 2, 3])).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.b.isin(set([1, 2, 3]))).toPandas(), + sdf.select(sdf.b.isin(set([1, 2, 3]))).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.d.isin([1.0, None, 3.5])).toPandas(), + sdf.select(sdf.d.isin([1.0, None, 3.5])).toPandas(), + ) + + # test columns + self.assert_eq( + cdf.select(cdf.a.isin(cdf.b)).toPandas(), + sdf.select(sdf.a.isin(sdf.b)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.a.isin(cdf.b, cdf.c)).toPandas(), + sdf.select(sdf.a.isin(sdf.b, sdf.c)).toPandas(), + ) + + # test columns mixed with literals + self.assert_eq( + cdf.select(cdf.a.isin(cdf.b, 4, 5, 6)).toPandas(), + sdf.select(sdf.a.isin(sdf.b, 4, 5, 6)).toPandas(), + ) + + def test_between(self): + query = """ + SELECT * FROM VALUES + (TIMESTAMP('2022-12-22 15:50:00'), DATE('2022-12-25'), 1.1), + (TIMESTAMP('2022-12-22 18:50:00'), NULL, 2.2), + (TIMESTAMP('2022-12-23 15:50:00'), DATE('2022-12-24'), 3.3), + (NULL, DATE('2022-12-22'), NULL) + AS tab(a, b, c) + """ + + # +-------------------+----------+----+ + # | a| b| c| + # +-------------------+----------+----+ + # |2022-12-22 15:50:00|2022-12-25| 1.1| + # |2022-12-22 18:50:00| NULL| 2.2| + # |2022-12-23 15:50:00|2022-12-24| 3.3| + # | NULL|2022-12-22|NULL| + # +-------------------+----------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(cdf.c.between(0, 2)).toPandas(), + sdf.select(sdf.c.between(0, 2)).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.c.between(1.1, 2.2)).toPandas(), + sdf.select(sdf.c.between(1.1, 2.2)).toPandas(), + ) + + self.assert_eq( + cdf.select(cdf.c.between(decimal.Decimal(0), decimal.Decimal(2))).toPandas(), + sdf.select(sdf.c.between(decimal.Decimal(0), decimal.Decimal(2))).toPandas(), + ) + + self.assert_eq( + cdf.select( + cdf.a.between( + datetime.datetime(2022, 12, 22, 17, 0, 0), + datetime.datetime(2022, 12, 23, 6, 0, 0), + ) + ).toPandas(), + sdf.select( + sdf.a.between( + datetime.datetime(2022, 12, 22, 17, 0, 0), + datetime.datetime(2022, 12, 23, 6, 0, 0), + ) + ).toPandas(), + ) + self.assert_eq( + cdf.select( + cdf.b.between(datetime.date(2022, 12, 23), datetime.date(2022, 12, 24)) + ).toPandas(), + sdf.select( + sdf.b.between(datetime.date(2022, 12, 23), datetime.date(2022, 12, 24)) + ).toPandas(), + ) + + def test_column_bitwise_ops(self): + # SPARK-41751: test bitwiseAND, bitwiseOR, bitwiseXOR + query = """ + SELECT * FROM VALUES + (1, 1, 0), (2, NULL, 1), (3, 3, 4) + AS tab(a, b, c) + """ + + # +---+----+---+ + # | a| b| c| + # +---+----+---+ + # | 1| 1| 0| + # | 2|NULL| 1| + # | 3| 3| 4| + # +---+----+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test bitwiseAND + self.assert_eq( + cdf.select(cdf.a.bitwiseAND(cdf.b), cdf["a"].bitwiseAND(CF.col("c"))).toPandas(), + sdf.select(sdf.a.bitwiseAND(sdf.b), sdf["a"].bitwiseAND(SF.col("c"))).toPandas(), + ) + + # test bitwiseOR + self.assert_eq( + cdf.select(cdf.a.bitwiseOR(cdf.b), cdf["a"].bitwiseOR(CF.col("c"))).toPandas(), + sdf.select(sdf.a.bitwiseOR(sdf.b), sdf["a"].bitwiseOR(SF.col("c"))).toPandas(), + ) + + # test bitwiseXOR + self.assert_eq( + cdf.select(cdf.a.bitwiseXOR(cdf.b), cdf["a"].bitwiseXOR(CF.col("c"))).toPandas(), + sdf.select(sdf.a.bitwiseXOR(sdf.b), sdf["a"].bitwiseXOR(SF.col("c"))).toPandas(), + ) + + def test_column_accessor(self): + query = """ + SELECT STRUCT(a, b, c) AS x, y, z, c FROM VALUES + (float(1.0), double(1.0), '2022', MAP('b', '123', 'a', 'kk'), ARRAY(1, 2, 3)), + (float(2.0), double(2.0), '2018', MAP('a', 'xy'), ARRAY(-1, -2, -3)), + (float(3.0), double(3.0), NULL, MAP('a', 'ab'), ARRAY(-1, 0, 1)) + AS tab(a, b, c, y, z) + """ + + # +----------------+-------------------+------------+----+ + # | x| y| z| c| + # +----------------+-------------------+------------+----+ + # |{1.0, 1.0, 2022}|{b -> 123, a -> kk}| [1, 2, 3]|2022| + # |{2.0, 2.0, 2018}| {a -> xy}|[-1, -2, -3]|2018| + # |{3.0, 3.0, null}| {a -> ab}| [-1, 0, 1]|NULL| + # +----------------+-------------------+------------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test struct + self.assert_eq( + cdf.select(cdf.x.a, cdf.x["b"], cdf["x"].c).toPandas(), + sdf.select(sdf.x.a, sdf.x["b"], sdf["x"].c).toPandas(), + ) + self.assert_eq( + cdf.select(CF.col("x").a, cdf.x.b, CF.col("x")["c"]).toPandas(), + sdf.select(SF.col("x").a, sdf.x.b, SF.col("x")["c"]).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.x.getItem("a"), cdf.x.getItem("b"), cdf["x"].getField("c")).toPandas(), + sdf.select(sdf.x.getItem("a"), sdf.x.getItem("b"), sdf["x"].getField("c")).toPandas(), + ) + + # test map + self.assert_eq( + cdf.select(cdf.y.a, cdf.y["b"], cdf["y"].c).toPandas(), + sdf.select(sdf.y.a, sdf.y["b"], sdf["y"].c).toPandas(), + ) + self.assert_eq( + cdf.select(CF.col("y").a, cdf.y.b, CF.col("y")["c"]).toPandas(), + sdf.select(SF.col("y").a, sdf.y.b, SF.col("y")["c"]).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.y.getItem("a"), cdf.y.getItem("b"), cdf["y"].getField("c")).toPandas(), + sdf.select(sdf.y.getItem("a"), sdf.y.getItem("b"), sdf["y"].getField("c")).toPandas(), + ) + + # test array + self.assert_eq( + cdf.select(cdf.z[0], cdf.z[1], cdf["z"][2]).toPandas(), + sdf.select(sdf.z[0], sdf.z[1], sdf["z"][2]).toPandas(), + ) + self.assert_eq( + cdf.select(CF.col("z")[0], cdf.z[10], CF.col("z")[-10]).toPandas(), + sdf.select(SF.col("z")[0], sdf.z[10], SF.col("z")[-10]).toPandas(), + ) + self.assert_eq( + cdf.select(cdf.z.getItem(0), cdf.z.getItem(1), cdf["z"].getField(2)).toPandas(), + sdf.select(sdf.z.getItem(0), sdf.z.getItem(1), sdf["z"].getField(2)).toPandas(), + ) + + # test string with slice + self.assert_eq( + cdf.select(cdf.c[0:1], cdf["c"][2:10]).toPandas(), + sdf.select(sdf.c[0:1], sdf["c"][2:10]).toPandas(), + ) + + def test_column_arithmetic_ops(self): + # SPARK-41761: test arithmetic ops + query = """ + SELECT * FROM VALUES + (1, 1, 0, NULL), (2, NULL, 1, 2.0), (3, 3, 4, 3.5) + AS tab(a, b, c, d) + """ + # +---+----+---+----+ + # | a| b| c| d| + # +---+----+---+----+ + # | 1| 1| 0|NULL| + # | 2|NULL| 1| 2.0| + # | 3| 3| 4| 3.5| + # +---+----+---+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select( + cdf.a + cdf["b"] - 1, cdf.a - cdf["b"] * cdf["c"] / 2, cdf.d / cdf.b / 3 + ).toPandas(), + sdf.select( + sdf.a + sdf["b"] - 1, sdf.a - sdf["b"] * sdf["c"] / 2, sdf.d / sdf.b / 3 + ).toPandas(), + ) + + # TODO(SPARK-41762): make __neg__ return the correct column name + # [left]: Index(['negative(a)'], dtype='object') + # [right]: Index(['(- a)'], dtype='object') + self.assert_eq( + cdf.select((-cdf.a).alias("x")).toPandas(), + sdf.select((-sdf.a).alias("x")).toPandas(), + ) + + self.assert_eq( + cdf.select(3 - cdf.a + cdf["b"] * cdf["c"] - cdf.d / cdf.b).toPandas(), + sdf.select(3 - sdf.a + sdf["b"] * sdf["c"] - sdf.d / sdf.b).toPandas(), + ) + + self.assert_eq( + cdf.select(cdf.a % cdf["b"], cdf["a"] % 2, 12 % cdf.c).toPandas(), + sdf.select(sdf.a % sdf["b"], sdf["a"] % 2, 12 % sdf.c).toPandas(), + ) + + self.assert_eq( + cdf.select(cdf.a ** cdf["b"], cdf.d**2, 2**cdf.c).toPandas(), + sdf.select(sdf.a ** sdf["b"], sdf.d**2, 2**sdf.c).toPandas(), + ) + + def test_column_field_ops(self): + # SPARK-41767: test withField, dropFields + query = """ + SELECT STRUCT(a, b, c, d) AS x, e FROM VALUES + (float(1.0), double(1.0), '2022', 1, 0), + (float(2.0), double(2.0), '2018', NULL, 2), + (float(3.0), double(3.0), NULL, 3, NULL) + AS tab(a, b, c, d, e) + """ + + # +----------------------+----+ + # | x| e| + # +----------------------+----+ + # | {1.0, 1.0, 2022, 1}| 0| + # |{2.0, 2.0, 2018, null}| 2| + # | {3.0, 3.0, null, 3}|NULL| + # +----------------------+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # add field + self.compare_by_show( + cdf.select(cdf.x.withField("z", cdf.e)), + sdf.select(sdf.x.withField("z", sdf.e)), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.withField("z", CF.col("e"))), + sdf.select(sdf.x.withField("z", SF.col("e"))), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.withField("z", CF.lit("xyz"))), + sdf.select(sdf.x.withField("z", SF.lit("xyz"))), + truncate=100, + ) + + # replace field + self.compare_by_show( + cdf.select(cdf.x.withField("a", cdf.e)), + sdf.select(sdf.x.withField("a", sdf.e)), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.withField("a", CF.col("e"))), + sdf.select(sdf.x.withField("a", SF.col("e"))), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.withField("a", CF.lit("xyz"))), + sdf.select(sdf.x.withField("a", SF.lit("xyz"))), + truncate=100, + ) + + # drop field + self.compare_by_show( + cdf.select(cdf.x.dropFields("a")), + sdf.select(sdf.x.dropFields("a")), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.dropFields("z")), + sdf.select(sdf.x.dropFields("z")), + truncate=100, + ) + self.compare_by_show( + cdf.select(cdf.x.dropFields("a", "b", "z")), + sdf.select(sdf.x.dropFields("a", "b", "z")), + truncate=100, + ) + + # check error + # invalid column: not a struct column + with self.assertRaises(SparkConnectException): + cdf.select(cdf.e.withField("a", CF.lit(1))).show() + + # invalid column: not a struct column + with self.assertRaises(SparkConnectException): + cdf.select(cdf.e.dropFields("a")).show() + + # cannot drop all fields in struct + with self.assertRaises(SparkConnectException): + cdf.select(cdf.x.dropFields("a", "b", "c", "d")).show() + + with self.assertRaises(PySparkTypeError) as pe: + cdf.select(cdf.x.withField(CF.col("a"), cdf.e)).show() + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "fieldName", "arg_type": "Column"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.select(cdf.x.withField("a", 2)).show() + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.select(cdf.x.dropFields("a", 1, 2)).show() + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "fieldName", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkValueError) as pe: + cdf.select(cdf.x.dropFields()).show() + + self.check_error( + exception=pe.exception, + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "dropFields"}, + ) + + def test_column_string_ops(self): + # SPARK-41764: test string ops + query = """ + SELECT * FROM VALUES + (1, 'abcdef', 'ghij', 'hello world', 'a'), + (2, 'abcd', 'efghij', 'how are you', 'd') + AS tab(a, b, c, d, e) + """ + + # +---+------+------+-----------+---+ + # | a| b| c| d| e| + # +---+------+------+-----------+---+ + # | 1|abcdef| ghij|hello world| a| + # | 2| abcd|efghij|how are you| d| + # +---+------+------+-----------+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select( + cdf.b.startswith("a"), cdf["c"].startswith("g"), cdf["b"].startswith(cdf.e) + ).toPandas(), + sdf.select( + sdf.b.startswith("a"), sdf["c"].startswith("g"), sdf["b"].startswith(sdf.e) + ).toPandas(), + ) + + self.assert_eq( + cdf.select( + cdf.b.endswith("a"), cdf["c"].endswith("j"), cdf["b"].endswith(cdf.e) + ).toPandas(), + sdf.select( + sdf.b.endswith("a"), sdf["c"].endswith("j"), sdf["b"].endswith(sdf.e) + ).toPandas(), + ) + + self.assert_eq( + cdf.select( + cdf.b.contains("a"), cdf["c"].contains("j"), cdf["b"].contains(cdf.e) + ).toPandas(), + sdf.select( + sdf.b.contains("a"), sdf["c"].contains("j"), sdf["b"].contains(sdf.e) + ).toPandas(), + ) + + def test_with_field_column_name(self): + data = [Row(a=Row(b=1, c=2))] + + cdf = self.connect.createDataFrame(data) + cdf1 = cdf.withColumn("a", cdf["a"].withField("b", CF.lit(3))).select("a.b") + + sdf = self.spark.createDataFrame(data) + sdf1 = sdf.withColumn("a", sdf["a"].withField("b", SF.lit(3))).select("a.b") + + self.assertEqual(cdf1.schema, sdf1.schema) + self.assertEqual(cdf1.collect(), sdf1.collect()) + + def test_distributed_sequence_id(self): + cdf = self.connect.range(10) + expected = self.connect.range(0, 10).selectExpr("id as index", "id") + self.assertEqual( + cdf.select(Column(DistributedSequenceID()).alias("index"), "*").collect(), + expected.collect(), + ) + + +if __name__ == "__main__": + import os + import unittest + from pyspark.sql.tests.connect.test_connect_column import * # noqa: F401 + + # TODO(SPARK-41794): Enable ANSI mode in this file. + os.environ["SPARK_ANSI_SQL_MODE"] = "false" + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_function.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_function.py new file mode 100644 index 0000000..a5d330f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_function.py @@ -0,0 +1,2380 @@ +# +# 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 os +import unittest + +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql import SparkSession as PySparkSession +from pyspark.sql.types import StringType, StructType, StructField, ArrayType, IntegerType +from pyspark.testing import assertDataFrameEqual +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect +from pyspark.testing.sqlutils import SQLTestUtils +from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException + +if should_test_connect: + from pyspark.sql.connect.column import Column + from pyspark.sql import functions as SF + from pyspark.sql.window import Window as SW + from pyspark.sql.dataframe import DataFrame as SDF + from pyspark.sql.connect import functions as CF + from pyspark.sql.connect.window import Window as CW + from pyspark.sql.connect.dataframe import DataFrame as CDF + + +class SparkConnectFunctionTests(ReusedConnectTestCase, PandasOnSparkTestUtils, SQLTestUtils): + """These test cases exercise the interface to the proto plan + generation but do not call Spark.""" + + @classmethod + def setUpClass(cls): + super(SparkConnectFunctionTests, cls).setUpClass() + # Disable the shared namespace so pyspark.sql.functions, etc point the regular + # PySpark libraries. + os.environ["PYSPARK_NO_NAMESPACE_SHARE"] = "1" + cls.connect = cls.spark # Switch Spark Connect session and regular PySpark sesion. + cls.spark = PySparkSession._instantiatedSession + assert cls.spark is not None + + @classmethod + def tearDownClass(cls): + cls.spark = cls.connect # Stopping Spark Connect closes the session in JVM at the server. + super(SparkConnectFunctionTests, cls).setUpClass() + del os.environ["PYSPARK_NO_NAMESPACE_SHARE"] + + def compare_by_show(self, df1, df2, n: int = 20, truncate: int = 20): + assert isinstance(df1, (SDF, CDF)) + if isinstance(df1, SDF): + str1 = df1._jdf.showString(n, truncate, False) + else: + str1 = df1._show_string(n, truncate, False) + + assert isinstance(df2, (SDF, CDF)) + if isinstance(df2, SDF): + str2 = df2._jdf.showString(n, truncate, False) + else: + str2 = df2._show_string(n, truncate, False) + + self.assertEqual(str1, str2) + + def test_count_star(self): + # SPARK-42099: test count(*), count(col(*)) and count(expr(*)) + data = [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")] + + cdf = self.connect.createDataFrame(data, schema=["age", "name"]) + sdf = self.spark.createDataFrame(data, schema=["age", "name"]) + + self.assertEqual( + cdf.select(CF.count(CF.expr("*")), CF.count(cdf.age)).collect(), + sdf.select(SF.count(SF.expr("*")), SF.count(sdf.age)).collect(), + ) + + self.assertEqual( + cdf.select(CF.count(CF.col("*")), CF.count(cdf.age)).collect(), + sdf.select(SF.count(SF.col("*")), SF.count(sdf.age)).collect(), + ) + + self.assertEqual( + cdf.select(CF.count("*"), CF.count(cdf.age)).collect(), + sdf.select(SF.count("*"), SF.count(sdf.age)).collect(), + ) + + self.assertEqual( + cdf.groupby("name").agg({"*": "count"}).sort("name").collect(), + sdf.groupby("name").agg({"*": "count"}).sort("name").collect(), + ) + + self.assertEqual( + cdf.groupby("name") + .agg(CF.count(CF.expr("*")), CF.count(cdf.age)) + .sort("name") + .collect(), + sdf.groupby("name") + .agg(SF.count(SF.expr("*")), SF.count(sdf.age)) + .sort("name") + .collect(), + ) + + self.assertEqual( + cdf.groupby("name") + .agg(CF.count(CF.col("*")), CF.count(cdf.age)) + .sort("name") + .collect(), + sdf.groupby("name") + .agg(SF.count(SF.col("*")), SF.count(sdf.age)) + .sort("name") + .collect(), + ) + + self.assertEqual( + cdf.groupby("name").agg(CF.count("*"), CF.count(cdf.age)).sort("name").collect(), + sdf.groupby("name").agg(SF.count("*"), SF.count(sdf.age)).sort("name").collect(), + ) + + def test_broadcast(self): + query = """ + SELECT * FROM VALUES + (0, float("NAN"), NULL), (1, NULL, 2.0), (2, 2.1, 3.5) + AS tab(a, b, c) + """ + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 0| NaN|NULL| + # | 1|NULL| 2.0| + # | 2| 2.1| 3.5| + # +---+----+----+ + + cdf = self.connect.sql(query) + cdf1 = cdf.select(cdf.a, "b") + cdf2 = cdf.select(cdf.a, "c") + + sdf = self.spark.sql(query) + sdf1 = sdf.select(sdf.a, "b") + sdf2 = sdf.select(sdf.a, "c") + + self.assert_eq( + cdf1.join(cdf2, on="a").toPandas(), + sdf1.join(sdf2, on="a").toPandas(), + ) + self.assert_eq( + cdf1.join(CF.broadcast(cdf2), on="a").toPandas(), + sdf1.join(SF.broadcast(sdf2), on="a").toPandas(), + ) + self.assert_eq( + CF.broadcast(cdf1).join(cdf2, on="a").toPandas(), + SF.broadcast(sdf1).join(sdf2, on="a").toPandas(), + ) + self.assert_eq( + CF.broadcast(cdf1).join(CF.broadcast(cdf2), on="a").toPandas(), + SF.broadcast(sdf1).join(SF.broadcast(sdf2), on="a").toPandas(), + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.broadcast(cdf.a) + + self.check_error( + exception=pe.exception, + error_class="NOT_DATAFRAME", + message_parameters={"arg_name": "df", "arg_type": "Column"}, + ) + + def test_normal_functions(self): + query = """ + SELECT * FROM VALUES + (0, float("NAN"), NULL), (1, NULL, 2.0), (2, 2.1, 3.5) + AS tab(a, b, c) + """ + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 0| NaN|NULL| + # | 1|NULL| 2.0| + # | 2| 2.1| 3.5| + # +---+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(CF.bitwise_not(cdf.a)).toPandas(), + sdf.select(SF.bitwise_not(sdf.a)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.bitwiseNOT(cdf.a)).toPandas(), + sdf.select(SF.bitwiseNOT(sdf.a)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.coalesce(cdf.a, "b", cdf.c)).toPandas(), + sdf.select(SF.coalesce(sdf.a, "b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.expr("a + b - c")).toPandas(), + sdf.select(SF.expr("a + b - c")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.greatest(cdf.a, "b", cdf.c)).toPandas(), + sdf.select(SF.greatest(sdf.a, "b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.isnan(cdf.a), CF.isnan("b")).toPandas(), + sdf.select(SF.isnan(sdf.a), SF.isnan("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.isnull(cdf.a), CF.isnull("b")).toPandas(), + sdf.select(SF.isnull(sdf.a), SF.isnull("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.input_file_name()).toPandas(), + sdf.select(SF.input_file_name()).toPandas(), + ) + self.assert_eq( + cdf.select(CF.least(cdf.a, "b", cdf.c)).toPandas(), + sdf.select(SF.least(sdf.a, "b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.monotonically_increasing_id()).toPandas(), + sdf.select(SF.monotonically_increasing_id()).toPandas(), + ) + self.assert_eq( + cdf.select(CF.nanvl("b", cdf.c)).toPandas(), + sdf.select(SF.nanvl("b", sdf.c)).toPandas(), + ) + # Can not compare the values due to the random seed + self.assertEqual( + cdf.select(CF.rand()).count(), + sdf.select(SF.rand()).count(), + ) + self.assert_eq( + cdf.select(CF.rand(100)).toPandas(), + sdf.select(SF.rand(100)).toPandas(), + ) + # Can not compare the values due to the random seed + self.assertEqual( + cdf.select(CF.randn()).count(), + sdf.select(SF.randn()).count(), + ) + self.assert_eq( + cdf.select(CF.randn(100)).toPandas(), + sdf.select(SF.randn(100)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.spark_partition_id()).toPandas(), + sdf.select(SF.spark_partition_id()).toPandas(), + ) + + def test_when_otherwise(self): + query = """ + SELECT * FROM VALUES + (0, float("NAN"), NULL), (1, NULL, 2.0), (2, 2.1, 3.5), (3, 3.1, float("NAN")) + AS tab(a, b, c) + """ + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 0| NaN|NULL| + # | 1|NULL| 2.0| + # | 2| 2.1| 3.5| + # | 3| 3.1| NaN| + # +---+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(CF.when(cdf.a == 0, 1.0).otherwise(2.0)).toPandas(), + sdf.select(SF.when(sdf.a == 0, 1.0).otherwise(2.0)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.when(cdf.a < 1, cdf.b).otherwise(cdf.c)).toPandas(), + sdf.select(SF.when(sdf.a < 1, sdf.b).otherwise(sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.when(cdf.a == 0, 1.0) + .when(CF.col("a") == 1, 2.0) + .when(cdf.a == 2, -1.0) + .otherwise(cdf.c) + ).toPandas(), + sdf.select( + SF.when(sdf.a == 0, 1.0) + .when(SF.col("a") == 1, 2.0) + .when(sdf.a == 2, -1.0) + .otherwise(sdf.c) + ).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.when(cdf.a < cdf.b, 1.0) + .when(CF.col("a") == 1, CF.abs("c") + cdf.b) + .otherwise(cdf.c + CF.col("a")) + ).toPandas(), + sdf.select( + SF.when(sdf.a < sdf.b, 1.0) + .when(SF.col("a") == 1, SF.abs("c") + sdf.b) + .otherwise(sdf.c + SF.col("a")) + ).toPandas(), + ) + + # when without otherwise + self.assert_eq( + cdf.select(CF.when(cdf.a < 1, cdf.b)).toPandas(), + sdf.select(SF.when(sdf.a < 1, sdf.b)).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.when(cdf.a == 0, 1.0) + .when(CF.col("a") == 1, cdf.b + CF.col("c")) + .when(cdf.a == 2, CF.abs(cdf.b)) + ).toPandas(), + sdf.select( + SF.when(sdf.a == 0, 1.0) + .when(SF.col("a") == 1, sdf.b + SF.col("c")) + .when(sdf.a == 2, SF.abs(sdf.b)) + ).toPandas(), + ) + + # check error + with self.assertRaisesRegex( + TypeError, + "when.* can only be applied on a Column previously generated by when.* function", + ): + cdf.a.when(cdf.a == 0, 1.0) + + with self.assertRaisesRegex( + TypeError, + "when.* can only be applied on a Column previously generated by when.* function", + ): + CF.col("c").when(cdf.a == 0, 1.0) + + with self.assertRaisesRegex( + TypeError, + "otherwise.* can only be applied on a Column previously generated by when", + ): + cdf.a.otherwise(1.0) + + with self.assertRaisesRegex( + TypeError, + "otherwise.* can only be applied on a Column previously generated by when", + ): + CF.col("c").otherwise(1.0) + + with self.assertRaisesRegex( + TypeError, + "otherwise.* can only be applied once on a Column previously generated by when", + ): + CF.when(cdf.a == 0, 1.0).otherwise(1.0).otherwise(1.0) + + with self.assertRaises(PySparkTypeError) as pe: + CF.when(True, 1.0).otherwise(1.0) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": "bool"}, + ) + + def test_sorting_functions_with_column(self): + funs = [ + CF.asc_nulls_first, + CF.asc_nulls_last, + CF.desc_nulls_first, + CF.desc_nulls_last, + ] + exprs = [CF.col("x"), "x"] + + for fun in funs: + for _expr in exprs: + res = fun(_expr) + self.assertIsInstance(res, Column) + self.assertIn(f"""{fun.__name__.replace("_", " ").upper()}'""", str(res)) + + for _expr in exprs: + res = CF.asc(_expr) + self.assertIsInstance(res, Column) + self.assertIn("""ASC NULLS FIRST'""", str(res)) + + for _expr in exprs: + res = CF.desc(_expr) + self.assertIsInstance(res, Column) + self.assertIn("""DESC NULLS LAST'""", str(res)) + + def test_sort_with_nulls_order(self): + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (true, NULL, 2.0), (NULL, 3, 3.0) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # | true|NULL| 2.0| + # | NULL| 3| 3.0| + # +-----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + for c in ["a", "b", "c"]: + self.assert_eq( + cdf.orderBy(CF.asc(c)).toPandas(), + sdf.orderBy(SF.asc(c)).toPandas(), + ) + self.assert_eq( + cdf.orderBy(CF.asc_nulls_first(c)).toPandas(), + sdf.orderBy(SF.asc_nulls_first(c)).toPandas(), + ) + self.assert_eq( + cdf.orderBy(CF.asc_nulls_last(c)).toPandas(), + sdf.orderBy(SF.asc_nulls_last(c)).toPandas(), + ) + self.assert_eq( + cdf.orderBy(CF.desc(c)).toPandas(), + sdf.orderBy(SF.desc(c)).toPandas(), + ) + self.assert_eq( + cdf.orderBy(CF.desc_nulls_first(c)).toPandas(), + sdf.orderBy(SF.desc_nulls_first(c)).toPandas(), + ) + self.assert_eq( + cdf.orderBy(CF.desc_nulls_last(c)).toPandas(), + sdf.orderBy(SF.desc_nulls_last(c)).toPandas(), + ) + + def test_math_functions(self): + query = """ + SELECT * FROM VALUES + (false, 1, NULL), (true, NULL, 2.0), (NULL, 3, 3.5) + AS tab(a, b, c) + """ + # +-----+----+----+ + # | a| b| c| + # +-----+----+----+ + # |false| 1|NULL| + # | true|NULL| 2.0| + # | NULL| 3| 3.5| + # +-----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + for cfunc, sfunc in [ + (CF.abs, SF.abs), + (CF.acos, SF.acos), + (CF.acosh, SF.acosh), + (CF.asin, SF.asin), + (CF.asinh, SF.asinh), + (CF.atan, SF.atan), + (CF.atanh, SF.atanh), + (CF.bin, SF.bin), + (CF.cbrt, SF.cbrt), + (CF.ceil, SF.ceil), + (CF.cos, SF.cos), + (CF.cosh, SF.cosh), + (CF.cot, SF.cot), + (CF.csc, SF.csc), + (CF.degrees, SF.degrees), + (CF.toDegrees, SF.toDegrees), + (CF.exp, SF.exp), + (CF.expm1, SF.expm1), + (CF.factorial, SF.factorial), + (CF.floor, SF.floor), + (CF.hex, SF.hex), + (CF.log, SF.log), + (CF.log10, SF.log10), + (CF.log1p, SF.log1p), + (CF.log2, SF.log2), + (CF.radians, SF.radians), + (CF.toRadians, SF.toRadians), + (CF.rint, SF.rint), + (CF.sec, SF.sec), + (CF.signum, SF.signum), + (CF.sin, SF.sin), + (CF.sinh, SF.sinh), + (CF.sqrt, SF.sqrt), + (CF.tan, SF.tan), + (CF.tanh, SF.tanh), + (CF.unhex, SF.unhex), + ]: + self.assert_eq( + cdf.select(cfunc("b"), cfunc(cdf.c)).toPandas(), + sdf.select(sfunc("b"), sfunc(sdf.c)).toPandas(), + ) + + # test log(arg1, arg2) + self.assert_eq( + cdf.select(CF.log(1.1, "b"), CF.log(1.2, cdf.c)).toPandas(), + sdf.select(SF.log(1.1, "b"), SF.log(1.2, sdf.c)).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.atan2("b", cdf.c)).toPandas(), + sdf.select(SF.atan2("b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.bround("b", 1)).toPandas(), + sdf.select(SF.bround("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.conv("b", 2, 16)).toPandas(), + sdf.select(SF.conv("b", 2, 16)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.hypot("b", cdf.c)).toPandas(), + sdf.select(SF.hypot("b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.pmod("b", cdf.c)).toPandas(), + sdf.select(SF.pmod("b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.pow("b", cdf.c)).toPandas(), + sdf.select(SF.pow("b", sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.round("b", 1)).toPandas(), + sdf.select(SF.round("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftleft("b", 1)).toPandas(), + sdf.select(SF.shiftleft("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftLeft("b", 1)).toPandas(), + sdf.select(SF.shiftLeft("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftright("b", 1)).toPandas(), + sdf.select(SF.shiftright("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftRight("b", 1)).toPandas(), + sdf.select(SF.shiftRight("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftrightunsigned("b", 1)).toPandas(), + sdf.select(SF.shiftrightunsigned("b", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.shiftRightUnsigned("b", 1)).toPandas(), + sdf.select(SF.shiftRightUnsigned("b", 1)).toPandas(), + ) + + def test_aggregation_functions(self): + query = """ + SELECT * FROM VALUES + (0, float("NAN"), NULL), (1, NULL, 2.0), (1, 2.1, 3.5), (0, 0.5, 1.0) + AS tab(a, b, c) + """ + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 0| NaN|NULL| + # | 1|NULL| 2.0| + # | 1| 2.1| 3.5| + # | 0| 0.5| 1.0| + # +---+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + for cfunc, sfunc in [ + (CF.approx_count_distinct, SF.approx_count_distinct), + (CF.approxCountDistinct, SF.approxCountDistinct), + (CF.avg, SF.avg), + (CF.collect_list, SF.collect_list), + (CF.collect_set, SF.collect_set), + (CF.count, SF.count), + (CF.first, SF.first), + (CF.kurtosis, SF.kurtosis), + (CF.last, SF.last), + (CF.max, SF.max), + (CF.mean, SF.mean), + (CF.median, SF.median), + (CF.min, SF.min), + (CF.mode, SF.mode), + (CF.product, SF.product), + (CF.skewness, SF.skewness), + (CF.stddev, SF.stddev), + (CF.stddev_pop, SF.stddev_pop), + (CF.stddev_samp, SF.stddev_samp), + (CF.sum, SF.sum), + (CF.sum_distinct, SF.sum_distinct), + (CF.sumDistinct, SF.sumDistinct), + (CF.var_pop, SF.var_pop), + (CF.var_samp, SF.var_samp), + (CF.variance, SF.variance), + ]: + self.assert_eq( + cdf.select(cfunc("b"), cfunc(cdf.c)).toPandas(), + sdf.select(sfunc("b"), sfunc(sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.groupBy("a").agg(cfunc("b"), cfunc(cdf.c)).toPandas(), + sdf.groupBy("a").agg(sfunc("b"), sfunc(sdf.c)).toPandas(), + ) + + for cfunc, sfunc in [ + (CF.corr, SF.corr), + (CF.covar_pop, SF.covar_pop), + (CF.covar_samp, SF.covar_samp), + (CF.max_by, SF.max_by), + (CF.min_by, SF.min_by), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.b, "c")).toPandas(), + sdf.select(sfunc(sdf.b, "c")).toPandas(), + ) + self.assert_eq( + cdf.groupBy("a").agg(cfunc(cdf.b, "c")).toPandas(), + sdf.groupBy("a").agg(sfunc(sdf.b, "c")).toPandas(), + ) + + # test grouping + self.assert_eq( + cdf.cube("a").agg(CF.grouping("a"), CF.sum("c")).orderBy("a").toPandas(), + sdf.cube("a").agg(SF.grouping("a"), SF.sum("c")).orderBy("a").toPandas(), + ) + + # test grouping_id + self.assert_eq( + cdf.cube("a").agg(CF.grouping_id(), CF.sum("c")).orderBy("a").toPandas(), + sdf.cube("a").agg(SF.grouping_id(), SF.sum("c")).orderBy("a").toPandas(), + ) + + # test percentile_approx + self.assert_eq( + cdf.select(CF.percentile_approx(cdf.b, 0.5, 1000)).toPandas(), + sdf.select(SF.percentile_approx(sdf.b, 0.5, 1000)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.percentile_approx(cdf.b, [0.1, 0.9])).toPandas(), + sdf.select(SF.percentile_approx(sdf.b, [0.1, 0.9])).toPandas(), + ) + self.assert_eq( + cdf.groupBy("a").agg(CF.percentile_approx("b", 0.5)).toPandas(), + sdf.groupBy("a").agg(SF.percentile_approx("b", 0.5)).toPandas(), + ) + self.assert_eq( + cdf.groupBy("a").agg(CF.percentile_approx(cdf.b, [0.1, 0.9])).toPandas(), + sdf.groupBy("a").agg(SF.percentile_approx(sdf.b, [0.1, 0.9])).toPandas(), + ) + + # test count_distinct + self.assert_eq( + cdf.select(CF.count_distinct("b"), CF.count_distinct(cdf.c)).toPandas(), + sdf.select(SF.count_distinct("b"), SF.count_distinct(sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.countDistinct("b"), CF.countDistinct(cdf.c)).toPandas(), + sdf.select(SF.countDistinct("b"), SF.countDistinct(sdf.c)).toPandas(), + ) + # The output column names of 'groupBy.agg(count_distinct)' in PySpark + # are incorrect, see SPARK-41391. + self.assert_eq( + cdf.groupBy("a") + .agg(CF.count_distinct("b").alias("x"), CF.count_distinct(cdf.c).alias("y")) + .toPandas(), + sdf.groupBy("a") + .agg(SF.count_distinct("b").alias("x"), SF.count_distinct(sdf.c).alias("y")) + .toPandas(), + ) + + def test_window_functions(self): + self.assertEqual(CW.unboundedPreceding, SW.unboundedPreceding) + + self.assertEqual(CW.unboundedFollowing, SW.unboundedFollowing) + + self.assertEqual(CW.currentRow, SW.currentRow) + + query = """ + SELECT * FROM VALUES + (0, float("NAN"), NULL), (1, NULL, 2.0), (1, 2.1, 3.5), (0, 0.5, 1.0), + (0, 1.5, 1.1), (1, 2.2, -1.0), (1, 0.1, -0.1), (0, 0.0, 5.0) + AS tab(a, b, c) + """ + # +---+----+----+ + # | a| b| c| + # +---+----+----+ + # | 0| NaN|NULL| + # | 1|NULL| 2.0| + # | 1| 2.1| 3.5| + # | 0| 0.5| 1.0| + # | 0| 1.5| 1.1| + # | 1| 2.2|-1.0| + # | 1| 0.1|-0.1| + # | 0| 0.0| 5.0| + # +---+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test window functions + for ccol, scol in [ + (CF.row_number(), SF.row_number()), + (CF.rank(), SF.rank()), + (CF.dense_rank(), SF.dense_rank()), + (CF.percent_rank(), SF.percent_rank()), + (CF.cume_dist(), SF.cume_dist()), + (CF.lag("c", 1), SF.lag("c", 1)), + (CF.lag("c", 1, -1.0), SF.lag("c", 1, -1.0)), + (CF.lag(cdf.c, -1), SF.lag(sdf.c, -1)), + (CF.lag(cdf.c, -1, float("nan")), SF.lag(sdf.c, -1, float("nan"))), + (CF.lead("c", 1), SF.lead("c", 1)), + (CF.lead("c", 1, -1.0), SF.lead("c", 1, -1.0)), + (CF.lead(cdf.c, -1), SF.lead(sdf.c, -1)), + (CF.lead(cdf.c, -1, float("nan")), SF.lead(sdf.c, -1, float("nan"))), + (CF.nth_value("c", 1), SF.nth_value("c", 1)), + (CF.nth_value(cdf.c, 2), SF.nth_value(sdf.c, 2)), + (CF.nth_value(cdf.c, 2, True), SF.nth_value(sdf.c, 2, True)), + (CF.nth_value(cdf.c, 2, False), SF.nth_value(sdf.c, 2, False)), + (CF.ntile(1), SF.ntile(1)), + (CF.ntile(2), SF.ntile(2)), + (CF.ntile(4), SF.ntile(4)), + ]: + + for cwin, swin in [ + (CW.orderBy("b"), SW.orderBy("b")), + (CW.partitionBy("a").orderBy("b"), SW.partitionBy("a").orderBy("b")), + ( + CW.partitionBy("a").orderBy(CF.col("b").desc()), + SW.partitionBy("a").orderBy(SF.col("b").desc()), + ), + (CW.partitionBy("a", cdf.c).orderBy("b"), SW.partitionBy("a", sdf.c).orderBy("b")), + (CW.partitionBy("a").orderBy("b", cdf.c), SW.partitionBy("a").orderBy("b", sdf.c)), + ( + CW.partitionBy("a").orderBy("b", cdf.c.desc()), + SW.partitionBy("a").orderBy("b", sdf.c.desc()), + ), + ]: + + self.assert_eq( + cdf.select(ccol.over(cwin)).toPandas(), + sdf.select(scol.over(swin)).toPandas(), + ) + + # test aggregation functions + for ccol, scol in [ + (CF.count("c"), SF.count("c")), + (CF.sum("c"), SF.sum("c")), + (CF.max(cdf.c), SF.max(sdf.c)), + (CF.min(cdf.c), SF.min(sdf.c)), + ]: + + for cwin, swin in [ + (CW.orderBy("b"), SW.orderBy("b")), + ( + CW.orderBy("b").rowsBetween(CW.currentRow, CW.currentRow), + SW.orderBy("b").rowsBetween(SW.currentRow, SW.currentRow), + ), + ( + CW.orderBy(cdf.b.desc()).rowsBetween(CW.currentRow - 1, CW.currentRow + 2), + SW.orderBy(sdf.b.desc()).rowsBetween(SW.currentRow - 1, SW.currentRow + 2), + ), + ( + CW.orderBy("b").rowsBetween(CW.unboundedPreceding, CW.currentRow), + SW.orderBy("b").rowsBetween(SW.unboundedPreceding, SW.currentRow), + ), + ( + CW.orderBy(cdf.b.desc()).rowsBetween(CW.currentRow, CW.unboundedFollowing), + SW.orderBy(sdf.b.desc()).rowsBetween(SW.currentRow, SW.unboundedFollowing), + ), + ( + CW.orderBy("b").rangeBetween(CW.currentRow, CW.currentRow), + SW.orderBy("b").rangeBetween(SW.currentRow, SW.currentRow), + ), + ( + CW.orderBy("b").rangeBetween(CW.currentRow - 1, CW.currentRow + 2), + SW.orderBy("b").rangeBetween(SW.currentRow - 1, SW.currentRow + 2), + ), + ( + CW.orderBy("b").rangeBetween(CW.unboundedPreceding, CW.currentRow), + SW.orderBy("b").rangeBetween(SW.unboundedPreceding, SW.currentRow), + ), + ( + CW.orderBy("b").rangeBetween(CW.currentRow, CW.unboundedFollowing), + SW.orderBy("b").rangeBetween(SW.currentRow, SW.unboundedFollowing), + ), + (CW.partitionBy("a").orderBy("b"), SW.partitionBy("a").orderBy("b")), + ( + CW.partitionBy(cdf.a) + .orderBy(CF.asc_nulls_last("b")) + .rowsBetween(CW.currentRow, CW.currentRow), + SW.partitionBy(sdf.a) + .orderBy(SF.asc_nulls_last("b")) + .rowsBetween(SW.currentRow, SW.currentRow), + ), + ( + CW.partitionBy("a") + .orderBy(cdf.b.desc()) + .rowsBetween(CW.currentRow - 1, CW.currentRow + 2), + SW.partitionBy("a") + .orderBy(sdf.b.desc()) + .rowsBetween(SW.currentRow - 1, SW.currentRow + 2), + ), + ( + CW.partitionBy("a") + .orderBy("b") + .rowsBetween(CW.unboundedPreceding, CW.currentRow), + SW.partitionBy("a") + .orderBy("b") + .rowsBetween(SW.unboundedPreceding, SW.currentRow), + ), + ( + CW.partitionBy("a") + .orderBy("b") + .rowsBetween(CW.currentRow, CW.unboundedFollowing), + SW.partitionBy("a") + .orderBy("b") + .rowsBetween(SW.currentRow, SW.unboundedFollowing), + ), + ( + CW.partitionBy(cdf.a) + .orderBy(cdf.b.desc(), "c") + .rangeBetween(CW.currentRow, CW.currentRow), + SW.partitionBy(sdf.a) + .orderBy(sdf.b.desc(), "c") + .rangeBetween(SW.currentRow, SW.currentRow), + ), + ( + CW.partitionBy("a") + .orderBy("b") + .rangeBetween(CW.currentRow - 1, CW.currentRow + 2), + SW.partitionBy("a") + .orderBy("b") + .rangeBetween(SW.currentRow - 1, SW.currentRow + 2), + ), + ( + CW.partitionBy("a") + .orderBy(CF.desc_nulls_last("b")) + .rangeBetween(CW.unboundedPreceding, CW.currentRow), + SW.partitionBy("a") + .orderBy(SF.desc_nulls_last("b")) + .rangeBetween(SW.unboundedPreceding, SW.currentRow), + ), + ( + CW.partitionBy("a") + .orderBy("b") + .rangeBetween(CW.currentRow, CW.unboundedFollowing), + SW.partitionBy("a") + .orderBy("b") + .rangeBetween(SW.currentRow, SW.unboundedFollowing), + ), + ]: + + self.assert_eq( + cdf.select(ccol.over(cwin)).toPandas(), + sdf.select(scol.over(swin)).toPandas(), + ) + + # check error + with self.assertRaises(PySparkValueError) as pe: + cdf.select(CF.sum("a").over(CW.orderBy("b").rowsBetween(0, (1 << 33)))).show() + + self.check_error( + exception=pe.exception, + error_class="VALUE_NOT_BETWEEN", + message_parameters={"arg_name": "end", "min": "-2147483648", "max": "2147483647"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + cdf.select(CF.rank().over(cdf.a)) + + self.check_error( + exception=pe.exception, + error_class="NOT_WINDOWSPEC", + message_parameters={"arg_name": "window", "arg_type": "Column"}, + ) + + # invalid window function + with self.assertRaises(AnalysisException): + cdf.select(cdf.b.over(CW.orderBy("b"))).show() + + # invalid window frame + # following functions require Windowframe(RowFrame, UnboundedPreceding, CurrentRow) + for ccol in [ + CF.row_number(), + CF.rank(), + CF.dense_rank(), + CF.percent_rank(), + CF.lag("c", 1), + CF.lead("c", 1), + CF.ntile(1), + ]: + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rowsBetween(CW.currentRow, CW.currentRow + 123)) + ).show() + + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rangeBetween(CW.currentRow, CW.currentRow + 123)) + ).show() + + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rangeBetween(CW.unboundedPreceding, CW.currentRow)) + ).show() + + # Function 'cume_dist' requires Windowframe(RangeFrame, UnboundedPreceding, CurrentRow) + ccol = CF.cume_dist() + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rangeBetween(CW.currentRow, CW.currentRow + 123)) + ).show() + + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rowsBetween(CW.currentRow, CW.currentRow + 123)) + ).show() + + with self.assertRaises(AnalysisException): + cdf.select( + ccol.over(CW.orderBy("b").rowsBetween(CW.unboundedPreceding, CW.currentRow)) + ).show() + + def test_window_order(self): + # SPARK-41773: test window function with order + data = [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")] + # +---+--------+ + # | id|category| + # +---+--------+ + # | 1| a| + # | 1| a| + # | 2| a| + # | 1| b| + # | 2| b| + # | 3| b| + # +---+--------+ + + cdf = self.connect.createDataFrame(data, ["id", "category"]) + sdf = self.spark.createDataFrame(data, ["id", "category"]) + + cw = CW.partitionBy("id").orderBy("category") + sw = SW.partitionBy("id").orderBy("category") + self.assert_eq( + cdf.withColumn("row_number", CF.row_number().over(cw)).toPandas(), + sdf.withColumn("row_number", SF.row_number().over(sw)).toPandas(), + ) + + cw = CW.partitionBy("category").orderBy("id") + sw = SW.partitionBy("category").orderBy("id") + self.assert_eq( + cdf.withColumn("row_number", CF.row_number().over(cw)).toPandas(), + sdf.withColumn("row_number", SF.row_number().over(sw)).toPandas(), + ) + + cw = CW.partitionBy("category").orderBy("id").rowsBetween(CW.currentRow, 1) + sw = SW.partitionBy("category").orderBy("id").rowsBetween(SW.currentRow, 1) + self.assert_eq( + cdf.withColumn("sum", CF.sum("id").over(cw)).sort("id", "category", "sum").toPandas(), + sdf.withColumn("sum", SF.sum("id").over(sw)).sort("id", "category", "sum").toPandas(), + ) + + cw = CW.partitionBy("category").orderBy("id").rangeBetween(CW.currentRow, 1) + sw = SW.partitionBy("category").orderBy("id").rangeBetween(SW.currentRow, 1) + self.assert_eq( + cdf.withColumn("sum", CF.sum("id").over(cw)).sort("id", "category").toPandas(), + sdf.withColumn("sum", SF.sum("id").over(sw)).sort("id", "category").toPandas(), + ) + + def test_collection_functions(self): + query = """ + SELECT * FROM VALUES + (ARRAY('a', 'ab'), ARRAY(1, 2, 3), ARRAY(1, NULL, 3), 1, 2, 'a'), + (ARRAY('x', NULL), NULL, ARRAY(1, 3), 3, 4, 'x'), + (NULL, ARRAY(-1, -2, -3), Array(), 5, 6, NULL) + AS tab(a, b, c, d, e, f) + """ + # +---------+------------+------------+---+---+----+ + # | a| b| c| d| e| f| + # +---------+------------+------------+---+---+----+ + # | [a, ab]| [1, 2, 3]|[1, null, 3]| 1| 2| a| + # |[x, null]| NULL| [1, 3]| 3| 4| x| + # | NULL|[-1, -2, -3]| []| 5| 6|NULL| + # +---------+------------+------------+---+---+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + for cfunc, sfunc in [ + (CF.array_distinct, SF.array_distinct), + (CF.array_compact, SF.array_compact), + (CF.array_max, SF.array_max), + (CF.array_min, SF.array_min), + (CF.reverse, SF.reverse), + (CF.size, SF.size), + ]: + self.assert_eq( + cdf.select(cfunc("a"), cfunc(cdf.b)).toPandas(), + sdf.select(sfunc("a"), sfunc(sdf.b)).toPandas(), + ) + + for cfunc, sfunc in [ + (CF.array_except, SF.array_except), + (CF.array_intersect, SF.array_intersect), + (CF.array_union, SF.array_union), + (CF.arrays_overlap, SF.arrays_overlap), + ]: + self.assert_eq( + cdf.select(cfunc("b", cdf.c)).toPandas(), + sdf.select(sfunc("b", sdf.c)).toPandas(), + ) + + for cfunc, sfunc in [ + (CF.array_position, SF.array_position), + (CF.array_remove, SF.array_remove), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.a, "ab")).toPandas(), + sdf.select(sfunc(sdf.a, "ab")).toPandas(), + ) + + # test array + self.assert_eq( + cdf.select(CF.array(cdf.d, "e")).toPandas(), + sdf.select(SF.array(sdf.d, "e")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array(cdf.d, "e", CF.lit(99))).toPandas(), + sdf.select(SF.array(sdf.d, "e", SF.lit(99))).toPandas(), + ) + + # test array_contains + self.assert_eq( + cdf.select(CF.array_contains(cdf.a, "ab")).toPandas(), + sdf.select(SF.array_contains(sdf.a, "ab")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_contains(cdf.a, cdf.f)).toPandas(), + sdf.select(SF.array_contains(sdf.a, sdf.f)).toPandas(), + ) + + # test array_append + self.assert_eq( + cdf.select(CF.array_append(cdf.a, "xyz")).toPandas(), + sdf.select(SF.array_append(sdf.a, "xyz")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_append(cdf.a, CF.lit("ab"))).toPandas(), + sdf.select(SF.array_append(sdf.a, SF.lit("ab"))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_append(cdf.a, cdf.f)).toPandas(), + sdf.select(SF.array_append(sdf.a, sdf.f)).toPandas(), + ) + + # test array_prepend + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, "xyz")).toPandas(), + sdf.select(SF.array_prepend(sdf.a, "xyz")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, CF.lit("ab"))).toPandas(), + sdf.select(SF.array_prepend(sdf.a, SF.lit("ab"))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, cdf.f)).toPandas(), + sdf.select(SF.array_prepend(sdf.a, sdf.f)).toPandas(), + ) + + # test array_insert + self.assert_eq( + cdf.select(CF.array_insert(cdf.a, -5, "ab")).toPandas(), + sdf.select(SF.array_insert(sdf.a, -5, "ab")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_insert(cdf.a, 3, cdf.f)).toPandas(), + sdf.select(SF.array_insert(sdf.a, 3, sdf.f)).toPandas(), + ) + + # test array_join + self.assert_eq( + cdf.select( + CF.array_join(cdf.a, ","), CF.array_join("b", ":"), CF.array_join("c", "~") + ).toPandas(), + sdf.select( + SF.array_join(sdf.a, ","), SF.array_join("b", ":"), SF.array_join("c", "~") + ).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.array_join(cdf.a, ",", "_null_"), + CF.array_join("b", ":", ".null."), + CF.array_join("c", "~", "NULL"), + ).toPandas(), + sdf.select( + SF.array_join(sdf.a, ",", "_null_"), + SF.array_join("b", ":", ".null."), + SF.array_join("c", "~", "NULL"), + ).toPandas(), + ) + + # test array_repeat + self.assert_eq( + cdf.select(CF.array_repeat(cdf.f, "d")).toPandas(), + sdf.select(SF.array_repeat(sdf.f, "d")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_repeat("f", cdf.d)).toPandas(), + sdf.select(SF.array_repeat("f", sdf.d)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_repeat("f", 3)).toPandas(), + sdf.select(SF.array_repeat("f", 3)).toPandas(), + ) + + # test arrays_zip + # TODO: Make toPandas support complex nested types like Array + # DataFrame.iloc[:, 0] (column name="arrays_zip(b, c)") values are different (66.66667 %) + # [index]: [0, 1, 2] + # [left]: [[{'b': 1, 'c': 1.0}, {'b': 2, 'c': None}, {'b': 3, 'c': 3.0}], None, + # [{'b': -1, 'c': None}, {'b': -2, 'c': None}, {'b': -3, 'c': None}]] + # [right]: [[(1, 1), (2, None), (3, 3)], None, [(-1, None), (-2, None), (-3, None)]] + self.compare_by_show( + cdf.select(CF.arrays_zip(cdf.b, "c")), + sdf.select(SF.arrays_zip(sdf.b, "c")), + ) + + # test concat + self.assert_eq( + cdf.select(CF.concat("d", cdf.e, CF.lit(-1))).toPandas(), + sdf.select(SF.concat("d", sdf.e, SF.lit(-1))).toPandas(), + ) + + # test create_map + self.compare_by_show( + cdf.select(CF.create_map(cdf.d, cdf.e)), sdf.select(SF.create_map(sdf.d, sdf.e)) + ) + self.compare_by_show( + cdf.select(CF.create_map(cdf.d, "e", "e", CF.lit(1))), + sdf.select(SF.create_map(sdf.d, "e", "e", SF.lit(1))), + ) + + # test element_at + self.assert_eq( + cdf.select(CF.element_at("a", 1)).toPandas(), + sdf.select(SF.element_at("a", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.element_at(cdf.a, 1)).toPandas(), + sdf.select(SF.element_at(sdf.a, 1)).toPandas(), + ) + + # test get + self.assert_eq( + cdf.select(CF.get("a", 1)).toPandas(), + sdf.select(SF.get("a", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.get(cdf.a, 1)).toPandas(), + sdf.select(SF.get(sdf.a, 1)).toPandas(), + ) + + # test shuffle + # Can not compare the values due to the random permutation + self.assertEqual( + cdf.select(CF.shuffle(cdf.a), CF.shuffle("b")).count(), + sdf.select(SF.shuffle(sdf.a), SF.shuffle("b")).count(), + ) + + # test slice + self.assert_eq( + cdf.select(CF.slice(cdf.a, 1, 2), CF.slice("c", 2, 3)).toPandas(), + sdf.select(SF.slice(sdf.a, 1, 2), SF.slice("c", 2, 3)).toPandas(), + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.slice(cdf.a, 1.0, 2) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "start", "arg_type": "float"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.slice(cdf.a, 1, 2.0) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "length", "arg_type": "float"}, + ) + + # test sort_array + self.assert_eq( + cdf.select(CF.sort_array(cdf.a, True), CF.sort_array("c", False)).toPandas(), + sdf.select(SF.sort_array(sdf.a, True), SF.sort_array("c", False)).toPandas(), + ) + + # test struct + self.compare_by_show( + cdf.select(CF.struct(cdf.a, "d", "e", cdf.f)), + sdf.select(SF.struct(sdf.a, "d", "e", sdf.f)), + ) + + # test sequence + self.assert_eq( + cdf.select(CF.sequence(CF.lit(1), CF.lit(5))).toPandas(), + sdf.select(SF.sequence(SF.lit(1), SF.lit(5))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.sequence(CF.lit(1), CF.lit(5), CF.lit(1))).toPandas(), + sdf.select(SF.sequence(SF.lit(1), SF.lit(5), SF.lit(1))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.sequence(cdf.d, "e")).toPandas(), + sdf.select(SF.sequence(sdf.d, "e")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.sequence(cdf.d, "e", CF.lit(1))).toPandas(), + sdf.select(SF.sequence(sdf.d, "e", SF.lit(1))).toPandas(), + ) + + def test_map_collection_functions(self): + query = """ + SELECT * FROM VALUES + (MAP('a', 'ab'), MAP('x', 'ab'), MAP(1, 2, 3, 4), 1, 'a', ARRAY(1, 2), ARRAY('X', 'Y')), + (MAP('x', 'yz'), MAP('c', NULL), NULL, 2, 'x', ARRAY(3, 4), ARRAY('A', 'B')), + (MAP('c', 'de'), NULL, MAP(-1, NULL, -3, -4), -3, 'c', NULL, ARRAY('Z')) + AS tab(a, b, c, e, f, g, h) + """ + # +---------+-----------+----------------------+---+---+------+------+ + # | a| b| c| e| f| g| h| + # +---------+-----------+----------------------+---+---+------+------+ + # |{a -> ab}| {x -> ab}| {1 -> 2, 3 -> 4}| 1| a|[1, 2]|[X, Y]| + # |{x -> yz}|{c -> null}| NULL| 2| x|[3, 4]|[A, B]| + # |{c -> de}| NULL|{-1 -> null, -3 -> -4}| -3| c| NULL| [Z]| + # +---------+-----------+----------------------+---+---+------+------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test map_concat + self.compare_by_show( + cdf.select(CF.map_concat(cdf.a, "b")), + sdf.select(SF.map_concat(sdf.a, "b")), + ) + + # test map_contains_key + self.compare_by_show( + cdf.select(CF.map_contains_key(cdf.a, "a"), CF.map_contains_key("c", 3)), + sdf.select(SF.map_contains_key(sdf.a, "a"), SF.map_contains_key("c", 3)), + ) + + # test map_entries + self.compare_by_show( + cdf.select(CF.map_entries(cdf.a), CF.map_entries("b")), + sdf.select(SF.map_entries(sdf.a), SF.map_entries("b")), + ) + + # test map_from_arrays + self.compare_by_show( + cdf.select(CF.map_from_arrays(cdf.g, "h")), + sdf.select(SF.map_from_arrays(sdf.g, "h")), + ) + + # test map_keys and map_values + self.compare_by_show( + cdf.select(CF.map_keys(cdf.a), CF.map_values("b")), + sdf.select(SF.map_keys(sdf.a), SF.map_values("b")), + ) + + # test size + self.assert_eq( + cdf.select(CF.size(cdf.a), CF.size("c")).toPandas(), + sdf.select(SF.size(sdf.a), SF.size("c")).toPandas(), + ) + + def test_generator_functions(self): + query = """ + SELECT * FROM VALUES + (ARRAY('a', 'ab'), ARRAY(1, 2, 3), ARRAY(1, NULL, 3), + MAP(1, 2, 3, 4), 1, FLOAT(2.0), 3), + (ARRAY('x', NULL), NULL, ARRAY(1, 3), + NULL, 3, FLOAT(4.0), 5), + (NULL, ARRAY(-1, -2, -3), Array(), + MAP(-1, NULL, -3, -4), 7, FLOAT('NAN'), 9) + AS tab(a, b, c, d, e, f, g) + """ + # +---------+------------+------------+----------------------+---+---+---+ + # | a| b| c| d| e| f| g| + # +---------+------------+------------+----------------------+---+---+---+ + # | [a, ab]| [1, 2, 3]|[1, null, 3]| {1 -> 2, 3 -> 4}| 1|2.0| 3| + # |[x, null]| NULL| [1, 3]| NULL| 3|4.0| 5| + # | NULL|[-1, -2, -3]| []|{-1 -> null, -3 -> -4}| 7|NaN| 9| + # +---------+------------+------------+----------------------+---+---+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test explode with arrays + self.assert_eq( + cdf.select(CF.explode(cdf.a), CF.col("b")).toPandas(), + sdf.select(SF.explode(sdf.a), SF.col("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.explode("a"), "b").toPandas(), + sdf.select(SF.explode("a"), "b").toPandas(), + ) + # test explode with maps + self.assert_eq( + cdf.select(CF.explode(cdf.d), CF.col("c")).toPandas(), + sdf.select(SF.explode(sdf.d), SF.col("c")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.explode("d"), "c").toPandas(), + sdf.select(SF.explode("d"), "c").toPandas(), + ) + + # test explode_outer with arrays + self.assert_eq( + cdf.select(CF.explode_outer(cdf.a), CF.col("b")).toPandas(), + sdf.select(SF.explode_outer(sdf.a), SF.col("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.explode_outer("a"), "b").toPandas(), + sdf.select(SF.explode_outer("a"), "b").toPandas(), + ) + # test explode_outer with maps + self.assert_eq( + cdf.select(CF.explode_outer(cdf.d), CF.col("c")).toPandas(), + sdf.select(SF.explode_outer(sdf.d), SF.col("c")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.explode_outer("d"), "c").toPandas(), + sdf.select(SF.explode_outer("d"), "c").toPandas(), + ) + + # test flatten + self.assert_eq( + cdf.select(CF.flatten(CF.array("b", cdf.c)), CF.col("b")).toPandas(), + sdf.select(SF.flatten(SF.array("b", sdf.c)), SF.col("b")).toPandas(), + ) + + # test inline + self.assert_eq( + cdf.select(CF.expr("ARRAY(STRUCT(e, f), STRUCT(g AS e, f))").alias("X")) + .select(CF.inline("X")) + .toPandas(), + sdf.select(SF.expr("ARRAY(STRUCT(e, f), STRUCT(g AS e, f))").alias("X")) + .select(SF.inline("X")) + .toPandas(), + ) + + # test inline_outer + self.assert_eq( + cdf.select(CF.expr("ARRAY(STRUCT(e, f), STRUCT(g AS e, f))").alias("X")) + .select(CF.inline_outer("X")) + .toPandas(), + sdf.select(SF.expr("ARRAY(STRUCT(e, f), STRUCT(g AS e, f))").alias("X")) + .select(SF.inline_outer("X")) + .toPandas(), + ) + + # test posexplode with arrays + self.assert_eq( + cdf.select(CF.posexplode(cdf.a), CF.col("b")).toPandas(), + sdf.select(SF.posexplode(sdf.a), SF.col("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.posexplode("a"), "b").toPandas(), + sdf.select(SF.posexplode("a"), "b").toPandas(), + ) + # test posexplode with maps + self.assert_eq( + cdf.select(CF.posexplode(cdf.d), CF.col("c")).toPandas(), + sdf.select(SF.posexplode(sdf.d), SF.col("c")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.posexplode("d"), "c").toPandas(), + sdf.select(SF.posexplode("d"), "c").toPandas(), + ) + + # test posexplode_outer with arrays + self.assert_eq( + cdf.select(CF.posexplode_outer(cdf.a), CF.col("b")).toPandas(), + sdf.select(SF.posexplode_outer(sdf.a), SF.col("b")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.posexplode_outer("a"), "b").toPandas(), + sdf.select(SF.posexplode_outer("a"), "b").toPandas(), + ) + # test posexplode_outer with maps + self.assert_eq( + cdf.select(CF.posexplode_outer(cdf.d), CF.col("c")).toPandas(), + sdf.select(SF.posexplode_outer(sdf.d), SF.col("c")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.posexplode_outer("d"), "c").toPandas(), + sdf.select(SF.posexplode_outer("d"), "c").toPandas(), + ) + + def test_lambda_functions(self): + query = """ + SELECT * FROM VALUES + (ARRAY('a', 'ab'), ARRAY(1, 2, 3), ARRAY(1, NULL, 3), 1, 2, 'a', NULL, MAP(0, 0)), + (ARRAY('x', NULL), NULL, ARRAY(1, 3), 3, 4, 'x', MAP(2, 0), MAP(-1, 1)), + (NULL, ARRAY(-1, -2, -3), Array(), 5, 6, NULL, MAP(-1, 2, -3, -4), NULL) + AS tab(a, b, c, d, e, f, g, h) + """ + # +---------+------------+------------+---+---+----+-------------------+---------+ + # | a| b| c| d| e| f| g| h| + # +---------+------------+------------+---+---+----+-------------------+---------+ + # | [a, ab]| [1, 2, 3]|[1, null, 3]| 1| 2| a| NULL| {0 -> 0}| + # |[x, null]| NULL| [1, 3]| 3| 4| x| {2 -> 0}|{-1 -> 1}| + # | NULL|[-1, -2, -3]| []| 5| 6|NULL|{-1 -> 2, -3 -> -4}| NULL| + # +---------+------------+------------+---+---+----+-------------------+---------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test exists + self.assert_eq( + cdf.select(CF.exists(cdf.b, lambda x: x < 0)).toPandas(), + sdf.select(SF.exists(sdf.b, lambda x: x < 0)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.exists("a", lambda x: CF.isnull(x))).toPandas(), + sdf.select(SF.exists("a", lambda x: SF.isnull(x))).toPandas(), + ) + + # test aggregate + # aggregate without finish + self.assert_eq( + cdf.select(CF.aggregate(cdf.b, "d", lambda acc, x: acc + x)).toPandas(), + sdf.select(SF.aggregate(sdf.b, "d", lambda acc, x: acc + x)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.aggregate("b", cdf.d, lambda acc, x: acc + x)).toPandas(), + sdf.select(SF.aggregate("b", sdf.d, lambda acc, x: acc + x)).toPandas(), + ) + + # aggregate with finish + self.assert_eq( + cdf.select( + CF.aggregate(cdf.b, "d", lambda acc, x: acc + x, lambda acc: acc + 100) + ).toPandas(), + sdf.select( + SF.aggregate(sdf.b, "d", lambda acc, x: acc + x, lambda acc: acc + 100) + ).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.aggregate("b", cdf.d, lambda acc, x: acc + x, lambda acc: acc + 100) + ).toPandas(), + sdf.select( + SF.aggregate("b", sdf.d, lambda acc, x: acc + x, lambda acc: acc + 100) + ).toPandas(), + ) + + # test array_sort + self.assert_eq( + cdf.select(CF.array_sort(cdf.b, lambda x, y: CF.abs(x) - CF.abs(y))).toPandas(), + sdf.select(SF.array_sort(sdf.b, lambda x, y: SF.abs(x) - SF.abs(y))).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.array_sort( + "a", + lambda x, y: CF.when(x.isNull() | y.isNull(), CF.lit(0)).otherwise( + CF.length(y) - CF.length(x) + ), + ) + ).toPandas(), + sdf.select( + SF.array_sort( + "a", + lambda x, y: SF.when(x.isNull() | y.isNull(), SF.lit(0)).otherwise( + SF.length(y) - SF.length(x) + ), + ) + ).toPandas(), + ) + + # test filter + self.assert_eq( + cdf.select(CF.filter(cdf.b, lambda x: x < 0)).toPandas(), + sdf.select(SF.filter(sdf.b, lambda x: x < 0)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.filter("a", lambda x: ~CF.isnull(x))).toPandas(), + sdf.select(SF.filter("a", lambda x: ~SF.isnull(x))).toPandas(), + ) + + # test forall + self.assert_eq( + cdf.select(CF.filter(cdf.b, lambda x: x != 0)).toPandas(), + sdf.select(SF.filter(sdf.b, lambda x: x != 0)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.filter("a", lambda x: ~CF.isnull(x))).toPandas(), + sdf.select(SF.filter("a", lambda x: ~SF.isnull(x))).toPandas(), + ) + + # test transform + # transform without index + self.assert_eq( + cdf.select(CF.transform(cdf.b, lambda x: x + 1)).toPandas(), + sdf.select(SF.transform(sdf.b, lambda x: x + 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.transform("b", lambda x: x + 1)).toPandas(), + sdf.select(SF.transform("b", lambda x: x + 1)).toPandas(), + ) + + # transform with index + self.assert_eq( + cdf.select(CF.transform(cdf.b, lambda x, i: x + 1 - i)).toPandas(), + sdf.select(SF.transform(sdf.b, lambda x, i: x + 1 - i)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.transform("b", lambda x, i: x + 1 - i)).toPandas(), + sdf.select(SF.transform("b", lambda x, i: x + 1 - i)).toPandas(), + ) + + # test zip_with + self.assert_eq( + cdf.select(CF.zip_with(cdf.b, "c", lambda v1, v2: v1 - CF.abs(v2))).toPandas(), + sdf.select(SF.zip_with(sdf.b, "c", lambda v1, v2: v1 - SF.abs(v2))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.zip_with("b", cdf.c, lambda v1, v2: v1 - CF.abs(v2))).toPandas(), + sdf.select(SF.zip_with("b", sdf.c, lambda v1, v2: v1 - SF.abs(v2))).toPandas(), + ) + + # test map_filter + self.compare_by_show( + cdf.select(CF.map_filter(cdf.g, lambda k, v: k > v)), + sdf.select(SF.map_filter(sdf.g, lambda k, v: k > v)), + ) + self.compare_by_show( + cdf.select(CF.map_filter("g", lambda k, v: k > v)), + sdf.select(SF.map_filter("g", lambda k, v: k > v)), + ) + + # test map_zip_with + self.compare_by_show( + cdf.select(CF.map_zip_with(cdf.g, "h", lambda k, v1, v2: v1 + v2)), + sdf.select(SF.map_zip_with(sdf.g, "h", lambda k, v1, v2: v1 + v2)), + ) + self.compare_by_show( + cdf.select(CF.map_zip_with("g", cdf.h, lambda k, v1, v2: v1 + v2)), + sdf.select(SF.map_zip_with("g", sdf.h, lambda k, v1, v2: v1 + v2)), + ) + + # test transform_keys + self.compare_by_show( + cdf.select(CF.transform_keys(cdf.g, lambda k, v: k - 1)), + sdf.select(SF.transform_keys(sdf.g, lambda k, v: k - 1)), + ) + self.compare_by_show( + cdf.select(CF.transform_keys("g", lambda k, v: k - 1)), + sdf.select(SF.transform_keys("g", lambda k, v: k - 1)), + ) + + # test transform_values + self.compare_by_show( + cdf.select(CF.transform_values(cdf.g, lambda k, v: CF.abs(v) + 1)), + sdf.select(SF.transform_values(sdf.g, lambda k, v: SF.abs(v) + 1)), + ) + self.compare_by_show( + cdf.select(CF.transform_values("g", lambda k, v: CF.abs(v) + 1)), + sdf.select(SF.transform_values("g", lambda k, v: SF.abs(v) + 1)), + ) + + def test_nested_lambda_function(self): + # SPARK-42089: test nested lambda function + query = "SELECT array(1, 2, 3) as numbers, array('a', 'b', 'c') as letters" + + cdf = self.connect.sql(query).select( + CF.flatten( + CF.transform( + "numbers", + lambda number: CF.transform( + "letters", lambda letter: CF.struct(number.alias("n"), letter.alias("l")) + ), + ) + ) + ) + + sdf = self.spark.sql(query).select( + SF.flatten( + SF.transform( + "numbers", + lambda number: SF.transform( + "letters", lambda letter: SF.struct(number.alias("n"), letter.alias("l")) + ), + ) + ) + ) + + # TODO: 'cdf.schema' has an extra metadata '{'__autoGeneratedAlias': 'true'}' + # self.assertEqual(cdf.schema, sdf.schema) + self.assertEqual(cdf.collect(), sdf.collect()) + + def test_csv_functions(self): + query = """ + SELECT * FROM VALUES + ('1,2,3', 'a,b,5.0'), + ('3,4,5', 'x,y,6.0') + AS tab(a, b) + """ + # +-----+-------+ + # | a| b| + # +-----+-------+ + # |1,2,3|a,b,5.0| + # |3,4,5|x,y,6.0| + # +-----+-------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test from_csv + self.compare_by_show( + cdf.select( + CF.from_csv(cdf.a, "a INT, b INT, c INT"), + CF.from_csv("b", "x STRING, y STRING, z DOUBLE"), + ), + sdf.select( + SF.from_csv(sdf.a, "a INT, b INT, c INT"), + SF.from_csv("b", "x STRING, y STRING, z DOUBLE"), + ), + ) + self.compare_by_show( + cdf.select( + CF.from_csv(cdf.a, CF.lit("a INT, b INT, c INT")), + CF.from_csv("b", CF.lit("x STRING, y STRING, z DOUBLE")), + ), + sdf.select( + SF.from_csv(sdf.a, SF.lit("a INT, b INT, c INT")), + SF.from_csv("b", SF.lit("x STRING, y STRING, z DOUBLE")), + ), + ) + self.compare_by_show( + cdf.select( + CF.from_csv(cdf.a, CF.lit("a INT, b INT, c INT"), {"maxCharsPerColumn": "3"}), + CF.from_csv( + "b", CF.lit("x STRING, y STRING, z DOUBLE"), {"maxCharsPerColumn": "3"} + ), + ), + sdf.select( + SF.from_csv(sdf.a, SF.lit("a INT, b INT, c INT"), {"maxCharsPerColumn": "3"}), + SF.from_csv( + "b", SF.lit("x STRING, y STRING, z DOUBLE"), {"maxCharsPerColumn": "3"} + ), + ), + ) + + # test schema_of_csv + self.assert_eq( + cdf.select(CF.schema_of_csv(CF.lit('{"a": 0}'))).toPandas(), + sdf.select(SF.schema_of_csv(SF.lit('{"a": 0}'))).toPandas(), + ) + self.assert_eq( + cdf.select( + CF.schema_of_csv(CF.lit('{"a": 0}'), {"maxCharsPerColumn": "10"}) + ).toPandas(), + sdf.select( + SF.schema_of_csv(SF.lit('{"a": 0}'), {"maxCharsPerColumn": "10"}) + ).toPandas(), + ) + + # test to_csv + self.compare_by_show( + cdf.select(CF.to_csv(CF.struct(CF.lit("a"), CF.lit("b")))), + sdf.select(SF.to_csv(SF.struct(SF.lit("a"), SF.lit("b")))), + ) + self.compare_by_show( + cdf.select(CF.to_csv(CF.struct(CF.lit("a"), CF.lit("b")), {"maxCharsPerColumn": "10"})), + sdf.select(SF.to_csv(SF.struct(SF.lit("a"), SF.lit("b")), {"maxCharsPerColumn": "10"})), + ) + + def test_json_functions(self): + query = """ + SELECT * FROM VALUES + ('{"a": 1}', '[1, 2, 3]', '{"f1": "value1", "f2": "value2"}'), + ('{"a": 0}', '[4, 5, 6]', '{"f1": "value12"}') + AS tab(a, b, c) + """ + # +--------+---------+--------------------------------+ + # | a| b| c| + # +--------+---------+--------------------------------+ + # |{"a": 1}|[1, 2, 3]|{"f1": "value1", "f2": "value2"}| + # |{"a": 0}|[4, 5, 6]| {"f1": "value12"}| + # +--------+---------+--------------------------------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test from_json + for schema in [ + "a INT", + "MAP", + StructType([StructField("a", IntegerType())]), + ArrayType(StructType([StructField("a", IntegerType())])), + ]: + self.compare_by_show( + cdf.select(CF.from_json(cdf.a, schema)), + sdf.select(SF.from_json(sdf.a, schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json("a", schema)), + sdf.select(SF.from_json("a", schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json(cdf.a, schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json(sdf.a, schema, {"mode": "FAILFAST"})), + ) + self.compare_by_show( + cdf.select(CF.from_json("a", schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json("a", schema, {"mode": "FAILFAST"})), + ) + + for schema in [ + "ARRAY", + ArrayType(IntegerType()), + ]: + self.compare_by_show( + cdf.select(CF.from_json(cdf.b, schema)), + sdf.select(SF.from_json(sdf.b, schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json("b", schema)), + sdf.select(SF.from_json("b", schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json(cdf.b, schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json(sdf.b, schema, {"mode": "FAILFAST"})), + ) + self.compare_by_show( + cdf.select(CF.from_json("b", schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json("b", schema, {"mode": "FAILFAST"})), + ) + + # SPARK-41880: from_json support non-literal expression + c_schema = CF.schema_of_json(CF.lit("""{"a": 2}""")) + s_schema = SF.schema_of_json(SF.lit("""{"a": 2}""")) + + self.compare_by_show( + cdf.select(CF.from_json(cdf.a, c_schema)), + sdf.select(SF.from_json(sdf.a, s_schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json("a", c_schema)), + sdf.select(SF.from_json("a", s_schema)), + ) + self.compare_by_show( + cdf.select(CF.from_json(cdf.a, c_schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json(sdf.a, s_schema, {"mode": "FAILFAST"})), + ) + self.compare_by_show( + cdf.select(CF.from_json("a", c_schema, {"mode": "FAILFAST"})), + sdf.select(SF.from_json("a", s_schema, {"mode": "FAILFAST"})), + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.from_json("a", [c_schema]) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_DATATYPE_OR_STR", + message_parameters={"arg_name": "schema", "arg_type": "list"}, + ) + + # test get_json_object + self.assert_eq( + cdf.select( + CF.get_json_object("c", "$.f1"), + CF.get_json_object(cdf.c, "$.f2"), + ).toPandas(), + sdf.select( + SF.get_json_object("c", "$.f1"), + SF.get_json_object(sdf.c, "$.f2"), + ).toPandas(), + ) + + # test json_tuple + self.assert_eq( + cdf.select(CF.json_tuple("c", "f1", "f2")).toPandas(), + sdf.select(SF.json_tuple("c", "f1", "f2")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.json_tuple(cdf.c, "f1", "f2")).toPandas(), + sdf.select(SF.json_tuple(sdf.c, "f1", "f2")).toPandas(), + ) + + # test schema_of_json + self.assert_eq( + cdf.select(CF.schema_of_json(CF.lit('{"a": 0}'))).toPandas(), + sdf.select(SF.schema_of_json(SF.lit('{"a": 0}'))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.schema_of_json(CF.lit('{"a": 0}'), {"mode": "FAILFAST"})).toPandas(), + sdf.select(SF.schema_of_json(SF.lit('{"a": 0}'), {"mode": "FAILFAST"})).toPandas(), + ) + + # test to_json + self.compare_by_show( + cdf.select(CF.to_json(CF.struct(CF.lit("a"), CF.lit("b")))), + sdf.select(SF.to_json(SF.struct(SF.lit("a"), SF.lit("b")))), + ) + self.compare_by_show( + cdf.select(CF.to_json(CF.struct(CF.lit("a"), CF.lit("b")), {"mode": "FAILFAST"})), + sdf.select(SF.to_json(SF.struct(SF.lit("a"), SF.lit("b")), {"mode": "FAILFAST"})), + ) + + def test_string_functions_one_arg(self): + query = """ + SELECT * FROM VALUES + (' ab ', 'ab ', NULL), (' ab', NULL, 'ab') + AS tab(a, b, c) + """ + # +--------+-----+----+ + # | a| b| c| + # +--------+-----+----+ + # | ab |ab |NULL| + # | ab| NULL| ab| + # +--------+-----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + for cfunc, sfunc in [ + (CF.upper, SF.upper), + (CF.lower, SF.lower), + (CF.ascii, SF.ascii), + (CF.base64, SF.base64), + (CF.unbase64, SF.unbase64), + (CF.ltrim, SF.ltrim), + (CF.rtrim, SF.rtrim), + (CF.trim, SF.trim), + (CF.sentences, SF.sentences), + (CF.initcap, SF.initcap), + (CF.soundex, SF.soundex), + (CF.bin, SF.bin), + (CF.hex, SF.hex), + (CF.unhex, SF.unhex), + (CF.length, SF.length), + (CF.octet_length, SF.octet_length), + (CF.bit_length, SF.bit_length), + (CF.reverse, SF.reverse), + ]: + self.assert_eq( + cdf.select(cfunc("a"), cfunc(cdf.b)).toPandas(), + sdf.select(sfunc("a"), sfunc(sdf.b)).toPandas(), + ) + + def test_string_functions_multi_args(self): + query = """ + SELECT * FROM VALUES + (1, 'abcdef', 'ghij', 'hello world', 'a.b.c.d'), + (2, 'abcd', 'efghij', 'how are you', 'a.b.c') + AS tab(a, b, c, d, e) + """ + # +---+------+------+-----------+-------+ + # | a| b| c| d| e| + # +---+------+------+-----------+-------+ + # | 1|abcdef| ghij|hello world|a.b.c.d| + # | 2| abcd|efghij|how are you| a.b.c| + # +---+------+------+-----------+-------+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select(CF.format_number(cdf.a, 2)).toPandas(), + sdf.select(SF.format_number(sdf.a, 2)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.format_number("a", 5)).toPandas(), + sdf.select(SF.format_number("a", 5)).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.concat_ws("-", cdf.b, "c")).toPandas(), + sdf.select(SF.concat_ws("-", sdf.b, "c")).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.decode("c", "UTF-8")).toPandas(), + sdf.select(SF.decode("c", "UTF-8")).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.encode("c", "UTF-8")).toPandas(), + sdf.select(SF.encode("c", "UTF-8")).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.format_string("%d %s", cdf.a, cdf.b)).toPandas(), + sdf.select(SF.format_string("%d %s", sdf.a, sdf.b)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.instr(cdf.b, "b")).toPandas(), sdf.select(SF.instr(sdf.b, "b")).toPandas() + ) + self.assert_eq( + cdf.select(CF.overlay(cdf.b, cdf.c, 2)).toPandas(), + sdf.select(SF.overlay(sdf.b, sdf.c, 2)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.substring(cdf.b, 1, 2)).toPandas(), + sdf.select(SF.substring(sdf.b, 1, 2)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.substring_index(cdf.e, ".", 2)).toPandas(), + sdf.select(SF.substring_index(sdf.e, ".", 2)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.levenshtein(cdf.b, cdf.c)).toPandas(), + sdf.select(SF.levenshtein(sdf.b, sdf.c)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.levenshtein(cdf.b, cdf.c, 1)).toPandas(), + sdf.select(SF.levenshtein(sdf.b, sdf.c, 1)).toPandas(), + ) + + self.assert_eq( + cdf.select(CF.locate("e", cdf.b)).toPandas(), + sdf.select(SF.locate("e", sdf.b)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.lpad(cdf.b, 10, "#")).toPandas(), + sdf.select(SF.lpad(sdf.b, 10, "#")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.rpad(cdf.b, 10, "#")).toPandas(), + sdf.select(SF.rpad(sdf.b, 10, "#")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.repeat(cdf.b, 2)).toPandas(), sdf.select(SF.repeat(sdf.b, 2)).toPandas() + ) + self.assert_eq( + cdf.select(CF.split(cdf.b, "[bd]")).toPandas(), + sdf.select(SF.split(sdf.b, "[bd]")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.regexp_extract(cdf.b, "(a+)(b)?(c)", 1)).toPandas(), + sdf.select(SF.regexp_extract(sdf.b, "(a+)(b)?(c)", 1)).toPandas(), + ) + self.assert_eq( + cdf.select(CF.regexp_replace(cdf.b, "(a+)(b)?(c)", "--")).toPandas(), + sdf.select(SF.regexp_replace(sdf.b, "(a+)(b)?(c)", "--")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.translate(cdf.b, "abc", "xyz")).toPandas(), + sdf.select(SF.translate(sdf.b, "abc", "xyz")).toPandas(), + ) + + # TODO(SPARK-41283): To compare toPandas for test cases with dtypes marked + def test_date_ts_functions(self): + query = """ + SELECT * FROM VALUES + ('1997/02/28 10:30:00', '2023/03/01 06:00:00', 'JST', 1428476400, 2020, 12, 6), + ('2000/01/01 04:30:05', '2020/05/01 12:15:00', 'PST', 1403892395, 2022, 12, 6) + AS tab(ts1, ts2, tz, seconds, Y, M, D) + """ + # +-------------------+-------------------+---+----------+----+---+---+ + # | ts1| ts2| tz| seconds| Y| M| D| + # +-------------------+-------------------+---+----------+----+---+---+ + # |1997/02/28 10:30:00|2023/03/01 06:00:00|JST|1428476400|2020| 12| 6| + # |2000/01/01 04:30:05|2020/05/01 12:15:00|PST|1403892395|2022| 12| 6| + # +-------------------+-------------------+---+----------+----+---+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # With no parameters + for cfunc, sfunc in [ + (CF.current_date, SF.current_date), + ]: + self.assert_eq( + cdf.select(cfunc()).toPandas(), + sdf.select(sfunc()).toPandas(), + ) + + # current_timestamp + # [left]: datetime64[ns, America/Los_Angeles] + # [right]: datetime64[ns] + # TODO: compare the return values after resolving dtypes difference + self.assertEqual( + cdf.select(CF.current_timestamp()).count(), + sdf.select(SF.current_timestamp()).count(), + ) + + # localtimestamp + s_pdf0 = sdf.select(SF.localtimestamp()).toPandas() + c_pdf = cdf.select(CF.localtimestamp()).toPandas() + s_pdf1 = sdf.select(SF.localtimestamp()).toPandas() + self.assert_eq(s_pdf0 < c_pdf, c_pdf < s_pdf1) + + # With only column parameter + for cfunc, sfunc in [ + (CF.year, SF.year), + (CF.quarter, SF.quarter), + (CF.month, SF.month), + (CF.dayofweek, SF.dayofweek), + (CF.dayofmonth, SF.dayofmonth), + (CF.dayofyear, SF.dayofyear), + (CF.hour, SF.hour), + (CF.minute, SF.minute), + (CF.second, SF.second), + (CF.weekofyear, SF.weekofyear), + (CF.last_day, SF.last_day), + (CF.unix_timestamp, SF.unix_timestamp), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.ts1)).toPandas(), + sdf.select(sfunc(sdf.ts1)).toPandas(), + ) + + # With format parameter + for cfunc, sfunc in [ + (CF.date_format, SF.date_format), + (CF.to_date, SF.to_date), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.ts1, format="yyyy-MM-dd")).toPandas(), + sdf.select(sfunc(sdf.ts1, format="yyyy-MM-dd")).toPandas(), + ) + self.compare_by_show( + # [left]: datetime64[ns, America/Los_Angeles] + # [right]: datetime64[ns] + cdf.select(CF.to_timestamp(cdf.ts1, format="yyyy-MM-dd")), + sdf.select(SF.to_timestamp(sdf.ts1, format="yyyy-MM-dd")), + ) + + # With tz parameter + for cfunc, sfunc in [ + (CF.from_utc_timestamp, SF.from_utc_timestamp), + (CF.to_utc_timestamp, SF.to_utc_timestamp), + # [left]: datetime64[ns, America/Los_Angeles] + # [right]: datetime64[ns] + ]: + self.compare_by_show( + cdf.select(cfunc(cdf.ts1, tz=cdf.tz)), + sdf.select(sfunc(sdf.ts1, tz=sdf.tz)), + ) + + # With numeric parameter + for cfunc, sfunc in [ + (CF.date_add, SF.date_add), + (CF.date_sub, SF.date_sub), + (CF.add_months, SF.add_months), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.ts1, cdf.D)).toPandas(), + sdf.select(sfunc(sdf.ts1, sdf.D)).toPandas(), + ) + + # With another timestamp as parameter + for cfunc, sfunc in [ + (CF.datediff, SF.datediff), + (CF.months_between, SF.months_between), + ]: + self.assert_eq( + cdf.select(cfunc(cdf.ts1, cdf.ts2)).toPandas(), + sdf.select(sfunc(sdf.ts1, sdf.ts2)).toPandas(), + ) + + # With seconds parameter + self.compare_by_show( + # [left]: datetime64[ns, America/Los_Angeles] + # [right]: datetime64[ns] + cdf.select(CF.timestamp_seconds(cdf.seconds)), + sdf.select(SF.timestamp_seconds(sdf.seconds)), + ) + + # make_date + self.assert_eq( + cdf.select(CF.make_date(cdf.Y, cdf.M, cdf.D)).toPandas(), + sdf.select(SF.make_date(sdf.Y, sdf.M, sdf.D)).toPandas(), + ) + + # date_trunc + self.compare_by_show( + # [left]: datetime64[ns, America/Los_Angeles] + # [right]: datetime64[ns] + cdf.select(CF.date_trunc("day", cdf.ts1)), + sdf.select(SF.date_trunc("day", sdf.ts1)), + ) + + # trunc + self.assert_eq( + cdf.select(CF.trunc(cdf.ts1, "year")).toPandas(), + sdf.select(SF.trunc(sdf.ts1, "year")).toPandas(), + ) + + # next_day + self.assert_eq( + cdf.select(CF.next_day(cdf.ts1, "Mon")).toPandas(), + sdf.select(SF.next_day(sdf.ts1, "Mon")).toPandas(), + ) + + def test_time_window_functions(self): + query = """ + SELECT * FROM VALUES + (TIMESTAMP('2022-12-25 10:30:00'), 1), + (TIMESTAMP('2022-12-25 10:31:00'), 2), + (TIMESTAMP('2022-12-25 10:32:00'), 1), + (TIMESTAMP('2022-12-25 10:33:00'), 2), + (TIMESTAMP('2022-12-26 09:30:00'), 1), + (TIMESTAMP('2022-12-26 09:35:00'), 3) + AS tab(date, val) + """ + + # +-------------------+---+ + # | date|val| + # +-------------------+---+ + # |2022-12-25 10:30:00| 1| + # |2022-12-25 10:31:00| 2| + # |2022-12-25 10:32:00| 1| + # |2022-12-25 10:33:00| 2| + # |2022-12-26 09:30:00| 1| + # |2022-12-26 09:35:00| 3| + # +-------------------+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test window + self.compare_by_show( + cdf.select(CF.window("date", "15 seconds")), + sdf.select(SF.window("date", "15 seconds")), + truncate=100, + ) + self.compare_by_show( + cdf.select(CF.window(cdf.date, "1 minute")), + sdf.select(SF.window(sdf.date, "1 minute")), + truncate=100, + ) + + self.compare_by_show( + cdf.select(CF.window("date", "15 seconds", "5 seconds")), + sdf.select(SF.window("date", "15 seconds", "5 seconds")), + truncate=100, + ) + self.compare_by_show( + cdf.select(CF.window(cdf.date, "1 minute", "10 seconds")), + sdf.select(SF.window(sdf.date, "1 minute", "10 seconds")), + truncate=100, + ) + + self.compare_by_show( + cdf.select(CF.window("date", "15 seconds", "10 seconds", "5 seconds")), + sdf.select(SF.window("date", "15 seconds", "10 seconds", "5 seconds")), + truncate=100, + ) + self.compare_by_show( + cdf.select(CF.window(cdf.date, "1 minute", "10 seconds", "5 seconds")), + sdf.select(SF.window(sdf.date, "1 minute", "10 seconds", "5 seconds")), + truncate=100, + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.window("date", "15 seconds", 10, "5 seconds") + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "slideDuration", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + CF.window("date", "15 seconds", "10 seconds", 5) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "startTime", "arg_type": "int"}, + ) + + # test session_window + self.compare_by_show( + cdf.select(CF.session_window("date", "15 seconds")), + sdf.select(SF.session_window("date", "15 seconds")), + truncate=100, + ) + self.compare_by_show( + cdf.select(CF.session_window(cdf.date, "1 minute")), + sdf.select(SF.session_window(sdf.date, "1 minute")), + truncate=100, + ) + + # test window_time + self.compare_by_show( + cdf.groupBy(CF.window("date", "5 seconds")) + .agg(CF.sum("val").alias("sum")) + .select(CF.window_time("window")), + sdf.groupBy(SF.window("date", "5 seconds")) + .agg(SF.sum("val").alias("sum")) + .select(SF.window_time("window")), + truncate=100, + ) + + def test_misc_functions(self): + query = """ + SELECT a, b, c, BINARY(c) as d FROM VALUES + (0, float("NAN"), 'x'), (1, NULL, 'y'), (1, 2.1, 'z'), (0, 0.5, NULL) + AS tab(a, b, c) + """ + # +---+----+----+----+ + # | a| b| c| d| + # +---+----+----+----+ + # | 0| NaN| x|[78]| + # | 1|NULL| y|[79]| + # | 1| 2.1| z|[7A]| + # | 0| 0.5|NULL|NULL| + # +---+----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # test assert_true + with self.assertRaises(SparkConnectException): + cdf.select(CF.assert_true(cdf.a > 0, "a should be positive!")).show() + + # test raise_error + with self.assertRaises(SparkConnectException): + cdf.select(CF.raise_error("a should be positive!")).show() + + # test crc32 + self.assert_eq( + cdf.select(CF.crc32(cdf.d)).toPandas(), + sdf.select(SF.crc32(sdf.d)).toPandas(), + ) + + # test hash + self.assert_eq( + cdf.select(CF.hash(cdf.a, "b", cdf.c)).toPandas(), + sdf.select(SF.hash(sdf.a, "b", sdf.c)).toPandas(), + ) + + # test xxhash64 + self.assert_eq( + cdf.select(CF.xxhash64(cdf.a, "b", cdf.c)).toPandas(), + sdf.select(SF.xxhash64(sdf.a, "b", sdf.c)).toPandas(), + ) + + # test md5 + self.assert_eq( + cdf.select(CF.md5(cdf.d), CF.md5("c")).toPandas(), + sdf.select(SF.md5(sdf.d), SF.md5("c")).toPandas(), + ) + + # test sha1 + self.assert_eq( + cdf.select(CF.sha1(cdf.d), CF.sha1("c")).toPandas(), + sdf.select(SF.sha1(sdf.d), SF.sha1("c")).toPandas(), + ) + + # test sha2 + self.assert_eq( + cdf.select(CF.sha2(cdf.c, 256), CF.sha2("d", 512)).toPandas(), + sdf.select(SF.sha2(sdf.c, 256), SF.sha2("d", 512)).toPandas(), + ) + + def test_call_udf(self): + query = """ + SELECT a, b, c, BINARY(c) as d FROM VALUES + (-1.0, float("NAN"), 'x'), (-2.1, NULL, 'y'), (1, 2.1, 'z'), (0, 0.5, NULL) + AS tab(a, b, c) + """ + + # +----+----+----+----+ + # | a| b| c| d| + # +----+----+----+----+ + # |-1.0| NaN| x|[78]| + # |-2.1|NULL| y|[79]| + # | 1.0| 2.1| z|[7A]| + # | 0.0| 0.5|NULL|NULL| + # +----+----+----+----+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + self.assert_eq( + cdf.select( + CF.call_udf("abs", cdf.a), CF.call_udf("xxhash64", "b", cdf.c, "d") + ).toPandas(), + sdf.select( + SF.call_udf("abs", sdf.a), SF.call_udf("xxhash64", "b", sdf.c, "d") + ).toPandas(), + ) + + def test_udf(self): + query = """ + SELECT a, b, c FROM VALUES + (1, 1.0, 'x'), (2, 2.0, 'y'), (3, 3.0, 'z') + AS tab(a, b, c) + """ + # +---+---+---+ + # | a| b| c| + # +---+---+---+ + # | 1|1.0| x| + # | 2|2.0| y| + # | 3|3.0| z| + # +---+---+---+ + + cdf = self.connect.sql(query) + sdf = self.spark.sql(query) + + # as a normal function + self.assert_eq( + cdf.withColumn("A", CF.udf(lambda x: x + 1)(cdf.a)).toPandas(), + sdf.withColumn("A", SF.udf(lambda x: x + 1)(sdf.a)).toPandas(), + ) + self.assert_eq( # returnType as DDL strings + cdf.withColumn("C", CF.udf(lambda x: len(x), "int")(cdf.c)).toPandas(), + sdf.withColumn("C", SF.udf(lambda x: len(x), "int")(sdf.c)).toPandas(), + ) + self.assert_eq( # returnType as DataType + cdf.withColumn("C", CF.udf(lambda x: len(x), IntegerType())(cdf.c)).toPandas(), + sdf.withColumn("C", SF.udf(lambda x: len(x), IntegerType())(sdf.c)).toPandas(), + ) + + # as a decorator + @CF.udf(StringType()) + def cfun(x): + return x + "a" + + @SF.udf(StringType()) + def sfun(x): + return x + "a" + + self.assert_eq( + cdf.withColumn("A", cfun(cdf.c)).toPandas(), + sdf.withColumn("A", sfun(sdf.c)).toPandas(), + ) + + def test_udtf(self): + class TestUDTF: + def eval(self, x: int, y: int): + yield x, x + 1 + yield y, y + 1 + + sfunc = SF.udtf(TestUDTF, returnType="a: int, b: int") + cfunc = CF.udtf(TestUDTF, returnType="a: int, b: int") + + assertDataFrameEqual(sfunc(SF.lit(1), SF.lit(1)), cfunc(CF.lit(1), CF.lit(1))) + + self.spark.udtf.register("test_udtf", sfunc) + self.connect.udtf.register("test_udtf", cfunc) + + query = "select * from test_udtf(1, 2)" + assertDataFrameEqual(self.spark.sql(query), self.connect.sql(query)) + + def test_pandas_udf_import(self): + self.assert_eq(getattr(CF, "pandas_udf"), getattr(SF, "pandas_udf")) + + +if __name__ == "__main__": + import os + from pyspark.sql.tests.connect.test_connect_function import * # noqa: F401 + + # TODO(SPARK-41547): Enable ANSI mode in this file. + os.environ["SPARK_ANSI_SQL_MODE"] = "false" + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_plan.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_plan.py new file mode 100644 index 0000000..c39fb6b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_connect_plan.py @@ -0,0 +1,1060 @@ +# +# 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 unittest +import uuid +import datetime +import decimal +import math + +from pyspark.testing.connectutils import ( + PlanOnlyTestFixture, + should_test_connect, + connect_requirement_message, +) +from pyspark.errors import PySparkValueError + +if should_test_connect: + import pyspark.sql.connect.proto as proto + from pyspark.sql.connect.column import Column + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.plan import WriteOperation, Read + from pyspark.sql.connect.readwriter import DataFrameReader + from pyspark.sql.connect.expressions import LiteralExpression + from pyspark.sql.connect.functions import col, lit, max, min, sum + from pyspark.sql.connect.types import pyspark_types_to_proto_types + from pyspark.sql.types import ( + StringType, + StructType, + StructField, + IntegerType, + MapType, + ArrayType, + DoubleType, + ) + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class SparkConnectPlanTests(PlanOnlyTestFixture): + """These test cases exercise the interface to the proto plan + generation but do not call Spark.""" + + def test_sql_project(self): + plan = self.connect.sql("SELECT 1")._plan.to_proto(self.connect) + self.assertEqual(plan.root.sql.query, "SELECT 1") + + def test_simple_project(self): + plan = self.connect.readTable(table_name=self.tbl_name)._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root, "Root relation must be set") + self.assertIsNotNone(plan.root.read) + + def test_join_using_columns(self): + left_input = self.connect.readTable(table_name=self.tbl_name) + right_input = self.connect.readTable(table_name=self.tbl_name) + plan = left_input.join(other=right_input, on="join_column")._plan.to_proto(self.connect) + self.assertEqual(len(plan.root.join.using_columns), 1) + + plan2 = left_input.join(other=right_input, on=["col1", "col2"])._plan.to_proto(self.connect) + self.assertEqual(len(plan2.root.join.using_columns), 2) + + def test_join_condition(self): + left_input = self.connect.readTable(table_name=self.tbl_name) + right_input = self.connect.readTable(table_name=self.tbl_name) + plan = left_input.join( + other=right_input, on=left_input.name == right_input.name + )._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root.join.join_condition) + plan = left_input.join( + other=right_input, + on=[left_input.name == right_input.name, left_input.age == right_input.age], + )._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root.join.join_condition) + + def test_crossjoin(self): + # SPARK-41227: Test CrossJoin + left_input = self.connect.readTable(table_name=self.tbl_name) + right_input = self.connect.readTable(table_name=self.tbl_name) + crossJoin_plan = left_input.crossJoin(other=right_input)._plan.to_proto(self.connect) + join_plan = left_input.join(other=right_input, how="cross")._plan.to_proto(self.connect) + self.assertEqual( + crossJoin_plan.root.join.left.read.named_table, + join_plan.root.join.left.read.named_table, + ) + self.assertEqual( + crossJoin_plan.root.join.right.read.named_table, + join_plan.root.join.right.read.named_table, + ) + self.assertEqual( + crossJoin_plan.root.join.join_type, + join_plan.root.join.join_type, + ) + + def test_filter(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3)._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root.filter) + self.assertTrue( + isinstance( + plan.root.filter.condition.unresolved_function, proto.Expression.UnresolvedFunction + ) + ) + self.assertEqual(plan.root.filter.condition.unresolved_function.function_name, ">") + self.assertEqual(len(plan.root.filter.condition.unresolved_function.arguments), 2) + + def test_filter_with_string_expr(self): + """SPARK-41297: filter supports SQL expression""" + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter("id < 10")._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root.filter) + self.assertIsNotNone(plan.root.filter.condition.expression_string) + self.assertEqual(plan.root.filter.condition.expression_string.expression, "id < 10") + + def test_fill_na(self): + # SPARK-41128: Test fill na + df = self.connect.readTable(table_name=self.tbl_name) + + plan = df.fillna(value=1)._plan.to_proto(self.connect) + self.assertEqual(len(plan.root.fill_na.values), 1) + self.assertEqual(plan.root.fill_na.values[0].long, 1) + self.assertEqual(plan.root.fill_na.cols, []) + + plan = df.na.fill(value="xyz")._plan.to_proto(self.connect) + self.assertEqual(len(plan.root.fill_na.values), 1) + self.assertEqual(plan.root.fill_na.values[0].string, "xyz") + self.assertEqual(plan.root.fill_na.cols, []) + + plan = df.na.fill(value="xyz", subset=["col_a", "col_b"])._plan.to_proto(self.connect) + self.assertEqual(len(plan.root.fill_na.values), 1) + self.assertEqual(plan.root.fill_na.values[0].string, "xyz") + self.assertEqual(plan.root.fill_na.cols, ["col_a", "col_b"]) + + plan = df.na.fill(value=True, subset=("col_a", "col_b", "col_c"))._plan.to_proto( + self.connect + ) + self.assertEqual(len(plan.root.fill_na.values), 1) + self.assertEqual(plan.root.fill_na.values[0].boolean, True) + self.assertEqual(plan.root.fill_na.cols, ["col_a", "col_b", "col_c"]) + + plan = df.fillna({"col_a": 1.5, "col_b": "abc"})._plan.to_proto(self.connect) + self.assertEqual(len(plan.root.fill_na.values), 2) + self.assertEqual(plan.root.fill_na.values[0].double, 1.5) + self.assertEqual(plan.root.fill_na.values[1].string, "abc") + self.assertEqual(plan.root.fill_na.cols, ["col_a", "col_b"]) + + def test_drop_na(self): + # SPARK-41148: Test drop na + df = self.connect.readTable(table_name=self.tbl_name) + + plan = df.dropna()._plan.to_proto(self.connect) + self.assertEqual(plan.root.drop_na.cols, []) + self.assertEqual(plan.root.drop_na.HasField("min_non_nulls"), False) + + plan = df.na.drop(thresh=2, subset=("col_a", "col_b"))._plan.to_proto(self.connect) + self.assertEqual(plan.root.drop_na.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.drop_na.min_non_nulls, 2) + + plan = df.dropna(how="all", subset="col_c")._plan.to_proto(self.connect) + self.assertEqual(plan.root.drop_na.cols, ["col_c"]) + self.assertEqual(plan.root.drop_na.min_non_nulls, 1) + + def test_replace(self): + # SPARK-41315: Test replace + df = self.connect.readTable(table_name=self.tbl_name) + + plan = df.replace(10, 20)._plan.to_proto(self.connect) + self.assertEqual(plan.root.replace.cols, []) + self.assertEqual(plan.root.replace.replacements[0].old_value.double, 10.0) + self.assertEqual(plan.root.replace.replacements[0].new_value.double, 20.0) + + plan = df.na.replace((1, 2, 3), (4, 5, 6), subset=("col_a", "col_b"))._plan.to_proto( + self.connect + ) + self.assertEqual(plan.root.replace.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.replace.replacements[0].old_value.double, 1.0) + self.assertEqual(plan.root.replace.replacements[0].new_value.double, 4.0) + self.assertEqual(plan.root.replace.replacements[1].old_value.double, 2.0) + self.assertEqual(plan.root.replace.replacements[1].new_value.double, 5.0) + self.assertEqual(plan.root.replace.replacements[2].old_value.double, 3.0) + self.assertEqual(plan.root.replace.replacements[2].new_value.double, 6.0) + + plan = df.replace(["Alice", "Bob"], ["A", "B"], subset="col_x")._plan.to_proto(self.connect) + self.assertEqual(plan.root.replace.cols, ["col_x"]) + self.assertEqual(plan.root.replace.replacements[0].old_value.string, "Alice") + self.assertEqual(plan.root.replace.replacements[0].new_value.string, "A") + self.assertEqual(plan.root.replace.replacements[1].old_value.string, "Bob") + self.assertEqual(plan.root.replace.replacements[1].new_value.string, "B") + + def test_unpivot(self): + df = self.connect.readTable(table_name=self.tbl_name) + + plan = ( + df.filter(df.col_name > 3) + .unpivot(["id"], ["name"], "variable", "value") + ._plan.to_proto(self.connect) + ) + self.assertTrue(all(isinstance(c, proto.Expression) for c in plan.root.unpivot.ids)) + self.assertEqual(plan.root.unpivot.ids[0].unresolved_attribute.unparsed_identifier, "id") + self.assertEqual(plan.root.unpivot.HasField("values"), True) + self.assertTrue( + all(isinstance(c, proto.Expression) for c in plan.root.unpivot.values.values) + ) + self.assertEqual( + plan.root.unpivot.values.values[0].unresolved_attribute.unparsed_identifier, "name" + ) + self.assertEqual(plan.root.unpivot.variable_column_name, "variable") + self.assertEqual(plan.root.unpivot.value_column_name, "value") + + plan = ( + df.filter(df.col_name > 3) + .unpivot(["id"], None, "variable", "value") + ._plan.to_proto(self.connect) + ) + self.assertTrue(len(plan.root.unpivot.ids) == 1) + self.assertTrue(all(isinstance(c, proto.Expression) for c in plan.root.unpivot.ids)) + self.assertEqual(plan.root.unpivot.ids[0].unresolved_attribute.unparsed_identifier, "id") + self.assertEqual(plan.root.unpivot.HasField("values"), False) + self.assertEqual(plan.root.unpivot.variable_column_name, "variable") + self.assertEqual(plan.root.unpivot.value_column_name, "value") + + def test_melt(self): + df = self.connect.readTable(table_name=self.tbl_name) + + plan = ( + df.filter(df.col_name > 3) + .melt(["id"], ["name"], "variable", "value") + ._plan.to_proto(self.connect) + ) + self.assertTrue(all(isinstance(c, proto.Expression) for c in plan.root.unpivot.ids)) + self.assertEqual(plan.root.unpivot.ids[0].unresolved_attribute.unparsed_identifier, "id") + self.assertEqual(plan.root.unpivot.HasField("values"), True) + self.assertTrue( + all(isinstance(c, proto.Expression) for c in plan.root.unpivot.values.values) + ) + self.assertEqual( + plan.root.unpivot.values.values[0].unresolved_attribute.unparsed_identifier, "name" + ) + self.assertEqual(plan.root.unpivot.variable_column_name, "variable") + self.assertEqual(plan.root.unpivot.value_column_name, "value") + + plan = ( + df.filter(df.col_name > 3) + .melt(["id"], [], "variable", "value") + ._plan.to_proto(self.connect) + ) + self.assertTrue(len(plan.root.unpivot.ids) == 1) + self.assertTrue(all(isinstance(c, proto.Expression) for c in plan.root.unpivot.ids)) + self.assertEqual(plan.root.unpivot.ids[0].unresolved_attribute.unparsed_identifier, "id") + self.assertEqual(plan.root.unpivot.HasField("values"), True) + self.assertTrue(len(plan.root.unpivot.values.values) == 0) + self.assertEqual(plan.root.unpivot.variable_column_name, "variable") + self.assertEqual(plan.root.unpivot.value_column_name, "value") + + def test_random_split(self): + # SPARK-41440: test randomSplit(weights, seed). + from typing import List + + df = self.connect.readTable(table_name=self.tbl_name) + + def checkRelations(relations: List["DataFrame"]): + self.assertTrue(len(relations) == 3) + + plan = relations[0]._plan.to_proto(self.connect) + self.assertEqual(plan.root.sample.lower_bound, 0.0) + self.assertEqual(plan.root.sample.upper_bound, 0.16666666666666666) + self.assertEqual(plan.root.sample.with_replacement, False) + self.assertEqual(plan.root.sample.HasField("seed"), True) + self.assertEqual(plan.root.sample.deterministic_order, True) + + plan = relations[1]._plan.to_proto(self.connect) + self.assertEqual(plan.root.sample.lower_bound, 0.16666666666666666) + self.assertEqual(plan.root.sample.upper_bound, 0.5) + self.assertEqual(plan.root.sample.with_replacement, False) + self.assertEqual(plan.root.sample.HasField("seed"), True) + self.assertEqual(plan.root.sample.deterministic_order, True) + + plan = relations[2]._plan.to_proto(self.connect) + self.assertEqual(plan.root.sample.lower_bound, 0.5) + self.assertEqual(plan.root.sample.upper_bound, 1.0) + self.assertEqual(plan.root.sample.with_replacement, False) + self.assertEqual(plan.root.sample.HasField("seed"), True) + self.assertEqual(plan.root.sample.deterministic_order, True) + + relations = df.filter(df.col_name > 3).randomSplit([1.0, 2.0, 3.0], 1) + checkRelations(relations) + + relations = df.filter(df.col_name > 3).randomSplit([1.0, 2.0, 3.0]) + checkRelations(relations) + + def test_observe(self): + # SPARK-41527: test DataFrame.observe() + df = self.connect.readTable(table_name=self.tbl_name) + + plan = ( + df.filter(df.col_name > 3) + .observe("my_metric", min("id"), max("id"), sum("id")) + ._plan.to_proto(self.connect) + ) + self.assertEqual(plan.root.collect_metrics.name, "my_metric") + self.assertTrue( + all(isinstance(c, proto.Expression) for c in plan.root.collect_metrics.metrics) + ) + self.assertEqual( + plan.root.collect_metrics.metrics[0].unresolved_function.function_name, "min" + ) + self.assertTrue( + len(plan.root.collect_metrics.metrics[0].unresolved_function.arguments) == 1 + ) + self.assertTrue( + all( + isinstance(c, proto.Expression) + for c in plan.root.collect_metrics.metrics[0].unresolved_function.arguments + ) + ) + self.assertEqual( + plan.root.collect_metrics.metrics[0] + .unresolved_function.arguments[0] + .unresolved_attribute.unparsed_identifier, + "id", + ) + + from pyspark.sql.observation import Observation + + plan = ( + df.filter(df.col_name > 3) + .observe(Observation("my_metric"), min("id"), max("id"), sum("id")) + ._plan.to_proto(self.connect) + ) + self.assertEqual(plan.root.collect_metrics.name, "my_metric") + self.assertTrue( + all(isinstance(c, proto.Expression) for c in plan.root.collect_metrics.metrics) + ) + self.assertEqual( + plan.root.collect_metrics.metrics[0].unresolved_function.function_name, "min" + ) + self.assertTrue( + len(plan.root.collect_metrics.metrics[0].unresolved_function.arguments) == 1 + ) + self.assertTrue( + all( + isinstance(c, proto.Expression) + for c in plan.root.collect_metrics.metrics[0].unresolved_function.arguments + ) + ) + self.assertEqual( + plan.root.collect_metrics.metrics[0] + .unresolved_function.arguments[0] + .unresolved_attribute.unparsed_identifier, + "id", + ) + + def test_summary(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3).summary()._plan.to_proto(self.connect) + self.assertEqual(plan.root.summary.statistics, []) + + plan = ( + df.filter(df.col_name > 3) + .summary("count", "mean", "stddev", "min", "25%") + ._plan.to_proto(self.connect) + ) + self.assertEqual( + plan.root.summary.statistics, + ["count", "mean", "stddev", "min", "25%"], + ) + + def test_describe(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3).describe()._plan.to_proto(self.connect) + self.assertEqual(plan.root.describe.cols, []) + + plan = df.filter(df.col_name > 3).describe("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual( + plan.root.describe.cols, + ["col_a", "col_b"], + ) + + def test_crosstab(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3).crosstab("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual(plan.root.crosstab.col1, "col_a") + self.assertEqual(plan.root.crosstab.col2, "col_b") + + plan = df.stat.crosstab("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual(plan.root.crosstab.col1, "col_a") + self.assertEqual(plan.root.crosstab.col2, "col_b") + + def test_freqItems(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = ( + df.filter(df.col_name > 3).freqItems(["col_a", "col_b"], 1)._plan.to_proto(self.connect) + ) + self.assertEqual(plan.root.freq_items.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.freq_items.support, 1) + plan = df.filter(df.col_name > 3).freqItems(["col_a", "col_b"])._plan.to_proto(self.connect) + self.assertEqual(plan.root.freq_items.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.freq_items.support, 0.01) + + plan = df.stat.freqItems(["col_a", "col_b"], 1)._plan.to_proto(self.connect) + self.assertEqual(plan.root.freq_items.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.freq_items.support, 1) + plan = df.stat.freqItems(["col_a", "col_b"])._plan.to_proto(self.connect) + self.assertEqual(plan.root.freq_items.cols, ["col_a", "col_b"]) + self.assertEqual(plan.root.freq_items.support, 0.01) + + def test_limit(self): + df = self.connect.readTable(table_name=self.tbl_name) + limit_plan = df.limit(10)._plan.to_proto(self.connect) + self.assertEqual(limit_plan.root.limit.limit, 10) + + def test_offset(self): + df = self.connect.readTable(table_name=self.tbl_name) + offset_plan = df.offset(10)._plan.to_proto(self.connect) + self.assertEqual(offset_plan.root.offset.offset, 10) + + def test_sample(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3).sample(fraction=0.3)._plan.to_proto(self.connect) + self.assertEqual(plan.root.sample.lower_bound, 0.0) + self.assertEqual(plan.root.sample.upper_bound, 0.3) + self.assertEqual(plan.root.sample.with_replacement, False) + self.assertEqual(plan.root.sample.HasField("seed"), False) + self.assertEqual(plan.root.sample.deterministic_order, False) + + plan = ( + df.filter(df.col_name > 3) + .sample(withReplacement=True, fraction=0.4, seed=-1) + ._plan.to_proto(self.connect) + ) + self.assertEqual(plan.root.sample.lower_bound, 0.0) + self.assertEqual(plan.root.sample.upper_bound, 0.4) + self.assertEqual(plan.root.sample.with_replacement, True) + self.assertEqual(plan.root.sample.seed, -1) + self.assertEqual(plan.root.sample.deterministic_order, False) + + def test_sort(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.filter(df.col_name > 3).sort("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual( + [f.child.unresolved_attribute.unparsed_identifier for f in plan.root.sort.order], + ["col_a", "col_b"], + ) + self.assertEqual(plan.root.sort.is_global, True) + self.assertEqual( + plan.root.sort.order[0].direction, + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_ASCENDING, + ) + self.assertEqual( + plan.root.sort.order[0].direction, + proto.Expression.SortOrder.NullOrdering.SORT_NULLS_FIRST, + ) + self.assertEqual( + plan.root.sort.order[1].direction, + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_ASCENDING, + ) + self.assertEqual( + plan.root.sort.order[1].direction, + proto.Expression.SortOrder.NullOrdering.SORT_NULLS_FIRST, + ) + + plan = df.filter(df.col_name > 3).orderBy("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual( + [f.child.unresolved_attribute.unparsed_identifier for f in plan.root.sort.order], + ["col_a", "col_b"], + ) + self.assertEqual(plan.root.sort.is_global, True) + + plan = ( + df.filter(df.col_name > 3) + .sortWithinPartitions(df.col_a.desc(), df.col_b.asc()) + ._plan.to_proto(self.connect) + ) + self.assertEqual( + [f.child.unresolved_attribute.unparsed_identifier for f in plan.root.sort.order], + ["col_a", "col_b"], + ) + self.assertEqual(plan.root.sort.is_global, False) + self.assertEqual( + plan.root.sort.order[0].direction, + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_DESCENDING, + ) + self.assertEqual( + plan.root.sort.order[0].direction, + proto.Expression.SortOrder.NullOrdering.SORT_NULLS_LAST, + ) + self.assertEqual( + plan.root.sort.order[1].direction, + proto.Expression.SortOrder.SortDirection.SORT_DIRECTION_ASCENDING, + ) + self.assertEqual( + plan.root.sort.order[1].direction, + proto.Expression.SortOrder.NullOrdering.SORT_NULLS_FIRST, + ) + + def test_drop(self): + # SPARK-41169: test drop + df = self.connect.readTable(table_name=self.tbl_name) + + plan = df.filter(df.col_name > 3).drop("col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual( + plan.root.drop.column_names, + ["col_a", "col_b"], + ) + + plan = df.filter(df.col_name > 3).drop(df.col_x, "col_b")._plan.to_proto(self.connect) + self.assertEqual( + [f.unresolved_attribute.unparsed_identifier for f in plan.root.drop.columns], + ["col_x"], + ) + self.assertEqual( + plan.root.drop.column_names, + ["col_b"], + ) + + def test_deduplicate(self): + df = self.connect.readTable(table_name=self.tbl_name) + + distinct_plan = df.distinct()._plan.to_proto(self.connect) + self.assertTrue(distinct_plan.root.deduplicate.HasField("input"), "input must be set") + + self.assertEqual(distinct_plan.root.deduplicate.all_columns_as_keys, True) + self.assertEqual(len(distinct_plan.root.deduplicate.column_names), 0) + + deduplicate_on_all_columns_plan = df.dropDuplicates()._plan.to_proto(self.connect) + self.assertEqual(deduplicate_on_all_columns_plan.root.deduplicate.all_columns_as_keys, True) + self.assertEqual(len(deduplicate_on_all_columns_plan.root.deduplicate.column_names), 0) + + deduplicate_on_all_columns_plan = df.drop_duplicates()._plan.to_proto(self.connect) + self.assertEqual(deduplicate_on_all_columns_plan.root.deduplicate.all_columns_as_keys, True) + self.assertEqual(len(deduplicate_on_all_columns_plan.root.deduplicate.column_names), 0) + + deduplicate_on_subset_columns_plan = df.dropDuplicates(["name", "height"])._plan.to_proto( + self.connect + ) + self.assertEqual( + deduplicate_on_subset_columns_plan.root.deduplicate.all_columns_as_keys, False + ) + self.assertEqual(len(deduplicate_on_subset_columns_plan.root.deduplicate.column_names), 2) + + def test_relation_alias(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.alias("table_alias")._plan.to_proto(self.connect) + self.assertEqual(plan.root.subquery_alias.alias, "table_alias") + self.assertIsNotNone(plan.root.subquery_alias.input) + + def test_range(self): + plan = self.connect.range(start=10, end=20, step=3, num_partitions=4)._plan.to_proto( + self.connect + ) + self.assertEqual(plan.root.range.start, 10) + self.assertEqual(plan.root.range.end, 20) + self.assertEqual(plan.root.range.step, 3) + self.assertEqual(plan.root.range.num_partitions, 4) + + plan = self.connect.range(start=10, end=20)._plan.to_proto(self.connect) + self.assertEqual(plan.root.range.start, 10) + self.assertEqual(plan.root.range.end, 20) + self.assertEqual(plan.root.range.step, 1) + self.assertFalse(plan.root.range.HasField("num_partitions")) + + def test_datasource_read(self): + reader = DataFrameReader(self.connect) + df = reader.load(path="test_path", format="text", schema="id INT", op1="opv", op2="opv2") + plan = df._plan.to_proto(self.connect) + data_source = plan.root.read.data_source + self.assertEqual(data_source.format, "text") + self.assertEqual(data_source.schema, "id INT") + self.assertEqual(len(data_source.options), 2) + self.assertEqual(data_source.options.get("op1"), "opv") + self.assertEqual(data_source.options.get("op2"), "opv2") + self.assertEqual(len(data_source.paths), 1) + self.assertEqual(data_source.paths[0], "test_path") + + def test_all_the_plans(self): + df = self.connect.readTable(table_name=self.tbl_name) + df = df.select(df.col1).filter(df.col2 == 2).sort(df.col3.asc()) + plan = df._plan.to_proto(self.connect) + self.assertIsNotNone(plan.root, "Root relation must be set") + self.assertIsNotNone(plan.root.read) + + def test_union(self): + df1 = self.connect.readTable(table_name=self.tbl_name) + df2 = self.connect.readTable(table_name=self.tbl_name) + plan1 = df1.union(df2)._plan.to_proto(self.connect) + self.assertTrue(plan1.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_UNION, plan1.root.set_op.set_op_type) + plan2 = df1.union(df2)._plan.to_proto(self.connect) + self.assertTrue(plan2.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_UNION, plan2.root.set_op.set_op_type) + plan3 = df1.unionByName(df2, True)._plan.to_proto(self.connect) + self.assertTrue(plan3.root.set_op.by_name) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_UNION, plan3.root.set_op.set_op_type) + + def test_subtract(self): + # SPARK-41453: test `subtract` API for Python client. + df1 = self.connect.readTable(table_name=self.tbl_name) + df2 = self.connect.readTable(table_name=self.tbl_name) + plan1 = df1.subtract(df2)._plan.to_proto(self.connect) + self.assertTrue(not plan1.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_EXCEPT, plan1.root.set_op.set_op_type) + + def test_except(self): + # SPARK-41010: test `except` API for Python client. + df1 = self.connect.readTable(table_name=self.tbl_name) + df2 = self.connect.readTable(table_name=self.tbl_name) + plan1 = df1.exceptAll(df2)._plan.to_proto(self.connect) + self.assertTrue(plan1.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_EXCEPT, plan1.root.set_op.set_op_type) + + def test_intersect(self): + # SPARK-41010: test `intersect` API for Python client. + df1 = self.connect.readTable(table_name=self.tbl_name) + df2 = self.connect.readTable(table_name=self.tbl_name) + plan1 = df1.intersect(df2)._plan.to_proto(self.connect) + self.assertFalse(plan1.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_INTERSECT, plan1.root.set_op.set_op_type) + plan2 = df1.intersectAll(df2)._plan.to_proto(self.connect) + self.assertTrue(plan2.root.set_op.is_all) + self.assertEqual(proto.SetOperation.SET_OP_TYPE_INTERSECT, plan2.root.set_op.set_op_type) + + def test_coalesce_and_repartition(self): + # SPARK-41026: test Coalesce and Repartition API in Python client. + df = self.connect.readTable(table_name=self.tbl_name) + plan1 = df.coalesce(10)._plan.to_proto(self.connect) + self.assertEqual(10, plan1.root.repartition.num_partitions) + self.assertFalse(plan1.root.repartition.shuffle) + plan2 = df.repartition(20)._plan.to_proto(self.connect) + self.assertTrue(plan2.root.repartition.shuffle) + + with self.assertRaises(PySparkValueError) as pe: + df.coalesce(-1)._plan.to_proto(self.connect) + + self.check_error( + exception=pe.exception, + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "numPartitions", "arg_value": "-1"}, + ) + + with self.assertRaises(PySparkValueError) as pe: + df.repartition(-1)._plan.to_proto(self.connect) + + self.check_error( + exception=pe.exception, + error_class="VALUE_NOT_POSITIVE", + message_parameters={"arg_name": "numPartitions", "arg_value": "-1"}, + ) + + def test_repartition_by_expression(self): + # SPARK-41354: test dataframe.repartition(expressions) + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.repartition(10, "col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual(10, plan.root.repartition_by_expression.num_partitions) + self.assertEqual( + [ + f.unresolved_attribute.unparsed_identifier + for f in plan.root.repartition_by_expression.partition_exprs + ], + ["col_a", "col_b"], + ) + + def test_repartition_by_range(self): + # SPARK-41354: test dataframe.repartitionByRange(expressions) + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.repartitionByRange(10, "col_a", "col_b")._plan.to_proto(self.connect) + self.assertEqual(10, plan.root.repartition_by_expression.num_partitions) + self.assertEqual( + [ + f.sort_order.child.unresolved_attribute.unparsed_identifier + for f in plan.root.repartition_by_expression.partition_exprs + ], + ["col_a", "col_b"], + ) + + def test_to(self): + # SPARK-41464: test `to` API in Python client. + df = self.connect.readTable(table_name=self.tbl_name) + + schema = StructType( + [ + StructField("col1", IntegerType(), True), + StructField("col2", StringType(), True), + StructField("map1", MapType(StringType(), IntegerType(), True), True), + StructField("array1", ArrayType(IntegerType(), True), True), + ] + ) + new_plan = df.to(schema)._plan.to_proto(self.connect) + self.assertEqual(pyspark_types_to_proto_types(schema), new_plan.root.to_schema.schema) + + def test_write_operation(self): + wo = WriteOperation(self.connect.readTable("name")._plan) + wo.mode = "overwrite" + wo.source = "parquet" + + p = wo.command(None) + self.assertIsNotNone(p) + self.assertFalse(p.write_operation.HasField("path")) + self.assertFalse(p.write_operation.HasField("table")) + + wo.path = "path" + p = wo.command(None) + self.assertIsNotNone(p) + self.assertTrue(p.write_operation.HasField("path")) + self.assertFalse(p.write_operation.HasField("table")) + + wo.path = None + wo.table_name = "table" + wo.table_save_method = "save_as_table" + p = wo.command(None) + self.assertFalse(p.write_operation.HasField("path")) + self.assertTrue(p.write_operation.HasField("table")) + + wo.bucket_cols = ["a", "b", "c"] + p = wo.command(None) + self.assertFalse(p.write_operation.HasField("bucket_by")) + + wo.num_buckets = 10 + p = wo.command(None) + self.assertTrue(p.write_operation.HasField("bucket_by")) + + # Unsupported save mode + wo.mode = "unknown" + with self.assertRaises(ValueError): + wo.command(None) + + def test_column_regexp(self): + # SPARK-41438: test colRegex + df = self.connect.readTable(table_name=self.tbl_name) + col = df.colRegex("col_name") + self.assertIsInstance(col, Column) + self.assertEqual("Column<'UnresolvedRegex(col_name)'>", str(col)) + + col_plan = col.to_plan(self.session.client) + self.assertIsNotNone(col_plan) + self.assertEqual(col_plan.unresolved_regex.col_name, "col_name") + + def test_print(self): + # SPARK-41717: test print + self.assertEqual( + self.connect.sql("SELECT 1")._plan.print().strip(), + "", + ) + self.assertEqual( + self.connect.range(1, 10)._plan.print().strip(), + "", + ) + + def test_repr(self): + # SPARK-41717: test __repr_html__ + self.assertIn("query: SELECT 1", self.connect.sql("SELECT 1")._plan._repr_html_().strip()) + + expected = ( + "Range
    ", + "start: 1
    ", + "end: 10
    ", + "step: 1
    ", + "num_partitions: None
    ", + ) + actual = self.connect.range(1, 10)._plan._repr_html_().strip() + for line in expected: + self.assertIn(line, actual) + + def test_select_with_columns_and_strings(self): + df = self.connect.with_plan(Read("table")) + self.assertIsNotNone(df.select(col("name"))._plan.to_proto(self.connect)) + self.assertIsNotNone(df.select("name")) + self.assertIsNotNone(df.select("name", "name2")) + self.assertIsNotNone(df.select(col("name"), col("name2"))) + self.assertIsNotNone(df.select(col("name"), "name2")) + self.assertIsNotNone(df.select("*")) + + def test_join_with_join_type(self): + df_left = self.connect.with_plan(Read("table")) + df_right = self.connect.with_plan(Read("table")) + for (join_type_str, join_type) in [ + (None, proto.Join.JoinType.JOIN_TYPE_INNER), + ("inner", proto.Join.JoinType.JOIN_TYPE_INNER), + ("outer", proto.Join.JoinType.JOIN_TYPE_FULL_OUTER), + ("full", proto.Join.JoinType.JOIN_TYPE_FULL_OUTER), + ("fullouter", proto.Join.JoinType.JOIN_TYPE_FULL_OUTER), + ("full_outer", proto.Join.JoinType.JOIN_TYPE_FULL_OUTER), + ("left", proto.Join.JoinType.JOIN_TYPE_LEFT_OUTER), + ("leftouter", proto.Join.JoinType.JOIN_TYPE_LEFT_OUTER), + ("left_outer", proto.Join.JoinType.JOIN_TYPE_LEFT_OUTER), + ("right", proto.Join.JoinType.JOIN_TYPE_RIGHT_OUTER), + ("rightouter", proto.Join.JoinType.JOIN_TYPE_RIGHT_OUTER), + ("right_outer", proto.Join.JoinType.JOIN_TYPE_RIGHT_OUTER), + ("semi", proto.Join.JoinType.JOIN_TYPE_LEFT_SEMI), + ("leftsemi", proto.Join.JoinType.JOIN_TYPE_LEFT_SEMI), + ("left_semi", proto.Join.JoinType.JOIN_TYPE_LEFT_SEMI), + ("anti", proto.Join.JoinType.JOIN_TYPE_LEFT_ANTI), + ("leftanti", proto.Join.JoinType.JOIN_TYPE_LEFT_ANTI), + ("left_anti", proto.Join.JoinType.JOIN_TYPE_LEFT_ANTI), + ("cross", proto.Join.JoinType.JOIN_TYPE_CROSS), + ]: + joined_df = df_left.join(df_right, on=col("name"), how=join_type_str)._plan.to_proto( + self.connect + ) + self.assertEqual(joined_df.root.join.join_type, join_type) + + def test_simple_column_expressions(self): + df = self.connect.with_plan(Read("table")) + + c1 = df.col_name + self.assertIsInstance(c1, Column) + c2 = df["col_name"] + self.assertIsInstance(c2, Column) + c3 = col("col_name") + self.assertIsInstance(c3, Column) + + # All Protos should be identical + cp1 = c1.to_plan(None) + cp2 = c2.to_plan(None) + cp3 = c3.to_plan(None) + + self.assertIsNotNone(cp1) + self.assertEqual(cp1, cp2) + self.assertEqual( + cp2.unresolved_attribute.unparsed_identifier, + cp3.unresolved_attribute.unparsed_identifier, + ) + self.assertTrue(cp2.unresolved_attribute.HasField("plan_id")) + self.assertFalse(cp3.unresolved_attribute.HasField("plan_id")) + + def test_null_literal(self): + null_lit = lit(None) + null_lit_p = null_lit.to_plan(None) + self.assertEqual(null_lit_p.literal.HasField("null"), True) + + def test_binary_literal(self): + val = b"binary\0\0asas" + bin_lit = lit(val) + bin_lit_p = bin_lit.to_plan(None) + self.assertEqual(bin_lit_p.literal.binary, val) + + def test_uuid_literal(self): + + val = uuid.uuid4() + with self.assertRaises(TypeError): + lit(val) + + def test_column_literals(self): + df = self.connect.with_plan(Read("table")) + lit_df = df.select(lit(10)) + self.assertIsNotNone(lit_df._plan.to_proto(None)) + + self.assertIsNotNone(lit(10).to_plan(None)) + plan = lit(10).to_plan(None) + self.assertIs(plan.literal.integer, 10) + + plan = lit(1 << 33).to_plan(None) + self.assertEqual(plan.literal.long, 1 << 33) + + def test_numeric_literal_types(self): + int_lit = lit(10) + float_lit = lit(10.1) + decimal_lit = lit(decimal.Decimal(99)) + + self.assertIsNotNone(int_lit.to_plan(None)) + self.assertIsNotNone(float_lit.to_plan(None)) + self.assertIsNotNone(decimal_lit.to_plan(None)) + + def test_float_nan_inf(self): + na_lit = lit(float("nan")) + self.assertIsNotNone(na_lit.to_plan(None)) + + inf_lit = lit(float("inf")) + self.assertIsNotNone(inf_lit.to_plan(None)) + + inf_lit = lit(float("-inf")) + self.assertIsNotNone(inf_lit.to_plan(None)) + + def test_datetime_literal_types(self): + """Test the different timestamp, date, and timedelta types.""" + datetime_lit = lit(datetime.datetime.now()) + + p = datetime_lit.to_plan(None) + self.assertIsNotNone(datetime_lit.to_plan(None)) + self.assertGreater(p.literal.timestamp, 10000000000000) + + date_lit = lit(datetime.date.today()) + time_delta = lit(datetime.timedelta(days=1, seconds=2, microseconds=3)) + + self.assertIsNotNone(date_lit.to_plan(None)) + self.assertIsNotNone(time_delta.to_plan(None)) + # (24 * 3600 + 2) * 1000000 + 3 + self.assertEqual(86402000003, time_delta.to_plan(None).literal.day_time_interval) + + def test_list_to_literal(self): + """Test conversion of lists to literals""" + empty_list = [] + single_type = [1, 2, 3, 4] + multi_type = ["ooo", 1, "asas", 2.3] + + empty_list_lit = lit(empty_list) + single_type_lit = lit(single_type) + multi_type_lit = lit(multi_type) + + p = empty_list_lit.to_plan(None) + self.assertIsNotNone(p) + + p = single_type_lit.to_plan(None) + self.assertIsNotNone(p) + + p = multi_type_lit.to_plan(None) + self.assertIsNotNone(p) + + def test_column_alias(self) -> None: + # SPARK-40809: Support for Column Aliases + col0 = col("a").alias("martin") + self.assertEqual("Column<'a AS martin'>", str(col0)) + + col0 = col("a").alias("martin", metadata={"pii": True}) + plan = col0.to_plan(self.session.client) + self.assertIsNotNone(plan) + self.assertEqual(plan.alias.metadata, '{"pii": true}') + + def test_column_expressions(self): + """Test a more complex combination of expressions and their translation into + the protobuf structure.""" + df = self.connect.with_plan(Read("table")) + + expr = lit(10) < lit(10) + expr_plan = expr.to_plan(None) + self.assertIsNotNone(expr_plan.unresolved_function) + self.assertEqual(expr_plan.unresolved_function.function_name, "<") + + expr = df.id % lit(10) == lit(10) + expr_plan = expr.to_plan(None) + self.assertIsNotNone(expr_plan.unresolved_function) + self.assertEqual(expr_plan.unresolved_function.function_name, "==") + + lit_fun = expr_plan.unresolved_function.arguments[1] + self.assertIsInstance(lit_fun, proto.Expression) + self.assertIsInstance(lit_fun.literal, proto.Expression.Literal) + self.assertEqual(lit_fun.literal.integer, 10) + + mod_fun = expr_plan.unresolved_function.arguments[0] + self.assertIsInstance(mod_fun, proto.Expression) + self.assertIsInstance(mod_fun.unresolved_function, proto.Expression.UnresolvedFunction) + self.assertEqual(len(mod_fun.unresolved_function.arguments), 2) + self.assertIsInstance(mod_fun.unresolved_function.arguments[0], proto.Expression) + self.assertIsInstance( + mod_fun.unresolved_function.arguments[0].unresolved_attribute, + proto.Expression.UnresolvedAttribute, + ) + self.assertEqual( + mod_fun.unresolved_function.arguments[0].unresolved_attribute.unparsed_identifier, "id" + ) + + def test_literal_expression_with_arrays(self): + l0 = LiteralExpression._from_value(["x", "y", "z"]).to_plan(None).literal + self.assertTrue(l0.array.element_type.HasField("string")) + self.assertEqual(len(l0.array.elements), 3) + self.assertEqual(l0.array.elements[0].string, "x") + self.assertEqual(l0.array.elements[1].string, "y") + self.assertEqual(l0.array.elements[2].string, "z") + + l1 = LiteralExpression._from_value([3, -3]).to_plan(None).literal + self.assertTrue(l1.array.element_type.HasField("integer")) + self.assertEqual(len(l1.array.elements), 2) + self.assertEqual(l1.array.elements[0].integer, 3) + self.assertEqual(l1.array.elements[1].integer, -3) + + l2 = LiteralExpression._from_value([float("nan"), -3.0, 0.0]).to_plan(None).literal + self.assertTrue(l2.array.element_type.HasField("double")) + self.assertEqual(len(l2.array.elements), 3) + self.assertTrue(math.isnan(l2.array.elements[0].double)) + self.assertEqual(l2.array.elements[1].double, -3.0) + self.assertEqual(l2.array.elements[2].double, 0.0) + + l3 = LiteralExpression._from_value([[3, 4], [5, 6, 7]]).to_plan(None).literal + self.assertTrue(l3.array.element_type.HasField("array")) + self.assertTrue(l3.array.element_type.array.element_type.HasField("integer")) + self.assertEqual(len(l3.array.elements), 2) + self.assertEqual(len(l3.array.elements[0].array.elements), 2) + self.assertEqual(len(l3.array.elements[1].array.elements), 3) + + l4 = ( + LiteralExpression._from_value([[float("inf"), 0.4], [0.5, float("nan")], []]) + .to_plan(None) + .literal + ) + self.assertTrue(l4.array.element_type.HasField("array")) + self.assertTrue(l4.array.element_type.array.element_type.HasField("double")) + self.assertEqual(len(l4.array.elements), 3) + self.assertEqual(len(l4.array.elements[0].array.elements), 2) + self.assertEqual(len(l4.array.elements[1].array.elements), 2) + self.assertEqual(len(l4.array.elements[2].array.elements), 0) + + def test_literal_to_any_conversion(self): + for value in [ + b"binary\0\0asas", + True, + False, + 0, + 12, + -1, + 0.0, + 1.234567, + decimal.Decimal(0.0), + decimal.Decimal(1.234567), + "sss", + datetime.date(2022, 12, 13), + datetime.datetime.now(), + datetime.timedelta(1, 2, 3), + [1, 2, 3, 4, 5, 6], + [-1.0, 2.0, 3.0], + ["x", "y", "z"], + [[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]], + ]: + lit = LiteralExpression._from_value(value) + proto_lit = lit.to_plan(None).literal + value2 = LiteralExpression._to_value(proto_lit) + self.assertEqual(value, value2) + + with self.assertRaises(AssertionError): + lit = LiteralExpression._from_value(1.234567) + proto_lit = lit.to_plan(None).literal + LiteralExpression._to_value(proto_lit, StringType()) + + with self.assertRaises(AssertionError): + lit = LiteralExpression._from_value("1.234567") + proto_lit = lit.to_plan(None).literal + LiteralExpression._to_value(proto_lit, DoubleType()) + + with self.assertRaises(AssertionError): + # build a array proto literal, but with incorrect elements + proto_lit = proto.Expression().literal + proto_lit.array.element_type.CopyFrom(pyspark_types_to_proto_types(StringType())) + proto_lit.array.elements.append( + LiteralExpression("string", StringType()).to_plan(None).literal + ) + proto_lit.array.elements.append( + LiteralExpression(1.234, DoubleType()).to_plan(None).literal + ) + + LiteralExpression._to_value(proto_lit, DoubleType) + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_connect_plan import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow.py new file mode 100644 index 0000000..a92ef97 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow.py @@ -0,0 +1,158 @@ +# +# 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 unittest +from distutils.version import LooseVersion + +import pandas as pd + +from pyspark.sql.tests.test_arrow import ArrowTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + +class ArrowParityTests(ArrowTestsMixin, ReusedConnectTestCase, PandasOnSparkTestUtils): + @unittest.skip("Spark Connect does not support Spark Context but the test depends on that.") + def test_createDataFrame_empty_partition(self): + super().test_createDataFrame_empty_partition() + + @unittest.skip("Spark Connect does not support fallback.") + def test_createDataFrame_fallback_disabled(self): + super().test_createDataFrame_fallback_disabled() + + @unittest.skip("Spark Connect does not support fallback.") + def test_createDataFrame_fallback_enabled(self): + super().test_createDataFrame_fallback_enabled() + + def test_createDataFrame_with_incorrect_schema(self): + self.check_createDataFrame_with_incorrect_schema() + + def test_createDataFrame_with_map_type(self): + self.check_createDataFrame_with_map_type(True) + + def test_createDataFrame_with_ndarray(self): + self.check_createDataFrame_with_ndarray(True) + + def test_createDataFrame_with_single_data_type(self): + self.check_createDataFrame_with_single_data_type() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_no_partition_frame(self): + super().test_no_partition_frame() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_no_partition_toPandas(self): + super().test_no_partition_toPandas() + + def test_pandas_self_destruct(self): + df = self.spark.range(100).select("id", "id", "id") + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": True}): + self_destruct_pdf = df.toPandas() + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": False}): + no_self_destruct_pdf = df.toPandas() + + self.assert_eq(self_destruct_pdf, no_self_destruct_pdf) + + def test_propagates_spark_exception(self): + self.check_propagates_spark_exception() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_toPandas_batch_order(self): + super().test_toPandas_batch_order() + + def test_toPandas_empty_df_arrow_enabled(self): + self.check_toPandas_empty_df_arrow_enabled(True) + + def test_create_data_frame_to_pandas_timestamp_ntz(self): + self.check_create_data_frame_to_pandas_timestamp_ntz(True) + + def test_create_data_frame_to_pandas_day_time_internal(self): + self.check_create_data_frame_to_pandas_day_time_internal(True) + + def test_toPandas_respect_session_timezone(self): + self.check_toPandas_respect_session_timezone(True) + + def test_toPandas_with_array_type(self): + self.check_toPandas_with_array_type(True) + + @unittest.skip("Spark Connect does not support fallback.") + def test_toPandas_fallback_disabled(self): + super().test_toPandas_fallback_disabled() + + @unittest.skip("Spark Connect does not support fallback.") + def test_toPandas_fallback_enabled(self): + super().test_toPandas_fallback_enabled() + + def test_toPandas_with_map_type(self): + self.check_toPandas_with_map_type(True) + + def test_toPandas_with_map_type_nulls(self): + self.check_toPandas_with_map_type_nulls(True) + + def test_createDataFrame_with_array_type(self): + self.check_createDataFrame_with_array_type(True) + + def test_createDataFrame_with_int_col_names(self): + self.check_createDataFrame_with_int_col_names(True) + + def test_timestamp_nat(self): + self.check_timestamp_nat(True) + + def test_toPandas_error(self): + self.check_toPandas_error(True) + + def test_toPandas_duplicate_field_names(self): + self.check_toPandas_duplicate_field_names(True) + + def test_createDataFrame_duplicate_field_names(self): + self.check_createDataFrame_duplicate_field_names(True) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43506): Enable ArrowTests.test_toPandas_empty_columns for pandas 2.0.0.", + ) + def test_toPandas_empty_columns(self): + self.check_toPandas_empty_columns(True) + + def test_createDataFrame_nested_timestamp(self): + self.check_createDataFrame_nested_timestamp(True) + + def test_toPandas_nested_timestamp(self): + self.check_toPandas_nested_timestamp(True) + + def test_createDataFrame_udt(self): + self.check_createDataFrame_udt(True) + + def test_toPandas_udt(self): + self.check_toPandas_udt(True) + + def test_create_dataframe_namedtuples(self): + self.check_create_dataframe_namedtuples(True) + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_arrow import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_map.py new file mode 100644 index 0000000..868aeae --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_map.py @@ -0,0 +1,38 @@ +# +# 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 unittest + +from pyspark.sql.tests.test_arrow_map import MapInArrowTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ArrowMapParityTests(MapInArrowTestsMixin, ReusedConnectTestCase): + def test_other_than_recordbatch_iter(self): + self.check_other_than_recordbatch_iter() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_arrow_map import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_python_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_python_udf.py new file mode 100644 index 0000000..e4a64a7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_arrow_python_udf.py @@ -0,0 +1,48 @@ +# +# 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 unittest + +from pyspark.sql.tests.connect.test_parity_udf import UDFParityTests +from pyspark.sql.tests.test_arrow_python_udf import PythonUDFArrowTestsMixin + + +class ArrowPythonUDFParityTests(UDFParityTests, PythonUDFArrowTestsMixin): + @classmethod + def setUpClass(cls): + super(ArrowPythonUDFParityTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "true") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDF.arrow.enabled") + finally: + super(ArrowPythonUDFParityTests, cls).tearDownClass() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_arrow_python_udf import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_catalog.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_catalog.py new file mode 100644 index 0000000..2b8a9d7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_catalog.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.test_catalog import CatalogTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class CatalogParityTests(CatalogTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_catalog import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_column.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_column.py new file mode 100644 index 0000000..5cce063 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_column.py @@ -0,0 +1,55 @@ +# +# 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 unittest + +from pyspark.testing.connectutils import should_test_connect + +if should_test_connect: + from pyspark import sql + from pyspark.sql.connect.column import Column + + # This is a hack to make the Column instance comparison works in `ColumnTestsMixin`. + # e.g., `isinstance(col, pyspark.sql.Column)`. + sql.Column = Column + +from pyspark.sql.tests.test_column import ColumnTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ColumnParityTests(ColumnTestsMixin, ReusedConnectTestCase): + # TODO(SPARK-42017): df["bad_key"] does not raise AnalysisException + @unittest.skip("Fails in Spark Connect, should enable.") + def test_access_column(self): + super().test_access_column() + + @unittest.skip("Requires JVM access.") + def test_validate_column_types(self): + super().test_validate_column_types() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_column import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_conf.py new file mode 100644 index 0000000..554f05f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_conf.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.test_conf import ConfTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ConfParityTests(ConfTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_conf import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_dataframe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_dataframe.py new file mode 100644 index 0000000..cc9f71f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_dataframe.py @@ -0,0 +1,96 @@ +# +# 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 unittest + +from pyspark.sql.tests.test_dataframe import DataFrameTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class DataFrameParityTests(DataFrameTestsMixin, ReusedConnectTestCase): + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_help_command(self): + super().test_help_command() + + # TODO(SPARK-41527): Implement DataFrame.observe + @unittest.skip("Fails in Spark Connect, should enable.") + def test_observe(self): + super().test_observe() + + # TODO(SPARK-41625): Support Structured Streaming + @unittest.skip("Fails in Spark Connect, should enable.") + def test_observe_str(self): + super().test_observe_str() + + # TODO(SPARK-41873): Implement DataFrame `pandas_api` + @unittest.skip("Fails in Spark Connect, should enable.") + def test_pandas_api(self): + super().test_pandas_api() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_repartitionByRange_dataframe(self): + super().test_repartitionByRange_dataframe() + + @unittest.skip("Spark Connect does not SparkContext but the tests depend on them.") + def test_same_semantics_error(self): + super().test_same_semantics_error() + + # Spark Connect throws `IllegalArgumentException` when calling `collect` instead of `sample`. + def test_sample(self): + super().test_sample() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_toDF_with_schema_string(self): + super().test_toDF_with_schema_string() + + def test_to_local_iterator_not_fully_consumed(self): + self.check_to_local_iterator_not_fully_consumed() + + def test_to_pandas_for_array_of_struct(self): + # Spark Connect's implementation is based on Arrow. + super().check_to_pandas_for_array_of_struct(True) + + def test_to_pandas_from_null_dataframe(self): + self.check_to_pandas_from_null_dataframe() + + def test_to_pandas_on_cross_join(self): + self.check_to_pandas_on_cross_join() + + def test_to_pandas_from_empty_dataframe(self): + self.check_to_pandas_from_empty_dataframe() + + def test_to_pandas_with_duplicated_column_names(self): + self.check_to_pandas_with_duplicated_column_names() + + def test_to_pandas_from_mixed_dataframe(self): + self.check_to_pandas_from_mixed_dataframe() + + def test_toDF_with_string(self): + super().test_toDF_with_string() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_dataframe import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_datasources.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_datasources.py new file mode 100644 index 0000000..9d4ab15 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_datasources.py @@ -0,0 +1,44 @@ +# +# 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 unittest + +from pyspark.sql.tests.test_datasources import DataSourcesTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class DataSourcesParityTests(DataSourcesTestsMixin, ReusedConnectTestCase): + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_csv_sampling_ratio(self): + super().test_csv_sampling_ratio() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_json_sampling_ratio(self): + super().test_json_sampling_ratio() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_datasources import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_errors.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_errors.py new file mode 100644 index 0000000..37f5b90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_errors.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.test_errors import ErrorsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ErrorsParityTests(ErrorsTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_errors import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_functions.py new file mode 100644 index 0000000..35ddf96 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_functions.py @@ -0,0 +1,59 @@ +# +# 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 unittest + +from pyspark.errors.exceptions.connect import SparkConnectException +from pyspark.sql.connect.column import Column +from pyspark.sql.tests.test_functions import FunctionsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class FunctionsParityTests(FunctionsTestsMixin, ReusedConnectTestCase): + def test_assert_true(self): + self.check_assert_true(SparkConnectException) + + @unittest.skip("Spark Connect does not support Spark Context but the test depends on that.") + def test_basic_functions(self): + super().test_basic_functions() + + @unittest.skip("Spark Connect does not support Spark Context but the test depends on that.") + def test_function_parity(self): + super().test_function_parity() + + @unittest.skip("Spark Connect does not support Spark Context but the test depends on that.") + def test_input_file_name_reset_for_rdd(self): + super().test_input_file_name_reset_for_rdd() + + def test_raise_error(self): + self.check_raise_error(SparkConnectException) + + def test_sorting_functions_with_column(self): + self.check_sorting_functions_with_column(Column) + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_functions import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_group.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_group.py new file mode 100644 index 0000000..1dba236 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_group.py @@ -0,0 +1,36 @@ +# +# 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.sql.tests.test_group import GroupTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GroupParityTests(GroupTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_group import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_cogrouped_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_cogrouped_map.py new file mode 100644 index 0000000..e4f387a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_cogrouped_map.py @@ -0,0 +1,59 @@ +# +# 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 unittest + +from pyspark.sql.tests.pandas.test_pandas_cogrouped_map import CogroupedApplyInPandasTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class CogroupedApplyInPandasTests(CogroupedApplyInPandasTestsMixin, ReusedConnectTestCase): + @unittest.skip("Fails in Spark Connect, should enable.") + def test_different_group_key_cardinality(self): + self.check_different_group_key_cardinality() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_wrong_args(self): + self.check_wrong_args() + + def test_apply_in_pandas_not_returning_pandas_dataframe(self): + self.check_apply_in_pandas_not_returning_pandas_dataframe() + + def test_apply_in_pandas_returning_wrong_column_names(self): + self.check_apply_in_pandas_returning_wrong_column_names() + + def test_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + self.check_apply_in_pandas_returning_no_column_names_and_wrong_amount() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_apply_in_pandas_returning_incompatible_type(self): + self.check_apply_in_pandas_returning_incompatible_type() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_wrong_return_type(self): + self.check_wrong_return_type() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_cogrouped_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map.py new file mode 100644 index 0000000..fbf564b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map.py @@ -0,0 +1,74 @@ +# +# 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 unittest + +from pyspark.sql.tests.pandas.test_pandas_grouped_map import GroupedApplyInPandasTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GroupedApplyInPandasTests(GroupedApplyInPandasTestsMixin, ReusedConnectTestCase): + # TODO(SPARK-42857): Support CreateDataFrame from Decimal128 + @unittest.skip("Fails in Spark Connect, should enable.") + def test_supported_types(self): + super().test_supported_types() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_wrong_return_type(self): + self.check_wrong_return_type() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_wrong_args(self): + self.check_wrong_args() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_unsupported_types(self): + self.check_unsupported_types() + + def test_register_grouped_map_udf(self): + self.check_register_grouped_map_udf() + + def test_column_order(self): + self.check_column_order() + + def test_apply_in_pandas_returning_wrong_column_names(self): + self.check_apply_in_pandas_returning_wrong_column_names() + + def test_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + self.check_apply_in_pandas_returning_no_column_names_and_wrong_amount() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_apply_in_pandas_returning_incompatible_type(self): + self.check_apply_in_pandas_returning_incompatible_type() + + def test_apply_in_pandas_not_returning_pandas_dataframe(self): + self.check_apply_in_pandas_not_returning_pandas_dataframe() + + @unittest.skip("Spark Connect doesn't support RDD but the test depends on it.") + def test_grouped_with_empty_partition(self): + super().test_grouped_with_empty_partition() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_grouped_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map_with_state.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map_with_state.py new file mode 100644 index 0000000..dc3bdf2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_grouped_map_with_state.py @@ -0,0 +1,62 @@ +# +# 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 unittest + +from pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state import ( + GroupedApplyInPandasWithStateTestsMixin, +) +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GroupedApplyInPandasWithStateTests( + GroupedApplyInPandasWithStateTestsMixin, ReusedConnectTestCase +): + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic(self): + super().test_apply_in_pandas_with_state_basic() + + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic_no_state(self): + super().test_apply_in_pandas_with_state_basic() + + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic_no_state_no_data(self): + super().test_apply_in_pandas_with_state_basic() + + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic_more_data(self): + super().test_apply_in_pandas_with_state_basic() + + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic_fewer_data(self): + super().test_apply_in_pandas_with_state_basic() + + @unittest.skip("foreachBatch will be supported in SPARK-42944.") + def test_apply_in_pandas_with_state_basic_with_null(self): + super().test_apply_in_pandas_with_state_basic() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_grouped_map_with_state import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_map.py new file mode 100644 index 0000000..6ff9b0c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_map.py @@ -0,0 +1,58 @@ +# +# 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 unittest + + +from pyspark.sql.tests.pandas.test_pandas_map import MapInPandasTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class MapInPandasParityTests(MapInPandasTestsMixin, ReusedConnectTestCase): + def test_other_than_dataframe_iter(self): + self.check_other_than_dataframe_iter() + + def test_dataframes_with_other_column_names(self): + self.check_dataframes_with_other_column_names() + + def test_dataframes_with_duplicate_column_names(self): + self.check_dataframes_with_duplicate_column_names() + + def test_dataframes_with_less_columns(self): + self.check_dataframes_with_less_columns() + + @unittest.skip("Fails in Spark Connect, should enable.") + def test_dataframes_with_incompatible_types(self): + self.check_dataframes_with_incompatible_types() + + def test_empty_dataframes_with_less_columns(self): + self.check_empty_dataframes_with_less_columns() + + def test_empty_dataframes_with_other_columns(self): + self.check_empty_dataframes_with_other_columns() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_pandas_map import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf.py new file mode 100644 index 0000000..09d1ab3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf.py @@ -0,0 +1,70 @@ +# +# 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.sql.connect.types import UnparsedDataType +from pyspark.sql.functions import pandas_udf, PandasUDFType +from pyspark.sql.tests.pandas.test_pandas_udf import PandasUDFTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class PandasUDFParityTests(PandasUDFTestsMixin, ReusedConnectTestCase): + def test_udf_wrong_arg(self): + self.check_udf_wrong_arg() + + def test_pandas_udf_decorator_with_return_type_string(self): + @pandas_udf("v double", PandasUDFType.GROUPED_MAP) + def foo(x): + return x + + self.assertEqual(foo.returnType, UnparsedDataType("v double")) + self.assertEqual(foo.evalType, PandasUDFType.GROUPED_MAP) + + @pandas_udf(returnType="double", functionType=PandasUDFType.SCALAR) + def foo(x): + return x + + self.assertEqual(foo.returnType, UnparsedDataType("double")) + self.assertEqual(foo.evalType, PandasUDFType.SCALAR) + + def test_pandas_udf_basic_with_return_type_string(self): + udf = pandas_udf(lambda x: x, "double", PandasUDFType.SCALAR) + self.assertEqual(udf.returnType, UnparsedDataType("double")) + self.assertEqual(udf.evalType, PandasUDFType.SCALAR) + + udf = pandas_udf(lambda x: x, "v double", PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, UnparsedDataType("v double")) + self.assertEqual(udf.evalType, PandasUDFType.GROUPED_MAP) + + udf = pandas_udf(lambda x: x, "v double", functionType=PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, UnparsedDataType("v double")) + self.assertEqual(udf.evalType, PandasUDFType.GROUPED_MAP) + + udf = pandas_udf(lambda x: x, returnType="v double", functionType=PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, UnparsedDataType("v double")) + self.assertEqual(udf.evalType, PandasUDFType.GROUPED_MAP) + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_pandas_udf import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_grouped_agg.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_grouped_agg.py new file mode 100644 index 0000000..98fad42 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_grouped_agg.py @@ -0,0 +1,50 @@ +# +# 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 unittest + +from pyspark.sql.tests.pandas.test_pandas_udf_grouped_agg import GroupedAggPandasUDFTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class PandasUDFGroupedAggParityTests(GroupedAggPandasUDFTestsMixin, ReusedConnectTestCase): + # TODO(SPARK-43727): Parity returnType check in Spark Connect + @unittest.skip("Fails in Spark Connect, should enable.") + def test_unsupported_types(self): + self.check_unsupported_types() + + def test_invalid_args(self): + self.check_invalid_args() + + @unittest.skip("Spark Connect doesn't support RDD but the test depends on it.") + def test_grouped_with_empty_partition(self): + super().test_grouped_with_empty_partition() + + @unittest.skip("Spark Connect does not support convert UNPARSED to catalyst types.") + def test_manual(self): + super().test_manual() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_udf_grouped_agg import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_scalar.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_scalar.py new file mode 100644 index 0000000..c950ca2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_scalar.py @@ -0,0 +1,64 @@ +# +# 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 unittest +from pyspark.sql.tests.pandas.test_pandas_udf_scalar import ScalarPandasUDFTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class PandasUDFScalarParityTests(ScalarPandasUDFTestsMixin, ReusedConnectTestCase): + def test_nondeterministic_vectorized_udf_in_aggregate(self): + self.check_nondeterministic_analysis_exception() + + @unittest.skip("Spark Connect doesn't support RDD but the test depends on it.") + def test_vectorized_udf_empty_partition(self): + super().test_vectorized_udf_empty_partition() + + @unittest.skip("Spark Connect doesn't support RDD but the test depends on it.") + def test_vectorized_udf_struct_with_empty_partition(self): + super().test_vectorized_udf_struct_with_empty_partition() + + def test_vectorized_udf_exception(self): + self.check_vectorized_udf_exception() + + def test_vectorized_udf_nested_struct(self): + self.check_vectorized_udf_nested_struct() + + def test_vectorized_udf_return_scalar(self): + self.check_vectorized_udf_return_scalar() + + def test_scalar_iter_udf_close(self): + self.check_scalar_iter_udf_close() + + # TODO(SPARK-43727): Parity returnType check in Spark Connect + @unittest.skip("Fails in Spark Connect, should enable.") + def test_vectorized_udf_wrong_return_type(self): + self.check_vectorized_udf_wrong_return_type() + + def test_vectorized_udf_invalid_length(self): + self.check_vectorized_udf_invalid_length() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_udf_scalar import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_window.py new file mode 100644 index 0000000..af7ff14 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_pandas_udf_window.py @@ -0,0 +1,40 @@ +# +# 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 unittest + +from pyspark.sql.tests.pandas.test_pandas_udf_window import WindowPandasUDFTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class PandasUDFWindowParityTests(WindowPandasUDFTestsMixin, ReusedConnectTestCase): + # TODO(SPARK-43734): Expression "(v)" within a window function doesn't raise a + # AnalysisException + @unittest.skip("Fails in Spark Connect, should enable.") + def test_invalid_args(self): + self.check_invalid_args() + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_pandas_udf_window import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_readwriter.py new file mode 100644 index 0000000..2fa3f79 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_readwriter.py @@ -0,0 +1,46 @@ +# +# 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 unittest + +from pyspark.sql.connect.readwriter import DataFrameWriterV2 +from pyspark.sql.tests.test_readwriter import ReadwriterTestsMixin, ReadwriterV2TestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ReadwriterParityTests(ReadwriterTestsMixin, ReusedConnectTestCase): + pass + + +class ReadwriterV2ParityTests(ReadwriterV2TestsMixin, ReusedConnectTestCase): + def test_api(self): + self.check_api(DataFrameWriterV2) + + def test_partitioning_functions(self): + self.check_partitioning_functions(DataFrameWriterV2) + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_readwriter import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_serde.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_serde.py new file mode 100644 index 0000000..dc9b437 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_serde.py @@ -0,0 +1,44 @@ +# +# 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 unittest + +from pyspark.sql.tests.test_serde import SerdeTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class SerdeParityTests(SerdeTestsMixin, ReusedConnectTestCase): + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_int_array_serialization(self): + super().test_int_array_serialization() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_serialize_nested_array_and_map(self): + super().test_serialize_nested_array_and_map() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_serde import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_types.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_types.py new file mode 100644 index 0000000..807c295 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_types.py @@ -0,0 +1,100 @@ +# +# 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 unittest + +from pyspark.sql.tests.test_types import TypesTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class TypesParityTests(TypesTestsMixin, ReusedConnectTestCase): + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_apply_schema(self): + super().test_apply_schema() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_apply_schema_to_dict_and_rows(self): + super().test_apply_schema_to_dict_and_rows() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_apply_schema_to_row(self): + super().test_apply_schema_to_dict_and_rows() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_create_dataframe_schema_mismatch(self): + super().test_create_dataframe_schema_mismatch() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_array_element_type_empty(self): + super().test_infer_array_element_type_empty() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_array_element_type_with_struct(self): + super().test_infer_array_element_type_with_struct() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_array_merge_element_types_with_rdd(self): + super().test_infer_array_merge_element_types_with_rdd() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_binary_type(self): + super().test_infer_binary_type() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_long_type(self): + super().test_infer_long_type() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_nested_dict_as_struct_with_rdd(self): + super().test_infer_nested_dict_as_struct_with_rdd() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_nested_schema(self): + super().test_infer_nested_schema() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_schema(self): + super().test_infer_schema() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_schema_to_local(self): + super().test_infer_schema_to_local() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_infer_schema_upcast_int_to_string(self): + super().test_infer_schema_upcast_int_to_string() + + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") + def test_rdd_with_udt(self): + super().test_rdd_with_udt() + + @unittest.skip("Requires JVM access.") + def test_udt(self): + super().test_udt() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_types import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udf.py new file mode 100644 index 0000000..1be7d69 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udf.py @@ -0,0 +1,100 @@ +# +# 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 unittest + +from pyspark.testing.connectutils import should_test_connect + +if should_test_connect: + from pyspark import sql + from pyspark.sql.connect.udf import UserDefinedFunction + + sql.udf.UserDefinedFunction = UserDefinedFunction + +from pyspark.sql.tests.test_udf import BaseUDFTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class UDFParityTests(BaseUDFTestsMixin, ReusedConnectTestCase): + @unittest.skip("Spark Connect does not support mapPartitions() but the test depends on it.") + def test_worker_original_stdin_closed(self): + super().test_worker_original_stdin_closed() + + @unittest.skip( + "Spark Connect does not support reading from Hadoop RDD but the test depends on it." + ) + def test_udf_with_input_file_name_for_hadooprdd(self): + super().test_udf_with_input_file_name_for_hadooprdd() + + @unittest.skip("Spark Connect does not support accumulator but the test depends on it.") + def test_same_accumulator_in_udfs(self): + super().test_same_accumulator_in_udfs() + + @unittest.skip("Spark Connect does not support spark.conf but the test depends on it.") + def test_udf_timestamp_ntz(self): + super().test_udf_timestamp_ntz() + + @unittest.skip("Spark Connect does not support broadcast but the test depends on it.") + def test_broadcast_in_udf(self): + super().test_broadcast_in_udf() + + @unittest.skip("Spark Connect does not support cache() but the test depends on it.") + def test_udf_cache(self): + super().test_udf_cache() + + @unittest.skip("Requires JVM access.") + def test_udf_defers_judf_initialization(self): + super().test_udf_defers_judf_initialization() + + @unittest.skip("Requires JVM access.") + def test_nondeterministic_udf3(self): + super().test_nondeterministic_udf3() + + def test_nondeterministic_udf_in_aggregate(self): + self.check_nondeterministic_udf_in_aggregate() + + def test_udf_registration_return_type_not_none(self): + self.check_udf_registration_return_type_not_none() + + @unittest.skip("Spark Connect doesn't support RDD but the test depends on it.") + def test_worker_original_stdin_closed(self): + super().test_worker_original_stdin_closed() + + @unittest.skip("Spark Connect does not support SQLContext but the test depends on it.") + def test_udf_on_sql_context(self): + super().test_udf_on_sql_context() + + @unittest.skip("Spark Connect does not support SQLContext but the test depends on it.") + def test_non_existed_udf_with_sql_context(self): + super().test_non_existed_udf_with_sql_context() + + @unittest.skip("Spark Connect does not support SQLContext but the test depends on it.") + def test_udf_registration_returns_udf_on_sql_context(self): + super().test_udf_registration_returns_udf_on_sql_context() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_udf import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udtf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udtf.py new file mode 100644 index 0000000..1222b1b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_parity_udtf.py @@ -0,0 +1,85 @@ +# +# 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.testing.connectutils import should_test_connect + +if should_test_connect: + from pyspark import sql + from pyspark.sql.connect.udtf import UserDefinedTableFunction + + sql.udtf.UserDefinedTableFunction = UserDefinedTableFunction + +from pyspark.sql.connect.functions import lit, udtf +from pyspark.sql.tests.test_udtf import BaseUDTFTestsMixin, UDTFArrowTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.errors.exceptions.connect import SparkConnectGrpcException + + +class UDTFParityTests(BaseUDTFTestsMixin, ReusedConnectTestCase): + @classmethod + def setUpClass(cls): + super(UDTFParityTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "false") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled") + finally: + super(UDTFParityTests, cls).tearDownClass() + + # TODO: use PySpark error classes instead of SparkConnectGrpcException + + def test_struct_output_type_casting_row(self): + self.check_struct_output_type_casting_row(SparkConnectGrpcException) + + def test_udtf_with_invalid_return_type(self): + @udtf(returnType="int") + class TestUDTF: + def eval(self, a: int): + yield a + 1, + + with self.assertRaisesRegex( + SparkConnectGrpcException, "Invalid Python user-defined table function return type." + ): + TestUDTF(lit(1)).collect() + + +class ArrowUDTFParityTests(UDTFArrowTestsMixin, UDTFParityTests): + @classmethod + def setUpClass(cls): + super(ArrowUDTFParityTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled") + finally: + super(ArrowUDTFParityTests, cls).tearDownClass() + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_udtf import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_session.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_session.py new file mode 100644 index 0000000..131b1b8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_session.py @@ -0,0 +1,181 @@ +# +# 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 +import time +import unittest +from typing import Optional + +from pyspark.sql.connect.client import ChannelBuilder +from pyspark.sql.connect.session import SparkSession as RemoteSparkSession +from pyspark.testing.connectutils import should_test_connect + +if should_test_connect: + from pyspark.testing.connectutils import ReusedConnectTestCase + + +class CustomChannelBuilder(ChannelBuilder): + @property + def userId(self) -> Optional[str]: + return "abc" + + +class SparkSessionTestCase(unittest.TestCase): + def test_fails_to_create_session_without_remote_and_channel_builder(self): + with self.assertRaises(ValueError): + RemoteSparkSession.builder.getOrCreate() + + def test_fails_to_create_when_both_remote_and_channel_builder_are_specified(self): + with self.assertRaises(ValueError): + ( + RemoteSparkSession.builder.channelBuilder(CustomChannelBuilder("sc://localhost")) + .remote("sc://localhost") + .getOrCreate() + ) + + def test_creates_session_with_channel_builder(self): + test_session = RemoteSparkSession.builder.channelBuilder( + CustomChannelBuilder("sc://other") + ).getOrCreate() + host = test_session.client.host + test_session.stop() + + self.assertEqual("other", host) + + def test_creates_session_with_remote(self): + test_session = RemoteSparkSession.builder.remote("sc://other").getOrCreate() + host = test_session.client.host + test_session.stop() + + self.assertEqual("other", host) + + def test_session_stop(self): + session = RemoteSparkSession.builder.remote("sc://other").getOrCreate() + + self.assertFalse(session.is_stopped) + session.stop() + self.assertTrue(session.is_stopped) + + def test_session_create_sets_active_session(self): + session = RemoteSparkSession.builder.remote("sc://abc").create() + session2 = RemoteSparkSession.builder.remote("sc://other").getOrCreate() + + self.assertIs(session, session2) + session.stop() + + +class JobCancellationTests(ReusedConnectTestCase): + def test_tags(self): + self.spark.clearTags() + self.spark.addTag("a") + self.assertEqual(self.spark.getTags(), {"a"}) + self.spark.addTag("b") + self.spark.removeTag("a") + self.assertEqual(self.spark.getTags(), {"b"}) + self.spark.addTag("c") + self.spark.clearTags() + self.assertEqual(self.spark.getTags(), set()) + self.spark.clearTags() + + def test_interrupt_tag(self): + thread_ids = range(4) + self.check_job_cancellation( + lambda job_group: self.spark.addTag(job_group), + lambda job_group: self.spark.interruptTag(job_group), + thread_ids, + [i for i in thread_ids if i % 2 == 0], + [i for i in thread_ids if i % 2 != 0], + ) + self.spark.clearTags() + + def test_interrupt_all(self): + thread_ids = range(4) + self.check_job_cancellation( + lambda job_group: None, + lambda job_group: self.spark.interruptAll(), + thread_ids, + thread_ids, + [], + ) + self.spark.clearTags() + + def check_job_cancellation( + self, setter, canceller, thread_ids, thread_ids_to_cancel, thread_ids_to_run + ): + + job_id_a = "job_ids_to_cancel" + job_id_b = "job_ids_to_run" + threads = [] + + # A list which records whether job is cancelled. + # The index of the array is the thread index which job run in. + is_job_cancelled = [False for _ in thread_ids] + + def run_job(job_id, index): + """ + Executes a job with the group ``job_group``. Each job waits for 3 seconds + and then exits. + """ + try: + setter(job_id) + + def func(itr): + for pdf in itr: + time.sleep(pdf._1.iloc[0]) + yield pdf + + self.spark.createDataFrame([[20]]).repartition(1).mapInPandas( + func, schema="_1 LONG" + ).collect() + is_job_cancelled[index] = False + except Exception: + # Assume that exception means job cancellation. + is_job_cancelled[index] = True + + # Test if job succeeded when not cancelled. + run_job(job_id_a, 0) + self.assertFalse(is_job_cancelled[0]) + self.spark.clearTags() + + # Run jobs + for i in thread_ids_to_cancel: + t = threading.Thread(target=run_job, args=(job_id_a, i)) + t.start() + threads.append(t) + + for i in thread_ids_to_run: + t = threading.Thread(target=run_job, args=(job_id_b, i)) + t.start() + threads.append(t) + + # Wait to make sure all jobs are executed. + time.sleep(10) + # And then, cancel one job group. + canceller(job_id_a) + + # Wait until all threads launching jobs are finished. + for t in threads: + t.join() + + for i in thread_ids_to_cancel: + self.assertTrue( + is_job_cancelled[i], "Thread {i}: Job in group A was not cancelled.".format(i=i) + ) + + for i in thread_ids_to_run: + self.assertFalse( + is_job_cancelled[i], "Thread {i}: Job in group B did not succeeded.".format(i=i) + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_utils.py new file mode 100644 index 0000000..917cb58 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/connect/test_utils.py @@ -0,0 +1,36 @@ +# +# 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.testing.connectutils import ReusedConnectTestCase +from pyspark.sql.tests.test_utils import UtilsTestsMixin + + +class ConnectUtilsTests(ReusedConnectTestCase, UtilsTestsMixin): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_utils import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py new file mode 100644 index 0000000..b867156 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py @@ -0,0 +1,588 @@ +# +# 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 unittest +from typing import cast + +from pyspark.sql.functions import array, explode, col, lit, udf, pandas_udf, sum +from pyspark.sql.types import ( + ArrayType, + DoubleType, + LongType, + StructType, + StructField, + YearMonthIntervalType, + Row, +) +from pyspark.sql.window import Window +from pyspark.errors import IllegalArgumentException, PythonException +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + +if have_pandas: + import pandas as pd + from pandas.testing import assert_frame_equal + +if have_pyarrow: + import pyarrow as pa # noqa: F401 + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class CogroupedApplyInPandasTestsMixin: + @property + def data1(self): + return ( + self.spark.range(10) + .withColumn("ks", array([lit(i) for i in range(20, 30)])) + .withColumn("k", explode(col("ks"))) + .withColumn("v", col("k") * 10) + .drop("ks") + ) + + @property + def data2(self): + return ( + self.spark.range(10) + .withColumn("ks", array([lit(i) for i in range(20, 30)])) + .withColumn("k", explode(col("ks"))) + .withColumn("v2", col("k") * 100) + .drop("ks") + ) + + def test_simple(self): + self._test_merge(self.data1, self.data2) + + def test_left_group_empty(self): + left = self.data1.where(col("id") % 2 == 0) + self._test_merge(left, self.data2) + + def test_right_group_empty(self): + right = self.data2.where(col("id") % 2 == 0) + self._test_merge(self.data1, right) + + def test_different_schemas(self): + right = self.data2.withColumn("v3", lit("a")) + self._test_merge( + self.data1, right, output_schema="id long, k int, v int, v2 int, v3 string" + ) + + def test_different_keys(self): + left = self.data1 + right = self.data2 + + def merge_pandas(lft, rgt): + return pd.merge(lft.rename(columns={"id2": "id"}), rgt, on=["id", "k"]) + + result = ( + left.withColumnRenamed("id", "id2") + .groupby("id2") + .cogroup(right.groupby("id")) + .applyInPandas(merge_pandas, "id long, k int, v int, v2 int") + .sort(["id", "k"]) + .toPandas() + ) + + left = left.toPandas() + right = right.toPandas() + + expected = pd.merge(left, right, on=["id", "k"]).sort_values(by=["id", "k"]) + + assert_frame_equal(expected, result) + + def test_complex_group_by(self): + left = pd.DataFrame.from_dict({"id": [1, 2, 3], "k": [5, 6, 7], "v": [9, 10, 11]}) + + right = pd.DataFrame.from_dict({"id": [11, 12, 13], "k": [5, 6, 7], "v2": [90, 100, 110]}) + + left_gdf = self.spark.createDataFrame(left).groupby(col("id") % 2 == 0) + + right_gdf = self.spark.createDataFrame(right).groupby(col("id") % 2 == 0) + + def merge_pandas(lft, rgt): + return pd.merge(lft[["k", "v"]], rgt[["k", "v2"]], on=["k"]) + + result = ( + left_gdf.cogroup(right_gdf) + .applyInPandas(merge_pandas, "k long, v long, v2 long") + .sort(["k"]) + .toPandas() + ) + + expected = pd.DataFrame.from_dict({"k": [5, 6, 7], "v": [9, 10, 11], "v2": [90, 100, 110]}) + + assert_frame_equal(expected, result) + + def test_empty_group_by(self): + self._test_merge(self.data1, self.data2, by=[]) + + def test_different_group_key_cardinality(self): + with QuietTest(self.sc): + self.check_different_group_key_cardinality() + + def check_different_group_key_cardinality(self): + left = self.data1 + right = self.data2 + + def merge_pandas(lft, _): + return lft + + with self.assertRaisesRegex( + IllegalArgumentException, + "requirement failed: Cogroup keys must have same size: 2 != 1", + ): + (left.groupby("id", "k").cogroup(right.groupby("id"))).applyInPandas( + merge_pandas, "id long, k int, v int" + ) + + def test_apply_in_pandas_not_returning_pandas_dataframe(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_not_returning_pandas_dataframe() + + def check_apply_in_pandas_not_returning_pandas_dataframe(self): + self._test_merge_error( + fn=lambda lft, rgt: lft.size + rgt.size, + error_class=PythonException, + error_message_regex="Return type of the user-defined function " + "should be pandas.DataFrame, but is int64.", + ) + + def test_apply_in_pandas_returning_column_names(self): + self._test_merge(fn=lambda lft, rgt: pd.merge(lft, rgt, on=["id", "k"])) + + def test_apply_in_pandas_returning_no_column_names(self): + def merge_pandas(lft, rgt): + res = pd.merge(lft, rgt, on=["id", "k"]) + res.columns = range(res.columns.size) + return res + + self._test_merge(fn=merge_pandas) + + def test_apply_in_pandas_returning_column_names_sometimes(self): + def merge_pandas(lft, rgt): + res = pd.merge(lft, rgt, on=["id", "k"]) + if 0 in lft["id"] and lft["id"][0] % 2 == 0: + return res + res.columns = range(res.columns.size) + return res + + self._test_merge(fn=merge_pandas) + + def test_apply_in_pandas_returning_wrong_column_names(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_wrong_column_names() + + def check_apply_in_pandas_returning_wrong_column_names(self): + def merge_pandas(lft, rgt): + if 0 in lft["id"] and lft["id"][0] % 2 == 0: + lft["add"] = 0 + if 0 in rgt["id"] and rgt["id"][0] % 3 == 0: + rgt["more"] = 1 + return pd.merge(lft, rgt, on=["id", "k"]) + + self._test_merge_error( + fn=merge_pandas, + error_class=PythonException, + error_message_regex="Column names of the returned pandas.DataFrame " + "do not match specified schema. Unexpected: add, more.\n", + ) + + def test_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_no_column_names_and_wrong_amount() + + def check_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + def merge_pandas(lft, rgt): + if 0 in lft["id"] and lft["id"][0] % 2 == 0: + lft[3] = 0 + if 0 in rgt["id"] and rgt["id"][0] % 3 == 0: + rgt[3] = 1 + res = pd.merge(lft, rgt, on=["id", "k"]) + res.columns = range(res.columns.size) + return res + + self._test_merge_error( + fn=merge_pandas, + error_class=PythonException, + error_message_regex="Number of columns of the returned pandas.DataFrame " + "doesn't match specified schema. Expected: 4 Actual: 6\n", + ) + + def test_apply_in_pandas_returning_empty_dataframe(self): + def merge_pandas(lft, rgt): + if 0 in lft["id"] and lft["id"][0] % 2 == 0: + return pd.DataFrame() + if 0 in rgt["id"] and rgt["id"][0] % 3 == 0: + return pd.DataFrame() + return pd.merge(lft, rgt, on=["id", "k"]) + + self._test_merge_empty(fn=merge_pandas) + + def test_apply_in_pandas_returning_incompatible_type(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_incompatible_type() + + def check_apply_in_pandas_returning_incompatible_type(self): + for safely in [True, False]: + with self.subTest(convertToArrowArraySafely=safely), self.sql_conf( + {"spark.sql.execution.pandas.convertToArrowArraySafely": safely} + ): + # sometimes we see ValueErrors + with self.subTest(convert="string to double"): + expected = ( + r"ValueError: Exception thrown when converting pandas.Series \(object\) " + r"with name 'k' to Arrow Array \(double\)." + ) + if safely: + expected = expected + ( + " It can be caused by overflows or other " + "unsafe conversions warned by Arrow. Arrow safe type check " + "can be disabled by using SQL config " + "`spark.sql.execution.pandas.convertToArrowArraySafely`." + ) + self._test_merge_error( + fn=lambda lft, rgt: pd.DataFrame({"id": [1], "k": ["2.0"]}), + output_schema="id long, k double", + error_class=PythonException, + error_message_regex=expected, + ) + + # sometimes we see TypeErrors + with self.subTest(convert="double to string"): + expected = ( + r"TypeError: Exception thrown when converting pandas.Series \(float64\) " + r"with name 'k' to Arrow Array \(string\).\n" + ) + self._test_merge_error( + fn=lambda lft, rgt: pd.DataFrame({"id": [1], "k": [2.0]}), + output_schema="id long, k string", + error_class=PythonException, + error_message_regex=expected, + ) + + def test_mixed_scalar_udfs_followed_by_cogrouby_apply(self): + df = self.spark.range(0, 10).toDF("v1") + df = df.withColumn("v2", udf(lambda x: x + 1, "int")(df["v1"])).withColumn( + "v3", pandas_udf(lambda x: x + 2, "int")(df["v1"]) + ) + + result = ( + df.groupby() + .cogroup(df.groupby()) + .applyInPandas( + lambda x, y: pd.DataFrame([(x.sum().sum(), y.sum().sum())]), "sum1 int, sum2 int" + ) + .collect() + ) + + self.assertEqual(result[0]["sum1"], 165) + self.assertEqual(result[0]["sum2"], 165) + + def test_with_key_left(self): + self._test_with_key(self.data1, self.data1, isLeft=True) + + def test_with_key_right(self): + self._test_with_key(self.data1, self.data1, isLeft=False) + + def test_with_key_left_group_empty(self): + left = self.data1.where(col("id") % 2 == 0) + self._test_with_key(left, self.data1, isLeft=True) + + def test_with_key_right_group_empty(self): + right = self.data1.where(col("id") % 2 == 0) + self._test_with_key(self.data1, right, isLeft=False) + + def test_with_key_complex(self): + def left_assign_key(key, lft, _): + return lft.assign(key=key[0]) + + result = ( + self.data1.groupby(col("id") % 2 == 0) + .cogroup(self.data2.groupby(col("id") % 2 == 0)) + .applyInPandas(left_assign_key, "id long, k int, v int, key boolean") + .sort(["id", "k"]) + .toPandas() + ) + + expected = self.data1.toPandas() + expected = expected.assign(key=expected.id % 2 == 0) + + assert_frame_equal(expected, result) + + def test_wrong_return_type(self): + with QuietTest(self.sc): + self.check_wrong_return_type() + + def check_wrong_return_type(self): + # Test that we get a sensible exception invalid values passed to apply + self._test_merge_error( + fn=lambda l, r: l, + output_schema=( + StructType().add("id", LongType()).add("v", ArrayType(YearMonthIntervalType())) + ), + error_class=NotImplementedError, + error_message_regex="Invalid return type.*ArrayType.*YearMonthIntervalType", + ) + + def test_wrong_args(self): + with QuietTest(self.sc): + self.check_wrong_args() + + def check_wrong_args(self): + self.__test_merge_error( + fn=lambda: 1, + output_schema=StructType([StructField("d", DoubleType())]), + error_class=ValueError, + error_message_regex="Invalid function", + ) + + def test_case_insensitive_grouping_column(self): + # SPARK-31915: case-insensitive grouping column should work. + df1 = self.spark.createDataFrame([(1, 1)], ("column", "value")) + + row = ( + df1.groupby("ColUmn") + .cogroup(df1.groupby("COLUMN")) + .applyInPandas(lambda r, l: r + l, "column long, value long") + .first() + ) + self.assertEqual(row.asDict(), Row(column=2, value=2).asDict()) + + df2 = self.spark.createDataFrame([(1, 1)], ("column", "value")) + + row = ( + df1.groupby("ColUmn") + .cogroup(df2.groupby("COLUMN")) + .applyInPandas(lambda r, l: r + l, "column long, value long") + .first() + ) + self.assertEqual(row.asDict(), Row(column=2, value=2).asDict()) + + def test_self_join(self): + # SPARK-34319: self-join with FlatMapCoGroupsInPandas + df = self.spark.createDataFrame([(1, 1)], ("column", "value")) + + row = ( + df.groupby("ColUmn") + .cogroup(df.groupby("COLUMN")) + .applyInPandas(lambda r, l: r + l, "column long, value long") + ) + + row = row.join(row).first() + + self.assertEqual(row.asDict(), Row(column=2, value=2).asDict()) + + def test_with_window_function(self): + # SPARK-42168: a window function with same partition keys but differing key order + ids = 2 + days = 100 + vals = 10000 + parts = 10 + + id_df = self.spark.range(ids) + day_df = self.spark.range(days).withColumnRenamed("id", "day") + vals_df = self.spark.range(vals).withColumnRenamed("id", "value") + df = id_df.join(day_df).join(vals_df) + + left_df = df.withColumnRenamed("value", "left").repartition(parts).cache() + # SPARK-42132: this bug requires us to alias all columns from df here + right_df = ( + df.select(col("id").alias("id"), col("day").alias("day"), col("value").alias("right")) + .repartition(parts) + .cache() + ) + + # note the column order is different to the groupBy("id", "day") column order below + window = Window.partitionBy("day", "id") + + left_grouped_df = left_df.groupBy("id", "day") + right_grouped_df = right_df.withColumn("day_sum", sum(col("day")).over(window)).groupBy( + "id", "day" + ) + + def cogroup(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame: + return pd.DataFrame( + [ + { + "id": left["id"][0] + if not left.empty + else (right["id"][0] if not right.empty else None), + "day": left["day"][0] + if not left.empty + else (right["day"][0] if not right.empty else None), + "lefts": len(left.index), + "rights": len(right.index), + } + ] + ) + + df = left_grouped_df.cogroup(right_grouped_df).applyInPandas( + cogroup, schema="id long, day long, lefts integer, rights integer" + ) + + actual = df.orderBy("id", "day").take(days) + self.assertEqual(actual, [Row(0, day, vals, vals) for day in range(days)]) + + @staticmethod + def _test_with_key(left, right, isLeft): + def right_assign_key(key, lft, rgt): + return lft.assign(key=key[0]) if isLeft else rgt.assign(key=key[0]) + + result = ( + left.groupby("id") + .cogroup(right.groupby("id")) + .applyInPandas(right_assign_key, "id long, k int, v int, key long") + .toPandas() + ) + + expected = left.toPandas() if isLeft else right.toPandas() + expected = expected.assign(key=expected.id) + + assert_frame_equal(expected, result) + + def _test_merge_empty(self, fn): + left = self.data1.toPandas() + right = self.data2.toPandas() + + expected = pd.merge( + left[left["id"] % 2 != 0], right[right["id"] % 3 != 0], on=["id", "k"] + ).sort_values(by=["id", "k"]) + + self._test_merge(self.data1, self.data2, fn=fn, expected=expected) + + def _test_merge( + self, + left=None, + right=None, + by=["id"], + fn=lambda lft, rgt: pd.merge(lft, rgt, on=["id", "k"]), + output_schema="id long, k int, v int, v2 int", + expected=None, + ): + def fn_with_key(_, lft, rgt): + return fn(lft, rgt) + + # Test fn with and without key argument + with self.subTest("without key"): + self.__test_merge(left, right, by, fn, output_schema, expected) + with self.subTest("with key"): + self.__test_merge(left, right, by, fn_with_key, output_schema, expected) + + def __test_merge( + self, + left=None, + right=None, + by=["id"], + fn=lambda lft, rgt: pd.merge(lft, rgt, on=["id", "k"]), + output_schema="id long, k int, v int, v2 int", + expected=None, + ): + # Test fn as is, cf. _test_merge + left = self.data1 if left is None else left + right = self.data2 if right is None else right + + result = ( + left.groupby(*by) + .cogroup(right.groupby(*by)) + .applyInPandas(fn, output_schema) + .sort(["id", "k"]) + .toPandas() + ) + + left = left.toPandas() + right = right.toPandas() + + expected = ( + pd.merge(left, right, on=["id", "k"]).sort_values(by=["id", "k"]) + if expected is None + else expected + ) + + assert_frame_equal(expected, result) + + def _test_merge_error( + self, + error_class, + error_message_regex, + left=None, + right=None, + by=["id"], + fn=lambda lft, rgt: pd.merge(lft, rgt, on=["id", "k"]), + output_schema="id long, k int, v int, v2 int", + ): + def fn_with_key(_, lft, rgt): + return fn(lft, rgt) + + # Test fn with and without key argument + with self.subTest("without key"): + self.__test_merge_error( + left=left, + right=right, + by=by, + fn=fn, + output_schema=output_schema, + error_class=error_class, + error_message_regex=error_message_regex, + ) + with self.subTest("with key"): + self.__test_merge_error( + left=left, + right=right, + by=by, + fn=fn_with_key, + output_schema=output_schema, + error_class=error_class, + error_message_regex=error_message_regex, + ) + + def __test_merge_error( + self, + error_class, + error_message_regex, + left=None, + right=None, + by=["id"], + fn=lambda lft, rgt: pd.merge(lft, rgt, on=["id", "k"]), + output_schema="id long, k int, v int, v2 int", + ): + # Test fn as is, cf. _test_merge_error + with self.assertRaisesRegex(error_class, error_message_regex): + self.__test_merge(left, right, by, fn, output_schema) + + +class CogroupedApplyInPandasTests(CogroupedApplyInPandasTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_cogrouped_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py new file mode 100644 index 0000000..742b365 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py @@ -0,0 +1,869 @@ +# +# 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 datetime +import unittest + +from collections import OrderedDict +from decimal import Decimal +from typing import cast + +from pyspark.sql import Row +from pyspark.sql.functions import ( + array, + explode, + col, + lit, + udf, + sum, + pandas_udf, + PandasUDFType, + window, +) +from pyspark.sql.types import ( + IntegerType, + DoubleType, + ArrayType, + BinaryType, + ByteType, + LongType, + DecimalType, + ShortType, + FloatType, + StringType, + BooleanType, + StructType, + StructField, + NullType, + MapType, + YearMonthIntervalType, +) +from pyspark.errors import PythonException, PySparkTypeError +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + +if have_pandas: + import pandas as pd + from pandas.testing import assert_frame_equal + +if have_pyarrow: + import pyarrow as pa # noqa: F401 + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class GroupedApplyInPandasTestsMixin: + @property + def data(self): + return ( + self.spark.range(10) + .withColumn("vs", array([lit(i) for i in range(20, 30)])) + .withColumn("v", explode(col("vs"))) + .drop("vs") + ) + + def test_supported_types(self): + + values = [ + 1, + 2, + 3, + 4, + 5, + 1.1, + 2.2, + Decimal(1.123), + [1, 2, 2], + True, + "hello", + bytearray([0x01, 0x02]), + None, + ] + output_fields = [ + ("id", IntegerType()), + ("byte", ByteType()), + ("short", ShortType()), + ("int", IntegerType()), + ("long", LongType()), + ("float", FloatType()), + ("double", DoubleType()), + ("decim", DecimalType(10, 3)), + ("array", ArrayType(IntegerType())), + ("bool", BooleanType()), + ("str", StringType()), + ("bin", BinaryType()), + ("null", NullType()), + ] + + output_schema = StructType([StructField(*x) for x in output_fields]) + df = self.spark.createDataFrame([values], schema=output_schema) + + # Different forms of group map pandas UDF, results of these are the same + udf1 = pandas_udf( + lambda pdf: pdf.assign( + byte=pdf.byte * 2, + short=pdf.short * 2, + int=pdf.int * 2, + long=pdf.long * 2, + float=pdf.float * 2, + double=pdf.double * 2, + decim=pdf.decim * 2, + bool=False if pdf.bool else True, + str=pdf.str + "there", + array=pdf.array, + bin=pdf.bin, + null=pdf.null, + ), + output_schema, + PandasUDFType.GROUPED_MAP, + ) + + udf2 = pandas_udf( + lambda _, pdf: pdf.assign( + byte=pdf.byte * 2, + short=pdf.short * 2, + int=pdf.int * 2, + long=pdf.long * 2, + float=pdf.float * 2, + double=pdf.double * 2, + decim=pdf.decim * 2, + bool=False if pdf.bool else True, + str=pdf.str + "there", + array=pdf.array, + bin=pdf.bin, + null=pdf.null, + ), + output_schema, + PandasUDFType.GROUPED_MAP, + ) + + udf3 = pandas_udf( + lambda key, pdf: pdf.assign( + id=key[0], + byte=pdf.byte * 2, + short=pdf.short * 2, + int=pdf.int * 2, + long=pdf.long * 2, + float=pdf.float * 2, + double=pdf.double * 2, + decim=pdf.decim * 2, + bool=False if pdf.bool else True, + str=pdf.str + "there", + array=pdf.array, + bin=pdf.bin, + null=pdf.null, + ), + output_schema, + PandasUDFType.GROUPED_MAP, + ) + + result1 = df.groupby("id").apply(udf1).sort("id").toPandas() + expected1 = df.toPandas().groupby("id").apply(udf1.func).reset_index(drop=True) + + result2 = df.groupby("id").apply(udf2).sort("id").toPandas() + expected2 = expected1 + + result3 = df.groupby("id").apply(udf3).sort("id").toPandas() + expected3 = expected1 + + assert_frame_equal(expected1, result1) + assert_frame_equal(expected2, result2) + assert_frame_equal(expected3, result3) + + def test_array_type_correct(self): + df = self.data.withColumn("arr", array(col("id"))).repartition(1, "id") + + output_schema = StructType( + [ + StructField("id", LongType()), + StructField("v", IntegerType()), + StructField("arr", ArrayType(LongType())), + ] + ) + + udf = pandas_udf(lambda pdf: pdf, output_schema, PandasUDFType.GROUPED_MAP) + + result = df.groupby("id").apply(udf).sort("id").toPandas() + expected = df.toPandas().groupby("id").apply(udf.func).reset_index(drop=True) + assert_frame_equal(expected, result) + + def test_register_grouped_map_udf(self): + with QuietTest(self.sc): + self.check_register_grouped_map_udf() + + def check_register_grouped_map_udf(self): + foo_udf = pandas_udf(lambda x: x, "id long", PandasUDFType.GROUPED_MAP) + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.catalog.registerFunction("foo_udf", foo_udf) + + self.check_error( + exception=pe.exception, + error_class="INVALID_UDF_EVAL_TYPE", + message_parameters={ + "eval_type": "SQL_BATCHED_UDF, SQL_ARROW_BATCHED_UDF, SQL_SCALAR_PANDAS_UDF, " + "SQL_SCALAR_PANDAS_ITER_UDF or SQL_GROUPED_AGG_PANDAS_UDF" + }, + ) + + def test_decorator(self): + df = self.data + + @pandas_udf("id long, v int, v1 double, v2 long", PandasUDFType.GROUPED_MAP) + def foo(pdf): + return pdf.assign(v1=pdf.v * pdf.id * 1.0, v2=pdf.v + pdf.id) + + result = df.groupby("id").apply(foo).sort("id").toPandas() + expected = df.toPandas().groupby("id").apply(foo.func).reset_index(drop=True) + assert_frame_equal(expected, result) + + def test_coerce(self): + df = self.data + + foo = pandas_udf(lambda pdf: pdf, "id long, v double", PandasUDFType.GROUPED_MAP) + + result = df.groupby("id").apply(foo).sort("id").toPandas() + expected = df.toPandas().groupby("id").apply(foo.func).reset_index(drop=True) + expected = expected.assign(v=expected.v.astype("float64")) + assert_frame_equal(expected, result) + + def test_complex_groupby(self): + df = self.data + + @pandas_udf("id long, v int, norm double", PandasUDFType.GROUPED_MAP) + def normalize(pdf): + v = pdf.v + return pdf.assign(norm=(v - v.mean()) / v.std()) + + result = df.groupby(col("id") % 2 == 0).apply(normalize).sort("id", "v").toPandas() + pdf = df.toPandas() + expected = pdf.groupby(pdf["id"] % 2 == 0, as_index=False).apply(normalize.func) + expected = expected.sort_values(["id", "v"]).reset_index(drop=True) + expected = expected.assign(norm=expected.norm.astype("float64")) + assert_frame_equal(expected, result) + + def test_empty_groupby(self): + df = self.data + + @pandas_udf("id long, v int, norm double", PandasUDFType.GROUPED_MAP) + def normalize(pdf): + v = pdf.v + return pdf.assign(norm=(v - v.mean()) / v.std()) + + result = df.groupby().apply(normalize).sort("id", "v").toPandas() + pdf = df.toPandas() + expected = normalize.func(pdf) + expected = expected.sort_values(["id", "v"]).reset_index(drop=True) + expected = expected.assign(norm=expected.norm.astype("float64")) + assert_frame_equal(expected, result) + + def test_apply_in_pandas_not_returning_pandas_dataframe(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_not_returning_pandas_dataframe() + + def check_apply_in_pandas_not_returning_pandas_dataframe(self): + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be pandas.DataFrame, but is tuple.", + ): + self._test_apply_in_pandas(lambda key, pdf: key) + + @staticmethod + def stats_with_column_names(key, pdf): + # order of column can be different to applyInPandas schema when column names are given + return pd.DataFrame([(pdf.v.mean(),) + key], columns=["mean", "id"]) + + @staticmethod + def stats_with_no_column_names(key, pdf): + # columns must be in order of applyInPandas schema when no columns given + return pd.DataFrame([key + (pdf.v.mean(),)]) + + def test_apply_in_pandas_returning_column_names(self): + self._test_apply_in_pandas(GroupedApplyInPandasTestsMixin.stats_with_column_names) + + def test_apply_in_pandas_returning_no_column_names(self): + self._test_apply_in_pandas(GroupedApplyInPandasTestsMixin.stats_with_no_column_names) + + def test_apply_in_pandas_returning_column_names_sometimes(self): + def stats(key, pdf): + if key[0] % 2: + return GroupedApplyInPandasTestsMixin.stats_with_column_names(key, pdf) + else: + return GroupedApplyInPandasTestsMixin.stats_with_no_column_names(key, pdf) + + self._test_apply_in_pandas(stats) + + def test_apply_in_pandas_returning_wrong_column_names(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_wrong_column_names() + + def check_apply_in_pandas_returning_wrong_column_names(self): + with self.assertRaisesRegex( + PythonException, + "Column names of the returned pandas.DataFrame do not match specified schema. " + "Missing: mean. Unexpected: median, std.\n", + ): + self._test_apply_in_pandas( + lambda key, pdf: pd.DataFrame( + [key + (pdf.v.median(), pdf.v.std())], columns=["id", "median", "std"] + ) + ) + + def test_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_no_column_names_and_wrong_amount() + + def check_apply_in_pandas_returning_no_column_names_and_wrong_amount(self): + with self.assertRaisesRegex( + PythonException, + "Number of columns of the returned pandas.DataFrame doesn't match " + "specified schema. Expected: 2 Actual: 3\n", + ): + self._test_apply_in_pandas( + lambda key, pdf: pd.DataFrame([key + (pdf.v.mean(), pdf.v.std())]) + ) + + def test_apply_in_pandas_returning_empty_dataframe(self): + self._test_apply_in_pandas_returning_empty_dataframe(pd.DataFrame()) + + def test_apply_in_pandas_returning_incompatible_type(self): + with QuietTest(self.sc): + self.check_apply_in_pandas_returning_incompatible_type() + + def check_apply_in_pandas_returning_incompatible_type(self): + for safely in [True, False]: + with self.subTest(convertToArrowArraySafely=safely), self.sql_conf( + {"spark.sql.execution.pandas.convertToArrowArraySafely": safely} + ): + # sometimes we see ValueErrors + with self.subTest(convert="string to double"): + expected = ( + r"ValueError: Exception thrown when converting pandas.Series \(object\) " + r"with name 'mean' to Arrow Array \(double\)." + ) + if safely: + expected = expected + ( + " It can be caused by overflows or other " + "unsafe conversions warned by Arrow. Arrow safe type check " + "can be disabled by using SQL config " + "`spark.sql.execution.pandas.convertToArrowArraySafely`." + ) + with self.assertRaisesRegex(PythonException, expected + "\n"): + self._test_apply_in_pandas( + lambda key, pdf: pd.DataFrame([key + (str(pdf.v.mean()),)]), + output_schema="id long, mean double", + ) + + # sometimes we see TypeErrors + with self.subTest(convert="double to string"): + with self.assertRaisesRegex( + PythonException, + r"TypeError: Exception thrown when converting pandas.Series \(float64\) " + r"with name 'mean' to Arrow Array \(string\).\n", + ): + self._test_apply_in_pandas( + lambda key, pdf: pd.DataFrame([key + (pdf.v.mean(),)]), + output_schema="id long, mean string", + ) + + def test_datatype_string(self): + df = self.data + + foo_udf = pandas_udf( + lambda pdf: pdf.assign(v1=pdf.v * pdf.id * 1.0, v2=pdf.v + pdf.id), + "id long, v int, v1 double, v2 long", + PandasUDFType.GROUPED_MAP, + ) + + result = df.groupby("id").apply(foo_udf).sort("id").toPandas() + expected = df.toPandas().groupby("id").apply(foo_udf.func).reset_index(drop=True) + assert_frame_equal(expected, result) + + def test_wrong_return_type(self): + with QuietTest(self.sc): + self.check_wrong_return_type() + + def check_wrong_return_type(self): + with self.assertRaisesRegex( + NotImplementedError, + "Invalid return type.*grouped map Pandas UDF.*ArrayType.*YearMonthIntervalType", + ): + pandas_udf( + lambda pdf: pdf, + StructType().add("id", LongType()).add("v", ArrayType(YearMonthIntervalType())), + PandasUDFType.GROUPED_MAP, + ) + + def test_wrong_args(self): + with QuietTest(self.sc): + self.check_wrong_args() + + def check_wrong_args(self): + df = self.data + + with self.assertRaisesRegex(ValueError, "Invalid udf"): + df.groupby("id").apply(lambda x: x) + with self.assertRaisesRegex(ValueError, "Invalid udf"): + df.groupby("id").apply(udf(lambda x: x, DoubleType())) + with self.assertRaisesRegex(ValueError, "Invalid udf"): + df.groupby("id").apply(sum(df.v)) + with self.assertRaisesRegex(ValueError, "Invalid udf"): + df.groupby("id").apply(df.v + 1) + with self.assertRaisesRegex(ValueError, "Invalid function"): + df.groupby("id").apply( + pandas_udf(lambda: 1, StructType([StructField("d", DoubleType())])) + ) + with self.assertRaisesRegex(ValueError, "Invalid udf"): + df.groupby("id").apply(pandas_udf(lambda x, y: x, DoubleType())) + with self.assertRaisesRegex(ValueError, "Invalid udf.*GROUPED_MAP"): + df.groupby("id").apply(pandas_udf(lambda x, y: x, DoubleType(), PandasUDFType.SCALAR)) + + def test_unsupported_types(self): + with QuietTest(self.sc): + self.check_unsupported_types() + + def check_unsupported_types(self): + common_err_msg = "Invalid return type.*grouped map Pandas UDF.*" + unsupported_types = [ + StructField("array_struct", ArrayType(YearMonthIntervalType())), + StructField("map", MapType(StringType(), YearMonthIntervalType())), + ] + + for unsupported_type in unsupported_types: + with self.subTest(unsupported_type=unsupported_type.name): + schema = StructType([StructField("id", LongType(), True), unsupported_type]) + + with self.assertRaisesRegex(NotImplementedError, common_err_msg): + pandas_udf(lambda x: x, schema, PandasUDFType.GROUPED_MAP) + + # Regression test for SPARK-23314 + def test_timestamp_dst(self): + # Daylight saving time for Los Angeles for 2015 is Sun, Nov 1 at 2:00 am + dt = [ + datetime.datetime(2015, 11, 1, 0, 30), + datetime.datetime(2015, 11, 1, 1, 30), + datetime.datetime(2015, 11, 1, 2, 30), + ] + df = self.spark.createDataFrame(dt, "timestamp").toDF("time") + foo_udf = pandas_udf(lambda pdf: pdf, "time timestamp", PandasUDFType.GROUPED_MAP) + result = df.groupby("time").apply(foo_udf).sort("time") + assert_frame_equal(df.toPandas(), result.toPandas()) + + def test_udf_with_key(self): + import numpy as np + + df = self.data + pdf = df.toPandas() + + def foo1(key, pdf): + assert type(key) == tuple + assert type(key[0]) == np.int64 + + return pdf.assign( + v1=key[0], v2=pdf.v * key[0], v3=pdf.v * pdf.id, v4=pdf.v * pdf.id.mean() + ) + + def foo2(key, pdf): + assert type(key) == tuple + assert type(key[0]) == np.int64 + assert type(key[1]) == np.int32 + + return pdf.assign(v1=key[0], v2=key[1], v3=pdf.v * key[0], v4=pdf.v + key[1]) + + def foo3(key, pdf): + assert type(key) == tuple + assert len(key) == 0 + return pdf.assign(v1=pdf.v * pdf.id) + + # v2 is int because numpy.int64 * pd.Series results in pd.Series + # v3 is long because pd.Series * pd.Series results in pd.Series + udf1 = pandas_udf( + foo1, "id long, v int, v1 long, v2 int, v3 long, v4 double", PandasUDFType.GROUPED_MAP + ) + + udf2 = pandas_udf( + foo2, "id long, v int, v1 long, v2 int, v3 int, v4 int", PandasUDFType.GROUPED_MAP + ) + + udf3 = pandas_udf(foo3, "id long, v int, v1 long", PandasUDFType.GROUPED_MAP) + + # Test groupby column + result1 = df.groupby("id").apply(udf1).sort("id", "v").toPandas() + expected1 = ( + pdf.groupby("id", as_index=False) + .apply(lambda x: udf1.func((x.id.iloc[0],), x)) + .sort_values(["id", "v"]) + .reset_index(drop=True) + ) + assert_frame_equal(expected1, result1) + + # Test groupby expression + result2 = df.groupby(df.id % 2).apply(udf1).sort("id", "v").toPandas() + expected2 = ( + pdf.groupby(pdf.id % 2, as_index=False) + .apply(lambda x: udf1.func((x.id.iloc[0] % 2,), x)) + .sort_values(["id", "v"]) + .reset_index(drop=True) + ) + assert_frame_equal(expected2, result2) + + # Test complex groupby + result3 = df.groupby(df.id, df.v % 2).apply(udf2).sort("id", "v").toPandas() + expected3 = ( + pdf.groupby([pdf.id, pdf.v % 2], as_index=False) + .apply( + lambda x: udf2.func( + ( + x.id.iloc[0], + (x.v % 2).iloc[0], + ), + x, + ) + ) + .sort_values(["id", "v"]) + .reset_index(drop=True) + ) + assert_frame_equal(expected3, result3) + + # Test empty groupby + result4 = df.groupby().apply(udf3).sort("id", "v").toPandas() + expected4 = udf3.func((), pdf) + assert_frame_equal(expected4, result4) + + def test_column_order(self): + with QuietTest(self.sc): + self.check_column_order() + + def check_column_order(self): + + # Helper function to set column names from a list + def rename_pdf(pdf, names): + pdf.rename( + columns={old: new for old, new in zip(pd_result.columns, names)}, inplace=True + ) + + df = self.data + grouped_df = df.groupby("id") + grouped_pdf = df.toPandas().groupby("id", as_index=False) + + # Function returns a pdf with required column names, but order could be arbitrary using dict + def change_col_order(pdf): + # Constructing a DataFrame from a dict should result in the same order, + # but use OrderedDict to ensure the pdf column order is different than schema + return pd.DataFrame.from_dict( + OrderedDict([("id", pdf.id), ("u", pdf.v * 2), ("v", pdf.v)]) + ) + + ordered_udf = pandas_udf( + change_col_order, "id long, v int, u int", PandasUDFType.GROUPED_MAP + ) + + # The UDF result should assign columns by name from the pdf + result = grouped_df.apply(ordered_udf).sort("id", "v").select("id", "u", "v").toPandas() + pd_result = grouped_pdf.apply(change_col_order) + expected = pd_result.sort_values(["id", "v"]).reset_index(drop=True) + assert_frame_equal(expected, result) + + # Function returns a pdf with positional columns, indexed by range + def range_col_order(pdf): + # Create a DataFrame with positional columns, fix types to long + return pd.DataFrame(list(zip(pdf.id, pdf.v * 3, pdf.v)), dtype="int64") + + range_udf = pandas_udf( + range_col_order, "id long, u long, v long", PandasUDFType.GROUPED_MAP + ) + + # The UDF result uses positional columns from the pdf + result = grouped_df.apply(range_udf).sort("id", "v").select("id", "u", "v").toPandas() + pd_result = grouped_pdf.apply(range_col_order) + rename_pdf(pd_result, ["id", "u", "v"]) + expected = pd_result.sort_values(["id", "v"]).reset_index(drop=True) + assert_frame_equal(expected, result) + + # Function returns a pdf with columns indexed with integers + def int_index(pdf): + return pd.DataFrame(OrderedDict([(0, pdf.id), (1, pdf.v * 4), (2, pdf.v)])) + + int_index_udf = pandas_udf(int_index, "id long, u int, v int", PandasUDFType.GROUPED_MAP) + + # The UDF result should assign columns by position of integer index + result = grouped_df.apply(int_index_udf).sort("id", "v").select("id", "u", "v").toPandas() + pd_result = grouped_pdf.apply(int_index) + rename_pdf(pd_result, ["id", "u", "v"]) + expected = pd_result.sort_values(["id", "v"]).reset_index(drop=True) + assert_frame_equal(expected, result) + + @pandas_udf("id long, v int", PandasUDFType.GROUPED_MAP) + def column_name_typo(pdf): + return pd.DataFrame({"iid": pdf.id, "v": pdf.v}) + + @pandas_udf("id long, v decimal", PandasUDFType.GROUPED_MAP) + def invalid_positional_types(pdf): + return pd.DataFrame([(1, datetime.date(2020, 10, 5))]) + + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}): + with self.assertRaisesRegex( + PythonException, + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: id. Unexpected: iid.\n", + ): + grouped_df.apply(column_name_typo).collect() + with self.assertRaisesRegex(Exception, "[D|d]ecimal.*got.*date"): + grouped_df.apply(invalid_positional_types).collect() + + def test_positional_assignment_conf(self): + with self.sql_conf( + {"spark.sql.legacy.execution.pandas.groupedMap.assignColumnsByName": False} + ): + + @pandas_udf("a string, b float", PandasUDFType.GROUPED_MAP) + def foo(_): + return pd.DataFrame([("hi", 1)], columns=["x", "y"]) + + df = self.data + result = df.groupBy("id").apply(foo).select("a", "b").collect() + for r in result: + self.assertEqual(r.a, "hi") + self.assertEqual(r.b, 1) + + def test_self_join_with_pandas(self): + @pandas_udf("key long, col string", PandasUDFType.GROUPED_MAP) + def dummy_pandas_udf(df): + return df[["key", "col"]] + + df = self.spark.createDataFrame( + [Row(key=1, col="A"), Row(key=1, col="B"), Row(key=2, col="C")] + ) + df_with_pandas = df.groupBy("key").apply(dummy_pandas_udf) + + # this was throwing an AnalysisException before SPARK-24208 + res = df_with_pandas.alias("temp0").join( + df_with_pandas.alias("temp1"), col("temp0.key") == col("temp1.key") + ) + self.assertEqual(res.count(), 5) + + def test_mixed_scalar_udfs_followed_by_groupby_apply(self): + df = self.spark.range(0, 10).toDF("v1") + df = df.withColumn("v2", udf(lambda x: x + 1, "int")(df["v1"])).withColumn( + "v3", pandas_udf(lambda x: x + 2, "int")(df["v1"]) + ) + + result = df.groupby().apply( + pandas_udf( + lambda x: pd.DataFrame([x.sum().sum()]), "sum int", PandasUDFType.GROUPED_MAP + ) + ) + + self.assertEqual(result.collect()[0]["sum"], 165) + + def test_grouped_with_empty_partition(self): + data = [Row(id=1, x=2), Row(id=1, x=3), Row(id=2, x=4)] + expected = [Row(id=1, x=5), Row(id=1, x=5), Row(id=2, x=4)] + num_parts = len(data) + 1 + df = self.spark.createDataFrame(self.sc.parallelize(data, numSlices=num_parts)) + + f = pandas_udf( + lambda pdf: pdf.assign(x=pdf["x"].sum()), "id long, x int", PandasUDFType.GROUPED_MAP + ) + + result = df.groupBy("id").apply(f).collect() + self.assertEqual(result, expected) + + def test_grouped_over_window(self): + + data = [ + (0, 1, "2018-03-10T00:00:00+00:00", [0]), + (1, 2, "2018-03-11T00:00:00+00:00", [0]), + (2, 2, "2018-03-12T00:00:00+00:00", [0]), + (3, 3, "2018-03-15T00:00:00+00:00", [0]), + (4, 3, "2018-03-16T00:00:00+00:00", [0]), + (5, 3, "2018-03-17T00:00:00+00:00", [0]), + (6, 3, "2018-03-21T00:00:00+00:00", [0]), + ] + + expected = {0: [0], 1: [1, 2], 2: [1, 2], 3: [3, 4, 5], 4: [3, 4, 5], 5: [3, 4, 5], 6: [6]} + + df = self.spark.createDataFrame(data, ["id", "group", "ts", "result"]) + df = df.select(col("id"), col("group"), col("ts").cast("timestamp"), col("result")) + + def f(pdf): + # Assign each result element the ids of the windowed group + pdf["result"] = [pdf["id"]] * len(pdf) + return pdf + + result = ( + df.groupby("group", window("ts", "5 days")) + .applyInPandas(f, df.schema) + .select("id", "result") + .orderBy("id") + .collect() + ) + + self.assertListEqual([Row(id=key, result=val) for key, val in expected.items()], result) + + def test_grouped_over_window_with_key(self): + + data = [ + (0, 1, "2018-03-10T00:00:00+00:00", [0]), + (1, 2, "2018-03-11T00:00:00+00:00", [0]), + (2, 2, "2018-03-12T00:00:00+00:00", [0]), + (3, 3, "2018-03-15T00:00:00+00:00", [0]), + (4, 3, "2018-03-16T00:00:00+00:00", [0]), + (5, 3, "2018-03-17T00:00:00+00:00", [0]), + (6, 3, "2018-03-21T00:00:00+00:00", [0]), + ] + + timezone = self.spark.conf.get("spark.sql.session.timeZone") + expected_window = [ + { + key: ( + pd.Timestamp(ts) + .tz_localize(datetime.timezone.utc) + .tz_convert(timezone) + .tz_localize(None) + ) + for key, ts in w.items() + } + for w in [ + { + "start": datetime.datetime(2018, 3, 10, 0, 0), + "end": datetime.datetime(2018, 3, 15, 0, 0), + }, + { + "start": datetime.datetime(2018, 3, 15, 0, 0), + "end": datetime.datetime(2018, 3, 20, 0, 0), + }, + { + "start": datetime.datetime(2018, 3, 20, 0, 0), + "end": datetime.datetime(2018, 3, 25, 0, 0), + }, + ] + ] + + expected_key = { + 0: (1, expected_window[0]), + 1: (2, expected_window[0]), + 2: (2, expected_window[0]), + 3: (3, expected_window[1]), + 4: (3, expected_window[1]), + 5: (3, expected_window[1]), + 6: (3, expected_window[2]), + } + + # id -> array of group with len of num records in window + expected = {0: [1], 1: [2, 2], 2: [2, 2], 3: [3, 3, 3], 4: [3, 3, 3], 5: [3, 3, 3], 6: [3]} + + df = self.spark.createDataFrame(data, ["id", "group", "ts", "result"]) + df = df.select(col("id"), col("group"), col("ts").cast("timestamp"), col("result")) + + def f(key, pdf): + group = key[0] + window_range = key[1] + + # Make sure the key with group and window values are correct + for _, i in pdf.id.items(): + assert expected_key[i][0] == group, "{} != {}".format(expected_key[i][0], group) + assert expected_key[i][1] == window_range, "{} != {}".format( + expected_key[i][1], window_range + ) + + return pdf.assign(result=[[group] * len(pdf)] * len(pdf)) + + result = ( + df.groupby("group", window("ts", "5 days")) + .applyInPandas(f, df.schema) + .select("id", "result") + .orderBy("id") + .collect() + ) + + self.assertListEqual([Row(id=key, result=val) for key, val in expected.items()], result) + + def test_case_insensitive_grouping_column(self): + # SPARK-31915: case-insensitive grouping column should work. + def my_pandas_udf(pdf): + return pdf.assign(score=0.5) + + df = self.spark.createDataFrame([[1, 1]], ["column", "score"]) + row = ( + df.groupby("COLUMN") + .applyInPandas(my_pandas_udf, schema="column integer, score float") + .first() + ) + self.assertEqual(row.asDict(), Row(column=1, score=0.5).asDict()) + + def _test_apply_in_pandas(self, f, output_schema="id long, mean double"): + df = self.data + + result = ( + df.groupby("id").applyInPandas(f, schema=output_schema).sort("id", "mean").toPandas() + ) + expected = df.select("id").distinct().withColumn("mean", lit(24.5)).toPandas() + + assert_frame_equal(expected, result) + + def _test_apply_in_pandas_returning_empty_dataframe(self, empty_df): + """Tests some returned DataFrames are empty.""" + df = self.data + + def stats(key, pdf): + if key[0] % 2 == 0: + return GroupedApplyInPandasTestsMixin.stats_with_no_column_names(key, pdf) + return empty_df + + result = ( + df.groupby("id") + .applyInPandas(stats, schema="id long, mean double") + .sort("id", "mean") + .collect() + ) + + actual_ids = {row[0] for row in result} + expected_ids = {row[0] for row in self.data.collect() if row[0] % 2 == 0} + self.assertSetEqual(expected_ids, actual_ids) + self.assertEqual(len(expected_ids), len(result)) + for row in result: + self.assertEqual(24.5, row[1]) + + def _test_apply_in_pandas_returning_empty_dataframe_error(self, empty_df, error): + with QuietTest(self.sc): + with self.assertRaisesRegex(PythonException, error): + self._test_apply_in_pandas_returning_empty_dataframe(empty_df) + + +class GroupedApplyInPandasTests(GroupedApplyInPandasTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_grouped_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py new file mode 100644 index 0000000..e1ec979 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py @@ -0,0 +1,342 @@ +# +# 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 random +import shutil +import string +import sys +import tempfile + +import unittest +from typing import cast + +from pyspark import SparkConf +from pyspark.sql.streaming.state import GroupStateTimeout, GroupState +from pyspark.sql.types import ( + LongType, + StringType, + StructType, + StructField, + Row, +) +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import eventually + +if have_pandas: + import pandas as pd + +if have_pyarrow: + import pyarrow as pa # noqa: F401 + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class GroupedApplyInPandasWithStateTestsMixin: + @classmethod + def conf(cls): + cfg = SparkConf() + cfg.set("spark.sql.shuffle.partitions", "5") + return cfg + + def _test_apply_in_pandas_with_state_basic(self, func, check_results): + input_path = tempfile.mkdtemp() + + def prepare_test_resource(): + with open(input_path + "/text-test.txt", "w") as fw: + fw.write("hello\n") + fw.write("this\n") + + prepare_test_resource() + + df = self.spark.readStream.format("text").load(input_path) + + for q in self.spark.streams.active: + q.stop() + self.assertTrue(df.isStreaming) + + output_type = StructType( + [StructField("key", StringType()), StructField("countAsString", StringType())] + ) + state_type = StructType([StructField("c", LongType())]) + + q = ( + df.groupBy(df["value"]) + .applyInPandasWithState( + func, output_type, state_type, "Update", GroupStateTimeout.NoTimeout + ) + .writeStream.queryName("this_query") + .foreachBatch(check_results) + .outputMode("update") + .start() + ) + + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + q.processAllAvailable() + + def test_apply_in_pandas_with_state_basic(self): + def func(key, pdf_iter, state): + assert isinstance(state, GroupState) + + total_len = 0 + for pdf in pdf_iter: + total_len += len(pdf) + + state.update((total_len,)) + assert state.get[0] == 1 + yield pd.DataFrame({"key": [key[0]], "countAsString": [str(total_len)]}) + + def check_results(batch_df, _): + self.assertEqual( + set(batch_df.sort("key").collect()), + {Row(key="hello", countAsString="1"), Row(key="this", countAsString="1")}, + ) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_basic_no_state(self): + def func(key, pdf_iter, state): + assert isinstance(state, GroupState) + # 2 data rows + yield pd.DataFrame({"key": [key[0], "foo"], "countAsString": ["100", "222"]}) + + def check_results(batch_df, _): + self.assertEqual( + set(batch_df.sort("key").collect()), + { + Row(key="hello", countAsString="100"), + Row(key="this", countAsString="100"), + Row(key="foo", countAsString="222"), + }, + ) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_basic_no_state_no_data(self): + def func(key, pdf_iter, state): + assert isinstance(state, GroupState) + # 2 data rows + yield pd.DataFrame({"key": [], "countAsString": []}) + + def check_results(batch_df, _): + self.assertTrue(len(set(batch_df.sort("key").collect())) == 0) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_basic_more_data(self): + # Test data rows returned are more or fewer than state. + def func(key, pdf_iter, state): + state.update((1,)) + assert isinstance(state, GroupState) + # 3 rows + yield pd.DataFrame( + {"key": [key[0], "foo", key[0] + "_2"], "countAsString": ["1", "666", "2"]} + ) + + def check_results(batch_df, _): + self.assertEqual( + set(batch_df.sort("key").collect()), + { + Row(key="hello", countAsString="1"), + Row(key="foo", countAsString="666"), + Row(key="hello_2", countAsString="2"), + Row(key="this", countAsString="1"), + Row(key="this_2", countAsString="2"), + }, + ) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_basic_fewer_data(self): + # Test data rows returned are more or fewer than state. + def func(key, pdf_iter, state): + state.update((1,)) + assert isinstance(state, GroupState) + yield pd.DataFrame({"key": [], "countAsString": []}) + + def check_results(batch_df, _): + self.assertTrue(len(set(batch_df.sort("key").collect())) == 0) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_basic_with_null(self): + def func(key, pdf_iter, state): + assert isinstance(state, GroupState) + + total_len = 0 + for pdf in pdf_iter: + total_len += len(pdf) + + state.update((total_len,)) + assert state.get[0] == 1 + yield pd.DataFrame({"key": [None], "countAsString": [str(total_len)]}) + + def check_results(batch_df, _): + self.assertEqual( + set(batch_df.sort("key").collect()), + {Row(key=None, countAsString="1")}, + ) + + self._test_apply_in_pandas_with_state_basic(func, check_results) + + def test_apply_in_pandas_with_state_python_worker_random_failure(self): + input_path = tempfile.mkdtemp() + output_path = tempfile.mkdtemp() + checkpoint_loc = tempfile.mkdtemp() + + shutil.rmtree(output_path) + shutil.rmtree(checkpoint_loc) + + def prepare_test_resource(): + data_range = list(string.ascii_lowercase) + for i in range(5): + picked_data = [ + data_range[random.randrange(0, len(data_range) - 1)] for x in range(100) + ] + + with open(input_path + "/part-%i.txt" % i, "w") as fw: + for data in picked_data: + fw.write(data + "\n") + + def run_query(): + df = ( + self.spark.readStream.format("text") + .option("maxFilesPerTrigger", "1") + .load(input_path) + ) + + for q in self.spark.streams.active: + q.stop() + self.assertTrue(df.isStreaming) + + output_type = StructType( + [StructField("value", StringType()), StructField("count", LongType())] + ) + state_type = StructType([StructField("cnt", LongType())]) + + def func(key, pdf_iter, state): + assert isinstance(state, GroupState) + + # user function call will happen at most 26 times + # should be huge enough to not trigger kill in every batches + # but should be also reasonable to trigger kill multiple times across batches + if random.randrange(30) == 1: + sys.exit(1) + + count = state.getOption + if count is None: + count = 0 + else: + count = count[0] + + for pdf in pdf_iter: + count += len(pdf) + + state.update((count,)) + yield pd.DataFrame({"value": [key[0]], "count": [count]}) + + query = ( + df.groupBy(df["value"]) + .applyInPandasWithState( + func, output_type, state_type, "Append", GroupStateTimeout.NoTimeout + ) + .writeStream.queryName("this_query") + .format("json") + .outputMode("append") + .option("path", output_path) + .option("checkpointLocation", checkpoint_loc) + .start() + ) + + return query + + prepare_test_resource() + + expected = ( + self.spark.read.format("text") + .load(input_path) + .groupBy("value") + .count() + .sort("value") + .collect() + ) + + q = run_query() + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + + def assert_test(): + nonlocal q + if not q.isActive: + print("query has been terminated, rerunning query...") + + # rerunning query as the query may have been killed by killed python worker + q = run_query() + + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + + curr_status = q.status + if not curr_status["isDataAvailable"] and not curr_status["isTriggerActive"]: + # The query is active but not running due to no further data available + # Check the output now. + result = ( + self.spark.read.schema("value string, count int") + .format("json") + .load(output_path) + .groupBy("value") + .max("count") + .selectExpr("value", "`max(count)` AS count") + .sort("value") + .collect() + ) + + return result == expected + else: + # still processing the data, defer checking the output. + return False + + try: + eventually(assert_test, timeout=120) + finally: + q.stop() + + +class GroupedApplyInPandasWithStateTests( + GroupedApplyInPandasWithStateTestsMixin, ReusedSQLTestCase +): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_map.py new file mode 100644 index 0000000..fb2f921 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_map.py @@ -0,0 +1,436 @@ +# +# 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 os +import shutil +import tempfile +import time +import unittest +from typing import cast + +from pyspark.sql import Row +from pyspark.sql.functions import col, encode, lit +from pyspark.errors import PythonException +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + +if have_pandas: + import pandas as pd + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class MapInPandasTestsMixin: + @staticmethod + def identity_dataframes_iter(*columns: str): + def func(iterator): + for pdf in iterator: + assert isinstance(pdf, pd.DataFrame) + assert pdf.columns.tolist() == list(columns) + yield pdf + + return func + + @staticmethod + def identity_dataframes_wo_column_names_iter(*columns: str): + def func(iterator): + for pdf in iterator: + assert isinstance(pdf, pd.DataFrame) + assert pdf.columns.tolist() == list(columns) + yield pdf.rename(columns=list(pdf.columns).index) + + return func + + @staticmethod + def dataframes_and_empty_dataframe_iter(*columns: str): + def func(iterator): + for pdf in iterator: + yield pdf + # after yielding all elements, also yield an empty dataframe with given columns + yield pd.DataFrame([], columns=list(columns)) + + return func + + def test_map_in_pandas(self): + # test returning iterator of DataFrames + df = self.spark.range(10, numPartitions=3) + actual = df.mapInPandas(self.identity_dataframes_iter("id"), "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + # test returning list of DataFrames + df = self.spark.range(10, numPartitions=3) + actual = df.mapInPandas(lambda it: [pdf for pdf in it], "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_multiple_columns(self): + data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] + df = self.spark.createDataFrame(data, "a int, b string") + + def func(iterator): + for pdf in iterator: + assert isinstance(pdf, pd.DataFrame) + assert [d.name for d in list(pdf.dtypes)] == ["int32", "object"] + yield pdf + + actual = df.mapInPandas(func, df.schema).collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_large_variable_types(self): + with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes": True}): + + def func(iterator): + for pdf in iterator: + assert isinstance(pdf, pd.DataFrame) + yield pdf + + df = ( + self.spark.range(10, numPartitions=3) + .select(col("id").cast("string").alias("str")) + .withColumn("bin", encode(col("str"), "utf8")) + ) + actual = df.mapInPandas(func, "str string, bin binary").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_no_column_names(self): + data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] + df = self.spark.createDataFrame(data, "a int, b string") + + def func(iterator): + for pdf in iterator: + yield pdf.rename(columns=list(pdf.columns).index) + + actual = df.mapInPandas(func, df.schema).collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_different_output_length(self): + def func(iterator): + for _ in iterator: + yield pd.DataFrame({"a": list(range(100))}) + + df = self.spark.range(10) + actual = df.repartition(1).mapInPandas(func, "a long").collect() + self.assertEqual(set((r.a for r in actual)), set(range(100))) + + def test_other_than_dataframe_iter(self): + with QuietTest(self.sc): + self.check_other_than_dataframe_iter() + + def check_other_than_dataframe_iter(self): + def no_iter(_): + return 1 + + def bad_iter_elem(_): + return iter([1]) + + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be iterator of pandas.DataFrame, " + "but is int.", + ): + (self.spark.range(10, numPartitions=3).mapInPandas(no_iter, "a int").count()) + + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be iterator of pandas.DataFrame, " + "but is iterator of int.", + ): + (self.spark.range(10, numPartitions=3).mapInPandas(bad_iter_elem, "a int").count()) + + def test_dataframes_with_other_column_names(self): + with QuietTest(self.sc): + self.check_dataframes_with_other_column_names() + + def check_dataframes_with_other_column_names(self): + def dataframes_with_other_column_names(iterator): + for pdf in iterator: + yield pdf.rename(columns={"id": "iid"}) + + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF\\] " + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: id. Unexpected: iid.\n", + ): + ( + self.spark.range(10, numPartitions=3) + .withColumn("value", lit(0)) + .mapInPandas(dataframes_with_other_column_names, "id int, value int") + .collect() + ) + + def test_dataframes_with_duplicate_column_names(self): + with QuietTest(self.sc): + self.check_dataframes_with_duplicate_column_names() + + def check_dataframes_with_duplicate_column_names(self): + def dataframes_with_other_column_names(iterator): + for pdf in iterator: + yield pdf.rename(columns={"id2": "id"}) + + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF\\] " + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: id2.\n", + ): + ( + self.spark.range(10, numPartitions=3) + .withColumn("id2", lit(0)) + .withColumn("value", lit(1)) + .mapInPandas(dataframes_with_other_column_names, "id int, id2 long, value int") + .collect() + ) + + def test_dataframes_with_less_columns(self): + with QuietTest(self.sc): + self.check_dataframes_with_less_columns() + + def check_dataframes_with_less_columns(self): + df = self.spark.range(10, numPartitions=3).withColumn("value", lit(0)) + + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF\\] " + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: id2.\n", + ): + f = self.identity_dataframes_iter("id", "value") + (df.mapInPandas(f, "id int, id2 long, value int").collect()) + + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_LENGTH_MISMATCH_FOR_PANDAS_UDF\\] " + "Number of columns of the returned pandas.DataFrame doesn't match " + "specified schema. Expected: 3 Actual: 2\n", + ): + f = self.identity_dataframes_wo_column_names_iter("id", "value") + (df.mapInPandas(f, "id int, id2 long, value int").collect()) + + def test_dataframes_with_more_columns(self): + df = self.spark.range(10, numPartitions=3).select( + "id", col("id").alias("value"), col("id").alias("extra") + ) + expected = df.select("id", "value").collect() + + f = self.identity_dataframes_iter("id", "value", "extra") + actual = df.repartition(1).mapInPandas(f, "id long, value long").collect() + self.assertEqual(actual, expected) + + f = self.identity_dataframes_wo_column_names_iter("id", "value", "extra") + actual = df.repartition(1).mapInPandas(f, "id long, value long").collect() + self.assertEqual(actual, expected) + + def test_dataframes_with_incompatible_types(self): + with QuietTest(self.sc): + self.check_dataframes_with_incompatible_types() + + def check_dataframes_with_incompatible_types(self): + def func(iterator): + for pdf in iterator: + yield pdf.assign(id=pdf["id"].apply(str)) + + for safely in [True, False]: + with self.subTest(convertToArrowArraySafely=safely), self.sql_conf( + {"spark.sql.execution.pandas.convertToArrowArraySafely": safely} + ): + # sometimes we see ValueErrors + with self.subTest(convert="string to double"): + expected = ( + r"ValueError: Exception thrown when converting pandas.Series " + r"\(object\) with name 'id' to Arrow Array \(double\)." + ) + if safely: + expected = expected + ( + " It can be caused by overflows or other " + "unsafe conversions warned by Arrow. Arrow safe type check " + "can be disabled by using SQL config " + "`spark.sql.execution.pandas.convertToArrowArraySafely`." + ) + with self.assertRaisesRegex(PythonException, expected + "\n"): + ( + self.spark.range(10, numPartitions=3) + .mapInPandas(func, "id double") + .collect() + ) + + # sometimes we see TypeErrors + with self.subTest(convert="double to string"): + with self.assertRaisesRegex( + PythonException, + r"TypeError: Exception thrown when converting pandas.Series " + r"\(float64\) with name 'id' to Arrow Array \(string\).\n", + ): + ( + self.spark.range(10, numPartitions=3) + .select(col("id").cast("double")) + .mapInPandas(self.identity_dataframes_iter("id"), "id string") + .collect() + ) + + def test_empty_iterator(self): + def empty_iter(_): + return iter([]) + + mapped = self.spark.range(10, numPartitions=3).mapInPandas(empty_iter, "a int, b string") + self.assertEqual(mapped.count(), 0) + + def test_empty_dataframes(self): + def empty_dataframes(_): + return iter([pd.DataFrame({"a": []})]) + + mapped = self.spark.range(10, numPartitions=3).mapInPandas(empty_dataframes, "a int") + self.assertEqual(mapped.count(), 0) + + def test_empty_dataframes_without_columns(self): + mapped = self.spark.range(10, numPartitions=3).mapInPandas( + self.dataframes_and_empty_dataframe_iter(), "id int" + ) + self.assertEqual(mapped.count(), 10) + + def test_empty_dataframes_with_less_columns(self): + with QuietTest(self.sc): + self.check_empty_dataframes_with_less_columns() + + def check_empty_dataframes_with_less_columns(self): + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF\\] " + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: value.\n", + ): + f = self.dataframes_and_empty_dataframe_iter("id") + ( + self.spark.range(10, numPartitions=3) + .withColumn("value", lit(0)) + .mapInPandas(f, "id int, value int") + .collect() + ) + + def test_empty_dataframes_with_more_columns(self): + mapped = self.spark.range(10, numPartitions=3).mapInPandas( + self.dataframes_and_empty_dataframe_iter("id", "extra"), "id int" + ) + self.assertEqual(mapped.count(), 10) + + def test_empty_dataframes_with_other_columns(self): + with QuietTest(self.sc): + self.check_empty_dataframes_with_other_columns() + + def check_empty_dataframes_with_other_columns(self): + def empty_dataframes_with_other_columns(iterator): + for _ in iterator: + yield pd.DataFrame({"iid": [], "value": []}) + + with self.assertRaisesRegex( + PythonException, + "PySparkRuntimeError: \\[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF\\] " + "Column names of the returned pandas.DataFrame do not match " + "specified schema. Missing: id. Unexpected: iid.\n", + ): + ( + self.spark.range(10, numPartitions=3) + .withColumn("value", lit(0)) + .mapInPandas(empty_dataframes_with_other_columns, "id int, value int") + .collect() + ) + + def test_chain_map_partitions_in_pandas(self): + def func(iterator): + for pdf in iterator: + assert isinstance(pdf, pd.DataFrame) + assert pdf.columns == ["id"] + yield pdf + + df = self.spark.range(10, numPartitions=3) + actual = df.mapInPandas(func, "id long").mapInPandas(func, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_self_join(self): + # SPARK-34319: self-join with MapInPandas + df1 = self.spark.range(10, numPartitions=3) + df2 = df1.mapInPandas(lambda iter: iter, "id long") + actual = df2.join(df2).collect() + expected = df1.join(df1).collect() + self.assertEqual(sorted(actual), sorted(expected)) + + # SPARK-33277 + def test_map_in_pandas_with_column_vector(self): + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + self.spark.range(0, 200000, 1, 1).write.parquet(path) + + def func(iterator): + for pdf in iterator: + yield pd.DataFrame({"id": [0] * len(pdf)}) + + for offheap in ["true", "false"]: + with self.sql_conf({"spark.sql.columnVector.offheap.enabled": offheap}): + self.assertEquals( + self.spark.read.parquet(path).mapInPandas(func, "id long").head(), Row(0) + ) + finally: + shutil.rmtree(path) + + +class MapInPandasTests(ReusedSQLTestCase, MapInPandasTestsMixin): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + + # Synchronize default timezone between Python and Java + cls.tz_prev = os.environ.get("TZ", None) # save current tz if set + tz = "America/Los_Angeles" + os.environ["TZ"] = tz + time.tzset() + + cls.sc.environment["TZ"] = tz + cls.spark.conf.set("spark.sql.session.timeZone", tz) + + @classmethod + def tearDownClass(cls): + del os.environ["TZ"] + if cls.tz_prev is not None: + os.environ["TZ"] = cls.tz_prev + time.tzset() + ReusedSQLTestCase.tearDownClass() + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf.py new file mode 100644 index 0000000..34cd9c2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf.py @@ -0,0 +1,357 @@ +# +# 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 unittest +import datetime +from typing import cast + +from pyspark.sql.functions import udf, pandas_udf, PandasUDFType, assert_true, lit +from pyspark.sql.types import DoubleType, StructType, StructField, LongType, DayTimeIntervalType +from pyspark.errors import ParseException, PythonException, PySparkTypeError +from pyspark.rdd import PythonEvalType +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class PandasUDFTestsMixin: + def test_pandas_udf_basic(self): + udf = pandas_udf(lambda x: x, DoubleType()) + self.assertEqual(udf.returnType, DoubleType()) + self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + udf = pandas_udf(lambda x: x, DoubleType(), PandasUDFType.SCALAR) + self.assertEqual(udf.returnType, DoubleType()) + self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + udf = pandas_udf( + lambda x: x, StructType([StructField("v", DoubleType())]), PandasUDFType.GROUPED_MAP + ) + self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())])) + self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + def test_pandas_udf_basic_with_return_type_string(self): + udf = pandas_udf(lambda x: x, "double", PandasUDFType.SCALAR) + self.assertEqual(udf.returnType, DoubleType()) + self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + udf = pandas_udf(lambda x: x, "v double", PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())])) + self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + udf = pandas_udf(lambda x: x, "v double", functionType=PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())])) + self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + udf = pandas_udf(lambda x: x, returnType="v double", functionType=PandasUDFType.GROUPED_MAP) + self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())])) + self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + def test_pandas_udf_decorator(self): + @pandas_udf(DoubleType()) + def foo(x): + return x + + self.assertEqual(foo.returnType, DoubleType()) + self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + @pandas_udf(returnType=DoubleType()) + def foo(x): + return x + + self.assertEqual(foo.returnType, DoubleType()) + self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + schema = StructType([StructField("v", DoubleType())]) + + @pandas_udf(schema, PandasUDFType.GROUPED_MAP) + def foo(x): + return x + + self.assertEqual(foo.returnType, schema) + self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + @pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP) + def foo(x): + return x + + self.assertEqual(foo.returnType, schema) + self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + @pandas_udf(returnType=schema, functionType=PandasUDFType.GROUPED_MAP) + def foo(x): + return x + + self.assertEqual(foo.returnType, schema) + self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + def test_pandas_udf_decorator_with_return_type_string(self): + schema = StructType([StructField("v", DoubleType())]) + + @pandas_udf("v double", PandasUDFType.GROUPED_MAP) + def foo(x): + return x + + self.assertEqual(foo.returnType, schema) + self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) + + @pandas_udf(returnType="double", functionType=PandasUDFType.SCALAR) + def foo(x): + return x + + self.assertEqual(foo.returnType, DoubleType()) + self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + def test_udf_wrong_arg(self): + with QuietTest(self.sc): + self.check_udf_wrong_arg() + + with self.assertRaises(ParseException): + + @pandas_udf("blah") + def foo(x): + return x + + with self.assertRaises(PySparkTypeError) as pe: + + @pandas_udf(returnType="double", functionType=PandasUDFType.GROUPED_MAP) + def foo(df): + return df + + self.check_error( + exception=pe.exception, + error_class="INVALID_RETURN_TYPE_FOR_PANDAS_UDF", + message_parameters={ + "eval_type": "SQL_GROUPED_MAP_PANDAS_UDF " + "or SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE", + "return_type": "DoubleType()", + }, + ) + + with self.assertRaisesRegex(ValueError, "Invalid function"): + + @pandas_udf(returnType="k int, v double", functionType=PandasUDFType.GROUPED_MAP) + def foo(k, v, w): + return k + + def check_udf_wrong_arg(self): + with self.assertRaises(PySparkTypeError) as pe: + + @pandas_udf(functionType=PandasUDFType.SCALAR) + def foo(x): + return x + + self.check_error( + exception=pe.exception, + error_class="CANNOT_BE_NONE", + message_parameters={"arg_name": "returnType"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + + @pandas_udf("double", 100) + def foo(x): + return x + + self.check_error( + exception=pe.exception, + error_class="INVALID_PANDAS_UDF_TYPE", + message_parameters={"arg_name": "functionType", "arg_type": "100"}, + ) + + with self.assertRaisesRegex(ValueError, "0-arg pandas_udfs.*not.*supported"): + pandas_udf(lambda: 1, LongType(), PandasUDFType.SCALAR) + with self.assertRaisesRegex(ValueError, "0-arg pandas_udfs.*not.*supported"): + + @pandas_udf(LongType(), PandasUDFType.SCALAR) + def zero_with_type(): + return 1 + + with self.assertRaises(PySparkTypeError) as pe: + + @pandas_udf(returnType=PandasUDFType.GROUPED_MAP) + def foo(df): + return df + + self.check_error( + exception=pe.exception, + error_class="NOT_DATATYPE_OR_STR", + message_parameters={"arg_name": "returnType", "arg_type": "int"}, + ) + + def test_stopiteration_in_udf(self): + def foo(x): + raise StopIteration() + + exc_message = "StopIteration" + df = self.spark.range(0, 100) + + # plain udf (test for SPARK-23754) + self.assertRaisesRegex( + PythonException, exc_message, df.withColumn("v", udf(foo)("id")).collect + ) + + # pandas scalar udf + self.assertRaisesRegex( + PythonException, + exc_message, + df.withColumn("v", pandas_udf(foo, "double", PandasUDFType.SCALAR)("id")).collect, + ) + + def test_stopiteration_in_grouped_map(self): + def foo(x): + raise StopIteration() + + def foofoo(x, y): + raise StopIteration() + + exc_message = "StopIteration" + df = self.spark.range(0, 100) + + # pandas grouped map + self.assertRaisesRegex( + PythonException, + exc_message, + df.groupBy("id").apply(pandas_udf(foo, df.schema, PandasUDFType.GROUPED_MAP)).collect, + ) + + self.assertRaisesRegex( + PythonException, + exc_message, + df.groupBy("id") + .apply(pandas_udf(foofoo, df.schema, PandasUDFType.GROUPED_MAP)) + .collect, + ) + + def test_stopiteration_in_grouped_agg(self): + def foo(x): + raise StopIteration() + + exc_message = "StopIteration" + df = self.spark.range(0, 100) + + # pandas grouped agg + self.assertRaisesRegex( + PythonException, + exc_message, + df.groupBy("id") + .agg(pandas_udf(foo, "double", PandasUDFType.GROUPED_AGG)("id")) + .collect, + ) + + def test_pandas_udf_detect_unsafe_type_conversion(self): + import pandas as pd + import numpy as np + + values = [1.0] * 3 + pdf = pd.DataFrame({"A": values}) + df = self.spark.createDataFrame(pdf).repartition(1) + + @pandas_udf(returnType="int") + def udf(column): + return pd.Series(np.linspace(0, 1, len(column))) + + # Since 0.11.0, PyArrow supports the feature to raise an error for unsafe cast. + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": True}): + with self.assertRaisesRegex( + Exception, "Exception thrown when converting pandas.Series" + ): + df.select(["A"]).withColumn("udf", udf("A")).collect() + + # Disabling Arrow safe type check. + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}): + df.select(["A"]).withColumn("udf", udf("A")).collect() + + def test_pandas_udf_arrow_overflow(self): + import pandas as pd + + df = self.spark.range(0, 1) + + @pandas_udf(returnType="byte") + def udf(column): + return pd.Series([128] * len(column)) + + # When enabling safe type check, Arrow 0.11.0+ disallows overflow cast. + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": True}): + with self.assertRaisesRegex( + Exception, "Exception thrown when converting pandas.Series" + ): + df.withColumn("udf", udf("id")).collect() + + # Disabling safe type check, let Arrow do the cast anyway. + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}): + df.withColumn("udf", udf("id")).collect() + + def test_pandas_udf_timestamp_ntz(self): + # SPARK-36626: Test TimestampNTZ in pandas UDF + @pandas_udf(returnType="timestamp_ntz") + def noop(s): + assert s.iloc[0] == datetime.datetime(1970, 1, 1, 0, 0) + return s + + with self.sql_conf({"spark.sql.session.timeZone": "Asia/Hong_Kong"}): + df = self.spark.createDataFrame( + [(datetime.datetime(1970, 1, 1, 0, 0),)], schema="dt timestamp_ntz" + ).select(noop("dt").alias("dt")) + + df.selectExpr("assert_true('1970-01-01 00:00:00' == CAST(dt AS STRING))").collect() + self.assertEqual(df.schema[0].dataType.typeName(), "timestamp_ntz") + self.assertEqual(df.first()[0], datetime.datetime(1970, 1, 1, 0, 0)) + + def test_pandas_udf_day_time_interval_type(self): + # SPARK-37277: Test DayTimeIntervalType in pandas UDF + import pandas as pd + + @pandas_udf(DayTimeIntervalType(DayTimeIntervalType.DAY, DayTimeIntervalType.SECOND)) + def noop(s: pd.Series) -> pd.Series: + assert s.iloc[0] == datetime.timedelta(microseconds=123) + return s + + df = self.spark.createDataFrame( + [(datetime.timedelta(microseconds=123),)], schema="td interval day to second" + ).select(noop("td").alias("td")) + + df.select( + assert_true(lit("INTERVAL '0 00:00:00.000123' DAY TO SECOND") == df.td.cast("string")) + ).collect() + self.assertEqual(df.schema[0].dataType.simpleString(), "interval day to second") + self.assertEqual(df.first()[0], datetime.timedelta(microseconds=123)) + + +class PandasUDFTests(PandasUDFTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py new file mode 100644 index 0000000..f434489 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py @@ -0,0 +1,592 @@ +# +# 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 unittest +from typing import cast + +from pyspark.rdd import PythonEvalType +from pyspark.sql import Row +from pyspark.sql.functions import ( + array, + explode, + col, + lit, + mean, + sum, + udf, + pandas_udf, + PandasUDFType, +) +from pyspark.sql.types import ArrayType, YearMonthIntervalType +from pyspark.errors import AnalysisException, PySparkNotImplementedError +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + + +if have_pandas: + import pandas as pd + from pandas.testing import assert_frame_equal + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class GroupedAggPandasUDFTestsMixin: + @property + def data(self): + return ( + self.spark.range(10) + .toDF("id") + .withColumn("vs", array([lit(i * 1.0) + col("id") for i in range(20, 30)])) + .withColumn("v", explode(col("vs"))) + .drop("vs") + .withColumn("w", lit(1.0)) + ) + + @property + def python_plus_one(self): + @udf("double") + def plus_one(v): + assert isinstance(v, (int, float)) + return float(v + 1) + + return plus_one + + @property + def pandas_scalar_plus_two(self): + @pandas_udf("double", PandasUDFType.SCALAR) + def plus_two(v): + assert isinstance(v, pd.Series) + return v + 2 + + return plus_two + + @property + def pandas_agg_mean_udf(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def avg(v): + return v.mean() + + return avg + + @property + def pandas_agg_sum_udf(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def sum(v): + return v.sum() + + return sum + + @property + def pandas_agg_weighted_mean_udf(self): + import numpy as np + + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def weighted_mean(v, w): + return np.average(v, weights=w) + + return weighted_mean + + def test_manual(self): + df = self.data + sum_udf = self.pandas_agg_sum_udf + mean_udf = self.pandas_agg_mean_udf + mean_arr_udf = pandas_udf( + self.pandas_agg_mean_udf.func, + ArrayType(self.pandas_agg_mean_udf.returnType), + self.pandas_agg_mean_udf.evalType, + ) + + result1 = ( + df.groupby("id") + .agg(sum_udf(df.v), mean_udf(df.v), mean_arr_udf(array(df.v))) + .sort("id") + ) + expected1 = self.spark.createDataFrame( + [ + [0, 245.0, 24.5, [24.5]], + [1, 255.0, 25.5, [25.5]], + [2, 265.0, 26.5, [26.5]], + [3, 275.0, 27.5, [27.5]], + [4, 285.0, 28.5, [28.5]], + [5, 295.0, 29.5, [29.5]], + [6, 305.0, 30.5, [30.5]], + [7, 315.0, 31.5, [31.5]], + [8, 325.0, 32.5, [32.5]], + [9, 335.0, 33.5, [33.5]], + ], + ["id", "sum(v)", "avg(v)", "avg(array(v))"], + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_basic(self): + df = self.data + weighted_mean_udf = self.pandas_agg_weighted_mean_udf + + # Groupby one column and aggregate one UDF with literal + result1 = df.groupby("id").agg(weighted_mean_udf(df.v, lit(1.0))).sort("id") + expected1 = df.groupby("id").agg(mean(df.v).alias("weighted_mean(v, 1.0)")).sort("id") + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + # Groupby one expression and aggregate one UDF with literal + result2 = df.groupby((col("id") + 1)).agg(weighted_mean_udf(df.v, lit(1.0))).sort(df.id + 1) + expected2 = ( + df.groupby((col("id") + 1)) + .agg(mean(df.v).alias("weighted_mean(v, 1.0)")) + .sort(df.id + 1) + ) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + + # Groupby one column and aggregate one UDF without literal + result3 = df.groupby("id").agg(weighted_mean_udf(df.v, df.w)).sort("id") + expected3 = df.groupby("id").agg(mean(df.v).alias("weighted_mean(v, w)")).sort("id") + assert_frame_equal(expected3.toPandas(), result3.toPandas()) + + # Groupby one expression and aggregate one UDF without literal + result4 = ( + df.groupby((col("id") + 1).alias("id")).agg(weighted_mean_udf(df.v, df.w)).sort("id") + ) + expected4 = ( + df.groupby((col("id") + 1).alias("id")) + .agg(mean(df.v).alias("weighted_mean(v, w)")) + .sort("id") + ) + assert_frame_equal(expected4.toPandas(), result4.toPandas()) + + def test_unsupported_types(self): + with QuietTest(self.sc): + self.check_unsupported_types() + + def check_unsupported_types(self): + with self.assertRaises(PySparkNotImplementedError) as pe: + pandas_udf( + lambda x: x, + ArrayType(ArrayType(YearMonthIntervalType())), + PandasUDFType.GROUPED_AGG, + ) + + self.check_error( + exception=pe.exception, + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": "Invalid return type with grouped aggregate Pandas UDFs: " + "ArrayType(ArrayType(YearMonthIntervalType(0, 1), True), True)" + }, + ) + + with self.assertRaises(PySparkNotImplementedError) as pe: + + @pandas_udf("mean double, std double", PandasUDFType.GROUPED_AGG) + def mean_and_std_udf(v): + return v.mean(), v.std() + + self.check_error( + exception=pe.exception, + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": "Invalid return type with grouped aggregate Pandas UDFs: " + "StructType([StructField('mean', DoubleType(), True), " + "StructField('std', DoubleType(), True)])" + }, + ) + + with self.assertRaises(PySparkNotImplementedError) as pe: + + @pandas_udf(ArrayType(YearMonthIntervalType()), PandasUDFType.GROUPED_AGG) + def mean_and_std_udf(v): # noqa: F811 + return {v.mean(): v.std()} + + self.check_error( + exception=pe.exception, + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": "Invalid return type with grouped aggregate Pandas UDFs: " + "ArrayType(YearMonthIntervalType(0, 1), True)" + }, + ) + + def test_alias(self): + df = self.data + mean_udf = self.pandas_agg_mean_udf + + result1 = df.groupby("id").agg(mean_udf(df.v).alias("mean_alias")) + expected1 = df.groupby("id").agg(mean(df.v).alias("mean_alias")) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_mixed_sql(self): + """ + Test mixing group aggregate pandas UDF with sql expression. + """ + df = self.data + sum_udf = self.pandas_agg_sum_udf + + # Mix group aggregate pandas UDF with sql expression + result1 = df.groupby("id").agg(sum_udf(df.v) + 1).sort("id") + expected1 = df.groupby("id").agg(sum(df.v) + 1).sort("id") + + # Mix group aggregate pandas UDF with sql expression (order swapped) + result2 = df.groupby("id").agg(sum_udf(df.v + 1)).sort("id") + + expected2 = df.groupby("id").agg(sum(df.v + 1)).sort("id") + + # Wrap group aggregate pandas UDF with two sql expressions + result3 = df.groupby("id").agg(sum_udf(df.v + 1) + 2).sort("id") + expected3 = df.groupby("id").agg(sum(df.v + 1) + 2).sort("id") + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + assert_frame_equal(expected3.toPandas(), result3.toPandas()) + + def test_mixed_udfs(self): + """ + Test mixing group aggregate pandas UDF with python UDF and scalar pandas UDF. + """ + df = self.data + plus_one = self.python_plus_one + plus_two = self.pandas_scalar_plus_two + sum_udf = self.pandas_agg_sum_udf + + # Mix group aggregate pandas UDF and python UDF + result1 = df.groupby("id").agg(plus_one(sum_udf(df.v))).sort("id") + expected1 = df.groupby("id").agg(plus_one(sum(df.v))).sort("id") + + # Mix group aggregate pandas UDF and python UDF (order swapped) + result2 = df.groupby("id").agg(sum_udf(plus_one(df.v))).sort("id") + expected2 = df.groupby("id").agg(sum(plus_one(df.v))).sort("id") + + # Mix group aggregate pandas UDF and scalar pandas UDF + result3 = df.groupby("id").agg(sum_udf(plus_two(df.v))).sort("id") + expected3 = df.groupby("id").agg(sum(plus_two(df.v))).sort("id") + + # Mix group aggregate pandas UDF and scalar pandas UDF (order swapped) + result4 = df.groupby("id").agg(plus_two(sum_udf(df.v))).sort("id") + expected4 = df.groupby("id").agg(plus_two(sum(df.v))).sort("id") + + # Wrap group aggregate pandas UDF with two python UDFs and use python UDF in groupby + result5 = ( + df.groupby(plus_one(df.id)).agg(plus_one(sum_udf(plus_one(df.v)))).sort("plus_one(id)") + ) + expected5 = ( + df.groupby(plus_one(df.id)).agg(plus_one(sum(plus_one(df.v)))).sort("plus_one(id)") + ) + + # Wrap group aggregate pandas UDF with two scala pandas UDF and user scala pandas UDF in + # groupby + result6 = ( + df.groupby(plus_two(df.id)).agg(plus_two(sum_udf(plus_two(df.v)))).sort("plus_two(id)") + ) + expected6 = ( + df.groupby(plus_two(df.id)).agg(plus_two(sum(plus_two(df.v)))).sort("plus_two(id)") + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + assert_frame_equal(expected3.toPandas(), result3.toPandas()) + assert_frame_equal(expected4.toPandas(), result4.toPandas()) + assert_frame_equal(expected5.toPandas(), result5.toPandas()) + assert_frame_equal(expected6.toPandas(), result6.toPandas()) + + def test_multiple_udfs(self): + """ + Test multiple group aggregate pandas UDFs in one agg function. + """ + df = self.data + mean_udf = self.pandas_agg_mean_udf + sum_udf = self.pandas_agg_sum_udf + weighted_mean_udf = self.pandas_agg_weighted_mean_udf + + result1 = ( + df.groupBy("id") + .agg(mean_udf(df.v), sum_udf(df.v), weighted_mean_udf(df.v, df.w)) + .sort("id") + .toPandas() + ) + expected1 = ( + df.groupBy("id") + .agg(mean(df.v), sum(df.v), mean(df.v).alias("weighted_mean(v, w)")) + .sort("id") + .toPandas() + ) + + assert_frame_equal(expected1, result1) + + def test_complex_groupby(self): + df = self.data + sum_udf = self.pandas_agg_sum_udf + plus_one = self.python_plus_one + plus_two = self.pandas_scalar_plus_two + + # groupby one expression + result1 = df.groupby(df.v % 2).agg(sum_udf(df.v)) + expected1 = df.groupby(df.v % 2).agg(sum(df.v)) + + # empty groupby + result2 = df.groupby().agg(sum_udf(df.v)) + expected2 = df.groupby().agg(sum(df.v)) + + # groupby one column and one sql expression + result3 = df.groupby(df.id, df.v % 2).agg(sum_udf(df.v)).orderBy(df.id, df.v % 2) + expected3 = df.groupby(df.id, df.v % 2).agg(sum(df.v)).orderBy(df.id, df.v % 2) + + # groupby one python UDF + result4 = df.groupby(plus_one(df.id)).agg(sum_udf(df.v)).sort("plus_one(id)") + expected4 = df.groupby(plus_one(df.id)).agg(sum(df.v)).sort("plus_one(id)") + + # groupby one scalar pandas UDF + result5 = df.groupby(plus_two(df.id)).agg(sum_udf(df.v)).sort("sum(v)") + expected5 = df.groupby(plus_two(df.id)).agg(sum(df.v)).sort("sum(v)") + + # groupby one expression and one python UDF + result6 = ( + df.groupby(df.v % 2, plus_one(df.id)) + .agg(sum_udf(df.v)) + .sort(["(v % 2)", "plus_one(id)"]) + ) + expected6 = ( + df.groupby(df.v % 2, plus_one(df.id)).agg(sum(df.v)).sort(["(v % 2)", "plus_one(id)"]) + ) + + # groupby one expression and one scalar pandas UDF + result7 = ( + df.groupby(df.v % 2, plus_two(df.id)) + .agg(sum_udf(df.v)) + .sort(["sum(v)", "plus_two(id)"]) + ) + expected7 = ( + df.groupby(df.v % 2, plus_two(df.id)).agg(sum(df.v)).sort(["sum(v)", "plus_two(id)"]) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + assert_frame_equal(expected3.toPandas(), result3.toPandas()) + assert_frame_equal(expected4.toPandas(), result4.toPandas()) + assert_frame_equal(expected5.toPandas(), result5.toPandas()) + assert_frame_equal(expected6.toPandas(), result6.toPandas()) + assert_frame_equal(expected7.toPandas(), result7.toPandas()) + + def test_complex_expressions(self): + df = self.data + plus_one = self.python_plus_one + plus_two = self.pandas_scalar_plus_two + sum_udf = self.pandas_agg_sum_udf + + # Test complex expressions with sql expression, python UDF and + # group aggregate pandas UDF + result1 = ( + df.withColumn("v1", plus_one(df.v)) + .withColumn("v2", df.v + 2) + .groupby(df.id, df.v % 2) + .agg( + sum_udf(col("v")), + sum_udf(col("v1") + 3), + sum_udf(col("v2")) + 5, + plus_one(sum_udf(col("v1"))), + sum_udf(plus_one(col("v2"))), + ) + .sort(["id", "(v % 2)"]) + .toPandas() + .sort_values(by=["id", "(v % 2)"]) + ) + + expected1 = ( + df.withColumn("v1", df.v + 1) + .withColumn("v2", df.v + 2) + .groupby(df.id, df.v % 2) + .agg( + sum(col("v")), + sum(col("v1") + 3), + sum(col("v2")) + 5, + plus_one(sum(col("v1"))), + sum(plus_one(col("v2"))), + ) + .sort(["id", "(v % 2)"]) + .toPandas() + .sort_values(by=["id", "(v % 2)"]) + ) + + # Test complex expressions with sql expression, scala pandas UDF and + # group aggregate pandas UDF + result2 = ( + df.withColumn("v1", plus_one(df.v)) + .withColumn("v2", df.v + 2) + .groupby(df.id, df.v % 2) + .agg( + sum_udf(col("v")), + sum_udf(col("v1") + 3), + sum_udf(col("v2")) + 5, + plus_two(sum_udf(col("v1"))), + sum_udf(plus_two(col("v2"))), + ) + .sort(["id", "(v % 2)"]) + .toPandas() + .sort_values(by=["id", "(v % 2)"]) + ) + + expected2 = ( + df.withColumn("v1", df.v + 1) + .withColumn("v2", df.v + 2) + .groupby(df.id, df.v % 2) + .agg( + sum(col("v")), + sum(col("v1") + 3), + sum(col("v2")) + 5, + plus_two(sum(col("v1"))), + sum(plus_two(col("v2"))), + ) + .sort(["id", "(v % 2)"]) + .toPandas() + .sort_values(by=["id", "(v % 2)"]) + ) + + # Test sequential groupby aggregate + result3 = ( + df.groupby("id") + .agg(sum_udf(df.v).alias("v")) + .groupby("id") + .agg(sum_udf(col("v"))) + .sort("id") + .toPandas() + ) + + expected3 = ( + df.groupby("id") + .agg(sum(df.v).alias("v")) + .groupby("id") + .agg(sum(col("v"))) + .sort("id") + .toPandas() + ) + + assert_frame_equal(expected1, result1) + assert_frame_equal(expected2, result2) + assert_frame_equal(expected3, result3) + + def test_retain_group_columns(self): + with self.sql_conf({"spark.sql.retainGroupColumns": False}): + df = self.data + sum_udf = self.pandas_agg_sum_udf + + result1 = df.groupby(df.id).agg(sum_udf(df.v)) + expected1 = df.groupby(df.id).agg(sum(df.v)) + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_array_type(self): + df = self.data + + array_udf = pandas_udf(lambda x: [1.0, 2.0], "array", PandasUDFType.GROUPED_AGG) + result1 = df.groupby("id").agg(array_udf(df["v"]).alias("v2")) + self.assertEqual(result1.first()["v2"], [1.0, 2.0]) + + def test_invalid_args(self): + with QuietTest(self.sc): + self.check_invalid_args() + + def check_invalid_args(self): + df = self.data + plus_one = self.python_plus_one + mean_udf = self.pandas_agg_mean_udf + with self.assertRaisesRegex(AnalysisException, "[MISSING_AGGREGATION]"): + df.groupby(df.id).agg(plus_one(df.v)).collect() + with self.assertRaisesRegex( + AnalysisException, "aggregate function.*argument.*aggregate function" + ): + df.groupby(df.id).agg(mean_udf(mean_udf(df.v))).collect() + with self.assertRaisesRegex( + AnalysisException, + "The group aggregate pandas UDF `avg` cannot be invoked together with as other, " + "non-pandas aggregate functions.", + ): + df.groupby(df.id).agg(mean_udf(df.v), mean(df.v)).collect() + + def test_register_vectorized_udf_basic(self): + sum_pandas_udf = pandas_udf( + lambda v: v.sum(), "integer", PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF + ) + + self.assertEqual(sum_pandas_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) + group_agg_pandas_udf = self.spark.udf.register("sum_pandas_udf", sum_pandas_udf) + self.assertEqual(group_agg_pandas_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) + q = "SELECT sum_pandas_udf(v1) FROM VALUES (3, 0), (2, 0), (1, 1) tbl(v1, v2) GROUP BY v2" + actual = sorted(map(lambda r: r[0], self.spark.sql(q).collect())) + expected = [1, 5] + self.assertEqual(actual, expected) + + def test_grouped_with_empty_partition(self): + data = [Row(id=1, x=2), Row(id=1, x=3), Row(id=2, x=4)] + expected = [Row(id=1, sum=5), Row(id=2, x=4)] + num_parts = len(data) + 1 + df = self.spark.createDataFrame(self.sc.parallelize(data, numSlices=num_parts)) + + f = pandas_udf(lambda x: x.sum(), "int", PandasUDFType.GROUPED_AGG) + + result = df.groupBy("id").agg(f(df["x"]).alias("sum")).collect() + self.assertEqual(result, expected) + + def test_grouped_without_group_by_clause(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def max_udf(v): + return v.max() + + df = self.spark.range(0, 100) + self.spark.udf.register("max_udf", max_udf) + + with self.tempView("table"): + df.createTempView("table") + + agg1 = df.agg(max_udf(df["id"])) + agg2 = self.spark.sql("select max_udf(id) from table") + assert_frame_equal(agg1.toPandas(), agg2.toPandas()) + + def test_no_predicate_pushdown_through(self): + # SPARK-30921: We should not pushdown predicates of PythonUDFs through Aggregate. + import numpy as np + + @pandas_udf("float", PandasUDFType.GROUPED_AGG) + def mean(x): + return np.mean(x) + + df = self.spark.createDataFrame([Row(id=1, foo=42), Row(id=2, foo=1), Row(id=2, foo=2)]) + + agg = df.groupBy("id").agg(mean("foo").alias("mean")) + filtered = agg.filter(agg["mean"] > 40.0) + + assert filtered.collect()[0]["mean"] == 42.0 + + +class GroupedAggPandasUDFTests(GroupedAggPandasUDFTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf_grouped_agg import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py new file mode 100644 index 0000000..7a80547 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py @@ -0,0 +1,1503 @@ +# +# 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 os +import random +import shutil +import tempfile +import time +import unittest +from datetime import date, datetime +from decimal import Decimal +from distutils.version import LooseVersion +from typing import cast + +from pyspark import TaskContext +from pyspark.rdd import PythonEvalType +from pyspark.sql import Column +from pyspark.sql.functions import array, col, expr, lit, sum, struct, udf, pandas_udf, PandasUDFType +from pyspark.sql.pandas.utils import pyarrow_version_less_than_minimum +from pyspark.sql.types import ( + IntegerType, + ByteType, + StructType, + ShortType, + BooleanType, + LongType, + FloatType, + DoubleType, + DecimalType, + StringType, + ArrayType, + StructField, + Row, + TimestampType, + MapType, + DateType, + BinaryType, + YearMonthIntervalType, +) +from pyspark.errors import AnalysisException +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + test_compiled, + test_not_compiled_message, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest, assertDataFrameEqual + +if have_pandas: + import pandas as pd + +if have_pyarrow: + import pyarrow as pa # noqa: F401 + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class ScalarPandasUDFTestsMixin: + @property + def nondeterministic_vectorized_udf(self): + import numpy as np + + @pandas_udf("double") + def random_udf(v): + return pd.Series(np.random.random(len(v))) + + random_udf = random_udf.asNondeterministic() + return random_udf + + @property + def nondeterministic_vectorized_iter_udf(self): + import numpy as np + + @pandas_udf("double", PandasUDFType.SCALAR_ITER) + def random_udf(it): + for v in it: + yield pd.Series(np.random.random(len(v))) + + random_udf = random_udf.asNondeterministic() + return random_udf + + @property + def df_with_nested_structs(self): + schema = StructType( + [ + StructField("id", IntegerType(), False), + StructField( + "info", + StructType( + [ + StructField("name", StringType(), False), + StructField("age", IntegerType(), False), + StructField( + "details", + StructType( + [ + StructField("field1", StringType(), False), + StructField("field2", IntegerType(), False), + ] + ), + False, + ), + ] + ), + False, + ), + ] + ) + data = [(1, ("John", 30, ("Value1", 10)))] + df = self.spark.createDataFrame(data, schema) + struct_df = df.select(struct(df.columns).alias("struct")) + # struct_df.dtype: + # [( + # 'struct', + # 'struct>>' + # )] + return struct_df + + @property + def df_with_nested_maps(self): + schema = StructType( + [ + StructField("id", StringType(), True), + StructField( + "attributes", MapType(StringType(), MapType(StringType(), StringType())), True + ), + ] + ) + data = [("1", {"personal": {"name": "John", "city": "New York"}})] + return self.spark.createDataFrame(data, schema) + + @property + def df_with_nested_arrays(self): + schema = StructType( + [ + StructField("id", IntegerType(), nullable=False), + StructField("nested_array", ArrayType(ArrayType(IntegerType())), nullable=False), + ] + ) + data = [(1, [[1, 2, 3], [4, 5]])] + return self.spark.createDataFrame(data, schema) + + def test_pandas_udf_tokenize(self): + tokenize = pandas_udf( + lambda s: s.apply(lambda str: str.split(" ")), ArrayType(StringType()) + ) + self.assertEqual(tokenize.returnType, ArrayType(StringType())) + df = self.spark.createDataFrame([("hi boo",), ("bye boo",)], ["vals"]) + result = df.select(tokenize("vals").alias("hi")) + self.assertEqual([Row(hi=["hi", "boo"]), Row(hi=["bye", "boo"])], result.collect()) + + def test_pandas_udf_nested_arrays(self): + tokenize = pandas_udf( + lambda s: s.apply(lambda str: [str.split(" ")]), ArrayType(ArrayType(StringType())) + ) + self.assertEqual(tokenize.returnType, ArrayType(ArrayType(StringType()))) + df = self.spark.createDataFrame([("hi boo",), ("bye boo",)], ["vals"]) + result = df.select(tokenize("vals").alias("hi")) + self.assertEqual([Row(hi=[["hi", "boo"]]), Row(hi=[["bye", "boo"]])], result.collect()) + + def test_input_nested_structs(self): + df = self.df_with_nested_structs + + mirror = pandas_udf(lambda s: s, df.dtypes[0][1]) + + self.assertEquals( + df.select(mirror(df.struct).alias("res")).first(), + Row( + res=Row( + id=1, info=Row(name="John", age=30, details=Row(field1="Value1", field2=10)) + ) + ), + ) + + def test_input_nested_maps(self): + df = self.df_with_nested_maps + + str_repr = pandas_udf(lambda s: s.astype(str), StringType()) + self.assertEquals( + df.select(str_repr(df.attributes).alias("res")).first(), + Row(res="{'personal': {'name': 'John', 'city': 'New York'}}"), + ) + + extract_name = pandas_udf(lambda s: s.apply(lambda x: x["personal"]["name"]), StringType()) + self.assertEquals( + df.select(extract_name(df.attributes).alias("res")).first(), + Row(res="John"), + ) + + def test_input_nested_arrays(self): + df = self.df_with_nested_arrays + + str_repr = pandas_udf(lambda s: s.astype(str), StringType()) + self.assertEquals( + df.select(str_repr(df.nested_array).alias("res")).first(), + Row(res="[array([1, 2, 3], dtype=int32) array([4, 5], dtype=int32)]"), + ) + + @unittest.skipIf( + pyarrow_version_less_than_minimum("2.0.0"), + "Pyarrow version must be 2.0.0 or higher", + ) + def test_pandas_array_struct(self): + # SPARK-38098: Support Array of Struct for Pandas UDFs and toPandas + import numpy as np + + @pandas_udf("Array>") + def return_cols(cols): + assert type(cols) == pd.Series + assert type(cols[0]) == np.ndarray + assert type(cols[0][0]) == dict + return cols + + df = self.spark.createDataFrame( + [[[("a", 2, 3.0), ("a", 2, 3.0)]], [[("b", 5, 6.0), ("b", 5, 6.0)]]], + "array_struct_col Array>", + ) + result = df.select(return_cols("array_struct_col")) + self.assertEqual( + [ + Row(output=[Row(col1="a", col2=2, col3=3.0), Row(col1="a", col2=2, col3=3.0)]), + Row(output=[Row(col1="b", col2=5, col3=6.0), Row(col1="b", col2=5, col3=6.0)]), + ], + result.collect(), + ) + + def test_vectorized_udf_basic(self): + df = self.spark.range(10).select( + col("id").cast("string").alias("str"), + col("id").cast("int").alias("int"), + col("id").alias("long"), + col("id").cast("float").alias("float"), + col("id").cast("double").alias("double"), + col("id").cast("decimal").alias("decimal"), + col("id").cast("boolean").alias("bool"), + array(col("id")).alias("array_long"), + ) + + def f(x): + return x + + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + str_f = pandas_udf(f, StringType(), udf_type) + int_f = pandas_udf(f, IntegerType(), udf_type) + long_f = pandas_udf(f, LongType(), udf_type) + float_f = pandas_udf(f, FloatType(), udf_type) + double_f = pandas_udf(f, DoubleType(), udf_type) + decimal_f = pandas_udf(f, DecimalType(), udf_type) + bool_f = pandas_udf(f, BooleanType(), udf_type) + array_long_f = pandas_udf(f, ArrayType(LongType()), udf_type) + res = df.select( + str_f(col("str")), + int_f(col("int")), + long_f(col("long")), + float_f(col("float")), + double_f(col("double")), + decimal_f("decimal"), + bool_f(col("bool")), + array_long_f("array_long"), + ) + self.assertEqual(df.collect(), res.collect()) + + def test_register_nondeterministic_vectorized_udf_basic(self): + random_pandas_udf = pandas_udf( + lambda x: random.randint(6, 6) + x, IntegerType() + ).asNondeterministic() + self.assertEqual(random_pandas_udf.deterministic, False) + self.assertEqual(random_pandas_udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + nondeterministic_pandas_udf = self.spark.catalog.registerFunction( + "randomPandasUDF", random_pandas_udf + ) + self.assertEqual(nondeterministic_pandas_udf.deterministic, False) + self.assertEqual(nondeterministic_pandas_udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + [row] = self.spark.sql("SELECT randomPandasUDF(1)").collect() + self.assertEqual(row[0], 7) + + def random_iter_udf(it): + for i in it: + yield random.randint(6, 6) + i + + random_pandas_iter_udf = pandas_udf( + random_iter_udf, IntegerType(), PandasUDFType.SCALAR_ITER + ).asNondeterministic() + self.assertEqual(random_pandas_iter_udf.deterministic, False) + self.assertEqual(random_pandas_iter_udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF) + nondeterministic_pandas_iter_udf = self.spark.catalog.registerFunction( + "randomPandasIterUDF", random_pandas_iter_udf + ) + self.assertEqual(nondeterministic_pandas_iter_udf.deterministic, False) + self.assertEqual( + nondeterministic_pandas_iter_udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + ) + [row] = self.spark.sql("SELECT randomPandasIterUDF(1)").collect() + self.assertEqual(row[0], 7) + + def test_vectorized_udf_null_boolean(self): + data = [(True,), (True,), (None,), (False,)] + schema = StructType().add("bool", BooleanType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + bool_f = pandas_udf(lambda x: x, BooleanType(), udf_type) + res = df.select(bool_f(col("bool"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_byte(self): + data = [(None,), (2,), (3,), (4,)] + schema = StructType().add("byte", ByteType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + byte_f = pandas_udf(lambda x: x, ByteType(), udf_type) + res = df.select(byte_f(col("byte"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_short(self): + data = [(None,), (2,), (3,), (4,)] + schema = StructType().add("short", ShortType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + short_f = pandas_udf(lambda x: x, ShortType(), udf_type) + res = df.select(short_f(col("short"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_int(self): + data = [(None,), (2,), (3,), (4,)] + schema = StructType().add("int", IntegerType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + int_f = pandas_udf(lambda x: x, IntegerType(), udf_type) + res = df.select(int_f(col("int"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_long(self): + data = [(None,), (2,), (3,), (4,)] + schema = StructType().add("long", LongType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + long_f = pandas_udf(lambda x: x, LongType(), udf_type) + res = df.select(long_f(col("long"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_float(self): + data = [(3.0,), (5.0,), (-1.0,), (None,)] + schema = StructType().add("float", FloatType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + float_f = pandas_udf(lambda x: x, FloatType(), udf_type) + res = df.select(float_f(col("float"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_double(self): + data = [(3.0,), (5.0,), (-1.0,), (None,)] + schema = StructType().add("double", DoubleType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + double_f = pandas_udf(lambda x: x, DoubleType(), udf_type) + res = df.select(double_f(col("double"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_decimal(self): + data = [(Decimal(3.0),), (Decimal(5.0),), (Decimal(-1.0),), (None,)] + schema = StructType().add("decimal", DecimalType(38, 18)) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + decimal_f = pandas_udf(lambda x: x, DecimalType(38, 18), udf_type) + res = df.select(decimal_f(col("decimal"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_string(self): + data = [("foo",), (None,), ("bar",), ("bar",)] + schema = StructType().add("str", StringType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + str_f = pandas_udf(lambda x: x, StringType(), udf_type) + res = df.select(str_f(col("str"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_string_in_udf(self): + df = self.spark.range(10) + + def scalar_f(x): + return pd.Series(map(str, x)) + + def iter_f(it): + for i in it: + yield scalar_f(i) + + for f, udf_type in [(scalar_f, PandasUDFType.SCALAR), (iter_f, PandasUDFType.SCALAR_ITER)]: + str_f = pandas_udf(f, StringType(), udf_type) + actual = df.select(str_f(col("id"))) + expected = df.select(col("id").cast("string")) + self.assertEqual(expected.collect(), actual.collect()) + + def test_vectorized_udf_datatype_string(self): + df = self.spark.range(10).select( + col("id").cast("string").alias("str"), + col("id").cast("int").alias("int"), + col("id").alias("long"), + col("id").cast("float").alias("float"), + col("id").cast("double").alias("double"), + col("id").cast("decimal").alias("decimal"), + col("id").cast("boolean").alias("bool"), + ) + + def f(x): + return x + + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + str_f = pandas_udf(f, "string", udf_type) + int_f = pandas_udf(f, "integer", udf_type) + long_f = pandas_udf(f, "long", udf_type) + float_f = pandas_udf(f, "float", udf_type) + double_f = pandas_udf(f, "double", udf_type) + decimal_f = pandas_udf(f, "decimal(38, 18)", udf_type) + bool_f = pandas_udf(f, "boolean", udf_type) + res = df.select( + str_f(col("str")), + int_f(col("int")), + long_f(col("long")), + float_f(col("float")), + double_f(col("double")), + decimal_f("decimal"), + bool_f(col("bool")), + ) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_null_binary(self): + data = [(bytearray(b"a"),), (None,), (bytearray(b"bb"),), (bytearray(b"ccc"),)] + schema = StructType().add("binary", BinaryType()) + df = self.spark.createDataFrame(data, schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + str_f = pandas_udf(lambda x: x, BinaryType(), udf_type) + res = df.select(str_f(col("binary"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_array_type(self): + data = [([1, 2],), ([3, 4],)] + array_schema = StructType([StructField("array", ArrayType(IntegerType()))]) + df = self.spark.createDataFrame(data, schema=array_schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + array_f = pandas_udf(lambda x: x, ArrayType(IntegerType()), udf_type) + result = df.select(array_f(col("array"))) + self.assertEqual(df.collect(), result.collect()) + + def test_vectorized_udf_null_array(self): + data = [([1, 2],), (None,), (None,), ([3, 4],), (None,)] + array_schema = StructType([StructField("array", ArrayType(IntegerType()))]) + df = self.spark.createDataFrame(data, schema=array_schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + array_f = pandas_udf(lambda x: x, ArrayType(IntegerType()), udf_type) + result = df.select(array_f(col("array"))) + self.assertEqual(df.collect(), result.collect()) + + def test_vectorized_udf_struct_type(self): + df = self.spark.range(10) + return_type = StructType([StructField("id", LongType()), StructField("str", StringType())]) + + def scalar_func(id): + return pd.DataFrame({"id": id, "str": id.apply(str)}) + + def iter_func(it): + for id in it: + yield scalar_func(id) + + for func, udf_type in [ + (scalar_func, PandasUDFType.SCALAR), + (iter_func, PandasUDFType.SCALAR_ITER), + ]: + f = pandas_udf(func, returnType=return_type, functionType=udf_type) + + expected = df.select( + struct(col("id"), col("id").cast("string").alias("str")).alias("struct") + ).collect() + + actual = df.select(f(col("id")).alias("struct")).collect() + self.assertEqual(expected, actual) + + g = pandas_udf(func, "id: long, str: string", functionType=udf_type) + actual = df.select(g(col("id")).alias("struct")).collect() + self.assertEqual(expected, actual) + + struct_f = pandas_udf(lambda x: x, return_type, functionType=udf_type) + actual = df.select(struct_f(struct(col("id"), col("id").cast("string").alias("str")))) + self.assertEqual(expected, actual.collect()) + + def test_vectorized_udf_struct_complex(self): + df = self.spark.range(10) + return_type = StructType( + [StructField("ts", TimestampType()), StructField("arr", ArrayType(LongType()))] + ) + + def _scalar_f(id): + return pd.DataFrame( + {"ts": id.apply(lambda i: pd.Timestamp(i)), "arr": id.apply(lambda i: [i, i + 1])} + ) + + scalar_f = pandas_udf(_scalar_f, returnType=return_type) + + @pandas_udf(returnType=return_type, functionType=PandasUDFType.SCALAR_ITER) + def iter_f(it): + for id in it: + yield _scalar_f(id) + + for f, udf_type in [(scalar_f, PandasUDFType.SCALAR), (iter_f, PandasUDFType.SCALAR_ITER)]: + actual = df.withColumn("f", f(col("id"))).collect() + for i, row in enumerate(actual): + id, f = row + self.assertEqual(i, id) + self.assertEqual(pd.Timestamp(i).to_pydatetime(), f[0]) + self.assertListEqual([i, i + 1], f[1]) + + def test_vectorized_udf_struct_empty(self): + df = self.spark.range(3) + return_type = StructType() + + def _scalar_f(id): + return pd.DataFrame(index=id) + + scalar_f = pandas_udf(_scalar_f, returnType=return_type) + + @pandas_udf(returnType=return_type, functionType=PandasUDFType.SCALAR_ITER) + def iter_f(it): + for id in it: + yield _scalar_f(id) + + for f, udf_type in [(scalar_f, "SCALAR"), (iter_f, "SCALAR_ITER")]: + with self.subTest(udf_type=udf_type): + assertDataFrameEqual( + df.withColumn("f", f(col("id"))), + [Row(id=0, f=Row()), Row(id=1, f=Row()), Row(id=2, f=Row())], + ) + + def test_vectorized_udf_nested_struct(self): + with QuietTest(self.sc): + self.check_vectorized_udf_nested_struct() + + def check_vectorized_udf_nested_struct(self): + df = self.spark.range(2) + + nested_type = StructType( + [ + StructField("id", IntegerType()), + StructField( + "nested", + StructType([StructField("foo", StringType()), StructField("bar", FloatType())]), + ), + ] + ) + + def func_dict(pser: pd.Series) -> pd.DataFrame: + return pd.DataFrame( + {"id": pser, "nested": pser.apply(lambda x: {"foo": str(x), "bar": float(x)})} + ) + + def func_row(pser: pd.Series) -> pd.DataFrame: + return pd.DataFrame( + {"id": pser, "nested": pser.apply(lambda x: Row(foo=str(x), bar=float(x)))} + ) + + expected = [ + Row(udf=Row(id=0, nested=Row(foo="0", bar=0.0))), + Row(udf=Row(id=1, nested=Row(foo="1", bar=1.0))), + ] + + for f in [func_dict, func_row]: + for udf_type, func in [ + (PandasUDFType.SCALAR, f), + (PandasUDFType.SCALAR_ITER, lambda iter: (f(pser) for pser in iter)), + ]: + with self.subTest(udf_type=udf_type, udf=f.__name__): + result = df.select( + pandas_udf(func, returnType=nested_type, functionType=udf_type)( + col("id") + ).alias("udf") + ).collect() + self.assertEqual(result, expected) + + def test_vectorized_udf_map_type(self): + data = [({},), ({"a": 1},), ({"a": 1, "b": 2},), ({"a": 1, "b": 2, "c": 3},)] + schema = StructType([StructField("map", MapType(StringType(), LongType()))]) + df = self.spark.createDataFrame(data, schema=schema) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + if LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + with QuietTest(self.sc): + with self.assertRaisesRegex(Exception, "MapType.*not supported"): + pandas_udf(lambda x: x, MapType(StringType(), LongType()), udf_type) + else: + map_f = pandas_udf(lambda x: x, MapType(StringType(), LongType()), udf_type) + result = df.select(map_f(col("map"))) + self.assertEqual(df.collect(), result.collect()) + + def test_vectorized_udf_complex(self): + df = self.spark.range(10).select( + col("id").cast("int").alias("a"), + col("id").cast("int").alias("b"), + col("id").cast("double").alias("c"), + ) + scalar_add = pandas_udf(lambda x, y: x + y, IntegerType()) + scalar_power2 = pandas_udf(lambda x: 2**x, IntegerType()) + scalar_mul = pandas_udf(lambda x, y: x * y, DoubleType()) + + @pandas_udf(IntegerType(), PandasUDFType.SCALAR_ITER) + def iter_add(it): + for x, y in it: + yield x + y + + @pandas_udf(IntegerType(), PandasUDFType.SCALAR_ITER) + def iter_power2(it): + for x in it: + yield 2**x + + @pandas_udf(DoubleType(), PandasUDFType.SCALAR_ITER) + def iter_mul(it): + for x, y in it: + yield x * y + + for add, power2, mul in [ + (scalar_add, scalar_power2, scalar_mul), + (iter_add, iter_power2, iter_mul), + ]: + res = df.select(add(col("a"), col("b")), power2(col("a")), mul(col("b"), col("c"))) + expected = df.select(expr("a + b"), expr("power(2, a)"), expr("b * c")) + self.assertEqual(expected.collect(), res.collect()) + + def test_vectorized_udf_exception(self): + with QuietTest(self.sc): + self.check_vectorized_udf_exception() + + def check_vectorized_udf_exception(self): + df = self.spark.range(10) + scalar_raise_exception = pandas_udf(lambda x: x * (1 / 0), LongType()) + + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_raise_exception(it): + for x in it: + yield x * (1 / 0) + + for raise_exception in [scalar_raise_exception, iter_raise_exception]: + with self.assertRaisesRegex(Exception, "division( or modulo)? by zero"): + df.select(raise_exception(col("id"))).collect() + + def test_vectorized_udf_invalid_length(self): + with QuietTest(self.sc): + self.check_vectorized_udf_invalid_length() + + def check_vectorized_udf_invalid_length(self): + df = self.spark.range(10) + raise_exception = pandas_udf(lambda _: pd.Series(1), LongType()) + with self.assertRaisesRegex( + Exception, "Result vector from pandas_udf was not the required length" + ): + df.select(raise_exception(col("id"))).collect() + + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_udf_wong_output_size(it): + for _ in it: + yield pd.Series(1) + + with self.assertRaisesRegex( + Exception, "The length of output in Scalar iterator.*" "the length of output was 1" + ): + df.select(iter_udf_wong_output_size(col("id"))).collect() + + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_udf_not_reading_all_input(it): + for batch in it: + batch_len = len(batch) + yield pd.Series([1] * batch_len) + break + + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 3}): + df1 = self.spark.range(10).repartition(1) + with self.assertRaisesRegex(Exception, "pandas iterator UDF should exhaust"): + df1.select(iter_udf_not_reading_all_input(col("id"))).collect() + + def test_vectorized_udf_chained(self): + df = self.spark.range(10) + scalar_f = pandas_udf(lambda x: x + 1, LongType()) + scalar_g = pandas_udf(lambda x: x - 1, LongType()) + + iter_f = pandas_udf( + lambda it: map(lambda x: x + 1, it), LongType(), PandasUDFType.SCALAR_ITER + ) + iter_g = pandas_udf( + lambda it: map(lambda x: x - 1, it), LongType(), PandasUDFType.SCALAR_ITER + ) + + for f, g in [(scalar_f, scalar_g), (iter_f, iter_g)]: + res = df.select(g(f(col("id")))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_chained_struct_type(self): + df = self.spark.range(10) + return_type = StructType([StructField("id", LongType()), StructField("str", StringType())]) + + @pandas_udf(return_type) + def scalar_f(id): + return pd.DataFrame({"id": id, "str": id.apply(str)}) + + scalar_g = pandas_udf(lambda x: x, return_type) + + @pandas_udf(return_type, PandasUDFType.SCALAR_ITER) + def iter_f(it): + for id in it: + yield pd.DataFrame({"id": id, "str": id.apply(str)}) + + iter_g = pandas_udf(lambda x: x, return_type, PandasUDFType.SCALAR_ITER) + + expected = df.select( + struct(col("id"), col("id").cast("string").alias("str")).alias("struct") + ).collect() + + for f, g in [(scalar_f, scalar_g), (iter_f, iter_g)]: + actual = df.select(g(f(col("id"))).alias("struct")).collect() + self.assertEqual(expected, actual) + + def test_vectorized_udf_wrong_return_type(self): + with QuietTest(self.sc): + self.check_vectorized_udf_wrong_return_type() + + def check_vectorized_udf_wrong_return_type(self): + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + with self.assertRaisesRegex( + NotImplementedError, + "Invalid return type.*scalar Pandas UDF.*ArrayType.*YearMonthIntervalType", + ): + pandas_udf(lambda x: x, ArrayType(YearMonthIntervalType()), udf_type) + + def test_vectorized_udf_return_scalar(self): + with QuietTest(self.sc): + self.check_vectorized_udf_return_scalar() + + def check_vectorized_udf_return_scalar(self): + df = self.spark.range(10) + scalar_f = pandas_udf(lambda x: 1.0, DoubleType()) + iter_f = pandas_udf( + lambda it: map(lambda x: 1.0, it), DoubleType(), PandasUDFType.SCALAR_ITER + ) + for f in [scalar_f, iter_f]: + with self.assertRaisesRegex(Exception, "Return.*type.*Series"): + df.select(f(col("id"))).collect() + + def test_vectorized_udf_decorator(self): + df = self.spark.range(10) + + @pandas_udf(returnType=LongType()) + def scalar_identity(x): + return x + + @pandas_udf(returnType=LongType(), functionType=PandasUDFType.SCALAR_ITER) + def iter_identity(x): + return x + + for identity in [scalar_identity, iter_identity]: + res = df.select(identity(col("id"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_empty_partition(self): + df = self.spark.createDataFrame(self.sc.parallelize([Row(id=1)], 2)) + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + f = pandas_udf(lambda x: x, LongType(), udf_type) + res = df.select(f(col("id"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_struct_with_empty_partition(self): + df = self.spark.createDataFrame(self.sc.parallelize([Row(id=1)], 2)).withColumn( + "name", lit("John Doe") + ) + + @pandas_udf("first string, last string") + def scalar_split_expand(n): + return n.str.split(expand=True) + + @pandas_udf("first string, last string", PandasUDFType.SCALAR_ITER) + def iter_split_expand(it): + for n in it: + yield n.str.split(expand=True) + + for split_expand in [scalar_split_expand, iter_split_expand]: + result = df.select(split_expand("name")).collect() + self.assertEqual(1, len(result)) + row = result[0] + self.assertEqual("John", row[0]["first"]) + self.assertEqual("Doe", row[0]["last"]) + + def test_vectorized_udf_varargs(self): + df = self.spark.range(start=1, end=2) + scalar_f = pandas_udf(lambda *v: v[0], LongType()) + + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_f(it): + for v in it: + yield v[0] + + for f in [scalar_f, iter_f]: + res = df.select(f(col("id"), col("id"))) + self.assertEqual(df.collect(), res.collect()) + + def test_vectorized_udf_dates(self): + schema = StructType().add("idx", LongType()).add("date", DateType()) + data = [ + ( + 0, + date(1969, 1, 1), + ), + ( + 1, + date(2012, 2, 2), + ), + ( + 2, + None, + ), + ( + 3, + date(2100, 4, 4), + ), + ( + 4, + date(2262, 4, 12), + ), + ] + df = self.spark.createDataFrame(data, schema=schema) + + def scalar_check_data(idx, date, date_copy): + msgs = [] + is_equal = date.isnull() + for i in range(len(idx)): + if (is_equal[i] and data[idx[i]][1] is None) or date[i] == data[idx[i]][1]: + msgs.append(None) + else: + msgs.append( + "date values are not equal (date='%s': data[%d][1]='%s')" + % (date[i], idx[i], data[idx[i]][1]) + ) + return pd.Series(msgs) + + def iter_check_data(it): + for idx, test_date, date_copy in it: + yield scalar_check_data(idx, test_date, date_copy) + + pandas_scalar_check_data = pandas_udf(scalar_check_data, StringType()) + pandas_iter_check_data = pandas_udf( + iter_check_data, StringType(), PandasUDFType.SCALAR_ITER + ) + + for check_data, udf_type in [ + (pandas_scalar_check_data, PandasUDFType.SCALAR), + (pandas_iter_check_data, PandasUDFType.SCALAR_ITER), + ]: + date_copy = pandas_udf(lambda t: t, returnType=DateType(), functionType=udf_type) + df = df.withColumn("date_copy", date_copy(col("date"))) + result = df.withColumn( + "check_data", check_data(col("idx"), col("date"), col("date_copy")) + ).collect() + + self.assertEqual(len(data), len(result)) + for i in range(len(result)): + self.assertEqual(data[i][1], result[i][1]) # "date" col + self.assertEqual(data[i][1], result[i][2]) # "date_copy" col + self.assertIsNone(result[i][3]) # "check_data" col + + def test_vectorized_udf_timestamps(self): + schema = StructType( + [StructField("idx", LongType(), True), StructField("timestamp", TimestampType(), True)] + ) + data = [ + (0, datetime(1969, 1, 1, 1, 1, 1)), + (1, datetime(2012, 2, 2, 2, 2, 2)), + (2, None), + (3, datetime(2100, 3, 3, 3, 3, 3)), + ] + + df = self.spark.createDataFrame(data, schema=schema) + + def scalar_check_data(idx, timestamp, timestamp_copy): + msgs = [] + is_equal = timestamp.isnull() # use this array to check values are equal + for i in range(len(idx)): + # Check that timestamps are as expected in the UDF + if (is_equal[i] and data[idx[i]][1] is None) or timestamp[ + i + ].to_pydatetime() == data[idx[i]][1]: + msgs.append(None) + else: + msgs.append( + "timestamp values are not equal (timestamp='%s': data[%d][1]='%s')" + % (timestamp[i], idx[i], data[idx[i]][1]) + ) + return pd.Series(msgs) + + def iter_check_data(it): + for idx, timestamp, timestamp_copy in it: + yield scalar_check_data(idx, timestamp, timestamp_copy) + + pandas_scalar_check_data = pandas_udf(scalar_check_data, StringType()) + pandas_iter_check_data = pandas_udf( + iter_check_data, StringType(), PandasUDFType.SCALAR_ITER + ) + + for check_data, udf_type in [ + (pandas_scalar_check_data, PandasUDFType.SCALAR), + (pandas_iter_check_data, PandasUDFType.SCALAR_ITER), + ]: + # Check that a timestamp passed through a pandas_udf will not be altered by timezone + # calc + f_timestamp_copy = pandas_udf( + lambda t: t, returnType=TimestampType(), functionType=udf_type + ) + df = df.withColumn("timestamp_copy", f_timestamp_copy(col("timestamp"))) + result = df.withColumn( + "check_data", check_data(col("idx"), col("timestamp"), col("timestamp_copy")) + ).collect() + # Check that collection values are correct + self.assertEqual(len(data), len(result)) + for i in range(len(result)): + self.assertEqual(data[i][1], result[i][1]) # "timestamp" col + self.assertEqual(data[i][1], result[i][2]) # "timestamp_copy" col + self.assertIsNone(result[i][3]) # "check_data" col + + def test_vectorized_udf_return_timestamp_tz(self): + df = self.spark.range(10) + + @pandas_udf(returnType=TimestampType()) + def scalar_gen_timestamps(id): + ts = [pd.Timestamp(i, unit="D", tz="America/Los_Angeles") for i in id] + return pd.Series(ts) + + @pandas_udf(returnType=TimestampType(), functionType=PandasUDFType.SCALAR_ITER) + def iter_gen_timestamps(it): + for id in it: + ts = [pd.Timestamp(i, unit="D", tz="America/Los_Angeles") for i in id] + yield pd.Series(ts) + + for gen_timestamps in [scalar_gen_timestamps, iter_gen_timestamps]: + result = df.withColumn("ts", gen_timestamps(col("id"))).collect() + spark_ts_t = TimestampType() + for r in result: + i, ts = r + ts_tz = pd.Timestamp(i, unit="D", tz="America/Los_Angeles").to_pydatetime() + expected = spark_ts_t.fromInternal(spark_ts_t.toInternal(ts_tz)) + self.assertEqual(expected, ts) + + def test_vectorized_udf_check_config(self): + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 3}): + df = self.spark.range(10, numPartitions=1) + + @pandas_udf(returnType=LongType()) + def scalar_check_records_per_batch(x): + return pd.Series(x.size).repeat(x.size) + + @pandas_udf(returnType=LongType(), functionType=PandasUDFType.SCALAR_ITER) + def iter_check_records_per_batch(it): + for x in it: + yield pd.Series(x.size).repeat(x.size) + + for check_records_per_batch in [ + scalar_check_records_per_batch, + iter_check_records_per_batch, + ]: + result = df.select(check_records_per_batch(col("id"))).collect() + for (r,) in result: + self.assertTrue(r <= 3) + + def test_vectorized_udf_timestamps_respect_session_timezone(self): + schema = StructType( + [StructField("idx", LongType(), True), StructField("timestamp", TimestampType(), True)] + ) + data = [ + (1, datetime(1969, 1, 1, 1, 1, 1)), + (2, datetime(2012, 2, 2, 2, 2, 2)), + (3, None), + (4, datetime(2100, 3, 3, 3, 3, 3)), + ] + df = self.spark.createDataFrame(data, schema=schema) + + scalar_internal_value = pandas_udf( + lambda ts: ts.apply(lambda ts: ts.value if ts is not pd.NaT else None), LongType() + ) + + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_internal_value(it): + for ts in it: + yield ts.apply(lambda ts: ts.value if ts is not pd.NaT else None) + + for internal_value, udf_type in [ + (scalar_internal_value, PandasUDFType.SCALAR), + (iter_internal_value, PandasUDFType.SCALAR_ITER), + ]: + f_timestamp_copy = pandas_udf(lambda ts: ts, TimestampType(), udf_type) + timezone = "America/Los_Angeles" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + df_la = df.withColumn("tscopy", f_timestamp_copy(col("timestamp"))).withColumn( + "internal_value", internal_value(col("timestamp")) + ) + result_la = df_la.select(col("idx"), col("internal_value")).collect() + # Correct result_la by adjusting 3 hours difference between Los Angeles and New York + diff = 3 * 60 * 60 * 1000 * 1000 * 1000 + result_la_corrected = df_la.select( + col("idx"), col("tscopy"), col("internal_value") + diff + ).collect() + + timezone = "America/New_York" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + df_ny = df.withColumn("tscopy", f_timestamp_copy(col("timestamp"))).withColumn( + "internal_value", internal_value(col("timestamp")) + ) + result_ny = df_ny.select(col("idx"), col("tscopy"), col("internal_value")).collect() + + self.assertNotEqual(result_ny, result_la) + self.assertEqual(result_ny, result_la_corrected) + + def test_nondeterministic_vectorized_udf(self): + # Test that nondeterministic UDFs are evaluated only once in chained UDF evaluations + @pandas_udf("double") + def scalar_plus_ten(v): + return v + 10 + + @pandas_udf("double", PandasUDFType.SCALAR_ITER) + def iter_plus_ten(it): + for v in it: + yield v + 10 + + for plus_ten in [scalar_plus_ten, iter_plus_ten]: + random_udf = self.nondeterministic_vectorized_udf + + df = self.spark.range(10).withColumn("rand", random_udf(col("id"))) + result1 = df.withColumn("plus_ten(rand)", plus_ten(df["rand"])).toPandas() + + self.assertEqual(random_udf.deterministic, False) + self.assertTrue(result1["plus_ten(rand)"].equals(result1["rand"] + 10)) + + def test_nondeterministic_vectorized_udf_in_aggregate(self): + with QuietTest(self.sc): + self.check_nondeterministic_analysis_exception() + + def check_nondeterministic_analysis_exception(self): + df = self.spark.range(10) + for random_udf in [ + self.nondeterministic_vectorized_udf, + self.nondeterministic_vectorized_iter_udf, + ]: + with self.assertRaisesRegex(AnalysisException, "Non-deterministic"): + df.groupby(df.id).agg(sum(random_udf(df.id))).collect() + with self.assertRaisesRegex(AnalysisException, "Non-deterministic"): + df.agg(sum(random_udf(df.id))).collect() + + def test_register_vectorized_udf_basic(self): + df = self.spark.range(10).select( + col("id").cast("int").alias("a"), col("id").cast("int").alias("b") + ) + scalar_original_add = pandas_udf(lambda x, y: x + y, IntegerType()) + self.assertEqual(scalar_original_add.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF) + + @pandas_udf(IntegerType(), PandasUDFType.SCALAR_ITER) + def iter_original_add(it): + for x, y in it: + yield x + y + + self.assertEqual(iter_original_add.evalType, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF) + + for original_add in [scalar_original_add, iter_original_add]: + self.assertEqual(original_add.deterministic, True) + new_add = self.spark.catalog.registerFunction("add1", original_add) + res1 = df.select(new_add(col("a"), col("b"))) + res2 = self.spark.sql( + "SELECT add1(t.a, t.b) FROM (SELECT id as a, id as b FROM range(10)) t" + ) + expected = df.select(expr("a + b")) + self.assertEqual(expected.collect(), res1.collect()) + self.assertEqual(expected.collect(), res2.collect()) + + def test_scalar_iter_udf_init(self): + import numpy as np + + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def rng(batch_iter): + context = TaskContext.get() + part = context.partitionId() + np.random.seed(part) + for batch in batch_iter: + yield pd.Series(np.random.randint(100, size=len(batch))) + + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 2}): + df = self.spark.range(10, numPartitions=2).select(rng(col("id").alias("v"))) + result1 = df.collect() + result2 = df.collect() + self.assertEqual( + result1, + result2, + "SCALAR ITER UDF can initialize state and produce deterministic RNG", + ) + + def test_scalar_iter_udf_close(self): + with QuietTest(self.sc): + self.check_scalar_iter_udf_close() + + def check_scalar_iter_udf_close(self): + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def test_close(batch_iter): + try: + for batch in batch_iter: + yield batch + finally: + raise RuntimeError("reached finally block") + + with self.assertRaisesRegex(Exception, "reached finally block"): + self.spark.range(1).select(test_close(col("id"))).collect() + + @unittest.skip("LimitPushDown should push limits through Python UDFs so this won't occur") + def test_scalar_iter_udf_close_early(self): + tmp_dir = tempfile.mkdtemp() + try: + tmp_file = tmp_dir + "/reach_finally_block" + + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def test_close(batch_iter): + generator_exit_caught = False + try: + for batch in batch_iter: + yield batch + time.sleep(1.0) # avoid the function finish too fast. + except GeneratorExit as ge: + generator_exit_caught = True + raise ge + finally: + assert generator_exit_caught, "Generator exit exception was not caught." + open(tmp_file, "a").close() + + with QuietTest(self.sc): + with self.sql_conf( + { + "spark.sql.execution.arrow.maxRecordsPerBatch": 1, + "spark.sql.execution.pandas.udf.buffer.size": 4, + } + ): + self.spark.range(10).repartition(1).select(test_close(col("id"))).limit( + 2 + ).collect() + # wait here because python udf worker will take some time to detect + # jvm side socket closed and then will trigger `GenerateExit` raised. + # wait timeout is 10s. + for i in range(100): + time.sleep(0.1) + if os.path.exists(tmp_file): + break + + assert os.path.exists(tmp_file), "finally block not reached." + + finally: + shutil.rmtree(tmp_dir) + + # Regression test for SPARK-23314 + def test_timestamp_dst(self): + # Daylight saving time for Los Angeles for 2015 is Sun, Nov 1 at 2:00 am + dt = [ + datetime(2015, 11, 1, 0, 30), + datetime(2015, 11, 1, 1, 30), + datetime(2015, 11, 1, 2, 30), + ] + df = self.spark.createDataFrame(dt, "timestamp").toDF("time") + + for udf_type in [PandasUDFType.SCALAR, PandasUDFType.SCALAR_ITER]: + foo_udf = pandas_udf(lambda x: x, "timestamp", udf_type) + result = df.withColumn("time", foo_udf(df.time)) + self.assertEqual(df.collect(), result.collect()) + + def test_udf_category_type(self): + @pandas_udf("string") + def to_category_func(x): + return x.astype("category") + + pdf = pd.DataFrame({"A": ["a", "b", "c", "a"]}) + df = self.spark.createDataFrame(pdf) + df = df.withColumn("B", to_category_func(df["A"])) + result_spark = df.toPandas() + + spark_type = df.dtypes[1][1] + # spark data frame and arrow execution mode enabled data frame type must match pandas + self.assertEqual(spark_type, "string") + + # Check result of column 'B' must be equal to column 'A' in type and values + pd.testing.assert_series_equal(result_spark["A"], result_spark["B"], check_names=False) + + def test_type_annotation(self): + # Regression test to check if type hints can be used. See SPARK-23569. + def noop(col: pd.Series) -> pd.Series: + return col + + df = self.spark.range(1).select(pandas_udf(f=noop, returnType="bigint")("id")) + self.assertEqual(df.first()[0], 0) + + def test_mixed_udf(self): + df = self.spark.range(0, 1).toDF("v") + + # Test mixture of multiple UDFs and Pandas UDFs. + + @udf("int") + def f1(x): + assert type(x) == int + return x + 1 + + @pandas_udf("int") + def f2_scalar(x): + assert type(x) == pd.Series + return x + 10 + + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def f2_iter(it): + for x in it: + assert type(x) == pd.Series + yield x + 10 + + @udf("int") + def f3(x): + assert type(x) == int + return x + 100 + + @pandas_udf("int") + def f4_scalar(x): + assert type(x) == pd.Series + return x + 1000 + + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def f4_iter(it): + for x in it: + assert type(x) == pd.Series + yield x + 1000 + + expected_chained_1 = df.withColumn("f2_f1", df["v"] + 11).collect() + expected_chained_2 = df.withColumn("f3_f2_f1", df["v"] + 111).collect() + expected_chained_3 = df.withColumn("f4_f3_f2_f1", df["v"] + 1111).collect() + expected_chained_4 = df.withColumn("f4_f2_f1", df["v"] + 1011).collect() + expected_chained_5 = df.withColumn("f4_f3_f1", df["v"] + 1101).collect() + + expected_multi = ( + df.withColumn("f1", df["v"] + 1) + .withColumn("f2", df["v"] + 10) + .withColumn("f3", df["v"] + 100) + .withColumn("f4", df["v"] + 1000) + .withColumn("f2_f1", df["v"] + 11) + .withColumn("f3_f1", df["v"] + 101) + .withColumn("f4_f1", df["v"] + 1001) + .withColumn("f3_f2", df["v"] + 110) + .withColumn("f4_f2", df["v"] + 1010) + .withColumn("f4_f3", df["v"] + 1100) + .withColumn("f3_f2_f1", df["v"] + 111) + .withColumn("f4_f2_f1", df["v"] + 1011) + .withColumn("f4_f3_f1", df["v"] + 1101) + .withColumn("f4_f3_f2", df["v"] + 1110) + .withColumn("f4_f3_f2_f1", df["v"] + 1111) + .collect() + ) + + for f2, f4 in [ + (f2_scalar, f4_scalar), + (f2_scalar, f4_iter), + (f2_iter, f4_scalar), + (f2_iter, f4_iter), + ]: + # Test single expression with chained UDFs + df_chained_1 = df.withColumn("f2_f1", f2(f1(df["v"]))) + df_chained_2 = df.withColumn("f3_f2_f1", f3(f2(f1(df["v"])))) + df_chained_3 = df.withColumn("f4_f3_f2_f1", f4(f3(f2(f1(df["v"]))))) + df_chained_4 = df.withColumn("f4_f2_f1", f4(f2(f1(df["v"])))) + df_chained_5 = df.withColumn("f4_f3_f1", f4(f3(f1(df["v"])))) + + self.assertEqual(expected_chained_1, df_chained_1.collect()) + self.assertEqual(expected_chained_2, df_chained_2.collect()) + self.assertEqual(expected_chained_3, df_chained_3.collect()) + self.assertEqual(expected_chained_4, df_chained_4.collect()) + self.assertEqual(expected_chained_5, df_chained_5.collect()) + + # Test multiple mixed UDF expressions in a single projection + df_multi_1 = ( + df.withColumn("f1", f1(col("v"))) + .withColumn("f2", f2(col("v"))) + .withColumn("f3", f3(col("v"))) + .withColumn("f4", f4(col("v"))) + .withColumn("f2_f1", f2(col("f1"))) + .withColumn("f3_f1", f3(col("f1"))) + .withColumn("f4_f1", f4(col("f1"))) + .withColumn("f3_f2", f3(col("f2"))) + .withColumn("f4_f2", f4(col("f2"))) + .withColumn("f4_f3", f4(col("f3"))) + .withColumn("f3_f2_f1", f3(col("f2_f1"))) + .withColumn("f4_f2_f1", f4(col("f2_f1"))) + .withColumn("f4_f3_f1", f4(col("f3_f1"))) + .withColumn("f4_f3_f2", f4(col("f3_f2"))) + .withColumn("f4_f3_f2_f1", f4(col("f3_f2_f1"))) + ) + + # Test mixed udfs in a single expression + df_multi_2 = ( + df.withColumn("f1", f1(col("v"))) + .withColumn("f2", f2(col("v"))) + .withColumn("f3", f3(col("v"))) + .withColumn("f4", f4(col("v"))) + .withColumn("f2_f1", f2(f1(col("v")))) + .withColumn("f3_f1", f3(f1(col("v")))) + .withColumn("f4_f1", f4(f1(col("v")))) + .withColumn("f3_f2", f3(f2(col("v")))) + .withColumn("f4_f2", f4(f2(col("v")))) + .withColumn("f4_f3", f4(f3(col("v")))) + .withColumn("f3_f2_f1", f3(f2(f1(col("v"))))) + .withColumn("f4_f2_f1", f4(f2(f1(col("v"))))) + .withColumn("f4_f3_f1", f4(f3(f1(col("v"))))) + .withColumn("f4_f3_f2", f4(f3(f2(col("v"))))) + .withColumn("f4_f3_f2_f1", f4(f3(f2(f1(col("v")))))) + ) + + self.assertEqual(expected_multi, df_multi_1.collect()) + self.assertEqual(expected_multi, df_multi_2.collect()) + + def test_mixed_udf_and_sql(self): + from pyspark.sql.connect.column import Column as ConnectColumn + + df = self.spark.range(0, 1).toDF("v") + + # Test mixture of UDFs, Pandas UDFs and SQL expression. + + @udf("int") + def f1(x): + assert type(x) == int + return x + 1 + + def f2(x): + assert type(x) in (Column, ConnectColumn) + return x + 10 + + @pandas_udf("int") + def f3s(x): + assert type(x) == pd.Series + return x + 100 + + @pandas_udf("int", PandasUDFType.SCALAR_ITER) + def f3i(it): + for x in it: + assert type(x) == pd.Series + yield x + 100 + + expected = ( + df.withColumn("f1", df["v"] + 1) + .withColumn("f2", df["v"] + 10) + .withColumn("f3", df["v"] + 100) + .withColumn("f1_f2", df["v"] + 11) + .withColumn("f1_f3", df["v"] + 101) + .withColumn("f2_f1", df["v"] + 11) + .withColumn("f2_f3", df["v"] + 110) + .withColumn("f3_f1", df["v"] + 101) + .withColumn("f3_f2", df["v"] + 110) + .withColumn("f1_f2_f3", df["v"] + 111) + .withColumn("f1_f3_f2", df["v"] + 111) + .withColumn("f2_f1_f3", df["v"] + 111) + .withColumn("f2_f3_f1", df["v"] + 111) + .withColumn("f3_f1_f2", df["v"] + 111) + .withColumn("f3_f2_f1", df["v"] + 111) + .collect() + ) + + for f3 in [f3s, f3i]: + df1 = ( + df.withColumn("f1", f1(df["v"])) + .withColumn("f2", f2(df["v"])) + .withColumn("f3", f3(df["v"])) + .withColumn("f1_f2", f1(f2(df["v"]))) + .withColumn("f1_f3", f1(f3(df["v"]))) + .withColumn("f2_f1", f2(f1(df["v"]))) + .withColumn("f2_f3", f2(f3(df["v"]))) + .withColumn("f3_f1", f3(f1(df["v"]))) + .withColumn("f3_f2", f3(f2(df["v"]))) + .withColumn("f1_f2_f3", f1(f2(f3(df["v"])))) + .withColumn("f1_f3_f2", f1(f3(f2(df["v"])))) + .withColumn("f2_f1_f3", f2(f1(f3(df["v"])))) + .withColumn("f2_f3_f1", f2(f3(f1(df["v"])))) + .withColumn("f3_f1_f2", f3(f1(f2(df["v"])))) + .withColumn("f3_f2_f1", f3(f2(f1(df["v"])))) + ) + + self.assertEqual(expected, df1.collect()) + + # SPARK-24721 + @unittest.skipIf(not test_compiled, test_not_compiled_message) # type: ignore + def test_datasource_with_udf(self): + # Same as SQLTests.test_datasource_with_udf, but with Pandas UDF + # This needs to a separate test because Arrow dependency is optional + import numpy as np + + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + self.spark.range(1).write.mode("overwrite").format("csv").save(path) + filesource_df = self.spark.read.option("inferSchema", True).csv(path).toDF("i") + datasource_df = ( + self.spark.read.format("org.apache.spark.sql.sources.SimpleScanSource") + .option("from", 0) + .option("to", 1) + .load() + .toDF("i") + ) + datasource_v2_df = ( + self.spark.read.format("org.apache.spark.sql.connector.SimpleDataSourceV2") + .load() + .toDF("i", "j") + ) + + c1 = pandas_udf(lambda x: x + 1, "int")(lit(1)) + c2 = pandas_udf(lambda x: x + 1, "int")(col("i")) + + f1 = pandas_udf(lambda x: pd.Series(np.repeat(False, len(x))), "boolean")(lit(1)) + f2 = pandas_udf(lambda x: pd.Series(np.repeat(False, len(x))), "boolean")(col("i")) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + result = df.withColumn("c", c1) + expected = df.withColumn("c", lit(2)) + self.assertEqual(expected.collect(), result.collect()) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + result = df.withColumn("c", c2) + expected = df.withColumn("c", col("i") + 1) + self.assertEqual(expected.collect(), result.collect()) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + for f in [f1, f2]: + result = df.filter(f) + self.assertEqual(0, result.count()) + finally: + shutil.rmtree(path) + + # SPARK-33277 + def test_pandas_udf_with_column_vector(self): + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + self.spark.range(0, 200000, 1, 1).write.parquet(path) + + @pandas_udf(LongType()) + def udf(x): + return pd.Series([0] * len(x)) + + for offheap in ["true", "false"]: + with self.sql_conf({"spark.sql.columnVector.offheap.enabled": offheap}): + self.assertEquals( + self.spark.read.parquet(path).select(udf("id")).head(), Row(0) + ) + finally: + shutil.rmtree(path) + + +class ScalarPandasUDFTests(ScalarPandasUDFTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + + # Synchronize default timezone between Python and Java + cls.tz_prev = os.environ.get("TZ", None) # save current tz if set + tz = "America/Los_Angeles" + os.environ["TZ"] = tz + time.tzset() + + cls.sc.environment["TZ"] = tz + cls.spark.conf.set("spark.sql.session.timeZone", tz) + + @classmethod + def tearDownClass(cls): + del os.environ["TZ"] + if cls.tz_prev is not None: + os.environ["TZ"] = cls.tz_prev + time.tzset() + ReusedSQLTestCase.tearDownClass() + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf_scalar import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py new file mode 100644 index 0000000..bfb874f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py @@ -0,0 +1,392 @@ +# +# 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 sys +import unittest +from inspect import signature +from typing import Union, Iterator, Tuple, cast, get_type_hints + +from pyspark.sql.functions import mean, lit +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.sql.pandas.typehints import infer_eval_type +from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType +from pyspark.sql import Row + +if have_pandas: + import pandas as pd + import numpy as np + from pandas.testing import assert_frame_equal + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class PandasUDFTypeHintsTests(ReusedSQLTestCase): + def test_type_annotation_scalar(self): + def func(col: pd.Series) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.DataFrame, col1: pd.Series) -> pd.DataFrame: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.DataFrame, *args: pd.Series) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.Series, *, col2: pd.DataFrame) -> pd.DataFrame: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def test_type_annotation_scalar_iter(self): + def func(iter: Iterator[pd.Series]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[pd.DataFrame, pd.Series]]) -> Iterator[pd.DataFrame]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[pd.DataFrame, ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[Union[pd.DataFrame, pd.Series], ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + @unittest.skipIf(sys.version_info < (3, 9), "Type hinting generics require Python 3.9.") + def test_type_annotation_tuple_generics(self): + def func(iter: Iterator[tuple[pd.DataFrame, pd.Series]]) -> Iterator[pd.DataFrame]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[tuple[pd.DataFrame, ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[tuple[Union[pd.DataFrame, pd.Series], ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def test_type_annotation_group_agg(self): + def func(col: pd.Series) -> str: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.DataFrame, col1: pd.Series) -> int: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.DataFrame, *args: pd.Series) -> Row: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> str: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.Series, *, col2: pd.DataFrame) -> float: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> float: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def test_type_annotation_negative(self): + def func(col: str) -> pd.Series: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*str", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.DataFrame, col1: int) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*int", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: Union[pd.DataFrame, str], col1: int) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*str", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series) -> Tuple[pd.DataFrame]: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*Tuple", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col, *args: pd.Series) -> pd.Series: + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame): + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series, *, col2) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def test_scalar_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def plus_one(v: Union[pd.Series, pd.DataFrame]) -> pd.Series: + return v + 1 + + plus_one = pandas_udf("long")(plus_one) + actual = df.select(plus_one(df.v).alias("plus_one")) + expected = df.selectExpr("(v + 1) as plus_one") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_scalar_iter_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def plus_one(itr: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in itr: + yield s + 1 + + plus_one = pandas_udf("long")(plus_one) + + actual = df.select(plus_one(df.v).alias("plus_one")) + expected = df.selectExpr("(v + 1) as plus_one") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_group_agg_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def weighted_mean(v: pd.Series, w: pd.Series) -> np.float64: + return np.average(v, weights=w) + + weighted_mean = pandas_udf("double")(weighted_mean) + + actual = df.groupby("id").agg(weighted_mean(df.v, lit(1.0))).sort("id") + expected = df.groupby("id").agg(mean(df.v).alias("weighted_mean(v, 1.0)")).sort("id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_group_apply_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(v: pd.DataFrame) -> pd.DataFrame: + return v + 1 + + actual = df.groupby("id").applyInPandas(pandas_plus_one, schema=df.schema).sort("id") + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_cogroup_apply_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame: + return left + 1 + + actual = ( + df.groupby("id") + .cogroup(self.spark.range(10).groupby("id")) + .applyInPandas(pandas_plus_one, schema=df.schema) + .sort("id") + ) + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_map_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(iter: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: + return map(lambda v: v + 1, iter) + + actual = df.mapInPandas(pandas_plus_one, schema=df.schema) + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_string_type_annotation(self): + def func(col: "pd.Series") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.DataFrame", col1: "pd.Series") -> "pd.DataFrame": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.DataFrame", *args: "pd.Series") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.Series", *args: "pd.Series", **kwargs: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.Series", *, col2: "pd.DataFrame") -> "pd.DataFrame": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: Union["pd.Series", "pd.DataFrame"], *, col2: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "Union[pd.Series, pd.DataFrame]", *, col2: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf_typehints import * # noqa: #401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py new file mode 100644 index 0000000..9b67515 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py @@ -0,0 +1,375 @@ +# +# 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 sys +import unittest +from inspect import signature +from typing import Union, Iterator, Tuple, cast, get_type_hints + +from pyspark.sql.functions import mean, lit +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.sql.pandas.typehints import infer_eval_type +from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType +from pyspark.sql import Row + +if have_pandas: + import pandas as pd + import numpy as np + from pandas.testing import assert_frame_equal + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class PandasUDFTypeHintsWithFutureAnnotationsTests(ReusedSQLTestCase): + def test_type_annotation_scalar(self): + def func(col: pd.Series) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.DataFrame, col1: pd.Series) -> pd.DataFrame: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.DataFrame, *args: pd.Series) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: pd.Series, *, col2: pd.DataFrame) -> pd.DataFrame: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> pd.Series: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def test_type_annotation_scalar_iter(self): + def func(iter: Iterator[pd.Series]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[pd.DataFrame, pd.Series]]) -> Iterator[pd.DataFrame]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[pd.DataFrame, ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def func(iter: Iterator[Tuple[Union[pd.DataFrame, pd.Series], ...]]) -> Iterator[pd.Series]: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR_ITER + ) + + def test_type_annotation_group_agg(self): + def func(col: pd.Series) -> str: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.DataFrame, col1: pd.Series) -> int: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.DataFrame, *args: pd.Series) -> Row: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> str: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: pd.Series, *, col2: pd.DataFrame) -> float: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> float: + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.GROUPED_AGG + ) + + def test_type_annotation_negative(self): + def func(col: str) -> pd.Series: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*str", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.DataFrame, col1: int) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*int", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: Union[pd.DataFrame, str], col1: int) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*str", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series) -> Tuple[pd.DataFrame]: + pass + + self.assertRaisesRegex( + NotImplementedError, + "Unsupported signature.*Tuple", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col, *args: pd.Series) -> pd.Series: + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame): + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def func(col: pd.Series, *, col2) -> pd.DataFrame: + pass + + self.assertRaisesRegex( + ValueError, + "should be specified.*Series", + infer_eval_type, + signature(func), + get_type_hints(func), + ) + + def test_scalar_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def plus_one(v: Union[pd.Series, pd.DataFrame]) -> pd.Series: + return v + 1 + + plus_one = pandas_udf("long")(plus_one) + actual = df.select(plus_one(df.v).alias("plus_one")) + expected = df.selectExpr("(v + 1) as plus_one") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_scalar_iter_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def plus_one(itr: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in itr: + yield s + 1 + + plus_one = pandas_udf("long")(plus_one) + + actual = df.select(plus_one(df.v).alias("plus_one")) + expected = df.selectExpr("(v + 1) as plus_one") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_group_agg_udf_type_hint(self): + df = self.spark.range(10).selectExpr("id", "id as v") + + def weighted_mean(v: pd.Series, w: pd.Series) -> np.float64: + return np.average(v, weights=w) + + weighted_mean = pandas_udf("double")(weighted_mean) + + actual = df.groupby("id").agg(weighted_mean(df.v, lit(1.0))).sort("id") + expected = df.groupby("id").agg(mean(df.v).alias("weighted_mean(v, 1.0)")).sort("id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_group_apply_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(v: pd.DataFrame) -> pd.DataFrame: + return v + 1 + + actual = df.groupby("id").applyInPandas(pandas_plus_one, schema=df.schema).sort("id") + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_cogroup_apply_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame: + return left + 1 + + actual = ( + df.groupby("id") + .cogroup(self.spark.range(10).groupby("id")) + .applyInPandas(pandas_plus_one, schema=df.schema) + .sort("id") + ) + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + def test_ignore_type_hint_in_map_in_pandas(self): + df = self.spark.range(10) + + def pandas_plus_one(iter: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: + return map(lambda v: v + 1, iter) + + actual = df.mapInPandas(pandas_plus_one, schema=df.schema) + expected = df.selectExpr("id + 1 as id") + assert_frame_equal(expected.toPandas(), actual.toPandas()) + + @unittest.skipIf( + sys.version_info < (3, 9), + "string annotations with future annotations do not work under Python<3.9", + ) + def test_string_type_annotation(self): + def func(col: "pd.Series") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.DataFrame", col1: "pd.Series") -> "pd.DataFrame": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.DataFrame", *args: "pd.Series") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.Series", *args: "pd.Series", **kwargs: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "pd.Series", *, col2: "pd.DataFrame") -> "pd.DataFrame": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: Union["pd.Series", "pd.DataFrame"], *, col2: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + def func(col: "Union[pd.Series, pd.DataFrame]", *, col2: "pd.DataFrame") -> "pd.Series": + pass + + self.assertEqual( + infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR + ) + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf_typehints_with_future_annotations import * # noqa: #401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py new file mode 100644 index 0000000..e74e378 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py @@ -0,0 +1,411 @@ +# +# 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 unittest +from typing import cast + +from pyspark.errors import AnalysisException +from pyspark.sql.functions import ( + array, + explode, + col, + lit, + mean, + min, + max, + rank, + udf, + pandas_udf, + PandasUDFType, +) +from pyspark.sql.window import Window +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + +if have_pandas: + from pandas.testing import assert_frame_equal + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class WindowPandasUDFTestsMixin: + @property + def data(self): + return ( + self.spark.range(10) + .toDF("id") + .withColumn("vs", array([lit(i * 1.0) + col("id") for i in range(20, 30)])) + .withColumn("v", explode(col("vs"))) + .drop("vs") + .withColumn("w", lit(1.0)) + ) + + @property + def python_plus_one(self): + @udf("double") + def plus_one(v): + assert isinstance(v, float) + return v + 1 + + return plus_one + + @property + def pandas_scalar_time_two(self): + return pandas_udf(lambda v: v * 2, "double") + + @property + def pandas_agg_count_udf(self): + @pandas_udf("long", PandasUDFType.GROUPED_AGG) + def count(v): + return len(v) + + return count + + @property + def pandas_agg_mean_udf(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def avg(v): + return v.mean() + + return avg + + @property + def pandas_agg_max_udf(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def max(v): + return v.max() + + return max + + @property + def pandas_agg_min_udf(self): + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def min(v): + return v.min() + + return min + + @property + def unbounded_window(self): + return ( + Window.partitionBy("id") + .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + .orderBy("v") + ) + + @property + def ordered_window(self): + return Window.partitionBy("id").orderBy("v") + + @property + def unpartitioned_window(self): + return Window.partitionBy() + + @property + def sliding_row_window(self): + return Window.partitionBy("id").orderBy("v").rowsBetween(-2, 1) + + @property + def sliding_range_window(self): + return Window.partitionBy("id").orderBy("v").rangeBetween(-2, 4) + + @property + def growing_row_window(self): + return Window.partitionBy("id").orderBy("v").rowsBetween(Window.unboundedPreceding, 3) + + @property + def growing_range_window(self): + return Window.partitionBy("id").orderBy("v").rangeBetween(Window.unboundedPreceding, 4) + + @property + def shrinking_row_window(self): + return Window.partitionBy("id").orderBy("v").rowsBetween(-2, Window.unboundedFollowing) + + @property + def shrinking_range_window(self): + return Window.partitionBy("id").orderBy("v").rangeBetween(-3, Window.unboundedFollowing) + + def test_simple(self): + df = self.data + w = self.unbounded_window + + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("mean_v", mean_udf(df["v"]).over(w)) + expected1 = df.withColumn("mean_v", mean(df["v"]).over(w)) + + result2 = df.select(mean_udf(df["v"]).over(w)) + expected2 = df.select(mean(df["v"]).over(w)) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + + def test_multiple_udfs(self): + df = self.data + w = self.unbounded_window + + result1 = ( + df.withColumn("mean_v", self.pandas_agg_mean_udf(df["v"]).over(w)) + .withColumn("max_v", self.pandas_agg_max_udf(df["v"]).over(w)) + .withColumn("min_w", self.pandas_agg_min_udf(df["w"]).over(w)) + ) + + expected1 = ( + df.withColumn("mean_v", mean(df["v"]).over(w)) + .withColumn("max_v", max(df["v"]).over(w)) + .withColumn("min_w", min(df["w"]).over(w)) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_replace_existing(self): + df = self.data + w = self.unbounded_window + + result1 = df.withColumn("v", self.pandas_agg_mean_udf(df["v"]).over(w)) + expected1 = df.withColumn("v", mean(df["v"]).over(w)) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_mixed_sql(self): + df = self.data + w = self.unbounded_window + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("v", mean_udf(df["v"] * 2).over(w) + 1) + expected1 = df.withColumn("v", mean(df["v"] * 2).over(w) + 1) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_mixed_udf(self): + df = self.data + w = self.unbounded_window + + plus_one = self.python_plus_one + time_two = self.pandas_scalar_time_two + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("v2", plus_one(mean_udf(plus_one(df["v"])).over(w))) + expected1 = df.withColumn("v2", plus_one(mean(plus_one(df["v"])).over(w))) + + result2 = df.withColumn("v2", time_two(mean_udf(time_two(df["v"])).over(w))) + expected2 = df.withColumn("v2", time_two(mean(time_two(df["v"])).over(w))) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + + def test_without_partitionBy(self): + df = self.data + w = self.unpartitioned_window + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("v2", mean_udf(df["v"]).over(w)) + expected1 = df.withColumn("v2", mean(df["v"]).over(w)) + + result2 = df.select(mean_udf(df["v"]).over(w)) + expected2 = df.select(mean(df["v"]).over(w)) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + + def test_mixed_sql_and_udf(self): + df = self.data + w = self.unbounded_window + ow = self.ordered_window + max_udf = self.pandas_agg_max_udf + min_udf = self.pandas_agg_min_udf + + result1 = df.withColumn("v_diff", max_udf(df["v"]).over(w) - min_udf(df["v"]).over(w)) + expected1 = df.withColumn("v_diff", max(df["v"]).over(w) - min(df["v"]).over(w)) + + # Test mixing sql window function and window udf in the same expression + result2 = df.withColumn("v_diff", max_udf(df["v"]).over(w) - min(df["v"]).over(w)) + expected2 = expected1 + + # Test chaining sql aggregate function and udf + result3 = ( + df.withColumn("max_v", max_udf(df["v"]).over(w)) + .withColumn("min_v", min(df["v"]).over(w)) + .withColumn("v_diff", col("max_v") - col("min_v")) + .drop("max_v", "min_v") + ) + expected3 = expected1 + + # Test mixing sql window function and udf + result4 = df.withColumn("max_v", max_udf(df["v"]).over(w)).withColumn( + "rank", rank().over(ow) + ) + expected4 = df.withColumn("max_v", max(df["v"]).over(w)).withColumn("rank", rank().over(ow)) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + assert_frame_equal(expected2.toPandas(), result2.toPandas()) + assert_frame_equal(expected3.toPandas(), result3.toPandas()) + assert_frame_equal(expected4.toPandas(), result4.toPandas()) + + def test_array_type(self): + df = self.data + w = self.unbounded_window + + array_udf = pandas_udf(lambda x: [1.0, 2.0], "array", PandasUDFType.GROUPED_AGG) + result1 = df.withColumn("v2", array_udf(df["v"]).over(w)) + self.assertEqual(result1.first()["v2"], [1.0, 2.0]) + + def test_invalid_args(self): + with QuietTest(self.sc): + self.check_invalid_args() + + def check_invalid_args(self): + df = self.data + w = self.unbounded_window + + with self.assertRaisesRegex(AnalysisException, ".*not supported within a window function"): + foo_udf = pandas_udf(lambda x: x, "v double", PandasUDFType.GROUPED_MAP) + df.withColumn("v2", foo_udf(df["v"]).over(w)) + + def test_bounded_simple(self): + from pyspark.sql.functions import mean, max, min, count + + df = self.data + w1 = self.sliding_row_window + w2 = self.shrinking_range_window + + plus_one = self.python_plus_one + count_udf = self.pandas_agg_count_udf + mean_udf = self.pandas_agg_mean_udf + max_udf = self.pandas_agg_max_udf + min_udf = self.pandas_agg_min_udf + + result1 = ( + df.withColumn("mean_v", mean_udf(plus_one(df["v"])).over(w1)) + .withColumn("count_v", count_udf(df["v"]).over(w2)) + .withColumn("max_v", max_udf(df["v"]).over(w2)) + .withColumn("min_v", min_udf(df["v"]).over(w1)) + ) + + expected1 = ( + df.withColumn("mean_v", mean(plus_one(df["v"])).over(w1)) + .withColumn("count_v", count(df["v"]).over(w2)) + .withColumn("max_v", max(df["v"]).over(w2)) + .withColumn("min_v", min(df["v"]).over(w1)) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_growing_window(self): + from pyspark.sql.functions import mean + + df = self.data + w1 = self.growing_row_window + w2 = self.growing_range_window + + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("m1", mean_udf(df["v"]).over(w1)).withColumn( + "m2", mean_udf(df["v"]).over(w2) + ) + + expected1 = df.withColumn("m1", mean(df["v"]).over(w1)).withColumn( + "m2", mean(df["v"]).over(w2) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_sliding_window(self): + from pyspark.sql.functions import mean + + df = self.data + w1 = self.sliding_row_window + w2 = self.sliding_range_window + + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("m1", mean_udf(df["v"]).over(w1)).withColumn( + "m2", mean_udf(df["v"]).over(w2) + ) + + expected1 = df.withColumn("m1", mean(df["v"]).over(w1)).withColumn( + "m2", mean(df["v"]).over(w2) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_shrinking_window(self): + from pyspark.sql.functions import mean + + df = self.data + w1 = self.shrinking_row_window + w2 = self.shrinking_range_window + + mean_udf = self.pandas_agg_mean_udf + + result1 = df.withColumn("m1", mean_udf(df["v"]).over(w1)).withColumn( + "m2", mean_udf(df["v"]).over(w2) + ) + + expected1 = df.withColumn("m1", mean(df["v"]).over(w1)).withColumn( + "m2", mean(df["v"]).over(w2) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + def test_bounded_mixed(self): + from pyspark.sql.functions import mean, max + + df = self.data + w1 = self.sliding_row_window + w2 = self.unbounded_window + + mean_udf = self.pandas_agg_mean_udf + max_udf = self.pandas_agg_max_udf + + result1 = ( + df.withColumn("mean_v", mean_udf(df["v"]).over(w1)) + .withColumn("max_v", max_udf(df["v"]).over(w2)) + .withColumn("mean_unbounded_v", mean_udf(df["v"]).over(w1)) + ) + + expected1 = ( + df.withColumn("mean_v", mean(df["v"]).over(w1)) + .withColumn("max_v", max(df["v"]).over(w2)) + .withColumn("mean_unbounded_v", mean(df["v"]).over(w1)) + ) + + assert_frame_equal(expected1.toPandas(), result1.toPandas()) + + +class WindowPandasUDFTests(WindowPandasUDFTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.pandas.test_pandas_udf_window import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/__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/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming.py new file mode 100644 index 0000000..0eea86d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming.py @@ -0,0 +1,400 @@ +# +# 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 os +import shutil +import tempfile +import time + +from pyspark.sql import Row +from pyspark.sql.functions import lit +from pyspark.sql.types import StructType, StructField, IntegerType, StringType +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.errors.exceptions.connect import SparkConnectException + + +class StreamingTestsMixin: + def test_streaming_query_functions_basic(self): + df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load() + query = ( + df.writeStream.format("memory") + .queryName("test_streaming_query_functions_basic") + .start() + ) + try: + self.assertEquals(query.name, "test_streaming_query_functions_basic") + self.assertTrue(isinstance(query.id, str)) + self.assertTrue(isinstance(query.runId, str)) + self.assertTrue(query.isActive) + self.assertEqual(query.exception(), None) + self.assertFalse(query.awaitTermination(1)) + query.processAllAvailable() + recentProgress = query.recentProgress + lastProgress = query.lastProgress + self.assertEqual(lastProgress["name"], query.name) + self.assertEqual(lastProgress["id"], query.id) + self.assertTrue(any(p == lastProgress for p in recentProgress)) + query.explain() + + except Exception as e: + self.fail( + "Streaming query functions sanity check shouldn't throw any error. " + "Error message: " + str(e) + ) + + finally: + query.stop() + + def test_stream_trigger(self): + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + + # Should take at least one arg + try: + df.writeStream.trigger() + except ValueError: + pass + + # Should not take multiple args + try: + df.writeStream.trigger(once=True, processingTime="5 seconds") + except ValueError: + pass + + # Should not take multiple args + try: + df.writeStream.trigger(processingTime="5 seconds", continuous="1 second") + except ValueError: + pass + + # Should take only keyword args + try: + df.writeStream.trigger("5 seconds") + self.fail("Should have thrown an exception") + except TypeError: + pass + + def test_stream_read_options(self): + schema = StructType([StructField("data", StringType(), False)]) + df = ( + self.spark.readStream.format("text") + .option("path", "python/test_support/sql/streaming") + .schema(schema) + .load() + ) + self.assertTrue(df.isStreaming) + self.assertEqual(df.schema.simpleString(), "struct") + + def test_stream_read_options_overwrite(self): + bad_schema = StructType([StructField("test", IntegerType(), False)]) + schema = StructType([StructField("data", StringType(), False)]) + # SPARK-32516 disables the overwrite behavior by default. + with self.sql_conf({"spark.sql.legacy.pathOptionBehavior.enabled": True}): + df = ( + self.spark.readStream.format("csv") + .option("path", "python/test_support/sql/fake") + .schema(bad_schema) + .load(path="python/test_support/sql/streaming", schema=schema, format="text") + ) + self.assertTrue(df.isStreaming) + self.assertEqual(df.schema.simpleString(), "struct") + + def test_stream_save_options(self): + df = ( + self.spark.readStream.format("text") + .load("python/test_support/sql/streaming") + .withColumn("id", lit(1)) + ) + for q in self.spark.streams.active: + q.stop() + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.assertTrue(df.isStreaming) + out = os.path.join(tmpPath, "out") + chk = os.path.join(tmpPath, "chk") + q = ( + df.writeStream.option("checkpointLocation", chk) + .queryName("this_query") + .format("parquet") + .partitionBy("id") + .outputMode("append") + .option("path", out) + .start() + ) + try: + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + q.processAllAvailable() + output_files = [] + for _, _, files in os.walk(out): + output_files.extend([f for f in files if not f.startswith(".")]) + self.assertTrue(len(output_files) > 0) + self.assertTrue(len(os.listdir(chk)) > 0) + finally: + q.stop() + shutil.rmtree(tmpPath) + + def test_stream_save_options_overwrite(self): + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + for q in self.spark.streams.active: + q.stop() + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.assertTrue(df.isStreaming) + out = os.path.join(tmpPath, "out") + chk = os.path.join(tmpPath, "chk") + fake1 = os.path.join(tmpPath, "fake1") + fake2 = os.path.join(tmpPath, "fake2") + # SPARK-32516 disables the overwrite behavior by default. + with self.sql_conf({"spark.sql.legacy.pathOptionBehavior.enabled": True}): + q = ( + df.writeStream.option("checkpointLocation", fake1) + .format("memory") + .option("path", fake2) + .queryName("fake_query") + .outputMode("append") + .start(path=out, format="parquet", queryName="this_query", checkpointLocation=chk) + ) + + try: + self.assertEqual(q.name, "this_query") + self.assertTrue(q.isActive) + q.processAllAvailable() + output_files = [] + for _, _, files in os.walk(out): + output_files.extend([f for f in files if not f.startswith(".")]) + self.assertTrue(len(output_files) > 0) + self.assertTrue(len(os.listdir(chk)) > 0) + self.assertFalse(os.path.isdir(fake1)) # should not have been created + self.assertFalse(os.path.isdir(fake2)) # should not have been created + finally: + q.stop() + shutil.rmtree(tmpPath) + + def test_stream_status_and_progress(self): + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + for q in self.spark.streams.active: + q.stop() + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.assertTrue(df.isStreaming) + out = os.path.join(tmpPath, "out") + chk = os.path.join(tmpPath, "chk") + + def func(x): + time.sleep(1) + return x + + from pyspark.sql.functions import col, udf + + sleep_udf = udf(func) + + # Use "sleep_udf" to delay the progress update so that we can test `lastProgress` when there + # were no updates. + q = df.select(sleep_udf(col("value")).alias("value")).writeStream.start( + path=out, format="parquet", queryName="this_query", checkpointLocation=chk + ) + try: + # "lastProgress" will return None in most cases. However, as it may be flaky when + # Jenkins is very slow, we don't assert it. If there is something wrong, "lastProgress" + # may throw error with a high chance and make this test flaky, so we should still be + # able to detect broken codes. + q.lastProgress + + q.processAllAvailable() + lastProgress = q.lastProgress + recentProgress = q.recentProgress + status = q.status + self.assertEqual(lastProgress["name"], q.name) + self.assertEqual(lastProgress["id"], q.id) + self.assertTrue(any(p == lastProgress for p in recentProgress)) + self.assertTrue( + "message" in status and "isDataAvailable" in status and "isTriggerActive" in status + ) + finally: + q.stop() + shutil.rmtree(tmpPath) + + def test_stream_await_termination(self): + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + for q in self.spark.streams.active: + q.stop() + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.assertTrue(df.isStreaming) + out = os.path.join(tmpPath, "out") + chk = os.path.join(tmpPath, "chk") + q = df.writeStream.start( + path=out, format="parquet", queryName="this_query", checkpointLocation=chk + ) + try: + self.assertTrue(q.isActive) + try: + q.awaitTermination("hello") + self.fail("Expected a value exception") + except ValueError: + pass + now = time.time() + # test should take at least 2 seconds + res = q.awaitTermination(2.6) + duration = time.time() - now + self.assertTrue(duration >= 2) + self.assertFalse(res) + + q.processAllAvailable() + q.stop() + # Sanity check when no parameter is set + q.awaitTermination() + self.assertFalse(q.isActive) + finally: + q.stop() + shutil.rmtree(tmpPath) + + def test_stream_exception(self): + sdf = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + sq = sdf.writeStream.format("memory").queryName("query_explain").start() + try: + sq.processAllAvailable() + self.assertEqual(sq.exception(), None) + finally: + sq.stop() + + from pyspark.sql.functions import col, udf + from pyspark.errors import StreamingQueryException + + bad_udf = udf(lambda x: 1 / 0) + sq = ( + sdf.select(bad_udf(col("value"))) + .writeStream.format("memory") + .queryName("this_query") + .start() + ) + try: + # Process some data to fail the query + sq.processAllAvailable() + self.fail("bad udf should fail the query") + except StreamingQueryException as e: + # This is expected + self._assert_exception_tree_contains_msg(e, "ZeroDivisionError") + finally: + exception = sq.exception() + sq.stop() + self.assertIsInstance(exception, StreamingQueryException) + self._assert_exception_tree_contains_msg(exception, "ZeroDivisionError") + + def _assert_exception_tree_contains_msg(self, exception, msg): + if isinstance(exception, SparkConnectException): + self._assert_exception_tree_contains_msg_connect(exception, msg) + else: + self._assert_exception_tree_contains_msg_default(exception, msg) + + def _assert_exception_tree_contains_msg_connect(self, exception, msg): + self.assertTrue( + msg in exception.message, + "Exception tree doesn't contain the expected message: %s" % msg, + ) + + def _assert_exception_tree_contains_msg_default(self, exception, msg): + e = exception + contains = msg in e.desc + while e.cause is not None and not contains: + e = e.cause + contains = msg in e.desc + self.assertTrue(contains, "Exception tree doesn't contain the expected message: %s" % msg) + + def test_query_manager_get(self): + df = self.spark.readStream.format("rate").load() + for q in self.spark.streams.active: + q.stop() + q = df.writeStream.format("noop").start() + + self.assertTrue(q.isActive) + self.assertTrue(q.id == self.spark.streams.get(q.id).id) + + q.stop() + + self.assertIsNone(self.spark.streams.get(q.id)) + + def test_query_manager_await_termination(self): + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + for q in self.spark.streams.active: + q.stop() + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.assertTrue(df.isStreaming) + out = os.path.join(tmpPath, "out") + chk = os.path.join(tmpPath, "chk") + q = df.writeStream.start( + path=out, format="parquet", queryName="this_query", checkpointLocation=chk + ) + try: + self.assertTrue(q.isActive) + try: + self.spark.streams.awaitAnyTermination("hello") + self.fail("Expected a value exception") + except ValueError: + pass + now = time.time() + # test should take at least 2 seconds + res = self.spark.streams.awaitAnyTermination(2.6) + duration = time.time() - now + self.assertTrue(duration >= 2) + self.assertFalse(res) + finally: + q.processAllAvailable() + q.stop() + shutil.rmtree(tmpPath) + + def test_streaming_read_from_table(self): + with self.table("input_table", "this_query"): + self.spark.sql("CREATE TABLE input_table (value string) USING parquet") + self.spark.sql("INSERT INTO input_table VALUES ('aaa'), ('bbb'), ('ccc')") + df = self.spark.readStream.table("input_table") + self.assertTrue(df.isStreaming) + q = df.writeStream.format("memory").queryName("this_query").start() + q.processAllAvailable() + q.stop() + result = self.spark.sql("SELECT * FROM this_query ORDER BY value").collect() + self.assertEqual( + set([Row(value="aaa"), Row(value="bbb"), Row(value="ccc")]), set(result) + ) + + def test_streaming_write_to_table(self): + with self.table("output_table"), tempfile.TemporaryDirectory() as tmpdir: + df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load() + q = df.writeStream.toTable("output_table", format="parquet", checkpointLocation=tmpdir) + self.assertTrue(q.isActive) + time.sleep(10) + q.stop() + result = self.spark.sql("SELECT value FROM output_table").collect() + self.assertTrue(len(result) > 0) + + +class StreamingTests(StreamingTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.streaming.test_streaming import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach.py new file mode 100644 index 0000000..0c5347d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach.py @@ -0,0 +1,295 @@ +# +# 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 os +import tempfile + +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class StreamingTestsForeachMixin: + class ForeachWriterTester: + def __init__(self, spark): + self.spark = spark + + def write_open_event(self, partitionId, epochId): + self._write_event(self.open_events_dir, {"partition": partitionId, "epoch": epochId}) + + def write_process_event(self, row): + self._write_event(self.process_events_dir, {"value": "text"}) + + def write_close_event(self, error): + self._write_event(self.close_events_dir, {"error": str(error)}) + + def write_input_file(self): + self._write_event(self.input_dir, "text") + + def open_events(self): + return self._read_events(self.open_events_dir, "partition INT, epoch INT") + + def process_events(self): + return self._read_events(self.process_events_dir, "value STRING") + + def close_events(self): + return self._read_events(self.close_events_dir, "error STRING") + + def run_streaming_query_on_writer(self, writer, num_files): + self._reset() + try: + sdf = self.spark.readStream.format("text").load(self.input_dir) + sq = sdf.writeStream.foreach(writer).start() + for i in range(num_files): + self.write_input_file() + sq.processAllAvailable() + finally: + self.stop_all() + + def assert_invalid_writer(self, writer, msg=None): + self._reset() + try: + sdf = self.spark.readStream.format("text").load(self.input_dir) + sq = sdf.writeStream.foreach(writer).start() + self.write_input_file() + sq.processAllAvailable() + self.fail("invalid writer %s did not fail the query" % str(writer)) # not expected + except Exception as e: + if msg: + assert msg in str(e), "%s not in %s" % (msg, str(e)) + + finally: + self.stop_all() + + def stop_all(self): + for q in self.spark.streams.active: + q.stop() + + def _reset(self): + self.input_dir = tempfile.mkdtemp() + self.open_events_dir = tempfile.mkdtemp() + self.process_events_dir = tempfile.mkdtemp() + self.close_events_dir = tempfile.mkdtemp() + + def _read_events(self, dir, json): + rows = self.spark.read.schema(json).json(dir).collect() + dicts = [row.asDict() for row in rows] + return dicts + + def _write_event(self, dir, event): + import uuid + + with open(os.path.join(dir, str(uuid.uuid4())), "w") as f: + f.write("%s\n" % str(event)) + + def __getstate__(self): + return (self.open_events_dir, self.process_events_dir, self.close_events_dir) + + def __setstate__(self, state): + self.open_events_dir, self.process_events_dir, self.close_events_dir = state + + # Those foreach tests are failed in macOS High Sierra by defined rules + # at http://sealiesoftware.com/blog/archive/2017/6/5/Objective-C_and_fork_in_macOS_1013.html + # To work around this, OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES. + def test_streaming_foreach_with_simple_function(self): + tester = self.ForeachWriterTester(self.spark) + + def foreach_func(row): + tester.write_process_event(row) + + tester.run_streaming_query_on_writer(foreach_func, 2) + self.assertEqual(len(tester.process_events()), 2) + + def test_streaming_foreach_with_basic_open_process_close(self): + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def open(self, partitionId, epochId): + tester.write_open_event(partitionId, epochId) + return True + + def process(self, row): + tester.write_process_event(row) + + def close(self, error): + tester.write_close_event(error) + + tester.run_streaming_query_on_writer(ForeachWriter(), 2) + + open_events = tester.open_events() + self.assertEqual(len(open_events), 2) + self.assertSetEqual(set([e["epoch"] for e in open_events]), {0, 1}) + + self.assertEqual(len(tester.process_events()), 2) + + close_events = tester.close_events() + self.assertEqual(len(close_events), 2) + self.assertSetEqual(set([e["error"] for e in close_events]), {"None"}) + + def test_streaming_foreach_with_open_returning_false(self): + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def open(self, partition_id, epoch_id): + tester.write_open_event(partition_id, epoch_id) + return False + + def process(self, row): + tester.write_process_event(row) + + def close(self, error): + tester.write_close_event(error) + + tester.run_streaming_query_on_writer(ForeachWriter(), 2) + + self.assertEqual(len(tester.open_events()), 2) + + self.assertEqual(len(tester.process_events()), 0) # no row was processed + + close_events = tester.close_events() + self.assertEqual(len(close_events), 2) + self.assertSetEqual(set([e["error"] for e in close_events]), {"None"}) + + def test_streaming_foreach_without_open_method(self): + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def process(self, row): + tester.write_process_event(row) + + def close(self, error): + tester.write_close_event(error) + + tester.run_streaming_query_on_writer(ForeachWriter(), 2) + self.assertEqual(len(tester.open_events()), 0) # no open events + self.assertEqual(len(tester.process_events()), 2) + self.assertEqual(len(tester.close_events()), 2) + + def test_streaming_foreach_without_close_method(self): + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def open(self, partition_id, epoch_id): + tester.write_open_event(partition_id, epoch_id) + return True + + def process(self, row): + tester.write_process_event(row) + + tester.run_streaming_query_on_writer(ForeachWriter(), 2) + self.assertEqual(len(tester.open_events()), 2) # no open events + self.assertEqual(len(tester.process_events()), 2) + self.assertEqual(len(tester.close_events()), 0) + + def test_streaming_foreach_without_open_and_close_methods(self): + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def process(self, row): + tester.write_process_event(row) + + tester.run_streaming_query_on_writer(ForeachWriter(), 2) + self.assertEqual(len(tester.open_events()), 0) # no open events + self.assertEqual(len(tester.process_events()), 2) + self.assertEqual(len(tester.close_events()), 0) + + def test_streaming_foreach_with_process_throwing_error(self): + from pyspark.errors import StreamingQueryException + + tester = self.ForeachWriterTester(self.spark) + + class ForeachWriter: + def process(self, row): + raise RuntimeError("test error") + + def close(self, error): + tester.write_close_event(error) + + try: + tester.run_streaming_query_on_writer(ForeachWriter(), 1) + self.fail("bad writer did not fail the query") # this is not expected + except StreamingQueryException: + # TODO: Verify whether original error message is inside the exception + pass + + self.assertEqual(len(tester.process_events()), 0) # no row was processed + close_events = tester.close_events() + self.assertEqual(len(close_events), 1) + # TODO: Verify whether original error message is inside the exception + + def test_streaming_foreach_with_invalid_writers(self): + + tester = self.ForeachWriterTester(self.spark) + + def func_with_iterator_input(iter): + for x in iter: + print(x) + + tester.assert_invalid_writer(func_with_iterator_input) + + class WriterWithoutProcess: + def open(self, partition): + pass + + tester.assert_invalid_writer(WriterWithoutProcess(), "ATTRIBUTE_NOT_CALLABLE") + + class WriterWithNonCallableProcess: + process = True + + tester.assert_invalid_writer(WriterWithNonCallableProcess(), "ATTRIBUTE_NOT_CALLABLE") + + class WriterWithNoParamProcess: + def process(self): + pass + + tester.assert_invalid_writer(WriterWithNoParamProcess()) + + # Abstract class for tests below + class WithProcess: + def process(self, row): + pass + + class WriterWithNonCallableOpen(WithProcess): + open = True + + tester.assert_invalid_writer(WriterWithNonCallableOpen(), "ATTRIBUTE_NOT_CALLABLE") + + class WriterWithNoParamOpen(WithProcess): + def open(self): + pass + + tester.assert_invalid_writer(WriterWithNoParamOpen()) + + class WriterWithNonCallableClose(WithProcess): + close = True + + tester.assert_invalid_writer(WriterWithNonCallableClose(), "ATTRIBUTE_NOT_CALLABLE") + + +class StreamingTestsForeach(StreamingTestsForeachMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.streaming.test_streaming_foreach import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py new file mode 100644 index 0000000..393101a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py @@ -0,0 +1,106 @@ +# +# 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 time + +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class StreamingTestsForeachBatchMixin: + def test_streaming_foreach_batch(self): + q = None + + def collectBatch(batch_df, batch_id): + batch_df.createOrReplaceGlobalTempView("test_view") + + try: + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + q = df.writeStream.foreachBatch(collectBatch).start() + q.processAllAvailable() + collected = self.spark.sql("select * from global_temp.test_view").collect() + self.assertTrue(len(collected), 2) + finally: + if q: + q.stop() + + def test_streaming_foreach_batch_tempview(self): + q = None + + def collectBatch(batch_df, batch_id): + batch_df.createOrReplaceTempView("updates") + # it should use the spark session within given DataFrame, as microbatch execution will + # clone the session which is no longer same with the session used to start the + # streaming query + assert len(batch_df.sparkSession.sql("SELECT * FROM updates").collect()) == 2 + # Write to a global view verify on the repl/client side. + batch_df.createOrReplaceGlobalTempView("temp_view") + + try: + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + q = df.writeStream.foreachBatch(collectBatch).start() + q.processAllAvailable() + collected = self.spark.sql("SELECT * FROM global_temp.temp_view").collect() + self.assertTrue(len(collected[0]), 2) + finally: + if q: + q.stop() + + def test_streaming_foreach_batch_propagates_python_errors(self): + from pyspark.errors import StreamingQueryException + + q = None + + def collectBatch(df, id): + raise RuntimeError("this should fail the query") + + try: + df = self.spark.readStream.format("text").load("python/test_support/sql/streaming") + q = df.writeStream.foreachBatch(collectBatch).start() + q.processAllAvailable() + self.fail("Expected a failure") + except StreamingQueryException as e: + self.assertTrue("this should fail" in str(e)) + finally: + if q: + q.stop() + + def test_streaming_foreach_batch_graceful_stop(self): + # SPARK-39218: Make foreachBatch streaming query stop gracefully + def func(batch_df, _): + batch_df.sparkSession._jvm.java.lang.Thread.sleep(10000) + + q = self.spark.readStream.format("rate").load().writeStream.foreachBatch(func).start() + time.sleep(3) # 'rowsPerSecond' defaults to 1. Waits 3 secs out for the input. + q.stop() + self.assertIsNone(q.exception(), "No exception has to be propagated.") + + +class StreamingTestsForeachBatch(StreamingTestsForeachBatchMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.streaming.test_streaming_foreach_batch import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_listener.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_listener.py new file mode 100644 index 0000000..87d0dae --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/streaming/test_streaming_listener.py @@ -0,0 +1,517 @@ +# +# 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 json +import time +import uuid +from datetime import datetime + +from pyspark import Row +from pyspark.sql.streaming import StreamingQueryListener +from pyspark.sql.streaming.listener import ( + QueryStartedEvent, + QueryProgressEvent, + QueryTerminatedEvent, + SinkProgress, + SourceProgress, + StateOperatorProgress, + StreamingQueryProgress, +) +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class StreamingListenerTestsMixin: + def check_start_event(self, event): + """Check QueryStartedEvent""" + self.assertTrue(isinstance(event, QueryStartedEvent)) + self.assertTrue(isinstance(event.id, uuid.UUID)) + self.assertTrue(isinstance(event.runId, uuid.UUID)) + self.assertTrue(event.name is None or event.name == "test") + try: + datetime.strptime(event.timestamp, "%Y-%m-%dT%H:%M:%S.%fZ") + except ValueError: + self.fail("'%s' is not in ISO 8601 format.") + + def check_progress_event(self, event): + """Check QueryProgressEvent""" + self.assertTrue(isinstance(event, QueryProgressEvent)) + self.check_streaming_query_progress(event.progress) + + def check_terminated_event(self, event, exception=None, error_class=None): + """Check QueryTerminatedEvent""" + self.assertTrue(isinstance(event, QueryTerminatedEvent)) + self.assertTrue(isinstance(event.id, uuid.UUID)) + self.assertTrue(isinstance(event.runId, uuid.UUID)) + if exception: + self.assertTrue(exception in event.exception) + else: + self.assertEquals(event.exception, None) + + if error_class: + self.assertTrue(error_class in event.errorClassOnException) + else: + self.assertEquals(event.errorClassOnException, None) + + def check_streaming_query_progress(self, progress): + """Check StreamingQueryProgress""" + self.assertTrue(isinstance(progress, StreamingQueryProgress)) + self.assertTrue(isinstance(progress.id, uuid.UUID)) + self.assertTrue(isinstance(progress.runId, uuid.UUID)) + self.assertEquals(progress.name, "test") + try: + json.loads(progress.json) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(progress.prettyJson) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(str(progress)) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + datetime.strptime(progress.timestamp, "%Y-%m-%dT%H:%M:%S.%fZ") + except Exception: + self.fail("'%s' is not in ISO 8601 format.") + self.assertTrue(isinstance(progress.batchId, int)) + self.assertTrue(progress.batchDuration is None or isinstance(progress.batchDuration, int)) + self.assertTrue(isinstance(progress.durationMs, dict)) + self.assertTrue( + set(progress.durationMs.keys()).issubset( + { + "triggerExecution", + "queryPlanning", + "getBatch", + "commitOffsets", + "latestOffset", + "addBatch", + "walCommit", + } + ) + ) + self.assertTrue(all(map(lambda v: isinstance(v, int), progress.durationMs.values()))) + + self.assertTrue(all(map(lambda v: isinstance(v, str), progress.eventTime.values()))) + + self.assertTrue(isinstance(progress.stateOperators, list)) + self.assertTrue(len(progress.stateOperators) >= 1) + for so in progress.stateOperators: + self.check_state_operator_progress(so) + + self.assertTrue(isinstance(progress.sources, list)) + self.assertTrue(len(progress.sources) >= 1) + for so in progress.sources: + self.check_source_progress(so) + + self.assertTrue(isinstance(progress.sink, SinkProgress)) + self.check_sink_progress(progress.sink) + self.assertTrue(isinstance(progress.observedMetrics, dict)) + + def check_state_operator_progress(self, progress): + """Check StateOperatorProgress""" + self.assertTrue(isinstance(progress, StateOperatorProgress)) + try: + json.loads(progress.json) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(progress.prettyJson) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(str(progress)) + except Exception: + self.fail("'%s' is not a valid JSON.") + self.assertTrue(isinstance(progress.operatorName, str)) + self.assertTrue(isinstance(progress.numRowsTotal, int)) + self.assertTrue(isinstance(progress.numRowsUpdated, int)) + self.assertTrue(isinstance(progress.allUpdatesTimeMs, int)) + self.assertTrue(isinstance(progress.numRowsRemoved, int)) + self.assertTrue(isinstance(progress.allRemovalsTimeMs, int)) + self.assertTrue(isinstance(progress.commitTimeMs, int)) + self.assertTrue(isinstance(progress.memoryUsedBytes, int)) + self.assertTrue(isinstance(progress.numRowsDroppedByWatermark, int)) + self.assertTrue(isinstance(progress.numShufflePartitions, int)) + self.assertTrue(isinstance(progress.numStateStoreInstances, int)) + self.assertTrue(isinstance(progress.customMetrics, dict)) + + def check_source_progress(self, progress): + """Check SourceProgress""" + self.assertTrue(isinstance(progress, SourceProgress)) + try: + json.loads(progress.json) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(progress.prettyJson) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(str(progress)) + except Exception: + self.fail("'%s' is not a valid JSON.") + self.assertTrue(isinstance(progress.description, str)) + self.assertTrue(isinstance(progress.startOffset, (str, type(None)))) + self.assertTrue(isinstance(progress.endOffset, (str, type(None)))) + self.assertTrue(isinstance(progress.latestOffset, (str, type(None)))) + self.assertTrue(isinstance(progress.numInputRows, int)) + self.assertTrue(isinstance(progress.inputRowsPerSecond, float)) + self.assertTrue(isinstance(progress.processedRowsPerSecond, float)) + self.assertTrue(isinstance(progress.metrics, dict)) + + def check_sink_progress(self, progress): + """Check SinkProgress""" + self.assertTrue(isinstance(progress, SinkProgress)) + try: + json.loads(progress.json) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(progress.prettyJson) + except Exception: + self.fail("'%s' is not a valid JSON.") + try: + json.loads(str(progress)) + except Exception: + self.fail("'%s' is not a valid JSON.") + self.assertTrue(isinstance(progress.description, str)) + self.assertTrue(isinstance(progress.numOutputRows, int)) + self.assertTrue(isinstance(progress.metrics, dict)) + + +class StreamingListenerTests(StreamingListenerTestsMixin, ReusedSQLTestCase): + def test_number_of_public_methods(self): + msg = ( + "New field or method was detected in JVM side. If you added a new public " + "field or method, implement that in the corresponding Python class too." + "Otherwise, fix the number on the assert here." + ) + + def get_number_of_public_methods(clz): + return len( + self.spark.sparkContext._jvm.org.apache.spark.util.Utils.classForName( + clz, True, False + ).getMethods() + ) + + self.assertEquals( + get_number_of_public_methods( + "org.apache.spark.sql.streaming.StreamingQueryListener$QueryStartedEvent" + ), + 15, + msg, + ) + self.assertEquals( + get_number_of_public_methods( + "org.apache.spark.sql.streaming.StreamingQueryListener$QueryProgressEvent" + ), + 12, + msg, + ) + self.assertEquals( + get_number_of_public_methods( + "org.apache.spark.sql.streaming.StreamingQueryListener$QueryTerminatedEvent" + ), + 15, + msg, + ) + self.assertEquals( + get_number_of_public_methods("org.apache.spark.sql.streaming.StreamingQueryProgress"), + 38, + msg, + ) + self.assertEquals( + get_number_of_public_methods("org.apache.spark.sql.streaming.StateOperatorProgress"), + 27, + msg, + ) + self.assertEquals( + get_number_of_public_methods("org.apache.spark.sql.streaming.SourceProgress"), 21, msg + ) + self.assertEquals( + get_number_of_public_methods("org.apache.spark.sql.streaming.SinkProgress"), 19, msg + ) + + def test_listener_events(self): + start_event = None + progress_event = None + terminated_event = None + + class TestListener(StreamingQueryListener): + def onQueryStarted(self, event): + nonlocal start_event + start_event = event + + def onQueryProgress(self, event): + nonlocal progress_event + progress_event = event + + def onQueryIdle(self, event): + pass + + def onQueryTerminated(self, event): + nonlocal terminated_event + terminated_event = event + + test_listener = TestListener() + + try: + self.spark.streams.addListener(test_listener) + + df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load() + + # check successful stateful query + df_stateful = df.groupBy().count() # make query stateful + q = ( + df_stateful.writeStream.format("noop") + .queryName("test") + .outputMode("complete") + .start() + ) + self.assertTrue(q.isActive) + time.sleep(10) + q.stop() + + # Make sure all events are empty + self.spark.sparkContext._jsc.sc().listenerBus().waitUntilEmpty() + + self.check_start_event(start_event) + self.check_progress_event(progress_event) + self.check_terminated_event(terminated_event) + + # Check query terminated with exception + from pyspark.sql.functions import col, udf + + bad_udf = udf(lambda x: 1 / 0) + q = df.select(bad_udf(col("value"))).writeStream.format("noop").start() + time.sleep(5) + q.stop() + self.spark.sparkContext._jsc.sc().listenerBus().waitUntilEmpty() + self.check_terminated_event(terminated_event, "ZeroDivisionError") + + finally: + self.spark.streams.removeListener(test_listener) + + def test_remove_listener(self): + # SPARK-38804: Test StreamingQueryManager.removeListener + class TestListener(StreamingQueryListener): + def onQueryStarted(self, event): + pass + + def onQueryProgress(self, event): + pass + + def onQueryIdle(self, event): + pass + + def onQueryTerminated(self, event): + pass + + test_listener = TestListener() + + num_listeners = len(self.spark.streams._jsqm.listListeners()) + self.spark.streams.addListener(test_listener) + self.assertEqual(num_listeners + 1, len(self.spark.streams._jsqm.listListeners())) + self.spark.streams.removeListener(test_listener) + self.assertEqual(num_listeners, len(self.spark.streams._jsqm.listListeners())) + + def test_query_started_event_fromJson(self): + start_event = """ + { + "id" : "78923ec2-8f4d-4266-876e-1f50cf3c283b", + "runId" : "55a95d45-e932-4e08-9caa-0a8ecd9391e8", + "name" : null, + "timestamp" : "2023-06-09T18:13:29.741Z" + } + """ + start_event = QueryStartedEvent.fromJson(json.loads(start_event)) + self.check_start_event(start_event) + self.assertEqual(start_event.id, uuid.UUID("78923ec2-8f4d-4266-876e-1f50cf3c283b")) + self.assertEqual(start_event.runId, uuid.UUID("55a95d45-e932-4e08-9caa-0a8ecd9391e8")) + self.assertIsNone(start_event.name) + self.assertEqual(start_event.timestamp, "2023-06-09T18:13:29.741Z") + + def test_query_terminated_event_fromJson(self): + terminated_json = """ + { + "id" : "78923ec2-8f4d-4266-876e-1f50cf3c283b", + "runId" : "55a95d45-e932-4e08-9caa-0a8ecd9391e8", + "exception" : "org.apache.spark.SparkException: Job aborted due to stage failure", + "errorClassOnException" : null} + """ + terminated_event = QueryTerminatedEvent.fromJson(json.loads(terminated_json)) + self.check_terminated_event(terminated_event, "SparkException") + self.assertEqual(terminated_event.id, uuid.UUID("78923ec2-8f4d-4266-876e-1f50cf3c283b")) + self.assertEqual(terminated_event.runId, uuid.UUID("55a95d45-e932-4e08-9caa-0a8ecd9391e8")) + self.assertIn("SparkException", terminated_event.exception) + self.assertIsNone(terminated_event.errorClassOnException) + + def test_streaming_query_progress_fromJson(self): + progress_json = """ + { + "id" : "00000000-0000-0001-0000-000000000001", + "runId" : "00000000-0000-0001-0000-000000000002", + "name" : "test", + "timestamp" : "2016-12-05T20:54:20.827Z", + "batchId" : 2, + "numInputRows" : 678, + "inputRowsPerSecond" : 10.0, + "processedRowsPerSecond" : 5.4, + "batchDuration": 5, + "durationMs" : { + "getBatch" : 0 + }, + "eventTime" : { + "min" : "2016-12-05T20:54:20.827Z", + "avg" : "2016-12-05T20:54:20.827Z", + "watermark" : "2016-12-05T20:54:20.827Z", + "max" : "2016-12-05T20:54:20.827Z" + }, + "stateOperators" : [ { + "operatorName" : "op1", + "numRowsTotal" : 0, + "numRowsUpdated" : 1, + "allUpdatesTimeMs" : 1, + "numRowsRemoved" : 2, + "allRemovalsTimeMs" : 34, + "commitTimeMs" : 23, + "memoryUsedBytes" : 3, + "numRowsDroppedByWatermark" : 0, + "numShufflePartitions" : 2, + "numStateStoreInstances" : 2, + "customMetrics" : { + "loadedMapCacheHitCount" : 1, + "loadedMapCacheMissCount" : 0, + "stateOnCurrentVersionSizeBytes" : 2 + } + } ], + "sources" : [ { + "description" : "source", + "startOffset" : 123, + "endOffset" : 456, + "latestOffset" : 789, + "numInputRows" : 678, + "inputRowsPerSecond" : 10.0, + "processedRowsPerSecond" : 5.4, + "metrics": {} + } ], + "sink" : { + "description" : "sink", + "numOutputRows" : -1, + "metrics": {} + }, + "observedMetrics" : { + "event1" : { + "c1" : 1, + "c2" : 3.0 + }, + "event2" : { + "rc" : 1, + "min_q" : "hello", + "max_q" : "world" + } + } + } + """ + progress = StreamingQueryProgress.fromJson(json.loads(progress_json)) + + self.check_streaming_query_progress(progress) + + # checks for progress + self.assertEqual(progress.id, uuid.UUID("00000000-0000-0001-0000-000000000001")) + self.assertEqual(progress.runId, uuid.UUID("00000000-0000-0001-0000-000000000002")) + self.assertEqual(progress.name, "test") + self.assertEqual(progress.timestamp, "2016-12-05T20:54:20.827Z") + self.assertEqual(progress.batchId, 2) + self.assertEqual(progress.numInputRows, 678) + self.assertEqual(progress.inputRowsPerSecond, 10.0) + self.assertEqual(progress.batchDuration, 5) + self.assertEqual(progress.durationMs, {"getBatch": 0}) + self.assertEqual( + progress.eventTime, + { + "min": "2016-12-05T20:54:20.827Z", + "avg": "2016-12-05T20:54:20.827Z", + "watermark": "2016-12-05T20:54:20.827Z", + "max": "2016-12-05T20:54:20.827Z", + }, + ) + self.assertEqual( + progress.observedMetrics, + { + "event1": Row("c1", "c2")(1, 3.0), + "event2": Row("rc", "min_q", "max_q")(1, "hello", "world"), + }, + ) + + # Check stateOperators list + self.assertEqual(len(progress.stateOperators), 1) + state_operator = progress.stateOperators[0] + self.assertTrue(isinstance(state_operator, StateOperatorProgress)) + self.assertEqual(state_operator.operatorName, "op1") + self.assertEqual(state_operator.numRowsTotal, 0) + self.assertEqual(state_operator.numRowsUpdated, 1) + self.assertEqual(state_operator.allUpdatesTimeMs, 1) + self.assertEqual(state_operator.numRowsRemoved, 2) + self.assertEqual(state_operator.allRemovalsTimeMs, 34) + self.assertEqual(state_operator.commitTimeMs, 23) + self.assertEqual(state_operator.memoryUsedBytes, 3) + self.assertEqual(state_operator.numRowsDroppedByWatermark, 0) + self.assertEqual(state_operator.numShufflePartitions, 2) + self.assertEqual(state_operator.numStateStoreInstances, 2) + self.assertEqual( + state_operator.customMetrics, + { + "loadedMapCacheHitCount": 1, + "loadedMapCacheMissCount": 0, + "stateOnCurrentVersionSizeBytes": 2, + }, + ) + + # Check sources list + self.assertEqual(len(progress.sources), 1) + source = progress.sources[0] + self.assertTrue(isinstance(source, SourceProgress)) + self.assertEqual(source.description, "source") + self.assertEqual(source.startOffset, "123") + self.assertEqual(source.endOffset, "456") + self.assertEqual(source.latestOffset, "789") + self.assertEqual(source.numInputRows, 678) + self.assertEqual(source.inputRowsPerSecond, 10.0) + self.assertEqual(source.processedRowsPerSecond, 5.4) + self.assertEqual(source.metrics, {}) + + # Check sink + sink = progress.sink + self.assertTrue(isinstance(sink, SinkProgress)) + self.assertEqual(sink.description, "sink") + self.assertEqual(sink.numOutputRows, -1) + self.assertEqual(sink.metrics, {}) + + +if __name__ == "__main__": + import unittest + + from pyspark.sql.tests.streaming.test_streaming_listener import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow.py new file mode 100644 index 0000000..73b6067 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow.py @@ -0,0 +1,1307 @@ +# +# 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 datetime +import os +import threading +import time +import unittest +import warnings +from distutils.version import LooseVersion +from typing import cast +from collections import namedtuple + +from pyspark import SparkContext, SparkConf +from pyspark.sql import Row, SparkSession +from pyspark.sql.functions import rand, udf, assert_true, lit +from pyspark.sql.types import ( + StructType, + StringType, + IntegerType, + LongType, + FloatType, + DoubleType, + DecimalType, + DateType, + TimestampType, + TimestampNTZType, + BinaryType, + StructField, + ArrayType, + MapType, + NullType, + DayTimeIntervalType, +) +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, + ExamplePoint, + ExamplePointUDT, +) +from pyspark.testing.utils import QuietTest +from pyspark.errors import ArithmeticException, PySparkTypeError, UnsupportedOperationException + +if have_pandas: + import pandas as pd + from pandas.testing import assert_frame_equal + +if have_pyarrow: + import pyarrow as pa # noqa: F401 + + +class ArrowTestsMixin: + @classmethod + def setUpClass(cls): + from datetime import date, datetime + from decimal import Decimal + + super().setUpClass() + cls.warnings_lock = threading.Lock() + + # Synchronize default timezone between Python and Java + cls.tz_prev = os.environ.get("TZ", None) # save current tz if set + tz = "America/Los_Angeles" + os.environ["TZ"] = tz + time.tzset() + + cls.spark.conf.set("spark.sql.session.timeZone", tz) + + # Test fallback + cls.spark.conf.set("spark.sql.execution.arrow.enabled", "false") + assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.enabled") == "false" + cls.spark.conf.set("spark.sql.execution.arrow.enabled", "true") + assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.enabled") == "true" + + cls.spark.conf.set("spark.sql.execution.arrow.fallback.enabled", "true") + assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.fallback.enabled") == "true" + cls.spark.conf.set("spark.sql.execution.arrow.fallback.enabled", "false") + assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.fallback.enabled") == "false" + + # Enable Arrow optimization in this tests. + cls.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true") + # Disable fallback by default to easily detect the failures. + cls.spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "false") + + cls.schema_wo_null = StructType( + [ + StructField("1_str_t", StringType(), True), + StructField("2_int_t", IntegerType(), True), + StructField("3_long_t", LongType(), True), + StructField("4_float_t", FloatType(), True), + StructField("5_double_t", DoubleType(), True), + StructField("6_decimal_t", DecimalType(38, 18), True), + StructField("7_date_t", DateType(), True), + StructField("8_timestamp_t", TimestampType(), True), + StructField("9_binary_t", BinaryType(), True), + ] + ) + cls.schema = cls.schema_wo_null.add("10_null_t", NullType(), True) + cls.data_wo_null = [ + ( + "a", + 1, + 10, + 0.2, + 2.0, + Decimal("2.0"), + date(1969, 1, 1), + datetime(1969, 1, 1, 1, 1, 1), + bytearray(b"a"), + ), + ( + "b", + 2, + 20, + 0.4, + 4.0, + Decimal("4.0"), + date(2012, 2, 2), + datetime(2012, 2, 2, 2, 2, 2), + bytearray(b"bb"), + ), + ( + "c", + 3, + 30, + 0.8, + 6.0, + Decimal("6.0"), + date(2100, 3, 3), + datetime(2100, 3, 3, 3, 3, 3), + bytearray(b"ccc"), + ), + ( + "d", + 4, + 40, + 1.0, + 8.0, + Decimal("8.0"), + date(2262, 4, 12), + datetime(2262, 3, 3, 3, 3, 3), + bytearray(b"dddd"), + ), + ] + cls.data = [tuple(list(d) + [None]) for d in cls.data_wo_null] + + @classmethod + def tearDownClass(cls): + del os.environ["TZ"] + if cls.tz_prev is not None: + os.environ["TZ"] = cls.tz_prev + time.tzset() + super().tearDownClass() + + def create_pandas_data_frame(self): + import numpy as np + + data_dict = {} + for j, name in enumerate(self.schema.names): + data_dict[name] = [self.data[i][j] for i in range(len(self.data))] + # need to convert these to numpy types first + data_dict["2_int_t"] = np.int32(data_dict["2_int_t"]) + data_dict["4_float_t"] = np.float32(data_dict["4_float_t"]) + return pd.DataFrame(data=data_dict) + + @property + def create_np_arrs(self): + import numpy as np + + int_dtypes = ["int8", "int16", "int32", "int64"] + float_dtypes = ["float32", "float64"] + return ( + [np.array([1, 2]).astype(t) for t in int_dtypes] + + [np.array([0.1, 0.2]).astype(t) for t in float_dtypes] + + [np.array([[1], [2]]).astype(t) for t in int_dtypes] + + [np.array([[0.1], [0.2]]).astype(t) for t in float_dtypes] + + [np.array([[1, 1, 1], [2, 2, 2]]).astype(t) for t in int_dtypes] + + [np.array([[0.1, 0.1, 0.1], [0.2, 0.2, 0.2]]).astype(t) for t in float_dtypes] + ) + + @unittest.skipIf( + not have_pyarrow or LooseVersion(pa.__version__) >= "2.0", + "will not fallback with pyarrow>=2.0", + ) + def test_toPandas_fallback_enabled(self): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.fallback.enabled": True}): + schema = StructType([StructField("a", ArrayType(StructType()), True)]) + df = self.spark.createDataFrame([([Row()],)], schema=schema) + with QuietTest(self.sc): + with self.warnings_lock: + with warnings.catch_warnings(record=True) as warns: + # we want the warnings to appear even if this test is run from a subclass + warnings.simplefilter("always") + pdf = df.toPandas() + # Catch and check the last UserWarning. + user_warns = [ + warn.message for warn in warns if isinstance(warn.message, UserWarning) + ] + self.assertTrue(len(user_warns) > 0) + self.assertTrue("Attempting non-optimization" in str(user_warns[-1])) + assert_frame_equal(pdf, pd.DataFrame({"a": [[Row()]]})) + + @unittest.skipIf( + not have_pyarrow or LooseVersion(pa.__version__) >= "2.0", + "will not fallback with pyarrow>=2.0", + ) + def test_toPandas_fallback_disabled(self): + schema = StructType([StructField("a", ArrayType(StructType()), True)]) + df = self.spark.createDataFrame([(None,)], schema=schema) + with QuietTest(self.sc): + with self.warnings_lock: + with self.assertRaises(PySparkTypeError) as pe: + df.toPandas() + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION", + message_parameters={"data_type": "Array of StructType"}, + ) + + def test_toPandas_empty_df_arrow_enabled(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_empty_df_arrow_enabled(arrow_enabled) + + def check_toPandas_empty_df_arrow_enabled(self, arrow_enabled): + # SPARK-30537 test that toPandas() on an empty dataframe has the correct dtypes + # when arrow is enabled + from datetime import date + from decimal import Decimal + + schema = StructType( + [ + StructField("a", StringType(), True), + StructField("a", IntegerType(), True), + StructField("c", TimestampType(), True), + StructField("d", NullType(), True), + StructField("e", LongType(), True), + StructField("f", FloatType(), True), + StructField("g", DateType(), True), + StructField("h", BinaryType(), True), + StructField("i", DecimalType(38, 18), True), + StructField("k", TimestampNTZType(), True), + StructField("L", DayTimeIntervalType(0, 3), True), + ] + ) + df = self.spark.createDataFrame([], schema=schema) + non_empty_df = self.spark.createDataFrame( + [ + ( + "a", + 1, + datetime.datetime(1969, 1, 1, 1, 1, 1), + None, + 10, + 0.2, + date(1969, 1, 1), + bytearray(b"a"), + Decimal("2.0"), + datetime.datetime(1969, 1, 1, 1, 1, 1), + datetime.timedelta(microseconds=123), + ) + ], + schema=schema, + ) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + pdf = df.toPandas() + pdf_non_empty = non_empty_df.toPandas() + self.assertTrue(pdf.dtypes.equals(pdf_non_empty.dtypes)) + + def test_null_conversion(self): + df_null = self.spark.createDataFrame( + [tuple([None for _ in range(len(self.data_wo_null[0]))])] + self.data_wo_null + ) + pdf = df_null.toPandas() + null_counts = pdf.isnull().sum().tolist() + self.assertTrue(all([c == 1 for c in null_counts])) + + def _toPandas_arrow_toggle(self, df): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + pdf = df.toPandas() + + pdf_arrow = df.toPandas() + + return pdf, pdf_arrow + + def test_toPandas_arrow_toggle(self): + df = self.spark.createDataFrame(self.data, schema=self.schema) + pdf, pdf_arrow = self._toPandas_arrow_toggle(df) + expected = self.create_pandas_data_frame() + assert_frame_equal(expected, pdf) + assert_frame_equal(expected, pdf_arrow) + + def test_create_data_frame_to_pandas_timestamp_ntz(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_create_data_frame_to_pandas_timestamp_ntz(arrow_enabled) + + def check_create_data_frame_to_pandas_timestamp_ntz(self, arrow_enabled): + # SPARK-36626: Test TimestampNTZ in createDataFrame and toPandas + with self.sql_conf({"spark.sql.session.timeZone": "America/Los_Angeles"}): + origin = pd.DataFrame({"a": [datetime.datetime(2012, 2, 2, 2, 2, 2)]}) + df = self.spark.createDataFrame( + origin, schema=StructType([StructField("a", TimestampNTZType(), True)]) + ) + df.selectExpr("assert_true('2012-02-02 02:02:02' == CAST(a AS STRING))").collect() + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + pdf = df.toPandas() + assert_frame_equal(origin, pdf) + + def test_create_data_frame_to_pandas_day_time_internal(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_create_data_frame_to_pandas_day_time_internal(arrow_enabled) + + def check_create_data_frame_to_pandas_day_time_internal(self, arrow_enabled): + # SPARK-37279: Test DayTimeInterval in createDataFrame and toPandas + origin = pd.DataFrame({"a": [datetime.timedelta(microseconds=123)]}) + df = self.spark.createDataFrame(origin) + df.select( + assert_true(lit("INTERVAL '0 00:00:00.000123' DAY TO SECOND") == df.a.cast("string")) + ).collect() + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + pdf = df.toPandas() + assert_frame_equal(origin, pdf) + + def test_toPandas_respect_session_timezone(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_respect_session_timezone(arrow_enabled) + + def check_toPandas_respect_session_timezone(self, arrow_enabled): + df = self.spark.createDataFrame(self.data, schema=self.schema) + + timezone = "America/Los_Angeles" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + pdf_la = df.toPandas() + + timezone = "America/New_York" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + pdf_ny = df.toPandas() + + self.assertFalse(pdf_ny.equals(pdf_la)) + + from pyspark.sql.pandas.types import _check_series_convert_timestamps_local_tz + + pdf_la_corrected = pdf_la.copy() + for field in self.schema: + if isinstance(field.dataType, TimestampType): + pdf_la_corrected[field.name] = _check_series_convert_timestamps_local_tz( + pdf_la_corrected[field.name], timezone + ) + assert_frame_equal(pdf_ny, pdf_la_corrected) + + def test_pandas_round_trip(self): + pdf = self.create_pandas_data_frame() + df = self.spark.createDataFrame(self.data, schema=self.schema) + pdf_arrow = df.toPandas() + assert_frame_equal(pdf_arrow, pdf) + + def test_pandas_self_destruct(self): + import pyarrow as pa + + rows = 2**10 + cols = 4 + expected_bytes = rows * cols * 8 + df = self.spark.range(0, rows).select(*[rand() for _ in range(cols)]) + # Test the self_destruct behavior by testing _collect_as_arrow directly + allocation_before = pa.total_allocated_bytes() + batches = df._collect_as_arrow(split_batches=True) + table = pa.Table.from_batches(batches) + del batches + pdf_split = table.to_pandas(self_destruct=True, split_blocks=True, use_threads=False) + allocation_after = pa.total_allocated_bytes() + difference = allocation_after - allocation_before + # Should be around 1x the data size (table should not hold on to any memory) + self.assertGreaterEqual(difference, 0.9 * expected_bytes) + self.assertLessEqual(difference, 1.1 * expected_bytes) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": False}): + no_self_destruct_pdf = df.toPandas() + # Note while memory usage is 2x data size here (both table and pdf hold on to + # memory), in this case Arrow still only tracks 1x worth of memory (since the + # batches are not allocated by Arrow in this case), so we can't make any + # assertions here + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": True}): + self_destruct_pdf = df.toPandas() + + assert_frame_equal(pdf_split, no_self_destruct_pdf) + assert_frame_equal(pdf_split, self_destruct_pdf) + + def test_filtered_frame(self): + df = self.spark.range(3).toDF("i") + pdf = df.filter("i < 0").toPandas() + self.assertEqual(len(pdf.columns), 1) + self.assertEqual(pdf.columns[0], "i") + self.assertTrue(pdf.empty) + + def test_no_partition_frame(self): + schema = StructType([StructField("field1", StringType(), True)]) + df = self.spark.createDataFrame(self.sc.emptyRDD(), schema) + pdf = df.toPandas() + self.assertEqual(len(pdf.columns), 1) + self.assertEqual(pdf.columns[0], "field1") + self.assertTrue(pdf.empty) + + def test_propagates_spark_exception(self): + with QuietTest(self.sc): + self.check_propagates_spark_exception() + + def check_propagates_spark_exception(self): + df = self.spark.range(3).toDF("i") + + def raise_exception(): + raise RuntimeError("My error") + + exception_udf = udf(raise_exception, IntegerType()) + df = df.withColumn("error", exception_udf()) + + with self.assertRaisesRegex(Exception, "My error"): + df.toPandas() + + def _createDataFrame_toggle(self, data, schema=None): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + df_no_arrow = self.spark.createDataFrame(data, schema=schema) + + df_arrow = self.spark.createDataFrame(data, schema=schema) + + return df_no_arrow, df_arrow + + def test_createDataFrame_toggle(self): + pdf = self.create_pandas_data_frame() + df_no_arrow, df_arrow = self._createDataFrame_toggle(pdf, schema=self.schema) + self.assertEqual(df_no_arrow.collect(), df_arrow.collect()) + + def test_createDataFrame_respect_session_timezone(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_respect_session_timezone(arrow_enabled) + + def check_createDataFrame_respect_session_timezone(self, arrow_enabled): + from datetime import timedelta + + pdf = self.create_pandas_data_frame() + timezone = "America/Los_Angeles" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df_la = self.spark.createDataFrame(pdf, schema=self.schema) + result_la = df_la.collect() + + timezone = "America/New_York" + with self.sql_conf({"spark.sql.session.timeZone": timezone}): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df_ny = self.spark.createDataFrame(pdf, schema=self.schema) + result_ny = df_ny.collect() + + self.assertNotEqual(result_ny, result_la) + + # Correct result_la by adjusting 3 hours difference between Los Angeles and New York + result_la_corrected = [ + Row( + **{ + k: v - timedelta(hours=3) if k == "8_timestamp_t" else v + for k, v in row.asDict().items() + } + ) + for row in result_la + ] + self.assertEqual(result_ny, result_la_corrected) + + def test_createDataFrame_with_schema(self): + pdf = self.create_pandas_data_frame() + df = self.spark.createDataFrame(pdf, schema=self.schema) + self.assertEqual(self.schema, df.schema) + pdf_arrow = df.toPandas() + assert_frame_equal(pdf_arrow, pdf) + + def test_createDataFrame_with_incorrect_schema(self): + with QuietTest(self.sc): + self.check_createDataFrame_with_incorrect_schema() + + def check_createDataFrame_with_incorrect_schema(self): + pdf = self.create_pandas_data_frame() + fields = list(self.schema) + fields[5], fields[6] = fields[6], fields[5] # swap decimal with date + wrong_schema = StructType(fields) + with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}): + with self.assertRaises(Exception) as context: + self.spark.createDataFrame(pdf, schema=wrong_schema) + + # the exception provides us with the column that is incorrect + exception = context.exception + self.assertTrue(hasattr(exception, "args")) + self.assertEqual(len(exception.args), 1) + self.assertRegex( + exception.args[0], + "with name '7_date_t' " "to Arrow Array \\(decimal128\\(38, 18\\)\\)", + ) + + # the inner exception provides us with the incorrect types + exception = exception.__context__ + self.assertTrue(hasattr(exception, "args")) + self.assertEqual(len(exception.args), 1) + self.assertRegex(exception.args[0], "[D|d]ecimal.*got.*date") + + def test_createDataFrame_with_names(self): + pdf = self.create_pandas_data_frame() + new_names = list(map(str, range(len(self.schema.fieldNames())))) + # Test that schema as a list of column names gets applied + df = self.spark.createDataFrame(pdf, schema=list(new_names)) + self.assertEqual(df.schema.fieldNames(), new_names) + # Test that schema as tuple of column names gets applied + df = self.spark.createDataFrame(pdf, schema=tuple(new_names)) + self.assertEqual(df.schema.fieldNames(), new_names) + + def test_createDataFrame_column_name_encoding(self): + pdf = pd.DataFrame({"a": [1]}) + columns = self.spark.createDataFrame(pdf).columns + self.assertTrue(isinstance(columns[0], str)) + self.assertEqual(columns[0], "a") + columns = self.spark.createDataFrame(pdf, ["b"]).columns + self.assertTrue(isinstance(columns[0], str)) + self.assertEqual(columns[0], "b") + + def test_createDataFrame_with_single_data_type(self): + with QuietTest(self.sc): + self.check_createDataFrame_with_single_data_type() + + def check_createDataFrame_with_single_data_type(self): + for schema in ["int", IntegerType()]: + with self.subTest(schema=schema): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame(pd.DataFrame({"a": [1]}), schema=schema).collect() + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW", + message_parameters={"data_type": "IntegerType()"}, + ) + + def test_createDataFrame_does_not_modify_input(self): + # Some series get converted for Spark to consume, this makes sure input is unchanged + pdf = self.create_pandas_data_frame() + # Use a nanosecond value to make sure it is not truncated + pdf.iloc[0, 7] = pd.Timestamp(1) + # Integers with nulls will get NaNs filled with 0 and will be casted + pdf.iloc[1, 1] = None + pdf_copy = pdf.copy(deep=True) + self.spark.createDataFrame(pdf, schema=self.schema) + self.assertTrue(pdf.equals(pdf_copy)) + + def test_schema_conversion_roundtrip(self): + from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema + + arrow_schema = to_arrow_schema(self.schema) + schema_rt = from_arrow_schema(arrow_schema, prefer_timestamp_ntz=True) + self.assertEqual(self.schema, schema_rt) + + def test_createDataFrame_with_ndarray(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_with_ndarray(arrow_enabled) + + def check_createDataFrame_with_ndarray(self, arrow_enabled): + import numpy as np + + dtypes = ["tinyint", "smallint", "int", "bigint", "float", "double"] + expected_dtypes = ( + [[("value", t)] for t in dtypes] + + [[("value", t)] for t in dtypes] + + [[("_1", t), ("_2", t), ("_3", t)] for t in dtypes] + ) + arrs = self.create_np_arrs + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + for arr, dtypes in zip(arrs, expected_dtypes): + df = self.spark.createDataFrame(arr) + self.assertEqual(df.dtypes, dtypes) + np.array_equal(np.array(df.collect()), arr) + + with self.assertRaisesRegex( + ValueError, "NumPy array input should be of 1 or 2 dimensions" + ): + self.spark.createDataFrame(np.array(0)) + + def test_createDataFrame_with_array_type(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_with_array_type(arrow_enabled) + + def check_createDataFrame_with_array_type(self, arrow_enabled): + pdf = pd.DataFrame({"a": [[1, 2], [3, 4]], "b": [["x", "y"], ["y", "z"]]}) + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf) + result = df.collect() + expected = [tuple(list(e) for e in rec) for rec in pdf.to_records(index=False)] + for r in range(len(expected)): + for e in range(len(expected[r])): + self.assertTrue(expected[r][e] == result[r][e]) + + def test_toPandas_with_array_type(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_with_array_type(arrow_enabled) + + def check_toPandas_with_array_type(self, arrow_enabled): + expected = [([1, 2], ["x", "y"]), ([3, 4], ["y", "z"])] + array_schema = StructType( + [StructField("a", ArrayType(IntegerType())), StructField("b", ArrayType(StringType()))] + ) + df = self.spark.createDataFrame(expected, schema=array_schema) + pdf = df.toPandas() + result = [tuple(list(e) for e in rec) for rec in pdf.to_records(index=False)] + for r in range(len(expected)): + for e in range(len(expected[r])): + self.assertTrue(expected[r][e] == result[r][e]) + + def test_createDataFrame_with_map_type(self): + with QuietTest(self.sc): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_with_map_type(arrow_enabled) + + def check_createDataFrame_with_map_type(self, arrow_enabled): + map_data = [{"a": 1}, {"b": 2, "c": 3}, {}, None, {"d": None}] + + pdf = pd.DataFrame({"id": [0, 1, 2, 3, 4], "m": map_data}) + for schema in ( + "id long, m map", + StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())), + ): + with self.subTest(schema=schema): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + if arrow_enabled and LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + with self.assertRaisesRegex(Exception, "MapType.*only.*pyarrow 2.0.0"): + self.spark.createDataFrame(pdf, schema=schema).collect() + else: + df = self.spark.createDataFrame(pdf, schema=schema) + + result = df.collect() + + for row in result: + i, m = row + self.assertEqual(m, map_data[i]) + + def test_createDataFrame_with_struct_type(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_with_struct_type(arrow_enabled) + + def check_createDataFrame_with_struct_type(self, arrow_enabled): + pdf = pd.DataFrame( + {"a": [Row(1, "a"), Row(2, "b")], "b": [{"s": 3, "t": "x"}, {"s": 4, "t": "y"}]} + ) + for schema in ( + "a struct, b struct", + StructType() + .add("a", StructType().add("x", LongType()).add("y", StringType())) + .add("b", StructType().add("s", LongType()).add("t", StringType())), + ): + with self.subTest(schema=schema): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf, schema) + result = df.collect() + expected = [(rec[0], Row(**rec[1])) for rec in pdf.to_records(index=False)] + for r in range(len(expected)): + for e in range(len(expected[r])): + self.assertTrue( + expected[r][e] == result[r][e], f"{expected[r][e]} == {result[r][e]}" + ) + + def test_createDataFrame_with_string_dtype(self): + # SPARK-34521: spark.createDataFrame does not support Pandas StringDtype extension type + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): + data = [["abc"], ["def"], [None], ["ghi"], [None]] + pandas_df = pd.DataFrame(data, columns=["col"], dtype="string") + schema = StructType([StructField("col", StringType(), True)]) + df = self.spark.createDataFrame(pandas_df, schema=schema) + + # dtypes won't match. Pandas has two different ways to store string columns: + # using ndarray (when dtype isn't specified) or using a StringArray when dtype="string". + # When calling dataframe#toPandas() it will use the ndarray version. + # Changing that to use a StringArray would be backwards incompatible. + assert_frame_equal(pandas_df, df.toPandas(), check_dtype=False) + + def test_createDataFrame_with_int64(self): + # SPARK-34521: spark.createDataFrame does not support Pandas StringDtype extension type + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): + pandas_df = pd.DataFrame({"col": [1, 2, 3, None]}, dtype="Int64") + df = self.spark.createDataFrame(pandas_df) + assert_frame_equal(pandas_df, df.toPandas(), check_dtype=False) + + def test_toPandas_with_map_type(self): + with QuietTest(self.sc): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_with_map_type(arrow_enabled) + + def check_toPandas_with_map_type(self, arrow_enabled): + origin = pd.DataFrame( + {"id": [0, 1, 2, 3], "m": [{}, {"a": 1}, {"a": 1, "b": 2}, {"a": 1, "b": 2, "c": 3}]} + ) + + for schema in [ + "id long, m map", + StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())), + ]: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + df = self.spark.createDataFrame(origin, schema=schema) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + if arrow_enabled and LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + with self.assertRaisesRegex(Exception, "MapType.*only.*pyarrow 2.0.0"): + df.toPandas() + else: + pdf = df.toPandas() + assert_frame_equal(origin, pdf) + + def test_toPandas_with_map_type_nulls(self): + with QuietTest(self.sc): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_with_map_type_nulls(arrow_enabled) + + def check_toPandas_with_map_type_nulls(self, arrow_enabled): + origin = pd.DataFrame( + {"id": [0, 1, 2, 3, 4], "m": [{"a": 1}, {"b": 2, "c": 3}, {}, None, {"d": None}]} + ) + + for schema in [ + "id long, m map", + StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())), + ]: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + df = self.spark.createDataFrame(origin, schema=schema) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + if arrow_enabled and LooseVersion(pa.__version__) < LooseVersion("2.0.0"): + with self.assertRaisesRegex(Exception, "MapType.*only.*pyarrow 2.0.0"): + df.toPandas() + else: + pdf = df.toPandas() + assert_frame_equal(origin, pdf) + + def test_createDataFrame_with_int_col_names(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_with_int_col_names(arrow_enabled) + + def check_createDataFrame_with_int_col_names(self, arrow_enabled): + import numpy as np + + pdf = pd.DataFrame(np.random.rand(4, 2)) + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf) + pdf_col_names = [str(c) for c in pdf.columns] + self.assertEqual(pdf_col_names, df.columns) + + @unittest.skipIf( + not have_pyarrow or LooseVersion(pa.__version__) >= "2.0", + "will not fallback with pyarrow>=2.0", + ) + def test_createDataFrame_fallback_enabled(self): + with QuietTest(self.sc): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.fallback.enabled": True}): + with warnings.catch_warnings(record=True) as warns: + # we want the warnings to appear even if this test is run from a subclass + warnings.simplefilter("always") + df = self.spark.createDataFrame( + pd.DataFrame({"a": [[Row()]]}), "a: array>" + ) + # Catch and check the last UserWarning. + user_warns = [ + warn.message for warn in warns if isinstance(warn.message, UserWarning) + ] + self.assertTrue(len(user_warns) > 0) + self.assertTrue("Attempting non-optimization" in str(user_warns[-1])) + self.assertEqual(df.collect(), [Row(a=[Row()])]) + + @unittest.skipIf( + not have_pyarrow or LooseVersion(pa.__version__) >= "2.0", + "will not fallback with pyarrow>=2.0", + ) + def test_createDataFrame_fallback_disabled(self): + with QuietTest(self.sc): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame(pd.DataFrame({"a": [[Row()]]}), "a: array>") + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW_VERSION", + message_parameters={"data_type": "Array of StructType"}, + ) + + # Regression test for SPARK-23314 + def test_timestamp_dst(self): + # Daylight saving time for Los Angeles for 2015 is Sun, Nov 1 at 2:00 am + dt = [ + datetime.datetime(2015, 11, 1, 0, 30), + datetime.datetime(2015, 11, 1, 1, 30), + datetime.datetime(2015, 11, 1, 2, 30), + ] + pdf = pd.DataFrame({"time": dt}) + + df_from_python = self.spark.createDataFrame(dt, "timestamp").toDF("time") + df_from_pandas = self.spark.createDataFrame(pdf) + + assert_frame_equal(pdf, df_from_python.toPandas()) + assert_frame_equal(pdf, df_from_pandas.toPandas()) + + # Regression test for SPARK-28003 + def test_timestamp_nat(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_timestamp_nat(arrow_enabled) + + def check_timestamp_nat(self, arrow_enabled): + dt = [pd.NaT, pd.Timestamp("2019-06-11"), None] * 100 + pdf = pd.DataFrame({"time": dt}) + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf) + + assert_frame_equal(pdf, df.toPandas()) + + def test_toPandas_batch_order(self): + def delay_first_part(partition_index, iterator): + if partition_index == 0: + time.sleep(0.1) + return iterator + + # Collects Arrow RecordBatches out of order in driver JVM then re-orders in Python + def run_test(num_records, num_parts, max_records, use_delay=False): + df = self.spark.range(num_records, numPartitions=num_parts).toDF("a") + if use_delay: + df = df.rdd.mapPartitionsWithIndex(delay_first_part).toDF() + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": max_records}): + pdf, pdf_arrow = self._toPandas_arrow_toggle(df) + assert_frame_equal(pdf, pdf_arrow) + + cases = [ + (1024, 512, 2), # Use large num partitions for more likely collecting out of order + (64, 8, 2, True), # Use delay in first partition to force collecting out of order + (64, 64, 1), # Test single batch per partition + (64, 1, 64), # Test single partition, single batch + (64, 1, 8), # Test single partition, multiple batches + (30, 7, 2), # Test different sized partitions + ] + + for case in cases: + run_test(*case) + + def test_createDataFrame_with_category_type(self): + pdf = pd.DataFrame({"A": ["a", "b", "c", "a"]}) + pdf["B"] = pdf["A"].astype("category") + category_first_element = dict(enumerate(pdf["B"].cat.categories))[0] + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): + arrow_df = self.spark.createDataFrame(pdf) + arrow_type = arrow_df.dtypes[1][1] + result_arrow = arrow_df.toPandas() + arrow_first_category_element = result_arrow["B"][0] + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): + df = self.spark.createDataFrame(pdf) + spark_type = df.dtypes[1][1] + result_spark = df.toPandas() + spark_first_category_element = result_spark["B"][0] + + assert_frame_equal(result_spark, result_arrow) + + # ensure original category elements are string + self.assertIsInstance(category_first_element, str) + # spark data frame and arrow execution mode enabled data frame type must match pandas + self.assertEqual(spark_type, "string") + self.assertEqual(arrow_type, "string") + self.assertIsInstance(arrow_first_category_element, str) + self.assertIsInstance(spark_first_category_element, str) + + def test_createDataFrame_with_float_index(self): + # SPARK-32098: float index should not produce duplicated or truncated Spark DataFrame + self.assertEqual( + self.spark.createDataFrame(pd.DataFrame({"a": [1, 2, 3]}, index=[2.0, 3.0, 4.0])) + .distinct() + .count(), + 3, + ) + + def test_no_partition_toPandas(self): + # SPARK-32301: toPandas should work from a Spark DataFrame with no partitions + # Forward-ported from SPARK-32300. + pdf = self.spark.sparkContext.emptyRDD().toDF("col1 int").toPandas() + self.assertEqual(len(pdf), 0) + self.assertEqual(list(pdf.columns), ["col1"]) + + def test_createDataFrame_empty_partition(self): + pdf = pd.DataFrame({"c1": [1], "c2": ["string"]}) + df = self.spark.createDataFrame(pdf) + self.assertEqual([Row(c1=1, c2="string")], df.collect()) + self.assertGreater(self.spark.sparkContext.defaultParallelism, len(pdf)) + + def test_toPandas_error(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_error(arrow_enabled) + + def check_toPandas_error(self, arrow_enabled): + with self.sql_conf( + { + "spark.sql.ansi.enabled": True, + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + } + ): + with self.assertRaises(ArithmeticException): + self.spark.sql("select 1/0").toPandas() + + def test_toPandas_duplicate_field_names(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_duplicate_field_names(arrow_enabled) + + def check_toPandas_duplicate_field_names(self, arrow_enabled): + data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))] + schema = ( + StructType() + .add("struct", StructType().add("x", StringType()).add("x", IntegerType())) + .add( + "struct", + StructType() + .add("a", IntegerType()) + .add("x", IntegerType()) + .add("x", StringType()) + .add("y", IntegerType()) + .add("y", StringType()), + ) + ) + for struct_in_pandas in ["legacy", "row", "dict"]: + df = self.spark.createDataFrame(data, schema=schema) + + with self.subTest(struct_in_pandas=struct_in_pandas): + with self.sql_conf( + { + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + "spark.sql.execution.pandas.structHandlingMode": struct_in_pandas, + } + ): + if arrow_enabled and struct_in_pandas == "legacy": + with self.assertRaisesRegexp( + UnsupportedOperationException, "DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT" + ): + df.toPandas() + else: + if struct_in_pandas == "dict": + expected = pd.DataFrame( + [ + [ + {"x_0": "a", "x_1": 1}, + {"a": 2, "x_0": 3, "x_1": "b", "y_0": 4, "y_1": "c"}, + ], + [ + {"x_0": "x", "x_1": 6}, + {"a": 7, "x_0": 8, "x_1": "y", "y_0": 9, "y_1": "z"}, + ], + ], + columns=schema.names, + ) + else: + expected = pd.DataFrame.from_records(data, columns=schema.names) + assert_frame_equal(df.toPandas(), expected) + + def test_createDataFrame_duplicate_field_names(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_duplicate_field_names(arrow_enabled) + + def check_createDataFrame_duplicate_field_names(self, arrow_enabled): + schema = ( + StructType() + .add("struct", StructType().add("x", StringType()).add("x", IntegerType())) + .add( + "struct", + StructType() + .add("a", IntegerType()) + .add("x", IntegerType()) + .add("x", StringType()) + .add("y", IntegerType()) + .add("y", StringType()), + ) + ) + + data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))] + pdf = pd.DataFrame.from_records(data, columns=schema.names) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf, schema) + + self.assertEqual(df.collect(), data) + + @unittest.skipIf( + LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), + "TODO(SPARK-43506): Enable ArrowTests.test_toPandas_empty_columns for pandas 2.0.0.", + ) + def test_toPandas_empty_columns(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_empty_columns(arrow_enabled) + + def check_toPandas_empty_columns(self, arrow_enabled): + df = self.spark.range(2).select([]) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + assert_frame_equal(df.toPandas(), pd.DataFrame(columns=[], index=range(2))) + + def test_createDataFrame_nested_timestamp(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_nested_timestamp(arrow_enabled) + + def check_createDataFrame_nested_timestamp(self, arrow_enabled): + schema = ( + StructType() + .add("ts", TimestampType()) + .add("ts_ntz", TimestampNTZType()) + .add( + "struct", StructType().add("ts", TimestampType()).add("ts_ntz", TimestampNTZType()) + ) + .add("array", ArrayType(TimestampType())) + .add("array_ntz", ArrayType(TimestampNTZType())) + .add("map", MapType(StringType(), TimestampType())) + .add("map_ntz", MapType(StringType(), TimestampNTZType())) + ) + data = [ + Row( + datetime.datetime(2023, 1, 1, 0, 0, 0), + datetime.datetime(2023, 1, 1, 0, 0, 0), + Row( + datetime.datetime(2023, 1, 1, 0, 0, 0), + datetime.datetime(2023, 1, 1, 0, 0, 0), + ), + [datetime.datetime(2023, 1, 1, 0, 0, 0)], + [datetime.datetime(2023, 1, 1, 0, 0, 0)], + dict(ts=datetime.datetime(2023, 1, 1, 0, 0, 0)), + dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)), + ) + ] + pdf = pd.DataFrame.from_records(data, columns=schema.names) + + with self.sql_conf( + { + "spark.sql.session.timeZone": "America/New_York", + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + } + ): + df = self.spark.createDataFrame(pdf, schema) + + expected = Row( + ts=datetime.datetime(2022, 12, 31, 21, 0, 0), + ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0), + struct=Row( + ts=datetime.datetime(2022, 12, 31, 21, 0, 0), + ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0), + ), + array=[datetime.datetime(2022, 12, 31, 21, 0, 0)], + array_ntz=[datetime.datetime(2023, 1, 1, 0, 0, 0)], + map=dict(ts=datetime.datetime(2022, 12, 31, 21, 0, 0)), + map_ntz=dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)), + ) + + self.assertEqual(df.first(), expected) + + def test_toPandas_nested_timestamp(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_nested_timestamp(arrow_enabled) + + def check_toPandas_nested_timestamp(self, arrow_enabled): + schema = ( + StructType() + .add("ts", TimestampType()) + .add("ts_ntz", TimestampNTZType()) + .add( + "struct", StructType().add("ts", TimestampType()).add("ts_ntz", TimestampNTZType()) + ) + .add("array", ArrayType(TimestampType())) + .add("array_ntz", ArrayType(TimestampNTZType())) + .add("map", MapType(StringType(), TimestampType())) + .add("map_ntz", MapType(StringType(), TimestampNTZType())) + ) + data = [ + Row( + datetime.datetime(2023, 1, 1, 0, 0, 0), + datetime.datetime(2023, 1, 1, 0, 0, 0), + Row( + datetime.datetime(2023, 1, 1, 0, 0, 0), + datetime.datetime(2023, 1, 1, 0, 0, 0), + ), + [datetime.datetime(2023, 1, 1, 0, 0, 0)], + [datetime.datetime(2023, 1, 1, 0, 0, 0)], + dict(ts=datetime.datetime(2023, 1, 1, 0, 0, 0)), + dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)), + ) + ] + df = self.spark.createDataFrame(data, schema) + + with self.sql_conf( + { + "spark.sql.session.timeZone": "America/New_York", + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + "spark.sql.execution.pandas.structHandlingMode": "row", + } + ): + pdf = df.toPandas() + + expected = pd.DataFrame( + { + "ts": [datetime.datetime(2023, 1, 1, 3, 0, 0)], + "ts_ntz": [datetime.datetime(2023, 1, 1, 0, 0, 0)], + "struct": [ + Row( + datetime.datetime(2023, 1, 1, 3, 0, 0), + datetime.datetime(2023, 1, 1, 0, 0, 0), + ) + ], + "array": [[datetime.datetime(2023, 1, 1, 3, 0, 0)]], + "array_ntz": [[datetime.datetime(2023, 1, 1, 0, 0, 0)]], + "map": [dict(ts=datetime.datetime(2023, 1, 1, 3, 0, 0))], + "map_ntz": [dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0))], + } + ) + + assert_frame_equal(pdf, expected) + + def test_createDataFrame_udt(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_createDataFrame_udt(arrow_enabled) + + def check_createDataFrame_udt(self, arrow_enabled): + schema = ( + StructType() + .add("point", ExamplePointUDT()) + .add("struct", StructType().add("point", ExamplePointUDT())) + .add("array", ArrayType(ExamplePointUDT())) + .add("map", MapType(StringType(), ExamplePointUDT())) + ) + data = [ + Row( + ExamplePoint(1.0, 2.0), + Row(ExamplePoint(3.0, 4.0)), + [ExamplePoint(5.0, 6.0)], + dict(point=ExamplePoint(7.0, 8.0)), + ) + ] + pdf = pd.DataFrame.from_records(data, columns=schema.names) + + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + df = self.spark.createDataFrame(pdf, schema) + + self.assertEqual(df.collect(), data) + + def test_toPandas_udt(self): + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_toPandas_udt(arrow_enabled) + + def check_toPandas_udt(self, arrow_enabled): + schema = ( + StructType() + .add("point", ExamplePointUDT()) + .add("struct", StructType().add("point", ExamplePointUDT())) + .add("array", ArrayType(ExamplePointUDT())) + .add("map", MapType(StringType(), ExamplePointUDT())) + ) + data = [ + Row( + ExamplePoint(1.0, 2.0), + Row(ExamplePoint(3.0, 4.0)), + [ExamplePoint(5.0, 6.0)], + dict(point=ExamplePoint(7.0, 8.0)), + ) + ] + df = self.spark.createDataFrame(data, schema) + + with self.sql_conf( + { + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + "spark.sql.execution.pandas.structHandlingMode": "row", + } + ): + pdf = df.toPandas() + + expected = pd.DataFrame.from_records(data, columns=schema.names) + + assert_frame_equal(pdf, expected) + + def test_create_dataframe_namedtuples(self): + # SPARK-44980: Inherited namedtuples in createDataFrame + for arrow_enabled in [True, False]: + with self.subTest(arrow_enabled=arrow_enabled): + self.check_create_dataframe_namedtuples(arrow_enabled) + + def check_create_dataframe_namedtuples(self, arrow_enabled): + MyTuple = namedtuple("MyTuple", ["a", "b", "c"]) + + class MyInheritedTuple(MyTuple): + pass + + with self.sql_conf( + { + "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled, + } + ): + df = self.spark.createDataFrame([MyInheritedTuple(1, 2, 3)]) + self.assertEqual(df.first(), Row(a=1, b=2, c=3)) + + df = self.spark.createDataFrame([MyInheritedTuple(1, 2, MyInheritedTuple(1, 2, 3))]) + self.assertEqual(df.first(), Row(a=1, b=2, c=Row(a=1, b=2, c=3))) + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class ArrowTests(ArrowTestsMixin, ReusedSQLTestCase): + pass + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class MaxResultArrowTests(unittest.TestCase): + # These tests are separate as 'spark.driver.maxResultSize' configuration + # is a static configuration to Spark context. + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession( + SparkContext( + "local[4]", cls.__name__, conf=SparkConf().set("spark.driver.maxResultSize", "10k") + ) + ) + + # Explicitly enable Arrow and disable fallback. + cls.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true") + cls.spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "false") + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "spark"): + cls.spark.stop() + + def test_exception_by_max_results(self): + with self.assertRaisesRegex(Exception, "is bigger than"): + self.spark.range(0, 10000, 1, 100).toPandas() + + +class EncryptionArrowTests(ArrowTests): + @classmethod + def conf(cls): + return super(EncryptionArrowTests, cls).conf().set("spark.io.encryption.enabled", "true") + + +class RDDBasedArrowTests(ArrowTests): + @classmethod + def conf(cls): + return ( + super(RDDBasedArrowTests, cls) + .conf() + .set("spark.sql.execution.arrow.localRelationThreshold", "0") + # to test multiple partitions + .set("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + ) + + +if __name__ == "__main__": + from pyspark.sql.tests.test_arrow import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_map.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_map.py new file mode 100644 index 0000000..1536774 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_map.py @@ -0,0 +1,182 @@ +# +# 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 os +import time +import unittest + +from pyspark.sql.utils import PythonException +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + +if have_pyarrow: + import pyarrow as pa + +if have_pandas: + import pandas as pd + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + pandas_requirement_message or pyarrow_requirement_message, +) +class MapInArrowTestsMixin(object): + def test_map_in_arrow(self): + def func(iterator): + for batch in iterator: + assert isinstance(batch, pa.RecordBatch) + assert batch.schema.names == ["id"] + yield batch + + df = self.spark.range(10) + actual = df.mapInArrow(func, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_multiple_columns(self): + data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] + df = self.spark.createDataFrame(data, "a int, b string") + + def func(iterator): + for batch in iterator: + assert isinstance(batch, pa.RecordBatch) + assert batch.schema.types == [pa.int32(), pa.string()] + yield batch + + actual = df.mapInArrow(func, df.schema).collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_large_variable_width_types(self): + with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes": True}): + data = [("foo", b"foo"), (None, None), ("bar", b"bar")] + df = self.spark.createDataFrame(data, "a string, b binary") + + def func(iterator): + for batch in iterator: + assert isinstance(batch, pa.RecordBatch) + assert batch.schema.types == [pa.large_string(), pa.large_binary()] + yield batch + + actual = df.mapInArrow(func, df.schema).collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_different_output_length(self): + def func(iterator): + for _ in iterator: + yield pa.RecordBatch.from_pandas(pd.DataFrame({"a": list(range(100))})) + + df = self.spark.range(10) + actual = df.repartition(1).mapInArrow(func, "a long").collect() + self.assertEqual(set((r.a for r in actual)), set(range(100))) + + def test_other_than_recordbatch_iter(self): + with QuietTest(self.sc): + self.check_other_than_recordbatch_iter() + + def check_other_than_recordbatch_iter(self): + def not_iter(_): + return 1 + + def bad_iter_elem(_): + return iter([1]) + + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be iterator " + "of pyarrow.RecordBatch, but is int.", + ): + (self.spark.range(10, numPartitions=3).mapInArrow(not_iter, "a int").count()) + + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be iterator " + "of pyarrow.RecordBatch, but is iterator of int.", + ): + (self.spark.range(10, numPartitions=3).mapInArrow(bad_iter_elem, "a int").count()) + + def test_empty_iterator(self): + def empty_iter(_): + return iter([]) + + self.assertEqual(self.spark.range(10).mapInArrow(empty_iter, "a int, b string").count(), 0) + + def test_empty_rows(self): + def empty_rows(_): + return iter([pa.RecordBatch.from_pandas(pd.DataFrame({"a": []}))]) + + self.assertEqual(self.spark.range(10).mapInArrow(empty_rows, "a int").count(), 0) + + def test_chain_map_in_arrow(self): + def func(iterator): + for batch in iterator: + assert isinstance(batch, pa.RecordBatch) + assert batch.schema.names == ["id"] + yield batch + + df = self.spark.range(10) + actual = df.mapInArrow(func, "id long").mapInArrow(func, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_self_join(self): + df1 = self.spark.range(10) + df2 = df1.mapInArrow(lambda iter: iter, "id long") + actual = df2.join(df2).collect() + expected = df1.join(df1).collect() + self.assertEqual(sorted(actual), sorted(expected)) + + +class MapInArrowTests(MapInArrowTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + + # Synchronize default timezone between Python and Java + cls.tz_prev = os.environ.get("TZ", None) # save current tz if set + tz = "America/Los_Angeles" + os.environ["TZ"] = tz + time.tzset() + + cls.sc.environment["TZ"] = tz + cls.spark.conf.set("spark.sql.session.timeZone", tz) + + @classmethod + def tearDownClass(cls): + del os.environ["TZ"] + if cls.tz_prev is not None: + os.environ["TZ"] = cls.tz_prev + time.tzset() + ReusedSQLTestCase.tearDownClass() + + +if __name__ == "__main__": + from pyspark.sql.tests.test_arrow_map import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_python_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_python_udf.py new file mode 100644 index 0000000..f48f076 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_arrow_python_udf.py @@ -0,0 +1,202 @@ +# +# 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 unittest + +from pyspark.errors import PythonException +from pyspark.sql import Row +from pyspark.sql.functions import udf +from pyspark.sql.tests.test_udf import BaseUDFTestsMixin +from pyspark.testing.sqlutils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, + ReusedSQLTestCase, +) +from pyspark.rdd import PythonEvalType + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message +) +class PythonUDFArrowTestsMixin(BaseUDFTestsMixin): + @unittest.skip("Unrelated test, and it fails when it runs duplicatedly.") + def test_broadcast_in_udf(self): + super(PythonUDFArrowTests, self).test_broadcast_in_udf() + + @unittest.skip("Unrelated test, and it fails when it runs duplicatedly.") + def test_register_java_function(self): + super(PythonUDFArrowTests, self).test_register_java_function() + + @unittest.skip("Unrelated test, and it fails when it runs duplicatedly.") + def test_register_java_udaf(self): + super(PythonUDFArrowTests, self).test_register_java_udaf() + + def test_complex_input_types(self): + row = ( + self.spark.range(1) + .selectExpr("array(1, 2, 3) as array", "map('a', 'b') as map", "struct(1, 2) as struct") + .select( + udf(lambda x: str(x))("array"), + udf(lambda x: str(x))("map"), + udf(lambda x: str(x))("struct"), + ) + .first() + ) + + self.assertEquals(row[0], "[1, 2, 3]") + self.assertEquals(row[1], "{'a': 'b'}") + self.assertEquals(row[2], "Row(col1=1, col2=2)") + + def test_use_arrow(self): + # useArrow=True + row_true = ( + self.spark.range(1) + .selectExpr( + "array(1, 2, 3) as array", + ) + .select( + udf(lambda x: str(x), useArrow=True)("array"), + ) + .first() + ) + + # useArrow=None + row_none = ( + self.spark.range(1) + .selectExpr( + "array(1, 2, 3) as array", + ) + .select( + udf(lambda x: str(x), useArrow=None)("array"), + ) + .first() + ) + + self.assertEquals(row_true[0], row_none[0]) # "[1, 2, 3]" + + # useArrow=False + row_false = ( + self.spark.range(1) + .selectExpr( + "array(1, 2, 3) as array", + ) + .select( + udf(lambda x: str(x), useArrow=False)("array"), + ) + .first() + ) + self.assertEquals(row_false[0], "[1, 2, 3]") + + def test_eval_type(self): + self.assertEquals( + udf(lambda x: str(x), useArrow=True).evalType, PythonEvalType.SQL_ARROW_BATCHED_UDF + ) + self.assertEquals( + udf(lambda x: str(x), useArrow=False).evalType, PythonEvalType.SQL_BATCHED_UDF + ) + + def test_register(self): + df = self.spark.range(1).selectExpr( + "array(1, 2, 3) as array", + ) + str_repr_func = self.spark.udf.register("str_repr", udf(lambda x: str(x), useArrow=True)) + + # To verify that Arrow optimization is on + self.assertEquals( + df.selectExpr("str_repr(array) AS str_id").first()[0], + "[1, 2, 3]", # The input is a NumPy array when the Arrow optimization is on + ) + + # To verify that a UserDefinedFunction is returned + self.assertListEqual( + df.selectExpr("str_repr(array) AS str_id").collect(), + df.select(str_repr_func("array").alias("str_id")).collect(), + ) + + def test_nested_array_input(self): + df = self.spark.range(1).selectExpr("array(array(1, 2), array(3, 4)) as nested_array") + self.assertEquals( + df.select( + udf(lambda x: str(x), returnType="string", useArrow=True)("nested_array") + ).first()[0], + "[[1, 2], [3, 4]]", + ) + + def test_type_coercion_string_to_numeric(self): + df_int_value = self.spark.createDataFrame(["1", "2"], schema="string") + df_floating_value = self.spark.createDataFrame(["1.1", "2.2"], schema="string") + + int_ddl_types = ["tinyint", "smallint", "int", "bigint"] + floating_ddl_types = ["double", "float"] + + for ddl_type in int_ddl_types: + # df_int_value + res = df_int_value.select(udf(lambda x: x, ddl_type)("value").alias("res")) + self.assertEquals(res.collect(), [Row(res=1), Row(res=2)]) + self.assertEquals(res.dtypes[0][1], ddl_type) + + floating_results = [ + [Row(res=1.1), Row(res=2.2)], + [Row(res=1.100000023841858), Row(res=2.200000047683716)], + ] + for ddl_type, floating_res in zip(floating_ddl_types, floating_results): + # df_int_value + res = df_int_value.select(udf(lambda x: x, ddl_type)("value").alias("res")) + self.assertEquals(res.collect(), [Row(res=1.0), Row(res=2.0)]) + self.assertEquals(res.dtypes[0][1], ddl_type) + # df_floating_value + res = df_floating_value.select(udf(lambda x: x, ddl_type)("value").alias("res")) + self.assertEquals(res.collect(), floating_res) + self.assertEquals(res.dtypes[0][1], ddl_type) + + # invalid + with self.assertRaises(PythonException): + df_floating_value.select(udf(lambda x: x, "int")("value").alias("res")).collect() + + with self.assertRaises(PythonException): + df_int_value.select(udf(lambda x: x, "decimal")("value").alias("res")).collect() + + with self.assertRaises(PythonException): + df_floating_value.select(udf(lambda x: x, "decimal")("value").alias("res")).collect() + + +class PythonUDFArrowTests(PythonUDFArrowTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + super(PythonUDFArrowTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "true") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDF.arrow.enabled") + finally: + super(PythonUDFArrowTests, cls).tearDownClass() + + +if __name__ == "__main__": + from pyspark.sql.tests.test_arrow_python_udf import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_catalog.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_catalog.py new file mode 100644 index 0000000..cafffdc --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_catalog.py @@ -0,0 +1,503 @@ +# +# 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 import StorageLevel +from pyspark.errors import AnalysisException +from pyspark.sql.types import StructType, StructField, IntegerType +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class CatalogTestsMixin: + def test_current_database(self): + spark = self.spark + with self.database("some_db"): + self.assertEqual(spark.catalog.currentDatabase(), "default") + spark.sql("CREATE DATABASE some_db") + spark.catalog.setCurrentDatabase("some_db") + self.assertEqual(spark.catalog.currentDatabase(), "some_db") + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.setCurrentDatabase("does_not_exist"), + ) + + def test_list_databases(self): + spark = self.spark + with self.database("some_db"): + databases = [db.name for db in spark.catalog.listDatabases()] + self.assertEqual(databases, ["default"]) + spark.sql("CREATE DATABASE some_db") + databases = [db.name for db in spark.catalog.listDatabases()] + self.assertEqual(sorted(databases), ["default", "some_db"]) + databases = [db.name for db in spark.catalog.listDatabases("def*")] + self.assertEqual(sorted(databases), ["default"]) + databases = [db.name for db in spark.catalog.listDatabases("def2*")] + self.assertEqual(sorted(databases), []) + + def test_database_exists(self): + # SPARK-36207: testing that database_exists returns correct boolean + spark = self.spark + with self.database("some_db"): + self.assertFalse(spark.catalog.databaseExists("some_db")) + spark.sql("CREATE DATABASE some_db") + self.assertTrue(spark.catalog.databaseExists("some_db")) + self.assertTrue(spark.catalog.databaseExists("spark_catalog.some_db")) + self.assertFalse(spark.catalog.databaseExists("spark_catalog.some_db2")) + + def test_get_database(self): + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + db = spark.catalog.getDatabase("spark_catalog.some_db") + self.assertEqual(db.name, "some_db") + self.assertEqual(db.catalog, "spark_catalog") + + def test_list_tables(self): + from pyspark.sql.catalog import Table + + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + with self.table("tab1", "some_db.tab2", "tab3_via_catalog"): + with self.tempView("temp_tab"): + self.assertEqual(spark.catalog.listTables(), []) + self.assertEqual(spark.catalog.listTables("some_db"), []) + spark.createDataFrame([(1, 1)]).createOrReplaceTempView("temp_tab") + spark.sql("CREATE TABLE tab1 (name STRING, age INT) USING parquet") + spark.sql("CREATE TABLE some_db.tab2 (name STRING, age INT) USING parquet") + + schema = StructType([StructField("a", IntegerType(), True)]) + description = "this a table created via Catalog.createTable()" + spark.catalog.createTable( + "tab3_via_catalog", schema=schema, description=description + ) + + tables = sorted(spark.catalog.listTables(), key=lambda t: t.name) + tablesWithPattern = sorted( + spark.catalog.listTables(pattern="tab*"), key=lambda t: t.name + ) + tablesDefault = sorted( + spark.catalog.listTables("default"), key=lambda t: t.name + ) + tablesDefaultWithPattern = sorted( + spark.catalog.listTables("default", "tab*"), key=lambda t: t.name + ) + tablesSomeDb = sorted(spark.catalog.listTables("some_db"), key=lambda t: t.name) + tablesSomeDbWithPattern = sorted( + spark.catalog.listTables("some_db", "tab*"), key=lambda t: t.name + ) + self.assertEqual(tables, tablesDefault) + self.assertEqual(tablesWithPattern, tablesDefaultWithPattern) + self.assertEqual(len(tables), 3) + self.assertEqual(len(tablesWithPattern), 2) + self.assertEqual(len(tablesSomeDb), 2) + self.assertEqual(len(tablesSomeDbWithPattern), 1) + + # make table in old fashion + def makeTable( + name, + database, + description, + tableType, + isTemporary, + ): + return Table( + name=name, + catalog=None, + namespace=[database] if database is not None else None, + description=description, + tableType=tableType, + isTemporary=isTemporary, + ) + + # compare tables in old fashion + def compareTables(t1, t2): + return ( + t1.name == t2.name + and t1.database == t2.database + and t1.description == t2.description + and t1.tableType == t2.tableType + and t1.isTemporary == t2.isTemporary + ) + + self.assertTrue( + compareTables( + tables[0], + makeTable( + name="tab1", + database="default", + description=None, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertTrue( + compareTables( + tables[1], + makeTable( + name="tab3_via_catalog", + database="default", + description=description, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertTrue( + compareTables( + tables[2], + makeTable( + name="temp_tab", + database=None, + description=None, + tableType="TEMPORARY", + isTemporary=True, + ), + ) + ) + self.assertTrue( + compareTables( + tablesWithPattern[0], + makeTable( + name="tab1", + database="default", + description=None, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertTrue( + compareTables( + tablesWithPattern[1], + makeTable( + name="tab3_via_catalog", + database="default", + description=description, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertTrue( + compareTables( + tablesSomeDb[0], + makeTable( + name="tab2", + database="some_db", + description=None, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertTrue( + compareTables( + tablesSomeDb[1], + makeTable( + name="temp_tab", + database=None, + description=None, + tableType="TEMPORARY", + isTemporary=True, + ), + ) + ) + self.assertTrue( + compareTables( + tablesSomeDbWithPattern[0], + makeTable( + name="tab2", + database="some_db", + description=None, + tableType="MANAGED", + isTemporary=False, + ), + ) + ) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.listTables("does_not_exist"), + ) + + def test_list_functions(self): + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + functions = dict((f.name, f) for f in spark.catalog.listFunctions()) + functionsDefault = dict((f.name, f) for f in spark.catalog.listFunctions("default")) + self.assertTrue(len(functions) > 200) + self.assertTrue("+" in functions) + self.assertTrue("like" in functions) + self.assertTrue("month" in functions) + self.assertTrue("to_date" in functions) + self.assertTrue("to_timestamp" in functions) + self.assertTrue("to_unix_timestamp" in functions) + self.assertTrue("current_database" in functions) + self.assertEqual(functions["+"].name, "+") + self.assertEqual(functions["+"].description, "expr1 + expr2 - Returns `expr1`+`expr2`.") + self.assertEqual( + functions["+"].className, "org.apache.spark.sql.catalyst.expressions.Add" + ) + self.assertTrue(functions["+"].isTemporary) + self.assertEqual(functions, functionsDefault) + + functionsWithPattern = dict( + (f.name, f) for f in spark.catalog.listFunctions(pattern="to*") + ) + functionsDefaultWithPattern = dict( + (f.name, f) for f in spark.catalog.listFunctions("default", "to*") + ) + self.assertTrue(len(functionsWithPattern) > 10) + self.assertFalse("+" in functionsWithPattern) + self.assertFalse("like" in functionsWithPattern) + self.assertFalse("month" in functionsWithPattern) + self.assertTrue("to_date" in functionsWithPattern) + self.assertTrue("to_timestamp" in functionsWithPattern) + self.assertTrue("to_unix_timestamp" in functionsWithPattern) + self.assertEqual(functionsWithPattern, functionsDefaultWithPattern) + functionsWithPattern = dict( + (f.name, f) for f in spark.catalog.listFunctions(pattern="*not_existing_func*") + ) + self.assertTrue(len(functionsWithPattern) == 0) + + with self.function("func1", "some_db.func2"): + try: + spark.udf + support_udf = True + except Exception: + support_udf = False + + if support_udf: + spark.udf.register("temp_func", lambda x: str(x)) + spark.sql("CREATE FUNCTION func1 AS 'org.apache.spark.data.bricks'") + spark.sql("CREATE FUNCTION some_db.func2 AS 'org.apache.spark.data.bricks'") + newFunctions = dict((f.name, f) for f in spark.catalog.listFunctions()) + newFunctionsSomeDb = dict( + (f.name, f) for f in spark.catalog.listFunctions("some_db") + ) + self.assertTrue(set(functions).issubset(set(newFunctions))) + self.assertTrue(set(functions).issubset(set(newFunctionsSomeDb))) + if support_udf: + self.assertTrue("temp_func" in newFunctions) + self.assertTrue("func1" in newFunctions) + self.assertTrue("func2" not in newFunctions) + if support_udf: + self.assertTrue("temp_func" in newFunctionsSomeDb) + self.assertTrue("func1" not in newFunctionsSomeDb) + self.assertTrue("func2" in newFunctionsSomeDb) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.listFunctions("does_not_exist"), + ) + + def test_function_exists(self): + # SPARK-36258: testing that function_exists returns correct boolean + spark = self.spark + with self.function("func1"): + self.assertFalse(spark.catalog.functionExists("func1")) + self.assertFalse(spark.catalog.functionExists("default.func1")) + self.assertFalse(spark.catalog.functionExists("spark_catalog.default.func1")) + self.assertFalse(spark.catalog.functionExists("func1", "default")) + spark.sql("CREATE FUNCTION func1 AS 'org.apache.spark.data.bricks'") + self.assertTrue(spark.catalog.functionExists("func1")) + self.assertTrue(spark.catalog.functionExists("default.func1")) + self.assertTrue(spark.catalog.functionExists("spark_catalog.default.func1")) + self.assertTrue(spark.catalog.functionExists("func1", "default")) + + def test_get_function(self): + spark = self.spark + with self.function("func1"): + spark.sql("CREATE FUNCTION func1 AS 'org.apache.spark.data.bricks'") + func1 = spark.catalog.getFunction("spark_catalog.default.func1") + self.assertTrue(func1.name == "func1") + self.assertTrue(func1.namespace == ["default"]) + self.assertTrue(func1.catalog == "spark_catalog") + self.assertTrue(func1.className == "org.apache.spark.data.bricks") + self.assertFalse(func1.isTemporary) + + def test_list_columns(self): + from pyspark.sql.catalog import Column + + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + with self.table("tab1", "some_db.tab2"): + spark.sql("CREATE TABLE tab1 (name STRING, age INT) USING parquet") + spark.sql( + "CREATE TABLE some_db.tab2 (nickname STRING, tolerance FLOAT) USING parquet" + ) + columns = sorted( + spark.catalog.listColumns("spark_catalog.default.tab1"), key=lambda c: c.name + ) + columnsDefault = sorted( + spark.catalog.listColumns("tab1", "default"), key=lambda c: c.name + ) + self.assertEqual(columns, columnsDefault) + self.assertEqual(len(columns), 2) + self.assertEqual( + columns[0], + Column( + name="age", + description=None, + dataType="int", + nullable=True, + isPartition=False, + isBucket=False, + ), + ) + self.assertEqual( + columns[1], + Column( + name="name", + description=None, + dataType="string", + nullable=True, + isPartition=False, + isBucket=False, + ), + ) + columns2 = sorted( + spark.catalog.listColumns("tab2", "some_db"), key=lambda c: c.name + ) + self.assertEqual(len(columns2), 2) + self.assertEqual( + columns2[0], + Column( + name="nickname", + description=None, + dataType="string", + nullable=True, + isPartition=False, + isBucket=False, + ), + ) + self.assertEqual( + columns2[1], + Column( + name="tolerance", + description=None, + dataType="float", + nullable=True, + isPartition=False, + isBucket=False, + ), + ) + self.assertRaisesRegex( + AnalysisException, "tab2", lambda: spark.catalog.listColumns("tab2") + ) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.listColumns("does_not_exist"), + ) + + def test_table_cache(self): + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + with self.table("tab1"): + spark.sql("CREATE TABLE some_db.tab1 (name STRING, age INT) USING parquet") + + def if_cached(x): + return spark.catalog.isCached(x) + + names = ["some_db.tab1", "spark_catalog.some_db.tab1"] + + def assert_cached(c: bool): + if c: + self.assertTrue(all(map(if_cached, names))) + else: + self.assertFalse(any(map(if_cached, names))) + + assert_cached(False) + spark.catalog.cacheTable("spark_catalog.some_db.tab1") + assert_cached(True) + spark.catalog.uncacheTable("spark_catalog.some_db.tab1") + assert_cached(False) + spark.catalog.cacheTable("spark_catalog.some_db.tab1", StorageLevel.MEMORY_ONLY) + assert_cached(True) + spark.catalog.clearCache() + assert_cached(False) + + def test_table_exists(self): + # SPARK-36176: testing that table_exists returns correct boolean + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + with self.table("tab1", "some_db.tab2"): + self.assertFalse(spark.catalog.tableExists("tab1")) + self.assertFalse(spark.catalog.tableExists("tab2", "some_db")) + spark.sql("CREATE TABLE tab1 (name STRING, age INT) USING parquet") + self.assertTrue(spark.catalog.tableExists("tab1")) + self.assertTrue(spark.catalog.tableExists("default.tab1")) + self.assertTrue(spark.catalog.tableExists("spark_catalog.default.tab1")) + self.assertTrue(spark.catalog.tableExists("tab1", "default")) + spark.sql("CREATE TABLE some_db.tab2 (name STRING, age INT) USING parquet") + self.assertFalse(spark.catalog.tableExists("tab2")) + self.assertTrue(spark.catalog.tableExists("some_db.tab2")) + self.assertTrue(spark.catalog.tableExists("spark_catalog.some_db.tab2")) + self.assertTrue(spark.catalog.tableExists("tab2", "some_db")) + + def test_get_table(self): + spark = self.spark + with self.database("some_db"): + spark.sql("CREATE DATABASE some_db") + with self.table("tab1"): + spark.sql("CREATE TABLE tab1 (name STRING, age INT) USING parquet") + self.assertEqual(spark.catalog.getTable("tab1").database, "default") + self.assertEqual(spark.catalog.getTable("default.tab1").catalog, "spark_catalog") + self.assertEqual(spark.catalog.getTable("spark_catalog.default.tab1").name, "tab1") + + def test_refresh_table(self): + import os + import tempfile + + spark = self.spark + with tempfile.TemporaryDirectory() as tmp_dir: + with self.table("my_tab"): + spark.sql( + "CREATE TABLE my_tab (col STRING) USING TEXT LOCATION '{}'".format(tmp_dir) + ) + spark.sql("INSERT INTO my_tab SELECT 'abc'") + spark.catalog.cacheTable("my_tab") + self.assertEqual(spark.table("my_tab").count(), 1) + + os.system("rm -rf {}/*".format(tmp_dir)) + self.assertEqual(spark.table("my_tab").count(), 1) + + spark.catalog.refreshTable("spark_catalog.default.my_tab") + self.assertEqual(spark.table("my_tab").count(), 0) + + +class CatalogTests(ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_catalog import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_column.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_column.py new file mode 100644 index 0000000..db6cd32 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_column.py @@ -0,0 +1,225 @@ +# -*- encoding: utf-8 -*- +# +# 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.sql import Column, Row +from pyspark.sql.types import StructType, StructField, LongType +from pyspark.errors import AnalysisException, PySparkTypeError +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class ColumnTestsMixin: + def test_column_name_encoding(self): + """Ensure that created columns has `str` type consistently.""" + columns = self.spark.createDataFrame([("Alice", 1)], ["name", "age"]).columns + self.assertEqual(columns, ["name", "age"]) + self.assertTrue(isinstance(columns[0], str)) + self.assertTrue(isinstance(columns[1], str)) + + def test_and_in_expression(self): + self.assertEqual(4, self.df.filter((self.df.key <= 10) & (self.df.value <= "2")).count()) + self.assertRaises(ValueError, lambda: (self.df.key <= 10) and (self.df.value <= "2")) + self.assertEqual(14, self.df.filter((self.df.key <= 3) | (self.df.value < "2")).count()) + self.assertRaises(ValueError, lambda: self.df.key <= 3 or self.df.value < "2") + self.assertEqual(99, self.df.filter(~(self.df.key == 1)).count()) + self.assertRaises(ValueError, lambda: not self.df.key == 1) + + def test_validate_column_types(self): + from pyspark.sql.functions import udf, to_json + from pyspark.sql.column import _to_java_column + + self.assertTrue("Column" in _to_java_column("a").getClass().toString()) + self.assertTrue("Column" in _to_java_column("a").getClass().toString()) + self.assertTrue("Column" in _to_java_column(self.spark.range(1).id).getClass().toString()) + + with self.assertRaises(PySparkTypeError) as pe: + _to_java_column(1) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": "int"}, + ) + + class A: + pass + + self.assertRaises(TypeError, lambda: _to_java_column(A())) + self.assertRaises(TypeError, lambda: _to_java_column([])) + + with self.assertRaises(PySparkTypeError) as pe: + udf(lambda x: x)(None) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": "NoneType"}, + ) + self.assertRaises(TypeError, lambda: to_json(1)) + + def test_column_operators(self): + ci = self.df.key + cs = self.df.value + ci == cs + self.assertTrue(isinstance((-ci - 1 - 2) % 3 * 2.5 / 3.5, Column)) + rcc = (1 + ci), (1 - ci), (1 * ci), (1 / ci), (1 % ci), (1**ci), (ci**1) + self.assertTrue(all(isinstance(c, Column) for c in rcc)) + cb = [ci == 5, ci != 0, ci > 3, ci < 4, ci >= 0, ci <= 7] + self.assertTrue(all(isinstance(c, Column) for c in cb)) + cbool = (ci & ci), (ci | ci), (~ci) + self.assertTrue(all(isinstance(c, Column) for c in cbool)) + css = ( + cs.contains("a"), + cs.like("a"), + cs.rlike("a"), + cs.ilike("A"), + cs.asc(), + cs.desc(), + cs.startswith("a"), + cs.endswith("a"), + ci.eqNullSafe(cs), + ) + self.assertTrue(all(isinstance(c, Column) for c in css)) + self.assertTrue(isinstance(ci.cast(LongType()), Column)) + self.assertRaisesRegex( + ValueError, "Cannot apply 'in' operator against a column", lambda: 1 in cs + ) + + def test_column_accessor(self): + from pyspark.sql.functions import col + + self.assertIsInstance(col("foo")[1:3], Column) + self.assertIsInstance(col("foo")[0], Column) + self.assertIsInstance(col("foo")["bar"], Column) + self.assertRaises(ValueError, lambda: col("foo")[0:10:2]) + + def test_column_select(self): + df = self.df + self.assertEqual(self.testData, df.select("*").collect()) + self.assertEqual(self.testData, df.select(df.key, df.value).collect()) + self.assertEqual([Row(value="1")], df.where(df.key == 1).select(df.value).collect()) + + def test_access_column(self): + df = self.df + self.assertTrue(isinstance(df.key, Column)) + self.assertTrue(isinstance(df["key"], Column)) + self.assertTrue(isinstance(df[0], Column)) + self.assertRaises(IndexError, lambda: df[2]) + self.assertRaises(AnalysisException, lambda: df["bad_key"]) + self.assertRaises(TypeError, lambda: df[{}]) + + def test_column_name_with_non_ascii(self): + columnName = "数量" + self.assertTrue(isinstance(columnName, str)) + schema = StructType([StructField(columnName, LongType(), True)]) + df = self.spark.createDataFrame([(1,)], schema) + self.assertEqual(schema, df.schema) + self.assertEqual("DataFrame[数量: bigint]", str(df)) + self.assertEqual([("数量", "bigint")], df.dtypes) + self.assertEqual(1, df.select("数量").first()[0]) + self.assertEqual(1, df.select(df["数量"]).first()[0]) + self.assertTrue(columnName in repr(df[columnName])) + + def test_field_accessor(self): + df = self.spark.createDataFrame([Row(l=[1], r=Row(a=1, b="b"), d={"k": "v"})]) + self.assertEqual(1, df.select(df.l[0]).first()[0]) + self.assertEqual(1, df.select(df.r["a"]).first()[0]) + self.assertEqual(1, df.select(df["r.a"]).first()[0]) + self.assertEqual("b", df.select(df.r["b"]).first()[0]) + self.assertEqual("b", df.select(df["r.b"]).first()[0]) + self.assertEqual("v", df.select(df.d["k"]).first()[0]) + + def test_bitwise_operations(self): + from pyspark.sql import functions + + row = Row(a=170, b=75) + df = self.spark.createDataFrame([row]) + result = df.select(df.a.bitwiseAND(df.b)).collect()[0].asDict() + self.assertEqual(170 & 75, result["(a & b)"]) + result = df.select(df.a.bitwiseOR(df.b)).collect()[0].asDict() + self.assertEqual(170 | 75, result["(a | b)"]) + result = df.select(df.a.bitwiseXOR(df.b)).collect()[0].asDict() + self.assertEqual(170 ^ 75, result["(a ^ b)"]) + result = df.select(functions.bitwiseNOT(df.b)).collect()[0].asDict() + self.assertEqual(~75, result["~b"]) + result = df.select(functions.bitwise_not(df.b)).collect()[0].asDict() + self.assertEqual(~75, result["~b"]) + + def test_with_field(self): + from pyspark.sql.functions import lit, col + + df = self.spark.createDataFrame([Row(a=Row(b=1, c=2))]) + self.assertIsInstance(df["a"].withField("b", lit(3)), Column) + self.assertIsInstance(df["a"].withField("d", lit(3)), Column) + result = df.withColumn("a", df["a"].withField("d", lit(3))).collect()[0].asDict() + self.assertEqual(3, result["a"]["d"]) + result = df.withColumn("a", df["a"].withField("b", lit(3))).collect()[0].asDict() + self.assertEqual(3, result["a"]["b"]) + + with self.assertRaises(PySparkTypeError) as pe: + df["a"].withField("b", 3) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN", + message_parameters={"arg_name": "col", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df["a"].withField(col("b"), lit(3)) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "fieldName", "arg_type": "Column"}, + ) + + def test_drop_fields(self): + df = self.spark.createDataFrame([Row(a=Row(b=1, c=2, d=Row(e=3, f=4)))]) + self.assertIsInstance(df["a"].dropFields("b"), Column) + self.assertIsInstance(df["a"].dropFields("b", "c"), Column) + self.assertIsInstance(df["a"].dropFields("d.e"), Column) + + result = ( + df.select( + df["a"].dropFields("b").alias("a1"), + df["a"].dropFields("d.e").alias("a2"), + ) + .first() + .asDict(True) + ) + + self.assertTrue("b" not in result["a1"] and "c" in result["a1"] and "d" in result["a1"]) + + self.assertTrue("e" not in result["a2"]["d"] and "f" in result["a2"]["d"]) + + +class ColumnTests(ColumnTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_column import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_conf.py new file mode 100644 index 0000000..15722c2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_conf.py @@ -0,0 +1,83 @@ +# +# 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 decimal import Decimal + +from pyspark.errors import IllegalArgumentException +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class ConfTestsMixin: + def test_conf(self): + spark = self.spark + spark.conf.set("bogo", "sipeo") + self.assertEqual(spark.conf.get("bogo"), "sipeo") + spark.conf.set("bogo", "ta") + self.assertEqual(spark.conf.get("bogo"), "ta") + self.assertEqual(spark.conf.get("bogo", "not.read"), "ta") + self.assertEqual(spark.conf.get("not.set", "ta"), "ta") + self.assertRaisesRegex(Exception, "not.set", lambda: spark.conf.get("not.set")) + spark.conf.unset("bogo") + self.assertEqual(spark.conf.get("bogo", "colombia"), "colombia") + + self.assertEqual(spark.conf.get("hyukjin", None), None) + + # This returns 'STATIC' because it's the default value of + # 'spark.sql.sources.partitionOverwriteMode', and `defaultValue` in + # `spark.conf.get` is unset. + self.assertEqual(spark.conf.get("spark.sql.sources.partitionOverwriteMode"), "STATIC") + + # This returns None because 'spark.sql.sources.partitionOverwriteMode' is unset, but + # `defaultValue` in `spark.conf.get` is set to None. + self.assertEqual(spark.conf.get("spark.sql.sources.partitionOverwriteMode", None), None) + + self.assertTrue(spark.conf.isModifiable("spark.sql.execution.arrow.maxRecordsPerBatch")) + self.assertFalse(spark.conf.isModifiable("spark.sql.warehouse.dir")) + + def test_conf_with_python_objects(self): + spark = self.spark + + for value, expected in [(True, "true"), (False, "false")]: + spark.conf.set("foo", value) + self.assertEqual(spark.conf.get("foo"), expected) + + spark.conf.set("foo", 1) + self.assertEqual(spark.conf.get("foo"), "1") + + with self.assertRaises(IllegalArgumentException): + spark.conf.set("foo", None) + + with self.assertRaises(Exception): + spark.conf.set("foo", Decimal(1)) + + spark.conf.unset("foo") + + +class ConfTests(ConfTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_conf import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_context.py new file mode 100644 index 0000000..b381833 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_context.py @@ -0,0 +1,201 @@ +# +# 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 os +import shutil +import sys +import tempfile +import unittest +from importlib import reload + +import py4j + +from pyspark import SparkContext, SQLContext +from pyspark.sql import Row, SparkSession +from pyspark.sql.types import StructType, StringType, StructField +from pyspark.testing.utils import ReusedPySparkTestCase + + +class HiveContextSQLTests(ReusedPySparkTestCase): + @classmethod + def setUpClass(cls): + ReusedPySparkTestCase.setUpClass() + cls.tempdir = tempfile.NamedTemporaryFile(delete=False) + cls.hive_available = True + cls.spark = None + try: + cls.sc._jvm.org.apache.hadoop.hive.conf.HiveConf() + except py4j.protocol.Py4JError: + cls.tearDownClass() + cls.hive_available = False + except TypeError: + cls.tearDownClass() + cls.hive_available = False + if cls.hive_available: + cls.spark = SparkSession.builder.enableHiveSupport().getOrCreate() + + os.unlink(cls.tempdir.name) + if cls.hive_available: + cls.testData = [Row(key=i, value=str(i)) for i in range(100)] + cls.df = cls.sc.parallelize(cls.testData).toDF() + + def setUp(self): + if not self.hive_available: + self.skipTest("Hive is not available.") + + @classmethod + def tearDownClass(cls): + ReusedPySparkTestCase.tearDownClass() + shutil.rmtree(cls.tempdir.name, ignore_errors=True) + if cls.spark is not None: + cls.spark.stop() + cls.spark = None + + def test_save_and_load_table(self): + df = self.df + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + df.write.saveAsTable("savedJsonTable", "json", "append", path=tmpPath) + actual = self.spark.catalog.createTable("externalJsonTable", tmpPath, "json") + self.assertEqual( + sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect()) + ) + self.assertEqual( + sorted(df.collect()), + sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()), + ) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + self.spark.sql("DROP TABLE externalJsonTable") + + df.write.saveAsTable("savedJsonTable", "json", "overwrite", path=tmpPath) + schema = StructType([StructField("value", StringType(), True)]) + actual = self.spark.catalog.createTable( + "externalJsonTable", + source="json", + schema=schema, + path=tmpPath, + noUse="this options will not be used", + ) + self.assertEqual( + sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect()) + ) + self.assertEqual( + sorted(df.select("value").collect()), + sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()), + ) + self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect())) + self.spark.sql("DROP TABLE savedJsonTable") + self.spark.sql("DROP TABLE externalJsonTable") + + defaultDataSourceName = self.spark.conf.get( + "spark.sql.sources.default", "org.apache.spark.sql.parquet" + ) + self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json") + df.write.saveAsTable("savedJsonTable", path=tmpPath, mode="overwrite") + actual = self.spark.catalog.createTable("externalJsonTable", path=tmpPath) + self.assertEqual( + sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect()) + ) + self.assertEqual( + sorted(df.collect()), + sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()), + ) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + self.spark.sql("DROP TABLE savedJsonTable") + self.spark.sql("DROP TABLE externalJsonTable") + self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName) + + shutil.rmtree(tmpPath) + + def test_limit_and_take(self): + df = self.spark.range(1, 1000, numPartitions=10) + + def assert_runs_only_one_job_stage_and_task(job_group_name, f): + tracker = self.sc.statusTracker() + self.sc.setJobGroup(job_group_name, description="") + f() + jobs = tracker.getJobIdsForGroup(job_group_name) + self.assertEqual(1, len(jobs)) + stages = tracker.getJobInfo(jobs[0]).stageIds + self.assertEqual(1, len(stages)) + self.assertEqual(1, tracker.getStageInfo(stages[0]).numTasks) + + # Regression test for SPARK-10731: take should delegate to Scala implementation + assert_runs_only_one_job_stage_and_task("take", lambda: df.take(1)) + # Regression test for SPARK-17514: limit(n).collect() should the perform same as take(n) + assert_runs_only_one_job_stage_and_task("collect_limit", lambda: df.limit(1).collect()) + + def test_unbounded_frames(self): + from pyspark.sql import functions as F + from pyspark.sql import window + + df = self.spark.range(0, 3) + + def rows_frame_match(): + return ( + "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING" + in df.select( + F.count("*").over(window.Window.rowsBetween(-sys.maxsize, sys.maxsize)) + ).columns[0] + ) + + def range_frame_match(): + return ( + "RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING" + in df.select( + F.count("*").over(window.Window.rangeBetween(-sys.maxsize, sys.maxsize)) + ).columns[0] + ) + + for new_maxsize in [2**31 - 1, 2**63 - 1, 2**127 - 1]: + old_maxsize = sys.maxsize + sys.maxsize = new_maxsize + try: + # Manually reload window module to use monkey-patched sys.maxsize. + reload(window) + self.assertTrue(rows_frame_match()) + self.assertTrue(range_frame_match()) + finally: + sys.maxsize = old_maxsize + + reload(window) + + +class SQLContextTests(unittest.TestCase): + def test_get_or_create(self): + sc = None + sql_context = None + try: + sc = SparkContext("local[4]", "SQLContextTests") + sql_context = SQLContext.getOrCreate(sc) + assert isinstance(sql_context, SQLContext) + finally: + if sql_context is not None: + sql_context.sparkSession.stop() + if sc is not None: + sc.stop() + + +if __name__ == "__main__": + from pyspark.sql.tests.test_context import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_dataframe.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_dataframe.py new file mode 100644 index 0000000..3304923 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_dataframe.py @@ -0,0 +1,1886 @@ +# +# 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 platform +from decimal import Decimal +import os +import pydoc +import shutil +import tempfile +import time +import unittest +from typing import cast +import io +from contextlib import redirect_stdout + +from pyspark import StorageLevel +from pyspark.sql import SparkSession, Row +from pyspark.sql.functions import col, lit, count, sum, mean, struct +from pyspark.sql.pandas.utils import pyarrow_version_less_than_minimum +from pyspark.sql.types import ( + StringType, + IntegerType, + DoubleType, + LongType, + StructType, + StructField, + BooleanType, + DateType, + TimestampType, + TimestampNTZType, + FloatType, + DayTimeIntervalType, +) +from pyspark.storagelevel import StorageLevel +from pyspark.errors import ( + AnalysisException, + IllegalArgumentException, + PySparkTypeError, + PySparkValueError, +) +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + SQLTestUtils, + have_pyarrow, + have_pandas, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.testing.utils import QuietTest + + +class DataFrameTestsMixin: + def test_range(self): + self.assertEqual(self.spark.range(1, 1).count(), 0) + self.assertEqual(self.spark.range(1, 0, -1).count(), 1) + self.assertEqual(self.spark.range(0, 1 << 40, 1 << 39).count(), 2) + self.assertEqual(self.spark.range(-2).count(), 0) + self.assertEqual(self.spark.range(3).count(), 3) + + def test_duplicated_column_names(self): + df = self.spark.createDataFrame([(1, 2)], ["c", "c"]) + row = df.select("*").first() + self.assertEqual(1, row[0]) + self.assertEqual(2, row[1]) + self.assertEqual("Row(c=1, c=2)", str(row)) + # Cannot access columns + self.assertRaises(AnalysisException, lambda: df.select(df[0]).first()) + self.assertRaises(AnalysisException, lambda: df.select(df.c).first()) + self.assertRaises(AnalysisException, lambda: df.select(df["c"]).first()) + + def test_freqItems(self): + vals = [Row(a=1, b=-2.0) if i % 2 == 0 else Row(a=i, b=i * 1.0) for i in range(100)] + df = self.spark.createDataFrame(vals) + items = df.stat.freqItems(("a", "b"), 0.4).collect()[0] + self.assertTrue(1 in items[0]) + self.assertTrue(-2.0 in items[1]) + + def test_help_command(self): + # Regression test for SPARK-5464 + rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}']) + df = self.spark.read.json(rdd) + # render_doc() reproduces the help() exception without printing output + pydoc.render_doc(df) + pydoc.render_doc(df.foo) + pydoc.render_doc(df.take(1)) + + def test_drop(self): + df = self.spark.createDataFrame([("A", 50, "Y"), ("B", 60, "Y")], ["name", "age", "active"]) + self.assertEqual(df.drop("active").columns, ["name", "age"]) + self.assertEqual(df.drop("active", "nonexistent_column").columns, ["name", "age"]) + self.assertEqual(df.drop("name", "age", "active").columns, []) + self.assertEqual(df.drop(col("name")).columns, ["age", "active"]) + self.assertEqual(df.drop(col("name"), col("age")).columns, ["active"]) + self.assertEqual(df.drop(col("name"), col("age"), col("random")).columns, ["active"]) + + def test_with_columns_renamed(self): + df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"]) + + # rename both columns + renamed_df1 = df.withColumnsRenamed({"name": "naam", "age": "leeftijd"}) + self.assertEqual(renamed_df1.columns, ["naam", "leeftijd"]) + + # rename one column with one missing name + renamed_df2 = df.withColumnsRenamed({"name": "naam", "address": "adres"}) + self.assertEqual(renamed_df2.columns, ["naam", "age"]) + + # negative test for incorrect type + with self.assertRaises(PySparkTypeError) as pe: + df.withColumnsRenamed(("name", "x")) + + self.check_error( + exception=pe.exception, + error_class="NOT_DICT", + message_parameters={"arg_name": "colsMap", "arg_type": "tuple"}, + ) + + def test_drop_duplicates(self): + # SPARK-36034 test that drop duplicates throws a type error when in correct type provided + df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"]) + + # shouldn't drop a non-null row + self.assertEqual(df.dropDuplicates().count(), 2) + + self.assertEqual(df.dropDuplicates(["name"]).count(), 1) + + self.assertEqual(df.dropDuplicates(["name", "age"]).count(), 2) + + with self.assertRaises(PySparkTypeError) as pe: + df.dropDuplicates("name") + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": "str"}, + ) + + def test_drop_duplicates_with_ambiguous_reference(self): + df1 = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + df2 = self.spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")]) + df3 = df1.join(df2, df1.name == df2.name, "inner") + + self.assertEqual(df3.drop("name", "age").columns, ["height"]) + self.assertEqual(df3.drop("name", df3.age, "unknown").columns, ["height"]) + self.assertEqual(df3.drop("name", "age", df3.height).columns, []) + + def test_drop_empty_column(self): + df = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + + self.assertEqual(df.drop().columns, ["age", "name"]) + self.assertEqual(df.drop(*[]).columns, ["age", "name"]) + + def test_drop_column_name_with_dot(self): + df = ( + self.spark.range(1, 3) + .withColumn("first.name", lit("Peter")) + .withColumn("city.name", lit("raleigh")) + .withColumn("state", lit("nc")) + ) + + self.assertEqual(df.drop("first.name").columns, ["id", "city.name", "state"]) + self.assertEqual(df.drop("city.name").columns, ["id", "first.name", "state"]) + self.assertEqual(df.drop("first.name", "city.name").columns, ["id", "state"]) + self.assertEqual( + df.drop("first.name", "city.name", "unknown.unknown").columns, ["id", "state"] + ) + self.assertEqual( + df.drop("unknown.unknown").columns, ["id", "first.name", "city.name", "state"] + ) + + def test_dropna(self): + schema = StructType( + [ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True), + StructField("height", DoubleType(), True), + ] + ) + + # shouldn't drop a non-null row + self.assertEqual( + self.spark.createDataFrame([("Alice", 50, 80.1)], schema).dropna().count(), 1 + ) + + # dropping rows with a single null value + self.assertEqual( + self.spark.createDataFrame([("Alice", None, 80.1)], schema).dropna().count(), 0 + ) + self.assertEqual( + self.spark.createDataFrame([("Alice", None, 80.1)], schema).dropna(how="any").count(), 0 + ) + + # if how = 'all', only drop rows if all values are null + self.assertEqual( + self.spark.createDataFrame([("Alice", None, 80.1)], schema).dropna(how="all").count(), 1 + ) + self.assertEqual( + self.spark.createDataFrame([(None, None, None)], schema).dropna(how="all").count(), 0 + ) + + # how and subset + self.assertEqual( + self.spark.createDataFrame([("Alice", 50, None)], schema) + .dropna(how="any", subset=["name", "age"]) + .count(), + 1, + ) + self.assertEqual( + self.spark.createDataFrame([("Alice", None, None)], schema) + .dropna(how="any", subset=["name", "age"]) + .count(), + 0, + ) + + # threshold + self.assertEqual( + self.spark.createDataFrame([("Alice", None, 80.1)], schema).dropna(thresh=2).count(), 1 + ) + self.assertEqual( + self.spark.createDataFrame([("Alice", None, None)], schema).dropna(thresh=2).count(), 0 + ) + + # threshold and subset + self.assertEqual( + self.spark.createDataFrame([("Alice", 50, None)], schema) + .dropna(thresh=2, subset=["name", "age"]) + .count(), + 1, + ) + self.assertEqual( + self.spark.createDataFrame([("Alice", None, 180.9)], schema) + .dropna(thresh=2, subset=["name", "age"]) + .count(), + 0, + ) + + # thresh should take precedence over how + self.assertEqual( + self.spark.createDataFrame([("Alice", 50, None)], schema) + .dropna(how="any", thresh=2, subset=["name", "age"]) + .count(), + 1, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame([("Alice", 50, None)], schema).dropna(subset=10) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": "int"}, + ) + + def test_fillna(self): + schema = StructType( + [ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True), + StructField("height", DoubleType(), True), + StructField("spy", BooleanType(), True), + ] + ) + + # fillna shouldn't change non-null values + row = self.spark.createDataFrame([("Alice", 10, 80.1, True)], schema).fillna(50).first() + self.assertEqual(row.age, 10) + + # fillna with int + row = self.spark.createDataFrame([("Alice", None, None, None)], schema).fillna(50).first() + self.assertEqual(row.age, 50) + self.assertEqual(row.height, 50.0) + + # fillna with double + row = self.spark.createDataFrame([("Alice", None, None, None)], schema).fillna(50.1).first() + self.assertEqual(row.age, 50) + self.assertEqual(row.height, 50.1) + + # fillna with bool + row = self.spark.createDataFrame([("Alice", None, None, None)], schema).fillna(True).first() + self.assertEqual(row.age, None) + self.assertEqual(row.spy, True) + + # fillna with string + row = self.spark.createDataFrame([(None, None, None, None)], schema).fillna("hello").first() + self.assertEqual(row.name, "hello") + self.assertEqual(row.age, None) + + # fillna with subset specified for numeric cols + row = ( + self.spark.createDataFrame([(None, None, None, None)], schema) + .fillna(50, subset=["name", "age"]) + .first() + ) + self.assertEqual(row.name, None) + self.assertEqual(row.age, 50) + self.assertEqual(row.height, None) + self.assertEqual(row.spy, None) + + # fillna with subset specified for string cols + row = ( + self.spark.createDataFrame([(None, None, None, None)], schema) + .fillna("haha", subset=["name", "age"]) + .first() + ) + self.assertEqual(row.name, "haha") + self.assertEqual(row.age, None) + self.assertEqual(row.height, None) + self.assertEqual(row.spy, None) + + # fillna with subset specified for bool cols + row = ( + self.spark.createDataFrame([(None, None, None, None)], schema) + .fillna(True, subset=["name", "spy"]) + .first() + ) + self.assertEqual(row.name, None) + self.assertEqual(row.age, None) + self.assertEqual(row.height, None) + self.assertEqual(row.spy, True) + + # fillna with dictionary for boolean types + row = self.spark.createDataFrame([Row(a=None), Row(a=True)]).fillna({"a": True}).first() + self.assertEqual(row.a, True) + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame([Row(a=None), Row(a=True)]).fillna(["a", True]) + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_STR", + message_parameters={"arg_name": "value", "arg_type": "list"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame([Row(a=None), Row(a=True)]).fillna(50, subset=10) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OR_TUPLE", + message_parameters={"arg_name": "subset", "arg_type": "int"}, + ) + + def test_repartitionByRange_dataframe(self): + schema = StructType( + [ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True), + StructField("height", DoubleType(), True), + ] + ) + + df1 = self.spark.createDataFrame( + [("Bob", 27, 66.0), ("Alice", 10, 10.0), ("Bob", 10, 66.0)], schema + ) + df2 = self.spark.createDataFrame( + [("Alice", 10, 10.0), ("Bob", 10, 66.0), ("Bob", 27, 66.0)], schema + ) + + # test repartitionByRange(numPartitions, *cols) + df3 = df1.repartitionByRange(2, "name", "age") + self.assertEqual(df3.rdd.getNumPartitions(), 2) + self.assertEqual(df3.rdd.first(), df2.rdd.first()) + self.assertEqual(df3.rdd.take(3), df2.rdd.take(3)) + + # test repartitionByRange(numPartitions, *cols) + df4 = df1.repartitionByRange(3, "name", "age") + self.assertEqual(df4.rdd.getNumPartitions(), 3) + self.assertEqual(df4.rdd.first(), df2.rdd.first()) + self.assertEqual(df4.rdd.take(3), df2.rdd.take(3)) + + # test repartitionByRange(*cols) + df5 = df1.repartitionByRange(5, "name", "age") + self.assertEqual(df5.rdd.first(), df2.rdd.first()) + self.assertEqual(df5.rdd.take(3), df2.rdd.take(3)) + + with self.assertRaises(PySparkTypeError) as pe: + df1.repartitionByRange([10], "name", "age") + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "numPartitions", "arg_type": "list"}, + ) + + def test_replace(self): + schema = StructType( + [ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True), + StructField("height", DoubleType(), True), + ] + ) + + # replace with int + row = self.spark.createDataFrame([("Alice", 10, 10.0)], schema).replace(10, 20).first() + self.assertEqual(row.age, 20) + self.assertEqual(row.height, 20.0) + + # replace with double + row = self.spark.createDataFrame([("Alice", 80, 80.0)], schema).replace(80.0, 82.1).first() + self.assertEqual(row.age, 82) + self.assertEqual(row.height, 82.1) + + # replace with string + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace("Alice", "Ann") + .first() + ) + self.assertEqual(row.name, "Ann") + self.assertEqual(row.age, 10) + + # replace with subset specified by a string of a column name w/ actual change + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace(10, 20, subset="age") + .first() + ) + self.assertEqual(row.age, 20) + + # replace with subset specified by a string of a column name w/o actual change + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace(10, 20, subset="height") + .first() + ) + self.assertEqual(row.age, 10) + + # replace with subset specified with one column replaced, another column not in subset + # stays unchanged. + row = ( + self.spark.createDataFrame([("Alice", 10, 10.0)], schema) + .replace(10, 20, subset=["name", "age"]) + .first() + ) + self.assertEqual(row.name, "Alice") + self.assertEqual(row.age, 20) + self.assertEqual(row.height, 10.0) + + # replace with subset specified but no column will be replaced + row = ( + self.spark.createDataFrame([("Alice", 10, None)], schema) + .replace(10, 20, subset=["name", "height"]) + .first() + ) + self.assertEqual(row.name, "Alice") + self.assertEqual(row.age, 10) + self.assertEqual(row.height, None) + + # replace with lists + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace(["Alice"], ["Ann"]) + .first() + ) + self.assertTupleEqual(row, ("Ann", 10, 80.1)) + + # replace with dict + row = self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace({10: 11}).first() + self.assertTupleEqual(row, ("Alice", 11, 80.1)) + + # test backward compatibility with dummy value + dummy_value = 1 + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace({"Alice": "Bob"}, dummy_value) + .first() + ) + self.assertTupleEqual(row, ("Bob", 10, 80.1)) + + # test dict with mixed numerics + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace({10: -10, 80.1: 90.5}) + .first() + ) + self.assertTupleEqual(row, ("Alice", -10, 90.5)) + + # replace with tuples + row = ( + self.spark.createDataFrame([("Alice", 10, 80.1)], schema) + .replace(("Alice",), ("Bob",)) + .first() + ) + self.assertTupleEqual(row, ("Bob", 10, 80.1)) + + # replace multiple columns + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .replace((10, 80.0), (20, 90)) + .first() + ) + self.assertTupleEqual(row, ("Alice", 20, 90.0)) + + # test for mixed numerics + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .replace((10, 80), (20, 90.5)) + .first() + ) + self.assertTupleEqual(row, ("Alice", 20, 90.5)) + + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .replace({10: 20, 80: 90.5}) + .first() + ) + self.assertTupleEqual(row, ("Alice", 20, 90.5)) + + # replace with boolean + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .selectExpr("name = 'Bob'", "age <= 15") + .replace(False, True) + .first() + ) + self.assertTupleEqual(row, (True, True)) + + # replace string with None and then drop None rows + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .replace("Alice", None) + .dropna() + ) + self.assertEqual(row.count(), 0) + + # replace with number and None + row = ( + self.spark.createDataFrame([("Alice", 10, 80.0)], schema) + .replace([10, 80], [20, None]) + .first() + ) + self.assertTupleEqual(row, ("Alice", 20, None)) + + # should fail if subset is not list, tuple or None + with self.assertRaises(TypeError): + self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace( + {10: 11}, subset=1 + ).first() + + # should fail if to_replace and value have different length + with self.assertRaises(ValueError): + self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace( + ["Alice", "Bob"], ["Eve"] + ).first() + + # should fail if when received unexpected type + with self.assertRaises(TypeError): + from datetime import datetime + + self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace( + datetime.now(), datetime.now() + ).first() + + # should fail if provided mixed type replacements + with self.assertRaises(ValueError): + self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace( + ["Alice", 10], ["Eve", 20] + ).first() + + with self.assertRaises(ValueError): + self.spark.createDataFrame([("Alice", 10, 80.1)], schema).replace( + {"Alice": "Bob", 10: 20} + ).first() + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame([("Alice", 10, 80.0)], schema).replace(["Alice", "Bob"]) + + self.check_error( + exception=pe.exception, + error_class="ARGUMENT_REQUIRED", + message_parameters={"arg_name": "value", "condition": "`to_replace` is dict"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + self.spark.createDataFrame([("Alice", 10, 80.0)], schema).replace(lambda x: x + 1, 10) + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL_OR_DICT_OR_FLOAT_OR_INT_OR_LIST_OR_STR_OR_TUPLE", + message_parameters={"arg_name": "to_replace", "arg_type": "function"}, + ) + + def test_with_column_with_existing_name(self): + keys = self.df.withColumn("key", self.df.key).select("key").collect() + self.assertEqual([r.key for r in keys], list(range(100))) + + # regression test for SPARK-10417 + def test_column_iterator(self): + def foo(): + for x in self.df.key: + break + + self.assertRaises(TypeError, foo) + + def test_with_columns(self): + # With single column + keys = self.df.withColumns({"key": self.df.key}).select("key").collect() + self.assertEqual([r.key for r in keys], list(range(100))) + + # With key and value columns + kvs = ( + self.df.withColumns({"key": self.df.key, "value": self.df.value}) + .select("key", "value") + .collect() + ) + self.assertEqual([(r.key, r.value) for r in kvs], [(i, str(i)) for i in range(100)]) + + # Columns rename + kvs = ( + self.df.withColumns({"key_alias": self.df.key, "value_alias": self.df.value}) + .select("key_alias", "value_alias") + .collect() + ) + self.assertEqual( + [(r.key_alias, r.value_alias) for r in kvs], [(i, str(i)) for i in range(100)] + ) + + # Type check + self.assertRaises(TypeError, self.df.withColumns, ["key"]) + self.assertRaises(Exception, self.df.withColumns) + + def test_generic_hints(self): + df1 = self.spark.range(10e10).toDF("id") + df2 = self.spark.range(10e10).toDF("id") + + self.assertIsInstance(df1.hint("broadcast"), type(df1)) + + # Dummy rules + self.assertIsInstance(df1.hint("broadcast", "foo", "bar"), type(df1)) + + with io.StringIO() as buf, redirect_stdout(buf): + df1.join(df2.hint("broadcast"), "id").explain(True) + self.assertEqual(1, buf.getvalue().count("BroadcastHashJoin")) + + # add tests for SPARK-23647 (test more types for hint) + def test_extended_hint_types(self): + df = self.spark.range(10e10).toDF("id") + such_a_nice_list = ["itworks1", "itworks2", "itworks3"] + hinted_df = df.hint("my awesome hint", 1.2345, "what", such_a_nice_list) + + self.assertIsInstance(df.hint("broadcast", []), type(df)) + self.assertIsInstance(df.hint("broadcast", ["foo", "bar"]), type(df)) + + with io.StringIO() as buf, redirect_stdout(buf): + hinted_df.explain(True) + explain_output = buf.getvalue() + self.assertGreaterEqual(explain_output.count("1.2345"), 1) + self.assertGreaterEqual(explain_output.count("what"), 1) + self.assertGreaterEqual(explain_output.count("itworks"), 1) + + def test_unpivot(self): + # SPARK-39877: test the DataFrame.unpivot method + df = self.spark.createDataFrame( + [ + (1, 10, 1.0, "one"), + (2, 20, 2.0, "two"), + (3, 30, 3.0, "three"), + ], + ["id", "int", "double", "str"], + ) + + with self.subTest(desc="with none identifier"): + with self.assertRaisesRegex(AssertionError, "ids must not be None"): + df.unpivot(None, ["int", "double"], "var", "val") + + with self.subTest(desc="with no identifier"): + for id in [[], ()]: + with self.subTest(ids=id): + actual = df.unpivot(id, ["int", "double"], "var", "val") + self.assertEqual(actual.schema.simpleString(), "struct") + self.assertEqual( + actual.collect(), + [ + Row(var="int", value=10.0), + Row(var="double", value=1.0), + Row(var="int", value=20.0), + Row(var="double", value=2.0), + Row(var="int", value=30.0), + Row(var="double", value=3.0), + ], + ) + + with self.subTest(desc="with single identifier column"): + for id in ["id", ["id"], ("id",)]: + with self.subTest(ids=id): + actual = df.unpivot(id, ["int", "double"], "var", "val") + self.assertEqual( + actual.schema.simpleString(), + "struct", + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, var="int", value=10.0), + Row(id=1, var="double", value=1.0), + Row(id=2, var="int", value=20.0), + Row(id=2, var="double", value=2.0), + Row(id=3, var="int", value=30.0), + Row(id=3, var="double", value=3.0), + ], + ) + + with self.subTest(desc="with multiple identifier columns"): + for ids in [["id", "double"], ("id", "double")]: + with self.subTest(ids=ids): + actual = df.unpivot(ids, ["int", "double"], "var", "val") + self.assertEqual( + actual.schema.simpleString(), + "struct", + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, double=1.0, var="int", value=10.0), + Row(id=1, double=1.0, var="double", value=1.0), + Row(id=2, double=2.0, var="int", value=20.0), + Row(id=2, double=2.0, var="double", value=2.0), + Row(id=3, double=3.0, var="int", value=30.0), + Row(id=3, double=3.0, var="double", value=3.0), + ], + ) + + with self.subTest(desc="with no identifier columns but none value columns"): + # select only columns that have common data type (double) + actual = df.select("id", "int", "double").unpivot([], None, "var", "val") + self.assertEqual(actual.schema.simpleString(), "struct") + self.assertEqual( + actual.collect(), + [ + Row(var="id", value=1.0), + Row(var="int", value=10.0), + Row(var="double", value=1.0), + Row(var="id", value=2.0), + Row(var="int", value=20.0), + Row(var="double", value=2.0), + Row(var="id", value=3.0), + Row(var="int", value=30.0), + Row(var="double", value=3.0), + ], + ) + + with self.subTest(desc="with single identifier columns but none value columns"): + for ids in ["id", ["id"], ("id",)]: + with self.subTest(ids=ids): + # select only columns that have common data type (double) + actual = df.select("id", "int", "double").unpivot(ids, None, "var", "val") + self.assertEqual( + actual.schema.simpleString(), "struct" + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, var="int", value=10.0), + Row(id=1, var="double", value=1.0), + Row(id=2, var="int", value=20.0), + Row(id=2, var="double", value=2.0), + Row(id=3, var="int", value=30.0), + Row(id=3, var="double", value=3.0), + ], + ) + + with self.subTest(desc="with multiple identifier columns but none given value columns"): + for ids in [["id", "str"], ("id", "str")]: + with self.subTest(ids=ids): + actual = df.unpivot(ids, None, "var", "val") + self.assertEqual( + actual.schema.simpleString(), + "struct", + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, str="one", var="int", val=10.0), + Row(id=1, str="one", var="double", val=1.0), + Row(id=2, str="two", var="int", val=20.0), + Row(id=2, str="two", var="double", val=2.0), + Row(id=3, str="three", var="int", val=30.0), + Row(id=3, str="three", var="double", val=3.0), + ], + ) + + with self.subTest(desc="with single value column"): + for values in ["int", ["int"], ("int",)]: + with self.subTest(values=values): + actual = df.unpivot("id", values, "var", "val") + self.assertEqual( + actual.schema.simpleString(), "struct" + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, var="int", val=10), + Row(id=2, var="int", val=20), + Row(id=3, var="int", val=30), + ], + ) + + with self.subTest(desc="with multiple value columns"): + for values in [["int", "double"], ("int", "double")]: + with self.subTest(values=values): + actual = df.unpivot("id", values, "var", "val") + self.assertEqual( + actual.schema.simpleString(), "struct" + ) + self.assertEqual( + actual.collect(), + [ + Row(id=1, var="int", val=10.0), + Row(id=1, var="double", val=1.0), + Row(id=2, var="int", val=20.0), + Row(id=2, var="double", val=2.0), + Row(id=3, var="int", val=30.0), + Row(id=3, var="double", val=3.0), + ], + ) + + with self.subTest(desc="with columns"): + for id in [df.id, [df.id], (df.id,)]: + for values in [[df.int, df.double], (df.int, df.double)]: + with self.subTest(ids=id, values=values): + self.assertEqual( + df.unpivot(id, values, "var", "val").collect(), + df.unpivot("id", ["int", "double"], "var", "val").collect(), + ) + + with self.subTest(desc="with column names and columns"): + for ids in [[df.id, "str"], (df.id, "str")]: + for values in [[df.int, "double"], (df.int, "double")]: + with self.subTest(ids=ids, values=values): + self.assertEqual( + df.unpivot(ids, values, "var", "val").collect(), + df.unpivot(["id", "str"], ["int", "double"], "var", "val").collect(), + ) + + with self.subTest(desc="melt alias"): + self.assertEqual( + df.unpivot("id", ["int", "double"], "var", "val").collect(), + df.melt("id", ["int", "double"], "var", "val").collect(), + ) + + def test_unpivot_negative(self): + # SPARK-39877: test the DataFrame.unpivot method + df = self.spark.createDataFrame( + [ + (1, 10, 1.0, "one"), + (2, 20, 2.0, "two"), + (3, 30, 3.0, "three"), + ], + ["id", "int", "double", "str"], + ) + + with self.subTest(desc="with no value columns"): + for values in [[], ()]: + with self.subTest(values=values): + with self.assertRaisesRegex( + AnalysisException, + r"\[UNPIVOT_REQUIRES_VALUE_COLUMNS] At least one value column " + r"needs to be specified for UNPIVOT, all columns specified as ids.*", + ): + df.unpivot("id", values, "var", "val").collect() + + with self.subTest(desc="with value columns without common data type"): + with self.assertRaisesRegex( + AnalysisException, + r"\[UNPIVOT_VALUE_DATA_TYPE_MISMATCH\] Unpivot value columns must share " + r"a least common type, some types do not: .*", + ): + df.unpivot("id", ["int", "str"], "var", "val").collect() + + def test_observe(self): + # SPARK-36263: tests the DataFrame.observe(Observation, *Column) method + from pyspark.sql import Observation + + df = self.spark.createDataFrame( + [ + (1, 1.0, "one"), + (2, 2.0, "two"), + (3, 3.0, "three"), + ], + ["id", "val", "label"], + ) + + unnamed_observation = Observation() + named_observation = Observation("metric") + observed = ( + df.orderBy("id") + .observe( + named_observation, + count(lit(1)).alias("cnt"), + sum(col("id")).alias("sum"), + mean(col("val")).alias("mean"), + ) + .observe(unnamed_observation, count(lit(1)).alias("rows")) + ) + + # test that observe works transparently + actual = observed.collect() + self.assertEqual( + [ + {"id": 1, "val": 1.0, "label": "one"}, + {"id": 2, "val": 2.0, "label": "two"}, + {"id": 3, "val": 3.0, "label": "three"}, + ], + [row.asDict() for row in actual], + ) + + # test that we retrieve the metrics + self.assertEqual(named_observation.get, dict(cnt=3, sum=6, mean=2.0)) + self.assertEqual(unnamed_observation.get, dict(rows=3)) + + # observation requires name (if given) to be non empty string + with self.assertRaisesRegex(TypeError, "name should be a string"): + Observation(123) + with self.assertRaisesRegex(ValueError, "name should not be empty"): + Observation("") + + # dataframe.observe requires at least one expr + with self.assertRaises(PySparkValueError) as pe: + df.observe(Observation()) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_BE_EMPTY", + message_parameters={"item": "exprs"}, + ) + + # dataframe.observe requires non-None Columns + for args in [(None,), ("id",), (lit(1), None), (lit(1), "id")]: + with self.subTest(args=args): + with self.assertRaises(PySparkTypeError) as pe: + df.observe(Observation(), *args) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OF_COLUMN", + message_parameters={"arg_name": "exprs"}, + ) + + def test_observe_str(self): + # SPARK-38760: tests the DataFrame.observe(str, *Column) method + from pyspark.sql.streaming import StreamingQueryListener + + observed_metrics = None + + class TestListener(StreamingQueryListener): + def onQueryStarted(self, event): + pass + + def onQueryProgress(self, event): + nonlocal observed_metrics + observed_metrics = event.progress.observedMetrics + + def onQueryIdle(self, event): + pass + + def onQueryTerminated(self, event): + pass + + self.spark.streams.addListener(TestListener()) + + df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load() + df = df.observe("metric", count(lit(1)).alias("cnt"), sum(col("value")).alias("sum")) + q = df.writeStream.format("noop").queryName("test").start() + self.assertTrue(q.isActive) + time.sleep(10) + q.stop() + + self.assertTrue(isinstance(observed_metrics, dict)) + self.assertTrue("metric" in observed_metrics) + row = observed_metrics["metric"] + self.assertTrue(isinstance(row, Row)) + self.assertTrue(hasattr(row, "cnt")) + self.assertTrue(hasattr(row, "sum")) + self.assertGreaterEqual(row.cnt, 0) + self.assertGreaterEqual(row.sum, 0) + + def test_sample(self): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.range(1).sample() + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL_OR_FLOAT_OR_INT", + message_parameters={ + "arg_name": "withReplacement (optional), fraction (required) and seed (optional)", + "arg_type": "NoneType, NoneType, NoneType", + }, + ) + + self.assertRaises(TypeError, lambda: self.spark.range(1).sample("a")) + + self.assertRaises(TypeError, lambda: self.spark.range(1).sample(seed="abc")) + + self.assertRaises( + IllegalArgumentException, lambda: self.spark.range(1).sample(-1.0).count() + ) + + def test_toDF_with_string(self): + df = self.spark.createDataFrame([("John", 30), ("Alice", 25), ("Bob", 28)]) + data = [("John", 30), ("Alice", 25), ("Bob", 28)] + + result = df.toDF("key", "value") + self.assertEqual(result.schema.simpleString(), "struct") + self.assertEqual(result.collect(), data) + + with self.assertRaises(PySparkTypeError) as pe: + df.toDF("key", None) + + self.check_error( + exception=pe.exception, + error_class="NOT_LIST_OF_STR", + message_parameters={"arg_name": "cols", "arg_type": "NoneType"}, + ) + + def test_toDF_with_schema_string(self): + data = [Row(key=i, value=str(i)) for i in range(100)] + rdd = self.sc.parallelize(data, 5) + + df = rdd.toDF("key: int, value: string") + self.assertEqual(df.schema.simpleString(), "struct") + self.assertEqual(df.collect(), data) + + # different but compatible field types can be used. + df = rdd.toDF("key: string, value: string") + self.assertEqual(df.schema.simpleString(), "struct") + self.assertEqual(df.collect(), [Row(key=str(i), value=str(i)) for i in range(100)]) + + # field names can differ. + df = rdd.toDF(" a: int, b: string ") + self.assertEqual(df.schema.simpleString(), "struct") + self.assertEqual(df.collect(), data) + + # number of fields must match. + self.assertRaisesRegex( + Exception, "LENGTH_SHOULD_BE_THE_SAME", lambda: rdd.toDF("key: int").collect() + ) + + # field types mismatch will cause exception at runtime. + self.assertRaisesRegex( + Exception, + "CANNOT_ACCEPT_OBJECT_IN_TYPE", + lambda: rdd.toDF("key: float, value: string").collect(), + ) + + # flat schema values will be wrapped into row. + df = rdd.map(lambda row: row.key).toDF("int") + self.assertEqual(df.schema.simpleString(), "struct") + self.assertEqual(df.collect(), [Row(key=i) for i in range(100)]) + + # users can use DataType directly instead of data type string. + df = rdd.map(lambda row: row.key).toDF(IntegerType()) + self.assertEqual(df.schema.simpleString(), "struct") + self.assertEqual(df.collect(), [Row(key=i) for i in range(100)]) + + def test_print_schema(self): + df = self.spark.createDataFrame([(1, (2, 2))], ["a", "b"]) + + with io.StringIO() as buf, redirect_stdout(buf): + df.printSchema(1) + self.assertEqual(1, buf.getvalue().count("long")) + self.assertEqual(0, buf.getvalue().count("_1")) + self.assertEqual(0, buf.getvalue().count("_2")) + + buf.truncate(0) + buf.seek(0) + + df.printSchema(2) + self.assertEqual(3, buf.getvalue().count("long")) + self.assertEqual(1, buf.getvalue().count("_1")) + self.assertEqual(1, buf.getvalue().count("_2")) + + def test_join_without_on(self): + df1 = self.spark.range(1).toDF("a") + df2 = self.spark.range(1).toDF("b") + + with self.sql_conf({"spark.sql.crossJoin.enabled": False}): + self.assertRaises(AnalysisException, lambda: df1.join(df2, how="inner").collect()) + + with self.sql_conf({"spark.sql.crossJoin.enabled": True}): + actual = df1.join(df2, how="inner").collect() + expected = [Row(a=0, b=0)] + self.assertEqual(actual, expected) + + # Regression test for invalid join methods when on is None, Spark-14761 + def test_invalid_join_method(self): + df1 = self.spark.createDataFrame([("Alice", 5), ("Bob", 8)], ["name", "age"]) + df2 = self.spark.createDataFrame([("Alice", 80), ("Bob", 90)], ["name", "height"]) + self.assertRaises(IllegalArgumentException, lambda: df1.join(df2, how="invalid-join-type")) + + # Cartesian products require cross join syntax + def test_require_cross(self): + + df1 = self.spark.createDataFrame([(1, "1")], ("key", "value")) + df2 = self.spark.createDataFrame([(1, "1")], ("key", "value")) + + with self.sql_conf({"spark.sql.crossJoin.enabled": False}): + # joins without conditions require cross join syntax + self.assertRaises(AnalysisException, lambda: df1.join(df2).collect()) + + # works with crossJoin + self.assertEqual(1, df1.crossJoin(df2).count()) + + def test_cache_dataframe(self): + df = self.spark.createDataFrame([(2, 2), (3, 3)]) + try: + self.assertEqual(df.storageLevel, StorageLevel.NONE) + + df.cache() + self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER) + + df.unpersist() + self.assertEqual(df.storageLevel, StorageLevel.NONE) + + df.persist() + self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER) + + df.unpersist(blocking=True) + self.assertEqual(df.storageLevel, StorageLevel.NONE) + + df.persist(StorageLevel.DISK_ONLY) + self.assertEqual(df.storageLevel, StorageLevel.DISK_ONLY) + finally: + df.unpersist() + self.assertEqual(df.storageLevel, StorageLevel.NONE) + + def test_cache_table(self): + spark = self.spark + tables = ["tab1", "tab2", "tab3"] + with self.tempView(*tables): + for i, tab in enumerate(tables): + spark.createDataFrame([(2, i), (3, i)]).createOrReplaceTempView(tab) + self.assertFalse(spark.catalog.isCached(tab)) + spark.catalog.cacheTable("tab1") + spark.catalog.cacheTable("tab3", StorageLevel.OFF_HEAP) + self.assertTrue(spark.catalog.isCached("tab1")) + self.assertFalse(spark.catalog.isCached("tab2")) + self.assertTrue(spark.catalog.isCached("tab3")) + spark.catalog.cacheTable("tab2") + spark.catalog.uncacheTable("tab1") + spark.catalog.uncacheTable("tab3") + self.assertFalse(spark.catalog.isCached("tab1")) + self.assertTrue(spark.catalog.isCached("tab2")) + self.assertFalse(spark.catalog.isCached("tab3")) + spark.catalog.clearCache() + self.assertFalse(spark.catalog.isCached("tab1")) + self.assertFalse(spark.catalog.isCached("tab2")) + self.assertFalse(spark.catalog.isCached("tab3")) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.isCached("does_not_exist"), + ) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.cacheTable("does_not_exist"), + ) + self.assertRaisesRegex( + AnalysisException, + "does_not_exist", + lambda: spark.catalog.uncacheTable("does_not_exist"), + ) + + def _to_pandas(self): + from datetime import datetime, date, timedelta + + schema = ( + StructType() + .add("a", IntegerType()) + .add("b", StringType()) + .add("c", BooleanType()) + .add("d", FloatType()) + .add("dt", DateType()) + .add("ts", TimestampType()) + .add("ts_ntz", TimestampNTZType()) + .add("dt_interval", DayTimeIntervalType()) + ) + data = [ + ( + 1, + "foo", + True, + 3.0, + date(1969, 1, 1), + datetime(1969, 1, 1, 1, 1, 1), + datetime(1969, 1, 1, 1, 1, 1), + timedelta(days=1), + ), + (2, "foo", True, 5.0, None, None, None, None), + ( + 3, + "bar", + False, + -1.0, + date(2012, 3, 3), + datetime(2012, 3, 3, 3, 3, 3), + datetime(2012, 3, 3, 3, 3, 3), + timedelta(hours=-1, milliseconds=421), + ), + ( + 4, + "bar", + False, + 6.0, + date(2100, 4, 4), + datetime(2100, 4, 4, 4, 4, 4), + datetime(2100, 4, 4, 4, 4, 4), + timedelta(microseconds=123), + ), + ] + df = self.spark.createDataFrame(data, schema) + return df.toPandas() + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas(self): + import numpy as np + + pdf = self._to_pandas() + types = pdf.dtypes + self.assertEqual(types[0], np.int32) + self.assertEqual(types[1], object) + self.assertEqual(types[2], bool) + self.assertEqual(types[3], np.float32) + self.assertEqual(types[4], object) # datetime.date + self.assertEqual(types[5], "datetime64[ns]") + self.assertEqual(types[6], "datetime64[ns]") + self.assertEqual(types[7], "timedelta64[ns]") + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_with_duplicated_column_names(self): + for arrow_enabled in [False, True]: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + self.check_to_pandas_with_duplicated_column_names() + + def check_to_pandas_with_duplicated_column_names(self): + import numpy as np + + sql = "select 1 v, 1 v" + df = self.spark.sql(sql) + pdf = df.toPandas() + types = pdf.dtypes + self.assertEqual(types.iloc[0], np.int32) + self.assertEqual(types.iloc[1], np.int32) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_on_cross_join(self): + for arrow_enabled in [False, True]: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}): + self.check_to_pandas_on_cross_join() + + def check_to_pandas_on_cross_join(self): + import numpy as np + + sql = """ + select t1.*, t2.* from ( + select explode(sequence(1, 3)) v + ) t1 left join ( + select explode(sequence(1, 3)) v + ) t2 + """ + with self.sql_conf({"spark.sql.crossJoin.enabled": True}): + df = self.spark.sql(sql) + pdf = df.toPandas() + types = pdf.dtypes + self.assertEqual(types.iloc[0], np.int32) + self.assertEqual(types.iloc[1], np.int32) + + @unittest.skipIf(have_pandas, "Required Pandas was found.") + def test_to_pandas_required_pandas_not_found(self): + with QuietTest(self.sc): + with self.assertRaisesRegex(ImportError, "Pandas >= .* must be installed"): + self._to_pandas() + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_avoid_astype(self): + import numpy as np + + schema = StructType().add("a", IntegerType()).add("b", StringType()).add("c", IntegerType()) + data = [(1, "foo", 16777220), (None, "bar", None)] + df = self.spark.createDataFrame(data, schema) + types = df.toPandas().dtypes + self.assertEqual(types[0], np.float64) # doesn't convert to np.int32 due to NaN value. + self.assertEqual(types[1], object) + self.assertEqual(types[2], np.float64) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_from_empty_dataframe(self): + is_arrow_enabled = [True, False] + for value in is_arrow_enabled: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": value}): + self.check_to_pandas_from_empty_dataframe() + + def check_to_pandas_from_empty_dataframe(self): + # SPARK-29188 test that toPandas() on an empty dataframe has the correct dtypes + # SPARK-30537 test that toPandas() on an empty dataframe has the correct dtypes + # when arrow is enabled + import numpy as np + + sql = """ + SELECT CAST(1 AS TINYINT) AS tinyint, + CAST(1 AS SMALLINT) AS smallint, + CAST(1 AS INT) AS int, + CAST(1 AS BIGINT) AS bigint, + CAST(0 AS FLOAT) AS float, + CAST(0 AS DOUBLE) AS double, + CAST(1 AS BOOLEAN) AS boolean, + CAST('foo' AS STRING) AS string, + CAST('2019-01-01' AS TIMESTAMP) AS timestamp, + CAST('2019-01-01' AS TIMESTAMP_NTZ) AS timestamp_ntz, + INTERVAL '1563:04' MINUTE TO SECOND AS day_time_interval + """ + dtypes_when_nonempty_df = self.spark.sql(sql).toPandas().dtypes + dtypes_when_empty_df = self.spark.sql(sql).filter("False").toPandas().dtypes + self.assertTrue(np.all(dtypes_when_empty_df == dtypes_when_nonempty_df)) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_from_null_dataframe(self): + is_arrow_enabled = [True, False] + for value in is_arrow_enabled: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": value}): + self.check_to_pandas_from_null_dataframe() + + def check_to_pandas_from_null_dataframe(self): + # SPARK-29188 test that toPandas() on a dataframe with only nulls has correct dtypes + # SPARK-30537 test that toPandas() on a dataframe with only nulls has correct dtypes + # using arrow + import numpy as np + + sql = """ + SELECT CAST(NULL AS TINYINT) AS tinyint, + CAST(NULL AS SMALLINT) AS smallint, + CAST(NULL AS INT) AS int, + CAST(NULL AS BIGINT) AS bigint, + CAST(NULL AS FLOAT) AS float, + CAST(NULL AS DOUBLE) AS double, + CAST(NULL AS BOOLEAN) AS boolean, + CAST(NULL AS STRING) AS string, + CAST(NULL AS TIMESTAMP) AS timestamp, + CAST(NULL AS TIMESTAMP_NTZ) AS timestamp_ntz, + INTERVAL '1563:04' MINUTE TO SECOND AS day_time_interval + """ + pdf = self.spark.sql(sql).toPandas() + types = pdf.dtypes + self.assertEqual(types[0], np.float64) + self.assertEqual(types[1], np.float64) + self.assertEqual(types[2], np.float64) + self.assertEqual(types[3], np.float64) + self.assertEqual(types[4], np.float32) + self.assertEqual(types[5], np.float64) + self.assertEqual(types[6], object) + self.assertEqual(types[7], object) + self.assertTrue(np.can_cast(np.datetime64, types[8])) + self.assertTrue(np.can_cast(np.datetime64, types[9])) + self.assertTrue(np.can_cast(np.timedelta64, types[10])) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_to_pandas_from_mixed_dataframe(self): + is_arrow_enabled = [True, False] + for value in is_arrow_enabled: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": value}): + self.check_to_pandas_from_mixed_dataframe() + + def check_to_pandas_from_mixed_dataframe(self): + # SPARK-29188 test that toPandas() on a dataframe with some nulls has correct dtypes + # SPARK-30537 test that toPandas() on a dataframe with some nulls has correct dtypes + # using arrow + import numpy as np + + sql = """ + SELECT CAST(col1 AS TINYINT) AS tinyint, + CAST(col2 AS SMALLINT) AS smallint, + CAST(col3 AS INT) AS int, + CAST(col4 AS BIGINT) AS bigint, + CAST(col5 AS FLOAT) AS float, + CAST(col6 AS DOUBLE) AS double, + CAST(col7 AS BOOLEAN) AS boolean, + CAST(col8 AS STRING) AS string, + timestamp_seconds(col9) AS timestamp, + timestamp_seconds(col10) AS timestamp_ntz, + INTERVAL '1563:04' MINUTE TO SECOND AS day_time_interval + FROM VALUES (1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + """ + pdf_with_some_nulls = self.spark.sql(sql).toPandas() + pdf_with_only_nulls = self.spark.sql(sql).filter("tinyint is null").toPandas() + self.assertTrue(np.all(pdf_with_only_nulls.dtypes == pdf_with_some_nulls.dtypes)) + + @unittest.skipIf( + not have_pandas or not have_pyarrow or pyarrow_version_less_than_minimum("2.0.0"), + pandas_requirement_message + or pyarrow_requirement_message + or "Pyarrow version must be 2.0.0 or higher", + ) + def test_to_pandas_for_array_of_struct(self): + for is_arrow_enabled in [True, False]: + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": is_arrow_enabled}): + self.check_to_pandas_for_array_of_struct(is_arrow_enabled) + + def check_to_pandas_for_array_of_struct(self, is_arrow_enabled): + # SPARK-38098: Support Array of Struct for Pandas UDFs and toPandas + import numpy as np + import pandas as pd + + df = self.spark.createDataFrame( + [[[("a", 2, 3.0), ("a", 2, 3.0)]], [[("b", 5, 6.0), ("b", 5, 6.0)]]], + "array_struct_col Array>", + ) + + pdf = df.toPandas() + self.assertEqual(type(pdf), pd.DataFrame) + self.assertEqual(type(pdf["array_struct_col"]), pd.Series) + if is_arrow_enabled: + self.assertEqual(type(pdf["array_struct_col"][0]), np.ndarray) + else: + self.assertEqual(type(pdf["array_struct_col"][0]), list) + + def test_create_dataframe_from_array_of_long(self): + import array + + data = [Row(longarray=array.array("l", [-9223372036854775808, 0, 9223372036854775807]))] + df = self.spark.createDataFrame(data) + self.assertEqual(df.first(), Row(longarray=[-9223372036854775808, 0, 9223372036854775807])) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_create_dataframe_from_pandas_with_timestamp(self): + import pandas as pd + from datetime import datetime + + pdf = pd.DataFrame( + {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]}, + columns=["d", "ts"], + ) + # test types are inferred correctly without specifying schema + df = self.spark.createDataFrame(pdf) + self.assertIsInstance(df.schema["ts"].dataType, TimestampType) + self.assertIsInstance(df.schema["d"].dataType, DateType) + # test with schema will accept pdf as input + df = self.spark.createDataFrame(pdf, schema="d date, ts timestamp") + self.assertIsInstance(df.schema["ts"].dataType, TimestampType) + self.assertIsInstance(df.schema["d"].dataType, DateType) + df = self.spark.createDataFrame(pdf, schema="d date, ts timestamp_ntz") + self.assertIsInstance(df.schema["ts"].dataType, TimestampNTZType) + self.assertIsInstance(df.schema["d"].dataType, DateType) + + @unittest.skipIf(have_pandas, "Required Pandas was found.") + def test_create_dataframe_required_pandas_not_found(self): + with QuietTest(self.sc): + with self.assertRaisesRegex( + ImportError, "(Pandas >= .* must be installed|No module named '?pandas'?)" + ): + import pandas as pd + from datetime import datetime + + pdf = pd.DataFrame( + {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]} + ) + self.spark.createDataFrame(pdf) + + # Regression test for SPARK-23360 + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore + def test_create_dataframe_from_pandas_with_dst(self): + import pandas as pd + from pandas.testing import assert_frame_equal + from datetime import datetime + + pdf = pd.DataFrame({"time": [datetime(2015, 10, 31, 22, 30)]}) + + df = self.spark.createDataFrame(pdf) + assert_frame_equal(pdf, df.toPandas()) + + orig_env_tz = os.environ.get("TZ", None) + try: + tz = "America/Los_Angeles" + os.environ["TZ"] = tz + time.tzset() + with self.sql_conf({"spark.sql.session.timeZone": tz}): + df = self.spark.createDataFrame(pdf) + assert_frame_equal(pdf, df.toPandas()) + finally: + del os.environ["TZ"] + if orig_env_tz is not None: + os.environ["TZ"] = orig_env_tz + time.tzset() + + # TODO(SPARK-43354): Re-enable test_create_dataframe_from_pandas_with_day_time_interval + @unittest.skipIf( + "pypy" in platform.python_implementation().lower(), + "Fails in PyPy Python 3.8, should enable.", + ) + def test_create_dataframe_from_pandas_with_day_time_interval(self): + # SPARK-37277: Test DayTimeIntervalType in createDataFrame without Arrow. + import pandas as pd + from datetime import timedelta + + df = self.spark.createDataFrame(pd.DataFrame({"a": [timedelta(microseconds=123)]})) + self.assertEqual(df.toPandas().a.iloc[0], timedelta(microseconds=123)) + + def test_repr_behaviors(self): + import re + + pattern = re.compile(r"^ *\|", re.MULTILINE) + df = self.spark.createDataFrame([(1, "1"), (22222, "22222")], ("key", "value")) + + # test when eager evaluation is enabled and _repr_html_ will not be called + with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}): + expected1 = """+-----+-----+ + || key|value| + |+-----+-----+ + || 1| 1| + ||22222|22222| + |+-----+-----+ + |""" + self.assertEqual(re.sub(pattern, "", expected1), df.__repr__()) + with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}): + expected2 = """+---+-----+ + ||key|value| + |+---+-----+ + || 1| 1| + ||222| 222| + |+---+-----+ + |""" + self.assertEqual(re.sub(pattern, "", expected2), df.__repr__()) + with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}): + expected3 = """+---+-----+ + ||key|value| + |+---+-----+ + || 1| 1| + |+---+-----+ + |only showing top 1 row + |""" + self.assertEqual(re.sub(pattern, "", expected3), df.__repr__()) + + # test when eager evaluation is enabled and _repr_html_ will be called + with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}): + expected1 = """ + | + | + | + |
    keyvalue
    11
    2222222222
    + |""" + self.assertEqual(re.sub(pattern, "", expected1), df._repr_html_()) + with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}): + expected2 = """ + | + | + | + |
    keyvalue
    11
    222222
    + |""" + self.assertEqual(re.sub(pattern, "", expected2), df._repr_html_()) + with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}): + expected3 = """ + | + | + |
    keyvalue
    11
    + |only showing top 1 row + |""" + self.assertEqual(re.sub(pattern, "", expected3), df._repr_html_()) + + # test when eager evaluation is disabled and _repr_html_ will be called + with self.sql_conf({"spark.sql.repl.eagerEval.enabled": False}): + expected = "DataFrame[key: bigint, value: string]" + self.assertEqual(None, df._repr_html_()) + self.assertEqual(expected, df.__repr__()) + with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}): + self.assertEqual(None, df._repr_html_()) + self.assertEqual(expected, df.__repr__()) + with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}): + self.assertEqual(None, df._repr_html_()) + self.assertEqual(expected, df.__repr__()) + + def test_to_local_iterator(self): + df = self.spark.range(8, numPartitions=4) + expected = df.collect() + it = df.toLocalIterator() + self.assertEqual(expected, list(it)) + + # Test DataFrame with empty partition + df = self.spark.range(3, numPartitions=4) + it = df.toLocalIterator() + expected = df.collect() + self.assertEqual(expected, list(it)) + + def test_to_local_iterator_prefetch(self): + df = self.spark.range(8, numPartitions=4) + expected = df.collect() + it = df.toLocalIterator(prefetchPartitions=True) + self.assertEqual(expected, list(it)) + + def test_to_local_iterator_not_fully_consumed(self): + with QuietTest(self.sc): + self.check_to_local_iterator_not_fully_consumed() + + def check_to_local_iterator_not_fully_consumed(self): + # SPARK-23961: toLocalIterator throws exception when not fully consumed + # Create a DataFrame large enough so that write to socket will eventually block + df = self.spark.range(1 << 20, numPartitions=2) + it = df.toLocalIterator() + self.assertEqual(df.take(1)[0], next(it)) + it = None # remove iterator from scope, socket is closed when cleaned up + # Make sure normal df operations still work + result = [] + for i, row in enumerate(df.toLocalIterator()): + result.append(row) + if i == 7: + break + self.assertEqual(df.take(8), result) + + def test_same_semantics_error(self): + with QuietTest(self.sc): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.range(10).sameSemantics(1) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "other", "arg_type": "int"}, + ) + + def test_input_files(self): + tpath = tempfile.mkdtemp() + shutil.rmtree(tpath) + try: + self.spark.range(1, 100, 1, 10).write.parquet(tpath) + # read parquet file and get the input files list + input_files_list = self.spark.read.parquet(tpath).inputFiles() + + # input files list should contain 10 entries + self.assertEqual(len(input_files_list), 10) + # all file paths in list must contain tpath + for file_path in input_files_list: + self.assertTrue(tpath in file_path) + finally: + shutil.rmtree(tpath) + + def test_df_show(self): + # SPARK-35408: ensure better diagnostics if incorrect parameters are passed + # to DataFrame.show + + df = self.spark.createDataFrame([("foo",)]) + df.show(5) + df.show(5, True) + df.show(5, 1, True) + df.show(n=5, truncate="1", vertical=False) + df.show(n=5, truncate=1.5, vertical=False) + + with self.assertRaises(PySparkTypeError) as pe: + df.show(True) + + self.check_error( + exception=pe.exception, + error_class="NOT_INT", + message_parameters={"arg_name": "n", "arg_type": "bool"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.show(vertical="foo") + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL", + message_parameters={"arg_name": "vertical", "arg_type": "str"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.show(truncate="foo") + + self.check_error( + exception=pe.exception, + error_class="NOT_BOOL", + message_parameters={"arg_name": "truncate", "arg_type": "str"}, + ) + + @unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), + ) + def test_pandas_api(self): + import pandas as pd + from pandas.testing import assert_frame_equal + + sdf = self.spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["Col1", "Col2"]) + psdf_from_sdf = sdf.pandas_api() + psdf_from_sdf_with_index = sdf.pandas_api(index_col="Col1") + pdf = pd.DataFrame({"Col1": ["a", "b", "c"], "Col2": [1, 2, 3]}) + pdf_with_index = pdf.set_index("Col1") + + assert_frame_equal(pdf, psdf_from_sdf.to_pandas()) + assert_frame_equal(pdf_with_index, psdf_from_sdf_with_index.to_pandas()) + + # test for SPARK-36337 + def test_create_nan_decimal_dataframe(self): + self.assertEqual( + self.spark.createDataFrame(data=[Decimal("NaN")], schema="decimal").collect(), + [Row(value=None)], + ) + + def test_to(self): + schema = StructType( + [StructField("i", StringType(), True), StructField("j", IntegerType(), True)] + ) + df = self.spark.createDataFrame([("a", 1)], schema) + + schema1 = StructType([StructField("j", StringType()), StructField("i", StringType())]) + df1 = df.to(schema1) + self.assertEqual(schema1, df1.schema) + self.assertEqual(df.count(), df1.count()) + + schema2 = StructType([StructField("j", LongType())]) + df2 = df.to(schema2) + self.assertEqual(schema2, df2.schema) + self.assertEqual(df.count(), df2.count()) + + schema3 = StructType([StructField("struct", schema1, False)]) + df3 = df.select(struct("i", "j").alias("struct")).to(schema3) + self.assertEqual(schema3, df3.schema) + self.assertEqual(df.count(), df3.count()) + + # incompatible field nullability + schema4 = StructType([StructField("j", LongType(), False)]) + self.assertRaisesRegex( + AnalysisException, "NULLABLE_COLUMN_OR_FIELD", lambda: df.to(schema4).count() + ) + + # field cannot upcast + schema5 = StructType([StructField("i", LongType())]) + self.assertRaisesRegex( + AnalysisException, "INVALID_COLUMN_OR_FIELD_DATA_TYPE", lambda: df.to(schema5).count() + ) + + def test_repartition(self): + df = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + with self.assertRaises(PySparkTypeError) as pe: + df.repartition([10], "name", "age").rdd.getNumPartitions() + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "numPartitions", "arg_type": "list"}, + ) + + def test_colregex(self): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.range(10).colRegex(10) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "colName", "arg_type": "int"}, + ) + + def test_where(self): + with self.assertRaises(PySparkTypeError) as pe: + self.spark.range(10).where(10) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "condition", "arg_type": "int"}, + ) + + def test_duplicate_field_names(self): + data = [ + Row(Row("a", 1), Row(2, 3, "b", 4, "c", "d")), + Row(Row("w", 6), Row(7, 8, "x", 9, "y", "z")), + ] + schema = ( + StructType() + .add("struct", StructType().add("x", StringType()).add("x", IntegerType())) + .add( + "struct", + StructType() + .add("a", IntegerType()) + .add("x", IntegerType()) + .add("x", StringType()) + .add("y", IntegerType()) + .add("y", StringType()) + .add("x", StringType()), + ) + ) + df = self.spark.createDataFrame(data, schema=schema) + + self.assertEqual(df.schema, schema) + self.assertEqual(df.collect(), data) + + +class QueryExecutionListenerTests(unittest.TestCase, SQLTestUtils): + # These tests are separate because it uses 'spark.sql.queryExecutionListeners' which is + # static and immutable. This can't be set or unset, for example, via `spark.conf`. + + @classmethod + def setUpClass(cls): + import glob + from pyspark.find_spark_home import _find_spark_home + + SPARK_HOME = _find_spark_home() + filename_pattern = ( + "sql/core/target/scala-*/test-classes/org/apache/spark/sql/" + "TestQueryExecutionListener.class" + ) + cls.has_listener = bool(glob.glob(os.path.join(SPARK_HOME, filename_pattern))) + + if cls.has_listener: + # Note that 'spark.sql.queryExecutionListeners' is a static immutable configuration. + cls.spark = ( + SparkSession.builder.master("local[4]") + .appName(cls.__name__) + .config( + "spark.sql.queryExecutionListeners", + "org.apache.spark.sql.TestQueryExecutionListener", + ) + .getOrCreate() + ) + + def setUp(self): + if not self.has_listener: + raise self.skipTest( + "'org.apache.spark.sql.TestQueryExecutionListener' is not " + "available. Will skip the related tests." + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "spark"): + cls.spark.stop() + + def tearDown(self): + self.spark._jvm.OnSuccessCall.clear() + + def test_query_execution_listener_on_collect(self): + self.assertFalse( + self.spark._jvm.OnSuccessCall.isCalled(), + "The callback from the query execution listener should not be called before 'collect'", + ) + self.spark.sql("SELECT * FROM range(1)").collect() + self.spark.sparkContext._jsc.sc().listenerBus().waitUntilEmpty(10000) + self.assertTrue( + self.spark._jvm.OnSuccessCall.isCalled(), + "The callback from the query execution listener should be called after 'collect'", + ) + + @unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), + ) + def test_query_execution_listener_on_collect_with_arrow(self): + with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): + self.assertFalse( + self.spark._jvm.OnSuccessCall.isCalled(), + "The callback from the query execution listener should not be " + "called before 'toPandas'", + ) + self.spark.sql("SELECT * FROM range(1)").toPandas() + self.spark.sparkContext._jsc.sc().listenerBus().waitUntilEmpty(10000) + self.assertTrue( + self.spark._jvm.OnSuccessCall.isCalled(), + "The callback from the query execution listener should be called after 'toPandas'", + ) + + +class DataFrameTests(DataFrameTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.test_dataframe import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_datasources.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_datasources.py new file mode 100644 index 0000000..6418983 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_datasources.py @@ -0,0 +1,252 @@ +# +# 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 shutil +import tempfile +import uuid + +from pyspark.sql import Row +from pyspark.sql.types import IntegerType, StructField, StructType, LongType, StringType +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class DataSourcesTestsMixin: + def test_linesep_text(self): + df = self.spark.read.text("python/test_support/sql/ages_newlines.csv", lineSep=",") + expected = [ + Row(value="Joe"), + Row(value="20"), + Row(value='"Hi'), + Row(value='\nI am Jeo"\nTom'), + Row(value="30"), + Row(value='"My name is Tom"\nHyukjin'), + Row(value="25"), + Row(value='"I am Hyukjin\n\nI love Spark!"\n'), + ] + self.assertEqual(df.collect(), expected) + + tpath = tempfile.mkdtemp() + shutil.rmtree(tpath) + try: + df.write.text(tpath, lineSep="!") + expected = [ + Row(value='Joe!20!"Hi!'), + Row(value='I am Jeo"'), + Row(value='Tom!30!"My name is Tom"'), + Row(value='Hyukjin!25!"I am Hyukjin'), + Row(value=""), + Row(value='I love Spark!"'), + Row(value="!"), + ] + readback = self.spark.read.text(tpath) + self.assertEqual(readback.collect(), expected) + finally: + shutil.rmtree(tpath) + + def test_multiline_json(self): + people1 = self.spark.read.json("python/test_support/sql/people.json") + people_array = self.spark.read.json( + "python/test_support/sql/people_array.json", multiLine=True + ) + self.assertEqual(people1.collect(), people_array.collect()) + + def test_encoding_json(self): + people_array = self.spark.read.json( + "python/test_support/sql/people_array_utf16le.json", multiLine=True, encoding="UTF-16LE" + ) + expected = [Row(age=30, name="Andy"), Row(age=19, name="Justin")] + self.assertEqual(people_array.collect(), expected) + + def test_linesep_json(self): + df = self.spark.read.json("python/test_support/sql/people.json", lineSep=",") + expected = [ + Row(_corrupt_record=None, name="Michael"), + Row(_corrupt_record=' "age":30}\n{"name":"Justin"', name=None), + Row(_corrupt_record=' "age":19}\n', name=None), + ] + self.assertEqual(df.collect(), expected) + + tpath = tempfile.mkdtemp() + shutil.rmtree(tpath) + try: + df = self.spark.read.json("python/test_support/sql/people.json") + df.write.json(tpath, lineSep="!!") + readback = self.spark.read.json(tpath, lineSep="!!") + self.assertEqual(readback.collect(), df.collect()) + finally: + shutil.rmtree(tpath) + + def test_multiline_csv(self): + ages_newlines = self.spark.read.csv( + "python/test_support/sql/ages_newlines.csv", multiLine=True + ) + expected = [ + Row(_c0="Joe", _c1="20", _c2="Hi,\nI am Jeo"), + Row(_c0="Tom", _c1="30", _c2="My name is Tom"), + Row(_c0="Hyukjin", _c1="25", _c2="I am Hyukjin\n\nI love Spark!"), + ] + self.assertEqual(ages_newlines.collect(), expected) + + def test_ignorewhitespace_csv(self): + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + self.spark.createDataFrame([[" a", "b ", " c "]]).write.csv( + tmpPath, ignoreLeadingWhiteSpace=False, ignoreTrailingWhiteSpace=False + ) + + expected = [Row(value=" a,b , c ")] + readback = self.spark.read.text(tmpPath) + self.assertEqual(readback.collect(), expected) + shutil.rmtree(tmpPath) + + def test_read_multiple_orc_file(self): + df = self.spark.read.orc( + [ + "python/test_support/sql/orc_partitioned/b=0/c=0", + "python/test_support/sql/orc_partitioned/b=1/c=1", + ] + ) + self.assertEqual(2, df.count()) + + def test_read_text_file_list(self): + df = self.spark.read.text( + ["python/test_support/sql/text-test.txt", "python/test_support/sql/text-test.txt"] + ) + count = df.count() + self.assertEqual(count, 4) + + def test_json_sampling_ratio(self): + rdd = self.spark.sparkContext.range(0, 100, 1, 1).map( + lambda x: '{"a":0.1}' if x == 1 else '{"a":%s}' % str(x) + ) + schema = ( + self.spark.read.option("inferSchema", True) + .option("samplingRatio", 0.5) + .json(rdd) + .schema + ) + self.assertEqual(schema, StructType([StructField("a", LongType(), True)])) + + def test_csv_sampling_ratio(self): + rdd = self.spark.sparkContext.range(0, 100, 1, 1).map(lambda x: "0.1" if x == 1 else str(x)) + schema = self.spark.read.option("inferSchema", True).csv(rdd, samplingRatio=0.5).schema + self.assertEqual(schema, StructType([StructField("_c0", IntegerType(), True)])) + + def test_checking_csv_header(self): + path = tempfile.mkdtemp() + shutil.rmtree(path) + try: + self.spark.createDataFrame([[1, 1000], [2000, 2]]).toDF("f1", "f2").write.option( + "header", "true" + ).csv(path) + schema = StructType( + [ + StructField("f2", IntegerType(), nullable=True), + StructField("f1", IntegerType(), nullable=True), + ] + ) + df = ( + self.spark.read.option("header", "true") + .schema(schema) + .csv(path, enforceSchema=False) + ) + self.assertRaisesRegex( + Exception, "CSV header does not conform to the schema", lambda: df.collect() + ) + finally: + shutil.rmtree(path) + + def test_ignore_column_of_all_nulls(self): + path = tempfile.mkdtemp() + shutil.rmtree(path) + try: + df = self.spark.createDataFrame( + [ + ["""{"a":null, "b":1, "c":3.0}"""], + ["""{"a":null, "b":null, "c":"string"}"""], + ["""{"a":null, "b":null, "c":null}"""], + ] + ) + df.write.text(path) + schema = StructType( + [ + StructField("b", LongType(), nullable=True), + StructField("c", StringType(), nullable=True), + ] + ) + readback = self.spark.read.json(path, dropFieldIfAllNull=True) + self.assertEqual(readback.schema, schema) + finally: + shutil.rmtree(path) + + def test_jdbc(self): + db = f"memory:{uuid.uuid4()}" + url = f"jdbc:derby:{db}" + dbtable = "test_table" + + try: + df = self.spark.range(10) + df.write.jdbc(url=f"{url};create=true", table=dbtable) + + readback = self.spark.read.jdbc(url=url, table=dbtable) + self.assertEqual(sorted(df.collect()), sorted(readback.collect())) + + additional_arguments = dict(column="id", lowerBound=3, upperBound=8, numPartitions=10) + readback = self.spark.read.jdbc(url=url, table=dbtable, **additional_arguments) + self.assertEqual(sorted(df.collect()), sorted(readback.collect())) + + additional_arguments = dict(predicates=['"id" < 5']) + readback = self.spark.read.jdbc(url=url, table=dbtable, **additional_arguments) + self.assertEqual(sorted(df.filter("id < 5").collect()), sorted(readback.collect())) + finally: + # Clean up. + with self.assertRaisesRegex(Exception, f"Database '{db}' dropped."): + self.spark.read.jdbc(url=f"{url};drop=true", table=dbtable).collect() + + def test_jdbc_format(self): + db = f"memory:{uuid.uuid4()}" + url = f"jdbc:derby:{db}" + dbtable = "test_table" + + try: + df = self.spark.range(10) + df.write.format("jdbc").options(url=f"{url};create=true", dbtable=dbtable).save() + readback = self.spark.read.format("jdbc").options(url=url, dbtable=dbtable).load() + self.assertEqual(sorted(df.collect()), sorted(readback.collect())) + finally: + # Clean up. + with self.assertRaisesRegex(Exception, f"Database '{db}' dropped."): + self.spark.read.format("jdbc").options( + url=f"{url};drop=true", dbtable=dbtable + ).load().collect() + + +class DataSourcesTests(DataSourcesTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_datasources import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_errors.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_errors.py new file mode 100644 index 0000000..2ae6ef5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_errors.py @@ -0,0 +1,69 @@ +# +# 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.errors import ( + ArithmeticException, + ArrayIndexOutOfBoundsException, + DateTimeException, + NumberFormatException, + SparkRuntimeException, +) +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class ErrorsTestsMixin: + def test_arithmetic_exception(self): + with self.assertRaises(ArithmeticException): + with self.sql_conf({"spark.sql.ansi.enabled": True}): + self.spark.sql("select 1/0").show() + + def test_array_index_out_of_bounds_exception(self): + with self.assertRaises(ArrayIndexOutOfBoundsException): + with self.sql_conf({"spark.sql.ansi.enabled": True}): + self.spark.sql("select array(1, 2)[2]").show() + + def test_date_time_exception(self): + with self.assertRaises(DateTimeException): + with self.sql_conf({"spark.sql.ansi.enabled": True}): + self.spark.sql("select unix_timestamp('2023-01-01', 'dd-MM-yyyy')").show() + + def test_number_format_exception(self): + with self.assertRaises(NumberFormatException): + with self.sql_conf({"spark.sql.ansi.enabled": True}): + self.spark.sql("select cast('abc' as double)").show() + + def test_spark_runtime_exception(self): + with self.assertRaises(SparkRuntimeException): + with self.sql_conf({"spark.sql.ansi.enabled": True}): + self.spark.sql("select cast('abc' as boolean)").show() + + +class ErrorsTests(ReusedSQLTestCase, ErrorsTestsMixin): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_errors import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_functions.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_functions.py new file mode 100644 index 0000000..d80a3cf --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_functions.py @@ -0,0 +1,1349 @@ +# +# 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 contextlib import redirect_stdout +import datetime +from inspect import getmembers, isfunction +import io +from itertools import chain +import math +import re +import unittest + +from py4j.protocol import Py4JJavaError + +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql import Row, Window, functions as F, types +from pyspark.sql.column import Column +from pyspark.testing.sqlutils import ReusedSQLTestCase, SQLTestUtils +from pyspark.testing.utils import have_numpy + + +class FunctionsTestsMixin: + def test_function_parity(self): + # This test compares the available list of functions in pyspark.sql.functions with those + # available in the Scala/Java DataFrame API in org.apache.spark.sql.functions. + # + # NOTE FOR DEVELOPERS: + # If this test fails one of the following needs to happen + # * If a function was added to org.apache.spark.sql.functions it either needs to be added to + # pyspark.sql.functions or added to the below expected_missing_in_py set. + # * If a function was added to pyspark.sql.functions that was already in + # org.apache.spark.sql.functions then it needs to be removed from expected_missing_in_py + # below. If the function has a different name it needs to be added to py_equiv_jvm + # mapping. + # * If it's not related to an added/removed function then likely the exclusion list + # jvm_excluded_fn needs to be updated. + + jvm_fn_set = {name for (name, value) in getmembers(self.sc._jvm.functions)} + py_fn_set = {name for (name, value) in getmembers(F, isfunction) if name[0] != "_"} + + # Functions on the JVM side we do not expect to be available in python because they are + # depreciated, irrelevant to python, or have equivalents. + jvm_excluded_fn = [ + "callUDF", # depreciated, use call_udf + "typedlit", # Scala only + "typedLit", # Scala only + "monotonicallyIncreasingId", # depreciated, use monotonically_increasing_id + "not", # equivalent to python ~expression + "any", # equivalent to python ~some + "len", # equivalent to python ~length + "udaf", # used for creating UDAF's which are not supported in PySpark + "random", # namespace conflict with python built-in module + "uuid", # namespace conflict with python built-in module + "chr", # namespace conflict with python built-in function + ] + + jvm_fn_set.difference_update(jvm_excluded_fn) + + # For functions that are named differently in pyspark this is the mapping of their + # python name to the JVM equivalent + py_equiv_jvm = {"create_map": "map"} + for py_name, jvm_name in py_equiv_jvm.items(): + if py_name in py_fn_set: + py_fn_set.remove(py_name) + py_fn_set.add(jvm_name) + + missing_in_py = jvm_fn_set.difference(py_fn_set) + + # Functions that we expect to be missing in python until they are added to pyspark + expected_missing_in_py = set() + + self.assertEqual( + expected_missing_in_py, missing_in_py, "Missing functions in pyspark not as expected" + ) + + def test_explode(self): + d = [ + Row(a=1, intlist=[1, 2, 3], mapfield={"a": "b"}), + Row(a=1, intlist=[], mapfield={}), + Row(a=1, intlist=None, mapfield=None), + ] + data = self.spark.createDataFrame(d) + + result = data.select(F.explode(data.intlist).alias("a")).select("a").collect() + self.assertEqual(result[0][0], 1) + self.assertEqual(result[1][0], 2) + self.assertEqual(result[2][0], 3) + + result = data.select(F.explode(data.mapfield).alias("a", "b")).select("a", "b").collect() + self.assertEqual(result[0][0], "a") + self.assertEqual(result[0][1], "b") + + result = [tuple(x) for x in data.select(F.posexplode_outer("intlist")).collect()] + self.assertEqual(result, [(0, 1), (1, 2), (2, 3), (None, None), (None, None)]) + + result = [tuple(x) for x in data.select(F.posexplode_outer("mapfield")).collect()] + self.assertEqual(result, [(0, "a", "b"), (None, None, None), (None, None, None)]) + + result = [x[0] for x in data.select(F.explode_outer("intlist")).collect()] + self.assertEqual(result, [1, 2, 3, None, None]) + + result = [tuple(x) for x in data.select(F.explode_outer("mapfield")).collect()] + self.assertEqual(result, [("a", "b"), (None, None), (None, None)]) + + def test_inline(self): + d = [ + Row(structlist=[Row(b=1, c=2), Row(b=3, c=4)]), + Row(structlist=[Row(b=None, c=5), None]), + Row(structlist=[]), + ] + data = self.spark.createDataFrame(d) + + result = [tuple(x) for x in data.select(F.inline(data.structlist)).collect()] + self.assertEqual(result, [(1, 2), (3, 4), (None, 5), (None, None)]) + + result = [tuple(x) for x in data.select(F.inline_outer(data.structlist)).collect()] + self.assertEqual(result, [(1, 2), (3, 4), (None, 5), (None, None), (None, None)]) + + def test_basic_functions(self): + rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}']) + df = self.spark.read.json(rdd) + df.count() + df.collect() + df.schema + + # cache and checkpoint + self.assertFalse(df.is_cached) + df.persist() + df.unpersist(True) + df.cache() + self.assertTrue(df.is_cached) + self.assertEqual(2, df.count()) + + with self.tempView("temp"): + df.createOrReplaceTempView("temp") + df = self.spark.sql("select foo from temp") + df.count() + df.collect() + + def test_corr(self): + df = self.spark.createDataFrame([Row(a=i, b=math.sqrt(i)) for i in range(10)]) + corr = df.stat.corr("a", "b") + self.assertTrue(abs(corr - 0.95734012) < 1e-6) + + def test_sampleby(self): + df = self.spark.createDataFrame([Row(a=i, b=(i % 3)) for i in range(100)]) + sampled = df.stat.sampleBy("b", fractions={0: 0.5, 1: 0.5}, seed=0) + self.assertTrue(35 <= sampled.count() <= 36) + + with self.assertRaises(PySparkTypeError) as pe: + df.sampleBy(10, fractions={0: 0.5, 1: 0.5}) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "col", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.sampleBy("b", fractions=[0.5, 0.5]) + + self.check_error( + exception=pe.exception, + error_class="NOT_DICT", + message_parameters={"arg_name": "fractions", "arg_type": "list"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.sampleBy("b", fractions={None: 0.5, 1: 0.5}) + + self.check_error( + exception=pe.exception, + error_class="DISALLOWED_TYPE_FOR_CONTAINER", + message_parameters={ + "arg_name": "fractions", + "arg_type": "dict", + "allowed_types": "float, int, str", + "return_type": "NoneType", + }, + ) + + def test_cov(self): + df = self.spark.createDataFrame([Row(a=i, b=2 * i) for i in range(10)]) + cov = df.stat.cov("a", "b") + self.assertTrue(abs(cov - 55.0 / 3) < 1e-6) + + with self.assertRaises(PySparkTypeError) as pe: + df.stat.cov(10, "b") + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "col1", "arg_type": "int"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.stat.cov("a", True) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "col2", "arg_type": "bool"}, + ) + + def test_crosstab(self): + df = self.spark.createDataFrame([Row(a=i % 3, b=i % 2) for i in range(1, 7)]) + ct = df.stat.crosstab("a", "b").collect() + ct = sorted(ct, key=lambda x: x[0]) + for i, row in enumerate(ct): + self.assertEqual(row[0], str(i)) + self.assertTrue(row[1], 1) + self.assertTrue(row[2], 1) + + def test_math_functions(self): + df = self.spark.createDataFrame([Row(a=i, b=2 * i) for i in range(10)]) + + SQLTestUtils.assert_close( + [math.cos(i) for i in range(10)], df.select(F.cos(df.a)).collect() + ) + SQLTestUtils.assert_close([math.cos(i) for i in range(10)], df.select(F.cos("a")).collect()) + SQLTestUtils.assert_close( + [math.sin(i) for i in range(10)], df.select(F.sin(df.a)).collect() + ) + SQLTestUtils.assert_close( + [math.sin(i) for i in range(10)], df.select(F.sin(df["a"])).collect() + ) + SQLTestUtils.assert_close( + [math.pow(i, 2 * i) for i in range(10)], df.select(F.pow(df.a, df.b)).collect() + ) + SQLTestUtils.assert_close( + [math.pow(i, 2) for i in range(10)], df.select(F.pow(df.a, 2)).collect() + ) + SQLTestUtils.assert_close( + [math.pow(i, 2) for i in range(10)], df.select(F.pow(df.a, 2.0)).collect() + ) + SQLTestUtils.assert_close( + [math.hypot(i, 2 * i) for i in range(10)], df.select(F.hypot(df.a, df.b)).collect() + ) + SQLTestUtils.assert_close( + [math.hypot(i, 2 * i) for i in range(10)], df.select(F.hypot("a", "b")).collect() + ) + SQLTestUtils.assert_close( + [math.hypot(i, 2) for i in range(10)], df.select(F.hypot("a", 2)).collect() + ) + SQLTestUtils.assert_close( + [math.hypot(i, 2) for i in range(10)], df.select(F.hypot(df.a, 2)).collect() + ) + + def test_inverse_trig_functions(self): + df = self.spark.createDataFrame([Row(a=i * 0.2, b=i * -0.2) for i in range(10)]) + + def check(trig, inv, y_axis_symmetrical): + SQLTestUtils.assert_close( + [n * 0.2 for n in range(10)], + df.select(inv(trig(df.a))).collect(), + ) + if y_axis_symmetrical: + SQLTestUtils.assert_close( + [n * 0.2 for n in range(10)], + df.select(inv(trig(df.b))).collect(), + ) + else: + SQLTestUtils.assert_close( + [n * -0.2 for n in range(10)], + df.select(inv(trig(df.b))).collect(), + ) + + check(F.cosh, F.acosh, y_axis_symmetrical=True) + check(F.sinh, F.asinh, y_axis_symmetrical=False) + check(F.tanh, F.atanh, y_axis_symmetrical=False) + + def test_reciprocal_trig_functions(self): + # SPARK-36683: Tests for reciprocal trig functions (SEC, CSC and COT) + lst = [ + 0.0, + math.pi / 6, + math.pi / 4, + math.pi / 3, + math.pi / 2, + math.pi, + 3 * math.pi / 2, + 2 * math.pi, + ] + + df = self.spark.createDataFrame(lst, types.DoubleType()) + + def to_reciprocal_trig(func): + return [1.0 / func(i) if func(i) != 0 else math.inf for i in lst] + + SQLTestUtils.assert_close( + to_reciprocal_trig(math.cos), df.select(F.sec(df.value)).collect() + ) + SQLTestUtils.assert_close( + to_reciprocal_trig(math.sin), df.select(F.csc(df.value)).collect() + ) + SQLTestUtils.assert_close( + to_reciprocal_trig(math.tan), df.select(F.cot(df.value)).collect() + ) + + def test_rand_functions(self): + df = self.spark.createDataFrame([Row(key=i, value=str(i)) for i in range(100)]) + + rnd = df.select("key", F.rand()).collect() + for row in rnd: + assert row[1] >= 0.0 and row[1] <= 1.0, "got: %s" % row[1] + rndn = df.select("key", F.randn(5)).collect() + for row in rndn: + assert row[1] >= -4.0 and row[1] <= 4.0, "got: %s" % row[1] + + # If the specified seed is 0, we should use it. + # https://issues.apache.org/jira/browse/SPARK-9691 + rnd1 = df.select("key", F.rand(0)).collect() + rnd2 = df.select("key", F.rand(0)).collect() + self.assertEqual(sorted(rnd1), sorted(rnd2)) + + rndn1 = df.select("key", F.randn(0)).collect() + rndn2 = df.select("key", F.randn(0)).collect() + self.assertEqual(sorted(rndn1), sorted(rndn2)) + + def test_string_functions(self): + string_functions = [ + "upper", + "lower", + "ascii", + "base64", + "unbase64", + "ltrim", + "rtrim", + "trim", + ] + + df = self.spark.createDataFrame([["nick"]], schema=["name"]) + with self.assertRaises(PySparkTypeError) as pe: + df.select(F.col("name").substr(0, F.lit(1))) + + self.check_error( + exception=pe.exception, + error_class="NOT_SAME_TYPE", + message_parameters={ + "arg_name1": "startPos", + "arg_name2": "length", + "arg_type1": "int", + "arg_type2": "Column", + }, + ) + + for name in string_functions: + self.assertEqual( + df.select(getattr(F, name)("name")).first()[0], + df.select(getattr(F, name)(F.col("name"))).first()[0], + ) + + def test_octet_length_function(self): + # SPARK-36751: add octet length api for python + df = self.spark.createDataFrame([("cat",), ("\U0001F408",)], ["cat"]) + actual = df.select(F.octet_length("cat")).collect() + self.assertEqual([Row(3), Row(4)], actual) + + def test_bit_length_function(self): + # SPARK-36751: add bit length api for python + df = self.spark.createDataFrame([("cat",), ("\U0001F408",)], ["cat"]) + actual = df.select(F.bit_length("cat")).collect() + self.assertEqual([Row(24), Row(32)], actual) + + def test_array_contains_function(self): + df = self.spark.createDataFrame([(["1", "2", "3"],), ([],)], ["data"]) + actual = df.select(F.array_contains(df.data, "1").alias("b")).collect() + self.assertEqual([Row(b=True), Row(b=False)], actual) + + def test_levenshtein_function(self): + df = self.spark.createDataFrame([("kitten", "sitting")], ["l", "r"]) + actual_without_threshold = df.select(F.levenshtein(df.l, df.r).alias("b")).collect() + self.assertEqual([Row(b=3)], actual_without_threshold) + actual_with_threshold = df.select(F.levenshtein(df.l, df.r, 2).alias("b")).collect() + self.assertEqual([Row(b=-1)], actual_with_threshold) + + def test_between_function(self): + df = self.spark.createDataFrame( + [Row(a=1, b=2, c=3), Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)] + ) + self.assertEqual( + [Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)], df.filter(df.a.between(df.b, df.c)).collect() + ) + + def test_dayofweek(self): + dt = datetime.datetime(2017, 11, 6) + df = self.spark.createDataFrame([Row(date=dt)]) + row = df.select(F.dayofweek(df.date)).first() + self.assertEqual(row[0], 2) + + # Test added for SPARK-37738; change Python API to accept both col & int as input + def test_date_add_function(self): + dt = datetime.date(2021, 12, 27) + + # Note; number var in Python gets converted to LongType column; + # this is not supported by the function, so cast to Integer explicitly + df = self.spark.createDataFrame([Row(date=dt, add=2)], "date date, add integer") + + self.assertTrue( + all( + df.select( + F.date_add(df.date, df.add) == datetime.date(2021, 12, 29), + F.date_add(df.date, "add") == datetime.date(2021, 12, 29), + F.date_add(df.date, 3) == datetime.date(2021, 12, 30), + ).first() + ) + ) + + # Test added for SPARK-37738; change Python API to accept both col & int as input + def test_date_sub_function(self): + dt = datetime.date(2021, 12, 27) + + # Note; number var in Python gets converted to LongType column; + # this is not supported by the function, so cast to Integer explicitly + df = self.spark.createDataFrame([Row(date=dt, sub=2)], "date date, sub integer") + + self.assertTrue( + all( + df.select( + F.date_sub(df.date, df.sub) == datetime.date(2021, 12, 25), + F.date_sub(df.date, "sub") == datetime.date(2021, 12, 25), + F.date_sub(df.date, 3) == datetime.date(2021, 12, 24), + ).first() + ) + ) + + # Test added for SPARK-37738; change Python API to accept both col & int as input + def test_add_months_function(self): + dt = datetime.date(2021, 12, 27) + + # Note; number in Python gets converted to LongType column; + # this is not supported by the function, so cast to Integer explicitly + df = self.spark.createDataFrame([Row(date=dt, add=2)], "date date, add integer") + + self.assertTrue( + all( + df.select( + F.add_months(df.date, df.add) == datetime.date(2022, 2, 27), + F.add_months(df.date, "add") == datetime.date(2022, 2, 27), + F.add_months(df.date, 3) == datetime.date(2022, 3, 27), + ).first() + ) + ) + + def test_make_date(self): + # SPARK-36554: expose make_date expression + df = self.spark.createDataFrame([(2020, 6, 26)], ["Y", "M", "D"]) + row_from_col = df.select(F.make_date(df.Y, df.M, df.D)).first() + self.assertEqual(row_from_col[0], datetime.date(2020, 6, 26)) + row_from_name = df.select(F.make_date("Y", "M", "D")).first() + self.assertEqual(row_from_name[0], datetime.date(2020, 6, 26)) + + def test_expr(self): + row = Row(a="length string", b=75) + df = self.spark.createDataFrame([row]) + result = df.select(F.expr("length(a)")).collect()[0].asDict() + self.assertEqual(13, result["length(a)"]) + + # add test for SPARK-10577 (test broadcast join hint) + def test_functions_broadcast(self): + df1 = self.spark.createDataFrame([(1, "1"), (2, "2")], ("key", "value")) + df2 = self.spark.createDataFrame([(1, "1"), (2, "2")], ("key", "value")) + + # equijoin - should be converted into broadcast join + with io.StringIO() as buf, redirect_stdout(buf): + df1.join(F.broadcast(df2), "key").explain(True) + self.assertGreaterEqual(buf.getvalue().count("Broadcast"), 1) + + # no join key -- should not be a broadcast join + with io.StringIO() as buf, redirect_stdout(buf): + df1.crossJoin(F.broadcast(df2)).explain(True) + self.assertGreaterEqual(buf.getvalue().count("Broadcast"), 1) + + # planner should not crash without a join + F.broadcast(df1).explain(True) + + def test_first_last_ignorenulls(self): + df = self.spark.range(0, 100) + df2 = df.select(F.when(df.id % 3 == 0, None).otherwise(df.id).alias("id")) + df3 = df2.select( + F.first(df2.id, False).alias("a"), + F.first(df2.id, True).alias("b"), + F.last(df2.id, False).alias("c"), + F.last(df2.id, True).alias("d"), + ) + self.assertEqual([Row(a=None, b=1, c=None, d=98)], df3.collect()) + + def test_approxQuantile(self): + df = self.spark.createDataFrame([Row(a=i, b=i + 10) for i in range(10)]) + for f in ["a", "a"]: + aq = df.stat.approxQuantile(f, [0.1, 0.5, 0.9], 0.1) + self.assertTrue(isinstance(aq, list)) + self.assertEqual(len(aq), 3) + self.assertTrue(all(isinstance(q, float) for q in aq)) + aqs = df.stat.approxQuantile(["a", "b"], [0.1, 0.5, 0.9], 0.1) + self.assertTrue(isinstance(aqs, list)) + self.assertEqual(len(aqs), 2) + self.assertTrue(isinstance(aqs[0], list)) + self.assertEqual(len(aqs[0]), 3) + self.assertTrue(all(isinstance(q, float) for q in aqs[0])) + self.assertTrue(isinstance(aqs[1], list)) + self.assertEqual(len(aqs[1]), 3) + self.assertTrue(all(isinstance(q, float) for q in aqs[1])) + aqt = df.stat.approxQuantile(("a", "b"), [0.1, 0.5, 0.9], 0.1) + self.assertTrue(isinstance(aqt, list)) + self.assertEqual(len(aqt), 2) + self.assertTrue(isinstance(aqt[0], list)) + self.assertEqual(len(aqt[0]), 3) + self.assertTrue(all(isinstance(q, float) for q in aqt[0])) + self.assertTrue(isinstance(aqt[1], list)) + self.assertEqual(len(aqt[1]), 3) + self.assertTrue(all(isinstance(q, float) for q in aqt[1])) + self.assertRaises(TypeError, lambda: df.stat.approxQuantile(123, [0.1, 0.9], 0.1)) + self.assertRaises(TypeError, lambda: df.stat.approxQuantile(("a", 123), [0.1, 0.9], 0.1)) + self.assertRaises(TypeError, lambda: df.stat.approxQuantile(["a", 123], [0.1, 0.9], 0.1)) + + def test_sorting_functions_with_column(self): + self.check_sorting_functions_with_column(Column) + + def check_sorting_functions_with_column(self, tpe): + funs = [F.asc_nulls_first, F.asc_nulls_last, F.desc_nulls_first, F.desc_nulls_last] + exprs = [F.col("x"), "x"] + + for fun in funs: + for _expr in exprs: + res = fun(_expr) + self.assertIsInstance(res, tpe) + self.assertIn(f"""'x {fun.__name__.replace("_", " ").upper()}'""", str(res)) + + for _expr in exprs: + res = F.asc(_expr) + self.assertIsInstance(res, tpe) + self.assertIn("""'x ASC NULLS FIRST'""", str(res)) + + for _expr in exprs: + res = F.desc(_expr) + self.assertIsInstance(res, tpe) + self.assertIn("""'x DESC NULLS LAST'""", str(res)) + + def test_sort_with_nulls_order(self): + df = self.spark.createDataFrame( + [("Tom", 80), (None, 60), ("Alice", 50)], ["name", "height"] + ) + self.assertEqual( + df.select(df.name).orderBy(F.asc_nulls_first("name")).collect(), + [Row(name=None), Row(name="Alice"), Row(name="Tom")], + ) + self.assertEqual( + df.select(df.name).orderBy(F.asc_nulls_last("name")).collect(), + [Row(name="Alice"), Row(name="Tom"), Row(name=None)], + ) + self.assertEqual( + df.select(df.name).orderBy(F.desc_nulls_first("name")).collect(), + [Row(name=None), Row(name="Tom"), Row(name="Alice")], + ) + self.assertEqual( + df.select(df.name).orderBy(F.desc_nulls_last("name")).collect(), + [Row(name="Tom"), Row(name="Alice"), Row(name=None)], + ) + + def test_input_file_name_reset_for_rdd(self): + rdd = self.sc.textFile("python/test_support/hello/hello.txt").map(lambda x: {"data": x}) + df = self.spark.createDataFrame(rdd, "data STRING") + df.select(F.input_file_name().alias("file")).collect() + + non_file_df = self.spark.range(100).select(F.input_file_name()) + + results = non_file_df.collect() + self.assertTrue(len(results) == 100) + + # [SPARK-24605]: if everything was properly reset after the last job, this should return + # empty string rather than the file read in the last job. + for result in results: + self.assertEqual(result[0], "") + + def test_slice(self): + df = self.spark.createDataFrame( + [ + ( + [1, 2, 3], + 2, + 2, + ), + ( + [4, 5], + 2, + 2, + ), + ], + ["x", "index", "len"], + ) + + expected = [Row(sliced=[2, 3]), Row(sliced=[5])] + self.assertEqual(df.select(F.slice(df.x, 2, 2).alias("sliced")).collect(), expected) + self.assertEqual( + df.select(F.slice(df.x, F.lit(2), F.lit(2)).alias("sliced")).collect(), expected + ) + self.assertEqual( + df.select(F.slice("x", "index", "len").alias("sliced")).collect(), expected + ) + + self.assertEqual( + df.select(F.slice(df.x, F.size(df.x) - 1, F.lit(1)).alias("sliced")).collect(), + [Row(sliced=[2]), Row(sliced=[4])], + ) + self.assertEqual( + df.select(F.slice(df.x, F.lit(1), F.size(df.x) - 1).alias("sliced")).collect(), + [Row(sliced=[1, 2]), Row(sliced=[4])], + ) + + def test_array_repeat(self): + df = self.spark.range(1) + df = df.withColumn("repeat_n", F.lit(3)) + + expected = [Row(val=[0, 0, 0])] + self.assertEqual(df.select(F.array_repeat("id", 3).alias("val")).collect(), expected) + self.assertEqual(df.select(F.array_repeat("id", F.lit(3)).alias("val")).collect(), expected) + self.assertEqual( + df.select(F.array_repeat("id", "repeat_n").alias("val")).collect(), expected + ) + + def test_input_file_name_udf(self): + df = self.spark.read.text("python/test_support/hello/hello.txt") + df = df.select(F.udf(lambda x: x)("value"), F.input_file_name().alias("file")) + file_name = df.collect()[0].file + self.assertTrue("python/test_support/hello/hello.txt" in file_name) + + def test_least(self): + df = self.spark.createDataFrame([(1, 4, 3)], ["a", "b", "c"]) + + expected = [Row(least=1)] + self.assertEqual(df.select(F.least(df.a, df.b, df.c).alias("least")).collect(), expected) + self.assertEqual( + df.select(F.least(F.lit(3), F.lit(5), F.lit(1)).alias("least")).collect(), expected + ) + self.assertEqual(df.select(F.least("a", "b", "c").alias("least")).collect(), expected) + + with self.assertRaises(PySparkValueError) as pe: + df.select(F.least(df.a).alias("least")).collect() + + self.check_error( + exception=pe.exception, + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "least", "num_cols": "2"}, + ) + + def test_overlay(self): + actual = list( + chain.from_iterable( + [ + re.findall("(overlay\\(.*\\))", str(x)) + for x in [ + F.overlay(F.col("foo"), F.col("bar"), 1), + F.overlay("x", "y", 3), + F.overlay(F.col("x"), F.col("y"), 1, 3), + F.overlay("x", "y", 2, 5), + F.overlay("x", "y", F.lit(11)), + F.overlay("x", "y", F.lit(2), F.lit(5)), + ] + ] + ) + ) + + expected = [ + "overlay(foo, bar, 1, -1)", + "overlay(x, y, 3, -1)", + "overlay(x, y, 1, 3)", + "overlay(x, y, 2, 5)", + "overlay(x, y, 11, -1)", + "overlay(x, y, 2, 5)", + ] + + self.assertListEqual(actual, expected) + + df = self.spark.createDataFrame([("SPARK_SQL", "CORE", 7, 0)], ("x", "y", "pos", "len")) + + exp = [Row(ol="SPARK_CORESQL")] + self.assertEqual(df.select(F.overlay(df.x, df.y, 7, 0).alias("ol")).collect(), exp) + self.assertEqual( + df.select(F.overlay(df.x, df.y, F.lit(7), F.lit(0)).alias("ol")).collect(), exp + ) + self.assertEqual(df.select(F.overlay("x", "y", "pos", "len").alias("ol")).collect(), exp) + + with self.assertRaises(PySparkTypeError) as pe: + df.select(F.overlay(df.x, df.y, 7.5, 0).alias("ol")).collect() + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "pos", "arg_type": "float"}, + ) + + with self.assertRaises(PySparkTypeError) as pe: + df.select(F.overlay(df.x, df.y, 7, 0.5).alias("ol")).collect() + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT_OR_STR", + message_parameters={"arg_name": "len", "arg_type": "float"}, + ) + + def test_percentile(self): + actual = list( + chain.from_iterable( + [ + re.findall("(percentile\\(.*\\))", str(x)) + for x in [ + F.percentile(F.col("foo"), F.lit(0.5)), + F.percentile(F.col("bar"), 0.25, 2), + F.percentile(F.col("bar"), [0.25, 0.5, 0.75]), + F.percentile(F.col("foo"), (0.05, 0.95), 100), + F.percentile("foo", 0.5), + F.percentile("bar", [0.1, 0.9], F.lit(10)), + ] + ] + ) + ) + + expected = [ + "percentile(foo, 0.5, 1)", + "percentile(bar, 0.25, 2)", + "percentile(bar, array(0.25, 0.5, 0.75), 1)", + "percentile(foo, array(0.05, 0.95), 100)", + "percentile(foo, 0.5, 1)", + "percentile(bar, array(0.1, 0.9), 10)", + ] + + self.assertListEqual(actual, expected) + + def test_median(self): + actual = list( + chain.from_iterable( + [ + re.findall("(median\\(.*\\))", str(x)) + for x in [ + F.median(F.col("foo")), + ] + ] + ) + ) + + expected = [ + "median(foo)", + ] + + self.assertListEqual(actual, expected) + + def test_percentile_approx(self): + actual = list( + chain.from_iterable( + [ + re.findall("(percentile_approx\\(.*\\))", str(x)) + for x in [ + F.percentile_approx(F.col("foo"), F.lit(0.5)), + F.percentile_approx(F.col("bar"), 0.25, 42), + F.percentile_approx(F.col("bar"), [0.25, 0.5, 0.75]), + F.percentile_approx(F.col("foo"), (0.05, 0.95), 100), + F.percentile_approx("foo", 0.5), + F.percentile_approx("bar", [0.1, 0.9], F.lit(10)), + ] + ] + ) + ) + + expected = [ + "percentile_approx(foo, 0.5, 10000)", + "percentile_approx(bar, 0.25, 42)", + "percentile_approx(bar, array(0.25, 0.5, 0.75), 10000)", + "percentile_approx(foo, array(0.05, 0.95), 100)", + "percentile_approx(foo, 0.5, 10000)", + "percentile_approx(bar, array(0.1, 0.9), 10)", + ] + + self.assertListEqual(actual, expected) + + def test_nth_value(self): + df = self.spark.createDataFrame( + [ + ("a", 0, None), + ("a", 1, "x"), + ("a", 2, "y"), + ("a", 3, "z"), + ("a", 4, None), + ("b", 1, None), + ("b", 2, None), + ], + schema=("key", "order", "value"), + ) + w = Window.partitionBy("key").orderBy("order") + + rs = df.select( + df.key, + df.order, + F.nth_value("value", 2).over(w), + F.nth_value("value", 2, False).over(w), + F.nth_value("value", 2, True).over(w), + ).collect() + + expected = [ + ("a", 0, None, None, None), + ("a", 1, "x", "x", None), + ("a", 2, "x", "x", "y"), + ("a", 3, "x", "x", "y"), + ("a", 4, "x", "x", "y"), + ("b", 1, None, None, None), + ("b", 2, None, None, None), + ] + + for r, ex in zip(sorted(rs), sorted(expected)): + self.assertEqual(tuple(r), ex[: len(r)]) + + def test_higher_order_function_failures(self): + # Should fail with varargs + with self.assertRaises(PySparkValueError) as pe: + F.transform(F.col("foo"), lambda *x: F.lit(1)) + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": ""}, + ) + + # Should fail with kwargs + with self.assertRaises(PySparkValueError) as pe: + F.transform(F.col("foo"), lambda **x: F.lit(1)) + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_PARAM_TYPE_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": ""}, + ) + + # Should fail with nullary function + with self.assertRaises(PySparkValueError) as pe: + F.transform(F.col("foo"), lambda: F.lit(1)) + + self.check_error( + exception=pe.exception, + error_class="WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": "", "num_args": "0"}, + ) + + # Should fail with quaternary function + with self.assertRaises(PySparkValueError) as pe: + F.transform(F.col("foo"), lambda x1, x2, x3, x4: F.lit(1)) + + self.check_error( + exception=pe.exception, + error_class="WRONG_NUM_ARGS_FOR_HIGHER_ORDER_FUNCTION", + message_parameters={"func_name": "", "num_args": "4"}, + ) + + # Should fail if function doesn't return Column + with self.assertRaises(PySparkValueError) as pe: + F.transform(F.col("foo"), lambda x: 1) + + self.check_error( + exception=pe.exception, + error_class="HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN", + message_parameters={"func_name": "", "return_type": "int"}, + ) + + def test_nested_higher_order_function(self): + # SPARK-35382: lambda vars must be resolved properly in nested higher order functions + df = self.spark.sql("SELECT array(1, 2, 3) as numbers, array('a', 'b', 'c') as letters") + + actual = df.select( + F.flatten( + F.transform( + "numbers", + lambda number: F.transform( + "letters", lambda letter: F.struct(number.alias("n"), letter.alias("l")) + ), + ) + ) + ).first()[0] + + expected = [ + (1, "a"), + (1, "b"), + (1, "c"), + (2, "a"), + (2, "b"), + (2, "c"), + (3, "a"), + (3, "b"), + (3, "c"), + ] + + self.assertEquals(actual, expected) + + def test_window_functions(self): + df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) + w = Window.partitionBy("value").orderBy("key") + + sel = df.select( + df.value, + df.key, + F.max("key").over(w.rowsBetween(0, 1)), + F.min("key").over(w.rowsBetween(0, 1)), + F.count("key").over(w.rowsBetween(float("-inf"), float("inf"))), + F.row_number().over(w), + F.rank().over(w), + F.dense_rank().over(w), + F.ntile(2).over(w), + ) + rs = sorted(sel.collect()) + expected = [ + ("1", 1, 1, 1, 1, 1, 1, 1, 1), + ("2", 1, 1, 1, 3, 1, 1, 1, 1), + ("2", 1, 2, 1, 3, 2, 1, 1, 1), + ("2", 2, 2, 2, 3, 3, 3, 2, 2), + ] + for r, ex in zip(rs, expected): + self.assertEqual(tuple(r), ex[: len(r)]) + + def test_window_functions_without_partitionBy(self): + df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) + w = Window.orderBy("key", df.value) + + sel = df.select( + df.value, + df.key, + F.max("key").over(w.rowsBetween(0, 1)), + F.min("key").over(w.rowsBetween(0, 1)), + F.count("key").over(w.rowsBetween(float("-inf"), float("inf"))), + F.row_number().over(w), + F.rank().over(w), + F.dense_rank().over(w), + F.ntile(2).over(w), + ) + rs = sorted(sel.collect()) + expected = [ + ("1", 1, 1, 1, 4, 1, 1, 1, 1), + ("2", 1, 1, 1, 4, 2, 2, 2, 1), + ("2", 1, 2, 1, 4, 3, 2, 2, 2), + ("2", 2, 2, 2, 4, 4, 4, 3, 2), + ] + for r, ex in zip(rs, expected): + self.assertEqual(tuple(r), ex[: len(r)]) + + def test_window_functions_cumulative_sum(self): + df = self.spark.createDataFrame([("one", 1), ("two", 2)], ["key", "value"]) + + # Test cumulative sum + sel = df.select( + df.key, F.sum(df.value).over(Window.rowsBetween(Window.unboundedPreceding, 0)) + ) + rs = sorted(sel.collect()) + expected = [("one", 1), ("two", 3)] + for r, ex in zip(rs, expected): + self.assertEqual(tuple(r), ex[: len(r)]) + + # Test boundary values less than JVM's Long.MinValue and make sure we don't overflow + sel = df.select( + df.key, F.sum(df.value).over(Window.rowsBetween(Window.unboundedPreceding - 1, 0)) + ) + rs = sorted(sel.collect()) + expected = [("one", 1), ("two", 3)] + for r, ex in zip(rs, expected): + self.assertEqual(tuple(r), ex[: len(r)]) + + # Test boundary values greater than JVM's Long.MaxValue and make sure we don't overflow + frame_end = Window.unboundedFollowing + 1 + sel = df.select( + df.key, F.sum(df.value).over(Window.rowsBetween(Window.currentRow, frame_end)) + ) + rs = sorted(sel.collect()) + expected = [("one", 3), ("two", 2)] + for r, ex in zip(rs, expected): + self.assertEqual(tuple(r), ex[: len(r)]) + + def test_window_time(self): + df = self.spark.createDataFrame( + [(datetime.datetime(2016, 3, 11, 9, 0, 7), 1)], ["date", "val"] + ) + + w = df.groupBy(F.window("date", "5 seconds")).agg(F.sum("val").alias("sum")) + r = w.select( + w.window.end.cast("string").alias("end"), + F.window_time(w.window).cast("string").alias("window_time"), + "sum", + ).collect() + self.assertEqual( + r[0], Row(end="2016-03-11 09:00:10", window_time="2016-03-11 09:00:09.999999", sum=1) + ) + + def test_collect_functions(self): + df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) + + self.assertEqual(sorted(df.select(F.collect_set(df.key).alias("r")).collect()[0].r), [1, 2]) + self.assertEqual( + sorted(df.select(F.collect_list(df.key).alias("r")).collect()[0].r), [1, 1, 1, 2] + ) + self.assertEqual( + sorted(df.select(F.collect_set(df.value).alias("r")).collect()[0].r), ["1", "2"] + ) + self.assertEqual( + sorted(df.select(F.collect_list(df.value).alias("r")).collect()[0].r), + ["1", "2", "2", "2"], + ) + + def test_datetime_functions(self): + df = self.spark.range(1).selectExpr("'2017-01-22' as dateCol") + parse_result = df.select(F.to_date(F.col("dateCol"))).first() + self.assertEqual(datetime.date(2017, 1, 22), parse_result["to_date(dateCol)"]) + + def test_assert_true(self): + self.check_assert_true(Py4JJavaError) + + def check_assert_true(self, tpe): + df = self.spark.range(3) + + self.assertEqual( + df.select(F.assert_true(df.id < 3)).toDF("val").collect(), + [Row(val=None), Row(val=None), Row(val=None)], + ) + + with self.assertRaisesRegex(tpe, "too big"): + df.select(F.assert_true(df.id < 2, "too big")).toDF("val").collect() + + with self.assertRaisesRegex(tpe, "2000000"): + df.select(F.assert_true(df.id < 2, df.id * 1e6)).toDF("val").collect() + + with self.assertRaises(PySparkTypeError) as pe: + df.select(F.assert_true(df.id < 2, 5)) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": "int"}, + ) + + def test_raise_error(self): + self.check_raise_error(Py4JJavaError) + + def check_raise_error(self, tpe): + df = self.spark.createDataFrame([Row(id="foobar")]) + + with self.assertRaisesRegex(tpe, "foobar"): + df.select(F.raise_error(df.id)).collect() + + with self.assertRaisesRegex(tpe, "barfoo"): + df.select(F.raise_error("barfoo")).collect() + + with self.assertRaises(PySparkTypeError) as pe: + df.select(F.raise_error(None)) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "errMsg", "arg_type": "NoneType"}, + ) + + def test_sum_distinct(self): + self.spark.range(10).select( + F.assert_true(F.sum_distinct(F.col("id")) == F.sumDistinct(F.col("id"))) + ).collect() + + def test_shiftleft(self): + self.spark.range(10).select( + F.assert_true(F.shiftLeft(F.col("id"), 2) == F.shiftleft(F.col("id"), 2)) + ).collect() + + def test_shiftright(self): + self.spark.range(10).select( + F.assert_true(F.shiftRight(F.col("id"), 2) == F.shiftright(F.col("id"), 2)) + ).collect() + + def test_shiftrightunsigned(self): + self.spark.range(10).select( + F.assert_true( + F.shiftRightUnsigned(F.col("id"), 2) == F.shiftrightunsigned(F.col("id"), 2) + ) + ).collect() + + def test_lit_day_time_interval(self): + td = datetime.timedelta(days=1, hours=12, milliseconds=123) + actual = self.spark.range(1).select(F.lit(td)).first()[0] + self.assertEqual(actual, td) + + def test_lit_list(self): + # SPARK-40271: added list type supporting + test_list = [1, 2, 3] + expected = [1, 2, 3] + actual = self.spark.range(1).select(F.lit(test_list)).first()[0] + self.assertEqual(actual, expected) + + test_list = [[1, 2, 3], [3, 4]] + expected = [[1, 2, 3], [3, 4]] + actual = self.spark.range(1).select(F.lit(test_list)).first()[0] + self.assertEqual(actual, expected) + + with self.sql_conf({"spark.sql.ansi.enabled": False}): + test_list = ["a", 1, None, 1.0] + expected = ["a", "1", None, "1.0"] + actual = self.spark.range(1).select(F.lit(test_list)).first()[0] + self.assertEqual(actual, expected) + + test_list = [["a", 1, None, 1.0], [1, None, "b"]] + expected = [["a", "1", None, "1.0"], ["1", None, "b"]] + actual = self.spark.range(1).select(F.lit(test_list)).first()[0] + self.assertEqual(actual, expected) + + df = self.spark.range(10) + with self.assertRaises(PySparkValueError) as pe: + F.lit([df.id, df.id]) + + self.check_error( + exception=pe.exception, + error_class="COLUMN_IN_LIST", + message_parameters={"func_name": "lit"}, + ) + + # Test added for SPARK-39832; change Python API to accept both col & str as input + def test_regexp_replace(self): + df = self.spark.createDataFrame( + [("100-200", r"(\d+)", "--")], ["str", "pattern", "replacement"] + ) + self.assertTrue( + all( + df.select( + F.regexp_replace("str", r"(\d+)", "--") == "-----", + F.regexp_replace("str", F.col("pattern"), F.col("replacement")) == "-----", + ).first() + ) + ) + + @unittest.skipIf(not have_numpy, "NumPy not installed") + def test_lit_np_scalar(self): + import numpy as np + + dtype_to_spark_dtypes = [ + (np.int8, [("1", "tinyint")]), + (np.int16, [("1", "smallint")]), + (np.int32, [("1", "int")]), + (np.int64, [("1", "bigint")]), + (np.float32, [("1.0", "float")]), + (np.float64, [("1.0", "double")]), + (np.bool_, [("true", "boolean")]), + ] + for dtype, spark_dtypes in dtype_to_spark_dtypes: + with self.subTest(dtype): + self.assertEqual(self.spark.range(1).select(F.lit(dtype(1))).dtypes, spark_dtypes) + + @unittest.skipIf(not have_numpy, "NumPy not installed") + def test_np_scalar_input(self): + import numpy as np + + df = self.spark.createDataFrame([([1, 2, 3],), ([],)], ["data"]) + for dtype in [np.int8, np.int16, np.int32, np.int64]: + res = df.select(F.array_contains(df.data, dtype(1)).alias("b")).collect() + self.assertEqual([Row(b=True), Row(b=False)], res) + res = df.select(F.array_position(df.data, dtype(1)).alias("c")).collect() + self.assertEqual([Row(c=1), Row(c=0)], res) + + df = self.spark.createDataFrame([([1.0, 2.0, 3.0],), ([],)], ["data"]) + for dtype in [np.float32, np.float64]: + res = df.select(F.array_contains(df.data, dtype(1)).alias("b")).collect() + self.assertEqual([Row(b=True), Row(b=False)], res) + res = df.select(F.array_position(df.data, dtype(1)).alias("c")).collect() + self.assertEqual([Row(c=1), Row(c=0)], res) + + @unittest.skipIf(not have_numpy, "NumPy not installed") + def test_ndarray_input(self): + import numpy as np + + arr_dtype_to_spark_dtypes = [ + ("int8", [("b", "array")]), + ("int16", [("b", "array")]), + ("int32", [("b", "array")]), + ("int64", [("b", "array")]), + ("float32", [("b", "array")]), + ("float64", [("b", "array")]), + ] + for t, expected_spark_dtypes in arr_dtype_to_spark_dtypes: + arr = np.array([1, 2]).astype(t) + self.assertEqual( + expected_spark_dtypes, self.spark.range(1).select(F.lit(arr).alias("b")).dtypes + ) + arr = np.array([1, 2]).astype(np.uint) + with self.assertRaises(PySparkTypeError) as pe: + self.spark.range(1).select(F.lit(arr).alias("b")) + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_NUMPY_ARRAY_SCALAR", + message_parameters={ + "dtype": "uint64", + }, + ) + + def test_binary_math_function(self): + funcs, expected = zip( + *[(F.atan2, 0.13664), (F.hypot, 8.07527), (F.pow, 2.14359), (F.pmod, 1.1)] + ) + df = self.spark.range(1).select(*(func(1.1, 8) for func in funcs)) + for a, e in zip(df.first(), expected): + self.assertAlmostEqual(a, e, 5) + + def test_map_functions(self): + # SPARK-38496: Check basic functionality of all "map" type related functions + expected = {"a": 1, "b": 2} + expected2 = {"c": 3, "d": 4} + df = self.spark.createDataFrame( + [(list(expected.keys()), list(expected.values()))], ["k", "v"] + ) + actual = ( + df.select( + F.expr("map('c', 3, 'd', 4) as dict2"), + F.map_from_arrays(df.k, df.v).alias("dict"), + "*", + ) + .select( + F.map_contains_key("dict", "a").alias("one"), + F.map_contains_key("dict", "d").alias("not_exists"), + F.map_keys("dict").alias("keys"), + F.map_values("dict").alias("values"), + F.map_entries("dict").alias("items"), + "*", + ) + .select( + F.map_concat("dict", "dict2").alias("merged"), + F.map_from_entries(F.arrays_zip("keys", "values")).alias("from_items"), + "*", + ) + .first() + ) + self.assertEqual(expected, actual["dict"]) + self.assertTrue(actual["one"]) + self.assertFalse(actual["not_exists"]) + self.assertEqual(list(expected.keys()), actual["keys"]) + self.assertEqual(list(expected.values()), actual["values"]) + self.assertEqual(expected, dict(actual["items"])) + self.assertEqual({**expected, **expected2}, dict(actual["merged"])) + self.assertEqual(expected, actual["from_items"]) + + def test_schema_of_json(self): + with self.assertRaises(PySparkTypeError) as pe: + F.schema_of_json(1) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "json", "arg_type": "int"}, + ) + + def test_schema_of_csv(self): + with self.assertRaises(PySparkTypeError) as pe: + F.schema_of_csv(1) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "csv", "arg_type": "int"}, + ) + + def test_from_csv(self): + df = self.spark.range(10) + with self.assertRaises(PySparkTypeError) as pe: + F.from_csv(df.id, 1) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "schema", "arg_type": "int"}, + ) + + def test_greatest(self): + df = self.spark.range(10) + with self.assertRaises(PySparkValueError) as pe: + F.greatest(df.id) + + self.check_error( + exception=pe.exception, + error_class="WRONG_NUM_COLUMNS", + message_parameters={"func_name": "greatest", "num_cols": "2"}, + ) + + def test_when(self): + with self.assertRaises(PySparkTypeError) as pe: + F.when("id", 1) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN", + message_parameters={"arg_name": "condition", "arg_type": "str"}, + ) + + def test_window(self): + with self.assertRaises(PySparkTypeError) as pe: + F.window("date", 5) + + self.check_error( + exception=pe.exception, + error_class="NOT_STR", + message_parameters={"arg_name": "windowDuration", "arg_type": "int"}, + ) + + def test_session_window(self): + with self.assertRaises(PySparkTypeError) as pe: + F.session_window("date", 5) + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_STR", + message_parameters={"arg_name": "gapDuration", "arg_type": "int"}, + ) + + def test_bucket(self): + with self.assertRaises(PySparkTypeError) as pe: + F.bucket("5", "id") + + self.check_error( + exception=pe.exception, + error_class="NOT_COLUMN_OR_INT", + message_parameters={"arg_name": "numBuckets", "arg_type": "str"}, + ) + + +class FunctionsTests(ReusedSQLTestCase, FunctionsTestsMixin): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_functions import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_group.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_group.py new file mode 100644 index 0000000..2715571 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_group.py @@ -0,0 +1,53 @@ +# +# 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.sql import Row +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class GroupTestsMixin: + def test_aggregator(self): + df = self.df + g = df.groupBy() + self.assertEqual([99, 100], sorted(g.agg({"key": "max", "value": "count"}).collect()[0])) + self.assertEqual([Row(**{"AVG(key#0)": 49.5})], g.mean().collect()) + + from pyspark.sql import functions + + self.assertEqual( + (0, "99"), tuple(g.agg(functions.first(df.key), functions.last(df.value)).first()) + ) + self.assertTrue(95 < g.agg(functions.approx_count_distinct(df.key)).first()[0]) + # test deprecated countDistinct + self.assertEqual(100, g.agg(functions.countDistinct(df.value)).first()[0]) + + +class GroupTests(GroupTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_group import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_pandas_sqlmetrics.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_pandas_sqlmetrics.py new file mode 100644 index 0000000..22a0e92 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_pandas_sqlmetrics.py @@ -0,0 +1,68 @@ +# +# 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 unittest +from typing import cast + +from pyspark.sql.functions import pandas_udf +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, + cast(str, pandas_requirement_message or pyarrow_requirement_message), +) +class PandasSQLMetrics(ReusedSQLTestCase): + def test_pandas_sql_metrics_basic(self): + # SPARK-34265: Instrument Python UDFs using SQL metrics + + python_sql_metrics = [ + "data sent to Python workers", + "data returned from Python workers", + "number of output rows", + ] + + @pandas_udf("long") + def test_pandas(col1): + return col1 * col1 + + self.spark.range(10).select(test_pandas("id")).collect() + + statusStore = self.spark._jsparkSession.sharedState().statusStore() + lastExecId = statusStore.executionsList().last().executionId() + executionMetrics = statusStore.execution(lastExecId).get().metrics().mkString() + + for metric in python_sql_metrics: + self.assertIn(metric, executionMetrics) + + +if __name__ == "__main__": + from pyspark.sql.tests.test_pandas_sqlmetrics import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_readwriter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_readwriter.py new file mode 100644 index 0000000..6bcef51 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_readwriter.py @@ -0,0 +1,253 @@ +# +# 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 os +import shutil +import tempfile + +from pyspark.errors import AnalysisException +from pyspark.sql.functions import col +from pyspark.sql.readwriter import DataFrameWriterV2 +from pyspark.sql.types import StructType, StructField, StringType +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class ReadwriterTestsMixin: + def test_save_and_load(self): + df = self.df + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + try: + df.write.json(tmpPath) + actual = self.spark.read.json(tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + schema = StructType([StructField("value", StringType(), True)]) + actual = self.spark.read.json(tmpPath, schema) + self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect())) + + df.write.json(tmpPath, "overwrite") + actual = self.spark.read.json(tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + df.write.save( + format="json", + mode="overwrite", + path=tmpPath, + noUse="this options will not be used in save.", + ) + actual = self.spark.read.load( + format="json", path=tmpPath, noUse="this options will not be used in load." + ) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + try: + self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json") + actual = self.spark.read.load(path=tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + finally: + self.spark.sql("RESET spark.sql.sources.default") + + csvpath = os.path.join(tempfile.mkdtemp(), "data") + df.write.option("quote", None).format("csv").save(csvpath) + finally: + shutil.rmtree(tmpPath) + + def test_save_and_load_builder(self): + df = self.df + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + try: + df.write.json(tmpPath) + actual = self.spark.read.json(tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + schema = StructType([StructField("value", StringType(), True)]) + actual = self.spark.read.json(tmpPath, schema) + self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect())) + + df.write.mode("overwrite").json(tmpPath) + actual = self.spark.read.json(tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + df.write.mode("overwrite").options( + noUse="this options will not be used in save." + ).option("noUse", "this option will not be used in save.").format("json").save( + path=tmpPath + ) + actual = self.spark.read.format("json").load( + path=tmpPath, noUse="this options will not be used in load." + ) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + + try: + self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json") + actual = self.spark.read.load(path=tmpPath) + self.assertEqual(sorted(df.collect()), sorted(actual.collect())) + finally: + self.spark.sql("RESET spark.sql.sources.default") + finally: + shutil.rmtree(tmpPath) + + def test_bucketed_write(self): + data = [ + (1, "foo", 3.0), + (2, "foo", 5.0), + (3, "bar", -1.0), + (4, "bar", 6.0), + ] + df = self.spark.createDataFrame(data, ["x", "y", "z"]) + + def count_bucketed_cols(names, table="pyspark_bucket"): + """Given a sequence of column names and a table name + query the catalog and return number o columns which are + used for bucketing + """ + cols = self.spark.catalog.listColumns(table) + num = len([c for c in cols if c.name in names and c.isBucket]) + return num + + with self.table("pyspark_bucket"): + # Test write with one bucketing column + df.write.bucketBy(3, "x").mode("overwrite").saveAsTable("pyspark_bucket") + self.assertEqual(count_bucketed_cols(["x"]), 1) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + # Test write two bucketing columns + df.write.bucketBy(3, "x", "y").mode("overwrite").saveAsTable("pyspark_bucket") + self.assertEqual(count_bucketed_cols(["x", "y"]), 2) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + # Test write with bucket and sort + df.write.bucketBy(2, "x").sortBy("z").mode("overwrite").saveAsTable("pyspark_bucket") + self.assertEqual(count_bucketed_cols(["x"]), 1) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + # Test write with a list of columns + df.write.bucketBy(3, ["x", "y"]).mode("overwrite").saveAsTable("pyspark_bucket") + self.assertEqual(count_bucketed_cols(["x", "y"]), 2) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + # Test write with bucket and sort with a list of columns + ( + df.write.bucketBy(2, "x") + .sortBy(["y", "z"]) + .mode("overwrite") + .saveAsTable("pyspark_bucket") + ) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + # Test write with bucket and sort with multiple columns + ( + df.write.bucketBy(2, "x") + .sortBy("y", "z") + .mode("overwrite") + .saveAsTable("pyspark_bucket") + ) + self.assertSetEqual(set(data), set(self.spark.table("pyspark_bucket").collect())) + + def test_insert_into(self): + df = self.spark.createDataFrame([("a", 1), ("b", 2)], ["C1", "C2"]) + with self.table("test_table"): + df.write.saveAsTable("test_table") + self.assertEqual(2, self.spark.sql("select * from test_table").count()) + + df.write.insertInto("test_table") + self.assertEqual(4, self.spark.sql("select * from test_table").count()) + + df.write.mode("overwrite").insertInto("test_table") + self.assertEqual(2, self.spark.sql("select * from test_table").count()) + + df.write.insertInto("test_table", True) + self.assertEqual(2, self.spark.sql("select * from test_table").count()) + + df.write.insertInto("test_table", False) + self.assertEqual(4, self.spark.sql("select * from test_table").count()) + + df.write.mode("overwrite").insertInto("test_table", False) + self.assertEqual(6, self.spark.sql("select * from test_table").count()) + + +class ReadwriterV2TestsMixin: + def test_api(self): + self.check_api(DataFrameWriterV2) + + def check_api(self, tpe): + df = self.df + writer = df.writeTo("testcat.t") + self.assertIsInstance(writer, tpe) + self.assertIsInstance(writer.option("property", "value"), tpe) + self.assertIsInstance(writer.options(property="value"), tpe) + self.assertIsInstance(writer.using("source"), tpe) + self.assertIsInstance(writer.partitionedBy("id"), tpe) + self.assertIsInstance(writer.partitionedBy(col("id")), tpe) + self.assertIsInstance(writer.tableProperty("foo", "bar"), tpe) + + def test_partitioning_functions(self): + self.check_partitioning_functions(DataFrameWriterV2) + + def check_partitioning_functions(self, tpe): + import datetime + from pyspark.sql.functions import years, months, days, hours, bucket + + df = self.spark.createDataFrame( + [(1, datetime.datetime(2000, 1, 1), "foo")], ("id", "ts", "value") + ) + + writer = df.writeTo("testcat.t") + + self.assertIsInstance(writer.partitionedBy(years("ts")), tpe) + self.assertIsInstance(writer.partitionedBy(months("ts")), tpe) + self.assertIsInstance(writer.partitionedBy(days("ts")), tpe) + self.assertIsInstance(writer.partitionedBy(hours("ts")), tpe) + self.assertIsInstance(writer.partitionedBy(bucket(11, "id")), tpe) + self.assertIsInstance(writer.partitionedBy(bucket(11, col("id"))), tpe) + self.assertIsInstance(writer.partitionedBy(bucket(3, "id"), hours(col("ts"))), tpe) + + def test_create(self): + df = self.df + with self.table("test_table"): + df.writeTo("test_table").using("parquet").create() + self.assertEqual(100, self.spark.sql("select * from test_table").count()) + + def test_create_without_provider(self): + df = self.df + with self.assertRaisesRegex( + AnalysisException, "NOT_SUPPORTED_COMMAND_WITHOUT_HIVE_SUPPORT" + ): + df.writeTo("test_table").create() + + +class ReadwriterTests(ReadwriterTestsMixin, ReusedSQLTestCase): + pass + + +class ReadwriterV2Tests(ReadwriterV2TestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_readwriter import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_serde.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_serde.py new file mode 100644 index 0000000..ef8bbd2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_serde.py @@ -0,0 +1,157 @@ +# +# 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 datetime +import shutil +import tempfile +import time + +from pyspark.sql import Row +from pyspark.sql.functions import lit +from pyspark.sql.types import StructType, StructField, DecimalType, BinaryType +from pyspark.testing.sqlutils import ReusedSQLTestCase, UTCOffsetTimezone + + +class SerdeTestsMixin: + def test_serialize_nested_array_and_map(self): + d = [Row(lst=[Row(a=1, b="s")], d={"key": Row(c=1.0, d="2")})] + rdd = self.sc.parallelize(d) + df = self.spark.createDataFrame(rdd) + row = df.head() + self.assertEqual(1, len(row.lst)) + self.assertEqual(1, row.lst[0].a) + self.assertEqual("2", row.d["key"].d) + + lst = df.rdd.map(lambda x: x.lst).first() + self.assertEqual(1, len(lst)) + self.assertEqual("s", lst[0].b) + + d = df.rdd.map(lambda x: x.d).first() + self.assertEqual(1, len(d)) + self.assertEqual(1.0, d["key"].c) + + row = df.rdd.map(lambda x: x.d["key"]).first() + self.assertEqual(1.0, row.c) + self.assertEqual("2", row.d) + + def test_select_null_literal(self): + df = self.spark.sql("select null as col") + self.assertEqual(Row(col=None), df.first()) + + def test_struct_in_map(self): + d = [Row(m={Row(i=1): Row(s="")})] + df = self.spark.createDataFrame(d) + k, v = list(df.head().m.items())[0] + self.assertEqual(1, k.i) + self.assertEqual("", v.s) + + def test_filter_with_datetime(self): + time = datetime.datetime(2015, 4, 17, 23, 1, 2, 3000) + date = time.date() + row = Row(date=date, time=time) + df = self.spark.createDataFrame([row]) + self.assertEqual(1, df.filter(df.date == date).count()) + self.assertEqual(1, df.filter(df.time == time).count()) + self.assertEqual(0, df.filter(df.date > date).count()) + self.assertEqual(0, df.filter(df.time > time).count()) + + def test_filter_with_datetime_timezone(self): + dt1 = datetime.datetime(2015, 4, 17, 23, 1, 2, 3000, tzinfo=UTCOffsetTimezone(0)) + dt2 = datetime.datetime(2015, 4, 17, 23, 1, 2, 3000, tzinfo=UTCOffsetTimezone(1)) + row = Row(date=dt1) + df = self.spark.createDataFrame([row]) + self.assertEqual(0, df.filter(df.date == dt2).count()) + self.assertEqual(1, df.filter(df.date > dt2).count()) + self.assertEqual(0, df.filter(df.date < dt2).count()) + + def test_time_with_timezone(self): + day = datetime.date.today() + now = datetime.datetime.now() + ts = time.mktime(now.timetuple()) + # class in __main__ is not serializable + from pyspark.testing.sqlutils import UTCOffsetTimezone + + utc = UTCOffsetTimezone() + utcnow = datetime.datetime.utcfromtimestamp(ts) # without microseconds + # add microseconds to utcnow (keeping year,month,day,hour,minute,second) + utcnow = datetime.datetime(*(utcnow.timetuple()[:6] + (now.microsecond, utc))) + df = self.spark.createDataFrame([(day, now, utcnow)]) + day1, now1, utcnow1 = df.first() + self.assertEqual(day1, day) + self.assertEqual(now, now1) + self.assertEqual(now, utcnow1) + + # regression test for SPARK-19561 + def test_datetime_at_epoch(self): + epoch = datetime.datetime.fromtimestamp(0) + df = self.spark.createDataFrame([Row(date=epoch)]) + first = df.select("date", lit(epoch).alias("lit_date")).first() + self.assertEqual(first["date"], epoch) + self.assertEqual(first["lit_date"], epoch) + + def test_decimal(self): + from decimal import Decimal + + schema = StructType([StructField("decimal", DecimalType(10, 5))]) + df = self.spark.createDataFrame([(Decimal("3.14159"),)], schema) + row = df.select(df.decimal + 1).first() + self.assertEqual(row[0], Decimal("4.14159")) + tmpPath = tempfile.mkdtemp() + shutil.rmtree(tmpPath) + df.write.parquet(tmpPath) + df2 = self.spark.read.parquet(tmpPath) + row = df2.first() + self.assertEqual(row[0], Decimal("3.14159")) + + def test_BinaryType_serialization(self): + # Pyrolite version <= 4.9 could not serialize BinaryType with Python3 SPARK-17808 + # The empty bytearray is test for SPARK-21534. + schema = StructType([StructField("mybytes", BinaryType())]) + data = [ + [bytearray(b"here is my data")], + [bytearray(b"and here is some more")], + [bytearray(b"")], + ] + df = self.spark.createDataFrame(data, schema=schema) + df.collect() + + def test_int_array_serialization(self): + # Note that this test seems dependent on parallelism. + data = self.spark.sparkContext.parallelize([[1, 2, 3, 4]] * 100, numSlices=12) + df = self.spark.createDataFrame(data, "array") + self.assertEqual(len(list(filter(lambda r: None in r.value, df.collect()))), 0) + + def test_bytes_as_binary_type(self): + df = self.spark.createDataFrame([[b"abcd"]], "col binary") + self.assertEqual(df.first().col, bytearray(b"abcd")) + + +class SerdeTests(SerdeTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_serde import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_session.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_session.py new file mode 100644 index 0000000..dacaff4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_session.py @@ -0,0 +1,412 @@ +# +# 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 os +import unittest + +from pyspark import SparkConf, SparkContext +from pyspark.sql import SparkSession, SQLContext, Row +from pyspark.sql.functions import col +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import PySparkTestCase + + +class SparkSessionTests(ReusedSQLTestCase): + def test_sqlcontext_reuses_sparksession(self): + sqlContext1 = SQLContext(self.sc) + sqlContext2 = SQLContext(self.sc) + self.assertTrue(sqlContext1.sparkSession is sqlContext2.sparkSession) + + +class SparkSessionTests1(ReusedSQLTestCase): + + # We can't include this test into SQLTests because we will stop class's SparkContext and cause + # other tests failed. + def test_sparksession_with_stopped_sparkcontext(self): + self.sc.stop() + sc = SparkContext("local[4]", self.sc.appName) + spark = SparkSession.builder.getOrCreate() + try: + df = spark.createDataFrame([(1, 2)], ["c", "c"]) + df.collect() + finally: + spark.stop() + sc.stop() + + +class SparkSessionTests2(PySparkTestCase): + + # This test is separate because it's closely related with session's start and stop. + # See SPARK-23228. + def test_set_jvm_default_session(self): + spark = SparkSession.builder.getOrCreate() + try: + self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isDefined()) + finally: + spark.stop() + self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isEmpty()) + + def test_jvm_default_session_already_set(self): + # Here, we assume there is the default session already set in JVM. + jsession = self.sc._jvm.SparkSession(self.sc._jsc.sc()) + self.sc._jvm.SparkSession.setDefaultSession(jsession) + + spark = SparkSession.builder.getOrCreate() + try: + self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isDefined()) + # The session should be the same with the exiting one. + self.assertTrue(jsession.equals(spark._jvm.SparkSession.getDefaultSession().get())) + finally: + spark.stop() + + +class SparkSessionTests3(unittest.TestCase): + def test_active_session(self): + spark = SparkSession.builder.master("local").getOrCreate() + try: + activeSession = SparkSession.getActiveSession() + df = activeSession.createDataFrame([(1, "Alice")], ["age", "name"]) + self.assertEqual(df.collect(), [Row(age=1, name="Alice")]) + with self.assertRaises(ValueError): + activeSession.createDataFrame(activeSession._sc.parallelize([[], []])) + finally: + spark.stop() + + def test_get_active_session_when_no_active_session(self): + active = SparkSession.getActiveSession() + self.assertEqual(active, None) + spark = SparkSession.builder.master("local").getOrCreate() + active = SparkSession.getActiveSession() + self.assertEqual(active, spark) + spark.stop() + active = SparkSession.getActiveSession() + self.assertEqual(active, None) + + def test_spark_session(self): + spark = SparkSession.builder.master("local").config("some-config", "v2").getOrCreate() + try: + self.assertEqual(spark.conf.get("some-config"), "v2") + self.assertEqual(spark.sparkContext._conf.get("some-config"), "v2") + self.assertEqual(spark.version, spark.sparkContext.version) + spark.sql("CREATE DATABASE test_db") + spark.catalog.setCurrentDatabase("test_db") + self.assertEqual(spark.catalog.currentDatabase(), "test_db") + spark.sql("CREATE TABLE table1 (name STRING, age INT) USING parquet") + self.assertEqual(spark.table("table1").columns, ["name", "age"]) + self.assertEqual(spark.range(3).count(), 3) + + # SPARK-37516: Only plain column references work as variable in SQL. + self.assertEqual( + spark.sql("select {c} from range(1)", c=col("id")).first(), spark.range(1).first() + ) + with self.assertRaisesRegex(ValueError, "Column"): + spark.sql("select {c} from range(10)", c=col("id") + 1) + finally: + spark.sql("DROP DATABASE test_db CASCADE") + spark.stop() + + def test_global_default_session(self): + spark = SparkSession.builder.master("local").getOrCreate() + try: + self.assertEqual(SparkSession.builder.getOrCreate(), spark) + finally: + spark.stop() + + def test_default_and_active_session(self): + spark = SparkSession.builder.master("local").getOrCreate() + activeSession = spark._jvm.SparkSession.getActiveSession() + defaultSession = spark._jvm.SparkSession.getDefaultSession() + try: + self.assertEqual(activeSession, defaultSession) + finally: + spark.stop() + + def test_config_option_propagated_to_existing_session(self): + session1 = SparkSession.builder.master("local").config("spark-config1", "a").getOrCreate() + self.assertEqual(session1.conf.get("spark-config1"), "a") + session2 = SparkSession.builder.config("spark-config1", "b").getOrCreate() + try: + self.assertEqual(session1, session2) + self.assertEqual(session1.conf.get("spark-config1"), "b") + finally: + session1.stop() + + def test_new_session(self): + session = SparkSession.builder.master("local").getOrCreate() + newSession = session.newSession() + try: + self.assertNotEqual(session, newSession) + finally: + session.stop() + newSession.stop() + + def test_create_new_session_if_old_session_stopped(self): + session = SparkSession.builder.master("local").getOrCreate() + session.stop() + newSession = SparkSession.builder.master("local").getOrCreate() + try: + self.assertNotEqual(session, newSession) + finally: + newSession.stop() + + def test_active_session_with_None_and_not_None_context(self): + from pyspark.context import SparkContext + from pyspark.conf import SparkConf + + sc = None + session = None + try: + sc = SparkContext._active_spark_context + self.assertEqual(sc, None) + activeSession = SparkSession.getActiveSession() + self.assertEqual(activeSession, None) + sparkConf = SparkConf() + sc = SparkContext.getOrCreate(sparkConf) + activeSession = sc._jvm.SparkSession.getActiveSession() + self.assertFalse(activeSession.isDefined()) + session = SparkSession(sc) + activeSession = sc._jvm.SparkSession.getActiveSession() + self.assertTrue(activeSession.isDefined()) + activeSession2 = SparkSession.getActiveSession() + self.assertNotEqual(activeSession2, None) + finally: + if session is not None: + session.stop() + if sc is not None: + sc.stop() + + +class SparkSessionTests4(ReusedSQLTestCase): + def test_get_active_session_after_create_dataframe(self): + session2 = None + try: + activeSession1 = SparkSession.getActiveSession() + session1 = self.spark + self.assertEqual(session1, activeSession1) + session2 = self.spark.newSession() + activeSession2 = SparkSession.getActiveSession() + self.assertEqual(session1, activeSession2) + self.assertNotEqual(session2, activeSession2) + session2.createDataFrame([(1, "Alice")], ["age", "name"]) + activeSession3 = SparkSession.getActiveSession() + self.assertEqual(session2, activeSession3) + session1.createDataFrame([(1, "Alice")], ["age", "name"]) + activeSession4 = SparkSession.getActiveSession() + self.assertEqual(session1, activeSession4) + finally: + if session2 is not None: + session2.stop() + + +class SparkSessionTests5(unittest.TestCase): + def setUp(self): + # These tests require restarting the Spark context so we set up a new one for each test + # rather than at the class level. + self.sc = SparkContext("local[4]", self.__class__.__name__, conf=SparkConf()) + self.spark = SparkSession(self.sc) + + def tearDown(self): + self.sc.stop() + self.spark.stop() + + def test_sqlcontext_with_stopped_sparksession(self): + # SPARK-30856: test that SQLContext.getOrCreate() returns a usable instance after + # the SparkSession is restarted. + sql_context = SQLContext.getOrCreate(self.spark.sparkContext) + self.spark.stop() + spark = SparkSession.builder.master("local[4]").appName(self.sc.appName).getOrCreate() + new_sql_context = SQLContext.getOrCreate(spark.sparkContext) + + self.assertIsNot(new_sql_context, sql_context) + self.assertIs(SQLContext.getOrCreate(spark.sparkContext).sparkSession, spark) + try: + df = spark.createDataFrame([(1, 2)], ["c", "c"]) + df.collect() + finally: + spark.stop() + self.assertIsNone(SQLContext._instantiatedContext) + + def test_sqlcontext_with_stopped_sparkcontext(self): + # SPARK-30856: test initialization via SparkSession when only the SparkContext is stopped + self.sc.stop() + spark = SparkSession.builder.master("local[4]").appName(self.sc.appName).getOrCreate() + self.sc = spark.sparkContext + self.assertIs(SQLContext.getOrCreate(self.sc).sparkSession, spark) + + def test_get_sqlcontext_with_stopped_sparkcontext(self): + # SPARK-30856: test initialization via SQLContext.getOrCreate() when only the SparkContext + # is stopped + self.sc.stop() + self.sc = SparkContext("local[4]", self.sc.appName) + self.assertIs(SQLContext.getOrCreate(self.sc)._sc, self.sc) + + +class SparkSessionBuilderTests(unittest.TestCase): + def test_create_spark_context_first_then_spark_session(self): + sc = None + session = None + try: + conf = SparkConf().set("key1", "value1") + sc = SparkContext("local[4]", "SessionBuilderTests", conf=conf) + session = SparkSession.builder.config("key2", "value2").getOrCreate() + + self.assertEqual(session.conf.get("key1"), "value1") + self.assertEqual(session.conf.get("key2"), "value2") + self.assertEqual(session.sparkContext, sc) + + self.assertFalse(sc.getConf().contains("key2")) + self.assertEqual(sc.getConf().get("key1"), "value1") + finally: + if session is not None: + session.stop() + if sc is not None: + sc.stop() + + def test_another_spark_session(self): + session1 = None + session2 = None + try: + session1 = SparkSession.builder.config("key1", "value1").getOrCreate() + session2 = SparkSession.builder.config( + "spark.sql.codegen.comments", "true" + ).getOrCreate() + + self.assertEqual(session1.conf.get("key1"), "value1") + self.assertEqual(session2.conf.get("key1"), "value1") + self.assertEqual(session1.conf.get("spark.sql.codegen.comments"), "false") + self.assertEqual(session2.conf.get("spark.sql.codegen.comments"), "false") + self.assertEqual(session1.sparkContext, session2.sparkContext) + + self.assertEqual(session1.sparkContext.getConf().get("key1"), "value1") + self.assertFalse(session1.sparkContext.getConf().contains("key2")) + finally: + if session1 is not None: + session1.stop() + if session2 is not None: + session2.stop() + + def test_create_spark_context_with_initial_session_options(self): + sc = None + session = None + try: + conf = SparkConf().set("key1", "value1") + sc = SparkContext("local[4]", "SessionBuilderTests", conf=conf) + session = ( + SparkSession.builder.config("spark.sql.codegen.comments", "true") + .enableHiveSupport() + .getOrCreate() + ) + + self.assertEqual(session._jsparkSession.sharedState().conf().get("key1"), "value1") + self.assertEqual( + session._jsparkSession.sharedState().conf().get("spark.sql.codegen.comments"), + "true", + ) + self.assertEqual( + session._jsparkSession.sharedState().conf().get("spark.sql.catalogImplementation"), + "hive", + ) + self.assertEqual(session.sparkContext, sc) + finally: + if session is not None: + session.stop() + if sc is not None: + sc.stop() + + def test_create_spark_context_with_initial_session_options_bool(self): + session = None + # Test if `True` is set as "true". + try: + session = SparkSession.builder.config( + "spark.sql.pyspark.jvmStacktrace.enabled", True + ).getOrCreate() + self.assertEqual(session.conf.get("spark.sql.pyspark.jvmStacktrace.enabled"), "true") + finally: + if session is not None: + session.stop() + # Test if `False` is set as "false". + try: + session = SparkSession.builder.config( + "spark.sql.pyspark.jvmStacktrace.enabled", False + ).getOrCreate() + self.assertEqual(session.conf.get("spark.sql.pyspark.jvmStacktrace.enabled"), "false") + finally: + if session is not None: + session.stop() + + +class SparkExtensionsTest(unittest.TestCase): + # These tests are separate because it uses 'spark.sql.extensions' which is + # static and immutable. This can't be set or unset, for example, via `spark.conf`. + + @classmethod + def setUpClass(cls): + import glob + from pyspark.find_spark_home import _find_spark_home + + SPARK_HOME = _find_spark_home() + filename_pattern = ( + "sql/core/target/scala-*/test-classes/org/apache/spark/sql/" + "SparkSessionExtensionSuite.class" + ) + if not glob.glob(os.path.join(SPARK_HOME, filename_pattern)): + raise unittest.SkipTest( + "'org.apache.spark.sql.SparkSessionExtensionSuite' is not " + "available. Will skip the related tests." + ) + + # Note that 'spark.sql.extensions' is a static immutable configuration. + cls.spark = ( + SparkSession.builder.master("local[4]") + .appName(cls.__name__) + .config("spark.sql.extensions", "org.apache.spark.sql.MyExtensions") + .getOrCreate() + ) + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def test_use_custom_class_for_extensions(self): + self.assertTrue( + self.spark._jsparkSession.sessionState() + .planner() + .strategies() + .contains( + self.spark._jvm.org.apache.spark.sql.MySparkStrategy(self.spark._jsparkSession) + ), + "MySparkStrategy not found in active planner strategies", + ) + self.assertTrue( + self.spark._jsparkSession.sessionState() + .analyzer() + .extendedResolutionRules() + .contains(self.spark._jvm.org.apache.spark.sql.MyRule(self.spark._jsparkSession)), + "MyRule not found in extended resolution rules", + ) + + +if __name__ == "__main__": + from pyspark.sql.tests.test_session import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_types.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_types.py new file mode 100644 index 0000000..90ecfd6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_types.py @@ -0,0 +1,1568 @@ +# -*- encoding: utf-8 -*- +# +# 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 array +import ctypes +import datetime +import os +import pickle +import sys +import unittest + +from pyspark.sql import Row +from pyspark.sql import functions as F +from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError +from pyspark.sql.types import ( + ByteType, + ShortType, + IntegerType, + FloatType, + DateType, + TimestampType, + DayTimeIntervalType, + YearMonthIntervalType, + MapType, + StringType, + CharType, + VarcharType, + StructType, + StructField, + ArrayType, + DoubleType, + LongType, + DecimalType, + BinaryType, + BooleanType, + NullType, +) +from pyspark.sql.types import ( + _array_signed_int_typecode_ctype_mappings, + _array_type_mappings, + _array_unsigned_int_typecode_ctype_mappings, + _infer_type, + _make_type_verifier, + _merge_type, +) +from pyspark.testing.sqlutils import ( + ReusedSQLTestCase, + ExamplePointUDT, + PythonOnlyUDT, + ExamplePoint, + PythonOnlyPoint, + MyObject, +) +from pyspark.testing.utils import PySparkErrorTestUtils + + +class TypesTestsMixin: + def test_apply_schema_to_row(self): + df = self.spark.read.json(self.sc.parallelize(["""{"a":2}"""])) + df2 = self.spark.createDataFrame(df.rdd.map(lambda x: x), df.schema) + self.assertEqual(df.collect(), df2.collect()) + + rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x)) + df3 = self.spark.createDataFrame(rdd, df.schema) + self.assertEqual(10, df3.count()) + + def test_infer_schema_to_local(self): + input = [{"a": 1}, {"b": "coffee"}] + rdd = self.sc.parallelize(input) + df = self.spark.createDataFrame(input) + df2 = self.spark.createDataFrame(rdd, samplingRatio=1.0) + self.assertEqual(df.schema, df2.schema) + + rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x, b=None)) + df3 = self.spark.createDataFrame(rdd, df.schema) + self.assertEqual(10, df3.count()) + + def test_apply_schema_to_dict_and_rows(self): + schema = StructType().add("a", IntegerType()).add("b", StringType()) + input = [{"a": 1}, {"b": "coffee"}] + rdd = self.sc.parallelize(input) + for verify in [False, True]: + df = self.spark.createDataFrame(input, schema, verifySchema=verify) + df2 = self.spark.createDataFrame(rdd, schema, verifySchema=verify) + self.assertEqual(df.schema, df2.schema) + + rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x, b=None)) + df3 = self.spark.createDataFrame(rdd, schema, verifySchema=verify) + self.assertEqual(10, df3.count()) + input = [Row(a=x, b=str(x)) for x in range(10)] + df4 = self.spark.createDataFrame(input, schema, verifySchema=verify) + self.assertEqual(10, df4.count()) + + def test_create_dataframe_schema_mismatch(self): + rdd = self.sc.parallelize(range(3)).map(lambda i: Row(a=i)) + schema = StructType([StructField("a", IntegerType()), StructField("b", StringType())]) + df = self.spark.createDataFrame(rdd, schema) + self.assertRaises(Exception, lambda: df.show()) + + def test_infer_schema(self): + d = [Row(l=[], d={}, s=None), Row(l=[Row(a=1, b="s")], d={"key": Row(c=1.0, d="2")}, s="")] + rdd = self.sc.parallelize(d) + df = self.spark.createDataFrame(rdd) + self.assertEqual([], df.rdd.map(lambda r: r.l).first()) + self.assertEqual([None, ""], df.rdd.map(lambda r: r.s).collect()) + + with self.tempView("test"): + df.createOrReplaceTempView("test") + result = self.spark.sql("SELECT l from test") + self.assertEqual([], result.head()[0]) + # We set `spark.sql.ansi.enabled` to False for this case + # since it occurs an error in ANSI mode if there is a list index + # or key that does not exist. + with self.sql_conf({"spark.sql.ansi.enabled": False}): + result = self.spark.sql("SELECT l[0].a from test where d['key'].d = '2'") + self.assertEqual(1, result.head()[0]) + + df2 = self.spark.createDataFrame(rdd, samplingRatio=1.0) + self.assertEqual(df.schema, df2.schema) + self.assertEqual({}, df2.rdd.map(lambda r: r.d).first()) + self.assertEqual([None, ""], df2.rdd.map(lambda r: r.s).collect()) + + with self.tempView("test2"): + df2.createOrReplaceTempView("test2") + result = self.spark.sql("SELECT l from test2") + self.assertEqual([], result.head()[0]) + # We set `spark.sql.ansi.enabled` to False for this case + # since it occurs an error in ANSI mode if there is a list index + # or key that does not exist. + with self.sql_conf({"spark.sql.ansi.enabled": False}): + result = self.spark.sql("SELECT l[0].a from test2 where d['key'].d = '2'") + self.assertEqual(1, result.head()[0]) + + def test_infer_schema_specification(self): + from decimal import Decimal + + class A: + def __init__(self): + self.a = 1 + + data = [ + True, + 1, + "a", + "a", + datetime.date(1970, 1, 1), + datetime.datetime(1970, 1, 1, 0, 0), + datetime.timedelta(microseconds=123456678), + 1.0, + array.array("d", [1]), + [1], + (1,), + {"a": 1}, + bytearray(1), + Decimal(1), + Row(a=1), + Row("a")(1), + A(), + ] + + df = self.spark.createDataFrame([data]) + actual = list(map(lambda x: x.dataType.simpleString(), df.schema)) + expected = [ + "boolean", + "bigint", + "string", + "string", + "date", + "timestamp", + "interval day to second", + "double", + "array", + "array", + "struct<_1:bigint>", + "map", + "binary", + "decimal(38,18)", + "struct", + "struct", + "struct", + ] + self.assertEqual(actual, expected) + + actual = list(df.first()) + expected = [ + True, + 1, + "a", + "a", + datetime.date(1970, 1, 1), + datetime.datetime(1970, 1, 1, 0, 0), + datetime.timedelta(microseconds=123456678), + 1.0, + [1.0], + [1], + Row(_1=1), + {"a": 1}, + bytearray(b"\x00"), + Decimal("1.000000000000000000"), + Row(a=1), + Row(a=1), + Row(a=1), + ] + self.assertEqual(actual, expected) + + with self.sql_conf({"spark.sql.timestampType": "TIMESTAMP_NTZ"}): + with self.sql_conf({"spark.sql.session.timeZone": "America/Sao_Paulo"}): + df = self.spark.createDataFrame([(datetime.datetime(1970, 1, 1, 0, 0),)]) + self.assertEqual(list(df.schema)[0].dataType.simpleString(), "timestamp_ntz") + self.assertEqual(df.first()[0], datetime.datetime(1970, 1, 1, 0, 0)) + + df = self.spark.createDataFrame( + [ + (datetime.datetime(1970, 1, 1, 0, 0),), + (datetime.datetime(1970, 1, 1, 0, 0, tzinfo=datetime.timezone.utc),), + ] + ) + self.assertEqual(list(df.schema)[0].dataType.simpleString(), "timestamp") + + def test_infer_schema_not_enough_names(self): + df = self.spark.createDataFrame([["a", "b"]], ["col1"]) + self.assertEqual(df.columns, ["col1", "_2"]) + + def test_infer_schema_upcast_int_to_string(self): + df = self.spark.createDataFrame( + self.spark.sparkContext.parallelize([[1, 1], ["x", 1]]), + schema=["a", "b"], + samplingRatio=0.99, + ) + self.assertEqual([Row(a="1", b=1), Row(a="x", b=1)], df.collect()) + + def test_infer_schema_upcast_float_to_string(self): + df = self.spark.createDataFrame([[1.33, 1], ["2.1", 1]], schema=["a", "b"]) + self.assertEqual([Row(a="1.33", b=1), Row(a="2.1", b=1)], df.collect()) + + def test_infer_schema_upcast_boolean_to_string(self): + df = self.spark.createDataFrame([[True, 1], ["false", 1]], schema=["a", "b"]) + self.assertEqual([Row(a="true", b=1), Row(a="false", b=1)], df.collect()) + + def test_infer_nested_schema(self): + NestedRow = Row("f1", "f2") + nestedRdd1 = self.sc.parallelize( + [NestedRow([1, 2], {"row1": 1.0}), NestedRow([2, 3], {"row2": 2.0})] + ) + df = self.spark.createDataFrame(nestedRdd1) + self.assertEqual(Row(f1=[1, 2], f2={"row1": 1.0}), df.collect()[0]) + + nestedRdd2 = self.sc.parallelize( + [NestedRow([[1, 2], [2, 3]], [1, 2]), NestedRow([[2, 3], [3, 4]], [2, 3])] + ) + df = self.spark.createDataFrame(nestedRdd2) + self.assertEqual(Row(f1=[[1, 2], [2, 3]], f2=[1, 2]), df.collect()[0]) + + from collections import namedtuple + + CustomRow = namedtuple("CustomRow", "field1 field2") + rdd = self.sc.parallelize( + [ + CustomRow(field1=1, field2="row1"), + CustomRow(field1=2, field2="row2"), + CustomRow(field1=3, field2="row3"), + ] + ) + df = self.spark.createDataFrame(rdd) + self.assertEqual(Row(field1=1, field2="row1"), df.first()) + + def test_infer_nested_dict_as_struct(self): + # SPARK-35929: Test inferring nested dict as a struct type. + NestedRow = Row("f1", "f2") + + with self.sql_conf({"spark.sql.pyspark.inferNestedDictAsStruct.enabled": True}): + data = [ + NestedRow([{"payment": 200.5, "name": "A"}], [1, 2]), + NestedRow([{"payment": 100.5, "name": "B"}], [2, 3]), + ] + + df = self.spark.createDataFrame(data) + self.assertEqual(Row(f1=[Row(payment=200.5, name="A")], f2=[1, 2]), df.first()) + + def test_infer_nested_dict_as_struct_with_rdd(self): + # SPARK-35929: Test inferring nested dict as a struct type. + NestedRow = Row("f1", "f2") + + with self.sql_conf({"spark.sql.pyspark.inferNestedDictAsStruct.enabled": True}): + data = [ + NestedRow([{"payment": 200.5, "name": "A"}], [1, 2]), + NestedRow([{"payment": 100.5, "name": "B"}], [2, 3]), + ] + + nestedRdd = self.sc.parallelize(data) + df = self.spark.createDataFrame(nestedRdd) + self.assertEqual(Row(f1=[Row(payment=200.5, name="A")], f2=[1, 2]), df.first()) + + def test_infer_array_merge_element_types(self): + # SPARK-39168: Test inferring array element type from all values in array + ArrayRow = Row("f1", "f2") + + data = [ArrayRow([1, None], [None, 2])] + + df = self.spark.createDataFrame(data) + self.assertEqual(Row(f1=[1, None], f2=[None, 2]), df.first()) + + # Test legacy behavior inferring only from the first element + with self.sql_conf( + {"spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled": True} + ): + # Legacy: f2 schema inferred as an array of nulls, should raise error + self.assertRaises(ValueError, lambda: self.spark.createDataFrame(data)) + + # an array with only null values should raise an error + data2 = [ArrayRow([1], [None])] + self.assertRaises(ValueError, lambda: self.spark.createDataFrame(data2)) + + # an array with no values should raise an error + data3 = [ArrayRow([1], [])] + self.assertRaises(ValueError, lambda: self.spark.createDataFrame(data3)) + + # an array with conflicting types should raise an error + # in this case this is ArrayType(StringType) and ArrayType(NullType) + data4 = [ArrayRow([1, "1"], [None])] + with self.assertRaisesRegex(ValueError, "types cannot be determined after inferring"): + self.spark.createDataFrame(data4) + + def test_infer_array_merge_element_types_with_rdd(self): + # SPARK-39168: Test inferring array element type from all values in array + ArrayRow = Row("f1", "f2") + + data = [ArrayRow([1, None], [None, 2])] + + rdd = self.sc.parallelize(data) + df = self.spark.createDataFrame(rdd) + self.assertEqual(Row(f1=[1, None], f2=[None, 2]), df.first()) + + def test_infer_array_element_type_empty(self): + # SPARK-39168: Test inferring array element type from all rows + ArrayRow = Row("f1") + + data = [ArrayRow([]), ArrayRow([None]), ArrayRow([1])] + + rdd = self.sc.parallelize(data) + df = self.spark.createDataFrame(rdd) + rows = df.collect() + self.assertEqual(Row(f1=[]), rows[0]) + self.assertEqual(Row(f1=[None]), rows[1]) + self.assertEqual(Row(f1=[1]), rows[2]) + + df = self.spark.createDataFrame(data) + rows = df.collect() + self.assertEqual(Row(f1=[]), rows[0]) + self.assertEqual(Row(f1=[None]), rows[1]) + self.assertEqual(Row(f1=[1]), rows[2]) + + def test_infer_array_element_type_with_struct(self): + # SPARK-39168: Test inferring array of struct type from all struct values + NestedRow = Row("f1") + + with self.sql_conf({"spark.sql.pyspark.inferNestedDictAsStruct.enabled": True}): + data = [NestedRow([{"payment": 200.5}, {"name": "A"}])] + + nestedRdd = self.sc.parallelize(data) + df = self.spark.createDataFrame(nestedRdd) + self.assertEqual( + Row(f1=[Row(payment=200.5, name=None), Row(payment=None, name="A")]), df.first() + ) + + df = self.spark.createDataFrame(data) + self.assertEqual( + Row(f1=[Row(payment=200.5, name=None), Row(payment=None, name="A")]), df.first() + ) + + # Test legacy behavior inferring only from the first element; excludes "name" field + with self.sql_conf( + {"spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled": True} + ): + df = self.spark.createDataFrame(data) + self.assertEqual(Row(f1=[Row(payment=200.5), Row(payment=None)]), df.first()) + + def test_create_dataframe_from_dict_respects_schema(self): + df = self.spark.createDataFrame([{"a": 1}], ["b"]) + self.assertEqual(df.columns, ["b"]) + + def test_negative_decimal(self): + try: + self.spark.sql("set spark.sql.legacy.allowNegativeScaleOfDecimal=true") + df = self.spark.createDataFrame([(1,), (11,)], ["value"]) + ret = df.select(F.col("value").cast(DecimalType(1, -1))).collect() + actual = list(map(lambda r: int(r.value), ret)) + self.assertEqual(actual, [0, 10]) + finally: + self.spark.sql("set spark.sql.legacy.allowNegativeScaleOfDecimal=false") + + def test_create_dataframe_from_objects(self): + data = [MyObject(1, "1"), MyObject(2, "2")] + df = self.spark.createDataFrame(data) + self.assertEqual(df.dtypes, [("key", "bigint"), ("value", "string")]) + self.assertEqual(df.first(), Row(key=1, value="1")) + + def test_apply_schema(self): + from datetime import date, datetime, timedelta + + rdd = self.sc.parallelize( + [ + ( + 127, + -128, + -32768, + 32767, + 2147483647, + 1.0, + date(2010, 1, 1), + datetime(2010, 1, 1, 1, 1, 1), + timedelta(days=1), + {"a": 1}, + (2,), + [1, 2, 3], + None, + ) + ] + ) + schema = StructType( + [ + StructField("byte1", ByteType(), False), + StructField("byte2", ByteType(), False), + StructField("short1", ShortType(), False), + StructField("short2", ShortType(), False), + StructField("int1", IntegerType(), False), + StructField("float1", FloatType(), False), + StructField("date1", DateType(), False), + StructField("time1", TimestampType(), False), + StructField("daytime1", DayTimeIntervalType(), False), + StructField("map1", MapType(StringType(), IntegerType(), False), False), + StructField("struct1", StructType([StructField("b", ShortType(), False)]), False), + StructField("list1", ArrayType(ByteType(), False), False), + StructField("null1", DoubleType(), True), + ] + ) + df = self.spark.createDataFrame(rdd, schema) + results = df.rdd.map( + lambda x: ( + x.byte1, + x.byte2, + x.short1, + x.short2, + x.int1, + x.float1, + x.date1, + x.time1, + x.daytime1, + x.map1["a"], + x.struct1.b, + x.list1, + x.null1, + ) + ) + r = ( + 127, + -128, + -32768, + 32767, + 2147483647, + 1.0, + date(2010, 1, 1), + datetime(2010, 1, 1, 1, 1, 1), + timedelta(days=1), + 1, + 2, + [1, 2, 3], + None, + ) + self.assertEqual(r, results.first()) + + with self.tempView("table2"): + df.createOrReplaceTempView("table2") + r = self.spark.sql( + "SELECT byte1 - 1 AS byte1, byte2 + 1 AS byte2, " + + "short1 + 1 AS short1, short2 - 1 AS short2, int1 - 1 AS int1, " + + "float1 + 1.5 as float1 FROM table2" + ).first() + + self.assertEqual((126, -127, -32767, 32766, 2147483646, 2.5), tuple(r)) + + def test_convert_row_to_dict(self): + row = Row(l=[Row(a=1, b="s")], d={"key": Row(c=1.0, d="2")}) + self.assertEqual(1, row.asDict()["l"][0].a) + df = self.spark.createDataFrame([row]) + + with self.tempView("test"): + df.createOrReplaceTempView("test") + row = self.spark.sql("select l, d from test").head() + self.assertEqual(1, row.asDict()["l"][0].a) + self.assertEqual(1.0, row.asDict()["d"]["key"].c) + + def test_udt(self): + from pyspark.sql.types import _parse_datatype_json_string, _infer_type, _make_type_verifier + + def check_datatype(datatype): + pickled = pickle.loads(pickle.dumps(datatype)) + assert datatype == pickled + scala_datatype = self.spark._jsparkSession.parseDataType(datatype.json()) + python_datatype = _parse_datatype_json_string(scala_datatype.json()) + assert datatype == python_datatype + + check_datatype(ExamplePointUDT()) + structtype_with_udt = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", ExamplePointUDT(), False), + ] + ) + check_datatype(structtype_with_udt) + p = ExamplePoint(1.0, 2.0) + self.assertEqual(_infer_type(p), ExamplePointUDT()) + _make_type_verifier(ExamplePointUDT())(ExamplePoint(1.0, 2.0)) + self.assertRaises(ValueError, lambda: _make_type_verifier(ExamplePointUDT())([1.0, 2.0])) + + check_datatype(PythonOnlyUDT()) + structtype_with_udt = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", PythonOnlyUDT(), False), + ] + ) + check_datatype(structtype_with_udt) + p = PythonOnlyPoint(1.0, 2.0) + self.assertEqual(_infer_type(p), PythonOnlyUDT()) + _make_type_verifier(PythonOnlyUDT())(PythonOnlyPoint(1.0, 2.0)) + self.assertRaises(ValueError, lambda: _make_type_verifier(PythonOnlyUDT())([1.0, 2.0])) + + def test_simple_udt_in_df(self): + schema = StructType().add("key", LongType()).add("val", PythonOnlyUDT()) + df = self.spark.createDataFrame( + [(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)], schema=schema + ) + df.collect() + + def test_nested_udt_in_df(self): + schema = StructType().add("key", LongType()).add("val", ArrayType(PythonOnlyUDT())) + df = self.spark.createDataFrame( + [(i % 3, [PythonOnlyPoint(float(i), float(i))]) for i in range(10)], schema=schema + ) + df.collect() + + schema = ( + StructType().add("key", LongType()).add("val", MapType(LongType(), PythonOnlyUDT())) + ) + df = self.spark.createDataFrame( + [(i % 3, {i % 3: PythonOnlyPoint(float(i + 1), float(i + 1))}) for i in range(10)], + schema=schema, + ) + df.collect() + + def test_complex_nested_udt_in_df(self): + schema = StructType().add("key", LongType()).add("val", PythonOnlyUDT()) + df = self.spark.createDataFrame( + [(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)], schema=schema + ) + df.collect() + + gd = df.groupby("key").agg({"val": "collect_list"}) + gd.collect() + udf = F.udf(lambda k, v: [(k, v[0])], ArrayType(df.schema)) + gd.select(udf(*gd)).collect() + + def test_udt_with_none(self): + df = self.spark.range(0, 10, 1, 1) + + def myudf(x): + if x > 0: + return PythonOnlyPoint(float(x), float(x)) + + self.spark.catalog.registerFunction("udf", myudf, PythonOnlyUDT()) + rows = [r[0] for r in df.selectExpr("udf(id)").take(2)] + self.assertEqual(rows, [None, PythonOnlyPoint(1, 1)]) + + def test_infer_schema_with_udt(self): + row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + schema = df.schema + field = [f for f in schema.fields if f.name == "point"][0] + self.assertEqual(type(field.dataType), ExamplePointUDT) + + with self.tempView("labeled_point"): + df.createOrReplaceTempView("labeled_point") + point = self.spark.sql("SELECT point FROM labeled_point").head().point + self.assertEqual(point, ExamplePoint(1.0, 2.0)) + + row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + schema = df.schema + field = [f for f in schema.fields if f.name == "point"][0] + self.assertEqual(type(field.dataType), PythonOnlyUDT) + + with self.tempView("labeled_point"): + df.createOrReplaceTempView("labeled_point") + point = self.spark.sql("SELECT point FROM labeled_point").head().point + self.assertEqual(point, PythonOnlyPoint(1.0, 2.0)) + + def test_infer_schema_with_udt_with_column_names(self): + row = (1.0, ExamplePoint(1.0, 2.0)) + df = self.spark.createDataFrame([row], ["label", "point"]) + schema = df.schema + field = [f for f in schema.fields if f.name == "point"][0] + self.assertEqual(type(field.dataType), ExamplePointUDT) + + with self.tempView("labeled_point"): + df.createOrReplaceTempView("labeled_point") + point = self.spark.sql("SELECT point FROM labeled_point").head().point + self.assertEqual(point, ExamplePoint(1.0, 2.0)) + + row = (1.0, PythonOnlyPoint(1.0, 2.0)) + df = self.spark.createDataFrame([row], ["label", "point"]) + schema = df.schema + field = [f for f in schema.fields if f.name == "point"][0] + self.assertEqual(type(field.dataType), PythonOnlyUDT) + + with self.tempView("labeled_point"): + df.createOrReplaceTempView("labeled_point") + point = self.spark.sql("SELECT point FROM labeled_point").head().point + self.assertEqual(point, PythonOnlyPoint(1.0, 2.0)) + + def test_apply_schema_with_udt(self): + row = (1.0, ExamplePoint(1.0, 2.0)) + schema = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", ExamplePointUDT(), False), + ] + ) + df = self.spark.createDataFrame([row], schema) + point = df.head().point + self.assertEqual(point, ExamplePoint(1.0, 2.0)) + + row = (1.0, PythonOnlyPoint(1.0, 2.0)) + schema = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", PythonOnlyUDT(), False), + ] + ) + df = self.spark.createDataFrame([row], schema) + point = df.head().point + self.assertEqual(point, PythonOnlyPoint(1.0, 2.0)) + + def test_apply_schema_with_nullable_udt(self): + rows = [(1.0, ExamplePoint(1.0, 2.0)), (2.0, None)] + schema = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", ExamplePointUDT(), True), + ] + ) + df = self.spark.createDataFrame(rows, schema) + points = [row.point for row in df.collect()] + self.assertEqual(points, [ExamplePoint(1.0, 2.0), None]) + + rows = [(1.0, PythonOnlyPoint(1.0, 2.0)), (2.0, None)] + schema = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", PythonOnlyUDT(), True), + ] + ) + df = self.spark.createDataFrame(rows, schema) + points = [row.point for row in df.collect()] + self.assertEqual(points, [PythonOnlyPoint(1.0, 2.0), None]) + + def test_udf_with_udt(self): + row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + udf = F.udf(lambda p: p.y, DoubleType()) + self.assertEqual(2.0, df.select(udf(df.point)).first()[0]) + udf2 = F.udf(lambda p: ExamplePoint(p.x + 1, p.y + 1), ExamplePointUDT()) + self.assertEqual(ExamplePoint(2.0, 3.0), df.select(udf2(df.point)).first()[0]) + + row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + udf = F.udf(lambda p: p.y, DoubleType()) + self.assertEqual(2.0, df.select(udf(df.point)).first()[0]) + udf2 = F.udf(lambda p: PythonOnlyPoint(p.x + 1, p.y + 1), PythonOnlyUDT()) + self.assertEqual(PythonOnlyPoint(2.0, 3.0), df.select(udf2(df.point)).first()[0]) + + def test_rdd_with_udt(self): + row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + self.assertEqual(1.0, df.rdd.map(lambda r: r.point.x).first()) + + row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + self.assertEqual(1.0, df.rdd.map(lambda r: r.point.x).first()) + + def test_parquet_with_udt(self): + row = Row(label=1.0, point=ExamplePoint(1.0, 2.0)) + df0 = self.spark.createDataFrame([row]) + output_dir = os.path.join(self.tempdir.name, "labeled_point") + df0.write.parquet(output_dir) + df1 = self.spark.read.parquet(output_dir) + point = df1.head().point + self.assertEqual(point, ExamplePoint(1.0, 2.0)) + + row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0)) + df0 = self.spark.createDataFrame([row]) + df0.write.parquet(output_dir, mode="overwrite") + df1 = self.spark.read.parquet(output_dir) + point = df1.head().point + self.assertEqual(point, PythonOnlyPoint(1.0, 2.0)) + + def test_union_with_udt(self): + row1 = (1.0, ExamplePoint(1.0, 2.0)) + row2 = (2.0, ExamplePoint(3.0, 4.0)) + schema = StructType( + [ + StructField("label", DoubleType(), False), + StructField("point", ExamplePointUDT(), False), + ] + ) + df1 = self.spark.createDataFrame([row1], schema) + df2 = self.spark.createDataFrame([row2], schema) + + result = df1.union(df2).orderBy("label").collect() + self.assertEqual( + result, + [ + Row(label=1.0, point=ExamplePoint(1.0, 2.0)), + Row(label=2.0, point=ExamplePoint(3.0, 4.0)), + ], + ) + + def test_cast_to_string_with_udt(self): + row = (ExamplePoint(1.0, 2.0), PythonOnlyPoint(3.0, 4.0)) + schema = StructType( + [ + StructField("point", ExamplePointUDT(), False), + StructField("pypoint", PythonOnlyUDT(), False), + ] + ) + df = self.spark.createDataFrame([row], schema) + + result = df.select(F.col("point").cast("string"), F.col("pypoint").cast("string")).head() + self.assertEqual(result, Row(point="(1.0, 2.0)", pypoint="[3.0, 4.0]")) + + def test_cast_to_udt_with_udt(self): + row = Row(point=ExamplePoint(1.0, 2.0), python_only_point=PythonOnlyPoint(1.0, 2.0)) + df = self.spark.createDataFrame([row]) + with self.assertRaises(AnalysisException): + df.select(F.col("point").cast(PythonOnlyUDT())).collect() + with self.assertRaises(AnalysisException): + df.select(F.col("python_only_point").cast(ExamplePointUDT())).collect() + + def test_struct_type(self): + struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + struct2 = StructType( + [StructField("f1", StringType(), True), StructField("f2", StringType(), True, None)] + ) + self.assertEqual(struct1.fieldNames(), struct2.names) + self.assertEqual(struct1, struct2) + + struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + struct2 = StructType([StructField("f1", StringType(), True)]) + self.assertNotEqual(struct1.fieldNames(), struct2.names) + self.assertNotEqual(struct1, struct2) + + struct1 = ( + StructType() + .add(StructField("f1", StringType(), True)) + .add(StructField("f2", StringType(), True, None)) + ) + struct2 = StructType( + [StructField("f1", StringType(), True), StructField("f2", StringType(), True, None)] + ) + self.assertEqual(struct1.fieldNames(), struct2.names) + self.assertEqual(struct1, struct2) + + struct1 = ( + StructType() + .add(StructField("f1", StringType(), True)) + .add(StructField("f2", StringType(), True, None)) + ) + struct2 = StructType([StructField("f1", StringType(), True)]) + self.assertNotEqual(struct1.fieldNames(), struct2.names) + self.assertNotEqual(struct1, struct2) + + # Catch exception raised during improper construction + self.assertRaises(ValueError, lambda: StructType().add("name")) + + struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + for field in struct1: + self.assertIsInstance(field, StructField) + + struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + self.assertEqual(len(struct1), 2) + + struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + self.assertIs(struct1["f1"], struct1.fields[0]) + self.assertIs(struct1[0], struct1.fields[0]) + self.assertEqual(struct1[0:1], StructType(struct1.fields[0:1])) + self.assertRaises(KeyError, lambda: struct1["f9"]) + self.assertRaises(IndexError, lambda: struct1[9]) + self.assertRaises(TypeError, lambda: struct1[9.9]) + + def test_parse_datatype_string(self): + from pyspark.sql.types import _all_atomic_types, _parse_datatype_string + + for k, t in _all_atomic_types.items(): + if k != "varchar" and k != "char": + self.assertEqual(t(), _parse_datatype_string(k)) + self.assertEqual(IntegerType(), _parse_datatype_string("int")) + self.assertEqual(CharType(1), _parse_datatype_string("char(1)")) + self.assertEqual(CharType(10), _parse_datatype_string("char( 10 )")) + self.assertEqual(CharType(11), _parse_datatype_string("char( 11)")) + self.assertEqual(VarcharType(1), _parse_datatype_string("varchar(1)")) + self.assertEqual(VarcharType(10), _parse_datatype_string("varchar( 10 )")) + self.assertEqual(VarcharType(11), _parse_datatype_string("varchar( 11)")) + self.assertEqual(DecimalType(1, 1), _parse_datatype_string("decimal(1 ,1)")) + self.assertEqual(DecimalType(10, 1), _parse_datatype_string("decimal( 10,1 )")) + self.assertEqual(DecimalType(11, 1), _parse_datatype_string("decimal(11,1)")) + self.assertEqual(ArrayType(IntegerType()), _parse_datatype_string("array")) + self.assertEqual( + MapType(IntegerType(), DoubleType()), _parse_datatype_string("map< int, double >") + ) + self.assertEqual( + StructType([StructField("a", IntegerType()), StructField("c", DoubleType())]), + _parse_datatype_string("struct"), + ) + self.assertEqual( + StructType([StructField("a", IntegerType()), StructField("c", DoubleType())]), + _parse_datatype_string("a:int, c:double"), + ) + self.assertEqual( + StructType([StructField("a", IntegerType()), StructField("c", DoubleType())]), + _parse_datatype_string("a INT, c DOUBLE"), + ) + + def test_metadata_null(self): + schema = StructType( + [ + StructField("f1", StringType(), True, None), + StructField("f2", StringType(), True, {"a": None}), + ] + ) + self.spark.createDataFrame([["a", "b"], ["c", "d"]], schema) + + def test_access_nested_types(self): + df = self.spark.createDataFrame([Row(l=[1], r=Row(a=1, b="b"), d={"k": "v"})]) + self.assertEqual(1, df.select(df.l[0]).first()[0]) + self.assertEqual(1, df.select(df.l.getItem(0)).first()[0]) + self.assertEqual(1, df.select(df.r.a).first()[0]) + self.assertEqual("b", df.select(df.r.getField("b")).first()[0]) + self.assertEqual("v", df.select(df.d["k"]).first()[0]) + self.assertEqual("v", df.select(df.d.getItem("k")).first()[0]) + + def test_infer_long_type(self): + longrow = [Row(f1="a", f2=100000000000000)] + df = self.sc.parallelize(longrow).toDF() + self.assertEqual(df.schema.fields[1].dataType, LongType()) + + # this saving as Parquet caused issues as well. + output_dir = os.path.join(self.tempdir.name, "infer_long_type") + df.write.parquet(output_dir) + df1 = self.spark.read.parquet(output_dir) + self.assertEqual("a", df1.first().f1) + self.assertEqual(100000000000000, df1.first().f2) + + self.assertEqual(_infer_type(1), LongType()) + self.assertEqual(_infer_type(2**10), LongType()) + self.assertEqual(_infer_type(2**20), LongType()) + self.assertEqual(_infer_type(2**31 - 1), LongType()) + self.assertEqual(_infer_type(2**31), LongType()) + self.assertEqual(_infer_type(2**61), LongType()) + self.assertEqual(_infer_type(2**71), LongType()) + + def test_infer_binary_type(self): + binaryrow = [Row(f1="a", f2=b"abcd")] + df = self.sc.parallelize(binaryrow).toDF() + self.assertEqual(df.schema.fields[1].dataType, BinaryType()) + + # this saving as Parquet caused issues as well. + output_dir = os.path.join(self.tempdir.name, "infer_binary_type") + df.write.parquet(output_dir) + df1 = self.spark.read.parquet(output_dir) + self.assertEqual("a", df1.first().f1) + self.assertEqual(b"abcd", df1.first().f2) + + self.assertEqual(_infer_type(b""), BinaryType()) + self.assertEqual(_infer_type(b"1234"), BinaryType()) + + def test_merge_type(self): + self.assertEqual(_merge_type(LongType(), NullType()), LongType()) + self.assertEqual(_merge_type(NullType(), LongType()), LongType()) + + self.assertEqual(_merge_type(LongType(), LongType()), LongType()) + + self.assertEqual( + _merge_type(ArrayType(LongType()), ArrayType(LongType())), ArrayType(LongType()) + ) + with self.assertRaises(PySparkTypeError) as pe: + _merge_type(ArrayType(LongType()), ArrayType(DoubleType())) + self.check_error( + exception=pe.exception, + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": "LongType", "data_type2": "DoubleType"}, + ) + + self.assertEqual( + _merge_type(MapType(StringType(), LongType()), MapType(StringType(), LongType())), + MapType(StringType(), LongType()), + ) + + self.assertEqual( + _merge_type(MapType(StringType(), LongType()), MapType(DoubleType(), LongType())), + MapType(StringType(), LongType()), + ) + + with self.assertRaises(PySparkTypeError) as pe: + _merge_type(MapType(StringType(), LongType()), MapType(StringType(), DoubleType())) + self.check_error( + exception=pe.exception, + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": "LongType", "data_type2": "DoubleType"}, + ) + + self.assertEqual( + _merge_type( + StructType([StructField("f1", LongType()), StructField("f2", StringType())]), + StructType([StructField("f1", LongType()), StructField("f2", StringType())]), + ), + StructType([StructField("f1", LongType()), StructField("f2", StringType())]), + ) + with self.assertRaises(PySparkTypeError) as pe: + _merge_type( + StructType([StructField("f1", LongType()), StructField("f2", StringType())]), + StructType([StructField("f1", DoubleType()), StructField("f2", StringType())]), + ) + self.check_error( + exception=pe.exception, + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": "LongType", "data_type2": "DoubleType"}, + ) + + self.assertEqual( + _merge_type( + StructType([StructField("f1", StructType([StructField("f2", LongType())]))]), + StructType([StructField("f1", StructType([StructField("f2", LongType())]))]), + ), + StructType([StructField("f1", StructType([StructField("f2", LongType())]))]), + ) + self.assertEqual( + _merge_type( + StructType([StructField("f1", StructType([StructField("f2", LongType())]))]), + StructType([StructField("f1", StructType([StructField("f2", StringType())]))]), + ), + StructType([StructField("f1", StructType([StructField("f2", StringType())]))]), + ) + + self.assertEqual( + _merge_type( + StructType( + [StructField("f1", ArrayType(LongType())), StructField("f2", StringType())] + ), + StructType( + [StructField("f1", ArrayType(LongType())), StructField("f2", StringType())] + ), + ), + StructType([StructField("f1", ArrayType(LongType())), StructField("f2", StringType())]), + ) + with self.assertRaises(PySparkTypeError) as pe: + _merge_type( + StructType( + [StructField("f1", ArrayType(LongType())), StructField("f2", StringType())] + ), + StructType( + [StructField("f1", ArrayType(DoubleType())), StructField("f2", StringType())] + ), + ) + self.check_error( + exception=pe.exception, + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": "LongType", "data_type2": "DoubleType"}, + ) + + self.assertEqual( + _merge_type( + StructType( + [ + StructField("f1", MapType(StringType(), LongType())), + StructField("f2", StringType()), + ] + ), + StructType( + [ + StructField("f1", MapType(StringType(), LongType())), + StructField("f2", StringType()), + ] + ), + ), + StructType( + [ + StructField("f1", MapType(StringType(), LongType())), + StructField("f2", StringType()), + ] + ), + ) + with self.assertRaises(PySparkTypeError) as pe: + _merge_type( + StructType( + [ + StructField("f1", MapType(StringType(), LongType())), + StructField("f2", StringType()), + ] + ), + StructType( + [ + StructField("f1", MapType(StringType(), DoubleType())), + StructField("f2", StringType()), + ] + ), + ) + self.check_error( + exception=pe.exception, + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": "LongType", "data_type2": "DoubleType"}, + ) + + self.assertEqual( + _merge_type( + StructType([StructField("f1", ArrayType(MapType(StringType(), LongType())))]), + StructType([StructField("f1", ArrayType(MapType(StringType(), LongType())))]), + ), + StructType([StructField("f1", ArrayType(MapType(StringType(), LongType())))]), + ) + self.assertEqual( + _merge_type( + StructType([StructField("f1", ArrayType(MapType(StringType(), LongType())))]), + StructType([StructField("f1", ArrayType(MapType(DoubleType(), LongType())))]), + ), + StructType([StructField("f1", ArrayType(MapType(StringType(), LongType())))]), + ) + + # test for SPARK-16542 + def test_array_types(self): + # This test need to make sure that the Scala type selected is at least + # as large as the python's types. This is necessary because python's + # array types depend on C implementation on the machine. Therefore there + # is no machine independent correspondence between python's array types + # and Scala types. + # See: https://docs.python.org/2/library/array.html + + def assertCollectSuccess(typecode, value): + row = Row(myarray=array.array(typecode, [value])) + df = self.spark.createDataFrame([row]) + self.assertEqual(df.first()["myarray"][0], value) + + # supported string types + # + # String types in python's array are "u" for Py_UNICODE and "c" for char. + # "u" will be removed in python 4, and "c" is not supported in python 3. + supported_string_types = [] + if sys.version_info[0] < 4: + supported_string_types += ["u"] + # test unicode + assertCollectSuccess("u", "a") + + # supported float and double + # + # Test max, min, and precision for float and double, assuming IEEE 754 + # floating-point format. + supported_fractional_types = ["f", "d"] + assertCollectSuccess("f", ctypes.c_float(1e38).value) + assertCollectSuccess("f", ctypes.c_float(1e-38).value) + assertCollectSuccess("f", ctypes.c_float(1.123456).value) + assertCollectSuccess("d", sys.float_info.max) + assertCollectSuccess("d", sys.float_info.min) + assertCollectSuccess("d", sys.float_info.epsilon) + + # supported signed int types + # + # The size of C types changes with implementation, we need to make sure + # that there is no overflow error on the platform running this test. + supported_signed_int_types = list( + set(_array_signed_int_typecode_ctype_mappings.keys()).intersection( + set(_array_type_mappings.keys()) + ) + ) + for t in supported_signed_int_types: + ctype = _array_signed_int_typecode_ctype_mappings[t] + max_val = 2 ** (ctypes.sizeof(ctype) * 8 - 1) + assertCollectSuccess(t, max_val - 1) + assertCollectSuccess(t, -max_val) + + # supported unsigned int types + # + # JVM does not have unsigned types. We need to be very careful to make + # sure that there is no overflow error. + supported_unsigned_int_types = list( + set(_array_unsigned_int_typecode_ctype_mappings.keys()).intersection( + set(_array_type_mappings.keys()) + ) + ) + for t in supported_unsigned_int_types: + ctype = _array_unsigned_int_typecode_ctype_mappings[t] + assertCollectSuccess(t, 2 ** (ctypes.sizeof(ctype) * 8) - 1) + + # all supported types + # + # Make sure the types tested above: + # 1. are all supported types + # 2. cover all supported types + supported_types = ( + supported_string_types + + supported_fractional_types + + supported_signed_int_types + + supported_unsigned_int_types + ) + self.assertEqual(set(supported_types), set(_array_type_mappings.keys())) + + # all unsupported types + # + # Keys in _array_type_mappings is a complete list of all supported types, + # and types not in _array_type_mappings are considered unsupported. + # PyPy seems not having array.typecodes. + all_types = set(["b", "B", "u", "h", "H", "i", "I", "l", "L", "q", "Q", "f", "d"]) + unsupported_types = all_types - set(supported_types) + # test unsupported types + for t in unsupported_types: + with self.assertRaises(PySparkTypeError) as pe: + a = array.array(t) + self.spark.createDataFrame([Row(myarray=a)]).collect() + + self.check_error( + exception=pe.exception, + error_class="CANNOT_INFER_TYPE_FOR_FIELD", + message_parameters={"field_name": "myarray"}, + ) + + def test_repr(self): + instances = [ + NullType(), + StringType(), + CharType(10), + VarcharType(10), + BinaryType(), + BooleanType(), + DateType(), + TimestampType(), + DecimalType(), + DoubleType(), + FloatType(), + ByteType(), + IntegerType(), + LongType(), + ShortType(), + ArrayType(StringType()), + MapType(StringType(), IntegerType()), + StructField("f1", StringType(), True), + StructType([StructField("f1", StringType(), True)]), + ] + for instance in instances: + self.assertEqual(eval(repr(instance)), instance) + + def test_daytime_interval_type_constructor(self): + # SPARK-37277: Test constructors in day time interval. + self.assertEqual(DayTimeIntervalType().simpleString(), "interval day to second") + self.assertEqual( + DayTimeIntervalType(DayTimeIntervalType.DAY).simpleString(), "interval day" + ) + self.assertEqual( + DayTimeIntervalType( + DayTimeIntervalType.HOUR, DayTimeIntervalType.SECOND + ).simpleString(), + "interval hour to second", + ) + + with self.assertRaisesRegex(RuntimeError, "interval None to 3 is invalid"): + DayTimeIntervalType(endField=DayTimeIntervalType.SECOND) + + with self.assertRaisesRegex(RuntimeError, "interval 123 to 123 is invalid"): + DayTimeIntervalType(123) + + with self.assertRaisesRegex(RuntimeError, "interval 0 to 321 is invalid"): + DayTimeIntervalType(DayTimeIntervalType.DAY, 321) + + def test_daytime_interval_type(self): + # SPARK-37277: Support DayTimeIntervalType in createDataFrame + timedetlas = [ + (datetime.timedelta(microseconds=123),), + ( + datetime.timedelta( + days=1, seconds=23, microseconds=123, milliseconds=4, minutes=5, hours=11 + ), + ), + (datetime.timedelta(microseconds=-123),), + (datetime.timedelta(days=-1),), + (datetime.timedelta(microseconds=388629894454999981),), + (datetime.timedelta(days=-1, seconds=86399, microseconds=999999),), # -1 microsecond + ] + df = self.spark.createDataFrame(timedetlas, schema="td interval day to second") + self.assertEqual(set(r.td for r in df.collect()), set(set(r[0] for r in timedetlas))) + + exprs = [ + "INTERVAL '1 02:03:04' DAY TO SECOND AS a", + "INTERVAL '1 02:03' DAY TO MINUTE AS b", + "INTERVAL '1 02' DAY TO HOUR AS c", + "INTERVAL '1' DAY AS d", + "INTERVAL '26:03:04' HOUR TO SECOND AS e", + "INTERVAL '26:03' HOUR TO MINUTE AS f", + "INTERVAL '26' HOUR AS g", + "INTERVAL '1563:04' MINUTE TO SECOND AS h", + "INTERVAL '1563' MINUTE AS i", + "INTERVAL '93784' SECOND AS j", + ] + df = self.spark.range(1).selectExpr(exprs) + + actual = list(df.first()) + expected = [ + datetime.timedelta(days=1, hours=2, minutes=3, seconds=4), + datetime.timedelta(days=1, hours=2, minutes=3), + datetime.timedelta(days=1, hours=2), + datetime.timedelta(days=1), + datetime.timedelta(hours=26, minutes=3, seconds=4), + datetime.timedelta(hours=26, minutes=3), + datetime.timedelta(hours=26), + datetime.timedelta(minutes=1563, seconds=4), + datetime.timedelta(minutes=1563), + datetime.timedelta(seconds=93784), + ] + + for n, (a, e) in enumerate(zip(actual, expected)): + self.assertEqual(a, e, "%s does not match with %s" % (exprs[n], expected[n])) + + def test_yearmonth_interval_type_constructor(self): + self.assertEqual(YearMonthIntervalType().simpleString(), "interval year to month") + self.assertEqual( + YearMonthIntervalType(YearMonthIntervalType.YEAR).simpleString(), "interval year" + ) + self.assertEqual( + YearMonthIntervalType( + YearMonthIntervalType.YEAR, YearMonthIntervalType.MONTH + ).simpleString(), + "interval year to month", + ) + + with self.assertRaisesRegex(RuntimeError, "interval None to 3 is invalid"): + YearMonthIntervalType(endField=3) + + with self.assertRaisesRegex(RuntimeError, "interval 123 to 123 is invalid"): + YearMonthIntervalType(123) + + with self.assertRaisesRegex(RuntimeError, "interval 0 to 321 is invalid"): + YearMonthIntervalType(YearMonthIntervalType.YEAR, 321) + + def test_yearmonth_interval_type(self): + schema1 = self.spark.sql("SELECT INTERVAL '10-8' YEAR TO MONTH AS interval").schema + self.assertEqual(schema1.fields[0].dataType, YearMonthIntervalType(0, 1)) + + schema2 = self.spark.sql("SELECT INTERVAL '10' YEAR AS interval").schema + self.assertEqual(schema2.fields[0].dataType, YearMonthIntervalType(0, 0)) + + schema3 = self.spark.sql("SELECT INTERVAL '8' MONTH AS interval").schema + self.assertEqual(schema3.fields[0].dataType, YearMonthIntervalType(1, 1)) + + +class DataTypeTests(unittest.TestCase): + # regression test for SPARK-6055 + def test_data_type_eq(self): + lt = LongType() + lt2 = pickle.loads(pickle.dumps(LongType())) + self.assertEqual(lt, lt2) + + # regression test for SPARK-7978 + def test_decimal_type(self): + t1 = DecimalType() + t2 = DecimalType(10, 2) + self.assertTrue(t2 is not t1) + self.assertNotEqual(t1, t2) + t3 = DecimalType(8) + self.assertNotEqual(t2, t3) + + def test_char_type(self): + v1 = CharType(10) + v2 = CharType(20) + self.assertTrue(v2 is not v1) + self.assertNotEqual(v1, v2) + v3 = CharType(10) + self.assertEqual(v1, v3) + self.assertFalse(v1 is v3) + + def test_varchar_type(self): + v1 = VarcharType(10) + v2 = VarcharType(20) + self.assertTrue(v2 is not v1) + self.assertNotEqual(v1, v2) + v3 = VarcharType(10) + self.assertEqual(v1, v3) + self.assertFalse(v1 is v3) + + # regression test for SPARK-10392 + def test_datetype_equal_zero(self): + dt = DateType() + self.assertEqual(dt.fromInternal(0), datetime.date(1970, 1, 1)) + + # regression test for SPARK-17035 + def test_timestamp_microsecond(self): + tst = TimestampType() + self.assertEqual(tst.toInternal(datetime.datetime.max) % 1000000, 999999) + + # regression test for SPARK-23299 + def test_row_without_column_name(self): + row = Row("Alice", 11) + self.assertEqual(repr(row), "") + + # test __repr__ with unicode values + self.assertEqual(repr(Row("数", "量")), "") + + # SPARK-44643: test __repr__ with empty Row + def test_row_repr_with_empty_row(self): + self.assertEqual(repr(Row(a=Row())), "Row(a=)") + self.assertEqual(repr(Row(Row())), ")>") + + EmptyRow = Row() + self.assertEqual(repr(Row(a=EmptyRow())), "Row(a=Row())") + self.assertEqual(repr(Row(EmptyRow())), "") + + def test_empty_row(self): + row = Row() + self.assertEqual(len(row), 0) + + def test_struct_field_type_name(self): + struct_field = StructField("a", IntegerType()) + self.assertRaises(TypeError, struct_field.typeName) + + def test_invalid_create_row(self): + row_class = Row("c1", "c2") + self.assertRaises(ValueError, lambda: row_class(1, 2, 3)) + + +class DataTypeVerificationTests(unittest.TestCase, PySparkErrorTestUtils): + def test_verify_type_exception_msg(self): + with self.assertRaises(PySparkValueError) as pe: + _make_type_verifier(StringType(), nullable=False, name="test_name")(None) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_BE_NONE", + message_parameters={ + "arg_name": "obj", + }, + ) + + schema = StructType([StructField("a", StructType([StructField("b", IntegerType())]))]) + with self.assertRaises(PySparkTypeError) as pe: + _make_type_verifier(schema)([["data"]]) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_ACCEPT_OBJECT_IN_TYPE", + message_parameters={ + "data_type": "IntegerType()", + "obj_name": "data", + "obj_type": "str", + }, + ) + + def test_verify_type_ok_nullable(self): + obj = None + types = [IntegerType(), FloatType(), StringType(), StructType([])] + for data_type in types: + try: + _make_type_verifier(data_type, nullable=True)(obj) + except Exception: + self.fail("verify_type(%s, %s, nullable=True)" % (obj, data_type)) + + def test_verify_type_not_nullable(self): + import array + import datetime + import decimal + + schema = StructType( + [ + StructField("s", StringType(), nullable=False), + StructField("i", IntegerType(), nullable=True), + ] + ) + + class MyObj: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + # obj, data_type + success_spec = [ + # String + ("", StringType()), + (1, StringType()), + (1.0, StringType()), + ([], StringType()), + ({}, StringType()), + # Char + ("", CharType(10)), + (1, CharType(10)), + (1.0, CharType(10)), + ([], CharType(10)), + ({}, CharType(10)), + # Varchar + ("", VarcharType(10)), + (1, VarcharType(10)), + (1.0, VarcharType(10)), + ([], VarcharType(10)), + ({}, VarcharType(10)), + # UDT + (ExamplePoint(1.0, 2.0), ExamplePointUDT()), + # Boolean + (True, BooleanType()), + # Byte + (-(2**7), ByteType()), + (2**7 - 1, ByteType()), + # Short + (-(2**15), ShortType()), + (2**15 - 1, ShortType()), + # Integer + (-(2**31), IntegerType()), + (2**31 - 1, IntegerType()), + # Long + (-(2**63), LongType()), + (2**63 - 1, LongType()), + # Float & Double + (1.0, FloatType()), + (1.0, DoubleType()), + # Decimal + (decimal.Decimal("1.0"), DecimalType()), + # Binary + (bytearray([1, 2]), BinaryType()), + # Date/Timestamp + (datetime.date(2000, 1, 2), DateType()), + (datetime.datetime(2000, 1, 2, 3, 4), DateType()), + (datetime.datetime(2000, 1, 2, 3, 4), TimestampType()), + # Array + ([], ArrayType(IntegerType())), + (["1", None], ArrayType(StringType(), containsNull=True)), + ([1, 2], ArrayType(IntegerType())), + ((1, 2), ArrayType(IntegerType())), + (array.array("h", [1, 2]), ArrayType(IntegerType())), + # Map + ({}, MapType(StringType(), IntegerType())), + ({"a": 1}, MapType(StringType(), IntegerType())), + ({"a": None}, MapType(StringType(), IntegerType(), valueContainsNull=True)), + # Struct + ({"s": "a", "i": 1}, schema), + ({"s": "a", "i": None}, schema), + ({"s": "a"}, schema), + ({"s": "a", "f": 1.0}, schema), + (Row(s="a", i=1), schema), + (Row(s="a", i=None), schema), + (["a", 1], schema), + (["a", None], schema), + (("a", 1), schema), + (MyObj(s="a", i=1), schema), + (MyObj(s="a", i=None), schema), + (MyObj(s="a"), schema), + ] + + # obj, data_type, exception class + failure_spec = [ + # String (match anything but None) + (None, StringType(), ValueError), + # CharType (match anything but None) + (None, CharType(10), ValueError), + # VarcharType (match anything but None) + (None, VarcharType(10), ValueError), + # UDT + (ExamplePoint(1.0, 2.0), PythonOnlyUDT(), ValueError), + # Boolean + (1, BooleanType(), TypeError), + ("True", BooleanType(), TypeError), + ([1], BooleanType(), TypeError), + # Byte + (-(2**7) - 1, ByteType(), ValueError), + (2**7, ByteType(), ValueError), + ("1", ByteType(), TypeError), + (1.0, ByteType(), TypeError), + # Short + (-(2**15) - 1, ShortType(), ValueError), + (2**15, ShortType(), ValueError), + # Integer + (-(2**31) - 1, IntegerType(), ValueError), + (2**31, IntegerType(), ValueError), + # Float & Double + (1, FloatType(), TypeError), + (1, DoubleType(), TypeError), + # Decimal + (1.0, DecimalType(), TypeError), + (1, DecimalType(), TypeError), + ("1.0", DecimalType(), TypeError), + # Binary + (1, BinaryType(), TypeError), + # Date/Timestamp + ("2000-01-02", DateType(), TypeError), + (946811040, TimestampType(), TypeError), + # Array + (["1", None], ArrayType(StringType(), containsNull=False), ValueError), + ([1, "2"], ArrayType(IntegerType()), TypeError), + # Map + ({"a": 1}, MapType(IntegerType(), IntegerType()), TypeError), + ({"a": "1"}, MapType(StringType(), IntegerType()), TypeError), + ( + {"a": None}, + MapType(StringType(), IntegerType(), valueContainsNull=False), + ValueError, + ), + # Struct + ({"s": "a", "i": "1"}, schema, TypeError), + (Row(s="a"), schema, ValueError), # Row can't have missing field + (Row(s="a", i="1"), schema, TypeError), + (["a"], schema, ValueError), + (["a", "1"], schema, TypeError), + (MyObj(s="a", i="1"), schema, TypeError), + (MyObj(s=None, i="1"), schema, ValueError), + ] + + # Check success cases + for obj, data_type in success_spec: + try: + _make_type_verifier(data_type, nullable=False)(obj) + except Exception: + self.fail("verify_type(%s, %s, nullable=False)" % (obj, data_type)) + + # Check failure cases + for obj, data_type, exp in failure_spec: + msg = "verify_type(%s, %s, nullable=False) == %s" % (obj, data_type, exp) + with self.assertRaises(exp, msg=msg): + _make_type_verifier(data_type, nullable=False)(obj) + + def test_row_without_field_sorting(self): + r = Row(b=1, a=2) + TestRow = Row("b", "a") + expected = TestRow(1, 2) + + self.assertEqual(r, expected) + self.assertEqual(repr(r), "Row(b=1, a=2)") + + +class TypesTests(TypesTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.test_types import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf.py new file mode 100644 index 0000000..2f8c1cd --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf.py @@ -0,0 +1,961 @@ +# +# 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 functools +import pydoc +import shutil +import tempfile +import unittest +import datetime + +from pyspark import SparkContext, SQLContext +from pyspark.sql import SparkSession, Column, Row +from pyspark.sql.functions import col, udf, assert_true, lit, rand +from pyspark.sql.udf import UserDefinedFunction +from pyspark.sql.types import ( + StringType, + IntegerType, + BooleanType, + DoubleType, + LongType, + ArrayType, + StructType, + StructField, + TimestampNTZType, + DayTimeIntervalType, +) +from pyspark.errors import AnalysisException, PythonException, PySparkTypeError +from pyspark.testing.sqlutils import ReusedSQLTestCase, test_compiled, test_not_compiled_message +from pyspark.testing.utils import QuietTest, assertDataFrameEqual + + +class BaseUDFTestsMixin(object): + def test_udf_with_callable(self): + data = self.spark.createDataFrame([(i, i**2) for i in range(10)], ["number", "squared"]) + + class PlusFour: + def __call__(self, col): + if col is not None: + return col + 4 + + call = PlusFour() + pudf = UserDefinedFunction(call, LongType()) + res = data.select(pudf(data["number"]).alias("plus_four")) + self.assertEqual(res.agg({"plus_four": "sum"}).collect()[0][0], 85) + + def test_udf_with_partial_function(self): + data = self.spark.createDataFrame([(i, i**2) for i in range(10)], ["number", "squared"]) + + def some_func(col, param): + if col is not None: + return col + param + + pfunc = functools.partial(some_func, param=4) + pudf = UserDefinedFunction(pfunc, LongType()) + res = data.select(pudf(data["number"]).alias("plus_four")) + self.assertEqual(res.agg({"plus_four": "sum"}).collect()[0][0], 85) + + def test_udf(self): + self.spark.catalog.registerFunction("twoArgs", lambda x, y: len(x) + y, IntegerType()) + [row] = self.spark.sql("SELECT twoArgs('test', 1)").collect() + self.assertEqual(row[0], 5) + + def test_udf_on_sql_context(self): + # This is to check if a deprecated 'SQLContext.registerFunction' can call its alias. + sqlContext = SQLContext.getOrCreate(self.spark.sparkContext) + sqlContext.registerFunction("oneArg", lambda x: len(x), IntegerType()) + [row] = sqlContext.sql("SELECT oneArg('test')").collect() + self.assertEqual(row[0], 4) + + def test_udf2(self): + with self.tempView("test"): + self.spark.catalog.registerFunction("strlen", lambda string: len(string), IntegerType()) + self.spark.createDataFrame([("test",)], ["a"]).createOrReplaceTempView("test") + [res] = self.spark.sql("SELECT strlen(a) FROM test WHERE strlen(a) > 1").collect() + self.assertEqual(4, res[0]) + + def test_udf3(self): + two_args = self.spark.catalog.registerFunction( + "twoArgs", UserDefinedFunction(lambda x, y: len(x) + y) + ) + self.assertEqual(two_args.deterministic, True) + [row] = self.spark.sql("SELECT twoArgs('test', 1)").collect() + self.assertEqual(row[0], "5") + + def test_udf_registration_return_type_none(self): + two_args = self.spark.catalog.registerFunction( + "twoArgs", UserDefinedFunction(lambda x, y: len(x) + y, "integer"), None + ) + self.assertEqual(two_args.deterministic, True) + [row] = self.spark.sql("SELECT twoArgs('test', 1)").collect() + self.assertEqual(row[0], 5) + + def test_udf_registration_return_type_not_none(self): + with QuietTest(self.sc): + self.check_udf_registration_return_type_not_none() + + def check_udf_registration_return_type_not_none(self): + # negative test for incorrect type + with self.assertRaises(PySparkTypeError) as pe: + self.spark.catalog.registerFunction( + "f", UserDefinedFunction(lambda x, y: len(x) + y, StringType()), StringType() + ) + + self.check_error( + exception=pe.exception, + error_class="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF", + message_parameters={"arg_name": "f", "return_type": "StringType()"}, + ) + + def test_nondeterministic_udf(self): + # Test that nondeterministic UDFs are evaluated only once in chained UDF evaluations + import random + + udf_random_col = udf(lambda: int(100 * random.random()), IntegerType()).asNondeterministic() + self.assertEqual(udf_random_col.deterministic, False) + df = self.spark.createDataFrame([Row(1)]).select(udf_random_col().alias("RAND")) + udf_add_ten = udf(lambda rand: rand + 10, IntegerType()) + [row] = df.withColumn("RAND_PLUS_TEN", udf_add_ten("RAND")).collect() + self.assertEqual(row[0] + 10, row[1]) + + def test_nondeterministic_udf2(self): + import random + + random_udf = udf(lambda: random.randint(6, 6), IntegerType()).asNondeterministic() + self.assertEqual(random_udf.deterministic, False) + random_udf1 = self.spark.catalog.registerFunction("randInt", random_udf) + self.assertEqual(random_udf1.deterministic, False) + [row] = self.spark.sql("SELECT randInt()").collect() + self.assertEqual(row[0], 6) + [row] = self.spark.range(1).select(random_udf1()).collect() + self.assertEqual(row[0], 6) + [row] = self.spark.range(1).select(random_udf()).collect() + self.assertEqual(row[0], 6) + # render_doc() reproduces the help() exception without printing output + pydoc.render_doc(udf(lambda: random.randint(6, 6), IntegerType())) + pydoc.render_doc(random_udf) + pydoc.render_doc(random_udf1) + pydoc.render_doc(udf(lambda x: x).asNondeterministic) + + def test_nondeterministic_udf3(self): + # regression test for SPARK-23233 + f = udf(lambda x: x) + # Here we cache the JVM UDF instance. + self.spark.range(1).select(f("id")) + # This should reset the cache to set the deterministic status correctly. + f = f.asNondeterministic() + # Check the deterministic status of udf. + df = self.spark.range(1).select(f("id")) + deterministic = df._jdf.logicalPlan().projectList().head().deterministic() + self.assertFalse(deterministic) + + def test_nondeterministic_udf_in_aggregate(self): + with QuietTest(self.sc): + self.check_nondeterministic_udf_in_aggregate() + + def check_nondeterministic_udf_in_aggregate(self): + from pyspark.sql.functions import sum + import random + + udf_random_col = udf(lambda: int(100 * random.random()), "int").asNondeterministic() + df = self.spark.range(10) + + with self.assertRaisesRegex(AnalysisException, "Non-deterministic"): + df.groupby("id").agg(sum(udf_random_col())).collect() + with self.assertRaisesRegex(AnalysisException, "Non-deterministic"): + df.agg(sum(udf_random_col())).collect() + + def test_chained_udf(self): + self.spark.catalog.registerFunction("double", lambda x: x + x, IntegerType()) + [row] = self.spark.sql("SELECT double(1)").collect() + self.assertEqual(row[0], 2) + [row] = self.spark.sql("SELECT double(double(1))").collect() + self.assertEqual(row[0], 4) + [row] = self.spark.sql("SELECT double(double(1) + 1)").collect() + self.assertEqual(row[0], 6) + + def test_single_udf_with_repeated_argument(self): + # regression test for SPARK-20685 + self.spark.catalog.registerFunction("add", lambda x, y: x + y, IntegerType()) + row = self.spark.sql("SELECT add(1, 1)").first() + self.assertEqual(tuple(row), (2,)) + + def test_multiple_udfs(self): + self.spark.catalog.registerFunction("double", lambda x: x * 2, IntegerType()) + [row] = self.spark.sql("SELECT double(1), double(2)").collect() + self.assertEqual(tuple(row), (2, 4)) + [row] = self.spark.sql("SELECT double(double(1)), double(double(2) + 2)").collect() + self.assertEqual(tuple(row), (4, 12)) + self.spark.catalog.registerFunction("add", lambda x, y: x + y, IntegerType()) + [row] = self.spark.sql("SELECT double(add(1, 2)), add(double(2), 1)").collect() + self.assertEqual(tuple(row), (6, 5)) + + def test_udf_in_filter_on_top_of_outer_join(self): + left = self.spark.createDataFrame([Row(a=1)]) + right = self.spark.createDataFrame([Row(a=1)]) + df = left.join(right, on="a", how="left_outer") + df = df.withColumn("b", udf(lambda x: "x")(df.a)) + self.assertEqual(df.filter('b = "x"').collect(), [Row(a=1, b="x")]) + + def test_udf_in_filter_on_top_of_join(self): + # regression test for SPARK-18589 + left = self.spark.createDataFrame([Row(a=1)]) + right = self.spark.createDataFrame([Row(b=1)]) + f = udf(lambda a, b: a == b, BooleanType()) + df = left.crossJoin(right).filter(f("a", "b")) + self.assertEqual(df.collect(), [Row(a=1, b=1)]) + + def test_udf_in_join_condition(self): + # regression test for SPARK-25314 + left = self.spark.createDataFrame([Row(a=1)]) + right = self.spark.createDataFrame([Row(b=1)]) + f = udf(lambda a, b: a == b, BooleanType()) + # The udf uses attributes from both sides of join, so it is pulled out as Filter + + # Cross join. + df = left.join(right, f("a", "b")) + with self.sql_conf({"spark.sql.crossJoin.enabled": False}): + with self.assertRaisesRegex(AnalysisException, "Detected implicit cartesian product"): + df.collect() + with self.sql_conf({"spark.sql.crossJoin.enabled": True}): + self.assertEqual(df.collect(), [Row(a=1, b=1)]) + + def test_udf_in_left_outer_join_condition(self): + # regression test for SPARK-26147 + from pyspark.sql.functions import col + + left = self.spark.createDataFrame([Row(a=1)]) + right = self.spark.createDataFrame([Row(b=1)]) + f = udf(lambda a: str(a), StringType()) + # The join condition can't be pushed down, as it refers to attributes from both sides. + # The Python UDF only refer to attributes from one side, so it's evaluable. + df = left.join(right, f("a") == col("b").cast("string"), how="left_outer") + with self.sql_conf({"spark.sql.crossJoin.enabled": True}): + self.assertEqual(df.collect(), [Row(a=1, b=1)]) + + def test_udf_and_common_filter_in_join_condition(self): + # regression test for SPARK-25314 + # test the complex scenario with both udf and common filter + left = self.spark.createDataFrame([Row(a=1, a1=1, a2=1), Row(a=2, a1=2, a2=2)]) + right = self.spark.createDataFrame([Row(b=1, b1=1, b2=1), Row(b=1, b1=3, b2=1)]) + f = udf(lambda a, b: a == b, BooleanType()) + df = left.join(right, [f("a", "b"), left.a1 == right.b1]) + # do not need spark.sql.crossJoin.enabled=true for udf is not the only join condition. + self.assertEqual(df.collect(), [Row(a=1, a1=1, a2=1, b=1, b1=1, b2=1)]) + + def test_udf_not_supported_in_join_condition(self): + # regression test for SPARK-25314 + # test python udf is not supported in join type except inner join. + left = self.spark.createDataFrame([Row(a=1, a1=1, a2=1), Row(a=2, a1=2, a2=2)]) + right = self.spark.createDataFrame([Row(b=1, b1=1, b2=1), Row(b=1, b1=3, b2=1)]) + f = udf(lambda a, b: a == b, BooleanType()) + + def runWithJoinType(join_type, type_string): + with self.assertRaisesRegex( + AnalysisException, + """Python UDF in the ON clause of a %s JOIN.""" % type_string, + ): + left.join(right, [f("a", "b"), left.a1 == right.b1], join_type).collect() + + runWithJoinType("full", "FULL OUTER") + runWithJoinType("left", "LEFT OUTER") + runWithJoinType("right", "RIGHT OUTER") + runWithJoinType("leftanti", "LEFT ANTI") + runWithJoinType("leftsemi", "LEFT SEMI") + + def test_udf_as_join_condition(self): + left = self.spark.createDataFrame([Row(a=1, a1=1, a2=1), Row(a=2, a1=2, a2=2)]) + right = self.spark.createDataFrame([Row(b=1, b1=1, b2=1), Row(b=1, b1=3, b2=1)]) + f = udf(lambda a: a, IntegerType()) + + df = left.join(right, [f("a") == f("b"), left.a1 == right.b1]) + self.assertEqual(df.collect(), [Row(a=1, a1=1, a2=1, b=1, b1=1, b2=1)]) + + def test_udf_without_arguments(self): + self.spark.catalog.registerFunction("foo", lambda: "bar") + [row] = self.spark.sql("SELECT foo()").collect() + self.assertEqual(row[0], "bar") + + def test_udf_with_array_type(self): + with self.tempView("test"): + self.spark.createDataFrame( + [ + ([0, 1, 2], {"key": [0, 1, 2, 3, 4]}), + ], + ["l", "d"], + ).createOrReplaceTempView("test") + self.spark.catalog.registerFunction( + "copylist", lambda l: list(l), ArrayType(IntegerType()) + ) + self.spark.catalog.registerFunction("maplen", lambda d: len(d), IntegerType()) + [(l1, l2)] = self.spark.sql("select copylist(l), maplen(d) from test").collect() + self.assertEqual(list(range(3)), l1) + self.assertEqual(1, l2) + + def test_broadcast_in_udf(self): + bar = {"a": "aa", "b": "bb", "c": "abc"} + foo = self.sc.broadcast(bar) + self.spark.catalog.registerFunction("MYUDF", lambda x: foo.value[x] if x else "") + [res] = self.spark.sql("SELECT MYUDF('c')").collect() + self.assertEqual("abc", res[0]) + [res] = self.spark.sql("SELECT MYUDF('')").collect() + self.assertEqual("", res[0]) + + def test_udf_with_filter_function(self): + df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) + from pyspark.sql.functions import col + + my_filter = udf(lambda a: a < 2, BooleanType()) + sel = df.select(col("key"), col("value")).filter((my_filter(col("key"))) & (df.value < "2")) + self.assertEqual(sel.collect(), [Row(key=1, value="1")]) + + def test_udf_with_aggregate_function(self): + df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) + from pyspark.sql.functions import col, sum + + my_filter = udf(lambda a: a == 1, BooleanType()) + sel = df.select(col("key")).distinct().filter(my_filter(col("key"))) + self.assertEqual(sel.collect(), [Row(key=1)]) + + my_copy = udf(lambda x: x, IntegerType()) + my_add = udf(lambda a, b: int(a + b), IntegerType()) + my_strlen = udf(lambda x: len(x), IntegerType()) + sel = ( + df.groupBy(my_copy(col("key")).alias("k")) + .agg(sum(my_strlen(col("value"))).alias("s")) + .select(my_add(col("k"), col("s")).alias("t")) + ) + self.assertEqual(sel.collect(), [Row(t=4), Row(t=3)]) + + def test_udf_in_generate(self): + from pyspark.sql.functions import explode + + df = self.spark.range(5) + f = udf(lambda x: list(range(x)), ArrayType(LongType())) + row = df.select(explode(f(*df))).groupBy().sum().first() + self.assertEqual(row[0], 10) + + df = self.spark.range(3) + res = df.select("id", explode(f(df.id))).collect() + self.assertEqual(res[0][0], 1) + self.assertEqual(res[0][1], 0) + self.assertEqual(res[1][0], 2) + self.assertEqual(res[1][1], 0) + self.assertEqual(res[2][0], 2) + self.assertEqual(res[2][1], 1) + + range_udf = udf(lambda value: list(range(value - 1, value + 1)), ArrayType(IntegerType())) + res = df.select("id", explode(range_udf(df.id))).collect() + self.assertEqual(res[0][0], 0) + self.assertEqual(res[0][1], -1) + self.assertEqual(res[1][0], 0) + self.assertEqual(res[1][1], 0) + self.assertEqual(res[2][0], 1) + self.assertEqual(res[2][1], 0) + self.assertEqual(res[3][0], 1) + self.assertEqual(res[3][1], 1) + + def test_udf_with_order_by_and_limit(self): + my_copy = udf(lambda x: x, IntegerType()) + df = self.spark.range(10).orderBy("id") + res = df.select(df.id, my_copy(df.id).alias("copy")).limit(1) + self.assertEqual(res.collect(), [Row(id=0, copy=0)]) + + def test_udf_registration_returns_udf(self): + df = self.spark.range(10) + add_three = self.spark.udf.register("add_three", lambda x: x + 3, IntegerType()) + + self.assertListEqual( + df.selectExpr("add_three(id) AS plus_three").collect(), + df.select(add_three("id").alias("plus_three")).collect(), + ) + + def test_udf_registration_returns_udf_on_sql_context(self): + df = self.spark.range(10) + + # This is to check if a 'SQLContext.udf' can call its alias. + sqlContext = SQLContext.getOrCreate(self.spark.sparkContext) + add_four = sqlContext.udf.register("add_four", lambda x: x + 4, IntegerType()) + + self.assertListEqual( + df.selectExpr("add_four(id) AS plus_four").collect(), + df.select(add_four("id").alias("plus_four")).collect(), + ) + + @unittest.skipIf(not test_compiled, test_not_compiled_message) # type: ignore + def test_register_java_function(self): + self.spark.udf.registerJavaFunction( + "javaStringLength", "test.org.apache.spark.sql.JavaStringLength", IntegerType() + ) + [value] = self.spark.sql("SELECT javaStringLength('test')").first() + self.assertEqual(value, 4) + + self.spark.udf.registerJavaFunction( + "javaStringLength2", "test.org.apache.spark.sql.JavaStringLength" + ) + [value] = self.spark.sql("SELECT javaStringLength2('test')").first() + self.assertEqual(value, 4) + + self.spark.udf.registerJavaFunction( + "javaStringLength3", "test.org.apache.spark.sql.JavaStringLength", "integer" + ) + [value] = self.spark.sql("SELECT javaStringLength3('test')").first() + self.assertEqual(value, 4) + + @unittest.skipIf(not test_compiled, test_not_compiled_message) # type: ignore + def test_register_java_udaf(self): + self.spark.udf.registerJavaUDAF("javaUDAF", "test.org.apache.spark.sql.MyDoubleAvg") + df = self.spark.createDataFrame([(1, "a"), (2, "b"), (3, "a")], ["id", "name"]) + df.createOrReplaceTempView("df") + row = self.spark.sql( + "SELECT name, javaUDAF(id) as avg from df group by name order by name desc" + ).first() + self.assertEqual(row.asDict(), Row(name="b", avg=102.0).asDict()) + + def test_non_existed_udf(self): + spark = self.spark + self.assertRaisesRegex( + AnalysisException, + "Can not load class non_existed_udf", + lambda: spark.udf.registerJavaFunction("udf1", "non_existed_udf"), + ) + + def test_non_existed_udf_with_sql_context(self): + # This is to check if a deprecated 'SQLContext.registerJavaFunction' can call its alias. + sqlContext = SQLContext.getOrCreate(self.spark.sparkContext) + self.assertRaisesRegex( + AnalysisException, + "Can not load class non_existed_udf", + lambda: sqlContext.registerJavaFunction("udf1", "non_existed_udf"), + ) + + def test_non_existed_udaf(self): + spark = self.spark + self.assertRaisesRegex( + AnalysisException, + "Can not load class non_existed_udaf", + lambda: spark.udf.registerJavaUDAF("udaf1", "non_existed_udaf"), + ) + + def test_udf_with_input_file_name(self): + from pyspark.sql.functions import input_file_name + + sourceFile = udf(lambda path: path, StringType()) + filePath = "python/test_support/sql/people1.json" + row = self.spark.read.json(filePath).select(sourceFile(input_file_name())).first() + self.assertTrue(row[0].find("people1.json") != -1) + + def test_udf_with_input_file_name_for_hadooprdd(self): + from pyspark.sql.functions import input_file_name + + def filename(path): + return path + + sameText = udf(filename, StringType()) + + rdd = self.sc.textFile("python/test_support/sql/people.json") + df = self.spark.read.json(rdd).select(input_file_name().alias("file")) + row = df.select(sameText(df["file"])).first() + self.assertTrue(row[0].find("people.json") != -1) + + rdd2 = self.sc.newAPIHadoopFile( + "python/test_support/sql/people.json", + "org.apache.hadoop.mapreduce.lib.input.TextInputFormat", + "org.apache.hadoop.io.LongWritable", + "org.apache.hadoop.io.Text", + ) + + df2 = self.spark.read.json(rdd2).select(input_file_name().alias("file")) + row2 = df2.select(sameText(df2["file"])).first() + self.assertTrue(row2[0].find("people.json") != -1) + + def test_udf_defers_judf_initialization(self): + # This is separate of UDFInitializationTests + # to avoid context initialization + # when udf is called + f = UserDefinedFunction(lambda x: x, StringType()) + + self.assertIsNone( + f._judf_placeholder, "judf should not be initialized before the first call." + ) + + self.assertIsInstance(f("foo"), Column, "UDF call should return a Column.") + + self.assertIsNotNone( + f._judf_placeholder, "judf should be initialized after UDF has been called." + ) + + def test_udf_with_string_return_type(self): + add_one = UserDefinedFunction(lambda x: x + 1, "integer") + make_pair = UserDefinedFunction(lambda x: (-x, x), "struct") + make_array = UserDefinedFunction( + lambda x: [float(x) for x in range(x, x + 3)], "array" + ) + + expected = (2, Row(x=-1, y=1), [1.0, 2.0, 3.0]) + actual = ( + self.spark.range(1, 2) + .toDF("x") + .select(add_one("x"), make_pair("x"), make_array("x")) + .first() + ) + + self.assertTupleEqual(expected, actual) + + def test_udf_should_not_accept_noncallable_object(self): + non_callable = None + self.assertRaises(TypeError, UserDefinedFunction, non_callable, StringType()) + + def test_udf_with_decorator(self): + from pyspark.sql.functions import lit + + @udf(IntegerType()) + def add_one(x): + if x is not None: + return x + 1 + + @udf(returnType=DoubleType()) + def add_two(x): + if x is not None: + return float(x + 2) + + @udf + def to_upper(x): + if x is not None: + return x.upper() + + @udf() + def to_lower(x): + if x is not None: + return x.lower() + + @udf + def substr(x, start, end): + if x is not None: + return x[start:end] + + @udf("long") + def trunc(x): + return int(x) + + @udf(returnType="double") + def as_double(x): + return float(x) + + df = self.spark.createDataFrame( + [(1, "Foo", "foobar", 3.0)], ("one", "Foo", "foobar", "float") + ).select( + add_one("one"), + add_two("one"), + to_upper("Foo"), + to_lower("Foo"), + substr("foobar", lit(0), lit(3)), + trunc("float"), + as_double("one"), + ) + + self.assertListEqual( + [tpe for _, tpe in df.dtypes], + ["int", "double", "string", "string", "string", "bigint", "double"], + ) + + self.assertListEqual(list(df.first()), [2, 3.0, "FOO", "foo", "foo", 3, 1.0]) + + def test_udf_wrapper(self): + def f(x): + """Identity""" + return x + + return_type = IntegerType() + f_ = udf(f, return_type) + + self.assertTrue(f.__doc__ in f_.__doc__) + self.assertEqual(f, f_.func) + self.assertEqual(return_type, f_.returnType) + + class F: + """Identity""" + + def __call__(self, x): + return x + + f = F() + return_type = IntegerType() + f_ = udf(f, return_type) + + self.assertTrue(f.__doc__ in f_.__doc__) + self.assertEqual(f, f_.func) + self.assertEqual(return_type, f_.returnType) + + f = functools.partial(f, x=1) + return_type = IntegerType() + f_ = udf(f, return_type) + + self.assertTrue(f.__doc__ in f_.__doc__) + self.assertEqual(f, f_.func) + self.assertEqual(return_type, f_.returnType) + + def test_udf_timestamp_ntz(self): + # SPARK-36626: Test TimestampNTZ in Python UDF + @udf(TimestampNTZType()) + def noop(x): + assert x == datetime.datetime(1970, 1, 1, 0, 0) + return x + + with self.sql_conf({"spark.sql.session.timeZone": "Pacific/Honolulu"}): + df = self.spark.createDataFrame( + [(datetime.datetime(1970, 1, 1, 0, 0),)], schema="dt timestamp_ntz" + ).select(noop("dt").alias("dt")) + + df.selectExpr("assert_true('1970-01-01 00:00:00' == CAST(dt AS STRING))").collect() + self.assertEqual(df.schema[0].dataType.typeName(), "timestamp_ntz") + self.assertEqual(df.first()[0], datetime.datetime(1970, 1, 1, 0, 0)) + + def test_udf_daytime_interval(self): + # SPARK-37277: Support DayTimeIntervalType in Python UDF + @udf(DayTimeIntervalType(DayTimeIntervalType.DAY, DayTimeIntervalType.SECOND)) + def noop(x): + assert x == datetime.timedelta(microseconds=123) + return x + + df = self.spark.createDataFrame( + [(datetime.timedelta(microseconds=123),)], schema="td interval day to second" + ).select(noop("td").alias("td")) + + df.select( + assert_true(lit("INTERVAL '0 00:00:00.000123' DAY TO SECOND") == df.td.cast("string")) + ).collect() + self.assertEqual(df.schema[0].dataType.simpleString(), "interval day to second") + self.assertEqual(df.first()[0], datetime.timedelta(microseconds=123)) + + def test_nonparam_udf_with_aggregate(self): + import pyspark.sql.functions as f + + df = self.spark.createDataFrame([(1, 2), (1, 2)]) + f_udf = f.udf(lambda: "const_str") + rows = df.distinct().withColumn("a", f_udf()).collect() + self.assertEqual(rows, [Row(_1=1, _2=2, a="const_str")]) + + # SPARK-24721 + @unittest.skipIf(not test_compiled, test_not_compiled_message) # type: ignore + def test_datasource_with_udf(self): + from pyspark.sql.functions import lit, col + + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + self.spark.range(1).write.mode("overwrite").format("csv").save(path) + filesource_df = self.spark.read.option("inferSchema", True).csv(path).toDF("i") + datasource_df = ( + self.spark.read.format("org.apache.spark.sql.sources.SimpleScanSource") + .option("from", 0) + .option("to", 1) + .load() + .toDF("i") + ) + datasource_v2_df = ( + self.spark.read.format("org.apache.spark.sql.connector.SimpleDataSourceV2") + .load() + .toDF("i", "j") + ) + + c1 = udf(lambda x: x + 1, "int")(lit(1)) + c2 = udf(lambda x: x + 1, "int")(col("i")) + + f1 = udf(lambda x: False, "boolean")(lit(1)) + f2 = udf(lambda x: False, "boolean")(col("i")) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + result = df.withColumn("c", c1) + expected = df.withColumn("c", lit(2)) + self.assertEqual(expected.collect(), result.collect()) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + result = df.withColumn("c", c2) + expected = df.withColumn("c", col("i") + 1) + self.assertEqual(expected.collect(), result.collect()) + + for df in [filesource_df, datasource_df, datasource_v2_df]: + for f in [f1, f2]: + result = df.filter(f) + self.assertEqual(0, result.count()) + finally: + shutil.rmtree(path) + + # SPARK-42134 + def test_file_dsv2_with_udf_filter(self): + from pyspark.sql.functions import lit + + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + with self.sql_conf({"spark.sql.sources.useV1SourceList": ""}): + self.spark.range(1).write.mode("overwrite").format("parquet").save(path) + df = self.spark.read.parquet(path).toDF("i") + f = udf(lambda x: False, "boolean")(lit(1)) + result = df.filter(f) + self.assertEqual(0, result.count()) + + finally: + shutil.rmtree(path) + + # SPARK-25591 + def test_same_accumulator_in_udfs(self): + data_schema = StructType( + [StructField("a", IntegerType(), True), StructField("b", IntegerType(), True)] + ) + data = self.spark.createDataFrame([[1, 2]], schema=data_schema) + + test_accum = self.sc.accumulator(0) + + def first_udf(x): + test_accum.add(1) + return x + + def second_udf(x): + test_accum.add(100) + return x + + func_udf = udf(first_udf, IntegerType()) + func_udf2 = udf(second_udf, IntegerType()) + data = data.withColumn("out1", func_udf(data["a"])) + data = data.withColumn("out2", func_udf2(data["b"])) + data.collect() + self.assertEqual(test_accum.value, 101) + + # SPARK-26293 + def test_udf_in_subquery(self): + f = udf(lambda x: x, "long") + with self.tempView("v"): + self.spark.range(1).filter(f("id") >= 0).createTempView("v") + result = self.spark.sql( + "select i from values(0L) as data(i) where i in (select id from v)" + ) + self.assertEqual(result.collect(), [Row(i=0)]) + + def test_udf_globals_not_overwritten(self): + @udf("string") + def f(): + assert "itertools" not in str(map) + + self.spark.range(1).select(f()).collect() + + def test_worker_original_stdin_closed(self): + # Test if it closes the original standard input of worker inherited from the daemon, + # and replaces it with '/dev/null'. See SPARK-26175. + def task(iterator): + import sys + + res = sys.stdin.read() + # Because the standard input is '/dev/null', it reaches to EOF. + assert res == "", "Expect read EOF from stdin." + return iterator + + self.sc.parallelize(range(1), 1).mapPartitions(task).count() + + def test_udf_with_256_args(self): + N = 256 + data = [["data-%d" % i for i in range(N)]] * 5 + df = self.spark.createDataFrame(data) + + def f(*a): + return "success" + + fUdf = udf(f, StringType()) + + r = df.select(fUdf(*df.columns)) + self.assertEqual(r.first()[0], "success") + + def test_udf_cache(self): + def func(x): + return x + + df = self.spark.range(1) + df.select(udf(func)("id")).cache() + + self.assertEqual( + df.select(udf(func)("id")) + ._jdf.queryExecution() + .withCachedData() + .getClass() + .getSimpleName(), + "InMemoryRelation", + ) + + # SPARK-34545 + def test_udf_input_serialization_valuecompare_disabled(self): + def f(e): + return e[0] + + df = self.spark.createDataFrame([((1.0, 1.0), (1, 1))], ["c1", "c2"]) + result = df.select( + "*", udf(f, DoubleType())("c1").alias("c3"), udf(f, IntegerType())("c2").alias("c4") + ) + self.assertEqual( + result.collect(), [Row(c1=Row(_1=1.0, _2=1.0), c2=Row(_1=1, _2=1), c3=1.0, c4=1)] + ) + + # SPARK-33277 + def test_udf_with_column_vector(self): + path = tempfile.mkdtemp() + shutil.rmtree(path) + + try: + self.spark.range(0, 100000, 1, 1).write.parquet(path) + + def f(x): + return 0 + + fUdf = udf(f, LongType()) + + for offheap in ["true", "false"]: + with self.sql_conf({"spark.sql.columnVector.offheap.enabled": offheap}): + self.assertEquals( + self.spark.read.parquet(path).select(fUdf("id")).head(), Row(0) + ) + finally: + shutil.rmtree(path) + + def test_udf_with_rand(self): + # SPARK-40121: rand() with Python UDF. + self.assertEqual( + len(self.spark.range(10).select(udf(lambda x: x, DoubleType())(rand())).collect()), 10 + ) + + def test_nested_struct(self): + df = self.spark.range(1).selectExpr( + "struct(1, struct('John', 30, ('value', 10))) as nested_struct" + ) + # Input + row = df.select(udf(lambda x: str(x))("nested_struct")).first() + self.assertEquals( + row[0], "Row(col1=1, col2=Row(col1='John', col2=30, col3=Row(col1='value', col2=10)))" + ) + # Output + row = df.select(udf(lambda x: x, returnType=df.dtypes[0][1])("nested_struct")).first() + self.assertEquals( + row[0], Row(col1=1, col2=Row(col1="John", col2=30, col3=Row(col1="value", col2=10))) + ) + + def test_nested_map(self): + df = self.spark.range(1).selectExpr("map('a', map('b', 'c')) as nested_map") + # Input + row = df.select(udf(lambda x: str(x))("nested_map")).first() + self.assertEquals(row[0], "{'a': {'b': 'c'}}") + # Output + + @udf(returnType=df.dtypes[0][1]) + def f(x): + x["a"]["b"] = "d" + return x + + row = df.select(f("nested_map")).first() + self.assertEquals(row[0], {"a": {"b": "d"}}) + + def test_nested_array(self): + df = self.spark.range(1).selectExpr("array(array(1, 2), array(3, 4)) as nested_array") + # Input + row = df.select(udf(lambda x: str(x))("nested_array")).first() + self.assertEquals(row[0], "[[1, 2], [3, 4]]") + # Output + + @udf(returnType=df.dtypes[0][1]) + def f(x): + x.append([4, 5]) + return x + + row = df.select(f("nested_array")).first() + self.assertEquals(row[0], [[1, 2], [3, 4], [4, 5]]) + + def test_complex_return_types(self): + row = ( + self.spark.range(1) + .selectExpr("array(1, 2, 3) as array", "map('a', 'b') as map", "struct(1, 2) as struct") + .select( + udf(lambda x: x, "array")("array"), + udf(lambda x: x, "map")("map"), + udf(lambda x: x, "struct")("struct"), + ) + .first() + ) + + self.assertEquals(row[0], [1, 2, 3]) + self.assertEquals(row[1], {"a": "b"}) + self.assertEquals(row[2], Row(col1=1, col2=2)) + + def test_raise_stop_iteration(self): + @udf("int") + def test_udf(a): + if a < 5: + return a + else: + raise StopIteration() + + assertDataFrameEqual( + self.spark.range(5).select(test_udf(col("id"))), [Row(i) for i in range(5)] + ) + + with self.assertRaisesRegex(PythonException, "StopIteration"): + self.spark.range(10).select(test_udf(col("id"))).show() + + +class UDFTests(BaseUDFTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + super(BaseUDFTestsMixin, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false") + + +class UDFInitializationTests(unittest.TestCase): + def tearDown(self): + if SparkSession._instantiatedSession is not None: + SparkSession._instantiatedSession.stop() + + if SparkContext._active_spark_context is not None: + SparkContext._active_spark_context.stop() + + def test_udf_init_should_not_initialize_context(self): + UserDefinedFunction(lambda x: x, StringType()) + + self.assertIsNone( + SparkContext._active_spark_context, + "SparkContext shouldn't be initialized when UserDefinedFunction is created.", + ) + self.assertIsNone( + SparkSession._instantiatedSession, + "SparkSession shouldn't be initialized when UserDefinedFunction is created.", + ) + + def test_err_parse_type_when_no_sc(self): + with self.assertRaisesRegex( + RuntimeError, + "SparkContext or SparkSession should be created first", + ): + udf(lambda x: x, "integer") + + +if __name__ == "__main__": + from pyspark.sql.tests.test_udf import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf_profiler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf_profiler.py new file mode 100644 index 0000000..136f423 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udf_profiler.py @@ -0,0 +1,114 @@ +# +# 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 tempfile +import unittest +import os +import sys +from io import StringIO + +from pyspark import SparkConf +from pyspark.sql import SparkSession +from pyspark.sql.functions import udf +from pyspark.profiler import UDFBasicProfiler + + +class UDFProfilerTests(unittest.TestCase): + def setUp(self): + self._old_sys_path = list(sys.path) + class_name = self.__class__.__name__ + conf = SparkConf().set("spark.python.profile", "true") + self.spark = ( + SparkSession.builder.master("local[4]") + .config(conf=conf) + .appName(class_name) + .getOrCreate() + ) + self.sc = self.spark.sparkContext + + def tearDown(self): + self.spark.stop() + sys.path = self._old_sys_path + + def test_udf_profiler(self): + self.do_computation() + + profilers = self.sc.profiler_collector.profilers + self.assertEqual(3, len(profilers)) + + old_stdout = sys.stdout + try: + sys.stdout = io = StringIO() + self.sc.show_profiles() + finally: + sys.stdout = old_stdout + + d = tempfile.gettempdir() + self.sc.dump_profiles(d) + + for i, udf_name in enumerate(["add1", "add2", "add1"]): + id, profiler, _ = profilers[i] + with self.subTest(id=id, udf_name=udf_name): + stats = profiler.stats() + self.assertTrue(stats is not None) + width, stat_list = stats.get_print_list([]) + func_names = [func_name for fname, n, func_name in stat_list] + self.assertTrue(udf_name in func_names) + + self.assertTrue(udf_name in io.getvalue()) + self.assertTrue("udf_%d.pstats" % id in os.listdir(d)) + + def test_custom_udf_profiler(self): + class TestCustomProfiler(UDFBasicProfiler): + def show(self, id): + self.result = "Custom formatting" + + self.sc.profiler_collector.udf_profiler_cls = TestCustomProfiler + + self.do_computation() + + profilers = self.sc.profiler_collector.profilers + self.assertEqual(3, len(profilers)) + _, profiler, _ = profilers[0] + self.assertTrue(isinstance(profiler, TestCustomProfiler)) + + self.sc.show_profiles() + self.assertEqual("Custom formatting", profiler.result) + + def do_computation(self): + @udf + def add1(x): + return x + 1 + + @udf + def add2(x): + return x + 2 + + df = self.spark.range(10) + df.select(add1("id"), add2("id"), add1("id")).collect() + + +if __name__ == "__main__": + from pyspark.sql.tests.test_udf_profiler import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udtf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udtf.py new file mode 100644 index 0000000..944ce6d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_udtf.py @@ -0,0 +1,1456 @@ +# +# 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 os +import tempfile +import unittest +from typing import Iterator + +from py4j.protocol import Py4JJavaError + +from pyspark.errors import ( + PySparkAttributeError, + PythonException, + PySparkTypeError, + AnalysisException, + PySparkRuntimeError, +) +from pyspark.rdd import PythonEvalType +from pyspark.sql.functions import lit, udf, udtf +from pyspark.sql.types import IntegerType, MapType, Row, StringType, StructType +from pyspark.testing import assertDataFrameEqual +from pyspark.testing.sqlutils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, + ReusedSQLTestCase, +) + + +class BaseUDTFTestsMixin: + def test_simple_udtf(self): + class TestUDTF: + def eval(self): + yield "hello", "world" + + func = udtf(TestUDTF, returnType="c1: string, c2: string") + rows = func().collect() + self.assertEqual(rows, [Row(c1="hello", c2="world")]) + + def test_udtf_yield_single_row_col(self): + class TestUDTF: + def eval(self, a: int): + yield a, + + func = udtf(TestUDTF, returnType="a: int") + rows = func(lit(1)).collect() + self.assertEqual(rows, [Row(a=1)]) + + def test_udtf_yield_multi_cols(self): + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + func = udtf(TestUDTF, returnType="a: int, b: int") + rows = func(lit(1)).collect() + self.assertEqual(rows, [Row(a=1, b=2)]) + + def test_udtf_yield_multi_rows(self): + class TestUDTF: + def eval(self, a: int): + yield a, + yield a + 1, + + func = udtf(TestUDTF, returnType="a: int") + rows = func(lit(1)).collect() + self.assertEqual(rows, [Row(a=1), Row(a=2)]) + + def test_udtf_yield_multi_row_col(self): + class TestUDTF: + def eval(self, a: int, b: int): + yield a, b, a + b + yield a, b, a - b + yield a, b, b - a + + func = udtf(TestUDTF, returnType="a: int, b: int, c: int") + rows = func(lit(1), lit(2)).collect() + self.assertEqual(rows, [Row(a=1, b=2, c=3), Row(a=1, b=2, c=-1), Row(a=1, b=2, c=1)]) + + def test_udtf_decorator(self): + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + rows = TestUDTF(lit(1)).collect() + self.assertEqual(rows, [Row(a=1, b=2)]) + + def test_udtf_registration(self): + class TestUDTF: + def eval(self, a: int, b: int): + yield a, b, a + b + yield a, b, a - b + yield a, b, b - a + + func = udtf(TestUDTF, returnType="a: int, b: int, c: int") + self.spark.udtf.register("testUDTF", func) + df = self.spark.sql("SELECT * FROM testUDTF(1, 2)") + self.assertEqual( + df.collect(), [Row(a=1, b=2, c=3), Row(a=1, b=2, c=-1), Row(a=1, b=2, c=1)] + ) + + def test_udtf_with_lateral_join(self): + class TestUDTF: + def eval(self, a: int, b: int) -> Iterator: + yield a, b, a + b + yield a, b, a - b + + func = udtf(TestUDTF, returnType="a: int, b: int, c: int") + self.spark.udtf.register("testUDTF", func) + df = self.spark.sql( + "SELECT f.* FROM values (0, 1), (1, 2) t(a, b), LATERAL testUDTF(a, b) f" + ) + expected = self.spark.createDataFrame( + [(0, 1, 1), (0, 1, -1), (1, 2, 3), (1, 2, -1)], schema=["a", "b", "c"] + ) + self.assertEqual(df.collect(), expected.collect()) + + def test_udtf_eval_with_return_stmt(self): + class TestUDTF: + def eval(self, a: int, b: int): + return [(a, a + 1), (b, b + 1)] + + func = udtf(TestUDTF, returnType="a: int, b: int") + rows = func(lit(1), lit(2)).collect() + self.assertEqual(rows, [Row(a=1, b=2), Row(a=2, b=3)]) + + def test_udtf_eval_returning_non_tuple(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield a + + with self.assertRaisesRegex(PythonException, "UDTF_INVALID_OUTPUT_ROW_TYPE"): + TestUDTF(lit(1)).collect() + + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + return (a,) + + with self.assertRaisesRegex(PythonException, "UDTF_INVALID_OUTPUT_ROW_TYPE"): + TestUDTF(lit(1)).collect() + + def test_udtf_with_invalid_return_value(self): + @udtf(returnType="x: int") + class TestUDTF: + def eval(self, a): + return a + + with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"): + TestUDTF(lit(1)).collect() + + def test_udtf_with_zero_arg_and_invalid_return_value(self): + @udtf(returnType="x: int") + class TestUDTF: + def eval(self): + return 1 + + with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"): + TestUDTF().collect() + + def test_udtf_with_invalid_return_value_in_terminate(self): + @udtf(returnType="x: int") + class TestUDTF: + def eval(self, a): + ... + + def terminate(self): + return 1 + + with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"): + TestUDTF(lit(1)).collect() + + def test_udtf_eval_with_no_return(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + ... + + self.assertEqual(TestUDTF(lit(1)).collect(), []) + + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + return + + self.assertEqual(TestUDTF(lit(1)).collect(), []) + + def test_udtf_with_conditional_return(self): + class TestUDTF: + def eval(self, a: int): + if a > 5: + yield a, + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql("SELECT * FROM range(0, 8) JOIN LATERAL test_udtf(id)").collect(), + [Row(id=6, a=6), Row(id=7, a=7)], + ) + + def test_udtf_with_empty_yield(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield + + assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=None)]) + + def test_udtf_with_none_output(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield a, + yield None, + + self.assertEqual(TestUDTF(lit(1)).collect(), [Row(a=1), Row(a=None)]) + df = self.spark.createDataFrame([(0, 1), (1, 2)], schema=["a", "b"]) + self.assertEqual(TestUDTF(lit(1)).join(df, "a", "inner").collect(), [Row(a=1, b=2)]) + assertDataFrameEqual( + TestUDTF(lit(1)).join(df, "a", "left"), [Row(a=None, b=None), Row(a=1, b=2)] + ) + + def test_udtf_with_none_input(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield a, + + self.assertEqual(TestUDTF(lit(None)).collect(), [Row(a=None)]) + self.spark.udtf.register("testUDTF", TestUDTF) + df = self.spark.sql("SELECT * FROM testUDTF(null)") + self.assertEqual(df.collect(), [Row(a=None)]) + + def test_udtf_with_wrong_num_input(self): + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + with self.assertRaisesRegex( + PythonException, r"eval\(\) missing 1 required positional argument: 'a'" + ): + TestUDTF().collect() + + with self.assertRaisesRegex( + PythonException, r"eval\(\) takes 2 positional arguments but 3 were given" + ): + TestUDTF(lit(1), lit(2)).collect() + + def test_udtf_init_with_additional_args(self): + @udtf(returnType="x int") + class TestUDTF: + def __init__(self, a: int): + ... + + def eval(self, a: int): + yield a, + + with self.assertRaisesRegex( + PythonException, r"__init__\(\) missing 1 required positional argument: 'a'" + ): + TestUDTF(lit(1)).show() + + def test_udtf_terminate_with_additional_args(self): + @udtf(returnType="x int") + class TestUDTF: + def eval(self, a: int): + yield a, + + def terminate(self, a: int): + ... + + with self.assertRaisesRegex( + PythonException, r"terminate\(\) missing 1 required positional argument: 'a'" + ): + TestUDTF(lit(1)).show() + + def test_udtf_with_wrong_num_output(self): + err_msg = ( + r"\[UDTF_RETURN_SCHEMA_MISMATCH\] The number of columns in the " + "result does not match the specified schema." + ) + + # Output less columns than specified return schema + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, + + with self.assertRaisesRegex(PythonException, err_msg): + TestUDTF(lit(1)).collect() + + # Output more columns than specified return schema + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + with self.assertRaisesRegex(PythonException, err_msg): + TestUDTF(lit(1)).collect() + + def test_udtf_with_empty_output_schema_and_non_empty_output(self): + @udtf(returnType=StructType()) + class TestUDTF: + def eval(self): + yield 1, + + with self.assertRaisesRegex(PythonException, "UDTF_RETURN_SCHEMA_MISMATCH"): + TestUDTF().collect() + + def test_udtf_with_non_empty_output_schema_and_empty_output(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self): + yield tuple() + + with self.assertRaisesRegex(PythonException, "UDTF_RETURN_SCHEMA_MISMATCH"): + TestUDTF().collect() + + def test_udtf_init(self): + @udtf(returnType="a: int, b: int, c: string") + class TestUDTF: + def __init__(self): + self.key = "test" + + def eval(self, a: int): + yield a, a + 1, self.key + + rows = TestUDTF(lit(1)).collect() + self.assertEqual(rows, [Row(a=1, b=2, c="test")]) + + def test_udtf_terminate(self): + @udtf(returnType="key: string, value: float") + class TestUDTF: + def __init__(self): + self._count = 0 + self._sum = 0 + + def eval(self, x: int): + self._count += 1 + self._sum += x + yield "input", float(x) + + def terminate(self): + yield "count", float(self._count) + yield "avg", self._sum / self._count + + self.assertEqual( + TestUDTF(lit(1)).collect(), + [Row(key="input", value=1), Row(key="count", value=1.0), Row(key="avg", value=1.0)], + ) + + self.spark.udtf.register("test_udtf", TestUDTF) + df = self.spark.sql( + "SELECT id, key, value FROM range(0, 10, 1, 2), " + "LATERAL test_udtf(id) WHERE key != 'input'" + ) + self.assertEqual( + df.collect(), + [ + Row(id=4, key="count", value=5.0), + Row(id=4, key="avg", value=2.0), + Row(id=9, key="count", value=5.0), + Row(id=9, key="avg", value=7.0), + ], + ) + + def test_init_with_exception(self): + @udtf(returnType="x: int") + class TestUDTF: + def __init__(self): + raise Exception("error") + + def eval(self): + yield 1, + + with self.assertRaisesRegex( + PythonException, + r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error " + r"in the '__init__' method: error", + ): + TestUDTF().show() + + def test_eval_with_exception(self): + @udtf(returnType="x: int") + class TestUDTF: + def eval(self): + raise Exception("error") + + with self.assertRaisesRegex( + PythonException, + r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error " + r"in the 'eval' method: error", + ): + TestUDTF().show() + + def test_terminate_with_exceptions(self): + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + def terminate(self): + raise ValueError("terminate error") + + with self.assertRaisesRegex( + PythonException, + r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error " + r"in the 'terminate' method: terminate error", + ): + TestUDTF(lit(1)).collect() + + def test_udtf_terminate_with_wrong_num_output(self): + err_msg = ( + r"\[UDTF_RETURN_SCHEMA_MISMATCH\] The number of columns in the result " + "does not match the specified schema." + ) + + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + def terminate(self): + yield 1, 2, 3 + + with self.assertRaisesRegex(PythonException, err_msg): + TestUDTF(lit(1)).show() + + @udtf(returnType="a: int, b: int") + class TestUDTF: + def eval(self, a: int): + yield a, a + 1 + + def terminate(self): + yield 1, + + with self.assertRaisesRegex(PythonException, err_msg): + TestUDTF(lit(1)).show() + + def test_udtf_determinism(self): + class TestUDTF: + def eval(self, a: int): + yield a, + + func = udtf(TestUDTF, returnType="x: int") + # The UDTF is marked as non-deterministic by default. + self.assertFalse(func.deterministic) + func = func.asDeterministic() + self.assertTrue(func.deterministic) + + def test_nondeterministic_udtf(self): + import random + + class RandomUDTF: + def eval(self, a: int): + yield a + int(random.random()), + + random_udtf = udtf(RandomUDTF, returnType="x: int") + assertDataFrameEqual(random_udtf(lit(1)), [Row(x=1)]) + self.spark.udtf.register("random_udtf", random_udtf) + assertDataFrameEqual(self.spark.sql("select * from random_udtf(1)"), [Row(x=1)]) + + def test_udtf_with_nondeterministic_input(self): + from pyspark.sql.functions import rand + + @udtf(returnType="x: int") + class TestUDTF: + def eval(self, a: int): + yield 1 if a > 100 else 0, + + assertDataFrameEqual(TestUDTF(rand(0) * 100), [Row(x=0)]) + + def test_udtf_with_invalid_return_type(self): + @udtf(returnType="int") + class TestUDTF: + def eval(self, a: int): + yield a + 1, + + with self.assertRaises(PySparkTypeError) as e: + TestUDTF(lit(1)).collect() + + self.check_error( + exception=e.exception, + error_class="UDTF_RETURN_TYPE_MISMATCH", + message_parameters={"name": "TestUDTF", "return_type": "IntegerType()"}, + ) + + @udtf(returnType=MapType(StringType(), IntegerType())) + class TestUDTF: + def eval(self, a: int): + yield a + 1, + + with self.assertRaises(PySparkTypeError) as e: + TestUDTF(lit(1)).collect() + + self.check_error( + exception=e.exception, + error_class="UDTF_RETURN_TYPE_MISMATCH", + message_parameters={ + "name": "TestUDTF", + "return_type": "MapType(StringType(), IntegerType(), True)", + }, + ) + + def test_udtf_with_struct_input_type(self): + @udtf(returnType="x: string") + class TestUDTF: + def eval(self, person): + yield f"{person.name}: {person.age}", + + self.spark.udtf.register("test_udtf", TestUDTF) + self.assertEqual( + self.spark.sql( + "select * from test_udtf(named_struct('name', 'Alice', 'age', 1))" + ).collect(), + [Row(x="Alice: 1")], + ) + + def test_udtf_with_array_input_type(self): + @udtf(returnType="x: string") + class TestUDTF: + def eval(self, args): + yield str(args), + + self.spark.udtf.register("test_udtf", TestUDTF) + self.assertEqual( + self.spark.sql("select * from test_udtf(array(1, 2, 3))").collect(), + [Row(x="[1, 2, 3]")], + ) + + def test_udtf_with_map_input_type(self): + @udtf(returnType="x: string") + class TestUDTF: + def eval(self, m): + yield str(m), + + self.spark.udtf.register("test_udtf", TestUDTF) + self.assertEqual( + self.spark.sql("select * from test_udtf(map('key', 'value'))").collect(), + [Row(x="{'key': 'value'}")], + ) + + def test_udtf_with_struct_output_types(self): + @udtf(returnType="x: struct") + class TestUDTF: + def eval(self, x: int): + yield {"a": x, "b": x + 1}, + + self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=Row(a=1, b=2))]) + + def test_udtf_with_array_output_types(self): + @udtf(returnType="x: array") + class TestUDTF: + def eval(self, x: int): + yield [x, x + 1, x + 2], + + self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=[1, 2, 3])]) + + def test_udtf_with_map_output_types(self): + @udtf(returnType="x: map") + class TestUDTF: + def eval(self, x: int): + yield {x: str(x)}, + + self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x={1: "1"})]) + + def test_udtf_with_empty_output_types(self): + @udtf(returnType=StructType()) + class TestUDTF: + def eval(self): + yield tuple() + + assertDataFrameEqual(TestUDTF(), [Row()]) + + def _check_result_or_exception( + self, func_handler, ret_type, expected, *, err_type=PythonException + ): + func = udtf(func_handler, returnType=ret_type) + if not isinstance(expected, str): + assertDataFrameEqual(func(), expected) + else: + with self.assertRaisesRegex(err_type, expected): + func().collect() + + def test_numeric_output_type_casting(self): + class TestUDTF: + def eval(self): + yield 1, + + for i, (ret_type, expected) in enumerate( + [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=1)]), + ("x: smallint", [Row(x=1)]), + ("x: int", [Row(x=1)]), + ("x: bigint", [Row(x=1)]), + ("x: string", [Row(x="1")]), # int to string is ok, but string to int is None + ( + "x: date", + "AttributeError", + ), # AttributeError: 'int' object has no attribute 'toordinal' + ( + "x: timestamp", + "AttributeError", + ), # AttributeError: 'int' object has no attribute 'tzinfo' + ("x: byte", [Row(x=1)]), + ("x: binary", [Row(x=None)]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=None)]), + ("x: map", [Row(x=None)]), + ("x: struct", "UNEXPECTED_TUPLE_WITH_STRUCT"), + ] + ): + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_numeric_string_output_type_casting(self): + class TestUDTF: + def eval(self): + yield "1", + + for ret_type, expected in [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=None)]), + ("x: smallint", [Row(x=None)]), + ("x: int", [Row(x=None)]), + ("x: bigint", [Row(x=None)]), + ("x: string", [Row(x="1")]), + ("x: date", "AttributeError"), + ("x: timestamp", "AttributeError"), + ("x: byte", [Row(x=None)]), + ("x: binary", [Row(x=bytearray(b"1"))]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=None)]), + ("x: map", [Row(x=None)]), + ("x: struct", "UNEXPECTED_TUPLE_WITH_STRUCT"), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_string_output_type_casting(self): + class TestUDTF: + def eval(self): + yield "hello", + + for ret_type, expected in [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=None)]), + ("x: smallint", [Row(x=None)]), + ("x: int", [Row(x=None)]), + ("x: bigint", [Row(x=None)]), + ("x: string", [Row(x="hello")]), + ("x: date", "AttributeError"), + ("x: timestamp", "AttributeError"), + ("x: byte", [Row(x=None)]), + ("x: binary", [Row(x=bytearray(b"hello"))]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=None)]), + ("x: map", [Row(x=None)]), + ("x: struct", "UNEXPECTED_TUPLE_WITH_STRUCT"), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_array_output_type_casting(self): + class TestUDTF: + def eval(self): + yield [0, 1.1, 2], + + for ret_type, expected in [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=None)]), + ("x: smallint", [Row(x=None)]), + ("x: int", [Row(x=None)]), + ("x: bigint", [Row(x=None)]), + ("x: string", [Row(x="[0, 1.1, 2]")]), + ("x: date", "AttributeError"), + ("x: timestamp", "AttributeError"), + ("x: byte", [Row(x=None)]), + ("x: binary", [Row(x=None)]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=[0, None, 2])]), + ("x: array", [Row(x=[None, 1.1, None])]), + ("x: array", [Row(x=["0", "1.1", "2"])]), + ("x: array", [Row(x=[None, None, None])]), + ("x: array>", [Row(x=[None, None, None])]), + ("x: map", [Row(x=None)]), + ("x: struct", [Row(x=Row(a=0, b=None, c=2))]), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_map_output_type_casting(self): + class TestUDTF: + def eval(self): + yield {"a": 0, "b": 1.1, "c": 2}, + + for ret_type, expected in [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=None)]), + ("x: smallint", [Row(x=None)]), + ("x: int", [Row(x=None)]), + ("x: bigint", [Row(x=None)]), + ("x: string", [Row(x="{a=0, b=1.1, c=2}")]), + ("x: date", "AttributeError"), + ("x: timestamp", "AttributeError"), + ("x: byte", [Row(x=None)]), + ("x: binary", [Row(x=None)]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=None)]), + ("x: map", [Row(x={"a": "0", "b": "1.1", "c": "2"})]), + ("x: map", [Row(x={"a": None, "b": None, "c": None})]), + ("x: map", [Row(x={"a": 0, "b": None, "c": 2})]), + ("x: map", [Row(x={"a": None, "b": 1.1, "c": None})]), + ("x: map>", [Row(x={"a": None, "b": None, "c": None})]), + ("x: struct", [Row(x=Row(a=0))]), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_struct_output_type_casting_dict(self): + class TestUDTF: + def eval(self): + yield {"a": 0, "b": 1.1, "c": 2}, + + for ret_type, expected in [ + ("x: boolean", [Row(x=None)]), + ("x: tinyint", [Row(x=None)]), + ("x: smallint", [Row(x=None)]), + ("x: int", [Row(x=None)]), + ("x: bigint", [Row(x=None)]), + ("x: string", [Row(x="{a=0, b=1.1, c=2}")]), + ("x: date", "AttributeError"), + ("x: timestamp", "AttributeError"), + ("x: byte", [Row(x=None)]), + ("x: binary", [Row(x=None)]), + ("x: float", [Row(x=None)]), + ("x: double", [Row(x=None)]), + ("x: decimal(10, 0)", [Row(x=None)]), + ("x: array", [Row(x=None)]), + ("x: map", [Row(x={"a": "0", "b": "1.1", "c": "2"})]), + ("x: struct", [Row(Row(a="0", b="1.1", c="2"))]), + ("x: struct", [Row(Row(a=0, b=None, c=2))]), + ("x: struct", [Row(Row(a=None, b=1.1, c=None))]), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_struct_output_type_casting_row(self): + self.check_struct_output_type_casting_row(Py4JJavaError) + + def check_struct_output_type_casting_row(self, error_type): + class TestUDTF: + def eval(self): + yield Row(a=0, b=1.1, c=2), + + err = ("PickleException", error_type) + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", err), + ("x: date", "ValueError"), + ("x: timestamp", "ValueError"), + ("x: byte", err), + ("x: binary", err), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", err), + ("x: map", err), + ("x: struct", [Row(Row(a="0", b="1.1", c="2"))]), + ("x: struct", [Row(Row(a=0, b=None, c=2))]), + ("x: struct", [Row(Row(a=None, b=1.1, c=None))]), + ]: + with self.subTest(ret_type=ret_type): + if isinstance(expected, tuple): + self._check_result_or_exception( + TestUDTF, ret_type, expected[0], err_type=expected[1] + ) + else: + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_inconsistent_output_types(self): + class TestUDTF: + def eval(self): + yield 1, + yield [1, 2], + + for ret_type, expected in [ + ("x: int", [Row(x=1), Row(x=None)]), + ("x: array", [Row(x=None), Row(x=[1, 2])]), + ]: + with self.subTest(ret_type=ret_type): + assertDataFrameEqual(udtf(TestUDTF, returnType=ret_type)(), expected) + + @unittest.skipIf(not have_pandas, pandas_requirement_message) + def test_udtf_with_pandas_input_type(self): + import pandas as pd + + @udtf(returnType="corr: double") + class TestUDTF: + def eval(self, s1: pd.Series, s2: pd.Series): + yield s1.corr(s2) + + self.spark.udtf.register("test_udtf", TestUDTF) + # TODO(SPARK-43968): check during compile time instead of runtime + with self.assertRaisesRegex( + PythonException, "AttributeError: 'int' object has no attribute 'corr'" + ): + self.spark.sql( + "select * from values (1, 2), (2, 3) t(a, b), lateral test_udtf(a, b)" + ).collect() + + def test_udtf_register_error(self): + @udf + def upper(s: str): + return s.upper() + + with self.assertRaises(PySparkTypeError) as e: + self.spark.udtf.register("test_udf", upper) + + self.check_error( + exception=e.exception, + error_class="INVALID_UDTF_EVAL_TYPE", + message_parameters={ + "name": "test_udf", + "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF", + }, + ) + + def test_udtf_pickle_error(self): + with tempfile.TemporaryDirectory() as d: + file = os.path.join(d, "file.txt") + file_obj = open(file, "w") + + @udtf(returnType="x: int") + class TestUDTF: + def eval(self): + file_obj + yield 1, + + with self.assertRaisesRegex(PySparkRuntimeError, "UDTF_SERIALIZATION_ERROR"): + TestUDTF().collect() + + def test_udtf_access_spark_session(self): + df = self.spark.range(10) + + @udtf(returnType="x: int") + class TestUDTF: + def eval(self): + df.collect() + yield 1, + + with self.assertRaisesRegex(PySparkRuntimeError, "UDTF_SERIALIZATION_ERROR"): + TestUDTF().collect() + + def test_udtf_no_eval(self): + with self.assertRaises(PySparkAttributeError) as e: + + @udtf(returnType="a: int, b: int") + class TestUDTF: + def run(self, a: int): + yield a, a + 1 + + self.check_error( + exception=e.exception, + error_class="INVALID_UDTF_NO_EVAL", + message_parameters={"name": "TestUDTF"}, + ) + + def test_udtf_with_no_handler_class(self): + with self.assertRaises(PySparkTypeError) as e: + + @udtf(returnType="a: int") + def test_udtf(a: int): + yield a, + + self.check_error( + exception=e.exception, + error_class="INVALID_UDTF_HANDLER_TYPE", + message_parameters={"type": "function"}, + ) + + with self.assertRaises(PySparkTypeError) as e: + udtf(1, returnType="a: int") + + self.check_error( + exception=e.exception, + error_class="INVALID_UDTF_HANDLER_TYPE", + message_parameters={"type": "int"}, + ) + + def test_udtf_with_table_argument_query(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))").collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_int_and_table_argument_query(self): + class TestUDTF: + def eval(self, i: int, row: Row): + if row["id"] > i: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql( + "SELECT * FROM test_udtf(5, TABLE (SELECT id FROM range(0, 8)))" + ).collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_table_argument_identifier(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + + with self.tempView("v"): + self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)") + self.assertEqual( + self.spark.sql("SELECT * FROM test_udtf(TABLE (v))").collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_int_and_table_argument_identifier(self): + class TestUDTF: + def eval(self, i: int, row: Row): + if row["id"] > i: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + + with self.tempView("v"): + self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)") + self.assertEqual( + self.spark.sql("SELECT * FROM test_udtf(5, TABLE (v))").collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_table_argument_unknown_identifier(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + + with self.assertRaisesRegex(AnalysisException, "TABLE_OR_VIEW_NOT_FOUND"): + self.spark.sql("SELECT * FROM test_udtf(TABLE (v))").collect() + + def test_udtf_with_table_argument_malformed_query(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + + with self.assertRaisesRegex(AnalysisException, "TABLE_OR_VIEW_NOT_FOUND"): + self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT * FROM v))").collect() + + def test_udtf_with_table_argument_cte_inside(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql( + """ + SELECT * FROM test_udtf(TABLE ( + WITH t AS ( + SELECT id FROM range(0, 8) + ) + SELECT * FROM t + )) + """ + ).collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_table_argument_cte_outside(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql( + """ + WITH t AS ( + SELECT id FROM range(0, 8) + ) + SELECT * FROM test_udtf(TABLE (SELECT id FROM t)) + """ + ).collect(), + [Row(a=6), Row(a=7)], + ) + + self.assertEqual( + self.spark.sql( + """ + WITH t AS ( + SELECT id FROM range(0, 8) + ) + SELECT * FROM test_udtf(TABLE (t)) + """ + ).collect(), + [Row(a=6), Row(a=7)], + ) + + # TODO(SPARK-44233): Fix the subquery resolution. + @unittest.skip("Fails to resolve the subquery.") + def test_udtf_with_table_argument_lateral_join(self): + class TestUDTF: + def eval(self, row: Row): + if row["id"] > 5: + yield row["id"], + + func = udtf(TestUDTF, returnType="a: int") + self.spark.udtf.register("test_udtf", func) + self.assertEqual( + self.spark.sql( + """ + SELECT * FROM + range(0, 8) AS t, + LATERAL test_udtf(TABLE (t)) + """ + ).collect(), + [Row(a=6), Row(a=7)], + ) + + def test_udtf_with_table_argument_multiple(self): + class TestUDTF: + def eval(self, a: Row, b: Row): + yield a[0], b[0] + + func = udtf(TestUDTF, returnType="a: int, b: int") + self.spark.udtf.register("test_udtf", func) + + query = """ + SELECT * FROM test_udtf( + TABLE (SELECT id FROM range(0, 2)), + TABLE (SELECT id FROM range(0, 3))) + """ + + with self.sql_conf({"spark.sql.tvf.allowMultipleTableArguments.enabled": False}): + with self.assertRaisesRegex( + AnalysisException, "TABLE_VALUED_FUNCTION_TOO_MANY_TABLE_ARGUMENTS" + ): + self.spark.sql(query).collect() + + with self.sql_conf({"spark.sql.tvf.allowMultipleTableArguments.enabled": True}): + self.assertEqual( + self.spark.sql(query).collect(), + [ + Row(a=0, b=0), + Row(a=1, b=0), + Row(a=0, b=1), + Row(a=1, b=1), + Row(a=0, b=2), + Row(a=1, b=2), + ], + ) + + def test_docstring(self): + class TestUDTF: + """A UDTF for test.""" + + def __init__(self): + """Initialize the UDTF""" + ... + + def eval(self, x: int): + """Evaluate the input row.""" + yield x + 1, + + def terminate(self): + """Terminate the UDTF.""" + ... + + cls = udtf(TestUDTF, returnType="y: int").func + self.assertIn("A UDTF for test", cls.__doc__) + self.assertIn("Initialize the UDTF", cls.__init__.__doc__) + self.assertIn("Evaluate the input row", cls.eval.__doc__) + self.assertIn("Terminate the UDTF", cls.terminate.__doc__) + + +class UDTFTests(BaseUDTFTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + super(UDTFTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "false") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled") + finally: + super(UDTFTests, cls).tearDownClass() + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message +) +class UDTFArrowTestsMixin(BaseUDTFTestsMixin): + def test_eval_type(self): + def upper(x: str): + return upper(x) + + class TestUDTF: + def eval(self, x: str): + return upper(x) + + self.assertEqual( + udtf(TestUDTF, returnType="x: string", useArrow=False).evalType, + PythonEvalType.SQL_TABLE_UDF, + ) + + self.assertEqual( + udtf(TestUDTF, returnType="x: string", useArrow=True).evalType, + PythonEvalType.SQL_ARROW_TABLE_UDF, + ) + + def test_udtf_arrow_sql_conf(self): + class TestUDTF: + def eval(self): + yield 1, + + # We do not use `self.sql_conf` here to test the SQL SET command + # instead of using PySpark's `spark.conf.set`. + old_value = self.spark.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled") + self.spark.sql("SET spark.sql.execution.pythonUDTF.arrow.enabled=False") + self.assertEqual(udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_TABLE_UDF) + self.spark.sql("SET spark.sql.execution.pythonUDTF.arrow.enabled=True") + self.assertEqual( + udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_ARROW_TABLE_UDF + ) + self.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", old_value) + + def test_udtf_eval_returning_non_tuple(self): + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + yield a + + # When arrow is enabled, it can handle non-tuple return value. + assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1)]) + + @udtf(returnType="a: int") + class TestUDTF: + def eval(self, a: int): + return [a] + + assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1)]) + + def test_numeric_output_type_casting(self): + class TestUDTF: + def eval(self): + yield 1, + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", [Row(x=True)]), + ("x: tinyint", [Row(x=1)]), + ("x: smallint", [Row(x=1)]), + ("x: int", [Row(x=1)]), + ("x: bigint", [Row(x=1)]), + ("x: string", [Row(x="1")]), # require arrow.cast + ("x: date", err), + ("x: byte", [Row(x=1)]), + ("x: binary", [Row(x=bytearray(b"\x01"))]), + ("x: float", [Row(x=1.0)]), + ("x: double", [Row(x=1.0)]), + ("x: decimal(10, 0)", err), + ("x: array", err), + ("x: map", err), + ("x: struct", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_numeric_string_output_type_casting(self): + class TestUDTF: + def eval(self): + yield "1", + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", [Row(x=True)]), + ("x: tinyint", [Row(x=1)]), + ("x: smallint", [Row(x=1)]), + ("x: int", [Row(x=1)]), + ("x: bigint", [Row(x=1)]), + ("x: string", [Row(x="1")]), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", [Row(x=1)]), + ("x: binary", [Row(x=bytearray(b"1"))]), + ("x: float", [Row(x=1.0)]), + ("x: double", [Row(x=1.0)]), + ("x: decimal(10, 0)", [Row(x=1)]), + ("x: array", [Row(x=["1"])]), + ("x: array", [Row(x=[1])]), + ("x: map", err), + ("x: struct", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_string_output_type_casting(self): + class TestUDTF: + def eval(self): + yield "hello", + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", [Row(x="hello")]), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", err), + ("x: binary", [Row(x=bytearray(b"hello"))]), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", [Row(x=["h", "e", "l", "l", "o"])]), + ("x: array", err), + ("x: map", err), + ("x: struct", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_array_output_type_casting(self): + class TestUDTF: + def eval(self): + yield [0, 1.1, 2], + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", err), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", err), + ("x: binary", err), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", [Row(x=["0", "1.1", "2"])]), + ("x: array", [Row(x=[False, True, True])]), + ("x: array", [Row(x=[0, 1, 2])]), + ("x: array", [Row(x=[0, 1.1, 2])]), + ("x: array>", err), + ("x: map", err), + ("x: struct", err), + ("x: struct", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_map_output_type_casting(self): + class TestUDTF: + def eval(self): + yield {"a": 0, "b": 1.1, "c": 2}, + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", err), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", err), + ("x: binary", err), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", [Row(x=["a", "b", "c"])]), + ("x: map", err), + ("x: map", err), + ("x: map", [Row(x={"a": 0, "b": 1, "c": 2})]), + ("x: map", [Row(x={"a": 0, "b": 1.1, "c": 2})]), + ("x: map>", err), + ("x: struct", [Row(x=Row(a=0))]), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_struct_output_type_casting_dict(self): + class TestUDTF: + def eval(self): + yield {"a": 0, "b": 1.1, "c": 2}, + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", err), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", err), + ("x: binary", err), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", [Row(x=["a", "b", "c"])]), + ("x: map", err), + ("x: struct", [Row(Row(a="0", b="1.1", c="2"))]), + ("x: struct", [Row(Row(a=0, b=1, c=2))]), + ("x: struct", [Row(Row(a=0, b=1.1, c=2))]), + ("x: struct,b:struct<>,c:struct<>>", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_struct_output_type_casting_row(self): + class TestUDTF: + def eval(self): + yield Row(a=0, b=1.1, c=2), + + err = "UDTF_ARROW_TYPE_CAST_ERROR" + + for ret_type, expected in [ + ("x: boolean", err), + ("x: tinyint", err), + ("x: smallint", err), + ("x: int", err), + ("x: bigint", err), + ("x: string", err), + ("x: date", err), + ("x: timestamp", err), + ("x: byte", err), + ("x: binary", err), + ("x: float", err), + ("x: double", err), + ("x: decimal(10, 0)", err), + ("x: array", [Row(x=["0", "1.1", "2"])]), + ("x: map", err), + ("x: struct", [Row(Row(a="0", b="1.1", c="2"))]), + ("x: struct", [Row(Row(a=0, b=1, c=2))]), + ("x: struct", [Row(Row(a=0, b=1.1, c=2))]), + ("x: struct,b:struct<>,c:struct<>>", err), + ]: + with self.subTest(ret_type=ret_type): + self._check_result_or_exception(TestUDTF, ret_type, expected) + + def test_inconsistent_output_types(self): + class TestUDTF: + def eval(self): + yield 1, + yield [1, 2], + + for ret_type in [ + "x: int", + "x: array", + ]: + with self.subTest(ret_type=ret_type): + with self.assertRaisesRegex(PythonException, "UDTF_ARROW_TYPE_CAST_ERROR"): + udtf(TestUDTF, returnType=ret_type)().collect() + + +class UDTFArrowTests(UDTFArrowTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + super(UDTFArrowTests, cls).setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true") + + @classmethod + def tearDownClass(cls): + try: + cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled") + finally: + super(UDTFArrowTests, cls).tearDownClass() + + +if __name__ == "__main__": + from pyspark.sql.tests.test_udtf import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_utils.py new file mode 100644 index 0000000..e1b7f29 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/test_utils.py @@ -0,0 +1,1674 @@ +# -*- encoding: utf-8 -*- +# +# 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 unittest +import difflib +from itertools import zip_longest + +from pyspark.sql.functions import sha2, to_timestamp +from pyspark.errors import ( + AnalysisException, + ParseException, + PySparkAssertionError, + IllegalArgumentException, + SparkUpgradeException, +) +from pyspark.testing.utils import assertDataFrameEqual, assertSchemaEqual, _context_diff, have_numpy +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.sql import Row +import pyspark.sql.functions as F +from pyspark.sql.functions import to_date, unix_timestamp, from_unixtime +from pyspark.sql.types import ( + StringType, + ArrayType, + LongType, + StructType, + MapType, + FloatType, + DoubleType, + StructField, + IntegerType, + BooleanType, +) +from pyspark.testing.sqlutils import have_pandas + + +class UtilsTestsMixin: + def test_assert_equal_inttype(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 3000), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 3000), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_arraytype(self): + df1 = self.spark.createDataFrame( + data=[ + ("john", ["Python", "Java"]), + ("jane", ["Scala", "SQL", "Java"]), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField("languages", ArrayType(StringType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("john", ["Python", "Java"]), + ("jane", ["Scala", "SQL", "Java"]), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField("languages", ArrayType(StringType()), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_approx_equal_arraytype_float(self): + df1 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.34]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.339999]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_approx_equal_arraytype_float_default_rtol_fail(self): + # fails with default rtol, 1e-5 + df1 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.34]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.341]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), df2.collect())): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=2 + ) + + error_msg = "Results do not match: " + percent_diff = (1 / 2) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_assert_approx_equal_arraytype_float_custom_rtol_pass(self): + # passes with custom rtol, 1e-2 + df1 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.34]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", [97.01, 89.23]), + ("student2", [91.86, 84.341]), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", ArrayType(FloatType()), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, rtol=1e-2) + + def test_assert_approx_equal_doubletype_custom_rtol_pass(self): + # passes with custom rtol, 1e-2 + df1 = self.spark.createDataFrame( + data=[ + ("student1", 97.01), + ("student2", 84.34), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grade", DoubleType(), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", 97.01), + ("student2", 84.341), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grade", DoubleType(), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, rtol=1e-2) + + def test_assert_approx_equal_decimaltype_custom_rtol_pass(self): + # passes with custom rtol, 1e-2 + df1 = self.spark.createDataFrame( + data=[ + ("student1", 83.14), + ("student2", 97.12), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grade", DoubleType(), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", 83.14), + ("student2", 97.111), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grade", DoubleType(), True), + ] + ), + ) + + # cast to DecimalType + df1 = df1.withColumn("col_1", F.col("grade").cast("decimal(4,3)")) + df2 = df2.withColumn("col_1", F.col("grade").cast("decimal(4,3)")) + + assertDataFrameEqual(df1, df2, rtol=1e-1) + + def test_assert_notequal_arraytype(self): + df1 = self.spark.createDataFrame( + data=[ + ("Amy", ["C++", "Rust"]), + ("John", ["Python", "Java"]), + ("Jane", ["Scala", "SQL", "Java"]), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField("languages", ArrayType(StringType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("Amy", ["C++", "Rust"]), + ("John", ["Python", "Java"]), + ("Jane", ["Scala", "Java"]), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField("languages", ArrayType(StringType()), True), + ] + ), + ) + + rows_str1 = "" + rows_str2 = "" + + sorted_list1 = sorted(df1.collect(), key=lambda x: str(x)) + sorted_list2 = sorted(df2.collect(), key=lambda x: str(x)) + + # count different rows + for r1, r2 in list(zip_longest(sorted_list1, sorted_list2)): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=3 + ) + + error_msg = "Results do not match: " + percent_diff = (1 / 3) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), df2.collect())): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=3 + ) + + error_msg = "Results do not match: " + percent_diff = (1 / 3) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_assert_equal_maptype(self): + df1 = self.spark.createDataFrame( + data=[ + ("student1", {"id": 222342203655477580}), + ("student2", {"id": 422322203155477692}), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("properties", MapType(StringType(), LongType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", {"id": 222342203655477580}), + ("student2", {"id": 422322203155477692}), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("properties", MapType(StringType(), LongType()), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_approx_equal_maptype_double(self): + df1 = self.spark.createDataFrame( + data=[ + ("student1", {"math": 76.23, "english": 92.64}), + ("student2", {"math": 87.89, "english": 84.48}), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", MapType(StringType(), DoubleType()), True), + ] + ), + ) + df2 = self.spark.createDataFrame( + data=[ + ("student1", {"math": 76.23, "english": 92.63999999}), + ("student2", {"math": 87.89, "english": 84.48}), + ], + schema=StructType( + [ + StructField("student", StringType(), True), + StructField("grades", MapType(StringType(), DoubleType()), True), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_approx_equal_nested_struct_double(self): + df1 = self.spark.createDataFrame( + data=[ + ("jane", (64.57, 76.63, 97.81)), + ("john", (93.92, 91.57, 84.36)), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField( + "grades", + StructType( + [ + StructField("math", DoubleType(), True), + StructField("english", DoubleType(), True), + StructField("biology", DoubleType(), True), + ] + ), + ), + ] + ), + ) + + df2 = self.spark.createDataFrame( + data=[ + ("jane", (64.57, 76.63, 97.81000001)), + ("john", (93.92, 91.57, 84.36)), + ], + schema=StructType( + [ + StructField("name", StringType(), True), + StructField( + "grades", + StructType( + [ + StructField("math", DoubleType(), True), + StructField("english", DoubleType(), True), + StructField("biology", DoubleType(), True), + ] + ), + ), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_nested_struct_str(self): + df1 = self.spark.createDataFrame( + data=[ + (1, ("jane", "anne", "doe")), + (2, ("john", "bob", "smith")), + ], + schema=StructType( + [ + StructField("id", IntegerType(), True), + StructField( + "name", + StructType( + [ + StructField("first", StringType(), True), + StructField("middle", StringType(), True), + StructField("last", StringType(), True), + ] + ), + ), + ] + ), + ) + + df2 = self.spark.createDataFrame( + data=[ + (1, ("jane", "anne", "doe")), + (2, ("john", "bob", "smith")), + ], + schema=StructType( + [ + StructField("id", IntegerType(), True), + StructField( + "name", + StructType( + [ + StructField("first", StringType(), True), + StructField("middle", StringType(), True), + StructField("last", StringType(), True), + ] + ), + ), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_nested_struct_str_duplicate(self): + df1 = self.spark.createDataFrame( + data=[ + (1, ("jane doe", "jane doe")), + (2, ("john smith", "john smith")), + ], + schema=StructType( + [ + StructField("id", IntegerType(), True), + StructField( + "full name", + StructType( + [ + StructField("name", StringType(), True), + StructField("name", StringType(), True), + ] + ), + ), + ] + ), + ) + + df2 = self.spark.createDataFrame( + data=[ + (1, ("jane doe", "jane doe")), + (2, ("john smith", "john smith")), + ], + schema=StructType( + [ + StructField("id", IntegerType(), True), + StructField( + "full name", + StructType( + [ + StructField("name", StringType(), True), + StructField("name", StringType(), True), + ] + ), + ), + ] + ), + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_duplicate_col(self): + df1 = self.spark.createDataFrame( + data=[ + (1, "Python", 1, 1), + (2, "Scala", 2, 2), + ], + schema=["number", "language", "number", "number"], + ) + df2 = self.spark.createDataFrame( + data=[ + (1, "Python", 1, 1), + (2, "Scala", 2, 2), + ], + schema=["number", "language", "number", "number"], + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_timestamp(self): + df1 = self.spark.createDataFrame( + data=[("1", "2023-01-01 12:01:01.000")], schema=["id", "timestamp"] + ) + + df2 = self.spark.createDataFrame( + data=[("1", "2023-01-01 12:01:01.000")], schema=["id", "timestamp"] + ) + + df1 = df1.withColumn("timestamp", to_timestamp("timestamp")) + df2 = df2.withColumn("timestamp", to_timestamp("timestamp")) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_equal_nullrow(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000), + (None, None), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000), + (None, None), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_notequal_nullval(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 2000), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", None), + ], + schema=["id", "amount"], + ) + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), df2.collect())): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=2 + ) + + error_msg = "Results do not match: " + percent_diff = (1 / 2) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_assert_equal_nulldf(self): + df1 = None + df2 = None + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_unequal_null_actual(self): + df1 = None + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 3000), + ], + schema=["id", "amount"], + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": None, + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": None, + }, + ) + + def test_assert_unequal_null_expected(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 3000), + ], + schema=["id", "amount"], + ) + df2 = None + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "expected", + "actual_type": None, + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "expected", + "actual_type": None, + }, + ) + + @unittest.skipIf(not have_pandas or not have_numpy, "no pandas or numpy dependency") + def test_assert_equal_exact_pandas_df(self): + import pandas as pd + import numpy as np + + df1 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 9)]), columns=["a", "b", "c"] + ) + df2 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 9)]), columns=["a", "b", "c"] + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + @unittest.skipIf(not have_pandas or not have_numpy, "no pandas or numpy dependency") + def test_assert_approx_equal_pandas_df(self): + import pandas as pd + import numpy as np + + # test that asserts close enough equality for pandas df + df1 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 59)]), columns=["a", "b", "c"] + ) + df2 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 59.0001)]), columns=["a", "b", "c"] + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + @unittest.skipIf(not have_pandas or not have_numpy, "no pandas or numpy dependency") + def test_assert_approx_equal_fail_exact_pandas_df(self): + import pandas as pd + import numpy as np + + # test that asserts close enough equality for pandas df + df1 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 59)]), columns=["a", "b", "c"] + ) + df2 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 59.0001)]), columns=["a", "b", "c"] + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=False, rtol=0, atol=0) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True, rtol=0, atol=0) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + @unittest.skipIf(not have_pandas or not have_numpy, "no pandas or numpy dependency") + def test_assert_unequal_pandas_df(self): + import pandas as pd + import numpy as np + + df1 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (6, 5, 4)]), columns=["a", "b", "c"] + ) + df2 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 9)]), columns=["a", "b", "c"] + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=False) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + @unittest.skipIf(not have_pandas or not have_numpy, "no pandas or numpy dependency") + def test_assert_type_error_pandas_df(self): + import pyspark.pandas as ps + import pandas as pd + import numpy as np + + df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) + df2 = pd.DataFrame( + data=np.array([(1, 2, 3), (4, 5, 6), (6, 5, 4)]), columns=["a", "b", "c"] + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=False) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": df1.to_string(), + "left_dtype": str(df1.dtypes), + "right": df2.to_string(), + "right_dtype": str(df2.dtypes), + }, + ) + + @unittest.skipIf(not have_pandas, "no pandas dependency") + def test_assert_equal_exact_pandas_on_spark_df(self): + import pyspark.pandas as ps + + df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) + df2 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + @unittest.skipIf(not have_pandas, "no pandas dependency") + def test_assert_equal_exact_pandas_on_spark_df(self): + import pyspark.pandas as ps + + df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) + df2 = ps.DataFrame(data=[30, 20, 10], columns=["Numbers"]) + + assertDataFrameEqual(df1, df2) + + @unittest.skipIf(not have_pandas, "no pandas dependency") + def test_assert_equal_approx_pandas_on_spark_df(self): + import pyspark.pandas as ps + + df1 = ps.DataFrame(data=[10.0001, 20.32, 30.1], columns=["Numbers"]) + df2 = ps.DataFrame(data=[10.0, 20.32, 30.1], columns=["Numbers"]) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + @unittest.skipIf(not have_pandas, "no pandas dependency") + def test_assert_error_pandas_pyspark_df(self): + import pyspark.pandas as ps + import pandas as pd + + df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) + df2 = self.spark.createDataFrame([(10,), (11,), (13,)], ["Numbers"]) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=False) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{ps.DataFrame.__name__}, " + f"{pd.DataFrame.__name__}, " + f"{ps.Series.__name__}, " + f"{pd.Series.__name__}, " + f"{ps.Index.__name__}" + f"{pd.Index.__name__}, ", + "arg_name": "expected", + "actual_type": type(df2), + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{ps.DataFrame.__name__}, " + f"{pd.DataFrame.__name__}, " + f"{ps.Series.__name__}, " + f"{pd.Series.__name__}, " + f"{ps.Index.__name__}" + f"{pd.Index.__name__}, ", + "arg_name": "expected", + "actual_type": type(df2), + }, + ) + + def test_assert_error_non_pyspark_df(self): + dict1 = {"a": 1, "b": 2} + dict2 = {"a": 1, "b": 2} + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(dict1, dict2) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": type(dict1), + }, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(dict1, dict2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": type(dict1), + }, + ) + + def test_row_order_ignored(self): + # test that row order is ignored (not checked) by default + df1 = self.spark.createDataFrame( + data=[ + ("2", 3000.00), + ("1", 1000.00), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000.00), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2) + + def test_check_row_order_error(self): + # test checkRowOrder=True + df1 = self.spark.createDataFrame( + data=[ + ("2", 3000.00), + ("1", 1000.00), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000.00), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), df2.collect())): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=2 + ) + + error_msg = "Results do not match: " + percent_diff = (2 / 2) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_remove_non_word_characters_long(self): + def remove_non_word_characters(col): + return F.regexp_replace(col, "[^\\w\\s]+", "") + + source_data = [("jo&&se",), ("**li**",), ("#::luisa",), (None,)] + source_df = self.spark.createDataFrame(source_data, ["name"]) + + actual_df = source_df.withColumn("clean_name", remove_non_word_characters(F.col("name"))) + + expected_data = [("jo&&se", "jose"), ("**li**", "li"), ("#::luisa", "luisa"), (None, None)] + expected_df = self.spark.createDataFrame(expected_data, ["name", "clean_name"]) + + assertDataFrameEqual(actual_df, expected_df) + + def test_assert_pyspark_approx_equal(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000.00), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000.0000001), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_assert_pyspark_approx_equal_custom_rtol(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000.00), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000.01), + ("2", 3000.00), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2, rtol=1e-2) + + def test_assert_pyspark_df_not_equal(self): + df1 = self.spark.createDataFrame( + data=[ + ("1", 1000.00), + ("2", 3000.00), + ("3", 2000.00), + ], + schema=["id", "amount"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1001.00), + ("2", 3000.00), + ("3", 2003.00), + ], + schema=["id", "amount"], + ) + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), df2.collect())): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=3 + ) + + error_msg = "Results do not match: " + percent_diff = (2 / 3) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_assert_notequal_schema(self): + df1 = self.spark.createDataFrame( + data=[ + (1, 1000), + (2, 3000), + ], + schema=["id", "number"], + ) + df2 = self.spark.createDataFrame( + data=[ + ("1", 1000), + ("2", 5000), + ], + schema=["id", "amount"], + ) + + generated_diff = difflib.ndiff(str(df1.schema).splitlines(), str(df2.schema).splitlines()) + + expected_error_msg = "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": expected_error_msg}, + ) + + def test_diff_schema_lens(self): + df1 = self.spark.createDataFrame( + data=[ + (1, 3000), + (2, 1000), + ], + schema=["id", "amount"], + ) + + df2 = self.spark.createDataFrame( + data=[ + (1, 3000, "a"), + (2, 1000, "b"), + ], + schema=["id", "amount", "letter"], + ) + + generated_diff = difflib.ndiff(str(df1.schema).splitlines(), str(df2.schema).splitlines()) + + expected_error_msg = "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, df2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": expected_error_msg}, + ) + + def test_schema_ignore_nullable(self): + s1 = StructType( + [StructField("id", IntegerType(), True), StructField("name", StringType(), True)] + ) + + df1 = self.spark.createDataFrame([(1, "jane"), (2, "john")], s1) + + s2 = StructType( + [StructField("id", IntegerType(), True), StructField("name", StringType(), False)] + ) + + df2 = self.spark.createDataFrame([(1, "jane"), (2, "john")], s2) + + assertDataFrameEqual(df1, df2) + + def test_schema_ignore_nullable_array_equal(self): + s1 = StructType([StructField("names", ArrayType(DoubleType(), True), True)]) + s2 = StructType([StructField("names", ArrayType(DoubleType(), False), False)]) + + assertSchemaEqual(s1, s2) + + def test_schema_ignore_nullable_struct_equal(self): + s1 = StructType( + [StructField("names", StructType([StructField("age", IntegerType(), True)]), True)] + ) + s2 = StructType( + [StructField("names", StructType([StructField("age", IntegerType(), False)]), False)] + ) + assertSchemaEqual(s1, s2) + + def test_schema_array_unequal(self): + s1 = StructType([StructField("names", ArrayType(IntegerType(), True), True)]) + s2 = StructType([StructField("names", ArrayType(DoubleType(), False), False)]) + + generated_diff = difflib.ndiff(str(s1).splitlines(), str(s2).splitlines()) + + expected_error_msg = "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertSchemaEqual(s1, s2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": expected_error_msg}, + ) + + def test_schema_struct_unequal(self): + s1 = StructType( + [StructField("names", StructType([StructField("age", DoubleType(), True)]), True)] + ) + s2 = StructType( + [StructField("names", StructType([StructField("age", IntegerType(), True)]), True)] + ) + + generated_diff = difflib.ndiff(str(s1).splitlines(), str(s2).splitlines()) + + expected_error_msg = "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertSchemaEqual(s1, s2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": expected_error_msg}, + ) + + def test_schema_more_nested_struct_unequal(self): + s1 = StructType( + [ + StructField( + "name", + StructType( + [ + StructField("firstname", StringType(), True), + StructField("middlename", StringType(), True), + StructField("lastname", StringType(), True), + ] + ), + ), + ] + ) + + s2 = StructType( + [ + StructField( + "name", + StructType( + [ + StructField("firstname", StringType(), True), + StructField("middlename", BooleanType(), True), + StructField("lastname", StringType(), True), + ] + ), + ), + ] + ) + + generated_diff = difflib.ndiff(str(s1).splitlines(), str(s2).splitlines()) + + expected_error_msg = "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertSchemaEqual(s1, s2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": expected_error_msg}, + ) + + def test_schema_unsupported_type(self): + s1 = "names: int" + s2 = "names: int" + + with self.assertRaises(PySparkAssertionError) as pe: + assertSchemaEqual(s1, s2) + + self.check_error( + exception=pe.exception, + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(s1)}, + ) + + def test_spark_sql(self): + assertDataFrameEqual(self.spark.sql("select 1 + 2 AS x"), self.spark.sql("select 3 AS x")) + assertDataFrameEqual( + self.spark.sql("select 1 + 2 AS x"), + self.spark.sql("select 3 AS x"), + checkRowOrder=True, + ) + + def test_spark_sql_sort_rows(self): + df1 = self.spark.createDataFrame( + data=[ + (1, 3000), + (2, 1000), + ], + schema=["id", "amount"], + ) + + df2 = self.spark.createDataFrame( + data=[ + (2, 1000), + (1, 3000), + ], + schema=["id", "amount"], + ) + + df1.createOrReplaceTempView("df1") + df2.createOrReplaceTempView("df2") + + assertDataFrameEqual( + self.spark.sql("select * from df1 order by amount"), self.spark.sql("select * from df2") + ) + + assertDataFrameEqual( + self.spark.sql("select * from df1 order by amount"), + self.spark.sql("select * from df2"), + checkRowOrder=True, + ) + + def test_empty_dataset(self): + df1 = self.spark.range(0, 10).limit(0) + + df2 = self.spark.range(0, 10).limit(0) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_no_column(self): + df1 = self.spark.range(0, 10).drop("id") + + df2 = self.spark.range(0, 10).drop("id") + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_empty_no_column(self): + df1 = self.spark.range(0, 10).drop("id").limit(0) + + df2 = self.spark.range(0, 10).drop("id").limit(0) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_empty_expected_list(self): + df1 = self.spark.range(0, 5).drop("id") + + df2 = [Row(), Row(), Row(), Row(), Row()] + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_no_column_expected_list(self): + df1 = self.spark.range(0, 10).limit(0) + + df2 = [] + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_empty_no_column_expected_list(self): + df1 = self.spark.range(0, 10).drop("id").limit(0) + + df2 = [] + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_special_vals(self): + df1 = self.spark.createDataFrame( + data=[ + (1, float("nan")), + (2, float("inf")), + (2, float("-inf")), + ], + schema=["id", "amount"], + ) + + df2 = self.spark.createDataFrame( + data=[ + (1, float("nan")), + (2, float("inf")), + (2, float("-inf")), + ], + schema=["id", "amount"], + ) + + assertDataFrameEqual(df1, df2, checkRowOrder=False) + assertDataFrameEqual(df1, df2, checkRowOrder=True) + + def test_df_list_row_equal(self): + df1 = self.spark.createDataFrame( + data=[ + (1, 3000), + (2, 1000), + ], + schema=["id", "amount"], + ) + + list_of_rows = [Row(1, 3000), Row(2, 1000)] + + assertDataFrameEqual(df1, list_of_rows, checkRowOrder=False) + assertDataFrameEqual(df1, list_of_rows, checkRowOrder=True) + + def test_list_rows_equal(self): + list_of_rows1 = [Row(1, "abc", 5000), Row(2, "def", 1000)] + list_of_rows2 = [Row(1, "abc", 5000), Row(2, "def", 1000)] + + assertDataFrameEqual(list_of_rows1, list_of_rows2, checkRowOrder=False) + assertDataFrameEqual(list_of_rows1, list_of_rows2, checkRowOrder=True) + + def test_list_rows_unequal(self): + list_of_rows1 = [Row(1, "abc", 5000), Row(2, "def", 1000)] + list_of_rows2 = [Row(1, "abc", 5000), Row(2, "defg", 1000)] + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(list_of_rows1, list_of_rows2)): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=2 + ) + + error_msg = "Results do not match: " + percent_diff = (1 / 2) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(list_of_rows1, list_of_rows2) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(list_of_rows1, list_of_rows2, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_list_row_unequal_schema(self): + df1 = self.spark.createDataFrame( + data=[ + (1, 3000), + (2, 1000), + (3, 10), + ], + schema=["id", "amount"], + ) + + list_of_rows = [Row(id=1, amount=300), Row(id=2, amount=100), Row(id=3, amount=10)] + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1, list_of_rows)): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=3 + ) + + error_msg = "Results do not match: " + percent_diff = (2 / 3) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, list_of_rows) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, list_of_rows, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + def test_list_row_unequal_schema(self): + from pyspark.sql import Row + + df1 = self.spark.createDataFrame( + data=[ + (1, 3000), + (2, 1000), + ], + schema=["id", "amount"], + ) + + list_of_rows = [Row(1, "3000"), Row(2, "1000")] + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in list(zip_longest(df1.collect(), list_of_rows)): + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=2 + ) + + error_msg = "Results do not match: " + percent_diff = (2 / 2) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, list_of_rows) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + with self.assertRaises(PySparkAssertionError) as pe: + assertDataFrameEqual(df1, list_of_rows, checkRowOrder=True) + + self.check_error( + exception=pe.exception, + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + +class UtilsTests(ReusedSQLTestCase, UtilsTestsMixin): + def test_capture_analysis_exception(self): + self.assertRaises(AnalysisException, lambda: self.spark.sql("select abc")) + self.assertRaises(AnalysisException, lambda: self.df.selectExpr("a + b")) + + def test_capture_user_friendly_exception(self): + try: + self.spark.sql("select `中文字段`") + except AnalysisException as e: + self.assertRegex(str(e), ".*UNRESOLVED_COLUMN.*`中文字段`.*") + + def test_spark_upgrade_exception(self): + # SPARK-32161 : Test case to Handle SparkUpgradeException in pythonic way + df = self.spark.createDataFrame([("2014-31-12",)], ["date_str"]) + df2 = df.select( + "date_str", to_date(from_unixtime(unix_timestamp("date_str", "yyyy-dd-aa"))) + ) + self.assertRaises(SparkUpgradeException, df2.collect) + + def test_capture_parse_exception(self): + self.assertRaises(ParseException, lambda: self.spark.sql("abc")) + + def test_capture_illegalargument_exception(self): + self.assertRaisesRegex( + IllegalArgumentException, + "Setting negative mapred.reduce.tasks", + lambda: self.spark.sql("SET mapred.reduce.tasks=-1"), + ) + df = self.spark.createDataFrame([(1, 2)], ["a", "b"]) + self.assertRaisesRegex( + IllegalArgumentException, + "1024 is not in the permitted values", + lambda: df.select(sha2(df.a, 1024)).collect(), + ) + try: + df.select(sha2(df.a, 1024)).collect() + except IllegalArgumentException as e: + self.assertRegex(e.desc, "1024 is not in the permitted values") + self.assertRegex(e.stackTrace, "org.apache.spark.sql.functions") + + def test_get_error_class_state(self): + # SPARK-36953: test CapturedException.getErrorClass and getSqlState (from SparkThrowable) + try: + self.spark.sql("""SELECT a""") + except AnalysisException as e: + self.assertEquals(e.getErrorClass(), "UNRESOLVED_COLUMN.WITHOUT_SUGGESTION") + self.assertEquals(e.getSqlState(), "42703") + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.test_utils import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_column.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_column.yml new file mode 100644 index 0000000..26eeb61 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_column.yml @@ -0,0 +1,37 @@ +# +# 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. +# + +- case: colDateTimeCompare + main: | + import datetime + from pyspark.sql.functions import col + + today = datetime.date.today() + now = datetime.datetime.now() + a_col = col("") + + a_col < today + a_col <= today + a_col == today + a_col >= today + a_col > today + + a_col < now + a_col <= now + a_col == now + a_col >= now + a_col > now diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_dataframe.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_dataframe.yml new file mode 100644 index 0000000..d32a09c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_dataframe.yml @@ -0,0 +1,140 @@ +# +# 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. +# + +- case: sampling + main: | + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + df = spark.range(1) + df.sample(1.0) + df.sample(0.5, 3) + df.sample(fraction=0.5, seed=3) + df.sample(withReplacement=True, fraction=0.5, seed=3) + df.sample(fraction=1.0) + df.sample(False, fraction=1.0) + + # Will raise a runtime error, though not a typecheck error, as bool is + # duck-type compatible with float. + df.sample(True) + + df.sample(withReplacement=False) + + out: | + main:16: error: No overload variant of "sample" of "DataFrame" matches argument type "bool" [call-overload] + main:16: note: Possible overload variants: + main:16: note: def sample(self, fraction: float, seed: Optional[int] = ...) -> DataFrame + main:16: note: def sample(self, withReplacement: Optional[bool], fraction: float, seed: Optional[int] = ...) -> DataFrame + + +- case: selectColumns + main: | + from pyspark.sql import SparkSession + from pyspark.sql.functions import col + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + df = spark.createDataFrame(data, schema="name str, age int") + + df.select(["name", "age"]) + df.select([col("name"), col("age")]) + + df.select(["name", col("age")]) # E: Argument 1 to "select" of "DataFrame" has incompatible type "List[object]"; expected "Union[List[Column], List[str]]" [arg-type] + + +- case: groupBy + main: | + from pyspark.sql import SparkSession + from pyspark.sql.functions import col + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + df = spark.createDataFrame(data, schema="name str, age int") + + df.groupBy(["name", "age"]) + df.groupby(["name", "age"]) + df.groupBy([col("name"), col("age")]) + df.groupby([col("name"), col("age")]) + df.groupBy(["name", col("age")]) # E: Argument 1 to "groupBy" of "DataFrame" has incompatible type "List[object]"; expected "Union[List[Column], List[str]]" [arg-type] + + +- case: rollup + main: | + from pyspark.sql import SparkSession + from pyspark.sql.functions import col + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + df = spark.createDataFrame(data, schema="name str, age int") + + df.rollup(["name", "age"]) + df.rollup([col("name"), col("age")]) + + + df.rollup(["name", col("age")]) # E: Argument 1 to "rollup" of "DataFrame" has incompatible type "List[object]"; expected "Union[List[Column], List[str]]" [arg-type] + + +- case: cube + main: | + from pyspark.sql import SparkSession + from pyspark.sql.functions import col + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + df = spark.createDataFrame(data, schema="name str, age int") + + df.cube(["name", "age"]) + df.cube([col("name"), col("age")]) + + + df.cube(["name", col("age")]) # E: Argument 1 to "cube" of "DataFrame" has incompatible type "List[object]"; expected "Union[List[Column], List[str]]" [arg-type] + + +- case: dropColumns + main: | + from pyspark.sql import SparkSession + from pyspark.sql.functions import lit, col + + spark = SparkSession.builder.getOrCreate() + df = spark.range(1) + df.drop("id") + df.drop("id", "foo") + df.drop(df.id) + + df.drop(col("id"), col("foo")) + + out: | + main:10: error: No overload variant of "drop" of "DataFrame" matches argument types "Column", "Column" [call-overload] + main:10: note: Possible overload variants: + main:10: note: def drop(self, cols: Union[Column, str]) -> DataFrame + main:10: note: def drop(self, *cols: str) -> DataFrame + + +- case: fillNullValues + main: | + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + df = spark.createDataFrame([(1,2)], schema=("id1", "id2")) + + df.fillna(value=1, subset="id1") + df.fillna(value=1, subset=("id1", "id2")) + df.fillna(value=1, subset=["id1"]) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_functions.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_functions.yml new file mode 100644 index 0000000..6c80420 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_functions.yml @@ -0,0 +1,101 @@ +# +# 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. +# + +- case: varargFunctionsOverloads + main: | + from pyspark.sql.functions import ( + array, + col, + create_map, + map_concat, + struct, + ) + + array(col("foo"), col("bar")) + array([col("foo"), col("bar")]) + array("foo", "bar") + array(["foo", "bar"]) + + create_map(col("foo"), col("bar")) + create_map([col("foo"), col("bar")]) + create_map("foo", "bar") + create_map(["foo", "bar"]) + + map_concat(col("foo"), col("bar")) + map_concat([col("foo"), col("bar")]) + map_concat("foo", "bar") + map_concat(["foo", "bar"]) + + struct(col("foo"), col("bar")) + struct([col("foo"), col("bar")]) + struct("foo", "bar") + struct(["foo", "bar"]) + + array([col("foo")], [col("bar")]) + create_map([col("foo")], [col("bar")]) + map_concat([col("foo")], [col("bar")]) + struct(["foo"], ["bar"]) + array(["foo"], ["bar"]) + create_map(["foo"], ["bar"]) + map_concat(["foo"], ["bar"]) + struct(["foo"], ["bar"]) + + col_names = ["foo", "bar"] + array(col_names) + struct(col_names) + create_map(col_names) + map_concat (col_names) + + col_objs = [col(x) for x in col_names] + array(col_objs) + struct(col_objs) + create_map(col_objs) + map_concat(col_objs) + + out: | + main:29: error: No overload variant of "array" matches argument types "List[Column]", "List[Column]" [call-overload] + main:29: note: Possible overload variants: + main:29: note: def array(*cols: Union[Column, str]) -> Column + main:29: note: def [ColumnOrName_] array(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:30: error: No overload variant of "create_map" matches argument types "List[Column]", "List[Column]" [call-overload] + main:30: note: Possible overload variants: + main:30: note: def create_map(*cols: Union[Column, str]) -> Column + main:30: note: def [ColumnOrName_] create_map(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:31: error: No overload variant of "map_concat" matches argument types "List[Column]", "List[Column]" [call-overload] + main:31: note: Possible overload variants: + main:31: note: def map_concat(*cols: Union[Column, str]) -> Column + main:31: note: def [ColumnOrName_] map_concat(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:32: error: No overload variant of "struct" matches argument types "List[str]", "List[str]" [call-overload] + main:32: note: Possible overload variants: + main:32: note: def struct(*cols: Union[Column, str]) -> Column + main:32: note: def [ColumnOrName_] struct(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:33: error: No overload variant of "array" matches argument types "List[str]", "List[str]" [call-overload] + main:33: note: Possible overload variants: + main:33: note: def array(*cols: Union[Column, str]) -> Column + main:33: note: def [ColumnOrName_] array(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:34: error: No overload variant of "create_map" matches argument types "List[str]", "List[str]" [call-overload] + main:34: note: Possible overload variants: + main:34: note: def create_map(*cols: Union[Column, str]) -> Column + main:34: note: def [ColumnOrName_] create_map(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:35: error: No overload variant of "map_concat" matches argument types "List[str]", "List[str]" [call-overload] + main:35: note: Possible overload variants: + main:35: note: def map_concat(*cols: Union[Column, str]) -> Column + main:35: note: def [ColumnOrName_] map_concat(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column + main:36: error: No overload variant of "struct" matches argument types "List[str]", "List[str]" [call-overload] + main:36: note: Possible overload variants: + main:36: note: def struct(*cols: Union[Column, str]) -> Column + main:36: note: def [ColumnOrName_] struct(Union[List[ColumnOrName_], Tuple[ColumnOrName_, ...]]) -> Column diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_readwriter.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_readwriter.yml new file mode 100644 index 0000000..2ce3637 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_readwriter.yml @@ -0,0 +1,45 @@ +# +# 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. +# + +- case: readWriterOptions + main: | + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + + spark.read.option("foo", 1) + spark.createDataFrame([(1, 2)], ["foo", "bar"]).write.option("bar", True) + + spark.read.load(foo=True) + + spark.read.load(foo=["a"]) # E: Argument "foo" to "load" of "DataFrameReader" has incompatible type "List[str]"; expected "Union[bool, float, int, str, None]" [arg-type] + spark.read.option("foo", (1, )) # E: Argument 2 to "option" of "DataFrameReader" has incompatible type "Tuple[int]"; expected "Union[bool, float, int, str, None]" [arg-type] + + +- case: readStreamOptions + main: | + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + + spark.read.option("foo", True).option("foo", 1).option("foo", 1.0).option("foo", "1").option("foo", None) + spark.readStream.option("foo", True).option("foo", 1).option("foo", 1.0).option("foo", "1").option("foo", None) + + spark.read.options(foo=True, bar=1).options(foo=1.0, bar="1", baz=None) + spark.readStream.options(foo=True, bar=1).options(foo=1.0, bar="1", baz=None) + + spark.readStream.load(foo=True) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_session.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_session.yml new file mode 100644 index 0000000..5c9fd9f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_session.yml @@ -0,0 +1,118 @@ +# +# 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. +# + +- case: createDataFrameStructsValid + main: | + from pyspark.sql import SparkSession + from pyspark.sql.types import StructType, StructField, StringType, IntegerType + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + schema = StructType([ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True) + ]) + + # Valid structs + spark.createDataFrame(data) + spark.createDataFrame(data, samplingRatio=0.1) + spark.createDataFrame(data, ("name", "age")) + spark.createDataFrame(data, schema) + spark.createDataFrame(data, "name string, age integer") + spark.createDataFrame([(1, ("foo", "bar"))], ("_1", "_2")) + spark.createDataFrame(data, ("name", "age"), samplingRatio=0.1) + + +- case: createDataFrameScalarsValid + main: | + + from pyspark.sql import SparkSession + from pyspark.sql.types import StructType, StructField, StringType, IntegerType + + spark = SparkSession.builder.getOrCreate() + + # Scalars + spark.createDataFrame([1, 2, 3], IntegerType()) + spark.createDataFrame(["foo", "bar"], "string") + + +- case: createDataFrameScalarsInvalid + main: | + from pyspark.sql import SparkSession + from pyspark.sql.types import StructType, StructField, StringType, IntegerType + + spark = SparkSession.builder.getOrCreate() + + schema = StructType([ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True) + ]) + + # Invalid - scalars require schema + spark.createDataFrame(["foo", "bar"]) # E: Value of type variable "RowLike" of "createDataFrame" of "SparkSession" cannot be "str" [type-var] + + # Invalid - data has to match schema (either product -> struct or scalar -> atomic) + spark.createDataFrame([1, 2, 3], schema) # E: Value of type variable "RowLike" of "createDataFrame" of "SparkSession" cannot be "int" [type-var] + + +- case: createDataFrameStructsInvalid + main: | + from pyspark.sql import SparkSession + from pyspark.sql.types import StructType, StructField, StringType, IntegerType + + spark = SparkSession.builder.getOrCreate() + + data = [('Alice', 1)] + + schema = StructType([ + StructField("name", StringType(), True), + StructField("age", IntegerType(), True) + ]) + + # Invalid product should have StructType schema + spark.createDataFrame(data, IntegerType()) + + # This shouldn't type check, though is technically speaking valid + # because samplingRatio is ignored + spark.createDataFrame(data, schema, samplingRatio=0.1) + + out: | + main:14: error: Value of type variable "AtomicValue" of "createDataFrame" of "SparkSession" cannot be "Tuple[str, int]" [type-var] + main:18: error: No overload variant of "createDataFrame" of "SparkSession" matches argument types "List[Tuple[str, int]]", "StructType", "float" [call-overload] + main:18: note: Possible overload variants: + main:18: note: def [RowLike in (List[Any], Tuple[Any, ...], Row)] createDataFrame(self, data: Iterable[RowLike], schema: Union[List[str], Tuple[str, ...]] = ..., samplingRatio: Optional[float] = ...) -> DataFrame + main:18: note: def [RowLike in (List[Any], Tuple[Any, ...], Row)] createDataFrame(self, data: RDD[RowLike], schema: Union[List[str], Tuple[str, ...]] = ..., samplingRatio: Optional[float] = ...) -> DataFrame + main:18: note: def [RowLike in (List[Any], Tuple[Any, ...], Row)] createDataFrame(self, data: Iterable[RowLike], schema: Union[StructType, str], *, verifySchema: bool = ...) -> DataFrame + main:18: note: def [RowLike in (List[Any], Tuple[Any, ...], Row)] createDataFrame(self, data: RDD[RowLike], schema: Union[StructType, str], *, verifySchema: bool = ...) -> DataFrame + main:18: note: def [AtomicValue in (datetime, date, Decimal, bool, str, int, float)] createDataFrame(self, data: RDD[AtomicValue], schema: Union[AtomicType, str], verifySchema: bool = ...) -> DataFrame + main:18: note: def [AtomicValue in (datetime, date, Decimal, bool, str, int, float)] createDataFrame(self, data: Iterable[AtomicValue], schema: Union[AtomicType, str], verifySchema: bool = ...) -> DataFrame + main:18: note: def createDataFrame(self, data: DataFrame, samplingRatio: Optional[float] = ...) -> DataFrame + main:18: note: def createDataFrame(self, data: DataFrame, schema: Union[StructType, str], verifySchema: bool = ...) -> DataFrame + + +- case: createDataFrameFromEmptyRdd + main: | + from pyspark.sql import SparkSession + from pyspark.sql.types import StructType + + spark = SparkSession.builder.getOrCreate() + + spark.createDataFrame( + spark.sparkContext.emptyRDD(), + schema=StructType(), + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_udf.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_udf.yml new file mode 100644 index 0000000..ef74b31 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/tests/typing/test_udf.yml @@ -0,0 +1,170 @@ +# +# 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. +# + +- case: scalarUDF + main: | + from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType + import pandas.core.series + import pandas.core.frame + + @pandas_udf("string", PandasUDFType.SCALAR) + def f(x: pandas.core.series.Series) -> pandas.core.series.Series: + return x + + @pandas_udf("string", PandasUDFType.SCALAR) + def g(x: pandas.core.series.Series, y: pandas.core.series.Series) -> pandas.core.series.Series: + return x + + @pandas_udf("string", PandasUDFType.SCALAR) + def h(*xs: pandas.core.series.Series) -> pandas.core.series.Series: + return xs[0] + + @pandas_udf("x string, y string", PandasUDFType.SCALAR) + def k(x: pandas.core.frame.DataFrame, y: pandas.core.series.Series) -> pandas.core.frame.DataFrame: + return x + + pandas_udf(lambda x: x, "string", PandasUDFType.SCALAR) + pandas_udf(lambda x, y: x, "string", PandasUDFType.SCALAR) + pandas_udf(lambda *xs: xs[0], "string", PandasUDFType.SCALAR) + + +- case: scalarIterUDF + main: | + from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType + from pyspark.sql.types import IntegerType + import pandas.core.series + from typing import Iterable + + @pandas_udf(IntegerType(), PandasUDFType.SCALAR_ITER) + def f(xs: Iterable[pandas.core.series.Series]) -> Iterable[pandas.core.series.Series]: + for x in xs: + yield x + 1 + + +- case: groupedMapUdf + main: | + from typing import Any + + from pyspark.sql.session import SparkSession + from pyspark.sql.types import StructField, StructType, LongType + from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType + import pandas.core.frame + + @pandas_udf("id long", PandasUDFType.GROUPED_MAP) + def f(pdf: pandas.core.frame.DataFrame) -> pandas.core.frame.DataFrame: + return pdf + + spark = SparkSession.builder.getOrCreate() + + dfg = spark.range(1).groupBy("id") + dfg.apply(f) + + @pandas_udf("id long", PandasUDFType.GROUPED_MAP) + def g(key: Any, pdf: pandas.core.frame.DataFrame) -> pandas.core.frame.DataFrame: + return pdf + + dfg.apply(g) + + + def h(pdf: pandas.core.frame.DataFrame) -> pandas.core.frame.DataFrame: + return pdf + + dfg.applyInPandas(h, "id long") + dfg.applyInPandas(h, StructType([StructField("id", LongType())])) + + +- case: groupedAggUDF + main: | + # Let's keep this one to make sure compatibility imports work + from pyspark.sql.functions import pandas_udf, PandasUDFType + from pyspark.sql.types import IntegerType + import pandas.core.series + + @pandas_udf(IntegerType(), PandasUDFType.GROUPED_AGG) + def f(x: pandas.core.series.Series) -> int: + return 42 + + @pandas_udf("int", PandasUDFType.GROUPED_AGG) + def g(x: pandas.core.series.Series, y: pandas.core.series.Series) -> int: + return 42 + + @pandas_udf("int", PandasUDFType.GROUPED_AGG) + def h(*xs: pandas.core.series.Series) -> int: + return 42 + + pandas_udf(lambda x: 42, "str", PandasUDFType.GROUPED_AGG) + pandas_udf(lambda x, y: 42, "str", PandasUDFType.GROUPED_AGG) + pandas_udf(lambda *xs: 42, "str", PandasUDFType.GROUPED_AGG) + + +- case: mapIterUdf + main: | + from pyspark.sql.session import SparkSession + from typing import Iterable + import pandas.core.frame + + spark = SparkSession.builder.getOrCreate() + + def f(batch_iter: Iterable[pandas.core.frame.DataFrame]) -> Iterable[pandas.core.frame.DataFrame]: + for pdf in batch_iter: + yield pdf[pdf.id == 1] + + spark.range(1).mapInPandas(f, "id long").show() + + +- case: legacyUDF + main: | + from pyspark.sql.functions import udf + from pyspark.sql.types import IntegerType + + udf(lambda x: x, "string") + + udf(lambda x: x) + + @udf("string") + def f(x: str) -> str: + return x + + @udf(returnType="string") + def g(x: str) -> str: + return x + + @udf(returnType=IntegerType()) + def h(x: int) -> int: + return x + + @udf + def i(x: str) -> str: + return x + + +- case: cogroupedAggUdf + main: | + from pyspark.sql.session import SparkSession + import pandas.core.frame + from pyspark.sql.types import StructType, StructField, LongType + + spark = SparkSession.builder.getOrCreate() + + dfg1 = spark.range(1).groupBy("id") + dfg2 = spark.range(1).groupBy("id") + + def f(x: pandas.core.frame.DataFrame, y: pandas.core.frame.DataFrame) -> pandas.core.frame.DataFrame: + return x + + dfg1.cogroup(dfg2).applyInPandas(f, "id int") + dfg1.cogroup(dfg2).applyInPandas(f, StructType([StructField("id", LongType())])) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/types.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/types.py new file mode 100644 index 0000000..24964c5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/types.py @@ -0,0 +1,2544 @@ +# +# 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 sys +import decimal +import time +import math +import datetime +import calendar +import json +import re +import base64 +from array import array +import ctypes +from collections.abc import Iterable +from functools import reduce +from typing import ( + cast, + overload, + Any, + Callable, + ClassVar, + Dict, + Iterator, + List, + Optional, + Union, + Tuple, + Type, + TypeVar, + TYPE_CHECKING, +) + +from py4j.protocol import register_input_converter +from py4j.java_gateway import GatewayClient, JavaClass, JavaGateway, JavaObject, JVMView + +from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.utils import has_numpy, get_active_spark_context +from pyspark.errors import PySparkNotImplementedError, PySparkTypeError, PySparkValueError + +if has_numpy: + import numpy as np + +T = TypeVar("T") +U = TypeVar("U") + +__all__ = [ + "DataType", + "NullType", + "CharType", + "StringType", + "VarcharType", + "BinaryType", + "BooleanType", + "DateType", + "TimestampType", + "TimestampNTZType", + "DecimalType", + "DoubleType", + "FloatType", + "ByteType", + "IntegerType", + "LongType", + "DayTimeIntervalType", + "YearMonthIntervalType", + "Row", + "ShortType", + "ArrayType", + "MapType", + "StructField", + "StructType", +] + + +if TYPE_CHECKING: + import numpy as np + + +class DataType: + """Base class for data types.""" + + def __repr__(self) -> str: + return self.__class__.__name__ + "()" + + def __hash__(self) -> int: + return hash(str(self)) + + def __eq__(self, other: Any) -> bool: + return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ + + def __ne__(self, other: Any) -> bool: + return not self.__eq__(other) + + @classmethod + def typeName(cls) -> str: + return cls.__name__[:-4].lower() + + def simpleString(self) -> str: + return self.typeName() + + def jsonValue(self) -> Union[str, Dict[str, Any]]: + return self.typeName() + + def json(self) -> str: + return json.dumps(self.jsonValue(), separators=(",", ":"), sort_keys=True) + + def needConversion(self) -> bool: + """ + Does this type needs conversion between Python object and internal SQL object. + + This is used to avoid the unnecessary conversion for ArrayType/MapType/StructType. + """ + return False + + def toInternal(self, obj: Any) -> Any: + """ + Converts a Python object into an internal SQL object. + """ + return obj + + def fromInternal(self, obj: Any) -> Any: + """ + Converts an internal SQL object into a native Python object. + """ + return obj + + +# This singleton pattern does not work with pickle, you will get +# another object after pickle and unpickle +class DataTypeSingleton(type): + """Metaclass for DataType""" + + _instances: ClassVar[Dict[Type["DataTypeSingleton"], "DataTypeSingleton"]] = {} + + def __call__(cls: Type[T]) -> T: + if cls not in cls._instances: # type: ignore[attr-defined] + cls._instances[cls] = super( # type: ignore[misc, attr-defined] + DataTypeSingleton, cls + ).__call__() + return cls._instances[cls] # type: ignore[attr-defined] + + +class NullType(DataType, metaclass=DataTypeSingleton): + """Null type. + + The data type representing None, used for the types that cannot be inferred. + """ + + @classmethod + def typeName(cls) -> str: + return "void" + + +class AtomicType(DataType): + """An internal type used to represent everything that is not + null, UDTs, arrays, structs, and maps.""" + + +class NumericType(AtomicType): + """Numeric data types.""" + + +class IntegralType(NumericType, metaclass=DataTypeSingleton): + """Integral data types.""" + + pass + + +class FractionalType(NumericType): + """Fractional data types.""" + + +class StringType(AtomicType, metaclass=DataTypeSingleton): + """String data type.""" + + pass + + +class CharType(AtomicType): + """Char data type + + Parameters + ---------- + length : int + the length limitation. + """ + + def __init__(self, length: int): + self.length = length + + def simpleString(self) -> str: + return "char(%d)" % (self.length) + + def jsonValue(self) -> str: + return "char(%d)" % (self.length) + + def __repr__(self) -> str: + return "CharType(%d)" % (self.length) + + +class VarcharType(AtomicType): + """Varchar data type + + Parameters + ---------- + length : int + the length limitation. + """ + + def __init__(self, length: int): + self.length = length + + def simpleString(self) -> str: + return "varchar(%d)" % (self.length) + + def jsonValue(self) -> str: + return "varchar(%d)" % (self.length) + + def __repr__(self) -> str: + return "VarcharType(%d)" % (self.length) + + +class BinaryType(AtomicType, metaclass=DataTypeSingleton): + """Binary (byte array) data type.""" + + pass + + +class BooleanType(AtomicType, metaclass=DataTypeSingleton): + """Boolean data type.""" + + pass + + +class DateType(AtomicType, metaclass=DataTypeSingleton): + """Date (datetime.date) data type.""" + + EPOCH_ORDINAL = datetime.datetime(1970, 1, 1).toordinal() + + def needConversion(self) -> bool: + return True + + def toInternal(self, d: datetime.date) -> int: + if d is not None: + return d.toordinal() - self.EPOCH_ORDINAL + + def fromInternal(self, v: int) -> datetime.date: + if v is not None: + return datetime.date.fromordinal(v + self.EPOCH_ORDINAL) + + +class TimestampType(AtomicType, metaclass=DataTypeSingleton): + """Timestamp (datetime.datetime) data type.""" + + def needConversion(self) -> bool: + return True + + def toInternal(self, dt: datetime.datetime) -> int: + if dt is not None: + seconds = ( + calendar.timegm(dt.utctimetuple()) if dt.tzinfo else time.mktime(dt.timetuple()) + ) + return int(seconds) * 1000000 + dt.microsecond + + def fromInternal(self, ts: int) -> datetime.datetime: + if ts is not None: + # using int to avoid precision loss in float + return datetime.datetime.fromtimestamp(ts // 1000000).replace(microsecond=ts % 1000000) + + +class TimestampNTZType(AtomicType, metaclass=DataTypeSingleton): + """Timestamp (datetime.datetime) data type without timezone information.""" + + def needConversion(self) -> bool: + return True + + @classmethod + def typeName(cls) -> str: + return "timestamp_ntz" + + def toInternal(self, dt: datetime.datetime) -> int: + if dt is not None: + seconds = calendar.timegm(dt.timetuple()) + return int(seconds) * 1000000 + dt.microsecond + + def fromInternal(self, ts: int) -> datetime.datetime: + if ts is not None: + # using int to avoid precision loss in float + return datetime.datetime.utcfromtimestamp(ts // 1000000).replace( + microsecond=ts % 1000000 + ) + + +class DecimalType(FractionalType): + """Decimal (decimal.Decimal) data type. + + The DecimalType must have fixed precision (the maximum total number of digits) + and scale (the number of digits on the right of dot). For example, (5, 2) can + support the value from [-999.99 to 999.99]. + + The precision can be up to 38, the scale must be less or equal to precision. + + When creating a DecimalType, the default precision and scale is (10, 0). When inferring + schema from decimal.Decimal objects, it will be DecimalType(38, 18). + + Parameters + ---------- + precision : int, optional + the maximum (i.e. total) number of digits (default: 10) + scale : int, optional + the number of digits on right side of dot. (default: 0) + """ + + def __init__(self, precision: int = 10, scale: int = 0): + self.precision = precision + self.scale = scale + self.hasPrecisionInfo = True # this is a public API + + def simpleString(self) -> str: + return "decimal(%d,%d)" % (self.precision, self.scale) + + def jsonValue(self) -> str: + return "decimal(%d,%d)" % (self.precision, self.scale) + + def __repr__(self) -> str: + return "DecimalType(%d,%d)" % (self.precision, self.scale) + + +class DoubleType(FractionalType, metaclass=DataTypeSingleton): + """Double data type, representing double precision floats.""" + + pass + + +class FloatType(FractionalType, metaclass=DataTypeSingleton): + """Float data type, representing single precision floats.""" + + pass + + +class ByteType(IntegralType): + """Byte data type, i.e. a signed integer in a single byte.""" + + def simpleString(self) -> str: + return "tinyint" + + +class IntegerType(IntegralType): + """Int data type, i.e. a signed 32-bit integer.""" + + def simpleString(self) -> str: + return "int" + + +class LongType(IntegralType): + """Long data type, i.e. a signed 64-bit integer. + + If the values are beyond the range of [-9223372036854775808, 9223372036854775807], + please use :class:`DecimalType`. + """ + + def simpleString(self) -> str: + return "bigint" + + +class ShortType(IntegralType): + """Short data type, i.e. a signed 16-bit integer.""" + + def simpleString(self) -> str: + return "smallint" + + +class AnsiIntervalType(AtomicType): + """The interval type which conforms to the ANSI SQL standard.""" + + pass + + +class DayTimeIntervalType(AnsiIntervalType): + """DayTimeIntervalType (datetime.timedelta).""" + + DAY = 0 + HOUR = 1 + MINUTE = 2 + SECOND = 3 + + _fields = { + DAY: "day", + HOUR: "hour", + MINUTE: "minute", + SECOND: "second", + } + + _inverted_fields = dict(zip(_fields.values(), _fields.keys())) + + def __init__(self, startField: Optional[int] = None, endField: Optional[int] = None): + if startField is None and endField is None: + # Default matched to scala side. + startField = DayTimeIntervalType.DAY + endField = DayTimeIntervalType.SECOND + elif startField is not None and endField is None: + endField = startField + + fields = DayTimeIntervalType._fields + if startField not in fields.keys() or endField not in fields.keys(): + raise RuntimeError("interval %s to %s is invalid" % (startField, endField)) + self.startField = cast(int, startField) + self.endField = cast(int, endField) + + def _str_repr(self) -> str: + fields = DayTimeIntervalType._fields + start_field_name = fields[self.startField] + end_field_name = fields[self.endField] + if start_field_name == end_field_name: + return "interval %s" % start_field_name + else: + return "interval %s to %s" % (start_field_name, end_field_name) + + simpleString = _str_repr + + jsonValue = _str_repr + + def __repr__(self) -> str: + return "%s(%d, %d)" % (type(self).__name__, self.startField, self.endField) + + def needConversion(self) -> bool: + return True + + def toInternal(self, dt: datetime.timedelta) -> Optional[int]: + if dt is not None: + return (((dt.days * 86400) + dt.seconds) * 1_000_000) + dt.microseconds + + def fromInternal(self, micros: int) -> Optional[datetime.timedelta]: + if micros is not None: + return datetime.timedelta(microseconds=micros) + + +class YearMonthIntervalType(AnsiIntervalType): + """YearMonthIntervalType, represents year-month intervals of the SQL standard""" + + YEAR = 0 + MONTH = 1 + + _fields = { + YEAR: "year", + MONTH: "month", + } + + _inverted_fields = dict(zip(_fields.values(), _fields.keys())) + + def __init__(self, startField: Optional[int] = None, endField: Optional[int] = None): + if startField is None and endField is None: + # Default matched to scala side. + startField = YearMonthIntervalType.YEAR + endField = YearMonthIntervalType.MONTH + elif startField is not None and endField is None: + endField = startField + + fields = YearMonthIntervalType._fields + if startField not in fields.keys() or endField not in fields.keys(): + raise RuntimeError("interval %s to %s is invalid" % (startField, endField)) + self.startField = cast(int, startField) + self.endField = cast(int, endField) + + def _str_repr(self) -> str: + fields = YearMonthIntervalType._fields + start_field_name = fields[self.startField] + end_field_name = fields[self.endField] + if start_field_name == end_field_name: + return "interval %s" % start_field_name + else: + return "interval %s to %s" % (start_field_name, end_field_name) + + simpleString = _str_repr + + jsonValue = _str_repr + + def __repr__(self) -> str: + return "%s(%d, %d)" % (type(self).__name__, self.startField, self.endField) + + +class ArrayType(DataType): + """Array data type. + + Parameters + ---------- + elementType : :class:`DataType` + :class:`DataType` of each element in the array. + containsNull : bool, optional + whether the array can contain null (None) values. + + Examples + -------- + >>> from pyspark.sql.types import ArrayType, StringType, StructField, StructType + + The below example demonstrates how to create class:`ArrayType`: + + >>> arr = ArrayType(StringType()) + + The array can contain null (None) values by default: + + >>> ArrayType(StringType()) == ArrayType(StringType(), True) + True + >>> ArrayType(StringType(), False) == ArrayType(StringType()) + False + """ + + def __init__(self, elementType: DataType, containsNull: bool = True): + assert isinstance(elementType, DataType), "elementType %s should be an instance of %s" % ( + elementType, + DataType, + ) + self.elementType = elementType + self.containsNull = containsNull + + def simpleString(self) -> str: + return "array<%s>" % self.elementType.simpleString() + + def __repr__(self) -> str: + return "ArrayType(%s, %s)" % (self.elementType, str(self.containsNull)) + + def jsonValue(self) -> Dict[str, Any]: + return { + "type": self.typeName(), + "elementType": self.elementType.jsonValue(), + "containsNull": self.containsNull, + } + + @classmethod + def fromJson(cls, json: Dict[str, Any]) -> "ArrayType": + return ArrayType(_parse_datatype_json_value(json["elementType"]), json["containsNull"]) + + def needConversion(self) -> bool: + return self.elementType.needConversion() + + def toInternal(self, obj: List[Optional[T]]) -> List[Optional[T]]: + if not self.needConversion(): + return obj + return obj and [self.elementType.toInternal(v) for v in obj] + + def fromInternal(self, obj: List[Optional[T]]) -> List[Optional[T]]: + if not self.needConversion(): + return obj + return obj and [self.elementType.fromInternal(v) for v in obj] + + +class MapType(DataType): + """Map data type. + + Parameters + ---------- + keyType : :class:`DataType` + :class:`DataType` of the keys in the map. + valueType : :class:`DataType` + :class:`DataType` of the values in the map. + valueContainsNull : bool, optional + indicates whether values can contain null (None) values. + + Notes + ----- + Keys in a map data type are not allowed to be null (None). + + Examples + -------- + >>> from pyspark.sql.types import IntegerType, FloatType, MapType, StringType + + The below example demonstrates how to create class:`MapType`: + + >>> map_type = MapType(StringType(), IntegerType()) + + The values of the map can contain null (``None``) values by default: + + >>> (MapType(StringType(), IntegerType()) + ... == MapType(StringType(), IntegerType(), True)) + True + >>> (MapType(StringType(), IntegerType(), False) + ... == MapType(StringType(), FloatType())) + False + """ + + def __init__(self, keyType: DataType, valueType: DataType, valueContainsNull: bool = True): + assert isinstance(keyType, DataType), "keyType %s should be an instance of %s" % ( + keyType, + DataType, + ) + assert isinstance(valueType, DataType), "valueType %s should be an instance of %s" % ( + valueType, + DataType, + ) + self.keyType = keyType + self.valueType = valueType + self.valueContainsNull = valueContainsNull + + def simpleString(self) -> str: + return "map<%s,%s>" % (self.keyType.simpleString(), self.valueType.simpleString()) + + def __repr__(self) -> str: + return "MapType(%s, %s, %s)" % (self.keyType, self.valueType, str(self.valueContainsNull)) + + def jsonValue(self) -> Dict[str, Any]: + return { + "type": self.typeName(), + "keyType": self.keyType.jsonValue(), + "valueType": self.valueType.jsonValue(), + "valueContainsNull": self.valueContainsNull, + } + + @classmethod + def fromJson(cls, json: Dict[str, Any]) -> "MapType": + return MapType( + _parse_datatype_json_value(json["keyType"]), + _parse_datatype_json_value(json["valueType"]), + json["valueContainsNull"], + ) + + def needConversion(self) -> bool: + return self.keyType.needConversion() or self.valueType.needConversion() + + def toInternal(self, obj: Dict[T, Optional[U]]) -> Dict[T, Optional[U]]: + if not self.needConversion(): + return obj + return obj and dict( + (self.keyType.toInternal(k), self.valueType.toInternal(v)) for k, v in obj.items() + ) + + def fromInternal(self, obj: Dict[T, Optional[U]]) -> Dict[T, Optional[U]]: + if not self.needConversion(): + return obj + return obj and dict( + (self.keyType.fromInternal(k), self.valueType.fromInternal(v)) for k, v in obj.items() + ) + + +class StructField(DataType): + """A field in :class:`StructType`. + + Parameters + ---------- + name : str + name of the field. + dataType : :class:`DataType` + :class:`DataType` of the field. + nullable : bool, optional + whether the field can be null (None) or not. + metadata : dict, optional + a dict from string to simple type that can be toInternald to JSON automatically + + Examples + -------- + >>> from pyspark.sql.types import StringType, StructField + >>> (StructField("f1", StringType(), True) + ... == StructField("f1", StringType(), True)) + True + >>> (StructField("f1", StringType(), True) + ... == StructField("f2", StringType(), True)) + False + """ + + def __init__( + self, + name: str, + dataType: DataType, + nullable: bool = True, + metadata: Optional[Dict[str, Any]] = None, + ): + assert isinstance(dataType, DataType), "dataType %s should be an instance of %s" % ( + dataType, + DataType, + ) + assert isinstance(name, str), "field name %s should be a string" % (name) + self.name = name + self.dataType = dataType + self.nullable = nullable + self.metadata = metadata or {} + + def simpleString(self) -> str: + return "%s:%s" % (self.name, self.dataType.simpleString()) + + def __repr__(self) -> str: + return "StructField('%s', %s, %s)" % (self.name, self.dataType, str(self.nullable)) + + def jsonValue(self) -> Dict[str, Any]: + return { + "name": self.name, + "type": self.dataType.jsonValue(), + "nullable": self.nullable, + "metadata": self.metadata, + } + + @classmethod + def fromJson(cls, json: Dict[str, Any]) -> "StructField": + return StructField( + json["name"], + _parse_datatype_json_value(json["type"]), + json["nullable"], + json["metadata"], + ) + + def needConversion(self) -> bool: + return self.dataType.needConversion() + + def toInternal(self, obj: T) -> T: + return self.dataType.toInternal(obj) + + def fromInternal(self, obj: T) -> T: + return self.dataType.fromInternal(obj) + + def typeName(self) -> str: # type: ignore[override] + raise PySparkTypeError( + error_class="INVALID_TYPENAME_CALL", + message_parameters={}, + ) + + +class StructType(DataType): + """Struct type, consisting of a list of :class:`StructField`. + + This is the data type representing a :class:`Row`. + + Iterating a :class:`StructType` will iterate over its :class:`StructField`\\s. + A contained :class:`StructField` can be accessed by its name or position. + + Examples + -------- + >>> from pyspark.sql.types import * + >>> struct1 = StructType([StructField("f1", StringType(), True)]) + >>> struct1["f1"] + StructField('f1', StringType(), True) + >>> struct1[0] + StructField('f1', StringType(), True) + + >>> struct1 = StructType([StructField("f1", StringType(), True)]) + >>> struct2 = StructType([StructField("f1", StringType(), True)]) + >>> struct1 == struct2 + True + >>> struct1 = StructType([StructField("f1", CharType(10), True)]) + >>> struct2 = StructType([StructField("f1", CharType(10), True)]) + >>> struct1 == struct2 + True + >>> struct1 = StructType([StructField("f1", VarcharType(10), True)]) + >>> struct2 = StructType([StructField("f1", VarcharType(10), True)]) + >>> struct1 == struct2 + True + >>> struct1 = StructType([StructField("f1", StringType(), True)]) + >>> struct2 = StructType([StructField("f1", StringType(), True), + ... StructField("f2", IntegerType(), False)]) + >>> struct1 == struct2 + False + + The below example demonstrates how to create a DataFrame based on a struct created + using class:`StructType` and class:`StructField`: + + >>> data = [("Alice", ["Java", "Scala"]), ("Bob", ["Python", "Scala"])] + >>> schema = StructType([ + ... StructField("name", StringType()), + ... StructField("languagesSkills", ArrayType(StringType())), + ... ]) + >>> df = spark.createDataFrame(data=data, schema=schema) + >>> df.printSchema() + root + |-- name: string (nullable = true) + |-- languagesSkills: array (nullable = true) + | |-- element: string (containsNull = true) + >>> df.show() + +-----+---------------+ + | name|languagesSkills| + +-----+---------------+ + |Alice| [Java, Scala]| + | Bob|[Python, Scala]| + +-----+---------------+ + """ + + def __init__(self, fields: Optional[List[StructField]] = None): + if not fields: + self.fields = [] + self.names = [] + else: + self.fields = fields + self.names = [f.name for f in fields] + assert all( + isinstance(f, StructField) for f in fields + ), "fields should be a list of StructField" + # Precalculated list of fields that need conversion with fromInternal/toInternal functions + self._needConversion = [f.needConversion() for f in self] + self._needSerializeAnyField = any(self._needConversion) + + @overload + def add( + self, + field: str, + data_type: Union[str, DataType], + nullable: bool = True, + metadata: Optional[Dict[str, Any]] = None, + ) -> "StructType": + ... + + @overload + def add(self, field: StructField) -> "StructType": + ... + + def add( + self, + field: Union[str, StructField], + data_type: Optional[Union[str, DataType]] = None, + nullable: bool = True, + metadata: Optional[Dict[str, Any]] = None, + ) -> "StructType": + """ + Construct a :class:`StructType` by adding new elements to it, to define the schema. + The method accepts either: + + a) A single parameter which is a :class:`StructField` object. + b) Between 2 and 4 parameters as (name, data_type, nullable (optional), + metadata(optional). The data_type parameter may be either a String or a + :class:`DataType` object. + + Parameters + ---------- + field : str or :class:`StructField` + Either the name of the field or a :class:`StructField` object + data_type : :class:`DataType`, optional + If present, the DataType of the :class:`StructField` to create + nullable : bool, optional + Whether the field to add should be nullable (default True) + metadata : dict, optional + Any additional metadata (default None) + + Returns + ------- + :class:`StructType` + + Examples + -------- + >>> from pyspark.sql.types import IntegerType, StringType, StructField, StructType + >>> struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) + >>> struct2 = StructType([StructField("f1", StringType(), True), + ... StructField("f2", StringType(), True, None)]) + >>> struct1 == struct2 + True + >>> struct1 = StructType().add(StructField("f1", StringType(), True)) + >>> struct2 = StructType([StructField("f1", StringType(), True)]) + >>> struct1 == struct2 + True + >>> struct1 = StructType().add("f1", "string", True) + >>> struct2 = StructType([StructField("f1", StringType(), True)]) + >>> struct1 == struct2 + True + """ + if isinstance(field, StructField): + self.fields.append(field) + self.names.append(field.name) + else: + if isinstance(field, str) and data_type is None: + raise PySparkValueError( + error_class="ARGUMENT_REQUIRED", + message_parameters={ + "arg_name": "data_type", + "condition": "passing name of struct_field to create", + }, + ) + + if isinstance(data_type, str): + data_type_f = _parse_datatype_json_value(data_type) + else: + data_type_f = data_type + self.fields.append(StructField(field, data_type_f, nullable, metadata)) + self.names.append(field) + # Precalculated list of fields that need conversion with fromInternal/toInternal functions + self._needConversion = [f.needConversion() for f in self] + self._needSerializeAnyField = any(self._needConversion) + return self + + def __iter__(self) -> Iterator[StructField]: + """Iterate the fields""" + return iter(self.fields) + + def __len__(self) -> int: + """Return the number of fields.""" + return len(self.fields) + + def __getitem__(self, key: Union[str, int]) -> StructField: + """Access fields by name or slice.""" + if isinstance(key, str): + for field in self: + if field.name == key: + return field + raise KeyError("No StructField named {0}".format(key)) + elif isinstance(key, int): + try: + return self.fields[key] + except IndexError: + raise IndexError("StructType index out of range") + elif isinstance(key, slice): + return StructType(self.fields[key]) + else: + raise PySparkTypeError( + error_class="NOT_INT_OR_SLICE_OR_STR", + message_parameters={"arg_name": "key", "arg_type": type(key).__name__}, + ) + + def simpleString(self) -> str: + return "struct<%s>" % (",".join(f.simpleString() for f in self)) + + def __repr__(self) -> str: + return "StructType([%s])" % ", ".join(str(field) for field in self) + + def jsonValue(self) -> Dict[str, Any]: + return {"type": self.typeName(), "fields": [f.jsonValue() for f in self]} + + @classmethod + def fromJson(cls, json: Dict[str, Any]) -> "StructType": + """ + Constructs :class:`StructType` from a schema defined in JSON format. + + Below is a JSON schema it must adhere to:: + + { + "title":"StructType", + "description":"Schema of StructType in json format", + "type":"object", + "properties":{ + "fields":{ + "description":"Array of struct fields", + "type":"array", + "items":{ + "type":"object", + "properties":{ + "name":{ + "description":"Name of the field", + "type":"string" + }, + "type":{ + "description": "Type of the field. Can either be + another nested StructType or primitive type", + "type":"object/string" + }, + "nullable":{ + "description":"If nulls are allowed", + "type":"boolean" + }, + "metadata":{ + "description":"Additional metadata to supply", + "type":"object" + }, + "required":[ + "name", + "type", + "nullable", + "metadata" + ] + } + } + } + } + } + + Parameters + ---------- + json : dict or a dict-like object e.g. JSON object + This "dict" must have "fields" key that returns an array of fields + each of which must have specific keys (name, type, nullable, metadata). + + Returns + ------- + :class:`StructType` + + Examples + -------- + >>> json_str = ''' + ... { + ... "fields": [ + ... { + ... "metadata": {}, + ... "name": "Person", + ... "nullable": true, + ... "type": { + ... "fields": [ + ... { + ... "metadata": {}, + ... "name": "name", + ... "nullable": false, + ... "type": "string" + ... }, + ... { + ... "metadata": {}, + ... "name": "surname", + ... "nullable": false, + ... "type": "string" + ... } + ... ], + ... "type": "struct" + ... } + ... } + ... ], + ... "type": "struct" + ... } + ... ''' + >>> import json + >>> scheme = StructType.fromJson(json.loads(json_str)) + >>> scheme.simpleString() + 'struct>' + """ + return StructType([StructField.fromJson(f) for f in json["fields"]]) + + def fieldNames(self) -> List[str]: + """ + Returns all field names in a list. + + Examples + -------- + >>> from pyspark.sql.types import StringType, StructField, StructType + >>> struct = StructType([StructField("f1", StringType(), True)]) + >>> struct.fieldNames() + ['f1'] + """ + return list(self.names) + + def needConversion(self) -> bool: + # We need convert Row()/namedtuple into tuple() + return True + + def toInternal(self, obj: Tuple) -> Tuple: + if obj is None: + return + + if self._needSerializeAnyField: + # Only calling toInternal function for fields that need conversion + if isinstance(obj, dict): + return tuple( + f.toInternal(obj.get(n)) if c else obj.get(n) + for n, f, c in zip(self.names, self.fields, self._needConversion) + ) + elif isinstance(obj, (tuple, list)): + return tuple( + f.toInternal(v) if c else v + for f, v, c in zip(self.fields, obj, self._needConversion) + ) + elif hasattr(obj, "__dict__"): + d = obj.__dict__ + return tuple( + f.toInternal(d.get(n)) if c else d.get(n) + for n, f, c in zip(self.names, self.fields, self._needConversion) + ) + else: + raise PySparkValueError( + error_class="UNEXPECTED_TUPLE_WITH_STRUCT", + message_parameters={"tuple": str(obj)}, + ) + else: + if isinstance(obj, dict): + return tuple(obj.get(n) for n in self.names) + elif isinstance(obj, (list, tuple)): + return tuple(obj) + elif hasattr(obj, "__dict__"): + d = obj.__dict__ + return tuple(d.get(n) for n in self.names) + else: + raise PySparkValueError( + error_class="UNEXPECTED_TUPLE_WITH_STRUCT", + message_parameters={"tuple": str(obj)}, + ) + + def fromInternal(self, obj: Tuple) -> "Row": + if obj is None: + return + if isinstance(obj, Row): + # it's already converted by pickler + return obj + + values: Union[Tuple, List] + if self._needSerializeAnyField: + # Only calling fromInternal function for fields that need conversion + values = [ + f.fromInternal(v) if c else v + for f, v, c in zip(self.fields, obj, self._needConversion) + ] + else: + values = obj + return _create_row(self.names, values) + + +class UserDefinedType(DataType): + """User-defined type (UDT). + + .. note:: WARN: Spark Internal Use Only + """ + + @classmethod + def typeName(cls) -> str: + return cls.__name__.lower() + + @classmethod + def sqlType(cls) -> DataType: + """ + Underlying SQL storage type for this UDT. + """ + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": "sqlType()"}, + ) + + @classmethod + def module(cls) -> str: + """ + The Python module of the UDT. + """ + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": "module()"}, + ) + + @classmethod + def scalaUDT(cls) -> str: + """ + The class name of the paired Scala UDT (could be '', if there + is no corresponding one). + """ + return "" + + def needConversion(self) -> bool: + return True + + @classmethod + def _cachedSqlType(cls) -> DataType: + """ + Cache the sqlType() into class, because it's heavily used in `toInternal`. + """ + if not hasattr(cls, "_cached_sql_type"): + cls._cached_sql_type = cls.sqlType() # type: ignore[attr-defined] + return cls._cached_sql_type # type: ignore[attr-defined] + + def toInternal(self, obj: Any) -> Any: + if obj is not None: + return self._cachedSqlType().toInternal(self.serialize(obj)) + + def fromInternal(self, obj: Any) -> Any: + v = self._cachedSqlType().fromInternal(obj) + if v is not None: + return self.deserialize(v) + + def serialize(self, obj: Any) -> Any: + """ + Converts a user-type object into a SQL datum. + """ + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": "toInternal()"}, + ) + + def deserialize(self, datum: Any) -> Any: + """ + Converts a SQL datum into a user-type object. + """ + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": "fromInternal()"}, + ) + + def simpleString(self) -> str: + return "udt" + + def json(self) -> str: + return json.dumps(self.jsonValue(), separators=(",", ":"), sort_keys=True) + + def jsonValue(self) -> Dict[str, Any]: + if self.scalaUDT(): + assert self.module() != "__main__", "UDT in __main__ cannot work with ScalaUDT" + schema = { + "type": "udt", + "class": self.scalaUDT(), + "pyClass": "%s.%s" % (self.module(), type(self).__name__), + "sqlType": self.sqlType().jsonValue(), + } + else: + ser = CloudPickleSerializer() + b = ser.dumps(type(self)) + schema = { + "type": "udt", + "pyClass": "%s.%s" % (self.module(), type(self).__name__), + "serializedClass": base64.b64encode(b).decode("utf8"), + "sqlType": self.sqlType().jsonValue(), + } + return schema + + @classmethod + def fromJson(cls, json: Dict[str, Any]) -> "UserDefinedType": + pyUDT = str(json["pyClass"]) # convert unicode to str + split = pyUDT.rfind(".") + pyModule = pyUDT[:split] + pyClass = pyUDT[split + 1 :] + m = __import__(pyModule, globals(), locals(), [pyClass]) + if not hasattr(m, pyClass): + s = base64.b64decode(json["serializedClass"].encode("utf-8")) + UDT = CloudPickleSerializer().loads(s) + else: + UDT = getattr(m, pyClass) + return UDT() + + def __eq__(self, other: Any) -> bool: + return type(self) == type(other) + + +_atomic_types: List[Type[DataType]] = [ + StringType, + CharType, + VarcharType, + BinaryType, + BooleanType, + DecimalType, + FloatType, + DoubleType, + ByteType, + ShortType, + IntegerType, + LongType, + DateType, + TimestampType, + TimestampNTZType, + NullType, +] +_all_atomic_types: Dict[str, Type[DataType]] = dict((t.typeName(), t) for t in _atomic_types) + +_complex_types: List[Type[Union[ArrayType, MapType, StructType]]] = [ArrayType, MapType, StructType] +_all_complex_types: Dict[str, Type[Union[ArrayType, MapType, StructType]]] = dict( + (v.typeName(), v) for v in _complex_types +) + +_LENGTH_CHAR = re.compile(r"char\(\s*(\d+)\s*\)") +_LENGTH_VARCHAR = re.compile(r"varchar\(\s*(\d+)\s*\)") +_FIXED_DECIMAL = re.compile(r"decimal\(\s*(\d+)\s*,\s*(-?\d+)\s*\)") +_INTERVAL_DAYTIME = re.compile(r"interval (day|hour|minute|second)( to (day|hour|minute|second))?") +_INTERVAL_YEARMONTH = re.compile(r"interval (year|month)( to (year|month))?") + + +def _parse_datatype_string(s: str) -> DataType: + """ + Parses the given data type string to a :class:`DataType`. The data type string format equals + :class:`DataType.simpleString`, except that the top level struct type can omit + the ``struct<>``. Since Spark 2.3, this also supports a schema in a DDL-formatted + string and case-insensitive strings. + + Examples + -------- + >>> _parse_datatype_string("int ") + IntegerType() + >>> _parse_datatype_string("INT ") + IntegerType() + >>> _parse_datatype_string("a: byte, b: decimal( 16 , 8 ) ") + StructType([StructField('a', ByteType(), True), StructField('b', DecimalType(16,8), True)]) + >>> _parse_datatype_string("a DOUBLE, b STRING") + StructType([StructField('a', DoubleType(), True), StructField('b', StringType(), True)]) + >>> _parse_datatype_string("a DOUBLE, b CHAR( 50 )") + StructType([StructField('a', DoubleType(), True), StructField('b', CharType(50), True)]) + >>> _parse_datatype_string("a DOUBLE, b VARCHAR( 50 )") + StructType([StructField('a', DoubleType(), True), StructField('b', VarcharType(50), True)]) + >>> _parse_datatype_string("a: array< short>") + StructType([StructField('a', ArrayType(ShortType(), True), True)]) + >>> _parse_datatype_string(" map ") + MapType(StringType(), StringType(), True) + + >>> # Error cases + >>> _parse_datatype_string("blabla") # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ParseException:... + >>> _parse_datatype_string("a: int,") # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ParseException:... + >>> _parse_datatype_string("array>> _parse_datatype_string("map>") # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ParseException:... + """ + sc = get_active_spark_context() + + def from_ddl_schema(type_str: str) -> DataType: + return _parse_datatype_json_string( + cast(JVMView, sc._jvm).org.apache.spark.sql.types.StructType.fromDDL(type_str).json() + ) + + def from_ddl_datatype(type_str: str) -> DataType: + return _parse_datatype_json_string( + cast(JVMView, sc._jvm) + .org.apache.spark.sql.api.python.PythonSQLUtils.parseDataType(type_str) + .json() + ) + + try: + # DDL format, "fieldname datatype, fieldname datatype". + return from_ddl_schema(s) + except Exception as e: + try: + # For backwards compatibility, "integer", "struct" and etc. + return from_ddl_datatype(s) + except BaseException: + try: + # For backwards compatibility, "fieldname: datatype, fieldname: datatype" case. + return from_ddl_datatype("struct<%s>" % s.strip()) + except BaseException: + raise e + + +def _parse_datatype_json_string(json_string: str) -> DataType: + """Parses the given data type JSON string. + + Examples + -------- + >>> import pickle + >>> def check_datatype(datatype): + ... pickled = pickle.loads(pickle.dumps(datatype)) + ... assert datatype == pickled + ... scala_datatype = spark._jsparkSession.parseDataType(datatype.json()) + ... python_datatype = _parse_datatype_json_string(scala_datatype.json()) + ... assert datatype == python_datatype + ... + >>> for cls in _all_atomic_types.values(): + ... if cls is not VarcharType and cls is not CharType: + ... check_datatype(cls()) + ... else: + ... check_datatype(cls(1)) + + >>> # Simple ArrayType. + >>> simple_arraytype = ArrayType(StringType(), True) + >>> check_datatype(simple_arraytype) + + >>> # Simple MapType. + >>> simple_maptype = MapType(StringType(), LongType()) + >>> check_datatype(simple_maptype) + + >>> # Simple StructType. + >>> simple_structtype = StructType([ + ... StructField("a", DecimalType(), False), + ... StructField("b", BooleanType(), True), + ... StructField("c", LongType(), True), + ... StructField("d", BinaryType(), False)]) + >>> check_datatype(simple_structtype) + + >>> # Complex StructType. + >>> complex_structtype = StructType([ + ... StructField("simpleArray", simple_arraytype, True), + ... StructField("simpleMap", simple_maptype, True), + ... StructField("simpleStruct", simple_structtype, True), + ... StructField("boolean", BooleanType(), False), + ... StructField("chars", CharType(10), False), + ... StructField("words", VarcharType(10), False), + ... StructField("withMeta", DoubleType(), False, {"name": "age"})]) + >>> check_datatype(complex_structtype) + + >>> # Complex ArrayType. + >>> complex_arraytype = ArrayType(complex_structtype, True) + >>> check_datatype(complex_arraytype) + + >>> # Complex MapType. + >>> complex_maptype = MapType(complex_structtype, + ... complex_arraytype, False) + >>> check_datatype(complex_maptype) + """ + return _parse_datatype_json_value(json.loads(json_string)) + + +def _parse_datatype_json_value(json_value: Union[dict, str]) -> DataType: + if not isinstance(json_value, dict): + if json_value in _all_atomic_types.keys(): + return _all_atomic_types[json_value]() + elif json_value == "decimal": + return DecimalType() + elif _FIXED_DECIMAL.match(json_value): + m = _FIXED_DECIMAL.match(json_value) + return DecimalType(int(m.group(1)), int(m.group(2))) # type: ignore[union-attr] + elif _INTERVAL_DAYTIME.match(json_value): + m = _INTERVAL_DAYTIME.match(json_value) + inverted_fields = DayTimeIntervalType._inverted_fields + first_field = inverted_fields.get(m.group(1)) # type: ignore[union-attr] + second_field = inverted_fields.get(m.group(3)) # type: ignore[union-attr] + if first_field is not None and second_field is None: + return DayTimeIntervalType(first_field) + return DayTimeIntervalType(first_field, second_field) + elif _INTERVAL_YEARMONTH.match(json_value): + m = _INTERVAL_YEARMONTH.match(json_value) + inverted_fields = YearMonthIntervalType._inverted_fields + first_field = inverted_fields.get(m.group(1)) # type: ignore[union-attr] + second_field = inverted_fields.get(m.group(3)) # type: ignore[union-attr] + if first_field is not None and second_field is None: + return YearMonthIntervalType(first_field) + return YearMonthIntervalType(first_field, second_field) + elif _LENGTH_CHAR.match(json_value): + m = _LENGTH_CHAR.match(json_value) + return CharType(int(m.group(1))) # type: ignore[union-attr] + elif _LENGTH_VARCHAR.match(json_value): + m = _LENGTH_VARCHAR.match(json_value) + return VarcharType(int(m.group(1))) # type: ignore[union-attr] + else: + raise PySparkValueError( + error_class="CANNOT_PARSE_DATATYPE", + message_parameters={"error": str(json_value)}, + ) + else: + tpe = json_value["type"] + if tpe in _all_complex_types: + return _all_complex_types[tpe].fromJson(json_value) + elif tpe == "udt": + return UserDefinedType.fromJson(json_value) + else: + raise PySparkValueError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": str(tpe)}, + ) + + +# Mapping Python types to Spark SQL DataType +_type_mappings = { + type(None): NullType, + bool: BooleanType, + int: LongType, + float: DoubleType, + str: StringType, + bytearray: BinaryType, + decimal.Decimal: DecimalType, + datetime.date: DateType, + datetime.datetime: TimestampType, # can be TimestampNTZType + datetime.time: TimestampType, # can be TimestampNTZType + datetime.timedelta: DayTimeIntervalType, + bytes: BinaryType, +} + +# Mapping Python array types to Spark SQL DataType +# We should be careful here. The size of these types in python depends on C +# implementation. We need to make sure that this conversion does not lose any +# precision. Also, JVM only support signed types, when converting unsigned types, +# keep in mind that it require 1 more bit when stored as signed types. +# +# Reference for C integer size, see: +# ISO/IEC 9899:201x specification, chapter 5.2.4.2.1 Sizes of integer types . +# Reference for python array typecode, see: +# https://docs.python.org/2/library/array.html +# https://docs.python.org/3.6/library/array.html +# Reference for JVM's supported integral types: +# http://docs.oracle.com/javase/specs/jvms/se8/html/jvms-2.html#jvms-2.3.1 + +_array_signed_int_typecode_ctype_mappings = { + "b": ctypes.c_byte, + "h": ctypes.c_short, + "i": ctypes.c_int, + "l": ctypes.c_long, +} + +_array_unsigned_int_typecode_ctype_mappings = { + "B": ctypes.c_ubyte, + "H": ctypes.c_ushort, + "I": ctypes.c_uint, + "L": ctypes.c_ulong, +} + + +def _int_size_to_type( + size: int, +) -> Optional[Union[Type[ByteType], Type[ShortType], Type[IntegerType], Type[LongType]]]: + """ + Return the Catalyst datatype from the size of integers. + """ + if size <= 8: + return ByteType + elif size <= 16: + return ShortType + elif size <= 32: + return IntegerType + elif size <= 64: + return LongType + else: + return None + + +# The list of all supported array typecodes, is stored here +_array_type_mappings: Dict[str, Type[DataType]] = { + # Warning: Actual properties for float and double in C is not specified in C. + # On almost every system supported by both python and JVM, they are IEEE 754 + # single-precision binary floating-point format and IEEE 754 double-precision + # binary floating-point format. And we do assume the same thing here for now. + "f": FloatType, + "d": DoubleType, +} + +# compute array typecode mappings for signed integer types +for _typecode in _array_signed_int_typecode_ctype_mappings.keys(): + size = ctypes.sizeof(_array_signed_int_typecode_ctype_mappings[_typecode]) * 8 + dt = _int_size_to_type(size) + if dt is not None: + _array_type_mappings[_typecode] = dt + +# compute array typecode mappings for unsigned integer types +for _typecode in _array_unsigned_int_typecode_ctype_mappings.keys(): + # JVM does not have unsigned types, so use signed types that is at least 1 + # bit larger to store + size = ctypes.sizeof(_array_unsigned_int_typecode_ctype_mappings[_typecode]) * 8 + 1 + dt = _int_size_to_type(size) + if dt is not None: + _array_type_mappings[_typecode] = dt + +# Type code 'u' in Python's array is deprecated since version 3.3, and will be +# removed in version 4.0. See: https://docs.python.org/3/library/array.html +if sys.version_info[0] < 4: + _array_type_mappings["u"] = StringType + + +def _from_numpy_type(nt: "np.dtype") -> Optional[DataType]: + """Convert NumPy type to Spark data type.""" + import numpy as np + + if nt == np.dtype("int8"): + return ByteType() + elif nt == np.dtype("int16"): + return ShortType() + elif nt == np.dtype("int32"): + return IntegerType() + elif nt == np.dtype("int64"): + return LongType() + elif nt == np.dtype("float32"): + return FloatType() + elif nt == np.dtype("float64"): + return DoubleType() + + return None + + +def _infer_type( + obj: Any, + infer_dict_as_struct: bool = False, + infer_array_from_first_element: bool = False, + prefer_timestamp_ntz: bool = False, +) -> DataType: + """Infer the DataType from obj""" + if obj is None: + return NullType() + + if hasattr(obj, "__UDT__"): + return obj.__UDT__ + + dataType = _type_mappings.get(type(obj)) + if dataType is DecimalType: + # the precision and scale of `obj` may be different from row to row. + return DecimalType(38, 18) + if dataType is TimestampType and prefer_timestamp_ntz and obj.tzinfo is None: + return TimestampNTZType() + if dataType is DayTimeIntervalType: + return DayTimeIntervalType() + if dataType is YearMonthIntervalType: + return YearMonthIntervalType() + elif dataType is not None: + return dataType() + + if isinstance(obj, dict): + if infer_dict_as_struct: + struct = StructType() + for key, value in obj.items(): + if key is not None and value is not None: + struct.add( + key, + _infer_type( + value, + infer_dict_as_struct, + infer_array_from_first_element, + prefer_timestamp_ntz, + ), + True, + ) + return struct + else: + for key, value in obj.items(): + if key is not None and value is not None: + return MapType( + _infer_type( + key, + infer_dict_as_struct, + infer_array_from_first_element, + prefer_timestamp_ntz, + ), + _infer_type( + value, + infer_dict_as_struct, + infer_array_from_first_element, + prefer_timestamp_ntz, + ), + True, + ) + return MapType(NullType(), NullType(), True) + elif isinstance(obj, list): + if len(obj) > 0: + if infer_array_from_first_element: + return ArrayType( + _infer_type(obj[0], infer_dict_as_struct, prefer_timestamp_ntz), True + ) + else: + return ArrayType( + reduce( + _merge_type, + (_infer_type(v, infer_dict_as_struct, prefer_timestamp_ntz) for v in obj), + ), + True, + ) + return ArrayType(NullType(), True) + elif isinstance(obj, array): + if obj.typecode in _array_type_mappings: + return ArrayType(_array_type_mappings[obj.typecode](), False) + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": f"array({obj.typecode})"}, + ) + else: + try: + return _infer_schema( + obj, + infer_dict_as_struct=infer_dict_as_struct, + infer_array_from_first_element=infer_array_from_first_element, + ) + except TypeError: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(obj).__name__}, + ) + + +def _infer_schema( + row: Any, + names: Optional[List[str]] = None, + infer_dict_as_struct: bool = False, + infer_array_from_first_element: bool = False, + prefer_timestamp_ntz: bool = False, +) -> StructType: + """Infer the schema from dict/namedtuple/object""" + items: Iterable[Tuple[str, Any]] + if isinstance(row, dict): + items = sorted(row.items()) + + elif isinstance(row, (tuple, list)): + if hasattr(row, "__fields__"): # Row + items = zip(row.__fields__, tuple(row)) # type: ignore[union-attr] + elif hasattr(row, "_fields"): # namedtuple + items = zip(row._fields, tuple(row)) # type: ignore[union-attr] + else: + if names is None: + names = ["_%d" % i for i in range(1, len(row) + 1)] + elif len(names) < len(row): + names.extend("_%d" % i for i in range(len(names) + 1, len(row) + 1)) + items = zip(names, row) + + elif hasattr(row, "__dict__"): # object + items = sorted(row.__dict__.items()) + + else: + raise PySparkTypeError( + error_class="CANNOT_INFER_SCHEMA_FOR_TYPE", + message_parameters={"data_type": type(row).__name__}, + ) + + fields = [] + for k, v in items: + try: + fields.append( + StructField( + k, + _infer_type( + v, + infer_dict_as_struct, + infer_array_from_first_element, + prefer_timestamp_ntz, + ), + True, + ) + ) + except TypeError: + raise PySparkTypeError( + error_class="CANNOT_INFER_TYPE_FOR_FIELD", + message_parameters={"field_name": k}, + ) + return StructType(fields) + + +def _has_nulltype(dt: DataType) -> bool: + """Return whether there is a NullType in `dt` or not""" + if isinstance(dt, StructType): + return any(_has_nulltype(f.dataType) for f in dt.fields) + elif isinstance(dt, ArrayType): + return _has_nulltype((dt.elementType)) + elif isinstance(dt, MapType): + return _has_nulltype(dt.keyType) or _has_nulltype(dt.valueType) + else: + return isinstance(dt, NullType) + + +def _has_type(dt: DataType, dts: Union[type, Tuple[type, ...]]) -> bool: + """Return whether there are specified types""" + if isinstance(dt, dts): + return True + elif isinstance(dt, StructType): + return any(_has_type(f.dataType, dts) for f in dt.fields) + elif isinstance(dt, ArrayType): + return _has_type(dt.elementType, dts) + elif isinstance(dt, MapType): + return _has_type(dt.keyType, dts) or _has_type(dt.valueType, dts) + else: + return False + + +@overload +def _merge_type(a: StructType, b: StructType, name: Optional[str] = None) -> StructType: + ... + + +@overload +def _merge_type(a: ArrayType, b: ArrayType, name: Optional[str] = None) -> ArrayType: + ... + + +@overload +def _merge_type(a: MapType, b: MapType, name: Optional[str] = None) -> MapType: + ... + + +@overload +def _merge_type(a: DataType, b: DataType, name: Optional[str] = None) -> DataType: + ... + + +def _merge_type( + a: Union[StructType, ArrayType, MapType, DataType], + b: Union[StructType, ArrayType, MapType, DataType], + name: Optional[str] = None, +) -> Union[StructType, ArrayType, MapType, DataType]: + if name is None: + + def new_msg(msg: str) -> str: + return msg + + def new_name(n: str) -> str: + return "field %s" % n + + else: + + def new_msg(msg: str) -> str: + return "%s: %s" % (name, msg) + + def new_name(n: str) -> str: + return "field %s in %s" % (n, name) + + if isinstance(a, NullType): + return b + elif isinstance(b, NullType): + return a + elif isinstance(a, TimestampType) and isinstance(b, TimestampNTZType): + return a + elif isinstance(a, TimestampNTZType) and isinstance(b, TimestampType): + return b + elif isinstance(a, AtomicType) and isinstance(b, StringType): + return b + elif isinstance(a, StringType) and isinstance(b, AtomicType): + return a + elif type(a) is not type(b): + # TODO: type cast (such as int -> long) + raise PySparkTypeError( + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": type(a).__name__, "data_type2": type(b).__name__}, + ) + + # same type + if isinstance(a, StructType): + nfs = dict((f.name, f.dataType) for f in cast(StructType, b).fields) + fields = [ + StructField( + f.name, _merge_type(f.dataType, nfs.get(f.name, NullType()), name=new_name(f.name)) + ) + for f in a.fields + ] + names = set([f.name for f in fields]) + for n in nfs: + if n not in names: + fields.append(StructField(n, nfs[n])) + return StructType(fields) + + elif isinstance(a, ArrayType): + return ArrayType( + _merge_type( + a.elementType, cast(ArrayType, b).elementType, name="element in array %s" % name + ), + True, + ) + + elif isinstance(a, MapType): + return MapType( + _merge_type(a.keyType, cast(MapType, b).keyType, name="key of map %s" % name), + _merge_type(a.valueType, cast(MapType, b).valueType, name="value of map %s" % name), + True, + ) + else: + return a + + +def _need_converter(dataType: DataType) -> bool: + if isinstance(dataType, StructType): + return True + elif isinstance(dataType, ArrayType): + return _need_converter(dataType.elementType) + elif isinstance(dataType, MapType): + return _need_converter(dataType.keyType) or _need_converter(dataType.valueType) + elif isinstance(dataType, NullType): + return True + else: + return False + + +def _create_converter(dataType: DataType) -> Callable: + """Create a converter to drop the names of fields in obj""" + if not _need_converter(dataType): + return lambda x: x + + if isinstance(dataType, ArrayType): + conv = _create_converter(dataType.elementType) + return lambda row: [conv(v) for v in row] + + elif isinstance(dataType, MapType): + kconv = _create_converter(dataType.keyType) + vconv = _create_converter(dataType.valueType) + return lambda row: dict((kconv(k), vconv(v)) for k, v in row.items()) + + elif isinstance(dataType, NullType): + return lambda x: None + + elif not isinstance(dataType, StructType): + return lambda x: x + + # dataType must be StructType + names = [f.name for f in dataType.fields] + converters = [_create_converter(f.dataType) for f in dataType.fields] + convert_fields = any(_need_converter(f.dataType) for f in dataType.fields) + + def convert_struct(obj: Any) -> Optional[Tuple]: + if obj is None: + return None + + if isinstance(obj, (tuple, list)): + if convert_fields: + return tuple(conv(v) for v, conv in zip(obj, converters)) + else: + return tuple(obj) + + if isinstance(obj, dict): + d = obj + elif hasattr(obj, "__dict__"): # object + d = obj.__dict__ + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(obj).__name__}, + ) + + if convert_fields: + return tuple([conv(d.get(name)) for name, conv in zip(names, converters)]) + else: + return tuple([d.get(name) for name in names]) + + return convert_struct + + +_acceptable_types = { + BooleanType: (bool,), + ByteType: (int,), + ShortType: (int,), + IntegerType: (int,), + LongType: (int,), + FloatType: (float,), + DoubleType: (float,), + DecimalType: (decimal.Decimal,), + StringType: (str,), + CharType: (str,), + VarcharType: (str,), + BinaryType: (bytearray, bytes), + DateType: (datetime.date, datetime.datetime), + TimestampType: (datetime.datetime,), + TimestampNTZType: (datetime.datetime,), + DayTimeIntervalType: (datetime.timedelta,), + ArrayType: (list, tuple, array), + MapType: (dict,), + StructType: (tuple, list, dict), +} + + +def _make_type_verifier( + dataType: DataType, + nullable: bool = True, + name: Optional[str] = None, +) -> Callable: + """ + Make a verifier that checks the type of obj against dataType and raises a TypeError if they do + not match. + + This verifier also checks the value of obj against datatype and raises a ValueError if it's not + within the allowed range, e.g. using 128 as ByteType will overflow. Note that, Python float is + not checked, so it will become infinity when cast to Java float, if it overflows. + + Examples + -------- + >>> _make_type_verifier(StructType([]))(None) + >>> _make_type_verifier(StringType())("") + >>> _make_type_verifier(LongType())(0) + >>> _make_type_verifier(LongType())(1 << 64) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> _make_type_verifier(ArrayType(ShortType()))(list(range(3))) + >>> _make_type_verifier(ArrayType(StringType()))(set()) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkTypeError:... + >>> _make_type_verifier(MapType(StringType(), IntegerType()))({}) + >>> _make_type_verifier(StructType([]))(()) + >>> _make_type_verifier(StructType([]))([]) + >>> _make_type_verifier(StructType([]))([1]) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> # Check if numeric values are within the allowed range. + >>> _make_type_verifier(ByteType())(12) + >>> _make_type_verifier(ByteType())(1234) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> _make_type_verifier(ByteType(), False)(None) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> _make_type_verifier( + ... ArrayType(ShortType(), False))([1, None]) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> _make_type_verifier( # doctest: +IGNORE_EXCEPTION_DETAIL + ... MapType(StringType(), IntegerType()) + ... )({None: 1}) + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + >>> schema = StructType().add("a", IntegerType()).add("b", StringType(), False) + >>> _make_type_verifier(schema)((1, None)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + pyspark.errors.exceptions.base.PySparkValueError:... + """ + + if name is None: + + def new_msg(msg: str) -> str: + return msg + + def new_name(n: str) -> str: + return "field %s" % n + + else: + + def new_msg(msg: str) -> str: + return "%s: %s" % (name, msg) + + def new_name(n: str) -> str: + return "field %s in %s" % (n, name) + + def verify_nullability(obj: Any) -> bool: + if obj is None: + if nullable: + return True + else: + raise PySparkValueError( + error_class="CANNOT_BE_NONE", + message_parameters={"arg_name": "obj"}, + ) + else: + return False + + _type = type(dataType) + + def assert_acceptable_types(obj: Any) -> None: + assert _type in _acceptable_types, new_msg( + "unknown datatype: %s for object %r" % (dataType, obj) + ) + + def verify_acceptable_types(obj: Any) -> None: + # subclass of them can not be fromInternal in JVM + if type(obj) not in _acceptable_types[_type]: + raise PySparkTypeError( + error_class="CANNOT_ACCEPT_OBJECT_IN_TYPE", + message_parameters={ + "data_type": str(dataType), + "obj_name": str(obj), + "obj_type": type(obj).__name__, + }, + ) + + if isinstance(dataType, (StringType, CharType, VarcharType)): + # StringType, CharType and VarcharType can work with any types + def verify_value(obj: Any) -> None: + pass + + elif isinstance(dataType, UserDefinedType): + verifier = _make_type_verifier(dataType.sqlType(), name=name) + + def verify_udf(obj: Any) -> None: + if not (hasattr(obj, "__UDT__") and obj.__UDT__ == dataType): + raise PySparkValueError( + error_class="NOT_INSTANCE_OF", + message_parameters={ + "value": str(obj), + "data_type": str(dataType), + }, + ) + verifier(dataType.toInternal(obj)) + + verify_value = verify_udf + + elif isinstance(dataType, ByteType): + + def verify_byte(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + if obj < -128 or obj > 127: + raise PySparkValueError( + error_class="VALUE_OUT_OF_BOUND", + message_parameters={ + "arg_name": "obj", + "lower_bound": "127", + "upper_bound": "-127", + "actual": str(obj), + }, + ) + + verify_value = verify_byte + + elif isinstance(dataType, ShortType): + + def verify_short(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + if obj < -32768 or obj > 32767: + raise PySparkValueError( + error_class="VALUE_OUT_OF_BOUND", + message_parameters={ + "arg_name": "obj", + "lower_bound": "32767", + "upper_bound": "-32768", + "actual": str(obj), + }, + ) + + verify_value = verify_short + + elif isinstance(dataType, IntegerType): + + def verify_integer(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + if obj < -2147483648 or obj > 2147483647: + raise PySparkValueError( + error_class="VALUE_OUT_OF_BOUND", + message_parameters={ + "arg_name": "obj", + "lower_bound": "2147483647", + "upper_bound": "-2147483648", + "actual": str(obj), + }, + ) + + verify_value = verify_integer + + elif isinstance(dataType, LongType): + + def verify_long(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + if obj < -9223372036854775808 or obj > 9223372036854775807: + raise PySparkValueError( + error_class="VALUE_OUT_OF_BOUND", + message_parameters={ + "arg_name": "obj", + "lower_bound": "9223372036854775807", + "upper_bound": "-9223372036854775808", + "actual": str(obj), + }, + ) + + verify_value = verify_long + + elif isinstance(dataType, ArrayType): + element_verifier = _make_type_verifier( + dataType.elementType, dataType.containsNull, name="element in array %s" % name + ) + + def verify_array(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + for i in obj: + element_verifier(i) + + verify_value = verify_array + + elif isinstance(dataType, MapType): + key_verifier = _make_type_verifier(dataType.keyType, False, name="key of map %s" % name) + value_verifier = _make_type_verifier( + dataType.valueType, dataType.valueContainsNull, name="value of map %s" % name + ) + + def verify_map(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + for k, v in obj.items(): + key_verifier(k) + value_verifier(v) + + verify_value = verify_map + + elif isinstance(dataType, StructType): + verifiers = [] + for f in dataType.fields: + verifier = _make_type_verifier(f.dataType, f.nullable, name=new_name(f.name)) + verifiers.append((f.name, verifier)) + + def verify_struct(obj: Any) -> None: + assert_acceptable_types(obj) + + if isinstance(obj, dict): + for f, verifier in verifiers: + verifier(obj.get(f)) + elif isinstance(obj, (tuple, list)): + if len(obj) != len(verifiers): + raise PySparkValueError( + error_class="LENGTH_SHOULD_BE_THE_SAME", + message_parameters={ + "arg1": "obj", + "arg2": "fields", + "arg1_length": str(len(obj)), + "arg2_length": str(len(verifiers)), + }, + ) + for v, (_, verifier) in zip(obj, verifiers): + verifier(v) + elif hasattr(obj, "__dict__"): + d = obj.__dict__ + for f, verifier in verifiers: + verifier(d.get(f)) + else: + raise PySparkTypeError( + error_class="CANNOT_ACCEPT_OBJECT_IN_TYPE", + message_parameters={ + "data_type": "StructType", + "obj_name": str(obj), + "obj_type": type(obj).__name__, + }, + ) + + verify_value = verify_struct + + else: + + def verify_default(obj: Any) -> None: + assert_acceptable_types(obj) + verify_acceptable_types(obj) + + verify_value = verify_default + + def verify(obj: Any) -> None: + if not verify_nullability(obj): + verify_value(obj) + + return verify + + +# This is used to unpickle a Row from JVM +def _create_row_inbound_converter(dataType: DataType) -> Callable: + return lambda *a: dataType.fromInternal(a) + + +def _create_row( + fields: Union["Row", List[str]], values: Union[Tuple[Any, ...], List[Any]] +) -> "Row": + row = Row(*values) + row.__fields__ = fields + return row + + +class Row(tuple): + + """ + A row in :class:`DataFrame`. + The fields in it can be accessed: + + * like attributes (``row.key``) + * like dictionary values (``row[key]``) + + ``key in row`` will search through row keys. + + Row can be used to create a row object by using named arguments. + It is not allowed to omit a named argument to represent that the value is + None or missing. This should be explicitly set to None in this case. + + .. versionchanged:: 3.0.0 + Rows created from named arguments no longer have + field names sorted alphabetically and will be ordered in the position as + entered. + + Examples + -------- + >>> from pyspark.sql import Row + >>> row = Row(name="Alice", age=11) + >>> row + Row(name='Alice', age=11) + >>> row['name'], row['age'] + ('Alice', 11) + >>> row.name, row.age + ('Alice', 11) + >>> 'name' in row + True + >>> 'wrong_key' in row + False + + Row also can be used to create another Row like class, then it + could be used to create Row objects, such as + + >>> Person = Row("name", "age") + >>> Person + + >>> 'name' in Person + True + >>> 'wrong_key' in Person + False + >>> Person("Alice", 11) + Row(name='Alice', age=11) + + This form can also be used to create rows as tuple values, i.e. with unnamed + fields. + + >>> row1 = Row("Alice", 11) + >>> row2 = Row(name="Alice", age=11) + >>> row1 == row2 + True + """ + + @overload + def __new__(cls, *args: str) -> "Row": + ... + + @overload + def __new__(cls, **kwargs: Any) -> "Row": + ... + + def __new__(cls, *args: Optional[str], **kwargs: Optional[Any]) -> "Row": + if args and kwargs: + raise PySparkValueError( + error_class="CANNOT_SET_TOGETHER", + message_parameters={"arg_list": "args and kwargs"}, + ) + if kwargs: + # create row objects + row = tuple.__new__(cls, list(kwargs.values())) + row.__fields__ = list(kwargs.keys()) + return row + else: + # create row class or objects + return tuple.__new__(cls, args) + + def asDict(self, recursive: bool = False) -> Dict[str, Any]: + """ + Return as a dict + + Parameters + ---------- + recursive : bool, optional + turns the nested Rows to dict (default: False). + + Notes + ----- + If a row contains duplicate field names, e.g., the rows of a join + between two :class:`DataFrame` that both have the fields of same names, + one of the duplicate fields will be selected by ``asDict``. ``__getitem__`` + will also return one of the duplicate fields, however returned value might + be different to ``asDict``. + + Examples + -------- + >>> from pyspark.sql import Row + >>> Row(name="Alice", age=11).asDict() == {'name': 'Alice', 'age': 11} + True + >>> row = Row(key=1, value=Row(name='a', age=2)) + >>> row.asDict() == {'key': 1, 'value': Row(name='a', age=2)} + True + >>> row.asDict(True) == {'key': 1, 'value': {'name': 'a', 'age': 2}} + True + """ + if not hasattr(self, "__fields__"): + raise PySparkTypeError( + error_class="CANNOT_CONVERT_TYPE", + message_parameters={ + "from_type": "Row", + "to_type": "dict", + }, + ) + + if recursive: + + def conv(obj: Any) -> Any: + if isinstance(obj, Row): + return obj.asDict(True) + elif isinstance(obj, list): + return [conv(o) for o in obj] + elif isinstance(obj, dict): + return dict((k, conv(v)) for k, v in obj.items()) + else: + return obj + + return dict(zip(self.__fields__, (conv(o) for o in self))) + else: + return dict(zip(self.__fields__, self)) + + def __contains__(self, item: Any) -> bool: + if hasattr(self, "__fields__"): + return item in self.__fields__ + else: + return super(Row, self).__contains__(item) + + # let object acts like class + def __call__(self, *args: Any) -> "Row": + """create new Row object""" + if len(args) > len(self): + raise PySparkValueError( + error_class="TOO_MANY_VALUES", + message_parameters={ + "expected": str(len(self)), + "item": "fields", + "actual": str(len(args)), + }, + ) + return _create_row(self, args) + + def __getitem__(self, item: Any) -> Any: + if isinstance(item, (int, slice)): + return super(Row, self).__getitem__(item) + try: + # it will be slow when it has many fields, + # but this will not be used in normal cases + idx = self.__fields__.index(item) + return super(Row, self).__getitem__(idx) + except IndexError: + raise KeyError(item) + except ValueError: + raise PySparkValueError(item) + + def __getattr__(self, item: str) -> Any: + if item.startswith("__"): + raise AttributeError(item) + try: + # it will be slow when it has many fields, + # but this will not be used in normal cases + idx = self.__fields__.index(item) + return self[idx] + except IndexError: + raise AttributeError(item) + except ValueError: + raise AttributeError(item) + + def __setattr__(self, key: Any, value: Any) -> None: + if key != "__fields__": + raise RuntimeError("Row is read-only") + self.__dict__[key] = value + + def __reduce__( + self, + ) -> Union[str, Tuple[Any, ...]]: + """Returns a tuple so Python knows how to pickle Row.""" + if hasattr(self, "__fields__"): + return (_create_row, (self.__fields__, tuple(self))) + else: + return tuple.__reduce__(self) + + def __repr__(self) -> str: + """Printable representation of Row used in Python REPL.""" + if hasattr(self, "__fields__"): + return "Row(%s)" % ", ".join( + "%s=%r" % (k, v) for k, v in zip(self.__fields__, tuple(self)) + ) + else: + return "" % ", ".join(repr(field) for field in self) + + +class DateConverter: + def can_convert(self, obj: Any) -> bool: + return isinstance(obj, datetime.date) + + def convert(self, obj: datetime.date, gateway_client: GatewayClient) -> JavaObject: + Date = JavaClass("java.sql.Date", gateway_client) + return Date.valueOf(obj.strftime("%Y-%m-%d")) + + +class DatetimeConverter: + def can_convert(self, obj: Any) -> bool: + return isinstance(obj, datetime.datetime) + + def convert(self, obj: datetime.datetime, gateway_client: GatewayClient) -> JavaObject: + Timestamp = JavaClass("java.sql.Timestamp", gateway_client) + seconds = ( + calendar.timegm(obj.utctimetuple()) if obj.tzinfo else time.mktime(obj.timetuple()) + ) + t = Timestamp(int(seconds) * 1000) + t.setNanos(obj.microsecond * 1000) + return t + + +class DatetimeNTZConverter: + def can_convert(self, obj: Any) -> bool: + from pyspark.sql.utils import is_timestamp_ntz_preferred + + return ( + isinstance(obj, datetime.datetime) + and obj.tzinfo is None + and is_timestamp_ntz_preferred() + ) + + def convert(self, obj: datetime.datetime, gateway_client: GatewayClient) -> JavaObject: + seconds = calendar.timegm(obj.utctimetuple()) + DateTimeUtils = JavaClass( + "org.apache.spark.sql.catalyst.util.DateTimeUtils", + gateway_client, + ) + return DateTimeUtils.microsToLocalDateTime(int(seconds) * 1000000 + obj.microsecond) + + +class DayTimeIntervalTypeConverter: + def can_convert(self, obj: Any) -> bool: + return isinstance(obj, datetime.timedelta) + + def convert(self, obj: datetime.timedelta, gateway_client: GatewayClient) -> JavaObject: + IntervalUtils = JavaClass( + "org.apache.spark.sql.catalyst.util.IntervalUtils", + gateway_client, + ) + return IntervalUtils.microsToDuration( + (math.floor(obj.total_seconds()) * 1000000) + obj.microseconds + ) + + +class NumpyScalarConverter: + def can_convert(self, obj: Any) -> bool: + return has_numpy and isinstance(obj, np.generic) + + def convert(self, obj: "np.generic", gateway_client: GatewayClient) -> Any: + return obj.item() + + +class NumpyArrayConverter: + def _from_numpy_type_to_java_type( + self, nt: "np.dtype", gateway: JavaGateway + ) -> Optional[JavaClass]: + """Convert NumPy type to Py4J Java type.""" + if nt in [np.dtype("int8"), np.dtype("int16")]: + # Mapping int8 to gateway.jvm.byte causes + # TypeError: 'bytes' object does not support item assignment + return gateway.jvm.short + elif nt == np.dtype("int32"): + return gateway.jvm.int + elif nt == np.dtype("int64"): + return gateway.jvm.long + elif nt == np.dtype("float32"): + return gateway.jvm.float + elif nt == np.dtype("float64"): + return gateway.jvm.double + elif nt == np.dtype("bool"): + return gateway.jvm.boolean + + return None + + def can_convert(self, obj: Any) -> bool: + return has_numpy and isinstance(obj, np.ndarray) and obj.ndim == 1 + + def convert(self, obj: "np.ndarray", gateway_client: GatewayClient) -> JavaObject: + from pyspark import SparkContext + + gateway = SparkContext._gateway + assert gateway is not None + plist = obj.tolist() + + if len(obj) > 0 and isinstance(plist[0], str): + jtpe = gateway.jvm.String + else: + jtpe = self._from_numpy_type_to_java_type(obj.dtype, gateway) + if jtpe is None: + raise PySparkTypeError( + error_class="UNSUPPORTED_NUMPY_ARRAY_SCALAR", + message_parameters={"dtype": str(obj.dtype)}, + ) + jarr = gateway.new_array(jtpe, len(obj)) + for i in range(len(plist)): + jarr[i] = plist[i] + return jarr + + +# datetime is a subclass of date, we should register DatetimeConverter first +register_input_converter(DatetimeNTZConverter()) +register_input_converter(DatetimeConverter()) +register_input_converter(DateConverter()) +register_input_converter(DayTimeIntervalTypeConverter()) +register_input_converter(NumpyScalarConverter()) +# NumPy array satisfies py4j.java_collections.ListConverter, +# so prepend NumpyArrayConverter +register_input_converter(NumpyArrayConverter(), prepend=True) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + + globs = globals() + globs["spark"] = SparkSession.builder.getOrCreate() + (failure_count, test_count) = doctest.testmod( + globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udf.py new file mode 100644 index 0000000..7d7784d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udf.py @@ -0,0 +1,716 @@ +# +# 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. +# +""" +User-defined function related classes and functions +""" +from inspect import getfullargspec + +import functools +import inspect +import sys +import warnings +from typing import Callable, Any, TYPE_CHECKING, Optional, cast, Union + +from py4j.java_gateway import JavaObject + +from pyspark import SparkContext +from pyspark.profiler import Profiler +from pyspark.rdd import _prepare_for_python_RDD, PythonEvalType +from pyspark.sql.column import Column, _to_java_column, _to_java_expr, _to_seq +from pyspark.sql.types import ( + DataType, + StringType, + StructType, + _parse_datatype_string, +) +from pyspark.sql.utils import get_active_spark_context +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version +from pyspark.errors import PySparkTypeError, PySparkNotImplementedError + +if TYPE_CHECKING: + from pyspark.sql._typing import DataTypeOrString, ColumnOrName, UserDefinedFunctionLike + from pyspark.sql.session import SparkSession + +__all__ = ["UDFRegistration"] + + +def _wrap_function( + sc: SparkContext, func: Callable[..., Any], returnType: Optional[DataType] = None +) -> JavaObject: + command: Any + if returnType is None: + command = func + else: + command = (func, returnType) + pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command) + assert sc._jvm is not None + return sc._jvm.SimplePythonFunction( + bytearray(pickled_command), + env, + includes, + sc.pythonExec, + sc.pythonVer, + broadcast_vars, + sc._javaAccumulator, + ) + + +def _create_udf( + f: Callable[..., Any], + returnType: "DataTypeOrString", + evalType: int, + name: Optional[str] = None, + deterministic: bool = True, +) -> "UserDefinedFunctionLike": + """Create a regular(non-Arrow-optimized) Python UDF.""" + # Set the name of the UserDefinedFunction object to be the name of function f + udf_obj = UserDefinedFunction( + f, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + ) + return udf_obj._wrapped() + + +def _create_py_udf( + f: Callable[..., Any], + returnType: "DataTypeOrString", + useArrow: Optional[bool] = None, +) -> "UserDefinedFunctionLike": + """Create a regular/Arrow-optimized Python UDF.""" + # The following table shows the results when the type coercion in Arrow is needed, that is, + # when the user-specified return type(SQL Type) of the UDF and the actual instance(Python + # Value(Type)) that the UDF returns are different. + # Arrow and Pickle have different type coercion rules, so a UDF might have a different result + # with/without Arrow optimization. That's the main reason the Arrow optimization for Python + # UDFs is disabled by default. + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+ # noqa + # |SQL Type \ Python Value(Type)|None(NoneType)|True(bool)|1(int)| a(str)| 1970-01-01(date)|1970-01-01 00:00:00(datetime)|1.0(float)|array('i', [1])(array)|[1](list)| (1,)(tuple)|bytearray(b'ABC')(bytearray)| 1(Decimal)|{'a': 1}(dict)| # noqa + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+ # noqa + # | boolean| None| True| None| None| None| None| None| None| None| None| None| None| None| # noqa + # | tinyint| None| None| 1| None| None| None| None| None| None| None| None| None| None| # noqa + # | smallint| None| None| 1| None| None| None| None| None| None| None| None| None| None| # noqa + # | int| None| None| 1| None| None| None| None| None| None| None| None| None| None| # noqa + # | bigint| None| None| 1| None| None| None| None| None| None| None| None| None| None| # noqa + # | string| None| 'true'| '1'| 'a'|'java.util.Gregor...| 'java.util.Gregor...| '1.0'| '[I@120d813a'| '[1]'|'[Ljava.lang.Obje...| '[B@48571878'| '1'| '{a=1}'| # noqa + # | date| None| X| X| X|datetime.date(197...| datetime.date(197...| X| X| X| X| X| X| X| # noqa + # | timestamp| None| X| X| X| X| datetime.datetime...| X| X| X| X| X| X| X| # noqa + # | float| None| None| None| None| None| None| 1.0| None| None| None| None| None| None| # noqa + # | double| None| None| None| None| None| None| 1.0| None| None| None| None| None| None| # noqa + # | binary| None| None| None|bytearray(b'a')| None| None| None| None| None| None| bytearray(b'ABC')| None| None| # noqa + # | decimal(10,0)| None| None| None| None| None| None| None| None| None| None| None|Decimal('1')| None| # noqa + # +-----------------------------+--------------+----------+------+---------------+--------------------+-----------------------------+----------+----------------------+---------+--------------------+----------------------------+------------+--------------+ # noqa + # Note: Python 3.9.15, Pandas 1.5.2 and PyArrow 10.0.1 are used. + # Note: The values of 'SQL Type' are DDL formatted strings, which can be used as `returnType`s. + # Note: The values inside the table are generated by `repr`. X' means it throws an exception + # during the conversion. + + if useArrow is None: + from pyspark.sql import SparkSession + + session = SparkSession._instantiatedSession + is_arrow_enabled = ( + False + if session is None + else session.conf.get("spark.sql.execution.pythonUDF.arrow.enabled") == "true" + ) + else: + is_arrow_enabled = useArrow + + eval_type: int = PythonEvalType.SQL_BATCHED_UDF + + if is_arrow_enabled: + try: + is_func_with_args = len(getfullargspec(f).args) > 0 + except TypeError: + is_func_with_args = False + if is_func_with_args: + require_minimum_pandas_version() + require_minimum_pyarrow_version() + eval_type = PythonEvalType.SQL_ARROW_BATCHED_UDF + else: + warnings.warn( + "Arrow optimization for Python UDFs cannot be enabled.", + UserWarning, + ) + + return _create_udf(f, returnType, eval_type) + + +class UserDefinedFunction: + """ + User defined function in Python + + .. versionadded:: 1.3 + + Notes + ----- + The constructor of this class is not supposed to be directly called. + Use :meth:`pyspark.sql.functions.udf` or :meth:`pyspark.sql.functions.pandas_udf` + to create this instance. + """ + + def __init__( + self, + func: Callable[..., Any], + returnType: "DataTypeOrString" = StringType(), + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_BATCHED_UDF, + deterministic: bool = True, + ): + if not callable(func): + raise PySparkTypeError( + error_class="NOT_CALLABLE", + message_parameters={"arg_name": "func", "arg_type": type(func).__name__}, + ) + + if not isinstance(returnType, (DataType, str)): + raise PySparkTypeError( + error_class="NOT_DATATYPE_OR_STR", + message_parameters={ + "arg_name": "returnType", + "arg_type": type(returnType).__name__, + }, + ) + + if not isinstance(evalType, int): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "evalType", "arg_type": type(evalType).__name__}, + ) + + self.func = func + self._returnType = returnType + # Stores UserDefinedPythonFunctions jobj, once initialized + self._returnType_placeholder: Optional[DataType] = None + self._judf_placeholder = None + self._name = name or ( + func.__name__ if hasattr(func, "__name__") else func.__class__.__name__ + ) + self.evalType = evalType + self.deterministic = deterministic + + @property + def returnType(self) -> DataType: + # This makes sure this is called after SparkContext is initialized. + # ``_parse_datatype_string`` accesses to JVM for parsing a DDL formatted string. + # TODO: PythonEvalType.SQL_BATCHED_UDF + if self._returnType_placeholder is None: + if isinstance(self._returnType, DataType): + self._returnType_placeholder = self._returnType + else: + self._returnType_placeholder = _parse_datatype_string(self._returnType) + + if ( + self.evalType == PythonEvalType.SQL_SCALAR_PANDAS_UDF + or self.evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + ): + try: + to_arrow_type(self._returnType_placeholder) + except TypeError: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": f"Invalid return type with scalar Pandas UDFs: " + f"{self._returnType_placeholder}" + }, + ) + elif ( + self.evalType == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF + or self.evalType == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE + ): + if isinstance(self._returnType_placeholder, StructType): + try: + to_arrow_type(self._returnType_placeholder) + except TypeError: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": f"Invalid return type with grouped map Pandas UDFs or " + f"at groupby.applyInPandas(WithState): {self._returnType_placeholder}" + }, + ) + else: + raise PySparkTypeError( + error_class="INVALID_RETURN_TYPE_FOR_PANDAS_UDF", + message_parameters={ + "eval_type": "SQL_GROUPED_MAP_PANDAS_UDF or " + "SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE", + "return_type": str(self._returnType_placeholder), + }, + ) + elif ( + self.evalType == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF + or self.evalType == PythonEvalType.SQL_MAP_ARROW_ITER_UDF + ): + if isinstance(self._returnType_placeholder, StructType): + try: + to_arrow_type(self._returnType_placeholder) + except TypeError: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": f"Invalid return type in mapInPandas: " + f"{self._returnType_placeholder}" + }, + ) + else: + raise PySparkTypeError( + error_class="INVALID_RETURN_TYPE_FOR_PANDAS_UDF", + message_parameters={ + "eval_type": "SQL_MAP_PANDAS_ITER_UDF or SQL_MAP_ARROW_ITER_UDF", + "return_type": str(self._returnType_placeholder), + }, + ) + elif self.evalType == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: + if isinstance(self._returnType_placeholder, StructType): + try: + to_arrow_type(self._returnType_placeholder) + except TypeError: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": f"Invalid return type in cogroup.applyInPandas: " + f"{self._returnType_placeholder}" + }, + ) + else: + raise PySparkTypeError( + error_class="INVALID_RETURN_TYPE_FOR_PANDAS_UDF", + message_parameters={ + "eval_type": "SQL_COGROUPED_MAP_PANDAS_UDF", + "return_type": str(self._returnType_placeholder), + }, + ) + elif self.evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF: + try: + # StructType is not yet allowed as a return type, explicitly check here to fail fast + if isinstance(self._returnType_placeholder, StructType): + raise TypeError + to_arrow_type(self._returnType_placeholder) + except TypeError: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={ + "feature": f"Invalid return type with grouped aggregate Pandas UDFs: " + f"{self._returnType_placeholder}" + }, + ) + + return self._returnType_placeholder + + @property + def _judf(self) -> JavaObject: + # It is possible that concurrent access, to newly created UDF, + # will initialize multiple UserDefinedPythonFunctions. + # This is unlikely, doesn't affect correctness, + # and should have a minimal performance impact. + if self._judf_placeholder is None: + self._judf_placeholder = self._create_judf(self.func) + return self._judf_placeholder + + def _create_judf(self, func: Callable[..., Any]) -> JavaObject: + from pyspark.sql import SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + sc = spark.sparkContext + + wrapped_func = _wrap_function(sc, func, self.returnType) + jdt = spark._jsparkSession.parseDataType(self.returnType.json()) + assert sc._jvm is not None + judf = sc._jvm.org.apache.spark.sql.execution.python.UserDefinedPythonFunction( + self._name, wrapped_func, jdt, self.evalType, self.deterministic + ) + return judf + + def __call__(self, *cols: "ColumnOrName") -> Column: + sc = get_active_spark_context() + profiler: Optional[Profiler] = None + memory_profiler: Optional[Profiler] = None + if sc.profiler_collector: + profiler_enabled = sc._conf.get("spark.python.profile", "false") == "true" + memory_profiler_enabled = sc._conf.get("spark.python.profile.memory", "false") == "true" + + # Disable profiling Pandas UDFs with iterators as input/output. + if profiler_enabled or memory_profiler_enabled: + if self.evalType in [ + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_ARROW_ITER_UDF, + ]: + profiler_enabled = memory_profiler_enabled = False + warnings.warn( + "Profiling UDFs with iterators input/output is not supported.", + UserWarning, + ) + + # Disallow enabling two profilers at the same time. + if profiler_enabled and memory_profiler_enabled: + # When both profilers are enabled, they interfere with each other, + # that makes the result profile misleading. + raise RuntimeError( + "'spark.python.profile' and 'spark.python.profile.memory' configuration" + " cannot be enabled together." + ) + elif profiler_enabled: + f = self.func + profiler = sc.profiler_collector.new_udf_profiler(sc) + + @functools.wraps(f) + def func(*args: Any, **kwargs: Any) -> Any: + assert profiler is not None + return profiler.profile(f, *args, **kwargs) + + func.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + judf = self._create_judf(func) + jUDFExpr = judf.builder(_to_seq(sc, cols, _to_java_expr)) + jPythonUDF = judf.fromUDFExpr(jUDFExpr) + id = jUDFExpr.resultId().id() + sc.profiler_collector.add_profiler(id, profiler) + else: # memory_profiler_enabled + f = self.func + memory_profiler = sc.profiler_collector.new_memory_profiler(sc) + (sub_lines, start_line) = inspect.getsourcelines(f.__code__) + + @functools.wraps(f) + def func(*args: Any, **kwargs: Any) -> Any: + assert memory_profiler is not None + return memory_profiler.profile( + sub_lines, start_line, f, *args, **kwargs # type: ignore[arg-type] + ) + + func.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + judf = self._create_judf(func) + jUDFExpr = judf.builder(_to_seq(sc, cols, _to_java_expr)) + jPythonUDF = judf.fromUDFExpr(jUDFExpr) + id = jUDFExpr.resultId().id() + sc.profiler_collector.add_profiler(id, memory_profiler) + else: + judf = self._judf + jPythonUDF = judf.apply(_to_seq(sc, cols, _to_java_column)) + return Column(jPythonUDF) + + # This function is for improving the online help system in the interactive interpreter. + # For example, the built-in help / pydoc.help. It wraps the UDF with the docstring and + # argument annotation. (See: SPARK-19161) + def _wrapped(self) -> "UserDefinedFunctionLike": + """ + Wrap this udf with a function and attach docstring from func + """ + + # It is possible for a callable instance without __name__ attribute or/and + # __module__ attribute to be wrapped here. For example, functools.partial. In this case, + # we should avoid wrapping the attributes from the wrapped function to the wrapper + # function. So, we take out these attribute names from the default names to set and + # then manually assign it after being wrapped. + assignments = tuple( + a for a in functools.WRAPPER_ASSIGNMENTS if a != "__name__" and a != "__module__" + ) + + @functools.wraps(self.func, assigned=assignments) + def wrapper(*args: "ColumnOrName") -> Column: + return self(*args) + + wrapper.__name__ = self._name + wrapper.__module__ = ( + self.func.__module__ + if hasattr(self.func, "__module__") + else self.func.__class__.__module__ + ) + + wrapper.func = self.func # type: ignore[attr-defined] + wrapper.returnType = self.returnType # type: ignore[attr-defined] + wrapper.evalType = self.evalType # type: ignore[attr-defined] + wrapper.deterministic = self.deterministic # type: ignore[attr-defined] + wrapper.asNondeterministic = functools.wraps( # type: ignore[attr-defined] + self.asNondeterministic + )(lambda: self.asNondeterministic()._wrapped()) + wrapper._unwrapped = self # type: ignore[attr-defined] + return wrapper # type: ignore[return-value] + + def asNondeterministic(self) -> "UserDefinedFunction": + """ + Updates UserDefinedFunction to nondeterministic. + + .. versionadded:: 2.3 + """ + # Here, we explicitly clean the cache to create a JVM UDF instance + # with 'deterministic' updated. See SPARK-23233. + self._judf_placeholder = None + self.deterministic = False + return self + + +class UDFRegistration: + """ + Wrapper for user-defined function registration. This instance can be accessed by + :attr:`spark.udf` or :attr:`sqlContext.udf`. + + .. versionadded:: 1.3.1 + """ + + def __init__(self, sparkSession: "SparkSession"): + self.sparkSession = sparkSession + + def register( + self, + name: str, + f: Union[Callable[..., Any], "UserDefinedFunctionLike"], + returnType: Optional["DataTypeOrString"] = None, + ) -> "UserDefinedFunctionLike": + """Register a Python function (including lambda function) or a user-defined function + as a SQL function. + + .. versionadded:: 1.3.1 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str, + name of the user-defined function in SQL statements. + f : function, :meth:`pyspark.sql.functions.udf` or :meth:`pyspark.sql.functions.pandas_udf` + a Python function, or a user-defined function. The user-defined function can + be either row-at-a-time or vectorized. See :meth:`pyspark.sql.functions.udf` and + :meth:`pyspark.sql.functions.pandas_udf`. + returnType : :class:`pyspark.sql.types.DataType` or str, optional + the return type of the registered user-defined function. The value can + be either a :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + `returnType` can be optionally specified when `f` is a Python function but not + when `f` is a user-defined function. Please see the examples below. + + Returns + ------- + function + a user-defined function + + Notes + ----- + To register a nondeterministic Python function, users need to first build + a nondeterministic user-defined function for the Python function and then register it + as a SQL function. + + Examples + -------- + 1. When `f` is a Python function: + + `returnType` defaults to string type and can be optionally specified. The produced + object must match the specified type. In this case, this API works as if + `register(name, f, returnType=StringType())`. + + >>> strlen = spark.udf.register("stringLengthString", lambda x: len(x)) + >>> spark.sql("SELECT stringLengthString('test')").collect() + [Row(stringLengthString(test)='4')] + + >>> spark.sql("SELECT 'foo' AS text").select(strlen("text")).collect() + [Row(stringLengthString(text)='3')] + + >>> from pyspark.sql.types import IntegerType + >>> _ = spark.udf.register("stringLengthInt", lambda x: len(x), IntegerType()) + >>> spark.sql("SELECT stringLengthInt('test')").collect() + [Row(stringLengthInt(test)=4)] + + >>> from pyspark.sql.types import IntegerType + >>> _ = spark.udf.register("stringLengthInt", lambda x: len(x), IntegerType()) + >>> spark.sql("SELECT stringLengthInt('test')").collect() + [Row(stringLengthInt(test)=4)] + + 2. When `f` is a user-defined function (from Spark 2.3.0): + + Spark uses the return type of the given user-defined function as the return type of + the registered user-defined function. `returnType` should not be specified. + In this case, this API works as if `register(name, f)`. + + >>> from pyspark.sql.types import IntegerType + >>> from pyspark.sql.functions import udf + >>> slen = udf(lambda s: len(s), IntegerType()) + >>> _ = spark.udf.register("slen", slen) + >>> spark.sql("SELECT slen('test')").collect() + [Row(slen(test)=4)] + + >>> import random + >>> from pyspark.sql.functions import udf + >>> from pyspark.sql.types import IntegerType + >>> random_udf = udf(lambda: random.randint(0, 100), IntegerType()).asNondeterministic() + >>> new_random_udf = spark.udf.register("random_udf", random_udf) + >>> spark.sql("SELECT random_udf()").collect() # doctest: +SKIP + [Row(random_udf()=82)] + + >>> import pandas as pd # doctest: +SKIP + >>> from pyspark.sql.functions import pandas_udf + >>> @pandas_udf("integer") # doctest: +SKIP + ... def add_one(s: pd.Series) -> pd.Series: + ... return s + 1 + ... + >>> _ = spark.udf.register("add_one", add_one) # doctest: +SKIP + >>> spark.sql("SELECT add_one(id) FROM range(3)").collect() # doctest: +SKIP + [Row(add_one(id)=1), Row(add_one(id)=2), Row(add_one(id)=3)] + + >>> @pandas_udf("integer") # doctest: +SKIP + ... def sum_udf(v: pd.Series) -> int: + ... return v.sum() + ... + >>> _ = spark.udf.register("sum_udf", sum_udf) # doctest: +SKIP + >>> q = "SELECT sum_udf(v1) FROM VALUES (3, 0), (2, 0), (1, 1) tbl(v1, v2) GROUP BY v2" + >>> spark.sql(q).collect() # doctest: +SKIP + [Row(sum_udf(v1)=1), Row(sum_udf(v1)=5)] + + """ + + # This is to check whether the input function is from a user-defined function or + # Python function. + if hasattr(f, "asNondeterministic"): + if returnType is not None: + raise PySparkTypeError( + error_class="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF", + message_parameters={"arg_name": "f", "return_type": str(returnType)}, + ) + f = cast("UserDefinedFunctionLike", f) + if f.evalType not in [ + PythonEvalType.SQL_BATCHED_UDF, + PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + ]: + raise PySparkTypeError( + error_class="INVALID_UDF_EVAL_TYPE", + message_parameters={ + "eval_type": "SQL_BATCHED_UDF, SQL_ARROW_BATCHED_UDF, " + "SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_PANDAS_ITER_UDF or " + "SQL_GROUPED_AGG_PANDAS_UDF" + }, + ) + source_udf = _create_udf( + f.func, + returnType=f.returnType, + name=name, + evalType=f.evalType, + deterministic=f.deterministic, + ) + register_udf = source_udf._unwrapped # type: ignore[attr-defined] + return_udf = register_udf + else: + if returnType is None: + returnType = StringType() + return_udf = _create_udf( + f, returnType=returnType, evalType=PythonEvalType.SQL_BATCHED_UDF, name=name + ) + register_udf = return_udf._unwrapped + self.sparkSession._jsparkSession.udf().registerPython(name, register_udf._judf) + return return_udf + + def registerJavaFunction( + self, + name: str, + javaClassName: str, + returnType: Optional["DataTypeOrString"] = None, + ) -> None: + """Register a Java user-defined function as a SQL function. + + In addition to a name and the function itself, the return type can be optionally specified. + When the return type is not specified we would infer it via reflection. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Parameters + ---------- + name : str + name of the user-defined function + javaClassName : str + fully qualified name of java class + returnType : :class:`pyspark.sql.types.DataType` or str, optional + the return type of the registered Java function. The value can be either + a :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + + Examples + -------- + >>> from pyspark.sql.types import IntegerType + >>> spark.udf.registerJavaFunction( + ... "javaStringLength", "test.org.apache.spark.sql.JavaStringLength", IntegerType()) + ... # doctest: +SKIP + >>> spark.sql("SELECT javaStringLength('test')").collect() # doctest: +SKIP + [Row(javaStringLength(test)=4)] + + >>> spark.udf.registerJavaFunction( + ... "javaStringLength2", "test.org.apache.spark.sql.JavaStringLength") + ... # doctest: +SKIP + >>> spark.sql("SELECT javaStringLength2('test')").collect() # doctest: +SKIP + [Row(javaStringLength2(test)=4)] + + >>> spark.udf.registerJavaFunction( + ... "javaStringLength3", "test.org.apache.spark.sql.JavaStringLength", "integer") + ... # doctest: +SKIP + >>> spark.sql("SELECT javaStringLength3('test')").collect() # doctest: +SKIP + [Row(javaStringLength3(test)=4)] + """ + + jdt = None + if returnType is not None: + if not isinstance(returnType, DataType): + returnType = _parse_datatype_string(returnType) + jdt = self.sparkSession._jsparkSession.parseDataType(returnType.json()) + self.sparkSession._jsparkSession.udf().registerJava(name, javaClassName, jdt) + + def registerJavaUDAF(self, name: str, javaClassName: str) -> None: + """Register a Java user-defined aggregate function as a SQL function. + + .. versionadded:: 2.3.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + name : str + name of the user-defined aggregate function + javaClassName : str + fully qualified name of java class + + Examples + -------- + >>> spark.udf.registerJavaUDAF("javaUDAF", "test.org.apache.spark.sql.MyDoubleAvg") + ... # doctest: +SKIP + >>> df = spark.createDataFrame([(1, "a"),(2, "b"), (3, "a")],["id", "name"]) + >>> df.createOrReplaceTempView("df") + >>> q = "SELECT name, javaUDAF(id) as avg from df group by name order by name desc" + >>> spark.sql(q).collect() # doctest: +SKIP + [Row(name='b', avg=102.0), Row(name='a', avg=102.0)] + """ + + self.sparkSession._jsparkSession.udf().registerJavaUDAF(name, javaClassName) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.udf + + globs = pyspark.sql.udf.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.udf tests").getOrCreate() + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.udf, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udtf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udtf.py new file mode 100644 index 0000000..7e05658 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/udtf.py @@ -0,0 +1,325 @@ +# +# 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. +# +""" +User-defined table function related classes and functions +""" +import pickle +import sys +import warnings +from typing import Any, Type, TYPE_CHECKING, Optional, Union + +from py4j.java_gateway import JavaObject + +from pyspark.errors import PySparkAttributeError, PySparkRuntimeError, PySparkTypeError +from pyspark.rdd import PythonEvalType +from pyspark.sql.column import _to_java_column, _to_seq +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version +from pyspark.sql.types import StructType, _parse_datatype_string +from pyspark.sql.udf import _wrap_function + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.session import SparkSession + +__all__ = ["UDTFRegistration"] + + +def _create_udtf( + cls: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_TABLE_UDF, + deterministic: bool = False, +) -> "UserDefinedTableFunction": + """Create a Python UDTF with the given eval type.""" + udtf_obj = UserDefinedTableFunction( + cls, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + ) + + return udtf_obj + + +def _create_py_udtf( + cls: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + deterministic: bool = False, + useArrow: Optional[bool] = None, +) -> "UserDefinedTableFunction": + """Create a regular or an Arrow-optimized Python UDTF.""" + # Determine whether to create Arrow-optimized UDTFs. + if useArrow is not None: + arrow_enabled = useArrow + else: + from pyspark.sql import SparkSession + + session = SparkSession._instantiatedSession + arrow_enabled = False + if session is not None: + value = session.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled") + if isinstance(value, str) and value.lower() == "true": + arrow_enabled = True + + eval_type: int = PythonEvalType.SQL_TABLE_UDF + + if arrow_enabled: + # Return the regular UDTF if the required dependencies are not satisfied. + try: + require_minimum_pandas_version() + require_minimum_pyarrow_version() + eval_type = PythonEvalType.SQL_ARROW_TABLE_UDF + except ImportError as e: + warnings.warn( + f"Arrow optimization for Python UDTFs cannot be enabled: {str(e)}. " + f"Falling back to using regular Python UDTFs.", + UserWarning, + ) + + return _create_udtf( + cls=cls, + returnType=returnType, + name=name, + evalType=eval_type, + deterministic=deterministic, + ) + + +def _validate_udtf_handler(cls: Any) -> None: + """Validate the handler class of a UDTF.""" + + if not isinstance(cls, type): + raise PySparkTypeError( + error_class="INVALID_UDTF_HANDLER_TYPE", message_parameters={"type": type(cls).__name__} + ) + + if not hasattr(cls, "eval"): + raise PySparkAttributeError( + error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__} + ) + + +class UserDefinedTableFunction: + """ + User-defined table function in Python + + .. versionadded:: 3.5.0 + + Notes + ----- + The constructor of this class is not supposed to be directly called. + Use :meth:`pyspark.sql.functions.udtf` to create this instance. + + This API is evolving. + """ + + def __init__( + self, + func: Type, + returnType: Union[StructType, str], + name: Optional[str] = None, + evalType: int = PythonEvalType.SQL_TABLE_UDF, + deterministic: bool = False, + ): + _validate_udtf_handler(func) + + self.func = func + self._returnType = returnType + self._returnType_placeholder: Optional[StructType] = None + self._inputTypes_placeholder = None + self._judtf_placeholder = None + self._name = name or func.__name__ + self.evalType = evalType + self.deterministic = deterministic + + @property + def returnType(self) -> StructType: + # `_parse_datatype_string` accesses to JVM for parsing a DDL formatted string. + # This makes sure this is called after SparkContext is initialized. + if self._returnType_placeholder is None: + if isinstance(self._returnType, str): + parsed = _parse_datatype_string(self._returnType) + else: + parsed = self._returnType + if not isinstance(parsed, StructType): + raise PySparkTypeError( + error_class="UDTF_RETURN_TYPE_MISMATCH", + message_parameters={ + "name": self._name, + "return_type": f"{parsed}", + }, + ) + self._returnType_placeholder = parsed + return self._returnType_placeholder + + @property + def _judtf(self) -> JavaObject: + if self._judtf_placeholder is None: + self._judtf_placeholder = self._create_judtf(self.func) + return self._judtf_placeholder + + def _create_judtf(self, func: Type) -> JavaObject: + from pyspark.sql import SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + sc = spark.sparkContext + + try: + wrapped_func = _wrap_function(sc, func) + except pickle.PicklingError as e: + if "CONTEXT_ONLY_VALID_ON_DRIVER" in str(e): + raise PySparkRuntimeError( + error_class="UDTF_SERIALIZATION_ERROR", + message_parameters={ + "name": self._name, + "message": "it appears that you are attempting to reference SparkSession " + "inside a UDTF. SparkSession can only be used on the driver, " + "not in code that runs on workers. Please remove the reference " + "and try again.", + }, + ) from None + raise PySparkRuntimeError( + error_class="UDTF_SERIALIZATION_ERROR", + message_parameters={ + "name": self._name, + "message": "Please check the stack trace and make sure the " + "function is serializable.", + }, + ) + + jdt = spark._jsparkSession.parseDataType(self.returnType.json()) + assert sc._jvm is not None + judtf = sc._jvm.org.apache.spark.sql.execution.python.UserDefinedPythonTableFunction( + self._name, wrapped_func, jdt, self.evalType, self.deterministic + ) + return judtf + + def __call__(self, *cols: "ColumnOrName") -> "DataFrame": + from pyspark.sql import DataFrame, SparkSession + + spark = SparkSession._getActiveSessionOrCreate() + sc = spark.sparkContext + + judtf = self._judtf + jPythonUDTF = judtf.apply(spark._jsparkSession, _to_seq(sc, cols, _to_java_column)) + return DataFrame(jPythonUDTF, spark) + + def asDeterministic(self) -> "UserDefinedTableFunction": + """ + Updates UserDefinedTableFunction to deterministic. + """ + # Explicitly clean the cache to create a JVM UDTF instance. + self._judtf_placeholder = None + self.deterministic = True + return self + + +class UDTFRegistration: + """ + Wrapper for user-defined table function registration. This instance can be accessed by + :attr:`spark.udtf` or :attr:`sqlContext.udtf`. + + .. versionadded:: 3.5.0 + """ + + def __init__(self, sparkSession: "SparkSession"): + self.sparkSession = sparkSession + + def register( + self, + name: str, + f: "UserDefinedTableFunction", + ) -> "UserDefinedTableFunction": + """Register a Python user-defined table function as a SQL table function. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + name : str + The name of the user-defined table function in SQL statements. + f : function or :meth:`pyspark.sql.functions.udtf` + The user-defined table function. + + Returns + ------- + function + The registered user-defined table function. + + Notes + ----- + Spark uses the return type of the given user-defined table function as the return + type of the registered user-defined function. + + To register a nondeterministic Python table function, users need to first build + a nondeterministic user-defined table function and then register it as a SQL function. + + Examples + -------- + >>> from pyspark.sql.functions import udtf + >>> @udtf(returnType="c1: int, c2: int") + ... class PlusOne: + ... def eval(self, x: int): + ... yield x, x + 1 + ... + >>> _ = spark.udtf.register(name="plus_one", f=PlusOne) + >>> spark.sql("SELECT * FROM plus_one(1)").collect() + [Row(c1=1, c2=2)] + + Use it with lateral join + + >>> spark.sql("SELECT * FROM VALUES (0, 1), (1, 2) t(x, y), LATERAL plus_one(x)").collect() + [Row(x=0, y=1, c1=0, c2=1), Row(x=1, y=2, c1=1, c2=2)] + """ + if f.evalType not in [PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF]: + raise PySparkTypeError( + error_class="INVALID_UDTF_EVAL_TYPE", + message_parameters={ + "name": name, + "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF", + }, + ) + + register_udtf = _create_udtf( + cls=f.func, + returnType=f.returnType, + name=name, + evalType=f.evalType, + deterministic=f.deterministic, + ) + self.sparkSession._jsparkSession.udtf().registerPython(name, register_udtf._judtf) + return register_udtf + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.udf + + globs = pyspark.sql.udtf.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.udtf tests").getOrCreate() + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.sql.udtf, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/utils.py new file mode 100644 index 0000000..02b67d9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/utils.py @@ -0,0 +1,339 @@ +# +# 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 inspect +import functools +import os +from typing import Any, Callable, Optional, Sequence, TYPE_CHECKING, cast, TypeVar, Union, Type + +from py4j.java_collections import JavaArray +from py4j.java_gateway import ( + JavaClass, + JavaGateway, + JavaObject, +) + +from pyspark import SparkContext + +# For backward compatibility. +from pyspark.errors import ( # noqa: F401 + AnalysisException, + ParseException, + IllegalArgumentException, + StreamingQueryException, + QueryExecutionException, + PythonException, + UnknownException, + SparkUpgradeException, + PySparkNotImplementedError, +) +from pyspark.errors.exceptions.captured import CapturedException # noqa: F401 +from pyspark.find_spark_home import _find_spark_home + +if TYPE_CHECKING: + from pyspark.sql.session import SparkSession + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.column import Column + from pyspark.sql.window import Window + from pyspark.pandas._typing import IndexOpsLike, SeriesOrIndex + +has_numpy = False +try: + import numpy as np # noqa: F401 + + has_numpy = True +except ImportError: + pass + + +FuncT = TypeVar("FuncT", bound=Callable[..., Any]) + + +def toJArray(gateway: JavaGateway, jtype: JavaClass, arr: Sequence[Any]) -> JavaArray: + """ + Convert python list to java type array + + Parameters + ---------- + gateway : + Py4j Gateway + jtype : + java type of element in array + arr : + python type list + """ + jarray: JavaArray = gateway.new_array(jtype, len(arr)) + for i in range(0, len(arr)): + jarray[i] = arr[i] + return jarray + + +def require_test_compiled() -> None: + """Raise Exception if test classes are not compiled""" + import os + import glob + + test_class_path = os.path.join(_find_spark_home(), "sql", "core", "target", "*", "test-classes") + paths = glob.glob(test_class_path) + + if len(paths) == 0: + raise RuntimeError( + "%s doesn't exist. Spark sql test classes are not compiled." % test_class_path + ) + + +class ForeachBatchFunction: + """ + This is the Python implementation of Java interface 'ForeachBatchFunction'. This wraps + the user-defined 'foreachBatch' function such that it can be called from the JVM when + the query is active. + """ + + def __init__(self, session: "SparkSession", func: Callable[["DataFrame", int], None]): + self.func = func + self.session = session + + def call(self, jdf: JavaObject, batch_id: int) -> None: + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.session import SparkSession + + try: + session_jdf = jdf.sparkSession() + # assuming that spark context is still the same between JVM and PySpark + wrapped_session_jdf = SparkSession(self.session.sparkContext, session_jdf) + self.func(DataFrame(jdf, wrapped_session_jdf), batch_id) + except Exception as e: + self.error = e + raise e + + class Java: + implements = ["org.apache.spark.sql.execution.streaming.sources.PythonForeachBatchFunction"] + + +def to_str(value: Any) -> Optional[str]: + """ + A wrapper over str(), but converts bool values to lower case strings. + If None is given, just returns None, instead of converting it to string "None". + """ + if isinstance(value, bool): + return str(value).lower() + elif value is None: + return value + else: + return str(value) + + +def is_timestamp_ntz_preferred() -> bool: + """ + Return a bool if TimestampNTZType is preferred according to the SQL configuration set. + """ + if is_remote(): + from pyspark.sql.connect.session import SparkSession as ConnectSparkSession + + session = ConnectSparkSession.getActiveSession() + if session is None: + return False + else: + return session.conf.get("spark.sql.timestampType", None) == "TIMESTAMP_NTZ" + else: + jvm = SparkContext._jvm + return jvm is not None and jvm.PythonSQLUtils.isTimestampNTZPreferred() + + +def is_remote() -> bool: + """ + Returns if the current running environment is for Spark Connect. + """ + return "SPARK_CONNECT_MODE_ENABLED" in os.environ + + +def try_remote_functions(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect import functions + + return getattr(functions, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def try_remote_avro_functions(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect.avro import functions + + return getattr(functions, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def try_remote_protobuf_functions(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect.protobuf import functions + + return getattr(functions, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def try_remote_window(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect.window import Window # type: ignore[misc] + + return getattr(Window, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def try_remote_windowspec(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect.window import WindowSpec + + return getattr(WindowSpec, f.__name__)(*args, **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def get_active_spark_context() -> SparkContext: + """Raise RuntimeError if SparkContext is not initialized, + otherwise, returns the active SparkContext.""" + sc = SparkContext._active_spark_context + if sc is None or sc._jvm is None: + raise RuntimeError("SparkContext or SparkSession should be created first.") + return sc + + +def try_remote_observation(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + # TODO(SPARK-41527): Add the support of Observation. + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + raise PySparkNotImplementedError( + error_class="NOT_IMPLEMENTED", + message_parameters={"feature": "Observation support for Spark Connect"}, + ) + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def try_remote_session_classmethod(f: FuncT) -> FuncT: + """Mark API supported from Spark Connect.""" + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + + if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: + from pyspark.sql.connect.session import SparkSession # type: ignore[misc] + + assert inspect.isclass(args[0]) + return getattr(SparkSession, f.__name__)(*args[1:], **kwargs) + else: + return f(*args, **kwargs) + + return cast(FuncT, wrapped) + + +def pyspark_column_op( + func_name: str, left: "IndexOpsLike", right: Any, fillna: Any = None +) -> Union["SeriesOrIndex", None]: + """ + Wrapper function for column_op to get proper Column class. + """ + from pyspark.pandas.base import column_op + from pyspark.sql.column import Column as PySparkColumn + from pyspark.pandas.data_type_ops.base import _is_extension_dtypes + + if is_remote(): + from pyspark.sql.connect.column import Column as ConnectColumn + + Column = ConnectColumn + else: + Column = PySparkColumn # type: ignore[assignment] + result = column_op(getattr(Column, func_name))(left, right) + # It works as expected on extension dtype, so we don't need to call `fillna` for this case. + if (fillna is not None) and (_is_extension_dtypes(left) or _is_extension_dtypes(right)): + fillna = None + # TODO(SPARK-43877): Fix behavior difference for compare binary functions. + return result.fillna(fillna) if fillna is not None else result + + +def get_column_class() -> Type["Column"]: + from pyspark.sql.column import Column as PySparkColumn + + if is_remote(): + from pyspark.sql.connect.column import Column as ConnectColumn + + return ConnectColumn # type: ignore[return-value] + else: + return PySparkColumn + + +def get_dataframe_class() -> Type["DataFrame"]: + from pyspark.sql.dataframe import DataFrame as PySparkDataFrame + + if is_remote(): + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + + return ConnectDataFrame # type: ignore[return-value] + else: + return PySparkDataFrame + + +def get_window_class() -> Type["Window"]: + from pyspark.sql.window import Window as PySparkWindow + + if is_remote(): + from pyspark.sql.connect.window import Window as ConnectWindow + + return ConnectWindow # type: ignore[return-value] + else: + return PySparkWindow diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/sql/window.py b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/window.py new file mode 100644 index 0000000..ca05cb0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/sql/window.py @@ -0,0 +1,492 @@ +# +# 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 sys +from typing import cast, Iterable, List, Tuple, TYPE_CHECKING, Union + +from py4j.java_gateway import JavaObject, JVMView + +from pyspark.sql.column import _to_seq, _to_java_column +from pyspark.sql.utils import ( + try_remote_window, + try_remote_windowspec, + get_active_spark_context, +) + +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName, ColumnOrName_ + +__all__ = ["Window", "WindowSpec"] + + +def _to_java_cols(cols: Tuple[Union["ColumnOrName", List["ColumnOrName_"]], ...]) -> JavaObject: + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] # type: ignore[assignment] + sc = get_active_spark_context() + return _to_seq(sc, cast(Iterable["ColumnOrName"], cols), _to_java_column) + + +class Window: + """ + Utility functions for defining window in DataFrames. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + + Notes + ----- + When ordering is not defined, an unbounded window frame (rowFrame, + unboundedPreceding, unboundedFollowing) is used by default. When ordering is defined, + a growing window frame (rangeFrame, unboundedPreceding, currentRow) is used by default. + + Examples + -------- + >>> # ORDER BY date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + >>> window = Window.orderBy("date").rowsBetween(Window.unboundedPreceding, Window.currentRow) + + >>> # PARTITION BY country ORDER BY date RANGE BETWEEN 3 PRECEDING AND 3 FOLLOWING + >>> window = Window.orderBy("date").partitionBy("country").rangeBetween(-3, 3) + """ + + _JAVA_MIN_LONG = -(1 << 63) # -9223372036854775808 + _JAVA_MAX_LONG = (1 << 63) - 1 # 9223372036854775807 + _PRECEDING_THRESHOLD = max(-sys.maxsize, _JAVA_MIN_LONG) + _FOLLOWING_THRESHOLD = min(sys.maxsize, _JAVA_MAX_LONG) + + unboundedPreceding: int = _JAVA_MIN_LONG + + unboundedFollowing: int = _JAVA_MAX_LONG + + currentRow: int = 0 + + @staticmethod + @try_remote_window + def partitionBy(*cols: Union["ColumnOrName", List["ColumnOrName_"]]) -> "WindowSpec": + """ + Creates a :class:`WindowSpec` with the partitioning defined. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + cols : str, :class:`Column` or list + names of columns or expressions + + Returns + ------- + :class: `WindowSpec` + A :class:`WindowSpec` with the partitioning defined. + + Examples + -------- + >>> from pyspark.sql import Window + >>> from pyspark.sql.functions import row_number + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Show row number order by ``id`` in partition ``category``. + + >>> window = Window.partitionBy("category").orderBy("id") + >>> df.withColumn("row_number", row_number().over(window)).show() + +---+--------+----------+ + | id|category|row_number| + +---+--------+----------+ + | 1| a| 1| + | 1| a| 2| + | 2| a| 3| + | 1| b| 1| + | 2| b| 2| + | 3| b| 3| + +---+--------+----------+ + """ + sc = get_active_spark_context() + jspec = cast(JVMView, sc._jvm).org.apache.spark.sql.expressions.Window.partitionBy( + _to_java_cols(cols) + ) + return WindowSpec(jspec) + + @staticmethod + @try_remote_window + def orderBy(*cols: Union["ColumnOrName", List["ColumnOrName_"]]) -> "WindowSpec": + """ + Creates a :class:`WindowSpec` with the ordering defined. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + cols : str, :class:`Column` or list + names of columns or expressions + + Returns + ------- + :class: `WindowSpec` + A :class:`WindowSpec` with the ordering defined. + + Examples + -------- + >>> from pyspark.sql import Window + >>> from pyspark.sql.functions import row_number + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Show row number order by ``category`` in partition ``id``. + + >>> window = Window.partitionBy("id").orderBy("category") + >>> df.withColumn("row_number", row_number().over(window)).show() + +---+--------+----------+ + | id|category|row_number| + +---+--------+----------+ + | 1| a| 1| + | 1| a| 2| + | 1| b| 3| + | 2| a| 1| + | 2| b| 2| + | 3| b| 1| + +---+--------+----------+ + """ + sc = get_active_spark_context() + jspec = cast(JVMView, sc._jvm).org.apache.spark.sql.expressions.Window.orderBy( + _to_java_cols(cols) + ) + return WindowSpec(jspec) + + @staticmethod + @try_remote_window + def rowsBetween(start: int, end: int) -> "WindowSpec": + """ + Creates a :class:`WindowSpec` with the frame boundaries defined, + from `start` (inclusive) to `end` (inclusive). + + Both `start` and `end` are relative positions from the current row. + For example, "0" means "current row", while "-1" means the row before + the current row, and "5" means the fifth row after the current row. + + We recommend users use ``Window.unboundedPreceding``, ``Window.unboundedFollowing``, + and ``Window.currentRow`` to specify special boundary values, rather than using integral + values directly. + + A row based boundary is based on the position of the row within the partition. + An offset indicates the number of rows above or below the current row, the frame for the + current row starts or ends. For instance, given a row based sliding frame with a lower bound + offset of -1 and a upper bound offset of +2. The frame for row with index 5 would range from + index 4 to index 7. + + .. versionadded:: 2.1.0 + + Parameters + ---------- + start : int + boundary start, inclusive. + The frame is unbounded if this is ``Window.unboundedPreceding``, or + any value less than or equal to -9223372036854775808. + end : int + boundary end, inclusive. + The frame is unbounded if this is ``Window.unboundedFollowing``, or + any value greater than or equal to 9223372036854775807. + + Returns + ------- + :class: `WindowSpec` + A :class:`WindowSpec` with the frame boundaries defined, + from `start` (inclusive) to `end` (inclusive). + + Examples + -------- + >>> from pyspark.sql import Window + >>> from pyspark.sql import functions as func + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Calculate sum of ``id`` in the range from currentRow to currentRow + 1 + in partition ``category`` + + >>> window = Window.partitionBy("category").orderBy("id").rowsBetween(Window.currentRow, 1) + >>> df.withColumn("sum", func.sum("id").over(window)).sort("id", "category", "sum").show() + +---+--------+---+ + | id|category|sum| + +---+--------+---+ + | 1| a| 2| + | 1| a| 3| + | 1| b| 3| + | 2| a| 2| + | 2| b| 5| + | 3| b| 3| + +---+--------+---+ + + """ + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + sc = get_active_spark_context() + jspec = cast(JVMView, sc._jvm).org.apache.spark.sql.expressions.Window.rowsBetween( + start, end + ) + return WindowSpec(jspec) + + @staticmethod + @try_remote_window + def rangeBetween(start: int, end: int) -> "WindowSpec": + """ + Creates a :class:`WindowSpec` with the frame boundaries defined, + from `start` (inclusive) to `end` (inclusive). + + Both `start` and `end` are relative from the current row. For example, + "0" means "current row", while "-1" means one off before the current row, + and "5" means the five off after the current row. + + We recommend users use ``Window.unboundedPreceding``, ``Window.unboundedFollowing``, + and ``Window.currentRow`` to specify special boundary values, rather than using integral + values directly. + + A range-based boundary is based on the actual value of the ORDER BY + expression(s). An offset is used to alter the value of the ORDER BY expression, for + instance if the current ORDER BY expression has a value of 10 and the lower bound offset + is -3, the resulting lower bound for the current row will be 10 - 3 = 7. This however puts a + number of constraints on the ORDER BY expressions: there can be only one expression and this + expression must have a numerical data type. An exception can be made when the offset is + unbounded, because no value modification is needed, in this case multiple and non-numeric + ORDER BY expression are allowed. + + .. versionadded:: 2.1.0 + + Parameters + ---------- + start : int + boundary start, inclusive. + The frame is unbounded if this is ``Window.unboundedPreceding``, or + any value less than or equal to max(-sys.maxsize, -9223372036854775808). + end : int + boundary end, inclusive. + The frame is unbounded if this is ``Window.unboundedFollowing``, or + any value greater than or equal to min(sys.maxsize, 9223372036854775807). + + Returns + ------- + :class: `WindowSpec` + A :class:`WindowSpec` with the frame boundaries defined, + from `start` (inclusive) to `end` (inclusive). + + Examples + -------- + >>> from pyspark.sql import Window + >>> from pyspark.sql import functions as func + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Calculate sum of ``id`` in the range from ``id`` of currentRow to ``id`` of currentRow + 1 + in partition ``category`` + + >>> window = Window.partitionBy("category").orderBy("id").rangeBetween(Window.currentRow, 1) + >>> df.withColumn("sum", func.sum("id").over(window)).sort("id", "category").show() + +---+--------+---+ + | id|category|sum| + +---+--------+---+ + | 1| a| 4| + | 1| a| 4| + | 1| b| 3| + | 2| a| 2| + | 2| b| 5| + | 3| b| 3| + +---+--------+---+ + + """ + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + sc = get_active_spark_context() + jspec = cast(JVMView, sc._jvm).org.apache.spark.sql.expressions.Window.rangeBetween( + start, end + ) + return WindowSpec(jspec) + + +class WindowSpec: + """ + A window specification that defines the partitioning, ordering, + and frame boundaries. + + Use the static methods in :class:`Window` to create a :class:`WindowSpec`. + + .. versionadded:: 1.4.0 + + .. versionchanged:: 3.4.0 + Supports Spark Connect. + """ + + def __init__(self, jspec: JavaObject) -> None: + self._jspec = jspec + + @try_remote_windowspec + def partitionBy(self, *cols: Union["ColumnOrName", List["ColumnOrName_"]]) -> "WindowSpec": + """ + Defines the partitioning columns in a :class:`WindowSpec`. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + cols : str, :class:`Column` or list + names of columns or expressions + """ + return WindowSpec(self._jspec.partitionBy(_to_java_cols(cols))) + + @try_remote_windowspec + def orderBy(self, *cols: Union["ColumnOrName", List["ColumnOrName_"]]) -> "WindowSpec": + """ + Defines the ordering columns in a :class:`WindowSpec`. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + cols : str, :class:`Column` or list + names of columns or expressions + """ + return WindowSpec(self._jspec.orderBy(_to_java_cols(cols))) + + @try_remote_windowspec + def rowsBetween(self, start: int, end: int) -> "WindowSpec": + """ + Defines the frame boundaries, from `start` (inclusive) to `end` (inclusive). + + Both `start` and `end` are relative positions from the current row. + For example, "0" means "current row", while "-1" means the row before + the current row, and "5" means the fifth row after the current row. + + We recommend users use ``Window.unboundedPreceding``, ``Window.unboundedFollowing``, + and ``Window.currentRow`` to specify special boundary values, rather than using integral + values directly. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + start : int + boundary start, inclusive. + The frame is unbounded if this is ``Window.unboundedPreceding``, or + any value less than or equal to max(-sys.maxsize, -9223372036854775808). + end : int + boundary end, inclusive. + The frame is unbounded if this is ``Window.unboundedFollowing``, or + any value greater than or equal to min(sys.maxsize, 9223372036854775807). + """ + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + return WindowSpec(self._jspec.rowsBetween(start, end)) + + @try_remote_windowspec + def rangeBetween(self, start: int, end: int) -> "WindowSpec": + """ + Defines the frame boundaries, from `start` (inclusive) to `end` (inclusive). + + Both `start` and `end` are relative from the current row. For example, + "0" means "current row", while "-1" means one off before the current row, + and "5" means the five off after the current row. + + We recommend users use ``Window.unboundedPreceding``, ``Window.unboundedFollowing``, + and ``Window.currentRow`` to specify special boundary values, rather than using integral + values directly. + + .. versionadded:: 1.4.0 + + Parameters + ---------- + start : int + boundary start, inclusive. + The frame is unbounded if this is ``Window.unboundedPreceding``, or + any value less than or equal to max(-sys.maxsize, -9223372036854775808). + end : int + boundary end, inclusive. + The frame is unbounded if this is ``Window.unboundedFollowing``, or + any value greater than or equal to min(sys.maxsize, 9223372036854775807). + """ + if start <= Window._PRECEDING_THRESHOLD: + start = Window.unboundedPreceding + if end >= Window._FOLLOWING_THRESHOLD: + end = Window.unboundedFollowing + return WindowSpec(self._jspec.rangeBetween(start, end)) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.sql.window + + globs = pyspark.sql.window.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("sql.window tests").getOrCreate() + globs["spark"] = spark + + (failure_count, test_count) = doctest.testmod( + pyspark.sql.window, globs=globs, optionflags=doctest.NORMALIZE_WHITESPACE + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/statcounter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/statcounter.py new file mode 100644 index 0000000..758d0a6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/statcounter.py @@ -0,0 +1,165 @@ +# +# 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. +# + +# This file is ported from spark/util/StatCounter.scala + +import copy +import math +from typing import Dict, Iterable, Optional + +try: + from numpy import maximum, minimum, sqrt +except ImportError: + maximum = max # type: ignore[assignment] + minimum = min # type: ignore[assignment] + sqrt = math.sqrt # type: ignore[assignment] + + +class StatCounter: + def __init__(self, values: Optional[Iterable[float]] = None): + if values is None: + values = list() + self.n = 0 # Running count of our values + self.mu = 0.0 # Running mean of our values + self.m2 = 0.0 # Running variance numerator (sum of (x - mean)^2) + self.maxValue = float("-inf") + self.minValue = float("inf") + + for v in values: + self.merge(v) + + # Add a value into this StatCounter, updating the internal statistics. + def merge(self, value: float) -> "StatCounter": + delta = value - self.mu + self.n += 1 + self.mu += delta / self.n + self.m2 += delta * (value - self.mu) + self.maxValue = maximum(self.maxValue, value) + self.minValue = minimum(self.minValue, value) + + return self + + # Merge another StatCounter into this one, adding up the internal statistics. + def mergeStats(self, other: "StatCounter") -> "StatCounter": + if not isinstance(other, StatCounter): + raise TypeError("Can only merge StatCounter but got %s" % type(other)) + + if other is self: # reference equality holds + self.mergeStats(other.copy()) # Avoid overwriting fields in a weird order + else: + if self.n == 0: + self.mu = other.mu + self.m2 = other.m2 + self.n = other.n + self.maxValue = other.maxValue + self.minValue = other.minValue + + elif other.n != 0: + delta = other.mu - self.mu + if other.n * 10 < self.n: + self.mu = self.mu + (delta * other.n) / (self.n + other.n) + elif self.n * 10 < other.n: + self.mu = other.mu - (delta * self.n) / (self.n + other.n) + else: + self.mu = (self.mu * self.n + other.mu * other.n) / (self.n + other.n) + + self.maxValue = maximum(self.maxValue, other.maxValue) + self.minValue = minimum(self.minValue, other.minValue) + + self.m2 += other.m2 + (delta * delta * self.n * other.n) / (self.n + other.n) + self.n += other.n + return self + + # Clone this StatCounter + def copy(self) -> "StatCounter": + return copy.deepcopy(self) + + def count(self) -> int: + return int(self.n) + + def mean(self) -> float: + return self.mu + + def sum(self) -> float: + return self.n * self.mu + + def min(self) -> float: + return self.minValue + + def max(self) -> float: + return self.maxValue + + # Return the variance of the values. + def variance(self) -> float: + if self.n == 0: + return float("nan") + else: + return self.m2 / self.n + + # + # Return the sample variance, which corrects for bias in estimating the variance by dividing + # by N-1 instead of N. + # + def sampleVariance(self) -> float: + if self.n <= 1: + return float("nan") + else: + return self.m2 / (self.n - 1) + + # Return the standard deviation of the values. + def stdev(self) -> float: + return sqrt(self.variance()) + + # + # Return the sample standard deviation of the values, which corrects for bias in estimating the + # variance by dividing by N-1 instead of N. + # + def sampleStdev(self) -> float: + return sqrt(self.sampleVariance()) + + def asDict(self, sample: bool = False) -> Dict[str, float]: + """Returns the :class:`StatCounter` members as a ``dict``. + + Examples + -------- + >>> sc.parallelize([1., 2., 3., 4.]).stats().asDict() + {'count': 4L, + 'max': 4.0, + 'mean': 2.5, + 'min': 1.0, + 'stdev': 1.2909944487358056, + 'sum': 10.0, + 'variance': 1.6666666666666667} + """ + return { + "count": self.count(), + "mean": self.mean(), + "sum": self.sum(), + "min": self.min(), + "max": self.max(), + "stdev": self.stdev() if sample else self.sampleStdev(), + "variance": self.variance() if sample else self.sampleVariance(), + } + + def __repr__(self) -> str: + return "(count: %s, mean: %s, stdev: %s, max: %s, min: %s)" % ( + self.count(), + self.mean(), + self.stdev(), + self.max(), + self.min(), + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/status.py b/spark-3.5.0-bin-hadoop3/python/pyspark/status.py new file mode 100644 index 0000000..7e64c41 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/status.py @@ -0,0 +1,112 @@ +# +# 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. +# + +__all__ = ["SparkJobInfo", "SparkStageInfo", "StatusTracker"] + +from typing import List, NamedTuple, Optional + +from py4j.java_collections import JavaArray +from py4j.java_gateway import JavaObject + + +class SparkJobInfo(NamedTuple): + """ + Exposes information about Spark Jobs. + """ + + jobId: int + stageIds: JavaArray + status: str + + +class SparkStageInfo(NamedTuple): + """ + Exposes information about Spark Stages. + """ + + stageId: int + currentAttemptId: int + name: str + numTasks: int + numActiveTasks: int + numCompletedTasks: int + numFailedTasks: int + + +class StatusTracker: + """ + Low-level status reporting APIs for monitoring job and stage progress. + + These APIs intentionally provide very weak consistency semantics; + consumers of these APIs should be prepared to handle empty / missing + information. For example, a job's stage ids may be known but the status + API may not have any information about the details of those stages, so + `getStageInfo` could potentially return `None` for a valid stage id. + + To limit memory usage, these APIs only provide information on recent + jobs / stages. These APIs will provide information for the last + `spark.ui.retainedStages` stages and `spark.ui.retainedJobs` jobs. + """ + + def __init__(self, jtracker: JavaObject): + self._jtracker = jtracker + + def getJobIdsForGroup(self, jobGroup: Optional[str] = None) -> List[int]: + """ + Return a list of all known jobs in a particular job group. If + `jobGroup` is None, then returns all known jobs that are not + associated with a job group. + + The returned list may contain running, failed, and completed jobs, + and may vary across invocations of this method. This method does + not guarantee the order of the elements in its result. + """ + return list(self._jtracker.getJobIdsForGroup(jobGroup)) + + def getActiveStageIds(self) -> List[int]: + """ + Returns an array containing the ids of all active stages. + """ + return sorted(list(self._jtracker.getActiveStageIds())) + + def getActiveJobsIds(self) -> List[int]: + """ + Returns an array containing the ids of all active jobs. + """ + return sorted((list(self._jtracker.getActiveJobIds()))) + + def getJobInfo(self, jobId: int) -> Optional[SparkJobInfo]: + """ + Returns a :class:`SparkJobInfo` object, or None if the job info + could not be found or was garbage collected. + """ + job = self._jtracker.getJobInfo(jobId) + if job is not None: + return SparkJobInfo(jobId, job.stageIds(), str(job.status())) + return None + + def getStageInfo(self, stageId: int) -> Optional[SparkStageInfo]: + """ + Returns a :class:`SparkStageInfo` object, or None if the stage + info could not be found or was garbage collected. + """ + stage = self._jtracker.getStageInfo(stageId) + if stage is not None: + # TODO: fetch them in batch for better performance + attrs = [getattr(stage, f)() for f in SparkStageInfo._fields[1:]] + return SparkStageInfo(stageId, *attrs) + return None diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/storagelevel.py b/spark-3.5.0-bin-hadoop3/python/pyspark/storagelevel.py new file mode 100644 index 0000000..dabc0d2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/storagelevel.py @@ -0,0 +1,97 @@ +# +# 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. +# + +__all__ = ["StorageLevel"] + +from typing import Any, ClassVar + + +class StorageLevel: + + """ + Flags for controlling the storage of an RDD. Each StorageLevel records whether to use memory, + whether to drop the RDD to disk if it falls out of memory, whether to keep the data in memory + in a JAVA-specific serialized format, and whether to replicate the RDD partitions on multiple + nodes. Also contains static constants for some commonly used storage levels, MEMORY_ONLY. + Since the data is always serialized on the Python side, all the constants use the serialized + formats. + """ + + NONE: ClassVar["StorageLevel"] + DISK_ONLY: ClassVar["StorageLevel"] + DISK_ONLY_2: ClassVar["StorageLevel"] + DISK_ONLY_3: ClassVar["StorageLevel"] + MEMORY_ONLY: ClassVar["StorageLevel"] + MEMORY_ONLY_2: ClassVar["StorageLevel"] + MEMORY_AND_DISK: ClassVar["StorageLevel"] + MEMORY_AND_DISK_2: ClassVar["StorageLevel"] + OFF_HEAP: ClassVar["StorageLevel"] + MEMORY_AND_DISK_DESER: ClassVar["StorageLevel"] + + def __init__( + self, + useDisk: bool, + useMemory: bool, + useOffHeap: bool, + deserialized: bool, + replication: int = 1, + ): + self.useDisk = useDisk + self.useMemory = useMemory + self.useOffHeap = useOffHeap + self.deserialized = deserialized + self.replication = replication + + def __repr__(self) -> str: + return "StorageLevel(%s, %s, %s, %s, %s)" % ( + self.useDisk, + self.useMemory, + self.useOffHeap, + self.deserialized, + self.replication, + ) + + def __str__(self) -> str: + result = "" + result += "Disk " if self.useDisk else "" + result += "Memory " if self.useMemory else "" + result += "OffHeap " if self.useOffHeap else "" + result += "Deserialized " if self.deserialized else "Serialized " + result += "%sx Replicated" % self.replication + return result + + def __eq__(self, other: Any) -> bool: + return ( + isinstance(other, StorageLevel) + and self.useMemory == other.useMemory + and self.useDisk == other.useDisk + and self.useOffHeap == other.useOffHeap + and self.deserialized == other.deserialized + and self.replication == other.replication + ) + + +StorageLevel.NONE = StorageLevel(False, False, False, False) +StorageLevel.DISK_ONLY = StorageLevel(True, False, False, False) +StorageLevel.DISK_ONLY_2 = StorageLevel(True, False, False, False, 2) +StorageLevel.DISK_ONLY_3 = StorageLevel(True, False, False, False, 3) +StorageLevel.MEMORY_ONLY = StorageLevel(False, True, False, False) +StorageLevel.MEMORY_ONLY_2 = StorageLevel(False, True, False, False, 2) +StorageLevel.MEMORY_AND_DISK = StorageLevel(True, True, False, False) +StorageLevel.MEMORY_AND_DISK_2 = StorageLevel(True, True, False, False, 2) +StorageLevel.OFF_HEAP = StorageLevel(True, True, True, False, 1) +StorageLevel.MEMORY_AND_DISK_DESER = StorageLevel(True, True, False, True) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/__init__.py new file mode 100644 index 0000000..7a7dc98 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/__init__.py @@ -0,0 +1,22 @@ +# +# 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.streaming.context import StreamingContext +from pyspark.streaming.dstream import DStream +from pyspark.streaming.listener import StreamingListener + +__all__ = ["StreamingContext", "DStream", "StreamingListener"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/context.py new file mode 100644 index 0000000..84e7cd7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/context.py @@ -0,0 +1,471 @@ +# +# 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 Any, Callable, List, Optional, TypeVar + +from py4j.java_gateway import java_import, is_instance_of, JavaObject + +from pyspark import RDD, SparkConf +from pyspark.serializers import NoOpSerializer, UTF8Deserializer, CloudPickleSerializer +from pyspark.context import SparkContext +from pyspark.storagelevel import StorageLevel +from pyspark.streaming.dstream import DStream +from pyspark.streaming.listener import StreamingListener +from pyspark.streaming.util import TransformFunction, TransformFunctionSerializer + +import warnings + +__all__ = ["StreamingContext"] + +T = TypeVar("T") + + +class StreamingContext: + """ + Main entry point for Spark Streaming functionality. A StreamingContext + represents the connection to a Spark cluster, and can be used to create + :class:`DStream` various input sources. It can be from an existing :class:`SparkContext`. + After creating and transforming DStreams, the streaming computation can + be started and stopped using `context.start()` and `context.stop()`, + respectively. `context.awaitTermination()` allows the current thread + to wait for the termination of the context by `stop()` or by an exception. + + .. deprecated:: Spark 3.4.0 + This is deprecated as of Spark 3.4.0. + There are no longer updates to DStream and it's a legacy project. + There is a newer and easier to use streaming engine in Spark called Structured Streaming. + You should use Spark Structured Streaming for your streaming applications. + + Parameters + ---------- + sparkContext : :class:`SparkContext` + SparkContext object. + batchDuration : int, optional + the time interval (in seconds) at which streaming + data will be divided into batches + """ + + _transformerSerializer = None + + # Reference to a currently active StreamingContext + _activeContext = None + + def __init__( + self, + sparkContext: SparkContext, + batchDuration: Optional[int] = None, + jssc: Optional[JavaObject] = None, + ): + warnings.warn( + "DStream is deprecated as of Spark 3.4.0. Migrate to Structured Streaming.", + FutureWarning, + ) + self._sc = sparkContext + self._jvm = self._sc._jvm + self._jssc = jssc or self._initialize_context(self._sc, batchDuration) + + def _initialize_context(self, sc: SparkContext, duration: Optional[int]) -> JavaObject: + self._ensure_initialized() + assert self._jvm is not None and duration is not None + return self._jvm.JavaStreamingContext(sc._jsc, self._jduration(duration)) + + def _jduration(self, seconds: int) -> JavaObject: + """ + Create Duration object given number of seconds + """ + assert self._jvm is not None + return self._jvm.Duration(int(seconds * 1000)) + + @classmethod + def _ensure_initialized(cls) -> None: + SparkContext._ensure_initialized() + gw = SparkContext._gateway + + assert gw is not None + + java_import(gw.jvm, "org.apache.spark.streaming.*") + java_import(gw.jvm, "org.apache.spark.streaming.api.java.*") + java_import(gw.jvm, "org.apache.spark.streaming.api.python.*") + + from pyspark.java_gateway import ensure_callback_server_started + + ensure_callback_server_started(gw) + + # register serializer for TransformFunction + # it happens before creating SparkContext when loading from checkpointing + cls._transformerSerializer = TransformFunctionSerializer( + SparkContext._active_spark_context, + CloudPickleSerializer(), + gw, + ) + + @classmethod + def getOrCreate( + cls, checkpointPath: str, setupFunc: Callable[[], "StreamingContext"] + ) -> "StreamingContext": + """ + Either recreate a StreamingContext from checkpoint data or create a new StreamingContext. + If checkpoint data exists in the provided `checkpointPath`, then StreamingContext will be + recreated from the checkpoint data. If the data does not exist, then the provided setupFunc + will be used to create a new context. + + Parameters + ---------- + checkpointPath : str + Checkpoint directory used in an earlier streaming program + setupFunc : function + Function to create a new context and setup DStreams + """ + cls._ensure_initialized() + gw = SparkContext._gateway + + assert gw is not None + + # Check whether valid checkpoint information exists in the given path + ssc_option = gw.jvm.StreamingContextPythonHelper().tryRecoverFromCheckpoint(checkpointPath) + if ssc_option.isEmpty(): + ssc = setupFunc() + ssc.checkpoint(checkpointPath) + return ssc + + jssc = gw.jvm.JavaStreamingContext(ssc_option.get()) + + # If there is already an active instance of Python SparkContext use it, or create a new one + if not SparkContext._active_spark_context: + jsc = jssc.sparkContext() + conf = SparkConf(_jconf=jsc.getConf()) + SparkContext(conf=conf, gateway=gw, jsc=jsc) + + sc = SparkContext._active_spark_context + + assert sc is not None + + # update ctx in serializer + assert cls._transformerSerializer is not None + cls._transformerSerializer.ctx = sc + return StreamingContext(sc, None, jssc) + + @classmethod + def getActive(cls) -> Optional["StreamingContext"]: + """ + Return either the currently active StreamingContext (i.e., if there is a context started + but not stopped) or None. + """ + activePythonContext = cls._activeContext + if activePythonContext is not None: + # Verify that the current running Java StreamingContext is active and is the same one + # backing the supposedly active Python context + activePythonContextJavaId = activePythonContext._jssc.ssc().hashCode() + activeJvmContextOption = activePythonContext._jvm.StreamingContext.getActive() + + if activeJvmContextOption.isEmpty(): + cls._activeContext = None + elif activeJvmContextOption.get().hashCode() != activePythonContextJavaId: + cls._activeContext = None + raise RuntimeError( + "JVM's active JavaStreamingContext is not the JavaStreamingContext " + "backing the action Python StreamingContext. This is unexpected." + ) + return cls._activeContext + + @classmethod + def getActiveOrCreate( + cls, checkpointPath: str, setupFunc: Callable[[], "StreamingContext"] + ) -> "StreamingContext": + """ + Either return the active StreamingContext (i.e. currently started but not stopped), + or recreate a StreamingContext from checkpoint data or create a new StreamingContext + using the provided setupFunc function. If the checkpointPath is None or does not contain + valid checkpoint data, then setupFunc will be called to create a new context and setup + DStreams. + + Parameters + ---------- + checkpointPath : str + Checkpoint directory used in an earlier streaming program. Can be + None if the intention is to always create a new context when there + is no active context. + setupFunc : function + Function to create a new JavaStreamingContext and setup DStreams + """ + + if not callable(setupFunc): + raise TypeError("setupFunc should be callable.") + activeContext = cls.getActive() + if activeContext is not None: + return activeContext + elif checkpointPath is not None: + return cls.getOrCreate(checkpointPath, setupFunc) + else: + return setupFunc() + + @property + def sparkContext(self) -> SparkContext: + """ + Return SparkContext which is associated with this StreamingContext. + """ + return self._sc + + def start(self) -> None: + """ + Start the execution of the streams. + """ + self._jssc.start() + StreamingContext._activeContext = self + + def awaitTermination(self, timeout: Optional[int] = None) -> None: + """ + Wait for the execution to stop. + + Parameters + ---------- + timeout : int, optional + time to wait in seconds + """ + if timeout is None: + self._jssc.awaitTermination() + else: + self._jssc.awaitTerminationOrTimeout(int(timeout * 1000)) + + def awaitTerminationOrTimeout(self, timeout: int) -> None: + """ + Wait for the execution to stop. Return `true` if it's stopped; or + throw the reported error during the execution; or `false` if the + waiting time elapsed before returning from the method. + + Parameters + ---------- + timeout : int + time to wait in seconds + """ + return self._jssc.awaitTerminationOrTimeout(int(timeout * 1000)) + + def stop(self, stopSparkContext: bool = True, stopGraceFully: bool = False) -> None: + """ + Stop the execution of the streams, with option of ensuring all + received data has been processed. + + Parameters + ---------- + stopSparkContext : bool, optional + Stop the associated SparkContext or not + stopGracefully : bool, optional + Stop gracefully by waiting for the processing of all received + data to be completed + """ + self._jssc.stop(stopSparkContext, stopGraceFully) + StreamingContext._activeContext = None + if stopSparkContext: + self._sc.stop() + + def remember(self, duration: int) -> None: + """ + Set each DStreams in this context to remember RDDs it generated + in the last given duration. DStreams remember RDDs only for a + limited duration of time and releases them for garbage collection. + This method allows the developer to specify how long to remember + the RDDs (if the developer wishes to query old data outside the + DStream computation). + + Parameters + ---------- + duration : int + Minimum duration (in seconds) that each DStream should remember its RDDs + """ + self._jssc.remember(self._jduration(duration)) + + def checkpoint(self, directory: str) -> None: + """ + Sets the context to periodically checkpoint the DStream operations for master + fault-tolerance. The graph will be checkpointed every batch interval. + + Parameters + ---------- + directory : str + HDFS-compatible directory where the checkpoint data will be reliably stored + """ + self._jssc.checkpoint(directory) + + def socketTextStream( + self, hostname: str, port: int, storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK_2 + ) -> "DStream[str]": + """ + Create an input from TCP source hostname:port. Data is received using + a TCP socket and receive byte is interpreted as UTF8 encoded ``\\n`` delimited + lines. + + Parameters + ---------- + hostname : str + Hostname to connect to for receiving data + port : int + Port to connect to for receiving data + storageLevel : :class:`pyspark.StorageLevel`, optional + Storage level to use for storing the received objects + """ + jlevel = self._sc._getJavaStorageLevel(storageLevel) + return DStream( + self._jssc.socketTextStream(hostname, port, jlevel), self, UTF8Deserializer() + ) + + def textFileStream(self, directory: str) -> "DStream[str]": + """ + Create an input stream that monitors a Hadoop-compatible file system + for new files and reads them as text files. Files must be written to the + monitored directory by "moving" them from another location within the same + file system. File names starting with . are ignored. + The text files must be encoded as UTF-8. + """ + return DStream(self._jssc.textFileStream(directory), self, UTF8Deserializer()) + + def binaryRecordsStream(self, directory: str, recordLength: int) -> "DStream[bytes]": + """ + Create an input stream that monitors a Hadoop-compatible file system + for new files and reads them as flat binary files with records of + fixed length. Files must be written to the monitored directory by "moving" + them from another location within the same file system. + File names starting with . are ignored. + + Parameters + ---------- + directory : str + Directory to load data from + recordLength : int + Length of each record in bytes + """ + return DStream( + self._jssc.binaryRecordsStream(directory, recordLength), self, NoOpSerializer() + ) + + def _check_serializers(self, rdds: List[RDD[T]]) -> None: + # make sure they have same serializer + if len(set(rdd._jrdd_deserializer for rdd in rdds)) > 1: + for i in range(len(rdds)): + # reset them to sc.serializer + rdds[i] = rdds[i]._reserialize() + + def queueStream( + self, + rdds: List[RDD[T]], + oneAtATime: bool = True, + default: Optional[RDD[T]] = None, + ) -> "DStream[T]": + """ + Create an input stream from a queue of RDDs or list. In each batch, + it will process either one or all of the RDDs returned by the queue. + + Parameters + ---------- + rdds : list + Queue of RDDs + oneAtATime : bool, optional + pick one rdd each time or pick all of them once. + default : :class:`pyspark.RDD`, optional + The default rdd if no more in rdds + + Notes + ----- + Changes to the queue after the stream is created will not be recognized. + """ + if default and not isinstance(default, RDD): + default = self._sc.parallelize(default) + + if not rdds and default: + rdds = [rdds] # type: ignore[list-item] + + if rdds and not isinstance(rdds[0], RDD): + rdds = [self._sc.parallelize(input) for input in rdds] + self._check_serializers(rdds) + + assert self._jvm is not None + queue = self._jvm.PythonDStream.toRDDQueue([r._jrdd for r in rdds]) + if default: + default = default._reserialize(rdds[0]._jrdd_deserializer) + assert default is not None + jdstream = self._jssc.queueStream(queue, oneAtATime, default._jrdd) + else: + jdstream = self._jssc.queueStream(queue, oneAtATime) + return DStream(jdstream, self, rdds[0]._jrdd_deserializer) + + def transform( + self, dstreams: List["DStream[Any]"], transformFunc: Callable[..., RDD[T]] + ) -> "DStream[T]": + """ + Create a new DStream in which each RDD is generated by applying + a function on RDDs of the DStreams. The order of the JavaRDDs in + the transform function parameter will be the same as the order + of corresponding DStreams in the list. + """ + jdstreams = [d._jdstream for d in dstreams] + # change the final serializer to sc.serializer + func = TransformFunction( + self._sc, + lambda t, *rdds: transformFunc(rdds), + *[d._jrdd_deserializer for d in dstreams], + ) + + assert self._jvm is not None + jfunc = self._jvm.TransformFunction(func) + jdstream = self._jssc.transform(jdstreams, jfunc) + return DStream(jdstream, self, self._sc.serializer) + + def union(self, *dstreams: "DStream[T]") -> "DStream[T]": + """ + Create a unified DStream from multiple DStreams of the same + type and same slide duration. + """ + if not dstreams: + raise ValueError("should have at least one DStream to union") + if len(dstreams) == 1: + return dstreams[0] + if len(set(s._jrdd_deserializer for s in dstreams)) > 1: + raise ValueError("All DStreams should have same serializer") + if len(set(s._slideDuration for s in dstreams)) > 1: + raise ValueError("All DStreams should have same slide duration") + + assert SparkContext._jvm is not None + jdstream_cls = SparkContext._jvm.org.apache.spark.streaming.api.java.JavaDStream + jpair_dstream_cls = SparkContext._jvm.org.apache.spark.streaming.api.java.JavaPairDStream + gw = SparkContext._gateway + if is_instance_of(gw, dstreams[0]._jdstream, jdstream_cls): + cls = jdstream_cls + elif is_instance_of(gw, dstreams[0]._jdstream, jpair_dstream_cls): + cls = jpair_dstream_cls + else: + cls_name = dstreams[0]._jdstream.getClass().getCanonicalName() + raise TypeError("Unsupported Java DStream class %s" % cls_name) + + assert gw is not None + jdstreams = gw.new_array(cls, len(dstreams)) + for i in range(0, len(dstreams)): + jdstreams[i] = dstreams[i]._jdstream + return DStream( + self._jssc.union(jdstreams), + self, + dstreams[0]._jrdd_deserializer, + ) + + def addStreamingListener(self, streamingListener: StreamingListener) -> None: + """ + Add a [[org.apache.spark.streaming.scheduler.StreamingListener]] object for + receiving system events related to streaming. + """ + assert self._jvm is not None + self._jssc.addStreamingListener( + self._jvm.JavaStreamingListenerWrapper( + self._jvm.PythonStreamingListenerWrapper(streamingListener) + ) + ) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/dstream.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/dstream.py new file mode 100644 index 0000000..934b3ae --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/dstream.py @@ -0,0 +1,933 @@ +# +# 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 operator +import time +from itertools import chain +from datetime import datetime +from typing import ( + Any, + Callable, + Generic, + Hashable, + Iterable, + List, + Optional, + Tuple, + TypeVar, + Union, + TYPE_CHECKING, + cast, + overload, +) + +from py4j.protocol import Py4JJavaError + +from pyspark.storagelevel import StorageLevel +from pyspark.streaming.util import rddToFileName, TransformFunction +from pyspark.rdd import portable_hash, RDD +from pyspark.resultiterable import ResultIterable +from py4j.java_gateway import JavaObject + +if TYPE_CHECKING: + from pyspark.serializers import Serializer + from pyspark.streaming.context import StreamingContext + +__all__ = ["DStream"] + +S = TypeVar("S") +T = TypeVar("T") +T_co = TypeVar("T_co", covariant=True) +U = TypeVar("U") +K = TypeVar("K", bound=Hashable) +V = TypeVar("V") + + +class DStream(Generic[T_co]): + """ + A Discretized Stream (DStream), the basic abstraction in Spark Streaming, + is a continuous sequence of RDDs (of the same type) representing a + continuous stream of data (see :class:`RDD` in the Spark core documentation + for more details on RDDs). + + DStreams can either be created from live data (such as, data from TCP + sockets, etc.) using a :class:`StreamingContext` or it can be + generated by transforming existing DStreams using operations such as + `map`, `window` and `reduceByKeyAndWindow`. While a Spark Streaming + program is running, each DStream periodically generates a RDD, either + from live data or by transforming the RDD generated by a parent DStream. + + DStreams internally is characterized by a few basic properties: + - A list of other DStreams that the DStream depends on + - A time interval at which the DStream generates an RDD + - A function that is used to generate an RDD after each time interval + """ + + def __init__( + self, + jdstream: JavaObject, + ssc: "StreamingContext", + jrdd_deserializer: "Serializer", + ): + self._jdstream = jdstream + self._ssc = ssc + self._sc = ssc._sc + self._jrdd_deserializer = jrdd_deserializer + self.is_cached = False + self.is_checkpointed = False + + def context(self) -> "StreamingContext": + """ + Return the StreamingContext associated with this DStream + """ + return self._ssc + + def count(self) -> "DStream[int]": + """ + Return a new DStream in which each RDD has a single element + generated by counting each RDD of this DStream. + """ + return self.mapPartitions(lambda i: [sum(1 for _ in i)]).reduce(operator.add) + + def filter(self: "DStream[T]", f: Callable[[T], bool]) -> "DStream[T]": + """ + Return a new DStream containing only the elements that satisfy predicate. + """ + + def func(iterator: Iterable[T]) -> Iterable[T]: + return filter(f, iterator) + + return self.mapPartitions(func, True) + + def flatMap( + self: "DStream[T]", + f: Callable[[T], Iterable[U]], + preservesPartitioning: bool = False, + ) -> "DStream[U]": + """ + Return a new DStream by applying a function to all elements of + this DStream, and then flattening the results + """ + + def func(s: int, iterator: Iterable[T]) -> Iterable[U]: + return chain.from_iterable(map(f, iterator)) + + return self.mapPartitionsWithIndex(func, preservesPartitioning) + + def map( + self: "DStream[T]", f: Callable[[T], U], preservesPartitioning: bool = False + ) -> "DStream[U]": + """ + Return a new DStream by applying a function to each element of DStream. + """ + + def func(iterator: Iterable[T]) -> Iterable[U]: + return map(f, iterator) + + return self.mapPartitions(func, preservesPartitioning) + + def mapPartitions( + self: "DStream[T]", + f: Callable[[Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + ) -> "DStream[U]": + """ + Return a new DStream in which each RDD is generated by applying + mapPartitions() to each RDDs of this DStream. + """ + + def func(s: int, iterator: Iterable[T]) -> Iterable[U]: + return f(iterator) + + return self.mapPartitionsWithIndex(func, preservesPartitioning) + + def mapPartitionsWithIndex( + self: "DStream[T]", + f: Callable[[int, Iterable[T]], Iterable[U]], + preservesPartitioning: bool = False, + ) -> "DStream[U]": + """ + Return a new DStream in which each RDD is generated by applying + mapPartitionsWithIndex() to each RDDs of this DStream. + """ + return self.transform(lambda rdd: rdd.mapPartitionsWithIndex(f, preservesPartitioning)) + + def reduce(self: "DStream[T]", func: Callable[[T, T], T]) -> "DStream[T]": + """ + Return a new DStream in which each RDD has a single element + generated by reducing each RDD of this DStream. + """ + return self.map(lambda x: (None, x)).reduceByKey(func, 1).map(lambda x: x[1]) + + def reduceByKey( + self: "DStream[Tuple[K, V]]", + func: Callable[[V, V], V], + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, V]]": + """ + Return a new DStream by applying reduceByKey to each RDD. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.combineByKey(lambda x: x, func, func, numPartitions) + + def combineByKey( + self: "DStream[Tuple[K, V]]", + createCombiner: Callable[[V], U], + mergeValue: Callable[[U, V], U], + mergeCombiners: Callable[[U, U], U], + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, U]]": + """ + Return a new DStream by applying combineByKey to each RDD. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + + def func(rdd: RDD[Tuple[K, V]]) -> RDD[Tuple[K, U]]: + return rdd.combineByKey(createCombiner, mergeValue, mergeCombiners, numPartitions) + + return self.transform(func) + + def partitionBy( + self: "DStream[Tuple[K, V]]", + numPartitions: int, + partitionFunc: Callable[[K], int] = portable_hash, + ) -> "DStream[Tuple[K, V]]": + """ + Return a copy of the DStream in which each RDD are partitioned + using the specified partitioner. + """ + return self.transform(lambda rdd: rdd.partitionBy(numPartitions, partitionFunc)) + + @overload + def foreachRDD(self: "DStream[T]", func: Callable[[RDD[T]], None]) -> None: + ... + + @overload + def foreachRDD(self: "DStream[T]", func: Callable[[datetime, RDD[T]], None]) -> None: + ... + + def foreachRDD( + self: "DStream[T]", + func: Union[Callable[[RDD[T]], None], Callable[[datetime, RDD[T]], None]], + ) -> None: + """ + Apply a function to each RDD in this DStream. + """ + if func.__code__.co_argcount == 1: + old_func = func + + def func(_: datetime, rdd: "RDD[T]") -> None: + return old_func(rdd) # type: ignore[call-arg, arg-type] + + jfunc = TransformFunction(self._sc, func, self._jrdd_deserializer) + assert self._ssc._jvm is not None + api = self._ssc._jvm.PythonDStream + api.callForeachRDD(self._jdstream, jfunc) + + def pprint(self, num: int = 10) -> None: + """ + Print the first num elements of each RDD generated in this DStream. + + Parameters + ---------- + num : int, optional + the number of elements from the first will be printed. + """ + + def takeAndPrint(time: datetime, rdd: RDD[T]) -> None: + taken = rdd.take(num + 1) + print("-------------------------------------------") + print("Time: %s" % time) + print("-------------------------------------------") + for record in taken[:num]: + print(record) + if len(taken) > num: + print("...") + print("") + + self.foreachRDD(takeAndPrint) + + def mapValues(self: "DStream[Tuple[K, V]]", f: Callable[[V], U]) -> "DStream[Tuple[K, U]]": + """ + Return a new DStream by applying a map function to the value of + each key-value pairs in this DStream without changing the key. + """ + + def map_values_fn(kv: Tuple[K, V]) -> Tuple[K, U]: + return kv[0], f(kv[1]) + + return self.map(map_values_fn, preservesPartitioning=True) + + def flatMapValues( + self: "DStream[Tuple[K, V]]", f: Callable[[V], Iterable[U]] + ) -> "DStream[Tuple[K, U]]": + """ + Return a new DStream by applying a flatmap function to the value + of each key-value pairs in this DStream without changing the key. + """ + + def flat_map_fn(kv: Tuple[K, V]) -> Iterable[Tuple[K, U]]: + return ((kv[0], x) for x in f(kv[1])) + + return self.flatMap(flat_map_fn, preservesPartitioning=True) + + def glom(self: "DStream[T]") -> "DStream[List[T]]": + """ + Return a new DStream in which RDD is generated by applying glom() + to RDD of this DStream. + """ + + def func(iterator: Iterable[T]) -> Iterable[List[T]]: + yield list(iterator) + + return self.mapPartitions(func) + + def cache(self: "DStream[T]") -> "DStream[T]": + """ + Persist the RDDs of this DStream with the default storage level + (`MEMORY_ONLY`). + """ + self.is_cached = True + self.persist(StorageLevel.MEMORY_ONLY) + return self + + def persist(self: "DStream[T]", storageLevel: StorageLevel) -> "DStream[T]": + """ + Persist the RDDs of this DStream with the given storage level + """ + self.is_cached = True + javaStorageLevel = self._sc._getJavaStorageLevel(storageLevel) + self._jdstream.persist(javaStorageLevel) + return self + + def checkpoint(self: "DStream[T]", interval: int) -> "DStream[T]": + """ + Enable periodic checkpointing of RDDs of this DStream + + Parameters + ---------- + interval : int + time in seconds, after each period of that, generated + RDD will be checkpointed + """ + self.is_checkpointed = True + self._jdstream.checkpoint(self._ssc._jduration(interval)) + return self + + def groupByKey( + self: "DStream[Tuple[K, V]]", numPartitions: Optional[int] = None + ) -> "DStream[Tuple[K, Iterable[V]]]": + """ + Return a new DStream by applying groupByKey on each RDD. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transform(lambda rdd: rdd.groupByKey(numPartitions)) + + def countByValue(self: "DStream[K]") -> "DStream[Tuple[K, int]]": + """ + Return a new DStream in which each RDD contains the counts of each + distinct value in each RDD of this DStream. + """ + return self.map(lambda x: (x, 1)).reduceByKey(lambda x, y: x + y) + + def saveAsTextFiles(self, prefix: str, suffix: Optional[str] = None) -> None: + """ + Save each RDD in this DStream as at text file, using string + representation of elements. + """ + + def saveAsTextFile(t: Optional[datetime], rdd: RDD[T]) -> None: + path = rddToFileName(prefix, suffix, t) + try: + rdd.saveAsTextFile(path) + except Py4JJavaError as e: + # after recovered from checkpointing, the foreachRDD may + # be called twice + if "FileAlreadyExistsException" not in str(e): + raise + + return self.foreachRDD(saveAsTextFile) + + # TODO: uncomment this until we have ssc.pickleFileStream() + # def saveAsPickleFiles(self, prefix, suffix=None): + # """ + # Save each RDD in this DStream as at binary file, the elements are + # serialized by pickle. + # """ + # def saveAsPickleFile(t, rdd): + # path = rddToFileName(prefix, suffix, t) + # try: + # rdd.saveAsPickleFile(path) + # except Py4JJavaError as e: + # # after recovered from checkpointing, the foreachRDD may + # # be called twice + # if 'FileAlreadyExistsException' not in str(e): + # raise + # return self.foreachRDD(saveAsPickleFile) + + @overload + def transform(self: "DStream[T]", func: Callable[[RDD[T]], RDD[U]]) -> "TransformedDStream[U]": + ... + + @overload + def transform( + self: "DStream[T]", func: Callable[[datetime, RDD[T]], RDD[U]] + ) -> "TransformedDStream[U]": + ... + + def transform( + self: "DStream[T]", + func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]], + ) -> "TransformedDStream[U]": + """ + Return a new DStream in which each RDD is generated by applying a function + on each RDD of this DStream. + + `func` can have one argument of `rdd`, or have two arguments of + (`time`, `rdd`) + """ + if func.__code__.co_argcount == 1: + oldfunc = func + + def func(_: datetime, rdd: RDD[T]) -> RDD[U]: + return oldfunc(rdd) # type: ignore[arg-type, call-arg] + + assert func.__code__.co_argcount == 2, "func should take one or two arguments" + return TransformedDStream(self, func) + + @overload + def transformWith( + self: "DStream[T]", + func: Callable[[RDD[T], RDD[U]], RDD[V]], + other: "DStream[U]", + keepSerializer: bool = ..., + ) -> "DStream[V]": + ... + + @overload + def transformWith( + self: "DStream[T]", + func: Callable[[datetime, RDD[T], RDD[U]], RDD[V]], + other: "DStream[U]", + keepSerializer: bool = ..., + ) -> "DStream[V]": + ... + + def transformWith( + self: "DStream[T]", + func: Union[ + Callable[[RDD[T], RDD[U]], RDD[V]], + Callable[[datetime, RDD[T], RDD[U]], RDD[V]], + ], + other: "DStream[U]", + keepSerializer: bool = False, + ) -> "DStream[V]": + """ + Return a new DStream in which each RDD is generated by applying a function + on each RDD of this DStream and 'other' DStream. + + `func` can have two arguments of (`rdd_a`, `rdd_b`) or have three + arguments of (`time`, `rdd_a`, `rdd_b`) + """ + if func.__code__.co_argcount == 2: + oldfunc = func + + def func(_: datetime, a: RDD[T], b: RDD[U]) -> RDD[V]: + return oldfunc(a, b) # type: ignore[call-arg, arg-type] + + assert func.__code__.co_argcount == 3, "func should take two or three arguments" + jfunc = TransformFunction( + self._sc, + func, + self._jrdd_deserializer, + other._jrdd_deserializer, + ) + assert self._sc._jvm is not None + dstream = self._sc._jvm.PythonTransformed2DStream( + self._jdstream.dstream(), other._jdstream.dstream(), jfunc + ) + jrdd_serializer = self._jrdd_deserializer if keepSerializer else self._sc.serializer + return DStream(dstream.asJavaDStream(), self._ssc, jrdd_serializer) + + def repartition(self: "DStream[T]", numPartitions: int) -> "DStream[T]": + """ + Return a new DStream with an increased or decreased level of parallelism. + """ + return self.transform(lambda rdd: rdd.repartition(numPartitions)) + + @property + def _slideDuration(self) -> None: + """ + Return the slideDuration in seconds of this DStream + """ + return self._jdstream.dstream().slideDuration().milliseconds() / 1000.0 + + def union(self: "DStream[T]", other: "DStream[U]") -> "DStream[Union[T, U]]": + """ + Return a new DStream by unifying data of another DStream with this DStream. + + Parameters + ---------- + other : :class:`DStream` + Another DStream having the same interval (i.e., slideDuration) + as this DStream. + """ + if self._slideDuration != other._slideDuration: + raise ValueError("the two DStream should have same slide duration") + return self.transformWith(lambda a, b: a.union(b), other, True) + + def cogroup( + self: "DStream[Tuple[K, V]]", + other: "DStream[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Tuple[ResultIterable[V], ResultIterable[U]]]]": + """ + Return a new DStream by applying 'cogroup' between RDDs of this + DStream and `other` DStream. + + Hash partitioning is used to generate the RDDs with `numPartitions` partitions. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transformWith( + lambda a, b: a.cogroup(b, numPartitions), + other, + ) + + def join( + self: "DStream[Tuple[K, V]]", + other: "DStream[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Tuple[V, U]]]": + """ + Return a new DStream by applying 'join' between RDDs of this DStream and + `other` DStream. + + Hash partitioning is used to generate the RDDs with `numPartitions` + partitions. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transformWith(lambda a, b: a.join(b, numPartitions), other) + + def leftOuterJoin( + self: "DStream[Tuple[K, V]]", + other: "DStream[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Tuple[V, Optional[U]]]]": + """ + Return a new DStream by applying 'left outer join' between RDDs of this DStream and + `other` DStream. + + Hash partitioning is used to generate the RDDs with `numPartitions` + partitions. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transformWith(lambda a, b: a.leftOuterJoin(b, numPartitions), other) + + def rightOuterJoin( + self: "DStream[Tuple[K, V]]", + other: "DStream[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Tuple[Optional[V], U]]]": + """ + Return a new DStream by applying 'right outer join' between RDDs of this DStream and + `other` DStream. + + Hash partitioning is used to generate the RDDs with `numPartitions` + partitions. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transformWith(lambda a, b: a.rightOuterJoin(b, numPartitions), other) + + def fullOuterJoin( + self: "DStream[Tuple[K, V]]", + other: "DStream[Tuple[K, U]]", + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Tuple[Optional[V], Optional[U]]]]": + """ + Return a new DStream by applying 'full outer join' between RDDs of this DStream and + `other` DStream. + + Hash partitioning is used to generate the RDDs with `numPartitions` + partitions. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + return self.transformWith(lambda a, b: a.fullOuterJoin(b, numPartitions), other) + + def _jtime(self, timestamp: Union[datetime, int, float]) -> JavaObject: + """Convert datetime or unix_timestamp into Time""" + if isinstance(timestamp, datetime): + timestamp = time.mktime(timestamp.timetuple()) + assert self._sc._jvm is not None + return self._sc._jvm.Time(int(timestamp * 1000)) + + def slice(self, begin: Union[datetime, int], end: Union[datetime, int]) -> List[RDD[T]]: + """ + Return all the RDDs between 'begin' to 'end' (both included) + + `begin`, `end` could be datetime.datetime() or unix_timestamp + """ + jrdds = self._jdstream.slice(self._jtime(begin), self._jtime(end)) + return [RDD(jrdd, self._sc, self._jrdd_deserializer) for jrdd in jrdds] + + def _validate_window_param(self, window: int, slide: Optional[int]) -> None: + duration = self._jdstream.dstream().slideDuration().milliseconds() + if int(window * 1000) % duration != 0: + raise ValueError( + "windowDuration must be multiple of the parent " + "dstream's slide (batch) duration (%d ms)" % duration + ) + if slide and int(slide * 1000) % duration != 0: + raise ValueError( + "slideDuration must be multiple of the parent " + "dstream's slide (batch) duration (%d ms)" % duration + ) + + def window(self, windowDuration: int, slideDuration: Optional[int] = None) -> "DStream[T]": + """ + Return a new DStream in which each RDD contains all the elements in seen in a + sliding window of time over this DStream. + + Parameters + ---------- + windowDuration : int + width of the window; must be a multiple of this DStream's + batching interval + slideDuration : int, optional + sliding interval of the window (i.e., the interval after which + the new DStream will generate RDDs); must be a multiple of this + DStream's batching interval + """ + self._validate_window_param(windowDuration, slideDuration) + d = self._ssc._jduration(windowDuration) + if slideDuration is None: + return DStream(self._jdstream.window(d), self._ssc, self._jrdd_deserializer) + s = self._ssc._jduration(slideDuration) + return DStream(self._jdstream.window(d, s), self._ssc, self._jrdd_deserializer) + + def reduceByWindow( + self: "DStream[T]", + reduceFunc: Callable[[T, T], T], + invReduceFunc: Optional[Callable[[T, T], T]], + windowDuration: int, + slideDuration: int, + ) -> "DStream[T]": + """ + Return a new DStream in which each RDD has a single element generated by reducing all + elements in a sliding window over this DStream. + + if `invReduceFunc` is not None, the reduction is done incrementally + using the old window's reduced value : + + 1. reduce the new values that entered the window (e.g., adding new counts) + + 2. "inverse reduce" the old values that left the window (e.g., subtracting old counts) + This is more efficient than `invReduceFunc` is None. + + Parameters + ---------- + reduceFunc : function + associative and commutative reduce function + invReduceFunc : function + inverse reduce function of `reduceFunc`; such that for all y, + and invertible x: + `invReduceFunc(reduceFunc(x, y), x) = y` + windowDuration : int + width of the window; must be a multiple of this DStream's + batching interval + slideDuration : int + sliding interval of the window (i.e., the interval after which + the new DStream will generate RDDs); must be a multiple of this + DStream's batching interval + """ + keyed = self.map(lambda x: (1, x)) + reduced = keyed.reduceByKeyAndWindow( + reduceFunc, invReduceFunc, windowDuration, slideDuration, 1 + ) + return reduced.map(lambda kv: kv[1]) + + def countByWindow( + self: "DStream[T]", windowDuration: int, slideDuration: int + ) -> "DStream[int]": + """ + Return a new DStream in which each RDD has a single element generated + by counting the number of elements in a window over this DStream. + windowDuration and slideDuration are as defined in the window() operation. + + This is equivalent to window(windowDuration, slideDuration).count(), + but will be more efficient if window is large. + """ + return self.map(lambda x: 1).reduceByWindow( + operator.add, operator.sub, windowDuration, slideDuration + ) + + def countByValueAndWindow( + self: "DStream[T]", + windowDuration: int, + slideDuration: int, + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[T, int]]": + """ + Return a new DStream in which each RDD contains the count of distinct elements in + RDDs in a sliding window over this DStream. + + Parameters + ---------- + windowDuration : int + width of the window; must be a multiple of this DStream's + batching interval + slideDuration : int + sliding interval of the window (i.e., the interval after which + the new DStream will generate RDDs); must be a multiple of this + DStream's batching interval + numPartitions : int, optional + number of partitions of each RDD in the new DStream. + """ + keyed = self.map(lambda x: (x, 1)) + counted = keyed.reduceByKeyAndWindow( + operator.add, operator.sub, windowDuration, slideDuration, numPartitions + ) + return counted.filter(lambda kv: kv[1] > 0) + + def groupByKeyAndWindow( + self: "DStream[Tuple[K, V]]", + windowDuration: int, + slideDuration: int, + numPartitions: Optional[int] = None, + ) -> "DStream[Tuple[K, Iterable[V]]]": + """ + Return a new DStream by applying `groupByKey` over a sliding window. + Similar to `DStream.groupByKey()`, but applies it over a sliding window. + + Parameters + ---------- + windowDuration : int + width of the window; must be a multiple of this DStream's + batching interval + slideDuration : int + sliding interval of the window (i.e., the interval after which + the new DStream will generate RDDs); must be a multiple of this + DStream's batching interval + numPartitions : int, optional + Number of partitions of each RDD in the new DStream. + """ + ls = self.mapValues(lambda x: [x]) + grouped = ls.reduceByKeyAndWindow( + lambda a, b: a.extend(b) or a, # type: ignore[func-returns-value] + lambda a, b: a[len(b) :], + windowDuration, + slideDuration, + numPartitions, + ) + return grouped.mapValues(ResultIterable) + + def reduceByKeyAndWindow( + self: "DStream[Tuple[K, V]]", + func: Callable[[V, V], V], + invFunc: Optional[Callable[[V, V], V]], + windowDuration: int, + slideDuration: Optional[int] = None, + numPartitions: Optional[int] = None, + filterFunc: Optional[Callable[[Tuple[K, V]], bool]] = None, + ) -> "DStream[Tuple[K, V]]": + """ + Return a new DStream by applying incremental `reduceByKey` over a sliding window. + + The reduced value of over a new window is calculated using the old window's reduce value : + 1. reduce the new values that entered the window (e.g., adding new counts) + 2. "inverse reduce" the old values that left the window (e.g., subtracting old counts) + + `invFunc` can be None, then it will reduce all the RDDs in window, could be slower + than having `invFunc`. + + Parameters + ---------- + func : function + associative and commutative reduce function + invFunc : function + inverse function of `reduceFunc` + windowDuration : int + width of the window; must be a multiple of this DStream's + batching interval + slideDuration : int, optional + sliding interval of the window (i.e., the interval after which + the new DStream will generate RDDs); must be a multiple of this + DStream's batching interval + numPartitions : int, optional + number of partitions of each RDD in the new DStream. + filterFunc : function, optional + function to filter expired key-value pairs; + only pairs that satisfy the function are retained + set this to null if you do not want to filter + """ + self._validate_window_param(windowDuration, slideDuration) + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + + reduced = self.reduceByKey(func, numPartitions) + + if invFunc: + + def reduceFunc(t: datetime, a: Any, b: Any) -> Any: + b = b.reduceByKey(func, numPartitions) + r = a.union(b).reduceByKey(func, numPartitions) if a else b + if filterFunc: + r = r.filter(filterFunc) + return r + + def invReduceFunc(t: datetime, a: Any, b: Any) -> Any: + b = b.reduceByKey(func, numPartitions) + joined = a.leftOuterJoin(b, numPartitions) + return joined.mapValues( + lambda kv: invFunc(kv[0], kv[1]) # type: ignore[misc] + if kv[1] is not None + else kv[0] + ) + + jreduceFunc = TransformFunction(self._sc, reduceFunc, reduced._jrdd_deserializer) + jinvReduceFunc = TransformFunction(self._sc, invReduceFunc, reduced._jrdd_deserializer) + if slideDuration is None: + slideDuration = self._slideDuration + assert self._sc._jvm is not None + dstream = self._sc._jvm.PythonReducedWindowedDStream( + reduced._jdstream.dstream(), + jreduceFunc, + jinvReduceFunc, + self._ssc._jduration(windowDuration), + self._ssc._jduration(slideDuration), # type: ignore[arg-type] + ) + return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer) + else: + return reduced.window(windowDuration, slideDuration).reduceByKey( + func, numPartitions # type: ignore[arg-type] + ) + + def updateStateByKey( + self: "DStream[Tuple[K, V]]", + updateFunc: Callable[[Iterable[V], Optional[S]], S], + numPartitions: Optional[int] = None, + initialRDD: Optional[Union[RDD[Tuple[K, S]], Iterable[Tuple[K, S]]]] = None, + ) -> "DStream[Tuple[K, S]]": + """ + Return a new "state" DStream where the state for each key is updated by applying + the given function on the previous state of the key and the new values of the key. + + Parameters + ---------- + updateFunc : function + State update function. If this function returns None, then + corresponding state key-value pair will be eliminated. + """ + if numPartitions is None: + numPartitions = self._sc.defaultParallelism + + if initialRDD and not isinstance(initialRDD, RDD): + initialRDD = self._sc.parallelize(initialRDD) + + def reduceFunc(t: datetime, a: Any, b: Any) -> Any: + if a is None: + g = b.groupByKey(numPartitions).mapValues(lambda vs: (list(vs), None)) + else: + g = a.cogroup(b.partitionBy(cast(int, numPartitions)), numPartitions) + g = g.mapValues(lambda ab: (list(ab[1]), list(ab[0])[0] if len(ab[0]) else None)) + state = g.mapValues(lambda vs_s: updateFunc(vs_s[0], vs_s[1])) + return state.filter(lambda k_v: k_v[1] is not None) + + jreduceFunc = TransformFunction( + self._sc, + reduceFunc, + self._sc.serializer, + self._jrdd_deserializer, + ) + if initialRDD: + initialRDD = cast(RDD[Tuple[K, S]], initialRDD)._reserialize(self._jrdd_deserializer) + assert self._sc._jvm is not None + dstream = self._sc._jvm.PythonStateDStream( + self._jdstream.dstream(), + jreduceFunc, + initialRDD._jrdd, + ) + else: + assert self._sc._jvm is not None + dstream = self._sc._jvm.PythonStateDStream(self._jdstream.dstream(), jreduceFunc) + + return DStream(dstream.asJavaDStream(), self._ssc, self._sc.serializer) + + +class TransformedDStream(DStream[U]): + """ + TransformedDStream is a DStream generated by an Python function + transforming each RDD of a DStream to another RDDs. + + Multiple continuous transformations of DStream can be combined into + one transformation. + """ + + @overload + def __init__(self: DStream[U], prev: DStream[T], func: Callable[[RDD[T]], RDD[U]]): + ... + + @overload + def __init__( + self: DStream[U], + prev: DStream[T], + func: Callable[[datetime, RDD[T]], RDD[U]], + ): + ... + + def __init__( + self, + prev: DStream[T], + func: Union[Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]]], + ): + self._ssc = prev._ssc + self._sc = self._ssc._sc + self._jrdd_deserializer = self._sc.serializer + self.is_cached = False + self.is_checkpointed = False + self._jdstream_val = None + + # Using type() to avoid folding the functions and compacting the DStreams which is not + # not strictly an object of TransformedDStream. + if type(prev) is TransformedDStream and not prev.is_cached and not prev.is_checkpointed: + prev_func: Callable = prev.func + func = cast(Callable[[datetime, RDD[T]], RDD[U]], func) + self.func: Union[ + Callable[[RDD[T]], RDD[U]], Callable[[datetime, RDD[T]], RDD[U]] + ] = lambda t, rdd: func(t, prev_func(t, rdd)) + self.prev: DStream[T] = prev.prev + else: + self.prev = prev + self.func = func + + @property + def _jdstream(self) -> JavaObject: + if self._jdstream_val is not None: + return self._jdstream_val + + jfunc = TransformFunction(self._sc, self.func, self.prev._jrdd_deserializer) + assert self._sc._jvm is not None + dstream = self._sc._jvm.PythonTransformedDStream(self.prev._jdstream.dstream(), jfunc) + self._jdstream_val = dstream.asJavaDStream() + return self._jdstream_val diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/kinesis.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/kinesis.py new file mode 100644 index 0000000..0eede34 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/kinesis.py @@ -0,0 +1,205 @@ +# +# 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 overload, Callable, Optional, TypeVar, Union + +from pyspark.serializers import NoOpSerializer +from pyspark.storagelevel import StorageLevel +from pyspark.streaming import DStream +from pyspark.streaming.context import StreamingContext +from pyspark.util import _print_missing_jar + + +__all__ = ["KinesisUtils", "InitialPositionInStream", "MetricsLevel", "utf8_decoder"] + + +class InitialPositionInStream: + LATEST, TRIM_HORIZON = (0, 1) + + +class MetricsLevel: + DETAILED, SUMMARY, NONE = (0, 1, 2) + + +T = TypeVar("T") + + +def utf8_decoder(s: Optional[bytes]) -> Optional[str]: + """Decode the unicode as UTF-8""" + if s is None: + return None + return s.decode("utf-8") + + +class KinesisUtils: + @staticmethod + @overload + def createStream( + ssc: StreamingContext, + kinesisAppName: str, + streamName: str, + endpointUrl: str, + regionName: str, + initialPositionInStream: str, + checkpointInterval: int, + metricsLevel: int = MetricsLevel.DETAILED, + storageLevel: StorageLevel = ..., + awsAccessKeyId: Optional[str] = ..., + awsSecretKey: Optional[str] = ..., + *, + stsAssumeRoleArn: Optional[str] = ..., + stsSessionName: Optional[str] = ..., + stsExternalId: Optional[str] = ..., + ) -> "DStream[Optional[str]]": + ... + + @staticmethod + @overload + def createStream( + ssc: StreamingContext, + kinesisAppName: str, + streamName: str, + endpointUrl: str, + regionName: str, + initialPositionInStream: str, + checkpointInterval: int, + metricsLevel: int = MetricsLevel.DETAILED, + storageLevel: StorageLevel = ..., + awsAccessKeyId: Optional[str] = ..., + awsSecretKey: Optional[str] = ..., + decoder: Callable[[Optional[bytes]], T] = ..., + stsAssumeRoleArn: Optional[str] = ..., + stsSessionName: Optional[str] = ..., + stsExternalId: Optional[str] = ..., + ) -> "DStream[T]": + ... + + @staticmethod + def createStream( + ssc: StreamingContext, + kinesisAppName: str, + streamName: str, + endpointUrl: str, + regionName: str, + initialPositionInStream: str, + checkpointInterval: int, + metricsLevel: int = MetricsLevel.DETAILED, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK_2, + awsAccessKeyId: Optional[str] = None, + awsSecretKey: Optional[str] = None, + decoder: Union[ + Callable[[Optional[bytes]], T], Callable[[Optional[bytes]], Optional[str]] + ] = utf8_decoder, + stsAssumeRoleArn: Optional[str] = None, + stsSessionName: Optional[str] = None, + stsExternalId: Optional[str] = None, + ) -> Union["DStream[Union[T, Optional[str]]]", "DStream[T]"]: + """ + Create an input stream that pulls messages from a Kinesis stream. This uses the + Kinesis Client Library (KCL) to pull messages from Kinesis. + + Parameters + ---------- + ssc : :class:`StreamingContext` + StreamingContext object + kinesisAppName : str + Kinesis application name used by the Kinesis Client Library (KCL) to + update DynamoDB + streamName : str + Kinesis stream name + endpointUrl : str + Url of Kinesis service (e.g., https://kinesis.us-east-1.amazonaws.com) + regionName : str + Name of region used by the Kinesis Client Library (KCL) to update + DynamoDB (lease coordination and checkpointing) and CloudWatch (metrics) + initialPositionInStream : int + In the absence of Kinesis checkpoint info, this is the + worker's initial starting position in the stream. The + values are either the beginning of the stream per Kinesis' + limit of 24 hours (InitialPositionInStream.TRIM_HORIZON) or + the tip of the stream (InitialPositionInStream.LATEST). + checkpointInterval : int + Checkpoint interval(in seconds) for Kinesis checkpointing. See the Kinesis + Spark Streaming documentation for more details on the different + types of checkpoints. + metricsLevel : int + Level of CloudWatch PutMetrics. + Can be set to either DETAILED, SUMMARY, or NONE. (default is DETAILED) + storageLevel : :class:`pyspark.StorageLevel`, optional + Storage level to use for storing the received objects (default is + StorageLevel.MEMORY_AND_DISK_2) + awsAccessKeyId : str, optional + AWS AccessKeyId (default is None. If None, will use + DefaultAWSCredentialsProviderChain) + awsSecretKey : str, optional + AWS SecretKey (default is None. If None, will use + DefaultAWSCredentialsProviderChain) + decoder : function, optional + A function used to decode value (default is utf8_decoder) + stsAssumeRoleArn : str, optional + ARN of IAM role to assume when using STS sessions to read from + the Kinesis stream (default is None). + stsSessionName : str, optional + Name to uniquely identify STS sessions used to read from Kinesis + stream, if STS is being used (default is None). + stsExternalId : str, optional + External ID that can be used to validate against the assumed IAM + role's trust policy, if STS is being used (default is None). + + Returns + ------- + A DStream object + + Notes + ----- + The given AWS credentials will get saved in DStream checkpoints if checkpointing + is enabled. Make sure that your checkpoint directory is secure. + """ + jlevel = ssc._sc._getJavaStorageLevel(storageLevel) + jduration = ssc._jduration(checkpointInterval) + + jvm = ssc._jvm + assert jvm is not None + + try: + helper = jvm.org.apache.spark.streaming.kinesis.KinesisUtilsPythonHelper() + except TypeError as e: + if str(e) == "'JavaPackage' object is not callable": + _print_missing_jar( + "Streaming's Kinesis", + "streaming-kinesis-asl", + "streaming-kinesis-asl-assembly", + ssc.sparkContext.version, + ) + raise + jstream = helper.createStream( + ssc._jssc, + kinesisAppName, + streamName, + endpointUrl, + regionName, + initialPositionInStream, + jduration, + metricsLevel, + jlevel, + awsAccessKeyId, + awsSecretKey, + stsAssumeRoleArn, + stsSessionName, + stsExternalId, + ) + stream: DStream = DStream(jstream, ssc, NoOpSerializer()) + return stream.map(lambda v: decoder(v)) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/listener.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/listener.py new file mode 100644 index 0000000..75c4aa2 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/listener.py @@ -0,0 +1,83 @@ +# +# 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 Any + + +__all__ = ["StreamingListener"] + + +class StreamingListener: + def __init__(self) -> None: + pass + + def onStreamingStarted(self, streamingStarted: Any) -> None: + """ + Called when the streaming has been started. + """ + pass + + def onReceiverStarted(self, receiverStarted: Any) -> None: + """ + Called when a receiver has been started + """ + pass + + def onReceiverError(self, receiverError: Any) -> None: + """ + Called when a receiver has reported an error + """ + pass + + def onReceiverStopped(self, receiverStopped: Any) -> None: + """ + Called when a receiver has been stopped + """ + pass + + def onBatchSubmitted(self, batchSubmitted: Any) -> None: + """ + Called when a batch of jobs has been submitted for processing. + """ + pass + + def onBatchStarted(self, batchSubmitted: Any) -> None: + """ + Called when processing of a batch of jobs has started. + """ + pass + + def onBatchCompleted(self, batchCompleted: Any) -> None: + """ + Called when processing of a batch of jobs has completed. + """ + pass + + def onOutputOperationStarted(self, outputOperationStarted: Any) -> None: + """ + Called when processing of a job of a batch has started. + """ + pass + + def onOutputOperationCompleted(self, outputOperationCompleted: Any) -> None: + """ + Called when processing of a job of a batch has completed + """ + pass + + class Java: + implements = ["org.apache.spark.streaming.api.java.PythonStreamingListener"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/__init__.py new file mode 100644 index 0000000..cce3aca --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_context.py new file mode 100644 index 0000000..1afcc90 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_context.py @@ -0,0 +1,184 @@ +# +# 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 os +import struct +import tempfile +import time + +from pyspark.streaming import StreamingContext +from pyspark.testing.streamingutils import PySparkStreamingTestCase + + +class StreamingContextTests(PySparkStreamingTestCase): + + duration = 0.1 + setupCalled = False + + def _add_input_stream(self): + inputs = [range(1, x) for x in range(101)] + stream = self.ssc.queueStream(inputs) + self._collect(stream, 1, block=False) + + def test_stop_only_streaming_context(self): + self._add_input_stream() + self.ssc.start() + self.ssc.stop(False) + self.assertEqual(len(self.sc.parallelize(range(5), 5).glom().collect()), 5) + + def test_stop_multiple_times(self): + self._add_input_stream() + self.ssc.start() + self.ssc.stop(False) + self.ssc.stop(False) + + def test_queue_stream(self): + input = [list(range(i + 1)) for i in range(3)] + dstream = self.ssc.queueStream(input) + result = self._collect(dstream, 3) + self.assertEqual(input, result) + + def test_text_file_stream(self): + d = tempfile.mkdtemp() + self.ssc = StreamingContext(self.sc, self.duration) + dstream2 = self.ssc.textFileStream(d).map(int) + result = self._collect(dstream2, 2, block=False) + self.ssc.start() + for name in ("a", "b"): + time.sleep(1) + with open(os.path.join(d, name), "w") as f: + f.writelines(["%d\n" % i for i in range(10)]) + self.wait_for(result, 2) + self.assertEqual([list(range(10)), list(range(10))], result) + + def test_binary_records_stream(self): + d = tempfile.mkdtemp() + self.ssc = StreamingContext(self.sc, self.duration) + dstream = self.ssc.binaryRecordsStream(d, 10).map(lambda v: struct.unpack("10b", bytes(v))) + result = self._collect(dstream, 2, block=False) + self.ssc.start() + for name in ("a", "b"): + time.sleep(1) + with open(os.path.join(d, name), "wb") as f: + f.write(bytearray(range(10))) + self.wait_for(result, 2) + self.assertEqual([list(range(10)), list(range(10))], [list(v[0]) for v in result]) + + def test_union(self): + input = [list(range(i + 1)) for i in range(3)] + dstream = self.ssc.queueStream(input) + dstream2 = self.ssc.queueStream(input) + dstream3 = self.ssc.union(dstream, dstream2) + result = self._collect(dstream3, 3) + expected = [i * 2 for i in input] + self.assertEqual(expected, result) + + def test_transform(self): + dstream1 = self.ssc.queueStream([[1]]) + dstream2 = self.ssc.queueStream([[2]]) + dstream3 = self.ssc.queueStream([[3]]) + + def func(rdds): + rdd1, rdd2, rdd3 = rdds + return rdd2.union(rdd3).union(rdd1) + + dstream = self.ssc.transform([dstream1, dstream2, dstream3], func) + + self.assertEqual([2, 3, 1], self._take(dstream, 3)) + + def test_transform_pairrdd(self): + # This regression test case is for SPARK-17756. + dstream = self.ssc.queueStream([[1], [2], [3]]).transform(lambda rdd: rdd.cartesian(rdd)) + self.assertEqual([(1, 1), (2, 2), (3, 3)], self._take(dstream, 3)) + + def test_get_active(self): + self.assertEqual(StreamingContext.getActive(), None) + + # Verify that getActive() returns the active context + self.ssc.queueStream([[1]]).foreachRDD(lambda rdd: rdd.count()) + self.ssc.start() + self.assertEqual(StreamingContext.getActive(), self.ssc) + + # Verify that getActive() returns None + self.ssc.stop(False) + self.assertEqual(StreamingContext.getActive(), None) + + # Verify that if the Java context is stopped, then getActive() returns None + self.ssc = StreamingContext(self.sc, self.duration) + self.ssc.queueStream([[1]]).foreachRDD(lambda rdd: rdd.count()) + self.ssc.start() + self.assertEqual(StreamingContext.getActive(), self.ssc) + self.ssc._jssc.stop(False) + self.assertEqual(StreamingContext.getActive(), None) + + def test_get_active_or_create(self): + # Test StreamingContext.getActiveOrCreate() without checkpoint data + # See CheckpointTests for tests with checkpoint data + self.ssc = None + self.assertEqual(StreamingContext.getActive(), None) + + def setupFunc(): + ssc = StreamingContext(self.sc, self.duration) + ssc.queueStream([[1]]).foreachRDD(lambda rdd: rdd.count()) + self.setupCalled = True + return ssc + + # Verify that getActiveOrCreate() (w/o checkpoint) calls setupFunc when no context is active + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(None, setupFunc) + self.assertTrue(self.setupCalled) + + # Verify that getActiveOrCreate() returns active context and does not call the setupFunc + self.ssc.start() + self.setupCalled = False + self.assertEqual(StreamingContext.getActiveOrCreate(None, setupFunc), self.ssc) + self.assertFalse(self.setupCalled) + + # Verify that getActiveOrCreate() calls setupFunc after active context is stopped + self.ssc.stop(False) + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(None, setupFunc) + self.assertTrue(self.setupCalled) + + # Verify that if the Java context is stopped, then getActive() returns None + self.ssc = StreamingContext(self.sc, self.duration) + self.ssc.queueStream([[1]]).foreachRDD(lambda rdd: rdd.count()) + self.ssc.start() + self.assertEqual(StreamingContext.getActive(), self.ssc) + self.ssc._jssc.stop(False) + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(None, setupFunc) + self.assertTrue(self.setupCalled) + + def test_await_termination_or_timeout(self): + self._add_input_stream() + self.ssc.start() + self.assertFalse(self.ssc.awaitTerminationOrTimeout(0.001)) + self.ssc.stop(False) + self.assertTrue(self.ssc.awaitTerminationOrTimeout(0.001)) + + +if __name__ == "__main__": + import unittest + from pyspark.streaming.tests.test_context import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_dstream.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_dstream.py new file mode 100644 index 0000000..d37e64a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_dstream.py @@ -0,0 +1,706 @@ +# +# 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 operator +import os +import shutil +import tempfile +import time +import unittest +from functools import reduce +from itertools import chain +import platform + +from pyspark import SparkConf, SparkContext +from pyspark.streaming import StreamingContext +from pyspark.testing.streamingutils import PySparkStreamingTestCase + + +@unittest.skipIf( + "pypy" in platform.python_implementation().lower(), + "The tests fail in PyPy3 implementation for an unknown reason.", +) +class BasicOperationTests(PySparkStreamingTestCase): + def test_map(self): + """Basic operation test for DStream.map.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + + def func(dstream): + return dstream.map(str) + + expected = [list(map(str, x)) for x in input] + self._test_func(input, func, expected) + + def test_flatMap(self): + """Basic operation test for DStream.flatMap.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + + def func(dstream): + return dstream.flatMap(lambda x: (x, x * 2)) + + expected = [list(chain.from_iterable((map(lambda y: [y, y * 2], x)))) for x in input] + self._test_func(input, func, expected) + + def test_filter(self): + """Basic operation test for DStream.filter.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + + def func(dstream): + return dstream.filter(lambda x: x % 2 == 0) + + expected = [[y for y in x if y % 2 == 0] for x in input] + self._test_func(input, func, expected) + + def test_count(self): + """Basic operation test for DStream.count.""" + input = [range(5), range(10), range(20)] + + def func(dstream): + return dstream.count() + + expected = [[len(x)] for x in input] + self._test_func(input, func, expected) + + def test_slice(self): + """Basic operation test for DStream.slice.""" + import datetime as dt + + self.ssc = StreamingContext(self.sc, 1.0) + self.ssc.remember(4.0) + input = [[1], [2], [3], [4]] + stream = self.ssc.queueStream([self.sc.parallelize(d, 1) for d in input]) + + time_vals = [] + + def get_times(t, rdd): + if rdd and len(time_vals) < len(input): + time_vals.append(t) + + stream.foreachRDD(get_times) + + self.ssc.start() + self.wait_for(time_vals, 4) + begin_time = time_vals[0] + + def get_sliced(begin_delta, end_delta): + begin = begin_time + dt.timedelta(seconds=begin_delta) + end = begin_time + dt.timedelta(seconds=end_delta) + rdds = stream.slice(begin, end) + result_list = [rdd.collect() for rdd in rdds] + return [r for result in result_list for r in result] + + self.assertEqual(set([1]), set(get_sliced(0, 0))) + self.assertEqual(set([2, 3]), set(get_sliced(1, 2))) + self.assertEqual(set([2, 3, 4]), set(get_sliced(1, 4))) + self.assertEqual(set([1, 2, 3, 4]), set(get_sliced(0, 4))) + + def test_reduce(self): + """Basic operation test for DStream.reduce.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + + def func(dstream): + return dstream.reduce(operator.add) + + expected = [[reduce(operator.add, x)] for x in input] + self._test_func(input, func, expected) + + def test_reduceByKey(self): + """Basic operation test for DStream.reduceByKey.""" + input = [ + [("a", 1), ("a", 1), ("b", 1), ("b", 1)], + [("", 1), ("", 1), ("", 1), ("", 1)], + [(1, 1), (1, 1), (2, 1), (2, 1), (3, 1)], + ] + + def func(dstream): + return dstream.reduceByKey(operator.add) + + expected = [[("a", 2), ("b", 2)], [("", 4)], [(1, 2), (2, 2), (3, 1)]] + self._test_func(input, func, expected, sort=True) + + def test_mapValues(self): + """Basic operation test for DStream.mapValues.""" + input = [ + [("a", 2), ("b", 2), ("c", 1), ("d", 1)], + [(0, 4), (1, 1), (2, 2), (3, 3)], + [(1, 1), (2, 1), (3, 1), (4, 1)], + ] + + def func(dstream): + return dstream.mapValues(lambda x: x + 10) + + expected = [ + [("a", 12), ("b", 12), ("c", 11), ("d", 11)], + [(0, 14), (1, 11), (2, 12), (3, 13)], + [(1, 11), (2, 11), (3, 11), (4, 11)], + ] + self._test_func(input, func, expected, sort=True) + + def test_flatMapValues(self): + """Basic operation test for DStream.flatMapValues.""" + input = [ + [("a", 2), ("b", 2), ("c", 1), ("d", 1)], + [(0, 4), (1, 1), (2, 1), (3, 1)], + [(1, 1), (2, 1), (3, 1), (4, 1)], + ] + + def func(dstream): + return dstream.flatMapValues(lambda x: (x, x + 10)) + + expected = [ + [("a", 2), ("a", 12), ("b", 2), ("b", 12), ("c", 1), ("c", 11), ("d", 1), ("d", 11)], + [(0, 4), (0, 14), (1, 1), (1, 11), (2, 1), (2, 11), (3, 1), (3, 11)], + [(1, 1), (1, 11), (2, 1), (2, 11), (3, 1), (3, 11), (4, 1), (4, 11)], + ] + self._test_func(input, func, expected) + + def test_glom(self): + """Basic operation test for DStream.glom.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + rdds = [self.sc.parallelize(r, 2) for r in input] + + def func(dstream): + return dstream.glom() + + expected = [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]] + self._test_func(rdds, func, expected) + + def test_mapPartitions(self): + """Basic operation test for DStream.mapPartitions.""" + input = [range(1, 5), range(5, 9), range(9, 13)] + rdds = [self.sc.parallelize(r, 2) for r in input] + + def func(dstream): + def f(iterator): + yield sum(iterator) + + return dstream.mapPartitions(f) + + expected = [[3, 7], [11, 15], [19, 23]] + self._test_func(rdds, func, expected) + + def test_countByValue(self): + """Basic operation test for DStream.countByValue.""" + input = [list(range(1, 5)) * 2, list(range(5, 7)) + list(range(5, 9)), ["a", "a", "b", ""]] + + def func(dstream): + return dstream.countByValue() + + expected = [ + [(1, 2), (2, 2), (3, 2), (4, 2)], + [(5, 2), (6, 2), (7, 1), (8, 1)], + [("a", 2), ("b", 1), ("", 1)], + ] + self._test_func(input, func, expected, sort=True) + + def test_groupByKey(self): + """Basic operation test for DStream.groupByKey.""" + input = [ + [(1, 1), (2, 1), (3, 1), (4, 1)], + [(1, 1), (1, 1), (1, 1), (2, 1), (2, 1), (3, 1)], + [("a", 1), ("a", 1), ("b", 1), ("", 1), ("", 1), ("", 1)], + ] + + def func(dstream): + return dstream.groupByKey().mapValues(list) + + expected = [ + [(1, [1]), (2, [1]), (3, [1]), (4, [1])], + [(1, [1, 1, 1]), (2, [1, 1]), (3, [1])], + [("a", [1, 1]), ("b", [1]), ("", [1, 1, 1])], + ] + self._test_func(input, func, expected, sort=True) + + def test_combineByKey(self): + """Basic operation test for DStream.combineByKey.""" + input = [ + [(1, 1), (2, 1), (3, 1), (4, 1)], + [(1, 1), (1, 1), (1, 1), (2, 1), (2, 1), (3, 1)], + [("a", 1), ("a", 1), ("b", 1), ("", 1), ("", 1), ("", 1)], + ] + + def func(dstream): + def add(a, b): + return a + str(b) + + return dstream.combineByKey(str, add, add) + + expected = [ + [(1, "1"), (2, "1"), (3, "1"), (4, "1")], + [(1, "111"), (2, "11"), (3, "1")], + [("a", "11"), ("b", "1"), ("", "111")], + ] + self._test_func(input, func, expected, sort=True) + + def test_repartition(self): + input = [range(1, 5), range(5, 9)] + rdds = [self.sc.parallelize(r, 2) for r in input] + + def func(dstream): + return dstream.repartition(1).glom() + + expected = [[[1, 2, 3, 4]], [[5, 6, 7, 8]]] + self._test_func(rdds, func, expected) + + def test_union(self): + input1 = [range(3), range(5), range(6)] + input2 = [range(3, 6), range(5, 6)] + + def func(d1, d2): + return d1.union(d2) + + expected = [list(range(6)), list(range(6)), list(range(6))] + self._test_func(input1, func, expected, input2=input2) + + def test_cogroup(self): + input = [ + [(1, 1), (2, 1), (3, 1)], + [(1, 1), (1, 1), (1, 1), (2, 1)], + [("a", 1), ("a", 1), ("b", 1), ("", 1), ("", 1)], + ] + input2 = [[(1, 2)], [(4, 1)], [("a", 1), ("a", 1), ("b", 1), ("", 1), ("", 2)]] + + def func(d1, d2): + return d1.cogroup(d2).mapValues(lambda vs: tuple(map(list, vs))) + + expected = [ + [(1, ([1], [2])), (2, ([1], [])), (3, ([1], []))], + [(1, ([1, 1, 1], [])), (2, ([1], [])), (4, ([], [1]))], + [("a", ([1, 1], [1, 1])), ("b", ([1], [1])), ("", ([1, 1], [1, 2]))], + ] + self._test_func(input, func, expected, sort=True, input2=input2) + + def test_join(self): + input = [[("a", 1), ("b", 2)]] + input2 = [[("b", 3), ("c", 4)]] + + def func(a, b): + return a.join(b) + + expected = [[("b", (2, 3))]] + self._test_func(input, func, expected, True, input2) + + def test_left_outer_join(self): + input = [[("a", 1), ("b", 2)]] + input2 = [[("b", 3), ("c", 4)]] + + def func(a, b): + return a.leftOuterJoin(b) + + expected = [[("a", (1, None)), ("b", (2, 3))]] + self._test_func(input, func, expected, True, input2) + + def test_right_outer_join(self): + input = [[("a", 1), ("b", 2)]] + input2 = [[("b", 3), ("c", 4)]] + + def func(a, b): + return a.rightOuterJoin(b) + + expected = [[("b", (2, 3)), ("c", (None, 4))]] + self._test_func(input, func, expected, True, input2) + + def test_full_outer_join(self): + input = [[("a", 1), ("b", 2)]] + input2 = [[("b", 3), ("c", 4)]] + + def func(a, b): + return a.fullOuterJoin(b) + + expected = [[("a", (1, None)), ("b", (2, 3)), ("c", (None, 4))]] + self._test_func(input, func, expected, True, input2) + + def test_update_state_by_key(self): + def updater(vs, s): + if not s: + s = [] + s.extend(vs) + return s + + input = [[("k", i)] for i in range(5)] + + def func(dstream): + return dstream.updateStateByKey(updater) + + expected = [[0], [0, 1], [0, 1, 2], [0, 1, 2, 3], [0, 1, 2, 3, 4]] + expected = [[("k", v)] for v in expected] + self._test_func(input, func, expected) + + def test_update_state_by_key_initial_rdd(self): + def updater(vs, s): + if not s: + s = [] + s.extend(vs) + return s + + initial = [("k", [0, 1])] + initial = self.sc.parallelize(initial, 1) + + input = [[("k", i)] for i in range(2, 5)] + + def func(dstream): + return dstream.updateStateByKey(updater, initialRDD=initial) + + expected = [[0, 1, 2], [0, 1, 2, 3], [0, 1, 2, 3, 4]] + expected = [[("k", v)] for v in expected] + self._test_func(input, func, expected) + + def test_failed_func(self): + # Test failure in + # TransformFunction.apply(rdd: Option[RDD[_]], time: Time) + input = [self.sc.parallelize([d], 1) for d in range(4)] + input_stream = self.ssc.queueStream(input) + + def failed_func(i): + raise ValueError("This is a special error") + + input_stream.map(failed_func).pprint() + self.ssc.start() + try: + self.ssc.awaitTerminationOrTimeout(10) + except BaseException: + import traceback + + failure = traceback.format_exc() + self.assertTrue("This is a special error" in failure) + return + + self.fail("a failed func should throw an error") + + def test_failed_func2(self): + # Test failure in + # TransformFunction.apply(rdd: Option[RDD[_]], rdd2: Option[RDD[_]], time: Time) + input = [self.sc.parallelize([d], 1) for d in range(4)] + input_stream1 = self.ssc.queueStream(input) + input_stream2 = self.ssc.queueStream(input) + + def failed_func(rdd1, rdd2): + raise ValueError("This is a special error") + + input_stream1.transformWith(failed_func, input_stream2, True).pprint() + self.ssc.start() + try: + self.ssc.awaitTerminationOrTimeout(10) + except BaseException: + import traceback + + failure = traceback.format_exc() + self.assertTrue("This is a special error" in failure) + return + + self.fail("a failed func should throw an error") + + def test_failed_func_with_reseting_failure(self): + input = [self.sc.parallelize([d], 1) for d in range(4)] + input_stream = self.ssc.queueStream(input) + + def failed_func(i): + if i == 1: + # Make it fail in the second batch + raise ValueError("This is a special error") + else: + return i + + # We should be able to see the results of the 3rd and 4th batches even if the second batch + # fails + expected = [[0], [2], [3]] + self.assertEqual(expected, self._collect(input_stream.map(failed_func), 3)) + try: + self.ssc.awaitTerminationOrTimeout(10) + except BaseException: + import traceback + + failure = traceback.format_exc() + self.assertTrue("This is a special error" in failure) + return + + self.fail("a failed func should throw an error") + + +@unittest.skipIf( + "pypy" in platform.python_implementation().lower(), + "The tests fail in PyPy3 implementation for an unknown reason.", +) +class WindowFunctionTests(PySparkStreamingTestCase): + + timeout = 15 + + def test_window(self): + input = [range(1), range(2), range(3), range(4), range(5)] + + def func(dstream): + return dstream.window(1.5, 0.5).count() + + expected = [[1], [3], [6], [9], [12], [9], [5]] + self._test_func(input, func, expected) + + def test_count_by_window(self): + input = [range(1), range(2), range(3), range(4), range(5)] + + def func(dstream): + return dstream.countByWindow(1.5, 0.5) + + expected = [[1], [3], [6], [9], [12], [9], [5]] + self._test_func(input, func, expected) + + def test_count_by_window_large(self): + input = [range(1), range(2), range(3), range(4), range(5), range(6)] + + def func(dstream): + return dstream.countByWindow(2.5, 0.5) + + expected = [[1], [3], [6], [10], [15], [20], [18], [15], [11], [6]] + self._test_func(input, func, expected) + + def test_count_by_value_and_window(self): + input = [range(1), range(2), range(3), range(4), range(5), range(6)] + + def func(dstream): + return dstream.countByValueAndWindow(2.5, 0.5) + + expected = [ + [(0, 1)], + [(0, 2), (1, 1)], + [(0, 3), (1, 2), (2, 1)], + [(0, 4), (1, 3), (2, 2), (3, 1)], + [(0, 5), (1, 4), (2, 3), (3, 2), (4, 1)], + [(0, 5), (1, 5), (2, 4), (3, 3), (4, 2), (5, 1)], + [(0, 4), (1, 4), (2, 4), (3, 3), (4, 2), (5, 1)], + [(0, 3), (1, 3), (2, 3), (3, 3), (4, 2), (5, 1)], + [(0, 2), (1, 2), (2, 2), (3, 2), (4, 2), (5, 1)], + [(0, 1), (1, 1), (2, 1), (3, 1), (4, 1), (5, 1)], + ] + self._test_func(input, func, expected) + + def test_group_by_key_and_window(self): + input = [[("a", i)] for i in range(5)] + + def func(dstream): + return dstream.groupByKeyAndWindow(1.5, 0.5).mapValues(list) + + expected = [ + [("a", [0])], + [("a", [0, 1])], + [("a", [0, 1, 2])], + [("a", [1, 2, 3])], + [("a", [2, 3, 4])], + [("a", [3, 4])], + [("a", [4])], + ] + self._test_func(input, func, expected) + + def test_reduce_by_invalid_window(self): + input1 = [range(3), range(5), range(1), range(6)] + d1 = self.ssc.queueStream(input1) + self.assertRaises(ValueError, lambda: d1.reduceByKeyAndWindow(None, None, 0.1, 0.1)) + self.assertRaises(ValueError, lambda: d1.reduceByKeyAndWindow(None, None, 1, 0.1)) + + def test_reduce_by_key_and_window_with_none_invFunc(self): + input = [range(1), range(2), range(3), range(4), range(5), range(6)] + + def func(dstream): + return ( + dstream.map(lambda x: (x, 1)) + .reduceByKeyAndWindow(operator.add, None, 5, 1) + .filter(lambda kv: kv[1] > 0) + .count() + ) + + expected = [[2], [4], [6], [6], [6], [6]] + self._test_func(input, func, expected) + + +@unittest.skipIf( + "pypy" in platform.python_implementation().lower(), + "The tests fail in PyPy3 implementation for an unknown reason.", +) +class CheckpointTests(unittest.TestCase): + + setupCalled = False + + @staticmethod + def tearDownClass(): + # Clean up in the JVM just in case there has been some issues in Python API + if SparkContext._jvm is not None: + jStreamingContextOption = ( + SparkContext._jvm.org.apache.spark.streaming.StreamingContext.getActive() + ) + if jStreamingContextOption.nonEmpty(): + jStreamingContextOption.get().stop() + + def setUp(self): + self.ssc = None + self.sc = None + self.cpd = None + + def tearDown(self): + if self.ssc is not None: + self.ssc.stop(True) + if self.sc is not None: + self.sc.stop() + if self.cpd is not None: + shutil.rmtree(self.cpd) + + def test_transform_function_serializer_failure(self): + inputd = tempfile.mkdtemp() + self.cpd = tempfile.mkdtemp("test_transform_function_serializer_failure") + + def setup(): + conf = SparkConf().set("spark.default.parallelism", 1) + sc = SparkContext(conf=conf) + ssc = StreamingContext(sc, 0.5) + + # A function that cannot be serialized + def process(time, rdd): + sc.parallelize(range(1, 10)) + + ssc.textFileStream(inputd).foreachRDD(process) + return ssc + + self.ssc = StreamingContext.getOrCreate(self.cpd, setup) + try: + self.ssc.start() + except BaseException: + import traceback + + failure = traceback.format_exc() + self.assertTrue( + "It appears that you are attempting to reference SparkContext" in failure + ) + return + + self.fail("using SparkContext in process should fail because it's not Serializable") + + def test_get_or_create_and_get_active_or_create(self): + inputd = tempfile.mkdtemp() + outputd = tempfile.mkdtemp() + "/" + + def updater(vs, s): + return sum(vs, s or 0) + + def setup(): + conf = SparkConf().set("spark.default.parallelism", 1) + sc = SparkContext(conf=conf) + ssc = StreamingContext(sc, 2) + dstream = ssc.textFileStream(inputd).map(lambda x: (x, 1)) + wc = dstream.updateStateByKey(updater) + wc.map(lambda x: "%s,%d" % x).saveAsTextFiles(outputd + "test") + wc.checkpoint(2) + self.setupCalled = True + return ssc + + # Verify that getOrCreate() calls setup() in absence of checkpoint files + self.cpd = tempfile.mkdtemp("test_streaming_cps") + self.setupCalled = False + self.ssc = StreamingContext.getOrCreate(self.cpd, setup) + self.assertTrue(self.setupCalled) + + self.ssc.start() + + def check_output(n): + while not os.listdir(outputd): + if self.ssc.awaitTerminationOrTimeout(0.5): + raise RuntimeError("ssc stopped") + time.sleep(1) # make sure mtime is larger than the previous one + with open(os.path.join(inputd, str(n)), "w") as f: + f.writelines(["%d\n" % i for i in range(10)]) + + while True: + if self.ssc.awaitTerminationOrTimeout(0.5): + raise RuntimeError("ssc stopped") + p = os.path.join(outputd, max(os.listdir(outputd))) + if "_SUCCESS" not in os.listdir(p): + # not finished + continue + ordd = self.ssc.sparkContext.textFile(p).map(lambda line: line.split(",")) + d = ordd.values().map(int).collect() + if not d: + continue + self.assertEqual(10, len(d)) + s = set(d) + self.assertEqual(1, len(s)) + m = s.pop() + if n > m: + continue + self.assertEqual(n, m) + break + + check_output(1) + check_output(2) + + # Verify the getOrCreate() recovers from checkpoint files + self.ssc.stop(True, True) + time.sleep(1) + self.setupCalled = False + self.ssc = StreamingContext.getOrCreate(self.cpd, setup) + self.assertFalse(self.setupCalled) + self.ssc.start() + check_output(3) + + # Verify that getOrCreate() uses existing SparkContext + self.ssc.stop(True, True) + time.sleep(1) + self.sc = SparkContext(conf=SparkConf()) + self.setupCalled = False + self.ssc = StreamingContext.getOrCreate(self.cpd, setup) + self.assertFalse(self.setupCalled) + self.assertTrue(self.ssc.sparkContext == self.sc) + + # Verify the getActiveOrCreate() recovers from checkpoint files + self.ssc.stop(True, True) + time.sleep(1) + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(self.cpd, setup) + self.assertFalse(self.setupCalled) + self.ssc.start() + check_output(4) + + # Verify that getActiveOrCreate() returns active context + self.setupCalled = False + self.assertEqual(StreamingContext.getActiveOrCreate(self.cpd, setup), self.ssc) + self.assertFalse(self.setupCalled) + + # Verify that getActiveOrCreate() uses existing SparkContext + self.ssc.stop(True, True) + time.sleep(1) + self.sc = SparkContext(conf=SparkConf()) + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(self.cpd, setup) + self.assertFalse(self.setupCalled) + self.assertTrue(self.ssc.sparkContext == self.sc) + + # Verify that getActiveOrCreate() calls setup() in absence of checkpoint files + self.ssc.stop(True, True) + shutil.rmtree(self.cpd) # delete checkpoint directory + time.sleep(1) + self.setupCalled = False + self.ssc = StreamingContext.getActiveOrCreate(self.cpd, setup) + self.assertTrue(self.setupCalled) + + # Stop everything + self.ssc.stop(True, True) + + +if __name__ == "__main__": + from pyspark.streaming.tests.test_dstream import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_kinesis.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_kinesis.py new file mode 100644 index 0000000..7efd7a7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_kinesis.py @@ -0,0 +1,118 @@ +# +# 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 time +import unittest + +from pyspark import StorageLevel +from pyspark.streaming.kinesis import KinesisUtils, InitialPositionInStream, MetricsLevel +from pyspark.testing.streamingutils import ( + should_test_kinesis, + kinesis_requirement_message, + PySparkStreamingTestCase, +) + + +@unittest.skipIf(not should_test_kinesis, kinesis_requirement_message) +class KinesisStreamTests(PySparkStreamingTestCase): + def test_kinesis_stream_api(self): + # Don't start the StreamingContext because we cannot test it in Jenkins + KinesisUtils.createStream( + self.ssc, + "myAppNam", + "mySparkStream", + "https://kinesis.us-west-2.amazonaws.com", + "us-west-2", + InitialPositionInStream.LATEST, + 2, + MetricsLevel.DETAILED, + StorageLevel.MEMORY_AND_DISK_2, + ) + KinesisUtils.createStream( + self.ssc, + "myAppNam", + "mySparkStream", + "https://kinesis.us-west-2.amazonaws.com", + "us-west-2", + InitialPositionInStream.LATEST, + 2, + MetricsLevel.DETAILED, + StorageLevel.MEMORY_AND_DISK_2, + "awsAccessKey", + "awsSecretKey", + ) + + def test_kinesis_stream(self): + import random + + kinesisAppName = "KinesisStreamTests-%d" % abs(random.randint(0, 10000000)) + kinesisTestUtils = self.ssc._jvm.org.apache.spark.streaming.kinesis.KinesisTestUtils(2) + try: + kinesisTestUtils.createStream() + aWSCredentials = kinesisTestUtils.getAWSCredentials() + stream = KinesisUtils.createStream( + self.ssc, + kinesisAppName, + kinesisTestUtils.streamName(), + kinesisTestUtils.endpointUrl(), + kinesisTestUtils.regionName(), + InitialPositionInStream.LATEST, + 10, + MetricsLevel.DETAILED, + StorageLevel.MEMORY_ONLY, + aWSCredentials.getAWSAccessKeyId(), + aWSCredentials.getAWSSecretKey(), + ) + + outputBuffer = [] + + def get_output(_, rdd): + for e in rdd.collect(): + outputBuffer.append(e) + + stream.foreachRDD(get_output) + self.ssc.start() + + testData = [i for i in range(1, 11)] + expectedOutput = set([str(i) for i in testData]) + start_time = time.time() + while time.time() - start_time < 120: + kinesisTestUtils.pushData(testData) + if expectedOutput == set(outputBuffer): + break + time.sleep(10) + self.assertEqual(expectedOutput, set(outputBuffer)) + except BaseException: + import traceback + + traceback.print_exc() + raise + finally: + self.ssc.stop(False) + kinesisTestUtils.deleteStream() + kinesisTestUtils.deleteDynamoDBTable(kinesisAppName) + + +if __name__ == "__main__": + from pyspark.streaming.tests.test_kinesis import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_listener.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_listener.py new file mode 100644 index 0000000..aeec278 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/tests/test_listener.py @@ -0,0 +1,160 @@ +# +# 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.streaming import StreamingListener +from pyspark.testing.streamingutils import PySparkStreamingTestCase + + +class StreamingListenerTests(PySparkStreamingTestCase): + + duration = 0.5 + + class BatchInfoCollector(StreamingListener): + def __init__(self): + super(StreamingListener, self).__init__() + self.batchInfosCompleted = [] + self.batchInfosStarted = [] + self.batchInfosSubmitted = [] + self.streamingStartedTime = [] + + def onStreamingStarted(self, streamingStarted): + self.streamingStartedTime.append(streamingStarted.time) + + def onBatchSubmitted(self, batchSubmitted): + self.batchInfosSubmitted.append(batchSubmitted.batchInfo()) + + def onBatchStarted(self, batchStarted): + self.batchInfosStarted.append(batchStarted.batchInfo()) + + def onBatchCompleted(self, batchCompleted): + self.batchInfosCompleted.append(batchCompleted.batchInfo()) + + def test_batch_info_reports(self): + batch_collector = self.BatchInfoCollector() + self.ssc.addStreamingListener(batch_collector) + input = [[1], [2], [3], [4]] + + def func(dstream): + return dstream.map(int) + + expected = [[1], [2], [3], [4]] + self._test_func(input, func, expected) + + batchInfosSubmitted = batch_collector.batchInfosSubmitted + batchInfosStarted = batch_collector.batchInfosStarted + batchInfosCompleted = batch_collector.batchInfosCompleted + streamingStartedTime = batch_collector.streamingStartedTime + + self.wait_for(batchInfosCompleted, 4) + + self.assertEqual(len(streamingStartedTime), 1) + + self.assertGreaterEqual(len(batchInfosSubmitted), 4) + for info in batchInfosSubmitted: + self.assertGreaterEqual(info.batchTime().milliseconds(), 0) + self.assertGreaterEqual(info.submissionTime(), 0) + + for streamId in info.streamIdToInputInfo(): + streamInputInfo = info.streamIdToInputInfo()[streamId] + self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0) + self.assertGreaterEqual(streamInputInfo.numRecords, 0) + for key in streamInputInfo.metadata(): + self.assertIsNotNone(streamInputInfo.metadata()[key]) + self.assertIsNotNone(streamInputInfo.metadataDescription()) + + for outputOpId in info.outputOperationInfos(): + outputInfo = info.outputOperationInfos()[outputOpId] + self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0) + self.assertGreaterEqual(outputInfo.id(), 0) + self.assertIsNotNone(outputInfo.name()) + self.assertIsNotNone(outputInfo.description()) + self.assertGreaterEqual(outputInfo.startTime(), -1) + self.assertGreaterEqual(outputInfo.endTime(), -1) + self.assertIsNone(outputInfo.failureReason()) + + self.assertEqual(info.schedulingDelay(), -1) + self.assertEqual(info.processingDelay(), -1) + self.assertEqual(info.totalDelay(), -1) + self.assertEqual(info.numRecords(), 0) + + self.assertGreaterEqual(len(batchInfosStarted), 4) + for info in batchInfosStarted: + self.assertGreaterEqual(info.batchTime().milliseconds(), 0) + self.assertGreaterEqual(info.submissionTime(), 0) + + for streamId in info.streamIdToInputInfo(): + streamInputInfo = info.streamIdToInputInfo()[streamId] + self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0) + self.assertGreaterEqual(streamInputInfo.numRecords, 0) + for key in streamInputInfo.metadata(): + self.assertIsNotNone(streamInputInfo.metadata()[key]) + self.assertIsNotNone(streamInputInfo.metadataDescription()) + + for outputOpId in info.outputOperationInfos(): + outputInfo = info.outputOperationInfos()[outputOpId] + self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0) + self.assertGreaterEqual(outputInfo.id(), 0) + self.assertIsNotNone(outputInfo.name()) + self.assertIsNotNone(outputInfo.description()) + self.assertGreaterEqual(outputInfo.startTime(), -1) + self.assertGreaterEqual(outputInfo.endTime(), -1) + self.assertIsNone(outputInfo.failureReason()) + + self.assertGreaterEqual(info.schedulingDelay(), 0) + self.assertEqual(info.processingDelay(), -1) + self.assertEqual(info.totalDelay(), -1) + self.assertEqual(info.numRecords(), 0) + + self.assertGreaterEqual(len(batchInfosCompleted), 4) + for info in batchInfosCompleted: + self.assertGreaterEqual(info.batchTime().milliseconds(), 0) + self.assertGreaterEqual(info.submissionTime(), 0) + + for streamId in info.streamIdToInputInfo(): + streamInputInfo = info.streamIdToInputInfo()[streamId] + self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0) + self.assertGreaterEqual(streamInputInfo.numRecords, 0) + for key in streamInputInfo.metadata(): + self.assertIsNotNone(streamInputInfo.metadata()[key]) + self.assertIsNotNone(streamInputInfo.metadataDescription()) + + for outputOpId in info.outputOperationInfos(): + outputInfo = info.outputOperationInfos()[outputOpId] + self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0) + self.assertGreaterEqual(outputInfo.id(), 0) + self.assertIsNotNone(outputInfo.name()) + self.assertIsNotNone(outputInfo.description()) + self.assertGreaterEqual(outputInfo.startTime(), 0) + self.assertGreaterEqual(outputInfo.endTime(), 0) + self.assertIsNone(outputInfo.failureReason()) + + self.assertGreaterEqual(info.schedulingDelay(), 0) + self.assertGreaterEqual(info.processingDelay(), 0) + self.assertGreaterEqual(info.totalDelay(), 0) + self.assertEqual(info.numRecords(), 0) + + +if __name__ == "__main__": + import unittest + from pyspark.streaming.tests.test_listener import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/util.py new file mode 100644 index 0000000..5cbc5f6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/streaming/util.py @@ -0,0 +1,168 @@ +# +# 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 time +from datetime import datetime +import traceback +import sys + +from py4j.java_gateway import is_instance_of + +from pyspark import SparkContext, RDD + + +class TransformFunction: + """ + This class wraps a function RDD[X] -> RDD[Y] that was passed to + DStream.transform(), allowing it to be called from Java via Py4J's + callback server. + + Java calls this function with a sequence of JavaRDDs and this function + returns a single JavaRDD pointer back to Java. + """ + + _emptyRDD = None + + def __init__(self, ctx, func, *deserializers): + self.ctx = ctx + self.func = func + self.deserializers = deserializers + self.rdd_wrap_func = lambda jrdd, ctx, ser: RDD(jrdd, ctx, ser) + self.failure = None + + def rdd_wrapper(self, func): + self.rdd_wrap_func = func + return self + + def call(self, milliseconds, jrdds): + # Clear the failure + self.failure = None + try: + if self.ctx is None: + self.ctx = SparkContext._active_spark_context + if not self.ctx or not self.ctx._jsc: + # stopped + return + + # extend deserializers with the first one + sers = self.deserializers + if len(sers) < len(jrdds): + sers += (sers[0],) * (len(jrdds) - len(sers)) + + rdds = [ + self.rdd_wrap_func(jrdd, self.ctx, ser) if jrdd else None + for jrdd, ser in zip(jrdds, sers) + ] + t = datetime.fromtimestamp(milliseconds / 1000.0) + r = self.func(t, *rdds) + if r: + # Here, we work around to ensure `_jrdd` is `JavaRDD` by wrapping it by `map`. + # org.apache.spark.streaming.api.python.PythonTransformFunction requires to return + # `JavaRDD`; however, this could be `JavaPairRDD` by some APIs, for example, `zip`. + # See SPARK-17756. + if is_instance_of(self.ctx._gateway, r._jrdd, "org.apache.spark.api.java.JavaRDD"): + return r._jrdd + else: + return r.map(lambda x: x)._jrdd + except BaseException: + self.failure = traceback.format_exc() + + def getLastFailure(self): + return self.failure + + def __repr__(self): + return "TransformFunction(%s)" % self.func + + class Java: + implements = ["org.apache.spark.streaming.api.python.PythonTransformFunction"] + + +class TransformFunctionSerializer: + """ + This class implements a serializer for PythonTransformFunction Java + objects. + + This is necessary because the Java PythonTransformFunction objects are + actually Py4J references to Python objects and thus are not directly + serializable. When Java needs to serialize a PythonTransformFunction, + it uses this class to invoke Python, which returns the serialized function + as a byte array. + """ + + def __init__(self, ctx, serializer, gateway=None): + self.ctx = ctx + self.serializer = serializer + self.gateway = gateway or self.ctx._gateway + self.gateway.jvm.PythonDStream.registerSerializer(self) + self.failure = None + + def dumps(self, id): + # Clear the failure + self.failure = None + try: + func = self.gateway.gateway_property.pool[id] + return bytearray( + self.serializer.dumps((func.func, func.rdd_wrap_func, func.deserializers)) + ) + except BaseException: + self.failure = traceback.format_exc() + + def loads(self, data): + # Clear the failure + self.failure = None + try: + f, wrap_func, deserializers = self.serializer.loads(bytes(data)) + return TransformFunction(self.ctx, f, *deserializers).rdd_wrapper(wrap_func) + except BaseException: + self.failure = traceback.format_exc() + + def getLastFailure(self): + return self.failure + + def __repr__(self): + return "TransformFunctionSerializer(%s)" % self.serializer + + class Java: + implements = ["org.apache.spark.streaming.api.python.PythonTransformFunctionSerializer"] + + +def rddToFileName(prefix, suffix, timestamp): + """ + Return string prefix-time(.suffix) + + Examples + -------- + >>> rddToFileName("spark", None, 12345678910) + 'spark-12345678910' + >>> rddToFileName("spark", "tmp", 12345678910) + 'spark-12345678910.tmp' + """ + if isinstance(timestamp, datetime): + seconds = time.mktime(timestamp.timetuple()) + timestamp = int(seconds * 1000) + timestamp.microsecond // 1000 + if suffix is None: + return prefix + "-" + str(timestamp) + else: + return prefix + "-" + str(timestamp) + "." + suffix + + +if __name__ == "__main__": + import doctest + + (failure_count, test_count) = doctest.testmod() + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/taskcontext.py b/spark-3.5.0-bin-hadoop3/python/pyspark/taskcontext.py new file mode 100644 index 0000000..329820c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/taskcontext.py @@ -0,0 +1,502 @@ +# +# 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 ClassVar, Type, Dict, List, Optional, Union, cast + +from pyspark.java_gateway import local_connect_and_auth +from pyspark.resource import ResourceInformation +from pyspark.serializers import read_int, write_int, write_with_length, UTF8Deserializer +from pyspark.errors import PySparkRuntimeError + + +class TaskContext: + + """ + Contextual information about a task which can be read or mutated during + execution. To access the TaskContext for a running task, use: + :meth:`TaskContext.get`. + + .. versionadded:: 2.2.0 + + Examples + -------- + >>> from pyspark import TaskContext + + Get a task context instance from :class:`RDD`. + + >>> spark.sparkContext.setLocalProperty("key1", "value") + >>> taskcontext = spark.sparkContext.parallelize([1]).map(lambda _: TaskContext.get()).first() + >>> isinstance(taskcontext.attemptNumber(), int) + True + >>> isinstance(taskcontext.partitionId(), int) + True + >>> isinstance(taskcontext.stageId(), int) + True + >>> isinstance(taskcontext.taskAttemptId(), int) + True + >>> taskcontext.getLocalProperty("key1") + 'value' + >>> isinstance(taskcontext.cpus(), int) + True + + Get a task context instance from a dataframe via Python UDF. + + >>> from pyspark.sql import Row + >>> from pyspark.sql.functions import udf + >>> @udf("STRUCT") + ... def taskcontext_as_row(): + ... taskcontext = TaskContext.get() + ... return Row( + ... anum=taskcontext.attemptNumber(), + ... partid=taskcontext.partitionId(), + ... stageid=taskcontext.stageId(), + ... taskaid=taskcontext.taskAttemptId(), + ... prop=taskcontext.getLocalProperty("key2"), + ... cpus=taskcontext.cpus()) + ... + >>> spark.sparkContext.setLocalProperty("key2", "value") + >>> [(anum, partid, stageid, taskaid, prop, cpus)] = ( + ... spark.range(1).select(taskcontext_as_row()).first() + ... ) + >>> isinstance(anum, int) + True + >>> isinstance(partid, int) + True + >>> isinstance(stageid, int) + True + >>> isinstance(taskaid, int) + True + >>> prop + 'value' + >>> isinstance(cpus, int) + True + + Get a task context instance from a dataframe via Pandas UDF. + + >>> import pandas as pd # doctest: +SKIP + >>> from pyspark.sql.functions import pandas_udf + >>> @pandas_udf("STRUCT<" + ... "anum: INT, partid: INT, stageid: INT, taskaid: INT, prop: STRING, cpus: INT>") + ... def taskcontext_as_row(_): + ... taskcontext = TaskContext.get() + ... return pd.DataFrame({ + ... "anum": [taskcontext.attemptNumber()], + ... "partid": [taskcontext.partitionId()], + ... "stageid": [taskcontext.stageId()], + ... "taskaid": [taskcontext.taskAttemptId()], + ... "prop": [taskcontext.getLocalProperty("key3")], + ... "cpus": [taskcontext.cpus()] + ... }) # doctest: +SKIP + ... + >>> spark.sparkContext.setLocalProperty("key3", "value") # doctest: +SKIP + >>> [(anum, partid, stageid, taskaid, prop, cpus)] = ( + ... spark.range(1).select(taskcontext_as_row("id")).first() + ... ) # doctest: +SKIP + >>> isinstance(anum, int) + True + >>> isinstance(partid, int) + True + >>> isinstance(stageid, int) + True + >>> isinstance(taskaid, int) + True + >>> prop + 'value' + >>> isinstance(cpus, int) + True + """ + + _taskContext: ClassVar[Optional["TaskContext"]] = None + + _attemptNumber: Optional[int] = None + _partitionId: Optional[int] = None + _stageId: Optional[int] = None + _taskAttemptId: Optional[int] = None + _localProperties: Optional[Dict[str, str]] = None + _cpus: Optional[int] = None + _resources: Optional[Dict[str, ResourceInformation]] = None + + def __new__(cls: Type["TaskContext"]) -> "TaskContext": + """ + Even if users construct :class:`TaskContext` instead of using get, give them the singleton. + """ + taskContext = cls._taskContext + if taskContext is not None: + return taskContext + cls._taskContext = taskContext = object.__new__(cls) + return taskContext + + @classmethod + def _getOrCreate(cls: Type["TaskContext"]) -> "TaskContext": + """Internal function to get or create global :class:`TaskContext`.""" + if cls._taskContext is None: + cls._taskContext = TaskContext() + return cls._taskContext + + @classmethod + def _setTaskContext(cls: Type["TaskContext"], taskContext: "TaskContext") -> None: + cls._taskContext = taskContext + + @classmethod + def get(cls: Type["TaskContext"]) -> Optional["TaskContext"]: + """ + Return the currently active :class:`TaskContext`. This can be called inside of + user functions to access contextual information about running tasks. + + Returns + ------- + :class:`TaskContext`, optional + + Notes + ----- + Must be called on the worker, not the driver. Returns ``None`` if not initialized. + """ + return cls._taskContext + + def stageId(self) -> int: + """ + The ID of the stage that this task belong to. + + Returns + ------- + int + current stage id. + """ + return cast(int, self._stageId) + + def partitionId(self) -> int: + """ + The ID of the RDD partition that is computed by this task. + + Returns + ------- + int + current partition id. + """ + return cast(int, self._partitionId) + + def attemptNumber(self) -> int: + """ + How many times this task has been attempted. The first task attempt will be assigned + attemptNumber = 0, and subsequent attempts will have increasing attempt numbers. + + Returns + ------- + int + current attempt number. + """ + return cast(int, self._attemptNumber) + + def taskAttemptId(self) -> int: + """ + An ID that is unique to this task attempt (within the same :class:`SparkContext`, + no two task attempts will share the same attempt ID). This is roughly equivalent + to Hadoop's `TaskAttemptID`. + + Returns + ------- + int + current task attempt id. + """ + return cast(int, self._taskAttemptId) + + def getLocalProperty(self, key: str) -> Optional[str]: + """ + Get a local property set upstream in the driver, or None if it is missing. + + Parameters + ---------- + key : str + the key of the local property to get. + + Returns + ------- + int + the value of the local property. + """ + return cast(Dict[str, str], self._localProperties).get(key, None) + + def cpus(self) -> int: + """ + CPUs allocated to the task. + + Returns + ------- + int + the number of CPUs. + """ + return cast(int, self._cpus) + + def resources(self) -> Dict[str, ResourceInformation]: + """ + Resources allocated to the task. The key is the resource name and the value is information + about the resource. + + Returns + ------- + dict + a dictionary of a string resource name, and :class:`ResourceInformation`. + """ + return cast(Dict[str, ResourceInformation], self._resources) + + +BARRIER_FUNCTION = 1 +ALL_GATHER_FUNCTION = 2 + + +def _load_from_socket( + port: Optional[Union[str, int]], + auth_secret: str, + function: int, + all_gather_message: Optional[str] = None, +) -> List[str]: + """ + Load data from a given socket, this is a blocking method thus only return when the socket + connection has been closed. + """ + (sockfile, sock) = local_connect_and_auth(port, auth_secret) + + # The call may block forever, so no timeout + sock.settimeout(None) + + if function == BARRIER_FUNCTION: + # Make a barrier() function call. + write_int(function, sockfile) + elif function == ALL_GATHER_FUNCTION: + # Make a all_gather() function call. + write_int(function, sockfile) + write_with_length(cast(str, all_gather_message).encode("utf-8"), sockfile) + else: + raise ValueError("Unrecognized function type") + sockfile.flush() + + # Collect result. + len = read_int(sockfile) + res = [] + for i in range(len): + res.append(UTF8Deserializer().loads(sockfile)) + + # Release resources. + sockfile.close() + sock.close() + + return res + + +class BarrierTaskContext(TaskContext): + + """ + A :class:`TaskContext` with extra contextual info and tooling for tasks in a barrier stage. + Use :func:`BarrierTaskContext.get` to obtain the barrier context for a running barrier task. + + .. versionadded:: 2.4.0 + + Notes + ----- + This API is experimental + + Examples + -------- + Set a barrier, and execute it with RDD. + + >>> from pyspark import BarrierTaskContext + >>> def block_and_do_something(itr): + ... taskcontext = BarrierTaskContext.get() + ... # Do something. + ... + ... # Wait until all tasks finished. + ... taskcontext.barrier() + ... + ... return itr + ... + >>> rdd = spark.sparkContext.parallelize([1]) + >>> rdd.barrier().mapPartitions(block_and_do_something).collect() + [1] + """ + + _port: ClassVar[Optional[Union[str, int]]] = None + _secret: ClassVar[Optional[str]] = None + + @classmethod + def _getOrCreate(cls: Type["BarrierTaskContext"]) -> "BarrierTaskContext": + """ + Internal function to get or create global :class:`BarrierTaskContext`. We need to make sure + :class:`BarrierTaskContext` is returned from here because it is needed in python worker + reuse scenario, see SPARK-25921 for more details. + """ + if not isinstance(cls._taskContext, BarrierTaskContext): + cls._taskContext = object.__new__(cls) + return cls._taskContext + + @classmethod + def get(cls: Type["BarrierTaskContext"]) -> "BarrierTaskContext": + """ + Return the currently active :class:`BarrierTaskContext`. + This can be called inside of user functions to access contextual information about + running tasks. + + Notes + ----- + Must be called on the worker, not the driver. Returns ``None`` if not initialized. + An Exception will raise if it is not in a barrier stage. + + This API is experimental + """ + if not isinstance(cls._taskContext, BarrierTaskContext): + raise PySparkRuntimeError( + error_class="NOT_IN_BARRIER_STAGE", + message_parameters={}, + ) + return cls._taskContext + + @classmethod + def _initialize( + cls: Type["BarrierTaskContext"], port: Optional[Union[str, int]], secret: str + ) -> None: + """ + Initialize :class:`BarrierTaskContext`, other methods within :class:`BarrierTaskContext` + can only be called after BarrierTaskContext is initialized. + """ + cls._port = port + cls._secret = secret + + def barrier(self) -> None: + """ + Sets a global barrier and waits until all tasks in this stage hit this barrier. + Similar to `MPI_Barrier` function in MPI, this function blocks until all tasks + in the same stage have reached this routine. + + .. versionadded:: 2.4.0 + + Notes + ----- + This API is experimental + + In a barrier stage, each task much have the same number of `barrier()` + calls, in all possible code branches. Otherwise, you may get the job hanging + or a `SparkException` after timeout. + """ + if self._port is None or self._secret is None: + raise PySparkRuntimeError( + error_class="CALL_BEFORE_INITIALIZE", + message_parameters={ + "func_name": "barrier", + "object": "BarrierTaskContext", + }, + ) + else: + _load_from_socket(self._port, self._secret, BARRIER_FUNCTION) + + def allGather(self, message: str = "") -> List[str]: + """ + This function blocks until all tasks in the same stage have reached this routine. + Each task passes in a message and returns with a list of all the messages passed in + by each of those tasks. + + .. versionadded:: 3.0.0 + + Notes + ----- + This API is experimental + + In a barrier stage, each task much have the same number of `barrier()` + calls, in all possible code branches. Otherwise, you may get the job hanging + or a `SparkException` after timeout. + """ + if not isinstance(message, str): + raise TypeError("Argument `message` must be of type `str`") + elif self._port is None or self._secret is None: + raise PySparkRuntimeError( + error_class="CALL_BEFORE_INITIALIZE", + message_parameters={ + "func_name": "allGather", + "object": "BarrierTaskContext", + }, + ) + else: + return _load_from_socket(self._port, self._secret, ALL_GATHER_FUNCTION, message) + + def getTaskInfos(self) -> List["BarrierTaskInfo"]: + """ + Returns :class:`BarrierTaskInfo` for all tasks in this barrier stage, + ordered by partition ID. + + .. versionadded:: 2.4.0 + + Notes + ----- + This API is experimental + + Examples + -------- + >>> from pyspark import BarrierTaskContext + >>> rdd = spark.sparkContext.parallelize([1]) + >>> barrier_info = rdd.barrier().mapPartitions( + ... lambda _: [BarrierTaskContext.get().getTaskInfos()]).collect()[0][0] + >>> barrier_info.address + '...:...' + """ + if self._port is None or self._secret is None: + raise PySparkRuntimeError( + error_class="CALL_BEFORE_INITIALIZE", + message_parameters={ + "func_name": "getTaskInfos", + "object": "BarrierTaskContext", + }, + ) + else: + addresses = cast(Dict[str, str], self._localProperties).get("addresses", "") + return [BarrierTaskInfo(h.strip()) for h in addresses.split(",")] + + +class BarrierTaskInfo: + """ + Carries all task infos of a barrier task. + + .. versionadded:: 2.4.0 + + Attributes + ---------- + address : str + The IPv4 address (host:port) of the executor that the barrier task is running on + + Notes + ----- + This API is experimental + """ + + def __init__(self, address: str) -> None: + self.address = address + + +def _test() -> None: + import doctest + import sys + from pyspark.sql import SparkSession + + globs = globals().copy() + globs["spark"] = ( + SparkSession.builder.master("local[2]").appName("taskcontext tests").getOrCreate() + ) + (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) + globs["spark"].stop() + + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/__init__.py new file mode 100644 index 0000000..57c2066 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/__init__.py @@ -0,0 +1,21 @@ +# +# 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.testing.utils import assertDataFrameEqual, assertSchemaEqual + +from pyspark.testing.pandasutils import assertPandasOnSparkEqual + +__all__ = ["assertDataFrameEqual", "assertSchemaEqual", "assertPandasOnSparkEqual"] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/connectutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/connectutils.py new file mode 100644 index 0000000..ba81c78 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/connectutils.py @@ -0,0 +1,199 @@ +# +# 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 shutil +import tempfile +import types +import typing +import os +import functools +import unittest +import uuid + +from pyspark import Row, SparkConf +from pyspark.testing.utils import PySparkErrorTestUtils +from pyspark.testing.sqlutils import ( + have_pandas, + pandas_requirement_message, + pyarrow_requirement_message, + SQLTestUtils, +) +from pyspark.sql.session import SparkSession as PySparkSession + + +grpc_requirement_message = None +try: + import grpc +except ImportError as e: + grpc_requirement_message = str(e) +have_grpc = grpc_requirement_message is None + + +grpc_status_requirement_message = None +try: + import grpc_status +except ImportError as e: + grpc_status_requirement_message = str(e) +have_grpc_status = grpc_status_requirement_message is None + +googleapis_common_protos_requirement_message = None +try: + from google.rpc import error_details_pb2 +except ImportError as e: + googleapis_common_protos_requirement_message = str(e) +have_googleapis_common_protos = googleapis_common_protos_requirement_message is None + + +connect_requirement_message = ( + pandas_requirement_message + or pyarrow_requirement_message + or grpc_requirement_message + or googleapis_common_protos_requirement_message + or grpc_status_requirement_message +) +should_test_connect: str = typing.cast(str, connect_requirement_message is None) + +if should_test_connect: + from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.plan import Read, Range, SQL, LogicalPlan + from pyspark.sql.connect.session import SparkSession + + +class MockRemoteSession: + def __init__(self): + self.hooks = {} + self.session_id = str(uuid.uuid4()) + + def set_hook(self, name, hook): + self.hooks[name] = hook + + def drop_hook(self, name): + self.hooks.pop(name) + + def __getattr__(self, item): + if item not in self.hooks: + raise LookupError(f"{item} is not defined as a method hook in MockRemoteSession") + return functools.partial(self.hooks[item]) + + +class MockDF(DataFrame): + """Helper class that must only be used for the mock plan tests.""" + + def __init__(self, session: SparkSession, plan: LogicalPlan): + super().__init__(session) + self._plan = plan + + def __getattr__(self, name): + """All attributes are resolved to columns, because none really exist in the + mocked DataFrame.""" + return self[name] + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class PlanOnlyTestFixture(unittest.TestCase, PySparkErrorTestUtils): + @classmethod + def _read_table(cls, table_name): + return cls._df_mock(Read(table_name)) + + @classmethod + def _udf_mock(cls, *args, **kwargs): + return "internal_name" + + @classmethod + def _df_mock(cls, plan: LogicalPlan) -> MockDF: + return MockDF(cls.connect, plan) + + @classmethod + def _session_range( + cls, + start, + end, + step=1, + num_partitions=None, + ): + return cls._df_mock(Range(start, end, step, num_partitions)) + + @classmethod + def _session_sql(cls, query): + return cls._df_mock(SQL(query)) + + if have_pandas: + + @classmethod + def _with_plan(cls, plan): + return cls._df_mock(plan) + + @classmethod + def setUpClass(cls): + cls.connect = MockRemoteSession() + cls.session = SparkSession.builder.remote().getOrCreate() + cls.tbl_name = "test_connect_plan_only_table_1" + + cls.connect.set_hook("readTable", cls._read_table) + cls.connect.set_hook("range", cls._session_range) + cls.connect.set_hook("sql", cls._session_sql) + cls.connect.set_hook("with_plan", cls._with_plan) + + @classmethod + def tearDownClass(cls): + cls.connect.drop_hook("readTable") + cls.connect.drop_hook("range") + cls.connect.drop_hook("sql") + cls.connect.drop_hook("with_plan") + + +@unittest.skipIf(not should_test_connect, connect_requirement_message) +class ReusedConnectTestCase(unittest.TestCase, SQLTestUtils, PySparkErrorTestUtils): + """ + Spark Connect version of :class:`pyspark.testing.sqlutils.ReusedSQLTestCase`. + """ + + @classmethod + def conf(cls): + """ + Override this in subclasses to supply a more specific conf + """ + conf = SparkConf(loadDefaults=False) + # Disable JVM stack trace in Spark Connect tests to prevent the + # HTTP header size from exceeding the maximum allowed size. + conf.set("spark.sql.pyspark.jvmStacktrace.enabled", "false") + # Make the server terminate reattachable streams every 1 second and 123 bytes, + # to make the tests exercise reattach. + conf.set("spark.connect.execute.reattachable.senderMaxStreamDuration", "1s") + conf.set("spark.connect.execute.reattachable.senderMaxStreamSize", "123") + return conf + + @classmethod + def master(cls): + return "local[4]" + + @classmethod + def setUpClass(cls): + cls.spark = ( + PySparkSession.builder.config(conf=cls.conf()) + .appName(cls.__name__) + .remote(cls.master()) + .getOrCreate() + ) + cls.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(cls.tempdir.name) + cls.testData = [Row(key=i, value=str(i)) for i in range(100)] + cls.df = cls.spark.createDataFrame(cls.testData) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tempdir.name, ignore_errors=True) + cls.spark.stop() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mllibutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mllibutils.py new file mode 100644 index 0000000..3edb7f0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mllibutils.py @@ -0,0 +1,30 @@ +# +# 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 unittest + +from pyspark import SparkContext +from pyspark.sql import SparkSession + + +class MLlibTestCase(unittest.TestCase): + def setUp(self): + self.sc = SparkContext("local[4]", "MLlib tests") + self.spark = SparkSession(self.sc) + + def tearDown(self): + self.spark.stop() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mlutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mlutils.py new file mode 100644 index 0000000..503ba7c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/mlutils.py @@ -0,0 +1,275 @@ +# +# 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 numpy as np + +from pyspark import keyword_only +from pyspark.ml import Estimator, Model, Transformer, UnaryTransformer +from pyspark.ml.evaluation import Evaluator +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import HasMaxIter, HasRegParam +from pyspark.ml.classification import Classifier, ClassificationModel +from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWritable +from pyspark.ml.wrapper import _java2py +from pyspark.sql import DataFrame, SparkSession +from pyspark.sql.types import DoubleType +from pyspark.testing.utils import ReusedPySparkTestCase as PySparkTestCase + + +def check_params(test_self, py_stage, check_params_exist=True): + """ + Checks common requirements for :py:class:`PySpark.ml.Params.params`: + + - set of params exist in Java and Python and are ordered by names + - param parent has the same UID as the object's UID + - default param value from Java matches value in Python + - optionally check if all params from Java also exist in Python + """ + py_stage_str = "%s %s" % (type(py_stage), py_stage) + if not hasattr(py_stage, "_to_java"): + return + java_stage = py_stage._to_java() + if java_stage is None: + return + test_self.assertEqual(py_stage.uid, java_stage.uid(), msg=py_stage_str) + if check_params_exist: + param_names = [p.name for p in py_stage.params] + java_params = list(java_stage.params()) + java_param_names = [jp.name() for jp in java_params] + test_self.assertEqual( + param_names, + sorted(java_param_names), + "Param list in Python does not match Java for %s:\nJava = %s\nPython = %s" + % (py_stage_str, java_param_names, param_names), + ) + for p in py_stage.params: + test_self.assertEqual(p.parent, py_stage.uid) + java_param = java_stage.getParam(p.name) + py_has_default = py_stage.hasDefault(p) + java_has_default = java_stage.hasDefault(java_param) + test_self.assertEqual( + py_has_default, + java_has_default, + "Default value mismatch of param %s for Params %s" % (p.name, str(py_stage)), + ) + if py_has_default: + if p.name == "seed": + continue # Random seeds between Spark and PySpark are different + java_default = _java2py( + test_self.sc, java_stage.clear(java_param).getOrDefault(java_param) + ) + py_stage.clear(p) + py_default = py_stage.getOrDefault(p) + # equality test for NaN is always False + if isinstance(java_default, float) and np.isnan(java_default): + java_default = "NaN" + py_default = "NaN" if np.isnan(py_default) else "not NaN" + test_self.assertEqual( + java_default, + py_default, + "Java default %s != python default %s of param %s for Params %s" + % (str(java_default), str(py_default), p.name, str(py_stage)), + ) + + +class SparkSessionTestCase(PySparkTestCase): + @classmethod + def setUpClass(cls): + PySparkTestCase.setUpClass() + cls.spark = SparkSession(cls.sc) + + @classmethod + def tearDownClass(cls): + PySparkTestCase.tearDownClass() + cls.spark.stop() + + +class MockDataset(DataFrame): + def __init__(self): + self.index = 0 + + +class HasFake(Params): + def __init__(self): + super(HasFake, self).__init__() + self.fake = Param(self, "fake", "fake param") + + def getFake(self): + return self.getOrDefault(self.fake) + + +class MockTransformer(Transformer, HasFake): + def __init__(self): + super(MockTransformer, self).__init__() + self.dataset_index = None + + def _transform(self, dataset): + self.dataset_index = dataset.index + dataset.index += 1 + return dataset + + +class MockUnaryTransformer(UnaryTransformer, DefaultParamsReadable, DefaultParamsWritable): + + shift = Param( + Params._dummy(), + "shift", + "The amount by which to shift " + "data in a DataFrame", + typeConverter=TypeConverters.toFloat, + ) + + def __init__(self, shiftVal=1): + super(MockUnaryTransformer, self).__init__() + self._setDefault(shift=1) + self._set(shift=shiftVal) + + def getShift(self): + return self.getOrDefault(self.shift) + + def setShift(self, shift): + self._set(shift=shift) + + def createTransformFunc(self): + shiftVal = self.getShift() + return lambda x: x + shiftVal + + def outputDataType(self): + return DoubleType() + + def validateInputType(self, inputType): + if inputType != DoubleType(): + raise TypeError("Bad input type: {}. ".format(inputType) + "Requires Double.") + + +class MockEstimator(Estimator, HasFake): + def __init__(self): + super(MockEstimator, self).__init__() + self.dataset_index = None + + def _fit(self, dataset): + self.dataset_index = dataset.index + model = MockModel() + self._copyValues(model) + return model + + +class MockModel(MockTransformer, Model, HasFake): + pass + + +class _DummyLogisticRegressionParams(HasMaxIter, HasRegParam): + def setMaxIter(self, value): + return self._set(maxIter=value) + + def setRegParam(self, value): + return self._set(regParam=value) + + +# This is a dummy LogisticRegression used in test for python backend estimator/model +class DummyLogisticRegression( + Classifier, _DummyLogisticRegressionParams, DefaultParamsReadable, DefaultParamsWritable +): + @keyword_only + def __init__( + self, + *, + featuresCol="features", + labelCol="label", + predictionCol="prediction", + maxIter=100, + regParam=0.0, + rawPredictionCol="rawPrediction", + ): + super(DummyLogisticRegression, self).__init__() + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @keyword_only + def setParams( + self, + *, + featuresCol="features", + labelCol="label", + predictionCol="prediction", + maxIter=100, + regParam=0.0, + rawPredictionCol="rawPrediction", + ): + kwargs = self._input_kwargs + self._set(**kwargs) + return self + + def _fit(self, dataset): + # Do nothing but create a dummy model + return self._copyValues(DummyLogisticRegressionModel()) + + +class DummyLogisticRegressionModel( + ClassificationModel, + _DummyLogisticRegressionParams, + DefaultParamsReadable, + DefaultParamsWritable, +): + def __init__(self): + super(DummyLogisticRegressionModel, self).__init__() + + def _transform(self, dataset): + # A dummy transform impl which always predict label 1 + from pyspark.sql.functions import array, lit + from pyspark.ml.functions import array_to_vector + + rawPredCol = self.getRawPredictionCol() + if rawPredCol: + dataset = dataset.withColumn( + rawPredCol, array_to_vector(array(lit(-100.0), lit(100.0))) + ) + predCol = self.getPredictionCol() + if predCol: + dataset = dataset.withColumn(predCol, lit(1.0)) + + return dataset + + @property + def numClasses(self): + # a dummy implementation for test. + return 2 + + @property + def intercept(self): + # a dummy implementation for test. + return 0.0 + + # This class only used in test. The following methods/properties are not used in tests. + + @property + def coefficients(self): + raise NotImplementedError() + + def predictRaw(self, value): + raise NotImplementedError() + + def numFeatures(self): + raise NotImplementedError() + + def predict(self, value): + raise NotImplementedError() + + +class DummyEvaluator(Evaluator, DefaultParamsReadable, DefaultParamsWritable): + def _evaluate(self, dataset): + # a dummy implementation for test. + return 1.0 diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/pandasutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/pandasutils.py new file mode 100644 index 0000000..c80ffb7 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/pandasutils.py @@ -0,0 +1,707 @@ +# +# 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 functools +import shutil +import tempfile +import warnings +from contextlib import contextmanager +from distutils.version import LooseVersion +import decimal +from typing import Any, Union, TYPE_CHECKING + +import pyspark.pandas as ps +from pyspark.pandas.frame import DataFrame +from pyspark.pandas.indexes import Index +from pyspark.pandas.series import Series +from pyspark.pandas.utils import SPARK_CONF_ARROW_ENABLED +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.errors import PySparkAssertionError + +tabulate_requirement_message = None +try: + from tabulate import tabulate +except ImportError as e: + # If tabulate requirement is not satisfied, skip related tests. + tabulate_requirement_message = str(e) +have_tabulate = tabulate_requirement_message is None + +matplotlib_requirement_message = None +try: + import matplotlib +except ImportError as e: + # If matplotlib requirement is not satisfied, skip related tests. + matplotlib_requirement_message = str(e) +have_matplotlib = matplotlib_requirement_message is None + +plotly_requirement_message = None +try: + import plotly +except ImportError as e: + # If plotly requirement is not satisfied, skip related tests. + plotly_requirement_message = str(e) +have_plotly = plotly_requirement_message is None + +try: + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() + import pandas as pd +except ImportError: + pass + +__all__ = ["assertPandasOnSparkEqual"] + + +def _assert_pandas_equal( + left: Union[pd.DataFrame, pd.Series, pd.Index], + right: Union[pd.DataFrame, pd.Series, pd.Index], + checkExact: bool, +): + from pandas.core.dtypes.common import is_numeric_dtype + from pandas.testing import assert_frame_equal, assert_index_equal, assert_series_equal + + if isinstance(left, pd.DataFrame) and isinstance(right, pd.DataFrame): + try: + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + kwargs = dict(check_freq=False) + else: + kwargs = dict() + + if LooseVersion(pd.__version__) < LooseVersion("1.1.1"): + # Due to https://github.com/pandas-dev/pandas/issues/35446 + checkExact = ( + checkExact + and all([is_numeric_dtype(dtype) for dtype in left.dtypes]) + and all([is_numeric_dtype(dtype) for dtype in right.dtypes]) + ) + + assert_frame_equal( + left, + right, + check_index_type=("equiv" if len(left.index) > 0 else False), + check_column_type=("equiv" if len(left.columns) > 0 else False), + check_exact=checkExact, + **kwargs, + ) + except AssertionError: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + elif isinstance(left, pd.Series) and isinstance(right, pd.Series): + try: + if LooseVersion(pd.__version__) >= LooseVersion("1.1"): + kwargs = dict(check_freq=False) + else: + kwargs = dict() + if LooseVersion(pd.__version__) < LooseVersion("1.1.1"): + # Due to https://github.com/pandas-dev/pandas/issues/35446 + checkExact = ( + checkExact and is_numeric_dtype(left.dtype) and is_numeric_dtype(right.dtype) + ) + assert_series_equal( + left, + right, + check_index_type=("equiv" if len(left.index) > 0 else False), + check_exact=checkExact, + **kwargs, + ) + except AssertionError: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_SERIES", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtype), + "right": right.to_string(), + "right_dtype": str(right.dtype), + }, + ) + elif isinstance(left, pd.Index) and isinstance(right, pd.Index): + try: + if LooseVersion(pd.__version__) < LooseVersion("1.1.1"): + # Due to https://github.com/pandas-dev/pandas/issues/35446 + checkExact = ( + checkExact and is_numeric_dtype(left.dtype) and is_numeric_dtype(right.dtype) + ) + assert_index_equal(left, right, check_exact=checkExact) + except AssertionError: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_INDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + else: + raise ValueError("Unexpected values: (%s, %s)" % (left, right)) + + +def _assert_pandas_almost_equal( + left: Union[pd.DataFrame, pd.Series, pd.Index], + right: Union[pd.DataFrame, pd.Series, pd.Index], + rtol: float = 1e-5, + atol: float = 1e-8, +): + """ + This function checks if given pandas objects approximately same, + which means the conditions below: + - Both objects are nullable + - Compare decimals and floats, where two values a and b are approximately equal + if they satisfy the following formula: + absolute(a - b) <= (atol + rtol * absolute(b)) + where rtol=1e-5 and atol=1e-8 by default + """ + + def compare_vals_approx(val1, val2): + # compare vals for approximate equality + if isinstance(lval, (float, decimal.Decimal)) or isinstance(rval, (float, decimal.Decimal)): + if abs(float(lval) - float(rval)) > (atol + rtol * abs(float(rval))): + return False + elif val1 != val2: + return False + return True + + if isinstance(left, pd.DataFrame) and isinstance(right, pd.DataFrame): + if left.shape != right.shape: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + for lcol, rcol in zip(left.columns, right.columns): + if lcol != rcol: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + for lnull, rnull in zip(left[lcol].isnull(), right[rcol].isnull()): + if lnull != rnull: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + for lval, rval in zip(left[lcol].dropna(), right[rcol].dropna()): + if not compare_vals_approx(lval, rval): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + if left.columns.names != right.columns.names: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_DATAFRAME", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtypes), + "right": right.to_string(), + "right_dtype": str(right.dtypes), + }, + ) + elif isinstance(left, pd.Series) and isinstance(right, pd.Series): + if left.name != right.name or len(left) != len(right): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_SERIES", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtype), + "right": right.to_string(), + "right_dtype": str(right.dtype), + }, + ) + for lnull, rnull in zip(left.isnull(), right.isnull()): + if lnull != rnull: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_SERIES", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtype), + "right": right.to_string(), + "right_dtype": str(right.dtype), + }, + ) + for lval, rval in zip(left.dropna(), right.dropna()): + if not compare_vals_approx(lval, rval): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_SERIES", + message_parameters={ + "left": left.to_string(), + "left_dtype": str(left.dtype), + "right": right.to_string(), + "right_dtype": str(right.dtype), + }, + ) + elif isinstance(left, pd.MultiIndex) and isinstance(right, pd.MultiIndex): + if len(left) != len(right): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_MULTIINDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + for lval, rval in zip(left, right): + if not compare_vals_approx(lval, rval): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_MULTIINDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + elif isinstance(left, pd.Index) and isinstance(right, pd.Index): + if len(left) != len(right): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_INDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + for lnull, rnull in zip(left.isnull(), right.isnull()): + if lnull != rnull: + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_INDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + for lval, rval in zip(left.dropna(), right.dropna()): + if not compare_vals_approx(lval, rval): + raise PySparkAssertionError( + error_class="DIFFERENT_PANDAS_INDEX", + message_parameters={ + "left": left, + "left_dtype": str(left.dtype), + "right": right, + "right_dtype": str(right.dtype), + }, + ) + else: + if not isinstance(left, (pd.DataFrame, pd.Series, pd.Index)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{pd.DataFrame.__name__}, " + f"{pd.Series.__name__}, " + f"{pd.Index.__name__}, ", + "arg_name": "left", + "actual_type": type(left), + }, + ) + elif not isinstance(right, (pd.DataFrame, pd.Series, pd.Index)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{pd.DataFrame.__name__}, " + f"{pd.Series.__name__}, " + f"{pd.Index.__name__}, ", + "arg_name": "right", + "actual_type": type(right), + }, + ) + + +def assertPandasOnSparkEqual( + actual: Union[DataFrame, Series, Index], + expected: Union[DataFrame, pd.DataFrame, Series, pd.Series, Index, pd.Index], + checkExact: bool = True, + almost: bool = False, + rtol: float = 1e-5, + atol: float = 1e-8, + checkRowOrder: bool = True, +): + r""" + A util function to assert equality between actual (pandas-on-Spark object) and expected + (pandas-on-Spark or pandas object). + + .. versionadded:: 3.5.0 + + Parameters + ---------- + actual: pandas-on-Spark DataFrame, Series, or Index + The object that is being compared or tested. + expected: pandas-on-Spark or pandas DataFrame, Series, or Index + The expected object, for comparison with the actual result. + checkExact: bool, optional + A flag indicating whether to compare exact equality. + If set to 'True' (default), the data is compared exactly. + If set to 'False', the data is compared less precisely, following pandas assert_frame_equal + approximate comparison (see documentation for more details). + almost: bool, optional + A flag indicating whether to use unittest `assertAlmostEqual` or `assertEqual`. + If set to 'True', the comparison is delegated to `unittest`'s `assertAlmostEqual` + (see documentation for more details). + If set to 'False' (default), the data is compared exactly with `unittest`'s + `assertEqual`. + rtol : float, optional + The relative tolerance, used in asserting almost equality for float values in actual + and expected. Set to 1e-5 by default. (See Notes) + atol : float, optional + The absolute tolerance, used in asserting almost equality for float values in actual + and expected. Set to 1e-8 by default. (See Notes) + checkRowOrder : bool, optional + A flag indicating whether the order of rows should be considered in the comparison. + If set to `False`, the row order is not taken into account. + If set to `True` (default), the order of rows will be checked during comparison. + (See Notes) + + Notes + ----- + For `checkRowOrder`, note that pandas-on-Spark DataFrame ordering is non-deterministic, unless + explicitly sorted. + + When `almost` is set to True, approximate equality will be asserted, where two values + a and b are approximately equal if they satisfy the following formula: + + ``absolute(a - b) <= (atol + rtol * absolute(b))``. + + Examples + -------- + >>> import pyspark.pandas as ps + >>> psdf1 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> psdf2 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> assertPandasOnSparkEqual(psdf1, psdf2) # pass, ps.DataFrames are equal + >>> s1 = ps.Series([212.32, 100.0001]) + >>> s2 = ps.Series([212.32, 100.0]) + >>> assertPandasOnSparkEqual(s1, s2, checkExact=False) # pass, ps.Series are approx equal + >>> s1 = ps.Index([212.300001, 100.000]) + >>> s2 = ps.Index([212.3, 100.0001]) + >>> assertPandasOnSparkEqual(s1, s2, almost=True) # pass, ps.Index obj are almost equal + """ + if actual is None and expected is None: + return True + elif actual is None or expected is None: + return False + + if not isinstance(actual, (DataFrame, Series, Index)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{DataFrame.__name__}, {Series.__name__}, {Index.__name__}", + "arg_name": "actual", + "actual_type": type(actual), + }, + ) + elif not isinstance(expected, (DataFrame, pd.DataFrame, Series, pd.Series, Index, pd.Index)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": f"{DataFrame.__name__}, " + f"{pd.DataFrame.__name__}, " + f"{Series.__name__}, " + f"{pd.Series.__name__}, " + f"{Index.__name__}" + f"{pd.Index.__name__}, ", + "arg_name": "expected", + "actual_type": type(expected), + }, + ) + else: + if not isinstance(actual, (pd.DataFrame, pd.Index, pd.Series)): + actual = actual.to_pandas() + if not isinstance(expected, (pd.DataFrame, pd.Index, pd.Series)): + expected = expected.to_pandas() + + if not checkRowOrder: + if isinstance(actual, pd.DataFrame) and len(actual.columns) > 0: + actual = actual.sort_values(by=actual.columns[0], ignore_index=True) + if isinstance(expected, pd.DataFrame) and len(expected.columns) > 0: + expected = expected.sort_values(by=expected.columns[0], ignore_index=True) + + if almost: + _assert_pandas_almost_equal(actual, expected, rtol=rtol, atol=atol) + else: + _assert_pandas_equal(actual, expected, checkExact=checkExact) + + +class PandasOnSparkTestUtils: + def convert_str_to_lambda(self, func: str): + """ + This function converts `func` str to lambda call + """ + return lambda x: getattr(x, func)() + + def assertPandasEqual(self, left: Any, right: Any, check_exact: bool = True): + _assert_pandas_equal(left, right, check_exact) + + def assertPandasAlmostEqual( + self, + left: Any, + right: Any, + rtol: float = 1e-5, + atol: float = 1e-8, + ): + _assert_pandas_almost_equal(left, right, rtol=rtol, atol=atol) + + def assert_eq( + self, + left: Any, + right: Any, + check_exact: bool = True, + almost: bool = False, + rtol: float = 1e-5, + atol: float = 1e-8, + check_row_order: bool = True, + ): + """ + Asserts if two arbitrary objects are equal or not. If given objects are Koalas DataFrame + or Series, they are converted into pandas' and compared. + + :param left: object to compare + :param right: object to compare + :param check_exact: if this is False, the comparison is done less precisely. + :param almost: if this is enabled, the comparison asserts approximate equality + for float and decimal values, where two values a and b are approximately equal + if they satisfy the following formula: + absolute(a - b) <= (atol + rtol * absolute(b)) + :param rtol: The relative tolerance, used in asserting approximate equality for + float values. Set to 1e-5 by default. + :param atol: The absolute tolerance, used in asserting approximate equality for + float values in actual and expected. Set to 1e-8 by default. + :param check_row_order: A flag indicating whether the order of rows should be considered + in the comparison. If set to False, row order will be ignored. + """ + import pandas as pd + from pandas.api.types import is_list_like + + # for pandas-on-Spark DataFrames, allow choice to ignore row order + if isinstance(left, (ps.DataFrame, ps.Series, ps.Index)): + return assertPandasOnSparkEqual( + left, + right, + checkExact=check_exact, + almost=almost, + rtol=rtol, + atol=atol, + checkRowOrder=check_row_order, + ) + + lobj = self._to_pandas(left) + robj = self._to_pandas(right) + if isinstance(lobj, (pd.DataFrame, pd.Series, pd.Index)): + if almost: + _assert_pandas_almost_equal(lobj, robj, rtol=rtol, atol=atol) + else: + _assert_pandas_equal(lobj, robj, checkExact=check_exact) + elif is_list_like(lobj) and is_list_like(robj): + self.assertTrue(len(left) == len(right)) + for litem, ritem in zip(left, right): + self.assert_eq(litem, ritem, check_exact=check_exact, almost=almost) + elif (lobj is not None and pd.isna(lobj)) and (robj is not None and pd.isna(robj)): + pass + else: + if almost: + self.assertAlmostEqual(lobj, robj) + else: + self.assertEqual(lobj, robj) + + @staticmethod + def _to_pandas(obj: Any): + if isinstance(obj, (DataFrame, Series, Index)): + return obj.to_pandas() + else: + return obj + + +class PandasOnSparkTestCase(ReusedSQLTestCase, PandasOnSparkTestUtils): + @classmethod + def setUpClass(cls): + super(PandasOnSparkTestCase, cls).setUpClass() + cls.spark.conf.set(SPARK_CONF_ARROW_ENABLED, True) + + +class TestUtils: + @contextmanager + def temp_dir(self): + tmp = tempfile.mkdtemp() + try: + yield tmp + finally: + shutil.rmtree(tmp) + + @contextmanager + def temp_file(self): + with self.temp_dir() as tmp: + yield tempfile.mkstemp(dir=tmp)[1] + + +class ComparisonTestBase(PandasOnSparkTestCase): + @property + def psdf(self): + return ps.from_pandas(self.pdf) + + @property + def pdf(self): + return self.psdf.to_pandas() + + +def compare_both(f=None, almost=True): + + if f is None: + return functools.partial(compare_both, almost=almost) + elif isinstance(f, bool): + return functools.partial(compare_both, almost=f) + + @functools.wraps(f) + def wrapped(self): + if almost: + compare = self.assertPandasAlmostEqual + else: + compare = self.assertPandasEqual + + for result_pandas, result_spark in zip(f(self, self.pdf), f(self, self.psdf)): + compare(result_pandas, result_spark.to_pandas()) + + return wrapped + + +@contextmanager +def assert_produces_warning( + expected_warning=Warning, + filter_level="always", + check_stacklevel=True, + raise_on_extra_warnings=True, +): + """ + Context manager for running code expected to either raise a specific + warning, or not raise any warnings. Verifies that the code raises the + expected warning, and that it does not raise any other unexpected + warnings. It is basically a wrapper around ``warnings.catch_warnings``. + + Notes + ----- + Replicated from pandas/_testing/_warnings.py. + + Parameters + ---------- + expected_warning : {Warning, False, None}, default Warning + The type of Exception raised. ``exception.Warning`` is the base + class for all warnings. To check that no warning is returned, + specify ``False`` or ``None``. + filter_level : str or None, default "always" + Specifies whether warnings are ignored, displayed, or turned + into errors. + Valid values are: + * "error" - turns matching warnings into exceptions + * "ignore" - discard the warning + * "always" - always emit a warning + * "default" - print the warning the first time it is generated + from each location + * "module" - print the warning the first time it is generated + from each module + * "once" - print the warning the first time it is generated + check_stacklevel : bool, default True + If True, displays the line that called the function containing + the warning to show were the function is called. Otherwise, the + line that implements the function is displayed. + raise_on_extra_warnings : bool, default True + Whether extra warnings not of the type `expected_warning` should + cause the test to fail. + + Examples + -------- + >>> import warnings + >>> with assert_produces_warning(): + ... warnings.warn(UserWarning()) + ... + >>> with assert_produces_warning(False): # doctest: +SKIP + ... warnings.warn(RuntimeWarning()) + ... + Traceback (most recent call last): + ... + AssertionError: Caused unexpected warning(s): ['RuntimeWarning']. + >>> with assert_produces_warning(UserWarning): # doctest: +SKIP + ... warnings.warn(RuntimeWarning()) + Traceback (most recent call last): + ... + AssertionError: Did not see expected warning of class 'UserWarning' + ..warn:: This is *not* thread-safe. + """ + __tracebackhide__ = True + + with warnings.catch_warnings(record=True) as w: + + saw_warning = False + warnings.simplefilter(filter_level) + yield w + extra_warnings = [] + + for actual_warning in w: + if expected_warning and issubclass(actual_warning.category, expected_warning): + saw_warning = True + + if check_stacklevel and issubclass( + actual_warning.category, (FutureWarning, DeprecationWarning) + ): + from inspect import getframeinfo, stack + + caller = getframeinfo(stack()[2][0]) + msg = ( + "Warning not set with correct stacklevel. ", + "File where warning is raised: {} != ".format(actual_warning.filename), + "{}. Warning message: {}".format(caller.filename, actual_warning.message), + ) + assert actual_warning.filename == caller.filename, msg + else: + extra_warnings.append( + ( + actual_warning.category.__name__, + actual_warning.message, + actual_warning.filename, + actual_warning.lineno, + ) + ) + if expected_warning: + msg = "Did not see expected warning of class {}".format(repr(expected_warning.__name__)) + assert saw_warning, msg + if raise_on_extra_warnings and extra_warnings: + raise AssertionError("Caused unexpected warning(s): {}".format(repr(extra_warnings))) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/sqlutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/sqlutils.py new file mode 100644 index 0000000..077d854 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/sqlutils.py @@ -0,0 +1,271 @@ +# +# 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 datetime +import math +import os +import shutil +import tempfile +from contextlib import contextmanager + +from pyspark.sql import SparkSession +from pyspark.sql.types import ArrayType, DoubleType, UserDefinedType, Row +from pyspark.testing.utils import ReusedPySparkTestCase, PySparkErrorTestUtils + + +pandas_requirement_message = None +try: + from pyspark.sql.pandas.utils import require_minimum_pandas_version + + require_minimum_pandas_version() +except ImportError as e: + # If Pandas version requirement is not satisfied, skip related tests. + pandas_requirement_message = str(e) + +pyarrow_requirement_message = None +try: + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + + require_minimum_pyarrow_version() +except ImportError as e: + # If Arrow version requirement is not satisfied, skip related tests. + pyarrow_requirement_message = str(e) + +test_not_compiled_message = None +try: + from pyspark.sql.utils import require_test_compiled + + require_test_compiled() +except Exception as e: + test_not_compiled_message = str(e) + +have_pandas = pandas_requirement_message is None +have_pyarrow = pyarrow_requirement_message is None +test_compiled = test_not_compiled_message is None + + +class UTCOffsetTimezone(datetime.tzinfo): + """ + Specifies timezone in UTC offset + """ + + def __init__(self, offset=0): + self.ZERO = datetime.timedelta(hours=offset) + + def utcoffset(self, dt): + return self.ZERO + + def dst(self, dt): + return self.ZERO + + +class ExamplePointUDT(UserDefinedType): + """ + User-defined type (UDT) for ExamplePoint. + """ + + @classmethod + def sqlType(cls): + return ArrayType(DoubleType(), False) + + @classmethod + def module(cls): + return "pyspark.sql.tests" + + @classmethod + def scalaUDT(cls): + return "org.apache.spark.sql.test.ExamplePointUDT" + + def serialize(self, obj): + return [obj.x, obj.y] + + def deserialize(self, datum): + return ExamplePoint(datum[0], datum[1]) + + +class ExamplePoint: + """ + An example class to demonstrate UDT in Scala, Java, and Python. + """ + + __UDT__ = ExamplePointUDT() + + def __init__(self, x, y): + self.x = x + self.y = y + + def __repr__(self): + return "ExamplePoint(%s,%s)" % (self.x, self.y) + + def __str__(self): + return "(%s,%s)" % (self.x, self.y) + + def __eq__(self, other): + return isinstance(other, self.__class__) and other.x == self.x and other.y == self.y + + +class PythonOnlyUDT(UserDefinedType): + """ + User-defined type (UDT) for ExamplePoint. + """ + + @classmethod + def sqlType(cls): + return ArrayType(DoubleType(), False) + + @classmethod + def module(cls): + return "__main__" + + def serialize(self, obj): + return [obj.x, obj.y] + + def deserialize(self, datum): + return PythonOnlyPoint(datum[0], datum[1]) + + @staticmethod + def foo(): + pass + + @property + def props(self): + return {} + + +class PythonOnlyPoint(ExamplePoint): + """ + An example class to demonstrate UDT in only Python + """ + + __UDT__ = PythonOnlyUDT() # type: ignore + + +class MyObject: + def __init__(self, key, value): + self.key = key + self.value = value + + +class SQLTestUtils: + """ + This util assumes the instance of this to have 'spark' attribute, having a spark session. + It is usually used with 'ReusedSQLTestCase' class but can be used if you feel sure the + the implementation of this class has 'spark' attribute. + """ + + @contextmanager + def sql_conf(self, pairs): + """ + A convenient context manager to test some configuration specific logic. This sets + `value` to the configuration `key` and then restores it back when it exits. + """ + assert isinstance(pairs, dict), "pairs should be a dictionary." + assert hasattr(self, "spark"), "it should have 'spark' attribute, having a spark session." + + keys = pairs.keys() + new_values = pairs.values() + old_values = [self.spark.conf.get(key, None) for key in keys] + for key, new_value in zip(keys, new_values): + self.spark.conf.set(key, new_value) + try: + yield + finally: + for key, old_value in zip(keys, old_values): + if old_value is None: + self.spark.conf.unset(key) + else: + self.spark.conf.set(key, old_value) + + @contextmanager + def database(self, *databases): + """ + A convenient context manager to test with some specific databases. This drops the given + databases if it exists and sets current database to "default" when it exits. + """ + assert hasattr(self, "spark"), "it should have 'spark' attribute, having a spark session." + + try: + yield + finally: + for db in databases: + self.spark.sql("DROP DATABASE IF EXISTS %s CASCADE" % db) + self.spark.catalog.setCurrentDatabase("default") + + @contextmanager + def table(self, *tables): + """ + A convenient context manager to test with some specific tables. This drops the given tables + if it exists. + """ + assert hasattr(self, "spark"), "it should have 'spark' attribute, having a spark session." + + try: + yield + finally: + for t in tables: + self.spark.sql("DROP TABLE IF EXISTS %s" % t) + + @contextmanager + def tempView(self, *views): + """ + A convenient context manager to test with some specific views. This drops the given views + if it exists. + """ + assert hasattr(self, "spark"), "it should have 'spark' attribute, having a spark session." + + try: + yield + finally: + for v in views: + self.spark.catalog.dropTempView(v) + + @contextmanager + def function(self, *functions): + """ + A convenient context manager to test with some specific functions. This drops the given + functions if it exists. + """ + assert hasattr(self, "spark"), "it should have 'spark' attribute, having a spark session." + + try: + yield + finally: + for f in functions: + self.spark.sql("DROP FUNCTION IF EXISTS %s" % f) + + @staticmethod + def assert_close(a, b): + c = [j[0] for j in b] + diff = [abs(v - c[k]) < 1e-6 if math.isfinite(v) else v == c[k] for k, v in enumerate(a)] + assert sum(diff) == len(a), f"sum: {sum(diff)}, len: {len(a)}" + + +class ReusedSQLTestCase(ReusedPySparkTestCase, SQLTestUtils, PySparkErrorTestUtils): + @classmethod + def setUpClass(cls): + super(ReusedSQLTestCase, cls).setUpClass() + cls.spark = SparkSession(cls.sc) + cls.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(cls.tempdir.name) + cls.testData = [Row(key=i, value=str(i)) for i in range(100)] + cls.df = cls.spark.createDataFrame(cls.testData) + + @classmethod + def tearDownClass(cls): + super(ReusedSQLTestCase, cls).tearDownClass() + cls.spark.stop() + shutil.rmtree(cls.tempdir.name, ignore_errors=True) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/streamingutils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/streamingutils.py new file mode 100644 index 0000000..57f27ce --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/streamingutils.py @@ -0,0 +1,178 @@ +# +# 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 os +import tempfile +import time +import unittest + +from pyspark import SparkConf, SparkContext, RDD +from pyspark.streaming import StreamingContext +from pyspark.testing.utils import search_jar + + +# Must be same as the variable and condition defined in KinesisTestUtils.scala and modules.py +kinesis_test_environ_var = "ENABLE_KINESIS_TESTS" +should_skip_kinesis_tests = not os.environ.get(kinesis_test_environ_var) == "1" + +if should_skip_kinesis_tests: + kinesis_requirement_message = ( + "Skipping all Kinesis Python tests as environmental variable 'ENABLE_KINESIS_TESTS' " + "was not set." + ) +else: + kinesis_asl_assembly_jar = search_jar( + "connector/kinesis-asl-assembly", + "spark-streaming-kinesis-asl-assembly-", + "spark-streaming-kinesis-asl-assembly_", + ) + if kinesis_asl_assembly_jar is None: + kinesis_requirement_message = ( + "Skipping all Kinesis Python tests as the optional Kinesis project was " + "not compiled into a JAR. To run these tests, " + "you need to build Spark with 'build/sbt -Pkinesis-asl assembly/package " + "streaming-kinesis-asl-assembly/assembly' or " + "'build/mvn -Pkinesis-asl package' before running this test." + ) + else: + existing_args = os.environ.get("PYSPARK_SUBMIT_ARGS", "pyspark-shell") + jars_args = "--jars %s" % kinesis_asl_assembly_jar + os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) + kinesis_requirement_message = None # type: ignore + +should_test_kinesis = kinesis_requirement_message is None + + +class PySparkStreamingTestCase(unittest.TestCase): + + timeout = 30 # seconds + duration = 0.5 + + @classmethod + def setUpClass(cls): + class_name = cls.__name__ + conf = SparkConf().set("spark.default.parallelism", 1) + cls.sc = SparkContext(appName=class_name, conf=conf) + cls.sc.setCheckpointDir(tempfile.mkdtemp()) + + @classmethod + def tearDownClass(cls): + cls.sc.stop() + # Clean up in the JVM just in case there has been some issues in Python API + try: + jSparkContextOption = SparkContext._jvm.SparkContext.get() + if jSparkContextOption.nonEmpty(): + jSparkContextOption.get().stop() + except BaseException: + pass + + def setUp(self): + self.ssc = StreamingContext(self.sc, self.duration) + + def tearDown(self): + if self.ssc is not None: + self.ssc.stop(False) + # Clean up in the JVM just in case there has been some issues in Python API + try: + jStreamingContextOption = StreamingContext._jvm.SparkContext.getActive() + if jStreamingContextOption.nonEmpty(): + jStreamingContextOption.get().stop(False) + except BaseException: + pass + + def wait_for(self, result, n): + start_time = time.time() + while len(result) < n and time.time() - start_time < self.timeout: + time.sleep(0.01) + if len(result) < n: + print("timeout after", self.timeout) + + def _take(self, dstream, n): + """ + Return the first `n` elements in the stream (will start and stop). + """ + results = [] + + def take(_, rdd): + if rdd and len(results) < n: + results.extend(rdd.take(n - len(results))) + + dstream.foreachRDD(take) + + self.ssc.start() + self.wait_for(results, n) + return results + + def _collect(self, dstream, n, block=True): + """ + Collect each RDDs into the returned list. + + Returns + ------- + list + which will have the collected items. + """ + result = [] + + def get_output(_, rdd): + if rdd and len(result) < n: + r = rdd.collect() + if r: + result.append(r) + + dstream.foreachRDD(get_output) + + if not block: + return result + + self.ssc.start() + self.wait_for(result, n) + return result + + def _test_func(self, input, func, expected, sort=False, input2=None): + """ + Parameters + ---------- + input : list + dataset for the test. This should be list of lists. + func : function + wrapped function. This function should return PythonDStream object. + expected + expected output for this testcase. + """ + if not isinstance(input[0], RDD): + input = [self.sc.parallelize(d, 1) for d in input] + input_stream = self.ssc.queueStream(input) + if input2 and not isinstance(input2[0], RDD): + input2 = [self.sc.parallelize(d, 1) for d in input2] + input_stream2 = self.ssc.queueStream(input2) if input2 is not None else None + + # Apply test function to stream. + if input2: + stream = func(input_stream, input_stream2) + else: + stream = func(input_stream) + + result = self._collect(stream, len(expected)) + if sort: + self._sort_result_based_on_key(result) + self._sort_result_based_on_key(expected) + self.assertEqual(expected, result) + + def _sort_result_based_on_key(self, outputs): + """Sort the list based on first value.""" + for output in outputs: + output.sort(key=lambda x: x[0]) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/testing/utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/utils.py new file mode 100644 index 0000000..7dd7236 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/testing/utils.py @@ -0,0 +1,626 @@ +# +# 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 glob +import os +import struct +import sys +import unittest +import difflib +from time import time, sleep +from typing import ( + Any, + Optional, + Union, + Dict, + List, + Tuple, + Iterator, +) +from itertools import zip_longest + +from pyspark import SparkContext, SparkConf +from pyspark.errors import PySparkAssertionError, PySparkException +from pyspark.find_spark_home import _find_spark_home +from pyspark.sql.dataframe import DataFrame +from pyspark.sql import Row +from pyspark.sql.types import StructType, AtomicType, StructField + +have_scipy = False +have_numpy = False +try: + import scipy.sparse # noqa: F401 + + have_scipy = True +except ImportError: + # No SciPy, but that's okay, we'll skip those tests + pass +try: + import numpy as np # noqa: F401 + + have_numpy = True +except ImportError: + # No NumPy, but that's okay, we'll skip those tests + pass + +__all__ = ["assertDataFrameEqual", "assertSchemaEqual"] + +SPARK_HOME = _find_spark_home() + + +def read_int(b): + return struct.unpack("!i", b)[0] + + +def write_int(i): + return struct.pack("!i", i) + + +def eventually(condition, timeout=30.0, catch_assertions=False): + """ + Wait a given amount of time for a condition to pass, else fail with an error. + This is a helper utility for PySpark tests. + + Parameters + ---------- + condition : function + Function that checks for termination conditions. condition() can return: + - True: Conditions met. Return without error. + - other value: Conditions not met yet. Continue. Upon timeout, + include last such value in error message. + Note that this method may be called at any time during + streaming execution (e.g., even before any results + have been created). + timeout : int + Number of seconds to wait. Default 30 seconds. + catch_assertions : bool + If False (default), do not catch AssertionErrors. + If True, catch AssertionErrors; continue, but save + error to throw upon timeout. + """ + start_time = time() + lastValue = None + while time() - start_time < timeout: + if catch_assertions: + try: + lastValue = condition() + except AssertionError as e: + lastValue = e + else: + lastValue = condition() + if lastValue is True: + return + sleep(0.01) + if isinstance(lastValue, AssertionError): + raise lastValue + else: + raise AssertionError( + "Test failed due to timeout after %g sec, with last condition returning: %s" + % (timeout, lastValue) + ) + + +class QuietTest: + def __init__(self, sc): + self.log4j = sc._jvm.org.apache.log4j + + def __enter__(self): + self.old_level = self.log4j.LogManager.getRootLogger().getLevel() + self.log4j.LogManager.getRootLogger().setLevel(self.log4j.Level.FATAL) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.log4j.LogManager.getRootLogger().setLevel(self.old_level) + + +class PySparkTestCase(unittest.TestCase): + def setUp(self): + self._old_sys_path = list(sys.path) + class_name = self.__class__.__name__ + self.sc = SparkContext("local[4]", class_name) + + def tearDown(self): + self.sc.stop() + sys.path = self._old_sys_path + + +class ReusedPySparkTestCase(unittest.TestCase): + @classmethod + def conf(cls): + """ + Override this in subclasses to supply a more specific conf + """ + return SparkConf() + + @classmethod + def setUpClass(cls): + cls.sc = SparkContext("local[4]", cls.__name__, conf=cls.conf()) + + @classmethod + def tearDownClass(cls): + cls.sc.stop() + + +class ByteArrayOutput: + def __init__(self): + self.buffer = bytearray() + + def write(self, b): + self.buffer += b + + def close(self): + pass + + +def search_jar(project_relative_path, sbt_jar_name_prefix, mvn_jar_name_prefix): + # Note that 'sbt_jar_name_prefix' and 'mvn_jar_name_prefix' are used since the prefix can + # vary for SBT or Maven specifically. See also SPARK-26856 + project_full_path = os.path.join(SPARK_HOME, project_relative_path) + + # We should ignore the following jars + ignored_jar_suffixes = ("javadoc.jar", "sources.jar", "test-sources.jar", "tests.jar") + + # Search jar in the project dir using the jar name_prefix for both sbt build and maven + # build because the artifact jars are in different directories. + sbt_build = glob.glob( + os.path.join(project_full_path, "target/scala-*/%s*.jar" % sbt_jar_name_prefix) + ) + maven_build = glob.glob(os.path.join(project_full_path, "target/%s*.jar" % mvn_jar_name_prefix)) + jar_paths = sbt_build + maven_build + jars = [jar for jar in jar_paths if not jar.endswith(ignored_jar_suffixes)] + + if not jars: + return None + elif len(jars) > 1: + raise RuntimeError("Found multiple JARs: %s; please remove all but one" % (", ".join(jars))) + else: + return jars[0] + + +def _terminal_color_support(): + try: + # determine if environment supports color + script = "$(test $(tput colors)) && $(test $(tput colors) -ge 8) && echo true || echo false" + return os.popen(script).read() + except Exception: + return False + + +def _context_diff(actual: List[str], expected: List[str], n: int = 3): + """ + Modified from difflib context_diff API, + see original code here: https://github.com/python/cpython/blob/main/Lib/difflib.py#L1180 + """ + + def red(s: str) -> str: + red_color = "\033[31m" + no_color = "\033[0m" + return red_color + str(s) + no_color + + prefix = dict(insert="+ ", delete="- ", replace="! ", equal=" ") + for group in difflib.SequenceMatcher(None, actual, expected).get_grouped_opcodes(n): + yield "*** actual ***" + if any(tag in {"replace", "delete"} for tag, _, _, _, _ in group): + for tag, i1, i2, _, _ in group: + for line in actual[i1:i2]: + if tag != "equal" and _terminal_color_support(): + yield red(prefix[tag] + str(line)) + else: + yield prefix[tag] + str(line) + + yield "\n" + + yield "*** expected ***" + if any(tag in {"replace", "insert"} for tag, _, _, _, _ in group): + for tag, _, _, j1, j2 in group: + for line in expected[j1:j2]: + if tag != "equal" and _terminal_color_support(): + yield red(prefix[tag] + str(line)) + else: + yield prefix[tag] + str(line) + + +class PySparkErrorTestUtils: + """ + This util provide functions to accurate and consistent error testing + based on PySpark error classes. + """ + + def check_error( + self, + exception: PySparkException, + error_class: str, + message_parameters: Optional[Dict[str, str]] = None, + ): + # Test if given error is an instance of PySparkException. + self.assertIsInstance( + exception, + PySparkException, + f"checkError requires 'PySparkException', got '{exception.__class__.__name__}'.", + ) + + # Test error class + expected = error_class + actual = exception.getErrorClass() + self.assertEqual( + expected, actual, f"Expected error class was '{expected}', got '{actual}'." + ) + + # Test message parameters + expected = message_parameters + actual = exception.getMessageParameters() + self.assertEqual( + expected, actual, f"Expected message parameters was '{expected}', got '{actual}'" + ) + + +def assertSchemaEqual(actual: StructType, expected: StructType): + r""" + A util function to assert equality between DataFrame schemas `actual` and `expected`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + actual : StructType + The DataFrame schema that is being compared or tested. + expected : StructType + The expected schema, for comparison with the actual schema. + + Notes + ----- + When assertSchemaEqual fails, the error message uses the Python `difflib` library to display + a diff log of the `actual` and `expected` schemas. + + Examples + -------- + >>> from pyspark.sql.types import StructType, StructField, ArrayType, IntegerType, DoubleType + >>> s1 = StructType([StructField("names", ArrayType(DoubleType(), True), True)]) + >>> s2 = StructType([StructField("names", ArrayType(DoubleType(), True), True)]) + >>> assertSchemaEqual(s1, s2) # pass, schemas are identical + + >>> df1 = spark.createDataFrame(data=[(1, 1000), (2, 3000)], schema=["id", "number"]) + >>> df2 = spark.createDataFrame(data=[("1", 1000), ("2", 5000)], schema=["id", "amount"]) + >>> assertSchemaEqual(df1.schema, df2.schema) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + PySparkAssertionError: [DIFFERENT_SCHEMA] Schemas do not match. + --- actual + +++ expected + - StructType([StructField('id', LongType(), True), StructField('number', LongType(), True)]) + ? ^^ ^^^^^ + + StructType([StructField('id', StringType(), True), StructField('amount', LongType(), True)]) + ? ^^^^ ++++ ^ + """ + if not isinstance(actual, StructType): + raise PySparkAssertionError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(actual)}, + ) + if not isinstance(expected, StructType): + raise PySparkAssertionError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(expected)}, + ) + + def compare_schemas_ignore_nullable(s1: StructType, s2: StructType): + if len(s1) != len(s2): + return False + zipped = zip_longest(s1, s2) + for sf1, sf2 in zipped: + if not compare_structfields_ignore_nullable(sf1, sf2): + return False + return True + + def compare_structfields_ignore_nullable(actualSF: StructField, expectedSF: StructField): + if actualSF is None and expectedSF is None: + return True + elif actualSF is None or expectedSF is None: + return False + if actualSF.name != expectedSF.name: + return False + else: + return compare_datatypes_ignore_nullable(actualSF.dataType, expectedSF.dataType) + + def compare_datatypes_ignore_nullable(dt1: Any, dt2: Any): + # checks datatype equality, using recursion to ignore nullable + if dt1.typeName() == dt2.typeName(): + if dt1.typeName() == "array": + return compare_datatypes_ignore_nullable(dt1.elementType, dt2.elementType) + elif dt1.typeName() == "struct": + return compare_schemas_ignore_nullable(dt1, dt2) + else: + return True + else: + return False + + # ignore nullable flag by default + if not compare_schemas_ignore_nullable(actual, expected): + generated_diff = difflib.ndiff(str(actual).splitlines(), str(expected).splitlines()) + + error_msg = "\n".join(generated_diff) + + raise PySparkAssertionError( + error_class="DIFFERENT_SCHEMA", + message_parameters={"error_msg": error_msg}, + ) + + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import pandas + import pyspark.pandas + + +def assertDataFrameEqual( + actual: Union[DataFrame, "pandas.DataFrame", "pyspark.pandas.DataFrame", List[Row]], + expected: Union[DataFrame, "pandas.DataFrame", "pyspark.pandas.DataFrame", List[Row]], + checkRowOrder: bool = False, + rtol: float = 1e-5, + atol: float = 1e-8, +): + r""" + A util function to assert equality between `actual` and `expected` + (DataFrames or lists of Rows), with optional parameters `checkRowOrder`, `rtol`, and `atol`. + + Supports Spark, Spark Connect, pandas, and pandas-on-Spark DataFrames. + For more information about pandas-on-Spark DataFrame equality, see the docs for + `assertPandasOnSparkEqual`. + + .. versionadded:: 3.5.0 + + Parameters + ---------- + actual : DataFrame (Spark, Spark Connect, pandas, or pandas-on-Spark) or list of Rows + The DataFrame that is being compared or tested. + expected : DataFrame (Spark, Spark Connect, pandas, or pandas-on-Spark) or list of Rows + The expected result of the operation, for comparison with the actual result. + checkRowOrder : bool, optional + A flag indicating whether the order of rows should be considered in the comparison. + If set to `False` (default), the row order is not taken into account. + If set to `True`, the order of rows is important and will be checked during comparison. + (See Notes) + rtol : float, optional + The relative tolerance, used in asserting approximate equality for float values in actual + and expected. Set to 1e-5 by default. (See Notes) + atol : float, optional + The absolute tolerance, used in asserting approximate equality for float values in actual + and expected. Set to 1e-8 by default. (See Notes) + + Notes + ----- + When `assertDataFrameEqual` fails, the error message uses the Python `difflib` library to + display a diff log of each row that differs in `actual` and `expected`. + + For `checkRowOrder`, note that PySpark DataFrame ordering is non-deterministic, unless + explicitly sorted. + + Note that schema equality is checked only when `expected` is a DataFrame (not a list of Rows). + + For DataFrames with float values, assertDataFrame asserts approximate equality. + Two float values a and b are approximately equal if the following equation is True: + + ``absolute(a - b) <= (atol + rtol * absolute(b))``. + + Examples + -------- + >>> df1 = spark.createDataFrame(data=[("1", 1000), ("2", 3000)], schema=["id", "amount"]) + >>> df2 = spark.createDataFrame(data=[("1", 1000), ("2", 3000)], schema=["id", "amount"]) + >>> assertDataFrameEqual(df1, df2) # pass, DataFrames are identical + + >>> df1 = spark.createDataFrame(data=[("1", 0.1), ("2", 3.23)], schema=["id", "amount"]) + >>> df2 = spark.createDataFrame(data=[("1", 0.109), ("2", 3.23)], schema=["id", "amount"]) + >>> assertDataFrameEqual(df1, df2, rtol=1e-1) # pass, DataFrames are approx equal by rtol + + >>> df1 = spark.createDataFrame(data=[(1, 1000), (2, 3000)], schema=["id", "amount"]) + >>> list_of_rows = [Row(1, 1000), Row(2, 3000)] + >>> assertDataFrameEqual(df1, list_of_rows) # pass, actual and expected data are equal + + >>> import pyspark.pandas as ps + >>> df1 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> df2 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]}) + >>> assertDataFrameEqual(df1, df2) # pass, pandas-on-Spark DataFrames are equal + + >>> df1 = spark.createDataFrame( + ... data=[("1", 1000.00), ("2", 3000.00), ("3", 2000.00)], schema=["id", "amount"]) + >>> df2 = spark.createDataFrame( + ... data=[("1", 1001.00), ("2", 3000.00), ("3", 2003.00)], schema=["id", "amount"]) + >>> assertDataFrameEqual(df1, df2) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + PySparkAssertionError: [DIFFERENT_ROWS] Results do not match: ( 66.66667 % ) + *** actual *** + ! Row(id='1', amount=1000.0) + Row(id='2', amount=3000.0) + ! Row(id='3', amount=2000.0) + *** expected *** + ! Row(id='1', amount=1001.0) + Row(id='2', amount=3000.0) + ! Row(id='3', amount=2003.0) + """ + if actual is None and expected is None: + return True + elif actual is None: + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": None, + }, + ) + elif expected is None: + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "expected", + "actual_type": None, + }, + ) + + has_pandas = False + try: + # If pandas dependencies are available, allow pandas or pandas-on-Spark DataFrame + import pyspark.pandas as ps + import pandas as pd + from pyspark.testing.pandasutils import PandasOnSparkTestUtils + + has_pandas = True + except ImportError: + # no pandas, so we won't call pandasutils functions + pass + + if has_pandas: + if ( + isinstance(actual, pd.DataFrame) + or isinstance(expected, pd.DataFrame) + or isinstance(actual, ps.DataFrame) + or isinstance(expected, ps.DataFrame) + ): + # handle pandas DataFrames + # assert approximate equality for float data + return PandasOnSparkTestUtils().assert_eq( + actual, expected, almost=True, rtol=rtol, atol=atol, check_row_order=checkRowOrder + ) + + from pyspark.sql.utils import get_dataframe_class + + # if is_remote(), allow Connect DataFrame + SparkDataFrame = get_dataframe_class() + + if not isinstance(actual, (DataFrame, SparkDataFrame, list)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "actual", + "actual_type": type(actual), + }, + ) + elif not isinstance(expected, (DataFrame, SparkDataFrame, list)): + raise PySparkAssertionError( + error_class="INVALID_TYPE_DF_EQUALITY_ARG", + message_parameters={ + "expected_type": "Union[DataFrame, ps.DataFrame, List[Row]]", + "arg_name": "expected", + "actual_type": type(expected), + }, + ) + + def compare_rows(r1: Row, r2: Row): + def compare_vals(val1, val2): + if isinstance(val1, list) and isinstance(val2, list): + return len(val1) == len(val2) and all( + compare_vals(x, y) for x, y in zip(val1, val2) + ) + elif isinstance(val1, Row) and isinstance(val2, Row): + return all(compare_vals(x, y) for x, y in zip(val1, val2)) + elif isinstance(val1, dict) and isinstance(val2, dict): + return ( + len(val1.keys()) == len(val2.keys()) + and val1.keys() == val2.keys() + and all(compare_vals(val1[k], val2[k]) for k in val1.keys()) + ) + elif isinstance(val1, float) and isinstance(val2, float): + if abs(val1 - val2) > (atol + rtol * abs(val2)): + return False + else: + if val1 != val2: + return False + return True + + if r1 is None and r2 is None: + return True + elif r1 is None or r2 is None: + return False + + return compare_vals(r1, r2) + + def assert_rows_equal(rows1: List[Row], rows2: List[Row]): + zipped = list(zip_longest(rows1, rows2)) + diff_rows_cnt = 0 + diff_rows = False + + rows_str1 = "" + rows_str2 = "" + + # count different rows + for r1, r2 in zipped: + rows_str1 += str(r1) + "\n" + rows_str2 += str(r2) + "\n" + if not compare_rows(r1, r2): + diff_rows_cnt += 1 + diff_rows = True + + generated_diff = _context_diff( + actual=rows_str1.splitlines(), expected=rows_str2.splitlines(), n=len(zipped) + ) + + if diff_rows: + error_msg = "Results do not match: " + percent_diff = (diff_rows_cnt / len(zipped)) * 100 + error_msg += "( %.5f %% )" % percent_diff + error_msg += "\n" + "\n".join(generated_diff) + raise PySparkAssertionError( + error_class="DIFFERENT_ROWS", + message_parameters={"error_msg": error_msg}, + ) + + # convert actual and expected to list + if not isinstance(actual, list) and not isinstance(expected, list): + # only compare schema if expected is not a List + assertSchemaEqual(actual.schema, expected.schema) + + if not isinstance(actual, list): + actual_list = actual.collect() + else: + actual_list = actual + + if not isinstance(expected, list): + expected_list = expected.collect() + else: + expected_list = expected + + if not checkRowOrder: + # rename duplicate columns for sorting + actual_list = sorted(actual_list, key=lambda x: str(x)) + expected_list = sorted(expected_list, key=lambda x: str(x)) + + assert_rows_equal(actual_list, expected_list) + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.testing.utils + + globs = pyspark.testing.utils.__dict__.copy() + spark = SparkSession.builder.master("local[4]").appName("testing.utils tests").getOrCreate() + globs["spark"] = spark + (failure_count, test_count) = doctest.testmod( + pyspark.testing.utils, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/__init__.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/__init__.py new file mode 100644 index 0000000..12bdf0d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/__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/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_appsubmit.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_appsubmit.py new file mode 100644 index 0000000..79b6b4f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_appsubmit.py @@ -0,0 +1,306 @@ +# +# 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 os +import re +import shutil +import subprocess +import tempfile +import unittest +import zipfile + +from pyspark.testing.utils import SPARK_HOME + + +class SparkSubmitTests(unittest.TestCase): + def setUp(self): + self.programDir = tempfile.mkdtemp() + tmp_dir = tempfile.gettempdir() + self.sparkSubmit = [ + os.path.join(SPARK_HOME, "bin", "spark-submit"), + "--conf", + "spark.driver.extraJavaOptions=-Djava.io.tmpdir={0}".format(tmp_dir), + "--conf", + "spark.executor.extraJavaOptions=-Djava.io.tmpdir={0}".format(tmp_dir), + ] + + def tearDown(self): + shutil.rmtree(self.programDir) + + def createTempFile(self, name, content, dir=None): + """ + Create a temp file with the given name and content and return its path. + Strips leading spaces from content up to the first '|' in each line. + """ + pattern = re.compile(r"^ *\|", re.MULTILINE) + content = re.sub(pattern, "", content.strip()) + if dir is None: + path = os.path.join(self.programDir, name) + else: + os.makedirs(os.path.join(self.programDir, dir)) + path = os.path.join(self.programDir, dir, name) + with open(path, "w") as f: + f.write(content) + return path + + def createFileInZip(self, name, content, ext=".zip", dir=None, zip_name=None): + """ + Create a zip archive containing a file with the given content and return its path. + Strips leading spaces from content up to the first '|' in each line. + """ + pattern = re.compile(r"^ *\|", re.MULTILINE) + content = re.sub(pattern, "", content.strip()) + if dir is None: + path = os.path.join(self.programDir, name + ext) + else: + path = os.path.join(self.programDir, dir, zip_name + ext) + zip = zipfile.ZipFile(path, "w") + zip.writestr(name, content) + zip.close() + return path + + def create_spark_package(self, artifact_name): + group_id, artifact_id, version = artifact_name.split(":") + self.createTempFile( + "%s-%s.pom" % (artifact_id, version), + ( + """ + | + | + | 4.0.0 + | %s + | %s + | %s + | + """ + % (group_id, artifact_id, version) + ).lstrip(), + os.path.join(group_id, artifact_id, version), + ) + self.createFileInZip( + "%s.py" % artifact_id, + """ + |def myfunc(x): + | return x + 1 + """, + ".jar", + os.path.join(group_id, artifact_id, version), + "%s-%s" % (artifact_id, version), + ) + + def test_single_script(self): + """Submit and test a single script file""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(lambda x: x * 2).collect()) + """, + ) + proc = subprocess.Popen(self.sparkSubmit + [script], stdout=subprocess.PIPE) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 4, 6]", out.decode("utf-8")) + + def test_script_with_local_functions(self): + """Submit and test a single script file calling a global function""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + | + |def foo(x): + | return x * 3 + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(foo).collect()) + """, + ) + proc = subprocess.Popen(self.sparkSubmit + [script], stdout=subprocess.PIPE) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[3, 6, 9]", out.decode("utf-8")) + + def test_module_dependency(self): + """Submit and test a script with a dependency on another module""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + |from mylib import myfunc + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(myfunc).collect()) + """, + ) + zip = self.createFileInZip( + "mylib.py", + """ + |def myfunc(x): + | return x + 1 + """, + ) + proc = subprocess.Popen( + self.sparkSubmit + ["--py-files", zip, script], stdout=subprocess.PIPE + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 3, 4]", out.decode("utf-8")) + + def test_module_dependency_on_cluster(self): + """Submit and test a script with a dependency on another module on a cluster""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + |from mylib import myfunc + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(myfunc).collect()) + """, + ) + zip = self.createFileInZip( + "mylib.py", + """ + |def myfunc(x): + | return x + 1 + """, + ) + proc = subprocess.Popen( + self.sparkSubmit + ["--py-files", zip, "--master", "local-cluster[1,1,1024]", script], + stdout=subprocess.PIPE, + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 3, 4]", out.decode("utf-8")) + + def test_package_dependency(self): + """Submit and test a script with a dependency on a Spark Package""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + |from mylib import myfunc + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(myfunc).collect()) + """, + ) + self.create_spark_package("a:mylib:0.1") + proc = subprocess.Popen( + self.sparkSubmit + + ["--packages", "a:mylib:0.1", "--repositories", "file:" + self.programDir, script], + stdout=subprocess.PIPE, + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 3, 4]", out.decode("utf-8")) + + def test_package_dependency_on_cluster(self): + """Submit and test a script with a dependency on a Spark Package on a cluster""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + |from mylib import myfunc + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(myfunc).collect()) + """, + ) + self.create_spark_package("a:mylib:0.1") + proc = subprocess.Popen( + self.sparkSubmit + + [ + "--packages", + "a:mylib:0.1", + "--repositories", + "file:" + self.programDir, + "--master", + "local-cluster[1,1,1024]", + script, + ], + stdout=subprocess.PIPE, + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 3, 4]", out.decode("utf-8")) + + def test_single_script_on_cluster(self): + """Submit and test a single script on a cluster""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkContext + | + |def foo(x): + | return x * 2 + | + |sc = SparkContext() + |print(sc.parallelize([1, 2, 3]).map(foo).collect()) + """, + ) + # this will fail if you have different spark.executor.memory + # in conf/spark-defaults.conf + proc = subprocess.Popen( + self.sparkSubmit + ["--master", "local-cluster[1,1,1024]", script], + stdout=subprocess.PIPE, + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode) + self.assertIn("[2, 4, 6]", out.decode("utf-8")) + + def test_user_configuration(self): + """Make sure user configuration is respected (SPARK-19307)""" + script = self.createTempFile( + "test.py", + """ + |from pyspark import SparkConf, SparkContext + | + |conf = SparkConf().set("spark.test_config", "1") + |sc = SparkContext(conf = conf) + |try: + | if sc._conf.get("spark.test_config") != "1": + | raise RuntimeError("Cannot find spark.test_config in SparkContext's conf.") + |finally: + | sc.stop() + """, + ) + proc = subprocess.Popen( + self.sparkSubmit + ["--master", "local", script], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + out, err = proc.communicate() + self.assertEqual(0, proc.returncode, msg="Process failed with error:\n {0}".format(out)) + + +if __name__ == "__main__": + from pyspark.tests.test_appsubmit import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_broadcast.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_broadcast.py new file mode 100644 index 0000000..90d3caa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_broadcast.py @@ -0,0 +1,196 @@ +# +# 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 os +import pickle +import random +import time +import tempfile +import unittest + +from py4j.protocol import Py4JJavaError + +from pyspark import SparkConf, SparkContext, Broadcast +from pyspark.java_gateway import launch_gateway +from pyspark.serializers import ChunkedStream +from pyspark.sql import SparkSession, Row + + +class BroadcastTest(unittest.TestCase): + def tearDown(self): + if getattr(self, "sc", None) is not None: + self.sc.stop() + self.sc = None + + def _test_encryption_helper(self, vs): + """ + Creates a broadcast variables for each value in vs, and runs a simple job to make sure the + value is the same when it's read in the executors. Also makes sure there are no task + failures. + """ + bs = [self.sc.broadcast(value=v) for v in vs] + exec_values = self.sc.parallelize(range(2)).map(lambda x: [b.value for b in bs]).collect() + for ev in exec_values: + self.assertEqual(ev, vs) + # make sure there are no task failures + status = self.sc.statusTracker() + for jid in status.getJobIdsForGroup(): + for sid in status.getJobInfo(jid).stageIds: + stage_info = status.getStageInfo(sid) + self.assertEqual(0, stage_info.numFailedTasks) + + def _test_multiple_broadcasts(self, *extra_confs): + """ + Test broadcast variables make it OK to the executors. Tests multiple broadcast variables, + and also multiple jobs. + """ + conf = SparkConf() + for key, value in extra_confs: + conf.set(key, value) + conf.setMaster("local-cluster[2,1,1024]") + self.sc = SparkContext(conf=conf) + self._test_encryption_helper([5]) + self._test_encryption_helper([5, 10, 20]) + + def test_broadcast_with_encryption(self): + self._test_multiple_broadcasts(("spark.io.encryption.enabled", "true")) + + def test_broadcast_no_encryption(self): + self._test_multiple_broadcasts() + + def _test_broadcast_on_driver(self, *extra_confs): + conf = SparkConf() + for key, value in extra_confs: + conf.set(key, value) + conf.setMaster("local-cluster[2,1,1024]") + self.sc = SparkContext(conf=conf) + bs = self.sc.broadcast(value=5) + self.assertEqual(5, bs.value) + + def test_broadcast_value_driver_no_encryption(self): + self._test_broadcast_on_driver() + + def test_broadcast_value_driver_encryption(self): + self._test_broadcast_on_driver(("spark.io.encryption.enabled", "true")) + + def test_broadcast_value_against_gc(self): + # Test broadcast value against gc. + conf = SparkConf() + conf.setMaster("local[1,1]") + conf.set("spark.memory.fraction", "0.0001") + self.sc = SparkContext(conf=conf) + b = self.sc.broadcast([100]) + try: + res = self.sc.parallelize([0], 1).map(lambda x: 0 if x == 0 else b.value[0]).collect() + self.assertEqual([0], res) + self.sc._jvm.java.lang.System.gc() + time.sleep(5) + res = self.sc.parallelize([1], 1).map(lambda x: 0 if x == 0 else b.value[0]).collect() + self.assertEqual([100], res) + finally: + b.destroy() + + def test_broadcast_when_sc_none(self): + # SPARK-39029 : Test case to improve test coverage of broadcast.py + # It tests the case when SparkContext is none and Broadcast is called at executor + conf = SparkConf() + conf.setMaster("local-cluster[2,1,1024]") + self.sc = SparkContext(conf=conf) + bs = self.sc.broadcast([10]) + bs_sc_none = Broadcast(sc=None, path=bs._path) + self.assertEqual(bs_sc_none.value, [10]) + + def test_broadcast_for_error_condition(self): + # SPARK-39029: Test case to improve test coverage of broadcast.py + # It tests the case when broadcast should raise error . + conf = SparkConf() + conf.setMaster("local-cluster[2,1,1024]") + self.sc = SparkContext(conf=conf) + bs = self.sc.broadcast([1]) + with self.assertRaisesRegex(pickle.PickleError, "Could.*not.*serialize.*broadcast"): + self.sc.broadcast(self.sc) + with self.assertRaisesRegex(Py4JJavaError, "RuntimeError.*Broadcast.*destroyed.*driver"): + self.sc.parallelize([1]).map(lambda x: bs.destroy()).collect() + with self.assertRaisesRegex(Py4JJavaError, "RuntimeError.*Broadcast.*unpersisted.*driver"): + self.sc.parallelize([1]).map(lambda x: bs.unpersist()).collect() + + def test_broadcast_in_udfs_with_encryption(self): + conf = SparkConf() + conf.set("spark.io.encryption.enabled", "true") + conf.setMaster("local-cluster[2,1,1024]") + self.sc = SparkContext(conf=conf) + bar = {"a": "aa", "b": "bb"} + foo = self.sc.broadcast(bar) + spark = SparkSession(self.sc) + spark.udf.register("MYUDF", lambda x: foo.value[x] if x else "") + sel = spark.sql("SELECT MYUDF('a') AS a, MYUDF('b') AS b") + self.assertEqual(sel.collect(), [Row(a="aa", b="bb")]) + spark.stop() + + +class BroadcastFrameProtocolTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + gateway = launch_gateway(SparkConf()) + cls._jvm = gateway.jvm + cls.longMessage = True + random.seed(42) + + def _test_chunked_stream(self, data, py_buf_size): + # write data using the chunked protocol from python. + chunked_file = tempfile.NamedTemporaryFile(delete=False) + dechunked_file = tempfile.NamedTemporaryFile(delete=False) + dechunked_file.close() + try: + out = ChunkedStream(chunked_file, py_buf_size) + out.write(data) + out.close() + # now try to read it in java + jin = self._jvm.java.io.FileInputStream(chunked_file.name) + jout = self._jvm.java.io.FileOutputStream(dechunked_file.name) + self._jvm.DechunkedInputStream.dechunkAndCopyToOutput(jin, jout) + # java should have decoded it back to the original data + self.assertEqual(len(data), os.stat(dechunked_file.name).st_size) + with open(dechunked_file.name, "rb") as f: + byte = f.read(1) + idx = 0 + while byte: + self.assertEqual(data[idx], bytearray(byte)[0], msg="idx = " + str(idx)) + byte = f.read(1) + idx += 1 + finally: + os.unlink(chunked_file.name) + os.unlink(dechunked_file.name) + + def test_chunked_stream(self): + def random_bytes(n): + return bytearray(random.getrandbits(8) for _ in range(n)) + + for data_length in [1, 10, 100, 10000]: + for buffer_length in [1, 2, 5, 8192]: + self._test_chunked_stream(random_bytes(data_length), buffer_length) + + +if __name__ == "__main__": + from pyspark.tests.test_broadcast import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_conf.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_conf.py new file mode 100644 index 0000000..cc9ff82 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_conf.py @@ -0,0 +1,44 @@ +# +# 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 random +import unittest + +from pyspark import SparkContext, SparkConf + + +class ConfTests(unittest.TestCase): + def test_memory_conf(self): + memoryList = ["1T", "1G", "1M", "1024K"] + for memory in memoryList: + sc = SparkContext(conf=SparkConf().set("spark.python.worker.memory", memory)) + lst = list(range(1024)) + random.shuffle(lst) + rdd = sc.parallelize(lst, 4) + self.assertEqual(sorted(lst), rdd.sortBy(lambda x: x).collect()) + sc.stop() + + +if __name__ == "__main__": + from pyspark.tests.test_conf import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_context.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_context.py new file mode 100644 index 0000000..0a96289 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_context.py @@ -0,0 +1,346 @@ +# +# 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 os +import shutil +import stat +import tempfile +import threading +import time +import unittest +from collections import namedtuple + +from pyspark import SparkConf, SparkFiles, SparkContext +from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest, SPARK_HOME + + +class CheckpointTests(ReusedPySparkTestCase): + def setUp(self): + self.checkpointDir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(self.checkpointDir.name) + self.sc.setCheckpointDir(self.checkpointDir.name) + + def tearDown(self): + shutil.rmtree(self.checkpointDir.name) + + def test_basic_checkpointing(self): + parCollection = self.sc.parallelize([1, 2, 3, 4]) + flatMappedRDD = parCollection.flatMap(lambda x: range(1, x + 1)) + + self.assertFalse(flatMappedRDD.isCheckpointed()) + self.assertTrue(flatMappedRDD.getCheckpointFile() is None) + self.assertFalse(self.sc.getCheckpointDir() is None) + + flatMappedRDD.checkpoint() + result = flatMappedRDD.collect() + time.sleep(1) # 1 second + self.assertTrue(flatMappedRDD.isCheckpointed()) + self.assertEqual(flatMappedRDD.collect(), result) + self.assertEqual( + "file:" + self.checkpointDir.name, + os.path.dirname(os.path.dirname(flatMappedRDD.getCheckpointFile())), + ) + self.assertEqual( + self.sc.getCheckpointDir(), os.path.dirname(flatMappedRDD.getCheckpointFile()) + ) + + def test_checkpoint_and_restore(self): + parCollection = self.sc.parallelize([1, 2, 3, 4]) + flatMappedRDD = parCollection.flatMap(lambda x: [x]) + + self.assertFalse(flatMappedRDD.isCheckpointed()) + self.assertTrue(flatMappedRDD.getCheckpointFile() is None) + + flatMappedRDD.checkpoint() + flatMappedRDD.count() # forces a checkpoint to be computed + time.sleep(1) # 1 second + + self.assertTrue(flatMappedRDD.getCheckpointFile() is not None) + recovered = self.sc._checkpointFile( + flatMappedRDD.getCheckpointFile(), flatMappedRDD._jrdd_deserializer + ) + self.assertEqual([1, 2, 3, 4], recovered.collect()) + + +class LocalCheckpointTests(ReusedPySparkTestCase): + def test_basic_localcheckpointing(self): + parCollection = self.sc.parallelize([1, 2, 3, 4]) + flatMappedRDD = parCollection.flatMap(lambda x: range(1, x + 1)) + + self.assertFalse(flatMappedRDD.isCheckpointed()) + self.assertFalse(flatMappedRDD.isLocallyCheckpointed()) + + flatMappedRDD.localCheckpoint() + result = flatMappedRDD.collect() + time.sleep(1) # 1 second + self.assertTrue(flatMappedRDD.isCheckpointed()) + self.assertTrue(flatMappedRDD.isLocallyCheckpointed()) + self.assertEqual(flatMappedRDD.collect(), result) + + +class AddFileTests(PySparkTestCase): + def test_add_py_file(self): + # To ensure that we're actually testing addPyFile's effects, check that + # this job fails due to `userlibrary` not being on the Python path: + # disable logging in log4j temporarily + def func(x): + from userlibrary import UserClass + + return UserClass().hello() + + with QuietTest(self.sc): + self.assertRaises(Exception, self.sc.parallelize(range(2)).map(func).first) + + # Add the file, so the job should now succeed: + path = os.path.join(SPARK_HOME, "python/test_support/userlibrary.py") + self.sc.addPyFile(path) + res = self.sc.parallelize(range(2)).map(func).first() + self.assertEqual("Hello World!", res) + + def test_add_file_locally(self): + path = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + self.sc.addFile(path) + download_path = SparkFiles.get("hello.txt") + self.assertNotEqual(path, download_path) + with open(download_path) as test_file: + self.assertEqual("Hello World!\n", test_file.readline()) + + def test_add_file_recursively_locally(self): + path = os.path.join(SPARK_HOME, "python/test_support/hello") + self.sc.addFile(path, True) + download_path = SparkFiles.get("hello") + self.assertNotEqual(path, download_path) + with open(download_path + "/hello.txt") as test_file: + self.assertEqual("Hello World!\n", test_file.readline()) + with open(download_path + "/sub_hello/sub_hello.txt") as test_file: + self.assertEqual("Sub Hello World!\n", test_file.readline()) + + def test_add_py_file_locally(self): + # To ensure that we're actually testing addPyFile's effects, check that + # this fails due to `userlibrary` not being on the Python path: + def func(): + from userlibrary import UserClass # noqa: F401 + + self.assertRaises(ImportError, func) + path = os.path.join(SPARK_HOME, "python/test_support/userlibrary.py") + self.sc.addPyFile(path) + from userlibrary import UserClass + + self.assertEqual("Hello World!", UserClass().hello()) + + def test_add_egg_file_locally(self): + # To ensure that we're actually testing addPyFile's effects, check that + # this fails due to `userlibrary` not being on the Python path: + def func(): + from userlib import UserClass + + UserClass() + + self.assertRaises(ImportError, func) + path = os.path.join(SPARK_HOME, "python/test_support/userlib-0.1.zip") + self.sc.addPyFile(path) + from userlib import UserClass + + self.assertEqual("Hello World from inside a package!", UserClass().hello()) + + def test_overwrite_system_module(self): + self.sc.addPyFile(os.path.join(SPARK_HOME, "python/test_support/SimpleHTTPServer.py")) + + import SimpleHTTPServer + + self.assertEqual("My Server", SimpleHTTPServer.__name__) + + def func(x): + import SimpleHTTPServer + + return SimpleHTTPServer.__name__ + + self.assertEqual(["My Server"], self.sc.parallelize(range(1)).map(func).collect()) + + +class ContextTests(unittest.TestCase): + def test_failed_sparkcontext_creation(self): + # Regression test for SPARK-1550 + self.assertRaises(Exception, lambda: SparkContext("an-invalid-master-name")) + + def test_get_or_create(self): + with SparkContext.getOrCreate() as sc: + self.assertTrue(SparkContext.getOrCreate() is sc) + + def test_parallelize_eager_cleanup(self): + with SparkContext() as sc: + temp_files = os.listdir(sc._temp_dir) + sc.parallelize([0, 1, 2]) + post_parallelize_temp_files = os.listdir(sc._temp_dir) + self.assertEqual(temp_files, post_parallelize_temp_files) + + def test_set_conf(self): + # This is for an internal use case. When there is an existing SparkContext, + # SparkSession's builder needs to set configs into SparkContext's conf. + sc = SparkContext() + sc._conf.set("spark.test.SPARK16224", "SPARK16224") + self.assertEqual(sc._jsc.sc().conf().get("spark.test.SPARK16224"), "SPARK16224") + sc.stop() + + def test_stop(self): + sc = SparkContext() + self.assertNotEqual(SparkContext._active_spark_context, None) + sc.stop() + self.assertEqual(SparkContext._active_spark_context, None) + + def test_with(self): + with SparkContext(): + self.assertNotEqual(SparkContext._active_spark_context, None) + self.assertEqual(SparkContext._active_spark_context, None) + + def test_with_exception(self): + try: + with SparkContext(): + self.assertNotEqual(SparkContext._active_spark_context, None) + raise RuntimeError() + except BaseException: + pass + self.assertEqual(SparkContext._active_spark_context, None) + + def test_with_stop(self): + with SparkContext() as sc: + self.assertNotEqual(SparkContext._active_spark_context, None) + sc.stop() + self.assertEqual(SparkContext._active_spark_context, None) + + def test_progress_api(self): + with SparkContext() as sc: + sc.setJobGroup("test_progress_api", "", True) + rdd = sc.parallelize(range(10)).map(lambda x: time.sleep(100)) + + def run(): + # When thread is pinned, job group should be set for each thread for now. + # Local properties seem not being inherited like Scala side does. + if os.environ.get("PYSPARK_PIN_THREAD", "true").lower() == "true": + sc.setJobGroup("test_progress_api", "", True) + try: + rdd.count() + except Exception: + pass + + t = threading.Thread(target=run) + t.daemon = True + t.start() + # wait for scheduler to start + time.sleep(3) + + tracker = sc.statusTracker() + jobIds = tracker.getJobIdsForGroup("test_progress_api") + self.assertEqual(1, len(jobIds)) + job = tracker.getJobInfo(jobIds[0]) + self.assertEqual(1, len(job.stageIds)) + stage = tracker.getStageInfo(job.stageIds[0]) + self.assertEqual(rdd.getNumPartitions(), stage.numTasks) + + sc.cancelAllJobs() + t.join() + # wait for event listener to update the status + time.sleep(1) + + job = tracker.getJobInfo(jobIds[0]) + self.assertEqual("FAILED", job.status) + self.assertEqual([], tracker.getActiveJobsIds()) + self.assertEqual([], tracker.getActiveStageIds()) + + sc.stop() + + def test_startTime(self): + with SparkContext() as sc: + self.assertGreater(sc.startTime, 0) + + def test_forbid_insecure_gateway(self): + # Fail immediately if you try to create a SparkContext + # with an insecure gateway + parameters = namedtuple("MockGatewayParameters", "auth_token")(None) + mock_insecure_gateway = namedtuple("MockJavaGateway", "gateway_parameters")(parameters) + with self.assertRaises(ValueError) as context: + SparkContext(gateway=mock_insecure_gateway) + self.assertIn("insecure Py4j gateway", str(context.exception)) + + def test_resources(self): + """Test the resources are empty by default.""" + with SparkContext() as sc: + resources = sc.resources + self.assertEqual(len(resources), 0) + + def test_disallow_to_create_spark_context_in_executors(self): + # SPARK-32160: SparkContext should not be created in executors. + with SparkContext("local-cluster[3, 1, 1024]") as sc: + with self.assertRaises(Exception) as context: + sc.range(2).foreach(lambda _: SparkContext()) + self.assertIn( + "CONTEXT_ONLY_VALID_ON_DRIVER", + str(context.exception), + ) + + def test_allow_to_create_spark_context_in_executors(self): + # SPARK-32160: SparkContext can be created in executors if the config is set. + + def create_spark_context(): + conf = SparkConf().set("spark.executor.allowSparkContext", "true") + with SparkContext(conf=conf): + pass + + with SparkContext("local-cluster[3, 1, 1024]") as sc: + sc.range(2).foreach(lambda _: create_spark_context()) + + +class ContextTestsWithResources(unittest.TestCase): + def setUp(self): + class_name = self.__class__.__name__ + self.tempFile = tempfile.NamedTemporaryFile(delete=False) + self.tempFile.write(b'echo {\\"name\\": \\"gpu\\", \\"addresses\\": [\\"0\\"]}') + self.tempFile.close() + # create temporary directory for Worker resources coordination + self.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(self.tempdir.name) + os.chmod( + self.tempFile.name, + stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH, + ) + conf = SparkConf().set("spark.test.home", SPARK_HOME) + conf = conf.set("spark.driver.resource.gpu.amount", "1") + conf = conf.set("spark.driver.resource.gpu.discoveryScript", self.tempFile.name) + self.sc = SparkContext("local-cluster[2,1,1024]", class_name, conf=conf) + + def test_resources(self): + """Test the resources are available.""" + resources = self.sc.resources + self.assertEqual(len(resources), 1) + self.assertTrue("gpu" in resources) + self.assertEqual(resources["gpu"].name, "gpu") + self.assertEqual(resources["gpu"].addresses, ["0"]) + + def tearDown(self): + os.unlink(self.tempFile.name) + self.sc.stop() + + +if __name__ == "__main__": + from pyspark.tests.test_context import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_daemon.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_daemon.py new file mode 100644 index 0000000..22196e5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_daemon.py @@ -0,0 +1,89 @@ +# +# 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 os +import sys +import time +import unittest + +from pyspark.serializers import read_int + + +class DaemonTests(unittest.TestCase): + def connect(self, port): + from socket import socket, AF_INET, AF_INET6, SOCK_STREAM + + family, host = AF_INET, "127.0.0.1" + if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true": + family, host = AF_INET6, "::1" + sock = socket(family, SOCK_STREAM) + sock.connect((host, port)) + # send a split index of -1 to shutdown the worker + sock.send(b"\xFF\xFF\xFF\xFF") + sock.close() + return True + + def do_termination_test(self, terminator): + from subprocess import Popen, PIPE + from errno import ECONNREFUSED + + # start daemon + daemon_path = os.path.join(os.path.dirname(__file__), "..", "daemon.py") + python_exec = sys.executable or os.environ.get("PYSPARK_PYTHON") + daemon = Popen([python_exec, daemon_path], stdin=PIPE, stdout=PIPE) + + # read the port number + port = read_int(daemon.stdout) + + # daemon should accept connections + self.assertTrue(self.connect(port)) + + # wait worker process spawned from daemon exit. + time.sleep(1) + + # request shutdown + terminator(daemon) + time.sleep(1) + + # daemon should no longer accept connections + try: + self.connect(port) + except EnvironmentError as exception: + self.assertEqual(exception.errno, ECONNREFUSED) + else: + self.fail("Expected EnvironmentError to be raised") + + def test_termination_stdin(self): + """Ensure that daemon and workers terminate when stdin is closed.""" + self.do_termination_test(lambda daemon: daemon.stdin.close()) + + def test_termination_sigterm(self): + """Ensure that daemon and workers terminate on SIGTERM.""" + from signal import SIGTERM + + self.do_termination_test(lambda daemon: os.kill(daemon.pid, SIGTERM)) + + +if __name__ == "__main__": + from pyspark.tests.test_daemon import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_install_spark.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_install_spark.py new file mode 100644 index 0000000..e980a17 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_install_spark.py @@ -0,0 +1,124 @@ +# +# 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 os +import tempfile +import unittest + +from pyspark.install import ( + install_spark, + DEFAULT_HADOOP, + DEFAULT_HIVE, + UNSUPPORTED_COMBINATIONS, + checked_versions, + checked_package_name, +) + + +class SparkInstallationTestCase(unittest.TestCase): + def test_install_spark(self): + # Test only one case. Testing this is expensive because it needs to download + # the Spark distribution. + spark_version, hadoop_version, hive_version = checked_versions("3.0.1", "3", "2.3") + + with tempfile.TemporaryDirectory() as tmp_dir: + install_spark( + dest=tmp_dir, + spark_version=spark_version, + hadoop_version=hadoop_version, + hive_version=hive_version, + ) + + self.assertTrue(os.path.isdir("%s/jars" % tmp_dir)) + self.assertTrue(os.path.exists("%s/bin/spark-submit" % tmp_dir)) + self.assertTrue(os.path.exists("%s/RELEASE" % tmp_dir)) + + def test_package_name(self): + self.assertEqual( + "spark-3.0.0-bin-hadoop3.2", checked_package_name("spark-3.0.0", "hadoop3.2", "hive2.3") + ) + + spark_version, hadoop_version, hive_version = checked_versions("3.2.0", "3", "2.3") + self.assertEqual( + "spark-3.2.0-bin-hadoop3.2", + checked_package_name(spark_version, hadoop_version, hive_version), + ) + + spark_version, hadoop_version, hive_version = checked_versions("3.3.0", "3", "2.3") + self.assertEqual( + "spark-3.3.0-bin-hadoop3", + checked_package_name(spark_version, hadoop_version, hive_version), + ) + + def test_checked_versions(self): + test_version = "3.0.1" # Just pick one version to test. + + # Positive test cases + self.assertEqual( + ("spark-2.4.1", "without-hadoop", "hive2.3"), + checked_versions("2.4.1", "without", "2.3"), + ) + + self.assertEqual( + ("spark-3.0.1", "without-hadoop", "hive2.3"), + checked_versions("spark-3.0.1", "without-hadoop", "hive2.3"), + ) + + self.assertEqual( + ("spark-3.3.0", "hadoop3", "hive2.3"), + checked_versions("spark-3.3.0", "hadoop3", "hive2.3"), + ) + + # Negative test cases + for (hadoop_version, hive_version) in UNSUPPORTED_COMBINATIONS: + with self.assertRaisesRegex(RuntimeError, "Hive.*should.*Hadoop"): + checked_versions( + spark_version=test_version, + hadoop_version=hadoop_version, + hive_version=hive_version, + ) + + with self.assertRaisesRegex(RuntimeError, "Spark version should start with 'spark-'"): + checked_versions( + spark_version="malformed", hadoop_version=DEFAULT_HADOOP, hive_version=DEFAULT_HIVE + ) + + with self.assertRaisesRegex(RuntimeError, "Spark distribution.*malformed.*"): + checked_versions( + spark_version=test_version, hadoop_version="malformed", hive_version=DEFAULT_HIVE + ) + + with self.assertRaisesRegex(RuntimeError, "Spark distribution.*malformed.*"): + checked_versions( + spark_version=test_version, hadoop_version=DEFAULT_HADOOP, hive_version="malformed" + ) + + with self.assertRaisesRegex(RuntimeError, "Spark distribution of hive1.2 is not supported"): + checked_versions( + spark_version=test_version, hadoop_version="hadoop3", hive_version="hive1.2" + ) + + +if __name__ == "__main__": + from pyspark.tests.test_install_spark import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_join.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_join.py new file mode 100644 index 0000000..de1c260 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_join.py @@ -0,0 +1,69 @@ +# +# 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.testing.utils import ReusedPySparkTestCase + + +class JoinTests(ReusedPySparkTestCase): + def test_narrow_dependency_in_join(self): + rdd = self.sc.parallelize(range(10)).map(lambda x: (x, x)) + parted = rdd.partitionBy(2) + self.assertEqual(2, parted.union(parted).getNumPartitions()) + self.assertEqual(rdd.getNumPartitions() + 2, parted.union(rdd).getNumPartitions()) + self.assertEqual(rdd.getNumPartitions() + 2, rdd.union(parted).getNumPartitions()) + + tracker = self.sc.statusTracker() + + self.sc.setJobGroup("test1", "test", True) + d = sorted(parted.join(parted).collect()) + self.assertEqual(10, len(d)) + self.assertEqual((0, (0, 0)), d[0]) + jobId = tracker.getJobIdsForGroup("test1")[0] + self.assertEqual(2, len(tracker.getJobInfo(jobId).stageIds)) + + self.sc.setJobGroup("test2", "test", True) + d = sorted(parted.join(rdd).collect()) + self.assertEqual(10, len(d)) + self.assertEqual((0, (0, 0)), d[0]) + jobId = tracker.getJobIdsForGroup("test2")[0] + self.assertEqual(3, len(tracker.getJobInfo(jobId).stageIds)) + + self.sc.setJobGroup("test3", "test", True) + d = sorted(parted.cogroup(parted).collect()) + self.assertEqual(10, len(d)) + self.assertEqual([[0], [0]], list(map(list, d[0][1]))) + jobId = tracker.getJobIdsForGroup("test3")[0] + self.assertEqual(2, len(tracker.getJobInfo(jobId).stageIds)) + + self.sc.setJobGroup("test4", "test", True) + d = sorted(parted.cogroup(rdd).collect()) + self.assertEqual(10, len(d)) + self.assertEqual([[0], [0]], list(map(list, d[0][1]))) + jobId = tracker.getJobIdsForGroup("test4")[0] + self.assertEqual(3, len(tracker.getJobInfo(jobId).stageIds)) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_join import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_memory_profiler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_memory_profiler.py new file mode 100644 index 0000000..31ad62e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_memory_profiler.py @@ -0,0 +1,167 @@ +# +# 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 os +import sys +import tempfile +import unittest +import warnings +from io import StringIO +from typing import Iterator +from unittest import mock + +from pyspark import SparkConf, SparkContext +from pyspark.profiler import has_memory_profiler +from pyspark.sql import SparkSession +from pyspark.sql.functions import pandas_udf, udf +from pyspark.testing.sqlutils import have_pandas, pandas_requirement_message +from pyspark.testing.utils import PySparkTestCase + + +@unittest.skipIf( + "COVERAGE_PROCESS_START" in os.environ, "Flaky with coverage enabled, skipping for now." +) +@unittest.skipIf(not has_memory_profiler, "Must have memory-profiler installed.") +@unittest.skipIf(not have_pandas, pandas_requirement_message) +class MemoryProfilerTests(PySparkTestCase): + def setUp(self): + self._old_sys_path = list(sys.path) + class_name = self.__class__.__name__ + conf = SparkConf().set("spark.python.profile.memory", "true") + self.sc = SparkContext("local[4]", class_name, conf=conf) + self.spark = SparkSession(sparkContext=self.sc) + + def test_memory_profiler(self): + self.exec_python_udf() + + profilers = self.sc.profiler_collector.profilers + self.assertEqual(1, len(profilers)) + id, profiler, _ = profilers[0] + stats = profiler.stats() + self.assertTrue(stats is not None) + + with mock.patch("sys.stdout", new=StringIO()) as fake_out: + self.sc.show_profiles() + self.assertTrue("plus_one" in fake_out.getvalue()) + + d = tempfile.gettempdir() + self.sc.dump_profiles(d) + self.assertTrue("udf_%d_memory.txt" % id in os.listdir(d)) + + def test_profile_pandas_udf(self): + udfs = [self.exec_pandas_udf_ser_to_ser, self.exec_pandas_udf_ser_to_scalar] + udf_names = ["ser_to_ser", "ser_to_scalar"] + for f, f_name in zip(udfs, udf_names): + f() + with mock.patch("sys.stdout", new=StringIO()) as fake_out: + self.sc.show_profiles() + self.assertTrue(f_name in fake_out.getvalue()) + + with warnings.catch_warnings(record=True) as warns: + warnings.simplefilter("always") + self.exec_pandas_udf_iter_to_iter() + user_warns = [warn.message for warn in warns if isinstance(warn.message, UserWarning)] + self.assertTrue(len(user_warns) > 0) + self.assertTrue( + "Profiling UDFs with iterators input/output is not supported" in str(user_warns[0]) + ) + + def test_profile_pandas_function_api(self): + apis = [self.exec_grouped_map] + f_names = ["grouped_map"] + for api, f_name in zip(apis, f_names): + api() + with mock.patch("sys.stdout", new=StringIO()) as fake_out: + self.sc.show_profiles() + self.assertTrue(f_name in fake_out.getvalue()) + + with warnings.catch_warnings(record=True) as warns: + warnings.simplefilter("always") + self.exec_map() + user_warns = [warn.message for warn in warns if isinstance(warn.message, UserWarning)] + self.assertTrue(len(user_warns) > 0) + self.assertTrue( + "Profiling UDFs with iterators input/output is not supported" in str(user_warns[0]) + ) + + def exec_python_udf(self): + @udf("int") + def plus_one(v): + return v + 1 + + self.spark.range(10).select(plus_one("id")).collect() + + def exec_pandas_udf_ser_to_ser(self): + import pandas as pd + + @pandas_udf("int") + def ser_to_ser(ser: pd.Series) -> pd.Series: + return ser + 1 + + self.spark.range(10).select(ser_to_ser("id")).collect() + + def exec_pandas_udf_ser_to_scalar(self): + import pandas as pd + + @pandas_udf("int") + def ser_to_scalar(ser: pd.Series) -> float: + return ser.median() + + self.spark.range(10).select(ser_to_scalar("id")).collect() + + # Unsupported + def exec_pandas_udf_iter_to_iter(self): + import pandas as pd + + @pandas_udf("int") + def iter_to_iter(batch_ser: Iterator[pd.Series]) -> Iterator[pd.Series]: + for ser in batch_ser: + yield ser + 1 + + self.spark.range(10).select(iter_to_iter("id")).collect() + + def exec_grouped_map(self): + import pandas as pd + + def grouped_map(pdf: pd.DataFrame) -> pd.DataFrame: + return pdf.assign(v=pdf.v - pdf.v.mean()) + + df = self.spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0)], ("id", "v")) + df.groupby("id").applyInPandas(grouped_map, schema="id long, v double").collect() + + # Unsupported + def exec_map(self): + import pandas as pd + + def map(pdfs: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: + for pdf in pdfs: + yield pdf[pdf.id == 1] + + df = self.spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0)], ("id", "v")) + df.mapInPandas(map, schema=df.schema).collect() + + +if __name__ == "__main__": + from pyspark.tests.test_memory_profiler import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_pin_thread.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_pin_thread.py new file mode 100644 index 0000000..975b549 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_pin_thread.py @@ -0,0 +1,194 @@ +# +# 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 os +import time +import threading +import unittest + +from pyspark import SparkContext, SparkConf, InheritableThread + + +class PinThreadTests(unittest.TestCase): + # These tests are in a separate class because it uses + # 'PYSPARK_PIN_THREAD' environment variable to test thread pin feature. + + @classmethod + def setUpClass(cls): + cls.old_pin_thread = os.environ.get("PYSPARK_PIN_THREAD") + os.environ["PYSPARK_PIN_THREAD"] = "true" + cls.sc = SparkContext("local[4]", cls.__name__, conf=SparkConf()) + + @classmethod + def tearDownClass(cls): + cls.sc.stop() + if cls.old_pin_thread is not None: + os.environ["PYSPARK_PIN_THREAD"] = cls.old_pin_thread + else: + del os.environ["PYSPARK_PIN_THREAD"] + + def test_pinned_thread(self): + threads = [] + exceptions = [] + property_name = "test_property_%s" % PinThreadTests.__name__ + jvm_thread_ids = [] + + for i in range(10): + + def test_local_property(): + jvm_thread_id = self.sc._jvm.java.lang.Thread.currentThread().getId() + jvm_thread_ids.append(jvm_thread_id) + + # If a property is set in this thread, later it should get the same property + # within this thread. + self.sc.setLocalProperty(property_name, str(i)) + + # 5 threads, 1 second sleep. 5 threads without a sleep. + time.sleep(i % 2) + + try: + assert self.sc.getLocalProperty(property_name) == str(i) + + # Each command might create a thread in multi-threading mode in Py4J. + # This assert makes sure that the created thread is being reused. + assert jvm_thread_id == self.sc._jvm.java.lang.Thread.currentThread().getId() + except Exception as e: + exceptions.append(e) + + threads.append(threading.Thread(target=test_local_property)) + + for t in threads: + t.start() + + for t in threads: + t.join() + + for e in exceptions: + raise e + + # Created JVM threads should be 10 because Python thread are 10. + assert len(set(jvm_thread_ids)) == 10 + + def test_multiple_group_jobs(self): + # SPARK-22340: Add a mode to pin Python thread into JVM's + self.check_job_cancellation( + lambda job_group: self.sc.setJobGroup( + job_group, "test rdd collect with setting job group" + ), + lambda job_group: self.sc.cancelJobGroup(job_group), + ) + + def test_multiple_group_tags(self): + # SPARK-44194: Test pinned thread mode with job tags. + self.check_job_cancellation( + lambda job_tag: self.sc.addJobTag(job_tag), + lambda job_tag: self.sc.cancelJobsWithTag(job_tag), + ) + + def check_job_cancellation(self, setter, canceller): + + job_id_a = "job_ids_to_cancel" + job_id_b = "job_ids_to_run" + + threads = [] + thread_ids = range(4) + thread_ids_to_cancel = [i for i in thread_ids if i % 2 == 0] + thread_ids_to_run = [i for i in thread_ids if i % 2 != 0] + + # A list which records whether job is cancelled. + # The index of the array is the thread index which job run in. + is_job_cancelled = [False for _ in thread_ids] + + def run_job(job_id, index): + """ + Executes a job with the group ``job_group``. Each job waits for 3 seconds + and then exits. + """ + try: + setter(job_id) + self.sc.parallelize([15]).map(lambda x: time.sleep(x)).collect() + is_job_cancelled[index] = False + except Exception: + # Assume that exception means job cancellation. + is_job_cancelled[index] = True + + # Test if job succeeded when not cancelled. + run_job(job_id_a, 0) + self.assertFalse(is_job_cancelled[0]) + + # Run jobs + for i in thread_ids_to_cancel: + t = threading.Thread(target=run_job, args=(job_id_a, i)) + t.start() + threads.append(t) + + for i in thread_ids_to_run: + t = threading.Thread(target=run_job, args=(job_id_b, i)) + t.start() + threads.append(t) + + # Wait to make sure all jobs are executed. + time.sleep(3) + # And then, cancel one job group. + canceller(job_id_a) + + # Wait until all threads launching jobs are finished. + for t in threads: + t.join() + + for i in thread_ids_to_cancel: + self.assertTrue( + is_job_cancelled[i], "Thread {i}: Job in group A was not cancelled.".format(i=i) + ) + + for i in thread_ids_to_run: + self.assertFalse( + is_job_cancelled[i], "Thread {i}: Job in group B did not succeeded.".format(i=i) + ) + + def test_inheritable_local_property(self): + self.sc.setLocalProperty("a", "hi") + expected = [] + + def get_inner_local_prop(): + expected.append(self.sc.getLocalProperty("b")) + + def get_outer_local_prop(): + expected.append(self.sc.getLocalProperty("a")) + self.sc.setLocalProperty("b", "hello") + t2 = InheritableThread(target=get_inner_local_prop) + t2.start() + t2.join() + + t1 = InheritableThread(target=get_outer_local_prop) + t1.start() + t1.join() + + self.assertEqual(self.sc.getLocalProperty("b"), None) + self.assertEqual(expected, ["hi", "hello"]) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_pin_thread import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_profiler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_profiler.py new file mode 100644 index 0000000..0183870 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_profiler.py @@ -0,0 +1,168 @@ +# +# 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 os +import sys +import tempfile +import unittest +from io import StringIO + +from pyspark import SparkConf, SparkContext, BasicProfiler +from pyspark.profiler import has_memory_profiler +from pyspark.sql import SparkSession +from pyspark.sql.functions import udf +from pyspark.errors import PythonException, PySparkRuntimeError +from pyspark.testing.utils import PySparkTestCase, PySparkErrorTestUtils + + +class ProfilerTests(PySparkTestCase): + def setUp(self): + self._old_sys_path = list(sys.path) + class_name = self.__class__.__name__ + conf = SparkConf().set("spark.python.profile", "true") + self.sc = SparkContext("local[4]", class_name, conf=conf) + + def test_profiler(self): + self.do_computation() + + profilers = self.sc.profiler_collector.profilers + self.assertEqual(1, len(profilers)) + id, profiler, _ = profilers[0] + stats = profiler.stats() + self.assertTrue(stats is not None) + width, stat_list = stats.get_print_list([]) + func_names = [func_name for fname, n, func_name in stat_list] + self.assertTrue("heavy_foo" in func_names) + + old_stdout = sys.stdout + sys.stdout = io = StringIO() + self.sc.show_profiles() + self.assertTrue("heavy_foo" in io.getvalue()) + sys.stdout = old_stdout + + d = tempfile.gettempdir() + self.sc.dump_profiles(d) + self.assertTrue("rdd_%d.pstats" % id in os.listdir(d)) + + def test_custom_profiler(self): + class TestCustomProfiler(BasicProfiler): + def show(self, id): + self.result = "Custom formatting" + + self.sc.profiler_collector.profiler_cls = TestCustomProfiler + + self.do_computation() + + profilers = self.sc.profiler_collector.profilers + self.assertEqual(1, len(profilers)) + _, profiler, _ = profilers[0] + self.assertTrue(isinstance(profiler, TestCustomProfiler)) + + self.sc.show_profiles() + self.assertEqual("Custom formatting", profiler.result) + + def do_computation(self): + def heavy_foo(x): + for i in range(1 << 18): + x = 1 # noqa: F841 + + rdd = self.sc.parallelize(range(100)) + rdd.foreach(heavy_foo) + + +class ProfilerTests2(unittest.TestCase, PySparkErrorTestUtils): + def test_profiler_disabled(self): + sc = SparkContext( + conf=SparkConf() + .set("spark.python.profile", "false") + .set("spark.python.profile.memory", "false") + ) + try: + with self.assertRaises(PySparkRuntimeError) as pe: + sc.show_profiles() + self.check_error( + exception=pe.exception, + error_class="INCORRECT_CONF_FOR_PROFILE", + message_parameters={}, + ) + + with self.assertRaises(PySparkRuntimeError) as pe: + sc.dump_profiles("/tmp/abc") + self.check_error( + exception=pe.exception, + error_class="INCORRECT_CONF_FOR_PROFILE", + message_parameters={}, + ) + finally: + sc.stop() + + def test_profiler_all_enabled(self): + sc = SparkContext( + conf=SparkConf() + .set("spark.python.profile", "true") + .set("spark.python.profile.memory", "true") + ) + spark = SparkSession(sparkContext=sc) + + @udf("int") + def plus_one(v): + return v + 1 + + try: + self.assertRaisesRegex( + RuntimeError, + "'spark.python.profile' and 'spark.python.profile.memory' configuration" + " cannot be enabled together", + lambda: spark.range(10).select(plus_one("id")).collect(), + ) + finally: + sc.stop() + + @unittest.skipIf(has_memory_profiler, "Test when memory-profiler is not installed.") + def test_no_memory_profile_installed(self): + sc = SparkContext( + conf=SparkConf() + .set("spark.python.profile", "false") + .set("spark.python.profile.memory", "true") + ) + spark = SparkSession(sparkContext=sc) + + @udf("int") + def plus_one(v): + return v + 1 + + try: + self.assertRaisesRegex( + PythonException, + "Install the 'memory_profiler' library in the cluster to enable memory " + "profiling", + lambda: spark.range(10).select(plus_one("id")).collect(), + ) + finally: + sc.stop() + + +if __name__ == "__main__": + from pyspark.tests.test_profiler import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rdd.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rdd.py new file mode 100644 index 0000000..752b5d5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rdd.py @@ -0,0 +1,939 @@ +# +# 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 datetime import datetime, timedelta +import hashlib +import os +import random +import tempfile +import time +import unittest +from glob import glob + +from py4j.protocol import Py4JJavaError + +from pyspark import shuffle, RDD +from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests +from pyspark.serializers import ( + CloudPickleSerializer, + BatchedSerializer, + CPickleSerializer, + MarshalSerializer, + UTF8Deserializer, + NoOpSerializer, +) +from pyspark.sql import SparkSession +from pyspark.testing.utils import ReusedPySparkTestCase, SPARK_HOME, QuietTest, have_numpy +from pyspark.testing.sqlutils import have_pandas + + +global_func = lambda: "Hi" # noqa: E731 + + +class RDDTests(ReusedPySparkTestCase): + def test_range(self): + self.assertEqual(self.sc.range(1, 1).count(), 0) + self.assertEqual(self.sc.range(1, 0, -1).count(), 1) + self.assertEqual(self.sc.range(0, 1 << 40, 1 << 39).count(), 2) + + def test_id(self): + rdd = self.sc.parallelize(range(10)) + id = rdd.id() + self.assertEqual(id, rdd.id()) + rdd2 = rdd.map(str).filter(bool) + id2 = rdd2.id() + self.assertEqual(id + 1, id2) + self.assertEqual(id2, rdd2.id()) + + def test_empty_rdd(self): + rdd = self.sc.emptyRDD() + self.assertTrue(rdd.isEmpty()) + + def test_sum(self): + self.assertEqual(0, self.sc.emptyRDD().sum()) + self.assertEqual(6, self.sc.parallelize([1, 2, 3]).sum()) + + def test_to_localiterator(self): + rdd = self.sc.parallelize([1, 2, 3]) + it = rdd.toLocalIterator() + self.assertEqual([1, 2, 3], sorted(it)) + + rdd2 = rdd.repartition(1000) + it2 = rdd2.toLocalIterator() + self.assertEqual([1, 2, 3], sorted(it2)) + + def test_to_localiterator_prefetch(self): + # Test that we fetch the next partition in parallel + # We do this by returning the current time and: + # reading the first elem, waiting, and reading the second elem + # If not in parallel then these would be at different times + # But since they are being computed in parallel we see the time + # is "close enough" to the same. + rdd = self.sc.parallelize(range(2), 2) + times1 = rdd.map(lambda x: datetime.now()) + times2 = rdd.map(lambda x: datetime.now()) + times_iter_prefetch = times1.toLocalIterator(prefetchPartitions=True) + times_iter = times2.toLocalIterator(prefetchPartitions=False) + times_prefetch_head = next(times_iter_prefetch) + times_head = next(times_iter) + time.sleep(2) + times_next = next(times_iter) + times_prefetch_next = next(times_iter_prefetch) + self.assertTrue(times_next - times_head >= timedelta(seconds=2)) + self.assertTrue(times_prefetch_next - times_prefetch_head < timedelta(seconds=1)) + + def test_save_as_textfile_with_unicode(self): + # Regression test for SPARK-970 + x = "\u00A1Hola, mundo!" + data = self.sc.parallelize([x]) + tempFile = tempfile.NamedTemporaryFile(delete=True) + tempFile.close() + data.saveAsTextFile(tempFile.name) + raw_contents = b"".join(open(p, "rb").read() for p in glob(tempFile.name + "/part-0000*")) + self.assertEqual(x, raw_contents.strip().decode("utf-8")) + + def test_save_as_textfile_with_utf8(self): + x = "\u00A1Hola, mundo!" + data = self.sc.parallelize([x.encode("utf-8")]) + tempFile = tempfile.NamedTemporaryFile(delete=True) + tempFile.close() + data.saveAsTextFile(tempFile.name) + raw_contents = b"".join(open(p, "rb").read() for p in glob(tempFile.name + "/part-0000*")) + self.assertEqual(x, raw_contents.strip().decode("utf8")) + + def test_transforming_cartesian_result(self): + # Regression test for SPARK-1034 + rdd1 = self.sc.parallelize([1, 2]) + rdd2 = self.sc.parallelize([3, 4]) + cart = rdd1.cartesian(rdd2) + cart.map(lambda x_y3: x_y3[0] + x_y3[1]).collect() + + def test_transforming_pickle_file(self): + # Regression test for SPARK-2601 + data = self.sc.parallelize(["Hello", "World!"]) + tempFile = tempfile.NamedTemporaryFile(delete=True) + tempFile.close() + data.saveAsPickleFile(tempFile.name) + pickled_file = self.sc.pickleFile(tempFile.name) + pickled_file.map(lambda x: x).collect() + + def test_cartesian_on_textfile(self): + # Regression test for + path = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + a = self.sc.textFile(path) + result = a.cartesian(a).collect() + (x, y) = result[0] + self.assertEqual("Hello World!", x.strip()) + self.assertEqual("Hello World!", y.strip()) + + def test_cartesian_chaining(self): + # Tests for SPARK-16589 + rdd = self.sc.parallelize(range(10), 2) + self.assertSetEqual( + set(rdd.cartesian(rdd).cartesian(rdd).collect()), + set([((x, y), z) for x in range(10) for y in range(10) for z in range(10)]), + ) + + self.assertSetEqual( + set(rdd.cartesian(rdd.cartesian(rdd)).collect()), + set([(x, (y, z)) for x in range(10) for y in range(10) for z in range(10)]), + ) + + self.assertSetEqual( + set(rdd.cartesian(rdd.zip(rdd)).collect()), + set([(x, (y, y)) for x in range(10) for y in range(10)]), + ) + + def test_zip_chaining(self): + # Tests for SPARK-21985 + rdd = self.sc.parallelize("abc", 2) + self.assertSetEqual(set(rdd.zip(rdd).zip(rdd).collect()), set([((x, x), x) for x in "abc"])) + self.assertSetEqual(set(rdd.zip(rdd.zip(rdd)).collect()), set([(x, (x, x)) for x in "abc"])) + + def test_union_pair_rdd(self): + # SPARK-31788: test if pair RDDs can be combined by union. + rdd = self.sc.parallelize([1, 2]) + pair_rdd = rdd.zip(rdd) + unionRDD = self.sc.union([pair_rdd, pair_rdd]) + self.assertEqual(set(unionRDD.collect()), set([(1, 1), (2, 2), (1, 1), (2, 2)])) + self.assertEqual(unionRDD.count(), 4) + + def test_deleting_input_files(self): + # Regression test for SPARK-1025 + tempFile = tempfile.NamedTemporaryFile(delete=False) + tempFile.write(b"Hello World!") + tempFile.close() + data = self.sc.textFile(tempFile.name) + filtered_data = data.filter(lambda x: True) + self.assertEqual(1, filtered_data.count()) + os.unlink(tempFile.name) + with QuietTest(self.sc): + self.assertRaises(Exception, lambda: filtered_data.count()) + + def test_sampling_default_seed(self): + # Test for SPARK-3995 (default seed setting) + data = self.sc.parallelize(range(1000), 1) + subset = data.takeSample(False, 10) + self.assertEqual(len(subset), 10) + + def test_aggregate_mutable_zero_value(self): + # Test for SPARK-9021; uses aggregate and treeAggregate to build dict + # representing a counter of ints + from collections import defaultdict + + # Show that single or multiple partitions work + data1 = self.sc.range(10, numSlices=1) + data2 = self.sc.range(10, numSlices=2) + + def seqOp(x, y): + x[y] += 1 + return x + + def comboOp(x, y): + for key, val in y.items(): + x[key] += val + return x + + counts1 = data1.aggregate(defaultdict(int), seqOp, comboOp) + counts2 = data2.aggregate(defaultdict(int), seqOp, comboOp) + counts3 = data1.treeAggregate(defaultdict(int), seqOp, comboOp, 2) + counts4 = data2.treeAggregate(defaultdict(int), seqOp, comboOp, 2) + + ground_truth = defaultdict(int, dict((i, 1) for i in range(10))) + self.assertEqual(counts1, ground_truth) + self.assertEqual(counts2, ground_truth) + self.assertEqual(counts3, ground_truth) + self.assertEqual(counts4, ground_truth) + + def test_aggregate_by_key_mutable_zero_value(self): + # Test for SPARK-9021; uses aggregateByKey to make a pair RDD that + # contains lists of all values for each key in the original RDD + + # list(range(...)) for Python 3.x compatibility (can't use * operator + # on a range object) + # list(zip(...)) for Python 3.x compatibility (want to parallelize a + # collection, not a zip object) + tuples = list(zip(list(range(10)) * 2, [1] * 20)) + # Show that single or multiple partitions work + data1 = self.sc.parallelize(tuples, 1) + data2 = self.sc.parallelize(tuples, 2) + + def seqOp(x, y): + x.append(y) + return x + + def comboOp(x, y): + x.extend(y) + return x + + values1 = data1.aggregateByKey([], seqOp, comboOp).collect() + values2 = data2.aggregateByKey([], seqOp, comboOp).collect() + # Sort lists to ensure clean comparison with ground_truth + values1.sort() + values2.sort() + + ground_truth = [(i, [1] * 2) for i in range(10)] + self.assertEqual(values1, ground_truth) + self.assertEqual(values2, ground_truth) + + def test_fold_mutable_zero_value(self): + # Test for SPARK-9021; uses fold to merge an RDD of dict counters into + # a single dict + from collections import defaultdict + + counts1 = defaultdict(int, dict((i, 1) for i in range(10))) + counts2 = defaultdict(int, dict((i, 1) for i in range(3, 8))) + counts3 = defaultdict(int, dict((i, 1) for i in range(4, 7))) + counts4 = defaultdict(int, dict((i, 1) for i in range(5, 6))) + all_counts = [counts1, counts2, counts3, counts4] + # Show that single or multiple partitions work + data1 = self.sc.parallelize(all_counts, 1) + data2 = self.sc.parallelize(all_counts, 2) + + def comboOp(x, y): + for key, val in y.items(): + x[key] += val + return x + + fold1 = data1.fold(defaultdict(int), comboOp) + fold2 = data2.fold(defaultdict(int), comboOp) + + ground_truth = defaultdict(int) + for counts in all_counts: + for key, val in counts.items(): + ground_truth[key] += val + self.assertEqual(fold1, ground_truth) + self.assertEqual(fold2, ground_truth) + + def test_fold_by_key_mutable_zero_value(self): + # Test for SPARK-9021; uses foldByKey to make a pair RDD that contains + # lists of all values for each key in the original RDD + + tuples = [(i, range(i)) for i in range(10)] * 2 + # Show that single or multiple partitions work + data1 = self.sc.parallelize(tuples, 1) + data2 = self.sc.parallelize(tuples, 2) + + def comboOp(x, y): + x.extend(y) + return x + + values1 = data1.foldByKey([], comboOp).collect() + values2 = data2.foldByKey([], comboOp).collect() + # Sort lists to ensure clean comparison with ground_truth + values1.sort() + values2.sort() + + # list(range(...)) for Python 3.x compatibility + ground_truth = [(i, list(range(i)) * 2) for i in range(10)] + self.assertEqual(values1, ground_truth) + self.assertEqual(values2, ground_truth) + + def test_aggregate_by_key(self): + data = self.sc.parallelize([(1, 1), (1, 1), (3, 2), (5, 1), (5, 3)], 2) + + def seqOp(x, y): + x.add(y) + return x + + def combOp(x, y): + x |= y + return x + + sets = dict(data.aggregateByKey(set(), seqOp, combOp).collect()) + self.assertEqual(3, len(sets)) + self.assertEqual(set([1]), sets[1]) + self.assertEqual(set([2]), sets[3]) + self.assertEqual(set([1, 3]), sets[5]) + + def test_itemgetter(self): + rdd = self.sc.parallelize([range(10)]) + from operator import itemgetter + + self.assertEqual([1], rdd.map(itemgetter(1)).collect()) + self.assertEqual([(2, 3)], rdd.map(itemgetter(2, 3)).collect()) + + def test_namedtuple_in_rdd(self): + from collections import namedtuple + + Person = namedtuple("Person", "id firstName lastName") + jon = Person(1, "Jon", "Doe") + jane = Person(2, "Jane", "Doe") + theDoes = self.sc.parallelize([jon, jane]) + self.assertEqual([jon, jane], theDoes.collect()) + + def test_large_broadcast(self): + N = 10000 + data = [[float(i) for i in range(300)] for i in range(N)] + bdata = self.sc.broadcast(data) # 27MB + m = self.sc.parallelize(range(1), 1).map(lambda x: len(bdata.value)).sum() + self.assertEqual(N, m) + + def test_unpersist(self): + N = 1000 + data = [[float(i) for i in range(300)] for i in range(N)] + bdata = self.sc.broadcast(data) # 3MB + bdata.unpersist() + m = self.sc.parallelize(range(1), 1).map(lambda x: len(bdata.value)).sum() + self.assertEqual(N, m) + bdata.destroy(blocking=True) + try: + self.sc.parallelize(range(1), 1).map(lambda x: len(bdata.value)).sum() + except Exception: + pass + else: + raise AssertionError("job should fail after destroy the broadcast") + + def test_multiple_broadcasts(self): + N = 1 << 21 + b1 = self.sc.broadcast(set(range(N))) # multiple blocks in JVM + r = list(range(1 << 15)) + random.shuffle(r) + s = str(r).encode() + checksum = hashlib.md5(s).hexdigest() + b2 = self.sc.broadcast(s) + r = list( + set( + self.sc.parallelize(range(10), 10) + .map(lambda x: (len(b1.value), hashlib.md5(b2.value).hexdigest())) + .collect() + ) + ) + self.assertEqual(1, len(r)) + size, csum = r[0] + self.assertEqual(N, size) + self.assertEqual(checksum, csum) + + random.shuffle(r) + s = str(r).encode() + checksum = hashlib.md5(s).hexdigest() + b2 = self.sc.broadcast(s) + r = list( + set( + self.sc.parallelize(range(10), 10) + .map(lambda x: (len(b1.value), hashlib.md5(b2.value).hexdigest())) + .collect() + ) + ) + self.assertEqual(1, len(r)) + size, csum = r[0] + self.assertEqual(N, size) + self.assertEqual(checksum, csum) + + def test_multithread_broadcast_pickle(self): + import threading + + b1 = self.sc.broadcast(list(range(3))) + b2 = self.sc.broadcast(list(range(3))) + + def f1(): + return b1.value + + def f2(): + return b2.value + + funcs_num_pickled = {f1: None, f2: None} + + def do_pickle(f, sc): + command = (f, None, sc.serializer, sc.serializer) + ser = CloudPickleSerializer() + ser.dumps(command) + + def process_vars(sc): + broadcast_vars = list(sc._pickled_broadcast_vars) + num_pickled = len(broadcast_vars) + sc._pickled_broadcast_vars.clear() + return num_pickled + + def run(f, sc): + do_pickle(f, sc) + funcs_num_pickled[f] = process_vars(sc) + + # pickle f1, adds b1 to sc._pickled_broadcast_vars in main thread local storage + do_pickle(f1, self.sc) + + # run all for f2, should only add/count/clear b2 from worker thread local storage + t = threading.Thread(target=run, args=(f2, self.sc)) + t.start() + t.join() + + # count number of vars pickled in main thread, only b1 should be counted and cleared + funcs_num_pickled[f1] = process_vars(self.sc) + + self.assertEqual(funcs_num_pickled[f1], 1) + self.assertEqual(funcs_num_pickled[f2], 1) + self.assertEqual(len(list(self.sc._pickled_broadcast_vars)), 0) + + def test_large_closure(self): + N = 200000 + data = [float(i) for i in range(N)] + rdd = self.sc.parallelize(range(1), 1).map(lambda x: len(data)) + self.assertEqual(N, rdd.first()) + # regression test for SPARK-6886 + self.assertEqual(1, rdd.map(lambda x: (x, 1)).groupByKey().count()) + + def test_zip_with_different_serializers(self): + a = self.sc.parallelize(range(5)) + b = self.sc.parallelize(range(100, 105)) + self.assertEqual(a.zip(b).collect(), [(0, 100), (1, 101), (2, 102), (3, 103), (4, 104)]) + a = a._reserialize(BatchedSerializer(CPickleSerializer(), 2)) + b = b._reserialize(MarshalSerializer()) + self.assertEqual(a.zip(b).collect(), [(0, 100), (1, 101), (2, 102), (3, 103), (4, 104)]) + # regression test for SPARK-4841 + path = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + t = self.sc.textFile(path) + cnt = t.count() + self.assertEqual(cnt, t.zip(t).count()) + rdd = t.map(str) + self.assertEqual(cnt, t.zip(rdd).count()) + # regression test for bug in _reserializer() + self.assertEqual(cnt, t.zip(rdd).count()) + + def test_zip_with_different_object_sizes(self): + # regress test for SPARK-5973 + a = self.sc.parallelize(range(10000)).map(lambda i: "*" * i) + b = self.sc.parallelize(range(10000, 20000)).map(lambda i: "*" * i) + self.assertEqual(10000, a.zip(b).count()) + + def test_zip_with_different_number_of_items(self): + a = self.sc.parallelize(range(5), 2) + # different number of partitions + b = self.sc.parallelize(range(100, 106), 3) + self.assertRaises(ValueError, lambda: a.zip(b)) + with QuietTest(self.sc): + # different number of batched items in JVM + b = self.sc.parallelize(range(100, 104), 2) + self.assertRaises(Exception, lambda: a.zip(b).count()) + # different number of items in one pair + b = self.sc.parallelize(range(100, 106), 2) + self.assertRaises(Exception, lambda: a.zip(b).count()) + # same total number of items, but different distributions + a = self.sc.parallelize([2, 3], 2).flatMap(range) + b = self.sc.parallelize([3, 2], 2).flatMap(range) + self.assertEqual(a.count(), b.count()) + self.assertRaises(Exception, lambda: a.zip(b).count()) + + def test_count_approx_distinct(self): + rdd = self.sc.parallelize(range(1000)) + self.assertTrue(950 < rdd.countApproxDistinct(0.03) < 1050) + self.assertTrue(950 < rdd.map(float).countApproxDistinct(0.03) < 1050) + self.assertTrue(950 < rdd.map(str).countApproxDistinct(0.03) < 1050) + self.assertTrue(950 < rdd.map(lambda x: (x, -x)).countApproxDistinct(0.03) < 1050) + + rdd = self.sc.parallelize([i % 20 for i in range(1000)], 7) + self.assertTrue(18 < rdd.countApproxDistinct() < 22) + self.assertTrue(18 < rdd.map(float).countApproxDistinct() < 22) + self.assertTrue(18 < rdd.map(str).countApproxDistinct() < 22) + self.assertTrue(18 < rdd.map(lambda x: (x, -x)).countApproxDistinct() < 22) + + self.assertRaises(ValueError, lambda: rdd.countApproxDistinct(0.00000001)) + + def test_histogram(self): + # empty + rdd = self.sc.parallelize([]) + self.assertEqual([0], rdd.histogram([0, 10])[1]) + self.assertEqual([0, 0], rdd.histogram([0, 4, 10])[1]) + self.assertRaises(ValueError, lambda: rdd.histogram(1)) + + # out of range + rdd = self.sc.parallelize([10.01, -0.01]) + self.assertEqual([0], rdd.histogram([0, 10])[1]) + self.assertEqual([0, 0], rdd.histogram((0, 4, 10))[1]) + + # in range with one bucket + rdd = self.sc.parallelize(range(1, 5)) + self.assertEqual([4], rdd.histogram([0, 10])[1]) + self.assertEqual([3, 1], rdd.histogram([0, 4, 10])[1]) + + # in range with one bucket exact match + self.assertEqual([4], rdd.histogram([1, 4])[1]) + + # out of range with two buckets + rdd = self.sc.parallelize([10.01, -0.01]) + self.assertEqual([0, 0], rdd.histogram([0, 5, 10])[1]) + + # out of range with two uneven buckets + rdd = self.sc.parallelize([10.01, -0.01]) + self.assertEqual([0, 0], rdd.histogram([0, 4, 10])[1]) + + # in range with two buckets + rdd = self.sc.parallelize([1, 2, 3, 5, 6]) + self.assertEqual([3, 2], rdd.histogram([0, 5, 10])[1]) + + # in range with two bucket and None + rdd = self.sc.parallelize([1, 2, 3, 5, 6, None, float("nan")]) + self.assertEqual([3, 2], rdd.histogram([0, 5, 10])[1]) + + # in range with two uneven buckets + rdd = self.sc.parallelize([1, 2, 3, 5, 6]) + self.assertEqual([3, 2], rdd.histogram([0, 5, 11])[1]) + + # mixed range with two uneven buckets + rdd = self.sc.parallelize([-0.01, 0.0, 1, 2, 3, 5, 6, 11.0, 11.01]) + self.assertEqual([4, 3], rdd.histogram([0, 5, 11])[1]) + + # mixed range with four uneven buckets + rdd = self.sc.parallelize([-0.01, 0.0, 1, 2, 3, 5, 6, 11.01, 12.0, 199.0, 200.0, 200.1]) + self.assertEqual([4, 2, 1, 3], rdd.histogram([0.0, 5.0, 11.0, 12.0, 200.0])[1]) + + # mixed range with uneven buckets and NaN + rdd = self.sc.parallelize( + [-0.01, 0.0, 1, 2, 3, 5, 6, 11.01, 12.0, 199.0, 200.0, 200.1, None, float("nan")] + ) + self.assertEqual([4, 2, 1, 3], rdd.histogram([0.0, 5.0, 11.0, 12.0, 200.0])[1]) + + # out of range with infinite buckets + rdd = self.sc.parallelize([10.01, -0.01, float("nan"), float("inf")]) + self.assertEqual([1, 2], rdd.histogram([float("-inf"), 0, float("inf")])[1]) + + # invalid buckets + self.assertRaises(ValueError, lambda: rdd.histogram([])) + self.assertRaises(ValueError, lambda: rdd.histogram([1])) + self.assertRaises(ValueError, lambda: rdd.histogram(0)) + self.assertRaises(TypeError, lambda: rdd.histogram({})) + + # without buckets + rdd = self.sc.parallelize(range(1, 5)) + self.assertEqual(([1, 4], [4]), rdd.histogram(1)) + + # without buckets single element + rdd = self.sc.parallelize([1]) + self.assertEqual(([1, 1], [1]), rdd.histogram(1)) + + # without bucket no range + rdd = self.sc.parallelize([1] * 4) + self.assertEqual(([1, 1], [4]), rdd.histogram(1)) + + # without buckets basic two + rdd = self.sc.parallelize(range(1, 5)) + self.assertEqual(([1, 2.5, 4], [2, 2]), rdd.histogram(2)) + + # without buckets with more requested than elements + rdd = self.sc.parallelize([1, 2]) + buckets = [1 + 0.2 * i for i in range(6)] + hist = [1, 0, 0, 0, 1] + self.assertEqual((buckets, hist), rdd.histogram(5)) + + # invalid RDDs + rdd = self.sc.parallelize([1, float("inf")]) + self.assertRaises(ValueError, lambda: rdd.histogram(2)) + rdd = self.sc.parallelize([float("nan")]) + self.assertRaises(ValueError, lambda: rdd.histogram(2)) + + # string + rdd = self.sc.parallelize(["ab", "ac", "b", "bd", "ef"], 2) + self.assertEqual([2, 2], rdd.histogram(["a", "b", "c"])[1]) + self.assertEqual((["ab", "ef"], [5]), rdd.histogram(1)) + self.assertRaises(TypeError, lambda: rdd.histogram(2)) + + def test_repartitionAndSortWithinPartitions_asc(self): + rdd = self.sc.parallelize([(0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)], 2) + + repartitioned = rdd.repartitionAndSortWithinPartitions(2, lambda key: key % 2, True) + partitions = repartitioned.glom().collect() + self.assertEqual(partitions[0], [(0, 5), (0, 8), (2, 6)]) + self.assertEqual(partitions[1], [(1, 3), (3, 8), (3, 8)]) + + def test_repartitionAndSortWithinPartitions_desc(self): + rdd = self.sc.parallelize([(0, 5), (3, 8), (2, 6), (0, 8), (3, 8), (1, 3)], 2) + + repartitioned = rdd.repartitionAndSortWithinPartitions(2, lambda key: key % 2, False) + partitions = repartitioned.glom().collect() + self.assertEqual(partitions[0], [(2, 6), (0, 5), (0, 8)]) + self.assertEqual(partitions[1], [(3, 8), (3, 8), (1, 3)]) + + def test_repartition_no_skewed(self): + num_partitions = 20 + a = self.sc.parallelize(range(int(1000)), 2) + xs = a.repartition(num_partitions).glom().map(len).collect() + zeros = len([x for x in xs if x == 0]) + self.assertTrue(zeros == 0) + xs = a.coalesce(num_partitions, True).glom().map(len).collect() + zeros = len([x for x in xs if x == 0]) + self.assertTrue(zeros == 0) + + def test_repartition_on_textfile(self): + path = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + rdd = self.sc.textFile(path) + result = rdd.repartition(1).collect() + self.assertEqual("Hello World!", result[0]) + + def test_distinct(self): + rdd = self.sc.parallelize((1, 2, 3) * 10, 10) + self.assertEqual(rdd.getNumPartitions(), 10) + self.assertEqual(rdd.distinct().count(), 3) + result = rdd.distinct(5) + self.assertEqual(result.getNumPartitions(), 5) + self.assertEqual(result.count(), 3) + + def test_external_group_by_key(self): + self.sc._conf.set("spark.python.worker.memory", "1m") + N = 2000001 + kv = self.sc.parallelize(range(N)).map(lambda x: (x % 3, x)) + gkv = kv.groupByKey().cache() + self.assertEqual(3, gkv.count()) + filtered = gkv.filter(lambda kv: kv[0] == 1) + self.assertEqual(1, filtered.count()) + self.assertEqual([(1, N // 3)], filtered.mapValues(len).collect()) + self.assertEqual( + [(N // 3, N // 3)], filtered.values().map(lambda x: (len(x), len(list(x)))).collect() + ) + result = filtered.collect()[0][1] + self.assertEqual(N // 3, len(result)) + self.assertTrue(isinstance(result.data, shuffle.ExternalListOfList)) + + def test_sort_on_empty_rdd(self): + self.assertEqual([], self.sc.parallelize(zip([], [])).sortByKey().collect()) + + def test_sample(self): + rdd = self.sc.parallelize(range(0, 100), 4) + wo = rdd.sample(False, 0.1, 2).collect() + wo_dup = rdd.sample(False, 0.1, 2).collect() + self.assertSetEqual(set(wo), set(wo_dup)) + wr = rdd.sample(True, 0.2, 5).collect() + wr_dup = rdd.sample(True, 0.2, 5).collect() + self.assertSetEqual(set(wr), set(wr_dup)) + wo_s10 = rdd.sample(False, 0.3, 10).collect() + wo_s20 = rdd.sample(False, 0.3, 20).collect() + self.assertNotEqual(set(wo_s10), set(wo_s20)) + wr_s11 = rdd.sample(True, 0.4, 11).collect() + wr_s21 = rdd.sample(True, 0.4, 21).collect() + self.assertNotEqual(set(wr_s11), set(wr_s21)) + + def test_null_in_rdd(self): + jrdd = self.sc._jvm.PythonUtils.generateRDDWithNull(self.sc._jsc) + rdd = RDD(jrdd, self.sc, UTF8Deserializer()) + self.assertEqual(["a", None, "b"], rdd.collect()) + rdd = RDD(jrdd, self.sc, NoOpSerializer()) + self.assertEqual([b"a", None, b"b"], rdd.collect()) + + def test_multiple_python_java_RDD_conversions(self): + # Regression test for SPARK-5361 + data = [("1", {"director": "David Lean"}), ("2", {"director": "Andrew Dominik"})] + data_rdd = self.sc.parallelize(data) + data_java_rdd = data_rdd._to_java_object_rdd() + data_python_rdd = self.sc._jvm.SerDeUtil.javaToPython(data_java_rdd) + converted_rdd = RDD(data_python_rdd, self.sc) + self.assertEqual(2, converted_rdd.count()) + + # conversion between python and java RDD threw exceptions + data_java_rdd = converted_rdd._to_java_object_rdd() + data_python_rdd = self.sc._jvm.SerDeUtil.javaToPython(data_java_rdd) + converted_rdd = RDD(data_python_rdd, self.sc) + self.assertEqual(2, converted_rdd.count()) + + # Regression test for SPARK-6294 + def test_take_on_jrdd(self): + rdd = self.sc.parallelize(range(1 << 20)).map(lambda x: str(x)) + rdd._jrdd.first() + + @unittest.skipIf(not have_numpy or not have_pandas, "NumPy or Pandas not installed") + def test_take_on_jrdd_with_large_rows_should_not_cause_deadlock(self): + # Regression test for SPARK-38677. + # + # Create a DataFrame with many columns, call a Python function on each row, and take only + # the first result row. + # + # This produces large rows that trigger a deadlock involving the following three threads: + # + # 1. The Scala task executor thread. During task execution, this is responsible for reading + # output produced by the Python process. However, in this case the task has finished + # early, and this thread is no longer reading output produced by the Python process. + # Instead, it is waiting for the Scala WriterThread to exit so that it can finish the + # task. + # + # 2. The Scala WriterThread. This is trying to send a large row to the Python process, and + # is waiting for the Python process to read that row. + # + # 3. The Python process. This is trying to send a large output to the Scala task executor + # thread, and is waiting for that thread to read that output. + # + # For this test to succeed rather than hanging, the Scala MonitorThread must detect this + # deadlock and kill the Python worker. + import numpy as np + import pandas as pd + + num_rows = 100000 + num_columns = 134 + data = np.zeros((num_rows, num_columns)) + columns = map(str, range(num_columns)) + df = SparkSession(self.sc).createDataFrame(pd.DataFrame(data, columns=columns)) + actual = CPickleSerializer().loads(df.rdd.map(list)._jrdd.first()) + expected = [list(data[0])] + self.assertEqual(expected, actual) + + def test_sortByKey_uses_all_partitions_not_only_first_and_last(self): + # Regression test for SPARK-5969 + seq = [(i * 59 % 101, i) for i in range(101)] # unsorted sequence + rdd = self.sc.parallelize(seq) + for ascending in [True, False]: + sort = rdd.sortByKey(ascending=ascending, numPartitions=5) + self.assertEqual(sort.collect(), sorted(seq, reverse=not ascending)) + sizes = sort.glom().map(len).collect() + for size in sizes: + self.assertGreater(size, 0) + + def test_pipe_functions(self): + data = ["1", "2", "3"] + rdd = self.sc.parallelize(data) + with QuietTest(self.sc): + self.assertEqual([], rdd.pipe("java").collect()) + self.assertRaises(Py4JJavaError, rdd.pipe("java", checkCode=True).collect) + result = rdd.pipe("cat").collect() + result.sort() + for x, y in zip(data, result): + self.assertEqual(x, y) + self.assertRaises(Py4JJavaError, rdd.pipe("grep 4", checkCode=True).collect) + self.assertEqual([], rdd.pipe("grep 4").collect()) + + def test_pipe_unicode(self): + # Regression test for SPARK-20947 + data = ["\u6d4b\u8bd5", "1"] + rdd = self.sc.parallelize(data) + result = rdd.pipe("cat").collect() + self.assertEqual(data, result) + + def test_stopiteration_in_user_code(self): + def stopit(*x): + raise StopIteration() + + seq_rdd = self.sc.parallelize(range(10)) + keyed_rdd = self.sc.parallelize((x % 2, x) for x in range(10)) + msg = "Caught StopIteration thrown from user's code; failing the task" + + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.map(stopit).collect) + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.filter(stopit).collect) + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.foreach, stopit) + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.reduce, stopit) + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.fold, 0, stopit) + self.assertRaisesRegex(Py4JJavaError, msg, seq_rdd.foreach, stopit) + self.assertRaisesRegex( + Py4JJavaError, msg, seq_rdd.cartesian(seq_rdd).flatMap(stopit).collect + ) + + # these methods call the user function both in the driver and in the executor + # the exception raised is different according to where the StopIteration happens + # RuntimeError is raised if in the driver + # Py4JJavaError is raised if in the executor (wraps the RuntimeError raised in the worker) + self.assertRaisesRegex( + (Py4JJavaError, RuntimeError), msg, keyed_rdd.reduceByKeyLocally, stopit + ) + self.assertRaisesRegex( + (Py4JJavaError, RuntimeError), msg, seq_rdd.aggregate, 0, stopit, lambda *x: 1 + ) + self.assertRaisesRegex( + (Py4JJavaError, RuntimeError), msg, seq_rdd.aggregate, 0, lambda *x: 1, stopit + ) + + def test_overwritten_global_func(self): + # Regression test for SPARK-27000 + global global_func + self.assertEqual(self.sc.parallelize([1]).map(lambda _: global_func()).first(), "Hi") + global_func = lambda: "Yeah" # noqa: E731 + self.assertEqual(self.sc.parallelize([1]).map(lambda _: global_func()).first(), "Yeah") + + def test_to_local_iterator_failure(self): + # SPARK-27548 toLocalIterator task failure not propagated to Python driver + + def fail(_): + raise RuntimeError("local iterator error") + + rdd = self.sc.range(10).map(fail) + + with self.assertRaisesRegex(Exception, "local iterator error"): + for _ in rdd.toLocalIterator(): + pass + + def test_to_local_iterator_collects_single_partition(self): + # Test that partitions are not computed until requested by iteration + + def fail_last(x): + if x == 9: + raise RuntimeError("This should not be hit") + return x + + rdd = self.sc.range(12, numSlices=4).map(fail_last) + it = rdd.toLocalIterator() + + # Only consume first 4 elements from partitions 1 and 2, this should not collect the last + # partition which would trigger the error + for i in range(4): + self.assertEqual(i, next(it)) + + def test_resourceprofile(self): + rp_builder = ResourceProfileBuilder() + ereqs = ExecutorResourceRequests().cores(2).memory("6g").memoryOverhead("1g") + ereqs.pysparkMemory("2g").resource("gpu", 2, "testGpus", "nvidia.com") + treqs = TaskResourceRequests().cpus(2).resource("gpu", 2) + + def assert_request_contents(exec_reqs, task_reqs): + self.assertEqual(len(exec_reqs), 5) + self.assertEqual(exec_reqs["cores"].amount, 2) + self.assertEqual(exec_reqs["memory"].amount, 6144) + self.assertEqual(exec_reqs["memoryOverhead"].amount, 1024) + self.assertEqual(exec_reqs["pyspark.memory"].amount, 2048) + self.assertEqual(exec_reqs["gpu"].amount, 2) + self.assertEqual(exec_reqs["gpu"].discoveryScript, "testGpus") + self.assertEqual(exec_reqs["gpu"].resourceName, "gpu") + self.assertEqual(exec_reqs["gpu"].vendor, "nvidia.com") + self.assertEqual(len(task_reqs), 2) + self.assertEqual(task_reqs["cpus"].amount, 2.0) + self.assertEqual(task_reqs["gpu"].amount, 2.0) + + assert_request_contents(ereqs.requests, treqs.requests) + rp = rp_builder.require(ereqs).require(treqs).build + assert_request_contents(rp.executorResources, rp.taskResources) + rdd = self.sc.parallelize(range(10)).withResources(rp) + return_rp = rdd.getResourceProfile() + assert_request_contents(return_rp.executorResources, return_rp.taskResources) + rddWithoutRp = self.sc.parallelize(range(10)) + self.assertEqual(rddWithoutRp.getResourceProfile(), None) + + def test_multiple_group_jobs(self): + import threading + + group_a = "job_ids_to_cancel" + group_b = "job_ids_to_run" + + threads = [] + thread_ids = range(4) + thread_ids_to_cancel = [i for i in thread_ids if i % 2 == 0] + thread_ids_to_run = [i for i in thread_ids if i % 2 != 0] + + # A list which records whether job is cancelled. + # The index of the array is the thread index which job run in. + is_job_cancelled = [False for _ in thread_ids] + + def run_job(job_group, index): + """ + Executes a job with the group ``job_group``. Each job waits for 3 seconds + and then exits. + """ + try: + self.sc.parallelize([15]).map(lambda x: time.sleep(x)).collectWithJobGroup( + job_group, "test rdd collect with setting job group" + ) + is_job_cancelled[index] = False + except Exception: + # Assume that exception means job cancellation. + is_job_cancelled[index] = True + + # Test if job succeeded when not cancelled. + run_job(group_a, 0) + self.assertFalse(is_job_cancelled[0]) + + # Run jobs + for i in thread_ids_to_cancel: + t = threading.Thread(target=run_job, args=(group_a, i)) + t.start() + threads.append(t) + + for i in thread_ids_to_run: + t = threading.Thread(target=run_job, args=(group_b, i)) + t.start() + threads.append(t) + + # Wait to make sure all jobs are executed. + time.sleep(3) + # And then, cancel one job group. + self.sc.cancelJobGroup(group_a) + + # Wait until all threads launching jobs are finished. + for t in threads: + t.join() + + for i in thread_ids_to_cancel: + self.assertTrue( + is_job_cancelled[i], "Thread {i}: Job in group A was not cancelled.".format(i=i) + ) + + for i in thread_ids_to_run: + self.assertFalse( + is_job_cancelled[i], "Thread {i}: Job in group B did not succeeded.".format(i=i) + ) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_rdd import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddbarrier.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddbarrier.py new file mode 100644 index 0000000..dd3d2d6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddbarrier.py @@ -0,0 +1,52 @@ +# +# 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.testing.utils import ReusedPySparkTestCase + + +class RDDBarrierTests(ReusedPySparkTestCase): + def test_map_partitions(self): + """Test RDDBarrier.mapPartitions""" + rdd = self.sc.parallelize(range(12), 4) + self.assertFalse(rdd._is_barrier()) + + rdd1 = rdd.barrier().mapPartitions(lambda it: it) + self.assertTrue(rdd1._is_barrier()) + + def test_map_partitions_with_index(self): + """Test RDDBarrier.mapPartitionsWithIndex""" + rdd = self.sc.parallelize(range(12), 4) + self.assertFalse(rdd._is_barrier()) + + def f(index, iterator): + yield index + + rdd1 = rdd.barrier().mapPartitionsWithIndex(f) + self.assertTrue(rdd1._is_barrier()) + self.assertEqual(rdd1.collect(), [0, 1, 2, 3]) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_rddbarrier import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddsampler.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddsampler.py new file mode 100644 index 0000000..b98f266 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_rddsampler.py @@ -0,0 +1,66 @@ +# +# 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.testing.utils import ReusedPySparkTestCase +from pyspark.rddsampler import RDDSampler, RDDStratifiedSampler + + +class RDDSamplerTests(ReusedPySparkTestCase): + def test_rdd_sampler_func(self): + # SPARK-38879: Test case to improve test coverage for RDDSampler + # RDDSampler.func + rdd = self.sc.parallelize(range(20), 2) + sample_count = rdd.mapPartitionsWithIndex(RDDSampler(False, 0.4, 10).func).count() + self.assertGreater(sample_count, 3) + self.assertLess(sample_count, 10) + sample_data = rdd.mapPartitionsWithIndex(RDDSampler(True, 1, 10).func).collect() + sample_data.sort() + # check if at least one element is repeated. + self.assertTrue( + any(sample_data[i] == sample_data[i - 1] for i in range(1, len(sample_data))) + ) + + def test_rdd_stratified_sampler_func(self): + # SPARK-38879: Test case to improve test coverage for RDDSampler + # RDDStratifiedSampler.func + + fractions = {"a": 0.8, "b": 0.2} + rdd = self.sc.parallelize(fractions.keys()).cartesian(self.sc.parallelize(range(0, 100))) + sample_data = dict( + rdd.mapPartitionsWithIndex( + RDDStratifiedSampler(False, fractions, 10).func, True + ).countByKey() + ) + # Since a have higher sampling rate (0.8), + # it will occur more number of times than b. + self.assertGreater(sample_data["a"], sample_data["b"]) + self.assertGreater(sample_data["a"], 60) + self.assertLess(sample_data["a"], 90) + self.assertGreater(sample_data["b"], 15) + self.assertLess(sample_data["b"], 30) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_rddsampler import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_readwrite.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_readwrite.py new file mode 100644 index 0000000..73f1025 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_readwrite.py @@ -0,0 +1,368 @@ +# +# 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 os +import shutil +import tempfile +import unittest + +from pyspark.testing.utils import ReusedPySparkTestCase, SPARK_HOME + + +class InputFormatTests(ReusedPySparkTestCase): + @classmethod + def setUpClass(cls): + ReusedPySparkTestCase.setUpClass() + cls.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(cls.tempdir.name) + cls.sc._jvm.WriteInputFormatTestDataGenerator.generateData(cls.tempdir.name, cls.sc._jsc) + + @classmethod + def tearDownClass(cls): + ReusedPySparkTestCase.tearDownClass() + shutil.rmtree(cls.tempdir.name) + + def test_oldhadoop(self): + basepath = self.tempdir.name + ints = sorted( + self.sc.hadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapred.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ).collect() + ) + ei = [(1, "aa"), (1, "aa"), (2, "aa"), (2, "bb"), (2, "bb"), (3, "cc")] + self.assertEqual(ints, ei) + + hellopath = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + oldconf = {"mapreduce.input.fileinputformat.inputdir": hellopath} + hello = self.sc.hadoopRDD( + "org.apache.hadoop.mapred.TextInputFormat", + "org.apache.hadoop.io.LongWritable", + "org.apache.hadoop.io.Text", + conf=oldconf, + ).collect() + result = [(0, "Hello World!")] + self.assertEqual(hello, result) + + def test_newhadoop(self): + basepath = self.tempdir.name + ints = sorted( + self.sc.newAPIHadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ).collect() + ) + ei = [(1, "aa"), (1, "aa"), (2, "aa"), (2, "bb"), (2, "bb"), (3, "cc")] + self.assertEqual(ints, ei) + + hellopath = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt") + newconf = {"mapreduce.input.fileinputformat.inputdir": hellopath} + hello = self.sc.newAPIHadoopRDD( + "org.apache.hadoop.mapreduce.lib.input.TextInputFormat", + "org.apache.hadoop.io.LongWritable", + "org.apache.hadoop.io.Text", + conf=newconf, + ).collect() + result = [(0, "Hello World!")] + self.assertEqual(hello, result) + + def test_newolderror(self): + basepath = self.tempdir.name + self.assertRaises( + Exception, + lambda: self.sc.hadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ), + ) + + self.assertRaises( + Exception, + lambda: self.sc.newAPIHadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapred.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ), + ) + + def test_bad_inputs(self): + basepath = self.tempdir.name + self.assertRaises( + Exception, + lambda: self.sc.sequenceFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.io.NotValidWritable", + "org.apache.hadoop.io.Text", + ), + ) + self.assertRaises( + Exception, + lambda: self.sc.hadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapred.NotValidInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ), + ) + self.assertRaises( + Exception, + lambda: self.sc.newAPIHadoopFile( + basepath + "/sftestdata/sfint/", + "org.apache.hadoop.mapreduce.lib.input.NotValidInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ), + ) + + def test_converters(self): + # use of custom converters + basepath = self.tempdir.name + maps = sorted( + self.sc.sequenceFile( + basepath + "/sftestdata/sfmap/", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.MapWritable", + keyConverter="org.apache.spark.api.python.TestInputKeyConverter", + valueConverter="org.apache.spark.api.python.TestInputValueConverter", + ).collect() + ) + em = [("\x01", []), ("\x01", [3.0]), ("\x02", [1.0]), ("\x02", [1.0]), ("\x03", [2.0])] + self.assertEqual(maps, em) + + def test_binary_files(self): + path = os.path.join(self.tempdir.name, "binaryfiles") + os.mkdir(path) + data = b"short binary data" + with open(os.path.join(path, "part-0000"), "wb") as f: + f.write(data) + [(p, d)] = self.sc.binaryFiles(path).collect() + self.assertTrue(p.endswith("part-0000")) + self.assertEqual(d, data) + + def test_binary_records(self): + path = os.path.join(self.tempdir.name, "binaryrecords") + os.mkdir(path) + with open(os.path.join(path, "part-0000"), "w") as f: + for i in range(100): + f.write("%04d" % i) + result = self.sc.binaryRecords(path, 4).map(int).collect() + self.assertEqual(list(range(100)), result) + + +class OutputFormatTests(ReusedPySparkTestCase): + def setUp(self): + self.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(self.tempdir.name) + + def tearDown(self): + shutil.rmtree(self.tempdir.name, ignore_errors=True) + + def test_oldhadoop(self): + basepath = self.tempdir.name + dict_data = [(1, {}), (1, {"row1": 1.0}), (2, {"row2": 2.0})] + self.sc.parallelize(dict_data).saveAsHadoopFile( + basepath + "/oldhadoop/", + "org.apache.hadoop.mapred.SequenceFileOutputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.MapWritable", + ) + result = self.sc.hadoopFile( + basepath + "/oldhadoop/", + "org.apache.hadoop.mapred.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.MapWritable", + ).collect() + for v in result: + self.assertTrue(v, dict_data) + + conf = { + "mapred.output.format.class": "org.apache.hadoop.mapred.SequenceFileOutputFormat", + "mapreduce.job.output.key.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.job.output.value.class": "org.apache.hadoop.io.MapWritable", + "mapreduce.output.fileoutputformat.outputdir": basepath + "/olddataset/", + } + self.sc.parallelize(dict_data).saveAsHadoopDataset(conf) + input_conf = {"mapreduce.input.fileinputformat.inputdir": basepath + "/olddataset/"} + result = self.sc.hadoopRDD( + "org.apache.hadoop.mapred.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.MapWritable", + conf=input_conf, + ).collect() + for v in result: + self.assertTrue(v, dict_data) + + def test_newhadoop(self): + basepath = self.tempdir.name + data = [(1, ""), (1, "a"), (2, "bcdf")] + self.sc.parallelize(data).saveAsNewAPIHadoopFile( + basepath + "/newhadoop/", + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ) + result = sorted( + self.sc.newAPIHadoopFile( + basepath + "/newhadoop/", + "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + ).collect() + ) + self.assertEqual(result, data) + + conf = { + "mapreduce.job.outputformat.class": ( + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + ), + "mapreduce.job.output.key.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.job.output.value.class": "org.apache.hadoop.io.Text", + "mapreduce.output.fileoutputformat.outputdir": basepath + "/newdataset/", + } + self.sc.parallelize(data).saveAsNewAPIHadoopDataset(conf) + input_conf = {"mapreduce.input.fileinputformat.inputdir": basepath + "/newdataset/"} + new_dataset = sorted( + self.sc.newAPIHadoopRDD( + "org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat", + "org.apache.hadoop.io.IntWritable", + "org.apache.hadoop.io.Text", + conf=input_conf, + ).collect() + ) + self.assertEqual(new_dataset, data) + + def test_newolderror(self): + basepath = self.tempdir.name + rdd = self.sc.parallelize(range(1, 4)).map(lambda x: (x, "a" * x)) + self.assertRaises( + Exception, + lambda: rdd.saveAsHadoopFile( + basepath + "/newolderror/saveAsHadoopFile/", + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat", + ), + ) + self.assertRaises( + Exception, + lambda: rdd.saveAsNewAPIHadoopFile( + basepath + "/newolderror/saveAsNewAPIHadoopFile/", + "org.apache.hadoop.mapred.SequenceFileOutputFormat", + ), + ) + + def test_bad_inputs(self): + basepath = self.tempdir.name + rdd = self.sc.parallelize(range(1, 4)).map(lambda x: (x, "a" * x)) + self.assertRaises( + Exception, + lambda: rdd.saveAsHadoopFile( + basepath + "/badinputs/saveAsHadoopFile/", + "org.apache.hadoop.mapred.NotValidOutputFormat", + ), + ) + self.assertRaises( + Exception, + lambda: rdd.saveAsNewAPIHadoopFile( + basepath + "/badinputs/saveAsNewAPIHadoopFile/", + "org.apache.hadoop.mapreduce.lib.output.NotValidOutputFormat", + ), + ) + + def test_converters(self): + # use of custom converters + basepath = self.tempdir.name + data = [(1, {3.0: "bb"}), (2, {1.0: "aa"}), (3, {2.0: "dd"})] + self.sc.parallelize(data).saveAsNewAPIHadoopFile( + basepath + "/converters/", + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat", + keyConverter="org.apache.spark.api.python.TestOutputKeyConverter", + valueConverter="org.apache.spark.api.python.TestOutputValueConverter", + ) + converted = sorted(self.sc.sequenceFile(basepath + "/converters/").collect()) + expected = [("1", 3.0), ("2", 1.0), ("3", 2.0)] + self.assertEqual(converted, expected) + + def test_reserialization(self): + basepath = self.tempdir.name + x = range(1, 5) + y = range(1001, 1005) + data = list(zip(x, y)) + rdd = self.sc.parallelize(x).zip(self.sc.parallelize(y)) + rdd.saveAsSequenceFile(basepath + "/reserialize/sequence") + result1 = sorted(self.sc.sequenceFile(basepath + "/reserialize/sequence").collect()) + self.assertEqual(result1, data) + + rdd.saveAsHadoopFile( + basepath + "/reserialize/hadoop", "org.apache.hadoop.mapred.SequenceFileOutputFormat" + ) + result2 = sorted(self.sc.sequenceFile(basepath + "/reserialize/hadoop").collect()) + self.assertEqual(result2, data) + + rdd.saveAsNewAPIHadoopFile( + basepath + "/reserialize/newhadoop", + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat", + ) + result3 = sorted(self.sc.sequenceFile(basepath + "/reserialize/newhadoop").collect()) + self.assertEqual(result3, data) + + conf4 = { + "mapred.output.format.class": "org.apache.hadoop.mapred.SequenceFileOutputFormat", + "mapreduce.job.output.key.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.job.output.value.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.output.fileoutputformat.outputdir": basepath + "/reserialize/dataset", + } + rdd.saveAsHadoopDataset(conf4) + result4 = sorted(self.sc.sequenceFile(basepath + "/reserialize/dataset").collect()) + self.assertEqual(result4, data) + + conf5 = { + "mapreduce.job.outputformat.class": ( + "org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat" + ), + "mapreduce.job.output.key.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.job.output.value.class": "org.apache.hadoop.io.IntWritable", + "mapreduce.output.fileoutputformat.outputdir": basepath + "/reserialize/newdataset", + } + rdd.saveAsNewAPIHadoopDataset(conf5) + result5 = sorted(self.sc.sequenceFile(basepath + "/reserialize/newdataset").collect()) + self.assertEqual(result5, data) + + def test_malformed_RDD(self): + basepath = self.tempdir.name + # non-batch-serialized RDD[[(K, V)]] should be rejected + data = [[(1, "a")], [(2, "aa")], [(3, "aaa")]] + rdd = self.sc.parallelize(data, len(data)) + self.assertRaises( + Exception, lambda: rdd.saveAsSequenceFile(basepath + "/malformed/sequence") + ) + + +if __name__ == "__main__": + from pyspark.tests.test_readwrite import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_serializers.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_serializers.py new file mode 100644 index 0000000..230723e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_serializers.py @@ -0,0 +1,257 @@ +# +# 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 math +import sys +import unittest + +from pyspark import serializers +from pyspark.serializers import ( + CloudPickleSerializer, + CompressedSerializer, + AutoBatchedSerializer, + BatchedSerializer, + AutoSerializer, + NoOpSerializer, + PairDeserializer, + FlattenedValuesSerializer, + CartesianDeserializer, + CPickleSerializer, + UTF8Deserializer, + MarshalSerializer, +) +from pyspark.testing.utils import ( + PySparkTestCase, + read_int, + write_int, + ByteArrayOutput, + have_numpy, + have_scipy, +) + + +class SerializationTestCase(unittest.TestCase): + def test_namedtuple(self): + from collections import namedtuple + from pyspark.cloudpickle import dumps, loads + + P = namedtuple("P", "x y") + p1 = P(1, 3) + p2 = loads(dumps(p1, 2)) + self.assertEqual(p1, p2) + + P2 = loads(dumps(P)) + p3 = P2(1, 3) + self.assertEqual(p1, p3) + + def test_itemgetter(self): + from operator import itemgetter + + ser = CloudPickleSerializer() + d = range(10) + getter = itemgetter(1) + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + + getter = itemgetter(0, 3) + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + + def test_function_module_name(self): + ser = CloudPickleSerializer() + + def func(x): + return x + + func2 = ser.loads(ser.dumps(func)) + self.assertEqual(func.__module__, func2.__module__) + + def test_attrgetter(self): + from operator import attrgetter + + ser = CloudPickleSerializer() + + class C: + def __getattr__(self, item): + return item + + d = C() + getter = attrgetter("a") + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + getter = attrgetter("a", "b") + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + + d.e = C() + getter = attrgetter("e.a") + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + getter = attrgetter("e.a", "e.b") + getter2 = ser.loads(ser.dumps(getter)) + self.assertEqual(getter(d), getter2(d)) + + # Regression test for SPARK-3415 + def test_pickling_file_handles(self): + # to be corrected with SPARK-11160 + try: + import xmlrunner # noqa: F401 + except ImportError: + ser = CloudPickleSerializer() + out1 = sys.stderr + out2 = ser.loads(ser.dumps(out1)) + self.assertEqual(out1, out2) + + def test_func_globals(self): + class Unpicklable: + def __reduce__(self): + raise RuntimeError("not picklable") + + global exit + exit = Unpicklable() + + ser = CloudPickleSerializer() + self.assertRaises(Exception, lambda: ser.dumps(exit)) + + def foo(): + sys.exit(0) + + self.assertTrue("exit" in foo.__code__.co_names) + ser.dumps(foo) + + def test_compressed_serializer(self): + ser = CompressedSerializer(CPickleSerializer()) + from io import BytesIO as StringIO + + io = StringIO() + ser.dump_stream(["abc", "123", range(5)], io) + io.seek(0) + self.assertEqual(["abc", "123", range(5)], list(ser.load_stream(io))) + ser.dump_stream(range(1000), io) + io.seek(0) + self.assertEqual(["abc", "123", range(5)] + list(range(1000)), list(ser.load_stream(io))) + io.close() + + def test_hash_serializer(self): + hash(NoOpSerializer()) + hash(UTF8Deserializer()) + hash(CPickleSerializer()) + hash(MarshalSerializer()) + hash(AutoSerializer()) + hash(BatchedSerializer(CPickleSerializer())) + hash(AutoBatchedSerializer(MarshalSerializer())) + hash(PairDeserializer(NoOpSerializer(), UTF8Deserializer())) + hash(CartesianDeserializer(NoOpSerializer(), UTF8Deserializer())) + hash(CompressedSerializer(CPickleSerializer())) + hash(FlattenedValuesSerializer(CPickleSerializer())) + + +@unittest.skipIf(not have_scipy, "SciPy not installed") +class SciPyTests(PySparkTestCase): + + """General PySpark tests that depend on scipy""" + + def test_serialize(self): + from scipy.special import gammaln + + x = range(1, 5) + expected = list(map(gammaln, x)) + observed = self.sc.parallelize(x).map(gammaln).collect() + self.assertEqual(expected, observed) + + +@unittest.skipIf(not have_numpy, "NumPy not installed") +class NumPyTests(PySparkTestCase): + + """General PySpark tests that depend on numpy""" + + def test_statcounter_array(self): + import numpy as np + + x = self.sc.parallelize([np.array([1.0, 1.0]), np.array([2.0, 2.0]), np.array([3.0, 3.0])]) + s = x.stats() + self.assertSequenceEqual([2.0, 2.0], s.mean().tolist()) + self.assertSequenceEqual([1.0, 1.0], s.min().tolist()) + self.assertSequenceEqual([3.0, 3.0], s.max().tolist()) + self.assertSequenceEqual([1.0, 1.0], s.sampleStdev().tolist()) + + stats_dict = s.asDict() + self.assertEqual(3, stats_dict["count"]) + self.assertSequenceEqual([2.0, 2.0], stats_dict["mean"].tolist()) + self.assertSequenceEqual([1.0, 1.0], stats_dict["min"].tolist()) + self.assertSequenceEqual([3.0, 3.0], stats_dict["max"].tolist()) + self.assertSequenceEqual([6.0, 6.0], stats_dict["sum"].tolist()) + self.assertSequenceEqual([1.0, 1.0], stats_dict["stdev"].tolist()) + self.assertSequenceEqual([1.0, 1.0], stats_dict["variance"].tolist()) + + stats_sample_dict = s.asDict(sample=True) + self.assertEqual(3, stats_dict["count"]) + self.assertSequenceEqual([2.0, 2.0], stats_sample_dict["mean"].tolist()) + self.assertSequenceEqual([1.0, 1.0], stats_sample_dict["min"].tolist()) + self.assertSequenceEqual([3.0, 3.0], stats_sample_dict["max"].tolist()) + self.assertSequenceEqual([6.0, 6.0], stats_sample_dict["sum"].tolist()) + self.assertSequenceEqual( + [0.816496580927726, 0.816496580927726], stats_sample_dict["stdev"].tolist() + ) + self.assertSequenceEqual( + [0.6666666666666666, 0.6666666666666666], stats_sample_dict["variance"].tolist() + ) + + +class SerializersTest(unittest.TestCase): + def test_chunked_stream(self): + original_bytes = bytearray(range(100)) + for data_length in [1, 10, 100]: + for buffer_length in [1, 2, 3, 5, 20, 99, 100, 101, 500]: + dest = ByteArrayOutput() + stream_out = serializers.ChunkedStream(dest, buffer_length) + stream_out.write(original_bytes[:data_length]) + stream_out.close() + num_chunks = int(math.ceil(float(data_length) / buffer_length)) + # length for each chunk, and a final -1 at the very end + exp_size = (num_chunks + 1) * 4 + data_length + self.assertEqual(len(dest.buffer), exp_size) + dest_pos = 0 + data_pos = 0 + for chunk_idx in range(num_chunks): + chunk_length = read_int(dest.buffer[dest_pos : (dest_pos + 4)]) + if chunk_idx == num_chunks - 1: + exp_length = data_length % buffer_length + if exp_length == 0: + exp_length = buffer_length + else: + exp_length = buffer_length + self.assertEqual(chunk_length, exp_length) + dest_pos += 4 + dest_chunk = dest.buffer[dest_pos : dest_pos + chunk_length] + orig_chunk = original_bytes[data_pos : data_pos + chunk_length] + self.assertEqual(dest_chunk, orig_chunk) + dest_pos += chunk_length + data_pos += chunk_length + # ends with a -1 + self.assertEqual(dest.buffer[-4:], write_int(-1)) + + +if __name__ == "__main__": + from pyspark.tests.test_serializers import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_shuffle.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_shuffle.py new file mode 100644 index 0000000..4fb7360 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_shuffle.py @@ -0,0 +1,267 @@ +# +# 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 random +import unittest +import tempfile +import os + +from py4j.protocol import Py4JJavaError + +from pyspark import shuffle, CPickleSerializer, SparkConf, SparkContext +from pyspark.shuffle import ( + Aggregator, + ExternalMerger, + ExternalSorter, + SimpleAggregator, + Merger, + ExternalGroupBy, +) + + +class MergerTests(unittest.TestCase): + def setUp(self): + self.N = 1 << 12 + self.lst = [i for i in range(self.N)] + self.data = list(zip(self.lst, self.lst)) + self.agg = Aggregator( + lambda x: [x], lambda x, y: x.append(y) or x, lambda x, y: x.extend(y) or x + ) + + def test_small_dataset(self): + m = ExternalMerger(self.agg, 1000) + m.mergeValues(self.data) + self.assertEqual(m.spills, 0) + self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N))) + + m = ExternalMerger(self.agg, 1000) + m.mergeCombiners(map(lambda x_y1: (x_y1[0], [x_y1[1]]), self.data)) + self.assertEqual(m.spills, 0) + self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N))) + + def test_medium_dataset(self): + m = ExternalMerger(self.agg, 20) + m.mergeValues(self.data) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N))) + + m = ExternalMerger(self.agg, 10) + m.mergeCombiners(map(lambda x_y2: (x_y2[0], [x_y2[1]]), self.data * 3)) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)) * 3) + + def test_shuffle_data_with_multiple_locations(self): + # SPARK-39179: Test shuffle of data with multiple location also check + # shuffle locations get randomized + + with tempfile.TemporaryDirectory() as tempdir1, tempfile.TemporaryDirectory() as tempdir2: + original = os.environ.get("SPARK_LOCAL_DIRS", None) + os.environ["SPARK_LOCAL_DIRS"] = tempdir1 + "," + tempdir2 + try: + index_of_tempdir1 = [False, False] + for idx in range(10): + m = ExternalMerger(self.agg, 20) + if m.localdirs[0].startswith(tempdir1): + index_of_tempdir1[0] = True + elif m.localdirs[1].startswith(tempdir1): + index_of_tempdir1[1] = True + m.mergeValues(self.data) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N))) + self.assertTrue( + index_of_tempdir1[0] and (index_of_tempdir1[0] == index_of_tempdir1[1]) + ) + finally: + if original is not None: + os.environ["SPARK_LOCAL_DIRS"] = original + else: + del os.environ["SPARK_LOCAL_DIRS"] + + def test_simple_aggregator_with_medium_dataset(self): + # SPARK-39179: Test Simple aggregator + agg = SimpleAggregator(lambda x, y: x + y) + m = ExternalMerger(agg, 20) + m.mergeValues(self.data) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(v for k, v in m.items()), sum(range(self.N))) + + def test_merger_not_implemented_error(self): + # SPARK-39179: Test Merger for error scenarios + agg = SimpleAggregator(lambda x, y: x + y) + + class DummyMerger(Merger): + def __init__(self, agg): + Merger.__init__(self, agg) + + dummy_merger = DummyMerger(agg) + with self.assertRaises(NotImplementedError): + dummy_merger.mergeValues(self.data) + with self.assertRaises(NotImplementedError): + dummy_merger.mergeCombiners(self.data) + with self.assertRaises(NotImplementedError): + dummy_merger.items() + + def test_huge_dataset(self): + m = ExternalMerger(self.agg, 5, partitions=3) + m.mergeCombiners(map(lambda k_v: (k_v[0], [str(k_v[1])]), self.data * 10)) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(len(v) for k, v in m.items()), self.N * 10) + m._cleanup() + + def test_group_by_key(self): + def gen_data(N, step): + for i in range(1, N + 1, step): + for j in range(i): + yield (i, [j]) + + def gen_gs(N, step=1): + return shuffle.GroupByKey(gen_data(N, step)) + + self.assertEqual(1, len(list(gen_gs(1)))) + self.assertEqual(2, len(list(gen_gs(2)))) + self.assertEqual(100, len(list(gen_gs(100)))) + self.assertEqual(list(range(1, 101)), [k for k, _ in gen_gs(100)]) + self.assertTrue(all(list(range(k)) == list(vs) for k, vs in gen_gs(100))) + + for k, vs in gen_gs(50002, 10000): + self.assertEqual(k, len(vs)) + self.assertEqual(list(range(k)), list(vs)) + + ser = CPickleSerializer() + lst = ser.loads(ser.dumps(list(gen_gs(50002, 30000)))) + for k, vs in lst: + self.assertEqual(k, len(vs)) + self.assertEqual(list(range(k)), list(vs)) + + def test_stopiteration_is_raised(self): + def stopit(*args, **kwargs): + raise StopIteration() + + def legit_create_combiner(x): + return [x] + + def legit_merge_value(x, y): + return x.append(y) or x + + def legit_merge_combiners(x, y): + return x.extend(y) or x + + data = [(x % 2, x) for x in range(100)] + + # wrong create combiner + m = ExternalMerger(Aggregator(stopit, legit_merge_value, legit_merge_combiners), 20) + with self.assertRaises((Py4JJavaError, RuntimeError)): + m.mergeValues(data) + + # wrong merge value + m = ExternalMerger(Aggregator(legit_create_combiner, stopit, legit_merge_combiners), 20) + with self.assertRaises((Py4JJavaError, RuntimeError)): + m.mergeValues(data) + + # wrong merge combiners + m = ExternalMerger(Aggregator(legit_create_combiner, legit_merge_value, stopit), 20) + with self.assertRaises((Py4JJavaError, RuntimeError)): + m.mergeCombiners(map(lambda x_y1: (x_y1[0], [x_y1[1]]), data)) + + +class ExternalGroupByTests(unittest.TestCase): + def setUp(self): + self.N = 1 << 20 + values = [i for i in range(self.N)] + keys = [i for i in range(2)] + import itertools + + self.data = [value for value in itertools.product(keys, values)] + self.agg = Aggregator( + lambda x: [x], lambda x, y: x.append(y) or x, lambda x, y: x.extend(y) or x + ) + + def test_medium_dataset(self): + # SPARK-39179: Test external group by for medium dataset + m = ExternalGroupBy(self.agg, 5, partitions=3) + m.mergeValues(self.data) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(sum(v) for k, v in m.items()), 2 * sum(range(self.N))) + + def test_dataset_with_keys_are_unsorted(self): + # SPARK-39179: Test external group when numbers of keys are greater than SORT KEY Limit. + m = ExternalGroupBy(self.agg, 5, partitions=3) + original = m.SORT_KEY_LIMIT + try: + m.SORT_KEY_LIMIT = 1 + m.mergeValues(self.data) + self.assertTrue(m.spills >= 1) + self.assertEqual(sum(sum(v) for k, v in m.items()), 2 * sum(range(self.N))) + finally: + m.SORT_KEY_LIMIT = original + + +class SorterTests(unittest.TestCase): + def test_in_memory_sort(self): + lst = list(range(1024)) + random.shuffle(lst) + sorter = ExternalSorter(1024) + self.assertEqual(sorted(lst), list(sorter.sorted(lst))) + self.assertEqual(sorted(lst, reverse=True), list(sorter.sorted(lst, reverse=True))) + self.assertEqual(sorted(lst, key=lambda x: -x), list(sorter.sorted(lst, key=lambda x: -x))) + self.assertEqual( + sorted(lst, key=lambda x: -x, reverse=True), + list(sorter.sorted(lst, key=lambda x: -x, reverse=True)), + ) + + def test_external_sort(self): + class CustomizedSorter(ExternalSorter): + def _next_limit(self): + return self.memory_limit + + lst = list(range(1024)) + random.shuffle(lst) + sorter = CustomizedSorter(1) + self.assertEqual(sorted(lst), list(sorter.sorted(lst))) + self.assertGreater(shuffle.DiskBytesSpilled, 0) + last = shuffle.DiskBytesSpilled + self.assertEqual(sorted(lst, reverse=True), list(sorter.sorted(lst, reverse=True))) + self.assertGreater(shuffle.DiskBytesSpilled, last) + last = shuffle.DiskBytesSpilled + self.assertEqual(sorted(lst, key=lambda x: -x), list(sorter.sorted(lst, key=lambda x: -x))) + self.assertGreater(shuffle.DiskBytesSpilled, last) + last = shuffle.DiskBytesSpilled + self.assertEqual( + sorted(lst, key=lambda x: -x, reverse=True), + list(sorter.sorted(lst, key=lambda x: -x, reverse=True)), + ) + self.assertGreater(shuffle.DiskBytesSpilled, last) + + def test_external_sort_in_rdd(self): + conf = SparkConf().set("spark.python.worker.memory", "1m") + sc = SparkContext(conf=conf) + lst = list(range(10240)) + random.shuffle(lst) + rdd = sc.parallelize(lst, 4) + self.assertEqual(sorted(lst), rdd.sortBy(lambda x: x).collect()) + sc.stop() + + +if __name__ == "__main__": + from pyspark.tests.test_shuffle import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_stage_sched.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_stage_sched.py new file mode 100644 index 0000000..56cc0a0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_stage_sched.py @@ -0,0 +1,153 @@ +# +# 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 os +import tempfile +import unittest +import time +import shutil +import json + +from pyspark import SparkConf, SparkContext +from pyspark.resource.profile import ResourceProfileBuilder +from pyspark.resource.requests import TaskResourceRequests +from pyspark.taskcontext import TaskContext + + +class StageSchedulingTest(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + if getattr(self, "sc", None) is not None: + self.sc.stop() + self.sc = None + + def _test_stage_scheduling( + self, + cpus_per_worker, + gpus_per_worker, + num_tasks, + resource_profile, + expected_max_concurrent_tasks, + ): + conf = SparkConf() + conf.setMaster(f"local-cluster[1,{cpus_per_worker},1024]").set( + "spark.task.maxFailures", "1" + ) + + if gpus_per_worker: + worker_res_config_file = os.path.join(self.temp_dir, "worker_res.json") + worker_res = [ + { + "id": { + "componentName": "spark.worker", + "resourceName": "gpu", + }, + "addresses": [str(i) for i in range(gpus_per_worker)], + } + ] + with open(worker_res_config_file, "w") as fp: + json.dump(worker_res, fp) + + conf.set("spark.worker.resource.gpu.amount", str(gpus_per_worker)) + conf.set("spark.worker.resourcesFile", worker_res_config_file) + conf.set("spark.executor.resource.gpu.amount", str(gpus_per_worker)) + + self.sc = SparkContext(conf=conf) + pids_output_dir = os.path.join(self.temp_dir, "pids") + os.mkdir(pids_output_dir) + + def mapper(_): + task_id = TaskContext.get().partitionId() + pid_file_path = os.path.join(pids_output_dir, str(task_id)) + with open(pid_file_path, mode="w"): + pass + time.sleep(0.1) + num_concurrent_tasks = len(os.listdir(pids_output_dir)) + time.sleep(1) + os.remove(pid_file_path) + return num_concurrent_tasks + + results = ( + self.sc.parallelize(range(num_tasks), num_tasks) + .withResources(resource_profile) + .map(mapper) + .collect() + ) + self.assertEqual(max(results), expected_max_concurrent_tasks) + + def test_stage_scheduling_3_cpu_per_task(self): + rp = ResourceProfileBuilder().require(TaskResourceRequests().cpus(3)).build + self._test_stage_scheduling( + cpus_per_worker=4, + gpus_per_worker=0, + num_tasks=2, + resource_profile=rp, + expected_max_concurrent_tasks=1, + ) + + def test_stage_scheduling_2_cpu_per_task(self): + rp = ResourceProfileBuilder().require(TaskResourceRequests().cpus(2)).build + self._test_stage_scheduling( + cpus_per_worker=4, + gpus_per_worker=0, + num_tasks=4, + resource_profile=rp, + expected_max_concurrent_tasks=2, + ) + + def test_stage_scheduling_2_cpus_2_gpus_per_task(self): + rp = ( + ResourceProfileBuilder() + .require(TaskResourceRequests().cpus(2).resource("gpu", 2)) + .build + ) + self._test_stage_scheduling( + cpus_per_worker=4, + gpus_per_worker=4, + num_tasks=4, + resource_profile=rp, + expected_max_concurrent_tasks=2, + ) + + def test_stage_scheduling_2_cpus_3_gpus_per_task(self): + rp = ( + ResourceProfileBuilder() + .require(TaskResourceRequests().cpus(2).resource("gpu", 3)) + .build + ) + self._test_stage_scheduling( + cpus_per_worker=4, + gpus_per_worker=4, + num_tasks=2, + resource_profile=rp, + expected_max_concurrent_tasks=1, + ) + + +if __name__ == "__main__": + from pyspark.tests.test_stage_sched import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_statcounter.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_statcounter.py new file mode 100644 index 0000000..747f42e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_statcounter.py @@ -0,0 +1,130 @@ +# +# 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.statcounter import StatCounter +from pyspark.testing.utils import ReusedPySparkTestCase +import math + + +class StatCounterTests(ReusedPySparkTestCase): + def test_base(self): + stats = self.sc.parallelize([1.0, 2.0, 3.0, 4.0]).stats() + self.assertEqual(stats.count(), 4) + self.assertEqual(stats.max(), 4.0) + self.assertEqual(stats.mean(), 2.5) + self.assertEqual(stats.min(), 1.0) + self.assertAlmostEqual(stats.stdev(), 1.118033988749895) + self.assertAlmostEqual(stats.sampleStdev(), 1.2909944487358056) + self.assertEqual(stats.sum(), 10.0) + self.assertAlmostEqual(stats.variance(), 1.25) + self.assertAlmostEqual(stats.sampleVariance(), 1.6666666666666667) + + def test_as_dict(self): + stats = self.sc.parallelize([1.0, 2.0, 3.0, 4.0]).stats().asDict() + self.assertEqual(stats["count"], 4) + self.assertEqual(stats["max"], 4.0) + self.assertEqual(stats["mean"], 2.5) + self.assertEqual(stats["min"], 1.0) + self.assertAlmostEqual(stats["stdev"], 1.2909944487358056) + self.assertEqual(stats["sum"], 10.0) + self.assertAlmostEqual(stats["variance"], 1.6666666666666667) + + stats = self.sc.parallelize([1.0, 2.0, 3.0, 4.0]).stats().asDict(sample=True) + self.assertEqual(stats["count"], 4) + self.assertEqual(stats["max"], 4.0) + self.assertEqual(stats["mean"], 2.5) + self.assertEqual(stats["min"], 1.0) + self.assertAlmostEqual(stats["stdev"], 1.118033988749895) + self.assertEqual(stats["sum"], 10.0) + self.assertAlmostEqual(stats["variance"], 1.25) + + def test_merge(self): + stats = StatCounter([1.0, 2.0, 3.0, 4.0]) + stats.merge(5.0) + self.assertEqual(stats.count(), 5) + self.assertEqual(stats.max(), 5.0) + self.assertEqual(stats.mean(), 3.0) + self.assertEqual(stats.min(), 1.0) + self.assertAlmostEqual(stats.stdev(), 1.414213562373095) + self.assertAlmostEqual(stats.sampleStdev(), 1.5811388300841898) + self.assertEqual(stats.sum(), 15.0) + self.assertAlmostEqual(stats.variance(), 2.0) + self.assertAlmostEqual(stats.sampleVariance(), 2.5) + + def test_merge_stats(self): + stats1 = StatCounter([1.0, 2.0, 3.0, 4.0]) + stats2 = StatCounter([1.0, 2.0, 3.0, 4.0]) + stats = stats1.mergeStats(stats2) + self.assertEqual(stats.count(), 8) + self.assertEqual(stats.max(), 4.0) + self.assertEqual(stats.mean(), 2.5) + self.assertEqual(stats.min(), 1.0) + self.assertAlmostEqual(stats.stdev(), 1.118033988749895) + self.assertAlmostEqual(stats.sampleStdev(), 1.1952286093343936) + self.assertEqual(stats.sum(), 20.0) + self.assertAlmostEqual(stats.variance(), 1.25) + self.assertAlmostEqual(stats.sampleVariance(), 1.4285714285714286) + execution_statements = [ + StatCounter([1.0, 2.0]).mergeStats(StatCounter(range(1, 301))), + StatCounter(range(1, 301)).mergeStats(StatCounter([1.0, 2.0])), + ] + for stats in execution_statements: + self.assertEqual(stats.count(), 302) + self.assertEqual(stats.max(), 300.0) + self.assertEqual(stats.min(), 1.0) + self.assertAlmostEqual(stats.mean(), 149.51324503311) + self.assertAlmostEqual(stats.variance(), 7596.302804701549) + self.assertAlmostEqual(stats.sampleVariance(), 7621.539691095905) + + def test_variance_when_size_zero(self): + # SPARK-38854: Test case to improve test coverage when + # StatCounter argument is empty list or None + arguments = [[], None] + + for arg in arguments: + stats = StatCounter(arg) + self.assertTrue(math.isnan(stats.variance())) + self.assertTrue(math.isnan(stats.sampleVariance())) + self.assertEqual(stats.count(), 0) + self.assertTrue(math.isinf(stats.max())) + self.assertTrue(math.isinf(stats.min())) + self.assertEqual(stats.mean(), 0.0) + + def test_merge_stats_with_self(self): + stats = StatCounter([1.0, 2.0, 3.0, 4.0]) + stats.mergeStats(stats) + self.assertEqual(stats.count(), 8) + self.assertEqual(stats.max(), 4.0) + self.assertEqual(stats.mean(), 2.5) + self.assertEqual(stats.min(), 1.0) + self.assertAlmostEqual(stats.stdev(), 1.118033988749895) + self.assertAlmostEqual(stats.sampleStdev(), 1.1952286093343936) + self.assertEqual(stats.sum(), 20.0) + self.assertAlmostEqual(stats.variance(), 1.25) + self.assertAlmostEqual(stats.sampleVariance(), 1.4285714285714286) + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_statcounter import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_taskcontext.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_taskcontext.py new file mode 100644 index 0000000..5d410aa --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_taskcontext.py @@ -0,0 +1,350 @@ +# +# 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 os +import random +import stat +import sys +import tempfile +import time +import unittest + +from pyspark import SparkConf, SparkContext, TaskContext, BarrierTaskContext +from pyspark.testing.utils import PySparkTestCase, SPARK_HOME, eventually + + +class TaskContextTests(PySparkTestCase): + def setUp(self): + self._old_sys_path = list(sys.path) + class_name = self.__class__.__name__ + # Allow retries even though they are normally disabled in local mode + self.sc = SparkContext("local[4, 2]", class_name) + + def test_stage_id(self): + """Test the stage ids are available and incrementing as expected.""" + rdd = self.sc.parallelize(range(10)) + stage1 = rdd.map(lambda x: TaskContext.get().stageId()).take(1)[0] + stage2 = rdd.map(lambda x: TaskContext.get().stageId()).take(1)[0] + # Test using the constructor directly rather than the get() + stage3 = rdd.map(lambda x: TaskContext().stageId()).take(1)[0] + self.assertEqual(stage1 + 1, stage2) + self.assertEqual(stage1 + 2, stage3) + self.assertEqual(stage2 + 1, stage3) + + def test_resources(self): + """Test the resources are empty by default.""" + rdd = self.sc.parallelize(range(10)) + resources1 = rdd.map(lambda x: TaskContext.get().resources()).take(1)[0] + # Test using the constructor directly rather than the get() + resources2 = rdd.map(lambda x: TaskContext().resources()).take(1)[0] + self.assertEqual(len(resources1), 0) + self.assertEqual(len(resources2), 0) + + def test_partition_id(self): + """Test the partition id.""" + rdd1 = self.sc.parallelize(range(10), 1) + rdd2 = self.sc.parallelize(range(10), 2) + pids1 = rdd1.map(lambda x: TaskContext.get().partitionId()).collect() + pids2 = rdd2.map(lambda x: TaskContext.get().partitionId()).collect() + self.assertEqual(0, pids1[0]) + self.assertEqual(0, pids1[9]) + self.assertEqual(0, pids2[0]) + self.assertEqual(1, pids2[9]) + + def test_attempt_number(self): + """Verify the attempt numbers are correctly reported.""" + rdd = self.sc.parallelize(range(10)) + # Verify a simple job with no failures + attempt_numbers = rdd.map(lambda x: TaskContext.get().attemptNumber()).collect() + map(lambda attempt: self.assertEqual(0, attempt), attempt_numbers) + + def fail_on_first(x): + """Fail on the first attempt so we get a positive attempt number""" + tc = TaskContext.get() + attempt_number = tc.attemptNumber() + partition_id = tc.partitionId() + attempt_id = tc.taskAttemptId() + if attempt_number == 0 and partition_id == 0: + raise RuntimeError("Failing on first attempt") + else: + return [x, partition_id, attempt_number, attempt_id] + + result = rdd.map(fail_on_first).collect() + # We should re-submit the first partition to it but other partitions should be attempt 0 + self.assertEqual([0, 0, 1], result[0][0:3]) + self.assertEqual([9, 3, 0], result[9][0:3]) + first_partition = filter(lambda x: x[1] == 0, result) + map(lambda x: self.assertEqual(1, x[2]), first_partition) + other_partitions = filter(lambda x: x[1] != 0, result) + map(lambda x: self.assertEqual(0, x[2]), other_partitions) + # The task attempt id should be different + self.assertTrue(result[0][3] != result[9][3]) + + def test_tc_on_driver(self): + """Verify that getting the TaskContext on the driver returns None.""" + tc = TaskContext.get() + self.assertTrue(tc is None) + + def test_get_local_property(self): + """Verify that local properties set on the driver are available in TaskContext.""" + key = "testkey" + value = "testvalue" + self.sc.setLocalProperty(key, value) + try: + rdd = self.sc.parallelize(range(1), 1) + prop1 = rdd.map(lambda _: TaskContext.get().getLocalProperty(key)).collect()[0] + self.assertEqual(prop1, value) + prop2 = rdd.map(lambda _: TaskContext.get().getLocalProperty("otherkey")).collect()[0] + self.assertTrue(prop2 is None) + finally: + self.sc.setLocalProperty(key, None) + + def test_barrier(self): + """ + Verify that BarrierTaskContext.barrier() performs global sync among all barrier tasks + within a stage. + """ + rdd = self.sc.parallelize(range(10), 4) + + def f(iterator): + yield sum(iterator) + + def context_barrier(x): + tc = BarrierTaskContext.get() + time.sleep(random.randint(1, 5) * 2) + tc.barrier() + return time.time() + + times = rdd.barrier().mapPartitions(f).map(context_barrier).collect() + self.assertTrue(max(times) - min(times) < 2) + + def test_all_gather(self): + """ + Verify that BarrierTaskContext.allGather() performs global sync among all barrier tasks + within a stage and passes messages properly. + """ + rdd = self.sc.parallelize(range(10), 4) + + def f(iterator): + yield sum(iterator) + + def context_barrier(x): + tc = BarrierTaskContext.get() + time.sleep(random.randint(1, 10)) + out = tc.allGather(str(tc.partitionId())) + pids = [int(e) for e in out] + return pids + + pids = rdd.barrier().mapPartitions(f).map(context_barrier).collect()[0] + self.assertEqual(pids, [0, 1, 2, 3]) + + def test_barrier_infos(self): + """ + Verify that BarrierTaskContext.getTaskInfos() returns a list of all task infos in the + barrier stage. + """ + rdd = self.sc.parallelize(range(10), 4) + + def f(iterator): + yield sum(iterator) + + taskInfos = ( + rdd.barrier() + .mapPartitions(f) + .map(lambda x: BarrierTaskContext.get().getTaskInfos()) + .collect() + ) + self.assertTrue(len(taskInfos) == 4) + self.assertTrue(len(taskInfos[0]) == 4) + + def test_context_get(self): + """ + Verify that TaskContext.get() works both in or not in a barrier stage. + """ + rdd = self.sc.parallelize(range(10), 4) + + def f(iterator): + taskContext = TaskContext.get() + if isinstance(taskContext, BarrierTaskContext): + yield taskContext.partitionId() + 1 + elif isinstance(taskContext, TaskContext): + yield taskContext.partitionId() + 2 + else: + yield -1 + + # for normal stage + result1 = rdd.mapPartitions(f).collect() + self.assertTrue(result1 == [2, 3, 4, 5]) + # for barrier stage + result2 = rdd.barrier().mapPartitions(f).collect() + self.assertTrue(result2 == [1, 2, 3, 4]) + + def test_barrier_context_get(self): + """ + Verify that BarrierTaskContext.get() should only works in a barrier stage. + """ + rdd = self.sc.parallelize(range(10), 4) + + def f(iterator): + try: + taskContext = BarrierTaskContext.get() + except Exception: + yield -1 + else: + yield taskContext.partitionId() + + # for normal stage + result1 = rdd.mapPartitions(f).collect() + self.assertTrue(result1 == [-1, -1, -1, -1]) + # for barrier stage + result2 = rdd.barrier().mapPartitions(f).collect() + self.assertTrue(result2 == [0, 1, 2, 3]) + + +class TaskContextTestsWithWorkerReuse(unittest.TestCase): + def setUp(self): + class_name = self.__class__.__name__ + conf = SparkConf().set("spark.python.worker.reuse", "true") + self.sc = SparkContext("local[2]", class_name, conf=conf) + + def test_barrier_with_python_worker_reuse(self): + """ + Regression test for SPARK-25921: verify that BarrierTaskContext.barrier() with + reused python worker. + """ + # start a normal job first to start all workers and get all worker pids + worker_pids = self.sc.parallelize(range(2), 2).map(lambda x: os.getpid()).collect() + # the worker will reuse in this barrier job + rdd = self.sc.parallelize(range(10), 2) + + def f(iterator): + yield sum(iterator) + + def context_barrier(x): + tc = BarrierTaskContext.get() + time.sleep(random.randint(1, 5) * 2) + tc.barrier() + return (time.time(), os.getpid()) + + result = rdd.barrier().mapPartitions(f).map(context_barrier).collect() + times = list(map(lambda x: x[0], result)) + pids = list(map(lambda x: x[1], result)) + # check both barrier and worker reuse effect + self.assertTrue(max(times) - min(times) < 2) + for pid in pids: + self.assertTrue(pid in worker_pids) + + def check_task_context_correct_with_python_worker_reuse(self): + """Verify the task context correct when reused python worker""" + # start a normal job first to start all workers and get all worker pids + worker_pids = self.sc.parallelize(range(2), 2).map(lambda x: os.getpid()).collect() + # the worker will reuse in this barrier job + rdd = self.sc.parallelize(range(10), 2) + + def context(iterator): + tp = TaskContext.get().partitionId() + try: + bp = BarrierTaskContext.get().partitionId() + except Exception: + bp = -1 + + yield (tp, bp, os.getpid()) + + # normal stage after normal stage + normal_result = rdd.mapPartitions(context).collect() + tps, bps, pids = zip(*normal_result) + self.assertTrue(tps == (0, 1)) + self.assertTrue(bps == (-1, -1)) + for pid in pids: + self.assertTrue(pid in worker_pids) + # barrier stage after normal stage + barrier_result = rdd.barrier().mapPartitions(context).collect() + tps, bps, pids = zip(*barrier_result) + self.assertTrue(tps == (0, 1)) + self.assertTrue(bps == (0, 1)) + for pid in pids: + self.assertTrue(pid in worker_pids) + # normal stage after barrier stage + normal_result2 = rdd.mapPartitions(context).collect() + tps, bps, pids = zip(*normal_result2) + self.assertTrue(tps == (0, 1)) + self.assertTrue(bps == (-1, -1)) + for pid in pids: + self.assertTrue(pid in worker_pids) + return True + + def test_task_context_correct_with_python_worker_reuse(self): + # Retrying the check as the PIDs from Python workers might be different even + # when reusing Python workers is enabled if a Python worker is dead for some reasons + # (e.g., socket connection failure) and new Python worker is created. + eventually(self.check_task_context_correct_with_python_worker_reuse, catch_assertions=True) + + def tearDown(self): + self.sc.stop() + + +class TaskContextTestsWithResources(unittest.TestCase): + def setUp(self): + class_name = self.__class__.__name__ + self.tempFile = tempfile.NamedTemporaryFile(delete=False) + self.tempFile.write(b'echo {\\"name\\": \\"gpu\\", \\"addresses\\": [\\"0\\"]}') + self.tempFile.close() + # create temporary directory for Worker resources coordination + self.tempdir = tempfile.NamedTemporaryFile(delete=False) + os.unlink(self.tempdir.name) + os.chmod( + self.tempFile.name, + stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH, + ) + conf = SparkConf().set("spark.test.home", SPARK_HOME) + conf = conf.set("spark.worker.resource.gpu.discoveryScript", self.tempFile.name) + conf = conf.set("spark.worker.resource.gpu.amount", 1) + conf = conf.set("spark.task.cpus", 2) + conf = conf.set("spark.task.resource.gpu.amount", "1") + conf = conf.set("spark.executor.resource.gpu.amount", "1") + self.sc = SparkContext("local-cluster[2,2,1024]", class_name, conf=conf) + + def test_cpus(self): + """Test the cpus are available.""" + rdd = self.sc.parallelize(range(10)) + cpus = rdd.map(lambda x: TaskContext.get().cpus()).take(1)[0] + self.assertEqual(cpus, 2) + + def test_resources(self): + """Test the resources are available.""" + rdd = self.sc.parallelize(range(10)) + resources = rdd.map(lambda x: TaskContext.get().resources()).take(1)[0] + self.assertEqual(len(resources), 1) + self.assertTrue("gpu" in resources) + self.assertEqual(resources["gpu"].name, "gpu") + self.assertEqual(resources["gpu"].addresses, ["0"]) + + def tearDown(self): + os.unlink(self.tempFile.name) + self.sc.stop() + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_taskcontext import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_util.py new file mode 100644 index 0000000..77f0672 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_util.py @@ -0,0 +1,97 @@ +# +# 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 os +import unittest + +from py4j.protocol import Py4JJavaError + +from pyspark import keyword_only +from pyspark.testing.utils import PySparkTestCase +from pyspark.find_spark_home import _find_spark_home + + +class KeywordOnlyTests(unittest.TestCase): + class Wrapped: + @keyword_only + def set(self, x=None, y=None): + if "x" in self._input_kwargs: + self._x = self._input_kwargs["x"] + if "y" in self._input_kwargs: + self._y = self._input_kwargs["y"] + return x, y + + def test_keywords(self): + w = self.Wrapped() + x, y = w.set(y=1) + self.assertEqual(y, 1) + self.assertEqual(y, w._y) + self.assertIsNone(x) + self.assertFalse(hasattr(w, "_x")) + + def test_non_keywords(self): + w = self.Wrapped() + self.assertRaises(TypeError, lambda: w.set(0, y=1)) + + def test_kwarg_ownership(self): + # test _input_kwargs is owned by each class instance and not a shared static variable + class Setter: + @keyword_only + def set(self, x=None, other=None, other_x=None): + if "other" in self._input_kwargs: + self._input_kwargs["other"].set(x=self._input_kwargs["other_x"]) + self._x = self._input_kwargs["x"] + + a = Setter() + b = Setter() + a.set(x=1, other=b, other_x=2) + self.assertEqual(a._x, 1) + self.assertEqual(b._x, 2) + + +class UtilTests(PySparkTestCase): + def test_py4j_str(self): + with self.assertRaises(Py4JJavaError) as context: + # This attempts java.lang.String(null) which throws an NPE. + self.sc._jvm.java.lang.String(None) + + self.assertTrue("NullPointerException" in str(context.exception)) + + def test_parsing_version_string(self): + from pyspark.util import VersionUtils + + self.assertRaises(ValueError, lambda: VersionUtils.majorMinorVersion("abced")) + + def test_find_spark_home(self): + # SPARK-38827: Test find_spark_home without `SPARK_HOME` environment variable set. + origin = os.environ["SPARK_HOME"] + try: + del os.environ["SPARK_HOME"] + self.assertEquals(origin, _find_spark_home()) + finally: + os.environ["SPARK_HOME"] = origin + + +if __name__ == "__main__": + from pyspark.tests.test_util import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_worker.py b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_worker.py new file mode 100644 index 0000000..703690b --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/test_worker.py @@ -0,0 +1,271 @@ +# -*- encoding: utf-8 -*- +# +# 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 os +import sys +import tempfile +import threading +import time +import unittest + +has_resource_module = True +try: + import resource +except ImportError: + has_resource_module = False + +from py4j.protocol import Py4JJavaError + +from pyspark import SparkConf, SparkContext +from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest, eventually + + +class WorkerTests(ReusedPySparkTestCase): + def test_cancel_task(self): + temp = tempfile.NamedTemporaryFile(delete=True) + temp.close() + path = temp.name + + def sleep(x): + import os + import time + + with open(path, "w") as f: + f.write("%d %d" % (os.getppid(), os.getpid())) + time.sleep(100) + + # start job in background thread + def run(): + try: + self.sc.parallelize(range(1), 1).foreach(sleep) + except Exception: + pass + + import threading + + t = threading.Thread(target=run) + t.daemon = True + t.start() + + daemon_pid, worker_pid = 0, 0 + cnt = 0 + while True: + if os.path.exists(path): + with open(path) as f: + data = f.read().split(" ") + try: + daemon_pid, worker_pid = map(int, data) + except ValueError: + # In case the value is not written yet. + cnt += 1 + if cnt == 10: + raise + else: + break + time.sleep(1) + + # cancel jobs + self.sc.cancelAllJobs() + t.join() + + for i in range(50): + try: + os.kill(worker_pid, 0) + time.sleep(0.1) + except OSError: + break # worker was killed + else: + self.fail("worker has not been killed after 5 seconds") + + try: + os.kill(daemon_pid, 0) + except OSError: + self.fail("daemon had been killed") + + # run a normal job + rdd = self.sc.parallelize(range(100), 1) + self.assertEqual(100, rdd.map(str).count()) + + def test_after_exception(self): + def raise_exception(_): + raise RuntimeError() + + rdd = self.sc.parallelize(range(100), 1) + with QuietTest(self.sc): + self.assertRaises(Py4JJavaError, lambda: rdd.foreach(raise_exception)) + self.assertEqual(100, rdd.map(str).count()) + + def test_after_non_exception_error(self): + # SPARK-33339: Pyspark application will hang due to non Exception + def raise_system_exit(_): + raise SystemExit() + + rdd = self.sc.parallelize(range(100), 1) + with QuietTest(self.sc): + self.assertRaises(Exception, lambda: rdd.foreach(raise_system_exit)) + self.assertEqual(100, rdd.map(str).count()) + + def test_after_jvm_exception(self): + tempFile = tempfile.NamedTemporaryFile(delete=False) + tempFile.write(b"Hello World!") + tempFile.close() + data = self.sc.textFile(tempFile.name, 1) + filtered_data = data.filter(lambda x: True) + self.assertEqual(1, filtered_data.count()) + os.unlink(tempFile.name) + with QuietTest(self.sc): + self.assertRaises(Exception, lambda: filtered_data.count()) + + rdd = self.sc.parallelize(range(100), 1) + self.assertEqual(100, rdd.map(str).count()) + + def test_accumulator_when_reuse_worker(self): + from pyspark.accumulators import INT_ACCUMULATOR_PARAM + + acc1 = self.sc.accumulator(0, INT_ACCUMULATOR_PARAM) + self.sc.parallelize(range(100), 20).foreach(lambda x: acc1.add(x)) + self.assertEqual(sum(range(100)), acc1.value) + + acc2 = self.sc.accumulator(0, INT_ACCUMULATOR_PARAM) + self.sc.parallelize(range(100), 20).foreach(lambda x: acc2.add(x)) + self.assertEqual(sum(range(100)), acc2.value) + self.assertEqual(sum(range(100)), acc1.value) + + def test_reuse_worker_after_take(self): + rdd = self.sc.parallelize(range(100000), 1) + self.assertEqual(0, rdd.first()) + + def count(): + try: + rdd.count() + except Exception: + pass + + t = threading.Thread(target=count) + t.daemon = True + t.start() + t.join(5) + self.assertTrue(not t.is_alive()) + self.assertEqual(100000, rdd.count()) + + def test_with_different_versions_of_python(self): + rdd = self.sc.parallelize(range(10)) + rdd.count() + version = self.sc.pythonVer + self.sc.pythonVer = "2.0" + try: + with QuietTest(self.sc): + self.assertRaises(Py4JJavaError, lambda: rdd.count()) + finally: + self.sc.pythonVer = version + + def test_python_exception_non_hanging(self): + # SPARK-21045: exceptions with no ascii encoding shall not hanging PySpark. + try: + + def f(): + raise RuntimeError("exception with 中 and \xd6\xd0") + + self.sc.parallelize([1]).map(lambda x: f()).count() + except Py4JJavaError as e: + self.assertRegex(str(e), "exception with 中") + + +class WorkerReuseTest(PySparkTestCase): + def test_reuse_worker_of_parallelize_range(self): + def check_reuse_worker_of_parallelize_range(): + rdd = self.sc.parallelize(range(20), 8) + previous_pids = rdd.map(lambda x: os.getpid()).collect() + current_pids = rdd.map(lambda x: os.getpid()).collect() + for pid in current_pids: + self.assertTrue(pid in previous_pids) + return True + + eventually(check_reuse_worker_of_parallelize_range, catch_assertions=True) + + +@unittest.skipIf( + not has_resource_module or sys.platform != "linux", + "Memory limit feature in Python worker is dependent on " + "Python's 'resource' module on Linux; however, not found or not on Linux.", +) +class WorkerMemoryTest(unittest.TestCase): + def setUp(self): + class_name = self.__class__.__name__ + conf = SparkConf().set("spark.executor.pyspark.memory", "2g") + self.sc = SparkContext("local[4]", class_name, conf=conf) + + def test_memory_limit(self): + rdd = self.sc.parallelize(range(1), 1) + + def getrlimit(): + return resource.getrlimit(resource.RLIMIT_AS) + + actual = rdd.map(lambda _: getrlimit()).collect() + self.assertTrue(len(actual) == 1) + self.assertTrue(len(actual[0]) == 2) + [(soft_limit, hard_limit)] = actual + self.assertEqual(soft_limit, 2 * 1024 * 1024 * 1024) + self.assertEqual(hard_limit, 2 * 1024 * 1024 * 1024) + + def tearDown(self): + self.sc.stop() + + +class WorkerSegfaultTest(ReusedPySparkTestCase): + @classmethod + def conf(cls): + _conf = super(WorkerSegfaultTest, cls).conf() + _conf.set("spark.python.worker.faulthandler.enabled", "true") + return _conf + + def test_python_segfault(self): + try: + + def f(): + import ctypes + + ctypes.string_at(0) + + self.sc.parallelize([1]).map(lambda x: f()).count() + except Py4JJavaError as e: + self.assertRegex(str(e), "Segmentation fault") + + +@unittest.skipIf( + "COVERAGE_PROCESS_START" in os.environ, + "Flaky with coverage enabled, skipping for now.", +) +class WorkerSegfaultNonDaemonTest(WorkerSegfaultTest): + @classmethod + def conf(cls): + _conf = super(WorkerSegfaultNonDaemonTest, cls).conf() + _conf.set("spark.python.use.daemon", "false") + return _conf + + +if __name__ == "__main__": + import unittest + from pyspark.tests.test_worker import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_context.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_context.yml new file mode 100644 index 0000000..1217651 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_context.yml @@ -0,0 +1,21 @@ +# +# 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. +# + +- case: contextInitalization + main: | + from pyspark import SparkContext + sc: SparkContext = SparkContext() diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_core.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_core.yml new file mode 100644 index 0000000..ff58613 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_core.yml @@ -0,0 +1,20 @@ +# +# 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. +# + +- case: coreImports + main: | + from pyspark import keyword_only, Row, SQLContext diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_rdd.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_rdd.yml new file mode 100644 index 0000000..3585533 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_rdd.yml @@ -0,0 +1,129 @@ +# +# 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. +# + +- case: toDF + main: | + from pyspark.sql.types import ( + IntegerType, + Row, + StructType, + StringType, + StructField, + ) + from collections import namedtuple + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + sc = spark.sparkContext + + struct = StructType([ + StructField("a", IntegerType()), + StructField("b", StringType()) + ]) + + AB = namedtuple("AB", ["a", "b"]) + + rdd_row = sc.parallelize([Row(a=1, b="foo")]) + rdd_row.toDF() + rdd_row.toDF(sampleRatio=0.4) + rdd_row.toDF(["a", "b"], sampleRatio=0.4) + rdd_row.toDF(struct) + + rdd_tuple = sc.parallelize([(1, "foo")]) + rdd_tuple.toDF() + rdd_tuple.toDF(sampleRatio=0.4) + rdd_tuple.toDF(["a", "b"], sampleRatio=0.4) + rdd_tuple.toDF(struct) + + rdd_list = sc.parallelize([[1, "foo"]]) + rdd_list.toDF() + rdd_list.toDF(sampleRatio=0.4) + rdd_list.toDF(["a", "b"], sampleRatio=0.4) + rdd_list.toDF(struct) + + rdd_named_tuple = sc.parallelize([AB(1, "foo")]) + rdd_named_tuple.toDF() + rdd_named_tuple.toDF(sampleRatio=0.4) + rdd_named_tuple.toDF(["a", "b"], sampleRatio=0.4) + rdd_named_tuple.toDF(struct) + + +- case: rddMethods + main: | + from operator import add + from typing import Iterable, Set, Tuple + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + sc = spark.sparkContext + + def f1(x: int) -> str: + return str(x) + + reveal_type(sc.range(10).map(f1)) + + def f2(x: int) -> Iterable[int]: + return range(x) + + reveal_type(sc.range(10).flatMap(f2)) + + reveal_type(sc.parallelize([("a", 1), ("b", 0)]).filter(lambda x: x[1] != 0)) + + reveal_type(sc.parallelize([("a", 1), ("b", 0)]).max()) + + reveal_type(sc.range(10).reduce(add)) + + def seq_func(xs: Set[str], x: int) -> Set[str]: + xs.add(str(x % 11)) + return xs + + def comb_func(xs: Set[str], ys: Set[str]) -> Set[str]: + xs.update(ys) + return xs + + zero: Set[str] = set() + + reveal_type(sc.parallelize([("a", 1)]).aggregateByKey(zero, seq_func, comb_func)) + + out: | + main:11: note: Revealed type is "pyspark.rdd.RDD[builtins.str]" + main:16: note: Revealed type is "pyspark.rdd.RDD[builtins.int]" + main:18: note: Revealed type is "pyspark.rdd.RDD[Tuple[builtins.str, builtins.int]]" + main:20: note: Revealed type is "Tuple[builtins.str, builtins.int]" + main:22: note: Revealed type is "builtins.int" + main:34: note: Revealed type is "pyspark.rdd.RDD[Tuple[builtins.str, builtins.set[builtins.str]]]" + +- case: rddMethodsErrors + main: | + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + sc = spark.sparkContext + + def f1(x: str) -> str: + return x + + sc.range(10).map(f1) + + def f2(x: int) -> str: + return str(x) + + sc.range(10).reduce(f2) + + out: | + main:9: error: Argument 1 to "map" of "RDD" has incompatible type "Callable[[str], str]"; expected "Callable[[int], str]" [arg-type] + main:14: error: Argument 1 to "reduce" of "RDD" has incompatible type "Callable[[int], str]"; expected "Callable[[int, int], int]" [arg-type] diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_resultiterable.yml b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_resultiterable.yml new file mode 100644 index 0000000..4261a29 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/tests/typing/test_resultiterable.yml @@ -0,0 +1,22 @@ +# +# 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. +# + +- case: resultIterable + main: | + from pyspark.resultiterable import ResultIterable + + ResultIterable([]) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/traceback_utils.py b/spark-3.5.0-bin-hadoop3/python/pyspark/traceback_utils.py new file mode 100644 index 0000000..af4169e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/traceback_utils.py @@ -0,0 +1,81 @@ +# +# 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 collections import namedtuple +import os +import traceback + + +CallSite = namedtuple("CallSite", "function file linenum") + + +def first_spark_call(): + """ + Return a CallSite representing the first Spark call in the current call stack. + """ + tb = traceback.extract_stack() + if len(tb) == 0: + return None + file, line, module, what = tb[len(tb) - 1] + sparkpath = os.path.dirname(file) + first_spark_frame = len(tb) - 1 + for i in range(0, len(tb)): + file, line, fun, what = tb[i] + if file.startswith(sparkpath): + first_spark_frame = i + break + if first_spark_frame == 0: + file, line, fun, what = tb[0] + return CallSite(function=fun, file=file, linenum=line) + sfile, sline, sfun, swhat = tb[first_spark_frame] + ufile, uline, ufun, uwhat = tb[first_spark_frame - 1] + return CallSite(function=sfun, file=ufile, linenum=uline) + + +class SCCallSiteSync: + """ + Helper for setting the spark context call site. + + Example usage: + from pyspark.context import SCCallSiteSync + with SCCallSiteSync() as css: + + """ + + _spark_stack_depth = 0 + + def __init__(self, sc): + call_site = first_spark_call() + if call_site is not None: + self._call_site = "%s at %s:%s" % ( + call_site.function, + call_site.file, + call_site.linenum, + ) + else: + self._call_site = "Error! Could not extract traceback info" + self._context = sc + + def __enter__(self): + if SCCallSiteSync._spark_stack_depth == 0: + self._context._jsc.setCallSite(self._call_site) + SCCallSiteSync._spark_stack_depth += 1 + + def __exit__(self, type, value, tb): + SCCallSiteSync._spark_stack_depth -= 1 + if SCCallSiteSync._spark_stack_depth == 0: + self._context._jsc.setCallSite(None) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/util.py new file mode 100644 index 0000000..5232c92 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/util.py @@ -0,0 +1,416 @@ +# -*- coding: utf-8 -*- +# +# 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 functools +import itertools +import os +import platform +import re +import sys +import threading +import traceback +from types import TracebackType +from typing import Any, Callable, Iterator, List, Optional, TextIO, Tuple + +from pyspark.errors import PySparkRuntimeError + +from py4j.clientserver import ClientServer + +__all__: List[str] = [] + +from py4j.java_gateway import JavaObject + + +def print_exec(stream: TextIO) -> None: + ei = sys.exc_info() + traceback.print_exception(ei[0], ei[1], ei[2], None, stream) + + +class VersionUtils: + """ + Provides utility method to determine Spark versions with given input string. + """ + + @staticmethod + def majorMinorVersion(sparkVersion: str) -> Tuple[int, int]: + """ + Given a Spark version string, return the (major version number, minor version number). + E.g., for 2.0.1-SNAPSHOT, return (2, 0). + + Examples + -------- + >>> sparkVersion = "2.4.0" + >>> VersionUtils.majorMinorVersion(sparkVersion) + (2, 4) + >>> sparkVersion = "2.3.0-SNAPSHOT" + >>> VersionUtils.majorMinorVersion(sparkVersion) + (2, 3) + """ + m = re.search(r"^(\d+)\.(\d+)(\..*)?$", sparkVersion) + if m is not None: + return (int(m.group(1)), int(m.group(2))) + else: + raise ValueError( + "Spark tried to parse '%s' as a Spark" % sparkVersion + + " version string, but it could not find the major and minor" + + " version numbers." + ) + + +def fail_on_stopiteration(f: Callable) -> Callable: + """ + Wraps the input function to fail on 'StopIteration' by raising a 'RuntimeError' + prevents silent loss of data when 'f' is used in a for loop in Spark code + """ + + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return f(*args, **kwargs) + except StopIteration as exc: + raise PySparkRuntimeError( + error_class="STOP_ITERATION_OCCURRED", + message_parameters={ + "exc": str(exc), + }, + ) + + return wrapper + + +def walk_tb(tb: Optional[TracebackType]) -> Iterator[TracebackType]: + while tb is not None: + yield tb + tb = tb.tb_next + + +def try_simplify_traceback(tb: TracebackType) -> Optional[TracebackType]: + """ + Simplify the traceback. It removes the tracebacks in the current package, and only + shows the traceback that is related to the thirdparty and user-specified codes. + + Returns + ------- + TracebackType or None + Simplified traceback instance. It returns None if it fails to simplify. + + Notes + ----- + This keeps the tracebacks once it sees they are from a different file even + though the following tracebacks are from the current package. + + Examples + -------- + >>> import importlib + >>> import sys + >>> import traceback + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as tmp_dir: + ... with open("%s/dummy_module.py" % tmp_dir, "w") as f: + ... _ = f.write( + ... 'def raise_stop_iteration():\\n' + ... ' raise StopIteration()\\n\\n' + ... 'def simple_wrapper(f):\\n' + ... ' def wrapper(*a, **k):\\n' + ... ' return f(*a, **k)\\n' + ... ' return wrapper\\n') + ... f.flush() + ... spec = importlib.util.spec_from_file_location( + ... "dummy_module", "%s/dummy_module.py" % tmp_dir) + ... dummy_module = importlib.util.module_from_spec(spec) + ... spec.loader.exec_module(dummy_module) + >>> def skip_doctest_traceback(tb): + ... import pyspark + ... root = os.path.dirname(pyspark.__file__) + ... pairs = zip(walk_tb(tb), traceback.extract_tb(tb)) + ... for cur_tb, cur_frame in pairs: + ... if cur_frame.filename.startswith(root): + ... return cur_tb + + Regular exceptions should show the file name of the current package as below. + + >>> exc_info = None + >>> try: + ... fail_on_stopiteration(dummy_module.raise_stop_iteration)() + ... except Exception as e: + ... tb = sys.exc_info()[-1] + ... e.__cause__ = None + ... exc_info = "".join( + ... traceback.format_exception(type(e), e, tb)) + >>> print(exc_info) # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS + Traceback (most recent call last): + File ... + ... + File "/.../pyspark/util.py", line ... + ... + pyspark.errors.exceptions.base.PySparkRuntimeError: ... + >>> "pyspark/util.py" in exc_info + True + + If the traceback is simplified with this method, it hides the current package file name: + + >>> exc_info = None + >>> try: + ... fail_on_stopiteration(dummy_module.raise_stop_iteration)() + ... except Exception as e: + ... tb = try_simplify_traceback(sys.exc_info()[-1]) + ... e.__cause__ = None + ... exc_info = "".join( + ... traceback.format_exception( + ... type(e), e, try_simplify_traceback(skip_doctest_traceback(tb)))) + >>> print(exc_info) # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS + pyspark.errors.exceptions.base.PySparkRuntimeError: ... + >>> "pyspark/util.py" in exc_info + False + + In the case below, the traceback contains the current package in the middle. + In this case, it just hides the top occurrence only. + + >>> exc_info = None + >>> try: + ... fail_on_stopiteration(dummy_module.simple_wrapper( + ... fail_on_stopiteration(dummy_module.raise_stop_iteration)))() + ... except Exception as e: + ... tb = sys.exc_info()[-1] + ... e.__cause__ = None + ... exc_info_a = "".join( + ... traceback.format_exception(type(e), e, tb)) + ... exc_info_b = "".join( + ... traceback.format_exception( + ... type(e), e, try_simplify_traceback(skip_doctest_traceback(tb)))) + >>> exc_info_a.count("pyspark/util.py") + 2 + >>> exc_info_b.count("pyspark/util.py") + 1 + """ + if "pypy" in platform.python_implementation().lower(): + # Traceback modification is not supported with PyPy in PySpark. + return None + if sys.version_info[:2] < (3, 7): + # Traceback creation is not supported Python < 3.7. + # See https://bugs.python.org/issue30579. + return None + + import pyspark + + root = os.path.dirname(pyspark.__file__) + tb_next = None + new_tb = None + pairs = zip(walk_tb(tb), traceback.extract_tb(tb)) + last_seen = [] + + for cur_tb, cur_frame in pairs: + if not cur_frame.filename.startswith(root): + # Filter the stacktrace from the PySpark source itself. + last_seen = [(cur_tb, cur_frame)] + break + + for cur_tb, cur_frame in reversed(list(itertools.chain(last_seen, pairs))): + # Once we have seen the file names outside, don't skip. + new_tb = TracebackType( + tb_next=tb_next, + tb_frame=cur_tb.tb_frame, + tb_lasti=cur_tb.tb_frame.f_lasti, + tb_lineno=cur_tb.tb_frame.f_lineno if cur_tb.tb_frame.f_lineno is not None else -1, + ) + tb_next = new_tb + return new_tb + + +def _print_missing_jar(lib_name: str, pkg_name: str, jar_name: str, spark_version: str) -> None: + print( + """ +________________________________________________________________________________________________ + + Spark %(lib_name)s libraries not found in class path. Try one of the following. + + 1. Include the %(lib_name)s library and its dependencies with in the + spark-submit command as + + $ bin/spark-submit --packages org.apache.spark:spark-%(pkg_name)s:%(spark_version)s ... + + 2. Download the JAR of the artifact from Maven Central http://search.maven.org/, + Group Id = org.apache.spark, Artifact Id = spark-%(jar_name)s, Version = %(spark_version)s. + Then, include the jar in the spark-submit command as + + $ bin/spark-submit --jars ... + +________________________________________________________________________________________________ + +""" + % { + "lib_name": lib_name, + "pkg_name": pkg_name, + "jar_name": jar_name, + "spark_version": spark_version, + } + ) + + +def _parse_memory(s: str) -> int: + """ + Parse a memory string in the format supported by Java (e.g. 1g, 200m) and + return the value in MiB + + Examples + -------- + >>> _parse_memory("256m") + 256 + >>> _parse_memory("2g") + 2048 + """ + units = {"g": 1024, "m": 1, "t": 1 << 20, "k": 1.0 / 1024} + if s[-1].lower() not in units: + raise ValueError("invalid format: " + s) + return int(float(s[:-1]) * units[s[-1].lower()]) + + +def inheritable_thread_target(f: Callable) -> Callable: + """ + Return thread target wrapper which is recommended to be used in PySpark when the + pinned thread mode is enabled. The wrapper function, before calling original + thread target, it inherits the inheritable properties specific + to JVM thread such as ``InheritableThreadLocal``. + + Also, note that pinned thread mode does not close the connection from Python + to JVM when the thread is finished in the Python side. With this wrapper, Python + garbage-collects the Python thread instance and also closes the connection + which finishes JVM thread correctly. + + When the pinned thread mode is off, it return the original ``f``. + + .. versionadded:: 3.2.0 + + Parameters + ---------- + f : function + the original thread target. + + Notes + ----- + This API is experimental. + + It is important to know that it captures the local properties when you decorate it + whereas :class:`InheritableThread` captures when the thread is started. + Therefore, it is encouraged to decorate it when you want to capture the local + properties. + + For example, the local properties from the current Spark context is captured + when you define a function here instead of the invocation: + + >>> @inheritable_thread_target + ... def target_func(): + ... pass # your codes. + + If you have any updates on local properties afterwards, it would not be reflected to + the Spark context in ``target_func()``. + + The example below mimics the behavior of JVM threads as close as possible: + + >>> Thread(target=inheritable_thread_target(target_func)).start() # doctest: +SKIP + """ + from pyspark import SparkContext + + if isinstance(SparkContext._gateway, ClientServer): + # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. + + # NOTICE the internal difference vs `InheritableThread`. `InheritableThread` + # copies local properties when the thread starts but `inheritable_thread_target` + # copies when the function is wrapped. + assert SparkContext._active_spark_context is not None + properties = SparkContext._active_spark_context._jsc.sc().getLocalProperties().clone() + + @functools.wraps(f) + def wrapped(*args: Any, **kwargs: Any) -> Any: + # Set local properties in child thread. + assert SparkContext._active_spark_context is not None + SparkContext._active_spark_context._jsc.sc().setLocalProperties(properties) + return f(*args, **kwargs) + + return wrapped + else: + return f + + +class InheritableThread(threading.Thread): + """ + Thread that is recommended to be used in PySpark instead of :class:`threading.Thread` + when the pinned thread mode is enabled. The usage of this class is exactly same as + :class:`threading.Thread` but correctly inherits the inheritable properties specific + to JVM thread such as ``InheritableThreadLocal``. + + Also, note that pinned thread mode does not close the connection from Python + to JVM when the thread is finished in the Python side. With this class, Python + garbage-collects the Python thread instance and also closes the connection + which finishes JVM thread correctly. + + When the pinned thread mode is off, this works as :class:`threading.Thread`. + + .. versionadded:: 3.1.0 + + Notes + ----- + This API is experimental. + """ + + _props: JavaObject + + def __init__(self, target: Callable, *args: Any, **kwargs: Any): + from pyspark import SparkContext + + if isinstance(SparkContext._gateway, ClientServer): + # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. + def copy_local_properties(*a: Any, **k: Any) -> Any: + # self._props is set before starting the thread to match the behavior with JVM. + assert hasattr(self, "_props") + assert SparkContext._active_spark_context is not None + SparkContext._active_spark_context._jsc.sc().setLocalProperties(self._props) + return target(*a, **k) + + super(InheritableThread, self).__init__( + target=copy_local_properties, *args, **kwargs # type: ignore[misc] + ) + else: + super(InheritableThread, self).__init__( + target=target, *args, **kwargs # type: ignore[misc] + ) + + def start(self) -> None: + from pyspark import SparkContext + + if isinstance(SparkContext._gateway, ClientServer): + # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. + + # Local property copy should happen in Thread.start to mimic JVM's behavior. + assert SparkContext._active_spark_context is not None + self._props = SparkContext._active_spark_context._jsc.sc().getLocalProperties().clone() + return super(InheritableThread, self).start() + + +if __name__ == "__main__": + if "pypy" not in platform.python_implementation().lower() and sys.version_info[:2] >= (3, 7): + import doctest + import pyspark.util + from pyspark.context import SparkContext + + globs = pyspark.util.__dict__.copy() + globs["sc"] = SparkContext("local[4]", "PythonTest") + (failure_count, test_count) = doctest.testmod(pyspark.util, globs=globs) + globs["sc"].stop() + + if failure_count: + sys.exit(-1) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/version.py b/spark-3.5.0-bin-hadoop3/python/pyspark/version.py new file mode 100644 index 0000000..b1f7662 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/version.py @@ -0,0 +1 @@ +__version__: str = '3.5.0' diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/worker.py b/spark-3.5.0-bin-hadoop3/python/pyspark/worker.py new file mode 100644 index 0000000..90b11d0 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/worker.py @@ -0,0 +1,1307 @@ +# +# 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. +# + +""" +Worker that receives input from Piped RDD. +""" +import os +import sys +import time +from inspect import currentframe, getframeinfo, getfullargspec +import importlib +import json +from typing import Any, Iterable, Iterator + +# 'resource' is a Unix specific module. +has_resource_module = True +try: + import resource +except ImportError: + has_resource_module = False +import traceback +import warnings +import faulthandler + +from pyspark.accumulators import _accumulatorRegistry +from pyspark.broadcast import Broadcast, _broadcastRegistry +from pyspark.java_gateway import local_connect_and_auth +from pyspark.taskcontext import BarrierTaskContext, TaskContext +from pyspark.files import SparkFiles +from pyspark.resource import ResourceInformation +from pyspark.rdd import PythonEvalType +from pyspark.serializers import ( + write_with_length, + write_int, + read_long, + read_bool, + write_long, + read_int, + SpecialLengths, + UTF8Deserializer, + CPickleSerializer, + BatchedSerializer, +) +from pyspark.sql.pandas.serializers import ( + ArrowStreamPandasUDFSerializer, + ArrowStreamPandasUDTFSerializer, + CogroupUDFSerializer, + ArrowStreamUDFSerializer, + ApplyInPandasWithStateSerializer, +) +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.types import BinaryType, StringType, StructType, _parse_datatype_json_string +from pyspark.util import fail_on_stopiteration, try_simplify_traceback +from pyspark import shuffle +from pyspark.errors import PySparkRuntimeError, PySparkTypeError + +pickleSer = CPickleSerializer() +utf8_deserializer = UTF8Deserializer() + + +def report_times(outfile, boot, init, finish): + write_int(SpecialLengths.TIMING_DATA, outfile) + write_long(int(1000 * boot), outfile) + write_long(int(1000 * init), outfile) + write_long(int(1000 * finish), outfile) + + +def add_path(path): + # worker can be used, so do not add path multiple times + if path not in sys.path: + # overwrite system packages + sys.path.insert(1, path) + + +def read_command(serializer, file): + command = serializer._read_with_length(file) + if isinstance(command, Broadcast): + command = serializer.loads(command.value) + return command + + +def chain(f, g): + """chain two functions together""" + return lambda *a: g(f(*a)) + + +def wrap_udf(f, return_type): + if return_type.needConversion(): + toInternal = return_type.toInternal + return lambda *a: toInternal(f(*a)) + else: + return lambda *a: f(*a) + + +def wrap_scalar_pandas_udf(f, return_type): + arrow_return_type = to_arrow_type(return_type) + + def verify_result_type(result): + if not hasattr(result, "__len__"): + pd_type = "pandas.DataFrame" if type(return_type) == StructType else "pandas.Series" + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": pd_type, + "actual": type(result).__name__, + }, + ) + return result + + def verify_result_length(result, length): + if len(result) != length: + raise PySparkRuntimeError( + error_class="SCHEMA_MISMATCH_FOR_PANDAS_UDF", + message_parameters={ + "expected": str(length), + "actual": str(len(result)), + }, + ) + return result + + return lambda *a: ( + verify_result_length(verify_result_type(f(*a)), len(a[0])), + arrow_return_type, + ) + + +def wrap_arrow_batch_udf(f, return_type): + import pandas as pd + + arrow_return_type = to_arrow_type(return_type) + + # "result_func" ensures the result of a Python UDF to be consistent with/without Arrow + # optimization. + # Otherwise, an Arrow-optimized Python UDF raises "pyarrow.lib.ArrowTypeError: Expected a + # string or bytes dtype, got ..." whereas a non-Arrow-optimized Python UDF returns + # successfully. + result_func = lambda pdf: pdf # noqa: E731 + if type(return_type) == StringType: + result_func = lambda r: str(r) if r is not None else r # noqa: E731 + elif type(return_type) == BinaryType: + result_func = lambda r: bytes(r) if r is not None else r # noqa: E731 + + @fail_on_stopiteration + def evaluate(*args: pd.Series) -> pd.Series: + return pd.Series([result_func(f(*a)) for a in zip(*args)]) + + def verify_result_length(result, length): + if len(result) != length: + raise PySparkRuntimeError( + error_class="SCHEMA_MISMATCH_FOR_PANDAS_UDF", + message_parameters={ + "expected": str(length), + "actual": str(len(result)), + }, + ) + return result + + return lambda *a: ( + verify_result_length(evaluate(*a), len(a[0])), + arrow_return_type, + ) + + +def wrap_pandas_batch_iter_udf(f, return_type): + arrow_return_type = to_arrow_type(return_type) + iter_type_label = "pandas.DataFrame" if type(return_type) == StructType else "pandas.Series" + + def verify_result(result): + if not isinstance(result, Iterator) and not hasattr(result, "__iter__"): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterator of {}".format(iter_type_label), + "actual": type(result).__name__, + }, + ) + return result + + def verify_element(elem): + import pandas as pd + + if not isinstance(elem, pd.DataFrame if type(return_type) == StructType else pd.Series): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterator of {}".format(iter_type_label), + "actual": "iterator of {}".format(type(elem).__name__), + }, + ) + + verify_pandas_result( + elem, return_type, assign_cols_by_name=True, truncate_return_schema=True + ) + + return elem + + return lambda *iterator: map( + lambda res: (res, arrow_return_type), map(verify_element, verify_result(f(*iterator))) + ) + + +def verify_pandas_result(result, return_type, assign_cols_by_name, truncate_return_schema): + import pandas as pd + + if type(return_type) == StructType: + if not isinstance(result, pd.DataFrame): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "pandas.DataFrame", + "actual": type(result).__name__, + }, + ) + + # check the schema of the result only if it is not empty or has columns + if not result.empty or len(result.columns) != 0: + # if any column name of the result is a string + # the column names of the result have to match the return type + # see create_array in pyspark.sql.pandas.serializers.ArrowStreamPandasSerializer + field_names = set([field.name for field in return_type.fields]) + # only the first len(field_names) result columns are considered + # when truncating the return schema + result_columns = ( + result.columns[: len(field_names)] if truncate_return_schema else result.columns + ) + column_names = set(result_columns) + if ( + assign_cols_by_name + and any(isinstance(name, str) for name in result.columns) + and column_names != field_names + ): + missing = sorted(list(field_names.difference(column_names))) + missing = f" Missing: {', '.join(missing)}." if missing else "" + + extra = sorted(list(column_names.difference(field_names))) + extra = f" Unexpected: {', '.join(extra)}." if extra else "" + + raise PySparkRuntimeError( + error_class="RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF", + message_parameters={ + "missing": missing, + "extra": extra, + }, + ) + # otherwise the number of columns of result have to match the return type + elif len(result_columns) != len(return_type): + raise PySparkRuntimeError( + error_class="RESULT_LENGTH_MISMATCH_FOR_PANDAS_UDF", + message_parameters={ + "expected": str(len(return_type)), + "actual": str(len(result.columns)), + }, + ) + else: + if not isinstance(result, pd.Series): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={"expected": "pandas.Series", "actual": type(result).__name__}, + ) + + +def wrap_arrow_batch_iter_udf(f, return_type): + arrow_return_type = to_arrow_type(return_type) + + def verify_result(result): + if not isinstance(result, Iterator) and not hasattr(result, "__iter__"): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterator of pyarrow.RecordBatch", + "actual": type(result).__name__, + }, + ) + return result + + def verify_element(elem): + import pyarrow as pa + + if not isinstance(elem, pa.RecordBatch): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterator of pyarrow.RecordBatch", + "actual": "iterator of {}".format(type(elem).__name__), + }, + ) + + return elem + + return lambda *iterator: map( + lambda res: (res, arrow_return_type), map(verify_element, verify_result(f(*iterator))) + ) + + +def wrap_cogrouped_map_pandas_udf(f, return_type, argspec, runner_conf): + _assign_cols_by_name = assign_cols_by_name(runner_conf) + + def wrapped(left_key_series, left_value_series, right_key_series, right_value_series): + import pandas as pd + + left_df = pd.concat(left_value_series, axis=1) + right_df = pd.concat(right_value_series, axis=1) + + if len(argspec.args) == 2: + result = f(left_df, right_df) + elif len(argspec.args) == 3: + key_series = left_key_series if not left_df.empty else right_key_series + key = tuple(s[0] for s in key_series) + result = f(key, left_df, right_df) + verify_pandas_result( + result, return_type, _assign_cols_by_name, truncate_return_schema=False + ) + + return result + + return lambda kl, vl, kr, vr: [(wrapped(kl, vl, kr, vr), to_arrow_type(return_type))] + + +def wrap_grouped_map_pandas_udf(f, return_type, argspec, runner_conf): + _assign_cols_by_name = assign_cols_by_name(runner_conf) + + def wrapped(key_series, value_series): + import pandas as pd + + if len(argspec.args) == 1: + result = f(pd.concat(value_series, axis=1)) + elif len(argspec.args) == 2: + key = tuple(s[0] for s in key_series) + result = f(key, pd.concat(value_series, axis=1)) + verify_pandas_result( + result, return_type, _assign_cols_by_name, truncate_return_schema=False + ) + + return result + + return lambda k, v: [(wrapped(k, v), to_arrow_type(return_type))] + + +def wrap_grouped_map_pandas_udf_with_state(f, return_type): + """ + Provides a new lambda instance wrapping user function of applyInPandasWithState. + + The lambda instance receives (key series, iterator of value series, state) and performs + some conversion to be adapted with the signature of user function. + + See the function doc of inner function `wrapped` for more details on what adapter does. + See the function doc of `mapper` function for + `eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE` for more details on + the input parameters of lambda function. + + Along with the returned iterator, the lambda instance will also produce the return_type as + converted to the arrow schema. + """ + + def wrapped(key_series, value_series_gen, state): + """ + Provide an adapter of the user function performing below: + + - Extract the first value of all columns in key series and produce as a tuple. + - If the state has timed out, call the user function with empty pandas DataFrame. + - If not, construct a new generator which converts each element of value series to + pandas DataFrame (lazy evaluation), and call the user function with the generator + - Verify each element of returned iterator to check the schema of pandas DataFrame. + """ + import pandas as pd + + key = tuple(s[0] for s in key_series) + + if state.hasTimedOut: + # Timeout processing pass empty iterator. Here we return an empty DataFrame instead. + values = [ + pd.DataFrame(columns=pd.concat(next(value_series_gen), axis=1).columns), + ] + else: + values = (pd.concat(x, axis=1) for x in value_series_gen) + + result_iter = f(key, values, state) + + def verify_element(result): + if not isinstance(result, pd.DataFrame): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterator of pandas.DataFrame", + "actual": "iterator of {}".format(type(result).__name__), + }, + ) + # the number of columns of result have to match the return type + # but it is fine for result to have no columns at all if it is empty + if not ( + len(result.columns) == len(return_type) + or (len(result.columns) == 0 and result.empty) + ): + raise PySparkRuntimeError( + error_class="RESULT_LENGTH_MISMATCH_FOR_PANDAS_UDF", + message_parameters={ + "expected": str(len(return_type)), + "actual": str(len(result.columns)), + }, + ) + + return result + + if isinstance(result_iter, pd.DataFrame): + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={ + "expected": "iterable of pandas.DataFrame", + "actual": type(result_iter).__name__, + }, + ) + + try: + iter(result_iter) + except TypeError: + raise PySparkTypeError( + error_class="UDF_RETURN_TYPE", + message_parameters={"expected": "iterable", "actual": type(result_iter).__name__}, + ) + + result_iter_with_validation = (verify_element(x) for x in result_iter) + + return ( + result_iter_with_validation, + state, + ) + + return lambda k, v, s: [(wrapped(k, v, s), to_arrow_type(return_type))] + + +def wrap_grouped_agg_pandas_udf(f, return_type): + arrow_return_type = to_arrow_type(return_type) + + def wrapped(*series): + import pandas as pd + + result = f(*series) + return pd.Series([result]) + + return lambda *a: (wrapped(*a), arrow_return_type) + + +def wrap_window_agg_pandas_udf(f, return_type, runner_conf, udf_index): + window_bound_types_str = runner_conf.get("pandas_window_bound_types") + window_bound_type = [t.strip().lower() for t in window_bound_types_str.split(",")][udf_index] + if window_bound_type == "bounded": + return wrap_bounded_window_agg_pandas_udf(f, return_type) + elif window_bound_type == "unbounded": + return wrap_unbounded_window_agg_pandas_udf(f, return_type) + else: + raise PySparkRuntimeError( + error_class="INVALID_WINDOW_BOUND_TYPE", + message_parameters={ + "window_bound_type": window_bound_type, + }, + ) + + +def wrap_unbounded_window_agg_pandas_udf(f, return_type): + # This is similar to grouped_agg_pandas_udf, the only difference + # is that window_agg_pandas_udf needs to repeat the return value + # to match window length, where grouped_agg_pandas_udf just returns + # the scalar value. + arrow_return_type = to_arrow_type(return_type) + + def wrapped(*series): + import pandas as pd + + result = f(*series) + return pd.Series([result]).repeat(len(series[0])) + + return lambda *a: (wrapped(*a), arrow_return_type) + + +def wrap_bounded_window_agg_pandas_udf(f, return_type): + arrow_return_type = to_arrow_type(return_type) + + def wrapped(begin_index, end_index, *series): + import pandas as pd + + result = [] + + # Index operation is faster on np.ndarray, + # So we turn the index series into np array + # here for performance + begin_array = begin_index.values + end_array = end_index.values + + for i in range(len(begin_array)): + # Note: Create a slice from a series for each window is + # actually pretty expensive. However, there + # is no easy way to reduce cost here. + # Note: s.iloc[i : j] is about 30% faster than s[i: j], with + # the caveat that the created slices shares the same + # memory with s. Therefore, user are not allowed to + # change the value of input series inside the window + # function. It is rare that user needs to modify the + # input series in the window function, and therefore, + # it is be a reasonable restriction. + # Note: Calling reset_index on the slices will increase the cost + # of creating slices by about 100%. Therefore, for performance + # reasons we don't do it here. + series_slices = [s.iloc[begin_array[i] : end_array[i]] for s in series] + result.append(f(*series_slices)) + return pd.Series(result) + + return lambda *a: (wrapped(*a), arrow_return_type) + + +def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index): + num_arg = read_int(infile) + arg_offsets = [read_int(infile) for i in range(num_arg)] + chained_func = None + for i in range(read_int(infile)): + f, return_type = read_command(pickleSer, infile) + if chained_func is None: + chained_func = f + else: + chained_func = chain(chained_func, f) + + if eval_type in ( + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_ARROW_BATCHED_UDF, + ): + func = chained_func + else: + # make sure StopIteration's raised in the user code are not ignored + # when they are processed in a for loop, raise them as RuntimeError's instead + func = fail_on_stopiteration(chained_func) + + # the last returnType will be the return type of UDF + if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF: + return arg_offsets, wrap_scalar_pandas_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF: + return arg_offsets, wrap_arrow_batch_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF: + return arg_offsets, wrap_pandas_batch_iter_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF: + return arg_offsets, wrap_pandas_batch_iter_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF: + return arg_offsets, wrap_arrow_batch_iter_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF: + argspec = getfullargspec(chained_func) # signature was lost when wrapping it + return arg_offsets, wrap_grouped_map_pandas_udf(func, return_type, argspec, runner_conf) + elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: + return arg_offsets, wrap_grouped_map_pandas_udf_with_state(func, return_type) + elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: + argspec = getfullargspec(chained_func) # signature was lost when wrapping it + return arg_offsets, wrap_cogrouped_map_pandas_udf(func, return_type, argspec, runner_conf) + elif eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF: + return arg_offsets, wrap_grouped_agg_pandas_udf(func, return_type) + elif eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF: + return arg_offsets, wrap_window_agg_pandas_udf(func, return_type, runner_conf, udf_index) + elif eval_type == PythonEvalType.SQL_BATCHED_UDF: + return arg_offsets, wrap_udf(func, return_type) + else: + raise ValueError("Unknown eval type: {}".format(eval_type)) + + +# Used by SQL_GROUPED_MAP_PANDAS_UDF and SQL_SCALAR_PANDAS_UDF and SQL_ARROW_BATCHED_UDF when +# returning StructType +def assign_cols_by_name(runner_conf): + return ( + runner_conf.get( + "spark.sql.legacy.execution.pandas.groupedMap.assignColumnsByName", "true" + ).lower() + == "true" + ) + + +# Read and process a serialized user-defined table function (UDTF) from a socket. +# It expects the UDTF to be in a specific format and performs various checks to +# ensure the UDTF is valid. This function also prepares a mapper function for applying +# the UDTF logic to input rows. +def read_udtf(pickleSer, infile, eval_type): + if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF: + runner_conf = {} + # Load conf used for arrow evaluation. + num_conf = read_int(infile) + for i in range(num_conf): + k = utf8_deserializer.loads(infile) + v = utf8_deserializer.loads(infile) + runner_conf[k] = v + + # NOTE: if timezone is set here, that implies respectSessionTimeZone is True + timezone = runner_conf.get("spark.sql.session.timeZone", None) + safecheck = ( + runner_conf.get("spark.sql.execution.pandas.convertToArrowArraySafely", "false").lower() + == "true" + ) + ser = ArrowStreamPandasUDTFSerializer(timezone, safecheck) + else: + # Each row is a group so do not batch but send one by one. + ser = BatchedSerializer(CPickleSerializer(), 1) + + # See `PythonUDTFRunner.PythonUDFWriterThread.writeCommand' + num_arg = read_int(infile) + arg_offsets = [read_int(infile) for _ in range(num_arg)] + handler = read_command(pickleSer, infile) + if not isinstance(handler, type): + raise PySparkRuntimeError( + f"Invalid UDTF handler type. Expected a class (type 'type'), but " + f"got an instance of {type(handler).__name__}." + ) + + return_type = _parse_datatype_json_string(utf8_deserializer.loads(infile)) + if not type(return_type) == StructType: + raise PySparkRuntimeError( + f"The return type of a UDTF must be a struct type, but got {type(return_type)}." + ) + + # Instantiate the UDTF class. + try: + udtf = handler() + except Exception as e: + raise PySparkRuntimeError( + error_class="UDTF_EXEC_ERROR", + message_parameters={"method_name": "__init__", "error": str(e)}, + ) + + # Validate the UDTF + if not hasattr(udtf, "eval"): + raise PySparkRuntimeError( + "Failed to execute the user defined table function because it has not " + "implemented the 'eval' method. Please add the 'eval' method and try " + "the query again." + ) + + if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF: + + def wrap_arrow_udtf(f, return_type): + import pandas as pd + + arrow_return_type = to_arrow_type(return_type) + return_type_size = len(return_type) + + def verify_result(result): + if not isinstance(result, pd.DataFrame): + raise PySparkTypeError( + error_class="INVALID_ARROW_UDTF_RETURN_TYPE", + message_parameters={ + "type_name": type(result).__name__, + "value": str(result), + "func": f.__name__, + }, + ) + + # Validate the output schema when the result dataframe has either output + # rows or columns. Note that we avoid using `df.empty` here because the + # result dataframe may contain an empty row. For example, when a UDTF is + # defined as follows: def eval(self): yield tuple(). + if len(result) > 0 or len(result.columns) > 0: + if len(result.columns) != return_type_size: + raise PySparkRuntimeError( + error_class="UDTF_RETURN_SCHEMA_MISMATCH", + message_parameters={ + "expected": str(return_type_size), + "actual": str(len(result.columns)), + "func": f.__name__, + }, + ) + + # Verify the type and the schema of the result. + verify_pandas_result( + result, return_type, assign_cols_by_name=False, truncate_return_schema=False + ) + return result + + # Wrap the exception thrown from the UDTF in a PySparkRuntimeError. + def func(*args: Any) -> Any: + try: + return f(*args) + except Exception as e: + raise PySparkRuntimeError( + error_class="UDTF_EXEC_ERROR", + message_parameters={"method_name": f.__name__, "error": str(e)}, + ) + + def check_return_value(res): + # Check whether the result of an arrow UDTF is iterable before + # using it to construct a pandas DataFrame. + if res is not None and not isinstance(res, Iterable): + raise PySparkRuntimeError( + error_class="UDTF_RETURN_NOT_ITERABLE", + message_parameters={ + "type": type(res).__name__, + "func": f.__name__, + }, + ) + + def evaluate(*args: pd.Series): + if len(args) == 0: + res = func() + check_return_value(res) + yield verify_result(pd.DataFrame(res)), arrow_return_type + else: + # Create tuples from the input pandas Series, each tuple + # represents a row across all Series. + row_tuples = zip(*args) + for row in row_tuples: + res = func(*row) + check_return_value(res) + yield verify_result(pd.DataFrame(res)), arrow_return_type + + return evaluate + + eval = wrap_arrow_udtf(getattr(udtf, "eval"), return_type) + + if hasattr(udtf, "terminate"): + terminate = wrap_arrow_udtf(getattr(udtf, "terminate"), return_type) + else: + terminate = None + + def mapper(_, it): + try: + for a in it: + # The eval function yields an iterator. Each element produced by this + # iterator is a tuple in the form of (pandas.DataFrame, arrow_return_type). + yield from eval(*[a[o] for o in arg_offsets]) + finally: + if terminate is not None: + yield from terminate() + + return mapper, None, ser, ser + + else: + + def wrap_udtf(f, return_type): + assert return_type.needConversion() + toInternal = return_type.toInternal + return_type_size = len(return_type) + + def verify_and_convert_result(result): + if result is not None: + if hasattr(result, "__len__") and len(result) != return_type_size: + raise PySparkRuntimeError( + error_class="UDTF_RETURN_SCHEMA_MISMATCH", + message_parameters={ + "expected": str(return_type_size), + "actual": str(len(result)), + "func": f.__name__, + }, + ) + + if not (isinstance(result, (list, dict, tuple)) or hasattr(result, "__dict__")): + raise PySparkRuntimeError( + error_class="UDTF_INVALID_OUTPUT_ROW_TYPE", + message_parameters={ + "type": type(result).__name__, + "func": f.__name__, + }, + ) + + return toInternal(result) + + # Evaluate the function and return a tuple back to the executor. + def evaluate(*a) -> tuple: + try: + res = f(*a) + except Exception as e: + raise PySparkRuntimeError( + error_class="UDTF_EXEC_ERROR", + message_parameters={"method_name": f.__name__, "error": str(e)}, + ) + + if res is None: + # If the function returns None or does not have an explicit return statement, + # an empty tuple is returned to the executor. + # This is because directly constructing tuple(None) results in an exception. + return tuple() + + if not isinstance(res, Iterable): + raise PySparkRuntimeError( + error_class="UDTF_RETURN_NOT_ITERABLE", + message_parameters={ + "type": type(res).__name__, + "func": f.__name__, + }, + ) + + # If the function returns a result, we map it to the internal representation and + # returns the results as a tuple. + return tuple(map(verify_and_convert_result, res)) + + return evaluate + + eval = wrap_udtf(getattr(udtf, "eval"), return_type) + + if hasattr(udtf, "terminate"): + terminate = wrap_udtf(getattr(udtf, "terminate"), return_type) + else: + terminate = None + + # Return an iterator of iterators. + def mapper(_, it): + try: + for a in it: + yield eval(*[a[o] for o in arg_offsets]) + finally: + if terminate is not None: + yield terminate() + + return mapper, None, ser, ser + + +def read_udfs(pickleSer, infile, eval_type): + runner_conf = {} + + if eval_type in ( + PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_UDF, + PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, + PythonEvalType.SQL_MAP_ARROW_ITER_UDF, + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, + PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, + PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, + ): + + # Load conf used for pandas_udf evaluation + num_conf = read_int(infile) + for i in range(num_conf): + k = utf8_deserializer.loads(infile) + v = utf8_deserializer.loads(infile) + runner_conf[k] = v + + state_object_schema = None + if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: + state_object_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile))) + + # NOTE: if timezone is set here, that implies respectSessionTimeZone is True + timezone = runner_conf.get("spark.sql.session.timeZone", None) + safecheck = ( + runner_conf.get("spark.sql.execution.pandas.convertToArrowArraySafely", "false").lower() + == "true" + ) + + if eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: + ser = CogroupUDFSerializer(timezone, safecheck, assign_cols_by_name(runner_conf)) + elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: + arrow_max_records_per_batch = runner_conf.get( + "spark.sql.execution.arrow.maxRecordsPerBatch", 10000 + ) + arrow_max_records_per_batch = int(arrow_max_records_per_batch) + + ser = ApplyInPandasWithStateSerializer( + timezone, + safecheck, + assign_cols_by_name(runner_conf), + state_object_schema, + arrow_max_records_per_batch, + ) + elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF: + ser = ArrowStreamUDFSerializer() + else: + # Scalar Pandas UDF handles struct type arguments as pandas DataFrames instead of + # pandas Series. See SPARK-27240. + df_for_struct = ( + eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF + or eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + or eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF + ) + # Arrow-optimized Python UDF takes a struct type argument as a Row + struct_in_pandas = ( + "row" if eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF else "dict" + ) + ndarray_as_list = eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF + # Arrow-optimized Python UDF uses explicit Arrow cast for type coercion + arrow_cast = eval_type == PythonEvalType.SQL_ARROW_BATCHED_UDF + ser = ArrowStreamPandasUDFSerializer( + timezone, + safecheck, + assign_cols_by_name(runner_conf), + df_for_struct, + struct_in_pandas, + ndarray_as_list, + arrow_cast, + ) + else: + ser = BatchedSerializer(CPickleSerializer(), 100) + + num_udfs = read_int(infile) + + is_scalar_iter = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF + is_map_pandas_iter = eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF + is_map_arrow_iter = eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF + + if is_scalar_iter or is_map_pandas_iter or is_map_arrow_iter: + if is_scalar_iter: + assert num_udfs == 1, "One SCALAR_ITER UDF expected here." + if is_map_pandas_iter: + assert num_udfs == 1, "One MAP_PANDAS_ITER UDF expected here." + if is_map_arrow_iter: + assert num_udfs == 1, "One MAP_ARROW_ITER UDF expected here." + + arg_offsets, udf = read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=0) + + def func(_, iterator): + num_input_rows = 0 + + def map_batch(batch): + nonlocal num_input_rows + + udf_args = [batch[offset] for offset in arg_offsets] + num_input_rows += len(udf_args[0]) + if len(udf_args) == 1: + return udf_args[0] + else: + return tuple(udf_args) + + iterator = map(map_batch, iterator) + result_iter = udf(iterator) + + num_output_rows = 0 + for result_batch, result_type in result_iter: + num_output_rows += len(result_batch) + # This assert is for Scalar Iterator UDF to fail fast. + # The length of the entire input can only be explicitly known + # by consuming the input iterator in user side. Therefore, + # it's very unlikely the output length is higher than + # input length. + assert ( + is_map_pandas_iter or is_map_arrow_iter or num_output_rows <= num_input_rows + ), "Pandas SCALAR_ITER UDF outputted more rows than input rows." + yield (result_batch, result_type) + + if is_scalar_iter: + try: + next(iterator) + except StopIteration: + pass + else: + raise PySparkRuntimeError( + error_class="STOP_ITERATION_OCCURRED_FROM_SCALAR_ITER_PANDAS_UDF", + message_parameters={}, + ) + + if num_output_rows != num_input_rows: + raise PySparkRuntimeError( + error_class="RESULT_LENGTH_MISMATCH_FOR_SCALAR_ITER_PANDAS_UDF", + message_parameters={ + "output_length": str(num_output_rows), + "input_length": str(num_input_rows), + }, + ) + + # profiling is not supported for UDF + return func, None, ser, ser + + def extract_key_value_indexes(grouped_arg_offsets): + """ + Helper function to extract the key and value indexes from arg_offsets for the grouped and + cogrouped pandas udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. + + Parameters + ---------- + grouped_arg_offsets: list + List containing the key and value indexes of columns of the + DataFrames to be passed to the udf. It consists of n repeating groups where n is the + number of DataFrames. Each group has the following format: + group[0]: length of group + group[1]: length of key indexes + group[2.. group[1] +2]: key attributes + group[group[1] +3 group[0]]: value attributes + """ + parsed = [] + idx = 0 + while idx < len(grouped_arg_offsets): + offsets_len = grouped_arg_offsets[idx] + idx += 1 + offsets = grouped_arg_offsets[idx : idx + offsets_len] + split_index = offsets[0] + 1 + offset_keys = offsets[1:split_index] + offset_values = offsets[split_index:] + parsed.append([offset_keys, offset_values]) + idx += offsets_len + return parsed + + if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF: + # We assume there is only one UDF here because grouped map doesn't + # support combining multiple UDFs. + assert num_udfs == 1 + + # See FlatMapGroupsInPandasExec for how arg_offsets are used to + # distinguish between grouping attributes and data attributes + arg_offsets, f = read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=0) + parsed_offsets = extract_key_value_indexes(arg_offsets) + + # Create function like this: + # mapper a: f([a[0]], [a[0], a[1]]) + def mapper(a): + keys = [a[o] for o in parsed_offsets[0][0]] + vals = [a[o] for o in parsed_offsets[0][1]] + return f(keys, vals) + + elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: + # We assume there is only one UDF here because grouped map doesn't + # support combining multiple UDFs. + assert num_udfs == 1 + + # See FlatMapGroupsInPandas(WithState)Exec for how arg_offsets are used to + # distinguish between grouping attributes and data attributes + arg_offsets, f = read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=0) + parsed_offsets = extract_key_value_indexes(arg_offsets) + + def mapper(a): + """ + The function receives (iterator of data, state) and performs extraction of key and + value from the data, with retaining lazy evaluation. + + See `load_stream` in `ApplyInPandasWithStateSerializer` for more details on the input + and see `wrap_grouped_map_pandas_udf_with_state` for more details on how output will + be used. + """ + from itertools import tee + + state = a[1] + data_gen = (x[0] for x in a[0]) + + # We know there should be at least one item in the iterator/generator. + # We want to peek the first element to construct the key, hence applying + # tee to construct the key while we retain another iterator/generator + # for values. + keys_gen, values_gen = tee(data_gen) + keys_elem = next(keys_gen) + keys = [keys_elem[o] for o in parsed_offsets[0][0]] + + # This must be generator comprehension - do not materialize. + vals = ([x[o] for o in parsed_offsets[0][1]] for x in values_gen) + + return f(keys, vals, state) + + elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: + # We assume there is only one UDF here because cogrouped map doesn't + # support combining multiple UDFs. + assert num_udfs == 1 + arg_offsets, f = read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=0) + + parsed_offsets = extract_key_value_indexes(arg_offsets) + + def mapper(a): + df1_keys = [a[0][o] for o in parsed_offsets[0][0]] + df1_vals = [a[0][o] for o in parsed_offsets[0][1]] + df2_keys = [a[1][o] for o in parsed_offsets[1][0]] + df2_vals = [a[1][o] for o in parsed_offsets[1][1]] + return f(df1_keys, df1_vals, df2_keys, df2_vals) + + else: + udfs = [] + for i in range(num_udfs): + udfs.append(read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index=i)) + + def mapper(a): + result = tuple(f(*[a[o] for o in arg_offsets]) for (arg_offsets, f) in udfs) + # In the special case of a single UDF this will return a single result rather + # than a tuple of results; this is the format that the JVM side expects. + if len(result) == 1: + return result[0] + else: + return result + + def func(_, it): + return map(mapper, it) + + # profiling is not supported for UDF + return func, None, ser, ser + + +def main(infile, outfile): + faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None) + try: + if faulthandler_log_path: + faulthandler_log_path = os.path.join(faulthandler_log_path, str(os.getpid())) + faulthandler_log_file = open(faulthandler_log_path, "w") + faulthandler.enable(file=faulthandler_log_file) + + boot_time = time.time() + split_index = read_int(infile) + if split_index == -1: # for unit tests + sys.exit(-1) + + version = utf8_deserializer.loads(infile) + if version != "%d.%d" % sys.version_info[:2]: + raise PySparkRuntimeError( + error_class="PYTHON_VERSION_MISMATCH", + message_parameters={ + "worker_version": str(sys.version_info[:2]), + "driver_version": str(version), + }, + ) + + # read inputs only for a barrier task + isBarrier = read_bool(infile) + boundPort = read_int(infile) + secret = UTF8Deserializer().loads(infile) + + # set up memory limits + memory_limit_mb = int(os.environ.get("PYSPARK_EXECUTOR_MEMORY_MB", "-1")) + if memory_limit_mb > 0 and has_resource_module: + total_memory = resource.RLIMIT_AS + try: + (soft_limit, hard_limit) = resource.getrlimit(total_memory) + msg = "Current mem limits: {0} of max {1}\n".format(soft_limit, hard_limit) + print(msg, file=sys.stderr) + + # convert to bytes + new_limit = memory_limit_mb * 1024 * 1024 + + if soft_limit == resource.RLIM_INFINITY or new_limit < soft_limit: + msg = "Setting mem limits to {0} of max {1}\n".format(new_limit, new_limit) + print(msg, file=sys.stderr) + resource.setrlimit(total_memory, (new_limit, new_limit)) + + except (resource.error, OSError, ValueError) as e: + # not all systems support resource limits, so warn instead of failing + lineno = ( + getframeinfo(currentframe()).lineno + 1 if currentframe() is not None else 0 + ) + if "__file__" in globals(): + print( + warnings.formatwarning( + "Failed to set memory limit: {0}".format(e), + ResourceWarning, + __file__, + lineno, + ), + file=sys.stderr, + ) + + # initialize global state + taskContext = None + if isBarrier: + taskContext = BarrierTaskContext._getOrCreate() + BarrierTaskContext._initialize(boundPort, secret) + # Set the task context instance here, so we can get it by TaskContext.get for + # both TaskContext and BarrierTaskContext + TaskContext._setTaskContext(taskContext) + else: + taskContext = TaskContext._getOrCreate() + # read inputs for TaskContext info + taskContext._stageId = read_int(infile) + taskContext._partitionId = read_int(infile) + taskContext._attemptNumber = read_int(infile) + taskContext._taskAttemptId = read_long(infile) + taskContext._cpus = read_int(infile) + taskContext._resources = {} + for r in range(read_int(infile)): + key = utf8_deserializer.loads(infile) + name = utf8_deserializer.loads(infile) + addresses = [] + taskContext._resources = {} + for a in range(read_int(infile)): + addresses.append(utf8_deserializer.loads(infile)) + taskContext._resources[key] = ResourceInformation(name, addresses) + + taskContext._localProperties = dict() + for i in range(read_int(infile)): + k = utf8_deserializer.loads(infile) + v = utf8_deserializer.loads(infile) + taskContext._localProperties[k] = v + + shuffle.MemoryBytesSpilled = 0 + shuffle.DiskBytesSpilled = 0 + _accumulatorRegistry.clear() + + # fetch name of workdir + spark_files_dir = utf8_deserializer.loads(infile) + SparkFiles._root_directory = spark_files_dir + SparkFiles._is_running_on_worker = True + + # fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH + 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) + add_path(os.path.join(spark_files_dir, filename)) + + importlib.invalidate_caches() + + # fetch names and values of broadcast variables + needs_broadcast_decryption_server = read_bool(infile) + num_broadcast_variables = read_int(infile) + if needs_broadcast_decryption_server: + # read the decrypted data from a server in the jvm + port = read_int(infile) + auth_secret = utf8_deserializer.loads(infile) + (broadcast_sock_file, _) = local_connect_and_auth(port, auth_secret) + + for _ in range(num_broadcast_variables): + bid = read_long(infile) + if bid >= 0: + if needs_broadcast_decryption_server: + 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) + + if needs_broadcast_decryption_server: + broadcast_sock_file.write(b"1") + broadcast_sock_file.close() + + _accumulatorRegistry.clear() + eval_type = read_int(infile) + if eval_type == PythonEvalType.NON_UDF: + func, profiler, deserializer, serializer = read_command(pickleSer, infile) + elif eval_type in (PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF): + func, profiler, deserializer, serializer = read_udtf(pickleSer, infile, eval_type) + else: + func, profiler, deserializer, serializer = read_udfs(pickleSer, infile, eval_type) + + init_time = time.time() + + def process(): + iterator = deserializer.load_stream(infile) + out_iter = func(split_index, iterator) + try: + serializer.dump_stream(out_iter, outfile) + finally: + if hasattr(out_iter, "close"): + out_iter.close() + + if profiler: + profiler.profile(process) + else: + process() + + # Reset task context to None. This is a guard code to avoid residual context when worker + # reuse. + TaskContext._setTaskContext(None) + BarrierTaskContext._setTaskContext(None) + except BaseException as e: + try: + exc_info = None + if os.environ.get("SPARK_SIMPLIFIED_TRACEBACK", False): + tb = try_simplify_traceback(sys.exc_info()[-1]) + if tb is not None: + e.__cause__ = None + exc_info = "".join(traceback.format_exception(type(e), e, tb)) + if exc_info is None: + exc_info = traceback.format_exc() + + write_int(SpecialLengths.PYTHON_EXCEPTION_THROWN, outfile) + write_with_length(exc_info.encode("utf-8"), outfile) + except IOError: + # JVM close the socket + pass + except BaseException: + # Write the error to stderr if it happened while serializing + print("PySpark worker failed with exception:", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) + sys.exit(-1) + finally: + if faulthandler_log_path: + faulthandler.disable() + faulthandler_log_file.close() + os.remove(faulthandler_log_path) + finish_time = time.time() + report_times(outfile, boot_time, init_time, finish_time) + write_long(shuffle.MemoryBytesSpilled, outfile) + write_long(shuffle.DiskBytesSpilled, outfile) + + # Mark the beginning of the accumulators section of the output + write_int(SpecialLengths.END_OF_DATA_SECTION, outfile) + write_int(len(_accumulatorRegistry), outfile) + for (aid, accum) in _accumulatorRegistry.items(): + pickleSer._write_with_length((aid, accum._value), outfile) + + # check end of stream + if read_int(infile) == SpecialLengths.END_OF_STREAM: + write_int(SpecialLengths.END_OF_STREAM, outfile) + else: + # write a different value to tell JVM to not reuse this worker + write_int(SpecialLengths.END_OF_DATA_SECTION, outfile) + sys.exit(-1) + + +if __name__ == "__main__": + # Read information about how to connect back to the JVM from the environment. + java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"]) + auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"] + (sock_file, _) = local_connect_and_auth(java_port, auth_secret) + # TODO: Remove the following two lines and use `Process.pid()` when we drop JDK 8. + write_int(os.getpid(), sock_file) + sock_file.flush() + main(sock_file, sock_file) diff --git a/spark-3.5.0-bin-hadoop3/python/pyspark/worker_util.py b/spark-3.5.0-bin-hadoop3/python/pyspark/worker_util.py new file mode 100644 index 0000000..0c119a5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/pyspark/worker_util.py @@ -0,0 +1,46 @@ +# +# 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. +# + +""" +Util functions for workers. +""" +import sys +from typing import IO + +from pyspark.errors import PySparkRuntimeError +from pyspark.serializers import ( + UTF8Deserializer, + CPickleSerializer, +) + +pickleSer = CPickleSerializer() +utf8_deserializer = UTF8Deserializer() + + +def check_python_version(infile: IO) -> None: + """ + Check the Python version between the running process and the one used to serialize the command. + """ + version = utf8_deserializer.loads(infile) + if version != "%d.%d" % sys.version_info[:2]: + raise PySparkRuntimeError( + error_class="PYTHON_VERSION_MISMATCH", + message_parameters={ + "worker_version": str(sys.version_info[:2]), + "driver_version": str(version), + }, + ) diff --git a/spark-3.5.0-bin-hadoop3/python/run-tests b/spark-3.5.0-bin-hadoop3/python/run-tests new file mode 100644 index 0000000..401fcae --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/run-tests @@ -0,0 +1,30 @@ +#!/usr/bin/env bash + +# +# 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. +# + + +FWDIR="$(cd "`dirname $0`"/..; pwd)" +cd "$FWDIR" + +PYTHON_VERSION_CHECK=$(python3 -c 'import sys; print(sys.version_info < (3, 8, 0))') +if [[ "$PYTHON_VERSION_CHECK" == "True" ]]; then + echo "Python versions prior to 3.8 are not supported." + exit -1 +fi + +exec python3 -u ./python/run-tests.py "$@" diff --git a/spark-3.5.0-bin-hadoop3/python/run-tests-with-coverage b/spark-3.5.0-bin-hadoop3/python/run-tests-with-coverage new file mode 100644 index 0000000..0ca054d --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/run-tests-with-coverage @@ -0,0 +1,69 @@ +#!/usr/bin/env bash + +# +# 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. +# + +set -o pipefail +set -e + +# This variable indicates which coverage executable to run to combine coverages +# and generate HTMLs, for example, 'coverage3' in Python 3. +COV_EXEC="${COV_EXEC:-coverage}" +FWDIR="$(cd "`dirname $0`"; pwd)" +pushd "$FWDIR" > /dev/null + +# Ensure that coverage executable is installed. +if ! hash $COV_EXEC 2>/dev/null; then + echo "Missing coverage executable in your path, skipping PySpark coverage" + exit 1 +fi + +# Set up the directories for coverage results. +export COVERAGE_DIR="$FWDIR/test_coverage" +rm -fr "$COVERAGE_DIR/coverage_data" +rm -fr "$COVERAGE_DIR/htmlcov" +mkdir -p "$COVERAGE_DIR/coverage_data" + +# Current directory are added in the python path so that it doesn't refer our built +# pyspark zip library first. +export PYTHONPATH="$FWDIR:$PYTHONPATH" +# Also, our sitecustomize.py and coverage_daemon.py are included in the path. +export PYTHONPATH="$COVERAGE_DIR:$PYTHONPATH" + +# We use 'spark.python.daemon.module' configuration to insert the coverage supported workers. +export SPARK_CONF_DIR="$COVERAGE_DIR/conf" + +# This environment variable enables the coverage. +export COVERAGE_PROCESS_START="$FWDIR/.coveragerc" + +./run-tests "$@" + +# Don't run coverage for the coverage command itself +unset COVERAGE_PROCESS_START + +# Coverage could generate empty coverage data files. Remove it to get rid of warnings when combining. +find $COVERAGE_DIR/coverage_data -size 0 -print0 | xargs -0 rm -fr +echo "Combining collected coverage data under $COVERAGE_DIR/coverage_data" +$COV_EXEC combine +echo "Creating XML report file at python/coverage.xml" +$COV_EXEC xml --ignore-errors --include "pyspark/*" --omit "pyspark/cloudpickle/*" +echo "Reporting the coverage data at $COVERAGE_DIR/coverage_data/coverage" +$COV_EXEC report --include "pyspark/*" --omit "pyspark/cloudpickle/*" +echo "Generating HTML files for PySpark coverage under $COVERAGE_DIR/htmlcov" +$COV_EXEC html --ignore-errors --include "pyspark/*" --directory "$COVERAGE_DIR/htmlcov" --omit "pyspark/cloudpickle/*" + +popd diff --git a/spark-3.5.0-bin-hadoop3/python/run-tests.py b/spark-3.5.0-bin-hadoop3/python/run-tests.py new file mode 100644 index 0000000..19e39c8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/run-tests.py @@ -0,0 +1,374 @@ +#!/usr/bin/env python3 + +# +# 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 logging +from argparse import ArgumentParser +import os +import platform +import re +import shutil +import subprocess +import sys +import tempfile +from threading import Thread, Lock +import time +import uuid +import queue as Queue +from multiprocessing import Manager + + +# Append `SPARK_HOME/dev` to the Python path so that we can import the sparktestsupport module +sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../dev/")) + + +from sparktestsupport import SPARK_HOME # noqa (suppress pep8 warnings) +from sparktestsupport.shellutils import which, subprocess_check_output # noqa +from sparktestsupport.modules import all_modules, pyspark_sql # noqa + + +python_modules = dict((m.name, m) for m in all_modules if m.python_test_goals if m.name != 'root') + + +def print_red(text): + print('\033[31m' + text + '\033[0m') + + +def get_valid_filename(s): + """Replace whitespaces and special characters in the given string to get a valid file name.""" + s = s.strip().replace(' ', '_').replace(os.sep, '_') + return re.sub(r'(?u)[^-\w.]', '', s) + + +SKIPPED_TESTS = None +LOG_FILE = os.path.join(SPARK_HOME, "python/unit-tests.log") +FAILURE_REPORTING_LOCK = Lock() +LOGGER = logging.getLogger() + +# Find out where the assembly jars are located. +# TODO: revisit for Scala 2.13 +for scala in ["2.12", "2.13"]: + build_dir = os.path.join(SPARK_HOME, "assembly", "target", "scala-" + scala) + if os.path.isdir(build_dir): + SPARK_DIST_CLASSPATH = os.path.join(build_dir, "jars", "*") + break +else: + raise RuntimeError("Cannot find assembly build directory, please build Spark first.") + + +def run_individual_python_test(target_dir, test_name, pyspark_python, keep_test_output): + """ + Runs an individual test. This function is called by the multi-process runner of all tests. + + Parameters + ---------- + target_dir + Destination for the Hive and log directory. + test_name + Test name. + pyspark_python + Python version used to run the test. + keep_test_output + Flag indicating if the test output should be retained after successful execution. + + """ + env = dict(os.environ) + env.update({ + 'SPARK_DIST_CLASSPATH': SPARK_DIST_CLASSPATH, + 'SPARK_TESTING': '1', + 'SPARK_PREPEND_CLASSES': '1', + 'PYSPARK_PYTHON': which(pyspark_python), + 'PYSPARK_DRIVER_PYTHON': which(pyspark_python), + # Preserve legacy nested timezone behavior for pyarrow>=2, remove after SPARK-32285 + 'PYARROW_IGNORE_TIMEZONE': '1', + }) + + # Create a unique temp directory under 'target/' for each run. The TMPDIR variable is + # recognized by the tempfile module to override the default system temp directory. + tmp_dir = os.path.join(target_dir, str(uuid.uuid4())) + while os.path.isdir(tmp_dir): + tmp_dir = os.path.join(target_dir, str(uuid.uuid4())) + os.mkdir(tmp_dir) + env["TMPDIR"] = tmp_dir + metastore_dir = os.path.join(tmp_dir, str(uuid.uuid4())) + while os.path.isdir(metastore_dir): + metastore_dir = os.path.join(metastore_dir, str(uuid.uuid4())) + os.mkdir(metastore_dir) + + # Also override the JVM's temp directory by setting driver and executor options. + java_options = "-Djava.io.tmpdir={0}".format(tmp_dir) + java_options = java_options + " -Dio.netty.tryReflectionSetAccessible=true -Xss4M" + spark_args = [ + "--conf", "spark.driver.extraJavaOptions='{0}'".format(java_options), + "--conf", "spark.executor.extraJavaOptions='{0}'".format(java_options), + "--conf", "spark.sql.warehouse.dir='{0}'".format(metastore_dir), + "pyspark-shell", + ] + + env["PYSPARK_SUBMIT_ARGS"] = " ".join(spark_args) + + output_prefix = get_valid_filename(pyspark_python + "__" + test_name + "__").lstrip("_") + # Delete is always set to False since the cleanup will be either done by removing the + # whole test dir, or the test output is retained. + per_test_output = tempfile.NamedTemporaryFile(prefix=output_prefix, dir=tmp_dir, + suffix=".log", delete=False) + LOGGER.info( + "Starting test(%s): %s (temp output: %s)", pyspark_python, test_name, per_test_output.name) + start_time = time.time() + try: + retcode = subprocess.Popen( + [os.path.join(SPARK_HOME, "bin/pyspark")] + test_name.split(), + stderr=per_test_output, stdout=per_test_output, env=env).wait() + if not keep_test_output: + # There exists a race condition in Python and it causes flakiness in MacOS + # https://github.com/python/cpython/issues/73885 + if platform.system() == "Darwin": + os.system("rm -rf " + tmp_dir) + else: + shutil.rmtree(tmp_dir, ignore_errors=True) + except BaseException: + LOGGER.exception("Got exception while running %s with %s", test_name, pyspark_python) + # Here, we use os._exit() instead of sys.exit() in order to force Python to exit even if + # this code is invoked from a thread other than the main thread. + os._exit(1) + duration = time.time() - start_time + # Exit on the first failure. + if retcode != 0: + try: + with FAILURE_REPORTING_LOCK: + with open(LOG_FILE, 'ab') as log_file: + per_test_output.seek(0) + log_file.writelines(per_test_output) + per_test_output.seek(0) + for line in per_test_output: + decoded_line = line.decode("utf-8", "replace") + if not re.match('[0-9]+', decoded_line): + print(decoded_line, end='') + per_test_output.close() + except BaseException: + LOGGER.exception("Got an exception while trying to print failed test output") + finally: + print_red("\nHad test failures in %s with %s; see logs." % (test_name, pyspark_python)) + # Here, we use os._exit() instead of sys.exit() in order to force Python to exit even if + # this code is invoked from a thread other than the main thread. + os._exit(-1) + else: + skipped_counts = 0 + try: + per_test_output.seek(0) + # Here expects skipped test output from unittest when verbosity level is + # 2 (or --verbose option is enabled). + decoded_lines = map(lambda line: line.decode("utf-8", "replace"), iter(per_test_output)) + skipped_tests = list(filter( + lambda line: re.search(r'test_.* \(pyspark\..*\) ... (skip|SKIP)', line), + decoded_lines)) + skipped_counts = len(skipped_tests) + if skipped_counts > 0: + key = (pyspark_python, test_name) + assert SKIPPED_TESTS is not None + SKIPPED_TESTS[key] = skipped_tests + per_test_output.close() + except BaseException: + import traceback + print_red("\nGot an exception while trying to store " + "skipped test output:\n%s" % traceback.format_exc()) + # Here, we use os._exit() instead of sys.exit() in order to force Python to exit even if + # this code is invoked from a thread other than the main thread. + os._exit(-1) + if skipped_counts != 0: + LOGGER.info( + "Finished test(%s): %s (%is) ... %s tests were skipped", pyspark_python, test_name, + duration, skipped_counts) + else: + LOGGER.info( + "Finished test(%s): %s (%is)", pyspark_python, test_name, duration) + + +def get_default_python_executables(): + python_execs = [x for x in ["python3.9", "pypy3"] if which(x)] + + if "python3.9" not in python_execs: + p = which("python3") + if not p: + LOGGER.error("No python3 executable found. Exiting!") + os._exit(1) + else: + python_execs.insert(0, p) + return python_execs + + +def parse_opts(): + parser = ArgumentParser( + prog="run-tests" + ) + parser.add_argument( + "--python-executables", type=str, default=','.join(get_default_python_executables()), + help="A comma-separated list of Python executables to test against (default: %(default)s)" + ) + parser.add_argument( + "--modules", type=str, + default=",".join(sorted(python_modules.keys())), + help="A comma-separated list of Python modules to test (default: %(default)s)" + ) + parser.add_argument( + "-p", "--parallelism", type=int, default=4, + help="The number of suites to test in parallel (default %(default)d)" + ) + parser.add_argument( + "--verbose", action="store_true", + help="Enable additional debug logging" + ) + + group = parser.add_argument_group("Developer Options") + group.add_argument( + "--testnames", type=str, + default=None, + help=( + "A comma-separated list of specific modules, classes and functions of doctest " + "or unittest to test. " + "For example, 'pyspark.sql.foo' to run the module as unittests or doctests, " + "'pyspark.sql.tests FooTests' to run the specific class of unittests, " + "'pyspark.sql.tests FooTests.test_foo' to run the specific unittest in the class. " + "'--modules' option is ignored if they are given.") + ) + group.add_argument( + "-k", "--keep-test-output", action='store_true', + default=False, + help=("If set to true will retain the temporary test directories. In addition, the " + "standard output and standard error are redirected to a file in the target " + "directory.") + ) + + args, unknown = parser.parse_known_args() + if unknown: + parser.error("Unsupported arguments: %s" % ' '.join(unknown)) + if args.parallelism < 1: + parser.error("Parallelism cannot be less than 1") + return args + + +def _check_coverage(python_exec): + # Make sure if coverage is installed. + try: + subprocess_check_output( + [python_exec, "-c", "import coverage"], + stderr=open(os.devnull, 'w')) + except BaseException: + print_red("Coverage is not installed in Python executable '%s' " + "but 'COVERAGE_PROCESS_START' environment variable is set, " + "exiting." % python_exec) + sys.exit(-1) + + +def main(): + opts = parse_opts() + if opts.verbose: + log_level = logging.DEBUG + else: + log_level = logging.INFO + should_test_modules = opts.testnames is None + logging.basicConfig(stream=sys.stdout, level=log_level, format="%(message)s") + LOGGER.info("Running PySpark tests. Output is in %s", LOG_FILE) + if os.path.exists(LOG_FILE): + os.remove(LOG_FILE) + python_execs = opts.python_executables.split(',') + LOGGER.info("Will test against the following Python executables: %s", python_execs) + + if should_test_modules: + modules_to_test = [] + for module_name in opts.modules.split(','): + if module_name in python_modules: + modules_to_test.append(python_modules[module_name]) + else: + print("Error: unrecognized module '%s'. Supported modules: %s" % + (module_name, ", ".join(python_modules))) + sys.exit(-1) + LOGGER.info("Will test the following Python modules: %s", [x.name for x in modules_to_test]) + else: + testnames_to_test = opts.testnames.split(',') + LOGGER.info("Will test the following Python tests: %s", testnames_to_test) + + task_queue = Queue.PriorityQueue() + for python_exec in python_execs: + # Check if the python executable has coverage installed when 'COVERAGE_PROCESS_START' + # environmental variable is set. + if "COVERAGE_PROCESS_START" in os.environ: + _check_coverage(python_exec) + + python_implementation = subprocess_check_output( + [python_exec, "-c", "import platform; print(platform.python_implementation())"], + universal_newlines=True).strip() + LOGGER.info("%s python_implementation is %s", python_exec, python_implementation) + LOGGER.info("%s version is: %s", python_exec, subprocess_check_output( + [python_exec, "--version"], stderr=subprocess.STDOUT, universal_newlines=True).strip()) + if should_test_modules: + for module in modules_to_test: + if python_implementation not in module.excluded_python_implementations: + for test_goal in module.python_test_goals: + heavy_tests = ['pyspark.streaming.tests', 'pyspark.mllib.tests', + 'pyspark.tests', 'pyspark.sql.tests', 'pyspark.ml.tests', + 'pyspark.pandas.tests'] + if any(map(lambda prefix: test_goal.startswith(prefix), heavy_tests)): + priority = 0 + else: + priority = 100 + task_queue.put((priority, (python_exec, test_goal))) + else: + for test_goal in testnames_to_test: + task_queue.put((0, (python_exec, test_goal))) + + # Create the target directory before starting tasks to avoid races. + target_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), 'target')) + if not os.path.isdir(target_dir): + os.mkdir(target_dir) + + def process_queue(task_queue): + while True: + try: + (priority, (python_exec, test_goal)) = task_queue.get_nowait() + except Queue.Empty: + break + try: + run_individual_python_test(target_dir, test_goal, + python_exec, opts.keep_test_output) + finally: + task_queue.task_done() + + start_time = time.time() + for _ in range(opts.parallelism): + worker = Thread(target=process_queue, args=(task_queue,)) + worker.daemon = True + worker.start() + try: + task_queue.join() + except (KeyboardInterrupt, SystemExit): + print_red("Exiting due to interrupt") + sys.exit(-1) + total_duration = time.time() - start_time + LOGGER.info("Tests passed in %i seconds", total_duration) + + for key, lines in sorted(SKIPPED_TESTS.items()): + pyspark_python, test_name = key + LOGGER.info("\nSkipped tests in %s with %s:" % (test_name, pyspark_python)) + for line in lines: + LOGGER.info(" %s" % line.rstrip()) + + +if __name__ == "__main__": + SKIPPED_TESTS = Manager().dict() + main() diff --git a/spark-3.5.0-bin-hadoop3/python/setup.cfg b/spark-3.5.0-bin-hadoop3/python/setup.cfg new file mode 100644 index 0000000..d100b93 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/setup.cfg @@ -0,0 +1,22 @@ +# +# 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. +# + +[bdist_wheel] +universal = 1 + +[metadata] +description-file = README.md diff --git a/spark-3.5.0-bin-hadoop3/python/setup.py b/spark-3.5.0-bin-hadoop3/python/setup.py new file mode 100644 index 0000000..b8e4c9a --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/setup.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 + +# +# 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 importlib.util +import glob +import os +import sys +import ctypes +from setuptools import setup +from setuptools.command.install import install +from shutil import copyfile, copytree, rmtree + +try: + exec(open("pyspark/version.py").read()) +except IOError: + print( + "Failed to load PySpark version file for packaging. You must be in Spark's python dir.", + file=sys.stderr, + ) + sys.exit(-1) +try: + spec = importlib.util.spec_from_file_location("install", "pyspark/install.py") + install_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(install_module) +except IOError: + print( + "Failed to load the installing module (pyspark/install.py) which had to be " + "packaged together.", + file=sys.stderr, + ) + sys.exit(-1) +VERSION = __version__ # noqa +# A temporary path so we can access above the Python project root and fetch scripts and jars we need +TEMP_PATH = "deps" +SPARK_HOME = os.path.abspath("../") + +# Provide guidance about how to use setup.py +incorrect_invocation_message = """ +If you are installing pyspark from spark source, you must first build Spark and +run sdist. + + To build Spark with maven you can run: + ./build/mvn -DskipTests clean package + Building the source dist is done in the Python directory: + cd python + python setup.py sdist + pip install dist/*.tar.gz""" + +# Figure out where the jars are we need to package with PySpark. +JARS_PATH = glob.glob(os.path.join(SPARK_HOME, "assembly/target/scala-*/jars/")) + +if len(JARS_PATH) == 1: + JARS_PATH = JARS_PATH[0] +elif os.path.isfile("../RELEASE") and len(glob.glob("../jars/spark*core*.jar")) == 1: + # Release mode puts the jars in a jars directory + JARS_PATH = os.path.join(SPARK_HOME, "jars") +elif len(JARS_PATH) > 1: + print( + "Assembly jars exist for multiple scalas ({0}), please cleanup assembly/target".format( + JARS_PATH + ), + file=sys.stderr, + ) + sys.exit(-1) +elif len(JARS_PATH) == 0 and not os.path.exists(TEMP_PATH): + print(incorrect_invocation_message, file=sys.stderr) + sys.exit(-1) + +EXAMPLES_PATH = os.path.join(SPARK_HOME, "examples/src/main/python") +SCRIPTS_PATH = os.path.join(SPARK_HOME, "bin") +USER_SCRIPTS_PATH = os.path.join(SPARK_HOME, "sbin") +DATA_PATH = os.path.join(SPARK_HOME, "data") +LICENSES_PATH = os.path.join(SPARK_HOME, "licenses") + +SCRIPTS_TARGET = os.path.join(TEMP_PATH, "bin") +USER_SCRIPTS_TARGET = os.path.join(TEMP_PATH, "sbin") +JARS_TARGET = os.path.join(TEMP_PATH, "jars") +EXAMPLES_TARGET = os.path.join(TEMP_PATH, "examples") +DATA_TARGET = os.path.join(TEMP_PATH, "data") +LICENSES_TARGET = os.path.join(TEMP_PATH, "licenses") + +# Check and see if we are under the spark path in which case we need to build the symlink farm. +# This is important because we only want to build the symlink farm while under Spark otherwise we +# want to use the symlink farm. And if the symlink farm exists under while under Spark (e.g. a +# partially built sdist) we should error and have the user sort it out. +in_spark = os.path.isfile("../core/src/main/scala/org/apache/spark/SparkContext.scala") or ( + os.path.isfile("../RELEASE") and len(glob.glob("../jars/spark*core*.jar")) == 1 +) + + +def _supports_symlinks(): + """Check if the system supports symlinks (e.g. *nix) or not.""" + return hasattr(os, "symlink") and ( + (not hasattr(ctypes, "windll")) # Non-Windows + or ( + # In some Windows, `os.symlink` works but only for admins. + hasattr(ctypes.windll, "shell32") + and hasattr(ctypes.windll.shell32, "IsUserAnAdmin") + and bool(ctypes.windll.shell32.IsUserAnAdmin()) + ) + ) + + +if in_spark: + # Construct links for setup + try: + os.mkdir(TEMP_PATH) + except BaseException: + print( + "Temp path for symlink to parent already exists {0}".format(TEMP_PATH), file=sys.stderr + ) + sys.exit(-1) + +# If you are changing the versions here, please also change ./python/pyspark/sql/pandas/utils.py +# For Arrow, you should also check ./pom.xml and ensure there are no breaking changes in the +# binary format protocol with the Java version, see ARROW_HOME/format/* for specifications. +# Also don't forget to update python/docs/source/getting_started/install.rst. +_minimum_pandas_version = "1.0.5" +_minimum_pyarrow_version = "4.0.0" +_minimum_grpc_version = "1.56.0" +_minimum_googleapis_common_protos_version = "1.56.4" + + +class InstallCommand(install): + # TODO(SPARK-32837) leverage pip's custom options + + def run(self): + install.run(self) + + # Make sure the destination is always clean. + spark_dist = os.path.join(self.install_lib, "pyspark", "spark-distribution") + rmtree(spark_dist, ignore_errors=True) + + if ("PYSPARK_HADOOP_VERSION" in os.environ) or ("PYSPARK_HIVE_VERSION" in os.environ): + # Note that PYSPARK_VERSION environment is just a testing purpose. + # PYSPARK_HIVE_VERSION environment variable is also internal for now in case + # we support another version of Hive in the future. + spark_version, hadoop_version, hive_version = install_module.checked_versions( + os.environ.get("PYSPARK_VERSION", VERSION).lower(), + os.environ.get("PYSPARK_HADOOP_VERSION", install_module.DEFAULT_HADOOP).lower(), + os.environ.get("PYSPARK_HIVE_VERSION", install_module.DEFAULT_HIVE).lower(), + ) + + if "PYSPARK_VERSION" not in os.environ and ( + (install_module.DEFAULT_HADOOP, install_module.DEFAULT_HIVE) + == (hadoop_version, hive_version) + ): + # Do not download and install if they are same as default. + return + + install_module.install_spark( + dest=spark_dist, + spark_version=spark_version, + hadoop_version=hadoop_version, + hive_version=hive_version, + ) + + +try: + # We copy the shell script to be under pyspark/python/pyspark so that the launcher scripts + # find it where expected. The rest of the files aren't copied because they are accessed + # using Python imports instead which will be resolved correctly. + try: + os.makedirs("pyspark/python/pyspark") + except OSError: + # Don't worry if the directory already exists. + pass + copyfile("pyspark/shell.py", "pyspark/python/pyspark/shell.py") + + if in_spark: + # Construct the symlink farm - this is necessary since we can't refer to the path above the + # package root and we need to copy the jars and scripts which are up above the python root. + if _supports_symlinks(): + os.symlink(JARS_PATH, JARS_TARGET) + os.symlink(SCRIPTS_PATH, SCRIPTS_TARGET) + os.symlink(USER_SCRIPTS_PATH, USER_SCRIPTS_TARGET) + os.symlink(EXAMPLES_PATH, EXAMPLES_TARGET) + os.symlink(DATA_PATH, DATA_TARGET) + os.symlink(LICENSES_PATH, LICENSES_TARGET) + else: + # For windows fall back to the slower copytree + copytree(JARS_PATH, JARS_TARGET) + copytree(SCRIPTS_PATH, SCRIPTS_TARGET) + copytree(USER_SCRIPTS_PATH, USER_SCRIPTS_TARGET) + copytree(EXAMPLES_PATH, EXAMPLES_TARGET) + copytree(DATA_PATH, DATA_TARGET) + copytree(LICENSES_PATH, LICENSES_TARGET) + else: + # If we are not inside of SPARK_HOME verify we have the required symlink farm + if not os.path.exists(JARS_TARGET): + print( + "To build packaging must be in the python directory under the SPARK_HOME.", + file=sys.stderr, + ) + + if not os.path.isdir(SCRIPTS_TARGET): + print(incorrect_invocation_message, file=sys.stderr) + sys.exit(-1) + + # Scripts directive requires a list of each script path and does not take wild cards. + script_names = os.listdir(SCRIPTS_TARGET) + scripts = list(map(lambda script: os.path.join(SCRIPTS_TARGET, script), script_names)) + # We add find_spark_home.py to the bin directory we install so that pip installed PySpark + # will search for SPARK_HOME with Python. + scripts.append("pyspark/find_spark_home.py") + + with open("README.md") as f: + long_description = f.read() + + setup( + name="pyspark", + version=VERSION, + description="Apache Spark Python API", + long_description=long_description, + long_description_content_type="text/markdown", + author="Spark Developers", + author_email="dev@spark.apache.org", + url="https://github.com/apache/spark/tree/master/python", + packages=[ + "pyspark", + "pyspark.cloudpickle", + "pyspark.mllib", + "pyspark.ml.connect", + "pyspark.mllib.linalg", + "pyspark.mllib.stat", + "pyspark.ml", + "pyspark.ml.linalg", + "pyspark.ml.param", + "pyspark.ml.torch", + "pyspark.ml.deepspeed", + "pyspark.sql", + "pyspark.sql.avro", + "pyspark.sql.connect", + "pyspark.sql.connect.avro", + "pyspark.sql.connect.client", + "pyspark.sql.connect.proto", + "pyspark.sql.connect.streaming", + "pyspark.sql.pandas", + "pyspark.sql.protobuf", + "pyspark.sql.streaming", + "pyspark.streaming", + "pyspark.sql.connect.streaming.worker", + "pyspark.bin", + "pyspark.sbin", + "pyspark.jars", + "pyspark.pandas", + "pyspark.pandas.data_type_ops", + "pyspark.pandas.indexes", + "pyspark.pandas.missing", + "pyspark.pandas.plot", + "pyspark.pandas.spark", + "pyspark.pandas.typedef", + "pyspark.pandas.usage_logging", + "pyspark.python.pyspark", + "pyspark.python.lib", + "pyspark.testing", + "pyspark.data", + "pyspark.licenses", + "pyspark.resource", + "pyspark.errors", + "pyspark.errors.exceptions", + "pyspark.examples.src.main.python", + ], + include_package_data=True, + package_dir={ + "pyspark.jars": "deps/jars", + "pyspark.bin": "deps/bin", + "pyspark.sbin": "deps/sbin", + "pyspark.python.lib": "lib", + "pyspark.data": "deps/data", + "pyspark.licenses": "deps/licenses", + "pyspark.examples.src.main.python": "deps/examples", + }, + package_data={ + "pyspark.jars": ["*.jar"], + "pyspark.bin": ["*"], + "pyspark.sbin": [ + "spark-config.sh", + "spark-daemon.sh", + "start-history-server.sh", + "stop-history-server.sh", + ], + "pyspark.python.lib": ["*.zip"], + "pyspark.data": ["*.txt", "*.data"], + "pyspark.licenses": ["*.txt"], + "pyspark.examples.src.main.python": ["*.py", "*/*.py"], + }, + scripts=scripts, + license="http://www.apache.org/licenses/LICENSE-2.0", + # Don't forget to update python/docs/source/getting_started/install.rst + # if you're updating the versions or dependencies. + install_requires=["py4j==0.10.9.7"], + extras_require={ + "ml": ["numpy>=1.15"], + "mllib": ["numpy>=1.15"], + "sql": [ + "pandas>=%s" % _minimum_pandas_version, + "pyarrow>=%s" % _minimum_pyarrow_version, + "numpy>=1.15", + ], + "pandas_on_spark": [ + "pandas>=%s" % _minimum_pandas_version, + "pyarrow>=%s" % _minimum_pyarrow_version, + "numpy>=1.15", + ], + "connect": [ + "pandas>=%s" % _minimum_pandas_version, + "pyarrow>=%s" % _minimum_pyarrow_version, + "grpcio>=%s" % _minimum_grpc_version, + "grpcio-status>=%s" % _minimum_grpc_version, + "googleapis-common-protos>=%s" % _minimum_googleapis_common_protos_version, + "numpy>=1.15", + ], + }, + python_requires=">=3.8", + classifiers=[ + "Development Status :: 5 - Production/Stable", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", + "Typing :: Typed", + ], + cmdclass={ + "install": InstallCommand, + }, + ) +finally: + # We only cleanup the symlink farm if we were in Spark, otherwise we are installing rather than + # packaging. + if in_spark: + # Depending on cleaning up the symlink farm or copied version + if _supports_symlinks(): + os.remove(os.path.join(TEMP_PATH, "jars")) + os.remove(os.path.join(TEMP_PATH, "bin")) + os.remove(os.path.join(TEMP_PATH, "sbin")) + os.remove(os.path.join(TEMP_PATH, "examples")) + os.remove(os.path.join(TEMP_PATH, "data")) + os.remove(os.path.join(TEMP_PATH, "licenses")) + else: + rmtree(os.path.join(TEMP_PATH, "jars")) + rmtree(os.path.join(TEMP_PATH, "bin")) + rmtree(os.path.join(TEMP_PATH, "sbin")) + rmtree(os.path.join(TEMP_PATH, "examples")) + rmtree(os.path.join(TEMP_PATH, "data")) + rmtree(os.path.join(TEMP_PATH, "licenses")) + os.rmdir(TEMP_PATH) diff --git a/spark-3.5.0-bin-hadoop3/python/test_coverage/conf/spark-defaults.conf b/spark-3.5.0-bin-hadoop3/python/test_coverage/conf/spark-defaults.conf new file mode 100644 index 0000000..bf44ea6 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_coverage/conf/spark-defaults.conf @@ -0,0 +1,21 @@ +# +# 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. +# + +# This is used to generate PySpark coverage results. Seems there's no way to +# add a configuration when SPARK_TESTING environment variable is set because +# we will directly execute modules by python -m. +spark.python.daemon.module coverage_daemon diff --git a/spark-3.5.0-bin-hadoop3/python/test_coverage/coverage_daemon.py b/spark-3.5.0-bin-hadoop3/python/test_coverage/coverage_daemon.py new file mode 100644 index 0000000..4372135 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_coverage/coverage_daemon.py @@ -0,0 +1,48 @@ +# +# 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 os +import imp +import platform + + +# This is a hack to always refer the main code rather than built zip. +main_code_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) +daemon = imp.load_source("daemon", "%s/pyspark/daemon.py" % main_code_dir) + +if "COVERAGE_PROCESS_START" in os.environ: + # PyPy with coverage makes the tests flaky, and CPython is enough for coverage report. + if "pypy" not in platform.python_implementation().lower(): + worker = imp.load_source("worker", "%s/pyspark/worker.py" % main_code_dir) + + def _cov_wrapped(*args, **kwargs): + import coverage + cov = coverage.coverage( + config_file=os.environ["COVERAGE_PROCESS_START"]) + cov.start() + try: + worker.main(*args, **kwargs) + finally: + cov.stop() + cov.save() + daemon.worker_main = _cov_wrapped +else: + raise RuntimeError("COVERAGE_PROCESS_START environment variable is not set, exiting.") + + +if __name__ == '__main__': + daemon.manager() diff --git a/spark-3.5.0-bin-hadoop3/python/test_coverage/sitecustomize.py b/spark-3.5.0-bin-hadoop3/python/test_coverage/sitecustomize.py new file mode 100644 index 0000000..1f31860 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_coverage/sitecustomize.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. +# + +# Note that this 'sitecustomize' module is a built-in feature in Python. +# If this module is defined, it's executed when the Python session begins. +# `coverage.process_startup()` seeks if COVERAGE_PROCESS_START environment +# variable is set or not. If set, it starts to run the coverage. +try: + import coverage + coverage.process_startup() +except ImportError: + pass diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/SimpleHTTPServer.py b/spark-3.5.0-bin-hadoop3/python/test_support/SimpleHTTPServer.py new file mode 100644 index 0000000..eddbd58 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/SimpleHTTPServer.py @@ -0,0 +1,22 @@ +# +# 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. +# + +""" +Used to test override standard SimpleHTTPServer module. +""" + +__name__ = "My Server" diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/hello/hello.txt b/spark-3.5.0-bin-hadoop3/python/test_support/hello/hello.txt new file mode 100644 index 0000000..980a0d5 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/hello/hello.txt @@ -0,0 +1 @@ +Hello World! diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/hello/sub_hello/sub_hello.txt b/spark-3.5.0-bin-hadoop3/python/test_support/hello/sub_hello/sub_hello.txt new file mode 100644 index 0000000..ce2d435 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/hello/sub_hello/sub_hello.txt @@ -0,0 +1 @@ +Sub Hello World! diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/ages_newlines.csv b/spark-3.5.0-bin-hadoop3/python/test_support/sql/ages_newlines.csv new file mode 100644 index 0000000..d19f673 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/ages_newlines.csv @@ -0,0 +1,6 @@ +Joe,20,"Hi, +I am Jeo" +Tom,30,"My name is Tom" +Hyukjin,25,"I am Hyukjin + +I love Spark!" diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/_SUCCESS b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/_SUCCESS new file mode 100644 index 0000000..e69de29 diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc new file mode 100644 index 0000000..834cf0b Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc new file mode 100644 index 0000000..4943801 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=0/c=0/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc new file mode 100644 index 0000000..693dcee Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/.part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc.crc differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc new file mode 100644 index 0000000..4cbb95a Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/sql/orc_partitioned/b=1/c=1/part-r-00000-829af031-b970-49d6-ad39-30460a0be2c8.orc differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/people.json b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people.json new file mode 100644 index 0000000..50a859c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people.json @@ -0,0 +1,3 @@ +{"name":"Michael"} +{"name":"Andy", "age":30} +{"name":"Justin", "age":19} diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/people1.json b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people1.json new file mode 100644 index 0000000..6d217da --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people1.json @@ -0,0 +1,2 @@ +{"name":"Jonathan", "aka": "John"} + diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array.json b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array.json new file mode 100644 index 0000000..c27c48f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array.json @@ -0,0 +1,13 @@ +[ + { + "name": "Michael" + }, + { + "name": "Andy", + "age": 30 + }, + { + "name": "Justin", + "age": 19 + } +] diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array_utf16le.json b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array_utf16le.json new file mode 100644 index 0000000..9c657fa Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/sql/people_array_utf16le.json differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/streaming/text-test.txt b/spark-3.5.0-bin-hadoop3/python/test_support/sql/streaming/text-test.txt new file mode 100644 index 0000000..ae1e76c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/streaming/text-test.txt @@ -0,0 +1,2 @@ +hello +this \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/sql/text-test.txt b/spark-3.5.0-bin-hadoop3/python/test_support/sql/text-test.txt new file mode 100644 index 0000000..ae1e76c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/sql/text-test.txt @@ -0,0 +1,2 @@ +hello +this \ No newline at end of file diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/test_pytorch_training_file.py b/spark-3.5.0-bin-hadoop3/python/test_support/test_pytorch_training_file.py new file mode 100644 index 0000000..4107197 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/test_pytorch_training_file.py @@ -0,0 +1,115 @@ +# +# 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. +# +# type: ignore + +batch_size = 100 +num_epochs = 3 +momentum = 0.5 +log_interval = 100 + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torchvision import datasets, transforms +import tempfile +import shutil + + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x) + + +def train_one_epoch(model, data_loader, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(data_loader): + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % log_interval == 0: + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(data_loader) * len(data), + 100.0 * batch_idx / len(data_loader), + loss.item(), + ) + ) + + +def train(learning_rate): + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel as DDP + from torch.utils.data.distributed import DistributedSampler + + print("Running distributed training") + dist.init_process_group("gloo") + + temp_dir = tempfile.mkdtemp() + + train_dataset = datasets.MNIST( + temp_dir, + train=True, + download=True, + transform=transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ), + ) + + train_sampler = DistributedSampler(dataset=train_dataset) + data_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, sampler=train_sampler + ) + + model = Net() + ddp_model = DDP(model) + + optimizer = optim.SGD(ddp_model.parameters(), lr=learning_rate, momentum=momentum) + for epoch in range(1, num_epochs + 1): + train_one_epoch(ddp_model, data_loader, optimizer, epoch) + + dist.destroy_process_group() + + shutil.rmtree(temp_dir) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("lr", help="learning_rate", default=0.001) + args = parser.parse_args() + print("learning rate chosen: ", float(args.lr)) + train(float(args.lr)) diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/userlib-0.1.zip b/spark-3.5.0-bin-hadoop3/python/test_support/userlib-0.1.zip new file mode 100644 index 0000000..496e134 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/python/test_support/userlib-0.1.zip differ diff --git a/spark-3.5.0-bin-hadoop3/python/test_support/userlibrary.py b/spark-3.5.0-bin-hadoop3/python/test_support/userlibrary.py new file mode 100644 index 0000000..2d262ba --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/python/test_support/userlibrary.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. +# + +""" +Used to test shipping of code dependencies with SparkContext.addPyFile(). +""" + + +class UserClass: + + def hello(self): + return "Hello World!" diff --git a/spark-3.5.0-bin-hadoop3/sbin/decommission-slave.sh b/spark-3.5.0-bin-hadoop3/sbin/decommission-slave.sh new file mode 100644 index 0000000..858bede --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/decommission-slave.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use decommission-worker.sh" +"${DIR}/decommission-worker.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/decommission-worker.sh b/spark-3.5.0-bin-hadoop3/sbin/decommission-worker.sh new file mode 100644 index 0000000..07e1e17 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/decommission-worker.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# A shell script to decommission all workers on a single worker +# +# Environment variables +# +# SPARK_WORKER_INSTANCES The number of worker instances that should be +# running on this worker machine. Default is 1. + +# Usage: decommission-worker.sh [--block-until-exit] +# Decommissions all workers on this worker machine. + +set -ex + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$SPARK_WORKER_INSTANCES" = "" ]; then + "${SPARK_HOME}/sbin"/spark-daemon.sh decommission org.apache.spark.deploy.worker.Worker 1 +else + for ((i=0; i<$SPARK_WORKER_INSTANCES; i++)); do + "${SPARK_HOME}/sbin"/spark-daemon.sh decommission org.apache.spark.deploy.worker.Worker $(( $i + 1 )) + done +fi + +# Check if --block-until-exit is set. +# This is done for systems which block on the decommissioning script and on exit +# shut down the entire system (e.g. K8s). +if [ "$1" == "--block-until-exit" ]; then + shift + # For now we only block on the 0th instance if there multiple instances. + instance=$1 + pid="$SPARK_PID_DIR/spark-$SPARK_IDENT_STRING-$command-$instance.pid" + wait $pid +fi diff --git a/spark-3.5.0-bin-hadoop3/sbin/slaves.sh b/spark-3.5.0-bin-hadoop3/sbin/slaves.sh new file mode 100644 index 0000000..b92007e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/slaves.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use workers.sh" +"${DIR}/workers.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/spark-config.sh b/spark-3.5.0-bin-hadoop3/sbin/spark-config.sh new file mode 100644 index 0000000..0bea4a4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/spark-config.sh @@ -0,0 +1,33 @@ +# +# 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. +# + +# included in all the spark scripts with source command +# should not be executable directly +# also should not be passed any arguments, since we need original $* + +# symlink and absolute path should rely on SPARK_HOME to resolve +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +export SPARK_CONF_DIR="${SPARK_CONF_DIR:-"${SPARK_HOME}/conf"}" +# Add the PySpark classes to the PYTHONPATH: +if [ -z "${PYSPARK_PYTHONPATH_SET}" ]; then + export PYTHONPATH="${SPARK_HOME}/python:${PYTHONPATH}" + export PYTHONPATH="${SPARK_HOME}/python/lib/py4j-0.10.9.7-src.zip:${PYTHONPATH}" + export PYSPARK_PYTHONPATH_SET=1 +fi diff --git a/spark-3.5.0-bin-hadoop3/sbin/spark-daemon.sh b/spark-3.5.0-bin-hadoop3/sbin/spark-daemon.sh new file mode 100644 index 0000000..3cfd5ac --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/spark-daemon.sh @@ -0,0 +1,264 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Runs a Spark command as a daemon. +# +# Environment Variables +# +# SPARK_CONF_DIR Alternate conf dir. Default is ${SPARK_HOME}/conf. +# SPARK_LOG_DIR Where log files are stored. ${SPARK_HOME}/logs by default. +# SPARK_LOG_MAX_FILES Max log files of Spark daemons can rotate to. Default is 5. +# SPARK_MASTER host:path where spark code should be rsync'd from +# SPARK_PID_DIR The pid files are stored. /tmp by default. +# SPARK_IDENT_STRING A string representing this instance of spark. $USER by default +# SPARK_NICENESS The scheduling priority for daemons. Defaults to 0. +# SPARK_NO_DAEMONIZE If set, will run the proposed command in the foreground. It will not output a PID file. +## + +usage="Usage: spark-daemon.sh [--config ] (start|stop|submit|status) " + +# if no args specified, show usage +if [ $# -le 1 ]; then + echo $usage + exit 1 +fi + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +# get arguments + +# Check if --config is passed as an argument. It is an optional parameter. +# Exit if the argument is not a directory. + +if [ "$1" == "--config" ] +then + shift + conf_dir="$1" + if [ ! -d "$conf_dir" ] + then + echo "ERROR : $conf_dir is not a directory" + echo $usage + exit 1 + else + export SPARK_CONF_DIR="$conf_dir" + fi + shift +fi + +option=$1 +shift +command=$1 +shift +instance=$1 +shift + +spark_rotate_log () +{ + log=$1; + + if [[ -z ${SPARK_LOG_MAX_FILES} ]]; then + num=5 + elif [[ ${SPARK_LOG_MAX_FILES} -gt 0 ]]; then + num=${SPARK_LOG_MAX_FILES} + else + echo "Error: SPARK_LOG_MAX_FILES must be a positive number, but got ${SPARK_LOG_MAX_FILES}" + exit -1 + fi + + if [ -f "$log" ]; then # rotate logs + while [ $num -gt 1 ]; do + prev=`expr $num - 1` + [ -f "$log.$prev" ] && mv "$log.$prev" "$log.$num" + num=$prev + done + mv "$log" "$log.$num"; + fi +} + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$SPARK_IDENT_STRING" = "" ]; then + export SPARK_IDENT_STRING="$USER" +fi + + +export SPARK_PRINT_LAUNCH_COMMAND="1" + +# get log directory +if [ "$SPARK_LOG_DIR" = "" ]; then + export SPARK_LOG_DIR="${SPARK_HOME}/logs" +fi +mkdir -p "$SPARK_LOG_DIR" +touch "$SPARK_LOG_DIR"/.spark_test > /dev/null 2>&1 +TEST_LOG_DIR=$? +if [ "${TEST_LOG_DIR}" = "0" ]; then + rm -f "$SPARK_LOG_DIR"/.spark_test +else + chown "$SPARK_IDENT_STRING" "$SPARK_LOG_DIR" +fi + +if [ "$SPARK_PID_DIR" = "" ]; then + SPARK_PID_DIR=/tmp +fi + +# some variables +log="$SPARK_LOG_DIR/spark-$SPARK_IDENT_STRING-$command-$instance-$HOSTNAME.out" +pid="$SPARK_PID_DIR/spark-$SPARK_IDENT_STRING-$command-$instance.pid" + +# Set default scheduling priority +if [ "$SPARK_NICENESS" = "" ]; then + export SPARK_NICENESS=0 +fi + +execute_command() { + if [ -z ${SPARK_NO_DAEMONIZE+set} ]; then + nohup -- "$@" >> $log 2>&1 < /dev/null & + newpid="$!" + + echo "$newpid" > "$pid" + + # Poll for up to 5 seconds for the java process to start + for i in {1..10} + do + if [[ $(ps -p "$newpid" -o comm=) =~ "java" ]]; then + break + fi + sleep 0.5 + done + + sleep 2 + # Check if the process has died; in that case we'll tail the log so the user can see + if [[ ! $(ps -p "$newpid" -o comm=) =~ "java" ]]; then + echo "failed to launch: $@" + tail -10 "$log" | sed 's/^/ /' + echo "full log in $log" + fi + else + "$@" + fi +} + +run_command() { + mode="$1" + shift + + mkdir -p "$SPARK_PID_DIR" + + if [ -f "$pid" ]; then + TARGET_ID="$(cat "$pid")" + if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then + echo "$command running as process $TARGET_ID. Stop it first." + exit 1 + fi + fi + + if [ "$SPARK_MASTER" != "" ]; then + echo rsync from "$SPARK_MASTER" + rsync -a -e ssh --delete --exclude=.svn --exclude='logs/*' --exclude='contrib/hod/logs/*' "$SPARK_MASTER/" "${SPARK_HOME}" + fi + + spark_rotate_log "$log" + echo "starting $command, logging to $log" + + case "$mode" in + (class) + execute_command nice -n "$SPARK_NICENESS" "${SPARK_HOME}"/bin/spark-class "$command" "$@" + ;; + + (submit) + execute_command nice -n "$SPARK_NICENESS" bash "${SPARK_HOME}"/bin/spark-submit --class "$command" "$@" + ;; + + (*) + echo "unknown mode: $mode" + exit 1 + ;; + esac + +} + +case $option in + + (submit) + run_command submit "$@" + ;; + + (start) + run_command class "$@" + ;; + + (stop) + + if [ -f $pid ]; then + TARGET_ID="$(cat "$pid")" + if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then + echo "stopping $command" + kill "$TARGET_ID" && rm -f "$pid" + else + echo "no $command to stop" + fi + else + echo "no $command to stop" + fi + ;; + + (decommission) + + if [ -f $pid ]; then + TARGET_ID="$(cat "$pid")" + if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then + echo "decommissioning $command" + kill -s SIGPWR "$TARGET_ID" + else + echo "no $command to decommission" + fi + else + echo "no $command to decommission" + fi + ;; + + (status) + + if [ -f $pid ]; then + TARGET_ID="$(cat "$pid")" + if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then + echo $command is running. + exit 0 + else + echo $pid file is present but $command not running + exit 1 + fi + else + echo $command not running. + exit 2 + fi + ;; + + (*) + echo $usage + exit 1 + ;; + +esac + + diff --git a/spark-3.5.0-bin-hadoop3/sbin/spark-daemons.sh b/spark-3.5.0-bin-hadoop3/sbin/spark-daemons.sh new file mode 100644 index 0000000..9a5e5f3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/spark-daemons.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Run a Spark command on all worker hosts. + +usage="Usage: spark-daemons.sh [--config ] [start|stop] command instance-number args..." + +# if no args specified, show usage +if [ $# -le 1 ]; then + echo $usage + exit 1 +fi + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +exec "${SPARK_HOME}/sbin/workers.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin/spark-daemon.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-all.sh b/spark-3.5.0-bin-hadoop3/sbin/start-all.sh new file mode 100644 index 0000000..064074e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-all.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Start all spark daemons. +# Starts the master on this node. +# Starts a worker on each node specified in conf/workers + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# Load the Spark configuration +. "${SPARK_HOME}/sbin/spark-config.sh" + +# Start Master +"${SPARK_HOME}/sbin"/start-master.sh + +# Start Workers +"${SPARK_HOME}/sbin"/start-workers.sh diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-connect-server.sh b/spark-3.5.0-bin-hadoop3/sbin/start-connect-server.sh new file mode 100644 index 0000000..a347f43 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-connect-server.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Enter posix mode for bash +set -o posix + +# Shell script for starting the Spark Connect server +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# NOTE: This exact class name is matched downstream by SparkSubmit. +# Any changes need to be reflected there. +CLASS="org.apache.spark.sql.connect.service.SparkConnectServer" + +if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then + echo "Usage: ./sbin/start-connect-server.sh [options]" + + "${SPARK_HOME}"/bin/spark-submit --help 2>&1 | grep -v Usage 1>&2 + exit 1 +fi + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +exec "${SPARK_HOME}"/sbin/spark-daemon.sh submit $CLASS 1 --name "Spark Connect server" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-history-server.sh b/spark-3.5.0-bin-hadoop3/sbin/start-history-server.sh new file mode 100644 index 0000000..71dace4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-history-server.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Starts the history server on the machine this script is executed on. +# +# Usage: start-history-server.sh +# +# Use the SPARK_HISTORY_OPTS environment variable to set history server configuration. +# + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# NOTE: This exact class name is matched downstream by SparkSubmit. +# Any changes need to be reflected there. +CLASS="org.apache.spark.deploy.history.HistoryServer" + +if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then + echo "Usage: ./sbin/start-history-server.sh [options]" + pattern="Usage:" + pattern+="\|Using Spark's default log4j profile:" + pattern+="\|Started daemon with process name" + pattern+="\|Registered signal handler for" + + "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2 + exit 1 +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" +. "${SPARK_HOME}/bin/load-spark-env.sh" + +exec "${SPARK_HOME}/sbin"/spark-daemon.sh start $CLASS 1 "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-master.sh b/spark-3.5.0-bin-hadoop3/sbin/start-master.sh new file mode 100644 index 0000000..36fe4b4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-master.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Starts the master on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# NOTE: This exact class name is matched downstream by SparkSubmit. +# Any changes need to be reflected there. +CLASS="org.apache.spark.deploy.master.Master" + +if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then + echo "Usage: ./sbin/start-master.sh [options]" + pattern="Usage:" + pattern+="\|Using Spark's default log4j profile:" + pattern+="\|Started daemon with process name" + pattern+="\|Registered signal handler for" + + "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2 + exit 1 +fi + +ORIGINAL_ARGS="$@" + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$SPARK_MASTER_PORT" = "" ]; then + SPARK_MASTER_PORT=7077 +fi + +if [ "$SPARK_MASTER_HOST" = "" ]; then + case `uname` in + (SunOS) + SPARK_MASTER_HOST="`/usr/sbin/check-hostname | awk '{print $NF}'`" + ;; + (*) + SPARK_MASTER_HOST="`hostname -f`" + ;; + esac +fi + +if [ "$SPARK_MASTER_WEBUI_PORT" = "" ]; then + SPARK_MASTER_WEBUI_PORT=8080 +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh start $CLASS 1 \ + --host $SPARK_MASTER_HOST --port $SPARK_MASTER_PORT --webui-port $SPARK_MASTER_WEBUI_PORT \ + $ORIGINAL_ARGS diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-mesos-dispatcher.sh b/spark-3.5.0-bin-hadoop3/sbin/start-mesos-dispatcher.sh new file mode 100644 index 0000000..c2e30d8 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-mesos-dispatcher.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash + +# +# 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. +# +# Starts the Mesos Cluster Dispatcher on the machine this script is executed on. +# The Mesos Cluster Dispatcher is responsible for launching the Mesos framework and +# Rest server to handle driver requests for Mesos cluster mode. +# Only one cluster dispatcher is needed per Mesos cluster. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$SPARK_MESOS_DISPATCHER_PORT" = "" ]; then + SPARK_MESOS_DISPATCHER_PORT=7077 +fi + +if [ "$SPARK_MESOS_DISPATCHER_HOST" = "" ]; then + case `uname` in + (SunOS) + SPARK_MESOS_DISPATCHER_HOST="`/usr/sbin/check-hostname | awk '{print $NF}'`" + ;; + (*) + SPARK_MESOS_DISPATCHER_HOST="`hostname -f`" + ;; + esac +fi + +if [ "$SPARK_MESOS_DISPATCHER_NUM" = "" ]; then + SPARK_MESOS_DISPATCHER_NUM=1 +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosClusterDispatcher $SPARK_MESOS_DISPATCHER_NUM --host $SPARK_MESOS_DISPATCHER_HOST --port $SPARK_MESOS_DISPATCHER_PORT "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-mesos-shuffle-service.sh b/spark-3.5.0-bin-hadoop3/sbin/start-mesos-shuffle-service.sh new file mode 100644 index 0000000..7b7dc9c --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-mesos-shuffle-service.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Starts the Mesos external shuffle service on the machine this script is executed on. +# The Mesos external shuffle service detects when an application exits and automatically +# cleans up its shuffle files. +# +# Usage: start-mesos-shuffle-service.sh +# +# Use the SPARK_SHUFFLE_OPTS environment variable to set shuffle service configuration. +# + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" +. "${SPARK_HOME}/bin/load-spark-env.sh" + +exec "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosExternalShuffleService 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-slave.sh b/spark-3.5.0-bin-hadoop3/sbin/start-slave.sh new file mode 100644 index 0000000..6868253 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-slave.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use start-worker.sh" +"${DIR}/start-worker.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-slaves.sh b/spark-3.5.0-bin-hadoop3/sbin/start-slaves.sh new file mode 100644 index 0000000..9b113d9 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-slaves.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use start-workers.sh" +"${DIR}/start-workers.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-thriftserver.sh b/spark-3.5.0-bin-hadoop3/sbin/start-thriftserver.sh new file mode 100644 index 0000000..b1d3871 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-thriftserver.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# +# Shell script for starting the Spark SQL Thrift server + +# Enter posix mode for bash +set -o posix + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# NOTE: This exact class name is matched downstream by SparkSubmit. +# Any changes need to be reflected there. +CLASS="org.apache.spark.sql.hive.thriftserver.HiveThriftServer2" + +function usage { + echo "Usage: ./sbin/start-thriftserver [options] [thrift server options]" + pattern="usage" + pattern+="\|Spark assembly has been built with Hive" + pattern+="\|NOTE: SPARK_PREPEND_CLASSES is set" + pattern+="\|Spark Command: " + pattern+="\|=======" + pattern+="\|--help" + pattern+="\|Using Spark's default log4j profile:" + pattern+="\|^log4j:" + pattern+="\|Started daemon with process name" + pattern+="\|Registered signal handler for" + + "${SPARK_HOME}"/bin/spark-submit --help 2>&1 | grep -v Usage 1>&2 + echo + echo "Thrift server options:" + "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2 +} + +if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then + usage + exit 1 +fi + +export SUBMIT_USAGE_FUNCTION=usage + +exec "${SPARK_HOME}"/sbin/spark-daemon.sh submit $CLASS 1 --name "Thrift JDBC/ODBC Server" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-worker.sh b/spark-3.5.0-bin-hadoop3/sbin/start-worker.sh new file mode 100644 index 0000000..fd58f01 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-worker.sh @@ -0,0 +1,92 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Starts a worker on the machine this script is executed on. +# +# Environment Variables +# +# SPARK_WORKER_INSTANCES The number of worker instances to run on this +# worker. Default is 1. Note it has been deprecate since Spark 3.0. +# SPARK_WORKER_PORT The base port number for the first worker. If set, +# subsequent workers will increment this number. If +# unset, Spark will find a valid port number, but +# with no guarantee of a predictable pattern. +# SPARK_WORKER_WEBUI_PORT The base port for the web interface of the first +# worker. Subsequent workers will increment this +# number. Default is 8081. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# NOTE: This exact class name is matched downstream by SparkSubmit. +# Any changes need to be reflected there. +CLASS="org.apache.spark.deploy.worker.Worker" + +if [[ $# -lt 1 ]] || [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then + echo "Usage: ./sbin/start-worker.sh [options]" + pattern="Usage:" + pattern+="\|Using Spark's default log4j profile:" + pattern+="\|Started daemon with process name" + pattern+="\|Registered signal handler for" + + "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2 + exit 1 +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +# First argument should be the master; we need to store it aside because we may +# need to insert arguments between it and the other arguments +MASTER=$1 +shift + +# Determine desired worker port +if [ "$SPARK_WORKER_WEBUI_PORT" = "" ]; then + SPARK_WORKER_WEBUI_PORT=8081 +fi + +# Start up the appropriate number of workers on this machine. +# quick local function to start a worker +function start_instance { + WORKER_NUM=$1 + shift + + if [ "$SPARK_WORKER_PORT" = "" ]; then + PORT_FLAG= + PORT_NUM= + else + PORT_FLAG="--port" + PORT_NUM=$(( $SPARK_WORKER_PORT + $WORKER_NUM - 1 )) + fi + WEBUI_PORT=$(( $SPARK_WORKER_WEBUI_PORT + $WORKER_NUM - 1 )) + + "${SPARK_HOME}/sbin"/spark-daemon.sh start $CLASS $WORKER_NUM \ + --webui-port "$WEBUI_PORT" $PORT_FLAG $PORT_NUM $MASTER "$@" +} + +if [ "$SPARK_WORKER_INSTANCES" = "" ]; then + start_instance 1 "$@" +else + for ((i=0; i<$SPARK_WORKER_INSTANCES; i++)); do + start_instance $(( 1 + $i )) "$@" + done +fi diff --git a/spark-3.5.0-bin-hadoop3/sbin/start-workers.sh b/spark-3.5.0-bin-hadoop3/sbin/start-workers.sh new file mode 100644 index 0000000..3867ef3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/start-workers.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Starts a worker instance on each machine specified in the conf/workers file. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" +. "${SPARK_HOME}/bin/load-spark-env.sh" + +# Find the port number for the master +if [ "$SPARK_MASTER_PORT" = "" ]; then + SPARK_MASTER_PORT=7077 +fi + +if [ "$SPARK_MASTER_HOST" = "" ]; then + case `uname` in + (SunOS) + SPARK_MASTER_HOST="`/usr/sbin/check-hostname | awk '{print $NF}'`" + ;; + (*) + SPARK_MASTER_HOST="`hostname -f`" + ;; + esac +fi + +# Launch the workers +"${SPARK_HOME}/sbin/workers.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin/start-worker.sh" "spark://$SPARK_MASTER_HOST:$SPARK_MASTER_PORT" diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-all.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-all.sh new file mode 100644 index 0000000..2c40905 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-all.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stop all spark daemons. +# Run this on the master node. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +# Load the Spark configuration +. "${SPARK_HOME}/sbin/spark-config.sh" + +# Stop the workers, then the master +"${SPARK_HOME}/sbin"/stop-workers.sh +"${SPARK_HOME}/sbin"/stop-master.sh + +if [ "$1" == "--wait" ] +then + printf "Waiting for workers to shut down..." + while true + do + running=`${SPARK_HOME}/sbin/workers.sh ps -ef | grep -v grep | grep deploy.worker.Worker` + if [ -z "$running" ] + then + printf "\nAll workers successfully shut down.\n" + break + else + printf "." + sleep 10 + fi + done +fi diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-connect-server.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-connect-server.sh new file mode 100644 index 0000000..7cf7440 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-connect-server.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stops the connect server on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.sql.connect.service.SparkConnectServer 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-history-server.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-history-server.sh new file mode 100644 index 0000000..14e3af4 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-history-server.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stops the history server on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +"${SPARK_HOME}/sbin/spark-daemon.sh" stop org.apache.spark.deploy.history.HistoryServer 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-master.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-master.sh new file mode 100644 index 0000000..14644ea --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-master.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stops the master on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.master.Master 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-dispatcher.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-dispatcher.sh new file mode 100644 index 0000000..b13e018 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-dispatcher.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash + +# +# 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. +# +# Stop the Mesos Cluster dispatcher on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +if [ "$SPARK_MESOS_DISPATCHER_NUM" = "" ]; then + SPARK_MESOS_DISPATCHER_NUM=1 +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosClusterDispatcher \ + $SPARK_MESOS_DISPATCHER_NUM + diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-shuffle-service.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-shuffle-service.sh new file mode 100644 index 0000000..d23cad3 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-mesos-shuffle-service.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stops the Mesos external shuffle service on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosExternalShuffleService 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-slave.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-slave.sh new file mode 100644 index 0000000..71ed299 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-slave.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use stop-worker.sh" +"${DIR}/stop-worker.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-slaves.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-slaves.sh new file mode 100644 index 0000000..c0aca68 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-slaves.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# +# 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. +# + +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +>&2 echo "This script is deprecated, use stop-workers.sh" +"${DIR}/stop-workers.sh" "$@" diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-thriftserver.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-thriftserver.sh new file mode 100644 index 0000000..cf45058 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-thriftserver.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Stops the thrift server on the machine this script is executed on. + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.sql.hive.thriftserver.HiveThriftServer2 1 diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-worker.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-worker.sh new file mode 100644 index 0000000..112b62e --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-worker.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# A shell script to stop all workers on a single worker +# +# Environment variables +# +# SPARK_WORKER_INSTANCES The number of worker instances that should be +# running on this worker machine. Default is 1. + +# Usage: stop-worker.sh +# Stops all workers on this worker machine + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$SPARK_WORKER_INSTANCES" = "" ]; then + "${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker 1 +else + for ((i=0; i<$SPARK_WORKER_INSTANCES; i++)); do + "${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker $(( $i + 1 )) + done +fi diff --git a/spark-3.5.0-bin-hadoop3/sbin/stop-workers.sh b/spark-3.5.0-bin-hadoop3/sbin/stop-workers.sh new file mode 100644 index 0000000..552800f --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/stop-workers.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash + +# +# 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. +# + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +"${SPARK_HOME}/sbin/workers.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin"/stop-worker.sh diff --git a/spark-3.5.0-bin-hadoop3/sbin/workers.sh b/spark-3.5.0-bin-hadoop3/sbin/workers.sh new file mode 100644 index 0000000..cab0330 --- /dev/null +++ b/spark-3.5.0-bin-hadoop3/sbin/workers.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash + +# +# 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. +# + +# Run a shell command on all worker hosts. +# +# Environment Variables +# +# SPARK_WORKERS File naming remote hosts. +# Default is ${SPARK_CONF_DIR}/workers. +# SPARK_CONF_DIR Alternate conf dir. Default is ${SPARK_HOME}/conf. +# SPARK_WORKER_SLEEP Seconds to sleep between spawning remote commands. +# SPARK_SSH_OPTS Options passed to ssh when running remote commands. +## + +usage="Usage: workers.sh [--config ] command..." + +# if no args specified, show usage +if [ $# -le 0 ]; then + echo $usage + exit 1 +fi + +if [ -z "${SPARK_HOME}" ]; then + export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)" +fi + +. "${SPARK_HOME}/sbin/spark-config.sh" + +# If the workers file is specified in the command line, +# then it takes precedence over the definition in +# spark-env.sh. Save it here. +if [ -f "$SPARK_WORKERS" ]; then + HOSTLIST=`cat "$SPARK_WORKERS"` +fi +if [ -f "$SPARK_SLAVES" ]; then + >&2 echo "SPARK_SLAVES is deprecated, use SPARK_WORKERS" + HOSTLIST=`cat "$SPARK_SLAVES"` +fi + + +# Check if --config is passed as an argument. It is an optional parameter. +# Exit if the argument is not a directory. +if [ "$1" == "--config" ] +then + shift + conf_dir="$1" + if [ ! -d "$conf_dir" ] + then + echo "ERROR : $conf_dir is not a directory" + echo $usage + exit 1 + else + export SPARK_CONF_DIR="$conf_dir" + fi + shift +fi + +. "${SPARK_HOME}/bin/load-spark-env.sh" + +if [ "$HOSTLIST" = "" ]; then + if [ "$SPARK_SLAVES" = "" ] && [ "$SPARK_WORKERS" = "" ]; then + if [ -f "${SPARK_CONF_DIR}/workers" ]; then + HOSTLIST=`cat "${SPARK_CONF_DIR}/workers"` + elif [ -f "${SPARK_CONF_DIR}/slaves" ]; then + HOSTLIST=`cat "${SPARK_CONF_DIR}/slaves"` + else + HOSTLIST=localhost + fi + else + if [ -f "$SPARK_WORKERS" ]; then + HOSTLIST=`cat "$SPARK_WORKERS"` + fi + if [ -f "$SPARK_SLAVES" ]; then + >&2 echo "SPARK_SLAVES is deprecated, use SPARK_WORKERS" + HOSTLIST=`cat "$SPARK_SLAVES"` + fi + fi +fi + + + +# By default disable strict host key checking +if [ "$SPARK_SSH_OPTS" = "" ]; then + SPARK_SSH_OPTS="-o StrictHostKeyChecking=no" +fi + +for host in `echo "$HOSTLIST"|sed "s/#.*$//;/^$/d"`; do + if [ -n "${SPARK_SSH_FOREGROUND}" ]; then + ssh $SPARK_SSH_OPTS "$host" $"${@// /\\ }" \ + 2>&1 | sed "s/^/$host: /" + else + ssh $SPARK_SSH_OPTS "$host" $"${@// /\\ }" \ + 2>&1 | sed "s/^/$host: /" & + fi + if [ "$SPARK_WORKER_SLEEP" != "" ]; then + sleep $SPARK_WORKER_SLEEP + fi + if [ "$SPARK_SLAVE_SLEEP" != "" ]; then + >&2 echo "SPARK_SLAVE_SLEEP is deprecated, use SPARK_WORKER_SLEEP" + sleep $SPARK_SLAVE_SLEEP + fi +done + +wait diff --git a/spark-3.5.0-bin-hadoop3/yarn/spark-3.5.0-yarn-shuffle.jar b/spark-3.5.0-bin-hadoop3/yarn/spark-3.5.0-yarn-shuffle.jar new file mode 100644 index 0000000..7255a72 Binary files /dev/null and b/spark-3.5.0-bin-hadoop3/yarn/spark-3.5.0-yarn-shuffle.jar differ